[FIX] fix wrong indexing for hidden states when prefix-cache take effect

[Dutch-voyage]

Motivation

return_hidden_states=True does not return prefilled hidden_states as expected.
see relevant issue #4997

Modifications

In scheduler_output_processor_mixin.py:116
change

req.hidden_states.append(
    logits_output.hidden_states[
        hidden_state_offset : (
            hidden_state_offset := hidden_state_offset
            + len(req.origin_input_ids)
        )
    ]
    .cpu()
    .clone()
    .tolist()
)

to

req.hidden_states.append(
    logits_output.hidden_states[
        hidden_state_offset : (
            hidden_state_offset := hidden_state_offset
            + req.extend_input_len
        )
    ]
    .cpu()
    .clone()
    .tolist()
)

Checklist

• Format your code according to the Code Formatting with Pre-Commit.
• Add unit tests as outlined in the Running Unit Tests.
• Update documentation / docstrings / example tutorials as needed, according to Writing Documentation.
• Provide throughput / latency benchmark results and accuracy evaluation results as needed, according to Benchmark and Profiling and Accuracy Results.
• For reviewers: If you haven't made any contributions to this PR and are only assisting with merging the main branch, please remove yourself as a co-author when merging the PR.
• Please feel free to join our Slack channel at https://slack.sglang.ai to discuss your PR.

[xiangchensong]

The proposed modification may still face the same issue. After patching the code, if you try to pass the same query multiple times, the hidden states shape is still not correct.

Try the following code and send the same request multiple times, and the prefilled hidden_states shape will become [1, llm_dim].

python -m sglang.launch_server \
  --model-path Qwen/QwQ-32B \
  --enable-return-hidden-states \
  --host 0.0.0.0 \
  --port 30000 \
  --tp 4 

PORT = 30000
prompts = [
    "Hello, my name is",
    "The president of the United States is",
    "The capital of France is",
    "The future of AI is",
]

sampling_params = {
    "temperature": 0.6,
    "top_p": 0.95,
    "max_new_tokens": 10,
}

json_data = {
    "text": prompts,
    "sampling_params": sampling_params,
    "return_hidden_states": True,
}
for _ in range(3):
    print("Sending request to the server...")
    response = requests.post(
        f"http://localhost:{PORT}/generate",
        json=json_data,
    )
    outputs = response.json()
    for prompt, output in zip(prompts, outputs):
        for i in range(len(output["meta_info"]["hidden_states"])):
            output["meta_info"]["hidden_states"][i] = torch.tensor(
                output["meta_info"]["hidden_states"][i]
            )
        print("===============================")
        print(
            f"Prompt: {prompt}\n"
            f"Prompt_Tokens: {output['meta_info']['prompt_tokens']}\t"
            f"Completion_tokens: {output['meta_info']['completion_tokens']}"
        )
        shapes = [torch.tensor(i.shape).tolist() for i in output["meta_info"]["hidden_states"]]
        print(f"Hidden States Shape: {shapes}")