From f8e945233b780fd8b94c48435e2bf57963d5ffa6 Mon Sep 17 00:00:00 2001
From: root <root@gpu-h100-0058.cm.cluster>
Date: Wed, 29 Oct 2025 08:22:48 -0700
Subject: [PATCH 1/5] initial commit

---
 .../inference/gpt/gpt_dynamic_inference.py    |  6 ++
 .../inference/contexts/dynamic_context.py     | 83 ++++++++++++++++++-
 .../core/inference/engines/dynamic_engine.py  | 72 ++++++++++++++++
 megatron/rl/inference/megatron.py             | 23 ++++-
 megatron/training/arguments.py                |  6 +-
 5 files changed, 185 insertions(+), 5 deletions(-)

diff --git a/examples/inference/gpt/gpt_dynamic_inference.py b/examples/inference/gpt/gpt_dynamic_inference.py
index a3b82a69e68..84868934bc9 100644
--- a/examples/inference/gpt/gpt_dynamic_inference.py
+++ b/examples/inference/gpt/gpt_dynamic_inference.py
@@ -133,6 +133,10 @@ def get_inference_context(
     else:
         max_sequence_length = args.inference_max_seq_length
 
+    metrics_writer = None
+    if args.inference_wandb_logging_step_interval > 0:
+        metrics_writer = get_wandb_writer()
+
     # Inference context.
     context = DynamicInferenceContext(
         params_dtype=args.params_dtype,
@@ -161,6 +165,7 @@ def get_inference_context(
         use_cuda_graphs_for_non_decode_steps=not args.decode_only_cuda_graphs,
         use_flashinfer_fused_rope=args.use_flashinfer_fused_rope,
         unified_memory_level=args.inference_dynamic_batching_unified_memory_level,
+        metrics_writer=metrics_writer,
     )
 
     return context
@@ -388,6 +393,7 @@ def main():
         random_seed=args.seed,
         track_paused_request_events=args.inference_dynamic_batching_track_paused_request_events,
         enable_chunked_prefill=not args.disable_chunked_prefill,
+        inference_logging_step_interval=args.inference_wandb_logging_step_interval,
     )
 
     setup_prefix = build_dynamic_engine_setup_prefix(args, model, context, requests)
diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py
index db4866d1142..e330e67d552 100644
--- a/megatron/core/inference/contexts/dynamic_context.py
+++ b/megatron/core/inference/contexts/dynamic_context.py
@@ -4,7 +4,7 @@
 import warnings
 from contextlib import nullcontext
 from enum import Enum
-from typing import List, Optional, Tuple
+from typing import TYPE_CHECKING, List, Optional, Tuple
 
 import torch
 import torch.nn.functional as F
@@ -45,6 +45,17 @@
 except ImportError:
     HAVE_FLASHINFER = False
 
+try:
+    import wandb  # pylint: disable=unused-import
+
+    HAVE_WANDB = True
+except ImportError:
+    HAVE_WANDB = False
+    wandb = None
+
+if TYPE_CHECKING:
+    import wandb as WandbModule
+
 
 class ContextOverflowError(Exception):
     """Base exception for when a new request does not fit.
@@ -222,6 +233,7 @@ class DynamicInferenceContext(BaseInferenceContext):
             levels will be included to control other tensors within the context.
         use_flashinfer_fused_rope (bool): If True, use flashinfer's fused rope implementation.
         If None, defaults to using flash-infer if available.
+        metrics_writer (Optional['WandbModule']): Wandb module for writing metrics.
     """
 
     def __init__(
@@ -247,6 +259,7 @@ def __init__(
         use_cuda_graphs_for_non_decode_steps: bool = True,
         use_flashinfer_fused_rope: bool = False,
         unified_memory_level: Optional[int] = 0,
+        metrics_writer: Optional['WandbModule'] = None,
     ):
         super().__init__(materialize_only_last_token_logits=materialize_only_last_token_logits)
 
@@ -256,6 +269,8 @@ def __init__(
                 block_size_tokens == 64
             ), "Flash MLA requires a block size of 64. Set --inference-dynamic-batching-block-size 64 to fix this assert"
 
+        self.metrics_writer = metrics_writer
+
         # Per partition num heads and hidden size.
         projection_size = kv_channels * num_attention_heads
         if tensor_model_parallel_size is None:
@@ -1618,3 +1633,69 @@ def calculate_log_probs(
 
         # Convert each log prob tensor into a list
         return [lp.tolist() for lp in selected_log_probs_list]
+
+    def get_kvcache_utilization_stats(self) -> dict:
+        """Compute KV cache buffer utilization stats for the current step.
+
+        Returns a dictionary with counts and percentages for both allocated chunk
+        usage (overall buffer occupancy) and active usage (chunks referenced by
+        currently active requests this step).
+
+        Return:
+            {
+            'total_chunks': int,
+            'allocated_chunks': int,
+            'active_unique_chunks': int,
+            'allocated_utilization': float,  # 0..1
+            'active_utilization': float,     # 0..1
+            'active_request_count': int,
+            'paused_request_count': int,
+            'gtd_chunk_count': int,
+            }
+        """
+        # Total usable chunks exclude the reserved dummy chunk.
+        total_chunks = max(self.chunk_allocator.chunk_count_total - 1, 1)
+        chunk_count_avail = int(self.chunk_allocator.chunk_count_avail)
+
+        # Overall allocated chunks in the buffer right now.
+        allocated_chunks = (self.chunk_allocator.chunk_count_total - 1) - chunk_count_avail
+        allocated_chunks = int(max(0, allocated_chunks))
+
+        # Active unique chunks referenced by current active requests only.
+        # Restrict to active rows to avoid padded rows in cudagraph-only tensors.
+        active_start = self.paused_request_count
+        active_end = self.total_request_count
+        if active_end > active_start:
+            active_rows = self.request_to_kv_chunk_ids[active_start:active_end]
+            # Filter valid chunk ids (>= 0) and count unique ids.
+            valid_ids = active_rows[active_rows >= 0]
+            if valid_ids.numel() > 0:
+                # torch.unique supports CUDA tensors; keep computation on device.
+                unique_ids = torch.unique(valid_ids)
+                active_unique_chunks = int(unique_ids.numel())
+            else:
+                active_unique_chunks = 0
+        else:
+            active_unique_chunks = 0
+
+        allocated_utilization = float(allocated_chunks) / float(total_chunks)
+        active_utilization = float(active_unique_chunks) / float(total_chunks)
+
+        # Diagnostic helpers
+        num_non_gtd_chunks = max(0, chunk_count_avail - int(self.gtd_chunk_count))
+        total_request_count = int(self.total_request_count)
+        return {
+            'total_chunks': int(total_chunks),
+            'allocated_chunks': int(allocated_chunks),
+            'active_unique_chunks': int(active_unique_chunks),
+            'allocated_utilization': allocated_utilization,
+            'active_utilization': active_utilization,
+            'active_request_count': int(self.get_active_request_count()),
+            'paused_request_count': int(self.paused_request_count),
+            'gtd_chunk_count': int(self.gtd_chunk_count),
+            'chunk_count_avail': int(chunk_count_avail),
+            'num_non_gtd_chunks': int(num_non_gtd_chunks),
+            'active_token_count': int(self.active_token_count),
+            'total_request_count': int(total_request_count),
+            'max_requests': int(self.max_requests),
+        }
diff --git a/megatron/core/inference/engines/dynamic_engine.py b/megatron/core/inference/engines/dynamic_engine.py
index 23a122f1c5f..bc513e68f4f 100644
--- a/megatron/core/inference/engines/dynamic_engine.py
+++ b/megatron/core/inference/engines/dynamic_engine.py
@@ -57,6 +57,14 @@
 except:
     HAVE_MSGPACK = False
 
+try:
+    import wandb
+
+    HAVE_WANDB = True
+except ImportError:
+    HAVE_WANDB = False
+    wandb = None
+
 
 def format_mem_bytes(mem_bytes):
     """Convert a byte count to a human-readable string in tb, gb, mb, kb, or bytes."""
@@ -89,6 +97,8 @@ class DynamicInferenceEngine(AbstractEngine):
         static_sampling (bool): If True, all requests are assumed to have the same
             sampling parameters. This avoids needing to loop through all requests and
             their sampling parameters every generation step, improving latency.
+        inference_logging_step_interval (int): The step interval at which to log
+        inference metrics to wandb. Defaults to 0, which means no logging.
     """
 
     def __init__(
@@ -101,6 +111,7 @@ def __init__(
         track_paused_request_events: bool = False,
         enable_chunked_prefill: bool = True,
         static_sampling: bool = False,
+        inference_logging_step_interval: int = 0,
     ):
 
         if enable_cuda_graph is not None:
@@ -137,6 +148,32 @@ def __init__(
         self.enable_chunked_prefill = enable_chunked_prefill
         self.static_sampling = static_sampling
 
+        self.inference_logging_step_interval = inference_logging_step_interval
+        # Configure wandb to use separate step counter for inference metrics (only once)
+        if self.inference_logging_step_interval > 0 and self.context.metrics_writer is not None:
+            logging.info(
+                f"\033[1;93m[INFERENCE]\033[0m "
+                f"\033[1;95mLogging inference metrics to wandb (rank {torch.distributed.get_rank()})\033[0m"
+            )
+            if HAVE_WANDB and self.context.metrics_writer.__name__ == "wandb":
+                # Make all inference/* metrics use inference_step as their x-axis
+                # This allows inference and training to have independent step counters
+                context.metrics_writer.define_metric(
+                    "inference/*", step_metric="inference/inference_step"
+                )
+                # Initialize inference step offset by querying existing run history
+                self.inference_step_offset = 0
+                if wandb.run is not None:
+                    api_run = wandb.Api().run(
+                        f"{wandb.run.entity}/{wandb.run.project}/{wandb.run.id}"
+                    )
+                    max_step = 0
+                    for row in api_run.scan_history(keys=["inference/inference_step"]):
+                        val = row.get("inference/inference_step")
+                        if isinstance(val, (int, float)) and int(val) > max_step:
+                            max_step = int(val)
+                    self.inference_step_offset = int(max_step)
+
         # Initialize the asyncio loop if it has not already been initialized.
         # TODO: Start the engine loop here.
         self._loop = get_asyncio_loop()
@@ -771,6 +808,41 @@ async def async_step(
             self.request_completion_futures[failed_request_id].set_result(failed_request)
         self.failed_request_ids.clear()
 
+        # Log KV cache utilization stats to TensorBoard or W&B (only on logging rank)
+        if (
+            self.inference_logging_step_interval > 0
+            and self.step_count > 0
+            and self.step_count % self.inference_logging_step_interval == 0
+            and self.context.metrics_writer is not None
+        ):
+
+            # Get KV cache utilization stats from dynamic context
+            kv_stats = self.context.get_kvcache_utilization_stats()
+
+            # Prepare metrics dictionary with all stats
+            # Use 'inference/' prefix for all metrics to separate from training metrics
+            metrics = {
+                'inference/inference_step': int(self.inference_step_offset + int(self.step_count)),
+                'inference/step_time_s': float(step_time),
+                'inference/waiting_queue_len': int(len(self.waiting_request_ids)),
+                'inference/total_requests_dict_size': int(len(self.requests)),
+            }
+            # Add KV stats with inference/ prefix
+            # Convert utilization metrics from 0-1 range to 0-100 percentage range for better visualization
+            for key, value in kv_stats.items():
+                if 'utilization' in key:
+                    # Convert to percentage (0-100) and group under kvcache_utilization
+                    metrics[f'inference/{key}'] = float(value * 100.0)
+                else:
+                    metrics[f'inference/{key}'] = value
+
+            if HAVE_WANDB and self.context.metrics_writer.__name__ == "wandb":
+                self.context.metrics_writer.log(metrics, commit=True)
+            else:
+                raise ValueError(
+                    f"Unsupported metrics writer type: {type(self.context.metrics_writer)}"
+                )
+
         # Print context state.
         if verbose:
             context = self.context
diff --git a/megatron/rl/inference/megatron.py b/megatron/rl/inference/megatron.py
index 4c739b709c1..6d8fac32044 100644
--- a/megatron/rl/inference/megatron.py
+++ b/megatron/rl/inference/megatron.py
@@ -80,7 +80,8 @@ def get_static_inference_engine(args: Namespace, model: MegatronModule) -> Abstr
 
 
 ## This code is copied from tools/run_text_generation_server.py
-def get_dynamic_inference_engine(args: Namespace, model: MegatronModule) -> AbstractEngine:
+def get_dynamic_inference_engine(args: Namespace, model: MegatronModule, inference_logging_step_interval: int = 0,
+    metrics_writer = None) -> AbstractEngine:
     """Get the relevant backend for running inference.
 
     This function will automatically choose the TRTLLMBackend when possible, and default to Mcore backend if the user does not specify any backends. TRTLLMBackend is not implmented yet.
@@ -88,6 +89,8 @@ def get_dynamic_inference_engine(args: Namespace, model: MegatronModule) -> Abst
     Args:
         args (Namespace): The user arguments parsed from command line
         model (MegatronModule): The megatron model.
+        inference_logging_step_interval (int): Step interval for logging inference metrics.
+        metrics_writer: Metrics writer (wandb module) for logging.
 
     Returns:
         AbstractBackend: The chosen backend
@@ -126,6 +129,7 @@ def get_dynamic_inference_engine(args: Namespace, model: MegatronModule) -> Abst
         mamba_d_model=args.hidden_size,
         mamba_d_conv=4 if args.is_hybrid_model else None,
         mamba_d_state=args.mamba_state_dim,
+        metrics_writer=metrics_writer,
     )
 
     inference_wrapped_model = GPTInferenceWrapper(model, args, inference_context)
@@ -143,6 +147,7 @@ def get_dynamic_inference_engine(args: Namespace, model: MegatronModule) -> Abst
         context=inference_context,
         enable_cuda_graph=args.enable_cuda_graph,
         random_seed=args.seed,
+        inference_logging_step_interval=inference_logging_step_interval,
     )
 
 
@@ -195,8 +200,20 @@ async def launch(cls, model: GPTModel, **kwargs):
                 logging.WARNING,
                 "WARNING: Tokenizer has no BOS token so prompt will not have BOS token",
             )
