[Kernel][XPU] Tensor-descriptor operand loads for Triton W8A8 scaled_mm (#47205)
Signed-off-by: oonyshch <xonyshch@gmail.com> Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
Lena Onyshchenko committed
Aug 10, 2026 at 02:56 UTC
751f2ccdd3b7ed2ba65dcb04dbd93187edc25b2d
5 files changed
+121
-8
tests/kernels/quantization/test_triton_scaled_mm.py
+37
@@ -123,3 +123,40 @@ def test_scaled_mm(
123
c_actual = torch_scaled_mm(a, b, scale_a, scale_b, out_dtype, bias)
124
125
torch.testing.assert_close(c_check, c_actual, rtol=1e-1, atol=1e-1)
126
+
127
+
128
+# TD operand loads must be bit-exact vs the plain masked-load path.
129
+@pytest.mark.skipif(
130
+ not (current_platform.is_cuda_alike() or current_platform.is_xpu()),
131
+ reason="Triton scaled_mm runs on CUDA-alike or XPU.",
132
+)
133
+@pytest.mark.parametrize(
134
+ "M,N,K", [(1, 4096, 4096), (64, 4096, 4096), (256, 2048, 4096)]
135
+)
136
+@pytest.mark.parametrize("in_dtype", get_8bit_types())
137
+@pytest.mark.parametrize("use_scalar_scale_a", [True, False])
138
+@pytest.mark.parametrize("use_bias", [True, False])
139
+def test_scaled_mm_td_matches_plain(M, N, K, in_dtype, use_scalar_scale_a, use_bias):
140
+ dev = current_platform.device_type
141
+ out_dtype = torch.bfloat16
142
+ set_random_seed(0)
143
+
144
+ is_fp = torch.tensor([1, 1], dtype=in_dtype).is_floating_point()
145
+ if is_fp:
146
+ a = (0.25 * torch.rand((M, K), dtype=torch.float32, device=dev)).to(in_dtype)
147
+ b = (0.25 * torch.rand((K, N), dtype=torch.float32, device=dev)).to(in_dtype)
148
+ else:
149
+ a = torch.randint(-32, 32, (M, K), dtype=in_dtype, device=dev)
150
+ b = torch.randint(-32, 32, (K, N), dtype=in_dtype, device=dev)
151
+
152
+ scale_a = (
153
+ torch.rand((1, 1), device=dev)
154
+ if use_scalar_scale_a
155
+ else 0.25 * torch.rand((M, 1), device=dev)
156
+ )
157
+ scale_b = 0.25 * torch.rand((N, 1), device=dev)
158
+ bias = torch.rand((N,), device=dev, dtype=out_dtype) if use_bias else None
159
+
160
+ out_plain = triton_scaled_mm(a, b, scale_a, scale_b, out_dtype, bias, use_td=False)
161
+ out_td = triton_scaled_mm(a, b, scale_a, scale_b, out_dtype, bias, use_td=True)
162
+ torch.testing.assert_close(out_td, out_plain, rtol=0, atol=0)
vllm/_custom_ops.py
+17
@@ -2003,6 +2003,23 @@ def scaled_int8_quant(
2003
Returns:
2004
tuple[torch.Tensor, torch.Tensor, torch.Tensor | None] : Output int8 tensor, scales, and optionally azp.
2005
"""
2006
+ if current_platform.is_xpu():
2007
+ # XPU has no _C int8 quant op; use the torch.compile reference.
2008
+ if not symmetric:
2009
+ raise NotImplementedError(
2010
+ "asymmetric int8 activation quantization is unsupported on XPU"
2011
+ )
2012
+ if scale is not None:
2013
+ q = (input.to(torch.float32) / scale).round().clamp(-128, 127)
2014
+ return q.to(torch.int8), scale, None
2015
+
2016
+ from vllm._xpu_ops import xpu_ops
2017
+
2018
+ q, scales, _ = xpu_ops.dynamic_per_token_int8_quant_ref(
2019
+ input.contiguous(), True, 8
2020
+ )
2021
+ return q, scales.reshape(-1, 1).to(torch.float32), None
2022
+
2023
output = torch.empty_like(input, dtype=torch.int8)
2024
if scale is not None:
2025
# static-per-tensor quantization.
vllm/model_executor/kernels/linear/__init__.py
+1
@@ -368,6 +368,7 @@ _POSSIBLE_INT8_KERNELS: dict[PlatformEnum, list[type[Int8ScaledMMLinearKernel]]]
368
HummingInt8ScaledMMLinearKernel,
369
],
370
PlatformEnum.ROCM: [AiterInt8ScaledMMLinearKernel, TritonInt8ScaledMMLinearKernel],
371
+ PlatformEnum.XPU: [TritonInt8ScaledMMLinearKernel],
372
}
373
374
# in priority/performance order (when available)
vllm/model_executor/kernels/linear/scaled_mm/triton.py
+2
-2
@@ -29,9 +29,9 @@ class TritonInt8ScaledMMLinearKernel(CutlassInt8ScaledMMLinearKernel):
29
def is_supported(
30
cls, compute_capability: int | None = None
31
) -> tuple[bool, str | None]:
32
- if current_platform.is_cuda_alike():
32
+ if current_platform.is_cuda_alike() or current_platform.is_xpu():
33
return True, None
34
- return False, "requires ROCm or CUDA."
34
+ return False, "requires ROCm, CUDA or XPU."
35
36
@classmethod
37
def can_implement(cls, c: Int8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:
vllm/model_executor/layers/quantization/compressed_tensors/triton_scaled_mm.py
+64
-6
@@ -4,7 +4,10 @@
4
5
import torch
6
7
-from vllm.triton_utils import tl, triton
7
+from vllm.triton_utils import tl, triton, use_tensor_descriptor
8
+from vllm.triton_utils.allocation import set_triton_allocator
9
+
10
+_TD_ALLOCATOR_DEVICES: set[torch.device] = set()
11
12
13
def is_weak_contiguous(x: torch.Tensor):
@@ -38,6 +41,8 @@ def scaled_mm_kernel(
41
BLOCK_SIZE_K: tl.constexpr,
42
BLOCK_SIZE_SCALE_A: tl.constexpr,
43
BLOCK_SIZE_SCALE_B: tl.constexpr,
44
+ USE_TD: tl.constexpr = False,
45
+ B_T: tl.constexpr = False,
46
):
47
pid = tl.program_id(axis=0)
48
@@ -85,13 +90,44 @@ def scaled_mm_kernel(
90
scale_a_ptrs = scale_a_ptr + offsets_scale_am
91
scale_b_ptrs = scale_b_ptr + offsets_scale_bn
92
93
+ if USE_TD:
94
+ a_desc = tl.make_tensor_descriptor(
95
+ a_ptr,
96
+ shape=[M, K],
97
+ strides=[stride_am, 1],
98
+ block_shape=[BLOCK_SIZE_M, BLOCK_SIZE_K],
99
+ )
100
+ if B_T:
101
+ # Checkpoint weights arrive as a transposed [N, K] view, so describe
102
+ # the underlying N-major buffer and transpose each tile.
103
+ b_desc = tl.make_tensor_descriptor(
104
+ b_ptr,
105
+ shape=[N, K],
106
+ strides=[stride_bn, 1],
107
+ block_shape=[BLOCK_SIZE_N, BLOCK_SIZE_K],
108
+ )
109
+ else:
110
+ b_desc = tl.make_tensor_descriptor(
111
+ b_ptr,
112
+ shape=[K, N],
113
+ strides=[stride_bk, 1],
114
+ block_shape=[BLOCK_SIZE_K, BLOCK_SIZE_N],
115
+ )
116
+
117
for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
89
- masks_k = offsets_k < K
90
- masks_a = masks_am[:, None] & masks_k[None, :]
91
- a = tl.load(a_ptrs, mask=masks_a)
118
+ if USE_TD:
119
+ a = a_desc.load([pid_m * BLOCK_SIZE_M, k * BLOCK_SIZE_K])
120
+ if B_T:
121
+ b = tl.trans(b_desc.load([pid_n * BLOCK_SIZE_N, k * BLOCK_SIZE_K]))
122
+ else:
123
+ b = b_desc.load([k * BLOCK_SIZE_K, pid_n * BLOCK_SIZE_N])
124
+ else:
125
+ masks_k = offsets_k < K
126
+ masks_a = masks_am[:, None] & masks_k[None, :]
127
+ a = tl.load(a_ptrs, mask=masks_a)
128
93
- masks_b = masks_k[:, None] & masks_bn[None, :]
94
- b = tl.load(b_ptrs, mask=masks_b)
129
+ masks_b = masks_k[:, None] & masks_bn[None, :]
130
+ b = tl.load(b_ptrs, mask=masks_b)
131
132
# Accumulate results.
133
accumulator = tl.dot(a, b, accumulator, out_dtype=accumulator_dtype)
@@ -149,6 +185,7 @@ def triton_scaled_mm(
185
block_size_n: int = 32,
186
block_size_k: int = 32,
187
use_heuristic=True,
188
+ use_td: bool | None = None,
189
) -> torch.Tensor:
190
M, K = input.shape
191
N = weight.shape[1]
@@ -195,6 +232,25 @@ def triton_scaled_mm(
232
233
accumulator_dtype = tl.float32 if input.is_floating_point() else tl.int32
234
235
+ # TD operand loads; gated on inner-dim contiguity and 16-byte tile alignment.
236
+ # Checkpoint weights are a transposed [N, K] view, so accept either
237
+ # orientation and let the kernel transpose per tile.
238
+ b_t = weight.stride(1) != 1 and weight.stride(0) == 1
239
+ b_inner = K if b_t else N
240
+ use_td = (
241
+ use_tensor_descriptor(use_td)
242
+ and input.stride(1) == 1
243
+ and (weight.stride(1) == 1 or b_t)
244
+ and (K * input.element_size()) % 16 == 0
245
+ and (b_inner * weight.element_size()) % 16 == 0
246
+ and (block_size_m & (block_size_m - 1)) == 0
247
+ and (block_size_n & (block_size_n - 1)) == 0
248
+ and (block_size_k & (block_size_k - 1)) == 0
249
+ )
250
+ if use_td and input.device not in _TD_ALLOCATOR_DEVICES:
251
+ set_triton_allocator(input.device)
252
+ _TD_ALLOCATOR_DEVICES.add(input.device)
253
+
254
# A = input, B = weight, C = result
255
# A = M x K, B = K x N, C = M x N
256
scaled_mm_kernel[grid](
@@ -219,6 +275,8 @@ def triton_scaled_mm(
275
BLOCK_SIZE_K=block_size_k,
276
BLOCK_SIZE_SCALE_A=block_size_sa,
277
BLOCK_SIZE_SCALE_B=block_size_sb,
278
+ USE_TD=use_td,
279
+ B_T=b_t,
280
)
281
282
return result.to(out_dtype)