                "peak_memory_gb": self.metrics.peak_memory_gb,
            },
            "detailed_metrics": {
                "iteration_times": self.metrics.iteration_times,
                "ttft_times": self.metrics.ttft_times,
                "prefill_times": self.metrics.prefill_times,
                "decode_times": self.metrics.decode_times,
            },
        }

        with open(filename, "w") as f:
            json.dump(results, f, indent=2)

        print(f"\nResults saved to {filename}")

class CustomFormatter(argparse.RawDescriptionHelpFormatter, argparse.ArgumentDefaultsHelpFormatter):
    pass

def parse_args() -> argparse.Namespace:
    """Command line interface for the benchmark"""
    parser = argparse.ArgumentParser(
        description="Thunder Inference Benchmark",
        formatter_class=CustomFormatter,
        epilog="""
Examples:
  python benchmark_inference.py --input-length 2048 --output-length 512 --model-name meta-llama/Llama-4-Maverick-17B-128E --mode eager
        """,
    )

    parser.add_argument(
        "--model-name",
        type=str,
        default=LLAMA4_MAVERICK_MODEL_ID,
        help="Model to benchmark",
    )
    parser.add_argument("--batch-size", type=int, default=1, help="Batch size for inference")
    parser.add_argument(
        "--input-length",
        type=int,
        default=2048,
        help="Input sequence length",
    )
    parser.add_argument(
        "--output-length",
        type=int,
        default=128,
        help="Output sequence length",
    )
    parser.add_argument("--num-iterations", type=int, default=100, help="Number of benchmark iterations")
    parser.add_argument("--warmup-iterations", type=int, default=10, help="Number of warmup iterations")
    parser.add_argument(
        "--num-layers",
        default=2,
        type=int,
        help="Number of layers of the moddel. Llama4 Maverick has 48 hidden layers, which could be too memory hungry",
    )
    parser.add_argument(
        "--disable-moe-replacement",
        action="store_true",
        help="Disallow replacement of Llama4TextMoe with our custom Llama4MoE which uses grouped gemm.",
    )

    # Execution configuration
    parser.add_argument(
        "--mode",
        type=str,
        default="eager",
        choices=("thunder", "eager", "inductor", "thunderjit"),
        help="Compilation mode: thunder, eager (default), or inductor. thunder runs thunderfx.",
    )
    parser.add_argument(
        "--fx-report-folder",
        default=None,
        type=str,
        help="Specify the folder for thunderfx_benchmark_report.",
    )

    parser.add_argument(
        "--enable-nvfp4",
        action="store_true",
        help="Enable NVFP4 quantization for MoE GroupedSwiGLU layers (has nvfuser grouped_mm support)",
    )
    parser.add_argument(
        "--enable-nv-linear",
        action="store_true",
        help="let nvfuser take care of linear and matmul, note that this might fail with distributed run. See: https://github.com/NVIDIA/Fuser/issues/4507",
    )

    parser.add_argument(
        "--thunder-trace",
        action="store_true",
        help="Enable debug dump of thunder trace",
    )
    parser.add_argument(
        "--debug-moe",
        action="store_true",
        help="Enable debug of MoE, limiting profile and compilation only to the hand written MoE layer",
    )
    parser.add_argument("--save-results", action="store_true", help="Save results to JSON file")
    parser.add_argument("--output-dir", type=str, default="./results", help="Directory to save results")
    parser.add_argument(
        "--thunder-cache",
        type=str,
        default=None,
        help="Cache option: no caching, same input, constant values, symbolic values. See `cache` argument of `thunder.jit` for more details.",
    )
    parser.add_argument("--enable-cudagraph", action="store_true", help="Pass CUDAGraphTransform to Thunder, or use reduce-overhead for torch.compile")
    parser.add_argument("--attn-implementation", type=str, default=None, help="Attention implementation")
    parser.add_argument("--use-hardcoded-model", action="store_true", help="Use the hardcoded Llama4 config, rather than pulling from huggingfaces")

    args = parser.parse_args()
    return args

def main():
    args = parse_args()
    if args.save_results:
        os.makedirs(args.output_dir, exist_ok=True)

    # Register NVFP4 custom ops with nvfuser translators when enabled
    if args.enable_nvfp4:
        try:
            _register_nvfp4_ops()
        except Exception as e:
            # If registration fails (e.g., nvfuser not available), warn and continue
            warnings.warn(f"Failed to register nvfp4 custom ops: {e}")

    config = InferenceBenchmarkConfig(
        model_name=args.model_name,
        batch_size=args.batch_size,
        input_length=args.input_length,
        output_length=args.output_length,


main

        "--debug-moe",
        action="store_true",
        help="Enable debug of MoE, limiting profile and compilation only to the hand written MoE layer",
    )
    parser.add_argument("--save-results", action="store_true", help="Save results to JSON file")
    parser.add_argument("--output-dir", type=str, default="./results", help="Directory to save results")
    parser.add_argument(
        "--thunder-cache",
        type=str,
        default=None,
        help="Cache option: no caching, same input, constant values, symbolic values. See `cache` argument of `thunder.jit` for more details.",
    )
    parser.add_argument("--enable-cudagraph", action="store_true", help="Pass CUDAGraphTransform to Thunder, or use reduce-overhead for torch.compile")
    parser.add_argument("--attn-implementation", type=str, default=None, help="Attention implementation")
    parser.add_argument("--use-hardcoded-model", action="store_true", help="Use the hardcoded Llama4 config, rather than pulling from huggingfaces")

    args = parser.parse_args()
    return args

def main():
    args = parse_args()
    if args.save_results:
        os.makedirs(args.output_dir, exist_ok=True)

    # Register NVFP4 custom ops with nvfuser translators when enabled
    if args.enable_nvfp4:
        try:
            _register_nvfp4_ops()
        except Exception as e:
            # If registration fails (e.g., nvfuser not available), warn and continue
            warnings.warn(f"Failed to register nvfp4 custom ops: {e}")

    config = InferenceBenchmarkConfig(
        model_name=args.model_name,
        batch_size=args.batch_size,
        input_length=args.input_length,
        output_length=args.output_length,
        num_layers=args.num_layers,
        num_iterations=args.num_iterations,
        warmup_iterations=args.warmup_iterations,
        mode=args.mode,
        enable_nvfp4=args.enable_nvfp4,
        fx_report_folder=args.fx_report_folder,
        enable_nv_linear=args.enable_nv_linear,
        disable_moe_replacement=args.disable_moe_replacement,
        attn_implementation=args.attn_implementation,
        thunder_cache=args.thunder_cache,
        enable_cudagraph=args.enable_cudagraph,
        debug_moe=args.debug_moe,
        use_hardcoded_model=args.use_hardcoded_model,
    )
    benchmark = InferenceBenchmark(config)

    benchmark.run_benchmark()
    benchmark.print_results()

    if args.thunder_trace and args.mode == "thunder":
        backend = benchmark.model._backend
        for subgraph_info in backend.subgraph_infos:
            assert isinstance(subgraph_info.original_graph_module, torch.fx.GraphModule)
            assert len(subgraph_info.thunder_compiled_fns)
            for thunder_fn in subgraph_info.thunder_compiled_fns:
                print(thunder.last_traces(thunder_fn)[-1])

    if args.save_results:
        timestamp = time.strftime("%Y%m%d_%H%M%S")
        filename = f"thunder_inference_{args.model_name.replace('/', '_')}_{timestamp}.json"
        path = os.path.join(args.output_dir, filename)
        benchmark.save_results(path)

