graph, benchdnn: support implicit causal mask in graph API #2330

ElaineBao · 2025-01-02T04:27:45Z

Description

The implementation of option 1.1 (top-left aligned causal mask with subgraph approach) in rfcs: graph api: support implicit causal mask in SDPA

include/oneapi/dnnl/dnnl_graph_types.h

src/graph/backend/dnnl/kernels/gen_index.cpp

src/graph/backend/dnnl/dnnl_op_def.hpp

tests/benchdnn/graph/ref_primitive.cpp

src/graph/backend/dnnl/passes/lower.cpp

src/graph/backend/dnnl/kernels/gen_index.cpp

tests/benchdnn/graph/custom_driver.cpp

ElaineBao · 2025-01-03T01:34:47Z

make test
enable benchdnn_nightly
disable benchdnn_all
enable benchdnn_graph

include/oneapi/dnnl/dnnl_graph.hpp

src/graph/backend/dnnl/kernels/gen_index.hpp

TaoLv · 2025-01-03T07:24:10Z

src/graph/backend/dnnl/patterns/sdp.cpp

@@ -146,6 +146,37 @@ DNNL_BACKEND_REGISTER_PATTERN_MATCHER_PASS(dnnl, float_sdp_fusion)
            return std::make_shared<sdp_base_t<>>();
        });

+DNNL_BACKEND_REGISTER_PATTERN_MATCHER_PASS(dnnl, float_sdp_implicit_mask_fusion)
+        .set_priority(21.0f)
+        .set_engine_kind(engine_kind::cpu)


the currently implementation of gen_index only focus on CPU

do you have any plan to implement it for GPU also?

Yes, but most likely in a separate PR, and not target v3.7.

Given that this pattern should work for optimized GPU version, I don't think it make sense to limit to engine kind.
I'd expect it returns unimplemented later in the flow (it won't pick up any other patterns anyway).

The concern is that: if we don't limit the engine kind here, then for GPU case, the unimplemented status returned to user will be at execution stage, which may be too late for them to make necessary modifications.

In contrast, if we adopt the current CPU handling (and plan to remove this constraint once GPU support is implemented), then unsupported partitions will be returned earlier at get_partitions stage, this approach allows users to handle these unsupported partitions earlier by themselves.

src/graph/backend/dnnl/utils.hpp

tests/benchdnn/graph/custom_driver.cpp

tests/gtests/graph/api/test_cpp_api_op.cpp

ElaineBao · 2025-01-06T01:42:17Z

make test
enable benchdnn_nightly
disable benchdnn_all
enable benchdnn_graph

doc/graph/operations/GenIndex.md

doc/graph/operations/GreaterEqual.md

tests/benchdnn/inputs/graph/op/harness_f32_ci

ranukund

I've left a few comments, please incorporate as you see fit, thanks!

doc/graph/fusion_patterns/sdpa.md

doc/graph/operations/GenIndex.md

doc/graph/operations/GreaterEqual.md

gyhintel · 2025-01-08T03:22:02Z

make test
enable benchdnn_nightly
disable benchdnn_all
enable benchdnn_graph

ElaineBao · 2025-01-08T06:37:57Z

make test
enable benchdnn_nightly
disable benchdnn_all
enable benchdnn_graph

ranukund

Few more edits suggested, please incorporate, thanks!

doc/graph/fusion_patterns/sdpa.md

doc/graph/operations/GreaterEqual.md

TaoLv

LGTM. I reviewed API, document, interface folder, part of the DNNL backend implementation.

ElaineBao · 2025-01-08T15:46:45Z

make test
enable benchdnn_nightly
disable benchdnn_all
enable benchdnn_graph

ranukund

Minor edits suggested. Rest looks good, thanks

doc/graph/operations/GenIndex.md

ranukund · 2025-01-08T16:01:01Z

doc/graph/operations/GenIndex.md