-
-        inference_engine: DynamicInferenceEngine = get_dynamic_inference_engine(args, model)
+        
+        # Get inference logging configuration from args
+        inference_logging_step_interval = args.inference_wandb_logging_step_interval
+        
+        # Get metrics writer if logging is enabled and on the logging rank
+        # Use the same rank convention as training (last rank logs)
+        metrics_writer = None
+        if inference_logging_step_interval > 0 and args.rank == (args.world_size - 1):
+            metrics_writer = get_wandb_writer()
+            if metrics_writer is None:
+                log_single_rank(logger, logging.WARNING, "WARNING: --rl-inference-logging-step-interval is set but no metrics writer "
+                           "wandb module is available. Inference logging will be disabled.")
+
+        inference_engine: DynamicInferenceEngine = get_dynamic_inference_engine(args, model, inference_logging_step_interval, metrics_writer)
         coordinator = DynamicEngineCoordinator(
             inference_engine,
             inference_max_requests=inference_engine.context.max_requests,
diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py
index 5d1bc2e40a3..271e92c634a 100644
--- a/megatron/training/arguments.py
+++ b/megatron/training/arguments.py
@@ -1497,7 +1497,11 @@ def _add_inference_args(parser):
                        help='Number of chunks along sequence dimension for MLP '
                        'computation during prefill')
     group.add_argument('--disable-chunked-prefill', default=False, action="store_true",
-                       help='Disable chunked prefill (chunked prefill is enabled by default).')  
+                       help='Disable chunked prefill (chunked prefill is enabled by default).')
+    group.add_argument('--inference-wandb-logging-step-interval', type=int, default=0,
+                       help='Step interval for logging inference metrics to wandb. '
+                            'Default to 0 to disable inference wandb logging.')
+
     return parser
 
 

