Skip to content

Commit d1fcd34

Browse files
GregoryComerGregoryComer
committed
Support bf16/fp16 activations in CPU SDPA
1 parent 39c0df6 commit d1fcd34

8 files changed

Lines changed: 361 additions & 59 deletions

File tree

‎extension/llm/custom_ops/custom_ops.py‎

Lines changed: 12 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -75,13 +75,16 @@ def _validate_params(
7575
value.dim() == 4
7676
), f"Expected value to be 4 dimensional but got {value.dim()} dimensions."
7777

78+
supported_dtypes = (torch.float32, torch.bfloat16, torch.float16)
7879
assert (
79-
query.dtype == torch.float32
80-
), f"Expected query to be float32 but got {query.dtype}"
81-
assert key.dtype == torch.float32, f"Expected key to be float32 but got {key.dtype}"
80+
query.dtype in supported_dtypes
81+
), f"Expected query to be float32, bfloat16, or float16 but got {query.dtype}"
8282
assert (
83-
value.dtype == torch.float32
84-
), f"Expected value to be float32 but got {value.dtype}"
83+
key.dtype == query.dtype
84+
), f"Expected key to have the same dtype as query but got {key.dtype}"
85+
assert (
86+
value.dtype == query.dtype
87+
), f"Expected value to have the same dtype as query but got {value.dtype}"
8588

8689
assert (
8790
key_cache.dim() == 4
@@ -91,11 +94,11 @@ def _validate_params(
9194
), f"Expected value_cache to be 4 dimensional but got {value_cache.dim()}"
9295

9396
assert (
94-
key_cache.dtype == torch.float32
95-
), f"Expected key_cache to be float32 but got {key_cache.dtype}"
97+
key_cache.dtype == query.dtype
98+
), f"Expected key_cache to have the same dtype as query but got {key_cache.dtype}"
9699
assert (
97-
value_cache.dtype == torch.float32
98-
), f"Expected value_cache to be float32 but got {value_cache.dtype}"
100+
value_cache.dtype == query.dtype
101+
), f"Expected value_cache to have the same dtype as query but got {value_cache.dtype}"
99102

