Making softmax_lse required output parameter for sdpa.

liuliu · Jun 11, 2024 · f4dd846 · f4dd846
1 parent 7f464ef
commit f4dd846
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 8 deletions.
diff --git a/lib/nnc/ccv_cnnp_model_addons.c b/lib/nnc/ccv_cnnp_model_addons.c
@@ -3886,7 +3886,7 @@ static void _ccv_cnnp_scaled_dot_product_attention_build(ccv_cnnp_model_t* const
 	cmd.info.scaled_dot_product_attention.upcast = self->upcast;
 	ccv_nnc_tensor_param_t output_params[3];
 	ccv_nnc_tensor_symbol_t output;
-	ccv_nnc_tensor_symbol_t saved_softmax;
+	ccv_nnc_tensor_symbol_t saved_softmax_lse;
 	ccv_nnc_tensor_symbol_t saved_v_proj = NO_TENSOR_SYMBOL;
 	ccv_nnc_tensor_symbol_t attn_mask = NO_TENSOR_SYMBOL;
 	ccv_nnc_tensor_symbol_t weights = NO_TENSOR_SYMBOL;
@@ -3913,7 +3913,7 @@ static void _ccv_cnnp_scaled_dot_product_attention_build(ccv_cnnp_model_t* const
 				bias_params,
 			}, 6, ccv_nnc_no_hint, output_params, 3);
 		output = ccv_nnc_tensor_symbol_new(graph, output_params[0], 0);
-		saved_softmax = ccv_nnc_tensor_symbol_new(graph, output_params[1], 0);
+		saved_softmax_lse = ccv_nnc_tensor_symbol_new(graph, output_params[1], 0);
 		saved_v_proj = ccv_nnc_tensor_symbol_new(graph, output_params[2], 0);
 	} else {
 		ccv_nnc_hint_tensor_auto(cmd, (ccv_nnc_tensor_param_t []){
@@ -3922,9 +3922,9 @@ static void _ccv_cnnp_scaled_dot_product_attention_build(ccv_cnnp_model_t* const
 				v_params,
 			}, 3, ccv_nnc_no_hint, output_params, 2);
 		output = ccv_nnc_tensor_symbol_new(graph, output_params[0], 0);
-		saved_softmax = ccv_nnc_tensor_symbol_new(graph, output_params[1], 0);
+		saved_softmax_lse = ccv_nnc_tensor_symbol_new(graph, output_params[1], 0);
 	}
-	ccv_nnc_graph_exec_symbol_new(graph, cmd, TENSOR_SYMBOL_LIST(inputs[0], inputs[1], inputs[2], attn_mask, weights, bias), TENSOR_SYMBOL_LIST(output, saved_softmax, saved_v_proj), "scaled_dot_product_attention");
+	ccv_nnc_graph_exec_symbol_new(graph, cmd, TENSOR_SYMBOL_LIST(inputs[0], inputs[1], inputs[2], attn_mask, weights, bias), TENSOR_SYMBOL_LIST(output, saved_softmax_lse, saved_v_proj), "scaled_dot_product_attention");
 	outputs[0] = output;
 }
 

diff --git a/lib/nnc/cmd/scaled_dot_product_attention/ccv_nnc_scaled_dot_product_attention.c b/lib/nnc/cmd/scaled_dot_product_attention/ccv_nnc_scaled_dot_product_attention.c
@@ -5,12 +5,12 @@
 static int _ccv_nnc_scaled_dot_product_attention_forw_bitmask(const ccv_nnc_cmd_param_t cmd, const int input_size, const int output_size, const uint64_t* const input_bitmasks, const int input_bitmask_size, const uint64_t* const output_bitmasks, const int output_bitmask_size)
 {
 	// 6 inputs (query, key, value, [attn_mask], [unify head weight], [unify head bias])
-	// 3 outputs (y, [softmax_lse], [qkv])
-	if (input_size == 6 && (input_bitmasks[0] & 55u) == 55u && (output_bitmasks[0] & 5u) == 5u)
+	// 3 outputs (y, softmax_lse, [qkv])
+	if (input_size == 6 && (input_bitmasks[0] & 55u) == 55u && (output_bitmasks[0] & 7u) == 7u)
 		return 1;
-	if (input_size == 5 && (input_bitmasks[0] & 23u) == 23u && (output_bitmasks[0] & 5u) == 5u)
+	if (input_size == 5 && (input_bitmasks[0] & 23u) == 23u && (output_bitmasks[0] & 7u) == 7u)
 		return 1;
-	if ((input_bitmasks[0] & 55u) == 7u && (output_bitmasks[0] & 5u) == 1u)
+	if ((input_bitmasks[0] & 55u) == 7u && (output_bitmasks[0] & 3u) == 3u)
 		return 1;
 	return 0;
 }