graph: api, doc, interface, utils, backend: support bottom-right implicit causal mask #2967

ElaineBao · 2025-03-27T10:02:46Z

Description

Implementation of RFC: #2885.
Support bottom-right implicit causal mask (large partition kernel) in Graph API. SDPA primitive kernel will be supported after GPU primitive support.

Currently this PR depends on :

GPU binary s32 data type enabling: gpu: intel: ocl: add s32 support for binary #2876.
GPU binary select broadcast support: MFDNN-13247

src/graph/backend/dnnl/dnnl_op_def.hpp

examples/graph/sdpa_bottom_right_implicit_mask.cpp

src/graph/interface/tensor.cpp

include/oneapi/dnnl/dnnl_graph.hpp

src/graph/interface/logical_tensor.cpp

TaoLv · 2025-03-31T14:16:53Z

src/graph/interface/logical_tensor.cpp

-    std::fill(val.dims, val.dims + DNNL_MAX_NDIMS, DNNL_GRAPH_UNKNOWN_DIM);
-    std::fill(val.layout.strides, val.layout.strides + DNNL_MAX_NDIMS,
-            DNNL_GRAPH_UNKNOWN_DIM);
+    if (ndims == 0) {


why the change?

To align the API behavior with other logical_tensor_init_... APIs.

We need to make sure that the uninitialized fields will not affect the hash results.

If ndims = 0, the hash result only combines ndims, without touch to the dims and strides, so this change doesn't affect the hash result.

src/graph/interface/tensor.cpp

doc/graph/operations/Subtract.md

rongzha1 · 2025-04-01T01:37:06Z

include/oneapi/dnnl/dnnl_graph.hpp

@@ -306,6 +306,7 @@ class logical_tensor {
        /// the library. For example, constant weight tensors in inference
        /// scenarios.
        constant = dnnl_graph_tensor_property_constant,
+        host_scalar = dnnl_graph_tensor_property_host_scalar,


is host_scalar constant? how to set const property for host_scalar?

host_scalar is not a constant, we don't have a way to set it as a const currently. There was some discussion here, but no decision yet.

src/graph/backend/dnnl/dnnl_op_def.hpp

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

src/graph/backend/dnnl/passes/transform.hpp

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

src/graph/backend/dnnl/op_executable.hpp

ElaineBao · 2025-04-01T15:04:15Z

make test
set test_scope=NIGHTLY
disable benchdnn_all
enable benchdnn_graph

TaoLv · 2025-04-02T02:33:11Z

src/graph/interface/logical_tensor.cpp

-    std::fill(val.dims, val.dims + DNNL_MAX_NDIMS, DNNL_GRAPH_UNKNOWN_DIM);
-    std::fill(val.layout.strides, val.layout.strides + DNNL_MAX_NDIMS,
-            DNNL_GRAPH_UNKNOWN_DIM);
+    if (ndims == 0) {


We need to make sure that the uninitialized fields will not affect the hash results.

src/graph/interface/tensor.hpp

TaoLv · 2025-04-02T03:12:04Z

@ElaineBao I want to confirm that the fusion and implementation still work when s_kv and s_q are device-side scalars.

ElaineBao · 2025-04-02T05:41:43Z

@ElaineBao I want to confirm that the fusion and implementation still work when s_kv and s_q are device-side scalars.

Yes, I just tried to modify the shape to 1D and also reset the property_type, it can pass.

onednn_verbose,v1,graph,exec,gpu,100002,sdp,bmm1;scale_mul_op;gen_index_row_op;mask_add_op;mask_sub_op;gen_index_col_op;mask_ge_op;mask_select_op;softmax;bmm2,,in0_f16:0:strided:undef:32x16x384x64:393216s24576s64s1 in1_f16:1:strided:undef:32x16x384x64:393216s24576s64s1 in2_f16:4:strided:undef:1:1 in3_s32:10:strided:undef:1:1 in4_s32:13:strided:undef:1:1 in5_f32:20:strided:undef:1:1 in6_f16:24:strided:undef:32x16x384x64:393216s24576s64s1 out0_f16:25:strided:undef:32x16x384x64:393216s24576s64s1,fpm:strict,larger_partition_kernel_t,dnnl_backend,118.229

ElaineBao · 2025-04-02T05:51:24Z

make test
set test_scope=NIGHTLY
disable benchdnn_all
enable benchdnn_graph

src/graph/backend/dnnl/op_executable.hpp

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

in sdp_primitive_kernel

which is used in SDPA bottom-right causal mask

ElaineBao · 2025-04-02T09:06:34Z

make test
set test_scope=NIGHTLY
disable benchdnn_all
enable benchdnn_graph

dzarukin · 2025-04-02T17:07:27Z

include/oneapi/dnnl/dnnl_graph_types.h

@@ -76,6 +76,7 @@ typedef enum {
    /// optimizations for constant tensors or cache constant tensors inside the
    /// library. For example, constant weight tensors in inference scenarios.
    dnnl_graph_tensor_property_constant = 2,
+    dnnl_graph_tensor_property_host_scalar = 3,


Documentation?

dzarukin · 2025-04-02T17:09:40Z

doc/programming_model/data_types.md

@@ -16,6 +16,7 @@ in comparison to fp32.
 | f16       | [IEEE half precision floating-point](https://en.wikipedia.org/wiki/Half-precision_floating-point_format#IEEE_754_half-precision_binary_floating-point_format:_binary16)                 |
 | s8/u8     | signed/unsigned 8-bit integer                                                                                                                                                           |
 | s4/u4     | signed/unsigned 4-bit integer                                                                                                                                                           |
+| s32       | signed/unsigned 32-bit integer                                                                                                                                                          |


It probably could use a wiki page link. @vpirogov, could you share your thoughts, please?

dzarukin · 2025-04-02T17:16:03Z

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

+            // check unimplemented bottom-right causal mask
+            auto post_op = get_post_op(cur_op);
+            if (post_op && post_op->get_kind() == graph::op_kind::Add)
+                return status::unimplemented;


Verbose message instead?

dzarukin · 2025-04-02T17:19:04Z

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

-            auto mem = make_dnnl_memory(md, p_engine, nullptr);
-            exec_args_set_.add_value_mem_map({in.get(), mem});
-            classify_mem(mem, in.get());
+            logical_tensor_t in_lt = in->get_logical_tensor();


dzarukin · 2025-04-02T17:25:55Z

src/graph/interface/tensor.hpp

+        // property is host_scalar
+        int32_t s32_value = 0;
+        // TODO: add more dtype support
+    } scalar;


Suggested change

} scalar;

} scalar_;

dzarukin · 2025-04-02T17:29:37Z

tests/benchdnn/graph/graph.cpp

@@ -407,11 +407,6 @@ void skip_unimplemented_ops(const dnnl::graph::partition &partition,
        const deserialized_graph_t &dg, res_t *res) {
    // A list of ops that don't have DNNL backend support so far.
    static const std::vector<std::string> unimplemented_ops {"Pow"};
-    // A list of ops that don't have DNNL backend support so far on GPU.
-    static const std::vector<std::string> unimplemented_ops_gpu {
-            "GreaterEqual"};


I'd just remove the op from unimplemented_ops_gpu and kept rest infrastructure intact for future unsupported ops.

dzarukin · 2025-04-02T17:35:34Z

tests/benchdnn/graph/graph_memory.cpp

@@ -112,9 +114,14 @@ dnnl::graph::tensor dnn_graph_mem_t::make_graph_tensor(
    dnnl_memory_get_data_handle(mem_.m_, &data_handle);
    dnnl::graph::logical_tensor graph_lt(lt.id_, lt.get_data_type(), lt.shape_,
            str2layout(lt.layout_type_), lt.get_property_type());
-    dnnl::graph::tensor ret(graph_lt, get_graph_engine(), data_handle);


Maybe use

const auto &g_eng = is_host_scalar ? get_graph_host_engine().operator const dnnl::engine &() : get_graph_engine().operator const dnnl::engine &();

instead?

dzarukin · 2025-04-02T17:36:46Z

tests/benchdnn/graph/input_displacer.cpp

@@ -309,7 +309,7 @@ int partition_data_displacer_t::displace_input_data(
    } else if (filling_type == filling_type_t::causal_mask) {
        SAFE(gen_causal_mask_filling(mem_replace, mem.md_, res), WARN);
    } else if (filling_type == filling_type_t::minus_infinity) {
-        static const std::vector<float> user_set {-INFINITY};
+        static const std::vector<float> user_set {-1e4};


dzarukin · 2025-04-02T17:40:48Z

tests/benchdnn/graph/utils.hpp

+}
+
+inline const cpp_engine_t &get_graph_host_engine() {
+    const dnnl::engine &g_eng


Suggested change

const dnnl::engine &g_eng

// Return `get_graph_engine` for `is_cpu` to avoid different engines.

const dnnl::engine &g_eng

ElaineBao added the component:graph-api label Mar 27, 2025

ElaineBao self-assigned this Mar 27, 2025

ElaineBao requested review from a team as code owners March 27, 2025 10:02

github-actions bot added documentation platform:gpu-intel component:api component:tests component:examples labels Mar 27, 2025

TaoLv reviewed Mar 27, 2025

View reviewed changes

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

examples/graph/sdpa_bottom_right_implicit_mask.cpp Outdated Show resolved Hide resolved

src/graph/interface/tensor.cpp Outdated Show resolved Hide resolved

ElaineBao force-pushed the yixin/bottom-right-causal-ref branch from 31a23c0 to c29fc8f Compare March 27, 2025 15:38

github-actions bot removed the component:examples label Mar 27, 2025

ElaineBao force-pushed the yixin/bottom-right-causal-ref branch from c29fc8f to a653332 Compare March 28, 2025 06:00

github-actions bot removed the platform:gpu-intel label Mar 28, 2025

ElaineBao removed the request for review from a team March 28, 2025 06:01

ElaineBao force-pushed the yixin/bottom-right-causal-ref branch from a653332 to 91f287f Compare March 31, 2025 13:43

TaoLv reviewed Mar 31, 2025

View reviewed changes

rongzha1 reviewed Apr 1, 2025

View reviewed changes

xiang1guo reviewed Apr 1, 2025

View reviewed changes

src/graph/backend/dnnl/passes/transform.hpp Outdated Show resolved Hide resolved

ElaineBao force-pushed the yixin/bottom-right-causal-ref branch from 91f287f to ef08d4c Compare April 1, 2025 02:32

xiang1guo reviewed Apr 1, 2025

View reviewed changes

src/graph/backend/dnnl/passes/insert_ops.cpp Outdated Show resolved Hide resolved

src/graph/backend/dnnl/passes/insert_ops.cpp Outdated Show resolved Hide resolved

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

ElaineBao force-pushed the yixin/bottom-right-causal-ref branch from ef08d4c to a4aa4c6 Compare April 1, 2025 05:44

TaoLv reviewed Apr 2, 2025

View reviewed changes

ElaineBao added 3 commits April 1, 2025 22:35

api: graph: add host_scalar property for logical tensor

78b7604

graph: interface, utils: support host_scalar property

3157c2b

graph: interface: extend Add and Subtract to support s32 dtype

9d7f7ea

ElaineBao added 2 commits April 1, 2025 22:40

doc: graph: extend Add and Subtract to support s32 dtype

d30d6a9

graph: backend: dnnl: enable bottom-right causal mask pattern

fc0af57

ElaineBao force-pushed the yixin/bottom-right-causal-ref branch from 2ce1c5b to c69976f Compare April 2, 2025 05:41

gyhintel reviewed Apr 2, 2025

View reviewed changes

src/graph/backend/dnnl/op_executable.hpp Outdated Show resolved Hide resolved

gyhintel reviewed Apr 2, 2025

View reviewed changes

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

ElaineBao added 9 commits April 2, 2025 02:03

graph: backend: dnnl: add internal op dnnl_host_scalar

2ca91af

graph: backend: dnnl: handling host_scalar in large_partition_kernel

1129631

graph: backend: dnnl: filter out bottom-right causal mask case

a7c6a59

in sdp_primitive_kernel

benchdnn: graph: support host_scalar graph memory allocation

0267867

benchdnn: graph: support dtype rewrite for s32 Add/Sub

b5171bb

which is used in SDPA bottom-right causal mask

benchdnn: graph: remove unnecessary skip of GreaterEqual

7b06d0f

benchdnn: graph: add test case for bottom-right implicit causal mask

6cc38ba

doc: graph: update sdpa fusion pattern doc

fad8841

temporary: set -inf to -1e4 to bypass decomposed select NaN issue

Loading
Loading status checks…

1987805

ElaineBao force-pushed the yixin/bottom-right-causal-ref branch from c69976f to 1987805 Compare April 2, 2025 09:06

gyhintel approved these changes Apr 2, 2025

View reviewed changes

dzarukin reviewed Apr 2, 2025

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

graph: api, doc, interface, utils, backend: support bottom-right implicit causal mask #2967

graph: api, doc, interface, utils, backend: support bottom-right implicit causal mask #2967

ElaineBao commented Mar 27, 2025

TaoLv Mar 31, 2025

ElaineBao Mar 31, 2025

TaoLv Apr 2, 2025

ElaineBao Apr 2, 2025

rongzha1 Apr 1, 2025

ElaineBao Apr 1, 2025

ElaineBao commented Apr 1, 2025

TaoLv Apr 2, 2025

TaoLv commented Apr 2, 2025

ElaineBao commented Apr 2, 2025

ElaineBao commented Apr 2, 2025

ElaineBao commented Apr 2, 2025

dzarukin Apr 2, 2025

dzarukin Apr 2, 2025

dzarukin Apr 2, 2025

dzarukin Apr 2, 2025

dzarukin Apr 2, 2025

dzarukin Apr 2, 2025

dzarukin Apr 2, 2025

dzarukin Apr 2, 2025

dzarukin Apr 2, 2025

	const dnnl::engine &g_eng
	// Return `get_graph_engine` for `is_cpu` to avoid different engines.
	const dnnl::engine &g_eng

graph: api, doc, interface, utils, backend: support bottom-right implicit causal mask #2967

Are you sure you want to change the base?

graph: api, doc, interface, utils, backend: support bottom-right implicit causal mask #2967

Conversation

ElaineBao commented Mar 27, 2025

Description

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 Apr 1, 2025

Choose a reason for hiding this comment

TaoLv commented Apr 2, 2025

ElaineBao commented Apr 2, 2025

ElaineBao commented Apr 2, 2025

ElaineBao commented Apr 2, 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

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment