From 7e2f90e2019eba16c08eb152f9f288bbadb452e9 Mon Sep 17 00:00:00 2001 From: Masaki Kozuki Date: Thu, 18 Sep 2025 01:24:57 -0700 Subject: [PATCH] Have seed and offset in int64 for cudnn-frontend SDPA Signed-off-by: Masaki Kozuki --- thunder/executors/cudnn_sdpa.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/thunder/executors/cudnn_sdpa.py b/thunder/executors/cudnn_sdpa.py index 00fb76424a..aef74b3225 100644 --- a/thunder/executors/cudnn_sdpa.py +++ b/thunder/executors/cudnn_sdpa.py @@ -97,10 +97,10 @@ def _make_cudnn_sdpa_forward_graph( Offset = None if dropout_p != 0.0: Seed = graph.tensor( - name="Seed", dim=scalar_dim_stride, stride=scalar_dim_stride, data_type=cudnn.data_type.INT32 + name="Seed", dim=scalar_dim_stride, stride=scalar_dim_stride, data_type=cudnn.data_type.INT64 ) Offset = graph.tensor( - name="Offset", dim=scalar_dim_stride, stride=scalar_dim_stride, data_type=cudnn.data_type.INT32 + name="Offset", dim=scalar_dim_stride, stride=scalar_dim_stride, data_type=cudnn.data_type.INT64 ) dropout_tuple = (dropout_p, Seed, Offset) @@ -450,10 +450,10 @@ def _make_cudnn_sdpa_backward_graph( Offset = None if dropout_p != 0.0: Seed = graph.tensor( - name="Seed", dim=scalar_dim_stride, stride=scalar_dim_stride, data_type=cudnn.data_type.INT32 + name="Seed", dim=scalar_dim_stride, stride=scalar_dim_stride, data_type=cudnn.data_type.INT64 ) Offset = graph.tensor( - name="Offset", dim=scalar_dim_stride, stride=scalar_dim_stride, data_type=cudnn.data_type.INT32 + name="Offset", dim=scalar_dim_stride, stride=scalar_dim_stride, data_type=cudnn.data_type.INT64 ) dropout_tuple = (dropout_p, Seed, Offset)