Add subclass based method for inference w/ MXFP8

[drisspg]

Stacked PRs:

• ->Add subclass based method for inference w/ MXFP8 #2132

---

add subclass based method for inference

Perf comparisons

Micro

https://fburl.com/whh557d1

I am seeing extra overhead then expected

Macro

Runnng

python benchmarks/benchmark_serving.py \
 --backend vllm \
 --model "Qwen/Qwen2-7B-Instruct" \
 --endpoint /v1/completions \
 --dataset-name sharegpt \
 --dataset-path ShareGPT_V3_unfiltered_cleaned_split.json \
 --num-prompts 1024

 ============ Serving Benchmark Result ============
Successful requests:                     1024      
Benchmark duration (s):                  13.14     
Total input tokens:                      225502    
Total generated tokens:                  185804    
Request throughput (req/s):              77.91     
Output token throughput (tok/s):         14137.59  
Total Token throughput (tok/s):          31295.75  
---------------Time to First Token----------------
Mean TTFT (ms):                          2659.95   
Median TTFT (ms):                        2613.38   
P99 TTFT (ms):                           4457.53   
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          51.95     
Median TPOT (ms):                        28.56     
P99 TPOT (ms):                           167.73    
---------------Inter-token Latency----------------
Mean ITL (ms):                           23.52     
Median ITL (ms):                         15.42     
P99 ITL (ms):                            148.78    
==================================================

Running against MXFP8:

python benchmarks/benchmark_serving.py \  
  --backend vllm \
  --model "/home/drisspg/meta/scripts/data/mxfp8-Qwen2-7B-Instruct" \
  --endpoint /v1/completions \
  --dataset-name sharegpt \
  --dataset-path ShareGPT_V3_unfiltered_cleaned_split.json \
  --num-prompts 1024

============ Serving Benchmark Result ============
Successful requests:                     1024      
Benchmark duration (s):                  13.43     
Total input tokens:                      225502    
Total generated tokens:                  185297    
Request throughput (req/s):              76.26     
Output token throughput (tok/s):         13800.30  
Total Token throughput (tok/s):          30594.93  
---------------Time to First Token----------------
Mean TTFT (ms):                          1119.68   
Median TTFT (ms):                        1100.86   
P99 TTFT (ms):                           1721.80   
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          29.11     
Median TPOT (ms):                        27.07     
P99 TPOT (ms):                           46.91     
---------------Inter-token Latency----------------
Mean ITL (ms):                           23.38     
Median ITL (ms):                         17.28     
P99 ITL (ms):                            49.45     
==================================================

Trace NonQuant:
https://fburl.com/sput3bmn

Trace Quant:
https://fburl.com/0pgmyrge

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

• CI workflows being skipped on PR

❌ 1 New Failure

As of commit 44a878b with merge base 955ebb0 ():

NEW FAILURE - The following job has failed:

• Run TorchAO Experimental Tests / test-mps-ops (macos-m1-stable) (gh)
  Process completed with exit code 1.

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

[jerryzh168]

why is swizzle top level?

[drisspg]

Thats a really good question

[jerryzh168]

nit: add to safe global for serialization?

ao/torchao/quantization/quant_api.py

Line 1993 in 58502e3

torch.serialization.add_safe_globals(

[drisspg]

Sorry, this should be a draft still, I am just dumping all my commits for now

@jerryzh168