feat(pd_router): implement prompt-length-aware routing for pd router #1800

googs1025 · 2025-11-25T02:39:49Z

Pull Request Description

[Please provide a clear and concise description of your changes here]

Related Issues

Resolves: #1798

Important: Before submitting, please complete the description above and review the checklist below.

Contribution Guidelines (Expand for Details)

We appreciate your contribution to aibrix! To ensure a smooth review process and maintain high code quality, please adhere to the following guidelines:

Pull Request Title Format

Your PR title should start with one of these prefixes to indicate the nature of the change:

[Bug]: Corrections to existing functionality
[CI]: Changes to build process or CI pipeline
[Docs]: Updates or additions to documentation
[API]: Modifications to aibrix's API or interface
[CLI]: Changes or additions to the Command Line Interface
[Misc]: For changes not covered above (use sparingly)

Note: For changes spanning multiple categories, use multiple prefixes in order of importance.

Submission Checklist

PR title includes appropriate prefix(es)
Changes are clearly explained in the PR description
New and existing tests pass successfully
Code adheres to project style and best practices
Documentation updated to reflect changes (if applicable)
Thorough testing completed, no regressions introduced

By submitting this PR, you confirm that you've read these guidelines and your changes align with the project's contribution standards.

gemini-code-assist · 2025-11-25T02:40:09Z

Summary of Changes

Hello @googs1025, 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 significantly enhances the pd_router by implementing intelligent dual routing strategies. This new approach optimizes the processing of Large Language Model (LLM) requests based on prompt length. Short prompts can now be routed directly to a decode pod, performing both prefill and decoding locally, thereby reducing latency. Longer prompts will continue to follow the established two-stage prefill-then-decode pipeline. This change aims to improve overall system efficiency and responsiveness for diverse LLM workloads.

Highlights

Dual Routing Strategies: Introduced two distinct routing strategies within the pd_router: 'Short-Prompt Optimization' for prompts below a configurable threshold, and the 'Standard Two-Stage Pipeline' for longer prompts.
Short-Prompt Optimization: Prompts with token counts less than or equal to AIBRIX_SHORT_PROMPT_THRESHOLD (default 2048) are now routed directly to a decode pod. This bypasses the remote prefill step, significantly reducing latency for short requests.
Configuration: A new environment variable, AIBRIX_SHORT_PROMPT_THRESHOLD, has been added to allow dynamic configuration of the token length threshold that determines whether a prompt is considered 'short'.
Core Logic Implementation: The pdRouter struct and its Route method in pd_disaggregation.go have been updated to incorporate this new conditional routing logic, complete with a detailed ASCII diagram illustrating the decision flow.
Enhanced Testing: Comprehensive new test cases have been added to pd_disaggregation_test.go to thoroughly validate both short-prompt and long-prompt routing scenarios, covering various pod configurations and different LLM engines (vLLM, SGLang).

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.

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 introduces a dual routing strategy for the PD router, optimizing for short prompts by routing them directly to a decode pod, bypassing the prefill stage. The changes include adding a configurable threshold for short prompts and updating the routing logic. The test suite has been significantly expanded to cover these new scenarios.

My review identifies a critical correctness issue in the new short-prompt routing path where it incorrectly depends on the availability of prefill pods, which could lead to failures. I've also suggested an improvement for logging to enhance observability when tokenization fails. The tests are well-written but miss a key scenario that would expose the bug I found.

gemini-code-assist · 2025-11-25T02:42:26Z

pkg/plugins/gateway/algorithms/pd_disaggregation.go

+			_, decodePod, err := r.filterPrefillDecodePods(ctx, readyPodList.All())
+			if err != nil {
+				klog.Warning("Failed to select decode pod for direct inference; falling back to prefill-decode flow",
+					"request_id", ctx.RequestID, "error", err)
+			} else if decodePod != nil {
+				ctx.SetTargetPod(decodePod)
+				return ctx.TargetAddress(), nil
+			}


The current implementation for the short-prompt optimization path reuses filterPrefillDecodePods, which is designed to find a pair of prefill and decode pods. This creates a bug: if no prefill pods are available, filterPrefillDecodePods will return an error, causing the short-prompt path to fail and fall back to the standard prefill-decode flow. The standard flow will then also fail for the same reason (no prefill pods).

The short-prompt path should be able to select a decode pod independently of prefill pods. A new test case like the one below would fail with the current implementation:

{ name: "short prompt with only decode pods available: should succeed", readyPods: []*v1.Pod{ // No prefill pod! {ObjectMeta: metav1.ObjectMeta{ Labels: map[string]string{PDRoleSetIdentifier: "test", PDRoleIdentifier: "decode"}, Name: "decode-1", }, Status: v1.PodStatus{ PodIP: "127.0.0.2", Conditions: []v1.PodCondition{{Type: v1.PodReady, Status: v1.ConditionTrue}}, }}, }, message: "hi", shortPromptThreshold: 2048, expectPrefillCall: false, serverCode: http.StatusOK, llmEngine: VLLMEngine, expectError: false, // This would fail expectTargetAddr: "127.0.0.2:8000", },

I recommend refactoring the pod selection logic. A new function, say selectDecodePod, should be created to encapsulate the logic for selecting only a decode pod. This would involve extracting the decode-pod-specific selection logic from filterPrefillDecodePods. The Route function can then be updated to call this new function for the short-prompt path.

Suggested change

_, decodePod, err := r.filterPrefillDecodePods(ctx, readyPodList.All())

if err != nil {

klog.Warning("Failed to select decode pod for direct inference; falling back to prefill-decode flow",

"request_id", ctx.RequestID, "error", err)

} else if decodePod != nil {

ctx.SetTargetPod(decodePod)

return ctx.TargetAddress(), nil

}

decodePod, err := r.selectDecodePod(ctx, readyPodList.All())

if err != nil {

klog.Warning("Failed to select decode pod for direct inference; falling back to prefill-decode flow",

"request_id", ctx.RequestID, "error", err)

} else {

ctx.SetTargetPod(decodePod)

return ctx.TargetAddress(), nil

}

I think we can ignore this comment, because filterPrefillDecodePods check both prefill and decode pod, this make sense in every round request

pkg/plugins/gateway/algorithms/pd_disaggregation.go

Jeffwan · 2025-12-01T11:21:27Z

pkg/plugins/gateway/algorithms/pd_disaggregation.go

 	prefillRequestTimeout         int     = utils.LoadEnvInt("AIBRIX_PREFILL_REQUEST_TIMEOUT", defaultPrefillRequestTimeout)
 	aibrixDecodeMaxRequest        float64 = utils.LoadEnvFloat("AIBRIX_DECODE_MAX_REQUEST", defaultMaxRequest)
 	aibrixDecodeMaxThroughputDiff float64 = utils.LoadEnvFloat("AIBRIX_DECODE_MAX_THROUGHPUT", defaultMaxTokenThroughputDiff)
+	shortPromptThreshold                  = utils.LoadEnvInt("AIBRIX_SHORT_PROMPT_THRESHOLD", defaultShortPromptThreshold)


em. this is to support remote prefill feature.

Instead of using SHORT_PROMPT etc, let's be specific and use REMOTE_PREFILL term instead

Jeffwan

overall looks good to me. I also suggest to update the doc page with detail P/D explaination. Let's replace "dual" with term instead in PR title

varungup90 · 2025-12-01T21:03:26Z

pkg/plugins/gateway/algorithms/pd_disaggregation.go

+// 1. **Short-Prompt Optimization (Decode-Only)**:
+//   - If the input prompt length (in tokens) ≤ AIBRIX_SHORT_PROMPT_THRESHOLD (default value = 2048 tokens),
+//     the request is routed directly to a decode pod.
+//   - The decode pod performs *both* prefill (KV cache computation) and decoding locally,


How will gateway know that decode pod supports both prefill/decode role? vLLM supports it, but not sure if every engine supports this.

Oh, that's a good question. I'll do some check on that, such as whether sglang supports it, and if not, I'll add the condition that it's limited to VLLM engine.

check in sglang community sgl-project/sglang#14284

googs1025 · 2025-12-03T03:03:29Z

pkg/plugins/gateway/algorithms/pd_disaggregation.go

 		return "", fmt.Errorf("engine validation failed for request %s: %w", ctx.RequestID, err)
 	}

+	// NOTE: Short-prompt optimization (bypassing remote prefill) is currently ONLY supported for vLLM.


add comment "ONLY supported for vLLM" for now

if r.remotePrefillThreshold > 0 && llmEngine == VLLMEngine {

In vLLM it supports multiple roles, but it is not necessary that user configures kv_both. Will need a flag to guide gateway.

Can you also run a benchmark to compare TTFT/TPOT/E2E P90/P99.

In vLLM it supports multiple roles, but it is not necessary that user configures kv_both. Will need a flag to guide gateway.

sorry, I don't know much about this part. Could you explain it in more detail or give me more relevant info?

Signed-off-by: CYJiang <[email protected]>

googs1025 requested review from Jeffwan and varungup90 November 25, 2025 02:39

gemini-code-assist bot reviewed Nov 25, 2025

View reviewed changes

googs1025 force-pushed the pd_disagg/feat branch 4 times, most recently from 9f43e3d to 8a1e47b Compare November 25, 2025 05:22

Jeffwan reviewed Dec 1, 2025

View reviewed changes

varungup90 reviewed Dec 1, 2025

View reviewed changes

googs1025 force-pushed the pd_disagg/feat branch 2 times, most recently from b4004a8 to 74fcad5 Compare December 3, 2025 03:02

googs1025 commented Dec 3, 2025

View reviewed changes

googs1025 force-pushed the pd_disagg/feat branch from 74fcad5 to 975ff83 Compare December 3, 2025 05:06

feat(pd_router): implement dual routing strategies for pd router

197e7e9

Signed-off-by: CYJiang <[email protected]>

googs1025 force-pushed the pd_disagg/feat branch from 975ff83 to 197e7e9 Compare December 5, 2025 00:09

googs1025 changed the title ~~feat(pd_router): implement dual routing strategies for pd router~~ feat(pd_router): implement prompt-length-aware routing for pd router Dec 5, 2025

feat(pd_router): implement prompt-length-aware routing for pd router #1800

Are you sure you want to change the base?

feat(pd_router): implement prompt-length-aware routing for pd router #1800

Conversation

googs1025 commented Nov 25, 2025

Pull Request Description

Related Issues

Pull Request Title Format

Submission Checklist

Uh oh!

gemini-code-assist bot commented Nov 25, 2025

Summary of Changes

Highlights

Footnotes

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist bot Nov 25, 2025

Choose a reason for hiding this comment

Uh oh!

googs1025 Nov 25, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Jeffwan left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

googs1025 Nov 25, 2025 •

edited

Loading