diff --git a/benchmark/bench_jit_per_token_group_quant_8bit.py b/benchmark/bench_jit_per_token_group_quant_8bit.py index 9668b804..d1b2b171 100644 --- a/benchmark/bench_jit_per_token_group_quant_8bit.py +++ b/benchmark/bench_jit_per_token_group_quant_8bit.py @@ -63,10 +63,15 @@ def aot_per_token_group_quant_8bit(x, group_size, dst_dtype, eps=1e-10): x_q = torch.empty_like(x, dtype=dst_dtype) x_s = torch.empty((m, n // group_size), device=x.device, dtype=torch.float32) min_8bit, max_8bit = _minmax(dst_dtype) + # Pin v1; enable_v2=None would import sglang (not in bench container). if dst_dtype == torch.int8: - sgl_per_token_group_quant_int8(x, x_q, x_s, group_size, eps, min_8bit, max_8bit) + sgl_per_token_group_quant_int8( + x, x_q, x_s, group_size, eps, min_8bit, max_8bit, enable_v2=False + ) else: - sgl_per_token_group_quant_fp8(x, x_q, x_s, group_size, eps, min_8bit, max_8bit) + sgl_per_token_group_quant_fp8( + x, x_q, x_s, group_size, eps, min_8bit, max_8bit, enable_v2=False + ) return x_q, x_s