[XLA:GPU/TMA] Disable TMA on B200 due to a timeout failure.

Moerafaat · Google-ML-Automation · commit ee43bb06cbb2 · 2025-11-26T09:01:56.000-08:00
PiperOrigin-RevId: 837146719
diff --git a/xla/stream_executor/gpu/tma_metadata.cc b/xla/stream_executor/gpu/tma_metadata.cc
@@ -460,11 +460,12 @@ absl::StatusOr<TmaMetadata> TmaMetadata::FromProto(
   return metadata;
 }
 
+// TODO(b/463912789): Re-enable TMA for Blackwell once the bug is fixed.
 bool IsTmaAvailableForDevice(
     const stream_executor::DeviceDescription& device_info) {
   if (auto* cuda_cc =
           device_info.gpu_compute_capability().cuda_compute_capability()) {
-    return cuda_cc->IsAtLeastHopper();
+    return cuda_cc->IsHopper();
   }
   return false;
 }

Original file line number	Diff line number	Diff line change
`@@ -460,11 +460,12 @@ absl::StatusOr<TmaMetadata> TmaMetadata::FromProto(`
`460`	`460`	`return metadata;`
`461`	`461`	`}`
`462`	`462`
	`463`	`+// TODO(b/463912789): Re-enable TMA for Blackwell once the bug is fixed.`
`463`	`464`	`bool IsTmaAvailableForDevice(`
`464`	`465`	`const stream_executor::DeviceDescription& device_info) {`
`465`	`466`	`if (auto* cuda_cc =`
`466`	`467`	`device_info.gpu_compute_capability().cuda_compute_capability()) {`
`467`		`- return cuda_cc->IsAtLeastHopper();`
	`468`	`+ return cuda_cc->IsHopper();`
`468`	`469`	`}`
`469`	`470`	`return false;`
`470`	`471`	`}`