From f9a267b7d6c0f36d33b26883fe821a90645e35f9 Mon Sep 17 00:00:00 2001
From: William Dykas <wdykas@cw-pdx-cs-001-vscode-02.cm.cluster>
Date: Wed, 29 Oct 2025 08:30:05 -0700
Subject: [PATCH 2/5] clean up

---
 megatron/core/inference/contexts/dynamic_context.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py
index e330e67d552..68b5831d0a2 100644
--- a/megatron/core/inference/contexts/dynamic_context.py
+++ b/megatron/core/inference/contexts/dynamic_context.py
@@ -1646,8 +1646,8 @@ def get_kvcache_utilization_stats(self) -> dict:
             'total_chunks': int,
             'allocated_chunks': int,
             'active_unique_chunks': int,
-            'allocated_utilization': float,  # 0..1
-            'active_utilization': float,     # 0..1
+            'allocated_utilization': float,
+            'active_utilization': float,
             'active_request_count': int,
             'paused_request_count': int,
             'gtd_chunk_count': int,

From e75d3d45e6329a9046c44b35edfdadfaed8a43e2 Mon Sep 17 00:00:00 2001
From: William Dykas <wdykas@cw-pdx-cs-001-vscode-02.cm.cluster>
Date: Wed, 29 Oct 2025 09:45:50 -0700
Subject: [PATCH 3/5] tests

