adding profiling context

atom/model_engine/engine_core.py

@@ -89,7 +89,6 @@ def __init__(self, config: Config, input_address: str, output_address: str):
         )
         self.input_thread.start()
 
-        self.profile_enbaled = config.torch_profiler_dir is not None
         self.mark_trace = getattr(config, "mark_trace", False)
         init_exit_handler(self)
         self._init_data_parallel(config)
@@ -128,20 +127,12 @@ def __init__(self, config: Config, input_address: str, output_address: str):
 
             config.num_kvcache_blocks = num_blocks
             if not config.enforce_eager:
-                # Start profiler before cudagraph capture only if mark-trace is enabled.
-                if self.profile_enbaled and self.mark_trace:
-                    self.runner_mgr.call_func(
-                        "start_profiler", "capture_graph", wait_out=True
-                    )
                 cap_cost, bs = self.runner_mgr.call_func(
                     "capture_cudagraph", wait_out=True
                 )
                 logger.info(
                     f"{self.label}: cudagraph capture{bs} cost: {cap_cost:.2f} seconds"
                 )
-                if self.profile_enbaled and self.mark_trace:
-                    # Persist a dedicated capture-graph trace immediately.
-                    self.runner_mgr.call_func("stop_profiler", wait_out=True)
             good = True
         finally:
             logger.info(
@@ -308,6 +299,7 @@ def _process_engine_step_inner(self):
         # Run the model forward pass if there are actual sequences
         has_seqs = len(scheduled_batch.req_ids) > 0
         if has_seqs:
+            self.scheduler.compute_roofline_aggregates(scheduled_batch, seqs)
             fwd_out = self.runner_mgr.call_func(
                 "forward", scheduled_batch, wait_out=True
             )

atom/model_engine/engine_utility.py

@@ -247,6 +247,10 @@ def _handle_configure_hidden_states(self, args: dict):
 
     def _handle_start_profile(self, args: dict):
         result = self.runner_mgr.call_func("start_profiler", wait_out=True)
+        # Flip the scheduler flag so per-iteration roofline aggregates
+        # (compute_roofline_aggregates) are emitted while profiling is active.
+        if self.scheduler is not None:
+            self.scheduler.profile_active = True
         logger.info(f"{self.label}: profiler started")
         self.output_queue.put_nowait(
             ("UTILITY_RESPONSE", {"cmd": "start_profile", "result": result})
@@ -255,6 +259,8 @@ def _handle_start_profile(self, args: dict):
     def _handle_stop_profile(self, args: dict):
         logger.info(f"{self.label}: stopping profiler...")
         result = self.runner_mgr.call_func("stop_profiler", wait_out=True)
+        if self.scheduler is not None:
+            self.scheduler.profile_active = False
         logger.info(f"{self.label}: profiler stopped, result={result}")
         self.output_queue.put_nowait(
             ("UTILITY_RESPONSE", {"cmd": "stop_profile", "result": result})

atom/model_engine/model_runner.py

@@ -2042,6 +2042,26 @@ def prepare_model(self, batch: ScheduledBatch):
             needs_independent_noise,
         )
 
+    @staticmethod
+    def _roofline_label_suffix(batch: Optional[ScheduledBatch]) -> str:
+        """Roofline FLOP aggregates for the trace label, or ``""``.
+
+        These fields are only populated by
+        :meth:`Scheduler.compute_roofline_aggregates` when profiling is active
+        and ``ATOM_ENABLE_DETAILED_ANNOTATION`` is set, so on the normal
+        (unprofiled) path this returns an empty string without any extra work.
+        Appending here keeps the annotation on the ``prefill[]``/``decode[]``
+        ``record_function`` (a GPU-recognized layer) instead of nesting an
+        extra span above ``run_model``.
+        """
+        if batch is None or batch.roofline_sqsq is None:
+            return ""
+        return (
+            f" sqsq={batch.roofline_sqsq}"
+            f" sqsk={batch.roofline_sqsk}"
+            f" sk={batch.roofline_sk}"
+        )
+
     def run_model(
         self,
         input_ids: torch.Tensor,
@@ -2077,6 +2097,7 @@ def run_model(
             tbo_on=forward_context.ubatch_slices is not None,
             bs=bs,
             batch=batch,
+            roofline_suffix=self._roofline_label_suffix(batch),
         )
 
         if not forward_mode.use_cudagraph:
@@ -2333,6 +2354,7 @@ def forward(self, batch: ScheduledBatch) -> ScheduledBatchOutput:
             needs_independent_noise=needs_independent_noise,
         )
         reset_forward_context()
+
         return fwd_output
 
     @torch.inference_mode()
@@ -2397,6 +2419,60 @@ def propose_draft_token_ids(
         )
         return self.tokenID_processor.prepare_draft_ids(batch, draft_token)
 
+    def start_capture_profiler(self):
+        """Set up the per-bs CUDA graph capture profiler (profiles in place).
+
+        Profiles the capture phase as graphs are captured and writes one trace
+        per batch size, per rank (``bs_<bs>_rank<rank>.json.gz``). Enabled on
+        every rank when a torch profiler dir is set and mark-trace is on.
+        """
+        self._capture_profile_enabled = (
+            self.profiler_dir is not None and self.mark_trace
+        )
+        if self._capture_profile_enabled:
+            self._profile_bs_idx = 0
+            self.capture_traces_dir = os.path.join(self.profiler_dir, "capture_traces")
+            os.makedirs(self.capture_traces_dir, exist_ok=True)
+            logger.info(f"{self.label}: Starting CUDA graph capture profiler...")
+
+            def on_trace_ready(prof):
+                # Invariant: exactly two prof.step() calls happen per captured
+                # batch size (schedule wait=1 + active=1, repeat=0), so
+                # on_trace_ready fires once per bs, in self.graph_bs order.
+                # This is a profiling-only diagnostic; log-and-skip rather than
+                # assert so a cadence mismatch can never abort CUDA-graph
+                # capture at server startup (and isn't stripped under python -O).
+                if self._profile_bs_idx >= len(self.graph_bs):
+                    logger.warning(
+                        "capture profiler fired %d times but only %d batch "
+                        "sizes were captured; skipping extra trace. Check the "
+                        "prof.step() cadence in capture_cudagraph.",
+                        self._profile_bs_idx + 1,
+                        len(self.graph_bs),
+                    )
+                    return
+                bs = self.graph_bs[self._profile_bs_idx]
+                trace_file = os.path.join(
+                    self.capture_traces_dir, f"bs_{bs}_rank{self.rank}.json.gz"
+                )
+                prof.export_chrome_trace(trace_file)
+                logger.info(f"Saved trace for bs={bs} to {trace_file}")
+                self._profile_bs_idx += 1
+
+            self.capture_profiler = torch_profiler.profile(
+                activities=[
+                    torch_profiler.ProfilerActivity.CUDA,
+                    torch_profiler.ProfilerActivity.CPU,
+                ],
+                schedule=torch_profiler.schedule(wait=1, warmup=0, active=1, repeat=0),
+                record_shapes=True,
+                with_stack=True,
+                profile_memory=False,
+                on_trace_ready=on_trace_ready,
+            )
+        else:
+            self.capture_profiler = nullcontext()
+
     @torch.inference_mode()
     def _piecewise_cg_active(self) -> bool:
         """True when the compiled model's dense pieces self-capture PIECEWISE
@@ -2511,6 +2587,9 @@ def capture_cudagraph(self):
         # TBO graphs don't capture compute_logits, so disable logits_in_graph.
         self.logits_in_graph = self.world_size == 1 and not is_tbo
 
+        # start capture profiler
+        self.start_capture_profiler()
+
         @contextmanager
         def pause_gc():
             # No GC during capture: a finalizer's hipModuleUnload aborts it (HIP 900).
@@ -2523,7 +2602,7 @@ def pause_gc():
                 gc.collect()
 
         _rsv_before_capture = torch.cuda.memory_reserved()
-        with pause_gc(), graph_capture() as capture_ctx:
+        with pause_gc(), graph_capture() as capture_ctx, self.capture_profiler as prof:
             capture_range = (
                 tqdm.tqdm(self.graph_bs) if self.rank == 0 else self.graph_bs
             )
@@ -2611,6 +2690,8 @@ def pause_gc():
                     outputs[:num_tokens] = model_output
                 if self.logits_in_graph:
                     self.model.compute_logits(outputs[:num_tokens])
+                if prof is not None:
+                    prof.step()
 
                 if _piecewise:
                     fc = get_forward_context()
@@ -2672,6 +2753,8 @@ def pause_gc():
                 self.graphs[(bs, max_q_len)] = graph
                 if self.logits_in_graph and ubatch_slices is None:
                     self.graph_logits[(bs, max_q_len)] = graph_logits
+                if prof is not None:
+                    prof.step()
                 if graph_aux is not None:
                     self.graph_aux_hidden[(bs, max_q_len)] = graph_aux
                 torch.cuda.synchronize()

atom/model_engine/run_labels.py

@@ -40,10 +40,15 @@ def build_run_label(
     tbo_on: bool,
     bs: int,
     batch: Optional["ScheduledBatch"],
+    roofline_suffix: str = "",
 ) -> str:
     """Build the ``record_function`` label for one forward pass.
 
     Pure function (no runtime state) — see module docstring for the taxonomy.
+
+    ``roofline_suffix`` is an already-formatted ``" key=value ..."`` fragment
+    (empty on the normal path) appended inside the brackets; see
+    ``ModelRunner._roofline_label_suffix``.
     """
     if use_cudagraph:
         kind = "dummy_decode" if is_dummy else "decode"
@@ -70,6 +75,7 @@ def build_run_label(
             else:
                 ctx_str = f"{ctx[:3].tolist()}...+{len(ctx) - 3}"
             label += f" tok={batch.total_tokens_num} ctx={ctx_str}"
+    label += roofline_suffix
     if tbo_on:
         label += " tbo=1"
     label += "]"

atom/model_engine/scheduler.py

@@ -29,6 +29,7 @@
 from atom.model_engine.block_manager import BlockManager
 from atom.model_engine.request import RequestOutput
 from atom.model_engine.sequence import Sequence, SequenceStatus, SequenceType
+from atom.utils import envs
 
 logger = logging.getLogger("atom")
 
@@ -343,6 +344,14 @@ def __init__(
         self.is_dummy_run = is_dummy_run
         self.num_spec_step = num_spec_step
 
+        # Roofline FLOP aggregates (set by Scheduler.compute_roofline_aggregates
+        # when profiling is active and ATOM_ENABLE_DETAILED_ANNOTATION is set).
+        # None on the normal path; consumed by ModelRunner.run_model to extend
+        # the prefill[]/decode[] trace labels.
+        self.roofline_sqsq: int | None = None  # sum N_Q^2
+        self.roofline_sqsk: int | None = None  # sum N_Q*N_KV
+        self.roofline_sk: int | None = None  # sum N_KV
+
         # Collect multimodal data from prefill sequences
         self.multimodal_data = {}
         for seq in seqs.values():
@@ -459,6 +468,11 @@ def __init__(self, config: Config):
         self.cache_stats: Optional[CacheStats] = (
             CacheStats() if config.enable_prefix_caching else None
         )
+        self.profile_active = False
+        # Cache the env flag once (env vars are fixed at process start) so the
+        # per-iteration compute_roofline_aggregates never pays an os.getenv.
+        self._detailed_annotation_enabled = envs.ATOM_ENABLE_DETAILED_ANNOTATION
+
         self.enable_chunked_prefill = config.enable_chunked_prefill
         # V4 SWA correctness on a prefix-cache hit is now handled entirely in
         # BlockManager: `_swa_bounded_hit` bounds the hit so the boundary's
@@ -1707,6 +1721,61 @@ def postprocess(
             self.running.remove(seq)
         return finished_seqs
 
+    def compute_roofline_aggregates(
+        self,
+        scheduled_batch: ScheduledBatch,
+        seqs: dict[int, Sequence],
+    ) -> None:
+        """Attach roofline FLOP aggregates to *scheduled_batch* in place.
+
+        Only the quadratic terms genuinely needed for a roofline FLOP estimate
+        are computed here. The request counts and total query tokens are
+        already emitted by the ``prefill[]``/``decode[]`` labels in
+        :meth:`ModelRunner.run_model`, so this avoids duplicating them.
+
+        The following batch-level sums are stored on the batch and appended to
+        those labels by the runner:
+
+            sqsq  — sum of N_Q^2      (per request)
+            sqsk  — sum of N_Q*N_KV   (per request)
+            sk    — sum of N_KV       (per request)
+
+        where ``N_Q`` is the number of query tokens scheduled for a request and
+        ``N_KV`` is its KV length (cached + new tokens for prefill, full
+        sequence length for decode). Aggregating over every request in the
+        batch gives the total for that single forward, which is exactly the
+        quantity a per-iteration roofline point needs. For MTP/spec-decode a
+        decode step schedules ``mtp_k + 1`` query tokens, so the scheduled
+        token count is used as ``N_Q`` for both branches (rather than a
+        hardcoded 1) to avoid undercounting.
+
+        This is a no-op (leaves the fields ``None``) unless profiling is active
+        and ``ATOM_ENABLE_DETAILED_ANNOTATION`` is set.
+        """
+        if not self.profile_active or not self._detailed_annotation_enabled:
+            return
+
+        sqsq = 0  # sum N_Q^2
+        sqsk = 0  # sum N_Q*N_KV
+        sk = 0  # sum N_KV
+        for seq, num_tokens in zip(seqs.values(), scheduled_batch.num_scheduled_tokens):
+            # Cast to Python int: num_scheduled_tokens is np.int32, so nq*nq /
+            # nq*nkv would overflow once a prefill/chunk exceeds ~46341 tokens
+            # (e.g. np.int32(65536)**2 == 0), silently corrupting the estimate.
+            nq = int(num_tokens)  # query tokens scheduled this forward
+            if seq.type == SequenceType.DECODE:
+                nkv = int(seq.num_tokens)  # full sequence length
+            else:
+                # PREFILL: KV length = cached + new tokens.
+                nkv = int(seq.num_cached_tokens) + nq
+            sqsq += nq * nq
+            sqsk += nq * nkv
+            sk += nkv
+
+        scheduled_batch.roofline_sqsq = sqsq
+        scheduled_batch.roofline_sqsk = sqsk
+        scheduled_batch.roofline_sk = sk
+
     def _connector_flag(self, name: str) -> bool:
         return bool(getattr(self.kv_connector, name, False))
 

atom/utils/envs.py

@@ -130,6 +130,11 @@
     # --- Profiling & Logging ---
     "ATOM_TORCH_PROFILER_DIR": lambda: os.getenv("ATOM_TORCH_PROFILER_DIR", None),
     "ATOM_PROFILER_MORE": lambda: os.getenv("ATOM_PROFILER_MORE", "0") == "1",
+    # When profiling is active, append roofline FLOP aggregates (sqsq, sqsk, sk)
+    # to the prefill[]/decode[] trace labels emitted by ModelRunner.run_model.
+    "ATOM_ENABLE_DETAILED_ANNOTATION": lambda: (
+        os.getenv("ATOM_ENABLE_DETAILED_ANNOTATION", "0") == "1"
+    ),
     "ATOM_PROFILER_TIMEOUT": lambda: float(os.getenv("ATOM_PROFILER_TIMEOUT", "300")),
     "ATOM_LOG_MORE": lambda: int(os.getenv("ATOM_LOG_MORE", "0")) != 0,
     # RTL (rocm-trace-lite) GPU kernel tracing — set to output directory to enable.

docs/environment_variables.md

@@ -125,6 +125,7 @@ land. See `atom/model_ops/v4_backend_gate.py` for the selector.
 |----------|------|---------|-------------|
 | **ATOM_TORCH_PROFILER_DIR** | str | — | When set, enables PyTorch profiler and writes traces to this directory. Create subdirectories per rank (e.g., `rank_0`, `dp0_tp0`). |
 | **ATOM_PROFILER_MORE** | bool | 0 (false) | When `ATOM_TORCH_PROFILER_DIR` is set and this is `1`, enables detailed profiling: `record_shapes`, `with_stack`, and `profile_memory`. |
+| **ATOM_ENABLE_DETAILED_ANNOTATION** | bool | 0 (false) | When profiling is active, appends roofline FLOP aggregates to the `prefill[]`/`decode[]` trace labels: `sqsq` (Σ N_Q²), `sqsk` (Σ N_Q·N_KV), and `sk` (Σ N_KV), where N_Q is the scheduled query tokens and N_KV the KV length per request. Used to estimate attention FLOPs for roofline analysis. |
 | **ATOM_LOG_MORE** | bool | 0 (false) | If set to `1`, use verbose logging format (includes process name, PID, path, line number, function name). |
 
 ### Debug Dump (`atom.utils.debug_helper`)

docs/serving_benchmarking_guide.md

@@ -362,6 +362,7 @@ the programmatic API.
 | `ATOM_TORCH_PROFILER_DIR` env var | Sets the default `torch_profiler_dir` in `Config` |
 | `ATOM_PROFILER_MORE=1` env var | Enables detailed profiling: `record_shapes`, `with_stack`, `profile_memory` |
 | `ATOM_PROFILER_TIMEOUT=<seconds>` env var | Overrides the `stop_profile` timeout; default is 300 seconds |
+| `ATOM_ENABLE_DETAILED_ANNOTATION=1` env var | Appends attention FLOP aggregates (`sqsq`, `sqsk`, `sk`) to the `prefill[]`/`decode[]` trace labels while profiling is active (see [5.6](#56-cuda-graph-capture-traces)) |
 
 When a profiler directory is configured, each worker saves traces to a
 rank-specific subdirectory:
@@ -443,6 +444,32 @@ python -m atom.benchmarks.benchmark_serving \
 This sends `POST /start_profile` before the benchmark and
 `POST /stop_profile` after completion.
 
+### 5.6 CUDA-Graph Capture Traces
+
+During CUDA-graph capture (server bring-up), ATOM can emit one trace file per
+captured batch size instead of a single combined blob. This makes each graph's
+capture cost easy to inspect in isolation and keeps individual trace files
+small. Capture-trace profiling is gated on `--mark-trace` (with
+`--torch-profiler-dir`/`ATOM_TORCH_PROFILER_DIR` set); the warmup phase is
+omitted so only meaningful capture work is retained.
+
+The per-batch-size traces are written to:
+
+```
+{profiler_dir}/capture_traces/bs_<bs>_rank<rank>.json.gz
+```
+
+where `<bs>` is the captured batch size and `<rank>` the worker rank. Each file
+is a gzip-compressed Chrome trace viewable with `chrome://tracing` or
+TensorBoard.
+
+To additionally annotate the run-phase `prefill[]`/`decode[]` labels with the
+attention FLOP aggregates used for roofline analysis, set
+`ATOM_ENABLE_DETAILED_ANNOTATION=1` (see [5.1](#51-configuration)). The added
+fields are `sqsq` (Σ N_Q²), `sqsk` (Σ N_Q·N_KV), and `sk` (Σ N_KV), summed over
+every request in the forward. These are attention-quadratic terms only — a full
+roofline still requires GEMM FLOPs and bytes moved.
+
 ---
 
 ## 6. Speculative Decoding (MTP)