From 751f2ccdd3b7ed2ba65dcb04dbd93187edc25b2d Mon Sep 17 00:00:00 2001 From: Lena Onyshchenko Date: Mon, 10 Aug 2026 02:56:11 +0200 Subject: [PATCH] [Kernel][XPU] Tensor-descriptor operand loads for Triton W8A8 scaled_mm (#47205) Signed-off-by: oonyshch Co-authored-by: Kunshang Ji --- .../quantization/test_triton_scaled_mm.py | 37 ++++++++++ vllm/_custom_ops.py | 17 +++++ .../model_executor/kernels/linear/__init__.py | 1 + .../kernels/linear/scaled_mm/triton.py | 4 +- .../compressed_tensors/triton_scaled_mm.py | 72 +++++++++++++++++-- 5 files changed, 122 insertions(+), 9 deletions(-) diff --git a/tests/kernels/quantization/test_triton_scaled_mm.py b/tests/kernels/quantization/test_triton_scaled_mm.py index e0f585c76a..e829c63a84 100644 --- a/tests/kernels/quantization/test_triton_scaled_mm.py +++ b/tests/kernels/quantization/test_triton_scaled_mm.py @@ -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) diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index 8c1cc105e2..c4ff05c75d 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -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. diff --git a/vllm/model_executor/kernels/linear/__init__.py b/vllm/model_executor/kernels/linear/__init__.py index 4fe7291bd3..90e983653e 100644 --- a/vllm/model_executor/kernels/linear/__init__.py +++ b/vllm/model_executor/kernels/linear/__init__.py @@ -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) diff --git a/vllm/model_executor/kernels/linear/scaled_mm/triton.py b/vllm/model_executor/kernels/linear/scaled_mm/triton.py index 78dad87295..d6169d2cfc 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/triton.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/triton.py @@ -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]: diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/triton_scaled_mm.py b/vllm/model_executor/layers/quantization/compressed_tensors/triton_scaled_mm.py index 25c7d335da..e0d1fdf271 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/triton_scaled_mm.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/triton_scaled_mm.py @@ -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)