if __name__ == "__main__":
    try:
        main()
    except Exception:
        raise
    finally:
        if mesh:
            for process_group in mesh.get_all_groups():
                dist.destroy_process_group(process_group)


  "rms_norm_eps": 1e-05,
  "rope_scaling": null,
  "rope_theta": 500000.0,
  "router_aux_loss_coef": 0.001,
  "router_jitter_noise": 0.0,
  "torch_dtype": "bfloat16",
  "use_cache": true,
  "use_qk_norm": false,
  "vocab_size": 202048
}
"""

# TODO: Add mm quantization once nvfuser implements nvfp4 gemm
# Register nvfp4 custom ops with Thunder and nvFuser
def _register_nvfp4_ops():
    """Register nvfp4 custom operations with Thunder."""
    # Register f16a_nvfp4weight_scaled_grouped_mm with nvfuser translator
    _nvfp4_grouped_mm_symbol = _register_custom_op(nvfuser_f16a_nvfp4weight_scaled_grouped_mm)

    def nvfp4_grouped_mm_translator(
        activation,
        fp4_weight,
        weight_scaling_factor,
        global_scale,
        offsets,
        blockscale_offsets,
        problem_sizes,
        *,
        fd,
        lc_to_nv_map,
    ):
        from nvfuser_direct import DataType
        from thunder.executors.nvfuserex_impl import getnv

        nv_act = getnv(activation, fd, lc_to_nv_map)
        nv_fp4_w = getnv(fp4_weight, fd, lc_to_nv_map)
        nv_sf_w = getnv(weight_scaling_factor, fd, lc_to_nv_map)
        nv_alpha = getnv(global_scale, fd, lc_to_nv_map)
        nv_offsets = getnv(offsets, fd, lc_to_nv_map)
        nv_blocksf_offsets = getnv(blockscale_offsets, fd, lc_to_nv_map)
        nv_problem_sizes = getnv(problem_sizes, fd, lc_to_nv_map)
        fp4_mat1, layout_fp8_scale1 = fd.ops.nv_grouped_block_quantize(nv_act, nv_offsets, nv_blocksf_offsets)
        out = fd.ops.cutlass_nvfp4_grouped_mm(
            fp4_mat1,
            nv_fp4_w,
            layout_fp8_scale1,
            nv_sf_w,
            nv_alpha,
            # NOTE: we might need to call contiguous on problem_sizes
            nv_problem_sizes,
            nv_offsets,
            nv_blocksf_offsets,
            DataType.BFloat16,
        )
        return out

    _register_nvfuser_translator(_nvfp4_grouped_mm_symbol, nvfp4_grouped_mm_translator)

# The logic is based on https://github.com/pytorch/ao/blob/b34c1037/torchao/quantization/quant_api.py#L230
def _replace_with_custom_fn_if_matches_filter_with_name(
    model,
    replacement_fn: Callable[[torch.nn.Module, str], torch.nn.Module],
    filter_fn: Callable[[torch.nn.Module, str], bool],
    cur_fqn="",
) -> None:
    """
    Recursively replaces each child module in `model` with the result of `replacement_fn(child)`

    replacement_fn (Callable[[torch.nn.Module, str], torch.nn.Module]): The function to replace matching modules.
    filter_fn (Callable[[torch.nn.Module, str], bool]): The function to filter matching modules.
    cur_fqn (str): The current fully qualified name of the module.

    Returns:
        None


        # Required as thunder doesn't understand inference mode
        # And some prims like `prims._grouped_mm` don't have grad rule defined yet.
        for p in model.parameters():
            p.requires_grad_(False)

        assert all(not p.requires_grad for p in model.parameters())

        # `thunderfx` seems to hide the access to vocab_size somewhere so
        # store it here before any compiler is applied.
        self.vocab_size = model.vocab_size

        if self.config.enable_nvfp4:
            _quantize_llama4(model)

        # If debug_moe is on, we'll only compile the moe section, which is much easier to digest.
        if self.config.debug_moe:
            def compile_wrapper(model, jitter):
                class ModelWrapper(torch.nn.Module):

                    def __init__(self, model, jitter):
                        super().__init__()
                        self.model = jitter._compile_model(model)
                        assert not hasattr(jitter.model, "_backend")
                        # store a reference to the compiled backend, so we can print the trace later.
                        if jitter.config.mode == "thunder":
                            jitter.model._backend = self.model._backend
                        self.jitter = jitter

                    def forward(self, *args, **kwargs):
                        if self.jitter.profiler_toggle:
                            torch.cuda.cudart().cudaProfilerStart()
                        out = self.model(*args, **kwargs)
                        if self.jitter.profiler_toggle:
                            torch.cuda.cudart().cudaProfilerStop()
                        return out

                return ModelWrapper(model, jitter)

            self.model = model
            # reuse the `replace` function with compilation.
            _replace_with_custom_fn_if_matches_filter_with_name(
                model,
                lambda model, cur_fqn: compile_wrapper(model, self),
                lambda model, cur_fqn: isinstance(model, Llama4MoE),
            )
        else:
            self.model = self._compile_model(model)


                            torch.cuda.cudart().cudaProfilerStart()
                        out = self.model(*args, **kwargs)
                        if self.jitter.profiler_toggle:
                            torch.cuda.cudart().cudaProfilerStop()
                        return out

                return ModelWrapper(model, jitter)

            self.model = model
            # reuse the `replace` function with compilation.
            _replace_with_custom_fn_if_matches_filter_with_name(
                model,
                lambda model, cur_fqn: compile_wrapper(model, self),
                lambda model, cur_fqn: isinstance(model, Llama4MoE),
            )
        else:
            self.model = self._compile_model(model)

    @property
    def _thunder_jit_options(self) -> dict[str, Any]:
        # `nv_enable_linear=True` might fail with distributed run
        # ref: https://github.com/NVIDIA/Fuser/issues/4507
        res = {"transforms": []}
        if self.config.enable_nv_linear:
            res["nv_enable_linear"] = True
            res["nv_enable_matmul"] = True
        if self.config.mode == "thunderjit":
            from thunder.recipes.hf_transformers import SDPAMaskTransform

            if not hasattr(self, "_mask_transform"):
                self._mask_transform = SDPAMaskTransform()
            res["transforms"].append(self._mask_transform)
            res["executors"] = [self._mask_transform.get_executor(), *thunder.get_default_executors()]
        if self.config.enable_cudagraph:
            res["transforms"].append(CUDAGraphTransform())
        if self.config.thunder_cache is not None:
            res["cache"] = self.config.thunder_cache

        return res

    def _compile_model(self, model):
        match self.config.mode:
            case "eager":
                return model
            case "inductor":
                if (self.config.enable_cudagraph):
                    return torch.compile(model, mode="reduce-overhead")
                else:
                    return torch.compile(model, mode="default")
            case "thunder":
                return thunderfx(model, **self._thunder_jit_options)
            case "thunderjit":
                return thunder.jit(model, **self._thunder_jit_options)
            case _:
                raise ValueError(f"Unknown mode: {self.config.mode}")

    def _load_model(self) -> torch.nn.Module:
        """Load the model based on configuration"""
        model_id = self.config.model_name


        # `nv_enable_linear=True` might fail with distributed run
        # ref: https://github.com/NVIDIA/Fuser/issues/4507
        res = {"transforms": []}
        if self.config.enable_nv_linear:
            res["nv_enable_linear"] = True
            res["nv_enable_matmul"] = True
        if self.config.mode == "thunderjit":
            from thunder.recipes.hf_transformers import SDPAMaskTransform

            if not hasattr(self, "_mask_transform"):
                self._mask_transform = SDPAMaskTransform()
            res["transforms"].append(self._mask_transform)
            res["executors"] = [self._mask_transform.get_executor(), *thunder.get_default_executors()]
        if self.config.enable_cudagraph:
            res["transforms"].append(CUDAGraphTransform())
        if self.config.thunder_cache is not None:
            res["cache"] = self.config.thunder_cache

        return res

    def _compile_model(self, model):
        match self.config.mode:
            case "eager":
                return model
            case "inductor":
                if (self.config.enable_cudagraph):
                    return torch.compile(model, mode="reduce-overhead")
                else:
                    return torch.compile(model, mode="default")
            case "thunder":
                return thunderfx(model, **self._thunder_jit_options)
            case "thunderjit":
                return thunder.jit(model, **self._thunder_jit_options)
            case _:
                raise ValueError(f"Unknown mode: {self.config.mode}")

    def _load_model(self) -> torch.nn.Module:
        """Load the model based on configuration"""
        model_id = self.config.model_name

        config = None
        if self.config.use_hardcoded_model:
            config = Llama4TextConfig.from_dict(
                json.loads(llama_4_Maverick_17B_128E_cfg_str)
            )
        else:
            config = AutoConfig.from_pretrained(model_id)

        if hasattr(config, "text_config"):
            config = config.text_config
        if self.config.num_layers:
            config.num_hidden_layers = self.config.num_layers

        self.hf_config = config


            res["cache"] = self.config.thunder_cache

        return res

    def _compile_model(self, model):
        match self.config.mode:
            case "eager":
                return model
            case "inductor":
                if (self.config.enable_cudagraph):
                    return torch.compile(model, mode="reduce-overhead")
                else:
                    return torch.compile(model, mode="default")
            case "thunder":
                return thunderfx(model, **self._thunder_jit_options)
            case "thunderjit":
                return thunder.jit(model, **self._thunder_jit_options)
            case _:
                raise ValueError(f"Unknown mode: {self.config.mode}")

    def _load_model(self) -> torch.nn.Module:
        """Load the model based on configuration"""
        model_id = self.config.model_name

        config = None
        if self.config.use_hardcoded_model:
            config = Llama4TextConfig.from_dict(
                json.loads(llama_4_Maverick_17B_128E_cfg_str)
            )
        else:
            config = AutoConfig.from_pretrained(model_id)

        if hasattr(config, "text_config"):
            config = config.text_config
        if self.config.num_layers:
            config.num_hidden_layers = self.config.num_layers

        self.hf_config = config

        with torch.device("meta"):
            model = AutoModelForCausalLM.from_config(
                config, torch_dtype=torch.bfloat16, attn_implementation=self.config.attn_implementation
            )

        return model

    def generate_batch(self) -> tuple[torch.Tensor, HybridChunkedCache]:
        """Generate a batch of input tokens"""
        batch_size = self.config.batch_size
        input_length = self.config.input_length

        input_ids = torch.randint(0, self.vocab_size, (batch_size, input_length), device=DEVICE)
        if LooseVersion(transformers.__version__) >= LooseVersion("4.55"):
            # Transformers deprecated HybridChunkedCache in favour of static in 4.55.x
            # NOTE: tp_size should only reflect tensor parallelism size, not total world size.
            # If 2D parallelism (data parallel + tensor parallel) is added in the future,
            # tp_size should be set to the tensor parallelism size only (e.g., WORLD_SIZE // DP_SIZE),
            # not WORLD_SIZE, so that StaticCache correctly handles sharded KV heads.
            past_key_values = StaticCache(
                config=self.hf_config,
                max_batch_size=input_ids.shape[0],
                max_cache_len=input_ids.shape[1] + self.config.output_length,
                device=DEVICE,
                dtype=torch.bfloat16,
                tp_size=WORLD_SIZE,


            config = Llama4TextConfig.from_dict(
                json.loads(llama_4_Maverick_17B_128E_cfg_str)
            )
        else:
            config = AutoConfig.from_pretrained(model_id)

        if hasattr(config, "text_config"):
            config = config.text_config
        if self.config.num_layers:
            config.num_hidden_layers = self.config.num_layers

        self.hf_config = config

        with torch.device("meta"):
            model = AutoModelForCausalLM.from_config(
                config, torch_dtype=torch.bfloat16, attn_implementation=self.config.attn_implementation
            )

        return model

    def generate_batch(self) -> tuple[torch.Tensor, HybridChunkedCache]:
        """Generate a batch of input tokens"""
        batch_size = self.config.batch_size
        input_length = self.config.input_length

        input_ids = torch.randint(0, self.vocab_size, (batch_size, input_length), device=DEVICE)
        if LooseVersion(transformers.__version__) >= LooseVersion("4.55"):
            # Transformers deprecated HybridChunkedCache in favour of static in 4.55.x
            # NOTE: tp_size should only reflect tensor parallelism size, not total world size.
            # If 2D parallelism (data parallel + tensor parallel) is added in the future,
            # tp_size should be set to the tensor parallelism size only (e.g., WORLD_SIZE // DP_SIZE),
            # not WORLD_SIZE, so that StaticCache correctly handles sharded KV heads.
            past_key_values = StaticCache(
                config=self.hf_config,
                max_batch_size=input_ids.shape[0],
                max_cache_len=input_ids.shape[1] + self.config.output_length,
                device=DEVICE,
                dtype=torch.bfloat16,
                tp_size=WORLD_SIZE,
            )
        else:
            past_key_values = HybridChunkedCache(
                self.hf_config, input_ids.shape[0], input_ids.shape[1] + self.config.output_length
            )
            for layer_idx in range(self.hf_config.num_hidden_layers):
                # key_states.shape[1] is used to retrieve the number of key value heads, all other dimensions can be 1 and ignored
                # https://github.com/huggingface/transformers/blob/9300728665aaeb0ebf4db99f9d9fbce916b4a183/src/transformers/cache_utils.py#L1822
                dummy_key_states = torch.empty(1, self.hf_config.num_key_value_heads // WORLD_SIZE, 1, 1, device=DEVICE)
                past_key_values.initialise_cache_layer(layer_idx, dummy_key_states)

        return input_ids, past_key_values

    def get_next_token(
        self, input_ids: torch.Tensor, past_key_values: HybridChunkedCache | StaticCache
    ) -> torch.Tensor:
        with torch.no_grad():
            outputs = self.model(input_ids, past_key_values=past_key_values, use_cache=True)
        logits = outputs.logits  # [B, seq_len, vocab_size]
        next_token_logits = logits[:, -1, :]
        next_token = torch.argmax(next_token_logits, dim=-1, keepdim=True)
        return next_token

    def prefill(self, input_ids: torch.Tensor, past_key_values: HybridChunkedCache) -> torch.Tensor:
        """
        Prefill phase: Process the entire input prompt at once.
        Returns the next token.

        Similar to: https://github.com/pytorch-labs/gpt-fast/blob/main/generate.py#L68-L82
        """
        return self.get_next_token(input_ids, past_key_values)


            past_key_values = StaticCache(
                config=self.hf_config,
                max_batch_size=input_ids.shape[0],
                max_cache_len=input_ids.shape[1] + self.config.output_length,
                device=DEVICE,
                dtype=torch.bfloat16,
                tp_size=WORLD_SIZE,
            )
        else:
            past_key_values = HybridChunkedCache(
                self.hf_config, input_ids.shape[0], input_ids.shape[1] + self.config.output_length
            )
            for layer_idx in range(self.hf_config.num_hidden_layers):
                # key_states.shape[1] is used to retrieve the number of key value heads, all other dimensions can be 1 and ignored
                # https://github.com/huggingface/transformers/blob/9300728665aaeb0ebf4db99f9d9fbce916b4a183/src/transformers/cache_utils.py#L1822
                dummy_key_states = torch.empty(1, self.hf_config.num_key_value_heads // WORLD_SIZE, 1, 1, device=DEVICE)
                past_key_values.initialise_cache_layer(layer_idx, dummy_key_states)

        return input_ids, past_key_values

    def get_next_token(
        self, input_ids: torch.Tensor, past_key_values: HybridChunkedCache | StaticCache
    ) -> torch.Tensor:
        with torch.no_grad():
            outputs = self.model(input_ids, past_key_values=past_key_values, use_cache=True)
        logits = outputs.logits  # [B, seq_len, vocab_size]
        next_token_logits = logits[:, -1, :]
        next_token = torch.argmax(next_token_logits, dim=-1, keepdim=True)
        return next_token

    def prefill(self, input_ids: torch.Tensor, past_key_values: HybridChunkedCache) -> torch.Tensor:
        """
        Prefill phase: Process the entire input prompt at once.
        Returns the next token.

        Similar to: https://github.com/pytorch-labs/gpt-fast/blob/main/generate.py#L68-L82
        """
        return self.get_next_token(input_ids, past_key_values)

    def decode_one_token(self, input_ids: torch.Tensor, past_key_values: HybridChunkedCache) -> torch.Tensor:
        """
        Decode phase: Generate a single token given the current sequence.
        Returns the next token.
        """
        # input_pos: [B, 1] One token at the time
        assert input_ids.shape[-1] == 1, f"Expected shape (B, 1), but found {input_ids.shape}"
        return self.get_next_token(input_ids, past_key_values)

    # TODO: Running `torchrun --nproc-per-node 2 thunder/benchmarks/benchmark_inference.py --input-length 32 --output-length 32 --mode eager --num-iterations 10`


prefill

                self.hf_config, input_ids.shape[0], input_ids.shape[1] + self.config.output_length
            )
            for layer_idx in range(self.hf_config.num_hidden_layers):
                # key_states.shape[1] is used to retrieve the number of key value heads, all other dimensions can be 1 and ignored
                # https://github.com/huggingface/transformers/blob/9300728665aaeb0ebf4db99f9d9fbce916b4a183/src/transformers/cache_utils.py#L1822
                dummy_key_states = torch.empty(1, self.hf_config.num_key_value_heads // WORLD_SIZE, 1, 1, device=DEVICE)
                past_key_values.initialise_cache_layer(layer_idx, dummy_key_states)

        return input_ids, past_key_values

    def get_next_token(
        self, input_ids: torch.Tensor, past_key_values: HybridChunkedCache | StaticCache
    ) -> torch.Tensor:
        with torch.no_grad():
            outputs = self.model(input_ids, past_key_values=past_key_values, use_cache=True)
        logits = outputs.logits  # [B, seq_len, vocab_size]
        next_token_logits = logits[:, -1, :]
        next_token = torch.argmax(next_token_logits, dim=-1, keepdim=True)
        return next_token

    def prefill(self, input_ids: torch.Tensor, past_key_values: HybridChunkedCache) -> torch.Tensor:
        """
        Prefill phase: Process the entire input prompt at once.
        Returns the next token.

        Similar to: https://github.com/pytorch-labs/gpt-fast/blob/main/generate.py#L68-L82
        """
        return self.get_next_token(input_ids, past_key_values)

    def decode_one_token(self, input_ids: torch.Tensor, past_key_values: HybridChunkedCache) -> torch.Tensor:
        """
        Decode phase: Generate a single token given the current sequence.
        Returns the next token.
        """
        # input_pos: [B, 1] One token at the time
        assert input_ids.shape[-1] == 1, f"Expected shape (B, 1), but found {input_ids.shape}"
        return self.get_next_token(input_ids, past_key_values)

    # TODO: Running `torchrun --nproc-per-node 2 thunder/benchmarks/benchmark_inference.py --input-length 32 --output-length 32 --mode eager --num-iterations 10`
    # with inference mode results in
    # [rank1]:   File "/opt/pytorch/lightning-thunder/thunder/benchmarks/layers_for_inference_benchmark.py", line 358, in grouped_mm
    # [rank1]:     group_outs.append(group_a @ b[idx])
    # [rank1]:                                 ~^^^^^
    # [rank1]: RuntimeError: Cannot set version_counter for inference tensor
    # @torch.inference_mode()
    def generate(
        self, input_ids: torch.Tensor, max_new_tokens: int, past_key_values: HybridChunkedCache
    ) -> dict[str, Any]:


    def get_next_token(
        self, input_ids: torch.Tensor, past_key_values: HybridChunkedCache | StaticCache
    ) -> torch.Tensor:
        with torch.no_grad():
            outputs = self.model(input_ids, past_key_values=past_key_values, use_cache=True)
        logits = outputs.logits  # [B, seq_len, vocab_size]
        next_token_logits = logits[:, -1, :]
        next_token = torch.argmax(next_token_logits, dim=-1, keepdim=True)
        return next_token

    def prefill(self, input_ids: torch.Tensor, past_key_values: HybridChunkedCache) -> torch.Tensor:
        """
        Prefill phase: Process the entire input prompt at once.
        Returns the next token.

        Similar to: https://github.com/pytorch-labs/gpt-fast/blob/main/generate.py#L68-L82
        """
        return self.get_next_token(input_ids, past_key_values)

    def decode_one_token(self, input_ids: torch.Tensor, past_key_values: HybridChunkedCache) -> torch.Tensor:
        """
        Decode phase: Generate a single token given the current sequence.
        Returns the next token.
        """
        # input_pos: [B, 1] One token at the time
        assert input_ids.shape[-1] == 1, f"Expected shape (B, 1), but found {input_ids.shape}"
        return self.get_next_token(input_ids, past_key_values)

    # TODO: Running `torchrun --nproc-per-node 2 thunder/benchmarks/benchmark_inference.py --input-length 32 --output-length 32 --mode eager --num-iterations 10`
    # with inference mode results in
    # [rank1]:   File "/opt/pytorch/lightning-thunder/thunder/benchmarks/layers_for_inference_benchmark.py", line 358, in grouped_mm
    # [rank1]:     group_outs.append(group_a @ b[idx])
    # [rank1]:                                 ~^^^^^
    # [rank1]: RuntimeError: Cannot set version_counter for inference tensor
    # @torch.inference_mode()
    def generate(
        self, input_ids: torch.Tensor, max_new_tokens: int, past_key_values: HybridChunkedCache
    ) -> dict[str, Any]:
        """
        Generate tokens using separate prefill and decode phases.
        Returns detailed metrics for both phases.
        """
        # Prefill phase - process the entire prompt
        with timer() as prefill_timer:
            torch.cuda.nvtx.range_push("prefill")
            first_token = self.prefill(input_ids, past_key_values)
            torch.cuda.nvtx.range_pop()


generate

        Similar to: https://github.com/pytorch-labs/gpt-fast/blob/main/generate.py#L68-L82
        """
        return self.get_next_token(input_ids, past_key_values)

    def decode_one_token(self, input_ids: torch.Tensor, past_key_values: HybridChunkedCache) -> torch.Tensor:
        """
        Decode phase: Generate a single token given the current sequence.
        Returns the next token.
        """
        # input_pos: [B, 1] One token at the time
        assert input_ids.shape[-1] == 1, f"Expected shape (B, 1), but found {input_ids.shape}"
        return self.get_next_token(input_ids, past_key_values)

    # TODO: Running `torchrun --nproc-per-node 2 thunder/benchmarks/benchmark_inference.py --input-length 32 --output-length 32 --mode eager --num-iterations 10`
    # with inference mode results in
    # [rank1]:   File "/opt/pytorch/lightning-thunder/thunder/benchmarks/layers_for_inference_benchmark.py", line 358, in grouped_mm
    # [rank1]:     group_outs.append(group_a @ b[idx])
    # [rank1]:                                 ~^^^^^
    # [rank1]: RuntimeError: Cannot set version_counter for inference tensor
    # @torch.inference_mode()
    def generate(
        self, input_ids: torch.Tensor, max_new_tokens: int, past_key_values: HybridChunkedCache
    ) -> dict[str, Any]:
        """
        Generate tokens using separate prefill and decode phases.
        Returns detailed metrics for both phases.
        """
        # Prefill phase - process the entire prompt
        with timer() as prefill_timer:
            torch.cuda.nvtx.range_push("prefill")
            first_token = self.prefill(input_ids, past_key_values)
            torch.cuda.nvtx.range_pop()
        prefill_time = prefill_timer()
        generated_tokens = [first_token]

        # Decode phase - generate remaining tokens one by one
        next_token = first_token
        with timer() as decode_timer:
            for idx in range(max_new_tokens - 1):
                torch.cuda.nvtx.range_push("decode:" + str(idx))
                next_token = self.decode_one_token(next_token, past_key_values)
                generated_tokens.append(next_token)
                torch.cuda.nvtx.range_pop()

        total_decode_time = decode_timer()

        return {
            "prefill_time_ms": prefill_time,
            "decode_time_ms": total_decode_time,
            "generated_tokens": generated_tokens,
            "total_tokens": max_new_tokens,
        }

    def measure_inference_step(
        self, input_ids: torch.Tensor, past_key_values: HybridChunkedCache, max_new_tokens: int
    ) -> dict[str, float]:
        """Measure a single inference step with detailed timing using separate prefill/decode"""
        # Generate tokens with separate prefill/decode tracking
        generation_result = self.generate(input_ids, max_new_tokens, past_key_values)
        total_time = generation_result["prefill_time_ms"] + generation_result["decode_time_ms"]

        # Extract metrics
        ttft = generation_result["prefill_time_ms"]  # Time to first token is the prefill time
        total_decode_time = generation_result["decode_time_ms"]
        avg_tbot = total_decode_time / (max_new_tokens - 1) if max_new_tokens > 1 else 0

        # Calculate throughput
        total_tokens = self.config.output_length * self.config.batch_size
        throughput = (total_tokens / total_time) * 1000  # tokens/second

        # Calculate separate prefill and decode throughput
        prefill_tokens = self.config.input_length * self.config.batch_size


        generated_tokens = [first_token]

        # Decode phase - generate remaining tokens one by one
        next_token = first_token
        with timer() as decode_timer:
            for idx in range(max_new_tokens - 1):
                torch.cuda.nvtx.range_push("decode:" + str(idx))
                next_token = self.decode_one_token(next_token, past_key_values)
                generated_tokens.append(next_token)
                torch.cuda.nvtx.range_pop()

        total_decode_time = decode_timer()

        return {
            "prefill_time_ms": prefill_time,
            "decode_time_ms": total_decode_time,
            "generated_tokens": generated_tokens,
            "total_tokens": max_new_tokens,
        }

    def measure_inference_step(
        self, input_ids: torch.Tensor, past_key_values: HybridChunkedCache, max_new_tokens: int
    ) -> dict[str, float]:
        """Measure a single inference step with detailed timing using separate prefill/decode"""
        # Generate tokens with separate prefill/decode tracking
        generation_result = self.generate(input_ids, max_new_tokens, past_key_values)
        total_time = generation_result["prefill_time_ms"] + generation_result["decode_time_ms"]

        # Extract metrics
        ttft = generation_result["prefill_time_ms"]  # Time to first token is the prefill time
        total_decode_time = generation_result["decode_time_ms"]
        avg_tbot = total_decode_time / (max_new_tokens - 1) if max_new_tokens > 1 else 0

        # Calculate throughput
        total_tokens = self.config.output_length * self.config.batch_size
        throughput = (total_tokens / total_time) * 1000  # tokens/second

        # Calculate separate prefill and decode throughput
        prefill_tokens = self.config.input_length * self.config.batch_size
        prefill_throughput = (prefill_tokens / generation_result["prefill_time_ms"]) * 1000

        decode_tokens = (self.config.output_length - 1) * self.config.batch_size
        decode_throughput = (decode_tokens / total_decode_time) * 1000 if total_decode_time > 0 else 0

        return {
            "ttft": ttft,
            "avg_tbot": avg_tbot,
            "total_time": total_time,
            "throughput": throughput,
            "prefill_throughput": prefill_throughput,
            "decode_throughput": decode_throughput,
            "prefill_time": generation_result["prefill_time_ms"],
            "total_decode_time": total_decode_time,
        }

    def run_benchmark(self) -> InferenceMetrics:
        """Run the full benchmark and collect metrics"""
        print(f"Running inference benchmark for {self.config.model_name}")

        print(f"Batch size: {self.config.batch_size}")
        print(f"Input length: {self.config.input_length}")
        print(f"Output length: {self.config.output_length}")
        print(f"Mode: {self.config.mode}")

        print(f"\nWarming up with {self.config.warmup_iterations} iterations...")
        input_ids, past_key_values = self.generate_batch()

        for idx in tqdm(range(self.config.warmup_iterations), disable=LOCAL_RANK != 0):
            past_key_values.reset()
            # Use output_length to warm up sufficiently. Otherwise, Thunder's
            # first-run latency is terribly slow due to lack of dynamic shape
            # support.
            torch.cuda.nvtx.range_push("warmup:" + str(idx))
            _ = self.measure_inference_step(input_ids, past_key_values, self.config.output_length)


        throughput = (total_tokens / total_time) * 1000  # tokens/second

        # Calculate separate prefill and decode throughput
        prefill_tokens = self.config.input_length * self.config.batch_size
        prefill_throughput = (prefill_tokens / generation_result["prefill_time_ms"]) * 1000

        decode_tokens = (self.config.output_length - 1) * self.config.batch_size
        decode_throughput = (decode_tokens / total_decode_time) * 1000 if total_decode_time > 0 else 0

        return {
            "ttft": ttft,
            "avg_tbot": avg_tbot,
            "total_time": total_time,
            "throughput": throughput,
            "prefill_throughput": prefill_throughput,
            "decode_throughput": decode_throughput,
            "prefill_time": generation_result["prefill_time_ms"],
            "total_decode_time": total_decode_time,
        }

    def run_benchmark(self) -> InferenceMetrics:
        """Run the full benchmark and collect metrics"""
        print(f"Running inference benchmark for {self.config.model_name}")

        print(f"Batch size: {self.config.batch_size}")
        print(f"Input length: {self.config.input_length}")
        print(f"Output length: {self.config.output_length}")
        print(f"Mode: {self.config.mode}")

        print(f"\nWarming up with {self.config.warmup_iterations} iterations...")
        input_ids, past_key_values = self.generate_batch()

        for idx in tqdm(range(self.config.warmup_iterations), disable=LOCAL_RANK != 0):
            past_key_values.reset()
            # Use output_length to warm up sufficiently. Otherwise, Thunder's
            # first-run latency is terribly slow due to lack of dynamic shape
            # support.
            torch.cuda.nvtx.range_push("warmup:" + str(idx))
            _ = self.measure_inference_step(input_ids, past_key_values, self.config.output_length)
            torch.cuda.nvtx.range_pop()

        print(f"\nRunning {self.config.num_iterations} benchmark iterations...")
        all_metrics = []

        if torch.cuda.is_available():
            torch.cuda.reset_peak_memory_stats()

        for idx in tqdm(range(self.config.num_iterations), disable=LOCAL_RANK != 0):
            past_key_values.reset()

            is_under_nsys = bool(os.environ.get("NSYS_PROFILING_SESSION_ID"))
            # Wrap each non-warmup iteration with cudaProfilerStart() and
            # cudaProfilerStop(). This allows the user to run
            # ```shell
            # nsys profile --capture-range=cudaProfilerApi --capture-range-end=repeat:<N> ...
            # ```
            # to record only the non-warmup iterations.
            # Note when debug_moe is on, we'll defer the flag inside the moe wrapper.
            self.profiler_toggle = is_under_nsys
            if is_under_nsys and not self.config.debug_moe:
                torch.cuda.cudart().cudaProfilerStart()
            torch.cuda.nvtx.range_push("step:" + str(idx))
            iter_metrics = self.measure_inference_step(input_ids, past_key_values, self.config.output_length)
            torch.cuda.nvtx.range_pop()
            if is_under_nsys and not self.config.debug_moe:
                torch.cuda.cudart().cudaProfilerStop()

            all_metrics.append(iter_metrics)

            # Track metrics
            self.metrics.iteration_times.append(iter_metrics["total_time"])
            self.metrics.ttft_times.append(iter_metrics["ttft"])
            self.metrics.prefill_times.append(iter_metrics["prefill_time"])
            self.metrics.decode_times.append(iter_metrics["total_decode_time"])

        self._calculate_aggregate_metrics(all_metrics)

        if torch.cuda.is_available():
            self.metrics.memory_used_gb = torch.cuda.memory_allocated() / 1e9
            self.metrics.peak_memory_gb = torch.cuda.max_memory_allocated() / 1e9

        if self.config.fx_report_folder is not None and self.config.mode == "thunder":
            self.model._backend.save_reproducer_to_folder(self.config.fx_report_folder)
            return

        return self.metrics

    def _calculate_aggregate_metrics(self, all_metrics: list[dict[str, Any]]):
        """Calculate aggregate metrics from individual iterations"""
        # Average throughput
        throughputs = [m["throughput"] for m in all_metrics]
        self.metrics.throughput_tokens_per_sec = statistics.mean(throughputs)

        # Average latency
        total_times = [m["total_time"] for m in all_metrics]
        total_tokens = self.config.output_length * self.config.batch_size
        self.metrics.latency_ms_per_token = statistics.mean(total_times) / total_tokens

        # TTFT
        ttfts = [m["ttft"] for m in all_metrics]
        self.metrics.time_to_first_token_ms = statistics.mean(ttfts)

        # TBOT
        self.metrics.time_between_output_tokens_ms = statistics.mean([m["avg_tbot"] for m in all_metrics])

        # Total time


            all_metrics.append(iter_metrics)

            # Track metrics
            self.metrics.iteration_times.append(iter_metrics["total_time"])
            self.metrics.ttft_times.append(iter_metrics["ttft"])
            self.metrics.prefill_times.append(iter_metrics["prefill_time"])
            self.metrics.decode_times.append(iter_metrics["total_decode_time"])

        self._calculate_aggregate_metrics(all_metrics)

        if torch.cuda.is_available():
            self.metrics.memory_used_gb = torch.cuda.memory_allocated() / 1e9
            self.metrics.peak_memory_gb = torch.cuda.max_memory_allocated() / 1e9

        if self.config.fx_report_folder is not None and self.config.mode == "thunder":
            self.model._backend.save_reproducer_to_folder(self.config.fx_report_folder)
            return

        return self.metrics

    def _calculate_aggregate_metrics(self, all_metrics: list[dict[str, Any]]):
        """Calculate aggregate metrics from individual iterations"""
        # Average throughput
        throughputs = [m["throughput"] for m in all_metrics]
        self.metrics.throughput_tokens_per_sec = statistics.mean(throughputs)

        # Average latency
        total_times = [m["total_time"] for m in all_metrics]
        total_tokens = self.config.output_length * self.config.batch_size
        self.metrics.latency_ms_per_token = statistics.mean(total_times) / total_tokens

        # TTFT
        ttfts = [m["ttft"] for m in all_metrics]
        self.metrics.time_to_first_token_ms = statistics.mean(ttfts)

        # TBOT
        self.metrics.time_between_output_tokens_ms = statistics.mean([m["avg_tbot"] for m in all_metrics])

        # Total time
        self.metrics.total_time_ms = statistics.mean(total_times)

        # Prefill metrics
        prefill_throughputs = [m["prefill_throughput"] for m in all_metrics]
        self.metrics.prefill_throughput_tokens_per_sec = statistics.mean(prefill_throughputs)
        prefill_times = [m["prefill_time"] for m in all_metrics]
        self.metrics.prefill_time_ms = statistics.mean(prefill_times)

        # Decode metrics
        decode_throughputs = [m["decode_throughput"] for m in all_metrics]
        self.metrics.decode_throughput_tokens_per_sec = statistics.mean(decode_throughputs)
        decode_times = [m["total_decode_time"] for m in all_metrics]
        self.metrics.decode_time_ms = statistics.mean(decode_times)

    def print_results(self):
        """Print benchmark results in a formatted way"""
        print("\n" + "=" * 60)
        print(f"BENCHMARK RESULTS - {self.config.model_name} {self.config.mode}")
        print("=" * 60)

        print("\nThroughput Metrics:")
        print(f"  Overall Throughput: {self.metrics.throughput_tokens_per_sec:.2f} tokens/sec")
        print(f"  Prefill Throughput: {self.metrics.prefill_throughput_tokens_per_sec:.2f} tokens/sec")
        print(f"  Decode Throughput: {self.metrics.decode_throughput_tokens_per_sec:.2f} tokens/sec")
        print(f"  Latency: {self.metrics.latency_ms_per_token:.2f} ms/token")

        print("\nLatency Breakdown:")
        print(f"  Time to First Token (TTFT): {self.metrics.time_to_first_token_ms:.2f} ms")
        print(f"  Time Between Output Tokens (TBOT): {self.metrics.time_between_output_tokens_ms:.2f} ms")
        print(f"  Prefill Time: {self.metrics.prefill_time_ms:.2f} ms")
        print(f"  Decode Time: {self.metrics.decode_time_ms:.2f} ms")
        print(f"  Total Generation Time: {self.metrics.total_time_ms:.2f} ms")


        self.metrics.time_to_first_token_ms = statistics.mean(ttfts)

        # TBOT
        self.metrics.time_between_output_tokens_ms = statistics.mean([m["avg_tbot"] for m in all_metrics])

        # Total time
        self.metrics.total_time_ms = statistics.mean(total_times)

        # Prefill metrics
        prefill_throughputs = [m["prefill_throughput"] for m in all_metrics]
        self.metrics.prefill_throughput_tokens_per_sec = statistics.mean(prefill_throughputs)
        prefill_times = [m["prefill_time"] for m in all_metrics]
        self.metrics.prefill_time_ms = statistics.mean(prefill_times)

        # Decode metrics
        decode_throughputs = [m["decode_throughput"] for m in all_metrics]
        self.metrics.decode_throughput_tokens_per_sec = statistics.mean(decode_throughputs)
        decode_times = [m["total_decode_time"] for m in all_metrics]
        self.metrics.decode_time_ms = statistics.mean(decode_times)

    def print_results(self):
        """Print benchmark results in a formatted way"""
        print("\n" + "=" * 60)
        print(f"BENCHMARK RESULTS - {self.config.model_name} {self.config.mode}")
        print("=" * 60)

        print("\nThroughput Metrics:")
        print(f"  Overall Throughput: {self.metrics.throughput_tokens_per_sec:.2f} tokens/sec")
        print(f"  Prefill Throughput: {self.metrics.prefill_throughput_tokens_per_sec:.2f} tokens/sec")
        print(f"  Decode Throughput: {self.metrics.decode_throughput_tokens_per_sec:.2f} tokens/sec")
        print(f"  Latency: {self.metrics.latency_ms_per_token:.2f} ms/token")

        print("\nLatency Breakdown:")
        print(f"  Time to First Token (TTFT): {self.metrics.time_to_first_token_ms:.2f} ms")
        print(f"  Time Between Output Tokens (TBOT): {self.metrics.time_between_output_tokens_ms:.2f} ms")
        print(f"  Prefill Time: {self.metrics.prefill_time_ms:.2f} ms")
        print(f"  Decode Time: {self.metrics.decode_time_ms:.2f} ms")
        print(f"  Total Generation Time: {self.metrics.total_time_ms:.2f} ms")

        print("\nMemory Usage:")
        print(f"  Current Memory: {self.metrics.memory_used_gb:.2f} GB")
        print(f"  Peak Memory: {self.metrics.peak_memory_gb:.2f} GB")

        if len(self.metrics.iteration_times) > 1:
            print("\nVariance Analysis:")
            print(f"  Throughput Std Dev: {statistics.stdev(self.metrics.iteration_times):.2f} ms")
            print(f"  TTFT Std Dev: {statistics.stdev(self.metrics.ttft_times):.2f} ms")

    def save_results(self, filename: str):
        """Save results to JSON file"""
        results = {
            "config": self.config.__dict__,
            "metrics": {
                "throughput_tokens_per_sec": self.metrics.throughput_tokens_per_sec,
                "prefill_throughput_tokens_per_sec": self.metrics.prefill_throughput_tokens_per_sec,
                "decode_throughput_tokens_per_sec": self.metrics.decode_throughput_tokens_per_sec,
                "latency_ms_per_token": self.metrics.latency_ms_per_token,
                "time_to_first_token_ms": self.metrics.time_to_first_token_ms,
                "time_between_output_tokens_ms": self.metrics.time_between_output_tokens_ms,
                "prefill_time_ms": self.metrics.prefill_time_ms,
                "decode_time_ms": self.metrics.decode_time_ms,
                "total_time_ms": self.metrics.total_time_ms,
                "memory_used_gb": self.metrics.memory_used_gb,
                "peak_memory_gb": self.metrics.peak_memory_gb,
            },
            "detailed_metrics": {
                "iteration_times": self.metrics.iteration_times,


        print(f"  Prefill Throughput: {self.metrics.prefill_throughput_tokens_per_sec:.2f} tokens/sec")
        print(f"  Decode Throughput: {self.metrics.decode_throughput_tokens_per_sec:.2f} tokens/sec")
        print(f"  Latency: {self.metrics.latency_ms_per_token:.2f} ms/token")

        print("\nLatency Breakdown:")
        print(f"  Time to First Token (TTFT): {self.metrics.time_to_first_token_ms:.2f} ms")
        print(f"  Time Between Output Tokens (TBOT): {self.metrics.time_between_output_tokens_ms:.2f} ms")
        print(f"  Prefill Time: {self.metrics.prefill_time_ms:.2f} ms")
        print(f"  Decode Time: {self.metrics.decode_time_ms:.2f} ms")
        print(f"  Total Generation Time: {self.metrics.total_time_ms:.2f} ms")

        print("\nMemory Usage:")
        print(f"  Current Memory: {self.metrics.memory_used_gb:.2f} GB")
        print(f"  Peak Memory: {self.metrics.peak_memory_gb:.2f} GB")

        if len(self.metrics.iteration_times) > 1:
            print("\nVariance Analysis:")
            print(f"  Throughput Std Dev: {statistics.stdev(self.metrics.iteration_times):.2f} ms")
            print(f"  TTFT Std Dev: {statistics.stdev(self.metrics.ttft_times):.2f} ms")

    def save_results(self, filename: str):
        """Save results to JSON file"""
        results = {
            "config": self.config.__dict__,
            "metrics": {
                "throughput_tokens_per_sec": self.metrics.throughput_tokens_per_sec,
                "prefill_throughput_tokens_per_sec": self.metrics.prefill_throughput_tokens_per_sec,
                "decode_throughput_tokens_per_sec": self.metrics.decode_throughput_tokens_per_sec,
                "latency_ms_per_token": self.metrics.latency_ms_per_token,
                "time_to_first_token_ms": self.metrics.time_to_first_token_ms,
                "time_between_output_tokens_ms": self.metrics.time_between_output_tokens_ms,
                "prefill_time_ms": self.metrics.prefill_time_ms,
                "decode_time_ms": self.metrics.decode_time_ms,
                "total_time_ms": self.metrics.total_time_ms,
                "memory_used_gb": self.metrics.memory_used_gb,
                "peak_memory_gb": self.metrics.peak_memory_gb,
            },
            "detailed_metrics": {
                "iteration_times": self.metrics.iteration_times,
                "ttft_times": self.metrics.ttft_times,
                "prefill_times": self.metrics.prefill_times,
                "decode_times": self.metrics.decode_times,
            },
        }

        with open(filename, "w") as f:
            json.dump(results, f, indent=2)

        print(f"\nResults saved to {filename}")

class CustomFormatter(argparse.RawDescriptionHelpFormatter, argparse.ArgumentDefaultsHelpFormatter):
    pass

def parse_args() -> argparse.Namespace:
    """Command line interface for the benchmark"""
    parser = argparse.ArgumentParser(
        description="Thunder Inference Benchmark",
        formatter_class=CustomFormatter,
        epilog="""
Examples:
  python benchmark_inference.py --input-length 2048 --output-length 512 --model-name meta-llama/Llama-4-Maverick-17B-128E --mode eager
        """,
    )

    parser.add_argument(
        "--model-name",
        type=str,


        # Materialize the model on the device (after Llama4MoE replacement and sharding)
        model.to_empty(device=DEVICE)
        assert all(p.device == DEVICE for p in model.parameters())

        # Required as thunder doesn't understand inference mode
        # And some prims like `prims._grouped_mm` don't have grad rule defined yet.
        for p in model.parameters():
            p.requires_grad_(False)

        assert all(not p.requires_grad for p in model.parameters())

        # `thunderfx` seems to hide the access to vocab_size somewhere so
        # store it here before any compiler is applied.
        self.vocab_size = model.vocab_size

        if self.config.enable_nvfp4:
            _quantize_llama4(model)

        # If debug_moe is on, we'll only compile the moe section, which is much easier to digest.
        if self.config.debug_moe:
            def compile_wrapper(model, jitter):
                class ModelWrapper(torch.nn.Module):

                    def __init__(self, model, jitter):
                        super().__init__()
                        self.model = jitter._compile_model(model)
                        assert not hasattr(jitter.model, "_backend")
                        # store a reference to the compiled backend, so we can print the trace later.
                        if jitter.config.mode == "thunder":
                            jitter.model._backend = self.model._backend
                        self.jitter = jitter

                    def forward(self, *args, **kwargs):
                        if self.jitter.profiler_toggle:
                            torch.cuda.cudart().cudaProfilerStart()
                        out = self.model(*args, **kwargs)
                        if self.jitter.profiler_toggle:
                            torch.cuda.cudart().cudaProfilerStop()
                        return out

                return ModelWrapper(model, jitter)

            self.model = model
            # reuse the `replace` function with compilation.
            _replace_with_custom_fn_if_matches_filter_with_name(
                model,
                lambda model, cur_fqn: compile_wrapper(model, self),
                lambda model, cur_fqn: isinstance(model, Llama4MoE),
            )
        else:
            self.model = self._compile_model(model)

    @property
    def _thunder_jit_options(self) -> dict[str, Any]:
        # `nv_enable_linear=True` might fail with distributed run
        # ref: https://github.com/NVIDIA/Fuser/issues/4507
        res = {"transforms": []}
        if self.config.enable_nv_linear:
            res["nv_enable_linear"] = True
            res["nv_enable_matmul"] = True


ModelWrapper

        model.to_empty(device=DEVICE)
        assert all(p.device == DEVICE for p in model.parameters())

        # Required as thunder doesn't understand inference mode
        # And some prims like `prims._grouped_mm` don't have grad rule defined yet.
        for p in model.parameters():
            p.requires_grad_(False)

        assert all(not p.requires_grad for p in model.parameters())

        # `thunderfx` seems to hide the access to vocab_size somewhere so
        # store it here before any compiler is applied.
        self.vocab_size = model.vocab_size

        if self.config.enable_nvfp4:
            _quantize_llama4(model)

        # If debug_moe is on, we'll only compile the moe section, which is much easier to digest.
        if self.config.debug_moe:
            def compile_wrapper(model, jitter):
                class ModelWrapper(torch.nn.Module):

                    def __init__(self, model, jitter):
                        super().__init__()
                        self.model = jitter._compile_model(model)
                        assert not hasattr(jitter.model, "_backend")
                        # store a reference to the compiled backend, so we can print the trace later.
                        if jitter.config.mode == "thunder":
                            jitter.model._backend = self.model._backend
                        self.jitter = jitter

                    def forward(self, *args, **kwargs):
                        if self.jitter.profiler_toggle:
                            torch.cuda.cudart().cudaProfilerStart()
                        out = self.model(*args, **kwargs)
                        if self.jitter.profiler_toggle:
                            torch.cuda.cudart().cudaProfilerStop()
                        return out

                return ModelWrapper(model, jitter)

            self.model = model
            # reuse the `replace` function with compilation.
            _replace_with_custom_fn_if_matches_filter_with_name(
                model,
                lambda model, cur_fqn: compile_wrapper(model, self),
                lambda model, cur_fqn: isinstance(model, Llama4MoE),
            )
        else:
            self.model = self._compile_model(model)

    @property
    def _thunder_jit_options(self) -> dict[str, Any]:
        # `nv_enable_linear=True` might fail with distributed run
        # ref: https://github.com/NVIDIA/Fuser/issues/4507
        res = {"transforms": []}
        if self.config.enable_nv_linear:


forward

        # store it here before any compiler is applied.
        self.vocab_size = model.vocab_size

        if self.config.enable_nvfp4:
            _quantize_llama4(model)

        # If debug_moe is on, we'll only compile the moe section, which is much easier to digest.
        if self.config.debug_moe:
            def compile_wrapper(model, jitter):
                class ModelWrapper(torch.nn.Module):

                    def __init__(self, model, jitter):
                        super().__init__()
                        self.model = jitter._compile_model(model)
                        assert not hasattr(jitter.model, "_backend")
                        # store a reference to the compiled backend, so we can print the trace later.
                        if jitter.config.mode == "thunder":
                            jitter.model._backend = self.model._backend
                        self.jitter = jitter

                    def forward(self, *args, **kwargs):
                        if self.jitter.profiler_toggle:
                            torch.cuda.cudart().cudaProfilerStart()
                        out = self.model(*args, **kwargs)
                        if self.jitter.profiler_toggle:
                            torch.cuda.cudart().cudaProfilerStop()
                        return out

                return ModelWrapper(model, jitter)

            self.model = model
            # reuse the `replace` function with compilation.
            _replace_with_custom_fn_if_matches_filter_with_name(
                model,
                lambda model, cur_fqn: compile_wrapper(model, self),
                lambda model, cur_fqn: isinstance(model, Llama4MoE),
            )
        else:
            self.model = self._compile_model(model)

    @property
    def _thunder_jit_options(self) -> dict[str, Any]:
        # `nv_enable_linear=True` might fail with distributed run
        # ref: https://github.com/NVIDIA/Fuser/issues/4507
        res = {"transforms": []}
        if self.config.enable_nv_linear:


# SPDX-FileCopyrightText: Copyright (c) 2025-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
import pytest

import torch

from .core import (
    run_benchmark,
    clear_dynamo_cache,
    unary_bwd_torch,
    with_executor,
    DEFAULT_EXECUTORS,
)
from .global_params import FLOAT_DTYPES
from .torch_ops import embedding

# (vocab, hidden) configurations seen in models.
EMBEDDING_CONFIGS = [
    (152064, 3584),  # hf_qwen2
    (32064, 3072),  # hf_phi3
    (131072, 5120),  # hf_mistral_nemo
]

SEQ_LENGTHS = [
    1024,
    2048,
    3072,
    4096,
    8192,
    12288,
    16384,
    20480,
    24576,
    28672,
    32768,
]

@pytest.mark.parametrize("executor", DEFAULT_EXECUTORS)
@pytest.mark.parametrize("seq_length", SEQ_LENGTHS)
@pytest.mark.parametrize("vocab_hidden", EMBEDDING_CONFIGS)
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
def test_embedding_fwd_baseline_benchmark(
    benchmark,
    seq_length: int,
    vocab_hidden: tuple,
    dtype: torch.dtype,
    executor: str,
):
    kwargs = {}
    if executor == "torchcompile":
        clear_dynamo_cache()
    if executor == "thunder":
        kwargs["nv_enable_embedding"] = True

    indices = torch.randint(0, vocab_hidden[0], (seq_length,), device="cuda")
    embedding_table = torch.randn(
        vocab_hidden, device="cuda", dtype=dtype, requires_grad=True
    )

    benchmark_fn = with_executor(executor, embedding, **kwargs)
    run_benchmark(
        benchmark,
        benchmark_fn,
        [indices, embedding_table],
    )

@pytest.mark.parametrize("executor", DEFAULT_EXECUTORS)
@pytest.mark.parametrize("seq_length", SEQ_LENGTHS)
@pytest.mark.parametrize("vocab_hidden", EMBEDDING_CONFIGS)
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
def test_embedding_bwd_baseline_benchmark(
    benchmark,
    seq_length: int,
    vocab_hidden: tuple,
    dtype: torch.dtype,
    executor: str,
):
    kwargs = {}
    if executor == "torchcompile":
        clear_dynamo_cache()
    if executor == "thunder":
        kwargs["nv_enable_embedding"] = True

    indices = torch.randint(0, vocab_hidden[0], (seq_length,), device="cuda")
    grads = torch.randn((seq_length, vocab_hidden[1]), device="cuda", dtype=dtype)
    embedding_table = torch.randn(
        vocab_hidden, device="cuda", dtype=dtype, requires_grad=True
    )

    # Compile the fwd fn for torchcompile
    fwd_fn = with_executor(executor, embedding, **kwargs)
    fwd_inputs = [indices, embedding_table]
    outputs = fwd_fn(fwd_inputs)

    run_benchmark(
        benchmark,
        unary_bwd_torch,
        [outputs, grads, *fwd_inputs],
    )


SEQ_LENGTHS = [
    1024,
    2048,
    3072,
    4096,
    8192,
    12288,
    16384,
    20480,
    24576,
    28672,
    32768,
]

@pytest.mark.parametrize("executor", DEFAULT_EXECUTORS)
@pytest.mark.parametrize("seq_length", SEQ_LENGTHS)
@pytest.mark.parametrize("vocab_hidden", EMBEDDING_CONFIGS)
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
def test_embedding_fwd_baseline_benchmark(
    benchmark,
    seq_length: int,
    vocab_hidden: tuple,
    dtype: torch.dtype,
    executor: str,
):
    kwargs = {}
    if executor == "torchcompile":
        clear_dynamo_cache()
    if executor == "thunder":
        kwargs["nv_enable_embedding"] = True

    indices = torch.randint(0, vocab_hidden[0], (seq_length,), device="cuda")
    embedding_table = torch.randn(
        vocab_hidden, device="cuda", dtype=dtype, requires_grad=True
    )

    benchmark_fn = with_executor(executor, embedding, **kwargs)
    run_benchmark(
        benchmark,
        benchmark_fn,
        [indices, embedding_table],
    )

@pytest.mark.parametrize("executor", DEFAULT_EXECUTORS)
@pytest.mark.parametrize("seq_length", SEQ_LENGTHS)
@pytest.mark.parametrize("vocab_hidden", EMBEDDING_CONFIGS)
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
def test_embedding_bwd_baseline_benchmark(
    benchmark,
    seq_length: int,
    vocab_hidden: tuple,
    dtype: torch.dtype,
    executor: str,
):
    kwargs = {}
    if executor == "torchcompile":
        clear_dynamo_cache()
    if executor == "thunder":
        kwargs["nv_enable_embedding"] = True

    indices = torch.randint(0, vocab_hidden[0], (seq_length,), device="cuda")


    if executor == "thunder":
        kwargs["nv_enable_embedding"] = True

    indices = torch.randint(0, vocab_hidden[0], (seq_length,), device="cuda")
    embedding_table = torch.randn(
        vocab_hidden, device="cuda", dtype=dtype, requires_grad=True
    )

    benchmark_fn = with_executor(executor, embedding, **kwargs)
    run_benchmark(
        benchmark,
        benchmark_fn,
        [indices, embedding_table],
    )

@pytest.mark.parametrize("executor", DEFAULT_EXECUTORS)
@pytest.mark.parametrize("seq_length", SEQ_LENGTHS)
@pytest.mark.parametrize("vocab_hidden", EMBEDDING_CONFIGS)
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
def test_embedding_bwd_baseline_benchmark(
    benchmark,
    seq_length: int,
    vocab_hidden: tuple,
    dtype: torch.dtype,
    executor: str,
):
    kwargs = {}
    if executor == "torchcompile":
        clear_dynamo_cache()
    if executor == "thunder":
        kwargs["nv_enable_embedding"] = True

    indices = torch.randint(0, vocab_hidden[0], (seq_length,), device="cuda")
    grads = torch.randn((seq_length, vocab_hidden[1]), device="cuda", dtype=dtype)
    embedding_table = torch.randn(
        vocab_hidden, device="cuda", dtype=dtype, requires_grad=True
    )

    # Compile the fwd fn for torchcompile
    fwd_fn = with_executor(executor, embedding, **kwargs)
    fwd_inputs = [indices, embedding_table]
    outputs = fwd_fn(fwd_inputs)

    run_benchmark(
        benchmark,
        unary_bwd_torch,
        [outputs, grads, *fwd_inputs],
    )


# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
import pytest
from nvfuser_direct import FusionDefinition, DataType
from nvfuser_direct.pytorch_utils import torch_dtype_to_nvfuser_dtype
from .core import run_benchmark, clear_dynamo_cache, with_executor, DEFAULT_EXECUTORS
import torch
from .global_params import generate_attn_inputs, FLOAT_DTYPES, PROMOTE_DTYPES
from .torch_ops import nanogpt_attn

# Fusion from nanogpt attention module
# The nvFuser defintion only includes the non-matmul computation (masked_fill + softmax + dropout)
def nanogpt_attn_fwd_fusion(
    fd: FusionDefinition, dtype: DataType, head_size: int, dropout_p: float
):
    T0 = fd.define_tensor(
        shape=[-1, -1, -1, -1],
        contiguity=[True, True, True, True],
        dtype=dtype,
        is_cpu=False,
    )
    T1 = fd.define_tensor(
        shape=[1, 1, -1, -1],
        contiguity=[None, None, False, True],
        dtype=DataType.Float,
        is_cpu=False,
    )

    if dtype in PROMOTE_DTYPES:
        T0 = fd.ops.cast(T0, dtype=DataType.Float)

    S2 = fd.define_scalar(1 / head_size**0.5, dtype=DataType.Double)
    T3 = fd.ops.mul(T0, S2)
    S4 = fd.define_scalar(0.00000, dtype=DataType.Double)
    T5 = fd.ops.eq(T1, S4)
    V10 = T0.shape()
    T11 = fd.ops.broadcast_in_dim(T5, shape=V10, broadcast_dims=[0, 1, 2, 3])
    S12 = fd.define_scalar(float("-inf"), dtype=DataType.Double)
    T13 = fd.ops.where(T11, S12, T3)
    T14 = fd.ops.max(T13, dims=[3], keepdim=False, dtype=DataType.Null)
    T20 = fd.ops.broadcast_in_dim(
        T14, shape=[T0.size(0), T0.size(1), T0.size(2), 1], broadcast_dims=[0, 1, 2]
    )
    T26 = fd.ops.broadcast_in_dim(T20, shape=V10, broadcast_dims=[0, 1, 2, 3])
    T27 = fd.ops.sub(T13, T26)
    T28 = fd.ops.exp(T27)
    T29 = fd.ops.sum(T28, dims=[3], keepdim=False, dtype=DataType.Null)
    T35 = fd.ops.broadcast_in_dim(
        T29, shape=[T0.size(0), T0.size(1), T0.size(2), 1], broadcast_dims=[0, 1, 2]
    )
    T41 = fd.ops.broadcast_in_dim(T35, shape=V10, broadcast_dims=[0, 1, 2, 3])
    T42 = fd.ops.reciprocal(T41)
    T43 = fd.ops.mul(T28, T42)
    S44 = fd.define_scalar(0.00000, dtype=DataType.Double)
    S45 = fd.define_scalar(1.00000, dtype=DataType.Double)
    T51 = fd.ops.uniform(S44, S45, shape=V10, dtype=DataType.Float)
    S52 = fd.define_scalar(1 - dropout_p, dtype=DataType.Double)
    T53 = fd.ops.lt(T51, S52)
    T55 = fd.ops.mul(T43, T53)
    S56 = fd.define_scalar(1 / (1 - dropout_p), dtype=DataType.Double)
    T57 = fd.ops.mul(T55, S56)

    if dtype in PROMOTE_DTYPES:
        T57 = fd.ops.cast(T57, dtype=dtype)
        T43 = fd.ops.cast(T43, dtype=dtype)

    fd.add_output(T57)
    fd.add_output(T53)
    fd.add_output(T43)
    fd.add_output(T11)

def nanogpt_attn_fwd_iobytes(size: tuple, dtype: torch.dtype):
    # Manual IOByte computation is required since nvFuser outputs (out, dropout_mask, attn, bias_mask) differ from baseline outputs (out).

    # Total IO bytes = in_tensor ([bs, nh, seq_len, seq_len], dtype) + bias ([seq_len, seq_len], float) +
    #   output ([bs, nh, seq_len, seq_len], dtype) + attn ([bs, nh, seq_len, seq_len], dtype) + dropout_mask ([bs, nh, seq_len, seq_len], bool) + bias_mask ([bs, nh, seq_len, seq_len], bool)

    bs, seq_len, nh, n_embd = size

    return int(
        bs * nh * seq_len * seq_len * (3 * dtype.itemsize + 2 * torch.bool.itemsize)
        + seq_len * seq_len * torch.float.itemsize
    )

@pytest.mark.parametrize("size", generate_attn_inputs())
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.inner_persistent
def test_nanogpt_attn_fwd_nvf_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    disable_validation: bool,
    disable_benchmarking: bool,
):
    batch_size, seq_len, nh, n_embd = size
    hs = n_embd // nh
    dropout_p = 0.2
    inputs = torch.randn(batch_size, nh, seq_len, seq_len, device="cuda", dtype=dtype)
    bias = torch.tril(torch.ones(seq_len, seq_len, device="cuda")).view(
        1, 1, seq_len, seq_len
    )

    bias_mask = bias[:, :, :seq_len, :seq_len] == 0
    bias_mask = bias_mask.broadcast_to(batch_size, nh, seq_len, seq_len)

    if not disable_validation:
        # For validating use a fusion definition with dropout_p=0.0
        with FusionDefinition() as val_fd:
            nanogpt_attn_fwd_fusion(
                val_fd, torch_dtype_to_nvfuser_dtype(dtype), hs, dropout_p=0.0
            )
        attn = inputs / (hs**0.5)
        dropout_mask = torch.ones(
            batch_size, nh, seq_len, seq_len, dtype=torch.bool, device="cuda"
        )
        attn = attn.masked_fill(bias[:, :, :seq_len, :seq_len] == 0, float("-inf"))
        attn = torch.nn.functional.softmax(attn, dim=-1)
        out = torch.nn.functional.dropout(attn, p=0.0)

        val_fd.validate([inputs, bias], [out, dropout_mask, attn, bias_mask])

    if not disable_benchmarking:
        with FusionDefinition() as fd:
            nanogpt_attn_fwd_fusion(
                fd, torch_dtype_to_nvfuser_dtype(dtype), hs, dropout_p
            )
        run_benchmark(benchmark, fd.execute, [inputs, bias])

@pytest.mark.parametrize("executor", DEFAULT_EXECUTORS)
@pytest.mark.parametrize("size", generate_attn_inputs())
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.inner_persistent
def test_nanogpt_attn_fwd_baseline_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    executor: str,
):
    if executor == "torchcompile":
        clear_dynamo_cache()
    batch_size, seq_len, nh, n_embd = size
    dropout_p = 0.2
    inputs = torch.randn(
        batch_size, nh, seq_len, seq_len, device="cuda", dtype=dtype, requires_grad=True
    )
    bias = torch.tril(torch.ones(seq_len, seq_len, device="cuda")).view(
        1, 1, seq_len, seq_len
    )

    benchmark_fn = with_executor(executor, nanogpt_attn)

    # Manually compute IOBytes: See PR #1725
    run_benchmark(
        benchmark,
        benchmark_fn,
        [inputs, bias, size, dropout_p],
        iobytes=nanogpt_attn_fwd_iobytes(size, dtype),
    )


# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
import pytest
from nvfuser_direct import FusionDefinition, DataType
from nvfuser_direct.pytorch_utils import torch_dtype_to_nvfuser_dtype
from .core import run_benchmark, clear_dynamo_cache, with_executor, DEFAULT_EXECUTORS
import torch
from .global_params import generate_attn_inputs, FLOAT_DTYPES, PROMOTE_DTYPES
from .torch_ops import nanogpt_attn

# Fusion from nanogpt attention module
# The nvFuser defintion only includes the non-matmul computation (masked_fill + softmax + dropout)
def nanogpt_attn_fwd_fusion(
    fd: FusionDefinition, dtype: DataType, head_size: int, dropout_p: float
):
    T0 = fd.define_tensor(
        shape=[-1, -1, -1, -1],
        contiguity=[True, True, True, True],
        dtype=dtype,
        is_cpu=False,
    )
    T1 = fd.define_tensor(
        shape=[1, 1, -1, -1],
        contiguity=[None, None, False, True],
        dtype=DataType.Float,
        is_cpu=False,
    )

    if dtype in PROMOTE_DTYPES:
        T0 = fd.ops.cast(T0, dtype=DataType.Float)

    S2 = fd.define_scalar(1 / head_size**0.5, dtype=DataType.Double)
    T3 = fd.ops.mul(T0, S2)
    S4 = fd.define_scalar(0.00000, dtype=DataType.Double)
    T5 = fd.ops.eq(T1, S4)
    V10 = T0.shape()
    T11 = fd.ops.broadcast_in_dim(T5, shape=V10, broadcast_dims=[0, 1, 2, 3])
    S12 = fd.define_scalar(float("-inf"), dtype=DataType.Double)
    T13 = fd.ops.where(T11, S12, T3)
    T14 = fd.ops.max(T13, dims=[3], keepdim=False, dtype=DataType.Null)
    T20 = fd.ops.broadcast_in_dim(
        T14, shape=[T0.size(0), T0.size(1), T0.size(2), 1], broadcast_dims=[0, 1, 2]
    )
    T26 = fd.ops.broadcast_in_dim(T20, shape=V10, broadcast_dims=[0, 1, 2, 3])
    T27 = fd.ops.sub(T13, T26)
    T28 = fd.ops.exp(T27)
    T29 = fd.ops.sum(T28, dims=[3], keepdim=False, dtype=DataType.Null)
    T35 = fd.ops.broadcast_in_dim(
        T29, shape=[T0.size(0), T0.size(1), T0.size(2), 1], broadcast_dims=[0, 1, 2]
    )
    T41 = fd.ops.broadcast_in_dim(T35, shape=V10, broadcast_dims=[0, 1, 2, 3])
    T42 = fd.ops.reciprocal(T41)
    T43 = fd.ops.mul(T28, T42)
    S44 = fd.define_scalar(0.00000, dtype=DataType.Double)
    S45 = fd.define_scalar(1.00000, dtype=DataType.Double)
    T51 = fd.ops.uniform(S44, S45, shape=V10, dtype=DataType.Float)
    S52 = fd.define_scalar(1 - dropout_p, dtype=DataType.Double)
    T53 = fd.ops.lt(T51, S52)
    T55 = fd.ops.mul(T43, T53)
    S56 = fd.define_scalar(1 / (1 - dropout_p), dtype=DataType.Double)
    T57 = fd.ops.mul(T55, S56)

    if dtype in PROMOTE_DTYPES:
        T57 = fd.ops.cast(T57, dtype=dtype)
        T43 = fd.ops.cast(T43, dtype=dtype)

    fd.add_output(T57)
    fd.add_output(T53)
    fd.add_output(T43)
    fd.add_output(T11)

def nanogpt_attn_fwd_iobytes(size: tuple, dtype: torch.dtype):
    # Manual IOByte computation is required since nvFuser outputs (out, dropout_mask, attn, bias_mask) differ from baseline outputs (out).

    # Total IO bytes = in_tensor ([bs, nh, seq_len, seq_len], dtype) + bias ([seq_len, seq_len], float) +
    #   output ([bs, nh, seq_len, seq_len], dtype) + attn ([bs, nh, seq_len, seq_len], dtype) + dropout_mask ([bs, nh, seq_len, seq_len], bool) + bias_mask ([bs, nh, seq_len, seq_len], bool)

    bs, seq_len, nh, n_embd = size

    return int(
        bs * nh * seq_len * seq_len * (3 * dtype.itemsize + 2 * torch.bool.itemsize)
        + seq_len * seq_len * torch.float.itemsize
    )

@pytest.mark.parametrize("size", generate_attn_inputs())
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.inner_persistent
def test_nanogpt_attn_fwd_nvf_benchmark(


    T43 = fd.ops.mul(T28, T42)
    S44 = fd.define_scalar(0.00000, dtype=DataType.Double)
    S45 = fd.define_scalar(1.00000, dtype=DataType.Double)
    T51 = fd.ops.uniform(S44, S45, shape=V10, dtype=DataType.Float)
    S52 = fd.define_scalar(1 - dropout_p, dtype=DataType.Double)
    T53 = fd.ops.lt(T51, S52)
    T55 = fd.ops.mul(T43, T53)
    S56 = fd.define_scalar(1 / (1 - dropout_p), dtype=DataType.Double)
    T57 = fd.ops.mul(T55, S56)

    if dtype in PROMOTE_DTYPES:
        T57 = fd.ops.cast(T57, dtype=dtype)
        T43 = fd.ops.cast(T43, dtype=dtype)

    fd.add_output(T57)
    fd.add_output(T53)
    fd.add_output(T43)
    fd.add_output(T11)

def nanogpt_attn_fwd_iobytes(size: tuple, dtype: torch.dtype):
    # Manual IOByte computation is required since nvFuser outputs (out, dropout_mask, attn, bias_mask) differ from baseline outputs (out).

    # Total IO bytes = in_tensor ([bs, nh, seq_len, seq_len], dtype) + bias ([seq_len, seq_len], float) +
    #   output ([bs, nh, seq_len, seq_len], dtype) + attn ([bs, nh, seq_len, seq_len], dtype) + dropout_mask ([bs, nh, seq_len, seq_len], bool) + bias_mask ([bs, nh, seq_len, seq_len], bool)

    bs, seq_len, nh, n_embd = size

    return int(
        bs * nh * seq_len * seq_len * (3 * dtype.itemsize + 2 * torch.bool.itemsize)
        + seq_len * seq_len * torch.float.itemsize
    )

@pytest.mark.parametrize("size", generate_attn_inputs())
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.inner_persistent
def test_nanogpt_attn_fwd_nvf_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    disable_validation: bool,
    disable_benchmarking: bool,
):
    batch_size, seq_len, nh, n_embd = size
    hs = n_embd // nh
    dropout_p = 0.2
    inputs = torch.randn(batch_size, nh, seq_len, seq_len, device="cuda", dtype=dtype)
    bias = torch.tril(torch.ones(seq_len, seq_len, device="cuda")).view(
        1, 1, seq_len, seq_len
    )


    fd.add_output(T11)

def nanogpt_attn_fwd_iobytes(size: tuple, dtype: torch.dtype):
    # Manual IOByte computation is required since nvFuser outputs (out, dropout_mask, attn, bias_mask) differ from baseline outputs (out).

    # Total IO bytes = in_tensor ([bs, nh, seq_len, seq_len], dtype) + bias ([seq_len, seq_len], float) +
    #   output ([bs, nh, seq_len, seq_len], dtype) + attn ([bs, nh, seq_len, seq_len], dtype) + dropout_mask ([bs, nh, seq_len, seq_len], bool) + bias_mask ([bs, nh, seq_len, seq_len], bool)

    bs, seq_len, nh, n_embd = size

    return int(
        bs * nh * seq_len * seq_len * (3 * dtype.itemsize + 2 * torch.bool.itemsize)
        + seq_len * seq_len * torch.float.itemsize
    )

@pytest.mark.parametrize("size", generate_attn_inputs())
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.inner_persistent
def test_nanogpt_attn_fwd_nvf_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    disable_validation: bool,
    disable_benchmarking: bool,
):
    batch_size, seq_len, nh, n_embd = size
    hs = n_embd // nh
    dropout_p = 0.2
    inputs = torch.randn(batch_size, nh, seq_len, seq_len, device="cuda", dtype=dtype)
    bias = torch.tril(torch.ones(seq_len, seq_len, device="cuda")).view(
        1, 1, seq_len, seq_len
    )

    bias_mask = bias[:, :, :seq_len, :seq_len] == 0
    bias_mask = bias_mask.broadcast_to(batch_size, nh, seq_len, seq_len)

    if not disable_validation:
        # For validating use a fusion definition with dropout_p=0.0
        with FusionDefinition() as val_fd:
            nanogpt_attn_fwd_fusion(
                val_fd, torch_dtype_to_nvfuser_dtype(dtype), hs, dropout_p=0.0
            )
        attn = inputs / (hs**0.5)
        dropout_mask = torch.ones(
            batch_size, nh, seq_len, seq_len, dtype=torch.bool, device="cuda"
        )
        attn = attn.masked_fill(bias[:, :, :seq_len, :seq_len] == 0, float("-inf"))
        attn = torch.nn.functional.softmax(attn, dim=-1)
        out = torch.nn.functional.dropout(attn, p=0.0)

        val_fd.validate([inputs, bias], [out, dropout_mask, attn, bias_mask])

    if not disable_benchmarking:
        with FusionDefinition() as fd:
            nanogpt_attn_fwd_fusion(
                fd, torch_dtype_to_nvfuser_dtype(dtype), hs, dropout_p
            )
        run_benchmark(benchmark, fd.execute, [inputs, bias])

@pytest.mark.parametrize("executor", DEFAULT_EXECUTORS)
@pytest.mark.parametrize("size", generate_attn_inputs())
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.inner_persistent
def test_nanogpt_attn_fwd_baseline_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    executor: str,
):
    if executor == "torchcompile":
        clear_dynamo_cache()
    batch_size, seq_len, nh, n_embd = size
    dropout_p = 0.2
    inputs = torch.randn(
        batch_size, nh, seq_len, seq_len, device="cuda", dtype=dtype, requires_grad=True
    )
    bias = torch.tril(torch.ones(seq_len, seq_len, device="cuda")).view(


            batch_size, nh, seq_len, seq_len, dtype=torch.bool, device="cuda"
        )
        attn = attn.masked_fill(bias[:, :, :seq_len, :seq_len] == 0, float("-inf"))
        attn = torch.nn.functional.softmax(attn, dim=-1)
        out = torch.nn.functional.dropout(attn, p=0.0)

        val_fd.validate([inputs, bias], [out, dropout_mask, attn, bias_mask])

    if not disable_benchmarking:
        with FusionDefinition() as fd:
            nanogpt_attn_fwd_fusion(
                fd, torch_dtype_to_nvfuser_dtype(dtype), hs, dropout_p
            )
        run_benchmark(benchmark, fd.execute, [inputs, bias])

@pytest.mark.parametrize("executor", DEFAULT_EXECUTORS)
@pytest.mark.parametrize("size", generate_attn_inputs())
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.inner_persistent
def test_nanogpt_attn_fwd_baseline_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    executor: str,
):
    if executor == "torchcompile":
        clear_dynamo_cache()
    batch_size, seq_len, nh, n_embd = size
    dropout_p = 0.2
    inputs = torch.randn(
        batch_size, nh, seq_len, seq_len, device="cuda", dtype=dtype, requires_grad=True
    )
    bias = torch.tril(torch.ones(seq_len, seq_len, device="cuda")).view(
        1, 1, seq_len, seq_len
    )

    benchmark_fn = with_executor(executor, nanogpt_attn)

    # Manually compute IOBytes: See PR #1725
    run_benchmark(
        benchmark,
        benchmark_fn,
        [inputs, bias, size, dropout_p],
        iobytes=nanogpt_attn_fwd_iobytes(size, dtype),
    )


# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
import pytest
from nvfuser_direct import FusionDefinition, DataType
from nvfuser_direct.pytorch_utils import torch_dtype_to_nvfuser_dtype
from .core import run_benchmark, clear_dynamo_cache, with_executor, DEFAULT_EXECUTORS
import torch
from .global_params import generate_attn_inputs, FLOAT_DTYPES, PROMOTE_DTYPES
from .torch_ops import huggingface_attn

# Fusion from huggingface attention implementation.
# The nvFuser defintion only includes the non-matmul computation (add + reshape + softmax + dropout)
def huggingface_attn_fwd_fusion(
    fd: FusionDefinition,
    dtype: DataType,
    dropout_p: float,
):
    T0 = fd.define_tensor(
        shape=[-1, -1, -1, -1],
        contiguity=[True, True, True, True],
        dtype=dtype,
        is_cpu=False,
    )
    T1 = fd.define_tensor(
        shape=[-1, -1, -1, -1],
        contiguity=[True, True, True, True],
        dtype=dtype,
        is_cpu=False,
    )

    if dtype in PROMOTE_DTYPES:
        T0 = fd.ops.cast(T0, dtype=DataType.Float)
        T1 = fd.ops.cast(T1, dtype=DataType.Float)

    T4 = fd.ops.add(T1, T0)

    T10 = fd.ops.reshape(
        T4, new_shape=[fd.ops.mul(T0.size(0), T0.size(1)), T0.size(2), T0.size(3)]
    )
    T12 = fd.ops.max(T10, dims=[2], keepdim=False, dtype=DataType.Null)

    T17 = fd.ops.broadcast_in_dim(
        T12,
        shape=[fd.ops.mul(T0.size(0), T0.size(1)), T0.size(2), 1],
        broadcast_dims=[0, 1],
    )
    T22 = fd.ops.broadcast_in_dim(
        T17,
        shape=[fd.ops.mul(T0.size(0), T0.size(1)), T0.size(2), T0.size(3)],
        broadcast_dims=[0, 1, 2],
    )
    T23 = fd.ops.sub(T10, T22)
    T24 = fd.ops.exp(T23)
    T25 = fd.ops.sum(T24, dims=[2], keepdim=False, dtype=DataType.Null)

    T30 = fd.ops.broadcast_in_dim(
        T25,
        shape=[fd.ops.mul(T0.size(0), T0.size(1)), T0.size(2), 1],
        broadcast_dims=[0, 1],
    )
    T35 = fd.ops.broadcast_in_dim(
        T30,
        shape=[fd.ops.mul(T0.size(0), T0.size(1)), T0.size(2), T0.size(3)],
        broadcast_dims=[0, 1, 2],
    )

    T36 = fd.ops.reciprocal(T35)
    T37 = fd.ops.mul(T24, T36)

    S39 = fd.define_scalar(0.00000, dtype=DataType.Double)
    S40 = fd.define_scalar(1.00000, dtype=DataType.Double)
    T45 = fd.ops.uniform(
        S39,
        S40,
        shape=[fd.ops.mul(T0.size(0), T0.size(1)), T0.size(2), T0.size(3)],
        dtype=DataType.Float,
    )
    S46 = fd.define_scalar(1 - dropout_p, dtype=DataType.Double)
    T47 = fd.ops.lt(T45, S46)

    T48 = fd.ops.cast(T47, dtype=DataType.Float)
    T49 = fd.ops.mul(T37, T48)
    S50 = fd.define_scalar(1 / (1 - dropout_p), dtype=DataType.Double)
    T51 = fd.ops.mul(T49, S50)

    if dtype in PROMOTE_DTYPES:
        T37 = fd.ops.cast(T37, dtype=dtype)
        T51 = fd.ops.cast(T51, dtype=dtype)

    fd.add_output(T51)
    fd.add_output(T37)
    fd.add_output(T47)

def huggingface_attn_fwd_iobytes(size: tuple, dtype: torch.dtype):
    # Manual IOByte computation is required since nvFuser outputs (output, attn, dropout_mask) differ from baseline outputs (output).

    # Total IO bytes = attention_mask ([bs, nh, seq_len, seq_len], dtype) + in_tensor ([bs, nh, seq_len, seq_len], dtype) +
    #   output ([bs, nh, seq_len, seq_len], dtype) + attn ([bs * nh, seq_len, seq_len], dtype) + dropout_mask ([bs * nh, seq_len, seq_len], bool)

    bs, seq_len, nh, n_embd = size
    return int(bs * nh * seq_len * seq_len * (dtype.itemsize * 4 + torch.bool.itemsize))

@pytest.mark.parametrize("size", generate_attn_inputs())
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.inner_persistent
def test_huggingface_attn_fwd_nvf_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    disable_validation: bool,
    disable_benchmarking: bool,
):
    batch_size, seq_len, nh, n_embd = size
    dropout_p = 0.2
    inputs = torch.randn(batch_size, nh, seq_len, seq_len, device="cuda", dtype=dtype)

    attention_mask = torch.zeros(
        batch_size, nh, seq_len, seq_len, device="cuda", dtype=dtype
    )

    if not disable_validation:
        # For validating use a fusion definition with dropout_p=0.0
        with FusionDefinition() as val_fd:
            huggingface_attn_fwd_fusion(
                val_fd, torch_dtype_to_nvfuser_dtype(dtype), dropout_p=0.0
            )

        dropout_mask = torch.ones(
            batch_size * nh, seq_len, seq_len, dtype=torch.bool, device="cuda"
        )
        attn = (inputs + attention_mask).view(batch_size * nh, seq_len, seq_len)
        attn = torch.nn.functional.softmax(attn, dim=-1)
        out = torch.nn.functional.dropout(attn, p=0.0)
        val_fd.validate([attention_mask, inputs], [out, attn, dropout_mask])

    if not disable_benchmarking:
        with FusionDefinition() as fd:
            huggingface_attn_fwd_fusion(
                fd, torch_dtype_to_nvfuser_dtype(dtype), dropout_p
            )
        run_benchmark(benchmark, fd.execute, [attention_mask, inputs])

@pytest.mark.parametrize("executor", DEFAULT_EXECUTORS)
@pytest.mark.parametrize("size", generate_attn_inputs())
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.inner_persistent
def test_huggingface_attn_fwd_baseline_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    executor: str,
):
    if executor == "torchcompile":
        clear_dynamo_cache()
    batch_size, seq_len, nh, n_embd = size
    dropout_p = 0.2
    inputs = torch.randn(
        batch_size, nh, seq_len, seq_len, device="cuda", dtype=dtype, requires_grad=True
    )
    attention_mask = torch.zeros(
        batch_size, nh, seq_len, seq_len, device="cuda", dtype=dtype
    )

    benchmark_fn = with_executor(executor, huggingface_attn)

    # Manually compute IOBytes: See PR #1725
    run_benchmark(
        benchmark,
        benchmark_fn,
        [inputs, attention_mask, size, dropout_p],
        iobytes=huggingface_attn_fwd_iobytes(size, dtype),
    )


# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
import pytest
from nvfuser_direct import FusionDefinition, DataType
from nvfuser_direct.pytorch_utils import torch_dtype_to_nvfuser_dtype
from .core import run_benchmark, clear_dynamo_cache, with_executor, DEFAULT_EXECUTORS
import torch
from .global_params import generate_attn_inputs, FLOAT_DTYPES, PROMOTE_DTYPES
from .torch_ops import huggingface_attn

# Fusion from huggingface attention implementation.
# The nvFuser defintion only includes the non-matmul computation (add + reshape + softmax + dropout)
def huggingface_attn_fwd_fusion(
    fd: FusionDefinition,
    dtype: DataType,
    dropout_p: float,
):
    T0 = fd.define_tensor(
        shape=[-1, -1, -1, -1],
        contiguity=[True, True, True, True],
        dtype=dtype,
        is_cpu=False,
    )
    T1 = fd.define_tensor(
        shape=[-1, -1, -1, -1],
        contiguity=[True, True, True, True],
        dtype=dtype,
        is_cpu=False,
    )

    if dtype in PROMOTE_DTYPES:
        T0 = fd.ops.cast(T0, dtype=DataType.Float)
        T1 = fd.ops.cast(T1, dtype=DataType.Float)

    T4 = fd.ops.add(T1, T0)

    T10 = fd.ops.reshape(
        T4, new_shape=[fd.ops.mul(T0.size(0), T0.size(1)), T0.size(2), T0.size(3)]
    )
    T12 = fd.ops.max(T10, dims=[2], keepdim=False, dtype=DataType.Null)

    T17 = fd.ops.broadcast_in_dim(
        T12,
        shape=[fd.ops.mul(T0.size(0), T0.size(1)), T0.size(2), 1],
        broadcast_dims=[0, 1],
    )
    T22 = fd.ops.broadcast_in_dim(
        T17,
        shape=[fd.ops.mul(T0.size(0), T0.size(1)), T0.size(2), T0.size(3)],
        broadcast_dims=[0, 1, 2],
    )
    T23 = fd.ops.sub(T10, T22)
    T24 = fd.ops.exp(T23)
    T25 = fd.ops.sum(T24, dims=[2], keepdim=False, dtype=DataType.Null)

    T30 = fd.ops.broadcast_in_dim(
        T25,
        shape=[fd.ops.mul(T0.size(0), T0.size(1)), T0.size(2), 1],
        broadcast_dims=[0, 1],
    )
    T35 = fd.ops.broadcast_in_dim(
        T30,
        shape=[fd.ops.mul(T0.size(0), T0.size(1)), T0.size(2), T0.size(3)],
        broadcast_dims=[0, 1, 2],
    )

    T36 = fd.ops.reciprocal(T35)
    T37 = fd.ops.mul(T24, T36)

    S39 = fd.define_scalar(0.00000, dtype=DataType.Double)
    S40 = fd.define_scalar(1.00000, dtype=DataType.Double)
    T45 = fd.ops.uniform(
        S39,
        S40,
        shape=[fd.ops.mul(T0.size(0), T0.size(1)), T0.size(2), T0.size(3)],
        dtype=DataType.Float,
    )
    S46 = fd.define_scalar(1 - dropout_p, dtype=DataType.Double)
    T47 = fd.ops.lt(T45, S46)

    T48 = fd.ops.cast(T47, dtype=DataType.Float)
    T49 = fd.ops.mul(T37, T48)
    S50 = fd.define_scalar(1 / (1 - dropout_p), dtype=DataType.Double)
    T51 = fd.ops.mul(T49, S50)

    if dtype in PROMOTE_DTYPES:
        T37 = fd.ops.cast(T37, dtype=dtype)
        T51 = fd.ops.cast(T51, dtype=dtype)

    fd.add_output(T51)
    fd.add_output(T37)
    fd.add_output(T47)

def huggingface_attn_fwd_iobytes(size: tuple, dtype: torch.dtype):
    # Manual IOByte computation is required since nvFuser outputs (output, attn, dropout_mask) differ from baseline outputs (output).

    # Total IO bytes = attention_mask ([bs, nh, seq_len, seq_len], dtype) + in_tensor ([bs, nh, seq_len, seq_len], dtype) +
    #   output ([bs, nh, seq_len, seq_len], dtype) + attn ([bs * nh, seq_len, seq_len], dtype) + dropout_mask ([bs * nh, seq_len, seq_len], bool)

    bs, seq_len, nh, n_embd = size
    return int(bs * nh * seq_len * seq_len * (dtype.itemsize * 4 + torch.bool.itemsize))

@pytest.mark.parametrize("size", generate_attn_inputs())
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.inner_persistent
def test_huggingface_attn_fwd_nvf_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    disable_validation: bool,


        shape=[fd.ops.mul(T0.size(0), T0.size(1)), T0.size(2), T0.size(3)],
        dtype=DataType.Float,
    )
    S46 = fd.define_scalar(1 - dropout_p, dtype=DataType.Double)
    T47 = fd.ops.lt(T45, S46)

    T48 = fd.ops.cast(T47, dtype=DataType.Float)
    T49 = fd.ops.mul(T37, T48)
    S50 = fd.define_scalar(1 / (1 - dropout_p), dtype=DataType.Double)
    T51 = fd.ops.mul(T49, S50)

    if dtype in PROMOTE_DTYPES:
        T37 = fd.ops.cast(T37, dtype=dtype)
        T51 = fd.ops.cast(T51, dtype=dtype)

    fd.add_output(T51)
    fd.add_output(T37)
    fd.add_output(T47)

def huggingface_attn_fwd_iobytes(size: tuple, dtype: torch.dtype):
    # Manual IOByte computation is required since nvFuser outputs (output, attn, dropout_mask) differ from baseline outputs (output).

    # Total IO bytes = attention_mask ([bs, nh, seq_len, seq_len], dtype) + in_tensor ([bs, nh, seq_len, seq_len], dtype) +
    #   output ([bs, nh, seq_len, seq_len], dtype) + attn ([bs * nh, seq_len, seq_len], dtype) + dropout_mask ([bs * nh, seq_len, seq_len], bool)

    bs, seq_len, nh, n_embd = size
    return int(bs * nh * seq_len * seq_len * (dtype.itemsize * 4 + torch.bool.itemsize))

@pytest.mark.parametrize("size", generate_attn_inputs())
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.inner_persistent
def test_huggingface_attn_fwd_nvf_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    disable_validation: bool,
    disable_benchmarking: bool,
):
    batch_size, seq_len, nh, n_embd = size
    dropout_p = 0.2
    inputs = torch.randn(batch_size, nh, seq_len, seq_len, device="cuda", dtype=dtype)

    attention_mask = torch.zeros(
        batch_size, nh, seq_len, seq_len, device="cuda", dtype=dtype
    )


        T51 = fd.ops.cast(T51, dtype=dtype)

    fd.add_output(T51)
    fd.add_output(T37)
    fd.add_output(T47)

def huggingface_attn_fwd_iobytes(size: tuple, dtype: torch.dtype):
    # Manual IOByte computation is required since nvFuser outputs (output, attn, dropout_mask) differ from baseline outputs (output).

    # Total IO bytes = attention_mask ([bs, nh, seq_len, seq_len], dtype) + in_tensor ([bs, nh, seq_len, seq_len], dtype) +
    #   output ([bs, nh, seq_len, seq_len], dtype) + attn ([bs * nh, seq_len, seq_len], dtype) + dropout_mask ([bs * nh, seq_len, seq_len], bool)

    bs, seq_len, nh, n_embd = size
    return int(bs * nh * seq_len * seq_len * (dtype.itemsize * 4 + torch.bool.itemsize))

@pytest.mark.parametrize("size", generate_attn_inputs())
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.inner_persistent
def test_huggingface_attn_fwd_nvf_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    disable_validation: bool,
    disable_benchmarking: bool,
):
    batch_size, seq_len, nh, n_embd = size
    dropout_p = 0.2
    inputs = torch.randn(batch_size, nh, seq_len, seq_len, device="cuda", dtype=dtype)

    attention_mask = torch.zeros(
        batch_size, nh, seq_len, seq_len, device="cuda", dtype=dtype
    )

    if not disable_validation:
        # For validating use a fusion definition with dropout_p=0.0
        with FusionDefinition() as val_fd:
            huggingface_attn_fwd_fusion(
                val_fd, torch_dtype_to_nvfuser_dtype(dtype), dropout_p=0.0
            )

        dropout_mask = torch.ones(
            batch_size * nh, seq_len, seq_len, dtype=torch.bool, device="cuda"
        )
        attn = (inputs + attention_mask).view(batch_size * nh, seq_len, seq_len)
        attn = torch.nn.functional.softmax(attn, dim=-1)
        out = torch.nn.functional.dropout(attn, p=0.0)
        val_fd.validate([attention_mask, inputs], [out, attn, dropout_mask])

    if not disable_benchmarking:
        with FusionDefinition() as fd:
            huggingface_attn_fwd_fusion(
                fd, torch_dtype_to_nvfuser_dtype(dtype), dropout_p
            )
        run_benchmark(benchmark, fd.execute, [attention_mask, inputs])

@pytest.mark.parametrize("executor", DEFAULT_EXECUTORS)
@pytest.mark.parametrize("size", generate_attn_inputs())
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.inner_persistent
def test_huggingface_attn_fwd_baseline_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    executor: str,
):
    if executor == "torchcompile":
        clear_dynamo_cache()
    batch_size, seq_len, nh, n_embd = size
    dropout_p = 0.2
    inputs = torch.randn(
        batch_size, nh, seq_len, seq_len, device="cuda", dtype=dtype, requires_grad=True
    )
    attention_mask = torch.zeros(


        dropout_mask = torch.ones(
            batch_size * nh, seq_len, seq_len, dtype=torch.bool, device="cuda"
        )
        attn = (inputs + attention_mask).view(batch_size * nh, seq_len, seq_len)
        attn = torch.nn.functional.softmax(attn, dim=-1)
        out = torch.nn.functional.dropout(attn, p=0.0)
        val_fd.validate([attention_mask, inputs], [out, attn, dropout_mask])

    if not disable_benchmarking:
        with FusionDefinition() as fd:
            huggingface_attn_fwd_fusion(
                fd, torch_dtype_to_nvfuser_dtype(dtype), dropout_p
            )
        run_benchmark(benchmark, fd.execute, [attention_mask, inputs])

@pytest.mark.parametrize("executor", DEFAULT_EXECUTORS)
@pytest.mark.parametrize("size", generate_attn_inputs())
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
@pytest.mark.inner_persistent
def test_huggingface_attn_fwd_baseline_benchmark(
    benchmark,
    size: tuple,
    dtype: torch.dtype,
    executor: str,
):
    if executor == "torchcompile":
        clear_dynamo_cache()
    batch_size, seq_len, nh, n_embd = size
    dropout_p = 0.2
    inputs = torch.randn(
        batch_size, nh, seq_len, seq_len, device="cuda", dtype=dtype, requires_grad=True
    )
    attention_mask = torch.zeros(
        batch_size, nh, seq_len, seq_len, device="cuda", dtype=dtype
    )

    benchmark_fn = with_executor(executor, huggingface_attn)

    # Manually compute IOBytes: See PR #1725
    run_benchmark(
        benchmark,
        benchmark_fn,
        [inputs, attention_mask, size, dropout_p],
        iobytes=huggingface_attn_fwd_iobytes(size, dtype),
    )