@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 57b1ce94f7
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 14, 05:30 UTC

[CPU] Refactor CPU unquantized linear (#24150)

Signed-off-by: jiang1.li <jiang1.li@intel.com>

Li, Jiang committed Sep 4, 2025 at 14:28 UTC 57b1ce94f7ac13624a32da90658b36b60a276a60
9 files changed +466 -26
csrc/cpu/dnnl_helper.cpp
+177
@@ -22,6 +22,23 @@ void release_dnnl_matmul_handler(int64_t handler) {
22 delete ptr;
23 }
24
25 +DNNLScratchPadManager::DNNLScratchPadManager() : size_(0), ptr_(nullptr) {
26 + this->realloc(allocation_unit * 128);
27 +}
28 +
29 +void DNNLScratchPadManager::realloc(size_t new_size) {
30 + new_size = round(new_size);
31 + if (new_size > size_) {
32 + ptr_ = std::aligned_alloc(64, new_size);
33 + size_ = new_size;
34 + }
35 +}
36 +
37 +DNNLScratchPadManager* DNNLScratchPadManager::get_dnnl_scratchpad_manager() {
38 + static DNNLScratchPadManager manager;
39 + return &manager;
40 +}
41 +
42 template <typename KT, typename VT>
43 class DNNLPrimitiveCache {
44 public:
@@ -166,6 +183,23 @@ struct hash<W8A8MatMulPrimitiveHandler::MSizeCacheKey> {
183 hash<int>()(static_cast<int>(val.bias_type));
184 }
185 };
186 +
187 +template <>
188 +struct hash<MatMulPrimitiveHandler::ClassMatmulCacheKey> {
189 + size_t operator()(
190 + const MatMulPrimitiveHandler::ClassMatmulCacheKey& val) const {
191 + return hash<dnnl_dim_t>()(val.b_n_size) ^ hash<dnnl_dim_t>()(val.b_k_size);
192 + }
193 +};
194 +
195 +template <>
196 +struct hash<MatMulPrimitiveHandler::MSizeCacheKey> {
197 + size_t operator()(const MatMulPrimitiveHandler::MSizeCacheKey& val) const {
198 + return hash<dnnl_dim_t>()(val.a_m_size) ^
199 + hash<dnnl_dim_t>()(val.a_m_stride) ^ hash<bool>()(val.use_bias) ^
200 + hash<int>()(static_cast<int>(val.bias_type));
201 + }
202 +};
203 } // namespace std
204
205 bool operator==(const W8A8MatMulPrimitiveHandler::ClassMatmulCacheKey& l,
@@ -181,6 +215,17 @@ bool operator==(const W8A8MatMulPrimitiveHandler::MSizeCacheKey& l,
215 l.bias_type == r.bias_type;
216 }
217
218 +bool operator==(const MatMulPrimitiveHandler::ClassMatmulCacheKey& l,
219 + const MatMulPrimitiveHandler::ClassMatmulCacheKey& r) {
220 + return l.b_n_size == r.b_n_size && l.b_k_size == r.b_k_size;
221 +}
222 +
223 +bool operator==(const MatMulPrimitiveHandler::MSizeCacheKey& l,
224 + const MatMulPrimitiveHandler::MSizeCacheKey& r) {
225 + return l.a_m_size == r.a_m_size && l.a_m_stride == r.a_m_stride &&
226 + l.use_bias == r.use_bias && l.bias_type == r.bias_type;
227 +}
228 +
229 static std::shared_ptr<W8A8MatMulPrimitiveHandler::MSizeCache>
230 get_w8a8_class_primitive_cache(
231 const W8A8MatMulPrimitiveHandler::ClassMatmulCacheKey& key,
@@ -239,6 +284,11 @@ void W8A8MatMulPrimitiveHandler::execute(ExecArgs& args) {
284 }
285
286 dnnl::matmul matmul = get_matmul_cache(args);
287 +
288 + auto&& [scratchpad_storage, scratchpad_mem_desc] = get_runtime_memory_ptr(5);
289 + scratchpad_storage->set_data_handle(
290 + DNNLScratchPadManager::get_dnnl_scratchpad_manager()->get_data<void>());
291 +
292 matmul.execute(default_stream(), memory_cache_);
293 default_stream().wait();
294 }
@@ -257,6 +307,8 @@ dnnl::matmul W8A8MatMulPrimitiveHandler::get_matmul_cache(
307
308 return m_size_cache_->get_or_create(key, [&]() {
309 dnnl::matmul::primitive_desc desc = this->create_primitive_desc(key, false);
310 + auto manager = DNNLScratchPadManager::get_dnnl_scratchpad_manager();
311 + manager->realloc(desc.scratchpad_desc().get_size());
312 return dnnl::matmul(desc);
313 });
314 }
@@ -300,6 +352,11 @@ void W8A8MatMulPrimitiveHandler::init_runtime_memory_cache(const Args& args) {
352 dnnl::memory({{b_n_size_}, dnnl::memory::data_type::f32, {1}},
353 default_engine(), nullptr);
354 set_runtime_memory_ptr(4, memory_cache_[DNNL_ARG_BIAS].get());
355 +
356 + memory_cache_[DNNL_ARG_SCRATCHPAD] =
357 + dnnl::memory({{b_n_size_}, dnnl::memory::data_type::f32, {1}},
358 + default_engine(), nullptr);
359 + set_runtime_memory_ptr(5, memory_cache_[DNNL_ARG_SCRATCHPAD].get());
360 }
361
362 dnnl::matmul::primitive_desc W8A8MatMulPrimitiveHandler::create_primitive_desc(
@@ -319,6 +376,9 @@ dnnl::matmul::primitive_desc W8A8MatMulPrimitiveHandler::create_primitive_desc(
376 dnnl::memory::format_tag::ab);
377
378 dnnl::primitive_attr attr;
379 +
380 + attr.set_scratchpad_mode(dnnl::scratchpad_mode::user);
381 +
382 // For PER_TOKEN, scales will be applied in outside epilogue
383 if (a_qs_ == QuantizationStrategy::PER_TENSOR) {
384 attr.set_scales_mask(DNNL_ARG_SRC, 0);
@@ -344,3 +404,120 @@ dnnl::matmul::primitive_desc W8A8MatMulPrimitiveHandler::create_primitive_desc(
404 attr);
405 }
406 }
407 +
408 +MatMulPrimitiveHandler::MatMulPrimitiveHandler(const Args& args)
409 + : DNNLMatMulPrimitiveHandler(
410 + static_cast<DNNLMatMulPrimitiveHandler::Args>(args), args.ab_type),
411 + m_size_cache_(nullptr) {
412 + assert(ab_type_ == dnnl::memory::data_type::f32 ||
413 + ab_type_ == dnnl::memory::data_type::bf16 ||
414 + ab_type_ == dnnl::memory::data_type::f16);
415 + prepack_weight(args.b_ptr,
416 + create_primitive_desc(
417 + MSizeCacheKey{.a_m_size = DNNL_RUNTIME_DIM_VAL,
418 + .a_m_stride = DNNL_RUNTIME_DIM_VAL,
419 + .use_bias = false,
420 + .bias_type = dnnl::memory::data_type::undef},
421 + true)
422 + .weights_desc());
423 + init_runtime_memory_cache(args);
424 +}
425 +
426 +static std::shared_ptr<MatMulPrimitiveHandler::MSizeCache>
427 +get_matul_class_primitive_cache(
428 + const MatMulPrimitiveHandler::ClassMatmulCacheKey& key,
429 + int64_t cache_size) {
430 + static MatMulPrimitiveHandler::ClassMatmulCache cache(128);
431 + assert(cache_size > 0);
432 + return cache.get_or_create(key, [&]() {
433 + return std::make_shared<MatMulPrimitiveHandler::MSizeCache>(cache_size);
434 + });
435 +}
436 +
437 +void MatMulPrimitiveHandler::execute(ExecArgs& args) {
438 + auto&& [a_storage, a_mem_desc] = get_runtime_memory_ptr(0);
439 + auto&& [c_storage, c_mem_desc] = get_runtime_memory_ptr(1);
440 + a_storage->set_data_handle((void*)args.a_ptr);
441 + a_mem_desc->dims[0] = args.a_m_size;
442 + a_mem_desc->format_desc.blocking.strides[0] = args.a_m_stride;
443 + c_storage->set_data_handle((void*)args.c_ptr);
444 + c_mem_desc->dims[0] = args.a_m_size;
445 +
446 + if (args.use_bias) {
447 + auto&& [bias_storage, bias_mem_desc] = get_runtime_memory_ptr(2);
448 + bias_storage->set_data_handle((void*)args.bias_ptr);
449 + }
450 +
451 + dnnl::matmul matmul = get_matmul_cache(args);
452 +
453 + auto&& [scratchpad_storage, scratchpad_mem_desc] = get_runtime_memory_ptr(3);
454 + scratchpad_storage->set_data_handle(
455 + DNNLScratchPadManager::get_dnnl_scratchpad_manager()->get_data<void>());
456 +
457 + matmul.execute(default_stream(), memory_cache_);
458 + default_stream().wait();
459 +}
460 +
461 +dnnl::matmul MatMulPrimitiveHandler::get_matmul_cache(
462 + const MSizeCacheKey& key) {
463 + if (m_size_cache_.get() == nullptr) {
464 + ClassMatmulCacheKey key = {.b_n_size = b_n_size_, .b_k_size = b_k_size_};
465 + m_size_cache_ = get_matul_class_primitive_cache(key, primitive_cache_size_);
466 + }
467 + return m_size_cache_->get_or_create(key, [&]() {
468 + dnnl::matmul::primitive_desc desc = this->create_primitive_desc(key, false);
469 + auto manager = DNNLScratchPadManager::get_dnnl_scratchpad_manager();
470 + manager->realloc(desc.scratchpad_desc().get_size());
471 + return dnnl::matmul(desc);
472 + });
473 +}
474 +
475 +dnnl::matmul::primitive_desc MatMulPrimitiveHandler::create_primitive_desc(
476 + const MSizeCacheKey& key, bool first_time) {
477 + dnnl::memory::desc a_md;
478 + dnnl::memory::desc b_md;
479 + if (first_time) {
480 + a_md = dnnl::memory::desc({key.a_m_size, b_k_size_}, b_type_,
481 + dnnl::memory::format_tag::ab);
482 + b_md = dnnl::memory::desc({b_k_size_, b_n_size_}, b_type_,
483 + dnnl::memory::format_tag::any);
484 + } else {
485 + a_md = dnnl::memory::desc({key.a_m_size, b_k_size_}, b_type_,
486 + {key.a_m_stride, 1});
487 + b_md = b_target_mem_desc_;
488 + }
489 + dnnl::memory::desc c_md({key.a_m_size, b_n_size_}, c_type_,
490 + dnnl::memory::format_tag::ab);
491 +
492 + dnnl::primitive_attr attr;
493 + attr.set_scratchpad_mode(dnnl::scratchpad_mode::user);
494 +
495 + if (key.use_bias) {
496 + dnnl::memory::desc bias_md({1, b_n_size_}, key.bias_type, {b_n_size_, 1});
497 + return dnnl::matmul::primitive_desc(default_engine(), a_md, b_md, bias_md,
498 + c_md, attr);
499 + } else {
500 + return dnnl::matmul::primitive_desc(default_engine(), a_md, b_md, c_md,
501 + attr);
502 + }
503 +}
504 +
505 +void MatMulPrimitiveHandler::init_runtime_memory_cache(const Args& args) {
506 + memory_cache_[DNNL_ARG_SRC] = dnnl::memory(
507 + {{1, b_k_size_}, b_type_, {b_k_size_, 1}}, default_engine(), nullptr);
508 + set_runtime_memory_ptr(0, memory_cache_[DNNL_ARG_SRC].get());
509 + memory_cache_[DNNL_ARG_DST] =
510 + dnnl::memory({{1, b_n_size_}, c_type_, dnnl::memory::format_tag::ab},
511 + default_engine(), nullptr);
512 + set_runtime_memory_ptr(1, memory_cache_[DNNL_ARG_DST].get());
513 +
514 + memory_cache_[DNNL_ARG_BIAS] =
515 + dnnl::memory({{b_n_size_}, dnnl::memory::data_type::f32, {1}},
516 + default_engine(), nullptr);
517 + set_runtime_memory_ptr(2, memory_cache_[DNNL_ARG_BIAS].get());
518 +
519 + memory_cache_[DNNL_ARG_SCRATCHPAD] =
520 + dnnl::memory({{b_n_size_}, dnnl::memory::data_type::f32, {1}},
521 + default_engine(), nullptr);
522 + set_runtime_memory_ptr(3, memory_cache_[DNNL_ARG_SCRATCHPAD].get());
523 +}
csrc/cpu/dnnl_helper.h
+74
@@ -59,6 +59,30 @@ constexpr inline dnnl::memory::data_type get_dnnl_type() {
59 return DNNLType<std::decay_t<T>>::type;
60 }
61
62 +class DNNLScratchPadManager {
63 + public:
64 + static constexpr size_t allocation_unit = 4 * 1024 * 1024; // 4KB
65 +
66 + static DNNLScratchPadManager* get_dnnl_scratchpad_manager();
67 +
68 + DNNLScratchPadManager();
69 +
70 + template <typename T>
71 + T* get_data() {
72 + return reinterpret_cast<T*>(ptr_);
73 + }
74 +
75 + static size_t round(size_t size) {
76 + return ((size + allocation_unit - 1) / allocation_unit) * allocation_unit;
77 + }
78 +
79 + void realloc(size_t new_size);
80 +
81 + private:
82 + size_t size_;
83 + void* ptr_;
84 +};
85 +
86 class DNNLMatMulPrimitiveHandler {
87 public:
88 virtual ~DNNLMatMulPrimitiveHandler() = default;
@@ -166,4 +190,54 @@ class W8A8MatMulPrimitiveHandler : public DNNLMatMulPrimitiveHandler {
190 std::shared_ptr<MSizeCache> m_size_cache_;
191 };
192
193 +class MatMulPrimitiveHandler : public DNNLMatMulPrimitiveHandler {
194 + public:
195 + struct Args : public DNNLMatMulPrimitiveHandler::Args {
196 + dnnl::memory::data_type ab_type;
197 + };
198 +
199 + struct ClassMatmulCacheKey {
200 + dnnl_dim_t b_n_size;
201 + dnnl_dim_t b_k_size;
202 +
203 + friend bool operator==(const ClassMatmulCacheKey& l,
204 + const ClassMatmulCacheKey& r);
205 + };
206 +
207 + struct MSizeCacheKey {
208 + dnnl_dim_t a_m_size;
209 + dnnl_dim_t a_m_stride;
210 + bool use_bias;
211 + dnnl::memory::data_type bias_type;
212 +
213 + friend bool operator==(const MSizeCacheKey& l, const MSizeCacheKey& r);
214 + };
215 +
216 + using MSizeCache = DNNLPrimitiveCache<MSizeCacheKey, dnnl::matmul>;
217 + using ClassMatmulCache =
218 + DNNLPrimitiveCache<ClassMatmulCacheKey, std::shared_ptr<MSizeCache>>;
219 +
220 + struct ExecArgs : public MSizeCacheKey {
221 + const void* a_ptr;
222 + const void* bias_ptr;
223 + void* c_ptr;
224 + };
225 +
226 + public:
227 + MatMulPrimitiveHandler(const Args& args);
228 +
229 + void execute(ExecArgs& args);
230 +
231 + private:
232 + dnnl::matmul::primitive_desc create_primitive_desc(const MSizeCacheKey& key,
233 + bool first_time);
234 +
235 + void init_runtime_memory_cache(const Args& args);
236 +
237 + dnnl::matmul get_matmul_cache(const MSizeCacheKey& key);
238 +
239 + private:
240 + std::shared_ptr<MSizeCache> m_size_cache_;
241 +};
242 +
243 #endif
csrc/cpu/dnnl_kernels.cpp
+54
@@ -379,6 +379,7 @@ void onednn_scaled_mm(
379 exec_args.a_ptr = a.data_ptr<int8_t>();
380 exec_args.a_m_size = a.size(0);
381 exec_args.bias_ptr = nullptr;
382 + exec_args.bias_type = get_dnnl_type<void>();
383 exec_args.use_bias = false;
384 exec_args.a_scales_ptr = nullptr;
385 exec_args.a_zero_points_ptr = nullptr;
@@ -492,3 +493,56 @@ void dynamic_scaled_int8_quant(
493 }
494 });
495 }
496 +
497 +int64_t create_onednn_mm_handler(const torch::Tensor& b,
498 + int64_t primitive_cache_size) {
499 + TORCH_CHECK(b.dim() == 2);
500 +
501 + MatMulPrimitiveHandler::Args args;
502 + args.primitive_cache_size = primitive_cache_size;
503 +
504 + args.b_k_size = b.size(0);
505 + args.b_k_stride = b.stride(0);
506 + args.b_n_size = b.size(1);
507 + args.b_n_stride = b.stride(1);
508 + args.b_ptr = b.data_ptr();
509 +
510 + VLLM_DISPATCH_FLOATING_TYPES(b.scalar_type(), "create_onednn_mm_handler",
511 + [&] {
512 + args.c_type = get_dnnl_type<scalar_t>();
513 + args.ab_type = get_dnnl_type<scalar_t>();
514 + });
515 +
516 + return reinterpret_cast<int64_t>(new MatMulPrimitiveHandler(args));
517 +}
518 +
519 +void onednn_mm(torch::Tensor& c, // [M, OC], row-major
520 + const torch::Tensor& a, // [M, IC], row-major
521 + const std::optional<torch::Tensor>& bias, int64_t handler) {
522 + CPU_KERNEL_GUARD_IN(onednn_mm)
523 + TORCH_CHECK(a.dim() == 2);
524 + TORCH_CHECK(a.stride(-1) == 1);
525 + TORCH_CHECK(c.is_contiguous());
526 + MatMulPrimitiveHandler* ptr =
527 + reinterpret_cast<MatMulPrimitiveHandler*>(handler);
528 +
529 + MatMulPrimitiveHandler::ExecArgs exec_args;
530 + exec_args.a_m_size = a.size(0);
531 + exec_args.a_m_stride = a.stride(0);
532 +
533 + VLLM_DISPATCH_FLOATING_TYPES(a.scalar_type(), "onednn_mm", [&] {
534 + if (bias.has_value()) {
535 + exec_args.use_bias = true;
536 + exec_args.bias_type = get_dnnl_type<scalar_t>();
537 + exec_args.bias_ptr = bias->data_ptr<scalar_t>();
538 + } else {
539 + exec_args.use_bias = false;
540 + exec_args.bias_type = get_dnnl_type<void>();
541 + exec_args.bias_ptr = nullptr;
542 + }
543 + exec_args.a_ptr = a.data_ptr<scalar_t>();
544 + exec_args.c_ptr = c.data_ptr<scalar_t>();
545 +
546 + ptr->execute(exec_args);
547 + });
548 +}
csrc/cpu/torch_bindings.cpp
+18
@@ -21,6 +21,12 @@ void onednn_scaled_mm(torch::Tensor& c, const torch::Tensor& a,
21 const std::optional<torch::Tensor>& bias,
22 int64_t handler);
23
24 +int64_t create_onednn_mm_handler(const torch::Tensor& b,
25 + int64_t primitive_cache_size);
26 +
27 +void onednn_mm(torch::Tensor& c, const torch::Tensor& a,
28 + const std::optional<torch::Tensor>& bias, int64_t handler);
29 +
30 void mla_decode_kvcache(torch::Tensor& out, torch::Tensor& query,
31 torch::Tensor& kv_cache, double scale,
32 torch::Tensor& block_tables, torch::Tensor& seq_lens);
@@ -153,6 +159,18 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
159 ops.def("release_dnnl_matmul_handler(int handler) -> ()",
160 &release_dnnl_matmul_handler);
161
162 + // Create oneDNN GEMM handler
163 + ops.def(
164 + "create_onednn_mm_handler(Tensor b, int "
165 + "primitive_cache_size) -> int",
166 + &create_onednn_mm_handler);
167 +
168 + // oneDNN GEMM
169 + ops.def(
170 + "onednn_mm(Tensor! c, Tensor a, Tensor? bias, "
171 + "int handler) -> ()");
172 + ops.impl("onednn_mm", torch::kCPU, &onednn_mm);
173 +
174 // Create oneDNN W8A8 handler
175 ops.def(
176 "create_onednn_scaled_mm_handler(Tensor b, Tensor b_scales, ScalarType "
tests/kernels/test_onednn.py
+70
@@ -111,6 +111,49 @@ def onednn_int8_gemm_test_helper(primitive_cache_size: int,
111 torch.testing.assert_close(out, baseline, rtol=1e-1, atol=1e0)
112
113
114 +def onednn_gemm_test_helper(primitive_cache_size: int,
115 + m: int,
116 + n: int,
117 + k: int,
118 + use_bias: bool,
119 + use_stride: bool,
120 + dtype: torch.dtype = torch.bfloat16,
121 + device: str = "cpu"):
122 + if use_stride:
123 + a = torch.rand((m, 2 * k), dtype=dtype, device=device) * 1.5
124 + a = a[:, :k]
125 + else:
126 + a = torch.rand((m, k), dtype=dtype, device=device) * 1.5
127 +
128 + b = torch.rand((n, k), dtype=dtype, device=device) * 1.5
129 +
130 + if use_bias:
131 + bias = torch.rand((n, ), device=device, dtype=dtype) * 5
132 + bias_f32 = bias.float()
133 + else:
134 + bias = None
135 + bias_f32 = None
136 +
137 + handler = ops.create_onednn_mm(
138 + b.t(),
139 + primitive_cache_size,
140 + )
141 +
142 + out = ops.onednn_mm(handler, a, bias)
143 + baseline = torch.nn.functional.linear(a.float(), b.float(),
144 + bias_f32).to(dtype=a.dtype)
145 +
146 + torch.testing.assert_close(out, baseline)
147 +
148 + if use_bias:
149 + # To test runtime bias setting
150 + out = ops.onednn_mm(handler, a, None)
151 + baseline = torch.nn.functional.linear(a.float(), b.float(),
152 + None).to(dtype=a.dtype)
153 +
154 + torch.testing.assert_close(out, baseline)
155 +
156 +
157 @pytest.mark.parametrize("n,k", NK_FACTORS)
158 @pytest.mark.parametrize("m_list", M_FACTORS)
159 @pytest.mark.parametrize("per_tensor_a_scale", [True, False])
@@ -142,3 +185,30 @@ def test_onednn_int8_scaled_gemm(
185 use_azp=use_azp,
186 out_dtype=output_type,
187 )
188 +
189 +
190 +@pytest.mark.parametrize("n,k", NK_FACTORS)
191 +@pytest.mark.parametrize("m_list", M_FACTORS)
192 +@pytest.mark.parametrize("use_bias", [True, False])
193 +@pytest.mark.parametrize("use_stride", [True, False])
194 +@pytest.mark.parametrize("dtype", DTYPE)
195 +@pytest.mark.parametrize("primitive_cache_size", CACHE_SIZES)
196 +def test_onednn_gemm(
197 + n: int,
198 + k: int,
199 + m_list: tuple[int],
200 + use_bias: bool,
201 + use_stride: bool,
202 + dtype: torch.dtype,
203 + primitive_cache_size: int,
204 +):
205 + for m in m_list:
206 + onednn_gemm_test_helper(
207 + primitive_cache_size=primitive_cache_size,
208 + m=m,
209 + n=n,
210 + k=k,
211 + use_bias=use_bias,
212 + use_stride=use_stride,
213 + dtype=dtype,
214 + )
vllm/_custom_ops.py
+29
@@ -1928,6 +1928,35 @@ class CPUDNNLGEMMHandler:
1928 torch.ops._C.release_dnnl_matmul_handler(self.handler)
1929
1930
1931 +if hasattr(torch.ops._C, "create_onednn_mm_handler"):
1932 + _supports_onednn = True
1933 +else:
1934 + _supports_onednn = False
1935 +
1936 +
1937 +def create_onednn_mm(
1938 + weight: torch.Tensor, # [K, N]
1939 + primitive_cache_size: int = 128,
1940 +) -> CPUDNNLGEMMHandler:
1941 + handler = CPUDNNLGEMMHandler()
1942 + handler.k, handler.n = weight.size()
1943 + handler.handler = torch.ops._C.create_onednn_mm_handler(
1944 + weight, primitive_cache_size)
1945 + return handler
1946 +
1947 +
1948 +def onednn_mm(
1949 + dnnl_handler: CPUDNNLGEMMHandler,
1950 + x: torch.Tensor,
1951 + bias: Optional[torch.Tensor],
1952 +) -> torch.Tensor:
1953 + output = torch.empty((*x.shape[0:-1], dnnl_handler.n), dtype=x.dtype)
1954 + torch.ops._C.onednn_mm(output, x.reshape(-1, dnnl_handler.k), bias,
1955 + dnnl_handler.handler)
1956 +
1957 + return output
1958 +
1959 +
1960 def create_onednn_scaled_mm(
1961 weight: torch.Tensor, # [K, N]
1962 weight_scales: torch.Tensor,
vllm/model_executor/layers/linear.py
+4 -21
@@ -9,7 +9,6 @@ import torch
9 import torch.nn as nn
10 from torch.nn.parameter import Parameter, UninitializedParameter
11
12 -from vllm import envs
12 from vllm.distributed import (divide, get_tensor_model_parallel_rank,
13 get_tensor_model_parallel_world_size,
14 split_tensor_along_last_dim,
@@ -200,26 +199,10 @@ class UnquantizedLinearMethod(LinearMethodBase):
199 set_weight_attrs(weight, extra_weight_attrs)
200
201 def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
203 - # special postprocessing for CPU SGL
204 - if current_platform.is_cpu() and envs.VLLM_CPU_SGL_KERNEL:
205 - from vllm.model_executor.layers.utils import check_cpu_sgl_kernel
206 - N, K = layer.weight.size()
207 - dtype = layer.weight.dtype
208 - if check_cpu_sgl_kernel(N, K, dtype):
209 - packed_weight = torch.ops._C.convert_weight_packed(
210 - layer.weight)
211 - assert packed_weight.size() == layer.weight.size()
212 - layer.weight.copy_(packed_weight)
213 - if layer.bias is not None:
214 - layer.bias = Parameter(layer.bias.to(torch.float32),
215 - requires_grad=False)
216 - layer.use_cpu_sgl = True
217 - else:
218 - logger.warning(
219 - "CPU SGL kernels require Intel AMX support,"
220 - " bf16/fp16/int8 weight, IC and OC are divisible by "
221 - "32 and 16.")
222 - layer.use_cpu_sgl = False
202 + if current_platform.is_cpu():
203 + from vllm.model_executor.layers.utils import (
204 + dispatch_cpu_unquantized_gemm)
205 + dispatch_cpu_unquantized_gemm(layer, remove_weight=True)
206
207 def apply(self,
208 layer: torch.nn.Module,
vllm/model_executor/layers/utils.py
+34 -5
@@ -142,20 +142,49 @@ direct_register_custom_op(
142 )
143
144
145 -def check_cpu_sgl_kernel(n: int, k: int, dtype: torch.dtype):
145 +def check_cpu_sgl_kernel(n: int, k: int, dtype: torch.dtype) -> bool:
146 return (torch._C._cpu._is_amx_tile_supported()
147 and (dtype in (torch.bfloat16, torch.int8)) and k % 32 == 0
148 and n % 16 == 0)
149
150
151 +def dispatch_cpu_unquantized_gemm(
152 + layer: torch.nn.Module,
153 + remove_weight: bool,
154 +) -> None:
155 + N, K = layer.weight.size()
156 + dtype = layer.weight.dtype
157 + if envs.VLLM_CPU_SGL_KERNEL and check_cpu_sgl_kernel(N, K, dtype):
158 + packed_weight = torch.ops._C.convert_weight_packed(layer.weight)
159 + if getattr(layer, "bias", None) is not None:
160 + bias_f32 = layer.bias.to(torch.float32)
161 + else:
162 + bias_f32 = None
163 + layer.cpu_linear = (
164 + lambda x, weight, bias: torch.ops._C.weight_packed_linear(
165 + x, packed_weight, bias_f32
166 + if bias is not None else None, True))
167 + if remove_weight:
168 + layer.weight = torch.nn.Parameter(torch.empty(0),
169 + requires_grad=False)
170 + elif ops._supports_onednn:
171 + origin_weight = layer.weight
172 + if remove_weight:
173 + layer.weight = torch.nn.Parameter(torch.empty(0),
174 + requires_grad=False)
175 + handler = ops.create_onednn_mm(origin_weight.t(), 32)
176 + layer.cpu_linear = lambda x, weight, bias: ops.onednn_mm(
177 + handler, x, bias)
178 + else:
179 + layer.cpu_linear = lambda x, weight, bias: torch.nn.functional.linear(
180 + x, weight, bias)
181 +
182 +
183 def cpu_unquantized_gemm(layer: torch.nn.Module,
184 x: torch.Tensor,
185 weight: torch.Tensor,
186 bias: Optional[torch.Tensor] = None):
155 - if getattr(layer, "use_cpu_sgl", False):
156 - return torch.ops._C.weight_packed_linear(x, weight, bias, True)
157 - else:
158 - return torch.nn.functional.linear(x, weight, bias)
187 + return layer.cpu_linear(x, weight, bias)
188
189
190 def dispatch_unquantized_gemm() -> Callable[..., torch.Tensor]:
vllm/model_executor/layers/vocab_parallel_embedding.py
+6
@@ -40,6 +40,12 @@ class UnquantizedEmbeddingMethod(QuantizeMethodBase):
40 layer.register_parameter("weight", weight)
41 set_weight_attrs(weight, extra_weight_attrs)
42
43 + def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
44 + if current_platform.is_cpu():
45 + from vllm.model_executor.layers.utils import (
46 + dispatch_cpu_unquantized_gemm)
47 + dispatch_cpu_unquantized_gemm(layer, remove_weight=False)
48 +
49 def apply(self,
50 layer: torch.nn.Module,
51 x: torch.Tensor,