#!/usr/bin/env python3
"""GDN per-layer glue: strided gates (measured neutral) and the core_attn_out zero-fill.

Edit 3 (RADIANCE_GDN_EMPTY_OUT): allocate core_attn_out with torch.empty. vLLM zero-fills it
because its kernels write only real rows and cudagraph pad rows would leak garbage (PR 28182);
the rx5 r4d fused_update zeroes rows [cu[N], o_rows) itself, radiance_gdn zeroes the tail on
its non-fused paths, and the Triton fallback below zero-fills the whole buffer. 48 launches/step.
VERDICT 2026-09-02: neutral (22.28-22.32 vs 22.29-22.33 ms/step, byte-identical, GSM8K @conc 8
98.00%). Kept dark at 0.

Edit 1-2: skip the two `.contiguous()` copies on the GDN gate tensors when the R4D path serves the step.

In QwenGDN's forward_cuda (Qwen3.5 layout) `b, a = self.split_ba(ba)` are column slices of the
in_proj_ba output, and vLLM copies both to contiguous before the core op. The R4D kernels take a
row stride (radiance_gdn passes `a.stride(0)`, and _plan only requires the head axis to be unit
stride), so on that path the copies are dead work: one launch + one ~3 us gap per linear-attention
layer, 48 layers per decode step (census 2026-09-02). The Triton fallback body still gets
contiguous tensors: the copies move to just after the radiance hook's return, where only a step
the R4D path declined reaches them. Gated by RADIANCE_GDN_STRIDED_GATES at runtime (default 0) so
the same build A/Bs; it needs patch_r4d.py to have run first (the hook text is the anchor).

VERDICT 2026-09-02: neutral at serve level (22.34-22.41 vs 22.31-22.33 ms/step on top of the fused
GDN norm, output byte-identical, GSM8K 97.80%). Kept dark at 0; the copies are evidently not on
the critical path, or inductor re-packs the custom-op inputs regardless.
"""
import sysconfig
from pathlib import Path

from _patchlib import apply, apply_any

PURELIB = Path(sysconfig.get_paths()["purelib"])
L = PURELIB / "vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py"

SPLIT_OLD = (
    "            b, a = self.split_ba(ba)\n"
    "            b = b.contiguous()\n"
    "            a = a.contiguous()\n"
)
SPLIT_NEW = (
    "            b, a = self.split_ba(ba)\n"
    "            # --- RADIANCE (patch_gdn_glue.py): the R4D core reads strided gates ---\n"
    "            if not (_radiance_gdn is not None and _radiance_gdn.STRIDED_GATES):\n"
    "                b = b.contiguous()\n"
    "                a = a.contiguous()\n"
)
HOOK_OLD = (
    "            if _radiance_gdn.forward_core_fused(self, mixed_qkv, b, a, core_attn_out):\n"
    "                return\n"
)
HOOK_NEW = (
    HOOK_OLD
    + "            b = b.contiguous()   # patch_gdn_glue.py: the Triton body below wants them packed\n"
    + "            a = a.contiguous()\n"
    + "            if _radiance_gdn.EMPTY_OUT:   # the Triton body writes only real rows\n"
    + "                core_attn_out.zero_()\n"
)
ZEROS_OLD = (
    "        core_attn_out = torch.zeros(\n"
    "            (num_tokens, self.num_v_heads // self.tp_size, self.head_v_dim),\n"
    "            dtype=hidden_states.dtype,\n"
    "            device=hidden_states.device,\n"
    "        )\n"
    "\n"
    "        torch.ops.vllm.qwen_gdn_attention_core(\n"
)
ZEROS_NEW = (
    "        # --- RADIANCE (patch_gdn_glue.py): the rx5 fused_update zeroes the pad rows itself ---\n"
    "        _alloc = (torch.empty if (_radiance_gdn is not None and _radiance_gdn.EMPTY_OUT)\n"
    "                  else torch.zeros)\n"
    "        core_attn_out = _alloc(\n"
    "            (num_tokens, self.num_v_heads // self.tp_size, self.head_v_dim),\n"
    "            dtype=hidden_states.dtype,\n"
    "            device=hidden_states.device,\n"
    "        )\n"
    "\n"
    "        torch.ops.vllm.qwen_gdn_attention_core(\n"
)


# --- vLLM 0.29 shape --------------------------------------------------------------------------
# 0.29 dropped the two .contiguous() calls outright, so upstream now hands the core strided gates
# unconditionally. We re-add them under the same knob rather than treating the hunk as superseded:
# with RADIANCE_GDN_STRIDED_GATES=0 (the launcher default) the R4D core is entitled to assume
# contiguous b/a, and inheriting upstream's new behaviour would quietly break that assumption.
# With the knob on, this matches upstream exactly.
SPLIT_OLD_029 = (
    "            b, a = self.split_ba(ba)\n"
)

SPLIT_NEW_029 = (
    "            b, a = self.split_ba(ba)\n"
    "            # --- RADIANCE (patch_gdn_glue.py): the R4D core reads strided gates ---\n"
    "            if not (_radiance_gdn is not None and _radiance_gdn.STRIDED_GATES):\n"
    "                b = b.contiguous()\n"
    "                a = a.contiguous()\n"
)


apply_any(L, [(SPLIT_OLD, SPLIT_NEW), (SPLIT_OLD_029, SPLIT_NEW_029)],
          "patch_gdn_glue.py): the R4D core reads strided gates", "gdn strided gates")
apply(L, HOOK_OLD, HOOK_NEW, "patch_gdn_glue.py: the Triton body below", "gdn fallback re-pack")
apply(L, ZEROS_OLD, ZEROS_NEW, "the rx5 fused_update zeroes the pad rows itself", "gdn core_attn_out alloc")