Comrade Chief Designer
71db0debcc
Root-caused the moving failures: they were my own scratch launcher. Its compile branches were restructured and the '-e VLLM_COMPILE_FULLGRAPH=0' line was dropped, so four consecutive boots mounted the wrapper patch while _FULLGRAPH still defaulted to True — and with fullgraph=True EVERY break is an error, so each boot reported whichever break came first (cell -> workspace -> compiler-disable). I had blamed TORCH_LOGS=graph_breaks for the difference; that was wrong, the flag only logs. The launcher sets the env var again, and patch 04's header, the runbook and the work log now carry the true run history as a table instead of the bad theory. New artifact: scripts/patches/glm53-compile-05-workspace-acquisition-opaque.patch — the indexer acquires its shared workspace INSIDE the forward (its profile-run branch literally says 'Reserve workspace for indexer during profiling run'), and v1/worker/workspace.py:_ensure_workspace_size reads ContextVar/ubatch state, which dynamo cannot trace. Wrap both call sites in a torch.compiler.disable'd helper (narrower than decorating the workspace manager, which every model would go through; a no-op when not compiling). All five artifacts verified against the pinned image: applied in sequence with patch -p1 they reproduce their recorded patched sha256 exactly (6e9cb81b…, 0751852e…, 732363bb…, 84969993…, c154115a…). tests/run.sh: 362 seams pass.
113 lines
5.6 KiB
Diff
113 lines
5.6 KiB
Diff
# glm53-compile-04-fused-qkv-rmsnorm-optout.patch — take GLM's fused-norm JIT kernel
|
|
# out of the traced path.
|
|
#
|
|
# STATUS: WIP — fourth of the "compile the GLM trunk" changelist. Needs 01 (the
|
|
# decorator) to matter; not part of the deployed shape.
|
|
#
|
|
# MEASUREMENT NOTE — read the run history before trusting any claim below, because a
|
|
# launcher bug made four consecutive boots misleading. The scratch launcher
|
|
# (serve_glm_autoround.sh) was restructured on 2026-09-30 and the
|
|
# `-e VLLM_COMPILE_FULLGRAPH=0` line was DROPPED, so those boots mounted the wrapper
|
|
# patch while `_FULLGRAPH` still defaulted to True. With fullgraph=True *every* break
|
|
# is an error, which is why the reported failure kept moving as breaks were removed:
|
|
#
|
|
# fullgraph=False, AOT off -> breaks tolerated; dies later on the dynamo name-map
|
|
# `KeyError: 'in_proj_qkvbfg_a'`
|
|
# fullgraph=True, AOT off -> dies on whichever break comes first
|
|
# (cell -> workspace -> compiler-disable)
|
|
#
|
|
# An earlier version of this note blamed `TORCH_LOGS=graph_breaks` for the difference.
|
|
# That was WRONG — the flag only logs; the launcher silently ran fullgraph=True.
|
|
# Launcher fixed; the env var is set again.
|
|
#
|
|
# What it means for THIS patch: under fullgraph=True the fused-norm call was the first
|
|
# break dynamo reported, and after this patch the report moved on to the workspace
|
|
# acquisition — so it removes a real break. It is not by itself what unblocks a boot:
|
|
# with breaks genuinely tolerated the blocker is the name-map KeyError.
|
|
#
|
|
# IT IS STILL THE RIGHT PATTERN TO REMOVE, either way: the launcher family
|
|
# (@kernel_launcher, 39 call sites) is exactly the shape that produces breaks, and the
|
|
# object this one closes over is bound after use (see below).
|
|
#
|
|
# THE BREAK (measured, TORCH_LOGS=graph_breaks, rank0, 2026-09-30). With 01+02+03 the
|
|
# trunk traces, and with AOT off + fullgraph=True dynamo reports a FATAL break — fatal
|
|
# even with VLLM_COMPILE_FULLGRAPH=0, because dynamo cannot build a resume function
|
|
# around it:
|
|
#
|
|
# torch._dynamo.exc.Unsupported: Read uninitialized cell
|
|
# Explanation: Attempted to read a cell variable that has not been populated yet.
|
|
# Developer debug context: CellVariable()
|
|
#
|
|
# User stack:
|
|
# glm5next/common/model.py:748 Glm5NextModel.forward layer(...)
|
|
# glm5next/common/model.py:526 Glm5NextDecoderLayer.forward
|
|
# glm5next/common/attention.py:590
|
|
# model_executor/layers/mla.py:208
|
|
# q_proj_input, kv_c_normed = fused_q_kv_rmsnorm(...)
|
|
# -> torch/_dynamo/polyfills/__init__.py:392 getattr_and_trace
|
|
#
|
|
# So dynamo is tracing the fork's JIT-kernel launcher framework
|
|
# (VllmTritonJitKernel + @kernel_launcher — 39 call sites) and finds a closure name
|
|
# that is not yet bound. In models/common/ops/fused_qk_rmsnorm.py the kernel object
|
|
# is bound AFTER the function that closes over it:
|
|
#
|
|
# def fused_q_kv_rmsnorm(...): # line 194
|
|
# ...
|
|
# _FUSED_Q_KV_RMSNORM_KERNEL(...) # line 211
|
|
# _FUSED_Q_KV_RMSNORM_KERNEL = FusedQKVRMSNormKernel() # line 223
|
|
#
|
|
# Two ways out, and this patch takes the cheap one:
|
|
# (a) WORKAROUND (here): stop calling the fused kernel from the traced region. The
|
|
# flag is already a constructor argument (mla.py:80, default False) that glm5next
|
|
# turns ON at attention.py:583, so it becomes an env knob.
|
|
# (b) ROOT FIX (not done): make the launcher trace-safe for all 39 sites. Any other
|
|
# @kernel_launcher-backed kernel still in the traced path can produce the same
|
|
# fatal break, so (b) is the durable fix and (a) is what unblocks iterating.
|
|
#
|
|
# COST of R4D_GLM_FUSE_QKV_RMSNORM=0: one extra small RMSNorm launch per MLA layer per
|
|
# token (q_a_layernorm + kv_a_layernorm instead of one fused kernel) — ~2 launches out
|
|
# of ~3,700 per token on this shape, against the 44 % of device time currently spent in
|
|
# unfused pointwise work. Correct trade for a compile experiment.
|
|
#
|
|
# TARGET: vllm/models/glm5next/common/attention.py (inside the pinned image)
|
|
# base sha256 362bbf0e6c50cdfa446c6551b9b1ddff425e08f88e3d4140fd8db88f8a619e0c
|
|
# patched sha256 84969993c250e9f8183a33a6e1ddd607272f48b950b560272eb33b7cdde31189
|
|
#
|
|
# APPLY: `patch -p1` from the vllm package root (verified: applies clean, reproduces
|
|
# the patched sha256), or bind-mount new/.../attention.py over site-packages.
|
|
# R4D_GLM_FUSE_QKV_RMSNORM=1 (default) restores the fused kernel.
|
|
--- a/vllm/models/glm5next/common/attention.py
|
|
+++ b/vllm/models/glm5next/common/attention.py
|
|
@@ -1,6 +1,7 @@
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
+import os
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from torch import nn
|
|
@@ -36,6 +37,14 @@
|
|
from vllm.utils.deep_gemm import PAGED_MQA_PAGE_SIZES
|
|
from vllm.v1.kv_cache_interface import KpoolTailSpec, MLAAttentionSpec
|
|
|
|
+
|
|
+# FORK-LOCAL (patch: keep the traced region trace-safe). R4D_GLM_FUSE_QKV_RMSNORM=0
|
|
+# replaces the fused `fused_q_kv_rmsnorm` JIT-kernel launch (one kernel) with two plain
|
|
+# RMSNorms, because dynamo cannot trace the launcher's closure ("Read uninitialized cell",
|
|
+# fatal even with fullgraph=False). The fused kernel object in
|
|
+# models/common/ops/fused_qk_rmsnorm.py is bound AFTER the function that uses it.
|
|
+_FUSE_QKV_RMSNORM = os.environ.get("R4D_GLM_FUSE_QKV_RMSNORM", "1") == "1"
|
|
+
|
|
logger = init_logger(__name__)
|
|
|
|
# Shared torch.compile config for the indexer's small-kernel leaves. The MLA
|
|
@@ -580,7 +589,7 @@
|
|
quant_config,
|
|
prefix,
|
|
skip_topk=False,
|
|
- fuse_qkv_rmsnorm=True,
|
|
+ fuse_qkv_rmsnorm=_FUSE_QKV_RMSNORM,
|
|
)
|
|
|
|
def forward(
|