Documentation
¶
Index ¶
Constants ¶
const ( // ConnectorTypeIdentifier is the pod label that specifies the KV transfer backend. ConnectorTypeIdentifier = "model.aibrix.ai/kv-connector-type" ConnectorTypeSHFS = "shfs" ConnectorTypeNIXL = "nixl" ConnectorTypeMooncake = "mooncake" )
Variables ¶
This section is empty.
Functions ¶
func Register ¶
func Register(name string, factory func() KVTransferAgent)
Register makes a KVTransferAgent factory available under name. Called from init() in each agent file.
func ResolveConnectorType ¶
ResolveConnectorType resolves the effective connector type from a pod label value, falling back to globalDefault when the label is absent or unrecognized.
func ValidAgentNames ¶
func ValidAgentNames() []string
ValidAgentNames returns a sorted list of registered connector type names.
Types ¶
type KVTransferAgent ¶
type KVTransferAgent interface {
// Type returns the connector identifier (e.g. "shfs", "nixl", "mooncake").
Type() string
// AugmentPrefillRequest mutates completionRequest with any fields the prefill
// pod requires before the HTTP request is sent (e.g. kv_transfer_params for SHFS).
AugmentPrefillRequest(
routingCtx *types.RoutingContext,
prefillPod *v1.Pod,
completionRequest map[string]any,
) error
// MergePrefillResponse injects connector-specific metadata from the prefill
// response into routingCtx.ReqBody before the request is forwarded to the decode pod.
MergePrefillResponse(
routingCtx *types.RoutingContext,
prefillResponse map[string]any,
prefillPod *v1.Pod,
) error
}
KVTransferAgent owns connector-specific request and response mutation for one KV transfer backend. Each implementation handles one backend (SHFS, NIXL, Mooncake, …).
func Resolve ¶
func Resolve(connectorType string) (KVTransferAgent, error)
Resolve returns a new KVTransferAgent for the given connector type. Returns an error for unknown types.
func ResolveAgentForPod ¶
func ResolveAgentForPod(pod *v1.Pod, globalDefault string) (KVTransferAgent, error)
ResolveAgentForPod resolves the KVTransferAgent for a pod. It checks the ConnectorTypeIdentifier pod label first, then falls back to globalDefault.
type MooncakeAgent ¶
type MooncakeAgent struct{}
MooncakeAgent implements KVTransferAgent for the Mooncake KV transfer backend. TODO: implement based on confirmed vLLM MooncakeConnector contract.
func (*MooncakeAgent) AugmentPrefillRequest ¶
func (a *MooncakeAgent) AugmentPrefillRequest( _ *types.RoutingContext, _ *v1.Pod, _ map[string]any, ) error
AugmentPrefillRequest adds any Mooncake-specific fields to the prefill request. TODO: implement once vLLM MooncakeConnector prefill request contract is confirmed.
func (*MooncakeAgent) MergePrefillResponse ¶
func (a *MooncakeAgent) MergePrefillResponse( _ *types.RoutingContext, _ map[string]any, _ *v1.Pod, ) error
MergePrefillResponse injects Mooncake-specific metadata from the prefill response into routingCtx.ReqBody before the request is forwarded to the decode pod. TODO: implement once vLLM MooncakeConnector response contract is confirmed.
func (*MooncakeAgent) Type ¶
func (a *MooncakeAgent) Type() string
type NIXLAgent ¶
type NIXLAgent struct{}
NIXLAgent implements KVTransferAgent for the NIXL backend (Neuron).
func (*NIXLAgent) AugmentPrefillRequest ¶
func (a *NIXLAgent) AugmentPrefillRequest( _ *types.RoutingContext, _ *v1.Pod, _ map[string]any, ) error
AugmentPrefillRequest is a no-op for NIXL: the backend manages KV transfer through its own mechanism and does not require a prefill request skeleton.
func (*NIXLAgent) MergePrefillResponse ¶
func (a *NIXLAgent) MergePrefillResponse( routingCtx *types.RoutingContext, prefillResponse map[string]any, prefillPod *v1.Pod, ) error
MergePrefillResponse wraps the entire prefill response under disagg_prefill_resp so the NixlConnector on the decode side can locate and pull the KV blocks.
type SHFSAgent ¶
type SHFSAgent struct{}
SHFSAgent implements KVTransferAgent for the AIBrix SHFS/KVCacheManager backend (GPU).
func (*SHFSAgent) AugmentPrefillRequest ¶
func (a *SHFSAgent) AugmentPrefillRequest( _ *types.RoutingContext, _ *v1.Pod, completionRequest map[string]any, ) error
AugmentPrefillRequest adds a kv_transfer_params skeleton so the prefill pod knows to populate remote block IDs for the decode side.
func (*SHFSAgent) MergePrefillResponse ¶
func (a *SHFSAgent) MergePrefillResponse( routingCtx *types.RoutingContext, prefillResponse map[string]any, prefillPod *v1.Pod, ) error
MergePrefillResponse extracts kv_transfer_params from the prefill response, sets remote_host to the prefill pod IP, and writes the merged params into routingCtx.ReqBody so the decode pod can pull the KV cache blocks.
Source Files