[XPU] Enable XPU blockfp8 for DSv3 - #49596
Conversation
Signed-off-by: frost-intel <frost.mitchell@intel.com>
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
|
@jikunshang could you please review and label for ci? |
|
@frost-intel can you verify whether this PR resolved your issue? |
|
This pull request has merge conflicts that must be resolved before it can be |
Signed-off-by: Frost Mitchell <frost.mitchell@intel.com>
Signed-off-by: frost-intel <frost.mitchell@intel.com>
|
@jikunshang That PR fixed one of the 3 issues described above (
I've also added a test to demonstrate the fix. However, I don't see many XPU-specific tests in the UTs, so let me know if you'd rather I remove that. |
|
@frost-intel for accuracy, do you mean the micro test shows oneDNN's result accuracy comparing to the ref result? Are you able to run full model? |
|
@xwu-intel The regular DeepSeek-V3 input for But in debugging I found that oneDNN does allow this op with some other input sizes where The problem is that since Additionally, trying to move padding lower in the stack instead of in vLLM would require padding to be done at runtime wasting additional memory. Here's a reproducer: repro_onednn_block_n.py
"""Minimal repro: oneDNN's fp8 block GEMM applies the wrong block scale along N
when N is not a multiple of block_n, silently.
Everything is picked so the right answer is obvious by inspection:
* activation = all ones
* weight block b = the constant 2**b, held flat across all K
* K = 512 = 4 * 128, so K is not the problem
Because each 128-row weight block is a single constant, block quantization
writes the *same* fp8 byte (448 = fp8 max) into every element of the weight --
all of the information ends up in the scale. So
out[:, n] == K * 2**(n // 128)
and every column in a block must show one flat value. Anything else is the
scale grid landing on the wrong rows.
What the output shows: oneDNN sizes the scale group as N / n_blocks rather than
using block_n, so at N=448 it switches scales every 112 rows instead of every
128. Padding N up to n_blocks * block_n makes the two agree.
python repro_onednn_block_n.py
"""
import torch
import vllm_xpu_kernels._xpu_C # noqa: F401 registers torch.ops._xpu_C
FP8 = torch.float8_e4m3fn
FMAX = torch.finfo(FP8).max # 448
BLOCK = 128
M, K, N = 16, 512, 448 # N = 3*128 + 64, deliberately not a multiple of BLOCK
#M, K, N = 16, 512, 2112 # N = 3*128 + 64, deliberately not a multiple of BLOCK
n_blk, k_blk = -(-N // BLOCK), K // BLOCK # 4 scale blocks along N
n_pad = n_blk * BLOCK # 512 rows of weight the scale grid describes
# Activation: all ones. amax = 1, so the scale is 1/FMAX and every byte is FMAX.
xq = torch.full((M, K), FMAX).to(FP8).xpu()
xs = torch.full((M, k_blk), 1.0 / FMAX).xpu()
# Weight: block b is the constant 2**b. Same story -- every byte is FMAX and
# the constant lives entirely in that block's scale.
vals = torch.tensor([2.0**b for b in range(n_blk)])
wq_pad = torch.full((n_pad, K), FMAX).to(FP8).xpu() # padded to the scale grid
wq = wq_pad[:N].contiguous() # what a real checkpoint hands you
ws = (vals / FMAX).view(n_blk, 1).repeat(1, k_blk).xpu() # [n_blocks, k_blocks]
def gemm(weight: torch.Tensor) -> torch.Tensor:
return torch.ops._xpu_C.fp8_gemm(
xq,
weight.t().contiguous(),
torch.bfloat16,
xs,
ws.t().contiguous(),
torch.Tensor(),
)[..., :N].float()
def show(t: torch.Tensor) -> str:
lo, hi = t.min().item(), t.max().item()
return f"{lo:.0f}" if lo == hi else f"{lo:.0f}..{hi:.0f}"
# out[:, n] / K identifies which block's scale oneDNN actually applied.
LUT = {K * v: i for i, v in enumerate(vals.tolist())}
def scale_blocks(out: torch.Tensor) -> list[int]:
return [LUT.get(out[0, n].item(), -1) for n in range(N)]
def runs(seq: list[int]):
lo = 0
for n in range(1, N + 1):
if n == N or (seq[n], n // BLOCK) != (seq[lo], lo // BLOCK):
yield lo, n, seq[lo]
lo = n
print(f"M={M} K={K} N={N} block_n={BLOCK} N % block_n = {N % BLOCK}")
print(f"scale has {n_blk} blocks along N -> oneDNN reads {n_pad} weight rows\n")
bad, good = gemm(wq), gemm(wq_pad)
print(f"{'cols':>12}{'want':>8}{'unpadded':>12}{'padded':>10} diagnosis")
for b in range(n_blk):
lo, hi = b * BLOCK, min((b + 1) * BLOCK, N)
want = K * vals[b].item()
tail = " <- ragged block" if hi - lo != BLOCK else ""
got = show(bad[:, lo:hi])
ok = "ok" if got == f"{want:.0f}" else "not even constant within the block"
print(f" [{lo:4d}:{hi:4d}]{want:8.0f}{got:>12}{show(good[:, lo:hi]):>10} {ok}{tail}")
print("\nWhich scale block oneDNN actually applied, per column (unpadded):\n")
for lo, hi, got in runs(scale_blocks(bad)):
want = lo // BLOCK
print(
f" cols [{lo:4d}:{hi:4d}] used scale block {got}, wanted {want}"
f" {'ok' if got == want else 'WRONG'}"
)
assert scale_blocks(good) == [n // BLOCK for n in range(N)]
print(f"\n (padded: all {N} columns get the right scale block)")
print(
f"\nThe switch points are multiples of {N // n_blk} = N / n_blocks, not of"
f" block_n = {BLOCK}.\nThe call does not raise -- it just returns wrong values."
) |
|
In vllm-xpu-kernels, the block scale is computed as bool is_block_quant = (m1_sc.dim() == 2) && (m1_sc.size(1) > 1);
int64_t wei_group_k = -1;
int64_t wei_group_n = -1;
if (is_block_quant) {
TORCH_CHECK(
m1_sc.size(1) == m2_sc.size(0),
"Mismatch group size in input and weight.",
m1_sc.size(1),
" vs ",
m2_sc.size(0));
wei_group_k = k / m2_sc.size(0);
wei_group_n = n / m2_sc.size(1);
}As long as Currently, there's no mechanism to fix this in vllm-xpu-kernels, except for if we added However, this is fixed by this PR, where we add padding to weights to ensure they are evenly divisible by the block size. |
|
@zufangzhu help to check this. looks good to me. Besides this PR, maybe add some checks in TORCH_CHECK to ensure it's divisible and easy to detect incorrect block. |
|
3 participants
Purpose
Enable DeepSeek-V3 on Intel XPU through the oneDNN block-scaled FP8 GEMM (XPUFp8BlockScaledMMKernel). Three issues blocked correct execution:
oneDNN requires N to be a multiple of block_n. Some DeepSeek-V3 block-FP8 projections have a per-partition N that isn't 128-aligned. We now pad the weight's N up to the next multiple of block_n in
process_weights_after_loading, record the unpadded size, and drop the padded output columns after the GEMM inapply_block_scaled_mm.Block-scale layout mismatch during weight dequant. The XPU kernel transposes the block scale to
[K/block_k, N/block_n]layout at load time.get_and_maybe_dequant_weights(used by MLA to recover kv_b_proj for weight absorption) assumes the checkpoint[N/block_n, K/block_k]layout and asserts on the mismatch. We recordlayer.weight_scale_transposed = Truewhen the kernel repacks the scale and undo the transpose in the dequant helper before scaled_dequantize.Make the shared block base padding-agnostic.
Fp8BlockScaledMMLinearKernel.apply_weightspreviously computed the output shape fromweight.shape[0], which is the padded N on XPU. It now derives the output shape from the actual GEMM output width (output.shape[-1]) after the matmul. This is a no-op for all non-padding backends (Triton / DeepGEMM / Cutlass / FlashInfer / Aiter-ROCm / CPU), where the returned width already equalsweight.shape[0], and it lets the XPU kernel return a sliced result that reshapes correctly.Test Plan
Run a block-FP8 DeepSeek-V3 checkpoint on Intel XPU Max 1550 and confirm it (a) loads without the oneDNN N-divisibility failure and without the scaled_dequantize shape assertion in the MLA process_weights_after_loading, and (b) produces coherent generations.
This did require using two workarounds which were outside the scope of this PR
torch.ops._C_cache_ops.getMemoryInfo(device)withtorch.xpu.get_mem_info(device)due to outdated L0 drivers on Max 1550 installation.Test Result
Before: DeepSeek-V3 block-FP8 failed to run on XPU.
After: The model loads and generates on XPU (TP×PP×EP, --enforce-eager); the padded kv_a_proj layers run through oneDNN and the MLA kv_b_proj dequant succeeds.
A partial run of gsm8k (limit: 40) shows high accuracy: