/*************************************************************************
 * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 *
 * See LICENSE for license information.
 ************************************************************************/

#include <cublasLt.h>
#include <cuda_bf16.h>
#include <cuda_runtime.h>
#include <gtest/gtest.h>

#include <algorithm>
#include <memory>
#include <numeric>
#include <optional>
#include <random>
#include <string>
#include <tuple>
#include <vector>

#include <transformer_engine/cast.h>
#include <transformer_engine/gemm.h>
#include <transformer_engine/recipe.h>
#include <transformer_engine/swizzle.h>
#include <transformer_engine/transformer_engine.h>

#include "../test_common.h"
#include "util/cuda_runtime.h"

using namespace transformer_engine;
using namespace test;

namespace {

enum class InputRecipe {
  kBF16,
  kFP8Current,
  kMXFP8,
  kNVFP4,
  kFP8BlockScaling1D1D,
  kFP8BlockScaling2D1D,
  kFP8BlockScaling1D2D,
};

inline const char* recipe_name(InputRecipe recipe) {
  switch (recipe) {
    case InputRecipe::kBF16:                 return "BF16";
    case InputRecipe::kFP8Current:           return "FP8Current";
    case InputRecipe::kMXFP8:                return "MXFP8";
    case InputRecipe::kNVFP4:                return "NVFP4";
    case InputRecipe::kFP8BlockScaling1D1D:  return "FP8BlockScaling1D1D";
    case InputRecipe::kFP8BlockScaling2D1D:  return "FP8BlockScaling2D1D";
    case InputRecipe::kFP8BlockScaling1D2D:  return "FP8BlockScaling1D2D";
    default:                                 return "Unknown";
  }
}

inline bool is_fp8_block_recipe(InputRecipe recipe) {
  return recipe == InputRecipe::kFP8BlockScaling1D1D ||
         recipe == InputRecipe::kFP8BlockScaling2D1D ||
         recipe == InputRecipe::kFP8BlockScaling1D2D;
}

inline NVTEScalingMode a_scaling_mode(InputRecipe recipe) {
  switch (recipe) {
    case InputRecipe::kFP8Current:           return NVTE_DELAYED_TENSOR_SCALING;
    case InputRecipe::kMXFP8:                return NVTE_MXFP8_1D_SCALING;
    case InputRecipe::kNVFP4:                return NVTE_NVFP4_1D_SCALING;
    case InputRecipe::kFP8BlockScaling2D1D:  return NVTE_BLOCK_SCALING_2D;
    case InputRecipe::kFP8BlockScaling1D1D:
    case InputRecipe::kFP8BlockScaling1D2D:  return NVTE_BLOCK_SCALING_1D;
    case InputRecipe::kBF16:                 return NVTE_DELAYED_TENSOR_SCALING;
    default:                                 return NVTE_DELAYED_TENSOR_SCALING;
  }
}

inline NVTEScalingMode b_scaling_mode(InputRecipe recipe) {
  switch (recipe) {
    case InputRecipe::kFP8BlockScaling1D2D:  return NVTE_BLOCK_SCALING_2D;
    case InputRecipe::kFP8BlockScaling2D1D:  return NVTE_BLOCK_SCALING_1D;
    default:                                 return a_scaling_mode(recipe);
  }
}

// Mul128 cases use dims that are multiples of 128 - full functionality across all recipes.
// kAllSameMul32 uses dims that are multiples of 32 but not 128; for recipes with padded
// grouped scale_inv storage, it exercises per-expert scale_inv offsets.
enum class ShapeCase {
  kAllSameMul128,
  kSameFirstMul128,
  kSameLastMul128,
  kAllDifferentMul128,
  kAllSameMul32,
};

Tensor make_fp8_operand(const std::string& name, const std::vector<size_t>& shape) {
  Tensor input_fp32(name + "_fp32", shape, DType::kFloat32);

  const size_t numel = shape[0] * shape[1];
  std::vector<float> data(numel);
  std::mt19937 gen(std::hash<std::string>{}(name));
  // Random mean and stddev -> different amax per tensor -> different scales
  std::uniform_real_distribution<float> param_dis(0.1f, 10.0f);
  float mean = param_dis(gen);
  float stddev = param_dis(gen);
  std::normal_distribution<float> dis(mean, stddev);
  for (size_t i = 0; i < numel; ++i) {
    data[i] = dis(gen);
  }
  NVTE_CHECK_CUDA(cudaMemcpy(input_fp32.rowwise_dptr(), data.data(),
                             numel * sizeof(float), cudaMemcpyHostToDevice));

  Tensor fp8(name, shape, TypeInfo<fp8e4m3>::dtype, true, true, NVTE_DELAYED_TENSOR_SCALING);

  nvte_compute_amax(input_fp32.data(), fp8.data(), 0);
  QuantizationConfigWrapper config;
  nvte_compute_scale_from_amax(fp8.data(), config, 0);
  nvte_quantize(input_fp32.data(), fp8.data(), 0);
  return fp8;
}

Tensor make_bf16_operand(const std::string& name, const std::vector<size_t>& shape) {
  Tensor t(name, shape, DType::kBFloat16);
  const size_t numel = shape[0] * shape[1];
  std::vector<__nv_bfloat16> ones(numel, __float2bfloat16(1.0f));
  NVTE_CHECK_CUDA(cudaMemcpy(t.rowwise_dptr(), ones.data(),
                             numel * sizeof(__nv_bfloat16), cudaMemcpyHostToDevice));
  return t;
}

// Creates an MXFP8 operand with the given single direction (scales along K dimension).
Tensor make_mxfp8_operand(const std::string& name, const std::vector<size_t>& shape,
                          bool use_rowwise) {
  Tensor input_bf16(name + "_bf16", shape, DType::kBFloat16);
  fillUniform(&input_bf16);

  Tensor mxfp8(name, shape, TypeInfo<fp8e4m3>::dtype, use_rowwise, !use_rowwise,
               NVTE_MXFP8_1D_SCALING);
  nvte_quantize(input_bf16.data(), mxfp8.data(), 0);

  Tensor mxfp8_swizzled(name + "_swizzled", shape, TypeInfo<fp8e4m3>::dtype,
                        use_rowwise, !use_rowwise, NVTE_MXFP8_1D_SCALING);
  mxfp8_swizzled.set_with_gemm_swizzled_scales(true);  // Must be set BEFORE swizzle call

  const size_t data_bytes = test::bytes(
      use_rowwise ? mxfp8.rowwise_shape() : mxfp8.columnwise_shape(), mxfp8.dtype());
  void* dst = use_rowwise ? mxfp8_swizzled.rowwise_dptr() : mxfp8_swizzled.columnwise_dptr();
  void* src = use_rowwise ? mxfp8.rowwise_dptr() : mxfp8.columnwise_dptr();
  NVTE_CHECK_CUDA(cudaMemcpy(dst, src, data_bytes, cudaMemcpyDeviceToDevice));

  nvte_swizzle_scaling_factors(mxfp8.data(), mxfp8_swizzled.data(), 0);
  NVTE_CHECK_CUDA(cudaDeviceSynchronize());
  return mxfp8_swizzled;
}

// Creates an NVFP4 operand with the given single direction, swizzled scales.
Tensor make_nvfp4_operand(const std::string& name, const std::vector<size_t>& shape,
                          bool use_rowwise) {
  Tensor input_bf16(name + "_bf16", shape, DType::kBFloat16);
  fillUniform(&input_bf16);

  Tensor nvfp4(name, shape, DType::kFloat4E2M1, use_rowwise, !use_rowwise,
               NVTE_NVFP4_1D_SCALING);
  QuantizationConfigWrapper quant_config;
  nvte_quantize_v2(input_bf16.data(), nvfp4.data(), quant_config, 0);

  Tensor nvfp4_sw(name + "_sw", shape, DType::kFloat4E2M1, use_rowwise, !use_rowwise,
                  NVTE_NVFP4_1D_SCALING);
  nvfp4_sw.set_with_gemm_swizzled_scales(true);

  // Copy quantized data + amax to swizzled tensor (swizzle only rewrites scale_inv).
  const auto amax_kind = use_rowwise ? kNVTEAmax : kNVTEColumnwiseAmax;
  const NVTEBasicTensor src_amax = nvte_get_tensor_param(nvfp4.data(), amax_kind);
  const NVTEBasicTensor dst_amax = nvte_get_tensor_param(nvfp4_sw.data(), amax_kind);
  NVTE_CHECK_CUDA(cudaMemcpy(dst_amax.data_ptr, src_amax.data_ptr, sizeof(float),
                             cudaMemcpyDeviceToDevice));
  const size_t data_bytes = test::bytes(
      use_rowwise ? nvfp4.rowwise_shape() : nvfp4.columnwise_shape(), nvfp4.dtype());
  void* dst_data = use_rowwise ? nvfp4_sw.rowwise_dptr() : nvfp4_sw.columnwise_dptr();
  void* src_data = use_rowwise ? nvfp4.rowwise_dptr() : nvfp4.columnwise_dptr();
  NVTE_CHECK_CUDA(cudaMemcpy(dst_data, src_data, data_bytes, cudaMemcpyDeviceToDevice));

  nvte_swizzle_scaling_factors(nvfp4.data(), nvfp4_sw.data(), 0);
  NVTE_CHECK_CUDA(cudaDeviceSynchronize());
  return nvfp4_sw;
}

// Creates an FP8 block-scaling operand with the given single direction (TN-only on Hopper).
Tensor make_fp8_block_scaling_operand(const std::string& name, const std::vector<size_t>& shape,
                                      bool use_rowwise,
                                      NVTEScalingMode scaling_mode = NVTE_BLOCK_SCALING_1D) {
  Tensor input_bf16(name + "_bf16", shape, DType::kBFloat16);
  fillUniform(&input_bf16);

  Tensor fp8_bs(name, shape, TypeInfo<fp8e4m3>::dtype, use_rowwise, !use_rowwise,
                scaling_mode);
  QuantizationConfigWrapper quant_config;
  quant_config.set_force_pow_2_scales(true);
  nvte_quantize_v2(input_bf16.data(), fp8_bs.data(), quant_config, 0);
  NVTE_CHECK_CUDA(cudaDeviceSynchronize());
  return fp8_bs;
}

struct TestParams {
  InputRecipe recipe;
  bool transa;
  bool transb;
  ShapeCase shape_case;
  bool use_null_c = false;  // When true, pass nullptr for C (valid when beta=0)
  DType output_dtype = DType::kBFloat16;  // Implementation also accepts FP16 / FP32.
};

// Returns a vector of (M, N, K) tuples for each GEMM in the group.
// M - number of rows in output D
// N - number of columns in output D
// K - reduction dimension shared between A and B
std::vector<std::tuple<size_t, size_t, size_t>> make_shapes(ShapeCase scase) {
  switch (scase) {
    case ShapeCase::kAllSameMul128:
      return {{128, 256, 384}, {128, 256, 384}, {128, 256, 384}};
    case ShapeCase::kSameFirstMul128:
      // Same M (first dim), varying N and K
      return {{128, 256, 384}, {128, 384, 512}, {128, 512, 640}};
    case ShapeCase::kSameLastMul128:
      // Same N (last dim), varying M and K
      return {{128, 256, 384}, {256, 256, 512}, {384, 256, 640}};
    case ShapeCase::kAllDifferentMul128:
      return {{128, 256, 384}, {256, 384, 512}, {384, 512, 640}};
    case ShapeCase::kAllSameMul32:
    default:
      return {{160, 288, 416}, {160, 288, 416}, {160, 288, 416}};
  }
}

constexpr size_t kCublasGroupedGemmVersion = 130300;        // Blackwell-only grouped GEMM
constexpr size_t kCublasGroupedGemmHopperVersion = 130400;  // adds Hopper support
// FP8 per-tensor scaling grouped GEMM on Hopper requires cuBLAS 13.5+
constexpr size_t kCublasFp8TensorScalingGroupedGemmHopperVersion = 130500;
constexpr size_t kCublasWorkspaceBytes = 32ull * 1024 * 1024;

inline std::string grouped_gemm_skip_reason(const TestParams& params) {
  const size_t cublas_ver = transformer_engine::cuda::cublas_version();
  if (cublas_ver < kCublasGroupedGemmVersion) {
    return "Grouped GEMM requires cuBLAS 13.3+, but run-time cuBLAS version is " +
           std::to_string(cublas_ver) + ".";
  }
  const int32_t cc = getDeviceComputeCapability();
  const std::string cc_suffix =
      "but device compute capability is " + std::to_string(cc) + ".";
  if (cc < hopperComputeCapability) {
    return "Grouped GEMM requires Hopper (SM90) or newer, " + cc_suffix;
  }
  if (cc < blackwellComputeCapability && cublas_ver < kCublasGroupedGemmHopperVersion) {
    return "Grouped GEMM on Hopper (SM90) requires cuBLAS 13.4+, but run-time cuBLAS "
           "version is " + std::to_string(cublas_ver) + ".";
  }
  if (params.recipe != InputRecipe::kBF16) {
    const bool is_blackwell_plus = cc >= blackwellComputeCapability;
    const bool fp8_block = is_fp8_block_recipe(params.recipe);
    if (!is_blackwell_plus && !fp8_block) {
      if (params.recipe == InputRecipe::kFP8Current) {
        if (cublas_ver < kCublasFp8TensorScalingGroupedGemmHopperVersion) {
          return "FP8 per-tensor scaling grouped GEMM on Hopper (SM90) requires cuBLAS 13.5+, "
                 "but run-time cuBLAS version is " + std::to_string(cublas_ver) + ".";
        }
      } else {
        return std::string(recipe_name(params.recipe)) +
               " grouped GEMM requires Blackwell (SM100) or newer, " + cc_suffix;
      }
    }
    if (is_blackwell_plus && fp8_block) {
      return "FP8 block scaling grouped GEMM is only supported on Hopper (SM90), " + cc_suffix;
    }
    if (params.recipe == InputRecipe::kNVFP4 && params.output_dtype == DType::kFloat16) {
      return "NVFP4 grouped GEMM does not support FP16 output.";
    }
  }
  return "";
}

// Reference setup shared by the three run_* variants: builds A/B/D tensors per recipe,
// runs nvte_multi_tensor_gemm to fill D_multi with reference results, and keeps the
// workspaces alive (returned in the struct so callers don't have to track them).
// Output dtype comes from TestParams::output_dtype (BF16 / FP16 / FP32).
struct GroupedGemmRefSetup {
  std::vector<std::tuple<size_t, size_t, size_t>> shapes;
  size_t num_gemms = 0;
  std::vector<Tensor> A_tensors;
  std::vector<Tensor> B_tensors;
  std::vector<Tensor> D_multi;
  std::vector<Tensor> workspaces;
  bool use_split_accum = false;
};

inline GroupedGemmRefSetup make_grouped_gemm_ref(const TestParams& params) {
  GroupedGemmRefSetup s;
  s.shapes = make_shapes(params.shape_case);
  s.num_gemms = s.shapes.size();
  s.A_tensors.reserve(s.num_gemms);
  s.B_tensors.reserve(s.num_gemms);
  s.D_multi.reserve(s.num_gemms);

  for (size_t i = 0; i < s.num_gemms; ++i) {
    const auto [M, N, K] = s.shapes[i];
    const std::vector<size_t> a_shape =
        params.transa ? std::vector<size_t>{N, K} : std::vector<size_t>{K, N};
    const std::vector<size_t> b_shape =
        params.transb ? std::vector<size_t>{K, M} : std::vector<size_t>{M, K};
    if (params.recipe == InputRecipe::kBF16) {
      s.A_tensors.emplace_back(make_bf16_operand("A" + std::to_string(i), a_shape));
      s.B_tensors.emplace_back(make_bf16_operand("B" + std::to_string(i), b_shape));
    } else {
      const bool a_use_rowwise = params.transa;
      const bool b_use_rowwise = !params.transb;
      switch (params.recipe) {
        case InputRecipe::kFP8Current:
          s.A_tensors.emplace_back(make_fp8_operand("A" + std::to_string(i), a_shape));
          s.B_tensors.emplace_back(make_fp8_operand("B" + std::to_string(i), b_shape));
          break;
        case InputRecipe::kMXFP8:
          s.A_tensors.emplace_back(make_mxfp8_operand("A" + std::to_string(i), a_shape,
                                                      a_use_rowwise));
          s.B_tensors.emplace_back(make_mxfp8_operand("B" + std::to_string(i), b_shape,
                                                      b_use_rowwise));
          break;
        case InputRecipe::kNVFP4:
          s.A_tensors.emplace_back(make_nvfp4_operand("A" + std::to_string(i), a_shape,
                                                      a_use_rowwise));
          s.B_tensors.emplace_back(make_nvfp4_operand("B" + std::to_string(i), b_shape,
                                                      b_use_rowwise));
          break;
        case InputRecipe::kFP8BlockScaling1D1D:
        case InputRecipe::kFP8BlockScaling2D1D:
        case InputRecipe::kFP8BlockScaling1D2D:
          s.A_tensors.emplace_back(make_fp8_block_scaling_operand("A" + std::to_string(i),
                                                                  a_shape, a_use_rowwise,
                                                                  a_scaling_mode(params.recipe)));
          s.B_tensors.emplace_back(make_fp8_block_scaling_operand("B" + std::to_string(i),
                                                                  b_shape, b_use_rowwise,
                                                                  b_scaling_mode(params.recipe)));
          break;
        default:
          NVTE_ERROR("Unsupported scaling mode in grouped GEMM test: " +
                     std::string(recipe_name(params.recipe)));
      }
    }
    s.D_multi.emplace_back(Tensor("D_multi" + std::to_string(i),
                                  std::vector<size_t>{M, N}, params.output_dtype));
  }

  // FP8 block scaling requires split accumulator (no fast accumulation).
  s.use_split_accum = is_fp8_block_recipe(params.recipe);

  std::vector<NVTETensor> A_ptrs(s.num_gemms), B_ptrs(s.num_gemms), D_ptrs(s.num_gemms);
  std::vector<NVTETensor> workspace_ptrs(s.num_gemms, nullptr);
  std::vector<NVTETensor> bias_ptrs(s.num_gemms, nullptr), gelu_ptrs(s.num_gemms, nullptr);
  s.workspaces.reserve(s.num_gemms);
  for (size_t i = 0; i < s.num_gemms; ++i) {
    A_ptrs[i] = s.A_tensors[i].data();
    B_ptrs[i] = s.B_tensors[i].data();
    D_ptrs[i] = s.D_multi[i].data();
    s.workspaces.emplace_back(Tensor("workspace" + std::to_string(i),
                                     std::vector<size_t>{kCublasWorkspaceBytes}, DType::kByte));
    workspace_ptrs[i] = s.workspaces.back().data();
  }
  nvte_multi_tensor_gemm(A_ptrs.data(), B_ptrs.data(), D_ptrs.data(), bias_ptrs.data(),
                         gelu_ptrs.data(), static_cast<int>(s.num_gemms),
                         params.transa, params.transb, false, workspace_ptrs.data(),
                         false, s.use_split_accum, 0, 0);
  NVTE_CHECK_CUDA(cudaDeviceSynchronize());
  return s;
}

// Allocate and initialize alpha/beta tensors for grouped GEMM.
// Hopper requires a single shared scalar; Blackwell+ uses per-matrix scalars.
struct AlphaBetaTensors {
  Tensor alpha;
  Tensor beta;
};

inline AlphaBetaTensors make_alpha_beta(size_t num_gemms) {
  const int32_t cc = getDeviceComputeCapability();
  const size_t n = cc < blackwellComputeCapability ? 1 : num_gemms;
  AlphaBetaTensors ab{Tensor("alpha", std::vector<size_t>{n}, DType::kFloat32),
                      Tensor("beta", std::vector<size_t>{n}, DType::kFloat32)};
  std::vector<float> a(n, 1.f);
  std::vector<float> b(n, 0.f);
  NVTE_CHECK_CUDA(cudaMemcpy(ab.alpha.rowwise_dptr(), a.data(), n * sizeof(float),
                             cudaMemcpyHostToDevice));
  NVTE_CHECK_CUDA(cudaMemcpy(ab.beta.rowwise_dptr(), b.data(), n * sizeof(float),
                             cudaMemcpyHostToDevice));
  return ab;
}

// Compare each tensor inside a grouped D buffer (with per-tensor offsets) against the
// reference D_multi[i] tensors.
// Capture `body(stream)` into a CUDA graph and replay it `n_replays` times, calling
// `verify(iter)` after each launch+sync. Used to assert that the grouped GEMM kernels
// are graph-safe (capture succeeds, replay is correct, and replays are deterministic).
template <typename Body, typename Verify>
inline void capture_and_replay_grouped_gemm(cudaStream_t stream, int n_replays,
                                            Body&& body, Verify&& verify) {
  // Warmup off-graph so cuBLASLt can initialize its internal heuristic / handle state
  // outside capture (the very first matmul on a handle may do host-side setup that
  // isn't capturable in some cuBLAS versions).
  body(stream);
  NVTE_CHECK_CUDA(cudaStreamSynchronize(stream));

  NVTE_CHECK_CUDA(cudaStreamBeginCapture(stream, cudaStreamCaptureModeRelaxed));
  body(stream);
  cudaGraph_t graph;
  NVTE_CHECK_CUDA(cudaStreamEndCapture(stream, &graph));

  cudaGraphExec_t exec;
  NVTE_CHECK_CUDA(cudaGraphInstantiate(&exec, graph, nullptr, nullptr, 0));

  for (int i = 0; i < n_replays; ++i) {
    NVTE_CHECK_CUDA(cudaGraphLaunch(exec, stream));
    NVTE_CHECK_CUDA(cudaStreamSynchronize(stream));
    verify(i);
  }

  NVTE_CHECK_CUDA(cudaGraphExecDestroy(exec));
  NVTE_CHECK_CUDA(cudaGraphDestroy(graph));
}

inline void compare_grouped_d_to_multi(
    const GroupedBuffers& grouped_D,
    const std::vector<std::tuple<size_t, size_t, size_t>>& shapes,
    std::vector<Tensor>& D_multi, const char* tag) {
  for (size_t i = 0; i < shapes.size(); ++i) {
    Tensor grouped_split("grouped_D" + std::to_string(i),
                         std::vector<size_t>{static_cast<size_t>(std::get<0>(shapes[i])),
                                             static_cast<size_t>(std::get<1>(shapes[i]))},
                         D_multi[i].dtype());
    const size_t offset_bytes =
        static_cast<size_t>(grouped_D.offsets_host[i]) * grouped_D.elem_size;
    NVTE_CHECK_CUDA(cudaMemcpy(grouped_split.rowwise_dptr(),
                               static_cast<char*>(grouped_D.get_data()) + offset_bytes,
                               grouped_D.tensor_bytes[i], cudaMemcpyDeviceToDevice));
    grouped_split.to_cpu();
    D_multi[i].to_cpu();
    auto [atol, rtol] = getTolerances(D_multi[i].dtype());
    switch (D_multi[i].dtype()) {
      case DType::kBFloat16:
        compareResults(tag, grouped_split, D_multi[i].rowwise_cpu_dptr<bf16>(), true, atol, rtol);
        break;
      case DType::kFloat16:
        compareResults(tag, grouped_split, D_multi[i].rowwise_cpu_dptr<fp16>(), true, atol, rtol);
        break;
      case DType::kFloat32:
        compareResults(tag, grouped_split, D_multi[i].rowwise_cpu_dptr<float>(), true, atol, rtol);
        break;
      default:
        NVTE_ERROR("Unsupported D dtype in test: " +
                   std::to_string(static_cast<int>(D_multi[i].dtype())));
    }
  }
}

void run_grouped_gemm_case(const TestParams& params) {
  if (auto reason = grouped_gemm_skip_reason(params); !reason.empty()) {
    GTEST_SKIP() << reason;
  }
  auto ref = make_grouped_gemm_ref(params);
  const auto& shapes = ref.shapes;
  const size_t num_gemms = ref.num_gemms;

  std::vector<Tensor*> A_views, B_views;
  A_views.reserve(num_gemms);
  B_views.reserve(num_gemms);
  for (size_t i = 0; i < num_gemms; ++i) {
    A_views.push_back(&ref.A_tensors[i]);
    B_views.push_back(&ref.B_tensors[i]);
  }

  GroupedBuffers grouped_A = build_grouped_tensor(A_views, ref.A_tensors[0].scaling_mode());
  GroupedBuffers grouped_B = build_grouped_tensor(B_views, ref.B_tensors[0].scaling_mode());

  std::vector<Tensor> C_tensors;
  std::vector<Tensor> D_group_tensors;
  C_tensors.reserve(num_gemms);
  D_group_tensors.reserve(num_gemms);
  for (size_t i = 0; i < num_gemms; ++i) {
    const auto [M, N, K] = shapes[i];
    (void)K;
    if (!params.use_null_c) {
      C_tensors.emplace_back(Tensor("C" + std::to_string(i),
                                    std::vector<size_t>{static_cast<size_t>(M), static_cast<size_t>(N)},
                                    params.output_dtype));
    }
    D_group_tensors.emplace_back(Tensor("D_group" + std::to_string(i),
                                        std::vector<size_t>{static_cast<size_t>(M), static_cast<size_t>(N)},
                                        params.output_dtype));
    NVTE_CHECK_CUDA(cudaMemset(D_group_tensors.back().rowwise_dptr(), 0, bytes(D_group_tensors.back().rowwise_shape(), D_group_tensors.back().dtype())));
  }

  std::vector<Tensor*> C_views, D_views;
  for (size_t i = 0; i < num_gemms; ++i) {
    if (!params.use_null_c) {
      C_views.push_back(&C_tensors[i]);
    }
    D_views.push_back(&D_group_tensors[i]);
  }

  std::optional<GroupedBuffers> grouped_C;
  if (!params.use_null_c) {
    grouped_C = build_grouped_tensor(C_views, NVTE_DELAYED_TENSOR_SCALING);
  }
  GroupedBuffers grouped_D = build_grouped_tensor(D_views, NVTE_DELAYED_TENSOR_SCALING);

  AlphaBetaTensors ab = make_alpha_beta(num_gemms);

  const size_t setup_ws_bytes = nvte_get_grouped_gemm_setup_workspace_size(num_gemms);
  Tensor setup_ws("setup_ws", std::vector<size_t>{setup_ws_bytes}, DType::kByte);
  Tensor cublas_ws("cublas_ws", std::vector<size_t>{kCublasWorkspaceBytes}, DType::kByte);

  GroupedMatmulConfigWrapper grouped_config;
  if (ref.use_split_accum) {
    grouped_config.set_use_split_accumulator(true);
  }

  nvte_grouped_gemm(grouped_A.get_handle(), params.transa, grouped_B.get_handle(), params.transb,
                    params.use_null_c ? nullptr : grouped_C->get_handle(), grouped_D.get_handle(),
                    ab.alpha.data(), ab.beta.data(), setup_ws.data(), cublas_ws.data(),
                    grouped_config, 0);
  NVTE_CHECK_CUDA(cudaDeviceSynchronize());

  compare_grouped_d_to_multi(grouped_D, shapes, ref.D_multi, "grouped_vs_multi");
}

void run_grouped_gemm_discrete_out_case(const TestParams& params) {
  if (auto reason = grouped_gemm_skip_reason(params); !reason.empty()) {
    GTEST_SKIP() << reason;
  }
  auto ref = make_grouped_gemm_ref(params);
  const auto& shapes = ref.shapes;
  const size_t num_gemms = ref.num_gemms;

  std::vector<Tensor*> A_views, B_views;
  A_views.reserve(num_gemms);
  B_views.reserve(num_gemms);
  for (size_t i = 0; i < num_gemms; ++i) {
    A_views.push_back(&ref.A_tensors[i]);
    B_views.push_back(&ref.B_tensors[i]);
  }

  GroupedBuffers grouped_A = build_grouped_tensor(A_views, ref.A_tensors[0].scaling_mode());
  GroupedBuffers grouped_B = build_grouped_tensor(B_views, ref.B_tensors[0].scaling_mode());

  std::vector<Tensor> C_tensors;
  std::vector<Tensor> D_list_tensors;
  C_tensors.reserve(num_gemms);
  D_list_tensors.reserve(num_gemms);
  for (size_t i = 0; i < num_gemms; ++i) {
    const auto [M, N, K] = shapes[i];
    (void)K;
    if (!params.use_null_c) {
      C_tensors.emplace_back(
          Tensor("C" + std::to_string(i), std::vector<size_t>{M, N}, params.output_dtype));
    }
    D_list_tensors.emplace_back(
        Tensor("D_list" + std::to_string(i), std::vector<size_t>{M, N}, params.output_dtype));
    NVTE_CHECK_CUDA(cudaMemset(D_list_tensors.back().rowwise_dptr(), 0,
                               bytes(D_list_tensors.back().rowwise_shape(),
                                     D_list_tensors.back().dtype())));
  }

  std::vector<NVTETensor> C_list_ptrs;
  std::vector<NVTETensor> D_list_ptrs;
  if (!params.use_null_c) {
    C_list_ptrs.reserve(num_gemms);
  }
  D_list_ptrs.reserve(num_gemms);
  for (size_t i = 0; i < num_gemms; ++i) {
    if (!params.use_null_c) {
      C_list_ptrs.push_back(C_tensors[i].data());
    }
    D_list_ptrs.push_back(D_list_tensors[i].data());
  }

  AlphaBetaTensors ab = make_alpha_beta(num_gemms);

  const size_t setup_ws_bytes = nvte_get_grouped_gemm_setup_workspace_size(num_gemms);
  Tensor setup_ws("setup_ws", std::vector<size_t>{setup_ws_bytes}, DType::kByte);
  Tensor cublas_ws("cublas_ws", std::vector<size_t>{kCublasWorkspaceBytes}, DType::kByte);

  GroupedMatmulConfigWrapper grouped_config;
  if (ref.use_split_accum) {
    grouped_config.set_use_split_accumulator(true);
  }

  nvte_grouped_gemm_with_discrete_out(
      grouped_A.get_handle(), params.transa, grouped_B.get_handle(), params.transb,
      params.use_null_c ? nullptr : C_list_ptrs.data(), params.use_null_c ? 0 : num_gemms,
      D_list_ptrs.data(), num_gemms, ab.alpha.data(), ab.beta.data(), setup_ws.data(),
      cublas_ws.data(), grouped_config, 0);
  NVTE_CHECK_CUDA(cudaDeviceSynchronize());

  for (size_t i = 0; i < num_gemms; ++i) {
    D_list_tensors[i].to_cpu();
    ref.D_multi[i].to_cpu();
    auto [atol, rtol] = getTolerances(ref.D_multi[i].dtype());
    switch (ref.D_multi[i].dtype()) {
      case DType::kBFloat16:
        compareResults("grouped_list_vs_multi", D_list_tensors[i],
                       ref.D_multi[i].rowwise_cpu_dptr<bf16>(), true, atol, rtol);
        break;
      case DType::kFloat16:
        compareResults("grouped_list_vs_multi", D_list_tensors[i],
                       ref.D_multi[i].rowwise_cpu_dptr<fp16>(), true, atol, rtol);
        break;
      case DType::kFloat32:
        compareResults("grouped_list_vs_multi", D_list_tensors[i],
                       ref.D_multi[i].rowwise_cpu_dptr<float>(), true, atol, rtol);
        break;
      default:
        NVTE_ERROR("Unsupported D dtype in test: " +
                   std::to_string(static_cast<int>(ref.D_multi[i].dtype())));
    }
  }
}

void run_grouped_gemm_discrete_in_case(const TestParams& params) {
  if (auto reason = grouped_gemm_skip_reason(params); !reason.empty()) {
    GTEST_SKIP() << reason;
  }
  auto ref = make_grouped_gemm_ref(params);
  const auto& shapes = ref.shapes;
  const size_t num_gemms = ref.num_gemms;

  std::vector<Tensor*> B_views;
  B_views.reserve(num_gemms);
  for (size_t i = 0; i < num_gemms; ++i) B_views.push_back(&ref.B_tensors[i]);

  GroupedBuffers grouped_B = build_grouped_tensor(B_views, ref.B_tensors[0].scaling_mode());

  std::vector<Tensor> C_tensors, D_group_tensors;
  C_tensors.reserve(num_gemms);
  D_group_tensors.reserve(num_gemms);
  for (size_t i = 0; i < num_gemms; ++i) {
    const auto [M, N, K] = shapes[i];
    (void)K;
    if (!params.use_null_c) {
      C_tensors.emplace_back(Tensor("C" + std::to_string(i), std::vector<size_t>{M, N},
                                    params.output_dtype));
    }
    D_group_tensors.emplace_back(Tensor("D_group" + std::to_string(i),
                                        std::vector<size_t>{M, N}, params.output_dtype));
    NVTE_CHECK_CUDA(cudaMemset(D_group_tensors.back().rowwise_dptr(), 0,
                               bytes(D_group_tensors.back().rowwise_shape(),
                                     D_group_tensors.back().dtype())));
  }

  std::vector<Tensor*> C_views, D_views;
  for (size_t i = 0; i < num_gemms; ++i) {
    if (!params.use_null_c) C_views.push_back(&C_tensors[i]);
    D_views.push_back(&D_group_tensors[i]);
  }

  std::optional<GroupedBuffers> grouped_C;
  if (!params.use_null_c) {
    grouped_C = build_grouped_tensor(C_views, NVTE_DELAYED_TENSOR_SCALING);
  }
  GroupedBuffers grouped_D = build_grouped_tensor(D_views, NVTE_DELAYED_TENSOR_SCALING);

  AlphaBetaTensors ab = make_alpha_beta(num_gemms);

  const size_t setup_ws_bytes = nvte_get_grouped_gemm_setup_workspace_size(num_gemms);
  Tensor setup_ws("setup_ws", std::vector<size_t>{setup_ws_bytes}, DType::kByte);
  Tensor cublas_ws("cublas_ws", std::vector<size_t>{kCublasWorkspaceBytes}, DType::kByte);

  std::vector<NVTETensor> A_list_ptrs;
  A_list_ptrs.reserve(num_gemms);
  for (size_t i = 0; i < num_gemms; ++i) A_list_ptrs.push_back(ref.A_tensors[i].data());

  GroupedMatmulConfigWrapper grouped_config;
  if (ref.use_split_accum) {
    grouped_config.set_use_split_accumulator(true);
  }

  nvte_grouped_gemm_with_discrete_inputA(
      A_list_ptrs.data(), num_gemms, params.transa, grouped_B.get_handle(), params.transb,
      params.use_null_c ? nullptr : grouped_C->get_handle(), grouped_D.get_handle(),
      ab.alpha.data(), ab.beta.data(), setup_ws.data(), cublas_ws.data(), grouped_config, 0);
  NVTE_CHECK_CUDA(cudaDeviceSynchronize());

  compare_grouped_d_to_multi(grouped_D, shapes, ref.D_multi, "grouped_discrete_in_vs_multi");
}

// Graph-capture variant of run_grouped_gemm_case. Captures nvte_grouped_gemm on a
// non-default stream and replays it twice, verifying outputs against the multi-tensor
// reference after each replay. Asserts capture succeeds, replay is correct, and the
// operation is deterministic across replays.
void run_grouped_gemm_graph_case(const TestParams& params) {
  if (auto reason = grouped_gemm_skip_reason(params); !reason.empty()) {
    GTEST_SKIP() << reason;
  }
  auto ref = make_grouped_gemm_ref(params);
  const auto& shapes = ref.shapes;
  const size_t num_gemms = ref.num_gemms;

  std::vector<Tensor*> A_views, B_views;
  A_views.reserve(num_gemms);
  B_views.reserve(num_gemms);
  for (size_t i = 0; i < num_gemms; ++i) {
    A_views.push_back(&ref.A_tensors[i]);
    B_views.push_back(&ref.B_tensors[i]);
  }
  GroupedBuffers grouped_A = build_grouped_tensor(A_views, ref.A_tensors[0].scaling_mode());
  GroupedBuffers grouped_B = build_grouped_tensor(B_views, ref.B_tensors[0].scaling_mode());

  std::vector<Tensor> C_tensors, D_group_tensors;
  C_tensors.reserve(num_gemms);
  D_group_tensors.reserve(num_gemms);
  for (size_t i = 0; i < num_gemms; ++i) {
    const auto [M, N, K] = shapes[i];
    (void)K;
    if (!params.use_null_c) {
      C_tensors.emplace_back(
          Tensor("C" + std::to_string(i), std::vector<size_t>{M, N}, params.output_dtype));
    }
    D_group_tensors.emplace_back(
        Tensor("D_group" + std::to_string(i), std::vector<size_t>{M, N}, params.output_dtype));
  }

  std::vector<Tensor*> C_views, D_views;
  for (size_t i = 0; i < num_gemms; ++i) {
    if (!params.use_null_c) C_views.push_back(&C_tensors[i]);
    D_views.push_back(&D_group_tensors[i]);
  }

  std::optional<GroupedBuffers> grouped_C;
  if (!params.use_null_c) {
    grouped_C = build_grouped_tensor(C_views, NVTE_DELAYED_TENSOR_SCALING);
  }
  GroupedBuffers grouped_D = build_grouped_tensor(D_views, NVTE_DELAYED_TENSOR_SCALING);

  AlphaBetaTensors ab = make_alpha_beta(num_gemms);
  const size_t setup_ws_bytes = nvte_get_grouped_gemm_setup_workspace_size(num_gemms);
  Tensor setup_ws("setup_ws", std::vector<size_t>{setup_ws_bytes}, DType::kByte);
  Tensor cublas_ws("cublas_ws", std::vector<size_t>{kCublasWorkspaceBytes}, DType::kByte);
  GroupedMatmulConfigWrapper grouped_config;
  if (ref.use_split_accum) grouped_config.set_use_split_accumulator(true);

  cudaStream_t stream;
  NVTE_CHECK_CUDA(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking));

