From 3085f501607a3590fefc6831d4a6d9ea6234a68d Mon Sep 17 00:00:00 2001 From: tzhouam Date: Thu, 10 Sep 2026 03:12:40 +0800 Subject: [PATCH] [Edge][CPU] Compile the per-sub-step lm head + top-k Gumbel sampling + embedding + projection into one graph Noise is drawn eagerly with the request generators (same RNG consumption as the eager loop) and passed in, so the compiled post-step is generator-free and shape-static; used on the fp32/torchao CPU path outside stored (top-p) mode. Co-Authored-By: Claude Fable 5.1 --- .../common/test_code_predictor_incremental.py | 28 ++++++++ .../models/common/qwen3_code_predictor.py | 65 +++++++++++++++++++ 2 files changed, 93 insertions(+) diff --git a/tests/model_executor/models/common/test_code_predictor_incremental.py b/tests/model_executor/models/common/test_code_predictor_incremental.py index 2e22ee2f2..12c4e7b76 100644 --- a/tests/model_executor/models/common/test_code_predictor_incremental.py +++ b/tests/model_executor/models/common/test_code_predictor_incremental.py @@ -97,3 +97,31 @@ def test_compiled_static_step_matches_eager(): mask = (torch.arange(max_seq) <= step).view(1, 1, 1, max_seq).expand(bsz, 1, 1, max_seq) out = step_fn(x[:, step : step + 1], pos[:, step : step + 1], caches, pos_t, mask) assert torch.allclose(out[:, 0], full[:, step], atol=1e-4, rtol=1e-4), step + + +@torch.inference_mode() +def test_post_step_sample_matches_eager_gumbel(): + from vllm_omni.model_executor.models.common.qwen3_code_predictor import _post_step_sample + + torch.manual_seed(3) + bsz, hidden, vocab = 2, 64, 32 + h = torch.randn(bsz, hidden) + head_w = torch.randn(vocab, hidden) + embed_w = torch.randn(vocab, hidden) + proj_w = torch.randn(hidden, hidden) + proj_b = torch.randn(hidden) + u = torch.rand(bsz, vocab).clamp_(1e-20, 1 - 1e-20) + # eager reference (the loop body of the predictor): top-k mask + Gumbel-max + logits = h @ head_w.T + scaled = logits * (1 / 0.9) + topk_vals, _ = scaled.topk(5, dim=-1) + scaled = scaled.masked_fill(scaled < topk_vals[:, -1:], float("-inf")) + ref_code = (scaled.float() - torch.log(-torch.log(u))).argmax(-1, keepdim=True) + ref_proj = torch.nn.functional.linear(embed_w[ref_code.view(-1)], proj_w, proj_b) + code, proj_row = _post_step_sample(h, head_w, embed_w, proj_w, proj_b, u, 1 / 0.9, 5, True) + assert torch.equal(code, ref_code) and torch.allclose(proj_row, ref_proj, atol=1e-5) + code_g, _ = _post_step_sample(h, head_w, embed_w, proj_w, proj_b, u, 0.0, 0, False) + assert torch.equal(code_g.view(-1), logits.argmax(-1)) + compiled = torch.compile(_post_step_sample, dynamic=False) + code_c, proj_c = compiled(h, head_w, embed_w, proj_w, proj_b, u, 1 / 0.9, 5, True) + assert torch.equal(code_c, ref_code) and torch.allclose(proj_c, ref_proj, atol=1e-4) 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 778759997..7584eaebe 100644 --- a/vllm_omni/model_executor/models/common/qwen3_code_predictor.py +++ b/vllm_omni/model_executor/models/common/qwen3_code_predictor.py @@ -462,6 +462,36 @@ class CodePredictorDecoderLayer(nn.Module): # =================================================================== +def _post_step_sample( + hidden: torch.Tensor, + head_w: torch.Tensor, + embed_w: torch.Tensor, + proj_w: torch.Tensor, + proj_b: torch.Tensor | None, + u: torch.Tensor, + inv_temperature: float, + top_k: int, + use_sampling: bool, +) -> tuple[torch.Tensor, torch.Tensor]: + """lm head -> top-k mask -> Gumbel-max sample (noise ``u`` drawn by the caller) -> embed -> project. + + Shape-static and generator-free so it compiles into one graph per (top_k, use_sampling). + Returns ``(code [B, 1], proj_row [B, H])``. + """ + logits = F.linear(hidden, head_w) + if use_sampling: + scaled = logits * inv_temperature + if top_k > 0: + topk_vals, _ = scaled.topk(top_k, dim=-1) + scaled = scaled.masked_fill(scaled < topk_vals[:, -1:], float("-inf")) + code = (scaled.float() - torch.log(-torch.log(u))).argmax(dim=-1, keepdim=True) + else: + code = logits.argmax(dim=-1, keepdim=True) + new_embed = F.embedding(code, embed_w) # [B, 1, H] + proj_row = F.linear(new_embed.reshape(code.shape[0], -1), proj_w, proj_b) + return code, proj_row + + class CodePredictorBaseModel(nn.Module): """Inner transformer for code predictor. @@ -1206,6 +1236,7 @@ class CodePredictorWrapper(nn.Module): quantize_(self.model, Int8WeightOnlyConfig()) quantize_(self.lm_head, Int8WeightOnlyConfig()) self._static_step_fn = torch.compile(self.model.forward_static_step, dynamic=False) + self._post_step_fn = torch.compile(_post_step_sample, dynamic=False) model = self.model def _prefill(x, pos_ids, caches): @@ -1446,6 +1477,40 @@ class CodePredictorWrapper(nn.Module): step, ) hidden_step = hidden_new[:, 0, :] + post_step = getattr(self, "_post_step_fn", None) + if ( + post_step is not None + and not stored_mode + and getattr(self, "_parity_ref", None) is None + and not self._wrapper_config.return_proj_buf + and isinstance(projection, nn.Linear) + ): + head = lm_heads[step - 1] + proj = projection + u = torch.empty(bsz, head.weight.shape[0], dtype=torch.float32, device=device) + row_generators = self._normalize_generators(sample_generator, bsz) + if isinstance(row_generators, list): + for row, row_generator in enumerate(row_generators): + u[row : row + 1].uniform_(_UNIFORM_EPS, 1.0 - _UNIFORM_EPS, generator=row_generator) + else: + u.uniform_(_UNIFORM_EPS, 1.0 - _UNIFORM_EPS, generator=row_generators) + code, proj_row = post_step( + hidden_step, + head.weight, + codec_embeds[step - 1].weight, + proj.weight, + proj.bias, + u, + float(sample_kwargs["inv_temperature"]), + int(sample_kwargs["top_k"]), + bool(sample_kwargs["use_sampling"]), + ) + all_codes[:, step] = code.reshape(bsz) + if step < num_groups - 1: + proj_buf[:bsz, step + 1, :] = proj_row + if _stats.enabled: + _stats.add("predictor.substep_ms", (time.perf_counter() - _t_sub) * 1000.0) + continue logits = lm_heads[step - 1](hidden_step) if getattr(self, "_parity_ref", None) is not None: if step == 1: