p

Signed-off-by: khluu <khluu000@gmail.com>

.buildkite/hardware_tests/amd.yaml

@@ -1,6 +1,7 @@
-group: Hardware
+group: Hardware - AMD Build 
 steps:
   - label: "AMD: :docker: build image"
+    key: image-build-amd
     depends_on: []
     device: amd_cpu
     no_plugin: true

.buildkite/test-amd.yaml

@@ -132,7 +132,7 @@ steps:
   - tests/entrypoints/
   commands:
   - pytest -v -s entrypoints/openai/tool_parsers
-  - pytest -v -s entrypoints/ --ignore=entrypoints/llm --ignore=entrypoints/openai --ignore=entrypoints/rpc --ignore=entrypoints/sleep --ignore=entrypoints/instrumentator --ignore=entrypoints/offline_mode --ignore=entrypoints/test_chat_utils.py  --ignore=entrypoints/pooling
+  - pytest -v -s entrypoints/ --ignore=entrypoints/llm --ignore=entrypoints/openai --ignore=entrypoints/rpc --ignore=entrypoints/instrumentator --ignore=entrypoints/offline_mode --ignore=entrypoints/test_chat_utils.py  --ignore=entrypoints/pooling
 
 - label: Entrypoints Integration Test (LLM) # 30min
   timeout_in_minutes: 40
@@ -179,14 +179,14 @@ steps:
   torch_nightly: true
   source_file_dependencies:
   - vllm/
-  - tests/entrypoints/sleep
   - tests/entrypoints/rpc
+  - tests/entrypoints/instrumentator
   - tests/tool_use
   commands:
   - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -v -s entrypoints/sleep
+  - pytest -v -s entrypoints/instrumentator
+  - PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/rpc
   - pytest -v -s tool_use
-  - PYTHONPATH=/vllm-workspace  pytest -v -s entrypoints/rpc
 
 - label: Entrypoints Integration Test (Pooling)
   timeout_in_minutes: 50
@@ -552,7 +552,7 @@ steps:
 - label: LoRA Test %N # 20min each
   timeout_in_minutes: 30
   mirror_hardwares: [amdexperimental]
-  agent_pool: mi325_8
+  agent_pool: mi325_1
   # grade: Blocking
   source_file_dependencies:
   - vllm/lora
@@ -648,7 +648,7 @@ steps:
 - label: Kernels Attention Test %N # 23min
   timeout_in_minutes: 35
   mirror_hardwares: [amdexperimental, amdproduction]
-  agent_pool: mi325_8
+  agent_pool: mi325_1
   # grade: Blocking
   source_file_dependencies:
   - csrc/attention/
@@ -663,7 +663,7 @@ steps:
 - label: Kernels Quantization Test %N # 64min
   timeout_in_minutes: 90
   mirror_hardwares: [amdexperimental]
-  agent_pool: mi325_8
+  agent_pool: mi325_1
   # grade: Blocking
   source_file_dependencies:
   - csrc/quantization/
@@ -676,7 +676,7 @@ steps:
 - label: Kernels MoE Test %N # 40min
   timeout_in_minutes: 60
   mirror_hardwares: [amdexperimental, amdproduction]
-  agent_pool: mi325_8
+  agent_pool: mi325_1
   # grade: Blocking
   source_file_dependencies:
   - csrc/quantization/cutlass_w8a8/moe/
@@ -839,7 +839,7 @@ steps:
 - label: Basic Models Tests (Extra Initialization) %N
   timeout_in_minutes: 45
   mirror_hardwares: [amdexperimental, amdproduction]
-  agent_pool: mi325_8
+  agent_pool: mi325_1
   # grade: Blocking
   torch_nightly: true
   source_file_dependencies:
@@ -901,7 +901,7 @@ steps:
 - label: Language Models Tests (Extra Standard) %N
   timeout_in_minutes: 45
   mirror_hardwares: [amdexperimental]
-  agent_pool: mi325_8
+  agent_pool: mi325_1
   # grade: Blocking
   torch_nightly: true
   source_file_dependencies:
@@ -922,7 +922,7 @@ steps:
 - label: Language Models Tests (Hybrid) %N
   timeout_in_minutes: 75
   mirror_hardwares: [amdexperimental]
-  agent_pool: mi325_8
+  agent_pool: mi325_1
   # grade: Blocking
   torch_nightly: true
   source_file_dependencies:
@@ -1334,7 +1334,7 @@ steps:
   - pytest -v -s ./compile/test_wrapper.py
   - VLLM_TEST_SAME_HOST=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed'
   - VLLM_TEST_SAME_HOST=1 VLLM_TEST_WITH_DEFAULT_DEVICE_SET=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed'
-  - pytest -v -s distributed/test_sequence_parallel.py
+  - pytest -v -s compile/correctness_e2e/test_sequence_parallel.py
   - CUDA_VISIBLE_DEVICES=0,1 pytest -v -s v1/shutdown
   - pytest -v -s v1/worker/test_worker_memory_snapshot.py
 

.buildkite/test-pipeline.yaml

@@ -118,7 +118,7 @@ steps:
   - tests/entrypoints/
   commands:
   - pytest -v -s entrypoints/openai/tool_parsers
-  - pytest -v -s entrypoints/ --ignore=entrypoints/llm --ignore=entrypoints/rpc --ignore=entrypoints/sleep --ignore=entrypoints/instrumentator --ignore=entrypoints/openai --ignore=entrypoints/offline_mode --ignore=entrypoints/test_chat_utils.py  --ignore=entrypoints/pooling
+  - pytest -v -s entrypoints/ --ignore=entrypoints/llm --ignore=entrypoints/rpc --ignore=entrypoints/instrumentator --ignore=entrypoints/openai --ignore=entrypoints/offline_mode --ignore=entrypoints/test_chat_utils.py  --ignore=entrypoints/pooling
 
 - label: Entrypoints Integration Test (LLM) # 30min
   timeout_in_minutes: 40
@@ -148,7 +148,7 @@ steps:
   - tests/entrypoints/test_chat_utils
   commands:
   - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -v -s entrypoints/openai --ignore=entrypoints/openai/test_chat_with_tool_reasoning.py --ignore=entrypoints/openai/test_oot_registration.py --ignore=entrypoints/openai/test_tensorizer_entrypoint.py --ignore=entrypoints/openai/correctness/  --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses
+  - pytest -v -s entrypoints/openai --ignore=entrypoints/openai/test_chat_with_tool_reasoning.py --ignore=entrypoints/instrumentator --ignore=entrypoints/openai/test_oot_registration.py --ignore=entrypoints/openai/test_tensorizer_entrypoint.py --ignore=entrypoints/openai/correctness/  --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses
   - pytest -v -s entrypoints/test_chat_utils.py
 
 - label: Entrypoints Integration Test (API Server 2)
@@ -159,13 +159,13 @@ steps:
   torch_nightly: true
   source_file_dependencies:
   - vllm/
-  - tests/entrypoints/sleep
   - tests/entrypoints/rpc
+  - tests/entrypoints/instrumentator
   - tests/tool_use
   commands:
   - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -v -s entrypoints/sleep
-  - PYTHONPATH=/vllm-workspace  pytest -v -s entrypoints/rpc
+  - pytest -v -s entrypoints/instrumentator
+  - PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/rpc
   - pytest -v -s tool_use
 
 - label: Entrypoints Integration Test (Pooling)
@@ -862,7 +862,7 @@ steps:
   commands:
     # Install fast path packages for testing against transformers
     # Note: also needed to run plamo2 model in vLLM
-    - uv pip install --system --no-build-isolation 'git+https://github.com/state-spaces/mamba@v2.2.5'
+    - uv pip install --system --no-build-isolation 'git+https://github.com/state-spaces/mamba@v2.3.0'
     - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.5.2'
     # Shard hybrid language model tests
     - pytest -v -s models/language/generation \
@@ -881,7 +881,7 @@ steps:
   commands:
     # Install fast path packages for testing against transformers
     # Note: also needed to run plamo2 model in vLLM
-    - uv pip install --system --no-build-isolation 'git+https://github.com/state-spaces/mamba@v2.2.5'
+    - uv pip install --system --no-build-isolation 'git+https://github.com/state-spaces/mamba@v2.3.0'
     - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.5.2'
     - pytest -v -s models/language/generation -m '(not core_model) and (not hybrid_model)'
 

.buildkite/test_areas/basic_correctness.yaml

@@ -14,3 +14,8 @@ steps:
   - pytest -v -s basic_correctness/test_cumem.py
   - pytest -v -s basic_correctness/test_basic_correctness.py
   - pytest -v -s basic_correctness/test_cpu_offload.py
+  mirror:
+    amd:
+      device: mi325_1
+      depends_on:
+      - image-build-amd

.buildkite/test_areas/benchmarks.yaml

@@ -17,3 +17,15 @@ steps:
   - tests/benchmarks/
   commands:
   - pytest -v -s benchmarks/
+
+- label: Attention Benchmarks Smoke Test (B200)
+  device: b200
+  num_gpus: 2
+  optional: true
+  working_dir: "/vllm-workspace/"
+  timeout_in_minutes: 10
+  source_file_dependencies:
+  - benchmarks/attention_benchmarks/
+  - vllm/v1/attention/
+  commands:
+  - python3 benchmarks/attention_benchmarks/benchmark.py --backends flash flashinfer --batch-specs "8q1s1k" --repeats 1 --warmup-iters 1

.buildkite/test_areas/compile.yaml

@@ -121,13 +121,10 @@ steps:
   optional: true
   commands:
     - nvidia-smi
-    # Run all models and attn backends but only Inductor partition and native custom ops
-    # -k "inductor_partition and not +rms_norm and not +quant_fp8"
+    # Run all models but only FLASHINFER, Inductor partition and native custom ops
     # Qwen requires +quant_fp8 as -quant_fp8 rms+quant fusion is not supported
-    # -k "inductor_partition and not +rms_norm and +quant_fp8 and qwen3"
-    # Run just llama3 (fp8 & fp4) for all config combinations
-    # -k "llama-3"
-    - pytest -v -s tests/compile/fusions_e2e/test_tp1_quant.py -k "inductor_partition and not +rms_norm and not +quant_fp8" -k "inductor_partition and not +rms_norm and +quant_fp8 and qwen3" -k "llama-3"
+    # Run just llama3 (fp8 & fp4) for all config combinations (only inductor partition)
+    - pytest -v -s tests/compile/fusions_e2e/test_tp1_quant.py -k "inductor_partition and (FLASHINFER and not +rms_norm and (not +quant_fp8 or +quant_fp8 and qwen3) or llama-3)"
 
 - label: Fusion E2E TP2 Quick (H100)
   timeout_in_minutes: 20
@@ -162,7 +159,7 @@ steps:
     - tests/compile/fusions_e2e/
   commands:
     - nvidia-smi
-    # Run just llama3 (fp4 & fp8 & bf16) for all config combinations
+    # Run just llama3 (fp8 & bf16) for all config combinations
     - pytest -v -s tests/compile/fusions_e2e/test_tp2_ar_rms.py -k "llama-3"
 
 - label: Fusion E2E TP2 AsyncTP Config Sweep (H100)
@@ -197,7 +194,8 @@ steps:
     - tests/compile/fusions_e2e/
   commands:
     - nvidia-smi
-    # Run all models and attn backends but only Inductor partition and native custom ops
+    # Run all models but only FLASHINFER, Inductor partition and native custom ops
+    # include qwen with +quant_fp8 as -quant_fp8 rms+quant fusion is not supported
     # for ar-rms-quant-fp4, also sweep llama3
-    - pytest -v -s tests/compile/fusions_e2e/test_tp2_ar_rms.py -k "inductor_partition and not +rms_norm and not +quant_fp8" -k "Llama-3.1-8B-Instruct-FP4"
-    - pytest -v -s tests/compile/fusions_e2e/test_tp2_async_tp.py -k "inductor_partition and not +rms_norm and not +quant_fp8"
+    - pytest -v -s tests/compile/fusions_e2e/test_tp2_ar_rms.py -k "(FLASHINFER and inductor_partition and not +rms_norm and (not +quant_fp8 or +quant_fp8 and qwen3)) or Llama-3.1-8B-Instruct-FP4"
+    - pytest -v -s tests/compile/fusions_e2e/test_tp2_async_tp.py -k "FLASHINFER and inductor_partition and not +rms_norm and (not +quant_fp8 or +quant_fp8 and qwen3)"

.buildkite/test_areas/entrypoints.yaml

@@ -24,6 +24,11 @@ steps:
   - pytest -v -s entrypoints/llm --ignore=entrypoints/llm/test_generate.py --ignore=entrypoints/llm/test_collective_rpc.py
   - pytest -v -s entrypoints/llm/test_generate.py # it needs a clean process
   - pytest -v -s entrypoints/offline_mode # Needs to avoid interference with other tests
+  mirror:
+    amd:
+      device: mi325_1
+      depends_on:
+      - image-build-amd
 
 - label: Entrypoints Integration (API Server 1)
   timeout_in_minutes: 130
@@ -42,15 +47,13 @@ steps:
   working_dir: "/vllm-workspace/tests"
   source_file_dependencies:
   - vllm/
-  - tests/tool_use
-  - tests/entrypoints/sleep
-  - tests/entrypoints/instrumentator
   - tests/entrypoints/rpc
+  - tests/entrypoints/instrumentator
+  - tests/tool_use
   commands:
   - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/rpc
   - pytest -v -s entrypoints/instrumentator
-  - pytest -v -s entrypoints/sleep
+  - PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/rpc
   - pytest -v -s tool_use
 
 - label: Entrypoints Integration (Pooling)

.buildkite/test_areas/models_basic.yaml

@@ -4,7 +4,6 @@ depends_on:
 steps:
 - label: Basic Models Tests (Initialization)
   timeout_in_minutes: 45
-  mirror_hardwares: [amdexperimental]
   torch_nightly: true
   source_file_dependencies:
   - vllm/
@@ -16,7 +15,6 @@ steps:
 
 - label: Basic Models Tests (Extra Initialization) %N
   timeout_in_minutes: 45
-  mirror_hardwares: [amdexperimental]
   torch_nightly: true
   source_file_dependencies:
   - vllm/model_executor/models/
@@ -38,6 +36,12 @@ steps:
   - tests/models/test_registry.py
   commands:
     - pytest -v -s models/test_terratorch.py models/test_transformers.py models/test_registry.py
+  mirror:
+    amd:
+      device: mi325_1
+      depends_on:
+      - image-build-amd
+    
 
 - label: Basic Models Test (Other CPU) # 5min
   depends_on: 

.buildkite/test_areas/models_language.yaml

@@ -4,7 +4,6 @@ depends_on:
 steps:
 - label: Language Models Tests (Standard)
   timeout_in_minutes: 25
-  mirror_hardwares: [amdexperimental]
   torch_nightly: true
   source_file_dependencies:
   - vllm/
@@ -16,7 +15,6 @@ steps:
 
 - label: Language Models Tests (Extra Standard) %N
   timeout_in_minutes: 45
-  mirror_hardwares: [amdexperimental]
   torch_nightly: true
   source_file_dependencies:
   - vllm/model_executor/models/
@@ -32,7 +30,6 @@ steps:
 
 - label: Language Models Tests (Hybrid) %N
   timeout_in_minutes: 75
-  mirror_hardwares: [amdexperimental]
   torch_nightly: true
   source_file_dependencies:
   - vllm/
@@ -40,7 +37,7 @@ steps:
   commands:
     # Install fast path packages for testing against transformers
     # Note: also needed to run plamo2 model in vLLM
-    - uv pip install --system --no-build-isolation 'git+https://github.com/state-spaces/mamba@v2.2.5'
+    - uv pip install --system --no-build-isolation 'git+https://github.com/state-spaces/mamba@v2.3.0'
     - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.5.2'
     # Shard hybrid language model tests
     - pytest -v -s models/language/generation -m hybrid_model --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
@@ -48,7 +45,6 @@ steps:
 
 - label: Language Models Test (Extended Generation) # 80min
   timeout_in_minutes: 110
-  mirror_hardwares: [amdexperimental]
   optional: true
   source_file_dependencies:
   - vllm/
@@ -56,13 +52,12 @@ steps:
   commands:
     # Install fast path packages for testing against transformers
     # Note: also needed to run plamo2 model in vLLM
-    - uv pip install --system --no-build-isolation 'git+https://github.com/state-spaces/mamba@v2.2.5'
+    - uv pip install --system --no-build-isolation 'git+https://github.com/state-spaces/mamba@v2.3.0'
     - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.5.2'
     - pytest -v -s models/language/generation -m '(not core_model) and (not hybrid_model)'
 
 - label: Language Models Test (PPL)
   timeout_in_minutes: 110
-  mirror_hardwares: [amdexperimental]
   optional: true
   source_file_dependencies:
   - vllm/
@@ -72,7 +67,6 @@ steps:
 
 - label: Language Models Test (Extended Pooling)  # 36min
   timeout_in_minutes: 50
-  mirror_hardwares: [amdexperimental]
   optional: true
   source_file_dependencies:
   - vllm/
@@ -82,7 +76,6 @@ steps:
 
 - label: Language Models Test (MTEB)
   timeout_in_minutes: 110
-  mirror_hardwares: [amdexperimental]
   optional: true
   source_file_dependencies:
   - vllm/

.buildkite/test_areas/samplers.yaml

@@ -12,3 +12,10 @@ steps:
   commands:
     - pytest -v -s samplers
     - VLLM_USE_FLASHINFER_SAMPLER=1 pytest -v -s samplers
+  mirror:
+    amd:
+      device: mi325_1
+      depends_on:
+      - image-build-amd
+      commands:
+      - pytest -v -s -m 'not skip_v1' samplers

.github/workflows/cleanup_pr_body.yml

@@ -19,6 +19,7 @@ jobs:
         uses: actions/setup-python@83679a892e2d95755f2dac6acb0bfd1e9ac5d548 # v6.1.0
         with:
           python-version: '3.12'
+          cache: 'pip'
 
       - name: Install Python dependencies
         run: |

.gitignore

@@ -238,3 +238,6 @@ ep_kernels_workspace/
 vllm/grpc/vllm_engine_pb2.py
 vllm/grpc/vllm_engine_pb2_grpc.py
 vllm/grpc/vllm_engine_pb2.pyi
+
+# Ignore generated cpu headers 
+csrc/cpu/cpu_attn_dispatch_generated.h

.pre-commit-config.yaml

@@ -143,6 +143,11 @@ repos:
     name: Check attention backend documentation is up to date
     entry: python tools/pre_commit/generate_attention_backend_docs.py --check
     language: python
+  - id: check-boolean-context-manager
+    name: Check for boolean ops in with-statements
+    entry: python tools/pre_commit/check_boolean_context_manager.py
+    language: python
+    types: [python]
   # Keep `suggestion` last
   - id: suggestion
     name: Suggestion

.readthedocs.yaml

@@ -9,13 +9,14 @@ build:
     python: "3.12"
   jobs:
     post_checkout:
-      - git fetch --unshallow || true
+      - git fetch origin main --unshallow --no-tags --filter=blob:none || true
+    pre_create_environment:
+      - pip install uv
+    create_environment:
+      - uv venv $READTHEDOCS_VIRTUALENV_PATH
+    install:
+      - uv pip install --python $READTHEDOCS_VIRTUALENV_PATH/bin/python --no-cache-dir -r requirements/docs.txt 
 
 mkdocs:
   configuration: mkdocs.yaml
   fail_on_warning: true
-
-# Optionally declare the Python requirements required to build your docs
-python:
-  install:
-    - requirements: requirements/docs.txt

CMakeLists.txt

@@ -293,6 +293,7 @@ set(VLLM_EXT_SRC
   "csrc/fused_qknorm_rope_kernel.cu"
   "csrc/layernorm_quant_kernels.cu"
   "csrc/sampler.cu"
+  "csrc/topk.cu"
   "csrc/cuda_view.cu"
   "csrc/quantization/gptq/q_gemm.cu"
   "csrc/quantization/w8a8/int8/scaled_quant.cu"

benchmarks/attention_benchmarks/batch_spec.py

@@ -229,3 +229,40 @@ def get_batch_stats(requests: list[BatchRequest]) -> dict:
             sum(r.kv_len for r in requests) / len(requests) if requests else 0
         ),
     }
+
+
+def get_batch_type(batch_spec: str, spec_decode_threshold: int = 8) -> str:
+    """
+    Classify a batch spec into a type string.
+
+    Args:
+        batch_spec: Batch specification string (e.g., "q2k", "8q1s1k", "2q2k_8q1s1k")
+        spec_decode_threshold: Max q_len to be considered spec-decode vs extend
+
+    Returns:
+        Type string: "prefill", "decode", "spec-decode", "extend", or "mixed (types...)"
+    """
+    requests = parse_batch_spec(batch_spec)
+
+    # Classify each request
+    types_present = set()
+    for req in requests:
+        if req.is_decode:
+            types_present.add("decode")
+        elif req.is_prefill:
+            types_present.add("prefill")
+        elif req.is_extend:
+            # Distinguish spec-decode (small q_len) from extend (chunked prefill)
+            if req.q_len <= spec_decode_threshold:
+                types_present.add("spec-decode")
+            else:
+                types_present.add("extend")
+
+    if len(types_present) == 1:
+        return types_present.pop()
+    elif len(types_present) > 1:
+        # Sort for consistent output
+        sorted_types = sorted(types_present)
+        return f"mixed ({'+'.join(sorted_types)})"
+    else:
+        return "unknown"

benchmarks/attention_benchmarks/common.py

@@ -12,6 +12,7 @@ from typing import Any
 
 import numpy as np
 import torch
+from batch_spec import get_batch_type, parse_batch_spec
 from rich.console import Console
 from rich.table import Table
 
@@ -316,12 +317,14 @@ class ResultsFormatter:
             backends: List of backend names being compared
             compare_to_fastest: Show percentage comparison to fastest
         """
-        # Group by batch spec
+        # Group by batch spec, preserving first-occurrence order
         by_spec = {}
+        specs_order = []
         for r in results:
             spec = r.config.batch_spec
             if spec not in by_spec:
                 by_spec[spec] = {}
+                specs_order.append(spec)
             by_spec[spec][r.config.backend] = r
 
         # Create shortened backend names for display
@@ -337,6 +340,8 @@ class ResultsFormatter:
 
         table = Table(title="Attention Benchmark Results")
         table.add_column("Batch\nSpec", no_wrap=True)
+        table.add_column("Type", no_wrap=True)
+        table.add_column("Batch\nSize", justify="right", no_wrap=True)
 
         multi = len(backends) > 1
         for backend in backends:
@@ -350,12 +355,14 @@ class ResultsFormatter:
                 table.add_column(col_rel, justify="right", no_wrap=False)
 
         # Add rows
-        for spec in sorted(by_spec.keys()):
+        for spec in specs_order:
             spec_results = by_spec[spec]
             times = {b: r.mean_time for b, r in spec_results.items() if r.success}
             best_time = min(times.values()) if times else 0.0
 
-            row = [spec]
+            batch_type = get_batch_type(spec)
+            batch_size = len(parse_batch_spec(spec))
+            row = [spec, batch_type, str(batch_size)]
             for backend in backends:
                 if backend in spec_results:
                     r = spec_results[backend]

benchmarks/attention_benchmarks/configs/standard_attention.yaml

@@ -25,10 +25,18 @@ batch_specs:
   - "4q1k_16q1s2k"     # 4 prefill + 16 decode
   - "2q4k_32q1s1k"     # 2 large prefill + 32 decode
 
-  # Context extension
-  - "q1ks2k"          # 1k query, 2k sequence (chunked prefill)
+  # Speculative decode (q <= 8)
+  - "16q2s1k"         # 16 requests, 2 spec tokens, 1k KV cache
+  - "16q4s1k"         # 16 requests, 4 spec tokens, 1k KV cache
+  - "16q8s1k"         # 16 requests, 8 spec tokens, 1k KV cache
+  - "32q4s2k"         # 32 requests, 4 spec tokens, 2k KV cache
+  - "8q8s4k"          # 8 requests, 8 spec tokens, 4k KV cache
+
+  # Context extension (chunked prefill)
+  - "q1ks2k"          # 1k query, 2k sequence
   - "2q1ks4k"         # 2 requests: 1k query, 4k sequence
 
+# Available backends: flash, triton, flashinfer
 backends:
   - flash
   - triton

benchmarks/attention_benchmarks/runner.py

@@ -8,7 +8,9 @@ This module provides helpers for running standard attention backends
 (FlashAttention, Triton, FlashInfer) with real vLLM integration.
 """
 
+import logging
 import types
+from contextlib import contextmanager
 
 import numpy as np
 import torch
@@ -24,8 +26,13 @@ from vllm.config import (
     ParallelConfig,
     SchedulerConfig,
     VllmConfig,
+    set_current_vllm_config,
+)
+from vllm.v1.attention.backends.utils import (
+    CommonAttentionMetadata,
+    get_kv_cache_layout,
+    set_kv_cache_layout,
 )
-from vllm.v1.attention.backends.utils import CommonAttentionMetadata
 from vllm.v1.kv_cache_interface import FullAttentionSpec
 
 # ============================================================================
@@ -37,22 +44,14 @@ _BACKEND_CONFIG = {
     "flash": {
         "module": "vllm.v1.attention.backends.flash_attn",
         "backend_class": "FlashAttentionBackend",
-        "dtype": torch.float16,
-        "cache_layout": "standard",
-        # ^ [2, num_blocks, block_size, num_kv_heads, head_dim]
     },
     "triton": {
         "module": "vllm.v1.attention.backends.triton_attn",
         "backend_class": "TritonAttentionBackend",
-        "dtype": torch.float32,
-        "cache_layout": "standard",
     },
     "flashinfer": {
         "module": "vllm.v1.attention.backends.flashinfer",
         "backend_class": "FlashInferBackend",
-        "dtype": torch.float16,
-        "cache_layout": "flashinfer",
-        # ^ [num_blocks, 2, block_size, num_kv_heads, head_dim]
     },
 }
 
@@ -66,6 +65,18 @@ def _get_backend_config(backend: str) -> dict:
     return _BACKEND_CONFIG[backend]
 
 
+@contextmanager
+def log_warnings_and_errors_only():
+    """Temporarily set vLLM logger to WARNING level."""
+    logger = logging.getLogger("vllm")
+    old_level = logger.level
+    logger.setLevel(logging.WARNING)
+    try:
+        yield
+    finally:
+        logger.setLevel(old_level)
+
+
 # ============================================================================
 # Metadata Building Helpers
 # ============================================================================
