Shanmugamr/fix logprobs stopwords

megatron/core/inference/batch_dimensions_utils.py

@@ -73,7 +73,9 @@ def is_applicable_for_batch_dim(
                 >= real_batch_dim.prefill_req_count + real_batch_dim.decode_req_count
             )
 
-    def is_valid(self, max_requests: int, max_sequence_length: int) -> bool:
+    def is_valid(
+        self, max_requests: int, max_sequence_length: int, num_speculative_tokens: int
+    ) -> bool:
         """
         Checks if the batch dimension is valid based on resource constraints.
 
@@ -92,11 +94,17 @@ def is_valid(self, max_requests: int, max_sequence_length: int) -> bool:
             return False
 
         # Check if token count is sufficient for requests
-        if self.token_count < self.prefill_req_count + self.decode_req_count:
+        if self.token_count < self.prefill_req_count + self.decode_req_count * (
+            num_speculative_tokens + 1
+        ):
             return False
 
         # Check if the prefill requests are shorter than the max sequence length
-        if self.token_count > self.prefill_req_count * max_sequence_length + self.decode_req_count:
+        if (
+            self.token_count
+            > self.prefill_req_count * max_sequence_length
+            + self.decode_req_count * (num_speculative_tokens + 1)
+        ):
             return False
 
         return True
@@ -308,6 +316,7 @@ def generate_cuda_graph_batch_dimensions_list(
         max_tokens: int,
         max_sequence_length: int,
         use_cuda_graphs_for_non_decode_steps: bool,
+        num_speculative_tokens: int = 0,
     ) -> Tuple[List[InferenceBatchDimensions], Optional[List[int]]]:
         """
         Generate CUDA graph batch dimensions.