+GenIndex{#dev_guide_op_genindex}
+================================
+
+## General


Suggested change

## General

## Overview

Thanks for the suggestion, as other operations are also using General, I suggest to change them all in a separate PR.

ranukund · 2025-01-08T16:01:19Z

doc/graph/operations/GreaterEqual.md

+GreaterEqual{#dev_guide_op_greaterequal}
+========================================
+
+## General


Suggested change

## General

## Overview

ElaineBao · 2025-01-08T16:18:44Z

make test
enable benchdnn_nightly
disable benchdnn_all
enable benchdnn_graph

tests/benchdnn/graph/custom_driver.cpp

ElaineBao added component:graph-api Codeowner: @oneapi-src/onednn-graph component:tests Codeowner: @oneapi-src/onednn-arch labels Jan 2, 2025

ElaineBao self-assigned this Jan 2, 2025

ElaineBao requested review from a team as code owners January 2, 2025 04:27

github-actions bot added the component:api Codeowner: @oneapi-src/onednn-arch label Jan 2, 2025

TaoLv reviewed Jan 2, 2025

View reviewed changes

include/oneapi/dnnl/dnnl_graph_types.h Outdated Show resolved Hide resolved

TaoLv reviewed Jan 2, 2025

View reviewed changes

src/graph/backend/dnnl/kernels/gen_index.cpp Outdated Show resolved Hide resolved

wzt1997 reviewed Jan 2, 2025

View reviewed changes

src/graph/backend/dnnl/dnnl_op_def.hpp Show resolved Hide resolved

tests/benchdnn/graph/ref_primitive.cpp Outdated Show resolved Hide resolved

src/graph/backend/dnnl/passes/lower.cpp Show resolved Hide resolved

ElaineBao force-pushed the graph-implicit-causal-mask branch from f444b05 to 4651fa6 Compare January 2, 2025 06:46

rongzha1 reviewed Jan 2, 2025

View reviewed changes

src/graph/backend/dnnl/kernels/gen_index.cpp Outdated Show resolved Hide resolved

gyhintel force-pushed the graph-implicit-causal-mask branch from 4651fa6 to 2a9ab89 Compare January 2, 2025 07:50

rongzha1 reviewed Jan 2, 2025

View reviewed changes

tests/benchdnn/graph/custom_driver.cpp Show resolved Hide resolved

rongzha1 reviewed Jan 2, 2025

View reviewed changes

tests/benchdnn/graph/custom_driver.cpp Outdated Show resolved Hide resolved

gyhintel force-pushed the graph-implicit-causal-mask branch 2 times, most recently from ab92e00 to 28f56e6 Compare January 2, 2025 13:27

TaoLv reviewed Jan 3, 2025

View reviewed changes

include/oneapi/dnnl/dnnl_graph.hpp Show resolved Hide resolved

TaoLv reviewed Jan 3, 2025

View reviewed changes

gyhintel force-pushed the graph-implicit-causal-mask branch from 28f56e6 to 4975a13 Compare January 3, 2025 08:44

ElaineBao requested a review from a team as a code owner January 4, 2025 07:35

github-actions bot added the documentation A request to change/fix/improve the documentation. Codeowner: @oneapi-src/onednn-doc label Jan 4, 2025

TaoLv reviewed Jan 6, 2025

View reviewed changes

doc/graph/operations/GenIndex.md Outdated Show resolved Hide resolved

doc/graph/operations/GreaterEqual.md Outdated Show resolved Hide resolved

TaoLv reviewed Jan 6, 2025

View reviewed changes

tests/benchdnn/inputs/graph/op/harness_f32_ci Show resolved Hide resolved

ElaineBao force-pushed the graph-implicit-causal-mask branch 2 times, most recently from 78c01cc to 0ac759c Compare January 6, 2025 08:29

ElaineBao changed the title ~~graph: api, interface, backend; benchdnn: graph: support implicit causal mask in graph API~~ graph, benchdnn: support implicit causal mask in graph API Jan 6, 2025

ranukund reviewed Jan 6, 2025

View reviewed changes

ElaineBao force-pushed the graph-implicit-causal-mask branch from 0ac759c to 4e81d51 Compare January 7, 2025 01:28

gyhintel force-pushed the graph-implicit-causal-mask branch from 3bbb517 to 1a555aa Compare January 8, 2025 03:14

ranukund reviewed Jan 8, 2025

View reviewed changes

TaoLv approved these changes Jan 8, 2025

View reviewed changes

dzarukin and others added 13 commits January 8, 2025 07:39

benchdnn: graph: styling

e8e9dfd

benchdnn: graph: move custom op on f32 completely

07589b9

graph: api: support GenIndex and GreaterEqual operations

173c983

graph: interface: support GenIndex and GreaterEqual operations

2c64f55

graph: backend: dnnl: add internal op of dnnl_gen_index

28441d5

graph: dnnl: enable gen_index in dnnl backend

2290e02

graph: dnnl: enable greater_equal in dnnl backend

524a847

benchdnn: graph: support genindex and greaterequal op in benchdnn graph

b481d78

benchdnn: graph: inputs: add gen_index and greater_equal cases

1272246

graph: utils: pm: support multi-consumers input for repetition

69181fc

graph: backend: dnnl: add pattern for sdp with implicit causal mask

df15f6b

benchdnn: graph: inputs: add a sdpa implicit causal mask case

a8c89aa

gtests: graph: api: add gtest for GenIndex and GreaterEqual

c9dab3b

ElaineBao force-pushed the graph-implicit-causal-mask branch from e133094 to 027e58b Compare January 8, 2025 15:40

ranukund approved these changes Jan 8, 2025

View reviewed changes

doc/graph/operations/GenIndex.md Outdated Show resolved Hide resolved

ranukund approved these changes Jan 8, 2025

View reviewed changes

ElaineBao added 3 commits January 8, 2025 08:03

doc: graph: add document for GenIndex and GreaterEqual

f30023b

doc: graph: update sdpa document to include implicit causal mask

6ee19e0

benchdnn: graph: skip unimplemented GenIndex op on gpu

9208bcf

ElaineBao force-pushed the graph-implicit-causal-mask branch from 027e58b to 9208bcf Compare January 8, 2025 16:03

dzarukin approved these changes Jan 8, 2025

View reviewed changes

tests/benchdnn/graph/custom_driver.cpp Show resolved Hide resolved

TaoLv merged commit 8b017b0 into main Jan 9, 2025
16 of 18 checks passed

TaoLv deleted the graph-implicit-causal-mask branch January 9, 2025 07:42

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

graph, benchdnn: support implicit causal mask in graph API #2330

graph, benchdnn: support implicit causal mask in graph API #2330

ElaineBao commented Jan 2, 2025 •

edited

Loading

ElaineBao commented Jan 3, 2025

TaoLv Jan 3, 2025

ElaineBao Jan 3, 2025

TaoLv Jan 6, 2025

ElaineBao Jan 7, 2025

dzarukin Jan 7, 2025

ElaineBao Jan 7, 2025

ElaineBao commented Jan 6, 2025

ranukund left a comment

gyhintel commented Jan 8, 2025

ElaineBao commented Jan 8, 2025

ranukund left a comment

TaoLv left a comment

ElaineBao commented Jan 8, 2025

ranukund left a comment •

edited

Loading

ranukund Jan 8, 2025

ElaineBao Jan 9, 2025

ranukund Jan 8, 2025

ElaineBao commented Jan 8, 2025

graph, benchdnn: support implicit causal mask in graph API #2330

graph, benchdnn: support implicit causal mask in graph API #2330

Conversation

ElaineBao commented Jan 2, 2025 • edited Loading

Description

ElaineBao commented Jan 3, 2025

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ElaineBao commented Jan 6, 2025

ranukund left a comment

Choose a reason for hiding this comment

gyhintel commented Jan 8, 2025

ElaineBao commented Jan 8, 2025

ranukund left a comment

Choose a reason for hiding this comment

TaoLv left a comment

Choose a reason for hiding this comment

ElaineBao commented Jan 8, 2025

ranukund left a comment • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ElaineBao commented Jan 8, 2025

ElaineBao commented Jan 2, 2025 •

edited

Loading

ranukund left a comment •

edited

Loading