diff --git a/.buildkite/hardware_tests/amd.yaml b/.buildkite/hardware_tests/amd.yaml index d47a7e394a46..bd73c8102c7f 100644 --- a/.buildkite/hardware_tests/amd.yaml +++ b/.buildkite/hardware_tests/amd.yaml @@ -18,6 +18,8 @@ steps: TERM: "xterm-256color" retry: automatic: + - exit_status: 1 # Transient Docker/BuildKit failure + limit: 1 - exit_status: -1 # Agent was lost limit: 1 - exit_status: -10 # Agent was lost @@ -46,6 +48,8 @@ steps: VLLM_BRANCH: "$BUILDKITE_COMMIT" retry: automatic: + - exit_status: 1 # Transient Docker/BuildKit failure + limit: 1 - exit_status: -1 # Agent was lost limit: 1 - exit_status: -10 # Agent was lost @@ -72,6 +76,8 @@ steps: VLLM_BRANCH: "$BUILDKITE_COMMIT" retry: automatic: + - exit_status: 1 # Transient Docker/BuildKit failure + limit: 1 - exit_status: -1 # Agent was lost limit: 1 - exit_status: -10 # Agent was lost diff --git a/.buildkite/scripts/ci-bake-rocm.sh b/.buildkite/scripts/ci-bake-rocm.sh index 2891291105f7..3b9a9a1d7136 100644 --- a/.buildkite/scripts/ci-bake-rocm.sh +++ b/.buildkite/scripts/ci-bake-rocm.sh @@ -285,7 +285,7 @@ get_content_arg_names() { fi | awk 'NF && !seen[$0]++' } -compute_ci_base_content_hash() { +compute_ci_base_content_hash_once() { local -a content_paths=() local -a content_args=() local dockerfile="${CI_BASE_DOCKERFILE:-}" @@ -301,7 +301,8 @@ compute_ci_base_content_hash() { if [[ -n "${dockerfile}" ]]; then printf 'dockerfile:%s\n' "${dockerfile}" printf 'resolved-build-args:\n' - hash_dockerfile_arg_values "${dockerfile}" "${content_args[@]}" + hash_dockerfile_arg_values "${dockerfile}" "${content_args[@]}" \ + || return 1 if [[ -n "${stages}" ]]; then printf 'dockerfile-stages:%s\n' "${stages}" if [[ -f "${dockerfile}" ]]; then @@ -314,6 +315,53 @@ compute_ci_base_content_hash() { } | sha256sum | cut -d' ' -f1 } +compute_ci_base_content_hash() { + local attempts="${CI_BASE_HASH_ATTEMPTS:-3}" + local delay_secs="${CI_BASE_HASH_RETRY_DELAY:-5}" + local attempt=0 + local hash="" + local failed=0 + local -a hashes=() + + if [[ ! "${attempts}" =~ ^[1-9][0-9]*$ ]]; then + echo "Invalid CI_BASE_HASH_ATTEMPTS: ${attempts}" >&2 + return 1 + fi + if [[ ! "${delay_secs}" =~ ^[0-9]+$ ]]; then + echo "Invalid CI_BASE_HASH_RETRY_DELAY: ${delay_secs}" >&2 + return 1 + fi + + for ((attempt = 1; attempt <= attempts; attempt++)); do + if ! hash=$(compute_ci_base_content_hash_once); then + echo "ci_base content hash calculation ${attempt}/${attempts} failed" >&2 + failed=1 + else + hashes+=("${hash}") + echo "ci_base content hash calculation ${attempt}/${attempts}: ${hash}" >&2 + fi + + if ((attempt < attempts)); then + sleep "${delay_secs}" + fi + done + + if ((failed)) || ((${#hashes[@]} != attempts)); then + echo "Could not calculate a reliable ci_base content hash" >&2 + return 1 + fi + + for hash in "${hashes[@]:1}"; do + if [[ "${hash}" != "${hashes[0]}" ]]; then + echo "ci_base content hash changed between calculations" >&2 + printf ' observed: %s\n' "${hashes[@]}" >&2 + return 1 + fi + done + + printf '%s\n' "${hashes[0]}" +} + extract_dockerfile_arg_default() { local dockerfile="$1" local arg_name="$2" @@ -366,7 +414,11 @@ hash_dockerfile_arg_values() { printf 'arg:%s=%s\n' "${arg_name}" "${arg_value:-}" if [[ "${arg_name}" == "BASE_IMAGE" && -n "${arg_value}" ]]; then digest=$(resolve_image_digest "${arg_value}") - printf 'arg:%s.digest=%s\n' "${arg_name}" "${digest:-unknown}" + if [[ -z "${digest}" ]]; then + echo "Failed to resolve digest for BASE_IMAGE=${arg_value}" >&2 + return 1 + fi + printf 'arg:%s.digest=%s\n' "${arg_name}" "${digest}" fi done } diff --git a/.buildkite/scripts/rocm/build-test-image.sh b/.buildkite/scripts/rocm/build-test-image.sh index 9803e20d02e8..9feb6c092402 100755 --- a/.buildkite/scripts/rocm/build-test-image.sh +++ b/.buildkite/scripts/rocm/build-test-image.sh @@ -13,6 +13,18 @@ metadata_get() { fi } +use_ci_base_if_present() { + local ci_base_image="" + + ci_base_image="$(metadata_get rocm-ci-base-image)" + if [[ -z "${ci_base_image}" ]]; then + return 1 + fi + + export CI_BASE_IMAGE="${ci_base_image}" + echo "Using ROCm ci_base image selected by the preceding build step: ${CI_BASE_IMAGE}" +} + use_refreshed_base_if_present() { local base_refreshed="" @@ -22,15 +34,12 @@ use_refreshed_base_if_present() { fi export BASE_IMAGE - export CI_BASE_IMAGE export IMAGE_TAG_LATEST BASE_IMAGE="$(metadata_get rocm-base-image)" - CI_BASE_IMAGE="$(metadata_get rocm-ci-base-image)" IMAGE_TAG_LATEST="$(metadata_get rocm-ci-image-descriptive)" echo "Using refreshed ROCm base image for test image: ${BASE_IMAGE}" - echo "Using refreshed ROCm ci_base image for test image: ${CI_BASE_IMAGE}" if [[ -n "${IMAGE_TAG_LATEST}" ]]; then echo "Also tagging full ROCm CI image as: ${IMAGE_TAG_LATEST}" fi @@ -41,6 +50,8 @@ use_refreshed_base_if_present() { main() { local base_refreshed=0 + use_ci_base_if_present || true + if use_refreshed_base_if_present; then base_refreshed=1 fi diff --git a/.buildkite/test_areas/disaggregated.yaml b/.buildkite/test_areas/disaggregated.yaml index b25884a96f54..4012f1ebd539 100644 --- a/.buildkite/test_areas/disaggregated.yaml +++ b/.buildkite/test_areas/disaggregated.yaml @@ -203,3 +203,25 @@ steps: commands: - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh - bash v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh + +# P TP 4 - D DPEP 4 test case for DSv4-Flash +- label: DSv4-Flash Disaggregated DP EP + key: dsv4-flash-disaggregated + timeout_in_minutes: 60 + device: h200 + optional: true + working_dir: "/vllm-workspace/tests" + num_devices: 8 + env: + ENABLE_HMA_FLAG: "1" + DP_EP: "1" + GPU_MEMORY_UTILIZATION: "0.85" + PREFILLER_TP_SIZE: "4" + DECODER_TP_SIZE: "4" + PREFILL_BLOCK_SIZE: "256" + DECODE_BLOCK_SIZE: "256" + MODEL_NAMES: "deepseek-ai/DeepSeek-V4-Flash" + VLLM_SERVE_EXTRA_ARGS: "--trust-remote-code,--kv-cache-dtype,fp8" + commands: + - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh + - bash v1/kv_connector/nixl_integration/run_accuracy_test.sh diff --git a/.github/workflows/issue_autolabel.yml b/.github/workflows/issue_autolabel.yml index 7a98ce7cc08a..e5d1accf477c 100644 --- a/.github/workflows/issue_autolabel.yml +++ b/.github/workflows/issue_autolabel.yml @@ -323,7 +323,7 @@ jobs: // {users} will be replaced with @mentions const ccConfig = { rocm: { - users: ['hongxiayang', 'tjtanaa', 'vllmellm'], + users: ['hongxiayang', 'tjtanaa', 'vllmellm', 'giuseppegrossi'], message: 'CC {users} for ROCm-related issue', }, mistral: { diff --git a/csrc/libtorch_stable/moe/topk_softplus_sqrt_kernels.cu b/csrc/libtorch_stable/moe/topk_softplus_sqrt_kernels.cu index 095a76678311..785bbf2f6e07 100644 --- a/csrc/libtorch_stable/moe/topk_softplus_sqrt_kernels.cu +++ b/csrc/libtorch_stable/moe/topk_softplus_sqrt_kernels.cu @@ -44,6 +44,12 @@ typedef __hip_bfloat162 __nv_bfloat162; namespace vllm { namespace moe { +template +__device__ __forceinline__ int64_t load_index_as_int64(const HashIndType* ptr, + int64_t offset) { + return static_cast(ptr[offset]); +} + /// Aligned array type template + typename HashIndType, typename InputType = float> __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__ void topkGatingSoftplusSqrt( const InputType* input, const bool* finished, float* output, const int num_rows, IndType* indices, int* source_rows, const int k, const int start_expert, const int end_expert, const bool renormalize, double routed_scaling_factor, const float* correction_bias, - const IndType* input_ids, const IndType* tid2eid) { + const HashIndType* input_ids, const HashIndType* tid2eid) { static_assert(std::is_same_v || std::is_same_v || std::is_same_v, @@ -240,8 +246,8 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__ // Hash MoE path: indices are predetermined from lookup table if constexpr (USE_HASH) { - const IndType token_id = input_ids[thread_row]; - const IndType* expert_indices_for_token = tid2eid + token_id * k; + const int64_t token_id = load_index_as_int64(input_ids, thread_row); + const int64_t token_expert_offset = token_id * static_cast(k); #pragma unroll for (int ii = 0; ii < VPT; ++ii) { float val = row_chunk[ii]; @@ -252,7 +258,8 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__ float selected_sum = 0.f; #pragma unroll for (int k_idx = 0; k_idx < k; ++k_idx) { - const int expert = expert_indices_for_token[k_idx]; + const int expert = static_cast( + load_index_as_int64(tid2eid, token_expert_offset + k_idx)); const int idx = k * thread_row + k_idx; for (int ii = 0; ii < VPT; ++ii) { const int group_id = ii / ELTS_PER_LDG; @@ -261,7 +268,7 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__ group_id * THREADS_PER_ROW * ELTS_PER_LDG + local_id; if (expert == expert_idx) { - indices[idx] = expert; + indices[idx] = static_cast(expert); selected_sum += row_chunk[ii]; break; } @@ -285,7 +292,8 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__ #pragma unroll for (int k_idx = 0; k_idx < k; ++k_idx) { - const int expert = expert_indices_for_token[k_idx]; + const int expert = static_cast( + load_index_as_int64(tid2eid, token_expert_offset + k_idx)); const int idx = k * thread_row + k_idx; for (int ii = 0; ii < VPT; ++ii) { const int group_id = ii / ELTS_PER_LDG; @@ -461,14 +469,15 @@ struct TopkConstants { } template + int MAX_BYTES_PER_LDG, typename IndType, typename HashIndType, + typename InputType> void topkGatingSoftplusSqrtLauncherHelper( const InputType* input, const bool* finished, float* output, IndType* indices, int* source_row, const int num_rows, const int k, const int start_expert, const int end_expert, const bool renormalize, double routed_scaling_factor, const float* correction_bias, - const bool use_hash, const IndType* input_ids, const IndType* tid2eid, - cudaStream_t stream) { + const bool use_hash, const HashIndType* input_ids, + const HashIndType* tid2eid, cudaStream_t stream) { static constexpr int BYTES_PER_LDG = MIN(MAX_BYTES_PER_LDG, sizeof(InputType) * EXPERTS); using Constants = @@ -481,7 +490,8 @@ void topkGatingSoftplusSqrtLauncherHelper( DISPATCH_HASH(use_hash, USE_HASH, { auto* kernel = &topkGatingSoftplusSqrt; + WARP_SIZE_PARAM, USE_HASH, IndType, HashIndType, + InputType>; #ifndef USE_ROCM cudaLaunchConfig_t config = {}; config.gridDim = num_blocks; @@ -538,13 +548,14 @@ void topkGatingSoftplusSqrtLauncherHelper( } #endif -template +template void topkGatingSoftplusSqrtKernelLauncher( const InputType* gating_output, float* topk_weights, IndType* topk_indices, int* token_expert_indices, const int num_tokens, const int num_experts, const int topk, const bool renormalize, double routed_scaling_factor, - const float* correction_bias, const bool use_hash, const IndType* input_ids, - const IndType* tid2eid, cudaStream_t stream) { + const float* correction_bias, const bool use_hash, + const HashIndType* input_ids, const HashIndType* tid2eid, + cudaStream_t stream) { static constexpr int WARPS_PER_TB = 4; static constexpr int BYTES_PER_LDG_POWER_OF_2 = 16; // for bfloat16 dtype, we need 4 bytes loading to make sure num_experts @@ -644,57 +655,55 @@ void dispatch_topk_softplus_sqrt_launch( if (correction_bias.has_value()) { bias_ptr = correction_bias.value().const_data_ptr(); } - bool use_hash = false; - if (tid2eid.has_value()) { - STD_TORCH_CHECK(input_ids.has_value(), - "input_ids is required for hash MoE"); - use_hash = true; - } - if (topk_indices.scalar_type() == torch::headeronly::ScalarType::Int) { - const int* input_ids_ptr = nullptr; - const int* tid2eid_ptr = nullptr; + + auto launch = [&](auto* topk_indices_ptr) { + using OutIndType = + typename std::remove_pointer::type; if (tid2eid.has_value()) { - input_ids_ptr = input_ids.value().const_data_ptr(); - tid2eid_ptr = tid2eid.value().const_data_ptr(); + STD_TORCH_CHECK(input_ids.has_value(), + "input_ids is required for hash MoE"); + STD_TORCH_CHECK( + input_ids.value().scalar_type() == tid2eid.value().scalar_type(), + "input_ids and tid2eid must have the same dtype"); + if (tid2eid.value().scalar_type() == + torch::headeronly::ScalarType::Long) { + vllm::moe::topkGatingSoftplusSqrtKernelLauncher( + gating_output, topk_weights.mutable_data_ptr(), + topk_indices_ptr, token_expert_indices.mutable_data_ptr(), + num_tokens, num_experts, topk, renormalize, routed_scaling_factor, + bias_ptr, true, input_ids.value().const_data_ptr(), + tid2eid.value().const_data_ptr(), stream); + } else { + STD_TORCH_CHECK(tid2eid.value().scalar_type() == + torch::headeronly::ScalarType::Int); + vllm::moe::topkGatingSoftplusSqrtKernelLauncher( + gating_output, topk_weights.mutable_data_ptr(), + topk_indices_ptr, token_expert_indices.mutable_data_ptr(), + num_tokens, num_experts, topk, renormalize, routed_scaling_factor, + bias_ptr, true, input_ids.value().const_data_ptr(), + tid2eid.value().const_data_ptr(), stream); + } + } else { + vllm::moe::topkGatingSoftplusSqrtKernelLauncher( + gating_output, topk_weights.mutable_data_ptr(), + topk_indices_ptr, token_expert_indices.mutable_data_ptr(), + num_tokens, num_experts, topk, renormalize, routed_scaling_factor, + bias_ptr, false, static_cast(nullptr), + static_cast(nullptr), stream); } + }; - vllm::moe::topkGatingSoftplusSqrtKernelLauncher( - gating_output, topk_weights.mutable_data_ptr(), - topk_indices.mutable_data_ptr(), - token_expert_indices.mutable_data_ptr(), num_tokens, num_experts, - topk, renormalize, routed_scaling_factor, bias_ptr, use_hash, - input_ids_ptr, tid2eid_ptr, stream); + if (topk_indices.scalar_type() == torch::headeronly::ScalarType::Int) { + launch(topk_indices.mutable_data_ptr()); } else if (topk_indices.scalar_type() == torch::headeronly::ScalarType::UInt32) { - const uint32_t* input_ids_ptr = nullptr; - const uint32_t* tid2eid_ptr = nullptr; - if (tid2eid.has_value()) { - input_ids_ptr = input_ids.value().const_data_ptr(); - tid2eid_ptr = tid2eid.value().const_data_ptr(); - } - vllm::moe::topkGatingSoftplusSqrtKernelLauncher( - gating_output, topk_weights.mutable_data_ptr(), - topk_indices.mutable_data_ptr(), - token_expert_indices.mutable_data_ptr(), num_tokens, num_experts, - topk, renormalize, routed_scaling_factor, bias_ptr, use_hash, - input_ids_ptr, tid2eid_ptr, stream); + launch(topk_indices.mutable_data_ptr()); } else { STD_TORCH_CHECK(topk_indices.scalar_type() == torch::headeronly::ScalarType::Long); - - const int64_t* input_ids_ptr = nullptr; - const int64_t* tid2eid_ptr = nullptr; - if (tid2eid.has_value()) { - input_ids_ptr = input_ids.value().const_data_ptr(); - tid2eid_ptr = tid2eid.value().const_data_ptr(); - } - - vllm::moe::topkGatingSoftplusSqrtKernelLauncher( - gating_output, topk_weights.mutable_data_ptr(), - topk_indices.mutable_data_ptr(), - token_expert_indices.mutable_data_ptr(), num_tokens, num_experts, - topk, renormalize, routed_scaling_factor, bias_ptr, use_hash, - input_ids_ptr, tid2eid_ptr, stream); + launch(topk_indices.mutable_data_ptr()); } } @@ -738,4 +747,4 @@ void topk_softplus_sqrt( STD_TORCH_CHECK(false, "Unsupported gating_output data type: ", gating_output.scalar_type()); } -} \ No newline at end of file +} diff --git a/docker/Dockerfile b/docker/Dockerfile index b47853a06c73..1263bff436fb 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -815,7 +815,7 @@ RUN --mount=type=cache,target=/opt/uv/cache \ # Install FlashInfer JIT cache (requires CUDA-version-specific index URL) # https://docs.flashinfer.ai/installation.html # From versions.json: .flashinfer.version -ARG FLASHINFER_VERSION=0.6.13 +ARG FLASHINFER_VERSION=0.6.14 RUN --mount=type=cache,target=/opt/uv/cache \ uv pip install --system flashinfer-jit-cache==${FLASHINFER_VERSION} \ --index-url https://flashinfer.ai/whl/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') diff --git a/docker/versions.json b/docker/versions.json index 4dffa00985c9..e6839bbb05cf 100644 --- a/docker/versions.json +++ b/docker/versions.json @@ -68,7 +68,7 @@ "default": "true" }, "FLASHINFER_VERSION": { - "default": "0.6.13" + "default": "0.6.14" }, "GDRCOPY_CUDA_VERSION": { "default": "12.8" diff --git a/docs/deployment/k8s.md b/docs/deployment/k8s.md index 7a92c99b2c4a..e7d0853e9f49 100644 --- a/docs/deployment/k8s.md +++ b/docs/deployment/k8s.md @@ -217,7 +217,7 @@ INFO: Uvicorn running on http://0.0.0.0:8000 (Press CTRL+C to quit) image: vllm/vllm-openai:latest command: ["/bin/sh", "-c"] args: [ - "vllm serve mistralai/Mistral-7B-Instruct-v0.3 --trust-remote-code --enable-chunked-prefill --max_num_batched_tokens 1024" + "vllm serve mistralai/Mistral-7B-Instruct-v0.3 --trust-remote-code --enable-chunked-prefill --max-num-batched-tokens 1024" ] env: - name: HF_TOKEN @@ -306,7 +306,7 @@ INFO: Uvicorn running on http://0.0.0.0:8000 (Press CTRL+C to quit) - SYS_PTRACE command: ["/bin/sh", "-c"] args: [ - "vllm serve mistralai/Mistral-7B-v0.3 --port 8000 --trust-remote-code --enable-chunked-prefill --max_num_batched_tokens 1024" + "vllm serve mistralai/Mistral-7B-v0.3 --port 8000 --trust-remote-code --enable-chunked-prefill --max-num-batched-tokens 1024" ] env: - name: HF_TOKEN diff --git a/docs/features/kv_offloading_usage.md b/docs/features/kv_offloading_usage.md index 8a48d11be73a..7042ef787c9f 100644 --- a/docs/features/kv_offloading_usage.md +++ b/docs/features/kv_offloading_usage.md @@ -152,16 +152,25 @@ Object keys follow the same run-configuration digest scheme as the filesystem ti The P2P tier (`type: "p2p"`) shares completed KV blocks between vLLM instances over RDMA via NIXL. Each instance binds a control socket on `host:port` and exchanges blocks directly with peers — no shared filesystem required. +PYTHONHASHSEED environment variable must be set to the same fixed value on all nodes. + | Key | Required | Default | Notes | | --- | --- | --- | --- | | `type` | yes | — | Must be `p2p`. | -| `host` | no | `0.0.0.0` | Address the control socket binds to. | -| `port` | no | `7777` | Port for the control socket. Must be reachable from peers. | +| `host` | no | `$VLLM_P2P_SIDE_CHANNEL_HOST` (`localhost`) | Address the control socket binds to, used verbatim as the identity peers dial back. When omitted, resolves from the env var below. The `localhost` default binds loopback only — for cross-host P2P you **must** set it to the node's routable IP (see below). | +| `port` | no | `$VLLM_P2P_SIDE_CHANNEL_PORT` (`5710`) | Base port for the control socket. Must be reachable from peers. The bound port is `base + data_parallel_index` (one socket per DP replica). When omitted, the base resolves from the env var below. | | `backends` | no | `["UCX"]` | NIXL transport backends. See [NixlConnector Usage Guide](nixl_connector_usage.md#selecting-a-nixl-transport-backend-plugin) for available backends and selection guidance. | | `num_threads` | no | `4` | NIXL agent worker threads. Only used when `backends` is UCX-only; ignored when any non-UCX backend is requested. | The `backends` and `num_threads` options mirror the conditional logic used by [`NixlConnector`](nixl_connector_usage.md#selecting-a-nixl-transport-backend-plugin): when any non-UCX backend is configured, NIXL is initialised with `backends=...`; otherwise it falls back to a UCX-only agent with the configured `num_threads`. This lets the P2P tier use a different transport (e.g. `MOONCAKE`, `GDS_MT`, `LIBFABRIC`) than the main `NixlConnector` running in the same process. +#### Environment Variables + +Rather than embedding `host`/`port` in each `secondary_tiers` entry, set them once at deploy time via environment variables (mirroring `VLLM_NIXL_SIDE_CHANNEL_HOST`/`VLLM_NIXL_SIDE_CHANNEL_PORT`). Explicit `host`/`port` config keys, when present, take precedence. + +- `VLLM_P2P_SIDE_CHANNEL_HOST` (default `localhost`): address the P2P control socket binds to. It is used **verbatim** as both the bind address and the identity peers dial back — there is no auto-detection (this mirrors `VLLM_NIXL_SIDE_CHANNEL_HOST`). The default binds the loopback interface only, so peers on another host cannot reach it. **For any cross-host P2P deployment you must set this explicitly to the node's routable IP** (e.g. the pod IP) before launching `vllm serve` — otherwise remote peers will fail to connect. The NIXL agent name is a separate per-process identifier, so peers sharing a `host:port` never collide. +- `VLLM_P2P_SIDE_CHANNEL_PORT` (default `5710`): base port for the P2P control socket. The port actually bound is `VLLM_P2P_SIDE_CHANNEL_PORT + data_parallel_index` — one socket per DP replica, matching NIXL (for DP=1 the offset is 0). The peer's port is passed as `remote_port` in `kv_transfer_params`; the router/EPP that selects the DP rank (e.g. via the `X-data-parallel-rank` header) computes `remote_port = base + rank`. The DP-index offset separates replicas *within* one deployment; two co-located *deployments* (a prefiller and a decoder on the same host) still need distinct base ports (e.g. decoder base `5711`) to avoid a bind collision. + ## Tuning Tips - `cpu_bytes_to_use`: a bigger CPU tier means fewer trips to slower secondary tiers and a higher hit rate. The value is total across all workers, not per-worker. Leave headroom for the rest of the host workload. diff --git a/docs/models/pooling_models/README.md b/docs/models/pooling_models/README.md index f8de9d437ad7..d69a4dc616eb 100644 --- a/docs/models/pooling_models/README.md +++ b/docs/models/pooling_models/README.md @@ -279,11 +279,66 @@ the pooler assigned to each task has the following attributes by default: | `embed` | `LAST` | ✅︎ | ❌ | | `classify` | `LAST` | ❌ | ✅︎ | -When loading [Sentence Transformers](https://huggingface.co/sentence-transformers) models, -its Sentence Transformers configuration file (`modules.json`) takes priority over the model's defaults. +#### Resolution precedence -You can further customize this via the `--pooler-config` option, -which takes priority over both the model's and Sentence Transformers' defaults. +The pooling method and `use_activation` are resolved per field. An explicitly +set field in `--pooler-config` takes precedence over Sentence Transformers +metadata, which in turn takes precedence over the model architecture or task +default. Fields left unset continue through the chain independently. + +The current `PoolerConfig` has no `normalize` or `activation` field. +`use_activation` controls whether the task's constructed normalization or +classification activation is applied. + +| Field | Source precedence | How to override | +| ----- | ----------------- | --------------- | +| Pooling method (`pooling_type`) | `--pooler-config` > boolean `pooling_mode_*` fields in the Pooling module referenced by Sentence Transformers `modules.json` > architecture default (`LAST` for sequence pooling and `ALL` for token pooling unless the architecture overrides it) | Set `{"pooling_type": "CLS"}`, or set `seq_pooling_type` / `tok_pooling_type` explicitly. | +| Embedding normalization (`use_activation`) | `--pooler-config` > Sentence Transformers modules (`true` when a Normalize module is present, otherwise `false`) > pooling-task default (`true`) when no Sentence Transformers Pooling module is found | Set `{"use_activation": false}` to return unnormalized embeddings. | +| Classification activation function | Hugging Face `problem_type` > Sentence Transformers activation metadata > sigmoid or softmax selected from the label count | The function cannot be selected through `--pooler-config`; set `{"use_activation": false}` to return logits instead. | + +Sentence Transformers configurations using the newer compact `pooling_mode` +string are not currently parsed; see [issue #45995](https://github.com/vllm-project/vllm/issues/45995). + +For converted models and predefined models using the standard DispatchPooler +adapters, `embed` and `token_embed` construct an L2-normalization head, while +`classify` and `token_classify` construct the selected classification activation. +In both cases, `use_activation` controls whether that head is applied. Models with +custom poolers can implement different behavior. + +To inspect the resolved fields without loading model weights: + +```python +from vllm.config import ModelConfig, PoolerConfig +from vllm.model_executor.layers.pooler.activations import get_act_fn + + +def inspect(requested: PoolerConfig) -> None: + model_config = ModelConfig( + "intfloat/e5-small", + runner="pooling", + pooler_config=requested, + ) + resolved = model_config.pooler_config + assert resolved is not None + print( + { + "seq_pooling_type": resolved.seq_pooling_type, + "tok_pooling_type": resolved.tok_pooling_type, + "use_activation": resolved.use_activation, + "sequence_classification_activation": type( + get_act_fn(model_config.hf_config) + ).__name__, + } + ) + + +inspect(PoolerConfig()) +inspect(PoolerConfig(pooling_type="CLS", use_activation=False)) +``` + +For `intfloat/e5-small`, the first result contains `MEAN`, `ALL`, and `True`. +The second contains `CLS`, `ALL`, and `False`. Both report the classification +activation that the standard sequence-classification adapter would construct. ## Removed Features diff --git a/docs/models/pooling_models/token_classify.md b/docs/models/pooling_models/token_classify.md index 6b2cefbde558..faada9c2cee4 100644 --- a/docs/models/pooling_models/token_classify.md +++ b/docs/models/pooling_models/token_classify.md @@ -47,6 +47,8 @@ The BAAI/bge-m3 model leverages token classification for sparse retrieval. For m | `ModernBertForTokenClassification` | ModernBERT-based | `disham993/electrical-ner-ModernBERT-base` | | | | `OpenAIPrivacyFilterForTokenClassification` | gpt-oss-based encoder | `openai/privacy-filter` | | | | `Qwen3ForTokenClassification`C | Qwen3-based | `bd2lcco/Qwen3-0.6B-finetuned` | | | +| `RobertaForTokenClassification` | RoBERTa-based | `Jean-Baptiste/roberta-large-ner-english` | | | +| `XLMRobertaForTokenClassification` | XLM-RoBERTa-based | `Davlan/xlm-roberta-base-ner-hrl` | | | | `*Model`C, `*ForCausalLM`C, etc. | Generative models | N/A | \* | \* | C Automatically converted into a classification model via `--convert classify`. ([details](./README.md#model-conversion)) diff --git a/requirements/cuda.txt b/requirements/cuda.txt index 91a57997684d..ab1df3ffe609 100644 --- a/requirements/cuda.txt +++ b/requirements/cuda.txt @@ -11,8 +11,11 @@ torchvision==0.26.0 # Required for phi3v processor. See https://github.com/pytor torchcodec >= 0.14 PyNvVideoCodec==2.0.4 # FlashInfer should be updated together with the Dockerfile -flashinfer-python==0.6.13 -flashinfer-cubin==0.6.13 +# flashinfer-cubin is not on PyPI since 0.6.14; setup.py excludes it from +# install_requires so the published wheel does not carry an unresolvable pin +--extra-index-url https://flashinfer.ai/whl/ +flashinfer-python==0.6.14 +flashinfer-cubin==0.6.14 apache-tvm-ffi==0.1.9 tilelang==0.1.9 nvidia-cudnn-frontend>=1.19.1 diff --git a/rust/.config/nextest.toml b/rust/.config/nextest.toml new file mode 100644 index 000000000000..3f006b74bf9b --- /dev/null +++ b/rust/.config/nextest.toml @@ -0,0 +1,3 @@ +# Kill a hung test (per-test 120s budget) so a broken wait fails instead of stalling. +[profile.default] +slow-timeout = { period = "60s", terminate-after = 2 } diff --git a/rust/src/chat/src/output/default/unified.rs b/rust/src/chat/src/output/default/unified.rs index 038b907a26af..f865e5a6ae0d 100644 --- a/rust/src/chat/src/output/default/unified.rs +++ b/rust/src/chat/src/output/default/unified.rs @@ -451,6 +451,57 @@ mod tests { output } + #[tokio::test] + async fn unified_stream_parses_formatted_tool_call_without_latch() { + use vllm_parser::tool::{HermesToolParser, ToolParser as _}; + + let hermes = HermesToolParser::create(&[]).unwrap(); + let parser = vllm_parser::unified::CombinedParser::new(None, Some(hermes)); + + // Regression guard: a formatted call (space before the outer `}`) must not + // trip the parse-error latch that turns it and every later call into text. + let d1 = decoded_delta( + r#"{"name":"get_weather","arguments":{"location":"Paris"} }"#, + ); + let d2 = finished_delta( + r#"{"name":"get_time","arguments":{"tz":"UTC"}}"#, + ); + + let stream = stream::iter(vec![d1, d2].into_iter().map(Ok)); + let events = unified_event_stream(stream, Box::new(parser)) + .collect::>() + .await + .into_iter() + .collect::>>() + .unwrap(); + + let names: Vec<&str> = events + .iter() + .filter_map(|event| match event { + AssistantEvent::ToolCallStart { name, .. } => Some(name.as_str()), + _ => None, + }) + .collect(); + assert_eq!(names, ["get_weather", "get_time"]); + + let text_leaks = events + .iter() + .filter(|event| { + matches!( + event, + AssistantEvent::TextDelta { + kind: AssistantBlockKind::Text, + .. + } + ) + }) + .count(); + assert_eq!( + text_leaks, 0, + "formatted tool call leaked into text: {events:#?}" + ); + } + #[tokio::test] async fn unified_stream_emits_reasoning_only_deltas() { let events = collect( diff --git a/rust/src/parser/src/tool/json/llama.rs b/rust/src/parser/src/tool/json/llama.rs index e898d07d863a..23e49724f07d 100644 --- a/rust/src/parser/src/tool/json/llama.rs +++ b/rust/src/parser/src/tool/json/llama.rs @@ -238,7 +238,9 @@ fn parse_llama_arguments_event( /// Parse the outer closing brace for one Llama JSON tool call. fn tool_call_close_event(input: &mut JsonToolInput<'_>) -> ModalResult { - literal("}").value(LlamaJsonEvent::ToolCallClose).parse_next(input) + seq!(_: ws0, _: literal("}")) + .value(LlamaJsonEvent::ToolCallClose) + .parse_next(input) } /// Parse a semicolon separator after one Llama JSON tool call. @@ -266,6 +268,28 @@ mod tests { format!(r#"{{"name":"{function_name}","parameters":{parameters}}}"#) } + #[test] + fn llama_tolerates_whitespace_before_outer_brace() { + // Whitespace between the parameters object's `}` and the outer `}` must + // still parse (json.loads / raw_decode parity). + let mut whole = Llama3JsonToolParser::new(&test_tools()); + let whole_output = whole.parse_complete(r#"{"name":"f","parameters":{"x":1} }"#).unwrap(); + assert_eq!(whole_output.calls().len(), 1); + assert_eq!(whole_output.calls()[0].name.as_deref(), Some("f")); + assert_eq!(whole_output.calls()[0].arguments, r#"{"x":1}"#); + + // Same input, whitespace before the outer `}` split across a chunk boundary. + let mut chunked = Llama3JsonToolParser::new(&test_tools()); + let mut output = ToolParserOutput::default(); + for chunk in [r#"{"name":"f","parameters":{"x":1}"#, " ", "}"] { + output.append(chunked.parse_chunk(chunk).unwrap()); + } + output.append(chunked.finish().unwrap()); + let output = output.coalesce(); + assert_eq!(output.calls().len(), 1); + assert_eq!(output.calls()[0].arguments, r#"{"x":1}"#); + } + #[test] fn llama_json_parse_complete_without_tool_call_keeps_text() { let mut parser = Llama3JsonToolParser::new(&test_tools()); diff --git a/rust/src/parser/src/tool/json/mod.rs b/rust/src/parser/src/tool/json/mod.rs index ee85ddf356fb..b707bff6fd19 100644 --- a/rust/src/parser/src/tool/json/mod.rs +++ b/rust/src/parser/src/tool/json/mod.rs @@ -317,7 +317,7 @@ fn tool_call_close_event( input: &mut JsonToolInput<'_>, config: JsonToolCallConfig, ) -> ModalResult { - let _ = literal("}").parse_next(input)?; + seq!(_: ws0, _: literal("}")).parse_next(input)?; match config.delimiter { Some(delimiter) => alt(( @@ -406,6 +406,34 @@ mod tests { output.coalesce() } + #[test] + fn json_tool_call_tolerates_whitespace_before_outer_brace() { + // Pretty-printed JSON puts whitespace between the arguments object's `}` + // and the outer object's `}`; it must still parse (json.loads parity). + let mut whole = JsonToolCallParser::new(DELIMITED_CONFIG); + let whole_output = collect_chunks( + &mut whole, + &[r#"{"function":"f","parameters":{"x":1} }"#], + ); + assert_eq!(whole_output.calls().len(), 1); + assert_eq!(whole_output.calls()[0].name.as_deref(), Some("f")); + assert_eq!(whole_output.calls()[0].arguments, r#"{"x":1}"#); + + // Same input, but the whitespace before the outer `}` is split across a + // chunk boundary (exercises `ws0` returning Incomplete on `Partial`). + let mut chunked = JsonToolCallParser::new(DELIMITED_CONFIG); + let chunked_output = collect_chunks( + &mut chunked, + &[ + r#"{"function":"f","parameters":{"x":1}"#, + " ", + "}", + ], + ); + assert_eq!(chunked_output.calls().len(), 1); + assert_eq!(chunked_output.calls()[0].arguments, r#"{"x":1}"#); + } + #[test] fn json_tool_call_delimiter_extracts_multiple_calls_in_one_block() { let input = build_tool_calls(&[ diff --git a/rust/src/server/src/tls_tests.rs b/rust/src/server/src/tls_tests.rs index 32d5582100c2..f553d4b5c345 100644 --- a/rust/src/server/src/tls_tests.rs +++ b/rust/src/server/src/tls_tests.rs @@ -526,16 +526,19 @@ async fn tls_handshake_timeout_drops_silent_client() { let (addr, shutdown) = spawn_server(Some(server_tls(&certs, 0))).await; let mut tcp = TcpStream::connect(&addr).await.expect("connect"); - tokio::task::yield_now().await; - tokio::time::advance(tls::TLS_HANDSHAKE_TIMEOUT + Duration::from_millis(1)).await; - tokio::task::yield_now().await; - + let start = tokio::time::Instant::now(); let mut buf = [0u8; 1]; - let read = tokio::time::timeout(Duration::from_secs(1), tcp.read(&mut buf)).await; + let read = tcp.read(&mut buf).await; assert!( - matches!(read, Ok(Ok(0)) | Ok(Err(_))), + matches!(read, Ok(0) | Err(_)), "server must drop a stalled TLS handshake (expected close, got {read:?})" ); + // The close is the handshake deadline, not an earlier one + assert!( + start.elapsed() >= tls::TLS_HANDSHAKE_TIMEOUT, + "closed too early to be the handshake deadline: {:?}", + start.elapsed() + ); shutdown.cancel(); } diff --git a/scripts/autotune_helion_kernels.py b/scripts/autotune_helion_kernels.py index c02d2a0206b3..864aabfc0964 100644 --- a/scripts/autotune_helion_kernels.py +++ b/scripts/autotune_helion_kernels.py @@ -37,6 +37,7 @@ get_kernel_by_name, get_registered_kernels, ) + from vllm.kernels.helion.ops import import_all_kernels from vllm.kernels.helion.utils import get_canonical_gpu_name from vllm.logger import init_logger from vllm.utils.import_utils import has_helion @@ -382,6 +383,9 @@ def main(): args = parser.parse_args() + # import all helion kernels + import_all_kernels() + import logging if args.verbose: diff --git a/scripts/benchmark_helion_kernels.py b/scripts/benchmark_helion_kernels.py new file mode 100644 index 000000000000..a2640aa52ffa --- /dev/null +++ b/scripts/benchmark_helion_kernels.py @@ -0,0 +1,457 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Benchmark a registered Helion kernel against a baseline. + +For each input case produced by the kernel's registered input generator, this +measures the latency of the Helion kernel and a chosen baseline, then reports +the speedup. + +Two baselines are supported (``--baseline``): + +- ``autotune`` (default): the kernel's autotuning baseline + (``helion_settings.autotune_baseline_fn``), wrapped in ``torch.compile`` + (inductor). This is the native-torch reference used by kernel autotuning + and correctness unit tests. +- ``cuda``: the corresponding hand-written CUDA op (``torch.ops._C.*``). The + mapping from Helion kernel name to CUDA op lives in ``CUDA_BASELINE_OPS`` + below. Every Helion kernel shares the same argument interface as its CUDA + counterpart, so inputs are forwarded verbatim. + +Usage: + # List available kernels + python scripts/benchmark_helion_kernels.py --list + + # Benchmark a kernel against the autotune baseline (default) + python scripts/benchmark_helion_kernels.py --kernel per_token_group_fp8_quant + + # Benchmark against the CUDA baseline + python scripts/benchmark_helion_kernels.py --kernel per_token_group_fp8_quant \\ + --baseline cuda + + # Disable CUDA graph capture and save results + python scripts/benchmark_helion_kernels.py --kernel per_token_group_fp8_quant \\ + --no-cudagraph --output results.json +""" + +import argparse +import copy +import gc +import json +import statistics +import sys +from collections.abc import Callable +from dataclasses import asdict, dataclass + +import torch + +from vllm.triton_utils import triton + +try: + from vllm.benchmarks.lib.utils import default_vllm_config + from vllm.kernels.helion import get_kernel_by_name, get_registered_kernels + from vllm.kernels.helion.ops import import_all_kernels + from vllm.logger import init_logger + from vllm.utils.import_utils import has_helion +except ImportError as e: + print(f"Error importing vLLM: {e}") + print("Please ensure vLLM is installed and in your Python path") + sys.exit(1) + +logger = init_logger("vllm.scripts.benchmark_helion_kernels") + + +# Maps a Helion kernel name to the CUDA op (attribute on ``torch.ops._C``) that +# implements the same operation. Helion kernels share the CUDA op's argument +# interface, so the kernel's input tuple is forwarded verbatim. Add an entry +# here when introducing a new kernel whose baseline should be the CUDA op. +CUDA_BASELINE_OPS: dict[str, str] = { + "dynamic_per_token_scaled_fp8_quant": "dynamic_per_token_scaled_fp8_quant", + "fused_qk_norm_rope": "fused_qk_norm_rope", + "per_token_group_fp8_quant": "per_token_group_fp8_quant", + "rms_norm_dynamic_per_token_quant": "rms_norm_dynamic_per_token_quant", + "rms_norm_per_block_quant": "rms_norm_per_block_quant", + "silu_and_mul_per_block_quant": "silu_and_mul_per_block_quant", + "scaled_mm": "cutlass_scaled_mm", +} + +# torch.compile options for the torch baseline, mirroring how these kernels are +# compiled inside vLLM. +_TORCH_COMPILE_OPTIONS: dict[str, bool] = { + "enable_auto_functionalized_v2": False, + "size_asserts": False, + "alignment_asserts": False, + "scalar_asserts": False, + "combo_kernels": True, + "benchmark_combo_kernel": True, +} + + +@dataclass +class Row: + case: str + baseline_ms: float + kernel_ms: float + speedup_x: float + + +def print_table(rows: list[Row]) -> None: + headers = ["case", "baseline_ms", "kernel_ms", "speedup(x)"] + + data = [ + [ + r.case, + f"{r.baseline_ms:.3f}", + f"{r.kernel_ms:.3f}", + f"{r.speedup_x:.3f}", + ] + for r in rows + ] + + cols = list(zip(*([headers] + data))) + widths = [max(len(cell) for cell in col) for col in cols] + + def fmt(row: list[str]) -> str: + return " | ".join(cell.ljust(w) for cell, w in zip(row, widths)) + + print(fmt(headers)) + print("-+-".join("-" * w for w in widths)) + for row in data: + print(fmt(row)) + + +def list_kernels() -> None: + kernels = get_registered_kernels() + + if not kernels: + print("No Helion kernels found in registry.") + return + + print("Available Helion kernels:") + print("=" * 50) + for name in sorted(kernels.keys()): + cuda = CUDA_BASELINE_OPS.get(name) + suffix = "" if cuda else " (no CUDA baseline mapping)" + print(f" {name}{suffix}") + print(f"\nTotal: {len(kernels)} kernels") + + +def check_requirements() -> bool: + if not torch.accelerator.is_available(): + logger.error("CUDA is not available. Helion benchmarking requires GPU.") + return False + if not has_helion(): + logger.error("Helion is not installed. Please install Helion package.") + return False + return True + + +def make_cuda_baseline(kernel_name: str) -> Callable: + """Return a callable invoking the CUDA op mapped to ``kernel_name``. + + The Helion kernel and its CUDA op share the same argument interface, so the + input tuple is forwarded verbatim. + """ + cuda_op_name = CUDA_BASELINE_OPS.get(kernel_name) + if cuda_op_name is None: + logger.error( + "No CUDA baseline mapping for kernel '%s'. Add an entry to " + "CUDA_BASELINE_OPS in %s (mapping the kernel name to its " + "torch.ops._C. name), or benchmark with --baseline torch.", + kernel_name, + __file__, + ) + sys.exit(1) + + cuda_op = getattr(torch.ops._C, cuda_op_name, None) + if cuda_op is None: + logger.error( + "torch.ops._C.%s is not available. Ensure the vLLM C extension is " + "built and loaded.", + cuda_op_name, + ) + sys.exit(1) + + return cuda_op + + +def make_autotune_baseline(kernel_name: str) -> Callable: + """Return the kernel's autotune baseline wrapped in ``torch.compile``. + + The baseline is the native-torch reference the kernel is tuned against, + registered via ``helion_settings.autotune_baseline_fn``. + """ + wrapper = get_kernel_by_name(kernel_name) + settings = wrapper.helion_settings + baseline_fn = getattr(settings, "autotune_baseline_fn", None) + if baseline_fn is None: + logger.error( + "Kernel '%s' has no autotune_baseline_fn in its helion_settings, so " + "the 'autotune' baseline is unavailable. Register one via " + "register_kernel(..., helion_settings=helion.Settings(" + "autotune_baseline_fn=...)), or benchmark with --baseline cuda.", + kernel_name, + ) + sys.exit(1) + + return torch.compile( + baseline_fn, + fullgraph=True, + dynamic=False, + backend="inductor", + options=_TORCH_COMPILE_OPTIONS, + ) + + +def cleanup_gpu_resources() -> None: + try: + torch.accelerator.empty_cache() + gc.collect() + if hasattr(torch, "_dynamo"): + torch._dynamo.reset() + torch.accelerator.synchronize() + except Exception as e: + logger.warning("Failed to cleanup GPU resources: %s", e) + + +_REDUCERS: dict[str, Callable[[list[float]], float]] = { + "min": min, + "max": max, + "mean": statistics.fmean, + "median": statistics.median, +} + + +def _reduce(times: list[float], return_mode: str) -> float: + return _REDUCERS[return_mode](times) + + +def do_bench_cudagraph_l2_clear( + fn: Callable, rep: int = 100, return_mode: str = "mean" +) -> float: + """CUDA-graph benchmark that flushes the L2 cache before every call. + + ``triton.testing.do_bench_cudagraph`` captures back-to-back kernel launches + with a warm L2 cache, which over-estimates performance for memory-bound + kernels. This clears L2 (via triton's benchmark cache buffer) before each + call and subtracts the isolated cache-clear cost from the measured time. + + Adapted from tritonbench's ``_do_bench_cudagraph_with_cache_clear`` using + only triton/torch primitives so no extra dependency is introduced. + """ + cache = triton.runtime.driver.active.get_empty_cache_for_benchmark() + clear_cache = cache.zero_ + + s = torch.Stream() + with s: + clear_cache() + fn() + + start_event = torch.Event(enable_timing=True) + end_event = torch.Event(enable_timing=True) + start_event.record() + for _ in range(5): + clear_cache() + fn() + end_event.record() + torch.accelerator.synchronize() + estimate_ms = start_event.elapsed_time(end_event) / 5 + n_repeat = 1000 if estimate_ms == 0 else max(1, int(rep / estimate_ms)) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + for _ in range(n_repeat): + clear_cache() + fn() + + clear_graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(clear_graph): + for _ in range(n_repeat): + clear_cache() + torch.accelerator.synchronize() + + kernel_times = [] + for _ in range(10): + start_event = torch.Event(enable_timing=True) + end_event = torch.Event(enable_timing=True) + start_event.record() + clear_graph.replay() + end_event.record() + torch.accelerator.synchronize() + clear_ms = start_event.elapsed_time(end_event) / n_repeat + + start_event = torch.Event(enable_timing=True) + end_event = torch.Event(enable_timing=True) + start_event.record() + graph.replay() + end_event.record() + torch.accelerator.synchronize() + total_ms = start_event.elapsed_time(end_event) / n_repeat + + kernel_times.append(total_ms - clear_ms) + + return _reduce(kernel_times, return_mode) + + +@torch.inference_mode() +def benchmark( + kernel_name: str, + baseline_fn: Callable, + repeat: int, + cudagraph: bool, + return_mode: str, +) -> list[Row]: + kernel = get_kernel_by_name(kernel_name) + # do_bench already flushes L2 per call; do_bench_cudagraph does not, so use + # the cache-clearing variant to avoid warm-L2 over-estimates. + benchmark_fn = do_bench_cudagraph_l2_clear if cudagraph else triton.testing.do_bench + + inputs_dict = kernel.get_inputs() + rows: list[Row] = [] + + for key, inputs in inputs_dict.items(): + logger.info("Benchmarking case %s", key) + + # Kernels may mutate their inputs in place; give each side its own copy. + kernel_inputs = copy.deepcopy(inputs) + baseline_inputs = copy.deepcopy(inputs) + + kernel_latency = benchmark_fn( + lambda kernel_inputs=kernel_inputs: kernel(*kernel_inputs), + rep=repeat, + return_mode=return_mode, + ) + baseline_latency = benchmark_fn( + lambda baseline_inputs=baseline_inputs: baseline_fn(*baseline_inputs), + rep=repeat, + return_mode=return_mode, + ) + + rows.append( + Row( + case=str(key), + baseline_ms=baseline_latency, + kernel_ms=kernel_latency, + speedup_x=baseline_latency / kernel_latency, + ) + ) + cleanup_gpu_resources() + + return rows + + +def main() -> None: + parser = argparse.ArgumentParser( + description="Benchmark a Helion kernel against a baseline", + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog=__doc__.split("Usage:")[1] if "Usage:" in __doc__ else "", + ) + parser.add_argument( + "--list", + action="store_true", + help="List available Helion kernels and exit", + ) + parser.add_argument( + "--kernel", + type=str, + help="Name of the single Helion kernel to benchmark", + ) + parser.add_argument( + "--repeat", + type=int, + default=100, + help="Number of benchmark repetitions (default: 100)", + ) + parser.add_argument( + "--no-cudagraph", + dest="cudagraph", + action="store_false", + help="Disable CUDA graph mode (enabled by default)", + ) + parser.add_argument( + "--baseline", + choices=["cuda", "autotune"], + default="autotune", + help=( + "Baseline to compare against: 'autotune' uses the kernel's " + "autotune_baseline_fn under torch.compile; 'cuda' uses the mapped " + "torch.ops._C op (default: autotune)" + ), + ) + parser.add_argument( + "--return-mode", + choices=["min", "max", "mean", "median"], + default="mean", + help="Statistic to report from the benchmark samples (default: mean)", + ) + parser.add_argument( + "--output", + type=str, + help="Path to save benchmark results as JSON (default: log only)", + ) + + args = parser.parse_args() + + import_all_kernels() + + if args.list: + list_kernels() + return + + if not args.kernel: + parser.error("--kernel is required (or use --list to see available kernels)") + + kernels = get_registered_kernels() + if args.kernel not in kernels: + logger.error("Kernel '%s' not found in registry.", args.kernel) + logger.error("Available kernels: %s", sorted(kernels.keys())) + sys.exit(1) + + wrapper = kernels[args.kernel] + if wrapper._disabled: + logger.error( + "Kernel '%s' is disabled: %s", + args.kernel, + wrapper._disabled_reason, + ) + sys.exit(1) + + if not check_requirements(): + sys.exit(1) + + with default_vllm_config(): + if args.baseline == "cuda": + baseline_fn = make_cuda_baseline(args.kernel) + else: + baseline_fn = make_autotune_baseline(args.kernel) + + rows = benchmark( + args.kernel, + baseline_fn, + args.repeat, + args.cudagraph, + args.return_mode, + ) + + print_table(rows) + + if args.output: + with open(args.output, "w") as f: + json.dump( + { + "kernel": args.kernel, + "baseline": args.baseline, + "cudagraph": args.cudagraph, + "repeat": args.repeat, + "return_mode": args.return_mode, + "results": [asdict(r) for r in rows], + }, + f, + indent=2, + ) + logger.info("Saved results to %s", args.output) + + +if __name__ == "__main__": + main() diff --git a/setup.py b/setup.py index e8f529701845..6419e3073ea8 100644 --- a/setup.py +++ b/setup.py @@ -1075,6 +1075,11 @@ def _read_requirements(filename: str) -> list[str]: # vllm-flash-attn is built only for CUDA 12.x. # Skip for other versions. continue + if "flashinfer-cubin" in req: + # Not on PyPI since 0.6.14 (only https://flashinfer.ai/whl), so + # it cannot be a wheel dependency; flashinfer falls back to + # fetching cubins at runtime when the package is absent. + continue if "nvidia-cutlass-dsl[cu13]" in req and cuda_major == "12": # [cu13] extra is the default; strip it on CUDA 12 builds. req = req.replace("nvidia-cutlass-dsl[cu13]", "nvidia-cutlass-dsl") diff --git a/tests/compile/test_graph_partition.py b/tests/compile/test_graph_partition.py index 8e20b704facc..bb2a6f2aee55 100644 --- a/tests/compile/test_graph_partition.py +++ b/tests/compile/test_graph_partition.py @@ -701,3 +701,67 @@ def test_decompose_size_with_getitem_user(): f"getitem node '{node.name}' has {len(node.args)} args " f"(expected 2): {node.args}" ) + + +def test_decompose_size_leaves_scalar_size_with_dim(): + """ + Regression test: _decompose_size_nodes must leave x.size(dim) alone. + + x.size() returns a torch.Size tuple that can't cross split boundaries and + must be decomposed. x.size(dim), however, already returns a scalar + SymInt/int that crosses fine, so the pass must not touch it. + + The punica LoRA path traces token_lora_mapping[:x.size(0)] under a dynamic + batch dim, so the size(0) node ends up nested inside a slice object: + + %size = call_method[target="size"](args = (%x, 0)) + %slice = call_function[target=getitem]( + args = (%mapping, slice(None, %size, None))) + + The old pass tried to decompose this scalar node too and then erase it, but + the slice still referenced it, raising "Tried to erase Node size but it + still had N users". The fix skips size calls that carry a dim argument. + """ + from torch._dynamo.source import LocalSource + from torch._subclasses.fake_tensor import FakeTensorMode + from torch.fx.experimental.symbolic_shapes import ShapeEnv + + # Build graph: + # %x = placeholder + # %mapping = placeholder + # %size = x.size(0) # scalar, with a dim arg + # %sliced = mapping[slice(None, %size, None)] # size node inside a slice + graph = fx.Graph() + x = graph.placeholder("x") + mapping = graph.placeholder("token_lora_mapping") + size_node = graph.call_method("size", args=(x, 0)) + sliced_node = graph.call_function( + operator.getitem, + args=(mapping, slice(None, size_node, None)), + ) + graph.output((sliced_node,)) + + # dim 0 dynamic (SymInt) — the realistic Unsloth + LoRA case. Without the + # skip, the pass would build per-dim replacements and then crash trying to + # erase the still-referenced size node. + shape_env = ShapeEnv() + src = LocalSource("tokens") + sym_tokens = shape_env.create_symintnode(shape_env.create_symbol(4, src), hint=4) + fake_mode = FakeTensorMode(shape_env=shape_env) + with fake_mode: + fake_x = torch.empty_strided((sym_tokens, 8), (8, 1)) + x.meta["example_value"] = fake_x + + gm = fx.GraphModule(torch.nn.Module(), graph) + + # Must not raise "Tried to erase Node ... still had N users". + _decompose_size_nodes(gm) + + # The scalar x.size(0) node is left in place, untouched. + remaining = list(gm.graph.find_nodes(op="call_method", target="size")) + assert len(remaining) == 1, ( + f"x.size(0) should be left untouched, found {len(remaining)} size nodes" + ) + assert remaining[0].args == (x, 0), ( + f"size node args changed: {remaining[0].args} (expected (x, 0))" + ) diff --git a/tests/kernels/moe/test_topk_softplus_sqrt.py b/tests/kernels/moe/test_topk_softplus_sqrt.py index 1b68213fafef..46ca934c1462 100644 --- a/tests/kernels/moe/test_topk_softplus_sqrt.py +++ b/tests/kernels/moe/test_topk_softplus_sqrt.py @@ -153,9 +153,9 @@ def test_fused_topk_softplus_sqrt_hash( # experts. hash_indices_table = torch.stack( [torch.randperm(num_experts)[:topk] for _ in range(vocab_size)] - ).to(device="cuda", dtype=torch.int32) + ).to(device="cuda", dtype=torch.long) input_ids = torch.randint( - 0, vocab_size, (num_tokens,), dtype=torch.int32, device="cuda" + 0, vocab_size, (num_tokens,), dtype=torch.long, device="cuda" ) topk_weights_ref, topk_ids_ref = _torch_topk_softplus_sqrt( diff --git a/tests/kernels/quantization/test_int4_emulation_moe.py b/tests/kernels/quantization/test_int4_emulation_moe.py new file mode 100644 index 000000000000..7016860d86f2 --- /dev/null +++ b/tests/kernels/quantization/test_int4_emulation_moe.py @@ -0,0 +1,1189 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Correctness tests for Int4EmulationTritonExperts MoE backend. + +Tests the weight dequantization helpers (_unpack_and_dequant_int4_gptq, +_unpack_and_dequant_int4_awq) and full MoE forward pass +(_process_weights_emulation_gptq, _process_weights_emulation_awq) +for both symmetric and asymmetric zero-point cases. + +Run `pytest tests/kernels/quantization/test_int4_emulation_moe.py`. +""" + +import numpy +import pytest +import torch +import torch.nn.functional as F + +from vllm.model_executor.layers.fused_moe.activation import MoEActivation +from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEConfig, + FusedMoEParallelConfig, + RoutingMethodType, + int4_w4a16_moe_quant_config, +) +from vllm.model_executor.layers.fused_moe.experts.int4_emulation_moe import ( + Int4EmulationTritonExperts, +) +from vllm.model_executor.layers.fused_moe.oracle.int_wna16 import ( + _process_weights_emulation_awq, + _process_weights_emulation_gptq, + _unpack_and_dequant_int4_awq, + _unpack_and_dequant_int4_gptq, +) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + awq_pack, + gptq_pack, +) +from vllm.platforms import current_platform + +pytestmark = pytest.mark.skipif( + not current_platform.is_cuda_alike(), + reason="Int4EmulationTritonExperts requires CUDA.", +) + +device = "cuda" + +# (E, K, N, group_size) +SHAPES = [ + pytest.param(2, 64, 32, 32, id="tiny-gs32"), + pytest.param(4, 128, 64, 64, id="small-gs64"), + pytest.param(4, 256, 128, 128, id="medium-gs128"), +] + +# (E, K, N, top_k, group_size, num_tokens) +E2E_CONFIGS = [ + pytest.param(4, 64, 32, 2, 32, 8, id="tiny"), + pytest.param(8, 128, 64, 2, 64, 16, id="small"), +] + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _quantize_sym(w_fp: torch.Tensor, group_size: int): + """Quantize [K, N] float to int4 symmetric (uint4b8), return q and scale.""" + K, N = w_fp.shape + assert K % group_size == 0 + n_groups = K // group_size + w_grouped = w_fp.reshape(n_groups, group_size, N) + scale = w_grouped.abs().amax(dim=1) / 7.0 + scale = scale.clamp(min=1e-6) + w_quant = (w_grouped / scale.unsqueeze(1)).round().clamp(-8, 7) + q = (w_quant + 8).to(torch.int32).reshape(K, N) + return q, scale + + +def _quantize_asym(w_fp: torch.Tensor, group_size: int): + """Quantize [K, N] float to uint4 asymmetric, return q, scale, zero.""" + K, N = w_fp.shape + assert K % group_size == 0 + n_groups = K // group_size + w_grouped = w_fp.reshape(n_groups, group_size, N) + wmin = w_grouped.amin(dim=1) + wmax = w_grouped.amax(dim=1) + scale = (wmax - wmin) / 15.0 + scale = scale.clamp(min=1e-6) + zero = (-wmin / scale).round().clamp(0, 15).to(torch.int32) + w_quant = ((w_grouped - wmin.unsqueeze(1)) / scale.unsqueeze(1)).round() + q = w_quant.clamp(0, 15).to(torch.int32).reshape(K, N) + return q, scale, zero + + +def _dequantize_ref( + w_uint: torch.Tensor, + scale: torch.Tensor, + zero=None, + output_dtype: torch.dtype = torch.bfloat16, +): + """Reference dequant for a single [K, N] slice.""" + K, N = w_uint.shape + n_groups = scale.shape[0] + group_size = K // n_groups + w = w_uint.reshape(n_groups, group_size, N).to(output_dtype) + s = scale.unsqueeze(1).to(output_dtype) + if zero is None: + return ((w - 8) * s).reshape(K, N) + z = zero.unsqueeze(1).to(output_dtype) + return ((w - z) * s).reshape(K, N) + + +def _pack_gptq_zeros(zero: torch.Tensor, N: int) -> torch.Tensor: + """Pack [n_groups, N] zeros into GPTQ format [n_groups, N//8] int32.""" + n_groups, _ = zero.shape + z = zero.to(torch.int32).cpu().numpy().astype(numpy.uint32) + packed = numpy.zeros((n_groups, N // 8), dtype=numpy.uint32) + for i in range(8): + packed |= z[:, i::8] << (i * 4) + return torch.from_numpy(packed.astype(numpy.int32)).to(device) + + +def _pack_awq_zeros(zero: torch.Tensor, N: int) -> torch.Tensor: + """Pack [n_groups, N] zeros into AWQ column format [n_groups, N//8] int32.""" + n_groups, _ = zero.shape + interleave = numpy.array([0, 2, 4, 6, 1, 3, 5, 7]) + z = zero.to(torch.int32).cpu().numpy().astype(numpy.uint32) + z_interleaved = z.reshape(-1, 8)[:, interleave].reshape(n_groups, N) + packed = numpy.zeros((n_groups, N // 8), dtype=numpy.uint32) + for i in range(8): + packed |= z_interleaved[:, i::8] << (i * 4) + return torch.from_numpy(packed.astype(numpy.int32)).to(device) + + +def _make_moe_config(E, K, N): + return FusedMoEConfig( + num_experts=E, + experts_per_token=2, + hidden_dim=K, + intermediate_size=N, + num_local_experts=E, + num_logical_experts=E, + moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), + activation=MoEActivation.SILU, + in_dtype=torch.bfloat16, + device=device, + routing_method=RoutingMethodType.TopK, + max_num_tokens=512, + ) + + +def _make_gptq_moe_weights(E, K, N, group_size, asym=False): + """Build GPTQ MoE weight tensors and per-expert float references.""" + torch.manual_seed(7) + w13_fp = torch.randn(E, K, 2 * N, dtype=torch.float16, device=device) + w2_fp = torch.randn(E, N, K, dtype=torch.float16, device=device) + + w13_list, w13s_list, w13z_list, w13_ref_list = [], [], [], [] + w2_list, w2s_list, w2z_list, w2_ref_list = [], [], [], [] + + for e in range(E): + if asym: + q13, s13, z13 = _quantize_asym(w13_fp[e], group_size) + w13_list.append(gptq_pack(q13, 4, K, 2 * N)) + w13s_list.append(s13) + w13z_list.append(_pack_gptq_zeros(z13, 2 * N)) + w13_ref_list.append(_dequantize_ref(q13, s13, z13)) + q2, s2, z2 = _quantize_asym(w2_fp[e], group_size) + w2_list.append(gptq_pack(q2, 4, N, K)) + w2s_list.append(s2) + w2z_list.append(_pack_gptq_zeros(z2, K)) + w2_ref_list.append(_dequantize_ref(q2, s2, z2)) + else: + q13, s13 = _quantize_sym(w13_fp[e], group_size) + w13_list.append(gptq_pack(q13, 4, K, 2 * N)) + w13s_list.append(s13) + w13z_list.append(None) + w13_ref_list.append(_dequantize_ref(q13, s13)) + q2, s2 = _quantize_sym(w2_fp[e], group_size) + w2_list.append(gptq_pack(q2, 4, N, K)) + w2s_list.append(s2) + w2z_list.append(None) + w2_ref_list.append(_dequantize_ref(q2, s2)) + + return ( + torch.stack(w13_list), + torch.stack(w13s_list), + torch.stack(w13z_list) if asym else None, + torch.stack(w13_ref_list), # [E, K, 2N] + torch.stack(w2_list), + torch.stack(w2s_list), + torch.stack(w2z_list) if asym else None, + torch.stack(w2_ref_list), # [E, N, K] + ) + + +def _make_awq_moe_weights(E, K, N, group_size, asym=False): + """Build AWQ MoE weight tensors and per-expert float references.""" + torch.manual_seed(8) + w13_fp = torch.randn(E, K, 2 * N, dtype=torch.float16, device=device) + w2_fp = torch.randn(E, N, K, dtype=torch.float16, device=device) + + w13_list, w13s_list, w13z_list, w13_ref_list = [], [], [], [] + w2_list, w2s_list, w2z_list, w2_ref_list = [], [], [], [] + + for e in range(E): + if asym: + q13, s13, z13 = _quantize_asym(w13_fp[e], group_size) + w13_list.append(awq_pack(q13, 4, K, 2 * N)) + w13s_list.append(s13) + w13z_list.append(_pack_awq_zeros(z13, 2 * N)) + w13_ref_list.append(_dequantize_ref(q13, s13, z13)) + q2, s2, z2 = _quantize_asym(w2_fp[e], group_size) + w2_list.append(awq_pack(q2, 4, N, K)) + w2s_list.append(s2) + w2z_list.append(_pack_awq_zeros(z2, K)) + w2_ref_list.append(_dequantize_ref(q2, s2, z2)) + else: + q13, s13 = _quantize_sym(w13_fp[e], group_size) + w13_list.append(awq_pack(q13, 4, K, 2 * N)) + w13s_list.append(s13) + w13z_list.append(None) + w13_ref_list.append(_dequantize_ref(q13, s13)) + q2, s2 = _quantize_sym(w2_fp[e], group_size) + w2_list.append(awq_pack(q2, 4, N, K)) + w2s_list.append(s2) + w2z_list.append(None) + w2_ref_list.append(_dequantize_ref(q2, s2)) + + return ( + torch.stack(w13_list), + torch.stack(w13s_list), + torch.stack(w13z_list) if asym else None, + torch.stack(w13_ref_list), # [E, K, 2N] + torch.stack(w2_list), + torch.stack(w2s_list), + torch.stack(w2z_list) if asym else None, + torch.stack(w2_ref_list), # [E, N, K] + ) + + +def _run_emulation_forward( + experts, w13_bf16, w2_bf16, hidden_states, topk_weights, topk_ids, E, K, N +): + ws13_size = hidden_states.shape[0] * topk_ids.shape[1] * max(N, K) + ws2_size = hidden_states.shape[0] * topk_ids.shape[1] * max(2 * N, K) + workspace13 = torch.zeros(ws13_size, dtype=hidden_states.dtype, device=device) + workspace2 = torch.zeros(ws2_size, dtype=hidden_states.dtype, device=device) + output = torch.zeros( + hidden_states.shape[0], K, dtype=hidden_states.dtype, device=device + ) + experts.apply( + output=output, + hidden_states=hidden_states, + w1=w13_bf16, + w2=w2_bf16, + topk_weights=topk_weights, + topk_ids=topk_ids, + activation=MoEActivation.SILU, + global_num_experts=E, + expert_map=None, + a1q_scale=None, + a2_scale=None, + workspace13=workspace13, + workspace2=workspace2, + expert_tokens_meta=None, + apply_router_weight_on_input=False, + ) + return output + + +# --------------------------------------------------------------------------- +# Tests: _unpack_and_dequant_int4_gptq +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("E, K, N, group_size", SHAPES) +def test_gptq_unpack_symmetric(E, K, N, group_size): + """GPTQ symmetric unpacker matches reference.""" + torch.manual_seed(0) + w_fp = torch.randn(E, K, N, dtype=torch.float16, device=device) + + packed_list, scale_list, ref_list = [], [], [] + for e in range(E): + q, s = _quantize_sym(w_fp[e], group_size) + packed_list.append(gptq_pack(q, 4, K, N)) + scale_list.append(s) + ref_list.append(_dequantize_ref(q, s, output_dtype=torch.float32)) + + w_packed = torch.stack(packed_list).to(device) + scale = torch.stack(scale_list).to(device) + ref = torch.stack(ref_list).to(device) + + out = _unpack_and_dequant_int4_gptq( + w_packed, scale, None, transpose_output=False, output_dtype=torch.float32 + ) + + assert out.shape == (E, K, N) + assert torch.allclose(out, ref, atol=0), ( + f"max diff: {(out - ref).abs().max().item()}" + ) + + +@pytest.mark.parametrize("E, K, N, group_size", SHAPES) +def test_gptq_unpack_asymmetric(E, K, N, group_size): + """GPTQ asymmetric unpacker matches reference.""" + torch.manual_seed(1) + w_fp = torch.randn(E, K, N, dtype=torch.float16, device=device) + + packed_list, scale_list, zero_list, ref_list = [], [], [], [] + for e in range(E): + q, s, z = _quantize_asym(w_fp[e], group_size) + packed_list.append(gptq_pack(q, 4, K, N)) + scale_list.append(s) + zero_list.append(_pack_gptq_zeros(z, N)) + ref_list.append(_dequantize_ref(q, s, z, output_dtype=torch.float32)) + + w_packed = torch.stack(packed_list).to(device) + scale = torch.stack(scale_list).to(device) + qzeros = torch.stack(zero_list).to(device) + ref = torch.stack(ref_list).to(device) + + out = _unpack_and_dequant_int4_gptq( + w_packed, scale, qzeros, transpose_output=False, output_dtype=torch.float32 + ) + + assert out.shape == (E, K, N) + assert torch.allclose(out, ref, atol=0), ( + f"max diff: {(out - ref).abs().max().item()}" + ) + + +@pytest.mark.parametrize("E, K, N, group_size", SHAPES) +def test_gptq_unpack_transpose(E, K, N, group_size): + """GPTQ transpose_output=True gives [E, N, K].""" + torch.manual_seed(2) + w_fp = torch.randn(E, K, N, dtype=torch.float16, device=device) + + packed_list, scale_list = [], [] + for e in range(E): + q, s = _quantize_sym(w_fp[e], group_size) + packed_list.append(gptq_pack(q, 4, K, N)) + scale_list.append(s) + + w_packed = torch.stack(packed_list).to(device) + scale = torch.stack(scale_list).to(device) + + out_normal = _unpack_and_dequant_int4_gptq(w_packed, scale, None, False) + out_transposed = _unpack_and_dequant_int4_gptq(w_packed, scale, None, True) + + assert out_transposed.shape == (E, N, K) + assert torch.allclose( + out_transposed, out_normal.permute(0, 2, 1).contiguous(), atol=0 + ) + + +# --------------------------------------------------------------------------- +# Tests: _unpack_and_dequant_int4_awq +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("E, K, N, group_size", SHAPES) +def test_awq_unpack_symmetric(E, K, N, group_size): + """AWQ symmetric unpacker matches reference.""" + torch.manual_seed(3) + w_fp = torch.randn(E, K, N, dtype=torch.float16, device=device) + + packed_list, scale_list, ref_list = [], [], [] + for e in range(E): + q, s = _quantize_sym(w_fp[e], group_size) + packed_list.append(awq_pack(q, 4, K, N)) + scale_list.append(s) + ref_list.append(_dequantize_ref(q, s, output_dtype=torch.float32)) + + w_packed = torch.stack(packed_list).to(device) + scale = torch.stack(scale_list).to(device) + ref = torch.stack(ref_list).to(device) + + out = _unpack_and_dequant_int4_awq( + w_packed, scale, None, transpose_output=False, output_dtype=torch.float32 + ) + + assert out.shape == (E, K, N) + assert torch.allclose(out, ref, atol=0), ( + f"max diff: {(out - ref).abs().max().item()}" + ) + + +@pytest.mark.parametrize("E, K, N, group_size", SHAPES) +def test_awq_unpack_asymmetric(E, K, N, group_size): + """AWQ asymmetric unpacker matches reference.""" + torch.manual_seed(4) + w_fp = torch.randn(E, K, N, dtype=torch.float16, device=device) + + packed_list, scale_list, zero_list, ref_list = [], [], [], [] + for e in range(E): + q, s, z = _quantize_asym(w_fp[e], group_size) + packed_list.append(awq_pack(q, 4, K, N)) + scale_list.append(s) + zero_list.append(_pack_awq_zeros(z, N)) + ref_list.append(_dequantize_ref(q, s, z, output_dtype=torch.float32)) + + w_packed = torch.stack(packed_list).to(device) + scale = torch.stack(scale_list).to(device) + qzeros = torch.stack(zero_list).to(device) + ref = torch.stack(ref_list).to(device) + + out = _unpack_and_dequant_int4_awq( + w_packed, scale, qzeros, transpose_output=False, output_dtype=torch.float32 + ) + + assert out.shape == (E, K, N) + assert torch.allclose(out, ref, atol=0), ( + f"max diff: {(out - ref).abs().max().item()}" + ) + + +@pytest.mark.parametrize("E, K, N, group_size", SHAPES) +def test_awq_unpack_transpose(E, K, N, group_size): + """AWQ transpose_output=True gives [E, N, K].""" + torch.manual_seed(5) + w_fp = torch.randn(E, K, N, dtype=torch.float16, device=device) + + packed_list, scale_list = [], [] + for e in range(E): + q, s = _quantize_sym(w_fp[e], group_size) + packed_list.append(awq_pack(q, 4, K, N)) + scale_list.append(s) + + w_packed = torch.stack(packed_list).to(device) + scale = torch.stack(scale_list).to(device) + + out_normal = _unpack_and_dequant_int4_awq(w_packed, scale, None, False) + out_transposed = _unpack_and_dequant_int4_awq(w_packed, scale, None, True) + + assert out_transposed.shape == (E, N, K) + assert torch.allclose( + out_transposed, out_normal.permute(0, 2, 1).contiguous(), atol=0 + ) + + +@pytest.mark.parametrize("E, K, N, group_size", SHAPES) +def test_awq_gptq_unpack_agree(E, K, N, group_size): + """AWQ and GPTQ unpackers produce identical values for the same weights.""" + torch.manual_seed(6) + w_fp = torch.randn(E, K, N, dtype=torch.float16, device=device) + + gptq_list, awq_list, scale_list = [], [], [] + for e in range(E): + q, s = _quantize_sym(w_fp[e], group_size) + gptq_list.append(gptq_pack(q, 4, K, N)) + awq_list.append(awq_pack(q, 4, K, N)) + scale_list.append(s) + + scale = torch.stack(scale_list).to(device) + out_gptq = _unpack_and_dequant_int4_gptq( + torch.stack(gptq_list).to(device), scale, None, False, torch.float32 + ) + out_awq = _unpack_and_dequant_int4_awq( + torch.stack(awq_list).to(device), scale, None, False, torch.float32 + ) + + assert torch.allclose(out_gptq, out_awq, atol=0), ( + f"max diff: {(out_gptq - out_awq).abs().max().item()}" + ) + + +# --------------------------------------------------------------------------- +# Tests: _process_weights_emulation_{gptq,awq} +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("E, K, N, group_size", SHAPES) +@pytest.mark.parametrize("asym", [False, True], ids=["sym", "asym"]) +def test_gptq_process_weights_shapes_and_values(E, K, N, group_size, asym): + """_process_weights_emulation_gptq shapes and values match reference.""" + w13, w13s, w13z, w13_ref, w2, w2s, w2z, w2_ref = _make_gptq_moe_weights( + E, K, N, group_size, asym + ) + result = _process_weights_emulation_gptq(w13, w2, w13s, w2s, w13z, w2z) + w13_out, w2_out = result[0], result[1] + + assert w13_out.shape == (E, 2 * N, K) + assert w2_out.shape == (E, K, N) + assert w13_out.dtype == torch.bfloat16 + assert w2_out.dtype == torch.bfloat16 + + expected_w13 = w13_ref.permute(0, 2, 1) + expected_w2 = w2_ref.permute(0, 2, 1) + + assert torch.allclose(w13_out.float(), expected_w13.float(), atol=0), ( + f"w13 max diff: {(w13_out.float() - expected_w13.float()).abs().max().item()}" + ) + assert torch.allclose(w2_out.float(), expected_w2.float(), atol=0), ( + f"w2 max diff: {(w2_out.float() - expected_w2.float()).abs().max().item()}" + ) + + +@pytest.mark.parametrize("E, K, N, group_size", SHAPES) +@pytest.mark.parametrize("asym", [False, True], ids=["sym", "asym"]) +def test_awq_process_weights_shapes_and_values(E, K, N, group_size, asym): + """_process_weights_emulation_awq shapes and values match reference.""" + w13, w13s, w13z, w13_ref, w2, w2s, w2z, w2_ref = _make_awq_moe_weights( + E, K, N, group_size, asym + ) + result = _process_weights_emulation_awq(w13, w2, w13s, w2s, w13z, w2z) + w13_out, w2_out = result[0], result[1] + + assert w13_out.shape == (E, 2 * N, K) + assert w2_out.shape == (E, K, N) + assert w13_out.dtype == torch.bfloat16 + assert w2_out.dtype == torch.bfloat16 + + expected_w13 = w13_ref.permute(0, 2, 1) + expected_w2 = w2_ref.permute(0, 2, 1) + + assert torch.allclose(w13_out.float(), expected_w13.float(), atol=0), ( + f"w13 max diff: {(w13_out.float() - expected_w13.float()).abs().max().item()}" + ) + assert torch.allclose(w2_out.float(), expected_w2.float(), atol=0), ( + f"w2 max diff: {(w2_out.float() - expected_w2.float()).abs().max().item()}" + ) + + +@pytest.mark.parametrize("E, K, N, group_size", SHAPES) +@pytest.mark.parametrize("asym", [False, True], ids=["sym", "asym"]) +def test_gptq_awq_process_weights_agree(E, K, N, group_size, asym): + """AWQ and GPTQ process_weights produce identical dequantized tensors.""" + torch.manual_seed(9) + w13_fp = torch.randn(E, K, 2 * N, dtype=torch.float16, device=device) + w2_fp = torch.randn(E, N, K, dtype=torch.float16, device=device) + + g13_list, g13s_list, g13z_list = [], [], [] + a13_list, a13s_list, a13z_list = [], [], [] + g2_list, g2s_list, g2z_list = [], [], [] + a2_list, a2s_list, a2z_list = [], [], [] + + for e in range(E): + if asym: + q13, s13, z13 = _quantize_asym(w13_fp[e], group_size) + q2, s2, z2 = _quantize_asym(w2_fp[e], group_size) + g13z_list.append(_pack_gptq_zeros(z13, 2 * N)) + a13z_list.append(_pack_awq_zeros(z13, 2 * N)) + g2z_list.append(_pack_gptq_zeros(z2, K)) + a2z_list.append(_pack_awq_zeros(z2, K)) + else: + q13, s13 = _quantize_sym(w13_fp[e], group_size) + q2, s2 = _quantize_sym(w2_fp[e], group_size) + g13z_list.append(None) + a13z_list.append(None) + g2z_list.append(None) + a2z_list.append(None) + + g13_list.append(gptq_pack(q13, 4, K, 2 * N)) + a13_list.append(awq_pack(q13, 4, K, 2 * N)) + g13s_list.append(s13) + a13s_list.append(s13) + g2_list.append(gptq_pack(q2, 4, N, K)) + a2_list.append(awq_pack(q2, 4, N, K)) + g2s_list.append(s2) + a2s_list.append(s2) + + gptq_res = _process_weights_emulation_gptq( + torch.stack(g13_list), + torch.stack(g2_list), + torch.stack(g13s_list), + torch.stack(g2s_list), + torch.stack(g13z_list) if asym else None, + torch.stack(g2z_list) if asym else None, + ) + awq_res = _process_weights_emulation_awq( + torch.stack(a13_list), + torch.stack(a2_list), + torch.stack(a13s_list), + torch.stack(a2s_list), + torch.stack(a13z_list) if asym else None, + torch.stack(a2z_list) if asym else None, + ) + + assert torch.allclose(gptq_res[0].float(), awq_res[0].float(), atol=1e-3), ( + f"w13 max diff: {(gptq_res[0] - awq_res[0]).float().abs().max().item()}" + ) + assert torch.allclose(gptq_res[1].float(), awq_res[1].float(), atol=1e-3), ( + f"w2 max diff: {(gptq_res[1] - awq_res[1]).float().abs().max().item()}" + ) + + +# --------------------------------------------------------------------------- +# End-to-end MoE forward pass tests +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("E, K, N, top_k, group_size, num_tokens", E2E_CONFIGS) +def test_gptq_vs_awq_forward_agree(E, K, N, top_k, group_size, num_tokens): + """GPTQ and AWQ emulation backends produce bit-identical forward outputs.""" + torch.manual_seed(42) + moe_config = _make_moe_config(E, K, N) + + w13_fp = torch.randn(E, K, 2 * N, dtype=torch.float16, device=device) + w2_fp = torch.randn(E, N, K, dtype=torch.float16, device=device) + + g13_list, g13s_list, a13_list, a13s_list = [], [], [], [] + g2_list, g2s_list, a2_list, a2s_list = [], [], [], [] + + for e in range(E): + q13, s13 = _quantize_sym(w13_fp[e], group_size) + q2, s2 = _quantize_sym(w2_fp[e], group_size) + g13_list.append(gptq_pack(q13, 4, K, 2 * N)) + a13_list.append(awq_pack(q13, 4, K, 2 * N)) + g13s_list.append(s13) + a13s_list.append(s13.clone()) + g2_list.append(gptq_pack(q2, 4, N, K)) + a2_list.append(awq_pack(q2, 4, N, K)) + g2s_list.append(s2) + a2s_list.append(s2.clone()) + + gptq_res = _process_weights_emulation_gptq( + torch.stack(g13_list), + torch.stack(g2_list), + torch.stack(g13s_list), + torch.stack(g2s_list), + None, + None, + ) + awq_res = _process_weights_emulation_awq( + torch.stack(a13_list), + torch.stack(a2_list), + torch.stack(a13s_list), + torch.stack(a2s_list), + None, + None, + ) + w13_gptq, w2_gptq = gptq_res[0], gptq_res[1] + w13_awq, w2_awq = awq_res[0], awq_res[1] + + dummy_scale = torch.ones(1, dtype=torch.float16, device=device) + experts_gptq = Int4EmulationTritonExperts( + moe_config, int4_w4a16_moe_quant_config(dummy_scale, dummy_scale) + ) + experts_awq = Int4EmulationTritonExperts( + moe_config, int4_w4a16_moe_quant_config(dummy_scale, dummy_scale) + ) + + hidden_states = torch.randn(num_tokens, K, dtype=torch.bfloat16, device=device) + topk_weights = torch.softmax( + torch.randn(num_tokens, top_k, dtype=torch.float32, device=device), dim=-1 + ) + topk_ids = torch.stack( + [torch.randperm(E, device=device)[:top_k] for _ in range(num_tokens)] + ).to(torch.int32) + + out_gptq = _run_emulation_forward( + experts_gptq, w13_gptq, w2_gptq, hidden_states, topk_weights, topk_ids, E, K, N + ) + out_awq = _run_emulation_forward( + experts_awq, w13_awq, w2_awq, hidden_states, topk_weights, topk_ids, E, K, N + ) + + assert torch.allclose(out_gptq, out_awq, atol=0), ( + f"max diff: {(out_gptq - out_awq).abs().max().item()}" + ) + + +# --------------------------------------------------------------------------- +# EP (Expert Parallelism) tests +# --------------------------------------------------------------------------- + +# (E, K, N, top_k, group_size, num_tokens, ep_size) +EP_CONFIGS = [ + pytest.param(4, 64, 32, 2, 32, 8, 2, id="E4-ep2"), + pytest.param(8, 64, 32, 2, 32, 16, 4, id="E8-ep4"), + pytest.param(8, 128, 64, 2, 64, 16, 2, id="E8-ep2"), +] + + +def _make_expert_map(global_num_experts: int, start: int, end: int) -> torch.Tensor: + """Build expert_map for a rank that owns experts [start, end).""" + expert_map = torch.full((global_num_experts,), -1, dtype=torch.int32, device=device) + expert_map[start:end] = torch.arange(end - start, dtype=torch.int32, device=device) + return expert_map + + +def _run_emulation_forward_ep( + experts, + w13_bf16, + w2_bf16, + hidden_states, + topk_weights, + topk_ids, + global_num_experts, + expert_map, +): + """Run forward with EP expert_map; returns output tensor.""" + T, K = hidden_states.shape + N = w2_bf16.shape[2] + ws13_size = T * topk_ids.shape[1] * max(N, K) + ws2_size = T * topk_ids.shape[1] * max(2 * N, K) + workspace13 = torch.zeros(ws13_size, dtype=hidden_states.dtype, device=device) + workspace2 = torch.zeros(ws2_size, dtype=hidden_states.dtype, device=device) + output = torch.zeros(T, K, dtype=hidden_states.dtype, device=device) + experts.apply( + output=output, + hidden_states=hidden_states, + w1=w13_bf16, + w2=w2_bf16, + topk_weights=topk_weights, + topk_ids=topk_ids, + activation=MoEActivation.SILU, + global_num_experts=global_num_experts, + expert_map=expert_map, + a1q_scale=None, + a2_scale=None, + workspace13=workspace13, + workspace2=workspace2, + expert_tokens_meta=None, + apply_router_weight_on_input=False, + ) + return output + + +@pytest.mark.parametrize("E, K, N, top_k, group_size, num_tokens, ep_size", EP_CONFIGS) +@pytest.mark.parametrize("fmt", ["gptq", "awq"]) +def test_ep_output_matches_no_ep(E, K, N, top_k, group_size, num_tokens, ep_size, fmt): + """EP simulation: sum of per-rank outputs equals the no-EP forward pass.""" + assert E % ep_size == 0 + num_local = E // ep_size + + torch.manual_seed(20) + + # Build all expert weights in BF16 (no-EP reference) + w13_fp = torch.randn(E, K, 2 * N, dtype=torch.float16, device=device) * 0.02 + w2_fp = torch.randn(E, N, K, dtype=torch.float16, device=device) * 0.02 + + packed13_list, scales13_list, packed2_list, scales2_list = [], [], [], [] + for e in range(E): + q13, s13 = _quantize_sym(w13_fp[e].float(), group_size) + q2, s2 = _quantize_sym(w2_fp[e].float(), group_size) + if fmt == "gptq": + packed13_list.append(gptq_pack(q13, 4, K, 2 * N)) + packed2_list.append(gptq_pack(q2, 4, N, K)) + else: + packed13_list.append(awq_pack(q13, 4, K, 2 * N)) + packed2_list.append(awq_pack(q2, 4, N, K)) + scales13_list.append(s13) + scales2_list.append(s2) + + process_fn = ( + _process_weights_emulation_gptq + if fmt == "gptq" + else _process_weights_emulation_awq + ) + res = process_fn( + torch.stack(packed13_list), + torch.stack(packed2_list), + torch.stack(scales13_list), + torch.stack(scales2_list), + None, + None, + ) + w13_all, w2_all = res[0], res[1] # [E, 2N, K], [E, K, N] + + hidden_states = torch.randn(num_tokens, K, dtype=torch.bfloat16, device=device) + topk_weights = torch.softmax( + torch.randn(num_tokens, top_k, dtype=torch.float32, device=device), dim=-1 + ) + topk_ids = torch.stack( + [torch.randperm(E, device=device)[:top_k] for _ in range(num_tokens)] + ).to(torch.int32) + + # No-EP reference + dummy_scale = torch.ones(1, dtype=torch.float16, device=device) + moe_config_full = FusedMoEConfig( + num_experts=E, + experts_per_token=top_k, + hidden_dim=K, + intermediate_size=N, + num_local_experts=E, + num_logical_experts=E, + moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), + activation=MoEActivation.SILU, + in_dtype=torch.bfloat16, + device=device, + routing_method=RoutingMethodType.TopK, + max_num_tokens=512, + ) + experts_ref = Int4EmulationTritonExperts( + moe_config_full, int4_w4a16_moe_quant_config(dummy_scale, dummy_scale) + ) + out_no_ep = _run_emulation_forward( + experts_ref, w13_all, w2_all, hidden_states, topk_weights, topk_ids, E, K, N + ) + + # EP simulation: sum contributions from each rank + out_ep_sum = torch.zeros(num_tokens, K, dtype=torch.bfloat16, device=device) + for rank in range(ep_size): + start = rank * num_local + end = start + num_local + expert_map = _make_expert_map(E, start, end) + moe_config_ep = FusedMoEConfig( + num_experts=E, + experts_per_token=top_k, + hidden_dim=K, + intermediate_size=N, + num_local_experts=num_local, + num_logical_experts=E, + moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), + activation=MoEActivation.SILU, + in_dtype=torch.bfloat16, + device=device, + routing_method=RoutingMethodType.TopK, + max_num_tokens=512, + ) + experts_ep = Int4EmulationTritonExperts( + moe_config_ep, int4_w4a16_moe_quant_config(dummy_scale, dummy_scale) + ) + out_rank = _run_emulation_forward_ep( + experts_ep, + w13_all[start:end], + w2_all[start:end], + hidden_states, + topk_weights, + topk_ids, + global_num_experts=E, + expert_map=expert_map, + ) + out_ep_sum = out_ep_sum + out_rank + + assert torch.allclose(out_ep_sum, out_no_ep, atol=1e-3), ( + f"[{fmt}] EP sum max diff: {(out_ep_sum - out_no_ep).abs().max().item():.6f}" + ) + + +@pytest.mark.parametrize("E, K, N, top_k, group_size, num_tokens, ep_size", EP_CONFIGS) +def test_ep_gptq_awq_agree(E, K, N, top_k, group_size, num_tokens, ep_size): + """With EP, GPTQ and AWQ emulation produce the same outputs per rank.""" + assert E % ep_size == 0 + num_local = E // ep_size + + torch.manual_seed(21) + w13_fp = torch.randn(E, K, 2 * N, dtype=torch.float16, device=device) * 0.02 + w2_fp = torch.randn(E, N, K, dtype=torch.float16, device=device) * 0.02 + + g13_list, g13s_list, g2_list, g2s_list = [], [], [], [] + a13_list, a13s_list, a2_list, a2s_list = [], [], [], [] + for e in range(E): + q13, s13 = _quantize_sym(w13_fp[e].float(), group_size) + q2, s2 = _quantize_sym(w2_fp[e].float(), group_size) + g13_list.append(gptq_pack(q13, 4, K, 2 * N)) + a13_list.append(awq_pack(q13, 4, K, 2 * N)) + g13s_list.append(s13) + a13s_list.append(s13.clone()) + g2_list.append(gptq_pack(q2, 4, N, K)) + a2_list.append(awq_pack(q2, 4, N, K)) + g2s_list.append(s2) + a2s_list.append(s2.clone()) + + gptq_res = _process_weights_emulation_gptq( + torch.stack(g13_list), + torch.stack(g2_list), + torch.stack(g13s_list), + torch.stack(g2s_list), + None, + None, + ) + awq_res = _process_weights_emulation_awq( + torch.stack(a13_list), + torch.stack(a2_list), + torch.stack(a13s_list), + torch.stack(a2s_list), + None, + None, + ) + w13_gptq, w2_gptq = gptq_res[0], gptq_res[1] + w13_awq, w2_awq = awq_res[0], awq_res[1] + + hidden_states = torch.randn(num_tokens, K, dtype=torch.bfloat16, device=device) + topk_weights = torch.softmax( + torch.randn(num_tokens, top_k, dtype=torch.float32, device=device), dim=-1 + ) + topk_ids = torch.stack( + [torch.randperm(E, device=device)[:top_k] for _ in range(num_tokens)] + ).to(torch.int32) + + dummy_scale = torch.ones(1, dtype=torch.float16, device=device) + + for rank in range(ep_size): + start = rank * num_local + end = start + num_local + expert_map = _make_expert_map(E, start, end) + moe_config_ep = FusedMoEConfig( + num_experts=E, + experts_per_token=top_k, + hidden_dim=K, + intermediate_size=N, + num_local_experts=num_local, + num_logical_experts=E, + moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), + activation=MoEActivation.SILU, + in_dtype=torch.bfloat16, + device=device, + routing_method=RoutingMethodType.TopK, + max_num_tokens=512, + ) + experts_gptq = Int4EmulationTritonExperts( + moe_config_ep, int4_w4a16_moe_quant_config(dummy_scale, dummy_scale) + ) + experts_awq = Int4EmulationTritonExperts( + moe_config_ep, int4_w4a16_moe_quant_config(dummy_scale, dummy_scale) + ) + out_gptq = _run_emulation_forward_ep( + experts_gptq, + w13_gptq[start:end], + w2_gptq[start:end], + hidden_states, + topk_weights, + topk_ids, + E, + expert_map, + ) + out_awq = _run_emulation_forward_ep( + experts_awq, + w13_awq[start:end], + w2_awq[start:end], + hidden_states, + topk_weights, + topk_ids, + E, + expert_map, + ) + assert torch.allclose(out_gptq, out_awq, atol=0), ( + f"rank={rank} max diff: {(out_gptq - out_awq).abs().max().item()}" + ) + + +@pytest.mark.parametrize("E, K, N, top_k, group_size, num_tokens, ep_size", EP_CONFIGS) +def test_ep_partial_rank_no_active_experts( + E, K, N, top_k, group_size, num_tokens, ep_size +): + """A rank that owns no token-selected experts produces an all-zero output.""" + assert E % ep_size == 0 + num_local = E // ep_size + + torch.manual_seed(22) + w13_fp = torch.randn(E, K, 2 * N, dtype=torch.float16, device=device) * 0.02 + w2_fp = torch.randn(E, N, K, dtype=torch.float16, device=device) * 0.02 + + packed13_list, scales13_list, packed2_list, scales2_list = [], [], [], [] + for e in range(E): + q13, s13 = _quantize_sym(w13_fp[e].float(), group_size) + q2, s2 = _quantize_sym(w2_fp[e].float(), group_size) + packed13_list.append(gptq_pack(q13, 4, K, 2 * N)) + packed2_list.append(gptq_pack(q2, 4, N, K)) + scales13_list.append(s13) + scales2_list.append(s2) + + res = _process_weights_emulation_gptq( + torch.stack(packed13_list), + torch.stack(packed2_list), + torch.stack(scales13_list), + torch.stack(scales2_list), + None, + None, + ) + w13_all, w2_all = res[0], res[1] + + # Force topk_ids to only use experts in [0, num_local) — rank 0's slice + topk_ids = torch.zeros(num_tokens, top_k, dtype=torch.int32, device=device) + topk_weights = torch.softmax( + torch.randn(num_tokens, top_k, dtype=torch.float32, device=device), dim=-1 + ) + hidden_states = torch.randn(num_tokens, K, dtype=torch.bfloat16, device=device) + + dummy_scale = torch.ones(1, dtype=torch.float16, device=device) + + # Last rank owns experts [E-num_local, E), tokens only route to [0, num_local) + last_rank = ep_size - 1 + start = last_rank * num_local + end = E + expert_map = _make_expert_map(E, start, end) + moe_config_ep = FusedMoEConfig( + num_experts=E, + experts_per_token=top_k, + hidden_dim=K, + intermediate_size=N, + num_local_experts=num_local, + num_logical_experts=E, + moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), + activation=MoEActivation.SILU, + in_dtype=torch.bfloat16, + device=device, + routing_method=RoutingMethodType.TopK, + max_num_tokens=512, + ) + experts_ep = Int4EmulationTritonExperts( + moe_config_ep, int4_w4a16_moe_quant_config(dummy_scale, dummy_scale) + ) + out = _run_emulation_forward_ep( + experts_ep, + w13_all[start:end], + w2_all[start:end], + hidden_states, + topk_weights, + topk_ids, + E, + expert_map, + ) + assert torch.all(out == 0), ( + f"Expected zeros for inactive rank, got max={out.abs().max().item()}" + ) + + +@pytest.mark.parametrize("E, K, N, top_k, group_size, num_tokens, ep_size", EP_CONFIGS) +def test_ep_sum_equals_full_forward(E, K, N, top_k, group_size, num_tokens, ep_size): + """With fixed routing, EP rank outputs sum to the single-rank full forward.""" + assert E % ep_size == 0 + num_local = E // ep_size + + torch.manual_seed(23) + w13_fp = torch.randn(E, K, 2 * N, dtype=torch.float16, device=device) * 0.02 + w2_fp = torch.randn(E, N, K, dtype=torch.float16, device=device) * 0.02 + + packed13_list, scales13_list, packed2_list, scales2_list = [], [], [], [] + for e in range(E): + q13, s13 = _quantize_sym(w13_fp[e].float(), group_size) + q2, s2 = _quantize_sym(w2_fp[e].float(), group_size) + packed13_list.append(gptq_pack(q13, 4, K, 2 * N)) + packed2_list.append(gptq_pack(q2, 4, N, K)) + scales13_list.append(s13) + scales2_list.append(s2) + + res = _process_weights_emulation_gptq( + torch.stack(packed13_list), + torch.stack(packed2_list), + torch.stack(scales13_list), + torch.stack(scales2_list), + None, + None, + ) + w13_all, w2_all = res[0], res[1] + + # Fix routing so every token uses exactly 2 consecutive experts (round-robin) + hidden_states = torch.randn(num_tokens, K, dtype=torch.bfloat16, device=device) + topk_ids = torch.stack( + [ + torch.tensor([(t * top_k + k) % E for k in range(top_k)], dtype=torch.int32) + for t in range(num_tokens) + ] + ).to(device) + topk_weights = torch.full((num_tokens, top_k), 1.0 / top_k, device=device) + + dummy_scale = torch.ones(1, dtype=torch.float16, device=device) + + # Full (no-EP) reference + moe_config_full = FusedMoEConfig( + num_experts=E, + experts_per_token=top_k, + hidden_dim=K, + intermediate_size=N, + num_local_experts=E, + num_logical_experts=E, + moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), + activation=MoEActivation.SILU, + in_dtype=torch.bfloat16, + device=device, + routing_method=RoutingMethodType.TopK, + max_num_tokens=512, + ) + experts_full = Int4EmulationTritonExperts( + moe_config_full, int4_w4a16_moe_quant_config(dummy_scale, dummy_scale) + ) + out_full = _run_emulation_forward( + experts_full, w13_all, w2_all, hidden_states, topk_weights, topk_ids, E, K, N + ) + + # EP sum + out_ep_sum = torch.zeros(num_tokens, K, dtype=torch.bfloat16, device=device) + for rank in range(ep_size): + start = rank * num_local + end = start + num_local + expert_map = _make_expert_map(E, start, end) + moe_config_ep = FusedMoEConfig( + num_experts=E, + experts_per_token=top_k, + hidden_dim=K, + intermediate_size=N, + num_local_experts=num_local, + num_logical_experts=E, + moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), + activation=MoEActivation.SILU, + in_dtype=torch.bfloat16, + device=device, + routing_method=RoutingMethodType.TopK, + max_num_tokens=512, + ) + experts_ep = Int4EmulationTritonExperts( + moe_config_ep, int4_w4a16_moe_quant_config(dummy_scale, dummy_scale) + ) + out_rank = _run_emulation_forward_ep( + experts_ep, + w13_all[start:end], + w2_all[start:end], + hidden_states, + topk_weights, + topk_ids, + E, + expert_map, + ) + out_ep_sum = out_ep_sum + out_rank + + assert torch.allclose(out_ep_sum, out_full, atol=1e-3), ( + f"EP sum max diff: {(out_ep_sum - out_full).abs().max().item():.6f}" + ) + + +@pytest.mark.parametrize("E, K, N, top_k, group_size, num_tokens", E2E_CONFIGS) +@pytest.mark.parametrize("fmt", ["gptq", "awq"]) +def test_emulation_output_close_to_bf16_reference( + E, K, N, top_k, group_size, num_tokens, fmt +): + """Emulation output is close to a direct BF16 MoE forward.""" + torch.manual_seed(11) + moe_config = _make_moe_config(E, K, N) + + w13_fp = torch.randn(E, K, 2 * N, dtype=torch.bfloat16, device=device) * 0.02 + w2_fp = torch.randn(E, N, K, dtype=torch.bfloat16, device=device) * 0.02 + + packed13_list, scales13_list, packed2_list, scales2_list = [], [], [], [] + for e in range(E): + q13, s13 = _quantize_sym(w13_fp[e].float(), group_size) + q2, s2 = _quantize_sym(w2_fp[e].float(), group_size) + if fmt == "gptq": + packed13_list.append(gptq_pack(q13, 4, K, 2 * N)) + packed2_list.append(gptq_pack(q2, 4, N, K)) + else: + packed13_list.append(awq_pack(q13, 4, K, 2 * N)) + packed2_list.append(awq_pack(q2, 4, N, K)) + scales13_list.append(s13) + scales2_list.append(s2) + + process_fn = ( + _process_weights_emulation_gptq + if fmt == "gptq" + else _process_weights_emulation_awq + ) + res = process_fn( + torch.stack(packed13_list), + torch.stack(packed2_list), + torch.stack(scales13_list), + torch.stack(scales2_list), + None, + None, + ) + w13_bf16, w2_bf16 = res[0], res[1] + + dummy_scale = torch.ones(1, dtype=torch.float16, device=device) + experts = Int4EmulationTritonExperts( + moe_config, int4_w4a16_moe_quant_config(dummy_scale, dummy_scale) + ) + + hidden_states = torch.randn(num_tokens, K, dtype=torch.bfloat16, device=device) + topk_weights = torch.softmax( + torch.randn(num_tokens, top_k, dtype=torch.float32, device=device), dim=-1 + ) + topk_ids = torch.stack( + [torch.randperm(E, device=device)[:top_k] for _ in range(num_tokens)] + ).to(torch.int32) + + out_emulation = _run_emulation_forward( + experts, w13_bf16, w2_bf16, hidden_states, topk_weights, topk_ids, E, K, N + ) + + ref = torch.zeros(num_tokens, K, dtype=torch.bfloat16, device=device) + for m in range(num_tokens): + acc = torch.zeros(K, dtype=torch.float32, device=device) + for k in range(top_k): + e = topk_ids[m, k].item() + w = topk_weights[m, k].item() + gate_up = hidden_states[m] @ w13_bf16[e].T + gate, up = gate_up.chunk(2) + act = F.silu(gate) * up + acc += w * (act @ w2_bf16[e].T).float() + ref[m] = acc.bfloat16() + + rel_l2 = ( + torch.norm(out_emulation.float() - ref.float()) + / torch.norm(ref.float()).clamp(min=1e-6) + ).item() + assert rel_l2 < 0.15, f"[{fmt}] relative L2 = {rel_l2:.4f} (threshold 0.15)" diff --git a/tests/kernels/quantization/test_triton_w4a16.py b/tests/kernels/quantization/test_triton_w4a16.py index 6502f5244292..42f163dea44a 100644 --- a/tests/kernels/quantization/test_triton_w4a16.py +++ b/tests/kernels/quantization/test_triton_w4a16.py @@ -302,3 +302,160 @@ class DummyLayer(torch.nn.Module): torch.testing.assert_close(layer.weight_packed, expected_w_kn8) torch.testing.assert_close(layer.weight_scale, expected_scales_gn) torch.testing.assert_close(layer.weight_zero_point, expected_zeros_gn8) + + +@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm only") +def test_triton_w4a16_process_weights_after_loading_keeps_gptq_qzeros_layout(): + if not torch.cuda.is_available(): + pytest.skip("CUDA/HIP device not available") + + from vllm.config import VllmConfig, set_current_vllm_config + from vllm.distributed import ( + ensure_model_parallel_initialized, + init_distributed_environment, + ) + from vllm.model_executor.kernels.linear.mixed_precision.MPLinearKernel import ( + MPLinearLayerConfig, + ) + from vllm.model_executor.parameter import ( + GroupQuantScaleParameter, + PackedvLLMParameter, + ) + from vllm.scalar_type import scalar_types + + with set_current_vllm_config(VllmConfig()): + init_distributed_environment( + world_size=1, + rank=0, + distributed_init_method="tcp://127.0.0.1:0", + local_rank=0, + ) + ensure_model_parallel_initialized(1, 1) + + set_random_seed(0) + + K, N = 256, 256 + G = 32 + + w_int4_kn = torch.randint(0, 16, (K, N), device=device, dtype=torch.int32) + w_gptq_k8n = _pack_int4_along_k_to_ckpt(w_int4_kn).t().contiguous() + scales_gptq_gn = 0.05 * torch.rand((K // G, N), device=device, dtype=torch.float16) + zeros_int4_gn = torch.randint(0, 16, (K // G, N), device=device, dtype=torch.int32) + zeros_gptq_gn8 = _pack_int4_along_n(zeros_int4_gn) + + config = MPLinearLayerConfig( + full_weight_shape=(K, N), + partition_weight_shape=(K, N), + weight_type=scalar_types.uint4, + act_type=torch.float16, + group_size=G, + zero_points=True, + has_g_idx=False, + ) + kernel = TritonW4A16LinearKernel( + config, + w_q_param_name="qweight", + w_s_param_name="scales", + w_zp_param_name="qzeros", + w_gidx_param_name=None, + ) + + weight_loader = lambda *args, **kwargs: None + + class DummyLayer(torch.nn.Module): + pass + + layer = DummyLayer() + layer.register_parameter( + "qweight", + PackedvLLMParameter( + data=w_gptq_k8n, + weight_loader=weight_loader, + input_dim=0, + output_dim=1, + packed_factor=8, + packed_dim=0, + ), + ) + layer.register_parameter( + "scales", + GroupQuantScaleParameter( + data=scales_gptq_gn, + weight_loader=weight_loader, + input_dim=0, + output_dim=1, + ), + ) + layer.register_parameter( + "qzeros", + PackedvLLMParameter( + data=zeros_gptq_gn8, + weight_loader=weight_loader, + input_dim=0, + output_dim=1, + packed_factor=8, + packed_dim=1, + ), + ) + + kernel.process_weights_after_loading(layer) + + expected_w_kn8 = _pack_int4_along_n(w_int4_kn) + + assert tuple(layer.qweight.shape) == (K, N // 8) + assert tuple(layer.scales.shape) == (K // G, N) + assert tuple(layer.qzeros.shape) == (K // G, N // 8) + + torch.testing.assert_close(layer.qweight, expected_w_kn8) + torch.testing.assert_close(layer.scales, scales_gptq_gn) + torch.testing.assert_close(layer.qzeros, zeros_gptq_gn8) + + +@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm only") +def test_triton_w4a16_symmetric_apply_ignores_qzeros(monkeypatch): + from vllm.model_executor.kernels.linear.mixed_precision.MPLinearKernel import ( + MPLinearLayerConfig, + ) + from vllm.scalar_type import scalar_types + + K, N, G = 256, 256, 32 + config = MPLinearLayerConfig( + full_weight_shape=(K, N), + partition_weight_shape=(K, N), + weight_type=scalar_types.uint4b8, + act_type=torch.float16, + group_size=G, + zero_points=False, + has_g_idx=False, + ) + kernel = TritonW4A16LinearKernel( + config, + w_q_param_name="qweight", + w_s_param_name="scales", + w_zp_param_name="qzeros", + w_gidx_param_name=None, + ) + + class DummyLayer(torch.nn.Module): + pass + + layer = DummyLayer() + layer.qweight = torch.empty((K, N // 8), device=device, dtype=torch.int32) + layer.scales = torch.empty((K // G, N), device=device, dtype=torch.float16) + layer.qzeros = torch.empty((1, 1), device=device, dtype=torch.int32) + + captured = {} + + def fake_gemm(*, a, b_q, scales, qzeros, group_size, zp_bias): + captured["qzeros"] = qzeros + captured["zp_bias"] = zp_bias + return torch.empty((a.shape[0], N), device=a.device, dtype=a.dtype) + + monkeypatch.setattr(triton_w4a16_module, "triton_w4a16_gemm", fake_gemm) + + x = torch.empty((2, K), device=device, dtype=torch.float16) + output = kernel.apply_weights(layer, x) + + assert output.shape == (2, N) + assert captured["qzeros"] is None + assert captured["zp_bias"] == scalar_types.uint4b8.bias diff --git a/tests/kernels/test_compressor_kv_cache.py b/tests/kernels/test_compressor_kv_cache.py index 74dc01472a8e..b8f1cb8bfa73 100644 --- a/tests/kernels/test_compressor_kv_cache.py +++ b/tests/kernels/test_compressor_kv_cache.py @@ -9,6 +9,7 @@ C) Indexer: head_dim=128 (all FP8), quant_block=128 D) DeepseekV4 Attention magnitude range: correctness across small/large values E) Indexer fused Triton kernel: compress+norm+rope+quant+insert + F) Indexer fused two-stage Triton kernel: head=512 cr>=128 (no-overlap) """ import math @@ -24,7 +25,9 @@ from vllm.models.deepseek_v4.common.ops.fused_compress_quant_cache import ( _fused_kv_compress_norm_rope_insert_indexer_attn, _fused_kv_compress_norm_rope_insert_indexer_mxfp4_attn, + _launch_two_stage_sparse_attn_compressor, ) +from vllm.platforms import current_platform from .test_fused_indexer_q_rope_quant import quantize_to_mxfp4 @@ -816,3 +819,130 @@ def test_cutedsl_full_cache_store(compress_ratio: int, store_fp8: bool): torch.testing.assert_close(actual.float(), ref_fp8.float(), rtol=0.0, atol=0.3) else: torch.testing.assert_close(actual.float(), ref.float(), rtol=3e-2, atol=3e-2) + + +# ── Test F: DeepseekV4 Attention two-stage split compressor (Triton) ───────── +# +# Same full pipeline as Test E (state-cache gather -> softmax-weighted compress +# -> RMSNorm -> GPT-J RoPE -> quant -> paged insert), but for the head=512 +# fp8_ds_mla layout via the two-stage split + + +@pytest.mark.skipif( + not current_platform.is_rocm(), + reason="two-stage split compressor is only enabled for ROCm at the moment", +) +@pytest.mark.parametrize("num_tokens", [1, 4, 8, 17]) +@pytest.mark.parametrize("kv_block_size", [16, 64]) +def test_fused_kv_insert_split(num_tokens: int, kv_block_size: int): + """Two-stage split compress+norm+rope+quant+insert for the head=512 KV cache.""" + HEAD_DIM = 512 + NOPE_DIM = 448 + ROPE_DIM = 64 + HEAD_BYTES = 584 # 448 fp8 + 128 bf16 + 8 uint8 scale + RMS_EPS = 1e-6 + FP8_MAX = 448.0 + QUANT_BLOCK = 64 + TOKEN_STRIDE = 576 + SCALE_DIM = 8 + STATE_BLOCK_SIZE = 8 # CompressorStateCache block_size for cr=128 + + device = "cuda" + torch.manual_seed(42) + compress_ratio = 128 + overlap = 0 # no overlap for cr=128 + coff = 1 + overlap + + num_pages = (compress_ratio * num_tokens - 1) // STATE_BLOCK_SIZE + 2 + state_cache = torch.randn( + num_pages, + STATE_BLOCK_SIZE, + 2 * coff * HEAD_DIM, # kv_state + score_state + dtype=torch.float32, + device=device, + ) + block_table = torch.arange(num_pages, dtype=torch.int32, device=device).unsqueeze(0) + token_to_req = torch.zeros(num_tokens, dtype=torch.int32, device=device) + slot_mapping = torch.arange(num_tokens, dtype=torch.int64, device=device) + positions = torch.arange( + compress_ratio - 1, + compress_ratio * num_tokens, + compress_ratio, + dtype=torch.int64, + device=device, + ) + rms_weight = torch.randn(HEAD_DIM, dtype=torch.bfloat16, device=device) + cos_sin_cache = torch.randn( + compress_ratio * num_tokens, ROPE_DIM, dtype=torch.float32, device=device + ) + + kv_n_blocks = (num_tokens + kv_block_size - 1) // kv_block_size + 1 + kv_cache = torch.zeros( + kv_n_blocks, kv_block_size, HEAD_BYTES, dtype=torch.uint8, device=device + ) + compress_scratch = torch.empty( + num_tokens, HEAD_DIM, dtype=torch.float32, device=device + ) + + _launch_two_stage_sparse_attn_compressor( + state_cache, + token_to_req, + positions, + slot_mapping, + block_table, + STATE_BLOCK_SIZE, + coff * HEAD_DIM, + compress_ratio, + cos_sin_cache, + kv_cache, + slot_mapping, + rms_weight, + RMS_EPS, + QUANT_BLOCK, + TOKEN_STRIDE, + SCALE_DIM, + HEAD_DIM, + ROPE_DIM, + num_tokens, + compress_scratch, + ) + + # PyTorch reference: compress -> RMSNorm -> GPT-J RoPE (pre-quant bf16 row). + ref = _reference_kv_compress_norm_rope( + state_cache, + block_table, + positions, + rms_weight, + cos_sin_cache, + compress_ratio, + overlap, + rms_eps=RMS_EPS, + fp8_max=FP8_MAX, + return_full_cache=True, + ) # [num_tokens, HEAD_DIM] bf16 + + # Dequant + gather the fp8_ds_mla cache back to bf16 (Test B op). + out = torch.zeros(1, num_tokens, HEAD_DIM, dtype=torch.bfloat16, device=device) + seq_lens = torch.tensor([num_tokens], dtype=torch.int32, device=device) + gather_block_table = torch.arange( + kv_n_blocks, dtype=torch.int32, device=device + ).unsqueeze(0) + dequantize_and_gather_k_cache( + out, kv_cache, seq_lens, None, gather_block_table, kv_block_size, offset=0 + ) + recovered = out[0, :num_tokens] + + # NoPE (first 448): FP8 quantized, expect UE8M0 error (same bound as Test A). + nope_diff = (recovered[:, :NOPE_DIM].float() - ref[:, :NOPE_DIM].float()).abs() + for t in range(num_tokens): + _, scales = _ue8m0_reference(ref[t, :NOPE_DIM].float(), QUANT_BLOCK, FP8_MAX) + max_allowed = 16.0 * scales.max().item() + token_diff = nope_diff[t].max().item() + assert token_diff <= max_allowed, ( + f"Token {t} nope diff {token_diff} exceeds max_allowed " + f"{max_allowed} (scale={scales.max().item()})" + ) + + # RoPE (last 64): stored as bf16. The kernel recomputes the rotation, so it + # is bf16-close to the reference rather than bit-exact (cf. test_cutedsl). + torch.testing.assert_close(recovered[:, NOPE_DIM:], ref[:, NOPE_DIM:]) diff --git a/tests/lora/test_gptoss_tp.py b/tests/lora/test_gptoss_tp.py index b9f99d7a8c2c..0f0d807441e9 100644 --- a/tests/lora/test_gptoss_tp.py +++ b/tests/lora/test_gptoss_tp.py @@ -42,6 +42,14 @@ ] +def reformat(text: str) -> str: + # Remove all spaces immediately before or after comma + text = ",".join(map(str.strip, text.split(","))) + # Remove duplicated blank spaces + text = " ".join(map(str.strip, text.split())) + return text + + def generate_and_test(llm: vllm.LLM, lora_path: str, lora_id: int) -> None: prompts = [ PROMPT_TEMPLATE.format( @@ -68,13 +76,18 @@ def generate_and_test(llm: vllm.LLM, lora_path: str, lora_id: int) -> None: generated_texts.append(generated_text) print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}") for i in range(len(EXPECTED_LORA_OUTPUT)): - generated = " ".join(generated_texts[i].split()) - expected = " ".join(EXPECTED_LORA_OUTPUT[i].split()) - assert generated.startswith(expected) + # The generated text may have different numbers of blank space, + # so reformat to compare. + compactGeneratedStr = reformat(generated_texts[i]) + compactExpectedStr = reformat(EXPECTED_LORA_OUTPUT[i]) + if not generated_texts[i].startswith( + EXPECTED_LORA_OUTPUT[i] + ) and not compactGeneratedStr.startswith(compactExpectedStr): + raise AssertionError( + f"Generated: {generated_texts[i]}, Expected: {EXPECTED_LORA_OUTPUT[i]}" + ) -# TODO: make the Mxfp4MoeBackend.TRITON spawn-safe. -# For now just use TRITON_UNFUSED kernel @pytest.mark.parametrize( "mxfp4_use_marlin", [ diff --git a/tests/lora/test_peft_helper.py b/tests/lora/test_peft_helper.py index e3035b00e9e0..23a8aad1a58c 100644 --- a/tests/lora/test_peft_helper.py +++ b/tests/lora/test_peft_helper.py @@ -22,6 +22,8 @@ {"modules_to_save": ["lm_head"]}, "only supports modules_to_save being None", ), + ("test_rank_zero", {"r": 0}, "must be a positive integer"), + ("test_rank_negative", {"r": -8}, "must be a positive integer"), ] @@ -97,3 +99,18 @@ def test_peft_helper_error( PEFTHelper.from_local_dir( test_dir, max_position_embeddings=4096 ).validate_legal(lora_config) + + +@pytest.mark.parametrize("bad_rank", [0, -1, -8]) +def test_peft_helper_invalid_rank_direct(bad_rank: int): + """Regression test: constructing a PEFTHelper with a non-positive rank + must raise a clear ValueError instead of crashing with an unrelated + ZeroDivisionError (r=0) or silently succeeding with a sign-flipped + scaling factor that validate_legal() never catches (r<0, since its only + rank check is the upper bound against max_lora_rank). + + Network-free: constructs PEFTHelper directly rather than going through + from_local_dir(), which needs an on-disk adapter_config.json. + """ + with pytest.raises(ValueError, match="must be a positive integer"): + PEFTHelper(r=bad_rank, lora_alpha=16, target_modules=["q_proj"]) diff --git a/tests/model_executor/layers/test_linear_load_weights.py b/tests/model_executor/layers/test_linear_load_weights.py new file mode 100644 index 000000000000..4fb95d5f106a --- /dev/null +++ b/tests/model_executor/layers/test_linear_load_weights.py @@ -0,0 +1,56 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Regression tests for MergedColumnParallelLinear/QKVParallelLinear +.load_weights() tolerating checkpoint tensors the layer never registered +as a param (e.g. bias when bias=False, or a GPTQ exporter's g_idx when +desc_act=False). + +Previously, ``getattr(self, name, self)`` used the layer itself as a +"not found" sentinel, so an unmatched ``bias``/``g_idx`` fell through to +``param.weight_loader(param, ...)`` with ``param`` bound to the whole +layer module -- crashing with ``AttributeError: ... has no attribute +'data'`` deep inside ``weight_loader`` instead of being skipped. +""" + +import torch + +from vllm.model_executor.layers.linear import ( + MergedColumnParallelLinear, + QKVParallelLinear, +) + + +def test_merged_column_parallel_load_weights_skips_unmatched_bias( + dist_init, default_vllm_config +): + layer = MergedColumnParallelLinear( + 4, [2, 2], bias=False, params_dtype=torch.float16 + ) + + loaded = list(layer.load_weights([("bias", torch.zeros(2))])) + + assert loaded == [] + + +def test_merged_column_parallel_load_weights_loads_matched_weight( + dist_init, default_vllm_config +): + layer = MergedColumnParallelLinear( + 4, [2, 2], bias=False, params_dtype=torch.float16 + ) + weight = torch.rand(2, 4, dtype=torch.float16) + weight.shard_id = 0 + + loaded = list(layer.load_weights([("weight", weight)])) + + assert loaded == ["weight"] + + +def test_qkv_parallel_load_weights_skips_unmatched_g_idx( + dist_init, default_vllm_config +): + layer = QKVParallelLinear(4, 2, 2, bias=False, params_dtype=torch.float16) + + loaded = list(layer.load_weights([("g_idx", torch.zeros(4, dtype=torch.int32))])) + + assert loaded == [] diff --git a/tests/models/language/pooling/test_token_classification.py b/tests/models/language/pooling/test_token_classification.py index 2524c7053ab7..2e95fb8f70a4 100644 --- a/tests/models/language/pooling/test_token_classification.py +++ b/tests/models/language/pooling/test_token_classification.py @@ -117,6 +117,46 @@ def test_modernbert_models( torch.testing.assert_close(hf_output, vllm_output, atol=3.2e-2, rtol=1e-3) +@pytest.mark.parametrize("model", ["Davlan/xlm-roberta-base-ner-hrl"]) +@pytest.mark.parametrize("dtype", ["float"]) +@torch.inference_mode +def test_xlm_roberta_models( + hf_runner, + vllm_runner, + example_prompts, + model: str, + dtype: str, +) -> None: + with vllm_runner(model, max_model_len=None, dtype=dtype) as vllm_model: + vllm_outputs = vllm_model.token_classify(example_prompts) + + # Use eager attention on ROCm to avoid HF Transformers flash attention + # accuracy issues: https://github.com/vllm-project/vllm/issues/30167 + hf_model_kwargs = {} + if current_platform.is_rocm(): + hf_model_kwargs["attn_implementation"] = "eager" + + with hf_runner( + model, + dtype=dtype, + auto_cls=AutoModelForTokenClassification, + model_kwargs=hf_model_kwargs, + ) as hf_model: + tokenizer = hf_model.tokenizer + hf_outputs = [] + for prompt in example_prompts: + inputs = tokenizer([prompt], return_tensors="pt") + inputs = hf_model.wrap_device(inputs) + output = hf_model.model(**inputs) + hf_outputs.append(softmax(output.logits[0])) + + # check logits difference + for hf_output, vllm_output in zip(hf_outputs, vllm_outputs): + hf_output = hf_output.detach().clone().cpu().float() + vllm_output = vllm_output.detach().clone().cpu().float() + torch.testing.assert_close(hf_output, vllm_output, atol=3.2e-2, rtol=1e-3) + + PRIVACY_FILTER_PROMPTS = [ "My name is Harry Potter.", "Email me at harry.potter@hogwarts.edu.", @@ -232,6 +272,11 @@ def test_bert_for_masked_lm( if current_platform.is_rocm(): hf_model_kwargs["attn_implementation"] = "eager" + # Run hf_runner reference with "highest" fp32 precision to match + # default behvior of vLLM. This is needed on ROCm since the + # pooling tests set matmul precision to "high" in conftest.py + prev_matmul_precision = torch.get_float32_matmul_precision() + torch.set_float32_matmul_precision("highest") with hf_runner( model, dtype=dtype, @@ -245,6 +290,7 @@ def test_bert_for_masked_lm( inputs = hf_model.wrap_device(inputs) output = hf_model.model(**inputs) hf_outputs.append(softmax(output.logits[0])) + torch.set_float32_matmul_precision(prev_matmul_precision) # Compare the per-token vocabulary distributions position by position. for hf_output, vllm_output in zip(hf_outputs, vllm_outputs): diff --git a/tests/models/registry.py b/tests/models/registry.py index 4e8fe96fdbf6..cecd2d965ae4 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -713,6 +713,12 @@ def check_available_online( "openai/privacy-filter", min_transformers_version="5.6.0.dev0", ), + "RobertaForTokenClassification": _HfExamplesInfo( + "Jean-Baptiste/roberta-large-ner-english" + ), + "XLMRobertaForTokenClassification": _HfExamplesInfo( + "Davlan/xlm-roberta-base-ner-hrl" + ), } _SEQUENCE_CLASSIFICATION_EXAMPLE_MODELS = { @@ -1360,7 +1366,15 @@ def check_available_online( "HuggingFaceTB/SmolVLM2-2.2B-Instruct" ), "Step3VLForConditionalGeneration": _HfExamplesInfo( - "stepfun-ai/step3", trust_remote_code=True + "stepfun-ai/step3", + trust_remote_code=True, + max_transformers_version="5.3", + transformers_version_reason={ + "hf": ( + "Transformers v5.4 removed the ignore_keys param from " + "validate_rope(); vLLM has vendored the config and is unaffected" + ) + }, ), "StepVLForConditionalGeneration": _HfExamplesInfo( "stepfun-ai/Step3-VL-10B", trust_remote_code=True diff --git a/tests/test_envs.py b/tests/test_envs.py index d4d120ecee51..56c04dd6f2e2 100644 --- a/tests/test_envs.py +++ b/tests/test_envs.py @@ -36,6 +36,18 @@ def test_nixl_side_channel_host_is_not_compile_factor( assert "VLLM_NIXL_SIDE_CHANNEL_HOST" not in envs.compile_factors() +def test_p2p_side_channel_defaults_and_override(monkeypatch: pytest.MonkeyPatch): + monkeypatch.delenv("VLLM_P2P_SIDE_CHANNEL_HOST", raising=False) + monkeypatch.delenv("VLLM_P2P_SIDE_CHANNEL_PORT", raising=False) + assert envs.VLLM_P2P_SIDE_CHANNEL_HOST == "localhost" + assert envs.VLLM_P2P_SIDE_CHANNEL_PORT == 5710 + + monkeypatch.setenv("VLLM_P2P_SIDE_CHANNEL_HOST", "10.0.0.20") + monkeypatch.setenv("VLLM_P2P_SIDE_CHANNEL_PORT", "5799") + assert envs.VLLM_P2P_SIDE_CHANNEL_HOST == "10.0.0.20" + assert envs.VLLM_P2P_SIDE_CHANNEL_PORT == 5799 + + def test_getattr_with_cache(monkeypatch: pytest.MonkeyPatch): monkeypatch.setenv("VLLM_HOST_IP", "1.1.1.1") monkeypatch.setenv("VLLM_PORT", "1234") diff --git a/tests/test_jit_monitor.py b/tests/test_jit_monitor.py index 9f3285ddec0e..50261a479d92 100644 --- a/tests/test_jit_monitor.py +++ b/tests/test_jit_monitor.py @@ -320,6 +320,26 @@ def compile_fn(*args, **kwargs): with pytest.raises(RuntimeError, match="CuTeDSL JIT compilation"): cute.compile(lambda: None, "arg", option=True) + def test_subscripted_compile_is_monitored(self): + """``cute.compile[options](...)`` (flashinfer >= 0.6.14) must work.""" + + class FakeCompileCallable: + def __getitem__(self, options): + return self + + def __call__(self, *args, **kwargs): + return "compiled" + + with _patch_jit_modules(_make_fake_knobs(), cute_compile=FakeCompileCallable()): + import cutlass.cute as cute + + jit_monitor.activate() + with mock.patch.object(jit_monitor.logger, "warning_once") as warning_once: + result = cute.compile[("opt_level", 3)](lambda: None, "arg") + + assert result == "compiled" + warning_once.assert_called_once() + class TestTileLangHook: def test_jit_kernel_logs_warning(self): diff --git a/tests/tool_use/test_responses_request_validations.py b/tests/tool_use/test_responses_request_validations.py index 63a1828c5009..59b156b76a18 100644 --- a/tests/tool_use/test_responses_request_validations.py +++ b/tests/tool_use/test_responses_request_validations.py @@ -4,7 +4,10 @@ import pytest from pydantic import ValidationError -from vllm.entrypoints.openai.responses.protocol import ResponsesRequest +from vllm.entrypoints.openai.responses.protocol import ( + ResponsesRequest, + ResponsesResponse, +) SAMPLE_TOOL = { "type": "function", @@ -182,3 +185,42 @@ def test_responses_request_empty_tools_named_tool_choice(): "tool_choice": NAMED_TOOL_CHOICE, } ) + + +# Regression tests for parallel_tool_calls=null crash in Responses API +# (from_request() passed None to ResponsesResponse.parallel_tool_calls, +# a non-optional bool field, causing a Pydantic ValidationError during +# response construction) +@pytest.mark.parametrize( + "value,expected", + [ + (True, True), + (False, False), + (None, True), # null must resolve to the documented default (true) + ], +) +def test_responses_response_parallel_tool_calls_null_resolves_to_default( + value, expected +): + request = ResponsesRequest.model_validate( + {"input": "Hello", "model": "test-model", "parallel_tool_calls": value} + ) + sampling_params = request.to_sampling_params(default_max_tokens=16) + r = ResponsesResponse.from_request( + request=request, + sampling_params=sampling_params, + model_name="test-model", + created_time=0, + output=[], + status="completed", + usage=None, + ) + assert r.parallel_tool_calls == expected + + +def test_responses_request_parallel_tool_calls_null_accepted(): + """Client sending null must be accepted at request validation time.""" + req = ResponsesRequest.model_validate( + {"input": "Hello", "model": "test-model", "parallel_tool_calls": None} + ) + assert req.parallel_tool_calls is None diff --git a/tests/v1/core/test_scheduler.py b/tests/v1/core/test_scheduler.py index 157170400b0d..da1e0c5e76c5 100644 --- a/tests/v1/core/test_scheduler.py +++ b/tests/v1/core/test_scheduler.py @@ -43,6 +43,65 @@ pytestmark = pytest.mark.cpu_test +def test_make_scheduled_encoder_input_stats_output_embeddings(): + scheduler = create_scheduler() + mm_features = [ + MultiModalFeatureSpec( + data=MultiModalKwargsItem.dummy(), + modality="image", + identifier="image-0", + mm_position=PlaceholderRange(offset=0, length=196), + ), + MultiModalFeatureSpec( + data=MultiModalKwargsItem.dummy(), + modality="video", + identifier="video-0", + mm_position=PlaceholderRange(offset=200, length=196), + ), + MultiModalFeatureSpec( + data=MultiModalKwargsItem.dummy(), + modality="audio", + identifier="audio-0", + mm_position=PlaceholderRange(offset=400, length=49), + ), + ] + scheduler.requests["req"] = Mock(mm_features=mm_features) + + stats = scheduler._make_scheduled_encoder_input_stats({"req": [0, 1, 2]}) + + assert stats is not None + assert stats.num_inputs == 3 + assert stats.output_tokens == 441 + + +def test_scheduled_encoder_input_stats_disabled_without_iteration_logging( + monkeypatch: pytest.MonkeyPatch, +): + scheduler = create_scheduler() + make_stats = Mock(side_effect=AssertionError("stats should not be computed")) + monkeypatch.setattr(scheduler, "_make_scheduled_encoder_input_stats", make_stats) + + scheduler_output = scheduler.schedule() + + make_stats.assert_not_called() + assert scheduler_output.scheduled_encoder_input_stats is None + + +def test_scheduled_encoder_input_stats_disabled_without_log_stats( + monkeypatch: pytest.MonkeyPatch, +): + scheduler = create_scheduler() + scheduler.log_stats = False + scheduler.observability_config.enable_logging_iteration_details = True + make_stats = Mock(side_effect=AssertionError("stats should not be computed")) + monkeypatch.setattr(scheduler, "_make_scheduled_encoder_input_stats", make_stats) + + scheduler_output = scheduler.schedule() + + make_stats.assert_not_called() + assert scheduler_output.scheduled_encoder_input_stats is None + + def test_add_requests(): scheduler = create_scheduler() requests = create_requests(num_requests=10) @@ -108,6 +167,29 @@ def test_schedule(enable_prefix_caching: bool, prompt_logprobs: int | None): assert scheduler.running[i] == request +def test_scheduler_stats_route_to_existing_output_client(): + scheduler = create_scheduler() + request = create_requests(num_requests=1)[0] + request.client_index = 1 + scheduler.add_request(request) + + scheduler_output = scheduler.schedule() + model_output = ModelRunnerOutput( + req_ids=[request.request_id], + req_id_to_index={request.request_id: 0}, + sampled_token_ids=[[1000]], + logprobs=None, + prompt_logprobs_dict={}, + pooler_output=[], + ) + + engine_core_outputs = scheduler.update_from_output(scheduler_output, model_output) + + assert 0 not in engine_core_outputs + assert engine_core_outputs[1].scheduler_stats is not None + assert len(engine_core_outputs[1].outputs) == 1 + + def test_schedule_multimodal_requests(): scheduler = create_scheduler(model="llava-hf/llava-1.5-7b-hf") mm_positions = [[PlaceholderRange(offset=i, length=100)] for i in range(10)] diff --git a/tests/v1/determinism/test_batch_invariance.py b/tests/v1/determinism/test_batch_invariance.py index fb12ffd17063..b2706ed89b7b 100644 --- a/tests/v1/determinism/test_batch_invariance.py +++ b/tests/v1/determinism/test_batch_invariance.py @@ -11,11 +11,13 @@ TEST_MODEL, _extract_step_logprobs, _random_prompt, + skip_if_not_cuda, skip_unsupported, ) import vllm.envs as envs from vllm import LLM, SamplingParams +from vllm.platforms import current_platform @skip_unsupported @@ -49,6 +51,11 @@ def test_v1_generation_is_deterministic_across_batch_sizes_with_needle( seed. - Keep max_tokens and max_model_len bounded for speed and memory use. """ + # Not all batch-invariant kernels are registered on XPU yet + # (e.g. attention, custom ops), so e2e determinism is not guaranteed. + if current_platform.is_xpu(): + pytest.xfail("Not all batch-invariant kernels registered on XPU yet") + seed = int(os.getenv("VLLM_TEST_SEED", "12345")) random.seed(seed) @@ -157,6 +164,11 @@ def test_logprobs_bitwise_batch_invariance_bs1_vs_bsN( block_m, block_n, ): + # Not all batch-invariant kernels are registered on XPU yet + # (e.g. attention, custom ops), so e2e determinism is not guaranteed. + if current_platform.is_xpu(): + pytest.xfail("Not all batch-invariant kernels registered on XPU yet") + seed = int(os.getenv("VLLM_TEST_SEED", "12345")) random.seed(seed) tp_size = int(os.getenv("VLLM_TEST_TP_SIZE", "1")) @@ -641,7 +653,7 @@ def test_logprobs_without_batch_invariance_should_fail( pytest.fail(fail_msg) -@skip_unsupported +@skip_if_not_cuda @pytest.mark.parametrize("backend", ["FLASH_ATTN"]) def test_decode_logprobs_match_prefill_logprobs( backend, diff --git a/tests/v1/determinism/test_nvfp4_batch_invariant.py b/tests/v1/determinism/test_nvfp4_batch_invariant.py index d7a1c9e84042..dafd1a16444d 100644 --- a/tests/v1/determinism/test_nvfp4_batch_invariant.py +++ b/tests/v1/determinism/test_nvfp4_batch_invariant.py @@ -9,7 +9,7 @@ from utils import ( _extract_step_logprobs, _random_prompt, - skip_unsupported, + skip_if_not_cuda, ) from vllm import LLM, SamplingParams @@ -40,7 +40,7 @@ def _make_llm(max_num_seqs: int, backend: str) -> LLM: ) -@skip_unsupported +@skip_if_not_cuda @pytest.mark.parametrize("backend", ["FLASH_ATTN"]) def test_dense_nvfp4_generation_is_deterministic_across_batch_sizes_e2e(backend): seed = int(os.getenv("VLLM_TEST_SEED", "12345")) diff --git a/tests/v1/determinism/test_online_batch_invariance.py b/tests/v1/determinism/test_online_batch_invariance.py index 2bebb2dca533..de80e4918c1e 100644 --- a/tests/v1/determinism/test_online_batch_invariance.py +++ b/tests/v1/determinism/test_online_batch_invariance.py @@ -17,7 +17,7 @@ import openai import pytest -from utils import BACKENDS, TEST_MODEL, _random_prompt, skip_unsupported +from utils import BACKENDS, TEST_MODEL, _random_prompt, skip_if_not_cuda from tests.utils import RemoteOpenAIServer @@ -133,7 +133,7 @@ def _compare_bs1_vs_bsn_single_process( ) -@skip_unsupported +@skip_if_not_cuda @pytest.mark.parametrize("backend", BACKENDS) def test_logprobs_bitwise_batch_invariance_bs1_vs_bsN( backend: str, diff --git a/tests/v1/determinism/test_rms_norm_batch_invariant.py b/tests/v1/determinism/test_rms_norm_batch_invariant.py index 7fbf8f046100..5b3b7a8758b1 100644 --- a/tests/v1/determinism/test_rms_norm_batch_invariant.py +++ b/tests/v1/determinism/test_rms_norm_batch_invariant.py @@ -9,7 +9,7 @@ import pytest import torch -from utils import skip_unsupported +from utils import skip_if_not_cuda, skip_unsupported from vllm.model_executor.layers.batch_invariant import ( rms_norm_batch_invariant, @@ -20,7 +20,7 @@ DEVICE_TYPE = current_platform.device_type -@skip_unsupported +@skip_if_not_cuda @pytest.mark.parametrize("batch_size", [1, 4, 16, 64]) @pytest.mark.parametrize("hidden_size", [512, 2048, 4096, 8192]) @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) @@ -73,7 +73,7 @@ def test_rms_norm_batch_invariant_vs_standard( ) -@skip_unsupported +@skip_if_not_cuda @pytest.mark.parametrize("hidden_size", [512, 4096]) @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) @pytest.mark.parametrize("eps", [1e-6]) @@ -166,7 +166,7 @@ def fused_add_rms_norm(x, residual, w, e) -> tuple[torch.Tensor, torch.Tensor]: ) -@skip_unsupported +@skip_if_not_cuda @pytest.mark.parametrize("batch_size", [1, 16, 128]) @pytest.mark.parametrize("seq_len", [1, 32, 512]) @pytest.mark.parametrize("hidden_size", [2048, 4096]) @@ -210,7 +210,7 @@ def test_rms_norm_3d_input( ) -@skip_unsupported +@skip_if_not_cuda def test_rms_norm_numerical_stability(default_vllm_config): """ Test RMS norm numerical stability with extreme values. @@ -303,7 +303,7 @@ def test_rms_norm_formula(default_vllm_config): ) -@skip_unsupported +@skip_if_not_cuda @pytest.mark.parametrize("hidden_size", [128, 1024, 4096, 16384]) def test_rms_norm_different_hidden_sizes(default_vllm_config, hidden_size: int): """ @@ -377,6 +377,35 @@ def test_rms_norm_determinism(default_vllm_config): ) +@skip_unsupported +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +def test_rms_norm_batch_invariance(dtype): + """Same row gives identical rms_norm result regardless of batch neighbors. + + This verifies that the output for a given row is independent of what other + rows are present in the batch — the core batch-invariance property. + """ + device = torch.device(DEVICE_TYPE) + torch.manual_seed(42) + hidden_size = 2048 + eps = 1e-6 + + weight = torch.randn(hidden_size, dtype=dtype, device=device) + row = torch.randn(1, hidden_size, dtype=dtype, device=device) + + # Compute rms_norm on the single row alone + out_single = rms_norm_batch_invariant(row, weight, eps=eps) + + # Embed the same row in a larger batch with random neighbors + batch = torch.randn(8, hidden_size, dtype=dtype, device=device) + batch[4] = row[0] + out_batch = rms_norm_batch_invariant(batch, weight, eps=eps) + + assert torch.equal(out_single[0], out_batch[4]), ( + "rms_norm output for a row differs when batch context changes" + ) + + if __name__ == "__main__": # Run a quick smoke test print("Running quick smoke test of RMS norm implementations...") diff --git a/tests/v1/determinism/utils.py b/tests/v1/determinism/utils.py index bbef61477232..f03ea05b4331 100644 --- a/tests/v1/determinism/utils.py +++ b/tests/v1/determinism/utils.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import os import random +from typing import NamedTuple import pytest import torch @@ -11,36 +12,61 @@ from vllm.transformers_utils.model_arch_config_convertor import ( ModelArchConfigConvertorBase, ) +from vllm.triton_utils import HAS_TRITON from vllm.v1.attention.backends.fa_utils import flash_attn_supports_mla -skip_unsupported = pytest.mark.skipif( - not (current_platform.is_cuda() and current_platform.has_device_capability(80)), - # Supports testing on Ampere and Ada Lovelace devices. - # Note: For devices with SM < 90, batch invariance does not support CUDA Graphs. - reason="Requires CUDA and >= Ampere (SM80)", -) -DEFAULT_MODEL = "Qwen/Qwen3-1.7B" -TEST_MODEL = os.getenv("VLLM_TEST_MODEL", DEFAULT_MODEL) +class DeviceConfig(NamedTuple): + available: bool + backends: list[str] -BACKENDS: list[str] = [ - "FLASH_ATTN", - "TRITON_ATTN", - "FLEX_ATTENTION", -] -# FlashInfer temporarily disabled due to invariant CTA sizes. -# See FlashInfer issue #2424 -# if has_flashinfer(): -# BACKENDS.append("FLASHINFER") +# Maps each device to its availability and supported backends. +DEVICE_BACKENDS: dict[str, DeviceConfig] = { + "cuda": DeviceConfig( + available=current_platform.is_cuda() + and current_platform.has_device_capability(80), + # FlashInfer backend temporarily disabled due to invariant CTA sizes. + # See FlashInfer issue #2424 + backends=["FLASH_ATTN", "TRITON_ATTN", "FLEX_ATTENTION"], + ), + "xpu": DeviceConfig( + available=current_platform.is_xpu() and HAS_TRITON, + backends=["TRITON_ATTN"], + ), +} -# only run MLA backends when the requested test model is itself an MLA model. +DEFAULT_MODEL = "Qwen/Qwen3-1.7B" +TEST_MODEL = os.getenv("VLLM_TEST_MODEL", DEFAULT_MODEL) + +# Override backends for MLA models (MLA only supported on CUDA). if os.getenv("VLLM_TEST_MODEL"): config = get_config(TEST_MODEL, trust_remote_code=False) if ModelArchConfigConvertorBase(config, config.get_text_config()).is_deepseek_mla(): - BACKENDS = ["TRITON_MLA"] - if flash_attn_supports_mla(): - BACKENDS.append("FLASH_ATTN_MLA") + DEVICE_BACKENDS["cuda"] = DeviceConfig( + available=DEVICE_BACKENDS["cuda"].available, + backends=["TRITON_MLA"] + + (["FLASH_ATTN_MLA"] if flash_attn_supports_mla() else []), + ) + DEVICE_BACKENDS["xpu"] = DeviceConfig( + available=DEVICE_BACKENDS["xpu"].available, + backends=[], + ) + +# Only include backends for devices that are actually available. +BACKENDS: list[str] = sorted( + {b for cfg in DEVICE_BACKENDS.values() if cfg.available for b in cfg.backends} +) + +skip_unsupported = pytest.mark.skipif( + not any(cfg.available for cfg in DEVICE_BACKENDS.values()), + reason="Requires CUDA >= Ampere (SM80) or Intel XPU with Triton", +) + +skip_if_not_cuda = pytest.mark.skipif( + not DEVICE_BACKENDS["cuda"].available, + reason="Requires CUDA >= Ampere (SM80)", +) def _random_prompt(min_words: int = 1024, max_words: int = 1024 * 2) -> str: diff --git a/tests/v1/engine/test_iteration_logging.py b/tests/v1/engine/test_iteration_logging.py new file mode 100644 index 000000000000..5a08308163ec --- /dev/null +++ b/tests/v1/engine/test_iteration_logging.py @@ -0,0 +1,88 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import time +from types import SimpleNamespace + +from vllm.v1.engine import EngineCoreOutputs +from vllm.v1.engine.core import EngineCore +from vllm.v1.metrics.stats import SchedulerIterationDetails, SchedulerStats + + +class FakeEngineCore: + def _make_iteration_details_stats( + self, iteration_details: SchedulerIterationDetails + ) -> SchedulerStats: + return SchedulerStats(iteration_details=iteration_details) + + +def make_iteration_details() -> SchedulerIterationDetails: + return SchedulerIterationDetails( + iteration_index=1, + num_ctx_requests=2, + num_ctx_tokens=3, + num_generation_requests=4, + num_generation_tokens=5, + elapsed_ms=6.7, + ) + + +def make_fake_engine(log_stats: bool = True) -> SimpleNamespace: + return SimpleNamespace( + log_stats=log_stats, + vllm_config=SimpleNamespace( + observability_config=SimpleNamespace( + enable_logging_iteration_details=True, + ) + ), + ) + + +def test_capture_iteration_details_disabled_without_log_stats(): + engine = make_fake_engine(log_stats=False) + + with EngineCore.capture_iteration_details(engine, None) as iteration_details: + assert iteration_details is None + + assert not hasattr(engine, "_iteration_index") + + +def test_capture_iteration_details_fills_elapsed_time(): + engine = make_fake_engine() + + with EngineCore.capture_iteration_details(engine, None) as iteration_details: + assert iteration_details is not None + assert iteration_details.elapsed_ms == 0.0 + assert iteration_details.is_dummy + time.sleep(0.001) + + assert iteration_details is not None + assert iteration_details.elapsed_ms > 0.0 + assert engine._iteration_index == 1 + + +def test_attach_iteration_details_uses_existing_output(): + iteration_details = make_iteration_details() + outputs = { + 2: EngineCoreOutputs(scheduler_stats=SchedulerStats()), + 1: EngineCoreOutputs(scheduler_stats=SchedulerStats()), + } + + EngineCore._attach_iteration_details(FakeEngineCore(), outputs, iteration_details) + + assert 0 not in outputs + assert outputs[2].scheduler_stats is not None + assert outputs[2].scheduler_stats.iteration_details == iteration_details + assert outputs[1].scheduler_stats is not None + assert outputs[1].scheduler_stats.iteration_details is None + + +def test_attach_iteration_details_falls_back_to_client_zero_without_outputs(): + iteration_details = make_iteration_details() + outputs: dict[int, EngineCoreOutputs] = {} + + EngineCore._attach_iteration_details(FakeEngineCore(), outputs, iteration_details) + + assert set(outputs) == {0} + assert outputs[0].scheduler_stats is not None + assert outputs[0].scheduler_stats.iteration_details == iteration_details diff --git a/tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh b/tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh index a357128d3ce4..77235582fe13 100755 --- a/tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh +++ b/tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh @@ -140,12 +140,13 @@ run_tests_for_model() { # Start prefill instances for i in $(seq 0 $((NUM_PREFILL_INSTANCES-1))); do # Calculate GPU ID - we'll distribute across available GPUs - GPU_ID=$((i % $(get_num_gpus))) - NEXT_GPU=${GPU_ID} + GPU_START=$((i % $(get_num_gpus))) + GPU_ID=$GPU_START + NEXT_GPU=$GPU_START # Reserve TP*PP GPUs for the prefiller (TP shards across PP stages). PREFILLER_WORLD_SIZE=$((PREFILLER_TP_SIZE * PREFILLER_PP_SIZE)) for (( j=1; j < PREFILLER_WORLD_SIZE; j++ )); do - NEXT_GPU=$(((GPU_ID + j) % $(get_num_gpus))) + NEXT_GPU=$(((GPU_START + j) % $(get_num_gpus))) GPU_ID="${GPU_ID},${NEXT_GPU}" done @@ -195,10 +196,12 @@ run_tests_for_model() { # Start decode instances for i in $(seq 0 $((NUM_DECODE_INSTANCES-1))); do # Calculate GPU ID - we'll distribute across available GPUs, starting from after prefill GPUs - GPU_ID=$(((i + NEXT_GPU + 1) % $(get_num_gpus))) + DECODE_START=$(((i + NEXT_GPU + 1) % $(get_num_gpus))) + GPU_ID=$DECODE_START + NEXT_GPU=$DECODE_START # If DECODER_TP_SIZE is more than 1 for (( j=1; j < DECODER_TP_SIZE; j++ )); do - NEXT_GPU=$(((GPU_ID + j) % $(get_num_gpus))) + NEXT_GPU=$(((DECODE_START + j) % $(get_num_gpus))) GPU_ID="${GPU_ID},${NEXT_GPU}" done # Calculate port number (base port + instance number) diff --git a/tests/v1/kv_connector/nixl_integration/test_accuracy.py b/tests/v1/kv_connector/nixl_integration/test_accuracy.py index bb68b7a57241..711f27d20945 100644 --- a/tests/v1/kv_connector/nixl_integration/test_accuracy.py +++ b/tests/v1/kv_connector/nixl_integration/test_accuracy.py @@ -26,6 +26,7 @@ "Qwen/Qwen3.5-0.8B": 0.33, "google/gemma-4-E2B-it": 0.485, "ai21labs/AI21-Jamba2-3B": 0.74, + "deepseek-ai/DeepSeek-V4-Flash": 0.95, } SIMPLE_PROMPT = ( diff --git a/tests/v1/kv_connector/unit/test_multi_connector.py b/tests/v1/kv_connector/unit/test_multi_connector.py index 441e8e3676d6..14ea6759f4d5 100644 --- a/tests/v1/kv_connector/unit/test_multi_connector.py +++ b/tests/v1/kv_connector/unit/test_multi_connector.py @@ -318,8 +318,12 @@ def test_multi_example_connector_consistency(): # connector (first nonzero match is chosen), so update_state_after_alloc # will report those external tokens on that one. Other connectors still # receive the request's real blocks but with 0 external tokens. - storage1_scheduler_events = _ignore_event_collection(events["storage1-SCHEDULER"]) - storage2_scheduler_events = _ignore_event_collection(events["storage2-SCHEDULER"]) + storage1_scheduler_events = _events_from_request( + _ignore_event_collection(events["storage1-SCHEDULER"]) + ) + storage2_scheduler_events = _events_from_request( + _ignore_event_collection(events["storage2-SCHEDULER"]) + ) assert storage1_scheduler_events[:4] == [ "on_new_request", "get_num_new_matched_tokens 0", @@ -348,8 +352,12 @@ def test_multi_example_connector_consistency(): # return 0 from the first connector, while the second connector has a hit. # Both connectors receive the request's real blocks, but only the chosen # (second) connector reports external tokens. - storage1_scheduler_events = _ignore_event_collection(events["storage1-SCHEDULER"]) - storage2_scheduler_events = _ignore_event_collection(events["storage2-SCHEDULER"]) + storage1_scheduler_events = _events_from_request( + _ignore_event_collection(events["storage1-SCHEDULER"]) + ) + storage2_scheduler_events = _events_from_request( + _ignore_event_collection(events["storage2-SCHEDULER"]) + ) assert storage1_scheduler_events[:4] == [ "on_new_request", "get_num_new_matched_tokens 0", @@ -375,6 +383,16 @@ def _ignore_event_collection(events: list[str]) -> list[str]: return [event for event in events if event not in ignored] +def _events_from_request(events: list[str]) -> list[str]: + # The async engine-core can emit a trailing build_connector_meta from the + # previous generation's final (idle) scheduler step. Depending on timing it + # may be flushed into this window ahead of on_new_request, so anchor the + # comparison on the new request's first event to avoid a flaky ordering. + if "on_new_request" in events: + return events[events.index("on_new_request") :] + return events + + def get_connector_events() -> dict[str, list[str]]: # Read in connector events and reset the files. import glob diff --git a/tests/v1/kv_offload/tiering/p2p/p2p_connector_proxy.py b/tests/v1/kv_offload/tiering/p2p/p2p_connector_proxy.py index 10566f40e6f0..2f0df8990233 100644 --- a/tests/v1/kv_offload/tiering/p2p/p2p_connector_proxy.py +++ b/tests/v1/kv_offload/tiering/p2p/p2p_connector_proxy.py @@ -14,7 +14,7 @@ --port 8192 \ --prefiller-host 127.0.0.1 --prefiller-port 8100 \ --decoder-host 127.0.0.1 --decoder-port 8200 \ - --p2p-connector-host 127.0.0.1 --p2p-connector-port 7777 + --p2p-connector-host 127.0.0.1 --p2p-connector-port 5710 """ import argparse @@ -73,14 +73,23 @@ async def lifespan(app: FastAPI): app.state.prefill_iterator = itertools.cycle(range(len(app.state.prefill_clients))) app.state.decode_iterator = itertools.cycle(range(len(app.state.decode_clients))) + # Round-robin over each role's data-parallel replicas (independent of the + # instance iterators above). For dp_size=1 these yield only rank 0. + app.state.prefill_dp_iterator = itertools.cycle( + range(global_args.prefiller_dp_size) + ) + app.state.decode_dp_iterator = itertools.cycle(range(global_args.decoder_dp_size)) + mode = "decoder-first" if global_args.decoder_first else "prefiller-first" pd_host = global_args.p2p_connector_host pd_port = global_args.p2p_connector_port + n_pref = len(app.state.prefill_clients) + n_dec = len(app.state.decode_clients) print( f"Proxy ready [{mode}]: " - f"{len(app.state.prefill_clients)} prefiller(s), " - f"{len(app.state.decode_clients)} decoder(s). " - f"P2PConnector at {pd_host}:{pd_port}" + f"{n_pref} prefiller(s) x dp={global_args.prefiller_dp_size}, " + f"{n_dec} decoder(s) x dp={global_args.decoder_dp_size}. " + f"P2PConnector base at {pd_host}:{pd_port}" ) yield @@ -111,8 +120,9 @@ def parse_args(): p.add_argument( "--p2p-connector-port", type=int, - default=7777, - help="Port of the prefiller's P2PConnector ZMQ socket", + default=int(os.getenv("VLLM_P2P_SIDE_CHANNEL_PORT", "5710")), + help="Port of the prefiller's P2PConnector ZMQ socket " + "(default: $VLLM_P2P_SIDE_CHANNEL_PORT or 5710)", ) # P2PConnector coordinates of the decoder — injected into prefill requests # so the prefiller's submit_store can resolve the peer to push KV to. @@ -125,8 +135,9 @@ def parse_args(): p.add_argument( "--decoder-p2p-connector-port", type=int, - default=7778, - help="Port of the decoder's P2PConnector ZMQ socket", + default=int(os.getenv("VLLM_P2P_SIDE_CHANNEL_PORT", "5710")) + 1, + help="Port of the decoder's P2PConnector ZMQ socket " + "(default: $VLLM_P2P_SIDE_CHANNEL_PORT + 1 or 5711)", ) p.add_argument( "--decoder-first", @@ -134,6 +145,23 @@ def parse_args(): help="Send decode request before prefill so decoder is already " "waiting when KV blocks arrive (decoder-first mode)", ) + p.add_argument( + "--prefiller-dp-size", + type=int, + default=1, + help="Data-parallel replica count of the prefiller. When >1 the proxy " + "round-robins prefill across ranks via the X-data-parallel-rank header " + "and injects remote_port = p2p-connector-port + rank into the decode " + "request so the decoder pulls KV from that replica's control socket. " + "Assumes a single prefiller instance (one host/port fronting N replicas).", + ) + p.add_argument( + "--decoder-dp-size", + type=int, + default=1, + help="Data-parallel replica count of the decoder. When >1 the proxy " + "round-robins decode across ranks via the X-data-parallel-rank header.", + ) args = p.parse_args() if len(args.prefiller_hosts) != len(args.prefiller_ports): raise ValueError("Prefiller host/port count mismatch") @@ -150,6 +178,23 @@ def _get_next(app, service: str): return app.state.decode_clients[next(app.state.decode_iterator)] +def _next_dp_rank(app, service: str): + """Advance the round-robin DP cursor for a role. + + Returns (rank, header_rank). ``rank`` (0..dp_size-1) is always used for the + P2P remote_port arithmetic; ``header_rank`` is the same value when the role + has dp_size>1 and None otherwise (so dp=1 sends no header, matching the + single-replica behavior). + """ + if service == "prefill": + rank = next(app.state.prefill_dp_iterator) + dp_size = global_args.prefiller_dp_size + else: + rank = next(app.state.decode_dp_iterator) + dp_size = global_args.decoder_dp_size + return rank, (rank if dp_size > 1 else None) + + def _auth_headers(request_id: str) -> dict: headers: dict = {"X-Request-Id": request_id} api_key = os.environ.get("OPENAI_API_KEY", "") @@ -158,7 +203,7 @@ def _auth_headers(request_id: str) -> dict: return headers -async def _prefill(client_info, endpoint, req_data, request_id): +async def _prefill(client_info, endpoint, req_data, request_id, dp_rank=None): """Send a prefill-only request (max_tokens=1) to the prefiller.""" data = req_data.copy() data["kv_transfer_params"] = { @@ -174,14 +219,21 @@ async def _prefill(client_info, endpoint, req_data, request_id): data.pop("min_completion_tokens", None) headers = _auth_headers(request_id) + if dp_rank is not None: + # Pin this prefill to a specific DP replica so its P2PConnector (bound + # at base_port + data_parallel_index) holds the produced KV blocks. + headers["X-data-parallel-rank"] = str(dp_rank) resp = await client_info["client"].post(endpoint, json=data, headers=headers) resp.raise_for_status() await resp.aread() return resp -async def _stream_decode(client_info, endpoint, req_data, request_id): +async def _stream_decode(client_info, endpoint, req_data, request_id, dp_rank=None): headers = _auth_headers(request_id) + if dp_rank is not None: + # Round-robin this decode onto a specific DP replica of the decoder. + headers["X-data-parallel-rank"] = str(dp_rank) async with client_info["client"].stream( "POST", endpoint, json=req_data, headers=headers ) as resp: @@ -195,24 +247,39 @@ async def _handle_completions(api: str, request: Request): req_data = await request.json() request_id = str(uuid.uuid4()) + # Pick DP replicas for both roles up front (before any await) so the + # two round-robin advances are atomic together. Header rank is None for + # a role with dp_size==1 (no header → single replica, unchanged). + prefill_rank, prefill_hdr = _next_dp_rank(request.app, "prefill") + decode_rank, decode_hdr = _next_dp_rank(request.app, "decode") + prefill_client = _get_next(request.app, "prefill") - await _prefill(prefill_client, api, req_data, request_id) + await _prefill(prefill_client, api, req_data, request_id, dp_rank=prefill_hdr) # Inject the prefiller's P2PConnector address so the decoder can pull - # KV blocks from it via the P2PConnector transport. + # KV blocks from it. remote_port = base + prefill_rank targets the + # replica that produced the KV (base+0 == base when dp=1). req_data["kv_transfer_params"] = { "prefill": { "kv_request_id": request_id, "remote_host": global_args.p2p_connector_host, - "remote_port": global_args.p2p_connector_port, + "remote_port": global_args.p2p_connector_port + prefill_rank, }, } decode_client = _get_next(request.app, "decode") - logger.debug("prefill=%s decode=%s", prefill_client, decode_client) + logger.debug( + "prefill=%s dp=%s decode=%s dp=%s", + prefill_client, + prefill_rank, + decode_client, + decode_rank, + ) async def generate(): - async for chunk in _stream_decode(decode_client, api, req_data, request_id): + async for chunk in _stream_decode( + decode_client, api, req_data, request_id, dp_rank=decode_hdr + ): yield chunk return StreamingResponse(generate(), media_type="application/json") @@ -237,6 +304,10 @@ async def _handle_completions_decoder_first(api: str, request: Request): req_data = await request.json() request_id = str(uuid.uuid4()) + # Pick DP replicas up front (see _handle_completions for rationale). + prefill_rank, prefill_hdr = _next_dp_rank(request.app, "prefill") + decode_rank, decode_hdr = _next_dp_rank(request.app, "decode") + prefill_client = _get_next(request.app, "prefill") decode_client = _get_next(request.app, "decode") @@ -245,7 +316,7 @@ async def _handle_completions_decoder_first(api: str, request: Request): "prefill": { "kv_request_id": request_id, "remote_host": global_args.p2p_connector_host, - "remote_port": global_args.p2p_connector_port, + "remote_port": global_args.p2p_connector_port + prefill_rank, }, } @@ -255,7 +326,7 @@ async def generate(): async def _run_decode(): try: async for chunk in _stream_decode( - decode_client, api, decode_data, request_id + decode_client, api, decode_data, request_id, dp_rank=decode_hdr ): await queue.put(("data", chunk)) except Exception as exc: @@ -268,7 +339,9 @@ async def _run_decode(): # 2. Send prefill — blocks are computed and pushed to the decoder try: - await _prefill(prefill_client, api, req_data, request_id) + await _prefill( + prefill_client, api, req_data, request_id, dp_rank=prefill_hdr + ) except Exception as exc: logger.warning("decoder-first: prefill failed: %s", exc) diff --git a/tests/v1/kv_offload/tiering/p2p/run_accuracy_test.sh b/tests/v1/kv_offload/tiering/p2p/run_accuracy_test.sh index 82053ebfad77..89300376a783 100755 --- a/tests/v1/kv_offload/tiering/p2p/run_accuracy_test.sh +++ b/tests/v1/kv_offload/tiering/p2p/run_accuracy_test.sh @@ -16,6 +16,12 @@ # NUM_DECODE_INSTANCES default 1 # PREFILLER_TP_SIZE default 1 # DECODER_TP_SIZE default 1 +# DP_EP when set, deploy the decoder in DP-EP attention +# mode (dp=DECODER_TP_SIZE, tp=1, +# --enable-expert-parallel) instead of tensor +# parallel; the proxy round-robins decode across +# the DP replicas. Mirrors +# nixl_integration/run_accuracy_test.sh. # GPU_MEMORY_UTILIZATION default 0.45 # MAX_MODEL_LEN default 512 # PREFILL_BLOCK_SIZE default 128 @@ -29,6 +35,8 @@ # NUM_PREFILL_INSTANCES=2 NUM_DECODE_INSTANCES=2 \ # bash tests/v1/kv_offload/tiering/p2p/run_accuracy_test.sh # bash tests/v1/kv_offload/tiering/p2p/run_accuracy_test.sh --decoder-first +# DP_EP=1 DECODER_TP_SIZE=2 \ +# bash tests/v1/kv_offload/tiering/p2p/run_accuracy_test.sh set -xe @@ -71,6 +79,14 @@ NUM_PREFILL_INSTANCES=${NUM_PREFILL_INSTANCES:-1} NUM_DECODE_INSTANCES=${NUM_DECODE_INSTANCES:-1} PREFILLER_TP_SIZE=${PREFILLER_TP_SIZE:-1} DECODER_TP_SIZE=${DECODER_TP_SIZE:-1} +# DP-EP attention mode (see header). When DP_EP is set the decoder uses +# data-parallel replicas (dp=DECODER_TP_SIZE) that the proxy round-robins over. +DP_EP=${DP_EP:-} +if [[ -n "$DP_EP" ]]; then + DECODER_DP_SIZE=${DECODER_TP_SIZE} +else + DECODER_DP_SIZE=1 +fi GPU_MEMORY_UTILIZATION=${GPU_MEMORY_UTILIZATION:-0.45} MAX_MODEL_LEN=${MAX_MODEL_LEN:-512} PREFILL_BLOCK_SIZE=${PREFILL_BLOCK_SIZE:-128} @@ -162,7 +178,11 @@ run_tests_for_model() { echo "================================" echo "Testing model: $model_name" echo " prefillers=${NUM_PREFILL_INSTANCES} (tp=${PREFILLER_TP_SIZE})" - echo " decoders=${NUM_DECODE_INSTANCES} (tp=${DECODER_TP_SIZE})" + if [[ -n "$DP_EP" ]]; then + echo " decoders=${NUM_DECODE_INSTANCES} (dp-ep=${DECODER_TP_SIZE}, tp=1)" + else + echo " decoders=${NUM_DECODE_INSTANCES} (tp=${DECODER_TP_SIZE})" + fi echo " decoder_first=${DECODER_FIRST}" echo "================================" @@ -241,9 +261,19 @@ run_tests_for_model() { --block-size ${DECODE_BLOCK_SIZE} \ --gpu-memory-utilization ${GPU_MEMORY_UTILIZATION} \ --max-model-len ${MAX_MODEL_LEN} \ - --tensor-parallel-size ${DECODER_TP_SIZE} \ --kv-transfer-config '${kv_cfg}'" + # DP-EP attention mode: data-parallel + expert-parallel decode replicas + # (dp=DECODER_TP_SIZE, tp=1) instead of tensor parallel. Mirrors + # nixl_integration/run_accuracy_test.sh. + if [[ -z "$DP_EP" ]]; then + BASE_CMD="${BASE_CMD} --tensor-parallel-size ${DECODER_TP_SIZE}" + else + echo "DP-EP Attention enabled, deploying decoder with dp=${DECODER_TP_SIZE} and tp=1" + BASE_CMD="${BASE_CMD} --data-parallel-size ${DECODER_TP_SIZE} \ + --tensor-parallel-size 1 --enable-expert-parallel" + fi + if [[ -n "$VLLM_SERVE_EXTRA_ARGS" ]]; then IFS=',' read -r -a extra_args <<< "$VLLM_SERVE_EXTRA_ARGS" for arg in "${extra_args[@]}"; do @@ -282,7 +312,8 @@ run_tests_for_model() { --p2p-connector-host ${P2P_HOST} \ --p2p-connector-port ${PREFILL_PD_PORTS[0]} \ --decoder-p2p-connector-host ${P2P_HOST} \ - --decoder-p2p-connector-port ${DECODE_PD_PORTS[0]}" + --decoder-p2p-connector-port ${DECODE_PD_PORTS[0]} \ + --decoder-dp-size ${DECODER_DP_SIZE}" if [[ "${DECODER_FIRST}" == "true" ]]; then PROXY_CMD="${PROXY_CMD} --decoder-first" diff --git a/tests/v1/kv_offload/tiering/p2p/test_data_transport.py b/tests/v1/kv_offload/tiering/p2p/test_data_transport.py index d515e102e21e..7eb812435247 100644 --- a/tests/v1/kv_offload/tiering/p2p/test_data_transport.py +++ b/tests/v1/kv_offload/tiering/p2p/test_data_transport.py @@ -183,6 +183,35 @@ def test_poll_ignores_in_progress(self): assert result.done == () assert result.failed == () + def test_poll_peer_id_scopes_to_peer(self): + """poll(peer_id) drains only that peer's transfers. + + Regression: the transport is shared across peer sessions (e.g. a + single prefiller serving a DP>1 decoder). An unscoped poll by one + session used to consume and discard sibling sessions' completions, + starving them until timeout. poll(peer_id) must leave other peers' + transfers inflight. + """ + transport = self._make_transport() + transport.add_remote_peer("peer:1", b"meta", 0x1000, 8, 1024) + transport.add_remote_peer("peer:2", b"meta", 0x2000, 8, 1024) + + tid1 = transport.write_blocks("peer:1", [0], [1]) + tid2 = transport.write_blocks("peer:2", [2], [3]) + transport._agent.check_xfer_state.return_value = "DONE" + + # Polling peer:1 must not consume peer:2's completed transfer. + result = transport.poll(peer_id="peer:1") + assert tid1 in result.done + assert tid2 not in result.done + assert tid1 not in transport._inflight + assert tid2 in transport._inflight + + # peer:2 sees its own completion when it polls. + result2 = transport.poll(peer_id="peer:2") + assert tid2 in result2.done + assert tid2 not in transport._inflight + def test_cancel_removes_inflight(self): """cancel removes transfers and releases handles.""" transport = self._make_transport() diff --git a/tests/v1/kv_offload/tiering/p2p/test_manager.py b/tests/v1/kv_offload/tiering/p2p/test_manager.py index 73e3655ecc99..52e2c8363232 100644 --- a/tests/v1/kv_offload/tiering/p2p/test_manager.py +++ b/tests/v1/kv_offload/tiering/p2p/test_manager.py @@ -9,9 +9,11 @@ from __future__ import annotations import time +import uuid from types import SimpleNamespace import numpy as np +import pytest from vllm.v1.kv_offload.base import LookupResult, ReqContext, ScheduleEndContext from vllm.v1.kv_offload.tiering.base import JobMetadata, JobResult @@ -598,7 +600,7 @@ def cancel(self, transfer_ids, mode: str = "immediate") -> list[int]: return list(self._still_queue[0]) return [] - def poll(self): + def poll(self, peer_id=None): self.poll_calls += 1 class _Empty: @@ -831,7 +833,7 @@ def write_blocks(self, peer_id, local_idxs, remote_idxs): self._inflight_done.append(tid) return tid - def poll(self): + def poll(self, peer_id=None): from vllm.v1.kv_offload.tiering.p2p.data.base import PollResult done = self._inflight_done[:] @@ -1373,3 +1375,115 @@ def test_dead_connection_with_pending_work_surfaces_failures(self): assert (900, False) not in finishes assert (900, True) not in finishes assert "req-store" in mgr_a._unbound_stores + + +# --------------------------------------------------------------------------- +# Tests for host/port resolution in __init__ (env-var defaults) +# --------------------------------------------------------------------------- + + +class TestBindHostPortDefaults: + """host/port fall back to VLLM_P2P_SIDE_CHANNEL_* when not in config.""" + + @staticmethod + def _construct(monkeypatch, dp_index=0, **kwargs) -> P2PSecondaryTierManager: + """Build a manager with the transports/file-mapper stubbed out. + + The host is used verbatim (no resolution). The transport constructor + args are recorded on ``mgr._test_calls`` so tests can assert the ZMQ + identity (``host:port``) stays decoupled from the NIXL agent name + (a uuid). + """ + monkeypatch.setattr( + manager_module, + "FileMapper", + SimpleNamespace( + from_offloading_spec=lambda **_: SimpleNamespace( + get_run_config=lambda: {} + ) + ), + ) + calls: dict = {} + monkeypatch.setattr( + manager_module, + "NixlTransport", + lambda agent_name, *a, **k: calls.update(nixl_name=agent_name) + or SimpleNamespace(), + ) + monkeypatch.setattr( + manager_module, + "ZmqTransport", + lambda local_id, host, port, *a, **k: calls.update( + zmq_id=local_id, zmq_host=host, zmq_port=port + ) + or SimpleNamespace(), + ) + spec = SimpleNamespace( + block_size_factor=1, + vllm_config=SimpleNamespace( + parallel_config=SimpleNamespace(data_parallel_index=dp_index) + ), + ) + mgr = P2PSecondaryTierManager(spec, memoryview(b""), **kwargs) + mgr._test_calls = calls + return mgr + + def test_defaults_from_env_unset(self, monkeypatch): + monkeypatch.delenv("VLLM_P2P_SIDE_CHANNEL_HOST", raising=False) + monkeypatch.delenv("VLLM_P2P_SIDE_CHANNEL_PORT", raising=False) + mgr = self._construct(monkeypatch) + # localhost default is used verbatim as the dial-back identity. + assert mgr._local_id == "localhost:5710" + + def test_env_override(self, monkeypatch): + monkeypatch.setenv("VLLM_P2P_SIDE_CHANNEL_HOST", "10.1.2.3") + monkeypatch.setenv("VLLM_P2P_SIDE_CHANNEL_PORT", "5799") + mgr = self._construct(monkeypatch) + assert mgr._local_id == "10.1.2.3:5799" + + def test_explicit_config_wins(self, monkeypatch): + monkeypatch.setenv("VLLM_P2P_SIDE_CHANNEL_HOST", "10.1.2.3") + monkeypatch.setenv("VLLM_P2P_SIDE_CHANNEL_PORT", "5799") + mgr = self._construct(monkeypatch, host="192.0.2.5", port=6001) + assert mgr._local_id == "192.0.2.5:6001" + + def test_dp_index_offsets_default_port(self, monkeypatch): + monkeypatch.delenv("VLLM_P2P_SIDE_CHANNEL_HOST", raising=False) + monkeypatch.delenv("VLLM_P2P_SIDE_CHANNEL_PORT", raising=False) + mgr = self._construct(monkeypatch, dp_index=2) + assert mgr._local_id == "localhost:5712" + + def test_dp_index_offsets_explicit_port(self, monkeypatch): + mgr = self._construct(monkeypatch, dp_index=1, host="192.0.2.5", port=6000) + assert mgr._local_id == "192.0.2.5:6001" + + @pytest.mark.parametrize( + "bind_host", ["localhost", "127.0.0.1", "::1", "0.0.0.0", "::", "192.0.2.5"] + ) + def test_host_used_verbatim(self, monkeypatch, bind_host): + # The host is never rewritten — loopbacks and wildcards alike become + # the dial-back identity verbatim (mirrors the NIXL connector). + mgr = self._construct(monkeypatch, host=bind_host, port=5710) + assert mgr._local_id == f"{bind_host}:5710" + + def test_nixl_name_decoupled_from_identity(self, monkeypatch): + # The ZMQ identity is the verbatim host:port; the NIXL agent name is a + # uuid, distinct from the identity and never used as an address. + mgr = self._construct(monkeypatch, host="127.0.0.1", port=5710) + assert mgr._test_calls["zmq_host"] == "127.0.0.1" + assert mgr._test_calls["zmq_port"] == 5710 + assert mgr._test_calls["zmq_id"] == "127.0.0.1:5710" + assert mgr._local_id == "127.0.0.1:5710" + # nixl_name is a valid uuid4 and not the host:port identity. + nixl_name = mgr._test_calls["nixl_name"] + assert nixl_name != mgr._local_id + assert uuid.UUID(nixl_name).version == 4 + + def test_same_host_port_gets_distinct_nixl_names(self, monkeypatch): + # The original collision: two peers sharing a host:port must still get + # distinct NIXL agent names so add_remote_agent doesn't reject a remote + # whose name equals the local. The per-process uuid guarantees this. + mgr_a = self._construct(monkeypatch, host="localhost", port=5710) + mgr_b = self._construct(monkeypatch, host="localhost", port=5710) + assert mgr_a._local_id == mgr_b._local_id == "localhost:5710" + assert mgr_a._nixl_agent_name != mgr_b._nixl_agent_name diff --git a/tests/v1/kv_offload/tiering/p2p/test_sessions.py b/tests/v1/kv_offload/tiering/p2p/test_sessions.py index 05c0a3abdd8b..d1c6ad0cda7b 100644 --- a/tests/v1/kv_offload/tiering/p2p/test_sessions.py +++ b/tests/v1/kv_offload/tiering/p2p/test_sessions.py @@ -110,7 +110,7 @@ def write_blocks(self, peer_id, local_idxs, remote_idxs) -> int | None: self._transfers[tid] = (peer_id, local_idxs, remote_idxs) return tid - def poll(self): + def poll(self, peer_id=None): from vllm.v1.kv_offload.tiering.p2p.data.base import PollResult result = PollResult(done=list(self._poll_done), failed=list(self._poll_failed)) diff --git a/tests/v1/kv_offload/tiering/test_fs_tier.py b/tests/v1/kv_offload/tiering/test_fs_tier.py index c90a4dad5e79..4ac734d957f4 100644 --- a/tests/v1/kv_offload/tiering/test_fs_tier.py +++ b/tests/v1/kv_offload/tiering/test_fs_tier.py @@ -91,24 +91,10 @@ def make_job( ) -def drain(tier: FileSystemTierManager, max_rounds: int = 100) -> list: - """ - Call get_finished_jobs() repeatedly until no new results arrive for 20 - consecutive rounds or max_rounds is reached. - """ - results = [] - idle = 0 - for _ in range(max_rounds): - time.sleep(0.01) - new = list(tier.get_finished_jobs()) - results.extend(new) - if new: - idle = 0 - else: - idle += 1 - if idle >= 20: - break - return results +def drain(tier: FileSystemTierManager) -> list: + """Block until all in-flight jobs finish, then collect results.""" + tier.drain_jobs() + return list(tier.get_finished_jobs()) def lookup_and_wait( diff --git a/tests/v1/metrics/test_stats.py b/tests/v1/metrics/test_stats.py index 21f496ea4aea..0de74f0faaa4 100644 --- a/tests/v1/metrics/test_stats.py +++ b/tests/v1/metrics/test_stats.py @@ -1,12 +1,17 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from vllm.v1.engine import FinishReason +from vllm.v1.core.sched.output import ScheduledEncoderInputStats, SchedulerOutput +from vllm.v1.engine import EngineCoreOutputs, FinishReason from vllm.v1.metrics.stats import ( IterationStats, PrefillStats, PromptTokenStats, RequestStateStats, + SchedulerIterationDetails, + SchedulerStats, ) +from vllm.v1.serial_utils import MsgpackDecoder, MsgpackEncoder +from vllm.v1.utils import compute_iteration_details def test_iteration_stats_repr(): @@ -14,6 +19,45 @@ def test_iteration_stats_repr(): assert repr(iteration_stats).startswith("IterationStats(") +def test_scheduler_iteration_details_serialization(): + iteration_details = SchedulerIterationDetails( + iteration_index=1, + num_ctx_requests=2, + num_ctx_tokens=3, + num_generation_requests=4, + num_generation_tokens=5, + elapsed_ms=6.7, + num_encoder_inputs=2, + num_encoder_output_tokens=392, + ) + outputs = EngineCoreOutputs( + scheduler_stats=SchedulerStats( + kv_cache_usage=0.5, + iteration_details=iteration_details, + ) + ) + + encoded = MsgpackEncoder().encode(outputs) + decoded = MsgpackDecoder(EngineCoreOutputs).decode(encoded) + + assert decoded.scheduler_stats is not None + assert decoded.scheduler_stats.kv_cache_usage == 0.5 + assert decoded.scheduler_stats.iteration_details == iteration_details + + +def test_compute_iteration_details_includes_encoder_stats(): + scheduler_output = SchedulerOutput.make_empty() + scheduler_output.scheduled_encoder_input_stats = ScheduledEncoderInputStats( + num_inputs=2, + output_tokens=392, + ) + + iteration_details = compute_iteration_details(scheduler_output) + + assert iteration_details.num_encoder_inputs == 2 + assert iteration_details.num_encoder_output_tokens == 392 + + def test_prefill_kv_computed_with_cache(): """Test that prefill KV compute correctly excludes cached tokens.""" iteration_stats = IterationStats() diff --git a/tests/v1/sample/test_head_dtype.py b/tests/v1/sample/test_head_dtype.py index 37f5a6e4531e..7531cdfc4501 100644 --- a/tests/v1/sample/test_head_dtype.py +++ b/tests/v1/sample/test_head_dtype.py @@ -95,6 +95,37 @@ def test_head_dtype_equal_to_model_dtype_uses_quant_method(default_vllm_config): assert logits.dtype == torch.bfloat16 +@pytest.mark.skipif( + not torch.cuda.is_available(), + reason="Exercises the torch.mm(out_dtype=...) device fast path, " + "available on CUDA and ROCm.", +) +def test_fp32_head_uses_mm_fast_path_on_device(default_vllm_config): + # On ROCm, current_platform.is_cuda() is False, so this previously fell + # through to the cast path (F.linear) instead of torch.mm(out_dtype=...), + # even though ROCm supports the out_dtype mm via its non-Lt GEMM path. + from unittest import mock + + vocab_size, hidden_size, num_tokens = 64, 16, 4 + lp = _build_processor(vocab_size) + lp.head_dtype = torch.float32 + + hidden_states = torch.randn( + num_tokens, hidden_size, dtype=torch.bfloat16, device="cuda" + ) + weight = torch.randn(vocab_size, hidden_size, dtype=torch.bfloat16, device="cuda") + + with mock.patch( + "vllm.model_executor.layers.logits_processor.F.linear" + ) as linear_mock: + logits = lp._get_logits(hidden_states, _FakeLmHead(weight), None) + + linear_mock.assert_not_called() + assert logits.dtype == torch.float32 + expected = torch.nn.functional.linear(hidden_states.float(), weight.float()) + torch.testing.assert_close(logits, expected) + + def test_fp32_head_rejects_quantized_lm_head(default_vllm_config): lp = _build_processor(64) lp.head_dtype = torch.float32 diff --git a/tests/v1/spec_decode/test_dflash_causality.py b/tests/v1/spec_decode/test_dflash_causality.py new file mode 100644 index 000000000000..310b3b6db86d --- /dev/null +++ b/tests/v1/spec_decode/test_dflash_causality.py @@ -0,0 +1,55 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Config-only resolution of DFlash draft attention causality. + +``dflash_has_any_non_causal`` decides pre-build whether the draft needs a +non-causal-capable backend, so its branch table (explicit override, SWA-derived +per-layer causality, and the no-``layer_types`` fallback) is worth pinning. +""" + +from types import SimpleNamespace + +import pytest + +from vllm.model_executor.models.qwen3_dflash import ( + _dflash_layer_causal, + dflash_has_any_non_causal, +) + + +def _config(num_hidden_layers, layer_types=None, causal_override=None): + dflash_config = None if causal_override is None else {"causal": causal_override} + return SimpleNamespace( + num_hidden_layers=num_hidden_layers, + layer_types=layer_types, + dflash_config=dflash_config, + ) + + +@pytest.mark.parametrize( + "config,expected", + [ + # Override forces causality on every layer, ignoring layer_types. + (_config(2, layer_types=["full_attention"] * 2, causal_override=True), False), + # Override forces non-causal on every layer. + ( + _config(2, layer_types=["sliding_attention"] * 2, causal_override=False), + True, + ), + # SWA-derived: full-attention layers are non-causal. + (_config(2, layer_types=["sliding_attention", "full_attention"]), True), + # SWA-derived: all-sliding is fully causal. + (_config(2, layer_types=["sliding_attention", "sliding_attention"]), False), + # No layer_types -> non-causal fallback. + (_config(2, layer_types=None), True), + (_config(2, layer_types=[]), True), + ], +) +def test_dflash_has_any_non_causal(config, expected): + assert dflash_has_any_non_causal(config) is expected + + +def test_dflash_layer_causal_is_per_layer(): + config = _config(2, layer_types=["sliding_attention", "full_attention"]) + assert _dflash_layer_causal(config, 0) is True + assert _dflash_layer_causal(config, 1) is False diff --git a/tools/pre_commit/generate_attention_backend_docs.py b/tools/pre_commit/generate_attention_backend_docs.py index 39d173f761a9..d44456530d4b 100644 --- a/tools/pre_commit/generate_attention_backend_docs.py +++ b/tools/pre_commit/generate_attention_backend_docs.py @@ -1480,7 +1480,9 @@ def _get_backends_from_return(stmts: list) -> list[str]: def _is_sm100_check(test: ast.expr) -> bool: - """Check if test is `something.major == 10`.""" + """Check if test is `something.major == 10`, possibly inside an `and`.""" + if isinstance(test, ast.BoolOp) and isinstance(test.op, ast.And): + return any(_is_sm100_check(value) for value in test.values) return ( isinstance(test, ast.Compare) and isinstance(test.left, ast.Attribute) diff --git a/vllm/compilation/backends.py b/vllm/compilation/backends.py index f80498c8135c..8fbf5b41f747 100644 --- a/vllm/compilation/backends.py +++ b/vllm/compilation/backends.py @@ -489,6 +489,11 @@ def _decompose_size_nodes(graph: fx.GraphModule) -> None: size_nodes = list(graph.graph.find_nodes(op="call_method", target="size")) for node in size_nodes: + # Only x.size() (no dim) returns a torch.Size tuple that can't cross + # split boundaries. x.size(dim) already returns a scalar SymInt/int, + # which crosses fine, so leave it untouched. + if len(node.args) > 1 or "dim" in node.kwargs: + continue tensor_node = node.args[0] ev = tensor_node.meta.get("example_value") assert ev is not None, ( diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index b9c612621982..0cc09f8295d0 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -10,6 +10,7 @@ from typing_extensions import Self from vllm.config import LoadConfig +from vllm.config.cache import CacheDType from vllm.config.kernel import MoEBackend from vllm.config.model import HfOverrides, ModelConfig from vllm.config.parallel import ParallelConfig @@ -118,6 +119,9 @@ class SpeculativeConfig: """Attention backend to use for the draft model. When `None`, the backend is automatically selected. Useful when the drafter requires a different attention backend (e.g. DFlash needs a non-causal-capable backend like FLASH_ATTN).""" + kv_cache_dtype: CacheDType | None = None + """KV cache dtype for the draft model. When `None`, the draft inherits the + target model's `--kv-cache-dtype`.""" max_model_len: int | None = Field(default=None, ge=1) """The maximum model length of the draft model. Used when testing the ability to skip speculation for some sequences.""" @@ -626,6 +630,18 @@ def compose_draft_hf_overrides( SpeculativeConfig._apply_composed_hf_override, target_hf_overrides ) + @staticmethod + def _is_custom_proposer_path(model: str | None) -> bool: + """True if ``model`` is a dotted import path (e.g. ``pkg.MyProposer``).""" + if model is None: + return False + if model.startswith(("http://", "https://", "file://")): + return False + if "/" in model: + return False + parts = model.split(".") + return len(parts) >= 2 and all(part.isidentifier() for part in parts) + def __post_init__(self): # Note: "method" is a new parameter that helps to extend the # configuration of non-model-based proposers, and the "model" parameter @@ -636,14 +652,9 @@ def __post_init__(self): # default. # infer method from user args - # Check if the model field contains a custom module path (e.g., 'pkg.Mod') - if ( - self.model is not None - and "." in self.model - and not self.model.startswith(("http://", "https://", "file://")) - and "/" not in self.model # not a HuggingFace repo (org/model) + if self.method is None and SpeculativeConfig._is_custom_proposer_path( + self.model ): - # Treat as a custom class path self.method = "custom_class" elif self.method is None: if self.model in ("ngram", "[ngram]"): diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py index 47243968936a..d48cb2a5ff37 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py @@ -1169,7 +1169,12 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): "All non-MLA kv cache tensors must have the same size" ) - if cache.shape[0] != num_blocks: + # When there's a mismatch between kbs<>bs, we rely on HMA to ensure + # caches are either [NB, PS] or [NB*r, PS/r] where r is bs/kbs. + if ( + self._physical_blocks_per_logical_kv_block == 1 + and cache.shape[0] != num_blocks + ): raise AssertionError( "All kv cache tensors must have the same number of " f"blocks; layer={layer_name}, " diff --git a/vllm/distributed/parallel_state.py b/vllm/distributed/parallel_state.py index 162ed03d23b2..e46ca1691c2f 100644 --- a/vllm/distributed/parallel_state.py +++ b/vllm/distributed/parallel_state.py @@ -1427,7 +1427,10 @@ def get_pcp_group() -> GroupCoordinator: @contextmanager -def graph_capture(device: torch.device): +def graph_capture( + device: torch.device, + graph_capture_context: GraphCaptureContext | None = None, +): """ `graph_capture` is a context manager which should surround the code that is capturing the CUDA graph. Its main purpose is to ensure that some @@ -1440,8 +1443,13 @@ def graph_capture(device: torch.device): the graph capture is running on a separate stream from the default stream, in order to explicitly distinguish the kernels to capture from other kernels possibly launched on background in the default stream. + + A caller may pass an explicit ``graph_capture_context`` to control the + stream used (e.g. to capture on the default stream). """ - context = GraphCaptureContext(torch.cuda.Stream(device=device)) + context = graph_capture_context or GraphCaptureContext( + torch.cuda.Stream(device=device) + ) with get_tp_group().graph_capture(context), get_pp_group().graph_capture(context): yield context diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 742d62ac3698..07877b0812c7 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -1754,6 +1754,10 @@ def create_speculative_config( if self.speculative_config is None: return None + self.speculative_config = { + k.replace("-", "_"): v for k, v in self.speculative_config.items() + } + # Note(Shangming): These parameters are not obtained from the cli arg # '--speculative-config' and must be passed in when creating the engine # config. diff --git a/vllm/entrypoints/openai/completion/protocol.py b/vllm/entrypoints/openai/completion/protocol.py index 6f8171f97ba7..1784d0f5364c 100644 --- a/vllm/entrypoints/openai/completion/protocol.py +++ b/vllm/entrypoints/openai/completion/protocol.py @@ -209,10 +209,10 @@ class CompletionRequest(OpenAIBaseModel): ), ) - vllm_xargs: dict[str, str | int | float] | None = Field( + vllm_xargs: dict[str, str | int | float | list[str | int | float]] | None = Field( default=None, description=( - "Additional request parameters with string or " + "Additional request parameters with (list of) string or " "numeric values, used by custom extensions." ), ) diff --git a/vllm/entrypoints/openai/responses/protocol.py b/vllm/entrypoints/openai/responses/protocol.py index fabc79677c2d..d4708a5fb3ee 100644 --- a/vllm/entrypoints/openai/responses/protocol.py +++ b/vllm/entrypoints/openai/responses/protocol.py @@ -749,7 +749,9 @@ def from_request( output=output, input_messages=input_messages, output_messages=output_messages, - parallel_tool_calls=request.parallel_tool_calls, + parallel_tool_calls=request.parallel_tool_calls + if request.parallel_tool_calls is not None + else ResponsesRequest.model_fields["parallel_tool_calls"].default, temperature=sampling_params.temperature, tool_choice=request.tool_choice, tools=request.tools, diff --git a/vllm/entrypoints/speech_to_text/transcription/protocol.py b/vllm/entrypoints/speech_to_text/transcription/protocol.py index abf1a11a0eea..3d6600fe3656 100644 --- a/vllm/entrypoints/speech_to_text/transcription/protocol.py +++ b/vllm/entrypoints/speech_to_text/transcription/protocol.py @@ -116,10 +116,10 @@ class TranscriptionRequest(OpenAIBaseModel): stream_include_usage: bool | None = False stream_continuous_usage_stats: bool | None = False - vllm_xargs: dict[str, str | int | float | bool] | None = Field( + vllm_xargs: dict[str, str | int | float | list[str | int | float]] | None = Field( default=None, description=( - "Additional request parameters with string or " + "Additional request parameters with (list of) string or " "numeric values, used by custom extensions." ), ) diff --git a/vllm/envs.py b/vllm/envs.py index adcc3c24f026..5b601cda8cb2 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -209,6 +209,8 @@ VLLM_DISABLE_REQUEST_ID_RANDOMIZATION: bool = False VLLM_NIXL_SIDE_CHANNEL_HOST: str = "localhost" VLLM_NIXL_SIDE_CHANNEL_PORT: int = 5600 + VLLM_P2P_SIDE_CHANNEL_HOST: str = "localhost" + VLLM_P2P_SIDE_CHANNEL_PORT: int = 5710 VLLM_EC_SIDE_CHANNEL_HOST: str = "localhost" VLLM_EC_SIDE_CHANNEL_PORT: int = 5601 VLLM_MOONCAKE_BOOTSTRAP_PORT: int = 8998 @@ -1586,6 +1588,16 @@ def _resolve_rust_frontend_path() -> str | None: "VLLM_NIXL_SIDE_CHANNEL_PORT": lambda: int( os.getenv("VLLM_NIXL_SIDE_CHANNEL_PORT", "5600") ), + # Address the P2P KV-offload control socket binds to. Defaults to + # ``localhost`` (loopback only); must be set to the node IP for + # cross-host P2P so remote peers can reach the socket. + "VLLM_P2P_SIDE_CHANNEL_HOST": lambda: os.getenv( + "VLLM_P2P_SIDE_CHANNEL_HOST", "localhost" + ), + # Port the P2P KV-offload control socket binds to. + "VLLM_P2P_SIDE_CHANNEL_PORT": lambda: int( + os.getenv("VLLM_P2P_SIDE_CHANNEL_PORT", "5710") + ), # IP address used for the EC connector's ZMQ side channel # (producer ROUTER bind, consumer DEALER dial). "VLLM_EC_SIDE_CHANNEL_HOST": lambda: os.getenv( diff --git a/vllm/kernels/helion/__init__.py b/vllm/kernels/helion/__init__.py index 8c05c428bb07..8a20c9eb9ef2 100644 --- a/vllm/kernels/helion/__init__.py +++ b/vllm/kernels/helion/__init__.py @@ -2,7 +2,6 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Helion integration for vLLM.""" -import vllm.kernels.helion.ops # noqa: F401 Auto-register all Helion ops from vllm.kernels.helion.case_key import CaseKey from vllm.kernels.helion.config_manager import ( ConfigManager, diff --git a/vllm/kernels/helion/ops/__init__.py b/vllm/kernels/helion/ops/__init__.py index eacd483bbb7d..588c070a1245 100644 --- a/vllm/kernels/helion/ops/__init__.py +++ b/vllm/kernels/helion/ops/__init__.py @@ -1,11 +1,33 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Auto-import all Helion op modules to trigger kernel registration.""" +"""Helion kernel implementation. + +Importing this package does NOT register any kernels. Runtime code imports the +specific op module it needs, e.g.:: + + from vllm.kernels.helion.ops import scaled_mm # noqa: F401 + +which triggers that op's ``@register_kernel`` as an import side effect. + +Tools that need the full registry (e.g. scripts/autotune_helion_kernels.py) +call ``import_all_ops()`` to force every op module to register. +""" import importlib import pkgutil -# Automatically import all submodules so that @register_kernel -# decorators execute and register ops with torch.ops.vllm_helion. -for _module_info in pkgutil.iter_modules(__path__): - importlib.import_module(f"{__name__}.{_module_info.name}") + +def import_all_kernels() -> list[str]: + """Import every kernel submodule so all ``@register_kernel`` decorators run. + + Returns: + The fully-qualified module names that were imported. + """ + imported: list[str] = [] + for module_info in pkgutil.iter_modules(__path__): + if module_info.ispkg: + continue + module_name = f"{__name__}.{module_info.name}" + importlib.import_module(module_name) + imported.append(module_name) + return imported diff --git a/vllm/kernels/helion/ops/dynamic_per_token_scaled_fp8_quant.py b/vllm/kernels/helion/ops/dynamic_per_token_scaled_fp8_quant.py index eef262dcfe25..45bd8f6fcf3e 100644 --- a/vllm/kernels/helion/ops/dynamic_per_token_scaled_fp8_quant.py +++ b/vllm/kernels/helion/ops/dynamic_per_token_scaled_fp8_quant.py @@ -109,7 +109,18 @@ def baseline( scale: torch.Tensor, # [num_tokens, 1] scale_ub: torch.Tensor | None = None, # scalar tensor ) -> None: - torch.ops._C.dynamic_per_token_scaled_fp8_quant(result, input, scale, scale_ub) + fp8_min, fp8_max = get_fp8_min_max() + min_scaling_factor = 1.0 / (fp8_max * 512.0) + + x = input.to(torch.float32) + s = torch.amax(torch.abs(x), dim=-1, keepdim=True) + if scale_ub is not None: + s = s.clamp(max=scale_ub) + s = (s * (1.0 / fp8_max)).clamp(min=min_scaling_factor) + y = (x / s).clamp(fp8_min, fp8_max) + + scale.copy_(s) + result.copy_(y.to(result.dtype)) # Overwrite autotune_baseline_atol and autotune_baseline_rtol diff --git a/vllm/kernels/helion/ops/per_token_group_fp8_quant.py b/vllm/kernels/helion/ops/per_token_group_fp8_quant.py index 8b73fac4b8eb..42fd0215f594 100644 --- a/vllm/kernels/helion/ops/per_token_group_fp8_quant.py +++ b/vllm/kernels/helion/ops/per_token_group_fp8_quant.py @@ -164,18 +164,17 @@ def baseline( dummy_is_scale_transposed: bool = False, dummy_is_tma_aligned: bool = False, ) -> None: - torch.ops._C.per_token_group_fp8_quant( - input, - output_q, - output_s, - group_size, - eps, - fp8_min, - fp8_max, - scale_ue8m0, - dummy_is_scale_transposed, - dummy_is_tma_aligned, - ) + num_tokens, hidden_size = input.shape + groups_per_row = hidden_size // group_size + + x = input.view(num_tokens, groups_per_row, group_size).to(torch.float32) + s = torch.clamp(torch.amax(torch.abs(x), dim=-1), min=eps) / fp8_max + if scale_ue8m0: + s = torch.exp2(torch.ceil(torch.log2(s))) + y = torch.clamp(x / s[:, :, None], fp8_min, fp8_max) + + output_s.copy_(s) + output_q.copy_(y.view(num_tokens, hidden_size).to(output_q.dtype)) @register_kernel( diff --git a/vllm/kernels/helion/ops/rms_norm_dynamic_per_token_quant.py b/vllm/kernels/helion/ops/rms_norm_dynamic_per_token_quant.py index f15132c27cfc..3e02169db3b8 100644 --- a/vllm/kernels/helion/ops/rms_norm_dynamic_per_token_quant.py +++ b/vllm/kernels/helion/ops/rms_norm_dynamic_per_token_quant.py @@ -131,9 +131,37 @@ def baseline( scale_ub: torch.Tensor | None = None, # [] residual: torch.Tensor | None = None, # [num_tokens, hidden_size] ) -> None: - torch.ops._C.rms_norm_dynamic_per_token_quant( - result, input, weight, scale, epsilon, scale_ub, residual - ) + _, hidden_size = input.shape + quant_dtype = result.dtype + qtype_min: int | float + qtype_max: int | float + + if quant_dtype == torch.int8: + qtype_min, qtype_max = get_int8_min_max() + min_scaling_factor = get_int8_min_scaling_factor() + else: + qtype_min, qtype_max = get_fp8_min_max() + min_scaling_factor = 1.0 / (qtype_max * 512.0) + + x = input.to(torch.float32) + if residual is not None: + x = x + residual + residual.copy_(x.to(residual.dtype)) + + rms = torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + epsilon) + x_norm = (x * rms).to(input.dtype) * weight + + s = torch.amax(torch.abs(x_norm), dim=-1, keepdim=True).to(torch.float32) + if scale_ub is not None: + s = s.clamp(max=scale_ub) + s = (s * (1.0 / qtype_max)).clamp(min=min_scaling_factor) + + y = x_norm / s + if quant_dtype == torch.int8: + y = y.round() + + scale.copy_(s) + result.copy_(y.clamp(qtype_min, qtype_max).to(result.dtype)) # Overwrite autotune_baseline_atol and autotune_baseline_rtol diff --git a/vllm/kernels/helion/ops/rms_norm_per_block_quant.py b/vllm/kernels/helion/ops/rms_norm_per_block_quant.py index e7df42f4dbed..da7651a1fa59 100644 --- a/vllm/kernels/helion/ops/rms_norm_per_block_quant.py +++ b/vllm/kernels/helion/ops/rms_norm_per_block_quant.py @@ -169,16 +169,40 @@ def baseline( group_size: int, is_scale_transposed: bool, ) -> None: - torch.ops._C.rms_norm_per_block_quant( - result, - input, - weight, - scale, - epsilon, - scale_ub, - residual, - group_size, - is_scale_transposed, + num_tokens, hidden_size = input.shape + groups_per_row = hidden_size // group_size + quant_dtype = result.dtype + qtype_min: int | float + qtype_max: int | float + + if quant_dtype == torch.int8: + qtype_min, qtype_max = get_int8_min_max() + min_scaling_factor = get_int8_min_scaling_factor() + else: + qtype_min, qtype_max = get_fp8_min_max() + min_scaling_factor = 1.0 / (qtype_max * 512.0) + + x = input.to(torch.float32) + if residual is not None: + x = x + residual + residual.copy_(x.to(residual.dtype)) + + rms = torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + epsilon) + x_norm = (x * rms).to(input.dtype) * weight + x_grouped = x_norm.view(num_tokens, groups_per_row, group_size).to(torch.float32) + + s = torch.amax(torch.abs(x_grouped), dim=-1).to(torch.float32) + if scale_ub is not None: + s = s.clamp(max=scale_ub) + s = (s * (1.0 / qtype_max)).clamp(min=min_scaling_factor) + + y = x_grouped / s[:, :, None] + if quant_dtype == torch.int8: + y = y.round() + + scale.copy_(s) + result.copy_( + y.clamp(qtype_min, qtype_max).view(num_tokens, hidden_size).to(result.dtype) ) diff --git a/vllm/kernels/helion/ops/silu_and_mul_per_block_quant.py b/vllm/kernels/helion/ops/silu_and_mul_per_block_quant.py index f3aaf226b047..06b7f10af2fd 100644 --- a/vllm/kernels/helion/ops/silu_and_mul_per_block_quant.py +++ b/vllm/kernels/helion/ops/silu_and_mul_per_block_quant.py @@ -13,6 +13,7 @@ get_int8_min_scaling_factor, ) from vllm.logger import init_logger +from vllm.model_executor.layers.activation import SiluAndMul from vllm.model_executor.layers.quantization.utils.quant_utils import ( get_fp8_min_max, ) @@ -150,8 +151,34 @@ def baseline( scale_ub: torch.Tensor | None = None, # scalar tensor is_scale_transposed: bool = False, ) -> None: - torch.ops._C.silu_and_mul_per_block_quant( - out, input, scales, group_size, scale_ub, is_scale_transposed + num_tokens, intermediate_size = out.shape + groups_per_row = intermediate_size // group_size + quant_dtype = out.dtype + qtype_min: int | float + qtype_max: int | float + + if quant_dtype == torch.int8: + qtype_min, qtype_max = get_int8_min_max() + min_scaling_factor = get_int8_min_scaling_factor() + else: + qtype_min, qtype_max = get_fp8_min_max() + min_scaling_factor = 1.0 / (qtype_max * 512.0) + + act = SiluAndMul.forward_native(input.to(torch.float32)) + x_grouped = act.view(num_tokens, groups_per_row, group_size) + + s = torch.amax(torch.abs(x_grouped), dim=-1) + if scale_ub is not None: + s = s.clamp(max=scale_ub) + s = (s * (1.0 / qtype_max)).clamp(min=min_scaling_factor) + + y = x_grouped / s[:, :, None] + if quant_dtype == torch.int8: + y = y.round() + + scales.copy_(s) + out.copy_( + y.clamp(qtype_min, qtype_max).view(num_tokens, intermediate_size).to(out.dtype) ) diff --git a/vllm/lora/peft_helper.py b/vllm/lora/peft_helper.py index 1443efd4f0cd..a0bb8bf3c283 100644 --- a/vllm/lora/peft_helper.py +++ b/vllm/lora/peft_helper.py @@ -51,6 +51,8 @@ def _validate_features(self) -> list[str]: return error_msg def __post_init__(self): + if self.r <= 0: + raise ValueError(f"LoRA rank `r` must be a positive integer, got {self.r}.") if self.use_rslora: logger.info_once("Loading LoRA weights trained with rsLoRA.") self.vllm_lora_scaling_factor = self.lora_alpha / math.sqrt(self.r) diff --git a/vllm/model_executor/kernels/linear/mixed_precision/triton_w4a16.py b/vllm/model_executor/kernels/linear/mixed_precision/triton_w4a16.py index 5ba1caabae66..bc0a587b6763 100644 --- a/vllm/model_executor/kernels/linear/mixed_precision/triton_w4a16.py +++ b/vllm/model_executor/kernels/linear/mixed_precision/triton_w4a16.py @@ -400,12 +400,29 @@ def repack_w_s(x: BasevLLMParameter) -> BasevLLMParameter: if self.w_zp_name is not None: zp = getattr(layer, self.w_zp_name, None) if zp is not None: - # Checkpoint: [N//8, K//G] int32 (N packed at dim 0, K//G at dim 1) - # Kernel needs: [K//G, N//8] — just transpose + c = self.config + K, N = c.partition_weight_shape + group_size = c.group_size if c.group_size != -1 else K + expected_shape = (K // group_size, N // 8) + transposed_shape = (N // 8, K // group_size) + + if tuple(zp.data.shape) == expected_shape: + # GPTQ/AutoGPTQ already stores qzeros in the kernel layout. + qzeros = zp.data.contiguous() + elif tuple(zp.data.shape) == transposed_shape: + # Compressed-tensors stores qzeros transposed from what the + # kernel needs. + qzeros = zp.data.t().contiguous() + else: + raise AssertionError( + f"{self.w_zp_name} shape mismatch: {zp.data.shape}; " + f"expected {expected_shape} or {transposed_shape}" + ) + replace_parameter( layer, self.w_zp_name, - torch.nn.Parameter(zp.data.t().contiguous(), requires_grad=False), + torch.nn.Parameter(qzeros, requires_grad=False), ) def apply_weights( @@ -420,14 +437,17 @@ def apply_weights( K = c.partition_weight_shape[0] group_size = c.group_size if c.group_size != -1 else K - # For symmetric types (uint4b8), use the scalar bias; no zeros tensor + # For symmetric types (uint4b8), use the scalar bias; no zeros tensor. + # Some checkpoint loaders still register qzeros parameters for GPTQ + # layers, but they are not part of the symmetric kernel contract. zp_bias = c.weight_type.bias if c.weight_type.has_bias() else 0 + qzeros = None if c.weight_type.has_bias() else w_zp output = triton_w4a16_gemm( a=x_2d, b_q=w_q, scales=w_s, - qzeros=w_zp, + qzeros=qzeros, group_size=group_size, zp_bias=zp_bias, ) diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 032f84984116..108e6ff1524d 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -393,6 +393,19 @@ def __init__( calculate_kv_scales = False self.quant_config = quant_config + if cache_config is not None and cache_config.kv_cache_dtype_skip_layers: + from vllm.model_executor.models.utils import extract_layer_index + + layer_idx = extract_layer_index(prefix) + if str(layer_idx) in cache_config.kv_cache_dtype_skip_layers: + kv_cache_dtype = "auto" + calculate_kv_scales = False + logger.debug( + "Layer %s: kv_cache_dtype=%s", + prefix, + kv_cache_dtype, + ) + dtype = torch.get_default_dtype() if attn_backend is not None: assert attn_backend.is_mla(), ( @@ -1025,7 +1038,7 @@ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec: num_kv_heads=1, head_size=self.head_size, dtype=kv_cache_dtype, - cache_dtype_str=vllm_config.cache_config.cache_dtype, + cache_dtype_str=self.kv_cache_dtype, kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype), ) diff --git a/vllm/model_executor/layers/batch_invariant.py b/vllm/model_executor/layers/batch_invariant.py index 917c72dee8c1..defca7c2c0bd 100644 --- a/vllm/model_executor/layers/batch_invariant.py +++ b/vllm/model_executor/layers/batch_invariant.py @@ -781,6 +781,7 @@ def _rms_norm_kernel( n_cols, eps, BLOCK_SIZE: tl.constexpr, + HAS_WEIGHT: tl.constexpr, ): """ Compute RMS normalization along the last dimension of a 2D tensor. @@ -813,18 +814,19 @@ def _rms_norm_kernel( col_idx = col_offset + tl.arange(0, BLOCK_SIZE) mask = col_idx < n_cols vals = tl.load(row_start_ptr + col_idx, mask=mask, other=0.0) - weight = tl.load(weight_ptr + col_idx, mask=mask, other=1.0) # Compute in float32 then convert back to input dtype vals_f32 = vals.to(tl.float32) - weight_f32 = weight.to(tl.float32) - output_f32 = vals_f32 * inv_rms * weight_f32 + output_f32 = vals_f32 * inv_rms + if HAS_WEIGHT: + weight = tl.load(weight_ptr + col_idx, mask=mask, other=1.0) + output_f32 = output_f32 * weight.to(tl.float32) output = output_f32.to(vals.dtype) tl.store(output_row_start_ptr + col_idx, output, mask=mask) def rms_norm_batch_invariant( input: torch.Tensor, - weight: torch.Tensor, + weight: torch.Tensor | None, eps: float = 1e-6, residual: torch.Tensor | None = None, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: @@ -834,7 +836,8 @@ def rms_norm_batch_invariant( Args: input: Input tensor of shape (..., hidden_size) - weight: Weight tensor of shape (hidden_size,) + weight: Weight tensor of shape (hidden_size,), or None to skip the + per-channel multiply (``RMSNorm(has_weight=False)``) eps: Small constant for numerical stability residual: Optional residual tensor fused into the normalization path @@ -851,17 +854,18 @@ def rms_norm_batch_invariant( ops.fused_add_rms_norm(input, residual, weight, eps) return input, residual - assert weight.dim() == 1, "Weight must be 1-dimensional" - assert input.shape[-1] == weight.shape[0], ( - f"Input last dimension ({input.shape[-1]}) must match " - f"weight dimension ({weight.shape[0]})" - ) + if weight is not None: + assert weight.dim() == 1, "Weight must be 1-dimensional" + assert input.shape[-1] == weight.shape[0], ( + f"Input last dimension ({input.shape[-1]}) must match " + f"weight dimension ({weight.shape[0]})" + ) + weight = weight.contiguous() # Flatten all dimensions except the last one original_shape = input.shape input_2d = input.reshape(-1, input.shape[-1]) input_2d = input_2d.contiguous() - weight = weight.contiguous() n_rows, n_cols = input_2d.shape @@ -870,13 +874,14 @@ def rms_norm_batch_invariant( grid = (n_rows,) _rms_norm_kernel[grid]( input_2d, - weight, + weight if weight is not None else input_2d, output, input_2d.stride(0), output.stride(0), n_cols, eps, BLOCK_SIZE=BLOCK_SIZE, + HAS_WEIGHT=weight is not None, ) return output.reshape(original_shape) @@ -904,37 +909,41 @@ def enable_batch_invariant_mode(): _batch_invariant_MODE = True _batch_invariant_LIB = torch.library.Library("aten", "IMPL") - if current_platform.is_device_capability_family(80): - # SM80 (Ampere) cannot rely on cuBLASLt-only determinism; install the - # triton persistent matmul overrides for mm/addmm/matmul/linear. - _batch_invariant_LIB.impl("aten::mm", mm_batch_invariant, "CUDA") - _batch_invariant_LIB.impl("aten::addmm", addmm_batch_invariant, "CUDA") - _batch_invariant_LIB.impl("aten::matmul", matmul_batch_invariant, "CUDA") - _batch_invariant_LIB.impl("aten::linear", linear_batch_invariant, "CUDA") - else: - # Hopper (SM90) and Blackwell (SM100): the only source of batch - # variance is split-k, which we disable via the cuBLAS workspace - # config. - os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":16:8" - os.environ["CUBLASLT_WORKSPACE_SIZE"] = "1" - - # Triton bmm/persistent-matmul kernels read this for the FP16 N-tile size; - # set unconditionally because bmm is overridden on all CUDA platforms. - if current_platform.is_cuda(): - _fp16_block_size_n = 256 if get_max_shared_memory_bytes() > 106496 else 128 + key = current_platform.dispatch_key - _batch_invariant_LIB.impl( - "aten::_log_softmax", _log_softmax_batch_invariant, "CUDA" - ) - _batch_invariant_LIB.impl("aten::softmax", softmax_batch_invariant, "CUDA") - _batch_invariant_LIB.impl("aten::_softmax", softmax_batch_invariant, "CUDA") - _batch_invariant_LIB.impl("aten::mean.dim", mean_batch_invariant, "CUDA") + if current_platform.is_cuda(): + if current_platform.is_device_capability_family(80): + # SM80 (Ampere) cannot rely on cuBLASLt-only determinism; install the + # triton persistent matmul overrides for mm/addmm/matmul/linear. + _batch_invariant_LIB.impl("aten::mm", mm_batch_invariant, key) + _batch_invariant_LIB.impl("aten::addmm", addmm_batch_invariant, key) + _batch_invariant_LIB.impl("aten::matmul", matmul_batch_invariant, key) + _batch_invariant_LIB.impl("aten::linear", linear_batch_invariant, key) + else: + # Hopper (SM90) and Blackwell (SM100): the only source of batch + # variance is split-k, which we disable via the cuBLAS workspace + # config. + os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":16:8" + os.environ["CUBLASLT_WORKSPACE_SIZE"] = "1" + _fp16_block_size_n = 256 if get_max_shared_memory_bytes() > 106496 else 128 + elif current_platform.is_xpu(): + _batch_invariant_LIB.impl("aten::mm", mm_batch_invariant, key) + _batch_invariant_LIB.impl("aten::addmm", addmm_batch_invariant, key) + # TODO: register matmul and linear for XPU + # once suitable Triton kernels are implemented + + _fp16_block_size_n = 128 + + _batch_invariant_LIB.impl("aten::_log_softmax", _log_softmax_batch_invariant, key) + _batch_invariant_LIB.impl("aten::softmax", softmax_batch_invariant, key) + _batch_invariant_LIB.impl("aten::_softmax", softmax_batch_invariant, key) + _batch_invariant_LIB.impl("aten::mean.dim", mean_batch_invariant, key) # torch 2.12+ registers a built-in Triton bmm kernel for CUDA # (torch._native.ops.bmm_outer_product), so we need allow_override # to replace it at the dispatcher level. _batch_invariant_LIB.impl( - "aten::bmm", bmm_batch_invariant, "CUDA", allow_override=True + "aten::bmm", bmm_batch_invariant, key, allow_override=True ) torch.bmm = bmm_batch_invariant @@ -947,7 +956,8 @@ def enable_batch_invariant_mode(): torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = ( reduced_precision_val ) - torch.backends.cuda.preferred_blas_library(backend="cublaslt") + if current_platform.is_cuda(): + torch.backends.cuda.preferred_blas_library(backend="cublaslt") def override_envs_for_invariance(): diff --git a/vllm/model_executor/layers/fused_moe/config.py b/vllm/model_executor/layers/fused_moe/config.py index ad2ed510b33f..39b14fc249e2 100644 --- a/vllm/model_executor/layers/fused_moe/config.py +++ b/vllm/model_executor/layers/fused_moe/config.py @@ -887,6 +887,9 @@ def int4_w4a16_moe_quant_config( block_shape: list[int] | None = None, a1_gscale: torch.Tensor | None = None, a2_gscale: torch.Tensor | None = None, + gemm1_clamp_limit: float | None = None, + gemm1_alpha: float | None = None, + gemm1_beta: float | None = None, ) -> FusedMoEQuantConfig: """ Construct a quant config for 16-bit float activations and int4 weights. @@ -897,6 +900,9 @@ def int4_w4a16_moe_quant_config( _a2=FusedMoEQuantDesc(shape=group_shape, alpha_or_gscale=a2_gscale), _w1=FusedMoEQuantDesc("int4", group_shape, w1_scale, None, w1_zp, w1_bias), _w2=FusedMoEQuantDesc("int4", group_shape, w2_scale, None, w2_zp, w2_bias), + gemm1_clamp_limit=gemm1_clamp_limit, + gemm1_alpha=gemm1_alpha, + gemm1_beta=gemm1_beta, ) @@ -950,6 +956,9 @@ def int8_w8a16_moe_quant_config( block_shape: list[int] | None = None, a1_gscale: torch.Tensor | None = None, a2_gscale: torch.Tensor | None = None, + gemm1_clamp_limit: float | None = None, + gemm1_alpha: float | None = None, + gemm1_beta: float | None = None, ) -> FusedMoEQuantConfig: """ Construct a quant config for 16-bit float activations and int8 weights. @@ -960,6 +969,9 @@ def int8_w8a16_moe_quant_config( _a2=FusedMoEQuantDesc(shape=group_shape, alpha_or_gscale=a2_gscale), _w1=FusedMoEQuantDesc(torch.int8, group_shape, w1_scale, None, w1_zp, w1_bias), _w2=FusedMoEQuantDesc(torch.int8, group_shape, w2_scale, None, w2_zp, w2_bias), + gemm1_clamp_limit=gemm1_clamp_limit, + gemm1_alpha=gemm1_alpha, + gemm1_beta=gemm1_beta, ) diff --git a/vllm/model_executor/layers/fused_moe/experts/int4_emulation_moe.py b/vllm/model_executor/layers/fused_moe/experts/int4_emulation_moe.py new file mode 100644 index 000000000000..7ec917e46984 --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/experts/int4_emulation_moe.py @@ -0,0 +1,129 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Int4 weight-only quantization emulation for MoE. + +Weights are dequantized from packed int4 to BF16 once at load time; +the forward pass then runs plain TritonExperts in BF16. +""" + +import torch + +import vllm.model_executor.layers.fused_moe.modular_kernel as mk +from vllm.logger import init_logger +from vllm.model_executor.layers.fused_moe.activation import MoEActivation +from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEConfig, + FusedMoEQuantConfig, +) +from vllm.model_executor.layers.fused_moe.experts.triton_moe import TritonExperts +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + QuantKey, + kInt4Static, + kInt4Static32, + kInt4Static32Asym, + kInt4StaticAsym, +) +from vllm.platforms import current_platform + +logger = init_logger(__name__) + + +class Int4EmulationTritonExperts(TritonExperts): + """Int4 W-only MoE that dequantizes weights to BF16 at load time. + + Weights arrive already dequantized (convert_to_wna16_moe_kernel_format + does the unpacking); apply() simply forwards to TritonExperts. + """ + + def __init__( + self, + moe_config: FusedMoEConfig, + quant_config: FusedMoEQuantConfig, + ): + super().__init__(moe_config, quant_config) + logger.warning_once( + "Using Int4EmulationTritonExperts MoE backend. Int4 weights are " + "dequantized to BF16 at load time " + ) + # Weights are dequantized to BF16 before apply() is called, so + # TritonExperts must see them as plain float — clear the int4 dtype + # and scales so the hidden-size assertion and kernel dispatch treat + # them as unquantized. + self.quant_config._w1.dtype = None + self.quant_config._w2.dtype = None + self.quant_config._w1.scale = None + self.quant_config._w2.scale = None + + @staticmethod + def _supports_current_device() -> bool: + return current_platform.is_cuda_alike() + + @staticmethod + def _supports_quant_scheme( + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + return ( + weight_key + in ( + kInt4Static, + kInt4Static32, + kInt4StaticAsym, + kInt4Static32Asym, + ) + and activation_key is None + ) + + @property + def quant_dtype(self) -> torch.dtype | str | None: + return None + + @property + def block_shape(self) -> list[int] | None: + return None + + @property + def expects_unquantized_inputs(self) -> bool: + return True + + def apply( + self, + output: torch.Tensor, + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + activation: MoEActivation, + global_num_experts: int, + expert_map: torch.Tensor | None, + a1q_scale: torch.Tensor | None, + a2_scale: torch.Tensor | None, + workspace13: torch.Tensor, + workspace2: torch.Tensor, + expert_tokens_meta: mk.ExpertTokensMetadata | None, + apply_router_weight_on_input: bool, + ): + if w1.element_size() < 2: + raise RuntimeError( + "Int4EmulationTritonExperts.apply() received packed int4 weights " + "(element_size < 2). Weights must be dequantized to BF16 before " + "the forward pass via convert_to_wna16_moe_kernel_format." + ) + return super().apply( + output=output, + hidden_states=hidden_states, + w1=w1, + w2=w2, + topk_weights=topk_weights, + topk_ids=topk_ids, + activation=activation, + global_num_experts=global_num_experts, + expert_map=expert_map, + a1q_scale=None, + a2_scale=None, + workspace13=workspace13, + workspace2=workspace2, + expert_tokens_meta=expert_tokens_meta, + apply_router_weight_on_input=apply_router_weight_on_input, + ) diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py index edcea5361ec1..55ddd1a2c964 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py @@ -122,7 +122,9 @@ def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bo return ( not moe_parallel_config.use_all2all_kernels or moe_parallel_config.use_ag_rs_all2all_kernels - ) and not moe_parallel_config.enable_eplb + ) and not ( + moe_parallel_config.enable_eplb or moe_parallel_config.is_sequence_parallel + ) class TrtLlmFp8ExpertsModular(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsModular): diff --git a/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py b/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py index d0cc08ea141f..a80013f65baa 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py +++ b/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py @@ -51,6 +51,7 @@ class WNA16MoEBackend(Enum): CPU = "CPU" FLASHINFER_TRTLLM = "FLASHINFER_TRTLLM" XPU = "XPU" + EMULATION = "EMULATION" def backend_to_kernel_cls( @@ -87,6 +88,12 @@ def backend_to_kernel_cls( ) return [CPUExpertsInt4] + elif backend == WNA16MoEBackend.EMULATION: + from vllm.model_executor.layers.fused_moe.experts.int4_emulation_moe import ( + Int4EmulationTritonExperts, + ) + + return [Int4EmulationTritonExperts] else: raise ValueError(f"Unknown WNA16 MoE backend: {backend.value}") @@ -105,6 +112,7 @@ def _get_priority_backends() -> list[WNA16MoEBackend]: WNA16MoEBackend.MARLIN, WNA16MoEBackend.BATCHED_MARLIN, WNA16MoEBackend.HUMMING, + WNA16MoEBackend.EMULATION, ] return _AVAILABLE_BACKENDS @@ -115,6 +123,7 @@ def map_wna16_backend(runner_backend: MoEBackend) -> WNA16MoEBackend: "marlin": WNA16MoEBackend.MARLIN, "humming": WNA16MoEBackend.HUMMING, "flashinfer_trtllm": WNA16MoEBackend.FLASHINFER_TRTLLM, + "emulation": WNA16MoEBackend.EMULATION, } if backend := mapping.get(runner_backend): return backend @@ -215,6 +224,9 @@ def make_wna16_moe_quant_config( w2_bias: torch.Tensor | None = None, a1_gscale: torch.Tensor | None = None, a2_gscale: torch.Tensor | None = None, + gemm1_clamp_limit: float | None = None, + gemm1_alpha: float | None = None, + gemm1_beta: float | None = None, ) -> FusedMoEQuantConfig: """Create the FusedMoEQuantConfig for 4 or 8-bit WNA16 MoE.""" if num_bits == 4: @@ -228,6 +240,9 @@ def make_wna16_moe_quant_config( block_shape=[0, group_size], a1_gscale=a1_gscale, a2_gscale=a2_gscale, + gemm1_clamp_limit=gemm1_clamp_limit, + gemm1_alpha=gemm1_alpha, + gemm1_beta=gemm1_beta, ) else: assert num_bits == 8 @@ -241,6 +256,9 @@ def make_wna16_moe_quant_config( block_shape=[0, group_size], a1_gscale=a1_gscale, a2_gscale=a2_gscale, + gemm1_clamp_limit=gemm1_clamp_limit, + gemm1_alpha=gemm1_alpha, + gemm1_beta=gemm1_beta, ) @@ -263,19 +281,23 @@ def make_wna16_moe_kernel( from vllm.model_executor.layers.fused_moe.experts.cpu_moe import ( CPUExpertsInt4, ) + from vllm.model_executor.layers.fused_moe.experts.int4_emulation_moe import ( + Int4EmulationTritonExperts, + ) from vllm.model_executor.layers.fused_moe.experts.xpu_moe import ( XPUExpertsWNA16, ) # Currently, we only support TrtLlmMxint4ExpertsMonolithic, MarlinExperts, - # BatchedMarlinExperts, XPUExpertsWNA16, CPUExpertsInt4, and the Humming - # grouped/indexed experts. + # BatchedMarlinExperts, XPUExpertsWNA16, CPUExpertsInt4, the Humming + # grouped/indexed experts, and Int4EmulationTritonExperts allowed_experts: tuple[type[mk.FusedMoEExperts], ...] = ( MarlinExperts, BatchedMarlinExperts, TrtLlmMxint4ExpertsMonolithic, XPUExpertsWNA16, CPUExpertsInt4, + Int4EmulationTritonExperts, ) if backend == WNA16MoEBackend.HUMMING: allowed_experts += tuple(backend_to_kernel_cls(WNA16MoEBackend.HUMMING)) @@ -1017,6 +1039,243 @@ def _humming_wna16_weight_schema( ) +def _unpack_and_dequant_int4_gptq( + w_int32: torch.Tensor, + scale: torch.Tensor, + qzeros: torch.Tensor | None, + transpose_output: bool, + output_dtype: torch.dtype = torch.bfloat16, +) -> torch.Tensor: + """Unpack GPTQ-packed int4 weights and dequantize to output_dtype. + + Args: + w_int32: packed weights, shape [E, K_packed, N] where K_packed = K//8 + (8 nibbles per int32, LSB-first in the K dimension). + scale: per-group scales, shape [E, K//group_size, N], float16. + qzeros: optional asymmetric zero-points, shape [E, K//gs, N//8], int32. + None for symmetric (uint4b8 with implicit bias 8). + transpose_output: if True return [E, N, K]; if False return [E, K, N]. + output_dtype: target floating-point dtype (bfloat16 or float16). + + Returns: + Dequantized weight tensor in the requested layout. + """ + E, K_packed, N = w_int32.shape + K = K_packed * 8 + + # Unpack: [E, K_packed, N] -> [E, K_packed, N, 8] via bit-shifts. + # The nibble index (last dim) enumerates K rows within each packed column, + # so we must fuse K_packed and the nibble dim, not N and the nibble dim. + # Permute to [E, K_packed, 8, N] before reshaping to [E, K, N]. + shifts = torch.arange(8, device=w_int32.device, dtype=torch.int32) * 4 + nibbles = (w_int32.unsqueeze(-1) >> shifts) & 0xF # [E, K_packed, N, 8] + + # Reshape to [E, K, N]: fuse K_packed and nibble index (dim 1 and 3) + w = nibbles.permute(0, 1, 3, 2).reshape(E, K, N).to(torch.int16) + + if qzeros is None: + # Symmetric uint4b8: subtract bias so the range is [-8, 7] + w = w - 8 + else: + # Asymmetric: unpack zero-points (same 8-nibble packing) and subtract + # qzeros shape: [E, K//gs, N//8] int32 + gs = K // scale.shape[1] + n_gs = scale.shape[1] + zp_shifts = torch.arange(8, device=qzeros.device, dtype=torch.int32) * 4 + zp_nibbles = (qzeros.unsqueeze(-1) >> zp_shifts) & 0xF # [E, n_gs, N//8, 8] + zp = zp_nibbles.reshape(E, n_gs, N).to(torch.int16) # [E, n_gs, N] + zp = zp.repeat_interleave(gs, dim=1) # [E, K, N] + w = w - zp + + # Broadcast scale [E, K//gs, N] -> [E, K, N] + gs = K // scale.shape[1] + scale_broadcast = scale.repeat_interleave(gs, dim=1).to(output_dtype) + + w_dequant = w.to(output_dtype) * scale_broadcast # [E, K, N] + + if transpose_output: + return w_dequant.permute(0, 2, 1).contiguous() # [E, N, K] + return w_dequant.contiguous() # [E, K, N] + + +def _unpack_and_dequant_int4_awq( + w_int32: torch.Tensor, + scale: torch.Tensor, + qzeros: torch.Tensor | None, + transpose_output: bool, + output_dtype: torch.dtype = torch.bfloat16, +) -> torch.Tensor: + """Unpack AWQ-packed int4 weights and dequantize to output_dtype. + + AWQ packs along the N (column) dimension with an interleave permutation + [0,2,4,6,1,3,5,7] applied before packing, so unpacking must undo that. + + Args: + w_int32: packed weights, shape [E, K, N_packed] where N_packed = N//8 + (8 nibbles per int32, packed along N with AWQ interleaving). + scale: per-group scales, shape [E, K//group_size, N], float16. + qzeros: asymmetric zero-points, shape [E, K//gs, N_packed], int32. + None for symmetric (uint4b8 with implicit bias 8). + transpose_output: if True return [E, N, K]; if False return [E, K, N]. + output_dtype: target floating-point dtype (bfloat16 or float16). + + Returns: + Dequantized weight tensor in the requested layout. + """ + E, K, N_packed = w_int32.shape + N = N_packed * 8 + + # Unpack 8 nibbles per int32 along the N dimension (LSB-first) + shifts = torch.arange(8, device=w_int32.device, dtype=torch.int32) * 4 + # [E, K, N_packed, 8] -> [E, K, N_packed*8] = [E, K, N_interleaved] + nibbles = (w_int32.unsqueeze(-1) >> shifts) & 0xF + w_interleaved = nibbles.reshape(E, K, N) # [E, K, N] but column-interleaved + + # Undo AWQ interleave: packed order is [0,2,4,6,1,3,5,7] within each group + # of 8. Inverse: position i in packed -> original column interleave[i]. + # To reverse: we need the inverse permutation so that + # w[:, :, inv_interleave] = w_interleaved gives the natural column order. + interleave = torch.tensor([0, 2, 4, 6, 1, 3, 5, 7], device=w_int32.device) + inv_interleave = torch.empty_like(interleave) + inv_interleave[interleave] = torch.arange(8, device=w_int32.device) + + # Apply inverse interleave within each group of 8 columns + w_reshaped = w_interleaved.reshape(E, K, N // 8, 8) # [E, K, groups, 8] + w_reordered = w_reshaped[:, :, :, inv_interleave] # undo interleave + w = w_reordered.reshape(E, K, N).to(torch.int16) # [E, K, N] + + if qzeros is None: + w = w - 8 + else: + # qzeros: [E, K//gs, N_packed] int32, same AWQ column packing + gs = K // scale.shape[1] + n_gs = scale.shape[1] + zp_nibbles = (qzeros.unsqueeze(-1) >> shifts) & 0xF # [E, n_gs, N_packed, 8] + zp_interleaved = zp_nibbles.reshape(E, n_gs, N) + zp_reshaped = zp_interleaved.reshape(E, n_gs, N // 8, 8) + zp_reordered = zp_reshaped[:, :, :, inv_interleave] + zp = zp_reordered.reshape(E, n_gs, N).to(torch.int16) # [E, n_gs, N] + zp = zp.repeat_interleave(gs, dim=1) # [E, K, N] + w = w - zp + + gs = K // scale.shape[1] + scale_broadcast = scale.repeat_interleave(gs, dim=1).to(output_dtype) # [E, K, N] + + w_dequant = w.to(output_dtype) * scale_broadcast # [E, K, N] + + if transpose_output: + return w_dequant.permute(0, 2, 1).contiguous() # [E, N, K] + return w_dequant.contiguous() # [E, K, N] + + +def _process_weights_emulation_gptq( + w13: torch.Tensor, + w2: torch.Tensor, + w13_scale: torch.Tensor, + w2_scale: torch.Tensor, + w13_qzeros: torch.Tensor | None, + w2_qzeros: torch.Tensor | None, +) -> tuple: + """Dequantize int4 weights to BF16 for the emulation backend. + + Inputs are in GPTQ packed format: + w13: [E, K//8, 2*N] int32 (gate+up proj stacked on dim 2) + w2: [E, N//8, K] int32 + w13_scale: [E, K//gs, 2*N] float16 + w2_scale: [E, N//gs, K] float16 + + Outputs (what TritonExperts expects): + w13_out: [E, 2*N, K] bfloat16 + w2_out: [E, K, N] bfloat16 + """ + # w13: packed along K (dim 1), output cols are 2*N (dim 2) + # transpose_output=True yields [E, 2*N, K] + w13_bf16 = _unpack_and_dequant_int4_gptq( + w13, w13_scale, w13_qzeros, transpose_output=True + ) + + # w2: packed along N (dim 1 is N//8), output cols are K (dim 2) + # After unpacking we get [E, N, K]; we want [E, K, N] for TritonExperts + # transpose_output=False gives [E, N, K], then we permute once more + w2_unpacked = _unpack_and_dequant_int4_gptq( + w2, w2_scale, w2_qzeros, transpose_output=False + ) # [E, N, K] + w2_bf16 = w2_unpacked.permute(0, 2, 1).contiguous() # [E, K, N] + + dummy = torch.ones(1, dtype=torch.float16, device=w13.device) + return ( + w13_bf16, # w13_qweight (now bf16, not int32) + w2_bf16, # w2_qweight (now bf16, not int32) + dummy, # w13_scales (unused; nulled out in Int4EmulationTritonExperts) + dummy, # w2_scales (unused) + None, # w13_g_idx + None, # w2_g_idx + None, # w13_g_idx_sort_indices + None, # w2_g_idx_sort_indices + None, # w13_qzeros + None, # w2_qzeros + None, # w13_input_global_scale + None, # w2_input_global_scale + None, # w13_bias + None, # w2_bias + ) + + +def _process_weights_emulation_awq( + w13: torch.Tensor, + w2: torch.Tensor, + w13_scale: torch.Tensor, + w2_scale: torch.Tensor, + w13_qzeros: torch.Tensor | None, + w2_qzeros: torch.Tensor | None, +) -> tuple: + """Dequantize AWQ int4 weights to BF16 for the emulation backend. + + AWQ inputs: + w13: [E, K, 2*N//8] int32 (packed along N, gate+up on dim 2) + w2: [E, N, K//8] int32 (packed along K) + w13_scale: [E, K//gs, 2*N] float16 + w2_scale: [E, N//gs, K] float16 + + Outputs (what TritonExperts expects): + w13_out: [E, 2*N, K] bfloat16 + w2_out: [E, K, N] bfloat16 + """ + # w13: AWQ-packed along N (dim 2), K is unpacked in dim 1 + # _unpack_and_dequant_int4_awq with transpose_output=True yields [E, 2*N, K] + w13_bf16 = _unpack_and_dequant_int4_awq( + w13, w13_scale, w13_qzeros, transpose_output=True + ) + + # w2: AWQ packs along K (dim 2 is K//8), N is unpacked in dim 1. + # AWQ w2 is [E, N, K//8] — same column-pack format applied to the K dim. + # _unpack_and_dequant_int4_awq expects [E, rows, N_packed] where the + # packed dim is columns. Treat dim 1 as rows and dim 2 as N_packed: + # unpacking gives [E, N, K]. Then permute to [E, K, N]. + w2_unpacked = _unpack_and_dequant_int4_awq( + w2, w2_scale, w2_qzeros, transpose_output=False + ) # [E, N, K] + w2_bf16 = w2_unpacked.permute(0, 2, 1).contiguous() # [E, K, N] + + dummy = torch.ones(1, dtype=torch.float16, device=w13.device) + return ( + w13_bf16, + w2_bf16, + dummy, + dummy, + None, + None, + None, + None, + None, + None, + None, + None, + None, + None, + ) + + def convert_to_wna16_moe_kernel_format( backend: WNA16MoEBackend, layer: torch.nn.Module, @@ -1203,5 +1462,25 @@ def convert_to_wna16_moe_kernel_format( w13_bias_out, w2_bias_out, ) + elif backend == WNA16MoEBackend.EMULATION: + from vllm.model_executor.layers.quantization.auto_awq import AutoAWQConfig + + if isinstance(quant_config, AutoAWQConfig): + return _process_weights_emulation_awq( + w13, + w2, + w13_scale, + w2_scale, + w13_qzeros, + w2_qzeros, + ) + return _process_weights_emulation_gptq( + w13, + w2, + w13_scale, + w2_scale, + w13_qzeros, + w2_qzeros, + ) else: raise ValueError(f"Unsupported wna16 MoE backend: {backend.value}") diff --git a/vllm/model_executor/layers/fused_moe/router/fused_topk_bias_router.py b/vllm/model_executor/layers/fused_moe/router/fused_topk_bias_router.py index 058a96ed6b51..1ddcaa50e83b 100644 --- a/vllm/model_executor/layers/fused_moe/router/fused_topk_bias_router.py +++ b/vllm/model_executor/layers/fused_moe/router/fused_topk_bias_router.py @@ -170,13 +170,12 @@ def fused_topk_bias( hash_indices_table: torch.Tensor | None = None, routed_scaling_factor: float = 1.0, ): - # The topk kernel dispatches dtype based on topk_ids (set by - # indices_type) and assumes input_tokens/hash_indices_table match. - if indices_type is not None: - if input_tokens is not None and input_tokens.dtype != indices_type: - input_tokens = input_tokens.to(dtype=indices_type) - if hash_indices_table is not None and hash_indices_table.dtype != indices_type: - hash_indices_table = hash_indices_table.to(dtype=indices_type) + if ( + input_tokens is not None + and hash_indices_table is not None + and input_tokens.dtype != hash_indices_table.dtype + ): + input_tokens = input_tokens.to(dtype=hash_indices_table.dtype) if not rocm_aiter_ops.is_fused_moe_enabled(): assert hidden_states.size(0) == gating_output.size(0), ( @@ -304,6 +303,7 @@ def fused_topk_bias( scores_for_choice = scores.view(-1, n_routed_experts) # For batch invariance, use sorted=True to ensure deterministic expert selection if hash_indices_table is not None: + assert input_tokens is not None topk_indices = hash_indices_table[input_tokens] else: use_sorted = envs.VLLM_BATCH_INVARIANT diff --git a/vllm/model_executor/layers/layernorm.py b/vllm/model_executor/layers/layernorm.py index 8418245b825b..e8e1f02a5b3a 100644 --- a/vllm/model_executor/layers/layernorm.py +++ b/vllm/model_executor/layers/layernorm.py @@ -102,9 +102,12 @@ def forward_cuda( assert self.variance_size_override is None, ( "Batch invariance is not supported for variance_size_override" ) + pass_weight = ( + self.pass_weight_add if residual is not None else self.pass_weight + ) return rms_norm_batch_invariant( x, - self.weight.data, + self.weight.data if pass_weight else None, self.variance_epsilon, residual=residual, ) diff --git a/vllm/model_executor/layers/linear.py b/vllm/model_executor/layers/linear.py index d92b9fc7d008..9345d1fbf9cf 100644 --- a/vllm/model_executor/layers/linear.py +++ b/vllm/model_executor/layers/linear.py @@ -923,11 +923,21 @@ def load_weights( param: Parameter if "." in name: submodule, _, attr = name.rpartition(".") - param = getattr(self.get_submodule(submodule), attr, self) + target = self.get_submodule(submodule) + param = getattr(target, attr, None) else: - param = getattr(self, name, self) - if param is None and name == "bias": - continue + target = self + param = getattr(self, name, None) + if param is None: + # Checkpoint provides a tensor this layer never registered + # as a param, e.g. bias when the layer was built with + # bias=False, or g_idx when desc_act=False (some GPTQ + # exporters still emit a trivial g_idx per layer regardless). + # Skip rather than crash. + leaf_name = name.rpartition(".")[-1] + if leaf_name in ("bias", "g_idx"): + continue + param = target param.weight_loader(param, loaded_weight, shard_id) logger.debug( "Loaded shard %s with shape %s into %s.%s", @@ -1332,11 +1342,21 @@ def load_weights( param: Parameter if "." in name: submodule, _, attr = name.rpartition(".") - param = getattr(self.get_submodule(submodule), attr, self) + target = self.get_submodule(submodule) + param = getattr(target, attr, None) else: - param = getattr(self, name, self) - if param is None and name == "bias": - continue + target = self + param = getattr(self, name, None) + if param is None: + # Checkpoint provides a tensor this layer never registered + # as a param, e.g. bias when the layer was built with + # bias=False, or g_idx when desc_act=False (some GPTQ + # exporters still emit a trivial g_idx per layer regardless). + # Skip rather than crash. + leaf_name = name.rpartition(".")[-1] + if leaf_name in ("bias", "g_idx"): + continue + param = target param.weight_loader(param, loaded_weight, shard_id) logger.debug( "Loaded shard %s with shape %s into %s.%s", diff --git a/vllm/model_executor/layers/logits_processor.py b/vllm/model_executor/layers/logits_processor.py index eb5b1082ce7e..496a1dd15304 100644 --- a/vllm/model_executor/layers/logits_processor.py +++ b/vllm/model_executor/layers/logits_processor.py @@ -115,14 +115,15 @@ def _apply_head( ) if ( self.head_dtype == torch.float32 - and current_platform.is_cuda() + and (current_platform.is_cuda() or current_platform.is_rocm()) and hidden_states.is_cuda ): # Accumulate the projection directly into fp32. This avoids # materializing an fp32 copy of the lm_head weight on every step, - # unlike casting both operands. `torch.mm(out_dtype=...)` is - # CUDA-only and only supports fp32 output for fp16/bf16 inputs, so - # other cases fall back to the cast path below. + # unlike casting both operands. `torch.mm(out_dtype=...)` only + # supports fp32 output for fp16/bf16 inputs, and is only + # implemented for CUDA and ROCm (the latter via the non-Lt GEMM + # path); other platforms fall back to the cast path below. flat = hidden_states.reshape(-1, hidden_states.shape[-1]) logits = torch.mm(flat, lm_head.weight.t(), out_dtype=self.head_dtype) if embedding_bias is not None: diff --git a/vllm/model_executor/layers/quantization/auto_awq.py b/vllm/model_executor/layers/quantization/auto_awq.py index 58104fa7d256..1e49ec3387e4 100644 --- a/vllm/model_executor/layers/quantization/auto_awq.py +++ b/vllm/model_executor/layers/quantization/auto_awq.py @@ -774,6 +774,9 @@ def get_fused_moe_quant_config(self, layer: RoutedExperts) -> FusedMoEQuantConfi w2_bias=getattr(layer, "w2_bias", None), a1_gscale=getattr(layer, "w13_input_global_scale", None), a2_gscale=getattr(layer, "w2_input_global_scale", None), + gemm1_clamp_limit=getattr(layer, "swiglu_limit", None), + gemm1_alpha=getattr(layer, "swiglu_alpha", None), + gemm1_beta=getattr(layer, "swiglu_beta", None), ) def select_gemm_impl( diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe.py index 00221485233f..7c03baf1e005 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe.py @@ -142,6 +142,23 @@ def get_moe_method( return CompressedTensorsW4A16FlydslMoEMethod( weight_quant, input_quant, layer.moe_config ) + elif moe_backend == "emulation": + # Although this is called 'Marlin', actually it selects + # emulation backend by calling select_wna16_moe_backend. + # TODO: we need to update CompressedTensorsWNA16MoeMethod + # to honor "--moe-backend" option + from .compressed_tensors_moe_wna16_marlin import ( + CompressedTensorsWNA16MarlinMoEMethod, + ) + + logger.info_once( + "Using CompressedTensorsWNA16MarlinMoEMethod " + "(emulation backend requested)" + ) + return CompressedTensorsWNA16MarlinMoEMethod( + weight_quant, input_quant, layer.moe_config, layer_name + ) + from .compressed_tensors_moe_wna16 import ( CompressedTensorsWNA16MoEMethod, ) diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a16_flydsl.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a16_flydsl.py index f8faddbd07bf..f2159b0eb2ac 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a16_flydsl.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a16_flydsl.py @@ -108,7 +108,8 @@ def __init__( # grouped actorder isn't supported by this kernel assert weight_quant.actorder != "group" assert weight_quant.symmetric, ( - "Only symmetric quantization is supported for MoE" + "Only symmetric quantization is supported for MoE. " + "Try --moe-backend emulation." ) def create_weights( diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py index cfeacc902f41..2dabf7a5154b 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py @@ -48,7 +48,8 @@ def __init__( # grouped actorder isn't supported by this kernel assert weight_quant.actorder != "group" assert weight_quant.symmetric, ( - "Only symmetric quantization is supported for MoE" + "Only symmetric quantization is supported for MoE. " + "Try --moe-backend emulation." ) def create_weights( diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_marlin.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_marlin.py index 8af36bcb1022..2b3317d00f3c 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_marlin.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_marlin.py @@ -467,10 +467,16 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: # Marlin-specific parameters (not needed for Flashinfer) if not is_flashinfer: - replace_parameter(layer, "w13_weight_g_idx", w13_g_idx_processed) - replace_parameter(layer, "w2_weight_g_idx", w2_g_idx_processed) - replace_parameter(layer, "w13_g_idx_sort_indices", w13_g_idx_sort_indices) - replace_parameter(layer, "w2_g_idx_sort_indices", w2_g_idx_sort_indices) + if w13_g_idx_processed is not None: + replace_parameter(layer, "w13_weight_g_idx", w13_g_idx_processed) + if w2_g_idx_processed is not None: + replace_parameter(layer, "w2_weight_g_idx", w2_g_idx_processed) + if w13_g_idx_sort_indices is not None: + replace_parameter( + layer, "w13_g_idx_sort_indices", w13_g_idx_sort_indices + ) + if w2_g_idx_sort_indices is not None: + replace_parameter(layer, "w2_g_idx_sort_indices", w2_g_idx_sort_indices) # Register input global scales if present if w13_input_global_scale is not None: @@ -484,8 +490,11 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: torch.nn.Parameter(w2_input_global_scale, requires_grad=False), ) - if self.experts_cls is not None and issubclass( - self.experts_cls, FusedMoEExpertsModular + # Marlin workspace — only needed for Marlin-family backends, not emulation. + if ( + self.experts_cls is not None + and issubclass(self.experts_cls, FusedMoEExpertsModular) + and self.wna16_backend != WNA16MoEBackend.EMULATION ): layer.workspace = marlin_make_workspace_new( layer.w13_weight_g_idx.device, 4 @@ -528,6 +537,9 @@ def get_fused_moe_quant_config( num_bits=self.num_bits, w1_zp=getattr(layer, "w13_weight_zero_point", None), w2_zp=getattr(layer, "w2_weight_zero_point", None), + gemm1_clamp_limit=getattr(layer, "swiglu_limit", None), + gemm1_alpha=getattr(layer, "swiglu_alpha", None), + gemm1_beta=getattr(layer, "swiglu_beta", None), ) def apply_monolithic( diff --git a/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py b/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py index 2fd21c2dff8d..6f0e237785e0 100644 --- a/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py +++ b/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py @@ -353,7 +353,13 @@ def prepare_nvfp4_moe_layer_for_fi_or_cutlass( layer.moe_config.hidden_dim = padded_hidden # Align weights for FI NVFP4 MoE kernels. - min_alignment = 16 if is_gated else 128 + # FlashInfer's TRT-LLM block-scale shuffle asserts the gate/up row dim + # (= up_mult * padded_intermediate, up_mult=2 when gated) is a multiple of + # 128. So gated needs padded_intermediate % 64 (2*64=128); the old value 16 + # left 2*intermediate a multiple of only 32, so an NVFP4 MoE whose rank-local + # intermediate is not 128-aligned at TP>1 (e.g. Gemma-4-26B-A4B at tp4) hit + # `assert M % 128 == 0`. Padded rows are zero -> outputs unchanged. + min_alignment = 64 if is_gated else 128 w13, w13_scale, w2, w2_scale, padded_intermediate = ( align_fp4_moe_weights_for_fi( w13, w13_scale, w2, w2_scale, is_act_and_mul, min_alignment diff --git a/vllm/model_executor/models/deepseek_mtp.py b/vllm/model_executor/models/deepseek_mtp.py index da7f77cae0f9..746f20e74618 100644 --- a/vllm/model_executor/models/deepseek_mtp.py +++ b/vllm/model_executor/models/deepseek_mtp.py @@ -45,9 +45,10 @@ def _restore_full_token_layout_if_needed( hidden_states: torch.Tensor, residual: torch.Tensor, num_tokens: int, + is_sequence_parallel: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: """Restore full token rows for the MTP proposer after SP MoE layers.""" - if hidden_states.shape[0] == num_tokens: + if not is_sequence_parallel and hidden_states.shape[0] == num_tokens: return hidden_states, residual combined_states = torch.cat([hidden_states, residual], dim=-1) @@ -142,6 +143,7 @@ def forward( hidden_states, residual, positions.shape[0], + is_sequence_parallel=self.mtp_block.use_sequence_parallel_moe, ) hidden_states = residual + hidden_states # pre-final-norm (logits hidden) # Recycle the post-final-norm hidden into the next draft step. diff --git a/vllm/model_executor/models/llama_eagle3.py b/vllm/model_executor/models/llama_eagle3.py index 2d859bd4918b..549e8b7bf63a 100644 --- a/vllm/model_executor/models/llama_eagle3.py +++ b/vllm/model_executor/models/llama_eagle3.py @@ -60,7 +60,7 @@ def __init__( self.self_attn.total_num_kv_heads, bias=qkv_bias, quant_config=quant_config, - prefix=maybe_prefix(prefix, "qkv_proj"), + prefix=maybe_prefix(prefix, "self_attn.qkv_proj"), ) self.hidden_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) diff --git a/vllm/model_executor/models/qwen3_dflash.py b/vllm/model_executor/models/qwen3_dflash.py index a03fec0234d2..7843ce72bd48 100644 --- a/vllm/model_executor/models/qwen3_dflash.py +++ b/vllm/model_executor/models/qwen3_dflash.py @@ -52,6 +52,26 @@ logger = init_logger(__name__) +_SLIDING_ATTENTION = "sliding_attention" + + +def _dflash_layer_causal(config: Qwen3Config, layer_idx: int) -> bool: + """``dflash_config.causal`` overrides all layers; else only SWA layers causal.""" + override = (getattr(config, "dflash_config", None) or {}).get("causal") + if override is not None: + return override + layer_types = getattr(config, "layer_types", None) + return bool(layer_types) and layer_types[layer_idx] == _SLIDING_ATTENTION + + +def dflash_has_any_non_causal(config: Qwen3Config) -> bool: + """Whether the draft needs a non-causal-capable backend, resolved from config + (config mirror of the model's ``get_draft_attn_causal``, usable pre-build).""" + return not all( + _dflash_layer_causal(config, i) for i in range(config.num_hidden_layers) + ) + + def _resolve_layer_attention( config: Qwen3Config, layer_idx: int ) -> tuple[int | None, bool]: @@ -79,12 +99,10 @@ def _resolve_layer_attention( dflash_config = getattr(config, "dflash_config", None) or {} layer_types = getattr(config, "layer_types", None) use_swa = dflash_config.get("use_swa", False) - config_causal = dflash_config.get("causal", None) - SLIDING_ATTENTION = "sliding_attention" any_sliding = False if layer_types is not None: - num_sliding = sum(lt == SLIDING_ATTENTION for lt in layer_types) + num_sliding = sum(lt == _SLIDING_ATTENTION for lt in layer_types) any_sliding = num_sliding > 0 # Mixed sliding/full attention needs multiple KV groups (V2 runner only). if ( @@ -97,16 +115,11 @@ def _resolve_layer_attention( "VLLM_USE_V2_MODEL_RUNNER=1." ) - default_causal = False + # ``use_swa`` forces SWA on every layer, even an all-full ``layer_types``. if layer_types is None or (use_swa and not any_sliding): - # An absent ``layer_types`` (or the all-"full_attention" one that may - # be synthesized when the checkpoint omits it) must not override - # ``dflash_config.use_swa``, which forces SWA on every layer. is_sliding = use_swa else: - is_sliding = layer_types[layer_idx] == SLIDING_ATTENTION - # Full-attention layers default non-causal; SWA layers default causal. - default_causal = is_sliding + is_sliding = layer_types[layer_idx] == _SLIDING_ATTENTION sliding_window = None if is_sliding: @@ -119,8 +132,7 @@ def _resolve_layer_attention( "dflash_config.swa_window_size or the top-level sliding_window." ) - causal = config_causal if config_causal is not None else default_causal - return sliding_window, causal + return sliding_window, _dflash_layer_causal(config, layer_idx) class DFlashQwen3Attention(nn.Module): diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index c8486dac5f99..cab03a8a6a26 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -297,6 +297,11 @@ "qwen3_asr_forced_aligner", "Qwen3ASRForcedAlignerForTokenClassification", ), + "RobertaForTokenClassification": ("roberta", "RobertaForTokenClassification"), + "XLMRobertaForTokenClassification": ( + "roberta", + "RobertaForTokenClassification", + ), } _SEQUENCE_CLASSIFICATION_MODELS = { diff --git a/vllm/model_executor/models/roberta.py b/vllm/model_executor/models/roberta.py index a2419fd1c617..948a939b9d58 100644 --- a/vllm/model_executor/models/roberta.py +++ b/vllm/model_executor/models/roberta.py @@ -42,7 +42,7 @@ from .bert_with_rope import BertWithRope, JinaRobertaModel from .interfaces import SupportsCrossEncoding -from .interfaces_base import default_pooling_type +from .interfaces_base import attn_type, default_pooling_type class RobertaEmbedding(nn.Module): @@ -332,3 +332,67 @@ def forward( inputs_embeds=inputs_embeds, intermediate_tensors=intermediate_tensors, ) + + +@attn_type("encoder_only") +@default_pooling_type(tok_pooling_type="ALL") +class RobertaForTokenClassification(nn.Module): + """A model that uses Roberta to provide token classification. + + Mirrors BertForTokenClassification, swapping in RobertaEmbedding for the + RoBERTa/XLM-RoBERTa position-embedding offset and weight layout. + + Also registered as XLMRobertaForTokenClassification since XLM-RoBERTa + checkpoints share RoBERTa's architecture and `roberta.*` weight prefix. + """ + + is_pooling_model = True + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + config = vllm_config.model_config.hf_config + self.head_dtype = vllm_config.model_config.head_dtype + self.num_labels = config.num_labels + self.roberta = BertModel( + vllm_config=vllm_config, + prefix=maybe_prefix(prefix, "roberta"), + embedding_class=RobertaEmbedding, + ) + self.classifier = nn.Linear( + config.hidden_size, config.num_labels, dtype=self.head_dtype + ) + + pooler_config = vllm_config.model_config.pooler_config + assert pooler_config is not None + + self.pooler = pooler_for_token_classify(pooler_config) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.roberta.embed_input_ids(input_ids) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): + loader = AutoWeightsLoader(self) + return loader.load_weights(weights) + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + token_type_ids: torch.Tensor | None = None, + ) -> torch.Tensor: + if token_type_ids is not None: + assert self.roberta.config.vocab_size < (1 << TOKEN_TYPE_SHIFT) + assert input_ids is not None + _encode_token_type_ids(input_ids, token_type_ids) + + hidden_states = self.roberta( + input_ids=input_ids, + positions=positions, + inputs_embeds=inputs_embeds, + intermediate_tensors=intermediate_tensors, + ) + + hidden_states = hidden_states.to(self.head_dtype) + return self.classifier(hidden_states) diff --git a/vllm/models/deepseek_v32/nvidia/model.py b/vllm/models/deepseek_v32/nvidia/model.py index 22ced9fa94d8..353aedc8ceed 100644 --- a/vllm/models/deepseek_v32/nvidia/model.py +++ b/vllm/models/deepseek_v32/nvidia/model.py @@ -7,7 +7,11 @@ import torch from vllm.config import VllmConfig -from vllm.distributed import get_pp_group +from vllm.distributed import ( + get_pp_group, + tensor_model_parallel_all_gather, + tensor_model_parallel_reduce_scatter, +) from vllm.model_executor.layers.fused_moe import ( fused_moe_make_expert_params_mapping, ) @@ -32,6 +36,7 @@ is_pp_missing_parameter, make_empty_intermediate_tensors_factory, make_layers, + sequence_parallel_chunk, ) from vllm.sequence import IntermediateTensors @@ -39,6 +44,18 @@ from .fused_ops import fused_allreduce_rms_norm +def _all_gather_sp_states( + hidden_states: torch.Tensor, + residual: torch.Tensor, + num_tokens: int, +) -> tuple[torch.Tensor, torch.Tensor]: + # combine hidden_states and residual and all gather once + combined_states = torch.cat([hidden_states, residual], dim=-1) + combined_states = tensor_model_parallel_all_gather(combined_states, 0)[:num_tokens] + hidden_states, residual = combined_states.chunk(2, dim=-1) + return hidden_states, residual.contiguous() + + class DeepseekV32DecoderLayer(torch.nn.Module): def __init__( self, @@ -92,6 +109,12 @@ def __init__( prefix=f"{prefix}.mlp", reduce_results=False, ) + self.use_sequence_parallel_moe = ( + parallel_config.use_sequence_parallel_moe + and parallel_config.pipeline_parallel_size == 1 + and isinstance(self.mlp, DeepseekV2MoE) + ) + self.tp_size = parallel_config.tensor_parallel_size self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.post_attention_layernorm = RMSNorm( config.hidden_size, eps=config.rms_norm_eps @@ -104,25 +127,53 @@ def forward( hidden_states: torch.Tensor, residual: torch.Tensor | None, ) -> tuple[torch.Tensor, torch.Tensor]: + full_num_tokens = positions.shape[0] + input_is_sequence_parallel = ( + self.use_sequence_parallel_moe + and residual is not None + and hidden_states.shape[0] != full_num_tokens + ) + if residual is None: # First layer: hidden_states is the (already reduced) embedding. residual = hidden_states hidden_states = self.input_layernorm(hidden_states) + elif input_is_sequence_parallel: + hidden_states, residual = self.input_layernorm(hidden_states, residual) else: # The previous layer's MLP/MoE output is left un-reduced; fuse its # all-reduce into this input_layernorm. hidden_states, residual = fused_allreduce_rms_norm( hidden_states, residual, self.input_layernorm ) - # self_attn's o_proj runs reduce_results=False; fuse its all-reduce with - # the post-attention RMSNorm. + if input_is_sequence_parallel: + hidden_states = tensor_model_parallel_all_gather(hidden_states, 0) + hidden_states = hidden_states[:full_num_tokens] + + # self_attn's o_proj runs reduce_results=False; reduce before RMSNorm. hidden_states = self.self_attn(positions=positions, hidden_states=hidden_states) - hidden_states, residual = fused_allreduce_rms_norm( - hidden_states, residual, self.post_attention_layernorm - ) + if self.use_sequence_parallel_moe: + # small trick using minus, eg. -17 % 8 = 7 + sp_pad = (-hidden_states.shape[0]) % self.tp_size + # pad if not divisible by world size + hidden_states = torch.nn.functional.pad(hidden_states, (0, 0, 0, sp_pad)) + hidden_states = tensor_model_parallel_reduce_scatter(hidden_states, 0) + if not input_is_sequence_parallel: + residual = sequence_parallel_chunk(residual) + hidden_states, residual = self.post_attention_layernorm( + hidden_states, residual + ) + else: + hidden_states, residual = fused_allreduce_rms_norm( + hidden_states, residual, self.post_attention_layernorm + ) + # MLP/MoE runs un-reduced; its all-reduce is fused into the next layer's # input_layernorm (or the model's final norm). - hidden_states = self.mlp(hidden_states) + if self.use_sequence_parallel_moe: + hidden_states = self.mlp(hidden_states, already_sequence_parallel=True) + else: + hidden_states = self.mlp(hidden_states) return hidden_states, residual @@ -206,12 +257,26 @@ def forward( residual = intermediate_tensors["residual"] aux_hidden_states = [] + full_num_tokens = positions.shape[0] for idx, layer in enumerate( islice(self.layers, self.start_layer, self.end_layer), start=self.start_layer, ): + if ( + hidden_states.shape[0] != full_num_tokens + and not layer.use_sequence_parallel_moe + ): + hidden_states, residual = _all_gather_sp_states( + hidden_states, residual, full_num_tokens + ) if idx in self.aux_hidden_state_layers: - aux_hidden_states.append(hidden_states + residual) + aux_hidden_state = hidden_states + residual + if aux_hidden_state.shape[0] != full_num_tokens: + aux_hidden_state = tensor_model_parallel_all_gather( + aux_hidden_state, 0 + ) + aux_hidden_state = aux_hidden_state[:full_num_tokens] + aux_hidden_states.append(aux_hidden_state) hidden_states, residual = layer(positions, hidden_states, residual) if not get_pp_group().is_last_rank: @@ -219,8 +284,15 @@ def forward( {"hidden_states": hidden_states, "residual": residual} ) - # Last layer's MoE output is un-reduced; fuse its all-reduce into norm. - hidden_states, _ = fused_allreduce_rms_norm(hidden_states, residual, self.norm) + if hidden_states.shape[0] != full_num_tokens: + hidden_states, residual = _all_gather_sp_states( + hidden_states, residual, full_num_tokens + ) + hidden_states, _ = self.norm(hidden_states, residual) + else: + hidden_states, _ = fused_allreduce_rms_norm( + hidden_states, residual, self.norm + ) if len(aux_hidden_states) > 0: return hidden_states, aux_hidden_states return hidden_states diff --git a/vllm/models/deepseek_v32/nvidia/mtp.py b/vllm/models/deepseek_v32/nvidia/mtp.py index 8edb3283d07d..d3d8e5aae9c5 100644 --- a/vllm/models/deepseek_v32/nvidia/mtp.py +++ b/vllm/models/deepseek_v32/nvidia/mtp.py @@ -96,11 +96,11 @@ def forward( hidden_states, residual, positions.shape[0], + is_sequence_parallel=self.mtp_block.use_sequence_parallel_moe, ) - # mtp_block's MoE output is left un-reduced (skip_final_all_reduce); the - # main model fuses that all-reduce into the next norm, but here the - # recycle hidden is consumed directly, so reduce it now. - hidden_states = tensor_model_parallel_all_reduce(hidden_states) + if not self.mtp_block.use_sequence_parallel_moe: + # Without sequence parallelism, the MoE output is left un-reduced. + hidden_states = tensor_model_parallel_all_reduce(hidden_states) # Recycle the POST-final-norm hidden into the next draft step. The # residual-add is fused into the final RMSNorm so it is computed # exactly once, and the result is returned for both tuple positions: diff --git a/vllm/models/deepseek_v4/__init__.py b/vllm/models/deepseek_v4/__init__.py index 05662bbd321a..c7dd948f79f8 100644 --- a/vllm/models/deepseek_v4/__init__.py +++ b/vllm/models/deepseek_v4/__init__.py @@ -21,10 +21,9 @@ from .amd.model import DeepseekV4ForCausalLM from .amd.mtp import DeepSeekV4MTP elif current_platform.is_xpu(): + from .xpu.dspark import DSparkDeepseekV4ForCausalLM # type: ignore[assignment] from .xpu.model import DeepseekV4ForCausalLM # type: ignore[assignment] from .xpu.mtp import DeepSeekV4MTP # type: ignore[assignment] - - DSparkDeepseekV4ForCausalLM = None # type: ignore[assignment, misc] else: from .nvidia.dspark import ( # type: ignore[assignment] DSparkDeepseekV4ForCausalLM, diff --git a/vllm/models/deepseek_v4/common/ops/fused_compress_quant_cache.py b/vllm/models/deepseek_v4/common/ops/fused_compress_quant_cache.py index 9a5e478e315f..a2085cd220f1 100644 --- a/vllm/models/deepseek_v4/common/ops/fused_compress_quant_cache.py +++ b/vllm/models/deepseek_v4/common/ops/fused_compress_quant_cache.py @@ -19,6 +19,7 @@ and N_QUANT_BLOCKS ue8m0 bytes. """ +from functools import lru_cache from typing import Any import torch @@ -296,6 +297,360 @@ def _fused_kv_compress_norm_rope_insert_sparse_attn( tl.store(bf16_ptr + rope_local, result.to(tl.bfloat16), mask=is_rope) +# ============================================================================= +# Split kernels variant of the head=512 compressor (deep cr=128 gather). +# - compress gather: instead of launching one program per token, split along +# the head dimension to maximize CU occupancy. The head dimension split +# does not require cross-group reduction +# - finalize norm rope quant store: same as the single pass kernel due to its +# per-token nature +# Mirrors the CUDA cutedsl split kernel where num_splits is occupancy-targeted. +# Currently only tested and validated on ROCm gfx950 +# ============================================================================= +@lru_cache(maxsize=1) +def _n_cu() -> int: + return torch.cuda.get_device_properties(0).multi_processor_count + + +def _pick_compress_num_splits( + num_actual: int, compress_ratio: int, head_dim: int +) -> int: + """Occupancy-targeted column splits for the cr>=128 head=512 compressor. + + Sizes the per-token fan-out so (estimated computing tokens) * num_splits ~ + #CU, capped by head tiling at a 32-wide min tile, as a power-of-2 divisor of + head_dim. + """ + max_splits = head_dim // 32 + est_compute = max(1, num_actual // compress_ratio) + target = -(-_n_cu() // est_compute) # ceil(#CU / est_compute) + ns = 1 + while ns * 2 <= min(target, max_splits) and head_dim % (ns * 2) == 0: + ns *= 2 + return ns + + +@triton.jit +def _compress_gather_split_sparse_attn( + state_cache_ptr, + state_cache_stride0, + state_cache_stride1, + positions_ptr, + slot_mapping_ptr, + token_to_req_indices_ptr, + block_table_ptr, + block_table_stride, + block_size, + scratch_ptr, + scratch_stride, + HEAD_SIZE: tl.constexpr, + STATE_WIDTH: tl.constexpr, + COMPRESS_RATIO: tl.constexpr, + NUM_SPLITS: tl.constexpr, + HEAD_TILE: tl.constexpr, # HEAD_SIZE // NUM_SPLITS +): + """Stage 1: per-(token, head-split) compress gather, write to fp32 scratch + + No-overlap gather (cr>=128) on rows [0, COMPRESS_RATIO) + """ + pid = tl.program_id(0) + token_idx = pid // NUM_SPLITS + split_idx = pid % NUM_SPLITS + + slot_id = tl.load(slot_mapping_ptr + token_idx) + if slot_id < 0: + return + position = tl.load(positions_ptr + token_idx) + if (position + 1) % COMPRESS_RATIO != 0: + return + req_idx = tl.load(token_to_req_indices_ptr + token_idx) + + start = position - COMPRESS_RATIO + 1 + rows = tl.arange(0, COMPRESS_RATIO) + pos = start + rows + mask_pos = pos >= 0 + block_numbers = tl.load( + block_table_ptr + req_idx * block_table_stride + pos // block_size, + mask=mask_pos, + other=0, + ).to(tl.int64) + block_offsets = pos % block_size + + col = split_idx * HEAD_TILE + tl.arange(0, HEAD_TILE) + row_base = ( + state_cache_ptr + + block_numbers * state_cache_stride0 + + block_offsets * state_cache_stride1 + ) + cmask = mask_pos[:, None] + + score = tl.load( + row_base[:, None] + STATE_WIDTH + col[None, :], + mask=cmask, + other=float("-inf"), + ) + score = tl.softmax(score, dim=0) + kv = tl.load(row_base[:, None] + col[None, :], mask=cmask, other=0.0) + compressed = tl.sum(kv * score, axis=0) # [HEAD_TILE] fp32 + tl.store(scratch_ptr + token_idx * scratch_stride + col, compressed) + + +@triton.jit +def _finalize_norm_rope_quant_store_sparse_attn( + scratch_ptr, + scratch_stride, + positions_ptr, + slot_mapping_ptr, + rms_norm_weight_ptr, + rms_norm_eps, + cos_sin_cache_ptr, + cos_sin_stride, + k_cache_ptr, + kv_slot_mapping_ptr, + kv_cache_block_size, + HEAD_SIZE: tl.constexpr, + TRITON_BLOCK_SIZE: tl.constexpr, + COMPRESS_RATIO: tl.constexpr, + ROPE_HEAD_DIM: tl.constexpr, + FP8_MAX: tl.constexpr, + QUANT_BLOCK: tl.constexpr, + TOKEN_STRIDE: tl.constexpr, + SCALE_DIM: tl.constexpr, + KV_BLOCK_STRIDE: tl.constexpr, +): + """Stage 2: read compressed_kv[512] from scratch buffer, then + RMSNorm + FP8 quant (nope) + RoPE + bf16 store + """ + token_idx = tl.program_id(0) + slot_id = tl.load(slot_mapping_ptr + token_idx) + if slot_id < 0: + return + position = tl.load(positions_ptr + token_idx) + if (position + 1) % COMPRESS_RATIO != 0: + return + + block = tl.arange(0, TRITON_BLOCK_SIZE) + mask = block < HEAD_SIZE + compressed_kv = tl.load( + scratch_ptr + token_idx * scratch_stride + block, mask=mask, other=0.0 + ) + + rms_w = tl.load(rms_norm_weight_ptr + block, mask=mask, other=0.0) + variance = tl.sum(compressed_kv * compressed_kv, axis=0) / HEAD_SIZE + rrms = tl.rsqrt(variance + rms_norm_eps) + normed = compressed_kv * rrms * rms_w + + kv_slot_idx = tl.load(kv_slot_mapping_ptr + token_idx) + if kv_slot_idx < 0: + return + kv_block_idx = kv_slot_idx // kv_cache_block_size + kv_pos_in_block = kv_slot_idx % kv_cache_block_size + cache_block_ptr = k_cache_ptr + kv_block_idx.to(tl.int64) * KV_BLOCK_STRIDE + fp8_ptr = cache_block_ptr + kv_pos_in_block * TOKEN_STRIDE + scale_ptr = ( + cache_block_ptr + + kv_cache_block_size * TOKEN_STRIDE + + kv_pos_in_block * SCALE_DIM + ) + + NOPE_HEAD_DIM: tl.constexpr = HEAD_SIZE - ROPE_HEAD_DIM + HALF_ROPE: tl.constexpr = ROPE_HEAD_DIM // 2 + N_QUANT_BLOCKS: tl.constexpr = TRITON_BLOCK_SIZE // QUANT_BLOCK + N_NOPE_BLOCKS: tl.constexpr = NOPE_HEAD_DIM // QUANT_BLOCK + INV_FP8_MAX: tl.constexpr = 1.0 / FP8_MAX + + quant_input = normed.to(tl.bfloat16).to(tl.float32) + quant_2d = tl.reshape(quant_input, (N_QUANT_BLOCKS, QUANT_BLOCK)) + block_absmax = tl.maximum(tl.max(tl.abs(quant_2d), axis=1), 1e-4) + raw_scales = block_absmax * INV_FP8_MAX + exponents = tl.ceil(tl.log2(raw_scales)) + inv_scales = tl.exp2(-exponents) + x_scaled = quant_2d * tl.reshape(inv_scales, (N_QUANT_BLOCKS, 1)) + x_clamped = tl.clamp(x_scaled, -FP8_MAX, FP8_MAX) + x_uint8 = tl.reshape( + x_clamped.to(tl.float8e4nv).to(tl.uint8, bitcast=True), + (TRITON_BLOCK_SIZE,), + ) + tl.store(fp8_ptr + block, x_uint8, mask=block < NOPE_HEAD_DIM) + + scale_idx = tl.arange(0, N_QUANT_BLOCKS) + encoded = tl.maximum(tl.minimum(exponents + 127.0, 255.0), 0.0) + tl.store( + scale_ptr + scale_idx, encoded.to(tl.uint8), mask=scale_idx < N_NOPE_BLOCKS + ) + tl.store(scale_ptr + N_NOPE_BLOCKS, tl.zeros((), dtype=tl.uint8)) + + NUM_PAIRS: tl.constexpr = TRITON_BLOCK_SIZE // 2 + NOPE_PAIRS: tl.constexpr = NOPE_HEAD_DIM // 2 + even, odd = tl.split(tl.reshape(normed, (NUM_PAIRS, 2))) + pair_idx = tl.arange(0, NUM_PAIRS) + rope_pair_local = pair_idx - NOPE_PAIRS + is_rope_pair = rope_pair_local >= 0 + cs_idx = tl.maximum(rope_pair_local, 0) + compressed_pos = (position // COMPRESS_RATIO) * COMPRESS_RATIO + cache_base = cos_sin_cache_ptr + compressed_pos * cos_sin_stride + cos_v = tl.load(cache_base + cs_idx, mask=is_rope_pair, other=1.0) + sin_v = tl.load(cache_base + HALF_ROPE + cs_idx, mask=is_rope_pair, other=0.0) + new_even = even * cos_v - odd * sin_v + new_odd = odd * cos_v + even * sin_v + result = tl.interleave(new_even, new_odd) + bf16_ptr = (fp8_ptr + NOPE_HEAD_DIM).to(tl.pointer_type(tl.bfloat16)) + rope_local = block - NOPE_HEAD_DIM + is_rope = (block >= NOPE_HEAD_DIM) & mask + tl.store(bf16_ptr + rope_local, result.to(tl.bfloat16), mask=is_rope) + + +def _launch_two_stage_sparse_attn_compressor( + state_cache: torch.Tensor, + token_to_req_indices: torch.Tensor, + positions: torch.Tensor, + slot_mapping: torch.Tensor, + block_table: torch.Tensor, + block_size: int, + state_width: int, + compress_ratio: int, + cos_sin_cache: torch.Tensor, + kv_cache: torch.Tensor, + kv_slot_mapping: torch.Tensor, + rms_norm_weight: torch.Tensor, + rms_norm_eps: float, + quant_block: int, + token_stride: int, + scale_dim: int, + head_dim: int, + rope_head_dim: int, + num_actual: int, + compress_scratch: torch.Tensor, +) -> None: + num_splits = _pick_compress_num_splits(num_actual, compress_ratio, head_dim) + head_tile = head_dim // num_splits + scratch = compress_scratch[:num_actual] + _compress_gather_split_sparse_attn[(num_actual * num_splits,)]( + state_cache, + state_cache.stride(0), + state_cache.stride(1), + positions, + slot_mapping, + token_to_req_indices, + block_table, + block_table.stride(0), + block_size, + scratch, + scratch.stride(0), + HEAD_SIZE=head_dim, + STATE_WIDTH=state_width, + COMPRESS_RATIO=compress_ratio, + NUM_SPLITS=num_splits, + HEAD_TILE=head_tile, + ) + _finalize_norm_rope_quant_store_sparse_attn[(num_actual,)]( + scratch, + scratch.stride(0), + positions, + slot_mapping, + rms_norm_weight, + rms_norm_eps, + cos_sin_cache, + cos_sin_cache.stride(0), + kv_cache, + kv_slot_mapping, + kv_cache.shape[1], + HEAD_SIZE=head_dim, + TRITON_BLOCK_SIZE=triton.next_power_of_2(head_dim), + COMPRESS_RATIO=compress_ratio, + ROPE_HEAD_DIM=rope_head_dim, + FP8_MAX=448.0, + QUANT_BLOCK=quant_block, + TOKEN_STRIDE=token_stride, + SCALE_DIM=scale_dim, + KV_BLOCK_STRIDE=kv_cache.stride(0), + ) + + +def compress_norm_rope_store_two_stage_triton( + state_cache: torch.Tensor, + num_actual: int, + token_to_req_indices: torch.Tensor, + positions: torch.Tensor, + slot_mapping: torch.Tensor, + block_table: torch.Tensor, + block_size: int, + state_width: int, + cos_sin_cache: torch.Tensor, + kv_cache: torch.Tensor, + k_cache_metadata: Any, + pdl_kwargs: dict, + head_dim: int, + rope_head_dim: int, + compress_ratio: int, + overlap: bool, + use_fp4_cache: bool, + rms_norm_weight: torch.Tensor, + rms_norm_eps: float, + quant_block: int, + token_stride: int, + scale_dim: int, + num_decode_tokens: int, + compress_scratch: torch.Tensor, +) -> None: + """Two-stage split compressor dispatch for head=512 cr>=128 (no-overlap) + + Run the occupancy-fanned two-stage split for prefill [num_decodee_tokens:] + to fill the CUs, and use the original single-pass launcher + for decode [0, num_decode_tokens) + """ + num_decodes = min(max(num_decode_tokens, 0), num_actual) + num_prefills = num_actual - num_decodes + if num_prefills > 0: + _launch_two_stage_sparse_attn_compressor( + state_cache=state_cache, + token_to_req_indices=token_to_req_indices[num_decodes:], + positions=positions[num_decodes:], + slot_mapping=slot_mapping[num_decodes:], + block_table=block_table, + block_size=block_size, + state_width=state_width, + compress_ratio=compress_ratio, + cos_sin_cache=cos_sin_cache, + kv_cache=kv_cache, + kv_slot_mapping=k_cache_metadata.slot_mapping[num_decodes:], + rms_norm_weight=rms_norm_weight, + rms_norm_eps=rms_norm_eps, + quant_block=quant_block, + token_stride=token_stride, + scale_dim=scale_dim, + head_dim=head_dim, + rope_head_dim=rope_head_dim, + num_actual=num_prefills, + compress_scratch=compress_scratch, + ) + if num_decodes > 0: + compress_norm_rope_store_triton( + state_cache=state_cache, + num_actual=num_decodes, + token_to_req_indices=token_to_req_indices, + positions=positions, + slot_mapping=slot_mapping, + block_table=block_table, + block_size=block_size, + state_width=state_width, + cos_sin_cache=cos_sin_cache, + kv_cache=kv_cache, + k_cache_metadata=k_cache_metadata, + pdl_kwargs=pdl_kwargs, + head_dim=head_dim, + rope_head_dim=rope_head_dim, + compress_ratio=compress_ratio, + overlap=overlap, + use_fp4_cache=use_fp4_cache, + rms_norm_weight=rms_norm_weight, + rms_norm_eps=rms_norm_eps, + quant_block=quant_block, + token_stride=token_stride, + scale_dim=scale_dim, + ) + + # ============================================================================= # Indexer path (head=128, all FP8, single quant block) # ============================================================================= diff --git a/vllm/models/deepseek_v4/compressor.py b/vllm/models/deepseek_v4/compressor.py index 24838c237ce5..13f327f6bc19 100644 --- a/vllm/models/deepseek_v4/compressor.py +++ b/vllm/models/deepseek_v4/compressor.py @@ -14,6 +14,7 @@ from vllm.model_executor.layers.linear import MergedColumnParallelLinear from vllm.models.deepseek_v4.common.ops.fused_compress_quant_cache import ( compress_norm_rope_store_triton, + compress_norm_rope_store_two_stage_triton, ) from vllm.models.deepseek_v4.common.ops.fused_indexer_q import MXFP4_BLOCK_SIZE from vllm.models.deepseek_v4.common.ops.save_partial_states import ( @@ -27,6 +28,7 @@ CommonAttentionMetadata, MultipleOf, ) +from vllm.v1.attention.backends.utils import split_decodes_and_prefills from vllm.v1.kv_cache_interface import ( KVCacheSpec, MLAAttentionSpec, @@ -34,6 +36,12 @@ ) +def _prefer_two_stage_compressor() -> bool: + # Platforms that favor the triton variant of two-stage compressor split. + # Currently only tested on ROCm + return current_platform.is_rocm() + + class CompressorBackend(AttentionBackend): def __init__(self): super().__init__() @@ -81,6 +89,7 @@ class CompressorMetadata: block_size: int token_to_req_indices: torch.Tensor | None = None # [num_tokens] + num_decode_tokens: int | None = None class CompressorMetadataBuilder(AttentionMetadataBuilder): @@ -107,11 +116,17 @@ def build( token_to_req_indices = common_attn_metadata.token_to_req_indices( self.token_to_req_indices ) + num_decode_tokens = None + if _prefer_two_stage_compressor(): + _, _, num_decode_tokens, _ = split_decodes_and_prefills( + common_attn_metadata, decode_threshold=1 + ) return CompressorMetadata( block_table=common_attn_metadata.block_table_tensor.clamp_(min=0), slot_mapping=common_attn_metadata.slot_mapping, block_size=self.block_size, token_to_req_indices=token_to_req_indices, + num_decode_tokens=num_decode_tokens, ) @@ -213,6 +228,25 @@ def __init__( self.overlap = compress_ratio == 4 self.coff = 1 + self.overlap + # The head=512 cr>=128 no-overlap deep gather uses the two-stage + # compressor, which needs an fp32 scratch [max_batched, 512] for + # the intermediate compressed_kv. + # Currently only tested on ROCm + self._use_two_stage_fused_compressor = ( + _prefer_two_stage_compressor() and head_dim == 512 and not self.overlap + ) + self.max_num_batched_tokens = ( + vllm_config.scheduler_config.max_num_batched_tokens + ) + self._compress_scratch: torch.Tensor | None = None + if self._use_two_stage_fused_compressor: + self._compress_scratch = torch.empty( + self.max_num_batched_tokens, + self.head_dim, + dtype=torch.float32, + device=self.device, + ) + state_dtype = torch.float32 self.ape = nn.Parameter( torch.empty( @@ -364,6 +398,15 @@ def forward( store_full_fp8=store_full_fp8, fp8_scale=fp8_scale, ) + elif self._use_two_stage_fused_compressor: + # head=512 cr>=128 (no overlap): two-pass split compressor on the + # prefill suffix, single-pass on the decode prefix. + assert state_metadata.num_decode_tokens is not None + compress_norm_rope_store_fn = compress_norm_rope_store_two_stage_triton + extra_kwargs = { + "num_decode_tokens": state_metadata.num_decode_tokens, + "compress_scratch": self._compress_scratch, + } else: # Indexer path (head_dim == 128) or non-CUDA GPUs (AMD, XPU, etc.). compress_norm_rope_store_fn = compress_norm_rope_store_triton diff --git a/vllm/models/deepseek_v4/xpu/dspark.py b/vllm/models/deepseek_v4/xpu/dspark.py new file mode 100644 index 000000000000..63333a9c0a2e --- /dev/null +++ b/vllm/models/deepseek_v4/xpu/dspark.py @@ -0,0 +1,436 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""DSpark draft model for DeepSeek-V4 on Intel XPU. + +Minimal XPU port of nvidia/dspark.py. Replaces tilelang MHC kernels with +the platform-agnostic custom ops (HCHeadOp, MHCPostOp) already used by the +XPU MTP path, and uses the XPU Triton-based qnorm_rope_kv_fp8_insert for +context KV precomputation. +""" + +from collections.abc import Iterable + +import regex as re +import torch +import torch.nn as nn + +from vllm.config import VllmConfig, get_current_vllm_config +from vllm.distributed import ( + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, +) +from vllm.logger import init_logger +from vllm.model_executor.layers.fused_moe import ( + fused_moe_make_expert_params_mapping, +) +from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.model_executor.layers.linear import ReplicatedLinear +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.mhc import HCHeadOp, MHCPostOp +from vllm.model_executor.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) +from vllm.model_executor.model_loader.weight_utils import default_weight_loader +from vllm.model_executor.models.qwen3_dspark import DSparkMarkovHead +from vllm.model_executor.models.utils import maybe_prefix + +from .model import ( + DeepseekV4DecoderLayer, + make_deepseek_v4_expert_params_mapping, +) + +logger = init_logger(__name__) + +_EXPERT_SCALE_RE = re.compile(r"\.experts\.\d+\.w[123]\.scale$") + + +class DSparkDeepseekV4Model(nn.Module): + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + super().__init__() + assert vllm_config.speculative_config is not None + config = vllm_config.speculative_config.draft_model_config.hf_config + self.config = config + self.hidden_size = config.hidden_size + self.hc_mult = config.hc_mult + self.hc_eps = config.hc_eps + self.rms_norm_eps = config.rms_norm_eps + self.num_hidden_layers = config.num_hidden_layers + self.target_layer_ids = tuple(config.dspark_target_layer_ids) + + self.num_dspark_layers = getattr(config, "n_mtp_layers", None) or 3 + + # Shared with target (aliased by speculator loading utility). + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + prefix=maybe_prefix(prefix, "embed_tokens"), + ) + + self.main_proj = ReplicatedLinear( + config.hidden_size * len(self.target_layer_ids), + config.hidden_size, + bias=False, + return_bias=False, + quant_config=vllm_config.quant_config, + prefix=maybe_prefix(prefix, "main_proj"), + ) + self.main_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + current_vllm_config = get_current_vllm_config() + self.layers = nn.ModuleList( + [ + DeepseekV4DecoderLayer( + current_vllm_config, + prefix=maybe_prefix(prefix, f"layers.{self.num_hidden_layers + i}"), + ) + for i in range(self.num_dspark_layers) + ] + ) + + # Heads: final norm + hc_head, and the Markov head + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + hc_dim = self.hc_mult * config.hidden_size + self.hc_head_fn = nn.Parameter( + torch.empty(self.hc_mult, hc_dim, dtype=torch.float32), + requires_grad=False, + ) + self.hc_head_base = nn.Parameter( + torch.empty(self.hc_mult, dtype=torch.float32), requires_grad=False + ) + self.hc_head_scale = nn.Parameter( + torch.empty(1, dtype=torch.float32), requires_grad=False + ) + draft_vocab_size = ( + getattr(config, "draft_vocab_size", None) or config.vocab_size + ) + self.markov_head = DSparkMarkovHead( + config.vocab_size, + draft_vocab_size, + config.dspark_markov_rank, + prefix=maybe_prefix(prefix, "markov_head"), + ) + + # XPU MHC ops (replaces tilelang) + self.mhc_post_op = MHCPostOp() + self.hc_head_op = HCHeadOp() + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.embed_tokens(input_ids) + + def combine_hidden_states(self, aux_hidden_states: torch.Tensor) -> torch.Tensor: + """main_x = main_norm(main_proj(concat of target aux hidden states)).""" + return self.main_norm(self.main_proj(aux_hidden_states)) + + @torch.inference_mode() + def precompute_and_store_context_kv( + self, + main_x: torch.Tensor, + context_positions: torch.Tensor, + context_slot_mappings: list[torch.Tensor | None] | None = None, + ) -> None: + """Insert the sliding-window context KV for every draft layer. + + Each layer derives its context KV from the SAME projected target hidden + ``main_x``, via that layer's own wkv + kv_norm + RoPE + quant, then + writes it at the layer's context slots. + """ + for i, layer in enumerate(self.layers): + slot_mapping = ( + None if context_slot_mappings is None else context_slot_mappings[i] + ) + attn = layer.attn + # wkv part of the fused wq_a|wkv projection (q_lora part discarded) + qr_kv, _ = attn.fused_wqa_wkv(main_x) + kv = qr_kv[..., attn.q_lora_rank :] + kv = attn.kv_norm(kv) + if slot_mapping is None: + continue + _insert_context_kv(attn, kv, context_positions, slot_mapping) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + inputs_embeds: torch.Tensor | None = None, + ) -> torch.Tensor: + if inputs_embeds is None: + inputs_embeds = self.embed_input_ids(input_ids) + # Expand to hc_mult copies for hyper-connections ([T, H] -> [T, hc, H]). + hidden_states = inputs_embeds.unsqueeze(-2).repeat(1, self.hc_mult, 1) + + residual = post_mix = res_mix = None + for layer in self.layers: + hidden_states, residual, post_mix, res_mix = layer( + hidden_states, + positions, + input_ids, + post_mix, + res_mix, + residual, + ) + # mhc_post: merge hyper-connection copies + hidden_states = self.mhc_post_op(hidden_states, residual, post_mix, res_mix) + # hc_head: reduces hc copies; return pre-norm head hidden + hidden_states = self.hc_head_op( + hidden_states, + self.hc_head_fn, + self.hc_head_scale, + self.hc_head_base, + self.rms_norm_eps, + self.hc_eps, + ) + return hidden_states + + +def _insert_context_kv( + attn: nn.Module, + kv: torch.Tensor, + positions: torch.Tensor, + slot_mapping: torch.Tensor, +) -> None: + """RoPE + quant + paged-cache insert of (already kv_norm'd) context KV. + + On XPU, we reuse the same xpu_qnorm_rope_kv_fp8_insert kernel used in + the forward path, passing a dummy q (result discarded). + """ + from vllm.models.deepseek_v4.xpu.xpu_qnorm_rope_kv_fp8_insert import ( + xpu_qnorm_rope_kv_fp8_insert, + ) + + swa_cache = attn.swa_cache_layer.kv_cache + block_size = attn.swa_cache_layer.block_size + cos_sin_cache = attn.rotary_emb.cos_sin_cache + n_ctx = kv.shape[0] + + # Dummy q — we only care about the KV insert side effect. + dummy_q = torch.empty( + (n_ctx, attn.n_local_heads, attn.head_dim), + dtype=kv.dtype, + device=kv.device, + ) + xpu_qnorm_rope_kv_fp8_insert( + dummy_q, + kv, + swa_cache, + slot_mapping, + positions, + cos_sin_cache, + attn.eps, + block_size, + ) + + +class DSparkDeepseekV4ForCausalLM(nn.Module): + """XPU DSpark draft model entry point for DeepSeek-V4.""" + + has_own_embed_tokens = False + has_own_lm_head = False + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + super().__init__() + assert vllm_config.speculative_config is not None + self.draft_model_config = vllm_config.speculative_config.draft_model_config + self.config = self.draft_model_config.hf_config + self.model = DSparkDeepseekV4Model( + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") + ) + # Shared with the target (aliased by the speculator's load utility). + self.lm_head = ParallelLMHead( + self.config.vocab_size, + self.config.hidden_size, + prefix=maybe_prefix(prefix, "lm_head"), + ) + self.logits_processor = LogitsProcessor(self.config.vocab_size) + + # --- Hooks used by the speculator --- + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.model.embed_input_ids(input_ids) + + def combine_hidden_states(self, aux_hidden_states: torch.Tensor) -> torch.Tensor: + return self.model.combine_hidden_states(aux_hidden_states) + + def get_draft_kv_cache_layer_names(self) -> list[str]: + return [layer.attn.swa_cache_layer.prefix for layer in self.model.layers] + + def precompute_and_store_context_kv( + self, + context_states: torch.Tensor, + context_positions: torch.Tensor, + context_slot_mappings: list[torch.Tensor | None] | None = None, + ) -> None: + self.model.precompute_and_store_context_kv( + context_states, context_positions, context_slot_mappings + ) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + inputs_embeds: torch.Tensor | None = None, + ) -> torch.Tensor: + return self.model(input_ids, positions, inputs_embeds) + + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor: + """Base logits U_k = lm_head(norm(head_hidden)).""" + return self.logits_processor(self.lm_head, self.model.norm(hidden_states)) + + def compute_draft_logits(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.compute_logits(hidden_states) + + def map_draft_to_target(self, draft_ids: torch.Tensor) -> torch.Tensor: + return draft_ids # full-vocab: draft ids are target ids + + def markov_embed(self, token_ids: torch.Tensor) -> torch.Tensor: + return self.model.markov_head.embed(token_ids) + + def markov_bias(self, markov_embed: torch.Tensor) -> torch.Tensor: + return self.model.markov_head.bias(markov_embed, self.logits_processor) + + # --- Weight loading --- + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + """Load ``mtp.{0,1,2}.*`` draft weights from the target checkpoint.""" + first_layer = self.model.layers[0] + use_mega_moe = first_layer.ffn.use_mega_moe + if use_mega_moe: + expert_mapping = make_deepseek_v4_expert_params_mapping( + self.config.n_routed_experts + ) + else: + expert_mapping = fused_moe_make_expert_params_mapping( + self, + ckpt_gate_proj_name="w1", + ckpt_down_proj_name="w2", + ckpt_up_proj_name="w3", + num_experts=self.config.n_routed_experts, + ) + expert_scale_suffix = ( + ".weight_scale" + if getattr(self.config, "expert_dtype", "fp4") == "fp4" + else ".weight_scale_inv" + ) + + stacked_params_mapping = [ + ("gate_up_proj", "w1", 0), + ("gate_up_proj", "w3", 1), + ("attn.fused_wqa_wkv", "attn.wq_a", 0), + ("attn.fused_wqa_wkv", "attn.wkv", 1), + ] + + params_dict = dict(self.named_parameters()) + loaded_params: set[str] = set() + + tp_size = get_tensor_model_parallel_world_size() + tp_rank = get_tensor_model_parallel_rank() + n_local_head = self.config.num_attention_heads // tp_size + head_start = n_local_head * tp_rank + head_end = n_local_head * (tp_rank + 1) + + for name, loaded_weight in weights: + mapped = self._remap_dspark_name(name) + if mapped is None: + continue + name = mapped + + # .scale -> per-method scale suffix + if name.endswith(".scale"): + suffix = ( + expert_scale_suffix + if _EXPERT_SCALE_RE.search(name) + else ".weight_scale_inv" + ) + name = name.removesuffix(".scale") + suffix + + # Expert weights + if ".experts." in name: + if ( + "weight_scale" in name + and loaded_weight.dtype == torch.float8_e8m0fnu + ): + loaded_weight = loaded_weight.view(torch.uint8) + for param_name, weight_name, expert_id, shard_id in expert_mapping: + if weight_name not in name: + continue + name_mapped = name.replace(weight_name, param_name) + if name_mapped not in params_dict: + continue + param = params_dict[name_mapped] + success = param.weight_loader( + param, + loaded_weight, + name_mapped, + shard_id=shard_id, + expert_id=expert_id, + return_success=True, + ) + if success: + loaded_params.add(name_mapped) + break + continue + + # Stacked params (decoder-layer only) + is_layer_param = name.startswith("model.layers.") + for param_name, weight_name, stacked_shard_id in stacked_params_mapping: + if not is_layer_param or weight_name not in name: + continue + name = name.replace(weight_name, param_name) + if name not in params_dict: + break + param = params_dict[name] + param.weight_loader(param, loaded_weight, stacked_shard_id) + loaded_params.add(name) + break + else: + if "attn_sink" in name: + if name not in params_dict: + continue + narrow = loaded_weight[head_start:head_end] + params_dict[name][: narrow.shape[0]].copy_(narrow) + loaded_params.add(name) + continue + if ".shared_experts.w2" in name: + name = name.replace( + ".shared_experts.w2", ".shared_experts.down_proj" + ) + if name.endswith(".ffn.gate.bias"): + name = name.replace( + ".ffn.gate.bias", ".ffn.gate.e_score_correction_bias" + ) + if name not in params_dict: + continue + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + loaded_params.add(name) + + self._finalize_moe() + logger.info_once("DSpark XPU draft model loaded: %d params", len(loaded_params)) + return loaded_params + + def _finalize_moe(self) -> None: + for layer in self.model.layers: + layer.ffn.finalize_mega_moe_weights() + + def _remap_dspark_name(self, name: str) -> str | None: + """Map checkpoint ``mtp.{i}.*`` name to this model's parameter path.""" + m = re.match(r"mtp\.(\d+)\.(.*)", name) + if m is None: + return None + stage = int(m.group(1)) + rest = m.group(2) + if rest.startswith("confidence_head."): + return None + head_prefixes = ( + "norm.", + "hc_head_fn", + "hc_head_base", + "hc_head_scale", + "markov_head.", + ) + if rest.startswith(("main_proj.", "main_norm.")) or rest.startswith( + head_prefixes + ): + return f"model.{rest}" + return f"model.layers.{stage}.{rest}" diff --git a/vllm/models/deepseek_v4/xpu/model.py b/vllm/models/deepseek_v4/xpu/model.py index 1e5a574bed4a..e8449b9c058b 100644 --- a/vllm/models/deepseek_v4/xpu/model.py +++ b/vllm/models/deepseek_v4/xpu/model.py @@ -44,7 +44,11 @@ VocabParallelEmbedding, ) from vllm.model_executor.model_loader.weight_utils import default_weight_loader -from vllm.model_executor.models.interfaces import SupportsPP +from vllm.model_executor.models.interfaces import ( + EagleModelMixin, + SupportsEagle3, + SupportsPP, +) from vllm.model_executor.models.utils import ( AutoWeightsLoader, PPMissingLayer, @@ -975,7 +979,7 @@ def forward( @support_torch_compile -class DeepseekV4Model(nn.Module): +class DeepseekV4Model(nn.Module, EagleModelMixin): def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() @@ -1113,7 +1117,11 @@ def forward( input_ids = input_ids.to(torch.int64) residual, post_mix, res_mix = None, None, None - for layer in islice(self.layers, self.start_layer, self.end_layer): + aux_hidden_states: list[torch.Tensor] = [] + for idx, layer in enumerate( + islice(self.layers, self.start_layer, self.end_layer), + start=self.start_layer, + ): hidden_states, residual, post_mix, res_mix = layer( hidden_states, positions, @@ -1122,6 +1130,9 @@ def forward( res_mix, residual, ) + if idx + 1 in self.aux_hidden_state_layers: + aux_recon = layer.hc_post(hidden_states, residual, post_mix, res_mix) + aux_hidden_states.append(aux_recon.mean(dim=1)) # The fused path defers the final hc_post to the next layer's # fused_post_pre. After the last layer we must apply it explicitly. if layer is not None: @@ -1143,6 +1154,8 @@ def forward( self.hc_eps, ) hidden_states = self.norm(hidden_states) + if len(aux_hidden_states) > 0: + return hidden_states, aux_hidden_states return hidden_states def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: @@ -1300,7 +1313,7 @@ def _make_deepseek_v4_weights_mapper(expert_dtype: str) -> WeightsMapper: ) -class DeepseekV4ForCausalLM(nn.Module, SupportsPP): +class DeepseekV4ForCausalLM(nn.Module, SupportsPP, SupportsEagle3): model_cls = DeepseekV4Model # Default mapper assumes the original FP4-expert checkpoint layout. diff --git a/vllm/platforms/cuda.py b/vllm/platforms/cuda.py index 9eac95e03249..4059a7e4f50d 100644 --- a/vllm/platforms/cuda.py +++ b/vllm/platforms/cuda.py @@ -85,6 +85,7 @@ def _get_backend_priorities( device_capability: DeviceCapability, num_heads: int | None = None, kv_cache_dtype: CacheDType | None = None, + use_non_causal: bool = False, ) -> list[AttentionBackendEnum]: """Get backend priorities with lazy import to avoid circular dependency.""" from vllm.utils.torch_utils import is_quantized_kv_cache @@ -141,7 +142,10 @@ def _get_backend_priorities( AttentionBackendEnum.FLASHMLA_SPARSE, ] else: - if device_capability.major == 10: + # SM100f defaults to FlashInfer for TRTLLM causal attention, but its non-causal + # cutlass path (used for dflash attention) is known to have problems. + # So prefer FlashAttention when non-causal on SM100f. + if device_capability.major == 10 and not use_non_causal: return [ AttentionBackendEnum.FLASHINFER, AttentionBackendEnum.FLASH_ATTN, @@ -368,6 +372,7 @@ def get_valid_backends( device_capability, num_heads, attn_selector_config.kv_cache_dtype, + attn_selector_config.use_non_causal, ) for priority, backend in enumerate(backend_priorities): try: diff --git a/vllm/transformers_utils/config.py b/vllm/transformers_utils/config.py index c0f3396b0bb9..bb72087a95d4 100644 --- a/vllm/transformers_utils/config.py +++ b/vllm/transformers_utils/config.py @@ -133,6 +133,14 @@ def __getitem__(self, key): _PATCH_HF_VALIDATE_ROPE: set[str] = {"sarvam_mla"} +# Model types whose checkpoints declare `layer_types` entries that upstream +# transformers has not added to `ALLOWED_LAYER_TYPES` yet, so its strict config +# validation rejects them (e.g. GLM-5.2 `glm_moe_dsa` use +# `deepseek_sparse_attention`). Extend the allowed set for these model types. +_PATCH_HF_ALLOWED_LAYER_TYPES: dict[str, tuple[str, ...]] = { + "glm_moe_dsa": ("deepseek_sparse_attention",), +} + _CONFIG_ATTRS_MAPPING: dict[str, str] = { "llm_config": "text_config", } @@ -206,6 +214,24 @@ def patched_validate_rope(self, *args, **kwargs): PretrainedConfig.validate_rope = patched_validate_rope +def _patch_hf_transformers_allowed_layer_types( + extra_layer_types: tuple[str, ...], +) -> None: + """Extend transformers' ``ALLOWED_LAYER_TYPES`` so its strict config + validation accepts layer types (e.g. ``deepseek_sparse_attention``) that a + checkpoint declares but upstream transformers has not registered yet. + """ + import transformers.configuration_utils as hf_configuration_utils + + missing = tuple( + layer_type + for layer_type in extra_layer_types + if layer_type not in hf_configuration_utils.ALLOWED_LAYER_TYPES + ) + if missing: + hf_configuration_utils.ALLOWED_LAYER_TYPES += missing + + class HFConfigParser(ConfigParserBase): def parse( self, @@ -248,6 +274,9 @@ def parse( if model_type in _PATCH_HF_VALIDATE_ROPE: _patch_hf_transformers_validate_rope() + if extra_layer_types := _PATCH_HF_ALLOWED_LAYER_TYPES.get(model_type): + _patch_hf_transformers_allowed_layer_types(extra_layer_types) + if model_type in _SPECULATIVE_DECODING_CONFIGS: config_class = _CONFIG_REGISTRY[model_type] config = config_class.from_pretrained( diff --git a/vllm/utils/jit_monitor.py b/vllm/utils/jit_monitor.py index 8228e24c1c3b..4c3327283a3c 100644 --- a/vllm/utils/jit_monitor.py +++ b/vllm/utils/jit_monitor.py @@ -267,6 +267,26 @@ def _log_cutedsl_jit_compile(fn_name: str) -> None: ) +class _MonitoredCuteCompile: + """Logs JIT compilations; a plain function would break ``cute.compile[opts]``.""" + + def __init__(self, inner): + self._inner = inner + + def __getitem__(self, options) -> "_MonitoredCuteCompile": + return _MonitoredCuteCompile(self._inner[options]) + + def __call__(self, *args, **kwargs): + kernel = args[0] if args else kwargs.get("function") + kernel_name = getattr(kernel, "__name__", None) + if kernel_name is None: + kernel_name = ( + kernel.__class__.__name__ if kernel is not None else "" + ) + _log_cutedsl_jit_compile(kernel_name) + return self._inner(*args, **kwargs) + + def _setup_cutedsl_jit_hook() -> None: """Wrap ``cutlass.cute.compile`` to warn on compilation.""" global _cutedsl_hook_installed @@ -279,20 +299,7 @@ def _setup_cutedsl_jit_hook() -> None: logger.debug("CuTeDSL is not available; skipping CuTeDSL JIT monitor.") return - original_compile = cute.compile - - @functools.wraps(original_compile) - def _compile_with_monitor(*args, **kwargs): - kernel = args[0] if args else kwargs.get("function") - kernel_name = getattr(kernel, "__name__", None) - if kernel_name is None: - kernel_name = ( - kernel.__class__.__name__ if kernel is not None else "" - ) - _log_cutedsl_jit_compile(kernel_name) - return original_compile(*args, **kwargs) - - cute.compile = _compile_with_monitor + cute.compile = _MonitoredCuteCompile(cute.compile) _cutedsl_hook_installed = True diff --git a/vllm/v1/attention/backends/flashinfer.py b/vllm/v1/attention/backends/flashinfer.py index da4e59d4eaaf..9eb1cdbef321 100755 --- a/vllm/v1/attention/backends/flashinfer.py +++ b/vllm/v1/attention/backends/flashinfer.py @@ -927,7 +927,8 @@ def get_cudagraph_support( has_trtllm_support = False break - if has_trtllm_support: + # trtllm-gen only supports causal attention. + if has_trtllm_support and not vllm_config.attention_config.use_non_causal: return AttentionCGSupport.UNIFORM_BATCH else: return AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE diff --git a/vllm/v1/attention/backends/mla/triton_mla.py b/vllm/v1/attention/backends/mla/triton_mla.py index acc9c9cb5010..5c10c16a78f4 100644 --- a/vllm/v1/attention/backends/mla/triton_mla.py +++ b/vllm/v1/attention/backends/mla/triton_mla.py @@ -180,6 +180,32 @@ def __init__( "TritonMLAImpl" ) + if current_platform.is_cuda(): + cap = current_platform.get_device_capability() + cap_str = cap.as_version_str() if cap is not None else "unknown" + dev = current_platform.get_device_name() + if self.kv_cache_dtype.startswith("fp8") and not ( + current_platform.has_device_capability(89) + ): + suggested = ( + "float16" if (cap is None or cap.to_int() < 80) else "bfloat16" + ) + raise ValueError( + f"FP8 KV cache is not supported by the Triton MLA backend " + f"on {dev} (compute capability {cap_str}); native FP8 " + f"(fp8e4nv) requires SM89+. Re-run with " + f"--kv-cache-dtype {suggested}." + ) + if self.kv_cache_dtype == "bfloat16" and not ( + current_platform.has_device_capability(80) + ): + raise ValueError( + f"bfloat16 KV cache is not supported by the Triton MLA " + f"backend on {dev} (compute capability {cap_str}); " + f"bfloat16 requires SM80+. Re-run with " + f"--kv-cache-dtype float16." + ) + # For FP8 KV cache, we dequantize to BF16 on load inside the # Triton kernel. Tell the common layer not to quantize queries # to FP8 — we handle FP8 KV cache with BF16 queries (Mode 1). diff --git a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py index 2153a460f696..c79dbfd5a80f 100644 --- a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py @@ -1187,10 +1187,12 @@ def _sparse_attn_prefill_ragged_kernel( kv_len = kv_end - kv_start k_offsets = tl.arange(0, BLOCK_K) + slot = tl.load( + kv_indices_ptr + kv_start + k_offsets, mask=k_offsets < kv_len, other=-1 + ) for k_start in tl.range(0, kv_len, BLOCK_K): k_pos = k_start + k_offsets in_range = k_pos < kv_len - slot = tl.load(kv_indices_ptr + kv_start + k_pos, mask=in_range, other=-1) valid = in_range & (slot >= 0) & (slot < num_kv) safe_slot = tl.where(valid, slot, 0) @@ -1201,7 +1203,11 @@ def _sparse_attn_prefill_ragged_kernel( mask=valid[:, None] & dim_mask[None, :], other=0.0, ) - kv = tl.where(valid[:, None] & dim_mask[None, :], kv, 0.0) + + next_k_pos = k_start + BLOCK_K + k_offsets + slot = tl.load( + kv_indices_ptr + kv_start + next_k_pos, mask=next_k_pos < kv_len, other=-1 + ) scores = tl.dot(q, tl.trans(kv)) * scale scores = tl.where(head_mask[:, None] & valid[None, :], scores, neg_large) @@ -1865,6 +1871,7 @@ def _rocm_sparse_attn_prefill_ragged_triton( block_h = 16 block_d = triton.next_power_of_2(head_dim) block_k = 16 if head_dim >= 256 else 32 + num_warps = 4 out = torch.empty_like(q, dtype=torch.bfloat16) _sparse_attn_prefill_ragged_kernel[(num_queries, triton.cdiv(num_heads, block_h))]( q, @@ -1889,7 +1896,7 @@ def _rocm_sparse_attn_prefill_ragged_triton( BLOCK_H=block_h, BLOCK_D=block_d, BLOCK_K=block_k, - num_warps=8, + num_warps=num_warps, ) return out diff --git a/vllm/v1/core/sched/output.py b/vllm/v1/core/sched/output.py index 4401fb050b3f..5667482e948b 100644 --- a/vllm/v1/core/sched/output.py +++ b/vllm/v1/core/sched/output.py @@ -179,6 +179,14 @@ def make_empty(cls) -> "CachedRequestData": ) +@dataclass +class ScheduledEncoderInputStats: + """Stats for encoder inputs scheduled in one iteration.""" + + num_inputs: int = 0 + output_tokens: int = 0 + + @dataclass class SchedulerOutput: # list of the requests that are scheduled for the first time. @@ -216,6 +224,8 @@ class SchedulerOutput: # freed from the encoder cache. free_encoder_mm_hashes: list[str] + scheduled_encoder_input_stats: ScheduledEncoderInputStats | None = None + # Request IDs that are preempted in this step. # Only used for v2 model runner. preempted_req_ids: set[str] | None = None diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 371139fa7226..bde6af0e66bb 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -44,6 +44,7 @@ CachedRequestData, GrammarOutput, NewRequestData, + ScheduledEncoderInputStats, SchedulerOutput, ) from vllm.v1.core.sched.request_queue import ( @@ -1121,6 +1122,15 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: len(num_scheduled_tokens) ] + scheduled_encoder_input_stats = None + if ( + self.log_stats + and self.observability_config.enable_logging_iteration_details + ): + scheduled_encoder_input_stats = self._make_scheduled_encoder_input_stats( + scheduled_encoder_inputs + ) + scheduler_output = SchedulerOutput( scheduled_new_reqs=new_reqs_data, scheduled_cached_reqs=cached_reqs_data, @@ -1128,6 +1138,7 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: total_num_scheduled_tokens=total_num_scheduled_tokens, scheduled_spec_decode_tokens=scheduled_spec_decode_tokens, scheduled_encoder_inputs=scheduled_encoder_inputs, + scheduled_encoder_input_stats=scheduled_encoder_input_stats, num_common_prefix_blocks=num_common_prefix_blocks, preempted_req_ids=self.reset_preempted_req_ids, # finished_req_ids is an existing state in the scheduler, @@ -1506,6 +1517,23 @@ def _try_schedule_encoder_inputs( external_load_encoder_input, ) + def _make_scheduled_encoder_input_stats( + self, scheduled_encoder_inputs: dict[str, list[int]] + ) -> ScheduledEncoderInputStats | None: + stats = ScheduledEncoderInputStats() + + for req_id, input_ids in scheduled_encoder_inputs.items(): + request = self.requests.get(req_id) + if request is None: + continue + + for input_id in input_ids: + mm_feature = request.mm_features[input_id] + stats.num_inputs += 1 + stats.output_tokens += mm_feature.mm_position.get_num_embeds() + + return stats if stats.num_inputs else None + def get_grammar_bitmask( self, scheduler_output: SchedulerOutput ) -> GrammarOutput | None: @@ -1877,7 +1905,10 @@ def update_from_output( if ( stats := self.make_stats( - spec_decoding_stats, kv_connector_stats, cudagraph_stats, perf_stats + spec_decoding_stats, + kv_connector_stats, + cudagraph_stats, + perf_stats, ) ) is not None: # Return stats to only one of the front-ends. diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index 301c5892a58a..383853807db6 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -79,16 +79,17 @@ ) from vllm.v1.executor import Executor from vllm.v1.kv_cache_interface import KVCacheConfig, get_kv_cache_spec_kind -from vllm.v1.metrics.stats import SchedulerStats +from vllm.v1.metrics.stats import SchedulerIterationDetails, SchedulerStats from vllm.v1.outputs import ModelRunnerOutput from vllm.v1.request import Request, RequestStatus from vllm.v1.serial_utils import MsgpackDecoder, MsgpackEncoder from vllm.v1.structured_output import StructuredOutputManager -from vllm.v1.utils import IterationDetails, compute_iteration_details +from vllm.v1.utils import compute_iteration_details from vllm.version import __version__ as VLLM_VERSION logger = init_logger(__name__) + HANDSHAKE_TIMEOUT_MINS = 5 _R = TypeVar("_R") # Return type for collective_rpc @@ -498,45 +499,74 @@ def log_error_detail(self, scheduler_output: SchedulerOutput): raise err @contextmanager - def log_iteration_details(self, scheduler_output: SchedulerOutput | None): - if not self.vllm_config.observability_config.enable_logging_iteration_details: - yield + def capture_iteration_details( + self, scheduler_output: SchedulerOutput | None + ) -> Generator[SchedulerIterationDetails | None, None, None]: + enable_details = ( + self.vllm_config.observability_config.enable_logging_iteration_details + ) + if not self.log_stats or not enable_details: + yield None return # 0-token step: let the dummy_batch wrapper log it (avoids double-log). - if scheduler_output and scheduler_output.total_num_scheduled_tokens == 0: - yield + if ( + scheduler_output is not None + and scheduler_output.total_num_scheduled_tokens == 0 + ): + yield None return - self._iteration_index = getattr(self, "_iteration_index", 0) + + iteration_index = getattr(self, "_iteration_index", 0) # scheduler_output=None marks a DP dummy iteration. if scheduler_output is None: - iteration_details = IterationDetails(0, 0, 0, 0) - is_dummy = True + iteration_details = SchedulerIterationDetails( + iteration_index=iteration_index, + num_ctx_requests=0, + num_ctx_tokens=0, + num_generation_requests=0, + num_generation_tokens=0, + elapsed_ms=0.0, + is_dummy=True, + ) else: - iteration_details = compute_iteration_details(scheduler_output) - is_dummy = False - before = time.monotonic() - yield - logger.info( - "".join( - [ - "Iteration(", - str(self._iteration_index), - "): ", - str(iteration_details.num_ctx_requests), - " context requests, ", - str(iteration_details.num_ctx_tokens), - " context tokens, ", - str(iteration_details.num_generation_requests), - " generation requests, ", - str(iteration_details.num_generation_tokens), - " generation tokens, iteration elapsed time: ", - format((time.monotonic() - before) * 1000, ".2f"), - " ms", - " (dummy)" if is_dummy else "", - ] + details = compute_iteration_details(scheduler_output) + iteration_details = SchedulerIterationDetails( + iteration_index=iteration_index, + num_ctx_requests=details.num_ctx_requests, + num_ctx_tokens=details.num_ctx_tokens, + num_generation_requests=details.num_generation_requests, + num_generation_tokens=details.num_generation_tokens, + elapsed_ms=0.0, + num_encoder_inputs=details.num_encoder_inputs, + num_encoder_output_tokens=details.num_encoder_output_tokens, ) - ) - self._iteration_index += 1 + + start_time = time.monotonic() + yield iteration_details + iteration_details.elapsed_ms = (time.monotonic() - start_time) * 1000 + self._iteration_index = iteration_index + 1 + + def _make_iteration_details_stats( + self, iteration_details: SchedulerIterationDetails + ) -> SchedulerStats: + stats = self.scheduler.make_stats() or SchedulerStats() + stats.iteration_details = iteration_details + return stats + + def _attach_iteration_details( + self, + outputs: dict[int, EngineCoreOutputs], + iteration_details: SchedulerIterationDetails | None, + ) -> None: + if iteration_details is None: + return + + if (eco := next(iter(outputs.values()), None)) is None: + outputs[0] = eco = EngineCoreOutputs() + if eco.scheduler_stats is None: + eco.scheduler_stats = self._make_iteration_details_stats(iteration_details) + else: + eco.scheduler_stats.iteration_details = iteration_details def _should_throttle_prefills(self) -> bool: """Whether to defer new prefills this step (DP prefill balancing). @@ -558,8 +588,8 @@ def step(self) -> tuple[dict[int, EngineCoreOutputs], bool]: future = self.model_executor.execute_model(scheduler_output, non_block=True) grammar_output = self.scheduler.get_grammar_bitmask(scheduler_output) with ( + self.capture_iteration_details(scheduler_output) as iteration_details, self.log_error_detail(scheduler_output), - self.log_iteration_details(scheduler_output), ): model_output = future.result() if model_output is None: @@ -571,6 +601,7 @@ def step(self) -> tuple[dict[int, EngineCoreOutputs], bool]: engine_core_outputs = self.scheduler.update_from_output( scheduler_output, model_output ) + self._attach_iteration_details(engine_core_outputs, iteration_details) return engine_core_outputs, scheduler_output.total_num_scheduled_tokens > 0 @@ -656,8 +687,8 @@ def step_with_batch_queue( # Block until the next result is available. future, scheduler_output, exec_model_fut = batch_queue.pop() with ( + self.capture_iteration_details(scheduler_output) as iteration_details, self.log_error_detail(scheduler_output), - self.log_iteration_details(scheduler_output), ): model_output = future.result() if model_output is None: @@ -672,6 +703,7 @@ def step_with_batch_queue( engine_core_outputs = self.scheduler.update_from_output( scheduler_output, model_output ) + self._attach_iteration_details(engine_core_outputs, iteration_details) # NOTE(nick): We can either handle the deferred tasks here or save # in a field and do it immediately once step_with_batch_queue is @@ -2019,8 +2051,13 @@ def run_busy_loop(self): # Execute a dummy pass when no ready requests ran, unless the # engine is sleeping. elif not self.model_executor.is_sleeping: - with self.log_iteration_details(None): + with self.capture_iteration_details(None) as iteration_details: self.execute_dummy_batch() + if iteration_details is not None and not self.has_coordinator: + stats = self._make_iteration_details_stats(iteration_details) + self.output_queue.put_nowait( + (0, EngineCoreOutputs(scheduler_stats=stats)) + ) # 3) All-reduce operation to determine global unfinished reqs. self.engines_running = self._has_global_unfinished_reqs( diff --git a/vllm/v1/kv_offload/tiering/p2p/data/base.py b/vllm/v1/kv_offload/tiering/p2p/data/base.py index a18b1ca0b3c9..ffe609d71c1e 100644 --- a/vllm/v1/kv_offload/tiering/p2p/data/base.py +++ b/vllm/v1/kv_offload/tiering/p2p/data/base.py @@ -207,8 +207,17 @@ def write_blocks( ... @abstractmethod - def poll(self) -> PollResult: - """Poll all inflight transfers for completion. + def poll(self, peer_id: str | None = None) -> PollResult: + """Poll inflight transfers for completion. + + Args: + peer_id: If given, only poll (and drain) transfers submitted for + this peer_id — the value passed to ``write_blocks``. This is + required when a single transport is shared across multiple + peer sessions: ``poll()`` pops completed handles, so an + unscoped poll by one session would consume and discard the + completions of its siblings, starving them. ``None`` polls + every peer's transfers (used only for the shutdown drain). Returns: PollResult with lists of completed and failed transfer_ids. diff --git a/vllm/v1/kv_offload/tiering/p2p/data/nixl.py b/vllm/v1/kv_offload/tiering/p2p/data/nixl.py index 5f283c4815c6..8d53b923be03 100644 --- a/vllm/v1/kv_offload/tiering/p2p/data/nixl.py +++ b/vllm/v1/kv_offload/tiering/p2p/data/nixl.py @@ -8,7 +8,7 @@ import itertools from collections.abc import Iterable -from typing import Any +from typing import Any, NamedTuple from vllm.distributed.nixl_utils import NixlWrapper as _NixlAgent from vllm.distributed.nixl_utils import nixl_agent_config as _NixlAgentConfig @@ -27,6 +27,17 @@ _EMPTY_POLL_RESULT: PollResult = PollResult(done=(), failed=()) +class _Inflight(NamedTuple): + """A submitted-but-not-yet-drained transfer. + + ``peer_id`` lets poll() scope to a single owning session, since the + transport is shared across all peer sessions of the engine. + """ + + peer_id: str + handle: object + + class NixlTransport(DataTransport): """Manages a NIXL agent, memory registration, and block transfers. @@ -37,14 +48,14 @@ class NixlTransport(DataTransport): def __init__( self, - local_id: str, + agent_name: str, view: memoryview, config_fields: dict | None = None, backends: list[str] | None = None, num_threads: int = 4, ) -> None: super().__init__(view, config_fields=config_fields) - self._local_id = local_id + self._agent_name = agent_name self._backends = list(backends) if backends else ["UCX"] self._num_threads = num_threads self._agent: Any = None @@ -52,7 +63,8 @@ def __init__( self._local_dlist: Any = None self._remote_dlists: dict[str, object] = {} self._peer_nixl_names: dict[str, str] = {} - self._inflight: dict[int, object] = {} # transfer_id → handle + # transfer_id → _Inflight(peer_id, handle). + self._inflight: dict[int, _Inflight] = {} self._next_id = itertools.count() self._init(view) @@ -70,7 +82,7 @@ def _init(self, view: memoryview) -> None: cfg = _NixlAgentConfig(backends=self._backends, capture_telemetry=True) logger.info( "NixlTransport %s: NIXL backends=%s", - self._local_id, + self._agent_name, self._backends, ) else: @@ -79,10 +91,10 @@ def _init(self, view: memoryview) -> None: ) logger.info( "NixlTransport %s: NIXL backends=[UCX] num_threads=%d", - self._local_id, + self._agent_name, self._num_threads, ) - self._agent = _NixlAgent(self._local_id, cfg) + self._agent = _NixlAgent(self._agent_name, cfg) total_size = self._num_blocks * self._block_len reg_descs = [(self._base_addr, total_size, 0, "")] @@ -95,7 +107,7 @@ def _init(self, view: memoryview) -> None: xfer_dlist = self._agent.get_xfer_descs(block_tuples, mem_type="DRAM") self._local_dlist = self._agent.prep_xfer_dlist("NIXL_INIT_AGENT", xfer_dlist) logger.info( - "NixlTransport %s: registered %d blocks", self._local_id, self._num_blocks + "NixlTransport %s: registered %d blocks", self._agent_name, self._num_blocks ) def get_agent_metadata(self) -> bytes: @@ -152,14 +164,14 @@ def write_blocks( logger.warning( "NixlTransport %s: write_blocks NO REMOTE DLIST for peer=%s " "(known peers=%s)", - self._local_id, + self._agent_name, peer_id, list(self._remote_dlists.keys()), ) return None logger.debug( "NixlTransport %s: write_blocks NIXL.transfer peer=%s blocks=%d", - self._local_id, + self._agent_name, peer_id, len(local_idxs), ) @@ -172,11 +184,16 @@ def write_blocks( ) self._agent.transfer(handle) transfer_id = next(self._next_id) - self._inflight[transfer_id] = handle + self._inflight[transfer_id] = _Inflight(peer_id, handle) return transfer_id - def poll(self) -> PollResult: - """Poll all inflight transfers. + def poll(self, peer_id: str | None = None) -> PollResult: + """Poll inflight transfers. + + When *peer_id* is given, only transfers submitted for that peer_id are + checked and drained — the transport is shared across peer sessions, so + an unscoped poll by one session would consume and discard siblings' + completions. *peer_id* None polls every peer (shutdown drain only). Returns PollResult(done=..., failed=...) with transfer IDs. Completed handles are released automatically. @@ -187,13 +204,15 @@ def poll(self) -> PollResult: done_ids: list[int] | None = None failed_ids: list[int] | None = None - for transfer_id, handle in self._inflight.items(): + for transfer_id, entry in self._inflight.items(): + if peer_id is not None and entry.peer_id != peer_id: + continue try: - state = self._agent.check_xfer_state(handle) + state = self._agent.check_xfer_state(entry.handle) except Exception as exc: logger.warning( "NixlTransport %s: check_xfer_state failed for transfer_id=%d: %s", - self._local_id, + self._agent_name, transfer_id, exc, ) @@ -212,9 +231,9 @@ def poll(self) -> PollResult: handles_to_release = [] for tid in done_ids or (): - handles_to_release.append(self._inflight.pop(tid)) + handles_to_release.append(self._inflight.pop(tid).handle) for tid in failed_ids or (): - handles_to_release.append(self._inflight.pop(tid)) + handles_to_release.append(self._inflight.pop(tid).handle) self._release_handles(handles_to_release) return PollResult( @@ -236,22 +255,24 @@ def cancel( """ if mode == "immediate": handles = [ - self._inflight.pop(tid) for tid in transfer_ids if tid in self._inflight + self._inflight.pop(tid).handle + for tid in transfer_ids + if tid in self._inflight ] self._release_handles(handles) return [] still_inflight: list[int] = [] for tid in transfer_ids: - handle = self._inflight.get(tid) - if handle is None: + entry = self._inflight.get(tid) + if entry is None: continue try: - self._agent.release_xfer_handle(handle) + self._agent.release_xfer_handle(entry.handle) except Exception as exc: logger.debug( "NixlTransport %s: cancel pending for transfer_id=%d: %s", - self._local_id, + self._agent_name, tid, exc, ) @@ -267,7 +288,7 @@ def cancel( def close(self) -> None: if self._agent is None: return - self._release_handles(list(self._inflight.values())) + self._release_handles([entry.handle for entry in self._inflight.values()]) self._inflight.clear() for peer_id in list(self._remote_dlists): self.remove_remote_peer(peer_id) @@ -292,6 +313,6 @@ def _release_handles(self, handles: list[object]) -> None: except Exception as exc: logger.warning( "NixlTransport %s: release_xfer_handle failed: %s", - self._local_id, + self._agent_name, exc, ) diff --git a/vllm/v1/kv_offload/tiering/p2p/manager.py b/vllm/v1/kv_offload/tiering/p2p/manager.py index 6949686f4482..0cb37dff0e08 100644 --- a/vllm/v1/kv_offload/tiering/p2p/manager.py +++ b/vllm/v1/kv_offload/tiering/p2p/manager.py @@ -9,12 +9,14 @@ from __future__ import annotations import time +import uuid from collections.abc import Iterable, Sequence from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any from typing_extensions import override +import vllm.envs as envs from vllm.logger import init_logger from vllm.v1.kv_offload.base import ( LookupResult, @@ -111,8 +113,8 @@ def __init__( offloading_spec: OffloadingSpec, primary_kv_view: memoryview, tier_type: str = "p2p", - host: str = "0.0.0.0", - port: int = 7777, + host: str | None = None, + port: int | None = None, backends: list[str] | None = None, num_threads: int = 4, **kwargs: Any, @@ -130,9 +132,19 @@ def __init__( primary_kv_view: Memoryview over the CPU primary tier; the NIXL agent registers this region for RDMA transfers. tier_type: Tier identifier (defaults to ``"p2p"``). - host: Address the ZMQ control socket binds to. - port: Port for the ZMQ control socket. Must be reachable - from peers. + host: Address the ZMQ control socket binds to, used verbatim + as both the bind address and the identity peers dial back + (mirrors the NIXL connector's ``VLLM_NIXL_SIDE_CHANNEL_HOST``; + no auto-detection). Defaults to + ``VLLM_P2P_SIDE_CHANNEL_HOST`` (``localhost``) when not set; + must be set to the node's routable IP for cross-host P2P so + remote peers can reach the socket. + port: Base port for the ZMQ control socket. Must be + reachable from peers. Defaults to + ``VLLM_P2P_SIDE_CHANNEL_PORT`` (``5710``) when not set. + The bound port is ``base + data_parallel_index`` so each + DP replica gets a distinct port (one socket per replica, + like NIXL); for DP=1 the offset is 0. backends: NIXL transport backends (e.g. ``["UCX"]``, ``["MOONCAKE"]``, ``["LIBFABRIC"]``). Defaults to ``["UCX"]``. When any non-UCX backend is requested, the @@ -145,8 +157,26 @@ def __init__( **kwargs: Reserved for future tier-specific options. """ super().__init__(offloading_spec, primary_kv_view, tier_type) - port = int(port) + if host is None: + host = envs.VLLM_P2P_SIDE_CHANNEL_HOST + if port is None: + port = envs.VLLM_P2P_SIDE_CHANNEL_PORT + # One control socket per DP replica: offset the base by the global + # data-parallel index so replicas on a host don't collide (mirrors + # NIXL). For DP=1 the index is 0, leaving the base port unchanged. + dp_index = offloading_spec.vllm_config.parallel_config.data_parallel_index + port = int(port) + dp_index + # Two decoupled identities: + # _local_id (``host:port``): the ZMQ control identity that peers + # dial back, used verbatim (the socket binds this host/port and + # the address is parsed back into host:port by the remote). + # _nixl_agent_name (uuid4): the NIXL agent name. It is never dialed + # — it travels opaquely inside the agent metadata blob — so it + # only needs to be globally unique. A per-process uuid guarantees + # that even for peers sharing a host:port (mirrors the NIXL + # connector; avoids the "remote agent name equals local" reject). self._local_id = f"{host}:{port}" + self._nixl_agent_name = str(uuid.uuid4()) config_fields = FileMapper.from_offloading_spec( root_dir="", @@ -155,7 +185,7 @@ def __init__( parallel_agnostic=True, ).get_run_config() self._data: DataTransport = NixlTransport( - self._local_id, + self._nixl_agent_name, primary_kv_view, config_fields=config_fields, backends=backends, diff --git a/vllm/v1/kv_offload/tiering/p2p/session/client.py b/vllm/v1/kv_offload/tiering/p2p/session/client.py index fce5008deb18..f080cb79a36e 100644 --- a/vllm/v1/kv_offload/tiering/p2p/session/client.py +++ b/vllm/v1/kv_offload/tiering/p2p/session/client.py @@ -139,6 +139,14 @@ def on_abort_ack(self, kv_request_id: str) -> None: """Handle an AbortAckMsg from the peer.""" req = self._inbound.pop(kv_request_id, None) if req is not None: + logger.warning( + "P2PSession %s: load request %s (job_id=%d) timed out; " + "load job completed with failure. If this recurs, ensure " + "PYTHONHASHSEED is set to the same value on all nodes.", + self._peer_id, + kv_request_id, + req.job_id, + ) self._completed_loads.append( LoadResult( job_id=req.job_id, diff --git a/vllm/v1/kv_offload/tiering/p2p/session/server.py b/vllm/v1/kv_offload/tiering/p2p/session/server.py index 6f5059e5503b..2a6c8c692446 100644 --- a/vllm/v1/kv_offload/tiering/p2p/session/server.py +++ b/vllm/v1/kv_offload/tiering/p2p/session/server.py @@ -282,7 +282,11 @@ def collect_results(self) -> list[StoreResult]: results.extend(self._pending_store_results) self._pending_store_results.clear() - poll_result = self._transport.poll() + # Scope the poll to this peer: the transport is shared across all peer + # sessions of the engine, and poll() drains completed handles. An + # unscoped poll here would consume sibling sessions' completions and + # report them as "unknown transfer_id", starving those sessions. + poll_result = self._transport.poll(self._peer_id) for tid in poll_result.done: xfer = self._inflight_pop(tid) diff --git a/vllm/v1/metrics/loggers.py b/vllm/v1/metrics/loggers.py index 021019dc1cdc..692106cf3967 100644 --- a/vllm/v1/metrics/loggers.py +++ b/vllm/v1/metrics/loggers.py @@ -160,6 +160,42 @@ def _get_throughput(self, tracked_stats: int, now: float) -> float: def log_prefix(self): return "Engine {:03d}: ".format(self.engine_index) + def _log_prefix_for_engine(self, engine_idx: int) -> str: + if self.engine_index == engine_idx: + return self.log_prefix + return "Engine {:03d}: ".format(engine_idx) + + def _log_iteration_details( + self, scheduler_stats: SchedulerStats, engine_idx: int + ) -> None: + details = scheduler_stats.iteration_details + if details is None: + return + + encoder_msg = "" + if details.num_encoder_inputs: + encoder_msg = ( + f", encoder inputs: {details.num_encoder_inputs}, " + f"encoder output embeddings: {details.num_encoder_output_tokens}" + ) + + logger.info( + "%sIteration(%d): %d context requests, %d context tokens, " + "%d generation requests, %d generation tokens, " + "iteration elapsed time: %.2f ms%s, " + "GPU KV cache usage: %.1f%%%s", + self._log_prefix_for_engine(engine_idx), + details.iteration_index, + details.num_ctx_requests, + details.num_ctx_tokens, + details.num_generation_requests, + details.num_generation_tokens, + details.elapsed_ms, + " (dummy)" if details.is_dummy else "", + scheduler_stats.kv_cache_usage * 100, + encoder_msg, + ) + def record( self, scheduler_stats: SchedulerStats | None, @@ -172,6 +208,7 @@ def record( self._track_iteration_stats(iteration_stats) if scheduler_stats is not None: + self._log_iteration_details(scheduler_stats, engine_idx) self.prefix_caching_metrics.observe(scheduler_stats.prefix_cache_stats) if scheduler_stats.connector_prefix_cache_stats is not None: diff --git a/vllm/v1/metrics/stats.py b/vllm/v1/metrics/stats.py index a7a5fb7a2d2f..20bb3e1caa64 100644 --- a/vllm/v1/metrics/stats.py +++ b/vllm/v1/metrics/stats.py @@ -167,6 +167,21 @@ class KVCacheEvictionEvent: reuse_gaps_seconds: tuple[float, ...] +@dataclass +class SchedulerIterationDetails: + """Scheduler-side details for one engine iteration.""" + + iteration_index: int + num_ctx_requests: int + num_ctx_tokens: int + num_generation_requests: int + num_generation_tokens: int + elapsed_ms: float + num_encoder_inputs: int = 0 + num_encoder_output_tokens: int = 0 + is_dummy: bool = False + + @dataclass class SchedulerStats: """Stats associated with the scheduler.""" @@ -181,6 +196,7 @@ class SchedulerStats: current_wave: int = 0 kv_cache_usage: float = 0.0 + iteration_details: SchedulerIterationDetails | None = None prefix_cache_stats: PrefixCacheStats = field(default_factory=PrefixCacheStats) connector_prefix_cache_stats: PrefixCacheStats | None = None diff --git a/vllm/v1/spec_decode/dflash.py b/vllm/v1/spec_decode/dflash.py index bae6935cef8f..5d1c0629218d 100644 --- a/vllm/v1/spec_decode/dflash.py +++ b/vllm/v1/spec_decode/dflash.py @@ -70,7 +70,11 @@ def __init__( # For DFlash we use the input embeddings to embed the mask token self.parallel_drafting_hidden_state_tensor = None - self.dflash_causal = self.dflash_config.get("causal", False) + from vllm.model_executor.models.qwen3_dflash import dflash_has_any_non_causal + + self.dflash_causal = not dflash_has_any_non_causal( + self.draft_model_config.hf_config + ) @override def _create_draft_vllm_config(self) -> VllmConfig: diff --git a/vllm/v1/spec_decode/llm_base_proposer.py b/vllm/v1/spec_decode/llm_base_proposer.py index 756c5f3b3717..f8b52d079a89 100644 --- a/vllm/v1/spec_decode/llm_base_proposer.py +++ b/vllm/v1/spec_decode/llm_base_proposer.py @@ -1300,6 +1300,15 @@ def _create_draft_vllm_config(self) -> VllmConfig: ), ) + if spec_cfg.kv_cache_dtype is not None: + base = replace( + base, + cache_config=replace( + base.cache_config, + cache_dtype=spec_cfg.kv_cache_dtype, + ), + ) + return base def _get_model(self) -> nn.Module: diff --git a/vllm/v1/utils.py b/vllm/v1/utils.py index b485d838f3e8..e17ffed93403 100644 --- a/vllm/v1/utils.py +++ b/vllm/v1/utils.py @@ -782,12 +782,16 @@ class IterationDetails: num_ctx_tokens: int num_generation_requests: int num_generation_tokens: int + num_encoder_inputs: int = 0 + num_encoder_output_tokens: int = 0 def __repr__(self) -> str: return f"IterationDetails(num_ctx_requests={self.num_ctx_requests},\ num_ctx_tokens={self.num_ctx_tokens}, \ num_generation_requests={self.num_generation_requests}, \ - num_generation_tokens={self.num_generation_tokens})" + num_generation_tokens={self.num_generation_tokens}, \ + num_encoder_inputs={self.num_encoder_inputs}, \ + num_encoder_output_tokens={self.num_encoder_output_tokens})" def compute_iteration_details(scheduler_output: SchedulerOutput) -> IterationDetails: @@ -818,9 +822,18 @@ def compute_iteration_details(scheduler_output: SchedulerOutput) -> IterationDet else: num_generation_requests += 1 num_generation_tokens += num_tokens + scheduled_encoder_input_stats = scheduler_output.scheduled_encoder_input_stats + num_encoder_inputs = 0 + num_encoder_output_tokens = 0 + if scheduled_encoder_input_stats is not None: + num_encoder_inputs = scheduled_encoder_input_stats.num_inputs + num_encoder_output_tokens = scheduled_encoder_input_stats.output_tokens + return IterationDetails( num_context_requests, num_context_tokens, num_generation_requests, num_generation_tokens, + num_encoder_inputs, + num_encoder_output_tokens, ) diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 1dee0d1b9553..2bb52e2fd89f 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -472,9 +472,6 @@ def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None: decode_query_len=self.decode_query_len, lora_capture_cases=self.lora_capture_cases, ) - if self.speculator is not None: - self.speculator.init_cudagraph_manager(cudagraph_mode) - check_attention_cp_compatibility(self.vllm_config) if isinstance(self.speculator, DraftModelSpeculator): # HACK(woosuk) @@ -485,6 +482,10 @@ def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None: self.input_buffers, self.attn_groups, ) + if self.speculator is not None: + # After set_attn, so the speculator can size its cudagraph mode + # to its own attention support. + self.speculator.init_cudagraph_manager(cudagraph_mode) self.kv_caches: list[torch.Tensor] = [] kv_caches_dict = init_kv_cache( diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py index 6fd99d55fc73..3cc88c49b2fa 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py @@ -6,11 +6,12 @@ import torch import torch.nn as nn -from vllm.config import VllmConfig +from vllm.config import VllmConfig, replace from vllm.config.compilation import CUDAGraphMode from vllm.forward_context import BatchDescriptor, set_forward_context from vllm.logger import init_logger from vllm.triton_utils import tl, triton +from vllm.v1.attention.backend import AttentionCGSupport from vllm.v1.attention.backends.utils import PAD_SLOT_ID from vllm.v1.kv_cache_interface import KVCacheConfig from vllm.v1.worker.gpu.attn_utils import build_slot_mappings_by_layer @@ -19,10 +20,7 @@ from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers from vllm.v1.worker.gpu.model_states.interface import ModelState from vllm.v1.worker.gpu.spec_decode.dflash.cudagraph import DFlashCudaGraphManager -from vllm.v1.worker.gpu.spec_decode.dflash.utils import ( - get_dflash_causal, - load_dflash_model, -) +from vllm.v1.worker.gpu.spec_decode.dflash.utils import load_dflash_model from vllm.v1.worker.gpu.spec_decode.speculator import DraftModelSpeculator from vllm.v1.worker.gpu.spec_decode.utils import get_parallel_drafting_token_id from vllm.v1.worker.utils import AttentionGroup @@ -50,7 +48,11 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): self.draft_model_config.hf_config ) - self.dflash_causal = get_dflash_causal(self.draft_model_config) + from vllm.model_executor.models.qwen3_dflash import dflash_has_any_non_causal + + self.requires_non_causal = dflash_has_any_non_causal( + self.draft_model_config.hf_config + ) # Whether the anchor query position is itself a prediction. DFlash default uses # the anchor as the bonus token (only mask tokens predict); DSpark samples from @@ -84,9 +86,32 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): self.query_cudagraph_manager: DFlashCudaGraphManager | None = None self.draft_kv_cache_group_id: int = -1 + @property + def attn_vllm_config(self) -> VllmConfig: + # The draft's attention differs from the target's in causality. + return replace( + self.vllm_config, + attention_config=replace( + self.vllm_config.attention_config, + use_non_causal=self.requires_non_causal, + ), + ) + def init_cudagraph_manager(self, cudagraph_mode: CUDAGraphMode) -> None: - # PIECEWISE cudagraphs are not supported for dflash - if cudagraph_mode.decode_mode() == CUDAGraphMode.FULL: + wants_full = cudagraph_mode.decode_mode() == CUDAGraphMode.FULL + supports_full = ( + self.attn_cg_support.min_cg_support.value + >= AttentionCGSupport.UNIFORM_BATCH.value + ) + if wants_full and not supports_full: + logger.warning( + "%s draft attention (%s) does not support full CUDA graphs; " + "running the draft eagerly.", + self._speculator_name, + self.attn_cg_support.min_cg_attn_backend, + ) + # PIECEWISE cudagraphs are not supported for dflash. + if wants_full and supports_full: cudagraph_mode = CUDAGraphMode.FULL_DECODE_ONLY else: cudagraph_mode = CUDAGraphMode.NONE @@ -158,8 +183,8 @@ def set_attn( # of the kv-cache group its cache belongs to. Models that share a single group # leave this as None and share one context slot mapping. self._layer_group_idx: list[int] | None = None - # Per-KV-group causal, falling back to the scalar dflash_causal. - self._group_causal: dict[int, bool] | bool = self.dflash_causal + # Per-KV-group causal, falling back to whether the drafter is all-causal. + self._group_causal: dict[int, bool] | bool = not self.requires_non_causal if hasattr(self.model, "get_draft_kv_cache_layer_names"): layer_names = self.model.get_draft_kv_cache_layer_names() name_to_gid = { diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/utils.py b/vllm/v1/worker/gpu/spec_decode/dflash/utils.py index 37fe693bbbdd..478274a05671 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/utils.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/utils.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import torch.nn as nn -from vllm.config import ModelConfig, VllmConfig, replace +from vllm.config import VllmConfig, replace from vllm.distributed.parallel_state import get_pp_group from vllm.model_executor.model_loader import get_model from vllm.v1.worker.gpu.spec_decode.eagle.utils import ( @@ -11,28 +11,30 @@ ) -def get_dflash_causal(draft_model_config: ModelConfig) -> bool: - """Whether the DFlash draft uses causal (vs non-causal) attention.""" - dflash_config = getattr(draft_model_config.hf_config, "dflash_config", None) or {} - return dflash_config.get("causal", False) - - def load_dflash_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Module: from vllm.compilation.backends import set_model_tag + from vllm.model_executor.models.qwen3_dflash import dflash_has_any_non_causal speculative_config = vllm_config.speculative_config assert speculative_config is not None draft_model_config = speculative_config.draft_model_config - # Modify the attention config so that we select an attention backend that matches - # the causal/non-causal mode of the dflash model. - causal = get_dflash_causal(draft_model_config) + # Select an attention backend that supports the drafter's attention: mixing + # a non-causal layer onto a causal-only backend would fail. draft_vllm_config = replace( vllm_config, attention_config=replace( vllm_config.attention_config, - use_non_causal=not causal, + use_non_causal=dflash_has_any_non_causal(draft_model_config.hf_config), backend=speculative_config.attention_backend, ), + cache_config=( + replace( + vllm_config.cache_config, + cache_dtype=speculative_config.kv_cache_dtype, + ) + if speculative_config.kv_cache_dtype is not None + else vllm_config.cache_config + ), ) with set_model_tag("dflash_head"): dflash_model = get_model( diff --git a/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py b/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py index 0236017cf2fc..9000ae878a43 100644 --- a/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py @@ -59,8 +59,6 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): self.max_num_tokens, draft_hidden, dtype=self.dtype, device=device ) - self.dflash_causal = False - self._step_cols = torch.arange( self.num_speculative_steps, dtype=torch.int32, device=device ) diff --git a/vllm/v1/worker/gpu/spec_decode/dspark/utils.py b/vllm/v1/worker/gpu/spec_decode/dspark/utils.py index 08ff30e5fdfd..9ea2ff0f7363 100644 --- a/vllm/v1/worker/gpu/spec_decode/dspark/utils.py +++ b/vllm/v1/worker/gpu/spec_decode/dspark/utils.py @@ -18,16 +18,23 @@ def load_dspark_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mo draft_model_config = speculative_config.draft_model_config from vllm.compilation.backends import set_model_tag + from vllm.model_executor.models.qwen3_dflash import dflash_has_any_non_causal - # DSpark uses non-causal attention. - causal = False draft_vllm_config = replace( vllm_config, attention_config=replace( vllm_config.attention_config, - use_non_causal=not causal, + use_non_causal=dflash_has_any_non_causal(draft_model_config.hf_config), backend=speculative_config.attention_backend, ), + cache_config=( + replace( + vllm_config.cache_config, + cache_dtype=speculative_config.kv_cache_dtype, + ) + if speculative_config.kv_cache_dtype is not None + else vllm_config.cache_config + ), ) with set_model_tag("dspark_head"): diff --git a/vllm/v1/worker/gpu/spec_decode/eagle/utils.py b/vllm/v1/worker/gpu/spec_decode/eagle/utils.py index bdd588e5786d..579652bc7d64 100644 --- a/vllm/v1/worker/gpu/spec_decode/eagle/utils.py +++ b/vllm/v1/worker/gpu/spec_decode/eagle/utils.py @@ -3,7 +3,7 @@ import torch import torch.nn as nn -from vllm.config import VllmConfig +from vllm.config import VllmConfig, replace from vllm.distributed.parallel_state import get_pp_group from vllm.lora.layers.base import BaseLayerWithLoRA from vllm.model_executor.model_loader import get_model @@ -39,6 +39,14 @@ def load_eagle_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mod speculative_config = vllm_config.speculative_config assert speculative_config is not None draft_model_config = speculative_config.draft_model_config + if speculative_config.kv_cache_dtype is not None: + vllm_config = replace( + vllm_config, + cache_config=replace( + vllm_config.cache_config, + cache_dtype=speculative_config.kv_cache_dtype, + ), + ) with set_model_tag("eagle_head"): eagle_model = get_model( vllm_config=vllm_config, model_config=draft_model_config diff --git a/vllm/v1/worker/gpu/spec_decode/speculator.py b/vllm/v1/worker/gpu/spec_decode/speculator.py index 36be47322115..f460f58fe090 100644 --- a/vllm/v1/worker/gpu/spec_decode/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/speculator.py @@ -175,6 +175,12 @@ def _prepare_eplb_forward(self, num_unpadded_tokens: int) -> None: num_unpadded_tokens, ) + @property + def attn_vllm_config(self) -> VllmConfig: + """Config for the draft's attention metadata builders. Overridden by + speculators whose attention mode differs from the target's.""" + return self.vllm_config + def set_attn( self, model_state: ModelState, @@ -185,9 +191,9 @@ def set_attn( ) -> None: self.model_state = model_state self.kv_cache_config = kv_cache_config - self.attn_groups, _, _ = init_attn_backend( + self.attn_groups, self.attn_cg_support, _ = init_attn_backend( kv_cache_config, - self.vllm_config, + self.attn_vllm_config, self.device, active_layer_names=self.draft_attn_layer_names, ) diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 377fb670e395..ea8980a6e558 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -42,6 +42,7 @@ from vllm.distributed.kv_transfer import get_kv_transfer_group, has_kv_transfer_group from vllm.distributed.kv_transfer.kv_connector.utils import copy_kv_blocks from vllm.distributed.parallel_state import ( + GraphCaptureContext, get_dcp_group, get_pp_group, get_tp_group, @@ -119,6 +120,7 @@ from vllm.utils.torch_utils import ( PIN_MEMORY, async_tensor_h2d, + current_stream, get_dtype_size, is_quantized_kv_cache, kv_cache_dtype_str_to_dtype, @@ -6607,11 +6609,31 @@ def profile_cudagraph_memory(self) -> int: per_graph_estimate = {} encoder_memory_estimate = 0 + # On ROCm, capture these throwaway profiling graphs on vLLM's dedicated + # compute stream instead of the fresh side stream graph_capture() + # allocates by default. torch's allocator pools free blocks per stream, + # so a side-stream forward strands a persistent aiter scratch buffer in + # a separate pool, shifting the physical placement of the real KV cache + # allocated afterward and slowing bandwidth-bound decode ~20%. The + # graphs are discarded, so a side stream is unnecessary here. + # Use current_stream(), not torch.cuda.current_stream(): before vLLM + # initializes its dedicated stream, torch returns the per-thread default + # stream (cuda_stream=0), which cannot be used for cudagraph capture. + # cap_ctx=None keeps the side-stream path on CUDA. + cap_ctx = ( + GraphCaptureContext(current_stream()) + if current_platform.is_rocm() + else None + ) + # Cleanup-only guard: CUDA graph capture errors should still propagate # because encoder graph capture is opt-in. try: set_cudagraph_capturing_enabled(True) - with self._freeze_gc(), graph_capture(device=self.device): + with ( + self._freeze_gc(), + graph_capture(device=self.device, graph_capture_context=cap_ctx), + ): torch.accelerator.synchronize() torch.accelerator.empty_cache() @@ -7042,12 +7064,14 @@ def _check_and_update_cudagraph_mode( # Initialize drafter's cudagraph dispatcher if using spec decode. if self.speculative_config and ( self.speculative_config.use_eagle() + or self.speculative_config.uses_draft_model() or self.speculative_config.uses_extract_hidden_states() ): assert isinstance( self.drafter, EagleProposer | DFlashProposer + | DraftModelProposer | ExtractHiddenStatesProposer | Gemma4Proposer, ) diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index 871f2f31c006..5fb0c387737a 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -492,11 +492,14 @@ def determine_available_memory(self) -> int: ) # Profile CUDA graph memory if graphs will be captured. - # Skip on ROCm/HIP/XPU as graph pool handles and get_memory_info - # behave differently and can produce incorrect/negative estimates. + # ROCm is included: #44825 moved the profiler to + # torch.accelerator.get_memory_info (reliable on ROCm, as used by + # the AMD-CI mem tests), and graph_pool_handle resolves to the same + # torch.cuda handle the live capture path already uses on ROCm. + # XPU stays excluded (see #39977). cudagraph_memory_estimate = 0 if ( - current_platform.is_cuda() + current_platform.is_cuda_alike() and self.vllm_config.compilation_config.cudagraph_mode != CUDAGraphMode.NONE ): @@ -512,8 +515,7 @@ def determine_available_memory(self) -> int: + profile_result.weights_memory ) - # On ROCm, cudagraph_memory_estimate is always 0 so this is a no-op. - # On CUDA, respect the opt-in flag as originally designed. + # Respect the opt-in flag as originally designed. cudagraph_memory_estimate_applied = ( cudagraph_memory_estimate if envs.VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS