Address review comments.

Signed-off-by: Revital Sur <eres@il.ibm.com>

Author: revit13

docs/communication.md

@@ -607,7 +607,7 @@ EPP-Phase: prefill
 `kwargs_data` carries the same per-image base64 tensors from the render step (same values sent to the encode stage). Each blob is a msgpack-serialized `MultiModalKwargsItem` containing both `pixel_values` and `image_grid_thw` (and any other model-specific keys). The prefill worker needs `image_grid_thw` to compute mRoPE (multimodal Rotary Position Embedding) positional encodings for the visual tokens.
 
 > [!NOTE]
-> Due to a bug in the `/inference/v1/generate` implementation, top-level `kv_transfer_params` and `ec_transfer_params` are not propagated to the engine: the endpoint reads transfer parameters only from `sampling_params.extra_args` (the top-level `ec_transfer_params` is used solely to echo back in the response). The coordinator nests both under `extra_args`:
+> Due to a bug in the `/inference/v1/generate` implementation, top-level `kv_transfer_params` and `ec_transfer_params` are not propagated to the engine: the endpoint reads transfer parameters only from `sampling_params.extra_args`. The coordinator nests both under `extra_args`:
 
 ```
 POST <gateway>/inference/v1/generate
@@ -647,15 +647,27 @@ EPP-Phase: prefill
 ```json
 {
   "request_id": "generate-tokens-abc123",
-  "choices": [],
+  "choices": [
+    {"index": 0, "logprobs": null, "finish_reason": "length", "token_ids": [28715], "routed_experts": null}
+  ],
+  "prompt_logprobs": null,
   "kv_transfer_params": {
-    "block_id": "block-999",
-    "peer_host": "10.0.0.42",
-    "peer_port": 7777
-  }
+    "do_remote_prefill": true,
+    "do_remote_decode": false,
+    "remote_block_ids": [[1, 2, 4]],
+    "remote_engine_id": "e1101616-bf27-4687-bd0b-05970390868e",
+    "remote_request_id": "generate-tokens-abc123",
+    "remote_host": "10.0.0.42",
+    "remote_port": 5600,
+    "tp_size": 1,
+    "remote_num_tokens": 318
+  },
+  "ec_transfer_params": null
 }
 ```
 
+`kv_transfer_params` carries the remote-prefill handoff the decode worker needs to pull the KV cache (`remote_engine_id`, `remote_block_ids`, `remote_host`/`remote_port`, `remote_num_tokens`). `ec_transfer_params` is `null` in the prefill response.
+
 ---
 
 ### Option B: /v1/chat/completions

pkg/coordinator/steps/decode.go

@@ -101,7 +101,7 @@ func (s *DecodeStep) prepareDecodeBody(ctx context.Context, reqCtx *pipeline.Req
 		if len(reqCtx.TokenIDs) > 0 {
 			reqCtx.Body["prompt"] = reqCtx.TokenIDs
 		}
-	default:
+	case gateway.FormatGenerate:
 		// The /inference/v1/generate engine reads transfer params only from
 		// sampling_params.extra_args; a top-level kv_transfer_params is ignored,
 		// so the decode worker never pulls the prefill KV over NIXL. Merge into
@@ -111,12 +111,7 @@ func (s *DecodeStep) prepareDecodeBody(ctx context.Context, reqCtx *pipeline.Req
 			sampling = map[string]any{}
 			reqCtx.Body[reqcommon.FieldSamplingParams] = sampling
 		}
-		extraArgs, ok := sampling[reqcommon.FieldExtraArgs].(map[string]any)
-		if !ok {
-			extraArgs = map[string]any{}
-			sampling[reqcommon.FieldExtraArgs] = extraArgs
-		}
-		extraArgs[reqcommon.FieldKVTransferParams] = kvParams
+		setGenerateTransferParams(sampling, kvParams, nil)
 	}
 }
 

pkg/coordinator/steps/prefill.go

@@ -175,29 +175,25 @@ func (s *PrefillStep) buildPrefillBody(ctx context.Context, reqCtx *pipeline.Req
 		}
 		return body, nil
 
-	default:
+	case gateway.FormatGenerate:
 		// The /inference/v1/generate engine reads transfer params only from
 		// sampling_params.extra_args; top-level fields are ignored on input.
-		extraArgs := map[string]any{
-			reqcommon.FieldKVTransferParams: kvParams,
-		}
-		if len(ecParams) > 0 {
-			extraArgs[reqcommon.FieldECTransferParams] = ecParams
-		}
+		sampling := map[string]any{reqcommon.FieldMaxTokens: 1}
+		setGenerateTransferParams(sampling, kvParams, ecParams)
 		body := map[string]any{
-			"request_id": reqCtx.RequestID,
-			"token_ids":  reqCtx.TokenIDs,
-			"model":      reqCtx.Model,
-			reqcommon.FieldSamplingParams: map[string]any{
-				reqcommon.FieldMaxTokens: 1,
-				reqcommon.FieldExtraArgs: extraArgs,
-			},
+			"request_id":                  reqCtx.RequestID,
+			"token_ids":                   reqCtx.TokenIDs,
+			"model":                       reqCtx.Model,
+			reqcommon.FieldSamplingParams: sampling,
 		}
 		if features != nil {
 			body["features"] = features
 		}
 		return body, nil
 	}
+	// resolveFormat only ever yields the three formats above; a new value
+	// reaching here is a programming error, not a client fault.
+	return nil, fmt.Errorf("prefill: unsupported request format %v", format)
 }
 
 type prefillResponse struct {

pkg/coordinator/steps/render.go

@@ -148,6 +148,10 @@ func (s *RenderStep) executeGenerate(ctx context.Context, reqCtx *pipeline.Reque
 	}
 	reqCtx.TokenIDs = tokenIDs
 
+	if err := validateSamplingParams(reqCtx.Body); err != nil {
+		return fmt.Errorf("render: %w", err)
+	}
+
 	rawFeatures := reqCtx.Body["features"]
 	var features map[string]any
 	if rawFeatures != nil {
@@ -161,6 +165,13 @@ func (s *RenderStep) executeGenerate(ctx context.Context, reqCtx *pipeline.Reque
 	if err != nil {
 		return fmt.Errorf("render: %w", err)
 	}
+	// The client supplies placeholder geometry directly on this path, so bound
+	// every span to the prompt before EncodeStep allocates from length. Unlike
+	// checkPlaceholderLimit, this guard is unconditional: it does not depend on
+	// the optional max_total_placeholder_tokens knob.
+	if err := validatePlaceholderBounds(entries, len(tokenIDs)); err != nil {
+		return fmt.Errorf("render: %w", err)
+	}
 	if err := s.checkPlaceholderLimit(entries); err != nil {
 		return err
 	}
@@ -356,9 +367,13 @@ func (s *RenderStep) checkPlaceholderLimit(entries []pipeline.MultimodalEntry) e
 	total := 0
 	for _, e := range entries {
 		total += e.Placeholder.Length
-	}
-	if total > s.maxTotalPlaceholderTokens {
-		return fmt.Errorf("too many placeholder tokens: got %d, max %d: %w", total, s.maxTotalPlaceholderTokens, pipeline.ErrBadRequest)
+		// total < 0 catches int overflow: lengths are non-negative, so a sum
+		// past MaxInt wraps negative and would otherwise slip past the
+		// total > max check, silently bypassing the limit. On the generate path
+		// lengths come straight from the client, so this must not be evadable.
+		if total < 0 || total > s.maxTotalPlaceholderTokens {
+			return fmt.Errorf("too many placeholder tokens: got %d, max %d: %w", total, s.maxTotalPlaceholderTokens, pipeline.ErrBadRequest)
+		}
 	}
 	return nil
 }

pkg/coordinator/steps/render_test.go

@@ -21,6 +21,7 @@ import (
 	"encoding/json"
 	"errors"
 	"io"
+	"math"
 	"net/http"
 	"net/http/httptest"
 	"strings"
@@ -394,6 +395,31 @@ func TestRenderStep_AllowsAtPlaceholderLimit(t *testing.T) {
 	}
 }
 
+func TestRenderStep_PlaceholderLimitOverflow(t *testing.T) {
+	step, err := NewRenderStep(nil, map[string]any{"max_total_placeholder_tokens": 5})
+	if err != nil {
+		t.Fatalf("NewRenderStep: %v", err)
+	}
+	rs := step.(*RenderStep)
+
+	// Two lengths whose sum overflows int and wraps negative. Without the
+	// overflow guard, total > max is false and the limit is silently bypassed.
+	entries := []pipeline.MultimodalEntry{
+		{Placeholder: pipeline.PlaceholderRange{Length: math.MaxInt}},
+		{Placeholder: pipeline.PlaceholderRange{Length: math.MaxInt}},
+	}
+	err = rs.checkPlaceholderLimit(entries)
+	if err == nil {
+		t.Fatal("expected error for overflowing placeholder length sum")
+	}
+	if !errors.Is(err, pipeline.ErrBadRequest) {
+		t.Fatalf("expected ErrBadRequest, got %v", err)
+	}
+	if !strings.Contains(err.Error(), "too many placeholder tokens") {
+		t.Fatalf("unexpected error: %v", err)
+	}
+}
+
 func TestRenderStep_RejectsNegativeLimits(t *testing.T) {
 	if _, err := NewRenderStep(nil, map[string]any{"max_total_tokens": -1}); err == nil {
 		t.Fatal("expected error for negative max_total_tokens")
@@ -555,6 +581,71 @@ func TestRenderStep_GenerateFormat_MalformedFeatures(t *testing.T) {
 	}
 }
 
+func TestRenderStep_GenerateFormat_PlaceholderOutOfBounds(t *testing.T) {
+	for _, tc := range []struct {
+		name   string
+		offset float64
+		length float64
+	}{
+		{"offset_out_of_range", 5, 1},
+		{"length_exceeds_prompt", 1, 9007199254740992},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			step, _ := NewRenderStep(nil, map[string]any{})
+			reqCtx := &pipeline.RequestContext{
+				OriginalPath: gateway.DefaultGeneratePath,
+				Body: map[string]any{
+					"model":     "test-model",
+					"token_ids": []any{float64(1), float64(32000), float64(32000), float64(2)},
+					"features": map[string]any{
+						"mm_hashes": map[string]any{"image": []any{"abc123"}},
+						"mm_placeholders": map[string]any{"image": []any{
+							map[string]any{"offset": tc.offset, "length": tc.length},
+						}},
+						"kwargs_data": map[string]any{"image": []any{"dGVuc29y"}},
+					},
+				},
+			}
+			err := step.Execute(context.Background(), reqCtx)
+			if err == nil {
+				t.Fatal("expected error for out-of-bounds placeholder")
+			}
+			if !errors.Is(err, pipeline.ErrBadRequest) {
+				t.Errorf("expected ErrBadRequest, got %v", err)
+			}
+		})
+	}
+}
+
+func TestRenderStep_GenerateFormat_InvalidSamplingParams(t *testing.T) {
+	for _, tc := range []struct {
+		name           string
+		samplingParams any
+	}{
+		{"array", []any{float64(1), float64(2)}},
+		{"string", "greedy"},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			step, _ := NewRenderStep(nil, map[string]any{})
+			reqCtx := &pipeline.RequestContext{
+				OriginalPath: gateway.DefaultGeneratePath,
+				Body: map[string]any{
+					"model":           "test-model",
+					"token_ids":       []any{float64(1), float64(2), float64(3)},
+					"sampling_params": tc.samplingParams,
+				},
+			}
+			err := step.Execute(context.Background(), reqCtx)
+			if err == nil {
+				t.Fatal("expected error for non-object sampling_params")
+			}
+			if !errors.Is(err, pipeline.ErrBadRequest) {
+				t.Errorf("expected ErrBadRequest, got %v", err)
+			}
+		})
+	}
+}
+
 func TestRenderStep_GenerateFormat_NullFeatures(t *testing.T) {
 	step, _ := NewRenderStep(nil, map[string]any{})
 	reqCtx := &pipeline.RequestContext{

pkg/coordinator/steps/utils.go

@@ -25,6 +25,7 @@ import (
 	"github.com/go-logr/logr"
 
 	logutil "github.com/llm-d/llm-d-router/pkg/common/observability/logging"
+	reqcommon "github.com/llm-d/llm-d-router/pkg/common/request"
 
 	"github.com/llm-d/llm-d-router/pkg/coordinator/gateway"
 	"github.com/llm-d/llm-d-router/pkg/coordinator/pipeline"
@@ -117,6 +118,24 @@ func mmKwargsField(kwargs []string) map[string][]any {
 	return map[string][]any{ModalityImage: items}
 }
 
+// setGenerateTransferParams nests the kv/ec transfer params under
+// sampling_params.extra_args, the only place the /inference/v1/generate engine
+// reads them (top-level kv_transfer_params/ec_transfer_params are ignored on
+// input). It get-or-creates extra_args on the given sampling map so a client's
+// existing generation fields survive. ecParams may be empty, in which case
+// ec_transfer_params is left unset.
+func setGenerateTransferParams(sampling map[string]any, kvParams any, ecParams map[string]any) {
+	extraArgs, ok := sampling[reqcommon.FieldExtraArgs].(map[string]any)
+	if !ok {
+		extraArgs = map[string]any{}
+		sampling[reqcommon.FieldExtraArgs] = extraArgs
+	}
+	extraArgs[reqcommon.FieldKVTransferParams] = kvParams
+	if len(ecParams) > 0 {
+		extraArgs[reqcommon.FieldECTransferParams] = ecParams
+	}
+}
+
 // coerceParamsMap coerces a transfer-params value from an upstream response to a
 // map: a non-object value is logged at debug and skipped (returns nil) rather
 // than failing the request. A missing or null value is already nil; an empty map
@@ -137,7 +156,8 @@ func coerceParamsMap(logger logr.Logger, v any, label string) map[string]any {
 
 // toIntSlice converts a JSON-unmarshalled []any of numeric elements to []int.
 // Each element must be a non-negative integer represented as float64 or json.Number.
-// Callers wrap the returned error with their own field context.
+// The returned error identifies the offending element by index and wraps
+// pipeline.ErrBadRequest.
 func toIntSlice(values []any) ([]int, error) {
 	out := make([]int, 0, len(values))
 	for i, v := range values {
@@ -325,3 +345,57 @@ func extractMultimodalEntries(features map[string]any) ([]pipeline.MultimodalEnt
 	}
 	return entries, nil
 }
+
+// validateSamplingParams checks that sampling_params and its nested extra_args,
+// when present, are JSON objects. Both are optional. The decode step merges
+// kv_transfer_params into sampling_params.extra_args; a non-object at either
+// level would fall into its fallback branch and be silently replaced with an
+// empty map, discarding client-requested generation parameters with no error.
+// Validating once at ingestion keeps that path fail-loud, consistent with
+// token_ids and features.
+func validateSamplingParams(body map[string]any) error {
+	raw, ok := body[reqcommon.FieldSamplingParams]
+	if !ok || raw == nil {
+		return nil
+	}
+	sampling, ok := raw.(map[string]any)
+	if !ok {
+		return fmt.Errorf("%s must be an object, got %T: %w",
+			reqcommon.FieldSamplingParams, raw, pipeline.ErrBadRequest)
+	}
+	ea, ok := sampling[reqcommon.FieldExtraArgs]
+	if !ok || ea == nil {
+		return nil
+	}
+	if _, ok := ea.(map[string]any); !ok {
+		return fmt.Errorf("%s.%s must be an object, got %T: %w",
+			reqcommon.FieldSamplingParams, reqcommon.FieldExtraArgs, ea, pipeline.ErrBadRequest)
+	}
+	return nil
+}
+
+// validatePlaceholderBounds checks that every placeholder span [offset,
+// offset+length) lies within a prompt of tokenCount tokens. It guards the
+// generate path, where the client supplies placeholder geometry directly:
+// EncodeStep.buildEncodeTokenIDs indexes token_ids[offset] and allocates
+// make([]int, 1+length), so an out-of-range offset reads the wrong token and an
+// unbounded length (a tiny request can claim billions) is a memory-exhaustion
+// vector. vLLM declares offset/length as plain unbounded ints on the generate
+// endpoint and does not enforce this, so the coordinator does. offset and
+// length are already guaranteed non-negative by extractMultimodalEntries.
+func validatePlaceholderBounds(entries []pipeline.MultimodalEntry, tokenCount int) error {
+	for i, e := range entries {
+		off := e.Placeholder.Offset
+		length := e.Placeholder.Length
+		if off >= tokenCount {
+			return fmt.Errorf("mm_placeholders[%d].offset %d out of range for %d token_ids: %w",
+				i, off, tokenCount, pipeline.ErrBadRequest)
+		}
+		// off < tokenCount, so tokenCount-off is positive and cannot overflow.
+		if length > tokenCount-off {
+			return fmt.Errorf("mm_placeholders[%d] span (offset %d + length %d) exceeds %d token_ids: %w",
+				i, off, length, tokenCount, pipeline.ErrBadRequest)
+		}
+	}
+	return nil
+}