feat: add sglang_trace_analyze task for GPU profiler trace analysis

metainfer/tasks/sglang_trace_analyze/__init__.py

@@ -0,0 +1,8 @@
+"""sglang_trace_analyze — auto-generate torch profiler traces via SGLang,
+analyze them (operator-to-structure mapping, kernel hotspots, TFLOPS / MFU,
+overlap opportunities, fuse suggestions), and surface results + LLM hints
+in the MetaInfer WebUI.
+"""
+
+from .orchestrator import plugin as _task_plugin  # noqa: F401
+from .server import plugin as _web_plugin          # noqa: F401

metainfer/tasks/sglang_trace_analyze/form.yaml

@@ -0,0 +1,67 @@
+- key: model_path
+  header: Model Path
+  question: "HuggingFace repo id or local path to the model weights."
+  required: true
+  form: text
+
+- key: version
+  header: Version
+  question: "Short tag for this run — used in trace directory naming and result labels."
+  required: true
+  form: text
+
+- key: batch_sizes
+  header: Batch Sizes
+  question: "Comma-separated list of decode batch sizes to profile, e.g. 1,4,8,16."
+  required: true
+  form: text
+
+- key: mapping_batch_size
+  header: Mapping BS
+  question: "Batch size for the mapping run (CUDA Graph disabled). One value is enough — kernel-to-layer mapping is independent of batch size."
+  required: true
+  default: "8"
+  form: number
+
+- key: input_len
+  header: Input Len
+  question: "Synthetic input sequence length."
+  required: true
+  default: "512"
+  form: number
+
+- key: output_len
+  header: Output Len
+  question: "Synthetic output sequence length."
+  required: true
+  default: "2000"
+  form: number
+
+- key: tp_size
+  header: TP Size
+  question: "Tensor-parallelism degree."
+  required: true
+  default: "1"
+  form: number
+
+- key: pp_size
+  header: PP Size
+  question: "Pipeline-parallelism degree."
+  required: true
+  default: "1"
+  form: number
+
+- key: gpu_model
+  header: GPU Model
+  question: "GPU model — used to look up theoretical peak TFLOPS and memory bandwidth."
+  required: true
+  form: select
+  options:
+    - label: "K100"
+      description: "FP32 49TF, TF32 98TF, BF16/FP16 192TF, INT8 392TOPS, BW 700GB/s"
+    - label: "A100_80G"
+      description: "FP32 19.5TF, TF32 156TF, BF16/FP16 312TF, INT8 624TOPS, BW 2039GB/s"
+    - label: "H100"
+      description: "FP32 67TF, TF32 989TF, BF16/FP16 989TF, INT8 1979TOPS, BW 3350GB/s"
+    - label: "B200"
+      description: "FP32 90TF, TF32 2250TF, BF16/FP16 2250TF, INT8 4500TOPS, BW 8000GB/s"

metainfer/tasks/sglang_trace_analyze/orchestrator/__init__.py

@@ -0,0 +1,6 @@
+"""Orchestrator (worker subprocess) for sglang_trace_analyze."""
+
+from metainfer.orchestrator.tasks import register
+from .plugin import PLUGIN
+
+register(PLUGIN)

metainfer/tasks/sglang_trace_analyze/orchestrator/cli.py

@@ -0,0 +1,46 @@
+"""CLI entry point for the sglang_trace_analyze orchestrator subprocess.
+
+The launcher spawns::
+
+    python -m <cli_module> run <requirements.json> --state-dir … --workspace-dir …
+
+Contract required by the framework (§6d): ``run`` subcommand + ``--state-dir``
+and ``--workspace-dir`` flags.
+"""
+
+from __future__ import annotations
+
+import argparse
+import sys
+from pathlib import Path
+
+
+def main(argv: list[str] | None = None) -> int:
+    parser = argparse.ArgumentParser(prog="metainfer-orchestrator")
+    sub = parser.add_subparsers(dest="command")
+
+    run_p = sub.add_parser("run")
+    run_p.add_argument("requirements", type=Path,
+                       help="Path to requirements.json")
+    run_p.add_argument("--state-dir", type=Path, required=True)
+    run_p.add_argument("--workspace-dir", type=Path, required=True)
+    # Task-specific flags
+    run_p.add_argument("--iter-limit", type=int, default=None,
+                       help="Override max iterations (default: derive from batch count)")
+
+    args = parser.parse_args(argv)
+    if args.command != "run":
+        parser.print_help()
+        return 1
+
+    from .orchestrator import run_with_requirements
+    return run_with_requirements(
+        requirements_path=args.requirements,
+        state_dir=args.state_dir,
+        workspace_dir=args.workspace_dir,
+        iter_limit=args.iter_limit,
+    )
+
+
+if __name__ == "__main__":
+    sys.exit(main())

metainfer/tasks/sglang_trace_analyze/orchestrator/flops_calculator.py

@@ -0,0 +1,195 @@
+"""Compute TFLOPS, bandwidth, and MFU for aggregated kernel entries.
+
+Uses:
+- ``gpu_specs.py`` for theoretical peak values
+- kernel ``input_dims`` (from MAPPING trace) or shape rules (for CUDA Graph
+  formal traces) to derive actual FLOP counts per invocation
+- kernel ``total_dur_us`` to compute actual TFLOPS/bandwidth
+"""
+
+from __future__ import annotations
+
+from typing import Any, Dict, List, Optional
+
+from .gpu_specs import GpuSpec
+
+
+def extract_ck_tile_dims(kernel_name: str) -> tuple | None:
+    """Extract (M, N, K) tile dimensions from a CK GEMM kernel name.
+
+    Example: ``Cijk_Alik_Bljk_SB_MT64x128x16_...`` → (64, 128, 16)
+    """
+    import re
+    m = re.search(r"MT(\d+)x(\d+)x(\d+)", kernel_name)
+    if m:
+        return int(m.group(1)), int(m.group(2)), int(m.group(3))
+    return None
+
+
+def calculate_mfu(
+    kernels: List[Dict[str, Any]],
+    gpu_spec: GpuSpec,
+    *,
+    batch_size: int,
+    dtype: str = "bf16",
+) -> List[Dict[str, Any]]:
+    """Augment each kernel entry with TFLOPS, bandwidth, MFU, and bound classification.
+
+    Args:
+        kernels: Aggregated kernel list. Each entry must have ``total_dur_us``
+            and ``count``. Entries from a non-CUDA Graph trace may also have
+            ``input_dims``, which are used for FLOP/byte estimation where available.
+        gpu_spec: GPU theoretical peak specification.
+        batch_size: Decode batch size used for this trace.
+        dtype: Compute dtype — determines which TFLOPS peak to use.
+            One of ``fp32``, ``tf32``, ``bf16``, ``fp16``, ``int8``.
+
+    Returns:
+        The same kernel list with added fields: ``tflops_actual``,
+        ``bandwidth_gb_s``, ``mfu``, ``bound``, ``flops_per_invocation``.
+    """
+    theoretical_tflops = _theoretical_peak(gpu_spec, dtype)
+    theoretical_bw = gpu_spec.bandwidth_gb_s
+
+    for k in kernels:
+        dur_s = k["total_dur_us"] / 1e6
+        count = k.get("count", 1)
+        dur_per_invocation_s = dur_s / count if count else dur_s
+        dims = k.get("input_dims", [])
+        op_type = k.get("op_type", "Other")
+        kernel_name = k.get("kernel_name", "")
+
+        flops = _estimate_flops(op_type, dims, batch_size)
+        bytes_moved = _estimate_bytes(op_type, dims, batch_size)
+
+        # For CK GEMM kernels without input dims, estimate from tile name
+        if flops == 0 and op_type == "GEMM":
+            tile = extract_ck_tile_dims(kernel_name)
+            if tile:
+                M, N, K_tile = tile
+                flops = 2 * M * N * K_tile * count
+                bytes_moved = (M * K_tile + K_tile * N + M * N) * 2 * count
+
+        tflops_actual = (flops / dur_s / 1e12) if dur_s > 0 else 0
+        bandwidth_gb_s = (bytes_moved / dur_s / 1e9) if dur_s > 0 else 0
+        mfu = (tflops_actual / theoretical_tflops * 100) if theoretical_tflops > 0 else 0
+
+        # Compute-bound vs memory-bound heuristic
+        ops_per_byte = flops / bytes_moved if bytes_moved > 0 else float("inf")
+        if theoretical_bw > 0:
+            crossover = theoretical_tflops * 1e12 / (theoretical_bw * 1e9)
+        else:
+            crossover = float("inf")
+        bound = "compute" if ops_per_byte > crossover else "memory"
+
+        k["tflops_actual"] = round(tflops_actual, 6) if flops > 0 else None
+        k["tflops_theoretical"] = theoretical_tflops
+        k["bandwidth_gb_s"] = round(bandwidth_gb_s, 1) if bytes_moved > 0 else None
+        k["bandwidth_theoretical"] = theoretical_bw
+        k["mfu"] = round(mfu, 3) if flops > 0 else None
+        k["bound"] = bound if (flops > 0 and bytes_moved > 0) else "unknown"
+        k["flops_per_invocation"] = int(flops)
+
+    return kernels
+
+
+def _theoretical_peak(spec: GpuSpec, dtype: str) -> float:
+    """Return theoretical peak TFLOPS for the given dtype."""
+    return {
+        "fp32": spec.fp32_tflops,
+        "tf32": spec.tf32_tflops,
+        "bf16": spec.bf16_tflops,
+        "fp16": spec.fp16_tflops,
+        "int8": spec.int8_tops,  # TOPS → TFLOPS approximate
+    }.get(dtype, spec.bf16_tflops)
+
+
+def _estimate_flops(
+    op_type: str,
+    dims: List[Any],
+    batch_size: int,
+) -> float:
+    """Estimate FLOPs for one kernel invocation.
+
+    For GEMM: 2*M*N*K (or 2*B*M*N*K for batched).
+    For Attention: approximately 4*B*seq_len*head_dim*num_heads^2.
+    For ElementWise: 2*num_elements.
+
+    Returns 0 if dims are unavailable (CUDA Graph trace).
+    """
+    if not dims:
+        return 0
+
+    # Use the first observed dim list
+    d = dims[0]
+
+    if op_type == "GEMM":
+        if isinstance(d, list) and len(d) >= 2:
+            if len(d) == 3:
+                M, K, N = int(d[0]), int(d[1]), int(d[2])
+                return 2 * M * N * K
+            B, M, N, K = _unpack_4d(d, batch_size)
+            return 2 * B * M * N * K
+
+    elif op_type == "Attention":
+        if isinstance(d, list) and len(d) >= 3:
+            seq_len = int(d[0])
+            num_heads = int(d[1])
+            head_dim = int(d[2])
+            return 4 * seq_len * head_dim * num_heads * num_heads * batch_size
+
+    elif op_type == "MoE":
+        if isinstance(d, list) and len(d) >= 3:
+            M, K, N = int(d[0]), int(d[1]), int(d[2])
+            return 2 * M * N * K
+
+    return 0
+
+
+def _estimate_bytes(
+    op_type: str,
+    dims: List[Any],
+    batch_size: int,
+) -> float:
+    """Estimate bytes moved (reads + writes) for one kernel invocation.
+
+    Simple heuristic: for GEMM, input_bytes ≈ (M*K + K*N) * dtype_size,
+    output_bytes ≈ M*N * dtype_size. For elementwise, ≈ 3 * num_elements.
+
+    Returns 0 if dims are unavailable.
+    """
+    if not dims:
+        return 0
+
+    d = dims[0]
+    dtype_size = 2  # bf16/fp16 default
+
+    if op_type == "GEMM":
+        if isinstance(d, list):
+            if len(d) == 3:
+                M, K, N = int(d[0]), int(d[1]), int(d[2])
+                return (M * K + K * N + M * N) * dtype_size
+            B, M, N, K = _unpack_4d(d, batch_size)
+            return B * (M * K + K * N + M * N) * dtype_size
+
+    elif op_type == "Attention":
+        if isinstance(d, list) and len(d) >= 3:
+            seq_len = int(d[0])
+            num_heads = int(d[1])
+            head_dim = int(d[2])
+            # Q, K, V reads + output write (approximate)
+            return batch_size * seq_len * num_heads * head_dim * 4 * dtype_size
+
+    return 0
+
+
+def _unpack_4d(
+    dims: list,
+    batch_size: int,
+) -> tuple:
+    """Unpack a 4-element dim list into (B, M, N, K), defaulting B to batch_size."""
+    if len(dims) >= 4:
+        return int(dims[0]), int(dims[1]), int(dims[2]), int(dims[3])
+    if len(dims) == 3:
+        return batch_size, int(dims[0]), int(dims[1]), int(dims[2])
+    return batch_size, int(dims[0]), 1, 1