[Docker] Install numactl CLI in CUDA runtime image (#41032 )

Signed-off-by: Zhewen Li <zhewenli@inferact.ai> Co-authored-by: Zhewen Li <zhewenli@inferact.ai> Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com> (cherry picked from commit 6232fb4b66)
[DSV4] Add silu clamp limit to shared expert (#40950 )
2026-04-27 12:04:32 -07:00 · 2026-04-27 00:40:46 -07:00 · 2026-04-26 21:43:00 -07:00 · 2026-04-26 21:42:15 -07:00
42 changed files with 757 additions and 312 deletions
@@ -11,29 +11,74 @@
 namespace vllm {

 template <typename scalar_t, scalar_t (*ACT_FN)(const scalar_t&),
-          bool act_first>
+          bool act_first, bool HAS_CLAMP>
 __device__ __forceinline__ scalar_t compute(const scalar_t& x,
-                                            const scalar_t& y) {
-  return act_first ? ACT_FN(x) * y : x * ACT_FN(y);
+                                            const scalar_t& y,
+                                            const float limit) {
+  if constexpr (act_first) {
+    scalar_t gate = x;
+    scalar_t up = y;
+    if constexpr (HAS_CLAMP) {
+      gate = (scalar_t)fminf((float)gate, limit);
+      up = (scalar_t)fmaxf(fminf((float)up, limit), -limit);
+    }
+    return ACT_FN(gate) * up;
+  } else {
+    scalar_t gate = x;
+    scalar_t up = y;
+    if constexpr (HAS_CLAMP) {
+      gate = (scalar_t)fmaxf(fminf((float)gate, limit), -limit);
+      up = (scalar_t)fminf((float)up, limit);
+    }
+    return gate * ACT_FN(up);
+  }
 }

 template <typename packed_t, packed_t (*PACKED_ACT_FN)(const packed_t&),
-          bool act_first>
+          bool act_first, bool HAS_CLAMP>
 __device__ __forceinline__ packed_t packed_compute(const packed_t& x,
-                                                   const packed_t& y) {
-  return act_first ? packed_mul(PACKED_ACT_FN(x), y)
-                   : packed_mul(x, PACKED_ACT_FN(y));
+                                                   const packed_t& y,
+                                                   const float limit) {
+  if constexpr (act_first) {
+    packed_t gate = x;
+    packed_t up = y;
+    if constexpr (HAS_CLAMP) {
+      float2 g = cast_to_float2(gate);
+      float2 u = cast_to_float2(up);
+      g.x = fminf(g.x, limit);
+      g.y = fminf(g.y, limit);
+      u.x = fmaxf(fminf(u.x, limit), -limit);
+      u.y = fmaxf(fminf(u.y, limit), -limit);
+      gate = cast_to_packed<packed_t>(g);
+      up = cast_to_packed<packed_t>(u);
+    }
+    return packed_mul(PACKED_ACT_FN(gate), up);
+  } else {
+    packed_t gate = x;
+    packed_t up = y;
+    if constexpr (HAS_CLAMP) {
+      float2 g = cast_to_float2(gate);
+      float2 u = cast_to_float2(up);
+      g.x = fmaxf(fminf(g.x, limit), -limit);
+      g.y = fmaxf(fminf(g.y, limit), -limit);
+      u.x = fminf(u.x, limit);
+      u.y = fminf(u.y, limit);
+      gate = cast_to_packed<packed_t>(g);
+      up = cast_to_packed<packed_t>(u);
+    }
+    return packed_mul(gate, PACKED_ACT_FN(up));
+  }
 }

 // Activation and gating kernel template.
 template <typename scalar_t, typename packed_t,
          scalar_t (*ACT_FN)(const scalar_t&),
          packed_t (*PACKED_ACT_FN)(const packed_t&), bool act_first,
-          bool use_vec, bool use_256b = false>
+          bool use_vec, bool HAS_CLAMP, bool use_256b = false>
 __global__ void act_and_mul_kernel(
    scalar_t* __restrict__ out,          // [..., d]
    const scalar_t* __restrict__ input,  // [..., 2, d]
-    const int d) {
+    const int d, const float limit) {
  const scalar_t* x_ptr = input + blockIdx.x * 2 * d;
  const scalar_t* y_ptr = x_ptr + d;
  scalar_t* out_ptr = out + blockIdx.x * d;
@@ -58,8 +103,9 @@ __global__ void act_and_mul_kernel(
      }
 #pragma unroll
      for (int j = 0; j < pvec_t::NUM_ELTS; j++) {
-        x.elts[j] = packed_compute<packed_t, PACKED_ACT_FN, act_first>(
-            x.elts[j], y.elts[j]);
+        x.elts[j] =
+            packed_compute<packed_t, PACKED_ACT_FN, act_first, HAS_CLAMP>(
+                x.elts[j], y.elts[j], limit);
      }
      if constexpr (use_256b) {
        st256(x, &out_vec[i]);
@@ -72,7 +118,8 @@ __global__ void act_and_mul_kernel(
    for (int64_t idx = threadIdx.x; idx < d; idx += blockDim.x) {
      const scalar_t x = VLLM_LDG(&x_ptr[idx]);
      const scalar_t y = VLLM_LDG(&y_ptr[idx]);
-      out_ptr[idx] = compute<scalar_t, ACT_FN, act_first>(x, y);
+      out_ptr[idx] =
+          compute<scalar_t, ACT_FN, act_first, HAS_CLAMP>(x, y, limit);
    }
  }
 }
@@ -151,8 +198,11 @@ packed_gelu_tanh_kernel(const packed_t& val) {

 // Launch activation and gating kernel.
 // Use ACT_FIRST (bool) indicating whether to apply the activation function
-// first.
-#define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL, PACKED_KERNEL, ACT_FIRST)        \
+// first. HAS_CLAMP (bool) enables pre-activation clamping: gate input is
+// clamped (max only) and up input is clamped (both sides) before the
+// activation function is applied.
+#define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL, PACKED_KERNEL, ACT_FIRST,        \
+                                      HAS_CLAMP, LIMIT)                        \
  auto dtype = input.scalar_type();                                            \
  int d = input.size(-1) / 2;                                                  \
  int64_t num_tokens = input.numel() / input.size(-1);                         \
@@ -177,8 +227,8 @@ packed_gelu_tanh_kernel(const packed_t& val) {
            scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type,      \
            KERNEL<scalar_t>,                                                  \
            PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>, \
-            ACT_FIRST, true, true><<<grid, block, 0, stream>>>(                \
-            out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d);          \
+            ACT_FIRST, true, HAS_CLAMP, true><<<grid, block, 0, stream>>>(     \
+            out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d, LIMIT);   \
      });                                                                      \
    } else {                                                                   \
      VLLM_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel", [&] {          \
@@ -186,8 +236,8 @@ packed_gelu_tanh_kernel(const packed_t& val) {
            scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type,      \
            KERNEL<scalar_t>,                                                  \
            PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>, \
-            ACT_FIRST, true, false><<<grid, block, 0, stream>>>(               \
-            out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d);          \
+            ACT_FIRST, true, HAS_CLAMP, false><<<grid, block, 0, stream>>>(    \
+            out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d, LIMIT);   \
      });                                                                      \
    }                                                                          \
  } else {                                                                     \
@@ -197,8 +247,8 @@ packed_gelu_tanh_kernel(const packed_t& val) {
          scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type,        \
          KERNEL<scalar_t>,                                                    \
          PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>,   \
-          ACT_FIRST, false><<<grid, block, 0, stream>>>(                       \
-          out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d);            \
+          ACT_FIRST, false, HAS_CLAMP><<<grid, block, 0, stream>>>(            \
+          out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d, LIMIT);     \
    });                                                                        \
  }

@@ -206,7 +256,14 @@ void silu_and_mul(torch::Tensor& out,    // [..., d]
                  torch::Tensor& input)  // [..., 2 * d]
 {
  LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel,
-                                true);
+                                true, false, 0.0f);
+}
+
+void silu_and_mul_clamp(torch::Tensor& out,    // [..., d]
+                        torch::Tensor& input,  // [..., 2 * d]
+                        double limit) {
+  LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel,
+                                true, true, (float)limit);
 }

 void mul_and_silu(torch::Tensor& out,    // [..., d]
@@ -215,21 +272,21 @@ void mul_and_silu(torch::Tensor& out,    // [..., d]
  // The difference between mul_and_silu and silu_and_mul is that mul_and_silu
  // applies the silu to the latter half of the input.
  LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel,
-                                false);
+                                false, false, 0.0f);
 }

 void gelu_and_mul(torch::Tensor& out,    // [..., d]
                  torch::Tensor& input)  // [..., 2 * d]
 {
  LAUNCH_ACTIVATION_GATE_KERNEL(vllm::gelu_kernel, vllm::packed_gelu_kernel,
-                                true);
+                                true, false, 0.0f);
 }

 void gelu_tanh_and_mul(torch::Tensor& out,    // [..., d]
                       torch::Tensor& input)  // [..., 2 * d]
 {
-  LAUNCH_ACTIVATION_GATE_KERNEL(vllm::gelu_tanh_kernel,
-                                vllm::packed_gelu_tanh_kernel, true);
+  LAUNCH_ACTIVATION_GATE_KERNEL(
+      vllm::gelu_tanh_kernel, vllm::packed_gelu_tanh_kernel, true, false, 0.0f);
 }

 namespace vllm {
@@ -163,6 +163,8 @@ void rotary_embedding(torch::Tensor& positions, torch::Tensor& query,

 void silu_and_mul(torch::Tensor& out, torch::Tensor& input);

+void silu_and_mul_clamp(torch::Tensor& out, torch::Tensor& input, double limit);
+
 void silu_and_mul_quant(torch::Tensor& out, torch::Tensor& input,
                        torch::Tensor& scale);

@@ -106,6 +106,12 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
  ops.def("silu_and_mul(Tensor! result, Tensor input) -> ()");
  ops.impl("silu_and_mul", torch::kCUDA, &silu_and_mul);

+  // SwiGLU activation with input clamping.
+  ops.def(
+      "silu_and_mul_with_clamp(Tensor! result, Tensor input, float limit) "
+      "-> ()");
+  ops.impl("silu_and_mul_with_clamp", torch::kCUDA, &silu_and_mul_clamp);
+
  ops.def(
      "silu_and_mul_quant(Tensor! result, Tensor input, Tensor scale) -> ()");
  ops.impl("silu_and_mul_quant", torch::kCUDA, &silu_and_mul_quant);
@@ -540,7 +540,9 @@ RUN CUDA_VERSION_DASH=$(echo $CUDA_VERSION | cut -d. -f1,2 | tr '.' '-') && \
        libcurand-dev-${CUDA_VERSION_DASH} \
        libcublas-${CUDA_VERSION_DASH} \
        # Required by fastsafetensors (fixes #20384)
-        libnuma-dev && \
+        libnuma-dev \
+        # numactl CLI for NUMA binding at runtime
+        numactl && \
    # Fixes nccl_allocator requiring nccl.h at runtime
    # https://github.com/vllm-project/vllm/blob/1336a1ea244fa8bfd7e72751cabbdb5b68a0c11a/vllm/distributed/device_communicators/pynccl_allocator.py#L22
    # NCCL packages don't use the cuda-MAJOR-MINOR naming convention,
@@ -292,10 +292,10 @@ Pooling models now support token-wise task.

 ### Score task

-`score` task have has been removed in v0.21, use `classify` instead. Only when a classification model outputs num_labels
-equal to 1 can it be used as a scoring model and have its scoring API enabled.
+`score` task is deprecated and will be removed in v0.20. Please use `classify` instead. Only when a
+classification model outputs num_labels equal to 1 can it be used as a scoring model and have its scoring API enabled.

 ### Pooling multitask support

-Pooling multitask support has been removed in v0.21. When the default pooling task is not what you want,
+Pooling multitask support is deprecated and will be removed in v0.20. When the default pooling task is not what you want,
 you need to manually specify it via `PoolerConfig(task=<task>)` offline or `--pooler-config.task <task>` online.
@@ -4,74 +4,68 @@
 import torch

 from vllm import LLM
-from vllm.config import PoolerConfig
 from vllm.inputs import TextPrompt
 from vllm.multimodal.utils import fetch_image

+# Initialize model
+model = LLM(
+    model="jinaai/jina-embeddings-v4-vllm-text-matching",
+    runner="pooling",
+    max_model_len=1024,
+    gpu_memory_utilization=0.8,
+)

-def main():
-    # Initialize model
-    model = LLM(
-        model="jinaai/jina-embeddings-v4-vllm-text-matching",
-        pooler_config=PoolerConfig(task="token_embed"),
-        runner="pooling",
-        max_model_len=1024,
-        gpu_memory_utilization=0.8,
-    )
+# Create text prompts
+text1 = "Ein wunderschöner Sonnenuntergang am Strand"
+text1_prompt = TextPrompt(prompt=f"Query: {text1}")

-    # Create text prompts
-    text1 = "Ein wunderschöner Sonnenuntergang am Strand"
-    text1_prompt = TextPrompt(prompt=f"Query: {text1}")
+text2 = "浜辺に沈む美しい夕日"
+text2_prompt = TextPrompt(prompt=f"Query: {text2}")

-    text2 = "浜辺に沈む美しい夕日"
-    text2_prompt = TextPrompt(prompt=f"Query: {text2}")
+# Create image prompt
+image = fetch_image(
+    "https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/eskimo.jpg"  # noqa: E501
+)
+image_prompt = TextPrompt(
+    prompt="<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>Describe the image.<|im_end|>\n",  # noqa: E501
+    multi_modal_data={"image": image},
+)

-    # Create image prompt
-    image = fetch_image(
-        "https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/eskimo.jpg"  # noqa: E501
-    )
-    image_prompt = TextPrompt(
-        prompt="<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>Describe the image.<|im_end|>\n",  # noqa: E501
-        multi_modal_data={"image": image},
-    )
-
-    # Encode all prompts
-    prompts = [text1_prompt, text2_prompt, image_prompt]
-    outputs = model.encode(prompts, pooling_task="token_embed")
-
-    def get_embeddings(outputs):
-        VISION_START_TOKEN_ID, VISION_END_TOKEN_ID = 151652, 151653
-
-        embeddings = []
-        for output in outputs:
-            if VISION_START_TOKEN_ID in output.prompt_token_ids:
-                # Gather only vision tokens
-                img_start_pos = torch.where(
-                    torch.tensor(output.prompt_token_ids) == VISION_START_TOKEN_ID
-                )[0][0]
-                img_end_pos = torch.where(
-                    torch.tensor(output.prompt_token_ids) == VISION_END_TOKEN_ID
-                )[0][0]
-                embeddings_tensor = output.outputs.data.detach().clone()[
-                    img_start_pos : img_end_pos + 1
-                ]
-            else:
-                # Use all tokens for text-only prompts
-                embeddings_tensor = output.outputs.data.detach().clone()
-
-            # Pool and normalize embeddings
-            pooled_output = (
-                embeddings_tensor.sum(dim=0, dtype=torch.float32)
-                / embeddings_tensor.shape[0]
-            )
-            embeddings.append(torch.nn.functional.normalize(pooled_output, dim=-1))
-        return embeddings
-
-    embeddings = get_embeddings(outputs)
-
-    for embedding in embeddings:
-        print(embedding.shape)
+# Encode all prompts
+prompts = [text1_prompt, text2_prompt, image_prompt]
+outputs = model.encode(prompts, pooling_task="token_embed")


-if __name__ == "__main__":
-    main()
+def get_embeddings(outputs):
+    VISION_START_TOKEN_ID, VISION_END_TOKEN_ID = 151652, 151653
+
+    embeddings = []
+    for output in outputs:
+        if VISION_START_TOKEN_ID in output.prompt_token_ids:
+            # Gather only vision tokens
+            img_start_pos = torch.where(
+                torch.tensor(output.prompt_token_ids) == VISION_START_TOKEN_ID
+            )[0][0]
+            img_end_pos = torch.where(
+                torch.tensor(output.prompt_token_ids) == VISION_END_TOKEN_ID
+            )[0][0]
+            embeddings_tensor = output.outputs.data.detach().clone()[
+                img_start_pos : img_end_pos + 1
+            ]
+        else:
+            # Use all tokens for text-only prompts
+            embeddings_tensor = output.outputs.data.detach().clone()
+
+        # Pool and normalize embeddings
+        pooled_output = (
+            embeddings_tensor.sum(dim=0, dtype=torch.float32)
+            / embeddings_tensor.shape[0]
+        )
+        embeddings.append(torch.nn.functional.normalize(pooled_output, dim=-1))
+    return embeddings
+
+
+embeddings = get_embeddings(outputs)
+
+for embedding in embeddings:
+    print(embedding.shape)
@@ -4,7 +4,6 @@
 from argparse import Namespace

 from vllm import LLM, EngineArgs
-from vllm.config import PoolerConfig
 from vllm.utils.argparse_utils import FlexibleArgumentParser


@@ -14,7 +13,6 @@ def parse_args():
    # Set example specific arguments
    parser.set_defaults(
        model="BAAI/bge-m3",
-        pooler_config=PoolerConfig(task="token_embed"),
        runner="pooling",
        enforce_eager=True,
    )
@@ -34,6 +32,15 @@ def main(args: Namespace):
    # You should pass runner="pooling" for embedding models
    llm = LLM(**vars(args))

+    # Generate embedding. The output is a list of EmbeddingRequestOutputs.
+    outputs = llm.embed(prompts)
+
+    # Print the outputs.
+    print("\nGenerated Outputs:\n" + "-" * 60)
+    for prompt, output in zip(prompts, outputs):
+        embeds = output.outputs.embedding
+        print(len(embeds))
+
    # Generate embedding for each token. The output is a list of PoolingRequestOutput.
    outputs = llm.encode(prompts, pooling_task="token_embed")

@@ -43,20 +50,6 @@ def main(args: Namespace):
        multi_vector = output.outputs.data
        print(multi_vector.shape)

-    query = "What is the capital of France?"
-    documents = [
-        "The capital of Brazil is Brasilia.",
-        "The capital of France is Paris.",
-    ]
-    # Generate scores.
-    outputs = llm.score(query, documents)
-    # Print the outputs.
-    print("\nGenerated Outputs:\n" + "-" * 60)
-    for document, output in zip(documents, outputs):
-        score = output.outputs.score
-        print(f"Pair: {[query, document]!r} \nScore: {score}")
-        print("-" * 60)
-

 if __name__ == "__main__":
    args = parse_args()
@@ -7,11 +7,10 @@ Example online usage of Pooling API for multi vector retrieval.
 Run `vllm serve <model> --runner pooling`
 to start up the server in vLLM. e.g.

-vllm serve BAAI/bge-m3 --pooler-config.task token_embed
+vllm serve BAAI/bge-m3
 """

 import argparse
-import pprint

 import requests
 import torch
@@ -33,8 +32,7 @@ def parse_args():


 def main(args):
-    pooling_url = f"http://{args.host}:{args.port}/pooling"
-    score_url = f"http://{args.host}:{args.port}/score"
+    api_url = f"http://{args.host}:{args.port}/pooling"
    model_name = args.model

    prompts = [
@@ -45,23 +43,11 @@ def main(args):
    ]
    prompt = {"model": model_name, "input": prompts}

-    pooling_response = post_http_request(prompt=prompt, api_url=pooling_url)
+    pooling_response = post_http_request(prompt=prompt, api_url=api_url)
    for output in pooling_response.json()["data"]:
        multi_vector = torch.tensor(output["data"])
        print(multi_vector.shape)

-    queries = "What is the capital of France?"
-    documents = [
-        "The capital of Brazil is Brasilia.",
-        "The capital of France is Paris.",
-    ]
-    prompt = {"model": model_name, "queries": queries, "documents": documents}
-    score_response = post_http_request(prompt=prompt, api_url=score_url)
-    print("\nPrompt when queries is string and documents is a list:")
-    pprint.pprint(prompt)
-    print("\nScore Response:")
-    pprint.pprint(score_response.json())
-

 if __name__ == "__main__":
    args = parse_args()
@@ -1,12 +1,13 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import logging
 import weakref

 import pytest
 import torch

 from tests.models.utils import softmax
-from vllm import LLM, ClassificationRequestOutput, PoolingParams
+from vllm import LLM, ClassificationRequestOutput, PoolingParams, PoolingRequestOutput
 from vllm.distributed import cleanup_dist_env_and_memory
 from vllm.tasks import PoolingTask

@@ -65,6 +66,18 @@ def test_list_prompts(llm: LLM):
        assert len(outputs[i].outputs.probs) == num_labels


+@pytest.mark.skip_global_cleanup
+def test_token_classify(llm: LLM, caplog_vllm):
+    with caplog_vllm.at_level(level=logging.WARNING, logger="vllm"):
+        outputs = llm.encode(prompt, pooling_task="token_classify", use_tqdm=False)
+        assert "deprecated" in caplog_vllm.text
+
+    assert len(outputs) == 1
+    assert isinstance(outputs[0], PoolingRequestOutput)
+    assert outputs[0].prompt_token_ids == prompt_token_ids
+    assert outputs[0].outputs.data.shape == (len(prompt_token_ids), num_labels)
+
+
@pytest.mark.skip_global_cleanup
 def test_pooling_params(llm: LLM):
    def get_outputs(use_activation):
@@ -97,12 +110,10 @@ def test_score_api(llm: LLM):
        llm.score("ping", "pong", use_tqdm=False)


-@pytest.mark.parametrize("task", ["embed", "token_embed", "token_classify", "plugin"])
+@pytest.mark.parametrize("task", ["embed", "token_embed", "plugin"])
 def test_unsupported_tasks(llm: LLM, task: PoolingTask):
    if task == "plugin":
        err_msg = "No IOProcessor plugin installed."
-    elif task == "token_classify":
-        err_msg = "Try switching the model's pooling_task via.+"
    else:
        err_msg = "Embedding API is not supported by this model.+"
    with pytest.raises(ValueError, match=err_msg):
@@ -436,7 +436,26 @@ async def test_pooling_classify(server: RemoteOpenAIServer, model_name: str):

@pytest.mark.asyncio
@pytest.mark.parametrize("model_name", [MODEL_NAME])
-@pytest.mark.parametrize("task", ["embed", "token_embed", "token_classify", "plugin"])
+async def test_pooling_token_classify(server: RemoteOpenAIServer, model_name: str):
+    task = "token_classify"
+    response = requests.post(
+        server.url_for("pooling"),
+        json={
+            "model": model_name,
+            "input": input_text,
+            "encoding_format": "float",
+            "task": task,
+        },
+    )
+    poolings = PoolingResponse.model_validate(response.json())
+    assert len(poolings.data) == 1
+    assert len(poolings.data[0].data) == 8
+    assert len(poolings.data[0].data[0]) == 2
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("model_name", [MODEL_NAME])
+@pytest.mark.parametrize("task", ["embed", "token_embed", "plugin"])
 async def test_pooling_not_supported(
    server: RemoteOpenAIServer, model_name: str, task: str
 ):
@@ -450,11 +469,8 @@ async def test_pooling_not_supported(
        },
    )
    assert response.json()["error"]["type"] == "BadRequestError"
-
    if task == "plugin":
        err_msg = "No IOProcessor plugin installed."
-    elif task == "token_classify":
-        err_msg = "Try switching the model's pooling_task via"
    else:
        err_msg = f"Unsupported task: {task!r}"
    assert response.json()["error"]["message"].startswith(err_msg)
@@ -1,5 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import logging
 import weakref

 import pytest
@@ -37,11 +38,11 @@ def llm():
        seed=0,
        attention_config=attention_config,
    )
-    assert embedding_size == llm.model_config.embedding_size

    yield weakref.proxy(llm)

    del llm
+
    cleanup_dist_env_and_memory()


@@ -73,6 +74,16 @@ def test_list_prompts(llm: LLM):
        assert len(outputs[i].outputs.embedding) == embedding_size


+@pytest.mark.skip_global_cleanup
+def test_token_embed(llm: LLM, caplog_vllm):
+    with caplog_vllm.at_level(level=logging.WARNING, logger="vllm"):
+        outputs = llm.encode(prompt, pooling_task="token_embed", use_tqdm=False)
+        assert "deprecated" in caplog_vllm.text
+
+    multi_vector = outputs[0].outputs.data
+    assert multi_vector.shape == (11, 384)
+
+
@pytest.mark.skip_global_cleanup
 def test_pooling_params(llm: LLM):
    def get_outputs(normalize):
@@ -96,14 +107,10 @@ def test_pooling_params(llm: LLM):
    )


-@pytest.mark.parametrize(
-    "task", ["token_classify", "classify", "token_embed", "plugin"]
-)
+@pytest.mark.parametrize("task", ["token_classify", "classify", "plugin"])
 def test_unsupported_tasks(llm: LLM, task: PoolingTask):
    if task == "plugin":
        err_msg = "No IOProcessor plugin installed."
-    elif task == "token_embed":
-        err_msg = "Try switching the model's pooling_task via.+"
    else:
        err_msg = "Classification API is not supported by this model.+"
    with pytest.raises(ValueError, match=err_msg):
@@ -732,9 +732,28 @@ async def test_pooling_embed(server: RemoteOpenAIServer, model_name: str):

@pytest.mark.asyncio
@pytest.mark.parametrize("model_name", [MODEL_NAME])
-@pytest.mark.parametrize(
-    "task", ["classify", "token_classify", "token_embed", "plugin"]
-)
+async def test_pooling_token_embed(server: RemoteOpenAIServer, model_name: str):
+    task = "token_embed"
+    response = requests.post(
+        server.url_for("pooling"),
+        json={
+            "model": model_name,
+            "input": input_text,
+            "encoding_format": "float",
+            "task": task,
+        },
+    )
+
+    poolings = PoolingResponse.model_validate(response.json())
+
+    assert len(poolings.data) == 1
+    assert len(poolings.data[0].data) == len(input_tokens)
+    assert len(poolings.data[0].data[0]) == 384
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("model_name", [MODEL_NAME])
+@pytest.mark.parametrize("task", ["classify", "token_classify", "plugin"])
 async def test_pooling_not_supported(
    server: RemoteOpenAIServer, model_name: str, task: str
 ):
@@ -750,8 +769,6 @@ async def test_pooling_not_supported(
    assert response.json()["error"]["type"] == "BadRequestError"
    if task == "plugin":
        err_msg = "No IOProcessor plugin installed."
-    elif task == "token_embed":
-        err_msg = "Try switching the model's pooling_task via"
    else:
        err_msg = f"Unsupported task: {task!r}"
    assert response.json()["error"]["message"].startswith(err_msg)
@@ -452,6 +452,25 @@ async def test_pooling_classify(server: RemoteOpenAIServer):
    assert len(poolings.data[0].data) == 1


+@pytest.mark.asyncio
+async def test_pooling_token_classify(server: RemoteOpenAIServer):
+    response = requests.post(
+        server.url_for("pooling"),
+        json={
+            "model": MODEL_NAME,
+            "task": "token_classify",
+            "input": input_text,
+            "encoding_format": "float",
+        },
+    )
+
+    poolings = PoolingResponse.model_validate(response.json())
+
+    assert len(poolings.data) == 1
+    assert len(poolings.data[0].data) == len(input_tokens)
+    assert len(poolings.data[0].data[0]) == 1
+
+
@pytest.mark.asyncio
 async def test_rerank_max_tokens_per_doc(
    server: RemoteOpenAIServer,
@@ -525,7 +544,7 @@ async def test_rerank_max_tokens_per_doc_validation(


@pytest.mark.asyncio
-@pytest.mark.parametrize("task", ["embed", "token_embed", "token_classify", "plugin"])
+@pytest.mark.parametrize("task", ["embed", "token_embed", "plugin"])
 async def test_pooling_not_supported(server: RemoteOpenAIServer, task: str):
    response = requests.post(
        server.url_for("pooling"),
@@ -539,8 +558,6 @@ async def test_pooling_not_supported(server: RemoteOpenAIServer, task: str):
    assert response.json()["error"]["type"] == "BadRequestError"
    if task == "plugin":
        err_msg = "No IOProcessor plugin installed."
-    elif task == "token_classify":
-        err_msg = "Try switching the model's pooling_task via"
    else:
        err_msg = f"Unsupported task: {task!r}"
    assert response.json()["error"]["message"].startswith(err_msg)
@@ -1,5 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import logging
 import weakref

 import pytest
@@ -59,19 +60,22 @@ def test_token_ids_prompts(llm: LLM):

@pytest.mark.skip_global_cleanup
 def test_score_api(llm: LLM):
-    err_msg = "This model does not support the Scoring API."
+    err_msg = "Scoring API is only enabled for num_labels == 1."
    with pytest.raises(ValueError, match=err_msg):
        llm.score("ping", "pong", use_tqdm=False)


@pytest.mark.parametrize("task", ["classify", "embed", "token_embed", "plugin"])
 def test_unsupported_tasks(llm: LLM, task: PoolingTask, caplog_vllm):
-    if task == "plugin":
-        err_msg = "No IOProcessor plugin installed."
-    elif task == "classify":
-        err_msg = "Try switching the model's pooling_task via.+"
+    if task == "classify":
+        with caplog_vllm.at_level(level=logging.WARNING, logger="vllm"):
+            llm.encode(prompt, pooling_task=task, use_tqdm=False)
+        assert "deprecated" in caplog_vllm.text
    else:
-        err_msg = "Embedding API is not supported by this model.+"
+        if task == "plugin":
+            err_msg = "No IOProcessor plugin installed."
+        else:
+            err_msg = "Embedding API is not supported by this model.+"

-    with pytest.raises(ValueError, match=err_msg):
-        llm.encode(prompt, pooling_task=task, use_tqdm=False)
+        with pytest.raises(ValueError, match=err_msg):
+            llm.encode(prompt, pooling_task=task, use_tqdm=False)
@@ -50,7 +50,7 @@ async def test_pooling_token_classify(server: RemoteOpenAIServer, model_name: st

@pytest.mark.asyncio
@pytest.mark.parametrize("model_name", [MODEL_NAME])
-@pytest.mark.parametrize("task", ["classify", "embed", "token_embed", "plugin"])
+@pytest.mark.parametrize("task", ["embed", "token_embed", "plugin"])
 async def test_pooling_not_supported(
    server: RemoteOpenAIServer, model_name: str, task: str
 ):
@@ -63,12 +63,9 @@ async def test_pooling_not_supported(
            "task": task,
        },
    )
-    assert response.json()["error"]["type"] == "BadRequestError"

    if task == "plugin":
        err_msg = "No IOProcessor plugin installed."
-    elif task == "classify":
-        err_msg = "Try switching the model's pooling_task via"
    else:
        err_msg = f"Unsupported task: {task!r}"
    assert response.json()["error"]["message"].startswith(err_msg)
@@ -1,5 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import logging
 import weakref

 import pytest
@@ -63,12 +64,15 @@ def test_token_ids_prompts(llm: LLM):

@pytest.mark.parametrize("task", ["embed", "classify", "token_classify", "plugin"])
 def test_unsupported_tasks(llm: LLM, task: PoolingTask, caplog_vllm):
-    if task == "plugin":
-        err_msg = "No IOProcessor plugin installed."
-    elif task == "embed":
-        err_msg = "Try switching the model's pooling_task via.+"
+    if task == "embed":
+        with caplog_vllm.at_level(level=logging.WARNING, logger="vllm"):
+            llm.encode(prompt, pooling_task=task, use_tqdm=False)
+        assert "deprecated" in caplog_vllm.text
    else:
-        err_msg = "Classification API is not supported by this model.+"
+        if task == "plugin":
+            err_msg = "No IOProcessor plugin installed."
+        else:
+            err_msg = "Classification API is not supported by this model.+"

-    with pytest.raises(ValueError, match=err_msg):
-        llm.encode(prompt, pooling_task=task, use_tqdm=False)
+        with pytest.raises(ValueError, match=err_msg):
+            llm.encode(prompt, pooling_task=task, use_tqdm=False)
@@ -73,7 +73,7 @@ async def test_pooling_token_embed(server: RemoteOpenAIServer, model_name: str):

@pytest.mark.asyncio
@pytest.mark.parametrize("model_name", [MODEL_NAME])
-@pytest.mark.parametrize("task", ["embed", "classify", "token_classify", "plugin"])
+@pytest.mark.parametrize("task", ["classify", "token_classify", "plugin"])
 async def test_pooling_not_supported(
    server: RemoteOpenAIServer, model_name: str, task: str
 ):
@@ -86,12 +86,9 @@ async def test_pooling_not_supported(
            "task": task,
        },
    )
-    assert response.json()["error"]["type"] == "BadRequestError"

    if task == "plugin":
        err_msg = "No IOProcessor plugin installed."
-    elif task == "embed":
-        err_msg = "Try switching the model's pooling_task via"
    else:
        err_msg = f"Unsupported task: {task!r}"
    assert response.json()["error"]["message"].startswith(err_msg)
@@ -16,6 +16,7 @@ from vllm.model_executor.layers.activation import (
    NewGELU,
    QuickGELU,
    SiluAndMul,
+    SiluAndMulWithClamp,
    SwigluOAIAndMul,
    SwigluStepAndMul,
    swiglustep_and_mul_triton,
@@ -116,6 +117,85 @@ def test_act_and_mul(
        opcheck(fn, (out, x))


+SWIGLU_LIMITS = [3.0, 7.0, 15.0]
+
+
+@pytest.mark.parametrize("swiglu_limit", SWIGLU_LIMITS)
+@pytest.mark.parametrize("num_tokens", NUM_TOKENS)
+@pytest.mark.parametrize("d", D)
+@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize("seed", SEEDS)
+@pytest.mark.parametrize("device", CUDA_DEVICES)
+@torch.inference_mode()
+def test_silu_and_mul_with_clamp(
+    default_vllm_config,
+    swiglu_limit: float,
+    num_tokens: int,
+    d: int,
+    dtype: torch.dtype,
+    seed: int,
+    device: str,
+) -> None:
+    """SiluAndMulWithClamp: cuda kernel must match native reference."""
+    set_random_seed(seed)
+    torch.set_default_device(device)
+    # Use large values to ensure clamping is exercised.
+    x = torch.randn(num_tokens, 2 * d, dtype=dtype) * swiglu_limit * 2
+
+    layer = SiluAndMulWithClamp(swiglu_limit, compile_native=False)
+    out = layer(x)
+    ref_out = layer.forward_native(x)
+
+    rtol = {
+        torch.float16: 2e-3,
+        torch.bfloat16: 2e-2,
+        torch.float: 1.3e-6,
+    }
+    torch.testing.assert_close(
+        out, ref_out, atol=get_default_atol(out), rtol=rtol[out.dtype]
+    )
+
+    # Verify clamping is actually being applied: the clamped output should
+    # differ from the unclamped SiluAndMul output when inputs are large.
+    unclamped_out = SiluAndMul.forward_native(x)
+    assert not torch.equal(ref_out.float(), unclamped_out.float()), (
+        "Input was not large enough to exercise the clamp; increase scale"
+    )
+
+    # Verify gate clamping semantics with a controlled scalar case.
+    # gate=large_val is clamped to limit first, then silu(limit) * 1.0.
+    x_gate = torch.tensor(
+        [[swiglu_limit * 20.0, 1.0]], dtype=torch.float32, device=device
+    )
+    out_gate = SiluAndMulWithClamp(swiglu_limit, compile_native=False)(x_gate)
+    expected_gate = torch.nn.functional.silu(
+        torch.tensor(swiglu_limit, dtype=torch.float32)
+    ).item()
+    torch.testing.assert_close(
+        out_gate,
+        torch.tensor([[expected_gate]], dtype=torch.float32, device=device),
+        atol=1e-3,
+        rtol=1e-3,
+    )
+
+    # Verify up clamping semantics: up >> limit gets clamped to limit.
+    x_up = torch.tensor(
+        [[1.0, swiglu_limit * 20.0]], dtype=torch.float32, device=device
+    )
+    out_up = SiluAndMulWithClamp(swiglu_limit, compile_native=False)(x_up)
+    silu_1 = torch.nn.functional.silu(torch.tensor(1.0)).item()
+    torch.testing.assert_close(
+        out_up,
+        torch.tensor([[silu_1 * swiglu_limit]], dtype=torch.float32, device=device),
+        atol=1e-3,
+        rtol=1e-3,
+    )
+
+    # opcheck
+    out_buf = torch.empty(x.shape[:-1] + (d,), dtype=dtype, device=device)
+    opcheck(torch.ops._C.silu_and_mul_with_clamp, (out_buf, x, swiglu_limit))
+
+
@pytest.mark.parametrize(
    "activation",
    [
@@ -6,7 +6,6 @@ from transformers import AutoModel

 from tests.models.utils import check_embeddings_close
 from vllm import TokensPrompt
-from vllm.config import PoolerConfig


@pytest.mark.parametrize(
@@ -22,7 +21,6 @@ def test_embed_models(hf_runner, vllm_runner, model: str):
    with vllm_runner(
        model,
        runner="pooling",
-        pooler_config=PoolerConfig(task="token_embed"),
        max_model_len=128,
        max_num_batched_tokens=chunk_size,
        enforce_eager=True,
@@ -3,6 +3,7 @@
 import httpx
 import openai
 import pytest
+import pytest_asyncio
 import torch

 from ....utils import RemoteOpenAIServer
@@ -24,42 +25,29 @@ sentences_2 = [
 similarity_reference = [[0.6259, 0.3474], [0.3309, 0.6734]]
 lexical_score_reference = [0.19554901123046875, 0.0]
 colbert_score_reference = [0.7797, 0.4620]
-SUPPORTED_TASKS = ["embed", "token_embed", "token_classify"]
-
-
-@pytest.fixture(scope="module", params=SUPPORTED_TASKS)
-def pooling_task(request):
-    yield request.param


@pytest.fixture(scope="module")
-def server(pooling_task):
+def server():
    args = [
        "--max-model-len",
        str(MAX_MODEL_LEN),
        "--hf-overrides",
        '{"architectures": ["BgeM3EmbeddingModel"]}',
-        "--pooler-config.task",
-        pooling_task,
    ]

    with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
        yield remote_server


+@pytest_asyncio.fixture
+async def client(server):
+    async with server.get_async_client() as async_client:
+        yield async_client
+
+
@pytest.mark.asyncio
-async def test_bge_m3_api_server_embedding(server, pooling_task):
-    client = server.get_async_client()
-
-    if pooling_task != "embed":
-        with pytest.raises(openai.InternalServerError):
-            await run_client_embeddings(
-                client,
-                MODEL_NAME,
-                sentences_1,
-            )
-        return
-
+async def test_bge_m3_api_server_embedding(client: openai.AsyncOpenAI):
    embeddings_list_1 = await run_client_embeddings(
        client,
        MODEL_NAME,
@@ -129,14 +117,7 @@ def compute_lexical_matching_score(


@pytest.mark.asyncio
-async def test_bge_m3_api_server_sparse_embedding(server, pooling_task):
-    client = server.get_async_client()
-
-    if pooling_task != "token_classify":
-        with pytest.raises(openai.BadRequestError):
-            await sparse_embeddings(client, sentences_1)
-        return
-
+async def test_bge_m3_api_server_sparse_embedding(client: openai.AsyncOpenAI):
    embeddings_1 = await sparse_embeddings(client, sentences_1)
    embeddings_2 = await sparse_embeddings(client, sentences_2)

@@ -156,11 +137,9 @@ async def test_bge_m3_api_server_sparse_embedding(server, pooling_task):


@pytest.mark.asyncio
-async def test_bge_m3_api_server_sparse_embedding_corner_case(server, pooling_task):
-    if pooling_task != "token_classify":
-        return
-
-    client = server.get_async_client()
+async def test_bge_m3_api_server_sparse_embedding_corner_case(
+    client: openai.AsyncOpenAI,
+):
    embeddings = await sparse_embeddings(client, ["Hi"])
    assert len(embeddings) == 1
    assert 2673 in embeddings[0]
@@ -176,18 +155,7 @@ def colbert_score(q_reps: torch.Tensor, p_reps: torch.Tensor) -> torch.Tensor:


@pytest.mark.asyncio
-async def test_bge_m3_api_server_multi_vector(server, pooling_task):
-    client = server.get_async_client()
-
-    if pooling_task != "token_embed":
-        with pytest.raises(openai.BadRequestError):
-            await client.post(
-                "../pooling",
-                body={"model": MODEL_NAME, "input": sentences_1, "task": "token_embed"},
-                cast_to=httpx.Response,
-            )
-        return
-
+async def test_bge_m3_api_server_multi_vector(client: openai.AsyncOpenAI):
    result_1 = await client.post(
        "../pooling",
        body={"model": MODEL_NAME, "input": sentences_1, "task": "token_embed"},
@@ -4,7 +4,6 @@ import pytest
 import torch

 from vllm import TokensPrompt
-from vllm.config import PoolerConfig


@pytest.mark.parametrize(
@@ -21,7 +20,6 @@ def test_extract_hidden_states(hf_runner, vllm_runner, model: str):
        max_model_len=128,
        enforce_eager=True,
        runner="pooling",
-        pooler_config=PoolerConfig(task="token_embed"),
        enable_prefix_caching=True,
    ) as vllm_model:
        pooling_outputs = vllm_model.llm.encode(
@@ -46,3 +44,14 @@ def test_extract_hidden_states(hf_runner, vllm_runner, model: str):
            assert len(output.prompt_token_ids) == n
            assert len(output.outputs.data) == n
            assert output.num_cached_tokens == 0
+
+        # skip_reading_prefix_cache can still write to cache
+        # to accelerate following requests
+        pooling_outputs = vllm_model.llm.encode(
+            [TokensPrompt(prompt_token_ids=t) for t in token_prompts],
+            pooling_task="embed",
+        )
+
+        for n, output in zip(n_prompt_tokens, pooling_outputs):
+            assert len(output.prompt_token_ids) == n
+            assert output.num_cached_tokens > 0
@@ -5,7 +5,6 @@ import torch
 from transformers import AutoModel

 from tests.models.utils import check_embeddings_close
-from vllm.config import PoolerConfig


@pytest.mark.parametrize(
@@ -18,7 +17,6 @@ def test_embed_models(hf_runner, vllm_runner, example_prompts, model: str, dtype
    with vllm_runner(
        model,
        runner="pooling",
-        pooler_config=PoolerConfig(task="token_embed"),
        max_model_len=None,
    ) as vllm_model:
        vllm_outputs = vllm_model.token_embed(example_prompts)
@@ -146,7 +146,7 @@ def test_multi_vector_retrieval_models_using_normalize(
        model,
        max_model_len=512,
        dtype=dtype,
-        pooler_config=PoolerConfig(use_activation=False, task="token_embed"),
+        pooler_config=PoolerConfig(use_activation=False),
    ) as vllm_model:
        wo_normalize = vllm_model.token_embed(example_prompts)

@@ -154,7 +154,7 @@ def test_multi_vector_retrieval_models_using_normalize(
        model,
        max_model_len=512,
        dtype=dtype,
-        pooler_config=PoolerConfig(use_activation=True, task="token_embed"),
+        pooler_config=PoolerConfig(use_activation=True),
    ) as vllm_model:
        w_normalize = vllm_model.token_embed(example_prompts)

@@ -2512,3 +2512,111 @@ def test_block_lookup_cache_multi_blocks_per_key():
    assert cache.pop(key1, 11) is block11
    assert cache.get_one_block(key1) is None
    assert cache.pop(key1, 12) is None
+
+
+def test_can_fit_full_sequence_swa_cap_admits_long_prompt():
+    """Hybrid full+SWA model with a pool sized at the startup minimum should
+    admit a prompt longer than the SWA cap, because SlidingWindowManager
+    recycles blocks during chunked prefill (issue #39734)."""
+    block_size = 16
+    sliding_window = 4 * block_size  # 64 tokens
+    max_num_batched_tokens = 8 * block_size  # 128 tokens
+    max_model_len = 64 * block_size  # 1024 tokens — much larger than the SWA cap
+    # Startup pool sizing: full demands cdiv(max_model_len, bs) = 64 blocks,
+    # SWA demands cdiv(SW-1+max_batched, bs) + 1 = cdiv(191, 16) + 1 = 13.
+    # Pool minimum = 64 + 13 = 77; +1 for the null block.
+    num_blocks = 64 + 13 + 1
+
+    config = KVCacheConfig(
+        num_blocks=num_blocks,
+        kv_cache_tensors=[],
+        kv_cache_groups=[
+            KVCacheGroupSpec(
+                ["layer_full"],
+                FullAttentionSpec(
+                    block_size=block_size,
+                    num_kv_heads=1,
+                    head_size=1,
+                    dtype=torch.float32,
+                ),
+            ),
+            KVCacheGroupSpec(
+                ["layer_swa"],
+                SlidingWindowSpec(
+                    block_size=block_size,
+                    num_kv_heads=1,
+                    head_size=1,
+                    dtype=torch.float32,
+                    sliding_window=sliding_window,
+                ),
+            ),
+        ],
+    )
+
+    manager = KVCacheManager(
+        config,
+        max_model_len=max_model_len,
+        max_num_batched_tokens=max_num_batched_tokens,
+        enable_caching=True,
+        hash_block_size=block_size,
+    )
+
+    # A prompt that is shorter than max_model_len but longer than SW + chunk:
+    # cdiv(prompt_len, bs) = 32 blocks. Without the cap, admission would
+    # demand 32 (full) + 32 (SWA) = 64 blocks. With the cap, SWA contributes
+    # only 13, so total = 32 + 13 = 45 ≤ pool size.
+    prompt_len = 32 * block_size
+    req = make_request("long", list(range(prompt_len)), block_size, sha256)
+
+    assert manager.can_fit_full_sequence(req)
+
+
+def test_can_fit_full_sequence_full_attention_still_gates_oversized():
+    """The cap only loosens the SWA group; a prompt that exceeds the
+    full-attention pool capacity must still be rejected."""
+    block_size = 16
+    sliding_window = 4 * block_size
+    max_num_batched_tokens = 8 * block_size
+    max_model_len = 64 * block_size
+    # Provide a tiny pool — even a small prompt should be rejected.
+    num_blocks = 5
+
+    config = KVCacheConfig(
+        num_blocks=num_blocks,
+        kv_cache_tensors=[],
+        kv_cache_groups=[
+            KVCacheGroupSpec(
+                ["layer_full"],
+                FullAttentionSpec(
+                    block_size=block_size,
+                    num_kv_heads=1,
+                    head_size=1,
+                    dtype=torch.float32,
+                ),
+            ),
+            KVCacheGroupSpec(
+                ["layer_swa"],
+                SlidingWindowSpec(
+                    block_size=block_size,
+                    num_kv_heads=1,
+                    head_size=1,
+                    dtype=torch.float32,
+                    sliding_window=sliding_window,
+                ),
+            ),
+        ],
+    )
+
+    manager = KVCacheManager(
+        config,
+        max_model_len=max_model_len,
+        max_num_batched_tokens=max_num_batched_tokens,
+        enable_caching=True,
+        hash_block_size=block_size,
+    )
+
+    # 16 blocks of full attention demand alone exceeds the 5-block pool.
+    prompt_len = 16 * block_size
+    req = make_request("oversized", list(range(prompt_len)), block_size, sha256)
+
+    assert not manager.can_fit_full_sequence(req)
@@ -22,11 +22,13 @@ pytestmark = pytest.mark.cpu_test


 def get_sliding_window_manager(sliding_window_spec, block_pool, enable_caching=True):
+    # Tests don't exercise admission gating; pass a large cap that is a no-op.
    return SlidingWindowManager(
        sliding_window_spec,
        block_pool=block_pool,
        enable_caching=enable_caching,
        kv_cache_group_id=0,
+        max_admission_blocks_per_request=10**9,
    )


@@ -38,6 +40,7 @@ def get_chunked_local_attention_manager(
        block_pool=block_pool,
        enable_caching=enable_caching,
        kv_cache_group_id=0,
+        max_admission_blocks_per_request=10**9,
    )


@@ -79,7 +79,7 @@ from vllm.renderers.inputs.preprocess import (
    prompt_to_seq,
 )
 from vllm.sampling_params import BeamSearchParams, RequestOutputKind, SamplingParams
-from vllm.tasks import SCORE_TYPE_MAP, PoolingTask
+from vllm.tasks import PoolingTask
 from vllm.tokenizers import TokenizerLike
 from vllm.usage.usage_lib import UsageContext
 from vllm.utils.counter import Counter
@@ -1204,9 +1204,12 @@ class LLM:
                    f"Supported tasks: {self.supported_tasks}"
                )
            else:
-                raise ValueError(
-                    f"Try switching the model's pooling_task "
-                    f'via `PoolerConfig(task="{pooling_task}")`'
+                logger.warning_once(
+                    "Pooling multitask support is deprecated and will "
+                    "be removed in v0.20. When the default pooling task is "
+                    "not what you want, you need to manually specify it "
+                    'via PoolerConfig(task="%s"). ',
+                    pooling_task,
                )

        if pooling_task == "plugin" and "plugin" not in self.pooling_io_processors:
@@ -1409,7 +1412,7 @@ class LLM:
                "pooling model."
            )

-        score_type: str | None = SCORE_TYPE_MAP.get(self.pooling_task, None)  # type: ignore[arg-type]
+        score_type = self.model_config.score_type
        if (
            score_type == "cross-encoder"
            and getattr(self.model_config.hf_config, "num_labels", 0) != 1
@@ -15,7 +15,10 @@ from starlette.datastructures import Headers
 from vllm import PoolingParams, PoolingRequestOutput, envs
 from vllm.config import VllmConfig
 from vllm.engine.protocol import EngineClient
-from vllm.entrypoints.chat_utils import ChatTemplateConfig
+from vllm.entrypoints.chat_utils import (
+    ChatTemplateConfig,
+    ChatTemplateContentFormatOption,
+)
 from vllm.entrypoints.logger import RequestLogger
 from vllm.entrypoints.openai.engine.protocol import ErrorResponse
 from vllm.entrypoints.openai.models.serving import OpenAIServingModels
@@ -45,7 +48,9 @@ class PoolingServingBase(ABC):
        models: OpenAIServingModels,
        *,
        request_logger: RequestLogger | None,
-        chat_template_config: ChatTemplateConfig,
+        chat_template: str | None = None,
+        chat_template_content_format: ChatTemplateContentFormatOption = "auto",
+        trust_request_chat_template: bool = False,
        return_tokens_as_token_ids: bool = False,
        log_error_stack: bool = False,
    ):
@@ -58,7 +63,11 @@ class PoolingServingBase(ABC):
        self.request_logger = request_logger
        self.return_tokens_as_token_ids = return_tokens_as_token_ids
        self.log_error_stack = log_error_stack
-        self.chat_template_config = chat_template_config
+        self.chat_template_config = ChatTemplateConfig(
+            chat_template=chat_template,
+            chat_template_content_format=chat_template_content_format,
+            trust_request_chat_template=trust_request_chat_template,
+        )

        # Shared thread pool executor for preprocessing and postprocessing.
        self._executor: Executor = models.renderer._executor
@@ -10,7 +10,7 @@ from vllm.entrypoints.chat_utils import ChatTemplateConfig
 from vllm.logger import init_logger
 from vllm.plugins.io_processors import has_io_processor
 from vllm.renderers import BaseRenderer
-from vllm.tasks import POOLING_TASKS, SCORE_TYPE_MAP, SupportedTask
+from vllm.tasks import POOLING_TASKS, SupportedTask

 from .base.io_processor import PoolingIOProcessor
 from .utils import enable_scoring_api
@@ -43,24 +43,23 @@ def init_pooling_io_processors(
 ) -> dict[str, PoolingIOProcessor]:
    model_config = vllm_config.model_config
    processors: dict[str, type[PoolingIOProcessor]] = {}
-    pooling_task = model_config.get_pooling_task(supported_tasks)

-    if pooling_task == "classify":
+    if "classify" in supported_tasks:
        from .classify.io_processor import ClassifyIOProcessor

        processors["classify"] = ClassifyIOProcessor

-    if pooling_task == "token_classify":
+    if "token_classify" in supported_tasks:
        from .classify.io_processor import TokenClassifyIOProcessor

        processors["token_classify"] = TokenClassifyIOProcessor

-    if pooling_task == "embed":
+    if "embed" in supported_tasks:
        from .embed.io_processor import EmbedIOProcessor

        processors["embed"] = EmbedIOProcessor

-    if pooling_task == "token_embed":
+    if "token_embed" in supported_tasks:
        from .embed.io_processor import TokenEmbedIOProcessor

        processors["token_embed"] = TokenEmbedIOProcessor
@@ -72,15 +71,15 @@ def init_pooling_io_processors(
        from .pooling.io_processor import PluginWithIOProcessorPlugins

        processors["plugin"] = PluginWithIOProcessorPlugins
-    elif pooling_task == "plugin":
+    elif "plugin" in supported_tasks:
        from .pooling.io_processor import PluginWithoutIOProcessorPlugins

        processors["plugin"] = PluginWithoutIOProcessorPlugins

    if enable_scoring_api(supported_tasks, model_config):
+        score_type = model_config.score_type
        from .scoring.io_processor import ScoringIOProcessors

-        score_type: str | None = SCORE_TYPE_MAP.get(pooling_task, None)  # type: ignore[arg-type]
        if score_type is not None and score_type in ScoringIOProcessors:
            processors[score_type] = ScoringIOProcessors[score_type]

@@ -141,10 +140,6 @@ def init_pooling_state(
    request_logger: RequestLogger | None,
    supported_tasks: tuple["SupportedTask", ...],
 ):
-    model_config = engine_client.model_config
-    if model_config is None:
-        return
-
    from vllm.entrypoints.chat_utils import load_chat_template
    from vllm.tasks import POOLING_TASKS

@@ -153,14 +148,8 @@ def init_pooling_state(
    from .pooling.serving import ServingPooling
    from .scoring.serving import ServingScores

+    model_config = engine_client.model_config
    resolved_chat_template = load_chat_template(args.chat_template)
-    pooling_task = model_config.get_pooling_task(supported_tasks)
-
-    chat_template_config = ChatTemplateConfig(
-        chat_template=resolved_chat_template,
-        chat_template_content_format=args.chat_template_content_format,
-        trust_request_chat_template=args.trust_request_chat_template,
-    )

    state.serving_pooling = (
        (
@@ -169,7 +158,9 @@ def init_pooling_state(
                state.openai_serving_models,
                supported_tasks=supported_tasks,
                request_logger=request_logger,
-                chat_template_config=chat_template_config,
+                chat_template=resolved_chat_template,
+                chat_template_content_format=args.chat_template_content_format,
+                trust_request_chat_template=args.trust_request_chat_template,
            )
        )
        if any(t in supported_tasks for t in POOLING_TASKS)
@@ -180,9 +171,11 @@ def init_pooling_state(
            engine_client,
            state.openai_serving_models,
            request_logger=request_logger,
-            chat_template_config=chat_template_config,
+            chat_template=resolved_chat_template,
+            chat_template_content_format=args.chat_template_content_format,
+            trust_request_chat_template=args.trust_request_chat_template,
        )
-        if pooling_task == "embed"
+        if "embed" in supported_tasks
        else None
    )
    state.serving_classification = (
@@ -190,18 +183,21 @@ def init_pooling_state(
            engine_client,
            state.openai_serving_models,
            request_logger=request_logger,
-            chat_template_config=chat_template_config,
+            chat_template=resolved_chat_template,
+            chat_template_content_format=args.chat_template_content_format,
+            trust_request_chat_template=args.trust_request_chat_template,
        )
-        if pooling_task == "classify"
+        if "classify" in supported_tasks
        else None
    )
    state.serving_scores = (
        ServingScores(
            engine_client,
            state.openai_serving_models,
-            supported_tasks=supported_tasks,
            request_logger=request_logger,
-            chat_template_config=chat_template_config,
+            chat_template=resolved_chat_template,
+            chat_template_content_format=args.chat_template_content_format,
+            trust_request_chat_template=args.trust_request_chat_template,
            enable_flash_late_interaction=getattr(
                args, "enable_flash_late_interaction", True
            ),
@@ -218,12 +214,7 @@ def get_pooling_invocation_types(
    # NOTE: Items defined earlier take higher priority
    invocation_types: list[tuple[RequestType, tuple[GetHandlerFn, EndpointFn]]] = []

-    if model_config is None:
-        return invocation_types
-
-    pooling_task = model_config.get_pooling_task(supported_tasks)
-
-    if pooling_task == "embed":
+    if "embed" in supported_tasks:
        from .embed.api_router import create_embedding, embedding
        from .embed.protocol import EmbeddingRequest

@@ -231,7 +222,7 @@ def get_pooling_invocation_types(
            (EmbeddingRequest, (embedding, create_embedding)),
        ]

-    if pooling_task == "classify":
+    if "classify" in supported_tasks:
        from .classify.api_router import classify, create_classify
        from .classify.protocol import ClassificationRequest

@@ -78,15 +78,17 @@ class ServingPooling(PoolingServingBase):

        # plugin task uses io_processor.parse_request to verify inputs
        if pooling_task != "plugin" and pooling_task != self.pooling_task:
-            if pooling_task not in self.supported_tasks:
+            if pooling_task not in self.io_processors:
                raise ValueError(
                    f"Unsupported task: {pooling_task!r} "
                    f"Supported tasks: {self.supported_tasks}"
                )
            else:
-                raise ValueError(
-                    "Try switching the model's pooling_task "
-                    f"via --pooler-config.task {request.task}."
+                logger.warning_once(
+                    "Pooling multitask support is deprecated and will be removed "
+                    "in v0.20. When the default pooling task is not what you want, you "
+                    "need to manually specify it via --pooler-config.task %s. ",
+                    pooling_task,
                )

        if pooling_task == "plugin" and "plugin" not in self.io_processors:
@@ -8,7 +8,6 @@ from vllm.engine.protocol import EngineClient
 from vllm.entrypoints.openai.engine.protocol import UsageInfo
 from vllm.logger import init_logger
 from vllm.outputs import PoolingRequestOutput, ScoringRequestOutput
-from vllm.tasks import SCORE_TYPE_MAP, SupportedTask
 from vllm.v1.pool.late_interaction import (
    build_late_interaction_doc_params,
    build_late_interaction_query_params,
@@ -39,15 +38,10 @@ class ServingScores(PoolingServing):
        self,
        engine_client: EngineClient,
        *args,
-        supported_tasks: tuple[SupportedTask, ...],
        enable_flash_late_interaction: bool = True,
        **kwargs,
    ):
-        pooling_task = engine_client.model_config.get_pooling_task(supported_tasks)
-        score_type = SCORE_TYPE_MAP.get(pooling_task, None)  # type: ignore[arg-type]
-        assert score_type is not None
-
-        self.io_processor_name: str = score_type
+        self.io_processor_name: str = engine_client.model_config.score_type
        self.enable_flash_late_interaction = (
            self.io_processor_name == "late-interaction"
            and enable_flash_late_interaction
@@ -141,14 +141,10 @@ def enable_scoring_api(
    supported_tasks: tuple["SupportedTask", ...],
    model_config: ModelConfig | None = None,
 ) -> bool:
-    if model_config is None:
-        return False
-
-    pooling_task = model_config.get_pooling_task(supported_tasks)
-    if pooling_task in ("embed", "token_embed"):
+    if any(t in supported_tasks for t in ("embed", "token_embed")):
        return True

-    if pooling_task == "classify":
+    if model_config is not None and "classify" in supported_tasks:
        num_labels = getattr(model_config.hf_config, "num_labels", 0)
        if num_labels != 1:
            logger.debug_once("Scoring API is only enabled for num_labels == 1.")
@@ -151,6 +151,46 @@ class SiluAndMul(CustomOp):
        return self.forward_cuda(x)


+@CustomOp.register("silu_and_mul_with_clamp")
+class SiluAndMulWithClamp(CustomOp):
+    """SwiGLU activation with input clamping (used by some MoE shared experts).
+
+    Computes:
+        gate = clamp(x[..., :d], max=swiglu_limit)
+        up   = clamp(x[..., d:], min=-swiglu_limit, max=swiglu_limit)
+        out  = silu(gate) * up
+    where d = x.shape[-1] // 2.
+
+    Shapes:
+        x: (num_tokens, 2 * d) or (batch_size, seq_len, 2 * d)
+        return: (num_tokens, d) or (batch_size, seq_len, d)
+    """
+
+    def __init__(self, swiglu_limit: float, *, compile_native: bool = True):
+        super().__init__(compile_native=compile_native)
+        self.swiglu_limit = float(swiglu_limit)
+        if current_platform.is_cuda_alike() or current_platform.is_xpu():
+            self.op = torch.ops._C.silu_and_mul_with_clamp
+        elif current_platform.is_cpu():
+            self._forward_method = self.forward_native
+
+    def forward_native(self, x: torch.Tensor) -> torch.Tensor:
+        d = x.shape[-1] // 2
+        gate = torch.clamp(x[..., :d], max=self.swiglu_limit)
+        up = torch.clamp(x[..., d:], min=-self.swiglu_limit, max=self.swiglu_limit)
+        return F.silu(gate) * up
+
+    def forward_cuda(self, x: torch.Tensor) -> torch.Tensor:
+        d = x.shape[-1] // 2
+        output_shape = x.shape[:-1] + (d,)
+        out = torch.empty(output_shape, dtype=x.dtype, device=x.device)
+        self.op(out, x, self.swiglu_limit)
+        return out
+
+    def forward_xpu(self, x: torch.Tensor) -> torch.Tensor:
+        return self.forward_cuda(x)
+
+
 # --8<-- [start:mul_and_silu]
@CustomOp.register("mul_and_silu")
 class MulAndSilu(CustomOp):
@@ -45,7 +45,7 @@ def _gelu_and_mul(
 # Uses static methods or standalone functions to avoid instantiating CustomOp
 # classes, which would call get_current_vllm_config() before config is set.
 _CPU_MOE_ACT_FN: dict[MoEActivation, Callable[[torch.Tensor], torch.Tensor]] = {
-    MoEActivation.SILU: SiluAndMul.forward_native,
+    MoEActivation.SILU: lambda x: SiluAndMul(compile_native=False).forward_native(x),
    MoEActivation.SWIGLUOAI: _swigluoai_forward_native,
    MoEActivation.GELU: _gelu_and_mul,
 }
@@ -17,6 +17,7 @@ from vllm.distributed import (
    get_tensor_model_parallel_world_size,
 )
 from vllm.forward_context import get_forward_context
+from vllm.model_executor.layers.activation import SiluAndMul, SiluAndMulWithClamp
 from vllm.model_executor.layers.deepseek_v4_attention import (
    DeepseekV4Indexer,
    DeepseekV4MLAModules,
@@ -34,7 +35,10 @@ from vllm.model_executor.layers.linear import (
    RowParallelLinear,
 )
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
-from vllm.model_executor.layers.quantization import QuantizationMethods
+from vllm.model_executor.layers.quantization import (
+    QuantizationConfig,
+    QuantizationMethods,
+)
 from vllm.model_executor.layers.quantization.fp8 import Fp8Config
 from vllm.model_executor.layers.quantization.mxfp4 import Mxfp4MoEMethod
 from vllm.model_executor.layers.quantization.utils.quant_utils import (
@@ -46,7 +50,6 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
    VocabParallelEmbedding,
 )
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
-from vllm.model_executor.models.deepseek_v2 import DeepseekV2MLP
 from vllm.model_executor.utils import set_weight_attrs
 from vllm.platforms import current_platform
 from vllm.sequence import IntermediateTensors
@@ -63,6 +66,57 @@ from .utils import (
 )


+class DeepseekV4MLP(nn.Module):
+    def __init__(
+        self,
+        hidden_size: int,
+        intermediate_size: int,
+        hidden_act: str,
+        swiglu_limit: float | None = None,
+        quant_config: QuantizationConfig | None = None,
+        reduce_results: bool = True,
+        is_sequence_parallel: bool = False,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+
+        # If is_sequence_parallel, the input and output tensors are sharded
+        # across the ranks within the tp_group. In this case the weights are
+        # replicated and no collective ops are needed.
+        # Otherwise we use standard TP with an allreduce at the end.
+        self.gate_up_proj = MergedColumnParallelLinear(
+            hidden_size,
+            [intermediate_size] * 2,
+            bias=False,
+            quant_config=quant_config,
+            disable_tp=is_sequence_parallel,
+            prefix=f"{prefix}.gate_up_proj",
+        )
+        self.down_proj = RowParallelLinear(
+            intermediate_size,
+            hidden_size,
+            bias=False,
+            quant_config=quant_config,
+            reduce_results=reduce_results,
+            disable_tp=is_sequence_parallel,
+            prefix=f"{prefix}.down_proj",
+        )
+        if hidden_act != "silu":
+            raise ValueError(
+                f"Unsupported activation: {hidden_act}. Only silu is supported for now."
+            )
+        if swiglu_limit is not None:
+            self.act_fn = SiluAndMulWithClamp(swiglu_limit)
+        else:
+            self.act_fn = SiluAndMul()
+
+    def forward(self, x):
+        gate_up, _ = self.gate_up_proj(x)
+        x = self.act_fn(gate_up)
+        x, _ = self.down_proj(x)
+        return x
+
+
 class DeepseekV4FP8Config(Fp8Config):
    """FP8 config that routes MoE layers to MXFP4 quantization.

@@ -672,10 +726,11 @@ class DeepseekV4MoE(nn.Module):
        else:
            intermediate_size = config.moe_intermediate_size * config.n_shared_experts

-            self.shared_experts = DeepseekV2MLP(
+            self.shared_experts = DeepseekV4MLP(
                hidden_size=config.hidden_size,
                intermediate_size=intermediate_size,
                hidden_act=config.hidden_act,
+                swiglu_limit=self.swiglu_limit,
                quant_config=quant_config,
                reduce_results=self.use_mega_moe,
                prefix=f"{prefix}.shared_experts",
@@ -87,6 +87,13 @@ class PoolingParams(
        return deepcopy(self)

    def verify(self, model_config: ModelConfig) -> None:
+        if self.task == "score":
+            logger.warning_once(
+                "`score` task is deprecated and will be removed in v0.20. "
+                "Please use `classify` instead."
+            )
+            self.task = "classify"
+
        # plugin task uses io_processor.parse_request to verify inputs,
        # skipping PoolingParams verify
        if self.task == "plugin":
@@ -16,11 +16,6 @@ PoolingTask = Literal[
 POOLING_TASKS: tuple[PoolingTask, ...] = get_args(PoolingTask)

 ScoreType = Literal["bi-encoder", "cross-encoder", "late-interaction"]
-SCORE_TYPE_MAP: dict[PoolingTask, ScoreType] = {
-    "embed": "bi-encoder",
-    "classify": "cross-encoder",
-    "token_embed": "late-interaction",
-}

 FrontendTask = Literal["render"]
 FRONTEND_TASKS: tuple[FrontendTask, ...] = get_args(FrontendTask)
@@ -34,6 +34,7 @@ class KVCacheCoordinator(ABC):
        self,
        kv_cache_config: KVCacheConfig,
        max_model_len: int,
+        max_num_batched_tokens: int,
        use_eagle: bool,
        enable_caching: bool,
        enable_kv_cache_events: bool,
@@ -65,6 +66,8 @@ class KVCacheCoordinator(ABC):
        self.single_type_managers = tuple(
            get_manager_for_kv_cache_spec(
                kv_cache_spec=kv_cache_group.kv_cache_spec,
+                max_num_batched_tokens=max_num_batched_tokens,
+                max_model_len=max_model_len,
                block_pool=self.block_pool,
                enable_caching=enable_caching,
                kv_cache_group_id=i,
@@ -271,6 +274,7 @@ class KVCacheCoordinatorNoPrefixCache(KVCacheCoordinator):
        self,
        kv_cache_config: KVCacheConfig,
        max_model_len: int,
+        max_num_batched_tokens: int,
        use_eagle: bool,
        enable_kv_cache_events: bool,
        dcp_world_size: int,
@@ -281,6 +285,7 @@ class KVCacheCoordinatorNoPrefixCache(KVCacheCoordinator):
        super().__init__(
            kv_cache_config,
            max_model_len,
+            max_num_batched_tokens,
            use_eagle,
            False,
            enable_kv_cache_events,
@@ -316,6 +321,7 @@ class UnitaryKVCacheCoordinator(KVCacheCoordinator):
        self,
        kv_cache_config: KVCacheConfig,
        max_model_len: int,
+        max_num_batched_tokens: int,
        use_eagle: bool,
        enable_caching: bool,
        enable_kv_cache_events: bool,
@@ -327,6 +333,7 @@ class UnitaryKVCacheCoordinator(KVCacheCoordinator):
        super().__init__(
            kv_cache_config,
            max_model_len,
+            max_num_batched_tokens,
            use_eagle,
            enable_caching,
            enable_kv_cache_events,
@@ -381,6 +388,7 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):
        self,
        kv_cache_config: KVCacheConfig,
        max_model_len: int,
+        max_num_batched_tokens: int,
        use_eagle: bool,
        enable_caching: bool,
        enable_kv_cache_events: bool,
@@ -392,6 +400,7 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):
        super().__init__(
            kv_cache_config,
            max_model_len,
+            max_num_batched_tokens,
            use_eagle,
            enable_caching,
            enable_kv_cache_events,
@@ -574,6 +583,7 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):
 def get_kv_cache_coordinator(
    kv_cache_config: KVCacheConfig,
    max_model_len: int,
+    max_num_batched_tokens: int,
    use_eagle: bool,
    enable_caching: bool,
    enable_kv_cache_events: bool,
@@ -586,6 +596,7 @@ def get_kv_cache_coordinator(
        return KVCacheCoordinatorNoPrefixCache(
            kv_cache_config,
            max_model_len,
+            max_num_batched_tokens,
            use_eagle,
            enable_kv_cache_events,
            dcp_world_size=dcp_world_size,
@@ -597,6 +608,7 @@ def get_kv_cache_coordinator(
        return UnitaryKVCacheCoordinator(
            kv_cache_config,
            max_model_len,
+            max_num_batched_tokens,
            use_eagle,
            enable_caching,
            enable_kv_cache_events,
@@ -608,6 +620,7 @@ def get_kv_cache_coordinator(
    return HybridKVCacheCoordinator(
        kv_cache_config,
        max_model_len,
+        max_num_batched_tokens,
        use_eagle,
        enable_caching,
        enable_kv_cache_events,
@@ -109,6 +109,7 @@ class KVCacheManager:
        kv_cache_config: KVCacheConfig,
        max_model_len: int,
        hash_block_size: int,
+        max_num_batched_tokens: int | None = None,
        enable_caching: bool = True,
        use_eagle: bool = False,
        log_stats: bool = False,
@@ -118,6 +119,11 @@ class KVCacheManager:
        metrics_collector: KVCacheMetricsCollector | None = None,
    ) -> None:
        self.max_model_len = max_model_len
+        # When unset, fall back to `max_model_len` so the recycling-aware cap
+        # collapses to the prior (uncapped) admission behavior. The scheduler
+        # always supplies the real value at runtime.
+        if max_num_batched_tokens is None:
+            max_num_batched_tokens = max_model_len

        self.enable_caching = enable_caching
        self.use_eagle = use_eagle
@@ -131,6 +137,7 @@ class KVCacheManager:
        self.coordinator = get_kv_cache_coordinator(
            kv_cache_config=kv_cache_config,
            max_model_len=self.max_model_len,
+            max_num_batched_tokens=max_num_batched_tokens,
            use_eagle=self.use_eagle,
            enable_caching=self.enable_caching,
            enable_kv_cache_events=enable_kv_cache_events,
@@ -228,6 +228,7 @@ class Scheduler(SchedulerInterface):
        self.kv_cache_manager = KVCacheManager(
            kv_cache_config=kv_cache_config,
            max_model_len=self.max_model_len,
+            max_num_batched_tokens=self.scheduler_config.max_num_batched_tokens,
            enable_caching=self.cache_config.enable_prefix_caching,
            use_eagle=self.use_eagle,
            log_stats=self.log_stats,
@@ -41,6 +41,7 @@ class SingleTypeKVCacheManager(ABC):
        kv_cache_group_id: int,
        dcp_world_size: int = 1,
        pcp_world_size: int = 1,
+        max_admission_blocks_per_request: int | None = None,
    ) -> None:
        """
        Initializes the SingleTypeKVCacheManager.
@@ -48,6 +49,12 @@ class SingleTypeKVCacheManager(ABC):
            kv_cache_spec: The kv_cache_spec for this manager.
            block_pool: The block pool.
            kv_cache_group_id: The id of the kv cache group of this manager.
+            max_admission_blocks_per_request: Recycling-aware per-request
+                block cap used by `get_num_blocks_to_allocate`. Only set for
+                spec types that recycle blocks across chunks (SWA,
+                chunked-local); `None` (the default) means no cap, which is
+                correct for full-attention-style specs that hold every
+                block until the request finishes.
        """
        self.block_size = kv_cache_spec.block_size
        self.dcp_world_size = dcp_world_size
@@ -57,6 +64,7 @@ class SingleTypeKVCacheManager(ABC):
        self.kv_cache_spec = kv_cache_spec
        self.block_pool = block_pool
        self.enable_caching = enable_caching
+        self._max_admission_blocks_per_request = max_admission_blocks_per_request
        self.new_block_ids: list[int] = []

        # Mapping from request ID to blocks to track the blocks allocated
@@ -105,6 +113,19 @@ class SingleTypeKVCacheManager(ABC):
        """

        num_required_blocks = cdiv(num_tokens, self.block_size)
+        if self._max_admission_blocks_per_request is not None:
+            # Recycling-aware specs (SWA, chunked-local) cap the per-request
+            # reservation here so admission matches the startup pool sizer
+            # (`SlidingWindowSpec.max_admission_blocks_per_request` / its
+            # chunked-local counterpart). `remove_skipped_blocks` runs from
+            # `allocate_slots` before each chunk's `get_num_blocks_to_allocate`,
+            # so per-request peak real-held blocks <= this cap, which keeps
+            # `sum(reservations) <= pool` <=> `sum(peak_real_held) <= pool`.
+            # Drift between the two would re-introduce the deadlock from
+            # issue #39734 or, worse, mid-prefill OOM.
+            num_required_blocks = min(
+                num_required_blocks, self._max_admission_blocks_per_request
+            )
        num_req_blocks = len(self.req_to_blocks.get(request_id, ()))

        if request_id in self.num_cached_block:
@@ -1126,8 +1147,21 @@ spec_manager_map: dict[type[KVCacheSpec], type[SingleTypeKVCacheManager]] = {


 def get_manager_for_kv_cache_spec(
-    kv_cache_spec: KVCacheSpec, **kwargs
+    kv_cache_spec: KVCacheSpec,
+    max_num_batched_tokens: int,
+    max_model_len: int,
+    **kwargs,
 ) -> SingleTypeKVCacheManager:
    manager_class = spec_manager_map[type(kv_cache_spec)]
+    # SlidingWindow / ChunkedLocalAttention managers recycle blocks across
+    # chunks; the runtime admission cap must match the recycling-aware bound
+    # the startup pool sizer uses (single source of truth: the spec method).
+    if isinstance(kv_cache_spec, (SlidingWindowSpec, ChunkedLocalAttentionSpec)):
+        kwargs["max_admission_blocks_per_request"] = (
+            kv_cache_spec.max_admission_blocks_per_request(
+                max_num_batched_tokens=max_num_batched_tokens,
+                max_model_len=max_model_len,
+            )
+        )
    manager = manager_class(kv_cache_spec, **kwargs)
    return manager
@@ -376,19 +376,28 @@ class MLAAttentionSpec(FullAttentionSpec):
 class ChunkedLocalAttentionSpec(AttentionSpec):
    attention_chunk_size: int

-    def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:
-        max_model_len = vllm_config.model_config.max_model_len
-        max_num_batched_tokens = vllm_config.scheduler_config.max_num_batched_tokens
+    def max_admission_blocks_per_request(
+        self, max_num_batched_tokens: int, max_model_len: int
+    ) -> int:
+        """Per-request admission cap, in blocks.

-        # During chunked prefill, we allocate KV cache for at most
-        # `self.attention_chunk_size` computed tokens plus the newly scheduled
-        # tokens. And we won't allocate KV cache for more than `max_model_len`
-        # tokens.
+        Single source of truth for both startup pool sizing
+        (`max_memory_usage_bytes`) and the runtime admission gate, so requests
+        admitted by startup can also be admitted at runtime.
+        """
+        # During chunked prefill, we hold KV for at most one chunk window.
        num_tokens = min(
            self.attention_chunk_size + max_num_batched_tokens, max_model_len
        )
+        return cdiv(num_tokens, self.block_size)

-        return cdiv(num_tokens, self.block_size) * self.page_size_bytes
+    def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:
+        max_model_len = vllm_config.model_config.max_model_len
+        max_num_batched_tokens = vllm_config.scheduler_config.max_num_batched_tokens
+        max_blocks = self.max_admission_blocks_per_request(
+            max_num_batched_tokens=max_num_batched_tokens, max_model_len=max_model_len
+        )
+        return max_blocks * self.page_size_bytes


@dataclass(frozen=True, kw_only=True)
@@ -409,26 +418,38 @@ class SlidingWindowSpec(AttentionSpec):
            * get_dtype_size(self.dtype)
        )

+    def max_admission_blocks_per_request(
+        self, max_num_batched_tokens: int, max_model_len: int
+    ) -> int:
+        """Per-request admission cap, in blocks.
+
+        Single source of truth for both startup pool sizing
+        (`max_memory_usage_bytes`) and the runtime admission gate. Per-request
+        real-held blocks plateau at this bound because
+        `SlidingWindowManager.remove_skipped_blocks` runs from `allocate_slots`
+        before each chunk's `get_num_blocks_to_allocate`.
+        """
+        # During chunked prefill, we hold KV for the last `sliding_window-1`
+        # computed tokens plus the newly scheduled tokens, and never more
+        # than `max_model_len`.
+        num_tokens = min(
+            self.sliding_window - 1 + max_num_batched_tokens, max_model_len
+        )
+        # +1 because the sliding window may not start from the beginning of
+        # the block. E.g. block size 4 and num_token 4 needs two blocks
+        # [XXCD][EF] to store the 6-token window [CDEF].
+        return cdiv(num_tokens, self.block_size) + 1
+
    def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:
        assert vllm_config.parallel_config.decode_context_parallel_size == 1, (
            "DCP not support sliding window."
        )
        max_model_len = vllm_config.model_config.max_model_len
        max_num_batched_tokens = vllm_config.scheduler_config.max_num_batched_tokens
-
-        # During chunked prefill, we allocate KV cache for the last
-        # `self.sliding_window-1` computed tokens plus the newly scheduled
-        # tokens. And we won't allocate KV cache for more than `max_model_len`
-        # tokens.
-        num_tokens = min(
-            self.sliding_window - 1 + max_num_batched_tokens, max_model_len
+        max_blocks = self.max_admission_blocks_per_request(
+            max_num_batched_tokens=max_num_batched_tokens, max_model_len=max_model_len
        )
-
-        # +1 here because the sliding window may not start from the beginning
-        # of the block. For example, if the block size is 4 and num_token
-        # is 4, we need two blocks [XXCD] [EF] to store the sliding
-        # window [CDEF] of 6 tokens.
-        return (cdiv(num_tokens, self.block_size) + 1) * self.page_size_bytes
+        return max_blocks * self.page_size_bytes


@dataclass(frozen=True, kw_only=True)
@@ -110,6 +110,9 @@ class SimpleCPUOffloadScheduler:
        self.cpu_coordinator: KVCacheCoordinator = get_kv_cache_coordinator(
            kv_cache_config=self.cpu_kv_cache_config,
            max_model_len=vllm_config.model_config.max_model_len,
+            max_num_batched_tokens=(
+                vllm_config.scheduler_config.max_num_batched_tokens
+            ),
            use_eagle=False,
            enable_caching=True,
            enable_kv_cache_events=self.enable_kv_cache_events,
Author	SHA1	Message	Date
Zhewen Liandkhluu	88d34c6409	[Docker] Install numactl CLI in CUDA runtime image (#41032 ) Signed-off-by: Zhewen Li <zhewenli@inferact.ai> Co-authored-by: Zhewen Li <zhewenli@inferact.ai> Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com> (cherry picked from commit `6232fb4b66`)	2026-04-27 12:04:32 -07:00
Yongye Zhuandkhluu	b8160878f0	[DSV4] Add silu clamp limit to shared expert (#40950 ) Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> (cherry picked from commit `706a04d34b`)	2026-04-27 00:40:46 -07:00
Dao007foreverandkhluu	84c276d7ea	[Bugfix] Cap SWA/chunked-local runtime admission to startup pool-sizing bound (#40946 ) Signed-off-by: Dao Le <Dao007forever@gmail.com> Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: Nick Hill <nickhill123@gmail.com> (cherry picked from commit `7b1bc0a3eb`)	2026-04-26 21:43:00 -07:00
khluu	5eb3657578	Revert "[Frontend] Remove frontend pooling multi task support. (#37861 )" This reverts commit `d2e2e856ad`.	2026-04-26 21:42:15 -07:00