#!/usr/bin/env python3
"""Make the b12x runner's multimodal `inputs_embeds` buffers follow the
model's native embed shape instead of a fixed 2-D [T, H].

At boot the warmup probed `x` shapes of [16, 16, 4096] (hc-stream run) and
[14, 4096] (plain run) from the SAME layer — the dsv4 b12x embedding emits
varying rank per request. Re-fit the buffer whenever the non-token tail
changes so both the mm merge and CUDA-graph capture see a stable tensor.
"""
from pathlib import Path

import vllm

p = (
    Path(vllm.__file__).parent
    / "v1/worker/gpu/mm/encoder_runner.py"
)
text = p.read_text()
old = """        x = self.model.embed_input_ids(
            input_ids, multimodal_embeddings=mm_embeds, is_multimodal=is_mm_embed
        )
        # Copy to the pre-allocated buffer for CUDA graphs.
        self.inputs_embeds[: x.shape[0]] = x
        return self.inputs_embeds
"""
new = """        x = self.model.embed_input_ids(
            input_ids, multimodal_embeddings=mm_embeds, is_multimodal=is_mm_embed
        )
        # dsv4 b12x embeds vary in rank per request ([T, K, H] hc-stream runs,
        # [T, H] plain): re-fit the buffer when the non-token tail changes.
        tail = tuple(x.shape[1:])
        if (
            self.inputs_embeds.dim() != x.dim()
            or tuple(self.inputs_embeds.shape[1:]) != tail
        ):
            self.inputs_embeds = torch.zeros(
                self.max_num_tokens, *tail, dtype=x.dtype, device=x.device
            )
        # Copy to the pre-allocated buffer for CUDA graphs.
        self.inputs_embeds[: x.shape[0]] = x
        return self.inputs_embeds
"""
if old not in text:
    raise SystemExit(f"encoder_runner patch target missing in {p}")
p.write_text(text.replace(old, new, 1))
print("encoder_runner mm shape-follow patch applied")

# Same shape-follow for the DSpark speculator's re-embed buffer: with the
# multimodal gate on, the MTP speculator re-embeds draft tokens through the
# target's embed and assigns [B, T, H] rows into a [max, hidden] buffer.
sp = (
    Path(vllm.__file__).parent
    / "v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py"
)
sp_text = sp.read_text()
sp_old = """        if self.inputs_embeds is not None:
            mm_embeds, is_mm_embed = mm_inputs or (None, None)
            self.inputs_embeds[:num_tokens] = self.model.embed_input_ids(
                self.input_buffers.input_ids[:num_tokens],
                multimodal_embeddings=mm_embeds,
                is_multimodal=is_mm_embed,
            )
"""
sp_new = """        if self.inputs_embeds is not None:
            mm_embeds, is_mm_embed = mm_inputs or (None, None)
            emb = self.model.embed_input_ids(
                self.input_buffers.input_ids[:num_tokens],
                multimodal_embeddings=mm_embeds,
                is_multimodal=is_mm_embed,
            )
            if emb.dim() != self.inputs_embeds.dim() or tuple(
                emb.shape[1:]
            ) != tuple(self.inputs_embeds.shape[1:]):
                self.inputs_embeds = torch.zeros_like(
                    emb, dtype=emb.dtype, device=emb.device
                )
            self.inputs_embeds[:num_tokens] = emb
"""
if sp_old in sp_text:
    sp.write_text(sp_text.replace(sp_old, sp_new, 1))
    print("speculator re-embed shape-follow patch applied")
else:
    print("speculator re-embed pattern not found — skipping (harmless)")