!11 feat(p800): enable operator-level timing in profiling experiment

Merge pull request !11 from yy-fighting/auto/main/13188872/d44301a2-1

experiments/p800/dsv4_p800_sglang_profile/start_sglang_docker.sh

@@ -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"

platforms/patches/kunlun_p800/p800_timing.py

@@ -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()

platforms/patches/kunlun_p800/sitecustomize_xpu.py

@@ -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