float8 inference weight-only quant should map to a fused kernel or explain why not

[vkuzo]

Today, Float8WeightOnlyConfig maps to a reference implementation of weight-only quant, which dequantized the tensor and then runs a high precision gemm:

ao/torchao/quantization/quantize_/workflows/float8/float8_tensor.py

Line 392 in ba3ac9f

return torch.nn.functional.linear(

Users have reported confusion about this, we should either clearly explain that no speedup is expected or map to a fast kernel.

[vipulSharma18]

Hi, I ran torchao's llama benchmark script with Float8WeightOnlyConfig on RTX 5090, and there's also abnormally high peak memory usage with this config, along with the initially reported slowdown:

20251108230606, tok/s=104.68, tok/s_decode=105.73, ttft=0.0186, mem/s=1571.23 GB/s, peak_mem=16.30 GB, model_size=15.01 GB quant: None, sparse: None, mod: Meta-Llama-3.1-8B, kv_quant: False, compile: True, compile_prefill: False, dtype: torch.bfloat16, device: cuda repro: python generate.py --checkpoint_path /root/checkpoints/unsloth/Meta-Llama-3.1-8B/model.pth --device cuda --precision torch.bfloat16 --compile --num_samples 5 --max_new_tokens 200 --batch_size 1 --top_k 200 --temperature 0.8 

20251108230802, tok/s=160.55, tok/s_decode=169.80, ttft=0.0675, mem/s=1204.97 GB/s, peak_mem= 9.21 GB, model_size= 7.51 GB quant: float8dq-tensor, sparse: None, mod: Meta-Llama-3.1-8B, kv_quant: False, compile: True, compile_prefill: False, dtype: torch.bfloat16, device: cuda repro: python generate.py --quantization float8dq-tensor --checkpoint_path /root/checkpoints/unsloth/Meta-Llama-3.1-8B/model.pth --device cuda --precision torch.bfloat16 --compile --num_samples 5 --max_new_tokens 200 --batch_size 1 --top_k 200 --temperature 0.8 

20251108231236, tok/s=  7.44, tok/s_decode=  7.48, ttft=0.1442, mem/s=  55.91 GB/s, peak_mem=26.00 GB, model_size= 7.51 GB quant: float8wo, sparse: None, mod: Meta-Llama-3.1-8B, kv_quant: False, compile: True, compile_prefill: False, dtype: torch.bfloat16, device: cuda repro: python generate.py --quantization float8wo --checkpoint_path /root/checkpoints/unsloth/Meta-Llama-3.1-8B/model.pth --device cuda --precision torch.bfloat16 --compile --num_samples 5 --max_new_tokens 200 --batch_size 1 --top_k 200 --temperature 0.8 

I'm trying to raise a PR now that I've been able to reproduce the behavior with torchao's benchmarking code instead of my own code.