@@ -88,11 +99,7 @@ def _build_common_attn_metadata(
     query_start_loc_cpu = query_start_loc.cpu()
 
     seq_lens = torch.tensor(kv_lens, dtype=torch.int32, device=device)
-    seq_lens_cpu = seq_lens.cpu()
-    max_seq_len = int(seq_lens_cpu.max())
-
-    context_lens = [kv - q for kv, q in zip(kv_lens, q_lens)]
-    num_computed_tokens_cpu = torch.tensor(context_lens, dtype=torch.int32)
+    max_seq_len = int(seq_lens.max().item())
 
     max_blocks = (max(kv_lens) + block_size - 1) // block_size
     num_blocks = batch_size * max_blocks
@@ -107,8 +114,6 @@ def _build_common_attn_metadata(
         query_start_loc=query_start_loc,
         query_start_loc_cpu=query_start_loc_cpu,
         seq_lens=seq_lens,
-        seq_lens_cpu=seq_lens_cpu,
-        num_computed_tokens_cpu=num_computed_tokens_cpu,
         num_reqs=batch_size,
         num_actual_tokens=total_tokens,
         max_query_len=max_query_len,
@@ -121,7 +126,6 @@ def _build_common_attn_metadata(
 
 def _create_vllm_config(
     config: BenchmarkConfig,
-    dtype: torch.dtype,
     max_num_blocks: int,
 ) -> VllmConfig:
     """Create a VllmConfig for benchmarking with mock model methods."""
@@ -129,7 +133,7 @@ def _create_vllm_config(
         model="meta-llama/Meta-Llama-3-8B",
         tokenizer="meta-llama/Meta-Llama-3-8B",
         trust_remote_code=False,
-        dtype=dtype,
+        dtype="auto",  # Use model's native dtype
         seed=0,
         max_model_len=1024,
     )
@@ -198,6 +202,7 @@ def _create_backend_impl(
     backend_cfg: dict,
     config: BenchmarkConfig,
     device: torch.device,
+    dtype: torch.dtype,
 ):
     """Create backend implementation instance."""
     import importlib
@@ -206,7 +211,6 @@ def _create_backend_impl(
     backend_class = getattr(backend_module, backend_cfg["backend_class"])
 
     scale = get_attention_scale(config.head_dim)
-    dtype = backend_cfg["dtype"]
 
     impl = backend_class.get_impl_cls()(
         num_heads=config.num_q_heads,
@@ -227,7 +231,7 @@ def _create_backend_impl(
 
     layer = MockLayer(device, kv_cache_spec=kv_cache_spec)
 
-    return backend_class, impl, layer, dtype
+    return backend_class, impl, layer
 
 
 def _create_metadata_builder(
@@ -235,11 +239,44 @@ def _create_metadata_builder(
     kv_cache_spec: FullAttentionSpec,
     vllm_config: VllmConfig,
     device: torch.device,
+    backend_name: str = "",
 ):
     """Create metadata builder instance."""
-    return backend_class.get_builder_cls()(
+    layer_names = ["layer_0"]
+    builder_cls = backend_class.get_builder_cls()
+
+    # Flashinfer needs get_per_layer_parameters mocked since we don't have
+    # real model layers registered
+    if backend_name == "flashinfer":
+        import unittest.mock
+
+        from vllm.v1.attention.backends.utils import PerLayerParameters
+
+        def mock_get_per_layer_parameters(vllm_config, layer_names, impl_cls):
+            head_size = vllm_config.model_config.get_head_size()
+            return {
+                layer_name: PerLayerParameters(
+                    window_left=-1,  # No sliding window
+                    logits_soft_cap=0.0,  # No soft cap
+                    sm_scale=1.0 / (head_size**0.5),  # Standard scale
+                )
+                for layer_name in layer_names
+            }
+
+        with unittest.mock.patch(
+            "vllm.v1.attention.backends.flashinfer.get_per_layer_parameters",
+            mock_get_per_layer_parameters,
+        ):
+            return builder_cls(
+                kv_cache_spec=kv_cache_spec,
+                layer_names=layer_names,
+                vllm_config=vllm_config,
+                device=device,
+            )
+
+    return builder_cls(
         kv_cache_spec=kv_cache_spec,
-        layer_names=["layer_0"],
+        layer_names=layer_names,
         vllm_config=vllm_config,
         device=device,
     )
@@ -281,39 +318,44 @@ def _create_input_tensors(
 def _create_kv_cache(
     config: BenchmarkConfig,
     max_num_blocks: int,
-    cache_layout: str,
+    backend_class,
     device: torch.device,
     dtype: torch.dtype,
 ) -> list:
-    """Create KV cache tensors for all layers."""
-    if cache_layout == "flashinfer":
-        # FlashInfer layout: [num_blocks, 2, block_size, num_kv_heads, head_dim]
-        cache_list = [
-            torch.zeros(
-                max_num_blocks,
-                2,
-                config.block_size,
-                config.num_kv_heads,
-                config.head_dim,
-                device=device,
-                dtype=dtype,
-            )
-            for _ in range(config.num_layers)
-        ]
-    else:
-        # Standard layout: [2, num_blocks, block_size, num_kv_heads, head_dim]
-        cache_list = [
-            torch.zeros(
-                2,
-                max_num_blocks,
-                config.block_size,
-                config.num_kv_heads,
-                config.head_dim,
-                device=device,
-                dtype=dtype,
-            )
-            for _ in range(config.num_layers)
-        ]
+    """Create KV cache tensors for all layers using the backend's methods.
+
+    Uses the backend's get_kv_cache_shape() and get_kv_cache_stride_order()
+    to create the cache with the correct shape and memory layout.
+    """
+    # Get the logical shape from the backend
+    cache_shape = backend_class.get_kv_cache_shape(
+        num_blocks=max_num_blocks,
+        block_size=config.block_size,
+        num_kv_heads=config.num_kv_heads,
+        head_size=config.head_dim,
+    )
+
+    # Get the stride order for custom memory layout
+    try:
+        stride_order = backend_class.get_kv_cache_stride_order()
+        assert len(stride_order) == len(cache_shape)
+    except (AttributeError, NotImplementedError):
+        stride_order = tuple(range(len(cache_shape)))
+
+    # Permute shape to physical layout order
+    physical_shape = tuple(cache_shape[i] for i in stride_order)
+
+    # Compute inverse permutation to get back to logical view
+    inv_order = [stride_order.index(i) for i in range(len(stride_order))]
+
+    cache_list = []
+    for _ in range(config.num_layers):
+        # Allocate in physical layout order (contiguous in memory)
+        cache = torch.zeros(*physical_shape, device=device, dtype=dtype)
+        # Permute to logical view
+        cache = cache.permute(*inv_order)
+        cache_list.append(cache)
+
     return cache_list
 
 
@@ -418,53 +460,72 @@ def run_attention_benchmark(config: BenchmarkConfig) -> BenchmarkResult:
     kv_lens = [r.kv_len for r in requests]
     total_q = sum(q_lens)
     max_kv = max(kv_lens)
+    batch_size = len(q_lens)
 
-    max_num_blocks = (max_kv + config.block_size - 1) // config.block_size
+    # Calculate total blocks needed: batch_size * max_blocks_per_request
+    max_blocks_per_request = (max_kv + config.block_size - 1) // config.block_size
+    max_num_blocks = batch_size * max_blocks_per_request
 
-    backend_class, impl, layer, dtype = _create_backend_impl(
-        backend_cfg, config, device
-    )
+    # Suppress vLLM logs during setup to reduce spam
+    with log_warnings_and_errors_only():
+        # Create vllm_config first - uses model's native dtype via "auto"
+        vllm_config = _create_vllm_config(config, max_num_blocks)
+        dtype = vllm_config.model_config.dtype
 
-    common_metadata = _build_common_attn_metadata(
-        q_lens, kv_lens, config.block_size, device
-    )
+        # Wrap everything in set_current_vllm_config context
+        # This is required for backends like flashinfer that need global config
+        with set_current_vllm_config(vllm_config):
+            backend_class, impl, layer = _create_backend_impl(
+                backend_cfg, config, device, dtype
+            )
 
-    kv_cache_spec = FullAttentionSpec(
-        block_size=config.block_size,
-        num_kv_heads=config.num_kv_heads,
-        head_size=config.head_dim,
-        dtype=dtype,
-    )
+            # Set KV cache layout if the backend requires a specific one
+            # (e.g., FlashInfer requires HND on SM100/Blackwell for TRTLLM attention)
+            required_layout = backend_class.get_required_kv_cache_layout()
+            if required_layout is not None:
+                set_kv_cache_layout(required_layout)
+                get_kv_cache_layout.cache_clear()
 
-    vllm_config = _create_vllm_config(config, dtype, max_num_blocks)
+            common_metadata = _build_common_attn_metadata(
+                q_lens, kv_lens, config.block_size, device
+            )
 
-    builder = _create_metadata_builder(
-        backend_class, kv_cache_spec, vllm_config, device
-    )
+            kv_cache_spec = FullAttentionSpec(
+                block_size=config.block_size,
+                num_kv_heads=config.num_kv_heads,
+                head_size=config.head_dim,
+                dtype=dtype,
+            )
 
-    attn_metadata = builder.build(
-        common_prefix_len=0,
-        common_attn_metadata=common_metadata,
-    )
+            builder = _create_metadata_builder(
+                backend_class, kv_cache_spec, vllm_config, device, config.backend
+            )
 
-    q_list, k_list, v_list = _create_input_tensors(config, total_q, device, dtype)
+            attn_metadata = builder.build(
+                common_prefix_len=0,
+                common_attn_metadata=common_metadata,
+            )
 
-    cache_list = _create_kv_cache(
-        config, max_num_blocks, backend_cfg["cache_layout"], device, dtype
-    )
+            q_list, k_list, v_list = _create_input_tensors(
+                config, total_q, device, dtype
+            )
 
-    times, mem_stats = _run_single_benchmark(
-        config,
-        impl,
-        layer,
-        q_list,
-        k_list,
-        v_list,
-        cache_list,
-        attn_metadata,
-        device,
-        dtype,
-    )
+            cache_list = _create_kv_cache(
+                config, max_num_blocks, backend_class, device, dtype
+            )
+
+            times, mem_stats = _run_single_benchmark(
+                config,
+                impl,
+                layer,
+                q_list,
+                k_list,
+                v_list,
+                cache_list,
+                attn_metadata,
+                device,
+                dtype,
+            )
 
     mean_time = np.mean(times)
     throughput = total_q / mean_time if mean_time > 0 else 0

benchmarks/kernels/benchmark_cutlass_moe_fp8.py

@@ -11,6 +11,7 @@ import torch
 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
 from tests.kernels.moe.utils import make_dummy_moe_config
 from vllm import _custom_ops as ops
+from vllm.model_executor.layers.fused_moe.activation import MoEActivation
 from vllm.model_executor.layers.fused_moe.config import fp8_w8a8_moe_quant_config
 from vllm.model_executor.layers.fused_moe.cutlass_moe import CutlassExpertsFp8
 from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts, fused_topk
@@ -161,7 +162,7 @@ def bench_run(
                 w2_fp8q_cutlass,
                 topk_weights,
                 topk_ids,
-                activation="silu",
+                activation=MoEActivation.SILU,
                 global_num_experts=num_experts,
             )
     torch.cuda.synchronize()

benchmarks/kernels/benchmark_fused_collective.py

@@ -5,7 +5,7 @@
 Benchmark for FlashInfer fused collective operations vs standard operations.
 
 This benchmark compares:
-1. FlashInfer's trtllm_allreduce_fusion (fused allreduce + rmsnorm + optional quant)
+1. FlashInfer's allreduce_fusion (fused allreduce + rmsnorm + optional quant)
 2. Standard tensor_model_parallel_all_reduce + separate rmsnorm/quant operations
 
 Usage with torchrun:
@@ -24,7 +24,6 @@ import torch.distributed as dist  # type: ignore
 
 from vllm.config.vllm import CompilationConfig, VllmConfig, set_current_vllm_config
 from vllm.distributed import (
-    get_tp_group,
     tensor_model_parallel_all_reduce,
 )
 from vllm.distributed.parallel_state import (
@@ -52,11 +51,12 @@ logger = init_logger(__name__)
 try:
     import flashinfer.comm as flashinfer_comm  # type: ignore
 
-    if not hasattr(flashinfer_comm, "trtllm_allreduce_fusion"):
+    if not (
+        hasattr(flashinfer_comm, "allreduce_fusion")
+        and hasattr(flashinfer_comm, "create_allreduce_fusion_workspace")
+    ):
         flashinfer_comm = None
-        logger.warning(
-            "FlashInfer comm module found but missing trtllm_allreduce_fusion"
-        )
+        logger.warning("FlashInfer comm module found but missing allreduce_fusion API")
 except ImportError:
     flashinfer_comm = None
     logger.warning("FlashInfer not found, only benchmarking standard operations")
@@ -75,7 +75,7 @@ _FI_MAX_SIZES = {
 }
 
 # Global workspace tensor for FlashInfer
-_FI_WORKSPACE_TENSOR = None
+_FI_WORKSPACE = None
 
 
 def setup_flashinfer_workspace(
@@ -83,10 +83,10 @@ def setup_flashinfer_workspace(
     rank: int,
     hidden_dim: int,
     max_token_num: int,
-    use_fp32_lamport: bool = False,
+    dtype: torch.dtype,
 ):
     """Setup FlashInfer workspace for fused allreduce operations."""
-    global _FI_WORKSPACE_TENSOR
+    global _FI_WORKSPACE
 
     if flashinfer_comm is None:
         return None, None
@@ -96,33 +96,29 @@ def setup_flashinfer_workspace(
         return None, None
 
     try:
-        # Create IPC workspace
-        ipc_handles, workspace_tensor = (
-            flashinfer_comm.trtllm_create_ipc_workspace_for_all_reduce_fusion(
-                tp_rank=rank,
-                tp_size=world_size,
-                max_token_num=max_token_num,
-                hidden_dim=hidden_dim,
-                group=get_tp_group().device_group,
-                use_fp32_lamport=use_fp32_lamport,
-            )
+        workspace = flashinfer_comm.create_allreduce_fusion_workspace(
+            backend="trtllm",
+            world_size=world_size,
+            rank=rank,
+            max_token_num=max_token_num,
+            hidden_dim=hidden_dim,
+            dtype=dtype,
         )
 
-        _FI_WORKSPACE_TENSOR = workspace_tensor
-        return ipc_handles, workspace_tensor
+        _FI_WORKSPACE = workspace
+        return workspace
     except Exception as e:
         logger.error("Failed to setup FlashInfer workspace: %s", e)
-        return None, None
+        return None
 
 
-def cleanup_flashinfer_workspace(ipc_handles):
+def cleanup_flashinfer_workspace(workspace):
     """Cleanup FlashInfer workspace."""
-    if flashinfer_comm is None or ipc_handles is None:
+    if flashinfer_comm is None or workspace is None:
         return
 
     try:
-        group = get_tp_group().device_group
-        flashinfer_comm.trtllm_destroy_ipc_workspace_for_all_reduce(ipc_handles, group)
+        workspace.destroy()
     except Exception as e:
         logger.error("Failed to cleanup FlashInfer workspace: %s", e)
 
@@ -132,25 +128,15 @@ class FlashInferFusedAllReduceParams:
 
     def __init__(
         self,
-        rank: int,
-        world_size: int,
-        use_fp32_lamport: bool = False,
         max_token_num: int = 1024,
     ):
-        self.rank = rank
-        self.world_size = world_size
-        self.use_fp32_lamport = use_fp32_lamport
-        self.trigger_completion_at_end = True
         self.launch_with_pdl = True
         self.fp32_acc = True
         self.max_token_num = max_token_num
 
     def get_trtllm_fused_allreduce_kwargs(self):
         return {
-            "world_rank": self.rank,
-            "world_size": self.world_size,
             "launch_with_pdl": self.launch_with_pdl,
-            "trigger_completion_at_end": self.trigger_completion_at_end,
             "fp32_acc": self.fp32_acc,
         }
 
@@ -165,7 +151,7 @@ def flashinfer_fused_allreduce_rmsnorm(
     norm_out: torch.Tensor | None = None,
 ):
     """FlashInfer fused allreduce + rmsnorm operation."""
-    if flashinfer_comm is None or _FI_WORKSPACE_TENSOR is None:
+    if flashinfer_comm is None or _FI_WORKSPACE is None:
         raise RuntimeError("FlashInfer not available or workspace not initialized")
 
     if norm_out is None:
@@ -174,18 +160,15 @@ def flashinfer_fused_allreduce_rmsnorm(
     else:
         residual_out = input_tensor
 
-    flashinfer_comm.trtllm_allreduce_fusion(
-        allreduce_in=input_tensor,
-        token_num=input_tensor.shape[0],
+    flashinfer_comm.allreduce_fusion(
+        input=input_tensor,
+        workspace=_FI_WORKSPACE,
+        pattern=flashinfer_comm.AllReduceFusionPattern.kARResidualRMSNorm,
         residual_in=residual,
         residual_out=residual_out,
         norm_out=norm_out,
         rms_gamma=rms_gamma,
         rms_eps=rms_eps,
-        hidden_dim=input_tensor.shape[-1],
-        workspace_ptrs=_FI_WORKSPACE_TENSOR,
-        pattern_code=flashinfer_comm.AllReduceFusionPattern.kARResidualRMSNorm,
-        allreduce_out=None,
         quant_out=None,
         scale_out=None,
         layout_code=flashinfer_comm.QuantizationSFLayout.SWIZZLED_128x4,
@@ -207,7 +190,7 @@ def flashinfer_fused_allreduce_rmsnorm_fp8_quant(
     quant_out: torch.Tensor | None = None,
 ):
     """FlashInfer fused allreduce + rmsnorm + FP8 quantization."""
-    if flashinfer_comm is None or _FI_WORKSPACE_TENSOR is None:
+    if flashinfer_comm is None or _FI_WORKSPACE is None:
         raise RuntimeError("FlashInfer not available or workspace not initialized")
 
     if norm_out is None:
@@ -216,18 +199,15 @@ def flashinfer_fused_allreduce_rmsnorm_fp8_quant(
     else:
         residual_out = input_tensor
 
-    flashinfer_comm.trtllm_allreduce_fusion(
-        allreduce_in=input_tensor,
-        token_num=input_tensor.shape[0],
+    flashinfer_comm.allreduce_fusion(
+        input=input_tensor,
+        workspace=_FI_WORKSPACE,
+        pattern=flashinfer_comm.AllReduceFusionPattern.kARResidualRMSNormFP8Quant,
         residual_in=residual,
         residual_out=residual_out,
         norm_out=norm_out,
         rms_gamma=rms_gamma,
         rms_eps=rms_eps,
-        hidden_dim=input_tensor.shape[-1],
-        workspace_ptrs=_FI_WORKSPACE_TENSOR,
-        pattern_code=flashinfer_comm.AllReduceFusionPattern.kARResidualRMSNormFP8Quant,
-        allreduce_out=None,
         quant_out=quant_out,
         scale_out=None,
         layout_code=flashinfer_comm.QuantizationSFLayout.SWIZZLED_128x4,
@@ -250,7 +230,7 @@ def flashinfer_fused_allreduce_rmsnorm_fp4_quant(
     norm_out: torch.Tensor | None = None,
 ):
     """FlashInfer fused allreduce + rmsnorm + FP4 quantization."""
-    if flashinfer_comm is None or _FI_WORKSPACE_TENSOR is None:
+    if flashinfer_comm is None or _FI_WORKSPACE is None:
         raise RuntimeError("FlashInfer not available or workspace not initialized")
 
     if norm_out is None:
@@ -259,18 +239,15 @@ def flashinfer_fused_allreduce_rmsnorm_fp4_quant(
     else:
         residual_out = input_tensor
 
-    flashinfer_comm.trtllm_allreduce_fusion(
-        allreduce_in=input_tensor,
-        token_num=input_tensor.shape[0],
+    flashinfer_comm.allreduce_fusion(
+        input=input_tensor,
+        workspace=_FI_WORKSPACE,
+        pattern=flashinfer_comm.AllReduceFusionPattern.kARResidualRMSNormFP4Quant,
         residual_in=residual,
         residual_out=residual_out,
         norm_out=norm_out,
         rms_gamma=rms_gamma,
         rms_eps=rms_eps,
-        hidden_dim=input_tensor.shape[-1],
-        workspace_ptrs=_FI_WORKSPACE_TENSOR,
-        pattern_code=flashinfer_comm.AllReduceFusionPattern.kARResidualRMSNormFP4Quant,
-        allreduce_out=None,
         quant_out=quant_out,
         scale_out=output_scale,
         layout_code=flashinfer_comm.QuantizationSFLayout.SWIZZLED_128x4,
@@ -1040,23 +1017,31 @@ def main():
     configs = list(itertools.product(args.num_tokens, dtypes, residual_options))
 
     # Setup FlashInfer workspace if available
-    ipc_handles = None
+    workspace = None
     allreduce_params = None
 
     if flashinfer_comm is not None:
         # Use the largest hidden dimension for workspace setup
+        max_element_size = max(torch.finfo(dt).bits // 8 for dt in dtypes)
+        workspace_dtype = (
+            torch.float32
+            if max_element_size == 4
+            else (torch.bfloat16 if torch.bfloat16 in dtypes else torch.float16)
+        )
         max_num_token = _FI_MAX_SIZES.get(world_size) // (
-            args.hidden_dim * world_size * 2
+            args.hidden_dim * max_element_size
         )
 
-        ipc_handles, workspace_tensor = setup_flashinfer_workspace(
-            world_size, rank, args.hidden_dim, max_num_token
+        workspace = setup_flashinfer_workspace(
+            world_size,
+            rank,
+            args.hidden_dim,
+            max_num_token,
+            dtype=workspace_dtype,
         )
 
-        if workspace_tensor is not None:
+        if workspace is not None:
             allreduce_params = FlashInferFusedAllReduceParams(
-                rank=rank,
-                world_size=world_size,
                 max_token_num=max_num_token,
             )
 
@@ -1119,8 +1104,8 @@ def main():
 
     finally:
         # Cleanup
-        if ipc_handles is not None:
-            cleanup_flashinfer_workspace(ipc_handles)
+        if workspace is not None:
+            cleanup_flashinfer_workspace(workspace)
 
         dist.barrier()
 

benchmarks/kernels/benchmark_moe.py

@@ -16,6 +16,7 @@ import torch
 from ray.experimental.tqdm_ray import tqdm
 
 from vllm.model_executor.layers.fused_moe import fused_topk
+from vllm.model_executor.layers.fused_moe.activation import MoEActivation
 from vllm.model_executor.layers.fused_moe.config import (
     FusedMoEConfig,
     FusedMoEParallelConfig,
@@ -211,7 +212,8 @@ def benchmark_config(
                         hidden_dim=hidden_size,
                         intermediate_size_per_partition=shard_intermediate_size,
                         num_local_experts=num_experts,
-                        activation="silu",
+                        num_logical_experts=num_experts,
+                        activation=MoEActivation.SILU,
                         moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),
                         in_dtype=init_dtype,
                         routing_method=RoutingMethodType.TopK,
@@ -226,9 +228,10 @@ def benchmark_config(
                 x, input_gating, topk, renormalize=not use_deep_gemm
             )
 
+            inplace = not disable_inplace()
             if use_deep_gemm:
                 return deep_gemm_experts(
-                    x, w1, w2, topk_weights, topk_ids, inplace=True
+                    x, w1, w2, topk_weights, topk_ids, inplace=inplace
                 )
             return fused_experts(
                 x,
@@ -236,7 +239,7 @@ def benchmark_config(
                 w2,
                 topk_weights,
                 topk_ids,
-                inplace=True,
+                inplace=inplace,
                 quant_config=quant_config,
             )
 
@@ -686,6 +689,7 @@ def get_model_params(config):
         "DeepseekV2ForCausalLM",
         "DeepseekV3ForCausalLM",
         "DeepseekV32ForCausalLM",
+        "GlmMoeDsaForCausalLM",
         "Glm4MoeForCausalLM",
         "Glm4MoeLiteForCausalLM",
         "NemotronHForCausalLM",

cmake/external_projects/vllm_flash_attn.cmake

@@ -38,7 +38,7 @@ else()
   FetchContent_Declare(
           vllm-flash-attn
           GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git
-          GIT_TAG 188be16520ceefdc625fdf71365585d2ee348fe2
+          GIT_TAG 5824e6e2008271063c3229ab3e7032bd74abbbc6
           GIT_PROGRESS TRUE
           # Don't share the vllm-flash-attn build between build types
           BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn

csrc/activation_kernels.cu

@@ -9,6 +9,111 @@
 
 namespace vllm {
 
+struct alignas(32) u32x8_t {
+  uint32_t u0, u1, u2, u3, u4, u5, u6, u7;
+};
+
+__device__ __forceinline__ void ld256(u32x8_t& val, const u32x8_t* ptr) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000
+  asm volatile("ld.global.nc.v8.u32 {%0,%1,%2,%3,%4,%5,%6,%7}, [%8];\n"
+               : "=r"(val.u0), "=r"(val.u1), "=r"(val.u2), "=r"(val.u3),
+                 "=r"(val.u4), "=r"(val.u5), "=r"(val.u6), "=r"(val.u7)
+               : "l"(ptr));
+#else
+  const uint4* uint_ptr = reinterpret_cast<const uint4*>(ptr);
+  uint4 top_half = __ldg(&uint_ptr[0]);
+  uint4 bottom_half = __ldg(&uint_ptr[1]);
+  val.u0 = top_half.x;
+  val.u1 = top_half.y;
+  val.u2 = top_half.z;
+  val.u3 = top_half.w;
+  val.u4 = bottom_half.x;
+  val.u5 = bottom_half.y;
+  val.u6 = bottom_half.z;
+  val.u7 = bottom_half.w;
+#endif
+}
+
+__device__ __forceinline__ void st256(u32x8_t& val, u32x8_t* ptr) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000
+  asm volatile("st.global.v8.u32 [%0], {%1,%2,%3,%4,%5,%6,%7,%8};\n"
+               :
+               : "l"(ptr), "r"(val.u0), "r"(val.u1), "r"(val.u2), "r"(val.u3),
+                 "r"(val.u4), "r"(val.u5), "r"(val.u6), "r"(val.u7)
+               : "memory");
+#else
+  uint4* uint_ptr = reinterpret_cast<uint4*>(ptr);
+  uint_ptr[0] = make_uint4(val.u0, val.u1, val.u2, val.u3);
+  uint_ptr[1] = make_uint4(val.u4, val.u5, val.u6, val.u7);
+#endif
+}
+
+template <bool support_256>
+struct VecTraits;
+
+template <>
+struct VecTraits<true> {
+  static constexpr int ARCH_MAX_VEC_SIZE = 32;
+  using vec_t = u32x8_t;
+};
+
+template <>
+struct VecTraits<false> {
+  static constexpr int ARCH_MAX_VEC_SIZE = 16;
+  using vec_t = int4;
+};
+
+template <typename T>
+struct PackedTraits;
+
+template <>
+struct PackedTraits<c10::BFloat16> {
+  using packed_t = __nv_bfloat162;
+};
+
+template <>
+struct PackedTraits<c10::Half> {
+  using packed_t = __half2;
+};
+
+template <>
+struct PackedTraits<float> {
+  using packed_t = float2;
+};
+
+template <typename packed_t>
+__device__ __forceinline__ float2 cast_to_float2(const packed_t& val) {
+  if constexpr (std::is_same_v<packed_t, __nv_bfloat162>) {
+    return __bfloat1622float2(val);
+  } else if constexpr (std::is_same_v<packed_t, __half2>) {
+    return __half22float2(val);
+  } else if constexpr (std::is_same_v<packed_t, float2>) {
+    return float2(val);
+  }
+}
+
+template <typename packed_t>
+__device__ __forceinline__ packed_t cast_to_packed(const float2& val) {
+  if constexpr (std::is_same_v<packed_t, __nv_bfloat162>) {
+    return __float22bfloat162_rn(val);
+  } else if constexpr (std::is_same_v<packed_t, __half2>) {
+    return __float22half2_rn(val);
+  } else if constexpr (std::is_same_v<packed_t, float2>) {
+    return float2(val);
+  }
+}
+
+template <typename packed_t>
+__device__ __forceinline__ packed_t packed_mul(const packed_t& x,
+                                               const packed_t& y) {
+  if constexpr (std::is_same_v<packed_t, __nv_bfloat162> ||
+                std::is_same_v<packed_t, __half2>) {
+    return __hmul2(x, y);
+  } else if constexpr (std::is_same_v<packed_t, float2>) {
+    return make_float2(x.x * y.x, x.y * y.y);
+  }
+}
+
 template <typename scalar_t, scalar_t (*ACT_FN)(const scalar_t&),
           bool act_first>
 __device__ __forceinline__ scalar_t compute(const scalar_t& x,
@@ -16,52 +121,69 @@ __device__ __forceinline__ scalar_t compute(const scalar_t& x,
   return act_first ? ACT_FN(x) * y : x * ACT_FN(y);
 }
 
+template <typename packed_t, packed_t (*PACKED_ACT_FN)(const packed_t&),
+          bool act_first>
+__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));
+}
+
 // Check if all pointers are 16-byte aligned for int4 vectorized access
-__device__ __forceinline__ bool is_16byte_aligned(const void* ptr) {
+__host__ __device__ __forceinline__ bool is_16byte_aligned(const void* ptr) {
   return (reinterpret_cast<uintptr_t>(ptr) & 15) == 0;
 }
 
+// Check if all pointers are 16-byte aligned for longlong4_32a vectorized access
+__host__ __device__ __forceinline__ bool is_32byte_aligned(const void* ptr) {
+  return (reinterpret_cast<uintptr_t>(ptr) & 31) == 0;
+}
+
 // Activation and gating kernel template.
-template <typename scalar_t, scalar_t (*ACT_FN)(const scalar_t&),
-          bool act_first>
+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>
 __global__ void act_and_mul_kernel(
     scalar_t* __restrict__ out,          // [..., d]
     const scalar_t* __restrict__ input,  // [..., 2, d]
     const int d) {
-  constexpr int VEC_SIZE = 16 / sizeof(scalar_t);
-  const int64_t token_idx = blockIdx.x;
-  const scalar_t* x_ptr = input + token_idx * 2 * d;
+  const scalar_t* x_ptr = input + blockIdx.x * 2 * d;
   const scalar_t* y_ptr = x_ptr + d;
-  scalar_t* out_ptr = out + token_idx * d;
+  scalar_t* out_ptr = out + blockIdx.x * d;
 
-  // Check alignment for 128-bit vectorized access.
-  // All three pointers must be 16-byte aligned for safe int4 operations.
-  const bool aligned = is_16byte_aligned(x_ptr) && is_16byte_aligned(y_ptr) &&
-                       is_16byte_aligned(out_ptr);
+  if constexpr (use_vec) {
+    // Fast path: 128-bit/256-bit vectorized loop
+    using vec_t = typename VecTraits<use_256b>::vec_t;
+    constexpr int ARCH_MAX_VEC_SIZE = VecTraits<use_256b>::ARCH_MAX_VEC_SIZE;
+    constexpr int VEC_SIZE = ARCH_MAX_VEC_SIZE / sizeof(packed_t);
 
-  if (aligned && d >= VEC_SIZE) {
-    // Fast path: 128-bit vectorized loop
-    const int4* x_vec = reinterpret_cast<const int4*>(x_ptr);
-    const int4* y_vec = reinterpret_cast<const int4*>(y_ptr);
-    int4* out_vec = reinterpret_cast<int4*>(out_ptr);
-    const int num_vecs = d / VEC_SIZE;
-    const int vec_end = num_vecs * VEC_SIZE;
+    const vec_t* x_vec = reinterpret_cast<const vec_t*>(x_ptr);
+    const vec_t* y_vec = reinterpret_cast<const vec_t*>(y_ptr);
+    vec_t* out_vec = reinterpret_cast<vec_t*>(out_ptr);
+    const int num_vecs = d / 2 / VEC_SIZE;
 
     for (int i = threadIdx.x; i < num_vecs; i += blockDim.x) {
-      int4 x = VLLM_LDG(&x_vec[i]), y = VLLM_LDG(&y_vec[i]), r;
-      auto* xp = reinterpret_cast<scalar_t*>(&x);
-      auto* yp = reinterpret_cast<scalar_t*>(&y);
-      auto* rp = reinterpret_cast<scalar_t*>(&r);
+      vec_t x, y;
+      if constexpr (use_256b) {
+        ld256(x, &x_vec[i]);
+        ld256(y, &y_vec[i]);
+      } else {
+        x = VLLM_LDG(&x_vec[i]);
+        y = VLLM_LDG(&y_vec[i]);
+      }
+      auto* xp = reinterpret_cast<packed_t*>(&x);
+      auto* yp = reinterpret_cast<packed_t*>(&y);
 #pragma unroll
       for (int j = 0; j < VEC_SIZE; j++) {
-        rp[j] = compute<scalar_t, ACT_FN, act_first>(xp[j], yp[j]);
+        xp[j] =
+            packed_compute<packed_t, PACKED_ACT_FN, act_first>(xp[j], yp[j]);
+      }
+      if constexpr (use_256b) {
+        st256(x, &out_vec[i]);
+      } else {
+        out_vec[i] = x;
       }
-      out_vec[i] = r;
-    }
-    // Scalar cleanup for remaining elements
-    for (int i = vec_end + threadIdx.x; i < d; i += blockDim.x) {
-      out_ptr[i] = compute<scalar_t, ACT_FN, act_first>(VLLM_LDG(&x_ptr[i]),
-                                                        VLLM_LDG(&y_ptr[i]));
     }
   } else {
     // Scalar fallback for unaligned data or small d
@@ -79,6 +201,15 @@ __device__ __forceinline__ T silu_kernel(const T& x) {
   return (T)(((float)x) / (1.0f + expf((float)-x)));
 }
 
+template <typename packed_t>
+__device__ __forceinline__ packed_t packed_silu_kernel(const packed_t& val) {
+  // x * sigmoid(x)
+  float2 fval = cast_to_float2(val);
+  fval.x = fval.x / (1.0f + expf(-fval.x));
+  fval.y = fval.y / (1.0f + expf(-fval.y));
+  return cast_to_packed<packed_t>(fval);
+}
+
 template <typename T>
 __device__ __forceinline__ T gelu_kernel(const T& x) {
   // Equivalent to PyTorch GELU with 'none' approximation.
@@ -89,6 +220,18 @@ __device__ __forceinline__ T gelu_kernel(const T& x) {
   return (T)(f * 0.5f * (1.0f + ::erf(f * ALPHA)));
 }
 
+template <typename packed_t>
+__device__ __forceinline__ packed_t packed_gelu_kernel(const packed_t& val) {
+  // Equivalent to PyTorch GELU with 'none' approximation.
+  // Refer to:
+  // https://github.com/pytorch/pytorch/blob/8ac9b20d4b090c213799e81acf48a55ea8d437d6/aten/src/ATen/native/cuda/ActivationGeluKernel.cu#L36-L38
+  constexpr float ALPHA = M_SQRT1_2;
+  float2 fval = cast_to_float2(val);
+  fval.x = fval.x * 0.5f * (1.0f + ::erf(fval.x * ALPHA));
+  fval.y = fval.y * 0.5f * (1.0f + ::erf(fval.y * ALPHA));
+  return cast_to_packed<packed_t>(fval);
+}
+
 template <typename T>
 __device__ __forceinline__ T gelu_tanh_kernel(const T& x) {
   // Equivalent to PyTorch GELU with 'tanh' approximation.
@@ -102,32 +245,83 @@ __device__ __forceinline__ T gelu_tanh_kernel(const T& x) {
   return (T)(0.5f * f * (1.0f + ::tanhf(inner)));
 }
 
+template <typename packed_t>
+__device__ __forceinline__ packed_t
+packed_gelu_tanh_kernel(const packed_t& val) {
+  // Equivalent to PyTorch GELU with 'tanh' approximation.
+  // Refer to:
+  // https://github.com/pytorch/pytorch/blob/8ac9b20d4b090c213799e81acf48a55ea8d437d6/aten/src/ATen/native/cuda/ActivationGeluKernel.cu#L25-L30
+  float2 fval = cast_to_float2(val);
+  constexpr float BETA = M_SQRT2 * M_2_SQRTPI * 0.5f;
+  constexpr float KAPPA = 0.044715;
+
+  float x_cube = fval.x * fval.x * fval.x;
+  float inner = BETA * (fval.x + KAPPA * x_cube);
+  fval.x = 0.5f * fval.x * (1.0f + ::tanhf(inner));
+
+  x_cube = fval.y * fval.y * fval.y;
+  inner = BETA * (fval.y + KAPPA * x_cube);
+  fval.y = 0.5f * fval.y * (1.0f + ::tanhf(inner));
+  return cast_to_packed<packed_t>(fval);
+}
+
 }  // namespace vllm
 
 // Launch activation and gating kernel.
 // Use ACT_FIRST (bool) indicating whether to apply the activation function
 // first.
-#define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL, ACT_FIRST)                 \
-  int d = input.size(-1) / 2;                                            \
-  int64_t num_tokens = input.numel() / input.size(-1);                   \
-  dim3 grid(num_tokens);                                                 \
-  dim3 block(std::min(d, 1024));                                         \
-  if (num_tokens == 0) {                                                 \
-    return;                                                              \
-  }                                                                      \
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));      \
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();          \
-  VLLM_DISPATCH_FLOATING_TYPES(                                          \
-      input.scalar_type(), "act_and_mul_kernel", [&] {                   \
-        vllm::act_and_mul_kernel<scalar_t, KERNEL<scalar_t>, ACT_FIRST>  \
-            <<<grid, block, 0, stream>>>(out.data_ptr<scalar_t>(),       \
-                                         input.data_ptr<scalar_t>(), d); \
-      });
+#define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL, PACKED_KERNEL, ACT_FIRST)     \
+  auto dtype = input.scalar_type();                                         \
+  int d = input.size(-1) / 2;                                               \
+  int64_t num_tokens = input.numel() / input.size(-1);                      \
+  if (num_tokens == 0) {                                                    \
+    return;                                                                 \
+  }                                                                         \
+  dim3 grid(num_tokens);                                                    \
+  int cc_major = at::cuda::getCurrentDeviceProperties()->major;             \
+  int support_vec = (cc_major >= 10 && num_tokens > 128) ? 32 : 16;         \
+  int vec_size = support_vec / at::elementSize(dtype);                      \
+  const bool use_vec = (d % vec_size == 0);                                 \
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));         \
+  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();             \
+  if (use_vec) {                                                            \
+    dim3 block(std::min(d / vec_size, 1024));                               \
+    if (cc_major >= 10 && num_tokens > 128) {                               \
+      VLLM_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel", [&] {       \
+        vllm::act_and_mul_kernel<                                           \
+            scalar_t, typename vllm::PackedTraits<scalar_t>::packed_t,      \
+            KERNEL<scalar_t>,                                               \
+            PACKED_KERNEL<typename vllm::PackedTraits<scalar_t>::packed_t>, \
+            ACT_FIRST, true, true><<<grid, block, 0, stream>>>(             \
+            out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d);       \
+      });                                                                   \
+    } else {                                                                \
+      VLLM_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel", [&] {       \
+        vllm::act_and_mul_kernel<                                           \
+            scalar_t, typename vllm::PackedTraits<scalar_t>::packed_t,      \
+            KERNEL<scalar_t>,                                               \
+            PACKED_KERNEL<typename vllm::PackedTraits<scalar_t>::packed_t>, \
+            ACT_FIRST, true, false><<<grid, block, 0, stream>>>(            \
+            out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d);       \
+      });                                                                   \
+    }                                                                       \
+  } else {                                                                  \
+    dim3 block(std::min(d, 1024));                                          \
+    VLLM_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel", [&] {         \
+      vllm::act_and_mul_kernel<                                             \
+          scalar_t, typename vllm::PackedTraits<scalar_t>::packed_t,        \
+          KERNEL<scalar_t>,                                                 \
+          PACKED_KERNEL<typename vllm::PackedTraits<scalar_t>::packed_t>,   \
+          ACT_FIRST, false><<<grid, block, 0, stream>>>(                    \
+          out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d);         \
+    });                                                                     \
+  }
 
 void silu_and_mul(torch::Tensor& out,    // [..., d]
                   torch::Tensor& input)  // [..., 2 * d]
 {
-  LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, true);
+  LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel,
+                                true);
 }
 
 void mul_and_silu(torch::Tensor& out,    // [..., d]
@@ -135,19 +329,22 @@ 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, false);
+  LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel,
+                                false);
 }
 
 void gelu_and_mul(torch::Tensor& out,    // [..., d]
                   torch::Tensor& input)  // [..., 2 * d]
 {
-  LAUNCH_ACTIVATION_GATE_KERNEL(vllm::gelu_kernel, true);
+  LAUNCH_ACTIVATION_GATE_KERNEL(vllm::gelu_kernel, vllm::packed_gelu_kernel,
+                                true);
 }
 
 void gelu_tanh_and_mul(torch::Tensor& out,    // [..., d]
                        torch::Tensor& input)  // [..., 2 * d]
 {
-  LAUNCH_ACTIVATION_GATE_KERNEL(vllm::gelu_tanh_kernel, true);
+  LAUNCH_ACTIVATION_GATE_KERNEL(vllm::gelu_tanh_kernel,
+                                vllm::packed_gelu_tanh_kernel, true);
 }
 
 namespace vllm {
@@ -158,42 +355,57 @@ __device__ __forceinline__ T fatrelu_kernel(const T& x, const float threshold) {
   return (T)(f > threshold ? f : 0.0f);
 }
 
-template <typename scalar_t, scalar_t (*ACT_FN)(const scalar_t&, const float)>
+template <typename packed_t>
+__device__ __forceinline__ packed_t
+packed_fatrelu_kernel(const packed_t& val, const float threshold) {
+  float2 fval = cast_to_float2(val);
+  fval.x = fval.x > threshold ? fval.x : 0.0f;
+  fval.y = fval.y > threshold ? fval.y : 0.0f;
+  return cast_to_packed<packed_t>(fval);
+}
+
+template <typename scalar_t, typename packed_t,
+          scalar_t (*ACT_FN)(const scalar_t&, const float),
+          packed_t (*PACKED_ACT_FN)(const packed_t&, const float), bool use_vec,
+          bool use_256b = false>
 __global__ void act_and_mul_kernel_with_param(
     scalar_t* __restrict__ out, const scalar_t* __restrict__ input, const int d,
     const float param) {
-  constexpr int VEC_SIZE = 16 / sizeof(scalar_t);
-  const int64_t token_idx = blockIdx.x;
-  const scalar_t* x_ptr = input + token_idx * 2 * d;
+  const scalar_t* x_ptr = input + blockIdx.x * 2 * d;
   const scalar_t* y_ptr = x_ptr + d;
-  scalar_t* out_ptr = out + token_idx * d;
+  scalar_t* out_ptr = out + blockIdx.x * d;
 
-  // Check alignment for 128-bit vectorized access
-  const bool aligned = is_16byte_aligned(x_ptr) && is_16byte_aligned(y_ptr) &&
-                       is_16byte_aligned(out_ptr);
+  if constexpr (use_vec) {
+    // Fast path: 128-bit/256-bit vectorized loop
+    using vec_t = typename VecTraits<use_256b>::vec_t;
+    constexpr int ARCH_MAX_VEC_SIZE = VecTraits<use_256b>::ARCH_MAX_VEC_SIZE;
+    constexpr int VEC_SIZE = ARCH_MAX_VEC_SIZE / sizeof(packed_t);
 
-  if (aligned && d >= VEC_SIZE) {
-    // Fast path: 128-bit vectorized loop
-    const int4* x_vec = reinterpret_cast<const int4*>(x_ptr);
-    const int4* y_vec = reinterpret_cast<const int4*>(y_ptr);
-    int4* out_vec = reinterpret_cast<int4*>(out_ptr);
-    const int num_vecs = d / VEC_SIZE;
-    const int vec_end = num_vecs * VEC_SIZE;
+    const vec_t* x_vec = reinterpret_cast<const vec_t*>(x_ptr);
+    const vec_t* y_vec = reinterpret_cast<const vec_t*>(y_ptr);
+    vec_t* out_vec = reinterpret_cast<vec_t*>(out_ptr);
+    const int num_vecs = d / 2 / VEC_SIZE;
 
     for (int i = threadIdx.x; i < num_vecs; i += blockDim.x) {
-      int4 x = VLLM_LDG(&x_vec[i]), y = VLLM_LDG(&y_vec[i]), r;
-      auto* xp = reinterpret_cast<scalar_t*>(&x);
-      auto* yp = reinterpret_cast<scalar_t*>(&y);
-      auto* rp = reinterpret_cast<scalar_t*>(&r);
+      vec_t x, y;
+      if constexpr (use_256b) {
+        ld256(x, &x_vec[i]);
+        ld256(y, &y_vec[i]);
+      } else {
+        x = VLLM_LDG(&x_vec[i]);
+        y = VLLM_LDG(&y_vec[i]);
+      }
+      auto* xp = reinterpret_cast<packed_t*>(&x);
+      auto* yp = reinterpret_cast<packed_t*>(&y);
 #pragma unroll
       for (int j = 0; j < VEC_SIZE; j++) {
-        rp[j] = ACT_FN(xp[j], param) * yp[j];
+        xp[j] = packed_mul(PACKED_ACT_FN(xp[j], param), yp[j]);
+      }
+      if constexpr (use_256b) {
+        st256(x, &out_vec[i]);
+      } else {
+        out_vec[i] = x;
       }
-      out_vec[i] = r;
-    }
-    // Scalar cleanup for remaining elements
-    for (int i = vec_end + threadIdx.x; i < d; i += blockDim.x) {
-      out_ptr[i] = ACT_FN(VLLM_LDG(&x_ptr[i]), param) * VLLM_LDG(&y_ptr[i]);
     }
   } else {
     // Scalar fallback for unaligned data or small d
@@ -276,20 +488,58 @@ __global__ void swigluoai_and_mul_kernel(
 
 }  // namespace vllm
 
-#define LAUNCH_ACTIVATION_GATE_KERNEL_WITH_PARAM(KERNEL, PARAM)         \
-  int d = input.size(-1) / 2;                                           \
-  int64_t num_tokens = input.numel() / input.size(-1);                  \
-  dim3 grid(num_tokens);                                                \
-  dim3 block(std::min(d, 1024));                                        \
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));     \
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();         \
-  VLLM_DISPATCH_FLOATING_TYPES(                                         \
-      input.scalar_type(), "act_and_mul_kernel_with_param", [&] {       \
-        vllm::act_and_mul_kernel_with_param<scalar_t, KERNEL<scalar_t>> \
-            <<<grid, block, 0, stream>>>(out.data_ptr<scalar_t>(),      \
-                                         input.data_ptr<scalar_t>(), d, \
-                                         PARAM);                        \
-      });
+#define LAUNCH_ACTIVATION_GATE_KERNEL_WITH_PARAM(KERNEL, PACKED_KERNEL, PARAM) \
+  auto dtype = input.scalar_type();                                            \
+  int d = input.size(-1) / 2;                                                  \
+  int64_t num_tokens = input.numel() / input.size(-1);                         \
+  if (num_tokens == 0) {                                                       \
+    return;                                                                    \
+  }                                                                            \
+  dim3 grid(num_tokens);                                                       \
+  int cc_major = at::cuda::getCurrentDeviceProperties()->major;                \
+  int support_vec = (cc_major >= 10 && num_tokens > 128) ? 32 : 16;            \
+  int vec_size = support_vec / at::elementSize(dtype);                         \
+  const bool use_vec = (d % vec_size == 0);                                    \
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));            \
+  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();                \
+  if (use_vec) {                                                               \
+    dim3 block(std::min(d / vec_size, 1024));                                  \
+    if (cc_major >= 10 && num_tokens > 128) {                                  \
+      VLLM_DISPATCH_FLOATING_TYPES(                                            \
+          dtype, "act_and_mul_kernel_with_param", [&] {                        \
+            vllm::act_and_mul_kernel_with_param<                               \
+                scalar_t, typename vllm::PackedTraits<scalar_t>::packed_t,     \
+                KERNEL<scalar_t>,                                              \
+                PACKED_KERNEL<                                                 \
+                    typename vllm::PackedTraits<scalar_t>::packed_t>,          \
+                true, true><<<grid, block, 0, stream>>>(                       \
+                out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d,       \
+                PARAM);                                                        \
+          });                                                                  \
+    } else {                                                                   \
+      VLLM_DISPATCH_FLOATING_TYPES(                                            \
+          dtype, "act_and_mul_kernel_with_param", [&] {                        \
+            vllm::act_and_mul_kernel_with_param<                               \
+                scalar_t, typename vllm::PackedTraits<scalar_t>::packed_t,     \
+                KERNEL<scalar_t>,                                              \
+                PACKED_KERNEL<                                                 \
+                    typename vllm::PackedTraits<scalar_t>::packed_t>,          \
+                true, false><<<grid, block, 0, stream>>>(                      \
+                out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d,       \
+                PARAM);                                                        \
+          });                                                                  \
+    }                                                                          \
+  } else {                                                                     \
+    dim3 block(std::min(d, 1024));                                             \
+    VLLM_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel_with_param", [&] { \
+      vllm::act_and_mul_kernel_with_param<                                     \
+          scalar_t, typename vllm::PackedTraits<scalar_t>::packed_t,           \
+          KERNEL<scalar_t>,                                                    \
+          PACKED_KERNEL<typename vllm::PackedTraits<scalar_t>::packed_t>,      \
+          false><<<grid, block, 0, stream>>>(                                  \
+          out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d, PARAM);     \
+    });                                                                        \
+  }
 
 #define LAUNCH_SIGLUOAI_AND_MUL(KERNEL, ALPHA, LIMIT)                          \
   int d = input.size(-1) / 2;                                                  \
@@ -309,7 +559,8 @@ __global__ void swigluoai_and_mul_kernel(
 void fatrelu_and_mul(torch::Tensor& out,    // [..., d],
                      torch::Tensor& input,  // [..., 2 * d]
                      double threshold) {
-  LAUNCH_ACTIVATION_GATE_KERNEL_WITH_PARAM(vllm::fatrelu_kernel, threshold);
+  LAUNCH_ACTIVATION_GATE_KERNEL_WITH_PARAM(
+      vllm::fatrelu_kernel, vllm::packed_fatrelu_kernel, threshold);
 }
 void swigluoai_and_mul(torch::Tensor& out,    // [..., d]
                        torch::Tensor& input,  // [..., 2 * d]
@@ -319,39 +570,41 @@ void swigluoai_and_mul(torch::Tensor& out,    // [..., d]
 namespace vllm {
 
 // Element-wise activation kernel template.
-template <typename scalar_t, scalar_t (*ACT_FN)(const scalar_t&)>
+template <typename scalar_t, scalar_t (*ACT_FN)(const scalar_t&), bool use_vec,
+          bool use_256b = false>
 __global__ void activation_kernel(
     scalar_t* __restrict__ out,          // [..., d]
     const scalar_t* __restrict__ input,  // [..., d]
     const int d) {
-  constexpr int VEC_SIZE = 16 / sizeof(scalar_t);
-  const int64_t token_idx = blockIdx.x;
-  const scalar_t* in_ptr = input + token_idx * d;
-  scalar_t* out_ptr = out + token_idx * d;
+  const scalar_t* in_ptr = input + blockIdx.x * d;
+  scalar_t* out_ptr = out + blockIdx.x * d;
 
-  // Check alignment for 128-bit vectorized access
-  const bool aligned = is_16byte_aligned(in_ptr) && is_16byte_aligned(out_ptr);
-
-  if (aligned && d >= VEC_SIZE) {
-    // Fast path: 128-bit vectorized loop
-    const int4* in_vec = reinterpret_cast<const int4*>(in_ptr);
-    int4* out_vec = reinterpret_cast<int4*>(out_ptr);
+  if constexpr (use_vec) {
+    // Fast path: 128-bit/256-bit vectorized loop
+    using vec_t = typename VecTraits<use_256b>::vec_t;
+    constexpr int ARCH_MAX_VEC_SIZE = VecTraits<use_256b>::ARCH_MAX_VEC_SIZE;
+    constexpr int VEC_SIZE = ARCH_MAX_VEC_SIZE / sizeof(scalar_t);
+    const vec_t* in_vec = reinterpret_cast<const vec_t*>(in_ptr);
+    vec_t* out_vec = reinterpret_cast<vec_t*>(out_ptr);
     const int num_vecs = d / VEC_SIZE;
-    const int vec_end = num_vecs * VEC_SIZE;
 
     for (int i = threadIdx.x; i < num_vecs; i += blockDim.x) {
-      int4 v = VLLM_LDG(&in_vec[i]), r;
+      vec_t v;
+      if constexpr (use_256b) {
+        ld256(v, &in_vec[i]);
+      } else {
+        v = VLLM_LDG(&in_vec[i]);
+      }
       auto* vp = reinterpret_cast<scalar_t*>(&v);
-      auto* rp = reinterpret_cast<scalar_t*>(&r);
 #pragma unroll
       for (int j = 0; j < VEC_SIZE; j++) {
-        rp[j] = ACT_FN(vp[j]);
+        vp[j] = ACT_FN(vp[j]);
+      }
+      if constexpr (use_256b) {
+        st256(v, &out_vec[i]);
+      } else {
+        out_vec[i] = v;
       }
-      out_vec[i] = r;
-    }
-    // Scalar cleanup for remaining elements
-    for (int i = vec_end + threadIdx.x; i < d; i += blockDim.x) {
-      out_ptr[i] = ACT_FN(VLLM_LDG(&in_ptr[i]));
     }
   } else {
     // Scalar fallback for unaligned data or small d
@@ -365,18 +618,43 @@ __global__ void activation_kernel(
 }  // namespace vllm
 
 // Launch element-wise activation kernel.
-#define LAUNCH_ACTIVATION_KERNEL(KERNEL)                                       \
-  int d = input.size(-1);                                                      \
-  int64_t num_tokens = input.numel() / d;                                      \
-  dim3 grid(num_tokens);                                                       \
-  dim3 block(std::min(d, 1024));                                               \
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));            \
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();                \
-  VLLM_DISPATCH_FLOATING_TYPES(input.scalar_type(), "activation_kernel", [&] { \
-    vllm::activation_kernel<scalar_t, KERNEL<scalar_t>>                        \
-        <<<grid, block, 0, stream>>>(out.data_ptr<scalar_t>(),                 \
-                                     input.data_ptr<scalar_t>(), d);           \
-  });
+#define LAUNCH_ACTIVATION_KERNEL(KERNEL)                                 \
+  auto dtype = input.scalar_type();                                      \
+  int d = input.size(-1);                                                \
+  int64_t num_tokens = input.numel() / input.size(-1);                   \
+  if (num_tokens == 0) {                                                 \
+    return;                                                              \
+  }                                                                      \
+  dim3 grid(num_tokens);                                                 \
+  int cc_major = at::cuda::getCurrentDeviceProperties()->major;          \
+  int support_vec = (cc_major >= 10 && num_tokens > 128) ? 32 : 16;      \
+  int vec_size = support_vec / at::elementSize(dtype);                   \
+  const bool use_vec = (d % vec_size == 0);                              \
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));      \
+  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();          \
+  if (use_vec) {                                                         \
+    dim3 block(std::min(d / vec_size, 1024));                            \
+    if (cc_major >= 10 && num_tokens > 128) {                            \
+      VLLM_DISPATCH_FLOATING_TYPES(dtype, "activation_kernel", [&] {     \
+        vllm::activation_kernel<scalar_t, KERNEL<scalar_t>, true, true>  \
+            <<<grid, block, 0, stream>>>(out.data_ptr<scalar_t>(),       \
+                                         input.data_ptr<scalar_t>(), d); \
+      });                                                                \
+    } else {                                                             \
+      VLLM_DISPATCH_FLOATING_TYPES(dtype, "activation_kernel", [&] {     \
+        vllm::activation_kernel<scalar_t, KERNEL<scalar_t>, true, false> \
+            <<<grid, block, 0, stream>>>(out.data_ptr<scalar_t>(),       \
+                                         input.data_ptr<scalar_t>(), d); \
+      });                                                                \
+    }                                                                    \
+  } else {                                                               \
+    dim3 block(std::min(d, 1024));                                       \
+    VLLM_DISPATCH_FLOATING_TYPES(dtype, "activation_kernel", [&] {       \
+      vllm::activation_kernel<scalar_t, KERNEL<scalar_t>, false>         \
+          <<<grid, block, 0, stream>>>(out.data_ptr<scalar_t>(),         \
+                                       input.data_ptr<scalar_t>(), d);   \
+    });                                                                  \
+  }
 
 namespace vllm {
 

csrc/cache_kernels.cu

@@ -1234,8 +1234,13 @@ void cp_gather_and_upconvert_fp8_kv_cache(
               "src_cache and seq_lens must be on the same device");
   TORCH_CHECK(src_cache.device() == workspace_starts.device(),
               "src_cache and workspace_starts must be on the same device");
-
-  TORCH_CHECK(src_cache.dtype() == torch::kUInt8, "src_cache must be uint8");
+  auto dtype = src_cache.scalar_type();
+  TORCH_CHECK(
+      dtype == at::ScalarType::Byte ||               // uint8
+          dtype == at::ScalarType::Float8_e4m3fn ||  // fp8 e4m3
+          dtype == at::ScalarType::Float8_e5m2,      // fp8 e5m2
+      "src_cache must be uint8, float8_e4m3fn, or float8_e5m2, but got ",
+      src_cache.dtype());
   TORCH_CHECK(dst.dtype() == torch::kBFloat16, "dst must be bfloat16");
   TORCH_CHECK(head_dim == 576, "head_dim must be 576 for MLA");
 
@@ -1244,14 +1249,21 @@ void cp_gather_and_upconvert_fp8_kv_cache(
   int64_t cache_entry_stride = src_cache.stride(1);
   int64_t dst_entry_stride = dst.stride(0);
 
+  const uint8_t* src_ptr = nullptr;
+  if (dtype == at::ScalarType::Byte) {
+    src_ptr = src_cache.data_ptr<uint8_t>();
+  } else {
+    // float8_e4m3fn or float8_e5m2
+    src_ptr = reinterpret_cast<const uint8_t*>(src_cache.data_ptr());
+  }
+
   // Decide on the number of splits based on the batch size
   int num_splits = batch_size > 128 ? 2 : batch_size > 64 ? 4 : 16;
   dim3 grid(batch_size, num_splits);
   dim3 block(576);
 
   vllm::cp_gather_and_upconvert_fp8_kv_cache<<<grid, block, 0, stream>>>(
-      src_cache.data_ptr<uint8_t>(),
-      reinterpret_cast<__nv_bfloat16*>(dst.data_ptr()),
+      src_ptr, reinterpret_cast<__nv_bfloat16*>(dst.data_ptr()),
       block_table.data_ptr<int32_t>(), seq_lens.data_ptr<int32_t>(),
       workspace_starts.data_ptr<int32_t>(), block_size, head_dim,
       block_table_stride, cache_block_stride, cache_entry_stride,

csrc/cpu/cpu_attn_impl.hpp

@@ -821,7 +821,7 @@ struct VecTypeTrait<c10::BFloat16> {
   using vec_t = vec_op::BF16Vec16;
 };
 
-#if !defined(__powerpc__) && !defined(__s390x__)
+#if !defined(__powerpc__)
 template <>
 struct VecTypeTrait<c10::Half> {
   using vec_t = vec_op::FP16Vec16;

csrc/cpu/cpu_fused_moe.cpp

@@ -147,7 +147,7 @@ void fused_moe_impl(scalar_t* __restrict__ output, scalar_t* __restrict__ input,
                     const int32_t token_num, const int32_t expert_num,
                     const int32_t topk_num, const int32_t input_size_13,
                     const int32_t output_size_13, const int32_t input_size_2,
-                    const int32_t output_size_2) {
+                    const int32_t output_size_2, const bool skip_weighted) {
   using scalar_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t;
   constexpr int32_t gemm_n_tile_size = gemm_t::NSize;
   constexpr int32_t gemm_m_tile_size = gemm_t::MaxMSize;
@@ -582,6 +582,11 @@ void fused_moe_impl(scalar_t* __restrict__ output, scalar_t* __restrict__ input,
         scalar_t* __restrict__ curr_output_buffer =
             output + token_id * output_size_2;
 
+        if (skip_weighted) {
+          // Only for topk_num == 1
+          *curr_weight = 1.0f;
+        }
+
         if (topk_num > 1) {
           {
             int32_t w2_output_idx = curr_expand_token_id_index_buffer[0];
@@ -699,7 +704,7 @@ void cpu_fused_moe(
     const std::optional<torch::Tensor>& w2_bias,  // [expert_num, output_size_2]
     const torch::Tensor& topk_weights,            // [token_num, k], float32
     const torch::Tensor& topk_id,                 // [token_num, k], int32
-    const std::string& act, const std::string& isa) {
+    const bool skip_weighted, const std::string& act, const std::string& isa) {
   const int32_t token_num = input.size(0);
   const int32_t input_size_13 = input.size(1);
   const int64_t input_stride = input.stride(0);
@@ -711,6 +716,8 @@ void cpu_fused_moe(
   const int32_t topk_num = topk_id.size(1);
   const FusedMOEAct act_type = get_act_type(act);
   cpu_utils::ISA isa_type = cpu_utils::get_isa(isa);
+  TORCH_CHECK(!skip_weighted || topk_num == 1,
+              "skip_weighted is only supported for topk=1 on CPU");
 
   VLLM_DISPATCH_FLOATING_TYPES(w13.scalar_type(), "cpu_fused_moe", [&]() {
     CPU_ISA_DISPATCH_IMPL(isa_type, [&]() {
@@ -721,7 +728,7 @@ void cpu_fused_moe(
           w2_bias.has_value() ? w2_bias->data_ptr<scalar_t>() : nullptr,
           topk_weights.data_ptr<float>(), topk_id.data_ptr<int32_t>(), act_type,
           token_num, expert_num, topk_num, input_size_13, output_size_13,
-          input_size_2, output_size_2);
+          input_size_2, output_size_2, skip_weighted);
     });
   });
 }

csrc/cpu/cpu_types_vxe.hpp

@@ -16,10 +16,12 @@ namespace vec_op {
 #define vec_sr(a, b) ((a) >> (b))  // Vector Shift Right Algebraic
 #define vec_sl(a, b) ((a) << (b))  // Vector Shift Left
 
-// FIXME: FP16 is not fully supported in Torch-CPU
-#define VLLM_DISPATCH_CASE_FLOATING_TYPES(...)         \
-  AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \
-  AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__)
+// NOTE: FP16 (Half) is supported on s390x via custom bit-manipulation
+// conversion. PyTorch itself lacks native s390x FP16 support.
+#define VLLM_DISPATCH_CASE_FLOATING_TYPES(...)            \
+  AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__)    \
+  AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__) \
+  AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__)
 
 #define VLLM_DISPATCH_FLOATING_TYPES(TYPE, NAME, ...) \
   AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__))
@@ -86,6 +88,39 @@ struct BF16Vec8 : public Vec<BF16Vec8> {
   }
 };
 
+struct FP16Vec8 : public Vec<FP16Vec8> {
+  constexpr static int VEC_ELEM_NUM = 8;
+
+  __vector signed short reg;
+
+  explicit FP16Vec8(const void* ptr) : reg(*(__vector signed short*)ptr) {}
+  explicit FP16Vec8(const FP32Vec8&);
+
+  void save(void* ptr) const {
+    *reinterpret_cast<__vector signed short*>(ptr) = reg;
+  }
+};
+
+struct FP16Vec16 : public Vec<FP16Vec16> {
+  constexpr static int VEC_ELEM_NUM = 16;
+
+  ss16x8x2_t reg;
+
+  explicit FP16Vec16(const void* ptr) {
+    // Load 256 bits (16 FP16 values) in two parts
+    reg.val[0] = (__vector signed short)vec_xl(0, (signed short*)ptr);
+    reg.val[1] = (__vector signed short)vec_xl(16, (signed short*)ptr);
+  }
+
+  explicit FP16Vec16(const FP32Vec16&);
+
+  void save(void* ptr) const {
+    // Save 256 bits in two parts
+    vec_xst(reg.val[0], 0, (signed short*)ptr);
+    vec_xst(reg.val[1], 16, (signed short*)ptr);
+  }
+};
+
 struct BF16Vec16 : public Vec<BF16Vec16> {
   constexpr static int VEC_ELEM_NUM = 16;
 
@@ -108,6 +143,92 @@ struct BF16Vec16 : public Vec<BF16Vec16> {
 
 const static __vector signed short zero = vec_splats((signed short)0);
 
+FORCE_INLINE __vector float fp16_to_fp32_bits(__vector unsigned int x) {
+  const __vector unsigned int mask_sign = {0x8000, 0x8000, 0x8000, 0x8000};
+  const __vector unsigned int mask_exp = {0x7C00, 0x7C00, 0x7C00, 0x7C00};
+  const __vector unsigned int mask_mant = {0x03FF, 0x03FF, 0x03FF, 0x03FF};
+  const __vector unsigned int bias_adj = {112, 112, 112, 112};
+  const __vector unsigned int exp_max_fp16 = {0x1F, 0x1F, 0x1F,
+                                              0x1F};  // FP16 NaN/Inf exponent
+  const __vector unsigned int exp_max_fp32 = {0xFF, 0xFF, 0xFF,
+                                              0xFF};  // FP32 NaN/Inf exponent
+
+  __vector unsigned int s = (x & mask_sign) << 16;
+  __vector unsigned int e = (x & mask_exp) >> 10;
+  __vector unsigned int m = (x & mask_mant) << 13;
+
+  // Check for NaN/Inf: exponent = 0x1F in FP16
+  __vector __bool int is_nan_inf = vec_cmpeq(e, exp_max_fp16);
+
+  // Normal: adjust bias; NaN/Inf: set to 0xFF
+  __vector unsigned int e_normal = e + bias_adj;
+  e = vec_sel(e_normal, exp_max_fp32, is_nan_inf);
+
+  return (__vector float)(s | (e << 23) | m);
+}
+
+FORCE_INLINE __vector unsigned int fp32_to_fp16_bits(__vector float f_in) {
+  __vector unsigned int in = (__vector unsigned int)f_in;
+
+  const __vector unsigned int mask_sign_32 = {0x80000000, 0x80000000,
+                                              0x80000000, 0x80000000};
+  const __vector unsigned int mask_exp_32 = {0x7F800000, 0x7F800000, 0x7F800000,
+                                             0x7F800000};
+  const __vector unsigned int mask_mant_32 = {0x007FFFFF, 0x007FFFFF,
+                                              0x007FFFFF, 0x007FFFFF};
+
+  // Use SIGNED integers for exponent math to handle underflow check
+  const __vector signed int bias_adj = {112, 112, 112, 112};
+  const __vector signed int zero = {0, 0, 0, 0};
+  const __vector signed int max_exp = {31, 31, 31, 31};  // Max FP16 exp
+  const __vector unsigned int exp_max_fp32 = {0xFF, 0xFF, 0xFF, 0xFF};
+  const __vector unsigned int exp_max_fp16 = {0x1F, 0x1F, 0x1F, 0x1F};
+
+  __vector unsigned int s = (in & mask_sign_32) >> 16;
+  __vector unsigned int e_u = (in & mask_exp_32) >> 23;
+
+  // Check for NaN/Inf: exponent = 0xFF in FP32
+  __vector __bool int is_nan_inf = vec_cmpeq(e_u, exp_max_fp32);
+
+  __vector signed int e_s = (__vector signed int)e_u;
+  e_s = vec_sub(e_s, bias_adj);
+  e_s = vec_max(e_s, zero);
+  e_s = vec_min(e_s, max_exp);
+  __vector unsigned int e_normal = (__vector unsigned int)e_s;
+
+  __vector unsigned int e_final = vec_sel(e_normal, exp_max_fp16, is_nan_inf);
+
+  const __vector unsigned int one_v = {1, 1, 1, 1};
+  const __vector unsigned int mask_sticky = {0xFFF, 0xFFF, 0xFFF, 0xFFF};
+
+  __vector unsigned int round_bit = (in >> 12) & one_v;
+  __vector unsigned int sticky = in & mask_sticky;
+  __vector unsigned int m = (in & mask_mant_32) >> 13;
+  __vector unsigned int lsb = m & one_v;  // LSB of mantissa for tie-breaking
+
+  // Round up if: round_bit && (sticky || lsb)
+  __vector __bool int sticky_nonzero =
+      vec_cmpgt(sticky, (__vector unsigned int){0, 0, 0, 0});
+  __vector __bool int lsb_set = vec_cmpeq(lsb, one_v);
+  __vector __bool int round_up =
+      vec_and(vec_cmpeq(round_bit, one_v), vec_or(sticky_nonzero, lsb_set));
+
+  m = vec_sel(m, m + one_v, round_up);
+
+  const __vector unsigned int mant_mask = {0x3FF, 0x3FF, 0x3FF, 0x3FF};
+  const __vector unsigned int max_normal_exp = {0x1E, 0x1E, 0x1E, 0x1E};
+  __vector __bool int mant_overflows = vec_cmpgt(m, mant_mask);
+  __vector __bool int would_overflow_to_inf =
+      vec_and(mant_overflows, vec_cmpeq(e_final, max_normal_exp));
+  __vector unsigned int e_inc = vec_min(e_final + one_v, exp_max_fp16);
+  e_final = vec_sel(e_final, e_inc, mant_overflows);
+  m = vec_and(m, mant_mask);
+  e_final = vec_sel(e_final, max_normal_exp, would_overflow_to_inf);
+  m = vec_sel(m, mant_mask, would_overflow_to_inf);
+
+  return s | (e_final << 10) | m;
+}
+
 struct BF16Vec32 : public Vec<BF16Vec32> {
   constexpr static int VEC_ELEM_NUM = 32;
 
@@ -180,6 +301,18 @@ struct FP32Vec8 : public Vec<FP32Vec8> {
     reg.val[1] = (__vector float)vec_mergel(v.reg, zero);
   }
 
+  explicit FP32Vec8(const FP16Vec8& v) {
+    // Cast to UNSIGNED short vector to prevent sign-extension during unpack
+    __vector unsigned short raw_u = (__vector unsigned short)v.reg;
+
+    // Unpack 8x16-bit to two 4x32-bit vectors (Zero extended)
+    __vector unsigned int raw_hi = (__vector unsigned int)vec_unpackh(raw_u);
+    __vector unsigned int raw_lo = (__vector unsigned int)vec_unpackl(raw_u);
+
+    reg.val[0] = fp16_to_fp32_bits(raw_hi);
+    reg.val[1] = fp16_to_fp32_bits(raw_lo);
+  }
+
   float reduce_sum() const {
     AliasReg ar;
     ar.reg = reg;
@@ -531,6 +664,22 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
     reg.val[3] = (__vector float)vec_mergel(v.reg.val[1], zero);
   }
 
+  explicit FP32Vec16(const FP16Vec16& v) {
+    __vector unsigned int raw_hi_0 =
+        (__vector unsigned int)vec_unpackh(v.reg.val[0]);
+    __vector unsigned int raw_lo_0 =
+        (__vector unsigned int)vec_unpackl(v.reg.val[0]);
+    reg.val[0] = fp16_to_fp32_bits(raw_hi_0);
+    reg.val[1] = fp16_to_fp32_bits(raw_lo_0);
+
+    __vector unsigned int raw_hi_1 =
+        (__vector unsigned int)vec_unpackh(v.reg.val[1]);
+    __vector unsigned int raw_lo_1 =
+        (__vector unsigned int)vec_unpackl(v.reg.val[1]);
+    reg.val[2] = fp16_to_fp32_bits(raw_hi_1);
+    reg.val[3] = fp16_to_fp32_bits(raw_lo_1);
+  }
+
   explicit FP32Vec16(const BF16Vec8& v) : FP32Vec16(FP32Vec8(v)) {}
 
   FP32Vec16 operator*(const FP32Vec16& b) const {
@@ -628,8 +777,10 @@ struct VecType<c10::BFloat16> {
   using vec_type = BF16Vec8;
 };
 
-// On s390x, FP16 (Half) is not natively supported, use FP32 vectors instead
-using FP16Vec16 = FP32Vec16;
+template <>
+struct VecType<c10::Half> {
+  using vec_type = FP16Vec8;
+};
 
 template <typename T>
 void storeFP32(float v, T* ptr) {
@@ -650,6 +801,52 @@ inline void storeFP32<c10::BFloat16>(float v, c10::BFloat16* ptr) {
   *ptr = *(v_ptr + 1);
 }
 
+template <>
+inline void storeFP32<::c10::Half>(float v, ::c10::Half* ptr) {
+  // Use bit-manipulation for IEEE FP32 to FP16 conversion since vector
+  // intrinsics for FP32 to FP16 conversion does not use IEEE rounding and can
+  // produce incorrect results for some inputs. Process each of the 4 vectors
+  // separately.
+  uint32_t in;
+  std::memcpy(&in, &v, sizeof(in));
+
+  uint32_t s = (in & 0x80000000) >> 16;  // Sign
+  uint32_t e = (in & 0x7F800000) >> 23;  // Exponent
+  uint32_t round_bit = (in >> 12) & 1;
+  uint32_t sticky = (in & 0xFFF) != 0;  // Any bits in [11..0]
+  uint32_t m = (in & 0x007FFFFF) >> 13;
+  uint32_t lsb = m & 1;  // LSB of mantissa for tie-breaking
+
+  // Check for NaN/Inf before rounding
+  bool is_nan_inf = (e == 0xFF);
+
+  if (round_bit && (sticky || lsb)) {
+    m++;
+    // Handle mantissa overflow: if m overflows 10 bits, increment exponent
+    if (m > 0x3FF) {
+      m = 0;
+      e++;
+    }
+  }
+
+  if (is_nan_inf) {
+    // NaN/Inf: preserve it
+    e = 0x1F;
+  } else {
+    // Normal: adjust bias (127 - 15), flush subnormals to zero
+    e = (e >= 112) ? (e - 112) : 0;
+    // If exponent overflows to Inf range, saturate to max normal FP16 value
+    if (e > 0x1E) {
+      e = 0x1E;   // Max normal exponent
+      m = 0x3FF;  // Max mantissa
+    }
+  }
+
+  uint16_t fp16 = (uint16_t)(s | (e << 10) | m);
+
+  *reinterpret_cast<uint16_t*>(ptr) = fp16;
+}
+
 #ifndef __VEC_CLASS_FP_NAN
   #define __VEC_CLASS_FP_NAN (1 << 6)
 #endif
@@ -803,6 +1000,44 @@ inline BF16Vec16::BF16Vec16(const FP32Vec16& v) {
   reg.val[1] = (__vector signed short)vec_perm(inp2, inp3, omask);
 }
 
+inline FP16Vec8::FP16Vec8(const FP32Vec8& v) {
+  // Use bit-manipulation for IEEE FP32 to FP16 conversion since vector
+  // intrinsics for FP32 to FP16 conversion does not use IEEE rounding and can
+  // produce incorrect results for some inputs. Process each of the 4 vectors
+  // separately.
+  __vector unsigned int res_hi = fp32_to_fp16_bits(v.reg.val[0]);
+  __vector unsigned int res_lo = fp32_to_fp16_bits(v.reg.val[1]);
+
+  const __vector unsigned char perm_pack = {
+      2,  3,  6,  7,  10, 11, 14, 15,  // Select lower 2 bytes from res_hi
+      18, 19, 22, 23, 26, 27, 30, 31   // Select lower 2 bytes from res_lo
+  };
+
+  reg = vec_perm((__vector signed short)res_hi, (__vector signed short)res_lo,
+                 perm_pack);
+}
+
+inline FP16Vec16::FP16Vec16(const FP32Vec16& v) {
+  // Use bit-manipulation for IEEE FP32 to FP16 conversion since vector
+  // intrinsics for FP32 to FP16 conversion does not use IEEE rounding and can
+  // produce incorrect results for some inputs. Process each of the 4 vectors
+  // separately.
+  __vector unsigned int res_0 = fp32_to_fp16_bits(v.reg.val[0]);
+  __vector unsigned int res_1 = fp32_to_fp16_bits(v.reg.val[1]);
+  __vector unsigned int res_2 = fp32_to_fp16_bits(v.reg.val[2]);
+  __vector unsigned int res_3 = fp32_to_fp16_bits(v.reg.val[3]);
+
+  const __vector unsigned char perm_pack = {
+      2,  3,  6,  7,  10, 11, 14, 15,  // Lower 2 bytes from first vector
+      18, 19, 22, 23, 26, 27, 30, 31   // Lower 2 bytes from second vector
+  };
+
+  reg.val[0] = vec_perm((__vector signed short)res_0,
+                        (__vector signed short)res_1, perm_pack);
+  reg.val[1] = vec_perm((__vector signed short)res_2,
+                        (__vector signed short)res_3, perm_pack);
+}
+
 // 1D softmax over `n` elements in `input`, writes result to `output`.
 // Uses FP32Vec8 for main body, scalar tail handling.
 // Requirement: n > 0

csrc/cpu/dnnl_helper.cpp

@@ -237,12 +237,20 @@ W8A8MatMulPrimitiveHandler::W8A8MatMulPrimitiveHandler(const Args& args)
   };
   dnnl::memory::desc original_b_md({b_k_size_, b_n_size_}, b_type_,
                                    {b_k_stride_, b_n_stride_});
+#ifdef __aarch64__
+  // dummy M size for prepacking weights
+  // Prepacking weights improves performance and avoid runtime reorders
+  constexpr dnnl_dim_t kProbeM = 128;
+#else
+  constexpr dnnl_dim_t kProbeM = DNNL_RUNTIME_DIM_VAL;
+#endif
+
   prepack_weight(args.b_ptr, original_b_md,
                  create_primitive_desc(
-                     MSizeCacheKey{.a_m_size = DNNL_RUNTIME_DIM_VAL,
+                     MSizeCacheKey{.a_m_size = kProbeM,
                                    .use_bias = false,
                                    .bias_type = dnnl::memory::data_type::undef},
-                     true)
+                     /*first_time=*/true)
                      .weights_desc());
   init_runtime_memory_cache(args);
 }

csrc/cpu/mla_decode.cpp

@@ -18,8 +18,8 @@ struct KernelVecType<float> {
 
 template <>
 struct KernelVecType<c10::Half> {
-#if defined(__powerpc64__) || defined(__s390x__)
-  // Power and s390x architecture-specific vector types
+#if defined(__powerpc64__)
+  // Power specific vector types
   using qk_load_vec_type = vec_op::FP32Vec16;
   using qk_vec_type = vec_op::FP32Vec16;
   using v_load_vec_type = vec_op::FP32Vec16;
@@ -38,16 +38,7 @@ struct KernelVecType<c10::BFloat16> {
   using qk_vec_type = vec_op::BF16Vec32;
   using v_load_vec_type = vec_op::BF16Vec16;
 };
-
-#elif defined(__s390x__)
-template <>
-struct KernelVecType<c10::BFloat16> {
-  using qk_load_vec_type = vec_op::BF16Vec16;
-  using qk_vec_type = vec_op::FP32Vec16;
-  using v_load_vec_type = vec_op::BF16Vec16;
-};
-
-#elif defined(__aarch64__)
+#else
 template <>
 struct KernelVecType<c10::BFloat16> {
   using qk_load_vec_type = vec_op::BF16Vec16;

csrc/cpu/torch_bindings.cpp

@@ -119,8 +119,8 @@ void cpu_fused_moe(torch::Tensor& output, const torch::Tensor& input,
                    const std::optional<torch::Tensor>& w13_bias,
                    const std::optional<torch::Tensor>& w2_bias,
                    const torch::Tensor& topk_weights,
-                   const torch::Tensor& topk_id, const std::string& act,
-                   const std::string& isa);
+                   const torch::Tensor& topk_id, const bool skip_weighted,
+                   const std::string& act, const std::string& isa);
 
 TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
   // vLLM custom ops
@@ -320,6 +320,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
   ops.def(
       "cpu_fused_moe(Tensor(a0!) output, Tensor input, Tensor w13, Tensor w2, "
       "Tensor? w13_bias, Tensor? w2_bias, Tensor topk_weights, Tensor topk_id, "
+      "bool skip_weighted, "
       "str act, str isa) -> ()");
   ops.impl("cpu_fused_moe", torch::kCPU, &cpu_fused_moe);
 #endif

csrc/ops.h

@@ -114,6 +114,10 @@ void top_k_per_row_decode(const torch::Tensor& logits, int64_t next_n,
                           int64_t numRows, int64_t stride0, int64_t stride1,
                           int64_t topK);
 
+void large_context_topk(const torch::Tensor& score, torch::Tensor& indices,
+                        const torch::Tensor& lengths,
+                        std::optional<torch::Tensor> row_starts_opt);
+
 void rms_norm_static_fp8_quant(torch::Tensor& out, torch::Tensor& input,
                                torch::Tensor& weight, torch::Tensor& scale,
                                double epsilon);

csrc/rocm/skinny_gemms.cu

@@ -1568,7 +1568,8 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
         {
   #endif
           unsigned int kOff = k + (thrd * A_CHUNK);
-          unsigned int kOffcp = min__(K - A_CHUNK, k_str + kOff);
+          unsigned int kOffcp =
+              k_str + kOff;  // min__(K - A_CHUNK, k_str + kOff);
           for (unsigned int n = 0; n < N; n += CHUNKK * sprdN) {
             __builtin_amdgcn_global_load_lds(
                 (int*)(&A[min__(

csrc/sampler.cu

@@ -725,4 +725,4 @@ void top_k_per_row_prefill(const torch::Tensor& logits,
                      static_cast<int>(stride0), static_cast<int>(stride1),
                      static_cast<int>(topK), kSortingAlgorithmThreshold);
   }
-}
+}

csrc/topk.cu

@@ -0,0 +1,373 @@
+// Portions of this file are adapted from SGLang PR:
+// https://github.com/sgl-project/sglang/pull/11194
+// and
+// https://github.com/sgl-project/sglang/pull/17747
+
+#include "cuda_compat.h"
+#include "dispatch_utils.h"
+
+#include <torch/cuda.h>
+#include <c10/cuda/CUDAGuard.h>
+
+#ifndef USE_ROCM
+  #include <cub/cub.cuh>
+#else
+  #include <hipcub/hipcub.hpp>
+#endif
+
+namespace vllm {
+
+constexpr int TopK = 2048;              // DeepSeek V3 sparse attention top-k
+constexpr int kThreadsPerBlock = 1024;  // Threads per block
+
+// Shared memory budget
+#if defined(USE_ROCM)
+constexpr size_t kSmem = 48 * 1024;  // ROCm default: 48KB
+#else
+// Reduced from 128KB to 32KB to improve occupancy.
+// Each radix pass needs at most ~TopK candidates in the threshold bin,
+// so 4K entries per round (2 rounds = 8K entries = 32KB) is sufficient.
+constexpr size_t kSmem = 8 * 1024 * sizeof(uint32_t);  // 32KB (bytes)
+#endif
+
+struct FastTopKParams {
+  const float* __restrict__ input;         // [batch, seq_len] Logits
+  const int32_t* __restrict__ row_starts;  // [batch] Offset into each row
+                                           // (optional)
+  int32_t* __restrict__ indices;           // [batch, TopK] Output top-k indices
+  int32_t* __restrict__ lengths;           // [batch] Sequence lengths per row
+  int64_t input_stride;                    // Stride between rows
+};
+
+__device__ __forceinline__ auto convert_to_uint32_v2(float x) -> uint32_t {
+  uint32_t bits = __float_as_uint(x);
+  return (bits & 0x80000000u) ? ~bits : (bits | 0x80000000u);
+}
+
+__device__ __forceinline__ auto convert_to_uint8(float x) -> uint8_t {
+  __half h = __float2half_rn(x);
+  uint16_t bits = __half_as_ushort(h);
+  uint16_t key = (bits & 0x8000) ? static_cast<uint16_t>(~bits)
+                                 : static_cast<uint16_t>(bits | 0x8000);
+  return static_cast<uint8_t>(key >> 8);
+}
+
+__device__ void naive_topk_cuda(const float* __restrict__ logits,
+                                int32_t* __restrict__ output_indices,
+                                int32_t seq_len) {
+  const int thread_id = threadIdx.x;
+  for (int i = thread_id; i < TopK; i += kThreadsPerBlock) {
+    output_indices[i] = (i < seq_len) ? i : -1;
+  }
+}
+
+// Adapted from:
+// https://github.com/sgl-project/sglang/blob/v0.5.8/sgl-kernel/csrc/elementwise/topk.cu#L87
+// by: DarkSharpness
+// which at the same time is an optimized topk kernel copied from tilelang
+// kernel
+__device__ void fast_topk_cuda_tl(
+    const float* __restrict__ logits,  // Input logits [seq_len]
+    int* __restrict__ output_indices,  // Output top-k indices [TopK]
+    int logits_offset,                 // Starting offset in logits array
+    int seq_len)                       // Number of valid logits to process
+{
+  constexpr int RADIX = 256;
+  constexpr int MAX_BUFFERED_ITEMS = kSmem / (2 * sizeof(int));
+
+  alignas(128) __shared__ int shared_histogram[2][RADIX + 128];
+  alignas(128) __shared__ int shared_output_count;
+  alignas(128) __shared__ int shared_threshold_bin;
+  alignas(128) __shared__ int shared_buffered_count[2];
+
+  extern __shared__ int buffered_indices[][MAX_BUFFERED_ITEMS];
+
+  const int thread_id = threadIdx.x;
+  int remaining_k = TopK;
+
+  // Pass 0: Build coarse 8-bit histogram using FP16 high bits
+  if (thread_id < RADIX + 1) {
+    shared_histogram[0][thread_id] = 0;
+  }
+  __syncthreads();
+
+  for (int idx = thread_id; idx < seq_len; idx += kThreadsPerBlock) {
+    const auto bin = convert_to_uint8(logits[idx + logits_offset]);
+    ::atomicAdd(&shared_histogram[0][bin], 1);
+  }
+  __syncthreads();
+
+  // Helper: Compute cumulative sum (suffix sum) over histogram using ping-pong
+  // buffers
+  auto compute_cumulative_sum = [&]() {
+    static_assert(1 << 8 == RADIX,
+                  "Radix must be 256 for 8 unrolled iterations");
+#pragma unroll 8
+    for (int i = 0; i < 8; ++i) {
+      if (C10_LIKELY(thread_id < RADIX)) {
+        const int stride = 1 << i;
+        const int src_buffer = i & 1;
+        const int dst_buffer = src_buffer ^ 1;
+
+        int value = shared_histogram[src_buffer][thread_id];
+        if (thread_id < RADIX - stride) {
+          value += shared_histogram[src_buffer][thread_id + stride];
+        }
+        shared_histogram[dst_buffer][thread_id] = value;
+      }
+      __syncthreads();
+    }
+  };
+
+  compute_cumulative_sum();
+
+  // Find threshold bin where cumsum crosses remaining_k
+  if (thread_id < RADIX && shared_histogram[0][thread_id] > remaining_k &&
+      shared_histogram[0][thread_id + 1] <= remaining_k) {
+    shared_threshold_bin = thread_id;
+    shared_buffered_count[0] = 0;
+    shared_output_count = 0;
+  }
+  __syncthreads();
+
+  const int threshold_bin = shared_threshold_bin;
+  remaining_k -= shared_histogram[0][threshold_bin + 1];
+
+  // Early exit if threshold bin perfectly matches remaining_k
+  if (remaining_k == 0) {
+    for (int idx = thread_id; idx < seq_len; idx += kThreadsPerBlock) {
+      const int bin = convert_to_uint8(logits[idx + logits_offset]);
+      if (bin > threshold_bin) {
+        const int output_pos = ::atomicAdd(&shared_output_count, 1);
+        output_indices[output_pos] = idx;
+      }
+    }
+    __syncthreads();
+    return;
+  }
+
+  // Prepare for refinement passes: Process threshold bin
+  __syncthreads();
+  if (thread_id < RADIX + 1) {
+    shared_histogram[0][thread_id] = 0;
+  }
+  __syncthreads();
+
+  // Scan all elements and:
+  // 1. Write indices > threshold_bin to output
+  // 2. Buffer indices == threshold_bin for refinement
+  // 3. Build histogram for next refinement pass (fused optimization)
+  for (int idx = thread_id; idx < seq_len; idx += kThreadsPerBlock) {
+    const float logit_value = logits[idx + logits_offset];
+    const int bin = convert_to_uint8(logit_value);
+
+    if (bin > threshold_bin) {
+      // in top-k, write to output
+      const int output_pos = ::atomicAdd(&shared_output_count, 1);
+      output_indices[output_pos] = idx;
+    } else if (bin == threshold_bin) {
+      // Candidate for top-k, needs refinement
+      const int buffer_pos = ::atomicAdd(&shared_buffered_count[0], 1);
+      if (C10_LIKELY(buffer_pos < MAX_BUFFERED_ITEMS)) {
+        buffered_indices[0][buffer_pos] = idx;
+        // Fused: Build histogram for next pass
+        const uint32_t fp32_bits = convert_to_uint32_v2(logit_value);
+        const int next_bin = (fp32_bits >> 24) & 0xFF;
+        ::atomicAdd(&shared_histogram[0][next_bin], 1);
+      }
+    }
+  }
+  __syncthreads();
+
+  // ============================================================================
+  // Passes 1-4: Refine using 8-bit passes over FP32 bits
+  // ============================================================================
+  // FP32 bits [31:0] split into 4 bytes processed MSB-first:
+  // Pass 1: bits [31:24], Pass 2: bits [23:16], Pass 3: bits [15:8], Pass 4:
+  // bits [7:0]
+#pragma unroll 4
+  for (int pass = 0; pass < 4; ++pass) {
+    __shared__ int shared_final_k;  // For final pass: remaining slots to fill
+    const int src_buffer = pass % 2;
+    const int dst_buffer = src_buffer ^ 1;
+
+    // Clamp buffered count to prevent overflow
+    const int raw_buffered = shared_buffered_count[src_buffer];
+    const int num_buffered =
+        (raw_buffered < MAX_BUFFERED_ITEMS) ? raw_buffered : MAX_BUFFERED_ITEMS;
+
+    compute_cumulative_sum();
+
+    // Find threshold bin for this pass
+    if (thread_id < RADIX && shared_histogram[0][thread_id] > remaining_k &&
+        shared_histogram[0][thread_id + 1] <= remaining_k) {
+      shared_threshold_bin = thread_id;
+      shared_buffered_count[dst_buffer] = 0;
+      shared_final_k = remaining_k - shared_histogram[0][thread_id + 1];
+    }
+    __syncthreads();
+
+    const int threshold_bin = shared_threshold_bin;
+    remaining_k -= shared_histogram[0][threshold_bin + 1];
+
+    // Bit offset for this pass: 24, 16, 8, 0
+    const int bit_offset = 24 - pass * 8;
+
+    // Early exit if threshold bin perfectly matches
+    if (remaining_k == 0) {
+      for (int i = thread_id; i < num_buffered; i += kThreadsPerBlock) {
+        const int idx = buffered_indices[src_buffer][i];
+        const uint32_t fp32_bits =
+            convert_to_uint32_v2(logits[idx + logits_offset]);
+        const int bin = (fp32_bits >> bit_offset) & 0xFF;
+        if (bin > threshold_bin) {
+          const int output_pos = ::atomicAdd(&shared_output_count, 1);
+          output_indices[output_pos] = idx;
+        }
+      }
+      __syncthreads();
+      break;
+    }
+
+    // Continue refinement
+    __syncthreads();
+    if (thread_id < RADIX + 1) {
+      shared_histogram[0][thread_id] = 0;
+    }
+    __syncthreads();
+
+    for (int i = thread_id; i < num_buffered; i += kThreadsPerBlock) {
+      const int idx = buffered_indices[src_buffer][i];
+      const float logit_value = logits[idx + logits_offset];
+      const uint32_t fp32_bits = convert_to_uint32_v2(logit_value);
+      const int bin = (fp32_bits >> bit_offset) & 0xFF;
+
+      if (bin > threshold_bin) {
+        // Definitely in top-k
+        const int output_pos = ::atomicAdd(&shared_output_count, 1);
+        output_indices[output_pos] = idx;
+      } else if (bin == threshold_bin) {
+        if (pass == 3) {
+          // Final pass (bits [7:0]): No more refinement possible
+          // Fill remaining slots in reverse order to maintain descending order
+          const int slot = ::atomicAdd(&shared_final_k, -1);
+          if (slot > 0) {
+            output_indices[TopK - slot] = idx;
+          }
+        } else {
+          // Buffer for next pass and build next histogram
+          const int buffer_pos =
+              ::atomicAdd(&shared_buffered_count[dst_buffer], 1);
+          if (C10_LIKELY(buffer_pos < MAX_BUFFERED_ITEMS)) {
+            buffered_indices[dst_buffer][buffer_pos] = idx;
+            // Fused: Build histogram for next pass
+            const int next_bit_offset = bit_offset - 8;
+            const int next_bin = (fp32_bits >> next_bit_offset) & 0xFF;
+            ::atomicAdd(&shared_histogram[0][next_bin], 1);
+          }
+        }
+      }
+    }
+    __syncthreads();
+  }
+}
+
+__global__ __launch_bounds__(kThreadsPerBlock) void topk_kernel(
+    const FastTopKParams params) {
+  const auto& [input, row_starts, indices, lengths, input_stride] = params;
+  const uint64_t batch_idx = blockIdx.x;
+  const int logits_offset = row_starts == nullptr ? 0 : row_starts[batch_idx];
+  const int seq_len = lengths[batch_idx];
+  int* output_indices = indices + batch_idx * TopK;
+  const float* logits = input + batch_idx * input_stride;
+
+  if (seq_len <= TopK) {
+    // Shortcut: All elements are in top-k
+    return naive_topk_cuda(logits, output_indices, seq_len);
+  } else {
+    return fast_topk_cuda_tl(logits, output_indices, logits_offset, seq_len);
+  }
+}
+
+FastTopKParams get_params(
+    const at::Tensor& score, const at::Tensor& lengths,
+    std::optional<at::Tensor> row_starts_opt = std::nullopt,
+    std::optional<at::Tensor> indices_opt = std::nullopt) {
+  const int64_t batch_size = score.size(0);
+
+  TORCH_CHECK(score.dim() == 2 && score.stride(1) == 1,
+              "score must be 2D with contiguous rows");
+  TORCH_CHECK(lengths.dim() == 1 && lengths.is_contiguous() &&
+                  lengths.size(0) == batch_size,
+              "lengths must be 1D contiguous with size matching batch");
+
+  const int32_t* row_starts_ptr = nullptr;
+  if (row_starts_opt.has_value()) {
+    const auto& row_starts = *row_starts_opt;
+    TORCH_CHECK(row_starts.dim() == 1 && row_starts.size(0) == batch_size,
+                "row_starts must be 1D with size matching batch");
+    row_starts_ptr = row_starts.data_ptr<int32_t>();
+  }
+
+  int32_t* indices_ptr = nullptr;
+  if (indices_opt.has_value()) {
+    const auto& indices = *indices_opt;
+    TORCH_CHECK(indices.dim() == 2 && indices.is_contiguous() &&
+                    indices.size(0) == batch_size && indices.size(1) == TopK,
+                "indices must be 2D contiguous [batch, TopK]");
+    indices_ptr = indices.data_ptr<int32_t>();
+  }
+
+  return FastTopKParams{
+      .input = score.data_ptr<float>(),
+      .row_starts = row_starts_ptr,
+      .indices = indices_ptr,
+      .lengths = lengths.data_ptr<int32_t>(),
+      .input_stride = score.stride(0),
+  };
+}
+
+template <auto* kernel_func, size_t smem_bytes>
+void setup_kernel_smem_once() {
+  static const cudaError_t result = []() -> cudaError_t {
+#ifdef USE_ROCM
+    auto func_ptr = reinterpret_cast<const void*>(kernel_func);
+#else
+    auto func_ptr = kernel_func;
+#endif
+    return cudaFuncSetAttribute(
+        func_ptr, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_bytes);
+  }();
+
+  TORCH_CHECK(
+      result == cudaSuccess,
+      "Failed to set kernel shared memory limit: ", cudaGetErrorString(result));
+}
+
+}  // namespace vllm
+
+void large_context_topk(
+    const torch::Tensor& logits, torch::Tensor& indices,
+    const torch::Tensor& seq_lens,
+    c10::optional<torch::Tensor> row_starts = c10::nullopt) {
+  TORCH_CHECK(logits.is_cuda(), "logits must be a CUDA tensor");
+  TORCH_CHECK(indices.is_cuda(), "indices must be a CUDA tensor");
+  TORCH_CHECK(seq_lens.is_cuda(), "seq_lens must be a CUDA tensor");
+  if (row_starts.has_value()) {
+    TORCH_CHECK(row_starts->is_cuda(), "row_starts must be a CUDA tensor");
+  }
+
+  const auto params = vllm::get_params(logits, seq_lens, row_starts, indices);
+  const int64_t batch_size = logits.size(0);
+
+  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+  const dim3 grid(static_cast<uint32_t>(batch_size));
+  const dim3 block(vllm::kThreadsPerBlock);
+
+  vllm::setup_kernel_smem_once<vllm::topk_kernel, vllm::kSmem>();
+  vllm::topk_kernel<<<grid, block, vllm::kSmem, stream>>>(params);
+
+  const cudaError_t result = cudaGetLastError();
+  TORCH_CHECK(result == cudaSuccess,
+              "large_context_topk kernel failed: ", cudaGetErrorString(result));
+}

csrc/torch_bindings.cpp

@@ -190,6 +190,12 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
       "int numRows, int stride0, int stride1, int topK) -> ()");
   ops.impl("top_k_per_row_decode", torch::kCUDA, &top_k_per_row_decode);
 
+  ops.def(
+      "large_context_topk(Tensor score, Tensor indices, Tensor lengths, "
+      "Tensor? "
+      "row_starts_opt) -> ()");
+  ops.impl("large_context_topk", torch::kCUDA, &large_context_topk);
+
   // Layernorm-quant
   // Apply Root Mean Square (RMS) Normalization to the input tensor.
   ops.def(

docker/Dockerfile

@@ -506,6 +506,7 @@ RUN apt-get update -y \
         curl \
         sudo \
         python3-pip \
+        git \
         ffmpeg \
         libsm6 \
         libxext6 \

docker/Dockerfile.rocm

@@ -1,7 +1,7 @@
 # default base image
 ARG REMOTE_VLLM="0"
 ARG COMMON_WORKDIR=/app
-ARG BASE_IMAGE=rocm/vllm-dev:base_custom_releases_rocm_v0.16.0_20260211
+ARG BASE_IMAGE=rocm/vllm-dev:base
 
 # Sccache configuration (only used in release pipeline)
 ARG USE_SCCACHE

docker/Dockerfile.rocm_base

@@ -1,5 +1,5 @@
-ARG BASE_IMAGE=rocm/dev-ubuntu-22.04:7.2-complete
-ARG TRITON_BRANCH="57c693b6"
+ARG BASE_IMAGE=rocm/dev-ubuntu-22.04:7.0-complete
+ARG TRITON_BRANCH="f332c492"
 ARG TRITON_REPO="https://github.com/ROCm/triton.git"
 ARG PYTORCH_BRANCH="89075173"
 ARG PYTORCH_REPO="https://github.com/ROCm/pytorch.git"
@@ -9,7 +9,7 @@ ARG PYTORCH_AUDIO_BRANCH="v2.9.0"
 ARG PYTORCH_AUDIO_REPO="https://github.com/pytorch/audio.git"
 ARG FA_BRANCH="0e60e394"
 ARG FA_REPO="https://github.com/Dao-AILab/flash-attention.git"
-ARG AITER_BRANCH="v0.1.10.post2"
+ARG AITER_BRANCH="6af8b687"
 ARG AITER_REPO="https://github.com/ROCm/aiter.git"
 ARG MORI_BRANCH="2d02c6a9"
 ARG MORI_REPO="https://github.com/ROCm/mori.git"
@@ -142,7 +142,6 @@ ARG PYTORCH_VISION_REPO
 ARG PYTORCH_AUDIO_REPO
 ARG USE_SCCACHE
 
-RUN apt-get update && apt-get install -y pkg-config liblzma-dev
 RUN git clone ${PYTORCH_REPO} pytorch
 RUN cd pytorch && git checkout ${PYTORCH_BRANCH} \
     && pip install -r requirements.txt && git submodule update --init --recursive \
@@ -240,7 +239,7 @@ RUN pip install pyyaml && cd aiter \
            export HIP_CLANG_PATH=/opt/sccache-wrappers \
            && sccache --show-stats; \
        fi \
-    && GPU_ARCHS=${AITER_ROCM_ARCH} python3 setup.py bdist_wheel --dist-dir=dist \
+    && PREBUILD_KERNELS=1 GPU_ARCHS=${AITER_ROCM_ARCH} python3 setup.py bdist_wheel --dist-dir=dist \
     && if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi \
     && ls /app/aiter/dist/*.whl
 RUN mkdir -p /app/install && cp /app/aiter/dist/*.whl /app/install

docs/benchmarking/cli.md

@@ -30,6 +30,7 @@ th {
 | HuggingFace-Other | ✅ | ✅ | `lmms-lab/LLaVA-OneVision-Data`, `Aeala/ShareGPT_Vicuna_unfiltered` |
 | HuggingFace-MTBench | ✅ | ✅ | `philschmid/mt-bench` |
 | HuggingFace-Blazedit | ✅ | ✅ | `vdaita/edit_5k_char`, `vdaita/edit_10k_char` |
+| HuggingFace-ASR | ✅ | ✅ | `openslr/librispeech_asr`, `facebook/voxpopuli`,  `LIUM/tedlium`, `edinburghcstr/ami`,        `speechcolab/gigaspeech`,        `kensho/spgispeech` |
 | Spec Bench | ✅ | ✅ | `wget https://raw.githubusercontent.com/hemingkx/Spec-Bench/refs/heads/main/data/spec_bench/question.jsonl` |
 | Custom | ✅ | ✅ | Local file: `data.jsonl` |
 | Custom MM | ✅ | ✅ | Local file: `mm_data.jsonl` |
@@ -299,6 +300,22 @@ vllm bench serve \
     --blazedit-max-distance 0.99
 ```
 
+`openslr/librispeech_asr`, `facebook/voxpopuli`, `LIUM/tedlium`, `edinburghcstr/ami`, `speechcolab/gigaspeech`, `kensho/spgispeech`
+
+```bash
+vllm bench serve \
+    --model openai/whisper-large-v3-turbo \
+    --backend openai-audio \
+    --dataset-name hf \
+    --dataset-path facebook/voxpopuli --hf-subset en --hf-split test --no-stream --trust-remote-code \
+    --num-prompts 99999999 \
+    --no-oversample \
+    --endpoint /v1/audio/transcriptions \
+    --ready-check-timeout-sec 600 \
+    --save-result \
+    --max-concurrency 512
+```
+
 #### Running With Sampling Parameters
 
 When using OpenAI-compatible backends such as `vllm`, optional sampling

docs/design/attention_backends.md

@@ -152,6 +152,7 @@ Priority is **1 = highest** (tried first).
 | **Sink** | Attention sink support (for StreamingLLM) |
 | **Sparse** | Sparse attention support (MLA only) |
 | **MM Prefix** | Multimodal prefix full attention support |
+| **DCP** | Decode Context Parallelism support (`--decode-context-parallel-size`) |
 | **Attention Types** | Supported attention patterns (Decoder, Encoder, Enc-Dec) |
 | **Compute Cap.** | Required CUDA compute capability (N/A for non-CUDA backends) |
 
@@ -159,20 +160,20 @@ Priority is **1 = highest** (tried first).
 
 ## Standard Attention (MHA, MQA, GQA) Backends
 
-| Backend | Version | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | MM Prefix | Attention Types | Compute Cap. |
-|---------|---------|--------|-----------|-------------|------------|------|-----------|-----------------|--------------|
-| `CPU_ATTN` |  | fp16, bf16, fp32 | `auto` | Any | 32, 64, 80, 96, 112, 128, 160, 192, 224, 256 | ❌ | ❌ | All | N/A |
-| `FLASHINFER` | Native† | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ❌ | ❌ | Decoder | 7.x-9.x |
-| `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ✅ | ❌ | Decoder | 10.x |
-| `FLASH_ATTN` | FA2* | fp16, bf16 | `auto`, `bfloat16` | %16 | Any | ❌ | ❌ | All | ≥8.0 |
-| `FLASH_ATTN` | FA3* | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ❌ | All | 9.x |
-| `FLASH_ATTN_DIFFKV` |  | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | Decoder | Any |
-| `FLEX_ATTENTION` |  | fp16, bf16, fp32 | `auto`, `bfloat16` | Any | Any | ❌ | ✅ | Decoder, Encoder Only | Any |
-| `ROCM_AITER_FA` |  | fp16, bf16 | `auto` | 16, 32 | 64, 128, 256 | ❌ | ❌ | Decoder | N/A |
-| `ROCM_AITER_UNIFIED_ATTN` |  | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | Decoder | N/A |
-| `ROCM_ATTN` |  | fp16, bf16, fp32 | `auto` | 16, 32, 544 | 32, 64, 96, 128, 160, 192, 224, 256 | ❌ | ❌ | Decoder | N/A |
-| `TREE_ATTN` |  | fp16, bf16 | `auto` | %16 | 32, 64, 96, 128, 160, 192, 224, 256 | ❌ | ❌ | Decoder | Any |
-| `TRITON_ATTN` |  | fp16, bf16, fp32 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ✅ | All | Any |
+| Backend | Version | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | MM Prefix | DCP | Attention Types | Compute Cap. |
+|---------|---------|--------|-----------|-------------|------------|------|-----------|-----|-----------------|--------------|
+| `CPU_ATTN` |  | fp16, bf16, fp32 | `auto` | Any | 32, 64, 80, 96, 112, 128, 160, 192, 224, 256 | ❌ | ❌ | ❌ | All | N/A |
+| `FLASHINFER` | Native† | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ❌ | ❌ | ✅ | Decoder | 7.x-9.x |
+| `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ✅ | ❌ | ✅ | Decoder | 10.x |
+| `FLASH_ATTN` | FA2* | fp16, bf16 | `auto`, `bfloat16` | %16 | Any | ❌ | ❌ | ✅ | All | ≥8.0 |
+| `FLASH_ATTN` | FA3* | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ❌ | ✅ | All | 9.x |
+| `FLASH_ATTN_DIFFKV` |  | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ✅ | Decoder | Any |
+| `FLEX_ATTENTION` |  | fp16, bf16, fp32 | `auto`, `bfloat16` | Any | Any | ❌ | ✅ | ❌ | Decoder, Encoder Only | Any |
+| `ROCM_AITER_FA` |  | fp16, bf16 | `auto` | 16, 32 | 64, 128, 256 | ❌ | ❌ | ❌ | Decoder | N/A |
+| `ROCM_AITER_UNIFIED_ATTN` |  | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | Decoder | N/A |
+| `ROCM_ATTN` |  | fp16, bf16, fp32 | `auto` | 16, 32, 544 | 32, 64, 96, 128, 160, 192, 224, 256 | ❌ | ❌ | ❌ | Decoder | N/A |
+| `TREE_ATTN` |  | fp16, bf16 | `auto` | %16 | 32, 64, 96, 128, 160, 192, 224, 256 | ❌ | ❌ | ❌ | Decoder | Any |
+| `TRITON_ATTN` |  | fp16, bf16, fp32 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ✅ | ❌ | All | Any |
 
 > **†** FlashInfer uses TRTLLM attention on Blackwell (SM100), which supports sinks. Disable via `--attention-config.use_trtllm_attention=0`.
 >
@@ -199,14 +200,14 @@ configuration.
 
 ### Decode Backends
 
-| Backend | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Sparse | MM Prefix | Attention Types | Compute Cap. |
-|---------|--------|-----------|-------------|------------|------|--------|-----------|-----------------|--------------|
-| `CUTLASS_MLA` | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | 128 | Any | ❌ | ❌ | ❌ | Decoder | 10.x |
-| `FLASHINFER_MLA` | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | Decoder | 10.x |
-| `FLASHMLA` | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | 64 | Any | ❌ | ❌ | ❌ | Decoder | 9.x-10.x |
-| `FLASHMLA_SPARSE` | bf16 | `auto`, `bfloat16`, `fp8_ds_mla` | 64 | 576 | ❌ | ✅ | ❌ | Decoder | 9.x-10.x |
-| `FLASH_ATTN_MLA` | fp16, bf16 | `auto`, `bfloat16` | %16 | Any | ❌ | ❌ | ❌ | Decoder | 9.x |
-| `ROCM_AITER_MLA` | fp16, bf16 | `auto` | 1 | Any | ❌ | ❌ | ❌ | Decoder | N/A |
-| `ROCM_AITER_MLA_SPARSE` | fp16, bf16 | `auto` | Any | 576 | ❌ | ❌ | ❌ | Decoder | N/A |
-| `ROCM_AITER_TRITON_MLA` | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | Decoder | N/A |
-| `TRITON_MLA` | fp16, bf16 | `auto`, `bfloat16` | Any | Any | ❌ | ❌ | ❌ | Decoder | Any |
+| Backend | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Sparse | MM Prefix | DCP | Attention Types | Compute Cap. |
+|---------|--------|-----------|-------------|------------|------|--------|-----------|-----|-----------------|--------------|
+| `CUTLASS_MLA` | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | 128 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 10.x |
+| `FLASHINFER_MLA` | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | 10.x |
+| `FLASHMLA` | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | 64 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x-10.x |
+| `FLASHMLA_SPARSE` | bf16 | `auto`, `bfloat16`, `fp8_ds_mla` | 64 | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x-10.x |
+| `FLASH_ATTN_MLA` | fp16, bf16 | `auto`, `bfloat16` | %16 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x |
+| `ROCM_AITER_MLA` | fp16, bf16 | `auto` | 1 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
+| `ROCM_AITER_MLA_SPARSE` | fp16, bf16 | `auto` | Any | 576 | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
+| `ROCM_AITER_TRITON_MLA` | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
+| `TRITON_MLA` | fp16, bf16 | `auto`, `bfloat16` | Any | Any | ❌ | ❌ | ❌ | ✅ | Decoder | Any |

docs/design/io_processor_plugins.md

@@ -14,8 +14,26 @@ IOProcessorOutput = TypeVar("IOProcessorOutput")
 
 class IOProcessor(ABC, Generic[IOProcessorInput, IOProcessorOutput]):
     def __init__(self, vllm_config: VllmConfig):
+        super().__init__()
+
         self.vllm_config = vllm_config
 
+    @abstractmethod
+    def parse_data(self, data: object) -> IOProcessorInput:
+        raise NotImplementedError
+
+    def merge_sampling_params(
+        self,
+        params: SamplingParams | None = None,
+    ) -> SamplingParams:
+        return params or SamplingParams()
+
+    def merge_pooling_params(
+        self,
+        params: PoolingParams | None = None,
+    ) -> PoolingParams:
+        return params or PoolingParams()
+
     @abstractmethod
     def pre_process(
         self,
@@ -55,29 +73,13 @@ class IOProcessor(ABC, Generic[IOProcessorInput, IOProcessorOutput]):
             [(i, item) async for i, item in model_output], key=lambda output: output[0]
         )
         collected_output = [output[1] for output in sorted_output]
-        return self.post_process(collected_output, request_id, **kwargs)
-
-    @abstractmethod
-    def parse_request(self, request: Any) -> IOProcessorInput:
-        raise NotImplementedError
-
-    def validate_or_generate_params(
-        self, params: SamplingParams | PoolingParams | None = None
-    ) -> SamplingParams | PoolingParams:
-        return params or PoolingParams()
-
-    @abstractmethod
-    def output_to_response(
-        self, plugin_output: IOProcessorOutput
-    ) -> IOProcessorResponse:
-        raise NotImplementedError
+        return self.post_process(collected_output, request_id=request_id, **kwargs)
 ```
 
-The `parse_request` method is used for validating the user prompt and converting it into the input expected by the `pre_process`/`pre_process_async` methods.
+The `parse_data` method is used for validating the user data and converting it into the input expected by the `pre_process*` methods.
+The `merge_sampling_params` and `merge_pooling_params` methods merge input `SamplingParams` or `PoolingParams` (if any) with the default one.
 The `pre_process*` methods take the validated plugin input to generate vLLM's model prompts for regular inference.
 The `post_process*` methods take `PoolingRequestOutput` objects as input and generate a custom plugin output.
-The `validate_or_generate_params` method is used for validating with the plugin any `SamplingParameters`/`PoolingParameters` received with the user request, or to generate new ones if none are specified. The function always returns the validated/generated parameters.
-The `output_to_response` method is used only for online serving and converts the plugin output to the `IOProcessorResponse` type that is then returned by the API Server. The implementation of the `/pooling` serving endpoint is available here [vllm/entrypoints/openai/serving_pooling.py](../../vllm/entrypoints/pooling/pooling/serving.py).
 
 An example implementation of a plugin that enables generating geotiff images with the PrithviGeospatialMAE model is available [here](https://github.com/IBM/terratorch/tree/main/terratorch/vllm/plugins/segmentation). Please, also refer to our online ([examples/pooling/plugin/prithvi_geospatial_mae_online.py](../../examples/pooling/plugin/prithvi_geospatial_mae_online.py)) and offline ([examples/pooling/plugin/prithvi_geospatial_mae_io_processor.py](../../examples/pooling/plugin/prithvi_geospatial_mae_io_processor.py)) inference examples.
 

docs/design/moe_kernel_features.md

@@ -32,7 +32,7 @@ th {
 
 | Backend | Output act. format | Quant. types | Quant. format | Async | Apply Weight On Input | Subclass |
 |---------|--------------------|--------------|---------------|-------|-----------------------|-----------|
-| naive | standard | all<sup>1</sup> | G,A,T | N | <sup>6</sup> | [layer.py][vllm.model_executor.layers.fused_moe.layer.FusedMoE.forward_impl] |
+| naive | standard | all<sup>1</sup> | G,A,T | N | <sup>6</sup> | [layer.py][vllm.model_executor.layers.fused_moe.layer.FusedMoE |
 | pplx | batched | fp8,int8 | G,A,T | Y | Y | [`PplxPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.pplx_prepare_finalize.PplxPrepareAndFinalize] |
 | deepep_high_throughput | standard | fp8 | G(128),A,T<sup>2</sup> | Y | Y | [`DeepEPLLPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.deepep_ll_prepare_finalize.DeepEPLLPrepareAndFinalize] |
 | deepep_low_latency | batched | fp8 | G(128),A,T<sup>3</sup> | Y | Y | [`DeepEPHTPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.deepep_ht_prepare_finalize.DeepEPHTPrepareAndFinalize] |

docs/features/multimodal_inputs.md

@@ -521,7 +521,7 @@ First, launch the OpenAI-compatible server:
 
 ```bash
 vllm serve microsoft/Phi-3.5-vision-instruct --runner generate \
-  --trust-remote-code --max-model-len 4096 --limit-mm-per-prompt '{"image":2}'
+  --trust-remote-code --max-model-len 4096 --limit-mm-per-prompt.image 2
 ```
 
 Then, you can use the OpenAI client as follows:

docs/features/quantization/README.md

@@ -48,7 +48,7 @@ th:not(:first-child) {
 |-----------------------|---------|----------|----------|-------|----------|-----------|-------------|-----------|
 | AWQ                   | ❌      | ✅︎       | ✅︎       | ✅︎    | ✅︎       | ❌         | ✅︎          | ✅︎        |
 | GPTQ                  | ✅︎      | ✅︎       | ✅︎       | ✅︎    | ✅︎       | ❌         | ✅︎          | ✅︎        |
-| Marlin (GPTQ/AWQ/FP8) | ❌      | ❌       | ✅︎       | ✅︎    | ✅︎       | ❌         | ❌          | ❌        |
+| Marlin (GPTQ/AWQ/FP8/FP4) | ❌      | ✅︎*       | ✅︎       | ✅︎    | ✅︎       | ❌         | ❌          | ❌        |
 | INT8 (W8A8)           | ❌      | ✅︎       | ✅︎       | ✅︎    | ✅︎       | ❌         | ❌          | ✅︎        |
 | FP8 (W8A8)            | ❌      | ❌       | ❌       | ✅︎    | ✅︎       | ✅︎         | ❌          | ❌        |
 | bitsandbytes          | ✅︎      | ✅︎       | ✅︎       | ✅︎    | ✅︎       | ❌         | ❌          | ❌        |
@@ -59,6 +59,7 @@ th:not(:first-child) {
 - ✅︎ indicates that the quantization method is supported on the specified hardware.
 - ❌ indicates that the quantization method is not supported on the specified hardware.
 - All Intel Gaudi quantization support has been migrated to [vLLM-Gaudi](https://github.com/vllm-project/vllm-gaudi).
+- *Turing does not support Marlin MXFP4.
 
 !!! note
     For information on quantization support on Google TPU, please refer to the [TPU-Inference Recommended Models and Features](https://docs.vllm.ai/projects/tpu/en/latest/recommended_models_features/) documentation.

docs/features/quantization/fp8.md

@@ -2,7 +2,7 @@
 
 vLLM supports FP8 (8-bit floating point) weight and activation quantization using hardware acceleration on GPUs such as Nvidia H100 and AMD MI300x.
 Currently, only Hopper and Ada Lovelace GPUs are officially supported for W8A8.
-Ampere GPUs are supported for W8A16 (weight-only FP8) utilizing Marlin kernels.
+Turing/Ampere GPUs are supported for W8A16 (weight-only FP8) utilizing Marlin kernels.
 Quantization of models with FP8 allows for a 2x reduction in model memory requirements and up to a 1.6x improvement in throughput with minimal impact on accuracy.
 
 Please visit the HF collection of [quantized FP8 checkpoints of popular LLMs ready to use with vLLM](https://huggingface.co/collections/neuralmagic/fp8-llms-for-vllm-666742ed2b78b7ac8df13127).
@@ -13,8 +13,8 @@ The FP8 types typically supported in hardware have two distinct representations,
 - **E5M2**: Consists of 1 sign bit, 5 exponent bits, and 2 bits of mantissa. It can store values up to +/-57344, +/- `inf`, and `nan`. The tradeoff for the increased dynamic range is lower precision of the stored values.
 
 !!! note
-    FP8 computation is supported on NVIDIA GPUs with compute capability > 8.9 (Ada Lovelace, Hopper).
-    FP8 models will run on compute capability > 8.0 (Ampere) as weight-only W8A16, utilizing FP8 Marlin.
+    FP8 computation is supported on NVIDIA GPUs with compute capability >= 8.9 (Ada Lovelace, Hopper).
+    FP8 models will run on compute capability >= 7.5 (Turing) as weight-only W8A16, utilizing FP8 Marlin.
 
 ## Installation
 

docs/models/supported_models.md

@@ -658,7 +658,7 @@ On the other hand, modalities separated by `/` are mutually exclusive.
 See [this page](../features/multimodal_inputs.md) on how to pass multi-modal inputs to the model.
 
 !!! tip
-    For hybrid-only models such as Llama-4, Step3 and Mistral-3, a text-only mode can be enabled by setting all supported multimodal modalities to 0 (e.g, `--limit-mm-per-prompt '{"image":0}`) so that their multimodal modules will not be loaded to free up more GPU memory for KV cache.
+    For hybrid-only models such as Llama-4, Step3, Mistral-3 and Qwen-3.5, a text-only mode can be enabled by setting all supported multimodal modalities to 0 (`--language-model-only`) so that their multimodal modules will not be loaded to free up more GPU memory for KV cache.
 
 !!! note
     vLLM currently supports adding LoRA adapters to the language backbone for most multimodal models. Additionally, vLLM now experimentally supports adding LoRA to the tower and connector modules for some multimodal models. See [this page](../features/lora.md).
@@ -738,6 +738,8 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
 | `Qwen2VLForConditionalGeneration` | QVQ, Qwen2-VL | T + I<sup>E+</sup> + V<sup>E+</sup> | `Qwen/QVQ-72B-Preview`, `Qwen/Qwen2-VL-7B-Instruct`, `Qwen/Qwen2-VL-72B-Instruct`, etc. | ✅︎ | ✅︎ |
 | `Qwen2_5_VLForConditionalGeneration` | Qwen2.5-VL | T + I<sup>E+</sup> + V<sup>E+</sup> | `Qwen/Qwen2.5-VL-3B-Instruct`, `Qwen/Qwen2.5-VL-72B-Instruct`, etc. | ✅︎ | ✅︎ |
 | `Qwen2_5OmniThinkerForConditionalGeneration` | Qwen2.5-Omni | T + I<sup>E+</sup> + V<sup>E+</sup> + A<sup>+</sup> | `Qwen/Qwen2.5-Omni-3B`, `Qwen/Qwen2.5-Omni-7B` | ✅︎ | ✅︎ |
+| `Qwen3_5ForConditionalGeneration` | Qwen3.5 | T + I<sup>E+</sup> + V<sup>E+</sup> | `Qwen/Qwen3.5-9B-Instruct`, etc. | ✅︎ | ✅︎ |
+| `Qwen3_5MoeForConditionalGeneration` | Qwen3.5-MOE | T + I<sup>E+</sup> + V<sup>E+</sup> | `Qwen/Qwen3.5-35B-A3B-Instruct`, etc. | ✅︎ | ✅︎ |
 | `Qwen3VLForConditionalGeneration` | Qwen3-VL | T + I<sup>E+</sup> + V<sup>E+</sup> | `Qwen/Qwen3-VL-4B-Instruct`, etc. | ✅︎ | ✅︎ |
 | `Qwen3VLMoeForConditionalGeneration` | Qwen3-VL-MOE | T + I<sup>E+</sup> + V<sup>E+</sup> | `Qwen/Qwen3-VL-30B-A3B-Instruct`, etc. | ✅︎ | ✅︎ |
 | `Qwen3OmniMoeThinkerForConditionalGeneration` | Qwen3-Omni | T + I<sup>E+</sup> + V<sup>E+</sup> + A<sup>+</sup> | `Qwen/Qwen3-Omni-30B-A3B-Instruct`, `Qwen/Qwen3-Omni-30B-A3B-Thinking` | ✅︎ | ✅︎ |
@@ -788,6 +790,7 @@ Speech2Text models trained specifically for Automatic Speech Recognition.
 
 | Architecture | Models | Example HF Models | [LoRA](../features/lora.md) | [PP](../serving/parallelism_scaling.md) |
 |--------------|--------|-------------------|----------------------|---------------------------|
+| `FunASRForConditionalGeneration` | FunASR | `allendou/Fun-ASR-Nano-2512-vllm`, etc. | | |
 | `Gemma3nForConditionalGeneration` | Gemma3n | `google/gemma-3n-E2B-it`, `google/gemma-3n-E4B-it`, etc. | | |
 | `GlmAsrForConditionalGeneration` | GLM-ASR | `zai-org/GLM-ASR-Nano-2512` | ✅︎ | ✅︎ |
 | `GraniteSpeechForConditionalGeneration` | Granite Speech | `ibm-granite/granite-speech-3.3-2b`, `ibm-granite/granite-speech-3.3-8b`, etc. | ✅︎ | ✅︎ |

examples/offline_inference/kv_load_failure_recovery/README.md

@@ -28,3 +28,4 @@ It demonstrates vLLM's ability to recover from KV load failures in both synchron
 
 ```bash
 ./run.sh
+```

examples/offline_inference/mistral-small.py

@@ -18,11 +18,11 @@ from vllm.assets.image import ImageAsset
 # # Mistral format
 # vllm serve mistralai/Mistral-Small-3.1-24B-Instruct-2503 \
 #   --tokenizer-mode mistral --config-format mistral --load-format mistral \
-#   --limit-mm-per-prompt '{"image":4}' --max-model-len 16384
+#   --limit-mm-per-prompt.image 4 --max-model-len 16384
 #
 # # HF format
 # vllm serve mistralai/Mistral-Small-3.1-24B-Instruct-2503 \
-#   --limit-mm-per-prompt '{"image":4}' --max-model-len 16384
+#   --limit-mm-per-prompt.image 4 --max-model-len 16384
 # ```
 #
 # - Client:

examples/offline_inference/run_one_batch.py

@@ -0,0 +1,112 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+from vllm import LLM, EngineArgs
+from vllm.config import ProfilerConfig
+from vllm.utils.argparse_utils import FlexibleArgumentParser
+
+DEFAULT_MAX_TOKENS = 16
+
+
+def create_parser() -> FlexibleArgumentParser:
+    parser = FlexibleArgumentParser()
+    EngineArgs.add_cli_args(parser)
+    parser.set_defaults(model="meta-llama/Llama-3.2-1B-Instruct")
+
+    batch_group = parser.add_argument_group("Batch parameters")
+    batch_group.add_argument("--batch-size", type=int, default=1)
+    batch_group.add_argument("--prompt-size", type=int, default=128)
+    batch_group.add_argument("--prompt-prefix", type=str, default="Hello, my name is")
+
+    profile_group = parser.add_argument_group("Profiling parameters")
+    profile_group.add_argument(
+        "--profile",
+        choices=["none", "prefill", "decode", "both"],
+        default="none",
+    )
+    profile_group.add_argument(
+        "--profile-dir",
+        type=str,
+        default="",
+        help="Required when --profile is not 'none'.",
+    )
+
+    return parser
+
+
+def _build_prompt(prefix: str, prompt_size: int) -> str:
+    if prompt_size <= 0:
+        return ""
+    if not prefix:
+        prefix = " "
+    if len(prefix) >= prompt_size:
+        return prefix[:prompt_size]
+    repeat_count = (prompt_size + len(prefix) - 1) // len(prefix)
+    return (prefix * repeat_count)[:prompt_size]
+
+
+def _build_profiler_config(
+    profile: str, profile_dir: str, max_tokens: int
+) -> ProfilerConfig | None:
+    if profile == "none":
+        return None
+    if not profile_dir:
+        raise ValueError("--profile-dir must be set when profiling is enabled.")
+    if profile == "prefill":
+        delay_iterations = 0
+        max_iterations = 1
+    elif profile == "decode":
+        delay_iterations = 1
+        max_iterations = max(1, max_tokens)
+    else:
+        delay_iterations = 0
+        max_iterations = 0
+
+    return ProfilerConfig(
+        profiler="torch",
+        torch_profiler_dir=profile_dir,
+        delay_iterations=delay_iterations,
+        max_iterations=max_iterations,
+    )
+
+
+def main(args: dict) -> None:
+    max_tokens = DEFAULT_MAX_TOKENS
+    batch_size = args.pop("batch_size")
+    prompt_size = args.pop("prompt_size")
+    prompt_prefix = args.pop("prompt_prefix")
+    profile = args.pop("profile")
+    profile_dir = args.pop("profile_dir")
+
+    profiler_config = _build_profiler_config(profile, profile_dir, max_tokens)
+    if profiler_config is not None:
+        args["profiler_config"] = profiler_config
+
+    llm = LLM(**args)
+
+    sampling_params = llm.get_default_sampling_params()
+    sampling_params.max_tokens = max_tokens
+    sampling_params.min_tokens = max_tokens
+    sampling_params.ignore_eos = True
+
+    prompt = _build_prompt(prompt_prefix, prompt_size)
+    prompts = [prompt] * batch_size
+
+    if profile != "none":
+        llm.start_profile()
+    outputs = llm.generate(prompts, sampling_params)
+    if profile != "none":
+        llm.stop_profile()
+
+    print("-" * 50)
+    for output in outputs:
+        generated_text = output.outputs[0].text
+        print(f"Prompt: {output.prompt!r}\nGenerated text: {generated_text!r}")
+        print("-" * 50)
+
+
+if __name__ == "__main__":
+    parser = create_parser()
+    main(vars(parser.parse_args()))

examples/offline_inference/spec_decode.py

@@ -5,14 +5,9 @@ from transformers import AutoTokenizer
 
 from vllm import LLM, SamplingParams
 from vllm.benchmarks.datasets import add_dataset_parser, get_samples
+from vllm.utils.argparse_utils import FlexibleArgumentParser
 from vllm.v1.metrics.reader import Counter, Vector
 
-try:
-    from vllm.utils.argparse_utils import FlexibleArgumentParser
-except ImportError:
-    from argparse import ArgumentParser as FlexibleArgumentParser
-
-
 QUESTION = "What is the content of each image?"
 IMAGE_URLS = [
     "https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/duck.jpg",

examples/online_serving/openai_chat_completion_client_for_multimodal.py

@@ -10,7 +10,7 @@ vllm serve llava-hf/llava-1.5-7b-hf
 
 (multi-image inference with Phi-3.5-vision-instruct)
 vllm serve microsoft/Phi-3.5-vision-instruct --runner generate \
-    --trust-remote-code --max-model-len 4096 --limit-mm-per-prompt '{"image":2}'
+    --trust-remote-code --max-model-len 4096 --limit-mm-per-prompt.image 2
 
 (audio inference with Ultravox)
 vllm serve fixie-ai/ultravox-v0_5-llama-3_2-1b \

examples/online_serving/openai_transcription_client.py

@@ -26,7 +26,9 @@ from openai import AsyncOpenAI, OpenAI
 from vllm.assets.audio import AudioAsset
 
 
-def sync_openai(audio_path: str, client: OpenAI, model: str):
+def sync_openai(
+    audio_path: str, client: OpenAI, model: str, *, repetition_penalty: float = 1.3
+):
     """
     Perform synchronous transcription using OpenAI-compatible API.
     """
@@ -40,7 +42,7 @@ def sync_openai(audio_path: str, client: OpenAI, model: str):
             # Additional sampling params not provided by OpenAI API.
             extra_body=dict(
                 seed=4419,
-                repetition_penalty=1.3,
+                repetition_penalty=repetition_penalty,
             ),
         )
         print("transcription result [sync]:", transcription.text)
@@ -129,7 +131,12 @@ def main(args):
     print(f"Using model: {model}")
 
     # Run the synchronous function
-    sync_openai(args.audio_path if args.audio_path else mary_had_lamb, client, model)
+    sync_openai(
+        audio_path=args.audio_path if args.audio_path else mary_had_lamb,
+        client=client,
+        model=model,
+        repetition_penalty=args.repetition_penalty,
+    )
 
     # Run the asynchronous function
     if "openai" in model:
@@ -161,5 +168,11 @@ if __name__ == "__main__":
         default=None,
         help="The path to the audio file to transcribe.",
     )
+    parser.add_argument(
+        "--repetition_penalty",
+        type=float,
+        default=1.3,
+        help="repetition penalty",
+    )
     args = parser.parse_args()
     main(args)

examples/pooling/classify/vision_classification_online.py

@@ -7,7 +7,7 @@ NOTE:
     vllm serve muziyongshixin/Qwen2.5-VL-7B-for-VideoCls \
          --runner pooling \
          --max-model-len 5000 \
-         --limit-mm-per-prompt '{"video": 1}' \
+         --limit-mm-per-prompt.video 1 \
          --hf-overrides '{"text_config": {"architectures": ["Qwen2_5_VLForSequenceClassification"]}}'
 """
 

mkdocs.yaml

@@ -63,8 +63,9 @@ plugins:
   - git-revision-date-localized:
       # exclude autogenerated files
       exclude:
-        - argparse/*
+        - api/*
         - examples/*
+        - generated/*
   - minify:
       minify_html: true
       minify_js: true
@@ -92,7 +93,6 @@ plugins:
               - "!.*_pb2_grpc"  # Exclude auto-generated gRPC stubs
             summary:
               modules: true
-            show_if_no_docstring: true
             show_signature_annotations: true
             separate_signature: true
             show_overloads: true

requirements/common.txt

@@ -7,7 +7,7 @@ requests >= 2.26.0
 tqdm
 blake3
 py-cpuinfo
-transformers >= 4.56.0, < 5
+transformers @ git+https://github.com/huggingface/transformers.git@main
 tokenizers >= 0.21.1  # Required for fast incremental detokenization.
 protobuf >= 5.29.6, !=6.30.*, !=6.31.*, !=6.32.*, !=6.33.0.*, !=6.33.1.*, !=6.33.2.*, !=6.33.3.*, !=6.33.4.* # Required by LlamaTokenizer, gRPC. CVE-2026-0994
 fastapi[standard] >= 0.115.0 # Required by FastAPI's form models in the OpenAI API server's audio transcriptions endpoint.

requirements/nightly_torch_test.txt

@@ -43,5 +43,5 @@ tritonclient>=2.51.0
 numba == 0.61.2 # Required for N-gram speculative decoding
 numpy
 runai-model-streamer[s3,gcs]==0.15.3
-fastsafetensors>=0.1.10
+fastsafetensors>=0.2.2
 pydantic>=2.12 # 2.11 leads to error on python 3.13

requirements/test.in

@@ -53,7 +53,7 @@ arctic-inference == 0.1.1 # Required for suffix decoding test
 numba == 0.61.2 # Required for N-gram speculative decoding
 numpy
 runai-model-streamer[s3,gcs]==0.15.3
-fastsafetensors>=0.1.10
+fastsafetensors>=0.2.2 # 0.2.2 contains important fixes for multi-GPU mem usage
 pydantic>=2.12 # 2.11 leads to error on python 3.13
 decord==0.6.0
 terratorch @ git+https://github.com/IBM/terratorch.git@1.1.rc3 # required for PrithviMAE test

requirements/test.txt

@@ -224,7 +224,7 @@ fastparquet==2024.11.0
     # via genai-perf
 fastrlock==0.8.2
     # via cupy-cuda12x
-fastsafetensors==0.1.10
+fastsafetensors==0.2.2
     # via -r requirements/test.in
 filelock==3.16.1
     # via
@@ -1174,7 +1174,6 @@ torch==2.10.0+cu129
     #   bitsandbytes
     #   efficientnet-pytorch
     #   encodec
-    #   fastsafetensors
     #   kornia
     #   lightly
     #   lightning

setup.py

@@ -1035,7 +1035,7 @@ setup(
     extras_require={
         "bench": ["pandas", "matplotlib", "seaborn", "datasets", "scipy"],
         "tensorizer": ["tensorizer==2.10.1"],
-        "fastsafetensors": ["fastsafetensors >= 0.1.10"],
+        "fastsafetensors": ["fastsafetensors >= 0.2.2"],
         "runai": ["runai-model-streamer[s3,gcs] >= 0.15.3"],
         "audio": [
             "librosa",

tests/benchmarks/test_serve_cli.py

@@ -1,15 +1,76 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import subprocess
+import tempfile
+import time
+from pathlib import Path
 
 import pytest
+import requests
+import urllib3
 
 from ..utils import RemoteOpenAIServer
 
 MODEL_NAME = "meta-llama/Llama-3.2-1B-Instruct"
 
 
-@pytest.fixture(scope="module")
+def generate_self_signed_cert(cert_dir: Path) -> tuple[Path, Path]:
+    """Generate a self-signed certificate for testing."""
+    cert_file = cert_dir / "cert.pem"
+    key_file = cert_dir / "key.pem"
+
+    # Generate self-signed certificate using openssl
+    subprocess.run(
+        [
+            "openssl",
+            "req",
+            "-x509",
+            "-newkey",
+            "rsa:2048",
+            "-keyout",
+            str(key_file),
+            "-out",
+            str(cert_file),
+            "-days",
+            "1",
+            "-nodes",
+            "-subj",
+            "/CN=localhost",
+        ],
+        check=True,
+        capture_output=True,
+    )
+    return cert_file, key_file
+
+
+class RemoteOpenAIServerSSL(RemoteOpenAIServer):
+    """RemoteOpenAIServer subclass that supports SSL with self-signed certs."""
+
+    @property
+    def url_root(self) -> str:
+        return f"https://{self.host}:{self.port}"
+
+    def _wait_for_server(self, *, url: str, timeout: float):
+        """Override to use HTTPS with SSL verification disabled."""
+        # Suppress InsecureRequestWarning for self-signed certs
+        urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
+
+        start = time.time()
+        while True:
+            try:
+                if requests.get(url, verify=False).status_code == 200:
+                    break
+            except Exception:
+                result = self._poll()
+                if result is not None and result != 0:
+                    raise RuntimeError("Server exited unexpectedly.") from None
+
+                time.sleep(0.5)
+                if time.time() - start > timeout:
+                    raise RuntimeError("Server failed to start in time.") from None
+
+
+@pytest.fixture(scope="function")
 def server():
     args = ["--max-model-len", "1024", "--enforce-eager", "--load-format", "dummy"]
 
@@ -17,6 +78,27 @@ def server():
         yield remote_server
 
 
+@pytest.fixture(scope="function")
+def ssl_server():
+    """Start a vLLM server with SSL enabled using a self-signed certificate."""
+    with tempfile.TemporaryDirectory() as cert_dir:
+        cert_file, key_file = generate_self_signed_cert(Path(cert_dir))
+        args = [
+            "--max-model-len",
+            "1024",
+            "--enforce-eager",
+            "--load-format",
+            "dummy",
+            "--ssl-certfile",
+            str(cert_file),
+            "--ssl-keyfile",
+            str(key_file),
+        ]
+
+        with RemoteOpenAIServerSSL(MODEL_NAME, args) as remote_server:
+            yield remote_server
+
+
 @pytest.mark.benchmark
 def test_bench_serve(server):
     # Test default model detection and input/output len
@@ -42,6 +124,31 @@ def test_bench_serve(server):
     assert result.returncode == 0, f"Benchmark failed: {result.stderr}"
 
 
+@pytest.mark.benchmark
+def test_bench_serve_insecure(ssl_server):
+    """Test --insecure flag with an HTTPS server using a self-signed certificate."""
+    base_url = f"https://{ssl_server.host}:{ssl_server.port}"
+    command = [
+        "vllm",
+        "bench",
+        "serve",
+        "--base-url",
+        base_url,
+        "--input-len",
+        "32",
+        "--output-len",
+        "4",
+        "--num-prompts",
+        "5",
+        "--insecure",
+    ]
+    result = subprocess.run(command, capture_output=True, text=True)
+    print(result.stdout)
+    print(result.stderr)
+
+    assert result.returncode == 0, f"Benchmark failed: {result.stderr}"
+
+
 @pytest.mark.benchmark
 def test_bench_serve_chat(server):
     command = [

tests/compile/fullgraph/test_simple.py

@@ -27,10 +27,29 @@ from ...utils import create_new_process_for_each_test
 from ..silly_attention import get_global_counter, reset_global_counter
 
 
+# Custom op that returns an unbacked symint during graph capture
+@torch.library.custom_op("mylib::foo", mutates_args=())
+def foo(x: torch.Tensor) -> int:
+    return 3
+
+
+@foo.register_fake
+def _(x):
+    return torch.library.get_ctx().new_dynamic_size()
+
+
 @support_torch_compile
 class SillyModel(nn.Module):
-    def __init__(self, *, vllm_config: VllmConfig, prefix: str = "", **kwargs) -> None:
+    def __init__(
+        self,
+        *,
+        vllm_config: VllmConfig,
+        prefix: str = "",
+        intermediate_unbacked=False,
+        **kwargs,
+    ) -> None:
         super().__init__()
+        self.intermediate_unbacked = intermediate_unbacked
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         """
@@ -44,6 +63,13 @@ class SillyModel(nn.Module):
         torch.ops.silly.attention(x, x, x, out)
         x = out
         x = x - 2
+
+        if self.intermediate_unbacked:
+            # Test for unbacked symints: the following is a fancy way to multiply by 1
+            u0 = foo(x)
+            ones = x.new_ones(x.shape[0], u0).sum(-1) / 3
+            x = x * ones
+
         x = x - 1
         out = torch.empty_like(x)
         torch.ops.silly.attention(x, x, x, out)
@@ -52,6 +78,7 @@ class SillyModel(nn.Module):
         return x
 
 
+@torch._dynamo.config.patch(capture_dynamic_output_shape_ops=True)
 def _run_simple_model(
     splitting_ops,
     use_inductor_graph_partition,
@@ -60,6 +87,8 @@ def _run_simple_model(
     expected_num_piecewise_capturable_graphs_seen,
     expected_num_backend_compilations,
     expected_num_cudagraph_captured,
+    *,
+    intermediate_unbacked=False,
 ):
     vllm_config = VllmConfig(
         compilation_config=CompilationConfig(
@@ -72,7 +101,11 @@ def _run_simple_model(
         )
     )
     with set_current_vllm_config(vllm_config):
-        model = SillyModel(vllm_config=vllm_config, prefix="")
+        model = SillyModel(
+            vllm_config=vllm_config,
+            prefix="",
+            intermediate_unbacked=intermediate_unbacked,
+        )
 
     inputs = torch.randn(100).cuda()
 
@@ -125,9 +158,10 @@ def _run_simple_model(
 
 
 @pytest.mark.parametrize("backend", ["inductor", "eager"])
+@pytest.mark.parametrize("intermediate_unbacked", [True, False])
 @torch.inference_mode()
 @create_new_process_for_each_test("spawn")
-def test_simple_piecewise_compile(backend):
+def test_simple_piecewise_compile(backend, intermediate_unbacked):
     _run_simple_model(
         splitting_ops=["silly::attention"],
         use_inductor_graph_partition=False,
@@ -140,6 +174,7 @@ def test_simple_piecewise_compile(backend):
         expected_num_backend_compilations=3,
         # num_cudagraph_sizes * num_piecewise_capturable_graphs_seen
         expected_num_cudagraph_captured=6,
+        intermediate_unbacked=intermediate_unbacked,
     )
 
 

tests/compile/passes/distributed/test_fusion_all_reduce.py

@@ -202,9 +202,10 @@ class TestAllReduceFusedAddRMSNormStaticQuantFP4Model(torch.nn.Module):
 @pytest.mark.skipif(envs.VLLM_TARGET_DEVICE not in ["cuda"], reason="Only test on CUDA")
 @pytest.mark.skipif(
     not find_spec("flashinfer")
-    or not has_module_attribute("flashinfer.comm", "trtllm_allreduce_fusion"),
+    or not has_module_attribute("flashinfer.comm", "allreduce_fusion")
+    or not has_module_attribute("flashinfer.comm", "create_allreduce_fusion_workspace"),
     reason="flashinfer is not found or flashinfer "
-    "is not compiled with trtllm_allreduce_fusion",
+    "is not compiled with allreduce_fusion",
 )
 def test_all_reduce_fusion_pass_replace(
     test_model: torch.nn.Module,

tests/compile/passes/test_functionalization.py

@@ -1,11 +1,10 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-import copy
-
 import pytest
 import torch
 
+import vllm.envs as envs
 from tests.compile.backend import TestBackend
 from tests.utils import TestFP8Layer
 from vllm.compilation.passes.fusion.act_quant_fusion import (
@@ -32,7 +31,6 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
 )
 from vllm.model_executor.layers.rotary_embedding import get_rope
 from vllm.platforms import current_platform
-from vllm.utils.torch_utils import direct_register_custom_op
 
 TEST_FP8 = current_platform.supports_fp8()
 FP8_DTYPE = current_platform.fp8_dtype()
@@ -200,82 +198,23 @@ class TestRotaryEmbeddingSliceScatter(torch.nn.Module):
         return [torch.ops.aten.slice_scatter.default]
 
 
-class TestFunctionWithMutatedArgsAndReturn(torch.nn.Module):
-    OP_REGISTERED = False
-
-    def __init__(self):
-        super().__init__()
-        self.register_test_custom_op()
-
-    @classmethod
-    def register_test_custom_op(cls):
-        if not cls.OP_REGISTERED:
-
-            def function_with_mutated_args_and_return_impl(
-                x: torch.Tensor,
-            ) -> torch.Tensor:
-                ret = x + 1
-                x.add_(2)
-                return ret
-
-            def function_with_mutated_args_and_return_fake(
-                x: torch.Tensor,
-            ) -> torch.Tensor:
-                return torch.empty_like(x)
-
-            direct_register_custom_op(
-                op_name="function_with_mutated_args_and_return",
-                op_func=function_with_mutated_args_and_return_impl,
-                mutates_args=["x"],
-                fake_impl=function_with_mutated_args_and_return_fake,
-            )
-
-            cls.OP_REGISTERED = True
-
-    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
-        # Clone x to avoid mutating the original tensor
-        ret = torch.ops.vllm.function_with_mutated_args_and_return(x)
-        return x, ret
-
-    def example_inputs(self, num_tokens=32):
-        hidden_states = torch.randn(num_tokens)
-        return (hidden_states,)
-
-    def ops_in_model(self, do_fusion):
-        return [torch.ops.vllm.function_with_mutated_args_and_return.default]
-
-    def ops_not_in_model(self):
-        return []
-
-
-MODELS_AND_DO_FUSION = {
-    TestSiluMul: [True, False],
-    TestFusedAddRMSNorm: [True, False],
-    TestRotaryEmbedding: [False],
-    TestRotaryEmbeddingSliceScatter: [False],
-    TestFunctionWithMutatedArgsAndReturn: [False],
-}
+MODELS = [
+    TestSiluMul,
+    TestFusedAddRMSNorm,
+    TestRotaryEmbedding,
+    TestRotaryEmbeddingSliceScatter,
+]
 
 
 @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
-@pytest.mark.parametrize(
-    "model_class, do_fusion",
-    [
-        (model_class, do_fusion)
-        for model_class, fusions in MODELS_AND_DO_FUSION.items()
-        for do_fusion in fusions
-    ],
-)
-@pytest.mark.skipif(
-    not current_platform.is_cuda_alike(),
-    reason="Only test on cuda and rocm platform",
-)
+@pytest.mark.parametrize("model_class", MODELS)
+@pytest.mark.parametrize("do_fusion", [True, False])
+@pytest.mark.skipif(envs.VLLM_TARGET_DEVICE != "cuda", reason="Only test on CUDA")
 def test_fix_functionalization(
     model_class: torch.nn.Module, do_fusion: bool, dtype: torch.dtype
 ):
     torch.set_default_device("cuda")
     torch.set_default_dtype(dtype)
-    torch.manual_seed(0)
 
     vllm_config = VllmConfig(
         model_config=ModelConfig(dtype=dtype),
@@ -307,14 +246,8 @@ def test_fix_functionalization(
         backend_no_func = TestBackend(*passes)
 
         model = model_class()
-        inputs_func = model.example_inputs()
-        inputs_no_func = copy.deepcopy(inputs_func)
-        model_func = model_class()
-        model_no_func = copy.deepcopy(model_func)
-        model_func = torch.compile(model_func, backend=backend_func)
-        model_no_func = torch.compile(model_no_func, backend=backend_no_func)
-        model_func(*inputs_func)
-        model_no_func(*inputs_no_func)
+        torch.compile(model, backend=backend_func)(*model.example_inputs())
+        torch.compile(model, backend=backend_no_func)(*model.example_inputs())
 
         # check if the functionalization pass is applied
         for op in model.ops_in_model(do_fusion):
@@ -332,8 +265,3 @@ def test_fix_functionalization(
                     found[op] = True
         assert all(found[op] for op in model.ops_in_model(do_fusion))
         assert all(not found.get(op) for op in model.ops_not_in_model())
-
-        # TODO (Rohan138): compare the outputs from model_func and model_no_func
-        # currently runs into errors while comparing `TestFusedAddRMSNorm`
-        # Linked issue: https://github.com/vllm-project/vllm/issues/34996
-        # torch.testing.assert_close(outputs_func, outputs_no_func)

tests/compile/passes/test_fusion_attn.py

@@ -92,6 +92,8 @@ class AttentionQuantPatternModel(torch.nn.Module):
     def build_attn_metadata(self, batch_size: int) -> AttentionMetadata:
         """Initialize attention metadata."""
 
+        # TODO (Rohan138) reuse utils from vllm/v1/worker/gpu/attn_utils.py
+
         # Create common attn metadata
         batch_spec = BatchSpec(seq_lens=[1] * batch_size, query_lens=[1] * batch_size)
         common_attn_metadata = create_common_attn_metadata(
@@ -100,58 +102,31 @@ class AttentionQuantPatternModel(torch.nn.Module):
 
         max_blocks = (max(batch_spec.seq_lens) + self.block_size - 1) // self.block_size
         num_blocks = batch_size * max_blocks
-        backend = self.attn.backend
 
-        # TODO(luka) use get_kv_cache_stride_order
-        # Create dummy KV cache for the selected backend
-        if backend == AttentionBackendEnum.ROCM_ATTN:
-            # k/v as 1st dimention
-            # HND: [num_blocks, num_kv_heads, block_size, head_size]
-            kv_cache = torch.zeros(
-                2,
-                num_blocks,
-                self.num_kv_heads,
-                self.block_size,
-                self.head_size,
-                dtype=self.kv_cache_dtype,
-                device=self.device,
-            )
-        elif backend == AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN:
-            # k/v as 1st dimention
-            # NHD: [num_blocks, block_size, num_kv_heads, head_size]
-            kv_cache = torch.zeros(
-                2,
-                num_blocks,
-                self.block_size,
-                self.num_kv_heads,
-                self.head_size,
-                dtype=self.kv_cache_dtype,
-                device=self.device,
-            )
-        elif backend == AttentionBackendEnum.TRITON_ATTN:
-            # k/v as 2nd dimention
-            # NHD: [num_blocks, block_size, num_kv_heads, head_size]
-            kv_cache = torch.zeros(
-                num_blocks,
-                2,
-                self.num_kv_heads,
-                self.block_size,
-                self.head_size,
-                dtype=self.kv_cache_dtype,
-                device=self.device,
-            )
-        elif backend == AttentionBackendEnum.FLASHINFER:
-            kv_cache = torch.zeros(
-                num_blocks,
-                2,
-                self.num_kv_heads,
-                self.block_size,
-                self.head_size,
-                dtype=self.kv_cache_dtype,
-                device=self.device,
-            ).permute(0, 1, 3, 2, 4)
-        else:
-            raise ValueError(f"Unsupported backend: {backend}")
+        # Fetch the attention backend and kv cache shape and stride order
+        attn_backend = self.attn.attn_backend
+        kv_cache_shape = attn_backend.get_kv_cache_shape(
+            num_blocks, self.block_size, self.num_kv_heads, self.head_size
+        )
+        try:
+            kv_cache_stride_order = attn_backend.get_kv_cache_stride_order()
+        except (AttributeError, NotImplementedError):
+            kv_cache_stride_order = tuple(range(len(kv_cache_shape)))
+
+        kv_cache_shape = tuple(kv_cache_shape[i] for i in kv_cache_stride_order)
+        inv_order = [
+            kv_cache_stride_order.index(i) for i in range(len(kv_cache_stride_order))
+        ]
+
+        # Create dummy KV cache
+        raw_tensor = torch.zeros(
+            2 * num_blocks * self.block_size * self.num_kv_heads * self.head_size,
+            dtype=self.kv_cache_dtype,
+            device=self.device,
+        )
+        raw_tensor = raw_tensor.view(kv_cache_shape)
+        kv_cache = raw_tensor.permute(*inv_order)
+
         self.attn.kv_cache = [kv_cache]
 
         # Build attn metadata

tests/compile/passes/test_rope_kvcache_fusion.py

@@ -1,325 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import pytest
-import torch
-
-import vllm.config
-from tests.compile.backend import TestBackend
-from tests.v1.attention.utils import BatchSpec, create_common_attn_metadata
-from vllm._aiter_ops import is_aiter_found_and_supported, rocm_aiter_ops
-from vllm.compilation.passes.fusion.matcher_utils import ROTARY_OP
-from vllm.compilation.passes.fusion.rope_kvcache_fusion import RopeKVCacheFusionPass
-from vllm.compilation.passes.utility.noop_elimination import NoOpEliminationPass
-from vllm.compilation.passes.utility.post_cleanup import PostCleanupPass
-from vllm.compilation.passes.utility.scatter_split_replace import (
-    ScatterSplitReplacementPass,
-)
-from vllm.compilation.passes.utility.split_coalescing import SplitCoalescingPass
-from vllm.config import (
-    CacheConfig,
-    CompilationConfig,
-    CompilationMode,
-    ModelConfig,
-    PassConfig,
-    VllmConfig,
-)
-from vllm.forward_context import get_forward_context, set_forward_context
-from vllm.model_executor.layers.attention import Attention
-from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding
-from vllm.platforms import current_platform
-from vllm.v1.attention.backend import (
-    AttentionBackend,
-    CommonAttentionMetadata,
-)
-from vllm.v1.attention.backends.registry import AttentionBackendEnum
-from vllm.v1.kv_cache_interface import AttentionSpec
-
-INDEX_SELECT_OP = torch.ops.aten.index.Tensor
-VLLM_UNIFIED_KV_CACHE_UPDATE_OP = torch.ops.vllm.unified_kv_cache_update
-FP8_DTYPE = current_platform.fp8_dtype()
-
-
-class QKRoPEKVCacheTestModel(torch.nn.Module):
-    def __init__(
-        self,
-        vllm_config: VllmConfig,
-        attn_backend: AttentionBackendEnum,
-        num_heads: int,
-        num_kv_heads: int,
-        head_size: int,
-        is_neox: bool,
-        dtype: torch.dtype,
-        device: torch.device,
-        prefix: str = "model.layers.0.self_attn.attn",
-    ):
-        super().__init__()
-        self.num_heads = num_heads
-        self.num_kv_heads = num_kv_heads
-        self.head_size = head_size
-        self.block_size = vllm_config.cache_config.block_size
-        self.q_size = num_heads * head_size
-        self.kv_size = num_kv_heads * head_size
-        self.is_neox = is_neox
-        self.dtype = dtype
-        self.device = device
-        self.layer_name = prefix
-
-        self.rotary_emb = RotaryEmbedding(
-            head_size,
-            rotary_dim=head_size,
-            max_position_embeddings=4096,
-            base=10000,
-            is_neox_style=is_neox,
-            dtype=self.dtype,
-        )
-
-        # Whether to check for the RoPE custom op or component index_select
-        self.enable_rope_custom_op = self.rotary_emb.enabled()
-
-        # Register layer metadata for the fusion pass via Attention.
-        self.attn = Attention(
-            num_heads=num_heads,
-            head_size=head_size,
-            scale=1.0 / head_size**0.5,
-            num_kv_heads=num_kv_heads,
-            cache_config=vllm_config.cache_config,
-            quant_config=vllm_config.quant_config,
-            prefix=prefix,
-            attn_backend=attn_backend.get_class(),
-        )
-        self.attn_backend: type[AttentionBackend] = self.attn.get_attn_backend()
-        assert not self.attn_backend.forward_includes_kv_cache_update, (
-            f"Attention backend {self.attn_backend} does not support fuse_rope_kvcache."
-        )
-        self.attn._k_scale = self.attn._k_scale.to(device)
-        self.attn._v_scale = self.attn._v_scale.to(device)
-
-        kv_cache_dtype_str = vllm_config.cache_config.cache_dtype
-        self.kv_cache_dtype = (
-            FP8_DTYPE if kv_cache_dtype_str.startswith("fp8") else self.dtype
-        )
-
-        # Initialize attn MetadataBuilder
-        self.builder = self.attn.attn_backend.get_builder_cls()(
-            kv_cache_spec=AttentionSpec(
-                block_size=self.block_size,
-                num_kv_heads=self.num_kv_heads,
-                head_size=head_size,
-                dtype=self.kv_cache_dtype,
-            ),
-            layer_names=[self.attn.layer_name],
-            vllm_config=vllm_config,
-            device=device,
-        )
-
-    def build_attn_metadata(self, batch_size: int) -> CommonAttentionMetadata:
-        """Initialize attention metadata."""
-        # Create common attn metadata
-        batch_spec = BatchSpec(seq_lens=[1] * batch_size, query_lens=[1] * batch_size)
-        common_attn_metadata = create_common_attn_metadata(
-            batch_spec, self.block_size, self.device, arange_block_indices=True
-        )
-
-        max_blocks = (max(batch_spec.seq_lens) + self.block_size - 1) // self.block_size
-        num_blocks = batch_size * max_blocks
-
-        # Fetch the attention backend and kv cache shape and stride order
-        attn_backend = self.attn.attn_backend
-        kv_cache_shape = attn_backend.get_kv_cache_shape(
-            num_blocks, self.block_size, self.num_kv_heads, self.head_size
-        )
-        try:
-            kv_cache_stride_order = attn_backend.get_kv_cache_stride_order()
-        except (AttributeError, NotImplementedError):
-            kv_cache_stride_order = tuple(range(len(kv_cache_shape)))
-
-        kv_cache_shape = tuple(kv_cache_shape[i] for i in kv_cache_stride_order)
-        inv_order = [
-            kv_cache_stride_order.index(i) for i in range(len(kv_cache_stride_order))
-        ]
-
-        # Create dummy KV cache
-        raw_tensor = torch.zeros(
-            2 * num_blocks * self.block_size * self.num_kv_heads * self.head_size,
-            dtype=self.kv_cache_dtype,
-            device=self.device,
-        )
-        raw_tensor = raw_tensor.view(kv_cache_shape)
-        kv_cache = raw_tensor.permute(*inv_order)
-
-        self.attn.kv_cache = [kv_cache]
-
-        # Build attn metadata
-        attn_metadata = self.builder.build(
-            common_prefix_len=0, common_attn_metadata=common_attn_metadata
-        )
-
-        return attn_metadata
-
-    def forward(
-        self, qkv: torch.Tensor, positions: torch.Tensor
-    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
-        # Create copy so inplace ops do not modify the original tensors
-        qkv = qkv.clone()
-        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
-        q, k = self.rotary_emb(positions, q, k)
-
-        # Instead of a full forward pass, match only the KV cache update op here
-        q = q.view(-1, self.num_heads, self.head_size)
-        k = k.view(-1, self.num_kv_heads, self.head_size)
-        v = v.view(-1, self.num_kv_heads, self.head_size)
-        kv_cache_dummy_dep = torch.ops.vllm.unified_kv_cache_update(
-            k, v, self.layer_name
-        )
-        return q, k, v, kv_cache_dummy_dep
-
-    def ops_in_model_before(self) -> list[torch._ops.OpOverload]:
-        ops = []
-        if self.enable_rope_custom_op:
-            ops.append(ROTARY_OP)
-        else:
-            ops.append(INDEX_SELECT_OP)
-        ops.append(torch.ops.vllm.unified_kv_cache_update.default)
-        return ops
-
-    def ops_in_model_after(self) -> list[torch._ops.OpOverload]:
-        return [torch.ops.vllm.fused_rope_and_unified_kv_cache_update.default]
-
-
-@pytest.mark.parametrize(
-    "attn_backend",
-    [
-        AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN,
-        AttentionBackendEnum.TRITON_ATTN,
-        AttentionBackendEnum.ROCM_ATTN,
-    ],
-)
-@pytest.mark.parametrize("enable_rope_custom_op", [True])  # [True, False])
-@pytest.mark.parametrize("num_heads", [64])
-@pytest.mark.parametrize("num_kv_heads", [8])
-@pytest.mark.parametrize("head_size", [64])
-@pytest.mark.parametrize("block_size", [16])
-@pytest.mark.parametrize("is_neox", [True, False])
-@pytest.mark.parametrize("dtype", [torch.bfloat16])
-@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"])
-@pytest.mark.skipif(
-    not is_aiter_found_and_supported(),
-    reason="Only test on ROCm with AITER installed and supported",
-)
-def test_rope_kvcache_fusion(
-    attn_backend: AttentionBackendEnum,
-    enable_rope_custom_op: bool,
-    num_heads: int,
-    num_kv_heads: int,
-    head_size: int,
-    block_size: int,
-    is_neox: bool,
-    dtype: torch.dtype,
-    kv_cache_dtype: str,
-    monkeypatch: pytest.MonkeyPatch,
-):
-    torch.set_default_device("cuda")
-    torch.set_default_dtype(dtype)
-    torch.manual_seed(0)
-
-    custom_ops: list[str] = []
-    if enable_rope_custom_op:
-        custom_ops.append("+rotary_embedding")
-
-    vllm_config = VllmConfig(
-        model_config=ModelConfig(dtype=dtype),
-        cache_config=CacheConfig(
-            block_size=block_size,
-            cache_dtype=kv_cache_dtype,
-        ),
-        compilation_config=CompilationConfig(
-            mode=CompilationMode.VLLM_COMPILE,
-            custom_ops=custom_ops,
-            pass_config=PassConfig(
-                fuse_rope_kvcache=True,
-                eliminate_noops=True,
-            ),
-        ),
-    )
-
-    with vllm.config.set_current_vllm_config(vllm_config), monkeypatch.context() as m:
-        m.setenv("VLLM_ROCM_USE_AITER", "1")
-        rocm_aiter_ops.refresh_env_variables()
-
-        model = QKRoPEKVCacheTestModel(
-            vllm_config=vllm_config,
-            attn_backend=attn_backend,
-            num_heads=num_heads,
-            num_kv_heads=num_kv_heads,
-            head_size=head_size,
-            is_neox=is_neox,
-            dtype=dtype,
-            device=torch.get_default_device(),
-        )
-
-        fusion_pass = RopeKVCacheFusionPass(vllm_config)
-        passes = [
-            NoOpEliminationPass(vllm_config),
-            SplitCoalescingPass(vllm_config),
-            ScatterSplitReplacementPass(vllm_config),
-            fusion_pass,
-            PostCleanupPass(vllm_config),
-        ]
-        backend = TestBackend(*passes)
-
-        T = 5
-
-        qkv = torch.randn(
-            T, num_heads * head_size + 2 * num_kv_heads * head_size, dtype=dtype
-        )
-        pos = torch.arange(T, dtype=torch.long)
-
-        qkv_unfused = qkv.clone()
-        pos_unfused = pos.clone()
-
-        with set_forward_context(None, vllm_config):
-            forward_context = get_forward_context()
-            attn_metadata = model.build_attn_metadata(T)
-            forward_context.slot_mapping = {
-                model.layer_name: attn_metadata.slot_mapping
-            }
-            q_unfused, k_unfused, v_unfused, dummy = model(qkv_unfused, pos_unfused)
-            attn_layer = forward_context.no_compile_layers[model.layer_name]
-            kv_cache_unfused = attn_layer.kv_cache[forward_context.virtual_engine]
-        del dummy
-
-        torch._dynamo.mark_dynamic(qkv, 0)
-        torch._dynamo.mark_dynamic(pos, 0)
-        with set_forward_context(None, vllm_config):
-            model_fused = torch.compile(model, backend=backend)
-            forward_context = get_forward_context()
-            attn_metadata = model_fused.build_attn_metadata(T)
-            forward_context.slot_mapping = {
-                model.layer_name: attn_metadata.slot_mapping
-            }
-            q_fused, k_fused, v_fused, dummy = model_fused(qkv, pos)
-            attn_layer = forward_context.no_compile_layers[model.layer_name]
-            kv_cache_fused = attn_layer.kv_cache[forward_context.virtual_engine]
-        del dummy
-
-        assert fusion_pass.matched_count == 1
-
-        backend.check_before_ops(model.ops_in_model_before())
-        backend.check_after_ops(model.ops_in_model_after())
-
-        if dtype == torch.float16:
-            ATOL, RTOL = (2e-3, 2e-3)
-        else:
-            ATOL, RTOL = (1e-2, 1e-2)
-
-        torch.testing.assert_close(q_unfused, q_fused, atol=ATOL, rtol=RTOL)
-        torch.testing.assert_close(k_unfused, k_fused, atol=ATOL, rtol=RTOL)
-        torch.testing.assert_close(v_unfused, v_fused, atol=ATOL, rtol=RTOL)
-        # Cannot compare fp8_* directly here, cast to model dtype instead
-        torch.testing.assert_close(
-            kv_cache_unfused.view(dtype),
-            kv_cache_fused.view(dtype),
-            atol=ATOL,
-            rtol=RTOL,
-        )

tests/compile/passes/test_scatter_split_replace.py

@@ -1,107 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import pytest
-import torch
-import torch.nn as nn
-
-import vllm
-from tests.compile.backend import TestBackend
-from vllm.compilation.passes.utility.scatter_split_replace import (
-    ScatterSplitReplacementPass,
-)
-from vllm.compilation.passes.utility.split_coalescing import SplitCoalescingPass
-from vllm.config import CompilationConfig, CompilationMode, VllmConfig
-from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding
-
-
-class ScatterSplitReplacementModel(nn.Module):
-    """Model with a rope+getitem+slice_scatter+split_with_sizes sequence."""
-
-    def __init__(
-        self,
-        num_heads: int,
-        num_kv_heads: int,
-        head_size: int,
-        dtype: torch.dtype,
-    ):
-        super().__init__()
-        self.q_size = num_heads * head_size
-        self.kv_size = num_kv_heads * head_size
-
-        self.rotary_emb = RotaryEmbedding(
-            head_size,
-            rotary_dim=head_size,
-            max_position_embeddings=4096,
-            base=10000,
-            is_neox_style=True,
-            dtype=dtype,
-        )
-
-    def forward(self, qkv: torch.Tensor, positions: torch.Tensor):
-        # Create copy so inplace ops do not modify the original tensors
-        qkv = qkv.clone()
-        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
-        q, k = self.rotary_emb(positions, q, k)
-        q = q + 1
-        k = k + 2
-        v = v + 3
-        return q, k, v
-
-    def ops_in_model_before(self) -> list[torch._ops.OpOverload]:
-        return [
-            torch.ops.aten.slice_scatter.default,
-            torch.ops.aten.split_with_sizes.default,
-            torch.ops.aten.getitem.default,
-        ]
-
-    def ops_in_model_after(self) -> list[torch._ops.OpOverload]:
-        return [torch.ops.aten.getitem.default]
-
-
-@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
-def test_scatter_split_replace(dtype):
-    torch.set_default_device("cuda")
-    torch.set_default_dtype(dtype)
-    torch.manual_seed(0)
-
-    num_heads = 8
-    num_kv_heads = 4
-    head_size = 64
-
-    vllm_config = VllmConfig(
-        compilation_config=CompilationConfig(
-            mode=CompilationMode.VLLM_COMPILE,
-            custom_ops=["+rotary_embedding"],
-        ),
-    )
-    with vllm.config.set_current_vllm_config(vllm_config):
-        # ScatterSplitReplacementPass requires SplitCoalescingPass to be run before it
-        coalesce_pass = SplitCoalescingPass(vllm_config)
-        replace_pass = ScatterSplitReplacementPass(vllm_config)
-        passes = [coalesce_pass, replace_pass]
-        backend = TestBackend(*passes)
-
-        model = ScatterSplitReplacementModel(num_heads, num_kv_heads, head_size, dtype)
-
-        T = 5
-        qkv = torch.randn(
-            T, num_heads * head_size + 2 * num_kv_heads * head_size, dtype=dtype
-        )
-        pos = torch.arange(T, dtype=torch.long)
-
-        qkv_eager = qkv.clone()
-        pos_eager = pos.clone()
-        result_eager = model(qkv_eager, pos_eager)
-
-        torch._dynamo.mark_dynamic(qkv, 0)
-        torch._dynamo.mark_dynamic(pos, 0)
-
-        model_compiled = torch.compile(model, backend=backend)
-        result_compiled = model_compiled(qkv, pos)
-
-        for eager, compiled in zip(result_eager, result_compiled):
-            torch.testing.assert_close(eager, compiled)
-
-        assert backend.op_count(torch.ops.aten.slice_scatter.default) == 0
-        assert backend.op_count(torch.ops.aten.split_with_sizes.default) == 1

tests/config/test_config_generation.py

@@ -80,12 +80,19 @@ def test_ray_runtime_env(monkeypatch: pytest.MonkeyPatch):
     ray.shutdown()
 
 
-def test_unrecognized_env():
+def test_unrecognized_env(monkeypatch):
     import os
 
+    from vllm.envs import environment_variables
+
+    # Remove any existing unrecognized VLLM env vars that might interfere
+    for env in list(os.environ):
+        if env.startswith("VLLM_") and env not in environment_variables:
+            monkeypatch.delenv(env, raising=False)
+
     # Test that if fail_on_environ_validation is True, then an error
     # is raised when an unrecognized vLLM environment variable is set
-    os.environ["VLLM_UNRECOGNIZED_ENV_VAR"] = "some_value"
+    monkeypatch.setenv("VLLM_UNRECOGNIZED_ENV_VAR", "some_value")
     engine_args = EngineArgs(
         fail_on_environ_validation=True,
     )
@@ -97,7 +104,7 @@ def test_unrecognized_env():
     engine_args.create_engine_config()
 
     # Test that when the unrecognized env var is removed, no error is raised
-    os.environ.pop("VLLM_UNRECOGNIZED_ENV_VAR", None)
+    monkeypatch.delenv("VLLM_UNRECOGNIZED_ENV_VAR")
     engine_args = EngineArgs(
         fail_on_environ_validation=True,
     )

tests/distributed/test_eplb_execute.py

@@ -295,12 +295,11 @@ def _test_async_transfer_layer_without_mtp_worker(
     for layer_idx in range(num_layers):
         is_unchanged, is_received_locally, recv_metadata = asyncio.run(
             transfer_layer(
-                old_global_expert_indices=old_indices_cpu,
-                new_global_expert_indices=new_indices_cpu,
-                expert_weights=expert_weights,
+                old_layer_indices=old_indices_cpu[layer_idx],
+                new_layer_indices=new_indices_cpu[layer_idx],
+                expert_weights=expert_weights[layer_idx],
                 expert_weights_buffer=expert_buffer,
                 ep_group=ep_group,
-                layer=layer_idx,
                 cuda_stream=cuda_stream,
             )
         )

tests/entrypoints/openai/responses/test_harmony.py

@@ -1302,16 +1302,17 @@ async def test_system_prompt_override(client: OpenAI, model_name: str):
         # Message structure may vary, skip this specific check
         pass
 
+    custom_system_prompt_2 = (
+        "You are a helpful assistant that always responds in exactly 5 words."
+    )
+
     # Test 3: Test with different custom system prompt
     response_2 = await client.responses.create(
         model=model_name,
         input=[
             {
                 "role": "system",
-                "content": (
-                    "You are a helpful assistant that always "
-                    "responds in exactly 5 words."
-                ),
+                "content": custom_system_prompt_2,
             },
             {"role": "user", "content": "What is the weather like?"},
         ],
@@ -1328,3 +1329,27 @@ async def test_system_prompt_override(client: OpenAI, model_name: str):
     assert 3 <= word_count <= 8, (
         f"Expected around 5 words, got {word_count} words: {response_2.output_text}"
     )
+
+    # Test 4: Test with structured content
+    response_3 = await client.responses.create(
+        model=model_name,
+        input=[
+            {
+                "role": "system",
+                "content": [{"type": "input_text", "text": custom_system_prompt_2}],
+            },
+            {"role": "user", "content": "What is the weather like?"},
+        ],
+        temperature=0.0,
+    )
+
+    assert response_3 is not None
+    assert response_3.status == "completed"
+    assert response_3.output_text is not None
+
+    # Count words in response (approximately, allowing for punctuation)
+    word_count = len(response_3.output_text.split())
+    # Allow some flexibility (4-7 words) since the model might not be perfectly precise
+    assert 3 <= word_count <= 8, (
+        f"Expected around 5 words, got {word_count} words: {response_3.output_text}"
+    )

tests/entrypoints/openai/responses/test_parsable_context.py

@@ -1,7 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-import importlib
+import importlib.util
 import json
 
 import pytest
@@ -179,12 +179,12 @@ async def test_mcp_tool_call(client: OpenAI, model_name: str):
     assert response.output[2].type == "reasoning"
     # make sure the correct math is in the final output
     assert response.output[3].type == "message"
-    assert "56088" in response.output[3].content[0].text
+    assert any(s in response.output[3].content[0].text for s in ("56088", "56,088"))
 
     # test raw input_messages / output_messages
     assert len(response.input_messages) == 1
     assert len(response.output_messages) == 3
-    assert "56088" in response.output_messages[2]["message"]
+    assert any(s in response.output_messages[2]["message"] for s in ("56088", "56,088"))
 
 
 @pytest.mark.asyncio

tests/entrypoints/openai/test_realtime_validation.py

@@ -129,5 +129,5 @@ async def test_multi_chunk_streaming(
                 " First words I spoke in the original phonograph."
                 " A little piece of practical poetry. Mary had a little lamb,"
                 " it sleeps with quite a flow, and everywhere that Mary went,"
-                " the lamb was sure to go"
+                " the lamb was sure to go."
             )

tests/entrypoints/test_context.py

@@ -8,6 +8,7 @@ from openai_harmony import Author, Message, Role, StreamState, TextContent
 
 from vllm.entrypoints.openai.responses.context import (
     HarmonyContext,
+    SimpleContext,
     StreamingHarmonyContext,
     TurnMetrics,
 )
@@ -597,3 +598,248 @@ def test_turn_metrics_copy_and_reset():
     assert copied_metrics.output_tokens == 20
     assert copied_metrics.cached_input_tokens == 5
     assert copied_metrics.tool_output_tokens == 3
+
+
+# ==================== SimpleContext Tests ====================
+
+
+def create_simple_context_output(
+    text="",
+    token_ids=None,
+    prompt="Test prompt",
+    prompt_token_ids=None,
+    num_cached_tokens=0,
+    logprobs=None,
+    finished=True,
+):
+    """Helper to create a RequestOutput with customizable text for
+    SimpleContext tests."""
+    if token_ids is None:
+        token_ids = []
+    return RequestOutput(
+        request_id="test-id",
+        prompt=prompt,
+        prompt_token_ids=prompt_token_ids,
+        prompt_logprobs=None,
+        outputs=[
+            CompletionOutput(
+                index=0,
+                text=text,
+                token_ids=token_ids,
+                cumulative_logprob=0.0,
+                logprobs=logprobs,
+                finish_reason=None,
+                stop_reason=None,
+            )
+        ],
+        finished=finished,
+        num_cached_tokens=num_cached_tokens,
+    )
+
+
+def test_simple_context_output_messages_empty():
+    """output_messages should be empty before any output is appended."""
+    context = SimpleContext()
+    assert context.output_messages == []
+
+
+def test_simple_context_output_messages_single_call():
+    """Non-streaming: single append_output produces a single output message."""
+    context = SimpleContext()
+    output = create_simple_context_output(
+        text="Hello world",
+        token_ids=[10, 20, 30],
+        prompt_token_ids=[1, 2, 3],
+    )
+    context.append_output(output)
+
+    messages = context.output_messages
+    assert len(messages) == 1
+    assert messages[0].message == "Hello world"
+    assert messages[0].tokens == [10, 20, 30]
+    assert messages[0].type == "raw_message_tokens"
+
+
+def test_simple_context_output_messages_streaming_consolidation():
+    """Streaming: multiple append_output calls consolidate into one message."""
+    context = SimpleContext()
+
+    # Simulate 3 streaming deltas
+    context.append_output(
+        create_simple_context_output(
+            text="Hello",
+            token_ids=[10],
+            prompt_token_ids=[1, 2, 3],
+        )
+    )
+    context.append_output(
+        create_simple_context_output(
+            text=" world",
+            token_ids=[20],
+            prompt_token_ids=[1, 2, 3],
+        )
+    )
+    context.append_output(
+        create_simple_context_output(
+            text="!",
+            token_ids=[30],
+            prompt_token_ids=[1, 2, 3],
+        )
+    )
+
+    messages = context.output_messages
+    assert len(messages) == 1
+    assert messages[0].message == "Hello world!"
+    assert messages[0].tokens == [10, 20, 30]
+
+
+def test_simple_context_output_messages_many_deltas():
+    """Streaming with many small deltas still produces a single message."""
+    context = SimpleContext()
+
+    words = ["The", " quick", " brown", " fox", " jumps"]
+    for i, word in enumerate(words):
+        context.append_output(
+            create_simple_context_output(
+                text=word,
+                token_ids=[100 + i],
+                prompt_token_ids=[1, 2],
+            )
+        )
+
+    messages = context.output_messages
+    assert len(messages) == 1
+    assert messages[0].message == "The quick brown fox jumps"
+    assert messages[0].tokens == [100, 101, 102, 103, 104]
+
+
+def test_simple_context_input_messages():
+    """input_messages is populated on the first append_output call."""
+    context = SimpleContext()
+    assert context.input_messages == []
+
+    context.append_output(
+        create_simple_context_output(
+            text="Hi",
+            token_ids=[10],
+            prompt="My prompt text",
+            prompt_token_ids=[1, 2, 3],
+        )
+    )
+
+    assert len(context.input_messages) == 1
+    assert context.input_messages[0].message == "My prompt text"
+    assert context.input_messages[0].tokens == [1, 2, 3]
+
+    # Second call should not add another input message
+    context.append_output(
+        create_simple_context_output(
+            text=" there",
+            token_ids=[20],
+            prompt="My prompt text",
+            prompt_token_ids=[1, 2, 3],
+        )
+    )
+
+    assert len(context.input_messages) == 1
+
+
+def test_simple_context_token_counting():
+    """Token counting accumulates across streaming deltas."""
+    context = SimpleContext()
+
+    context.append_output(
+        create_simple_context_output(
+            text="a",
+            token_ids=[10, 11],
+            prompt_token_ids=[1, 2, 3, 4, 5],
+            num_cached_tokens=2,
+        )
+    )
+    context.append_output(
+        create_simple_context_output(
+            text="b",
+            token_ids=[12],
+            prompt_token_ids=[1, 2, 3, 4, 5],
+            num_cached_tokens=2,
+        )
+    )
+
+    assert context.num_prompt_tokens == 5
+    assert context.num_output_tokens == 3  # 2 + 1
+    assert context.num_cached_tokens == 2
+
+
+def test_simple_context_final_output():
+    """final_output reconstructs accumulated text and token_ids."""
+    context = SimpleContext()
+
+    context.append_output(
+        create_simple_context_output(
+            text="foo",
+            token_ids=[1, 2],
+            prompt_token_ids=[10],
+        )
+    )
+    context.append_output(
+        create_simple_context_output(
+            text="bar",
+            token_ids=[3],
+            prompt_token_ids=[10],
+        )
+    )
+
+    final = context.final_output
+    assert final is not None
+    assert final.outputs[0].text == "foobar"
+    assert final.outputs[0].token_ids == (1, 2, 3)
+
+
+def test_simple_context_output_messages_empty_text_with_tokens():
+    """output_messages should be returned when tokens exist even if text is
+    empty (e.g. special tokens)."""
+    context = SimpleContext()
+    context.append_output(
+        create_simple_context_output(
+            text="",
+            token_ids=[99],
+            prompt_token_ids=[1],
+        )
+    )
+
+    messages = context.output_messages
+    assert len(messages) == 1
+    assert messages[0].message == ""
+    assert messages[0].tokens == [99]
+
+
+def test_simple_context_output_messages_no_mutation():
+    """Each call to output_messages returns a fresh list; callers can't
+    corrupt internal state."""
+    context = SimpleContext()
+    context.append_output(
+        create_simple_context_output(
+            text="hello",
+            token_ids=[1],
+            prompt_token_ids=[10],
+        )
+    )
+
+    msgs1 = context.output_messages
+    msgs2 = context.output_messages
+    assert msgs1 is not msgs2
+    assert msgs1[0].message == msgs2[0].message
+
+    # Appending more output updates the property
+    context.append_output(
+        create_simple_context_output(
+            text=" world",
+            token_ids=[2],
+            prompt_token_ids=[10],
+        )
+    )
+
+    msgs3 = context.output_messages
+    assert len(msgs3) == 1
+    assert msgs3[0].message == "hello world"
+    assert msgs3[0].tokens == [1, 2]

tests/evals/gsm8k/configs/moe-refactor/Llama-4-Scout-BF16-fi-cutlass.yaml

@@ -5,3 +5,5 @@ num_fewshot: 5
 server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2 --enable-expert-parallel"
 env:
   VLLM_USE_FLASHINFER_MOE_FP16: "1"
+  VLLM_FLASHINFER_MOE_BACKEND: "throughput"
+

tests/evals/gsm8k/configs/moe-refactor/Mixtral-8x7B-BF16-fi-cutlass.yaml

@@ -5,3 +5,4 @@ num_fewshot: 5
 server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2 --enable-expert-parallel"
 env:
   VLLM_USE_FLASHINFER_MOE_FP16: "1"
+  VLLM_FLASHINFER_MOE_BACKEND: "throughput"

tests/kernels/moe/modular_kernel_tools/common.py

@@ -22,6 +22,7 @@ from vllm.distributed import (
 )
 from vllm.forward_context import set_forward_context
 from vllm.model_executor.layers.fused_moe import fused_topk
+from vllm.model_executor.layers.fused_moe.activation import MoEActivation
 from vllm.model_executor.layers.fused_moe.all2all_utils import (
     maybe_make_prepare_finalize,
 )
@@ -585,6 +586,7 @@ def make_modular_kernel(
         tp_size_=get_tensor_model_parallel_world_size(),
         pcp_size_=get_pcp_group().world_size,
         dp_size_=get_dp_group().world_size,
+        sp_size_=1,
         vllm_parallel_config=vllm_config.parallel_config,
     )
 
@@ -594,10 +596,11 @@ def make_modular_kernel(
         hidden_dim=config.K,
         intermediate_size_per_partition=config.N,
         num_local_experts=config.num_local_experts,
+        num_logical_experts=config.E,
         moe_parallel_config=moe_parallel_config,
         in_dtype=config.dtype,
         max_num_tokens=next_power_of_2(config.M),
-        activation="silu",
+        activation=MoEActivation.SILU,
         device=vllm_config.device_config.device,
         routing_method=RoutingMethodType.DeepSeekV3,
     )

tests/kernels/moe/test_cpu_fused_moe.py

@@ -6,6 +6,7 @@ import torch
 
 from tests.kernels.allclose_default import get_default_atol, get_default_rtol
 from vllm._custom_ops import cpu_fused_moe, cpu_prepack_moe_weight
+from vllm.model_executor.layers.fused_moe.activation import MoEActivation
 from vllm.model_executor.layers.fused_moe.cpu_fused_moe import _CPU_MOE_ACT_FN
 from vllm.platforms import current_platform
 from vllm.utils.torch_utils import set_random_seed
@@ -19,7 +20,7 @@ EXPERT_NUM = [
 HIDDEN_DIM = [128, 2880]
 INTERMEDIATE_DIM = [128, 2880]
 BATCH_SIZE = [1, 64, 256]
-ACT = ["silu", "swigluoai"]
+ACT = [MoEActivation.SILU, MoEActivation.SWIGLUOAI]
 USE_BIAS = [True, False]
 ISA = ["amx", "vec"] if torch._C._cpu._is_amx_tile_supported() else ["vec"]
 DTYPE = [torch.bfloat16]
@@ -33,7 +34,7 @@ def ref_fused_moe(
     w2_bias: torch.Tensor | None,
     topk_weights: torch.Tensor,
     topk_ids: torch.Tensor,
-    activation: str,
+    activation: MoEActivation,
 ) -> torch.Tensor:
     len_experts = w13.size(0)
 
@@ -103,7 +104,7 @@ def test_cpu_fused_moe(
     intermediate_size: int,
     use_bias: bool,
     dtype: torch.dtype,
-    act: str,
+    act: MoEActivation,
     isa: str,
 ):
     set_random_seed(0)
@@ -153,7 +154,7 @@ def test_cpu_fused_moe(
         w2_bias,
         topk_weight,
         topk_ids,
-        act,
+        act.value,
         isa,
     )
 

tests/kernels/moe/test_cutlass_moe.py

@@ -12,6 +12,7 @@ from tests.kernels.moe.utils import make_dummy_moe_config
 from vllm import _custom_ops as ops
 from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config
 from vllm.model_executor.layers.fused_moe import fused_experts, fused_topk
+from vllm.model_executor.layers.fused_moe.activation import MoEActivation
 from vllm.model_executor.layers.fused_moe.config import (
     FUSED_MOE_UNQUANTIZED_CONFIG,
     FusedMoEQuantConfig,
@@ -531,7 +532,7 @@ def test_run_cutlass_moe_fp8(
         c_strides1 = torch.full((e,), 2 * n, device="cuda", dtype=torch.int64)
         c_strides2 = torch.full((e,), k, device="cuda", dtype=torch.int64)
 
-        activation = "silu"
+        activation = MoEActivation.SILU
         a1q, a1q_scale = moe_kernel_quantize_input(
             mt.a, mt.a_scale, torch.float8_e4m3fn, per_act_token
         )

tests/kernels/moe/test_deepep_deepgemm_moe.py

@@ -16,6 +16,7 @@ from typing_extensions import ParamSpec
 
 from vllm.config import VllmConfig, set_current_vllm_config
 from vllm.forward_context import set_forward_context
+from vllm.model_executor.layers.fused_moe.activation import MoEActivation
 from vllm.model_executor.layers.fused_moe.config import (
     FusedMoEQuantConfig,
     fp8_w8a8_moe_quant_config,
@@ -324,7 +325,7 @@ def deepep_deepgemm_moe_impl(
             w2=w2,
             topk_weights=test_tensors.topk_weights,
             topk_ids=test_tensors.topk,
-            activation="silu",
+            activation=MoEActivation.SILU,
             global_num_experts=num_experts,
             expert_map=build_expert_map(),
             apply_router_weight_on_input=False,

tests/kernels/moe/test_deepep_moe.py

@@ -15,6 +15,7 @@ from vllm import _custom_ops as ops
 from vllm.config import VllmConfig, set_current_vllm_config
 from vllm.model_executor.layers.activation import SiluAndMul
 from vllm.model_executor.layers.fused_moe import TritonExperts
+from vllm.model_executor.layers.fused_moe.activation import MoEActivation
 from vllm.model_executor.layers.fused_moe.config import (
     FusedMoEQuantConfig,
 )
@@ -260,7 +261,7 @@ def deep_ep_moe_impl(
             w2=w2,
             topk_weights=topk_weights_chunk,
             topk_ids=topk_chunk,
-            activation="silu",
+            activation=MoEActivation.SILU,
             global_num_experts=num_experts,
             expert_map=build_expert_map(),
             apply_router_weight_on_input=False,

tests/kernels/moe/test_flashinfer.py

@@ -7,6 +7,7 @@ import torch
 
 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
 from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config
+from vllm.model_executor.layers.fused_moe.activation import MoEActivation
 from vllm.model_executor.layers.fused_moe.config import (
     FusedMoEConfig,
     FusedMoEParallelConfig,
@@ -93,9 +94,14 @@ class TestData:
 
     @staticmethod
     def make_moe_tensors_8bit(
-        m: int, k: int, n: int, e: int, is_trtllm: bool, activation: str = "silu"
+        m: int,
+        k: int,
+        n: int,
+        e: int,
+        is_trtllm: bool,
+        activation: MoEActivation = MoEActivation.SILU,
     ) -> "TestData":
-        is_gated = activation != "relu2_no_mul"
+        is_gated = activation.is_gated
 
         hidden_states = torch.randn((m, k), device="cuda", dtype=torch.bfloat16) / 10
         w13 = torch.randn(
@@ -194,7 +200,7 @@ def test_flashinfer_per_tensor_moe_fp8_no_graph(
             topk_weights=topk_weights,
             topk_ids=topk_ids,
             inplace=False,
-            activation="silu",
+            activation=MoEActivation.SILU,
             global_num_experts=e,
             expert_map=None,
             apply_router_weight_on_input=True,
@@ -219,21 +225,19 @@ def test_flashinfer_per_tensor_moe_fp8_no_graph(
 @pytest.mark.parametrize("m,n,k", MNK_FACTORS)
 @pytest.mark.parametrize("e", NUM_EXPERTS)
 @pytest.mark.parametrize("topk", TOP_KS)
-@pytest.mark.parametrize("activation", ["silu", "relu2_no_mul"])
+@pytest.mark.parametrize("activation", [MoEActivation.SILU, MoEActivation.RELU2_NO_MUL])
 def test_flashinfer_cutlass_moe_fp8_no_graph(
     m: int,
     n: int,
     k: int,
     e: int,
     topk: int,
-    activation: str,
+    activation: MoEActivation,
     monkeypatch,
     workspace_init,
 ):
     set_random_seed(7)
     monkeypatch.setenv("VLLM_FUSED_MOE_CHUNK_SIZE", "8192")
-    assert activation in ["silu", "relu2_no_mul"]
-    is_act_and_mul = activation == "silu_and_mul"
     with set_current_vllm_config(vllm_config):
         td = TestData.make_moe_tensors_8bit(
             m, k, n, e, is_trtllm=False, activation=activation
@@ -287,11 +291,12 @@ def test_flashinfer_cutlass_moe_fp8_no_graph(
             hidden_dim=k,
             intermediate_size_per_partition=n,
             num_local_experts=e,
+            num_logical_experts=e,
             activation=activation,
             device="cuda",
             moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),
             in_dtype=torch.bfloat16,
-            is_act_and_mul=is_act_and_mul,
+            is_act_and_mul=activation.is_gated,
             routing_method=RoutingMethodType.TopK,
         )
 
@@ -318,3 +323,44 @@ def test_flashinfer_cutlass_moe_fp8_no_graph(
         torch.testing.assert_close(
             output, flashinfer_cutlass_output, atol=5.5e-2, rtol=1e-2
         )
+
+
+@pytest.mark.parametrize(
+    "num_experts,intermediate,hidden",
+    [
+        (8, 2048, 1536),
+        (64, 4096, 4096),
+    ],
+)
+def test_convert_moe_weights_to_flashinfer_trtllm_block_layout(
+    num_experts, intermediate, hidden
+):
+    from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
+        convert_moe_weights_to_flashinfer_trtllm_block_layout,
+    )
+
+    w13 = torch.randn(
+        (num_experts, 2 * intermediate, hidden), dtype=torch.bfloat16, device="cuda"
+    )
+    w2 = torch.randn(
+        (num_experts, hidden, intermediate), dtype=torch.bfloat16, device="cuda"
+    )
+
+    cache: dict[torch.Size, torch.Tensor] = {}
+    w13_converted, w2_converted = convert_moe_weights_to_flashinfer_trtllm_block_layout(
+        cache, w13, w2
+    )
+
+    assert w13_converted.ndim == 4, (
+        f"Expected 4D tensor, got shape {w13_converted.shape}"
+    )
+    assert w2_converted.ndim == 4, f"Expected 4D tensor, got shape {w2_converted.shape}"
+
+    assert w13_converted.numel() == w13.numel(), "W13 element count should be preserved"
+    assert w2_converted.numel() == w2.numel(), "W2 element count should be preserved"
+
+    assert w13_converted.dtype == torch.bfloat16
+    assert w2_converted.dtype == torch.bfloat16
+
+    assert w13_converted.shape[0] == num_experts
+    assert w2_converted.shape[0] == num_experts

tests/kernels/moe/test_flashinfer_moe.py

@@ -13,6 +13,7 @@ from tests.kernels.utils import torch_moe
 from vllm import _custom_ops as ops
 from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config
 from vllm.model_executor.layers.fused_moe import fused_topk
+from vllm.model_executor.layers.fused_moe.activation import MoEActivation
 from vllm.model_executor.layers.fused_moe.config import (
     FusedMoEConfig,
     FusedMoEParallelConfig,
@@ -54,7 +55,7 @@ MNK_FACTORS = [
 @pytest.mark.parametrize("e", [40, 64, 256])
 @pytest.mark.parametrize("topk", [1, 6, 8])
 @pytest.mark.parametrize("dtype", [torch.bfloat16])
-@pytest.mark.parametrize("activation", ["silu_and_mul", "relu2"])
+@pytest.mark.parametrize("activation", [MoEActivation.SILU, MoEActivation.RELU2_NO_MUL])
 @torch.inference_mode()
 def test_flashinfer_fp4_moe_no_graph(
     m: int,
@@ -63,7 +64,7 @@ def test_flashinfer_fp4_moe_no_graph(
     e: int,
     topk: int,
     dtype: torch.dtype,
-    activation: str,
+    activation: MoEActivation,
     workspace_init,
 ):
     set_random_seed(7)
@@ -73,7 +74,7 @@ def test_flashinfer_fp4_moe_no_graph(
         a = torch.randn((m, k), device="cuda", dtype=dtype) / 10
 
         quant_blocksize = 16
-        is_gated_act = activation == "silu_and_mul"
+        is_gated_act = activation.is_gated
 
         w1_q, w2_q, quant_config = make_test_quant_config(
             e,
@@ -97,6 +98,7 @@ def test_flashinfer_fp4_moe_no_graph(
             hidden_dim=k,
             intermediate_size_per_partition=n,
             num_local_experts=e,
+            num_logical_experts=e,
             activation=activation,
             device="cuda",
             moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),
@@ -111,15 +113,13 @@ def test_flashinfer_fp4_moe_no_graph(
             inplace=False,
         )
 
-        fi_activation = {"silu_and_mul": "silu", "relu2": "relu2_no_mul"}[activation]
-
         flashinfer_output = flashinfer_experts(
             hidden_states=a,
             w1=w1_q,
             w2=w2_q,
             topk_weights=topk_weights,
             topk_ids=topk_ids,
-            activation=fi_activation,
+            activation=activation,
         )
 
         # Reference check:

tests/kernels/moe/test_modular_oai_triton_moe.py

@@ -7,6 +7,7 @@ Test modular OAI Triton MoE
 import pytest
 import torch
 
+from vllm.model_executor.layers.fused_moe.activation import MoEActivation
 from vllm.utils.import_utils import has_triton_kernels
 
 if not has_triton_kernels():
@@ -192,7 +193,7 @@ def oai_triton_moe_impl(
         w2=w2,
         topk_weights=topk_weights,
         topk_ids=topk_ids,
-        activation="swigluoai",
+        activation=MoEActivation.SWIGLUOAI,
         global_num_experts=num_experts,
         expert_map=None,
         apply_router_weight_on_input=False,

tests/kernels/moe/test_moe.py

@@ -29,6 +29,7 @@ from vllm.config import VllmConfig, set_current_vllm_config
 from vllm.distributed.parallel_state import init_distributed_environment
 from vllm.forward_context import get_forward_context, set_forward_context
 from vllm.model_executor.layers.fused_moe import (
+    MoEActivation,
     fused_topk,
 )
 from vllm.model_executor.layers.fused_moe.config import (
@@ -1155,7 +1156,10 @@ def test_fused_marlin_moe_with_bias(m):
 @pytest.mark.parametrize("m", [1, 64, 256])
 @pytest.mark.parametrize("n,k", [(1024, 1024), (2048, 2048)])
 @pytest.mark.parametrize("e,topk", [(8, 2), (64, 4)])
-def test_fused_marlin_moe_non_gated(m: int, n: int, k: int, e: int, topk: int):
+@pytest.mark.parametrize("activation", [MoEActivation.RELU2_NO_MUL])
+def test_fused_marlin_moe_non_gated(
+    m: int, n: int, k: int, e: int, topk: int, activation: MoEActivation
+):
     """Test Marlin MoE with non-gated activation (relu2_no_mul).
 
     Non-gated activations like relu2 don't have the gate-up projection pattern,
@@ -1198,7 +1202,7 @@ def test_fused_marlin_moe_non_gated(m: int, n: int, k: int, e: int, topk: int):
             w2_data.w_ref,
             score,
             topk,
-            activation="relu2",
+            activation=activation,
         )
 
     marlin_output = fused_marlin_moe(
@@ -1223,7 +1227,7 @@ def test_fused_marlin_moe_non_gated(m: int, n: int, k: int, e: int, topk: int):
         w2_zeros=w2_data.zeros,
         quant_type_id=quant_type.id,
         is_k_full=is_k_full,
-        activation="relu2_no_mul",
+        activation=activation,
     )
 
     torch.testing.assert_close(marlin_output, torch_output, atol=1e-1, rtol=0)
@@ -1330,9 +1334,18 @@ def test_moe_sum(m: int, topk: int, k: int, dtype: torch.dtype):
 @pytest.mark.parametrize("topk", [2])
 @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
 @pytest.mark.parametrize("with_bias", [False, True])
-@pytest.mark.parametrize("activation", ["silu"])
+@pytest.mark.parametrize("activation", [MoEActivation.SILU])
 @pytest.mark.skipif(not current_platform.is_cpu(), reason="CPU only test")
-def test_cpu_fused_moe_basic(m, n, k, e, topk, dtype, with_bias, activation):
+def test_cpu_fused_moe_basic(
+    m: int,
+    n: int,
+    k: int,
+    e: int,
+    topk: int,
+    dtype: torch.dtype,
+    with_bias: bool,
+    activation: MoEActivation,
+):
     from vllm.model_executor.layers.fused_moe.cpu_fused_moe import CPUFusedMOE
 
     device = "cpu"
@@ -1558,3 +1571,104 @@ def test_batched_fused_marlin_moe(
     marlin_output = br.run(a, kwargs)
 
     torch.testing.assert_close(marlin_output, ref_marlin_output, atol=1e-3, rtol=0)
+
+
+@pytest.mark.parametrize("m,n,k", [(32, 1024, 1024)])
+@pytest.mark.parametrize("e,topk", [(8, 2)])
+@pytest.mark.parametrize("dtype", [torch.bfloat16])
+@pytest.mark.skipif(
+    not current_platform.is_device_capability_family(100),
+    reason="TRTLLM backend test only runs on Blackwell GPUs (SM10x).",
+)
+def test_unquantized_bf16_flashinfer_trtllm_backend(
+    m: int,
+    n: int,
+    k: int,
+    e: int,
+    topk: int,
+    dtype: torch.dtype,
+    monkeypatch,
+    workspace_init,
+):
+    """
+    Test BF16 unquantized MoE with FlashInfer TRTLLM backend.
+    """
+    set_random_seed(7)
+
+    monkeypatch.setenv("VLLM_USE_FLASHINFER_MOE_FP16", "1")
+
+    from vllm.model_executor.layers.fused_moe.config import (
+        FusedMoEConfig,
+        FusedMoEParallelConfig,
+        RoutingMethodType,
+    )
+    from vllm.model_executor.layers.fused_moe.oracle.unquantized import (
+        UnquantizedMoeBackend,
+    )
+    from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import (
+        UnquantizedFusedMoEMethod,
+    )
+
+    # Setup test data
+    a = torch.randn((m, k), device="cuda", dtype=dtype) / 10
+    w1 = torch.randn((e, 2 * n, k), device="cuda", dtype=dtype) / 10
+    w2 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 10
+    router_logits = torch.randn((m, e), device="cuda", dtype=dtype)
+
+    moe_config = FusedMoEConfig(
+        num_experts=e,
+        experts_per_token=topk,
+        hidden_dim=k,
+        intermediate_size_per_partition=n,
+        num_local_experts=e,
+        num_logical_experts=e,
+        activation="silu",
+        device="cuda",
+        moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),
+        in_dtype=dtype,
+        is_act_and_mul=True,
+        routing_method=RoutingMethodType.Renormalize,
+        max_num_tokens=m,
+    )
+
+    with set_current_vllm_config(vllm_config):
+        quant_method = UnquantizedFusedMoEMethod(moe_config)
+
+        # Verify TRTLLM backend was selected
+        assert (
+            quant_method.unquantized_backend == UnquantizedMoeBackend.FLASHINFER_TRTLLM
+        ), f"Expected FLASHINFER_TRTLLM backend, got {quant_method.unquantized_backend}"
+
+        # Verify it's using monolithic path
+        assert quant_method.is_monolithic, (
+            "FLASHINFER_TRTLLM backend should use monolithic forward"
+        )
+        layer = torch.nn.Module()
+        layer.w13_weight = Parameter(w1.clone(), requires_grad=False)
+        layer.w2_weight = Parameter(w2.clone(), requires_grad=False)
+        layer.global_num_experts = e
+        layer.local_num_experts = e
+        layer.top_k = topk
+        layer.num_expert_group = 1
+        layer.topk_group = 1
+        layer.intermediate_size_per_partition = n
+        layer.ep_rank = 0
+        layer.activation = "silu"
+        layer.e_score_correction_bias = None
+        layer.routing_method_type = RoutingMethodType.Renormalize
+
+        quant_method.process_weights_after_loading(layer)
+
+        trtllm_output = quant_method.forward_monolithic_cuda(
+            layer=layer,
+            x=a,
+            router_logits=router_logits,
+        )
+
+        # Compute torch baseline
+        w1_original = w1.clone()
+        w2_original = w2.clone()
+        baseline_output = torch_moe(a, w1_original, w2_original, router_logits, topk)
+
+    close = torch.isclose(trtllm_output, baseline_output, atol=1e-1, rtol=0.85)
+    assert close.float().mean() > 0.925

tests/kernels/moe/test_pplx_cutlass_moe.py

@@ -9,6 +9,7 @@ from tests.kernels.utils import torch_experts
 from vllm import _custom_ops as ops
 from vllm.config import VllmConfig, set_current_vllm_config
 from vllm.model_executor.layers.fused_moe import fused_topk
+from vllm.model_executor.layers.fused_moe.activation import MoEActivation
 from vllm.model_executor.layers.fused_moe.config import (
     FusedMoEConfig,
     FusedMoEParallelConfig,
@@ -147,8 +148,9 @@ def pplx_cutlass_moe(
             hidden_dim=hidden_dim,
             intermediate_size_per_partition=intermediate_dim,
             num_local_experts=num_local_experts,
+            num_logical_experts=num_experts,
             moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),
-            activation="silu",
+            activation=MoEActivation.SILU,
             in_dtype=torch.bfloat16,
             device="cuda",
             routing_method=RoutingMethodType.Llama4,

tests/kernels/moe/test_triton_moe_no_act_mul.py

@@ -11,15 +11,11 @@ import pytest
 import torch
 
 from tests.kernels.moe.utils import make_dummy_moe_config
+from vllm.model_executor.layers.fused_moe.activation import MoEActivation
 from vllm.model_executor.layers.fused_moe.config import (
     FUSED_MOE_UNQUANTIZED_CONFIG,
 )
 from vllm.model_executor.layers.fused_moe.fused_moe import TritonExperts
-from vllm.model_executor.layers.fused_moe.utils import (
-    GELU_NO_MUL,
-    RELU2_NO_MUL,
-    SILU_NO_MUL,
-)
 from vllm.platforms import current_platform
 
 # Test parameters
@@ -28,7 +24,11 @@ N_SIZES = [128, 256]
 K_SIZES = [64, 128]
 TOPK_VALUES = [1, 2]
 NUM_EXPERTS = 8
-NO_MUL_ACTIVATIONS = [SILU_NO_MUL, GELU_NO_MUL, RELU2_NO_MUL]
+NO_MUL_ACTIVATIONS = [
+    MoEActivation.SILU_NO_MUL,
+    MoEActivation.GELU_NO_MUL,
+    MoEActivation.RELU2_NO_MUL,
+]
 
 
 def make_test_tensors(
@@ -73,7 +73,7 @@ def test_triton_experts_no_mul_activation(
     n: int,
     k: int,
     topk: int,
-    activation: str,
+    activation: MoEActivation,
 ):
     hidden_states, w1, w2, topk_weights, topk_ids = make_test_tensors(
         m, n, k, NUM_EXPERTS, topk
@@ -161,11 +161,11 @@ def test_workspace_shapes_no_mul_vs_gated():
     )
 
     ws1_no_mul, _, out_no_mul = experts.workspace_shapes(
-        M, N, K, topk, 8, 8, None, SILU_NO_MUL
+        M, N, K, topk, 8, 8, None, MoEActivation.SILU_NO_MUL
     )
 
     ws1_gated, _, out_gated = experts.workspace_shapes(
-        M, N, K, topk, 8, 8, None, "silu"
+        M, N, K, topk, 8, 8, None, MoEActivation.SILU
     )
 
     # For no_mul: activation_out_dim = N
@@ -202,10 +202,10 @@ def test_adjust_n_for_activation():
     N = 256
 
     # Gated activations should return N // 2
-    assert experts.adjust_N_for_activation(N, "silu") == N // 2
-    assert experts.adjust_N_for_activation(N, "gelu") == N // 2
+    assert experts.adjust_N_for_activation(N, MoEActivation.SILU) == N // 2
+    assert experts.adjust_N_for_activation(N, MoEActivation.GELU) == N // 2
 
     # Non-gated activations should return N
-    assert experts.adjust_N_for_activation(N, SILU_NO_MUL) == N
-    assert experts.adjust_N_for_activation(N, GELU_NO_MUL) == N
-    assert experts.adjust_N_for_activation(N, RELU2_NO_MUL) == N
+    assert experts.adjust_N_for_activation(N, MoEActivation.SILU_NO_MUL) == N
+    assert experts.adjust_N_for_activation(N, MoEActivation.GELU_NO_MUL) == N
+    assert experts.adjust_N_for_activation(N, MoEActivation.RELU2_NO_MUL) == N

tests/kernels/moe/test_unquantized_backend_selection.py

@@ -0,0 +1,132 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from unittest.mock import patch
+
+import pytest
+
+from tests.kernels.moe.utils import make_dummy_moe_config
+from vllm.model_executor.layers.fused_moe.oracle.unquantized import (
+    UnquantizedMoeBackend,
+    select_unquantized_moe_backend,
+)
+
+
+@pytest.mark.parametrize(
+    "platform_method,expected_backend",
+    [
+        ("is_cuda", UnquantizedMoeBackend.TRITON),  # Default CUDA without FlashInfer
+        ("is_rocm", UnquantizedMoeBackend.TRITON),
+        ("is_cpu", UnquantizedMoeBackend.CPU),
+        ("is_xpu", UnquantizedMoeBackend.XPU),
+        ("is_tpu", UnquantizedMoeBackend.TPU),
+        ("is_out_of_tree", UnquantizedMoeBackend.OOT),
+    ],
+)
+@patch(
+    "vllm.model_executor.layers.fused_moe.oracle.unquantized.has_flashinfer",
+    return_value=False,
+)
+def test_select_default_backend_by_platform(
+    mock_has_flashinfer,
+    monkeypatch,
+    platform_method,
+    expected_backend,
+):
+    """Test backend selection for different platforms."""
+    with patch(
+        "vllm.model_executor.layers.fused_moe.oracle.unquantized.current_platform"
+    ) as mock_platform:
+        # Set all platform checks to False
+        mock_platform.is_cuda.return_value = False
+        mock_platform.is_rocm.return_value = False
+        mock_platform.is_cpu.return_value = False
+        mock_platform.is_xpu.return_value = False
+        mock_platform.is_tpu.return_value = False
+        mock_platform.is_out_of_tree.return_value = False
+
+        # Set only the specified platform to True
+        getattr(mock_platform, platform_method).return_value = True
+
+        moe_config = make_dummy_moe_config()
+        selected_backend = select_unquantized_moe_backend(
+            moe_config=moe_config,
+            use_ep=False,
+            use_dp=False,
+        )
+
+        assert selected_backend == expected_backend
+
+
+@patch(
+    "vllm.model_executor.layers.fused_moe.oracle.unquantized.has_flashinfer",
+    return_value=True,
+)
+@patch(
+    "vllm.model_executor.layers.fused_moe.oracle.unquantized.is_supported_config_trtllm_bf16",
+    return_value=(True, None),
+)
+def test_select_cuda_flashinfer_trtllm_backend(
+    mock_has_flashinfer, mock_is_supported_trtllm, monkeypatch
+):
+    """Test CUDA backend selection when FlashInfer TRTLLM is available and enabled."""
+    with patch(
+        "vllm.model_executor.layers.fused_moe.oracle.unquantized.current_platform"
+    ) as mock_platform:
+        # Set as CUDA platform
+        mock_platform.is_cuda.return_value = True
+        mock_platform.is_rocm.return_value = False
+        mock_platform.is_cpu.return_value = False
+        mock_platform.is_xpu.return_value = False
+        mock_platform.is_tpu.return_value = False
+        mock_platform.is_out_of_tree.return_value = False
+
+        monkeypatch.setenv("VLLM_USE_FLASHINFER_MOE_FP16", "1")
+
+        moe_config = make_dummy_moe_config()
+
+        selected_backend = select_unquantized_moe_backend(
+            moe_config=moe_config,
+            use_ep=True,
+            use_dp=False,
+        )
+
+        assert selected_backend == UnquantizedMoeBackend.FLASHINFER_TRTLLM
+
+
+@patch(
+    "vllm.model_executor.layers.fused_moe.oracle.unquantized.has_flashinfer",
+    return_value=True,
+)
+@patch(
+    "vllm.model_executor.layers.fused_moe.oracle.unquantized.is_supported_config_trtllm_bf16",
+    return_value=(False, None),
+)
+def test_select_cuda_flashinfer_cutlass_backend(
+    mock_has_flashinfer, mock_is_supported_trtllm, monkeypatch
+):
+    """Test CUDA backend selection when FlashInfer TRTLLM is not available
+    and FlashInfer CUTLASS is available."""
+    with patch(
+        "vllm.model_executor.layers.fused_moe.oracle.unquantized.current_platform"
+    ) as mock_platform:
+        # Set as CUDA platform with Hopper capability
+        mock_platform.is_cuda.return_value = True
+        mock_platform.is_rocm.return_value = False
+        mock_platform.is_cpu.return_value = False
+        mock_platform.is_xpu.return_value = False
+        mock_platform.is_tpu.return_value = False
+        mock_platform.is_out_of_tree.return_value = False
+        mock_platform.has_device_capability.return_value = True  # SM90+
+
+        # Enable FlashInfer via env var
+        monkeypatch.setenv("VLLM_USE_FLASHINFER_MOE_FP16", "1")
+
+        moe_config = make_dummy_moe_config()
+
+        selected_backend = select_unquantized_moe_backend(
+            moe_config=moe_config,
+            use_ep=True,  # CUTLASS requires EP
+            use_dp=False,  # CUTLASS doesn't support DP
+        )
+
+        assert selected_backend == UnquantizedMoeBackend.FLASHINFER_CUTLASS

tests/kernels/moe/utils.py

@@ -12,6 +12,7 @@ from vllm.model_executor.layers.fused_moe import (
     fused_experts,
     fused_topk,
 )
+from vllm.model_executor.layers.fused_moe.activation import MoEActivation
 from vllm.model_executor.layers.fused_moe.config import (
     FusedMoEConfig,
     FusedMoEParallelConfig,
@@ -52,8 +53,9 @@ def make_dummy_moe_config(
         hidden_dim=hidden_dim,
         intermediate_size_per_partition=intermediate_size_per_partition,
         num_local_experts=num_experts,
+        num_logical_experts=num_experts,
         moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),
-        activation="silu",
+        activation=MoEActivation.SILU,
         in_dtype=in_dtype,
         device="cuda",
         routing_method=RoutingMethodType.TopK,

tests/kernels/quantization/test_rocm_skinny_gemms.py

@@ -270,6 +270,9 @@ def test_rocm_wvsplitk_fp8_kernel(
     out = ops.wvSplitKQ(B, A, dtype, scale_a, scale_b, get_cu_count(), BIAS)
 
     if xnorm:
-        assert torch.allclose(out, ref_out, atol=1e-3, rtol=1e-8)
+        torch.testing.assert_close(out, ref_out, atol=1e-3, rtol=1e-8)
+    elif k >= 32 * 1024:
+        # wider pytrch thresh for large-K & no xnorm
+        torch.testing.assert_close(out, ref_out, atol=0.07, rtol=5e-2)
     else:
-        assert torch.allclose(out, ref_out, 0.01)
+        torch.testing.assert_close(out, ref_out, atol=0.01, rtol=0.01)

tests/kernels/test_top_k_per_row.py

@@ -275,3 +275,114 @@ def test_top_k_per_row_decode_large_vocab_size(clean_logits: bool) -> None:
     _run_top_k_per_row_decode_test(
         top_k, batch_size, next_n, vocab_size, clean_logits, data_generation
     )
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA")
+@pytest.mark.parametrize("clean_logits", [True, False])
+@torch.inference_mode()
+def test_deepseek_hybrid_topk(clean_logits: bool) -> None:
+    torch.set_default_device("cuda:0")
+
+    top_k = 2048
+
+    # Test case 1: Short sequences (< 8192)
+    batch_size_short = 4
+    next_n = 1
+    num_rows_short = batch_size_short * next_n
+
+    # Create sequences with max length < 8192
+    seq_lens_short = torch.randint(
+        4000, 8000, (batch_size_short,), dtype=torch.int32, device="cuda"
+    )
+
+    row_starts_short = torch.zeros(num_rows_short, dtype=torch.int32, device="cuda")
+    row_indices_short = torch.arange(num_rows_short, device="cuda") // next_n
+    next_n_offset_short = torch.arange(num_rows_short, device="cuda") % next_n
+    row_ends_short = (
+        seq_lens_short[row_indices_short] - next_n + next_n_offset_short + 1
+    )
+
+    logits_short = create_random_logits(
+        row_starts_short, row_ends_short, torch.float32, 42, clean_logits, "random"
+    )
+
+    indices_vllm = torch.empty(
+        (num_rows_short, top_k), dtype=torch.int32, device="cuda"
+    )
+
+    # Use vllm's kernel for short sequences
+    torch.ops._C.top_k_per_row_decode(
+        logits_short,
+        next_n,
+        seq_lens_short,
+        indices_vllm,
+        num_rows_short,
+        logits_short.stride(0),
+        logits_short.stride(1),
+        top_k,
+    )
+
+    # Test case 2: Long sequences (>= 8192) - should use large_context_topk kernel
+    batch_size_long = 4
+    num_rows_long = batch_size_long * next_n
+
+    # Create sequences with max length >= 8192
+    seq_lens_long = torch.randint(
+        8192, 16384, (batch_size_long,), dtype=torch.int32, device="cuda"
+    )
+
+    row_starts_long = torch.zeros(num_rows_long, dtype=torch.int32, device="cuda")
+    row_indices_long = torch.arange(num_rows_long, device="cuda") // next_n
+    next_n_offset_long = torch.arange(num_rows_long, device="cuda") % next_n
+    row_ends_long = seq_lens_long[row_indices_long] - next_n + next_n_offset_long + 1
+
+    logits_long = create_random_logits(
+        row_starts_long, row_ends_long, torch.float32, 43, clean_logits, "random"
+    )
+
+    indices = torch.empty((num_rows_long, top_k), dtype=torch.int32, device="cuda")
+
+    # Use large_context_topk kernel for long sequences
+    if next_n == 1:
+        lengths = seq_lens_long
+    else:
+        offsets = torch.arange(next_n, device=logits_long.device, dtype=torch.int32)
+        lengths = (seq_lens_long.unsqueeze(1) - next_n + 1 + offsets).flatten()
+
+    torch.ops._C.large_context_topk(
+        logits_long,
+        indices,
+        lengths,
+        None,
+    )
+
+    torch_indices_short = torch.empty(
+        (num_rows_short, top_k), dtype=torch.int32, device="cuda"
+    )
+    for i in range(num_rows_short):
+        row_end = int(row_ends_short[i])
+        k_i = min(top_k, row_end)
+        idx = logits_short[i, :row_end].topk(k_i, dim=-1)[1]
+        torch_indices_short[i, :k_i] = idx
+
+    assert compare_top_k_results(
+        logits_short,
+        indices_vllm,
+        torch_indices_short,
+        row_starts_short,
+        row_ends_short,
+        top_k,
+    ), "top_k_per_row_decode kernel (short sequences) doesn't match torch.topk"
+
+    torch_indices_long = torch.empty(
+        (num_rows_long, top_k), dtype=torch.int32, device="cuda"
+    )
+    for i in range(num_rows_long):
+        row_end = int(row_ends_long[i])
+        k_i = min(top_k, row_end)
+        idx = logits_long[i, :row_end].topk(k_i, dim=-1)[1]
+        torch_indices_long[i, :k_i] = idx
+
+    assert compare_top_k_results(
+        logits_long, indices, torch_indices_long, row_starts_long, row_ends_long, top_k
+    ), "large_context_topk kernel (long sequences) doesn't match torch.topk"

tests/kernels/utils.py

@@ -15,6 +15,7 @@ from torch._prims_common import TensorLikeType
 from tests.kernels.quant_utils import native_w8a8_block_matmul
 from vllm.model_executor.custom_op import op_registry
 from vllm.model_executor.layers.activation import SiluAndMul
+from vllm.model_executor.layers.fused_moe.activation import MoEActivation
 from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input
 from vllm.utils.torch_utils import make_tensor_with_pad
 from vllm.v1.attention.backend import AttentionType
@@ -840,7 +841,7 @@ def torch_experts(
     per_act_token_quant=False,
     block_shape: list[int] | None = None,
     apply_router_weights_on_input: bool = False,
-    activation: str = "silu_and_mul",
+    activation: MoEActivation = MoEActivation.SILU,
 ) -> torch.Tensor:
     assert (
         global_num_experts == -1
@@ -883,7 +884,7 @@ def torch_experts(
 
     f32 = torch.float32
 
-    act = op_registry[activation]
+    act = op_registry[activation.custom_op_name]
 
     for i in range(num_experts):
         mask = topk_ids == i
@@ -973,7 +974,7 @@ def torch_moe(
     b_bias2: torch.Tensor | None = None,
     global_num_experts: int = -1,
     expert_map: torch.Tensor | None = None,
-    activation: str = "silu_and_mul",
+    activation: MoEActivation = MoEActivation.SILU,
 ) -> torch.Tensor:
     score = torch.softmax(score, dim=-1, dtype=torch.float32)
     topk_weight, topk_ids = torch.topk(score, topk)

tests/models/language/generation/test_hybrid.py

@@ -7,6 +7,7 @@ import pytest
 
 from tests.models.registry import HF_EXAMPLE_MODELS
 from tests.utils import multi_gpu_test
+from vllm import LLM
 from vllm.engine.arg_utils import EngineArgs
 from vllm.platforms import current_platform
 from vllm.sampling_params import SamplingParams
@@ -769,3 +770,30 @@ def test_apc_multiple_prompts_partial_cached_outputs(
             name_0="vllm_no_cache",
             name_1=f"vllm_cache_it_{r_idx + 1}",
         )
+
+
+# we have to use a real large model to get reasonable results
+# the model can't be a hybrid model as we need block_size 16
+@pytest.mark.parametrize("model", ["tiiuae/falcon-mamba-7b"])
+def test_apc_common_prefix_same_batch(
+    model: str,
+    monkeypatch,
+) -> None:
+    # Required to put the two requests in the same batch
+    monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
+    llm = LLM(
+        model=model,
+        enforce_eager=True,
+        block_size=16,
+        mamba_block_size=16,
+        enable_prefix_caching=True,
+        seed=42,
+    )
+    prompts = [
+        "hello what is one plus one what is one plus one what is one plus one the answer is",  # noqa: E501
+        "hello what is one plus one what is one plus one what is one plus one the answer is",  # noqa: E501
+    ]
+    sampling_params = SamplingParams(temperature=0.8, top_p=0.95, max_tokens=20)
+    outputs = llm.generate(prompts, sampling_params)
+    for output in outputs:
+        assert "two" in output.outputs[0].text

tests/models/multimodal/generation/test_common.py

@@ -377,7 +377,7 @@ VLM_TEST_SETTINGS = {
         use_tokenizer_eos=True,
         vllm_output_post_proc=model_utils.fuyu_vllm_to_hf_output,
         num_logprobs=10,
-        image_size_factors=[(), (0.25,), (0.25, 0.25, 0.25), (0.25, 0.2, 0.15)],
+        image_size_factors=[(0.25,), (0.25, 0.25, 0.25), (0.25, 0.2, 0.15)],
         marks=[large_gpu_mark(min_gb=32)],
     ),
     "gemma3": VLMTestInfo(
@@ -437,7 +437,7 @@ VLM_TEST_SETTINGS = {
         max_num_seqs=2,
         get_stop_token_ids=lambda tok: [151329, 151336, 151338],
         num_logprobs=10,
-        image_size_factors=[(), (0.25,), (0.25, 0.25, 0.25), (0.25, 0.2, 0.15)],
+        image_size_factors=[(0.25,), (0.25, 0.25, 0.25), (0.25, 0.2, 0.15)],
         auto_cls=AutoModelForImageTextToText,
         marks=[large_gpu_mark(min_gb=32)],
     ),
@@ -468,7 +468,7 @@ VLM_TEST_SETTINGS = {
         max_num_seqs=2,
         get_stop_token_ids=lambda tok: [151329, 151336, 151338],
         num_logprobs=10,
-        image_size_factors=[(), (0.25,), (0.25, 0.25, 0.25), (0.25, 0.2, 0.15)],
+        image_size_factors=[(0.25,), (0.25, 0.25, 0.25), (0.25, 0.2, 0.15)],
         auto_cls=AutoModelForImageTextToText,
         marks=[large_gpu_mark(min_gb=32)],
     ),

tests/models/multimodal/processing/test_common.py

@@ -108,6 +108,7 @@ _ADD_SPECIAL_TOKENS_OVERRIDES = {
     "paligemma": False,
     "ultravox": False,
     "whisper": False,
+    "lfm2_vl": False,
 }
 
 _IGNORE_MM_KEYS = {

tests/models/registry.py

@@ -275,6 +275,9 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
         "zai-org/GLM-4.7-Flash",
         min_transformers_version="5.0.0",
     ),
+    "GlmMoeDsaForCausalLM": _HfExamplesInfo(
+        "zai-org/GLM-5", min_transformers_version="5.0.1", is_available_online=False
+    ),
     "GPT2LMHeadModel": _HfExamplesInfo("openai-community/gpt2", {"alias": "gpt2"}),
     "GPTBigCodeForCausalLM": _HfExamplesInfo(
         "bigcode/starcoder",
@@ -710,6 +713,10 @@ _MULTIMODAL_EXAMPLE_MODELS = {
         "baidu/ERNIE-4.5-VL-28B-A3B-PT",
         trust_remote_code=True,
     ),
+    "FunASRForConditionalGeneration": _HfExamplesInfo(
+        "allendou/Fun-ASR-Nano-2512-vllm",
+        is_available_online=False,
+    ),
     "FunAudioChatForConditionalGeneration": _HfExamplesInfo(
         "funaudiochat", is_available_online=False
     ),
@@ -718,7 +725,6 @@ _MULTIMODAL_EXAMPLE_MODELS = {
     "Gemma3nForConditionalGeneration": _HfExamplesInfo("google/gemma-3n-E2B-it"),
     "GlmAsrForConditionalGeneration": _HfExamplesInfo(
         "zai-org/GLM-ASR-Nano-2512",
-        trust_remote_code=True,
         min_transformers_version="5.0.0",
     ),
     "GraniteVision": _HfExamplesInfo("ibm-granite/granite-vision-3.3-2b"),
@@ -967,6 +973,26 @@ _MULTIMODAL_EXAMPLE_MODELS = {
         max_model_len=4096,
         min_transformers_version="4.57",
     ),
+    "Qwen3_5ForConditionalGeneration": _HfExamplesInfo(
+        "Qwen/Qwen3.5-9B-Instruct",
+        max_model_len=4096,
+        min_transformers_version="5.1.0",
+    ),
+    "Qwen3_5MoeForConditionalGeneration": _HfExamplesInfo(
+        "Qwen/Qwen3.5-35B-A3B-Instruct",
+        max_model_len=4096,
+        min_transformers_version="5.1.0",
+    ),
+    "Qwen3_5MTP": _HfExamplesInfo(
+        "Qwen/Qwen3.5-9B-Instruct",
+        speculative_model="Qwen/Qwen3.5-9B-Instruct",
+        min_transformers_version="5.1.0",
+    ),
+    "Qwen3_5MoeMTP": _HfExamplesInfo(
+        "Qwen/Qwen3.5-35B-A3B-Instruct",
+        speculative_model="Qwen/Qwen3.5-35B-A3B-Instruct",
+        min_transformers_version="5.1.0",
+    ),
     "Qwen3OmniMoeForConditionalGeneration": _HfExamplesInfo(
         "Qwen/Qwen3-Omni-30B-A3B-Instruct",
         max_model_len=4096,