[xpu] Route per-block FP8 GEMM through torch._scaled_mm on XPU#4548
[xpu] Route per-block FP8 GEMM through torch._scaled_mm on XPU#4548hoshibara wants to merge 6 commits into
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/4548
Note: Links to docs will display an error until the docs builds have been completed. This comment was automatically generated by Dr. CI and updates every 15 minutes. |
meta-cla
Bot
added
the
CLA Signed
On Intel BMG XPU, the original blockwise_fp8_gemm triton kernel was an order of magnitude slower than per-row scaled_mm. XPU's oneDNN _scaled_mm natively supports mixed-mode (act=1x128, weight=128x128) where b is interpreted as [K, N] and scale_b as [K/128, N/128]. For the XPU path only, call addmm_float8_unwrapped_inference with: - a = inpt_data [M, K] row-major - scale_a = input_tensor.scale.squeeze(-1) [M, K/128] - b = weight_tensor.qdata [K, N] column-major view - scale_b = weight_tensor.scale [K/128, N/128] The existing blockwise_fp8_gemm triton path is preserved as the default for non-XPU devices. Verified on BMG Arc B580: per-block eager GEMM now 0.61x of per-row at batch=1 and 1.26x at batch=1024; shape-correctness max_err < 0.08 across 6 configs. Co-Authored-Extended: Claude <noreply@anthropic.com>
5177cb9 to
fe6c37f
Compare
The XPU path previously passed raw input_tensor.scale/weight_tensor.qdata directly, bypassing preprocess_scale and preprocess_data. This meant pad_inner_dim and contiguity fixes were skipped, which is only safe when K is 128-aligned. Change to reuse the same preprocessed inpt_data/input_scale/w_data/w_scale that the CUDA triton path and per-row path already consume. The only intentional difference is b passed without .t(): oneDNN expects column-major weights (already arranged by preprocess_data), while blockwise_fp8_gemm expects row-major.
…wise path Use only inpt_data/input_scale/w_data/w_scale in the XPU _is_128_128_scaled path. Replace input_tensor.scale.shape[-1] -> input_scale.shape[-1] and input_tensor.dtype -> input_scale.dtype.
xiaowangintel
added
module: inference
|
|
|
@claude pls follow this skills #3823 (comment) to review this pr. |
|
Claude finished @xiaowangintel's task in 3m 19s —— View job PR Review — Route per-block FP8 GEMM through
|
Two related fixes for the XPU blockwise path (PerBlock[1,128] act, PerBlock[128,128] wt) routed through torch._scaled_mm: 1. input_scale reshape: preprocess_scale flattens the 3D scale (1, seq, K/128) into (M*K/128, 1), and .reshape(-1, scale.shape[-1]) was a no-op because shape[-1]==1. Use inpt_data.shape[0] to reconstruct the correct (M, K/128) shape. 2. output_dtype: the refactored code passed input_scale.dtype (float32 after preprocess_scale) to _scaled_mm, producing float32 output. Restore input_tensor.dtype (bfloat16) to match the non-blockwise path and the original fe6c37f behavior.
aa3b093 to
668816b
Compare
hoshibara
changed the title
2 participants
On Intel BMG XPU, the original blockwise_fp8_gemm triton kernel was an order of magnitude slower than per-row scaled_mm. XPU's oneDNN _scaled_mm natively supports mixed-mode (act=1x128, weight=128x128) where b is interpreted as [K, N] and scale_b as [K/128, N/128].
For the XPU path only, call addmm_float8_unwrapped_inference with:
The existing blockwise_fp8_gemm triton path is preserved as the default for non-XPU devices.
Verified on BMG Arc B580: per-block eager GEMM now 0.61x of per-row at batch=1 and 1.26x at batch=1024; shape-correctness max_err < 0.08 across 6 configs.