Compare commits
35
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
07782ff2fc | ||
|
|
0295b4555c | ||
|
|
702e82e422 | ||
|
|
c878b43b64 | ||
|
|
2b84ac669c | ||
|
|
11d3976b88 | ||
|
|
40da9625a1 | ||
|
|
8d9babd4de | ||
|
|
e99ba957ec | ||
|
|
64ac1395e8 | ||
|
|
61cf087680 | ||
|
|
847a57cd12 | ||
|
|
fcd6ac97ed | ||
|
|
95be2a7f22 | ||
|
|
0e60c925cf | ||
|
|
d7ff22204a | ||
|
|
c0bd8b13da | ||
|
|
caeb887bf6 | ||
|
|
6b3166a7c7 | ||
|
|
25e2e136ef | ||
|
|
6874638bc4 | ||
|
|
e24663c5a9 | ||
|
|
c50e105a88 | ||
|
|
a766b30349 | ||
|
|
1faa8cb73c | ||
|
|
e89a91d927 | ||
|
|
909b147197 | ||
|
|
a88b3be7c4 | ||
|
|
a49ea5a58f | ||
|
|
30ebe0dc3c | ||
|
|
cef65f0715 | ||
|
|
6f3b2047ab | ||
|
|
02e8f26cea | ||
|
|
4a00a511bb | ||
|
|
a0d8d944e2 |
@@ -8,6 +8,21 @@ steps:
|
||||
- group: "Build Python wheels"
|
||||
key: "build-wheels"
|
||||
steps:
|
||||
- label: "Build wheel - aarch64 - CUDA 12.8"
|
||||
depends_on: ~
|
||||
id: build-wheel-arm64-cuda-12-8
|
||||
agents:
|
||||
queue: arm64_cpu_queue_postmerge
|
||||
commands:
|
||||
# #NOTE: torch_cuda_arch_list is derived from upstream PyTorch build files here:
|
||||
# https://github.com/pytorch/pytorch/blob/main/.ci/aarch64_linux/aarch64_ci_build.sh#L7
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.8.1 --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX' --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
|
||||
- "mkdir artifacts"
|
||||
- "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
|
||||
- "bash .buildkite/scripts/upload-nightly-wheels.sh"
|
||||
env:
|
||||
DOCKER_BUILDKIT: "1"
|
||||
|
||||
- label: "Build wheel - aarch64 - CUDA 12.9"
|
||||
depends_on: ~
|
||||
id: build-wheel-arm64-cuda-12-9
|
||||
@@ -51,6 +66,19 @@ steps:
|
||||
env:
|
||||
DOCKER_BUILDKIT: "1"
|
||||
|
||||
- label: "Build wheel - x86_64 - CUDA 12.8"
|
||||
depends_on: ~
|
||||
id: build-wheel-x86-cuda-12-8
|
||||
agents:
|
||||
queue: cpu_queue_postmerge
|
||||
commands:
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.8.1 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
|
||||
- "mkdir artifacts"
|
||||
- "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
|
||||
- "bash .buildkite/scripts/upload-nightly-wheels.sh manylinux_2_31"
|
||||
env:
|
||||
DOCKER_BUILDKIT: "1"
|
||||
|
||||
- label: "Build wheel - x86_64 - CUDA 12.9"
|
||||
depends_on: ~
|
||||
id: build-wheel-x86-cuda-12-9
|
||||
|
||||
@@ -67,7 +67,7 @@ start_nodes() {
|
||||
# 3. map the huggingface cache directory to the container
|
||||
# 3. assign ip addresses to the containers (head node: 192.168.10.10, worker nodes:
|
||||
# starting from 192.168.10.11)
|
||||
docker run -d "$GPU_DEVICES" --shm-size=10.24gb -e HF_TOKEN \
|
||||
docker run -d $GPU_DEVICES --shm-size=10.24gb -e HF_TOKEN \
|
||||
-v ~/.cache/huggingface:/root/.cache/huggingface --name "node$node" \
|
||||
--network docker-net --ip 192.168.10.$((10 + $node)) --rm "$DOCKER_IMAGE" \
|
||||
/bin/bash -c "tail -f /dev/null"
|
||||
|
||||
@@ -55,9 +55,11 @@ steps:
|
||||
grade: Blocking
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/detokenizer
|
||||
- tests/multimodal
|
||||
- tests/utils_
|
||||
commands:
|
||||
- pytest -v -s detokenizer
|
||||
- pytest -v -s -m 'not cpu_test' multimodal
|
||||
- pytest -v -s utils_
|
||||
|
||||
@@ -547,7 +549,7 @@ steps:
|
||||
- tests/samplers
|
||||
- tests/conftest.py
|
||||
commands:
|
||||
- pytest -v -s -m 'not skip_v1' samplers
|
||||
- pytest -v -s -m samplers
|
||||
|
||||
- label: LoRA Test %N # 20min each
|
||||
timeout_in_minutes: 30
|
||||
@@ -2213,7 +2215,7 @@ steps:
|
||||
- tests/samplers
|
||||
- tests/conftest.py
|
||||
commands:
|
||||
- pytest -v -s -m 'not skip_v1' samplers
|
||||
- pytest -v -s -m samplers
|
||||
|
||||
- label: LoRA Test %N # 20min each
|
||||
timeout_in_minutes: 30
|
||||
|
||||
@@ -51,9 +51,11 @@ steps:
|
||||
mirror_hardwares: [amdexperimental]
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/detokenizer
|
||||
- tests/multimodal
|
||||
- tests/utils_
|
||||
commands:
|
||||
- pytest -v -s detokenizer
|
||||
- pytest -v -s -m 'not cpu_test' multimodal
|
||||
- pytest -v -s utils_
|
||||
|
||||
|
||||
@@ -165,6 +165,7 @@ steps:
|
||||
num_devices: 2
|
||||
num_nodes: 2
|
||||
no_plugin: true
|
||||
optional: true # TODO: revert once infra issue solved
|
||||
source_file_dependencies:
|
||||
- vllm/distributed/
|
||||
- vllm/engine/
|
||||
|
||||
@@ -108,9 +108,11 @@ steps:
|
||||
timeout_in_minutes: 50
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/detokenizer
|
||||
- tests/multimodal
|
||||
- tests/utils_
|
||||
commands:
|
||||
- pytest -v -s detokenizer
|
||||
- pytest -v -s -m 'not cpu_test' multimodal
|
||||
- pytest -v -s utils_
|
||||
|
||||
|
||||
@@ -18,4 +18,4 @@ steps:
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
commands:
|
||||
- pytest -v -s -m 'not skip_v1' samplers
|
||||
- pytest -v -s -m samplers
|
||||
|
||||
@@ -771,6 +771,25 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
|
||||
endif()
|
||||
endif()
|
||||
|
||||
# DeepSeek V3 fused A GEMM kernel (requires SM 9.0+, Hopper and later)
|
||||
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
|
||||
cuda_archs_loose_intersection(DSV3_FUSED_A_GEMM_ARCHS "9.0a;10.0f;11.0f" "${CUDA_ARCHS}")
|
||||
else()
|
||||
cuda_archs_loose_intersection(DSV3_FUSED_A_GEMM_ARCHS "9.0a;10.0a;10.1a;10.3a" "${CUDA_ARCHS}")
|
||||
endif()
|
||||
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND DSV3_FUSED_A_GEMM_ARCHS)
|
||||
set(DSV3_FUSED_A_GEMM_SRC "csrc/dsv3_fused_a_gemm.cu")
|
||||
set_gencode_flags_for_srcs(
|
||||
SRCS "${DSV3_FUSED_A_GEMM_SRC}"
|
||||
CUDA_ARCHS "${DSV3_FUSED_A_GEMM_ARCHS}")
|
||||
list(APPEND VLLM_EXT_SRC ${DSV3_FUSED_A_GEMM_SRC})
|
||||
list(APPEND VLLM_GPU_FLAGS "-DENABLE_DSV3_FUSED_A_GEMM=1")
|
||||
message(STATUS "Building dsv3_fused_a_gemm for archs: ${DSV3_FUSED_A_GEMM_ARCHS}")
|
||||
else()
|
||||
message(STATUS "Not building dsv3_fused_a_gemm as no compatible archs found "
|
||||
"in CUDA target architectures.")
|
||||
endif()
|
||||
|
||||
# moe_data.cu is used by all CUTLASS MoE kernels.
|
||||
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
|
||||
cuda_archs_loose_intersection(CUTLASS_MOE_DATA_ARCHS "9.0a;10.0f;11.0f;12.0f" "${CUDA_ARCHS}")
|
||||
|
||||
@@ -19,7 +19,7 @@ else()
|
||||
FetchContent_Declare(
|
||||
flashmla
|
||||
GIT_REPOSITORY https://github.com/vllm-project/FlashMLA
|
||||
GIT_TAG 692917b1cda61b93ac9ee2d846ec54e75afe87b1
|
||||
GIT_TAG f46984b4caf19409bf4ab94a540341918dd6a1a3
|
||||
GIT_PROGRESS TRUE
|
||||
CONFIGURE_COMMAND ""
|
||||
BUILD_COMMAND ""
|
||||
|
||||
@@ -13,8 +13,13 @@ struct alignas(32) u32x8_t {
|
||||
uint32_t u0, u1, u2, u3, u4, u5, u6, u7;
|
||||
};
|
||||
|
||||
// 256-bit vector loads require both sm_100+ and CUDA 12.9+
|
||||
#define VLLM_SUPPORTS_256BIT_VECTORS \
|
||||
(defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 && \
|
||||
__CUDACC_VER_MAJOR__ >= 12 && __CUDACC_VER_MINOR__ >= 9)
|
||||
|
||||
__device__ __forceinline__ void ld256(u32x8_t& val, const u32x8_t* ptr) {
|
||||
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000
|
||||
#if VLLM_SUPPORTS_256BIT_VECTORS
|
||||
asm volatile("ld.global.nc.v8.u32 {%0,%1,%2,%3,%4,%5,%6,%7}, [%8];\n"
|
||||
: "=r"(val.u0), "=r"(val.u1), "=r"(val.u2), "=r"(val.u3),
|
||||
"=r"(val.u4), "=r"(val.u5), "=r"(val.u6), "=r"(val.u7)
|
||||
@@ -35,7 +40,7 @@ __device__ __forceinline__ void ld256(u32x8_t& val, const u32x8_t* ptr) {
|
||||
}
|
||||
|
||||
__device__ __forceinline__ void st256(u32x8_t& val, u32x8_t* ptr) {
|
||||
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000
|
||||
#if VLLM_SUPPORTS_256BIT_VECTORS
|
||||
asm volatile("st.global.v8.u32 [%0], {%1,%2,%3,%4,%5,%6,%7,%8};\n"
|
||||
:
|
||||
: "l"(ptr), "r"(val.u0), "r"(val.u1), "r"(val.u2), "r"(val.u3),
|
||||
|
||||
@@ -1305,7 +1305,8 @@ void indexer_k_quant_and_cache(
|
||||
const at::cuda::OptionalCUDAGuard device_guard(device_of(k));
|
||||
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
|
||||
|
||||
DISPATCH_BY_KV_CACHE_DTYPE(k.dtype(), "fp8_e4m3",
|
||||
static const std::string kv_cache_dtype = "fp8_e4m3";
|
||||
DISPATCH_BY_KV_CACHE_DTYPE(k.dtype(), kv_cache_dtype,
|
||||
CALL_INDEXER_K_QUANT_AND_CACHE);
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,747 @@
|
||||
/*
|
||||
* Adapted from
|
||||
* https://github.com/sgl-project/sglang/blob/main/sgl-kernel/csrc/gemm/dsv3_fused_a_gemm.cu
|
||||
* which was adapted from
|
||||
* https://github.com/NVIDIA/TensorRT-LLM/blob/619709fc33bd5dc268f19d6a741fe7ed51c0f8f5/cpp/tensorrt_llm/kernels/dsv3MinLatencyKernels/dsv3FusedAGemm.cu
|
||||
*
|
||||
* Copyright (c) 2019-2024, NVIDIA CORPORATION. All rights reserved.
|
||||
* Copyright (c) 2021, NAVER Corp. Authored by CLOVA.
|
||||
*
|
||||
* 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
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* 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/ATen.h>
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
#include <cuda_bf16.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <torch/all.h>
|
||||
|
||||
#include "core/registration.h"
|
||||
|
||||
#include <cstdlib>
|
||||
#include <mutex>
|
||||
|
||||
namespace {
|
||||
|
||||
inline int getSMVersion() {
|
||||
auto* props = at::cuda::getCurrentDeviceProperties();
|
||||
return props->major * 10 + props->minor;
|
||||
}
|
||||
|
||||
inline bool getEnvEnablePDL() {
|
||||
static std::once_flag flag;
|
||||
static bool enablePDL = false;
|
||||
std::call_once(flag, [&]() {
|
||||
if (getSMVersion() >= 90) {
|
||||
char const* env = std::getenv("TRTLLM_ENABLE_PDL");
|
||||
enablePDL = env && env[0] == '1' && env[1] == '\0';
|
||||
}
|
||||
});
|
||||
return enablePDL;
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
using bf16_t = __nv_bfloat16;
|
||||
|
||||
__device__ void hmma_16_8_16_f32acc_bf16ab(float (&d_reg)[4],
|
||||
const bf16_t (&a_reg)[8],
|
||||
const bf16_t (&b_reg)[4],
|
||||
float const (&c_reg)[4]) {
|
||||
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
|
||||
uint32_t a0 = *reinterpret_cast<uint32_t const*>(a_reg + 0);
|
||||
uint32_t a1 = *reinterpret_cast<uint32_t const*>(a_reg + 2);
|
||||
uint32_t a2 = *reinterpret_cast<uint32_t const*>(a_reg + 4);
|
||||
uint32_t a3 = *reinterpret_cast<uint32_t const*>(a_reg + 6);
|
||||
uint32_t b0 = *reinterpret_cast<uint32_t const*>(b_reg + 0);
|
||||
uint32_t b1 = *reinterpret_cast<uint32_t const*>(b_reg + 2);
|
||||
asm volatile(
|
||||
"mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 "
|
||||
"{%0, %1, %2, %3},"
|
||||
"{%4, %5, %6, %7},"
|
||||
"{%8, %9},"
|
||||
"{%10, %11, %12, %13};\n"
|
||||
: "=f"(d_reg[0]), "=f"(d_reg[1]), "=f"(d_reg[2]), "=f"(d_reg[3])
|
||||
: "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1), "f"(d_reg[0]),
|
||||
"f"(d_reg[1]), "f"(d_reg[2]), "f"(d_reg[3]));
|
||||
#endif
|
||||
}
|
||||
|
||||
extern "C" {
|
||||
__device__ uint32_t __nvvm_get_smem_pointer(void*);
|
||||
}
|
||||
|
||||
__device__ void ldgsts_128(void const* gPtr, void* sPtr, uint32_t pred) {
|
||||
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
|
||||
if (pred) {
|
||||
uint32_t smemPtrAsUint32 = __nvvm_get_smem_pointer(sPtr);
|
||||
asm volatile("cp.async.cg.shared.global.L2::128B [%0], [%1], %2;\n" ::"r"(
|
||||
smemPtrAsUint32),
|
||||
"l"(gPtr), "n"(16));
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
__device__ void ldsm_x4(void* smem_ptr, uint32_t* reg_ptr) {
|
||||
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
|
||||
asm volatile(
|
||||
"ldmatrix.sync.aligned.x4.m8n8.shared.b16 {%0, %1, %2, %3}, [%4];\n"
|
||||
: "=r"(reg_ptr[0]), "=r"(reg_ptr[1]), "=r"(reg_ptr[2]), "=r"(reg_ptr[3])
|
||||
: "r"(__nvvm_get_smem_pointer(smem_ptr)));
|
||||
#endif
|
||||
}
|
||||
|
||||
template <class Type>
|
||||
__device__ int apply_swizzle_343_on_elem_row_col(int row_idx_, int col_idx_) {
|
||||
uint32_t row_idx = *reinterpret_cast<uint32_t*>(&row_idx_);
|
||||
uint32_t col_idx = *reinterpret_cast<uint32_t*>(&col_idx_);
|
||||
row_idx = row_idx % 8;
|
||||
row_idx = row_idx * (16 / sizeof(Type));
|
||||
col_idx = col_idx ^ row_idx;
|
||||
return *reinterpret_cast<int*>(&col_idx);
|
||||
}
|
||||
|
||||
__device__ void initialize_barrier(
|
||||
uint64_t* smem_barrier, // 64 bits user-manged barrier in smem
|
||||
int thread_count =
|
||||
1) // Thread count expected to arrive/wait on this barrier
|
||||
{
|
||||
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
|
||||
uint32_t smem_int_ptr = __nvvm_get_smem_pointer(smem_barrier);
|
||||
asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;\n" ::"r"(smem_int_ptr),
|
||||
"r"(thread_count));
|
||||
#endif
|
||||
}
|
||||
|
||||
// Barrier wait
|
||||
__device__ void wait_barrier(
|
||||
uint64_t* smem_barrier, // 64 bits user-manged barrier in smem
|
||||
int phase_bit) // Current phase bit the barrier waiting to flip
|
||||
{
|
||||
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
|
||||
uint32_t smem_int_ptr = __nvvm_get_smem_pointer(smem_barrier);
|
||||
asm volatile(
|
||||
"{\n"
|
||||
".reg .pred P1;\n"
|
||||
"LAB_WAIT:\n"
|
||||
"mbarrier.try_wait.parity.shared::cta.b64 P1, [%0], %1;\n"
|
||||
"@P1 bra DONE;\n"
|
||||
"bra LAB_WAIT;\n"
|
||||
"DONE:\n"
|
||||
"}\n" ::"r"(smem_int_ptr),
|
||||
"r"(phase_bit));
|
||||
#endif
|
||||
}
|
||||
|
||||
__device__ bool try_wait_barrier(uint64_t* smem_ptr, int phase_bit) {
|
||||
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
|
||||
uint32_t wait_complete;
|
||||
uint32_t smem_int_ptr = __nvvm_get_smem_pointer(smem_ptr);
|
||||
asm volatile(
|
||||
"{\n\t"
|
||||
".reg .pred P1; \n\t"
|
||||
"mbarrier.try_wait.parity.shared::cta.b64 P1, [%1], %2; \n\t"
|
||||
"selp.b32 %0, 1, 0, P1; \n\t"
|
||||
"}"
|
||||
: "=r"(wait_complete)
|
||||
: "r"(smem_int_ptr), "r"(phase_bit));
|
||||
return static_cast<bool>(wait_complete);
|
||||
#endif
|
||||
return false;
|
||||
}
|
||||
|
||||
// Barrier arrive
|
||||
__device__ void arrive_barrier(
|
||||
uint64_t* smem_barrier) // 64 bits user-manged barrier in smem
|
||||
{
|
||||
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
|
||||
uint32_t smem_int_ptr = __nvvm_get_smem_pointer(smem_barrier);
|
||||
asm volatile(
|
||||
"{\n"
|
||||
".reg .b64 state; \n"
|
||||
"mbarrier.arrive.shared::cta.b64 state, [%0];\n"
|
||||
"}\n" ::"r"(smem_int_ptr));
|
||||
#endif
|
||||
}
|
||||
|
||||
__device__ void ldgsts_arrive(uint64_t* smem_barrier) {
|
||||
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
|
||||
uint32_t smem_int_ptr = __nvvm_get_smem_pointer(smem_barrier);
|
||||
asm volatile("cp.async.mbarrier.arrive.noinc.shared.b64 [%0];"
|
||||
:
|
||||
: "r"(smem_int_ptr));
|
||||
#endif
|
||||
}
|
||||
|
||||
template <int gemm_k, int tile_m, int tile_k, int stage_cnt>
|
||||
struct GmemLoaderA {
|
||||
static constexpr int elem_bytes = 2;
|
||||
static constexpr int vec_bytes = 16;
|
||||
static constexpr int vec_elems = vec_bytes / elem_bytes;
|
||||
static constexpr int thread_cnt = 64;
|
||||
static_assert((tile_m * tile_k) % (vec_elems * thread_cnt) == 0);
|
||||
static constexpr int a_inst_cnt_per_iter =
|
||||
(tile_m * tile_k) / (vec_elems * thread_cnt);
|
||||
static_assert(gemm_k % tile_k == 0);
|
||||
static constexpr int k_iter_cnt = gemm_k / tile_k;
|
||||
|
||||
// Extra params to keep the order of k reduction...
|
||||
static constexpr int mma_warp_cnt = 4;
|
||||
static constexpr int per_mma_warp_k = tile_k / mma_warp_cnt;
|
||||
static constexpr int k_each_chunk = gemm_k / mma_warp_cnt;
|
||||
|
||||
private:
|
||||
__device__ int k_project(int tile_k_idx) {
|
||||
return (tile_k_idx / per_mma_warp_k * k_each_chunk) +
|
||||
(tile_k_idx % per_mma_warp_k);
|
||||
}
|
||||
|
||||
public:
|
||||
__device__ GmemLoaderA(bf16_t const* gmem_a_local_, bf16_t* smem_a_,
|
||||
uint64_t* smem_barrier_)
|
||||
: gmem_a(gmem_a_local_),
|
||||
smem_a(smem_a_),
|
||||
smem_barrier(smem_barrier_),
|
||||
local_tid(threadIdx.x % thread_cnt) {}
|
||||
|
||||
__device__ void prepare() {
|
||||
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
|
||||
// swizzle, that's what we want.
|
||||
#pragma unroll
|
||||
for (int i = 0; i < a_inst_cnt_per_iter; i++) {
|
||||
int linear_idx = local_tid * vec_elems + i * thread_cnt * vec_elems;
|
||||
int m_idx = linear_idx / tile_k;
|
||||
int k_idx = linear_idx % tile_k;
|
||||
k_idx = apply_swizzle_343_on_elem_row_col<bf16_t>(m_idx, k_idx);
|
||||
a_smem_offsets[i] = m_idx * tile_k + k_idx;
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
__device__ void issue_mainloop() {
|
||||
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
|
||||
#pragma unroll 1
|
||||
for (int loop_idx = 0; loop_idx < k_iter_cnt; loop_idx++) {
|
||||
if (need_wait) {
|
||||
wait_barrier(smem_barrier + 1 + stage_idx * 2, phase_bit);
|
||||
}
|
||||
int next_stage_idx = stage_idx + 1;
|
||||
int next_phase_bit =
|
||||
next_stage_idx == stage_cnt ? phase_bit ^ 1 : phase_bit;
|
||||
next_stage_idx = next_stage_idx == stage_cnt ? 0 : next_stage_idx;
|
||||
if (loop_idx != k_iter_cnt - 1) {
|
||||
need_wait = !try_wait_barrier(smem_barrier + 1 + next_stage_idx * 2,
|
||||
next_phase_bit);
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int i = 0; i < a_inst_cnt_per_iter; i++) {
|
||||
int smem_offset = a_smem_offsets[i];
|
||||
bf16_t* smem_ptr_this_iter =
|
||||
smem_a + stage_idx * tile_m * tile_k + smem_offset;
|
||||
int linear_idx = local_tid * vec_elems + i * thread_cnt * vec_elems;
|
||||
int m_idx = linear_idx / tile_k;
|
||||
int k_idx = linear_idx % tile_k;
|
||||
int gmem_offset = m_idx * gemm_k + k_project(k_idx);
|
||||
bf16_t const* gmem_ptr_this_iter = gmem_a + gmem_offset;
|
||||
ldgsts_128(gmem_ptr_this_iter, smem_ptr_this_iter, true);
|
||||
}
|
||||
ldgsts_arrive(smem_barrier + stage_idx * 2);
|
||||
|
||||
stage_idx = next_stage_idx;
|
||||
phase_bit = next_phase_bit;
|
||||
gmem_a += per_mma_warp_k;
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
bf16_t const* gmem_a;
|
||||
bf16_t* smem_a;
|
||||
uint64_t* smem_barrier;
|
||||
int local_tid;
|
||||
int stage_idx = 0;
|
||||
int phase_bit = 1;
|
||||
bool need_wait = true;
|
||||
|
||||
// per smem_stage, store with swizzle information
|
||||
int a_smem_offsets[a_inst_cnt_per_iter];
|
||||
};
|
||||
|
||||
template <int gemm_k, int tile_n, int tile_k, int stage_cnt>
|
||||
struct GmemLoaderB {
|
||||
static constexpr int elem_bytes = 2;
|
||||
static constexpr int vec_bytes = 16;
|
||||
static constexpr int vec_elems = vec_bytes / elem_bytes;
|
||||
static constexpr int thread_cnt = 64;
|
||||
static_assert((tile_n * tile_k) % (vec_elems * thread_cnt) == 0);
|
||||
static constexpr int b_inst_cnt_per_iter =
|
||||
(tile_n * tile_k) / (vec_elems * thread_cnt);
|
||||
static_assert(gemm_k % tile_k == 0);
|
||||
static constexpr int k_iter_cnt = gemm_k / tile_k;
|
||||
|
||||
// Extra params to keep the order of k reduction...
|
||||
static constexpr int mma_warp_cnt = 4;
|
||||
static constexpr int per_mma_warp_k = tile_k / mma_warp_cnt;
|
||||
static constexpr int k_each_chunk = gemm_k / mma_warp_cnt;
|
||||
|
||||
private:
|
||||
__device__ int k_project(int tile_k_idx) {
|
||||
return (tile_k_idx / per_mma_warp_k * k_each_chunk) +
|
||||
(tile_k_idx % per_mma_warp_k);
|
||||
}
|
||||
|
||||
public:
|
||||
__device__ GmemLoaderB(bf16_t const* gmem_b_local_, bf16_t* smem_b_,
|
||||
uint64_t* smem_barrier_, int gemm_n_)
|
||||
: gmem_b(gmem_b_local_),
|
||||
smem_b(smem_b_),
|
||||
smem_barrier(smem_barrier_),
|
||||
gemm_n(gemm_n_),
|
||||
local_tid(threadIdx.x % thread_cnt) {}
|
||||
|
||||
__device__ void prepare() {
|
||||
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
|
||||
// swizzle, that's what we want.
|
||||
#pragma unroll
|
||||
for (int i = 0; i < b_inst_cnt_per_iter; i++) {
|
||||
int linear_idx = local_tid * vec_elems + i * thread_cnt * vec_elems;
|
||||
int n_idx = linear_idx / tile_k;
|
||||
int k_idx = linear_idx % tile_k;
|
||||
k_idx = apply_swizzle_343_on_elem_row_col<bf16_t>(n_idx, k_idx);
|
||||
b_smem_offsets[i] = n_idx * tile_k + k_idx;
|
||||
preds[i] = n_idx < gemm_n;
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
__device__ void issue_mainloop() {
|
||||
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
|
||||
asm volatile("griddepcontrol.wait;");
|
||||
#pragma unroll 1
|
||||
for (int loop_idx = 0; loop_idx < k_iter_cnt; loop_idx++) {
|
||||
if (need_wait) {
|
||||
wait_barrier(smem_barrier + 1 + stage_idx * 2, phase_bit);
|
||||
}
|
||||
int next_stage_idx = stage_idx + 1;
|
||||
int next_phase_bit =
|
||||
next_stage_idx == stage_cnt ? phase_bit ^ 1 : phase_bit;
|
||||
next_stage_idx = next_stage_idx == stage_cnt ? 0 : next_stage_idx;
|
||||
if (loop_idx != k_iter_cnt - 1) {
|
||||
need_wait = !try_wait_barrier(smem_barrier + 1 + next_stage_idx * 2,
|
||||
next_phase_bit);
|
||||
}
|
||||
#pragma unroll
|
||||
for (int i = 0; i < b_inst_cnt_per_iter; i++) {
|
||||
int smem_offset = b_smem_offsets[i];
|
||||
bf16_t* smem_ptr_this_iter =
|
||||
smem_b + stage_idx * tile_n * tile_k + smem_offset;
|
||||
int linear_idx = local_tid * vec_elems + i * thread_cnt * vec_elems;
|
||||
int n_idx = linear_idx / tile_k;
|
||||
int k_idx = linear_idx % tile_k;
|
||||
int gmem_offset = n_idx * gemm_k + k_project(k_idx);
|
||||
bf16_t const* gmem_ptr_this_iter = gmem_b + gmem_offset;
|
||||
ldgsts_128(gmem_ptr_this_iter, smem_ptr_this_iter, preds[i]);
|
||||
}
|
||||
ldgsts_arrive(smem_barrier + stage_idx * 2);
|
||||
|
||||
stage_idx = next_stage_idx;
|
||||
phase_bit = next_phase_bit;
|
||||
gmem_b += per_mma_warp_k;
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
bf16_t const* gmem_b;
|
||||
bf16_t* smem_b;
|
||||
uint64_t* smem_barrier;
|
||||
int gemm_n;
|
||||
int local_tid;
|
||||
int stage_idx = 0;
|
||||
int phase_bit = 1;
|
||||
bool need_wait = true;
|
||||
|
||||
// per smem_stage, store with swizzle information
|
||||
int b_smem_offsets[b_inst_cnt_per_iter];
|
||||
uint32_t preds[b_inst_cnt_per_iter];
|
||||
};
|
||||
|
||||
template <int gemm_m, int gemm_k, int tile_m, int tile_n, int tile_k,
|
||||
int stage_cnt>
|
||||
struct MmaComputer {
|
||||
static constexpr int elem_bytes = 2;
|
||||
static constexpr int thread_cnt = 128;
|
||||
static_assert(gemm_k % tile_k == 0);
|
||||
static_assert(tile_k % (thread_cnt / 32) == 0);
|
||||
static constexpr int per_warp_tile_k = tile_k / (thread_cnt / 32);
|
||||
static constexpr int k_iter_cnt = gemm_k / tile_k;
|
||||
static constexpr int k_phase_cnt = per_warp_tile_k / 16;
|
||||
static constexpr int m_iter_cnt = (tile_m + 15) / 16;
|
||||
static constexpr int n_iter_cnt =
|
||||
(tile_n + 7) /
|
||||
8; // Possible to have non-1 n_iter_cnt for ab_swap m16 case.
|
||||
static_assert(m_iter_cnt == 1);
|
||||
static_assert(n_iter_cnt == 1 || n_iter_cnt == 2);
|
||||
|
||||
__device__ MmaComputer(bf16_t* gmem_c_local_, bf16_t* smem_a_,
|
||||
bf16_t* smem_b_, uint64_t* smem_barrier_,
|
||||
int warp_idx_, int gemm_n_)
|
||||
: gmem_c(gmem_c_local_),
|
||||
smem_a(smem_a_),
|
||||
smem_b(smem_b_),
|
||||
smem_barrier(smem_barrier_),
|
||||
warp_idx(warp_idx_ - (thread_cnt / 32)),
|
||||
gemm_n(gemm_n_) {}
|
||||
|
||||
private:
|
||||
__device__ constexpr int internal_b_atom_func(int tid) {
|
||||
if constexpr (tile_n < 8) {
|
||||
return (tid % tile_n) + ((tid % 8) / tile_n * 0) + tid / 8 * 8 * tile_n;
|
||||
} else {
|
||||
return (tid % 8) + ((tid % 32) / 8 * (tile_n * 8));
|
||||
}
|
||||
}
|
||||
|
||||
public:
|
||||
__device__ void prepare() {
|
||||
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
|
||||
#pragma unroll
|
||||
for (int i = 0; i < k_phase_cnt; i++) {
|
||||
int linear_idx = (lane_idx % 16) + (lane_idx / 16) * 128 + i * 256;
|
||||
int m_idx = linear_idx % tile_m;
|
||||
int k_idx = linear_idx / tile_m + warp_k_offset_in_tile_k;
|
||||
k_idx = apply_swizzle_343_on_elem_row_col<bf16_t>(m_idx, k_idx);
|
||||
a_smem_offsets[0][i] = m_idx * tile_k + k_idx;
|
||||
}
|
||||
#pragma unroll
|
||||
for (int n_iter_idx = 0; n_iter_idx < n_iter_cnt; n_iter_idx++) {
|
||||
#pragma unroll
|
||||
for (int i = 0; i < k_phase_cnt; i += 2) { // Special i+=2 for B.
|
||||
int linear_idx =
|
||||
internal_b_atom_func(lane_idx) + i * tile_n * 16 + n_iter_idx * 8;
|
||||
int n_idx = linear_idx % tile_n;
|
||||
int k_idx = linear_idx / tile_n + warp_k_offset_in_tile_k;
|
||||
k_idx = apply_swizzle_343_on_elem_row_col<bf16_t>(n_idx, k_idx);
|
||||
b_smem_offsets[n_iter_idx][i] = n_idx * tile_k + k_idx;
|
||||
}
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
__device__ void issue_mainloop() {
|
||||
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
|
||||
#pragma unroll 1
|
||||
for (int loop_idx = 0; loop_idx < k_iter_cnt; loop_idx++) {
|
||||
wait_barrier(smem_barrier + 0 + stage_idx * 2, phase_bit);
|
||||
|
||||
#pragma unroll
|
||||
for (int i = 0; i < k_phase_cnt; i++) {
|
||||
int smem_offset = a_smem_offsets[0][i];
|
||||
bf16_t* smem_ptr_this_iter =
|
||||
smem_a + stage_idx * tile_m * tile_k + smem_offset;
|
||||
ldsm_x4(smem_ptr_this_iter, reinterpret_cast<uint32_t*>(a_reg[0][i]));
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int n_iter_idx = 0; n_iter_idx < n_iter_cnt; n_iter_idx++) {
|
||||
#pragma unroll
|
||||
for (int i = 0; i < k_phase_cnt; i += 2) {
|
||||
int smem_offset = b_smem_offsets[n_iter_idx][i];
|
||||
bf16_t* smem_ptr_this_iter =
|
||||
smem_b + stage_idx * tile_n * tile_k + smem_offset;
|
||||
ldsm_x4(smem_ptr_this_iter,
|
||||
reinterpret_cast<uint32_t*>(b_reg[n_iter_idx][i]));
|
||||
}
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int k_iter_idx = 0; k_iter_idx < k_phase_cnt; k_iter_idx++) {
|
||||
#pragma unroll
|
||||
for (int n_iter_idx = 0; n_iter_idx < n_iter_cnt; n_iter_idx++) {
|
||||
hmma_16_8_16_f32acc_bf16ab(
|
||||
acc_reg[0][n_iter_idx], a_reg[0][k_iter_idx],
|
||||
b_reg[n_iter_idx][k_iter_idx], acc_reg[0][n_iter_idx]);
|
||||
}
|
||||
}
|
||||
::arrive_barrier(smem_barrier + 1 + stage_idx * 2);
|
||||
stage_idx += 1;
|
||||
phase_bit = stage_idx == stage_cnt ? phase_bit ^ 1 : phase_bit;
|
||||
stage_idx = stage_idx == stage_cnt ? 0 : stage_idx;
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
__device__ void epi() {
|
||||
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
|
||||
asm volatile("bar.sync %0, %1;" : : "r"(1), "r"(thread_cnt));
|
||||
// reorganize the acc_reg
|
||||
constexpr int thread_m = 2;
|
||||
constexpr int thread_n = 2 * n_iter_cnt;
|
||||
constexpr int cta_mma_n = n_iter_cnt * 8;
|
||||
float acc_reg_reorg[thread_m][thread_n];
|
||||
|
||||
for (int i = 0; i < thread_m; i++) {
|
||||
for (int j = 0; j < thread_n; j++) {
|
||||
acc_reg_reorg[i][j] = acc_reg[0][j / 2][(j % 2) + (i * 2)];
|
||||
}
|
||||
}
|
||||
|
||||
// 4 x cosize(smem_c_layout)
|
||||
float* smem_c = reinterpret_cast<float*>(smem_a);
|
||||
// coord -> index
|
||||
auto smem_c_index_func = [&](int m_idx, int n_idx) {
|
||||
int group_rows = 32 / cta_mma_n;
|
||||
int group_cnt = 2;
|
||||
return (m_idx % group_rows * cta_mma_n) +
|
||||
(m_idx / group_rows * (32 + group_cnt)) + n_idx;
|
||||
};
|
||||
constexpr int cosize_smem_c = ((tile_m * cta_mma_n) / 32) * (32 + 2);
|
||||
|
||||
// This should be optimized to STS.64 but can not be STS.128 due to the bank
|
||||
// index.
|
||||
#pragma unroll
|
||||
for (int m_idx_thread = 0; m_idx_thread < thread_m; m_idx_thread++) {
|
||||
#pragma unroll
|
||||
for (int n_idx_thread = 0; n_idx_thread < thread_n; n_idx_thread++) {
|
||||
int m_idx = (lane_idx / 4) + m_idx_thread * 8;
|
||||
int n_idx =
|
||||
((lane_idx % 4) * 2) + (n_idx_thread % 2) + (n_idx_thread / 2) * 8;
|
||||
smem_c[cosize_smem_c * warp_idx + smem_c_index_func(m_idx, n_idx)] =
|
||||
acc_reg_reorg[m_idx_thread][n_idx_thread];
|
||||
}
|
||||
}
|
||||
asm volatile("bar.sync %0, %1;" : : "r"(1), "r"(thread_cnt));
|
||||
|
||||
if (warp_idx == 0) {
|
||||
constexpr int final_acc_reg_cnt = (tile_m * tile_n + 31) / 32;
|
||||
float acc_final[final_acc_reg_cnt]{};
|
||||
|
||||
#pragma unroll
|
||||
for (int reg_idx = 0; reg_idx < final_acc_reg_cnt; reg_idx++) {
|
||||
int linear_idx = reg_idx * 32 + lane_idx;
|
||||
int m_idx = linear_idx % tile_m;
|
||||
int n_idx = linear_idx / tile_m;
|
||||
acc_final[reg_idx] +=
|
||||
smem_c[smem_c_index_func(m_idx, n_idx) + 0 * cosize_smem_c] +
|
||||
smem_c[smem_c_index_func(m_idx, n_idx) + 1 * cosize_smem_c] +
|
||||
smem_c[smem_c_index_func(m_idx, n_idx) + 2 * cosize_smem_c] +
|
||||
smem_c[smem_c_index_func(m_idx, n_idx) + 3 * cosize_smem_c];
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int reg_idx = 0; reg_idx < final_acc_reg_cnt; reg_idx++) {
|
||||
int linear_idx = reg_idx * 32 + lane_idx;
|
||||
int m_idx = linear_idx % tile_m;
|
||||
int n_idx = linear_idx / tile_m;
|
||||
if (m_idx < tile_m && n_idx < gemm_n) {
|
||||
gmem_c[n_idx * gemm_m + m_idx] = acc_final[reg_idx];
|
||||
}
|
||||
}
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
bf16_t* gmem_c;
|
||||
bf16_t* smem_a;
|
||||
bf16_t* smem_b;
|
||||
uint64_t* smem_barrier;
|
||||
int warp_idx;
|
||||
int gemm_n;
|
||||
int stage_idx = 0;
|
||||
int phase_bit = 0;
|
||||
int lane_idx = threadIdx.x % 32;
|
||||
int warp_k_offset_in_tile_k = warp_idx * per_warp_tile_k;
|
||||
|
||||
int a_smem_offsets[m_iter_cnt][k_phase_cnt];
|
||||
int b_smem_offsets[n_iter_cnt][k_phase_cnt];
|
||||
|
||||
bf16_t a_reg[m_iter_cnt][k_phase_cnt][8];
|
||||
bf16_t b_reg[n_iter_cnt][k_phase_cnt][4];
|
||||
float acc_reg[m_iter_cnt][n_iter_cnt][4]{};
|
||||
};
|
||||
|
||||
// AB swapped, kernel is k-major, k-major, m-major
|
||||
template <int batch_size, int gemm_m, int gemm_k, int tile_m, int tile_n,
|
||||
int tile_k, int stage_cnt>
|
||||
__global__ __launch_bounds__(256, 1) void fused_a_gemm_kernel(
|
||||
bf16_t* output, bf16_t const* mat_a, bf16_t const* mat_b, int gemm_n) {
|
||||
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
|
||||
constexpr int load_thread_cnt = 128;
|
||||
constexpr int compute_thread_cnt = 128;
|
||||
constexpr int thread_cnt = load_thread_cnt + compute_thread_cnt;
|
||||
(void)thread_cnt;
|
||||
static_assert(gemm_m % 16 == 0);
|
||||
static_assert(gemm_k % tile_k == 0);
|
||||
static_assert(gemm_m % tile_m == 0);
|
||||
static_assert(
|
||||
tile_k == 128 || tile_k == 256 || tile_k == 512 ||
|
||||
tile_k == 1024); // tile_k must be larger than 64 since 4 warp splitK.
|
||||
static_assert(tile_m == 16);
|
||||
constexpr int g2s_vec_bytes = 16;
|
||||
constexpr int a_elem_bytes = 2;
|
||||
constexpr int b_elem_bytes = 2;
|
||||
static_assert((tile_m * a_elem_bytes + tile_n * b_elem_bytes) * tile_k *
|
||||
stage_cnt <=
|
||||
225 * 1024);
|
||||
static_assert((tile_m * tile_k * a_elem_bytes) %
|
||||
(load_thread_cnt * g2s_vec_bytes) ==
|
||||
0);
|
||||
static_assert((tile_n * tile_k * b_elem_bytes) %
|
||||
(load_thread_cnt * g2s_vec_bytes) ==
|
||||
0);
|
||||
|
||||
extern __shared__ char smem[];
|
||||
uint64_t* smem_barrier = reinterpret_cast<uint64_t*>(
|
||||
smem); // producer,consumer; producer,consumer; ...
|
||||
bf16_t* smem_a = reinterpret_cast<bf16_t*>(smem + (stage_cnt * 8 * 2 + 1024) /
|
||||
1024 * 1024);
|
||||
bf16_t* smem_b = smem_a + tile_m * tile_k * stage_cnt;
|
||||
|
||||
int cta_m_idx = tile_m * blockIdx.x;
|
||||
int cta_n_idx = tile_n * blockIdx.y;
|
||||
bf16_t const* gmem_a_local = mat_a + cta_m_idx * gemm_k;
|
||||
bf16_t const* gmem_b_local = mat_b + cta_n_idx * gemm_k;
|
||||
bf16_t* gmem_c_local = output + cta_n_idx * gemm_m + cta_m_idx;
|
||||
|
||||
int warp_idx = __shfl_sync(0xffffffff, threadIdx.x / 32, 0);
|
||||
|
||||
if (warp_idx == 4) {
|
||||
for (int i = 0; i < stage_cnt; i++) {
|
||||
initialize_barrier(smem_barrier + i * 2 + 0,
|
||||
load_thread_cnt); // producer
|
||||
initialize_barrier(smem_barrier + i * 2 + 1,
|
||||
compute_thread_cnt); // consumer
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
if (warp_idx < 2) {
|
||||
GmemLoaderA<gemm_k, tile_m, tile_k, stage_cnt> a_loader(
|
||||
gmem_a_local, smem_a, smem_barrier);
|
||||
a_loader.prepare();
|
||||
a_loader.issue_mainloop();
|
||||
} else if (warp_idx < 4) {
|
||||
GmemLoaderB<gemm_k, tile_n, tile_k, stage_cnt> b_loader(
|
||||
gmem_b_local, smem_b, smem_barrier, gemm_n);
|
||||
b_loader.prepare();
|
||||
b_loader.issue_mainloop();
|
||||
} else {
|
||||
MmaComputer<gemm_m, gemm_k, tile_m, tile_n, tile_k, stage_cnt> mma_computer(
|
||||
gmem_c_local, smem_a, smem_b, smem_barrier, warp_idx, gemm_n);
|
||||
mma_computer.prepare();
|
||||
mma_computer.issue_mainloop();
|
||||
mma_computer.epi();
|
||||
}
|
||||
asm volatile("griddepcontrol.launch_dependents;");
|
||||
#endif
|
||||
}
|
||||
|
||||
template <typename T, int kHdIn, int kHdOut, int kTileN>
|
||||
void invokeFusedAGemm(T* output, T const* mat_a, T const* mat_b, int num_tokens,
|
||||
cudaStream_t const stream) {
|
||||
constexpr int gemm_m = kHdOut; // 2112
|
||||
int const gemm_n = num_tokens; // 1-16
|
||||
constexpr int gemm_k = kHdIn; // 7168
|
||||
constexpr int batch_size = 1;
|
||||
std::swap(mat_a, mat_b);
|
||||
constexpr int tile_m = 16;
|
||||
constexpr int tile_n = kTileN; // 8 or 16
|
||||
constexpr int tile_k = std::max(256, 1024 / tile_n); // 256
|
||||
constexpr int max_stage_cnt =
|
||||
1024 * 192 / ((tile_m + tile_n) * tile_k * sizeof(bf16_t));
|
||||
constexpr int k_iter_cnt = gemm_k / tile_k;
|
||||
constexpr int stage_cnt =
|
||||
k_iter_cnt > max_stage_cnt ? max_stage_cnt : k_iter_cnt;
|
||||
int cta_m_cnt = gemm_m / tile_m;
|
||||
int cta_n_cnt = (gemm_n + tile_n - 1) / tile_n;
|
||||
constexpr int barrier_bytes = (stage_cnt * 16 + 1023) / 1024 * 1024;
|
||||
constexpr int smem_bytes =
|
||||
((tile_m * 2 + tile_n * 2) * tile_k * stage_cnt + barrier_bytes + 1023) /
|
||||
1024 * 1024;
|
||||
|
||||
dim3 grid(cta_m_cnt, cta_n_cnt, 1);
|
||||
dim3 block_size(256);
|
||||
cudaLaunchConfig_t config;
|
||||
config.gridDim = grid;
|
||||
config.blockDim = block_size;
|
||||
config.dynamicSmemBytes = smem_bytes;
|
||||
config.stream = stream;
|
||||
cudaLaunchAttribute attrs[1];
|
||||
attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
|
||||
attrs[0].val.programmaticStreamSerializationAllowed = getEnvEnablePDL();
|
||||
config.numAttrs = 1;
|
||||
config.attrs = attrs;
|
||||
if (smem_bytes >= (48 * 1024)) {
|
||||
cudaFuncSetAttribute(fused_a_gemm_kernel<batch_size, gemm_m, gemm_k, tile_m,
|
||||
tile_n, tile_k, stage_cnt>,
|
||||
cudaFuncAttributeMaxDynamicSharedMemorySize,
|
||||
smem_bytes);
|
||||
}
|
||||
cudaLaunchKernelEx(&config,
|
||||
fused_a_gemm_kernel<batch_size, gemm_m, gemm_k, tile_m,
|
||||
tile_n, tile_k, stage_cnt>,
|
||||
output, mat_a, mat_b, gemm_n);
|
||||
}
|
||||
|
||||
template void invokeFusedAGemm<__nv_bfloat16, 7168, 2112, 8>(
|
||||
__nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, int num_tokens,
|
||||
cudaStream_t);
|
||||
|
||||
template void invokeFusedAGemm<__nv_bfloat16, 7168, 2112, 16>(
|
||||
__nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, int num_tokens,
|
||||
cudaStream_t);
|
||||
|
||||
void dsv3_fused_a_gemm(torch::Tensor& output, torch::Tensor const& mat_a,
|
||||
torch::Tensor const& mat_b) {
|
||||
TORCH_CHECK(mat_a.dim() == 2 && mat_b.dim() == 2 && output.dim() == 2);
|
||||
int const num_tokens = mat_a.size(0);
|
||||
int const hd_in = mat_a.size(1);
|
||||
int const hd_out = mat_b.size(1);
|
||||
|
||||
constexpr int kHdIn = 7168;
|
||||
constexpr int kHdOut = 2112;
|
||||
TORCH_CHECK(num_tokens >= 1 && num_tokens <= 16,
|
||||
"required 1 <= mat_a.shape[0] <= 16")
|
||||
TORCH_CHECK(hd_in == kHdIn, "required mat_a.shape[1] == 7168")
|
||||
TORCH_CHECK(hd_out == kHdOut, "required mat_b.shape[1] == 2112")
|
||||
TORCH_CHECK(output.size(0) == num_tokens,
|
||||
"required output.shape[0] == mat_a.shape[0]")
|
||||
TORCH_CHECK(output.size(1) == hd_out,
|
||||
"required output.shape[1] == mat_b.shape[1]")
|
||||
|
||||
TORCH_CHECK(mat_a.stride(1) == 1, "mat_a must be a row major tensor");
|
||||
TORCH_CHECK(output.stride(1) == 1, "output must be a row major tensor");
|
||||
TORCH_CHECK(mat_b.stride(0) == 1, "mat_b must be a column major tensor");
|
||||
|
||||
TORCH_CHECK(mat_a.scalar_type() == torch::kBFloat16 &&
|
||||
mat_b.scalar_type() == torch::kBFloat16,
|
||||
"Only BFloat16 input dtype is supported")
|
||||
TORCH_CHECK(output.scalar_type() == torch::kBFloat16,
|
||||
"Only BFloat16 output dtype is supported")
|
||||
|
||||
TORCH_CHECK(getSMVersion() >= 90, "required CUDA ARCH >= SM_90");
|
||||
|
||||
auto stream = at::cuda::getCurrentCUDAStream(mat_a.get_device());
|
||||
if (num_tokens <= 8) {
|
||||
invokeFusedAGemm<__nv_bfloat16, kHdIn, kHdOut, 8>(
|
||||
reinterpret_cast<__nv_bfloat16*>(output.mutable_data_ptr()),
|
||||
reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr()),
|
||||
reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr()), num_tokens,
|
||||
stream);
|
||||
} else {
|
||||
invokeFusedAGemm<__nv_bfloat16, kHdIn, kHdOut, 16>(
|
||||
reinterpret_cast<__nv_bfloat16*>(output.mutable_data_ptr()),
|
||||
reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr()),
|
||||
reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr()), num_tokens,
|
||||
stream);
|
||||
}
|
||||
}
|
||||
+8
-1
@@ -315,7 +315,9 @@ void silu_and_mul_scaled_fp4_experts_quant(
|
||||
void per_token_group_quant_fp8(const torch::Tensor& input,
|
||||
torch::Tensor& output_q, torch::Tensor& output_s,
|
||||
int64_t group_size, double eps, double fp8_min,
|
||||
double fp8_max, bool scale_ue8m0);
|
||||
double fp8_max, bool scale_ue8m0,
|
||||
bool dummy_is_scale_transposed,
|
||||
bool dummy_is_tma_aligned);
|
||||
|
||||
void per_token_group_quant_int8(const torch::Tensor& input,
|
||||
torch::Tensor& output_q,
|
||||
@@ -408,3 +410,8 @@ void qr_all_reduce(fptr_t _fa, torch::Tensor& inp, torch::Tensor& out,
|
||||
int64_t quant_level, bool cast_bf2half = false);
|
||||
int64_t qr_max_size();
|
||||
#endif
|
||||
|
||||
#ifndef USE_ROCM
|
||||
void dsv3_fused_a_gemm(torch::Tensor& output, torch::Tensor const& mat_a,
|
||||
torch::Tensor const& mat_b);
|
||||
#endif
|
||||
@@ -97,7 +97,7 @@ __global__ void rms_norm_per_block_quant_kernel(
|
||||
scalar_t const* __restrict__ input, // [..., hidden_size]
|
||||
scalar_t const* __restrict__ weight, // [hidden_size]
|
||||
float const* scale_ub, float const var_epsilon, int32_t const hidden_size,
|
||||
scalar_t* __restrict__ residual = nullptr) {
|
||||
scalar_t* __restrict__ residual = nullptr, int64_t outer_scale_stride = 1) {
|
||||
float rms;
|
||||
// Compute RMS
|
||||
// Always able to vectorize due to constraints on hidden_size
|
||||
@@ -108,7 +108,8 @@ __global__ void rms_norm_per_block_quant_kernel(
|
||||
// Always able to vectorize due to constraints on hidden_size and group_size
|
||||
vllm::vectorized::compute_dynamic_per_token_scales<
|
||||
scalar_t, scalar_out_t, has_residual, is_scale_transposed, group_size>(
|
||||
nullptr, scales, input, weight, rms, scale_ub, hidden_size, residual);
|
||||
nullptr, scales, input, weight, rms, scale_ub, hidden_size, residual,
|
||||
outer_scale_stride);
|
||||
|
||||
// RMS Norm + Quant
|
||||
// Always able to vectorize due to constraints on hidden_size
|
||||
@@ -119,7 +120,8 @@ __global__ void rms_norm_per_block_quant_kernel(
|
||||
vllm::vectorized::norm_and_quant<
|
||||
scalar_t, scalar_out_t, std::is_same_v<scalar_out_t, int8_t>,
|
||||
has_residual, is_scale_transposed, group_size>(
|
||||
out, input, weight, rms, scales, hidden_size, residual);
|
||||
out, input, weight, rms, scales, hidden_size, residual,
|
||||
outer_scale_stride);
|
||||
}
|
||||
|
||||
} // namespace vllm
|
||||
@@ -225,7 +227,8 @@ void rms_norm_per_block_quant_dispatch(
|
||||
: nullptr,
|
||||
var_epsilon, hidden_size,
|
||||
has_residual ? residual->data_ptr<scalar_in_t>()
|
||||
: nullptr);
|
||||
: nullptr,
|
||||
scales.stride(1));
|
||||
});
|
||||
});
|
||||
});
|
||||
@@ -257,6 +260,11 @@ void rms_norm_per_block_quant(torch::Tensor& out, torch::Tensor const& input,
|
||||
TORCH_CHECK(group_size == 128 || group_size == 64,
|
||||
"Unsupported group size: ", group_size);
|
||||
|
||||
if (scales.stride(1) > 1) {
|
||||
TORCH_CHECK(is_scale_transposed,
|
||||
"Outer scale stride must be 1 when scales are not transposed");
|
||||
}
|
||||
|
||||
rms_norm_per_block_quant_dispatch(out, input, weight, scales, group_size,
|
||||
var_epsilon, scale_ub, residual,
|
||||
is_scale_transposed);
|
||||
|
||||
@@ -74,7 +74,7 @@ __device__ void compute_dynamic_per_token_scales(
|
||||
scalar_t const* __restrict__ input, scalar_t const* __restrict__ weight,
|
||||
float const rms, float const* __restrict__ scale_ub,
|
||||
int32_t const hidden_size, scalar_t const* __restrict__ residual = nullptr,
|
||||
int32_t const group_size = 0) {
|
||||
int32_t const group_size = 0, int64_t outer_scale_stride = 1) {
|
||||
float block_absmax_val_maybe = 0.0f;
|
||||
constexpr scalar_out_t qmax{quant_type_max_v<scalar_out_t>};
|
||||
__syncthreads();
|
||||
@@ -133,7 +133,9 @@ __device__ void compute_dynamic_per_token_scales(
|
||||
scale = max(scale / qmax, min_scaling_factor<scalar_out_t>::val());
|
||||
// Global output store
|
||||
if constexpr (is_scale_transposed) {
|
||||
all_token_scales[(threadIdx.x / threads_per_group) * gridDim.x +
|
||||
int64_t const scale_rows = (gridDim.x + outer_scale_stride - 1) /
|
||||
outer_scale_stride * outer_scale_stride;
|
||||
all_token_scales[(threadIdx.x / threads_per_group) * scale_rows +
|
||||
blockIdx.x] = scale;
|
||||
} else {
|
||||
all_token_scales[blockIdx.x * num_groups +
|
||||
@@ -180,13 +182,11 @@ __device__ void compute_dynamic_per_token_scales(
|
||||
|
||||
template <typename scalar_t, typename scalar_out_t, bool is_scale_inverted,
|
||||
bool has_residual = false, bool is_scale_transposed = false>
|
||||
__device__ void norm_and_quant(scalar_out_t* __restrict__ output,
|
||||
scalar_t const* __restrict__ input,
|
||||
scalar_t const* __restrict__ weight,
|
||||
float const rms, float* const scale,
|
||||
int32_t const hidden_size,
|
||||
scalar_t* __restrict__ residual = nullptr,
|
||||
int32_t const group_size = 0) {
|
||||
__device__ void norm_and_quant(
|
||||
scalar_out_t* __restrict__ output, scalar_t const* __restrict__ input,
|
||||
scalar_t const* __restrict__ weight, float const rms, float* const scale,
|
||||
int32_t const hidden_size, scalar_t* __restrict__ residual = nullptr,
|
||||
int32_t const group_size = 0, int64_t outer_scale_stride = 1) {
|
||||
int64_t const token_offset = blockIdx.x * static_cast<int64_t>(hidden_size);
|
||||
|
||||
for (auto i = threadIdx.x; i < hidden_size; i += blockDim.x) {
|
||||
@@ -202,7 +202,9 @@ __device__ void norm_and_quant(scalar_out_t* __restrict__ output,
|
||||
int64_t scale_idx = 0;
|
||||
if (group_size > 0) {
|
||||
if constexpr (is_scale_transposed) {
|
||||
scale_idx = (i / group_size) * gridDim.x + blockIdx.x;
|
||||
int64_t const scale_rows = (gridDim.x + outer_scale_stride - 1) /
|
||||
outer_scale_stride * outer_scale_stride;
|
||||
scale_idx = (i / group_size) * scale_rows + blockIdx.x;
|
||||
} else {
|
||||
scale_idx = blockIdx.x * (hidden_size / group_size) + i / group_size;
|
||||
}
|
||||
@@ -286,8 +288,8 @@ __device__ void compute_dynamic_per_token_scales(
|
||||
float* __restrict__ token_scale, float* __restrict__ all_token_scales,
|
||||
scalar_t const* __restrict__ input, scalar_t const* __restrict__ weight,
|
||||
float const rms, float const* __restrict__ scale_ub,
|
||||
int32_t const hidden_size,
|
||||
scalar_t const* __restrict__ residual = nullptr) {
|
||||
int32_t const hidden_size, scalar_t const* __restrict__ residual = nullptr,
|
||||
int64_t outer_scale_stride = 1) {
|
||||
constexpr scalar_out_t qmax{quant_type_max_v<scalar_out_t>};
|
||||
|
||||
const int VEC_SIZE = 4;
|
||||
@@ -382,7 +384,9 @@ __device__ void compute_dynamic_per_token_scales(
|
||||
scale = max(scale / qmax, min_scaling_factor<scalar_out_t>::val());
|
||||
// Global output store
|
||||
if constexpr (is_scale_transposed) {
|
||||
all_token_scales[(threadIdx.x / threads_per_group) * gridDim.x +
|
||||
int64_t const scale_rows = (gridDim.x + outer_scale_stride - 1) /
|
||||
outer_scale_stride * outer_scale_stride;
|
||||
all_token_scales[(threadIdx.x / threads_per_group) * scale_rows +
|
||||
blockIdx.x] = scale;
|
||||
} else {
|
||||
all_token_scales[blockIdx.x * num_groups +
|
||||
@@ -463,7 +467,8 @@ __device__ void norm_and_quant(scalar_out_t* __restrict__ output,
|
||||
scalar_t const* __restrict__ weight,
|
||||
float const rms, float* const scale,
|
||||
int32_t const hidden_size,
|
||||
scalar_t* __restrict__ residual = nullptr) {
|
||||
scalar_t* __restrict__ residual = nullptr,
|
||||
int64_t outer_scale_stride = 1) {
|
||||
int64_t const token_offset = blockIdx.x * static_cast<int64_t>(hidden_size);
|
||||
|
||||
// Vectorized input/output/weight/residual to better utilize memory bandwidth.
|
||||
@@ -516,7 +521,9 @@ __device__ void norm_and_quant(scalar_out_t* __restrict__ output,
|
||||
int64_t const num_groups = hidden_size / group_size;
|
||||
int64_t scale_idx = 0;
|
||||
if constexpr (is_scale_transposed) {
|
||||
scale_idx = (i * VEC_SIZE / group_size) * gridDim.x + blockIdx.x;
|
||||
int64_t const scale_rows = (gridDim.x + outer_scale_stride - 1) /
|
||||
outer_scale_stride * outer_scale_stride;
|
||||
scale_idx = (i * VEC_SIZE / group_size) * scale_rows + blockIdx.x;
|
||||
} else {
|
||||
scale_idx = blockIdx.x * num_groups + i * VEC_SIZE / group_size;
|
||||
}
|
||||
|
||||
@@ -379,7 +379,9 @@ void per_token_group_quant_8bit_packed(const torch::Tensor& input,
|
||||
void per_token_group_quant_fp8(const torch::Tensor& input,
|
||||
torch::Tensor& output_q, torch::Tensor& output_s,
|
||||
int64_t group_size, double eps, double fp8_min,
|
||||
double fp8_max, bool scale_ue8m0) {
|
||||
double fp8_max, bool scale_ue8m0,
|
||||
bool dummy_is_scale_transposed = false,
|
||||
bool dummy_is_tma_aligned = false) {
|
||||
per_token_group_quant_8bit(input, output_q, output_s, group_size, eps,
|
||||
fp8_min, fp8_max, scale_ue8m0);
|
||||
}
|
||||
@@ -239,6 +239,11 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
|
||||
|
||||
// Quantization ops
|
||||
#ifndef USE_ROCM
|
||||
// DeepSeek V3 fused A GEMM (SM 9.0+, bf16 only, 1-16 tokens).
|
||||
ops.def(
|
||||
"dsv3_fused_a_gemm(Tensor! output, Tensor mat_a, Tensor mat_b) -> ()");
|
||||
ops.impl("dsv3_fused_a_gemm", torch::kCUDA, &dsv3_fused_a_gemm);
|
||||
|
||||
// Quantized GEMM for AWQ.
|
||||
ops.def(
|
||||
"awq_gemm(Tensor _in_feats, Tensor _kernel, Tensor _scaling_factors, "
|
||||
@@ -643,11 +648,13 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
|
||||
|
||||
#ifndef USE_ROCM
|
||||
// Compute per-token-group FP8 quantized tensor and scaling factor.
|
||||
// The dummy arguments are here so we can correctly fuse with RMSNorm.
|
||||
ops.def(
|
||||
"per_token_group_fp8_quant(Tensor input, Tensor! output_q, Tensor! "
|
||||
"output_s, "
|
||||
"int group_size, float eps, float fp8_min, float fp8_max, bool "
|
||||
"scale_ue8m0) -> ()");
|
||||
"scale_ue8m0, bool dummy_is_scale_transposed, bool dummy_is_tma_aligned "
|
||||
") -> ()");
|
||||
ops.impl("per_token_group_fp8_quant", torch::kCUDA,
|
||||
&per_token_group_quant_fp8);
|
||||
|
||||
|
||||
+4
-8
@@ -418,10 +418,6 @@ RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
fi && \
|
||||
python3 setup.py bdist_wheel --dist-dir=dist --py-limited-api=cp38
|
||||
|
||||
# Copy extension wheels from extensions-build stage for later use
|
||||
COPY --from=extensions-build /tmp/deepgemm/dist /tmp/deepgemm/dist
|
||||
COPY --from=extensions-build /tmp/ep_kernels_workspace/dist /tmp/ep_kernels_workspace/dist
|
||||
|
||||
# Check the size of the wheel if RUN_WHEEL_CHECK is true
|
||||
COPY .buildkite/check-wheel-size.py check-wheel-size.py
|
||||
# sync the default value with .buildkite/check-wheel-size.py
|
||||
@@ -660,9 +656,9 @@ RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
. /etc/environment && \
|
||||
uv pip list
|
||||
|
||||
# Install deepgemm wheel that has been built in the `build` stage
|
||||
# Install deepgemm wheel that has been built in the `extensions-build` stage
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
--mount=type=bind,from=build,source=/tmp/deepgemm/dist,target=/tmp/deepgemm/dist,ro \
|
||||
--mount=type=bind,from=extensions-build,source=/tmp/deepgemm/dist,target=/tmp/deepgemm/dist,ro \
|
||||
sh -c 'if ls /tmp/deepgemm/dist/*.whl >/dev/null 2>&1; then \
|
||||
uv pip install --system /tmp/deepgemm/dist/*.whl; \
|
||||
else \
|
||||
@@ -672,8 +668,8 @@ RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
# Pytorch now installs NVSHMEM, setting LD_LIBRARY_PATH
|
||||
ENV LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH
|
||||
|
||||
# Install EP kernels wheels (pplx-kernels and DeepEP) that have been built in the `build` stage
|
||||
RUN --mount=type=bind,from=build,src=/tmp/ep_kernels_workspace/dist,target=/vllm-workspace/ep_kernels/dist \
|
||||
# Install EP kernels wheels (pplx-kernels and DeepEP) that have been built in the `extensions-build` stage
|
||||
RUN --mount=type=bind,from=extensions-build,src=/tmp/ep_kernels_workspace/dist,target=/vllm-workspace/ep_kernels/dist \
|
||||
--mount=type=cache,target=/root/.cache/uv \
|
||||
uv pip install --system ep_kernels/dist/*.whl --verbose \
|
||||
--extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
|
||||
|
||||
Binary file not shown.
|
Before Width: | Height: | Size: 325 KiB After Width: | Height: | Size: 331 KiB |
File diff suppressed because one or more lines are too long
|
After Width: | Height: | Size: 339 KiB |
File diff suppressed because one or more lines are too long
|
After Width: | Height: | Size: 374 KiB |
@@ -155,3 +155,4 @@ The interface for the model/module may change during vLLM's development. If you
|
||||
- `use_v1` parameter in `Platform.get_attn_backend_cls` is deprecated. It has been removed in v0.13.0.
|
||||
- `_Backend` in `vllm.attention` is deprecated. It has been removed in v0.13.0. Please use `vllm.v1.attention.backends.registry.register_backend` to add new attention backend to `AttentionBackendEnum` instead.
|
||||
- `seed_everything` platform interface is deprecated. It has been removed in v0.16.0. Please use `vllm.utils.torch_utils.set_random_seed` instead.
|
||||
- `prompt` in `Platform.validate_request` is deprecated and will be removed in v0.18.0.
|
||||
|
||||
@@ -36,12 +36,12 @@ th:not(:first-child) {
|
||||
}
|
||||
</style>
|
||||
|
||||
| Feature | [CP](../configuration/optimization.md#chunked-prefill) | [APC](automatic_prefix_caching.md) | [LoRA](lora.md) | [SD](spec_decode/README.md) | CUDA graph | [pooling](../models/pooling_models.md) | <abbr title="Encoder-Decoder Models">enc-dec</abbr> | <abbr title="Logprobs">logP</abbr> | <abbr title="Prompt Logprobs">prmpt logP</abbr> | <abbr title="Async Output Processing">async output</abbr> | multi-step | <abbr title="Multimodal Inputs">mm</abbr> | best-of | beam-search | [prompt-embeds](prompt_embeds.md) |
|
||||
| Feature | [CP](../configuration/optimization.md#chunked-prefill) | [APC](automatic_prefix_caching.md) | [LoRA](lora.md) | [SD](speculative_decoding/README.md) | CUDA graph | [pooling](../models/pooling_models.md) | <abbr title="Encoder-Decoder Models">enc-dec</abbr> | <abbr title="Logprobs">logP</abbr> | <abbr title="Prompt Logprobs">prmpt logP</abbr> | <abbr title="Async Output Processing">async output</abbr> | multi-step | <abbr title="Multimodal Inputs">mm</abbr> | best-of | beam-search | [prompt-embeds](prompt_embeds.md) |
|
||||
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
|
||||
| [CP](../configuration/optimization.md#chunked-prefill) | ✅ | | | | | | | | | | | | | | |
|
||||
| [APC](automatic_prefix_caching.md) | ✅ | ✅ | | | | | | | | | | | | | |
|
||||
| [LoRA](lora.md) | ✅ | ✅ | ✅ | | | | | | | | | | | | |
|
||||
| [SD](spec_decode/README.md) | ✅ | ✅ | ❌ | ✅ | | | | | | | | | | | |
|
||||
| [SD](speculative_decoding/README.md) | ✅ | ✅ | ❌ | ✅ | | | | | | | | | | | |
|
||||
| CUDA graph | ✅ | ✅ | ✅ | ✅ | ✅ | | | | | | | | | | |
|
||||
| [pooling](../models/pooling_models.md) | 🟠\* | 🟠\* | ✅ | ❌ | ✅ | ✅ | | | | | | | | | |
|
||||
| <abbr title="Encoder-Decoder Models">enc-dec</abbr> | ❌ | [❌](https://github.com/vllm-project/vllm/issues/7366) | ❌ | [❌](https://github.com/vllm-project/vllm/issues/7366) | ✅ | ✅ | ✅ | | | | | | | | |
|
||||
@@ -64,7 +64,7 @@ th:not(:first-child) {
|
||||
| [CP](../configuration/optimization.md#chunked-prefill) | [❌](https://github.com/vllm-project/vllm/issues/2729) | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
|
||||
| [APC](automatic_prefix_caching.md) | [❌](https://github.com/vllm-project/vllm/issues/3687) | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
|
||||
| [LoRA](lora.md) | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
|
||||
| [SD](spec_decode/README.md) | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ |
|
||||
| [SD](speculative_decoding/README.md) | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ |
|
||||
| CUDA graph | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | [❌](https://github.com/vllm-project/vllm/issues/26970) |
|
||||
| [pooling](../models/pooling_models.md) | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
|
||||
| <abbr title="Encoder-Decoder Models">enc-dec</abbr> | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ |
|
||||
|
||||
@@ -1,330 +0,0 @@
|
||||
# Speculative Decoding
|
||||
|
||||
!!! warning
|
||||
Currently, speculative decoding in vLLM is not compatible with pipeline parallelism.
|
||||
|
||||
This document shows how to use [Speculative Decoding](https://x.com/karpathy/status/1697318534555336961) with vLLM.
|
||||
Speculative decoding is a technique which improves inter-token latency in memory-bound LLM inference.
|
||||
|
||||
!!! tip
|
||||
To train your own draft models for speculative decoding, see [Speculators](speculators.md), a library for training draft models that integrates seamlessly with vLLM.
|
||||
|
||||
## Speculating with a draft model
|
||||
|
||||
The following code configures vLLM in an offline mode to use speculative decoding with a draft model, speculating 5 tokens at a time.
|
||||
|
||||
!!! warning
|
||||
In vllm v0.10.0, speculative decoding with a draft model is not supported.
|
||||
If you use the following code, you will get a `NotImplementedError`.
|
||||
|
||||
??? code
|
||||
|
||||
```python
|
||||
from vllm import LLM, SamplingParams
|
||||
|
||||
prompts = [
|
||||
"The future of AI is",
|
||||
]
|
||||
sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
|
||||
|
||||
llm = LLM(
|
||||
model="facebook/opt-6.7b",
|
||||
tensor_parallel_size=1,
|
||||
speculative_config={
|
||||
"model": "facebook/opt-125m",
|
||||
"num_speculative_tokens": 5,
|
||||
},
|
||||
)
|
||||
outputs = llm.generate(prompts, sampling_params)
|
||||
|
||||
for output in outputs:
|
||||
prompt = output.prompt
|
||||
generated_text = output.outputs[0].text
|
||||
print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
|
||||
```
|
||||
|
||||
To perform the same with an online mode launch the server:
|
||||
|
||||
```bash
|
||||
vllm serve facebook/opt-6.7b \
|
||||
--host 0.0.0.0 \
|
||||
--port 8000 \
|
||||
--seed 42 \
|
||||
-tp 1 \
|
||||
--gpu_memory_utilization 0.8 \
|
||||
--speculative_config '{"model": "facebook/opt-125m", "num_speculative_tokens": 5}'
|
||||
```
|
||||
|
||||
!!! warning
|
||||
Note: Please use `--speculative_config` to set all configurations related to speculative decoding. The previous method of specifying the model through `--speculative_model` and adding related parameters (e.g., `--num_speculative_tokens`) separately has been deprecated now.
|
||||
|
||||
Then use a client:
|
||||
|
||||
??? code
|
||||
|
||||
```python
|
||||
from openai import OpenAI
|
||||
|
||||
# Modify OpenAI's API key and API base to use vLLM's API server.
|
||||
openai_api_key = "EMPTY"
|
||||
openai_api_base = "http://localhost:8000/v1"
|
||||
|
||||
client = OpenAI(
|
||||
# defaults to os.environ.get("OPENAI_API_KEY")
|
||||
api_key=openai_api_key,
|
||||
base_url=openai_api_base,
|
||||
)
|
||||
|
||||
models = client.models.list()
|
||||
model = models.data[0].id
|
||||
|
||||
# Completion API
|
||||
stream = False
|
||||
completion = client.completions.create(
|
||||
model=model,
|
||||
prompt="The future of AI is",
|
||||
echo=False,
|
||||
n=1,
|
||||
stream=stream,
|
||||
)
|
||||
|
||||
print("Completion results:")
|
||||
if stream:
|
||||
for c in completion:
|
||||
print(c)
|
||||
else:
|
||||
print(completion)
|
||||
```
|
||||
|
||||
## Speculating by matching n-grams in the prompt
|
||||
|
||||
The following code configures vLLM to use speculative decoding where proposals are generated by
|
||||
matching n-grams in the prompt. For more information read [this thread.](https://x.com/joao_gante/status/1747322413006643259)
|
||||
|
||||
??? code
|
||||
|
||||
```python
|
||||
from vllm import LLM, SamplingParams
|
||||
|
||||
prompts = [
|
||||
"The future of AI is",
|
||||
]
|
||||
sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
|
||||
|
||||
llm = LLM(
|
||||
model="facebook/opt-6.7b",
|
||||
tensor_parallel_size=1,
|
||||
speculative_config={
|
||||
"method": "ngram",
|
||||
"num_speculative_tokens": 5,
|
||||
"prompt_lookup_max": 4,
|
||||
},
|
||||
)
|
||||
outputs = llm.generate(prompts, sampling_params)
|
||||
|
||||
for output in outputs:
|
||||
prompt = output.prompt
|
||||
generated_text = output.outputs[0].text
|
||||
print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
|
||||
```
|
||||
|
||||
## Speculating using Suffix Decoding
|
||||
|
||||
The following code configures vLLM to use speculative decoding where proposals are generated using Suffix Decoding ([technical report](https://arxiv.org/abs/2411.04975)).
|
||||
|
||||
Like n-gram, Suffix Decoding can generate draft tokens by pattern-matching using the last `n` generated tokens. Unlike n-gram, Suffix Decoding (1) can pattern-match against both the prompt and previous generations, (2) uses frequency counts to propose the most likely continuations, and (3) speculates an adaptive number of tokens for each request at each iteration to get better acceptance rates.
|
||||
|
||||
Suffix Decoding can achieve better performance for tasks with high repetition, such as code-editing, agentic loops (e.g. self-reflection, self-consistency), and RL rollouts.
|
||||
|
||||
!!! tip "Install Arctic Inference"
|
||||
Suffix Decoding requires [Arctic Inference](https://github.com/snowflakedb/ArcticInference). You can install it with `pip install arctic-inference`.
|
||||
|
||||
!!! tip "Suffix Decoding Speculative Tokens"
|
||||
Suffix Decoding will speculate a dynamic number of tokens for each request at each decoding step, so the `num_speculative_tokens` configuration specifies the *maximum* number of speculative tokens. It is suggested to use a high number such as `16` or `32` (default).
|
||||
|
||||
??? code
|
||||
|
||||
```python
|
||||
from vllm import LLM, SamplingParams
|
||||
|
||||
prompts = [
|
||||
"The future of AI is",
|
||||
]
|
||||
sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
|
||||
|
||||
llm = LLM(
|
||||
model="facebook/opt-6.7b",
|
||||
tensor_parallel_size=1,
|
||||
speculative_config={
|
||||
"method": "suffix",
|
||||
"num_speculative_tokens": 32,
|
||||
},
|
||||
)
|
||||
outputs = llm.generate(prompts, sampling_params)
|
||||
|
||||
for output in outputs:
|
||||
prompt = output.prompt
|
||||
generated_text = output.outputs[0].text
|
||||
print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
|
||||
```
|
||||
|
||||
## Speculating using MLP speculators
|
||||
|
||||
The following code configures vLLM to use speculative decoding where proposals are generated by
|
||||
draft models that condition draft predictions on both context vectors and sampled tokens.
|
||||
For more information see [this blog](https://pytorch.org/blog/hitchhikers-guide-speculative-decoding/) or
|
||||
[this technical report](https://arxiv.org/abs/2404.19124).
|
||||
|
||||
??? code
|
||||
|
||||
```python
|
||||
from vllm import LLM, SamplingParams
|
||||
|
||||
prompts = [
|
||||
"The future of AI is",
|
||||
]
|
||||
sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
|
||||
|
||||
llm = LLM(
|
||||
model="meta-llama/Meta-Llama-3.1-70B-Instruct",
|
||||
tensor_parallel_size=4,
|
||||
speculative_config={
|
||||
"model": "ibm-ai-platform/llama3-70b-accelerator",
|
||||
"draft_tensor_parallel_size": 1,
|
||||
},
|
||||
)
|
||||
outputs = llm.generate(prompts, sampling_params)
|
||||
|
||||
for output in outputs:
|
||||
prompt = output.prompt
|
||||
generated_text = output.outputs[0].text
|
||||
print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
|
||||
```
|
||||
|
||||
Note that these speculative models currently need to be run without tensor parallelism, although
|
||||
it is possible to run the main model using tensor parallelism (see example above). Since the
|
||||
speculative models are relatively small, we still see significant speedups. However, this
|
||||
limitation will be fixed in a future release.
|
||||
|
||||
A variety of speculative models of this type are available on HF hub:
|
||||
|
||||
- [llama-13b-accelerator](https://huggingface.co/ibm-ai-platform/llama-13b-accelerator)
|
||||
- [llama3-8b-accelerator](https://huggingface.co/ibm-ai-platform/llama3-8b-accelerator)
|
||||
- [codellama-34b-accelerator](https://huggingface.co/ibm-ai-platform/codellama-34b-accelerator)
|
||||
- [llama2-70b-accelerator](https://huggingface.co/ibm-ai-platform/llama2-70b-accelerator)
|
||||
- [llama3-70b-accelerator](https://huggingface.co/ibm-ai-platform/llama3-70b-accelerator)
|
||||
- [granite-3b-code-instruct-accelerator](https://huggingface.co/ibm-granite/granite-3b-code-instruct-accelerator)
|
||||
- [granite-8b-code-instruct-accelerator](https://huggingface.co/ibm-granite/granite-8b-code-instruct-accelerator)
|
||||
- [granite-7b-instruct-accelerator](https://huggingface.co/ibm-granite/granite-7b-instruct-accelerator)
|
||||
- [granite-20b-code-instruct-accelerator](https://huggingface.co/ibm-granite/granite-20b-code-instruct-accelerator)
|
||||
|
||||
## Speculating using EAGLE based draft models
|
||||
|
||||
The following code configures vLLM to use speculative decoding where proposals are generated by
|
||||
an [EAGLE (Extrapolation Algorithm for Greater Language-model Efficiency)](https://arxiv.org/pdf/2401.15077) based draft model. A more detailed example for offline mode, including how to extract request level acceptance rate, can be found in [examples/offline_inference/spec_decode.py](../../../examples/offline_inference/spec_decode.py)
|
||||
|
||||
??? code
|
||||
|
||||
```python
|
||||
from vllm import LLM, SamplingParams
|
||||
|
||||
prompts = [
|
||||
"The future of AI is",
|
||||
]
|
||||
sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
|
||||
|
||||
llm = LLM(
|
||||
model="meta-llama/Meta-Llama-3-8B-Instruct",
|
||||
tensor_parallel_size=4,
|
||||
speculative_config={
|
||||
"model": "yuhuili/EAGLE-LLaMA3-Instruct-8B",
|
||||
"draft_tensor_parallel_size": 1,
|
||||
"num_speculative_tokens": 2,
|
||||
"method": "eagle",
|
||||
},
|
||||
)
|
||||
|
||||
outputs = llm.generate(prompts, sampling_params)
|
||||
|
||||
for output in outputs:
|
||||
prompt = output.prompt
|
||||
generated_text = output.outputs[0].text
|
||||
print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
|
||||
|
||||
```
|
||||
|
||||
A few important things to consider when using the EAGLE based draft models:
|
||||
|
||||
1. The EAGLE draft models available in the [HF repository for EAGLE models](https://huggingface.co/yuhuili) should
|
||||
be able to be loaded and used directly by vLLM after <https://github.com/vllm-project/vllm/pull/12304>.
|
||||
If you are using vllm version before <https://github.com/vllm-project/vllm/pull/12304>, please use the
|
||||
[script](https://gist.github.com/abhigoyal1997/1e7a4109ccb7704fbc67f625e86b2d6d) to convert the speculative model,
|
||||
and specify `"model": "path/to/modified/eagle/model"` in `speculative_config`. If weight-loading problems still occur when using the latest version of vLLM, please leave a comment or raise an issue.
|
||||
|
||||
2. The EAGLE based draft models need to be run without tensor parallelism
|
||||
(i.e. draft_tensor_parallel_size is set to 1 in `speculative_config`), although
|
||||
it is possible to run the main model using tensor parallelism (see example above).
|
||||
|
||||
3. When using EAGLE-based speculators with vLLM, the observed speedup is lower than what is
|
||||
reported in the reference implementation [here](https://github.com/SafeAILab/EAGLE). This issue is under
|
||||
investigation and tracked here: <https://github.com/vllm-project/vllm/issues/9565>.
|
||||
|
||||
4. When using EAGLE-3 based draft model, option "method" must be set to "eagle3".
|
||||
That is, to specify `"method": "eagle3"` in `speculative_config`.
|
||||
|
||||
A variety of EAGLE draft models are available on the Hugging Face hub:
|
||||
|
||||
| Base Model | EAGLE on Hugging Face | # EAGLE Parameters |
|
||||
|---------------------------------------------------------------------|-------------------------------------------|--------------------|
|
||||
| Vicuna-7B-v1.3 | yuhuili/EAGLE-Vicuna-7B-v1.3 | 0.24B |
|
||||
| Vicuna-13B-v1.3 | yuhuili/EAGLE-Vicuna-13B-v1.3 | 0.37B |
|
||||
| Vicuna-33B-v1.3 | yuhuili/EAGLE-Vicuna-33B-v1.3 | 0.56B |
|
||||
| LLaMA2-Chat 7B | yuhuili/EAGLE-llama2-chat-7B | 0.24B |
|
||||
| LLaMA2-Chat 13B | yuhuili/EAGLE-llama2-chat-13B | 0.37B |
|
||||
| LLaMA2-Chat 70B | yuhuili/EAGLE-llama2-chat-70B | 0.99B |
|
||||
| Mixtral-8x7B-Instruct-v0.1 | yuhuili/EAGLE-mixtral-instruct-8x7B | 0.28B |
|
||||
| LLaMA3-Instruct 8B | yuhuili/EAGLE-LLaMA3-Instruct-8B | 0.25B |
|
||||
| LLaMA3-Instruct 70B | yuhuili/EAGLE-LLaMA3-Instruct-70B | 0.99B |
|
||||
| Qwen2-7B-Instruct | yuhuili/EAGLE-Qwen2-7B-Instruct | 0.26B |
|
||||
| Qwen2-72B-Instruct | yuhuili/EAGLE-Qwen2-72B-Instruct | 1.05B |
|
||||
|
||||
## Lossless guarantees of Speculative Decoding
|
||||
|
||||
In vLLM, speculative decoding aims to enhance inference efficiency while maintaining accuracy. This section addresses the lossless guarantees of
|
||||
speculative decoding, breaking down the guarantees into three key areas:
|
||||
|
||||
1. **Theoretical Losslessness**
|
||||
\- Speculative decoding sampling is theoretically lossless up to the precision limits of hardware numerics. Floating-point errors might
|
||||
cause slight variations in output distributions, as discussed
|
||||
in [Accelerating Large Language Model Decoding with Speculative Sampling](https://arxiv.org/pdf/2302.01318)
|
||||
|
||||
2. **Algorithmic Losslessness**
|
||||
\- vLLM’s implementation of speculative decoding is algorithmically validated to be lossless. Key validation tests include:
|
||||
|
||||
> - **Rejection Sampler Convergence**: Ensures that samples from vLLM’s rejection sampler align with the target
|
||||
> distribution. [View Test Code](https://github.com/vllm-project/vllm/blob/47b65a550866c7ffbd076ecb74106714838ce7da/tests/samplers/test_rejection_sampler.py#L252)
|
||||
> - **Greedy Sampling Equality**: Confirms that greedy sampling with speculative decoding matches greedy sampling
|
||||
> without it. This verifies that vLLM's speculative decoding framework, when integrated with the vLLM forward pass and the vLLM rejection sampler,
|
||||
> provides a lossless guarantee. Almost all of the tests in [tests/spec_decode/e2e](../../tests/spec_decode/e2e).
|
||||
> verify this property using [this assertion implementation](https://github.com/vllm-project/vllm/blob/b67ae00cdbbe1a58ffc8ff170f0c8d79044a684a/tests/spec_decode/e2e/conftest.py#L291)
|
||||
|
||||
3. **vLLM Logprob Stability**
|
||||
\- vLLM does not currently guarantee stable token log probabilities (logprobs). This can result in different outputs for the
|
||||
same request across runs. For more details, see the FAQ section
|
||||
titled *Can the output of a prompt vary across runs in vLLM?* in the [FAQs](../../usage/faq.md).
|
||||
|
||||
While vLLM strives to ensure losslessness in speculative decoding, variations in generated outputs with and without speculative decoding
|
||||
can occur due to following factors:
|
||||
|
||||
- **Floating-Point Precision**: Differences in hardware numerical precision may lead to slight discrepancies in the output distribution.
|
||||
- **Batch Size and Numerical Stability**: Changes in batch size may cause variations in logprobs and output probabilities, potentially
|
||||
due to non-deterministic behavior in batched operations or numerical instability.
|
||||
|
||||
For mitigation strategies, please refer to the FAQ entry *Can the output of a prompt vary across runs in vLLM?* in the [FAQs](../../usage/faq.md).
|
||||
|
||||
## Resources for vLLM contributors
|
||||
|
||||
- [A Hacker's Guide to Speculative Decoding in vLLM](https://www.youtube.com/watch?v=9wNAgpX6z_4)
|
||||
- [What is Lookahead Scheduling in vLLM?](https://docs.google.com/document/d/1Z9TvqzzBPnh5WHcRwjvK2UEeFeq5zMZb5mFE8jR0HCs/edit#heading=h.1fjfb0donq5a)
|
||||
- [Information on batch expansion](https://docs.google.com/document/d/1T-JaS2T1NRfdP51qzqpyakoCXxSXTtORppiwaj5asxA/edit#heading=h.kk7dq05lc6q8)
|
||||
- [Dynamic speculative decoding](https://github.com/vllm-project/vllm/issues/4565)
|
||||
@@ -0,0 +1,62 @@
|
||||
# Speculative Decoding
|
||||
|
||||
This document shows how to use [Speculative Decoding](https://arxiv.org/pdf/2302.01318) with vLLM to reduce inter-token latency under medium-to-low QPS (query per second), memory-bound workloads.
|
||||
|
||||
To train your own draft models for optimized speculative decoding, see [vllm-project/speculators](speculators.md) for seamless training and integration with vLLM.
|
||||
|
||||
## vLLM Speculation Methods
|
||||
|
||||
vLLM supports a variety of methods of speculative decoding. Model-based methods such as EAGLE, draft models, and mlp provide the best latency reduction, while simpler methods such as n-gram and and suffix decoding provide modest speedups without increasing workload during peak traffic.
|
||||
|
||||
- [EAGLE](eagle.md)
|
||||
- [Draft Model](draft_model.md)
|
||||
- [Multi-Layer Perceptron](mlp.md)
|
||||
- [N-Gram](n_gram.md)
|
||||
- [Suffix Decoding](suffix.md)
|
||||
|
||||
## Lossless guarantees of Speculative Decoding
|
||||
|
||||
In vLLM, speculative decoding aims to enhance inference efficiency while maintaining accuracy. This section addresses the lossless guarantees of
|
||||
speculative decoding, breaking down the guarantees into three key areas:
|
||||
|
||||
1. **Theoretical Losslessness**
|
||||
\- Speculative decoding sampling is theoretically lossless up to the precision limits of hardware numerics. Floating-point errors might
|
||||
cause slight variations in output distributions, as discussed
|
||||
in [Accelerating Large Language Model Decoding with Speculative Sampling](https://arxiv.org/pdf/2302.01318)
|
||||
|
||||
2. **Algorithmic Losslessness**
|
||||
\- vLLM’s implementation of speculative decoding is algorithmically validated to be lossless. Key validation tests include:
|
||||
|
||||
> - **Rejection Sampler Convergence**: Ensures that samples from vLLM’s rejection sampler align with the target
|
||||
> distribution. [View Test Code](https://github.com/vllm-project/vllm/blob/47b65a550866c7ffbd076ecb74106714838ce7da/tests/samplers/test_rejection_sampler.py#L252)
|
||||
> - **Greedy Sampling Equality**: Confirms that greedy sampling with speculative decoding matches greedy sampling
|
||||
> without it. This verifies that vLLM's speculative decoding framework, when integrated with the vLLM forward pass and the vLLM rejection sampler,
|
||||
> provides a lossless guarantee. Almost all of the tests in [tests/spec_decode/e2e](/tests/v1/spec_decode).
|
||||
> verify this property using [this assertion implementation](https://github.com/vllm-project/vllm/blob/b67ae00cdbbe1a58ffc8ff170f0c8d79044a684a/tests/spec_decode/e2e/conftest.py#L291)
|
||||
|
||||
3. **vLLM Logprob Stability**
|
||||
\- vLLM does not currently guarantee stable token log probabilities (logprobs). This can result in different outputs for the
|
||||
same request across runs. For more details, see the FAQ section
|
||||
titled *Can the output of a prompt vary across runs in vLLM?* in the [FAQs](../../usage/faq.md).
|
||||
|
||||
While vLLM strives to ensure losslessness in speculative decoding, variations in generated outputs with and without speculative decoding
|
||||
can occur due to following factors:
|
||||
|
||||
- **Floating-Point Precision**: Differences in hardware numerical precision may lead to slight discrepancies in the output distribution.
|
||||
- **Batch Size and Numerical Stability**: Changes in batch size may cause variations in logprobs and output probabilities, potentially
|
||||
due to non-deterministic behavior in batched operations or numerical instability.
|
||||
|
||||
For mitigation strategies, please refer to the FAQ entry *Can the output of a prompt vary across runs in vLLM?* in the [FAQs](../../usage/faq.md).
|
||||
|
||||
## Known Feature Incompatibility
|
||||
|
||||
1. Pipeline parallelism is not composible with speculative decoding as of `vllm<=0.15.0`
|
||||
2. Speculative decoding with a draft models is not supported in `vllm<=0.10.0`
|
||||
|
||||
## Resources for vLLM contributors
|
||||
|
||||
- [[vLLM Office Hours #40] Intro to Speculators](https://www.youtube.com/watch?v=2ISAr_JVGLs)
|
||||
- [A Hacker's Guide to Speculative Decoding in vLLM](https://www.youtube.com/watch?v=9wNAgpX6z_4)
|
||||
- [What is Lookahead Scheduling in vLLM?](https://docs.google.com/document/d/1Z9TvqzzBPnh5WHcRwjvK2UEeFeq5zMZb5mFE8jR0HCs/edit#heading=h.1fjfb0donq5a)
|
||||
- [Information on batch expansion](https://docs.google.com/document/d/1T-JaS2T1NRfdP51qzqpyakoCXxSXTtORppiwaj5asxA/edit#heading=h.kk7dq05lc6q8)
|
||||
- [Dynamic speculative decoding](https://github.com/vllm-project/vllm/issues/4565)
|
||||
@@ -0,0 +1,80 @@
|
||||
# Draft Models
|
||||
|
||||
The following code configures vLLM in an offline mode to use speculative decoding with a draft model, speculating 5 tokens at a time.
|
||||
|
||||
```python
|
||||
from vllm import LLM, SamplingParams
|
||||
|
||||
prompts = ["The future of AI is"]
|
||||
sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
|
||||
|
||||
llm = LLM(
|
||||
model="Qwen/Qwen3-8B",
|
||||
tensor_parallel_size=1,
|
||||
speculative_config={
|
||||
"model": "Qwen/Qwen3-0.6B",
|
||||
"num_speculative_tokens": 5,
|
||||
"method": "draft_model",
|
||||
},
|
||||
)
|
||||
outputs = llm.generate(prompts, sampling_params)
|
||||
|
||||
for output in outputs:
|
||||
prompt = output.prompt
|
||||
generated_text = output.outputs[0].text
|
||||
print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
|
||||
```
|
||||
|
||||
To perform the equivalent launch in online mode, use the following server-side code:
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen3-4B-Thinking-2507 \
|
||||
--host 0.0.0.0 \
|
||||
--port 8000 \
|
||||
--seed 42 \
|
||||
-tp 1 \
|
||||
--max_model_len 2048 \
|
||||
--gpu_memory_utilization 0.8 \
|
||||
--speculative_config '{"model": "Qwen/Qwen3-0.6B", "num_speculative_tokens": 5, "method": "draft_model"}'
|
||||
```
|
||||
|
||||
The code used to request as completions as a client remains unchanged:
|
||||
|
||||
??? code
|
||||
|
||||
```python
|
||||
from openai import OpenAI
|
||||
|
||||
# Modify OpenAI's API key and API base to use vLLM's API server.
|
||||
openai_api_key = "EMPTY"
|
||||
openai_api_base = "http://localhost:8000/v1"
|
||||
|
||||
client = OpenAI(
|
||||
# defaults to os.environ.get("OPENAI_API_KEY")
|
||||
api_key=openai_api_key,
|
||||
base_url=openai_api_base,
|
||||
)
|
||||
|
||||
models = client.models.list()
|
||||
model = models.data[0].id
|
||||
|
||||
# Completion API
|
||||
stream = False
|
||||
completion = client.completions.create(
|
||||
model=model,
|
||||
prompt="The future of AI is",
|
||||
echo=False,
|
||||
n=1,
|
||||
stream=stream,
|
||||
)
|
||||
|
||||
print("Completion results:")
|
||||
if stream:
|
||||
for c in completion:
|
||||
print(c)
|
||||
else:
|
||||
print(completion)
|
||||
```
|
||||
|
||||
!!! warning
|
||||
Note: Please use `--speculative_config` to set all configurations related to speculative decoding. The previous method of specifying the model through `--speculative_model` and adding related parameters (e.g., `--num_speculative_tokens`) separately has been deprecated.
|
||||
@@ -0,0 +1,67 @@
|
||||
# EAGLE Draft Models
|
||||
|
||||
The following code configures vLLM to use speculative decoding where proposals are generated by an [EAGLE (Extrapolation Algorithm for Greater Language-model Efficiency)](https://arxiv.org/pdf/2401.15077) based draft model. A more detailed example for offline mode, including how to extract request level acceptance rate, can be found in [examples/offline_inference/spec_decode.py](../../../examples/offline_inference/spec_decode.py)
|
||||
|
||||
## Eagle Drafter Example
|
||||
|
||||
```python
|
||||
from vllm import LLM, SamplingParams
|
||||
|
||||
prompts = ["The future of AI is"]
|
||||
sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
|
||||
|
||||
llm = LLM(
|
||||
model="meta-llama/Meta-Llama-3-8B-Instruct",
|
||||
tensor_parallel_size=4,
|
||||
speculative_config={
|
||||
"model": "yuhuili/EAGLE-LLaMA3-Instruct-8B",
|
||||
"draft_tensor_parallel_size": 1,
|
||||
"num_speculative_tokens": 2,
|
||||
"method": "eagle",
|
||||
},
|
||||
)
|
||||
|
||||
outputs = llm.generate(prompts, sampling_params)
|
||||
|
||||
for output in outputs:
|
||||
prompt = output.prompt
|
||||
generated_text = output.outputs[0].text
|
||||
print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
|
||||
```
|
||||
|
||||
## Eagle3 Drafter Example
|
||||
|
||||
```python
|
||||
from vllm import LLM, SamplingParams
|
||||
|
||||
prompts = ["The future of AI is"]
|
||||
sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
|
||||
|
||||
llm = LLM(
|
||||
model="meta-llama/Meta-Llama-3-8B-Instruct",
|
||||
tensor_parallel_size=2,
|
||||
speculative_config={
|
||||
"model": "RedHatAI/Llama-3.1-8B-Instruct-speculator.eagle3",
|
||||
"draft_tensor_parallel_size": 2,
|
||||
"num_speculative_tokens": 2,
|
||||
"method": "eagle",
|
||||
},
|
||||
)
|
||||
|
||||
outputs = llm.generate(prompts, sampling_params)
|
||||
|
||||
for output in outputs:
|
||||
prompt = output.prompt
|
||||
generated_text = output.outputs[0].text
|
||||
print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
|
||||
```
|
||||
|
||||
## Pre-Trained Eagle Draft Models
|
||||
|
||||
A variety of EAGLE draft models are available on the Hugging Face hub:
|
||||
|
||||
* [RedHatAI/speculator-models](https://huggingface.co/collections/RedHatAI/speculator-models)
|
||||
* [yuhuili/models](https://huggingface.co/yuhuili/models?search=eagle)
|
||||
|
||||
!!! warning
|
||||
If you are using `vllm<0.7.0`, please use [this script](https://gist.github.com/abhigoyal1997/1e7a4109ccb7704fbc67f625e86b2d6d) to convert the speculative model and specify `"model": "path/to/modified/eagle/model"` in `speculative_config`.
|
||||
@@ -0,0 +1,42 @@
|
||||
# MLP Draft Models
|
||||
|
||||
The following code configures vLLM to use speculative decoding where proposals are generated by draft models that condition draft predictions on both context vectors and sampled tokens. For more information see [The Hitchhiker's Guide to Speculative Decoding](https://pytorch.org/blog/hitchhikers-guide-speculative-decoding/) and [IBM Research's Technical Report](https://arxiv.org/abs/2404.19124).
|
||||
|
||||
## MLP Drafter Example
|
||||
|
||||
```python
|
||||
from vllm import LLM, SamplingParams
|
||||
|
||||
prompts = ["The future of AI is"]
|
||||
sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
|
||||
|
||||
llm = LLM(
|
||||
model="meta-llama/Meta-Llama-3.1-70B-Instruct",
|
||||
tensor_parallel_size=4,
|
||||
speculative_config={
|
||||
"model": "ibm-ai-platform/llama3-70b-accelerator",
|
||||
"draft_tensor_parallel_size": 1,
|
||||
"method": "mlp_speculator",
|
||||
},
|
||||
)
|
||||
outputs = llm.generate(prompts, sampling_params)
|
||||
|
||||
for output in outputs:
|
||||
prompt = output.prompt
|
||||
generated_text = output.outputs[0].text
|
||||
print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
|
||||
```
|
||||
|
||||
## Pre-Trained MLP Drafter Models
|
||||
|
||||
A variety of speculative models of this type are available on HF hub:
|
||||
|
||||
- [llama-13b-accelerator](https://huggingface.co/ibm-ai-platform/llama-13b-accelerator)
|
||||
- [llama3-8b-accelerator](https://huggingface.co/ibm-ai-platform/llama3-8b-accelerator)
|
||||
- [codellama-34b-accelerator](https://huggingface.co/ibm-ai-platform/codellama-34b-accelerator)
|
||||
- [llama2-70b-accelerator](https://huggingface.co/ibm-ai-platform/llama2-70b-accelerator)
|
||||
- [llama3-70b-accelerator](https://huggingface.co/ibm-ai-platform/llama3-70b-accelerator)
|
||||
- [granite-3b-code-instruct-accelerator](https://huggingface.co/ibm-granite/granite-3b-code-instruct-accelerator)
|
||||
- [granite-8b-code-instruct-accelerator](https://huggingface.co/ibm-granite/granite-8b-code-instruct-accelerator)
|
||||
- [granite-7b-instruct-accelerator](https://huggingface.co/ibm-granite/granite-7b-instruct-accelerator)
|
||||
- [granite-20b-code-instruct-accelerator](https://huggingface.co/ibm-granite/granite-20b-code-instruct-accelerator)
|
||||
@@ -0,0 +1,27 @@
|
||||
# N-Gram Speculation
|
||||
|
||||
The following code configures vLLM to use speculative decoding where proposals are generated by
|
||||
matching n-grams in the prompt. For more information read [this thread.](https://x.com/joao_gante/status/1747322413006643259)
|
||||
|
||||
```python
|
||||
from vllm import LLM, SamplingParams
|
||||
|
||||
prompts = ["The future of AI is"]
|
||||
sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
|
||||
|
||||
llm = LLM(
|
||||
model="Qwen/Qwen3-8B",
|
||||
tensor_parallel_size=1,
|
||||
speculative_config={
|
||||
"method": "ngram",
|
||||
"num_speculative_tokens": 5,
|
||||
"prompt_lookup_max": 4,
|
||||
},
|
||||
)
|
||||
outputs = llm.generate(prompts, sampling_params)
|
||||
|
||||
for output in outputs:
|
||||
prompt = output.prompt
|
||||
generated_text = output.outputs[0].text
|
||||
print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
|
||||
```
|
||||
+4
-1
@@ -1,4 +1,7 @@
|
||||
# Speculators
|
||||
# vLLM-Project/Speculators
|
||||
|
||||

|
||||

|
||||
|
||||
[Speculators](https://docs.vllm.ai/projects/speculators/en/latest/) is a library for accelerating LLM inference through speculative decoding, providing efficient draft model training that integrates seamlessly with vLLM to reduce latency and improve throughput.
|
||||
|
||||
@@ -0,0 +1,35 @@
|
||||
# Suffix Decoding
|
||||
|
||||
The following code configures vLLM to use speculative decoding where proposals are generated using Suffix Decoding ([technical report](https://arxiv.org/abs/2411.04975)).
|
||||
|
||||
Like n-gram, Suffix Decoding can generate draft tokens by pattern-matching using the last `n` generated tokens. Unlike n-gram, Suffix Decoding (1) can pattern-match against both the prompt and previous generations, (2) uses frequency counts to propose the most likely continuations, and (3) speculates an adaptive number of tokens for each request at each iteration to get better acceptance rates.
|
||||
|
||||
Suffix Decoding can achieve better performance for tasks with high repetition, such as code-editing, agentic loops (e.g. self-reflection, self-consistency), and RL rollouts.
|
||||
|
||||
!!! tip "Install Arctic Inference"
|
||||
Suffix Decoding requires [Arctic Inference](https://github.com/snowflakedb/ArcticInference). You can install it with `pip install arctic-inference`.
|
||||
|
||||
!!! tip "Suffix Decoding Speculative Tokens"
|
||||
Suffix Decoding will speculate a dynamic number of tokens for each request at each decoding step, so the `num_speculative_tokens` configuration specifies the *maximum* number of speculative tokens. It is suggested to use a high number such as `16` or `32` (default).
|
||||
|
||||
```python
|
||||
from vllm import LLM, SamplingParams
|
||||
|
||||
prompts = ["The future of AI is"]
|
||||
sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
|
||||
|
||||
llm = LLM(
|
||||
model="Qwen/Qwen3-8B",
|
||||
tensor_parallel_size=1,
|
||||
speculative_config={
|
||||
"method": "suffix",
|
||||
"num_speculative_tokens": 32,
|
||||
},
|
||||
)
|
||||
outputs = llm.generate(prompts, sampling_params)
|
||||
|
||||
for output in outputs:
|
||||
prompt = output.prompt
|
||||
generated_text = output.outputs[0].text
|
||||
print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
|
||||
```
|
||||
@@ -105,6 +105,10 @@ plugins:
|
||||
- https://numpy.org/doc/stable/objects.inv
|
||||
- https://pytorch.org/docs/stable/objects.inv
|
||||
- https://psutil.readthedocs.io/en/stable/objects.inv
|
||||
- redirects:
|
||||
redirect_maps:
|
||||
features/spec_decode/README.md: features/speculative_decoding/README.md
|
||||
features/spec_decode/speculators.md: features/speculative_decoding/speculators.md
|
||||
|
||||
markdown_extensions:
|
||||
- attr_list
|
||||
|
||||
@@ -113,7 +113,6 @@ markers = [
|
||||
"cpu_test: mark test as CPU-only test",
|
||||
"split: run this test as part of a split",
|
||||
"distributed: run this test only in distributed GPU tests",
|
||||
"skip_v1: do not run this test with v1",
|
||||
"optional: optional tests that are automatically skipped, include --optional to run them",
|
||||
]
|
||||
|
||||
|
||||
@@ -7,6 +7,7 @@ mkdocs-awesome-nav
|
||||
mkdocs-glightbox
|
||||
mkdocs-git-revision-date-localized-plugin
|
||||
mkdocs-minify-plugin
|
||||
mkdocs-redirects
|
||||
regex
|
||||
ruff
|
||||
pydantic
|
||||
|
||||
@@ -1,2 +1,3 @@
|
||||
lmcache >= 0.3.9
|
||||
nixl >= 0.7.1 # Required for disaggregated prefill
|
||||
mooncake-transfer-engine >= 0.3.8
|
||||
|
||||
@@ -977,9 +977,9 @@ if _is_cuda():
|
||||
# FA3 requires CUDA 12.3 or later
|
||||
ext_modules.append(CMakeExtension(name="vllm.vllm_flash_attn._vllm_fa3_C"))
|
||||
if envs.VLLM_USE_PRECOMPILED or (
|
||||
CUDA_HOME and get_nvcc_cuda_version() >= Version("12.9")
|
||||
CUDA_HOME and get_nvcc_cuda_version() >= Version("12.8")
|
||||
):
|
||||
# FlashMLA requires CUDA 12.9 or later
|
||||
# FlashMLA requires CUDA 12.8 or later
|
||||
# Optional since this doesn't get built (produce an .so file) when
|
||||
# not targeting a hopper system
|
||||
ext_modules.append(CMakeExtension(name="vllm._flashmla_C", optional=True))
|
||||
|
||||
@@ -50,10 +50,9 @@ def test_tp1_fp8_fusions(
|
||||
run_e2e_fusion_test,
|
||||
monkeypatch,
|
||||
):
|
||||
if use_deepgemm:
|
||||
# TODO(luka/eliza) DeepGEMM uses different quants, matching not supported
|
||||
if use_deepgemm and is_blackwell():
|
||||
# TODO(luka) DeepGEMM uses different quants, matching not supported
|
||||
# - on Blackwell, uses a special quant fp8, currently not supported
|
||||
# - on Hopper, tma-aligned scales inhibit matching (fix WIP)
|
||||
pytest.skip("DeepGEMM & quant matching not currently supported")
|
||||
|
||||
matches = matches_fn(n_layers)
|
||||
@@ -66,7 +65,6 @@ def test_tp1_fp8_fusions(
|
||||
model_kwargs["hf_overrides"] = hf_overrides(n_layers)
|
||||
model_kwargs["load_format"] = "dummy"
|
||||
model_kwargs["max_model_len"] = 1024
|
||||
|
||||
compilation_config = dict(
|
||||
use_inductor_graph_partition=inductor_graph_partition,
|
||||
custom_ops=custom_ops.split(","),
|
||||
|
||||
@@ -7,7 +7,6 @@ from vllm.entrypoints.llm import LLM
|
||||
from vllm.sampling_params import SamplingParams
|
||||
|
||||
|
||||
@pytest.mark.skip_v1
|
||||
@pytest.mark.parametrize("model", ["distilbert/distilgpt2"])
|
||||
def test_computed_prefix_blocks(model: str):
|
||||
# This test checks if the engine generates completions both with and
|
||||
|
||||
@@ -124,6 +124,8 @@ def test_init_weight_transfer_engine_calls_engine():
|
||||
if torch.cuda.device_count() < 1:
|
||||
pytest.skip("Need at least 1 GPU for this test")
|
||||
|
||||
# Run in-process so mock.patch works (spawn won't inherit the mock)
|
||||
os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0"
|
||||
# Enable insecure serialization to allow pickling functions for collective_rpc
|
||||
os.environ["VLLM_ALLOW_INSECURE_SERIALIZATION"] = "1"
|
||||
|
||||
@@ -171,6 +173,8 @@ def test_update_weights_calls_engine():
|
||||
if torch.cuda.device_count() < 1:
|
||||
pytest.skip("Need at least 1 GPU for this test")
|
||||
|
||||
# Run in-process so mock.patch works (spawn won't inherit the mock)
|
||||
os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0"
|
||||
# Enable insecure serialization to allow pickling functions for collective_rpc
|
||||
os.environ["VLLM_ALLOW_INSECURE_SERIALIZATION"] = "1"
|
||||
|
||||
@@ -228,6 +232,8 @@ def test_full_weight_transfer_flow():
|
||||
if torch.cuda.device_count() < 1:
|
||||
pytest.skip("Need at least 1 GPU for this test")
|
||||
|
||||
# Run in-process so mock.patch works (spawn won't inherit the mock)
|
||||
os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0"
|
||||
# Enable insecure serialization to allow pickling functions for collective_rpc
|
||||
os.environ["VLLM_ALLOW_INSECURE_SERIALIZATION"] = "1"
|
||||
|
||||
|
||||
@@ -0,0 +1,8 @@
|
||||
model_name: "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8"
|
||||
accuracy_threshold: 0.29
|
||||
num_questions: 1319
|
||||
num_fewshot: 5
|
||||
server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2"
|
||||
env:
|
||||
VLLM_USE_FLASHINFER_MOE_FP8: "1"
|
||||
VLLM_FLASHINFER_MOE_BACKEND: "latency"
|
||||
+8
@@ -0,0 +1,8 @@
|
||||
model_name: "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4"
|
||||
accuracy_threshold: 0.29
|
||||
num_questions: 1319
|
||||
num_fewshot: 5
|
||||
server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2"
|
||||
env:
|
||||
VLLM_USE_FLASHINFER_MOE_FP4: "1"
|
||||
VLLM_FLASHINFER_MOE_BACKEND: "throughput"
|
||||
@@ -13,3 +13,5 @@ Llama-4-Scout-BF16-fi-cutlass.yaml
|
||||
Llama-4-Scout-BF16-triton.yaml
|
||||
Mixtral-8x7B-BF16-fi-cutlass.yaml
|
||||
Mixtral-8x7B-BF16-triton.yaml
|
||||
Nemotron-Nano-30B-Fp8-ModelOpt-fi-trtllm.yaml
|
||||
Nemotron-Nano-30B-NvFp4-ModelOpt-fi-cutlass.yaml
|
||||
|
||||
@@ -10,7 +10,7 @@ from vllm.utils.math_utils import next_power_of_2
|
||||
from vllm.utils.torch_utils import set_random_seed
|
||||
from vllm.v1.attention.ops.triton_unified_attention import unified_attention
|
||||
|
||||
NUM_HEADS = [(4, 4), (8, 2)]
|
||||
NUM_HEADS = [(4, 4), (8, 2), (5, 1)]
|
||||
HEAD_SIZES = [128, 256]
|
||||
BLOCK_SIZES = [16]
|
||||
|
||||
@@ -20,6 +20,8 @@ QDTYPES = (
|
||||
if not current_platform.is_rocm()
|
||||
else [None, torch.float8_e4m3fnuz]
|
||||
)
|
||||
FP8_DTYPE = current_platform.fp8_dtype()
|
||||
|
||||
# one value large enough to test overflow in index calculation.
|
||||
# one value small enough to test the schema op check
|
||||
NUM_BLOCKS = [32768, 2048]
|
||||
@@ -217,3 +219,127 @@ def test_triton_unified_attn(
|
||||
torch.testing.assert_close(output, ref_output, atol=atol, rtol=rtol),
|
||||
f"{torch.max(torch.abs(output - ref_output))}",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"seq_lens",
|
||||
[
|
||||
[(1, 1328), (5, 18), (129, 463)],
|
||||
[(1, 523), (1, 37), (1, 2011)],
|
||||
[(1, 1)] * 533,
|
||||
[(533, 533)] * 533,
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("num_heads", NUM_HEADS)
|
||||
@pytest.mark.parametrize("head_size", HEAD_SIZES)
|
||||
@pytest.mark.parametrize("block_size", BLOCK_SIZES)
|
||||
@pytest.mark.parametrize("sliding_window", [None, 64, 128, 256])
|
||||
@pytest.mark.parametrize("soft_cap", [None, 50.0])
|
||||
@pytest.mark.parametrize("num_blocks", NUM_BLOCKS)
|
||||
@pytest.mark.parametrize("seq_threshold_3D", SEQ_THRESHOLD_3D_VALUES)
|
||||
@torch.inference_mode()
|
||||
def test_triton_unified_attn_fp16_input_fp8_output(
|
||||
seq_lens: list[tuple[int, int]],
|
||||
num_heads: tuple[int, int],
|
||||
head_size: int,
|
||||
sliding_window: int | None,
|
||||
block_size: int,
|
||||
soft_cap: float | None,
|
||||
num_blocks: int,
|
||||
seq_threshold_3D: int,
|
||||
) -> None:
|
||||
"""Test with fp16 input and fp8 output using output_scale."""
|
||||
torch.set_default_device("cuda")
|
||||
|
||||
set_random_seed(0)
|
||||
num_seqs = len(seq_lens)
|
||||
query_lens = [x[0] for x in seq_lens]
|
||||
kv_lens = [x[1] for x in seq_lens]
|
||||
num_query_heads = num_heads[0]
|
||||
num_kv_heads = num_heads[1]
|
||||
assert num_query_heads % num_kv_heads == 0
|
||||
max_query_len = max(query_lens)
|
||||
max_kv_len = max(kv_lens)
|
||||
window_size = (sliding_window - 1, 0) if sliding_window is not None else (-1, -1)
|
||||
scale = head_size**-0.5
|
||||
|
||||
dtype = torch.float16
|
||||
query = torch.randn(sum(query_lens), num_query_heads, head_size, dtype=dtype)
|
||||
key_cache = torch.randn(
|
||||
num_blocks, block_size, num_kv_heads, head_size, dtype=dtype
|
||||
)
|
||||
value_cache = torch.randn_like(key_cache)
|
||||
cu_query_lens = torch.tensor([0] + query_lens, dtype=torch.int32).cumsum(
|
||||
dim=0, dtype=torch.int32
|
||||
)
|
||||
kv_lens_tensor = torch.tensor(kv_lens, dtype=torch.int32)
|
||||
|
||||
max_num_blocks_per_seq = (max_kv_len + block_size - 1) // block_size
|
||||
block_tables = torch.randint(
|
||||
0, num_blocks, (num_seqs, max_num_blocks_per_seq), dtype=torch.int32
|
||||
)
|
||||
|
||||
output = torch.empty(sum(query_lens), num_query_heads, head_size, dtype=FP8_DTYPE)
|
||||
|
||||
output_scale = torch.tensor(0.5, dtype=torch.float32)
|
||||
|
||||
num_par_softmax_segments = 16
|
||||
head_size_padded = next_power_of_2(head_size)
|
||||
softmax_segm_output = torch.empty(
|
||||
(seq_threshold_3D, num_query_heads, num_par_softmax_segments, head_size_padded),
|
||||
dtype=torch.float32,
|
||||
)
|
||||
softmax_segm_max = torch.empty(
|
||||
(seq_threshold_3D, num_query_heads, num_par_softmax_segments),
|
||||
dtype=torch.float32,
|
||||
)
|
||||
softmax_segm_expsum = torch.empty(
|
||||
(seq_threshold_3D, num_query_heads, num_par_softmax_segments),
|
||||
dtype=torch.float32,
|
||||
)
|
||||
|
||||
unified_attention(
|
||||
q=query,
|
||||
k=key_cache,
|
||||
v=value_cache,
|
||||
out=output,
|
||||
cu_seqlens_q=cu_query_lens,
|
||||
seqused_k=kv_lens_tensor,
|
||||
max_seqlen_q=max_query_len,
|
||||
max_seqlen_k=max_kv_len,
|
||||
softmax_scale=scale,
|
||||
causal=True,
|
||||
window_size=window_size,
|
||||
block_table=block_tables,
|
||||
softcap=soft_cap if soft_cap is not None else 0,
|
||||
q_descale=None,
|
||||
k_descale=None,
|
||||
v_descale=None,
|
||||
output_scale=output_scale,
|
||||
seq_threshold_3D=seq_threshold_3D,
|
||||
num_par_softmax_segments=num_par_softmax_segments,
|
||||
softmax_segm_output=softmax_segm_output,
|
||||
softmax_segm_max=softmax_segm_max,
|
||||
softmax_segm_expsum=softmax_segm_expsum,
|
||||
)
|
||||
|
||||
ref_output = ref_paged_attn(
|
||||
query=query,
|
||||
key_cache=key_cache,
|
||||
value_cache=value_cache,
|
||||
query_lens=query_lens,
|
||||
kv_lens=kv_lens,
|
||||
block_tables=block_tables,
|
||||
scale=scale,
|
||||
sliding_window=sliding_window,
|
||||
soft_cap=soft_cap,
|
||||
)
|
||||
|
||||
output_fp16 = output.to(torch.float32) * output_scale.item()
|
||||
output_fp16 = output_fp16.to(torch.float16)
|
||||
|
||||
atol, rtol = 2e-1, 2e-1
|
||||
(
|
||||
torch.testing.assert_close(output_fp16, ref_output, atol=atol, rtol=rtol),
|
||||
f"{torch.max(torch.abs(output_fp16 - ref_output))}",
|
||||
)
|
||||
|
||||
@@ -2,6 +2,8 @@
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
|
||||
import itertools
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
@@ -21,7 +23,7 @@ QUANT_DTYPES = [torch.int8, current_platform.fp8_dtype()]
|
||||
VEC_HIDDEN_SIZES = [1024, 1025, 1027, 1029]
|
||||
# Avoid combinatorial explosion with full Cartesian product
|
||||
NUM_TOKENS_HIDDEN_SIZES = [
|
||||
*[(1, i) for i in [1, 64, *VEC_HIDDEN_SIZES, 5120, 5137]],
|
||||
*[(1, i) for i in [1, 64, 128, *VEC_HIDDEN_SIZES, 5120, 5137]],
|
||||
*[(2048, i) for i in [1, 64, *VEC_HIDDEN_SIZES, 5137]],
|
||||
*[(4096, i) for i in [1, 64, 5137]],
|
||||
]
|
||||
@@ -29,6 +31,7 @@ NUM_TOKENS_HIDDEN_SIZES = [
|
||||
ADD_RESIDUAL = [False, True]
|
||||
SCALE_UBS = [True, False]
|
||||
GROUP_SIZES = [None, [1, 64], [1, 128]]
|
||||
TMA_ALIGNMENTS = [0, 4]
|
||||
SEEDS = [0]
|
||||
CUDA_DEVICES = [f"cuda:{i}" for i in range(1 if torch.cuda.device_count() == 1 else 2)]
|
||||
|
||||
@@ -110,12 +113,21 @@ def ops_dynamic_per_token_or_block_quant(
|
||||
residual: torch.Tensor | None,
|
||||
scale_ub: torch.Tensor | None,
|
||||
group_size: list[int] | None,
|
||||
tma_alignment: int,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
|
||||
if residual is not None:
|
||||
residual = residual.clone()
|
||||
if group_size is not None:
|
||||
out, scales = ops.rms_norm_per_block_quant(
|
||||
x, weight, EPS, quant_dtype, group_size, scale_ub, residual, True
|
||||
x,
|
||||
weight,
|
||||
EPS,
|
||||
quant_dtype,
|
||||
group_size,
|
||||
scale_ub,
|
||||
residual,
|
||||
True,
|
||||
tma_alignment,
|
||||
)
|
||||
scales = scales.contiguous()
|
||||
else:
|
||||
@@ -132,9 +144,10 @@ def ops_impl(
|
||||
residual: torch.Tensor | None,
|
||||
scale_ub: torch.Tensor | None,
|
||||
group_size: list[int] | None,
|
||||
tma_alignment: int,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
|
||||
return ops_dynamic_per_token_or_block_quant(
|
||||
weight, x, quant_dtype, residual, scale_ub, group_size
|
||||
weight, x, quant_dtype, residual, scale_ub, group_size, tma_alignment
|
||||
)
|
||||
|
||||
|
||||
@@ -143,7 +156,10 @@ def ops_impl(
|
||||
@pytest.mark.parametrize("has_scale_ub", SCALE_UBS)
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("quant_dtype", QUANT_DTYPES)
|
||||
@pytest.mark.parametrize("group_size", GROUP_SIZES)
|
||||
@pytest.mark.parametrize(
|
||||
"group_size, tma_alignment",
|
||||
[(None, 0), *itertools.product(GROUP_SIZES, TMA_ALIGNMENTS)],
|
||||
)
|
||||
@pytest.mark.parametrize("seed", SEEDS)
|
||||
@pytest.mark.parametrize("device", CUDA_DEVICES)
|
||||
@torch.inference_mode()
|
||||
@@ -156,6 +172,7 @@ def test_rms_norm(
|
||||
dtype: torch.dtype,
|
||||
quant_dtype: torch.dtype,
|
||||
group_size: list[int] | None,
|
||||
tma_alignment: int,
|
||||
seed: int,
|
||||
device: str,
|
||||
) -> None:
|
||||
@@ -173,6 +190,20 @@ def test_rms_norm(
|
||||
# blockwise baseline doesn't support scale_ub
|
||||
return
|
||||
|
||||
if (
|
||||
group_size is None or quant_dtype != current_platform.fp8_dtype()
|
||||
) and tma_alignment != 0:
|
||||
# TMA alignment is only supported for groupwise fp8 kernels
|
||||
return
|
||||
|
||||
if (
|
||||
group_size is not None
|
||||
and tma_alignment != 0
|
||||
and hidden_size // group_size[1] % tma_alignment == 0
|
||||
):
|
||||
# Skip tests where TMA alignment doesn't create extra padding to save time
|
||||
return
|
||||
|
||||
if has_scale_ub and quant_dtype != current_platform.fp8_dtype():
|
||||
# skip
|
||||
return
|
||||
@@ -196,7 +227,7 @@ def test_rms_norm(
|
||||
layer, x, quant_dtype, residual, scale_ub, group_size
|
||||
)
|
||||
ops_out, ops_scales, ops_residual = ops_impl(
|
||||
layer.weight, x, quant_dtype, residual, scale_ub, group_size
|
||||
layer.weight, x, quant_dtype, residual, scale_ub, group_size, tma_alignment
|
||||
)
|
||||
|
||||
assert ref_out.dtype == quant_dtype
|
||||
|
||||
@@ -398,80 +398,3 @@ def test_convert_moe_weights_to_flashinfer_trtllm_block_layout(
|
||||
|
||||
assert w13_converted.shape[0] == num_experts
|
||||
assert w2_converted.shape[0] == num_experts
|
||||
|
||||
|
||||
def test_flashinfer_blockscale_fp8_none_expert_group(monkeypatch):
|
||||
"""Test that flashinfer_fused_moe_blockscale_fp8 handles num_expert_group=None.
|
||||
|
||||
Regression test for https://github.com/vllm-project/vllm/issues/34477
|
||||
MiniMax-M2.1 uses sigmoid scoring with e_score_correction_bias but no
|
||||
grouped top-k, resulting in num_expert_group=None. This triggered a crash
|
||||
in the flashinfer kernel when DeepSeekV3 routing was selected.
|
||||
"""
|
||||
if not current_platform.has_device_capability(100):
|
||||
pytest.skip("Test requires SM >= 100 (Blackwell)")
|
||||
|
||||
import vllm.model_executor.layers.fused_moe.flashinfer_trtllm_moe # noqa: E501, F401
|
||||
from tests.kernels.quant_utils import native_per_token_group_quant_fp8
|
||||
|
||||
set_random_seed(7)
|
||||
monkeypatch.setenv("VLLM_FUSED_MOE_CHUNK_SIZE", "8192")
|
||||
|
||||
e = 16 # num_experts (must be divisible by 4)
|
||||
topk = 6 # top_k > 1 triggers DeepSeekV3 routing with sigmoid
|
||||
m, n, k = 10, 4096, 5120
|
||||
block_shape = [128, 128]
|
||||
block_k = block_shape[1]
|
||||
|
||||
with set_current_vllm_config(vllm_config):
|
||||
# Create BF16 hidden states
|
||||
x = torch.randn((m, k), device="cuda", dtype=torch.bfloat16) / 10
|
||||
|
||||
# Create FP8 block-scale quantized weights
|
||||
w13_bf16 = torch.randn((e, 2 * n, k), device="cuda", dtype=torch.bfloat16) / 10
|
||||
w2_bf16 = torch.randn((e, k, n), device="cuda", dtype=torch.bfloat16) / 10
|
||||
|
||||
# Quantize weights per-block to FP8
|
||||
w13_fp8_list, w13_scale_list = [], []
|
||||
w2_fp8_list, w2_scale_list = [], []
|
||||
for i in range(e):
|
||||
wq, ws = native_per_token_group_quant_fp8(w13_bf16[i], block_k)
|
||||
w13_fp8_list.append(wq)
|
||||
w13_scale_list.append(ws)
|
||||
|
||||
wq, ws = native_per_token_group_quant_fp8(w2_bf16[i], block_k)
|
||||
w2_fp8_list.append(wq)
|
||||
w2_scale_list.append(ws)
|
||||
|
||||
w13_fp8 = torch.stack(w13_fp8_list)
|
||||
w13_scale = torch.stack(w13_scale_list)
|
||||
w2_fp8 = torch.stack(w2_fp8_list)
|
||||
w2_scale = torch.stack(w2_scale_list)
|
||||
|
||||
# DeepSeekV3 routing uses float32 logits + optional bias
|
||||
routing_logits = torch.randn((m, e), device="cuda", dtype=torch.float32)
|
||||
routing_bias = torch.randn(e, device="cuda", dtype=torch.float32)
|
||||
|
||||
# This should NOT crash with num_expert_group=None
|
||||
output = torch.ops.vllm.flashinfer_fused_moe_blockscale_fp8(
|
||||
routing_logits=routing_logits,
|
||||
routing_bias=routing_bias,
|
||||
x=x,
|
||||
w13_weight=w13_fp8,
|
||||
w13_weight_scale_inv=w13_scale,
|
||||
w2_weight=w2_fp8,
|
||||
w2_weight_scale_inv=w2_scale,
|
||||
global_num_experts=e,
|
||||
top_k=topk,
|
||||
num_expert_group=None,
|
||||
topk_group=None,
|
||||
intermediate_size=n,
|
||||
expert_offset=0,
|
||||
local_num_experts=e,
|
||||
block_shape=block_shape,
|
||||
routing_method_type=RoutingMethodType.DeepSeekV3,
|
||||
routed_scaling=1.0,
|
||||
)
|
||||
|
||||
assert output is not None
|
||||
assert output.shape == (m, k)
|
||||
|
||||
@@ -9,6 +9,7 @@ from vllm.model_executor.layers.fused_moe.oracle.unquantized import (
|
||||
UnquantizedMoeBackend,
|
||||
select_unquantized_moe_backend,
|
||||
)
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
@@ -65,6 +66,9 @@ def test_select_default_backend_by_platform(
|
||||
"vllm.model_executor.layers.fused_moe.oracle.unquantized.is_supported_config_trtllm_bf16",
|
||||
return_value=(True, None),
|
||||
)
|
||||
@pytest.mark.skipif(
|
||||
not current_platform.is_cuda(), reason="Only supported on NVIDIA platforms."
|
||||
)
|
||||
def test_select_cuda_flashinfer_trtllm_backend(
|
||||
mock_has_flashinfer, mock_is_supported_trtllm, monkeypatch
|
||||
):
|
||||
@@ -101,6 +105,9 @@ def test_select_cuda_flashinfer_trtllm_backend(
|
||||
"vllm.model_executor.layers.fused_moe.oracle.unquantized.is_supported_config_trtllm_bf16",
|
||||
return_value=(False, None),
|
||||
)
|
||||
@pytest.mark.skipif(
|
||||
not current_platform.is_cuda(), reason="Only supported on NVIDIA platforms."
|
||||
)
|
||||
def test_select_cuda_flashinfer_cutlass_backend(
|
||||
mock_has_flashinfer, mock_is_supported_trtllm, monkeypatch
|
||||
):
|
||||
|
||||
@@ -155,9 +155,9 @@ def test_rocm_wvsplitkrc_kernel(xnorm, n, k, m, dtype, seed, bias_mode):
|
||||
out = ops.wvSplitKrc(B, A.view(-1, A.size(-1)), cu_count, BIAS)
|
||||
|
||||
if xnorm:
|
||||
assert torch.allclose(out, ref_out, atol=1e-3, rtol=1e-8)
|
||||
torch.testing.assert_close(out, ref_out, atol=1e-3, rtol=1e-8)
|
||||
else:
|
||||
assert torch.allclose(out, ref_out, atol=1e-3, rtol=1e-2)
|
||||
torch.testing.assert_close(out, ref_out, atol=1e-3, rtol=1e-2)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("n,k,m", NKM_FACTORS_LLMM1)
|
||||
@@ -177,7 +177,7 @@ def test_rocm_llmm1_kernel(n, k, m, dtype, rows_per_block, seed):
|
||||
ref_out = torch.matmul(A, B.t())
|
||||
out = ops.LLMM1(B, A, rows_per_block)
|
||||
|
||||
assert torch.allclose(out, ref_out, rtol=0.01)
|
||||
torch.testing.assert_close(out, ref_out, atol=1e-8, rtol=1e-2)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("n,k,m", NKM_FACTORS_WVSPLITK)
|
||||
@@ -194,7 +194,7 @@ def test_rocm_wvsplitk_kernel(n, k, m, dtype, seed):
|
||||
ref_out = torch.nn.functional.linear(A, B)
|
||||
out = ops.wvSplitK(B, A.view(-1, A.size(-1)), cu_count)
|
||||
|
||||
assert torch.allclose(out, ref_out, rtol=0.01)
|
||||
torch.testing.assert_close(out, ref_out, atol=1e-8, rtol=1e-2)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("n,k,m", NKM_FACTORS_WVSPLITK)
|
||||
@@ -213,7 +213,7 @@ def test_rocm_wvsplitk_bias1D_kernel(n, k, m, dtype, seed):
|
||||
ref_out = torch.nn.functional.linear(A, B, BIAS)
|
||||
out = ops.wvSplitK(B, A.view(-1, A.size(-1)), cu_count, BIAS)
|
||||
|
||||
assert torch.allclose(out, ref_out, rtol=0.01)
|
||||
torch.testing.assert_close(out, ref_out, atol=1e-8, rtol=1e-2)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("n,k,m", NKM_FACTORS_WVSPLITK)
|
||||
@@ -232,7 +232,7 @@ def test_rocm_wvsplitk_bias2D_kernel(n, k, m, dtype, seed):
|
||||
ref_out = torch.nn.functional.linear(A, B, BIAS)
|
||||
out = ops.wvSplitK(B, A.view(-1, A.size(-1)), cu_count, BIAS)
|
||||
|
||||
assert torch.allclose(out, ref_out, rtol=0.01)
|
||||
torch.testing.assert_close(out, ref_out, atol=1e-8, rtol=1e-2)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("xnorm", [False, True])
|
||||
@@ -275,4 +275,4 @@ def test_rocm_wvsplitk_fp8_kernel(
|
||||
# wider pytrch thresh for large-K & no xnorm
|
||||
torch.testing.assert_close(out, ref_out, atol=0.07, rtol=5e-2)
|
||||
else:
|
||||
torch.testing.assert_close(out, ref_out, atol=0.01, rtol=0.01)
|
||||
torch.testing.assert_close(out, ref_out, atol=1e-2, rtol=1e-2)
|
||||
|
||||
@@ -153,5 +153,5 @@ def test_default_mm_lora_does_not_expand_string_reqs(vllm_runner):
|
||||
# Then check to make sure the submitted lora request
|
||||
# and text prompt were zipped together correctly
|
||||
engine_args, engine_kwargs = mock_add_request.call_args
|
||||
assert engine_args[1]["prompt"] == AUDIO_PROMPT
|
||||
assert engine_kwargs["lora_request"] is None
|
||||
assert engine_kwargs["prompt_text"] == AUDIO_PROMPT
|
||||
|
||||
@@ -88,9 +88,8 @@ class Qwen2VLTester:
|
||||
# Validate outputs
|
||||
for generated, expected in zip(generated_texts, expected_outputs):
|
||||
assert expected.startswith(generated), (
|
||||
f"Generated text {generated} doesn't "
|
||||
f"Generated text {generated} doesn't match expected pattern {expected}"
|
||||
)
|
||||
f"match expected pattern {expected}"
|
||||
|
||||
def run_beam_search_test(
|
||||
self,
|
||||
@@ -118,11 +117,14 @@ class Qwen2VLTester:
|
||||
inputs, beam_search_params, lora_request=lora_request
|
||||
)
|
||||
|
||||
for output_obj, expected_outs in zip(outputs, expected_outputs):
|
||||
for output_obj, expected_texts in zip(outputs, expected_outputs):
|
||||
output_texts = [seq.text for seq in output_obj.sequences]
|
||||
assert output_texts == expected_outs, (
|
||||
f"Generated texts {output_texts} do not match expected {expected_outs}"
|
||||
) # noqa: E501
|
||||
|
||||
for output_text, expected_text in zip(output_texts, expected_texts):
|
||||
# NOTE beam search .text contains the whole text including inputs
|
||||
assert output_text.endswith(expected_text), (
|
||||
f"Generated {output_text} does not match expected {expected_text}"
|
||||
)
|
||||
|
||||
|
||||
TEST_IMAGES = [
|
||||
@@ -151,11 +153,10 @@ EXPECTED_OUTPUTS_VISION_NO_CONNECTOR = [
|
||||
"A closeup shot of the Tokyo Skytree with pink flowers in the foreground.",
|
||||
]
|
||||
|
||||
# NOTE - beam search .text contains the whole text
|
||||
EXPECTED_BEAM_SEARCH_OUTPUTS = [
|
||||
[
|
||||
"<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>What is in the image?<|im_end|>\n<|im_start|>assistant\nA majestic skyscraper stands", # noqa: E501
|
||||
"<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>What is in the image?<|im_end|>\n<|im_start|>assistant\nA majestic tower stands tall", # noqa: E501
|
||||
"A majestic skyscraper stands",
|
||||
"A majestic tower stands tall",
|
||||
],
|
||||
]
|
||||
|
||||
|
||||
@@ -169,13 +169,6 @@ VLM_TEST_SETTINGS = {
|
||||
auto_cls=AutoModelForImageTextToText,
|
||||
vllm_output_post_proc=model_utils.qwen2_vllm_to_hf_output,
|
||||
patch_hf_runner=model_utils.qwen3_vl_patch_hf_runner,
|
||||
vllm_runner_kwargs={
|
||||
"attention_config": {
|
||||
"backend": "ROCM_AITER_FA",
|
||||
},
|
||||
}
|
||||
if current_platform.is_rocm()
|
||||
else None,
|
||||
image_size_factors=[(0.25,), (0.25, 0.25, 0.25), (0.25, 0.2, 0.15)],
|
||||
marks=[
|
||||
pytest.mark.core_model,
|
||||
|
||||
@@ -7,8 +7,10 @@ import mimetypes
|
||||
import os
|
||||
from tempfile import NamedTemporaryFile, TemporaryDirectory
|
||||
|
||||
import aiohttp
|
||||
import numpy as np
|
||||
import pytest
|
||||
import requests
|
||||
import torch
|
||||
from PIL import Image, ImageChops
|
||||
|
||||
@@ -318,3 +320,58 @@ async def test_allowed_media_domains(video_url: str, num_frames: int):
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
_, _ = await connector.fetch_video_async(disallowed_url)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ssrf_bypass_backslash_in_url(local_asset_server):
|
||||
"""Verify that backslash-@ URL parsing confusion cannot bypass the
|
||||
allowed_media_domains check (GHSA-v359-jj2v-j536).
|
||||
|
||||
urllib3.parse_url() and aiohttp/yarl disagree on how to parse a
|
||||
backslash before ``@``. urllib3 treats ``\\`` as part of the path
|
||||
(encoding it as ``%5C``), while yarl treats it as a userinfo
|
||||
separator, changing the effective host. The fix normalises the URL
|
||||
through urllib3 *before* handing it to aiohttp so both layers agree.
|
||||
"""
|
||||
port = local_asset_server.port
|
||||
asset = TEST_IMAGE_ASSETS[0]
|
||||
|
||||
# Craft the bypass payload: urllib3 sees host=127.0.0.1, but an
|
||||
# un-patched aiohttp would see host=example.com.
|
||||
bypass_url = f"http://127.0.0.1:{port}\\@example.com/{asset}"
|
||||
|
||||
connector = MediaConnector(
|
||||
allowed_media_domains=["127.0.0.1"],
|
||||
)
|
||||
|
||||
# After the fix the request is made to 127.0.0.1 (the local asset
|
||||
# server) using the normalised URL. The normalised path will be
|
||||
# /%5C@example.com/<asset> which won't match any file the server
|
||||
# knows about, so we expect an HTTP error — but crucially NOT a
|
||||
# successful fetch from example.com.
|
||||
with pytest.raises(requests.exceptions.HTTPError):
|
||||
connector.fetch_image(bypass_url)
|
||||
|
||||
with pytest.raises(aiohttp.ClientResponseError):
|
||||
await connector.fetch_image_async(bypass_url)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ssrf_bypass_backslash_disallowed_domain():
|
||||
"""The reverse direction: even when the *attacker-controlled* host
|
||||
appears in the urllib3-parsed hostname position the allowlist must
|
||||
still block it.
|
||||
"""
|
||||
# urllib3.parse_url sees host=example.com which is NOT in the
|
||||
# allowlist, so this must be rejected before any request is made.
|
||||
bypass_url = "https://example.com\\@safe.example.org/image.png"
|
||||
|
||||
connector = MediaConnector(
|
||||
allowed_media_domains=["safe.example.org"],
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="allowed domains"):
|
||||
connector.fetch_image(bypass_url)
|
||||
|
||||
with pytest.raises(ValueError, match="allowed domains"):
|
||||
await connector.fetch_image_async(bypass_url)
|
||||
|
||||
@@ -6,6 +6,7 @@ import pytest
|
||||
from vllm.assets.image import ImageAsset
|
||||
from vllm.assets.video import VideoAsset
|
||||
from vllm.config import CacheConfig, ModelConfig, VllmConfig
|
||||
from vllm.multimodal.parse import parse_mm_uuids
|
||||
from vllm.renderers.hf import HfRenderer
|
||||
from vllm.tokenizers.registry import tokenizer_args_from_config
|
||||
|
||||
@@ -45,10 +46,11 @@ def test_multi_modal_uuids_length_mismatch_raises():
|
||||
mm_uuids = {"image": ["hash_cherry"]}
|
||||
|
||||
mm_processor = renderer.get_mm_processor()
|
||||
mm_items = mm_processor.info.parse_mm_data(mm_data)
|
||||
mm_data_items = mm_processor.info.parse_mm_data(mm_data)
|
||||
mm_uuid_items = parse_mm_uuids(mm_uuids)
|
||||
|
||||
with pytest.raises(ValueError, match="must have same length as"):
|
||||
renderer._process_mm_uuids(mm_data, mm_items, mm_uuids, "req-1")
|
||||
renderer._process_mm_uuids(mm_data, mm_data_items, mm_uuid_items, "req-1")
|
||||
|
||||
|
||||
def test_multi_modal_uuids_missing_modality_raises():
|
||||
@@ -63,10 +65,11 @@ def test_multi_modal_uuids_missing_modality_raises():
|
||||
mm_uuids = {"image": ["hash_cherry"]}
|
||||
|
||||
mm_processor = renderer.get_mm_processor()
|
||||
mm_items = mm_processor.info.parse_mm_data(mm_data)
|
||||
mm_data_items = mm_processor.info.parse_mm_data(mm_data)
|
||||
mm_uuid_items = parse_mm_uuids(mm_uuids)
|
||||
|
||||
with pytest.raises(ValueError, match="is empty but .* is missing"):
|
||||
renderer._process_mm_uuids(mm_data, mm_items, mm_uuids, "req-2")
|
||||
renderer._process_mm_uuids(mm_data, mm_data_items, mm_uuid_items, "req-2")
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
@@ -78,7 +81,7 @@ def test_multi_modal_uuids_missing_modality_raises():
|
||||
],
|
||||
)
|
||||
def test_multi_modal_uuids_accepts_none_and_passes_through(
|
||||
monkeypatch, mm_cache_gb: float, enable_prefix_caching: bool
|
||||
mm_cache_gb: float, enable_prefix_caching: bool
|
||||
):
|
||||
renderer = _build_renderer(
|
||||
mm_cache_gb=mm_cache_gb,
|
||||
@@ -94,9 +97,11 @@ def test_multi_modal_uuids_accepts_none_and_passes_through(
|
||||
mm_uuids = {"image": [None, "hash_stop"], "video": None}
|
||||
|
||||
mm_processor = renderer.get_mm_processor()
|
||||
mm_items = mm_processor.info.parse_mm_data(mm_data)
|
||||
mm_data_items = mm_processor.info.parse_mm_data(mm_data)
|
||||
mm_uuid_items = parse_mm_uuids(mm_uuids)
|
||||
|
||||
processed_mm_uuids = renderer._process_mm_uuids(
|
||||
mm_data, mm_items, mm_uuids, "req-3"
|
||||
mm_data, mm_data_items, mm_uuid_items, "req-3"
|
||||
)
|
||||
|
||||
assert processed_mm_uuids == mm_uuids
|
||||
@@ -111,7 +116,7 @@ def test_multi_modal_uuids_accepts_none_and_passes_through(
|
||||
],
|
||||
)
|
||||
def test_multi_modal_uuids_accepts_empty(
|
||||
monkeypatch, mm_cache_gb: float, enable_prefix_caching: bool
|
||||
mm_cache_gb: float, enable_prefix_caching: bool
|
||||
):
|
||||
renderer = _build_renderer(
|
||||
mm_cache_gb=mm_cache_gb,
|
||||
@@ -124,15 +129,17 @@ def test_multi_modal_uuids_accepts_empty(
|
||||
mm_uuids = {"image": [], "video": None} # type: ignore[var-annotated]
|
||||
|
||||
mm_processor = renderer.get_mm_processor()
|
||||
mm_items = mm_processor.info.parse_mm_data(mm_data)
|
||||
mm_data_items = mm_processor.info.parse_mm_data(mm_data)
|
||||
mm_uuid_items = parse_mm_uuids(mm_uuids)
|
||||
|
||||
processed_mm_uuids = renderer._process_mm_uuids(
|
||||
mm_data, mm_items, mm_uuids, "req-4"
|
||||
mm_data, mm_data_items, mm_uuid_items, "req-4"
|
||||
)
|
||||
|
||||
assert processed_mm_uuids == mm_uuids
|
||||
|
||||
|
||||
def test_multi_modal_uuids_ignored_when_caching_disabled(monkeypatch):
|
||||
def test_multi_modal_uuids_ignored_when_caching_disabled():
|
||||
# When both processor cache is 0 and prefix caching disabled, the
|
||||
# processor builds overrides from request id instead of using user UUIDs.
|
||||
renderer = _build_renderer(mm_cache_gb=0.0, enable_prefix_caching=False)
|
||||
@@ -145,9 +152,11 @@ def test_multi_modal_uuids_ignored_when_caching_disabled(monkeypatch):
|
||||
mm_uuids = {"image": ["hash_cherry", "hash_stop"], "video": ["hash_video"]}
|
||||
|
||||
mm_processor = renderer.get_mm_processor()
|
||||
mm_items = mm_processor.info.parse_mm_data(mm_data)
|
||||
mm_data_items = mm_processor.info.parse_mm_data(mm_data)
|
||||
mm_uuid_items = parse_mm_uuids(mm_uuids)
|
||||
|
||||
processed_mm_uuids = renderer._process_mm_uuids(
|
||||
mm_data, mm_items, mm_uuids, request_id
|
||||
mm_data, mm_data_items, mm_uuid_items, request_id
|
||||
)
|
||||
|
||||
# Expect request-id-based overrides are passed through
|
||||
|
||||
@@ -8,7 +8,7 @@ MODEL = "meta-llama/Llama-3.2-1B-Instruct"
|
||||
|
||||
|
||||
def test_unsupported_configs():
|
||||
with pytest.raises(NotImplementedError):
|
||||
with pytest.raises(ValueError):
|
||||
AsyncEngineArgs(
|
||||
model=MODEL,
|
||||
speculative_config={
|
||||
|
||||
+42
-5
@@ -450,15 +450,30 @@ def rms_norm_per_block_quant(
|
||||
scale_ub: torch.Tensor | None = None,
|
||||
residual: torch.Tensor | None = None,
|
||||
is_scale_transposed: bool = False,
|
||||
tma_alignment: int = 0,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
assert len(group_size) == 2
|
||||
output = torch.empty_like(input, dtype=quant_dtype)
|
||||
if is_scale_transposed:
|
||||
scales = torch.empty(
|
||||
(input.shape[-1] // group_size[1], input.numel() // input.shape[-1]),
|
||||
device=input.device,
|
||||
dtype=torch.float32,
|
||||
).transpose(0, 1)
|
||||
if tma_alignment == 0:
|
||||
scales = torch.empty(
|
||||
(input.shape[-1] // group_size[1], input.numel() // input.shape[-1]),
|
||||
device=input.device,
|
||||
dtype=torch.float32,
|
||||
).transpose(0, 1)
|
||||
else:
|
||||
m = input.shape[-2]
|
||||
sf_k = input.shape[-1] // group_size[1]
|
||||
tma_aligned_m = (m + tma_alignment - 1) // tma_alignment * tma_alignment
|
||||
shape = input.shape[:-2] + (m, sf_k)
|
||||
stride = (
|
||||
(1, tma_aligned_m)
|
||||
if input.dim() == 2
|
||||
else (tma_aligned_m * sf_k, 1, tma_aligned_m)
|
||||
)
|
||||
scales = torch.empty_strided(
|
||||
shape, stride, device=input.device, dtype=torch.float32
|
||||
)
|
||||
else:
|
||||
scales = torch.empty(
|
||||
(input.numel() // input.shape[-1], input.shape[-1] // group_size[1]),
|
||||
@@ -466,6 +481,10 @@ def rms_norm_per_block_quant(
|
||||
dtype=torch.float32,
|
||||
)
|
||||
|
||||
assert tma_alignment in [0, 4], "Expected TMA alignment 0 or 4, but got " + str(
|
||||
tma_alignment
|
||||
)
|
||||
|
||||
torch.ops._C.rms_norm_per_block_quant(
|
||||
output,
|
||||
input,
|
||||
@@ -2770,6 +2789,24 @@ def sm100_cutlass_mla_get_workspace_size(
|
||||
)
|
||||
|
||||
|
||||
def dsv3_fused_a_gemm(
|
||||
output: torch.Tensor,
|
||||
mat_a: torch.Tensor,
|
||||
mat_b: torch.Tensor,
|
||||
) -> None:
|
||||
"""DeepSeek V3 fused A GEMM (SM 9.0+, bf16 only, 1-16 tokens).
|
||||
|
||||
Computes output = mat_a @ mat_b.T where:
|
||||
mat_a: [num_tokens, 7168] row-major bf16 (hidden states)
|
||||
mat_b: [7168, 2112] column-major bf16 (weight transposed)
|
||||
output: [num_tokens, 2112] row-major bf16
|
||||
|
||||
Optimized for the DeepSeek V2/V3 QKV A-projection at small batch sizes.
|
||||
Requires SM 9.0+ (Hopper).
|
||||
"""
|
||||
torch.ops._C.dsv3_fused_a_gemm(output, mat_a, mat_b)
|
||||
|
||||
|
||||
if hasattr(torch.ops._C, "weight_packed_linear"):
|
||||
|
||||
@register_fake("_C::weight_packed_linear")
|
||||
|
||||
@@ -292,6 +292,7 @@ class MatcherQuantFP8(MatcherCustomOp):
|
||||
has_col_major_scales: bool = False,
|
||||
is_e8m0: bool = False,
|
||||
match_rocm_aiter: bool = False,
|
||||
is_tma_aligned: bool = False,
|
||||
) -> None:
|
||||
if enabled is None:
|
||||
enabled = QuantFP8.enabled()
|
||||
@@ -301,6 +302,7 @@ class MatcherQuantFP8(MatcherCustomOp):
|
||||
self.has_col_major_scales = has_col_major_scales
|
||||
self.is_e8m0 = is_e8m0
|
||||
self.match_rocm_aiter = match_rocm_aiter
|
||||
self.is_tma_aligned = is_tma_aligned
|
||||
|
||||
if match_rocm_aiter:
|
||||
assert not quant_key.scale.group_shape.is_per_tensor(), (
|
||||
@@ -336,6 +338,7 @@ class MatcherQuantFP8(MatcherCustomOp):
|
||||
quant_key.scale.group_shape,
|
||||
column_major_scales=has_col_major_scales,
|
||||
use_ue8m0=is_e8m0,
|
||||
tma_aligned_scales=self.is_tma_aligned,
|
||||
compile_native=False,
|
||||
)
|
||||
|
||||
@@ -367,8 +370,11 @@ class MatcherQuantFP8(MatcherCustomOp):
|
||||
)
|
||||
|
||||
if self.quant_key.scale.group_shape.is_per_group():
|
||||
assert scale is None
|
||||
scale = self.make_scale(input, transposed=self.has_col_major_scales)
|
||||
# for tma_aligned, the scale must be passed to forward_custom
|
||||
# tma_aligned fusion then matches by custom op arguments
|
||||
if not self.is_tma_aligned:
|
||||
assert scale is None
|
||||
scale = self.make_scale(input, transposed=self.has_col_major_scales)
|
||||
|
||||
finfo = torch.finfo(self.quant_key.dtype)
|
||||
fp8_min = finfo.min
|
||||
@@ -384,6 +390,8 @@ class MatcherQuantFP8(MatcherCustomOp):
|
||||
fp8_min=fp8_min,
|
||||
fp8_max=fp8_max,
|
||||
scale_ue8m0=self.is_e8m0,
|
||||
dummy_is_scale_transposed=self.has_col_major_scales,
|
||||
dummy_is_tma_aligned=self.is_tma_aligned,
|
||||
)
|
||||
return result, scale
|
||||
|
||||
|
||||
@@ -121,6 +121,7 @@ class RMSNormQuantPattern:
|
||||
key: FusedRMSQuantKey,
|
||||
has_col_major_scales: bool = False,
|
||||
is_e8m0: bool = False,
|
||||
is_tma_aligned: bool = False,
|
||||
) -> None:
|
||||
self.epsilon = epsilon
|
||||
self.quant_dtype = key.quant.dtype
|
||||
@@ -136,7 +137,10 @@ class RMSNormQuantPattern:
|
||||
else MatcherFusedAddRMSNorm(epsilon)
|
||||
)
|
||||
self.quant_matcher = MatcherQuantFP8(
|
||||
key.quant, has_col_major_scales=has_col_major_scales, is_e8m0=is_e8m0
|
||||
key.quant,
|
||||
has_col_major_scales=has_col_major_scales,
|
||||
is_e8m0=is_e8m0,
|
||||
is_tma_aligned=is_tma_aligned,
|
||||
)
|
||||
|
||||
|
||||
@@ -262,8 +266,9 @@ class FusedAddRMSNormGroupQuantPattern(RMSNormQuantPattern):
|
||||
quant_dtype: torch.dtype,
|
||||
group_shape: GroupShape,
|
||||
symmetric: bool = True,
|
||||
has_col_major_scales: bool = False,
|
||||
is_e8m0: bool = False,
|
||||
has_col_major_scales: bool = True,
|
||||
is_tma_aligned: bool = True,
|
||||
) -> None:
|
||||
scale = ScaleDesc(torch.float32, False, group_shape)
|
||||
key = FusedRMSQuantKey(
|
||||
@@ -271,29 +276,63 @@ class FusedAddRMSNormGroupQuantPattern(RMSNormQuantPattern):
|
||||
quant=QuantKey(dtype=quant_dtype, scale=scale, symmetric=symmetric),
|
||||
)
|
||||
self.group_shape = group_shape
|
||||
self.has_col_major_scales = has_col_major_scales
|
||||
self.is_e8m0 = is_e8m0
|
||||
self.has_col_major_scales = has_col_major_scales
|
||||
self.is_tma_aligned = is_tma_aligned
|
||||
super().__init__(
|
||||
epsilon, key, has_col_major_scales=has_col_major_scales, is_e8m0=is_e8m0
|
||||
epsilon,
|
||||
key,
|
||||
has_col_major_scales=has_col_major_scales,
|
||||
is_e8m0=is_e8m0,
|
||||
is_tma_aligned=is_tma_aligned,
|
||||
)
|
||||
|
||||
def register(self, pm_pass: PatternMatcherPass) -> None:
|
||||
def pattern(
|
||||
input: torch.Tensor, weight: torch.Tensor, residual: torch.Tensor
|
||||
input: torch.Tensor,
|
||||
weight: torch.Tensor,
|
||||
residual: torch.Tensor,
|
||||
scale: torch.Tensor,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
result_rms, residual = self.rmsnorm_matcher(input, weight, residual)
|
||||
result, scale = self.quant_matcher(result_rms)
|
||||
result = torch.empty(
|
||||
result_rms.shape,
|
||||
device=result_rms.device,
|
||||
dtype=self.quant_matcher.quant_key.dtype,
|
||||
)
|
||||
assert scale is not None
|
||||
finfo = torch.finfo(self.quant_matcher.quant_key.dtype)
|
||||
fp8_min = finfo.min
|
||||
fp8_max = finfo.max
|
||||
|
||||
_, result, scale = auto_functionalized(
|
||||
self.quant_matcher.QUANT_OP,
|
||||
input=result_rms,
|
||||
output_q=result,
|
||||
output_s=scale,
|
||||
group_size=self.quant_matcher.quant_key.scale.group_shape[1],
|
||||
eps=1e-10,
|
||||
fp8_min=fp8_min,
|
||||
fp8_max=fp8_max,
|
||||
scale_ue8m0=self.quant_matcher.is_e8m0,
|
||||
dummy_is_scale_transposed=self.has_col_major_scales,
|
||||
dummy_is_tma_aligned=self.is_tma_aligned,
|
||||
)
|
||||
|
||||
return result, residual, scale
|
||||
|
||||
def replacement(
|
||||
input: torch.Tensor, weight: torch.Tensor, residual: torch.Tensor
|
||||
input: torch.Tensor,
|
||||
weight: torch.Tensor,
|
||||
residual: torch.Tensor,
|
||||
scale: torch.Tensor,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
# In case we're matching native rms-norm, conversions might be
|
||||
# optimized out. We convert here just to be safe.
|
||||
input = input.to(dtype=self.model_dtype)
|
||||
|
||||
result = torch.empty_like(input, dtype=self.quant_dtype)
|
||||
scale = self.quant_matcher.make_scale(input, self.has_col_major_scales)
|
||||
|
||||
at = auto_functionalized(
|
||||
self.FUSED_OP,
|
||||
result=result,
|
||||
@@ -310,10 +349,12 @@ class FusedAddRMSNormGroupQuantPattern(RMSNormQuantPattern):
|
||||
# result, residual, scale
|
||||
return at[1], at[3], at[2]
|
||||
|
||||
scale = self.quant_matcher.empty_f32(1, 1)
|
||||
|
||||
pm.register_replacement(
|
||||
pattern,
|
||||
replacement,
|
||||
self.rmsnorm_matcher.inputs(),
|
||||
self.rmsnorm_matcher.inputs() + [scale],
|
||||
pm.fwd_only,
|
||||
pm_pass,
|
||||
)
|
||||
@@ -326,8 +367,9 @@ class RMSNormGroupQuantPattern(RMSNormQuantPattern):
|
||||
quant_dtype: torch.dtype,
|
||||
group_shape: GroupShape,
|
||||
symmetric: bool = True,
|
||||
has_col_major_scales: bool = False,
|
||||
is_e8m0: bool = False,
|
||||
has_col_major_scales: bool = True,
|
||||
is_tma_aligned: bool = True,
|
||||
) -> None:
|
||||
scale = ScaleDesc(torch.float32, False, group_shape)
|
||||
key = FusedRMSQuantKey(
|
||||
@@ -335,29 +377,55 @@ class RMSNormGroupQuantPattern(RMSNormQuantPattern):
|
||||
quant=QuantKey(dtype=quant_dtype, scale=scale, symmetric=symmetric),
|
||||
)
|
||||
self.group_shape = group_shape
|
||||
self.has_col_major_scales = has_col_major_scales
|
||||
self.is_tma_aligned = is_tma_aligned
|
||||
super().__init__(
|
||||
epsilon, key, has_col_major_scales=has_col_major_scales, is_e8m0=is_e8m0
|
||||
epsilon,
|
||||
key,
|
||||
has_col_major_scales=self.has_col_major_scales,
|
||||
is_e8m0=is_e8m0,
|
||||
is_tma_aligned=is_tma_aligned,
|
||||
)
|
||||
|
||||
def register(self, pm_pass: PatternMatcherPass) -> None:
|
||||
def pattern(
|
||||
input: torch.Tensor, weight: torch.Tensor
|
||||
input: torch.Tensor, weight: torch.Tensor, scale: torch.Tensor
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
result_rms = self.rmsnorm_matcher(input, weight)
|
||||
result, scale = self.quant_matcher(result_rms)
|
||||
result = torch.empty(
|
||||
result_rms.shape,
|
||||
device=result_rms.device,
|
||||
dtype=self.quant_matcher.quant_key.dtype,
|
||||
)
|
||||
assert scale is not None
|
||||
finfo = torch.finfo(self.quant_matcher.quant_key.dtype)
|
||||
fp8_min = finfo.min
|
||||
fp8_max = finfo.max
|
||||
|
||||
_, result, scale = auto_functionalized(
|
||||
self.quant_matcher.QUANT_OP,
|
||||
input=result_rms,
|
||||
output_q=result,
|
||||
output_s=scale,
|
||||
group_size=self.quant_matcher.quant_key.scale.group_shape[1],
|
||||
eps=1e-10,
|
||||
fp8_min=fp8_min,
|
||||
fp8_max=fp8_max,
|
||||
scale_ue8m0=self.quant_matcher.is_e8m0,
|
||||
dummy_is_scale_transposed=self.has_col_major_scales,
|
||||
dummy_is_tma_aligned=self.is_tma_aligned,
|
||||
)
|
||||
|
||||
return result, scale
|
||||
|
||||
def replacement(
|
||||
input: torch.Tensor, weight: torch.Tensor
|
||||
input: torch.Tensor, weight: torch.Tensor, scale: torch.Tensor
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
# In case we're matching native rms-norm, conversions might be
|
||||
# optimized out. We convert here just to be safe.
|
||||
input = input.to(dtype=self.model_dtype)
|
||||
|
||||
result = torch.empty_like(input, dtype=self.quant_dtype)
|
||||
scale = self.quant_matcher.make_scale(
|
||||
input, transposed=self.quant_matcher.has_col_major_scales
|
||||
)
|
||||
at = auto_functionalized(
|
||||
self.FUSED_OP,
|
||||
result=result,
|
||||
@@ -368,16 +436,18 @@ class RMSNormGroupQuantPattern(RMSNormQuantPattern):
|
||||
scale_ub=None,
|
||||
residual=None,
|
||||
group_size=self.group_shape[1],
|
||||
is_scale_transposed=self.quant_matcher.has_col_major_scales,
|
||||
is_scale_transposed=self.has_col_major_scales,
|
||||
)
|
||||
|
||||
# result, scale
|
||||
return at[1], at[2]
|
||||
|
||||
scale = self.quant_matcher.empty_f32(1, 1)
|
||||
|
||||
pm.register_replacement(
|
||||
pattern,
|
||||
replacement,
|
||||
self.rmsnorm_matcher.inputs(),
|
||||
self.rmsnorm_matcher.inputs() + [scale],
|
||||
pm.fwd_only,
|
||||
pm_pass,
|
||||
)
|
||||
@@ -532,23 +602,26 @@ class RMSNormQuantFusionPass(VllmPatternMatcherPass):
|
||||
for group_shape in [GroupShape(1, 128), GroupShape(1, 64)]:
|
||||
for has_col_major_scales in [True, False]:
|
||||
for is_e8m0 in [True, False]:
|
||||
# Fuse fused_add_rms_norm + fp8 group quant
|
||||
FusedAddRMSNormGroupQuantPattern(
|
||||
epsilon,
|
||||
FP8_DTYPE,
|
||||
group_shape=group_shape,
|
||||
has_col_major_scales=has_col_major_scales,
|
||||
is_e8m0=is_e8m0,
|
||||
).register(self.patterns)
|
||||
for is_tma_aligned in [False, True]:
|
||||
# Fuse fused_add_rms_norm + fp8 group quant
|
||||
FusedAddRMSNormGroupQuantPattern(
|
||||
epsilon,
|
||||
FP8_DTYPE,
|
||||
group_shape=group_shape,
|
||||
is_e8m0=is_e8m0,
|
||||
has_col_major_scales=has_col_major_scales,
|
||||
is_tma_aligned=is_tma_aligned,
|
||||
).register(self.patterns)
|
||||
|
||||
# Fuse rms_norm + fp8 group quant
|
||||
RMSNormGroupQuantPattern(
|
||||
epsilon,
|
||||
FP8_DTYPE,
|
||||
group_shape=group_shape,
|
||||
has_col_major_scales=has_col_major_scales,
|
||||
is_e8m0=is_e8m0,
|
||||
).register(self.patterns)
|
||||
# Fuse rms_norm + fp8 group quant
|
||||
RMSNormGroupQuantPattern(
|
||||
epsilon,
|
||||
FP8_DTYPE,
|
||||
group_shape=group_shape,
|
||||
is_e8m0=is_e8m0,
|
||||
has_col_major_scales=has_col_major_scales,
|
||||
is_tma_aligned=is_tma_aligned,
|
||||
).register(self.patterns)
|
||||
|
||||
self.dump_patterns(config, self.patterns)
|
||||
|
||||
|
||||
@@ -305,6 +305,13 @@ class SpeculativeConfig:
|
||||
# can not be detected, it will be considered as the "draft_model" by
|
||||
# default.
|
||||
|
||||
# infer method from user args
|
||||
if self.method is None:
|
||||
if self.model in ("ngram", "[ngram]"):
|
||||
self.method = "ngram"
|
||||
else:
|
||||
self.method = "draft_model"
|
||||
|
||||
if self.method in get_args(MTPModelTypes) and self.method != "mtp":
|
||||
logger.warning(
|
||||
"method `%s` is deprecated and replaced with mtp.", self.method
|
||||
@@ -334,13 +341,6 @@ class SpeculativeConfig:
|
||||
"num_speculative_tokens was provided but without speculative model."
|
||||
)
|
||||
|
||||
# Automatically configure the method for ngram when "model" is used
|
||||
# instead of "method"
|
||||
if self.method is None and (
|
||||
self.model is not None and self.model in ("ngram", "[ngram]")
|
||||
):
|
||||
self.method = "ngram"
|
||||
|
||||
if self.method in ("ngram", "[ngram]"):
|
||||
# Unified to "ngram" internally
|
||||
self.method = "ngram"
|
||||
@@ -505,6 +505,13 @@ class SpeculativeConfig:
|
||||
)
|
||||
|
||||
if self.speculative_token_tree is None:
|
||||
if self.num_speculative_tokens is None:
|
||||
raise ValueError(
|
||||
"A speculative model was provided, but neither "
|
||||
"`speculative_token_tree` nor `num_speculative_tokens` "
|
||||
"was provided"
|
||||
)
|
||||
|
||||
# Generate chain of tokens.
|
||||
self.speculative_token_tree = str(
|
||||
[(i + 1) * (0,) for i in range(self.num_speculative_tokens)]
|
||||
|
||||
+10
-5
@@ -95,11 +95,16 @@ def enable_norm_fusion(cfg: "VllmConfig") -> bool:
|
||||
|
||||
|
||||
def enable_act_fusion(cfg: "VllmConfig") -> bool:
|
||||
"""Enable if either SiLU+Mul or quant FP8 custom op is active;
|
||||
otherwise Inductor handles fusion."""
|
||||
return cfg.compilation_config.is_custom_op_enabled(
|
||||
"silu_and_mul"
|
||||
) or cfg.compilation_config.is_custom_op_enabled("quant_fp8")
|
||||
"""
|
||||
Enable if either SiLU+Mul or quant FP8 custom op is active;
|
||||
otherwise Inductor handles fusion.
|
||||
Also enable for FP4 models as FP4 quant is always custom so Inductor cannot fuse it.
|
||||
"""
|
||||
return (
|
||||
cfg.compilation_config.is_custom_op_enabled("silu_and_mul")
|
||||
or cfg.compilation_config.is_custom_op_enabled("quant_fp8")
|
||||
or (cfg.model_config is not None and cfg.model_config.is_nvfp4_quantized())
|
||||
)
|
||||
|
||||
|
||||
def enable_allreduce_rms_fusion(cfg: "VllmConfig") -> bool:
|
||||
|
||||
@@ -40,6 +40,7 @@ class AnthropicContentBlock(BaseModel):
|
||||
source: dict[str, Any] | None = None
|
||||
# For tool use/result
|
||||
id: str | None = None
|
||||
tool_use_id: str | None = None
|
||||
name: str | None = None
|
||||
input: dict[str, Any] | None = None
|
||||
content: str | list[dict[str, Any]] | None = None
|
||||
|
||||
@@ -139,7 +139,7 @@ class AnthropicServingMessages(OpenAIServingChat):
|
||||
openai_messages.append(
|
||||
{
|
||||
"role": "tool",
|
||||
"tool_call_id": block.id or "",
|
||||
"tool_call_id": block.tool_use_id or "",
|
||||
"content": str(block.content)
|
||||
if block.content
|
||||
else "",
|
||||
|
||||
+17
-58
@@ -519,7 +519,6 @@ class LLM:
|
||||
),
|
||||
params=seq_params,
|
||||
lora_requests=seq_lora_requests,
|
||||
tokenization_kwargs=tokenization_kwargs,
|
||||
priorities=seq_priority,
|
||||
)
|
||||
|
||||
@@ -543,51 +542,31 @@ class LLM:
|
||||
outputs = self._run_engine(use_tqdm=use_tqdm)
|
||||
return self.engine_class.validate_outputs(outputs, RequestOutput)
|
||||
|
||||
def _resolve_lora_reqs(
|
||||
self,
|
||||
prompts: Sequence[ProcessorInputs],
|
||||
lora_request: Sequence[LoRARequest | None] | LoRARequest | None,
|
||||
):
|
||||
lora_config = self.llm_engine.vllm_config.lora_config
|
||||
seq_lora_requests = self._lora_request_to_seq(lora_request, len(prompts))
|
||||
|
||||
if (
|
||||
lora_config is None
|
||||
or not self.model_config.is_multimodal_model
|
||||
or (lora_config and lora_config.default_mm_loras is None)
|
||||
):
|
||||
return seq_lora_requests
|
||||
|
||||
return [
|
||||
self._resolve_single_prompt_mm_lora(
|
||||
prompt,
|
||||
lora_req,
|
||||
lora_config.default_mm_loras,
|
||||
)
|
||||
for prompt, lora_req in zip(prompts, seq_lora_requests)
|
||||
]
|
||||
|
||||
def _resolve_single_prompt_mm_lora(
|
||||
def _resolve_mm_lora(
|
||||
self,
|
||||
prompt: ProcessorInputs,
|
||||
lora_request: LoRARequest | None,
|
||||
default_mm_loras: dict[str, str] | None,
|
||||
):
|
||||
if not default_mm_loras or prompt["type"] != "multimodal":
|
||||
) -> LoRARequest | None:
|
||||
if prompt["type"] != "multimodal":
|
||||
return lora_request
|
||||
|
||||
lora_config = self.llm_engine.vllm_config.lora_config
|
||||
default_mm_loras = None if lora_config is None else lora_config.default_mm_loras
|
||||
if not default_mm_loras:
|
||||
return lora_request
|
||||
|
||||
prompt_modalities = prompt["mm_placeholders"].keys()
|
||||
intersection = set(prompt_modalities).intersection(default_mm_loras.keys())
|
||||
if not intersection:
|
||||
return lora_request
|
||||
|
||||
if len(intersection) > 1:
|
||||
# TODO: Would be nice to be able to have multiple loras per prompt
|
||||
logger.warning(
|
||||
"Multiple modality specific loras were registered and would be"
|
||||
" used by a single prompt consuming several modalities; "
|
||||
" currently we only support one lora per request; as such,"
|
||||
" lora(s) registered with modalities: %s"
|
||||
" will be skipped",
|
||||
"Multiple modality specific loras were registered and would be "
|
||||
"used by a single prompt consuming several modalities; "
|
||||
"currently we only support one lora per request; as such, "
|
||||
"lora(s) registered with modalities: %s will be skipped",
|
||||
intersection,
|
||||
)
|
||||
return lora_request
|
||||
@@ -1813,7 +1792,6 @@ class LLM:
|
||||
params=seq_params,
|
||||
use_tqdm=use_tqdm,
|
||||
lora_requests=seq_lora_requests,
|
||||
tokenization_kwargs=tokenization_kwargs,
|
||||
priorities=seq_priority,
|
||||
)
|
||||
|
||||
@@ -1872,7 +1850,6 @@ class LLM:
|
||||
params=seq_params,
|
||||
lora_requests=seq_lora_requests,
|
||||
use_tqdm=use_tqdm,
|
||||
tokenization_kwargs=tokenization_kwargs,
|
||||
)
|
||||
|
||||
def _render_and_run_requests(
|
||||
@@ -1881,7 +1858,6 @@ class LLM:
|
||||
params: Sequence[SamplingParams | PoolingParams],
|
||||
*,
|
||||
lora_requests: Sequence[LoRARequest | None] | None = None,
|
||||
tokenization_kwargs: dict[str, Any] | None = None,
|
||||
priorities: Sequence[int] | None = None,
|
||||
use_tqdm: bool | Callable[..., tqdm] = True,
|
||||
):
|
||||
@@ -1899,7 +1875,6 @@ class LLM:
|
||||
prompts=prompts,
|
||||
params=params,
|
||||
lora_requests=lora_requests,
|
||||
tokenization_kwargs=tokenization_kwargs,
|
||||
priorities=priorities,
|
||||
)
|
||||
|
||||
@@ -1911,7 +1886,6 @@ class LLM:
|
||||
params: Sequence[SamplingParams | PoolingParams],
|
||||
*,
|
||||
lora_requests: Sequence[LoRARequest | None] | None = None,
|
||||
tokenization_kwargs: dict[str, Any] | None = None,
|
||||
priorities: Sequence[int] | None = None,
|
||||
) -> list[str]:
|
||||
added_request_ids: list[str] = []
|
||||
@@ -1921,8 +1895,10 @@ class LLM:
|
||||
request_id = self._add_request(
|
||||
prompt,
|
||||
params[i],
|
||||
lora_request=None if lora_requests is None else lora_requests[i],
|
||||
tokenization_kwargs=tokenization_kwargs,
|
||||
lora_request=self._resolve_mm_lora(
|
||||
prompt,
|
||||
None if lora_requests is None else lora_requests[i],
|
||||
),
|
||||
priority=0 if priorities is None else priorities[i],
|
||||
)
|
||||
added_request_ids.append(request_id)
|
||||
@@ -1938,7 +1914,6 @@ class LLM:
|
||||
prompt: ProcessorInputs,
|
||||
params: SamplingParams | PoolingParams,
|
||||
lora_request: LoRARequest | None = None,
|
||||
tokenization_kwargs: dict[str, Any] | None = None,
|
||||
priority: int = 0,
|
||||
) -> str:
|
||||
if isinstance(params, SamplingParams):
|
||||
@@ -1947,27 +1922,11 @@ class LLM:
|
||||
|
||||
request_id = str(next(self.request_counter))
|
||||
|
||||
if params.truncate_prompt_tokens is not None:
|
||||
params_type = type(params).__name__
|
||||
warnings.warn(
|
||||
f"The `truncate_prompt_tokens` parameter in `{params_type}` "
|
||||
"is deprecated and will be removed in v0.16. "
|
||||
"Please pass it via `tokenization_kwargs` instead.",
|
||||
DeprecationWarning,
|
||||
stacklevel=2,
|
||||
)
|
||||
|
||||
tokenization_kwargs = merge_kwargs(
|
||||
tokenization_kwargs,
|
||||
dict(truncate_prompt_tokens=params.truncate_prompt_tokens),
|
||||
)
|
||||
|
||||
return self.llm_engine.add_request(
|
||||
request_id,
|
||||
prompt,
|
||||
params,
|
||||
lora_request=lora_request,
|
||||
tokenization_kwargs=tokenization_kwargs,
|
||||
priority=priority,
|
||||
)
|
||||
|
||||
|
||||
@@ -91,7 +91,7 @@ class InputPreprocessor:
|
||||
self,
|
||||
prompt: str | list[int],
|
||||
mm_data: MultiModalDataDict,
|
||||
mm_processor_kwargs: Mapping[str, object] | None,
|
||||
mm_processor_kwargs: Mapping[str, object] | None = None,
|
||||
tokenization_kwargs: dict[str, Any] | None = None,
|
||||
*,
|
||||
mm_uuids: MultiModalUUIDDict | None = None,
|
||||
@@ -103,9 +103,9 @@ class InputPreprocessor:
|
||||
return self.renderer._process_multimodal(
|
||||
prompt,
|
||||
mm_data,
|
||||
mm_uuids=mm_uuids,
|
||||
mm_processor_kwargs=mm_processor_kwargs,
|
||||
tokenization_kwargs=tokenization_kwargs,
|
||||
mm_uuids=mm_uuids,
|
||||
)
|
||||
|
||||
def _process_embeds(
|
||||
@@ -144,7 +144,7 @@ class InputPreprocessor:
|
||||
inputs = self._process_multimodal(
|
||||
prompt_token_ids,
|
||||
multi_modal_data,
|
||||
parsed_content.get("mm_processor_kwargs") or {},
|
||||
parsed_content.get("mm_processor_kwargs"),
|
||||
tokenization_kwargs=tokenization_kwargs,
|
||||
mm_uuids=parsed_content.get("multi_modal_uuids"),
|
||||
)
|
||||
|
||||
@@ -8,11 +8,7 @@ import torch
|
||||
|
||||
import vllm.envs as envs
|
||||
from vllm.config import ParallelConfig
|
||||
from vllm.distributed import (
|
||||
get_dp_group,
|
||||
get_pcp_group,
|
||||
get_tensor_model_parallel_rank,
|
||||
)
|
||||
from vllm.distributed import get_dp_group, get_pcp_group, get_tensor_model_parallel_rank
|
||||
from vllm.logger import init_logger
|
||||
from vllm.model_executor.layers.fused_moe.activation import MoEActivation
|
||||
from vllm.model_executor.layers.quantization.utils.ocp_mx_utils import (
|
||||
@@ -126,20 +122,31 @@ class RoutingMethodType(IntEnum):
|
||||
|
||||
|
||||
def get_routing_method_type(
|
||||
scoring_func: str, top_k: int, renormalize: bool
|
||||
scoring_func: str,
|
||||
top_k: int,
|
||||
renormalize: bool,
|
||||
num_expert_group: int | None,
|
||||
has_e_score_bias: bool,
|
||||
) -> RoutingMethodType:
|
||||
if has_e_score_bias:
|
||||
if (num_expert_group or 0) > 0 and scoring_func == "sigmoid":
|
||||
return RoutingMethodType.DeepSeekV3
|
||||
else:
|
||||
return RoutingMethodType.Unspecified
|
||||
|
||||
if scoring_func == "sigmoid":
|
||||
if top_k == 1:
|
||||
return RoutingMethodType.Llama4
|
||||
else:
|
||||
return RoutingMethodType.DeepSeekV3
|
||||
elif scoring_func == "softmax":
|
||||
return RoutingMethodType.Unspecified
|
||||
|
||||
if scoring_func == "softmax":
|
||||
if renormalize:
|
||||
return RoutingMethodType.Renormalize
|
||||
else:
|
||||
return RoutingMethodType.Default
|
||||
else:
|
||||
return RoutingMethodType.Unspecified
|
||||
|
||||
return RoutingMethodType.Unspecified
|
||||
|
||||
|
||||
@dataclass
|
||||
|
||||
@@ -402,7 +402,7 @@ class CPUFusedMOE:
|
||||
input,
|
||||
topk_weights,
|
||||
topk_ids,
|
||||
activation,
|
||||
activation.value,
|
||||
global_num_experts,
|
||||
skip_weighted,
|
||||
)
|
||||
|
||||
@@ -30,6 +30,7 @@ from vllm.model_executor.layers.fused_moe.utils import (
|
||||
disable_inplace,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.utils.marlin_utils import (
|
||||
get_marlin_input_dtype,
|
||||
marlin_make_workspace_new,
|
||||
marlin_moe_intermediate_size,
|
||||
marlin_quant_input,
|
||||
@@ -550,6 +551,8 @@ class MarlinExpertsBase(mk.FusedMoEPermuteExpertsUnpermute):
|
||||
self.w13_g_idx_sort_indices = w13_g_idx_sort_indices
|
||||
self.w2_g_idx_sort_indices = w2_g_idx_sort_indices
|
||||
self.is_k_full = is_k_full
|
||||
self.input_dtype = get_marlin_input_dtype()
|
||||
|
||||
super().__init__(
|
||||
moe_config=moe_config,
|
||||
quant_config=quant_config,
|
||||
@@ -736,6 +739,7 @@ class MarlinExperts(MarlinExpertsBase):
|
||||
sort_indices1=self.w13_g_idx_sort_indices,
|
||||
sort_indices2=self.w2_g_idx_sort_indices,
|
||||
is_k_full=self.is_k_full,
|
||||
input_dtype=self.input_dtype,
|
||||
)
|
||||
|
||||
def moe_sum(self, input: torch.Tensor, output: torch.Tensor) -> None:
|
||||
|
||||
@@ -381,7 +381,6 @@ class AiterExperts(mk.FusedMoEPermuteExpertsUnpermute):
|
||||
# TODO(rob): rocm_aiter_fused_experts uses self.quant_config's
|
||||
# a_scales for static quantization. Update this to fit better
|
||||
# with the interface once all quant integrations are complete.
|
||||
assert a2_scale == self.quant_config.a2_scale
|
||||
|
||||
if expert_tokens_meta is not None:
|
||||
num_local_tokens = expert_tokens_meta.expert_num_tokens
|
||||
|
||||
@@ -165,6 +165,8 @@ class FusedTopKBiasRouter(BaseRouter):
|
||||
scoring_func=self.scoring_func,
|
||||
top_k=self.top_k,
|
||||
renormalize=self.renormalize,
|
||||
num_expert_group=None,
|
||||
has_e_score_bias=True,
|
||||
)
|
||||
|
||||
def _compute_routing(
|
||||
|
||||
@@ -142,6 +142,8 @@ class FusedTopKRouter(BaseRouter):
|
||||
scoring_func=self.scoring_func,
|
||||
top_k=self.top_k,
|
||||
renormalize=self.renormalize,
|
||||
num_expert_group=None,
|
||||
has_e_score_bias=False,
|
||||
)
|
||||
|
||||
def _compute_routing(
|
||||
|
||||
@@ -13,7 +13,10 @@ from vllm.model_executor.custom_op import CustomOp
|
||||
from vllm.model_executor.layers.batch_invariant import (
|
||||
vllm_is_batch_invariant,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.config import RoutingMethodType
|
||||
from vllm.model_executor.layers.fused_moe.config import (
|
||||
RoutingMethodType,
|
||||
get_routing_method_type,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
|
||||
rocm_aiter_grouped_topk,
|
||||
)
|
||||
@@ -277,16 +280,15 @@ class GroupedTopKRouter(BaseRouter):
|
||||
self.e_score_correction_bias = e_score_correction_bias
|
||||
self.num_fused_shared_experts = num_fused_shared_experts
|
||||
|
||||
if scoring_func == "sigmoid":
|
||||
self._routing_method_type = RoutingMethodType.DeepSeekV3
|
||||
else:
|
||||
# NOTE: this prohibits the FLASHINFER_TRTLLM kernels from
|
||||
# being selected, since they only support DeepSeek-style.
|
||||
self._routing_method_type = RoutingMethodType.Unspecified
|
||||
|
||||
@property
|
||||
def routing_method_type(self) -> RoutingMethodType:
|
||||
return self._routing_method_type
|
||||
return get_routing_method_type(
|
||||
scoring_func=self.scoring_func,
|
||||
top_k=self.top_k,
|
||||
renormalize=self.renormalize,
|
||||
num_expert_group=self.num_expert_group,
|
||||
has_e_score_bias=self.e_score_correction_bias is not None,
|
||||
)
|
||||
|
||||
def _compute_routing(
|
||||
self,
|
||||
|
||||
@@ -685,8 +685,13 @@ class MergedColumnParallelLinear(ColumnParallelLinear):
|
||||
self,
|
||||
param: Parameter,
|
||||
loaded_weight: torch.Tensor,
|
||||
loaded_shard_id: int | None = None,
|
||||
loaded_shard_id: tuple[int, ...] | int | None = None,
|
||||
):
|
||||
if isinstance(loaded_shard_id, tuple):
|
||||
raise NotImplementedError(
|
||||
"Shard id with multiple indices is not supported in weight_loader, "
|
||||
"please use weight_loader_v2 instead."
|
||||
)
|
||||
# Special case for GGUF
|
||||
# initialize GGUF param after we know the quantize type
|
||||
is_gguf_weight = getattr(param, "is_gguf_weight", False)
|
||||
@@ -770,6 +775,8 @@ class MergedColumnParallelLinear(ColumnParallelLinear):
|
||||
if output_dim is not None:
|
||||
shard_offset = sum(self.output_sizes[:loaded_shard_id])
|
||||
shard_size = self.output_sizes[loaded_shard_id]
|
||||
shard_offset //= self.tp_size
|
||||
shard_size //= self.tp_size
|
||||
|
||||
if isinstance(param, BlockQuantScaleParameter):
|
||||
weight_block_size = getattr(self, "weight_block_size", None)
|
||||
@@ -777,9 +784,6 @@ class MergedColumnParallelLinear(ColumnParallelLinear):
|
||||
weight_block_size, shard_size, shard_offset
|
||||
)
|
||||
|
||||
shard_offset //= self.tp_size
|
||||
shard_size //= self.tp_size
|
||||
|
||||
# Special case for quantization.
|
||||
# If quantized, we need to adjust the offset and size to account
|
||||
# for the packing.
|
||||
@@ -825,7 +829,10 @@ class MergedColumnParallelLinear(ColumnParallelLinear):
|
||||
param_data.copy_(loaded_weight)
|
||||
|
||||
def _load_fused_module_from_checkpoint(
|
||||
self, param: BasevLLMParameter, loaded_weight: torch.Tensor
|
||||
self,
|
||||
param: BasevLLMParameter,
|
||||
loaded_weight: torch.Tensor,
|
||||
output_sizes: list[int] | None = None,
|
||||
):
|
||||
"""
|
||||
Handle special case for models where MLP layers are already
|
||||
@@ -839,7 +846,8 @@ class MergedColumnParallelLinear(ColumnParallelLinear):
|
||||
|
||||
current_shard_offset = 0
|
||||
shard_offsets: list[tuple[int, int, int]] = []
|
||||
for i, output_size in enumerate(self.output_sizes):
|
||||
output_sizes = output_sizes or self.output_sizes
|
||||
for i, output_size in enumerate(output_sizes):
|
||||
shard_offsets.append((i, current_shard_offset, output_size))
|
||||
current_shard_offset += output_size
|
||||
|
||||
@@ -864,23 +872,38 @@ class MergedColumnParallelLinear(ColumnParallelLinear):
|
||||
self,
|
||||
param: BasevLLMParameter,
|
||||
loaded_weight: torch.Tensor,
|
||||
loaded_shard_id: int | None = None,
|
||||
loaded_shard_id: tuple[int, ...] | int | None = None,
|
||||
):
|
||||
if loaded_shard_id is None:
|
||||
if loaded_shard_id is None or isinstance(loaded_shard_id, tuple):
|
||||
if isinstance(param, PerTensorScaleParameter):
|
||||
param.load_merged_column_weight(loaded_weight=loaded_weight, shard_id=0)
|
||||
return
|
||||
elif type(param) in (RowvLLMParameter, BasevLLMParameter):
|
||||
param.load_merged_column_weight(loaded_weight=loaded_weight)
|
||||
return
|
||||
output_sizes = (
|
||||
[self.output_sizes[idx] for idx in loaded_shard_id]
|
||||
if loaded_shard_id
|
||||
else None
|
||||
)
|
||||
if isinstance(param, BlockQuantScaleParameter):
|
||||
weight_block_size = getattr(self, "weight_block_size", None)
|
||||
output_sizes = [
|
||||
adjust_block_scale_shard(weight_block_size, size, 0)[0]
|
||||
for size in (output_sizes or self.output_sizes)
|
||||
]
|
||||
# TODO: @dsikka - move to parameter.py
|
||||
self._load_fused_module_from_checkpoint(param, loaded_weight)
|
||||
self._load_fused_module_from_checkpoint(
|
||||
param, loaded_weight, output_sizes=output_sizes
|
||||
)
|
||||
return
|
||||
|
||||
assert loaded_shard_id < len(self.output_sizes)
|
||||
|
||||
shard_offset = sum(self.output_sizes[:loaded_shard_id])
|
||||
shard_size = self.output_sizes[loaded_shard_id]
|
||||
shard_offset //= self.tp_size
|
||||
shard_size //= self.tp_size
|
||||
|
||||
if isinstance(param, BlockQuantScaleParameter):
|
||||
weight_block_size = getattr(self, "weight_block_size", None)
|
||||
@@ -888,9 +911,6 @@ class MergedColumnParallelLinear(ColumnParallelLinear):
|
||||
weight_block_size, shard_size, shard_offset
|
||||
)
|
||||
|
||||
shard_offset //= self.tp_size
|
||||
shard_size //= self.tp_size
|
||||
|
||||
param.load_merged_column_weight(
|
||||
loaded_weight=loaded_weight,
|
||||
shard_id=loaded_shard_id,
|
||||
|
||||
@@ -129,6 +129,7 @@ class MultiHeadLatentAttentionWrapper(PluggableLayer):
|
||||
assert self.q_b_proj is not None, (
|
||||
"q_b_proj is required when q_lora_rank is not None"
|
||||
)
|
||||
|
||||
qkv_lora = self.fused_qkv_a_proj(hidden_states)[0]
|
||||
q_c, kv_lora = qkv_lora.split(
|
||||
[self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim],
|
||||
|
||||
@@ -18,6 +18,11 @@ else:
|
||||
class QuantizeMethodBase(ABC):
|
||||
"""Base class for different quantized methods."""
|
||||
|
||||
# Whether this method creates weights on meta device for online quantization.
|
||||
# When True, weights are created on meta device and quantized layer-wise
|
||||
# in process_weights_after_loading, reducing peak memory during loading.
|
||||
uses_meta_device: bool = False
|
||||
|
||||
@abstractmethod
|
||||
def create_weights(
|
||||
self, layer: torch.nn.Module, *weight_args, **extra_weight_attrs
|
||||
|
||||
@@ -527,6 +527,8 @@ class Fp8OnlineLinearMethod(Fp8LinearMethod):
|
||||
"""Online version of Fp8LinearMethod, loads the fp16/bf16 checkpoint
|
||||
and quantized the weights during loading."""
|
||||
|
||||
uses_meta_device: bool = True
|
||||
|
||||
def create_weights(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
@@ -1039,6 +1041,8 @@ class Fp8OnlineMoEMethod(Fp8MoEMethod):
|
||||
quant_config: The quantization config.
|
||||
"""
|
||||
|
||||
uses_meta_device: bool = True
|
||||
|
||||
def __init__(self, quant_config: Fp8Config, layer: torch.nn.Module):
|
||||
super().__init__(quant_config, layer)
|
||||
assert not quant_config.is_checkpoint_fp8_serialized
|
||||
|
||||
@@ -16,6 +16,9 @@ from vllm.model_executor.layers.fused_moe import (
|
||||
MoEActivation,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe import modular_kernel as mk
|
||||
from vllm.model_executor.layers.fused_moe.all2all_utils import (
|
||||
maybe_make_prepare_finalize,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.config import (
|
||||
FusedMoEQuantConfig,
|
||||
mxfp4_mxfp8_moe_quant_config,
|
||||
@@ -25,7 +28,6 @@ from vllm.model_executor.layers.fused_moe.config import (
|
||||
from vllm.model_executor.layers.fused_moe.fused_marlin_moe import (
|
||||
BatchedMarlinExperts,
|
||||
MarlinExperts,
|
||||
fused_marlin_moe,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe import (
|
||||
OAITritonExperts,
|
||||
@@ -52,7 +54,6 @@ from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (
|
||||
from vllm.model_executor.layers.quantization.utils.quant_utils import is_layer_skipped
|
||||
from vllm.model_executor.utils import set_weight_attrs
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.scalar_type import scalar_types
|
||||
from vllm.utils.flashinfer import has_flashinfer
|
||||
from vllm.utils.import_utils import has_triton_kernels
|
||||
from vllm.utils.math_utils import round_up
|
||||
@@ -219,7 +220,6 @@ class Mxfp4Config(QuantizationConfig):
|
||||
return XpuMxfp4MoEMethod(layer.moe_config)
|
||||
else:
|
||||
quant_method = Mxfp4MoEMethod(layer.moe_config)
|
||||
quant_method.marlin_input_dtype = get_marlin_input_dtype(prefix)
|
||||
return quant_method
|
||||
elif isinstance(layer, Attention):
|
||||
# TODO: Add support for MXFP4 Attention.
|
||||
@@ -243,7 +243,6 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
|
||||
self.weight_dtype = "mxfp4"
|
||||
self.mxfp4_backend = get_mxfp4_backend(moe.is_lora_enabled)
|
||||
|
||||
self.marlin_input_dtype = None
|
||||
self.max_capture_size = (
|
||||
get_current_vllm_config().compilation_config.max_cudagraph_capture_size
|
||||
)
|
||||
@@ -254,6 +253,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
|
||||
"Please check your environment and try again."
|
||||
)
|
||||
self._cache_permute_indices: dict[torch.Size, torch.Tensor] = {}
|
||||
self.moe_mk: mk.FusedMoEModularKernel | None = None
|
||||
|
||||
def create_weights(
|
||||
self,
|
||||
@@ -408,7 +408,30 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
|
||||
|
||||
def process_weights_after_loading(self, layer):
|
||||
if self.mxfp4_backend == Mxfp4Backend.MARLIN:
|
||||
prepare_moe_fp4_layer_for_marlin(layer, input_dtype=self.marlin_input_dtype)
|
||||
prepare_moe_fp4_layer_for_marlin(
|
||||
layer, input_dtype=get_marlin_input_dtype()
|
||||
)
|
||||
|
||||
self.moe_quant_config = self.get_fused_moe_quant_config(layer)
|
||||
assert self.moe_quant_config is not None
|
||||
|
||||
prepare_finalize = maybe_make_prepare_finalize(
|
||||
moe=self.moe,
|
||||
quant_config=self.moe_quant_config,
|
||||
routing_tables=layer._maybe_init_expert_routing_tables(),
|
||||
allow_new_interface=True,
|
||||
)
|
||||
assert prepare_finalize is not None
|
||||
|
||||
self.moe_mk = mk.FusedMoEModularKernel(
|
||||
prepare_finalize,
|
||||
MarlinExperts(
|
||||
self.moe,
|
||||
self.moe_quant_config,
|
||||
),
|
||||
inplace=not self.moe.disable_inplace,
|
||||
shared_experts=None,
|
||||
)
|
||||
elif (
|
||||
self.mxfp4_backend == Mxfp4Backend.SM100_FI_MXFP4_MXFP8_TRTLLM
|
||||
or self.mxfp4_backend == Mxfp4Backend.SM100_FI_MXFP4_BF16
|
||||
@@ -910,27 +933,19 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
|
||||
raise NotImplementedError("EPLB is not supported for mxfp4")
|
||||
|
||||
if self.mxfp4_backend == Mxfp4Backend.MARLIN:
|
||||
return fused_marlin_moe(
|
||||
x,
|
||||
layer.w13_weight,
|
||||
layer.w2_weight,
|
||||
layer.w13_bias,
|
||||
layer.w2_bias,
|
||||
layer.w13_weight_scale,
|
||||
layer.w2_weight_scale,
|
||||
topk_weights,
|
||||
topk_ids,
|
||||
global_scale1=None,
|
||||
global_scale2=None,
|
||||
quant_type_id=scalar_types.float4_e2m1f.id,
|
||||
apply_router_weight_on_input=layer.apply_router_weight_on_input,
|
||||
global_num_experts=layer.global_num_experts,
|
||||
activation=layer.activation,
|
||||
expert_map=layer.expert_map,
|
||||
input_dtype=self.marlin_input_dtype,
|
||||
inplace=not self.moe.disable_inplace,
|
||||
)
|
||||
assert self.moe_mk is not None
|
||||
|
||||
return self.moe_mk(
|
||||
hidden_states=x,
|
||||
w1=layer.w13_weight,
|
||||
w2=layer.w2_weight,
|
||||
topk_weights=topk_weights,
|
||||
topk_ids=topk_ids,
|
||||
activation=layer.activation,
|
||||
global_num_experts=layer.global_num_experts,
|
||||
expert_map=layer.expert_map,
|
||||
apply_router_weight_on_input=layer.apply_router_weight_on_input,
|
||||
)
|
||||
assert _can_support_mxfp4(
|
||||
layer.use_grouped_topk,
|
||||
layer.topk_group,
|
||||
|
||||
@@ -122,6 +122,8 @@ def is_supported_config_trtllm(
|
||||
return False, _make_reason("routing method")
|
||||
elif activation_format != mk.FusedMoEActivationFormat.Standard:
|
||||
return False, _make_reason("activation format")
|
||||
elif moe_config.hidden_dim % 512 != 0:
|
||||
return False, _make_reason("hidden_dim must be divisible by 512")
|
||||
|
||||
return True, None
|
||||
|
||||
|
||||
@@ -924,7 +924,16 @@ def per_token_group_quant_fp8(
|
||||
# TODO(bnell): this causes some fp8 moe test to fail.
|
||||
if current_platform.is_cuda() and x.is_contiguous():
|
||||
torch.ops._C.per_token_group_fp8_quant(
|
||||
x, x_q, x_s, group_size, eps, fp8_min, fp8_max, use_ue8m0
|
||||
x,
|
||||
x_q,
|
||||
x_s,
|
||||
group_size,
|
||||
eps,
|
||||
fp8_min,
|
||||
fp8_max,
|
||||
use_ue8m0,
|
||||
column_major_scales,
|
||||
tma_aligned_scales,
|
||||
)
|
||||
return x_q, x_s
|
||||
|
||||
|
||||
@@ -1092,16 +1092,20 @@ def initialize_dummy_weights(
|
||||
is fixed, the random values generated by this function only depends on
|
||||
the parameter's number of elements and its data type.
|
||||
"""
|
||||
# TODO(future PR): make the check below more generic as more online
|
||||
# quant backends are added
|
||||
is_fp8_py_quant = model_config.quantization == "fp8"
|
||||
|
||||
# Check if any module uses online quantization with meta device weights.
|
||||
# If so, we'll skip initializing params on meta device since they'll be
|
||||
# handled in `process_weights_after_loading`.
|
||||
def uses_meta_device(module: torch.nn.Module) -> bool:
|
||||
quant_method = getattr(module, "quant_method", None)
|
||||
return getattr(quant_method, "uses_meta_device", False)
|
||||
|
||||
has_online_quant = any(uses_meta_device(m) for m in model.modules())
|
||||
|
||||
for param in model.state_dict().values():
|
||||
if is_fp8_py_quant and param.device == torch.device("meta"):
|
||||
# for fp8.py's online quantization, dummy weight init will happen
|
||||
# in `process_weights_after_loading`.
|
||||
# TODO(future PR): consider refactoring dummy model init to compose
|
||||
# better with online quantization
|
||||
if has_online_quant and param.device == torch.device("meta"):
|
||||
# For online quantization, weights are created on meta device and
|
||||
# dummy weight init will happen in `process_weights_after_loading`.
|
||||
continue
|
||||
|
||||
initialize_single_dummy_weight(param, low, high, seed)
|
||||
|
||||
@@ -36,9 +36,13 @@ from vllm.multimodal.inputs import (
|
||||
MultiModalFieldConfig,
|
||||
MultiModalInputs,
|
||||
MultiModalKwargsItems,
|
||||
MultiModalUUIDDict,
|
||||
)
|
||||
from vllm.multimodal.parse import ImageProcessorItems, ImageSize, MultiModalDataItems
|
||||
from vllm.multimodal.parse import (
|
||||
ImageProcessorItems,
|
||||
ImageSize,
|
||||
MultiModalDataItems,
|
||||
MultiModalUUIDItems,
|
||||
)
|
||||
from vllm.multimodal.processing import (
|
||||
BaseDummyInputsBuilder,
|
||||
BaseMultiModalProcessor,
|
||||
@@ -203,10 +207,9 @@ class CLIPMultiModalProcessor(BaseMultiModalProcessor[CLIPProcessingInfo]):
|
||||
self,
|
||||
prompt: str | list[int],
|
||||
mm_items: MultiModalDataItems,
|
||||
hf_processor_mm_kwargs: Mapping[str, object],
|
||||
mm_uuid_items: MultiModalUUIDItems | None = None,
|
||||
hf_processor_mm_kwargs: Mapping[str, object] | None = None,
|
||||
tokenization_kwargs: Mapping[str, object] | None = None,
|
||||
*,
|
||||
mm_uuids: MultiModalUUIDDict | None = None,
|
||||
) -> MultiModalInputs:
|
||||
if mm_items:
|
||||
if isinstance(prompt, str):
|
||||
@@ -235,9 +238,9 @@ class CLIPMultiModalProcessor(BaseMultiModalProcessor[CLIPProcessingInfo]):
|
||||
return super().apply(
|
||||
prompt=prompt,
|
||||
mm_items=mm_items,
|
||||
mm_uuid_items=mm_uuid_items,
|
||||
hf_processor_mm_kwargs=hf_processor_mm_kwargs,
|
||||
tokenization_kwargs=tokenization_kwargs,
|
||||
mm_uuids=mm_uuids,
|
||||
)
|
||||
|
||||
def _hf_processor_applies_updates(
|
||||
|
||||
@@ -32,6 +32,7 @@ import torch
|
||||
from torch import nn
|
||||
from transformers import DeepseekV2Config, DeepseekV3Config
|
||||
|
||||
import vllm._custom_ops as ops
|
||||
from vllm._aiter_ops import rocm_aiter_ops
|
||||
from vllm.compilation.decorators import support_torch_compile
|
||||
from vllm.config import CacheConfig, ParallelConfig, VllmConfig, get_current_vllm_config
|
||||
@@ -711,6 +712,64 @@ class Indexer(nn.Module):
|
||||
return self.indexer_op(hidden_states, q_fp8, k, weights)
|
||||
|
||||
|
||||
class DeepSeekV2FusedQkvAProj(MergedColumnParallelLinear):
|
||||
def __init__(
|
||||
self,
|
||||
input_size: int,
|
||||
output_size: int,
|
||||
quant_config: QuantizationConfig | None = None,
|
||||
prefix: str = "",
|
||||
):
|
||||
super().__init__(
|
||||
input_size,
|
||||
output_size,
|
||||
bias=False,
|
||||
quant_config=quant_config,
|
||||
disable_tp=True,
|
||||
prefix=f"{prefix}.kv_a_proj_with_mqa",
|
||||
)
|
||||
|
||||
# Check if the DeepSeek V3 fused A GEMM kernel can be used.
|
||||
# This kernel supports PDL and is optimized for low batch size.
|
||||
self._use_min_latency_gemm = (
|
||||
hasattr(self, "weight")
|
||||
and self.weight.dtype == torch.bfloat16
|
||||
and self.weight.shape[0] == 2112
|
||||
and self.weight.shape[1] == 7168
|
||||
and current_platform.is_cuda()
|
||||
and (
|
||||
current_platform.is_device_capability(90)
|
||||
or current_platform.is_device_capability_family(100)
|
||||
)
|
||||
)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_,
|
||||
) -> torch.Tensor | tuple[torch.Tensor, torch.nn.Parameter | None]:
|
||||
num_tokens = input_.shape[0]
|
||||
if self._use_min_latency_gemm and (0 < num_tokens <= 16):
|
||||
output = torch.empty(
|
||||
num_tokens,
|
||||
2112,
|
||||
dtype=torch.bfloat16,
|
||||
device=input_.device,
|
||||
)
|
||||
ops.dsv3_fused_a_gemm(
|
||||
output,
|
||||
input_,
|
||||
self.weight.T,
|
||||
)
|
||||
if not self.return_bias:
|
||||
return output
|
||||
output_bias = self.bias if self.skip_bias_add else None
|
||||
return output, output_bias
|
||||
else:
|
||||
# Fallback to the standard forward method when
|
||||
# the fused A GEMM kernel cannot be used.
|
||||
return super().forward(input_)
|
||||
|
||||
|
||||
class DeepseekV2MLAAttention(nn.Module):
|
||||
"""
|
||||
Main reference: DeepseekV2 paper, and FlashInfer Implementation
|
||||
@@ -756,13 +815,11 @@ class DeepseekV2MLAAttention(nn.Module):
|
||||
self.max_position_embeddings = max_position_embeddings
|
||||
|
||||
if self.q_lora_rank is not None:
|
||||
self.fused_qkv_a_proj = MergedColumnParallelLinear(
|
||||
self.fused_qkv_a_proj = DeepSeekV2FusedQkvAProj(
|
||||
self.hidden_size,
|
||||
[self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim],
|
||||
bias=False,
|
||||
quant_config=quant_config,
|
||||
prefix=f"{prefix}.fused_qkv_a_proj",
|
||||
disable_tp=True,
|
||||
)
|
||||
else:
|
||||
self.kv_a_proj_with_mqa = ReplicatedLinear(
|
||||
|
||||
@@ -24,13 +24,13 @@ from vllm.multimodal.inputs import (
|
||||
MultiModalDataDict,
|
||||
MultiModalFieldConfig,
|
||||
MultiModalKwargsItems,
|
||||
MultiModalUUIDDict,
|
||||
)
|
||||
from vllm.multimodal.parse import (
|
||||
ImageEmbeddingItems,
|
||||
ImageProcessorItems,
|
||||
ImageSize,
|
||||
MultiModalDataItems,
|
||||
MultiModalUUIDItems,
|
||||
)
|
||||
from vllm.multimodal.processing import BaseDummyInputsBuilder
|
||||
from vllm.multimodal.processing.processor import (
|
||||
@@ -313,9 +313,9 @@ class DeepseekVL2MultiModalProcessor(
|
||||
self,
|
||||
prompt: str | list[int],
|
||||
mm_data_items: MultiModalDataItems,
|
||||
mm_uuid_items: MultiModalUUIDItems | None,
|
||||
hf_processor_mm_kwargs: Mapping[str, object],
|
||||
tokenization_kwargs: Mapping[str, object],
|
||||
mm_uuids: MultiModalUUIDDict | None = None,
|
||||
) -> tuple[list[int], MultiModalProcessingInfo, bool]:
|
||||
# The processor logic is different for len(images) <= 2 vs > 2
|
||||
# Since the processing cache assumes that the processor output is
|
||||
@@ -325,17 +325,17 @@ class DeepseekVL2MultiModalProcessor(
|
||||
return self._apply_hf_processor(
|
||||
prompt=prompt,
|
||||
mm_data_items=mm_data_items,
|
||||
mm_uuid_items=mm_uuid_items,
|
||||
hf_processor_mm_kwargs=hf_processor_mm_kwargs,
|
||||
tokenization_kwargs=tokenization_kwargs,
|
||||
mm_uuids=mm_uuids,
|
||||
)
|
||||
|
||||
return super()._cached_apply_hf_processor(
|
||||
prompt=prompt,
|
||||
mm_data_items=mm_data_items,
|
||||
mm_uuid_items=mm_uuid_items,
|
||||
hf_processor_mm_kwargs=hf_processor_mm_kwargs,
|
||||
tokenization_kwargs=tokenization_kwargs,
|
||||
mm_uuids=mm_uuids,
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -16,11 +16,12 @@ from transformers import PretrainedConfig
|
||||
|
||||
from vllm.model_executor.layers.quantization import QuantizationConfig
|
||||
from vllm.multimodal import MULTIMODAL_REGISTRY
|
||||
from vllm.multimodal.inputs import MultiModalKwargsItems, MultiModalUUIDDict
|
||||
from vllm.multimodal.inputs import MultiModalKwargsItems
|
||||
from vllm.multimodal.parse import (
|
||||
ImageEmbeddingItems,
|
||||
ImageProcessorItems,
|
||||
MultiModalDataItems,
|
||||
MultiModalUUIDItems,
|
||||
)
|
||||
from vllm.multimodal.processing.processor import (
|
||||
MultiModalProcessingInfo,
|
||||
@@ -491,9 +492,9 @@ class H2OVLMultiModalProcessor(BaseInternVLMultiModalProcessor[H2OVLProcessingIn
|
||||
self,
|
||||
prompt: str | list[int],
|
||||
mm_data_items: MultiModalDataItems,
|
||||
mm_uuid_items: MultiModalUUIDItems | None,
|
||||
hf_processor_mm_kwargs: Mapping[str, object],
|
||||
tokenization_kwargs: Mapping[str, object],
|
||||
mm_uuids: MultiModalUUIDDict | None = None,
|
||||
) -> tuple[list[int], MultiModalProcessingInfo, bool]:
|
||||
# The processor logic is different for len(images) <= 1 vs > 1
|
||||
# Since the processing cache assumes that the processor output is
|
||||
@@ -503,17 +504,17 @@ class H2OVLMultiModalProcessor(BaseInternVLMultiModalProcessor[H2OVLProcessingIn
|
||||
return self._apply_hf_processor(
|
||||
prompt=prompt,
|
||||
mm_data_items=mm_data_items,
|
||||
mm_uuid_items=mm_uuid_items,
|
||||
hf_processor_mm_kwargs=hf_processor_mm_kwargs,
|
||||
tokenization_kwargs=tokenization_kwargs,
|
||||
mm_uuids=mm_uuids,
|
||||
)
|
||||
|
||||
return super()._cached_apply_hf_processor(
|
||||
prompt=prompt,
|
||||
mm_data_items=mm_data_items,
|
||||
mm_uuid_items=mm_uuid_items,
|
||||
hf_processor_mm_kwargs=hf_processor_mm_kwargs,
|
||||
tokenization_kwargs=tokenization_kwargs,
|
||||
mm_uuids=mm_uuids,
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -30,7 +30,6 @@ from vllm.multimodal.inputs import (
|
||||
MultiModalFieldConfig,
|
||||
MultiModalInputs,
|
||||
MultiModalKwargsItems,
|
||||
MultiModalUUIDDict,
|
||||
mm_inputs,
|
||||
)
|
||||
from vllm.multimodal.parse import (
|
||||
@@ -38,6 +37,7 @@ from vllm.multimodal.parse import (
|
||||
ImageProcessorItems,
|
||||
ImageSize,
|
||||
MultiModalDataItems,
|
||||
MultiModalUUIDItems,
|
||||
)
|
||||
from vllm.multimodal.processing import (
|
||||
BaseDummyInputsBuilder,
|
||||
@@ -773,9 +773,9 @@ class MantisMultiModalProcessor(LlavaMultiModalProcessor):
|
||||
self,
|
||||
prompt: str | list[int],
|
||||
mm_items: MultiModalDataItems,
|
||||
hf_processor_mm_kwargs: Mapping[str, object],
|
||||
mm_uuid_items: MultiModalUUIDItems | None = None,
|
||||
hf_processor_mm_kwargs: Mapping[str, object] | None = None,
|
||||
tokenization_kwargs: Mapping[str, object] | None = None,
|
||||
mm_uuids: MultiModalUUIDDict | None = None,
|
||||
) -> MultiModalInputs:
|
||||
hf_config = self.info.get_hf_config()
|
||||
image_token_id = hf_config.image_token_index
|
||||
@@ -789,9 +789,9 @@ class MantisMultiModalProcessor(LlavaMultiModalProcessor):
|
||||
result = super().apply(
|
||||
prompt,
|
||||
mm_items,
|
||||
hf_processor_mm_kwargs,
|
||||
tokenization_kwargs,
|
||||
mm_uuids=mm_uuids,
|
||||
mm_uuid_items,
|
||||
hf_processor_mm_kwargs=hf_processor_mm_kwargs,
|
||||
tokenization_kwargs=tokenization_kwargs,
|
||||
)
|
||||
|
||||
mm_item_counts = mm_items.get_all_counts()
|
||||
|
||||
@@ -16,12 +16,12 @@ from vllm.multimodal.inputs import (
|
||||
MultiModalFieldConfig,
|
||||
MultiModalInputs,
|
||||
MultiModalKwargsItems,
|
||||
MultiModalUUIDDict,
|
||||
)
|
||||
from vllm.multimodal.parse import (
|
||||
ImageEmbeddingItems,
|
||||
ImageProcessorItems,
|
||||
MultiModalDataItems,
|
||||
MultiModalUUIDItems,
|
||||
)
|
||||
from vllm.multimodal.processing import (
|
||||
BaseDummyInputsBuilder,
|
||||
@@ -231,16 +231,16 @@ class PaliGemmaMultiModalProcessor(BaseMultiModalProcessor[PaliGemmaProcessingIn
|
||||
self,
|
||||
prompt: str | list[int],
|
||||
mm_items: MultiModalDataItems,
|
||||
hf_processor_mm_kwargs: Mapping[str, object],
|
||||
mm_uuid_items: MultiModalUUIDItems | None = None,
|
||||
hf_processor_mm_kwargs: Mapping[str, object] | None = None,
|
||||
tokenization_kwargs: Mapping[str, object] | None = None,
|
||||
mm_uuids: MultiModalUUIDDict | None = None,
|
||||
) -> MultiModalInputs:
|
||||
mm_inputs = super().apply(
|
||||
prompt,
|
||||
mm_items,
|
||||
hf_processor_mm_kwargs,
|
||||
tokenization_kwargs,
|
||||
mm_uuids=mm_uuids,
|
||||
mm_uuid_items,
|
||||
hf_processor_mm_kwargs=hf_processor_mm_kwargs,
|
||||
tokenization_kwargs=tokenization_kwargs,
|
||||
)
|
||||
prompt_token_ids = mm_inputs["prompt_token_ids"]
|
||||
|
||||
|
||||
@@ -44,10 +44,14 @@ from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalKwargsItems
|
||||
from vllm.multimodal.inputs import (
|
||||
MultiModalDataDict,
|
||||
MultiModalFieldConfig,
|
||||
MultiModalUUIDDict,
|
||||
NestedTensors,
|
||||
)
|
||||
from vllm.multimodal.parse import ImageProcessorItems, ImageSize, MultiModalDataItems
|
||||
from vllm.multimodal.parse import (
|
||||
ImageProcessorItems,
|
||||
ImageSize,
|
||||
MultiModalDataItems,
|
||||
MultiModalUUIDItems,
|
||||
)
|
||||
from vllm.multimodal.processing import BaseDummyInputsBuilder, ProcessorInputs
|
||||
from vllm.multimodal.processing.processor import (
|
||||
BaseMultiModalProcessor,
|
||||
@@ -344,16 +348,16 @@ class PixtralMultiModalProcessor(BaseMultiModalProcessor[PixtralProcessingInfo])
|
||||
self,
|
||||
prompt: str | list[int],
|
||||
mm_data_items: MultiModalDataItems,
|
||||
mm_uuid_items: MultiModalUUIDItems | None,
|
||||
hf_processor_mm_kwargs: Mapping[str, object],
|
||||
tokenization_kwargs: Mapping[str, object],
|
||||
mm_uuids: MultiModalUUIDDict | None = None,
|
||||
) -> tuple[list[int], MultiModalProcessingInfo, bool]:
|
||||
prompt_ids, mm_info, _ = super()._cached_apply_hf_processor(
|
||||
prompt=prompt,
|
||||
mm_data_items=mm_data_items,
|
||||
mm_uuid_items=mm_uuid_items,
|
||||
hf_processor_mm_kwargs=hf_processor_mm_kwargs,
|
||||
tokenization_kwargs=tokenization_kwargs,
|
||||
mm_uuids=mm_uuids,
|
||||
)
|
||||
|
||||
# NOTE: The tokens are already inserted by the chat template
|
||||
|
||||
@@ -30,36 +30,20 @@ from collections.abc import Callable, Iterable
|
||||
import torch
|
||||
from einops import rearrange
|
||||
from torch import nn
|
||||
from transformers.activations import ACT2FN
|
||||
|
||||
from vllm.compilation.decorators import support_torch_compile
|
||||
from vllm.config import (
|
||||
CacheConfig,
|
||||
ModelConfig,
|
||||
SpeculativeConfig,
|
||||
VllmConfig,
|
||||
get_current_vllm_config,
|
||||
)
|
||||
from vllm.distributed import (
|
||||
divide,
|
||||
get_pp_group,
|
||||
get_tensor_model_parallel_rank,
|
||||
get_tensor_model_parallel_world_size,
|
||||
)
|
||||
from vllm.logger import init_logger
|
||||
from vllm.model_executor.layers.layernorm import (
|
||||
GemmaRMSNorm as Qwen3_5RMSNorm,
|
||||
)
|
||||
from vllm.model_executor.layers.layernorm import RMSNormGated
|
||||
from vllm.model_executor.layers.linear import (
|
||||
ColumnParallelLinear,
|
||||
MergedColumnParallelLinear,
|
||||
RowParallelLinear,
|
||||
)
|
||||
from vllm.model_executor.layers.linear import MergedColumnParallelLinear
|
||||
from vllm.model_executor.layers.logits_processor import LogitsProcessor
|
||||
from vllm.model_executor.layers.mamba.mamba_mixer2 import (
|
||||
mamba_v2_sharded_weight_loader,
|
||||
)
|
||||
from vllm.model_executor.layers.mamba.mamba_utils import (
|
||||
MambaStateCopyFunc,
|
||||
MambaStateCopyFuncCalculator,
|
||||
@@ -73,11 +57,8 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
|
||||
)
|
||||
from vllm.model_executor.model_loader.weight_utils import (
|
||||
default_weight_loader,
|
||||
sharded_weight_loader,
|
||||
)
|
||||
from vllm.model_executor.utils import set_weight_attrs
|
||||
from vllm.multimodal import MULTIMODAL_REGISTRY
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.sequence import IntermediateTensors
|
||||
from vllm.transformers_utils.configs.qwen3_5 import (
|
||||
Qwen3_5Config,
|
||||
@@ -99,7 +80,6 @@ from .interfaces import (
|
||||
)
|
||||
from .qwen2_moe import Qwen2MoeMLP as Qwen3NextMLP
|
||||
from .qwen3_next import (
|
||||
ChunkGatedDeltaRule,
|
||||
Qwen3NextAttention,
|
||||
Qwen3NextDecoderLayer,
|
||||
Qwen3NextGatedDeltaNet,
|
||||
@@ -139,154 +119,31 @@ class Qwen3_5MoeProcessingInfo(Qwen3VLProcessingInfo):
|
||||
|
||||
|
||||
class Qwen3_5GatedDeltaNet(Qwen3NextGatedDeltaNet):
|
||||
def __init__(
|
||||
self,
|
||||
config: Qwen3_5TextConfig | Qwen3_5MoeTextConfig,
|
||||
model_config: ModelConfig | None = None,
|
||||
cache_config: CacheConfig | None = None,
|
||||
quant_config: QuantizationConfig | None = None,
|
||||
speculative_config: SpeculativeConfig | None = None,
|
||||
prefix: str = "",
|
||||
) -> None:
|
||||
super(Qwen3NextGatedDeltaNet, self).__init__()
|
||||
self.tp_size = get_tensor_model_parallel_world_size()
|
||||
self.tp_rank = get_tensor_model_parallel_rank()
|
||||
self.hidden_size = config.hidden_size
|
||||
self.num_v_heads = config.linear_num_value_heads
|
||||
self.num_k_heads = config.linear_num_key_heads
|
||||
self.head_k_dim = config.linear_key_head_dim
|
||||
self.head_v_dim = config.linear_value_head_dim
|
||||
self.key_dim = self.head_k_dim * self.num_k_heads
|
||||
self.value_dim = self.head_v_dim * self.num_v_heads
|
||||
|
||||
self.conv_kernel_size = config.linear_conv_kernel_dim
|
||||
self.layer_idx = extract_layer_index(prefix)
|
||||
self.activation = config.hidden_act
|
||||
self.act = ACT2FN[config.hidden_act]
|
||||
self.layer_norm_epsilon = config.rms_norm_eps
|
||||
self.prefix = prefix
|
||||
|
||||
self.config = config
|
||||
self.model_config = model_config
|
||||
self.cache_config = cache_config
|
||||
self.quant_config = quant_config
|
||||
self.speculative_config = speculative_config
|
||||
self.num_spec = (
|
||||
self.speculative_config.num_speculative_tokens
|
||||
if self.speculative_config
|
||||
else 0
|
||||
)
|
||||
|
||||
# QKV
|
||||
self.conv_dim = self.key_dim * 2 + self.value_dim
|
||||
self.conv1d = ColumnParallelLinear(
|
||||
input_size=self.conv_kernel_size,
|
||||
output_size=self.conv_dim,
|
||||
bias=False,
|
||||
prefix=f"{prefix}.conv1d",
|
||||
)
|
||||
self.conv1d.weight.data = self.conv1d.weight.data.unsqueeze(1)
|
||||
|
||||
self.in_proj_qkv = MergedColumnParallelLinear(
|
||||
input_size=self.hidden_size,
|
||||
output_sizes=[self.key_dim, self.key_dim, self.value_dim],
|
||||
bias=False,
|
||||
quant_config=quant_config,
|
||||
prefix=f"{prefix}.in_proj_qkv",
|
||||
)
|
||||
self.in_proj_z = ColumnParallelLinear(
|
||||
input_size=self.hidden_size,
|
||||
output_size=self.value_dim,
|
||||
bias=False,
|
||||
quant_config=quant_config,
|
||||
prefix=f"{prefix}.in_proj_z",
|
||||
)
|
||||
self.in_proj_b = ColumnParallelLinear(
|
||||
input_size=self.hidden_size,
|
||||
output_size=self.num_v_heads,
|
||||
bias=False,
|
||||
quant_config=quant_config,
|
||||
prefix=f"{prefix}.in_proj_b",
|
||||
)
|
||||
self.in_proj_a = ColumnParallelLinear(
|
||||
input_size=self.hidden_size,
|
||||
output_size=self.num_v_heads,
|
||||
bias=False,
|
||||
quant_config=quant_config,
|
||||
prefix=f"{prefix}.in_proj_a",
|
||||
)
|
||||
|
||||
query_key_settings = (self.key_dim, 0, False)
|
||||
value_settings = (self.value_dim, 0, False)
|
||||
|
||||
delattr(self.conv1d.weight, "weight_loader")
|
||||
set_weight_attrs(
|
||||
self.conv1d.weight,
|
||||
{
|
||||
"weight_loader": mamba_v2_sharded_weight_loader(
|
||||
[
|
||||
query_key_settings,
|
||||
query_key_settings,
|
||||
value_settings,
|
||||
],
|
||||
self.tp_size,
|
||||
self.tp_rank,
|
||||
)
|
||||
},
|
||||
)
|
||||
|
||||
# selective projection used to make dt, B and C input dependant
|
||||
|
||||
# time step projection (discretization)
|
||||
# instantiate once and copy inv_dt in init_weights of PretrainedModel
|
||||
self.dt_bias = nn.Parameter(
|
||||
torch.ones(self.num_v_heads // self.tp_size),
|
||||
)
|
||||
self.A_log = nn.Parameter(
|
||||
torch.empty(
|
||||
divide(self.num_v_heads, self.tp_size),
|
||||
)
|
||||
)
|
||||
|
||||
set_weight_attrs(self.A_log, {"weight_loader": sharded_weight_loader(0)})
|
||||
set_weight_attrs(self.dt_bias, {"weight_loader": sharded_weight_loader(0)})
|
||||
|
||||
self.norm = RMSNormGated(
|
||||
self.head_v_dim,
|
||||
eps=self.layer_norm_epsilon,
|
||||
group_size=None,
|
||||
norm_before_gate=True,
|
||||
device=current_platform.current_device(),
|
||||
dtype=config.dtype,
|
||||
)
|
||||
|
||||
self.out_proj = RowParallelLinear(
|
||||
self.value_dim,
|
||||
self.hidden_size,
|
||||
bias=False,
|
||||
input_is_parallel=True,
|
||||
quant_config=quant_config,
|
||||
prefix=f"{prefix}.out_proj",
|
||||
)
|
||||
|
||||
self.chunk_gated_delta_rule = ChunkGatedDeltaRule()
|
||||
|
||||
compilation_config = get_current_vllm_config().compilation_config
|
||||
if prefix in compilation_config.static_forward_context:
|
||||
raise ValueError(f"Duplicate layer name: {prefix}")
|
||||
compilation_config.static_forward_context[prefix] = self
|
||||
|
||||
def fix_query_key_value_ordering(
|
||||
self,
|
||||
mixed_qkv,
|
||||
z,
|
||||
b,
|
||||
a,
|
||||
mixed_qkvz: torch.Tensor,
|
||||
mixed_ba: torch.Tensor,
|
||||
):
|
||||
raise NotImplementedError(
|
||||
"Qwen3.5 Series dont need to fix query key value ordering"
|
||||
)
|
||||
|
||||
def create_qkvz_proj(
|
||||
self,
|
||||
hidden_size: int,
|
||||
key_dim: int,
|
||||
value_dim: int,
|
||||
quant_config: QuantizationConfig | None,
|
||||
prefix: str,
|
||||
) -> MergedColumnParallelLinear:
|
||||
return MergedColumnParallelLinear(
|
||||
input_size=hidden_size,
|
||||
output_sizes=[key_dim, key_dim, value_dim, value_dim],
|
||||
bias=False,
|
||||
quant_config=quant_config,
|
||||
prefix=prefix,
|
||||
)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
@@ -303,11 +160,13 @@ class Qwen3_5GatedDeltaNet(Qwen3NextGatedDeltaNet):
|
||||
# ============================================================
|
||||
# Part 1: Input Projection
|
||||
# ============================================================
|
||||
mixed_qkv, _ = self.in_proj_qkv(hidden_states)
|
||||
z, _ = self.in_proj_z(hidden_states)
|
||||
mixed_qkvz, _ = self.in_proj_qkvz(hidden_states)
|
||||
qkv_size = (self.key_dim * 2 + self.value_dim) // self.tp_size
|
||||
z_size = self.value_dim // self.tp_size
|
||||
mixed_qkv, z = mixed_qkvz.split([qkv_size, z_size], dim=-1)
|
||||
z = z.reshape(z.size(0), -1, self.head_v_dim)
|
||||
b, _ = self.in_proj_b(hidden_states)
|
||||
a, _ = self.in_proj_a(hidden_states)
|
||||
ba, _ = self.in_proj_ba(hidden_states)
|
||||
b, a = ba.chunk(2, dim=-1)
|
||||
|
||||
b = b.contiguous()
|
||||
a = a.contiguous()
|
||||
@@ -506,11 +365,18 @@ class Qwen3_5Model(Qwen3NextModel):
|
||||
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
|
||||
stacked_params_mapping = [
|
||||
# (param_name, shard_name, shard_id)
|
||||
# self attention
|
||||
("qkv_proj", "q_proj", "q"),
|
||||
("qkv_proj", "k_proj", "k"),
|
||||
("qkv_proj", "v_proj", "v"),
|
||||
# mlp
|
||||
("gate_up_proj", "gate_proj", 0),
|
||||
("gate_up_proj", "up_proj", 1),
|
||||
# GDN
|
||||
("in_proj_qkvz", "in_proj_qkv", (0, 1, 2)),
|
||||
("in_proj_qkvz", "in_proj_z", 3),
|
||||
("in_proj_ba", "in_proj_b", 0),
|
||||
("in_proj_ba", "in_proj_a", 1),
|
||||
]
|
||||
|
||||
params_dict = dict(self.named_parameters())
|
||||
@@ -657,6 +523,9 @@ class Qwen3_5ForCausalLMBase(
|
||||
"v_proj",
|
||||
],
|
||||
"gate_up_proj": ["gate_proj", "up_proj"],
|
||||
# GDN fused projections.
|
||||
"in_proj_qkvz": ["in_proj_qkv", "in_proj_z"],
|
||||
"in_proj_ba": ["in_proj_b", "in_proj_a"],
|
||||
}
|
||||
|
||||
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
|
||||
@@ -754,7 +623,12 @@ class Qwen3_5MoeForCausalLM(Qwen3_5ForCausalLMBase, QwenNextMixtureOfExperts):
|
||||
dummy_inputs=Qwen3VLDummyInputsBuilder,
|
||||
)
|
||||
class Qwen3_5ForConditionalGeneration(Qwen3VLForConditionalGeneration, IsHybrid):
|
||||
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
|
||||
packed_modules_mapping = Qwen3VLForConditionalGeneration.packed_modules_mapping | {
|
||||
"in_proj_qkvz": ["in_proj_qkv", "in_proj_z"],
|
||||
"in_proj_ba": ["in_proj_b", "in_proj_a"],
|
||||
}
|
||||
|
||||
def __init__(self, *, vllm_config: VllmConfig, prefix: str = "model"):
|
||||
# protocols have not __init__ method, so we need to use nn.Module.__init__
|
||||
nn.Module.__init__(self)
|
||||
config: Qwen3_5Config = vllm_config.model_config.hf_config
|
||||
@@ -962,7 +836,7 @@ class Qwen3_5_MoeMixtureOfExperts(MixtureOfExperts):
|
||||
class Qwen3_5MoeForConditionalGeneration(
|
||||
Qwen3_5ForConditionalGeneration, Qwen3_5_MoeMixtureOfExperts
|
||||
):
|
||||
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
|
||||
def __init__(self, *, vllm_config: VllmConfig, prefix: str = "model"):
|
||||
# protocols have not __init__ method, so we need to use nn.Module.__init__
|
||||
nn.Module.__init__(self)
|
||||
config: Qwen3_5MoeConfig = vllm_config.model_config.hf_config
|
||||
|
||||
@@ -44,6 +44,7 @@ from vllm.model_executor.layers.layernorm import (
|
||||
from vllm.model_executor.layers.layernorm import RMSNormGated
|
||||
from vllm.model_executor.layers.linear import (
|
||||
ColumnParallelLinear,
|
||||
MergedColumnParallelLinear,
|
||||
QKVParallelLinear,
|
||||
ReplicatedLinear,
|
||||
RowParallelLinear,
|
||||
@@ -406,19 +407,19 @@ class Qwen3NextGatedDeltaNet(nn.Module, MambaBase):
|
||||
self.conv1d.weight.data = self.conv1d.weight.data.unsqueeze(1)
|
||||
|
||||
# projection of the input hidden states
|
||||
self.projection_size_qkvz = self.key_dim * 2 + self.value_dim * 2
|
||||
self.projection_size_ba = self.num_v_heads * 2
|
||||
self.in_proj_qkvz = ColumnParallelLinear(
|
||||
input_size=self.hidden_size,
|
||||
output_size=self.projection_size_qkvz,
|
||||
bias=False,
|
||||
# Qwen3-Next and Qwen3.5 has a different qkv_proj layout,
|
||||
# we need to create qkvz_proj adaptively here.
|
||||
self.in_proj_qkvz = self.create_qkvz_proj(
|
||||
hidden_size=self.hidden_size,
|
||||
key_dim=self.key_dim,
|
||||
value_dim=self.value_dim,
|
||||
quant_config=quant_config,
|
||||
prefix=f"{prefix}.in_proj_qkvz",
|
||||
)
|
||||
# ba_proj doesn't support blockwise fp8 quantization.
|
||||
self.in_proj_ba = ColumnParallelLinear(
|
||||
self.in_proj_ba = MergedColumnParallelLinear(
|
||||
input_size=self.hidden_size,
|
||||
output_size=self.projection_size_ba,
|
||||
output_sizes=[self.num_v_heads] * 2,
|
||||
bias=False,
|
||||
quant_config=quant_config,
|
||||
prefix=f"{prefix}.in_proj_ba",
|
||||
@@ -484,10 +485,26 @@ class Qwen3NextGatedDeltaNet(nn.Module, MambaBase):
|
||||
raise ValueError(f"Duplicate layer name: {prefix}")
|
||||
compilation_config.static_forward_context[prefix] = self
|
||||
|
||||
def create_qkvz_proj(
|
||||
self,
|
||||
hidden_size: int,
|
||||
key_dim: int,
|
||||
value_dim: int,
|
||||
quant_config: QuantizationConfig | None,
|
||||
prefix: str,
|
||||
) -> MergedColumnParallelLinear:
|
||||
return MergedColumnParallelLinear(
|
||||
input_size=hidden_size,
|
||||
output_sizes=[sum((key_dim, key_dim, value_dim, value_dim))],
|
||||
bias=False,
|
||||
quant_config=quant_config,
|
||||
prefix=prefix,
|
||||
)
|
||||
|
||||
def fix_query_key_value_ordering(
|
||||
self,
|
||||
mixed_qkvz,
|
||||
mixed_ba,
|
||||
mixed_qkvz: torch.Tensor,
|
||||
mixed_ba: torch.Tensor,
|
||||
):
|
||||
"""
|
||||
Derives `query`, `key` and `value` tensors from `mixed_qkvzba`.
|
||||
|
||||
@@ -42,9 +42,13 @@ from vllm.multimodal.inputs import (
|
||||
MultiModalFieldConfig,
|
||||
MultiModalInputs,
|
||||
MultiModalKwargsItems,
|
||||
MultiModalUUIDDict,
|
||||
)
|
||||
from vllm.multimodal.parse import ImageProcessorItems, ImageSize, MultiModalDataItems
|
||||
from vllm.multimodal.parse import (
|
||||
ImageProcessorItems,
|
||||
ImageSize,
|
||||
MultiModalDataItems,
|
||||
MultiModalUUIDItems,
|
||||
)
|
||||
from vllm.multimodal.processing import (
|
||||
BaseDummyInputsBuilder,
|
||||
BaseMultiModalProcessor,
|
||||
@@ -189,10 +193,9 @@ class SiglipMultiModalProcessor(BaseMultiModalProcessor[SiglipProcessingInfo]):
|
||||
self,
|
||||
prompt: str | list[int],
|
||||
mm_items: MultiModalDataItems,
|
||||
hf_processor_mm_kwargs: Mapping[str, object],
|
||||
mm_uuid_items: MultiModalUUIDItems | None = None,
|
||||
hf_processor_mm_kwargs: Mapping[str, object] | None = None,
|
||||
tokenization_kwargs: Mapping[str, object] | None = None,
|
||||
*,
|
||||
mm_uuids: MultiModalUUIDDict | None = None,
|
||||
) -> MultiModalInputs:
|
||||
if mm_items:
|
||||
if isinstance(prompt, str):
|
||||
@@ -221,9 +224,9 @@ class SiglipMultiModalProcessor(BaseMultiModalProcessor[SiglipProcessingInfo]):
|
||||
return super().apply(
|
||||
prompt=prompt,
|
||||
mm_items=mm_items,
|
||||
mm_uuid_items=mm_uuid_items,
|
||||
hf_processor_mm_kwargs=hf_processor_mm_kwargs,
|
||||
tokenization_kwargs=tokenization_kwargs,
|
||||
mm_uuids=mm_uuids,
|
||||
)
|
||||
|
||||
def _hf_processor_applies_updates(
|
||||
|
||||
@@ -46,7 +46,6 @@ from vllm.multimodal.inputs import (
|
||||
MultiModalFieldConfig,
|
||||
MultiModalInputs,
|
||||
MultiModalKwargsItems,
|
||||
MultiModalUUIDDict,
|
||||
PlaceholderRange,
|
||||
mm_inputs,
|
||||
)
|
||||
@@ -55,6 +54,7 @@ from vllm.multimodal.parse import (
|
||||
ModalityDataItems,
|
||||
MultiModalDataItems,
|
||||
MultiModalDataParser,
|
||||
MultiModalUUIDItems,
|
||||
)
|
||||
from vllm.multimodal.processing import (
|
||||
BaseDummyInputsBuilder,
|
||||
@@ -196,15 +196,19 @@ class TerratorchMultiModalProcessor(BaseMultiModalProcessor[TerratorchProcessing
|
||||
self,
|
||||
prompt: str | list[int],
|
||||
mm_items: MultiModalDataItems,
|
||||
hf_processor_mm_kwargs: Mapping[str, object],
|
||||
mm_uuid_items: MultiModalUUIDItems | None = None,
|
||||
hf_processor_mm_kwargs: Mapping[str, object] | None = None,
|
||||
tokenization_kwargs: Mapping[str, object] | None = None,
|
||||
mm_uuids: MultiModalUUIDDict | None = None,
|
||||
) -> MultiModalInputs:
|
||||
if hf_processor_mm_kwargs is None:
|
||||
hf_processor_mm_kwargs = {}
|
||||
if tokenization_kwargs is None:
|
||||
tokenization_kwargs = {}
|
||||
|
||||
mm_hashes = self._hash_mm_items(
|
||||
mm_items, hf_processor_mm_kwargs, tokenization_kwargs, mm_uuids=mm_uuids
|
||||
mm_items,
|
||||
mm_uuid_items,
|
||||
hf_processor_mm_kwargs=hf_processor_mm_kwargs,
|
||||
)
|
||||
|
||||
_, passthrough_data = self._get_hf_mm_data(mm_items)
|
||||
|
||||
@@ -31,11 +31,14 @@ from vllm.multimodal.inputs import (
|
||||
MultiModalFeatureSpec,
|
||||
MultiModalFieldConfig,
|
||||
MultiModalInputs,
|
||||
MultiModalUUIDDict,
|
||||
PlaceholderRange,
|
||||
mm_inputs,
|
||||
)
|
||||
from vllm.multimodal.parse import ImageProcessorItems, MultiModalDataItems
|
||||
from vllm.multimodal.parse import (
|
||||
ImageProcessorItems,
|
||||
MultiModalDataItems,
|
||||
MultiModalUUIDItems,
|
||||
)
|
||||
from vllm.multimodal.processing import (
|
||||
BaseDummyInputsBuilder,
|
||||
BaseMultiModalProcessor,
|
||||
@@ -177,9 +180,9 @@ class MultiModalProcessor(BaseMultiModalProcessor[MultiModalProcessingInfo]):
|
||||
self,
|
||||
prompt: str | list[int],
|
||||
mm_items: MultiModalDataItems,
|
||||
hf_processor_mm_kwargs: Mapping[str, object],
|
||||
mm_uuid_items: MultiModalUUIDItems | None = None,
|
||||
hf_processor_mm_kwargs: Mapping[str, object] | None = None,
|
||||
tokenization_kwargs: Mapping[str, object] | None = None,
|
||||
mm_uuids: MultiModalUUIDDict | None = None,
|
||||
) -> MultiModalInputs:
|
||||
"""
|
||||
Process multi-modal inputs to be used in vLLM.
|
||||
@@ -187,6 +190,8 @@ class MultiModalProcessor(BaseMultiModalProcessor[MultiModalProcessingInfo]):
|
||||
Apply HF Processor on prompt text and multi-modal data together,
|
||||
outputting token IDs and processed tensors.
|
||||
"""
|
||||
if hf_processor_mm_kwargs is None:
|
||||
hf_processor_mm_kwargs = {}
|
||||
if tokenization_kwargs is None:
|
||||
tokenization_kwargs = {}
|
||||
|
||||
@@ -258,7 +263,9 @@ class MultiModalProcessor(BaseMultiModalProcessor[MultiModalProcessingInfo]):
|
||||
|
||||
# Use overrides if provided; fallback to data-dependent hashing.
|
||||
mm_hashes = self._hash_mm_items(
|
||||
mm_items, hf_processor_mm_kwargs, tokenization_kwargs, mm_uuids=mm_uuids
|
||||
mm_items,
|
||||
mm_uuid_items,
|
||||
hf_processor_mm_kwargs=hf_processor_mm_kwargs,
|
||||
)
|
||||
|
||||
return mm_inputs(
|
||||
|
||||
@@ -41,13 +41,13 @@ from vllm.multimodal.inputs import (
|
||||
MultiModalDataDict,
|
||||
MultiModalFieldConfig,
|
||||
MultiModalKwargsItems,
|
||||
MultiModalUUIDDict,
|
||||
NestedTensors,
|
||||
)
|
||||
from vllm.multimodal.parse import (
|
||||
AudioProcessorItems,
|
||||
MultiModalDataItems,
|
||||
MultiModalDataParser,
|
||||
MultiModalUUIDItems,
|
||||
)
|
||||
from vllm.multimodal.processing import BaseDummyInputsBuilder, ProcessorInputs
|
||||
from vllm.multimodal.processing.processor import (
|
||||
@@ -363,16 +363,16 @@ class VoxtralMultiModalProcessor(BaseMultiModalProcessor[VoxtralProcessingInfo])
|
||||
self,
|
||||
prompt: str | list[int],
|
||||
mm_data_items: MultiModalDataItems,
|
||||
mm_uuid_items: MultiModalUUIDItems | None,
|
||||
hf_processor_mm_kwargs: Mapping[str, object],
|
||||
tokenization_kwargs: Mapping[str, object],
|
||||
mm_uuids: MultiModalUUIDDict | None = None,
|
||||
) -> tuple[list[int], MultiModalProcessingInfo, bool]:
|
||||
prompt_ids, mm_info, _ = super()._cached_apply_hf_processor(
|
||||
prompt=prompt,
|
||||
mm_data_items=mm_data_items,
|
||||
mm_uuid_items=mm_uuid_items,
|
||||
hf_processor_mm_kwargs=hf_processor_mm_kwargs,
|
||||
tokenization_kwargs=tokenization_kwargs,
|
||||
mm_uuids=mm_uuids,
|
||||
)
|
||||
|
||||
# NOTE: The tokens are already inserted by the chat template
|
||||
|
||||
@@ -155,7 +155,7 @@ The built-in modalities are defined by
|
||||
[`MultiModalDataBuiltins`][vllm.multimodal.inputs.MultiModalDataBuiltins].
|
||||
"""
|
||||
|
||||
MultiModalUUIDDict: TypeAlias = Mapping[str, list[str | None] | str]
|
||||
MultiModalUUIDDict: TypeAlias = Mapping[str, Sequence[str | None] | str]
|
||||
"""
|
||||
A dictionary containing user-provided UUIDs for items in each modality.
|
||||
If a UUID for an item is not provided, its entry will be `None` and
|
||||
|
||||
@@ -146,7 +146,7 @@ class MediaConnector:
|
||||
|
||||
connection = self.connection
|
||||
data = connection.get_bytes(
|
||||
url,
|
||||
url_spec.url,
|
||||
timeout=fetch_timeout,
|
||||
allow_redirects=envs.VLLM_MEDIA_URL_ALLOW_REDIRECTS,
|
||||
)
|
||||
@@ -177,7 +177,7 @@ class MediaConnector:
|
||||
|
||||
connection = self.connection
|
||||
data = await connection.async_get_bytes(
|
||||
url,
|
||||
url_spec.url,
|
||||
timeout=fetch_timeout,
|
||||
allow_redirects=envs.VLLM_MEDIA_URL_ALLOW_REDIRECTS,
|
||||
)
|
||||
|
||||
+63
-48
@@ -3,7 +3,7 @@
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from collections import UserDict
|
||||
from collections.abc import Callable, Iterator, Mapping, Sequence
|
||||
from collections.abc import Callable, Iterator, Mapping, Sequence, Set
|
||||
from typing import (
|
||||
TYPE_CHECKING,
|
||||
Any,
|
||||
@@ -33,6 +33,7 @@ from .inputs import (
|
||||
MultiModalDataDict,
|
||||
MultiModalFieldConfig,
|
||||
MultiModalKwargsItems,
|
||||
MultiModalUUIDDict,
|
||||
VideoItem,
|
||||
)
|
||||
from .media import MediaWithBytes
|
||||
@@ -297,14 +298,15 @@ class DictEmbeddingItems(
|
||||
return self.data
|
||||
|
||||
|
||||
class AudioProcessorItems(ProcessorBatchItems[HfAudioItem]):
|
||||
def __init__(self, data: Sequence[HfAudioItem] | None) -> None:
|
||||
if data is None:
|
||||
data = [None]
|
||||
class AudioProcessorItems(ProcessorBatchItems[HfAudioItem | None]):
|
||||
def __init__(self, data: Sequence[HfAudioItem | None]) -> None:
|
||||
super().__init__(data, "audio")
|
||||
|
||||
def get_audio_length(self, item_idx: int) -> int:
|
||||
audio = self.get(item_idx)
|
||||
if audio is None:
|
||||
raise ValueError(f"Cannot get length of cached audio at {item_idx}")
|
||||
|
||||
return len(audio)
|
||||
|
||||
|
||||
@@ -322,14 +324,14 @@ class ImageSize(NamedTuple):
|
||||
height: int
|
||||
|
||||
|
||||
class ImageProcessorItems(ProcessorBatchItems[HfImageItem]):
|
||||
def __init__(self, data: Sequence[HfImageItem] | None) -> None:
|
||||
if data is None:
|
||||
data = [None]
|
||||
class ImageProcessorItems(ProcessorBatchItems[HfImageItem | None]):
|
||||
def __init__(self, data: Sequence[HfImageItem | None]) -> None:
|
||||
super().__init__(data, "image")
|
||||
|
||||
def get_image_size(self, item_idx: int) -> ImageSize:
|
||||
image = self.get(item_idx)
|
||||
if image is None:
|
||||
raise ValueError(f"Cannot get size of cached image at {item_idx}")
|
||||
|
||||
if isinstance(image, PILImage.Image):
|
||||
return ImageSize(*image.size)
|
||||
@@ -349,22 +351,31 @@ class ImageEmbeddingItems(EmbeddingItems):
|
||||
super().__init__(data, "image", expected_hidden_size)
|
||||
|
||||
|
||||
class VideoProcessorItems(ProcessorBatchItems[HfVideoItem]):
|
||||
class VideoProcessorItems(ProcessorBatchItems[HfVideoItem | None]):
|
||||
def __init__(
|
||||
self,
|
||||
data: Sequence[HfVideoItem] | None,
|
||||
data: Sequence[HfVideoItem | None],
|
||||
metadata: dict[str, Any] | list[dict[str, Any] | None] | None = None,
|
||||
) -> None:
|
||||
if data is None:
|
||||
data = [None]
|
||||
super().__init__(data, "video")
|
||||
|
||||
self.metadata = metadata
|
||||
|
||||
def get_num_frames(self, item_idx: int) -> int:
|
||||
return len(self.get(item_idx))
|
||||
video = self.get(item_idx)
|
||||
if video is None:
|
||||
raise ValueError(f"Cannot get length of cached video at {item_idx}")
|
||||
|
||||
return len(video)
|
||||
|
||||
def get_frame_size(self, item_idx: int) -> ImageSize:
|
||||
image = self.get(item_idx)[0] # Assume that the video isn't empty
|
||||
video = self.get(item_idx)
|
||||
if video is None:
|
||||
raise ValueError(f"Cannot get size of cached video at {item_idx}")
|
||||
if len(video) == 0:
|
||||
raise ValueError(f"Cannot get size of empty video at {item_idx}")
|
||||
|
||||
image = video[0]
|
||||
|
||||
if isinstance(image, PILImage.Image):
|
||||
return ImageSize(*image.size)
|
||||
@@ -400,6 +411,15 @@ class MultiModalDataItems(UserDict[str, ModalityDataItems[Any, Any]]):
|
||||
normalized such that each entry corresponds to a list.
|
||||
"""
|
||||
|
||||
def select(self, modalities: Set[str]):
|
||||
"""
|
||||
Construct a new `MultiModalDataItems` instance containing only the
|
||||
selected modalities.
|
||||
"""
|
||||
return MultiModalDataItems(
|
||||
{modality: self[modality] for modality in modalities}
|
||||
)
|
||||
|
||||
def get_count(self, modality: str, *, strict: bool = True) -> int:
|
||||
"""
|
||||
Get the number of data items belonging to a modality.
|
||||
@@ -497,19 +517,11 @@ class MultiModalDataParser:
|
||||
) -> TypeGuard[torch.Tensor | list[torch.Tensor]]:
|
||||
if isinstance(data, torch.Tensor):
|
||||
return data.ndim == 3
|
||||
if is_list_of(data, torch.Tensor):
|
||||
if is_list_of(data, torch.Tensor) and len(data) > 0:
|
||||
return data[0].ndim == 2 # type: ignore[index]
|
||||
|
||||
return False
|
||||
|
||||
def _is_empty(self, data: object) -> TypeGuard[None]:
|
||||
if isinstance(data, list):
|
||||
return len(data) == 0
|
||||
if isinstance(data, (np.ndarray, torch.Tensor)):
|
||||
return data.size == 0
|
||||
|
||||
return False
|
||||
|
||||
def _get_audio_with_sr(
|
||||
self,
|
||||
audio: AudioItem,
|
||||
@@ -545,12 +557,6 @@ class MultiModalDataParser:
|
||||
data: ModalityData[AudioItem],
|
||||
) -> ModalityDataItems[Any, Any] | None:
|
||||
if data is None:
|
||||
return AudioProcessorItems(None)
|
||||
|
||||
# also check single audio item with sampling rate
|
||||
if self._is_empty(data) or (
|
||||
isinstance(data, tuple) and self._is_empty(data[0])
|
||||
):
|
||||
return None
|
||||
|
||||
if self.is_embeddings(data):
|
||||
@@ -558,9 +564,8 @@ class MultiModalDataParser:
|
||||
|
||||
data_items: list[AudioItem]
|
||||
if (
|
||||
is_list_of(data, float)
|
||||
or isinstance(data, (np.ndarray, torch.Tensor))
|
||||
and data.ndim == 1
|
||||
(is_list_of(data, float) and len(data) > 0)
|
||||
or (isinstance(data, (np.ndarray, torch.Tensor)) and data.ndim == 1)
|
||||
or isinstance(data, tuple)
|
||||
):
|
||||
data_items = [data]
|
||||
@@ -591,18 +596,13 @@ class MultiModalDataParser:
|
||||
data: ModalityData[ImageItem],
|
||||
) -> ModalityDataItems[Any, Any] | None:
|
||||
if data is None:
|
||||
return ImageProcessorItems(None)
|
||||
|
||||
if self._is_empty(data):
|
||||
return None
|
||||
|
||||
if self.is_embeddings(data):
|
||||
return ImageEmbeddingItems(data, self.expected_hidden_size)
|
||||
|
||||
if (
|
||||
isinstance(data, (PILImage.Image, MediaWithBytes))
|
||||
or isinstance(data, (np.ndarray, torch.Tensor))
|
||||
and data.ndim == 3
|
||||
if isinstance(data, (PILImage.Image, MediaWithBytes)) or (
|
||||
isinstance(data, (np.ndarray, torch.Tensor)) and data.ndim == 3
|
||||
):
|
||||
data_items = [data]
|
||||
elif isinstance(data, (np.ndarray, torch.Tensor)):
|
||||
@@ -617,19 +617,14 @@ class MultiModalDataParser:
|
||||
data: ModalityData[VideoItem],
|
||||
) -> ModalityDataItems[Any, Any] | None:
|
||||
if data is None:
|
||||
return VideoProcessorItems(None)
|
||||
|
||||
if self._is_empty(data):
|
||||
return None
|
||||
|
||||
if self.is_embeddings(data):
|
||||
return VideoEmbeddingItems(data, self.expected_hidden_size)
|
||||
|
||||
data_items: list[VideoItem]
|
||||
if (
|
||||
is_list_of(data, PILImage.Image)
|
||||
or isinstance(data, (np.ndarray, torch.Tensor))
|
||||
and data.ndim == 4
|
||||
if (is_list_of(data, PILImage.Image) and len(data) > 0) or (
|
||||
isinstance(data, (np.ndarray, torch.Tensor)) and data.ndim == 4
|
||||
):
|
||||
data_items = [data]
|
||||
elif isinstance(data, (np.ndarray, torch.Tensor)):
|
||||
@@ -664,12 +659,15 @@ class MultiModalDataParser:
|
||||
data: ModalityData[Any],
|
||||
) -> ModalityDataItems[Any, Any] | None:
|
||||
"""Parse vision chunk data (unified image and video chunks)."""
|
||||
if data is None or self._is_empty(data):
|
||||
if data is None:
|
||||
return None
|
||||
|
||||
if self.is_embeddings(data):
|
||||
raise ValueError("Do not support embedding data for vision_chunk right now")
|
||||
|
||||
if isinstance(data, dict):
|
||||
data = [data]
|
||||
|
||||
return VisionChunkProcessorItems(data)
|
||||
|
||||
def _get_subparsers(self) -> Mapping[str, ModalityDataParser]:
|
||||
@@ -693,3 +691,20 @@ class MultiModalDataParser:
|
||||
mm_items[k] = parsed_data
|
||||
|
||||
return mm_items
|
||||
|
||||
|
||||
MultiModalUUIDItems: TypeAlias = dict[str, Sequence[str | None]]
|
||||
"""
|
||||
As [`MultiModalUUIDDict`][vllm.multimodal.inputs.MultiModalUUIDDict], but
|
||||
normalized such that each entry corresponds to a list.
|
||||
"""
|
||||
|
||||
|
||||
def parse_mm_uuids(mm_uuids: MultiModalUUIDDict | None) -> MultiModalUUIDItems:
|
||||
if mm_uuids is None:
|
||||
return {}
|
||||
|
||||
return {
|
||||
modality: [uuids] if isinstance(uuids, str) else uuids
|
||||
for modality, uuids in mm_uuids.items()
|
||||
}
|
||||
|
||||
@@ -32,7 +32,6 @@ from ..inputs import (
|
||||
MultiModalKwargsItem,
|
||||
MultiModalKwargsItems,
|
||||
MultiModalKwargsOptionalItems,
|
||||
MultiModalUUIDDict,
|
||||
PlaceholderRange,
|
||||
mm_enc_dec_inputs,
|
||||
mm_inputs,
|
||||
@@ -41,6 +40,7 @@ from ..parse import (
|
||||
DictEmbeddingItems,
|
||||
EmbeddingItems,
|
||||
MultiModalDataItems,
|
||||
MultiModalUUIDItems,
|
||||
)
|
||||
from .context import (
|
||||
BaseProcessingInfo,
|
||||
@@ -1014,11 +1014,15 @@ class BaseMultiModalProcessor(ABC, Generic[_I]):
|
||||
self,
|
||||
prompt: str,
|
||||
mm_items: MultiModalDataItems,
|
||||
hf_processor_mm_kwargs: Mapping[str, object],
|
||||
*,
|
||||
mm_uuids: MultiModalUUIDDict | None = None,
|
||||
mm_uuid_items: MultiModalUUIDItems | None = None,
|
||||
hf_processor_mm_kwargs: Mapping[str, object] | None = None,
|
||||
) -> MultiModalInputs:
|
||||
return self.apply(prompt, mm_items, hf_processor_mm_kwargs, mm_uuids=mm_uuids)
|
||||
return self.apply(
|
||||
prompt,
|
||||
mm_items,
|
||||
mm_uuid_items,
|
||||
hf_processor_mm_kwargs=hf_processor_mm_kwargs,
|
||||
)
|
||||
|
||||
@abstractmethod
|
||||
def _get_mm_fields_config(
|
||||
@@ -1174,7 +1178,10 @@ class BaseMultiModalProcessor(ABC, Generic[_I]):
|
||||
|
||||
In addition, return whether prompt updates have been applied.
|
||||
"""
|
||||
processor_data, passthrough_data = self._get_hf_mm_data(mm_items)
|
||||
valid_mm_items = mm_items.select(
|
||||
{k for k, c in mm_items.get_all_counts().items() if c > 0}
|
||||
)
|
||||
processor_data, passthrough_data = self._get_hf_mm_data(valid_mm_items)
|
||||
|
||||
processed_data = self._call_hf_processor(
|
||||
prompt=prompt_text,
|
||||
@@ -1301,69 +1308,57 @@ class BaseMultiModalProcessor(ABC, Generic[_I]):
|
||||
|
||||
def _hash_mm_items(
|
||||
self,
|
||||
mm_items: MultiModalDataItems,
|
||||
mm_data_items: MultiModalDataItems,
|
||||
mm_uuid_items: MultiModalUUIDItems | None,
|
||||
hf_processor_mm_kwargs: Mapping[str, object],
|
||||
tokenization_kwargs: Mapping[str, object],
|
||||
*,
|
||||
mm_uuids: MultiModalUUIDDict | None = None,
|
||||
) -> MultiModalHashes:
|
||||
"""Create MM hashes to be returned.
|
||||
|
||||
|
||||
Note: When overrides are provided via callers of `apply`,
|
||||
`_hash_mm_items` will be bypassed and the overrides will be used.
|
||||
"""
|
||||
model_id = self.info.model_id
|
||||
|
||||
hashes: MultiModalHashes = {}
|
||||
mm_uuids = mm_uuids or {}
|
||||
if mm_uuid_items is None:
|
||||
mm_uuid_items = {}
|
||||
|
||||
for modality, items in mm_items.items():
|
||||
if modality in mm_uuids:
|
||||
mm_uuids_per_modality = mm_uuids[modality]
|
||||
if isinstance(mm_uuids_per_modality, str):
|
||||
mm_uuids_per_modality = [mm_uuids_per_modality]
|
||||
mm_hashes: MultiModalHashes = {}
|
||||
hasher = MultiModalHasher
|
||||
|
||||
for modality, data_items in mm_data_items.items():
|
||||
if modality in mm_uuid_items:
|
||||
uuid_items = mm_uuid_items[modality]
|
||||
|
||||
# For None entries, compute a hash; otherwise, use provided ID.
|
||||
computed: list[str] = []
|
||||
for i, item in enumerate(items.get_all_items_for_hash()):
|
||||
item_uuid = mm_uuids_per_modality[i]
|
||||
hashes: list[str] = []
|
||||
for i, item in enumerate(data_items.get_all_items_for_hash()):
|
||||
uuid_item = uuid_items[i]
|
||||
|
||||
# NOTE: Even if a item_uuid is provided, we still compute a
|
||||
# hash if `hf_processor_mm_kwargs` or `tokenization_kwargs`
|
||||
# are provided. This is because the processed multimodal
|
||||
# inputs can be different depending on the processor kwargs.
|
||||
if (
|
||||
item_uuid is None
|
||||
or hf_processor_mm_kwargs
|
||||
or tokenization_kwargs
|
||||
):
|
||||
# NOTE: Even if a uuid_item is provided, we still compute a hash
|
||||
# if `hf_processor_mm_kwargs` is provided.
|
||||
# This is because the processed multimodal inputs can be different
|
||||
# depending on the processor kwargs.
|
||||
if uuid_item is None or hf_processor_mm_kwargs:
|
||||
# NOTE: use provided hash string to hash with kwargs
|
||||
# if available for better performance.
|
||||
item = item_uuid if item_uuid is not None else item
|
||||
computed.append(
|
||||
MultiModalHasher.hash_kwargs(
|
||||
item = uuid_item if uuid_item is not None else item
|
||||
hashes.append(
|
||||
hasher.hash_kwargs(
|
||||
model_id=model_id,
|
||||
**{modality: item},
|
||||
**hf_processor_mm_kwargs,
|
||||
**tokenization_kwargs,
|
||||
)
|
||||
)
|
||||
else:
|
||||
computed.append(item_uuid)
|
||||
hashes[modality] = computed
|
||||
hashes.append(uuid_item)
|
||||
|
||||
mm_hashes[modality] = hashes
|
||||
else:
|
||||
hashes[modality] = [
|
||||
MultiModalHasher.hash_kwargs(
|
||||
mm_hashes[modality] = [
|
||||
hasher.hash_kwargs(
|
||||
model_id=model_id,
|
||||
**{modality: item},
|
||||
**hf_processor_mm_kwargs,
|
||||
**tokenization_kwargs,
|
||||
)
|
||||
for item in items
|
||||
for item in data_items
|
||||
]
|
||||
|
||||
return hashes
|
||||
return mm_hashes
|
||||
|
||||
def _get_cache_missing_items(
|
||||
self,
|
||||
@@ -1468,10 +1463,9 @@ class BaseMultiModalProcessor(ABC, Generic[_I]):
|
||||
self,
|
||||
prompt: str | list[int],
|
||||
mm_data_items: MultiModalDataItems,
|
||||
mm_uuid_items: MultiModalUUIDItems | None,
|
||||
hf_processor_mm_kwargs: Mapping[str, object],
|
||||
tokenization_kwargs: Mapping[str, object],
|
||||
*,
|
||||
mm_uuids: MultiModalUUIDDict | None = None,
|
||||
) -> tuple[list[int], MultiModalProcessingInfo, bool]:
|
||||
(
|
||||
prompt_ids,
|
||||
@@ -1494,9 +1488,8 @@ class BaseMultiModalProcessor(ABC, Generic[_I]):
|
||||
with timed_preprocessor_operation(self.info.ctx, "hashing"):
|
||||
mm_hashes = self._hash_mm_items(
|
||||
mm_data_items,
|
||||
hf_processor_mm_kwargs,
|
||||
tokenization_kwargs,
|
||||
mm_uuids=mm_uuids,
|
||||
mm_uuid_items,
|
||||
hf_processor_mm_kwargs=hf_processor_mm_kwargs,
|
||||
)
|
||||
|
||||
mm_prompt_updates = self._get_mm_prompt_updates(
|
||||
@@ -1517,10 +1510,9 @@ class BaseMultiModalProcessor(ABC, Generic[_I]):
|
||||
self,
|
||||
prompt: str | list[int],
|
||||
mm_data_items: MultiModalDataItems,
|
||||
mm_uuid_items: MultiModalUUIDItems | None,
|
||||
hf_processor_mm_kwargs: Mapping[str, object],
|
||||
tokenization_kwargs: Mapping[str, object],
|
||||
*,
|
||||
mm_uuids: MultiModalUUIDDict | None = None,
|
||||
) -> tuple[list[int], MultiModalProcessingInfo, bool]:
|
||||
"""
|
||||
Apply the HF processor on the full prompt text,
|
||||
@@ -1533,17 +1525,16 @@ class BaseMultiModalProcessor(ABC, Generic[_I]):
|
||||
return self._apply_hf_processor(
|
||||
prompt=prompt,
|
||||
mm_data_items=mm_data_items,
|
||||
mm_uuid_items=mm_uuid_items,
|
||||
hf_processor_mm_kwargs=hf_processor_mm_kwargs,
|
||||
tokenization_kwargs=tokenization_kwargs,
|
||||
mm_uuids=mm_uuids,
|
||||
)
|
||||
|
||||
with timed_preprocessor_operation(self.info.ctx, "hashing"):
|
||||
mm_hashes = self._hash_mm_items(
|
||||
mm_data_items,
|
||||
hf_processor_mm_kwargs,
|
||||
tokenization_kwargs,
|
||||
mm_uuids=mm_uuids,
|
||||
mm_uuid_items,
|
||||
hf_processor_mm_kwargs=hf_processor_mm_kwargs,
|
||||
)
|
||||
|
||||
with timed_preprocessor_operation(self.info.ctx, "cache_lookup"):
|
||||
@@ -1753,10 +1744,9 @@ class BaseMultiModalProcessor(ABC, Generic[_I]):
|
||||
self,
|
||||
prompt: str | list[int],
|
||||
mm_items: MultiModalDataItems,
|
||||
hf_processor_mm_kwargs: Mapping[str, object],
|
||||
mm_uuid_items: MultiModalUUIDItems | None = None,
|
||||
hf_processor_mm_kwargs: Mapping[str, object] | None = None,
|
||||
tokenization_kwargs: Mapping[str, object] | None = None,
|
||||
*,
|
||||
mm_uuids: MultiModalUUIDDict | None = None,
|
||||
) -> MultiModalInputs:
|
||||
"""
|
||||
Process multi-modal inputs to be used in vLLM.
|
||||
@@ -1775,6 +1765,8 @@ class BaseMultiModalProcessor(ABC, Generic[_I]):
|
||||
if request_id is not None:
|
||||
self.info.ctx.create_timing_stats(request_id)
|
||||
|
||||
if hf_processor_mm_kwargs is None:
|
||||
hf_processor_mm_kwargs = {}
|
||||
if tokenization_kwargs is None:
|
||||
tokenization_kwargs = {}
|
||||
|
||||
@@ -1785,9 +1777,9 @@ class BaseMultiModalProcessor(ABC, Generic[_I]):
|
||||
) = self._cached_apply_hf_processor(
|
||||
prompt,
|
||||
mm_items,
|
||||
mm_uuid_items,
|
||||
hf_processor_mm_kwargs,
|
||||
tokenization_kwargs=tokenization_kwargs,
|
||||
mm_uuids=mm_uuids,
|
||||
)
|
||||
|
||||
# NOTE: tokenization_kwargs are not required to init processor
|
||||
@@ -1861,10 +1853,9 @@ class EncDecMultiModalProcessor(BaseMultiModalProcessor[_I]):
|
||||
self,
|
||||
prompt: str | list[int],
|
||||
mm_items: MultiModalDataItems,
|
||||
hf_processor_mm_kwargs: Mapping[str, object],
|
||||
mm_uuid_items: MultiModalUUIDItems | None = None,
|
||||
hf_processor_mm_kwargs: Mapping[str, object] | None = None,
|
||||
tokenization_kwargs: Mapping[str, object] | None = None,
|
||||
*,
|
||||
mm_uuids: MultiModalUUIDDict | None = None,
|
||||
) -> MultiModalEncDecInputs:
|
||||
"""
|
||||
Process multi-modal inputs to be used in vLLM.
|
||||
@@ -1877,9 +1868,9 @@ class EncDecMultiModalProcessor(BaseMultiModalProcessor[_I]):
|
||||
encoder_inputs = super().apply(
|
||||
encoder_prompt,
|
||||
mm_items,
|
||||
hf_processor_mm_kwargs,
|
||||
tokenization_kwargs,
|
||||
mm_uuids=mm_uuids,
|
||||
mm_uuid_items,
|
||||
hf_processor_mm_kwargs=hf_processor_mm_kwargs,
|
||||
tokenization_kwargs=tokenization_kwargs,
|
||||
)
|
||||
|
||||
return self._get_enc_dec_inputs(
|
||||
|
||||
@@ -17,7 +17,7 @@ if TYPE_CHECKING:
|
||||
from torch.distributed import PrefixStore, ProcessGroup
|
||||
|
||||
from vllm.config import VllmConfig
|
||||
from vllm.inputs import ProcessorInputs, PromptType
|
||||
from vllm.inputs import ProcessorInputs
|
||||
from vllm.pooling_params import PoolingParams
|
||||
from vllm.sampling_params import SamplingParams
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
@@ -568,9 +568,8 @@ class Platform:
|
||||
@classmethod
|
||||
def validate_request(
|
||||
cls,
|
||||
prompt: "PromptType | ProcessorInputs",
|
||||
params: "SamplingParams | PoolingParams",
|
||||
processed_inputs: "ProcessorInputs",
|
||||
params: "SamplingParams | PoolingParams",
|
||||
) -> None:
|
||||
"""Raises if this request is unsupported on this platform"""
|
||||
|
||||
|
||||
+37
-37
@@ -51,7 +51,7 @@ if TYPE_CHECKING:
|
||||
MultiModalInputs,
|
||||
MultiModalUUIDDict,
|
||||
)
|
||||
from vllm.multimodal.parse import MultiModalDataItems
|
||||
from vllm.multimodal.parse import MultiModalDataItems, MultiModalUUIDItems
|
||||
from vllm.multimodal.processing import BaseMultiModalProcessor
|
||||
|
||||
logger = init_logger(__name__)
|
||||
@@ -463,23 +463,25 @@ class BaseRenderer(ABC, Generic[_T]):
|
||||
def _validate_mm_uuids(
|
||||
self,
|
||||
mm_data: "MultiModalDataDict",
|
||||
mm_items: "MultiModalDataItems",
|
||||
mm_uuids: "MultiModalUUIDDict | None",
|
||||
mm_data_items: "MultiModalDataItems",
|
||||
mm_uuid_items: "MultiModalUUIDItems",
|
||||
) -> None:
|
||||
if mm_uuids is None:
|
||||
mm_uuids = {}
|
||||
|
||||
# NOTE: Keys corresponding to `None` in `mm_data` don't appear in `mm_items`
|
||||
modalities = mm_data.keys() | mm_uuids.keys()
|
||||
# NOTE: Keys corresponding to `None` in `mm_data` don't appear in
|
||||
# `mm_data_items`
|
||||
modalities = mm_data.keys() | mm_uuid_items.keys()
|
||||
|
||||
for modality in modalities:
|
||||
data_items = mm_items.get(modality) or list[Any]()
|
||||
data_items = mm_data_items.get(modality)
|
||||
uuid_items = mm_uuid_items.get(modality)
|
||||
|
||||
uuid_items = mm_uuids.get(modality) or list[str | None]()
|
||||
if isinstance(uuid_items, str):
|
||||
uuid_items = [uuid_items]
|
||||
if data_items is None:
|
||||
if uuid_items is None:
|
||||
raise ValueError(
|
||||
f"multi_modal_data[{modality!r}] is empty but "
|
||||
f"multi_modal_uuids[{modality!r}] is missing."
|
||||
)
|
||||
|
||||
if len(data_items) > 0:
|
||||
elif uuid_items is not None:
|
||||
if len(uuid_items) > 0 and len(data_items) != len(uuid_items):
|
||||
raise ValueError(
|
||||
f"If given, multi_modal_uuids[{modality!r}] must have "
|
||||
@@ -488,24 +490,17 @@ class BaseRenderer(ABC, Generic[_T]):
|
||||
)
|
||||
|
||||
for i, item in enumerate(data_items):
|
||||
if item is None:
|
||||
if not uuid_items:
|
||||
raise ValueError(
|
||||
f"multi_modal_data[{modality!r}][{i}] is empty but "
|
||||
f"multi_modal_uuids[{modality!r}] is missing."
|
||||
)
|
||||
|
||||
if uuid_items[i] is None:
|
||||
raise ValueError(
|
||||
f"multi_modal_data[{modality!r}][{i}] is empty but "
|
||||
f"multi_modal_uuids[{modality!r}][{i}] is missing."
|
||||
)
|
||||
if item is None and uuid_items[i] is None:
|
||||
raise ValueError(
|
||||
f"multi_modal_data[{modality!r}][{i}] is empty but "
|
||||
f"multi_modal_uuids[{modality!r}][{i}] is missing."
|
||||
)
|
||||
|
||||
def _process_mm_uuids(
|
||||
self,
|
||||
mm_data: "MultiModalDataDict",
|
||||
mm_items: "MultiModalDataItems",
|
||||
mm_uuids: "MultiModalUUIDDict | None",
|
||||
mm_data_items: "MultiModalDataItems",
|
||||
mm_uuid_items: "MultiModalUUIDItems",
|
||||
mm_req_id: str,
|
||||
):
|
||||
model_config = self.model_config
|
||||
@@ -520,40 +515,45 @@ class BaseRenderer(ABC, Generic[_T]):
|
||||
and model_config.multimodal_config.mm_processor_cache_gb == 0
|
||||
and not self.config.cache_config.enable_prefix_caching
|
||||
):
|
||||
mm_uuids = {
|
||||
mm_uuid_items = {
|
||||
modality: [f"{mm_req_id}-{modality}-{i}" for i in range(data_count)]
|
||||
for modality, data_count in mm_items.get_all_counts().items()
|
||||
for modality, data_count in mm_data_items.get_all_counts().items()
|
||||
}
|
||||
|
||||
self._validate_mm_uuids(mm_data, mm_items, mm_uuids)
|
||||
self._validate_mm_uuids(mm_data, mm_data_items, mm_uuid_items)
|
||||
|
||||
return mm_uuids
|
||||
return mm_uuid_items
|
||||
|
||||
# TODO: Remove str and tokenization_kwargs after deprecating InputPreprocessor
|
||||
def _process_multimodal(
|
||||
self,
|
||||
prompt: list[int] | str,
|
||||
mm_data: "MultiModalDataDict",
|
||||
mm_uuids: "MultiModalUUIDDict | None",
|
||||
mm_processor_kwargs: Mapping[str, object] | None,
|
||||
tokenization_kwargs: dict[str, Any] | None,
|
||||
mm_uuids: "MultiModalUUIDDict | None",
|
||||
) -> "MultiModalInputs":
|
||||
from vllm.multimodal.parse import parse_mm_uuids
|
||||
from vllm.multimodal.processing.context import set_request_id
|
||||
|
||||
mm_req_id = f"renderer-mm-{self._mm_req_counter.inc(1)}"
|
||||
|
||||
mm_processor = self.get_mm_processor()
|
||||
|
||||
mm_items = mm_processor.info.parse_mm_data(mm_data)
|
||||
mm_uuids = self._process_mm_uuids(mm_data, mm_items, mm_uuids, mm_req_id)
|
||||
mm_data_items = mm_processor.info.parse_mm_data(mm_data)
|
||||
mm_uuid_items = parse_mm_uuids(mm_uuids)
|
||||
|
||||
mm_uuids = self._process_mm_uuids(
|
||||
mm_data, mm_data_items, mm_uuid_items, mm_req_id
|
||||
)
|
||||
|
||||
with set_request_id(mm_req_id), set_default_torch_num_threads():
|
||||
mm_inputs = mm_processor.apply(
|
||||
prompt,
|
||||
mm_items,
|
||||
hf_processor_mm_kwargs=mm_processor_kwargs or {},
|
||||
mm_data_items,
|
||||
mm_uuid_items,
|
||||
hf_processor_mm_kwargs=mm_processor_kwargs,
|
||||
tokenization_kwargs=tokenization_kwargs,
|
||||
mm_uuids=mm_uuids,
|
||||
)
|
||||
|
||||
self.update_mm_cache_stats()
|
||||
|
||||
@@ -349,7 +349,7 @@ def _align(x: int, y: int) -> int:
|
||||
|
||||
|
||||
# Taken from https://github.com/deepseek-ai/DeepGEMM/blob/v2.1.1/csrc/utils/math.hpp#L19
|
||||
def get_tma_aligned_size(x: int, element_size: int):
|
||||
def get_tma_aligned_size(x: int, element_size: int) -> int:
|
||||
return _align(x, 16 // element_size)
|
||||
|
||||
|
||||
|
||||
@@ -27,7 +27,7 @@ from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry
|
||||
from vllm.outputs import STREAM_FINISHED, PoolingRequestOutput, RequestOutput
|
||||
from vllm.plugins.io_processors import get_io_processor
|
||||
from vllm.pooling_params import PoolingParams
|
||||
from vllm.renderers import merge_kwargs, renderer_from_config
|
||||
from vllm.renderers import renderer_from_config
|
||||
from vllm.renderers.inputs.preprocess import extract_prompt_components
|
||||
from vllm.sampling_params import RequestOutputKind, SamplingParams
|
||||
from vllm.tasks import SupportedTask
|
||||
@@ -319,21 +319,6 @@ class AsyncLLM(EngineClient):
|
||||
"prompt logprobs"
|
||||
)
|
||||
|
||||
if params.truncate_prompt_tokens is not None:
|
||||
params_type = type(params).__name__
|
||||
warnings.warn(
|
||||
f"The `truncate_prompt_tokens` parameter in `{params_type}` "
|
||||
"is deprecated and will be removed in v0.16. "
|
||||
"Please pass it via `tokenization_kwargs` instead.",
|
||||
DeprecationWarning,
|
||||
stacklevel=2,
|
||||
)
|
||||
|
||||
tokenization_kwargs = merge_kwargs(
|
||||
tokenization_kwargs,
|
||||
dict(truncate_prompt_tokens=params.truncate_prompt_tokens),
|
||||
)
|
||||
|
||||
if isinstance(prompt, AsyncGenerator):
|
||||
if reasoning_ended is not None:
|
||||
raise NotImplementedError
|
||||
@@ -353,6 +338,12 @@ class AsyncLLM(EngineClient):
|
||||
|
||||
# Convert Input --> Request.
|
||||
if isinstance(prompt, EngineCoreRequest):
|
||||
logger.warning_once(
|
||||
"Passing EngineCoreRequest to AsyncLLM.generate() and .add_requests() "
|
||||
"is deprecated and will be removed in v0.18. You should instead pass "
|
||||
"the outputs of Renderer.render_cmpl() or Renderer.render_chat()."
|
||||
)
|
||||
|
||||
request = prompt
|
||||
if request_id != request.request_id:
|
||||
logger.warning_once(
|
||||
|
||||
@@ -2,6 +2,7 @@
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import time
|
||||
import warnings
|
||||
from collections.abc import Mapping
|
||||
from typing import Any, Literal
|
||||
|
||||
@@ -28,6 +29,7 @@ from vllm.sampling_params import SamplingParams
|
||||
from vllm.tasks import POOLING_TASKS, SupportedTask
|
||||
from vllm.tokenizers import TokenizerLike
|
||||
from vllm.utils import length_from_prompt_token_ids_or_embeds, random_uuid
|
||||
from vllm.utils.func_utils import supports_kw
|
||||
from vllm.utils.jsontree import json_iter_leaves
|
||||
from vllm.v1.engine import EngineCoreRequest
|
||||
|
||||
@@ -72,6 +74,33 @@ class InputProcessor:
|
||||
mm_registry=mm_registry,
|
||||
)
|
||||
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
platform_validate_request = current_platform.validate_request
|
||||
if supports_kw(platform_validate_request, "prompt"):
|
||||
logger.warning_once(
|
||||
"The signature of Platform.validate_request has changed from "
|
||||
"`(cls, prompt, params, processed_inputs) -> None` to "
|
||||
"`(cls, processed_inputs, params) -> None`. The old signature "
|
||||
"will no longer be supported starting from v0.18."
|
||||
)
|
||||
|
||||
orig_validate_request = platform_validate_request
|
||||
|
||||
def compat_validate_request(
|
||||
processed_inputs: ProcessorInputs,
|
||||
params: SamplingParams | PoolingParams,
|
||||
):
|
||||
return orig_validate_request(
|
||||
processed_inputs,
|
||||
params,
|
||||
processed_inputs, # type: ignore
|
||||
) # type: ignore
|
||||
|
||||
platform_validate_request = compat_validate_request
|
||||
|
||||
self._platform_validate_request = platform_validate_request
|
||||
|
||||
@property
|
||||
def tokenizer(self) -> TokenizerLike | None:
|
||||
return self.renderer.tokenizer
|
||||
@@ -87,6 +116,16 @@ class InputProcessor:
|
||||
supported_tasks: tuple[SupportedTask, ...] | None,
|
||||
):
|
||||
"""Raise `ValueError` if SamplingParams or PoolingParams is not valid."""
|
||||
if params.truncate_prompt_tokens is not None:
|
||||
params_type = type(params).__name__
|
||||
warnings.warn(
|
||||
f"The `truncate_prompt_tokens` parameter in `{params_type}` "
|
||||
"is deprecated and will be removed in v0.17. "
|
||||
"Please pass it via `tokenization_kwargs` instead.",
|
||||
DeprecationWarning,
|
||||
stacklevel=2,
|
||||
)
|
||||
|
||||
if isinstance(params, SamplingParams):
|
||||
params.verify(
|
||||
self.model_config,
|
||||
@@ -211,11 +250,24 @@ class InputProcessor:
|
||||
)
|
||||
|
||||
if isinstance(prompt, dict) and "type" in prompt:
|
||||
if tokenization_kwargs:
|
||||
logger.warning_once(
|
||||
"Passing tokenization_kwargs to InputProcessor is deprecated "
|
||||
"and will be removed in v0.18. You should instead pass "
|
||||
"them to Renderer.render_cmpl() or Renderer.render_chat()."
|
||||
)
|
||||
|
||||
if arrival_time is None:
|
||||
arrival_time = prompt.get("arrival_time", time.time()) # type: ignore[assignment]
|
||||
|
||||
processed_inputs: ProcessorInputs = prompt # type: ignore[assignment]
|
||||
else:
|
||||
logger.warning_once(
|
||||
"Passing raw prompts to InputProcessor is deprecated "
|
||||
"and will be removed in v0.18. You should instead pass "
|
||||
"the outputs of Renderer.render_cmpl() or Renderer.render_chat()."
|
||||
)
|
||||
|
||||
if arrival_time is None:
|
||||
arrival_time = time.time()
|
||||
|
||||
@@ -224,13 +276,7 @@ class InputProcessor:
|
||||
tokenization_kwargs=tokenization_kwargs,
|
||||
)
|
||||
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
current_platform.validate_request(
|
||||
prompt=prompt,
|
||||
params=params,
|
||||
processed_inputs=processed_inputs,
|
||||
)
|
||||
self._platform_validate_request(processed_inputs, params)
|
||||
|
||||
encoder_inputs, decoder_inputs = split_enc_dec_inputs(processed_inputs)
|
||||
self._validate_model_inputs(encoder_inputs, decoder_inputs)
|
||||
|
||||
@@ -234,10 +234,16 @@ class LLMEngine:
|
||||
|
||||
# Process raw inputs into the request.
|
||||
if isinstance(prompt, EngineCoreRequest):
|
||||
logger.warning_once(
|
||||
"Passing EngineCoreRequest to LLMEngine.generate() and .add_requests() "
|
||||
"is deprecated and will be removed in v0.18. You should instead pass "
|
||||
"the outputs of Renderer.render_cmpl() or Renderer.render_chat()."
|
||||
)
|
||||
|
||||
request = prompt
|
||||
if request_id != request.request_id:
|
||||
logger.warning_once(
|
||||
"AsyncLLM.add_request() was passed a request_id parameter that "
|
||||
"LLMEngine.add_request() was passed a request_id parameter that "
|
||||
"does not match the EngineCoreRequest.request_id attribute. The "
|
||||
"latter will be used, and the former will be ignored."
|
||||
)
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user