[XLA:GPU]: Limit the amount of variadics we support for collectives

Caps variadic all-gather and reduce-scatter kernels at 16 operands.
This covers most cases and avoids pathological collectives with 512 operands or something.

PiperOrigin-RevId: 1003933042

third_party/xla/xla/backends/gpu/runtime/all_gather.cc

@@ -153,6 +153,11 @@ absl::StatusOr<AllGatherInfo> BuildAllGatherInfo(
   }
   const int64_t num_devices =
       all_gather->device_list()->num_devices_per_group();
+  if (all_gather->operand_count() > kAllGatherMaxNumOperands) {
+    return absl::UnimplementedError(absl::StrFormat(
+        "All-gather kernel supports at most %d operands. Got %d.",
+        kAllGatherMaxNumOperands, all_gather->operand_count()));
+  }
   if (device_info.device_interconnect_info().active_links <= 0) {
     return absl::UnimplementedError(
         "Collective kernels are only supported on devices with NVLink/UALink "

third_party/xla/xla/backends/gpu/runtime/all_gather.h

@@ -48,6 +48,10 @@ inline constexpr int64_t kAllGatherMaxBlocksPerGrid = 64;
 // Base on the experimental results.
 inline constexpr int64_t kMaxAllGatherSizeBytes = 16 * 1024 * 1024;  // 16 MB
 
+// Maximum number of operands supported by the variadic Triton all-gather
+// kernel.
+inline constexpr int64_t kAllGatherMaxNumOperands = 16;
+
 // Encapsulates the information needed to perform an all-gather via the Triton
 // collective kernel backend.
 struct AllGatherInfo {

third_party/xla/xla/backends/gpu/runtime/all_gather_build_info_test.cc

@@ -300,6 +300,50 @@ TEST_F(BuildAllGatherInfoTest, VariadicChecksSizePerOperandAndUsesMaxElements) {
                 Field(&AllGatherInfo::max_output_elements, 4 * 1024 * 1024))));
 }
 