  auto body = [&](cudaStream_t s) {
    for (auto& d : D_group_tensors) {
      NVTE_CHECK_CUDA(cudaMemsetAsync(d.rowwise_dptr(), 0,
                                      bytes(d.rowwise_shape(), d.dtype()), s));
    }
    nvte_grouped_gemm(grouped_A.get_handle(), params.transa, grouped_B.get_handle(),
                      params.transb,
                      params.use_null_c ? nullptr : grouped_C->get_handle(),
                      grouped_D.get_handle(), ab.alpha.data(), ab.beta.data(),
                      setup_ws.data(), cublas_ws.data(), grouped_config, s);
  };

  capture_and_replay_grouped_gemm(
      stream, /*n_replays=*/2, body,
      [&](int iter) {
        const std::string tag = "grouped_graph_replay_" + std::to_string(iter);
        compare_grouped_d_to_multi(grouped_D, shapes, ref.D_multi, tag.c_str());
      });

  NVTE_CHECK_CUDA(cudaStreamDestroy(stream));
}

// Graph-capture variant of run_grouped_gemm_discrete_out_case.
void run_grouped_gemm_discrete_out_graph_case(const TestParams& params) {
  if (auto reason = grouped_gemm_skip_reason(params); !reason.empty()) {
    GTEST_SKIP() << reason;
  }
  auto ref = make_grouped_gemm_ref(params);
  const auto& shapes = ref.shapes;
  const size_t num_gemms = ref.num_gemms;

  std::vector<Tensor*> A_views, B_views;
  A_views.reserve(num_gemms);
  B_views.reserve(num_gemms);
  for (size_t i = 0; i < num_gemms; ++i) {
    A_views.push_back(&ref.A_tensors[i]);
    B_views.push_back(&ref.B_tensors[i]);
  }
  GroupedBuffers grouped_A = build_grouped_tensor(A_views, ref.A_tensors[0].scaling_mode());
  GroupedBuffers grouped_B = build_grouped_tensor(B_views, ref.B_tensors[0].scaling_mode());

  std::vector<Tensor> C_tensors, D_list_tensors;
  C_tensors.reserve(num_gemms);
  D_list_tensors.reserve(num_gemms);
  for (size_t i = 0; i < num_gemms; ++i) {
    const auto [M, N, K] = shapes[i];
    (void)K;
    if (!params.use_null_c) {
      C_tensors.emplace_back(
          Tensor("C" + std::to_string(i), std::vector<size_t>{M, N}, params.output_dtype));
    }
    D_list_tensors.emplace_back(
        Tensor("D_list" + std::to_string(i), std::vector<size_t>{M, N}, params.output_dtype));
  }

  std::vector<NVTETensor> C_list_ptrs, D_list_ptrs;
  if (!params.use_null_c) C_list_ptrs.reserve(num_gemms);
  D_list_ptrs.reserve(num_gemms);
  for (size_t i = 0; i < num_gemms; ++i) {
    if (!params.use_null_c) C_list_ptrs.push_back(C_tensors[i].data());
    D_list_ptrs.push_back(D_list_tensors[i].data());
  }

  AlphaBetaTensors ab = make_alpha_beta(num_gemms);
  const size_t setup_ws_bytes = nvte_get_grouped_gemm_setup_workspace_size(num_gemms);
  Tensor setup_ws("setup_ws", std::vector<size_t>{setup_ws_bytes}, DType::kByte);
  Tensor cublas_ws("cublas_ws", std::vector<size_t>{kCublasWorkspaceBytes}, DType::kByte);
  GroupedMatmulConfigWrapper grouped_config;
  if (ref.use_split_accum) grouped_config.set_use_split_accumulator(true);

  cudaStream_t stream;
  NVTE_CHECK_CUDA(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking));

  auto body = [&](cudaStream_t s) {
    for (auto& d : D_list_tensors) {
      NVTE_CHECK_CUDA(cudaMemsetAsync(d.rowwise_dptr(), 0,
                                      bytes(d.rowwise_shape(), d.dtype()), s));
    }
    nvte_grouped_gemm_with_discrete_out(
        grouped_A.get_handle(), params.transa, grouped_B.get_handle(), params.transb,
        params.use_null_c ? nullptr : C_list_ptrs.data(),
        params.use_null_c ? 0 : num_gemms, D_list_ptrs.data(), num_gemms, ab.alpha.data(),
        ab.beta.data(), setup_ws.data(), cublas_ws.data(), grouped_config, s);
  };

  auto verify = [&](int iter) {
    const std::string tag = "discrete_out_graph_replay_" + std::to_string(iter);
    for (size_t i = 0; i < num_gemms; ++i) {
      D_list_tensors[i].to_cpu();
      ref.D_multi[i].to_cpu();
      auto [atol, rtol] = getTolerances(ref.D_multi[i].dtype());
      switch (ref.D_multi[i].dtype()) {
        case DType::kBFloat16:
          compareResults(tag.c_str(), D_list_tensors[i],
                         ref.D_multi[i].rowwise_cpu_dptr<bf16>(), true, atol, rtol);
          break;
        case DType::kFloat16:
          compareResults(tag.c_str(), D_list_tensors[i],
                         ref.D_multi[i].rowwise_cpu_dptr<fp16>(), true, atol, rtol);
          break;
        case DType::kFloat32:
          compareResults(tag.c_str(), D_list_tensors[i],
                         ref.D_multi[i].rowwise_cpu_dptr<float>(), true, atol, rtol);
          break;
        default:
          NVTE_ERROR("Unsupported D dtype in test: " +
                     std::to_string(static_cast<int>(ref.D_multi[i].dtype())));
      }
    }
  };

  capture_and_replay_grouped_gemm(stream, /*n_replays=*/2, body, verify);
  NVTE_CHECK_CUDA(cudaStreamDestroy(stream));
}

// Graph-capture variant of run_grouped_gemm_discrete_in_case.
void run_grouped_gemm_discrete_in_graph_case(const TestParams& params) {
  if (auto reason = grouped_gemm_skip_reason(params); !reason.empty()) {
    GTEST_SKIP() << reason;
  }
  auto ref = make_grouped_gemm_ref(params);
  const auto& shapes = ref.shapes;
  const size_t num_gemms = ref.num_gemms;

  std::vector<Tensor*> B_views;
  B_views.reserve(num_gemms);
  for (size_t i = 0; i < num_gemms; ++i) B_views.push_back(&ref.B_tensors[i]);
  GroupedBuffers grouped_B = build_grouped_tensor(B_views, ref.B_tensors[0].scaling_mode());

  std::vector<Tensor> C_tensors, D_group_tensors;
  C_tensors.reserve(num_gemms);
  D_group_tensors.reserve(num_gemms);
  for (size_t i = 0; i < num_gemms; ++i) {
    const auto [M, N, K] = shapes[i];
    (void)K;
    if (!params.use_null_c) {
      C_tensors.emplace_back(Tensor("C" + std::to_string(i),
                                    std::vector<size_t>{M, N}, params.output_dtype));
    }
    D_group_tensors.emplace_back(Tensor("D_group" + std::to_string(i),
                                        std::vector<size_t>{M, N}, params.output_dtype));
  }

  std::vector<Tensor*> C_views, D_views;
  for (size_t i = 0; i < num_gemms; ++i) {
    if (!params.use_null_c) C_views.push_back(&C_tensors[i]);
    D_views.push_back(&D_group_tensors[i]);
  }

  std::optional<GroupedBuffers> grouped_C;
  if (!params.use_null_c) {
    grouped_C = build_grouped_tensor(C_views, NVTE_DELAYED_TENSOR_SCALING);
  }
  GroupedBuffers grouped_D = build_grouped_tensor(D_views, NVTE_DELAYED_TENSOR_SCALING);

  AlphaBetaTensors ab = make_alpha_beta(num_gemms);
  const size_t setup_ws_bytes = nvte_get_grouped_gemm_setup_workspace_size(num_gemms);
  Tensor setup_ws("setup_ws", std::vector<size_t>{setup_ws_bytes}, DType::kByte);
  Tensor cublas_ws("cublas_ws", std::vector<size_t>{kCublasWorkspaceBytes}, DType::kByte);

  std::vector<NVTETensor> A_list_ptrs;
  A_list_ptrs.reserve(num_gemms);
  for (size_t i = 0; i < num_gemms; ++i) A_list_ptrs.push_back(ref.A_tensors[i].data());

  GroupedMatmulConfigWrapper grouped_config;
  if (ref.use_split_accum) grouped_config.set_use_split_accumulator(true);

  cudaStream_t stream;
  NVTE_CHECK_CUDA(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking));

  auto body = [&](cudaStream_t s) {
    for (auto& d : D_group_tensors) {
      NVTE_CHECK_CUDA(cudaMemsetAsync(d.rowwise_dptr(), 0,
                                      bytes(d.rowwise_shape(), d.dtype()), s));
    }
    nvte_grouped_gemm_with_discrete_inputA(
        A_list_ptrs.data(), num_gemms, params.transa, grouped_B.get_handle(), params.transb,
        params.use_null_c ? nullptr : grouped_C->get_handle(), grouped_D.get_handle(),
        ab.alpha.data(), ab.beta.data(), setup_ws.data(), cublas_ws.data(), grouped_config, s);
  };

  capture_and_replay_grouped_gemm(
      stream, /*n_replays=*/2, body,
      [&](int iter) {
        const std::string tag = "discrete_in_graph_replay_" + std::to_string(iter);
        compare_grouped_d_to_multi(grouped_D, shapes, ref.D_multi, tag.c_str());
      });

  NVTE_CHECK_CUDA(cudaStreamDestroy(stream));
}

class GroupedGemmTest : public ::testing::TestWithParam<TestParams> {};

TEST_P(GroupedGemmTest, CompareWithMultiTensorGemm) {
  run_grouped_gemm_case(GetParam());
}

TEST_P(GroupedGemmTest, CompareWithMultiTensorGemmDiscreteOut) {
  run_grouped_gemm_discrete_out_case(GetParam());
}

TEST_P(GroupedGemmTest, CompareWithMultiTensorGemmDiscreteIn) {
  run_grouped_gemm_discrete_in_case(GetParam());
}

TEST_P(GroupedGemmTest, CudaGraphCapture) {
  run_grouped_gemm_graph_case(GetParam());
}

TEST_P(GroupedGemmTest, CudaGraphCaptureDiscreteOut) {
  run_grouped_gemm_discrete_out_graph_case(GetParam());
}

TEST_P(GroupedGemmTest, CudaGraphCaptureDiscreteIn) {
  run_grouped_gemm_discrete_in_graph_case(GetParam());
}

std::string MakeGroupedGemmTestName(const testing::TestParamInfo<GroupedGemmTest::ParamType>& info) {
  constexpr const char* kShapeNames[] = {"AllSameMul128", "SameMMul128", "SameNMul128",
                                         "AllDiffMul128", "AllSameMul32"};
  const std::string layout = std::string("ta") + (info.param.transa ? "T" : "N") +
                             "tb" + (info.param.transb ? "T" : "N");
  const std::string null_c = info.param.use_null_c ? "_NullC" : "";
  std::string out_suffix;
  switch (info.param.output_dtype) {
    case DType::kBFloat16: break;  // default, no suffix
    case DType::kFloat16:  out_suffix = "_outFP16"; break;
    case DType::kFloat32:  out_suffix = "_outFP32"; break;
    default:               out_suffix = "_outUnknown"; break;
  }
  return std::string(recipe_name(info.param.recipe)) + "_" +
         kShapeNames[static_cast<int>(info.param.shape_case)] + "_" + layout + null_c + out_suffix;
}

