mirrors/tensorflow
0
0 Fork 0
mirror of https://github.com/tensorflow/tensorflow.git synced 2026-10-26 00:09:51 +08:00
Code Activity

Automated Code Change

Browse Source 
PiperOrigin-RevId: 1006165887
This commit is contained in:
A. Unique TensorFlower authored and TensorFlower Gardener committed 2026-10-24 01:25:34 -07:00
1 parent c0fc2ea20b
commit 14665b7f0b
1 file changed
+12 -11
Download Patch File Download Diff File Expand all files Collapse all files
third_party/xla/xla/backends/gpu/transforms/cudnn_custom_call_compiler.cc
Vendored
+12 -11
View File
@@ -358,7 +358,7 @@ absl::StatusOr<se::gpu::CudnnGraph> BuildGraphForCustomCallToBackwardFMHA(
auto computations = custom_call->called_computations();
const HloComputation* score_mod_fwd_comp = nullptr;
const HloComputation* score_mod_bwd_comp = nullptr;
stream_executor::gpu::ScoreModFunc* score_mod = nullptr;
std::optional<stream_executor::gpu::ScoreModFunc> score_mod;
TF_RET_CHECK(computations.size() <= 1);
if (computations.size() == 1) {
score_mod_bwd_comp = computations[0];
@@ -369,20 +369,21 @@ absl::StatusOr<se::gpu::CudnnGraph> BuildGraphForCustomCallToBackwardFMHA(
return absl::InternalError("Can't find fmha fwd custom call.");
}
score_mod_fwd_comp = fwd_custom_call->called_computations()[0];
auto smf = stream_executor::gpu::ScoreModFunc(score_mod_fwd_comp,
score_mod_bwd_comp);
score_mod = &smf;
score_mod.emplace(score_mod_fwd_comp, score_mod_bwd_comp);
input_index += score_mod_fwd_comp->num_parameters() - 1;
}
stream_executor::gpu::ScoreModFunc* score_mod_ptr =
score_mod.has_value() ? &score_mod.value() : nullptr;
TF_RET_CHECK(input_index == custom_call->operand_count());
ABSL_ASSIGN_OR_RETURN(se::gpu::CudnnGraph graph,
se::gpu::GetCudnnFlashAttentionBackwardOperationGraph(
dnn_support, gpu_device_info, q, k, p, v, d_output, dq,
dk, dv, bias, dbias, dropout_rate, config.seed(),
config.fmha_scale(), dropout_rate > 0.0,
bias.has_value(), dnn_mask_type, force_deterministic,
sliding_window_length, max_seg_per_batch, score_mod));
ABSL_ASSIGN_OR_RETURN(
se::gpu::CudnnGraph graph,
se::gpu::GetCudnnFlashAttentionBackwardOperationGraph(
dnn_support, gpu_device_info, q, k, p, v, d_output, dq, dk, dv, bias,
dbias, dropout_rate, config.seed(), config.fmha_scale(),
dropout_rate > 0.0, bias.has_value(), dnn_mask_type,
force_deterministic, sliding_window_length, max_seg_per_batch,
score_mod_ptr));
return graph;
}