iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Kernel][XPU] Tensor-descriptor operand loads for Triton W8A8 scaled_mm (#47205)

Browse Source 
Signed-off-by: oonyshch <xonyshch@gmail.com>
Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
This commit is contained in:
Lena Onyshchenko
2026-08-10 00:56:11 +00:00
committed by GitHub
co-authored by Kunshang Ji
parent d6941300fc
commit 751f2ccdd3
5 changed files with 122 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/kernels/quantization/test_triton_scaled_mm.py
+37
View File
@@ -123,3 +123,40 @@ def test_scaled_mm(
c_actual = torch_scaled_mm(a, b, scale_a, scale_b, out_dtype, bias)
torch.testing.assert_close(c_check, c_actual, rtol=1e-1, atol=1e-1)
# TD operand loads must be bit-exact vs the plain masked-load path.
@pytest.mark.skipif(
not (current_platform.is_cuda_alike() or current_platform.is_xpu()),
reason="Triton scaled_mm runs on CUDA-alike or XPU.",
)
@pytest.mark.parametrize(
"M,N,K", [(1, 4096, 4096), (64, 4096, 4096), (256, 2048, 4096)]
)
@pytest.mark.parametrize("in_dtype", get_8bit_types())
@pytest.mark.parametrize("use_scalar_scale_a", [True, False])
@pytest.mark.parametrize("use_bias", [True, False])
def test_scaled_mm_td_matches_plain(M, N, K, in_dtype, use_scalar_scale_a, use_bias):
dev = current_platform.device_type
out_dtype = torch.bfloat16
set_random_seed(0)
is_fp = torch.tensor([1, 1], dtype=in_dtype).is_floating_point()
if is_fp:
a = (0.25 * torch.rand((M, K), dtype=torch.float32, device=dev)).to(in_dtype)
b = (0.25 * torch.rand((K, N), dtype=torch.float32, device=dev)).to(in_dtype)
else:
a = torch.randint(-32, 32, (M, K), dtype=in_dtype, device=dev)
b = torch.randint(-32, 32, (K, N), dtype=in_dtype, device=dev)
scale_a = (
torch.rand((1, 1), device=dev)
if use_scalar_scale_a
else 0.25 * torch.rand((M, 1), device=dev)
)
scale_b = 0.25 * torch.rand((N, 1), device=dev)
bias = torch.rand((N,), device=dev, dtype=out_dtype) if use_bias else None
out_plain = triton_scaled_mm(a, b, scale_a, scale_b, out_dtype, bias, use_td=False)
out_td = triton_scaled_mm(a, b, scale_a, scale_b, out_dtype, bias, use_td=True)
torch.testing.assert_close(out_td, out_plain, rtol=0, atol=0)
vllm/_custom_ops.py
+17
View File
@@ -2003,6 +2003,23 @@ def scaled_int8_quant(
Returns:
tuple[torch.Tensor, torch.Tensor, torch.Tensor | None] : Output int8 tensor, scales, and optionally azp.
"""
if current_platform.is_xpu():
# XPU has no _C int8 quant op; use the torch.compile reference.
if not symmetric:
raise NotImplementedError(
"asymmetric int8 activation quantization is unsupported on XPU"
)
if scale is not None:
q = (input.to(torch.float32) / scale).round().clamp(-128, 127)
return q.to(torch.int8), scale, None
from vllm._xpu_ops import xpu_ops
q, scales, _ = xpu_ops.dynamic_per_token_int8_quant_ref(
input.contiguous(), True, 8
)
return q, scales.reshape(-1, 1).to(torch.float32), None
output = torch.empty_like(input, dtype=torch.int8)
if scale is not None:
# static-per-tensor quantization.
vllm/model_executor/kernels/linear/__init__.py
+1
View File
@@ -368,6 +368,7 @@ _POSSIBLE_INT8_KERNELS: dict[PlatformEnum, list[type[Int8ScaledMMLinearKernel]]]
HummingInt8ScaledMMLinearKernel,
],
PlatformEnum.ROCM: [AiterInt8ScaledMMLinearKernel, TritonInt8ScaledMMLinearKernel],
PlatformEnum.XPU: [TritonInt8ScaledMMLinearKernel],
}
# in priority/performance order (when available)
vllm/model_executor/kernels/linear/scaled_mm/triton.py
+2 -2
View File
@@ -29,9 +29,9 @@ class TritonInt8ScaledMMLinearKernel(CutlassInt8ScaledMMLinearKernel):
def is_supported(
cls, compute_capability: int | None = None
) -> tuple[bool, str | None]:
if current_platform.is_cuda_alike():
if current_platform.is_cuda_alike() or current_platform.is_xpu():
return True, None
return False, "requires ROCm or CUDA."
return False, "requires ROCm, CUDA or XPU."
@classmethod
def can_implement(cls, c: Int8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:
vllm/model_executor/layers/quantization/compressed_tensors/triton_scaled_mm.py
+65 -7
View File
@@ -4,7 +4,10 @@
import torch
from vllm.triton_utils import tl, triton
from vllm.triton_utils import tl, triton, use_tensor_descriptor
from vllm.triton_utils.allocation import set_triton_allocator
_TD_ALLOCATOR_DEVICES: set[torch.device] = set()
def is_weak_contiguous(x: torch.Tensor):
@@ -38,6 +41,8 @@ def scaled_mm_kernel(
BLOCK_SIZE_K: tl.constexpr,
BLOCK_SIZE_SCALE_A: tl.constexpr,
BLOCK_SIZE_SCALE_B: tl.constexpr,
USE_TD: tl.constexpr = False,
B_T: tl.constexpr = False,
):
pid = tl.program_id(axis=0)
@@ -85,13 +90,44 @@ def scaled_mm_kernel(
scale_a_ptrs = scale_a_ptr + offsets_scale_am
scale_b_ptrs = scale_b_ptr + offsets_scale_bn
for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
masks_k = offsets_k < K
masks_a = masks_am[:, None] & masks_k[None, :]
a = tl.load(a_ptrs, mask=masks_a)
if USE_TD:
a_desc = tl.make_tensor_descriptor(
a_ptr,
shape=[M, K],
strides=[stride_am, 1],
block_shape=[BLOCK_SIZE_M, BLOCK_SIZE_K],
)
if B_T:
# Checkpoint weights arrive as a transposed [N, K] view, so describe
# the underlying N-major buffer and transpose each tile.
b_desc = tl.make_tensor_descriptor(
b_ptr,
shape=[N, K],
strides=[stride_bn, 1],
block_shape=[BLOCK_SIZE_N, BLOCK_SIZE_K],
)
else:
b_desc = tl.make_tensor_descriptor(
b_ptr,
shape=[K, N],
strides=[stride_bk, 1],
block_shape=[BLOCK_SIZE_K, BLOCK_SIZE_N],
)
masks_b = masks_k[:, None] & masks_bn[None, :]
b = tl.load(b_ptrs, mask=masks_b)
for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
if USE_TD:
a = a_desc.load([pid_m * BLOCK_SIZE_M, k * BLOCK_SIZE_K])
if B_T:
b = tl.trans(b_desc.load([pid_n * BLOCK_SIZE_N, k * BLOCK_SIZE_K]))
else:
b = b_desc.load([k * BLOCK_SIZE_K, pid_n * BLOCK_SIZE_N])
else:
masks_k = offsets_k < K
masks_a = masks_am[:, None] & masks_k[None, :]
a = tl.load(a_ptrs, mask=masks_a)
masks_b = masks_k[:, None] & masks_bn[None, :]
b = tl.load(b_ptrs, mask=masks_b)
# Accumulate results.
accumulator = tl.dot(a, b, accumulator, out_dtype=accumulator_dtype)
@@ -149,6 +185,7 @@ def triton_scaled_mm(
block_size_n: int = 32,
block_size_k: int = 32,
use_heuristic=True,
use_td: bool | None = None,
) -> torch.Tensor:
M, K = input.shape
N = weight.shape[1]
@@ -195,6 +232,25 @@ def triton_scaled_mm(
accumulator_dtype = tl.float32 if input.is_floating_point() else tl.int32
# TD operand loads; gated on inner-dim contiguity and 16-byte tile alignment.
# Checkpoint weights are a transposed [N, K] view, so accept either
# orientation and let the kernel transpose per tile.
b_t = weight.stride(1) != 1 and weight.stride(0) == 1
b_inner = K if b_t else N
use_td = (
use_tensor_descriptor(use_td)
and input.stride(1) == 1
and (weight.stride(1) == 1 or b_t)
and (K * input.element_size()) % 16 == 0
and (b_inner * weight.element_size()) % 16 == 0
and (block_size_m & (block_size_m - 1)) == 0
and (block_size_n & (block_size_n - 1)) == 0
and (block_size_k & (block_size_k - 1)) == 0
)
if use_td and input.device not in _TD_ALLOCATOR_DEVICES:
set_triton_allocator(input.device)
_TD_ALLOCATOR_DEVICES.add(input.device)
# A = input, B = weight, C = result
# A = M x K, B = K x N, C = M x N
scaled_mm_kernel[grid](
@@ -219,6 +275,8 @@ def triton_scaled_mm(
BLOCK_SIZE_K=block_size_k,
BLOCK_SIZE_SCALE_A=block_size_sa,
BLOCK_SIZE_SCALE_B=block_size_sb,
USE_TD=use_td,
B_T=b_t,
)
return result.to(out_dtype)