Add FLUX - AMD Quark quantitation example

[amd-xiaoyu12]

No description provided.

[gemini-code-assist]

Code Review

This pull request introduces a self-contained example and documentation for running FLUX.1 with AMD Quark native inference, supporting FP8 and MXFP4 dynamic quantization frozen into real low-precision kernels. The review feedback highlights critical improvements for multi-GPU and resource-constrained environments: performing quantization before enabling sequential CPU offload to preserve offload hooks, guarding the safetensors export to prevent race conditions across ranks, and explicitly setting the active CUDA device to avoid device mismatch errors.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

If enable_sequential_cpu_offload is enabled, the CPU offload hooks are registered on the original pipe.transformer submodules. Replacing pipe.transformer with the quantized/frozen model after enabling CPU offload will discard these hooks. Furthermore, moving the entire quantized model to GPU via quantized.to(device) inside quantize_transformer_native will bypass CPU offload entirely, leading to potential Out-Of-Memory (OOM) errors on memory-constrained GPUs.

Performing the quantization and freezing before enabling CPU offload ensures that the hooks are correctly registered on the final quantized/frozen modules and that the model is correctly offloaded to CPU.

Suggested change

	if args.enable_sequential_cpu_offload:
	pipe.enable_sequential_cpu_offload(gpu_id=local_rank)
	logging.info(f"rank {local_rank} sequential CPU offload enabled")
	else:
	pipe = pipe.to(f"cuda:{local_rank}")
	# Apply Quark native-inference quantization to the transformer in place,
	# after the pipeline is on device but before prepare_run / inference.
	device = torch.device(f"cuda:{local_rank}")
	pipe.transformer = quantize_transformer_native(
	transformer=pipe.transformer,
	quant_mode=args.quark_quantization_mode,
	device=device,
	export_dir=args.quark_export_dir,
	)
	# Apply Quark native-inference quantization to the transformer in place,
	# before the pipeline is moved to device or sequential CPU offload is enabled.
	# This ensures that CPU offload hooks are correctly registered on the final quantized/frozen modules.
	device = torch.device(f"cuda:{local_rank}")
	pipe.transformer = quantize_transformer_native(
	transformer=pipe.transformer,
	quant_mode=args.quark_quantization_mode,
	device=device,
	export_dir=args.quark_export_dir,
	)
	if args.enable_sequential_cpu_offload:
	pipe.enable_sequential_cpu_offload(gpu_id=local_rank)
	logging.info(f"rank {local_rank} sequential CPU offload enabled")
	else:
	pipe = pipe.to(f"cuda:{local_rank}")

[gemini-code-assist]

In a multi-GPU environment (e.g., when running with torchrun), all ranks will attempt to write the exported safetensors to the same --quark_export_dir simultaneously. This can cause race conditions, file corruption, or permission errors. Guarding the export with get_world_group().rank == 0 ensures only the master rank writes the exported model.

Suggested change

	if export_dir is not None:
	if export_dir is not None and get_world_group().rank == 0:

[gemini-code-assist]

In multi-GPU scripts, failing to explicitly set the active CUDA device via torch.cuda.set_device(local_rank) can cause un-indexed CUDA operations (such as torch.cuda.reset_peak_memory_stats() or torch.Generator(device="cuda")) to default to cuda:0, leading to incorrect memory tracking or device mismatch errors.

Suggested change

	local_rank = get_world_group().local_rank
	local_rank = get_world_group().local_rank
	torch.cuda.set_device(local_rank)

[pds-amd]

Thanks for this contribution.

I tried running the example on a gfx1201 machine with 32GB VRAM and it OOMs - it looks like there's a hard assumption that the whole model and quantized weights can be resident on the GPU. Sequential CPU offload also failed.

On RDNA4, --use_fp8_gemms already dispatches to aiter block-128 (xFuserFP8BlockScaleLinear); on MI3xx it uses torchao per-tensor. I'm assuming you're testing this on MI3xx?

Could you please provide some details in your PR description, e.g.

• frame running this against --use_fp8_gemms/--use_fp4_gemms, and give benchmarks (latency, memory, etc) vs that baseline.
• output images

Please also see the Gemini comments.

[avjves]

Thanks!

And to add to @pds-amd's comments, we have generally deprecated the examples in favor of using the model runner. (see here). Any chance you could add the support to the model runner instead of a one-off example?