llm-d · roytman · Aug 20, 2026 · Aug 23, 2026 · Aug 24, 2026 · Aug 25, 2026
@@ -93,6 +93,15 @@ pipeline:
   # env: COORDINATOR_PIPELINE_USE_OPENAI_FORMAT
   use_openai_format: true
 
+  # dp_size is the kv-nixl model server's data-parallel world size. Only
+  # kv-nixl uses it: prefill and decode calls for one request are pinned to
+  # the same deterministic rank via the x-data-parallel-rank header, so a
+  # DP>1 backend that shares its HTTP port across ranks (e.g. via
+  # SO_REUSEPORT) serves both legs from the same rank. 1 (default) disables
+  # this and omits the header.
+  # env: COORDINATOR_PIPELINE_DP_SIZE
+  # dp_size: 1
+
   steps:
     # -------------------------------------------------------------------
     # replace-media-urls: first step. Downloads any http(s) image_url

@@ -369,6 +369,16 @@ overridden process-wide with the `SGLANG_BOOTSTRAP_PORT` environment variable. T
 is read once on the first prefill request that uses the connector; a non-integer value is
 rejected in favor of the default and logged at error level.
 
+`kv-nixl` additionally supports `pipeline.dp_size` (default `1`, disabled), the model
+server's data-parallel world size. When set above 1, the prefill and decode step hash the
+request ID into a deterministic rank and set it as the `x-data-parallel-rank` header on
+both requests, propagating a rank the prefill response echoes back in
+`kv_transfer_params.remote_dp_rank` if present. This pins both legs of one disaggregated
+request to the same rank so a DP>1 backend that shares its HTTP port across ranks (for
+example via `SO_REUSEPORT`) does not split them across ranks, mirroring the
+llm-d-router sidecar's `nixlv2` connector. `kv-sglang` and `kv-shared-storage` ignore
+`dp_size`.
+
 EC connector protocols ([pkg/coordinator/connectors/ec/](../pkg/coordinator/connectors/ec/)) ship encoder
 embeddings from encode pods to the prefill pod:
 

