From 82955b3a225c6c0a7d6f736cdbaeb09da30e67c4 Mon Sep 17 00:00:00 2001 From: tzhouam Date: Thu, 10 Sep 2026 02:54:53 +0800 Subject: [PATCH] [Edge][CPU] Default the code predictor's CPU fast path to fp32 + compiled static step; int8 opt-in In-situ parity against the bf16 model (VLLM_OMNI_PREDICTOR_PARITY_CHECK=1): fp32 control keeps 93.8 % top-1 / 95 % top-5 / KL 0.0015; torch.ao dynamic int8 falls to 22 % top-1 (max logit diff 7), so per-tensor activation quantization is not acceptable as a default. Adds the ao-nohead variant (int8 transformer body, fp32 lm heads) for measurement. Co-Authored-By: Claude Fable 5.1 --- .../models/common/qwen3_code_predictor.py | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/vllm_omni/model_executor/models/common/qwen3_code_predictor.py b/vllm_omni/model_executor/models/common/qwen3_code_predictor.py index 48f968dec..778759997 100644 --- a/vllm_omni/model_executor/models/common/qwen3_code_predictor.py +++ b/vllm_omni/model_executor/models/common/qwen3_code_predictor.py @@ -1173,7 +1173,11 @@ class CodePredictorWrapper(nn.Module): self._cpu_fast_path_done = True if device.type != "cpu" or current_omni_platform.is_npu(): return - mode = os.environ.get(self._INT8_ENV, "ao").strip().lower() # "ao": torch.ao dynamic int8 (validated) + # Default "fp32": fp32 weights + compiled shape-static step. Measured in situ against the + # bf16 model (VLLM_OMNI_PREDICTOR_PARITY_CHECK=1): fp32 keeps 93.8 % top-1 / KL 0.0015, + # while torch.ao dynamic int8 ("ao") drops to 22 % top-1 (per-tensor activation + # quantization of the residual stream), so int8 is opt-in only. + mode = os.environ.get(self._INT8_ENV, "fp32").strip().lower() if mode in ("0", "false", "off", "no"): return self._parity_ref = None @@ -1214,8 +1218,12 @@ class CodePredictorWrapper(nn.Module): from torch.ao.quantization import quantize_dynamic self.model = quantize_dynamic(self.model.float(), {nn.Linear}, dtype=torch.qint8) - self.lm_head = quantize_dynamic(self.lm_head.float(), {nn.Linear}, dtype=torch.qint8) - kind = "torch.ao dynamic int8 Linear (eager)" + if mode == "ao-nohead": # keep the 15 lm heads (logits) in fp32 + self.lm_head = self.lm_head.float() + kind = "torch.ao dynamic int8 Linear, fp32 lm heads (eager)" + else: + self.lm_head = quantize_dynamic(self.lm_head.float(), {nn.Linear}, dtype=torch.qint8) + kind = "torch.ao dynamic int8 Linear (eager)" self.small_to_mtp_projection = self.small_to_mtp_projection.float() self._model_dtype = torch.float32 self._lm_heads_list = list(self.lm_head)