+TEST_F(BuildAllGatherInfoTest, FailsForMoreThanMaxOperands) {
+  std::vector<std::string> params;
+  std::vector<std::string> param_names;
+  std::vector<std::string> output_shapes;
+  for (int i = 0; i <= kAllGatherMaxNumOperands; ++i) {
+    params.push_back(absl::StrFormat("param_%d = f32[8] parameter(%d)", i, i));
+    param_names.push_back(absl::StrFormat("param_%d", i));
+    output_shapes.push_back("f32[16]");
+  }
+  const std::string module_str = absl::StrFormat(
+      R"(
+  HloModule test
+  ENTRY test_computation {
+    %s
+    ROOT all-gather = (%s) all-gather(%s),
+        dimensions={0}, replica_groups={{0,1}}
+  }
+  )",
+      absl::StrJoin(params, "\n"), absl::StrJoin(output_shapes, ", "),
+      absl::StrJoin(param_names, ", "));
+  se::DeviceDescription device_info = TestGpuDeviceInfo::H100SXMDeviceInfo();
+  stream_executor::GpuTargetConfigProto target_config_proto;
+  *target_config_proto.mutable_gpu_device_info() = device_info.ToProto();
+  target_config_proto.mutable_gpu_device_info()
+      ->mutable_device_interconnect_info()
+      ->set_active_links(18);
+  target_config_proto.set_platform_name("CUDA");
+  ASSERT_OK_AND_ASSIGN(gpu::GpuTargetConfig target_config,
+                       gpu::GpuTargetConfig::FromProto(target_config_proto));
+  GpuTopology gpu_topology("platform_version", /*num_partitions=*/1,
+                           /*num_hosts_per_partition=*/1,
+                           /*num_devices_per_host=*/2, target_config);
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       ParseAndReturnVerifiedModule(module_str, 2));
+  const HloInstruction* hlo_instr =
+      HloHardwareIndependentTestBase::FindInstruction(module.get(),
+                                                      HloOpcode::kAllGather);
+  EXPECT_THAT(
+      BuildAllGatherInfo(gpu_topology, Cast<HloAllGatherInstruction>(hlo_instr),
+                         /*device_assignment=*/nullptr),
+      StatusIs(absl::StatusCode::kUnimplemented,
+               HasSubstr("supports at most 16 operands")));
+}
+
 TEST_F(BuildAllGatherInfoTest,
        CreateAllGatherKernelSpecMatchesAllReduceArgumentLayout) {
   constexpr absl::string_view kModuleStr = R"(

third_party/xla/xla/backends/gpu/runtime/reduce_scatter.cc

@@ -210,6 +210,11 @@ absl::StatusOr<ReduceScatterInfo> BuildReduceScatterInfo(
     return absl::UnimplementedError(
         "Reduce-scatter kernel requires at least one operand.");
   }
+  if (num_operands > kReduceScatterMaxNumOperands) {
+    return absl::UnimplementedError(absl::StrFormat(
+        "Reduce-scatter kernel supports at most %d operands. Got %d.",
+        kReduceScatterMaxNumOperands, num_operands));
+  }
   if (device_info.device_interconnect_info().active_links <= 0) {
     return absl::UnimplementedError(
         "Collective kernels are only supported on devices with NVLink "

third_party/xla/xla/backends/gpu/runtime/reduce_scatter.h

@@ -35,6 +35,10 @@ inline constexpr int64_t kMaxReduceScatterSizeBytes =
 // Maximum number of blocks per grid for the Triton reduce-scatter kernel.
 inline constexpr int64_t kReduceScatterMaxBlocksPerGrid = 64;
 
+// Maximum number of operands supported by the variadic Triton reduce-scatter
+// kernel.
+inline constexpr int64_t kReduceScatterMaxNumOperands = 16;
+
 // Encapsulates the information needed to perform a reduce-scatter via the
 // Triton collective kernel backend.
 struct ReduceScatterInfo {

third_party/xla/xla/backends/gpu/runtime/reduce_scatter_build_info_test.cc

@@ -407,6 +407,54 @@ TEST_F(BuildReduceScatterInfoTest, FailsForVariadicWithMixedElementTypes) {
                        HasSubstr("same element type")));
 }
 
+TEST_F(BuildReduceScatterInfoTest, FailsForMoreThanMaxOperands) {
+  std::vector<std::string> params;
+  std::vector<std::string> param_names;
+  std::vector<std::string> output_shapes;
+  for (int i = 0; i <= kReduceScatterMaxNumOperands; ++i) {
+    params.push_back(absl::StrFormat("param_%d = f32[16] parameter(%d)", i, i));
+    param_names.push_back(absl::StrFormat("param_%d", i));
+    output_shapes.push_back("f32[8]");
+  }
+  const std::string module_str = absl::StrFormat(
+      R"(
+  HloModule test
+  apply_op {
+    x = f32[] parameter(0)
+    y = f32[] parameter(1)
+    ROOT apply_op = f32[] add(x, y)
+  }
+  ENTRY test_computation {
+    %s
+    ROOT reduce-scatter = (%s) reduce-scatter(%s),
+        dimensions={0}, replica_groups={{0,1}}, to_apply=apply_op
+  }
+  )",
+      absl::StrJoin(params, "\n"), absl::StrJoin(output_shapes, ", "),
+      absl::StrJoin(param_names, ", "));
+  se::DeviceDescription device_info = TestGpuDeviceInfo::H100SXMDeviceInfo();
+  stream_executor::GpuTargetConfigProto target_config_proto;
+  *target_config_proto.mutable_gpu_device_info() = device_info.ToProto();
+  target_config_proto.mutable_gpu_device_info()
+      ->mutable_device_interconnect_info()
+      ->set_active_links(18);
+  target_config_proto.set_platform_name("CUDA");
+  ASSERT_OK_AND_ASSIGN(gpu::GpuTargetConfig target_config,
+                       gpu::GpuTargetConfig::FromProto(target_config_proto));
+  GpuTopology gpu_topology("platform_version", /*num_partitions=*/1,
+                           /*num_hosts_per_partition=*/1,
+                           /*num_devices_per_host=*/2, target_config);
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       ParseAndReturnVerifiedModule(module_str, 2));
+  const auto* rs = Cast<HloReduceScatterInstruction>(
+      HloHardwareIndependentTestBase::FindInstruction(
+          module.get(), HloOpcode::kReduceScatter));
+  EXPECT_THAT(BuildReduceScatterInfo(gpu_topology, rs,
+                                     /*device_assignment=*/nullptr),
+              StatusIs(absl::StatusCode::kUnimplemented,
+                       HasSubstr("supports at most 16 operands")));
+}
+
 TEST_F(BuildReduceScatterInfoTest, FailsForLargeInputs) {
   // 4 * 1024 * 1024 F32 output elements * 2 devices = 32 MB input > 16 MB
   // limit -> unimplemented.