feat(observability): add NVTX tracing and Prometheus metrics

csrc/ops.cpp

@@ -4,6 +4,8 @@
 #include <torch/extension.h>
 #include <cuda_bf16.h>
 
+#include "utils/nvtx_utils.h"
+
 // Fused LogP Declarations
 torch::Tensor fused_logp_forward(torch::Tensor logits, torch::Tensor token_ids);
 
@@ -125,49 +127,98 @@ at::Tensor prefix_shared_attention(
 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
     m.doc() = "RL-Kernel High-Performance Operator Extension Library";
 
-    m.def("fused_logp", &fused_logp_forward, "Fused logp forward fallback");
+    m.def("fused_logp", rlk::nvtx_wrap("rlk::fused_logp", &fused_logp_forward),
+          "Fused logp forward fallback");
 
 #if defined(__CUDACC__) || defined(KERNEL_ALIGN_WITH_SM90)
-    m.def("fused_logp_sm90", &fused_logp_sm90_forward, "TMA-accelerated Online Softmax Fused LogP");
-    m.def("fused_linear_logp_sm90", &fused_linear_logp_sm90_forward,
+    m.def("fused_logp_sm90", rlk::nvtx_wrap("rlk::fused_logp_sm90", &fused_logp_sm90_forward),
+          "TMA-accelerated Online Softmax Fused LogP");
+    m.def("fused_linear_logp_sm90",
+          rlk::nvtx_wrap("rlk::fused_linear_logp_sm90", &fused_linear_logp_sm90_forward),
           "TMA+WGMMA fused linear log-prob (hidden @ W^T -> selected-token logp), SM90");
-    m.def("fused_linear_logp_sm90_global_target", &fused_linear_logp_sm90_global_target_forward,
+    m.def("fused_linear_logp_sm90_global_target",
+          rlk::nvtx_wrap("rlk::fused_linear_logp_sm90_global_target",
+                         &fused_linear_logp_sm90_global_target_forward),
           "TMA+WGMMA local-shard target-logit/lse for vocab-parallel linear log-prob, SM90");
-    m.def("fused_linear_logp_sm90_backward", &fused_linear_logp_sm90_backward,
+    m.def("fused_linear_logp_sm90_backward",
+          rlk::nvtx_wrap("rlk::fused_linear_logp_sm90_backward",
+                         &fused_linear_logp_sm90_backward),
           "CUDA fused backward for linear log-prob, SM90 backend");
-    m.def("linear_logp_probs_bf16_forward", &linear_logp_probs_bf16_forward,
+    m.def("linear_logp_probs_bf16_forward",
+          rlk::nvtx_wrap("rlk::linear_logp_probs_bf16_forward", &linear_logp_probs_bf16_forward),
           "Build bf16 softmax probabilities and selected log-prob from bf16 logits");
-    m.def("linear_logp_bf16_forward", &linear_logp_bf16_forward,
+    m.def("linear_logp_bf16_forward",
+          rlk::nvtx_wrap("rlk::linear_logp_bf16_forward", &linear_logp_bf16_forward),
           "Build selected log-prob and lse from bf16 logits without saving probabilities");
-    m.def("linear_logp_local_probs_bf16_forward", &linear_logp_local_probs_bf16_forward,
+    m.def("linear_logp_local_probs_bf16_forward",
+          rlk::nvtx_wrap("rlk::linear_logp_local_probs_bf16_forward",
+                         &linear_logp_local_probs_bf16_forward),
           "Build local bf16 softmax probabilities, target logits, and lse from bf16 logits");
-    m.def("linear_logp_local_bf16_forward", &linear_logp_local_bf16_forward,
+    m.def("linear_logp_local_bf16_forward",
+          rlk::nvtx_wrap("rlk::linear_logp_local_bf16_forward", &linear_logp_local_bf16_forward),
           "Build local target logits and lse from bf16 logits without saving probabilities");
-    m.def("linear_logp_probs_bf16_to_dlogits_", &linear_logp_probs_bf16_to_dlogits_,
+    m.def("linear_logp_probs_bf16_to_dlogits_",
+          rlk::nvtx_wrap("rlk::linear_logp_probs_bf16_to_dlogits_",
+                         &linear_logp_probs_bf16_to_dlogits_),
           "In-place bf16 probs -> dlogits for selected log-prob backward");
     m.def("linear_logp_local_probs_bf16_to_dlogits_",
-          &linear_logp_local_probs_bf16_to_dlogits_,
+          rlk::nvtx_wrap("rlk::linear_logp_local_probs_bf16_to_dlogits_",
+                         &linear_logp_local_probs_bf16_to_dlogits_),
           "In-place local bf16 probs -> TP dlogits for selected log-prob backward");
-    m.def("linear_logp_logits_bf16_to_dlogits", &linear_logp_logits_bf16_to_dlogits,
+    m.def("linear_logp_logits_bf16_to_dlogits",
+          rlk::nvtx_wrap("rlk::linear_logp_logits_bf16_to_dlogits",
+                         &linear_logp_logits_bf16_to_dlogits),
           "Build bf16 dlogits from bf16 logits and fp32 lse");
 #endif
 
 #if defined(__CUDACC__) || defined(KERNEL_ALIGN_WITH_CUDA)
-    m.def("fused_logp_forward_out", &fused_logp_forward_out, "Fused logp out");
-    m.def("fused_logp_forward_fp32", &fused_logp_forward_fp32, "Fused logp fp32");
-    m.def("fused_logp_forward_indexed_out", &fused_logp_forward_indexed_out, "Fused logp indexed out");
-    m.def("fused_logp_forward_indexed_fp32", &fused_logp_forward_indexed_fp32, "Fused logp indexed fp32");
-    m.def("fused_logp_forward_online_out", &fused_logp_forward_online_out, "Fused logp online out");
-    m.def("fused_logp_forward_online_fp32", &fused_logp_forward_online_fp32, "Fused logp online fp32");
-    m.def("fused_logp_forward_online_indexed_out", &fused_logp_forward_online_indexed_out, "Fused logp online indexed out");
-    m.def("fused_logp_forward_online_indexed_fp32", &fused_logp_forward_online_indexed_fp32, "Fused logp online indexed fp32");
-    m.def("deterministic_logp", &deterministic_logp_forward, "Batch-invariant deterministic logp");
-    m.def("deterministic_logp_forward_out", &deterministic_logp_forward_out, "Batch-invariant deterministic logp out");
-    m.def("deterministic_logp_forward_fp32", &deterministic_logp_forward_fp32, "Batch-invariant deterministic logp fp32");
-    m.def("deterministic_logp_forward_indexed_out", &deterministic_logp_forward_indexed_out, "Batch-invariant deterministic logp indexed out");
-    m.def("deterministic_logp_forward_indexed_fp32", &deterministic_logp_forward_indexed_fp32, "Batch-invariant deterministic logp indexed fp32");
+    m.def("fused_logp_forward_out",
+          rlk::nvtx_wrap("rlk::fused_logp_forward_out", &fused_logp_forward_out),
+          "Fused logp out");
+    m.def("fused_logp_forward_fp32",
+          rlk::nvtx_wrap("rlk::fused_logp_forward_fp32", &fused_logp_forward_fp32),
+          "Fused logp fp32");
+    m.def("fused_logp_forward_indexed_out",
+          rlk::nvtx_wrap("rlk::fused_logp_forward_indexed_out", &fused_logp_forward_indexed_out),
+          "Fused logp indexed out");
+    m.def("fused_logp_forward_indexed_fp32",
+          rlk::nvtx_wrap("rlk::fused_logp_forward_indexed_fp32", &fused_logp_forward_indexed_fp32),
+          "Fused logp indexed fp32");
+    m.def("fused_logp_forward_online_out",
+          rlk::nvtx_wrap("rlk::fused_logp_forward_online_out", &fused_logp_forward_online_out),
+          "Fused logp online out");
+    m.def("fused_logp_forward_online_fp32",
+          rlk::nvtx_wrap("rlk::fused_logp_forward_online_fp32", &fused_logp_forward_online_fp32),
+          "Fused logp online fp32");
+    m.def("fused_logp_forward_online_indexed_out",
+          rlk::nvtx_wrap("rlk::fused_logp_forward_online_indexed_out",
+                         &fused_logp_forward_online_indexed_out),
+          "Fused logp online indexed out");
+    m.def("fused_logp_forward_online_indexed_fp32",
+          rlk::nvtx_wrap("rlk::fused_logp_forward_online_indexed_fp32",
+                         &fused_logp_forward_online_indexed_fp32),
+          "Fused logp online indexed fp32");
+    m.def("deterministic_logp",
+          rlk::nvtx_wrap("rlk::deterministic_logp", &deterministic_logp_forward),
+          "Batch-invariant deterministic logp");
+    m.def("deterministic_logp_forward_out",
+          rlk::nvtx_wrap("rlk::deterministic_logp_forward_out", &deterministic_logp_forward_out),
+          "Batch-invariant deterministic logp out");
+    m.def("deterministic_logp_forward_fp32",
+          rlk::nvtx_wrap("rlk::deterministic_logp_forward_fp32", &deterministic_logp_forward_fp32),
+          "Batch-invariant deterministic logp fp32");
+    m.def("deterministic_logp_forward_indexed_out",
+          rlk::nvtx_wrap("rlk::deterministic_logp_forward_indexed_out",
+                         &deterministic_logp_forward_indexed_out),
+          "Batch-invariant deterministic logp indexed out");
+    m.def("deterministic_logp_forward_indexed_fp32",
+          rlk::nvtx_wrap("rlk::deterministic_logp_forward_indexed_fp32",
+                         &deterministic_logp_forward_indexed_fp32),
+          "Batch-invariant deterministic logp indexed fp32");
 
     // registry Prefix-Shared Attention
-    m.def("prefix_shared_attention", &prefix_shared_attention, "Prefix-Shared Fused Attention for GRPO");
+    m.def("prefix_shared_attention",
+          rlk::nvtx_wrap("rlk::prefix_shared_attention", &prefix_shared_attention),
+          "Prefix-Shared Fused Attention for GRPO");
 #endif
 }

csrc/utils/nvtx_utils.h

@@ -0,0 +1,76 @@
+// SPDX-License-Identifier: Apache-2.0
+// Copyright (c) 2026 RL-Kernel Contributors
+
+// NVTX range helpers for the pybind operator bindings.
+//
+// NVTX v3 is header-only and ships with CUDA >= 10, so no linker change is
+// needed. ROCm builds and builds with KERNEL_ALIGN_DISABLE_NVTX take the
+// identity pass-through branch, which generates no code.
+
+#pragma once
+
+#include <cstdlib>
+#include <cstring>
+#include <utility>
+
+#if defined(KERNEL_ALIGN_WITH_CUDA) && !defined(KERNEL_ALIGN_DISABLE_NVTX)
+#include <nvtx3/nvToolsExt.h>
+
+namespace rlk {
+
+inline bool nvtx_enabled() {
+  static const bool enabled = [] {
+    const char* value = std::getenv("RL_KERNEL_NVTX");
+    return value != nullptr &&
+           (std::strcmp(value, "1") == 0 || std::strcmp(value, "true") == 0 ||
+            std::strcmp(value, "yes") == 0 || std::strcmp(value, "on") == 0);
+  }();
+  return enabled;
+}
+
+class NvtxRange {
+ public:
+  explicit NvtxRange(const char* name) : active_(nvtx_enabled()) {
+    if (active_) {
+      nvtxRangePushA(name);
+    }
+  }
+  ~NvtxRange() {
+    if (active_) {
+      nvtxRangePop();
+    }
+  }
+  NvtxRange(const NvtxRange&) = delete;
+  NvtxRange& operator=(const NvtxRange&) = delete;
+
+ private:
+  bool active_;
+};
+
+// Wraps a free function so the pybind-bound call is enclosed in an NVTX range.
+template <typename Ret, typename... Args>
+auto nvtx_wrap(const char* name, Ret (*fn)(Args...)) {
+  return [name, fn](Args... args) -> Ret {
+    NvtxRange guard{name};
+    return fn(std::forward<Args>(args)...);
+  };
+}
+
+}  // namespace rlk
+
+#define RLK_NVTX_RANGE(name) ::rlk::NvtxRange _rlk_nvtx_guard{(name)}
+
+#else  // ROCm / NVTX-disabled builds: identity pass-through, zero code generated.
+
+namespace rlk {
+
+template <typename F>
+auto nvtx_wrap(const char*, F fn) {
+  return fn;
+}
+
+}  // namespace rlk
+
+#define RLK_NVTX_RANGE(name)
+
+#endif

docs/.nav.yml

@@ -12,6 +12,7 @@ nav:
       - getting_started/installation.md
       - getting_started/faq.md
       - Hardware Profiling Guide: getting_started/hardware-profiling.md
+      - Observability: getting_started/observability.md
     - Operators:
       - operators/README.md
       - operators/activation.md

docs/getting_started/observability.md

@@ -0,0 +1,158 @@
+# Observability
+
+RL-Kernel ships a two-layer observability stack:
+
+- **NVTX tracing (micro):** opt-in per-operator `rlk::<op>` ranges around every C++ binding plus
+  Python stage ranges (`rlk::rollout.generate`, `rlk::score`, `rlk::train_step`, ...),
+  visible on an [Nsight Systems](https://developer.nvidia.com/nsight-systems) timeline.
+  Set `RL_KERNEL_NVTX=1` only for profiling runs; normal runs retain just a cached branch.
+- **Prometheus metrics (macro):** a per-process `/metrics` endpoint with token throughput,
+  rollout QPS, stage latencies, hardware fallback state, KV-cache fragmentation, peak GPU
+  memory, and training loss — plus a ready-made Grafana dashboard.
+
+## Installation
+
+Prometheus support is an optional extra; the engine runs unchanged without it:
+
+```bash
+pip install -e ".[observability]"     # adds prometheus-client
+```
+
+Without `prometheus-client` installed, every metrics call degrades to a silent no-op.
+
+## Prometheus metrics
+
+### Starting the exporter
+
+The exporter never starts as an import side effect. Either call it explicitly:
+
+```python
+from rl_engine.observability import start_metrics_server
+
+port = start_metrics_server()          # default 0.0.0.0:8000; port=0 -> ephemeral
+```
+
+or opt in via environment variables when running the examples:
+
+```bash
+RL_KERNEL_METRICS=1 python examples/grpo_single_gpu.py --steps 3 --device cpu &
+curl -s localhost:8000/metrics | grep rl_kernel_
+```
+
+| Variable | Default | Meaning |
+| --- | --- | --- |
+| `RL_KERNEL_METRICS` | unset | `1` → examples/entrypoints start the exporter |
+| `RL_KERNEL_METRICS_PORT` | `8000` | Exporter port; `0` = ephemeral |
+| `RL_KERNEL_METRICS_ADDR` | `0.0.0.0` | Bind address override |
+
+For multi-process (e.g. Ray) deployments each worker exposes its own endpoint; use
+`RL_KERNEL_METRICS_PORT=0` to avoid port collisions and let your scrape topology discover
+the workers.
+
+### Metric schema
+
+All metrics use the `rl_kernel_` prefix and Prometheus base units (seconds, bytes).
+
+| Metric | Type | Labels |
+| --- | --- | --- |
+| `rl_kernel_build_info` | Gauge (=1) | `version`, `device_type`, `backend_version`, `ext_available` |
+| `rl_kernel_stage_duration_seconds` | Histogram | `stage` |
+| `rl_kernel_requests_total` | Counter | `stage`, `status` (`ok`/`error`) |
+| `rl_kernel_output_tokens_total` | Counter | `stage` |
+| `rl_kernel_selected_backend_info` | Gauge (=1) | `op_type`, `backend`, `is_fallback` |
+| `rl_kernel_hardware_fallback_total` | Counter | `op_type`, `requested_backend`, `selected_backend`, `reason` |
+| `rl_kernel_kv_cache_blocks` | Gauge | `kind` (`required`/`reserved`) |
+| `rl_kernel_kv_cache_fragmentation_ratio` | Gauge | — |
+| `rl_kernel_gpu_peak_memory_bytes` | Gauge | `stage`, `kind` (`allocated`/`reserved`) |
+| `rl_kernel_training_loss` | Gauge | — |
+| `rl_kernel_weight_version` | Gauge | `role` (`published`/`consumed`) |
+
+Fallbacks are reported at two layers: `selected_backend_info` captures the *state* chosen
+once per process at kernel dispatch time (`is_fallback="true"` means a lower-priority
+backend was selected), while `hardware_fallback_total` counts *events* — backend
+import/instantiation failures and per-call attention-backend fallbacks
+(flash-attention → eager).
+
+### Useful queries
+
+```promql
+rate(rl_kernel_output_tokens_total[1m])                          # token throughput
+rate(rl_kernel_requests_total{stage="rollout_generate"}[1m])     # rollout QPS
+histogram_quantile(0.95, rate(rl_kernel_stage_duration_seconds_bucket[5m]))
+sum(rl_kernel_selected_backend_info{is_fallback="true"}) / sum(rl_kernel_selected_backend_info)
+```
+
+### Prometheus scrape config
+
+```yaml
+scrape_configs:
+  - job_name: rl-kernel
+    scrape_interval: 5s
+    static_configs:
+      - targets: ["localhost:8000"]
+```
+
+### Grafana dashboard
+
+Import [`examples/grafana-dashboard.json`](https://github.com/RL-Align/RL-Kernel/blob/main/examples/grafana-dashboard.json)
+(Grafana ≥ 10, *Dashboards → Import*) and select your Prometheus datasource when prompted —
+the dashboard binds it through the `${DS_PROMETHEUS}` input, no hardcoded UID.
+
+### Using the facade in your own code
+
+```python
+from rl_engine.observability import metrics
+
+with metrics.stage_timer("my_stage"):        # duration histogram + ok/error counter
+    run_stage()
+metrics.add_output_tokens("my_stage", n_tokens)
+```
+
+Rules the built-in instrumentation follows (and yours should too):
+
+- per step / per request granularity — never per token in hot loops;
+- label values only from finite enums in code, never user data;
+- exceptions propagate unchanged (recorded with `status="error"`);
+- never re-query something the stage already computed — see below.
+
+### Overhead
+
+Metrics use exact 16-call batched updates for the module-level registry: counters and
+histogram buckets remain exact, gauges retain the latest value, and pending updates flush
+before each scrape. This moves Prometheus locks and repeated fixed-label lookups off the
+hot path. On an RTX 5090 with torch 2.8/CUDA 12.8, the complete default score path measured
+below 1% overhead at the benchmark's toy shape (paired, randomly interleaved calls).
+
+| Workload | Baseline stage time | With metrics | Overhead |
+| --- | --- | --- | --- |
+| `benchmark_stateless_executor` defaults (toy model, hidden 256 / vocab 4096) | 1.038 ms | 1.047 ms | +0.90 % |
+
+With `RL_KERNEL_NVTX=1`, Python and C++ ranges intentionally add profiler annotation cost;
+do not use that mode for the normal metrics-overhead acceptance run.
+
+If you add instrumentation, be careful what you *read* to populate it. Querying
+`torch.cuda.max_memory_allocated()` / `max_memory_reserved()` costs ~85 µs for the pair —
+each one rebuilds the full `torch.cuda.memory_stats()` dict. Prefer values the stage has
+already computed, or read `memory_stats()` once and pull both peaks out of it.
+
+## NVTX tracing with Nsight Systems
+
+Capture a timeline on a CUDA machine:
+
+```bash
+RL_KERNEL_NVTX=1 nsys profile -t nvtx,cuda -o rlk_trace \
+  python examples/grpo_single_gpu.py --steps 3
+nsys stats --report nvtx_sum rlk_trace.nsys-rep      # lists rlk::* ranges
+```
+
+On the timeline the Python stage ranges (`rlk::train_step` with `rlk::train.forward` /
+`rlk::train.backward` / `rlk::train.optim_step` inside) contain the per-operator
+`rlk::<binding_name>` ranges emitted by the C++ extension, correlated with kernel
+execution on the CUDA streams.
+
+Notes:
+
+- NVTX v3 is header-only and ships with the CUDA Toolkit; no extra link dependency.
+- ROCm builds compile the tracing macros to no-ops.
+- To compile the extension without NVTX, build with `KERNEL_ALIGN_DISABLE_NVTX=1`.
+- Python stage ranges silently no-op when CUDA is unavailable.

envs.py

@@ -26,3 +26,4 @@ def env_flag(name: str, default: bool = False) -> bool:
 KERNEL_ALIGN_NCU_LINEINFO = "KERNEL_ALIGN_NCU_LINEINFO"
 KERNEL_ALIGN_ALLOW_UNSUPPORTED_MSVC = "KERNEL_ALIGN_ALLOW_UNSUPPORTED_MSVC"
 KERNEL_ALIGN_FORCE_SM90 = "KERNEL_ALIGN_FORCE_SM90"
+KERNEL_ALIGN_DISABLE_NVTX = "KERNEL_ALIGN_DISABLE_NVTX"