# glm53-compile-01-model-decorator.patch — put the GLM-5.3-Flash trunk under torch.compile.
#
# STATUS: WIP — first of three patches in the "compile the GLM trunk" changelist. The
# set does NOT yet boot the engine (see 03 for the next blocker), so it is NOT part of
# the deployed frank-primary shape. Nothing in scripts/ mounts these files.
#
# WHY THIS IS THE RIGHT LEVER (measured 2026-09-30, not inferred):
#
# The fork's own flagship trunk — vllm/models/qwen4_exp/amd/model.py — carries
# @support_torch_compile. vllm/models/glm5next/common/model.py does not, and the engine
# says so on every boot:
#
#   Cudagraph mode FULL_AND_PIECEWISE requires piecewise capture, but the loaded model
#   provides neither a compiled submodule nor breakable CUDA graphs.
#   Overriding to FULL_DECODE_ONLY.
#
# A torch-profiler window on the live endpoint (496 decode tokens, rank0, 3 x 160-token
# generations at the deployed shape) shows what the missing compilation costs. Total
# device kernel time 46.18 s / 496 tokens = ~93 ms per token, against a measured step
# time of ~104 ms (9.6 tok/s with the profiler attached):
#
#   16.07 s  34.8 %  708,500 launches  vectorized_elementwise_kernel      (22.7 us each)
#    4.60 s  10.0 %    42,000           r4d_gemm_moe_w4a16_nt_b16_m64     (the MoE GEMM)
#    3.11 s   6.7 %   174,708           elementwise_kernel_manual_unroll
#    2.94 s   6.4 %   138,252           wvSplitK_hf_sml_                  (split-K GEMM)
#    4.09 s   8.9 %    83,160           topk/router (sbtopk gatherTopK + warpMergeSortTopK
#                                       + bitonicSortKVInPlace)
#    3.01 s   6.5 %   133,920           r4d_mhc_* (hyper-connection pre/post/finish)
#    2.78 s   6.0 %   127,916           inductor triton_poi_* / triton_per_* (the only
#                                       fusions that DO happen)
#    0.49 s   1.05%    20,664           r4d_moe_lru_gather_kernel         <-- the suspect
#
# That last row retired the fleet's working hypothesis. Decode on this shape is NOT
# expert-gather bound: the gather is 1 % of device time, the same order as the telemetry
# kernel that was switched on to measure it. 44 % of the budget is pointwise glue —
# bf16<->fp32 conversions, adds, pows/rsqrts from fp32 RMSNorm, clamp/scale quant
# passes — spread over ~1,428 launches per token. Inductor fusion is the only lever
# that touches that mass; the router is second at 9 % but is binary-blocked (the
# libr4d route kernel exists only as r4d_moe_route_kernel<128|256|512> instantiations
# and GLM has 288 experts, with no HIP sources shipped to rebuild).
#
# TARGET: vllm/models/glm5next/common/model.py   (inside the pinned image)
#   base    sha256 008eeaa2a07e091465319f86900b5a321d921238723ccbe22e6b9e971e1d2b28
#   patched sha256 6e9cb81b84a9a88bd41403ca404aa04c3b791899dd77d661c3f25ae686cecd40
#
# APPLY: `patch -p1` from the vllm package root, or bind-mount new/.../model.py over
# site-packages. Then R4D_GLM_COMPILE=1 (default) enables it; =0 makes the decorator a
# no-op, so the file can stay mounted for an A/B without a second artifact.
# Requires 02 (trace-safe probes) — without it the worker dies during profile_run on
# the first probe call inside the traced region.
--- a/vllm/models/glm5next/common/model.py
+++ b/vllm/models/glm5next/common/model.py
@@ -1,12 +1,14 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import os
 from collections.abc import Iterable
 from typing import ClassVar, Literal
 
 import torch
 from torch import nn
 
+from vllm.compilation.decorators import support_torch_compile
 from vllm.config import ParallelConfig, VllmConfig
 from vllm.distributed import (
     get_ep_group,
@@ -623,6 +625,18 @@
         )
 
 
+_GLM_COMPILE = os.environ.get("R4D_GLM_COMPILE", "1") == "1"
+
+
+@support_torch_compile(
+    dynamic_arg_dims={
+        "input_ids": 0,
+        "positions": -1,
+        "intermediate_tensors": 0,
+        "inputs_embeds": 0,
+    },
+    enable_if=lambda _cfg: _GLM_COMPILE,
+)
 class Glm5NextModel(nn.Module):
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         super().__init__()