[Bugfix] Fix hidden state size in EAGLE PD disaggregation buffers

[michelemarzollo]

Motivation

In PD disaggregation currently the buffers to move the draft model's hidden states are allocated based on the target model's hidden state size. This works only if target and draft model have the same hidden state size.
While LLama3-8B the sizes match (so the bug is invisible), if I try to run https://huggingface.co/meta-llama/Llama-3.3-70B-Instruct with https://huggingface.co/yuhuili/EAGLE3-LLaMA3.3-Instruct-70B I get the following output:

[2025-11-19 14:49:50 TP0] Scheduler hit an exception: Traceback (most recent call last):
  File "/sgl-workspace/sglang/python/sglang/srt/managers/scheduler.py", line 2746, in run_scheduler_process
    scheduler.event_loop_normal_disagg_prefill()
  File "/usr/local/lib/python3.12/dist-packages/torch/utils/_contextlib.py", line 120, in decorate_context
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/sgl-workspace/sglang/python/sglang/srt/disaggregation/prefill.py", line 343, in event_loop_normal_disagg_prefill
    self.process_batch_result_disagg_prefill(batch, result)
  File "/sgl-workspace/sglang/python/sglang/srt/disaggregation/prefill.py", line 460, in process_batch_result_disagg_prefill
    self.send_kv_chunk(req, last_chunk=True)
  File "/sgl-workspace/sglang/python/sglang/srt/disaggregation/prefill.py", line 652, in send_kv_chunk
    self.disagg_metadata_buffers.set_buf(req)
  File "/sgl-workspace/sglang/python/sglang/srt/disaggregation/utils.py", line 222, in set_buf
    self.output_hidden_states[req.metadata_buffer_index].copy_(
RuntimeError: The size of tensor a (8192) must match the size of tensor b (6144) at non-singleton dimension 0

[2025-11-19 14:49:50] SIGQUIT received. signum=None, frame=None. It usually means one child failed.
[2025-11-19 14:49:50 TP1] Scheduler hit an exception: Traceback (most recent call last):
  File "/sgl-workspace/sglang/python/sglang/srt/managers/scheduler.py", line 2746, in run_scheduler_process
    scheduler.event_loop_normal_disagg_prefill()
  File "/usr/local/lib/python3.12/dist-packages/torch/utils/_contextlib.py", line 120, in decorate_context
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/sgl-workspace/sglang/python/sglang/srt/disaggregation/prefill.py", line 343, in event_loop_normal_disagg_prefill
    self.process_batch_result_disagg_prefill(batch, result)
  File "/sgl-workspace/sglang/python/sglang/srt/disaggregation/prefill.py", line 460, in process_batch_result_disagg_prefill
    self.send_kv_chunk(req, last_chunk=True)
  File "/sgl-workspace/sglang/python/sglang/srt/disaggregation/prefill.py", line 652, in send_kv_chunk
    self.disagg_metadata_buffers.set_buf(req)
  File "/sgl-workspace/sglang/python/sglang/srt/disaggregation/utils.py", line 222, in set_buf
    self.output_hidden_states[req.metadata_buffer_index].copy_(
RuntimeError: The size of tensor a (8192) must match the size of tensor b (6144) at non-singleton dimension 0

Where 8192 is the target model's hidden size, while 6144 is the draft model's hidden size.

Modifications

I just pick the hidden size from the draft models' config. Not sure if the draft model config's should be held by the scheduler as I did it. I am open to suggestions!

I also changed the MetadataBuffers to avoid using them when they are not needed. In my understanding they are only needed for eagle, but I am not familiar with this part of the codebase. Please, have a look! (I tested with NGRAM and it works, while I am not able to run the Standalone specualation.

Accuracy Tests

No need. Before it was crashing, now it works.

Benchmarking and Profiling

No need. Before it was crashing, now it works.
But to reproduce, I am running the following script to start the server:

#!/usr/bin/env bash
set -e


# EAGLE3

CUDA_VISIBLE_DEVICES=0,1 python3 -m sglang.launch_server \
    --model /storage/datasets/huggingface/models/Llama-3.3-70B-Instruct  \
    --cuda-graph-max-bs 64 \
    --max-running-requests 64 \
    --speculative-algorithm EAGLE3 \
    --speculative-draft-model-path /storage/datasets/huggingface/models/sglang-EAGLE3-LLaMA3.3-Instruct-70B \
    --disaggregation-mode prefill \
    --disaggregation-transfer-backend nixl \
    --tp 2 &


CUDA_VISIBLE_DEVICES=2,3 python3 -m sglang.launch_server \
    --model /storage/datasets/huggingface/models/Llama-3.3-70B-Instruct  \
    --cuda-graph-max-bs 2 \
    --max-running-requests 2 \
    --speculative-algorithm EAGLE3 \
    --speculative-draft-model-path /storage/datasets/huggingface/models/sglang-EAGLE3-LLaMA3.3-Instruct-70B \
    --disaggregation-mode decode \
    --port 30001 \
    --disaggregation-transfer-backend nixl \
    --tp 2 &


# TO RUN FOR ALL

python3 -m sglang_router.launch_router \
    --pd-disaggregation \
    --prefill http://127.0.0.1:30000 \
    --decode http://127.0.0.1:30001 \
    --host 0.0.0.0 \
    --port 8000

# Done when it says "Starting server on 0.0.0.0:8000"

Once the server is running, you can test with

python3 -m sglang.bench_serving \
    --model /storage/datasets/huggingface/models/Llama-3.3-70B-Instruct  \
    --base-url "http://127.0.0.1:8000" \
    --dataset-name "sharegpt" \
    --num-prompts 10 \
    --sharegpt-output-len 100 \
    --disable-ignore-eos \
    --max-concurrency 4 \
    --pd-separated

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.
• Work with maintainers to merge your PR. See the PR Merge Process

[gemini-code-assist]

Summary of Changes

Hello @michelemarzollo, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request resolves a critical bug in the PD disaggregation mechanism, specifically within the EAGLE/EAGLE3 speculative decoding algorithms. Previously, the system incorrectly allocated hidden state buffers based on the target model's configuration, leading to crashes when the draft and target models had differing hidden state dimensions. The fix ensures that these buffers are now correctly sized according to the draft model's specifications, enabling robust operation with diverse model architectures.

Highlights

• Bug Fix for Hidden State Size Mismatch: Addresses a RuntimeError occurring in PD disaggregation when the target and draft models have different hidden state sizes, specifically when allocating buffers for the draft model's hidden states. This previously led to crashes.
• Corrected Buffer Allocation Logic: Modifies the buffer allocation for hidden states in PD disaggregation to correctly use the hidden_size and dtype from the draft model's configuration (drafter_config) instead of the target model's configuration (model_config).
• Introduction of drafter_config: A new drafter_config attribute is introduced in the scheduler to store the ModelConfig of the speculative draft model, ensuring its properties are correctly referenced for operations like buffer allocation.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request correctly fixes a bug where the disaggregation buffers for EAGLE were allocated using the target model's hidden state size, causing a crash when it differed from the draft model's size. The fix involves introducing drafter_config to hold the draft model's configuration and using it to get the correct hidden state size.

Your approach of storing drafter_config on the scheduler is reasonable. It's the central component for managing configurations, so it's a logical place for it.

I have a few suggestions to improve the robustness and correctness of the implementation:

1. The initialization of drafter_config should be more comprehensive to include other server arguments, not just the model path.
2. The usage of drafter_config in init_disaggregation should be safer to handle cases where disaggregation might be used with non-EAGLE algorithms, preventing potential crashes.

Details are in the line comments. Overall, this is a good fix for the reported issue.

[gemini-code-assist]

The ModelConfig for the draft model is initialized with only the model path. This is incorrect as it misses other important configurations from server_args like trust_remote_code, revision, dtype, quantization, etc. This can lead to incorrect behavior or crashes. You should use ModelConfig.from_server_args to ensure all relevant arguments are passed, similar to how it's done for the draft worker in TpModelWorker.

Suggested change

	self.drafter_config = ModelConfig(
	self.server_args.speculative_draft_model_path
	)
	self.drafter_config = ModelConfig.from_server_args(
	self.server_args,
	model_path=self.server_args.speculative_draft_model_path,
	model_revision=self.server_args.speculative_draft_model_revision,
	is_draft_model=True,
	override_config_file=self.server_args.decrypted_draft_config_file,
	)

[michelemarzollo]

Done

[gemini-code-assist]

This will raise an AttributeError if disaggregation is enabled for a non-EAGLE speculative algorithm, because self.drafter_config would be None. To make this more robust, you should fall back to self.model_config if self.drafter_config is not available. This maintains the previous behavior for non-EAGLE cases while fixing the bug for EAGLE.

Suggested change

	hidden_size=self.drafter_config.hf_text_config.hidden_size,
	hidden_states_dtype=self.drafter_config.dtype,
	hidden_size=(self.drafter_config or self.model_config).hf_text_config.hidden_size,
	hidden_states_dtype=(self.drafter_config or self.model_config).dtype,

[michelemarzollo]

No need to pass some dimension, if we don't have any states to move: let's pass 0.

[gemini-code-assist]

Similar to the issue in the DisaggregationMode.DECODE branch, this will crash if disaggregation is enabled for a non-EAGLE speculative algorithm. You should fall back to self.model_config to prevent this.

Suggested change

	hidden_size=self.drafter_config.hf_text_config.hidden_size,
	hidden_states_dtype=self.drafter_config.dtype,
	hidden_size=(self.drafter_config or self.model_config).hf_text_config.hidden_size,
	hidden_states_dtype=(self.drafter_config or self.model_config).dtype,

[michelemarzollo]

Same as above.

[ZeldaHuang]

hi @michelemarzollo , you can try to change

sglang/python/sglang/srt/managers/scheduler.py

Line 867 in 7dcf910

hidden_size=self.model_config.hf_text_config.hidden_size,

and

sglang/python/sglang/srt/managers/scheduler.py

Line 912 in 7dcf910

hidden_size=self.model_config.hf_text_config.hidden_size,

to self.draft_worker.model_config.hidden_size when use eagle3

[ShangmingCai]

@michelemarzollo Thx for the bug report and the PR, will review this PR today.

[ShangmingCai]

Do these changes affect the correctness of get_buf_infos, get_buf as well? I see you only modified the set_buf interface.

[ShangmingCai]

hi @michelemarzollo , you can try to change

sglang/python/sglang/srt/managers/scheduler.py

Line 867 in 7dcf910

hidden_size=self.model_config.hf_text_config.hidden_size,

and

sglang/python/sglang/srt/managers/scheduler.py

Line 912 in 7dcf910

hidden_size=self.model_config.hf_text_config.hidden_size,

to self.draft_worker.model_config.hidden_size when use eagle3

@michelemarzollo Can you try @ZeldaHuang 's method? I think we should minimize the changes. We don't need to filter out the hidden state since the metadata buffer slot is way smaller compared to the kvcache, actually. But anyway, feel free to do some experiments to verify the overhead. It would be great if we could keep the code clean and reduce the unnecessary overhead.

[michelemarzollo]

Thanks a lot to both of you for the comments! Agree with all of them. Now I get the hidden state size from the drafter, and otherwise pass 0 (I don't see any reason for passing the model's hidden state, I find it even misleading). I don't modify anymore the utils file (very little overhead indeed, especially because the size is 0 if not eagle). Let me know if there is still something you don't like!
I ran a throughput test for eagle, ngram and autoregressive and things look as expected.

[ShangmingCai]

Thanks a lot to both of you for the comments! Agree with all of them. Now I get the hidden state size from the drafter, and otherwise pass 0 (I don't see any reason for passing the model's hidden state, I find it even misleading). I don't modify anymore the utils file (very little overhead indeed, especially because the size is 0 if not eagle). Let me know if there is still something you don't like! I ran a throughput test for eagle, ngram and autoregressive and things look as expected.

I am not quite sure whether pass 0 is safe for all scenes, maybe a padding value like 64 (RDMA granularity) sounds safer to me.

[michelemarzollo]

Sure, 64 looks good to me. I don't have such an understanding to know if there could be any issues, so I trust you. I can even leave the target's hidden size if you prefer, I just couldn't get why :)

[ShangmingCai]

/tag-and-rerun-ci

[ShangmingCai]

The only failed CI test is irrelevant:

This PR LGTM, and I think it is ready-to-merge now. I will bypass CI after a small comment fix.

[ShangmingCai]

Minor modifications:

• Since float32 is 4 Bytes, so padding the unused hidden size to 16 is enough.
• Add @ZeldaHuang as co-author.