docs(perf): profile CUDA graph coverage and PLE costs on Spark

experiments/profile128k/instrument.py

@@ -0,0 +1,29 @@
+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)