!11 feat(p800): enable operator-level timing in profiling experiment
Merge pull request !11 from yy-fighting/auto/main/13188872/d44301a2-1
This commit is contained in:
Quantong Qiu
Gitee
commit
680f2c31d6
@ -81,6 +81,7 @@ env_args=(
|
||||
-e BKCL_RDMA_VERBS=1
|
||||
-e XSGL_INT8_LM_HEAD=1
|
||||
-e SGLANG_P800_ALL_GATHER_FALLBACK=0
|
||||
-e SGLANG_TIMING_ENABLED=1
|
||||
-e XPU_ENABLE_PROFILER_TRACING=1
|
||||
)
|
||||
|
||||
@ -99,7 +100,7 @@ set -Eeuo pipefail
|
||||
cd /workspace
|
||||
find /root/miniconda/envs/python310_torch25_cuda/lib/python3.10/site-packages/sglang -type d -name __pycache__ -exec rm -rf {} + 2>/dev/null || true
|
||||
/root/miniconda/envs/python310_torch25_cuda/bin/pip install --upgrade safetensors -q 2>/dev/null
|
||||
/root/miniconda/envs/python310_torch25_cuda/bin/python -c "import torch; torch.float8_e8m0fnu = torch.uint8; import runpy, sys; sys.argv[0] = 'sglang.launch_server'; runpy.run_module('sglang.launch_server', run_name='__main__')" LAUNCH_ARGS_PLACEHOLDER
|
||||
/root/miniconda/envs/python310_torch25_cuda/bin/python -c "import p800_timing; import torch; torch.float8_e8m0fnu = torch.uint8; import runpy, sys; sys.argv[0] = 'sglang.launch_server'; runpy.run_module('sglang.launch_server', run_name='__main__')" LAUNCH_ARGS_PLACEHOLDER
|
||||
ENTRYEOF
|
||||
|
||||
# Insert the launch args
|
||||
@ -116,6 +117,7 @@ patch_mounts=(
|
||||
-v "${PATCH_ROOT}/config_backup_small_w8a8_int8.json:/root/miniconda/envs/python310_torch25_cuda/lib/python3.10/site-packages/sglang/srt/configs/config_backup_small.json:ro"
|
||||
-v "${PATCH_ROOT}/parallel_state.py:/root/miniconda/envs/python310_torch25_cuda/lib/python3.10/site-packages/sglang/srt/distributed/parallel_state.py:ro"
|
||||
-v "${PATCH_ROOT}/vocab_parallel_embedding.py:/root/miniconda/envs/python310_torch25_cuda/lib/python3.10/site-packages/sglang/srt/layers/vocab_parallel_embedding.py:ro"
|
||||
-v "${PATCH_ROOT}/p800_timing.py:/root/miniconda/envs/python310_torch25_cuda/lib/python3.10/site-packages/p800_timing.py:ro"
|
||||
-v "${ENTRYPOINT_FILE}:/workspace/entrypoint.sh:ro"
|
||||
)
|
||||
BENCH_SERVING_PATCH="${SCRIPT_DIR}/patches/bench_serving.py"
|
||||
|
||||
275
platforms/patches/kunlun_p800/p800_timing.py
Normal file
275
platforms/patches/kunlun_p800/p800_timing.py
Normal file
@ -0,0 +1,275 @@
|
||||
"""
|
||||
P800 Manual Timing Module for SGLang
|
||||
Monkey-patches DeepSeek-V4 decoder layers to record per-op timing.
|
||||
Bypasses broken PyTorch Profiler/CUPTI on P800 XPU.
|
||||
|
||||
Usage: set SGLANG_TIMING_ENABLED=1 in container env
|
||||
"""
|
||||
|
||||
import os
|
||||
import time
|
||||
import json
|
||||
import atexit
|
||||
import signal
|
||||
import functools
|
||||
import threading
|
||||
|
||||
_lock = threading.Lock()
|
||||
_timing = {}
|
||||
_write_counter = 0
|
||||
_WRITE_INTERVAL = 5000 # write every N records
|
||||
_phase = threading.local() # per-thread: "prefill" or "decode"
|
||||
|
||||
def _get_output_path():
|
||||
base = os.environ.get("SGLANG_TIMING_OUTPUT",
|
||||
"/tmp/p800_timing_results.json")
|
||||
pid = os.getpid()
|
||||
name, ext = os.path.splitext(base)
|
||||
return f"{name}_{pid}{ext}"
|
||||
|
||||
def _write_to_file(path, summary, raw_count):
|
||||
result = {
|
||||
"summary": summary,
|
||||
"config": {
|
||||
"timestamp": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
|
||||
"pid": os.getpid(),
|
||||
},
|
||||
"raw_count": raw_count,
|
||||
}
|
||||
try:
|
||||
os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
|
||||
with open(path, "w") as f:
|
||||
json.dump(result, f, indent=2)
|
||||
except Exception as exc:
|
||||
print(f"[P800_TIMING] Write error: {exc}", flush=True)
|
||||
|
||||
def _compute_summary():
|
||||
summary = {}
|
||||
raw_count = {}
|
||||
for key, samples in _timing.items():
|
||||
if not samples or not isinstance(samples, list):
|
||||
continue
|
||||
vals = [s.get("t", 0) for s in samples
|
||||
if isinstance(s.get("t"), (int, float))]
|
||||
if not vals:
|
||||
continue
|
||||
vals.sort()
|
||||
n = len(vals)
|
||||
raw_count[key] = n
|
||||
summary[key] = {
|
||||
"count": n,
|
||||
"total_ms": round(sum(vals), 3),
|
||||
"avg_ms": round(sum(vals) / n, 4),
|
||||
"p50_ms": round(vals[n // 2], 4),
|
||||
"p95_ms": round(vals[int(n * 0.95)], 4),
|
||||
"p99_ms": round(vals[int(n * 0.99)], 4),
|
||||
}
|
||||
return summary, raw_count
|
||||
|
||||
def _flush():
|
||||
with _lock:
|
||||
summary, raw_count = _compute_summary()
|
||||
if summary:
|
||||
_write_to_file(_get_output_path(), summary, raw_count)
|
||||
|
||||
def _record(key, t_ms, extra=None, phase=None):
|
||||
global _write_counter
|
||||
if phase is None:
|
||||
phase = getattr(_phase, "mode", "unknown")
|
||||
full_key = f"{key}_{phase}" if phase != "unknown" else key
|
||||
with _lock:
|
||||
e = {"t": round(t_ms, 4), "at": time.time(), "phase": phase}
|
||||
if extra:
|
||||
e.update(extra)
|
||||
_timing.setdefault(full_key, []).append(e)
|
||||
_write_counter += 1
|
||||
if _write_counter % _WRITE_INTERVAL == 0:
|
||||
_flush()
|
||||
|
||||
def _sync():
|
||||
try:
|
||||
import torch
|
||||
torch.cuda.synchronize()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
@atexit.register
|
||||
def _write():
|
||||
_flush()
|
||||
|
||||
def _sigterm_handler(signum, frame):
|
||||
_flush()
|
||||
signal.signal(signal.SIGTERM, signal.SIG_DFL)
|
||||
os.kill(os.getpid(), signal.SIGTERM)
|
||||
|
||||
try:
|
||||
signal.signal(signal.SIGTERM, _sigterm_handler)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
def _make_timed_fn(orig_fn, key):
|
||||
"""Create a timed wrapper for a given function.
|
||||
key can be a string or a callable(args, kwargs) -> string."""
|
||||
@functools.wraps(orig_fn)
|
||||
def timed_fn(*args, **kwargs):
|
||||
_sync()
|
||||
t0 = time.time()
|
||||
try:
|
||||
return orig_fn(*args, **kwargs)
|
||||
finally:
|
||||
_sync()
|
||||
k = key(*args, **kwargs) if callable(key) else key
|
||||
_record(k, (time.time() - t0) * 1000)
|
||||
return timed_fn
|
||||
|
||||
|
||||
def patch():
|
||||
"""Apply monkey-patches to DeepSeek-V4 model classes."""
|
||||
try:
|
||||
from sglang.srt.models import deepseek_v4
|
||||
except ImportError:
|
||||
return False
|
||||
|
||||
# ── Layer forward (split prefill vs decode) ──
|
||||
orig_layer_fwd = deepseek_v4.DeepseekV4DecoderLayer.forward
|
||||
|
||||
@functools.wraps(orig_layer_fwd)
|
||||
def timed_layer_fwd(self, *args, **kwargs):
|
||||
# Determine prefill vs decode from forward_batch or hidden_states shape
|
||||
fb = kwargs.get("forward_batch", args[3] if len(args) > 3 else None)
|
||||
hs = kwargs.get("hidden_states", args[1] if len(args) > 1 else None)
|
||||
# Determine prefill vs decode from forward_batch flags
|
||||
# is_prefill_only: true if ALL reqs in batch are prefill
|
||||
# is_extend_in_batch: true if any req in batch has extend tokens
|
||||
# Fallback: check extend_num_tokens > 0
|
||||
is_prefill = False
|
||||
if fb is not None:
|
||||
if getattr(fb, "is_prefill_only", False):
|
||||
is_prefill = True
|
||||
elif getattr(fb, "is_extend_in_batch", False):
|
||||
is_prefill = True
|
||||
elif getattr(fb, "extend_num_tokens", None) is not None and fb.extend_num_tokens > 0:
|
||||
is_prefill = True
|
||||
_phase.mode = "prefill" if is_prefill else "decode"
|
||||
|
||||
_sync()
|
||||
t0 = time.time()
|
||||
try:
|
||||
return orig_layer_fwd(self, *args, **kwargs)
|
||||
finally:
|
||||
_sync()
|
||||
phase = _phase.mode
|
||||
dt = (time.time() - t0) * 1000
|
||||
_record("layer_ms", dt, phase=phase, extra={"layer": self.layer_id})
|
||||
|
||||
deepseek_v4.DeepseekV4DecoderLayer.forward = timed_layer_fwd
|
||||
|
||||
# ── HC Pre / Post ──
|
||||
for hc_name in ["hc_pre", "hc_post"]:
|
||||
fn = getattr(deepseek_v4.DeepseekV4DecoderLayer, hc_name, None)
|
||||
if fn is not None:
|
||||
setattr(deepseek_v4.DeepseekV4DecoderLayer, hc_name,
|
||||
_make_timed_fn(fn, f"{hc_name}_ms"))
|
||||
print(f"[P800_TIMING] Patched {hc_name}", flush=True)
|
||||
|
||||
# ── Attention forward (MQALayer) ──
|
||||
for name in ["MQALayer", "RadixAttention", "DeepseekV4Attention", "DeepseekV4FlashAttention"]:
|
||||
cls = getattr(deepseek_v4, name, None)
|
||||
if cls is not None and hasattr(cls, "forward"):
|
||||
cls.forward = _make_timed_fn(cls.forward, "attention_ms")
|
||||
print(f"[P800_TIMING] Patched attention: {name}", flush=True)
|
||||
break
|
||||
|
||||
# ── MQALayer._forward_prepare (QKV projection) ──
|
||||
cls = getattr(deepseek_v4, "MQALayer", None)
|
||||
if cls is not None and hasattr(cls, "_forward_prepare"):
|
||||
cls._forward_prepare = _make_timed_fn(cls._forward_prepare, "qkv_prepare_ms")
|
||||
print(f"[P800_TIMING] Patched MQALayer._forward_prepare", flush=True)
|
||||
|
||||
# ── RMSNorm (both input_layernorm and post_attention_layernorm) ──
|
||||
for rms_name in ["RMSNorm", "DeepseekRefRMSNorm"]:
|
||||
cls = getattr(deepseek_v4, rms_name, None)
|
||||
if cls is not None and hasattr(cls, "forward"):
|
||||
cls.forward = _make_timed_fn(cls.forward, "rmsnorm_ms")
|
||||
print(f"[P800_TIMING] Patched RMSNorm: {rms_name}", flush=True)
|
||||
break
|
||||
|
||||
# ── MLP/MoE (deepseek_v2.DeepseekV2MoE) ──
|
||||
try:
|
||||
from sglang.srt.models import deepseek_v2
|
||||
cls = deepseek_v2.DeepseekV2MoE
|
||||
if hasattr(cls, "forward"):
|
||||
cls.forward = _make_timed_fn(cls.forward, "mlp_ms")
|
||||
print(f"[P800_TIMING] Patched MoE: DeepseekV2MoE", flush=True)
|
||||
moe_patched = True
|
||||
|
||||
# ── MoE sub-operators ──
|
||||
# Gate
|
||||
if hasattr(cls, "gate") and hasattr(cls.gate, "forward"):
|
||||
cls.gate.forward = _make_timed_fn(cls.gate.forward, "moe_gate_ms")
|
||||
print(f"[P800_TIMING] Patched MoE gate", flush=True)
|
||||
|
||||
# Shared experts
|
||||
if hasattr(cls, "_forward_shared_experts"):
|
||||
orig_shared = cls._forward_shared_experts
|
||||
cls._forward_shared_experts = _make_timed_fn(orig_shared, "moe_shared_ms")
|
||||
print(f"[P800_TIMING] Patched MoE shared_experts", flush=True)
|
||||
|
||||
# forward_normal (routed experts kernel)
|
||||
if hasattr(cls, "forward_normal"):
|
||||
# The actual name might be klx_forward_normal
|
||||
for fn_name in ["forward_normal", "klx_forward_normal"]:
|
||||
fn = getattr(cls, fn_name, None)
|
||||
if fn is not None:
|
||||
setattr(cls, fn_name, _make_timed_fn(fn, "moe_routed_ms"))
|
||||
print(f"[P800_TIMING] Patched MoE {fn_name}", flush=True)
|
||||
break
|
||||
elif hasattr(cls, "forward_mega_moe"):
|
||||
cls.forward_mega_moe = _make_timed_fn(cls.forward_mega_moe, "moe_routed_ms")
|
||||
print(f"[P800_TIMING] Patched MoE forward_mega_moe", flush=True)
|
||||
except Exception as e:
|
||||
print(f"[P800_TIMING] MoE patch error: {e}", flush=True)
|
||||
moe_patched = False
|
||||
|
||||
if not moe_patched:
|
||||
for name in ["FusedMoE", "DeepseekV4MoE", "MoE"]:
|
||||
cls = getattr(deepseek_v4, name, None)
|
||||
if cls is not None and hasattr(cls, "forward"):
|
||||
cls.forward = _make_timed_fn(cls.forward, "mlp_ms")
|
||||
print(f"[P800_TIMING] Patched MoE: {name} (fallback)", flush=True)
|
||||
break
|
||||
|
||||
# ── All-reduce ──
|
||||
try:
|
||||
from sglang.srt.layers.communicator import Communicator
|
||||
if hasattr(Communicator, "all_reduce"):
|
||||
Communicator.all_reduce = _make_timed_fn(Communicator.all_reduce, "all_reduce_ms")
|
||||
print(f"[P800_TIMING] Patched all_reduce", flush=True)
|
||||
except Exception as e:
|
||||
print(f"[P800_TIMING] all_reduce patch failed: {e}", flush=True)
|
||||
|
||||
print("[P800_TIMING] Patches applied: layer, hc_pre, hc_post, attention, "
|
||||
"qkv_prepare, rmsnorm, moe, gate, shared, routed", flush=True)
|
||||
return True
|
||||
|
||||
|
||||
def _wait_and_patch():
|
||||
"""Poll for deepseek_v4 module and apply patches."""
|
||||
import sys
|
||||
import threading as _thr
|
||||
enabled = os.environ.get("SGLANG_TIMING_ENABLED", "0")
|
||||
if enabled != "1":
|
||||
return
|
||||
|
||||
def _poll():
|
||||
while True:
|
||||
if 'sglang.srt.models.deepseek_v4' in sys.modules:
|
||||
patch()
|
||||
return
|
||||
_thr.Event().wait(0.5)
|
||||
|
||||
_thr.Thread(target=_poll, daemon=True).start()
|
||||
|
||||
|
||||
_wait_and_patch()
|
||||
@ -5,3 +5,6 @@ import torch
|
||||
# them on P800 via its own FP8 kernels.
|
||||
if not hasattr(torch, "float8_e8m0fnu"):
|
||||
torch.float8_e8m0fnu = torch.uint8
|
||||
|
||||
# P800 Timing module
|
||||
import p800_timing
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user