A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 10, 09:00 UTC
[Bugfix][Quantization] Ensure input contiguity in per_token_quant_int8 (#31637)
Signed-off-by: vensen <vensenmu@gmail.com>
Vensen committed
Jan 11, 2026 at 04:40 UTC
6ea001cfb7549265e9053569b4b17126ebcfc2d0
1 file changed
+8
-5
vllm/model_executor/layers/quantization/utils/int8_utils.py
+8
-5
@@ -122,15 +122,17 @@ def _per_token_quant_int8(
122
123
124
def per_token_quant_int8(x):
125
+ original_shape = x.shape
126
+ if x.dim() > 2:
127
+ x = x.view(-1, original_shape[-1])
128
M = x.numel() // x.shape[-1]
129
N = x.shape[-1]
127
- x_q = torch.empty_like(x, device=x.device, dtype=torch.int8)
128
- scales = torch.empty(x.shape[:-1] + (1,), device=x.device, dtype=torch.float32)
130
+ x_q = torch.empty((M, N), device=x.device, dtype=torch.int8)
131
+ scales = torch.empty((M, 1), device=x.device, dtype=torch.float32)
132
BLOCK = triton.next_power_of_2(N)
133
# heuristics for number of warps
134
num_warps = min(max(BLOCK // 256, 1), 8)
132
-
133
- assert x.is_contiguous()
135
+ x = x.contiguous()
136
_per_token_quant_int8[(M,)](
137
x,
138
x_q,
@@ -142,7 +144,8 @@ def per_token_quant_int8(x):
144
num_warps=num_warps,
145
num_stages=1,
146
)
145
-
147
+ x_q = x_q.view(*original_shape)
148
+ scales = scales.view(*original_shape[:-1], 1)
149
return x_q, scales
150
151