Operator Optimization Workflow

This workflow is for operator/kernel profiling, comparison, and ablation workloads such as fused MoE, RPA, attention, matmul, or normalization kernels. It keeps the Falcon workflow generic: the skill runs experiments, the workload writes a standard artifact layout, and the operator-analysis plugin summarizes the evidence.

Artifact Contract

Every operator optimization experiment should write this layout under the Falcon-provided artifact directory:

<artifact root>/
  manifest.json
  profiling/
    env.txt
    python-packages.txt
  rank-0/
    benchmark/
      metrics.jsonl
    profiling/
      xprof/
    compiler/
      hlo/
      llo/
  rank-1/
    benchmark/
      metrics.jsonl
    profiling/
      xprof/
    compiler/
      hlo/
      llo/

Keep the Falcon YAML thin. It should select hardware, set common profiling environment, establish rank-0 artifact ownership, and call the workload. Do not inline a full benchmark program into the manifest. Framework/operator-specific logic belongs in an upstream benchmark entrypoint, a small script, or a recipe snippet that can be reviewed and tested independently.

For TPU/JAX operator profiling, the workload should enable libtpu profiling metadata before running the benchmark command, then dump Pallas LLO into compiler/llo/ after compilation. The libtpu flags mirror the MaxText/ant-pretrain operator profiling workflow, where they make XProf show custom-call regions and LLO debug mapping:

OUT="${ARTIFACT_LOCAL_DIR:-/tmp/operator-artifact}"
mkdir -p "$OUT/compiler/llo"
LIBTPU_PROFILE_ARGS="--xla_enable_custom_call_region_trace=true --xla_xprof_register_llo_debug_info=true --xla_mosaic_dump_to=$OUT/compiler/llo"
export LIBTPU_INIT_ARGS="${LIBTPU_INIT_ARGS:+$LIBTPU_INIT_ARGS }$LIBTPU_PROFILE_ARGS"
export JAX_COMPILATION_CACHE_DIR="${JAX_COMPILATION_CACHE_DIR:-/tmp/jax_compilation_cache}"

Rank 0 writes root metadata: manifest.json, profiling/env.txt, and profiling/python-packages.txt. Raw observations are rank-scoped: metrics, LLO dumps, XProf traces, benchmark logs, and similar files go under rank-<n>/. A distributed benchmark may only write rank-0/benchmark/metrics.jsonl; that is valid. If multiple ranks write metrics, the analysis plugin decides how to aggregate or display them. A run without profiler output, custom-call region trace, or Pallas LLO dump can still be used as a benchmark datapoint, but it should not be used as full profiling evidence.

For distributed operator runs, all ranks may need to execute the benchmark, but only rank 0 should write root metadata. Every rank should write raw evidence only under its own rank-<n>/ directory. Do not let multiple ranks overwrite the same shared GCS FUSE artifact paths.

The agent-facing operator-profile.yaml template exports FALCON_OPERATOR_RANK, FALCON_OPERATOR_IS_LEADER, ROOT, and OUT for this pattern. OUT always points to $ROOT/rank-$FALCON_OPERATOR_RANK. Workload commands that pipe output through tee must explicitly preserve the benchmark process exit status so a failed benchmark cannot be hidden by a successful logging pipeline.

manifest.json follows schemas/operator_artifact_v1.json. The stable top-level classifier is:

{
  "schema_version": 1,
  "workflow": "operator-optimization",
  "operator_family": "fused_moe",
  "operator_name": "pallas_fused_moe",
  "hardware": {
    "device_type": "v7x",
    "device_count": 8,
    "device_topo": "2x2x1"
  },
  "dimensions": {
    "ep_size": 8,
    "top_k": 8,
    "num_experts": 256
  }
}

For RPA, the same contract uses a different family and dimensions:

{
  "schema_version": 1,
  "workflow": "operator-optimization",
  "operator_family": "rpa",
  "operator_name": "rpa_forward",
  "dimensions": {
    "block_size": 128,
    "num_heads": 32,
    "head_dim": 128,
    "causal": true
  }
}

rank-<n>/benchmark/metrics.jsonl is one JSON object per measurement. The plugin does not require a fixed set of dimensions; it groups rows by non-metric scalar fields and computes deltas for a chosen primary metric.

Example:

{"variant":"baseline","tokens":64,"dtype":"bf16","latency_ms":1.24,"tokens_per_sec":51612}
{"variant":"baseline","tokens":256,"dtype":"bf16","latency_ms":3.80,"tokens_per_sec":67368}

Analysis Plugin

Use operator-analysis for all operator families. It supports:

• Single artifact: summarize one run.
• Baseline plus candidate: compare each candidate against the baseline.
• Matrix/ablation: summarize all inputs and compute baseline deltas when a baseline role is present.

Single-artifact analysis:

schema_version: 1
exp_id: exp-xxx
spec:
  name: operator-analysis
  plugin_names: [operator-analysis]
  params:
    mode: single
    primary_metric: latency_ms

Multi-artifact analysis:

schema_version: 1
name: operator-ablation
owner_exp: exp-baseline
plugins:
  - operator-analysis
inputs:
  - name: baseline
    role: baseline
    exp_id: exp-baseline
  - name: candidate-a
    role: candidate
    exp_id: exp-candidate-a
  - name: candidate-b
    role: candidate
    exp_id: exp-candidate-b
config_refs:
  mode: ablation
  operator_family: fused_moe
  primary_metric: latency_ms

The plugin writes:

summary.json
operator-analysis.jsonl
report.md

For Pallas/Mosaic operators, run pallas-llo-analysis as a second, independent analysis against the same exp_id. Do not put operator-analysis and pallas-llo-analysis into one multi-plugin analysis: separate analysis records give each plugin its own params, result path, summary.json, and report.md.

schema_version: 1
exp_id: exp-xxx
spec:
  name: pallas-llo-analysis
  plugin_names: [pallas-llo-analysis]
  params:
    mode: single

Skill Behavior

The falcon-workflow skill should not create one bespoke workflow per operator. It should:

1. Select an operator scenario from the user prompt, then identify the upstream workload command before rendering YAML. For sglang-jax fused MoE microbenchmarks, that command is based on python -m benchmark.moe.bench_fused_moe or the upstream benchmark/moe/run_fused_moe_ablation.sh; for serving profiling it is sgl_jax.launch_server plus sgl_jax.bench_serving.
2. Render one or more generic operator profile manifests that write the standard artifact layout. The Falcon YAML is not operator-specific: only the workload command and manifest.json dimensions vary. For replica > 1, rank 0 writes root metadata and every rank writes raw evidence under rank-<n>/.
3. Keep workload commands short. They should call the benchmark and write raw outputs under OUT; distributed initialization, parameter matrices, and framework-specific validation should live in the benchmark command or a small script, not in a long YAML block.
4. For sglang-jax workloads, install the package without python[tpu], verify the selected TPU image's JAX/JAXLIB pair before importing jax, and upgrade only libtpu with --no-deps when needed. The upstream extra can pin a different jax[tpu] bundle and silently replace the image's known-good JAX/JAXLIB combination.
5. Run falcon workflow profile submit and profile collect.
6. For one artifact, create operator-analysis directly.
7. For Pallas/Mosaic operators, create a separate pallas-llo-analysis on the same exp_id.
8. For two or more artifacts, use falcon workflow multi-artifact create with operator-analysis.
9. Summarize each analysis's report.md / summary.json, not just Falcon metadata.

Family-specific logic belongs in small scenario adapters and future plugin family adapters. The workflow and artifact contract stay common across fused MoE, RPA, attention, and future operators.

The generic skill template is skills/falcon-workflow/references/operator-optimization/operator-profile.yaml. Do not add a new YAML file for every operator family unless the Falcon manifest schema itself needs a new capability.