Skip to content

Commit ff2930f

Browse files
ThomasJannaudThomasJannaud
authored
Generate packed causal masks directly (#23178)
Summary: Add `cadence::sdpa_bitwise_causal_mask_gen`, which creates an LSB-first packed `uint8` causal mask from one-dimensional token positions and a static key length without materializing a dense mask. Register the schema, fake implementation, Python reference, generic C++ kernel, and build target in both source cells. Reviewed By: mvartani-meta Differential Revision: D121921652 Pull Request resolved: #23178
1 parent 21f770c commit ff2930f

5 files changed

Lines changed: 209 additions & 0 deletions

File tree

‎backends/cadence/aot/ops_registrations.py‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -526,6 +526,12 @@ def register_fake(
526526
lib.define(
527527
"sdpa_bitwise_mask_gen.out(Tensor mask, float threshold, *, Tensor(a!) out) -> Tensor (a!)"
528528
)
529+
lib.define(
530+
"sdpa_bitwise_causal_mask_gen(Tensor positions, int key_length) -> (Tensor out)"
531+
)
532+
lib.define(
533+
"sdpa_bitwise_causal_mask_gen.out(Tensor positions, int key_length, *, Tensor(a!) out) -> Tensor(a!)"
534+
)
529535

530536
# Load/store with iDMA. These only exist before memory planning.
531537
# Post memory planning, we check that outputs/inputs for the load/store are in
@@ -3350,6 +3356,21 @@ def sdpa_bitwise_mask_gen_meta(
33503356
return mask.new_empty(mask_shape, dtype=torch.uint8)
33513357

33523358

3359+
@register_fake("cadence::sdpa_bitwise_causal_mask_gen")
3360+
def sdpa_bitwise_causal_mask_gen_meta(
3361+
positions: torch.Tensor,
3362+
key_length: int,
3363+
) -> torch.Tensor:
3364+
torch._check(
3365+
positions.dtype in (torch.int32, torch.int64),
3366+
lambda: "expected int32 or int64 positions",
3367+
)
3368+
torch._check(positions.dim() == 1, lambda: "positions must be one-dimensional")
3369+
torch._check(key_length > 0, lambda: "key_length must be positive")
3370+
torch._check(key_length % 8 == 0, lambda: "key_length must be a multiple of 8")
3371+
return positions.new_empty((positions.numel(), key_length // 8), dtype=torch.uint8)
3372+
3373+
33533374
@register_fake("cadence::quantized_w8a32_linear")
33543375
def quantized_w8a32_linear_meta(
33553376
src: torch.Tensor,

‎backends/cadence/aot/ref_implementations.py‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3000,6 +3000,38 @@ def sdpa_bitwise_mask_gen(mask: torch.Tensor, threshold: float) -> torch.Tensor:
30003000
return packed.view(*original_shape[:-1], packed_last)
30013001

30023002

3003+
@impl_tracked(m, "sdpa_bitwise_causal_mask_gen")
3004+
def sdpa_bitwise_causal_mask_gen(
3005+
positions: torch.Tensor, key_length: int
3006+
) -> torch.Tensor:
3007+
"""Generate an LSB-first packed causal mask directly from token positions."""
3008+
assert positions.dtype in (
3009+
torch.int32,
3010+
torch.int64,
3011+
), "Positions must use int32 or int64"
3012+
assert positions.dim() == 1, "Positions must be one-dimensional"
3013+
assert (
3014+
key_length > 0 and key_length % 8 == 0
3015+
), "Key length must be positive and divisible by 8"
3016+
assert bool(torch.all(positions >= 0)), "Positions must be nonnegative"
3017+
3018+
visible_lengths = (
3019+
torch.clamp(positions.to(torch.int64), max=key_length - 1).unsqueeze(1) + 1
3020+
)
3021+
byte_starts = torch.arange(
3022+
0,
3023+
key_length,
3024+
8,
3025+
dtype=torch.int64,
3026+
device=positions.device,
3027+
).unsqueeze(0)
3028+
visible_bits = torch.clamp(visible_lengths - byte_starts, min=0, max=8)
3029+
return torch.bitwise_and(
3030+
torch.bitwise_left_shift(torch.full_like(visible_bits, 0xFF), visible_bits),
3031+
0xFF,
3032+
).to(torch.uint8)
3033+
3034+
30033035
@impl_tracked(m, "slice_scatter_")
30043036
def slice_scatter_impl(
30053037
self: torch.Tensor,

‎backends/cadence/generic/operators/op_sdpa_bitwise_causal_mask_gen.cpp‎

Lines changed: 115 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,115 @@
1+
/*
2+
* Copyright (c) Meta Platforms, Inc. and affiliates.
3+
* All rights reserved.
4+
*
5+
* This source code is licensed under the BSD-style license found in the
6+
* LICENSE file in the root directory of this source tree.
7+
*/
8+
9+
#include <executorch/backends/cadence/generic/operators/op_sdpa_bitwise_causal_mask_gen.h>
10+
11+
#include <cstddef>
12+
#include <cstdint>
13+
#include <cstring>
14+
15+
#include <executorch/runtime/kernel/kernel_includes.h>
16+
17+
namespace impl {
18+
namespace generic {
19+
namespace native {
20+
namespace {
21+
22+
template <typename Position>
23+
bool positions_are_valid(const Position* positions, int64_t num_positions) {
24+
for (int64_t row = 0; row < num_positions; ++row) {
25+
if (positions[row] < 0) {
26+
return false;
27+
}
28+
}
29+
return true;
30+
}
31+
32+
template <typename Position>
33+
void generate_mask(
34+
const Position* positions,
35+
int64_t num_positions,
36+
int64_t packed_width,
37+
uint8_t* out) {
38+
for (int64_t row = 0; row < num_positions; ++row) {
39+
uint8_t* row_out = out + row * packed_width;
40+
const int64_t visible = positions[row] >= packed_width * 8
41+
? packed_width * 8
42+
: static_cast<int64_t>(positions[row]) + 1;
43+
const int64_t full_visible_bytes = visible / 8;
44+
const int64_t partial_visible_bits = visible % 8;
45+
46+
std::memset(row_out, 0, static_cast<std::size_t>(full_visible_bytes));
47+
int64_t next_byte = full_visible_bytes;
48+
if (partial_visible_bits != 0) {
49+
row_out[next_byte++] =
50+
static_cast<uint8_t>(0xFFu << partial_visible_bits);
51+
}
52+
std::memset(
53+
row_out + next_byte,
54+
0xFF,
55+
static_cast<std::size_t>(packed_width - next_byte));
56+
}
57+
}
58+
59+
} // namespace
60+
61+
::executorch::aten::Tensor& sdpa_bitwise_causal_mask_gen_out(
62+
::executorch::runtime::KernelRuntimeContext& ctx,
63+
const ::executorch::aten::Tensor& positions,
64+
int64_t key_length,
65+
::executorch::aten::Tensor& out) {
66+
const auto position_sizes = positions.sizes();
67+
ET_KERNEL_CHECK(ctx, position_sizes.size() == 1, InvalidArgument, out);
68+
ET_KERNEL_CHECK(
69+
ctx, key_length > 0 && key_length % 8 == 0, InvalidArgument, out);
70+
ET_KERNEL_CHECK(
71+
ctx,
72+
positions.dtype() == ::executorch::aten::ScalarType::Int ||
73+
positions.dtype() == ::executorch::aten::ScalarType::Long,
74+
InvalidArgument,
75+
out);
76+
ET_KERNEL_CHECK(
77+
ctx,
78+
out.dtype() == ::executorch::aten::ScalarType::Byte,
79+
InvalidArgument,
80+
out);
81+
82+
const int64_t num_positions = positions.numel();
83+
const int64_t packed_width = key_length / 8;
84+
ET_KERNEL_CHECK(
85+
ctx,
86+
out.dim() == 2 && out.size(0) == num_positions &&
87+
out.size(1) == packed_width,
88+
InvalidArgument,
89+
out);
90+
91+
const bool valid_positions =
92+
positions.dtype() == ::executorch::aten::ScalarType::Long
93+
? positions_are_valid(positions.const_data_ptr<int64_t>(), num_positions)
94+
: positions_are_valid(positions.const_data_ptr<int32_t>(), num_positions);
95+
ET_KERNEL_CHECK(ctx, valid_positions, InvalidArgument, out);
96+
97+
if (positions.dtype() == ::executorch::aten::ScalarType::Long) {
98+
generate_mask(
99+
positions.const_data_ptr<int64_t>(),
100+
num_positions,
101+
packed_width,
102+
out.mutable_data_ptr<uint8_t>());
103+
} else {
104+
generate_mask(
105+
positions.const_data_ptr<int32_t>(),
106+
num_positions,
107+
packed_width,
108+
out.mutable_data_ptr<uint8_t>());
109+
}
110+
return out;
111+
}
112+
113+
} // namespace native
114+
} // namespace generic
115+
} // namespace impl

‎backends/cadence/generic/operators/op_sdpa_bitwise_causal_mask_gen.h‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,28 @@
1+
/*
2+
* Copyright (c) Meta Platforms, Inc. and affiliates.
3+
* All rights reserved.
4+
*
5+
* This source code is licensed under the BSD-style license found in the
6+
* LICENSE file in the root directory of this source tree.
7+
*/
8+
9+
#pragma once
10+
11+
#include <cstdint>
12+
13+
#include <executorch/runtime/core/exec_aten/exec_aten.h>
14+
#include <executorch/runtime/kernel/kernel_runtime_context.h>
15+
16+
namespace impl {
17+
namespace generic {
18+
namespace native {
19+
20+
::executorch::aten::Tensor& sdpa_bitwise_causal_mask_gen_out(
21+
::executorch::runtime::KernelRuntimeContext& ctx,
22+
const ::executorch::aten::Tensor& positions,
23+
int64_t key_length,
24+
::executorch::aten::Tensor& out);
25+
26+
} // namespace native
27+
} // namespace generic
28+
} // namespace impl

‎backends/cadence/generic/operators/targets.bzl‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -471,6 +471,19 @@ def define_common_targets():
471471
visibility = ["PUBLIC"],
472472
)
473473

474+
runtime.cxx_library(
475+
name = "op_sdpa_bitwise_causal_mask_gen",
476+
srcs = ["op_sdpa_bitwise_causal_mask_gen.cpp"],
477+
exported_headers = ["op_sdpa_bitwise_causal_mask_gen.h"],
478+
platforms = CXX,
479+
deps = [
480+
"//executorch/runtime/kernel:kernel_includes",
481+
"//executorch/runtime/core/exec_aten:lib",
482+
"//executorch/runtime/kernel:kernel_runtime_context",
483+
],
484+
visibility = ["PUBLIC"],
485+
)
486+
474487
runtime.cxx_library(
475488
name = "op_sdpa_bitwise_mask_gen",
476489
srcs = ["op_sdpa_bitwise_mask_gen.cpp"],

0 commit comments

Comments
 (0)