feat(p800): add manual timing module to bypass broken PyTorch Profiler

PyTorch Profiler crashes on P800 XPU due to a CUPTI bug in
cuptiActivityDisable() that SIGKILLs the server process. This
module provides an alternative by monkey-patching SGLang's
DeepSeek-V4 decoder layer and its sub-operators to record
per-op wall-clock timing.

What it does:
- Patches DeepseekV4DecoderLayer.forward and sub-operators
  (hc_pre/hc_post, RMSNorm, MQALayer attention, QKV projection,
  DeepseekV2MoE and its gate/shared/routed experts) with timed
  wrappers that call torch.cuda.synchronize() before and after.
- Distinguishes prefill vs decode via forward_batch flags
  (is_prefill_only / is_extend_in_batch / extend_num_tokens)
  and propagates the phase to sub-operators via threading.local.
- Writes per-PID JSON summaries (count/avg/p50/p95/p99) every
  5000 records plus on atexit/SIGTERM, so data survives when
  the server is stopped.

Usage: set SGLANG_TIMING_ENABLED=1 in the container environment.
Output: /tmp/p800_timing_results_{pid}.json

Known limitations:
- torch.cuda.synchronize() adds overhead, inflating small-op times.
- moe_routed_ms only covers forward_normal (65% of calls); the
  dual_stream path taken during CUDA graph capture is not patched.
- All-reduce patch fails (Communicator import path mismatch).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

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