@@ -344,6 +353,7 @@ def generate_cuda_graph_batch_dimensions_list(
             max_tokens: Maximum total tokens
             max_sequence_length: Maximum sequence length
             use_cuda_graphs_for_non_decode_steps: Whether to use CUDA graphs for non-decode steps
+            num_speculative_tokens: Number of speculative tokens
 
         Returns:
             Tuple containing:
@@ -355,7 +365,7 @@ def generate_cuda_graph_batch_dimensions_list(
         def add_if_valid(token_count: int, prefill_req_count: int, decode_req_count: int) -> None:
             """Helper to create and append batch dimension to list only if it's valid."""
             batch_dim = InferenceBatchDimensions(token_count, prefill_req_count, decode_req_count)
-            if batch_dim.is_valid(max_requests, max_sequence_length):
+            if batch_dim.is_valid(max_requests, max_sequence_length, num_speculative_tokens):
                 cuda_graph_batch_dimensions_list.append(batch_dim)
 
         # Cuda graph token-counts
@@ -372,9 +382,10 @@ def add_if_valid(token_count: int, prefill_req_count: int, decode_req_count: int
             ):
                 cuda_graph_max_tokens = max_tokens
 
-            assert cuda_graph_max_tokens == max_requests, (
-                f"cuda_graph_max_tokens ({cuda_graph_max_tokens}) must equal max_requests "
-                f"({max_requests}). This is required for correctly syncing EP ranks: "
+            assert cuda_graph_max_tokens == max_requests * (num_speculative_tokens + 1), (
+                f"cuda_graph_max_tokens ({cuda_graph_max_tokens}) must equal max_requests *"
+                f"(num_speculative_tokens + 1) ({max_requests * (num_speculative_tokens + 1)}). "
+                "This is required for correctly syncing EP ranks: "
                 f"prefill and decode graph pools must have the same token count granularity."
             )
 
@@ -395,8 +406,9 @@ def add_if_valid(token_count: int, prefill_req_count: int, decode_req_count: int
             )
 
             # Calculate separate token counts for decode-only graphs.
-            # Decode graphs can be more conservative since each request uses exactly 1 token.
-            cuda_graph_max_tokens_decode = min(cuda_graph_max_tokens, max_requests)
+            cuda_graph_max_tokens_decode = min(
+                cuda_graph_max_tokens, max_requests * (num_speculative_tokens + 1)
+            )
             cuda_graph_decode_token_counts = (
                 CUDAGraphBatchDimensionBuilder._calculate_cuda_graph_token_counts(
                     tp_size=tp_size,
@@ -415,20 +427,29 @@ def add_if_valid(token_count: int, prefill_req_count: int, decode_req_count: int
         ):  # decode only
             # Use decode-specific token counts for decode-only graphs
             for size in cuda_graph_decode_token_counts:
+                decode_req_count = min(size // (num_speculative_tokens + 1), max_requests)
+                token_count = decode_req_count * (num_speculative_tokens + 1)
+                token_count = token_count // tp_size * tp_size
                 add_if_valid(
-                    token_count=min(size, max_requests),
-                    prefill_req_count=0,
-                    decode_req_count=min(size, max_requests),
+                    token_count=token_count, prefill_req_count=0, decode_req_count=decode_req_count
                 )
         else:
             # Mixed prefill and decode mode
             # Create prefill and mixed dimensions with full token counts
             for size in cuda_graph_prefill_token_counts:
+                assert size % tp_size == 0
+                prefill_req_count = min(cuda_graph_mixed_prefill_request_count, max_requests)
+                decode_req_count = max(
+                    0,
+                    min(
+                        (size - prefill_req_count) // (num_speculative_tokens + 1),
+                        max_requests - prefill_req_count,
+                    ),
+                )
                 add_if_valid(
                     token_count=size,
-                    prefill_req_count=min(cuda_graph_mixed_prefill_request_count, max_requests),
-                    decode_req_count=min(size, max_requests)
-                    - min(cuda_graph_mixed_prefill_request_count, max_requests),
+                    prefill_req_count=prefill_req_count,
+                    decode_req_count=decode_req_count,
                 )
                 # We need to ensure the prefill requests are shorter than the max sequence length,
                 # considering the one decode token is used for prefill request construction
@@ -445,16 +466,21 @@ def add_if_valid(token_count: int, prefill_req_count: int, decode_req_count: int
 
             # Create decode-only dimensions with optimized token counts
             for size in cuda_graph_decode_token_counts:
+                decode_req_count = min(size // (num_speculative_tokens + 1), max_requests)
+                token_count = decode_req_count * (num_speculative_tokens + 1)
+                token_count = token_count // tp_size * tp_size
                 add_if_valid(
-                    token_count=min(size, max_requests),
-                    prefill_req_count=0,
-                    decode_req_count=min(size, max_requests),
+                    token_count=token_count, prefill_req_count=0, decode_req_count=decode_req_count
                 )
 
         # Remove duplicates and sort by prefill token count
         cuda_graph_batch_dimensions_list = list(set(cuda_graph_batch_dimensions_list))
         cuda_graph_batch_dimensions_list.sort(
-            key=lambda x: ((x.token_count - x.decode_req_count), x.decode_req_count), reverse=True
+            key=lambda x: (
+                (x.token_count - x.decode_req_count * (num_speculative_tokens + 1)),
+                x.decode_req_count,
+            ),
+            reverse=True,
         )
 
         # Collect actual token counts from batch dimensions, then unique and sort

megatron/core/inference/config.py

@@ -228,6 +228,9 @@ class InferenceConfig:
     enable_chunked_prefill: bool = False
     """Whether to enable chunked prefill."""
 
+    num_speculative_tokens: int = 0
+    """The number of speculative tokens to generate for decode steps."""
+
     enable_prefix_caching: bool = False
     """Whether to enable prefix caching for KV cache block sharing."""
 

megatron/core/inference/contexts/attention_context/mamba_metadata.py

@@ -259,12 +259,17 @@ def update(
             self.cu_seqlens = self._cu_seqlens_buffer[: padded_prefill_count + 1]
 
         if padded_decode_count > 0 and padded_prefill_count > 0:
-            self._device_decode_prefill_buffer[0] = real_decode_count
+            self._device_decode_prefill_buffer[0] = cu_seqlens[real_decode_count]
             # This describes the number of items in the prefill tensor relative to the
             # decode tensor. If chunked prefill is present, it is included in the
             # "prefill" part of the main split.
-            self._device_decode_prefill_buffer[1] = regular_prefill_count + (
-                1 if has_chunked_prefill_req else 0
+            self._device_decode_prefill_buffer[1] = (
+                cu_seqlens[
+                    real_decode_count
+                    + regular_prefill_count
+                    + (1 if has_chunked_prefill_req else 0)
+                ]
+                - cu_seqlens[real_decode_count]
             )
             self.device_decode_prefill = self._device_decode_prefill_buffer
 

megatron/core/inference/contexts/attention_context/mha_metadata.py

@@ -41,6 +41,7 @@ def update(
         request_to_kv_block_ids: torch.Tensor,
         batch_dimensions: InferenceBatchDimensions,
         padded_batch_dimensions: InferenceBatchDimensions,
+        num_speculative_tokens: int = 0,
     ):
         """
         Args:
@@ -49,6 +50,7 @@ def update(
             request_to_kv_block_ids: (>real_batch_size, max_kv_blocks)
             batch_dimensions: Configuration object containing real batch settings
             padded_batch_dimensions: Configuration object containing padded batch settings
+            num_speculative_tokens: Number of speculative tokens
         """
         # Extract values from configs
         real_batch_size = batch_dimensions.req_count
@@ -99,7 +101,7 @@ def update(
         )
 
         if padded_batch_dimensions.prefill_req_count == 0:
-            self._max_seqlen_q = 1
+            self._max_seqlen_q = num_speculative_tokens + 1
         else:
             # Make sure we will launch the prefill kernel for prefill graphs
             self._max_seqlen_q = max(2, padded_batch_dimensions.token_count)
@@ -150,6 +152,7 @@ def update(
         request_to_kv_block_ids: torch.Tensor,
         batch_dimensions: InferenceBatchDimensions,
         padded_batch_dimensions: InferenceBatchDimensions,
+        num_speculative_tokens: int = 0,
     ):
         """
         Args:
@@ -158,13 +161,15 @@ def update(
             request_to_kv_block_ids: (>real_batch_size, max_kv_blocks)
             batch_dimensions: Configuration object containing real batch settings
             padded_batch_dimensions: Configuration object containing padded batch settings
+            num_speculative_tokens: Number of speculative tokens
         """
         super().update(
             request_query_lengths,
             request_kv_length_offsets,
             request_to_kv_block_ids,
             batch_dimensions,
             padded_batch_dimensions,
+            num_speculative_tokens,
         )
 
     def reset(self):
@@ -183,6 +188,7 @@ def update(
         request_to_kv_block_ids: torch.Tensor,
         batch_dimensions: InferenceBatchDimensions,
         padded_batch_dimensions: InferenceBatchDimensions,
+        num_speculative_tokens: int = 0,
     ):
         """
         Args:
@@ -191,17 +197,19 @@ def update(
             request_to_kv_block_ids: (>real_batch_size, max_kv_blocks)
             batch_dimensions: Configuration object containing real batch settings
             padded_batch_dimensions: Configuration object containing padded batch settings
+            num_speculative_tokens: Number of speculative tokens
         """
         super().update(
             request_query_lengths,
             request_kv_length_offsets,
             request_to_kv_block_ids,
             batch_dimensions,
             padded_batch_dimensions,
+            num_speculative_tokens,
         )
         if len(self.state_data["query_lengths"]) > 0:
             self.state_data["max_seqlen_q"] = torch.max(self.state_data["query_lengths"]).item()
             self.state_data["max_seqlen_k"] = torch.max(self.state_data["kv_seq_lengths"]).item()
         else:
-            self.state_data["max_seqlen_q"] = 1
+            self.state_data["max_seqlen_q"] = num_speculative_tokens + 1
             self.state_data["max_seqlen_k"] = 1

megatron/core/inference/contexts/dynamic_block_allocator.py

@@ -85,19 +85,39 @@ def get_total_used(self):
 
     def get_active_used(self):
         """Compute number of active blocks used."""
-        return (
-            self.context.request_kv_block_counts[
-                self.context.paused_request_count : self.context.total_request_count
-            ]
-            .sum()
-            .item()
-        )
+        if not self.enable_prefix_caching:
+            return (
+                self.context.request_kv_block_counts[
+                    self.context.paused_request_count : self.context.total_request_count
+                ]
+                .sum()
+                .item()
+            )
+
+        active_start = self.context.paused_request_count
+        active_end = self.context.total_request_count
+        if active_end > active_start:
+            active_rows = self.context.request_to_kv_block_ids[active_start:active_end]
+            valid_ids = active_rows[active_rows >= 0]
+            if valid_ids.numel() > 0:
+                return int(torch.unique(valid_ids).numel())
+        return 0
 
     def get_paused_used(self):
         """Compute number of paused blocks used."""
-        return (
-            self.context.request_kv_block_counts[: self.context.paused_request_count].sum().item()
-        )
+        if not self.enable_prefix_caching:
+            return (
+                self.context.request_kv_block_counts[: self.context.paused_request_count]
+                .sum()
+                .item()
+            )
+
+        if self.context.paused_request_count > 0:
+            paused_rows = self.context.request_to_kv_block_ids[: self.context.paused_request_count]
+            valid_ids = paused_rows[paused_rows >= 0]
+            if valid_ids.numel() > 0:
+                return int(torch.unique(valid_ids).numel())
+        return 0
 
     def get_active_avail(self):
         """Compute number of active blocks available."""