@@ -19,6 +19,12 @@ const (
 	// DataParallelEndpointHeader is the header name used to indicate the worker <ip:port> for Data Parallel
 	DataParallelEndpointHeader = "x-data-parallel-host-port"
 
+	// DataParallelRankHeader pins a request to a specific vLLM data-parallel
+	// rank. Set on both legs of a disaggregated pair so a DP>1 backend that
+	// shares its HTTP port across ranks (e.g. via SO_REUSEPORT) serves both
+	// legs from the same rank instead of splitting them across ranks.
+	DataParallelRankHeader = "x-data-parallel-rank"
+
 	// KVCacheSourceHeader is the header name used to indicate the worker <ip:port> holding
 	// the most cached prefix KV blocks for the request, to pull from over the P2P connector
 	// instead of recomputing them

@@ -58,6 +58,7 @@ type PipelineConfig struct {
 	KVConnector     string       `mapstructure:"kv_connector"`
 	ECConnector     string       `mapstructure:"ec_connector"`
 	UseOpenAIFormat bool         `mapstructure:"use_openai_format"`
+	DPSize          int          `mapstructure:"dp_size"`
 	Steps           []StepConfig `mapstructure:"steps"`
 }
 
@@ -83,6 +84,7 @@ func Load(path string) (*Config, error) {
 	v.SetDefault("gateway.idle_conn_timeout", 90*time.Second)
 	v.SetDefault("gateway.timeout", 60*time.Second)
 	v.SetDefault("pipeline.use_openai_format", true)
+	v.SetDefault("pipeline.dp_size", 1)
 
 	if err := v.ReadInConfig(); err != nil {
 		return nil, fmt.Errorf("reading config: %w", err)

@@ -55,6 +55,7 @@ func TestLoadDefaults(t *testing.T) {
 		{"gateway.idle_conn_timeout", cfg.Gateway.IdleConnTimeout, 90 * time.Second},
 		{"gateway.timeout", cfg.Gateway.Timeout, 60 * time.Second},
 		{"pipeline.use_openai_format", cfg.Pipeline.UseOpenAIFormat, true},
+		{"pipeline.dp_size", cfg.Pipeline.DPSize, 1},
 	}
 	for _, c := range checks {
 		if c.got != c.want {
@@ -86,6 +87,12 @@ func TestLoadEnvOverride(t *testing.T) {
 			envVal: "false",
 			check:  func(c *Config) (any, any) { return c.Pipeline.UseOpenAIFormat, false },
 		},
+		{
+			name:   "nested int",
+			envKey: "COORDINATOR_PIPELINE_DP_SIZE",
+			envVal: "4",
+			check:  func(c *Config) (any, any) { return c.Pipeline.DPSize, 4 },
+		},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {

@@ -0,0 +1,94 @@
+/*
+Copyright 2026 The llm-d Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package kv
+
+import (
+	"encoding/binary"
+	"encoding/json"
+	"math"
+
+	"golang.org/x/crypto/blake2s"
+)
+
+// remoteDPRankField is the kv_transfer_params key a prefill worker may echo
+// back to report which data-parallel rank it actually ran on. The plain NIXL
+// P2P protocol this package implements never sets it itself: the decode leg
+// addresses the prefill engine directly via remote_engine_id/remote_host/
+// remote_port, so which rank handled prefill does not matter for the KV pull.
+// resolveDecodeDPRank still checks for it so this connector stays correct if
+// a future prefill response ever starts returning one.
+const remoteDPRankField = "remote_dp_rank"
+
+// pickDPRank returns a deterministic DP rank for a request as
+// blake2s(requestID) mod dpSize, so pinning it as the same HTTP header value
+// on both legs of a disaggregated pair keeps a DP>1 backend that shares its
+// HTTP port across ranks (e.g. via SO_REUSEPORT) from splitting the pair
+// across ranks. dpSize <= 1 returns 0 so single-DP deployments are unaffected.
+func pickDPRank(requestID string, dpSize int) int {
+	if dpSize <= 1 {
+		return 0
+	}
+	h, err := blake2s.New256(nil)
+	if err != nil {
+		// Only fails on invalid key length, never for nil; fail safe to rank 0.
+		return 0
+	}
+	_, _ = h.Write([]byte(requestID))
+	sum := h.Sum(nil)
+	return int(binary.BigEndian.Uint64(sum[:8]) % uint64(dpSize))
+}
+
+// resolveDecodeDPRank picks the DP rank for the decode leg. It prefers the
+// rank the prefill leg reported in its kv_transfer_params (remoteDPRankField),
+// but only when that value is a valid integer in [0, dpSize); otherwise it
+// falls back to the deterministic hash of the request id. The second return
+// value reports whether the prefill-returned rank was used.
+func resolveDecodeDPRank(prefillKV any, requestID string, dpSize int) (rank int, usedReturned bool) {
+	fallback := pickDPRank(requestID, dpSize)
+	if dpSize <= 1 {
+		return fallback, false
+	}
+	pkv, ok := prefillKV.(map[string]any)
+	if !ok {
+		return fallback, false
+	}
+	rv, present := pkv[remoteDPRankField]
+	if !present {
+		return fallback, false
+	}
+	if f, ok := rv.(float64); ok && f != math.Trunc(f) {
+		return fallback, false
+	}
+	if ri, ok := toInt(rv); ok && ri >= 0 && ri < dpSize {
+		return ri, true
+	}
+	return fallback, false
+}
+
+// toInt converts a JSON number value (float64, int, or json.Number) to int.
+func toInt(v any) (int, bool) {
+	switch n := v.(type) {
+	case int:
+		return n, true
+	case float64:
+		return int(n), true
+	case json.Number:
+		i, err := n.Int64()
+		return int(i), err == nil
+	}
+	return 0, false
+}
@@ -0,0 +1,84 @@
+/*
+Copyright 2026 The llm-d Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package kv
+
+import "testing"
+
+// TestPickDPRankSingleDP verifies dpSize <= 1 short-circuits to 0 without
+// hashing.
+func TestPickDPRankSingleDP(t *testing.T) {
+	for _, dpSize := range []int{0, 1, -1} {
+		if got := pickDPRank("any-request-id", dpSize); got != 0 {
+			t.Errorf("pickDPRank(_, %d) = %d; want 0", dpSize, got)
+		}
+	}
+}
+
+// TestPickDPRankDeterministicAndInRange verifies the two load-bearing
+// invariants: the same requestID+dpSize always returns the same rank
+// (otherwise the prefill and decode legs of one pair could land on
+// different ranks), and the result is always in [0, dpSize).
+func TestPickDPRankDeterministicAndInRange(t *testing.T) {
+	requestIDs := []string{"req-1", "cmpl-abc-123", "00000000-0000-0000-0000-000000000000", ""}
+	for _, dpSize := range []int{2, 3, 8, 16} {
+		for _, rid := range requestIDs {
+			first := pickDPRank(rid, dpSize)
+			if first < 0 || first >= dpSize {
+				t.Errorf("pickDPRank(%q, %d) = %d; want in [0, %d)", rid, dpSize, first, dpSize)
+			}
+			for i := range 3 {
+				if got := pickDPRank(rid, dpSize); got != first {
+					t.Errorf("pickDPRank(%q, %d) = %d on call %d, want %d (must be deterministic)", rid, dpSize, got, i, first)
+				}
+			}
+		}
+	}
+}
+
+func TestResolveDecodeDPRank(t *testing.T) {
+	const dpSize = 8
+	const rid = "cmpl-rank-test"
+	hashFallback := pickDPRank(rid, dpSize)
+
+	cases := []struct {
+		name         string
+		prefillKV    any
+		dpSize       int
+		wantRank     int
+		wantReturned bool
+	}{
+		{name: "valid returned rank (float64)", prefillKV: map[string]any{remoteDPRankField: float64(3)}, dpSize: dpSize, wantRank: 3, wantReturned: true},
+		{name: "valid returned rank (int)", prefillKV: map[string]any{remoteDPRankField: 5}, dpSize: dpSize, wantRank: 5, wantReturned: true},
+		{name: "zero is valid", prefillKV: map[string]any{remoteDPRankField: float64(0)}, dpSize: dpSize, wantRank: 0, wantReturned: true},
+		{name: "omitted falls back to hash", prefillKV: map[string]any{}, dpSize: dpSize, wantRank: hashFallback, wantReturned: false},
+		{name: "nil kv falls back to hash", prefillKV: nil, dpSize: dpSize, wantRank: hashFallback, wantReturned: false},
+		{name: "non-numeric falls back to hash", prefillKV: map[string]any{remoteDPRankField: "two"}, dpSize: dpSize, wantRank: hashFallback, wantReturned: false},
+		{name: "fractional falls back to hash (no truncation)", prefillKV: map[string]any{remoteDPRankField: float64(3.5)}, dpSize: dpSize, wantRank: hashFallback, wantReturned: false},
+		{name: "out-of-range high falls back to hash", prefillKV: map[string]any{remoteDPRankField: float64(8)}, dpSize: dpSize, wantRank: hashFallback, wantReturned: false},
+		{name: "negative falls back to hash", prefillKV: map[string]any{remoteDPRankField: float64(-1)}, dpSize: dpSize, wantRank: hashFallback, wantReturned: false},
+		{name: "single-DP always rank 0", prefillKV: map[string]any{remoteDPRankField: float64(3)}, dpSize: 1, wantRank: 0, wantReturned: false},
+	}
+	for _, c := range cases {
+		t.Run(c.name, func(t *testing.T) {
+			gotRank, gotReturned := resolveDecodeDPRank(c.prefillKV, rid, c.dpSize)
+			if gotRank != c.wantRank || gotReturned != c.wantReturned {
+				t.Errorf("resolveDecodeDPRank(%v, %q, %d) = (%d, %t); want (%d, %t)",
+					c.prefillKV, rid, c.dpSize, gotRank, gotReturned, c.wantRank, c.wantReturned)
+			}
+		})
+	}
+}
@@ -50,14 +50,30 @@ type Connector interface {
 	PrepareDecodeKVParams(ctx context.Context, reqCtx *pipeline.RequestContext) map[string]any
 }
 
-// Build returns the KV connector for name. An empty name selects DefaultKVConnectorName.
-func Build(name string) (Connector, error) {
+// RankHeaderer is implemented by connectors whose wire protocol needs to pin
+// a request to a specific vLLM data-parallel rank via an HTTP header, in
+// addition to whatever Prepare*KVParams returns. PrefillStep and DecodeStep
+// type-assert for this; a connector that does not implement it is unaffected.
+type RankHeaderer interface {
+	// PrefillHeaders returns headers to set on the prefill request, or nil.
+	PrefillHeaders(ctx context.Context, reqCtx *pipeline.RequestContext) map[string]string
+	// DecodeHeaders returns headers to set on the decode request, or nil.
+	DecodeHeaders(ctx context.Context, reqCtx *pipeline.RequestContext) map[string]string
+}
+
+// Build returns the KV connector for name. An empty name selects
+// DefaultKVConnectorName. dpSize is the model server's data-parallel world
+// size (1 disables DP-rank pinning); only the nixl connector uses it.
+func Build(name string, dpSize int) (Connector, error) {
+	if dpSize < 1 {
+		return nil, fmt.Errorf("kv connector: dp_size must be a positive integer, got %d", dpSize)
+	}
 	if name == "" {
 		name = DefaultKVConnectorName
 	}
 	switch name {
 	case NIXL:
-		return nixlKV{}, nil
+		return nixlKV{dpSize: dpSize}, nil
 	case SharedStorage:
 		return sharedStorageKV{}, nil
 	case SGLang:

@@ -25,7 +25,7 @@ import (
 )
 
 func TestSGLangKV_Params(t *testing.T) {
-	c, err := Build(SGLang)
+	c, err := Build(SGLang, 1)
 	if err != nil {
 		t.Fatalf("Build(%q): %v", SGLang, err)
 	}
@@ -71,13 +71,13 @@ func TestSGLangKV_Params(t *testing.T) {
 }
 
 func TestBuild_UnknownReturnsError(t *testing.T) {
-	if _, err := Build("does-not-exist"); err == nil {
+	if _, err := Build("does-not-exist", 1); err == nil {
 		t.Fatal("expected error for unknown connector")
 	}
 }
 
 func TestBuild_EmptyReturnsDefault(t *testing.T) {
-	c, err := Build("")
+	c, err := Build("", 1)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -86,6 +86,14 @@ func TestBuild_EmptyReturnsDefault(t *testing.T) {
 	}
 }
 
+func TestBuild_RejectsInvalidDPSize(t *testing.T) {
+	for _, dpSize := range []int{0, -1} {
+		if _, err := Build(NIXL, dpSize); err == nil {
+			t.Errorf("Build(%q, %d): expected error for non-positive dp_size", NIXL, dpSize)
+		}
+	}
+}
+
 func TestConnectors_KVParams(t *testing.T) {
 	cases := []struct {
 		name           string
@@ -126,7 +134,7 @@ func TestConnectors_KVParams(t *testing.T) {
 
 	for _, tc := range cases {
 		t.Run(tc.name, func(t *testing.T) {
-			c, err := Build(tc.name)
+			c, err := Build(tc.name, 1)
 			if err != nil {
 				t.Fatalf("Build(%q): %v", tc.name, err)
 			}