# Sample configuration for inference benchmarks
benchmark_mode: "inference"
quantization_config_recipe_names: # Will run a baseline inference for model by default, without quantization for comparison
  - "int8wo"
  - "int8dq"
  - "float8dq-tensor"
  - "float8wo"
output_dir: "benchmarks/microbenchmarks/results"
model_params:
  - name: "small_bf16_linear"
    matrix_shapes:
      - name: "small_sweep"
        min_power: 14
        max_power: 16
    high_precision_dtype: "torch.bfloat16"
    torch_compile_mode: "max-autotune"
    device: "cuda"
    model_type: "linear"
    enable_profiler: true  # Enable profiling for this model
    enable_memory_profiler: true  # Enable memory profiling for this model

  - name: "ln_linear_sigmoid_cuda"
    matrix_shapes:
      - name: "custom"
        shapes: [
          [2048, 4096, 1024],
        ]
    high_precision_dtype: "torch.bfloat16"
    torch_compile_mode: "max-autotune"
    device: "cuda"
    model_type: "ln_linear_sigmoid"
    enable_profiler: true
    enable_memory_profiler: true

  - name: "bf16_transformer_block"
    matrix_shapes:
      - name: "custom"
        shapes: [
          [2048, 4096, 1024],  # For transformer_block, k is the hidden dimension
        ]
    high_precision_dtype: "torch.bfloat16"
    torch_compile_mode: "max-autotune"
    device: "cuda"
    model_type: "transformer_block" # TODO: Add a custom model (Figure out how to do this, maybe pass a .py file with model definition)
    enable_profiler: true
    enable_memory_profiler: true

  - name: "large_bf16_ln_linear"
    matrix_shapes:
      - name: "llama"  # Example of using LLaMa shapes
      - name: "pow2"  # Example of using power of 2 shapes
        min_power: 10  # 1024
        max_power: 12  # 4096
      - name: "pow2_extended"  # Example of using extended power of 2 shapes
        min_power: 10  # 1024
        max_power: 11  # 2048
    high_precision_dtype: "torch.bfloat16"
    torch_compile_mode: "max-autotune"
    device: "cuda"
    model_type: "linear"
    enable_profiler: true
    enable_memory_profiler: true