// TestParams: {recipe, transa, transb, shape_case, use_null_c}
const std::vector<TestParams> kTestParams = {
    // FP8 tests (each tensor has random mean/stddev -> different scales)
    {InputRecipe::kFP8Current, true, false, ShapeCase::kAllDifferentMul128, false},
    {InputRecipe::kFP8Current, false, true, ShapeCase::kAllDifferentMul128, false},
    {InputRecipe::kFP8Current, false, false, ShapeCase::kAllSameMul128, false},
    // BF16 tests
    {InputRecipe::kBF16, true, false, ShapeCase::kSameFirstMul128, false},
    {InputRecipe::kBF16, false, true, ShapeCase::kSameLastMul128, false},
    {InputRecipe::kBF16, false, false, ShapeCase::kAllSameMul128, false},
    {InputRecipe::kBF16, true, true, ShapeCase::kAllDifferentMul128, false},
    // Test NULL C (valid when beta=0)
    {InputRecipe::kBF16, false, false, ShapeCase::kAllSameMul128, true},
    // MXFP8 tests
    {InputRecipe::kMXFP8, true, false, ShapeCase::kAllSameMul128, false},
    {InputRecipe::kMXFP8, true, false, ShapeCase::kAllDifferentMul128, false},
    {InputRecipe::kMXFP8, false, true, ShapeCase::kAllSameMul128, false},
    {InputRecipe::kMXFP8, false, true, ShapeCase::kAllDifferentMul128, false},
    {InputRecipe::kMXFP8, false, false, ShapeCase::kAllSameMul128, false},
    {InputRecipe::kMXFP8, false, false, ShapeCase::kAllDifferentMul128, false},
    {InputRecipe::kMXFP8, false, false, ShapeCase::kSameFirstMul128, false},
    // MXFP8 with NULL C
    {InputRecipe::kMXFP8, true, false, ShapeCase::kAllSameMul128, true},
    // NVFP4 tests (all transpose combinations - GEMM internally forces TN)
    {InputRecipe::kNVFP4, true, false, ShapeCase::kAllSameMul128, false},
    {InputRecipe::kNVFP4, true, false, ShapeCase::kAllDifferentMul128, false},
    {InputRecipe::kNVFP4, true, false, ShapeCase::kSameFirstMul128, false},
    {InputRecipe::kNVFP4, true, false, ShapeCase::kSameLastMul128, false},
    {InputRecipe::kNVFP4, false, true, ShapeCase::kAllSameMul128, false},
    {InputRecipe::kNVFP4, false, true, ShapeCase::kAllDifferentMul128, false},
    {InputRecipe::kNVFP4, false, false, ShapeCase::kAllSameMul128, false},
    {InputRecipe::kNVFP4, false, false, ShapeCase::kAllDifferentMul128, false},
    // NVFP4 with NULL C
    {InputRecipe::kNVFP4, true, false, ShapeCase::kAllSameMul128, true},
    // Non-default output dtypes (BF16 covered everywhere else).
    {InputRecipe::kBF16,          false, false, ShapeCase::kAllSameMul128, false,
     /*output_dtype=*/DType::kFloat32},
    {InputRecipe::kFP8Current,    true,  false, ShapeCase::kAllSameMul128, false,
     /*output_dtype=*/DType::kFloat16},
    // FP8 Block Scaling tests (TN layout on Hopper, block size 128)
    {InputRecipe::kFP8BlockScaling1D1D, true, false, ShapeCase::kAllSameMul128, false},
    {InputRecipe::kFP8BlockScaling1D1D, true, false, ShapeCase::kAllDifferentMul128, false},
    {InputRecipe::kFP8BlockScaling1D1D, true, false, ShapeCase::kSameFirstMul128, false},
    {InputRecipe::kFP8BlockScaling1D1D, true, false, ShapeCase::kSameLastMul128, false},
    {InputRecipe::kFP8BlockScaling1D1D, false, true, ShapeCase::kAllSameMul128, false},
    {InputRecipe::kFP8BlockScaling1D1D, false, false, ShapeCase::kAllSameMul128, false},
    // FP8 Block Scaling with NULL C
    {InputRecipe::kFP8BlockScaling1D1D, true, false, ShapeCase::kAllSameMul128, true},
    // Dims multiples of 32 but not 128 exercise padded scale_inv offsets for recipes that use
    // padded grouped scale_inv storage.
    {InputRecipe::kMXFP8, true, false, ShapeCase::kAllSameMul32, false},
    {InputRecipe::kMXFP8, false, true, ShapeCase::kAllSameMul32, false},
    {InputRecipe::kMXFP8, false, false, ShapeCase::kAllSameMul32, false},
    {InputRecipe::kNVFP4, true, false, ShapeCase::kAllSameMul32, false},
    {InputRecipe::kNVFP4, false, true, ShapeCase::kAllSameMul32, false},
    {InputRecipe::kNVFP4, false, false, ShapeCase::kAllSameMul32, false},
    {InputRecipe::kFP8BlockScaling1D1D, true, false, ShapeCase::kAllSameMul32, false},
    {InputRecipe::kFP8BlockScaling1D1D, false, true, ShapeCase::kAllSameMul32, false},
    {InputRecipe::kFP8BlockScaling1D1D, false, false, ShapeCase::kAllSameMul32, false},
    // Mixed FP8 block scaling modes supported by cuBLASLt: 2D x 1D and 1D x 2D.
    {InputRecipe::kFP8BlockScaling2D1D, true, false, ShapeCase::kAllSameMul128, false},
    {InputRecipe::kFP8BlockScaling2D1D, true, false, ShapeCase::kAllDifferentMul128, false},
    {InputRecipe::kFP8BlockScaling1D2D, true, false, ShapeCase::kAllSameMul128, false},
    {InputRecipe::kFP8BlockScaling1D2D, true, false, ShapeCase::kAllDifferentMul128, false},
    {InputRecipe::kFP8BlockScaling1D2D, false, false, ShapeCase::kAllSameMul32, false},
};

INSTANTIATE_TEST_SUITE_P(OperatorTest,
                         GroupedGemmTest,
                         ::testing::ValuesIn(kTestParams),
                         MakeGroupedGemmTestName);

}  // namespace
