30 lines
1.8 KiB
Python
30 lines
1.8 KiB
Python
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)
|