@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 7be5d113d8
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 W8A8 scaled_mm (#23071)

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

Li, Jiang committed Aug 21, 2025 at 09:34 UTC 7be5d113d8784536b79f27f24cfa91958dc291b0
17 files changed +1527 -1275
.buildkite/scripts/hardware_ci/run-cpu-test.sh
+6 -1
@@ -46,6 +46,11 @@ function cpu_tests() {
46 set -e
47 python3 examples/offline_inference/basic/generate.py --model facebook/opt-125m"
48
49 + # Run kernel tests
50 + docker exec cpu-test-"$NUMA_NODE" bash -c "
51 + set -e
52 + pytest -v -s tests/kernels/test_onednn.py"
53 +
54 # Run basic model test
55 docker exec cpu-test-"$NUMA_NODE" bash -c "
56 set -e
@@ -99,4 +104,4 @@ function cpu_tests() {
104
105 # All of CPU tests are expected to be finished less than 40 mins.
106 export -f cpu_tests
102 -timeout 1.5h bash -c "cpu_tests $CORE_RANGE $NUMA_NODE"
107 +timeout 2h bash -c "cpu_tests $CORE_RANGE $NUMA_NODE"
cmake/cpu_extension.cmake
+20 -39
@@ -182,17 +182,17 @@ endif()
182 #
183 # Build oneDNN for W8A8 GEMM kernels (only for x86-AVX512 /ARM platforms)
184 # Flag to enable ACL kernels for AARCH64 platforms
185 -if ( VLLM_BUILD_ACL STREQUAL "ON")
185 +if (VLLM_BUILD_ACL STREQUAL "ON")
186 set(USE_ACL ON)
187 else()
188 set(USE_ACL OFF)
189 endif()
190
191 -if ((AVX512_FOUND AND NOT AVX512_DISABLED) OR ASIMD_FOUND)
191 +if ((AVX512_FOUND AND NOT AVX512_DISABLED) OR ASIMD_FOUND OR POWER9_FOUND OR POWER10_FOUND OR POWER11_FOUND)
192 FetchContent_Declare(
193 oneDNN
194 GIT_REPOSITORY https://github.com/oneapi-src/oneDNN.git
195 - GIT_TAG v3.8.1
195 + GIT_TAG v3.9
196 GIT_PROGRESS TRUE
197 GIT_SHALLOW TRUE
198 )
@@ -204,7 +204,7 @@ if ((AVX512_FOUND AND NOT AVX512_DISABLED) OR ASIMD_FOUND)
204 endif()
205 set(ONEDNN_AARCH64_USE_ACL "ON")
206 set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wl,-rpath,$ENV{ACL_ROOT_DIR}/build/")
207 - endif()
207 + endif()
208
209 set(ONEDNN_LIBRARY_TYPE "STATIC")
210 set(ONEDNN_BUILD_DOC "OFF")
@@ -217,38 +217,23 @@ if ((AVX512_FOUND AND NOT AVX512_DISABLED) OR ASIMD_FOUND)
217 set(ONEDNN_ENABLE_ITT_TASKS "OFF")
218 set(ONEDNN_ENABLE_MAX_CPU_ISA "OFF")
219 set(ONEDNN_ENABLE_CPU_ISA_HINTS "OFF")
220 + set(ONEDNN_VERBOSE "OFF")
221 set(CMAKE_POLICY_DEFAULT_CMP0077 NEW)
222
223 FetchContent_MakeAvailable(oneDNN)
223 -
224 - list(APPEND LIBS dnnl)
225 -elseif(POWER10_FOUND)
226 - FetchContent_Declare(
227 - oneDNN
228 - GIT_REPOSITORY https://github.com/oneapi-src/oneDNN.git
229 - GIT_TAG v3.7.2
230 - GIT_PROGRESS TRUE
231 - GIT_SHALLOW TRUE
224 + add_library(dnnl_ext OBJECT "csrc/cpu/dnnl_helper.cpp")
225 + target_include_directories(
226 + dnnl_ext
227 + PUBLIC ${oneDNN_SOURCE_DIR}/include
228 + PUBLIC ${oneDNN_BINARY_DIR}/include
229 + PRIVATE ${oneDNN_SOURCE_DIR}/src
230 )
233 -
234 - set(ONEDNN_LIBRARY_TYPE "STATIC")
235 - set(ONEDNN_BUILD_DOC "OFF")
236 - set(ONEDNN_BUILD_EXAMPLES "OFF")
237 - set(ONEDNN_BUILD_TESTS "OFF")
238 - set(ONEDNN_ENABLE_WORKLOAD "INFERENCE")
239 - set(ONEDNN_ENABLE_PRIMITIVE "MATMUL;REORDER")
240 - set(ONEDNN_BUILD_GRAPH "OFF")
241 - set(ONEDNN_ENABLE_JIT_PROFILING "OFF")
242 - set(ONEDNN_ENABLE_ITT_TASKS "OFF")
243 - set(ONEDNN_ENABLE_MAX_CPU_ISA "OFF")
244 - set(ONEDNN_ENABLE_CPU_ISA_HINTS "OFF")
245 - set(CMAKE_POLICY_DEFAULT_CMP0077 NEW)
246 -
247 - set(DNNL_CPU_RUNTIME "OMP")
248 -
249 - FetchContent_MakeAvailable(oneDNN)
250 -
251 - list(APPEND LIBS dnnl)
231 + target_link_libraries(dnnl_ext dnnl)
232 + target_compile_options(dnnl_ext PRIVATE ${CXX_COMPILE_FLAGS} -fPIC)
233 + list(APPEND LIBS dnnl_ext)
234 + set(USE_ONEDNN ON)
235 +else()
236 + set(USE_ONEDNN OFF)
237 endif()
238
239 message(STATUS "CPU extension compile flags: ${CXX_COMPILE_FLAGS}")
@@ -275,7 +260,6 @@ set(VLLM_EXT_SRC
260
261 if (AVX512_FOUND AND NOT AVX512_DISABLED)
262 set(VLLM_EXT_SRC
278 - "csrc/cpu/quant.cpp"
263 "csrc/cpu/shm.cpp"
264 ${VLLM_EXT_SRC})
265 if (ENABLE_AVX512BF16 AND ENABLE_AVX512VNNI)
@@ -289,14 +273,11 @@ if (AVX512_FOUND AND NOT AVX512_DISABLED)
273 ${VLLM_EXT_SRC})
274 add_compile_definitions(-DCPU_CAPABILITY_AVX512)
275 endif()
292 -elseif(POWER10_FOUND)
293 - set(VLLM_EXT_SRC
294 - "csrc/cpu/quant.cpp"
295 - ${VLLM_EXT_SRC})
276 endif()
297 -if (ASIMD_FOUND)
277 +
278 +if(USE_ONEDNN)
279 set(VLLM_EXT_SRC
299 - "csrc/cpu/quant.cpp"
280 + "csrc/cpu/dnnl_kernels.cpp"
281 ${VLLM_EXT_SRC})
282 endif()
283
csrc/cpu/cpu_types_x86.hpp
+4 -4
@@ -89,7 +89,7 @@ struct FP16Vec16 : public Vec<FP16Vec16> {
89
90 explicit FP16Vec16(const FP32Vec16&);
91
92 - void save(void* ptr) const { *reinterpret_cast<__m256i*>(ptr) = reg; }
92 + void save(void* ptr) const { _mm256_storeu_si256((__m256i*)ptr, reg); }
93
94 void save(void* ptr, const int elem_num) const {
95 constexpr uint32_t M = 0xFFFFFFFF;
@@ -126,7 +126,7 @@ struct BF16Vec16 : public Vec<BF16Vec16> {
126
127 explicit BF16Vec16(const FP32Vec16&);
128
129 - void save(void* ptr) const { *reinterpret_cast<__m256i*>(ptr) = reg; }
129 + void save(void* ptr) const { _mm256_storeu_si256((__m256i*)ptr, reg); }
130
131 void save(void* ptr, const int elem_num) const {
132 constexpr uint32_t M = 0xFFFFFFFF;
@@ -180,8 +180,8 @@ struct BF16Vec32 : public Vec<BF16Vec32> {
180 (__m128i)vec8_data.reg, 1)) {}
181
182 void save(void* ptr) const {
183 - *reinterpret_cast<__m256i*>(ptr) = reg_low;
184 - *reinterpret_cast<__m256i*>((__m256i*)ptr + 1) = reg_high;
183 + _mm256_storeu_si256((__m256i*)ptr, reg_low);
184 + _mm256_storeu_si256((__m256i*)ptr + 1, reg_high);
185 }
186 };
187 #endif
csrc/cpu/dnnl_helper.cpp new
+346
@@ -0,0 +1,346 @@
1 +#include <list>
2 +#include <optional>
3 +
4 +#include "common/memory_desc.hpp"
5 +#include "common/memory.hpp"
6 +
7 +#include "dnnl_helper.h"
8 +
9 +static dnnl::engine& default_engine() {
10 + static dnnl::engine engine(dnnl::engine::kind::cpu, 0);
11 + return engine;
12 +}
13 +
14 +static dnnl::stream& default_stream() {
15 + static dnnl::stream stream(default_engine());
16 + return stream;
17 +}
18 +
19 +void release_dnnl_matmul_handler(int64_t handler) {
20 + DNNLMatMulPrimitiveHandler* ptr =
21 + reinterpret_cast<DNNLMatMulPrimitiveHandler*>(handler);
22 + delete ptr;
23 +}
24 +
25 +template <typename KT, typename VT>
26 +class DNNLPrimitiveCache {
27 + public:
28 + using cache_value_t = std::pair<KT, VT>;
29 + using result_value_t = VT;
30 + using container_t = std::list<cache_value_t>;
31 + using value_iterator_t = typename container_t::iterator;
32 + using map_t = std::unordered_map<KT, value_iterator_t>;
33 + using creator_t = VT (*)();
34 +
35 + public:
36 + DNNLPrimitiveCache(size_t capacity)
37 + : capacity_(capacity),
38 + values_(),
39 + key_to_value_(std::min(256lu, capacity)) {
40 + assert(capacity > 0);
41 + }
42 +
43 + template <typename F>
44 + result_value_t get_or_create(const KT& key, F&& creator) {
45 + std::optional<value_iterator_t> value = get_value(key);
46 + if (value.has_value()) {
47 + return value.value()->second;
48 + } else {
49 + return add_value({key, creator()})->second;
50 + }
51 + }
52 +
53 + size_t size() const { return values_.size(); }
54 +
55 + private:
56 + void dump_data() {
57 + std::stringstream ss;
58 + ss << "table_id: " << std::hex << reinterpret_cast<size_t>(this) << std::dec
59 + << "\n";
60 + ss << "container: [";
61 + for (auto&& iter : values_) {
62 + ss << "(" << iter.first << ", " << std::hex
63 + << reinterpret_cast<size_t>(iter.second.get()) << "), " << std::dec;
64 + }
65 + ss << "]\n";
66 +
67 + ss << "map: [";
68 + for (auto&& iter : key_to_value_) {
69 + ss << "(" << iter.first << ", " << iter.second->first << ", " << std::hex
70 + << reinterpret_cast<size_t>(iter.second->second.get()) << std::dec
71 + << "), ";
72 + }
73 + ss << "]\n";
74 + std::printf("%s\n", ss.str().c_str());
75 + }
76 +
77 + value_iterator_t add_value(cache_value_t&& new_value) {
78 + if (size() == capacity_) {
79 + cache_value_t& last_item = values_.back();
80 + key_to_value_.erase(last_item.first);
81 + values_.pop_back();
82 + }
83 +
84 + auto& added_value_ = values_.emplace_front(std::move(new_value));
85 + key_to_value_.emplace(added_value_.first, values_.begin());
86 + return values_.begin();
87 + }
88 +
89 + std::optional<value_iterator_t> get_value(const KT& key) {
90 + if (key_to_value_.size() > 0 && key == values_.begin()->first) {
91 + return values_.begin();
92 + }
93 +
94 + auto value_map_iterator = key_to_value_.find(key);
95 + if (value_map_iterator != key_to_value_.end()) {
96 + values_.splice(values_.begin(), values_, value_map_iterator->second);
97 + return value_map_iterator->second;
98 + } else {
99 + return {};
100 + }
101 + }
102 +
103 + private:
104 + const size_t capacity_;
105 + container_t values_;
106 + map_t key_to_value_;
107 +};
108 +
109 +DNNLMatMulPrimitiveHandler::DNNLMatMulPrimitiveHandler(
110 + const Args& args, dnnl::memory::data_type b_type)
111 + : b_n_size_(args.b_n_size),
112 + b_n_stride_(args.b_n_stride),
113 + b_k_size_(args.b_k_size),
114 + b_k_stride_(args.b_k_stride),
115 + b_type_(b_type),
116 + c_type_(args.c_type),
117 + runtime_memory_ptrs_(8),
118 + primitive_cache_size_(args.primitive_cache_size) {
119 + assert(primitive_cache_size_ > 0);
120 +}
121 +
122 +void DNNLMatMulPrimitiveHandler::prepack_weight(
123 + void* original_b_ptr, dnnl::memory::desc b_target_mem_desc) {
124 + dnnl::memory::desc original_b_md({b_k_size_, b_n_size_}, b_type_,
125 + {b_k_stride_, b_n_stride_});
126 + dnnl::memory original_weight(original_b_md, default_engine(), original_b_ptr);
127 + dnnl::memory packed_weight(b_target_mem_desc, default_engine());
128 + {
129 + dnnl::reorder(original_weight, packed_weight)
130 + .execute(default_stream(), original_weight, packed_weight);
131 + default_stream().wait();
132 + }
133 + memory_cache_[DNNL_ARG_WEIGHTS] = packed_weight;
134 + b_target_mem_desc_ = b_target_mem_desc;
135 +}
136 +
137 +void DNNLMatMulPrimitiveHandler::set_runtime_memory_ptr(
138 + size_t index, dnnl_memory* memory_ptr) {
139 + dnnl::impl::memory_storage_t* mem_storage_ptr = memory_ptr->memory_storage();
140 + dnnl_memory_desc* mem_desc = const_cast<dnnl_memory_desc*>(memory_ptr->md());
141 + runtime_memory_ptrs_[index] = {mem_storage_ptr, mem_desc};
142 +}
143 +
144 +std::pair<dnnl::impl::memory_storage_t*, dnnl_memory_desc*>
145 +DNNLMatMulPrimitiveHandler::get_runtime_memory_ptr(size_t index) {
146 + return runtime_memory_ptrs_[index];
147 +}
148 +
149 +namespace std {
150 +template <>
151 +struct hash<W8A8MatMulPrimitiveHandler::ClassMatmulCacheKey> {
152 + size_t operator()(
153 + const W8A8MatMulPrimitiveHandler::ClassMatmulCacheKey& val) const {
154 + return hash<dnnl_dim_t>()(val.b_n_size) ^ hash<dnnl_dim_t>()(val.b_k_size) ^
155 + hash<int>()(static_cast<int>(val.a_qs)) ^
156 + hash<int>()(static_cast<int>(val.b_qs)) ^ hash<bool>()(val.use_azp) ^
157 + hash<int>()(static_cast<int>(val.c_type));
158 + }
159 +};
160 +
161 +template <>
162 +struct hash<W8A8MatMulPrimitiveHandler::MSizeCacheKey> {
163 + size_t operator()(
164 + const W8A8MatMulPrimitiveHandler::MSizeCacheKey& val) const {
165 + return hash<dnnl_dim_t>()(val.a_m_size) ^ hash<bool>()(val.use_bias) ^
166 + hash<int>()(static_cast<int>(val.bias_type));
167 + }
168 +};
169 +} // namespace std
170 +
171 +bool operator==(const W8A8MatMulPrimitiveHandler::ClassMatmulCacheKey& l,
172 + const W8A8MatMulPrimitiveHandler::ClassMatmulCacheKey& r) {
173 + return l.b_n_size == r.b_n_size && l.b_k_size == r.b_k_size &&
174 + l.a_qs == r.a_qs && l.b_qs == r.b_qs && l.use_azp == r.use_azp &&
175 + l.c_type == r.c_type;
176 +}
177 +
178 +bool operator==(const W8A8MatMulPrimitiveHandler::MSizeCacheKey& l,
179 + const W8A8MatMulPrimitiveHandler::MSizeCacheKey& r) {
180 + return l.use_bias == r.use_bias && l.a_m_size == r.a_m_size &&
181 + l.bias_type == r.bias_type;
182 +}
183 +
184 +static std::shared_ptr<W8A8MatMulPrimitiveHandler::MSizeCache>
185 +get_w8a8_class_primitive_cache(
186 + const W8A8MatMulPrimitiveHandler::ClassMatmulCacheKey& key,
187 + int64_t cache_size) {
188 + static W8A8MatMulPrimitiveHandler::ClassMatmulCache cache(128);
189 + assert(cache_size > 0);
190 + return cache.get_or_create(key, [&]() {
191 + return std::make_shared<W8A8MatMulPrimitiveHandler::MSizeCache>(cache_size);
192 + });
193 +}
194 +
195 +W8A8MatMulPrimitiveHandler::W8A8MatMulPrimitiveHandler(const Args& args)
196 + : DNNLMatMulPrimitiveHandler(
197 + static_cast<const DNNLMatMulPrimitiveHandler::Args&>(args),
198 + dnnl::memory::data_type::s8),
199 + use_azp_(args.use_a_zero_point),
200 + a_qs_(args.a_quantization_strategy),
201 + b_qs_(args.b_quantization_strategy),
202 + m_size_cache_(nullptr) {
203 + assert(a_qs_ != QuantizationStrategy::PER_OUTPUT_CHANNEL);
204 + assert(b_qs_ != QuantizationStrategy::PER_TOKEN);
205 + if (a_qs_ == QuantizationStrategy::PER_TOKEN) {
206 + assert(!use_azp_);
207 + };
208 + prepack_weight(args.b_ptr,
209 + create_primitive_desc(
210 + MSizeCacheKey{.a_m_size = DNNL_RUNTIME_DIM_VAL,
211 + .use_bias = false,
212 + .bias_type = dnnl::memory::data_type::undef},
213 + true)
214 + .weights_desc());
215 + init_runtime_memory_cache(args);
216 +}
217 +
218 +void W8A8MatMulPrimitiveHandler::execute(ExecArgs& args) {
219 + auto&& [a_storage, a_mem_desc] = get_runtime_memory_ptr(0);
220 + auto&& [c_storage, c_mem_desc] = get_runtime_memory_ptr(1);
221 + a_storage->set_data_handle((void*)args.a_ptr);
222 + a_mem_desc->dims[0] = args.a_m_size;
223 + c_storage->set_data_handle((void*)args.c_ptr);
224 + c_mem_desc->dims[0] = args.a_m_size;
225 +
226 + if (a_qs_ == QuantizationStrategy::PER_TENSOR) {
227 + auto&& [a_scale_storage, a_scale_mem_desc] = get_runtime_memory_ptr(2);
228 + a_scale_storage->set_data_handle((void*)args.a_scales_ptr);
229 + }
230 + if (use_azp_) {
231 + auto&& [a_zero_point_storage, a_zero_point_mem_desc] =
232 + get_runtime_memory_ptr(3);
233 + a_zero_point_storage->set_data_handle((void*)args.a_zero_points_ptr);
234 + }
235 +
236 + if (args.use_bias) {
237 + auto&& [bias_storage, bias_mem_desc] = get_runtime_memory_ptr(4);
238 + bias_storage->set_data_handle((void*)args.bias_ptr);
239 + }
240 +
241 + dnnl::matmul matmul = get_matmul_cache(args);
242 + matmul.execute(default_stream(), memory_cache_);
243 + default_stream().wait();
244 +}
245 +
246 +dnnl::matmul W8A8MatMulPrimitiveHandler::get_matmul_cache(
247 + const MSizeCacheKey& key) {
248 + if (m_size_cache_.get() == nullptr) {
249 + ClassMatmulCacheKey key = {.b_n_size = b_n_size_,
250 + .b_k_size = b_k_size_,
251 + .a_qs = a_qs_,
252 + .b_qs = b_qs_,
253 + .use_azp = use_azp_,
254 + .c_type = c_type_};
255 + m_size_cache_ = get_w8a8_class_primitive_cache(key, primitive_cache_size_);
256 + }
257 +
258 + return m_size_cache_->get_or_create(key, [&]() {
259 + dnnl::matmul::primitive_desc desc = this->create_primitive_desc(key, false);
260 + return dnnl::matmul(desc);
261 + });
262 +}
263 +
264 +void W8A8MatMulPrimitiveHandler::init_runtime_memory_cache(const Args& args) {
265 + memory_cache_[DNNL_ARG_SRC] = dnnl::memory({{1, b_k_size_},
266 + dnnl::memory::data_type::s8,
267 + dnnl::memory::format_tag::ab},
268 + default_engine(), nullptr);
269 + set_runtime_memory_ptr(0, memory_cache_[DNNL_ARG_SRC].get());
270 + memory_cache_[DNNL_ARG_DST] =
271 + dnnl::memory({{1, b_n_size_}, c_type_, dnnl::memory::format_tag::ab},
272 + default_engine(), nullptr);
273 + set_runtime_memory_ptr(1, memory_cache_[DNNL_ARG_DST].get());
274 +
275 + // For PER_TOKEN, scales will be applied in outside epilogue
276 + if (a_qs_ == QuantizationStrategy::PER_TENSOR) {
277 + memory_cache_[DNNL_ARG_ATTR_SCALES | DNNL_ARG_SRC] = dnnl::memory(
278 + {{1}, dnnl::memory::data_type::f32, {1}}, default_engine(), nullptr);
279 + set_runtime_memory_ptr(
280 + 2, memory_cache_[DNNL_ARG_ATTR_SCALES | DNNL_ARG_SRC].get());
281 + if (use_azp_) {
282 + memory_cache_[DNNL_ARG_ATTR_ZERO_POINTS | DNNL_ARG_SRC] = dnnl::memory(
283 + {{1}, dnnl::memory::data_type::s32, {1}}, default_engine(), nullptr);
284 + set_runtime_memory_ptr(
285 + 3, memory_cache_[DNNL_ARG_ATTR_ZERO_POINTS | DNNL_ARG_SRC].get());
286 + }
287 + }
288 +
289 + if (b_qs_ == QuantizationStrategy::PER_TENSOR) {
290 + memory_cache_[DNNL_ARG_ATTR_SCALES | DNNL_ARG_WEIGHTS] =
291 + dnnl::memory({{1}, dnnl::memory::data_type::f32, {1}}, default_engine(),
292 + (void*)args.b_scales_ptr);
293 + } else if (b_qs_ == QuantizationStrategy::PER_OUTPUT_CHANNEL) {
294 + memory_cache_[DNNL_ARG_ATTR_SCALES | DNNL_ARG_WEIGHTS] =
295 + dnnl::memory({{b_n_size_}, dnnl::memory::data_type::f32, {1}},
296 + default_engine(), (void*)args.b_scales_ptr);
297 + }
298 +
299 + memory_cache_[DNNL_ARG_BIAS] =
300 + dnnl::memory({{b_n_size_}, dnnl::memory::data_type::f32, {1}},
301 + default_engine(), nullptr);
302 + set_runtime_memory_ptr(4, memory_cache_[DNNL_ARG_BIAS].get());
303 +}
304 +
305 +dnnl::matmul::primitive_desc W8A8MatMulPrimitiveHandler::create_primitive_desc(
306 + const MSizeCacheKey& key, bool first_time) {
307 + dnnl::memory::desc a_md({key.a_m_size, b_k_size_},
308 + dnnl::memory::data_type::s8,
309 + dnnl::memory::format_tag::ab);
310 + dnnl::memory::desc b_md;
311 + if (first_time) {
312 + b_md =
313 + dnnl::memory::desc({b_k_size_, b_n_size_}, dnnl::memory::data_type::s8,
314 + dnnl::memory::format_tag::any);
315 + } else {
316 + b_md = b_target_mem_desc_;
317 + }
318 + dnnl::memory::desc c_md({key.a_m_size, b_n_size_}, c_type_,
319 + dnnl::memory::format_tag::ab);
320 +
321 + dnnl::primitive_attr attr;
322 + // For PER_TOKEN, scales will be applied in outside epilogue
323 + if (a_qs_ == QuantizationStrategy::PER_TENSOR) {
324 + attr.set_scales_mask(DNNL_ARG_SRC, 0);
325 + if (use_azp_) {
326 + attr.set_zero_points_mask(DNNL_ARG_SRC, 0);
327 + }
328 + }
329 +
330 + if (b_qs_ == QuantizationStrategy::PER_TENSOR) {
331 + attr.set_scales_mask(DNNL_ARG_WEIGHTS, 0);
332 + } else if (b_qs_ == QuantizationStrategy::PER_OUTPUT_CHANNEL) {
333 + attr.set_scales_mask(DNNL_ARG_WEIGHTS, 2);
334 + }
335 +
336 + if (key.use_bias) {
337 + // For PER_TOKEN, bias will be applied in epilogue
338 + assert(a_qs_ == QuantizationStrategy::PER_TENSOR);
339 + dnnl::memory::desc bias_md({1, b_n_size_}, key.bias_type, {b_n_size_, 1});
340 + return dnnl::matmul::primitive_desc(default_engine(), a_md, b_md, bias_md,
341 + c_md, attr);
342 + } else {
343 + return dnnl::matmul::primitive_desc(default_engine(), a_md, b_md, c_md,
344 + attr);
345 + }
346 +}
csrc/cpu/dnnl_helper.h new
+169
@@ -0,0 +1,169 @@
1 +#ifndef DNNL_HELPER_H
2 +#define DNNL_HELPER_H
3 +
4 +#include <optional>
5 +#include <cassert>
6 +
7 +#include "oneapi/dnnl/dnnl.hpp"
8 +
9 +namespace c10 {
10 +struct BFloat16;
11 +struct Half;
12 +} // namespace c10
13 +
14 +namespace dnnl {
15 +namespace impl {
16 +struct memory_storage_t;
17 +struct matmul_pd_t;
18 +struct matmul_desc_t;
19 +} // namespace impl
20 +} // namespace dnnl
21 +struct dnnl_memory_desc;
22 +
23 +template <typename KT, typename VT>
24 +class DNNLPrimitiveCache;
25 +
26 +template <typename T>
27 +struct DNNLType {
28 + static constexpr dnnl::memory::data_type type =
29 + dnnl::memory::data_type::undef;
30 +};
31 +
32 +template <>
33 +struct DNNLType<int8_t> {
34 + static constexpr dnnl::memory::data_type type = dnnl::memory::data_type::s8;
35 +};
36 +
37 +template <>
38 +struct DNNLType<int32_t> {
39 + static constexpr dnnl::memory::data_type type = dnnl::memory::data_type::s32;
40 +};
41 +
42 +template <>
43 +struct DNNLType<float> {
44 + static constexpr dnnl::memory::data_type type = dnnl::memory::data_type::f32;
45 +};
46 +
47 +template <>
48 +struct DNNLType<c10::BFloat16> {
49 + static constexpr dnnl::memory::data_type type = dnnl::memory::data_type::bf16;
50 +};
51 +
52 +template <>
53 +struct DNNLType<c10::Half> {
54 + static constexpr dnnl::memory::data_type type = dnnl::memory::data_type::f16;
55 +};
56 +
57 +template <typename T>
58 +constexpr inline dnnl::memory::data_type get_dnnl_type() {
59 + return DNNLType<std::decay_t<T>>::type;
60 +}
61 +
62 +class DNNLMatMulPrimitiveHandler {
63 + public:
64 + virtual ~DNNLMatMulPrimitiveHandler() = default;
65 +
66 + protected:
67 + struct Args {
68 + dnnl_dim_t b_n_size;
69 + dnnl_dim_t b_n_stride;
70 + dnnl_dim_t b_k_size;
71 + dnnl_dim_t b_k_stride;
72 + void* b_ptr;
73 + dnnl::memory::data_type c_type;
74 + size_t primitive_cache_size;
75 + };
76 +
77 + protected:
78 + DNNLMatMulPrimitiveHandler(const Args& args, dnnl::memory::data_type b_type);
79 +
80 + void prepack_weight(void* original_b_ptr,
81 + dnnl::memory::desc b_target_mem_desc);
82 +
83 + void set_runtime_memory_ptr(size_t index, dnnl_memory* memory_ptr);
84 +
85 + std::pair<dnnl::impl::memory_storage_t*, dnnl_memory_desc*>
86 + get_runtime_memory_ptr(size_t index);
87 +
88 + protected:
89 + const dnnl_dim_t b_n_size_;
90 + const dnnl_dim_t b_n_stride_;
91 + const dnnl_dim_t b_k_size_;
92 + const dnnl_dim_t b_k_stride_;
93 + dnnl::memory::data_type b_type_;
94 + dnnl::memory::data_type c_type_;
95 + std::unordered_map<int, dnnl::memory> memory_cache_;
96 + std::vector<std::pair<dnnl::impl::memory_storage_t*, dnnl_memory_desc*>>
97 + runtime_memory_ptrs_;
98 + dnnl::memory::desc b_target_mem_desc_;
99 + int64_t primitive_cache_size_;
100 +};
101 +
102 +class W8A8MatMulPrimitiveHandler : public DNNLMatMulPrimitiveHandler {
103 + public:
104 + enum class QuantizationStrategy { PER_TOKEN, PER_TENSOR, PER_OUTPUT_CHANNEL };
105 +
106 + struct Args : public DNNLMatMulPrimitiveHandler::Args {
107 + bool use_a_zero_point;
108 + QuantizationStrategy a_quantization_strategy;
109 + QuantizationStrategy b_quantization_strategy;
110 + float* b_scales_ptr;
111 + };
112 +
113 + struct ClassMatmulCacheKey {
114 + dnnl_dim_t b_n_size;
115 + dnnl_dim_t b_k_size;
116 + QuantizationStrategy a_qs;
117 + QuantizationStrategy b_qs;
118 + bool use_azp;
119 + dnnl::memory::data_type c_type;
120 +
121 + friend bool operator==(const ClassMatmulCacheKey& l,
122 + const ClassMatmulCacheKey& r);
123 + };
124 +
125 + struct MSizeCacheKey {
126 + dnnl_dim_t a_m_size;
127 + bool use_bias;
128 + dnnl::memory::data_type bias_type;
129 +
130 + friend bool operator==(const MSizeCacheKey& l, const MSizeCacheKey& r);
131 + };
132 +
133 + using MSizeCache = DNNLPrimitiveCache<MSizeCacheKey, dnnl::matmul>;
134 + using ClassMatmulCache =
135 + DNNLPrimitiveCache<ClassMatmulCacheKey, std::shared_ptr<MSizeCache>>;
136 +
137 + struct ExecArgs : public MSizeCacheKey {
138 + const int8_t* a_ptr;
139 + const float* a_scales_ptr;
140 + const int32_t* a_zero_points_ptr;
141 + const void* bias_ptr;
142 + void* c_ptr;
143 + };
144 +
145 + public:
146 + W8A8MatMulPrimitiveHandler(const Args& args);
147 +
148 + QuantizationStrategy get_input_scale_strategy() const { return a_qs_; }
149 +
150 + bool get_input_use_zero_point() const { return use_azp_; }
151 +
152 + void execute(ExecArgs& args);
153 +
154 + private:
155 + dnnl::matmul::primitive_desc create_primitive_desc(const MSizeCacheKey& key,
156 + bool first_time);
157 +
158 + void init_runtime_memory_cache(const Args& args);
159 +
160 + dnnl::matmul get_matmul_cache(const MSizeCacheKey& key);
161 +
162 + private:
163 + const bool use_azp_;
164 + const QuantizationStrategy a_qs_;
165 + const QuantizationStrategy b_qs_;
166 + std::shared_ptr<MSizeCache> m_size_cache_;
167 +};
168 +
169 +#endif
csrc/cpu/dnnl_helper.hpp deleted
-206
@@ -1,206 +0,0 @@
1 -#ifndef DNNL_HELPER_HPP
2 -#define DNNL_HELPER_HPP
3 -
4 -#include <c10/util/BFloat16.h>
5 -#include <c10/util/Half.h>
6 -
7 -#include "oneapi/dnnl/dnnl.hpp"
8 -
9 -namespace {
10 -template <typename T>
11 -struct DNNLType {
12 - static constexpr dnnl::memory::data_type type =
13 - dnnl::memory::data_type::undef;
14 -};
15 -
16 -template <>
17 -struct DNNLType<int8_t> {
18 - static constexpr dnnl::memory::data_type type = dnnl::memory::data_type::s8;
19 -};
20 -
21 -template <>
22 -struct DNNLType<int32_t> {
23 - static constexpr dnnl::memory::data_type type = dnnl::memory::data_type::s32;
24 -};
25 -
26 -template <>
27 -struct DNNLType<float> {
28 - static constexpr dnnl::memory::data_type type = dnnl::memory::data_type::f32;
29 -};
30 -
31 -template <>
32 -struct DNNLType<c10::BFloat16> {
33 - static constexpr dnnl::memory::data_type type = dnnl::memory::data_type::bf16;
34 -};
35 -
36 -template <>
37 -struct DNNLType<c10::Half> {
38 - static constexpr dnnl::memory::data_type type = dnnl::memory::data_type::f16;
39 -};
40 -
41 -template <typename T>
42 -constexpr inline dnnl::memory::data_type get_dnnl_type() {
43 - return DNNLType<std::decay_t<T>>::type;
44 -}
45 -}; // namespace
46 -
47 -template <bool InputNoScale>
48 -class DNNLPrimitiveHelper {
49 - public:
50 - // I8 input GEMM kernel (C = a_scales * A @ (b_scales * B^T) + bias)
51 - // A: [M, K], row-major
52 - // B: [K, N], column-major
53 - // C: [M, N], row-major
54 - // bias: [N], row-major, optional
55 - // a_scales: [MS]
56 - // b_scales: [NS]
57 - // Note: Due to the limitation of oneDNN
58 - // (https://github.com/oneapi-src/oneDNN/issues/1636), the quantized bias is
59 - // not supported.
60 -
61 - template <typename OutputT, typename BiasT>
62 - static void gemm_s8s8_jit(const int8_t* a, const int8_t* b, OutputT* c,
63 - const BiasT* bias, dnnl_dim_t M, dnnl_dim_t N,
64 - dnnl_dim_t K, const float* a_scales,
65 - const float* b_scales, dnnl_dim_t MS,
66 - dnnl_dim_t NS) {
67 - auto&& OutputType = get_dnnl_type<OutputT>();
68 - auto&& BiasType = get_dnnl_type<BiasT>();
69 -
70 - dnnl::memory::desc a_md({M, K}, dnnl::memory::data_type::s8, {K, 1});
71 - dnnl::memory::desc b_md({K, N}, dnnl::memory::data_type::s8, {1, K});
72 - dnnl::memory::desc c_md({M, N}, OutputType, {N, 1});
73 -
74 - dnnl::primitive_attr attr;
75 - if constexpr (!InputNoScale) {
76 - if (MS == 1) {
77 - // per-tensor
78 - attr.set_scales_mask(DNNL_ARG_SRC, 0);
79 - } else {
80 - // per-token
81 - TORCH_CHECK(false, "per-token quantization is unsupported.");
82 - }
83 - }
84 -
85 - if (NS == 1) {
86 - // per-tensor
87 - attr.set_scales_mask(DNNL_ARG_WEIGHTS, 0);
88 - } else {
89 - // per-channel
90 - attr.set_scales_mask(DNNL_ARG_WEIGHTS, 2);
91 - }
92 -
93 - dnnl::matmul::primitive_desc matmul_pd;
94 -// Create memory descriptors with format_tag::any for the primitive. This
95 -// enables the matmul primitive to choose memory layouts for an
96 -// optimized primitive implementation, and these layouts may differ from the
97 -// ones provided by the user.
98 -#ifdef __aarch64__
99 - auto mat_src_md = dnnl::memory::desc({M, K}, dnnl::memory::data_type::s8,
100 - dnnl::memory::format_tag::any);
101 - auto mat_weights_md = dnnl::memory::desc(
102 - {K, N}, dnnl::memory::data_type::s8, dnnl::memory::format_tag::any);
103 - auto mat_dst_md =
104 - dnnl::memory::desc({M, N}, OutputType, dnnl::memory::format_tag::any);
105 - if (bias) {
106 - dnnl::memory::desc bias_md({1, N}, BiasType, {N, 1});
107 - matmul_pd = dnnl::matmul::primitive_desc(default_engine(), mat_src_md,
108 - mat_weights_md, bias_md,
109 - mat_dst_md, attr);
110 - } else {
111 - matmul_pd = dnnl::matmul::primitive_desc(
112 - default_engine(), mat_src_md, mat_weights_md, mat_dst_md, attr);
113 - }
114 -#else
115 - if (bias) {
116 - dnnl::memory::desc bias_md({1, N}, BiasType, {N, 1});
117 - matmul_pd = dnnl::matmul::primitive_desc(default_engine(), a_md, b_md,
118 - bias_md, c_md, attr);
119 - } else {
120 - matmul_pd = dnnl::matmul::primitive_desc(default_engine(), a_md, b_md,
121 - c_md, attr);
122 - }
123 -#endif
124 - dnnl::matmul matmul(matmul_pd);
125 -
126 - auto& engine = default_engine();
127 -
128 - dnnl::memory a_m(a_md, engine, (void*)a);
129 - dnnl::memory b_m(b_md, engine, (void*)b);
130 - dnnl::memory c_m(c_md, engine, (void*)c);
131 - dnnl::memory a_scales_m({{MS}, dnnl::memory::data_type::f32, {1}}, engine,
132 - (void*)a_scales);
133 - dnnl::memory b_scales_m({{NS}, dnnl::memory::data_type::f32, {1}}, engine,
134 - (void*)b_scales);
135 -
136 - auto& stream = default_stream();
137 -
138 - auto mat_src_mem = a_m;
139 - auto mat_weights_mem = b_m;
140 - auto mat_dst_mem = c_m;
141 -#ifdef __aarch64__
142 - if (matmul_pd.weights_desc() != b_m.get_desc()) {
143 - mat_weights_mem = dnnl::memory(matmul_pd.weights_desc(), engine);
144 - dnnl::reorder(b_m, mat_weights_mem).execute(stream, b_m, mat_weights_mem);
145 - }
146 -#endif
147 - if constexpr (InputNoScale) {
148 - if (bias) {
149 - dnnl::memory::desc bias_md({N}, BiasType, {1});
150 - dnnl::memory bias_m(bias_md, engine, (void*)bias);
151 - matmul.execute(
152 - stream, {
153 - {DNNL_ARG_SRC, mat_src_mem},
154 - {DNNL_ARG_WEIGHTS, mat_weights_mem},
155 - {DNNL_ARG_BIAS, bias_m},
156 - {DNNL_ARG_DST, mat_dst_mem},
157 - {DNNL_ARG_ATTR_SCALES | DNNL_ARG_WEIGHTS, b_scales_m},
158 - });
159 - } else {
160 - matmul.execute(
161 - stream, {
162 - {DNNL_ARG_SRC, mat_src_mem},
163 - {DNNL_ARG_WEIGHTS, mat_weights_mem},
164 - {DNNL_ARG_DST, mat_dst_mem},
165 - {DNNL_ARG_ATTR_SCALES | DNNL_ARG_WEIGHTS, b_scales_m},
166 - });
167 - }
168 - } else {
169 - if (bias) {
170 - dnnl::memory::desc bias_md({N}, BiasType, {1});
171 - dnnl::memory bias_m(bias_md, engine, (void*)bias);
172 - matmul.execute(
173 - stream, {
174 - {DNNL_ARG_SRC, mat_src_mem},
175 - {DNNL_ARG_WEIGHTS, mat_weights_mem},
176 - {DNNL_ARG_BIAS, bias_m},
177 - {DNNL_ARG_DST, mat_dst_mem},
178 - {DNNL_ARG_ATTR_SCALES | DNNL_ARG_SRC, a_scales_m},
179 - {DNNL_ARG_ATTR_SCALES | DNNL_ARG_WEIGHTS, b_scales_m},
180 - });
181 - } else {
182 - matmul.execute(
183 - stream, {
184 - {DNNL_ARG_SRC, mat_src_mem},
185 - {DNNL_ARG_WEIGHTS, mat_weights_mem},
186 - {DNNL_ARG_DST, mat_dst_mem},
187 - {DNNL_ARG_ATTR_SCALES | DNNL_ARG_SRC, a_scales_m},
188 - {DNNL_ARG_ATTR_SCALES | DNNL_ARG_WEIGHTS, b_scales_m},
189 - });
190 - }
191 - }
192 - stream.wait();
193 - }
194 -
195 - private:
196 - static dnnl::engine& default_engine() {
197 - static dnnl::engine engine(dnnl::engine::kind::cpu, 0);
198 - return engine;
199 - }
200 -
201 - static dnnl::stream& default_stream() {
202 - static dnnl::stream stream(default_engine());
203 - return stream;
204 - }
205 -};
206 -#endif
csrc/cpu/dnnl_kernels.cpp new
+494
@@ -0,0 +1,494 @@
1 +#include "cpu_types.hpp"
2 +#include "dnnl_helper.h"
3 +
4 +namespace {
5 +template <typename scalar_t>
6 +struct KernelVecType {
7 + using load_vec_type = void;
8 + using cvt_vec_type = void;
9 +};
10 +
11 +template <>
12 +struct KernelVecType<float> {
13 + using load_vec_type = vec_op::FP32Vec16;
14 + using cvt_vec_type = vec_op::FP32Vec16;
15 +};
16 +
17 +#if !defined(__aarch64__) || defined(ARM_BF16_SUPPORT)
18 +template <>
19 +struct KernelVecType<c10::BFloat16> {
20 + using load_vec_type = vec_op::BF16Vec16;
21 + using cvt_vec_type = vec_op::FP32Vec16;
22 +};
23 +#endif
24 +
25 +template <>
26 +struct KernelVecType<c10::Half> {
27 +#if defined(__powerpc64__) || defined(__s390x__)
28 + // Power architecture-specific vector type
29 + using load_vec_type = vec_op::FP32Vec16;
30 +#else
31 + // Fallback for other architectures
32 + using load_vec_type = vec_op::FP16Vec16;
33 +#endif
34 + using cvt_vec_type = vec_op::FP32Vec16;
35 +};
36 +
37 +template <bool AZP, typename scalar_t>
38 +void static_scaled_int8_quant_impl(const scalar_t* input, int8_t* output,
39 + const float* scale, const int32_t* azp,
40 + const int64_t num_tokens,
41 + const int64_t input_stride,
42 + const int64_t hidden_size) {
43 + using load_vec_t = typename KernelVecType<scalar_t>::load_vec_type;
44 + using cvt_vec_t = typename KernelVecType<scalar_t>::cvt_vec_type;
45 + constexpr int64_t vec_elem_num = load_vec_t::VEC_ELEM_NUM;
46 +
47 + constexpr float i8_min =
48 + static_cast<float>(std::numeric_limits<int8_t>::min());
49 + constexpr float i8_max =
50 + static_cast<float>(std::numeric_limits<int8_t>::max());
51 + const cvt_vec_t inv_scale(1.0 / *scale);
52 + const cvt_vec_t i8_min_vec(i8_min);
53 + const cvt_vec_t i8_max_vec(i8_max);
54 +
55 + cvt_vec_t zp_vec;
56 + if constexpr (AZP) {
57 + zp_vec = cvt_vec_t(static_cast<float>(*azp));
58 + }
59 +
60 +#pragma omp parallel for
61 + for (int64_t i = 0; i < num_tokens; ++i) {
62 + int64_t j = 0;
63 + const scalar_t* input_ptr = input + i * input_stride;
64 + int8_t* output_ptr = output + i * hidden_size;
65 + for (; j < hidden_size - vec_elem_num; j += vec_elem_num) {
66 + load_vec_t elems(input_ptr + j);
67 + cvt_vec_t elems_fp32(elems);
68 + elems_fp32 = elems_fp32 * inv_scale;
69 +
70 + if constexpr (AZP) {
71 + elems_fp32 = elems_fp32 + zp_vec;
72 + }
73 +
74 + elems_fp32 = elems_fp32.clamp(i8_min_vec, i8_max_vec);
75 + vec_op::INT8Vec16 elems_int8(elems_fp32);
76 + elems_int8.save(output_ptr + j);
77 + }
78 +
79 + load_vec_t elems(input_ptr + j);
80 + cvt_vec_t elems_fp32(elems);
81 + elems_fp32 = elems_fp32 * inv_scale;
82 +
83 + if constexpr (AZP) {
84 + elems_fp32 = elems_fp32 + zp_vec;
85 + }
86 +
87 + elems_fp32 = elems_fp32.clamp(i8_min_vec, i8_max_vec);
88 + vec_op::INT8Vec16 elems_int8(elems_fp32);
89 + elems_int8.save(output_ptr + j, hidden_size - j);
90 + }
91 +}
92 +
93 +template <bool AZP, typename scalar_t>
94 +void dynamic_scaled_int8_quant_impl(const scalar_t* input, int8_t* output,
95 + float* scale, int32_t* azp,
96 + const int64_t num_tokens,
97 + const int64_t input_stride,
98 + const int64_t hidden_size) {
99 + using load_vec_t = typename KernelVecType<scalar_t>::load_vec_type;
100 + using cvt_vec_t = typename KernelVecType<scalar_t>::cvt_vec_type;
101 + constexpr int vec_elem_num = load_vec_t::VEC_ELEM_NUM;
102 +
103 + constexpr float i8_min =
104 + static_cast<float>(std::numeric_limits<int8_t>::min());
105 + constexpr float i8_max =
106 + static_cast<float>(std::numeric_limits<int8_t>::max());
107 + const cvt_vec_t i8_min_vec(i8_min);
108 + const cvt_vec_t i8_max_vec(i8_max);
109 +
110 +#pragma omp parallel for
111 + for (int64_t i = 0; i < num_tokens; ++i) {
112 + cvt_vec_t max_value(std::numeric_limits<float>::lowest());
113 + cvt_vec_t min_value(std::numeric_limits<float>::max());
114 + {
115 + int64_t j = 0;
116 + const scalar_t* input_ptr = input + i * input_stride;
117 + for (; j < hidden_size - vec_elem_num; j += vec_elem_num) {
118 + load_vec_t elems(input_ptr + j);
119 + cvt_vec_t elems_fp32(elems);
120 + if constexpr (AZP) {
121 + max_value = max_value.max(elems_fp32);
122 + min_value = min_value.min(elems_fp32);
123 + } else {
124 + max_value = max_value.max(elems_fp32.abs());
125 + }
126 + }
127 +
128 + load_vec_t elems(input_ptr + j);
129 + cvt_vec_t elems_fp32(elems);
130 +
131 + if (j + vec_elem_num == hidden_size) {
132 + if constexpr (AZP) {
133 + max_value = max_value.max(elems_fp32);
134 + min_value = min_value.min(elems_fp32);
135 + } else {
136 + max_value = max_value.max(elems_fp32.abs());
137 + }
138 + } else {
139 + if constexpr (AZP) {
140 + max_value = max_value.max(elems_fp32, hidden_size - j);
141 + min_value = min_value.min(elems_fp32, hidden_size - j);
142 + } else {
143 + max_value = max_value.max(elems_fp32.abs(), hidden_size - j);
144 + }
145 + }
146 + }
147 +
148 + float scale_val, azp_val;
149 + if constexpr (AZP) {
150 + float max_scalar = max_value.reduce_max();
151 + float min_scalar = min_value.reduce_min();
152 + scale_val = (max_scalar - min_scalar) / 255.0f;
153 + azp_val = std::nearbyint(-128.0f - min_scalar / scale_val);
154 + azp[i] = azp_val;
155 + scale[i] = scale_val;
156 + } else {
157 + scale_val = max_value.reduce_max() / 127.0f;
158 + scale[i] = scale_val;
159 + }
160 +
161 + const cvt_vec_t inv_scale(1.0 / scale_val);
162 + const cvt_vec_t azp_vec(azp_val);
163 +
164 + {
165 + int64_t j = 0;
166 + const scalar_t* input_ptr = input + i * input_stride;
167 + int8_t* output_ptr = output + i * hidden_size;
168 + for (; j < hidden_size - vec_elem_num; j += vec_elem_num) {
169 + load_vec_t elems(input_ptr + j);
170 + cvt_vec_t elems_fp32(elems);
171 + elems_fp32 = (elems_fp32 * inv_scale);
172 +
173 + if constexpr (AZP) {
174 + elems_fp32 = elems_fp32 + azp_vec;
175 + }
176 + elems_fp32 = elems_fp32.clamp(i8_min_vec, i8_max_vec);
177 + vec_op::INT8Vec16 elems_int8(elems_fp32);
178 + elems_int8.save(output_ptr + j);
179 + }
180 +
181 + load_vec_t elems(input_ptr + j);
182 + cvt_vec_t elems_fp32(elems);
183 + elems_fp32 = (elems_fp32 * inv_scale);
184 +
185 + if constexpr (AZP) {
186 + elems_fp32 = elems_fp32 + azp_vec;
187 + }
188 + elems_fp32 = elems_fp32.clamp(i8_min_vec, i8_max_vec);
189 + vec_op::INT8Vec16 elems_int8(elems_fp32);
190 + elems_int8.save(output_ptr + j, hidden_size - j);
191 + }
192 + }
193 +}
194 +
195 +template <bool AZP, bool Bias, typename scalar_t>
196 +void dynamic_quant_epilogue(const float* input, scalar_t* output,
197 + const float* a_scale, const int32_t* azp,
198 + const float* azp_adj, const scalar_t* bias,
199 + const int64_t num_tokens,
200 + const int64_t hidden_size) {
201 + CPU_KERNEL_GUARD_IN(dynamic_quant_epilogue)
202 + using load_vec_t = typename KernelVecType<scalar_t>::load_vec_type;
203 + using cvt_vec_t = typename KernelVecType<scalar_t>::cvt_vec_type;
204 + constexpr int vec_elem_num = load_vec_t::VEC_ELEM_NUM;
205 +
206 + const int64_t thread_num = omp_get_max_threads();
207 + if (num_tokens > thread_num) {
208 +#pragma omp parallel for
209 + for (int64_t i = 0; i < num_tokens; ++i) {
210 + const float* input_ptr = input + i * hidden_size;
211 + scalar_t* output_ptr = output + i * hidden_size;
212 + int64_t j = 0;
213 + cvt_vec_t token_scale_vec(a_scale[i]);
214 + cvt_vec_t token_zp_scale_vec;
215 + if constexpr (AZP) {
216 + float zp_scale_val = a_scale[i] * static_cast<float>(azp[i]);
217 + token_zp_scale_vec = cvt_vec_t(zp_scale_val);
218 + }
219 + for (; j < hidden_size - vec_elem_num; ++j) {
220 + cvt_vec_t elems_fp32(input_ptr + j);
221 + elems_fp32 = elems_fp32 * token_scale_vec;
222 + if constexpr (AZP) {
223 + cvt_vec_t azp_adj_fp32(azp_adj + j);
224 + elems_fp32 = elems_fp32 - azp_adj_fp32 * token_zp_scale_vec;
225 + }
226 + if constexpr (Bias) {
227 + load_vec_t bias_vec(bias + j);
228 + cvt_vec_t bias_vec_fp32(bias_vec);
229 + elems_fp32 = elems_fp32 + bias_vec_fp32;
230 + }
231 + load_vec_t elems_out(elems_fp32);
232 + elems_out.save(output_ptr + j);
233 + }
234 + cvt_vec_t elems_fp32(input_ptr + j);
235 + elems_fp32 = elems_fp32 * token_scale_vec;
236 + if constexpr (AZP) {
237 + cvt_vec_t azp_adj_fp32(azp_adj + j);
238 + elems_fp32 = elems_fp32 - azp_adj_fp32 * token_zp_scale_vec;
239 + }
240 + if constexpr (Bias) {
241 + load_vec_t bias_vec(bias + j);
242 + cvt_vec_t bias_vec_fp32(bias_vec);
243 + elems_fp32 = elems_fp32 + bias_vec_fp32;
244 + }
245 + load_vec_t elems_out(elems_fp32);
246 + elems_out.save(output_ptr + j, hidden_size - j);
247 + }
248 + } else {
249 + const int64_t vec_iteration =
250 + (hidden_size + vec_elem_num - 1) / vec_elem_num;
251 + const int64_t vec_iteration_per_thread =
252 + (vec_iteration + thread_num - 1) / thread_num;
253 + const int64_t elem_num_per_thread = vec_iteration_per_thread * vec_elem_num;
254 +#pragma omp parallel for schedule(static, 1)
255 + for (int64_t i = 0; i < thread_num; ++i) {
256 + const int64_t start = elem_num_per_thread * i;
257 + const int64_t end = std::min(hidden_size, elem_num_per_thread + start);
258 + for (int64_t j = 0; j < num_tokens; ++j) {
259 + cvt_vec_t token_scale_vec(a_scale[j]);
260 + cvt_vec_t token_zp_scale_vec;
261 + if constexpr (AZP) {
262 + float zp_scale_val = a_scale[j] * static_cast<float>(azp[j]);
263 + token_zp_scale_vec = cvt_vec_t(zp_scale_val);
264 + }
265 + int64_t k = start;
266 + const float* input_ptr = input + j * hidden_size;
267 + scalar_t* output_ptr = output + j * hidden_size;
268 + for (; k < end - vec_elem_num; k += vec_elem_num) {
269 + cvt_vec_t elems_fp32(input_ptr + k);
270 + elems_fp32 = elems_fp32 * token_scale_vec;
271 + if constexpr (AZP) {
272 + cvt_vec_t azp_adj_fp32(azp_adj + k);
273 + elems_fp32 = elems_fp32 - azp_adj_fp32 * token_zp_scale_vec;
274 + }
275 + if constexpr (Bias) {
276 + load_vec_t bias_vec(bias + k);
277 + cvt_vec_t bias_vec_fp32(bias_vec);
278 + elems_fp32 = elems_fp32 + bias_vec_fp32;
279 + }
280 + load_vec_t elems_out(elems_fp32);
281 + elems_out.save(output_ptr + k);
282 + }
283 + if (k < end) {
284 + cvt_vec_t elems_fp32(input_ptr + k);
285 + elems_fp32 = elems_fp32 * token_scale_vec;
286 + if constexpr (AZP) {
287 + cvt_vec_t azp_adj_fp32(azp_adj + k);
288 + elems_fp32 = elems_fp32 - azp_adj_fp32 * token_zp_scale_vec;
289 + }
290 + if constexpr (Bias) {
291 + load_vec_t bias_vec(bias + k);
292 + cvt_vec_t bias_vec_fp32(bias_vec);
293 + elems_fp32 = elems_fp32 + bias_vec_fp32;
294 + }
295 + load_vec_t elems_out(elems_fp32);
296 + elems_out.save(output_ptr + k, end - k);
297 + }
298 + }
299 + }
300 + }
301 +}
302 +} // namespace
303 +
304 +int64_t create_onednn_scaled_mm_handler(
305 + const torch::Tensor& b, // [IC, OC], column-major
306 + const torch::Tensor& b_scales, // [1] or [OC]
307 + at::ScalarType output_type, bool dynamic_act_quant, bool use_azp,
308 + int64_t primitive_cache_size) {
309 + TORCH_CHECK(b.dim() == 2);
310 + TORCH_CHECK(b.stride(0) == 1); // Column-major
311 + TORCH_CHECK(b_scales.is_contiguous());
312 +
313 + W8A8MatMulPrimitiveHandler::Args args;
314 + args.primitive_cache_size = primitive_cache_size;
315 +
316 + if (b_scales.numel() == 1) {
317 + args.b_quantization_strategy =
318 + W8A8MatMulPrimitiveHandler::QuantizationStrategy::PER_TENSOR;
319 + } else {
320 + TORCH_CHECK_EQ(b_scales.numel(), b.size(1));
321 + args.b_quantization_strategy =
322 + W8A8MatMulPrimitiveHandler::QuantizationStrategy::PER_OUTPUT_CHANNEL;
323 + }
324 + args.b_scales_ptr = b_scales.data_ptr<float>();
325 + args.b_k_size = b.size(0);
326 + args.b_k_stride = b.stride(0);
327 + args.b_n_size = b.size(1);
328 + args.b_n_stride = b.stride(1);
329 + args.b_ptr = b.data_ptr<int8_t>();
330 +
331 + if (dynamic_act_quant) {
332 + // dynamic per-token, bias, A scales and A zps will be applied in outside.
333 + args.a_quantization_strategy =
334 + W8A8MatMulPrimitiveHandler::QuantizationStrategy::PER_TOKEN;
335 + args.use_a_zero_point = false;
336 + } else {
337 + // static per-tensor
338 + args.a_quantization_strategy =
339 + W8A8MatMulPrimitiveHandler::QuantizationStrategy::PER_TENSOR;
340 + args.use_a_zero_point = use_azp;
341 + }
342 +
343 + VLLM_DISPATCH_FLOATING_TYPES(output_type, "create_onednn_scaled_mm_handler",
344 + [&] {
345 + if (dynamic_act_quant) {
346 + args.c_type = get_dnnl_type<float>();
347 + } else {
348 + args.c_type = get_dnnl_type<scalar_t>();
349 + }
350 + });
351 +
352 + return reinterpret_cast<int64_t>(new W8A8MatMulPrimitiveHandler(args));
353 +}
354 +
355 +void onednn_scaled_mm(
356 + torch::Tensor& c, // [M, OC], row-major
357 + const torch::Tensor& a, // [M, IC], row-major
358 + const torch::Tensor& a_scales, // [M] or [1]
359 + const std::optional<torch::Tensor>& azp, // [M] or [1]
360 + const std::optional<torch::Tensor>& azp_adj, // [M] or [1]
361 + const std::optional<torch::Tensor>& bias, // [N]
362 + int64_t handler) {
363 + CPU_KERNEL_GUARD_IN(onednn_scaled_mm)
364 + TORCH_CHECK(a.dim() == 2);
365 + TORCH_CHECK(a.is_contiguous());
366 + TORCH_CHECK(c.is_contiguous());
367 + W8A8MatMulPrimitiveHandler* ptr =
368 + reinterpret_cast<W8A8MatMulPrimitiveHandler*>(handler);
369 + const int32_t* azp_ptr = nullptr;
370 + if (azp.has_value()) {
371 + azp_ptr = azp->data_ptr<int32_t>();
372 + }
373 + if (ptr->get_input_scale_strategy() ==
374 + W8A8MatMulPrimitiveHandler::QuantizationStrategy::PER_TENSOR) {
375 + TORCH_CHECK_EQ(a_scales.numel(), 1);
376 + }
377 +
378 + W8A8MatMulPrimitiveHandler::ExecArgs exec_args;
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.use_bias = false;
383 + exec_args.a_scales_ptr = nullptr;
384 + exec_args.a_zero_points_ptr = nullptr;
385 +
386 + VLLM_DISPATCH_FLOATING_TYPES(c.scalar_type(), "onednn_scaled_mm", [&] {
387 + if (ptr->get_input_scale_strategy() ==
388 + W8A8MatMulPrimitiveHandler::QuantizationStrategy::PER_TENSOR) {
389 + if (bias.has_value()) {
390 + exec_args.bias_ptr = bias->data_ptr<scalar_t>();
391 + exec_args.bias_type = get_dnnl_type<scalar_t>();
392 + exec_args.use_bias = true;
393 + }
394 + exec_args.a_scales_ptr = a_scales.data_ptr<float>();
395 + exec_args.a_zero_points_ptr = azp_ptr;
396 + exec_args.c_ptr = c.data_ptr<scalar_t>();
397 + ptr->execute(exec_args);
398 + } else if (ptr->get_input_scale_strategy() ==
399 + W8A8MatMulPrimitiveHandler::QuantizationStrategy::PER_TOKEN) {
400 + torch::Tensor tmp_fp32_out =
401 + torch::empty_like(c, ::at::ScalarType::Float);
402 + exec_args.c_ptr = tmp_fp32_out.data_ptr<float>();
403 + ptr->execute(exec_args);
404 + if (bias.has_value()) {
405 + if (azp.has_value()) {
406 + dynamic_quant_epilogue<true, true>(
407 + tmp_fp32_out.data_ptr<float>(), c.data_ptr<scalar_t>(),
408 + a_scales.data_ptr<float>(), azp_ptr, azp_adj->data_ptr<float>(),
409 + bias->data_ptr<scalar_t>(), c.size(0), c.size(1));
410 + } else {
411 + dynamic_quant_epilogue<false, true>(
412 + tmp_fp32_out.data_ptr<float>(), c.data_ptr<scalar_t>(),
413 + a_scales.data_ptr<float>(), azp_ptr, nullptr,
414 + bias->data_ptr<scalar_t>(), c.size(0), c.size(1));
415 + }
416 + } else {
417 + if (azp.has_value()) {
418 + dynamic_quant_epilogue<true, false>(
419 + tmp_fp32_out.data_ptr<float>(), c.data_ptr<scalar_t>(),
420 + a_scales.data_ptr<float>(), azp_ptr, azp_adj->data_ptr<float>(),
421 + (scalar_t*)nullptr, c.size(0), c.size(1));
422 + } else {
423 + dynamic_quant_epilogue<false, false>(
424 + tmp_fp32_out.data_ptr<float>(), c.data_ptr<scalar_t>(),
425 + a_scales.data_ptr<float>(), azp_ptr, nullptr, (scalar_t*)nullptr,
426 + c.size(0), c.size(1));
427 + }
428 + }
429 + } else {
430 + TORCH_CHECK(false, "invalid act quant type.");
431 + }
432 + });
433 +}
434 +
435 +// static-per-tensor quantization.
436 +void static_scaled_int8_quant(
437 + torch::Tensor& out, // [batch, hidden_size]
438 + const torch::Tensor& input, // [batch, hidden_size]
439 + const torch::Tensor& scale, std::optional<torch::Tensor> const& azp) {
440 + CPU_KERNEL_GUARD_IN(static_scaled_int8_quant)
441 + TORCH_CHECK(out.is_contiguous());
442 + TORCH_CHECK_EQ(input.dim(), 2);
443 + TORCH_CHECK_EQ(input.stride(1), 1);
444 + TORCH_CHECK(scale.numel() == 1);
445 + TORCH_CHECK(!azp.has_value() || azp->numel() == 1);
446 +
447 + const int64_t stride = input.stride(0);
448 + const int64_t hidden_size = input.size(1);
449 + const int64_t num_tokens = input.size(0);
450 + VLLM_DISPATCH_FLOATING_TYPES(
451 + input.scalar_type(), "static_scaled_int8_quant_impl", [&] {
452 + if (azp.has_value()) {
453 + static_scaled_int8_quant_impl<true>(
454 + input.data_ptr<scalar_t>(), out.data_ptr<int8_t>(),
455 + scale.data_ptr<float>(), azp->data_ptr<int32_t>(), num_tokens,
456 + stride, hidden_size);
457 + } else {
458 + static_scaled_int8_quant_impl<false>(input.data_ptr<scalar_t>(),
459 + out.data_ptr<int8_t>(),
460 + scale.data_ptr<float>(), nullptr,
461 + num_tokens, stride, hidden_size);
462 + }
463 + });
464 +}
465 +
466 +// dynamic-per-token quantization.
467 +void dynamic_scaled_int8_quant(
468 + torch::Tensor& out, // [batch, hidden_size]
469 + const torch::Tensor& input, // [batch, hidden_size]
470 + torch::Tensor& scale, // [batch, 1]
471 + std::optional<torch::Tensor> const& azp) {
472 + CPU_KERNEL_GUARD_IN(dynamic_scaled_int8_quant)
473 + TORCH_CHECK(out.is_contiguous());
474 + TORCH_CHECK_EQ(input.dim(), 2);
475 + TORCH_CHECK_EQ(input.stride(1), 1);
476 +
477 + const int64_t hidden_size = input.size(1);
478 + const int64_t num_tokens = input.size(0);
479 + const int64_t stride = input.stride(0);
480 + VLLM_DISPATCH_FLOATING_TYPES(
481 + input.scalar_type(), "dynamic_scaled_int8_quant_impl", [&] {
482 + if (azp.has_value()) {
483 + dynamic_scaled_int8_quant_impl<true>(
484 + input.data_ptr<scalar_t>(), out.data_ptr<int8_t>(),
485 + scale.data_ptr<float>(), azp->data_ptr<int32_t>(), num_tokens,
486 + stride, hidden_size);
487 + } else {
488 + dynamic_scaled_int8_quant_impl<false>(
489 + input.data_ptr<scalar_t>(), out.data_ptr<int8_t>(),
490 + scale.data_ptr<float>(), nullptr, num_tokens, stride,
491 + hidden_size);
492 + }
493 + });
494 +}
csrc/cpu/quant.cpp deleted
-951
@@ -1,951 +0,0 @@
1 -#include "cpu_types.hpp"
2 -#include "dnnl_helper.hpp"
3 -
4 -namespace {
5 -template <typename scalar_t>
6 -struct KernelVecType {
7 - using load_vec_type = void;
8 - using azp_adj_load_vec_type = void;
9 - using cvt_vec_type = void;
10 -};
11 -
12 -template <>
13 -struct KernelVecType<float> {
14 - using load_vec_type = vec_op::FP32Vec16;
15 - using azp_adj_load_vec_type = vec_op::INT32Vec16;
16 - using cvt_vec_type = vec_op::FP32Vec16;
17 -};
18 -
19 -#if !defined(__aarch64__) || defined(ARM_BF16_SUPPORT)
20 -template <>
21 -struct KernelVecType<c10::BFloat16> {
22 - using load_vec_type = vec_op::BF16Vec16;
23 - using azp_adj_load_vec_type = vec_op::INT32Vec16;
24 - using cvt_vec_type = vec_op::FP32Vec16;
25 -};
26 -#endif
27 -
28 -template <>
29 -struct KernelVecType<c10::Half> {
30 -#if defined(__powerpc64__) || defined(__s390x__)
31 - // Power architecture-specific vector type
32 - using load_vec_type = vec_op::FP32Vec16;
33 -#else
34 - // Fallback for other architectures
35 - using load_vec_type = vec_op::FP16Vec16;
36 -#endif
37 - using azp_adj_load_vec_type = vec_op::INT32Vec16;
38 - using cvt_vec_type = vec_op::FP32Vec16;
39 -};
40 -
41 -#if defined(__AVX512F__) || defined(__aarch64__)
42 -template <bool AZP, typename scalar_t>
43 -void static_scaled_int8_quant_impl(const scalar_t* input, int8_t* output,
44 - const float* scale, const int32_t* azp,
45 - const int num_tokens,
46 - const int hidden_size) {
47 - using load_vec_t = typename KernelVecType<scalar_t>::load_vec_type;
48 - using cvt_vec_t = typename KernelVecType<scalar_t>::cvt_vec_type;
49 - constexpr int vec_elem_num = load_vec_t::VEC_ELEM_NUM;
50 -
51 - constexpr float i8_min =
52 - static_cast<float>(std::numeric_limits<int8_t>::min());
53 - constexpr float i8_max =
54 - static_cast<float>(std::numeric_limits<int8_t>::max());
55 - const cvt_vec_t inv_scale(1.0 / *scale);
56 - const cvt_vec_t i8_min_vec(i8_min);
57 - const cvt_vec_t i8_max_vec(i8_max);
58 -
59 - cvt_vec_t zp_vec;
60 - if constexpr (AZP) {
61 - zp_vec = cvt_vec_t(static_cast<float>(*azp));
62 - }
63 -
64 - #pragma omp parallel for
65 - for (int i = 0; i < num_tokens; ++i) {
66 - int j = 0;
67 - for (; j < hidden_size - vec_elem_num; j += vec_elem_num) {
68 - load_vec_t elems(input + i * hidden_size + j);
69 - cvt_vec_t elems_fp32(elems);
70 - elems_fp32 = elems_fp32 * inv_scale;
71 -
72 - if constexpr (AZP) {
73 - elems_fp32 = elems_fp32 + zp_vec;
74 - }
75 -
76 - elems_fp32 = elems_fp32.clamp(i8_min_vec, i8_max_vec);
77 - vec_op::INT8Vec16 elems_int8(elems_fp32);
78 - elems_int8.save(output + i * hidden_size + j);
79 - }
80 -
81 - load_vec_t elems(input + i * hidden_size + j);
82 - cvt_vec_t elems_fp32(elems);
83 - elems_fp32 = elems_fp32 * inv_scale;
84 -
85 - if constexpr (AZP) {
86 - elems_fp32 = elems_fp32 + zp_vec;
87 - }
88 -
89 - elems_fp32 = elems_fp32.clamp(i8_min_vec, i8_max_vec);
90 - vec_op::INT8Vec16 elems_int8(elems_fp32);
91 - elems_int8.save(output + i * hidden_size + j, hidden_size - j);
92 - }
93 -}
94 -
95 -template <bool AZP, typename scalar_t>
96 -void dynamic_scaled_int8_quant_impl(const scalar_t* input, int8_t* output,
97 - float* scale, int32_t* azp,
98 - const int num_tokens,
99 - const int hidden_size) {
100 - using load_vec_t = typename KernelVecType<scalar_t>::load_vec_type;
101 - using cvt_vec_t = typename KernelVecType<scalar_t>::cvt_vec_type;
102 - constexpr int vec_elem_num = load_vec_t::VEC_ELEM_NUM;
103 -
104 - constexpr float i8_min =
105 - static_cast<float>(std::numeric_limits<int8_t>::min());
106 - constexpr float i8_max =
107 - static_cast<float>(std::numeric_limits<int8_t>::max());
108 - const cvt_vec_t i8_min_vec(i8_min);
109 - const cvt_vec_t i8_max_vec(i8_max);
110 -
111 - #pragma omp parallel for
112 - for (int i = 0; i < num_tokens; ++i) {
113 - cvt_vec_t max_value(std::numeric_limits<float>::lowest());
114 - cvt_vec_t min_value(std::numeric_limits<float>::max());
115 - {
116 - int j = 0;
117 - for (; j < hidden_size - vec_elem_num; j += vec_elem_num) {
118 - load_vec_t elems(input + i * hidden_size + j);
119 - cvt_vec_t elems_fp32(elems);
120 - if constexpr (AZP) {
121 - max_value = max_value.max(elems_fp32);
122 - min_value = min_value.min(elems_fp32);
123 - } else {
124 - max_value = max_value.max(elems_fp32.abs());
125 - }
126 - }
127 -
128 - load_vec_t elems(input + i * hidden_size + j);
129 - cvt_vec_t elems_fp32(elems);
130 -
131 - if (j + vec_elem_num == hidden_size) {
132 - if constexpr (AZP) {
133 - max_value = max_value.max(elems_fp32);
134 - min_value = min_value.min(elems_fp32);
135 - } else {
136 - max_value = max_value.max(elems_fp32.abs());
137 - }
138 - } else {
139 - if constexpr (AZP) {
140 - max_value = max_value.max(elems_fp32, hidden_size - j);
141 - min_value = min_value.min(elems_fp32, hidden_size - j);
142 - } else {
143 - max_value = max_value.max(elems_fp32.abs(), hidden_size - j);
144 - }
145 - }
146 - }
147 -
148 - float scale_val, azp_val;
149 - if constexpr (AZP) {
150 - float max_scalar = max_value.reduce_max();
151 - float min_scalar = min_value.reduce_min();
152 - scale_val = (max_scalar - min_scalar) / 255.0f;
153 - azp_val = std::nearbyint(-128.0f - min_scalar / scale_val);
154 - azp[i] = static_cast<int32_t>(azp_val);
155 - scale[i] = scale_val;
156 - } else {
157 - scale_val = max_value.reduce_max() / 127.0f;
158 - scale[i] = scale_val;
159 - }
160 -
161 - const cvt_vec_t inv_scale(1.0 / scale_val);
162 - const cvt_vec_t azp_vec(azp_val);
163 -
164 - {
165 - int j = 0;
166 - for (; j < hidden_size - vec_elem_num; j += vec_elem_num) {
167 - load_vec_t elems(input + i * hidden_size + j);
168 - cvt_vec_t elems_fp32(elems);
169 - elems_fp32 = (elems_fp32 * inv_scale);
170 -
171 - if constexpr (AZP) {
172 - elems_fp32 = elems_fp32 + azp_vec;
173 - }
174 - elems_fp32 = elems_fp32.clamp(i8_min_vec, i8_max_vec);
175 - vec_op::INT8Vec16 elems_int8(elems_fp32);
176 - elems_int8.save(output + i * hidden_size + j);
177 - }
178 -
179 - load_vec_t elems(input + i * hidden_size + j);
180 - cvt_vec_t elems_fp32(elems);
181 - elems_fp32 = (elems_fp32 * inv_scale);
182 -
183 - if constexpr (AZP) {
184 - elems_fp32 = elems_fp32 + azp_vec;
185 - }
186 - elems_fp32 = elems_fp32.clamp(i8_min_vec, i8_max_vec);
187 - vec_op::INT8Vec16 elems_int8(elems_fp32);
188 - elems_int8.save(output + i * hidden_size + j, hidden_size - j);
189 - }
190 - }
191 -}
192 -
193 -template <bool PerChannel, typename scalar_t>
194 -void static_quant_epilogue(const float* input, scalar_t* output,
195 - const float a_scale, const float* b_scale,
196 - const int32_t* azp_with_adj, const int num_tokens,
197 - const int hidden_size) {
198 - CPU_KERNEL_GUARD_IN(dynamic_output_scale_impl)
199 - using load_vec_t = typename KernelVecType<scalar_t>::load_vec_type;
200 - using azp_adj_load_vec_t =
201 - typename KernelVecType<scalar_t>::azp_adj_load_vec_type;
202 - using cvt_vec_t = typename KernelVecType<scalar_t>::cvt_vec_type;
203 - constexpr int vec_elem_num = load_vec_t::VEC_ELEM_NUM;
204 -
205 - #pragma omp parallel for
206 - for (int i = 0; i < num_tokens; ++i) {
207 - cvt_vec_t a_scale_vec(a_scale);
208 - cvt_vec_t b_scale_vec(*b_scale);
209 - cvt_vec_t scale_vec = a_scale_vec * b_scale_vec;
210 -
211 - int j = 0;
212 - for (; j < hidden_size - vec_elem_num; j += vec_elem_num) {
213 - cvt_vec_t elems_fp32(input + i * hidden_size + j);
214 - azp_adj_load_vec_t azp_adj_vec(azp_with_adj + j);
215 - cvt_vec_t azp_adj_fp32(azp_adj_vec);
216 -
217 - if constexpr (PerChannel) {
218 - b_scale_vec = cvt_vec_t(b_scale + j);
219 - scale_vec = b_scale_vec * a_scale_vec;
220 - }
221 -
222 - elems_fp32 = elems_fp32 - scale_vec * azp_adj_fp32;
223 -
224 - load_vec_t elems_out(elems_fp32);
225 - elems_out.save(output + i * hidden_size + j);
226 - }
227 -
228 - cvt_vec_t elems_fp32(input + i * hidden_size + j);
229 - azp_adj_load_vec_t azp_adj_vec(azp_with_adj + j);
230 - cvt_vec_t azp_adj_fp32(azp_adj_vec);
231 -
232 - if constexpr (PerChannel) {
233 - b_scale_vec = cvt_vec_t(b_scale + j);
234 - scale_vec = b_scale_vec * a_scale_vec;
235 - }
236 -
237 - elems_fp32 = elems_fp32 - scale_vec * azp_adj_fp32;
238 -
239 - load_vec_t elems_out(elems_fp32);
240 - elems_out.save(output + i * hidden_size + j, hidden_size - j);
241 - }
242 -}
243 -
244 -template <bool AZP, bool PerChannel, bool Bias, typename scalar_t>
245 -void dynamic_quant_epilogue(const float* input, scalar_t* output,
246 - const float* a_scale, const float* b_scale,
247 - const int32_t* azp, const int32_t* azp_adj,
248 - const scalar_t* bias, const int num_tokens,
249 - const int hidden_size) {
250 - CPU_KERNEL_GUARD_IN(dynamic_quant_epilogue)
251 - using load_vec_t = typename KernelVecType<scalar_t>::load_vec_type;
252 - using azp_adj_load_vec_t =
253 - typename KernelVecType<scalar_t>::azp_adj_load_vec_type;
254 - using cvt_vec_t = typename KernelVecType<scalar_t>::cvt_vec_type;
255 - constexpr int vec_elem_num = load_vec_t::VEC_ELEM_NUM;
256 -
257 - #pragma omp parallel for
258 - for (int i = 0; i < num_tokens; ++i) {
259 - int j = 0;
260 - cvt_vec_t token_scale_vec(a_scale[i]);
261 - cvt_vec_t token_zp_scale_vec;
262 - if constexpr (AZP) {
263 - float zp_scale_val = a_scale[i] * static_cast<float>(azp[i]);
264 - if constexpr (!PerChannel) {
265 - zp_scale_val *= *b_scale;
266 - }
267 - token_zp_scale_vec = cvt_vec_t(zp_scale_val);
268 - }
269 -
270 - for (; j < hidden_size - vec_elem_num; j += vec_elem_num) {
271 - cvt_vec_t elems_fp32(input + i * hidden_size + j);
272 - elems_fp32 = elems_fp32 * token_scale_vec;
273 -
274 - if constexpr (AZP) {
275 - azp_adj_load_vec_t azp_adj_vec(azp_adj + j);
276 - cvt_vec_t azp_adj_fp32(azp_adj_vec);
277 - azp_adj_fp32 = azp_adj_fp32 * token_zp_scale_vec;
278 -
279 - if constexpr (PerChannel) {
280 - cvt_vec_t b_scale_vec(b_scale + j);
281 - azp_adj_fp32 = azp_adj_fp32 * b_scale_vec;
282 - }
283 -
284 - elems_fp32 = elems_fp32 - azp_adj_fp32;
285 - }
286 -
287 - if constexpr (Bias) {
288 - load_vec_t bias_vec(bias + j);
289 - cvt_vec_t bias_vec_fp32(bias_vec);
290 - elems_fp32 = elems_fp32 + bias_vec_fp32;
291 - }
292 -
293 - load_vec_t elems_out(elems_fp32);
294 - elems_out.save(output + i * hidden_size + j);
295 - }
296 -
297 - cvt_vec_t elems_fp32(input + i * hidden_size + j);
298 - elems_fp32 = elems_fp32 * token_scale_vec;
299 -
300 - if constexpr (AZP) {
301 - azp_adj_load_vec_t azp_adj_vec(azp_adj + j);
302 - cvt_vec_t azp_adj_fp32(azp_adj_vec);
303 - azp_adj_fp32 = azp_adj_fp32 * token_zp_scale_vec;
304 -
305 - if constexpr (PerChannel) {
306 - cvt_vec_t b_scale_vec(b_scale + j);
307 - azp_adj_fp32 = azp_adj_fp32 * b_scale_vec;
308 - }
309 -
310 - elems_fp32 = elems_fp32 - azp_adj_fp32;
311 - }
312 -
313 - if constexpr (Bias) {
314 - load_vec_t bias_vec(bias + j);
315 - cvt_vec_t bias_vec_fp32(bias_vec);
316 - elems_fp32 = elems_fp32 + bias_vec_fp32;
317 - }
318 -
319 - load_vec_t elems_out(elems_fp32);
320 - elems_out.save(output + i * hidden_size + j, hidden_size - j);
321 - }
322 -}
323 -#elif defined(__powerpc64__)
324 -template <bool AZP, typename scalar_t>
325 -void static_scaled_int8_quant_impl(const scalar_t* input, int8_t* output,
326 - const float* scale, const int32_t* azp,
327 - const int num_tokens,
328 - const int hidden_size) {
329 - using load_vec_t = typename KernelVecType<scalar_t>::load_vec_type;
330 - using cvt_vec_t = typename KernelVecType<scalar_t>::cvt_vec_type;
331 - constexpr int vec_elem_num = load_vec_t::VEC_ELEM_NUM;
332 -
333 - constexpr float i8_min =
334 - static_cast<float>(std::numeric_limits<int8_t>::min());
335 - constexpr float i8_max =
336 - static_cast<float>(std::numeric_limits<int8_t>::max());
337 -
338 - const cvt_vec_t inv_scale(1.0 / *scale);
339 - const cvt_vec_t i8_min_vec(i8_min);
340 - const cvt_vec_t i8_max_vec(i8_max);
341 -
342 - cvt_vec_t zp_vec;
343 - if constexpr (AZP) {
344 - zp_vec = cvt_vec_t(static_cast<float>(*azp));
345 - }
346 - #pragma omp parallel for
347 - for (int i = 0; i < num_tokens; ++i) {
348 - int j = 0;
349 - for (; j < hidden_size - vec_elem_num; j += vec_elem_num) {
350 - load_vec_t elems(input + i * hidden_size + j);
351 - cvt_vec_t elems_fp32(elems);
352 - elems_fp32 = elems_fp32 * inv_scale;
353 - if constexpr (AZP) {
354 - elems_fp32 = elems_fp32 + zp_vec;
355 - }
356 - elems_fp32 = elems_fp32.clamp(i8_min_vec, i8_max_vec);
357 - vec_op::INT8Vec16 elems_int8(elems_fp32);
358 - elems_int8.save(output + i * hidden_size + j);
359 - }
360 - load_vec_t elems(input + i * hidden_size + j);
361 - cvt_vec_t elems_fp32(elems);
362 - elems_fp32 = elems_fp32 * inv_scale;
363 -
364 - if constexpr (AZP) {
365 - elems_fp32 = elems_fp32 + zp_vec;
366 - }
367 -
368 - elems_fp32 = elems_fp32.clamp(i8_min_vec, i8_max_vec);
369 - vec_op::INT8Vec16 elems_int8(elems_fp32);
370 - elems_int8.save(output + i * hidden_size + j, hidden_size - j);
371 - }
372 -}
373 -template <bool AZP, typename scalar_t>
374 -void dynamic_scaled_int8_quant_impl(const scalar_t* input, int8_t* output,
375 - float* scale, int32_t* azp,
376 - const int num_tokens,
377 - const int hidden_size) {
378 - using load_vec_t = typename KernelVecType<scalar_t>::load_vec_type;
379 - using cvt_vec_t = typename KernelVecType<scalar_t>::cvt_vec_type;
380 - constexpr int vec_elem_num = load_vec_t::VEC_ELEM_NUM;
381 -
382 - constexpr float i8_min =
383 - static_cast<float>(std::numeric_limits<int8_t>::min());
384 - constexpr float i8_max =
385 - static_cast<float>(std::numeric_limits<int8_t>::max());
386 - const cvt_vec_t i8_min_vec(i8_min);
387 - const cvt_vec_t i8_max_vec(i8_max);
388 -
389 - #pragma omp parallel for
390 - for (int i = 0; i < num_tokens; ++i) {
391 - cvt_vec_t max_value(std::numeric_limits<float>::lowest());
392 - cvt_vec_t min_value(std::numeric_limits<float>::max());
393 - {
394 - int j = 0;
395 - for (; j < hidden_size - vec_elem_num; j += vec_elem_num) {
396 - load_vec_t elems(input + i * hidden_size + j);
397 - cvt_vec_t elems_fp32(elems);
398 - if constexpr (AZP) {
399 - max_value = max_value.max(elems_fp32);
400 - min_value = min_value.min(elems_fp32);
401 - } else {
402 - max_value = max_value.max(elems_fp32.abs());
403 - }
404 - }
405 -
406 - load_vec_t elems(input + i * hidden_size + j);
407 - cvt_vec_t elems_fp32(elems);
408 -
409 - if (j + vec_elem_num == hidden_size) {
410 - if constexpr (AZP) {
411 - max_value = max_value.max(elems_fp32);
412 - min_value = min_value.min(elems_fp32);
413 - } else {
414 - max_value = max_value.max(elems_fp32.abs());
415 - }
416 - } else {
417 - if constexpr (AZP) {
418 - max_value = max_value.max(elems_fp32, hidden_size - j);
419 - min_value = min_value.min(elems_fp32, hidden_size - j);
420 - } else {
421 - max_value = max_value.max(elems_fp32.abs(), hidden_size - j);
422 - }
423 - }
424 - }
425 -
426 - float scale_val, azp_val;
427 - if constexpr (AZP) {
428 - float max_scalar = max_value.reduce_max();
429 - float min_scalar = min_value.reduce_min();
430 - scale_val = (max_scalar - min_scalar) / 255.0f;
431 - azp_val = std::nearbyint(-128.0f - min_scalar / scale_val);
432 - azp[i] = static_cast<int32_t>(azp_val);
433 - scale[i] = scale_val;
434 - } else {
435 - scale_val = max_value.reduce_max() / 127.0f;
436 - scale[i] = scale_val;
437 - }
438 -
439 - const cvt_vec_t inv_scale(1.0 / scale_val);
440 - const cvt_vec_t azp_vec(azp_val);
441 -
442 - {
443 - int j = 0;
444 - for (; j < hidden_size - vec_elem_num; j += vec_elem_num) {
445 - load_vec_t elems(input + i * hidden_size + j);
446 - cvt_vec_t elems_fp32(elems);
447 - elems_fp32 = (elems_fp32 * inv_scale);
448 -
449 - if constexpr (AZP) {
450 - elems_fp32 = elems_fp32 + azp_vec;
451 - }
452 - elems_fp32 = elems_fp32.clamp(i8_min_vec, i8_max_vec);
453 - vec_op::INT8Vec16 elems_int8(elems_fp32);
454 - elems_int8.save(output + i * hidden_size + j);
455 - }
456 -
457 - load_vec_t elems(input + i * hidden_size + j);
458 - cvt_vec_t elems_fp32(elems);
459 - elems_fp32 = (elems_fp32 * inv_scale);
460 -
461 - if constexpr (AZP) {
462 - elems_fp32 = elems_fp32 + azp_vec;
463 - }
464 - elems_fp32 = elems_fp32.clamp(i8_min_vec, i8_max_vec);
465 - vec_op::INT8Vec16 elems_int8(elems_fp32);
466 - elems_int8.save(output + i * hidden_size + j, hidden_size - j);
467 - }
468 - }
469 -}
470 -template <bool PerChannel, typename scalar_t>
471 -void static_quant_epilogue(const float* input, scalar_t* output,
472 - const float a_scale, const float* b_scale,
473 - const int32_t* azp_with_adj, const int num_tokens,
474 - const int hidden_size) {
475 - CPU_KERNEL_GUARD_IN(dynamic_output_scale_impl)
476 - using load_vec_t = typename KernelVecType<scalar_t>::load_vec_type;
477 - using azp_adj_load_vec_t =
478 - typename KernelVecType<scalar_t>::azp_adj_load_vec_type;
479 - using cvt_vec_t = typename KernelVecType<scalar_t>::cvt_vec_type;
480 - constexpr int vec_elem_num = load_vec_t::VEC_ELEM_NUM;
481 -
482 - #pragma omp parallel for
483 - for (int i = 0; i < num_tokens; ++i) {
484 - cvt_vec_t a_scale_vec(a_scale);
485 - cvt_vec_t b_scale_vec(*b_scale);
486 - cvt_vec_t scale_vec = a_scale_vec * b_scale_vec;
487 -
488 - int j = 0;
489 - for (; j < hidden_size - vec_elem_num; j += vec_elem_num) {
490 - cvt_vec_t elems_fp32(input + i * hidden_size + j);
491 - azp_adj_load_vec_t azp_adj_vec(azp_with_adj + j);
492 - cvt_vec_t azp_adj_fp32(azp_adj_vec);
493 -
494 - if constexpr (PerChannel) {
495 - b_scale_vec = cvt_vec_t(b_scale + j);
496 - scale_vec = b_scale_vec * a_scale_vec;
497 - }
498 - elems_fp32 = elems_fp32 - scale_vec * azp_adj_fp32;
499 - load_vec_t elems_out(elems_fp32);
500 - elems_out.save(output + i * hidden_size + j);
501 - }
502 -
503 - cvt_vec_t elems_fp32(input + i * hidden_size + j);
504 - azp_adj_load_vec_t azp_adj_vec(azp_with_adj + j);
505 - cvt_vec_t azp_adj_fp32(azp_adj_vec);
506 -
507 - if constexpr (PerChannel) {
508 - b_scale_vec = cvt_vec_t(b_scale + j);
509 - scale_vec = b_scale_vec * a_scale_vec;
510 - }
511 -
512 - elems_fp32 = elems_fp32 - scale_vec * azp_adj_fp32;
513 -
514 - load_vec_t elems_out(elems_fp32);
515 - elems_out.save(output + i * hidden_size + j, hidden_size - j);
516 - }
517 -}
518 -template <bool AZP, bool PerChannel, bool Bias, typename scalar_t>
519 -void dynamic_quant_epilogue(const float* input, scalar_t* output,
520 - const float* a_scale, const float* b_scale,
521 - const int32_t* azp, const int32_t* azp_adj,
522 - const scalar_t* bias, const int num_tokens,
523 - const int hidden_size) {
524 - CPU_KERNEL_GUARD_IN(dynamic_quant_epilogue)
525 - using load_vec_t = typename KernelVecType<scalar_t>::load_vec_type;
526 - using azp_adj_load_vec_t =
527 - typename KernelVecType<scalar_t>::azp_adj_load_vec_type;
528 - using cvt_vec_t = typename KernelVecType<scalar_t>::cvt_vec_type;
529 - constexpr int vec_elem_num = load_vec_t::VEC_ELEM_NUM;
530 -
531 - #pragma omp parallel for
532 - for (int i = 0; i < num_tokens; ++i) {
533 - int j = 0;
534 - cvt_vec_t token_scale_vec(a_scale[i]);
535 - cvt_vec_t token_zp_scale_vec;
536 - if constexpr (AZP) {
537 - float zp_scale_val = a_scale[i] * static_cast<float>(azp[i]);
538 - if constexpr (!PerChannel) {
539 - zp_scale_val *= *b_scale;
540 - }
541 - token_zp_scale_vec = cvt_vec_t(zp_scale_val);
542 - }
543 -
544 - for (; j < hidden_size - vec_elem_num; j += vec_elem_num) {
545 - cvt_vec_t elems_fp32(input + i * hidden_size + j);
546 - elems_fp32 = elems_fp32 * token_scale_vec;
547 -
548 - if constexpr (AZP) {
549 - azp_adj_load_vec_t azp_adj_vec(azp_adj + j);
550 - cvt_vec_t azp_adj_fp32(azp_adj_vec);
551 - azp_adj_fp32 = azp_adj_fp32 * token_zp_scale_vec;
552 -
553 - if constexpr (PerChannel) {
554 - cvt_vec_t b_scale_vec(b_scale + j);
555 - azp_adj_fp32 = azp_adj_fp32 * b_scale_vec;
556 - }
557 -
558 - elems_fp32 = elems_fp32 - azp_adj_fp32;
559 - }
560 -
561 - if constexpr (Bias) {
562 - load_vec_t bias_vec(bias + j);
563 - cvt_vec_t bias_vec_fp32(bias_vec);
564 - elems_fp32 = elems_fp32 + bias_vec_fp32;
565 - }
566 -
567 - load_vec_t elems_out(elems_fp32);
568 - elems_out.save(output + i * hidden_size + j);
569 - }
570 -
571 - cvt_vec_t elems_fp32(input + i * hidden_size + j);
572 - elems_fp32 = elems_fp32 * token_scale_vec;
573 -
574 - if constexpr (AZP) {
575 - azp_adj_load_vec_t azp_adj_vec(azp_adj + j);
576 - cvt_vec_t azp_adj_fp32(azp_adj_vec);
577 - azp_adj_fp32 = azp_adj_fp32 * token_zp_scale_vec;
578 -
579 - if constexpr (PerChannel) {
580 - cvt_vec_t b_scale_vec(b_scale + j);
581 - azp_adj_fp32 = azp_adj_fp32 * b_scale_vec;
582 - }
583 -
584 - elems_fp32 = elems_fp32 - azp_adj_fp32;
585 - }
586 -
587 - if constexpr (Bias) {
588 - load_vec_t bias_vec(bias + j);
589 - cvt_vec_t bias_vec_fp32(bias_vec);
590 - elems_fp32 = elems_fp32 + bias_vec_fp32;
591 - }
592 -
593 - load_vec_t elems_out(elems_fp32);
594 - elems_out.save(output + i * hidden_size + j, hidden_size - j);
595 - }
596 -}
597 -#else
598 -template <typename scalar_t>
599 -void static_scaled_int8_quant_impl(const scalar_t* input, int8_t* output,
600 - const float* scale, const int32_t* azp,
601 - const int num_tokens,
602 - const int hidden_size) {
603 - TORCH_CHECK(false,
604 - "static_scaled_int8_quant_impl requires AVX512/powerpc64/AArch64 "
605 - "support.")
606 -}
607 -
608 -template <typename scalar_t>
609 -void dynamic_scaled_int8_quant_impl(const scalar_t* input, int8_t* output,
610 - float* scale, int32_t* azp,
611 - const int num_tokens,
612 - const int hidden_size) {
613 - TORCH_CHECK(false,
614 - "dynamic_scaled_int8_quant_impl requires "
615 - "AVX512/powerpc64/AArch64 support.")
616 -}
617 -
618 -template <bool PerChannel, typename scalar_t>
619 -void static_quant_epilogue(const float* input, scalar_t* output,
620 - const float a_scale, const float* b_scale,
621 - const int32_t* azp_with_adj, const int num_tokens,
622 - const int hidden_size) {
623 - TORCH_CHECK(
624 - false, "static_quant_epilogue requires AVX512/powerpc64/AArch64 support.")
625 -}
626 -
627 -template <typename scalar_t>
628 -void dynamic_quant_epilogue(const float* input, scalar_t* output,
629 - const float* a_scale, const float* b_scale,
630 - const int32_t* azp, const int32_t* azp_with_adj,
631 - const scalar_t* bias, const int num_tokens,
632 - const int hidden_size) {
633 - TORCH_CHECK(
634 - false,
635 - "dynamic_quant_epilogue requires AVX512/powerpc64/AArch64 support.")
636 -}
637 -#endif
638 -} // namespace
639 -
640 -void int8_scaled_mm(torch::Tensor& c, // [M, OC], row-major
641 - const torch::Tensor& a, // [M, IC], row-major
642 - const torch::Tensor& b, // [IC, OC], column-major
643 - const torch::Tensor& a_scales, // [1] or [M]
644 - const torch::Tensor& b_scales, // [1] or [OC]
645 - const std::optional<torch::Tensor>& bias // [OC]
646 -) {
647 - CPU_KERNEL_GUARD_IN(cutlass_scaled_mm)
648 - // Checks for conformality
649 - TORCH_CHECK(a.dtype() == torch::kInt8 && b.dtype() == torch::kInt8,
650 - "int8_scaled_mm only supports INT8 inputs.")
651 - TORCH_CHECK(a.dim() == 2 && b.dim() == 2 && c.dim() == 2);
652 - TORCH_CHECK(c.size(0) == a.size(0) && a.size(1) == b.size(0) &&
653 - b.size(1) == c.size(1));
654 - TORCH_CHECK(a_scales.numel() == 1 || a_scales.numel() == a.size(0));
655 - TORCH_CHECK(b_scales.numel() == 1 || b_scales.numel() == b.size(1));
656 -
657 - // Check for strides and alignment
658 - TORCH_CHECK(a.stride(1) == 1 && c.stride(1) == 1); // Row-major
659 - TORCH_CHECK(b.stride(0) == 1); // Column-major
660 - TORCH_CHECK(c.stride(0) % 16 == 0 &&
661 - b.stride(1) % 16 == 0); // 16 Byte Alignment
662 - TORCH_CHECK(a_scales.is_contiguous() && b_scales.is_contiguous());
663 -
664 - if (bias) {
665 - TORCH_CHECK(bias->numel() == b.size(1) && bias->is_contiguous() &&
666 - bias->dim() == 1);
667 - }
668 -
669 - VLLM_DISPATCH_FLOATING_TYPES(c.scalar_type(), "int8_scaled_mm", [&] {
670 - if (a_scales.numel() != 1) {
671 - // per-token
672 - // Note: oneDNN doesn't support per-token activation quantization
673 - // Ideally we want to fuse the GEMM and the scale procedure with oneDNN
674 - // JIT, the intermediate data is cached in registers or L1. But for now
675 - // the oneDNN GEMM code generation only supports two quantization
676 - // patterns: per-tensor or per-output-channel of weight.
677 - // So we have to apply the per-token scale with a 'epilogue'. In C=s_a *
678 - // s_b * (A@B) + bias, the C_inter = s_b * (A@B) is computed by oneDNN
679 - // GEMM, then the per-token scale (and bias) is applied with the epilogue
680 - // C=s_a * C_inter + bias.
681 - torch::Tensor tmp_fp32_out =
682 - torch::empty_like(c, ::at::ScalarType::Float);
683 - // Compute C_inter=s_b * (A@B)
684 - DNNLPrimitiveHelper<true>::gemm_s8s8_jit<float, void>(
685 - a.data_ptr<int8_t>(), b.data_ptr<int8_t>(),
686 - tmp_fp32_out.data_ptr<float>(), nullptr, a.size(0), b.size(1),
687 - a.size(1), nullptr, b_scales.data_ptr<float>(), 0, b_scales.numel());
688 - if (bias.has_value()) {
689 - // Compute C=s_a * C_inter + bias
690 - dynamic_quant_epilogue<false, true, true>(
691 - tmp_fp32_out.data_ptr<float>(), c.data_ptr<scalar_t>(),
692 - a_scales.data_ptr<float>(), nullptr, nullptr, nullptr,
693 - bias->data_ptr<scalar_t>(), c.size(0), c.size(1));
694 - } else {
695 - // Compute C=s_a * C_inter
696 - dynamic_quant_epilogue<false, true, false, scalar_t>(
697 - tmp_fp32_out.data_ptr<float>(), c.data_ptr<scalar_t>(),
698 - a_scales.data_ptr<float>(), nullptr, nullptr, nullptr, nullptr,
699 - c.size(0), c.size(1));
700 - }
701 - } else {
702 - // per-tensor
703 - if (bias.has_value()) {
704 - // Compute C=s_a * s_b * (A@B) + bias
705 - DNNLPrimitiveHelper<false>::gemm_s8s8_jit(
706 - a.data_ptr<int8_t>(), b.data_ptr<int8_t>(), c.data_ptr<scalar_t>(),
707 - bias->data_ptr<scalar_t>(), a.size(0), b.size(1), a.size(1),
708 - a_scales.data_ptr<float>(), b_scales.data_ptr<float>(),
709 - a_scales.numel(), b_scales.numel());
710 - } else {
711 - // Compute C=s_a * s_b * (A@B)
712 - DNNLPrimitiveHelper<false>::gemm_s8s8_jit<scalar_t, void>(
713 - a.data_ptr<int8_t>(), b.data_ptr<int8_t>(), c.data_ptr<scalar_t>(),
714 - nullptr, a.size(0), b.size(1), a.size(1),
715 - a_scales.data_ptr<float>(), b_scales.data_ptr<float>(),
716 - a_scales.numel(), b_scales.numel());
717 - }
718 - }
719 - });
720 -}
721 -
722 -void int8_scaled_mm_azp(torch::Tensor& c, // [M, OC], row-major
723 - const torch::Tensor& a, // [M, IC], row-major
724 - const torch::Tensor& b, // [IC, OC], column-major
725 - const torch::Tensor& a_scales, // [1] or [M]
726 - const torch::Tensor& b_scales, // [1] or [OC]
727 - const torch::Tensor& azp_adj, // [OC]
728 - const std::optional<torch::Tensor>& azp, // [1] or [M]
729 - const std::optional<torch::Tensor>& bias // [OC]
730 -) {
731 - CPU_KERNEL_GUARD_IN(cutlass_scaled_mm_azp)
732 - // Checks for conformality
733 - TORCH_CHECK(a.dtype() == torch::kInt8 && b.dtype() == torch::kInt8,
734 - "int8_scaled_mm_azp only supports INT8 inputs.")
735 - TORCH_CHECK(a.dim() == 2 && b.dim() == 2 && c.dim() == 2);
736 - TORCH_CHECK(c.size(0) == a.size(0) && a.size(1) == b.size(0) &&
737 - b.size(1) == c.size(1));
738 - TORCH_CHECK(a_scales.numel() == 1 || a_scales.numel() == a.size(0));
739 - TORCH_CHECK(b_scales.numel() == 1 || b_scales.numel() == b.size(1));
740 -
741 - // Check for strides and alignment
742 - TORCH_CHECK(a.stride(1) == 1 && c.stride(1) == 1); // Row-major
743 - TORCH_CHECK(b.stride(0) == 1); // Column-major
744 - TORCH_CHECK(c.stride(0) % 16 == 0 &&
745 - b.stride(1) % 16 == 0); // 16 Byte Alignment
746 - TORCH_CHECK(a_scales.is_contiguous() && b_scales.is_contiguous());
747 -
748 - if (bias) {
749 - TORCH_CHECK(bias->numel() == b.size(1) && bias->is_contiguous());
750 - }
751 - if (azp) {
752 - TORCH_CHECK(azp->numel() == a.size(0) && azp->is_contiguous());
753 - }
754 - TORCH_CHECK(azp_adj.numel() == b.size(1) && azp_adj.is_contiguous());
755 -
756 - // azp & bias types
757 - TORCH_CHECK(azp_adj.dtype() == torch::kInt32);
758 - TORCH_CHECK(!azp || azp->dtype() == torch::kInt32);
759 - TORCH_CHECK(!bias || bias->dtype() == c.dtype(),
760 - "currently bias dtype must match output dtype ", c.dtype());
761 -
762 - VLLM_DISPATCH_FLOATING_TYPES(c.scalar_type(), "int8_scaled_mm_azp", [&] {
763 - torch::Tensor tmp_fp32_out = torch::empty_like(c, ::at::ScalarType::Float);
764 - if (a_scales.numel() != 1) {
765 - // per-token
766 - // Note: oneDNN doesn't support per-token activation quantization
767 - // Compute C_inter=s_b * (A@B)
768 - DNNLPrimitiveHelper<true>::gemm_s8s8_jit<float, void>(
769 - a.data_ptr<int8_t>(), b.data_ptr<int8_t>(),
770 - tmp_fp32_out.data_ptr<float>(), nullptr, a.size(0), b.size(1),
771 - a.size(1), nullptr, b_scales.data_ptr<float>(), 0, b_scales.numel());
772 - if (bias.has_value()) {
773 - // Compute C=s_a * C_inter - s_a * s_b * azp * azp_adj + bias
774 - if (b_scales.numel() != 1) {
775 - // Per-Channel
776 - dynamic_quant_epilogue<true, true, true>(
777 - tmp_fp32_out.data_ptr<float>(), c.data_ptr<scalar_t>(),
778 - a_scales.data_ptr<float>(), b_scales.data_ptr<float>(),
779 - azp->data_ptr<int32_t>(), azp_adj.data_ptr<int32_t>(),
780 - bias->data_ptr<scalar_t>(), c.size(0), c.size(1));
781 - } else {
782 - // Per-Tensor
783 - dynamic_quant_epilogue<true, false, true>(
784 - tmp_fp32_out.data_ptr<float>(), c.data_ptr<scalar_t>(),
785 - a_scales.data_ptr<float>(), b_scales.data_ptr<float>(),
786 - azp->data_ptr<int32_t>(), azp_adj.data_ptr<int32_t>(),
787 - bias->data_ptr<scalar_t>(), c.size(0), c.size(1));
788 - }
789 - } else {
790 - // Compute C=s_a * C_inter - s_a * s_b * azp * azp_adj
791 - if (b_scales.numel() != 1) {
792 - // Per-Channel
793 - dynamic_quant_epilogue<true, true, false, scalar_t>(
794 - tmp_fp32_out.data_ptr<float>(), c.data_ptr<scalar_t>(),
795 - a_scales.data_ptr<float>(), b_scales.data_ptr<float>(),
796 - azp->data_ptr<int32_t>(), azp_adj.data_ptr<int32_t>(), nullptr,
797 - c.size(0), c.size(1));
798 - } else {
799 - // Per-Tensor
800 - dynamic_quant_epilogue<true, false, false, scalar_t>(
801 - tmp_fp32_out.data_ptr<float>(), c.data_ptr<scalar_t>(),
802 - a_scales.data_ptr<float>(), b_scales.data_ptr<float>(),
803 - azp->data_ptr<int32_t>(), azp_adj.data_ptr<int32_t>(), nullptr,
804 - c.size(0), c.size(1));
805 - }
806 - }
807 - } else {
808 - // per-tensor
809 - if (bias.has_value()) {
810 - // Compute C_inter=s_a * s_b * (A@B) + bias
811 - DNNLPrimitiveHelper<false>::gemm_s8s8_jit(
812 - a.data_ptr<int8_t>(), b.data_ptr<int8_t>(),
813 - tmp_fp32_out.data_ptr<float>(), bias->data_ptr<scalar_t>(),
814 - a.size(0), b.size(1), a.size(1), a_scales.data_ptr<float>(),
815 - b_scales.data_ptr<float>(), a_scales.numel(), b_scales.numel());
816 - } else {
817 - // Compute C_inter=s_a * s_b * (A@B)
818 - DNNLPrimitiveHelper<false>::gemm_s8s8_jit<float, void>(
819 - a.data_ptr<int8_t>(), b.data_ptr<int8_t>(),
820 - tmp_fp32_out.data_ptr<float>(), nullptr, a.size(0), b.size(1),
821 - a.size(1), a_scales.data_ptr<float>(), b_scales.data_ptr<float>(),
822 - a_scales.numel(), b_scales.numel());
823 - }
824 -
825 - // Compute C=C_inter - s_a * s_b * azp_adj
826 - if (b_scales.numel() != 1) {
827 - // Per-Channel
828 - static_quant_epilogue<true>(
829 - tmp_fp32_out.data_ptr<float>(), c.data_ptr<scalar_t>(),
830 - *a_scales.data_ptr<float>(), b_scales.data_ptr<float>(),
831 - azp_adj.data_ptr<int32_t>(), a.size(0), b.size(1));
832 - } else {
833 - // Per-Tensor
834 - static_quant_epilogue<false>(
835 - tmp_fp32_out.data_ptr<float>(), c.data_ptr<scalar_t>(),
836 - *a_scales.data_ptr<float>(), b_scales.data_ptr<float>(),
837 - azp_adj.data_ptr<int32_t>(), a.size(0), b.size(1));
838 - }
839 - }
840 - });
841 -}
842 -
843 -// static-per-tensor quantization.
844 -void static_scaled_int8_quant(torch::Tensor& out, // [..., hidden_size]
845 - const torch::Tensor& input, // [..., hidden_size]
846 - const torch::Tensor& scale,
847 - std::optional<torch::Tensor> const& azp) {
848 - CPU_KERNEL_GUARD_IN(static_scaled_int8_quant)
849 - TORCH_CHECK(input.is_contiguous());
850 - TORCH_CHECK(out.is_contiguous());
851 - TORCH_CHECK(scale.numel() == 1);
852 - TORCH_CHECK(!azp.has_value() || azp->numel() == 1);
853 -
854 - const int hidden_size = input.size(-1);
855 - const int num_tokens = input.numel() / hidden_size;
856 - VLLM_DISPATCH_FLOATING_TYPES(
857 - input.scalar_type(), "static_scaled_int8_quant_impl", [&] {
858 - if (azp.has_value()) {
859 - static_scaled_int8_quant_impl<true>(
860 - input.data_ptr<scalar_t>(), out.data_ptr<int8_t>(),
861 - scale.data_ptr<float>(), azp->data_ptr<int32_t>(), num_tokens,
862 - hidden_size);
863 - } else {
864 - static_scaled_int8_quant_impl<false>(
865 - input.data_ptr<scalar_t>(), out.data_ptr<int8_t>(),
866 - scale.data_ptr<float>(), nullptr, num_tokens, hidden_size);
867 - }
868 - });
869 -}
870 -
871 -// dynamic-per-token quantization.
872 -void dynamic_scaled_int8_quant(
873 - torch::Tensor& out, // [..., hidden_size]
874 - const torch::Tensor& input, // [..., hidden_size]
875 - torch::Tensor& scale, // [..., 1]
876 - std::optional<torch::Tensor> const& azp) {
877 - CPU_KERNEL_GUARD_IN(dynamic_scaled_int8_quant)
878 - TORCH_CHECK(input.is_contiguous());
879 - TORCH_CHECK(out.is_contiguous());
880 -
881 - int const hidden_size = input.size(-1);
882 - int const num_tokens = input.numel() / hidden_size;
883 - VLLM_DISPATCH_FLOATING_TYPES(
884 - input.scalar_type(), "dynamic_scaled_int8_quant_impl", [&] {
885 - if (azp.has_value()) {
886 - dynamic_scaled_int8_quant_impl<true>(
887 - input.data_ptr<scalar_t>(), out.data_ptr<int8_t>(),
888 - scale.data_ptr<float>(), azp->data_ptr<int32_t>(), num_tokens,
889 - hidden_size);
890 - } else {
891 - dynamic_scaled_int8_quant_impl<false>(
892 - input.data_ptr<scalar_t>(), out.data_ptr<int8_t>(),
893 - scale.data_ptr<float>(), nullptr, num_tokens, hidden_size);
894 - }
895 - });
896 -}
897 -
898 -#if defined(__powerpc64__)
899 -void int8_scaled_mm_ppc64le(torch::Tensor& c, // [M, OC], row-major
900 - const torch::Tensor& a, // [M, IC], row-major
901 - const torch::Tensor& b, // [IC, OC], column-major
902 - const torch::Tensor& a_scales,
903 - const torch::Tensor& b_scales,
904 - const std::optional<torch::Tensor>& bias // [OC]
905 -) {
906 - CPU_KERNEL_GUARD_IN(cutlass_scaled_mm)
907 - // Checks for conformality
908 - TORCH_CHECK(a.dtype() == torch::kInt8 && b.dtype() == torch::kInt8,
909 - "int8_scaled_mm_ppc64le only supports INT8 inputs.");
910 - TORCH_CHECK(a.dim() == 2 && b.dim() == 2 && c.dim() == 2);
911 - TORCH_CHECK(c.size(0) == a.size(0) && a.size(1) == b.size(0) &&
912 - b.size(1) == c.size(1));
913 - // We dont need this
914 - TORCH_CHECK(a_scales.numel() == 1 || a_scales.numel() == a.size(0));
915 - TORCH_CHECK(b_scales.numel() == 1 || b_scales.numel() == b.size(1));
916 -
917 - // Check for strides and alignment
918 - TORCH_CHECK(a.stride(1) == 1 && c.stride(1) == 1); // Row-major
919 - TORCH_CHECK(b.stride(0) == 1); // Column-major
920 - TORCH_CHECK(c.stride(0) % 16 == 0 &&
921 - b.stride(1) % 16 == 0); // 16 Byte Alignment
922 - TORCH_CHECK(a_scales.is_contiguous() && b_scales.is_contiguous());
923 -
924 - if (bias) {
925 - TORCH_CHECK(bias->numel() == b.size(1) && bias->is_contiguous() &&
926 - bias->dim() == 1);
927 - }
928 - VLLM_DISPATCH_FLOATING_TYPES(c.scalar_type(), "int8_scaled_mm_ppc64le", [&] {
929 - torch::Tensor tmp_fp32_out = torch::empty_like(c, ::at::ScalarType::Float);
930 - // Compute C_inter=s_b * (A@B)
931 - DNNLPrimitiveHelper<true>::gemm_s8s8_jit<float, void>(
932 - a.data_ptr<int8_t>(), b.data_ptr<int8_t>(),
933 - tmp_fp32_out.data_ptr<float>(), nullptr, a.size(0), b.size(1),
934 - a.size(1), nullptr, b_scales.data_ptr<float>(), 0, b_scales.numel());
935 - if (bias.has_value()) {
936 - // Compute C=s_a * C_inter + bias
937 - dynamic_quant_epilogue<false, true, true>(
938 - tmp_fp32_out.data_ptr<float>(), c.data_ptr<scalar_t>(),
939 - a_scales.data_ptr<float>(), nullptr, nullptr, nullptr,
940 - bias->data_ptr<scalar_t>(), c.size(0), c.size(1));
941 - } else {
942 - // Compute C=s_a * C_inter
943 - dynamic_quant_epilogue<false, true, false, scalar_t>(
944 - tmp_fp32_out.data_ptr<float>(), c.data_ptr<scalar_t>(),
945 - a_scales.data_ptr<float>(), nullptr, nullptr, nullptr, nullptr,
946 - c.size(0), c.size(1));
947 - }
948 - });
949 -}
950 -
951 -#endif
csrc/cpu/torch_bindings.cpp
+32 -64
@@ -6,25 +6,20 @@
6
7 std::string init_cpu_threads_env(const std::string& cpu_ids);
8
9 -void int8_scaled_mm(torch::Tensor& c, const torch::Tensor& a,
10 - const torch::Tensor& b, const torch::Tensor& a_scales,
11 - const torch::Tensor& b_scales,
12 - const std::optional<torch::Tensor>& bias);
13 -
14 -void int8_scaled_mm_azp(torch::Tensor& c, const torch::Tensor& a,
15 - const torch::Tensor& b, const torch::Tensor& a_scales,
16 - const torch::Tensor& b_scales,
17 - const torch::Tensor& azp_adj,
18 - const std::optional<torch::Tensor>& azp,
19 - const std::optional<torch::Tensor>& bias);
20 -
21 -#if defined(__powerpc64__)
22 -void int8_scaled_mm_ppc64le(torch::Tensor& c, const torch::Tensor& a,
23 - const torch::Tensor& b,
24 - const torch::Tensor& a_scales,
25 - const torch::Tensor& b_scales,
26 - const std::optional<torch::Tensor>& bias);
27 -#endif
9 +void release_dnnl_matmul_handler(int64_t handler);
10 +
11 +int64_t create_onednn_scaled_mm_handler(const torch::Tensor& b,
12 + const torch::Tensor& b_scales,
13 + at::ScalarType output_type,
14 + bool dynamic_act_quant, bool use_azp,
15 + int64_t primitive_cache_size);
16 +
17 +void onednn_scaled_mm(torch::Tensor& c, const torch::Tensor& a,
18 + const torch::Tensor& a_scales,
19 + const std::optional<torch::Tensor>& azp,
20 + const std::optional<torch::Tensor>& azp_adj,
21 + const std::optional<torch::Tensor>& bias,
22 + int64_t handler);
23
24 void mla_decode_kvcache(torch::Tensor& out, torch::Tensor& query,
25 torch::Tensor& kv_cache, double scale,
@@ -151,8 +146,25 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
146 ops.impl("rotary_embedding", torch::kCPU, &rotary_embedding);
147
148 // Quantization
154 -#if defined(__AVX512F__) || (defined(__aarch64__) && !defined(__APPLE__))
149 +#if defined(__AVX512F__) || (defined(__aarch64__) && !defined(__APPLE__)) || \
150 + defined(__powerpc64__)
151 at::Tag stride_tag = at::Tag::needs_fixed_stride_order;
152 + // Helper function to release oneDNN handlers
153 + ops.def("release_dnnl_matmul_handler(int handler) -> ()",
154 + &release_dnnl_matmul_handler);
155 +
156 + // Create oneDNN W8A8 handler
157 + ops.def(
158 + "create_onednn_scaled_mm_handler(Tensor b, Tensor b_scales, ScalarType "
159 + "output_type, bool dynamic_act_quant, bool use_azp, int "
160 + "primitive_cache_size) -> int",
161 + &create_onednn_scaled_mm_handler);
162 +
163 + // oneDNN scaled_mm for W8A8 with static per-tensor activation quantization
164 + ops.def(
165 + "onednn_scaled_mm(Tensor! c, Tensor a, Tensor a_scales, Tensor? azp, "
166 + "Tensor? azp_adj, Tensor? bias, int handler) -> ()");
167 + ops.impl("onednn_scaled_mm", torch::kCPU, &onednn_scaled_mm);
168
169 // Compute int8 quantized tensor for given scaling factor.
170 ops.def(
@@ -168,50 +180,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
180 {stride_tag});
181 ops.impl("dynamic_scaled_int8_quant", torch::kCPU,
182 &dynamic_scaled_int8_quant);
171 - // W8A8 GEMM, supporting symmetric per-tensor or per-row/column
172 - // quantization.
173 - ops.def(
174 - "cutlass_scaled_mm(Tensor! out, Tensor a,"
175 - " Tensor b, Tensor a_scales,"
176 - " Tensor b_scales, Tensor? bias) -> ()",
177 - {stride_tag});
178 - ops.impl("cutlass_scaled_mm", torch::kCPU, &int8_scaled_mm);
179 - // w8a8 GEMM, supporting asymmetric per-tensor or per-row/column
180 - // quantization.
181 - ops.def(
182 - "cutlass_scaled_mm_azp(Tensor! out, Tensor a,"
183 - " Tensor b, Tensor a_scales,"
184 - " Tensor b_scales, Tensor azp_adj,"
185 - " Tensor? azp, Tensor? bias) -> ()",
186 - {stride_tag});
187 - ops.impl("cutlass_scaled_mm_azp", torch::kCPU, &int8_scaled_mm_azp);
188 -#elif defined(__powerpc64__)
189 - // Compute int8 quantized tensor for given scaling factor.
190 - ops.def(
191 - "static_scaled_int8_quant(Tensor! out, Tensor input, Tensor scale,"
192 - "Tensor? azp) -> ()");
193 - ops.impl("static_scaled_int8_quant", torch::kCPU, &static_scaled_int8_quant);
194 -
195 - // Compute int8 quantized tensor and scaling factor
196 - ops.def(
197 - "dynamic_scaled_int8_quant(Tensor! out, Tensor input, Tensor! scale, "
198 - "Tensor!? azp) -> ()");
199 - ops.impl("dynamic_scaled_int8_quant", torch::kCPU,
200 - &dynamic_scaled_int8_quant);
201 - // W8A8 GEMM, supporting symmetric quantization.
202 - ops.def(
203 - "cutlass_scaled_mm(Tensor! out, Tensor a,"
204 - " Tensor b, Tensor a_scales,"
205 - " Tensor b_scales, Tensor? bias) -> ()");
206 - ops.impl("cutlass_scaled_mm", torch::kCPU, &int8_scaled_mm_ppc64le);
207 - // w8a8 GEMM, supporting asymmetric per-tensor or per-row/column
208 - // quantization.
209 - ops.def(
210 - "cutlass_scaled_mm_azp(Tensor! out, Tensor a,"
211 - " Tensor b, Tensor a_scales,"
212 - " Tensor b_scales, Tensor azp_adj,"
213 - " Tensor? azp, Tensor? bias) -> ()");
214 - ops.impl("cutlass_scaled_mm_azp", torch::kCPU, &int8_scaled_mm_azp);
183 #endif
184
185 // SHM CCL
tests/kernels/test_onednn.py new
+144
@@ -0,0 +1,144 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Integration tests for FlexAttention backend vs default backend"""
4 +
5 +from typing import Optional
6 +
7 +import pytest
8 +import torch
9 +
10 +from tests.kernels.utils import to_int8
11 +from vllm import _custom_ops as ops
12 +from vllm.platforms import current_platform
13 +
14 +if not current_platform.is_cpu():
15 + pytest.skip("skipping CPU-only tests", allow_module_level=True)
16 +
17 +NK_FACTORS = [
18 + (256, 128),
19 + (4096, 4096),
20 + (16384, 4096),
21 + (1023, 491),
22 + (1001, 15),
23 +]
24 +M_FACTORS = [
25 + (16, 1, 32, 128, 64),
26 + (1, 17, 1, 31, 17),
27 +]
28 +CACHE_SIZES = [2]
29 +DTYPE = [torch.bfloat16]
30 +
31 +
32 +def rand_int8(shape: tuple, device: str = "cpu"):
33 + return to_int8(torch.rand(shape, device=device) * 255 - 128)
34 +
35 +
36 +def ref_int8_scaled_mm(
37 + a: torch.Tensor,
38 + b: torch.Tensor,
39 + scale_a: torch.Tensor,
40 + scale_b: torch.Tensor,
41 + azp: Optional[torch.Tensor],
42 + bias: Optional[torch.Tensor],
43 + output_type: torch.dtype,
44 +):
45 + if azp is not None:
46 + a = a.to(dtype=torch.float32) - azp.to(dtype=torch.float32)
47 + output = torch.mm((scale_a * a.to(dtype=torch.float32)),
48 + (scale_b * b.to(dtype=torch.float32)))
49 + if bias is not None:
50 + output += bias.float()
51 +
52 + return output.to(dtype=output_type)
53 +
54 +
55 +def onednn_int8_gemm_test_helper(primitive_cache_size: int,
56 + m: int,
57 + n: int,
58 + k: int,
59 + per_tensor_a_quant: bool,
60 + per_tensor_b_quant: bool,
61 + use_azp: bool,
62 + use_bias: bool,
63 + out_dtype: torch.dtype = torch.bfloat16,
64 + device: str = "cpu"):
65 + # Test for a oneDNN kernel with per-tensor / per-token activation
66 + # quantization and per-tensor / per-output channel weight quantization.
67 + a = to_int8(torch.randn((m, k), device=device) * 5)
68 + b = to_int8(torch.randn((n, k), device=device).t() * 5)
69 +
70 + a_scales_shape = (1, 1) if per_tensor_a_quant else (m, 1)
71 + b_scales_shape = (1, 1) if per_tensor_b_quant else (1, n)
72 +
73 + scale_a = (torch.randn(a_scales_shape, device=device, dtype=torch.float32))
74 + scale_b = (torch.randn(b_scales_shape, device=device, dtype=torch.float32))
75 +
76 + if use_azp:
77 + azp = torch.rand(a_scales_shape, dtype=torch.float32) * 10 + 1.5
78 + azp = (azp / scale_a).round().to(dtype=torch.int32)
79 + azp_adj = scale_b * b.sum(dim=0, keepdim=True, dtype=torch.float32)
80 + else:
81 + azp = None
82 + azp_adj = None
83 +
84 + if use_bias:
85 + bias = torch.rand((n, ), device=device, dtype=out_dtype) * 10
86 + else:
87 + bias = None
88 +
89 + handler = ops.create_onednn_scaled_mm(
90 + b,
91 + scale_b,
92 + out_dtype,
93 + not per_tensor_a_quant,
94 + use_azp,
95 + primitive_cache_size,
96 + )
97 +
98 + out = torch.zeros((m, n), dtype=out_dtype)
99 + ops.onednn_scaled_mm(handler, a, out, scale_a, azp, azp_adj, bias)
100 + baseline = ref_int8_scaled_mm(a, b, scale_a, scale_b, azp, bias, out_dtype)
101 +
102 + torch.testing.assert_close(out, baseline, rtol=1e-1, atol=1e0)
103 +
104 + if use_bias:
105 + # To test runtime bias setting
106 + out = torch.zeros((m, n), dtype=out_dtype)
107 + ops.onednn_scaled_mm(handler, a, out, scale_a, azp, azp_adj, None)
108 + baseline = ref_int8_scaled_mm(a, b, scale_a, scale_b, azp, None,
109 + out_dtype)
110 +
111 + torch.testing.assert_close(out, baseline, rtol=1e-1, atol=1e0)
112 +
113 +
114 +@pytest.mark.parametrize("n,k", NK_FACTORS)
115 +@pytest.mark.parametrize("m_list", M_FACTORS)
116 +@pytest.mark.parametrize("per_tensor_a_scale", [True, False])
117 +@pytest.mark.parametrize("per_tensor_b_scale", [True, False])
118 +@pytest.mark.parametrize("use_bias", [True, False])
119 +@pytest.mark.parametrize("use_azp", [True, False])
120 +@pytest.mark.parametrize("output_type", DTYPE)
121 +@pytest.mark.parametrize("primitive_cache_size", CACHE_SIZES)
122 +def test_onednn_int8_scaled_gemm(
123 + n: int,
124 + k: int,
125 + m_list: tuple[int],
126 + per_tensor_a_scale: bool,
127 + per_tensor_b_scale: bool,
128 + use_bias: bool,
129 + use_azp: bool,
130 + output_type: torch.dtype,
131 + primitive_cache_size: int,
132 +):
133 + for m in m_list:
134 + onednn_int8_gemm_test_helper(
135 + primitive_cache_size=primitive_cache_size,
136 + m=m,
137 + n=n,
138 + k=k,
139 + per_tensor_a_quant=per_tensor_a_scale,
140 + per_tensor_b_quant=per_tensor_b_scale,
141 + use_bias=use_bias,
142 + use_azp=use_azp,
143 + out_dtype=output_type,
144 + )
vllm/_custom_ops.py
+83
@@ -1827,3 +1827,86 @@ if hasattr(torch.ops._C, "int8_scaled_mm_with_quant"):
1827 M = mat1.size(0)
1828 N = mat2.size(0)
1829 return torch.empty((M, N), dtype=out_dtype)
1830 +
1831 +
1832 +class CPUDNNLGEMMHandler:
1833 +
1834 + def __init__(self) -> None:
1835 + self.handler: Optional[int] = None
1836 + self.n = -1
1837 + self.k = -1
1838 +
1839 + def __del__(self):
1840 + if self.handler is not None:
1841 + torch.ops._C.release_dnnl_matmul_handler(self.handler)
1842 +
1843 +
1844 +def create_onednn_scaled_mm(
1845 + weight: torch.Tensor, # [K, N]
1846 + weight_scales: torch.Tensor,
1847 + output_type: torch.dtype,
1848 + dynamic_quant: bool,
1849 + use_azp: bool,
1850 + primitive_cache_size: int = 128,
1851 +) -> CPUDNNLGEMMHandler:
1852 + handler = CPUDNNLGEMMHandler()
1853 + handler.k, handler.n = weight.size()
1854 + handler.handler = torch.ops._C.create_onednn_scaled_mm_handler(
1855 + weight, weight_scales, output_type, dynamic_quant, use_azp,
1856 + primitive_cache_size)
1857 + return handler
1858 +
1859 +
1860 +def onednn_scaled_int8_quant(input: torch.Tensor,
1861 + scale: Optional[torch.Tensor] = None,
1862 + azp: Optional[torch.Tensor] = None,
1863 + symmetric: bool = True):
1864 + """
1865 + Quantize the input tensor to int8 and return the quantized tensor and scale, and maybe azp.
1866 +
1867 + Args:
1868 + input: The input tensor to be quantized to int8.
1869 + scale: Optional scaling factor for the int8 quantization.
1870 + When not provided, we invoke dynamic-per-token quantization.
1871 + azp: Optional zero-point for the int8 quantization.
1872 + Must be provided for asymmetric quantization if `scale` is provided.
1873 + symmetric: Whether to use symmetric quantization (scale only, azp ignored).
1874 +
1875 + Returns:
1876 + tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]] : Output int8 tensor, scales, and optionally azp.
1877 + """
1878 + output = torch.empty_like(input, dtype=torch.int8)
1879 + token_num = input.numel() // input.shape[-1]
1880 + input = input.view((token_num, input.shape[-1]))
1881 + if scale is not None:
1882 + # static-per-tensor quantization.
1883 + assert symmetric == (
1884 + azp
1885 + is None), "azp must only be provided for asymmetric quantization."
1886 + torch.ops._C.static_scaled_int8_quant(output, input, scale, azp)
1887 + return output, scale, azp
1888 +
1889 + # dynamic-per-token quantization.
1890 + input_scales = torch.empty((token_num, 1),
1891 + device=input.device,
1892 + dtype=torch.float32)
1893 + input_azp = None if symmetric else torch.empty_like(input_scales,
1894 + dtype=torch.int32)
1895 + torch.ops._C.dynamic_scaled_int8_quant(output, input, input_scales,
1896 + input_azp)
1897 + return output, input_scales, input_azp
1898 +
1899 +
1900 +def onednn_scaled_mm(
1901 + dnnl_handler: CPUDNNLGEMMHandler,
1902 + x: torch.Tensor,
1903 + output: torch.Tensor,
1904 + input_scale: Optional[torch.Tensor],
1905 + input_zp: Optional[torch.Tensor],
1906 + input_zp_adj: Optional[torch.Tensor],
1907 + bias: Optional[torch.Tensor],
1908 +) -> torch.Tensor:
1909 + torch.ops._C.onednn_scaled_mm(output, x, input_scale, input_zp,
1910 + input_zp_adj, bias, dnnl_handler.handler)
1911 +
1912 + return output
vllm/model_executor/layers/fused_moe/layer.py
+8 -3
@@ -360,10 +360,15 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
360 elif current_platform.is_cpu():
361 if current_platform.get_cpu_architecture() == CpuArchEnum.X86:
362 from vllm.model_executor.layers.fused_moe import cpu_fused_moe
363 - dtype = layer.w13_weight.dtype
363 + from vllm.model_executor.layers.utils import (
364 + check_cpu_sgl_kernel)
365 + dtype_w13 = layer.w13_weight.dtype
366 + _, n_w13, k_w13 = layer.w13_weight.size()
367 + dtype_w2 = layer.w2_weight.dtype
368 + _, n_w2, k_w2 = layer.w2_weight.size()
369 if (envs.VLLM_CPU_SGL_KERNEL
365 - and torch._C._cpu._is_amx_tile_supported()
366 - and dtype == torch.bfloat16):
370 + and check_cpu_sgl_kernel(n_w13, k_w13, dtype_w13)
371 + and check_cpu_sgl_kernel(n_w2, k_w2, dtype_w2)):
372 packed_w13_weight = torch.ops._C.convert_weight_packed(
373 layer.w13_weight)
374 assert packed_w13_weight.size() == layer.w13_weight.size()
vllm/model_executor/layers/linear.py
+4 -4
@@ -199,11 +199,10 @@ class UnquantizedLinearMethod(LinearMethodBase):
199
200 def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
201 if current_platform.is_cpu() and envs.VLLM_CPU_SGL_KERNEL:
202 + from vllm.model_executor.layers.utils import check_cpu_sgl_kernel
203 N, K = layer.weight.size()
204 dtype = layer.weight.dtype
204 - if (torch._C._cpu._is_amx_tile_supported()
205 - and dtype == torch.bfloat16 and N % 32 == 0
206 - and K % 32 == 0):
205 + if check_cpu_sgl_kernel(N, K, dtype):
206 packed_weight = torch.ops._C.convert_weight_packed(
207 layer.weight)
208 assert packed_weight.size() == layer.weight.size()
@@ -215,7 +214,8 @@ class UnquantizedLinearMethod(LinearMethodBase):
214 else:
215 logger.warning(
216 "CPU SGL kernels require Intel AMX support,"
218 - " bfloat16 weight, IC and OC are divisible by 32.")
217 + " bf16/fp16/int8 weight, IC and OC are divisible by "
218 + "32 and 16.")
219 layer.use_cpu_sgl = False
220
221 def apply(self,
vllm/model_executor/layers/quantization/kernels/scaled_mm/__init__.py
+3 -1
@@ -6,6 +6,8 @@ from typing import Optional
6
7 from vllm.model_executor.layers.quantization.kernels.scaled_mm.aiter import (
8 AiterScaledMMLinearKernel)
9 +from vllm.model_executor.layers.quantization.kernels.scaled_mm.cpu import (
10 + CPUScaledMMLinearKernel)
11 from vllm.model_executor.layers.quantization.kernels.scaled_mm.cutlass import (
12 CutlassScaledMMLinearKernel)
13 from vllm.model_executor.layers.quantization.kernels.scaled_mm.ScaledMMLinearKernel import ( # noqa: E501
@@ -18,7 +20,7 @@ from vllm.platforms import PlatformEnum, current_platform
20
21 # in priority/performance order (when available)
22 _POSSIBLE_KERNELS: dict[PlatformEnum, list[type[ScaledMMLinearKernel]]] = {
21 - PlatformEnum.CPU: [CutlassScaledMMLinearKernel],
23 + PlatformEnum.CPU: [CPUScaledMMLinearKernel],
24 PlatformEnum.CUDA: [CutlassScaledMMLinearKernel],
25 PlatformEnum.ROCM: [AiterScaledMMLinearKernel, TritonScaledMMLinearKernel],
26 PlatformEnum.TPU: [XLAScaledMMLinearKernel],
vllm/model_executor/layers/quantization/kernels/scaled_mm/cpu.py new
+206
@@ -0,0 +1,206 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from typing import Optional
5 +
6 +import torch
7 +
8 +from vllm import _custom_ops as ops
9 +from vllm import envs
10 +from vllm.model_executor.layers.quantization.utils import replace_parameter
11 +from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
12 + convert_to_channelwise)
13 +from vllm.model_executor.layers.utils import check_cpu_sgl_kernel
14 +from vllm.platforms import current_platform
15 +from vllm.platforms.interface import CpuArchEnum
16 +
17 +from .ScaledMMLinearKernel import (ScaledMMLinearKernel,
18 + ScaledMMLinearLayerConfig)
19 +
20 +
21 +class CPUScaledMMLinearKernel(ScaledMMLinearKernel):
22 +
23 + @classmethod
24 + def get_min_capability(cls) -> int:
25 + return 75
26 +
27 + @classmethod
28 + def can_implement(
29 + cls, c: ScaledMMLinearLayerConfig) -> tuple[bool, Optional[str]]:
30 + if not current_platform.is_cpu():
31 + return False, "CPUScaledMM requires running on CPU."
32 +
33 + return True, None
34 +
35 + def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
36 + weight = getattr(layer, self.w_q_name)
37 + dtype = weight.dtype
38 + N, K = weight.size()
39 + if (current_platform.get_cpu_architecture() == CpuArchEnum.X86
40 + and envs.VLLM_CPU_SGL_KERNEL and self.config.input_symmetric
41 + and check_cpu_sgl_kernel(N, K, dtype)):
42 + self.linear_method = self._apply_weights_sgl
43 + self.process_weights_for_sgl(layer)
44 + else:
45 + self.linear_method = self._apply_weights_onednn
46 + self.process_weights_for_onednn(layer)
47 +
48 + def process_weights_for_onednn(self, layer: torch.nn.Module) -> None:
49 + # WEIGHT
50 + # Transpose to [K, N] for convenience
51 + weight = getattr(layer, self.w_q_name)
52 + replace_parameter(
53 + layer, self.w_q_name,
54 + torch.nn.Parameter(weight.t().data, requires_grad=False))
55 +
56 + # WEIGHT SCALE
57 + # oneDNN kernels support only per-tensor and per-channel.
58 + # If we have a fused module (QKV, MLP) with per tensor scales (thus N
59 + # scales being passed to the kernel), convert to the per-channel case.
60 + is_fused_module = len(layer.logical_widths) > 1
61 + weight_scale = getattr(layer, self.w_s_name)
62 + if is_fused_module and not self.config.is_channelwise:
63 + weight_scale = convert_to_channelwise(weight_scale,
64 + layer.logical_widths)
65 + replace_parameter(
66 + layer, self.w_s_name,
67 + torch.nn.Parameter(weight_scale.data, requires_grad=False))
68 +
69 + # INPUT SCALE
70 + if self.config.is_static_input_scheme:
71 + input_scale = getattr(layer, self.i_s_name)
72 +
73 + if self.config.input_symmetric:
74 + replace_parameter(
75 + layer, self.i_s_name,
76 + torch.nn.Parameter(input_scale.max(), requires_grad=False))
77 + setattr(layer, self.i_zp_name, None)
78 + else:
79 + input_zero_point = getattr(layer, self.i_zp_name)
80 +
81 + # reconstruct the ranges
82 + int8_traits = torch.iinfo(torch.int8)
83 + azps = input_zero_point.to(dtype=torch.int32)
84 + range_max = (input_scale * (int8_traits.max - azps)).max()
85 + range_min = (input_scale * (int8_traits.min - azps)).min()
86 +
87 + scale = (range_max - range_min) / (int8_traits.max -
88 + int8_traits.min)
89 + replace_parameter(
90 + layer, self.i_s_name,
91 + torch.nn.Parameter(scale, requires_grad=False))
92 +
93 + azp = (int8_traits.min -
94 + range_min / scale).round().to(dtype=torch.int32)
95 + replace_parameter(layer, self.i_zp_name,
96 + torch.nn.Parameter(azp, requires_grad=False))
97 +
98 + else:
99 + setattr(layer, self.i_s_name, None)
100 + setattr(layer, self.i_zp_name, None)
101 +
102 + # Different from cutlass, oneDNN kernels only need the AZP adjustment
103 + # term for dynamic quantization. And s_b should be folded into the
104 + # term. Such as:
105 + # s_a * s_b * [(A - zp_a)B] + bias =
106 + # s_a * (s_b * AB) - s_a * s_b * zp_a * B + bias =
107 + # s_a * GEMM_output - s_a * zp_a * adj + bias
108 + if not (self.config.input_symmetric
109 + and self.config.is_static_input_scheme):
110 + weight = getattr(layer, self.w_q_name)
111 + weight_scale = getattr(layer, self.w_s_name)
112 + azp_adj = weight.sum(dim=0, keepdim=True, dtype=torch.float32)
113 + azp_adj = azp_adj * weight_scale.squeeze()
114 + setattr(layer, self.azp_adj_name,
115 + torch.nn.Parameter(azp_adj, requires_grad=False))
116 + else:
117 + setattr(layer, self.azp_adj_name, None)
118 +
119 + weight = getattr(layer, self.w_q_name)
120 + self.dnnl_handler = ops.create_onednn_scaled_mm(
121 + weight,
122 + getattr(layer, self.w_s_name),
123 + torch.get_default_dtype(),
124 + getattr(layer, self.i_s_name) is None,
125 + not self.config.input_symmetric,
126 + 32,
127 + )
128 + # weight is prepacked and maintained by the dnnl_handler,
129 + # release the original weight
130 + setattr(layer, self.w_q_name, None)
131 + del weight
132 +
133 + def process_weights_for_sgl(self, layer: torch.nn.Module) -> None:
134 + # WEIGHT
135 + weight = getattr(layer, self.w_q_name)
136 + packed_weight = torch.ops._C.convert_weight_packed(weight)
137 + replace_parameter(
138 + layer, self.w_q_name,
139 + torch.nn.Parameter(packed_weight, requires_grad=False))
140 +
141 + if layer.bias is not None:
142 + bias = layer.bias
143 + layer.register_parameter(
144 + "bias_fp32",
145 + torch.nn.Parameter(bias.float().data, requires_grad=False))
146 +
147 + # WEIGHT SCALE
148 + # CPU SGL kernels only support per-channel.
149 + # For per-tensor quant, convert to the per-channel case.
150 + weight_scale = getattr(layer, self.w_s_name)
151 + if not self.config.is_channelwise:
152 + weight_scale = convert_to_channelwise(weight_scale,
153 + layer.logical_widths)
154 + replace_parameter(
155 + layer, self.w_s_name,
156 + torch.nn.Parameter(weight_scale.data, requires_grad=False))
157 +
158 + setattr(layer, self.i_s_name, None)
159 + setattr(layer, self.i_zp_name, None)
160 + setattr(layer, self.azp_adj_name, None)
161 +
162 + def apply_weights(self,
163 + layer: torch.nn.Module,
164 + x: torch.Tensor,
165 + bias: Optional[torch.Tensor] = None) -> torch.Tensor:
166 + return self.linear_method(
167 + layer,
168 + x,
169 + bias,
170 + )
171 +
172 + def _apply_weights_onednn(
173 + self,
174 + layer: torch.nn.Module,
175 + x: torch.Tensor,
176 + bias: Optional[torch.Tensor] = None) -> torch.Tensor:
177 + w_q, w_s, i_s, i_zp, azp_adj = self._get_weight_params(layer)
178 +
179 + # ops.scaled_int8_quant supports both dynamic and static quant:
180 + # * dynamic, i_s is None and x_s computed from x.
181 + # * static, i_s is scalar and x_s is i_s.
182 + x_q, x_s, x_zp = ops.onednn_scaled_int8_quant(
183 + x, i_s, i_zp, self.config.input_symmetric)
184 +
185 + m = x.size(0)
186 + n = self.dnnl_handler.n
187 + out = torch.empty((m, n), dtype=x.dtype)
188 + ops.onednn_scaled_mm(self.dnnl_handler, x_q, out, x_s, x_zp, azp_adj,
189 + bias)
190 +
191 + return out
192 +
193 + def _apply_weights_sgl(
194 + self,
195 + layer: torch.nn.Module,
196 + x: torch.Tensor,
197 + bias: Optional[torch.Tensor] = None) -> torch.Tensor:
198 + w_q, w_s, _, _, _ = self._get_weight_params(layer)
199 + return torch.ops._C.int8_scaled_mm_with_quant(
200 + x,
201 + w_q,
202 + w_s,
203 + layer.bias_fp32 if bias is not None else None,
204 + x.dtype,
205 + True,
206 + )
vllm/model_executor/layers/quantization/kernels/scaled_mm/cutlass.py
+2 -2
@@ -25,8 +25,8 @@ class CutlassScaledMMLinearKernel(ScaledMMLinearKernel):
25 def can_implement(
26 cls, c: ScaledMMLinearLayerConfig) -> tuple[bool, Optional[str]]:
27
28 - if (not current_platform.is_cuda() and not current_platform.is_cpu()):
29 - return False, "CutlassScaledMM requires running on CUDA or CPU."
28 + if not current_platform.is_cuda():
29 + return False, "CutlassScaledMM requires running on CUDA."
30
31 return True, None
32
vllm/model_executor/layers/utils.py
+6
@@ -142,6 +142,12 @@ direct_register_custom_op(
142 )
143
144
145 +def check_cpu_sgl_kernel(n: int, k: int, dtype: torch.dtype):
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 cpu_unquantized_gemm(layer: torch.nn.Module,
152 x: torch.Tensor,
153 weight: torch.Tensor,