Added benchmark for LLaMA 3 model for attention tests

[howardzhang-cv]

Stack from ghstack (oldest at bottom):

• Add FA4 fp8 backend to low precision attention api #3947
• Add FA4 monkey-patch path for low-precision attention #3960
• -> Added benchmark for LLaMA 3 model for attention tests #3930
• Added benchmark for single attention layer across different sequence lengths #3929
• Added benchmarking for new torchao low precision attention api #3865
• Added new API for low precision fp8 attention using FA3 #3857
• Add FP8 FA3 low-precision attention with monkey-patch SDPA path #3959

Summary

• Added new benchmark for new low precision attention API
• uses Llama3-8b model, prefill
• 2-phase test
  • perplexity using the WikiText-2 test set
  • runtime over different sequence lengths (1024 -> 131072)
• Can set baseline and test models between different backends: (fa2, fa3, fa3_fp8, fa4, fa4_fp8)
• has options to control torch.compile usage, warmup_iters, sequence length, number of runtime test iterations, rope fusion

Example Run

python benchmarks/prototype/attention/eval_llama3_model.py --baseline fa3 --test fa3_fp8

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/3930

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit b6f072f with merge base 42bcdc4 ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[namgyu-youn]

We don't have to define wiki-text tokenizer here. Instead, we can run lm-eval directly like

ao/benchmarks/quantization/calibration_based/create_quantized_model.py

Lines 30 to 37 in 450699a

	# Calibrate
	evaluator.simple_evaluate(
	HFLM(pretrained=model, tokenizer=tokenizer),
	tasks=tasks,
	limit=limit,
	batch_size=1,
	)
	quantize_(model, config_class(base_config, step="convert"), filter_fn=filter_fn)

[howardzhang-cv]

Thanks! I changed it to use this instead

[namgyu-youn]

This is also not needed. See above comment

[howardzhang-cv]

Changed

[namgyu-youn]

Q. Can we compute forward pass latency using vLLM directly, similar to e2e?

ao/benchmarks/quantization/measure_accuracy_and_performance.sh

Line 176 in 3d3dd50

vllm bench throughput --model $OUTPUT_DIR --dtype bfloat16 $PREFILL_ARGS 2>&1 | tee -a "$LOG_FILE"

[howardzhang-cv]

Unfortunately not. Unlike the other quantization APIs in TorchAO, the low precision attention path requires replacing F.scaled_dot_product_attention with a specific attention backend capable of low precision attention (e.g. FA3/4). So we need to ensure that our model calls F.SDPA.

[namgyu-youn]

What does DEFAULT_MODEL_ID do? Should it be called by args (--model_id) default directly?

[howardzhang-cv]

Fixed

[namgyu-youn]

LGTM, thanks for addressing all the comments!