# glm53-compile-05-workspace-acquisition-opaque.patch — the indexer's workspace
# acquisition must not be traced.
#
# STATUS: WIP — fifth of the "compile the GLM trunk" changelist. This is a break worth
# removing, but NOT part of the confirmed blocker and NOT required for the
# break-tolerant route: `torch.compiler.disable` is itself a graph break, so this patch
# helps only a hypothetical fullgraph=True route, where the workspace acquisition would
# be better RESTRUCTURED out of the forward (allocate at init/profile time) than made
# opaque. Kept because the discovery is real and the wrapper is a no-op when nothing is
# compiling.
#
# THE BREAK — CONFIRMED with AND WITHOUT the logging flag (2026-09-30):
#
#   torch._dynamo.exc.Unsupported: Unsupported method call
#   user stack
#     glm5next/common/model.py:748      Glm5NextModel.forward
#     glm5next/common/model.py:526      Glm5NextDecoderLayer.forward
#     glm5next/common/attention.py:599 / :401
#     glm5next/amd/sparse_indexer.py:711  forward_native
#     glm5next/amd/sparse_indexer.py:132  sparse_attn_indexer_kpool
#     v1/worker/workspace.py:135          get_simultaneous
#     v1/worker/workspace.py:155          _ensure_workspace_size
#
# Line 132 is the PROFILE-RUN branch — the source comment says what it is:
#
#     if not isinstance(attn_metadata, dict):
#         # Reserve workspace for indexer during profiling run
#
# so the engine dies exactly where it sizes its memory.  `_ensure_workspace_size`
# reads ContextVar and ubatch state (`_workspace_lane.get()`,
# `dbo_current_ubatch_id()`) and may allocate; none of that is traceable.
#
# FIX: the documented escape hatch — `torch.compiler.disable` around a thin wrapper,
# so dynamo stops analysing the call and executes it eagerly from inside the compiled
# region.  Chosen over decorating the workspace manager's own methods because that
# would change workspace behaviour for every model, not just this one.  The wrapper is
# a no-op when nothing is compiling, so it needs no env gate and is safe in the
# stock (uncompiled) path — which is why this patch can be reasoned about
# independently of whether the changelist ever ships.
#
# CAVEAT, stated plainly: a compiler-disabled region is still a graph break, so this
# needs patch 03 (VLLM_COMPILE_FULLGRAPH=0) to be tolerable at all, and whether dynamo
# can build a resume function around it is exactly what the next boot decides — the
# difference between this failure and 04's `Read uninitialized cell`.
#
# TARGET: vllm/models/glm5next/amd/sparse_indexer.py   (inside the pinned image)
#   base    sha256 f773a99265d3f440c8fad22c8fb5a8b33f926162094cffc3a2ecdaa42763c5b1
#   patched sha256 c154115afe8253a35c652b7198929ef4c1874957d333fdb3340cacf88e612b78
#
# APPLY: `patch -p1` from the vllm package root, or bind-mount over site-packages.
--- a/vllm/models/glm5next/amd/sparse_indexer.py
+++ b/vllm/models/glm5next/amd/sparse_indexer.py
@@ -35,6 +35,21 @@
 from vllm.v1.attention.ops.common import pack_seq_triton, unpack_seq_triton
 from vllm.v1.worker.workspace import current_workspace_manager
 
+@torch.compiler.disable
+def _acquire_workspace(*shapes_and_dtypes):
+    """Acquire shared workspace buffers, untraced.
+
+    Workspace acquisition reads ContextVar/ubatch state and may allocate
+    (v1/worker/workspace.py: ``_ensure_workspace_size`` -> ``_workspace_lane.get()``,
+    ``dbo_current_ubatch_id()``), which dynamo cannot trace -- it raises
+    ``Unsupported: Unsupported method call`` and kills the engine during the profile
+    run.  ``torch.compiler.disable`` keeps the call as an opaque, eagerly-executed
+    region inside the compiled graph.  No-op when not compiling.
+    """
+
+    return current_workspace_manager().get_simultaneous(*shapes_and_dtypes)
+
+
 logger = init_logger(__name__)
 
 # kpool write helper: form pools from the current token batch and compress them
@@ -129,7 +144,7 @@
         values_spec, scales_spec = _gather_workspace_shapes(
             total_seq_lens, head_dim, fp8_dtype, use_fp4_cache
         )
-        current_workspace_manager().get_simultaneous(
+        _acquire_workspace(
             values_spec,
             scales_spec,
             ((RADIX_TOPK_WORKSPACE_SIZE,), torch.uint8),
@@ -275,11 +290,10 @@
         # Layout switches between FP8 (head_dim bytes + 4-byte fp32 scale) and
         # MXFP4 (head_dim/2 bytes packed + head_dim/MXFP4_BLOCK_SIZE ue8m0
         # scales) based on use_fp4_cache.
-        workspace_manager = current_workspace_manager()
         values_spec, scales_spec = _gather_workspace_shapes(
             total_seq_lens, head_dim, fp8_dtype, use_fp4_cache
         )
-        k_quant_full, k_scale_full = workspace_manager.get_simultaneous(
+        k_quant_full, k_scale_full = _acquire_workspace(
             values_spec,
             scales_spec,
         )