backed out of some weird changes that claude made

Author: kyle-pena-kuzco

python/sglang/srt/layers/logits_processor.py

@@ -366,15 +366,13 @@ def forward(
             """
             verification_hidden_states_to_store: Optional[torch.Tensor] = None
             if logits_metadata.verification_algorithm.is_toploc():
+                logger.debug(
+                    f"Capturing TopLoc verification hidden states with shape {pruned_states.shape if pruned_states is not None else 'None'}"
+                )
                 verification_hidden_states_to_store = (
                     pruned_states[sample_indices] if sample_indices else pruned_states
                 )
 
-            # For TOPLOC verification algorithm, capture hidden states and generate proof
-            verification_proof: Optional[list] = None
-            if logits_metadata.verification_algorithm.is_toploc():
-                verification_proof = self.generate_verification_proof(hidden_states)
-
         if not logits_metadata.extend_return_logprob:
             # Decode mode or extend mode without return_logprob.
             return LogitsProcessorOutput(
@@ -573,21 +571,6 @@ def compute_temp_top_p_normalized_logprobs(
         else:
             return torch.nn.functional.log_softmax(last_logits, dim=-1)
 
-    def generate_verification_proof(self, hidden_states: torch.Tensor) -> list:
-        """Generate a verification proof from hidden states.
-
-        The proof is a fingerprint or hash-like representation of the hidden states.
-        In this implementation, we use a simple mean of the hidden states as a proof,
-        but more sophisticated methods could be implemented.
-
-        Args:
-            hidden_states: The hidden states to generate proof from
-
-        Returns:
-            A list representation of the proof
-        """
-        return []
-
 
 @triton.jit
 def fused_softcap_kernel(

python/sglang/srt/managers/scheduler_output_processor_mixin.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import logging
 from typing import TYPE_CHECKING, List, Optional, Tuple, Union
 
 from sglang.srt.layers.logits_processor import LogitsProcessorOutput
@@ -14,6 +15,8 @@
         ScheduleBatch,
     )
 
+logger = logging.getLogger(__name__)
+
 
 class SchedulerOutputProcessorMixin:
     """
@@ -120,6 +123,9 @@ def process_batch_result_prefill(
                         and logits_output.verification_hidden_states is not None
                     ):
                         # Process verification hidden states for the current request
+                        logger.debug(
+                            f"Processing verification hidden states for prefill in req {req.req_id}"
+                        )
                         req.verification_proofs.append(
                             create_toploc_proofs(
                                 logits_output.verification_hidden_states[
@@ -132,6 +138,9 @@ def process_batch_result_prefill(
                                 .clone()
                             )
                         )
+                        logger.debug(
+                            f"Added verification proof #{len(req.verification_proofs)} to req {req.req_id} (prefill)"
+                        )
 
                     if req.grammar is not None:
                         req.grammar.accept_token(next_token_id)
@@ -270,11 +279,17 @@ def process_batch_result_decode(
                 )
 
             if logits_output.verification_hidden_states is not None:
+                logger.debug(
+                    f"Processing verification hidden states for decode in req {req.req_id}"
+                )
                 req.verification_proofs.append(
                     create_toploc_proofs(
                         logits_output.verification_hidden_states[i].cpu().clone()
                     )
                 )
+                logger.debug(
+                    f"Added verification proof #{len(req.verification_proofs)} to req {req.req_id} (decode)"
+                )
 
             if req.grammar is not None and batch.spec_algorithm.is_none():
                 req.grammar.accept_token(next_token_id)
@@ -589,6 +604,9 @@ def stream_output_generation(
                 if req.return_verification_proofs and hasattr(
                     req, "verification_proofs"
                 ):
+                    logger.debug(
+                        f"Collecting verification proofs for req {req.req_id}: {len(req.verification_proofs) if req.verification_proofs else 0} proofs"
+                    )
                     if verification_proofs is None:
                         verification_proofs = []
                     verification_proofs.append(req.verification_proofs)

python/sglang/srt/model_executor/model_runner.py

@@ -18,14 +18,12 @@
 import json
 import logging
 import os
-import threading
 import time
 from dataclasses import dataclass
 from typing import List, Optional, Tuple, Union
 
 import torch
 import torch.distributed as dist
-from toploc import build_proofs_base64
 
 from sglang.srt.configs.device_config import DeviceConfig
 from sglang.srt.configs.load_config import LoadConfig
@@ -57,7 +55,7 @@
 )
 from sglang.srt.mem_cache.paged_allocator import PagedTokenToKVPoolAllocator
 from sglang.srt.model_executor.cuda_graph_runner import CudaGraphRunner
-from sglang.srt.model_executor.forward_batch_info import CaptureHiddenMode, ForwardBatch
+from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 from sglang.srt.model_loader import get_model
 from sglang.srt.model_loader.loader import (
     DefaultModelLoader,
@@ -125,19 +123,6 @@ def __init__(
         self.req_to_token_pool = req_to_token_pool
         self.token_to_kv_pool_allocator = token_to_kv_pool_allocator
 
-        # Activation saving setup
-        self.save_activations = server_args.toploc_fingerprint
-        if self.save_activations:
-            self.capture_hidden_mode = (
-                CaptureHiddenMode.LAST
-            )  # Only capture final hidden state
-            self.verification_algorithm = (
-                VerificationAlgorithm.TOPLOC
-            )  # Set verification algorithm
-            self.is_cuda_graph_capturing = (
-                False  # Flag to track CUDA graph capturing state
-            )
-
         # Model-specific adjustment
         self.model_specific_adjustment()
 
@@ -921,9 +906,7 @@ def init_cuda_graphs(self):
         logger.info(
             f"Capture cuda graph begin. This can take up to several minutes. avail mem={before_mem:.2f} GB"
         )
-
         self.cuda_graph_runner = CudaGraphRunner(self)
-
         after_mem = get_available_gpu_memory(self.device, self.gpu_id)
         logger.info(
             f"Capture cuda graph end. Time elapsed: {time.time() - tic:.2f} s. "
@@ -978,8 +961,6 @@ def forward_idle(self, forward_batch: ForwardBatch):
     def forward(
         self, forward_batch: ForwardBatch, skip_attn_backend_init: bool = False
     ) -> LogitsProcessorOutput:
-        """Run the forward pass."""
-        # Run cuda graph if possible
         if (
             forward_batch.forward_mode.is_cuda_graph()
             and self.cuda_graph_runner

python/sglang/srt/openai_api/adapter.py

@@ -732,6 +732,32 @@ async def generate_stream_resp():
                     prompt_tokens[index] = content["meta_info"]["prompt_tokens"]
                     completion_tokens[index] = content["meta_info"]["completion_tokens"]
 
+                    if not stream_buffer:  # The first chunk
+                        if request.echo:
+                            if isinstance(request.prompt, str):
+                                # for the case of single str prompts
+                                prompts = request.prompt
+                            elif isinstance(request.prompt, list):
+                                if isinstance(request.prompt[0], str):
+                                    # for the case of multiple str prompts
+                                    prompts = request.prompt[index // request.n]
+                                elif isinstance(request.prompt[0], int):
+                                    # for the case of single token ids prompt
+                                    prompts = tokenizer_manager.tokenizer.decode(
+                                        request.prompt, skip_special_tokens=True
+                                    )
+                                elif isinstance(request.prompt[0], list) and isinstance(
+                                    request.prompt[0][0], int
+                                ):
+                                    # for the case of multiple token ids prompts
+                                    prompts = tokenizer_manager.tokenizer.decode(
+                                        request.prompt[index // request.n],
+                                        skip_special_tokens=True,
+                                    )
+
+                            # Prepend prompt in response text.
+                            text = prompts + text
+
                     if request.logprobs is not None:
                         # The first chunk and echo is enabled.
                         if not stream_buffer and request.echo:
@@ -1070,6 +1096,9 @@ def v1_chat_generate_response(
 
         finish_reason = ret_item["meta_info"]["finish_reason"]
 
+        tool_calls = None
+        text = ret_item["text"]
+
         if isinstance(request, list):
             tool_choice = request[idx].tool_choice
             tools = request[idx].tools
@@ -1084,7 +1113,7 @@ def v1_chat_generate_response(
                 parser = ReasoningParser(
                     model_type=reasoning_parser, stream_reasoning=False
                 )
-                reasoning_text, text = parser.parse_non_stream(ret_item["text"])
+                reasoning_text, text = parser.parse_non_stream(text)
             except Exception as e:
                 logger.error(f"Exception: {e}")
                 return create_error_response(
@@ -1093,10 +1122,8 @@ def v1_chat_generate_response(
                 )
         else:
             reasoning_text = None
-            text = ret_item["text"]
 
-        tool_calls = None
-        if tool_call_parser and tool_choice != "none" and tools:
+        if tool_choice != "none" and tools:
             parser = FunctionCallParser(tools, tool_call_parser)
             if parser.has_tool_call(text):
                 if finish_reason["type"] == "stop":
@@ -1122,6 +1149,12 @@ def v1_chat_generate_response(
 
         # Extract verification proofs if available
         verification_proofs = ret_item["meta_info"].get("verification_proofs", None)
+        if verification_proofs:
+            logger.debug(
+                f"Retrieved verification proofs from response: {len(verification_proofs)} proof sets"
+            )
+        else:
+            logger.debug("No verification proofs found in response")
 
         if to_file:
             # to make the choice data json serializable

python/sglang/srt/server_args.py

@@ -1134,6 +1134,10 @@ def prepare_server_args(argv: List[str]) -> ServerArgs:
     ServerArgs.add_cli_args(parser)
     raw_args = parser.parse_args(argv)
     server_args = ServerArgs.from_cli_args(raw_args)
+    if server_args.toploc_fingerprint:
+        logger.info(
+            f"TopLoc fingerprint verification enabled with topk={server_args.toploc_verification_topk}"
+        )
     return server_args
 
 

python/sglang/srt/verification/verification_utils.py

@@ -1,10 +1,13 @@
+import logging
 from typing import Optional
 
 import torch
 from toploc import build_proofs_base64
 
 from sglang.srt.managers.schedule_batch import global_server_args_dict
 
+logger = logging.getLogger(__name__)
+
 
 def create_toploc_proofs(
     verification_hidden_states: Optional[torch.Tensor],
@@ -18,16 +21,34 @@ def create_toploc_proofs(
     Returns:
         The hidden states tensor moved to CPU or None if input was None
     """
+    if verification_hidden_states is None:
+        logger.warning(
+            "Attempted to create TopLoc proofs with None verification_hidden_states"
+        )
+        return None
+
+    logger.debug(
+        f"Creating TopLoc proofs from tensor with shape {verification_hidden_states.shape}"
+    )
 
     # Move to CPU . Will have size [N,hidden] - each one should represent a "last token"
     verification_hidden_states = verification_hidden_states.detach().cpu()
 
     topk = global_server_args_dict["toploc_verification_topk"]
+    logger.debug(f"Using TopLoc verification topk={topk}")
 
     # Will return N proofs
-    return build_proofs_base64(
-        verification_hidden_states,
-        decode_batching_size=3,
-        topk=topk,
-        skip_prefill=False,
-    )
+    try:
+        proofs = build_proofs_base64(
+            verification_hidden_states,
+            decode_batching_size=3,
+            topk=topk,
+            skip_prefill=False,
+        )
+        logger.debug(
+            f"Successfully generated {len(proofs) if proofs else 0} TopLoc proofs"
+        )
+        return proofs
+    except Exception as e:
+        logger.error(f"Error generating TopLoc proofs: {str(e)}")
+        return None