diff --git a/CMakeLists.txt b/CMakeLists.txt index 1ef9d596aec5..b23a307d7d0c 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -49,7 +49,7 @@ install(CODE "set(CMAKE_INSTALL_LOCAL_ONLY TRUE)" ALL_COMPONENTS) set(PYTHON_SUPPORTED_VERSIONS "3.10" "3.11" "3.12" "3.13" "3.14") # Supported AMD GPU architectures. -set(HIP_SUPPORTED_ARCHS "gfx906;gfx908;gfx90a;gfx942;gfx950;gfx1030;gfx1100;gfx1101;gfx1102;gfx1103;gfx1150;gfx1151;gfx1152;gfx1153;gfx1200;gfx1201") +set(HIP_SUPPORTED_ARCHS "gfx906;gfx908;gfx90a;gfx942;gfx950;gfx1250;gfx1030;gfx1100;gfx1101;gfx1102;gfx1103;gfx1150;gfx1151;gfx1152;gfx1153;gfx1200;gfx1201") # ROCm installation prefix. Default to /opt/rocm but allow override via # -DROCM_PATH=/your/rocm/path when invoking cmake. @@ -1370,6 +1370,34 @@ if(VLLM_GPU_LANG STREQUAL "HIP") "csrc/rocm/torch_bindings.cpp" "csrc/rocm/skinny_gemms.cu" "csrc/rocm/attention.cu") + set(VLLM_ROCM_EXT_FLAGS ${VLLM_GPU_FLAGS}) + + # skinny_gemms.cu is built on gfx9/gfx11 ISA (MFMA, dot2/dot4, legacy + # s_waitcnt asm) that gfx1250 (gfx12) does not provide. Exclude it from the + # gfx1250 build and disable its op registrations (VLLM_SKIP_SKINNY_GEMMS); + # vLLM falls back to default/Triton GEMM for those ops on gfx1250. + list(REMOVE_ITEM VLLM_ROCM_EXT_SRC "csrc/rocm/skinny_gemms.cu") + set(VLLM_SKINNY_ARCHES ${VLLM_GPU_ARCHES}) + list(FILTER VLLM_SKINNY_ARCHES EXCLUDE REGEX "gfx1250") + if(VLLM_SKINNY_ARCHES) + message(STATUS "Building skinny_gemms for archs: ${VLLM_SKINNY_ARCHES}") + hipify_sources_target(VLLM_SKINNY_HIP_SRCS _rocm_C_skinny "csrc/rocm/skinny_gemms.cu") + unset(_VLLM_LAST_HIPIFY_TARGET) + add_library(_rocm_C_skinny OBJECT ${VLLM_SKINNY_HIP_SRCS}) + add_dependencies(_rocm_C_skinny hipify_all) + set_source_files_properties(${VLLM_SKINNY_HIP_SRCS} PROPERTIES LANGUAGE ${VLLM_GPU_LANG}) + set_target_properties(_rocm_C_skinny PROPERTIES + ${VLLM_GPU_LANG}_ARCHITECTURES "${VLLM_SKINNY_ARCHES}" + POSITION_INDEPENDENT_CODE ON) + target_include_directories(_rocm_C_skinny PRIVATE ${CMAKE_CURRENT_BINARY_DIR}/csrc) + target_compile_options(_rocm_C_skinny PRIVATE + $<$:${VLLM_ROCM_EXT_FLAGS}>) + target_compile_definitions(_rocm_C_skinny PRIVATE "-DTORCH_EXTENSION_NAME=_rocm_C") + target_link_libraries(_rocm_C_skinny PRIVATE torch) + else() + message(STATUS "Only gfx1250, skipping skinny_gemms") + list(APPEND VLLM_ROCM_EXT_FLAGS "-DVLLM_SKIP_SKINNY_GEMMS") + endif() set(VLLM_ROCM_HAS_GFX1100 OFF) if(VLLM_GPU_ARCHES MATCHES "gfx1100") @@ -1385,10 +1413,16 @@ if(VLLM_GPU_LANG STREQUAL "HIP") DESTINATION vllm LANGUAGE ${VLLM_GPU_LANG} SOURCES ${VLLM_ROCM_EXT_SRC} - COMPILE_FLAGS ${VLLM_GPU_FLAGS} + COMPILE_FLAGS ${VLLM_ROCM_EXT_FLAGS} ARCHITECTURES ${VLLM_GPU_ARCHES} USE_SABI 3 WITH_SOABI) + + if(TARGET _rocm_C_skinny) + target_link_libraries(_rocm_C PRIVATE _rocm_C_skinny) + else() + target_compile_definitions(_rocm_C PRIVATE VLLM_SKIP_SKINNY_GEMMS) + endif() if(VLLM_ROCM_HAS_GFX1100) target_compile_definitions(_rocm_C PRIVATE VLLM_ROCM_GFX1100) diff --git a/benchmarks/vllm_smoketest.sh b/benchmarks/vllm_smoketest.sh new file mode 100644 index 000000000000..facd4f0f98b2 --- /dev/null +++ b/benchmarks/vllm_smoketest.sh @@ -0,0 +1,110 @@ +#!/bin/bash +set -euo pipefail + +# add model args here +declare -A MODEL_PATHS=( + [gpt-oss-120b-mxfp4]="/data/amd/gpt-oss-120b-w-mxfp4-a-fp8" + [DeepSeek-R1-0528-MXFP4]="/data/amd/DeepSeek-R1-0528-MXFP4" +) +declare -A MODEL_SERVE_ARGS=( + [gpt-oss-120b-mxfp4]="--tensor-parallel-size 1 --gpu_memory_utilization 0.7 --attention-backend TRITON_ATTN" + [DeepSeek-R1-0528-MXFP4]="--tensor-parallel-size 1 --gpu_memory_utilization 0.9 --dtype auto --no-enable-prefix-caching --disable-uvicorn-access-log --trust-remote-code" +) +declare -A MODEL_ENV=( + [gpt-oss-120b-mxfp4]="HSA_OVERRIDE_GFX_VERSION=12.5.0 HSA_ENABLE_SDMA=0 VLLM_ROCM_AITER_FUSED_MOE_TRITON_GEMM_A4W4=1 VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_SKINNY_GEMM=0 VLLM_ROCM_USE_AITER_RMSNORM=0 " + [DeepSeek-R1-0528-MXFP4]="HSA_ENABLE_SDMA=0 USE_SVM=0 HSA_XNACK=0 VLLM_ROCM_AITER_FUSED_MOE_TRITON_GEMM_A4W4=1 VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION=1 VLLM_ROCM_USE_SKINNY_GEMM=0 VLLM_ROCM_USE_AITER_RMSNORM=0" +) + +MODEL_NAME="gpt-oss-120b-mxfp4" # default +MODEL_PATH_OVERRIDE="" +MODEL_ENV_OVERRIDE="" +MODEL_SET=0 # whether --model was explicitly passed +RUN_LM_EVAL=0 +PORT=8000 + +usage() { + cat <<'EOF' +Usage: ./vllm_smoketest.sh [--model NAME] [--model-path PATH] [--env ENV_VARS] [--lm-eval] [--port PORT] [--list] + --model NAME Registry key to run (default: gpt-oss-120b-mxfp4) + --model-path PATH Override model path for this run + --env ENV_VARS Override environment variables (e.g., "VAR1=val1 VAR2=val2") + --lm-eval Run lm_eval after curl check + --port PORT Server port (default: 8000) + --list List available models +EOF +} + +while [[ $# -gt 0 ]]; do + case "$1" in + --model) MODEL_NAME="$2"; MODEL_SET=1; shift 2 ;; + --model-path) MODEL_PATH_OVERRIDE="$2"; shift 2 ;; + --env) MODEL_ENV_OVERRIDE="$2"; shift 2 ;; + --lm-eval) RUN_LM_EVAL=1; shift ;; + --port) PORT="$2"; shift 2 ;; + --list) for n in "${!MODEL_PATHS[@]}"; do echo "$n -> ${MODEL_PATHS[$n]}"; done; exit 0 ;; + -h|--help) usage; exit 0 ;; + *) echo "Unknown argument: $1" >&2; usage; exit 1 ;; + esac +done + +if [[ -n "$MODEL_PATH_OVERRIDE" && $MODEL_SET -eq 0 ]]; then + echo "ERROR: --model-path requires --model (which model's serve args/env to use)." >&2 + usage; exit 1 +fi + +MODEL="${MODEL_PATHS[$MODEL_NAME]:-$MODEL_NAME}" +[[ -n "$MODEL_PATH_OVERRIDE" ]] && MODEL="$MODEL_PATH_OVERRIDE" +SERVE_ARGS="${MODEL_SERVE_ARGS[$MODEL_NAME]:-}" +MODEL_ENV_ARGS="${MODEL_ENV[$MODEL_NAME]:-}" +[[ -n "$MODEL_ENV_OVERRIDE" ]] && MODEL_ENV_ARGS="$MODEL_ENV_OVERRIDE" +echo "Model: $MODEL | Port: $PORT | lm_eval: $RUN_LM_EVAL" + +LOG_DIR="$(pwd)/vllm_smoke_test_logs" +mkdir -p "$LOG_DIR" +echo "Logs: $LOG_DIR" + +server_pid="" +cleanup() { [[ -n "$server_pid" ]] && kill "$server_pid" 2>/dev/null && wait "$server_pid" 2>/dev/null; true; } +trap cleanup EXIT +trap 'exit 130' INT TERM + +# --- start server --- +env $MODEL_ENV_ARGS \ +vllm serve --model "$MODEL" --host localhost --port "$PORT" $SERVE_ARGS & +server_pid=$! +echo "Waiting for server (pid $server_pid)..." + +until curl -s "http://localhost:$PORT/health" &>/dev/null; do + ps -p "$server_pid" >/dev/null || { echo "ERROR: server died"; tail -n 50 "$LOG_DIR/server.log"; exit 1; } + sleep 5 +done +echo "Server ready." + +# --- check if text was generated --- +echo "=== Curl completion check ===" +http_code=$(curl -s -o "$LOG_DIR/curl_completion.log" -w '%{http_code}' \ + "http://localhost:$PORT/v1/completions" \ + -H "Content-Type: application/json" \ + -d "{\"model\":\"$MODEL\",\"prompt\":\"The capital of France is \",\"max_tokens\":32,\"temperature\":0}") +cat "$LOG_DIR/curl_completion.log"; echo +[[ "$http_code" == "200" ]] || { echo "FAIL: HTTP $http_code" >&2; exit 1; } +grep -q '"text"' "$LOG_DIR/curl_completion.log" || { echo "FAIL: no text in response" >&2; exit 1; } +grep -qi "Paris" "$LOG_DIR/curl_completion.log" || { echo "FAIL: 'Paris' not found in response" >&2; exit 1; } +CURL_CHECK_PASSED=1 + +# --- lm_eval --- +if [[ $RUN_LM_EVAL -eq 1 ]]; then + echo "=== lm_eval ===" + pip install "lm-eval[api]" + lm_eval --model local-completions \ + --model_args "model=$MODEL,base_url=http://localhost:$PORT/v1/completions,num_concurrent=64,max_retries=3,tokenized_requests=False" \ + --tasks gsm8k --num_fewshot 3 --limit 100 2>&1 | tee "$LOG_DIR/lm_eval.log" +fi + +# Shutdown server +cleanup + +# Print PASS only if curl check succeeded +[[ ${CURL_CHECK_PASSED:-0} -eq 1 ]] && echo "PASS" + +echo "Done." diff --git a/cmake/external_projects/triton_kernels.cmake b/cmake/external_projects/triton_kernels.cmake index 2966c78030bd..dba2e36608d3 100644 --- a/cmake/external_projects/triton_kernels.cmake +++ b/cmake/external_projects/triton_kernels.cmake @@ -1,6 +1,6 @@ # Install OpenAI triton_kernels from https://github.com/triton-lang/triton/tree/main/python/triton_kernels -set(DEFAULT_TRITON_KERNELS_TAG "v3.5.1") +set(DEFAULT_TRITON_KERNELS_TAG "padroute") # Set TRITON_KERNELS_SRC_DIR for use with local development with vLLM. We expect TRITON_KERNELS_SRC_DIR to # be directly set to the triton_kernels python directory. @@ -11,13 +11,14 @@ if (DEFINED ENV{TRITON_KERNELS_SRC_DIR}) SOURCE_DIR $ENV{TRITON_KERNELS_SRC_DIR} ) +#TODO (JPVILLAM): Test new release of triton kernels and put back to normal else() - set(TRITON_GIT "https://github.com/triton-lang/triton.git") + set(TRITON_GIT "https://github.com/jpvillam-amd/triton.git") message (STATUS "[triton_kernels] Fetch from ${TRITON_GIT}:${DEFAULT_TRITON_KERNELS_TAG}") FetchContent_Declare( triton_kernels # TODO (varun) : Fetch just the triton_kernels directory from Triton - GIT_REPOSITORY https://github.com/triton-lang/triton.git + GIT_REPOSITORY https://github.com/jpvillam-amd/triton.git GIT_TAG ${DEFAULT_TRITON_KERNELS_TAG} GIT_PROGRESS TRUE SOURCE_SUBDIR python/triton_kernels/triton_kernels diff --git a/csrc/libtorch_stable/minimax_reduce_rms_kernel.cu b/csrc/libtorch_stable/minimax_reduce_rms_kernel.cu index d9af0f5efe0f..da7fee5670f5 100644 --- a/csrc/libtorch_stable/minimax_reduce_rms_kernel.cu +++ b/csrc/libtorch_stable/minimax_reduce_rms_kernel.cu @@ -804,35 +804,6 @@ void minimax_reduce_rms_op(MiniMaxReduceRMSParams const& params) { } // namespace tensorrt_llm } // namespace vllm -torch::stable::Tensor minimax_allreduce_rms( - torch::stable::Tensor const& input, - torch::stable::Tensor const& norm_weight, torch::stable::Tensor workspace, - int64_t const rank, int64_t const nranks, double const eps) { - const torch::stable::accelerator::DeviceGuard device_guard( - input.get_device_index()); - auto allreduce_params = vllm::tensorrt_llm::MiniMaxReduceRMSParams(); - - allreduce_params.nranks = static_cast(nranks); - allreduce_params.rank = static_cast(rank); - allreduce_params.dtype = input.scalar_type(); - allreduce_params.size_q = static_cast(input.numel()); - allreduce_params.hidden_dim = static_cast(input.size(-1)); - allreduce_params.stride_q = allreduce_params.hidden_dim; - allreduce_params.workspace = - reinterpret_cast(workspace.mutable_data_ptr()); - allreduce_params.allreduce_in = const_cast(input.const_data_ptr()); - allreduce_params.rms_gamma = const_cast(norm_weight.const_data_ptr()); - allreduce_params.rms_eps = static_cast(eps); - allreduce_params.stream = get_current_cuda_stream(input.get_device_index()); - - torch::stable::Tensor rms_norm_out = torch::stable::empty_like(input); - allreduce_params.rms_norm_out = rms_norm_out.mutable_data_ptr(); - - vllm::tensorrt_llm::minimax_reduce_rms_op(allreduce_params); - - return rms_norm_out; -} - std::tuple minimax_allreduce_rms_qk(torch::stable::Tensor qkv, torch::stable::Tensor const& norm_weight_q, diff --git a/csrc/libtorch_stable/ops.h b/csrc/libtorch_stable/ops.h index d60b68a5868d..7cf34d8b03ab 100644 --- a/csrc/libtorch_stable/ops.h +++ b/csrc/libtorch_stable/ops.h @@ -288,10 +288,6 @@ void fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert( int64_t cache_block_size); #ifndef USE_ROCM -torch::stable::Tensor minimax_allreduce_rms( - torch::stable::Tensor const& input, - torch::stable::Tensor const& norm_weight, torch::stable::Tensor workspace, - int64_t const rank, int64_t const nranks, double const eps); std::tuple minimax_allreduce_rms_qk(torch::stable::Tensor qkv, torch::stable::Tensor const& norm_weight_q, diff --git a/csrc/libtorch_stable/torch_bindings.cpp b/csrc/libtorch_stable/torch_bindings.cpp index 1be7217ce789..158999a6633b 100644 --- a/csrc/libtorch_stable/torch_bindings.cpp +++ b/csrc/libtorch_stable/torch_bindings.cpp @@ -449,10 +449,6 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { "int cache_block_size) -> ()"); #ifndef USE_ROCM - ops.def( - "minimax_allreduce_rms(" - "Tensor input, Tensor norm_weight, Tensor workspace, " - "int rank, int nranks, float eps) -> Tensor"); ops.def( "minimax_allreduce_rms_qk(" "Tensor qkv, Tensor norm_weight_q, Tensor norm_weight_k, " @@ -705,7 +701,6 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) { "fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert", TORCH_BOX(&fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert)); #ifndef USE_ROCM - ops.impl("minimax_allreduce_rms", TORCH_BOX(&minimax_allreduce_rms)); ops.impl("minimax_allreduce_rms_qk", TORCH_BOX(&minimax_allreduce_rms_qk)); #endif ops.impl("fused_minimax_m3_qknorm_rope_kv_insert", diff --git a/csrc/quickreduce/base.h b/csrc/quickreduce/base.h index 6c3456d06f20..e9af6d7e822c 100644 --- a/csrc/quickreduce/base.h +++ b/csrc/quickreduce/base.h @@ -79,6 +79,7 @@ union BufferResource { }; }; +#if !defined(__gfx1250__) __quickreduce_device_inline__ static int32x4_t buffer_load_dwordx4( int32x4_t srsrc, int32_t voffset, int32_t soffset, int32_t aux) __asm("llvm.amdgcn.raw.buffer.load.v4i32"); @@ -86,6 +87,15 @@ __quickreduce_device_inline__ static int32x4_t buffer_load_dwordx4( __quickreduce_device_inline__ static void buffer_store_dwordx4( int32x4_t data, int32x4_t srsrc, int32_t voffset, int32_t soffset, int32_t aux) __asm("llvm.amdgcn.raw.buffer.store.v4i32"); +#else +__quickreduce_device_inline__ static int32x4_t buffer_load_dwordx4( + int32x4_t srsrc, int32_t voffset, int32_t soffset, + int32_t aux) {} + +__quickreduce_device_inline__ static void buffer_store_dwordx4( + int32x4_t data, int32x4_t srsrc, int32_t voffset, int32_t soffset, + int32_t aux) {} +#endif __quickreduce_device_inline__ static void set_fp16_ovfl(bool const value) { #if defined(__gfx942__) diff --git a/csrc/rocm/attention.cu b/csrc/rocm/attention.cu index 4ac255d0a75f..86c5044ce981 100644 --- a/csrc/rocm/attention.cu +++ b/csrc/rocm/attention.cu @@ -2405,6 +2405,15 @@ template __device__ __forceinline__ floatx8 gcn_wmma16x16x16_instr(const bit16x8& inpA, const bit16x8& inpB, const floatx8& inpC) { +#if defined(__gfx1250__) + // gfx1250 (gfx12 family) does not provide the gfx12 WMMA variant used by + // gfx1200/1201 (needs wmma-128b-insts). This custom-attention WMMA path is + // unsupported on gfx1250; trap if ever launched (fail loud, not silent-wrong). + (void)inpA; + (void)inpB; + __builtin_trap(); + return inpC; +#else if constexpr (std::is_same::value) { return __builtin_amdgcn_wmma_f32_16x16x16_f16_w32_gfx12(inpA, inpB, inpC); } else if constexpr (std::is_same::value) { @@ -2412,6 +2421,7 @@ __device__ __forceinline__ floatx8 gcn_wmma16x16x16_instr(const bit16x8& inpA, } else { static_assert(false, "unsupported 16b dtype"); } +#endif } template diff --git a/csrc/rocm/torch_bindings.cpp b/csrc/rocm/torch_bindings.cpp index 03de6dcd1576..53943848f3f7 100644 --- a/csrc/rocm/torch_bindings.cpp +++ b/csrc/rocm/torch_bindings.cpp @@ -14,6 +14,10 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, rocm_ops) { // vLLM custom ops for rocm +// skinny_gemms.cu (LLMM1/wvSplitK/wvSplitKrc/wvSplitKQ) is excluded on gfx1250 +// (gfx9/gfx11 ISA, unsupported there); skip these registrations to avoid +// undefined symbols. vLLM uses default/Triton GEMM for these ops on gfx1250. +#ifndef VLLM_SKIP_SKINNY_GEMMS // Custom gemm op for matrix-vector multiplication rocm_ops.def( "LLMM1(Tensor in_a, Tensor in_b, int rows_per_block) -> " @@ -38,6 +42,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, rocm_ops) { "Tensor scale_a, " " Tensor scale_b, int CuCount) -> ()"); rocm_ops.impl("wvSplitKQ", torch::kCUDA, &wvSplitKQ); +#endif // VLLM_SKIP_SKINNY_GEMMS #ifdef VLLM_ROCM_GFX1100 // W4A16 GPTQ kernels for AMD RDNA3 (gfx1100). diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm index c17444217f02..23506e259b25 100644 --- a/docker/Dockerfile.rocm +++ b/docker/Dockerfile.rocm @@ -28,14 +28,14 @@ ARG SCCACHE_S3_NO_CREDENTIALS=0 FROM ${BASE_IMAGE} AS base -ARG ARG_PYTORCH_ROCM_ARCH +ARG ARG_PYTORCH_ROCM_ARCH=gfx942;gfx950;gfx1250 ENV PYTORCH_ROCM_ARCH=${ARG_PYTORCH_ROCM_ARCH:-${PYTORCH_ROCM_ARCH}} # Install build dependencies and utilities RUN apt-get update -q -y && apt-get install -q -y \ sqlite3 libsqlite3-dev libfmt-dev libmsgpack-dev libsuitesparse-dev \ apt-transport-https ca-certificates wget curl \ - libnuma-dev ccache mold + build-essential libnuma-dev ccache mold RUN --mount=type=cache,target=/root/.cache/pip \ python3 -m pip install --upgrade pip # Note: mold is installed but not set as the system default linker because @@ -61,6 +61,10 @@ ENV UV_HTTP_TIMEOUT=500 ENV UV_INDEX_STRATEGY="unsafe-best-match" # Use copy mode to avoid hardlink failures with Docker cache mounts ENV UV_LINK_MODE=copy +# python binary fall back for non venv builds +ENV UV_PYTHON=${VIRTUAL_ENV:-/usr}/bin/python3 +# Expose paths from wheel installation +ENV PKG_CONFIG_PATH=${ROCM_PATH}/lib/rocm_sysdeps/lib/pkgconfig:${PKG_CONFIG_PATH} # ccache directory - persisted across layer rebuilds via cache mounts. ENV CCACHE_DIR=/root/.cache/ccache ENV CCACHE_COMPILERCHECK=content @@ -110,8 +114,8 @@ WORKDIR ${COMMON_WORKDIR} FROM base AS fetch_vllm_0 ONBUILD COPY ./ vllm/ FROM base AS fetch_vllm_1 -ARG VLLM_REPO="https://github.com/vllm-project/vllm.git" -ARG VLLM_BRANCH="main" +ARG VLLM_REPO="https://github.com/ROCm/vllm.git" +ARG VLLM_BRANCH="455_wip" ENV VLLM_REPO=${VLLM_REPO} ENV VLLM_BRANCH=${VLLM_BRANCH} ONBUILD RUN git clone ${VLLM_REPO} \ @@ -242,7 +246,7 @@ ARG RIXL_BRANCH="39be1de8" ARG RIXL_REPO="https://github.com/ROCm/RIXL.git" ARG UCX_BRANCH="bfb51733" ARG UCX_REPO="https://github.com/openucx/ucx.git" -ENV ROCM_PATH=/opt/rocm +# ENV ROCM_PATH=/opt/rocm -> correct ROCM_PATH is set in base image ENV UCX_HOME=/usr/local/ucx ENV RIXL_HOME=/usr/local/rixl ENV RIXL_BENCH_HOME=/usr/local/rixl_bench @@ -298,6 +302,7 @@ RUN --mount=type=cache,target=/root/.cache/ccache \ git checkout ${RIXL_BRANCH} && \ CC="ccache gcc" CXX="ccache g++" \ meson setup build --prefix=${RIXL_HOME} \ + --force-fallback-for=abseil-cpp \ -Ducx_path=${UCX_HOME} \ -Drocm_path=${ROCM_PATH} && \ cd build && \ @@ -310,6 +315,10 @@ RUN --mount=type=cache,target=/root/.cache/ccache \ RUN cd /opt/rixl && \ sed -i "s/--exclude 'libamdhip64\*'/--exclude 'libamdhip64*' --exclude 'libcore*' --exclude 'libpull*'/" \ contrib/build-wheel.sh && \ + # The wheel build re-runs meson via meson-python; force the bundled abseil + sed -i 's|setup = \["-Dinstall_headers=false"\]|setup = ["-Dinstall_headers=false", "--force-fallback-for=abseil-cpp"]|' \ + pyproject.toml && \ + grep -q 'force-fallback-for' pyproject.toml && \ mkdir -p /app/install && \ _ucx_install_dir=${UCX_HOME} \ ./contrib/build-wheel.sh \ @@ -326,7 +335,7 @@ ARG ROCSHMEM_REPO="https://github.com/ROCm/rocm-systems.git" # DeepEP only supports gfx942 and gfx950; build ROCShmem for the same set so # it can be linked against DeepEP without arch mismatches. ARG DEEPEP_ROCM_ARCH="gfx942;gfx950" -ENV ROCM_PATH=/opt/rocm +# ENV ROCM_PATH=/opt/rocm -> Correct rocm_path is set in base image ENV ROCSHMEM_DIR=/opt/rocshmem RUN --mount=type=cache,target=/root/.cache/ccache \ @@ -575,9 +584,6 @@ RUN --mount=type=cache,target=/root/.cache/uv \ ENV HF_XET_HIGH_PERFORMANCE=1 ENV HF_HUB_DOWNLOAD_TIMEOUT=60 -# Keep torch.cuda.is_available() fork-safe (see vllm/env_override.py). -ENV PYTORCH_NVML_BASED_CUDA_CHECK=1 - # Pre-install vLLM test dependencies. COPY requirements/test/rocm.txt /tmp/rocm-test-reqs.txt RUN --mount=type=cache,target=/root/.cache/uv \ @@ -698,9 +704,6 @@ ENV SAFETENSORS_FAST_GPU=1 # Performance environment variable. ENV HIP_FORCE_DEV_KERNARG=1 -# Keep torch.cuda.is_available() fork-safe (see vllm/env_override.py). -ENV PYTORCH_NVML_BASED_CUDA_CHECK=1 - # Workaround for ROCm profiler limits RUN echo "ROCTRACER_MAX_EVENTS=10000000" > ${COMMON_WORKDIR}/libkineto.conf ENV KINETO_CONFIG="${COMMON_WORKDIR}/libkineto.conf" @@ -708,8 +711,16 @@ RUN echo "VLLM_BASE_IMAGE=${BASE_IMAGE}" >> ${COMMON_WORKDIR}/versions.txt \ && echo "MORI_NIC_BACKEND=${NIC_BACKEND}" >> ${COMMON_WORKDIR}/versions.txt \ && echo "AINIC_VERSION=${AINIC_VERSION}" >> ${COMMON_WORKDIR}/versions.txt + +### Install triton from upstream for AITER Deps TODO: (JPVILLAM) If possible to get this on whls it would be better +RUN pip3 uninstall -y triton && \ + git clone https://github.com/triton-lang/triton.git && \ + cd triton && \ + git checkout c517f38c && \ + TRITON_APPEND_CMAKE_ARGS="-DCMAKE_FIND_USE_CMAKE_ENVIRONMENT_PATH=FALSE" pip3 install . + CMD ["/bin/bash"] #Set entrypoint for vllm-openai official images FROM final AS vllm-openai -ENTRYPOINT ["vllm", "serve"] +ENTRYPOINT ["vllm", "serve"] \ No newline at end of file diff --git a/docker/Dockerfile.rocm_1250_ffm b/docker/Dockerfile.rocm_1250_ffm new file mode 100644 index 000000000000..686d3f7b0864 --- /dev/null +++ b/docker/Dockerfile.rocm_1250_ffm @@ -0,0 +1,14 @@ +# FFM overlay for gfx1250 builds +ARG BASE_IMAGE=rocm/vllm-private:juan_455_npi_test +FROM ${BASE_IMAGE} + +# extract tarball into /root/x/ffm/. +ADD ubuntu_24_04_rel.tar.gz /root/x/ffm/ +ENV HSA_MODEL_LIB=/root/x/ffm/libhsakmtmodel.so +ENV HSA_MODEL_TOML="/root/x/ffm/ffm_config.toml" +ENV HSA_MODEL_ARGS=ffm_enable_time_slicing +ENV HSA_MODEL_TOPOLOGY=/root/x/ffm/topology/mi450 +ENV HSA_MODEL_NUM_THREADS=64 + +### Remove AMD SMI +RUN pip uninstall -y amdsmi && rm -rf ${SITE_PACKAGES}/_rocm_sdk_core/share/amd_smi/amdsmi diff --git a/docker/Dockerfile.rocm_base b/docker/Dockerfile.rocm_base index fbd5e1e60e3b..20833d99e95f 100644 --- a/docker/Dockerfile.rocm_base +++ b/docker/Dockerfile.rocm_base @@ -1,15 +1,16 @@ -ARG BASE_IMAGE=rocm/dev-ubuntu-22.04:7.2.3-complete -ARG TRITON_BRANCH="0f380657" -ARG TRITON_REPO="https://github.com/ROCm/triton.git" -ARG PYTORCH_BRANCH="d0c8b1f3" # release/2.11 as of 6/09 -ARG PYTORCH_REPO="https://github.com/ROCm/pytorch.git" -ARG PYTORCH_VISION_BRANCH="v0.24.1" -ARG PYTORCH_VISION_REPO="https://github.com/pytorch/vision.git" -ARG PYTORCH_AUDIO_BRANCH="v2.9.0" -ARG PYTORCH_AUDIO_REPO="https://github.com/pytorch/audio.git" +ARG BASE_IMAGE=ubuntu:24.04 +ARG ROCM_WHEEL_INDEX=https://rocm.devreleases.amd.com/whl-multi-arch/ +ARG ROCM_SDK_VERSION=7.14.0a20260623 +ARG TORCH_VERSION=2.11.0+rocm7.14.0a20260623 +ARG TORCHVISION_VERSION=0.26.0+rocm7.14.0a20260623 +ARG TORCHAUDIO_VERSION=2.11.0+rocm7.14.0a20260623 +ARG TRITON_VERSION=3.7.1+git110cd8e2.rocm7.14.0a20260623 +ARG APEX_VERSION=1.11.0+rocm7.14.0a20260623 + + ARG FA_BRANCH="0e60e394" ARG FA_REPO="https://github.com/Dao-AILab/flash-attention.git" -ARG AITER_BRANCH="v0.1.16.post2" +ARG AITER_BRANCH="main" ARG AITER_REPO="https://github.com/ROCm/aiter.git" ARG MORI_BRANCH="v1.1.0" ARG MORI_REPO="https://github.com/ROCm/mori.git" @@ -24,13 +25,15 @@ ARG SCCACHE_S3_NO_CREDENTIALS=0 FROM ${BASE_IMAGE} AS base -ENV PATH=/opt/rocm/llvm/bin:/opt/rocm/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin -ENV ROCM_PATH=/opt/rocm -ENV LD_LIBRARY_PATH=/opt/rocm/lib:/usr/local/lib: -ARG PYTORCH_ROCM_ARCH=gfx90a;gfx942;gfx950;gfx1100;gfx1101;gfx1200;gfx1201;gfx1150;gfx1151 +ARG PYTORCH_ROCM_ARCH=gfx942;gfx950;gfx1250 ENV PYTORCH_ROCM_ARCH=${PYTORCH_ROCM_ARCH} -ENV AITER_ROCM_ARCH=gfx942;gfx950 +ENV AITER_ROCM_ARCH=${PYTORCH_ROCM_ARCH} ENV MORI_GPU_ARCHS=gfx942;gfx950 +ENV FA_GPU_ARCHS=gfx942;gfx950 + +# TODO: Unset these when support is available for gfx1250 +ENV ENABLE_CK=0 +ARG PREBUILD_KERNELS=0 # Required for RCCL in ROCm7.1 ENV HSA_NO_SCRATCH_RECLAIM=1 @@ -44,19 +47,23 @@ ENV DEBIAN_FRONTEND=noninteractive # Install Python and other dependencies RUN apt-get update -y \ - && apt-get install -y software-properties-common git curl sudo vim less libgfortran5 libopenmpi-dev libpci-dev liblzma-dev pkg-config \ + && apt-get install -y software-properties-common git curl sudo vim less libgfortran5 libopenmpi-dev libpci-dev liblzma-dev libnuma-dev libdrm-dev pkg-config g++ \ && for i in 1 2 3; do \ - add-apt-repository -y ppa:deadsnakes/ppa && break || \ - { echo "Attempt $i failed, retrying in 5s..."; sleep 5; }; \ + add-apt-repository -y ppa:deadsnakes/ppa && break || \ + { echo "Attempt $i failed, retrying in 5s..."; sleep 5; }; \ done \ && apt-get update -y \ && apt-get install -y python${PYTHON_VERSION} python${PYTHON_VERSION}-dev python${PYTHON_VERSION}-venv \ - python${PYTHON_VERSION}-lib2to3 python-is-python3 \ + python${PYTHON_VERSION}-lib2to3 python-is-python3 \ && update-alternatives --install /usr/bin/python3 python3 /usr/bin/python${PYTHON_VERSION} 1 \ && update-alternatives --set python3 /usr/bin/python${PYTHON_VERSION} \ && ln -sf /usr/bin/python${PYTHON_VERSION}-config /usr/bin/python3-config \ - && curl -sS https://bootstrap.pypa.io/get-pip.py | python${PYTHON_VERSION} \ - && python3 --version && python3 -m pip --version + && python3 --version + +ENV VIRTUAL_ENV=/opt/venv +RUN python${PYTHON_VERSION} -m venv "${VIRTUAL_ENV}" && \ + "${VIRTUAL_ENV}/bin/python" -m pip install --upgrade pip setuptools PyYAML +ENV PATH=${VIRTUAL_ENV}/bin:$PATH RUN pip install -U packaging 'cmake<4' ninja wheel 'setuptools<80' pybind11 Cython RUN apt-get update && apt-get install -y libjpeg-dev libsox-dev libsox-fmt-all sox && rm -rf /var/lib/apt/lists/* @@ -69,31 +76,105 @@ ARG SCCACHE_BUCKET_NAME ARG SCCACHE_REGION_NAME ARG SCCACHE_S3_NO_CREDENTIALS RUN if [ "$USE_SCCACHE" = "1" ]; then \ - echo "Installing sccache..." \ - && SCCACHE_ARCH="x86_64" \ - && SCCACHE_VERSION="v0.8.1" \ - && SCCACHE_DL_URL="${SCCACHE_DOWNLOAD_URL:-https://github.com/mozilla/sccache/releases/download/${SCCACHE_VERSION}/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl.tar.gz}" \ - && curl -L -o /tmp/sccache.tar.gz ${SCCACHE_DL_URL} \ - && tar -xzf /tmp/sccache.tar.gz -C /tmp \ - && mv /tmp/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl/sccache /usr/bin/sccache \ - && chmod +x /usr/bin/sccache \ - && rm -rf /tmp/sccache.tar.gz /tmp/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl \ - && sccache --version; \ + echo "Installing sccache..." \ + && SCCACHE_ARCH="x86_64" \ + && SCCACHE_VERSION="v0.8.1" \ + && SCCACHE_DL_URL="${SCCACHE_DOWNLOAD_URL:-https://github.com/mozilla/sccache/releases/download/${SCCACHE_VERSION}/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl.tar.gz}" \ + && curl -L -o /tmp/sccache.tar.gz ${SCCACHE_DL_URL} \ + && tar -xzf /tmp/sccache.tar.gz -C /tmp \ + && mv /tmp/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl/sccache /usr/bin/sccache \ + && chmod +x /usr/bin/sccache \ + && rm -rf /tmp/sccache.tar.gz /tmp/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl \ + && sccache --version; \ + fi + +## +## Install PyTorch w/ Triton + ROCM_SDK from ROCM wheel index +## +ARG ROCM_WHEEL_INDEX +ARG TORCH_VERSION +ARG TORCHVISION_VERSION +ARG TORCHAUDIO_VERSION +ARG ROCM_SDK_VERSION +ARG APEX_VERSION +ENV SITE_PACKAGES=${VIRTUAL_ENV}/lib/python${PYTHON_VERSION}/site-packages +ENV ROCM_PATH=${SITE_PACKAGES}/_rocm_sdk_devel +ENV ROCM_HOME=${ROCM_PATH} +ENV ROCM_SOURCE_DIR=${ROCM_PATH} +ENV ROCM_BIN=${ROCM_PATH}/bin +ENV ROCM_CMAKE_PREFIX=${ROCM_PATH}/lib/cmake +ENV HIP_DEVICE_LIB_PATH=${SITE_PACKAGES}/_rocm_sdk_core/lib/llvm/amdgcn/bitcode +ENV PATH=${ROCM_PATH}/bin:${ROCM_PATH}/llvm/bin:$PATH +ENV LD_LIBRARY_PATH=${ROCM_PATH}/lib:${SITE_PACKAGES}/_rocm_sdk_core/lib +ENV CMAKE_PREFIX_PATH=${ROCM_PATH}/lib/cmake:${SITE_PACKAGES}/torch/share/cmake +ENV PYTHONPATH=${SITE_PACKAGES}/_rocm_sdk_core/share/amd_smi + +# torch/torchvision/torchaudio must be pinned to mutually-consistent builds +# (same +rocm... suffix) or the C++ ops break at import (ABI skew). The rocm +# sdk version is derived from torch's own dependency pin unless overridden, +# which keeps the set consistent and avoids pip backtracking. +RUN pip install --pre --index-url "${ROCM_WHEEL_INDEX}" \ + --extra-index-url https://pypi.org/simple \ + "torch[device-all]==${TORCH_VERSION}" \ + "torchvision==${TORCHVISION_VERSION}" \ + "torchaudio==${TORCHAUDIO_VERSION}" \ + "rocm[libraries,devel,device-all]==${ROCM_SDK_VERSION}" && \ + rocm-sdk init + +# Torch runtime deps that may not be published on the ROCm wheel index; +# install them from PyPI afterwards. +RUN pip install filelock "typing-extensions>=4.10.0" "sympy>=1.13.3" \ + "networkx>=2.5.1" jinja2 "fsspec>=0.8.5" + + + +# Expose the rocm-sdk wheel as a conventional /opt/rocm install so downstream +# builds (Dockerfile.rocm: vLLM csrc, RIXL/UCX, ROCShmem/DeepEP) keep working. +RUN ln -sfn "${ROCM_PATH}" /opt/rocm; + +RUN if [ -f "${SITE_PACKAGES}/rocm_sdk/__init__.py" ]; then \ + sed -i 's/rtld_global: bool = True/rtld_global: bool = False/g' \ + "${SITE_PACKAGES}/rocm_sdk/__init__.py"; \ fi +# The ROCm SDK wheel ships a broken CMake export for hsakmt +# This patch includes the right paths for the numa build target +RUN <<'EOF' +set -eu +TARGETS="/opt/rocm/lib/cmake/hsakmt/hsakmtTargets.cmake" +[ -f "$TARGETS" ] || exit 0 # nothing to patch +grep -q NUMA_LIBRARY "$TARGETS" && exit 0 # already patched + +# 1. Point libdrm's -L at the copy bundled in the wheel, not the builder path. +sed -i 's|-L/__w/[^;"]*|-L${_IMPORT_PREFIX}/lib/rocm_sysdeps/lib|g' "$TARGETS" + +# 2. Drop the nonexistent RHEL libc path (libc is linked implicitly anyway). +sed -i 's|/usr/lib64/libc.so;||g' "$TARGETS" + +# 3. Define the numa::numa target the export references but forgot to create. +cat >> "$TARGETS" <<'CMAKE' + +if(NOT TARGET numa::numa) + find_library(NUMA_LIBRARY NAMES numa REQUIRED) + add_library(numa::numa UNKNOWN IMPORTED) + set_target_properties(numa::numa PROPERTIES IMPORTED_LOCATION "${NUMA_LIBRARY}") +endif() +CMAKE +EOF + # Setup sccache for HIP compilation via HIP_CLANG_PATH # This creates wrapper scripts in a separate directory and points HIP to use them # This avoids modifying the original ROCm binaries which can break detection # NOTE: HIP_CLANG_PATH is NOT set as ENV to avoid affecting downstream images (Dockerfile.rocm) # Instead, each build stage should export HIP_CLANG_PATH=/opt/sccache-wrappers if USE_SCCACHE=1 RUN if [ "$USE_SCCACHE" = "1" ]; then \ - echo "Setting up sccache wrappers for HIP compilation..." \ - && mkdir -p /opt/sccache-wrappers \ - && printf '#!/bin/bash\nexec sccache /opt/rocm/lib/llvm/bin/clang++ "$@"\n' > /opt/sccache-wrappers/clang++ \ - && chmod +x /opt/sccache-wrappers/clang++ \ - && printf '#!/bin/bash\nexec sccache /opt/rocm/lib/llvm/bin/clang "$@"\n' > /opt/sccache-wrappers/clang \ - && chmod +x /opt/sccache-wrappers/clang \ - && echo "sccache wrappers created in /opt/sccache-wrappers"; \ + echo "Setting up sccache wrappers for HIP compilation..." \ + && mkdir -p /opt/sccache-wrappers \ + && printf '#!/bin/bash\nexec sccache ${ROCM_PATH}/lib/llvm/bin/clang++ "$@"\n' > /opt/sccache-wrappers/clang++ \ + && chmod +x /opt/sccache-wrappers/clang++ \ + && printf '#!/bin/bash\nexec sccache ${ROCM_PATH}/lib/llvm/bin/clang "$@"\n' > /opt/sccache-wrappers/clang \ + && chmod +x /opt/sccache-wrappers/clang \ + && echo "sccache wrappers created in /opt/sccache-wrappers"; \ fi # Set sccache environment variables only when USE_SCCACHE=1 @@ -105,103 +186,29 @@ ENV SCCACHE_S3_NO_CREDENTIALS=${USE_SCCACHE:+${SCCACHE_S3_NO_CREDENTIALS}} ENV SCCACHE_IDLE_TIMEOUT=${USE_SCCACHE:+0} -### -### Triton Build -### -FROM base AS build_triton -ARG TRITON_BRANCH -ARG TRITON_REPO -RUN git clone ${TRITON_REPO} -# Cherry picking the following -# https://github.com/triton-lang/triton/pull/8991 -RUN cd triton \ - && git checkout ${TRITON_BRANCH} \ - && git config --global user.email "you@example.com" && git config --global user.name "Your Name" \ - && git cherry-pick 555d04f \ - && if [ ! -f setup.py ]; then cd python; fi \ - && python3 setup.py bdist_wheel --dist-dir=dist \ - && mkdir -p /app/install && cp dist/*.whl /app/install -RUN if [ -d triton/python/triton_kernels ]; then pip install build && cd triton/python/triton_kernels \ - && python3 -m build --wheel && cp dist/*.whl /app/install; fi - - ### ### AMD SMI Build ### FROM base AS build_amdsmi -RUN cd /opt/rocm/share/amd_smi \ +RUN cd ${SITE_PACKAGES}/_rocm_sdk_core/share/amd_smi \ && pip wheel . --wheel-dir=dist -RUN mkdir -p /app/install && cp /opt/rocm/share/amd_smi/dist/*.whl /app/install - - -### -### Pytorch build -### -FROM base AS build_pytorch -ARG PYTORCH_BRANCH -ARG PYTORCH_VISION_BRANCH -ARG PYTORCH_AUDIO_BRANCH -ARG PYTORCH_REPO -ARG PYTORCH_VISION_REPO -ARG PYTORCH_AUDIO_REPO -ARG USE_SCCACHE - -RUN apt-get update && apt-get install -y pkg-config liblzma-dev -RUN git clone ${PYTORCH_REPO} pytorch -RUN cd pytorch && git checkout ${PYTORCH_BRANCH} -RUN cd pytorch \ - && pip install -r requirements.txt && git submodule update --init --recursive -RUN cd pytorch && python3 tools/amd_build/build_amd.py \ - && if [ "$USE_SCCACHE" = "1" ]; then \ - export HIP_CLANG_PATH=/opt/sccache-wrappers \ - && export CMAKE_C_COMPILER_LAUNCHER=sccache \ - && export CMAKE_CXX_COMPILER_LAUNCHER=sccache \ - && sccache --show-stats; \ - fi \ - && CMAKE_PREFIX_PATH=$(python3 -c 'import sys; print(sys.prefix)') python3 setup.py bdist_wheel --dist-dir=dist \ - && if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi \ - && pip install dist/*.whl -RUN git clone ${PYTORCH_VISION_REPO} vision -RUN cd vision && git checkout ${PYTORCH_VISION_BRANCH} \ - && if [ "$USE_SCCACHE" = "1" ]; then \ - export HIP_CLANG_PATH=/opt/sccache-wrappers \ - && export CMAKE_C_COMPILER_LAUNCHER=sccache \ - && export CMAKE_CXX_COMPILER_LAUNCHER=sccache; \ - fi \ - && python3 setup.py bdist_wheel --dist-dir=dist \ - && if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi \ - && pip install dist/*.whl -RUN git clone ${PYTORCH_AUDIO_REPO} audio -RUN cd audio && git checkout ${PYTORCH_AUDIO_BRANCH} \ - && git submodule update --init --recursive \ - && pip install -r requirements.txt \ - && if [ "$USE_SCCACHE" = "1" ]; then \ - export HIP_CLANG_PATH=/opt/sccache-wrappers \ - && export CMAKE_C_COMPILER_LAUNCHER=sccache \ - && export CMAKE_CXX_COMPILER_LAUNCHER=sccache; \ - fi \ - && python3 setup.py bdist_wheel --dist-dir=dist \ - && if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi \ - && pip install dist/*.whl -RUN mkdir -p /app/install && cp /app/pytorch/dist/*.whl /app/install \ - && cp /app/vision/dist/*.whl /app/install \ - && cp /app/audio/dist/*.whl /app/install +RUN mkdir -p /app/install && cp ${SITE_PACKAGES}/_rocm_sdk_core/share/amd_smi/dist/*.whl /app/install ### -### MORI Build +### MORI Build TODO(Build needs fixing) ### FROM base AS build_mori ARG MORI_BRANCH ARG MORI_REPO -RUN --mount=type=bind,from=build_pytorch,src=/app/install/,target=/install \ - pip install /install/*.whl -RUN git clone ${MORI_REPO} -RUN cd mori \ +ARG MORI_GPU_ARCHS +RUN mkdir -p /app/install; \ + git clone ${MORI_REPO} \ + && cd mori \ && git checkout ${MORI_BRANCH} \ && git submodule update --init --recursive \ - && python3 setup.py bdist_wheel --dist-dir=dist && ls /app/mori/dist/*.whl -RUN mkdir -p /app/install && cp /app/mori/dist/*.whl /app/install + && python3 setup.py bdist_wheel --dist-dir=dist && ls /app/mori/dist/*.whl \ + && cp /app/mori/dist/*.whl /app/install; ### @@ -211,19 +218,18 @@ FROM base AS build_fa ARG FA_BRANCH ARG FA_REPO ARG USE_SCCACHE -RUN --mount=type=bind,from=build_pytorch,src=/app/install/,target=/install \ - pip install /install/*.whl -RUN git clone ${FA_REPO} -RUN cd flash-attention \ +RUN mkdir -p /app/install; \ + git clone ${FA_REPO} \ + && cd flash-attention \ && git checkout ${FA_BRANCH} \ && git submodule update --init \ && if [ "$USE_SCCACHE" = "1" ]; then \ - export HIP_CLANG_PATH=/opt/sccache-wrappers \ - && sccache --show-stats; \ - fi \ - && GPU_ARCHS=$(echo ${PYTORCH_ROCM_ARCH} | sed -e 's/;gfx1[0-9]\{3\}//g') python3 setup.py bdist_wheel --dist-dir=dist \ - && if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi -RUN mkdir -p /app/install && cp /app/flash-attention/dist/*.whl /app/install + export HIP_CLANG_PATH=/opt/sccache-wrappers \ + && sccache --show-stats; \ + fi \ + && GPU_ARCHS=$(echo ${FA_GPU_ARCHS} | sed -e 's/;gfx1[0-9]\{3\}//g') python3 setup.py bdist_wheel --dist-dir=dist \ + && if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi \ + && cp dist/*.whl /app/install; ### @@ -233,18 +239,16 @@ FROM base AS build_aiter ARG AITER_BRANCH ARG AITER_REPO ARG USE_SCCACHE -RUN --mount=type=bind,from=build_pytorch,src=/app/install/,target=/install \ - pip install /install/*.whl RUN git clone --recursive --branch ${AITER_BRANCH} ${AITER_REPO} RUN cd aiter \ && git submodule update --init --recursive \ && pip install -r requirements.txt RUN pip install pyyaml && cd aiter \ && if [ "$USE_SCCACHE" = "1" ]; then \ - export HIP_CLANG_PATH=/opt/sccache-wrappers \ - && sccache --show-stats; \ - fi \ - && PREBUILD_KERNELS=1 AITER_USE_SYSTEM_TRITON=1 GPU_ARCHS=${AITER_ROCM_ARCH} python3 setup.py bdist_wheel --dist-dir=dist \ + export HIP_CLANG_PATH=/opt/sccache-wrappers \ + && sccache --show-stats; \ + fi \ + && AITER_USE_SYSTEM_TRITON=1 PREBUILD_KERNELS=${PREBUILD_KERNELS} GPU_ARCHS=${AITER_ROCM_ARCH} python3 setup.py bdist_wheel --dist-dir=dist \ && if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi \ && ls /app/aiter/dist/*.whl RUN mkdir -p /app/install && cp /app/aiter/dist/*.whl /app/install @@ -258,46 +262,35 @@ RUN mkdir -p /app/install && cp /app/aiter/dist/*.whl /app/install # only includes dependencies used by wheel release pipeline FROM base AS debs_wheel_release RUN mkdir /app/debs -RUN --mount=type=bind,from=build_triton,src=/app/install/,target=/install \ - cp /install/*.whl /app/debs RUN --mount=type=bind,from=build_fa,src=/app/install/,target=/install \ - cp /install/*.whl /app/debs + if ls /install/*.whl >/dev/null 2>&1; then cp /install/*.whl /app/debs; fi RUN --mount=type=bind,from=build_amdsmi,src=/app/install/,target=/install \ cp /install/*.whl /app/debs -RUN --mount=type=bind,from=build_pytorch,src=/app/install/,target=/install \ - cp /install/*.whl /app/debs RUN --mount=type=bind,from=build_aiter,src=/app/install/,target=/install \ cp /install/*.whl /app/debs # Full debs stage - includes Mori (used by Docker releases) FROM base AS debs RUN mkdir /app/debs -RUN --mount=type=bind,from=build_triton,src=/app/install/,target=/install \ - cp /install/*.whl /app/debs RUN --mount=type=bind,from=build_fa,src=/app/install/,target=/install \ - cp /install/*.whl /app/debs + if ls /install/*.whl >/dev/null 2>&1; then cp /install/*.whl /app/debs; fi RUN --mount=type=bind,from=build_amdsmi,src=/app/install/,target=/install \ cp /install/*.whl /app/debs -RUN --mount=type=bind,from=build_pytorch,src=/app/install/,target=/install \ - cp /install/*.whl /app/debs RUN --mount=type=bind,from=build_aiter,src=/app/install/,target=/install \ cp /install/*.whl /app/debs RUN --mount=type=bind,from=build_mori,src=/app/install/,target=/install \ - cp /install/*.whl /app/debs + if ls /install/*.whl >/dev/null 2>&1; then cp /install/*.whl /app/debs; fi FROM base AS final RUN --mount=type=bind,from=debs,src=/app/debs,target=/install \ pip install /install/*.whl ARG BASE_IMAGE -ARG TRITON_BRANCH -ARG TRITON_REPO -ARG PYTORCH_BRANCH -ARG PYTORCH_VISION_BRANCH -ARG PYTORCH_REPO -ARG PYTORCH_VISION_REPO -ARG PYTORCH_AUDIO_BRANCH -ARG PYTORCH_AUDIO_REPO +ARG ROCM_WHEEL_INDEX +ARG ROCM_SDK_VERSION +ARG TORCH_VERSION +ARG TORCHVISION_VERSION +ARG TORCHAUDIO_VERSION ARG FA_BRANCH ARG FA_REPO ARG AITER_BRANCH @@ -305,14 +298,11 @@ ARG AITER_REPO ARG MORI_BRANCH ARG MORI_REPO RUN echo "BASE_IMAGE: ${BASE_IMAGE}" > /app/versions.txt \ - && echo "TRITON_BRANCH: ${TRITON_BRANCH}" >> /app/versions.txt \ - && echo "TRITON_REPO: ${TRITON_REPO}" >> /app/versions.txt \ - && echo "PYTORCH_BRANCH: ${PYTORCH_BRANCH}" >> /app/versions.txt \ - && echo "PYTORCH_VISION_BRANCH: ${PYTORCH_VISION_BRANCH}" >> /app/versions.txt \ - && echo "PYTORCH_REPO: ${PYTORCH_REPO}" >> /app/versions.txt \ - && echo "PYTORCH_VISION_REPO: ${PYTORCH_VISION_REPO}" >> /app/versions.txt \ - && echo "PYTORCH_AUDIO_BRANCH: ${PYTORCH_AUDIO_BRANCH}" >> /app/versions.txt \ - && echo "PYTORCH_AUDIO_REPO: ${PYTORCH_AUDIO_REPO}" >> /app/versions.txt \ + && echo "ROCM_WHEEL_INDEX: ${ROCM_WHEEL_INDEX}" >> /app/versions.txt \ + && echo "ROCM_SDK_VERSION: ${ROCM_SDK_VERSION}" >> /app/versions.txt \ + && echo "TORCH_VERSION: ${TORCH_VERSION}" >> /app/versions.txt \ + && echo "TORCHVISION_VERSION: ${TORCHVISION_VERSION}" >> /app/versions.txt \ + && echo "TORCHAUDIO_VERSION: ${TORCHAUDIO_VERSION}" >> /app/versions.txt \ && echo "FA_BRANCH: ${FA_BRANCH}" >> /app/versions.txt \ && echo "FA_REPO: ${FA_REPO}" >> /app/versions.txt \ && echo "AITER_BRANCH: ${AITER_BRANCH}" >> /app/versions.txt \ diff --git a/requirements/rocm.txt b/requirements/rocm.txt index 5179f6ee8d74..259016458866 100644 --- a/requirements/rocm.txt +++ b/requirements/rocm.txt @@ -6,7 +6,6 @@ grpcio==1.78.0 grpcio-reflection==1.78.0 numba == 0.65.0 # Required for N-gram speculative decoding - # Dependencies for AMD GPUs datasets peft @@ -22,7 +21,7 @@ timm>=1.0.17 # amd-quark: required for Quark quantization on ROCm # To be consistent with test_quark.py amd-quark>=0.8.99 -tilelang==0.1.10 +#tilelang==0.1.10 # Required apache-tvm-ffi matching tilelang version apache-tvm-ffi==0.1.10 # Required for faster safetensors model loading diff --git a/tests/entrypoints/speech_to_text/transcription/test_transcription_validation_whisper.py b/tests/entrypoints/speech_to_text/transcription/test_transcription_validation_whisper.py index 511179f7fcb1..c50980de2543 100644 --- a/tests/entrypoints/speech_to_text/transcription/test_transcription_validation_whisper.py +++ b/tests/entrypoints/speech_to_text/transcription/test_transcription_validation_whisper.py @@ -31,7 +31,7 @@ def _get_attention_backend_params() -> list[str | None]: falls back to ROCM_AITER_UNIFIED_ATTN or TRITON_ATTN for cross-attention since ROCM_ATTN doesn't support ENCODER_DECODER) - TRITON_ATTN: always available on ROCm - - ROCM_AITER_UNIFIED_ATTN: only on gfx942/gfx950 + - ROCM_AITER_UNIFIED_ATTN: only on gfx942/gfx950/gfx1250 On non-ROCm platforms, we just run with the default backend. """ @@ -40,9 +40,9 @@ def _get_attention_backend_params() -> list[str | None]: if current_platform.is_rocm(): backends: list[str | None] = [None, "TRITON_ATTN"] - from vllm.platforms.rocm import _ON_MI3XX + from vllm.platforms.rocm import get_cdna_version - if _ON_MI3XX: + if get_cdna_version() > 2: backends.append("ROCM_AITER_UNIFIED_ATTN") return backends except Exception: diff --git a/tests/entrypoints/speech_to_text/translation/test_translation_validation.py b/tests/entrypoints/speech_to_text/translation/test_translation_validation.py index ed3cff5f1c22..25f26404d0d9 100644 --- a/tests/entrypoints/speech_to_text/translation/test_translation_validation.py +++ b/tests/entrypoints/speech_to_text/translation/test_translation_validation.py @@ -37,13 +37,13 @@ def _get_rocm_attention_config(model_name): if "whisper" in model_name.lower(): try: - from vllm.platforms.rocm import _ON_MI3XX + from vllm.platforms.rocm import get_cdna_version - if _ON_MI3XX: + if get_cdna_version() > 2: return {"backend": "ROCM_AITER_UNIFIED_ATTN"} except ImportError: logger.warning( - "Could not import _ON_MI3XX from rocm platform, " + "Could not check cdna version from rocm platform, " "falling back to TRITON_ATTN for Whisper." ) return {"backend": "TRITON_ATTN"} diff --git a/tests/kernels/attention/test_rocm_aiter_unified_attn.py b/tests/kernels/attention/test_rocm_aiter_unified_attn.py index c02a457c98a4..9d538dba43f4 100644 --- a/tests/kernels/attention/test_rocm_aiter_unified_attn.py +++ b/tests/kernels/attention/test_rocm_aiter_unified_attn.py @@ -15,15 +15,15 @@ from vllm.platforms import current_platform from vllm.utils.torch_utils import set_random_seed -_SKIP_NON_MI3XX = True +_SKIP_NON_CDNA_2_PLUS = True if current_platform.is_rocm(): - from vllm.platforms.rocm import on_mi3xx + from vllm.platforms.rocm import get_cdna_version - _SKIP_NON_MI3XX = not on_mi3xx() + _SKIP_NON_CDNA_2_PLUS = get_cdna_version() < 2 pytestmark = [ pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm-specific tests"), - pytest.mark.skipif(_SKIP_NON_MI3XX, reason="MI300/MI350 ROCm only"), + pytest.mark.skipif(_SKIP_NON_CDNA_2_PLUS, reason="CDNA 2+ ROCm only"), ] NUM_Q_HEADS = 8 diff --git a/tests/kernels/moe/test_modular_oai_triton_moe.py b/tests/kernels/moe/test_modular_oai_triton_moe.py index 0315d8d89e56..585aa34d5966 100644 --- a/tests/kernels/moe/test_modular_oai_triton_moe.py +++ b/tests/kernels/moe/test_modular_oai_triton_moe.py @@ -4,6 +4,12 @@ Test modular OAI Triton MoE """ +from __future__ import annotations + +import json +import os +from pathlib import Path + import pytest import torch @@ -48,6 +54,52 @@ ] +def deepseek_v4_flash_moe_topology(): + """MoE sizes for the MODEL path used by ``launch_dsv4.sh``. + + Default weights path: ``/data/deepseek-ai/DeepSeek-V4-Flash/config.json``. + When that file is readable, values come from ``hidden_size``, + ``moe_intermediate_size``, ``n_routed_experts``, and ``num_experts_per_tok``. + Otherwise fall back to the same numeric constants. + """ + defaults = { + "hidden_size": 4096, + "moe_intermediate_size": 2048, + "n_routed_experts": 256, + "num_experts_per_tok": 6, + } + cfg_path = Path( + os.environ.get( + "DEEPSEEK_V4_FLASH_CONFIG", + "/data/deepseek-ai/DeepSeek-V4-Flash/config.json", + ) + ) + if cfg_path.is_file(): + with cfg_path.open() as f: + cfg = json.load(f) + return { + "hidden_size": int(cfg["hidden_size"]), + "moe_intermediate_size": int(cfg["moe_intermediate_size"]), + "n_routed_experts": int(cfg["n_routed_experts"]), + "num_experts_per_tok": int(cfg["num_experts_per_tok"]), + } + return defaults + + +def scaled_deepseek_v4_flash_problem( + *, + dim_scale: int = 8, + expert_scale: int = 8, +): + """Smaller K/N/E for kernel tests; keeps production top_k and K:N ratio (~2:1).""" + t = deepseek_v4_flash_moe_topology() + k = max(128, t["hidden_size"] // dim_scale) + n = max(64, t["moe_intermediate_size"] // dim_scale) + num_experts = max(t["num_experts_per_tok"], t["n_routed_experts"] // expert_scale) + topk = t["num_experts_per_tok"] + return k, n, num_experts, topk + + def unshuffle_weight(w: torch.Tensor): first = w[..., ::2] second = w[..., 1::2] @@ -261,3 +313,95 @@ def test_oai_triton_moe( ) assert_close(ref=out_ref, tri=out, maxtol=0.025, rmstol=0.005) + + +@pytest.mark.skipif( + not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform." +) +def test_unfused_oai_triton_experts_apply_direct_deepseek_v4_topology(workspace_init): + """Exercise ``UnfusedOAITritonExperts.apply`` with explicit workspaces. + + Same MoE topology as ``launch_dsv4.sh`` / DeepSeek-V4-Flash ``config.json``, + with linear dimensions and expert count scaled down for test GPU memory. + """ + wait_for_gpu_memory_to_clear(devices=[0], threshold_ratio=0.1) + set_random_seed(0) + + k, n, num_experts, topk = scaled_deepseek_v4_flash_problem() + m = 7 + dtype = torch.bfloat16 + + ( + w1, + w2, + w1_bias, + w2_bias, + w1_tri, + w2_tri, + w1_bias_tri, + w2_bias_tri, + w1_precision_config, + w2_precision_config, + ) = make_weights(dtype, k, n, num_experts) + + x = torch.randn((m, k), dtype=dtype, device="cuda") + router_logits = torch.randn(m, num_experts, device="cuda", dtype=dtype) + topk_weights, topk_ids = torch.topk(router_logits, k=topk, dim=-1, sorted=True) + topk_weights = torch.nn.functional.softmax(topk_weights, dim=-1) + + quant_config = mxfp4_w4a16_moe_quant_config( + w1_bias=w1_bias_tri, + w2_bias=w2_bias_tri, + w1_scale=w1_precision_config, + w2_scale=w2_precision_config, + ) + moe_config = make_dummy_moe_config( + num_experts=num_experts, + experts_per_token=topk, + hidden_dim=k, + intermediate_size_per_partition=n, + ) + experts = UnfusedOAITritonExperts(moe_config, quant_config) + + if not UnfusedOAITritonExperts._supports_current_device(): + pytest.skip("UnfusedOAITritonExperts does not support this device") + + _, _, N, K, top_k = experts.moe_problem_size(x, w1_tri, w2_tri, topk_ids) + assert top_k == topk + ws13_shape, ws2_shape, out_shape = experts.workspace_shapes( + m, + N, + K, + topk, + num_experts, + num_experts, + None, + MoEActivation.SWIGLUOAI, + ) + workspace13 = torch.empty(ws13_shape, dtype=dtype, device="cuda") + workspace2 = torch.empty(ws2_shape, dtype=dtype, device="cuda") + output = torch.empty(out_shape, dtype=dtype, device="cuda") + + with set_current_vllm_config(VllmConfig()): + out_ref = torch_moe_impl( + x, w1, w2, w1_bias, w2_bias, topk_weights, topk_ids + ) + experts.apply( + output=output, + hidden_states=x, + w1=w1_tri, + w2=w2_tri, + topk_weights=topk_weights, + topk_ids=topk_ids, + activation=MoEActivation.SWIGLUOAI, + global_num_experts=num_experts, + expert_map=None, + a1q_scale=None, + a2_scale=None, + workspace13=workspace13, + workspace2=workspace2, + expert_tokens_meta=None, + apply_router_weight_on_input=False, + ) + + assert_close(ref=out_ref, tri=output, maxtol=0.025, rmstol=0.005) diff --git a/tests/kernels/moe/test_routing.py b/tests/kernels/moe/test_routing.py index 62a4968a0d1f..59e93092281e 100644 --- a/tests/kernels/moe/test_routing.py +++ b/tests/kernels/moe/test_routing.py @@ -19,13 +19,13 @@ def _is_aiter_capable() -> bool: - """Check if the platform supports AITER (gfx942/gfx950).""" + """Check if the platform supports AITER (gfx942/gfx950/gfx1250).""" if not current_platform.is_rocm(): return False try: - from vllm.platforms.rocm import _ON_MI3XX + from vllm.platforms.rocm import get_cdna_version - return _ON_MI3XX + return get_cdna_version() > 2 except ImportError: return False diff --git a/tests/model_executor/layers/test_rocm_unquantized_gemm.py b/tests/model_executor/layers/test_rocm_unquantized_gemm.py index f4de9bc9038f..f20be7cfb618 100644 --- a/tests/model_executor/layers/test_rocm_unquantized_gemm.py +++ b/tests/model_executor/layers/test_rocm_unquantized_gemm.py @@ -26,6 +26,7 @@ def test_rocm_unquantized_gemm_gfx1x_wvsplitk_path(monkeypatch): monkeypatch.setattr("vllm.platforms.rocm.on_gfx1x", lambda: True) monkeypatch.setattr("vllm.platforms.rocm.on_gfx9", lambda: False) monkeypatch.setattr("vllm.platforms.rocm.on_gfx950", lambda: False) + monkeypatch.setattr("vllm.platforms.rocm.on_gfx1250", lambda: False) monkeypatch.setattr(utils, "num_compute_units", lambda: 120) wvsplitk_mock = MagicMock(side_effect=lambda w, x_view, _, __: x_view @ w.t()) @@ -52,6 +53,7 @@ def test_rocm_unquantized_gemm_gfx1x_n_gt_5_falls_back(monkeypatch): monkeypatch.setattr("vllm.platforms.rocm.on_gfx1x", lambda: True) monkeypatch.setattr("vllm.platforms.rocm.on_gfx9", lambda: False) monkeypatch.setattr("vllm.platforms.rocm.on_gfx950", lambda: False) + monkeypatch.setattr("vllm.platforms.rocm.on_gfx1250", lambda: False) monkeypatch.setattr(utils, "num_compute_units", lambda: 120) wvsplitk_mock = MagicMock(side_effect=lambda w, x_view, _, __: x_view @ w.t()) @@ -76,6 +78,7 @@ def test_rocm_unquantized_gemm_gfx950_wvsplitkrc_path(monkeypatch): monkeypatch.setattr("vllm.platforms.rocm.on_gfx1x", lambda: False) monkeypatch.setattr("vllm.platforms.rocm.on_gfx9", lambda: False) monkeypatch.setattr("vllm.platforms.rocm.on_gfx950", lambda: True) + monkeypatch.setattr("vllm.platforms.rocm.on_gfx1250", lambda: True) monkeypatch.setattr(utils, "num_compute_units", lambda: 120) wvsplitkrc_mock = MagicMock(side_effect=lambda x_view, w, _, __: x_view @ w.t()) diff --git a/tests/models/multimodal/generation/test_granite_speech.py b/tests/models/multimodal/generation/test_granite_speech.py index 3019f5f22d4b..de4c0ad2327c 100644 --- a/tests/models/multimodal/generation/test_granite_speech.py +++ b/tests/models/multimodal/generation/test_granite_speech.py @@ -45,9 +45,9 @@ def vllm_to_hf_output( def granite_speech_attention_config(): """Return attention config for Granite Speech tests on ROCm.""" if current_platform.is_rocm(): - from vllm.platforms.rocm import on_mi3xx + from vllm.platforms.rocm import get_cdna_version - if on_mi3xx(): + if get_cdna_version(): return {"backend": "ROCM_AITER_FA"} return {"backend": "TRITON_ATTN"} return None diff --git a/tests/models/quantization/test_bitsandbytes.py b/tests/models/quantization/test_bitsandbytes.py index 03c19b0bf62a..06745ae697e9 100644 --- a/tests/models/quantization/test_bitsandbytes.py +++ b/tests/models/quantization/test_bitsandbytes.py @@ -21,11 +21,11 @@ from ..utils import check_embeddings_close, check_logprobs_close if current_platform.is_rocm(): - from vllm.platforms.rocm import on_gfx9 + from vllm.platforms.rocm import on_cdna pytestmark = pytest.mark.skipif( - on_gfx9(), - reason="bitsandbytes not supported on gfx9 (warp size 64 limitation)", + on_cdna(), + reason="bitsandbytes not supported on CDNA (warp size 64 limitation)", ) models_4bit_to_test = [ diff --git a/tests/renderers/test_completions.py b/tests/renderers/test_completions.py index 76e88f4213e0..d184eb8621ce 100644 --- a/tests/renderers/test_completions.py +++ b/tests/renderers/test_completions.py @@ -64,12 +64,14 @@ class DummyTokenizer: def __post_init__(self) -> None: self._captured_encode_kwargs: dict = {} + self._captured_text_len: int = 0 def decode(self, tokens: list[int]): return str(tokens) def encode(self, text: str, **kwargs): self._captured_encode_kwargs = kwargs + self._captured_text_len = len(text) in_length = len(text) truncation = kwargs.get("truncation") @@ -366,6 +368,85 @@ def test_tokens_input_with_needs_detokenization(self): assert results[0]["prompt_token_ids"] == tokens assert results[0]["prompt"] == "[1, 2, 3, 4]" + def test_explicit_side_tokenizer_unbounded(self): + renderer = _build_renderer(MockModelConfig()) + + prompts = renderer.render_prompts( + _preprocess_prompt(renderer.model_config, "x" * 500) + ) + results = renderer.tokenize_prompts( + prompts, + TokenizeParams( + max_total_tokens=100, + truncate_prompt_tokens=4, + truncation_side="left", + ), + ) + + assert len(results) == 1 + assert len(results[0]["prompt_token_ids"]) == 4 + + kwargs = renderer.tokenizer._captured_encode_kwargs + assert kwargs["truncation"] is False + + def test_explicit_side_left_text(self): + renderer = _build_renderer(MockModelConfig()) + + prompts = renderer.render_prompts( + _preprocess_prompt(renderer.model_config, "x" * 50) + ) + results = renderer.tokenize_prompts( + prompts, + TokenizeParams( + max_total_tokens=100, + truncate_prompt_tokens=5, + truncation_side="left", + ), + ) + + assert len(results) == 1 + assert len(results[0]["prompt_token_ids"]) == 5 + assert results[0]["prompt_token_ids"] == list(range(45, 50)) + + def test_explicit_side_right_text(self): + renderer = _build_renderer(MockModelConfig()) + + prompts = renderer.render_prompts( + _preprocess_prompt(renderer.model_config, "x" * 50) + ) + results = renderer.tokenize_prompts( + prompts, + TokenizeParams( + max_total_tokens=100, + truncate_prompt_tokens=5, + truncation_side="right", + ), + ) + + assert len(results) == 1 + assert len(results[0]["prompt_token_ids"]) == 5 + assert results[0]["prompt_token_ids"] == list(range(5)) + + def test_explicit_side_text_pretokenization_guard(self): + renderer = _build_renderer(MockModelConfig(), max_chars_per_token=1) + + prompts = renderer.render_prompts( + _preprocess_prompt(renderer.model_config, "x" * 500) + ) + results = renderer.tokenize_prompts( + prompts, + TokenizeParams( + max_total_tokens=100, + truncate_prompt_tokens=4, + truncation_side="left", + ), + ) + + assert len(results) == 1 + assert len(results[0]["prompt_token_ids"]) == 4 + + assert renderer.tokenizer._captured_text_len <= 100 + class TestRenderEmbedPrompt: def _create_test_embed_bytes(self, tensor: torch.Tensor) -> bytes: diff --git a/tests/v1/attention/test_rocm_attention_backends_selection.py b/tests/v1/attention/test_rocm_attention_backends_selection.py index 48c6de8f8bd7..f1bd319215ec 100644 --- a/tests/v1/attention/test_rocm_attention_backends_selection.py +++ b/tests/v1/attention/test_rocm_attention_backends_selection.py @@ -28,16 +28,9 @@ def mock_vllm_config(): @pytest.fixture -def mock_on_gfx9(): - """Mock gfx9 arch detection to return True.""" - with patch("vllm.platforms.rocm.on_gfx9", return_value=True): - yield - - -@pytest.fixture -def mock_on_mi3xx(): - """Mock mi3xx arch detection to return True.""" - with patch("vllm.platforms.rocm.on_mi3xx", return_value=True): +def mock_get_cdna_version(): + """Mock cdna version arch detection to return True.""" + with patch("vllm.platforms.rocm.get_cdna_version", return_value=3): yield @@ -111,8 +104,7 @@ def test_standard_attention_backend_selection( selected_backend, expected_backend_path, mock_vllm_config, - mock_on_gfx9, - mock_on_mi3xx, + mock_get_cdna_version, monkeypatch, ): """Test standard attention backend selection with various configurations.""" @@ -313,12 +305,12 @@ def test_mla_backend_selection( def test_aiter_fa_requires_mi3xx(mock_vllm_config): - """Test that ROCM_AITER_FA requires mi3xx architecture.""" + """Test that ROCM_AITER_FA requires CDNA3+ architecture.""" from vllm.platforms.rocm import RocmPlatform - # Mock on_mi3xx to return False (used by supports_compute_capability) + # Mock cdna version to return 1 (used by supports_compute_capability) with ( - patch("vllm.platforms.rocm.on_mi3xx", return_value=False), + patch("vllm.platforms.rocm.get_cdna_version", return_value=1), pytest.raises( ValueError, match="compute capability not supported", diff --git a/tests/v1/kv_connector/nixl_integration/run_mamba_prefix_cache_test.sh b/tests/v1/kv_connector/nixl_integration/run_mamba_prefix_cache_test.sh index c7e65972004a..a34f07edc974 100755 --- a/tests/v1/kv_connector/nixl_integration/run_mamba_prefix_cache_test.sh +++ b/tests/v1/kv_connector/nixl_integration/run_mamba_prefix_cache_test.sh @@ -51,6 +51,7 @@ vllm serve $MODEL \ --trust-remote-code \ --enable-prefix-caching \ --mamba-cache-mode all \ + --attention-backend FLASHINFER \ --kv-transfer-config "$KV_CONFIG" & # Start decode instance @@ -68,6 +69,7 @@ vllm serve $MODEL \ --trust-remote-code \ --enable-prefix-caching \ --mamba-cache-mode all \ + --attention-backend FLASHINFER \ --kv-transfer-config "$KV_CONFIG" & echo "Waiting for prefill instance on port $PREFILL_PORT..." diff --git a/tests/v1/kv_connector/unit/test_hidden_states_connector.py b/tests/v1/kv_connector/unit/test_hidden_states_connector.py new file mode 100644 index 000000000000..ffd648289134 --- /dev/null +++ b/tests/v1/kv_connector/unit/test_hidden_states_connector.py @@ -0,0 +1,126 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""CPU-only unit tests for ExampleHiddenStatesConnector KV-cache-group logic.""" + +from types import SimpleNamespace + +import pytest +import torch + +from vllm.distributed.kv_transfer.kv_connector.v1.example_hidden_states_connector import ( # noqa: E501 + ExampleHiddenStatesConnector, +) +from vllm.v1.core.kv_cache_utils import get_kv_cache_groups +from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, + HiddenStateCacheSpec, + KVCacheGroupSpec, + MLAAttentionSpec, + SlidingWindowMLASpec, +) + + +def _full(block_size: int) -> FullAttentionSpec: + return FullAttentionSpec( + block_size=block_size, num_kv_heads=8, head_size=128, dtype=torch.bfloat16 + ) + + +def _hidden(block_size: int) -> HiddenStateCacheSpec: + return HiddenStateCacheSpec( + block_size=block_size, num_kv_heads=6, head_size=2048, dtype=torch.bfloat16 + ) + + +def _config(*specs): + """Minimal stand-in exposing only ``kv_cache_groups`` (all the helpers read).""" + return SimpleNamespace( + kv_cache_groups=[ + KVCacheGroupSpec(layer_names=[f"layer.{i}"], kv_cache_spec=spec) + for i, spec in enumerate(specs) + ] + ) + + +# ---- _find_cache_kv_group_id ------------------------------------------------ + + +def test_find_group_id_none_config_returns_zero(): + assert ExampleHiddenStatesConnector._find_cache_kv_group_id(None) == 0 + + +def test_find_group_id_single_non_hidden_group_returns_zero(): + # Uniform (dense) model: one group, no HiddenStateCacheSpec -> group 0. + cfg = _config(_full(16)) + assert ExampleHiddenStatesConnector._find_cache_kv_group_id(cfg) == 0 + + +def test_find_group_id_locates_hidden_group_when_not_first(): + # Hybrid layout: the hidden-states group is not group 0. + cfg = _config(_full(528), _hidden(22), _full(528)) + assert ExampleHiddenStatesConnector._find_cache_kv_group_id(cfg) == 1 + + +def test_find_group_id_locates_hidden_group_last(): + cfg = _config(_full(528), _full(528), _hidden(22)) + assert ExampleHiddenStatesConnector._find_cache_kv_group_id(cfg) == 2 + + +def test_find_group_id_raises_when_no_hidden_group_and_multiple_groups(): + cfg = _config(_full(16), _full(16)) + with pytest.raises(ValueError, match="Could not uniquely identify"): + ExampleHiddenStatesConnector._find_cache_kv_group_id(cfg) + + +def test_find_group_id_raises_when_multiple_hidden_groups(): + cfg = _config(_hidden(22), _hidden(22)) + with pytest.raises(ValueError, match="Could not uniquely identify"): + ExampleHiddenStatesConnector._find_cache_kv_group_id(cfg) + + +# ---- _get_cache_block_size -------------------------------------------------- + + +def test_get_block_size_reads_hidden_group_spec_not_global(): + # Hidden group keeps block size 22; the global is bumped to 528 for hybrids. + vllm_config = SimpleNamespace(cache_config=SimpleNamespace(block_size=528)) + cfg = _config(_full(528), _hidden(22)) + block_size = ExampleHiddenStatesConnector._get_cache_block_size( + vllm_config, cfg, cache_kv_group_id=1 + ) + assert block_size == 22 + + +def test_get_block_size_falls_back_to_cache_config_when_no_kv_cache_config(): + vllm_config = SimpleNamespace(cache_config=SimpleNamespace(block_size=16)) + block_size = ExampleHiddenStatesConnector._get_cache_block_size( + vllm_config, None, cache_kv_group_id=0 + ) + assert block_size == 16 + + +# ---- MLA-verifier absorption ------------------------------------------------ + + +def test_find_group_id_errors_clearly_when_absorbed_by_mla_swa_verifier(): + # HiddenStateCacheSpec subclasses MLAAttentionSpec, so an MLA + sliding- + # window MLA verifier absorbs it into the MLA group instead of isolating it. + dt = torch.bfloat16 + spec = { + "layers.0.mla": MLAAttentionSpec( + block_size=64, num_kv_heads=1, head_size=576, dtype=dt + ), + "layers.1.swa": SlidingWindowMLASpec( + block_size=64, num_kv_heads=1, head_size=576, dtype=dt, sliding_window=512 + ), + "cache_only_layers.61": _hidden(64), + } + vllm_config = SimpleNamespace( + scheduler_config=SimpleNamespace(disable_hybrid_kv_cache_manager=False), + speculative_config=None, + ) + groups = get_kv_cache_groups(vllm_config, spec) + assert not any(isinstance(g.kv_cache_spec, HiddenStateCacheSpec) for g in groups) + cfg = SimpleNamespace(kv_cache_groups=groups) + with pytest.raises(ValueError, match="MLA verifiers are unsupported"): + ExampleHiddenStatesConnector._find_cache_kv_group_id(cfg) diff --git a/tests/v1/worker/test_mixed_warmup_gate.py b/tests/v1/worker/test_mixed_warmup_gate.py new file mode 100644 index 000000000000..6941df6773c1 --- /dev/null +++ b/tests/v1/worker/test_mixed_warmup_gate.py @@ -0,0 +1,30 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for the max_num_reqs gate on the V2 mixed prefill+decode warmup.""" + +from types import SimpleNamespace + +import pytest + +from vllm.v1.worker.gpu.warmup import run_mixed_prefill_decode_warmup + + +def _fail(*args, **kwargs): + raise AssertionError("worker callback must not run when warmup is skipped") + + +@pytest.mark.parametrize("max_num_reqs", [1, 0]) +def test_mixed_warmup_skipped_for_single_seq(max_num_reqs): + """A mixed prefill+decode step needs >=2 requests; with max_num_reqs < 2 + the warmup must be skipped without touching the worker callbacks.""" + runner = SimpleNamespace(is_pooling_model=False, max_num_reqs=max_num_reqs) + + assert ( + run_mixed_prefill_decode_warmup( + runner, + worker_execute_model=_fail, + worker_sample_tokens=_fail, + num_tokens=128, + ) + is False + ) diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index 4a8b4209d874..d439acf34abe 100644 --- a/vllm/_aiter_ops.py +++ b/vllm/_aiter_ops.py @@ -6,6 +6,7 @@ from typing import Protocol import torch +import torch.nn.functional as F from torch._ops import OpOverload from torch.distributed import ProcessGroup @@ -91,7 +92,7 @@ def should_custom_ar(self, inp: torch.Tensor) -> bool: ... def is_aiter_found_and_supported() -> bool: """Check if AITER library is available and platform supports it. - Checks: platform (ROCm), device arch (gfx9), and library existence. + Checks: platform (ROCm), device arch is CDNA 3 or better, and library existence. Does NOT check environment variables - that's handled by rocm_aiter_ops.is_enabled(). This function determines if aiter CAN be used, not if it SHOULD be used. @@ -105,9 +106,9 @@ def is_aiter_found_and_supported() -> bool: VLLM_ROCM_USE_AITER=0, while preventing unwanted JIT warnings for auto-discovery. """ if current_platform.is_rocm() and IS_AITER_FOUND: - from vllm.platforms.rocm import on_mi3xx + from vllm.platforms.rocm import get_cdna_version - return on_mi3xx() + return get_cdna_version() > 2 return False @@ -149,6 +150,141 @@ def wrapper(*args, **kwargs): return wrapper +def _mxfp4_moe_w1_triton_gemm_view(w1: torch.Tensor, w2: torch.Tensor) -> torch.Tensor: + """Map expert w1 to moe_gemm_a4w4 layout (stride(-2)==1 on [E, K, N]). + + vLLM stores w13 as [E, 2*inter, hidden/2] (K along dim2). aiter downcast uses the + same axis order with K-packed stride-1 on dim1; do not compose transpose+as_strided + (that overruns storage). Standard layout only needs transpose(1, 2). + """ + if w1.stride(-2) == 1: + return w1 + h = w2.shape[1] + e, a, b = w1.shape + if b * 2 == h: + return w1.transpose(1, 2) + if a * 2 == h: + kp, n_out = a, b + return w1.as_strided((e, kp, n_out), (kp * n_out, 1, kp)) + return w1.transpose(1, 2) + + +def _mxfp4_moe_w2_triton_gemm_view(w2: torch.Tensor) -> torch.Tensor: + """Map vLLM w2 [E, model_dim, inter/2] to moe_gemm_a4w4 layout.""" + if w2.stride(-2) == 1: + return w2 + return w2.transpose(1, 2) + + +def _silu_and_mul_glu(x: torch.Tensor) -> torch.Tensor: + """SiLU(gate) * up for g1u1 tensors with last dim 2 * inter (matches CK ``silu_and_mul``).""" + d = x.shape[-1] // 2 + gate, up = x[..., :d], x[..., d:] + return (F.silu(gate.to(torch.float32)) * up.to(torch.float32)).to(dtype=x.dtype) + + +def _mxfp4_moe_weight_as_uint8(w: torch.Tensor) -> torch.Tensor: + """Triton moe_gemm_a4w4 JIT cannot canonicalize torch float4 dtypes (KeyError).""" + if w.dtype == torch.uint8: + return w + fp4 = getattr(torch, "float4_e2m1fn_x2", None) + if fp4 is not None and w.dtype == fp4: + return w.view(torch.uint8) + name = getattr(w.dtype, "name", str(w.dtype)) + if "float4" in name or "e2m1" in name: + return w.view(torch.uint8) + return w + + +def _rocm_aiter_fused_moe_triton_gemm_a4w4( + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_weight: torch.Tensor, + topk_ids: torch.Tensor, + activation_method: int, + w1_scale: torch.Tensor | None, + w2_scale: torch.Tensor | None, + output_dtype: torch.dtype | None, +) -> torch.Tensor: + from aiter import ActivationType + from aiter.fused_moe import get_inter_dim + from aiter.ops.triton.moe.moe_op_gemm_a4w4 import moe_gemm_a4w4, mxfp4_quant + from aiter.ops.triton.moe.moe_routing.routing import routing + + if activation_method != int(ActivationType.Silu): + raise RuntimeError( + "Triton moe_gemm_a4w4 path supports only Silu/SwiGLU (ActivationType.Silu)." + ) + + device = hidden_states.device + out_dtype = output_dtype or hidden_states.dtype + # FP32 matmul outputs for numerical headroom before the second mxfp4_quant. + gemm_out_dt = torch.float32 + m, _ = topk_ids.shape + num_experts = w1.shape[0] + + logits = torch.full((m, num_experts), -1e9, device=device, dtype=torch.float32) + tid = topk_ids.long().clamp(min=0, max=num_experts - 1) + logits.scatter_( + 1, + tid, + torch.log(topk_weight.to(torch.float32).clamp(min=1e-20)), + ) + + rdata, gather_idx, scatter_idx = routing(logits, topk_ids.shape[1]) + gate_scal = rdata.gate_scal + + e, model_dim, inter_dim = get_inter_dim(w1.shape, w2.shape) + is_g1u1 = inter_dim != w1.shape[1] + if not is_g1u1: + raise RuntimeError( + "Triton moe_gemm_a4w4 path expects SwiGLU (w1 dim1 == 2 * inter)." + ) + + w1_v = _mxfp4_moe_weight_as_uint8(_mxfp4_moe_w1_triton_gemm_view(w1, w2)) + w2_v = _mxfp4_moe_weight_as_uint8(_mxfp4_moe_w2_triton_gemm_view(w2)) + + x_q, x_s = mxfp4_quant(hidden_states.to(out_dtype)) + # Raw W1@x per (token, expert), then CK-style silu_and_mul (not Triton ``apply_swiglu``). + # ``scatter_indx=None`` keeps expanded rows until after activation; higher VRAM than fused swiglu. + mid_raw = moe_gemm_a4w4( + x_q, + w1_v, + x_s, + w1_scale, + x_static_scale=None, + quant_static_scale=None, + bias=None, + routing_data=rdata, + gather_indx=gather_idx, + scatter_indx=None, + gammas=None, + swizzle_mx_scale=None, + out_dtype=gemm_out_dt, + apply_swiglu=False, + ) + mid = _silu_and_mul_glu(mid_raw) + mid_q, mid_s = mxfp4_quant(mid.to(out_dtype)) + out = moe_gemm_a4w4( + mid_q, + w2_v, + mid_s, + w2_scale, + x_static_scale=None, + quant_static_scale=None, + bias=None, + routing_data=rdata, + gather_indx=None, + scatter_indx=scatter_idx, + gammas=gate_scal, + swizzle_mx_scale=None, + out_dtype=gemm_out_dt, + apply_swiglu=False, + ) + return out.to(out_dtype) + + def _rocm_aiter_fused_moe_impl( hidden_states: torch.Tensor, w1: torch.Tensor, @@ -179,6 +315,30 @@ def _rocm_aiter_fused_moe_impl( activation = ActivationType(activation_method) quant_type = QuantType(quant_method) + m_tokens = hidden_states.shape[0] + use_triton_a4w4 = ( + envs.VLLM_ROCM_AITER_FUSED_MOE_TRITON_GEMM_A4W4 + and quant_type == QuantType.per_1x32 + and expert_mask is None + and not doweight_stage1 + and num_local_tokens is None + and a1_scale is None + and a2_scale is None + and m_tokens >= envs.VLLM_ROCM_AITER_TRITON_MOE_MIN_TOKENS + ) + if use_triton_a4w4: + return _rocm_aiter_fused_moe_triton_gemm_a4w4( + hidden_states, + w1, + w2, + topk_weight, + topk_ids, + activation_method, + w1_scale, + w2_scale, + output_dtype, + ) + extra_kwargs: dict = {} if gate_mode and rocm_aiter_ops.fused_moe_supports_gate_mode(): extra_kwargs["gate_mode"] = gate_mode @@ -1731,16 +1891,22 @@ def is_fp8bmm_enabled(cls) -> bool: @classmethod @if_aiter_supported def is_fp4bmm_enabled(cls) -> bool: - from vllm.platforms.rocm import on_gfx950 + from vllm.platforms.rocm import get_cdna_version - return cls._AITER_ENABLED and cls._FP4BMM_ENABLED and on_gfx950() + return ( + cls._AITER_ENABLED and cls._FP4BMM_ENABLED and get_cdna_version() == 4 + ) # TODO GFX1250: Swap to > 3 @classmethod @if_aiter_supported def is_linear_hipbmm_enabled(cls) -> bool: - from vllm.platforms.rocm import on_mi3xx + from vllm.platforms.rocm import get_cdna_version - return cls.is_linear_enabled() and on_mi3xx() and cls._LINEAR_HIPBMM_ENABLED + return ( + cls.is_linear_enabled() + and (get_cdna_version() > 2) + and cls._LINEAR_HIPBMM_ENABLED + ) @classmethod @if_aiter_supported @@ -2554,7 +2720,9 @@ def group_fp8_quant( @staticmethod def is_triton_gemm_w8a8_tuned(n: int, k: int) -> bool: - return (n, k) in [ + from vllm.platforms.rocm import on_gfx1250 + + return on_gfx1250() or (n, k) in [ (1024, 8192), (2112, 7168), (3072, 1536), diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index 02404a2f5173..7dcb890aecef 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -3785,20 +3785,6 @@ def _hadacore_transform_fake(x: torch.Tensor, inplace: bool) -> torch.Tensor: return torch.empty_like(x) if not inplace else x -if hasattr(torch.ops._C, "minimax_allreduce_rms"): - - @register_fake("_C::minimax_allreduce_rms") - def _minimax_allreduce_rms_fake( - input: torch.Tensor, - norm_weight: torch.Tensor, - workspace: torch.Tensor, - rank: int, - nranks: int, - eps: float, - ) -> torch.Tensor: - return torch.empty_like(input) - - if hasattr(torch.ops._C, "minimax_allreduce_rms_qk"): @register_fake("_C::minimax_allreduce_rms_qk") diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/example_hidden_states_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/example_hidden_states_connector.py index 7e6c95bf8fb1..299ff037ad2c 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/example_hidden_states_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/example_hidden_states_connector.py @@ -109,6 +109,49 @@ def prefer_cross_layer_blocks(self) -> bool: # Must be False so that drafter kv cache isn't merged with verifier's return False + @classmethod + def _find_cache_kv_group_id(cls, kv_cache_config: "KVCacheConfig | None") -> int: + """Index of the KV cache group holding the extracted hidden states. + + Located by spec type so it resolves on both scheduler and worker side. + """ + if kv_cache_config is None: + return 0 + + from vllm.v1.kv_cache_interface import HiddenStateCacheSpec + + groups = kv_cache_config.kv_cache_groups + group_ids = [ + gid + for gid, group in enumerate(groups) + if isinstance(group.kv_cache_spec, HiddenStateCacheSpec) + ] + if len(group_ids) == 1: + return group_ids[0] + if not group_ids and len(groups) == 1: + return 0 + raise ValueError( + "Could not uniquely identify the extract-hidden-states KV cache " + f"group among {len(groups)} groups; the hidden-states layer must be " + "isolated in its own group (MLA verifiers are unsupported)." + ) + + @staticmethod + def _get_cache_block_size( + vllm_config: "VllmConfig", + kv_cache_config: "KVCacheConfig | None", + cache_kv_group_id: int, + ) -> int: + """Block size of the hidden-states group, read from its own spec. + + cache_config.block_size is bumped to a common multiple for hybrid + verifiers; the page-aligned hidden-states group keeps a smaller one. + """ + if kv_cache_config is None: + return vllm_config.cache_config.block_size + cache_group = kv_cache_config.kv_cache_groups[cache_kv_group_id] + return cache_group.kv_cache_spec.block_size + def __init__( self, vllm_config: "VllmConfig", @@ -120,7 +163,12 @@ def __init__( role=role, kv_cache_config=kv_cache_config, ) - self._block_size = vllm_config.cache_config.block_size + # Read the hidden-states group and its block size from the group spec; + # cache_config.block_size is bumped (wrong) for hybrid verifiers. + self._cache_kv_group_id = self._find_cache_kv_group_id(kv_cache_config) + self._block_size = self._get_cache_block_size( + vllm_config, kv_cache_config, self._cache_kv_group_id + ) self._storage_path = self._kv_transfer_config.get_from_extra_config( "shared_storage_path", "/tmp" ) @@ -151,13 +199,6 @@ def __init__( # Worker-side state (set by register_kv_caches). self._kv_cache: torch.Tensor | None = None - # Identify which KV cache group holds the hidden-states layer. - self._hs_group_idx: int = 0 - if self._kv_cache_config is not None: - for i, group in enumerate(self._kv_cache_config.kv_cache_groups): - if any("cache_only_layers" in n for n in group.layer_names): - self._hs_group_idx = i - break # Only TP rank 0 writes hidden states to disk; other TP ranks no-op. # Set in register_kv_caches (after distributed init). self._is_tp_rank_zero: bool = True @@ -267,12 +308,14 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): ) self._kv_cache = kv_caches[self.cache_layers[0]] - # Find the KV cache group index for hidden states - if self._kv_cache_config is not None: - for i, group in enumerate(self._kv_cache_config.kv_cache_groups): - if self.cache_layers[0] in group.layer_names: - self._hs_group_idx = i - break + # Block size must match the indexed buffer, else reads hit the wrong + # slots. Raise (not assert) so the check survives `python -O`. + if self._block_size != self._kv_cache.shape[1]: + raise ValueError( + f"Hidden-states block-size mismatch: derived {self._block_size} " + f"but buffer block size is {self._kv_cache.shape[1]}; read slots " + "would be wrong (likely a hybrid block-size resolution bug)." + ) @staticmethod def _write_tensors( @@ -543,7 +586,7 @@ def request_finished_all_groups( request: "Request", block_ids: tuple[list[int], ...], ) -> tuple[bool, dict[str, Any] | None]: - return self.request_finished(request, block_ids[self._hs_group_idx]) + return self.request_finished(request, block_ids[self._cache_kv_group_id]) @classmethod def get_required_kvcache_layout(cls, vllm_config: "VllmConfig") -> str | None: diff --git a/vllm/envs.py b/vllm/envs.py index 1f94be8ac5fa..f880858b711d 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -134,6 +134,7 @@ VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION: bool = False VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS: bool = False VLLM_ROCM_USE_AITER_TRITON_GEMM: bool = True + VLLM_ROCM_AITER_FUSED_MOE_TRITON_GEMM_A4W4: bool = False VLLM_ROCM_USE_SKINNY_GEMM: bool = True VLLM_ROCM_FP8_PADDING: bool = True VLLM_ROCM_MOE_PADDING: bool = True @@ -1227,6 +1228,11 @@ def _resolve_rust_frontend_path() -> str | None: "VLLM_ROCM_USE_AITER_TRITON_GEMM": lambda: ( os.getenv("VLLM_ROCM_USE_AITER_TRITON_GEMM", "True").lower() in ("true", "1") ), + # ROCm AITER fused MoE: use Triton moe_gemm_a4w4 (MXFP4) instead of aiter.fused_moe. + "VLLM_ROCM_AITER_FUSED_MOE_TRITON_GEMM_A4W4": lambda: ( + os.getenv("VLLM_ROCM_AITER_FUSED_MOE_TRITON_GEMM_A4W4", "False").lower() + in ("true", "1") + ), # use rocm skinny gemms "VLLM_ROCM_USE_SKINNY_GEMM": lambda: ( os.getenv("VLLM_ROCM_USE_SKINNY_GEMM", "True").lower() in ("true", "1") diff --git a/vllm/model_executor/kernels/linear/scaled_mm/pytorch.py b/vllm/model_executor/kernels/linear/scaled_mm/pytorch.py index 2b6d3ed73698..a04be2e097c0 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/pytorch.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/pytorch.py @@ -104,10 +104,10 @@ def is_supported( if not current_platform.is_rocm(): return False, "requires ROCm." - from vllm.platforms.rocm import on_mi3xx + from vllm.platforms.rocm import get_cdna_version - if not on_mi3xx(): - return False, "requires MI3xx." + if get_cdna_version() <= 2: + return False, "requires CDNA3+" if compute_capability is not None and compute_capability < 94: return False, "requires compute capability 94 and above." diff --git a/vllm/model_executor/kernels/linear/scaled_mm/rocm.py b/vllm/model_executor/kernels/linear/scaled_mm/rocm.py index 64bc5b6c8bbe..688338b31188 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/rocm.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/rocm.py @@ -79,10 +79,10 @@ def is_supported( if not current_platform.is_rocm(): return False, "requires ROCm." - from vllm.platforms.rocm import on_gfx12x, on_mi3xx + from vllm.platforms.rocm import get_cdna_version - if not (on_mi3xx() or on_gfx12x()): - return False, "requires MI3xx or gfx12x" + if get_cdna_version() <= 2: + return False, "requires CDNA capabilities greater than 2" if not envs.VLLM_ROCM_USE_SKINNY_GEMM: return False, "requires VLLM_ROCM_USE_SKINNY_GEMM to be enabled." diff --git a/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a8_moe.py b/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a8_moe.py index 7c3fe5831f3b..3c17b8c8b09e 100644 --- a/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a8_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a8_moe.py @@ -46,11 +46,32 @@ def aiter_triton_kernel_w4a8_moe_forward( and quant_config.use_mxfp4_w4a8 and rocm_aiter_ops.is_enabled() ) - from aiter.ops.triton.moe_routing.routing import routing as aiter_routing + from vllm.platforms.rocm import on_gfx1250 + + try: + from aiter.ops.triton.moe.moe_routing import routing as _routing_mod + except ImportError: + from aiter.ops.triton.moe_routing import routing as _routing_mod + + if on_gfx1250(): + _routing_mod.is_tdm_avail = lambda: False + aiter_routing = _routing_mod.routing routing_data, gather_idx, scatter_idx = aiter_routing( gating_output, topk, sm_first=not renormalize ) + + # gfx1250: aiter's in-kernel gather is numerically broken (validated on the + # FFM sim: do_gather=True -> maxrel ~2.4), so gather rows into expert-sorted + # order in torch and pass gather_indx=None. Per aiter's moe_gemm_torch, + # sorted row i reads source token gather_idx[i] // n_expts_act, so this + # reproduces the in-kernel gather exactly (manual gather -> maxrel ~5e-3). + # gfx950 keeps the (working) in-kernel gather. + if on_gfx1250(): + gather_src = gather_idx.to(torch.long) // topk + hidden_states = hidden_states[gather_src] + gather_idx = None + return triton_kernel_fused_mxfp4_w4a8_experts( None, hidden_states, @@ -130,6 +151,12 @@ def triton_kernel_fused_mxfp4_w4a8_experts( hidden_states, quant_config.w1_precision.flex_ctx.lhs_data.scale ) + # gfx1250 stores the MXFP4 weight scale unswizzled (StridedLayout, see + # mxfp4_utils._swizzle_mxfp4) because the gfx1250 moe_gemm_a8w4 reads a + # CDNA4-swizzled scale as garbage (validated on the FFM sim: CDNA4_SCALE -> + # maxrel ~7e4, plain/None -> ~6e-3); pass swizzle_mx_scale=None there. + # gfx950 uses the CDNA4 swizzle layout. + intermediate_cache1 = moe_gemm_a8w4( hidden_states, w1.storage.data, @@ -199,12 +226,16 @@ def activation_format() -> mk.FusedMoEActivationFormat: @staticmethod def _supports_current_device() -> bool: - # Requires AITER and GFX950 + # Requires AITER and a supported AMD arch. gfx950 (CDNA4) uses the + # in-kernel gather + CDNA4 scale swizzle; gfx1250 routes through the + # same moe_gemm_a8w4 kernel with a manual gather and unswizzled scales + # (see aiter_triton_kernel_w4a8_moe_forward / the swizzle handling in + # triton_kernel_fused_mxfp4_w4a8_experts). if not rocm_aiter_ops.is_enabled(): return False - from vllm.platforms.rocm import on_gfx950 + from vllm.platforms.rocm import on_gfx950, on_gfx1250 - return on_gfx950() + return on_gfx950() or on_gfx1250() @staticmethod def _supports_no_act_and_mul() -> bool: diff --git a/vllm/model_executor/layers/fused_moe/experts/fused_batched_moe.py b/vllm/model_executor/layers/fused_moe/experts/fused_batched_moe.py index 21bda8e173fd..c4863c7d623f 100644 --- a/vllm/model_executor/layers/fused_moe/experts/fused_batched_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/fused_batched_moe.py @@ -758,13 +758,13 @@ def _supports_quant_scheme( ) -> bool: p = current_platform if p.is_rocm(): - from vllm.platforms.rocm import on_gfx9 + from vllm.platforms.rocm import get_cdna_version - is_rocm_on_gfx9 = on_gfx9() + _rocm_support_fp8 = get_cdna_version() > 2 else: - is_rocm_on_gfx9 = False + _rocm_support_fp8 = False - device_supports_fp8 = is_rocm_on_gfx9 or ( + device_supports_fp8 = _rocm_support_fp8 or ( p.is_cuda() and p.has_device_capability((8, 9)) ) diff --git a/vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py b/vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py index 4b0a0b8ecad3..8b688697d167 100644 --- a/vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py @@ -1,9 +1,13 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from dataclasses import replace + import torch import vllm.model_executor.layers.fused_moe.modular_kernel as mk +from vllm import _custom_ops as ops +from vllm._aiter_ops import rocm_aiter_ops 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 ( @@ -47,13 +51,13 @@ def _triton_kernel_moe_supports_current_device() -> bool: # range was not validated. return cap is not None and (9, 0) <= (cap.major, cap.minor) < (11, 0) if p.is_rocm(): - from vllm.platforms.rocm import on_gfx1x, on_gfx9 + from vllm.platforms.rocm import on_gfx1x, on_gfx9, on_gfx1250 # gfx9 family: gfx90a (MI200), gfx942/gfx950 (MI3xx); # on_gfx9() already excludes gfx906/gfx908. # gfx1x family: gfx11xx (RDNA3/3.5) and gfx12xx (RDNA4); # on_gfx1x() excludes gfx10xx (RDNA1/RDNA2). - return on_gfx9() or on_gfx1x() + return on_gfx9() or on_gfx1x() or on_gfx1250() return False @@ -570,6 +574,44 @@ def triton_kernel_moe_forward( ) -> torch.Tensor: sm_first = not renormalize + # ROCm aiter-native MXFP4 fast path (mirrors ROCm/ATOM triton_kernel_moe_forward): + # native aiter routing on the raw logits + the two-GEMM aiter compute. Falls back + # to the vendored triton_kernels routing + matmul_ogs path below when unavailable + # or out of the ported scope (SiLU is FP8-only). + if ( + expert_map is None + and quant_config is not None + and rocm_aiter_ops.is_enabled() + and _aiter_act_quant_mode(quant_config) is not None + ): + _act_mode = _aiter_act_quant_mode(quant_config) + _aiter_ops = _import_aiter_moe_ops() + if _aiter_ops is not None and ( + activation == MoEActivation.SWIGLUOAI or _act_mode == _AITER_ACT_FP8 + ): + routing_data, gather_idx, scatter_idx = _aiter_ops["routing"]( + gating_output, topk, sm_first=sm_first + ) + swiglu_alpha, swiglu_limit = _aiter_swiglu_params(quant_config) + # Native aiter routing applies router weights in-kernel, so no post-scale. + return _aiter_moe_two_gemm( + _aiter_ops, + hidden_states, + w1, + w2, + routing_data, + gather_idx, + scatter_idx, + topk, + quant_config, + _act_mode, + activation, + swiglu_alpha, + swiglu_limit, + apply_router_weight_on_input, + None, + ) + # When no expert map is provided (no EP), call the fused `routing()` # kernel directly. It combines softmax, topk, bitmatrix packing, and # routing-metadata construction in a single launch, instead of the @@ -814,6 +856,503 @@ def make_routing_data( return routing_data, gather_indx, scatter_indx +def _maybe_downcast_bf16_hidden_to_fp8_mxfp4_unfused( + hidden_states: torch.Tensor, + quant_config: FusedMoEQuantConfig, + *, + gemm_num: int = 1, +) -> tuple[torch.Tensor, torch.Tensor | None]: + """BF16 activations -> FP8 before unfused ``matmul_ogs`` for MXFP4 w4a16 (ROCm). + + Uses ``aiter.ops.triton.moe.quant_moe.downcast_to_static_fp8`` with a scalar scale: + for GEMM1: ``w1_precision.flex_ctx.lhs_data.scale`` or ``a1_scale`` when + single-element; for GEMM2: ``w2_precision.flex_ctx.lhs_data.scale`` or + ``a2_scale``; otherwise ``max(|x|) / 448`` for e4m3 range (matches AITER static + FP8 path). + + Returns: + (activations, lhs_scale_or_none): ``lhs_scale`` is the scalar tensor used for + FP8 downcast when this path runs; otherwise ``None`` (activations unchanged). + """ + if gemm_num not in (1, 2): + raise ValueError(f"gemm_num must be 1 or 2, got {gemm_num}") + if not quant_config.use_mxfp4_w4a16: + return hidden_states, None + try: + from aiter.ops.triton.moe.quant_moe import downcast_to_static_fp8 + except ImportError: + return hidden_states, None + if not rocm_aiter_ops.is_enabled(): + return hidden_states, None + + qp = quant_config.w1_precision if gemm_num == 1 else quant_config.w2_precision + a_scale = quant_config.a1_scale if gemm_num == 1 else quant_config.a2_scale + + lhs_scale = None + if qp is not None and qp.flex_ctx.lhs_data.scale is not None: + s = qp.flex_ctx.lhs_data.scale + if s.numel() == 1: + lhs_scale = s + if lhs_scale is None and a_scale is not None: + s = a_scale + if s.numel() == 1: + lhs_scale = s + if lhs_scale is None: + amax = hidden_states.abs().max().clamp(min=1e-12) + lhs_scale = (amax / 448.0).to(dtype=torch.float32) + + return downcast_to_static_fp8(hidden_states, lhs_scale), lhs_scale + + +def _mxfp4_w4a8_unpadded_dims( + moe_cfg: FusedMoEConfig, +) -> tuple[int | None, int | None, int | None, int | None]: + """Logical unpadded GEMM sizes for padded MXFP4 checkpoints (e.g. GFX950 swizzle).""" + if ( + moe_cfg.intermediate_size_per_partition_unpadded is None + or moe_cfg.hidden_dim_unpadded is None + ): + return None, None, None, None + unpadded_n_w1 = moe_cfg.intermediate_size_per_partition_unpadded * 2 + unpadded_k_w1 = moe_cfg.hidden_dim_unpadded + unpadded_n_w2 = moe_cfg.hidden_dim_unpadded + unpadded_k_w2 = moe_cfg.intermediate_size_per_partition_unpadded + return unpadded_n_w1, unpadded_k_w1, unpadded_n_w2, unpadded_k_w2 + + +def _mxfp4_w4a8_resolve_lhs_scale( + quant_config: FusedMoEQuantConfig, + *, + gemm_num: int, + activations: torch.Tensor, +) -> torch.Tensor: + """Scalar FP8 LHS scale for ``downcast_to_static_fp8`` / ``moe_gemm_a8w4``. + + Prefer calibrated ``flex_ctx.lhs_data.scale`` or ``a{1,2}_scale``; otherwise + ``max(|activations|) / 448`` (e4m3), matching + ``_maybe_downcast_bf16_hidden_to_fp8_mxfp4_unfused``. + """ + if gemm_num not in (1, 2): + raise ValueError(f"gemm_num must be 1 or 2, got {gemm_num}") + qp = quant_config.w1_precision if gemm_num == 1 else quant_config.w2_precision + a_scale = quant_config.a1_scale if gemm_num == 1 else quant_config.a2_scale + lhs_scale = None + if qp is not None and qp.flex_ctx is not None: + ld = qp.flex_ctx.lhs_data + if ld is not None and ld.scale is not None and ld.scale.numel() == 1: + lhs_scale = ld.scale + if lhs_scale is None and a_scale is not None: + s = a_scale + if s.numel() == 1: + lhs_scale = s + if lhs_scale is None: + amax = activations.abs().max().clamp(min=1e-12) + lhs_scale = (amax / 448.0).to(dtype=torch.float32) + return lhs_scale + + +def _precision_config_with_mxfp4_unfused_fp8_lhs_scale( + precision_config, + lhs_scale: torch.Tensor | None, +): + """Ensure ``matmul_ogs`` sees the same LHS FP8 scale as ``downcast_to_static_fp8``.""" + if precision_config is None or lhs_scale is None: + return precision_config + from triton_kernels.numerics import InFlexData + + new_flex = replace( + precision_config.flex_ctx, + lhs_data=InFlexData(scale=lhs_scale), + ) + return replace(precision_config, flex_ctx=new_flex) + + +# --------------------------------------------------------------------------- +# aiter `.moe.*` MXFP4 MoE path (ported from ROCm/ATOM +# atom/model_ops/fused_moe_triton.py). +# +# Scope (agreed): the reference's *complete* paths only -- +# * SwiGLU (interleaved, gpt-oss): FP8 -> a8w4, FP4/BF16 -> a16w4, with the +# activation fused into the GEMM (apply_swiglu=True, swiglu_add_residual). +# * SiLU (half/half, e.g. DeepSeek-V4): FP8 -> a8w4 via dynamic MXFP8 quant + +# fused_clamp_act_mul. +# The reference's SiLU FP4 (a4w4) / BF16 (a16w4) paths feed an *unactivated* +# buffer to GEMM2 and are therefore intentionally NOT ported -- they raise. +# --------------------------------------------------------------------------- + +# aiter MXFP4 activation-quant modes (mirror ATOM MoEActivationQuant). +_AITER_ACT_FP8 = "fp8" # a8w4: FP8 activations (use_mxfp4_w4a8) +_AITER_ACT_FP4 = "fp4" # a4w4: MXFP4 activations (use_mxfp4_w4a4) +_AITER_ACT_BF16 = "bf16" # a16w4: BF16 activations (use_mxfp4_w4a16) + + +def _aiter_raw(t): + """Unwrap a triton_kernels ``Tensor`` to its backing ``torch.Tensor``. + + aiter's ``.moe.*`` kernels take plain ``torch.Tensor`` weights/scales, while + vLLM stores MXFP4 weights as triton_kernels ``Tensor`` wrappers + (``.storage.data``). Plain tensors pass through unchanged. + """ + return t.storage.data if hasattr(t, "storage") else t + + +def _aiter_act_quant_mode(quant_config: FusedMoEQuantConfig) -> str | None: + """Activation-quant mode for aiter MXFP4 MoE, or ``None`` if not MXFP4.""" + if quant_config.use_mxfp4_w4a8: + return _AITER_ACT_FP8 + if quant_config.use_mxfp4_w4a4: + return _AITER_ACT_FP4 + if quant_config.use_mxfp4_w4a16: + return _AITER_ACT_BF16 + return None + + +def _aiter_swiglu_params(quant_config: FusedMoEQuantConfig) -> tuple[float, float]: + """(alpha, clamp limit) for the SwiGLU/SiLU activation; DeepSeek-V4 defaults.""" + alpha = quant_config.gemm1_alpha if quant_config.gemm1_alpha is not None else 1.0 + limit = ( + quant_config.gemm1_clamp_limit + if quant_config.gemm1_clamp_limit is not None + else 7.0 + ) + return alpha, limit + + +def _import_aiter_moe_ops() -> dict | None: + """Import the aiter ``.moe.*`` triton MoE kernels (ROCm/ATOM API), or ``None``. + + Mirrors ROCm/ATOM ``atom/model_ops/fused_moe_triton.py``. Returns a dict of + callables on success; ``None`` if aiter (or this kernel family) is absent, so + the caller can fall back to the vendored ``matmul_ogs`` path. + """ + try: + from aiter.ops.triton.fusions.fused_clamp_act_mul import fused_clamp_act_mul + from aiter.ops.triton.moe.moe_op_gemm_a8w4 import moe_gemm_a8w4 + from aiter.ops.triton.moe.moe_op_gemm_a16w4 import moe_gemm_a16w4 + from aiter.ops.triton.moe.moe_routing.routing import routing as aiter_routing + from aiter.ops.triton.moe.quant_moe import downcast_to_static_fp8 + from aiter.ops.triton.quant.quant import dynamic_mxfp8_quant + from aiter.ops.triton.utils._triton.arch_info import get_arch + except ImportError: + return None + return { + "moe_gemm_a8w4": moe_gemm_a8w4, + "moe_gemm_a16w4": moe_gemm_a16w4, + "routing": aiter_routing, + "downcast_to_static_fp8": downcast_to_static_fp8, + "dynamic_mxfp8_quant": dynamic_mxfp8_quant, + "fused_clamp_act_mul": fused_clamp_act_mul, + "get_arch": get_arch, + } + + +def _aiter_routing_from_topk( + aiter_ops: dict, + topk_ids: torch.Tensor, + topk_weights: torch.Tensor, + num_experts: int, + device: torch.device, +) -> tuple: + """Build aiter ``RoutingData`` from a precomputed topk selection. + + vLLM selects experts before the expert kernel runs, but aiter's MXFP4 GEMMs + need aiter routing structs (distinct from the vendored triton_kernels ones). + aiter ``routing`` always softmaxes its logits, so we scatter ``log(weight)`` + into a dense ``-inf`` logit matrix: softmax-over-selected reproduces both the + exact expert selection and ``weight / sum(weight)``. The per-token weight sum + (e.g. ``routed_scaling_factor``) divided out here is restored by the caller. + + Returns ``(routing_data, gather_idx, scatter_idx, topk_weights_f32)``. + """ + routing = aiter_ops["routing"] + M = topk_weights.shape[0] + topk = topk_ids.size(1) + tw = topk_weights.to(torch.float32) + logits = torch.full((M, num_experts), -1e30, device=device, dtype=torch.float32) + logits.scatter_( + 1, + topk_ids.long().clamp(min=0, max=num_experts - 1), + torch.log(tw.clamp(min=1e-20)), + ) + routing_data, gather_idx, scatter_idx = routing(logits, topk, sm_first=False) + return routing_data, gather_idx, scatter_idx, tw + + +def _aiter_moe_two_gemm( + aiter_ops: dict, + hidden_states: torch.Tensor, + w1, + w2, + routing_data, + gather_idx, + scatter_idx, + topk: int, + quant_config: FusedMoEQuantConfig, + act_mode: str, + activation: MoEActivation, + swiglu_alpha: float, + swiglu_limit: float, + apply_router_weight_on_input: bool, + moe_config: FusedMoEConfig | None, +) -> torch.Tensor: + """Two-GEMM aiter MXFP4 MoE compute (ported from ROCm/ATOM fused_moe_triton). + + vLLM/gfx1250 adaptations vs the reference: weights/scales are unwrapped from + triton_kernels ``Tensor`` (``.storage.data``); scales are passed unswizzled + with logical ``unpadded_N/K`` (the gfx1250 convention used by the validated + aiter path); and GEMM1's in-kernel gather -- numerically broken on gfx1250 -- + is replaced by an explicit torch gather there. + + Returns the (M, K) bf16 result *before* any router-weight post-scale; the + caller applies that when routing was reconstructed from a topk selection. + """ + moe_gemm_a8w4 = aiter_ops["moe_gemm_a8w4"] + moe_gemm_a16w4 = aiter_ops["moe_gemm_a16w4"] + dynamic_mxfp8_quant = aiter_ops["dynamic_mxfp8_quant"] + fused_clamp_act_mul = aiter_ops["fused_clamp_act_mul"] + downcast_to_static_fp8 = aiter_ops["downcast_to_static_fp8"] + get_arch = aiter_ops["get_arch"] + + assert quant_config.w1_precision is not None + assert quant_config.w2_precision is not None + w1_data = _aiter_raw(w1) + w2_data = _aiter_raw(w2) + w1_wscale = _aiter_raw(quant_config.w1_precision.weight_scale) + w2_wscale = _aiter_raw(quant_config.w2_precision.weight_scale) + w1_bias = quant_config.w1_bias + w2_bias = quant_config.w2_bias + + if moe_config is not None: + unpadded_n_w1, unpadded_k_w1, unpadded_n_w2, unpadded_k_w2 = ( + _mxfp4_w4a8_unpadded_dims(moe_config) + ) + else: + unpadded_n_w1 = unpadded_k_w1 = unpadded_n_w2 = unpadded_k_w2 = None + + arch = get_arch() + quant_dtype = torch.float8_e4m3fnuz if arch == "gfx942" else torch.float8_e4m3fn + + # gfx1250's in-kernel gather miscomputes (validated on the FFM sim), so gather + # rows into expert-sorted order in torch and pass gather_indx=None instead. Per + # aiter's moe_gemm_torch, sorted row i reads source token gather_idx[i] // topk. + if arch == "gfx1250": + gather_src = gather_idx.to(torch.long) // topk + gemm1_input = hidden_states[gather_src] + gemm1_gather_indx = None + else: + gemm1_input = hidden_states + gemm1_gather_indx = gather_idx + + gammas = routing_data.gate_scal if routing_data else None + g1_gammas = gammas if apply_router_weight_on_input else None + g2_gammas = None if apply_router_weight_on_input else gammas + + if activation == MoEActivation.SWIGLUOAI: + # Interleaved [gate, up] (gpt-oss): activation fused into the GEMM. + if act_mode == _AITER_ACT_FP8: + a13_scale = _mxfp4_w4a8_resolve_lhs_scale( + quant_config, gemm_num=1, activations=gemm1_input + ) + a2_scale = _mxfp4_w4a8_resolve_lhs_scale( + quant_config, gemm_num=2, activations=gemm1_input + ) + hidden_fp8 = downcast_to_static_fp8(gemm1_input, a13_scale) + interm = moe_gemm_a8w4( + hidden_fp8, + w1_data, + None, + w1_wscale, + a13_scale, + a2_scale, + w1_bias, + routing_data, + gather_indx=gemm1_gather_indx, + gammas=g1_gammas, + swizzle_mx_scale=None, + out_dtype=quant_dtype, + apply_swiglu=True, + alpha=swiglu_alpha, + limit=swiglu_limit, + swiglu_add_residual=True, + unpadded_N=unpadded_n_w1, + unpadded_K=unpadded_k_w1, + ) + out = moe_gemm_a8w4( + interm, + w2_data, + None, + w2_wscale, + a2_scale, + None, + w2_bias, + routing_data, + scatter_indx=scatter_idx, + gammas=g2_gammas, + swizzle_mx_scale=None, + unpadded_N=unpadded_n_w2, + unpadded_K=unpadded_k_w2, + ) + else: + # FP4 and BF16 activations both route through the a16w4 kernel. + interm = moe_gemm_a16w4( + gemm1_input, + w1_data, + None, + w1_wscale, + None, + None, + w1_bias, + routing_data, + gather_indx=gemm1_gather_indx, + gammas=g1_gammas, + swizzle_mx_scale=None, + apply_swiglu=True, + alpha=swiglu_alpha, + limit=swiglu_limit, + swiglu_add_residual=True, + unpadded_N=unpadded_n_w1, + unpadded_K=unpadded_k_w1, + ) + out = moe_gemm_a16w4( + interm, + w2_data, + None, + w2_wscale, + None, + None, + w2_bias, + routing_data, + scatter_indx=scatter_idx, + gammas=g2_gammas, + swizzle_mx_scale=None, + unpadded_N=unpadded_n_w2, + unpadded_K=unpadded_k_w2, + ) + return out + + # SiLU (concatenated [gate | up], half/half): manual activation between GEMMs. + # FP8 (a8w4) and BF16 (w4a16 served as a8w4) both run through moe_gemm_a8w4 by + # dynamically quantizing activations to FP8 (dynamic_mxfp8_quant below). Only + # FP4 (a4w4) SiLU is unsupported (ATOM reference a4w4 SiLU path is incomplete). + if act_mode == _AITER_ACT_FP4: + raise NotImplementedError( + "aiter SiLU MoE is unsupported for FP4 (a4w4) activations because the " + "ATOM reference a4w4 SiLU path is incomplete; use FP8 or BF16 (w4a16)." + ) + + hidden_q, a13_scale = dynamic_mxfp8_quant(gemm1_input, quant_dtype=quant_dtype) + raw_gate_up = moe_gemm_a8w4( + hidden_q, + w1_data, + a13_scale, + w1_wscale, + None, + None, + w1_bias, + routing_data, + gather_indx=gemm1_gather_indx, + gammas=g1_gammas, + swizzle_mx_scale=None, + out_dtype=torch.bfloat16, + apply_swiglu=False, + unpadded_N=unpadded_n_w1, + unpadded_K=unpadded_k_w1, + ) + if unpadded_n_w1 is not None: + raw_gate_up = raw_gate_up[:, :unpadded_n_w1] + # SiLU + clamp + (gate * up), fused with dynamic MXFP8 quant of the result. + interm_fp8, a2_scale = fused_clamp_act_mul( + raw_gate_up, + swiglu_limit=swiglu_limit, + activation="silu", + dtype_quant=quant_dtype, + scale_dtype_fmt="ue8m0", + quant_block_size=32, + ) + out = moe_gemm_a8w4( + interm_fp8, + w2_data, + a2_scale, + w2_wscale, + None, + None, + w2_bias, + routing_data, + scatter_indx=scatter_idx, + gammas=g2_gammas, + swizzle_mx_scale=None, + unpadded_N=unpadded_n_w2, + unpadded_K=unpadded_k_w2, + ) + return out + + +def _aiter_mxfp4_fused_experts( + output: torch.Tensor, + hidden_states: torch.Tensor, + w1, + w2, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + quant_config: FusedMoEQuantConfig, + activation: MoEActivation, + apply_router_weight_on_input: bool, + moe_config: FusedMoEConfig | None, +) -> torch.Tensor | None: + """aiter-native MXFP4 MoE for the modular expert path (ROCm fast path). + + Builds aiter routing from ``topk_ids``/``topk_weights``, runs the two-GEMM + aiter compute, restores the per-token router-weight sum that routing + reconstruction normalized away, and writes ``output``. + + Returns ``output`` on success, or ``None`` (caller falls back to ``matmul_ogs``) + when aiter is unavailable, the weights are not MXFP4, or the requested + activation/quant combination is out of the ported scope. + """ + if not rocm_aiter_ops.is_enabled(): + return None + act_mode = _aiter_act_quant_mode(quant_config) + if act_mode is None: + return None + if quant_config.w1_precision is None or quant_config.w2_precision is None: + return None + # SiLU served via a8w4 for FP8 and BF16 (dynamic FP8 quant); only FP4 (a4w4) + # SiLU is out of scope (reference a4w4 SiLU path incomplete). This keeps w4a16 + # SiLU on the aiter-native fast path instead of the vendored-routing fallback. + if activation != MoEActivation.SWIGLUOAI and act_mode == _AITER_ACT_FP4: + return None + aiter_ops = _import_aiter_moe_ops() + if aiter_ops is None: + return None + + num_experts = w1.shape[0] + routing_data, gather_idx, scatter_idx, tw = _aiter_routing_from_topk( + aiter_ops, topk_ids, topk_weights, num_experts, hidden_states.device + ) + swiglu_alpha, swiglu_limit = _aiter_swiglu_params(quant_config) + out = _aiter_moe_two_gemm( + aiter_ops, + hidden_states, + w1, + w2, + routing_data, + gather_idx, + scatter_idx, + topk_ids.size(1), + quant_config, + act_mode, + activation, + swiglu_alpha, + swiglu_limit, + apply_router_weight_on_input, + moe_config, + ) + # Restore the per-token weight sum that routing reconstruction normalized out. + out = out.to(output.dtype) * tw.sum(dim=1, keepdim=True) + output.copy_(out.reshape(output.shape)) + return output + + @triton.jit def _masked_topk_sum_kernel( inp_ptr, # (M, topk, K) contiguous @@ -1024,9 +1563,49 @@ def apply( self.quant_config: FusedMoEQuantConfig = FUSED_MOE_UNQUANTIZED_CONFIG if expert_map is not None: - # Preserve -1 (invalid / non-local slots, e.g. from EP dispatch): - # make_routing_data treats -1 as the skip sentinel. - topk_ids = remap_topk_to_local(topk_ids, expert_map) + topk_ids = expert_map[topk_ids] + + # ROCm aiter-native MXFP4 fast path (bypasses the gfx1250-incompatible + # matmul_ogs unswizzle). Returns None -> fall through to matmul_ogs. + if ( + not apply_router_weight_on_input + and rocm_aiter_ops.is_enabled() + and _aiter_mxfp4_fused_experts( + output, + hidden_states, + w1, + w2, + topk_weights, + topk_ids, + self.quant_config, + activation, + apply_router_weight_on_input, + getattr(self, "moe_config", None), + ) + is not None + ): + return + + # ROCm aiter-native MXFP4 fast path (bypasses the gfx1250-incompatible + # matmul_ogs unswizzle). Returns None -> fall through to matmul_ogs. + if ( + not apply_router_weight_on_input + and rocm_aiter_ops.is_enabled() + and _aiter_mxfp4_fused_experts( + output, + hidden_states, + w1, + w2, + topk_weights, + topk_ids, + self.quant_config, + activation, + apply_router_weight_on_input, + getattr(self, "moe_config", None), + ) + is not None + ): + return local_num_experts = w1.shape[0] if global_num_experts == -1: @@ -1110,7 +1689,7 @@ def activation( alpha = ( quant_config.gemm1_alpha if quant_config.gemm1_alpha is not None - else 1.702 + else 1.0 # SILU (DeepSeek-V4); was 1.702 gpt-oss default ) limit = ( quant_config.gemm1_clamp_limit @@ -1143,6 +1722,44 @@ def activation( else: super().activation(activation, output, input) + def _try_apply_aiter_w4a8( + self, + output: torch.Tensor, + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + quant_config: FusedMoEQuantConfig, + activation: MoEActivation, + apply_router_weight_on_input: bool, + ) -> torch.Tensor | None: + """ROCm aiter-native MXFP4 MoE fast path; ``None`` -> caller uses matmul_ogs. + + The vendored ``triton_kernels`` ``matmul_ogs`` MXFP4 path calls + ``unswizzle_mx_scale_cdna4``, which has no gfx1250 scale-layout variant and + fails to compile on gfx1250. aiter's ``.moe.*`` gluon kernels run on gfx1250, + so we route the MoE through them. + + Delegates to :func:`_aiter_mxfp4_fused_experts` (ported from ROCm/ATOM + ``fused_moe_triton.py``): a8w4 SiLU via dynamic MXFP8 + ``fused_clamp_act_mul``, + and the fused-SwiGLU modes. Routing is built aiter-native from + ``topk_ids``/``topk_weights`` and the router-weight sum is restored afterward; + see that function for the routing/scaling details. + """ + return _aiter_mxfp4_fused_experts( + output, + hidden_states, + w1, + w2, + topk_weights, + topk_ids, + quant_config, + activation, + apply_router_weight_on_input, + self.moe_config, + ) + def apply( self, output: torch.Tensor, @@ -1203,6 +1820,45 @@ def apply( if global_num_experts == -1: global_num_experts = E + # gfx1250 fast path: aiter-native W4A8 moe_gemm_a8w4 (built straight from + # topk_ids/topk_weights), bypassing the gfx1250-incompatible matmul_ogs + # unswizzle. Output-side router weighting only (the post-scale below assumes + # it); LoRA still uses the matmul_ogs path further down. + if ( + ( + quant_config.use_mxfp4_w4a8 + or quant_config.use_mxfp4_w4a16 + or quant_config.use_mxfp4_w4a4 + ) + and not apply_router_weight_on_input + and self._lora_context is None + and rocm_aiter_ops.is_enabled() + ): + if ( + self._try_apply_aiter_w4a8( + output, + hidden_states, + w1, + w2, + topk_weights, + topk_ids, + quant_config, + activation, + apply_router_weight_on_input, + ) + is not None + ): + return + + # FP8 activations + AITER Triton moe_gemm_a8w4 (MXFP4 w4a8), then moe_sum — avoids + # matmul_ogs on ROCm. Without AITER the unfused path below uses matmul_ogs (and for + # MXFP4 w4a16, optional FP8 activations via _maybe_downcast_bf16_hidden_to_fp8_mxfp4_unfused). + # NOTE: the fused `triton_kernel_fused_mxfp4_w4a8_experts` helper is not + # present in this tree (it lives in the dsv4_455 WIP), so the no-LoRA + # fused branch was removed. The unfused AITER `moe_gemm_a8w4` path below + # (`use_aiter_unfused_a8w4`) handles MXFP4 w4a8 for both LoRA and non-LoRA + # and avoids the gfx1250-incompatible `matmul_ogs` unswizzle. + # Note that the output tensor might be in workspace13 intermediate_cache1 = _resize_cache(workspace2, (batch_dim, M * topk, N)) intermediate_cache3 = _resize_cache(workspace2, (batch_dim, M * topk, K)) @@ -1211,18 +1867,108 @@ def apply( gammas = routing_data.gate_scal if routing_data else None - matmul_ogs( - hidden_states, - w1, - quant_config.w1_bias, - routing_data, - gather_indx=gather_indx, - precision_config=quant_config.w1_precision, - gammas=gammas if apply_router_weight_on_input else None, - fused_activation=None, - y=intermediate_cache1, + hidden_bf16_for_lora = hidden_states + + _moe_gemm_a8w4_fn = None + _downcast_static_fp8_fn = None + # The weights are MXFP4 in both the w4a8 and w4a16 cases; aiter's + # moe_gemm_a8w4 is a W4A8 (FP8 act x MXFP4 weight) kernel, so we can also + # serve the w4a16 case by dynamically quantizing the bf16 activations to + # FP8 here and passing the scale. This routes the MoE through aiter's + # gfx1250 kernel instead of the gfx1250-incompatible matmul_ogs. + _mxfp4_weights = quant_config.use_mxfp4_w4a8 or quant_config.use_mxfp4_w4a16 + if _mxfp4_weights and rocm_aiter_ops.is_enabled(): + try: + from aiter.ops.triton.moe.moe_op_gemm_a8w4 import ( + moe_gemm_a8w4 as _moe_gemm_a8w4_fn, + ) + from aiter.ops.triton.moe.quant_moe import ( + downcast_to_static_fp8 as _downcast_static_fp8_fn, + ) + except ImportError: + pass + use_aiter_unfused_a8w4 = ( + _mxfp4_weights + and _moe_gemm_a8w4_fn is not None + and _downcast_static_fp8_fn is not None ) + if use_aiter_unfused_a8w4: + assert quant_config.w1_precision is not None + assert quant_config.w2_precision is not None + unpadded_n_w1, unpadded_k_w1, unpadded_n_w2, unpadded_k_w2 = ( + _mxfp4_w4a8_unpadded_dims(self.moe_config) + ) + swiglu_alpha = ( + quant_config.gemm1_alpha + if quant_config.gemm1_alpha is not None + else 1.0 # SILU (DeepSeek-V4); was 1.702 gpt-oss default + ) + swiglu_limit = ( + quant_config.gemm1_clamp_limit + if quant_config.gemm1_clamp_limit is not None + else 7.0 + ) + gemm1_lhs_scale = _mxfp4_w4a8_resolve_lhs_scale( + quant_config, + gemm_num=1, + activations=hidden_bf16_for_lora, + ) + # GEMM1 passes W2's static scale through to the kernel; resolve it before + # GEMM2 activations exist (fallback uses token hidden states for amax). + gemm2_lhs_scale_proxy = _mxfp4_w4a8_resolve_lhs_scale( + quant_config, + gemm_num=2, + activations=hidden_bf16_for_lora, + ) + hidden_fp8 = _downcast_static_fp8_fn( + hidden_bf16_for_lora, + gemm1_lhs_scale, + ) + gemm1_out = _moe_gemm_a8w4_fn( + hidden_fp8, + w1.storage.data, + None, + quant_config.w1_precision.weight_scale.storage.data, + gemm1_lhs_scale, + gemm2_lhs_scale_proxy, + quant_config.w1_bias, + routing_data, + gather_indx=gather_indx, + gammas=gammas if apply_router_weight_on_input else None, + swizzle_mx_scale="CDNA4_SCALE", + out_dtype=torch.bfloat16, + apply_swiglu=False, + alpha=swiglu_alpha, + limit=swiglu_limit, + unpadded_N=unpadded_n_w1, + unpadded_K=unpadded_k_w1, + ) + intermediate_cache1.copy_(gemm1_out.reshape(intermediate_cache1.shape)) + else: + hidden_for_gemm1, gemm1_lhs_scale_fp8 = ( + _maybe_downcast_bf16_hidden_to_fp8_mxfp4_unfused( + hidden_bf16_for_lora, + quant_config, + ) + ) + w1_precision_for_gemm = _precision_config_with_mxfp4_unfused_fp8_lhs_scale( + quant_config.w1_precision, + gemm1_lhs_scale_fp8, + ) + + matmul_ogs( + hidden_for_gemm1, + w1, + quant_config.w1_bias, + routing_data, + gather_indx=gather_indx, + precision_config=w1_precision_for_gemm, + gammas=gammas if apply_router_weight_on_input else None, + fused_activation=None, + y=intermediate_cache1, + ) + # w13 LoRA: gather the activation input from expert-sorted # intermediate_cache1, then add the LoRA delta in-place on that copy # before passing it to activation — exactly mirroring the old @@ -1243,7 +1989,7 @@ def apply( ) = self.apply_w13_lora( lora_context, y=act_input, - x=hidden_states, + x=hidden_bf16_for_lora, topk_ids=global_topk_ids, topk_weights=topk_weights, expert_map=expert_map, @@ -1264,16 +2010,56 @@ def apply( # Set n_expts_act to 1 to unfuse the sum so we can do it manually via moe_sum. routing_data.n_expts_act = 1 - matmul_ogs( - intermediate_cache2[gather_indx.src_indx], - w2, - quant_config.w2_bias, - routing_data, - scatter_indx=scatter_indx, - precision_config=quant_config.w2_precision, - gammas=None if apply_router_weight_on_input else gammas, - y=intermediate_cache3, - ) + gemm2_input_bf16 = intermediate_cache2[gather_indx.src_indx] + if use_aiter_unfused_a8w4: + gemm2_lhs_scale = _mxfp4_w4a8_resolve_lhs_scale( + quant_config, + gemm_num=2, + activations=gemm2_input_bf16, + ) + gemm2_fp8 = _downcast_static_fp8_fn( + gemm2_input_bf16, + gemm2_lhs_scale, + ) + gemm2_out = _moe_gemm_a8w4_fn( + gemm2_fp8, + w2.storage.data, + None, + quant_config.w2_precision.weight_scale.storage.data, + gemm2_lhs_scale, + None, + quant_config.w2_bias, + routing_data, + scatter_indx=scatter_indx, + gammas=None if apply_router_weight_on_input else gammas, + swizzle_mx_scale="CDNA4_SCALE", + unpadded_N=unpadded_n_w2, + unpadded_K=unpadded_k_w2, + ) + intermediate_cache3.copy_(gemm2_out.reshape(intermediate_cache3.shape)) + else: + gemm2_input, gemm2_lhs_scale_fp8 = ( + _maybe_downcast_bf16_hidden_to_fp8_mxfp4_unfused( + gemm2_input_bf16, + quant_config, + gemm_num=2, + ) + ) + w2_precision_for_gemm = _precision_config_with_mxfp4_unfused_fp8_lhs_scale( + quant_config.w2_precision, + gemm2_lhs_scale_fp8, + ) + + matmul_ogs( + gemm2_input, + w2, + quant_config.w2_bias, + routing_data, + scatter_indx=scatter_indx, + precision_config=w2_precision_for_gemm, + gammas=None if apply_router_weight_on_input else gammas, + y=intermediate_cache3, + ) # w2 LoRA: after matmul_ogs with scatter_indx, intermediate_cache3 is # in token-topk order, matching the (M, topk, K) layout add_lora_w2 expects. diff --git a/vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py.bak b/vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py.bak new file mode 100644 index 000000000000..189a5ca8fd9e --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py.bak @@ -0,0 +1,1548 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from dataclasses import replace + +import torch + +import vllm.model_executor.layers.fused_moe.modular_kernel as mk +from vllm import _custom_ops as ops +from vllm._aiter_ops import rocm_aiter_ops +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 ( + FUSED_MOE_UNQUANTIZED_CONFIG, + FusedMoEConfig, + FusedMoEParallelConfig, + FusedMoEQuantConfig, + RoutingMethodType, +) +from vllm.model_executor.layers.fused_moe.experts.lora_experts_mixin import ( + LoRAExpertsMixin, +) +from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import ( + TopKWeightAndReduceNoOP, +) +from vllm.model_executor.layers.fused_moe.utils import _resize_cache +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + QuantKey, + kMxfp4Static, +) +from vllm.platforms import current_platform +from vllm.triton_utils import tl, triton +from vllm.utils.import_utils import has_triton_kernels + +from ..utils import swiglu_limit_func + +logger = init_logger(__name__) + + +def _triton_kernel_moe_supports_current_device() -> bool: + # Shared device gate for the OAI Triton MoE expert classes. + # Platform-aware to avoid ROCm capability aliasing — cap (9, 0) + # matches both gfx90a (verified) and gfx906 (unverified), so we + # dispatch on gfx-string helpers instead of the cap tuple on ROCm. + p = current_platform + if p.is_cuda(): + cap = p.get_device_capability() + # Keep the original `(9, 0) <= cap < (11, 0)` window on + # CUDA (covers Hopper SM90 and Blackwell SM100, excludes + # SM120) — this PR is ROCm-scoped and the broader CUDA + # range was not validated. + return cap is not None and (9, 0) <= (cap.major, cap.minor) < (11, 0) + if p.is_rocm(): + from vllm.platforms.rocm import on_gfx1x, on_gfx9 + + # gfx9 family: gfx90a (MI200), gfx942/gfx950 (MI3xx); + # on_gfx9() already excludes gfx906/gfx908. + # gfx1x family: gfx11xx (RDNA3/3.5) and gfx12xx (RDNA4); + # on_gfx1x() excludes gfx10xx (RDNA1/RDNA2). + return on_gfx9() or on_gfx1x() + return False + + +def _patch_make_bitmatrix_metadata() -> None: + """Monkey-patch make_bitmatrix_metadata to support non-power-of-2 top_k. + + triton's tl.arange requires a power-of-2 range. The original kernel + computes BLOCK_SIZE = BLOCK_PER_TOK * TOKS_PER_ROW (= 32 * top_k). For + DeepSeek-V4 with top_k=6 this gives 192, which is not a power of 2 and + causes a compile error at the first forward pass. + + Fix: define a drop-in replacement kernel that accepts an extra constexpr + BLOCK_SIZE_PADDED (next power of 2 >= BLOCK_SIZE) and uses it for the + tl.arange call while keeping the actual BLOCK_SIZE as the stride between + thread-blocks so that all flat indices into NonzeroIndx stay correct. + Elements beyond BLOCK_SIZE are masked out (col_indx = 0xffff) and ignored. + + This function is called once at module load time and patches the function + inside the triton_kernels tensor module so that SparseMatrix.__post_init__ + picks up the fixed version transparently. + """ + import torch + import triton + import triton.language as tl + + try: + if current_platform.is_rocm(): + from triton_kernels.tensor_details import bitmatrix as _bm + from triton_kernels.tensor_details.bitmatrix import ( + BitmatrixMetadata, + _keyed_add, + cdiv, + ) + from triton_kernels.tensor_details.bitmatrix_details.sum_bitmatrix_rows import ( # noqa: E501 + sum_bitmatrix_rows, + ) + else: + from vllm.third_party.triton_kernels.tensor_details import ( + bitmatrix as _bm, + ) + from vllm.third_party.triton_kernels.tensor_details.bitmatrix import ( + BitmatrixMetadata, + _keyed_add, + cdiv, + ) + from vllm.third_party.triton_kernels.tensor_details.bitmatrix_details.sum_bitmatrix_rows import ( # noqa: E501 + sum_bitmatrix_rows, + ) + except ImportError: + return + + @triton.jit + def _stage2_pow2( + ColSortedIndx, + RowSortedIndx, + NonzeroIndx, + n_tokens, + ColPartialSum, + stride_pm, + stride_pn, + ColOffs, + TOKS_PER_ROW: tl.constexpr, + BLOCK_PER_TOK: tl.constexpr, + BLOCK_SIZE_PADDED: tl.constexpr, + ): + # Actual number of elements per block (may not be a power of 2). + BLOCK_SIZE: tl.constexpr = BLOCK_PER_TOK * TOKS_PER_ROW + tl.static_assert(BLOCK_SIZE_PADDED <= 32768) + if isinstance(n_tokens, tl.tensor) and n_tokens.dtype.is_ptr(): + n_tokens = tl.load(n_tokens) + nonzero_indx_size = n_tokens * TOKS_PER_ROW + pid_m = tl.program_id(0) + # Use BLOCK_SIZE_PADDED (a power of 2) for tl.arange, but stride by + # the actual BLOCK_SIZE so flat positions in NonzeroIndx are correct. + # Elements with offs_local >= BLOCK_SIZE have offs_global beyond the + # valid range, get col_indx = 0xffff, and are filtered by the mask + # below without producing any output. + offs_local = tl.arange(0, BLOCK_SIZE_PADDED) + offs_global = pid_m * BLOCK_SIZE + offs_local + mask = offs_global < nonzero_indx_size + col_indx = tl.load(NonzeroIndx + offs_global, mask=mask, other=-1).to(tl.uint32) + kv_pairs = ((col_indx << 16) | offs_local).to(tl.uint32) + kv_pairs = tl.sort(kv_pairs, 0) + col_indx = kv_pairs >> 16 + offs_global = pid_m * BLOCK_SIZE + (kv_pairs & 0xFFFF) + mask = col_indx != 0xFFFF + x = kv_pairs & 0xFFFF0000 | 0x00000001 + cols_and_inclusive_run_lengths = tl.associative_scan(x, 0, _keyed_add) + exclusive_run_lengths = (cols_and_inclusive_run_lengths - 1) & 0xFFFF + row_sorted_indx = tl.load( + ColPartialSum + pid_m * stride_pm + col_indx * stride_pn, mask=mask + ) + row_sorted_indx += tl.load(ColOffs + col_indx, mask=mask) + row_sorted_indx += exclusive_run_lengths + tl.store(RowSortedIndx + offs_global, row_sorted_indx, mask=mask) + tl.store(ColSortedIndx + row_sorted_indx, offs_global, mask=mask) + + def _make_bitmatrix_metadata_pow2_safe(nonzero_indx, bitmatrix): + assert nonzero_indx.ndim == 2 + PARTIAL_BLOCK_M = 32 + col_sum, col_partial_sum = sum_bitmatrix_rows( + bitmatrix, partials_block_size=PARTIAL_BLOCK_M + ) + device = bitmatrix.device + n_indx = nonzero_indx.numel() + n_cols = bitmatrix.shape[1] + col_offs = torch.empty(n_cols, dtype=torch.int32, device=device) + combined_indx = torch.empty(n_indx * 2, dtype=torch.int32, device=device) + col_sorted_indx = combined_indx[:n_indx] + row_sorted_indx = combined_indx[n_indx:] + MEMSET_BLOCK = 1024 + memset_grid = (cdiv(n_indx * 2, MEMSET_BLOCK) + n_cols + 1,) + _bm._bitmatrix_metadata_compute_stage1[memset_grid]( + combined_indx, + n_indx * 2, + -1, + MEMSET_BLOCK, + col_sum, + col_offs, + col_sum.shape[0], + col_partial_sum, + col_partial_sum.shape[0], + col_partial_sum.stride(0), + col_partial_sum.stride(1), + BLOCK_M=512, + BLOCK_N=512, + ) + toks_per_row = nonzero_indx.shape[-1] + block_size = PARTIAL_BLOCK_M * toks_per_row + # Next power of 2 >= block_size (required by tl.arange). + block_size_padded = 1 << (max(block_size, 1) - 1).bit_length() + compute_grid = (cdiv(bitmatrix.shape_max[0], PARTIAL_BLOCK_M),) + _stage2_pow2[compute_grid]( + col_sorted_indx, + row_sorted_indx, + nonzero_indx, + bitmatrix.shape[0], + col_partial_sum, + col_partial_sum.stride(0), + col_partial_sum.stride(1), + col_offs, + TOKS_PER_ROW=toks_per_row, + BLOCK_PER_TOK=PARTIAL_BLOCK_M, + BLOCK_SIZE_PADDED=block_size_padded, + ) + return BitmatrixMetadata( + col_sum=col_sum, + col_sorted_indx=col_sorted_indx, + row_sorted_indx=row_sorted_indx, + ) + + # The most reliable patch point: SparseMatrix.__post_init__ looks up + # make_bitmatrix_metadata via its own __globals__ dict (the tensor.py + # module dict). Patching through __globals__ works regardless of how + # sys.modules maps "triton_kernels.tensor" vs + # "vllm.third_party.triton_kernels.tensor". + from triton_kernels.tensor import SparseMatrix as _SparseMatrix + + _SparseMatrix.__post_init__.__globals__["make_bitmatrix_metadata"] = ( + _make_bitmatrix_metadata_pow2_safe + ) + # Also patch the bitmatrix module itself in case it is imported directly. + _bm.make_bitmatrix_metadata = _make_bitmatrix_metadata_pow2_safe + + +# Two API generations of triton_kernels are supported: +# - v3.5.1 (the version bundled with vLLM): exposes `routing()` and +# `routing_from_bitmatrix()` in triton_kernels.routing; the `Bitmatrix` +# constructor takes a `scratchpad` argument. +# - v3.6.0+: removes the `routing` module in favor of a `SparseMatrix` +# based path, and adds a `dtype=BIT` kwarg to `Bitmatrix`. Used only +# when the user has triton_kernels installed system-wide at v3.6.0+. +# +# `use_legacy_triton_kernels` selects between them at import time based on +# whether `SparseMatrix` is importable. +use_legacy_triton_kernels = False + +if has_triton_kernels(): + try: + import triton_kernels.swiglu + from triton_kernels.matmul_ogs import ( + FnSpecs, + FusedActivation, + GatherIndx, + RoutingData, + ScatterIndx, + matmul_ogs, + ) + from triton_kernels.tensor import ( + BIT, + Bitmatrix, + ) + + try: + from triton_kernels.tensor import ( + SparseMatrix, + make_ragged_tensor_metadata, + ) + except ImportError: + # TODO(mgoin): drop the v3.5.1 pin and remove this fallback once + # the gpt-oss perf regression in v3.6.0+ is resolved upstream. + # Tracking: https://github.com/triton-lang/triton/issues/9969 + use_legacy_triton_kernels = True + if not use_legacy_triton_kernels: + _patch_make_bitmatrix_metadata() + except (AttributeError, ImportError) as e: + logger.error( + "Failed to import Triton kernels. Please make sure your triton " + "version is compatible. Error: %s", + e, + ) + + +@triton.jit +def pack_bitmatrix( + bitmatrix, + topk_ids, + n_rows, # n_rows in bitmatrix / topk_ids + bm_cols: tl.constexpr, # n int32_t bitpacks in bitmatrix + n_expts_act, # num_topk + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, +): + """ + Packs topk_ids into a bitmatrix. + code reference: + https://github.com/triton-lang/triton/blob/dd1bbc52b34d202dfe5ffea1e04fb16166c5c04e/python/triton_kernels/bench/distributed.py#L264 + """ + pid_m = tl.program_id(0) + offsets_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offsets_k = tl.arange(0, BLOCK_SIZE_K) + offsets = offsets_m[:, None] * n_expts_act + offsets_k[None, :] + mask = (offsets_m < n_rows)[:, None] & (offsets_k < n_expts_act)[None, :] + indices = tl.load(topk_ids + offsets, mask=mask, other=-1) + valid = indices >= 0 + div = indices // 32 + rem = indices % 32 + one = tl.cast(1, tl.uint32) + + # Iterate through all the relevant bitmatrix columns. + for i in range(bm_cols): + # When BLOCK_SIZE_K=32, offs is just the column index. + offs = tl.arange(0, BLOCK_SIZE_K // 32) + i * (BLOCK_SIZE_K // 32) + # All topks that need to go into this column has the correct bit set. + # Other bits are 0. x is a 2D tensor. + # Guard with `valid` to prevent negative indices from producing + # spurious bits (on HIP, -1 // 32 == 0 and 1 << (-1 % 32) sets + # bit 31). + x = tl.where( + valid[:, :, None] & (div[:, :, None] == offs[None, None, :]), + (one << rem)[:, :, None], + 0, + ) + # Reduce x to get a single int32_t bitpack. + y = tl.reduce_or(x, axis=1) + bitmatrix_ptrs = bitmatrix + offsets_m[:, None] * bm_cols + offs[None, :] + tl.store(bitmatrix_ptrs, y, mask=offsets_m[:, None] < n_rows) + + +def triton_kernel_moe_forward( + hidden_states: torch.Tensor, + w1, # Tensor or triton_kernels.Tensor + w2, # Tensor or triton_kernels.Tensor + gating_output: torch.Tensor, + topk: int, + renormalize: bool, + activation: MoEActivation = MoEActivation.SWIGLUOAI, + quant_config: FusedMoEQuantConfig | None = None, + apply_router_weight_on_input: bool = False, + global_num_experts: int = -1, + expert_map: torch.Tensor | None = None, + unpadded_N_w1=None, + unpadded_K_w1=None, + unpadded_N_w2=None, + unpadded_K_w2=None, +) -> torch.Tensor: + sm_first = not renormalize + + # When no expert map is provided (no EP), call the fused `routing()` + # kernel directly. It combines softmax, topk, bitmatrix packing, and + # routing-metadata construction in a single launch, instead of the + # three separate kernels used by the generic path below. + # Only available in the legacy (v3.5.1) API; the v3.6.0+ path inlines + # equivalent logic via SparseMatrix in `make_routing_data`. + if use_legacy_triton_kernels and expert_map is None: + from triton_kernels.routing import routing as fused_routing + + routing_data, gather_idx, scatter_idx = fused_routing( + gating_output, topk, sm_first=sm_first + ) + effective_expert_map = None + effective_global_num_experts = global_num_experts + else: + from triton_kernels.topk import topk as topk_fn + + logits = gating_output + if sm_first: + logits = torch.softmax(logits, dim=-1) + topk_result = topk_fn(logits, topk, apply_softmax=not sm_first) + # topk may return a tuple (vals, indx, bitmatrix) or a + # SparseMatrix depending on the triton_kernels version. + if isinstance(topk_result, tuple): + topk_weights, topk_ids_raw, _ = topk_result + else: + topk_weights = topk_result.vals + topk_ids_raw = topk_result.indx + + if expert_map is not None: + # topk_ids_raw contains global expert IDs - remap to local. + topk_ids = expert_map[topk_ids_raw.to(torch.long)] + local_num_experts = w1.shape[0] + routing_data, gather_idx, scatter_idx = make_routing_data( + topk_ids, topk_weights, local_num_experts + ) + # expert_map already applied; pass None downstream. + effective_expert_map = None + effective_global_num_experts = local_num_experts + else: + topk_ids = topk_ids_raw.to(torch.long) + routing_data, gather_idx, scatter_idx = make_routing_data( + topk_ids, topk_weights, gating_output.shape[-1] + ) + effective_expert_map = expert_map + effective_global_num_experts = global_num_experts + + output = torch.empty_like(hidden_states) + effective_quant_config = ( + quant_config if quant_config is not None else FUSED_MOE_UNQUANTIZED_CONFIG + ) + + return triton_kernel_fused_experts( + output, + hidden_states, + w1, + w2, + routing_data, + gather_idx, + scatter_idx, + topk=topk, + activation=activation, + quant_config=effective_quant_config, + apply_router_weight_on_input=apply_router_weight_on_input, + global_num_experts=effective_global_num_experts, + expert_map=effective_expert_map, + ) + + +# This is a triton implementation of the fused_experts function +def triton_kernel_fused_experts( + output_tensor: torch.Tensor, + hidden_states: torch.Tensor, + w1, # Tensor or triton_kernels.Tensor + w2, # Tensor or triton_kernels.Tensor + routing_data, # RoutingData + gather_indx, # GatherIndx + scatter_indx, # ScatterIndx + topk: int, + activation: MoEActivation = MoEActivation.SWIGLUOAI, + quant_config: FusedMoEQuantConfig | None = None, + swiglu_alpha: float = 1.702, + swiglu_limit: float = 7.0, + apply_router_weight_on_input: bool = False, + global_num_experts: int = -1, + expert_map: torch.Tensor | None = None, + intermediate_cache: torch.Tensor | None = None, + a1q_scale: torch.Tensor | None = None, +) -> torch.Tensor: + """Triton implementation of fused expert computation using OAI kernels.""" + assert activation == MoEActivation.SWIGLUOAI, ( + "Only SWIGLUOAI activation is supported" + ) + assert quant_config is not None + + # type check, uint8 means mxfp4 + assert hidden_states.dtype == torch.bfloat16 + assert quant_config.w1_bias is None or quant_config.w1_bias.dtype == torch.float32 + assert quant_config.w2_bias is None or quant_config.w2_bias.dtype == torch.float32 + + # Shape check, only check non-mxfp4 + assert hidden_states.ndim == 2 + assert hidden_states.shape[-1] == w1.shape[-2] + assert w2.shape[-1] == w1.shape[1] + + batch_dim = 1 + M, K = hidden_states.shape[-2:] + E, _, N = w1.shape + + if global_num_experts == -1: + global_num_experts = E + + if intermediate_cache is None: + intermediate_cache = torch.empty( + (batch_dim, M * topk, N // 2), + device=hidden_states.device, + dtype=hidden_states.dtype, + ) + + # Add batch_dim to output buffer because matmul_ogs expects 3D output + intermediate_cache = _resize_cache( + intermediate_cache, (batch_dim, M * topk, N // 2) + ) + output_tensor = _resize_cache(output_tensor, (batch_dim, M, K)) + + act = ( + FusedActivation( + FnSpecs( + "swiglu", + triton_kernels.swiglu.swiglu_fn, + ("alpha", "limit"), + reduction_n=2, + ), + (swiglu_alpha, swiglu_limit), + ) + if not use_legacy_triton_kernels + else FusedActivation( + FnSpecs("swiglu", triton_kernels.swiglu.swiglu_fn, ("alpha", "limit")), + (swiglu_alpha, swiglu_limit), + 2, + ) + ) + gammas = routing_data.gate_scal if routing_data else None + + matmul_ogs( + hidden_states, + w1, + quant_config.w1_bias, + routing_data, + gather_indx=gather_indx, + precision_config=quant_config.w1_precision, + gammas=gammas if apply_router_weight_on_input else None, + fused_activation=act, + y=intermediate_cache, + ) + + matmul_ogs( + intermediate_cache.view(M * topk, N // 2), + w2, + quant_config.w2_bias, + routing_data, + scatter_indx=scatter_indx, + precision_config=quant_config.w2_precision, + gammas=None if apply_router_weight_on_input else gammas, + y=output_tensor, + ) + output_tensor = output_tensor.view(M, K) + return output_tensor + + +def make_routing_data( + topk_ids: torch.Tensor, + topk_weights: torch.Tensor, + num_local_experts: int, +) -> tuple["RoutingData", torch.Tensor, torch.Tensor]: + topk_ids = topk_ids.to(torch.int16) + topk_weights = topk_weights.to(torch.bfloat16) + + n_rows, num_topk = topk_ids.size() + + BLOCK_SIZE_M = 512 + BLOCK_SIZE_K = 32 + + bm_cols = triton.cdiv(num_local_experts, BLOCK_SIZE_K) # n_bitpacks + bitmatrix = torch.zeros( + (n_rows, bm_cols), dtype=torch.uint32, device=topk_ids.device + ) + + grid = (triton.cdiv(n_rows, BLOCK_SIZE_M),) + pack_bitmatrix[grid]( + bitmatrix, + topk_ids, + n_rows, + bm_cols, + num_topk, + BLOCK_SIZE_M=BLOCK_SIZE_M, + BLOCK_SIZE_K=BLOCK_SIZE_K, + ) + + bitmatrix_shape = [n_rows, bm_cols * 32] + bitmatrix_shape_max = [n_rows, None] + bitmatrix = ( + Bitmatrix( + bitmatrix, dtype=BIT, shape=bitmatrix_shape, shape_max=bitmatrix_shape_max + ) + if not use_legacy_triton_kernels + else Bitmatrix( + bitmatrix, + shape=bitmatrix_shape, + shape_max=bitmatrix_shape_max, + scratchpad=None, + ) + ) + + # matmul_ogs expects invalid topk_weights to be -1s + topk_weights = torch.where(topk_ids == -1, -1.0, topk_weights) + + if use_legacy_triton_kernels: + from triton_kernels.routing import routing_from_bitmatrix + + return routing_from_bitmatrix( + bitmatrix, topk_weights, topk_ids, num_local_experts, num_topk + ) + + sparse_logits = SparseMatrix(indx=topk_ids, vals=topk_weights, mask=bitmatrix) + dispatch_indx = sparse_logits.mask_metadata.row_sorted_indx + combine_indx = sparse_logits.mask_metadata.col_sorted_indx + ragged_batch_metadata = make_ragged_tensor_metadata( + sparse_logits.mask_metadata.col_sum, + dispatch_indx.shape[0], + ) + gate_scal = sparse_logits.vals.flatten()[combine_indx] + routing_data = RoutingData( + gate_scal, + ragged_batch_metadata.block_sizes, + num_local_experts, + num_topk, + ragged_batch_metadata, + ) + gather_indx = GatherIndx(combine_indx, dispatch_indx) + scatter_indx = ScatterIndx(dispatch_indx, combine_indx) + return routing_data, gather_indx, scatter_indx + + +def _maybe_downcast_bf16_hidden_to_fp8_mxfp4_unfused( + hidden_states: torch.Tensor, + quant_config: FusedMoEQuantConfig, + *, + gemm_num: int = 1, +) -> tuple[torch.Tensor, torch.Tensor | None]: + """BF16 activations -> FP8 before unfused ``matmul_ogs`` for MXFP4 w4a16 (ROCm). + + Uses ``aiter.ops.triton.quant_moe.downcast_to_static_fp8`` with a scalar scale: + for GEMM1: ``w1_precision.flex_ctx.lhs_data.scale`` or ``a1_scale`` when + single-element; for GEMM2: ``w2_precision.flex_ctx.lhs_data.scale`` or + ``a2_scale``; otherwise ``max(|x|) / 448`` for e4m3 range (matches AITER static + FP8 path). + + Returns: + (activations, lhs_scale_or_none): ``lhs_scale`` is the scalar tensor used for + FP8 downcast when this path runs; otherwise ``None`` (activations unchanged). + """ + if gemm_num not in (1, 2): + raise ValueError(f"gemm_num must be 1 or 2, got {gemm_num}") + if not quant_config.use_mxfp4_w4a16: + return hidden_states, None + try: + from aiter.ops.triton.quant_moe import downcast_to_static_fp8 + except ImportError: + return hidden_states, None + if not rocm_aiter_ops.is_enabled(): + return hidden_states, None + + qp = quant_config.w1_precision if gemm_num == 1 else quant_config.w2_precision + a_scale = quant_config.a1_scale if gemm_num == 1 else quant_config.a2_scale + + lhs_scale = None + if qp is not None and qp.flex_ctx.lhs_data.scale is not None: + s = qp.flex_ctx.lhs_data.scale + if s.numel() == 1: + lhs_scale = s + if lhs_scale is None and a_scale is not None: + s = a_scale + if s.numel() == 1: + lhs_scale = s + if lhs_scale is None: + amax = hidden_states.abs().max().clamp(min=1e-12) + lhs_scale = (amax / 448.0).to(dtype=torch.float32) + + return downcast_to_static_fp8(hidden_states, lhs_scale), lhs_scale + + +def _mxfp4_w4a8_unpadded_dims( + moe_cfg: FusedMoEConfig, +) -> tuple[int | None, int | None, int | None, int | None]: + """Logical unpadded GEMM sizes for padded MXFP4 checkpoints (e.g. GFX950 swizzle).""" + if ( + moe_cfg.intermediate_size_per_partition_unpadded is None + or moe_cfg.hidden_dim_unpadded is None + ): + return None, None, None, None + unpadded_n_w1 = moe_cfg.intermediate_size_per_partition_unpadded * 2 + unpadded_k_w1 = moe_cfg.hidden_dim_unpadded + unpadded_n_w2 = moe_cfg.hidden_dim_unpadded + unpadded_k_w2 = moe_cfg.intermediate_size_per_partition_unpadded + return unpadded_n_w1, unpadded_k_w1, unpadded_n_w2, unpadded_k_w2 + + +def _mxfp4_w4a8_resolve_lhs_scale( + quant_config: FusedMoEQuantConfig, + *, + gemm_num: int, + activations: torch.Tensor, +) -> torch.Tensor: + """Scalar FP8 LHS scale for ``downcast_to_static_fp8`` / ``moe_gemm_a8w4``. + + Prefer calibrated ``flex_ctx.lhs_data.scale`` or ``a{1,2}_scale``; otherwise + ``max(|activations|) / 448`` (e4m3), matching + ``_maybe_downcast_bf16_hidden_to_fp8_mxfp4_unfused``. + """ + if gemm_num not in (1, 2): + raise ValueError(f"gemm_num must be 1 or 2, got {gemm_num}") + qp = quant_config.w1_precision if gemm_num == 1 else quant_config.w2_precision + a_scale = quant_config.a1_scale if gemm_num == 1 else quant_config.a2_scale + lhs_scale = None + if qp is not None and qp.flex_ctx is not None: + ld = qp.flex_ctx.lhs_data + if ld is not None and ld.scale is not None and ld.scale.numel() == 1: + lhs_scale = ld.scale + if lhs_scale is None and a_scale is not None: + s = a_scale + if s.numel() == 1: + lhs_scale = s + if lhs_scale is None: + amax = activations.abs().max().clamp(min=1e-12) + lhs_scale = (amax / 448.0).to(dtype=torch.float32) + return lhs_scale + + +def _precision_config_with_mxfp4_unfused_fp8_lhs_scale( + precision_config, + lhs_scale: torch.Tensor | None, +): + """Ensure ``matmul_ogs`` sees the same LHS FP8 scale as ``downcast_to_static_fp8``.""" + if precision_config is None or lhs_scale is None: + return precision_config + from triton_kernels.numerics import InFlexData + + new_flex = replace( + precision_config.flex_ctx, + lhs_data=InFlexData(scale=lhs_scale), + ) + return replace(precision_config, flex_ctx=new_flex) + + +class BaseOAITritonExperts(mk.FusedMoEExpertsModular): + @property + def expects_unquantized_inputs(self) -> bool: + return True + + @staticmethod + def _supports_current_device() -> bool: + return _triton_kernel_moe_supports_current_device() and has_triton_kernels() + + @staticmethod + def _supports_no_act_and_mul() -> bool: + return False + + @staticmethod + def _supports_quant_scheme( + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + SUPPORTED_W_A = [ + (kMxfp4Static, None), + ] + return (weight_key, activation_key) in SUPPORTED_W_A + + @staticmethod + def _supports_activation(activation: MoEActivation) -> bool: + raise NotImplementedError + + @staticmethod + def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool: + return True + + def moe_problem_size( + self, + a1: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_ids: torch.Tensor, + ) -> tuple[int, int, int, int, int]: + """ + Extract the MoE problem size from the given tensor arguments: + - a: The hidden states, input to the MoE layer. + - w1: The first set of expert weights. + - w2: The second set of expert weights. + - topk_ids: The topk ids. + Note: extracting the problem shape from the weight and activation + tensors is not obvious. It needs to be done this way specifically + due to subtle issues with particular kernels, e.g. the int4 kernels + divide the trailing dimension by two, so it's not "correct" to + extract N or K from the trailing dimension of w1 or w2. Similarly, + some kernels transpose the weights, so this needs to be kept in mind. + Note: This implementation covers most cases. However, if experts + require a specialized implementation, like MarlinExperts, they are free + to override this function. + """ + assert len(w1.shape) == 3 and len(w2.shape) == 3 + E, _, N = w1.shape + K = a1.size(-1) + + assert a1.dim() == 2 + assert topk_ids.size(0) == a1.size(0), f"{topk_ids.size(0)} != {a1.size(0)}" + M = a1.size(0) + + assert topk_ids.dim() == 2 + topk = topk_ids.size(1) + + return E, M, N, K, topk + + def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: + # Weight application and reduction happens in the fused_experts kernel. + return TopKWeightAndReduceNoOP() + + def _make_routing_data( + self, + topk_ids: torch.Tensor, + topk_weights: torch.Tensor, + num_local_experts: int, + ) -> tuple["RoutingData", torch.Tensor, torch.Tensor]: + return make_routing_data(topk_ids, topk_weights, num_local_experts) + + +class OAITritonExperts(BaseOAITritonExperts): + """OAI Triton-based fused MoE expert implementation.""" + + @staticmethod + def _supports_activation(activation: MoEActivation) -> bool: + return activation == MoEActivation.SWIGLUOAI + + @staticmethod + def activation_format() -> mk.FusedMoEActivationFormat: + return mk.FusedMoEActivationFormat.Standard + + def workspace_shapes( + self, + M: int, + N: int, + K: int, + topk: int, + global_num_experts: int, + local_num_experts: int, + expert_tokens_meta: mk.ExpertTokensMetadata | None, + activation: MoEActivation, + ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]: + # workspace are allocated inside the kernel + activation_out_dim = self.adjust_N_for_activation(N, activation) + workspace1 = (0, 0) + workspace2 = (M * topk, activation_out_dim) + output = (M, K) + return (workspace1, workspace2, output) + + 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 self.quant_config is None: + self.quant_config: FusedMoEQuantConfig = FUSED_MOE_UNQUANTIZED_CONFIG + + if expert_map is not None: + topk_ids = expert_map[topk_ids] + + local_num_experts = w1.shape[0] + if global_num_experts == -1: + global_num_experts = local_num_experts + + routing_data, gather_indx, scatter_indx = self._make_routing_data( + topk_ids, topk_weights, local_num_experts + ) + + topk = topk_ids.size(1) + triton_kernel_fused_experts( + output, + hidden_states, + w1, + w2, + routing_data, + gather_indx, + scatter_indx, + topk=topk, + activation=activation, + quant_config=self.quant_config, + apply_router_weight_on_input=False, + global_num_experts=local_num_experts, + expert_map=None, # applied already + intermediate_cache=workspace2, + a1q_scale=a1q_scale, + ) + + +class UnfusedOAITritonExperts(LoRAExpertsMixin, BaseOAITritonExperts): + """ + A Triton based MoE expert class that operates on expert standard + format and explicitly keeps the activation and reduction (moe_sum) steps + unfused from the matmul_ogs kernel. This exposes injection points + for activation and moe_sum. + + One use case for it is to inject LoRA modules on the activation and moe_sum. + """ + + @staticmethod + def _supports_activation(activation: MoEActivation) -> bool: + return activation in [ + MoEActivation.SILU, + MoEActivation.GELU, + MoEActivation.SWIGLUOAI, + MoEActivation.SWIGLUSTEP, + MoEActivation.SWIGLUOAI_UNINTERLEAVE, + ] + + @staticmethod + def activation_format() -> mk.FusedMoEActivationFormat: + return mk.FusedMoEActivationFormat.Standard + + def workspace_shapes( + self, + M: int, + N: int, + K: int, + topk: int, + global_num_experts: int, + local_num_experts: int, + expert_tokens_meta: mk.ExpertTokensMetadata | None, + activation: MoEActivation, + ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]: + # workspace are allocated inside the kernel + activation_out_dim = self.adjust_N_for_activation(N, activation) + workspace1 = (M * topk, activation_out_dim) + workspace2 = (M * topk, max(N, K)) + output = (M, K) + return (workspace1, workspace2, output) + + def moe_sum(self, input: torch.Tensor, output: torch.Tensor): + ops.moe_sum(input, output) + + def activation( + self, + activation: MoEActivation, + output: torch.Tensor, + input: torch.Tensor, + **kwargs, + ) -> None: + quant_config = self.quant_config or FUSED_MOE_UNQUANTIZED_CONFIG + if activation == MoEActivation.SWIGLUOAI: + alpha = ( + quant_config.gemm1_alpha + if quant_config.gemm1_alpha is not None + else 1.702 + ) + limit = ( + quant_config.gemm1_clamp_limit + if quant_config.gemm1_clamp_limit is not None + else 7.0 + ) + torch.ops._C.swigluoai_and_mul(output, input, alpha, limit) + elif ( + activation == MoEActivation.SILU + and quant_config.gemm1_clamp_limit is not None + ): + swiglu_limit_func( + output, + input, + quant_config.gemm1_clamp_limit, + ) + elif activation == MoEActivation.SWIGLUOAI_UNINTERLEAVE: + assert quant_config.gemm1_clamp_limit is not None + alpha = ( + quant_config.gemm1_alpha + if quant_config.gemm1_alpha is not None + else 1.0 + ) + beta = ( + quant_config.gemm1_beta if quant_config.gemm1_beta is not None else 0.0 + ) + torch.ops._C.silu_and_mul_with_clamp( + output, input, quant_config.gemm1_clamp_limit, alpha, beta + ) + else: + super().activation(activation, output, input) + + def _try_apply_aiter_w4a8( + self, + output: torch.Tensor, + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + quant_config: FusedMoEQuantConfig, + apply_router_weight_on_input: bool, + ) -> torch.Tensor | None: + """ROCm/gfx1250 MXFP4 MoE via aiter's W4A8 ``moe_gemm_a8w4``. + + The vendored ``triton_kernels`` ``matmul_ogs`` MXFP4 path calls + ``unswizzle_mx_scale_cdna4``, which has no gfx1250 scale-layout variant and + fails to compile on gfx1250. aiter's gluon ``moe_gemm_a8w4`` (FP8 act x MXFP4 + weight) does run on gfx1250, so we route the MoE through it. + + Routing is built aiter-native directly from ``topk_ids``/``topk_weights`` so we + never touch the diverged vendored ``triton_kernels`` routing structs (whose + ``RoutingData``/``GatherIndx``/``ExptData`` field names and layout differ from + aiter's). This mirrors ``triton_kernel_fused_mxfp4_w4a8_experts`` + (aiter_mxfp4_w4a8_moe.py) but resolves the FP8 LHS scales dynamically because + DeepSeek-V4 is ``activation_scheme: dynamic`` (no calibrated static scale). + + aiter's ``routing`` always softmaxes the gate weights, so we feed + ``log(topk_weights)``: softmax-over-selected then recovers + ``topk_weights / sum(topk_weights)``. DeepSeek-V4 weights are + norm_topk_prob-normalized then scaled by ``routed_scaling_factor`` (so they sum + to that factor, not 1); we restore the exact weighting with a per-token output + rescale by the weight sum (== 1 for plain renormalize routing, a no-op there). + + Returns ``output`` on success, or ``None`` if aiter is unavailable (caller then + falls back to ``matmul_ogs``). + """ + try: + from aiter.ops.triton.moe.moe_routing.routing import ( + routing as aiter_routing, + ) + from aiter.ops.triton.moe_op_gemm_a8w4 import moe_gemm_a8w4 + from aiter.ops.triton.quant_moe import downcast_to_static_fp8 + except ImportError: + return None + + assert quant_config.w1_precision is not None + assert quant_config.w2_precision is not None + + M = hidden_states.shape[0] + num_experts = w1.shape[0] + topk = topk_ids.size(1) + + # Reconstruct dense gating logits from the sparse topk selection. log(weight) + # padded with a large-negative sentinel makes aiter's softmax reproduce both + # this exact expert selection and weight = topk_weights / per-token-sum. + tw = topk_weights.to(torch.float32) + logits = torch.full( + (M, num_experts), -1e30, device=hidden_states.device, dtype=torch.float32 + ) + logits.scatter_( + 1, + topk_ids.long().clamp(min=0, max=num_experts - 1), + torch.log(tw.clamp(min=1e-20)), + ) + logger.info( + "[aiter-w4a8] ENTER M=%d K=%d num_experts=%d topk=%d " + "w1=%s w2=%s w1_wscale=%s w2_wscale=%s", + M, + hidden_states.shape[1], + num_experts, + topk, + tuple(w1.storage.data.shape), + tuple(w2.storage.data.shape), + tuple(quant_config.w1_precision.weight_scale.storage.data.shape), + tuple(quant_config.w2_precision.weight_scale.storage.data.shape), + ) + routing_data, gather_idx, scatter_idx = aiter_routing( + logits, topk, sm_first=False + ) + gammas = routing_data.gate_scal + logger.info( + "[aiter-w4a8] routing OK block_m=%s gate_scal=%s gather=%s scatter=%s", + getattr(routing_data, "block_m", None), + tuple(gammas.shape), + tuple(gather_idx.shape), + tuple(scatter_idx.shape), + ) + + unpadded_n_w1, unpadded_k_w1, unpadded_n_w2, unpadded_k_w2 = ( + _mxfp4_w4a8_unpadded_dims(self.moe_config) + ) + swiglu_alpha = ( + quant_config.gemm1_alpha + if quant_config.gemm1_alpha is not None + else 1.702 + ) + swiglu_limit = ( + quant_config.gemm1_clamp_limit + if quant_config.gemm1_clamp_limit is not None + else 7.0 + ) + + # FP8 LHS scales (static if calibrated, else dynamic amax/448). GEMM1 emits FP8 + # directly (apply_swiglu=True, out_dtype fp8) quantized with GEMM2's LHS scale, + # so resolve GEMM2's scale up front (hidden states as the amax proxy) and reuse + # it as GEMM2's input scale so quantize/dequantize are consistent. + gemm1_lhs_scale = _mxfp4_w4a8_resolve_lhs_scale( + quant_config, gemm_num=1, activations=hidden_states + ) + gemm2_lhs_scale = _mxfp4_w4a8_resolve_lhs_scale( + quant_config, gemm_num=2, activations=hidden_states + ) + # aiter's in-kernel gather is numerically broken on gfx1250 (validated on the + # FFM sim: do_gather=True -> maxrel ~2.4), so gather rows into expert-sorted + # order in torch and pass gather_indx=None. Per aiter's moe_gemm_torch, sorted + # row i reads source token gather_idx[i] // n_expts_act, so this reproduces the + # in-kernel gather exactly (validated: manual gather -> maxrel ~5e-3). + gather_src = gather_idx.to(torch.long) // topk + hidden_sorted = hidden_states[gather_src] + hidden_fp8 = downcast_to_static_fp8(hidden_sorted, gemm1_lhs_scale) + logger.info( + "[aiter-w4a8] scales g1_lhs=%s g2_lhs=%s hidden_fp8=%s; calling gemm1 " + "(unpadded N/K w1=%s/%s w2=%s/%s, alpha=%.4f limit=%.2f)", + float(gemm1_lhs_scale), + float(gemm2_lhs_scale), + tuple(hidden_fp8.shape), + unpadded_n_w1, + unpadded_k_w1, + unpadded_n_w2, + unpadded_k_w2, + swiglu_alpha, + swiglu_limit, + ) + + intermediate_cache1 = moe_gemm_a8w4( + hidden_fp8, + w1.storage.data, + None, + quant_config.w1_precision.weight_scale.storage.data, + gemm1_lhs_scale, + gemm2_lhs_scale, + quant_config.w1_bias, + routing_data, + gather_indx=None, + gammas=gammas if apply_router_weight_on_input else None, + swizzle_mx_scale=None, + out_dtype=torch.float8_e4m3fn, + apply_swiglu=True, + alpha=swiglu_alpha, + limit=swiglu_limit, + unpadded_N=unpadded_n_w1, + unpadded_K=unpadded_k_w1, + ) + logger.info( + "[aiter-w4a8] gemm1 OK ic1=%s %s; calling gemm2", + tuple(intermediate_cache1.shape), + intermediate_cache1.dtype, + ) + intermediate_cache3 = moe_gemm_a8w4( + intermediate_cache1, + w2.storage.data, + None, + quant_config.w2_precision.weight_scale.storage.data, + gemm2_lhs_scale, + None, + quant_config.w2_bias, + routing_data, + scatter_indx=scatter_idx, + gammas=None if apply_router_weight_on_input else gammas, + swizzle_mx_scale=None, + unpadded_N=unpadded_n_w2, + unpadded_K=unpadded_k_w2, + ) + + # Restore the per-token weight sum (e.g. routed_scaling_factor) that the + # log+softmax normalization divided out. Output weighting (gammas on GEMM2) is + # linear in the gate, so a single post-scale is exact. + logger.info( + "[aiter-w4a8] gemm2 OK ic3=%s %s; rescaling + writing output=%s", + tuple(intermediate_cache3.shape), + intermediate_cache3.dtype, + tuple(output.shape), + ) + out = intermediate_cache3.to(output.dtype) * tw.sum(dim=1, keepdim=True) + output.copy_(out.reshape(output.shape)) + logger.info("[aiter-w4a8] DONE") + return output + + 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, + ): + # Use local variable to help mypy narrow the type after None check + quant_config = self.quant_config + if quant_config is None: + quant_config = FUSED_MOE_UNQUANTIZED_CONFIG + + global_topk_ids = topk_ids + if expert_map is not None: + topk_ids = expert_map[topk_ids] + + local_num_experts = w1.shape[0] + if global_num_experts == -1: + global_num_experts = local_num_experts + + routing_data, gather_indx, scatter_indx = self._make_routing_data( + topk_ids, topk_weights, local_num_experts + ) + + topk = topk_ids.size(1) + + # type check, uint8 means mxfp4 + assert hidden_states.dtype == torch.bfloat16 + assert ( + quant_config.w1_bias is None or quant_config.w1_bias.dtype == torch.float32 + ) + assert ( + quant_config.w2_bias is None or quant_config.w2_bias.dtype == torch.float32 + ) + + # Shape check, only check non-mxfp4 + assert hidden_states.ndim == 2 + assert hidden_states.shape[-1] == w1.shape[-2] + assert w2.shape[-1] == w1.shape[1] + + batch_dim = 1 + M, K = hidden_states.shape + E, _, N = w1.shape + + if global_num_experts == -1: + global_num_experts = E + + # gfx1250 fast path: aiter-native W4A8 moe_gemm_a8w4 (built straight from + # topk_ids/topk_weights), bypassing the gfx1250-incompatible matmul_ogs + # unswizzle. Output-side router weighting only (the post-scale below assumes + # it); LoRA still uses the matmul_ogs path further down. + if ( + (quant_config.use_mxfp4_w4a8 or quant_config.use_mxfp4_w4a16) + and not apply_router_weight_on_input + and self._lora_context is None + and rocm_aiter_ops.is_enabled() + ): + if ( + self._try_apply_aiter_w4a8( + output, + hidden_states, + w1, + w2, + topk_weights, + topk_ids, + quant_config, + apply_router_weight_on_input, + ) + is not None + ): + return + + # FP8 activations + AITER Triton moe_gemm_a8w4 (MXFP4 w4a8), then moe_sum — avoids + # matmul_ogs on ROCm. Without AITER the unfused path below uses matmul_ogs (and for + # MXFP4 w4a16, optional FP8 activations via _maybe_downcast_bf16_hidden_to_fp8_mxfp4_unfused). + # NOTE: the fused `triton_kernel_fused_mxfp4_w4a8_experts` helper is not + # present in this tree (it lives in the dsv4_455 WIP), so the no-LoRA + # fused branch was removed. The unfused AITER `moe_gemm_a8w4` path below + # (`use_aiter_unfused_a8w4`) handles MXFP4 w4a8 for both LoRA and non-LoRA + # and avoids the gfx1250-incompatible `matmul_ogs` unswizzle. + + # Note that the output tensor might be in workspace13 + intermediate_cache1 = _resize_cache(workspace2, (batch_dim, M * topk, N)) + intermediate_cache3 = _resize_cache(workspace2, (batch_dim, M * topk, K)) + activation_out_dim = self.adjust_N_for_activation(N, activation) + intermediate_cache2 = _resize_cache(workspace13, (M * topk, activation_out_dim)) + + gammas = routing_data.gate_scal if routing_data else None + + hidden_bf16_for_lora = hidden_states + + _moe_gemm_a8w4_fn = None + _downcast_static_fp8_fn = None + # The weights are MXFP4 in both the w4a8 and w4a16 cases; aiter's + # moe_gemm_a8w4 is a W4A8 (FP8 act x MXFP4 weight) kernel, so we can also + # serve the w4a16 case by dynamically quantizing the bf16 activations to + # FP8 here and passing the scale. This routes the MoE through aiter's + # gfx1250 kernel instead of the gfx1250-incompatible matmul_ogs. + _mxfp4_weights = quant_config.use_mxfp4_w4a8 or quant_config.use_mxfp4_w4a16 + if _mxfp4_weights and rocm_aiter_ops.is_enabled(): + try: + from aiter.ops.triton.moe_op_gemm_a8w4 import ( + moe_gemm_a8w4 as _moe_gemm_a8w4_fn, + ) + from aiter.ops.triton.quant_moe import ( + downcast_to_static_fp8 as _downcast_static_fp8_fn, + ) + except ImportError: + pass + use_aiter_unfused_a8w4 = ( + _mxfp4_weights + and _moe_gemm_a8w4_fn is not None + and _downcast_static_fp8_fn is not None + ) + + if use_aiter_unfused_a8w4: + assert quant_config.w1_precision is not None + assert quant_config.w2_precision is not None + unpadded_n_w1, unpadded_k_w1, unpadded_n_w2, unpadded_k_w2 = ( + _mxfp4_w4a8_unpadded_dims(self.moe_config) + ) + swiglu_alpha = ( + quant_config.gemm1_alpha + if quant_config.gemm1_alpha is not None + else 1.702 + ) + swiglu_limit = ( + quant_config.gemm1_clamp_limit + if quant_config.gemm1_clamp_limit is not None + else 7.0 + ) + gemm1_lhs_scale = _mxfp4_w4a8_resolve_lhs_scale( + quant_config, + gemm_num=1, + activations=hidden_bf16_for_lora, + ) + # GEMM1 passes W2's static scale through to the kernel; resolve it before + # GEMM2 activations exist (fallback uses token hidden states for amax). + gemm2_lhs_scale_proxy = _mxfp4_w4a8_resolve_lhs_scale( + quant_config, + gemm_num=2, + activations=hidden_bf16_for_lora, + ) + hidden_fp8 = _downcast_static_fp8_fn( + hidden_bf16_for_lora, + gemm1_lhs_scale, + ) + gemm1_out = _moe_gemm_a8w4_fn( + hidden_fp8, + w1.storage.data, + None, + quant_config.w1_precision.weight_scale.storage.data, + gemm1_lhs_scale, + gemm2_lhs_scale_proxy, + quant_config.w1_bias, + routing_data, + gather_indx=gather_indx, + gammas=gammas if apply_router_weight_on_input else None, + swizzle_mx_scale="CDNA4_SCALE", + out_dtype=torch.bfloat16, + apply_swiglu=False, + alpha=swiglu_alpha, + limit=swiglu_limit, + unpadded_N=unpadded_n_w1, + unpadded_K=unpadded_k_w1, + ) + intermediate_cache1.copy_(gemm1_out.reshape(intermediate_cache1.shape)) + else: + hidden_for_gemm1, gemm1_lhs_scale_fp8 = ( + _maybe_downcast_bf16_hidden_to_fp8_mxfp4_unfused( + hidden_bf16_for_lora, + quant_config, + ) + ) + w1_precision_for_gemm = _precision_config_with_mxfp4_unfused_fp8_lhs_scale( + quant_config.w1_precision, + gemm1_lhs_scale_fp8, + ) + + matmul_ogs( + hidden_for_gemm1, + w1, + quant_config.w1_bias, + routing_data, + gather_indx=gather_indx, + precision_config=w1_precision_for_gemm, + gammas=gammas if apply_router_weight_on_input else None, + fused_activation=None, + y=intermediate_cache1, + ) + + # w13 LoRA: gather the activation input from expert-sorted + # intermediate_cache1, then add the LoRA delta in-place on that copy + # before passing it to activation — exactly mirroring the old + # decorator approach which modified the gathered tensor in-place. + act_input = intermediate_cache1.view(-1, N)[gather_indx.dst_indx] + + sorted_token_ids_lora = None + expert_ids_lora = None + num_tokens_post_padded_lora = None + token_lora_mapping = None + lora_context = self._lora_context + if lora_context is not None: + ( + sorted_token_ids_lora, + expert_ids_lora, + num_tokens_post_padded_lora, + token_lora_mapping, + ) = self.apply_w13_lora( + lora_context, + y=act_input, + x=hidden_bf16_for_lora, + topk_ids=global_topk_ids, + topk_weights=topk_weights, + expert_map=expert_map, + w1=w1, + w2=w2, + num_tokens=M, + top_k_num=topk, + ) + + self.activation( + activation, + intermediate_cache2, + act_input, + ) + + # matmul_ogs grouped reduction fuses sum across multiple experts: + # y[dst_indx // n_expts_act, :] += x + # Set n_expts_act to 1 to unfuse the sum so we can do it manually via moe_sum. + routing_data.n_expts_act = 1 + + gemm2_input_bf16 = intermediate_cache2[gather_indx.src_indx] + if use_aiter_unfused_a8w4: + gemm2_lhs_scale = _mxfp4_w4a8_resolve_lhs_scale( + quant_config, + gemm_num=2, + activations=gemm2_input_bf16, + ) + gemm2_fp8 = _downcast_static_fp8_fn( + gemm2_input_bf16, + gemm2_lhs_scale, + ) + gemm2_out = _moe_gemm_a8w4_fn( + gemm2_fp8, + w2.storage.data, + None, + quant_config.w2_precision.weight_scale.storage.data, + gemm2_lhs_scale, + None, + quant_config.w2_bias, + routing_data, + scatter_indx=scatter_indx, + gammas=None if apply_router_weight_on_input else gammas, + swizzle_mx_scale="CDNA4_SCALE", + unpadded_N=unpadded_n_w2, + unpadded_K=unpadded_k_w2, + ) + intermediate_cache3.copy_(gemm2_out.reshape(intermediate_cache3.shape)) + else: + gemm2_input, gemm2_lhs_scale_fp8 = ( + _maybe_downcast_bf16_hidden_to_fp8_mxfp4_unfused( + gemm2_input_bf16, + quant_config, + gemm_num=2, + ) + ) + w2_precision_for_gemm = _precision_config_with_mxfp4_unfused_fp8_lhs_scale( + quant_config.w2_precision, + gemm2_lhs_scale_fp8, + ) + + matmul_ogs( + gemm2_input, + w2, + quant_config.w2_bias, + routing_data, + scatter_indx=scatter_indx, + precision_config=w2_precision_for_gemm, + gammas=None if apply_router_weight_on_input else gammas, + y=intermediate_cache3, + ) + + # w2 LoRA: after matmul_ogs with scatter_indx, intermediate_cache3 is + # in token-topk order, matching the (M, topk, K) layout add_lora_w2 expects. + if lora_context is not None: + self.apply_w2_lora( + lora_context, + y=intermediate_cache3.view(-1, topk, K), + x=intermediate_cache2, + topk_weights=topk_weights, + sorted_token_ids_lora=sorted_token_ids_lora, + expert_ids_lora=expert_ids_lora, + num_tokens_post_padded_lora=num_tokens_post_padded_lora, + token_lora_mapping=token_lora_mapping, + num_tokens=M, + w1=w1, + w2=w2, + top_k_num=topk, + ) + + self.moe_sum(intermediate_cache3.view(-1, topk, K), output) + + +class OAITritonMxfp4ExpertsMonolithic(mk.FusedMoEExpertsMonolithic): + """Monolithic Triton MXFP4 expert. Wraps triton_kernel_moe_forward().""" + + def __init__( + self, + moe_config: FusedMoEConfig, + quant_config: FusedMoEQuantConfig, + ): + super().__init__(moe_config, quant_config) + self.topk = moe_config.experts_per_token + self.renormalize = moe_config.routing_method in ( + RoutingMethodType.Renormalize, + RoutingMethodType.RenormalizeNaive, + ) + + @staticmethod + def activation_format() -> mk.FusedMoEActivationFormat: + return mk.FusedMoEActivationFormat.Standard + + @staticmethod + def _supports_current_device() -> bool: + return _triton_kernel_moe_supports_current_device() and has_triton_kernels() + + @staticmethod + def _supports_no_act_and_mul() -> bool: + return False + + @staticmethod + def _supports_quant_scheme( + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + SUPPORTED_W_A = [ + (kMxfp4Static, None), + ] + return (weight_key, activation_key) in SUPPORTED_W_A + + @staticmethod + def _supports_activation(activation: MoEActivation) -> bool: + return activation == MoEActivation.SWIGLUOAI + + @staticmethod + def _supports_parallel_config( + moe_parallel_config: FusedMoEParallelConfig, + ) -> bool: + return ( + not moe_parallel_config.use_all2all_kernels + and not moe_parallel_config.enable_eplb + and moe_parallel_config.dp_size <= 1 + ) + + @staticmethod + def _supports_routing_method( + routing_method: RoutingMethodType, + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + return routing_method in [ + RoutingMethodType.Renormalize, + RoutingMethodType.RenormalizeNaive, + ] + + @staticmethod + def _supports_router_logits_dtype( + router_logits_dtype: torch.dtype | None, + routing_method: RoutingMethodType, + ) -> bool: + return True + + @property + def expects_unquantized_inputs(self) -> bool: + return True + + def apply( + self, + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + router_logits: torch.Tensor, + activation: MoEActivation, + global_num_experts: int, + expert_map: torch.Tensor | None, + a1q_scale: torch.Tensor | None, + apply_router_weight_on_input: bool, + # grouped topk + fused topk bias parameters + num_expert_group: int | None = None, + e_score_correction_bias: torch.Tensor | None = None, + routed_scaling_factor: float | None = None, + topk_group: int | None = None, + ) -> torch.Tensor: + return triton_kernel_moe_forward( + hidden_states=hidden_states, + w1=w1, + w2=w2, + gating_output=router_logits, + topk=self.topk, + renormalize=self.renormalize, + global_num_experts=global_num_experts, + expert_map=expert_map, + quant_config=self.quant_config, + apply_router_weight_on_input=apply_router_weight_on_input, + ) diff --git a/vllm/model_executor/layers/quantization/quark/schemes/quark_w4a8_mxfp4_fp8.py b/vllm/model_executor/layers/quantization/quark/schemes/quark_w4a8_mxfp4_fp8.py index 29283c7bbda4..1ca7cbfb2fe5 100644 --- a/vllm/model_executor/layers/quantization/quark/schemes/quark_w4a8_mxfp4_fp8.py +++ b/vllm/model_executor/layers/quantization/quark/schemes/quark_w4a8_mxfp4_fp8.py @@ -64,9 +64,9 @@ def __init__( kernel_supported_gpu = False if current_platform.is_rocm(): - from vllm.platforms.rocm import on_gfx950 + from vllm.platforms.rocm import get_cdna_version - kernel_supported_gpu = on_gfx950() + kernel_supported_gpu = get_cdna_version() > 3 self.use_aiter_kernel = ( is_aiter_found_and_supported() diff --git a/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py b/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py index db88ba273cd2..5110bcb7be70 100644 --- a/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py +++ b/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py @@ -71,6 +71,12 @@ def _swizzle_mxfp4(quant_tensor, scale, num_warps=8): scale_layout = CDNA4MXScaleLayout else: + # gfx1250 (and other ROCm archs): the vendored triton_kernels has no + # gfx1250 scale-swizzle, and the gfx1250 aiter `moe_gemm_a8w4` kernel + # reads a CDNA4-swizzled scale as garbage (validated on the FFM sim: + # CDNA4_SCALE -> maxrel ~7e4, plain/None -> ~6e-3). Keep the scale + # unswizzled (StridedLayout) and pass swizzle_mx_scale=None in the + # W4A8 path (gpt_oss_triton_kernels_moe._try_apply_aiter_w4a8). scale_layout = StridedLayout else: value_layout, value_layout_opts = layout.make_default_matmul_mxfp4_w_layout( diff --git a/vllm/model_executor/layers/utils.py b/vllm/model_executor/layers/utils.py index 6ca42c0e7f0e..efe14377324f 100644 --- a/vllm/model_executor/layers/utils.py +++ b/vllm/model_executor/layers/utils.py @@ -122,7 +122,7 @@ def use_aiter_triton_gemm(n, m, k, dtype): def rocm_unquantized_gemm_impl( x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor | None = None ) -> torch.Tensor: - from vllm.platforms.rocm import on_gfx1x, on_gfx9, on_gfx950 + from vllm.platforms.rocm import on_gfx1250, on_gfx1x, on_gfx9, on_gfx950 n = x.numel() // x.size(-1) m = weight.shape[0] @@ -164,7 +164,13 @@ def rocm_unquantized_gemm_impl( if use_skinny_reduce_counting: return ops.wvSplitKrc(x, weight, cu_count, bias) - if use_aiter_triton_gemm(n, m, k, x.dtype): + # gfx1250's aiter gemm_a16w16 uses the gluon backend, which requires + # K % 256 == 0 (it walks K with fixed-size descriptors and won't pad a + # partial last tile). Some whitelisted shapes have K=2880 (e.g. gpt-oss-120b + # hidden), so skip aiter there and fall back to the torch GEMM path below. + if use_aiter_triton_gemm(n, m, k, x.dtype) and not ( + on_gfx1250() and k % 256 != 0 + ): from aiter.ops.triton.gemm_a16w16 import gemm_a16w16 return gemm_a16w16(x, weight, bias) @@ -172,6 +178,8 @@ def rocm_unquantized_gemm_impl( use_skinny = ( envs.VLLM_ROCM_USE_SKINNY_GEMM and (on_gfx9() or on_gfx1x()) + # build (gfx9/gfx11 ISA); fall back to torch GEMM there. + # TODO GFX1250: Include once skinny GEMM is supported on gfx1250 and x.dtype in [torch.float16, torch.bfloat16] and k % 8 == 0 ) diff --git a/vllm/model_executor/models/extract_hidden_states.py b/vllm/model_executor/models/extract_hidden_states.py index 8df4823b6973..87dba7b5c984 100644 --- a/vllm/model_executor/models/extract_hidden_states.py +++ b/vllm/model_executor/models/extract_hidden_states.py @@ -83,7 +83,10 @@ def basic_cache( kv_cache: torch.Tensor, # shape: [num_blocks, block_size, num_heads, head_size] slot_mapping: torch.Tensor, # shape: [seq_len] ): + # Padding slots are -1; redirect them to the null block (block 0, never + # allocated to a request) so the scatter stays branch-free and sync-free. block_size = kv_cache.shape[1] + slot_mapping = slot_mapping.clamp_min(0) kv_cache[slot_mapping // block_size, slot_mapping % block_size] = to_cache diff --git a/vllm/model_executor/models/minicpmv4_6.py b/vllm/model_executor/models/minicpmv4_6.py index 79ac79d709ac..6cdc3624151b 100644 --- a/vllm/model_executor/models/minicpmv4_6.py +++ b/vllm/model_executor/models/minicpmv4_6.py @@ -140,7 +140,7 @@ def get_video_prompt_texts( per_frame = image_start + video_token * source_tokens + image_end if grids[0] > 0 and grids[1] > 0 and patch_tokens > 0: slice_ph = slice_start + video_token * patch_tokens + slice_end - rows = [slice_ph * grids[0] for _ in range(grids[1])] + rows = [slice_ph * grids[1] for _ in range(grids[0])] per_frame += "\n".join(rows) body = per_frame * num_frames @@ -596,7 +596,7 @@ def get_slice_image_placeholder( if use_image_id: placeholder = f"{id_start}{image_idx}{id_end}" + placeholder - num_cols, num_rows = grids[0], grids[1] + num_rows, num_cols = grids[0], grids[1] if num_cols > 0 and num_rows > 0 and patch_tokens > 0: slice_ph = slice_start + image_token * patch_tokens + slice_end slices = [slice_ph * num_cols for _ in range(num_rows)] diff --git a/vllm/model_executor/warmup/flashinfer_sparse_mla_warmup.py b/vllm/model_executor/warmup/flashinfer_sparse_mla_warmup.py index 44be769e246e..80dfd5199614 100644 --- a/vllm/model_executor/warmup/flashinfer_sparse_mla_warmup.py +++ b/vllm/model_executor/warmup/flashinfer_sparse_mla_warmup.py @@ -130,7 +130,7 @@ def _run_flashinfer_sparse_mla_decode_autotune( with torch.inference_mode(): warmup_executed = True if is_leader: - if _uses_v2_model_runner(runner): + if _uses_v2_model_runner(runner) and runner.max_num_reqs >= 2: v2_runner = cast("V2GPUModelRunner", runner) warmup_executed = run_mixed_prefill_decode_warmup( v2_runner, @@ -144,7 +144,7 @@ def _run_flashinfer_sparse_mla_decode_autotune( with flashinfer_autotune(True, cache=str(cache_path)): runner._dummy_run(**dummy_run_kwargs) else: - if _uses_v2_model_runner(runner): + if _uses_v2_model_runner(runner) and runner.max_num_reqs >= 2: v2_runner = cast("V2GPUModelRunner", runner) warmup_executed = run_mixed_prefill_decode_warmup( v2_runner, @@ -236,7 +236,7 @@ def deepseek_v4_sparse_mla_attention_warmup(worker: "Worker") -> None: ) mixed_warmup_done = _deepseek_v4_sparse_mla_decode_autotune(worker, mixed_tokens) if not mixed_warmup_done: - if _uses_v2_model_runner(runner): + if _uses_v2_model_runner(runner) and runner.max_num_reqs >= 2: v2_runner = cast("V2GPUModelRunner", runner) run_mixed_prefill_decode_warmup( v2_runner, diff --git a/vllm/platforms/__init__.py b/vllm/platforms/__init__.py index ac536aff00c5..5f5bd78793e5 100644 --- a/vllm/platforms/__init__.py +++ b/vllm/platforms/__init__.py @@ -125,6 +125,20 @@ def rocm_platform_plugin() -> str | None: except Exception as e: logger.debug("ROCm platform is not available because: %s", str(e)) + # Fallback: amdsmi queries the kernel driver/sysfs, which does not see GPUs + # exposed only through the HSA layer (e.g. the FFM pre-silicon simulator). + # Trust torch's HIP runtime detection in that case. + if not is_rocm: + try: + import torch + + if torch.version.hip is not None and torch.cuda.device_count() > 0: + is_rocm = True + logger.debug("ROCm platform detected via torch.version.hip " + "(amdsmi unavailable).") + except Exception as e: + logger.debug("torch HIP ROCm fallback failed: %s", str(e)) + return "vllm.platforms.rocm.RocmPlatform" if is_rocm else None diff --git a/vllm/platforms/rocm.py b/vllm/platforms/rocm.py index 6c3a0fe96ecc..0884b874c60f 100644 --- a/vllm/platforms/rocm.py +++ b/vllm/platforms/rocm.py @@ -37,8 +37,14 @@ amdsmi_topo_get_link_type, amdsmi_topo_get_numa_node_number, ) + + _AMDSMI_AVAILABLE = True except ImportError as e: logger.warning("Failed to import from amdsmi with %r", e) + # amdsmi is unavailable on the FFM pre-silicon simulator (and would report + # the host's real GPUs, not the sim). amdsmi-decorated methods fall back to + # torch when this is False. + _AMDSMI_AVAILABLE = False try: import vllm._C # noqa: F401 @@ -154,6 +160,10 @@ def _sync_hip_cuda_env_vars(): def with_amdsmi_context(fn): @wraps(fn) def wrapper(*args, **kwargs): + if not _AMDSMI_AVAILABLE: + # amdsmi unavailable (e.g. FFM simulator): skip init/shutdown and let + # the wrapped function use its own non-amdsmi (torch) fallback. + return fn(*args, **kwargs) amdsmi_init() try: return fn(*args, **kwargs) @@ -194,11 +204,13 @@ def _get_gcn_arch() -> str: return _query_gcn_arch_from_amdsmi() except Exception as e: logger.debug("Failed to get GCN arch via amdsmi: %s", e) - logger.warning_once( - "Failed to get GCN arch via amdsmi, falling back to torch.cuda. " - "This will initialize CUDA and may cause " - "issues if CUDA_VISIBLE_DEVICES is not set yet." - ) + # NOTE: use logger.debug, not warning_once, here. This runs at module + # load while resolving the platform; warning_once imports + # vllm.distributed, which causes a circular import before + # current_platform is bound. This path is taken on the FFM simulator, + # where amdsmi is unavailable (and would report the host's real gfx950 + # cards rather than the simulated gfx1250) — torch.cuda below is correct. + logger.debug("Failed to get GCN arch via amdsmi, falling back to torch.cuda.") # Ultimate fallback: use torch.cuda (will initialize CUDA) return torch.cuda.get_device_properties("cuda").gcnArchName @@ -218,6 +230,14 @@ def _get_gcn_arch() -> str: _ON_GFX90A = "gfx90a" in _GCN_ARCH _ON_GFX942 = "gfx942" in _GCN_ARCH _ON_GFX950 = "gfx950" in _GCN_ARCH +_ON_GFX1250 = "gfx1250" in _GCN_ARCH + +_ON_CDNA = any(arch in _GCN_ARCH for arch in ["gfx9, gfx1250"]) +_ON_RDNA = any( + arch + for arch in _GCN_ARCH + if (arch.startswith("gfx11") or arch.startswith("gfx12")) and arch != "gfx1250" +) def _capability_from_gcn_arch(gcn_arch: str) -> tuple[int, int] | None: @@ -292,7 +312,7 @@ def _capability_from_gcn_arch(gcn_arch: str) -> tuple[int, int] | None: def on_gfx1x() -> bool: - return _ON_GFX1X + return _ON_GFX1X and not _ON_CDNA def on_gfx11() -> bool: @@ -308,7 +328,11 @@ def on_gfx1151() -> bool: def on_gfx12x() -> bool: - return _ON_GFX12X + return _ON_GFX12X and not _ON_CDNA + + +def on_gfx1250() -> bool: + return _ON_GFX1250 def on_mi3xx() -> bool: @@ -331,13 +355,34 @@ def on_gfx950() -> bool: return _ON_GFX950 +def on_cdna() -> bool: + return _ON_CDNA + + +def on_rdna() -> bool: + return _ON_RDNA + + +def get_cdna_version() -> int: + if on_gfx90a(): + return 2 + if on_gfx942(): + return 3 + if on_gfx950(): + return 4 + if on_gfx1250(): + return 5 + return -1 + + + # Enable HIP online tuning early, before hipBLASLt initializes. # Turn on hipBLASLt online tuning if use AITER hipBLASLt GEMM. if ( envs.VLLM_ROCM_USE_AITER and envs.VLLM_ROCM_USE_AITER_LINEAR and envs.VLLM_ROCM_USE_AITER_LINEAR_HIPBMM - and on_mi3xx() + and get_cdna_version() > 2 ): os.environ["HIP_ONLINE_TUNING"] = "1" @@ -356,7 +401,7 @@ def use_rocm_custom_paged_attention( ) -> bool: # custom paged attn always supported on V0. On V1, requires sliding window # disabled due to observed numerical discrepancy. - if _ON_GFX9: + if on_cdna(): return ( (sliding_window == 0 or sliding_window == (-1, -1)) and (qtype == torch.half or qtype == torch.bfloat16) @@ -643,12 +688,12 @@ def get_vit_attn_backend( from vllm._aiter_ops import rocm_aiter_ops - if rocm_aiter_ops.is_enabled() and on_gfx9(): + if rocm_aiter_ops.is_enabled() and on_cdna(): logger.info_once("Using AITER Flash Attention backend for ViT model.") return AttentionBackendEnum.ROCM_AITER_FA if ( - on_gfx9() + on_cdna() and find_spec("flash_attn") is not None and (dtype == torch.float16 or dtype == torch.bfloat16) ): @@ -719,6 +764,9 @@ def is_fully_connected(cls, physical_device_ids: list[int]) -> bool: @with_amdsmi_context @lru_cache(maxsize=8) def get_device_name(cls, device_id: int = 0) -> str: + if not _AMDSMI_AVAILABLE: + # FFM simulator: amdsmi can't see the simulated GPU; use torch. + return torch.cuda.get_device_name(device_id) physical_device_id = cls.device_id_to_physical_device_id(device_id) handle = amdsmi_get_processor_handles()[physical_device_id] asic_info = amdsmi_get_gpu_asic_info(handle) @@ -864,11 +912,11 @@ def get_device_communicator_cls(cls) -> str: @classmethod def supports_mx(cls) -> bool: - return any(gfx in _GCN_ARCH for gfx in ["gfx95"]) + return any(gfx in _GCN_ARCH for gfx in ["gfx95", "gfx1250"]) @classmethod def supports_fp8(cls) -> bool: - return on_gfx9() or on_gfx12x() + return on_cdna() or on_gfx12x() @classmethod def is_fp8_fnuz(cls) -> bool: diff --git a/vllm/renderers/params.py b/vllm/renderers/params.py index 7e3670c738d2..a07a49230676 100644 --- a/vllm/renderers/params.py +++ b/vllm/renderers/params.py @@ -314,10 +314,11 @@ def get_encode_kwargs(self) -> dict[str, Any]: # while still failing `self._token_len_check` as expected by users max_length = self.max_input_tokens + 1 - # Explicit truncation-side overrides require the full token sequence so - # we can slice from the requested side in _token_truncation. Disable - # tokenizer-level truncation because generation tokenizers default to - # left truncation while callers may request right truncation. + # Explicit truncation-side overrides require the full token sequence + # so we can slice from the requested side in _token_truncation. + # Disable tokenizer-level truncation because its default side may + # differ from the requested side. The defense against unbounded + # tokenization lives in _text_len_check (character-level pre-trim). if self.truncation_side is not None and self.truncate_prompt_tokens is not None: return dict( truncation=False, @@ -333,15 +334,13 @@ def get_encode_kwargs(self) -> dict[str, Any]: def _text_len_check(self, tokenizer: TokenizerLike | None, text: str) -> str: """Apply length checks to prompt text if necessary.""" max_input_tokens = self.max_input_tokens - if max_input_tokens is None: + if max_input_tokens is None or tokenizer is None: return text - if self.truncate_prompt_tokens is None and tokenizer is not None: - max_input_chars = max_input_tokens * tokenizer.max_chars_per_token + max_input_chars = max_input_tokens * tokenizer.max_chars_per_token + if self.truncate_prompt_tokens is None: if len(text) > max_input_chars: - # To save resources, fail the request outright without even - # attempting tokenization raise VLLMValidationError( f"This model's maximum context length is " f"{self.max_total_tokens} tokens. However, you requested " @@ -354,6 +353,11 @@ def _text_len_check(self, tokenizer: TokenizerLike | None, text: str) -> str: parameter="input_text", value=len(text), ) + elif self.truncation_side is not None and len(text) > max_input_chars: + if self.truncation_side == "left": + text = text[-max_input_chars:] + else: + text = text[:max_input_chars] return text diff --git a/vllm/utils/import_utils.py b/vllm/utils/import_utils.py index 043798a584b8..7cbf63e5040f 100644 --- a/vllm/utils/import_utils.py +++ b/vllm/utils/import_utils.py @@ -500,7 +500,9 @@ def has_triton_kernels() -> bool: @cache def has_tilelang() -> bool: """Whether the optional `tilelang` package is available.""" - return _has_module("tilelang") + from vllm.platforms.rocm import on_gfx1250 + + return _has_module("tilelang") and not on_gfx1250() def has_arctic_inference() -> bool: diff --git a/vllm/v1/attention/backends/fa_utils.py b/vllm/v1/attention/backends/fa_utils.py index 474523780ff7..b2b173858e07 100644 --- a/vllm/v1/attention/backends/fa_utils.py +++ b/vllm/v1/attention/backends/fa_utils.py @@ -30,17 +30,22 @@ flash_attn_varlen_func = xpu_ops.flash_attn_varlen_func # type: ignore[assignment] get_scheduler_metadata = xpu_ops.get_scheduler_metadata # type: ignore[assignment] elif current_platform.is_rocm(): + # On ROCm we use AITER's Triton flash-attention; the upstream flash-attn + # package is not installed/available. (Same source as aiter_triton_mla.py.) try: - from flash_attn import flash_attn_varlen_func # type: ignore[no-redef] + from aiter.ops.triton.mha import ( # type: ignore[no-redef] + flash_attn_varlen_func, + ) - # Mark that upstream flash-attn is available on ROCm + # A working flash_attn_varlen_func is available on ROCm (via AITER). _ROCM_FLASH_ATTN_AVAILABLE = True except ImportError: def flash_attn_varlen_func(*args: Any, **kwargs: Any) -> Any: # type: ignore[no-redef,misc] raise ImportError( - "ROCm platform requires upstream flash-attn " - "to be installed. Please install flash-attn first." + "ROCm platform requires AITER's Triton MHA " + "(aiter.ops.triton.mha.flash_attn_varlen_func). " + "Please install aiter." ) # ROCm doesn't use scheduler metadata (FA3 feature), provide stub @@ -237,7 +242,7 @@ def is_flash_attn_varlen_func_available() -> bool: Platform-specific sources: - CUDA: vllm.vllm_flash_attn.flash_attn_varlen_func - XPU: xpu_ops.flash_attn_varlen_func - - ROCm: upstream flash_attn.flash_attn_varlen_func (if available) + - ROCm: aiter.ops.triton.mha.flash_attn_varlen_func (if AITER available) Note: This is separate from the AITER flash attention backend (rocm_aiter_fa.py) which uses rocm_aiter_ops.flash_attn_varlen_func. The condition to use AITER is diff --git a/vllm/v1/attention/backends/rocm_aiter_fa.py b/vllm/v1/attention/backends/rocm_aiter_fa.py index 7e850af3e7bc..cea954b47643 100644 --- a/vllm/v1/attention/backends/rocm_aiter_fa.py +++ b/vllm/v1/attention/backends/rocm_aiter_fa.py @@ -768,12 +768,12 @@ def get_kv_cache_shape( @classmethod def supports_compute_capability(cls, capability: DeviceCapability) -> bool: - from vllm.platforms.rocm import on_mi3xx + from vllm.platforms.rocm import get_cdna_version # DeviceCapability is currently created using torch.cuda.get_device_capability() - # which is known to be buggy on rocm systems. on_mi3xx uses amd-smi which is + # which is known to be buggy on rocm systems. on CDNA uses amd-smi which is # more reliable. - return on_mi3xx() + return get_cdna_version() > 2 @classmethod def supports_non_causal(cls) -> bool: diff --git a/vllm/v1/sample/ops/topk_topp_sampler.py b/vllm/v1/sample/ops/topk_topp_sampler.py index 69b35830add2..344c53b04b3a 100644 --- a/vllm/v1/sample/ops/topk_topp_sampler.py +++ b/vllm/v1/sample/ops/topk_topp_sampler.py @@ -10,6 +10,7 @@ from vllm.config.model import LogprobsMode from vllm.logger import init_logger from vllm.platforms import CpuArchEnum, current_platform +from vllm.platforms.rocm import on_gfx1250 from vllm.triton_utils import HAS_TRITON if HAS_TRITON: @@ -110,6 +111,7 @@ def __init__( elif ( logprobs_mode not in ("processed_logits", "processed_logprobs") and rocm_aiter_ops.is_enabled() + and not on_gfx1250() # TODO (JPVILLAM): Enable this path ): self.aiter_ops = None self._aiter_ops_import_failed = False diff --git a/vllm/v1/worker/gpu/warmup.py b/vllm/v1/worker/gpu/warmup.py index ff9a75f0ef1a..de47cb8880a3 100644 --- a/vllm/v1/worker/gpu/warmup.py +++ b/vllm/v1/worker/gpu/warmup.py @@ -34,7 +34,7 @@ def run_mixed_prefill_decode_warmup( req_id_prefix: str = "_v2_mixed_warmup", ) -> bool: """Run a V2 mixed prefill+decode step through normal scheduler inputs.""" - if model_runner.is_pooling_model or num_tokens < 3: + if model_runner.is_pooling_model or model_runner.max_num_reqs < 2 or num_tokens < 3: return False decode_req_id = f"{req_id_prefix}_decode_" diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index 0c5512d5e15f..3db3fbf961c8 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -66,7 +66,6 @@ from vllm.utils.gpu_sync_debug import enable_gpu_sync_check, with_gpu_sync_check from vllm.utils.mem_constants import GiB_bytes from vllm.utils.mem_utils import MemorySnapshot, format_gib, memory_profiling -from vllm.utils.torch_utils import set_random_seed from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheSpec from vllm.v1.outputs import ( @@ -348,7 +347,7 @@ def init_device(self): logger.info_once("Using V2 Model Runner") # Set random seed. - set_random_seed(self.model_config.seed) + # set_random_seed(self.model_config.seed) # Now take memory snapshot after NCCL is initialized gc.collect() @@ -848,7 +847,7 @@ def compile_or_warm_up_model(self) -> CompilationTimes: # Reset the seed to ensure that the random state is not affected by # the model initialization and profiling. - set_random_seed(self.model_config.seed) + # set_random_seed(self.model_config.seed) # Eagerly trigger inductor's once-per-process lazy inits during # warmup (rather than on a later compile cache-miss at runtime).