forked from Karylab-cklius/vllm
Compare commits
6
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
0fc695fc6d | ||
|
|
91df0fad4d | ||
|
|
78743ab5bf | ||
|
|
b2d7294b0f | ||
|
|
741ba421d8 | ||
|
|
ac94893da3 |
@@ -218,6 +218,10 @@ COPY requirements/common.txt requirements/common.txt
|
|||||||
COPY requirements/cuda.txt requirements/cuda.txt
|
COPY requirements/cuda.txt requirements/cuda.txt
|
||||||
COPY use_existing_torch.py use_existing_torch.py
|
COPY use_existing_torch.py use_existing_torch.py
|
||||||
COPY pyproject.toml pyproject.toml
|
COPY pyproject.toml pyproject.toml
|
||||||
|
# nvidia-cutlass-dsl[cu13] installs -libs-base and -libs-cu13 wheels that
|
||||||
|
# share paths with different content. uv can extract them in either order,
|
||||||
|
# leaving base files that break CUDA 13 CuTe DSL JIT.
|
||||||
|
# TODO(mmangkad): Remove this after NVIDIA/cutlass#3259 is fixed.
|
||||||
RUN --mount=type=cache,target=/opt/uv/cache \
|
RUN --mount=type=cache,target=/opt/uv/cache \
|
||||||
if [ "$(echo $CUDA_VERSION | cut -d. -f1)" = "12" ]; then \
|
if [ "$(echo $CUDA_VERSION | cut -d. -f1)" = "12" ]; then \
|
||||||
sed -i 's/^nvidia-cutlass-dsl\[cu13\]/nvidia-cutlass-dsl/' requirements/cuda.txt; \
|
sed -i 's/^nvidia-cutlass-dsl\[cu13\]/nvidia-cutlass-dsl/' requirements/cuda.txt; \
|
||||||
@@ -234,6 +238,13 @@ RUN --mount=type=cache,target=/opt/uv/cache \
|
|||||||
else \
|
else \
|
||||||
uv pip install --python /opt/venv/bin/python3 -r requirements/cuda.txt \
|
uv pip install --python /opt/venv/bin/python3 -r requirements/cuda.txt \
|
||||||
--extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \
|
--extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \
|
||||||
|
fi \
|
||||||
|
&& if [ "$(echo $CUDA_VERSION | cut -d. -f1)" = "13" ]; then \
|
||||||
|
CUTLASS_DSL_VERSION=$(uv pip show --python /opt/venv/bin/python3 nvidia-cutlass-dsl 2>/dev/null | awk '/^Version:/{print $2}') && \
|
||||||
|
if [ -n "$CUTLASS_DSL_VERSION" ]; then \
|
||||||
|
uv pip install --python /opt/venv/bin/python3 --force-reinstall --no-deps \
|
||||||
|
"nvidia-cutlass-dsl-libs-cu13==${CUTLASS_DSL_VERSION}"; \
|
||||||
|
fi; \
|
||||||
fi
|
fi
|
||||||
|
|
||||||
# Track PyTorch lib versions used during build and match in downstream instances.
|
# Track PyTorch lib versions used during build and match in downstream instances.
|
||||||
@@ -745,6 +756,10 @@ ENV VLLM_ENABLE_CUDA_COMPATIBILITY=0
|
|||||||
ARG PYTORCH_CUDA_INDEX_BASE_URL
|
ARG PYTORCH_CUDA_INDEX_BASE_URL
|
||||||
COPY requirements/common.txt /tmp/common.txt
|
COPY requirements/common.txt /tmp/common.txt
|
||||||
COPY requirements/cuda.txt /tmp/requirements-cuda.txt
|
COPY requirements/cuda.txt /tmp/requirements-cuda.txt
|
||||||
|
# nvidia-cutlass-dsl[cu13] installs -libs-base and -libs-cu13 wheels that
|
||||||
|
# share paths with different content. uv can extract them in either order,
|
||||||
|
# leaving base files that break CUDA 13 CuTe DSL JIT.
|
||||||
|
# TODO(mmangkad): Remove this after NVIDIA/cutlass#3259 is fixed.
|
||||||
RUN --mount=type=cache,target=/opt/uv/cache \
|
RUN --mount=type=cache,target=/opt/uv/cache \
|
||||||
if [ "$(echo $CUDA_VERSION | cut -d. -f1)" = "12" ]; then \
|
if [ "$(echo $CUDA_VERSION | cut -d. -f1)" = "12" ]; then \
|
||||||
sed -i 's/^nvidia-cutlass-dsl\[cu13\]/nvidia-cutlass-dsl/' /tmp/requirements-cuda.txt; \
|
sed -i 's/^nvidia-cutlass-dsl\[cu13\]/nvidia-cutlass-dsl/' /tmp/requirements-cuda.txt; \
|
||||||
@@ -752,6 +767,13 @@ RUN --mount=type=cache,target=/opt/uv/cache \
|
|||||||
fi && \
|
fi && \
|
||||||
uv pip install --system -r /tmp/requirements-cuda.txt \
|
uv pip install --system -r /tmp/requirements-cuda.txt \
|
||||||
--extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') && \
|
--extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') && \
|
||||||
|
if [ "$(echo $CUDA_VERSION | cut -d. -f1)" = "13" ]; then \
|
||||||
|
CUTLASS_DSL_VERSION=$(uv pip show --system nvidia-cutlass-dsl 2>/dev/null | awk '/^Version:/{print $2}') && \
|
||||||
|
if [ -n "$CUTLASS_DSL_VERSION" ]; then \
|
||||||
|
uv pip install --system --force-reinstall --no-deps \
|
||||||
|
"nvidia-cutlass-dsl-libs-cu13==${CUTLASS_DSL_VERSION}"; \
|
||||||
|
fi; \
|
||||||
|
fi && \
|
||||||
rm /tmp/requirements-cuda.txt /tmp/common.txt
|
rm /tmp/requirements-cuda.txt /tmp/common.txt
|
||||||
|
|
||||||
# Install FlashInfer JIT cache (requires CUDA-version-specific index URL)
|
# Install FlashInfer JIT cache (requires CUDA-version-specific index URL)
|
||||||
@@ -842,6 +864,19 @@ RUN --mount=type=bind,from=build,src=/tmp/ep_kernels_workspace/dist,target=/vllm
|
|||||||
uv pip install --system ep_kernels/dist/*.whl --verbose \
|
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 '.')
|
--extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
|
||||||
|
|
||||||
|
# nvidia-cutlass-dsl[cu13] installs -libs-base and -libs-cu13 wheels that
|
||||||
|
# share paths with different content. Force -libs-cu13 last after runtime
|
||||||
|
# dependency installs so uv cannot leave base files behind.
|
||||||
|
# TODO(mmangkad): Remove this after NVIDIA/cutlass#3259 is fixed.
|
||||||
|
RUN --mount=type=cache,target=/opt/uv/cache \
|
||||||
|
if [ "$(echo $CUDA_VERSION | cut -d. -f1)" = "13" ]; then \
|
||||||
|
CUTLASS_DSL_VERSION=$(uv pip show --system nvidia-cutlass-dsl 2>/dev/null | awk '/^Version:/{print $2}') && \
|
||||||
|
if [ -n "$CUTLASS_DSL_VERSION" ]; then \
|
||||||
|
uv pip install --system --force-reinstall --no-deps \
|
||||||
|
"nvidia-cutlass-dsl-libs-cu13==${CUTLASS_DSL_VERSION}"; \
|
||||||
|
fi; \
|
||||||
|
fi
|
||||||
|
|
||||||
# Download FlashInfer precompiled cubins AFTER all pip installs are done.
|
# Download FlashInfer precompiled cubins AFTER all pip installs are done.
|
||||||
# This must run after the vLLM wheel and EP kernels installs above, because
|
# This must run after the vLLM wheel and EP kernels installs above, because
|
||||||
# those can reinstall/touch flashinfer packages. Downloading cubins earlier
|
# those can reinstall/touch flashinfer packages. Downloading cubins earlier
|
||||||
|
|||||||
@@ -168,6 +168,12 @@ RUN --mount=type=cache,target=/root/.cache/uv \
|
|||||||
######################### TRITON-CPU BUILD IMAGE #########################
|
######################### TRITON-CPU BUILD IMAGE #########################
|
||||||
FROM base AS vllm-triton-cpu-build
|
FROM base AS vllm-triton-cpu-build
|
||||||
|
|
||||||
|
# Support for cross-compilation with x86 ISA including AVX2 and AVX512: docker build --build-arg VLLM_CPU_X86="true" ...
|
||||||
|
# Re-declared here because this stage is `FROM base` (not `vllm-build`), so it
|
||||||
|
# does not inherit the ARG/ENV defined there. Without it, the guard below would
|
||||||
|
# see an empty value and build triton-cpu on non-x86 targets (e.g. arm64).
|
||||||
|
ARG VLLM_CPU_X86=0
|
||||||
|
|
||||||
WORKDIR /vllm-workspace
|
WORKDIR /vllm-workspace
|
||||||
|
|
||||||
RUN mkdir dist
|
RUN mkdir dist
|
||||||
@@ -269,6 +275,11 @@ ENV HF_HUB_DOWNLOAD_TIMEOUT 60
|
|||||||
######################### RELEASE IMAGE #########################
|
######################### RELEASE IMAGE #########################
|
||||||
FROM base AS vllm-openai
|
FROM base AS vllm-openai
|
||||||
|
|
||||||
|
# Re-declared here because this stage is `FROM base` (not `vllm-build`), so the
|
||||||
|
# RUN below that gates the triton-cpu wheel install on $VLLM_CPU_X86 would
|
||||||
|
# otherwise see an empty value and try to install it on non-x86 targets.
|
||||||
|
ARG VLLM_CPU_X86=0
|
||||||
|
|
||||||
WORKDIR /vllm-workspace
|
WORKDIR /vllm-workspace
|
||||||
|
|
||||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||||
|
|||||||
@@ -11,7 +11,7 @@ transformers >= 4.56.0, != 5.0.*, != 5.1.*, != 5.2.*, != 5.3.*, != 5.4.*, != 5.5
|
|||||||
tokenizers >= 0.21.1 # Required for fast incremental detokenization.
|
tokenizers >= 0.21.1 # Required for fast incremental detokenization.
|
||||||
safetensors >= 0.6.2 # MXFP4/MXFP6 dtype support (F8_E8M0, F4) added in 0.6.0: https://github.com/huggingface/safetensors/pull/611
|
safetensors >= 0.6.2 # MXFP4/MXFP6 dtype support (F8_E8M0, F4) added in 0.6.0: https://github.com/huggingface/safetensors/pull/611
|
||||||
protobuf >= 5.29.6, !=6.30.*, !=6.31.*, !=6.32.*, !=6.33.0.*, !=6.33.1.*, !=6.33.2.*, !=6.33.3.*, !=6.33.4.* # Required by LlamaTokenizer, gRPC. CVE-2026-0994
|
protobuf >= 5.29.6, !=6.30.*, !=6.31.*, !=6.32.*, !=6.33.0.*, !=6.33.1.*, !=6.33.2.*, !=6.33.3.*, !=6.33.4.* # Required by LlamaTokenizer, gRPC. CVE-2026-0994
|
||||||
fastapi[standard] >= 0.115.0 # Required by FastAPI's form models in the OpenAI API server's audio transcriptions endpoint.
|
fastapi[standard] >= 0.115.0, < 0.137 # Required by FastAPI's form models in the OpenAI API server's audio transcriptions endpoint; <0.137 due to prometheus-fastapi-instrumentator#370.
|
||||||
aiohttp >= 3.13.3
|
aiohttp >= 3.13.3
|
||||||
openai >= 2.0.0 # For Responses API with reasoning content
|
openai >= 2.0.0 # For Responses API with reasoning content
|
||||||
pydantic >= 2.12.0
|
pydantic >= 2.12.0
|
||||||
|
|||||||
@@ -23,6 +23,7 @@ timm>=1.0.17
|
|||||||
# To be consistent with test_quark.py
|
# To be consistent with test_quark.py
|
||||||
amd-quark>=0.8.99
|
amd-quark>=0.8.99
|
||||||
tilelang==0.1.10
|
tilelang==0.1.10
|
||||||
|
# Required apache-tvm-ffi matching tilelang version
|
||||||
|
apache-tvm-ffi==0.1.10
|
||||||
# Required for faster safetensors model loading
|
# Required for faster safetensors model loading
|
||||||
fastsafetensors >= 0.3.2
|
fastsafetensors >= 0.3.2
|
||||||
|
|||||||
@@ -44,6 +44,7 @@ anyio==4.13.0
|
|||||||
# watchfiles
|
# watchfiles
|
||||||
apache-tvm-ffi==0.1.10
|
apache-tvm-ffi==0.1.10
|
||||||
# via
|
# via
|
||||||
|
# -c requirements/rocm.txt
|
||||||
# tilelang
|
# tilelang
|
||||||
# xgrammar
|
# xgrammar
|
||||||
arctic-inference==0.1.1
|
arctic-inference==0.1.1
|
||||||
|
|||||||
@@ -340,6 +340,16 @@ def rocm_aiter_fused_experts(
|
|||||||
moe_config.intermediate_size_per_partition
|
moe_config.intermediate_size_per_partition
|
||||||
- moe_config.intermediate_size_per_partition_unpadded
|
- moe_config.intermediate_size_per_partition_unpadded
|
||||||
)
|
)
|
||||||
|
# Round hidden_pad/intermediate_pad to match AITER's CK/FlyDSL MoE
|
||||||
|
# dispatch (currently pinned to v0.1.13.post1):
|
||||||
|
# https://github.com/ROCm/aiter/blob/v0.1.13.post1/aiter/fused_moe.py#L1073
|
||||||
|
# https://github.com/ROCm/aiter/blob/v0.1.13.post1/aiter/fused_moe.py#L1099
|
||||||
|
# TODO: Revisit this once we bump AITER to 0.1.15 with padding fixes
|
||||||
|
# for CK/FlyDSL MoE GEMM e.g. https://github.com/ROCm/aiter/pull/3401
|
||||||
|
hidden_pad = hidden_pad // 128 * 128
|
||||||
|
intermediate_pad = (
|
||||||
|
intermediate_pad // 64 * 64 * (2 if moe_config.tp_size == 1 else 1)
|
||||||
|
)
|
||||||
|
|
||||||
return rocm_aiter_ops.fused_moe(
|
return rocm_aiter_ops.fused_moe(
|
||||||
hidden_states,
|
hidden_states,
|
||||||
@@ -357,8 +367,8 @@ def rocm_aiter_fused_experts(
|
|||||||
doweight_stage1=apply_router_weight_on_input,
|
doweight_stage1=apply_router_weight_on_input,
|
||||||
num_local_tokens=num_local_tokens,
|
num_local_tokens=num_local_tokens,
|
||||||
output_dtype=output_dtype,
|
output_dtype=output_dtype,
|
||||||
hidden_pad=hidden_pad // 128 * 128,
|
hidden_pad=hidden_pad,
|
||||||
intermediate_pad=intermediate_pad // 64 * 64 * 2,
|
intermediate_pad=intermediate_pad,
|
||||||
bias1=quant_config.w1_bias if quant_config.use_mxfp4_w4a16 else None,
|
bias1=quant_config.w1_bias if quant_config.use_mxfp4_w4a16 else None,
|
||||||
bias2=quant_config.w2_bias if quant_config.use_mxfp4_w4a16 else None,
|
bias2=quant_config.w2_bias if quant_config.use_mxfp4_w4a16 else None,
|
||||||
moe_sorting_dispatch_policy=moe_sorting_dispatch_policy,
|
moe_sorting_dispatch_policy=moe_sorting_dispatch_policy,
|
||||||
|
|||||||
@@ -242,7 +242,12 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
|
|||||||
vllm_config.model_config.max_model_len,
|
vllm_config.model_config.max_model_len,
|
||||||
vllm_config.scheduler_config.max_num_batched_tokens,
|
vllm_config.scheduler_config.max_num_batched_tokens,
|
||||||
)
|
)
|
||||||
self._init_fp8_prefill_ps_buffers(max_num_reqs, max_prefill_qlen, device)
|
self._init_fp8_prefill_ps_buffers(
|
||||||
|
max_num_reqs,
|
||||||
|
max_prefill_qlen,
|
||||||
|
vllm_config.scheduler_config.max_num_batched_tokens,
|
||||||
|
device,
|
||||||
|
)
|
||||||
|
|
||||||
if self.compilation_config.cudagraph_mode.has_full_cudagraphs():
|
if self.compilation_config.cudagraph_mode.has_full_cudagraphs():
|
||||||
self.paged_kv_indptr = torch.zeros(
|
self.paged_kv_indptr = torch.zeros(
|
||||||
@@ -257,21 +262,29 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
|
|||||||
self,
|
self,
|
||||||
max_num_reqs: int,
|
max_num_reqs: int,
|
||||||
max_prefill_qlen: int,
|
max_prefill_qlen: int,
|
||||||
|
max_num_batched_tokens: int,
|
||||||
device: torch.device,
|
device: torch.device,
|
||||||
) -> None:
|
) -> None:
|
||||||
"""Pre-allocate persistent buffers for FP8 MLA prefill PS metadata.
|
"""Pre-allocate persistent buffers for FP8 MLA prefill PS metadata.
|
||||||
|
|
||||||
Uses ``get_ps_metadata_info_v1`` with max values so the buffers are
|
Uses ``get_ps_metadata_info_v1`` with max values so the buffers are
|
||||||
large enough for any batch. ``get_ps_metadata_v1`` fills them
|
large enough for any batch. ``get_ps_metadata_v1`` fills them
|
||||||
per-batch in ``build()``.
|
per-batch in ``build()``. The FP8 prefill forward path also uses the
|
||||||
|
global workspace manager for per-call scratch, so reserve its maximum
|
||||||
|
shape here before the workspace manager is locked after warmup.
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
max_num_reqs: Maximum number of concurrent requests.
|
max_num_reqs: Maximum number of concurrent requests.
|
||||||
max_prefill_qlen: Maximum Q-length for a single request in one
|
max_prefill_qlen: Maximum Q-length for a single request in one
|
||||||
prefill batch. Should be ``min(max_model_len,
|
prefill batch. Should be ``min(max_model_len,
|
||||||
max_num_batched_tokens)`` — the chunked-prefill scheduler
|
max_num_batched_tokens)`` — a single request never exceeds
|
||||||
never emits more than ``max_num_batched_tokens`` new tokens
|
``max_model_len`` tokens, nor the per-batch token budget.
|
||||||
per batch.
|
max_num_batched_tokens: Maximum number of tokens scheduled in one
|
||||||
|
batch. The ``final_lse`` scratch is sized by ``total_q`` (the
|
||||||
|
summed Q-length over all prefill requests in the batch), which
|
||||||
|
is bounded by this budget rather than by a single request's
|
||||||
|
``max_prefill_qlen`` — concurrent requests can sum to more than
|
||||||
|
``max_model_len`` when ``max_model_len < max_num_batched_tokens``.
|
||||||
device: Target device for the buffers.
|
device: Target device for the buffers.
|
||||||
"""
|
"""
|
||||||
from aiter import get_ps_metadata_info_v1
|
from aiter import get_ps_metadata_info_v1
|
||||||
@@ -279,6 +292,7 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
|
|||||||
# After kv_b_proj decompression, K has num_heads heads (same as Q).
|
# After kv_b_proj decompression, K has num_heads heads (same as Q).
|
||||||
# So gqa_ratio=1 and num_head_k=num_heads for the PS kernel.
|
# So gqa_ratio=1 and num_head_k=num_heads for the PS kernel.
|
||||||
num_head_k = self.num_heads
|
num_head_k = self.num_heads
|
||||||
|
v_head_dim = self.mla_dims.v_head_dim
|
||||||
# gqa_ratio = 1
|
# gqa_ratio = 1
|
||||||
# qlen_granularity = _FP8_PREFILL_TILE_Q // max(gqa_ratio, 1)
|
# qlen_granularity = _FP8_PREFILL_TILE_Q // max(gqa_ratio, 1)
|
||||||
qlen_granularity = _FP8_PREFILL_TILE_Q
|
qlen_granularity = _FP8_PREFILL_TILE_Q
|
||||||
@@ -318,6 +332,21 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
|
|||||||
device=device,
|
device=device,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
from vllm.v1.worker.workspace import current_workspace_manager
|
||||||
|
|
||||||
|
max_num_partial_tiles = reduce_partial_map_size
|
||||||
|
current_workspace_manager().get_simultaneous(
|
||||||
|
(
|
||||||
|
(max_num_partial_tiles * _FP8_PREFILL_TILE_Q, num_head_k, v_head_dim),
|
||||||
|
torch.float32,
|
||||||
|
),
|
||||||
|
(
|
||||||
|
(max_num_partial_tiles * _FP8_PREFILL_TILE_Q, num_head_k),
|
||||||
|
torch.float32,
|
||||||
|
),
|
||||||
|
((max_num_batched_tokens, num_head_k), torch.float32),
|
||||||
|
)
|
||||||
|
|
||||||
logger.info(
|
logger.info(
|
||||||
"FP8 MLA prefill PS buffers allocated "
|
"FP8 MLA prefill PS buffers allocated "
|
||||||
"(max_batch=%d, max_qlen=%d, num_head_k=%d)",
|
"(max_batch=%d, max_qlen=%d, num_head_k=%d)",
|
||||||
|
|||||||
Reference in New Issue
Block a user