---
 .../inference/contexts/dynamic_context.py     |  56 ++--
 .../inference/test_wandb_logging.py           | 276 ++++++++++++++++++
 2 files changed, 303 insertions(+), 29 deletions(-)
 create mode 100644 tests/unit_tests/inference/test_wandb_logging.py

diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py
index 68b5831d0a2..676b661bba3 100644
--- a/megatron/core/inference/contexts/dynamic_context.py
+++ b/megatron/core/inference/contexts/dynamic_context.py
@@ -1637,64 +1637,62 @@ def calculate_log_probs(
     def get_kvcache_utilization_stats(self) -> dict:
         """Compute KV cache buffer utilization stats for the current step.
 
-        Returns a dictionary with counts and percentages for both allocated chunk
-        usage (overall buffer occupancy) and active usage (chunks referenced by
+        Returns a dictionary with counts and percentages for both allocated block
+        usage (overall buffer occupancy) and active usage (blocks referenced by
         currently active requests this step).
 
         Return:
             {
-            'total_chunks': int,
-            'allocated_chunks': int,
-            'active_unique_chunks': int,
+            'total_blocks': int,
+            'allocated_blocks': int,
+            'active_unique_blocks': int,
             'allocated_utilization': float,
             'active_utilization': float,
             'active_request_count': int,
             'paused_request_count': int,
-            'gtd_chunk_count': int,
+            'gtd_block_count': int,
             }
         """
-        # Total usable chunks exclude the reserved dummy chunk.
-        total_chunks = max(self.chunk_allocator.chunk_count_total - 1, 1)
-        chunk_count_avail = int(self.chunk_allocator.chunk_count_avail)
+        # Total usable blocks exclude the reserved dummy block.
+        total_blocks = max(self.block_allocator.block_count_total - 1, 1)
+        block_count_avail = int(self.block_allocator.block_count_avail)
 
-        # Overall allocated chunks in the buffer right now.
-        allocated_chunks = (self.chunk_allocator.chunk_count_total - 1) - chunk_count_avail
-        allocated_chunks = int(max(0, allocated_chunks))
+        # Overall allocated blocks in the buffer right now.
+        allocated_blocks = (self.block_allocator.block_count_total - 1) - block_count_avail
+        allocated_blocks = int(max(0, allocated_blocks))
 
-        # Active unique chunks referenced by current active requests only.
-        # Restrict to active rows to avoid padded rows in cudagraph-only tensors.
+        # Active unique blocks referenced by current active requests only.
         active_start = self.paused_request_count
         active_end = self.total_request_count
         if active_end > active_start:
-            active_rows = self.request_to_kv_chunk_ids[active_start:active_end]
-            # Filter valid chunk ids (>= 0) and count unique ids.
+            active_rows = self.request_to_kv_block_ids[active_start:active_end]
+            # Filter valid block ids (>= 0) and count unique ids.
             valid_ids = active_rows[active_rows >= 0]
             if valid_ids.numel() > 0:
-                # torch.unique supports CUDA tensors; keep computation on device.
                 unique_ids = torch.unique(valid_ids)
-                active_unique_chunks = int(unique_ids.numel())
+                active_unique_blocks = int(unique_ids.numel())
             else:
-                active_unique_chunks = 0
+                active_unique_blocks = 0
         else:
-            active_unique_chunks = 0
+            active_unique_blocks = 0
 
-        allocated_utilization = float(allocated_chunks) / float(total_chunks)
-        active_utilization = float(active_unique_chunks) / float(total_chunks)
+        allocated_utilization = float(allocated_blocks) / float(total_blocks)
+        active_utilization = float(active_unique_blocks) / float(total_blocks)
 
         # Diagnostic helpers
-        num_non_gtd_chunks = max(0, chunk_count_avail - int(self.gtd_chunk_count))
+        num_non_gtd_blocks = max(0, block_count_avail - int(self.gtd_block_count))
         total_request_count = int(self.total_request_count)
         return {
-            'total_chunks': int(total_chunks),
-            'allocated_chunks': int(allocated_chunks),
-            'active_unique_chunks': int(active_unique_chunks),
+            'total_blocks': int(total_blocks),
+            'allocated_blocks': int(allocated_blocks),
+            'active_unique_blocks': int(active_unique_blocks),
             'allocated_utilization': allocated_utilization,
             'active_utilization': active_utilization,
             'active_request_count': int(self.get_active_request_count()),
             'paused_request_count': int(self.paused_request_count),
-            'gtd_chunk_count': int(self.gtd_chunk_count),
-            'chunk_count_avail': int(chunk_count_avail),
-            'num_non_gtd_chunks': int(num_non_gtd_chunks),
+            'gtd_block_count': int(self.gtd_block_count),
+            'block_count_avail': int(block_count_avail),
+            'num_non_gtd_blocks': int(num_non_gtd_blocks),
             'active_token_count': int(self.active_token_count),
             'total_request_count': int(total_request_count),
             'max_requests': int(self.max_requests),
diff --git a/tests/unit_tests/inference/test_wandb_logging.py b/tests/unit_tests/inference/test_wandb_logging.py
new file mode 100644
index 00000000000..3fd2f2977a9
--- /dev/null
+++ b/tests/unit_tests/inference/test_wandb_logging.py
@@ -0,0 +1,276 @@
+# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+
+"""Unit tests for wandb logging functionality in inference."""
+
+import pytest
+import torch
+from unittest.mock import Mock, MagicMock, patch, create_autospec
+
+from megatron.core.inference.contexts.dynamic_context import DynamicInferenceContext
+from megatron.core.inference.engines import DynamicInferenceEngine
+from megatron.core.inference.inference_request import DynamicInferenceRequest
+from megatron.core.inference.sampling_params import SamplingParams
+from megatron.core.inference.text_generation_controllers.text_generation_controller import (
+    TextGenerationController,
+)
+from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed
+from tests.unit_tests.test_utilities import Utils
+
+
+def set_rounder(value):
+    """Utility function to set the DynamicInferenceContext rounder."""
+    DynamicInferenceContext.ROUNDER = value  # For backwards compatibility
+    DynamicInferenceContext.TOKEN_ROUNDER = value
+    DynamicInferenceContext.REQUEST_ROUNDER = value
+
+
+class TestInferenceWandbLogging:
+    """Test suite for wandb logging in inference."""
+
+    def setup_method(self):
+        """Set up test fixtures."""
+        Utils.initialize_model_parallel(
+            tensor_model_parallel_size=1,
+            pipeline_model_parallel_size=1,
+        )
+        model_parallel_cuda_manual_seed(123)
+        set_rounder(64)
+
+    def teardown_method(self):
+        """Clean up test fixtures."""
+        set_rounder(64)
+        Utils.destroy_model_parallel()
+
+    def _get_dynamic_context(
+        self,
+        params_dtype=torch.float32,
+        num_layers=4,
+        kv_channels=8,
+        num_attention_heads=2,
+        max_sequence_length=512,
+        buffer_size_gb=0.03,
+        block_size_tokens=128,
+        buffer_guaranteed_fraction=0.1,
+        metrics_writer=None,
+    ):
+        """Helper to create a DynamicInferenceContext."""
+        return DynamicInferenceContext(
+            params_dtype=params_dtype,
+            num_layers=num_layers,
+            kv_channels=kv_channels,
+            num_attention_heads=num_attention_heads,
+            max_sequence_length=max_sequence_length,
+            num_cuda_graphs=None,
+            buffer_size_gb=buffer_size_gb,
+            buffer_guaranteed_fraction=buffer_guaranteed_fraction,
+            block_size_tokens=block_size_tokens,
+            metrics_writer=metrics_writer,
+        )
+
+
+    @pytest.mark.internal
+    def test_get_kvcache_utilization_stats_with_requests(self):
+        """Test get_kvcache_utilization_stats() with empty context and then with active requests."""
+        dynamic_context = self._get_dynamic_context()
+        
+        # First, test with empty context
+        stats = dynamic_context.get_kvcache_utilization_stats()
+        
+        # Verify all required fields are present
+        assert 'total_blocks' in stats
+        assert 'allocated_blocks' in stats
+        assert 'active_unique_blocks' in stats
+        assert 'allocated_utilization' in stats
+        assert 'active_utilization' in stats
+        assert 'active_request_count' in stats
+        assert 'paused_request_count' in stats
+        assert 'gtd_block_count' in stats
+        assert 'block_count_avail' in stats
+        assert 'num_non_gtd_blocks' in stats
+        assert 'active_token_count' in stats
+        assert 'total_request_count' in stats
+        assert 'max_requests' in stats
+        
+        # Verify values for empty context
+        assert stats['allocated_blocks'] == 0
+        assert stats['active_unique_blocks'] == 0
+        assert stats['allocated_utilization'] == 0.0
+        assert stats['active_utilization'] == 0.0
+        assert stats['active_request_count'] == 0
+        assert stats['paused_request_count'] == 0
+        assert stats['active_token_count'] == 0
+        assert stats['total_request_count'] == 0
+        
+        # Now add a request and verify stats update correctly
+        context_length = 144
+        dynamic_context.add_request(
+            DynamicInferenceRequest(
+                request_id=0,
+                prompt_tokens=torch.arange(0, context_length, dtype=torch.long, device='cuda'),
+                sampling_params=SamplingParams(
+                    num_tokens_to_generate=dynamic_context.max_tokens - context_length
+                ),
+            )
+        )
+        
+        # Initialize attention state to populate block table
+        dynamic_context.initialize_attention_state()
+        
+        # Get stats after adding request
+        stats_after = dynamic_context.get_kvcache_utilization_stats()
+        
+        # Verify that we have allocated blocks
+        assert stats_after['allocated_blocks'] > 0
+        assert stats_after['active_unique_blocks'] > 0
+        assert stats_after['allocated_utilization'] > 0.0
+        assert stats_after['active_utilization'] > 0.0
+        
+        # Verify request counts
+        assert stats_after['active_request_count'] == 1
+        assert stats_after['total_request_count'] == 1
+        assert stats_after['active_token_count'] == context_length
+        assert stats_after['paused_request_count'] == 0
+        
+        # Verify that total_blocks remains constant
+        assert stats_after['total_blocks'] == stats['total_blocks']
+        assert stats_after['total_blocks'] > 0
+        
+        # Verify that gtd_block_count remains constant
+        assert stats_after['gtd_block_count'] == stats['gtd_block_count']
+        
+        # Verify that max_requests remains constant
+        assert stats_after['max_requests'] == stats['max_requests']
+        assert stats_after['max_requests'] > 0
+        
+        # Verify block availability decreased after allocation
+        assert stats_after['block_count_avail'] < stats['block_count_avail']
+        
+        # Verify relationship: allocated_blocks + block_count_avail + 1 (dummy) = total
+        assert stats_after['allocated_blocks'] + stats_after['block_count_avail'] + 1 == dynamic_context.block_allocator.block_count_total
+        
+        # Verify utilization bounds [0, 1]
+        assert 0.0 <= stats_after['allocated_utilization'] <= 1.0
+        assert 0.0 <= stats_after['active_utilization'] <= 1.0
+        
+        # Verify relationship: active_utilization <= allocated_utilization
+        # (active blocks are a subset of allocated blocks)
+        assert stats_after['active_utilization'] <= stats_after['allocated_utilization']
+        
+        # Verify relationship: active_unique_blocks <= allocated_blocks
+        assert stats_after['active_unique_blocks'] <= stats_after['allocated_blocks']
+        
+        # Calculate expected number of blocks needed for this request
+        expected_blocks_needed = (context_length + dynamic_context.block_size_tokens - 1) // dynamic_context.block_size_tokens
+        assert stats_after['allocated_blocks'] == expected_blocks_needed
+
+    @pytest.mark.internal
+    def test_kvcache_utilization_stats_types(self):
+        """Test that get_kvcache_utilization_stats() returns correct types."""
+        dynamic_context = self._get_dynamic_context()
+        stats = dynamic_context.get_kvcache_utilization_stats()
+        
+        # All integer fields
+        int_fields = [
+            'total_blocks', 'allocated_blocks', 'active_unique_blocks',
+            'active_request_count', 'paused_request_count', 'gtd_block_count',
+            'block_count_avail', 'num_non_gtd_blocks', 'active_token_count',
+            'total_request_count', 'max_requests'
+        ]
+        
+        for field in int_fields:
+            assert isinstance(stats[field], int), f"{field} should be int but is {type(stats[field])}"
+        
+        # All float fields
+        float_fields = ['allocated_utilization', 'active_utilization']
+        for field in float_fields:
+            assert isinstance(stats[field], float), f"{field} should be float but is {type(stats[field])}"
+
+
+    @pytest.mark.internal
+    @patch('megatron.core.inference.engines.dynamic_engine.HAVE_WANDB', True)
+    def test_engine_logging_step_interval_zero(self):
+        """Test that no logging occurs when inference_logging_step_interval is 0."""
+        mock_wandb = Mock()
+        mock_wandb.__name__ = "wandb"
+        mock_wandb.log = Mock()
+        
+        dynamic_context = self._get_dynamic_context(metrics_writer=mock_wandb)
+        
+        # Create mock controller with proper spec to pass isinstance checks
+        mock_controller = create_autospec(TextGenerationController, instance=True)
+        # Set up nested mock structure
+        mock_controller.inference_wrapped_model = Mock()
+        mock_controller.inference_wrapped_model.model = Mock()
+        mock_controller.inference_wrapped_model.model.config = Mock()
+        mock_controller.inference_wrapped_model.model.config.cuda_graph_impl = "none"
+        
+        engine = DynamicInferenceEngine(
+            controller=mock_controller,
+            context=dynamic_context,
+            random_seed=123,
+            inference_logging_step_interval=0,  # Disabled
+        )
+        
+        # Verify log was never called
+        mock_wandb.log.assert_not_called()
+
+
+    @pytest.mark.internal
+    def test_paused_requests_in_stats(self):
+        """Test that paused requests are correctly reflected in stats."""
+        set_rounder(1)
+        dynamic_context = DynamicInferenceContext(
+            params_dtype=torch.float32,
+            num_layers=2,
+            kv_channels=64,
+            num_attention_heads=8,
+            max_sequence_length=128,
+            num_cuda_graphs=None,
+            buffer_size_gb=0.01,  # Small buffer to force pausing
+            buffer_guaranteed_fraction=0.1,
+            block_size_tokens=32,
+        )
+        
+        # Add multiple requests to potentially trigger pausing
+        for i in range(5):
+            dynamic_context.add_request(
+                DynamicInferenceRequest(
+                    request_id=i,
+                    prompt_tokens=torch.zeros(10, device='cuda'),
+                    sampling_params=SamplingParams(num_tokens_to_generate=10),
+                )
+            )
+        
+        if dynamic_context.total_request_count > 0:
+            dynamic_context.initialize_attention_state()
+            stats = dynamic_context.get_kvcache_utilization_stats()
+            
+            # Verify paused request count is included
+            assert 'paused_request_count' in stats
+            assert stats['paused_request_count'] >= 0
+
+    @pytest.mark.internal
+    def test_metrics_writer_none_handling(self):
+        """Test that engine handles None metrics_writer gracefully."""
+        dynamic_context = self._get_dynamic_context(metrics_writer=None)
+        
+        # Create mock controller with proper spec to pass isinstance checks
+        mock_controller = create_autospec(TextGenerationController, instance=True)
+        # Set up nested mock structure
+        mock_controller.inference_wrapped_model = Mock()
+        mock_controller.inference_wrapped_model.model = Mock()
+        mock_controller.inference_wrapped_model.model.config = Mock()
+        mock_controller.inference_wrapped_model.model.config.cuda_graph_impl = "none"
+        
+        # Should not raise error even with logging interval set
+        engine = DynamicInferenceEngine(
+            controller=mock_controller,
+            context=dynamic_context,
+            random_seed=123,
+            inference_logging_step_interval=10,
+        )
+        
+        # Verify engine was created successfully
+        assert engine.inference_logging_step_interval == 10
+        assert engine.context.metrics_writer is None
+

From 04c85499ae49a84eaa3be4ace8124a501006fb6b Mon Sep 17 00:00:00 2001
From: William Dykas <wdykas@cw-pdx-cs-001-vscode-02.cm.cluster>
Date: Wed, 29 Oct 2025 09:49:21 -0700
Subject: [PATCH 4/5] typo

---
 megatron/core/inference/engines/dynamic_engine.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/megatron/core/inference/engines/dynamic_engine.py b/megatron/core/inference/engines/dynamic_engine.py
index bc513e68f4f..9a727569f68 100644
--- a/megatron/core/inference/engines/dynamic_engine.py
+++ b/megatron/core/inference/engines/dynamic_engine.py
@@ -808,7 +808,7 @@ async def async_step(
             self.request_completion_futures[failed_request_id].set_result(failed_request)
         self.failed_request_ids.clear()
 
-        # Log KV cache utilization stats to TensorBoard or W&B (only on logging rank)
+        # Log KV cache utilization stats to W&B
         if (
             self.inference_logging_step_interval > 0
             and self.step_count > 0

From f206dc9c13aa70dba549b3661ce7395db83b34c0 Mon Sep 17 00:00:00 2001
From: root <root@gpu-h100-0220.cm.cluster>
Date: Fri, 31 Oct 2025 12:11:25 -0700
Subject: [PATCH 5/5] formatting

---
 .../inference/test_wandb_logging.py           | 102 ++++++++++--------
 1 file changed, 57 insertions(+), 45 deletions(-)

diff --git a/tests/unit_tests/inference/test_wandb_logging.py b/tests/unit_tests/inference/test_wandb_logging.py
index 3fd2f2977a9..1512e805f9c 100644
--- a/tests/unit_tests/inference/test_wandb_logging.py
+++ b/tests/unit_tests/inference/test_wandb_logging.py
@@ -2,9 +2,10 @@
 
 """Unit tests for wandb logging functionality in inference."""
 
+from unittest.mock import MagicMock, Mock, create_autospec, patch
+
 import pytest
 import torch
-from unittest.mock import Mock, MagicMock, patch, create_autospec
 
 from megatron.core.inference.contexts.dynamic_context import DynamicInferenceContext
 from megatron.core.inference.engines import DynamicInferenceEngine
@@ -30,8 +31,7 @@ class TestInferenceWandbLogging:
     def setup_method(self):
         """Set up test fixtures."""
         Utils.initialize_model_parallel(
-            tensor_model_parallel_size=1,
-            pipeline_model_parallel_size=1,
+            tensor_model_parallel_size=1, pipeline_model_parallel_size=1
         )
         model_parallel_cuda_manual_seed(123)
         set_rounder(64)
@@ -67,15 +67,14 @@ def _get_dynamic_context(
             metrics_writer=metrics_writer,
         )
 
-
     @pytest.mark.internal
     def test_get_kvcache_utilization_stats_with_requests(self):
         """Test get_kvcache_utilization_stats() with empty context and then with active requests."""
         dynamic_context = self._get_dynamic_context()
-        
+
         # First, test with empty context
         stats = dynamic_context.get_kvcache_utilization_stats()
-        
+
         # Verify all required fields are present
         assert 'total_blocks' in stats
         assert 'allocated_blocks' in stats
@@ -90,7 +89,7 @@ def test_get_kvcache_utilization_stats_with_requests(self):
         assert 'active_token_count' in stats
         assert 'total_request_count' in stats
         assert 'max_requests' in stats
-        
+
         # Verify values for empty context
         assert stats['allocated_blocks'] == 0
         assert stats['active_unique_blocks'] == 0
@@ -100,7 +99,7 @@ def test_get_kvcache_utilization_stats_with_requests(self):
         assert stats['paused_request_count'] == 0
         assert stats['active_token_count'] == 0
         assert stats['total_request_count'] == 0
-        
+
         # Now add a request and verify stats update correctly
         context_length = 144
         dynamic_context.add_request(
@@ -112,55 +111,60 @@ def test_get_kvcache_utilization_stats_with_requests(self):
                 ),
             )
         )
-        
+
         # Initialize attention state to populate block table
         dynamic_context.initialize_attention_state()
-        
+
         # Get stats after adding request
         stats_after = dynamic_context.get_kvcache_utilization_stats()
-        
+
         # Verify that we have allocated blocks
         assert stats_after['allocated_blocks'] > 0
         assert stats_after['active_unique_blocks'] > 0
         assert stats_after['allocated_utilization'] > 0.0
         assert stats_after['active_utilization'] > 0.0
-        
+
         # Verify request counts
         assert stats_after['active_request_count'] == 1
         assert stats_after['total_request_count'] == 1
         assert stats_after['active_token_count'] == context_length
         assert stats_after['paused_request_count'] == 0
-        
+
         # Verify that total_blocks remains constant
         assert stats_after['total_blocks'] == stats['total_blocks']
         assert stats_after['total_blocks'] > 0
-        
+
         # Verify that gtd_block_count remains constant
         assert stats_after['gtd_block_count'] == stats['gtd_block_count']
-        
+
         # Verify that max_requests remains constant
         assert stats_after['max_requests'] == stats['max_requests']
         assert stats_after['max_requests'] > 0
-        
+
         # Verify block availability decreased after allocation
         assert stats_after['block_count_avail'] < stats['block_count_avail']
-        
+
         # Verify relationship: allocated_blocks + block_count_avail + 1 (dummy) = total
-        assert stats_after['allocated_blocks'] + stats_after['block_count_avail'] + 1 == dynamic_context.block_allocator.block_count_total
-        
+        assert (
+            stats_after['allocated_blocks'] + stats_after['block_count_avail'] + 1
+            == dynamic_context.block_allocator.block_count_total
+        )
+
         # Verify utilization bounds [0, 1]
         assert 0.0 <= stats_after['allocated_utilization'] <= 1.0
         assert 0.0 <= stats_after['active_utilization'] <= 1.0
-        
+
         # Verify relationship: active_utilization <= allocated_utilization
         # (active blocks are a subset of allocated blocks)
         assert stats_after['active_utilization'] <= stats_after['allocated_utilization']
-        
+
         # Verify relationship: active_unique_blocks <= allocated_blocks
         assert stats_after['active_unique_blocks'] <= stats_after['allocated_blocks']
-        
+
         # Calculate expected number of blocks needed for this request
-        expected_blocks_needed = (context_length + dynamic_context.block_size_tokens - 1) // dynamic_context.block_size_tokens
+        expected_blocks_needed = (
+            context_length + dynamic_context.block_size_tokens - 1
+        ) // dynamic_context.block_size_tokens
         assert stats_after['allocated_blocks'] == expected_blocks_needed
 
     @pytest.mark.internal
@@ -168,23 +172,33 @@ def test_kvcache_utilization_stats_types(self):
         """Test that get_kvcache_utilization_stats() returns correct types."""
         dynamic_context = self._get_dynamic_context()
         stats = dynamic_context.get_kvcache_utilization_stats()
-        
+
         # All integer fields
         int_fields = [
-            'total_blocks', 'allocated_blocks', 'active_unique_blocks',
-            'active_request_count', 'paused_request_count', 'gtd_block_count',
-            'block_count_avail', 'num_non_gtd_blocks', 'active_token_count',
-            'total_request_count', 'max_requests'
+            'total_blocks',
+            'allocated_blocks',
+            'active_unique_blocks',
+            'active_request_count',
+            'paused_request_count',
+            'gtd_block_count',
+            'block_count_avail',
+            'num_non_gtd_blocks',
+            'active_token_count',
+            'total_request_count',
+            'max_requests',
         ]
-        
+
         for field in int_fields:
-            assert isinstance(stats[field], int), f"{field} should be int but is {type(stats[field])}"
-        
+            assert isinstance(
+                stats[field], int
+            ), f"{field} should be int but is {type(stats[field])}"
+
         # All float fields
         float_fields = ['allocated_utilization', 'active_utilization']
         for field in float_fields:
-            assert isinstance(stats[field], float), f"{field} should be float but is {type(stats[field])}"
-
+            assert isinstance(
+                stats[field], float
+            ), f"{field} should be float but is {type(stats[field])}"
 
     @pytest.mark.internal
     @patch('megatron.core.inference.engines.dynamic_engine.HAVE_WANDB', True)
@@ -193,9 +207,9 @@ def test_engine_logging_step_interval_zero(self):
         mock_wandb = Mock()
         mock_wandb.__name__ = "wandb"
         mock_wandb.log = Mock()
-        
+
         dynamic_context = self._get_dynamic_context(metrics_writer=mock_wandb)
-        
+
         # Create mock controller with proper spec to pass isinstance checks
         mock_controller = create_autospec(TextGenerationController, instance=True)
         # Set up nested mock structure
@@ -203,18 +217,17 @@ def test_engine_logging_step_interval_zero(self):
         mock_controller.inference_wrapped_model.model = Mock()
         mock_controller.inference_wrapped_model.model.config = Mock()
         mock_controller.inference_wrapped_model.model.config.cuda_graph_impl = "none"
-        
+
         engine = DynamicInferenceEngine(
             controller=mock_controller,
             context=dynamic_context,
             random_seed=123,
             inference_logging_step_interval=0,  # Disabled
         )
-        
+
         # Verify log was never called
         mock_wandb.log.assert_not_called()
 
-
     @pytest.mark.internal
     def test_paused_requests_in_stats(self):
         """Test that paused requests are correctly reflected in stats."""
@@ -230,7 +243,7 @@ def test_paused_requests_in_stats(self):
             buffer_guaranteed_fraction=0.1,
             block_size_tokens=32,
         )
-        
+
         # Add multiple requests to potentially trigger pausing
         for i in range(5):
             dynamic_context.add_request(
@@ -240,11 +253,11 @@ def test_paused_requests_in_stats(self):
                     sampling_params=SamplingParams(num_tokens_to_generate=10),
                 )
             )
