from pathlib import Path
import ast
p=Path(__file__).resolve().parent
s=(p/'vllm_ple_mmap.original.py').read_text()
replacements=[
('            torch.cuda.current_stream(ids.device).synchronize()', '            with torch.cuda.nvtx.range("PLE_wait_for_GPU"):\n                torch.cuda.current_stream(ids.device).synchronize()'),
('        rows = table.gather(uniq)  # uint8 [U, row_bytes], fresh & writable','        with torch.cuda.nvtx.range("PLE_CPU_gather"):\n            rows = table.gather(uniq)  # uint8 [U, row_bytes], fresh & writable'),
('        uniq, inverse = np.unique(ids_np, return_inverse=True)', '        with torch.cuda.nvtx.range("PLE_CPU_dedup"):\n            uniq, inverse = np.unique(ids_np, return_inverse=True)')]
for before,after in replacements:
 assert s.count(before)==1,(before,s.count(before))
 s=s.replace(before,after)
s+='''\n# Temporary diagnostic annotations; no tensor or arithmetic changes.
_profile_original_lookup_ids = _lookup_ids_impl
def _lookup_ids_impl(ngram_ids: torch.Tensor, output: torch.Tensor, layer_name: str) -> None:
    with torch.cuda.nvtx.range("PLE_lookup_tokens=" + str(ngram_ids.shape[0])):
        return _profile_original_lookup_ids(ngram_ids, output, layer_name)
'''
ast.parse(s);(p/'vllm_ple_mmap.profile.py').write_text(s)
s=(p/'cudagraph_utils.original.py').read_text()
s+='''\n# Diagnostic dispatch marker; does not change graph selection.
_profile_original_dispatch = CudaGraphManager.dispatch
def _profile_dispatch(self, *args, **kwargs):
    desc = _profile_original_dispatch(self, *args, **kwargs)
    actual = args[1] if len(args) > 1 else kwargs.get("num_tokens", -1)
    torch.cuda.nvtx.mark(f"CG_DISPATCH:actual={actual}:mode={desc.cg_mode.name}:padded={desc.num_tokens}")
    return desc
CudaGraphManager.dispatch = _profile_dispatch
'''
ast.parse(s);(p/'cudagraph_utils.profile.py').write_text(s)