Tracing the Input: How a Single Grep Uncovered SGLang's API Contract for Tokenized Prompts

In the middle of a complex debugging session aimed at extracting hidden states from the Kimi-K2.5 model for EAGLE-3 speculative decoding training, the assistant fired off what appears to be a simple bash command — a grep for input_ids in a single SGLang source file. But this seemingly mundane query, captured in message <msg id=3377>, represents a critical pivot point in the investigation. It is the moment when the assistant stopped guessing about API behavior and started reading the source code to understand the exact contract between client and server.

The Message

ssh root@10.1.230.174 "grep -rn 'input_ids' /root/sglang/python/sglang/srt/managers/io_struct.py | head -10"
175:    # The token ids for text; one can specify either text or input_ids
176:    input_ids: Optional[Union[List[List[int]], List[int]]] = None
177:    # The embeddings for input_ids; one can specify either text or input_ids or input_embeds.
300:                       text, input_ids, input_embeds are provided)
314:            self.text is None and self.input_ids is None and self.input_embeds is None
317:            and self.input_ids is not None
321:                "Either text, input_ids or i...

The Context That Led Here

To understand why this message was written, we must trace back through the preceding conversation. The assistant had spent days building an EAGLE-3 training pipeline for the Kimi-K2.5 model — a massive Mixture-of-Experts language model running across 8 NVIDIA RTX PRO 6000 Blackwell GPUs. The core idea was to train a lightweight "drafter" model that could predict multiple tokens per step, accelerating inference through speculative decoding.

A critical prerequisite for this training was extracting intermediate hidden states from the base model at specific layers (3, 31, and 59). The assistant had developed a non-invasive server-side patch (Approach C) that captured these states during the prefill (EXTEND) phase of SGLang's forward pass and saved them as binary .pt files. This patch was verified to work with text prompts — a 24-token sentence produced correctly shaped [24, 7168] tensors.

However, the extraction script (02b_extract_hidden_states_sglang.py) was designed to send requests using raw token IDs via the prompt_token_ids field, not text strings. This was an intentional design choice: using pre-tokenized prompts avoids the overhead and potential inconsistency of server-side re-tokenization, and allows the script to work with the exact token sequences from the training dataset. When the assistant tested this approach in message <msg id=3374>, the server returned a 400 error. The prompt_token_ids field — which the assistant had assumed was a valid parameter for SGLang's OpenAI-compatible /v1/completions endpoint — was silently rejected.

The Reasoning Behind the Grep

The assistant's response to this failure reveals a methodical debugging mindset. Rather than guessing or consulting documentation, it went straight to the source code. The first step, in message <msg id=3375>, was to check whether prompt_token_ids appeared anywhere in the server's request-handling code. The grep found it only in response processing (context.py), where it reads prompt_token_ids from the engine output to populate usage statistics. This confirmed that prompt_token_ids is an output field, not an input parameter — the server returns it in the response metadata, but does not accept it in requests.

This discovery prompted the next logical question: if prompt_token_ids is not the right field name for sending raw token IDs, what is? The assistant then searched for input_ids in the OpenAI protocol definitions (<msg id=3376>), but that search was still running when the target message was issued. The grep in <msg id=3377> targets a different file — io_struct.py — which is the data structure that actually carries request information through SGLang's internal pipeline.

What the Grep Revealed

The output is illuminating. Lines 175-177 of io_struct.py define the input_ids field as an optional parameter that can be either a list of lists (for batch requests) or a flat list of integers. The comment explicitly states: "one can specify either text or input_ids." This is the definitive answer: the correct field name for sending raw token IDs is input_ids, not prompt_token_ids.

The grep also reveals the validation logic at lines 300-321, which checks that at least one of text, input_ids, or input_embeds is provided, and that they aren't all simultaneously set. This is the API contract the assistant needed to find.

Assumptions and Mistakes

This episode illuminates several assumptions that were silently guiding the assistant's work. The primary assumption was that SGLang's OpenAI-compatible API would follow the same prompt_token_ids convention used by other inference engines like vLLM. In vLLM's API, prompt_token_ids is indeed a valid request parameter — the assistant had previously used it successfully with vLLM in the same project (as seen in earlier segments). The assumption that SGLang would be API-compatible with vLLM was reasonable but incorrect for this particular field.

A secondary assumption was that the field name in the request would match the field name in the response metadata. The assistant found prompt_token_ids in the response processing code and initially suspected it might also work as a request parameter. This is a common pattern in REST APIs, but SGLang's designers chose different names for input (input_ids) and output (prompt_token_ids), breaking the symmetry.

There was also an implicit assumption that the OpenAI-compatible endpoint (/v1/completions) would accept raw token IDs at all. The grep results confirm it does — just under a different field name — but the assistant had not yet tested this when the target message was written. The 400 error with prompt_token_ids could have been interpreted as "this endpoint doesn't support raw token IDs," but the assistant correctly suspected it was a naming issue rather than a capability gap.

The Thinking Process Visible in the Reasoning

What makes this message particularly interesting is what it reveals about the assistant's debugging strategy. The progression from <msg id=3374> (the failing test) through <msg id=3375> (checking prompt_token_ids in server code) to <msg id=3376> (checking input_ids in protocol definitions) and finally <msg id=3377> (checking input_ids in io_struct.py) shows a systematic narrowing of the search space.

The assistant is effectively performing a form of "source code archaeology" — digging through layers of abstraction to find the actual data structure that carries request parameters. It starts at the API layer (OpenAI protocol definitions), then moves to the internal request processing (context.py), and finally reaches the core data structure (io_struct.py) that defines what fields are actually accepted. This is a classic debugging pattern: when a black-box API rejects your input, trace the input path through the system until you find where the validation happens.

The choice of io_struct.py is itself telling. The assistant had previously worked with this file during the hidden state dump patch development and knew it contained the request parsing logic. This prior knowledge, accumulated over dozens of earlier messages in the conversation, guided the search to the right file quickly.

Input and Output Knowledge

The input knowledge required to understand this message is substantial. One needs to know that SGLang is a serving framework for large language models, that it provides an OpenAI-compatible API, that the assistant had built a custom hidden state extraction patch, that the extraction script uses raw token IDs, and that a previous test with prompt_token_ids had failed. One also needs familiarity with Python type annotations (Optional, Union, List) to interpret the grep results.

The output knowledge created by this message is precise and actionable: the input_ids field in io_struct.py is the correct parameter for sending raw token IDs. The field accepts either List[List[int]] (for batch) or List[int] (for single requests). The validation logic ensures mutual exclusivity with text and input_embeds. This knowledge directly unblocks the extraction pipeline — the script can be updated to use input_ids instead of prompt_token_ids, and the 400 error should resolve.

Significance in the Larger Narrative

In the broader arc of the conversation, this message represents a turning point in the hidden state extraction effort. The assistant had already solved the hard problem — building a non-invasive server-side patch that captures hidden states without disrupting normal operation. The remaining challenge was purely about API compatibility: how to send requests in a format the server would accept. This grep closed that gap.

The message also exemplifies a recurring theme in the conversation: the assistant's preference for reading source code over consulting documentation or guessing. Throughout the session, whenever the assistant encounters an unexpected behavior, it immediately searches the relevant source files to understand the actual implementation. This approach is particularly valuable when working with rapidly evolving frameworks like SGLang, where documentation may lag behind the code.

The result of this investigation — the discovery that input_ids is the correct field — would be applied in subsequent messages to fix the extraction script and successfully run the 10K-sample hidden state extraction that ultimately enabled the EAGLE-3 drafter training. A single grep, in the right file at the right moment, saved hours of trial-and-error experimentation.