diff --git a/.github/workflows/ut.yaml b/.github/workflows/ut.yaml index 8b1c8932b..b95f3092c 100644 --- a/.github/workflows/ut.yaml +++ b/.github/workflows/ut.yaml @@ -5,12 +5,23 @@ on: branches: [main] push: branches: [main] + schedule: + # Nightly full-scope run at 02:00 UTC + - cron: '0 2 * * *' + workflow_dispatch: + inputs: + test_scope: + description: 'Test scope: full, ci, mini, ondemand:llama, ondemand:deepseek' + required: false + default: 'ci' permissions: contents: read env: REGISTRY: localhost:5000 + # Nightly (schedule) → full scope; manual dispatch → user choice; PR/push → ci scope + XPU_KERNEL_TEST_SCOPE: ${{ github.event_name == 'schedule' && 'full' || github.event.inputs.test_scope || 'ci' }} jobs: @@ -87,8 +98,9 @@ jobs: - name: test run: | - ZE_AFFINITY_MASK=0,1 SKIP_HANG_KERNEL=1 SKIP_ACC_ERROR_KERNEL=1 pytest -v -s tests/ - VLLM_XPU_FORCE_XE_DEFAULT_KERNEL=1 ZE_AFFINITY_MASK=0,1 pytest -v -s tests/fused_moe/test_grouped_gemm.py::test_grouped_gemm + echo "Running tests with XPU_KERNEL_TEST_SCOPE=${{ env.XPU_KERNEL_TEST_SCOPE }}" + XPU_KERNEL_TEST_SCOPE=${{ env.XPU_KERNEL_TEST_SCOPE }} ZE_AFFINITY_MASK=0,1 SKIP_HANG_KERNEL=1 SKIP_ACC_ERROR_KERNEL=1 pytest -v -s tests/ + VLLM_XPU_FORCE_XE_DEFAULT_KERNEL=1 XPU_KERNEL_TEST_SCOPE=${{ env.XPU_KERNEL_TEST_SCOPE }} ZE_AFFINITY_MASK=0,1 pytest -v -s tests/fused_moe/test_grouped_gemm.py::test_grouped_gemm clean-repo-pvc: runs-on: self-hosted-pvc @@ -129,11 +141,12 @@ jobs: - name: test run: | + echo "Running tests with XPU_KERNEL_TEST_SCOPE=${{ env.XPU_KERNEL_TEST_SCOPE }}" # tests/test_moe_align_block_size.py, tests/test_moe_lora_align_sum.py takes much time than expected. ignore it for now. - ZE_AFFINITY_MASK=0,1 pytest -v -s tests/ --ignore=tests/test_lora_ops.py --ignore=tests/test_fp8_quant.py --ignore=tests/test_moe_align_block_size.py --ignore=tests/test_moe_lora_align_sum.py --ignore=tests/test_cache.py::test_swap_blocks --ignore=tests/test_topk_per_row.py --ignore=tests/test_lora_ops.py + XPU_KERNEL_TEST_SCOPE=${{ env.XPU_KERNEL_TEST_SCOPE }} ZE_AFFINITY_MASK=0,1 pytest -v -s tests/ --ignore=tests/test_lora_ops.py --ignore=tests/test_fp8_quant.py --ignore=tests/test_moe_align_block_size.py --ignore=tests/test_moe_lora_align_sum.py --ignore=tests/test_cache.py::test_swap_blocks --ignore=tests/test_topk_per_row.py --ignore=tests/test_lora_ops.py # fixme: Running lora UT separately to avoid OOM when running together with other tests. - ZE_AFFINITY_MASK=0,1 pytest -v -s tests/test_lora_ops.py - VLLM_XPU_FORCE_XE_DEFAULT_KERNEL=1 ZE_AFFINITY_MASK=0,1 pytest -v -s tests/fused_moe/test_grouped_gemm.py::test_grouped_gemm + XPU_KERNEL_TEST_SCOPE=${{ env.XPU_KERNEL_TEST_SCOPE }} ZE_AFFINITY_MASK=0,1 pytest -v -s tests/test_lora_ops.py + VLLM_XPU_FORCE_XE_DEFAULT_KERNEL=1 XPU_KERNEL_TEST_SCOPE=${{ env.XPU_KERNEL_TEST_SCOPE }} ZE_AFFINITY_MASK=0,1 pytest -v -s tests/fused_moe/test_grouped_gemm.py::test_grouped_gemm clean-repo-bmg: runs-on: self-hosted-bmg diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 59e85cd82..3cb04a6e9 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -6,11 +6,6 @@ default_stages: - manual # Run in CI exclude: 'vllm/third_party/.*' repos: -- repo: https://github.com/google/yapf - rev: v0.43.0 - hooks: - - id: yapf - args: [--in-place, --verbose] - repo: https://github.com/astral-sh/ruff-pre-commit rev: v0.11.7 hooks: diff --git a/CMakeLists.txt b/CMakeLists.txt index d40193091..6ada68f3d 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -50,14 +50,35 @@ set(BUILD_SYCL_TLA_KERNELS ON CACHE BOOL "Build SYCL-TLA based kernels for XPU") # ARCHITECTURE OPTIONS -set(VLLM_XPU_ENABLE_XE2 ON) -set(VLLM_XPU_ENABLE_XE_DEFAULT ON) - -set(BASIC_KERNELS_ENABLED ON) -set(FA2_KERNELS_ENABLED ON) -set(GND_KERNELS_ENABLED ON) -set(MOE_KERNELS_ENABLED ON) -set(XPU_SPECIFIC_KERNELS_ENABLED ON) +option(VLLM_XPU_ENABLE_XE2 "Enable XE2 architecture kernels" ON) +option(VLLM_XPU_ENABLE_XE_DEFAULT "Enable XE Default architecture kernels" ON) + +# KERNEL OPTIONS — each controls whether the corresponding Python extension is +# built. Override from the command line with -D=OFF or via the +# identically-named environment variable (read by setup.py). +option(BASIC_KERNELS_ENABLED "Build basic kernels (_C extension)" ON) +option(FA2_KERNELS_ENABLED + "Build Flash Attention 2 kernels (_vllm_fa2_C extension)" ON) +option(MOE_KERNELS_ENABLED + "Build MoE kernels (_moe_C extension + grouped_gemm TLA)" ON) +option(GDN_KERNELS_ENABLED "Build GDN attention kernels (gdn_attn TLA)" ON) +option(XPU_SPECIFIC_KERNELS_ENABLED + "Build XPU-specific kernels (_xpu_C extension)" ON) +option(XPUMEM_ALLOCATOR_ENABLED "Build xpumem_allocator extension" ON) + +message(STATUS "") +message(STATUS "Kernel build configuration:") +message(STATUS " BUILD_SYCL_TLA_KERNELS = ${BUILD_SYCL_TLA_KERNELS}") +message(STATUS " VLLM_XPU_ENABLE_XE2 = ${VLLM_XPU_ENABLE_XE2}") +message(STATUS " VLLM_XPU_ENABLE_XE_DEFAULT = ${VLLM_XPU_ENABLE_XE_DEFAULT}") +message(STATUS " BASIC_KERNELS_ENABLED = ${BASIC_KERNELS_ENABLED}") +message(STATUS " FA2_KERNELS_ENABLED = ${FA2_KERNELS_ENABLED}") +message(STATUS " MOE_KERNELS_ENABLED = ${MOE_KERNELS_ENABLED}") +message(STATUS " GDN_KERNELS_ENABLED = ${GDN_KERNELS_ENABLED}") +message( + STATUS " XPU_SPECIFIC_KERNELS_ENABLED = ${XPU_SPECIFIC_KERNELS_ENABLED}") +message(STATUS " XPUMEM_ALLOCATOR_ENABLED = ${XPUMEM_ALLOCATOR_ENABLED}") +message(STATUS "") # # Try to find python package with an executable that exactly matches @@ -313,23 +334,40 @@ if(BUILD_SYCL_TLA_KERNELS) # extensions shared library set(SYCL_TLA_COMPILE_OPTIONS "") if(VLLM_XPU_ENABLE_XE_DEFAULT) - add_subdirectory(csrc/xpu/grouped_gemm/xe_default) - list(APPEND GROUPED_GEMM_LIB_NAME "grouped_gemm_xe_default") + if(MOE_KERNELS_ENABLED) + add_subdirectory(csrc/xpu/grouped_gemm/xe_default) + list(APPEND GROUPED_GEMM_LIB_NAME "grouped_gemm_xe_default") + endif() list(APPEND SYCL_TLA_COMPILE_OPTIONS -DVLLM_XPU_ENABLE_XE_DEFAULT) endif() if(VLLM_XPU_ENABLE_XE2) - add_subdirectory(csrc/xpu/grouped_gemm/xe_2) - add_subdirectory(csrc/xpu/attn/xe_2) - add_subdirectory(csrc/xpu/gdn_attn/xe_2) - list(APPEND GROUPED_GEMM_LIB_NAME "grouped_gemm_xe_2") - list(APPEND ATTN_KERNEL_LIB_NAME "attn_kernels_xe_2") - list(APPEND GDN_ATTN_LIB_NAME "gdn_attn_kernels_xe_2") + if(MOE_KERNELS_ENABLED) + add_subdirectory(csrc/xpu/grouped_gemm/xe_2) + list(APPEND GROUPED_GEMM_LIB_NAME "grouped_gemm_xe_2") + endif() + if(FA2_KERNELS_ENABLED) + add_subdirectory(csrc/xpu/attn/xe_2) + list(APPEND ATTN_KERNEL_LIB_NAME "attn_kernels_xe_2") + endif() + if(GDN_KERNELS_ENABLED) + add_subdirectory(csrc/xpu/gdn_attn/xe_2) + list(APPEND GDN_ATTN_LIB_NAME "gdn_attn_kernels_xe_2") + endif() list(APPEND SYCL_TLA_COMPILE_OPTIONS -DVLLM_XPU_ENABLE_XE2) endif() list(APPEND VLLM_GPU_COMPILE_FLAGS ${SYCL_TLA_COMPILE_OPTIONS}) endif() +# Feature compile defines — these guard op registrations and interface code so +# that disabled features don't pull in unbuilt TLA library symbols. +if(MOE_KERNELS_ENABLED) + list(APPEND VLLM_GPU_COMPILE_FLAGS -DVLLM_MOE_ENABLED) +endif() +if(GDN_KERNELS_ENABLED) + list(APPEND VLLM_GPU_COMPILE_FLAGS -DVLLM_GDN_ENABLED) +endif() + # define vLLM XPU cmake variables set(VLLM_XPU_INCLUDE_DIR "") @@ -352,9 +390,6 @@ list( # # xpumem_allocator extension - XPU memory allocator with Python callbacks # -set(XPUMEM_ALLOCATOR_ENABLED - ON - CACHE BOOL "Build xpumem_allocator extension") if(XPUMEM_ALLOCATOR_ENABLED) message(STATUS "Enabling xpumem_allocator extension.") @@ -404,8 +439,10 @@ if(BASIC_KERNELS_ENABLED) set(VLLM_EXT_SRC "csrc/cache.cpp" "csrc/layernorm.cpp" + "csrc/layernorm_quant.cpp" "csrc/activation.cpp" "csrc/pos_encoding_kernels.cpp" + "csrc/fused_qknorm_rope.cpp" "csrc/torch_bindings.cpp" "csrc/quantization/fp8/fp8_quant.cpp" "csrc/quantization/fp4/mxfp4_quant.cpp" @@ -441,6 +478,14 @@ endif() # # flash attention _C extension # +if(FA2_KERNELS_ENABLED AND NOT ATTN_KERNEL_LIB_NAME) + message( + WARNING + "FA2_KERNELS_ENABLED is ON but no attention kernel libraries are available. " + "The _vllm_fa2_C extension will be built without kernel libraries and may " + "fail at runtime. Enable BUILD_SYCL_TLA_KERNELS and VLLM_XPU_ENABLE_XE2 " + "for full FA2 functionality.") +endif() if(FA2_KERNELS_ENABLED) message(STATUS "Enabling fa2 extension.") file(GLOB FA2_GEN_SRCS "csrc/flash_attn/*.cpp" @@ -482,9 +527,14 @@ if(XPU_SPECIFIC_KERNELS_ENABLED) "csrc/xpu/sampler/topk_topp_sampler.cpp" "csrc/xpu/sycl/deepseek_scaling_rope.cpp" "csrc/xpu/rand/exponential.cpp" - "csrc/xpu/grouped_gemm/grouped_gemm_interface.cpp" - "csrc/xpu/utils.cpp" - "csrc/xpu/gdn_attn/gdn_attn_interface.cpp") + "csrc/xpu/utils.cpp") + if(MOE_KERNELS_ENABLED) + list(APPEND VLLM_EXT_XPU_SRC + "csrc/xpu/grouped_gemm/grouped_gemm_interface.cpp") + endif() + if(GDN_KERNELS_ENABLED) + list(APPEND VLLM_EXT_XPU_SRC "csrc/xpu/gdn_attn/gdn_attn_interface.cpp") + endif() include_directories("/usr/include") # TODO: check if we need this flags list(APPEND VLLM_GPU_FLAGS # "-gline-tables-only") diff --git a/csrc/activation.cpp b/csrc/activation.cpp index 9bef17807..f957136b8 100644 --- a/csrc/activation.cpp +++ b/csrc/activation.cpp @@ -1,9 +1,14 @@ #include #include #include +#include #include "utils.h" #include "dispatch_utils.h" +#include +#include +#include "quantization/fp8/quant_utils.h" + #define VLLM_LDG(arg) *(arg) namespace vllm { @@ -67,6 +72,12 @@ inline T gelu_tanh_kernel(const T& x) { return (T)(0.5f * f * (1.0f + sycl::tanh(inner))); } +template +inline T fatrelu_kernel(const T& x, const float threshold) { + const float f = (float)x; + return (T)(f > threshold ? f : 0.0f); +} + template < typename scalar_t, scalar_t (*ACT_FN)(const scalar_t&), @@ -172,6 +183,97 @@ class act_and_mul_vec_kernel { const int d_; }; +template < + typename scalar_t, + scalar_t (*ACT_FN)(const scalar_t&), + typename fp8_type, + int VEC_SIZE> +class act_and_mul_quant_vec_kernel { + public: + act_and_mul_quant_vec_kernel( + fp8_type* __restrict__ out, // [..., d] + const scalar_t* __restrict__ input, // [..., 2 * d] + const float* __restrict__ scale, // [1] + const int d) + : out_(out), input_(input), scale_(scale), d_(d) {} + + void operator()(sycl::nd_item<1> item) const { + using vec_t = vllm::xpu::aligned_vec; + + const int64_t token_idx = item.get_group(0); + const int64_t offset = item.get_local_linear_id(); + const int64_t step = item.get_local_range(0); + const int64_t bound = d_ / VEC_SIZE; + + const float inv_scale = 1.0f / (*scale_); + const float fp8_max = static_cast(fp8::quant_type_max_v); + + // x and y halves are laid out contiguously: [x0..xd-1, y0..yd-1] + const auto* v_x = + reinterpret_cast(input_) + token_idx * bound * 2; + const auto* v_y = v_x + bound; + + for (int64_t i = offset; i < bound; i += step) { + vec_t xv = v_x[i]; + vec_t yv = v_y[i]; +#pragma unroll + for (int j = 0; j < VEC_SIZE; j++) { + float val = static_cast(ACT_FN(xv[j]) * yv[j]) * inv_scale; + float clamped = sycl::fmax(-fp8_max, sycl::fmin(val, fp8_max)); + out_[token_idx * d_ + i * VEC_SIZE + j] = + static_cast(clamped); + } + } + } + + private: + fp8_type* __restrict__ out_; // [..., d] + const scalar_t* __restrict__ input_; // [..., 2 * d] + const float* __restrict__ scale_; // [1] + const int d_; +}; + +template < + typename scalar_t, + scalar_t (*ACT_FN)(const scalar_t&, const float), + int VEC_SIZE> +class act_and_mul_with_param_vec_kernel { + public: + act_and_mul_with_param_vec_kernel( + scalar_t* __restrict__ out, + const scalar_t* __restrict__ input, + const int d, + const float param) + : out_(out), input_(input), d_(d), param_(param) {} + + void operator()(sycl::nd_item<1> item) const { + using vec_t = vllm::xpu::aligned_vec; + const int64_t token_idx = item.get_group(0); + const int64_t offset = item.get_local_linear_id(); + const int64_t step = item.get_local_range(0); + const int64_t bound = d_ / VEC_SIZE; + + for (int64_t i = offset; i < bound; i += step) { + auto x_vec = + reinterpret_cast(input_)[token_idx * bound * 2 + i]; + auto y_vec = reinterpret_cast( + input_)[token_idx * bound * 2 + i + bound]; + vec_t out_vec; +#pragma unroll + for (int j = 0; j < VEC_SIZE; ++j) { + out_vec[j] = ACT_FN(x_vec[j], param_) * y_vec[j]; + } + reinterpret_cast(out_)[token_idx * bound + i] = out_vec; + } + } + + private: + scalar_t* __restrict__ out_; + const scalar_t* __restrict__ input_; + const int d_; + const float param_; +}; + template [[intel::device_indirectly_callable]] inline __attribute__((always_inline)) T swigluoai_and_mul(const T& gate, const T& up, float alpha, float limit) { @@ -311,6 +413,17 @@ class swiglustep_and_mul_kernel { break; \ } +#define VEC_LAUNCH_ACT_AND_MUL_WITH_PARAM(KERNEL, N) \ + case N: { \ + queue.submit([&](sycl::handler& cgh) { \ + cgh.parallel_for( \ + sycl::nd_range<1>(num_tokens * wg_size, wg_size), \ + vllm::act_and_mul_with_param_vec_kernel( \ + (sycl_t*)out_ptr, (sycl_t*)input_ptr, d, param)); \ + }); \ + break; \ + } + #define LAUNCH_ACTIVATION_GATE_KERNEL_VEC(KERNEL, ACT_FIRST) \ using sycl_t = vllm::xpu::SyclTypeTrait::Type; \ int d = input.size(-1) / 2; \ @@ -343,6 +456,39 @@ class swiglustep_and_mul_kernel { TORCH_CHECK(false, "Unsupported vector size: ", vec_size); \ } +#define LAUNCH_ACTIVATION_GATE_KERNEL_WITH_PARAM_VEC(KERNEL, PARAM) \ + using sycl_t = vllm::xpu::SyclTypeTrait::Type; \ + int d = input.size(-1) / 2; \ + int64_t num_tokens = input.numel() / input.size(-1); \ + if (num_tokens == 0) { \ + return; \ + } \ + auto out_ptr = out.data_ptr(); \ + auto input_ptr = input.data_ptr(); \ + const float param = static_cast(PARAM); \ + at::DeviceGuard device_guard(input.device()); \ + auto& queue = vllm::xpu::vllmGetQueue(); \ + int vec_size = static_cast(sizeof(float) * 4 / sizeof(scalar_t)); \ + { \ + int64_t tmp_wg = \ + std::min(static_cast(d), static_cast(1024)); \ + while (vec_size > 1 && (vec_size >> 1) * tmp_wg >= d) { \ + vec_size = vec_size >> 1; \ + } \ + } \ + if (d % vec_size != 0) vec_size = 1; \ + int64_t wg_size = std::min( \ + static_cast(d / vec_size), static_cast(1024)); \ + switch (vec_size) { \ + VEC_LAUNCH_ACT_AND_MUL_WITH_PARAM(KERNEL, 1); \ + VEC_LAUNCH_ACT_AND_MUL_WITH_PARAM(KERNEL, 2); \ + VEC_LAUNCH_ACT_AND_MUL_WITH_PARAM(KERNEL, 4); \ + VEC_LAUNCH_ACT_AND_MUL_WITH_PARAM(KERNEL, 8); \ + VEC_LAUNCH_ACT_AND_MUL_WITH_PARAM(KERNEL, 16); \ + default: \ + TORCH_CHECK(false, "Unsupported vector size: ", vec_size); \ + } + void silu_and_mul( torch::Tensor& out, // [..., d] torch::Tensor& input) // [..., 2 * d] @@ -352,6 +498,73 @@ void silu_and_mul( }); } +// Fused SiLU + Mul + FP8 Quantization +// Input: [..., 2*d] in FP16/BF16, Output: [..., d] in FP8 +// Dispatches to the vectorized kernel (VEC_SIZE=1..8) based on alignment. +#define LAUNCH_ACT_AND_MUL_QUANT_VEC(KERNEL, N) \ + case N: { \ + int64_t wg_size = \ + std::min(static_cast(d / N), static_cast(1024)); \ + VLLM_DISPATCH_FP8_TYPES( \ + out.scalar_type(), "act_and_mul_quant_vec_kernel_fp8", [&] { \ + auto out_ptr = out.data_ptr(); \ + queue.submit([&](sycl::handler& cgh) { \ + cgh.parallel_for( \ + sycl::nd_range<1>(num_tokens * wg_size, wg_size), \ + vllm::act_and_mul_quant_vec_kernel( \ + out_ptr, (sycl_t*)input_ptr, scale_ptr, d)); \ + }); \ + }); \ + break; \ + } + +#define LAUNCH_ACTIVATION_GATE_QUANT_KERNEL(KERNEL) \ + using sycl_t = vllm::xpu::SyclTypeTrait::Type; \ + int d = input.size(-1) / 2; \ + int64_t num_tokens = input.numel() / input.size(-1); \ + if (num_tokens == 0) { \ + return; \ + } \ + auto input_ptr = input.data_ptr(); \ + auto scale_ptr = scale.data_ptr(); \ + at::DeviceGuard device_guard(input.device()); \ + auto& queue = vllm::xpu::vllmGetQueue(); \ + /* Compute vec_size like non-quant path: gcd(4*sizeof(float)/sizeof, d) */ \ + int vec_size = static_cast(sizeof(float) * 4 / sizeof(sycl_t)); \ + { \ + int64_t tmp_wg = \ + std::min(static_cast(d), static_cast(1024)); \ + while (vec_size > 1 && (vec_size >> 1) * tmp_wg >= d) { \ + vec_size = vec_size >> 1; \ + } \ + } \ + if (d % vec_size != 0) vec_size = 1; \ + switch (vec_size) { \ + LAUNCH_ACT_AND_MUL_QUANT_VEC(KERNEL, 1); \ + LAUNCH_ACT_AND_MUL_QUANT_VEC(KERNEL, 2); \ + LAUNCH_ACT_AND_MUL_QUANT_VEC(KERNEL, 4); \ + LAUNCH_ACT_AND_MUL_QUANT_VEC(KERNEL, 8); \ + LAUNCH_ACT_AND_MUL_QUANT_VEC(KERNEL, 16); \ + default: \ + TORCH_CHECK(false, "Unsupported vector size: ", vec_size); \ + } + +void silu_and_mul_quant( + torch::Tensor& out, // [..., d] FP8 + torch::Tensor& input, // [..., 2 * d] FP16/BF16 + torch::Tensor& scale) // [1] FP32 +{ + TORCH_CHECK( + out.dtype() == torch::kFloat8_e4m3fn || + out.dtype() == torch::kFloat8_e5m2); + TORCH_CHECK( + input.dtype() == torch::kFloat16 || input.dtype() == torch::kBFloat16); + TORCH_CHECK(input.size(-1) % 2 == 0); + VLLM_DISPATCH_FLOATING_TYPES(input.scalar_type(), "silu_and_mul_quant", [&] { + LAUNCH_ACTIVATION_GATE_QUANT_KERNEL(vllm::silu_kernel); + }); +} + void mul_and_silu( torch::Tensor& out, // [..., d] torch::Tensor& input) // [..., 2 * d] @@ -379,6 +592,16 @@ void gelu_tanh_and_mul( }); } +void fatrelu_and_mul( + torch::Tensor& out, // [..., d] + torch::Tensor& input, // [..., 2 * d] + double threshold) { + VLLM_DISPATCH_FLOATING_TYPES(input.scalar_type(), "fatrelu_and_mul", [&] { + LAUNCH_ACTIVATION_GATE_KERNEL_WITH_PARAM_VEC( + vllm::fatrelu_kernel, threshold); + }); +} + // Launch element-wise activation kernel. #define LAUNCH_ACTIVATION_KERNEL(KERNEL) \ using sycl_t = vllm::xpu::SyclTypeTrait::Type; \ diff --git a/csrc/cache.cpp b/csrc/cache.cpp index 3dd01a9e2..afff3daa1 100644 --- a/csrc/cache.cpp +++ b/csrc/cache.cpp @@ -1175,6 +1175,37 @@ void swap_blocks( return; } +/** + * @brief Batch version of swap_blocks: copies N independent (src, dst, size) + * triples in a single call, amortising per-copy overhead. + * + * Thin wrapper that validates the CPU tensor inputs and delegates to + * vllm::xpu::xpuAsyncMemcpyBatch for the actual copy logic. + */ +void swap_blocks_batch( + const torch::Tensor& src_ptrs, + const torch::Tensor& dst_ptrs, + const torch::Tensor& sizes) { + TORCH_CHECK(src_ptrs.device().is_cpu(), "src_ptrs must be on CPU"); + TORCH_CHECK(dst_ptrs.device().is_cpu(), "dst_ptrs must be on CPU"); + TORCH_CHECK(sizes.device().is_cpu(), "sizes must be on CPU"); + TORCH_CHECK(src_ptrs.dtype() == torch::kUInt64, "src_ptrs must be uint64"); + TORCH_CHECK(dst_ptrs.dtype() == torch::kUInt64, "dst_ptrs must be uint64"); + TORCH_CHECK(sizes.dtype() == torch::kUInt64, "sizes must be uint64"); + + const int64_t n = src_ptrs.size(0); + TORCH_CHECK(dst_ptrs.size(0) == n, "dst_ptrs length must match src_ptrs"); + TORCH_CHECK(sizes.size(0) == n, "sizes length must match src_ptrs"); + + if (n == 0) return; + + vllm::xpu::xpuAsyncMemcpyBatch( + src_ptrs.data_ptr(), + dst_ptrs.data_ptr(), + sizes.data_ptr(), + n); +} + namespace vllm { // Kernel for FP8 conversion (matches CUDA convert_fp8_kernel pattern). diff --git a/csrc/flash_attn/flash_api.cpp b/csrc/flash_attn/flash_api.cpp index dade9f882..9eb694724 100644 --- a/csrc/flash_attn/flash_api.cpp +++ b/csrc/flash_attn/flash_api.cpp @@ -22,36 +22,33 @@ inline int get_num_splits( int cur_parallel = batch_size * num_heads_kv; int kv_blocks = (max_seqlen_k + block_size - 1) / block_size; - // Below 128 KV blocks the per-split FMHA compute is too small relative - // to the ReduceSplitK overhead, regardless of block size. - if (kv_blocks < 128) return 1; + // Minimum KV blocks to benefit from splitting, aligned with the kernel's + // kMinBlocksForSplit. Below this the kernel falls back to single-split + // regardless, and splitting only adds ReduceSplitK overhead. + int min_kv_blocks = (block_size <= 64) ? 32 : 128; + if (kv_blocks < min_kv_blocks) return 1; + // GPU well-utilized or saturated: splitting only adds ReduceSplitK + // overhead without improving FMHA throughput. + if (cur_parallel >= num_xe_cores) return 1; + + // Under-utilized (cur_parallel < num_xe_cores): split to fill the GPU. int target_splits; - if (cur_parallel < num_xe_cores) { - // Under-utilized: fill GPU cores. - // Scale by block_size since larger blocks mean more compute per WG. - int eff_parallel = cur_parallel * block_size / 64; - eff_parallel = std::max(1, eff_parallel); - target_splits = (num_xe_cores + eff_parallel - 1) / eff_parallel; - } else if (cur_parallel <= num_xe_cores * 2) { - // Well-utilized zone (1x-2x oversubscription): - // GPU is busy, splitting adds overhead without benefit. - return 1; + if (num_heads_kv >= 4) { + // Many KV heads: each split adds kv_heads WGs. Scale inversely + // with block_size — p64 gets ~20 splits, p128 gets ~10. + target_splits = std::max(4, num_xe_cores * 64 / block_size); } else { - // Heavily oversubscribed (>2x): shorter WGs help. - // But gate out when compute is already saturated. - int eff_parallel = cur_parallel * block_size / 64; - if (eff_parallel >= num_xe_cores * 8) return 1; - target_splits = std::max(1, kv_blocks / 64); - int par_cap = std::max(1, num_xe_cores * 8 / cur_parallel); - target_splits = std::min(target_splits, par_cap); + // Few KV heads: target at least num_xe_cores splits for parallelism; + // allow more for very long sequences (~10 blocks/split at p64). + int blocks_per_split = (block_size <= 64) ? 10 : 8; + target_splits = std::max(num_xe_cores, kv_blocks / blocks_per_split); } - // Each split must process at least 32 KV blocks. - int max_splits_blocks = std::max(1, kv_blocks / 32); - // Hard cap: more splits give diminishing returns and increase - // ReduceSplitK overhead and temporary buffer memory. - int num_splits = std::min({target_splits, max_splits_blocks, 8}); + // Each split must process at least 3 KV blocks to amortize overhead. + int max_splits_blocks = std::max(1, kv_blocks / 3); + // Hard cap: beyond 40 splits, diminishing returns. + int num_splits = std::min({target_splits, max_splits_blocks, 40}); return std::max(1, num_splits); } @@ -208,11 +205,10 @@ std::vector mha_varlen_fwd( : at::empty( {num_tokens, num_heads_q * num_kv_splits, v_head_dim}, q.options().device(q.device())); - at::Tensor max_logits = at::full( + at::Tensor max_logits = at::empty( {num_tokens, num_heads_q, num_kv_splits}, - -std::numeric_limits::infinity(), q.options().dtype(at::kFloat).device(q.device())); - at::Tensor exp_sums = at::zeros( + at::Tensor exp_sums = at::empty( {num_tokens, num_heads_q, num_kv_splits}, q.options().dtype(at::kFloat).device(q.device())); diff --git a/csrc/fused_qknorm_rope.cpp b/csrc/fused_qknorm_rope.cpp new file mode 100644 index 000000000..f8d76db62 --- /dev/null +++ b/csrc/fused_qknorm_rope.cpp @@ -0,0 +1,368 @@ +/* + * Ported from the CUDA implementation in csrc/fused_qknorm_rope_kernel.cu. + */ + +#include + +#include +#include + +#include "dispatch_utils.h" +#include "utils.h" + +namespace vllm { + +template < + typename scalar_t, + typename scalar_t_cache, + int head_dim, + bool IS_NEOX> +class fused_qk_norm_rope_kernel { + public: + static constexpr int SUB_GROUP_SIZE = 32; + static constexpr int NUM_ELEMS_PER_THREAD = head_dim / SUB_GROUP_SIZE; + + static_assert( + head_dim % (SUB_GROUP_SIZE * 2) == 0, "head_dim must be divisible by 64"); + + fused_qk_norm_rope_kernel( + scalar_t* __restrict__ qkv_, + const int num_heads_q_, + const int num_heads_k_, + const int num_heads_v_, + const float eps_, + const scalar_t* __restrict__ q_weight_, + const scalar_t* __restrict__ k_weight_, + const scalar_t_cache* __restrict__ cos_sin_cache_, + const int64_t* __restrict__ position_ids_, + const int num_tokens_, + const int rotary_dim_) + : qkv(qkv_), + num_heads_q(num_heads_q_), + num_heads_k(num_heads_k_), + num_heads_v(num_heads_v_), + eps(eps_), + q_weight(q_weight_), + k_weight(k_weight_), + cos_sin_cache(cos_sin_cache_), + position_ids(position_ids_), + num_tokens(num_tokens_), + rotary_dim(rotary_dim_) {} + + void operator() [[sycl::reqd_sub_group_size(SUB_GROUP_SIZE)]] ( + const sycl::nd_item<1>& item) const { + auto sg = item.get_sub_group(); + const int lane_id = sg.get_local_linear_id(); + const int sg_id_in_wg = sg.get_group_linear_id(); + const int sgs_per_wg = sg.get_group_linear_range(); + const int global_sg_id = item.get_group(0) * sgs_per_wg + sg_id_in_wg; + + const int total_qk_heads = num_heads_q + num_heads_k; + const int token_idx = global_sg_id / total_qk_heads; + const int local_head_idx = global_sg_id % total_qk_heads; + + if (token_idx >= num_tokens) return; + + const bool is_q = local_head_idx < num_heads_q; + const int head_idx = is_q ? local_head_idx : local_head_idx - num_heads_q; + const int num_heads = num_heads_q + num_heads_k + num_heads_v; + + // Compute the offset into the QKV tensor for this sub-group's head. + int offset_warp; + if (is_q) { + offset_warp = token_idx * num_heads * head_dim + head_idx * head_dim; + } else { + offset_warp = token_idx * num_heads * head_dim + num_heads_q * head_dim + + head_idx * head_dim; + } + int offset_thread = offset_warp + lane_id * NUM_ELEMS_PER_THREAD; + + // Load elements and compute sum of squares for RMSNorm. + float elements[NUM_ELEMS_PER_THREAD]; + float sum_of_squares = 0.0f; + +#pragma unroll + for (int i = 0; i < NUM_ELEMS_PER_THREAD; i++) { + float val = static_cast(qkv[offset_thread + i]); + elements[i] = val; + sum_of_squares += val * val; + } + + // Reduce sum across sub-group. + sum_of_squares = + sycl::reduce_over_group(sg, sum_of_squares, sycl::plus()); + + // Compute RMS normalization factor. + float rms_rcp = + sycl::rsqrt(sum_of_squares / static_cast(head_dim) + eps); + + // Apply RMSNorm with learned weights. + const scalar_t* weight = is_q ? q_weight : k_weight; +#pragma unroll + for (int i = 0; i < NUM_ELEMS_PER_THREAD; i++) { + int dim = lane_id * NUM_ELEMS_PER_THREAD + i; + float w = static_cast(weight[dim]); + elements[i] *= rms_rcp * w; + } + + // Apply RoPE. + const int64_t pos_id = position_ids[token_idx]; + const int embed_dim = rotary_dim / 2; + const scalar_t_cache* cache_ptr = cos_sin_cache + pos_id * rotary_dim; + const scalar_t_cache* cos_ptr = cache_ptr; + const scalar_t_cache* sin_ptr = cache_ptr + embed_dim; + + const int rotary_lanes = rotary_dim / NUM_ELEMS_PER_THREAD; + + if (lane_id < rotary_lanes) { + if constexpr (!IS_NEOX) { + // Interleaved-style RoPE (GPT-J style). +#pragma unroll + for (int i = 0; i < NUM_ELEMS_PER_THREAD / 2; ++i) { + const int idx0 = 2 * i; + const int idx1 = 2 * i + 1; + const int dim_idx = lane_id * NUM_ELEMS_PER_THREAD + idx0; + + const float val0 = elements[idx0]; + const float val1 = elements[idx1]; + + const int half_dim = dim_idx / 2; + const float cos_val = static_cast(cos_ptr[half_dim]); + const float sin_val = static_cast(sin_ptr[half_dim]); + + elements[idx0] = val0 * cos_val - val1 * sin_val; + elements[idx1] = val0 * sin_val + val1 * cos_val; + } + } else { + // Neox-style RoPE: exchange data with partner lane. + sycl::group_barrier(sg); + const int pair_offset = (rotary_dim / 2) / NUM_ELEMS_PER_THREAD; + +#pragma unroll + for (int i = 0; i < NUM_ELEMS_PER_THREAD; i++) { + float partner_val = + sycl::permute_group_by_xor(sg, elements[i], pair_offset); + + if (lane_id < pair_offset) { + partner_val = -partner_val; + } + + int dim_idx = lane_id * NUM_ELEMS_PER_THREAD + i; + dim_idx = (dim_idx * 2) % rotary_dim; + int half_dim = dim_idx / 2; + + const float cos_val = static_cast(cos_ptr[half_dim]); + const float sin_val = static_cast(sin_ptr[half_dim]); + + elements[i] = elements[i] * cos_val + partner_val * sin_val; + } + sycl::group_barrier(sg); + } + } + + // Store results back in-place. +#pragma unroll + for (int i = 0; i < NUM_ELEMS_PER_THREAD; i++) { + qkv[offset_thread + i] = static_cast(elements[i]); + } + } + + private: + scalar_t* __restrict__ qkv; + const int num_heads_q; + const int num_heads_k; + const int num_heads_v; + const float eps; + const scalar_t* __restrict__ q_weight; + const scalar_t* __restrict__ k_weight; + const scalar_t_cache* __restrict__ cos_sin_cache; + const int64_t* __restrict__ position_ids; + const int num_tokens; + const int rotary_dim; +}; + +template +void launch_fused_qk_norm_rope( + torch::Tensor& qkv, + int num_tokens, + int num_heads_q, + int num_heads_k, + int num_heads_v, + int head_dim, + int rotary_dim, + float eps, + torch::Tensor& q_weight, + torch::Tensor& k_weight, + torch::Tensor& cos_sin_cache, + bool is_neox, + torch::Tensor& position_ids) { + using sycl_t = typename vllm::xpu::SyclTypeTrait::Type; + using sycl_cache_t = typename vllm::xpu::SyclTypeTrait::Type; + + constexpr int block_size = 256; + constexpr int sgs_per_wg = block_size / 32; + + const int total_qk_heads = num_heads_q + num_heads_k; + const int total_sgs = num_tokens * total_qk_heads; + const int grid_size = (total_sgs + sgs_per_wg - 1) / sgs_per_wg; + + auto qkv_ptr = reinterpret_cast(qkv.data_ptr()); + auto q_weight_ptr = + reinterpret_cast(q_weight.data_ptr()); + auto k_weight_ptr = + reinterpret_cast(k_weight.data_ptr()); + auto cos_sin_cache_ptr = reinterpret_cast( + cos_sin_cache.data_ptr()); + auto position_ids_ptr = position_ids.data_ptr(); + + auto& queue = vllm::xpu::vllmGetQueue(); + +#define DISPATCH_HEAD_DIM(HEAD_DIM) \ + do { \ + if (is_neox) { \ + queue.submit([&](sycl::handler& cgh) { \ + cgh.parallel_for( \ + sycl::nd_range<1>(grid_size * block_size, block_size), \ + fused_qk_norm_rope_kernel( \ + qkv_ptr, \ + num_heads_q, \ + num_heads_k, \ + num_heads_v, \ + eps, \ + q_weight_ptr, \ + k_weight_ptr, \ + cos_sin_cache_ptr, \ + position_ids_ptr, \ + num_tokens, \ + rotary_dim)); \ + }); \ + } else { \ + queue.submit([&](sycl::handler& cgh) { \ + cgh.parallel_for( \ + sycl::nd_range<1>(grid_size * block_size, block_size), \ + fused_qk_norm_rope_kernel( \ + qkv_ptr, \ + num_heads_q, \ + num_heads_k, \ + num_heads_v, \ + eps, \ + q_weight_ptr, \ + k_weight_ptr, \ + cos_sin_cache_ptr, \ + position_ids_ptr, \ + num_tokens, \ + rotary_dim)); \ + }); \ + } \ + } while (0) + + switch (head_dim) { + case 64: + DISPATCH_HEAD_DIM(64); + break; + case 128: + DISPATCH_HEAD_DIM(128); + break; + case 256: + DISPATCH_HEAD_DIM(256); + break; + case 512: + DISPATCH_HEAD_DIM(512); + break; + default: + TORCH_CHECK( + false, + "Unsupported head dimension for fused_qk_norm_rope: ", + head_dim); + } + +#undef DISPATCH_HEAD_DIM +} + +} // namespace vllm + +void fused_qk_norm_rope( + torch::Tensor& qkv, + int64_t num_heads_q, + int64_t num_heads_k, + int64_t num_heads_v, + int64_t head_dim, + double eps, + torch::Tensor& q_weight, + torch::Tensor& k_weight, + torch::Tensor& cos_sin_cache, + bool is_neox, + torch::Tensor& position_ids) { + const at::DeviceGuard device_guard(qkv.device()); + + CHECK_DEVICE(qkv); + CHECK_CONTIGUOUS(qkv); + CHECK_DEVICE(position_ids); + CHECK_CONTIGUOUS(position_ids); + CHECK_DEVICE(q_weight); + CHECK_CONTIGUOUS(q_weight); + CHECK_DEVICE(k_weight); + CHECK_CONTIGUOUS(k_weight); + CHECK_DEVICE(cos_sin_cache); + CHECK_CONTIGUOUS(cos_sin_cache); + + TORCH_CHECK( + position_ids.scalar_type() == torch::kInt64, + "position_ids must be int64"); + TORCH_CHECK( + qkv.dim() == 2, + "QKV tensor must be 2D: [num_tokens, " + "(num_heads_q+num_heads_k+num_heads_v)*head_dim]"); + TORCH_CHECK(position_ids.dim() == 1, "Position IDs must be 1D: [num_tokens]"); + TORCH_CHECK(q_weight.dim() == 1, "Query weights must be 1D: [head_dim]"); + TORCH_CHECK(k_weight.dim() == 1, "Key weights must be 1D: [head_dim]"); + TORCH_CHECK( + cos_sin_cache.dim() == 2, + "Cos/sin cache must be 2D: [max_position, rotary_dim]"); + TORCH_CHECK( + q_weight.size(0) == head_dim, + "Query weights size must match head dimension"); + TORCH_CHECK( + k_weight.size(0) == head_dim, + "Key weights size must match head dimension"); + TORCH_CHECK(cos_sin_cache.size(1) % 2 == 0, "rotary_dim must be even"); + TORCH_CHECK( + cos_sin_cache.size(1) <= head_dim, + "rotary_dim must be less than or equal to head_dim"); + TORCH_CHECK( + qkv.scalar_type() == q_weight.scalar_type() && + qkv.scalar_type() == k_weight.scalar_type(), + "qkv, q_weight and k_weight must have the same dtype"); + + int64_t num_tokens = qkv.size(0); + TORCH_CHECK( + position_ids.size(0) == num_tokens, + "Number of tokens in position_ids must match QKV"); + + int64_t total_heads = num_heads_q + num_heads_k + num_heads_v; + TORCH_CHECK( + qkv.size(1) == total_heads * head_dim, + "QKV tensor size must match total number of heads and head dimension"); + + VLLM_DISPATCH_HALF_TYPES(qkv.scalar_type(), "fused_qk_norm_rope", [&] { + using qkv_scalar_t = scalar_t; + VLLM_DISPATCH_FLOATING_TYPES( + cos_sin_cache.scalar_type(), "fused_qk_norm_rope_cache", [&] { + using cache_scalar_t = scalar_t; + vllm::launch_fused_qk_norm_rope( + qkv, + static_cast(num_tokens), + static_cast(num_heads_q), + static_cast(num_heads_k), + static_cast(num_heads_v), + static_cast(head_dim), + static_cast(cos_sin_cache.size(1)), + static_cast(eps), + q_weight, + k_weight, + cos_sin_cache, + is_neox, + position_ids); + }); + }); +} diff --git a/csrc/layernorm.cpp b/csrc/layernorm.cpp index 8af7d8202..745b0e7c7 100644 --- a/csrc/layernorm.cpp +++ b/csrc/layernorm.cpp @@ -1,6 +1,7 @@ #include #include +#include #include "utils.h" #include "dispatch_utils.h" @@ -331,6 +332,7 @@ void rms_norm( torch::Tensor& input, torch::Tensor& weight, double epsilon) { + const at::DeviceGuard device_guard(input.device()); TORCH_CHECK(out.is_contiguous()); if (input.stride(-1) != 1) { input = input.contiguous(); @@ -348,6 +350,7 @@ void fused_add_rms_norm( torch::Tensor& residual, torch::Tensor& weight, double epsilon) { + const at::DeviceGuard device_guard(input.device()); int hidden_size = input.size(-1); int num_tokens = input.numel() / hidden_size; diff --git a/csrc/layernorm_quant.cpp b/csrc/layernorm_quant.cpp new file mode 100644 index 000000000..df9719434 --- /dev/null +++ b/csrc/layernorm_quant.cpp @@ -0,0 +1,866 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +#include +#include +#include + +#include +#include + +#include "dispatch_utils.h" +#include "ops.h" +#include "utils.h" +#include "quantization/fp8/quant_utils.h" +#include + +namespace vllm { + +template +struct alignas(8) vec4_t { + scalar_t val[4]; +}; + +template +class rms_norm_dynamic_per_token_quant_kernel { + public: + rms_norm_dynamic_per_token_quant_kernel( + out_t* __restrict__ out_, + scalar_t* __restrict__ residual_, + const scalar_t* __restrict__ input_, + const scalar_t* __restrict__ weight_, + const float* __restrict__ scale_ub_, + float* __restrict__ scales_, + const float epsilon_, + const int hidden_size_) + : out(out_), + residual(residual_), + input(input_), + weight(weight_), + scale_ub(scale_ub_), + scales(scales_), + epsilon(epsilon_), + hidden_size(hidden_size_) {} + + void operator()(sycl::nd_item<1> item) const { + const int tid = item.get_local_id(0); + const int local_range = item.get_local_range(0); + const int64_t token_idx = item.get_group(0); + + // s_local[0] = inv_rms, s_local[1] = scale + auto& s_local = + *sycl::ext::oneapi::group_local_memory_for_overwrite( + item.get_group()); + + const scalar_t* token_input = input + token_idx * hidden_size; + out_t* token_output = out + token_idx * hidden_size; + scalar_t* token_residual = nullptr; + if constexpr (has_residual) { + token_residual = residual + token_idx * hidden_size; + } + + // Pass 1: optional residual add + compute variance + float variance = 0.0f; + for (int i = tid; i < hidden_size; i += local_range) { + float x = static_cast(token_input[i]); + if constexpr (has_residual) { + x += static_cast(token_residual[i]); + token_residual[i] = static_cast(x); + // Read back the dtype-rounded value so variance is consistent with + // what pass 2 will use when reading token_residual. + x = static_cast(token_residual[i]); + } + variance += x * x; + } + variance = sycl::reduce_over_group( + item.get_group(), variance, sycl::plus()); + if (tid == 0) { + s_local[0] = sycl::rsqrt(variance / hidden_size + epsilon); + } + sycl::group_barrier(item.get_group()); + const float inv_rms = s_local[0]; + + // Pass 2: compute max |norm(x)| across the row → token scale + float absmax = 0.0f; + for (int i = tid; i < hidden_size; i += local_range) { + const float x = has_residual ? static_cast(token_residual[i]) + : static_cast(token_input[i]); + const float norm_x = x * inv_rms * static_cast(weight[i]); + absmax = sycl::max(absmax, sycl::fabs(norm_x)); + } + absmax = sycl::reduce_over_group( + item.get_group(), absmax, sycl::maximum()); + + if (tid == 0) { + float computed_scale; + if constexpr (std::is_same_v) { + computed_scale = (absmax > 0.0f) ? (absmax / 127.0f) : 1.0f; + } else { + // FP8 + const float fp8_max = static_cast(fp8::quant_type_max_v); + const float clamped_absmax = + (scale_ub != nullptr) ? sycl::min(absmax, *scale_ub) : absmax; + computed_scale = sycl::max( + clamped_absmax / fp8_max, fp8::min_scaling_factor::val()); + } + s_local[1] = computed_scale; + scales[token_idx] = computed_scale; + } + sycl::group_barrier(item.get_group()); + const float inv_scale = 1.0f / s_local[1]; + + // Pass 3: normalize and quantize + for (int i = tid; i < hidden_size; i += local_range) { + const float x = has_residual ? static_cast(token_residual[i]) + : static_cast(token_input[i]); + const float norm_x = x * inv_rms * static_cast(weight[i]); + const float q = norm_x * inv_scale; + + if constexpr (std::is_same_v) { + token_output[i] = static_cast( + sycl::max(sycl::min(sycl::rint(q), 127.0f), -128.0f)); + } else { + // FP8 + const float fp8_max = static_cast(fp8::quant_type_max_v); + token_output[i] = + static_cast(sycl::max(sycl::min(q, fp8_max), -fp8_max)); + } + } + } + + private: + out_t* __restrict__ out; + scalar_t* __restrict__ residual; + const scalar_t* __restrict__ input; + const scalar_t* __restrict__ weight; + const float* __restrict__ scale_ub; + float* __restrict__ scales; + const float epsilon; + const int hidden_size; +}; + +template +class rms_norm_per_block_quant_kernel { + public: + rms_norm_per_block_quant_kernel( + out_t* __restrict__ out_, + scalar_t* __restrict__ residual_, + const scalar_t* __restrict__ input_, + const scalar_t* __restrict__ weight_, + float* __restrict__ scales_, + const float epsilon_, + const int hidden_size_, + const int group_size_, + const int num_tokens_, + const int64_t scale_stride_token_, + const int64_t scale_stride_group_) + : out(out_), + residual(residual_), + input(input_), + weight(weight_), + scales(scales_), + epsilon(epsilon_), + hidden_size(hidden_size_), + group_size(group_size_), + num_tokens(num_tokens_), + scale_stride_token(scale_stride_token_), + scale_stride_group(scale_stride_group_) {} + + void operator()(sycl::nd_item<1> item) const { + const int tid = item.get_local_id(0); + const int local_range = item.get_local_range(0); + const int64_t token_idx = item.get_group(0); + + // s_local[0] = inv_rms, s_local[1] = current group scale + auto& s_local = + *sycl::ext::oneapi::group_local_memory_for_overwrite( + item.get_group()); + + const scalar_t* token_input = input + token_idx * hidden_size; + out_t* token_output = out + token_idx * hidden_size; + scalar_t* token_residual = nullptr; + if constexpr (has_residual) { + token_residual = residual + token_idx * hidden_size; + } + + // Pass 1: optional residual add + compute full-row variance + float variance = 0.0f; + for (int i = tid; i < hidden_size; i += local_range) { + float x = static_cast(token_input[i]); + if constexpr (has_residual) { + x += static_cast(token_residual[i]); + token_residual[i] = static_cast(x); + // Read back the dtype-rounded value so variance is consistent with + // what passes 2+3 will use when reading token_residual. + x = static_cast(token_residual[i]); + } + variance += x * x; + } + variance = sycl::reduce_over_group( + item.get_group(), variance, sycl::plus()); + if (tid == 0) { + s_local[0] = sycl::rsqrt(variance / hidden_size + epsilon); + } + sycl::group_barrier(item.get_group()); + const float inv_rms = s_local[0]; + + // Pass 2+3: for each column group, compute scale then quantize + const int num_groups = hidden_size / group_size; + for (int g = 0; g < num_groups; g++) { + const int group_start = g * group_size; + + // Find max |norm(x)| within this group + float group_absmax = 0.0f; + for (int i = tid; i < group_size; i += local_range) { + const int col = group_start + i; + const float x = has_residual ? static_cast(token_residual[col]) + : static_cast(token_input[col]); + const float norm_x = x * inv_rms * static_cast(weight[col]); + group_absmax = sycl::max(group_absmax, sycl::fabs(norm_x)); + } + group_absmax = sycl::reduce_over_group( + item.get_group(), group_absmax, sycl::maximum()); + + if (tid == 0) { + float group_scale; + if constexpr (std::is_same_v) { + group_scale = (group_absmax > 0.0f) ? (group_absmax / 127.0f) : 1.0f; + } else { + // FP8 + const float fp8_max = + static_cast(fp8::quant_type_max_v); + group_scale = sycl::max( + group_absmax / fp8_max, fp8::min_scaling_factor::val()); + } + s_local[1] = group_scale; + const int64_t scale_idx = token_idx * scale_stride_token + + static_cast(g) * scale_stride_group; + scales[scale_idx] = group_scale; + } + sycl::group_barrier(item.get_group()); + const float inv_scale = 1.0f / s_local[1]; + + // Quantize this group + for (int i = tid; i < group_size; i += local_range) { + const int col = group_start + i; + const float x = has_residual ? static_cast(token_residual[col]) + : static_cast(token_input[col]); + const float norm_x = x * inv_rms * static_cast(weight[col]); + const float q = norm_x * inv_scale; + + if constexpr (std::is_same_v) { + token_output[col] = static_cast( + sycl::max(sycl::min(sycl::rint(q), 127.0f), -128.0f)); + } else { + // FP8 + const float fp8_max = + static_cast(fp8::quant_type_max_v); + token_output[col] = + static_cast(sycl::max(sycl::min(q, fp8_max), -fp8_max)); + } + } + sycl::group_barrier(item.get_group()); + } + } + + private: + out_t* __restrict__ out; + scalar_t* __restrict__ residual; + const scalar_t* __restrict__ input; + const scalar_t* __restrict__ weight; + float* __restrict__ scales; + const float epsilon; + const int hidden_size; + const int group_size; + const int num_tokens; + const int64_t scale_stride_token; + const int64_t scale_stride_group; +}; + +template +class rms_norm_static_fp8_quant_kernel { + public: + rms_norm_static_fp8_quant_kernel( + out_t* __restrict__ out_, + const scalar_t* __restrict__ input_, + const int input_stride_, + const scalar_t* __restrict__ weight_, + const float* __restrict__ scale_, + const float epsilon_, + const int hidden_size_) + : out(out_), + input(input_), + input_stride(input_stride_), + weight(weight_), + scale(scale_), + epsilon(epsilon_), + hidden_size(hidden_size_) {} + + void operator()(sycl::nd_item<1> item) const { + using vec_t = vllm::xpu::aligned_vec; + + const int tid = item.get_local_id(0); + const int local_range = item.get_local_range(0); + const int64_t token_idx = item.get_group(0); + + auto& s_variance = + *sycl::ext::oneapi::group_local_memory_for_overwrite( + item.get_group()); + + const scalar_t* token_input = input + token_idx * input_stride; + out_t* token_output = out + token_idx * hidden_size; + const int nvec = hidden_size / VEC_SIZE; + + // Pass 1: compute variance — VEC_SIZE elements per work-item per iteration + float variance = 0.0f; + const auto* v_in = reinterpret_cast(token_input); + for (int i = tid; i < nvec; i += local_range) { + vec_t v = v_in[i]; +#pragma unroll + for (int j = 0; j < VEC_SIZE; j++) { + float x = static_cast(v[j]); + variance += x * x; + } + } + variance = sycl::reduce_over_group( + item.get_group(), variance, sycl::plus()); + if (tid == 0) { + s_variance = sycl::rsqrt(variance / hidden_size + epsilon); + } + sycl::group_barrier(item.get_group()); + const float inv_rms = s_variance; + + // Invert scale to avoid division + const float scale_inv = 1.0f / (*scale); + const float fp8_max = static_cast(fp8::quant_type_max_v); + + // Pass 2: normalize, apply weight, quantize — same vectorization as Pass 1 + const auto* v_w = reinterpret_cast(weight); + for (int i = tid; i < nvec; i += local_range) { + vec_t src = v_in[i]; + vec_t wgt = v_w[i]; +#pragma unroll + for (int j = 0; j < VEC_SIZE; j++) { + float x = static_cast(src[j]); + // Weight multiply in scalar_t precision to match unfused path + float norm_x = + static_cast(static_cast(x * inv_rms) * wgt[j]); + float q = norm_x * scale_inv; + token_output[i * VEC_SIZE + j] = + static_cast(sycl::max(sycl::min(q, fp8_max), -fp8_max)); + } + } + } + + private: + out_t* __restrict__ out; + const scalar_t* __restrict__ input; + const int input_stride; + const scalar_t* __restrict__ weight; + const float* __restrict__ scale; + const float epsilon; + const int hidden_size; +}; + +template +class fused_add_rms_norm_static_fp8_quant_kernel { + public: + fused_add_rms_norm_static_fp8_quant_kernel( + out_t* __restrict__ out_, + const scalar_t* __restrict__ input_, + const int input_stride_, + scalar_t* __restrict__ residual_, + const scalar_t* __restrict__ weight_, + const float* __restrict__ scale_, + const float epsilon_, + const int hidden_size_) + : out(out_), + input(input_), + input_stride(input_stride_), + residual(residual_), + weight(weight_), + scale(scale_), + epsilon(epsilon_), + hidden_size(hidden_size_) {} + + void operator()(sycl::nd_item<1> item) const { + using vec_t = vllm::xpu::aligned_vec; + + const int tid = item.get_local_id(0); + const int local_range = item.get_local_range(0); + const int64_t token_idx = item.get_group(0); + + auto& s_variance = + *sycl::ext::oneapi::group_local_memory_for_overwrite( + item.get_group()); + + const scalar_t* token_input = input + token_idx * input_stride; + scalar_t* token_residual = residual + token_idx * hidden_size; + out_t* token_output = out + token_idx * hidden_size; + const int nvec = hidden_size / VEC_SIZE; + + // Pass 1: add residual + compute variance, VEC_SIZE elements per iteration + float variance = 0.0f; + const auto* v_in = reinterpret_cast(token_input); + auto* v_res = reinterpret_cast(token_residual); + for (int i = tid; i < nvec; i += local_range) { + vec_t inp = v_in[i]; + vec_t res = v_res[i]; +#pragma unroll + for (int j = 0; j < VEC_SIZE; j++) { + // Add in scalar_t precision to match fused_add_rms_norm kernel + scalar_t z = inp[j] + res[j]; + res[j] = z; + float xf = static_cast(z); + variance += xf * xf; + } + v_res[i] = res; // write updated residual back + } + variance = sycl::reduce_over_group( + item.get_group(), variance, sycl::plus()); + if (tid == 0) { + s_variance = sycl::rsqrt(variance / hidden_size + epsilon); + } + sycl::group_barrier(item.get_group()); + const float inv_rms = s_variance; + + // Invert scale to avoid division + const float scale_inv = 1.0f / (*scale); + const float fp8_max = static_cast(fp8::quant_type_max_v); + + // Pass 2: normalize from residual, apply weight, quantize + const auto* v_w = reinterpret_cast(weight); + for (int i = tid; i < nvec; i += local_range) { + vec_t res = v_res[i]; + vec_t wgt = v_w[i]; +#pragma unroll + for (int j = 0; j < VEC_SIZE; j++) { + float x = static_cast(res[j]); + // Weight multiply in scalar_t precision to match unfused path + float norm_x = + static_cast(static_cast(x * inv_rms) * wgt[j]); + float q = norm_x * scale_inv; + token_output[i * VEC_SIZE + j] = + static_cast(sycl::max(sycl::min(q, fp8_max), -fp8_max)); + } + } + } + + private: + out_t* __restrict__ out; + const scalar_t* __restrict__ input; + const int input_stride; + scalar_t* __restrict__ residual; + const scalar_t* __restrict__ weight; + const float* __restrict__ scale; + const float epsilon; + const int hidden_size; +}; + +template +void call_rms_norm_dynamic_per_token_quant_kernel( + torch::Tensor& out, + std::optional& residual, + torch::Tensor const& input, + torch::Tensor const& weight, + std::optional const& scale_ub, + torch::Tensor& scales, + float epsilon) { + using sycl_t = typename vllm::xpu::SyclTypeTrait::Type; + + const int hidden_size = input.size(-1); + const int64_t num_tokens = input.numel() / hidden_size; + const int block_size = std::min(hidden_size, 1024); + + auto* out_ptr = out.data_ptr(); + auto* input_ptr = input.data_ptr(); + auto* weight_ptr = weight.data_ptr(); + auto* scales_ptr = scales.data_ptr(); + const float* scale_ub_ptr = + scale_ub.has_value() ? scale_ub->data_ptr() : nullptr; + + auto& queue = vllm::xpu::vllmGetQueue(); + + auto launch = [&](auto has_residual_tag) { + constexpr bool has_residual = decltype(has_residual_tag)::value; + sycl_t* residual_ptr = nullptr; + if constexpr (has_residual) { + residual_ptr = (sycl_t*)residual->data_ptr(); + } + queue.submit([&](sycl::handler& cgh) { + cgh.parallel_for( + sycl::nd_range<1>(num_tokens * block_size, block_size), + rms_norm_dynamic_per_token_quant_kernel( + out_ptr, + residual_ptr, + (const sycl_t*)input_ptr, + (const sycl_t*)weight_ptr, + scale_ub_ptr, + scales_ptr, + epsilon, + hidden_size)); + }); + }; + + if (residual.has_value()) { + launch(std::true_type{}); + } else { + launch(std::false_type{}); + } +} + +template +void call_rms_norm_per_block_quant_kernel( + torch::Tensor& out, + std::optional& residual, + torch::Tensor const& input, + torch::Tensor const& weight, + torch::Tensor& scales, + float epsilon, + int group_size, + bool is_scale_transposed) { + using sycl_t = typename vllm::xpu::SyclTypeTrait::Type; + + const int hidden_size = input.size(-1); + const int64_t num_tokens = input.numel() / hidden_size; + const int num_groups = hidden_size / group_size; + const int block_size = std::min(hidden_size, 1024); + + // Compute scale strides based on transposition + const int64_t scale_stride_token = + is_scale_transposed ? 1 : static_cast(num_groups); + const int64_t scale_stride_group = is_scale_transposed ? num_tokens : 1LL; + + auto* out_ptr = out.data_ptr(); + auto* input_ptr = input.data_ptr(); + auto* weight_ptr = weight.data_ptr(); + auto* scales_ptr = scales.data_ptr(); + + auto& queue = vllm::xpu::vllmGetQueue(); + + auto launch = [&](auto has_residual_tag) { + constexpr bool has_residual = decltype(has_residual_tag)::value; + sycl_t* residual_ptr = nullptr; + if constexpr (has_residual) { + residual_ptr = (sycl_t*)residual->data_ptr(); + } + queue.submit([&](sycl::handler& cgh) { + cgh.parallel_for( + sycl::nd_range<1>(num_tokens * block_size, block_size), + rms_norm_per_block_quant_kernel( + out_ptr, + residual_ptr, + (const sycl_t*)input_ptr, + (const sycl_t*)weight_ptr, + scales_ptr, + epsilon, + hidden_size, + group_size, + static_cast(num_tokens), + scale_stride_token, + scale_stride_group)); + }); + }; + + if (residual.has_value()) { + launch(std::true_type{}); + } else { + launch(std::false_type{}); + } +} + +template +void call_rms_norm_static_fp8_quant_kernel( + torch::Tensor& out, + torch::Tensor const& input, + torch::Tensor const& weight, + torch::Tensor const& scale, + float epsilon) { + using sycl_t = typename vllm::xpu::SyclTypeTrait::Type; + + const int hidden_size = input.size(-1); + const int input_stride = input.stride(-2); + const int64_t num_tokens = input.numel() / hidden_size; + + // Match CUDA: smaller blocks when num_tokens is large for better occupancy + const int max_block_size = (num_tokens < 256) ? 1024 : 256; + + // Dispatch VEC_SIZE = gcd(16 / sizeof(sycl_t), hidden_size) matching CUDA + const int candidate_vec_size = + std::gcd(static_cast(16 / sizeof(sycl_t)), hidden_size); + const int vec_size = + (input_stride % candidate_vec_size == 0) ? candidate_vec_size : 1; + const int block_size = std::min(hidden_size / vec_size, max_block_size); + + auto& queue = vllm::xpu::vllmGetQueue(); + + auto launch = [&](auto vec_tag) { + constexpr int VS = decltype(vec_tag)::value; + queue.submit([&](sycl::handler& cgh) { + cgh.parallel_for( + sycl::nd_range<1>(num_tokens * block_size, block_size), + rms_norm_static_fp8_quant_kernel( + out.data_ptr(), + (const sycl_t*)input.data_ptr(), + input_stride, + (const sycl_t*)weight.data_ptr(), + scale.data_ptr(), + epsilon, + hidden_size)); + }); + }; + + // Dispatch on vec_size (gcd guarantees hidden_size % vec_size == 0) + switch (vec_size) { + case 8: + launch(std::integral_constant{}); + break; + case 4: + launch(std::integral_constant{}); + break; + case 2: + launch(std::integral_constant{}); + break; + default: + launch(std::integral_constant{}); + break; + } +} + +template +void call_fused_add_rms_norm_static_fp8_quant_kernel( + torch::Tensor& out, + torch::Tensor& input, + torch::Tensor& residual, + torch::Tensor const& weight, + torch::Tensor const& scale, + float epsilon) { + using sycl_t = typename vllm::xpu::SyclTypeTrait::Type; + + const int hidden_size = input.size(-1); + const int input_stride = input.stride(-2); + const int64_t num_tokens = input.numel() / hidden_size; + + // Match CUDA: smaller blocks when num_tokens is large for better occupancy + const int max_block_size = (num_tokens < 256) ? 1024 : 256; + + // Dispatch VEC_SIZE = gcd(16 / sizeof(sycl_t), hidden_size) matching CUDA. + // Also require input_stride % vec_size == 0 for safe vectorized input reads. + int vec_size = std::gcd(static_cast(16 / sizeof(sycl_t)), hidden_size); + if (input_stride % vec_size != 0) { + vec_size = 1; + } + const int block_size = std::min(hidden_size / vec_size, max_block_size); + + auto& queue = vllm::xpu::vllmGetQueue(); + + auto launch = [&](auto vec_tag) { + constexpr int VS = decltype(vec_tag)::value; + queue.submit([&](sycl::handler& cgh) { + cgh.parallel_for( + sycl::nd_range<1>(num_tokens * block_size, block_size), + fused_add_rms_norm_static_fp8_quant_kernel( + out.data_ptr(), + (const sycl_t*)input.data_ptr(), + input_stride, + (sycl_t*)residual.data_ptr(), + (const sycl_t*)weight.data_ptr(), + scale.data_ptr(), + epsilon, + hidden_size)); + }); + }; + + // Dispatch on vec_size (gcd guarantees hidden_size % vec_size == 0) + switch (vec_size) { + case 8: + launch(std::integral_constant{}); + break; + case 4: + launch(std::integral_constant{}); + break; + case 2: + launch(std::integral_constant{}); + break; + default: + launch(std::integral_constant{}); + break; + } +} + +} // namespace vllm + +void rms_norm_dynamic_per_token_quant( + torch::Tensor& out, + torch::Tensor const& input, + torch::Tensor const& weight, + torch::Tensor& scales, + double const epsilon, + std::optional scale_ub, + std::optional residual) { + const at::DeviceGuard device_guard(input.device()); + TORCH_CHECK(input.is_contiguous(), "input must be contiguous"); + TORCH_CHECK(out.is_contiguous(), "out must be contiguous"); + TORCH_CHECK( + weight.dtype() == input.dtype(), + "weight and input must have the same dtype"); + TORCH_CHECK(scales.dtype() == torch::kFloat32, "scales must be float32"); + TORCH_CHECK( + out.dtype() == torch::kFloat8_e4m3fn || out.dtype() == torch::kInt8, + "output must be float8_e4m3fn or int8"); + if (scale_ub.has_value()) { + TORCH_CHECK( + out.dtype() == torch::kFloat8_e4m3fn, + "scale_ub is only supported for FP8 output"); + } + if (residual.has_value()) { + TORCH_CHECK( + residual->scalar_type() == input.scalar_type(), + "residual and input must have the same dtype"); + TORCH_CHECK(residual->is_contiguous(), "residual must be contiguous"); + } + + if (out.dtype() == torch::kFloat8_e4m3fn) { + VLLM_DISPATCH_FLOATING_TYPES( + input.scalar_type(), "rms_norm_dynamic_per_token_quant", [&] { + vllm::call_rms_norm_dynamic_per_token_quant_kernel< + scalar_t, + at::Float8_e4m3fn>( + out, + residual, + input, + weight, + scale_ub, + scales, + static_cast(epsilon)); + }); + } else { + VLLM_DISPATCH_FLOATING_TYPES( + input.scalar_type(), "rms_norm_dynamic_per_token_quant_int8", [&] { + vllm::call_rms_norm_dynamic_per_token_quant_kernel( + out, + residual, + input, + weight, + scale_ub, + scales, + static_cast(epsilon)); + }); + } +} + +void rms_norm_per_block_quant( + torch::Tensor& out, + torch::Tensor const& input, + torch::Tensor const& weight, + torch::Tensor& scales, + double const epsilon, + std::optional scale_ub, + std::optional residual, + int64_t group_size, + bool is_scale_transposed) { + const at::DeviceGuard device_guard(input.device()); + TORCH_CHECK(input.is_contiguous(), "input must be contiguous"); + TORCH_CHECK(out.is_contiguous(), "out must be contiguous"); + TORCH_CHECK( + weight.dtype() == input.dtype(), + "weight and input must have the same dtype"); + TORCH_CHECK(scales.dtype() == torch::kFloat32, "scales must be float32"); + TORCH_CHECK( + out.dtype() == torch::kFloat8_e4m3fn || out.dtype() == torch::kInt8, + "output must be float8_e4m3fn or int8"); + TORCH_CHECK( + input.size(-1) % group_size == 0, + "hidden_size must be divisible by group_size"); + if (residual.has_value()) { + TORCH_CHECK( + residual->scalar_type() == input.scalar_type(), + "residual and input must have the same dtype"); + TORCH_CHECK(residual->is_contiguous(), "residual must be contiguous"); + } + + if (out.dtype() == torch::kFloat8_e4m3fn) { + VLLM_DISPATCH_FLOATING_TYPES( + input.scalar_type(), "rms_norm_per_block_quant", [&] { + vllm:: + call_rms_norm_per_block_quant_kernel( + out, + residual, + input, + weight, + scales, + static_cast(epsilon), + static_cast(group_size), + is_scale_transposed); + }); + } else { + VLLM_DISPATCH_FLOATING_TYPES( + input.scalar_type(), "rms_norm_per_block_quant_int8", [&] { + vllm::call_rms_norm_per_block_quant_kernel( + out, + residual, + input, + weight, + scales, + static_cast(epsilon), + static_cast(group_size), + is_scale_transposed); + }); + } +} + +void rms_norm_static_fp8_quant( + torch::Tensor& out, + torch::Tensor& input, + torch::Tensor& weight, + torch::Tensor& scale, + double epsilon) { + const at::DeviceGuard device_guard(input.device()); + TORCH_CHECK(out.is_contiguous(), "out must be contiguous"); + TORCH_CHECK( + weight.dtype() == input.dtype(), + "weight and input must have the same dtype"); + + VLLM_DISPATCH_FLOATING_TYPES( + input.scalar_type(), "rms_norm_static_fp8_quant", [&] { + VLLM_DISPATCH_FP8_TYPES( + out.scalar_type(), "rms_norm_static_fp8_quant_fp8", [&] { + vllm::call_rms_norm_static_fp8_quant_kernel( + out, input, weight, scale, static_cast(epsilon)); + }); + }); +} + +void fused_add_rms_norm_static_fp8_quant( + torch::Tensor& out, + torch::Tensor& input, + torch::Tensor& residual, + torch::Tensor& weight, + torch::Tensor& scale, + double epsilon) { + const at::DeviceGuard device_guard(input.device()); + TORCH_CHECK(out.is_contiguous(), "out must be contiguous"); + TORCH_CHECK(residual.is_contiguous(), "residual must be contiguous"); + TORCH_CHECK( + residual.scalar_type() == input.scalar_type(), + "residual and input must have the same dtype"); + TORCH_CHECK( + weight.scalar_type() == input.scalar_type(), + "weight and input must have the same dtype"); + + VLLM_DISPATCH_FLOATING_TYPES( + input.scalar_type(), "fused_add_rms_norm_static_fp8_quant", [&] { + VLLM_DISPATCH_FP8_TYPES( + out.scalar_type(), "fused_add_rms_norm_static_fp8_quant_fp8", [&] { + vllm::call_fused_add_rms_norm_static_fp8_quant_kernel< + scalar_t, + fp8_t>( + out, + input, + residual, + weight, + scale, + static_cast(epsilon)); + }); + }); +} diff --git a/csrc/moe/remap_hidden_states.cpp b/csrc/moe/remap_hidden_states.cpp index 4226ac711..9e884336d 100644 --- a/csrc/moe/remap_hidden_states.cpp +++ b/csrc/moe/remap_hidden_states.cpp @@ -215,14 +215,16 @@ class RemapHiddenStates { item.barrier(sycl::access::fence_space::local_space); - auto hidden_states_base = - hidden_states + row * hidden_size + local_id * ElemsPerItem; + auto hidden_states_base = hidden_states + + row * static_cast(hidden_size) + + local_id * ElemsPerItem; TA* remapped_hidden_states_base[TopK]; #pragma unroll for (int i = 0; i < TopK; ++i) { - remapped_hidden_states_base[i] = remapped_hidden_states + - rows_offset[i] * hidden_size + - local_id * ElemsPerItem; + remapped_hidden_states_base[i] = + remapped_hidden_states + + rows_offset[i] * static_cast(hidden_size) + + local_id * ElemsPerItem; } int stride = local_range * ElemsPerItem; @@ -244,7 +246,7 @@ class RemapHiddenStates { if (hidden_states_scales != nullptr && remapped_hidden_states_scales != nullptr) { - int scaled_hidden_size = hidden_size / block_k; + int64_t scaled_hidden_size = hidden_size / block_k; loop_count = (scaled_hidden_size + stride - 1) / stride; for (int l = 0; l < loop_count; ++l) { diff --git a/csrc/ops.h b/csrc/ops.h index 25d38e67d..6618008ad 100644 --- a/csrc/ops.h +++ b/csrc/ops.h @@ -16,14 +16,57 @@ void fused_add_rms_norm( torch::Tensor& weight, double epsilon); +// Fused RMSNorm + dynamic per-token quantization (FP8 or INT8 output). +void rms_norm_dynamic_per_token_quant( + torch::Tensor& out, + torch::Tensor const& input, + torch::Tensor const& weight, + torch::Tensor& scales, + double const epsilon, + std::optional scale_ub, + std::optional residual); + +// Fused RMSNorm + per-column-block quantization (FP8 or INT8 output). +void rms_norm_per_block_quant( + torch::Tensor& out, + torch::Tensor const& input, + torch::Tensor const& weight, + torch::Tensor& scales, + double const epsilon, + std::optional scale_ub, + std::optional residual, + int64_t group_size, + bool is_scale_transposed); + +void rms_norm_static_fp8_quant( + torch::Tensor& out, + torch::Tensor& input, + torch::Tensor& weight, + torch::Tensor& scale, + double epsilon); + +void fused_add_rms_norm_static_fp8_quant( + torch::Tensor& out, + torch::Tensor& input, + torch::Tensor& residual, + torch::Tensor& weight, + torch::Tensor& scale, + double epsilon); + void silu_and_mul(torch::Tensor& out, torch::Tensor& input); +void silu_and_mul_quant( + torch::Tensor& out, torch::Tensor& input, torch::Tensor& scale); + void mul_and_silu(torch::Tensor& out, torch::Tensor& input); void gelu_and_mul(torch::Tensor& out, torch::Tensor& input); void gelu_tanh_and_mul(torch::Tensor& out, torch::Tensor& input); +void fatrelu_and_mul( + torch::Tensor& out, torch::Tensor& input, double threshold); + void gelu_fast(torch::Tensor& out, torch::Tensor& input); void gelu_new(torch::Tensor& out, torch::Tensor& input); @@ -38,6 +81,19 @@ void rotary_embedding( torch::Tensor& cos_sin_cache, bool is_neox); +void fused_qk_norm_rope( + torch::Tensor& qkv, + int64_t num_heads_q, + int64_t num_heads_k, + int64_t num_heads_v, + int64_t head_dim, + double eps, + torch::Tensor& q_weight, + torch::Tensor& k_weight, + torch::Tensor& cos_sin_cache, + bool is_neox, + torch::Tensor& position_ids); + void reshape_and_cache( torch::Tensor& key, torch::Tensor& value, @@ -158,6 +214,11 @@ void swap_blocks( int64_t block_size_in_bytes, const torch::Tensor& block_mapping); +void swap_blocks_batch( + const torch::Tensor& src_ptrs, + const torch::Tensor& dst_ptrs, + const torch::Tensor& sizes); + void top_k_per_row_decode( const torch::Tensor& logits, int64_t next_n, diff --git a/csrc/quantization/fp4/mxfp4_quant.cpp b/csrc/quantization/fp4/mxfp4_quant.cpp index fc0604fb9..e3c9e7138 100644 --- a/csrc/quantization/fp4/mxfp4_quant.cpp +++ b/csrc/quantization/fp4/mxfp4_quant.cpp @@ -48,8 +48,7 @@ void per_token_group_quant_mxfp4( return; } - at::Device curDevice = at::Device(at::kXPU, at::xpu::current_device()); - at::DeviceGuard device_guard(curDevice); + const at::DeviceGuard device_guard(input.device()); // Choose how many sub-groups to pack into one work-group. constexpr int THREADS_PER_GROUP = 32; diff --git a/csrc/quantization/fp8/fp8_quant.cpp b/csrc/quantization/fp8/fp8_quant.cpp index f5aa26b19..698726edf 100644 --- a/csrc/quantization/fp8/fp8_quant.cpp +++ b/csrc/quantization/fp8/fp8_quant.cpp @@ -504,8 +504,7 @@ void static_scaled_fp8_quant( const int64_t in_row_stride = input.stride(-2); const int64_t out_row_stride = out.stride(-2); - at::Device curDevice = at::Device(at::kXPU, at::xpu::current_device()); - at::DeviceGuard device_guard(curDevice); + const at::DeviceGuard device_guard(input.device()); auto& queue = vllm::xpu::vllmGetQueue(); VLLM_DISPATCH_FLOATING_TYPES( @@ -562,8 +561,7 @@ void dynamic_scaled_fp8_quant( int64_t in_row_stride = input.stride(-2); int64_t out_row_stride = out.stride(-2); - at::Device curDevice = at::Device(at::kXPU, at::xpu::current_device()); - at::DeviceGuard device_guard(curDevice); + const at::DeviceGuard device_guard(input.device()); auto& queue = vllm::xpu::vllmGetQueue(); VLLM_DISPATCH_FLOATING_TYPES( @@ -616,8 +614,7 @@ void per_token_group_quant_fp8( TORCH_CHECK(input.numel() % group_size == 0); TORCH_CHECK(output_s.dim() == 2); - at::Device curDevice = at::Device(at::kXPU, at::xpu::current_device()); - at::DeviceGuard device_guard(curDevice); + const at::DeviceGuard device_guard(input.device()); constexpr int THREADS_PER_GROUP = 32; @@ -684,8 +681,7 @@ void dynamic_per_token_scaled_fp8_quant( sycl::range<1> grid(num_tokens); sycl::range<1> block(std::min(hidden_size, 1024)); - at::Device curDevice = at::Device(at::kXPU, at::xpu::current_device()); - at::DeviceGuard device_guard(curDevice); + const at::DeviceGuard device_guard(input.device()); auto& queue = vllm::xpu::vllmGetQueue(); VLLM_DISPATCH_FLOATING_TYPES( diff --git a/csrc/torch_bindings.cpp b/csrc/torch_bindings.cpp index 04e3480fe..dc5847c11 100644 --- a/csrc/torch_bindings.cpp +++ b/csrc/torch_bindings.cpp @@ -21,9 +21,12 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { // Layernorm // Apply Root Mean Square (RMS) Normalization to the input tensor. + // FIXME: torch op check consider input & weight is mutable in some ut + // cases. so we make it mutable here. ops.def( "rms_norm(Tensor! result, Tensor input, Tensor weight, float epsilon) " - "-> ()"); + "-> " + "()"); ops.impl("rms_norm", torch::kXPU, &rms_norm); // In-place fused Add and RMS Normalization. @@ -32,26 +35,69 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { "float epsilon) -> ()"); ops.impl("fused_add_rms_norm", torch::kXPU, &fused_add_rms_norm); + // Fused RMSNorm + dynamic per-token quantization (FP8 or INT8). + ops.def( + "rms_norm_dynamic_per_token_quant(Tensor! result, Tensor input, " + "Tensor weight, Tensor! scale, float epsilon, " + "Tensor? scale_ub, Tensor!? residual) -> ()"); + ops.impl( + "rms_norm_dynamic_per_token_quant", + torch::kXPU, + &rms_norm_dynamic_per_token_quant); + + // Fused RMSNorm + per-column-block quantization (FP8 or INT8). + ops.def( + "rms_norm_per_block_quant(Tensor! result, Tensor input, " + "Tensor weight, Tensor! scale, float epsilon, " + "Tensor? scale_ub, Tensor!? residual, int group_size, " + "bool is_scale_transposed) -> ()"); + ops.impl("rms_norm_per_block_quant", torch::kXPU, &rms_norm_per_block_quant); + + // Fused RMSNorm + static FP8 quantization. + ops.def( + "rms_norm_static_fp8_quant(Tensor! result, Tensor input, Tensor weight, " + "Tensor scale, float epsilon) -> ()"); + ops.impl( + "rms_norm_static_fp8_quant", torch::kXPU, &rms_norm_static_fp8_quant); + + // In-place fused Add + RMSNorm + static FP8 quantization. + ops.def( + "fused_add_rms_norm_static_fp8_quant(Tensor! result, Tensor input, " + "Tensor! residual, Tensor weight, " + "Tensor scale, float epsilon) -> ()"); + ops.impl( + "fused_add_rms_norm_static_fp8_quant", + torch::kXPU, + &fused_add_rms_norm_static_fp8_quant); + // activation ops - ops.def("silu_and_mul(Tensor! out, Tensor! input) -> ()"); + ops.def("silu_and_mul(Tensor! result, Tensor input) -> ()"); ops.impl("silu_and_mul", torch::kXPU, &silu_and_mul); - ops.def("mul_and_silu(Tensor! out, Tensor! input) -> ()"); + // Fused SiLU + Mul + FP8 Quantization + ops.def( + "silu_and_mul_quant(Tensor! result, Tensor input, Tensor scale) -> ()"); + ops.impl("silu_and_mul_quant", torch::kXPU, &silu_and_mul_quant); + + ops.def("mul_and_silu(Tensor! out, Tensor input) -> ()"); ops.impl("mul_and_silu", torch::kXPU, &mul_and_silu); - ops.def("gelu_and_mul(Tensor! out, Tensor! input) -> ()"); + ops.def("gelu_and_mul(Tensor! out, Tensor input) -> ()"); ops.impl("gelu_and_mul", torch::kXPU, &gelu_and_mul); - ops.def("gelu_tanh_and_mul(Tensor! out, Tensor! input) -> ()"); + ops.def("gelu_tanh_and_mul(Tensor! out, Tensor input) -> ()"); ops.impl("gelu_tanh_and_mul", torch::kXPU, &gelu_tanh_and_mul); - ops.def("gelu_fast(Tensor! out, Tensor! input) -> ()"); + ops.def("fatrelu_and_mul(Tensor! out, Tensor! input, float threshold) -> ()"); + ops.impl("fatrelu_and_mul", torch::kXPU, &fatrelu_and_mul); + + ops.def("gelu_fast(Tensor! out, Tensor input) -> ()"); ops.impl("gelu_fast", torch::kXPU, &gelu_fast); - ops.def("gelu_new(Tensor! out, Tensor! input) -> ()"); + ops.def("gelu_new(Tensor! out, Tensor input) -> ()"); ops.impl("gelu_new", torch::kXPU, &gelu_new); - ops.def("gelu_quick(Tensor! out, Tensor! input) -> ()"); + ops.def("gelu_quick(Tensor! out, Tensor input) -> ()"); ops.impl("gelu_quick", torch::kXPU, &gelu_quick); // pos_embedding @@ -61,6 +107,14 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { " Tensor cos_sin_cache, bool is_neox) -> ()"); ops.impl("rotary_embedding", torch::kXPU, &rotary_embedding); + // Fused QK RMSNorm + RoPE + ops.def( + "fused_qk_norm_rope(Tensor! qkv, int num_heads_q, " + "int num_heads_k, int num_heads_v, int head_dim, float eps, " + "Tensor q_weight, Tensor k_weight, Tensor cos_sin_cache, " + "bool is_neox, Tensor position_ids) -> ()"); + ops.impl("fused_qk_norm_rope", torch::kXPU, &fused_qk_norm_rope); + // Compute FP8 quantized tensor for given scaling factor. ops.def( "static_scaled_fp8_quant(Tensor! result, Tensor input, Tensor scale, " @@ -206,6 +260,12 @@ TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _cache_ops), cache_ops) { "swap_blocks(Tensor src, Tensor! dst," " int block_size_in_bytes, Tensor block_mapping) -> ()"); cache_ops.impl("swap_blocks", torch::kXPU, &swap_blocks); + // Batch swap: copies N (src_ptr, dst_ptr, size) triples in one call. + // The target XPU device is auto-inferred from the device pointer. + cache_ops.def( + "swap_blocks_batch(Tensor src_ptrs, Tensor dst_ptrs," + " Tensor sizes) -> ()"); + cache_ops.impl("swap_blocks_batch", torch::kCPU, &swap_blocks_batch); cache_ops.def( "indexer_k_quant_and_cache(Tensor k, Tensor! kv_cache," "Tensor slot_mapping, int quant_block_size, str scale_fmt) -> ()"); diff --git a/csrc/utils/mem_cpy.cpp b/csrc/utils/mem_cpy.cpp index 88998d9d4..661da04e6 100644 --- a/csrc/utils/mem_cpy.cpp +++ b/csrc/utils/mem_cpy.cpp @@ -151,6 +151,136 @@ void xpuAsyncMemcpy( } } +// Infer which XPU device a USM device pointer was allocated on by probing +// each device's SYCL context. Returns the device index on success. +// This is O(num_xpu_devices) but avoids threading an explicit device argument +// through the entire call chain when all callers already have the pointer. +static at::DeviceIndex infer_xpu_device_from_ptr(const void* device_ptr) { + const int n_devs = c10::xpu::device_count(); + for (int i = 0; i < n_devs; i++) { + auto ctx = vllm::xpu::vllmGetQueue(i).get_context(); + auto type = sycl::get_pointer_type(device_ptr, ctx); + if (type == sycl::usm::alloc::device || type == sycl::usm::alloc::shared) { + return static_cast(i); + } + } + TORCH_CHECK(false, "Cannot determine XPU device from pointer"); + return -1; +} + +void xpuAsyncMemcpyBatch( + const uint64_t* src_ptrs, + const uint64_t* dst_ptrs, + const uint64_t* sizes, + int64_t n) { + if (n == 0) return; + + // Scan the first non-zero entry to determine copy direction. + // Also capture the device-side pointer so we can infer which XPU to use. + const void* device_probe = nullptr; + bool needs_staging = false; + bool dst_is_pageable = false; // D2H to pageable host -> sync copy + for (int64_t i = 0; i < n; i++) { + if (sizes[i] == 0) continue; + const void* first_src = reinterpret_cast(src_ptrs[i]); + const void* first_dst = reinterpret_cast(dst_ptrs[i]); + + // Use device 0's context as a probe: we only need pointer *type* here, + // and USM pointer types are consistent across all devices on the same + // platform (host/unknown are always host; device is always device on its + // own platform). The actual device index is resolved below via + // infer_xpu_device_from_ptr(). + auto probe_ctx = vllm::xpu::vllmGetQueue(0).get_context(); + auto src_type = sycl::get_pointer_type(first_src, probe_ctx); + auto dst_type = sycl::get_pointer_type(first_dst, probe_ctx); + bool src_is_host = + (src_type == sycl::usm::alloc::host || + src_type == sycl::usm::alloc::unknown); + bool dst_is_device = (dst_type == sycl::usm::alloc::device); + needs_staging = src_is_host && dst_is_device; + // D2H to pageable host requires synchronous copy to avoid corruption. + dst_is_pageable = !dst_is_device && (dst_type == sycl::usm::alloc::unknown); + // Device-side pointer: dst for H2D, src for D2H or D2D. + device_probe = needs_staging ? first_dst : first_src; + break; + } + + if (device_probe == nullptr) return; // all sizes are zero + + // Infer the target XPU device from the device pointer and set the guard so + // that vllmGetQueue() returns the correct in-order queue. + const at::DeviceIndex dev = infer_xpu_device_from_ptr(device_probe); + const at::DeviceGuard device_guard(at::Device(at::kXPU, dev)); + + auto& queue = vllm::xpu::vllmGetQueue(); + + // Compute total bytes needed for the H2D staging buffer. + uint64_t total_bytes = 0; + for (int64_t i = 0; i < n; i++) { + total_bytes += sizes[i]; + } + + if (needs_staging) { + // H2D: allocate one contiguous pinned staging buffer, snapshot all source + // blocks, then submit all async DMAs. This avoids N separate allocator + // round-trips and protects against caller mutation after return. + auto staging = at::getHostAllocator(at::kXPU)->allocate( + static_cast(total_bytes)); + char* staging_ptr = static_cast(staging.get()); + TORCH_CHECK(staging_ptr, "Failed to allocate pinned staging buffer"); + + // Phase 1: snapshot all source blocks into staging (pure CPU work). + size_t staging_offset = 0; + for (int64_t i = 0; i < n; i++) { + size_t sz = static_cast(sizes[i]); + if (sz == 0) continue; + std::memcpy( + staging_ptr + staging_offset, + reinterpret_cast(src_ptrs[i]), + sz); + staging_offset += sz; + } + + // Phase 2: submit async DMA from staging to device in a tight loop, + // maximising PCIe/copy-engine throughput without interleaved CPU work. + staging_offset = 0; + for (int64_t i = 0; i < n; i++) { + size_t sz = static_cast(sizes[i]); + if (sz == 0) continue; + queue.memcpy( + reinterpret_cast(dst_ptrs[i]), + staging_ptr + staging_offset, + sz); + staging_offset += sz; + } + + // Keep the staging buffer alive until all submitted DMAs complete. + if (staging.get_context() != nullptr) { + at::getHostAllocator(at::kXPU)->record_event( + staging_ptr, + const_cast(staging.get_context()), + at::xpu::getCurrentXPUStream()); + } + } else { + // D2H or D2D: dst_is_pageable was probed once from the first non-zero + // entry (all entries share the same direction and memory class). + // Pageable D2H is unsafe with async DMA; fall back to sync copy. + for (int64_t i = 0; i < n; i++) { + size_t sz = static_cast(sizes[i]); + if (sz == 0) continue; + + const void* src = reinterpret_cast(src_ptrs[i]); + void* dst = reinterpret_cast(dst_ptrs[i]); + + if (dst_is_pageable) { + queue.memcpy(dst, src, sz).wait(); + } else { + queue.memcpy(dst, src, sz); + } + } + } +} + } // namespace xpu } // namespace vllm diff --git a/csrc/utils/mem_cpy.h b/csrc/utils/mem_cpy.h index ab274dc12..465ef4c1d 100644 --- a/csrc/utils/mem_cpy.h +++ b/csrc/utils/mem_cpy.h @@ -1,5 +1,6 @@ #pragma once #include +#include namespace vllm { namespace xpu { @@ -32,5 +33,27 @@ void xpuAsyncMemcpy( const void* hctx, bool is_pinned); +/** + * @brief Batch async memcpy: copies N independent (src, dst, size) triples + * in a single call, amortising per-copy overhead. + * + * The copy direction is auto-detected from the first non-zero entry's USM + * pointer types. All entries must share the same direction. + * + * For H2D: snapshots all source blocks through a single contiguous pinned + * staging buffer so the caller may safely mutate host memory immediately. + * For D2H / D2D: direct async DMA without staging. + * + * @param src_ptrs Array of N raw source addresses + * @param dst_ptrs Array of N raw destination addresses + * @param sizes Array of N byte counts + * @param n Number of entries + */ +void xpuAsyncMemcpyBatch( + const uint64_t* src_ptrs, + const uint64_t* dst_ptrs, + const uint64_t* sizes, + int64_t n); + } // namespace xpu } // namespace vllm diff --git a/csrc/xpu/attn/xe_2/chunk_prefill_configure.cmake b/csrc/xpu/attn/xe_2/chunk_prefill_configure.cmake index 99f4f8254..76c977c45 100644 --- a/csrc/xpu/attn/xe_2/chunk_prefill_configure.cmake +++ b/csrc/xpu/attn/xe_2/chunk_prefill_configure.cmake @@ -6,7 +6,7 @@ function(fmha_forward_configure FILENAME_SUFFIX) set(BOOL_FLAG_true "t") set(policy_list "chunk_policy_head64" "chunk_policy_head96" "chunk_policy_head128" - "chunk_policy_head192" "chunk_policy_head256") + "chunk_policy_head192" "chunk_policy_head256" "chunk_policy_head512") set(IMPL_KV_T "fp16") diff --git a/csrc/xpu/attn/xe_2/chunk_prefill_extern.hpp b/csrc/xpu/attn/xe_2/chunk_prefill_extern.hpp index 619be7dc7..7e1876aea 100644 --- a/csrc/xpu/attn/xe_2/chunk_prefill_extern.hpp +++ b/csrc/xpu/attn/xe_2/chunk_prefill_extern.hpp @@ -12,7 +12,8 @@ X(chunk_policy_head96) \ X(chunk_policy_head128) \ X(chunk_policy_head192) \ - X(chunk_policy_head256) + X(chunk_policy_head256) \ + X(chunk_policy_head512) // ============================================================================= // Automatic extern template declarations for all policy combinations diff --git a/csrc/xpu/attn/xe_2/fmha_utils.hpp b/csrc/xpu/attn/xe_2/fmha_utils.hpp index 67ce7b073..1533001c9 100644 --- a/csrc/xpu/attn/xe_2/fmha_utils.hpp +++ b/csrc/xpu/attn/xe_2/fmha_utils.hpp @@ -7,6 +7,7 @@ #define HEAD_SIZE_LIMIT_2 128 #define HEAD_SIZE_LIMIT_3 192 #define HEAD_SIZE_LIMIT_4 256 +#define HEAD_SIZE_LIMIT_5 512 enum class CutlassDType { half, bfloat16, float8_e4m3, float8_e5m2 }; @@ -83,6 +84,13 @@ struct chunk_policy_head256 { using SubgroupLayoutQK = Layout>; }; +struct chunk_policy_head512 { + using ShapeQK = Shape<_256, _32, _32>; + using ShapePV = Shape<_256, _32, _32>; + using ShapeOut = Shape<_256, _256>; + using SubgroupLayoutQK = Layout>; +}; + // define decode policy template struct decode_policy_qpacked_head { diff --git a/csrc/xpu/attn/xe_2/fmha_xe2.cpp b/csrc/xpu/attn/xe_2/fmha_xe2.cpp index 6d136b34d..e82aaa974 100644 --- a/csrc/xpu/attn/xe_2/fmha_xe2.cpp +++ b/csrc/xpu/attn/xe_2/fmha_xe2.cpp @@ -148,7 +148,7 @@ void cutlass_chunk_prefill_impl( CutlassQKType cuQKType = aten_to_Cutlass_qk_dtype(query, key_cache); - static constexpr int max_head_size = 256; + static constexpr int max_head_size = 512; TORCH_CHECK( head_size <= max_head_size, "FMHA forward only supports head dimension at most " + @@ -169,6 +169,9 @@ void cutlass_chunk_prefill_impl( } else if (args.head_size <= HEAD_SIZE_LIMIT_4) { policy_dispatch_func( queue, cuQKType, args, is_paged, is_causal, is_local, is_sink); + } else if (args.head_size <= HEAD_SIZE_LIMIT_5) { + policy_dispatch_func( + queue, cuQKType, args, is_paged, is_causal, is_local, is_sink); } else { TORCH_CHECK(false, "Unsupported head size for fmha"); } diff --git a/csrc/xpu/attn/xe_2/kernel/paged_decode_kernel.hpp b/csrc/xpu/attn/xe_2/kernel/paged_decode_kernel.hpp index cde1266e3..eeeb6dd75 100644 --- a/csrc/xpu/attn/xe_2/kernel/paged_decode_kernel.hpp +++ b/csrc/xpu/attn/xe_2/kernel/paged_decode_kernel.hpp @@ -323,7 +323,10 @@ class XeFMHAFwdSplitKVKernel { // Per-sequence split decision: short sequences are treated as // single-split even when num_kv_splits > 1, avoiding precision // loss from the split-reduce roundtrip. - constexpr int kMinBlocksForSplit = 128; + // Scale threshold with tile width: smaller tiles (p64) have fewer + // SGs per WG, so splitting becomes beneficial at fewer blocks. + constexpr int tile_n = get<1>(TileShapeQK{}); + constexpr int kMinBlocksForSplit = (tile_n <= 64) ? 32 : 128; bool is_single_split = (num_kv_splits > 1) && (windowed_k_blocks < kMinBlocksForSplit); @@ -620,6 +623,14 @@ class ReduceSplitK { int num_blocks_per_split = cute::ceil_div(windowed_k_blocks, num_kv_splits); + // Mirror the FMHA kernel's is_single_split decision so we only + // read split slots that were actually written. + constexpr int tile_n = get<1>(typename FMHAKernel_::TileShapeQK{}); + constexpr int kMinBlocksForSplit = (tile_n <= 64) ? 32 : 128; + bool is_single_split = + (num_kv_splits > 1) && (windowed_k_blocks < kMinBlocksForSplit); + int effective_splits = is_single_split ? 1 : num_kv_splits; + int offset_o = 0, offset_o_accum = 0; int offset_exp_sums = 0, offset_max_logits = 0; @@ -693,7 +704,7 @@ class ReduceSplitK { cutlass::platform::numeric_limits::lowest()}; ElementLSE global_exp_sums{0}; // only first subgroup participates - if (thr_id < num_kv_splits && + if (thr_id < effective_splits && thr_id * num_blocks_per_split < windowed_k_blocks) { ElementLSE cur_max_logit = max_logits(seq_idx, thr_id, head_q, l_coord); global_max_logits = sycl::max(global_max_logits, cur_max_logit); @@ -718,7 +729,7 @@ class ReduceSplitK { idx += SGPerWG::value * intel::sg_size) { ElementLSE acc = 0; global_exp_sums = 0; - for (int i = 0; i < num_kv_splits; ++i) { + for (int i = 0; i < effective_splits; ++i) { if (i * num_blocks_per_split >= windowed_k_blocks) { break; } diff --git a/csrc/xpu/attn/xe_2/paged_decode.hpp b/csrc/xpu/attn/xe_2/paged_decode.hpp index c694b687a..219290d17 100644 --- a/csrc/xpu/attn/xe_2/paged_decode.hpp +++ b/csrc/xpu/attn/xe_2/paged_decode.hpp @@ -29,22 +29,26 @@ using decode_policy_q8_h96_p64 = decode_policy_qpacked_head<_8, _96, _64>; using decode_policy_q8_h128_p64 = decode_policy_qpacked_head<_8, _128, _64>; using decode_policy_q8_h192_p64 = decode_policy_qpacked_head<_8, _192, _64>; using decode_policy_q8_h256_p64 = decode_policy_qpacked_head<_8, _256, _64>; +using decode_policy_q8_h512_p64 = decode_policy_qpacked_head<_8, _512, _64>; using decode_policy_q16_h64_p64 = decode_policy_qpacked_head<_16, _64, _64>; using decode_policy_q16_h96_p64 = decode_policy_qpacked_head<_16, _96, _64>; using decode_policy_q16_h128_p64 = decode_policy_qpacked_head<_16, _128, _64>; using decode_policy_q16_h192_p64 = decode_policy_qpacked_head<_16, _192, _64>; using decode_policy_q16_h256_p64 = decode_policy_qpacked_head<_16, _256, _64>; +using decode_policy_q16_h512_p64 = decode_policy_qpacked_head<_16, _512, _64>; using decode_policy_q8_h64_p128 = decode_policy_qpacked_head<_8, _64, _128>; using decode_policy_q8_h96_p128 = decode_policy_qpacked_head<_8, _96, _128>; using decode_policy_q8_h128_p128 = decode_policy_qpacked_head<_8, _128, _128>; using decode_policy_q8_h192_p128 = decode_policy_qpacked_head<_8, _192, _128>; using decode_policy_q8_h256_p128 = decode_policy_qpacked_head<_8, _256, _128>; +using decode_policy_q8_h512_p128 = decode_policy_qpacked_head<_8, _512, _128>; using decode_policy_q16_h64_p128 = decode_policy_qpacked_head<_16, _64, _128>; using decode_policy_q16_h96_p128 = decode_policy_qpacked_head<_16, _96, _128>; using decode_policy_q16_h128_p128 = decode_policy_qpacked_head<_16, _128, _128>; using decode_policy_q16_h192_p128 = decode_policy_qpacked_head<_16, _192, _128>; using decode_policy_q16_h256_p128 = decode_policy_qpacked_head<_16, _256, _128>; +using decode_policy_q16_h512_p128 = decode_policy_qpacked_head<_16, _512, _128>; struct paged_decode_args_t { void* query; diff --git a/csrc/xpu/attn/xe_2/paged_decode_configure.cmake b/csrc/xpu/attn/xe_2/paged_decode_configure.cmake index 2131afd19..eee965bff 100644 --- a/csrc/xpu/attn/xe_2/paged_decode_configure.cmake +++ b/csrc/xpu/attn/xe_2/paged_decode_configure.cmake @@ -37,12 +37,14 @@ function(paged_decode_configure FILENAME_SUFFIX) set(policy_8_128_64 "decode_policy_q8_h128_p64") set(policy_8_192_64 "decode_policy_q8_h192_p64") set(policy_8_256_64 "decode_policy_q8_h256_p64") + set(policy_8_512_64 "decode_policy_q8_h512_p64") set(policy_8_64_128 "decode_policy_q8_h64_p128") set(policy_8_96_128 "decode_policy_q8_h96_p128") set(policy_8_128_128 "decode_policy_q8_h128_p128") set(policy_8_192_128 "decode_policy_q8_h192_p128") set(policy_8_256_128 "decode_policy_q8_h256_p128") + set(policy_8_512_128 "decode_policy_q8_h512_p128") # Q-group size 16 policies set(policy_16_64_64 "decode_policy_q16_h64_p64") @@ -50,16 +52,18 @@ function(paged_decode_configure FILENAME_SUFFIX) set(policy_16_128_64 "decode_policy_q16_h128_p64") set(policy_16_192_64 "decode_policy_q16_h192_p64") set(policy_16_256_64 "decode_policy_q16_h256_p64") + set(policy_16_512_64 "decode_policy_q16_h512_p64") set(policy_16_64_128 "decode_policy_q16_h64_p128") set(policy_16_96_128 "decode_policy_q16_h96_p128") set(policy_16_128_128 "decode_policy_q16_h128_p128") set(policy_16_192_128 "decode_policy_q16_h192_p128") set(policy_16_256_128 "decode_policy_q16_h256_p128") + set(policy_16_512_128 "decode_policy_q16_h512_p128") # Configuration space dimensions set(qgroup_list "8" "16") - set(headsize_list "64" "96" "128" "192" "256") + set(headsize_list "64" "96" "128" "192" "256" "512") set(pagesize_list "64" "128") # ============================================================================= diff --git a/csrc/xpu/attn/xe_2/paged_decode_extern.hpp b/csrc/xpu/attn/xe_2/paged_decode_extern.hpp index 3a1a70e78..8a9e1b7d2 100644 --- a/csrc/xpu/attn/xe_2/paged_decode_extern.hpp +++ b/csrc/xpu/attn/xe_2/paged_decode_extern.hpp @@ -14,21 +14,25 @@ X(decode_policy_q8_h128_p64) \ X(decode_policy_q8_h192_p64) \ X(decode_policy_q8_h256_p64) \ + X(decode_policy_q8_h512_p64) \ X(decode_policy_q16_h64_p64) \ X(decode_policy_q16_h96_p64) \ X(decode_policy_q16_h128_p64) \ X(decode_policy_q16_h192_p64) \ X(decode_policy_q16_h256_p64) \ + X(decode_policy_q16_h512_p64) \ X(decode_policy_q8_h64_p128) \ X(decode_policy_q8_h96_p128) \ X(decode_policy_q8_h128_p128) \ X(decode_policy_q8_h192_p128) \ X(decode_policy_q8_h256_p128) \ + X(decode_policy_q8_h512_p128) \ X(decode_policy_q16_h64_p128) \ X(decode_policy_q16_h96_p128) \ X(decode_policy_q16_h128_p128) \ X(decode_policy_q16_h192_p128) \ - X(decode_policy_q16_h256_p128) + X(decode_policy_q16_h256_p128) \ + X(decode_policy_q16_h512_p128) // ============================================================================= // Automatic extern template declarations for all policy combinations diff --git a/csrc/xpu/attn/xe_2/paged_decode_utils.hpp b/csrc/xpu/attn/xe_2/paged_decode_utils.hpp index dcb6f9912..88cc8c930 100644 --- a/csrc/xpu/attn/xe_2/paged_decode_utils.hpp +++ b/csrc/xpu/attn/xe_2/paged_decode_utils.hpp @@ -59,6 +59,11 @@ inline void dispatch_by_head_size( decode_policy_qpacked_head>( queue, cuQKType, args, args.is_causal, args.is_local, args.is_sink); break; + case 5: + decode_policy_dispatch_func< + decode_policy_qpacked_head>( + queue, cuQKType, args, args.is_causal, args.is_local, args.is_sink); + break; default: TORCH_CHECK(false, "Unsupported head size for fmha"); } diff --git a/csrc/xpu/attn/xe_2/paged_decode_xe2.cpp b/csrc/xpu/attn/xe_2/paged_decode_xe2.cpp index 615372e0d..c3020e656 100644 --- a/csrc/xpu/attn/xe_2/paged_decode_xe2.cpp +++ b/csrc/xpu/attn/xe_2/paged_decode_xe2.cpp @@ -210,7 +210,7 @@ void cutlass_paged_decode_impl( CutlassQKType cuQKType = aten_to_Cutlass_qk_dtype(query, key_cache); - static constexpr int max_head_size = 256; + static constexpr int max_head_size = 512; TORCH_CHECK( head_size <= max_head_size, "FMHA forward only supports head dimension at most " + @@ -222,6 +222,7 @@ void cutlass_paged_decode_impl( if (head_size <= HEAD_SIZE_LIMIT_2) return 2; if (head_size <= HEAD_SIZE_LIMIT_3) return 3; if (head_size <= HEAD_SIZE_LIMIT_4) return 4; + if (head_size <= HEAD_SIZE_LIMIT_5) return 5; return -1; }; diff --git a/csrc/xpu/grouped_gemm/xe_2/gemm_xe2.hpp b/csrc/xpu/grouped_gemm/xe_2/gemm_xe2.hpp index 8a4d85ffa..0f78d838c 100644 --- a/csrc/xpu/grouped_gemm/xe_2/gemm_xe2.hpp +++ b/csrc/xpu/grouped_gemm/xe_2/gemm_xe2.hpp @@ -52,6 +52,33 @@ namespace MoE { using namespace cute; +template +CUTE_DEVICE TB apply_scale(TB& x, float& y) { + static_assert( + is_any_of_v, "Only BF16 & FP16 are supported"); + uint16_t z = sycl::bit_cast(x); +#if defined(__SYCL_DEVICE_ONLY__) && defined(SYCL_INTEL_TARGET) + if constexpr (is_same_v) { + asm("{\n" + ".decl Z_FP16 v_type=G type=HF num_elts=16 alias=<%0,0>\n" + ".decl Y_FP32 v_type=G type=F num_elts=16 alias=<%1,0>\n" + "mul (M1, 16) Z_FP16(0,0)<1> Z_FP16(0,0)<1;1,0> Y_FP32(0,0)<1;1,0>\n" + "}\n" + : "+rw"(z) + : "rw"(y)); + } else { + asm("{\n" + ".decl Z_BF16 v_type=G type=BF num_elts=16 alias=<%0,0>\n" + ".decl Y_FP32 v_type=G type=F num_elts=16 alias=<%1,0>\n" + "mul (M1, 16) Z_BF16(0,0)<1> Z_BF16(0,0)<1;1,0> Y_FP32(0,0)<1;1,0>\n" + "}\n" + : "+rw"(z) + : "rw"(y)); + } +#endif + return sycl::bit_cast(z); +} + template < class GmemTiledCopyA, class GmemTiledCopyB, @@ -290,7 +317,7 @@ CUTE_DEVICE void xe_gemm_4bits( auto pAgA = thr_prefetch_A.partition_S(gA); auto pBgB = thr_prefetch_B.partition_S(gB); - const int prefetch_dist = 3; + const int prefetch_dist = 6; constexpr int barrier_scope = 2; @@ -322,7 +349,8 @@ CUTE_DEVICE void xe_gemm_4bits( int group_num = get<1>(A.shape()) / group_size; int x_idx = sg_local_id / channel_num; - TA scales[thr_N * channel_num]; + using scaleStoreType = conditional_t, half_t, float>; + scaleStoreType scales[thr_N * channel_num]; clear(tCrC); @@ -335,6 +363,23 @@ CUTE_DEVICE void xe_gemm_4bits( for (; k_tile_prefetch < prefetch_dist; k_tile_prefetch++) { prefetch(prefetch_a, pAgA(_, _, _, k_tile_prefetch)); prefetch(prefetch_b, pBgB(_, _, _, k_tile_prefetch)); + + if (k_tile_prefetch * group_size < shape<1>(A)) { + auto next_scales_tensor = make_tensor( + make_gmem_ptr( + reinterpret_cast( + Scales + (n_tile_start + n_sg_start) * group_num + + k_tile_prefetch)), + make_layout( + make_shape(Int{}, Int<1>{}), + make_stride(group_num, Int<1>{}))); + auto prefetch_scales = make_block_2d_prefetch<1>( + make_shape(Int{}, Int<1>{}), next_scales_tensor); + auto thr_prefetch_scales = prefetch_scales.get_slice(sg_local_id); + auto pSgS = thr_prefetch_scales.partition_S( + make_identity_tensor(make_shape(Int{}, Int<1>{}))); + prefetch(prefetch_scales, pSgS(_, 0, 0)); + } } for (int k_tile = 0; k_tile < k_tile_count; k_tile++, k_tile_prefetch++) { @@ -352,7 +397,7 @@ CUTE_DEVICE void xe_gemm_4bits( for (int c = 0; c < channel_num; ++c) { int real_idx = x_idx + c * (sg_local_range / channel_num); int sg_local_n = n * sg_local_range + real_idx; - TA scale; + scaleStoreType scale; if constexpr (std::is_same_v) { scale = Scales [(n_tile_start + n_sg_start + sg_local_n) * group_num + @@ -363,12 +408,30 @@ CUTE_DEVICE void xe_gemm_4bits( [(n_tile_start + n_sg_start + sg_local_n) * group_num + group_idx] << 23; - scale = static_cast(reinterpret_cast(scale_u32)); + scale = static_cast( + reinterpret_cast(scale_u32)); } scales[n * channel_num + c] = scale; } } + + if ((group_idx + prefetch_dist) * group_size < shape<1>(A)) { + auto next_scales_tensor = make_tensor( + make_gmem_ptr( + reinterpret_cast( + Scales + (n_tile_start + n_sg_start) * group_num + + group_idx + prefetch_dist)), + make_layout( + make_shape(Int{}, Int<1>{}), + make_stride(group_num, Int<1>{}))); + auto prefetch_scales = make_block_2d_prefetch<1>( + make_shape(Int{}, Int<1>{}), next_scales_tensor); + auto thr_prefetch_scales = prefetch_scales.get_slice(sg_local_id); + auto pSgS = thr_prefetch_scales.partition_S( + make_identity_tensor(make_shape(Int{}, Int<1>{}))); + prefetch(prefetch_scales, pSgS(_, 0, 0)); + } } if (k_tile_prefetch < k_tile_count) { @@ -385,7 +448,12 @@ CUTE_DEVICE void xe_gemm_4bits( for (int c = 0; c < channel_num; ++c) { CUTLASS_PRAGMA_UNROLL for (int i = 0; i < tCrB.size() / thr_N / channel_num; ++i) { - tCrB(cute::tuple(c, _), n, _)[i] *= scales[n * channel_num + c]; + if constexpr (std::is_same_v) { + tCrB(cute::tuple(c, _), n, _)[i] *= scales[n * channel_num + c]; + } else { + tCrB(cute::tuple(c, _), n, _)[i] = apply_scale( + tCrB(cute::tuple(c, _), n, _)[i], scales[n * channel_num + c]); + } } } } diff --git a/csrc/xpu/grouped_gemm/xe_2/gemm_xe2_policy.hpp b/csrc/xpu/grouped_gemm/xe_2/gemm_xe2_policy.hpp index 57f320725..6b6966c3f 100644 --- a/csrc/xpu/grouped_gemm/xe_2/gemm_xe2_policy.hpp +++ b/csrc/xpu/grouped_gemm/xe_2/gemm_xe2_policy.hpp @@ -19,6 +19,18 @@ class xe_gemm_policy_base { class w16a16_policy : public xe_gemm_policy_base {}; +class w16a16_policy_n_128 : public xe_gemm_policy_base { + public: + using WGTile = Shape<_256, _128, _32>; + using SGLayout = Layout, Stride<_2, _1, _0>>; +}; + +class w16a16_policy_n_64 : public xe_gemm_policy_base { + public: + using WGTile = Shape<_256, _64, _32>; + using SGLayout = Layout, Stride<_1, _1, _0>>; +}; + class w16a16_policy_m_8 : public xe_gemm_policy_base { public: using WGTile = Shape<_8, _64, _32>; diff --git a/csrc/xpu/grouped_gemm/xe_2/grouped_gemm_xe2_interface.hpp b/csrc/xpu/grouped_gemm/xe_2/grouped_gemm_xe2_interface.hpp index 310bd8cea..c09252afc 100644 --- a/csrc/xpu/grouped_gemm/xe_2/grouped_gemm_xe2_interface.hpp +++ b/csrc/xpu/grouped_gemm/xe_2/grouped_gemm_xe2_interface.hpp @@ -286,7 +286,10 @@ at::Tensor cutlass_grouped_gemm_xe2_impl( } \ } - if (A_avg_M <= 32) { + if (A_avg_M <= 4) { + using policy = w4a16_policy_m_8; + W4A16LauncherCallER(policy); + } else if (A_avg_M <= 8) { using policy = w4a16_policy_m_16; W4A16LauncherCallER(policy); } else if (A_avg_M <= 128) { @@ -322,10 +325,10 @@ at::Tensor cutlass_grouped_gemm_xe2_impl( MoEGEMMLauncherCallER('R', 'R', policy, scalar_t, float_e5m2_t, float); \ } - if (A_avg_M <= 32) { + if (A_avg_M <= 8) { using policy = w8a16_policy_m_16; W8A16LauncherCallER(policy); - } else if (A_avg_M <= 128) { + } else if (A_avg_M <= 32) { using policy = w8a16_policy_m_32; W8A16LauncherCallER(policy); } else { @@ -346,12 +349,23 @@ at::Tensor cutlass_grouped_gemm_xe2_impl( MoEGEMMLauncherCallER('R', 'R', policy, scalar_t, scalar_t, scalar_t); \ } - if (A_avg_M <= 4) { + if (A_avg_M <= 8) { using policy = w16a16_policy_m_16; W16A16LauncherCallER(policy); - } else { - using policy = w16a16_policy; + } else if (A_avg_M <= 16) { + using policy = w16a16_policy_m_32; W16A16LauncherCallER(policy); + } else { + if (B_N <= 64) { + using policy = w16a16_policy_n_64; + W16A16LauncherCallER(policy); + } else if (B_N <= 512) { + using policy = w16a16_policy_n_128; + W16A16LauncherCallER(policy); + } else { + using policy = w16a16_policy; + W16A16LauncherCallER(policy); + } } #undef W16A16LauncherCallER } diff --git a/csrc/xpu/onednn/lru_cache.h b/csrc/xpu/onednn/lru_cache.h new file mode 100644 index 000000000..20054421a --- /dev/null +++ b/csrc/xpu/onednn/lru_cache.h @@ -0,0 +1,95 @@ +#pragma once + +#include +#include +#include +#include + +namespace oneDNN { + +template < + class key_t, + class value_t, + template class map_t = std::unordered_map> +class lru_cache { + public: + using value_type = std::pair; + using list_type = std::list; + using list_iter = typename list_type::iterator; + using map_type = map_t; + using const_list_iter = typename list_type::const_iterator; + using size_type = typename list_type::size_type; + + explicit lru_cache(size_type capacity) : capacity_(capacity) {} + lru_cache() : capacity_(0) {} + + [[nodiscard]] size_type size() const noexcept { return map_.size(); } + [[nodiscard]] size_type max_size() const noexcept { return capacity_; } + [[nodiscard]] bool empty() const noexcept { return vlist_.empty(); } + + void resize(size_type new_capacity) { + capacity_ = new_capacity; + trim(); + } + + list_iter begin() noexcept { return vlist_.begin(); } + const_list_iter begin() const noexcept { return vlist_.begin(); } + list_iter end() noexcept { return vlist_.end(); } + const_list_iter end() const noexcept { return vlist_.end(); } + + void clear() noexcept { + map_.clear(); + vlist_.clear(); + } + + void swap(lru_cache& other) noexcept { + using std::swap; + swap(vlist_, other.vlist_); + swap(map_, other.map_); + swap(capacity_, other.capacity_); + } + + list_iter find(const key_t& key) { + auto it = map_.find(key); + if (it == map_.end()) return end(); + vlist_.splice(vlist_.begin(), vlist_, it->second); + return it->second; + } + + std::pair insert(const value_type& value) { + auto it = map_.find(value.first); + if (it != map_.end()) { + // Move existing to front + vlist_.splice(vlist_.begin(), vlist_, it->second); + return {it->second, false}; + } + + // Insert new at front + vlist_.emplace_front(value); + map_[value.first] = vlist_.begin(); + + trim(); + + return {vlist_.begin(), true}; + } + + list_iter erase(list_iter pos) { + map_.erase(pos->first); + return vlist_.erase(pos); + } + + private: + void trim() { + while (map_.size() > capacity_) { + auto last = std::prev(vlist_.end()); + map_.erase(last->first); + vlist_.pop_back(); + } + } + + list_type vlist_; + map_type map_; + size_type capacity_; +}; + +} // namespace oneDNN diff --git a/csrc/xpu/onednn/onednn_ext.h b/csrc/xpu/onednn/onednn_ext.h index fe024069c..4135617fb 100644 --- a/csrc/xpu/onednn/onednn_ext.h +++ b/csrc/xpu/onednn/onednn_ext.h @@ -1,10 +1,9 @@ #pragma once -#include -#include #include #include +#include "lru_cache.h" #include "onednn_runtime.h" namespace std { @@ -737,8 +736,7 @@ T1 concat(const T1& t1, const T2& t2, const Ts&... ts) { return concat(concat(t1, t2), ts...); } -using primitive_cache = - at::native::onednn::lru_cache; +using primitive_cache = lru_cache; template struct matmul_primitive_cache_t { diff --git a/csrc/xpu/ops.h b/csrc/xpu/ops.h index 1f26b7554..f8932ddbf 100644 --- a/csrc/xpu/ops.h +++ b/csrc/xpu/ops.h @@ -52,6 +52,7 @@ torch::Tensor int4_gemm_w4a8( const std::optional& g_idx, const std::optional& bias); +#ifdef VLLM_MOE_ENABLED torch::Tensor cutlass_grouped_gemm_interface( torch::Tensor ptr_A, torch::Tensor ptr_B, @@ -64,6 +65,7 @@ torch::Tensor cutlass_grouped_gemm_interface( int64_t num_experts, bool is_B_int4, bool is_B_mxfp4); +#endif std::tuple deepseek_scaling_rope( const at::Tensor& positions, @@ -74,6 +76,7 @@ std::tuple deepseek_scaling_rope( int64_t rotary_dim, bool is_neox); +#ifdef VLLM_GDN_ENABLED void gdn_attention( torch::Tensor& core_attn_out, torch::Tensor& z, @@ -98,6 +101,7 @@ void gdn_attention( const int64_t num_actual_tokens, const int64_t tp_size, const bool reorder_input); +#endif bool is_bmg(int64_t device_index); diff --git a/csrc/xpu/torch_bindings.cpp b/csrc/xpu/torch_bindings.cpp index 05075919b..0e90c99fa 100644 --- a/csrc/xpu/torch_bindings.cpp +++ b/csrc/xpu/torch_bindings.cpp @@ -1,6 +1,8 @@ #include "core/registration.h" #include "xpu/ops.h" -#include "xpu/grouped_gemm/grouped_gemm_interface.h" +#ifdef VLLM_MOE_ENABLED + #include "xpu/grouped_gemm/grouped_gemm_interface.h" +#endif #include "xpu/lora/lora_ops.h" #include @@ -35,6 +37,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, xpu_ops) { "bias) -> Tensor"); xpu_ops.impl("int4_gemm_w4a8", torch::kXPU, &int4_gemm_w4a8); +#ifdef VLLM_MOE_ENABLED xpu_ops.def( "cutlass_grouped_gemm_interface(Tensor ptr_A, Tensor ptr_B, Tensor? " "ptr_scales, " @@ -48,6 +51,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, xpu_ops) { "cutlass_grouped_gemm_interface", torch::kXPU, &cutlass_grouped_gemm_interface); +#endif xpu_ops.def( "deepseek_scaling_rope(Tensor! positions, Tensor! query, Tensor! key, " @@ -72,6 +76,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, xpu_ops) { "-> ()"); xpu_ops.impl("bgmv_expand_slice", torch::kXPU, &bgmv_expand_slice); +#ifdef VLLM_GDN_ENABLED xpu_ops.def( "gdn_attention(Tensor! core_attn_out, Tensor! z, Tensor " "projected_states_qkvz, Tensor projected_states_ba," @@ -83,6 +88,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, xpu_ops) { "Tensor non_spec_state_indices_tensor, int num_actual_tokens, int " "tp_size, bool reorder_input) -> ()"); xpu_ops.impl("gdn_attention", torch::kXPU, &gdn_attention); +#endif // for empty tensor functions, we don't need dispatch key like torch::kXPU xpu_ops.def("is_bmg(int device_index) -> bool"); diff --git a/csrc/xpu_view.cpp b/csrc/xpu_view.cpp index 6a85edca1..96bc7bd4a 100644 --- a/csrc/xpu_view.cpp +++ b/csrc/xpu_view.cpp @@ -2,6 +2,7 @@ #include "ops.h" #include #include +#include namespace vllm::xpu { @@ -17,13 +18,17 @@ namespace vllm::xpu { class XPUHostViewAllocator : public c10::Allocator { public: + struct OwnerContext { + torch::Tensor owner; + }; + /** * @brief Constructor * @param host_ptr Pre-allocated host memory pointer * @param size Size of the host memory (in bytes) */ - XPUHostViewAllocator(void* host_ptr, size_t size) - : host_ptr_(host_ptr), size_(size) {} + XPUHostViewAllocator(void* host_ptr, size_t size, torch::Tensor owner) + : host_ptr_(host_ptr), size_(size), owner_(std::move(owner)) {} /** * @brief Allocate memory (actually just validates and wraps existing host @@ -36,15 +41,20 @@ class XPUHostViewAllocator : public c10::Allocator { // Verify requested memory size doesn't exceed pre-allocated memory size TORCH_CHECK( n <= size_, "Requested size exceeds allocated host pointer size"); - // Return wrapped data pointer with no-op deleter since memory is externally - // managed + // Use unique_ptr for RAII: if current_device() or DataPtr construction + // throws, the OwnerContext is automatically cleaned up instead of leaked. + auto ctx = std::make_unique(OwnerContext{owner_}); auto device_id = c10::xpu::current_device(); - return { - host_ptr_, // Actual data pointer - host_ptr_, // Context pointer (same as data pointer here) - [](void*) {}, // No-op deleter, doesn't actually free memory - c10::Device(c10::DeviceType::XPU, device_id) // Device type set to XPU - }; + + c10::DataPtr data_ptr{ + host_ptr_, + ctx.get(), + [](void* ptr) { delete static_cast(ptr); }, + c10::Device(c10::DeviceType::XPU, device_id)}; + + // DataPtr now owns the context via its deleter — release from unique_ptr. + ctx.release(); + return data_ptr; } /** @@ -71,6 +81,7 @@ class XPUHostViewAllocator : public c10::Allocator { private: void* const host_ptr_; // Pre-allocated host memory pointer const size_t size_; // Size of pre-allocated memory + torch::Tensor owner_; // Keeps pinned host storage alive }; } // namespace vllm::xpu @@ -92,7 +103,8 @@ torch::Tensor get_xpu_view_from_cpu_tensor(torch::Tensor& cpu_tensor) { auto scalar_type = cpu_tensor.scalar_type(); size_t byte_size = cpu_tensor.numel() * cpu_tensor.element_size(); - vllm::xpu::XPUHostViewAllocator allocator(host_ptr, byte_size); + // Keep `cpu_tensor` storage alive through the view tensor's lifetime. + vllm::xpu::XPUHostViewAllocator allocator(host_ptr, byte_size, cpu_tensor); c10::DataPtr data_ptr = allocator.allocate(byte_size); c10::Storage storage( c10::Storage::use_byte_size_t(), byte_size, std::move(data_ptr)); diff --git a/docs/test_scope_design.md b/docs/test_scope_design.md new file mode 100644 index 000000000..c216ab650 --- /dev/null +++ b/docs/test_scope_design.md @@ -0,0 +1,134 @@ +# Test Scope Design Document + +## Overview + +This document describes the multi-scope test system for `vllm-xpu-kernels`. The system supports four test scopes, controlled by the environment variable `XPU_KERNEL_TEST_SCOPE`, to balance CI speed vs. coverage across different scenarios. + +## Test Scopes + +| Scope | Env Value | Description | Use Case | +|-------|-----------|-------------|----------| +| **full** | `full` (or unset) | All tests × all parameters | Nightly CI, pre-release validation | +| **ci** | `ci` | All tests × reduced parameters | PR CI, push-to-main CI | +| **mini** | `mini` | Subset of tests × minimal parameters | Simulator, quick smoke test | +| **on-demand** | `ondemand:` | Profile-specific tests × model-specific shapes | Model-targeted validation (e.g. `ondemand:llama`) | + +## Environment Variable + +```bash +# Full scope (default) — nightly +export XPU_KERNEL_TEST_SCOPE=full +# or simply unset it + +# CI scope — PR CI +export XPU_KERNEL_TEST_SCOPE=ci + +# Mini scope — simulator / quick validation +export XPU_KERNEL_TEST_SCOPE=mini + +# On-demand scope — model-specific +export XPU_KERNEL_TEST_SCOPE=ondemand:llama +export XPU_KERNEL_TEST_SCOPE=ondemand:deepseek +``` + +## Architecture + +### Per-Test-Module Configuration + +Each test module defines a `TEST_SCOPE_PARAMS` dictionary with scope-specific parameter overrides: + +```python +TEST_SCOPE_PARAMS = { + "ci": { # CI scope: reduced shapes + "default": { # applies to all functions unless overridden + "num_tokens": [1, 128], + "hidden_size": [64], + }, + "test_specific_fn": { # per-function override + "num_tokens": [1], + }, + }, + "mini": { # Mini scope: minimal shapes + "default": { + "num_tokens": [1], + "hidden_size": [32], + }, + }, +} +``` + +### On-Demand Profiles + +On-demand profiles are defined centrally in `tests/test_scope_profiles.py`. Each profile maps test functions to model-relevant parameter sets: + +```python +ONDEMAND_PROFILES = { + "llama3": { + "tests/test_activation.py": { + "test_act_and_mul": { + "num_tokens": [1, 128], + "d": [11008], + "activation": ["silu_and_mul"], + }, + }, + ... + }, + "deepseek": { ... }, +} +``` + +### conftest.py Hook + +The `pytest_generate_tests` hook in `conftest.py`: + +1. Reads `XPU_KERNEL_TEST_SCOPE` env var +2. For `full` scope: no modification (run everything) +3. For `ci` / `mini` scope: looks up `TEST_SCOPE_PARAMS[scope]` in the test module +4. For `ondemand:` scope: looks up the profile in `test_scope_profiles.py` +5. Replaces `@pytest.mark.parametrize` values with the scoped subset +6. For `mini` scope: also skips entire tests if marked with `skip_for_mini=True` + +### Test Skipping for Mini Scope + +Tests can be skipped entirely in mini scope using: + +```python +SKIP_IN_MINI_SCOPE = True # Module-level flag to skip all tests in this module +``` + +Or per-function via `TEST_SCOPE_PARAMS`: + +```python +TEST_SCOPE_PARAMS = { + "mini": { + "test_expensive_fn": None, # None means skip this test in mini scope + }, +} +``` + +## Backward Compatibility + +- The old `XPU_KERNEL_PYTEST_PROFILER=MINI` environment variable is still supported and maps to `XPU_KERNEL_TEST_SCOPE=mini`. +- The old `MINI_PYTEST_PARAMS` dictionary is still consumed as a fallback for `ci` and `mini` scopes when `TEST_SCOPE_PARAMS` is not defined. + +## CI Workflow Integration + +```yaml +# PR CI (ci scope) +- name: test + run: | + XPU_KERNEL_TEST_SCOPE=ci pytest -v -s tests/ + +# Nightly CI (full scope) +- name: test + run: | + pytest -v -s tests/ +``` + +## Migration Guide + +For existing test files, the `MINI_PYTEST_PARAMS` still works. To adopt the new system: + +1. Rename `MINI_PYTEST_PARAMS` → embedded into `TEST_SCOPE_PARAMS["ci"]` (and optionally `["mini"]`) +2. Add more granular scope definitions as needed +3. The old `SKIP_TEST_FOR_MINI_SCOPE` pattern is replaced by `SKIP_IN_MINI_SCOPE = True` diff --git a/pyproject.toml b/pyproject.toml index dfde296b0..560a62241 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -27,13 +27,14 @@ classifiers = [ "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.12", + "Programming Language :: Python :: 3.13", "Intended Audience :: Developers", "Intended Audience :: Information Technology", "Intended Audience :: Science/Research", "Topic :: Scientific/Engineering :: Artificial Intelligence", "Topic :: Scientific/Engineering :: Information Analysis", ] -requires-python = ">=3.9,<3.13" +requires-python = ">=3.9,<3.14" dynamic = [ "version", "dependencies", "optional-dependencies"] [project.urls] @@ -42,21 +43,14 @@ Documentation="https://docs.vllm.ai/en/latest/" [tool.setuptools_scm] -# no extra settings needed, presence enables setuptools-scm +tag_regex = '^(?:[\w-]+-)?(?P[vV]?\d+(?:\.\d+)*(?:[._-]?(?:dev|a|b|rc|alpha|beta)\d*)?)$' +git_describe_command = "git describe --dirty --tags --long --match 'v*'" [tool.setuptools] [tool.setuptools.packages.find] where = ["."] include = ["vllm_xpu_kernels*"] -[tool.yapfignore] -ignore_patterns = [ - ".buildkite/**", - "benchmarks/**", - "build/**", - "examples/**", -] - [tool.ruff] # Allow lines to be as long as 80. line-length = 80 @@ -115,12 +109,6 @@ files = [ "vllm/triton_utils", "vllm/usage", ] -# TODO(woosuk): Include the code from Megatron and HuggingFace. -exclude = [ - "vllm/model_executor/parallel_utils/|vllm/model_executor/models/", - # Ignore triton kernels in ops. - 'vllm/attention/ops/.*\.py$' -] [tool.codespell] ignore-words-list = "dout, te, indicies, subtile, ElementE" diff --git a/setup.py b/setup.py index 61206f217..a62858ba1 100755 --- a/setup.py +++ b/setup.py @@ -75,6 +75,12 @@ def _build_custom_ops() -> bool: return True +def _is_enabled(env_name: str) -> bool: + """Check if a build option env var is enabled (default ON).""" + val = os.environ.get(env_name, "ON").strip().upper() + return val not in ("0", "OFF", "FALSE", "NO") + + class CMakeExtension(Extension): def __init__(self, name: str, cmake_lists_dir: str = '.', **kwa) -> None: @@ -176,6 +182,23 @@ def configure(self, ext: CMakeExtension) -> None: # on subsequent calls to python. cmake_args += ['-DVLLM_PYTHON_PATH={}'.format(":".join(sys.path))] + # Forward kernel build options to cmake so option() defaults are + # overridden when the user sets environment variables. + _kernel_options = [ + "BUILD_SYCL_TLA_KERNELS", + "VLLM_XPU_ENABLE_XE2", + "VLLM_XPU_ENABLE_XE_DEFAULT", + "BASIC_KERNELS_ENABLED", + "FA2_KERNELS_ENABLED", + "MOE_KERNELS_ENABLED", + "GDN_KERNELS_ENABLED", + "XPU_SPECIFIC_KERNELS_ENABLED", + "XPUMEM_ALLOCATOR_ENABLED", + ] + for opt in _kernel_options: + cmake_args.append('-D{}={}'.format( + opt, "ON" if _is_enabled(opt) else "OFF")) + # Override the base directory for FetchContent downloads to $ROOT/.deps # This allows sharing dependencies between profiles, # and plays more nicely with sccache. @@ -431,6 +454,10 @@ def get_base_commit_in_main_branch() -> str: ] } +_SCM_TAG_REGEX = (r'^(?:[\w-]+-)?' + r'(?P[vV]?\d+(?:\.\d+)*' + r'(?:[._-]?(?:dev|a|b|rc|alpha|beta)\d*)?)$') +_SCM_DESCRIBE_CMD = "git describe --dirty --tags --long --match 'v*'" def get_vllm_version() -> str: # Allow overriding the version. @@ -439,7 +466,11 @@ def get_vllm_version() -> str: os.environ["SETUPTOOLS_SCM_PRETEND_VERSION"] = env_version return get_version(write_to="_version.py") - version = get_version(write_to="_version.py") + version = get_version( + write_to="_version.py", + git_describe_command=_SCM_DESCRIBE_CMD, + tag_regex=_SCM_TAG_REGEX, + ) return version @@ -498,21 +529,37 @@ def build_extensions(self) -> None: ext_modules = [] -# List of additional shared libraries to install (intermediate build artifacts) -additional_libraries = { - "attn_kernels_xe_2": "/csrc/xpu/attn/xe_2", - "gdn_attn_kernels_xe_2": "/csrc/xpu/gdn_attn/xe_2", - "grouped_gemm_xe_default": "/csrc/xpu/grouped_gemm/xe_default", - "grouped_gemm_xe_2": "/csrc/xpu/grouped_gemm/xe_2", -} +# List of additional shared libraries to install (intermediate build artifacts). +# Only include libraries whose cmake targets will actually be built, based on +# the architecture and TLA kernel options. +additional_libraries = {} +if _is_enabled("BUILD_SYCL_TLA_KERNELS"): + if _is_enabled("VLLM_XPU_ENABLE_XE2"): + if _is_enabled("FA2_KERNELS_ENABLED"): + additional_libraries["attn_kernels_xe_2"] = "/csrc/xpu/attn/xe_2" + if _is_enabled("GDN_KERNELS_ENABLED"): + additional_libraries["gdn_attn_kernels_xe_2"] = ( + "/csrc/xpu/gdn_attn/xe_2") + if _is_enabled("MOE_KERNELS_ENABLED"): + additional_libraries["grouped_gemm_xe_2"] = ( + "/csrc/xpu/grouped_gemm/xe_2") + if _is_enabled("VLLM_XPU_ENABLE_XE_DEFAULT") and _is_enabled( + "MOE_KERNELS_ENABLED"): + additional_libraries["grouped_gemm_xe_default"] = ( + "/csrc/xpu/grouped_gemm/xe_default") if _build_custom_ops(): - ext_modules.append(CMakeExtension(name="vllm_xpu_kernels._C")) - ext_modules.append(CMakeExtension(name="vllm_xpu_kernels._vllm_fa2_C")) - ext_modules.append(CMakeExtension(name="vllm_xpu_kernels._moe_C")) - ext_modules.append(CMakeExtension(name="vllm_xpu_kernels._xpu_C")) - ext_modules.append( - CMakeExtension(name="vllm_xpu_kernels.xpumem_allocator")) + if _is_enabled("BASIC_KERNELS_ENABLED"): + ext_modules.append(CMakeExtension(name="vllm_xpu_kernels._C")) + if _is_enabled("FA2_KERNELS_ENABLED"): + ext_modules.append(CMakeExtension(name="vllm_xpu_kernels._vllm_fa2_C")) + if _is_enabled("MOE_KERNELS_ENABLED"): + ext_modules.append(CMakeExtension(name="vllm_xpu_kernels._moe_C")) + if _is_enabled("XPU_SPECIFIC_KERNELS_ENABLED"): + ext_modules.append(CMakeExtension(name="vllm_xpu_kernels._xpu_C")) + if _is_enabled("XPUMEM_ALLOCATOR_ENABLED"): + ext_modules.append( + CMakeExtension(name="vllm_xpu_kernels.xpumem_allocator")) if ext_modules: cmdclass = { diff --git a/tests/conftest.py b/tests/conftest.py index bdeb376c4..c69d2ddbb 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -6,32 +6,178 @@ from tests.utils import create_kv_caches_with_random +def _get_test_scope(): + """Determine the active test scope from environment variables. + + Priority: + 1. XPU_KERNEL_TEST_SCOPE (new unified env var) + 2. XPU_KERNEL_PYTEST_PROFILER=MINI (legacy, maps to "mini") + 3. Default: "full" + + Returns one of: "full", "ci", "mini", or "ondemand:" + """ + scope = os.getenv("XPU_KERNEL_TEST_SCOPE", "").strip().lower() + if scope: + return scope + + # Legacy compatibility + if os.getenv("XPU_KERNEL_PYTEST_PROFILER", "").strip().upper() == "MINI": + return "mini" + + return "full" + + +def _resolve_scope_params(module, func_name, scope): + """Resolve parameter overrides for the given scope. + + Lookup order: + 1. ondemand: profile-defined function entry or default entry + 2. mini: MINI_PYTEST_PARAMS[func_name] + 3. mini: MINI_PYTEST_PARAMS["default"] + 4. ci/full: use original test parameters + + Returns: + - dict of {param_name: values} to override, or + - None to skip this test entirely (when the entry is explicitly None) + - empty dict {} to run with original params (no override) + """ + if scope.startswith("ondemand:"): + # On-demand profiles are loaded from tests/test_scope_profiles.py + # Support both "ondemand:llama3" and "ondemand::llama3" + profile_name = scope.split(":", 1)[1].lstrip(":") + try: + from tests.test_scope_profiles import ONDEMAND_PROFILES + except ImportError: + return {} + profile = ONDEMAND_PROFILES.get(profile_name, {}) + # Match by module file path (relative) + module_file = getattr(module, "__file__", "") + for path_key, funcs in profile.items(): + if module_file.endswith(path_key): + entry = funcs.get(func_name, funcs.get("default", None)) + return entry # may be None (skip) or dict + return None # module not in profile → skip + + scope_key = scope # "ci" or "mini" + + # ci scope always uses the original/default parametrization. + if scope_key == "ci": + return {} + + # mini scope always uses the current file's MINI_PYTEST_PARAMS. + if scope_key == "mini": + legacy = getattr(module, "MINI_PYTEST_PARAMS", {}) + entry = legacy.get(func_name) + if entry is not None: + return entry + entry = legacy.get("default") + if entry is not None: + return entry + + return {} + + +def _apply_param_overrides(metafunc, profile): + """Replace @pytest.mark.parametrize values in markers (not via + metafunc.parametrize). + + Instead of calling metafunc.parametrize() directly (which conflicts with + pytest's built-in marker processing), we replace the marker objects with + new ones carrying the overridden values. The built-in pytest_generate_tests + hook then processes all markers uniformly. + """ + new_markers = [] + for mark in metafunc.definition.own_markers: + if (mark.name == "parametrize" and mark.args[0] in profile): + param_name = mark.args[0] + split_names = [n.strip() for n in param_name.split(",")] + if all(n in metafunc.fixturenames for n in split_names): + # Replace with a new marker carrying override values + new_mark = pytest.mark.parametrize(param_name, + profile[param_name]) + new_markers.append(new_mark.mark) + continue + new_markers.append(mark) + metafunc.definition.own_markers = new_markers + + +def _skip_test(metafunc, reason): + """Skip a test by collapsing parametrize markers to single values and adding + skip. + + We cannot use ``pytest.skip()`` inside ``pytest_generate_tests`` because the + resulting ``Skipped`` exception propagates up through ``_genfunctions`` and + aborts the entire **module** collection, discarding items already generated + for other test functions in the same file. + + Instead we: + 1. Replace each ``@pytest.mark.parametrize`` marker with a single-value + version so pytest still maps the parameter names to fixtures (avoiding + "fixture not found" errors). + 2. Add a ``@pytest.mark.skip`` marker to the **Python function object** + itself (``metafunc.function``). Markers on + ``metafunc.definition.own_markers`` do NOT propagate to the + ``Function`` items that ``_genfunctions`` creates; those items read + markers from the function object's ``pytestmark`` attribute instead. + """ + new_markers = [] + for m in metafunc.definition.own_markers: + if m.name == "parametrize" and m.args: + # Keep the marker but collapse to a single placeholder value + param_name = m.args[0] + original_values = m.args[1] + single = [original_values[0]] if original_values else [None] + new_markers.append( + pytest.mark.parametrize(param_name, single).mark) + else: + new_markers.append(m) + metafunc.definition.own_markers = new_markers + + # Add skip marker to the actual function object so the resulting + # Function items inherit it via pytestmark. + skip_mark = pytest.mark.skip(reason=reason).mark + func = metafunc.function + existing = list(getattr(func, "pytestmark", [])) + existing.append(skip_mark) + func.pytestmark = existing + + def pytest_generate_tests(metafunc): - use_mini_pytest_profiler = os.getenv("XPU_KERNEL_PYTEST_PROFILER", - "") == "MINI" - if not use_mini_pytest_profiler: + """Hook to apply test scope parameter overrides. + + Controlled by XPU_KERNEL_TEST_SCOPE env var: + - "full" (or unset): run all tests with all parameters (no override) + - "ci": run all tests with reduced parameter sets + - "mini": run subset of tests with minimal parameters + - "ondemand:": run model-specific tests and shapes + + See docs/test_scope_design.md for details. + """ + scope = _get_test_scope() + + if scope == "full": return module = metafunc.module + func_name = metafunc.function.__name__ - func_pytest_params = getattr(module, "MINI_PYTEST_PARAMS", {}) - profile = func_pytest_params.get(metafunc.function.__name__, None) + # --- Module-level skip for mini scope --- + if scope == "mini" and getattr(module, "SKIP_IN_MINI_SCOPE", False): + _skip_test(metafunc, "Skipped in mini scope (SKIP_IN_MINI_SCOPE=True)") + return - if not profile: - profile = func_pytest_params.get('default', None) + profile = _resolve_scope_params(module, func_name, scope) + + # None means explicitly skip this test + if profile is None: + _skip_test(metafunc, f"Skipped in {scope} scope") + return + # Empty dict means no override (run with original params) if not profile: return - for param_name, values in profile.items(): - split_names = [name.strip() for name in param_name.split(",")] - if all(name in metafunc.fixturenames for name in split_names): - new_markers = [] - for mark in metafunc.definition.own_markers: - if mark.name == "parametrize" and mark.args[0] != param_name: - new_markers.append(mark) - metafunc.definition.own_markers = new_markers - metafunc.parametrize(param_name, values) + _apply_param_overrides(metafunc, profile) @pytest.fixture diff --git a/tests/flash_attn/test_flash_attn_varlen_func.py b/tests/flash_attn/test_flash_attn_varlen_func.py index 9264d8ccd..9ca1ab0be 100644 --- a/tests/flash_attn/test_flash_attn_varlen_func.py +++ b/tests/flash_attn/test_flash_attn_varlen_func.py @@ -11,7 +11,7 @@ from vllm_xpu_kernels.flash_attn_interface import flash_attn_varlen_func NUM_HEADS = [(8, 2)] -HEAD_SIZES = [64, 128, 192, 256] +HEAD_SIZES = [64, 128, 192, 256, 512] BLOCK_SIZES = [64, 128] DTYPES = [torch.bfloat16] QDTYPES = [None] @@ -365,6 +365,8 @@ def test_decode_with_paged_kv( # if q_dtype is not None and (dtype != torch.bfloat16 or fa_version == 2): # pytest.skip("Flash attention with quantized inputs is only " # "supported on version 3 with bfloat16 base type") + if head_size == 512 and block_size == 128: + pytest.skip("skip test cases that may run out of SLM.") if num_heads == (16, 1) and head_size == 256: pytest.skip("skip test cases that may run out of SLM.") if block_size == 128 and num_blocks == 32768 and head_size >= 192: diff --git a/tests/fused_moe/test_fused_moe.py b/tests/fused_moe/test_fused_moe.py index 64707fa5c..2682d85d4 100755 --- a/tests/fused_moe/test_fused_moe.py +++ b/tests/fused_moe/test_fused_moe.py @@ -22,14 +22,15 @@ EP_RANK = [0, 1, 2, 3] EP_SIZE = [4] +# override pytest parameters when enable mini pytest MINI_PYTEST_PARAMS = { "default": { "m,n,k": [(1, 256, 128)], "e": [2], "topk": [1], "dtype": [torch.bfloat16], - "has_bias": [True] - } + "has_bias": [True], + }, } diff --git a/tests/fused_moe/test_remap_hidden_states.py b/tests/fused_moe/test_remap_hidden_states.py index 5bde882f2..35a036aa4 100644 --- a/tests/fused_moe/test_remap_hidden_states.py +++ b/tests/fused_moe/test_remap_hidden_states.py @@ -226,6 +226,98 @@ def test_remap_hidden_states(num_rows, hidden_size, total_experts_num, topk, print("Mismatched ref:", ref_unpermuted_scales[mismatched_indices]) +@pytest.mark.parametrize("num_rows", [262144]) +@pytest.mark.parametrize("hidden_size", [2048]) +@pytest.mark.parametrize("total_experts_num", [128]) +@pytest.mark.parametrize("topk", [8]) +@pytest.mark.parametrize("has_expert_map", [False]) +@pytest.mark.parametrize("recipe", + ["bf16", "fp16", "mxfp8", "mxfp4", "fp8block"]) +def test_remap_hidden_states_overflow(num_rows, hidden_size, total_experts_num, + topk, has_expert_map, recipe): + seed_everything(7) + + data_dtype, scale_dtype = RECIPE_TO_DTYPE.get(recipe, (None, None)) + + if has_expert_map: + local_experts_num = total_experts_num // 2 + else: + local_experts_num = total_experts_num + + if data_dtype in [torch.bfloat16, torch.float16]: + hidden_states = torch.randn((num_rows, hidden_size), + dtype=data_dtype, + device=DEVICE) + scales = None + elif data_dtype is torch.float8_e4m3fn: + hidden_states_fp32 = torch.randn((num_rows, hidden_size), + dtype=torch.float32, + device=DEVICE) + hidden_states = hidden_states_fp32.to(torch.float8_e4m3fn) + if recipe == "fp8block": + block_k = 128 + scales = torch.randn((num_rows, hidden_size // block_k), + device=DEVICE, + dtype=torch.float32) + elif recipe == "mxfp8": + block_k = 16 + scales = torch.randint(1, + 256, (num_rows, hidden_size // block_k), + device=DEVICE, + dtype=torch.uint8).view( + torch.float8_e8m0fnu) + elif data_dtype is torch.float4_e2m1fn_x2: + block_k = 16 # two input elem in a 8bit + hidden_states_fp32 = torch.randn((num_rows, hidden_size // 2), + dtype=torch.float32, + device=DEVICE) + hidden_states = hidden_states_fp32.to(torch.uint8).view( + torch.float4_e2m1fn_x2) + scales = torch.randint(1, + 256, (num_rows, hidden_size // 2 // block_k), + device=DEVICE, + dtype=torch.uint8).view(torch.float8_e8m0fnu) + + remapped_hidden_states = torch.empty_like(hidden_states).repeat_interleave( + topk, dim=0) + remapped_scales = None + if scale_dtype is not None: + remapped_scales = torch.empty_like(scales).repeat_interleave(topk, + dim=0) + expert_first_token_offset = torch.zeros((local_experts_num + 1), + dtype=torch.int64, + device=DEVICE) + unpermuted_row_to_permuted_row = torch.empty((num_rows, topk), + dtype=torch.int32, + device=DEVICE) + + expert_map = None + if has_expert_map: + expert_map = torch.full((total_experts_num, ), + -1, + dtype=torch.int64, + device=DEVICE) + expert_map[torch.randperm( + total_experts_num, + device=DEVICE)[:local_experts_num]] = torch.randperm( + local_experts_num, device=DEVICE) + expert_map = expert_map.to(torch.int32) + + scores = torch.randn((num_rows, total_experts_num), + device=DEVICE, + dtype=torch.float32) + _, topk_ids = torch.topk(scores, k=topk, dim=-1, sorted=False) + topk_ids = topk_ids.to(torch.int64) + + torch.ops._moe_C.remap_hidden_states( + hidden_states, scales, remapped_hidden_states, remapped_scales, + expert_map, expert_first_token_offset, unpermuted_row_to_permuted_row, + topk_ids, total_experts_num, local_experts_num) + + print("remapped_hidden_states", remapped_hidden_states, flush=True) + print("remapped_scales", remapped_scales, flush=True) + + def ref_init_expert_map(expert_map, local_experts_num, ep_rank, ep_size): expert_map_tmp = torch.full((local_experts_num * ep_size, ), -1, diff --git a/tests/gdn_attn/test_gdn_attn.py b/tests/gdn_attn/test_gdn_attn.py index f965c0dcb..f2f1af026 100755 --- a/tests/gdn_attn/test_gdn_attn.py +++ b/tests/gdn_attn/test_gdn_attn.py @@ -404,6 +404,10 @@ def test_gdn_attention(num_actual_tokens, batch_size, num_k_heads, head_k_dim, rtol = 5e-2 torch.testing.assert_close(z, ref_z, atol=atol, rtol=rtol) + + if num_actual_tokens == 8192: + pytest.skip("FIXME, skip core_attn_out test because of random error") + torch.testing.assert_close(core_attn_out, ref_core_attn_out, atol=atol, diff --git a/tests/ops/activation_op.py b/tests/ops/activation_op.py index 02b8cf8fd..b973a5825 100644 --- a/tests/ops/activation_op.py +++ b/tests/ops/activation_op.py @@ -96,6 +96,31 @@ def extra_repr(self) -> str: return f'approximate={repr(self.approximate)}' +class FatreluAndMul(CustomOp): + + def __init__(self, threshold: float): + super().__init__() + self.threshold = threshold + self.op = ops.fatrelu_and_mul + + def forward_native(self, x: torch.Tensor) -> torch.Tensor: + d = x.shape[-1] // 2 + x1 = x[..., :d] + x2 = x[..., d:] + x1 = F.threshold(x1, self.threshold, 0.0) + return x1 * x2 + + def forward_xpu(self, x: torch.Tensor) -> torch.Tensor: + d = x.shape[-1] // 2 + output_shape = (x.shape[:-1] + (d, )) + out = torch.empty(output_shape, dtype=x.dtype, device=x.device) + self.op(out, x, self.threshold) + return out + + def extra_repr(self) -> str: + return f'threshold={self.threshold!r}' + + class FastGELU(CustomOp): def __init__(self): diff --git a/tests/register_ops.py b/tests/register_ops.py index c5ee99d21..06dc63f0b 100644 --- a/tests/register_ops.py +++ b/tests/register_ops.py @@ -26,6 +26,11 @@ def silu_and_mul(out: torch.Tensor, input: torch.Tensor) -> None: torch.ops._C.silu_and_mul(out, input) +def silu_and_mul_quant(out: torch.Tensor, input: torch.Tensor, + scale: torch.Tensor) -> None: + torch.ops._C.silu_and_mul_quant(out, input, scale) + + def gelu_fast(out: torch.Tensor, input: torch.Tensor) -> None: torch.ops._C.gelu_fast(out, input) @@ -50,6 +55,11 @@ def gelu_tanh_and_mul(out: torch.Tensor, input: torch.Tensor) -> None: torch.ops._C.gelu_tanh_and_mul(out, input) +def fatrelu_and_mul(out: torch.Tensor, input: torch.Tensor, + threshold: float) -> None: + torch.ops._C.fatrelu_and_mul(out, input, threshold) + + def rotary_embedding( positions: torch.Tensor, query: torch.Tensor, @@ -487,6 +497,17 @@ def swap_blocks( block_mapping) +def swap_blocks_batch( + src_ptrs: torch.Tensor, + dst_ptrs: torch.Tensor, + sizes: torch.Tensor, +) -> None: + """Batch version of swap_blocks: copies N independent (src, dst, size) + triples in a single call. The target XPU device is auto-inferred from the + device-side pointers in src_ptrs/dst_ptrs.""" + torch.ops._C_cache_ops.swap_blocks_batch(src_ptrs, dst_ptrs, sizes) + + def topk_sigmoid(topk_weights: torch.Tensor, topk_ids: torch.Tensor, token_expert_indices: torch.Tensor, gating_output: torch.Tensor, renormalize: bool, diff --git a/tests/test_activation.py b/tests/test_activation.py index a2fffc397..eb3da18d5 100644 --- a/tests/test_activation.py +++ b/tests/test_activation.py @@ -1,10 +1,13 @@ # SPDX-License-Identifier: Apache-2.0 +import random + import pytest import torch from tests.allclose_default import get_default_atol, get_default_rtol -from tests.ops.activation_op import (FastGELU, GeluAndMul, MulAndSilu, NewGELU, - QuickGELU, Relu2NoMul, SiluAndMul) +from tests.ops.activation_op import (FastGELU, FatreluAndMul, GeluAndMul, + MulAndSilu, NewGELU, QuickGELU, + Relu2NoMul, SiluAndMul) from tests.utils import opcheck, seed_everything DTYPES = [torch.half, torch.bfloat16, torch.float] @@ -24,8 +27,9 @@ } -@pytest.mark.parametrize("activation", - ["silu_and_mul", "mul_and_silu", "gelu", "gelu_tanh"]) +@pytest.mark.parametrize( + "activation", + ["silu_and_mul", "mul_and_silu", "gelu", "gelu_tanh", "fatrelu"]) @pytest.mark.parametrize("num_tokens", NUM_TOKENS) @pytest.mark.parametrize("d", D) @pytest.mark.parametrize("dtype", DTYPES) @@ -55,15 +59,25 @@ def test_act_and_mul( elif activation == "gelu_tanh": layer = GeluAndMul(approximate="tanh") fn = torch.ops._C.gelu_tanh_and_mul + elif activation == "fatrelu": + threshold = random.uniform(0, 1) + layer = FatreluAndMul(threshold) + fn = torch.ops._C.fatrelu_and_mul out = layer(x) ref_out = layer.forward_native(x) - torch.testing.assert_close(out, ref_out, atol=1e-3, rtol=1e-3) + if activation == "fatrelu": + torch.testing.assert_close(out, ref_out, atol=0.0, rtol=0.0) + else: + torch.testing.assert_close(out, ref_out, atol=1e-3, rtol=1e-3) d = x.shape[-1] // 2 output_shape = (x.shape[:-1] + (d, )) out = torch.empty(output_shape, dtype=x.dtype, device=x.device) - opcheck(fn, (out, x)) + if activation == "fatrelu": + opcheck(fn, (out, x, threshold)) + else: + opcheck(fn, (out, x)) @pytest.mark.parametrize("activation", diff --git a/tests/test_cache.py b/tests/test_cache.py index dfcef8cc8..1f727c7d3 100644 --- a/tests/test_cache.py +++ b/tests/test_cache.py @@ -2,6 +2,7 @@ import random +import numpy as np import pytest import torch @@ -101,6 +102,16 @@ "device": ["xpu:0"], "kv_cache_dtype": KV_CACHE_DTYPE, }, + "test_swap_blocks_batch": { + "direction": [("cpu", "xpu")], + "device": ["xpu:0"], + }, + "test_swap_blocks_batch_empty": { + "device": ["xpu:0"], + }, + "test_swap_blocks_batch_h2d_mutation_race": { + "device": ["xpu:0"], + }, } @@ -948,3 +959,152 @@ def test_swap_blocks_mla( msg=f"Block {src} from src should have been swapped to block " f"{dst} in dst_cache.", ) + + +# --------------------------------------------------------------------------- +# swap_blocks_batch tests +# --------------------------------------------------------------------------- + + +def _build_batch_args( + src_cache: torch.Tensor, + dst_cache: torch.Tensor, + block_mapping: list[tuple[int, int]], + block_size_in_bytes: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Build (src_ptrs, dst_ptrs, sizes) tensors for swap_blocks_batch.""" + n = len(block_mapping) + src_arr = np.empty(n, dtype=np.uint64) + dst_arr = np.empty(n, dtype=np.uint64) + sz_arr = np.full(n, block_size_in_bytes, dtype=np.uint64) + + src_base = src_cache.data_ptr() + dst_base = dst_cache.data_ptr() + stride = src_cache.stride(0) * src_cache.element_size() + + for i, (sb, db) in enumerate(block_mapping): + src_arr[i] = src_base + sb * stride + dst_arr[i] = dst_base + db * stride + + return (torch.from_numpy(src_arr), torch.from_numpy(dst_arr), + torch.from_numpy(sz_arr)) + + +@pytest.mark.parametrize("direction", COPYING_DIRECTION) +@pytest.mark.parametrize("device", DEVICES) +@torch.inference_mode() +def test_swap_blocks_batch( + direction: tuple[str, str], + device: str, +) -> None: + """Test swap_blocks_batch for H2D, D2H and D2D directions.""" + num_mappings = 64 + num_heads = 8 + head_size = 64 + block_size = 8 + num_blocks = 256 + dtype = torch.bfloat16 + seed = 0 + + seed_everything(seed) + + src_device = device if direction[0] == "xpu" else "cpu" + dst_device = device if direction[1] == "xpu" else "cpu" + if "xpu" in direction: + torch.xpu.set_device(device) + + src_blocks = random.sample(range(num_blocks), num_mappings) + if src_device == dst_device: + remaining = list(set(range(num_blocks)) - set(src_blocks)) + dst_blocks = random.sample(remaining, num_mappings) + else: + dst_blocks = random.sample(range(num_blocks), num_mappings) + block_mapping = list(zip(src_blocks, dst_blocks)) + + src_key, src_val = create_kv_caches_with_random(num_blocks, block_size, 1, + num_heads, head_size, + "auto", dtype, seed, + src_device) + dst_key, dst_val = create_kv_caches_with_random(num_blocks, block_size, 1, + num_heads, head_size, + "auto", dtype, seed, + dst_device) + + src_key_clone = src_key[0].clone() + src_val_clone = src_val[0].clone() + + block_size_in_bytes = src_key[0].element_size() * src_key[0].stride(0) + + # Build batch args and call + for src_cache, dst_cache in [(src_key[0], dst_key[0]), + (src_val[0], dst_val[0])]: + sp, dp, sz = _build_batch_args(src_cache, dst_cache, block_mapping, + block_size_in_bytes) + ops.swap_blocks_batch(sp, dp, sz) + + torch.xpu.synchronize() + + for sb, db in block_mapping: + torch.testing.assert_close(src_key_clone[sb].cpu(), + dst_key[0][db].cpu()) + torch.testing.assert_close(src_val_clone[sb].cpu(), + dst_val[0][db].cpu()) + + +@pytest.mark.parametrize("device", DEVICES) +@torch.inference_mode() +def test_swap_blocks_batch_h2d_mutation_race(device: str) -> None: + """Verify staging buffer protects against caller mutation for H2D batch.""" + num_mappings = 256 + num_heads = 8 + head_size = 128 + block_size = 32 + num_blocks = 512 + dtype = torch.bfloat16 + seed = 0 + + seed_everything(seed) + + src_blocks = random.sample(range(num_blocks), num_mappings) + dst_blocks = random.sample(range(num_blocks), num_mappings) + block_mapping = list(zip(src_blocks, dst_blocks)) + + # Source: pinned CPU memory + src_key, src_val = create_kv_caches_with_pinned(num_blocks, block_size, 1, + num_heads, head_size, + "auto", dtype, seed, "cpu") + assert src_key[0].is_pinned() + + # Destination: XPU + dst_key, dst_val = create_kv_caches_with_random(num_blocks, block_size, 1, + num_heads, head_size, + "auto", dtype, seed) + + src_key_clone = src_key[0].clone() + src_val_clone = src_val[0].clone() + + block_size_in_bytes = src_key[0].element_size() * src_key[0].stride(0) + + for src_cache, dst_cache in [(src_key[0], dst_key[0]), + (src_val[0], dst_val[0])]: + sp, dp, sz = _build_batch_args(src_cache, dst_cache, block_mapping, + block_size_in_bytes) + ops.swap_blocks_batch(sp, dp, sz) + + # Immediately mutate source — should not affect destination. + src_key[0].fill_(0) + src_val[0].fill_(0) + + torch.xpu.synchronize() + + for sb, db in block_mapping: + torch.testing.assert_close( + src_key_clone[sb].cpu(), + dst_key[0][db].cpu(), + msg=f"Key block {sb}→{db} corrupted by post-call mutation", + ) + torch.testing.assert_close( + src_val_clone[sb].cpu(), + dst_val[0][db].cpu(), + msg=f"Value block {sb}→{db} corrupted by post-call mutation", + ) diff --git a/tests/test_exponential_2d.py b/tests/test_exponential_2d.py index f6bdca294..9d4d838bd 100755 --- a/tests/test_exponential_2d.py +++ b/tests/test_exponential_2d.py @@ -12,6 +12,17 @@ VOCAB_SIZE = [100, 1000, 10000] DEVICE = "xpu" +# CI/mini scope parameter overrides +MINI_PYTEST_PARAMS = { + "default": { + "batch_size": [1, 4], + "vocab_size": [1000], + "seed": [42], + "offset": [0], + "lambda_param": [1.0], + }, +} + class ExponentialDistributionTester: diff --git a/tests/test_fp8_gemm_onednn.py b/tests/test_fp8_gemm_onednn.py index eaef35520..298ca1236 100644 --- a/tests/test_fp8_gemm_onednn.py +++ b/tests/test_fp8_gemm_onednn.py @@ -33,7 +33,7 @@ # override pytest parameters when enable mini pytest MINI_PYTEST_PARAMS = { "test_fp8_gemm_w8a16": { - "batch": 1, + "batch": [1], "mnk_factors": MINI_MNK_FACTORS[:1], }, "test_fp8_gemm_per_tensor": { diff --git a/tests/test_fp8_quant.py b/tests/test_fp8_quant.py index b5ba512b5..0c07aa733 100644 --- a/tests/test_fp8_quant.py +++ b/tests/test_fp8_quant.py @@ -12,7 +12,12 @@ scaled_fp8_quant, scaled_quantize) from tests.ops.mx_utils import to_mxfp -SKIP_TEST_FOR_MINI_SCOPE = os.getenv("XPU_KERNEL_PYTEST_PROFILER") == "MINI" +# Legacy compatibility: used by skipif below. Now handled by conftest.py +# for the general case via SKIP_IN_MINI_SCOPE or TEST_SCOPE_PARAMS. +_test_scope = os.getenv("XPU_KERNEL_TEST_SCOPE", "").strip().lower() +_is_mini_scope = (_test_scope == "mini" or os.getenv( + "XPU_KERNEL_PYTEST_PROFILER", "").strip().upper() == "MINI") +SKIP_TEST_FOR_MINI_SCOPE = _is_mini_scope def as_float32_tensor(x: Union[float, torch.tensor]) -> torch.tensor: diff --git a/tests/test_fused_norm_quant.py b/tests/test_fused_norm_quant.py new file mode 100644 index 000000000..681884e9f --- /dev/null +++ b/tests/test_fused_norm_quant.py @@ -0,0 +1,447 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +# UT for fused RMSNorm + Quantization kernels. +# Adapted from tests/kernels/core/test_fused_quant_layernorm.py (CUDA). + +import pytest +import torch + +import vllm_xpu_kernels._C # noqa: F401 – registers torch.ops._C.* +from tests.ops.layernorm_op import RMSNorm + +DTYPES = [torch.half, torch.bfloat16] +# XPU FP8 default is e4m3fn; INT8 is always tested +QUANT_DTYPES = [torch.float8_e4m3fn, torch.int8] + +NUM_TOKENS = [1, 7, 83, 2048] +HIDDEN_SIZES = [64, 128, 1024, 5120] +ADD_RESIDUAL = [False, True] +GROUP_SIZES = [64, 128] +SEEDS = [0] +EPS = 1e-6 + +XPU_DEVICES = [ + f"xpu:{i}" for i in range(1 if torch.xpu.device_count() == 1 else 2) +] + +MINI_PYTEST_PARAMS = { + "default": { + "num_tokens": [4], + "HIDDEN_SIZES": [64], + "GROUP_SIZES": [64], + }, +} + + +def _ref_rms_norm( + layer: RMSNorm, + x: torch.Tensor, + residual: torch.Tensor | None, +) -> tuple[torch.Tensor, torch.Tensor | None]: + x_f32 = x.float() + if residual is not None: + residual = residual.clone() + # Kernel: z = input + residual, stored as orig_dtype, then read back + z_half = (x_f32 + residual.float()).to(x.dtype) + residual = z_half.clone() + x_f32 = z_half.float() + variance = x_f32.pow(2).mean(dim=-1, keepdim=True) + inv_rms = torch.rsqrt(variance + layer.variance_epsilon) + normed = x_f32 * inv_rms * layer.weight.float() + return normed, residual + + +def _ref_per_token_quant( + normed: torch.Tensor, + quant_dtype: torch.dtype, + scale_ub: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Reference dynamic per-token quantization.""" + num_tokens = normed.numel() // normed.shape[-1] + normed_2d = normed.view(num_tokens, -1).float() + + if quant_dtype == torch.int8: + absmax = normed_2d.abs().amax(dim=1, keepdim=True) # [T, 1] + scale = torch.where(absmax > 0, absmax / 127.0, + torch.ones_like(absmax)) + q = (normed_2d / scale).round().clamp(-128, 127).to(torch.int8) + return q, scale.squeeze(1) + else: + # FP8 e4m3fn + fp8_max = torch.finfo(quant_dtype).max + absmax = normed_2d.abs().amax(dim=1, keepdim=True) # [T, 1] + if scale_ub is not None: + absmax = torch.min(absmax, scale_ub.float()) + min_sf = 1.0 / (fp8_max * 512.0) + scale = torch.clamp(absmax / fp8_max, min=min_sf) # [T, 1] + inv_scale = 1.0 / scale + q = (normed_2d * inv_scale).clamp(-fp8_max, fp8_max).to(quant_dtype) + return q, scale.squeeze(1) + + +def _ref_per_group_quant( + normed: torch.Tensor, + group_size: int, + quant_dtype: torch.dtype, +) -> tuple[torch.Tensor, torch.Tensor]: + """Reference dynamic per-column-group quantization.""" + num_tokens = normed.numel() // normed.shape[-1] + hidden = normed.shape[-1] + num_groups = hidden // group_size + normed_2d = normed.view(num_tokens, hidden).float() + + q_out = torch.empty_like(normed_2d, dtype=quant_dtype) + scales = torch.empty(num_tokens, + num_groups, + dtype=torch.float32, + device=normed.device) + + fp8_max = torch.finfo( + quant_dtype).max if quant_dtype != torch.int8 else 127.0 + min_sf = 1.0 / (fp8_max * 512.0) if quant_dtype != torch.int8 else 0.0 + + for g in range(num_groups): + chunk = normed_2d[:, g * group_size:(g + 1) * group_size] # [T, G] + absmax = chunk.abs().amax(dim=1, keepdim=True) # [T, 1] + + if quant_dtype == torch.int8: + scale = torch.where(absmax > 0, absmax / 127.0, + torch.ones_like(absmax)) + q = (chunk / scale).round().clamp(-128, 127).to(torch.int8) + else: + scale = torch.clamp(absmax / fp8_max, min=min_sf) + inv_scale = 1.0 / scale + q = (chunk * inv_scale).clamp(-fp8_max, fp8_max).to(quant_dtype) + + q_out[:, g * group_size:(g + 1) * group_size] = q + scales[:, g] = scale.squeeze(1) + + return q_out, scales + + +def _ops_per_token_quant( + weight: torch.Tensor, + x: torch.Tensor, + quant_dtype: torch.dtype, + residual: torch.Tensor | None, + scale_ub: torch.Tensor | None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: + x = x.contiguous() + out = torch.empty_like(x, dtype=quant_dtype) + num_tokens = x.numel() // x.shape[-1] + scales = torch.empty(num_tokens, dtype=torch.float32, device=x.device) + if residual is not None: + residual = residual.clone().contiguous() + torch.ops._C.rms_norm_dynamic_per_token_quant(out, x, weight, scales, EPS, + scale_ub, residual) + return out, scales, residual + + +def _ops_per_group_quant( + weight: torch.Tensor, + x: torch.Tensor, + quant_dtype: torch.dtype, + group_size: int, + residual: torch.Tensor | None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: + x = x.contiguous() + out = torch.empty_like(x, dtype=quant_dtype) + num_tokens = x.numel() // x.shape[-1] + num_groups = x.shape[-1] // group_size + scales = torch.empty(num_tokens, + num_groups, + dtype=torch.float32, + device=x.device) + if residual is not None: + residual = residual.clone().contiguous() + torch.ops._C.rms_norm_per_block_quant( + out, + x, + weight, + scales, + EPS, + None, # scale_ub (not used for per-block) + residual, + group_size, + False, # is_scale_transposed + ) + return out, scales, residual + + +@pytest.mark.parametrize("num_tokens", NUM_TOKENS) +@pytest.mark.parametrize("hidden_size", HIDDEN_SIZES) +@pytest.mark.parametrize("add_residual", ADD_RESIDUAL) +@pytest.mark.parametrize("dtype", DTYPES) +@pytest.mark.parametrize("quant_dtype", QUANT_DTYPES) +@pytest.mark.parametrize("seed", SEEDS) +@pytest.mark.parametrize("device", XPU_DEVICES) +@torch.inference_mode() +def test_rms_norm_dynamic_per_token_quant( + num_tokens: int, + hidden_size: int, + add_residual: bool, + dtype: torch.dtype, + quant_dtype: torch.dtype, + seed: int, + device: str, +) -> None: + torch.manual_seed(seed) + torch.set_default_device("xpu") + torch.xpu.set_device(device) + + layer = RMSNorm(hidden_size, eps=EPS).to(dtype=dtype) + layer.weight.data.normal_(mean=1.0, std=0.1) + + scale = 1.0 / hidden_size + x = torch.randn(num_tokens, hidden_size, dtype=dtype) * scale + residual = torch.randn_like(x) * scale if add_residual else None + + # Reference: native RMSNorm + per-token quant + ref_normed, ref_residual = _ref_rms_norm(layer, x, residual) + ref_q, ref_scales = _ref_per_token_quant(ref_normed, quant_dtype) + + # Kernel + ops_q, ops_scales, ops_residual = _ops_per_token_quant( + layer.weight.data, x, quant_dtype, residual, None) + + assert ops_q.dtype == quant_dtype + assert ops_scales.dtype == torch.float32 + + if quant_dtype == torch.int8: + torch.testing.assert_close(ref_scales, + ops_scales, + atol=1e-6, + rtol=1e-6) + torch.testing.assert_close(ref_q, ops_q, atol=1, rtol=0) + else: + # FP8: scales should be close; compare quantized values or dequantized + torch.testing.assert_close(ref_scales, + ops_scales, + atol=1e-5, + rtol=1e-5) + ref_qf = ref_q.float() + ops_qf = ops_q.float() + if not torch.allclose(ref_qf, ops_qf, atol=1e-6): + # Fallback: dequantize both with their own scales and compare. + # NOTE: It is possible that some future test cases trigger this + # max diff due to precision issues. If such an error is + # encountered, it's recommended to inspect the differences between + # all corresponding elements from each tensor (e.g. by looping over + # them) and checking how many the max diff error shows up on (just + # a few bad elements should still be considered acceptable). + ref_deq = ref_qf * ref_scales.view(-1, 1) + ops_deq = ops_qf * ops_scales.view(-1, 1) + torch.testing.assert_close(ref_deq, ops_deq, atol=0.2, rtol=0.15) + + if add_residual: + torch.testing.assert_close(ref_residual, + ops_residual, + atol=1e-2, + rtol=1e-2) + + +@pytest.mark.parametrize("num_tokens", [1, 7, 83, 2048]) +@pytest.mark.parametrize("hidden_size", [128, 1024, 5120]) +@pytest.mark.parametrize("group_size", GROUP_SIZES) +@pytest.mark.parametrize("add_residual", ADD_RESIDUAL) +@pytest.mark.parametrize("dtype", DTYPES) +@pytest.mark.parametrize("quant_dtype", [torch.float8_e4m3fn, torch.int8]) +@pytest.mark.parametrize("seed", SEEDS) +@pytest.mark.parametrize("device", XPU_DEVICES) +@torch.inference_mode() +def test_rms_norm_per_block_quant( + num_tokens: int, + hidden_size: int, + group_size: int, + add_residual: bool, + dtype: torch.dtype, + quant_dtype: torch.dtype, + seed: int, + device: str, +) -> None: + if hidden_size % group_size != 0: + pytest.skip(f"hidden_size {hidden_size} not divisible by \ + group_size {group_size}") + + torch.manual_seed(seed) + torch.set_default_device("xpu") + torch.xpu.set_device(device) + + layer = RMSNorm(hidden_size, eps=EPS).to(dtype=dtype) + layer.weight.data.normal_(mean=1.0, std=0.1) + + scale = 1.0 / hidden_size + x = torch.randn(num_tokens, hidden_size, dtype=dtype) * scale + residual = torch.randn_like(x) * scale if add_residual else None + + # Reference: native RMSNorm + per-group quant + ref_normed, ref_residual = _ref_rms_norm(layer, x, residual) + ref_q, ref_scales = _ref_per_group_quant(ref_normed, group_size, + quant_dtype) + + # Kernel + ops_q, ops_scales, ops_residual = _ops_per_group_quant( + layer.weight.data, x, quant_dtype, group_size, residual) + + assert ops_q.dtype == quant_dtype + assert ops_scales.dtype == torch.float32 + assert ops_scales.shape == (num_tokens, hidden_size // group_size) + + torch.testing.assert_close(ref_scales, ops_scales, atol=1e-4, rtol=1e-4) + + if quant_dtype == torch.int8: + torch.testing.assert_close(ref_q, ops_q, atol=1, rtol=0) + else: + num_groups = hidden_size // group_size + ref_qf = ref_q.float().view(num_tokens, num_groups, group_size) + ops_qf = ops_q.float().view(num_tokens, num_groups, group_size) + if not torch.allclose(ref_qf, ops_qf, atol=1e-6): + ref_scales_e = ref_scales.unsqueeze(-1) + ops_scales_e = ops_scales.unsqueeze(-1) + ref_deq = ref_qf * ref_scales_e + ops_deq = ops_qf * ops_scales_e + torch.testing.assert_close(ref_deq, ops_deq, atol=0.2, rtol=0.15) + + if add_residual: + torch.testing.assert_close(ref_residual, + ops_residual, + atol=1e-2, + rtol=1e-2) + + +def _ops_rms_norm_static_fp8_quant( + weight: torch.Tensor, + x: torch.Tensor, + scale: torch.Tensor, +) -> torch.Tensor: + out = torch.empty_like(x, dtype=torch.float8_e4m3fn) + torch.ops._C.rms_norm_static_fp8_quant(out, x, weight, scale, EPS) + return out + + +def _ops_fused_add_rms_norm_static_fp8_quant( + weight: torch.Tensor, + x: torch.Tensor, + residual: torch.Tensor, + scale: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + out = torch.empty(x.shape[-1], dtype=torch.float8_e4m3fn, + device=x.device).expand_as(x).contiguous() + residual = residual.clone().contiguous() + torch.ops._C.fused_add_rms_norm_static_fp8_quant(out, x, residual, weight, + scale, EPS) + return out, residual + + +def _ref_static_fp8_quant( + normed: torch.Tensor, + scale: torch.Tensor, +) -> torch.Tensor: + """Reference static FP8 quantization with pre-computed scale.""" + fp8_max = torch.finfo(torch.float8_e4m3fn).max + inv_scale = 1.0 / scale.float() + q = (normed.float() * inv_scale).clamp(-fp8_max, fp8_max) + return q.to(torch.float8_e4m3fn) + + +@pytest.mark.parametrize("num_tokens", NUM_TOKENS) +@pytest.mark.parametrize("hidden_size", HIDDEN_SIZES) +@pytest.mark.parametrize("dtype", DTYPES) +@pytest.mark.parametrize("seed", SEEDS) +@pytest.mark.parametrize("device", XPU_DEVICES) +@pytest.mark.parametrize("strided_input", [False, True]) +@torch.inference_mode() +def test_rms_norm_static_fp8_quant( + num_tokens: int, + hidden_size: int, + dtype: torch.dtype, + seed: int, + device: str, + strided_input: bool, +) -> None: + torch.manual_seed(seed) + torch.set_default_device("xpu") + torch.xpu.set_device(device) + + layer = RMSNorm(hidden_size, eps=EPS).to(dtype=dtype) + layer.weight.data.normal_(mean=1.0, std=0.1) + + scale_val = 1.0 + quant_scale = torch.tensor(scale_val, dtype=torch.float32) + input_scale = 1.0 / hidden_size + last_dim = 2 * hidden_size if strided_input else hidden_size + x = torch.randn(num_tokens, last_dim, dtype=dtype) * input_scale + x = x[..., :hidden_size] + if num_tokens > 1: + assert x.is_contiguous() != strided_input + + # Reference: native RMSNorm then static FP8 quant + ref_normed, _ = _ref_rms_norm(layer, x, None) + ref_q = _ref_static_fp8_quant(ref_normed, quant_scale) + + # Kernel + ops_q = _ops_rms_norm_static_fp8_quant(layer.weight.data, x, quant_scale) + + assert ops_q.dtype == torch.float8_e4m3fn + # FP8 has coarse quantization; boundary elements may differ by up to 1 LSB. + ref_qf = ref_q.float() + ops_qf = ops_q.float() + if not torch.allclose(ref_qf, ops_qf, atol=1e-6): + ref_deq = ref_qf * quant_scale + ops_deq = ops_qf * quant_scale + torch.testing.assert_close(ref_deq, ops_deq, atol=0.2, rtol=0.15) + + +@pytest.mark.parametrize("num_tokens", NUM_TOKENS) +@pytest.mark.parametrize("hidden_size", HIDDEN_SIZES) +@pytest.mark.parametrize("dtype", DTYPES) +@pytest.mark.parametrize("seed", SEEDS) +@pytest.mark.parametrize("device", XPU_DEVICES) +@pytest.mark.parametrize("strided_input", [False, True]) +@torch.inference_mode() +def test_fused_add_rms_norm_static_fp8_quant( + num_tokens: int, + hidden_size: int, + dtype: torch.dtype, + seed: int, + device: str, + strided_input: bool, +) -> None: + torch.manual_seed(seed) + torch.set_default_device("xpu") + torch.xpu.set_device(device) + + layer = RMSNorm(hidden_size, eps=EPS).to(dtype=dtype) + layer.weight.data.normal_(mean=1.0, std=0.1) + + scale_val = 1.0 + quant_scale = torch.tensor(scale_val, dtype=torch.float32) + input_scale = 1.0 / hidden_size + last_dim = 2 * hidden_size if strided_input else hidden_size + x = torch.randn(num_tokens, last_dim, dtype=dtype) * input_scale + x = x[..., :hidden_size] + if num_tokens > 1: + assert x.is_contiguous() != strided_input + residual = torch.randn(num_tokens, hidden_size, dtype=dtype) * input_scale + + # Reference: native RMSNorm (with residual) then static FP8 quant + ref_normed, ref_residual = _ref_rms_norm(layer, x, residual) + ref_q = _ref_static_fp8_quant(ref_normed, quant_scale) + + # Kernel + ops_q, ops_residual = _ops_fused_add_rms_norm_static_fp8_quant( + layer.weight.data, x, residual, quant_scale) + + assert ops_q.dtype == torch.float8_e4m3fn + # Compare quantized values first, then fall back to dequantized comparison. + ref_qf = ref_q.float() + ops_qf = ops_q.float() + if not torch.allclose(ref_qf, ops_qf, atol=1e-6): + ref_deq = ref_qf * quant_scale + ops_deq = ops_qf * quant_scale + torch.testing.assert_close(ref_deq, ops_deq, atol=0.2, rtol=0.15) + torch.testing.assert_close(ops_residual, + ref_residual, + atol=1e-2, + rtol=1e-2) diff --git a/tests/test_fused_qk_norm_rope.py b/tests/test_fused_qk_norm_rope.py new file mode 100644 index 000000000..95471febf --- /dev/null +++ b/tests/test_fused_qk_norm_rope.py @@ -0,0 +1,258 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch + +import vllm_xpu_kernels._C # noqa: F401 - registers torch.ops._C ops +from tests.utils import opcheck + +DTYPES = [torch.half, torch.bfloat16] +IS_NEOX = [True, False] +EPS_VALUES = [1e-5, 1e-6] +HEAD_DIMS = [64, 128] +NUM_TOKENS = [1, 4, 32] +HEAD_CONFIGS = [(16, 4), (32, 8)] # (num_heads_q, num_heads_kv) +ROTARY_RATIOS = [1.0, 0.5] +SEEDS = [13] +XPU_DEVICES = [ + f"xpu:{i}" for i in range(1 if torch.xpu.device_count() == 1 else 2) +] + +MINI_PYTEST_PARAMS = { + "default": { + "num_tokens": [4], + "head_dim": [128], + "head_config": [(16, 4)], + "rotary_ratio": [1.0], + }, +} + + +def _rms_norm(x: torch.Tensor, weight: torch.Tensor, + eps: float) -> torch.Tensor: + """Reference RMSNorm: x is [..., head_dim], weight is [head_dim].""" + variance = x.to(torch.float32).pow(2).mean(dim=-1, keepdim=True) + x_normed = x * torch.rsqrt(variance + eps) + return (x_normed * weight).to(x.dtype) + + +def _apply_rotary_emb(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, + is_neox: bool) -> torch.Tensor: + """Apply RoPE to x. x shape: [num_tokens, num_heads, head_dim]. + cos/sin shape: [num_tokens, rotary_dim/2].""" + rotary_dim = cos.shape[-1] * 2 + x_rot = x[..., :rotary_dim].float() + x_pass = x[..., rotary_dim:] + + if is_neox: + # Neox style: first half and second half + half = rotary_dim // 2 + x1 = x_rot[..., :half] + x2 = x_rot[..., half:] + cos_val = cos.unsqueeze(1) # [tokens, 1, rotary_dim/2] + sin_val = sin.unsqueeze(1) + o1 = x1 * cos_val - x2 * sin_val + o2 = x2 * cos_val + x1 * sin_val + x_rot = torch.cat([o1, o2], dim=-1) + else: + # Interleaved style: pairs of (x0, x1) + x_rot = x_rot.view(*x_rot.shape[:-1], -1, 2) + x0 = x_rot[..., 0] + x1 = x_rot[..., 1] + cos_val = cos.unsqueeze(1) + sin_val = sin.unsqueeze(1) + o0 = x0 * cos_val - x1 * sin_val + o1 = x0 * sin_val + x1 * cos_val + x_rot = torch.stack([o0, o1], dim=-1).flatten(-2) + + return torch.cat([x_rot.to(x.dtype), x_pass], dim=-1) + + +def _reference_fused_qk_norm_rope( + qkv: torch.Tensor, + num_heads_q: int, + num_heads_kv: int, + head_dim: int, + eps: float, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + cos_sin_cache: torch.Tensor, + is_neox: bool, + positions: torch.Tensor, +) -> torch.Tensor: + """Pure PyTorch reference implementation.""" + q_size = num_heads_q * head_dim + kv_size = num_heads_kv * head_dim + + q, k, v = qkv.split([q_size, kv_size, kv_size], dim=-1) + + # RMSNorm on Q heads + q_by_head = q.view(-1, num_heads_q, head_dim) + q_by_head = _rms_norm(q_by_head, q_weight, eps) + + # RMSNorm on K heads + k_by_head = k.view(-1, num_heads_kv, head_dim) + k_by_head = _rms_norm(k_by_head, k_weight, eps) + + # Get cos/sin for the given positions + rotary_dim = cos_sin_cache.shape[1] + embed_dim = rotary_dim // 2 + cos_sin = cos_sin_cache[positions] # [num_tokens, rotary_dim] + cos = cos_sin[:, :embed_dim].float() + sin = cos_sin[:, embed_dim:].float() + + # Apply RoPE + q_by_head = _apply_rotary_emb(q_by_head, cos, sin, is_neox) + k_by_head = _apply_rotary_emb(k_by_head, cos, sin, is_neox) + + q = q_by_head.view(-1, q_size) + k = k_by_head.view(-1, kv_size) + + return torch.cat([q, k, v], dim=-1) + + +@pytest.mark.parametrize("device", XPU_DEVICES) +@pytest.mark.parametrize("dtype", DTYPES) +@pytest.mark.parametrize("is_neox", IS_NEOX) +@pytest.mark.parametrize("eps", EPS_VALUES) +@pytest.mark.parametrize("head_dim", HEAD_DIMS) +@pytest.mark.parametrize("num_tokens", NUM_TOKENS) +@pytest.mark.parametrize("head_config", HEAD_CONFIGS) +@pytest.mark.parametrize("rotary_ratio", ROTARY_RATIOS) +@pytest.mark.parametrize("seed", SEEDS) +@torch.inference_mode() +def test_fused_qk_norm_rope( + device: str, + dtype: torch.dtype, + is_neox: bool, + eps: float, + head_dim: int, + num_tokens: int, + head_config: tuple, + rotary_ratio: float, + seed: int, +) -> None: + torch.manual_seed(seed) + torch.set_default_device("xpu") + torch.xpu.set_device(device) + + num_heads_q, num_heads_kv = head_config + rotary_dim = int(head_dim * rotary_ratio) + max_position = 4096 + + total_dim = (num_heads_q + 2 * num_heads_kv) * head_dim + qkv_base = torch.randn(num_tokens, total_dim, dtype=dtype, device=device) + qkv_fused = qkv_base.clone() + positions = torch.randint(0, + max_position, (num_tokens, ), + dtype=torch.long, + device=device) + + q_weight = torch.empty(head_dim, dtype=dtype, device=device) + q_weight.normal_(mean=1.0, std=0.1) + k_weight = torch.empty(head_dim, dtype=dtype, device=device) + k_weight.normal_(mean=1.0, std=0.1) + + # Build cos_sin_cache: [max_position, rotary_dim] + # Layout: [cos_0..cos_{rotary_dim/2-1}, sin_0..sin_{rotary_dim/2-1}] + inv_freq = 1.0 / (10000.0**( + torch.arange(0, rotary_dim // 2, dtype=torch.float32, device=device) / + (rotary_dim // 2))) + t = torch.arange(max_position, dtype=torch.float32, device=device) + freqs = torch.outer(t, inv_freq) + cos_sin_cache = torch.cat([freqs.cos(), freqs.sin()], dim=-1) + + # Reference (before in-place op) + ref_result = _reference_fused_qk_norm_rope( + qkv_base, + num_heads_q, + num_heads_kv, + head_dim, + eps, + q_weight, + k_weight, + cos_sin_cache, + is_neox, + positions, + ) + + # Run fused kernel (in-place on qkv_fused) + torch.ops._C.fused_qk_norm_rope( + qkv_fused, + num_heads_q, + num_heads_kv, + num_heads_kv, + head_dim, + eps, + q_weight, + k_weight, + cos_sin_cache, + is_neox, + positions, + ) + + if dtype == torch.float16: + ATOL, RTOL = (2e-3, 2e-3) + else: + ATOL, RTOL = (1e-2, 1e-2) + + torch.testing.assert_close( + qkv_fused, + ref_result, + atol=ATOL, + rtol=RTOL, + ) + + +@pytest.mark.parametrize("device", XPU_DEVICES) +@pytest.mark.parametrize("dtype", DTYPES) +@pytest.mark.parametrize("is_neox", IS_NEOX) +@torch.inference_mode() +def test_fused_qk_norm_rope_opcheck( + device: str, + dtype: torch.dtype, + is_neox: bool, +) -> None: + """Validate the op schema and registration with opcheck.""" + torch.manual_seed(42) + torch.set_default_device("xpu") + torch.xpu.set_device(device) + + num_heads_q, num_heads_kv = 16, 4 + head_dim = 128 + num_tokens = 4 + eps = 1e-5 + max_position = 4096 + rotary_dim = head_dim + + total_dim = (num_heads_q + 2 * num_heads_kv) * head_dim + qkv = torch.randn(num_tokens, total_dim, dtype=dtype, device=device) + positions = torch.arange(num_tokens, dtype=torch.long, device=device) + + q_weight = torch.ones(head_dim, dtype=dtype, device=device) + k_weight = torch.ones(head_dim, dtype=dtype, device=device) + + inv_freq = 1.0 / (10000.0**( + torch.arange(0, rotary_dim // 2, dtype=torch.float32, device=device) / + (rotary_dim // 2))) + t = torch.arange(max_position, dtype=torch.float32, device=device) + freqs = torch.outer(t, inv_freq) + cos_sin_cache = torch.cat([freqs.cos(), freqs.sin()], dim=-1) + + opcheck( + torch.ops._C.fused_qk_norm_rope, + ( + qkv, + num_heads_q, + num_heads_kv, + num_heads_kv, + head_dim, + eps, + q_weight, + k_weight, + cos_sin_cache, + is_neox, + positions, + ), + ) diff --git a/tests/test_fused_quant_activation.py b/tests/test_fused_quant_activation.py new file mode 100644 index 000000000..99b4c8a4f --- /dev/null +++ b/tests/test_fused_quant_activation.py @@ -0,0 +1,118 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch +import torch.nn.functional as F + +import tests.register_ops as ops +from tests.utils import seed_everything + +DTYPES = [torch.half, torch.bfloat16] +FP8_DTYPES = [torch.float8_e4m3fn, torch.float8_e5m2] +NUM_TOKENS = [1, 7, 83, 512] +HIDDEN_SIZES = [16, 128, 512, 4096] +SEEDS = [0] +XPU_DEVICES = [ + f"xpu:{i}" for i in range(1 if torch.xpu.device_count() == 1 else 2) +] + +MINI_PYTEST_PARAMS = { + "default": { + "num_tokens": [1], + "HIDDEN_SIZES": [16], + }, +} + + +def ref_silu_and_mul_quant( + x: torch.Tensor, + scale: torch.Tensor, + fp8_dtype: torch.dtype, +) -> torch.Tensor: + """Reference implementation: SiLU+Mul then static FP8 quant.""" + d = x.shape[-1] // 2 + silu_mul_out = F.silu(x[..., :d]) * x[..., d:] + + fp8_max = torch.finfo(fp8_dtype).max + inv_scale = 1.0 / scale.item() + result = (silu_mul_out.float() * inv_scale).clamp(-fp8_max, fp8_max) + return result.to(fp8_dtype) + + +@pytest.mark.parametrize("num_tokens", NUM_TOKENS) +@pytest.mark.parametrize("hidden_size", HIDDEN_SIZES) +@pytest.mark.parametrize("dtype", DTYPES) +@pytest.mark.parametrize("fp8_dtype", FP8_DTYPES) +@pytest.mark.parametrize("seed", SEEDS) +@pytest.mark.parametrize("device", XPU_DEVICES) +@torch.inference_mode() +def test_silu_and_mul_quant( + num_tokens: int, + hidden_size: int, + dtype: torch.dtype, + fp8_dtype: torch.dtype, + seed: int, + device: str, +) -> None: + seed_everything(seed) + torch.set_default_device(device) + + x = torch.randn(num_tokens, hidden_size * 2, dtype=dtype) + scale = torch.tensor([0.5], dtype=torch.float32, device=device) + + # Reference + ref_out = ref_silu_and_mul_quant(x, scale, fp8_dtype) + + # Fused kernel + d = x.shape[-1] // 2 + out = torch.empty(num_tokens, d, dtype=fp8_dtype, device=device) + ops.silu_and_mul_quant(out, x, scale) + + assert out.dtype == fp8_dtype + assert out.shape == ref_out.shape + torch.testing.assert_close( + out.to(dtype=torch.float32), + ref_out.to(dtype=torch.float32), + ) + + +@pytest.mark.parametrize("num_tokens", NUM_TOKENS) +@pytest.mark.parametrize("hidden_size", HIDDEN_SIZES) +@pytest.mark.parametrize("dtype", DTYPES) +@pytest.mark.parametrize("fp8_dtype", FP8_DTYPES) +@pytest.mark.parametrize("seed", SEEDS) +@pytest.mark.parametrize("device", XPU_DEVICES) +@torch.inference_mode() +def test_silu_and_mul_quant_vs_separate( + num_tokens: int, + hidden_size: int, + dtype: torch.dtype, + fp8_dtype: torch.dtype, + seed: int, + device: str, +) -> None: + seed_everything(seed) + torch.set_default_device(device) + + x = torch.randn(num_tokens, hidden_size * 2, dtype=dtype) + scale = torch.tensor([0.5], dtype=torch.float32, device=device) + + # Separate ops + d = x.shape[-1] // 2 + silu_mul_out = torch.empty(num_tokens, d, dtype=dtype, device=device) + ops.silu_and_mul(silu_mul_out, x) + + separate_out = torch.empty(num_tokens, d, dtype=fp8_dtype, device=device) + ops.static_scaled_fp8_quant(separate_out, silu_mul_out, scale) + + # Fused kernel + fused_out = torch.empty(num_tokens, d, dtype=fp8_dtype, device=device) + ops.silu_and_mul_quant(fused_out, x, scale) + + assert fused_out.dtype == fp8_dtype + assert fused_out.shape == separate_out.shape + torch.testing.assert_close( + fused_out.to(dtype=torch.float32), + separate_out.to(dtype=torch.float32), + ) diff --git a/tests/test_scope_profiles.py b/tests/test_scope_profiles.py new file mode 100644 index 000000000..057797fa7 --- /dev/null +++ b/tests/test_scope_profiles.py @@ -0,0 +1,409 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +On-demand test scope profiles for model-specific validation. + +Each profile maps test file paths (relative, matching suffix) to per-function +parameter overrides. Use ``XPU_KERNEL_TEST_SCOPE=ondemand:`` to +activate a profile. + +Setting a function entry to ``None`` skips that test entirely. +Setting it to ``{}`` runs the test with original (full) parameters. + +Example usage: + XPU_KERNEL_TEST_SCOPE=ondemand:llama3 pytest -v -s tests/ + +To add a new profile: + 1. Add a new key to ONDEMAND_PROFILES + 2. Map relevant test files → functions → parameter overrides + 3. Only include tests/shapes that the target model actually uses +""" +import torch + +# --------------------------------------------------------------------------- +# Llama-family models (Llama-3-70B, CodeLlama, etc.) +# - SiluAndMul activation, RMSNorm, Rotary Embedding, attention with 64 heads +# and head size 128 +# - MHA: 64 heads, head_size 128 +# - Hidden: 8192, intermediate:28672 (Llama-3) +# - FP8/BF16 quantization +# --------------------------------------------------------------------------- +LLAMA3_HEAD_SIZE = 128 +LLAMA3_NUM_HEADS = 64 +LLAMA3_HIDDEN_SIZE = LLAMA3_HEAD_SIZE * LLAMA3_NUM_HEADS +LLAMA3_INTERMEDIATE_SIZE = 28672 +LLAMA3_NUM_KV_HEADS = 8 +_LLAMA3_PROFILE = { + "tests/test_activation.py": { + "test_act_and_mul": { + "activation": ["silu_and_mul"], + "num_tokens": [1, 128, 2048], + "d": [LLAMA3_INTERMEDIATE_SIZE], + }, + "test_activation": None, # Llama doesn't use standalone activations + }, + "tests/test_layernorm.py": { + "test_rms_norm": { + "num_tokens": [1, 128, 2048], + "hidden_size": [LLAMA3_HIDDEN_SIZE], + }, + }, + "tests/test_rotary_embedding.py": { + "test_rotary_embedding_opcheck": { + "is_neox_style": [True], + "max_position": [1024], + "head_size": [LLAMA3_HEAD_SIZE], + "seq_len": [1, 128, 1024], + }, + }, + "tests/test_cache.py": { + "test_reshape_and_cache_flash": { + "num_tokens": [1, 128], + "num_heads": [LLAMA3_NUM_HEADS], + "head_size": [LLAMA3_HEAD_SIZE], + "block_size": [64], + "dtype": [torch.bfloat16], + }, + }, + "tests/test_fp8_quant.py": { + "test_dynamic_per_tensor_fp8_quant": { + "num_tokens": [1, 128], + "hidden_size": [LLAMA3_HIDDEN_SIZE], + }, + "test_dynamic_per_token_fp8_quant": { + "num_tokens": [1, 128], + "hidden_size": [LLAMA3_HIDDEN_SIZE], + }, + }, + "tests/test_fp8_gemm_onednn.py": { + "test_fp8_gemm_per_tensor": { + "mnk_factors": + [(1, LLAMA3_HIDDEN_SIZE, LLAMA3_INTERMEDIATE_SIZE), + (128, LLAMA3_HIDDEN_SIZE, LLAMA3_INTERMEDIATE_SIZE)], + }, + "test_fp8_gemm_per_channel": { + "mnk_factors": + [(1, LLAMA3_HIDDEN_SIZE, LLAMA3_INTERMEDIATE_SIZE), + (128, LLAMA3_HIDDEN_SIZE, LLAMA3_INTERMEDIATE_SIZE)], + }, + }, + "tests/flash_attn/test_flash_attn_varlen_func.py": { + "test_varlen_with_paged_kv": { + "seq_lens": [[(1, 1)]], + "num_heads": [(LLAMA3_NUM_HEADS, LLAMA3_NUM_KV_HEADS)], + "head_size": [LLAMA3_HEAD_SIZE], + "num_blocks": [2048], + "window_size": [(-1, -1)], + "is_paged": [True], + }, + "test_decode_with_paged_kv": { + "seq_lens": [[(1, 1)]], + "num_heads": [(LLAMA3_NUM_HEADS, LLAMA3_NUM_KV_HEADS)], + "head_size": [LLAMA3_HEAD_SIZE], + "num_blocks": [2048], + "window_size": [(-1, -1)], + }, + }, +} + +LLAMA4_HEAD_SIZE = 128 +LLAMA4_NUM_HEADS = 40 +LLAMA4_HIDDEN_SIZE = LLAMA4_HEAD_SIZE * LLAMA4_NUM_HEADS # 5120 +LLAMA4_INTERMEDIATE_SIZE = 8192 +LLAMA4_NUM_KV_HEADS = 8 +LLAMA4_NUM_EXPERTS = 16 +LLAMA4_TOPK = 1 +_LLAMA4_PROFILE = { + # ---- Activation: SiluAndMul (SwiGLU) ---- + "tests/test_activation.py": { + "test_act_and_mul": { + "activation": ["silu_and_mul"], + "num_tokens": [1, 128, 2048], + "d": [LLAMA4_INTERMEDIATE_SIZE], + }, + "test_activation": None, # Scout doesn't use standalone activations + }, + # ---- RMSNorm ---- + "tests/test_layernorm.py": { + "test_rms_norm": { + "num_tokens": [1, 128, 2048], + "hidden_size": [LLAMA4_HIDDEN_SIZE], + }, + }, + # ---- Rotary Embedding: interleaved style (iRoPE) ---- + "tests/test_rotary_embedding.py": { + "test_rotary_embedding_opcheck": { + "is_neox_style": [False], # Llama 4 uses interleaved RoPE + "max_position": [1024], + "head_size": [LLAMA4_HEAD_SIZE], + "seq_len": [1, 128, 1024], + }, + }, + # ---- KV Cache: GQA (not MLA) ---- + "tests/test_cache.py": { + "test_reshape_and_cache_flash": { + "num_tokens": [1, 128], + "num_heads": [LLAMA4_NUM_KV_HEADS], + "head_size": [LLAMA4_HEAD_SIZE], + "block_size": [16], + "dtype": [torch.bfloat16], + }, + }, + # ---- TopK routing: softmax, top-1, 16 experts, use torch.topk, ignore ---- + # ---- Fused MoE: 16 experts, top-1 ---- + "tests/fused_moe/test_fused_moe.py": { + "test_fused_moe": { + "m,n,k": [(1, LLAMA4_INTERMEDIATE_SIZE, LLAMA4_HIDDEN_SIZE), + (128, LLAMA4_INTERMEDIATE_SIZE, LLAMA4_HIDDEN_SIZE)], + "e": [LLAMA4_NUM_EXPERTS], + "topk": [LLAMA4_TOPK], + "dtype": [torch.bfloat16], #FIXME: add low precision + "has_bias": [True, False], + }, + }, + # ---- Grouped GEMM: 16 experts, top-1 ---- + "tests/fused_moe/test_grouped_gemm.py": { + "test_grouped_gemm": { + "m,n,k": [(1, LLAMA4_INTERMEDIATE_SIZE, LLAMA4_HIDDEN_SIZE), + (128, LLAMA4_INTERMEDIATE_SIZE, LLAMA4_HIDDEN_SIZE)], + "e": [LLAMA4_NUM_EXPERTS], + "topk": [LLAMA4_TOPK], + "dtype": [torch.bfloat16], #FIXME: add low precision + "has_bias": [True, False], + }, + }, + # ---- MoE prologue ---- + "tests/fused_moe/test_moe_prologue.py": { + "test_prologue": { + "m,n,k": [(1, LLAMA4_INTERMEDIATE_SIZE, LLAMA4_HIDDEN_SIZE), + (128, LLAMA4_INTERMEDIATE_SIZE, LLAMA4_HIDDEN_SIZE)], + "e": [LLAMA4_NUM_EXPERTS], + "topk": [LLAMA4_TOPK], + }, + }, + # ---- MoE remap hidden states ---- + "tests/fused_moe/test_remap_hidden_states.py": { + "test_remap_hidden_states": { + "total_experts_num": [LLAMA4_NUM_EXPERTS], + "topk": [LLAMA4_TOPK], + }, + }, + # ---- MoE align block size ---- + "tests/test_moe_align_block_size.py": { + "test_moe_align_block_size": { + "m": [1, 128, 2048], + "num_experts": [LLAMA4_NUM_EXPERTS], + "topk": [LLAMA4_TOPK], + "block_size": [128], + }, + }, + # ---- MoE gather ---- + "tests/test_moe_gather.py": { + "test_moe_gather": { + "input_len": [1, 128], + "hidden_size": [LLAMA4_HIDDEN_SIZE], + "num_experts": [LLAMA4_NUM_EXPERTS], + "topk": [LLAMA4_TOPK], + }, + }, + # ---- MoE sum ---- + "tests/test_moe_sum.py": { + "test_moe_sum": { + "m": [1, 128], + "topk": [1], # top-1 sum is trivial; use smallest available + "k": [LLAMA4_HIDDEN_SIZE], + }, + }, + # ---- Flash Attention: GQA 40q/8kv heads ---- + "tests/flash_attn/test_flash_attn_varlen_func.py": { + "test_varlen_with_paged_kv": { + "seq_lens": [[(1, 1328), (5, 18), (129, 463)]], + "num_heads": [(LLAMA4_NUM_HEADS, LLAMA4_NUM_KV_HEADS)], + "head_size": [LLAMA4_HEAD_SIZE], + "num_blocks": [2048], + "window_size": [(-1, -1)], + "is_paged": [True], + }, + "test_decode_with_paged_kv": { + "seq_lens": [[(1, 1025), (1, 523), (1, 37)]], + "num_heads": [(LLAMA4_NUM_HEADS, LLAMA4_NUM_KV_HEADS)], + "head_size": [LLAMA4_HEAD_SIZE], + "num_blocks": [2048], + "window_size": [(-1, -1)], + }, + "test_decode_with_paged_kv_mla": None, # Not MLA + }, + # ---- Merge attention states ---- + "tests/test_merge_attn_states.py": { + "test_merge_attn_states": { + "num_tokens": [1, 128], + "num_query_heads": [LLAMA4_NUM_HEADS], + "head_size": [LLAMA4_HEAD_SIZE], + "output_dtype": [torch.bfloat16], + }, + }, + # ---- FP8 quantization ---- + "tests/test_fp8_quant.py": { + "test_dynamic_per_tensor_fp8_quant": { + "num_tokens": [1, 128], + "hidden_size": [LLAMA4_HIDDEN_SIZE], + }, + "test_dynamic_per_token_fp8_quant": { + "num_tokens": [1, 128], + "hidden_size": [LLAMA4_HIDDEN_SIZE], + }, + }, + # ---- FP8 GEMM ---- + "tests/test_fp8_gemm_onednn.py": { + "test_fp8_gemm_per_tensor": { + "mnk_factors": [ + (1, LLAMA4_HIDDEN_SIZE, LLAMA4_INTERMEDIATE_SIZE), + (128, LLAMA4_HIDDEN_SIZE, LLAMA4_INTERMEDIATE_SIZE), + ], + }, + "test_fp8_gemm_per_channel": { + "mnk_factors": [ + (1, LLAMA4_HIDDEN_SIZE, LLAMA4_INTERMEDIATE_SIZE), + (128, LLAMA4_HIDDEN_SIZE, LLAMA4_INTERMEDIATE_SIZE), + ], + }, + }, +} + +# --------------------------------------------------------------------------- +# DeepSeek-V3/R1 MLA models +# - MLA attention (Multi-head Latent Attention) +# - MoE with grouped topk (256 experts, top-8) +# - SiluAndMul, RMSNorm +# - kv_lora_rank=512, qk_rope_head_dim=64, v_head_dim=128 +# --------------------------------------------------------------------------- +_DEEPSEEK_PROFILE = { + "tests/test_activation.py": { + "test_act_and_mul": { + "activation": ["silu_and_mul"], + "num_tokens": [1, 128, 2048], + "d": [13824], + "dtype": [torch.bfloat16], + }, + }, + "tests/test_layernorm.py": { + "test_rms_norm": { + "num_tokens": [1, 128, 2048], + "hidden_size": [7168], + }, + }, + "tests/test_rotary_embedding.py": { + "test_rotary_embedding_opcheck": { + "is_neox_style": [True], + "max_position": [1024], + "head_size": [192], + "seq_len": [1, 128, 1024], + }, + }, + "tests/test_cache.py": { + "test_reshape_and_cache_flash": { + "num_tokens": [1, 128], + "num_heads": [8], + "head_size": [128], + "block_size": [16], + "dtype": [torch.bfloat16], + }, + "test_concat_and_cache_mla": { + "num_tokens": [1, 128], + "num_blocks": [4], + "block_size": [16], + }, + }, + "tests/test_topk.py": { + "test_fused_topk_softmax": { + "topk": [8], + "n_expert": [256], + "n_token": [1, 128, 2048], + }, + }, + "tests/test_grouped_topk.py": { + "test_grouped_topk": { + "n_hidden": [256], + "n_token": [1, 128], + "topk": [8], + "n_group": [8], + "renormalize": [True], + "scoring_func": ["softmax"], + }, + }, + "tests/fused_moe/test_fused_moe.py": { + "test_fused_moe": { + "m,n,k": [(1, 5120, 7168), (128, 5120, 7168)], + "e": [256], + "topk": [8], + "dtype": [torch.bfloat16], + "has_bias": [True], + }, + }, + "tests/fused_moe/test_grouped_gemm.py": { + "test_grouped_gemm": { + "m,n,k": [(1, 5120, 7168), (128, 5120, 7168)], + "e": [256], + "topk": [8], + "dtype": [torch.bfloat16], + "has_bias": [True], + }, + }, + "tests/test_moe_align_block_size.py": { + "test_moe_align_block_size": { + "m": [1, 128, 2048], + "num_experts": [256], + "topk": [8], + "block_size": [128], + }, + }, + "tests/flash_attn/test_flash_attn_varlen_func.py": { + "test_decode_with_paged_kv_mla": { + "seq_lens": [[(1, 1025), (1, 523), (1, 37)]], + "num_heads": [(8, 1)], + "head_size_kv": [(192, 128)], + "num_blocks": [2048], + }, + "test_varlen_with_paged_kv": { + "seq_lens": [[(1, 1328), (5, 18), (129, 463)]], + "num_heads": [(8, 1)], + "head_size": [128], + "num_blocks": [2048], + "window_size": [(-1, -1)], + "is_paged": [True], + }, + }, + "tests/test_merge_attn_states.py": { + "test_merge_attn_states": { + "num_tokens": [1, 128], + "num_query_heads": [128], + "head_size": [128], + "output_dtype": [torch.bfloat16], + }, + }, + "tests/test_fp8_quant.py": { + "test_dynamic_per_token_fp8_quant": { + "num_tokens": [1, 128], + "hidden_size": [7168], + }, + "test_per_block_fp8_quant": { + "num_tokens_block_quant": [1, 128], + "hidden_size_block_quant": [7168], + }, + }, + "tests/test_swigluoai_and_mul.py": { + "test_act_and_mul": { + "num_tokens": [1, 128, 2048], + "d": [13824], + "dtype": [torch.bfloat16], + }, + }, +} + +# --------------------------------------------------------------------------- +# Registry of all on-demand profiles +# --------------------------------------------------------------------------- +ONDEMAND_PROFILES = { + "llama3": _LLAMA3_PROFILE, + "llama4": _LLAMA4_PROFILE, + "deepseek": _DEEPSEEK_PROFILE, +} diff --git a/tests/test_topk_topp_sampler.py b/tests/test_topk_topp_sampler.py index 7390f0bb0..be932ea74 100755 --- a/tests/test_topk_topp_sampler.py +++ b/tests/test_topk_topp_sampler.py @@ -13,6 +13,17 @@ P = [0.1, 0.2, 0.4, 0.8, 1.0, None] LOGPROBS_MODE = ["raw_logits", "processed_logits", "processed_logprobs"] +# CI/mini scope parameter overrides +MINI_PYTEST_PARAMS = { + "default": { + "batch_size": [1, 32], + "vocab_size": [1024], + "k": [1, 128, None], + "p": [0.5, None], + "logprobs_mode": ["raw_logits"], + }, +} + @pytest.mark.parametrize("batch_size", BATCH_SIZE) @pytest.mark.parametrize("vocab_size", VOCAB_SIZE) diff --git a/tests/test_uva.py b/tests/test_uva.py index 5c3c60fed..81050393f 100644 --- a/tests/test_uva.py +++ b/tests/test_uva.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import os +import gc import pytest import torch @@ -11,12 +11,18 @@ f"xpu:{i}" for i in range(1 if torch.xpu.device_count() == 1 else 2) ] -SKIP_TEST_FOR_MINI_SCOPE = os.getenv("XPU_KERNEL_PYTEST_PROFILER") == "MINI" +# Skip entire module in mini scope +SKIP_IN_MINI_SCOPE = True + +# CI scope parameter overrides +MINI_PYTEST_PARAMS = { + "default": { + "device": ["xpu:0"], + }, +} @pytest.mark.parametrize("device", XPU_DEVICES) -@pytest.mark.skipif(SKIP_TEST_FOR_MINI_SCOPE, - reason="Skip UVA tests for the mini pytest profiler.") def test_cpu_write(device): torch.set_default_device(device) cpu_tensor = torch.zeros(10, @@ -42,8 +48,6 @@ def test_cpu_write(device): @pytest.mark.parametrize("device", XPU_DEVICES) -@pytest.mark.skipif(SKIP_TEST_FOR_MINI_SCOPE, - reason="Skip UVA tests for the mini pytest profiler.") def test_gpu_write(device): torch.set_default_device(device) cpu_tensor = torch.zeros(10, @@ -66,3 +70,23 @@ def test_gpu_write(device): assert cpu_tensor[0, 0] == 2 assert cpu_tensor[2, 3] == 4 assert cpu_tensor[4, 5] == -2 + + +@pytest.mark.parametrize("device", XPU_DEVICES) +def test_view_lifetime_after_owner_drop(device): + torch.set_default_device(device) + cpu_tensor = torch.arange(100, + dtype=torch.int32, + device="cpu", + pin_memory=True).view(10, 10) + xpu_view = torch.ops._C.get_xpu_view_from_cpu_tensor(cpu_tensor) + + # Drop the original owner reference and force Python GC. + del cpu_tensor + gc.collect() + + # Exercise both read and write from the XPU view after owner drop. + assert xpu_view[2, 3].item() == 23 + xpu_view.add_(1) + assert xpu_view[0, 0].item() == 1 + assert xpu_view[9, 9].item() == 100 diff --git a/tests/test_xpu_memcpy_sync.py b/tests/test_xpu_memcpy_sync.py index 64ee426c4..8d1dc0953 100644 --- a/tests/test_xpu_memcpy_sync.py +++ b/tests/test_xpu_memcpy_sync.py @@ -14,6 +14,13 @@ f"xpu:{i}" for i in range(1 if torch.xpu.device_count() == 1 else 2) ] +# CI/mini scope parameter overrides +MINI_PYTEST_PARAMS = { + "default": { + "device": ["xpu:0"], + }, +} + @pytest.mark.parametrize("device", XPU_DEVICES) def test_xpu_memcpy_sync_host_to_device(device: str) -> None: diff --git a/tools/envs.py b/tools/envs.py index 94b5bfb5c..924992a09 100644 --- a/tools/envs.py +++ b/tools/envs.py @@ -100,6 +100,34 @@ def get_vllm_port() -> Optional[int]: "CMAKE_BUILD_TYPE": lambda: os.getenv("CMAKE_BUILD_TYPE"), + # ================== Kernel Build Options ================== + # These control which kernel extensions are built. Set to "0" or "OFF" + # to disable. They are forwarded to CMake as -D flags by setup.py. + + # Build SYCL-TLA based kernels (attention, grouped_gemm shared libs) + "BUILD_SYCL_TLA_KERNELS": + lambda: os.getenv("BUILD_SYCL_TLA_KERNELS", "ON"), + + # Architecture options + "VLLM_XPU_ENABLE_XE2": + lambda: os.getenv("VLLM_XPU_ENABLE_XE2", "ON"), + "VLLM_XPU_ENABLE_XE_DEFAULT": + lambda: os.getenv("VLLM_XPU_ENABLE_XE_DEFAULT", "ON"), + + # Individual kernel extension toggles + "BASIC_KERNELS_ENABLED": + lambda: os.getenv("BASIC_KERNELS_ENABLED", "ON"), + "FA2_KERNELS_ENABLED": + lambda: os.getenv("FA2_KERNELS_ENABLED", "ON"), + "MOE_KERNELS_ENABLED": + lambda: os.getenv("MOE_KERNELS_ENABLED", "ON"), + "GDN_KERNELS_ENABLED": + lambda: os.getenv("GDN_KERNELS_ENABLED", "ON"), + "XPU_SPECIFIC_KERNELS_ENABLED": + lambda: os.getenv("XPU_SPECIFIC_KERNELS_ENABLED", "ON"), + "XPUMEM_ALLOCATOR_ENABLED": + lambda: os.getenv("XPUMEM_ALLOCATOR_ENABLED", "ON"), + # If set, vllm will print verbose logs during installation "VERBOSE": lambda: bool(int(os.getenv('VERBOSE', '0'))),