# 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(