Skip to content

Commit 85971ff

Browse files
mfeliz-cruisemfeliz-cruise
authored
[feat] Support conversion of scaled_dot_product_attention (#2549)
1 parent 5dab7ca commit 85971ff

File tree

11 files changed

+276
-1
lines changed

11 files changed

+276
-1
lines changed

‎core/conversion/converters/BUILD‎

100755100644
Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -66,6 +66,7 @@ cc_library(
6666
"impl/einsum.cpp",
6767
"impl/element_wise.cpp",
6868
"impl/expand.cpp",
69+
"impl/internal_ops.cpp",
6970
"impl/interpolate.cpp",
7071
"impl/layer_norm.cpp",
7172
"impl/linear.cpp",

‎core/conversion/converters/impl/internal_ops.cpp‎

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,46 @@
1+
#include "core/conversion/converters/converters.h"
2+
#include "core/util/prelude.h"
3+
#include "torch/torch.h"
4+
5+
namespace torch_tensorrt {
6+
namespace core {
7+
namespace conversion {
8+
namespace converters {
9+
namespace impl {
10+
namespace {
11+
12+
auto linear_registrations TORCHTRT_UNUSED = RegisterNodeConversionPatterns().pattern(
13+
{"trt::attn_bias_from_attn_mask(Tensor attn_mask) -> Tensor",
14+
[](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
15+
// Converter for internal op used in unpack_scaled_dot_product_attention
16+
// We don't have visibility to check types during lowering and can't introduce conditionals so do type specific
17+
// specialization here
18+
auto in = args[0].ITensorOrFreeze(ctx);
19+
auto out = in;
20+
if (in->getType() == nvinfer1::DataType::kBOOL) {
21+
auto not_layer = ctx->net->addUnary(*in, nvinfer1::UnaryOperation::kNOT);
22+
TORCHTRT_CHECK(not_layer, "Unable to create not layer for attn_bias_from_attn_mask");
23+
not_layer->setName((util::node_info(n) + "_not").c_str());
24+
auto neg_inf = torch::tensor(-std::numeric_limits<float>::infinity());
25+
auto neg_inf_itensor = tensor_to_const(ctx, neg_inf);
26+
auto prod_layer = add_elementwise(
27+
ctx,
28+
nvinfer1::ElementWiseOperation::kPROD,
29+
not_layer->getOutput(0),
30+
neg_inf_itensor,
31+
util::node_info(n) + "_mul");
32+
auto add_layer = add_elementwise(
33+
ctx, nvinfer1::ElementWiseOperation::kSUM, prod_layer->getOutput(0), in, util::node_info(n) + "_add");
34+
out = add_layer->getOutput(0);
35+
}
36+
auto out_tensor = ctx->AssociateValueAndTensor(n->outputs()[0], out);
37+
LOG_DEBUG("Output tensor shape: " << out_tensor->getDimensions());
38+
LOG_DEBUG("Output tensor type: " << out_tensor->getType());
39+
return true;
40+
}});
41+
} // namespace
42+
} // namespace impl
43+
} // namespace converters
44+
} // namespace conversion
45+
} // namespace core
46+
} // namespace torch_tensorrt

‎core/conversion/converters/impl/unary.cpp‎

Lines changed: 16 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -79,6 +79,22 @@ auto logical_not_registration TORCHTRT_UNUSED = RegisterNodeConversionPatterns()
7979
return true;
8080
}});
8181

82+
auto sqrt_registration TORCHTRT_UNUSED = RegisterNodeConversionPatterns().pattern(
83+
{"aten::sqrt(Tensor self) -> Tensor", [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
84+
auto in = args[0].ITensorOrFreeze(ctx);
85+
if (in->getType() == nvinfer1::DataType::kINT32) {
86+
// unary sqrt layer only supports float inputs
87+
in = castITensor(ctx, in, nvinfer1::DataType::kFLOAT, util::node_info(n).c_str());
88+
}
89+
auto unary_layer = ctx->net->addUnary(*in, nvinfer1::UnaryOperation::kSQRT);
90+
TORCHTRT_CHECK(unary_layer, "Unable to create sqrt layer from node: " << *n);
91+
unary_layer->setName(util::node_info(n).c_str());
92+
unary_layer->setOutputType(0, in->getType());
93+
auto out_tensor = ctx->AssociateValueAndTensor(n->outputs()[0], unary_layer->getOutput(0));
94+
LOG_DEBUG("Output tensor shape: " << out_tensor->getDimensions());
95+
return true;
96+
}});
97+
8298
auto isfinite_registration TORCHTRT_UNUSED = RegisterNodeConversionPatterns().pattern(
8399
{"aten::isfinite(Tensor self) -> Tensor", [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
84100
auto in = args[0].ITensorOrFreeze(ctx);
@@ -126,7 +142,6 @@ convert(atan, kATAN);
126142
convert(floor, kFLOOR);
127143
convert(log, kLOG);
128144
convert(ceil, kCEIL);
129-
convert(sqrt, kSQRT);
130145
convert(exp, kEXP);
131146
convert(neg, kNEG);
132147
convert(erf, kERF);

‎core/lowering/lowering.cpp‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -146,6 +146,7 @@ void LowerGraph(std::shared_ptr<torch::jit::Graph>& g, std::vector<torch::jit::I
146146
if (lower_info.converting_to_trt_engine) {
147147
passes::RemoveCollectionCast(g);
148148
}
149+
passes::UnpackScaledDotProductAttention(g);
149150
passes::UnpackAndCastMaskedFill(g, lower_info.getGPUDeviceString());
150151
passes::UnpackAndCastNumToTensor(g, lower_info.getGPUDeviceString());
151152
passes::UnpackAndCastFull(g, lower_info.getGPUDeviceString());

‎core/lowering/passes/BUILD‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,7 @@ cc_library(
3838
"unpack_hardswish.cpp",
3939
"unpack_log_softmax.cpp",
4040
"unpack_rsqrt.cpp",
41+
"unpack_scaled_dot_product_attention.cpp",
4142
"unpack_std.cpp",
4243
"unpack_var.cpp",
4344
"view_to_reshape.cpp",

‎core/lowering/passes/passes.h‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,7 @@ void UnpackHardSigmoid(std::shared_ptr<torch::jit::Graph>& graph);
4949
void UnpackAndCastMaskedFill(std::shared_ptr<torch::jit::Graph>& graph, std::string target_device_name);
5050
void UnpackAndCastNumToTensor(std::shared_ptr<torch::jit::Graph>& graph, std::string target_device_name);
5151
void UnpackAndCastFull(std::shared_ptr<torch::jit::Graph>& graph, std::string target_device_name);
52+
void UnpackScaledDotProductAttention(std::shared_ptr<torch::jit::Graph>& graph);
5253
void ReplaceScalarImplicit(std::shared_ptr<torch::jit::Graph>& graph);
5354
void ReplaceAtenPad(std::shared_ptr<torch::jit::Graph>& graph);
5455
void ReplaceTileWithRepeat(std::shared_ptr<torch::jit::Graph>& graph);

‎core/lowering/passes/unpack_scaled_dot_product_attention.cpp‎

Lines changed: 94 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,94 @@
1+
#include "torch/csrc/jit/ir/subgraph_matcher.h"
2+
#include "torch/csrc/jit/passes/subgraph_rewrite.h"
3+
4+
#include "core/util/prelude.h"
5+
#include "torch/csrc/jit/ir/irparser.h"
6+
7+
namespace torch_tensorrt {
8+
namespace core {
9+
namespace lowering {
10+
namespace passes {
11+
12+
// https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html
13+
void UnpackScaledDotProductAttention(std::shared_ptr<torch::jit::Graph>& graph) {
14+
std::string sdpa_pattern = R"IR(
15+
graph(%query, %key, %value, %attn_mask, %dropout_p, %is_causal):
16+
%out: Tensor = aten::scaled_dot_product_attention(%query, %key, %value, %attn_mask, %dropout_p, %is_causal)
17+
return (%out))IR";
18+
19+
std::string unpacked_sdpa_pattern = R"IR(
20+
graph(%query, %key, %value, %attn_mask, %dropout_p, %is_causal):
21+
%none : NoneType = prim::Constant()
22+
%1 : int = prim::Constant[value=-1]()
23+
%2 : int = prim::Constant[value=-2]()
24+
%3 : int = aten::size(%query, %1)
25+
%q_size : Long() = prim::NumToTensor(%3)
26+
%sqrt : Tensor = aten::sqrt(%q_size)
27+
%scale_factor : Tensor = aten::reciprocal(%sqrt)
28+
%key_transpose : Tensor = aten::transpose(%key, %2, %1)
29+
%matmul : Tensor = aten::matmul(%query, %key_transpose)
30+
%attn_weight : Tensor = aten::mul(%matmul, %scale_factor)
31+
%softmax : Tensor = aten::softmax(%attn_weight, %1, %none)
32+
%out : Tensor = aten::matmul(%softmax, %value)
33+
return(%out))IR";
34+
35+
std::string unpacked_sdpa_attn_biased_pattern = R"IR(
36+
graph(%query, %key, %value, %attn_mask, %dropout_p, %is_causal):
37+
%none : NoneType = prim::Constant()
38+
%0 : int = prim::Constant[value=1]()
39+
%1 : int = prim::Constant[value=-1]()
40+
%2 : int = prim::Constant[value=-2]()
41+
%3 : int = aten::size(%query, %1)
42+
%q_size : Long() = prim::NumToTensor(%3)
43+
%sqrt : Tensor = aten::sqrt(%q_size)
44+
%scale_factor : Tensor = aten::reciprocal(%sqrt)
45+
%key_transpose : Tensor = aten::transpose(%key, %2, %1)
46+
%matmul : Tensor = aten::matmul(%query, %key_transpose)
47+
%attn_weight : Tensor = aten::mul(%matmul, %scale_factor)
48+
%attn_bias : Tensor = trt::attn_bias_from_attn_mask(%attn_mask)
49+
%attn_weight_with_bias : Tensor = aten::add(%attn_weight, %attn_bias, %0)
50+
%softmax : Tensor = aten::softmax(%attn_weight_with_bias, %1, %none)
51+
%out : Tensor = aten::matmul(%softmax, %value)
52+
return(%out))IR";
53+
54+
// rewrite with None attn_mask
55+
torch::jit::SubgraphRewriter sdpa_rewriter;
56+
sdpa_rewriter.RegisterRewritePattern(sdpa_pattern, unpacked_sdpa_pattern);
57+
sdpa_rewriter.runOnGraph(
58+
graph, [](const torch::jit::Match& match, const std::unordered_map<std::string, torch::jit::Value*>&) {
59+
auto is_causal_node = match.anchor->inputs().at(5)->node();
60+
if (is_causal_node->kind() != at::prim::Constant) {
61+
LOG_WARNING("Could not unpack scaled_dot_product_attention with non constant is_causal: " << *is_causal_node);
62+
return false;
63+
}
64+
if (is_causal_node->i(at::attr::value) == 1) {
65+
LOG_WARNING("Could not unpack scaled_dot_product_attention with is_causal = True: " << *is_causal_node);
66+
return false;
67+
}
68+
auto attn_mask_node = match.anchor->inputs().at(3)->node();
69+
if (attn_mask_node->kind() != at::prim::Constant || !attn_mask_node->mustBeNone()) {
70+
return false;
71+
}
72+
return true;
73+
});
74+
75+
// rewrite with float/bool attn_mask this uses a custom op to implement the divergent behavior between bool and float
76+
// masks without a conditional
77+
torch::jit::SubgraphRewriter sdpa_attn_mask_rewriter;
78+
sdpa_attn_mask_rewriter.RegisterRewritePattern(sdpa_pattern, unpacked_sdpa_attn_biased_pattern);
79+
sdpa_attn_mask_rewriter.runOnGraph(
80+
graph, [](const torch::jit::Match& match, const std::unordered_map<std::string, torch::jit::Value*>&) {
81+
auto is_causal_node = match.anchor->inputs().at(5)->node();
82+
if (is_causal_node->kind() != at::prim::Constant || is_causal_node->i(at::attr::value) == 1) {
83+
// messages already written in first pass, do not write again
84+
return false;
85+
}
86+
return true;
87+
});
88+
LOG_GRAPH("Post unpack scaled_dot_product_attention: " << *graph);
89+
}
90+
91+
} // namespace passes
92+
} // namespace lowering
93+
} // namespace core
94+
} // namespace torch_tensorrt

‎core/lowering/register_trt_placeholder_ops.cpp‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
#include <limits>
12
#include "torch/csrc/jit/runtime/custom_operator.h"
23

34
namespace torch {
@@ -14,6 +15,17 @@ RegisterOperators trt_placeholder_ops_reg({
1415
"trt::const(Tensor val) -> Tensor",
1516
[](Stack& stack) { /*noop*/ },
1617
aliasAnalysisFromSchema()),
18+
Operator(
19+
"trt::attn_bias_from_attn_mask(Tensor attn_mask) -> Tensor",
20+
[](Stack& stack) {
21+
auto attn_mask = pop(stack).to<at::Tensor>();
22+
if (attn_mask.scalar_type() == at::kBool) {
23+
attn_mask = attn_mask;
24+
attn_mask.masked_fill_(attn_mask.logical_not(), -std::numeric_limits<float>::infinity());
25+
}
26+
return attn_mask;
27+
},
28+
c10::AliasAnalysisKind::CONSERVATIVE),
1729
});
1830

1931
} // namespace jit

‎tests/core/conversion/converters/BUILD‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -203,6 +203,10 @@ converter_test(
203203
name = "test_where",
204204
)
205205

206+
converter_test(
207+
name = "test_scaled_dot_product_attention",
208+
)
209+
206210
test_suite(
207211
name = "converter_tests",
208212
tests = [
@@ -238,6 +242,7 @@ test_suite(
238242
":test_reduce",
239243
":test_replication_pad",
240244
":test_roll",
245+
":test_scaled_dot_product_attention",
241246
":test_scatter",
242247
":test_select",
243248
":test_shuffle",

‎tests/core/conversion/converters/test_scaled_dot_product_attention.cpp‎

Lines changed: 84 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,84 @@
1+
#include <string>
2+
#include "core/compiler.h"
3+
#include "core/lowering/passes/passes.h"
4+
#include "gtest/gtest.h"
5+
#include "tests/util/util.h"
6+
#include "torch/csrc/jit/ir/irparser.h"
7+
8+
TEST(Converters, ATenScaledDotProductAttentionConvertsCorrectly) {
9+
const auto graph = R"IR(
10+
graph(%query : Tensor, %key : Tensor, %value : Tensor):
11+
%none : NoneType = prim::Constant()
12+
%0 : float = prim::Constant[value=0.]()
13+
%false : bool = prim::Constant[value=0]()
14+
%3 : Tensor = aten::scaled_dot_product_attention(%query, %key, %value, %none, %0, %false)
15+
return (%3))IR";
16+
17+
auto g = std::make_shared<torch::jit::Graph>();
18+
torch::jit::parseIR(graph, &*g);
19+
20+
auto query = at::rand({32, 8, 128, 64}, {at::kCUDA});
21+
auto key = at::rand({32, 8, 128, 64}, {at::kCUDA});
22+
auto value = at::rand({32, 8, 128, 64}, {at::kCUDA});
23+
auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {});
24+
auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {query, key, value});
25+
26+
torch_tensorrt::core::lowering::passes::UnpackScaledDotProductAttention(g);
27+
28+
params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {});
29+
auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {query, key, value});
30+
31+
ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0], 1e-5));
32+
}
33+
34+
TEST(Converters, ATenScaledDotProductAttnMaskFloatConvertsCorrectly) {
35+
const auto graph = R"IR(
36+
graph(%query : Tensor, %key : Tensor, %value : Tensor, %attn_mask : Tensor):
37+
%0 : float = prim::Constant[value=0.]()
38+
%false : bool = prim::Constant[value=0]()
39+
%3 : Tensor = aten::scaled_dot_product_attention(%query, %key, %value, %attn_mask, %0, %false)
40+
return (%3))IR";
41+
42+
auto g = std::make_shared<torch::jit::Graph>();
43+
torch::jit::parseIR(graph, &*g);
44+
45+
auto query = at::rand({32, 8, 128, 64}, {at::kCUDA});
46+
auto key = at::rand({32, 8, 128, 64}, {at::kCUDA});
47+
auto value = at::rand({32, 8, 128, 64}, {at::kCUDA});
48+
auto attn_mask = at::rand({32, 8, 128, 128}, {at::kCUDA});
49+
auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {});
50+
auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {query, key, value, attn_mask});
51+
52+
torch_tensorrt::core::lowering::passes::UnpackScaledDotProductAttention(g);
53+
54+
params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {});
55+
auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {query, key, value, attn_mask});
56+
57+
ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0], 1e-5));
58+
}
59+
60+
TEST(Converters, ATenScaledDotProductAttnMaskBoolConvertsCorrectly) {
61+
const auto graph = R"IR(
62+
graph(%query : Tensor, %key : Tensor, %value : Tensor, %attn_mask : Tensor):
63+
%0 : float = prim::Constant[value=0.]()
64+
%false : bool = prim::Constant[value=0]()
65+
%3 : Tensor = aten::scaled_dot_product_attention(%query, %key, %value, %attn_mask, %0, %false)
66+
return (%3))IR";
67+
68+
auto g = std::make_shared<torch::jit::Graph>();
69+
torch::jit::parseIR(graph, &*g);
70+
71+
auto query = at::rand({32, 8, 128, 64}, {at::kCUDA});
72+
auto key = at::rand({32, 8, 128, 64}, {at::kCUDA});
73+
auto value = at::rand({32, 8, 128, 64}, {at::kCUDA});
74+
auto attn_mask = at::randint(0, 2, {32, 8, 128, 128}, at::kCUDA).to(at::kBool);
75+
auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {});
76+
auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {query, key, value, attn_mask});
77+
78+
torch_tensorrt::core::lowering::passes::UnpackScaledDotProductAttention(g);
79+
80+
params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {});
81+
auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {query, key, value, attn_mask});
82+
83+
ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0], 1e-5));
84+
}

0 commit comments

Comments
 (0)