100103
assert (
101104
key_cache.size() == value_cache.size()

‎extension/llm/custom_ops/op_sdpa.cpp‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -45,8 +45,10 @@ bool validate_flash_attention_args(
4545

4646
ET_CHECK_OR_RETURN_FALSE(
4747
(query.scalar_type() == ScalarType::Float) ||
48+
(query.scalar_type() == ScalarType::Half) ||
49+
(query.scalar_type() == ScalarType::BFloat16) ||
4850
(query.scalar_type() == ScalarType::Char),
49-
"Query must be Float type");
51+
"Query must be Float, Half, BFloat16, or Char type");
5052

5153
ET_CHECK_OR_RETURN_FALSE(
5254
(query.scalar_type() == key.scalar_type()) &&
@@ -266,7 +268,7 @@ Tensor& flash_attention_kernel_out(
266268

267269
auto seq_len = query.size(2);
268270

269-
ET_SWITCH_FLOAT_TYPES(
271+
ET_SWITCH_FLOATHBF16_TYPES(
270272
query.scalar_type(), ctx, "flash_attention", CTYPE, [&] {
271273
// TODO we need to re-evaluate this for ARM CPUs
272274
// And there can be many so instead of templatizing
@@ -414,7 +416,7 @@ Tensor& custom_sdpa_out_impl(
414416

415417
// TODO(task): replace the template param selection logic
416418
// with whatever apprpriately makes more sense for
417-
ET_SWITCH_FLOAT_TYPES(
419+
ET_SWITCH_FLOATHBF16_TYPES(
418420
output.scalar_type(), ctx, "flash_attention", CTYPE, [&] {
419421
// TODO we need to re-evaluate this for ARM CPUs
420422
// And there can be many so instead of templatizing

‎extension/llm/custom_ops/op_sdpa_impl.h‎

Lines changed: 113 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -74,8 +74,10 @@ void _q_at_k_gemm(
7474
accum_t* qk_data) {
7575
ET_CHECK_MSG(q_data.dtype == k_data.dtype, "q and k must have same dtype");
7676
ET_CHECK_MSG(
77-
q_data.dtype == ScalarType::Char || q_data.dtype == ScalarType::Float,
78-
"q and k must be either int8 or float");
77+
q_data.dtype == ScalarType::Char || q_data.dtype == ScalarType::Float ||
78+
q_data.dtype == ScalarType::Half ||
79+
q_data.dtype == ScalarType::BFloat16,
80+
"q and k must be int8, float, half, or bfloat16");
7981
if (q_data.dtype == ScalarType::Char) {
8082
if constexpr (std::is_same<accum_t, float>::value) {
8183
int a_stride_m_tmp, b_stride_n_tmp;
@@ -103,6 +105,35 @@ void _q_at_k_gemm(
103105
ET_CHECK_MSG(
104106
false, "Accumulation in dtype other than float not supported yet");
105107
}
108+
} else if (
109+
q_data.dtype == ScalarType::BFloat16 ||
110+
q_data.dtype == ScalarType::Half) {
111+
if constexpr (std::is_same<accum_t, float>::value) {
112+
auto do_gemm = [&](auto rt_tag) {
113+
using rt = decltype(rt_tag);
114+
::executorch::cpublas::gemm(
115+
::executorch::cpublas::TransposeType::Transpose,
116+
::executorch::cpublas::TransposeType::NoTranspose,
117+
k_n,
118+
q_m,
119+
qk_k,
120+
1.0f,
121+
static_cast<const rt*>(k_data.data),
122+
k_stride_n,
123+
static_cast<const rt*>(q_data.data),
124+
q_stride_m,
125+
0.0f,
126+
qk_data,
127+
k_n);
128+
};
129+
if (q_data.dtype == ScalarType::BFloat16) {
130+
do_gemm(::executorch::aten::BFloat16{});
131+
} else {
132+
do_gemm(::executorch::aten::Half{});
133+
}
134+
} else {
135+
ET_CHECK_MSG(false, "Reduced-precision q@k requires float accumulation");
136+
}
106137
} else {
107138
::executorch::cpublas::gemm(
108139
::executorch::cpublas::TransposeType::Transpose,
@@ -251,7 +282,7 @@ void _qk_at_v_gemm(
251282
const int64_t m,
252283
const int64_t n,
253284
const int64_t k,
254-
const accum_t* qk_data,
285+
const void* qk_data,
255286
const int64_t qk_stride_m,
256287
const MaybeQuantizedMatrixData& v_data,
257288
const int64_t v_stride_n,
@@ -261,14 +292,15 @@ void _qk_at_v_gemm(
261292
accum_t* buf_qdq_ptr) {
262293
if (v_data.dtype == ScalarType::Char) {
263294
if constexpr (std::is_same<accum_t, float>::value) {
295+
const float* qk = static_cast<const float*>(qk_data);
264296
if (m > 4) {
265297
// For larger batch sizes, dequantize and use BLAS for better
266298
// performance
267299
dequant_and_gemm(
268300
m,
269301
n,
270302
k,
271-
const_cast<float*>(qk_data),
303+
const_cast<float*>(qk),
272304
qk_stride_m,
273305
v_data,
274306
v_stride_n,
@@ -286,7 +318,7 @@ void _qk_at_v_gemm(
286318
m,
287319
n,
288320
k,
289-
qk_data,
321+
qk,
290322
qk_stride_m /*lhs_stride_m*/,
291323
static_cast<const int8_t*>(v_data.data),
292324
v_stride_n /*rhs_stride_n*/,
@@ -301,6 +333,37 @@ void _qk_at_v_gemm(
301333
ET_CHECK_MSG(
302334
false, "Accumulation in dtype other than float not supported yet");
303335
}
336+
} else if (
337+
v_data.dtype == ScalarType::BFloat16 ||
338+
v_data.dtype == ScalarType::Half) {
339+
// qk has been cast to the activation dtype (see qk_reduced_data); both
340+
// operands are reduced precision and accumulate into the float output.
341+
if constexpr (std::is_same<accum_t, float>::value) {
342+
auto do_gemm = [&](auto rt_tag) {
343+
using rt = decltype(rt_tag);
344+
::executorch::cpublas::gemm(
345+
::executorch::cpublas::TransposeType::NoTranspose,
346+
::executorch::cpublas::TransposeType::NoTranspose,
347+
n,
348+
m,
349+
k,
350+
1.0f,
351+
static_cast<const rt*>(v_data.data),
352+
v_stride_n,
353+
static_cast<const rt*>(qk_data),
354+
qk_stride_m,
355+
beta,
356+
o_data,
357+
o_stride_m);
358+
};
359+
if (v_data.dtype == ScalarType::BFloat16) {
360+
do_gemm(::executorch::aten::BFloat16{});
361+
} else {
362+
do_gemm(::executorch::aten::Half{});
363+
}
364+
} else {
365+
ET_CHECK_MSG(false, "Reduced-precision qk@v requires float accumulation");
366+
}
304367
} else {
305368
::executorch::cpublas::gemm(
306369
::executorch::cpublas::TransposeType::NoTranspose,
@@ -311,7 +374,7 @@ void _qk_at_v_gemm(
311374
static_cast<accum_t>(1),
312375
static_cast<const accum_t*>(v_data.data),
313376
v_stride_n,
314-
qk_data,
377+
static_cast<const accum_t*>(qk_data),
315378
qk_stride_m,
316379
beta,
317380
o_data,
@@ -572,11 +635,9 @@ void cpu_flash_attention(
572635
constexpr bool is_reduced_type =
573636
::executorch::runtime::is_reduced_floating_point_v<scalar_t>;
574637

575-
ET_CHECK_MSG(
576-
!is_reduced_type, "FlashAttention does not support reduced types.");
577-
// Figure out mixed precision a little later
578-
// using accum_t = at::opmath_type<scalar_t>;
579-
using accum_t = scalar_t;
638+
// Reduced-precision (bf16/fp16) activations accumulate in float: the two
639+
// matmuls run as reduced-in/float-out gemms and the softmax stays in float.
640+
using accum_t = std::conditional_t<is_reduced_type, float, scalar_t>;
580641
using Vec = vec::Vectorized<accum_t>;
581642
accum_t scaling_factor = static_cast<accum_t>(calculate_scale(query, scale));
582643

@@ -774,7 +835,19 @@ void cpu_flash_attention(
774835
} else {
775836
buf = scratch.get();
776837
}
838+
std::unique_ptr<char[]> allocated_buf_reduced;
777839
void* buf_reduced = nullptr;
840+
if (is_reduced_type) {
841+
int64_t size_reduced_bytes =
842+
qSplitSize * kvSplitSize * num_thread * sizeof(scalar_t);
843+
Result<void*> scratch_reduced = ctx.allocate_temp(size_reduced_bytes, 64);
844+
if (!scratch_reduced.ok()) {
845+
allocated_buf_reduced = std::make_unique<char[]>(size_reduced_bytes);
846+
buf_reduced = allocated_buf_reduced.get();
847+
} else {
848+
buf_reduced = scratch_reduced.get();
849+
}
850+
}
778851
int64_t size_per_thread_qdq_vec = kvSplitSize * headSize;
779852
// Lets align size_per_thread_qdq_vec to 64 bytes, for coalesced cache reads,
780853
// by padding with right number of per thread elements
@@ -1015,18 +1088,16 @@ void cpu_flash_attention(
10151088
if (tmp_max == -std::numeric_limits<accum_t>::infinity()) {
10161089
// to avoid `nan = exp2f(-inf - (-inf))`
10171090
fill_stub(
1018-
conditional_data_ptr(qk_data, qk_reduced_data) +
1019-
row * kvBlockSize,
1020-
static_cast<scalar_t>(0),
1091+
qk_data + row * kvBlockSize,
1092+
static_cast<accum_t>(0),
10211093
kvBlockSize);
10221094
} else {
10231095
// qk <- exp(qk - max) and sum per row
10241096
tmp_sum = tmp_max;
10251097
_exp_reduce_sum_fusion_kernel(
10261098
qk_data + row * kvBlockSize,
10271099
kvBlockSize,
1028-
conditional_data_ptr(qk_data, qk_reduced_data) +
1029-
row * kvBlockSize,
1100+
qk_data + row * kvBlockSize,
10301101
tmp_sum);
10311102
// exp_tmp <- exp(max[row] - max)
10321103
exp_tmp = std::exp(qk_max_data[row] - tmp_max);
@@ -1068,12 +1139,23 @@ void cpu_flash_attention(
10681139
headSize,
10691140
v_quant_params_StrideN,
10701141
value.scalar_type());
1142+
// For reduced-precision activations the attention weights are cast to
1143+
// the activation dtype so that Softmax(q @ k.T) @ v runs as a
1144+
// reduced-in/float-out gemm matching the value matrix.
1145+
const void* qk_gemm_data = qk_data;
1146+
if constexpr (is_reduced_type) {
1147+
if (!is_quantized_sdpa) {
1148+
vec::convert<accum_t, scalar_t>(
1149+
qk_data, qk_reduced_data, qBlockSize * kvBlockSize);
1150+
qk_gemm_data = qk_reduced_data;
1151+
}
1152+
}
10711153
// Calculate Softmax(q @ k.T) @ v
10721154
_qk_at_v_gemm<accum_t>(
10731155
qBlockSize,
10741156
headSize,
10751157
kvBlockSize,
1076-
qk_data,
1158+
qk_gemm_data,
10771159
kvBlockSize,
10781160
v_sub_matrix_data,
10791161
vStrideN,
@@ -1086,12 +1168,20 @@ void cpu_flash_attention(
10861168
// reorder MHA output with strides
10871169
for (int64_t row = 0; row < qBlockSize; ++row) {
10881170
accum_t sum_reciprocal = 1 / qk_sum_data[row];
1089-
vec::map<scalar_t>(
1090-
[sum_reciprocal](Vec x) { return x * Vec(sum_reciprocal); },
1091-
out_data + i * oStrideB + j * oStrideH + m * oStrideM +
1092-
row * oStrideM,
1093-
dst_data + row * headSize,
1094-
headSize);
1171+
scalar_t* out_row = out_data + i * oStrideB + j * oStrideH +
1172+
m * oStrideM + row * oStrideM;
1173+
const accum_t* dst_row = dst_data + row * headSize;
1174+
if constexpr (is_reduced_type) {
1175+
for (int64_t d = 0; d < headSize; ++d) {
1176+
out_row[d] = static_cast<scalar_t>(dst_row[d] * sum_reciprocal);
1177+
}
1178+
} else {
1179+
vec::map<scalar_t>(
1180+
[sum_reciprocal](Vec x) { return x * Vec(sum_reciprocal); },
1181+
out_row,
1182+
dst_row,
1183+
headSize);
1184+
}
10951185
}
10961186
// Move to the next query
10971187
data_index_step(i, batchSize, j, num_head, k, qSlice);

0 commit comments

Comments
 (0)