-        
+
         if dynamic_context.total_request_count > 0:
             dynamic_context.initialize_attention_state()
             stats = dynamic_context.get_kvcache_utilization_stats()
-            
+
             # Verify paused request count is included
             assert 'paused_request_count' in stats
             assert stats['paused_request_count'] >= 0
@@ -253,7 +266,7 @@ def test_paused_requests_in_stats(self):
     def test_metrics_writer_none_handling(self):
         """Test that engine handles None metrics_writer gracefully."""
         dynamic_context = self._get_dynamic_context(metrics_writer=None)
-        
+
         # Create mock controller with proper spec to pass isinstance checks
         mock_controller = create_autospec(TextGenerationController, instance=True)
         # Set up nested mock structure
@@ -261,7 +274,7 @@ def test_metrics_writer_none_handling(self):
         mock_controller.inference_wrapped_model.model = Mock()
         mock_controller.inference_wrapped_model.model.config = Mock()
         mock_controller.inference_wrapped_model.model.config.cuda_graph_impl = "none"
-        
+
         # Should not raise error even with logging interval set
         engine = DynamicInferenceEngine(
             controller=mock_controller,
@@ -269,8 +282,7 @@ def test_metrics_writer_none_handling(self):
             random_seed=123,
             inference_logging_step_interval=10,
         )
-        
+
         # Verify engine was created successfully
         assert engine.inference_logging_step_interval == 10
         assert engine.context.metrics_writer is None
-