Documentation
¶
Index ¶
- Constants
- func GetDisaggRequestID(machineID int64) int64
- func Register(h EngineHandler)
- func SetConnectorTypeFunc(fn func() string)
- func ValidEngineNames() []string
- func ValidateTRTMachineID(machineID int64) error
- type DefaultHandler
- type EngineHandler
- type SGLangHandler
- type TRTLLMHandler
- func (h *TRTLLMHandler) AugmentPrefillRequest(_ *types.RoutingContext, _ *v1.Pod, completionRequest map[string]any) error
- func (h *TRTLLMHandler) IsAsync() bool
- func (h *TRTLLMHandler) MergePrefillResponse(routingCtx *types.RoutingContext, responseData map[string]any, ...) error
- func (h *TRTLLMHandler) Name() string
- type VLLMHandler
Constants ¶
const ( // TRTMachineIDBits is the width of the machine-ID field in snowflake disagg request IDs. TRTMachineIDBits = 10 // TRTMinGlobalID is the minimum global disagg ID; values below this are treated as local by TRT-LLM. TRTMinGlobalID = int64(1) << 42 )
Snowflake-style disagg request ID constants for TensorRT-LLM PD routing. Layout: [timestamp(41b)][machineID(10b)][counter(12b)] The modulo rotation guarantees result >= TRTMinGlobalID so TRT-LLM's executor treats it as a global (cross-worker) disagg ID rather than a local one.
Variables ¶
This section is empty.
Functions ¶
func GetDisaggRequestID ¶
GetDisaggRequestID generates a snowflake-style ID shared between a prefill and its decode request so TRT-LLM can correlate the KV-cache entry.
func Register ¶
func Register(h EngineHandler)
Register adds h to the global engine registry. Called from init() in each engine file.
func SetConnectorTypeFunc ¶
func SetConnectorTypeFunc(fn func() string)
SetConnectorTypeFunc installs a function that returns the current global KV connector type. Call this once at startup with a closure over the routingalgorithms package-level var.
func ValidEngineNames ¶
func ValidEngineNames() []string
ValidEngineNames returns the sorted list of registered engine names.
func ValidateTRTMachineID ¶
ValidateTRTMachineID returns an error if machineID does not fit in TRTMachineIDBits bits.
Types ¶
type DefaultHandler ¶
type DefaultHandler struct{}
DefaultHandler is a no-op handler used for unknown or future engines. It performs a synchronous prefill with no request augmentation or response processing, which matches the original "default" branch behaviour.
func (*DefaultHandler) AugmentPrefillRequest ¶
func (h *DefaultHandler) AugmentPrefillRequest(_ *types.RoutingContext, _ *v1.Pod, _ map[string]any) error
func (*DefaultHandler) IsAsync ¶
func (h *DefaultHandler) IsAsync() bool
func (*DefaultHandler) MergePrefillResponse ¶
func (h *DefaultHandler) MergePrefillResponse(_ *types.RoutingContext, _ map[string]any, _ *v1.Pod) error
func (*DefaultHandler) Name ¶
func (h *DefaultHandler) Name() string
type EngineHandler ¶
type EngineHandler interface {
// Name returns the engine identifier (e.g. "vllm", "sglang", "trtllm").
Name() string
// IsAsync returns true when the engine's prefill handshake is
// self-coordinating and should be fired in a goroutine (e.g. SGLang).
IsAsync() bool
// AugmentPrefillRequest adds engine-specific fields to completionRequest
// before it is POSTed to the prefill pod.
AugmentPrefillRequest(routingCtx *types.RoutingContext, pod *v1.Pod, completionRequest map[string]any) error
// MergePrefillResponse injects engine-specific data from the prefill
// response into routingCtx.ReqBody before the decode pod receives it.
MergePrefillResponse(routingCtx *types.RoutingContext, responseData map[string]any, pod *v1.Pod) error
}
EngineHandler encapsulates engine-specific behaviour for PD-disaggregated prefill: how to augment the outbound prefill request body and how to merge the prefill response back into the decode request body.
func Resolve ¶
func Resolve(name string) EngineHandler
Resolve returns the registered EngineHandler for name. Returns DefaultHandler for unknown engines so the caller always gets a valid handler without needing to check for nil.
type SGLangHandler ¶
type SGLangHandler struct{}
SGLangHandler implements EngineHandler for SGLang. SGLang uses a bootstrap handshake (bootstrap_host/port/room) to coordinate KV transfer between prefill and decode pods out-of-band, so the prefill fires asynchronously (IsAsync returns true).
func (*SGLangHandler) AugmentPrefillRequest ¶
func (h *SGLangHandler) AugmentPrefillRequest( routingCtx *types.RoutingContext, pod *v1.Pod, completionRequest map[string]any, ) error
AugmentPrefillRequest injects bootstrap_host, bootstrap_port, and a random bootstrap_room. It also propagates these fields into routingCtx.ReqBody so the decode pod receives the bootstrap address.
func (*SGLangHandler) IsAsync ¶
func (h *SGLangHandler) IsAsync() bool
func (*SGLangHandler) MergePrefillResponse ¶
func (h *SGLangHandler) MergePrefillResponse( _ *types.RoutingContext, _ map[string]any, _ *v1.Pod, ) error
MergePrefillResponse is a no-op for SGLang: the bootstrap handshake is self-coordinating, so no data from the prefill response needs to be injected into the decode request.
func (*SGLangHandler) Name ¶
func (h *SGLangHandler) Name() string
type TRTLLMHandler ¶
type TRTLLMHandler struct{}
TRTLLMHandler implements EngineHandler for TensorRT-LLM.
func (*TRTLLMHandler) AugmentPrefillRequest ¶
func (h *TRTLLMHandler) AugmentPrefillRequest( _ *types.RoutingContext, _ *v1.Pod, completionRequest map[string]any, ) error
AugmentPrefillRequest adds disaggregated_params with request_type="context_only" and a unique disagg_request_id generated from the machine snowflake ID.
func (*TRTLLMHandler) IsAsync ¶
func (h *TRTLLMHandler) IsAsync() bool
func (*TRTLLMHandler) MergePrefillResponse ¶
func (h *TRTLLMHandler) MergePrefillResponse( routingCtx *types.RoutingContext, responseData map[string]any, prefillPod *v1.Pod, ) error
MergePrefillResponse injects TensorRT-LLM disaggregated_params from the prefill response into routingCtx.ReqBody so the decode worker can resume generation from the pre-filled KV cache.
func (*TRTLLMHandler) Name ¶
func (h *TRTLLMHandler) Name() string
type VLLMHandler ¶
type VLLMHandler struct{}
VLLMHandler implements EngineHandler for vLLM. KV-transfer augmentation is fully delegated to the KVTransferAgent resolved for the prefill pod, so adding a new connector type (e.g. Mooncake) requires no changes here.
func (*VLLMHandler) AugmentPrefillRequest ¶
func (h *VLLMHandler) AugmentPrefillRequest( routingCtx *types.RoutingContext, pod *v1.Pod, completionRequest map[string]any, ) error
func (*VLLMHandler) IsAsync ¶
func (h *VLLMHandler) IsAsync() bool
func (*VLLMHandler) MergePrefillResponse ¶
func (h *VLLMHandler) MergePrefillResponse( routingCtx *types.RoutingContext, responseData map[string]any, pod *v1.Pod, ) error
func (*VLLMHandler) Name ¶
func (h *VLLMHandler) Name() string
Source Files