# 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 BEFORE TRUSTING ANY CLAIM BELOW.  Two break-tolerance runs
# disagree, and the difference is a logging flag:
#   * 01+02+03, AOT off, breaks tolerated, NO TORCH_LOGS  -> gets PAST all graph-break
#     handling and dies on the dynamo name-map KeyError ('in_proj_qkvbfg_a').
#   * same + TORCH_LOGS=graph_breaks                       -> dies on the fatal
#     `Unsupported: Read uninitialized cell` at this very call site.
# The likely mechanism is that `graph_breaks` logging makes dynamo build a break
# *explanation* (walk the frames, name the reason), and that analysis is what trips on
# the uninitialized cell; without the flag the break is taken silently.  So the cell
# error is PROBABLY an artifact of the diagnostic, not necessarily a blocker for a
# break-tolerant compile.  Whether this patch is therefore necessary — as opposed to
# removing a call dynamo cannot *explain* — is being decided by a flag-free run of
# 01+02+03+04.  Do not repeat the earlier overclaim ("fatal even with fullgraph=False");
# it was written from the confounded run.
#
# IT IS STILL THE RIGHT PATTERN TO REMOVE, either way: the launcher family
# (@kernel_launcher, 39 call sites) is exactly what breaks dynamo's break analysis, and
# the object it 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(