From b3701c99681d1a30e37059f00046b388f02ecfff Mon Sep 17 00:00:00 2001 From: baodi Date: Wed, 1 Apr 2026 09:58:44 +0800 Subject: [PATCH 1/3] Support arbitrary KV cache strides in paged_decode for MLA (#165) * Support arbitrary KV cache strides in paged_decode for MLA - Remove CHECK_CONTIGUOUS for k/v in flash_api.cpp (stride(-1)==1 still enforced) - Add KV cache stride fields to paged_decode_args_t - Extract actual tensor strides in paged_decode_xe2.cpp - Use actual strides in DecodeKernelLauncher::initialize() instead of packed strides - Replace make_ordered_layout with make_layout using passed strides for K/V in kernel - Add test_decode_with_paged_kv_mla unit test with non-contiguous KV cache slices Signed-off-by: baodii * make format happy Signed-off-by: baodii * make format happy Signed-off-by: baodii * Apply suggestions from code review Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> Signed-off-by: baodi * add TORCH_CHECK to aviod potential bugs Signed-off-by: baodii * fix kv stride bug Signed-off-by: baodii * fix UT file Signed-off-by: baodii --------- Signed-off-by: baodii Signed-off-by: baodi Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- CMakeLists.txt | 5 +- csrc/flash_attn/flash_api.cpp | 18 ++- .../attn/xe_2/kernel/paged_decode_kernel.hpp | 6 +- csrc/xpu/attn/xe_2/paged_decode.hpp | 57 +++++++- csrc/xpu/attn/xe_2/paged_decode_xe2.cpp | 25 +++- .../flash_attn/test_flash_attn_varlen_func.py | 126 +++++++++++++++++- 6 files changed, 212 insertions(+), 25 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 249c2cbcd..19213a5da 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -187,8 +187,9 @@ if(VLLM_GPU_LANG STREQUAL "SYCL") set(SYCL_LINK_FLAGS "") list(APPEND SYCL_LINK_FLAGS "-fsycl") set(SYCL_DEVICE_LINK_FLAGS ${SYCL_LINK_FLAGS}) - set(SYCL_DEVICE_LINK_FLAGS ${SYCL_DEVICE_LINK_FLAGS} - -fsycl-max-parallel-link-jobs=16) + set(SYCL_DEVICE_LINK_FLAGS + ${SYCL_DEVICE_LINK_FLAGS} -fsycl-max-parallel-link-jobs=16 + -flink-huge-device-code) set(SYCL_DEVICE_LINK_FLAGS ${SYCL_DEVICE_LINK_FLAGS} "-Xspirv-translator;-spirv-ext=+SPV_INTEL_split_barrier,+SPV_INTEL_2d_block_io,+SPV_INTEL_subgroup_matrix_multiply_accumulate" diff --git a/csrc/flash_attn/flash_api.cpp b/csrc/flash_attn/flash_api.cpp index 972cd6f2c..dade9f882 100644 --- a/csrc/flash_attn/flash_api.cpp +++ b/csrc/flash_attn/flash_api.cpp @@ -110,8 +110,6 @@ std::vector mha_varlen_fwd( v.stride(-1) == 1, "Input tensor must have contiguous last dimension"); TORCH_CHECK(q.dim() == 3, "query must be in ragged format"); CHECK_CONTIGUOUS(q); - CHECK_CONTIGUOUS(k); - CHECK_CONTIGUOUS(v); at::Tensor block_table; bool is_paged = block_table_.has_value(); @@ -143,8 +141,6 @@ std::vector mha_varlen_fwd( at::Tensor out; if (out_.has_value()) { out = *out_; - } else { - out = torch::empty_like(q); } bool is_varlen = true; @@ -152,6 +148,9 @@ std::vector mha_varlen_fwd( bool is_sink = softmax_sink_.has_value(); if (max_seqlen_q > 1 || !is_paged) { + if (!out_.has_value()) { + out = torch::empty_like(q); + } at::Tensor seqlens_k = is_paged ? *seqused_k : cu_seqlens_k; cutlass_chunk_prefill_interface( @@ -189,10 +188,17 @@ std::vector mha_varlen_fwd( int num_tokens = q.size(0); int batch_size = static_cast(cu_seqlens_q.size(0)) - 1; int num_heads_q = q.size(1); - int head_dim = q.size(2); + int v_head_dim = v.size(-1); int num_heads_kv = k.size(2); int block_size = k.size(1); + // Output shape uses V's head_dim (may differ from Q/K for MLA) + if (!out_.has_value()) { + out = torch::empty( + {num_tokens, num_heads_q, v_head_dim}, + q.options().device(q.device())); + } + int num_kv_splits = num_splits.value_or(get_num_splits( queue, batch_size, num_heads_kv, effective_seqlen_k, block_size)); @@ -200,7 +206,7 @@ std::vector mha_varlen_fwd( num_kv_splits == 1 ? out : at::empty( - {num_tokens, num_heads_q * num_kv_splits, head_dim}, + {num_tokens, num_heads_q * num_kv_splits, v_head_dim}, q.options().device(q.device())); at::Tensor max_logits = at::full( {num_tokens, num_heads_q, num_kv_splits}, 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 e0a8e4dbe..cde1266e3 100644 --- a/csrc/xpu/attn/xe_2/kernel/paged_decode_kernel.hpp +++ b/csrc/xpu/attn/xe_2/kernel/paged_decode_kernel.hpp @@ -356,8 +356,10 @@ class XeFMHAFwdSplitKVKernel { auto ptrMax_logits = p.max_logits + offset_max_logits; auto layout_q = make_ordered_layout(shape_Q, Step<_1, _0, _2, _3>{}); - auto layout_k = make_ordered_layout(shape_K, Step<_2, _0, _1, _3>{}); - auto layout_v = make_ordered_layout(shape_V, Step<_0, _2, _1, _3>{}); + auto layout_k = make_layout( + shape_K, make_stride(get<0>(p.dK), _1{}, get<2>(p.dK), get<3>(p.dK))); + auto layout_v = make_layout( + shape_V, make_stride(_1{}, get<1>(p.dV), get<2>(p.dV), get<3>(p.dV))); auto layout_o = make_ordered_layout(shape_O, Step<_1, _0, _2, _3, _4>{}); auto layout_exp_sums = diff --git a/csrc/xpu/attn/xe_2/paged_decode.hpp b/csrc/xpu/attn/xe_2/paged_decode.hpp index b35a0d6db..c694b687a 100644 --- a/csrc/xpu/attn/xe_2/paged_decode.hpp +++ b/csrc/xpu/attn/xe_2/paged_decode.hpp @@ -69,6 +69,7 @@ struct paged_decode_args_t { int num_heads_q; int num_heads_k; int head_size; + int v_head_size; int max_blocks_per_seq; int block_size; int window_size_left = -1; @@ -79,6 +80,13 @@ struct paged_decode_args_t { bool is_local = false; bool is_sink = false; int num_kv_splits = 1; + // KV cache strides [num_blocks, block_size, num_heads_kv, head_size] + int64_t k_stride_page = 0; + int64_t k_stride_seq = 0; + int64_t k_stride_heads = 0; + int64_t v_stride_page = 0; + int64_t v_stride_seq = 0; + int64_t v_stride_heads = 0; }; template @@ -122,7 +130,7 @@ struct DecodeKernelLauncher { auto head_size_qk = shape.head_size_qk = shape_init.head_size_qk = args.head_size; auto head_size_vo = shape.head_size_vo = shape_init.head_size_vo = - args.head_size; + args.v_head_size; if constexpr (isVarLen) { batch = shape_init.batch = 1; @@ -150,12 +158,47 @@ struct DecodeKernelLauncher { stride_Q = cutlass::make_cute_packed_stride( StrideQ{}, cute::make_shape(seq_len_qo, head_size_qk, num_heads_q, batch)); - stride_K = cutlass::make_cute_packed_stride( - StrideK{}, - cute::make_shape(seq_len_kv, head_size_qk, num_heads_kv, batch)); - stride_V = cutlass::make_cute_packed_stride( - StrideV{}, - cute::make_shape(head_size_vo, seq_len_kv, num_heads_kv, batch)); + if (args.k_stride_seq > 0) { + // Use actual strides from KV cache tensors (supports non-contiguous + // layouts such as MLA combined KV cache) + constexpr int64_t kIntMax = + static_cast(std::numeric_limits::max()); + TORCH_CHECK( + args.k_stride_seq <= kIntMax && args.k_stride_heads <= kIntMax && + args.k_stride_page <= kIntMax && args.v_stride_seq <= kIntMax && + args.v_stride_heads <= kIntMax && args.v_stride_page <= kIntMax, + "KV cache stride exceeds int32 max (", + kIntMax, + "): k_stride_seq=", + args.k_stride_seq, + " k_stride_heads=", + args.k_stride_heads, + " k_stride_page=", + args.k_stride_page, + " v_stride_seq=", + args.v_stride_seq, + " v_stride_heads=", + args.v_stride_heads, + " v_stride_page=", + args.v_stride_page); + stride_K = StrideK{ + static_cast(args.k_stride_seq), + _1{}, + static_cast(args.k_stride_heads), + static_cast(args.k_stride_page)}; + stride_V = StrideV{ + _1{}, + static_cast(args.v_stride_seq), + static_cast(args.v_stride_heads), + static_cast(args.v_stride_page)}; + } else { + stride_K = cutlass::make_cute_packed_stride( + StrideK{}, + cute::make_shape(seq_len_kv, head_size_qk, num_heads_kv, batch)); + stride_V = cutlass::make_cute_packed_stride( + StrideV{}, + cute::make_shape(head_size_vo, seq_len_kv, num_heads_kv, batch)); + } stride_O = cutlass::make_cute_packed_stride( StrideO{}, cute::make_shape(seq_len_qo, head_size_vo, num_heads_q, batch)); diff --git a/csrc/xpu/attn/xe_2/paged_decode_xe2.cpp b/csrc/xpu/attn/xe_2/paged_decode_xe2.cpp index 6232b12c2..615372e0d 100644 --- a/csrc/xpu/attn/xe_2/paged_decode_xe2.cpp +++ b/csrc/xpu/attn/xe_2/paged_decode_xe2.cpp @@ -116,7 +116,7 @@ void cutlass_paged_decode_impl( } // general params - int batch_size, num_heads_q, num_heads_kv, head_size; + int batch_size, num_heads_q, num_heads_kv, head_size, v_head_size; // additional params int total_seqlen_q, total_seqlen_k; int num_blocks, block_size, max_blocks_per_seq; @@ -126,6 +126,7 @@ void cutlass_paged_decode_impl( num_heads_q = query.size(1); num_heads_kv = key_cache.size(1); head_size = query.size(2); + v_head_size = value_cache.size(-1); total_seqlen_q = query.size(0); total_seqlen_k = key_cache.size(0); } else { @@ -134,6 +135,7 @@ void cutlass_paged_decode_impl( num_heads_q = query.size(1); num_heads_kv = key_cache.size(1); head_size = query.size(3); + v_head_size = value_cache.size(-1); max_seqlen_q = query.size(2); max_seqlen_k = key_cache.size(2); } @@ -176,6 +178,7 @@ void cutlass_paged_decode_impl( num_heads_q, num_heads_kv, head_size, + v_head_size, max_blocks_per_seq, block_size, window_size_left, @@ -185,7 +188,25 @@ void cutlass_paged_decode_impl( is_causal, is_local, is_sink, - num_kv_splits}; + num_kv_splits, + // KV cache strides + key_cache.stride(0), + key_cache.stride(1), + key_cache.stride(2), + value_cache.stride(0), + value_cache.stride(1), + value_cache.stride(2)}; + + TORCH_CHECK( + key_cache.stride(-1) == 1, + "paged_decode_xe2: key_cache must be contiguous in the last dimension " + "(head_dim), got stride=", + key_cache.stride(-1)); + TORCH_CHECK( + value_cache.stride(-1) == 1, + "paged_decode_xe2: value_cache must be contiguous in the last dimension " + "(head_dim), got stride=", + value_cache.stride(-1)); CutlassQKType cuQKType = aten_to_Cutlass_qk_dtype(query, key_cache); diff --git a/tests/flash_attn/test_flash_attn_varlen_func.py b/tests/flash_attn/test_flash_attn_varlen_func.py index 8f6a68752..9264d8ccd 100644 --- a/tests/flash_attn/test_flash_attn_varlen_func.py +++ b/tests/flash_attn/test_flash_attn_varlen_func.py @@ -10,20 +10,20 @@ from tests.utils import format_tc from vllm_xpu_kernels.flash_attn_interface import flash_attn_varlen_func -NUM_HEADS = [(4, 4), (8, 2), (10, 2), (16, 1)] +NUM_HEADS = [(8, 2)] HEAD_SIZES = [64, 128, 192, 256] BLOCK_SIZES = [64, 128] -DTYPES = [torch.bfloat16, torch.half] +DTYPES = [torch.bfloat16] QDTYPES = [None] # one value large enough to test overflow in index calculation. # one value small enough to test the schema op check -NUM_BLOCKS = [32768, 2048] +NUM_BLOCKS = [2048] SOFT_CAPS = [None] SLIDING_WINDOWS = [(-1, 127), (127, -1), (64, 64), (-1, -1)] SINK = [False, True] CASUAL = [False, True] PAGED = [False, True] -FP8KV = [torch.float8_e5m2, torch.float8_e4m3fn, None] +FP8KV = [torch.float8_e4m3fn, None] def ref_paged_attn(query: torch.Tensor, @@ -49,8 +49,10 @@ def ref_paged_attn(query: torch.Tensor, block_tables = block_tables.cpu().numpy() if is_paged: _, block_size, num_kv_heads, head_size = key_cache.shape + v_head_size = value_cache.shape[-1] else: _, num_kv_heads, head_size = key_cache.shape + v_head_size = value_cache.shape[-1] if is_fp8_query: query = (query.to(torch.float32) * q_descale).to(dtype) @@ -69,7 +71,7 @@ def ref_paged_attn(query: torch.Tensor, k = key_cache[block_indices].view(-1, num_kv_heads, head_size) k = k[:kv_len] - v = value_cache[block_indices].view(-1, num_kv_heads, head_size) + v = value_cache[block_indices].view(-1, num_kv_heads, v_head_size) v = v[:kv_len] else: k = key_cache[start_idx_kv:start_idx_kv + kv_len] @@ -137,8 +139,13 @@ def ref_paged_attn(query: torch.Tensor, "num_heads": [(8, 2)], "head_size": [64, 128], "num_blocks": [64], - "fp8_dtype": [torch.float8_e4m3fn, None], "window_size": [(-1, -1), (127, -1)], + }, + "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], } } @@ -466,3 +473,110 @@ def test_decode_with_paged_kv( torch.testing.assert_close(output, ref_output, atol=atol, rtol=rtol), \ f"{torch.max(torch.abs(output - ref_output))}" torch.xpu.empty_cache() + + +@pytest.mark.parametrize( + "seq_lens", + [[(1, 1025)], [(1, 523), (1, 37), + (1, 2011)], [(1, 523), (1, 37), (1, 2011), (1, 5000)]]) +@pytest.mark.parametrize("num_heads", [(8, 1), (16, 1), (8, 2)]) +@pytest.mark.parametrize("head_size_kv", [(192, 128)]) +@pytest.mark.parametrize("block_size", BLOCK_SIZES) +@pytest.mark.parametrize("dtype", [torch.bfloat16]) +@pytest.mark.parametrize("num_blocks", [2048]) +@torch.inference_mode() +def test_decode_with_paged_kv_mla( + seq_lens: list[tuple[int, int]], + num_heads: tuple[int, int], + head_size_kv: tuple[int, int], + dtype: torch.dtype, + block_size: int, + num_blocks: int, +) -> None: + """Test paged decode with MLA-like KV cache layout. + + MLA stores K and V in a shared buffer of shape [..., k_head_size]. + K uses the full buffer (head_dim=192, rope + nope), while V is a + slice of the first v_head_size dims (head_dim=128). V is therefore + non-contiguous with stride(-1)==1. + + The kernel computes attention scores using K (head_dim=k_head_size) and + multiplies by V (head_dim=v_head_size), producing output with + head_dim=v_head_size. + """ + torch.set_default_device("xpu") + torch.xpu.set_device("xpu:0") + torch.manual_seed(42) + num_seqs = len(seq_lens) + query_lens = [x[0] for x in seq_lens] + kv_lens = [x[1] for x in seq_lens] + num_query_heads = num_heads[0] + num_kv_heads = num_heads[1] + assert num_query_heads % num_kv_heads == 0 + max_query_len = max(query_lens) + max_kv_len = max(kv_lens) + + k_head_size, v_head_size = head_size_kv + scale = k_head_size**-0.5 + + query = torch.randn(sum(query_lens), + num_query_heads, + k_head_size, + dtype=dtype) + + # MLA-like combined KV cache: buffer has k_head_size dims. + # K uses the full buffer (head_dim=192), V is the first v_head_size + # dims (head_dim=128) — a non-contiguous slice. + combined_kv_cache = torch.randn(num_blocks, + block_size, + num_kv_heads, + k_head_size, + dtype=dtype) + + key_cache = combined_kv_cache + value_cache = combined_kv_cache[..., :v_head_size] + + assert key_cache.is_contiguous(), "key_cache should be contiguous" + assert not value_cache.is_contiguous(), \ + "value_cache should be non-contiguous" + assert value_cache.stride(-1) == 1 + + cu_query_lens = torch.tensor([0] + query_lens, + dtype=torch.int32).cumsum(dim=0, + dtype=torch.int32) + seq_k = torch.tensor(kv_lens, dtype=torch.int32) + + max_num_blocks_per_seq = (max_kv_len + block_size - 1) // block_size + block_tables = torch.randint(0, + num_blocks, + (num_seqs, max_num_blocks_per_seq), + dtype=torch.int32) + + output = flash_attn_varlen_func(query, + key_cache, + value_cache, + max_query_len, + cu_query_lens, + max_kv_len, + seqused_k=seq_k, + softmax_scale=scale, + causal=False, + block_table=block_tables, + window_size=(-1, -1)) + + ref_output = ref_paged_attn(query=query, + key_cache=key_cache, + value_cache=value_cache, + query_lens=query_lens, + kv_lens=kv_lens, + block_tables=block_tables, + scale=scale, + casual=False, + is_paged=True, + sink=None, + window_size_left=-1, + window_size_right=-1) + atol, rtol = 1e-2, 1e-2 + torch.testing.assert_close(output, ref_output, atol=atol, rtol=rtol), \ + f"{torch.max(torch.abs(output - ref_output))}" + torch.xpu.empty_cache() From 74657ec68fba8412a63913177d476470aa1dc5b7 Mon Sep 17 00:00:00 2001 From: zhenwei-intel Date: Tue, 31 Mar 2026 08:21:39 +0000 Subject: [PATCH 2/3] ggml dequantize Signed-off-by: zhenwei-intel --- .github/workflows/ut.yaml | 2 + CMakeLists.txt | 1 + csrc/quantization/gguf/ggml_dequantize.cpp | 114 +++++++++++ csrc/quantization/gguf/ggml_dequantize.hpp | 175 +++++++++++++++++ csrc/xpu/ops.h | 7 + csrc/xpu/torch_bindings.cpp | 5 + tests/register_ops.py | 8 + tests/test_gguf.py | 214 +++++++++++++++++++++ 8 files changed, 526 insertions(+) create mode 100644 csrc/quantization/gguf/ggml_dequantize.cpp create mode 100644 csrc/quantization/gguf/ggml_dequantize.hpp create mode 100644 tests/test_gguf.py diff --git a/.github/workflows/ut.yaml b/.github/workflows/ut.yaml index 8b1c8932b..d71109169 100644 --- a/.github/workflows/ut.yaml +++ b/.github/workflows/ut.yaml @@ -82,6 +82,7 @@ jobs: ccache -p || true git submodule sync && git submodule update --init --recursive uv pip install -r requirements.txt + uv pip install gguf MAX_JOBS=128 uv pip install --no-build-isolation -e . -v ccache -s || true @@ -124,6 +125,7 @@ jobs: ccache -p || true git submodule sync && git submodule update --init --recursive uv pip install -r requirements.txt + uv pip install gguf MAX_JOBS=80 uv pip install --no-build-isolation -e . -v ccache -s || true diff --git a/CMakeLists.txt b/CMakeLists.txt index 19213a5da..d0d570cec 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -479,6 +479,7 @@ if(XPU_SPECIFIC_KERNELS_ENABLED) "csrc/xpu/lora/lora_shrink.cpp" "csrc/xpu/lora/lora_expand.cpp" "csrc/xpu/sampler/topk_topp_sampler.cpp" + "csrc/quantization/gguf/ggml_dequantize.cpp" "csrc/xpu/sycl/deepseek_scaling_rope.cpp" "csrc/xpu/rand/exponential.cpp" "csrc/xpu/grouped_gemm/grouped_gemm_interface.cpp" diff --git a/csrc/quantization/gguf/ggml_dequantize.cpp b/csrc/quantization/gguf/ggml_dequantize.cpp new file mode 100644 index 000000000..cdc541990 --- /dev/null +++ b/csrc/quantization/gguf/ggml_dequantize.cpp @@ -0,0 +1,114 @@ +#include "ggml_dequantize.hpp" +#include "utils.h" +#include "xpu/ops.h" + +#include +#include + +namespace ggml = vllm::ggml; + +torch::Tensor ggml_dequantize( + const torch::Tensor& W, + int64_t type, + int64_t m, + int64_t n, + std::optional out_dtype) { + CHECK_DEVICE(W); + CHECK_CONTIGUOUS(W); + + TORCH_CHECK( + type == ggml::GGML_TYPE_Q4_0 || type == ggml::GGML_TYPE_Q5_0 || + type == ggml::GGML_TYPE_Q8_0, + "XPU ggml_dequantize currently only supports Q4_0 (type=2), " + "Q5_0 (type=6) and Q8_0 (type=8), got ", + type); + TORCH_CHECK( + W.scalar_type() == at::ScalarType::Byte, + "XPU ggml_dequantize expects uint8 weights, got ", + W.scalar_type()); + TORCH_CHECK(m >= 0 && n >= 0, "m and n must be non-negative"); + + const int64_t numel = m * n; + TORCH_CHECK( + numel % ggml::QK4_0 == 0, + ggml::ggml_type_name(type), + " dequantize expects m * n to be divisible by ", + ggml::QK4_0, + ", got ", + numel); + + const int64_t expected_nbytes = ggml::get_expected_nbytes(type, numel); + const int64_t weight_nbytes = W.numel() * W.element_size(); + TORCH_CHECK( + weight_nbytes == expected_nbytes, + ggml::ggml_type_name(type), + " packed weight size mismatch: expected ", + expected_nbytes, + " bytes for shape (", + m, + ", ", + n, + "), got ", + weight_nbytes, + " bytes"); + + const auto dtype = out_dtype.value_or(torch::kFloat16); + TORCH_CHECK( + dtype == torch::kFloat16 || dtype == torch::kBFloat16 || + dtype == torch::kFloat32, + "XPU ggml_dequantize only supports fp16, bf16 or fp32 outputs, got ", + dtype); + + auto options = torch::TensorOptions().dtype(dtype).device(W.device()); + auto output = torch::empty({m, n}, options); + if (numel == 0) { + return output; + } + + at::DeviceGuard device_guard(W.device()); + auto& queue = vllm::xpu::vllmGetQueue(W.device().index()); + const auto* weight_ptr = W.data_ptr(); + + VLLM_DISPATCH_FLOATING_TYPES(output.scalar_type(), "ggml_dequantize", [&] { + using sycl_t = typename vllm::xpu::SyclTypeTrait::Type; + auto* out_ptr = reinterpret_cast(output.data_ptr()); + + switch (type) { + case ggml::GGML_TYPE_Q4_0: { + auto* blocks = reinterpret_cast(weight_ptr); + queue.submit([&](sycl::handler& cgh) { + cgh.parallel_for( + sycl::range<1>(static_cast(numel)), + ggml::ggml_dequantize_q4_0_kernel( + blocks, out_ptr, numel)); + }); + break; + } + case ggml::GGML_TYPE_Q5_0: { + auto* blocks = reinterpret_cast(weight_ptr); + queue.submit([&](sycl::handler& cgh) { + cgh.parallel_for( + sycl::range<1>(static_cast(numel)), + ggml::ggml_dequantize_q5_0_kernel( + blocks, out_ptr, numel)); + }); + break; + } + case ggml::GGML_TYPE_Q8_0: { + auto* blocks = reinterpret_cast(weight_ptr); + queue.submit([&](sycl::handler& cgh) { + cgh.parallel_for( + sycl::range<1>(static_cast(numel)), + ggml::ggml_dequantize_q8_0_kernel( + blocks, out_ptr, numel)); + }); + break; + } + default: + TORCH_CHECK( + false, "Unsupported GGML type for XPU ggml_dequantize: ", type); + } + }); + + return output; +} \ No newline at end of file diff --git a/csrc/quantization/gguf/ggml_dequantize.hpp b/csrc/quantization/gguf/ggml_dequantize.hpp new file mode 100644 index 000000000..9525ab0b0 --- /dev/null +++ b/csrc/quantization/gguf/ggml_dequantize.hpp @@ -0,0 +1,175 @@ +#pragma once + +#include "dispatch_utils.h" +#include "utils.h" + +#include +#include + +namespace vllm { +namespace ggml { + +constexpr int64_t GGML_TYPE_Q4_0 = 2; +constexpr int64_t GGML_TYPE_Q5_0 = 6; +constexpr int64_t GGML_TYPE_Q8_0 = 8; +constexpr int64_t QK4_0 = 32; +constexpr int64_t QK5_0 = 32; +constexpr int64_t QK8_0 = 32; + +struct block_q4_0 { + sycl::half d; + uint8_t qs[QK4_0 / 2]; +}; + +static_assert(sizeof(block_q4_0) == 18, "Unexpected Q4_0 block size"); + +struct block_q5_0 { + sycl::half d; + uint8_t qh[4]; + uint8_t qs[QK5_0 / 2]; +}; + +static_assert(sizeof(block_q5_0) == 22, "Unexpected Q5_0 block size"); + +struct block_q8_0 { + sycl::half d; + int8_t qs[QK8_0]; +}; + +static_assert(sizeof(block_q8_0) == 34, "Unexpected Q8_0 block size"); + +inline uint32_t load_u32_le(const uint8_t* bytes) { + return static_cast(bytes[0]) | + (static_cast(bytes[1]) << 8) | + (static_cast(bytes[2]) << 16) | + (static_cast(bytes[3]) << 24); +} + +template +class ggml_dequantize_q4_0_kernel { + public: + using sycl_t = typename vllm::xpu::SyclTypeTrait::Type; + + ggml_dequantize_q4_0_kernel( + const block_q4_0* blocks, sycl_t* out, int64_t numel) + : blocks_(blocks), out_(out), numel_(numel) {} + + void operator()(sycl::id<1> index) const { + const int64_t i = index[0]; + if (i >= numel_) { + return; + } + + const int64_t block_index = i / QK4_0; + const int64_t block_offset = i % QK4_0; + const block_q4_0& block = blocks_[block_index]; + const bool is_high_half = block_offset >= (QK4_0 / 2); + const int64_t quant_index = + is_high_half ? (block_offset - QK4_0 / 2) : block_offset; + const uint8_t packed = block.qs[quant_index]; + const int quant = is_high_half ? (packed >> 4) : (packed & 0x0F); + const float value = + (static_cast(quant) - 8.0f) * static_cast(block.d); + out_[i] = static_cast(value); + } + + private: + const block_q4_0* blocks_; + sycl_t* out_; + int64_t numel_; +}; + +template +class ggml_dequantize_q5_0_kernel { + public: + using sycl_t = typename vllm::xpu::SyclTypeTrait::Type; + + ggml_dequantize_q5_0_kernel( + const block_q5_0* blocks, sycl_t* out, int64_t numel) + : blocks_(blocks), out_(out), numel_(numel) {} + + void operator()(sycl::id<1> index) const { + const int64_t i = index[0]; + if (i >= numel_) { + return; + } + + const int64_t block_index = i / QK5_0; + const int64_t block_offset = i % QK5_0; + const block_q5_0& block = blocks_[block_index]; + const bool is_high_half = block_offset >= (QK5_0 / 2); + const int64_t quant_index = + is_high_half ? (block_offset - QK5_0 / 2) : block_offset; + const uint8_t packed = block.qs[quant_index]; + const uint32_t qh = load_u32_le(block.qh); + const int xh = is_high_half ? ((qh >> (quant_index + 12)) & 0x10) + : (((qh >> quant_index) << 4) & 0x10); + const int base_quant = is_high_half ? (packed >> 4) : (packed & 0x0F); + const float value = (static_cast(base_quant | xh) - 16.0f) * + static_cast(block.d); + out_[i] = static_cast(value); + } + + private: + const block_q5_0* blocks_; + sycl_t* out_; + int64_t numel_; +}; + +template +class ggml_dequantize_q8_0_kernel { + public: + using sycl_t = typename vllm::xpu::SyclTypeTrait::Type; + + ggml_dequantize_q8_0_kernel( + const block_q8_0* blocks, sycl_t* out, int64_t numel) + : blocks_(blocks), out_(out), numel_(numel) {} + + void operator()(sycl::id<1> index) const { + const int64_t i = index[0]; + if (i >= numel_) { + return; + } + + const int64_t block_index = i / QK8_0; + const int64_t block_offset = i % QK8_0; + const block_q8_0& block = blocks_[block_index]; + const float value = static_cast(block.qs[block_offset]) * + static_cast(block.d); + out_[i] = static_cast(value); + } + + private: + const block_q8_0* blocks_; + sycl_t* out_; + int64_t numel_; +}; + +inline int64_t get_expected_nbytes(int64_t type, int64_t numel) { + switch (type) { + case GGML_TYPE_Q4_0: + return (numel / QK4_0) * static_cast(sizeof(block_q4_0)); + case GGML_TYPE_Q5_0: + return (numel / QK5_0) * static_cast(sizeof(block_q5_0)); + case GGML_TYPE_Q8_0: + return (numel / QK8_0) * static_cast(sizeof(block_q8_0)); + default: + return -1; + } +} + +inline const char* ggml_type_name(int64_t type) { + switch (type) { + case GGML_TYPE_Q4_0: + return "Q4_0"; + case GGML_TYPE_Q5_0: + return "Q5_0"; + case GGML_TYPE_Q8_0: + return "Q8_0"; + default: + return "unknown"; + } +} + +} // namespace ggml +} // namespace vllm \ No newline at end of file diff --git a/csrc/xpu/ops.h b/csrc/xpu/ops.h index d7eaad965..0408a8a34 100644 --- a/csrc/xpu/ops.h +++ b/csrc/xpu/ops.h @@ -44,6 +44,13 @@ torch::Tensor int4_gemm_w4a8( const std::optional& g_idx, const std::optional& bias); +torch::Tensor ggml_dequantize( + const torch::Tensor& W, + int64_t type, + int64_t m, + int64_t n, + std::optional out_dtype); + torch::Tensor cutlass_grouped_gemm_interface( torch::Tensor ptr_A, torch::Tensor ptr_B, diff --git a/csrc/xpu/torch_bindings.cpp b/csrc/xpu/torch_bindings.cpp index cdfa12844..a98c05757 100644 --- a/csrc/xpu/torch_bindings.cpp +++ b/csrc/xpu/torch_bindings.cpp @@ -30,6 +30,11 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, xpu_ops) { "bias) -> Tensor"); xpu_ops.impl("int4_gemm_w4a8", torch::kXPU, &int4_gemm_w4a8); + xpu_ops.def( + "ggml_dequantize(Tensor W, int type, SymInt m, SymInt n, ScalarType? " + "dtype) -> Tensor"); + xpu_ops.impl("ggml_dequantize", torch::kXPU, &ggml_dequantize); + xpu_ops.def( "cutlass_grouped_gemm_interface(Tensor ptr_A, Tensor ptr_B, Tensor? " "ptr_scales, " diff --git a/tests/register_ops.py b/tests/register_ops.py index 9af0d8cd7..67a14b558 100644 --- a/tests/register_ops.py +++ b/tests/register_ops.py @@ -270,6 +270,14 @@ def int4_gemm_w4a8(input: torch.Tensor, group_size, g_idx, bias) +def ggml_dequantize(weight: torch.Tensor, + quant_type: int, + m: int, + n: int, + dtype: Optional[torch.dtype] = None): + return torch.ops._xpu_C.ggml_dequantize(weight, quant_type, m, n, dtype) + + def fp8_gemm(input: torch.Tensor, weight: torch.Tensor, out_dtype: Optional[torch.dtype], scale_act: Optional[torch.Tensor], diff --git a/tests/test_gguf.py b/tests/test_gguf.py new file mode 100644 index 000000000..221fc44bc --- /dev/null +++ b/tests/test_gguf.py @@ -0,0 +1,214 @@ +# SPDX-License-Identifier: Apache-2.0 +from pathlib import Path + +import pytest +import torch +from gguf import GGMLQuantizationType, GGUFReader, ReaderTensor, dequantize +from huggingface_hub import snapshot_download + +from tests.register_ops import ggml_dequantize + +GGUF_SAMPLE = snapshot_download("Isotr0py/test-gguf-sample") +GGUF_SAMPLE_MOE = snapshot_download("SzymonOzog/test-gguf-moe-sample") + + +def get_gguf_sample_tensors( + hidden_size: int, + quant_type: GGMLQuantizationType) -> list[ReaderTensor]: + sample_dir = GGUF_SAMPLE + filename = f"Quant_{quant_type.name}_{hidden_size}.gguf" + sample_file = Path(sample_dir) / filename + return GGUFReader(sample_file).tensors + + +def get_gguf_MoE_tensors( + hidden_size: int, + quant_type: GGMLQuantizationType) -> list[ReaderTensor]: + sample_dir = GGUF_SAMPLE_MOE + filename = f"Quant_{quant_type.name}_{hidden_size}.gguf" + sample_file = Path(sample_dir) / filename + return GGUFReader(sample_file).tensors + + +DTYPES = [torch.bfloat16] # [torch.half, torch.bfloat16, torch.float32] +# Hidden_size for testing, must match the sample file in HF repo, +# we have `hidden_size = 256, 1024` for test in HF repo currently. +HIDDEN_SIZES = [256, 1024] +NUM_TOKENS = [7, 2050] # Arbitrary values for testing +SEEDS = [0] +QUANT_TYPES = [ + # i-matrix + # GGMLQuantizationType.IQ1_M, + # GGMLQuantizationType.IQ1_S, + # GGMLQuantizationType.IQ2_S, + # GGMLQuantizationType.IQ2_XS, + # GGMLQuantizationType.IQ3_S, + # GGMLQuantizationType.IQ3_XXS, + # GGMLQuantizationType.IQ4_NL, + # GGMLQuantizationType.IQ4_XS, + # # k-quants + # GGMLQuantizationType.Q2_K, + # GGMLQuantizationType.Q3_K, + # GGMLQuantizationType.Q4_K, + # GGMLQuantizationType.Q5_K, + # GGMLQuantizationType.Q6_K, + # standard quantization + GGMLQuantizationType.Q4_0, + GGMLQuantizationType.Q5_0, + GGMLQuantizationType.Q8_0, +] + + +@pytest.mark.parametrize("hidden_size", HIDDEN_SIZES) +@pytest.mark.parametrize("dtype", DTYPES) +@pytest.mark.parametrize("quant_type", QUANT_TYPES) +@torch.inference_mode() +def test_dequantize( + hidden_size: int, + dtype: torch.dtype, + quant_type: GGMLQuantizationType, +): + device = "xpu" + + tensors = get_gguf_sample_tensors(hidden_size, quant_type) + for tensor in tensors: + shape_str = tensor.name.split("_")[-1] + shape = map(int, shape_str.split("x")) + + ref_output = torch.tensor(dequantize(tensor.data, quant_type), + device=device).to(dtype) + output = ggml_dequantize( + torch.tensor(tensor.data, device=device), + quant_type, + *list(shape), + dtype, + ) + + torch.testing.assert_close(output, ref_output, atol=1e-2, rtol=4e-2) + + +""" +@pytest.mark.parametrize("hidden_size", HIDDEN_SIZES) +@pytest.mark.parametrize("dtype", DTYPES) +@pytest.mark.parametrize("quant_type", QUANT_TYPES) +@torch.inference_mode() +def test_mmvq(hidden_size: int, dtype: torch.dtype, +quant_type: GGMLQuantizationType): + set_random_seed(0) + + tensors = get_gguf_sample_tensors(hidden_size, quant_type) + x = torch.rand((1, hidden_size), dtype=dtype, device="cuda") + for tensor in tensors: + weight = torch.tensor(dequantize(tensor.data, quant_type), + device="cuda").to( + dtype + ) + ref_output = x @ weight.T + + qweight = torch.tensor(tensor.data, device="cuda") + output = ops.ggml_mul_mat_vec_a8(qweight, x, quant_type, + qweight.shape[0]).to( + dtype + ) + + torch.testing.assert_close(output, ref_output, atol=1, rtol=1e-1) + +@pytest.mark.parametrize("num_tokens", NUM_TOKENS) +@pytest.mark.parametrize("hidden_size", HIDDEN_SIZES) +@pytest.mark.parametrize("dtype", DTYPES) +@pytest.mark.parametrize( + "quant_type", + [ + # k-quants + GGMLQuantizationType.Q2_K, + GGMLQuantizationType.Q3_K, + GGMLQuantizationType.Q4_K, + GGMLQuantizationType.Q5_K, + GGMLQuantizationType.Q6_K, + # standard quants + GGMLQuantizationType.Q4_0, + GGMLQuantizationType.Q5_0, + GGMLQuantizationType.Q8_0, + ], +) +@torch.inference_mode() +def test_mmq( + num_tokens: int, + hidden_size: int, + dtype: torch.dtype, + quant_type: GGMLQuantizationType, +): + set_random_seed(0) + + tensors = get_gguf_sample_tensors(hidden_size, quant_type) + x = torch.rand((num_tokens, hidden_size), dtype=dtype, device="cuda") + for tensor in tensors: + weight = torch.tensor(dequantize(tensor.data, quant_type), + device="cuda").to( + dtype + ) + ref_output = x @ weight.T + + qweight = torch.tensor(tensor.data, device="cuda") + output = ops.ggml_mul_mat_a8(qweight, x, quant_type, qweight.shape[0]) + atols = {torch.half: 1, torch.bfloat16: 1.5, torch.float: 1.2} + # test matrix has inputs centered around 0 and lower precision from + # bfloat16 tends to accumulate and can greatly inflate rtol + # since outputs are also very close to 0 + rtols = {torch.half: 1e-1, torch.bfloat16: 1e4, torch.float: 2e1} + torch.testing.assert_close( + output, ref_output, atol=atols[dtype], rtol=rtols[dtype] + ) + +@pytest.mark.parametrize("num_tokens", NUM_TOKENS) +@pytest.mark.parametrize("hidden_size", [512]) +@pytest.mark.parametrize("top_k", [4, 8]) +@pytest.mark.parametrize("dtype", DTYPES) +@pytest.mark.parametrize("quant_type", QUANT_TYPES) +@torch.inference_mode() +def test_moe( + num_tokens: int, + hidden_size: int, + dtype: torch.dtype, + quant_type: GGMLQuantizationType, + top_k: int, +): + set_random_seed(0) + H, E = 1024, 256 + + x = torch.rand((num_tokens, H), dtype=dtype, device="cuda") + + topk_weights = torch.rand(num_tokens, top_k, device="cuda", dtype=dtype) + topk_ids = torch.randint( + 0, E, (num_tokens, top_k), device="cuda", dtype=torch.int32 + ) + + tensors = get_gguf_MoE_tensors(hidden_size, quant_type) + + w13 = tensors[0] + w2 = tensors[1] + + w13_dequant = torch.tensor(dequantize(w13.data, quant_type), + device="cuda").to( + dtype + ) + + w2_dequant = torch.tensor(dequantize(w2.data, quant_type), + device="cuda").to(dtype) + + output = _fused_moe_gguf( + x, + torch.tensor(w13.data, device="cuda"), + torch.tensor(w2.data, device="cuda"), + topk_weights, + topk_ids, + quant_type, + quant_type, + "silu", + ) + + ref_output = fused_experts( + x, w13_dequant, w2_dequant, topk_weights, topk_ids + ).reshape(output.shape) + torch.testing.assert_close(output, ref_output, atol=1, rtol=1e-1) +""" From 02da13dd4ef9b2eec3291cda997c2a8e6396f59b Mon Sep 17 00:00:00 2001 From: zhenwei-intel Date: Wed, 1 Apr 2026 02:17:34 +0000 Subject: [PATCH 3/3] update Signed-off-by: zhenwei-intel --- csrc/quantization/gguf/ggml_dequantize.cpp | 30 +++++++--------------- csrc/quantization/gguf/ggml_dequantize.hpp | 13 ++++++++++ 2 files changed, 22 insertions(+), 21 deletions(-) diff --git a/csrc/quantization/gguf/ggml_dequantize.cpp b/csrc/quantization/gguf/ggml_dequantize.cpp index cdc541990..7d36a4689 100644 --- a/csrc/quantization/gguf/ggml_dequantize.cpp +++ b/csrc/quantization/gguf/ggml_dequantize.cpp @@ -20,37 +20,25 @@ torch::Tensor ggml_dequantize( type == ggml::GGML_TYPE_Q4_0 || type == ggml::GGML_TYPE_Q5_0 || type == ggml::GGML_TYPE_Q8_0, "XPU ggml_dequantize currently only supports Q4_0 (type=2), " - "Q5_0 (type=6) and Q8_0 (type=8), got ", - type); + "Q5_0 (type=6) and Q8_0 (type=8), got ", type); TORCH_CHECK( W.scalar_type() == at::ScalarType::Byte, - "XPU ggml_dequantize expects uint8 weights, got ", - W.scalar_type()); + "XPU ggml_dequantize expects uint8 weights, got ", W.scalar_type()); TORCH_CHECK(m >= 0 && n >= 0, "m and n must be non-negative"); const int64_t numel = m * n; + const int64_t quant_block_size = ggml::get_quant_block_size(type); TORCH_CHECK( - numel % ggml::QK4_0 == 0, - ggml::ggml_type_name(type), - " dequantize expects m * n to be divisible by ", - ggml::QK4_0, - ", got ", - numel); + numel % quant_block_size == 0, ggml::ggml_type_name(type), + " dequantize expects m * n to be divisible by ", quant_block_size, + ", got ", numel); const int64_t expected_nbytes = ggml::get_expected_nbytes(type, numel); const int64_t weight_nbytes = W.numel() * W.element_size(); TORCH_CHECK( - weight_nbytes == expected_nbytes, - ggml::ggml_type_name(type), - " packed weight size mismatch: expected ", - expected_nbytes, - " bytes for shape (", - m, - ", ", - n, - "), got ", - weight_nbytes, - " bytes"); + weight_nbytes == expected_nbytes, ggml::ggml_type_name(type), + " packed weight size mismatch: expected ", expected_nbytes, + " bytes for shape (", m, ", ", n, "), got ", weight_nbytes, " bytes"); const auto dtype = out_dtype.value_or(torch::kFloat16); TORCH_CHECK( diff --git a/csrc/quantization/gguf/ggml_dequantize.hpp b/csrc/quantization/gguf/ggml_dequantize.hpp index 9525ab0b0..6408c8ac4 100644 --- a/csrc/quantization/gguf/ggml_dequantize.hpp +++ b/csrc/quantization/gguf/ggml_dequantize.hpp @@ -158,6 +158,19 @@ inline int64_t get_expected_nbytes(int64_t type, int64_t numel) { } } +inline int64_t get_quant_block_size(int64_t type) { + switch (type) { + case GGML_TYPE_Q4_0: + return QK4_0; + case GGML_TYPE_Q5_0: + return QK5_0; + case GGML_TYPE_Q8_0: + return QK8_0; + default: + return -1; + } +} + inline const char* ggml_type_name(int64_t type) { switch (type) { case GGML_TYPE_Q4_0: