Ggml dequant

.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
 

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"
@@ -478,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"

csrc/flash_attn/flash_api.cpp

@@ -110,8 +110,6 @@ std::vector<at::Tensor> 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,15 +141,16 @@ std::vector<at::Tensor> mha_varlen_fwd(
   at::Tensor out;
   if (out_.has_value()) {
     out = *out_;
-  } else {
-    out = torch::empty_like(q);
   }
 
   bool is_varlen = true;
   bool is_local = (window_size_left != -1) | (window_size_right != -1);
   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,18 +188,25 @@ std::vector<at::Tensor> mha_varlen_fwd(
     int num_tokens = q.size(0);
     int batch_size = static_cast<int>(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));
 
     at::Tensor tmp_out =
         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},

csrc/quantization/gguf/ggml_dequantize.cpp

@@ -0,0 +1,102 @@
+#include "ggml_dequantize.hpp"
+#include "utils.h"
+#include "xpu/ops.h"
+
+#include <ATen/ATen.h>
+#include <ATen/DeviceGuard.h>
+
+namespace ggml = vllm::ggml;
+
+torch::Tensor ggml_dequantize(
+    const torch::Tensor& W,
+    int64_t type,
+    int64_t m,
+    int64_t n,
+    std::optional<c10::ScalarType> 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;
+  const int64_t quant_block_size = ggml::get_quant_block_size(type);
+  TORCH_CHECK(
+      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");
+
+  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<uint8_t>();
+
+  VLLM_DISPATCH_FLOATING_TYPES(output.scalar_type(), "ggml_dequantize", [&] {
+    using sycl_t = typename vllm::xpu::SyclTypeTrait<scalar_t>::Type;
+    auto* out_ptr = reinterpret_cast<sycl_t*>(output.data_ptr<scalar_t>());
+
+    switch (type) {
+      case ggml::GGML_TYPE_Q4_0: {
+        auto* blocks = reinterpret_cast<const ggml::block_q4_0*>(weight_ptr);
+        queue.submit([&](sycl::handler& cgh) {
+          cgh.parallel_for(
+              sycl::range<1>(static_cast<size_t>(numel)),
+              ggml::ggml_dequantize_q4_0_kernel<scalar_t>(
+                  blocks, out_ptr, numel));
+        });
+        break;
+      }
+      case ggml::GGML_TYPE_Q5_0: {
+        auto* blocks = reinterpret_cast<const ggml::block_q5_0*>(weight_ptr);
+        queue.submit([&](sycl::handler& cgh) {
+          cgh.parallel_for(
+              sycl::range<1>(static_cast<size_t>(numel)),
+              ggml::ggml_dequantize_q5_0_kernel<scalar_t>(
+                  blocks, out_ptr, numel));
+        });
+        break;
+      }
+      case ggml::GGML_TYPE_Q8_0: {
+        auto* blocks = reinterpret_cast<const ggml::block_q8_0*>(weight_ptr);
+        queue.submit([&](sycl::handler& cgh) {
+          cgh.parallel_for(
+              sycl::range<1>(static_cast<size_t>(numel)),
+              ggml::ggml_dequantize_q8_0_kernel<scalar_t>(
+                  blocks, out_ptr, numel));
+        });
+        break;
+      }
+      default:
+        TORCH_CHECK(
+            false, "Unsupported GGML type for XPU ggml_dequantize: ", type);
+    }
+  });
+
+  return output;
+}

csrc/quantization/gguf/ggml_dequantize.hpp

@@ -0,0 +1,188 @@
+#pragma once
+
+#include "dispatch_utils.h"
+#include "utils.h"
+
+#include <cstdint>
+#include <sycl/sycl.hpp>
+
+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<uint32_t>(bytes[0]) |
+         (static_cast<uint32_t>(bytes[1]) << 8) |
+         (static_cast<uint32_t>(bytes[2]) << 16) |
+         (static_cast<uint32_t>(bytes[3]) << 24);
+}
+
+template <typename scalar_t>
+class ggml_dequantize_q4_0_kernel {
+ public:
+  using sycl_t = typename vllm::xpu::SyclTypeTrait<scalar_t>::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<float>(quant) - 8.0f) * static_cast<float>(block.d);
+    out_[i] = static_cast<sycl_t>(value);
+  }
+
+ private:
+  const block_q4_0* blocks_;
+  sycl_t* out_;
+  int64_t numel_;
+};
+
+template <typename scalar_t>
+class ggml_dequantize_q5_0_kernel {
+ public:
+  using sycl_t = typename vllm::xpu::SyclTypeTrait<scalar_t>::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<float>(base_quant | xh) - 16.0f) *
+                        static_cast<float>(block.d);
+    out_[i] = static_cast<sycl_t>(value);
+  }
+
+ private:
+  const block_q5_0* blocks_;
+  sycl_t* out_;
+  int64_t numel_;
+};
+
+template <typename scalar_t>
+class ggml_dequantize_q8_0_kernel {
+ public:
+  using sycl_t = typename vllm::xpu::SyclTypeTrait<scalar_t>::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<float>(block.qs[block_offset]) *
+                        static_cast<float>(block.d);
+    out_[i] = static_cast<sycl_t>(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<int64_t>(sizeof(block_q4_0));
+    case GGML_TYPE_Q5_0:
+      return (numel / QK5_0) * static_cast<int64_t>(sizeof(block_q5_0));
+    case GGML_TYPE_Q8_0:
+      return (numel / QK8_0) * static_cast<int64_t>(sizeof(block_q8_0));
+    default:
+      return -1;
+  }
+}
+
+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:
+      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

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 =