Done

Signed-off-by: Jee Jee Li <pandaleefree@gmail.com>

.buildkite/test-amd.yaml

@@ -388,10 +388,10 @@ steps:
     - python3 basic/offline_inference/embed.py
     - python3 basic/offline_inference/score.py
     # Multi-modal models
-    - python3 generate/multimodal/audio_language_offline.py --seed 0
-    - python3 generate/multimodal/vision_language_offline.py --seed 0
-    - python3 generate/multimodal/vision_language_multi_image_offline.py --seed 0
-    - python3 generate/multimodal/encoder_decoder_multimodal_offline.py --model-type whisper --seed 0
+    - python3 offline_inference/audio_language.py --seed 0
+    - python3 offline_inference/vision_language.py --seed 0
+    - python3 offline_inference/vision_language_multi_image.py --seed 0
+    - python3 offline_inference/encoder_decoder_multimodal.py --model-type whisper --seed 0
     # Pooling models
     - python3 pooling/embed/vision_embedding_offline.py --seed 0
     # Features demo
@@ -1647,10 +1647,10 @@ steps:
     - python3 basic/offline_inference/embed.py
     - python3 basic/offline_inference/score.py
     # Multi-modal models
-    - python3 generate/multimodal/audio_language_offline.py --seed 0
-    - python3 generate/multimodal/vision_language_offline.py --seed 0
-    - python3 generate/multimodal/vision_language_multi_image_offline.py --seed 0
-    - python3 generate/multimodal/encoder_decoder_multimodal_offline.py --model-type whisper --seed 0
+    - python3 offline_inference/audio_language.py --seed 0
+    - python3 offline_inference/vision_language.py --seed 0
+    - python3 offline_inference/vision_language_multi_image.py --seed 0
+    - python3 offline_inference/encoder_decoder_multimodal.py --model-type whisper --seed 0
     # Pooling models
     - python3 pooling/embed/vision_embedding_offline.py --seed 0
     # Features demo
@@ -1951,8 +1951,8 @@ steps:
   - pytest -v -s tests/models/multimodal/processing/
   - pytest -v -s tests/models/multimodal/test_mapping.py
   - python3 examples/basic/offline_inference/chat.py
-  - python3 examples/generate/multimodal/vision_language_offline.py --model-type qwen2_5_vl
-  - VLLM_WORKER_MULTIPROC_METHOD=spawn python3 examples/generate/multimodal/audio_language_offline.py --model-type whisper
+  - python3 examples/offline_inference/vision_language.py --model-type qwen2_5_vl
+  - VLLM_WORKER_MULTIPROC_METHOD=spawn python3 examples/offline_inference/audio_language.py --model-type whisper
 
 #-------------------------------------------------------  mi300 · quantization  --------------------------------------------------------#
 
@@ -2930,10 +2930,10 @@ steps:
   - python3 basic/offline_inference/embed.py
   - python3 basic/offline_inference/score.py
   # Multi-modal models
-  - python3 generate/multimodal/audio_language_offline.py --seed 0
-  - python3 generate/multimodal/vision_language_offline.py --seed 0
-  - python3 generate/multimodal/vision_language_multi_image_offline.py --seed 0
-  - python3 generate/multimodal/encoder_decoder_multimodal_offline.py --model-type whisper --seed 0
+  - python3 offline_inference/audio_language.py --seed 0
+  - python3 offline_inference/vision_language.py --seed 0
+  - python3 offline_inference/vision_language_multi_image.py --seed 0
+  - python3 offline_inference/encoder_decoder_multimodal.py --model-type whisper --seed 0
   # Pooling models
   - python3 pooling/embed/vision_embedding_offline.py --seed 0
   # Features demo

.buildkite/test_areas/kernels.yaml

@@ -95,13 +95,11 @@ steps:
   - tests/kernels/moe/test_deepgemm.py
   - tests/kernels/moe/test_batched_deepgemm.py
   - tests/kernels/attention/test_deepgemm_attention.py
-  - tests/quantization/test_cutlass_w4a16.py
   commands:
     - pytest -v -s kernels/quantization/test_block_fp8.py
     - pytest -v -s kernels/moe/test_deepgemm.py
     - pytest -v -s kernels/moe/test_batched_deepgemm.py
     - pytest -v -s kernels/attention/test_deepgemm_attention.py
-    - pytest -v -s quantization/test_cutlass_w4a16.py
 
 - label: Kernels (B200)
   timeout_in_minutes: 30

.buildkite/test_areas/misc.yaml

@@ -113,10 +113,10 @@ steps:
     - python3 basic/offline_inference/embed.py
     - python3 basic/offline_inference/score.py
     # for multi-modal models
-    - python3 generate/multimodal/audio_language_offline.py --seed 0
-    - python3 generate/multimodal/vision_language_offline.py --seed 0
-    - python3 generate/multimodal/vision_language_multi_image_offline.py --seed 0
-    - python3 generate/multimodal/encoder_decoder_multimodal_offline.py --model-type whisper --seed 0
+    - python3 offline_inference/audio_language.py --seed 0
+    - python3 offline_inference/vision_language.py --seed 0
+    - python3 offline_inference/vision_language_multi_image.py --seed 0
+    - python3 offline_inference/encoder_decoder_multimodal.py --model-type whisper --seed 0
      # for pooling models
     - python3 pooling/embed/vision_embedding_offline.py --seed 0
     # for features demo

.buildkite/test_areas/model_runner_v2.yaml

@@ -44,10 +44,10 @@ steps:
     #- python3 basic/offline_inference/generate.py --model meta-llama/Llama-2-13b-chat-hf --cpu-offload-gb 10  # TODO
     #- python3 basic/offline_inference/embed.py   # TODO
     # for multi-modal models
-    - python3 generate/multimodal/audio_language_offline.py --seed 0
-    - python3 generate/multimodal/vision_language_offline.py --seed 0
-    - python3 generate/multimodal/vision_language_multi_image_offline.py --seed 0
-    - python3 generate/multimodal/encoder_decoder_multimodal_offline.py --model-type whisper --seed 0
+    - python3 offline_inference/audio_language.py --seed 0
+    - python3 offline_inference/vision_language.py --seed 0
+    - python3 offline_inference/vision_language_multi_image.py --seed 0
+    - python3 offline_inference/encoder_decoder_multimodal.py --model-type whisper --seed 0
     # for pooling models
     - python3 pooling/embed/vision_embedding_offline.py --seed 0
     # for features demo

.buildkite/test_areas/models_basic.yaml

@@ -69,9 +69,9 @@ steps:
     - pytest -v -s tests/models/multimodal/processing/
     - pytest -v -s tests/models/multimodal/test_mapping.py
     - python3 examples/basic/offline_inference/chat.py
-    - python3 examples/generate/multimodal/vision_language_offline.py --model-type qwen2_5_vl
+    - python3 examples/offline_inference/vision_language.py --model-type qwen2_5_vl
     # Whisper needs spawn method to avoid deadlock
-    - VLLM_WORKER_MULTIPROC_METHOD=spawn python3 examples/generate/multimodal/audio_language_offline.py --model-type whisper
+    - VLLM_WORKER_MULTIPROC_METHOD=spawn python3 examples/offline_inference/audio_language.py --model-type whisper
 
 - label: Transformers Backward Compatibility Models Test
   working_dir: "/vllm-workspace/"
@@ -83,7 +83,7 @@ steps:
     - pytest -v -s tests/models/test_transformers.py
     - pytest -v -s tests/models/multimodal/processing/
     - pytest -v -s tests/models/multimodal/test_mapping.py
-    - python3 examples/basic/offline_inference/chat.py
-    - python3 examples/generate/multimodal/vision_language_offline.py --model-type qwen2_5_vl
+    - python3 examples/offline_inference/basic/chat.py
+    - python3 examples/offline_inference/vision_language.py --model-type qwen2_5_vl
     # Whisper needs spawn method to avoid deadlock
-    - VLLM_WORKER_MULTIPROC_METHOD=spawn python3 examples/generate/multimodal/audio_language_offline.py --model-type whisper
+    - VLLM_WORKER_MULTIPROC_METHOD=spawn python3 examples/offline_inference/audio_language.py --model-type whisper

.github/mergify.yml

@@ -389,7 +389,11 @@ pull_request_rules:
       - files~=^tests/entrypoints/anthropic/.*tool.*
       - files~=^vllm/tool_parsers/
       - files=docs/features/tool_calling.md
-      - files~=^examples/tool_calling/
+      - files~=^examples/tool_chat_*
+      - files=examples/offline_inference/chat_with_tools.py
+      - files=examples/online_serving/openai_chat_completion_client_with_tools_required.py
+      - files=examples/online_serving/openai_chat_completion_tool_calls_with_reasoning.py
+      - files=examples/online_serving/openai_chat_completion_client_with_tools.py
   actions:
     label:
       add:

CMakeLists.txt

@@ -310,9 +310,7 @@ set(VLLM_EXT_SRC
   "csrc/torch_bindings.cpp")
 
 if(VLLM_GPU_LANG STREQUAL "CUDA")
-  list(APPEND VLLM_EXT_SRC
-    "csrc/minimax_reduce_rms_kernel.cu"
-    "csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu")
+  list(APPEND VLLM_EXT_SRC "csrc/minimax_reduce_rms_kernel.cu")
 
   SET(CUTLASS_ENABLE_HEADERS_ONLY ON CACHE BOOL "Enable only the header library")
 
@@ -1053,8 +1051,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
   list(APPEND VLLM_MOE_EXT_SRC
     "csrc/moe/moe_wna16.cu"
     "csrc/moe/grouped_topk_kernels.cu"
-    "csrc/moe/router_gemm.cu"
-    "csrc/moe/topk_softplus_sqrt_kernels.cu")
+    "csrc/moe/router_gemm.cu")
 endif()
 
 if(VLLM_GPU_LANG STREQUAL "CUDA")

benchmarks/kernels/benchmark_vit_fp8_attn.py

@@ -1,324 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-# Benchmarks FP8 vs BF16 ViT attention via FlashInfer cuDNN backend.
-#
-# == Usage Examples ==
-#
-# Benchmark mode (default, FlashInfer CUDAGraph Bench)
-#   python3 benchmark_vit_fp8_attn.py
-#
-# Profile mode (PyTorch profiler, saves TensorBoard traces):
-#   python3 benchmark_vit_fp8_attn.py --profile
-#   python3 benchmark_vit_fp8_attn.py --profile --profile-output-dir ./profile_traces
-#
-# Custom seq_lens:
-#   python3 benchmark_vit_fp8_attn.py --seq-lens 4096 8192 16384
-
-from functools import partial
-
-import numpy as np
-import torch
-from torch.profiler import ProfilerActivity, profile, record_function
-
-from vllm.utils.argparse_utils import FlexibleArgumentParser
-
-# Qwen3-VL defaults
-NUM_HEADS = 16
-HEAD_DIM = 72
-DEFAULT_SEQ_LENS = [2304, 4096, 8192, 16384]
-
-
-def _setup_fp8_attention(num_heads: int, head_dim: int) -> tuple:
-    """Create FP8 and BF16 attention modules + workspace."""
-    from types import SimpleNamespace
-    from unittest.mock import patch
-
-    from vllm.config import VllmConfig, set_current_vllm_config
-    from vllm.config.multimodal import MultiModalConfig
-    from vllm.model_executor.layers.attention.mm_encoder_attention import (
-        MMEncoderAttention,
-        _get_flashinfer_workspace_buffer,
-    )
-    from vllm.v1.attention.backends.registry import AttentionBackendEnum
-
-    old_dtype = torch.get_default_dtype()
-    torch.set_default_dtype(torch.bfloat16)
-
-    backend_patch = patch(
-        "vllm.model_executor.layers.attention.mm_encoder_attention"
-        ".get_vit_attn_backend",
-        return_value=AttentionBackendEnum.FLASHINFER,
-    )
-
-    # FP8 attention
-    mm_config_fp8 = MultiModalConfig(mm_encoder_attn_dtype="fp8")
-    vllm_config_fp8 = VllmConfig()
-    vllm_config_fp8.model_config = SimpleNamespace(multimodal_config=mm_config_fp8)
-    with set_current_vllm_config(vllm_config_fp8), backend_patch:
-        attn_fp8 = MMEncoderAttention(
-            num_heads=num_heads,
-            head_size=head_dim,
-            prefix="visual.blocks.0.attn",
-        ).to("cuda")
-
-    # BF16 attention (no FP8)
-    with set_current_vllm_config(VllmConfig()), backend_patch:
-        attn_bf16 = MMEncoderAttention(
-            num_heads=num_heads,
-            head_size=head_dim,
-            prefix="visual.blocks.0.attn",
-        ).to("cuda")
-
-    torch.set_default_dtype(old_dtype)
-
-    workspace = _get_flashinfer_workspace_buffer()
-    return attn_fp8, attn_bf16, workspace
-
-
-def _build_meta(
-    seq_len: int,
-    num_heads: int,
-    head_dim: int,
-    fp8: bool,
-):
-    """Build cu_seqlens, max_seqlen, sequence_lengths."""
-    from vllm.model_executor.layers.attention.mm_encoder_attention import (
-        MMEncoderAttention,
-    )
-    from vllm.utils.math_utils import round_up
-    from vllm.v1.attention.backends.registry import AttentionBackendEnum
-
-    cu_np = np.array([0, seq_len], dtype=np.int32)
-    fp8_padded = num_heads * round_up(head_dim, 16) if fp8 else None
-
-    seq_lengths = MMEncoderAttention.maybe_compute_seq_lens(
-        AttentionBackendEnum.FLASHINFER, cu_np, torch.device("cuda")
-    )
-    max_seqlen = torch.tensor(
-        MMEncoderAttention.compute_max_seqlen(AttentionBackendEnum.FLASHINFER, cu_np),
-        dtype=torch.int32,
-    )
-    cu_seqlens = MMEncoderAttention.maybe_recompute_cu_seqlens(
-        AttentionBackendEnum.FLASHINFER,
-        cu_np,
-        num_heads * head_dim,
-        1,
-        torch.device("cuda"),
-        fp8_padded_hidden_size=fp8_padded,
-    )
-    return cu_seqlens, max_seqlen, seq_lengths
-
-
-def run_benchmark(
-    seq_lens: list[int],
-    num_heads: int,
-    head_dim: int,
-    method: str,
-):
-    """Benchmark FP8 vs BF16 attention across seq_lens.
-
-    Uses FlashInfer GPU-level timing to measure pure kernel time,
-    excluding CPU launch overhead.
-    """
-    if method == "cupti":
-        from flashinfer.testing import bench_gpu_time_with_cupti as bench_fn
-
-        bench_fn = partial(bench_fn, use_cuda_graph=True, cold_l2_cache=False)
-    elif method == "cudagraph":
-        from flashinfer.testing import (
-            bench_gpu_time_with_cudagraph as bench_fn,
-        )
-
-        bench_fn = partial(bench_fn, cold_l2_cache=False)
-    else:
-        raise ValueError(f"Invalid method: {method}")
-
-    attn_fp8, attn_bf16, workspace = _setup_fp8_attention(num_heads, head_dim)
-
-    print(f"Timing method: {method}")
-    print(f"{'seq_len':>8} {'BF16 (us)':>12} {'FP8 (us)':>12} {'Speedup':>10}")
-    print("-" * 46)
-
-    for seq_len in seq_lens:
-        torch.manual_seed(42)
-
-        q = torch.randn(
-            seq_len,
-            num_heads,
-            head_dim,
-            device="cuda",
-            dtype=torch.bfloat16,
-        )
-        k = torch.randn_like(q)
-        v = torch.randn_like(q)
-
-        cu_fp8, max_s, seq_l = _build_meta(seq_len, num_heads, head_dim, fp8=True)
-        # we can reuse cu_fp8 for cu_bf16 since q, k, and v are contiguous
-        cu_bf16 = cu_fp8.clone()
-
-        def bf16_fn(q=q, k=k, v=v, cu=cu_bf16, ms=max_s, sl=seq_l):
-            attn_bf16._forward_flashinfer(q, k, v, cu, ms, sl)
-
-        def fp8_fn(q=q, k=k, v=v, cu=cu_fp8, ms=max_s, sl=seq_l):
-            attn_fp8._forward_flashinfer(q, k, v, cu, ms, sl)
-
-        # bench_fn returns List[float] of per-iteration times in ms
-        bf16_times = bench_fn(bf16_fn)
-        fp8_times = bench_fn(fp8_fn)
-
-        bf16_us = np.median(bf16_times) * 1e3  # ms -> us
-        fp8_us = np.median(fp8_times) * 1e3
-        speedup = bf16_us / fp8_us if fp8_us > 0 else float("inf")
-
-        print(f"{seq_len:>8} {bf16_us:>12.1f} {fp8_us:>12.1f} {speedup:>9.2f}x")
-
-
-def _make_trace_handler(output_dir: str, worker_name: str, label: str):
-    """Create a trace handler that saves to TensorBoard and prints summary."""
-
-    def handler(prof):
-        torch.profiler.tensorboard_trace_handler(output_dir, worker_name)(prof)
-        print(f"\n{'=' * 80}")
-        print(label)
-        print(f"{'=' * 80}")
-        print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=20))
-
-    return handler
-
-
-def run_profile(
-    seq_len: int,
-    num_heads: int,
-    head_dim: int,
-    warmup: int,
-    output_dir: str,
-):
-    """Profile FP8 vs BF16 attention with PyTorch profiler."""
-    attn_fp8, attn_bf16, workspace = _setup_fp8_attention(num_heads, head_dim)
-
-    torch.manual_seed(42)
-    q = torch.randn(
-        seq_len,
-        num_heads,
-        head_dim,
-        device="cuda",
-        dtype=torch.bfloat16,
-    )
-    k = torch.randn_like(q)
-    v = torch.randn_like(q)
-
-    cu_fp8, max_s, seq_l = _build_meta(seq_len, num_heads, head_dim, fp8=True)
-    # we can reuse cu_fp8 for cu_bf16 since q, k, and v are contiguous
-    cu_bf16 = cu_fp8.clone()
-
-    sched = torch.profiler.schedule(wait=0, warmup=warmup, active=1)
-
-    # Profile BF16 (warmup handled by profiler schedule)
-    with profile(
-        activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
-        schedule=sched,
-        on_trace_ready=_make_trace_handler(
-            output_dir,
-            f"bf16_h{head_dim}_s{seq_len}",
-            f"BF16 Attention (seq_len={seq_len}, heads={num_heads}, "
-            f"head_dim={head_dim})",
-        ),
-    ) as prof_bf16:
-        for _ in range(warmup + 1):
-            with record_function("bf16_attention"):
-                attn_bf16._forward_flashinfer(
-                    q.clone(), k.clone(), v.clone(), cu_bf16, max_s, seq_l
-                )
-                torch.accelerator.synchronize()
-            prof_bf16.step()
-
-    # Profile FP8 (warmup handled by profiler schedule)
-    with profile(
-        activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
-        schedule=sched,
-        on_trace_ready=_make_trace_handler(
-            output_dir,
-            f"fp8_h{head_dim}_s{seq_len}",
-            f"FP8 Attention (seq_len={seq_len}, heads={num_heads}, "
-            f"head_dim={head_dim})",
-        ),
-    ) as prof_fp8:
-        for _ in range(warmup + 1):
-            with record_function("fp8_attention"):
-                attn_fp8._forward_flashinfer(
-                    q.clone(), k.clone(), v.clone(), cu_fp8, max_s, seq_l
-                )
-                torch.accelerator.synchronize()
-            prof_fp8.step()
-
-    print(f"\nTensorBoard traces saved to: {output_dir}")
-    print(f"View with: tensorboard --logdir={output_dir}")
-
-
-if __name__ == "__main__":
-    parser = FlexibleArgumentParser(description="Benchmark FP8 vs BF16 ViT attention.")
-    parser.add_argument(
-        "--seq-lens",
-        type=int,
-        nargs="+",
-        default=DEFAULT_SEQ_LENS,
-        help="Sequence lengths to benchmark",
-    )
-    parser.add_argument(
-        "--num-heads",
-        type=int,
-        default=NUM_HEADS,
-    )
-    parser.add_argument(
-        "--head-dim",
-        type=int,
-        default=HEAD_DIM,
-    )
-    parser.add_argument(
-        "--method",
-        choices=["cupti", "cudagraph"],
-        default="cudagraph",
-        help="GPU timing method: cupti (CUPTI kernel timing) or "
-        "cudagraph (CUDA graph capture/replay). Default: cudagraph",
-    )
-    parser.add_argument(
-        "--warmup",
-        type=int,
-        default=10,
-        help="Warmup iterations (profile mode only)",
-    )
-    parser.add_argument(
-        "--profile",
-        action="store_true",
-        help="Run PyTorch profiler instead of benchmark",
-    )
-    parser.add_argument(
-        "--profile-seq-len",
-        type=int,
-        default=8192,
-        help="Sequence length for profiling (default: 8192)",
-    )
-    parser.add_argument(
-        "--profile-output-dir",
-        type=str,
-        default="./profile_traces",
-        help="Output directory for TensorBoard traces (default: ./profile_traces)",
-    )
-    args = parser.parse_args()
-
-    if args.profile:
-        run_profile(
-            args.profile_seq_len,
-            args.num_heads,
-            args.head_dim,
-            args.warmup,
-            args.profile_output_dir,
-        )
-    else:
-        run_benchmark(
-            args.seq_lens,
-            args.num_heads,
-            args.head_dim,
-            args.method,
-        )

cmake/external_projects/deepgemm.cmake

@@ -20,7 +20,7 @@ else()
   FetchContent_Declare(
     deepgemm
     GIT_REPOSITORY https://github.com/deepseek-ai/DeepGEMM.git
-    GIT_TAG 891d57b4db1071624b5c8fa0d1e51cb317fa709f
+    GIT_TAG 477618cd51baffca09c4b0b87e97c03fe827ef03
     GIT_SUBMODULES "third-party/cutlass" "third-party/fmt"
     GIT_PROGRESS TRUE
     CONFIGURE_COMMAND ""
@@ -120,11 +120,6 @@ if(DEEPGEMM_ARCHS)
     COMPONENT _deep_gemm_C
     FILES_MATCHING PATTERN "*.py")
 
-  install(DIRECTORY "${deepgemm_SOURCE_DIR}/deep_gemm/mega/"
-    DESTINATION vllm/third_party/deep_gemm/mega
-    COMPONENT _deep_gemm_C
-    FILES_MATCHING PATTERN "*.py")
-
   # Generate envs.py (normally generated by DeepGEMM's setup.py build step)
   file(WRITE "${CMAKE_CURRENT_BINARY_DIR}/deep_gemm_envs.py"
     "# Pre-installed environment variables\npersistent_envs = dict()\n")

cmake/external_projects/flashmla.cmake

@@ -19,7 +19,7 @@ else()
   FetchContent_Declare(
         flashmla
         GIT_REPOSITORY https://github.com/vllm-project/FlashMLA
-        GIT_TAG a6ec2ba7bd0a7dff98b3f4d3e6b52b159c48d78b
+        GIT_TAG 692917b1cda61b93ac9ee2d846ec54e75afe87b1
         GIT_PROGRESS TRUE
         CONFIGURE_COMMAND ""
         BUILD_COMMAND ""

csrc/activation_kernels.cu

@@ -11,74 +11,29 @@
 namespace vllm {
 
 template <typename scalar_t, scalar_t (*ACT_FN)(const scalar_t&),
-          bool act_first, bool HAS_CLAMP>
+          bool act_first>
 __device__ __forceinline__ scalar_t compute(const scalar_t& x,
-                                            const scalar_t& y,
-                                            const float limit) {
-  if constexpr (act_first) {
-    scalar_t gate = x;
-    scalar_t up = y;
-    if constexpr (HAS_CLAMP) {
-      gate = (scalar_t)fminf((float)gate, limit);
-      up = (scalar_t)fmaxf(fminf((float)up, limit), -limit);
-    }
-    return ACT_FN(gate) * up;
-  } else {
-    scalar_t gate = x;
-    scalar_t up = y;
-    if constexpr (HAS_CLAMP) {
-      gate = (scalar_t)fmaxf(fminf((float)gate, limit), -limit);
-      up = (scalar_t)fminf((float)up, limit);
-    }
-    return gate * ACT_FN(up);
-  }
+                                            const scalar_t& y) {
+  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, bool HAS_CLAMP>
+          bool act_first>
 __device__ __forceinline__ packed_t packed_compute(const packed_t& x,
-                                                   const packed_t& y,
-                                                   const float limit) {
-  if constexpr (act_first) {
-    packed_t gate = x;
-    packed_t up = y;
-    if constexpr (HAS_CLAMP) {
-      float2 g = cast_to_float2(gate);
-      float2 u = cast_to_float2(up);
-      g.x = fminf(g.x, limit);
-      g.y = fminf(g.y, limit);
-      u.x = fmaxf(fminf(u.x, limit), -limit);
-      u.y = fmaxf(fminf(u.y, limit), -limit);
-      gate = cast_to_packed<packed_t>(g);
-      up = cast_to_packed<packed_t>(u);
-    }
-    return packed_mul(PACKED_ACT_FN(gate), up);
-  } else {
-    packed_t gate = x;
-    packed_t up = y;
-    if constexpr (HAS_CLAMP) {
-      float2 g = cast_to_float2(gate);
-      float2 u = cast_to_float2(up);
-      g.x = fmaxf(fminf(g.x, limit), -limit);
-      g.y = fmaxf(fminf(g.y, limit), -limit);
-      u.x = fminf(u.x, limit);
-      u.y = fminf(u.y, limit);
-      gate = cast_to_packed<packed_t>(g);
-      up = cast_to_packed<packed_t>(u);
-    }
-    return packed_mul(gate, PACKED_ACT_FN(up));
-  }
+                                                   const packed_t& y) {
+  return act_first ? packed_mul(PACKED_ACT_FN(x), y)
+                   : packed_mul(x, PACKED_ACT_FN(y));
 }
 
 // Activation and gating kernel template.
 template <typename scalar_t, typename packed_t,
           scalar_t (*ACT_FN)(const scalar_t&),
           packed_t (*PACKED_ACT_FN)(const packed_t&), bool act_first,
-          bool use_vec, bool HAS_CLAMP, bool use_256b = false>
+          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, const float limit) {
+    const int d) {
   const scalar_t* x_ptr = input + blockIdx.x * 2 * d;
   const scalar_t* y_ptr = x_ptr + d;
   scalar_t* out_ptr = out + blockIdx.x * d;
@@ -103,9 +58,8 @@ __global__ void act_and_mul_kernel(
       }
 #pragma unroll
       for (int j = 0; j < pvec_t::NUM_ELTS; j++) {
-        x.elts[j] =
-            packed_compute<packed_t, PACKED_ACT_FN, act_first, HAS_CLAMP>(
-                x.elts[j], y.elts[j], limit);
+        x.elts[j] = packed_compute<packed_t, PACKED_ACT_FN, act_first>(
+            x.elts[j], y.elts[j]);
       }
       if constexpr (use_256b) {
         st256(x, &out_vec[i]);
@@ -118,8 +72,7 @@ __global__ void act_and_mul_kernel(
     for (int64_t idx = threadIdx.x; idx < d; idx += blockDim.x) {
       const scalar_t x = VLLM_LDG(&x_ptr[idx]);
       const scalar_t y = VLLM_LDG(&y_ptr[idx]);
-      out_ptr[idx] =
-          compute<scalar_t, ACT_FN, act_first, HAS_CLAMP>(x, y, limit);
+      out_ptr[idx] = compute<scalar_t, ACT_FN, act_first>(x, y);
     }
   }
 }
@@ -198,11 +151,8 @@ packed_gelu_tanh_kernel(const packed_t& val) {
 
 // Launch activation and gating kernel.
 // Use ACT_FIRST (bool) indicating whether to apply the activation function
-// first. HAS_CLAMP (bool) enables pre-activation clamping: gate input is
-// clamped (max only) and up input is clamped (both sides) before the
-// activation function is applied.
-#define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL, PACKED_KERNEL, ACT_FIRST,        \
-                                      HAS_CLAMP, LIMIT)                        \
+// first.
+#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);                         \
@@ -227,8 +177,8 @@ packed_gelu_tanh_kernel(const packed_t& val) {
             scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type,      \
             KERNEL<scalar_t>,                                                  \
             PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>, \
-            ACT_FIRST, true, HAS_CLAMP, true><<<grid, block, 0, stream>>>(     \
-            out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d, LIMIT);   \
+            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", [&] {          \
@@ -236,8 +186,8 @@ packed_gelu_tanh_kernel(const packed_t& val) {
             scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type,      \
             KERNEL<scalar_t>,                                                  \
             PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>, \
-            ACT_FIRST, true, HAS_CLAMP, false><<<grid, block, 0, stream>>>(    \
-            out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d, LIMIT);   \
+            ACT_FIRST, true, false><<<grid, block, 0, stream>>>(               \
+            out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d);          \
       });                                                                      \
     }                                                                          \
   } else {                                                                     \
@@ -247,8 +197,8 @@ packed_gelu_tanh_kernel(const packed_t& val) {
           scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type,        \
           KERNEL<scalar_t>,                                                    \
           PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>,   \
-          ACT_FIRST, false, HAS_CLAMP><<<grid, block, 0, stream>>>(            \
-          out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d, LIMIT);     \
+          ACT_FIRST, false><<<grid, block, 0, stream>>>(                       \
+          out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d);            \
     });                                                                        \
   }
 
@@ -256,14 +206,7 @@ void silu_and_mul(torch::Tensor& out,    // [..., d]
                   torch::Tensor& input)  // [..., 2 * d]
 {
   LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel,
-                                true, false, 0.0f);
-}
-
-void silu_and_mul_clamp(torch::Tensor& out,    // [..., d]
-                        torch::Tensor& input,  // [..., 2 * d]
-                        double limit) {
-  LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel,
-                                true, true, (float)limit);
+                                true);
 }
 
 void mul_and_silu(torch::Tensor& out,    // [..., d]
@@ -272,21 +215,21 @@ void mul_and_silu(torch::Tensor& out,    // [..., d]
   // The difference between mul_and_silu and silu_and_mul is that mul_and_silu
   // applies the silu to the latter half of the input.
   LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel,
-                                false, false, 0.0f);
+                                false);
 }
 
 void gelu_and_mul(torch::Tensor& out,    // [..., d]
                   torch::Tensor& input)  // [..., 2 * d]
 {
   LAUNCH_ACTIVATION_GATE_KERNEL(vllm::gelu_kernel, vllm::packed_gelu_kernel,
-                                true, false, 0.0f);
+                                true);
 }
 
 void gelu_tanh_and_mul(torch::Tensor& out,    // [..., d]
                        torch::Tensor& input)  // [..., 2 * d]
 {
-  LAUNCH_ACTIVATION_GATE_KERNEL(
-      vllm::gelu_tanh_kernel, vllm::packed_gelu_tanh_kernel, true, false, 0.0f);
+  LAUNCH_ACTIVATION_GATE_KERNEL(vllm::gelu_tanh_kernel,
+                                vllm::packed_gelu_tanh_kernel, true);
 }
 
 namespace vllm {

csrc/cpu/pos_encoding.cpp

@@ -178,12 +178,7 @@ void rotary_embedding_gptj_impl(
 
 void rotary_embedding(torch::Tensor& positions, torch::Tensor& query,
                       std::optional<torch::Tensor> key, int64_t head_size,
-                      torch::Tensor& cos_sin_cache, bool is_neox,
-                      int64_t rope_dim_offset, bool inverse) {
-  TORCH_CHECK(rope_dim_offset == 0,
-              "rope_dim_offset != 0 is not supported on CPU");
-  TORCH_CHECK(!inverse, "inverse rotary embedding is not supported on CPU");
-
+                      torch::Tensor& cos_sin_cache, bool is_neox) {
   int num_tokens = positions.numel();
   int rot_dim = cos_sin_cache.size(1);
   int num_heads = query.size(-1) / head_size;

csrc/cpu/torch_bindings.cpp

@@ -263,8 +263,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
   ops.def(
       "rotary_embedding(Tensor positions, Tensor! query,"
       "                 Tensor!? key, int head_size,"
-      "                 Tensor cos_sin_cache, bool is_neox, int "
-      "rope_dim_offset=0, bool inverse=False) -> ()");
+      "                 Tensor cos_sin_cache, bool is_neox) -> ()");
   ops.impl("rotary_embedding", torch::kCPU, &rotary_embedding);
 
   // Quantization

csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu

@@ -1,477 +0,0 @@
-/*
- * SPDX-License-Identifier: Apache-2.0
- * SPDX-FileCopyrightText: Copyright contributors to the vLLM project
- *
- * Horizontally-fused DeepseekV4-MLA kernel:
- *   - Q side:  per-head RMSNorm (no weight) + GPT-J RoPE on last ROPE_DIM
- *   - KV side: GPT-J RoPE on last ROPE_DIM + UE8M0 FP8 quant on NoPE + paged
- *              cache insert
- *
- * Structured after `applyMLARopeAndAssignQKVKernelGeneration` in
- * TensorRT-LLM's mlaKernels.cu: one kernel, one grid, with head-slot
- * dispatch choosing Q vs KV work per warp.  The per-warp RMSNorm/RoPE
- * skeleton is adapted from vllm-deepseek_v4's existing
- * `fusedQKNormRopeKernel` (csrc/fused_qknorm_rope_kernel.cu).
- *
- * Assumptions (hard-coded for DeepseekV4 attention):
- *   HEAD_DIM  = 512
- *   ROPE_DIM  = 64   (RoPE applied to dims [NOPE_DIM, HEAD_DIM))
- *   NOPE_DIM  = 448
- *   QUANT_BLOCK = 64 (UE8M0 FP8 quant block)
- *   FP8_MAX   = 448.0f
- *   is_neox=false (GPT-J interleaved pairs)
- *   cos_sin_cache layout [max_pos, rope_dim] = cos || sin (cos first, sin
- *     second along last dim; each half is rope_dim/2 = 32 values)
- *
- * Cache layout per paged-cache block (block_size tokens):
- *   [0,            bs*576):          token data, 448 fp8 + 128 bf16 each
- *   [bs*576,       bs*576 + bs*8):   UE8M0 scales, 7 real + 1 pad per token
- */
-
-#include <cmath>
-#include <cuda_fp8.h>
-#include <cuda_runtime.h>
-#include <type_traits>
-
-#include <ATen/cuda/CUDAContext.h>
-#include <c10/cuda/CUDAGuard.h>
-#include <torch/cuda.h>
-
-#include "cuda_compat.h"
-#include "dispatch_utils.h"
-#include "type_convert.cuh"
-
-#ifndef FINAL_MASK
-  #define FINAL_MASK 0xffffffffu
-#endif
-
-namespace vllm {
-namespace deepseek_v4_fused_ops {
-
-namespace {
-inline int getSMVersion() {
-  auto* props = at::cuda::getCurrentDeviceProperties();
-  return props->major * 10 + props->minor;
-}
-}  // namespace
-
-// ────────────────────────────────────────────────────────────────────────────
-// Constants
-// ────────────────────────────────────────────────────────────────────────────
-constexpr int kHeadDim = 512;
-constexpr int kRopeDim = 64;
-constexpr int kNopeDim = kHeadDim - kRopeDim;  // 448
-constexpr int kQuantBlock = 64;
-constexpr int kNumQuantBlocks = kNopeDim / kQuantBlock;   // 7
-constexpr int kScaleBytesPerToken = kNumQuantBlocks + 1;  // 8 (7 real + 1 pad)
-constexpr int kTokenDataBytes = kNopeDim + kRopeDim * 2;  // 448 + 128 = 576
-constexpr float kFp8Max = 448.0f;
-
-// Per-warp layout:  32 lanes × 16 elems/lane = 512 elems = HEAD_DIM.
-constexpr int kNumLanes = 32;
-constexpr int kElemsPerLane = kHeadDim / kNumLanes;  // 16
-
-// ────────────────────────────────────────────────────────────────────────────
-// Small inline helpers
-// ────────────────────────────────────────────────────────────────────────────
-__device__ __forceinline__ float warp4MaxAbs(float val) {
-  // Reduce absolute max across 4 consecutive lanes (lane id & 3 group).
-  float peer = __shfl_xor_sync(FINAL_MASK, val, 1);
-  val = fmaxf(val, peer);
-  peer = __shfl_xor_sync(FINAL_MASK, val, 2);
-  val = fmaxf(val, peer);
-  return val;
-}
-
-template <typename T>
-__device__ __forceinline__ float warpSum(float val) {
-#pragma unroll
-  for (int mask = 16; mask > 0; mask >>= 1) {
-    val += __shfl_xor_sync(FINAL_MASK, val, mask, 32);
-  }
-  return val;
-}
-
-// ────────────────────────────────────────────────────────────────────────────
-// Kernel
-// ────────────────────────────────────────────────────────────────────────────
-//
-// Grid: 1D, gridDim.x = ceil(num_tokens_full * (num_heads_q + 1) /
-// warps_per_block) Block: blockDim.x = 256 threads (8 warps per block) Each
-// warp handles one (token, head_slot) pair. head_slot < num_heads_q          →
-// Q branch (RMSNorm + RoPE, in place) head_slot == num_heads_q         → KV
-// branch (RoPE + UE8M0 quant + insert)
-//
-// With DP padding, q/kv/position_ids can have more rows than slot_mapping.
-// The Q branch covers all `num_tokens_full` rows (downstream attention uses
-// them).  The KV branch only inserts the first `num_tokens_insert` tokens
-// (= slot_mapping length) into the paged cache.
-//
-template <typename scalar_t_in>
-__global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel(
-    scalar_t_in* __restrict__ q_inout,         // [N, H, 512] bf16, in place
-    scalar_t_in const* __restrict__ kv_in,     // [N, 512] bf16
-    uint8_t* __restrict__ k_cache,             // [num_blocks, block_stride]
-    int64_t const* __restrict__ slot_mapping,  // [num_tokens_insert] i64
-    int64_t const* __restrict__ position_ids,  // [N] i64
-    float const* __restrict__ cos_sin_cache,   // [max_pos, 64] fp32
-    float const eps,
-    int const num_tokens_full,    // = q.size(0) = kv.size(0)
-    int const num_tokens_insert,  // = slot_mapping.size(0), ≤ num_tokens_full
-    int const num_heads_q,        // H
-    int const cache_block_size,   // tokens per paged-cache block
-    int const kv_block_stride) {  // bytes per paged-cache block
-#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM)
-  // BF16 _typeConvert specialization is unavailable on pre-Ampere.  The
-  // DeepseekV4 kernel only runs with bf16 inputs in practice, so compile a
-  // no-op stub for sm_70/sm_75 to keep multi-arch builds happy.
-  if constexpr (std::is_same_v<scalar_t_in, c10::BFloat16>) {
-    return;
-  } else {
-#endif
-    using Converter = vllm::_typeConvert<scalar_t_in>;
-
-    int const warpsPerBlock = blockDim.x / 32;
-    int const warpId = threadIdx.x / 32;
-    int const laneId = threadIdx.x % 32;
-    int const globalWarpIdx = blockIdx.x * warpsPerBlock + warpId;
-
-    int const total_slots_per_token = num_heads_q + 1;
-    int const tokenIdx = globalWarpIdx / total_slots_per_token;
-    int const slotIdx = globalWarpIdx % total_slots_per_token;
-    if (tokenIdx >= num_tokens_full) return;
-
-    bool const isKV = (slotIdx == num_heads_q);
-    // KV branch: skip DP-padded tokens (no slot reserved for them).
-    if (isKV && tokenIdx >= num_tokens_insert) return;
-
-    // PDL: wait for predecessor kernel (upstream q/kv producer) to signal
-    // before touching any global memory.  No-op when PDL is not enabled on
-    // the launch.  The CUDA runtime wrapper emits the griddepcontrol.wait
-    // PTX with the required memory clobber internally.
-#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
-    cudaGridDependencySynchronize();
-#endif
-
-    // Dim range this lane owns within the 512-wide head.
-    int const dim_base = laneId * kElemsPerLane;  // in [0, 512) step 16
-
-    // ── Load 16 bf16 → 16 fp32 registers (one 16-byte + one 16-byte LDG) ────
-    float elements[kElemsPerLane];
-    float sumOfSquares = 0.0f;
-
-    scalar_t_in const* src_ptr;
-    if (isKV) {
-      src_ptr = kv_in + static_cast<int64_t>(tokenIdx) * kHeadDim + dim_base;
-    } else {
-      int64_t const q_row_offset =
-          (static_cast<int64_t>(tokenIdx) * num_heads_q + slotIdx) * kHeadDim +
-          dim_base;
-      src_ptr = q_inout + q_row_offset;
-    }
-
-    // Two 16-byte loads per thread (8 bf16 each).  Use uint4 as the vector
-    // type and bitcast to scalar_t_in packed pairs for conversion.
-    uint4 v0 = *reinterpret_cast<uint4 const*>(src_ptr);
-    uint4 v1 = *reinterpret_cast<uint4 const*>(src_ptr + 8);
-
-    {
-      typename Converter::packed_hip_type const* p0 =
-          reinterpret_cast<typename Converter::packed_hip_type const*>(&v0);
-      typename Converter::packed_hip_type const* p1 =
-          reinterpret_cast<typename Converter::packed_hip_type const*>(&v1);
-// Each packed_hip_type holds 2 bf16 → 4 packed = 8 elems per uint4.
-#pragma unroll
-      for (int i = 0; i < 4; i++) {
-        float2 f2 = Converter::convert(p0[i]);
-        elements[2 * i] = f2.x;
-        elements[2 * i + 1] = f2.y;
-      }
-#pragma unroll
-      for (int i = 0; i < 4; i++) {
-        float2 f2 = Converter::convert(p1[i]);
-        elements[8 + 2 * i] = f2.x;
-        elements[8 + 2 * i + 1] = f2.y;
-      }
-    }
-
-    // ── Q branch: RMSNorm with no weight (has_weight=False) ─────────────────
-    // Variance + rsqrt + multiply all in fp32, no intermediate bf16 round.
-    // The downstream bf16 round only happens at the final store.
-    if (!isKV) {
-#pragma unroll
-      for (int i = 0; i < kElemsPerLane; i++) {
-        sumOfSquares += elements[i] * elements[i];
-      }
-      sumOfSquares = warpSum<float>(sumOfSquares);
-      float const rms_rcp =
-          rsqrtf(sumOfSquares / static_cast<float>(kHeadDim) + eps);
-#pragma unroll
-      for (int i = 0; i < kElemsPerLane; i++) {
-        elements[i] = elements[i] * rms_rcp;
-      }
-    }
-
-    // ── GPT-J RoPE on dims [NOPE_DIM, HEAD_DIM) ─────────────────────────────
-    // All math in fp32.  cos_sin_cache is loaded as fp32 (its native storage).
-    bool const is_rope_lane = dim_base >= kNopeDim;
-    if (is_rope_lane) {
-      int64_t const pos = position_ids[tokenIdx];
-      constexpr int kHalfRope = kRopeDim / 2;  // 32
-      float const* cos_ptr = cos_sin_cache + pos * kRopeDim;
-      float const* sin_ptr = cos_ptr + kHalfRope;
-
-      int const rope_local_base = dim_base - kNopeDim;  // in [0, 64) step 16
-#pragma unroll
-      for (int p = 0; p < kElemsPerLane / 2; p++) {
-        int const pair_dim = rope_local_base + 2 * p;
-        int const half_idx = pair_dim / 2;
-        float const cos_v = VLLM_LDG(cos_ptr + half_idx);
-        float const sin_v = VLLM_LDG(sin_ptr + half_idx);
-        float const x_even = elements[2 * p];
-        float const x_odd = elements[2 * p + 1];
-        elements[2 * p] = x_even * cos_v - x_odd * sin_v;
-        elements[2 * p + 1] = x_even * sin_v + x_odd * cos_v;
-      }
-    }
-
-    // ═══════════════════════════════════════════════════════════════════════
-    // Q branch: cast to bf16 and store back in place.
-    // ═══════════════════════════════════════════════════════════════════════
-    if (!isKV) {
-      uint4 out0, out1;
-      typename Converter::packed_hip_type* po0 =
-          reinterpret_cast<typename Converter::packed_hip_type*>(&out0);
-      typename Converter::packed_hip_type* po1 =
-          reinterpret_cast<typename Converter::packed_hip_type*>(&out1);
-#pragma unroll
-      for (int i = 0; i < 4; i++) {
-        po0[i] = Converter::convert(
-            make_float2(elements[2 * i], elements[2 * i + 1]));
-      }
-#pragma unroll
-      for (int i = 0; i < 4; i++) {
-        po1[i] = Converter::convert(
-            make_float2(elements[8 + 2 * i], elements[8 + 2 * i + 1]));
-      }
-      scalar_t_in* dst =
-          q_inout +
-          (static_cast<int64_t>(tokenIdx) * num_heads_q + slotIdx) * kHeadDim +
-          dim_base;
-      *reinterpret_cast<uint4*>(dst) = out0;
-      *reinterpret_cast<uint4*>(dst + 8) = out1;
-#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
-      cudaTriggerProgrammaticLaunchCompletion();
-#endif
-      return;
-    }
-
-    // ═══════════════════════════════════════════════════════════════════════
-    // KV branch.
-    // ═══════════════════════════════════════════════════════════════════════
-    int64_t const slot_id = slot_mapping[tokenIdx];
-    if (slot_id < 0) {
-#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
-      cudaTriggerProgrammaticLaunchCompletion();
-#endif
-      return;
-    }
-
-    int64_t const block_idx = slot_id / cache_block_size;
-    int64_t const pos_in_block = slot_id % cache_block_size;
-    uint8_t* block_base =
-        k_cache + block_idx * static_cast<int64_t>(kv_block_stride);
-    uint8_t* token_fp8_ptr = block_base + pos_in_block * kTokenDataBytes;
-    uint8_t* token_bf16_ptr = token_fp8_ptr + kNopeDim;
-    uint8_t* token_scale_ptr =
-        block_base + static_cast<int64_t>(cache_block_size) * kTokenDataBytes +
-        pos_in_block * kScaleBytesPerToken;
-
-    // Round K to bf16 first, matching the unfused reference path where K is
-    // materialized as bf16 before K quantization.  absmax, clamp, and FP8
-    // quant below all run on these bf16-rounded values.
-#pragma unroll
-    for (int i = 0; i < kElemsPerLane; i++) {
-      elements[i] = Converter::convert(Converter::convert(elements[i]));
-    }
-
-    // Per-quant-block absmax must be computed by ALL 32 lanes (warp-collective
-    // shuffle requires full participation).  RoPE lanes contribute garbage,
-    // but their values are gated out below via `!is_rope_lane`.
-    float local_absmax = 0.0f;
-#pragma unroll
-    for (int i = 0; i < kElemsPerLane; i++) {
-      local_absmax = fmaxf(local_absmax, fabsf(elements[i]));
-    }
-    float const absmax = fmaxf(warp4MaxAbs(local_absmax), 1e-4f);
-    float const exponent = ceilf(log2f(absmax / kFp8Max));
-    float const inv_scale = exp2f(-exponent);
-
-    if (!is_rope_lane) {
-      // ── NoPE lane: UE8M0 FP8 quant ───────────────────────────────────────
-      uint8_t out_bytes[kElemsPerLane];
-#pragma unroll
-      for (int i = 0; i < kElemsPerLane; i++) {
-        float scaled = elements[i] * inv_scale;
-        scaled = fminf(fmaxf(scaled, -kFp8Max), kFp8Max);
-        __nv_fp8_storage_t s =
-            __nv_cvt_float_to_fp8(scaled, __NV_SATFINITE, __NV_E4M3);
-        out_bytes[i] = static_cast<uint8_t>(s);
-      }
-      // One 16-byte STG per lane.
-      *reinterpret_cast<uint4*>(token_fp8_ptr + dim_base) =
-          *reinterpret_cast<uint4 const*>(out_bytes);
-
-      // Lane (4k) of each 4-lane group writes the scale byte for block k<7.
-      if ((laneId & 3) == 0) {
-        int const q_block_idx = laneId >> 2;  // 0..6 for NoPE lanes
-        float encoded = fmaxf(fminf(exponent + 127.0f, 255.0f), 0.0f);
-        token_scale_ptr[q_block_idx] = static_cast<uint8_t>(encoded);
-      }
-      // Lane 0 also writes the padding byte at index 7.
-      if (laneId == 0) {
-        token_scale_ptr[kNumQuantBlocks] = 0;  // pad
-      }
-    } else {
-      // ── RoPE lane: cast back to bf16 and store to cache bf16 tail ────────
-      uint4 out0, out1;
-      typename Converter::packed_hip_type* po0 =
-          reinterpret_cast<typename Converter::packed_hip_type*>(&out0);
-      typename Converter::packed_hip_type* po1 =
-          reinterpret_cast<typename Converter::packed_hip_type*>(&out1);
-#pragma unroll
-      for (int i = 0; i < 4; i++) {
-        po0[i] = Converter::convert(
-            make_float2(elements[2 * i], elements[2 * i + 1]));
-      }
-#pragma unroll
-      for (int i = 0; i < 4; i++) {
-        po1[i] = Converter::convert(
-            make_float2(elements[8 + 2 * i], elements[8 + 2 * i + 1]));
-      }
-      int const rope_local_base = dim_base - kNopeDim;  // in [0, 64)
-      scalar_t_in* bf16_dst =
-          reinterpret_cast<scalar_t_in*>(token_bf16_ptr) + rope_local_base;
-      *reinterpret_cast<uint4*>(bf16_dst) = out0;
-      *reinterpret_cast<uint4*>(bf16_dst + 8) = out1;
-    }
-#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
-    cudaTriggerProgrammaticLaunchCompletion();
-#endif
-#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM)
-  }
-#endif
-}
-
-// ────────────────────────────────────────────────────────────────────────────
-// Launch wrapper
-// ────────────────────────────────────────────────────────────────────────────
-template <typename scalar_t_in>
-void launchFusedDeepseekV4QNormRopeKVRopeQuantInsert(
-    scalar_t_in* q_inout, scalar_t_in const* kv_in, uint8_t* k_cache,
-    int64_t const* slot_mapping, int64_t const* position_ids,
-    float const* cos_sin_cache, float const eps, int const num_tokens_full,
-    int const num_tokens_insert, int const num_heads_q,
-    int const cache_block_size, int const kv_block_stride,
-    cudaStream_t stream) {
-  constexpr int kBlockSize = 256;
-  constexpr int kWarpsPerBlock = kBlockSize / 32;
-  int64_t const total_warps =
-      static_cast<int64_t>(num_tokens_full) * (num_heads_q + 1);
-  int const grid =
-      static_cast<int>((total_warps + kWarpsPerBlock - 1) / kWarpsPerBlock);
-
-  // PDL: enable programmatic stream serialization whenever the hardware
-  // supports it (SM90+).  On pre-Hopper GPUs the attribute is unavailable,
-  // so leave numAttrs = 0 and launch as a regular kernel.
-  static int const sm_version = getSMVersion();
-  // Host-side guard: the device kernel body is compiled as a no-op for
-  // bf16 on pre-Ampere (sm_70/sm_75) because _typeConvert<BFloat16> is
-  // unavailable there.  Refuse the launch loudly instead of silently
-  // skipping the work.
-  TORCH_CHECK(
-      sm_version >= 80,
-      "fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert requires sm_80+ "
-      "(Ampere or newer); got sm_",
-      sm_version);
-  cudaLaunchConfig_t config;
-  config.gridDim = dim3(grid);
-  config.blockDim = dim3(kBlockSize);
-  config.dynamicSmemBytes = 0;
-  config.stream = stream;
-  cudaLaunchAttribute attrs[1];
-  attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
-  attrs[0].val.programmaticStreamSerializationAllowed = 1;
-  config.attrs = attrs;
-  config.numAttrs = (sm_version >= 90) ? 1 : 0;
-
-  cudaLaunchKernelEx(
-      &config, fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel<scalar_t_in>,
-      q_inout, kv_in, k_cache, slot_mapping, position_ids, cos_sin_cache, eps,
-      num_tokens_full, num_tokens_insert, num_heads_q, cache_block_size,
-      kv_block_stride);
-}
-
-}  // namespace deepseek_v4_fused_ops
-}  // namespace vllm
-
-// ────────────────────────────────────────────────────────────────────────────
-// Torch op wrapper
-// ────────────────────────────────────────────────────────────────────────────
-void fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
-    torch::Tensor& q,                    // [N, H, 512] bf16, in place
-    torch::Tensor const& kv,             // [N, 512] bf16 (read-only)
-    torch::Tensor& k_cache,              // [num_blocks, block_bytes] uint8
-    torch::Tensor const& slot_mapping,   // [N] int64
-    torch::Tensor const& position_ids,   // [N] int64
-    torch::Tensor const& cos_sin_cache,  // [max_pos, rope_dim] bf16
-    double eps, int64_t cache_block_size) {
-  TORCH_CHECK(q.is_cuda() && q.is_contiguous(), "q must be contiguous CUDA");
-  TORCH_CHECK(kv.is_cuda() && kv.is_contiguous(), "kv must be contiguous CUDA");
-  TORCH_CHECK(k_cache.is_cuda(), "k_cache must be CUDA");
-  TORCH_CHECK(slot_mapping.is_cuda() && slot_mapping.dtype() == torch::kInt64,
-              "slot_mapping must be int64 CUDA");
-  TORCH_CHECK(position_ids.is_cuda() && position_ids.dtype() == torch::kInt64,
-              "position_ids must be int64 CUDA");
-  TORCH_CHECK(cos_sin_cache.is_cuda(), "cos_sin_cache must be CUDA");
-  TORCH_CHECK(q.dim() == 3 && q.size(2) == 512, "q shape [N, H, 512]");
-  TORCH_CHECK(kv.dim() == 2 && kv.size(1) == 512, "kv shape [N, 512]");
-  TORCH_CHECK(q.dtype() == kv.dtype(), "q and kv dtype must match");
-  TORCH_CHECK(k_cache.dtype() == torch::kUInt8, "k_cache must be uint8");
-  TORCH_CHECK(cos_sin_cache.dim() == 2 && cos_sin_cache.size(1) == 64,
-              "cos_sin_cache shape [max_pos, 64]");
-  TORCH_CHECK(cos_sin_cache.dtype() == torch::kFloat32,
-              "cos_sin_cache must be float32");
-
-  // With DP padding, slot_mapping can be shorter than q/kv/positions.
-  // Q-norm+RoPE runs on all q.size(0) rows (downstream attention uses them);
-  // KV quant+insert runs only on the first slot_mapping.size(0) rows.
-  int const num_tokens_full = static_cast<int>(q.size(0));
-  int const num_tokens_insert = static_cast<int>(slot_mapping.size(0));
-  TORCH_CHECK(static_cast<int>(kv.size(0)) == num_tokens_full &&
-                  static_cast<int>(position_ids.size(0)) == num_tokens_full,
-              "q/kv/position_ids row counts must match");
-  TORCH_CHECK(num_tokens_insert <= num_tokens_full,
-              "slot_mapping must not exceed q row count");
-  int const num_heads_q = static_cast<int>(q.size(1));
-  int const cache_block_size_i = static_cast<int>(cache_block_size);
-  int const kv_block_stride = static_cast<int>(k_cache.stride(0));
-
-  at::cuda::OptionalCUDAGuard device_guard(device_of(q));
-  auto stream = at::cuda::getCurrentCUDAStream();
-
-  VLLM_DISPATCH_HALF_TYPES(
-      q.scalar_type(), "fused_deepseek_v4_qnorm_rope_kv_insert", [&] {
-        using qkv_scalar_t = scalar_t;
-        vllm::deepseek_v4_fused_ops::
-            launchFusedDeepseekV4QNormRopeKVRopeQuantInsert<qkv_scalar_t>(
-                reinterpret_cast<qkv_scalar_t*>(q.data_ptr()),
-                reinterpret_cast<qkv_scalar_t const*>(kv.data_ptr()),
-                reinterpret_cast<uint8_t*>(k_cache.data_ptr()),
-                reinterpret_cast<int64_t const*>(slot_mapping.data_ptr()),
-                reinterpret_cast<int64_t const*>(position_ids.data_ptr()),
-                cos_sin_cache.data_ptr<float>(), static_cast<float>(eps),
-                num_tokens_full, num_tokens_insert, num_heads_q,
-                cache_block_size_i, kv_block_stride, stream);
-      });
-}

csrc/layernorm_kernels.cu

@@ -77,8 +77,7 @@ __global__ void rms_norm_kernel(
 #pragma unroll
     for (int j = 0; j < VEC_SIZE; j++) {
       float x = static_cast<float>(src1.val[j]);
-      float w = static_cast<float>(src2.val[j]);
-      dst.val[j] = static_cast<scalar_t>(x * s_variance * w);
+      dst.val[j] = ((scalar_t)(x * s_variance)) * src2.val[j];
     }
     v_out[i] = dst;
   }
@@ -135,17 +134,10 @@ fused_add_rms_norm_kernel(
   for (int idx = threadIdx.x; idx < vec_hidden_size; idx += blockDim.x) {
     int id = blockIdx.x * vec_hidden_size + idx;
     int64_t strided_id = blockIdx.x * vec_input_stride + idx;
-    _f16Vec<scalar_t, width> res = residual_v[id];
-    _f16Vec<scalar_t, width> w = weight_v[idx];
-    _f16Vec<scalar_t, width> out;
-    using Converter = _typeConvert<scalar_t>;
-#pragma unroll
-    for (int j = 0; j < width; ++j) {
-      float x = Converter::convert(res.data[j]);
-      float wf = Converter::convert(w.data[j]);
-      out.data[j] = Converter::convert(x * s_variance * wf);
-    }
-    input_v[strided_id] = out;
+    _f16Vec<scalar_t, width> temp = residual_v[id];
+    temp *= s_variance;
+    temp *= weight_v[idx];
+    input_v[strided_id] = temp;
   }
 }
 
@@ -182,8 +174,8 @@ fused_add_rms_norm_kernel(
 
   for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) {
     float x = (float)residual[blockIdx.x * hidden_size + idx];
-    float w = (float)weight[idx];
-    input[blockIdx.x * input_stride + idx] = (scalar_t)(x * s_variance * w);
+    input[blockIdx.x * input_stride + idx] =
+        ((scalar_t)(x * s_variance)) * weight[idx];
   }
 }
 

csrc/layernorm_quant_kernels.cu

@@ -65,16 +65,9 @@ __global__ void rms_norm_static_fp8_quant_kernel(
 #pragma unroll
     for (int j = 0; j < VEC_SIZE; j++) {
       float x = static_cast<float>(src1.val[j]);
-      float w = static_cast<float>(src2.val[j]);
-      // Round normalized result through scalar_t to match the precision of the
-      // unfused composite (rms_norm writes scalar_t, then
-      // static_scaled_fp8_quant re-loads it as float before FP8 conversion).
-      // Without this round, the fused path is strictly more accurate and
-      // disagrees with the composite at exact E4M3 quantization tie boundaries.
-      scalar_t out_norm = static_cast<scalar_t>(x * s_variance * w);
+      float const out_norm = ((scalar_t)(x * s_variance)) * src2.val[j];
       out[blockIdx.x * hidden_size + idx * VEC_SIZE + j] =
-          scaled_fp8_conversion<true, fp8_type>(static_cast<float>(out_norm),
-                                                scale_inv);
+          scaled_fp8_conversion<true, fp8_type>(out_norm, scale_inv);
     }
   }
 }
@@ -134,21 +127,13 @@ fused_add_rms_norm_static_fp8_quant_kernel(
 
   for (int idx = threadIdx.x; idx < vec_hidden_size; idx += blockDim.x) {
     int id = blockIdx.x * vec_hidden_size + idx;
-    _f16Vec<scalar_t, width> res = residual_v[id];
-    _f16Vec<scalar_t, width> w = weight_v[idx];
-    using Converter = _typeConvert<scalar_t>;
-    using HipT = typename Converter::hip_type;
+    _f16Vec<scalar_t, width> temp = residual_v[id];
+    temp *= s_variance;
+    temp *= weight_v[idx];
 #pragma unroll
     for (int i = 0; i < width; ++i) {
-      float x = Converter::convert(res.data[i]);
-      float wf = Converter::convert(w.data[i]);
-      // See note in rms_norm_static_fp8_quant_kernel: round through scalar_t
-      // to match the unfused composite path at FP8 boundaries. We use the
-      // backend's hip_type for the intermediate since c10::Half/BFloat16 has
-      // ambiguous conversions on CUDA and no implicit conversion on ROCm.
-      HipT out_norm_h = Converter::convert(x * s_variance * wf);
-      out[id * width + i] = scaled_fp8_conversion<true, fp8_type>(
-          Converter::convert(out_norm_h), scale_inv);
+      out[id * width + i] =
+          scaled_fp8_conversion<true, fp8_type>(float(temp.data[i]), scale_inv);
     }
   }
 }
@@ -191,12 +176,9 @@ fused_add_rms_norm_static_fp8_quant_kernel(
 
   for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) {
     float x = (float)residual[blockIdx.x * hidden_size + idx];
-    float w = (float)weight[idx];
-    // See note in rms_norm_static_fp8_quant_kernel: round through scalar_t
-    // to match the unfused composite path at FP8 boundaries.
-    scalar_t out_norm = static_cast<scalar_t>(x * s_variance * w);
-    out[blockIdx.x * hidden_size + idx] = scaled_fp8_conversion<true, fp8_type>(
-        static_cast<float>(out_norm), scale_inv);
+    float const out_norm = ((scalar_t)(x * s_variance)) * weight[idx];
+    out[blockIdx.x * hidden_size + idx] =
+        scaled_fp8_conversion<true, fp8_type>(out_norm, scale_inv);
   }
 }
 

csrc/moe/moe_ops.h

@@ -12,15 +12,6 @@ void topk_sigmoid(torch::Tensor& topk_weights, torch::Tensor& topk_indices,
                   torch::Tensor& gating_output, bool renormalize,
                   std::optional<torch::Tensor> bias);
 
-void topk_softplus_sqrt(torch::Tensor& topk_weights,
-                        torch::Tensor& topk_indices,
-                        torch::Tensor& token_expert_indices,
-                        torch::Tensor& gating_output, bool renormalize,
-                        double routed_scaling_factor,
-                        const c10::optional<torch::Tensor>& correction_bias,
-                        const c10::optional<torch::Tensor>& input_ids,
-                        const c10::optional<torch::Tensor>& tid2eid);
-
 void moe_sum(torch::Tensor& input, torch::Tensor& output);
 
 void moe_align_block_size(torch::Tensor topk_ids, int64_t num_experts,

csrc/moe/topk_softplus_sqrt_kernels.cu

@@ -1,715 +0,0 @@
-/*
- * Adapted from
- * https://github.com/NVIDIA/TensorRT-LLM/blob/v0.7.1/cpp/tensorrt_llm/kernels/mixtureOfExperts/moe_kernels.cu
- * Copyright (c) 2024, The vLLM team.
- * SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION &
- * AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-#include <type_traits>
-#include <torch/all.h>
-#include <ATen/cuda/CUDAContext.h>
-#include <c10/cuda/CUDAGuard.h>
-#include "../cuda_compat.h"
-#include "../cub_helpers.h"
-#ifndef USE_ROCM
-  #include <cuda_bf16.h>
-  #include <cuda_fp16.h>
-#else
-  #include <hip/hip_bf16.h>
-  #include <hip/hip_fp16.h>
-typedef __hip_bfloat16 __nv_bfloat16;
-typedef __hip_bfloat162 __nv_bfloat162;
-#endif
-
-#define MAX(a, b) ((a) > (b) ? (a) : (b))
-#define MIN(a, b) ((a) < (b) ? (a) : (b))
-
-namespace vllm {
-namespace moe {
-
-/// Aligned array type
-template <typename T,
-          /// Number of elements in the array
-          int N,
-          /// Alignment requirement in bytes
-          int Alignment = sizeof(T) * N>
-struct alignas(Alignment) AlignedArray {
-  T data[N];
-};
-
-template <typename T>
-__device__ __forceinline__ float toFloat(T value) {
-  if constexpr (std::is_same_v<T, float>) {
-    return value;
-  } else if constexpr (std::is_same_v<T, __nv_bfloat16>) {
-    return __bfloat162float(value);
-  } else if constexpr (std::is_same_v<T, __half>) {
-    return __half2float(value);
-  }
-}
-
-#define FINAL_MASK 0xffffffff
-template <typename T>
-__inline__ __device__ T warpReduceSum(T val) {
-#pragma unroll
-  for (int mask = 16; mask > 0; mask >>= 1)
-    val += __shfl_xor_sync(FINAL_MASK, val, mask, 32);
-  return val;
-}
-
-// ====================== TopK softplus_sqrt things
-// ===============================
-
-/*
-  A Top-K gating softplus_sqrt written to exploit when the number of experts in
-  the MoE layers are a small power of 2. This allows us to cleanly share the
-  rows among the threads in a single warp and eliminate communication between
-  warps (so no need to use shared mem).
-
-  It fuses the sigmoid, max and argmax into a single kernel.
-
-  Limitations:
-  1) This implementation is optimized for when the number of experts is a small
-  power of 2. Additionally it also supports when number of experts is multiple
-  of 64 which is still faster than the computing sigmoid and topK separately
-  (only tested on CUDA yet). 2) This implementation assumes k is small, but will
-  work for any k.
-*/
-
-template <int VPT, int NUM_EXPERTS, int WARPS_PER_CTA, int BYTES_PER_LDG,
-          int WARP_SIZE_PARAM, bool USE_HASH, typename IndType,
-          typename InputType = float>
-__launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
-    void topkGatingSoftplusSqrt(
-        const InputType* input, const bool* finished, float* output,
-        const int num_rows, IndType* indices, int* source_rows, const int k,
-        const int start_expert, const int end_expert, const bool renormalize,
-        double routed_scaling_factor, const float* correction_bias,
-        const IndType* input_ids, const IndType* tid2eid) {
-  static_assert(std::is_same_v<InputType, float> ||
-                    std::is_same_v<InputType, __nv_bfloat16> ||
-                    std::is_same_v<InputType, __half>,
-                "InputType must be float, __nv_bfloat16, or __half");
-
-  // We begin by enforcing compile time assertions and setting up compile time
-  // constants.
-  static_assert(BYTES_PER_LDG == (BYTES_PER_LDG & -BYTES_PER_LDG),
-                "BYTES_PER_LDG must be power of 2");
-  static_assert(BYTES_PER_LDG <= 16, "BYTES_PER_LDG must be leq 16");
-
-  // Number of bytes each thread pulls in per load
-  static constexpr int ELTS_PER_LDG = BYTES_PER_LDG / sizeof(InputType);
-  static constexpr int ELTS_PER_ROW = NUM_EXPERTS;
-  static constexpr int THREADS_PER_ROW = ELTS_PER_ROW / VPT;
-  static constexpr int LDG_PER_THREAD = VPT / ELTS_PER_LDG;
-
-  if constexpr (std::is_same_v<InputType, __nv_bfloat16> ||
-                std::is_same_v<InputType, __half>) {
-    static_assert(ELTS_PER_LDG == 1 || ELTS_PER_LDG % 2 == 0,
-                  "ELTS_PER_LDG must be 1 or even for 16-bit conversion");
-  }
-
-  // Restrictions based on previous section.
-  static_assert(
-      VPT % ELTS_PER_LDG == 0,
-      "The elements per thread must be a multiple of the elements per ldg");
-  static_assert(WARP_SIZE_PARAM % THREADS_PER_ROW == 0,
-                "The threads per row must cleanly divide the threads per warp");
-  static_assert(THREADS_PER_ROW == (THREADS_PER_ROW & -THREADS_PER_ROW),
-                "THREADS_PER_ROW must be power of 2");
-  static_assert(THREADS_PER_ROW <= WARP_SIZE_PARAM,
-                "THREADS_PER_ROW can be at most warp size");
-
-  // We have NUM_EXPERTS elements per row. We specialize for small #experts
-  static constexpr int ELTS_PER_WARP = WARP_SIZE_PARAM * VPT;
-  static constexpr int ROWS_PER_WARP = ELTS_PER_WARP / ELTS_PER_ROW;
-  static constexpr int ROWS_PER_CTA = WARPS_PER_CTA * ROWS_PER_WARP;
-
-  // Restrictions for previous section.
-  static_assert(ELTS_PER_WARP % ELTS_PER_ROW == 0,
-                "The elts per row must cleanly divide the total elt per warp");
-
-  // ===================== From this point, we finally start computing run-time
-  // variables. ========================
-
-  // Compute CTA and warp rows. We pack multiple rows into a single warp, and a
-  // block contains WARPS_PER_CTA warps. This, each block processes a chunk of
-  // rows. We start by computing the start row for each block.
-  const int cta_base_row = blockIdx.x * ROWS_PER_CTA;
-
-  // Now, using the base row per thread block, we compute the base row per warp.
-  const int warp_base_row = cta_base_row + threadIdx.y * ROWS_PER_WARP;
-
-  // The threads in a warp are split into sub-groups that will work on a row.
-  // We compute row offset for each thread sub-group
-  const int thread_row_in_warp = threadIdx.x / THREADS_PER_ROW;
-  const int thread_row = warp_base_row + thread_row_in_warp;
-
-  // Threads with indices out of bounds should early exit here.
-  if (thread_row >= num_rows) {
-    return;
-  }
-  const bool row_is_active = finished ? !finished[thread_row] : true;
-
-  // We finally start setting up the read pointers for each thread. First, each
-  // thread jumps to the start of the row it will read.
-  const InputType* thread_row_ptr = input + thread_row * ELTS_PER_ROW;
-
-  // Now, we compute the group each thread belong to in order to determine the
-  // first column to start loads.
-  const int thread_group_idx = threadIdx.x % THREADS_PER_ROW;
-  const int first_elt_read_by_thread = thread_group_idx * ELTS_PER_LDG;
-  const InputType* thread_read_ptr = thread_row_ptr + first_elt_read_by_thread;
-
-  // Finally, we pull in the data from global mem
-  float row_chunk[VPT];
-
-#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
-  asm volatile("griddepcontrol.wait;");
-#endif
-
-  // NOTE(zhuhaoran): dispatch different input types loading, BF16/FP16 convert
-  // to float
-  if constexpr (std::is_same_v<InputType, float>) {
-    using VecType = AlignedArray<float, ELTS_PER_LDG>;
-    VecType* row_chunk_vec_ptr = reinterpret_cast<VecType*>(&row_chunk);
-    const VecType* vec_thread_read_ptr =
-        reinterpret_cast<const VecType*>(thread_read_ptr);
-#pragma unroll
-    for (int ii = 0; ii < LDG_PER_THREAD; ++ii) {
-      row_chunk_vec_ptr[ii] = vec_thread_read_ptr[ii * THREADS_PER_ROW];
-    }
-  } else if constexpr (std::is_same_v<InputType, __nv_bfloat16>) {
-    if constexpr (ELTS_PER_LDG >= 2) {
-      using VecType = AlignedArray<__nv_bfloat16, ELTS_PER_LDG>;
-      float2* row_chunk_f2 = reinterpret_cast<float2*>(row_chunk);
-      const VecType* vec_thread_read_ptr =
-          reinterpret_cast<const VecType*>(thread_read_ptr);
-#pragma unroll
-      for (int ii = 0; ii < LDG_PER_THREAD; ++ii) {
-        VecType vec = vec_thread_read_ptr[ii * THREADS_PER_ROW];
-        int base_idx_f2 = ii * ELTS_PER_LDG / 2;
-#pragma unroll
-        for (int jj = 0; jj < ELTS_PER_LDG / 2; ++jj) {
-          row_chunk_f2[base_idx_f2 + jj] = __bfloat1622float2(
-              *reinterpret_cast<const __nv_bfloat162*>(vec.data + jj * 2));
-        }
-      }
-    } else {  // ELTS_PER_LDG == 1
-#pragma unroll
-      for (int ii = 0; ii < LDG_PER_THREAD; ++ii) {
-        const __nv_bfloat16* scalar_ptr =
-            thread_read_ptr + ii * THREADS_PER_ROW;
-        row_chunk[ii] = __bfloat162float(*scalar_ptr);
-      }
-    }
-  } else if constexpr (std::is_same_v<InputType, __half>) {
-    if constexpr (ELTS_PER_LDG >= 2) {
-      using VecType = AlignedArray<__half, ELTS_PER_LDG>;
-      float2* row_chunk_f2 = reinterpret_cast<float2*>(row_chunk);
-      const VecType* vec_thread_read_ptr =
-          reinterpret_cast<const VecType*>(thread_read_ptr);
-#pragma unroll
-      for (int ii = 0; ii < LDG_PER_THREAD; ++ii) {
-        VecType vec = vec_thread_read_ptr[ii * THREADS_PER_ROW];
-        int base_idx_f2 = ii * ELTS_PER_LDG / 2;
-#pragma unroll
-        for (int jj = 0; jj < ELTS_PER_LDG / 2; ++jj) {
-          row_chunk_f2[base_idx_f2 + jj] = __half22float2(
-              *reinterpret_cast<const __half2*>(vec.data + jj * 2));
-        }
-      }
-    } else {  // ELTS_PER_LDG == 1
-#pragma unroll
-      for (int ii = 0; ii < LDG_PER_THREAD; ++ii) {
-        const __half* scalar_ptr = thread_read_ptr + ii * THREADS_PER_ROW;
-        row_chunk[ii] = __half2float(*scalar_ptr);
-      }
-    }
-  }
-  constexpr float threshold = 20.0f;
-  constexpr float beta = 1.0f;
-
-  // Hash MoE path: indices are predetermined from lookup table
-  if constexpr (USE_HASH) {
-    const IndType token_id = input_ids[thread_row];
-    const IndType* expert_indices_for_token = tid2eid + token_id * k;
-#pragma unroll
-    for (int ii = 0; ii < VPT; ++ii) {
-      float val = row_chunk[ii];
-      float val_b = val * beta;
-      val = (val_b > threshold) ? val : (__logf(1.0f + __expf(val_b))) / beta;
-      row_chunk[ii] = sqrtf(val);
-    }
-    float selected_sum = 0.f;
-#pragma unroll
-    for (int k_idx = 0; k_idx < k; ++k_idx) {
-      const int expert = expert_indices_for_token[k_idx];
-      const int idx = k * thread_row + k_idx;
-      for (int ii = 0; ii < VPT; ++ii) {
-        const int group_id = ii / ELTS_PER_LDG;
-        const int local_id = ii % ELTS_PER_LDG;
-        const int expert_idx = first_elt_read_by_thread +
-                               group_id * THREADS_PER_ROW * ELTS_PER_LDG +
-                               local_id;
-        if (expert == expert_idx) {
-          indices[idx] = expert;
-          selected_sum += row_chunk[ii];
-          break;
-        }
-      }
-    }
-    // Compute per-thread scale (using warp reduction when renormalizing).
-    if (renormalize) {
-      selected_sum = warpReduceSum(selected_sum);
-    }
-    float scale = static_cast<float>(routed_scaling_factor);
-    if (renormalize) {
-      const float denom = selected_sum > 0.f ? selected_sum : 1.f;
-      scale /= denom;
-    }
-
-#pragma unroll
-    for (int k_idx = 0; k_idx < k; ++k_idx) {
-      const int expert = expert_indices_for_token[k_idx];
-      const int idx = k * thread_row + k_idx;
-      for (int ii = 0; ii < VPT; ++ii) {
-        const int group_id = ii / ELTS_PER_LDG;
-        const int local_id = ii % ELTS_PER_LDG;
-        const int expert_idx = first_elt_read_by_thread +
-                               group_id * THREADS_PER_ROW * ELTS_PER_LDG +
-                               local_id;
-        if (expert == expert_idx) {
-          output[idx] = row_chunk[ii] * scale;
-          break;
-        }
-      }
-    }
-#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
-    asm volatile("griddepcontrol.launch_dependents;");
-#endif
-    return;
-  }
-
-#pragma unroll
-  for (int ii = 0; ii < VPT; ++ii) {
-    float val = row_chunk[ii];
-    float val_b = val * beta;
-    // Compute softplus: log(1 + exp(val)) with numerical stability
-    // When val > threshold, softplus(x) ≈ x to avoid exp overflow
-    val = (val_b > threshold) ? val : (__logf(1.0f + __expf(val_b))) / beta;
-    val = sqrtf(val);
-    if (correction_bias) {
-      const int group_id = ii / ELTS_PER_LDG;
-      const int local_id = ii % ELTS_PER_LDG;
-      const int expert_idx = first_elt_read_by_thread +
-                             group_id * THREADS_PER_ROW * ELTS_PER_LDG +
-                             local_id;
-      val = val + correction_bias[expert_idx];
-    }
-    row_chunk[ii] = val;
-  }
-
-  // Original TopK path: find top-k experts by score
-  // Now, sigmoid_res contains the sigmoid of the row chunk. Now, I want to find
-  // the topk elements in each row, along with the max index.
-  int start_col = first_elt_read_by_thread;
-  static constexpr int COLS_PER_GROUP_LDG = ELTS_PER_LDG * THREADS_PER_ROW;
-
-  float selected_sum = 0.f;
-  for (int k_idx = 0; k_idx < k; ++k_idx) {
-    // First, each thread does the local argmax
-    float max_val = row_chunk[0];
-    int expert = start_col;
-#pragma unroll
-    for (int ldg = 0, col = start_col; ldg < LDG_PER_THREAD;
-         ++ldg, col += COLS_PER_GROUP_LDG) {
-#pragma unroll
-      for (int ii = 0; ii < ELTS_PER_LDG; ++ii) {
-        float val = row_chunk[ldg * ELTS_PER_LDG + ii];
-
-        // No check on the experts here since columns with the smallest index
-        // are processed first and only updated if > (not >=)
-        if (val > max_val) {
-          max_val = val;
-          expert = col + ii;
-        }
-      }
-    }
-
-// Now, we perform the argmax reduce. We use the butterfly pattern so threads
-// reach consensus about the max. This will be useful for K > 1 so that the
-// threads can agree on "who" had the max value. That thread can then blank out
-// their max with -inf and the warp can run more iterations...
-#pragma unroll
-    for (int mask = THREADS_PER_ROW / 2; mask > 0; mask /= 2) {
-      float other_max =
-          VLLM_SHFL_XOR_SYNC_WIDTH(max_val, mask, THREADS_PER_ROW);
-      int other_expert =
-          VLLM_SHFL_XOR_SYNC_WIDTH(expert, mask, THREADS_PER_ROW);
-
-      // We want lower indices to "win" in every thread so we break ties this
-      // way
-      if (other_max > max_val ||
-          (other_max == max_val && other_expert < expert)) {
-        max_val = other_max;
-        expert = other_expert;
-      }
-    }
-
-    // Write the max for this k iteration to global memory.
-    if (thread_group_idx == 0) {
-      // Add a guard to ignore experts not included by this node
-      const bool node_uses_expert =
-          expert >= start_expert && expert < end_expert;
-      const bool should_process_row = row_is_active && node_uses_expert;
-
-      // The lead thread from each sub-group will write out the final results to
-      // global memory. (This will be a single) thread per row of the
-      // input/output matrices.
-      const int idx = k * thread_row + k_idx;
-      if (correction_bias != nullptr) {
-        max_val -= correction_bias[expert];
-      }
-      output[idx] = max_val;
-      indices[idx] = should_process_row ? (expert - start_expert) : NUM_EXPERTS;
-      source_rows[idx] = k_idx * num_rows + thread_row;
-      if (renormalize) {
-        selected_sum += max_val;
-      }
-    }
-
-    // Finally, we clear the value in the thread with the current max if there
-    // is another iteration to run.
-    if (k_idx + 1 < k) {
-      const int ldg_group_for_expert = expert / COLS_PER_GROUP_LDG;
-      const int thread_to_clear_in_group =
-          (expert / ELTS_PER_LDG) % THREADS_PER_ROW;
-
-      // Only the thread in the group which produced the max will reset the
-      // "winning" value to -inf.
-      if (thread_group_idx == thread_to_clear_in_group) {
-        const int offset_for_expert = expert % ELTS_PER_LDG;
-        // Safe to set to any negative value since row_chunk values must be
-        // between 0 and 1.
-        row_chunk[ldg_group_for_expert * ELTS_PER_LDG + offset_for_expert] =
-            -10000.f;
-      }
-    }
-  }
-
-  // Apply renormalization and routed scaling factor to final weights.
-  if (thread_group_idx == 0) {
-    float scale = static_cast<float>(routed_scaling_factor);
-    if (renormalize) {
-      const float denom = selected_sum > 0.f ? selected_sum : 1.f;
-      scale /= denom;
-    }
-    for (int k_idx = 0; k_idx < k; ++k_idx) {
-      const int idx = k * thread_row + k_idx;
-      output[idx] = output[idx] * scale;
-    }
-  }
-#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
-  asm volatile("griddepcontrol.launch_dependents;");
-#endif
-}
-
-namespace detail {
-// Constructs some constants needed to partition the work across threads at
-// compile time.
-template <int EXPERTS, int BYTES_PER_LDG, int WARP_SIZE_PARAM,
-          typename InputType>
-struct TopkConstants {
-  static constexpr int ELTS_PER_LDG = BYTES_PER_LDG / sizeof(InputType);
-  static_assert(EXPERTS / (ELTS_PER_LDG * WARP_SIZE_PARAM) == 0 ||
-                    EXPERTS % (ELTS_PER_LDG * WARP_SIZE_PARAM) == 0,
-                "");
-  static constexpr int VECs_PER_THREAD =
-      MAX(1, EXPERTS / (ELTS_PER_LDG * WARP_SIZE_PARAM));
-  static constexpr int VPT = VECs_PER_THREAD * ELTS_PER_LDG;
-  static constexpr int THREADS_PER_ROW = EXPERTS / VPT;
-  static const int ROWS_PER_WARP = WARP_SIZE_PARAM / THREADS_PER_ROW;
-};
-}  // namespace detail
-
-#define DISPATCH_HASH(use_hash, USE_HASH, ...)                                 \
-  if (use_hash) {                                                              \
-    const bool USE_HASH = true;                                                \
-    static_assert(USE_HASH == true, "USE_HASH must be compile-time constant"); \
-    __VA_ARGS__                                                                \
-  } else {                                                                     \
-    const bool USE_HASH = false;                                               \
-    static_assert(USE_HASH == false,                                           \
-                  "USE_HASH must be compile-time constant");                   \
-    __VA_ARGS__                                                                \
-  }
-
-template <int EXPERTS, int WARPS_PER_TB, int WARP_SIZE_PARAM,
-          int MAX_BYTES_PER_LDG, typename IndType, typename InputType>
-void topkGatingSoftplusSqrtLauncherHelper(
-    const InputType* input, const bool* finished, float* output,
-    IndType* indices, int* source_row, const int num_rows, const int k,
-    const int start_expert, const int end_expert, const bool renormalize,
-    double routed_scaling_factor, const float* correction_bias,
-    const bool use_hash, const IndType* input_ids, const IndType* tid2eid,
-    cudaStream_t stream) {
-  static constexpr int BYTES_PER_LDG =
-      MIN(MAX_BYTES_PER_LDG, sizeof(InputType) * EXPERTS);
-  using Constants =
-      detail::TopkConstants<EXPERTS, BYTES_PER_LDG, WARP_SIZE_PARAM, InputType>;
-  static constexpr int VPT = Constants::VPT;
-  static constexpr int ROWS_PER_WARP = Constants::ROWS_PER_WARP;
-  const int num_warps = (num_rows + ROWS_PER_WARP - 1) / ROWS_PER_WARP;
-  const int num_blocks = (num_warps + WARPS_PER_TB - 1) / WARPS_PER_TB;
-  dim3 block_dim(WARP_SIZE_PARAM, WARPS_PER_TB);
-  DISPATCH_HASH(use_hash, USE_HASH, {
-    auto* kernel =
-        &topkGatingSoftplusSqrt<VPT, EXPERTS, WARPS_PER_TB, BYTES_PER_LDG,
-                                WARP_SIZE_PARAM, USE_HASH, IndType, InputType>;
-#ifndef USE_ROCM
-    cudaLaunchConfig_t config = {};
-    config.gridDim = num_blocks;
-    config.blockDim = block_dim;
-    config.dynamicSmemBytes = 0;
-    config.stream = stream;
-    cudaLaunchAttribute attrs[1];
-    attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
-    attrs[0].val.programmaticStreamSerializationAllowed = 1;
-    config.numAttrs = 1;
-    config.attrs = attrs;
-    cudaLaunchKernelEx(&config, kernel, input, finished, output, num_rows,
-                       indices, source_row, k, start_expert, end_expert,
-                       renormalize, routed_scaling_factor, correction_bias,
-                       input_ids, tid2eid);
-#else
-    kernel<<<num_blocks, block_dim, 0, stream>>>(
-        input, finished, output, num_rows, indices, source_row, k, start_expert,
-        end_expert, renormalize, routed_scaling_factor, correction_bias,
-        input_ids, tid2eid);
-#endif
-  })
-}
-
-#ifndef USE_ROCM
-  #define LAUNCH_SOFTPLUS_SQRT(NUM_EXPERTS, WARPS_PER_TB, MAX_BYTES)           \
-    static_assert(WARP_SIZE == 32,                                             \
-                  "Unsupported warp size. Only 32 is supported for CUDA");     \
-    topkGatingSoftplusSqrtLauncherHelper<NUM_EXPERTS, WARPS_PER_TB, WARP_SIZE, \
-                                         MAX_BYTES>(                           \
-        gating_output, nullptr, topk_weights, topk_indices,                    \
-        token_expert_indices, num_tokens, topk, 0, num_experts, renormalize,   \
-        routed_scaling_factor, correction_bias, use_hash, input_ids, tid2eid,  \
-        stream);
-#else
-  #define LAUNCH_SOFTPLUS_SQRT(NUM_EXPERTS, WARPS_PER_TB, MAX_BYTES)           \
-    if (WARP_SIZE == 64) {                                                     \
-      topkGatingSoftplusSqrtLauncherHelper<NUM_EXPERTS, WARPS_PER_TB, 64,      \
-                                           MAX_BYTES>(                         \
-          gating_output, nullptr, topk_weights, topk_indices,                  \
-          token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \
-          routed_scaling_factor, correction_bias, use_hash, input_ids,         \
-          tid2eid, stream);                                                    \
-    } else if (WARP_SIZE == 32) {                                              \
-      topkGatingSoftplusSqrtLauncherHelper<NUM_EXPERTS, WARPS_PER_TB, 32,      \
-                                           MAX_BYTES>(                         \
-          gating_output, nullptr, topk_weights, topk_indices,                  \
-          token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \
-          routed_scaling_factor, correction_bias, use_hash, input_ids,         \
-          tid2eid, stream);                                                    \
-    } else {                                                                   \
-      assert(false &&                                                          \
-             "Unsupported warp size. Only 32 and 64 are supported for ROCm");  \
-    }
-#endif
-
-template <typename IndType, typename InputType>
-void topkGatingSoftplusSqrtKernelLauncher(
-    const InputType* gating_output, float* topk_weights, IndType* topk_indices,
-    int* token_expert_indices, const int num_tokens, const int num_experts,
-    const int topk, const bool renormalize, double routed_scaling_factor,
-    const float* correction_bias, const bool use_hash, const IndType* input_ids,
-    const IndType* tid2eid, cudaStream_t stream) {
-  static constexpr int WARPS_PER_TB = 4;
-  static constexpr int BYTES_PER_LDG_POWER_OF_2 = 16;
-#ifndef USE_ROCM
-  // for bfloat16 dtype, we need 4 bytes loading to make sure num_experts
-  // elements can be loaded by a warp
-  static constexpr int BYTES_PER_LDG_MULTIPLE_64 =
-      (std::is_same_v<InputType, __nv_bfloat16> ||
-       std::is_same_v<InputType, __half>)
-          ? 4
-          : 8;
-#endif
-  switch (num_experts) {
-    case 1:
-      LAUNCH_SOFTPLUS_SQRT(1, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
-      break;
-    case 2:
-      LAUNCH_SOFTPLUS_SQRT(2, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
-      break;
-    case 4:
-      LAUNCH_SOFTPLUS_SQRT(4, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
-      break;
-    case 8:
-      LAUNCH_SOFTPLUS_SQRT(8, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
-      break;
-    case 16:
-      LAUNCH_SOFTPLUS_SQRT(16, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
-      break;
-    case 32:
-      LAUNCH_SOFTPLUS_SQRT(32, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
-      break;
-    case 64:
-      LAUNCH_SOFTPLUS_SQRT(64, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
-      break;
-    case 128:
-      LAUNCH_SOFTPLUS_SQRT(128, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
-      break;
-    case 256:
-      LAUNCH_SOFTPLUS_SQRT(256, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
-      break;
-    case 512:
-      LAUNCH_SOFTPLUS_SQRT(512, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
-      break;
-      // (CUDA only) support multiples of 64 when num_experts is not power of 2.
-      // ROCm uses WARP_SIZE 64 so 8 bytes loading won't fit for some of
-      // num_experts, alternatively we can test 4 bytes loading and enable it in
-      // future.
-#ifndef USE_ROCM
-    case 192:
-      LAUNCH_SOFTPLUS_SQRT(192, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
-      break;
-    case 320:
-      LAUNCH_SOFTPLUS_SQRT(320, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
-      break;
-    case 384:
-      LAUNCH_SOFTPLUS_SQRT(384, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
-      break;
-    case 448:
-      LAUNCH_SOFTPLUS_SQRT(448, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
-      break;
-    case 576:
-      LAUNCH_SOFTPLUS_SQRT(576, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
-      break;
-#endif
-    default: {
-      TORCH_CHECK(false, "Unsupported expert number: ", num_experts);
-    }
-  }
-}
-
-}  // namespace moe
-}  // namespace vllm
-
-template <typename ComputeType>
-void dispatch_topk_softplus_sqrt_launch(
-    const ComputeType* gating_output, torch::Tensor& topk_weights,
-    torch::Tensor& topk_indices, torch::Tensor& token_expert_indices,
-    int num_tokens, int num_experts, int topk, bool renormalize,
-    double routed_scaling_factor,
-    const c10::optional<torch::Tensor>& correction_bias,
-    const c10::optional<torch::Tensor>& input_ids,
-    const c10::optional<torch::Tensor>& tid2eid, cudaStream_t stream) {
-  const float* bias_ptr = nullptr;
-  if (correction_bias.has_value()) {
-    bias_ptr = correction_bias.value().data_ptr<float>();
-  }
-  bool use_hash = false;
-  if (tid2eid.has_value()) {
-    TORCH_CHECK(input_ids.has_value(), "input_ids is required for hash MoE");
-    use_hash = true;
-  }
-  if (topk_indices.scalar_type() == at::ScalarType::Int) {
-    const int* input_ids_ptr = nullptr;
-    const int* tid2eid_ptr = nullptr;
-    if (tid2eid.has_value()) {
-      input_ids_ptr = input_ids.value().data_ptr<int>();
-      tid2eid_ptr = tid2eid.value().data_ptr<int>();
-    }
-
-    vllm::moe::topkGatingSoftplusSqrtKernelLauncher<int, ComputeType>(
-        gating_output, topk_weights.data_ptr<float>(),
-        topk_indices.data_ptr<int>(), token_expert_indices.data_ptr<int>(),
-        num_tokens, num_experts, topk, renormalize, routed_scaling_factor,
-        bias_ptr, use_hash, input_ids_ptr, tid2eid_ptr, stream);
-  } else if (topk_indices.scalar_type() == at::ScalarType::UInt32) {
-    const uint32_t* input_ids_ptr = nullptr;
-    const uint32_t* tid2eid_ptr = nullptr;
-    if (tid2eid.has_value()) {
-      input_ids_ptr = input_ids.value().data_ptr<uint32_t>();
-      tid2eid_ptr = tid2eid.value().data_ptr<uint32_t>();
-    }
-    vllm::moe::topkGatingSoftplusSqrtKernelLauncher<uint32_t, ComputeType>(
-        gating_output, topk_weights.data_ptr<float>(),
-        topk_indices.data_ptr<uint32_t>(), token_expert_indices.data_ptr<int>(),
-        num_tokens, num_experts, topk, renormalize, routed_scaling_factor,
-        bias_ptr, use_hash, input_ids_ptr, tid2eid_ptr, stream);
-  } else {
-    TORCH_CHECK(topk_indices.scalar_type() == at::ScalarType::Long);
-
-    const int64_t* input_ids_ptr = nullptr;
-    const int64_t* tid2eid_ptr = nullptr;
-    if (tid2eid.has_value()) {
-      input_ids_ptr = input_ids.value().data_ptr<int64_t>();
-      tid2eid_ptr = tid2eid.value().data_ptr<int64_t>();
-    }
-
-    vllm::moe::topkGatingSoftplusSqrtKernelLauncher<int64_t, ComputeType>(
-        gating_output, topk_weights.data_ptr<float>(),
-        topk_indices.data_ptr<int64_t>(), token_expert_indices.data_ptr<int>(),
-        num_tokens, num_experts, topk, renormalize, routed_scaling_factor,
-        bias_ptr, use_hash, input_ids_ptr, tid2eid_ptr, stream);
-  }
-}
-
-void topk_softplus_sqrt(
-    torch::Tensor& topk_weights,          // [num_tokens, topk]
-    torch::Tensor& topk_indices,          // [num_tokens, topk]
-    torch::Tensor& token_expert_indices,  // [num_tokens, topk]
-    torch::Tensor& gating_output,         // [num_tokens, num_experts]
-    bool renormalize, double routed_scaling_factor,
-    const c10::optional<torch::Tensor>& correction_bias,
-    const c10::optional<torch::Tensor>& input_ids,
-    const c10::optional<torch::Tensor>& tid2eid) {
-  const int num_experts = gating_output.size(-1);
-  const auto num_tokens = gating_output.numel() / num_experts;
-  const int topk = topk_weights.size(-1);
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(gating_output));
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
-
-  if (gating_output.scalar_type() == at::ScalarType::Float) {
-    dispatch_topk_softplus_sqrt_launch<float>(
-        gating_output.data_ptr<float>(), topk_weights, topk_indices,
-        token_expert_indices, num_tokens, num_experts, topk, renormalize,
-        routed_scaling_factor, correction_bias, input_ids, tid2eid, stream);
-  } else if (gating_output.scalar_type() == at::ScalarType::Half) {
-    dispatch_topk_softplus_sqrt_launch<__half>(
-        reinterpret_cast<const __half*>(gating_output.data_ptr<at::Half>()),
-        topk_weights, topk_indices, token_expert_indices, num_tokens,
-        num_experts, topk, renormalize, routed_scaling_factor, correction_bias,
-        input_ids, tid2eid, stream);
-  } else if (gating_output.scalar_type() == at::ScalarType::BFloat16) {
-    dispatch_topk_softplus_sqrt_launch<__nv_bfloat16>(
-        reinterpret_cast<const __nv_bfloat16*>(
-            gating_output.data_ptr<at::BFloat16>()),
-        topk_weights, topk_indices, token_expert_indices, num_tokens,
-        num_experts, topk, renormalize, routed_scaling_factor, correction_bias,
-        input_ids, tid2eid, stream);
-  } else {
-    TORCH_CHECK(false, "Unsupported gating_output data type: ",
-                gating_output.scalar_type());
-  }
-}

csrc/moe/torch_bindings.cpp

@@ -16,14 +16,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) {
       "bias) -> ()");
   m.impl("topk_sigmoid", torch::kCUDA, &topk_sigmoid);
 
-#ifndef USE_ROCM
-  m.def(
-      "topk_softplus_sqrt(Tensor! topk_weights, Tensor! topk_indices, Tensor! "
-      "token_expert_indices, Tensor gating_output, bool renormalize, float "
-      "routed_scaling_factor, Tensor? "
-      "bias, Tensor? input_ids, Tensor? tid2eid) -> ()");
-  m.impl("topk_softplus_sqrt", torch::kCUDA, &topk_softplus_sqrt);
-#endif
   // Calculate the result of moe by summing up the partial results
   // from all selected experts.
   m.def("moe_sum(Tensor input, Tensor! output) -> ()");

csrc/ops.h

@@ -100,11 +100,6 @@ void fused_qk_norm_rope(torch::Tensor& qkv, int64_t num_heads_q,
                         bool is_neox, torch::Tensor& position_ids,
                         int64_t forced_token_heads_per_warp);
 
-void fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
-    torch::Tensor& q, torch::Tensor const& kv, torch::Tensor& k_cache,
-    torch::Tensor const& slot_mapping, torch::Tensor const& position_ids,
-    torch::Tensor const& cos_sin_cache, double eps, int64_t cache_block_size);
-
 void apply_repetition_penalties_(torch::Tensor& logits,
                                  const torch::Tensor& prompt_mask,
                                  const torch::Tensor& output_mask,
@@ -158,13 +153,10 @@ void silu_and_mul_per_block_quant(torch::Tensor& out,
 
 void rotary_embedding(torch::Tensor& positions, torch::Tensor& query,
                       std::optional<torch::Tensor> key, int64_t head_size,
-                      torch::Tensor& cos_sin_cache, bool is_neox,
-                      int64_t rope_dim_offset, bool inverse);
+                      torch::Tensor& cos_sin_cache, bool is_neox);
 
 void silu_and_mul(torch::Tensor& out, torch::Tensor& input);
 
-void silu_and_mul_clamp(torch::Tensor& out, torch::Tensor& input, double limit);
-
 void silu_and_mul_quant(torch::Tensor& out, torch::Tensor& input,
                         torch::Tensor& scale);
 

csrc/persistent_topk.cuh

@@ -18,6 +18,7 @@ namespace persistent {
 // Constants
 // ============================================================================
 
+constexpr int TopK = 2048;
 constexpr int kThreadsPerBlock = 1024;
 constexpr int RADIX = 256;
 
@@ -127,12 +128,11 @@ struct RadixRowState {
 
 struct PersistentTopKParams {
   const float* __restrict__ input;  // [num_rows, stride]
-  int32_t* __restrict__ output;     // [num_rows, top_k]
+  int32_t* __restrict__ output;     // [num_rows, TopK]
   int32_t* __restrict__ lengths;    // [num_rows]
   RadixRowState* row_states;        // large path: per-group state
   uint32_t num_rows;
   uint32_t stride;
-  uint32_t top_k;           // actual k value for output stride
   uint32_t chunk_size;      // large path: elements per CTA
   uint32_t ctas_per_group;  // 1=medium, >1=large
   uint32_t max_seq_len;     // max seq_len across all rows (for early CTA exit)
@@ -154,7 +154,6 @@ __device__ __forceinline__ uint32_t decode_bin(float x) {
   return key >> 5;
 }
 
-template <int TopK>
 __device__ __noinline__ void histogram_2048_topk(
     const float* __restrict__ logits, int32_t* __restrict__ output_indices,
     int32_t seq_len) {
@@ -419,7 +418,6 @@ __device__ __noinline__ void histogram_2048_topk(
 // by: DarkSharpness
 // which at the same time is an optimized topk kernel copied from tilelang
 // kernel
-template <int TopK>
 __device__ __noinline__ void histogram_256_topk(
     const float* __restrict__ logits, int* __restrict__ output_indices,
     int logits_offset, int seq_len) {
@@ -651,7 +649,7 @@ __device__ __forceinline__ void wait_ge(int* ptr, int target_val,
 // Adapted from https://github.com/flashinfer-ai/flashinfer/pull/2215
 // ============================================================================
 
-template <int TopK, uint32_t VEC_SIZE>
+template <uint32_t VEC_SIZE>
 __device__ void radix_topk(const float* __restrict__ row_input,
                            int32_t* __restrict__ row_output, uint32_t seq_len,
                            uint32_t my_chunk_start, uint32_t chunk_size,
@@ -859,7 +857,7 @@ __device__ void radix_topk(const float* __restrict__ row_input,
 // see filtered_topk.cuh)
 // ============================================================================
 
-template <int TopK = 2048, uint32_t VEC_SIZE = 1>
+template <uint32_t VEC_SIZE = 1>
 __global__ void __launch_bounds__(kThreadsPerBlock, 2)
     persistent_topk_kernel(PersistentTopKParams params) {
   const uint32_t tx = threadIdx.x;
@@ -917,7 +915,7 @@ __global__ void __launch_bounds__(kThreadsPerBlock, 2)
     if (row_idx >= params.num_rows) break;
 
     const uint32_t seq_len = params.lengths[row_idx];
-    int32_t* row_output = params.output + row_idx * params.top_k;
+    int32_t* row_output = params.output + row_idx * TopK;
     const float* row_input = params.input + row_idx * params.stride;
 
     if (seq_len <= RADIX_THRESHOLD) {
@@ -929,19 +927,19 @@ __global__ void __launch_bounds__(kThreadsPerBlock, 2)
             row_output[i] = (i < seq_len) ? static_cast<int32_t>(i) : -1;
           }
         } else if (seq_len <= static_cast<uint32_t>(HIST2048_THRESHOLD)) {
-          histogram_2048_topk<TopK>(row_input, row_output, seq_len);
+          histogram_2048_topk(row_input, row_output, seq_len);
         } else {
-          histogram_256_topk<TopK>(row_input, row_output, 0, seq_len);
+          histogram_256_topk(row_input, row_output, 0, seq_len);
         }
       }
       continue;
     }
 
     const uint32_t my_chunk_start = cta_in_group * chunk_size;
-    radix_topk<TopK, VEC_SIZE>(
-        row_input, row_output, seq_len, my_chunk_start, chunk_size,
-        local_histogram, suffix_sum, shared_scalars, shared_ordered, state,
-        cta_in_group, ctas_per_group, barrier_phase, iter, tx);
+    radix_topk<VEC_SIZE>(row_input, row_output, seq_len, my_chunk_start,
+                         chunk_size, local_histogram, suffix_sum,
+                         shared_scalars, shared_ordered, state, cta_in_group,
+                         ctas_per_group, barrier_phase, iter, tx);
   }
 }
 
@@ -1013,6 +1011,7 @@ struct FilteredTopKTraits<float> {
   }
 };
 
+constexpr uint32_t FILTERED_TOPK_MAX_K = 2048;
 constexpr uint32_t FILTERED_TOPK_BLOCK_THREADS = 1024;
 constexpr uint32_t FILTERED_TOPK_SMEM_INPUT_SIZE =
     16 * 1024;  // 16K indices per buffer
@@ -1026,7 +1025,7 @@ constexpr size_t FILTERED_TOPK_SMEM_DYNAMIC =
  * \tparam IdType Index type (int32_t)
  * \tparam VEC_SIZE Vector size for input loads (1, 2, 4, or 8)
  */
-template <typename DType, typename IdType, int VEC_SIZE, uint32_t MAX_K = 2048>
+template <typename DType, typename IdType, int VEC_SIZE>
 __global__ void __launch_bounds__(FILTERED_TOPK_BLOCK_THREADS)
     FilteredTopKUnifiedKernel(const DType* __restrict__ input,
                               IdType* __restrict__ output,
@@ -1060,7 +1059,7 @@ __global__ void __launch_bounds__(FILTERED_TOPK_BLOCK_THREADS)
   alignas(128) __shared__ int s_counter;
   alignas(128) __shared__ int s_threshold_bin_id;
   alignas(128) __shared__ int s_num_input[2];
-  alignas(128) __shared__ int s_indices[MAX_K];
+  alignas(128) __shared__ int s_indices[FILTERED_TOPK_MAX_K];
 
   auto& s_histogram = s_histogram_buf[0];
 
@@ -1281,7 +1280,7 @@ constexpr int ComputeFilteredTopKVecSize(uint32_t max_len) {
   return static_cast<int>(g);
 }
 
-template <typename DType, typename IdType, uint32_t MAX_K = 2048>
+template <typename DType, typename IdType>
 cudaError_t FilteredTopKRaggedTransform(DType* input, IdType* output_indices,
                                         IdType* lengths, uint32_t num_rows,
                                         uint32_t top_k_val, uint32_t max_len,
@@ -1298,7 +1297,7 @@ cudaError_t FilteredTopKRaggedTransform(DType* input, IdType* output_indices,
 
 #define DISPATCH_VEC_SIZE(VS)                                               \
   if (vec_size == VS) {                                                     \
-    auto kernel = FilteredTopKUnifiedKernel<DType, IdType, VS, MAX_K>;      \
+    auto kernel = FilteredTopKUnifiedKernel<DType, IdType, VS>;             \
     FLASHINFER_CUDA_CALL(cudaFuncSetAttribute(                              \
         kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size));   \
     FLASHINFER_CUDA_CALL(cudaLaunchKernel((void*)kernel, grid, block, args, \

csrc/pos_encoding_kernels.cu

@@ -9,29 +9,28 @@ namespace vllm {
 
 template <typename scalar_t, bool IS_NEOX>
 inline __device__ void apply_token_rotary_embedding(
-    scalar_t* __restrict__ arr, const float* __restrict__ cos_ptr,
-    const float* __restrict__ sin_ptr, int rot_offset, int embed_dim,
-    const bool inverse) {
+    scalar_t* __restrict__ arr, const scalar_t* __restrict__ cos_ptr,
+    const scalar_t* __restrict__ sin_ptr, int rot_offset, int embed_dim) {
   int x_index, y_index;
-  float cos_f, sin_f;
+  scalar_t cos, sin;
   if (IS_NEOX) {
+    // GPT-NeoX style rotary embedding.
     x_index = rot_offset;
     y_index = embed_dim + rot_offset;
-    cos_f = VLLM_LDG(cos_ptr + x_index);
-    sin_f = VLLM_LDG(sin_ptr + x_index);
+    cos = VLLM_LDG(cos_ptr + x_index);
+    sin = VLLM_LDG(sin_ptr + x_index);
   } else {
+    // GPT-J style rotary embedding.
     x_index = 2 * rot_offset;
     y_index = 2 * rot_offset + 1;
-    cos_f = VLLM_LDG(cos_ptr + x_index / 2);
-    sin_f = VLLM_LDG(sin_ptr + x_index / 2);
+    cos = VLLM_LDG(cos_ptr + x_index / 2);
+    sin = VLLM_LDG(sin_ptr + x_index / 2);
   }
-  if (inverse) {
-    sin_f = -sin_f;
-  }
-  const float x_f = static_cast<float>(arr[x_index]);
-  const float y_f = static_cast<float>(arr[y_index]);
-  arr[x_index] = static_cast<scalar_t>(x_f * cos_f - y_f * sin_f);
-  arr[y_index] = static_cast<scalar_t>(y_f * cos_f + x_f * sin_f);
+
+  const scalar_t x = arr[x_index];
+  const scalar_t y = arr[y_index];
+  arr[x_index] = x * cos - y * sin;
+  arr[y_index] = y * cos + x * sin;
 }
 
 template <typename scalar_t, bool IS_NEOX>
@@ -43,23 +42,22 @@ inline __device__ void apply_rotary_embedding(
                                    // [batch_size, seq_len, num_kv_heads,
                                    // head_size] or [num_tokens, num_kv_heads,
                                    // head_size]
-    const float* cache_ptr, const int head_size, const int num_heads,
+    const scalar_t* cache_ptr, const int head_size, const int num_heads,
     const int num_kv_heads, const int rot_dim, const int token_idx,
     const int64_t query_stride, const int64_t key_stride,
-    const int64_t head_stride, const int64_t rope_dim_offset,
-    const bool inverse) {
+    const int64_t head_stride) {
   const int embed_dim = rot_dim / 2;
-  const float* cos_ptr = cache_ptr;
-  const float* sin_ptr = cache_ptr + embed_dim;
+  const scalar_t* cos_ptr = cache_ptr;
+  const scalar_t* sin_ptr = cache_ptr + embed_dim;
 
   const int nq = num_heads * embed_dim;
   for (int i = threadIdx.x; i < nq; i += blockDim.x) {
     const int head_idx = i / embed_dim;
     const int64_t token_head =
-        token_idx * query_stride + head_idx * head_stride + rope_dim_offset;
+        token_idx * query_stride + head_idx * head_stride;
     const int rot_offset = i % embed_dim;
     apply_token_rotary_embedding<scalar_t, IS_NEOX>(
-        query + token_head, cos_ptr, sin_ptr, rot_offset, embed_dim, inverse);
+        query + token_head, cos_ptr, sin_ptr, rot_offset, embed_dim);
   }
 
   if (key != nullptr) {
@@ -67,10 +65,10 @@ inline __device__ void apply_rotary_embedding(
     for (int i = threadIdx.x; i < nk; i += blockDim.x) {
       const int head_idx = i / embed_dim;
       const int64_t token_head =
-          token_idx * key_stride + head_idx * head_stride + rope_dim_offset;
+          token_idx * key_stride + head_idx * head_stride;
       const int rot_offset = i % embed_dim;
       apply_token_rotary_embedding<scalar_t, IS_NEOX>(
-          key + token_head, cos_ptr, sin_ptr, rot_offset, embed_dim, inverse);
+          key + token_head, cos_ptr, sin_ptr, rot_offset, embed_dim);
     }
   }
 }
@@ -86,18 +84,19 @@ __global__ void rotary_embedding_kernel(
                                  // [batch_size, seq_len, num_kv_heads,
                                  // head_size] or [num_tokens, num_kv_heads,
                                  // head_size]
-    const float* __restrict__ cos_sin_cache,  // [max_position, rot_dim] fp32
+    const scalar_t* __restrict__ cos_sin_cache,  // [max_position, 2, rot_dim //
+                                                 // 2]
     const int rot_dim, const int64_t query_stride, const int64_t key_stride,
     const int64_t head_stride, const int num_heads, const int num_kv_heads,
-    const int head_size, const int64_t rope_dim_offset, const bool inverse) {
+    const int head_size) {
+  // Each thread block is responsible for one token.
   const int token_idx = blockIdx.x;
   int64_t pos = positions[token_idx];
-  const float* cache_ptr = cos_sin_cache + pos * rot_dim;
+  const scalar_t* cache_ptr = cos_sin_cache + pos * rot_dim;
 
   apply_rotary_embedding<scalar_t, IS_NEOX>(
       query, key, cache_ptr, head_size, num_heads, num_kv_heads, rot_dim,
-      token_idx, query_stride, key_stride, head_stride, rope_dim_offset,
-      inverse);
+      token_idx, query_stride, key_stride, head_stride);
 }
 
 }  // namespace vllm
@@ -116,7 +115,7 @@ void rotary_embedding(
     // [num_tokens, num_heads, head_size]
     int64_t head_size,
     torch::Tensor& cos_sin_cache,  // [max_position, rot_dim]
-    bool is_neox, int64_t rope_dim_offset, bool inverse) {
+    bool is_neox) {
   // num_tokens = batch_size * seq_len
   int64_t num_tokens = positions.numel();
   int positions_ndim = positions.dim();
@@ -155,8 +154,6 @@ void rotary_embedding(
   int seq_dim_idx = positions_ndim - 1;
   int64_t query_stride = query.stride(seq_dim_idx);
   int64_t key_stride = key.has_value() ? key->stride(seq_dim_idx) : 0;
-
-  TORCH_CHECK((rot_dim + rope_dim_offset) <= head_size);
   // Determine head stride: for [*, heads, head_size] use stride of last dim;
   // for flat [*, heads*head_size], heads blocks are contiguous of size
   // head_size
@@ -168,23 +165,20 @@ void rotary_embedding(
   dim3 block(std::min<int64_t>(num_heads * rot_dim / 2, 512));
   const at::cuda::OptionalCUDAGuard device_guard(device_of(query));
   const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
-  auto cache_f32 = cos_sin_cache.to(torch::kFloat32);
   VLLM_DISPATCH_FLOATING_TYPES(query.scalar_type(), "rotary_embedding", [&] {
     if (is_neox) {
       vllm::rotary_embedding_kernel<scalar_t, true><<<grid, block, 0, stream>>>(
           positions.data_ptr<int64_t>(), query.data_ptr<scalar_t>(),
           key.has_value() ? key->data_ptr<scalar_t>() : nullptr,
-          cache_f32.data_ptr<float>(), rot_dim, query_stride, key_stride,
-          head_stride, num_heads, num_kv_heads, head_size, rope_dim_offset,
-          inverse);
+          cos_sin_cache.data_ptr<scalar_t>(), rot_dim, query_stride, key_stride,
+          head_stride, num_heads, num_kv_heads, head_size);
     } else {
       vllm::rotary_embedding_kernel<scalar_t, false>
           <<<grid, block, 0, stream>>>(
               positions.data_ptr<int64_t>(), query.data_ptr<scalar_t>(),
               key.has_value() ? key->data_ptr<scalar_t>() : nullptr,
-              cache_f32.data_ptr<float>(), rot_dim, query_stride, key_stride,
-              head_stride, num_heads, num_kv_heads, head_size, rope_dim_offset,
-              inverse);
+              cos_sin_cache.data_ptr<scalar_t>(), rot_dim, query_stride,
+              key_stride, head_stride, num_heads, num_kv_heads, head_size);
     }
   });
 }

csrc/sampler.cu

@@ -258,13 +258,7 @@ __device__ bool processHistogramStep(
   auto processBins = [&](float logit, int idx) {
     if (isPartialMatch<patternShift>(logit, logitPattern)) {
       uint32_t binIdx = extractBinIdx<step>(logit);
-      // Only write elements with binIdx < thresholdBinIdx when:
-      // 1. This is step 0 and the threshold bin is small enough (no step 1)
-      // 2. This is step >= 1 (where pattern matching filters correctly)
-      // This prevents duplicates when step 0 and step 1 both run.
-      bool shouldWriteDirectly =
-          (step == 0 && smemFinalBinSize[0] <= kNumFinalItems) || (step >= 1);
-      if (binIdx < thresholdBinIdx && shouldWriteDirectly) {
+      if (binIdx < thresholdBinIdx) {
         // The element is part of the top-k selection
         int dstIdx = atomicAdd(&smemFoundTopKValues[0], 1);
 

csrc/topk.cu

@@ -10,17 +10,33 @@
   #include "persistent_topk.cuh"
 #endif
 
-namespace {
-
+void persistent_topk(const torch::Tensor& logits, const torch::Tensor& lengths,
+                     torch::Tensor& output, torch::Tensor& workspace, int64_t k,
+                     int64_t max_seq_len) {
 #ifndef USE_ROCM
-template <int TopK>
-void launch_persistent_topk(const torch::Tensor& logits,
-                            const torch::Tensor& lengths, torch::Tensor& output,
-                            torch::Tensor& workspace, int64_t max_seq_len) {
-  namespace P = vllm::persistent;
+  TORCH_CHECK(logits.is_cuda(), "logits must be CUDA tensor");
+  TORCH_CHECK(lengths.is_cuda(), "lengths must be CUDA tensor");
+  TORCH_CHECK(output.is_cuda(), "output must be CUDA tensor");
+  TORCH_CHECK(logits.dtype() == torch::kFloat32, "Only float32 supported");
+  TORCH_CHECK(lengths.dtype() == torch::kInt32, "lengths must be int32");
+  TORCH_CHECK(output.dtype() == torch::kInt32, "output must be int32");
+  TORCH_CHECK(logits.dim() == 2, "logits must be 2D");
+  TORCH_CHECK(lengths.dim() == 1 || lengths.dim() == 2,
+              "lengths must be 1D or 2D");
+  TORCH_CHECK(lengths.is_contiguous(), "lengths must be contiguous");
+  TORCH_CHECK(output.dim() == 2, "output must be 2D");
 
   const int64_t num_rows = logits.size(0);
   const int64_t stride = logits.size(1);
+
+  TORCH_CHECK(lengths.numel() == num_rows, "lengths size mismatch");
+  TORCH_CHECK(output.size(0) == num_rows && output.size(1) == k,
+              "output size mismatch");
+  namespace P = vllm::persistent;
+
+  TORCH_CHECK(k == P::TopK, "k must be 2048");
+  TORCH_CHECK(k <= stride, "k out of range");
+
   cudaStream_t stream = at::cuda::getCurrentCUDAStream();
 
   static int num_sms = 0;
@@ -34,17 +50,18 @@ void launch_persistent_topk(const torch::Tensor& logits,
   }
 
   if (num_rows > 32 && max_smem_per_block >= 128 * 1024) {
-    cudaError_t status =
-        vllm::FilteredTopKRaggedTransform<float, int32_t, TopK>(
-            logits.data_ptr<float>(), output.data_ptr<int32_t>(),
-            lengths.data_ptr<int32_t>(), static_cast<uint32_t>(num_rows),
-            static_cast<uint32_t>(TopK), static_cast<uint32_t>(stride), stream);
+    cudaError_t status = vllm::FilteredTopKRaggedTransform<float, int32_t>(
+        logits.data_ptr<float>(), output.data_ptr<int32_t>(),
+        lengths.data_ptr<int32_t>(), static_cast<uint32_t>(num_rows),
+        static_cast<uint32_t>(k), static_cast<uint32_t>(stride), stream);
     TORCH_CHECK(status == cudaSuccess,
                 "FilteredTopK failed: ", cudaGetErrorString(status));
   } else {
     TORCH_CHECK(workspace.is_cuda(), "workspace must be CUDA tensor");
     TORCH_CHECK(workspace.dtype() == torch::kUInt8, "workspace must be uint8");
 
+    // Smem cap: smaller smem → more CTAs/group → more per-row parallelism for
+    // large path. Empirically tuned.
     int effective_max_smem;
     if (num_rows <= 4) {
       effective_max_smem =
@@ -84,7 +101,7 @@ void launch_persistent_topk(const torch::Tensor& logits,
 
     int occupancy = 1;
     cudaOccupancyMaxActiveBlocksPerMultiprocessor(
-        &occupancy, P::persistent_topk_kernel<TopK, 4>, P::kThreadsPerBlock,
+        &occupancy, P::persistent_topk_kernel<4>, P::kThreadsPerBlock,
         smem_size);
     if (occupancy < 1) occupancy = 1;
 
@@ -104,16 +121,15 @@ void launch_persistent_topk(const torch::Tensor& logits,
     params.lengths = lengths.data_ptr<int32_t>();
     params.num_rows = static_cast<uint32_t>(num_rows);
     params.stride = static_cast<uint32_t>(stride);
-    params.top_k = static_cast<uint32_t>(TopK);
     params.chunk_size = chunk_size;
     params.row_states =
         reinterpret_cast<P::RadixRowState*>(workspace.data_ptr<uint8_t>());
     params.ctas_per_group = ctas_per_group;
     params.max_seq_len = static_cast<uint32_t>(max_seq_len);
 
-  #define LAUNCH_PERSISTENT(TOPK_VAL, VS)                                     \
+  #define LAUNCH_PERSISTENT(VS)                                               \
     do {                                                                      \
-      auto kernel = &P::persistent_topk_kernel<TOPK_VAL, VS>;                 \
+      auto kernel = &P::persistent_topk_kernel<VS>;                           \
       cudaError_t err = cudaFuncSetAttribute(                                 \
           kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);    \
       TORCH_CHECK(err == cudaSuccess,                                         \
@@ -122,11 +138,11 @@ void launch_persistent_topk(const torch::Tensor& logits,
     } while (0)
 
     if (vec_size == 4) {
-      LAUNCH_PERSISTENT(TopK, 4);
+      LAUNCH_PERSISTENT(4);
     } else if (vec_size == 2) {
-      LAUNCH_PERSISTENT(TopK, 2);
+      LAUNCH_PERSISTENT(2);
     } else {
-      LAUNCH_PERSISTENT(TopK, 1);
+      LAUNCH_PERSISTENT(1);
     }
   #undef LAUNCH_PERSISTENT
   }
@@ -134,46 +150,6 @@ void launch_persistent_topk(const torch::Tensor& logits,
   cudaError_t err = cudaGetLastError();
   TORCH_CHECK(err == cudaSuccess,
               "persistent_topk failed: ", cudaGetErrorString(err));
-}
-#endif
-
-}  // anonymous namespace
-
-void persistent_topk(const torch::Tensor& logits, const torch::Tensor& lengths,
-                     torch::Tensor& output, torch::Tensor& workspace, int64_t k,
-                     int64_t max_seq_len) {
-#ifndef USE_ROCM
-  TORCH_CHECK(logits.is_cuda(), "logits must be CUDA tensor");
-  TORCH_CHECK(lengths.is_cuda(), "lengths must be CUDA tensor");
-  TORCH_CHECK(output.is_cuda(), "output must be CUDA tensor");
-  TORCH_CHECK(logits.dtype() == torch::kFloat32, "Only float32 supported");
-  TORCH_CHECK(lengths.dtype() == torch::kInt32, "lengths must be int32");
-  TORCH_CHECK(output.dtype() == torch::kInt32, "output must be int32");
-  TORCH_CHECK(logits.dim() == 2, "logits must be 2D");
-  TORCH_CHECK(lengths.dim() == 1 || lengths.dim() == 2,
-              "lengths must be 1D or 2D");
-  TORCH_CHECK(lengths.is_contiguous(), "lengths must be contiguous");
-  TORCH_CHECK(output.dim() == 2, "output must be 2D");
-
-  const int64_t num_rows = logits.size(0);
-  const int64_t stride = logits.size(1);
-
-  TORCH_CHECK(lengths.numel() == num_rows, "lengths size mismatch");
-  TORCH_CHECK(output.size(0) == num_rows && output.size(1) == k,
-              "output size mismatch");
-  TORCH_CHECK(k == 512 || k == 1024 || k == 2048,
-              "persistent_topk supports k=512, k=1024, or k=2048, got k=", k);
-
-  if (k == 512) {
-    launch_persistent_topk<512>(logits, lengths, output, workspace,
-                                max_seq_len);
-  } else if (k == 1024) {
-    launch_persistent_topk<1024>(logits, lengths, output, workspace,
-                                 max_seq_len);
-  } else {
-    launch_persistent_topk<2048>(logits, lengths, output, workspace,
-                                 max_seq_len);
-  }
 #else
   TORCH_CHECK(false, "persistent_topk is not supported on ROCm");
 #endif

csrc/torch_bindings.cpp

@@ -106,12 +106,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
   ops.def("silu_and_mul(Tensor! result, Tensor input) -> ()");
   ops.impl("silu_and_mul", torch::kCUDA, &silu_and_mul);
 
-  // SwiGLU activation with input clamping.
-  ops.def(
-      "silu_and_mul_with_clamp(Tensor! result, Tensor input, float limit) "
-      "-> ()");
-  ops.impl("silu_and_mul_with_clamp", torch::kCUDA, &silu_and_mul_clamp);
-
   ops.def(
       "silu_and_mul_quant(Tensor! result, Tensor input, Tensor scale) -> ()");
   ops.impl("silu_and_mul_quant", torch::kCUDA, &silu_and_mul_quant);
@@ -183,19 +177,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
       "int forced_token_heads_per_warp=-1) -> ()");
   ops.impl("fused_qk_norm_rope", torch::kCUDA, &fused_qk_norm_rope);
 
-#ifndef USE_ROCM
-  // Horizontally-fused DeepseekV4-MLA: per-head RMSNorm + GPT-J RoPE for Q, and
-  // GPT-J RoPE + UE8M0 FP8 quant + paged cache insert for KV, all in one
-  // kernel launch.
-  ops.def(
-      "fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert("
-      "Tensor! q, Tensor kv, Tensor! k_cache, "
-      "Tensor slot_mapping, Tensor position_ids, Tensor cos_sin_cache, "
-      "float eps, int cache_block_size) -> ()");
-  ops.impl("fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert", torch::kCUDA,
-           &fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert);
-#endif
-
   // Apply repetition penalties to logits in-place
   ops.def(
       "apply_repetition_penalties_(Tensor! logits, Tensor prompt_mask, "
@@ -259,8 +240,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
   ops.def(
       "rotary_embedding(Tensor positions, Tensor! query,"
       "                 Tensor!? key, int head_size,"
-      "                 Tensor cos_sin_cache, bool is_neox, int "
-      "rope_dim_offset=0, bool inverse=False) -> ()");
+      "                 Tensor cos_sin_cache, bool is_neox) -> ()");
   ops.impl("rotary_embedding", torch::kCUDA, &rotary_embedding);
 
   // Quantization ops

docker/Dockerfile

@@ -538,7 +538,7 @@ RUN CUDA_VERSION_DASH=$(echo $CUDA_VERSION | cut -d. -f1,2 | tr '.' '-') && \
         cuda-nvrtc-${CUDA_VERSION_DASH} \
         cuda-cuobjdump-${CUDA_VERSION_DASH} \
         libcurand-dev-${CUDA_VERSION_DASH} \
-        libcublas-dev-${CUDA_VERSION_DASH} \
+        libcublas-${CUDA_VERSION_DASH} \
         # Required by fastsafetensors (fixes #20384)
         libnuma-dev && \
     # Fixes nccl_allocator requiring nccl.h at runtime

docs/design/attention_backends.md

@@ -213,7 +213,7 @@ configuration.
 | `FLASHINFER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | 10.x |
 | `FLASHINFER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | 10.x |
 | `FLASHMLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 64 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x-10.x |
-| `FLASHMLA_SPARSE` | bf16 | `auto`, `bfloat16`, `fp8_ds_mla` | 64 | 512, 576 | ❌ | ✅ | ❌ | ❌ | 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`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x |
 | `ROCM_AITER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %1 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
 | `ROCM_AITER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16` | 1 | Any | ❌ | ✅ | ❌ | ❌ | Decoder | N/A |

docs/features/multimodal_inputs.md

@@ -68,7 +68,7 @@ You can pass a single image to the `'image'` field of the multi-modal dictionary
         print(generated_text)
     ```
 
-Full example: [examples/generate/multimodal/vision_language_offline.py](../../examples/generate/multimodal/vision_language_offline.py)
+Full example: [examples/offline_inference/vision_language.py](../../examples/offline_inference/vision_language.py)
 
 To substitute multiple images inside the same text prompt, you can pass in a list of images instead:
 
@@ -101,7 +101,7 @@ To substitute multiple images inside the same text prompt, you can pass in a lis
         print(generated_text)
     ```
 
-Full example: [examples/generate/multimodal/vision_language_multi_image_offline.py](../../examples/generate/multimodal/vision_language_multi_image_offline.py)
+Full example: [examples/offline_inference/vision_language_multi_image.py](../../examples/offline_inference/vision_language_multi_image.py)
 
 If using the [LLM.chat](../models/generative_models.md#llmchat) method, you can pass images directly in the message content using various formats: image URLs, PIL Image objects, or pre-computed embeddings:
 
@@ -287,13 +287,13 @@ Instead of NumPy arrays, you can also pass `'torch.Tensor'` instances, as shown
     !!! note
         'process_vision_info' is only applicable to Qwen2.5-VL and similar models.
 
-Full example: [examples/generate/multimodal/vision_language_offline.py](../../examples/generate/multimodal/vision_language_offline.py)
+Full example: [examples/offline_inference/vision_language.py](../../examples/offline_inference/vision_language.py)
 
 ### Audio Inputs
 
 You can pass a tuple `(array, sampling_rate)` to the `'audio'` field of the multi-modal dictionary.
 
-Full example: [examples/generate/multimodal/audio_language_offline.py](../../examples/generate/multimodal/audio_language_offline.py)
+Full example: [examples/offline_inference/audio_language.py](../../examples/offline_inference/audio_language.py)
 
 #### Chunking Long Audio for Transcription
 
@@ -674,7 +674,7 @@ Then, you can use the OpenAI client as follows:
     print("Chat completion output:", chat_response.choices[0].message.content)
     ```
 
-Full example: [examples/generate/multimodal/openai_chat_completion_client_for_multimodal.py](../../examples/generate/multimodal/openai_chat_completion_client_for_multimodal.py)
+Full example: [examples/online_serving/openai_chat_completion_client_for_multimodal.py](../../examples/online_serving/openai_chat_completion_client_for_multimodal.py)
 
 !!! tip
     Loading from local file paths is also supported on vLLM: You can specify the allowed local media path via `--allowed-local-media-path` when launching the API server/engine,
@@ -745,7 +745,7 @@ Then, you can use the OpenAI client as follows:
     print("Chat completion output from image url:", result)
     ```
 
-Full example: [examples/generate/multimodal/openai_chat_completion_client_for_multimodal.py](../../examples/generate/multimodal/openai_chat_completion_client_for_multimodal.py)
+Full example: [examples/online_serving/openai_chat_completion_client_for_multimodal.py](../../examples/online_serving/openai_chat_completion_client_for_multimodal.py)
 
 !!! note
     By default, the timeout for fetching videos through HTTP URL is `30` seconds.
@@ -958,7 +958,7 @@ Alternatively, you can pass `audio_url`, which is the audio counterpart of `imag
     print("Chat completion output from audio url:", result)
     ```
 
-Full example: [examples/generate/multimodal/openai_chat_completion_client_for_multimodal.py](../../examples/generate/multimodal/openai_chat_completion_client_for_multimodal.py)
+Full example: [examples/online_serving/openai_chat_completion_client_for_multimodal.py](../../examples/online_serving/openai_chat_completion_client_for_multimodal.py)
 
 !!! note
     By default, the timeout for fetching audios through HTTP URL is `10` seconds.

docs/features/quantization/README.md

@@ -20,7 +20,6 @@ The following are the supported quantization formats for vLLM:
 - [AMD Quark](quark.md)
 - [Quantized KV Cache](quantized_kvcache.md)
 - [TorchAO](torchao.md)
-- [FP8 ViT Encoder Attention](fp8_vit_attn.md)
 
 ## Supported Hardware
 

docs/features/quantization/fp8_vit_attn.md

@@ -1,109 +0,0 @@
-# FP8 ViT Encoder Attention
-
-For visual understanding workloads with large images (e.g. QHD, 4K) and relatively
-short text prompts/generation, the ViT encoder attention can become a significant
-bottleneck, especially when the text model is quantized (e.g. NVFP4). vLLM
-supports optional FP8 quantization for the ViT encoder attention via the
-FlashInfer cuDNN backend. Q/K/V are quantized on-the-fly to FP8 before the
-cuDNN attention call.
-
-!!! note
-    - Currently supports Qwen3-VL family models only (`qwen3_vl`, `qwen3_vl_moe`,
-      `qwen3_5`, `qwen3_5_moe`, and other models using Qwen3 ViT).
-    - Dynamic scaling is not compatible with ViT full CUDA graphs.
-    - Performance gains are mostly visible at QHD/4K resolutions or multi-image
-      requests. Smaller images may see no speedup due to quantization overhead
-      (3 quantization kernel launches + un-padding).
-    - FP8 tensor-core speedup is more pronounced on GB300 than GB200.
-
-## Requirements
-
-- FlashInfer cuDNN backend with cuDNN >= 9.17.1.
-
-## Usage
-
-Enable FP8 ViT attention by passing `--mm-encoder-attn-dtype fp8` together
-with `--mm-encoder-attn-backend FLASHINFER`:
-
-```bash
-vllm serve $MODEL \
-    --mm-encoder-attn-backend FLASHINFER \
-    --mm-encoder-attn-dtype fp8
-```
-
-By default (no scale file), **dynamic scaling** is used: a 16-entry circular
-buffer of observed Q/K/V amax values drives per-forward scale updates. This
-matches BF16 accuracy without any calibration but adds a small per-forward
-overhead.
-
-## Calibrate-Once, Reuse Workflow (Recommended)
-
-For production, calibrate static scales on a representative dataset once and
-reuse them to avoid the dynamic overhead:
-
-```bash
-# Step 1: calibrate and save scales (runs dynamic scaling for 16 passes,
-# then dumps the learned scales to JSON).
-vllm bench mm-processor \
-    --model $MODEL --mm-encoder-attn-backend FLASHINFER \
-    --mm-encoder-attn-dtype fp8 \
-    --mm-encoder-fp8-scale-save-path /path/to/scales.json \
-    --dataset-name hf --dataset-path lmarena-ai/VisionArena-Chat \
-    --num-prompts 100
-
-# Step 2: serve with static scales (no dynamic overhead).
-vllm serve $MODEL \
-    --mm-encoder-attn-backend FLASHINFER \
-    --mm-encoder-attn-dtype fp8 \
-    --mm-encoder-fp8-scale-path /path/to/scales.json
-```
-
-Saved scales are multiplied by `--mm-encoder-fp8-scale-save-margin` (default
-`1.5`) to leave headroom against activation outliers not present in the
-calibration set. The default has been validated to generalize across datasets
-(e.g. VisionArena-Chat calibration maintains BF16 accuracy on ChartQA).
-
-## Scale File Format
-
-```json
-{
-    "visual.blocks.0.attn.attn": {"q": 224.0, "k": 198.0, "v": 210.0},
-    "visual.blocks.1.attn.attn": {"q": 218.0, "k": 195.0, "v": 207.0}
-}
-```
-
-Keys `q_scale` / `k_scale` / `v_scale` are accepted as aliases.
-
-## Performance
-
-**Core cuDNN attention kernel** (PyTorch profiler, `cudnn_generated_fort_native_sdpa_sm100_flash_fprop`, head_dim=128, seq_len=8192):
-
-| Hardware | BF16 | FP8 | Speedup |
-| -------- | ---- | ---- | ------- |
-| GB200 | 350 us | 312 us | **1.12x** |
-| GB300 | 300 us | 211 us | **1.42x** |
-
-**End-to-end encoder forward time** (Qwen3-VL-30B-A3B-Instruct on GB200, 3 images/request):
-
-| Resolution | BF16 median | FP8 median | Speedup |
-| ---------- | ----------- | ---------- | ------- |
-| HD (720x1280) | 31.77 ms | 36.39 ms | 0.87x |
-| FullHD (1080x1920) | 57.99 ms | 58.73 ms | ~same |
-| QHD (1440x2560) | 131.83 ms | 122.30 ms | **1.08x** |
-| 4K (2160x3840) | 543.44 ms | 460.31 ms | **1.18x** |
-
-Crossover is around FullHD with 3 images/request. At QHD and above, FP8 wins.
-
-## Accuracy
-
-ChartQA, Qwen3-VL-8B-Instruct, 500 samples. FP8 static uses scales calibrated
-on VisionArena-Chat (with default 1.5x margin):
-
-| Metric | BF16 | FP8 dynamic | FP8 static |
-| ------ | ---- | ----------- | ---------- |
-| relaxed_accuracy | 0.780 | 0.776 | 0.780 |
-| anywhere_accuracy | 0.806 | 0.816 | 0.814 |
-| exact_match | 0.584 | 0.582 | 0.578 |
-
-All three configurations match within statistical noise, confirming that
-static scales calibrated on one dataset generalize to another.

docs/features/reasoning_outputs.md

@@ -202,7 +202,7 @@ The reasoning content is also available when both tool calling and the reasoning
     print(f"Arguments: {tool_call.arguments}")
     ```
 
-For more examples, please refer to [examples/reasoning/openai_chat_completion_tool_calls_with_reasoning.py](../../examples/reasoning/openai_chat_completion_tool_calls_with_reasoning.py).
+For more examples, please refer to [examples/online_serving/openai_chat_completion_tool_calls_with_reasoning.py](../../examples/online_serving/openai_chat_completion_tool_calls_with_reasoning.py).
 
 ## Server-Level Default Chat Template Kwargs
 

docs/models/supported_models.md

@@ -384,7 +384,6 @@ th {
 | `DeepseekForCausalLM` | DeepSeek | `deepseek-ai/deepseek-llm-67b-base`, `deepseek-ai/deepseek-llm-7b-chat`, etc. | ✅︎ | ✅︎ |
 | `DeepseekV2ForCausalLM` | DeepSeek-V2 | `deepseek-ai/DeepSeek-V2`, `deepseek-ai/DeepSeek-V2-Chat`, etc. | ✅︎ | ✅︎ |
 | `DeepseekV3ForCausalLM` | DeepSeek-V3 | `deepseek-ai/DeepSeek-V3`, `deepseek-ai/DeepSeek-R1`, `deepseek-ai/DeepSeek-V3.1`, etc. | ✅︎ | ✅︎ |
-| `DeepseekV4ForCausalLM` | DeepSeek-V4 | `deepseek-ai/DeepSeek-V4-Flash`, `deepseek-ai/DeepSeek-V4-Pro`, etc. | | |
 | `Dots1ForCausalLM` | dots.llm1 | `rednote-hilab/dots.llm1.base`, `rednote-hilab/dots.llm1.inst`, etc. | | ✅︎ |
 | `DotsOCRForCausalLM` | dots_ocr | `rednote-hilab/dots.ocr` | ✅︎ | ✅︎ |
 | `Ernie4_5ForCausalLM` | Ernie4.5 | `baidu/ERNIE-4.5-0.3B-PT`, etc. | ✅︎ | ✅︎ |
@@ -439,7 +438,6 @@ th {
 | `Mamba2ForCausalLM` | Mamba2 | `mistralai/Mamba-Codestral-7B-v0.1`, etc. | | ✅︎ |
 | `MiMoForCausalLM` | MiMo | `XiaomiMiMo/MiMo-7B-RL`, etc. | ✅︎ | ✅︎ |
 | `MiMoV2FlashForCausalLM` | MiMoV2Flash | `XiaomiMiMo/MiMo-V2-Flash`, etc. | | ✅︎ |
-| `MiMoV2ProForCausalLM` | MiMoV2Pro | `XiaomiMiMo/MiMo-V2.5-Pro`, etc. | | ✅︎ |
 | `MiniCPMForCausalLM` | MiniCPM | `openbmb/MiniCPM-2B-sft-bf16`, `openbmb/MiniCPM-2B-dpo-bf16`, `openbmb/MiniCPM-S-1B-sft`, etc. | ✅︎ | ✅︎ |
 | `MiniCPM3ForCausalLM` | MiniCPM3 | `openbmb/MiniCPM3-4B`, etc. | ✅︎ | ✅︎ |
 | `MiniMaxForCausalLM` | MiniMax-Text | `MiniMaxAI/MiniMax-Text-01-hf`, etc. | | |
@@ -591,7 +589,6 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
 | `LlavaNextVideoForConditionalGeneration` | LLaVA-NeXT-Video | T + V | `llava-hf/LLaVA-NeXT-Video-7B-hf`, etc. | | ✅︎ |
 | `LlavaOnevisionForConditionalGeneration` | LLaVA-Onevision | T + I<sup>+</sup> + V<sup>+</sup> | `llava-hf/llava-onevision-qwen2-7b-ov-hf`, `llava-hf/llava-onevision-qwen2-0.5b-ov-hf`, etc. | | ✅︎ |
 | `MiDashengLMModel` | MiDashengLM | T + A<sup>+</sup> | `mispeech/midashenglm-7b` | | ✅︎ |
-| `MiMoV2OmniForCausalLM` | MiMo-V2.5-Omni | T + I<sup>E+</sup> + V<sup>E+</sup> + A<sup>+</sup> | `XiaomiMiMo/MiMo-V2.5-Omni` | | ✅︎ |
 | `MiniCPMO` | MiniCPM-O | T + I<sup>E+</sup> + V<sup>E+</sup> + A<sup>E+</sup> | `openbmb/MiniCPM-o-2_6`, etc. | ✅︎ | ✅︎ |
 | `MiniCPMV` | MiniCPM-V | T + I<sup>E+</sup> + V<sup>E+</sup> | `openbmb/MiniCPM-V-2` (see note), `openbmb/MiniCPM-Llama3-V-2_5`, `openbmb/MiniCPM-V-2_6`, `openbmb/MiniCPM-V-4`, `openbmb/MiniCPM-V-4_5`, etc. | ✅︎ | |
 | `MiniMaxVL01ForConditionalGeneration` | MiniMax-VL | T + I<sup>E+</sup> | `MiniMaxAI/MiniMax-VL-01`, etc. | | ✅︎ |
@@ -646,10 +643,10 @@ Some models are supported only via the [Transformers modeling backend](#transfor
 !!! note
     `Gemma3nForConditionalGeneration` is only supported on V1 due to shared KV caching and it depends on `timm>=1.0.17` to make use of its
     MobileNet-v5 vision backbone.
-
+  
     Performance is not yet fully optimized mainly due to:
-
-    - Both audio and vision MM encoders use `transformers.AutoModel` implementation.
+  
+    - Both audio and vision MM encoders use `transformers.AutoModel` implementation.  
     - There's no PLE caching or out-of-memory swapping support, as described in [Google's blog](https://developers.googleblog.com/en/introducing-gemma-3n/). These features might be too model-specific for vLLM, and swapping in particular may be better suited for constrained setups.
 
 !!! note

docs/serving/openai_compatible_server.md

@@ -251,7 +251,7 @@ The following extra parameters are supported:
 Our Responses API is compatible with [OpenAI's Responses API](https://platform.openai.com/docs/api-reference/responses);
 you can use the [official OpenAI Python client](https://github.com/openai/openai-python) to interact with it.
 
-Code example: [examples/online_serving/openai_responses_client_with_tools.py](../../examples/tool_calling/openai_responses_client_with_tools.py)
+Code example: [examples/online_serving/openai_responses_client_with_tools.py](../../examples/online_serving/openai_responses_client_with_tools.py)
 
 #### Extra parameters
 
@@ -279,7 +279,7 @@ you can use the [official OpenAI Python client](https://github.com/openai/openai
 !!! note
     To use the Transcriptions API, please install with extra audio dependencies using `pip install vllm[audio]`.
 
-Code example: [examples/speech_to_text/openai/openai_transcription_client.py](../../examples/speech_to_text/openai/openai_transcription_client.py)
+Code example: [examples/online_serving/openai_transcription_client.py](../../examples/online_serving/openai_transcription_client.py)
 
 NOTE: beam search is currently supported in the transcriptions endpoint for encoder-decoder multimodal models, e.g., whisper, but highly inefficient as work for handling the encoder/decoder cache is actively ongoing. This is an active point of ongoing optimization and will be handled properly in the very near future.
 
@@ -397,7 +397,7 @@ Please mind that the popular `openai/whisper-large-v3-turbo` model does not supp
 !!! note
     To use the Translation API, please install with extra audio dependencies using `pip install vllm[audio]`.
 
-Code example: [examples/speech_to_text/openai/openai_translation_client.py](../../examples/speech_to_text/openai/openai_translation_client.py)
+Code example: [examples/online_serving/openai_translation_client.py](../../examples/online_serving/openai_translation_client.py)
 
 #### Extra Parameters
 

examples/offline_inference/qwen2_5_omni/README.md

@@ -6,15 +6,15 @@ This folder provides several example scripts on how to inference Qwen2.5-Omni of
 
 ```bash
 # Audio + image + video
-python examples/generate/multimodal/qwen2_5_omni/only_thinker.py \
+python examples/offline_inference/qwen2_5_omni/only_thinker.py \
     -q mixed_modalities
 
 # Read vision and audio inputs from a single video file
-python examples/generate/multimodal/qwen2_5_omni/only_thinker.py \
+python examples/offline_inference/qwen2_5_omni/only_thinker.py \
     -q use_audio_in_video
 
 # Multiple audios
-python examples/generate/multimodal/qwen2_5_omni/only_thinker.py \
+python examples/offline_inference/qwen2_5_omni/only_thinker.py \
     -q multi_audios
 ```
 
@@ -24,16 +24,16 @@ You can also test Qwen2.5-Omni on a single modality:
 
 ```bash
 # Process audio inputs
-python examples/generate/multimodal/audio_language_offline.py \
+python examples/offline_inference/audio_language.py \
     --model-type qwen2_5_omni
 
 # Process image inputs
-python examples/generate/multimodal/vision_language_offline.py \
+python examples/offline_inference/vision_language.py \
     --modality image \
     --model-type qwen2_5_omni
 
 # Process video inputs
-python examples/generate/multimodal/vision_language_offline.py \
+python examples/offline_inference/vision_language.py \
     --modality video \
     --model-type qwen2_5_omni
 ```

examples/offline_inference/vision_language.py

@@ -1402,7 +1402,7 @@ def run_mantis(questions: list[str], modality: str) -> ModelRequestData:
 # MiniCPM-V
 def run_minicpmv_base(questions: list[str], modality: str, model_name):
     assert modality in ["image", "video", "image+video"]
-    # If you want to use `MiniCPM-o-2_6` with audio inputs, check `audio_language_offline.py` # noqa
+    # If you want to use `MiniCPM-o-2_6` with audio inputs, check `audio_language.py` # noqa
 
     # 2.0
     # The official repo doesn't work yet, so we need to use a fork for now

requirements/cuda.txt

@@ -11,8 +11,6 @@ torchvision==0.26.0 # Required for phi3v processor. See https://github.com/pytor
 # FlashInfer should be updated together with the Dockerfile
 flashinfer-python==0.6.8.post1
 flashinfer-cubin==0.6.8.post1
-apache-tvm-ffi==0.1.9
-tilelang==0.1.9
 # Cap nvidia-cudnn-frontend (transitive dep of flashinfer) due to
 # breaking changes in 1.19.0
 nvidia-cudnn-frontend>=1.13.0,<1.19.0

tests/compile/correctness_e2e/test_sequence_parallel.py

@@ -261,8 +261,6 @@ def _compare_sp(
         },
         "use_inductor_graph_partition": use_inductor_graph_partition,
     }
-    if not use_inductor_graph_partition:
-        compilation_config["splitting_ops"] = []
 
     tp_sp_args = [
         *common_args,

tests/compile/fusions_e2e/conftest.py

@@ -116,11 +116,6 @@ def run_e2e_fusion_test(monkeypatch, caplog_mp_spawn):
         model_kwargs["attention_config"] = {"backend": attn_backend.backend.name}
         model_kwargs["tensor_parallel_size"] = tp_size
 
-        # Cap warmup memory: tests use small max_model_len (1024) but the
-        # engine default max_num_batched_tokens is 16384. Warming up large
-        # models (e.g. Llama-4-Scout-FP8) at 16384 tokens may trigger OOM.
-        model_kwargs.setdefault("max_num_batched_tokens", 8192)
-
         # Sparse MLA models (DSv3.2) hit an over-strict inductor assertion in
         # decompose_auto_functionalized when +rotary_embedding is forced into
         # the compile graph. Disable qk_norm+rope fusion (which auto-enables

tests/compile/passes/distributed/test_async_tp.py

@@ -19,7 +19,6 @@ from vllm.config import (
     VllmConfig,
     set_current_vllm_config,
 )
-from vllm.config.utils import Range
 from vllm.distributed import (
     tensor_model_parallel_all_gather,
     tensor_model_parallel_reduce_scatter,
@@ -289,22 +288,6 @@ def test_async_tp_pass_replace(
     run_torch_spawn(async_tp_pass_on_test_model, num_processes)
 
 
-def test_async_tp_pass_requires_full_graph_compilation():
-    vllm_config = VllmConfig()
-    vllm_config.compilation_config.use_inductor_graph_partition = False
-    vllm_config.compilation_config.splitting_ops = [
-        "vllm::unified_attention_with_output"
-    ]
-
-    async_tp_pass = object.__new__(AsyncTPPass)
-    async_tp_pass.compilation_config = vllm_config.compilation_config
-
-    with pytest.raises(
-        AssertionError, match="AsyncTPPass requires full-graph compilation"
-    ):
-        async_tp_pass.is_applicable_for_range(Range(start=8, end=8))
-
-
 def async_tp_pass_on_test_model(
     local_rank: int,
     world_size: int,

tests/compile/passes/distributed/test_sequence_parallelism.py

@@ -22,7 +22,6 @@ from vllm.config import (
     get_current_vllm_config,
     set_current_vllm_config,
 )
-from vllm.config.utils import Range
 from vllm.distributed import tensor_model_parallel_all_reduce
 from vllm.distributed.parallel_state import (
     init_distributed_environment,
@@ -217,24 +216,6 @@ def test_sequence_parallelism_pass(
     run_torch_spawn(sequence_parallelism_pass_on_test_model, num_processes)
 
 
-def test_sequence_parallelism_pass_requires_full_graph_compilation():
-    vllm_config = VllmConfig()
-    vllm_config.compilation_config.use_inductor_graph_partition = False
-    vllm_config.compilation_config.splitting_ops = [
-        "vllm::unified_attention_with_output"
-    ]
-
-    sequence_parallelism_pass = object.__new__(SequenceParallelismPass)
-    sequence_parallelism_pass.compilation_config = vllm_config.compilation_config
-    sequence_parallelism_pass.min_token_num = 1
-
-    with pytest.raises(
-        AssertionError,
-        match="SequenceParallelismPass requires full-graph compilation",
-    ):
-        sequence_parallelism_pass.is_applicable_for_range(Range(start=8, end=8))
-
-
 def sequence_parallelism_pass_on_test_model(
     local_rank: int,
     world_size: int,

tests/compile/test_config.py

@@ -407,7 +407,7 @@ def test_should_split():
         (None, 257, 1, False, 2048, CUDAGraphMode.FULL_AND_PIECEWISE, 256),
         # max from list
         ([1, 2, 4, 15], None, 1, False, 2048, CUDAGraphMode.FULL_AND_PIECEWISE, 15),
-        # SP forces full-graph compilation, sizes are filtered by TP
+        # filtered out 15 due to SP
         ([1, 2, 4, 15], None, 2, True, 2048, CUDAGraphMode.FULL_AND_PIECEWISE, 4),
         # limited by the max_tokens
         ([1, 2, 4, 15], None, 1, False, 8, CUDAGraphMode.FULL_AND_PIECEWISE, 4),
@@ -465,123 +465,6 @@ def test_cudagraph_sizes_post_init(
         )
 
 
-@pytest.mark.skipif(
-    not current_platform.support_static_graph_mode(),
-    reason="Skip if not cudagraph mode supported",
-)
-@pytest.mark.parametrize(
-    (
-        "cudagraph_mode",
-        "use_inductor_graph_partition",
-        "expected_enable_sp",
-        "expected_cudagraph_mode",
-        "expected_piecewise_compile",
-        "expected_capture_sizes",
-        "expected_max_size",
-    ),
-    [
-        (CUDAGraphMode.PIECEWISE, False, True, CUDAGraphMode.FULL, False, [2, 4], 4),
-        (
-            CUDAGraphMode.FULL_DECODE_ONLY,
-            False,
-            True,
-            CUDAGraphMode.FULL_DECODE_ONLY,
-            False,
-            [2, 4],
-            4,
-        ),
-        (
-            CUDAGraphMode.FULL_AND_PIECEWISE,
-            False,
-            True,
-            CUDAGraphMode.FULL,
-            False,
-            [2, 4],
-            4,
-        ),
-        (
-            CUDAGraphMode.FULL_AND_PIECEWISE,
-            True,
-            True,
-            CUDAGraphMode.FULL_AND_PIECEWISE,
-            True,
-            [2, 4],
-            4,
-        ),
-    ],
-)
-def test_sequence_parallelism_requires_full_graph_compilation(
-    cudagraph_mode: CUDAGraphMode,
-    use_inductor_graph_partition: bool,
-    expected_enable_sp: bool,
-    expected_cudagraph_mode: CUDAGraphMode,
-    expected_piecewise_compile: bool,
-    expected_capture_sizes: list[int],
-    expected_max_size: int,
-):
-    with patch.object(current_platform, "device_count", return_value=2):
-        vllm_config = VllmConfig(
-            parallel_config=ParallelConfig(tensor_parallel_size=2),
-            scheduler_config=SchedulerConfig(
-                max_num_seqs=128,
-                max_num_batched_tokens=2048,
-                max_model_len=2048,
-                is_encoder_decoder=False,
-            ),
-        )
-        vllm_config.model_config = MagicMock(
-            dtype=torch.float16,
-            enforce_eager=False,
-            is_moe=False,
-            disable_cascade_attn=False,
-            get_hidden_size=MagicMock(return_value=4096),
-        )
-        vllm_config.compilation_config = CompilationConfig(
-            mode=CompilationMode.VLLM_COMPILE,
-            cudagraph_capture_sizes=[1, 2, 4, 15],
-            max_cudagraph_capture_size=None,
-            compile_sizes=["cudagraph_capture_sizes"],
-            use_inductor_graph_partition=use_inductor_graph_partition,
-            pass_config=PassConfig(
-                enable_sp=True,
-                fuse_gemm_comms=True,
-                fuse_norm_quant=True,
-                fuse_act_quant=True,
-                eliminate_noops=True,
-                sp_min_token_num=512,
-            ),
-            cudagraph_mode=cudagraph_mode,
-        )
-        vllm_config.compilation_config.set_splitting_ops_for_v1(
-            all2all_backend=vllm_config.parallel_config.all2all_backend,
-            data_parallel_size=1,
-        )
-        vllm_config._set_compile_ranges()
-        vllm_config._set_cudagraph_sizes()
-
-    assert (
-        vllm_config.compilation_config.use_inductor_graph_partition
-        == use_inductor_graph_partition
-    )
-    assert (
-        bool(vllm_config.compilation_config.splitting_ops) == expected_piecewise_compile
-    )
-    assert vllm_config.compilation_config.pass_config.enable_sp == expected_enable_sp
-    assert (
-        vllm_config.compilation_config.pass_config.fuse_gemm_comms == expected_enable_sp
-    )
-    assert vllm_config.compilation_config.cudagraph_mode == expected_cudagraph_mode
-    assert (
-        vllm_config.compilation_config.cudagraph_capture_sizes == expected_capture_sizes
-    )
-    assert (
-        vllm_config.compilation_config.max_cudagraph_capture_size == expected_max_size
-    )
-    assert (
-        511 in vllm_config.compilation_config.compile_ranges_endpoints
-    ) == expected_enable_sp
-
-
 def test_cached_compilation_config(default_vllm_config):
     import torch
     from torch._inductor.utils import run_and_get_code

tests/config/test_multimodal_config.py

@@ -41,21 +41,3 @@ def test_language_model_only_affects_model_hash():
     base_hash = ModelConfig(model).compute_hash()
     lm_only_hash = ModelConfig(model, language_model_only=True).compute_hash()
     assert base_hash != lm_only_hash
-
-
-def test_mm_encoder_fp8_scale_path_requires_fp8():
-    with pytest.raises(ValueError, match="mm_encoder_attn_dtype"):
-        MultiModalConfig(mm_encoder_fp8_scale_path="/tmp/scales.json")
-
-
-def test_mm_encoder_attn_dtype_hash_updates(tmp_path):
-    scale_file = tmp_path / "scales.json"
-    scale_file.write_text("{}")
-    base_hash = MultiModalConfig().compute_hash()
-    fp8_hash = MultiModalConfig(mm_encoder_attn_dtype="fp8").compute_hash()
-    fp8_static_hash = MultiModalConfig(
-        mm_encoder_attn_dtype="fp8",
-        mm_encoder_fp8_scale_path=str(scale_file),
-    ).compute_hash()
-    assert base_hash != fp8_hash
-    assert fp8_hash != fp8_static_hash

tests/entrypoints/openai/test_fingerprint.py

@@ -1,76 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Unit tests for ``system_fingerprint`` construction."""
-
-from types import SimpleNamespace
-
-import pytest
-
-from vllm.entrypoints.openai import fingerprint as fp
-
-
-def _cfg(tp=1, pp=1, dp=1, ep=False, digest="a3b21f94deadbeef"):
-    c = SimpleNamespace(
-        parallel_config=SimpleNamespace(
-            tensor_parallel_size=tp,
-            pipeline_parallel_size=pp,
-            data_parallel_size=dp,
-            enable_expert_parallel=ep,
-        )
-    )
-    c.compute_hash = lambda: digest  # type: ignore[attr-defined]
-    return c
-
-
-@pytest.fixture(autouse=True)
-def _reset():
-    fp.set_default_fingerprint_mode("full")
-    yield
-    fp.set_default_fingerprint_mode("full")
-
-
-def test_four_modes_produce_expected_shapes():
-    from vllm import __version__ as v
-
-    cfg = _cfg(tp=8, ep=True)
-
-    assert fp.build_system_fingerprint(cfg, "full") == (f"vllm-{v}-tp8-ep-a3b21f94")
-    assert fp.build_system_fingerprint(cfg, "hash") == f"vllm-{v}-a3b21f94"
-    assert fp.build_system_fingerprint(cfg, "custom", "my-fp") == "my-fp"
-    assert fp.build_system_fingerprint(cfg, "none") is None
-
-
-def test_full_mode_emits_only_non_trivial_parallelism():
-    from vllm import __version__ as v
-
-    # Single-GPU: nothing between version and hash.
-    assert fp.build_system_fingerprint(_cfg(), "full") == f"vllm-{v}-a3b21f94"
-    # All parallelism axes.
-    assert (
-        fp.build_system_fingerprint(_cfg(tp=8, pp=2, dp=4, ep=True), "full")
-        == f"vllm-{v}-tp8-pp2-dp4-ep-a3b21f94"
-    )
-
-
-def test_get_respects_set_default():
-    cfg = _cfg(tp=8)
-    full = fp.get_system_fingerprint(cfg)
-    assert full == fp.get_system_fingerprint(cfg)
-
-    fp.set_default_fingerprint_mode("hash")
-    hashed = fp.get_system_fingerprint(cfg)
-    assert hashed != full
-    assert "tp8" not in hashed
-
-    fp.set_default_fingerprint_mode("custom", "deploy-42")
-    assert fp.get_system_fingerprint(cfg) == "deploy-42"
-
-    fp.set_default_fingerprint_mode("none")
-    assert fp.get_system_fingerprint(cfg) is None
-
-
-def test_compute_hash_failure_does_not_raise():
-    cfg = _cfg()
-    cfg.compute_hash = lambda: (_ for _ in ()).throw(RuntimeError("boom"))
-    assert fp.build_system_fingerprint(cfg, "full").endswith("-nohash")
-    assert fp.build_system_fingerprint(cfg, "hash").endswith("-nohash")

tests/kernels/attention/test_deepgemm_attention.py

@@ -9,8 +9,8 @@ from vllm.platforms import current_platform
 from vllm.utils.deep_gemm import (
     _ceil_to_ue8m0,
     calc_diff,
-    fp8_fp4_mqa_logits,
-    fp8_fp4_paged_mqa_logits,
+    fp8_mqa_logits,
+    fp8_paged_mqa_logits,
     get_num_sms,
     get_paged_mqa_logits_metadata,
 )
@@ -127,8 +127,8 @@ def test_deepgemm_fp8_mqa_logits(clean_logits: bool):
 
                 q_fp8 = q.to(torch.float8_e4m3fn)
                 kv_fp8 = per_custom_dims_cast_to_fp8(kv, (0,), False)
-                logits = fp8_fp4_mqa_logits(
-                    (q_fp8, None), kv_fp8, weights, ks, ke, clean_logits=clean_logits
+                logits = fp8_mqa_logits(
+                    q_fp8, kv_fp8, weights, ks, ke, clean_logits=clean_logits
                 )
 
                 ref_logits = _ref_fp8_mqa_logits(
@@ -150,7 +150,7 @@ def test_deepgemm_fp8_mqa_logits(clean_logits: bool):
                 assert diff < 1e-3, f"{diff=}"
 
 
-def _ref_fp8_fp4_paged_mqa_logits(
+def _ref_fp8_paged_mqa_logits(
     q: torch.Tensor,
     kv_cache: torch.Tensor,
     weights: torch.Tensor,
@@ -205,10 +205,8 @@ def _ref_fp8_fp4_paged_mqa_logits(
 @pytest.mark.skipif(
     not current_platform.has_device_capability(90), reason="SM90 and SM100 only"
 )
-def test_deepgemm_fp8_fp4_paged_mqa_logits():
-    # NOTE: clean_logits=True is incompatible with the 2D context_lens
-    # required by csrc/apis/attention.hpp; only the False path is exercised.
-    clean_logits = False
+@pytest.mark.parametrize("clean_logits", [True, False])
+def test_deepgemm_fp8_paged_mqa_logits(clean_logits: bool):
     torch.manual_seed(0)
     random.seed(0)
 
@@ -260,29 +258,21 @@ def test_deepgemm_fp8_fp4_paged_mqa_logits():
                 q_fp8 = q.to(torch.float8_e4m3fn)
                 kv_cache_fp8 = kv_cache_cast_to_fp8(kv_cache)
 
-                # deep_gemm paged MQA logits requires 2D context_lens of
-                # shape (B, next_n) (csrc/apis/attention.hpp:332-335);
-                # see indexer.py:607-608. For each batch/next_n token, the
-                # effective context length is context_lens[b] - next_n + j + 1.
-                next_n_arange = torch.arange(next_n, device="cuda", dtype=torch.int32)
-                context_lens_2d = (
-                    context_lens.unsqueeze(-1) - next_n + 1 + next_n_arange
-                ).contiguous()
                 schedule_metadata = get_paged_mqa_logits_metadata(
-                    context_lens_2d, blocksize, get_num_sms()
+                    context_lens, blocksize, get_num_sms()
                 )
-                logits = fp8_fp4_paged_mqa_logits(
-                    (q_fp8, None),
+                logits = fp8_paged_mqa_logits(
+                    q_fp8,
                     kv_cache_fp8,
                     weights,
-                    context_lens_2d,
+                    context_lens,
                     block_tables,
                     schedule_metadata,
                     max_model_len,
                     clean_logits=clean_logits,
                 )
 
-                ref_logits = _ref_fp8_fp4_paged_mqa_logits(
+                ref_logits = _ref_fp8_paged_mqa_logits(
                     q,
                     kv_cache,
                     weights,

tests/kernels/core/test_activation.py

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

tests/kernels/core/test_fused_q_kv_rmsnorm.py

@@ -1,81 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Correctness + large-token-count launch tests for fused_q_kv_rmsnorm.
-
-Before the grid-dim fix the kernel used grid ``(2, num_tokens)``, which hit
-CUDA's 65535 grid-y cap for ``num_tokens >= 65536`` and failed with
-``Triton Error [CUDA]: invalid argument`` at every large chunked-prefill
-profile run. These tests pin the new grid layout.
-"""
-
-from __future__ import annotations
-
-import pytest
-import torch
-
-from vllm.platforms import current_platform
-from vllm.v1.attention.ops.deepseek_v4_ops import fused_q_kv_rmsnorm
-
-pytestmark = pytest.mark.skipif(
-    not current_platform.is_cuda_alike(),
-    reason="fused_q_kv_rmsnorm requires a CUDA/ROCm device",
-)
-
-
-def _ref_rmsnorm(x: torch.Tensor, w: torch.Tensor, eps: float) -> torch.Tensor:
-    x_f32 = x.to(torch.float32)
-    variance = x_f32.pow(2).mean(dim=-1, keepdim=True)
-    y = x_f32 * torch.rsqrt(variance + eps) * w.to(torch.float32)
-    return y.to(x.dtype)
-
-
-@pytest.mark.parametrize("num_tokens", [1, 17, 1024, 8192])
-@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
-def test_fused_q_kv_rmsnorm_correctness(num_tokens: int, dtype: torch.dtype):
-    torch.manual_seed(0)
-    device = "cuda"
-    q_size, kv_size = 192, 576
-    qr = torch.randn(num_tokens, q_size, dtype=dtype, device=device)
-    kv = torch.randn(num_tokens, kv_size, dtype=dtype, device=device)
-    qw = torch.randn(q_size, dtype=dtype, device=device)
-    kvw = torch.randn(kv_size, dtype=dtype, device=device)
-    eps = 1e-6
-
-    qr_out, kv_out = fused_q_kv_rmsnorm(qr, kv, qw, kvw, eps)
-
-    qr_ref = _ref_rmsnorm(qr, qw, eps)
-    kv_ref = _ref_rmsnorm(kv, kvw, eps)
-
-    tol = dict(rtol=1e-2, atol=1e-2)
-    torch.testing.assert_close(qr_out, qr_ref, **tol)
-    torch.testing.assert_close(kv_out, kv_ref, **tol)
-
-
-@pytest.mark.parametrize("num_tokens", [65535, 65536, 131072])
-def test_fused_q_kv_rmsnorm_launches_past_grid_y_cap(num_tokens: int):
-    """Regression guard: grid used to be (2, num_tokens), hitting CUDA's
-    65535 grid-y cap at num_tokens >= 65536. The new grid (num_tokens, 2)
-    lifts that bound to 2**31-1."""
-    device = "cuda"
-    dtype = torch.bfloat16
-    q_size, kv_size = 192, 576
-    qr = torch.randn(num_tokens, q_size, dtype=dtype, device=device)
-    kv = torch.randn(num_tokens, kv_size, dtype=dtype, device=device)
-    qw = torch.randn(q_size, dtype=dtype, device=device)
-    kvw = torch.randn(kv_size, dtype=dtype, device=device)
-
-    qr_out, kv_out = fused_q_kv_rmsnorm(qr, kv, qw, kvw, 1e-6)
-    # spot-check a couple of rows against the torch reference
-    for row in (0, num_tokens // 2, num_tokens - 1):
-        torch.testing.assert_close(
-            qr_out[row],
-            _ref_rmsnorm(qr[row : row + 1], qw, 1e-6)[0],
-            rtol=1e-2,
-            atol=1e-2,
-        )
-        torch.testing.assert_close(
-            kv_out[row],
-            _ref_rmsnorm(kv[row : row + 1], kvw, 1e-6)[0],
-            rtol=1e-2,
-            atol=1e-2,
-        )

tests/kernels/core/test_vit_fp8_attn.py

@@ -1,279 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Tests for the full FP8 ViT attention path (quantize -> cuDNN -> un-pad)."""
-
-import contextlib
-
-import pytest
-import torch
-
-from vllm.triton_utils import HAS_TRITON
-from vllm.utils.flashinfer import (
-    is_flashinfer_cudnn_fp8_prefill_attn_supported,
-)
-from vllm.v1.attention.backends.registry import AttentionBackendEnum
-
-
-def _has_flashinfer_cudnn() -> bool:
-    """Check if FlashInfer cuDNN backend is available."""
-    try:
-        from flashinfer.prefill import (
-            cudnn_batch_prefill_with_kv_cache,  # noqa: F401
-        )
-
-        return True
-    except ImportError:
-        return False
-
-
-HEAD_DIMS = [72, 80]
-SEQ_LENS = [256]
-NUM_HEADS = [16]
-
-
-@pytest.fixture
-def _fp8_attention():
-    """Create FP8-enabled MMEncoderAttention via config."""
-    from types import SimpleNamespace
-    from unittest.mock import patch
-
-    from vllm.config import VllmConfig, set_current_vllm_config
-    from vllm.config.multimodal import MultiModalConfig
-
-    if not is_flashinfer_cudnn_fp8_prefill_attn_supported():
-        pytest.skip("FlashInfer cuDNN FP8 prefill attention not supported")
-
-    mm_config = MultiModalConfig(mm_encoder_attn_dtype="fp8")
-    vllm_config = VllmConfig()
-    vllm_config.model_config = SimpleNamespace(multimodal_config=mm_config)
-
-    # MMEncoderAttention reads torch.get_default_dtype() during init
-    # to determine the output dtype. In real model loading this is bf16.
-    old_dtype = torch.get_default_dtype()
-    torch.set_default_dtype(torch.bfloat16)
-
-    with (
-        set_current_vllm_config(vllm_config),
-        patch(
-            "vllm.model_executor.layers.attention.mm_encoder_attention"
-            ".get_vit_attn_backend",
-            return_value=AttentionBackendEnum.FLASHINFER,
-        ),
-    ):
-        yield
-
-    torch.set_default_dtype(old_dtype)
-
-
-def _build_cu_seqlens_and_meta(
-    seq_len: int,
-    num_heads: int,
-    head_dim: int,
-    fp8_padded_hidden_size: int | None = None,
-) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-    """Build cu_seqlens, max_seqlen, sequence_lengths for a single sequence."""
-    import numpy as np
-
-    from vllm.model_executor.layers.attention.mm_encoder_attention import (
-        MMEncoderAttention,
-    )
-
-    cu_seqlens_np = np.array([0, seq_len], dtype=np.int32)
-
-    sequence_lengths = MMEncoderAttention.maybe_compute_seq_lens(
-        AttentionBackendEnum.FLASHINFER,
-        cu_seqlens_np,
-        torch.device("cuda"),
-    )
-
-    max_seqlen = torch.tensor(
-        MMEncoderAttention.compute_max_seqlen(
-            AttentionBackendEnum.FLASHINFER, cu_seqlens_np
-        ),
-        dtype=torch.int32,
-    )
-
-    cu_seqlens = MMEncoderAttention.maybe_recompute_cu_seqlens(
-        AttentionBackendEnum.FLASHINFER,
-        cu_seqlens_np,
-        num_heads * head_dim,
-        1,  # tp_size
-        torch.device("cuda"),
-        fp8_padded_hidden_size=fp8_padded_hidden_size,
-    )
-
-    return cu_seqlens, max_seqlen, sequence_lengths
-
-
-@pytest.mark.skipif(
-    not (HAS_TRITON and _has_flashinfer_cudnn()),
-    reason="Triton and FlashInfer cuDNN required",
-)
-@pytest.mark.parametrize("head_dim", HEAD_DIMS)
-@pytest.mark.parametrize("seq_len", SEQ_LENS)
-@pytest.mark.parametrize("num_heads", NUM_HEADS)
-def test_fp8_attn_output_shape(
-    head_dim: int,
-    seq_len: int,
-    num_heads: int,
-    _fp8_attention,
-) -> None:
-    """Verify FP8 attention produces correct output shape after un-padding."""
-    from vllm.model_executor.layers.attention.mm_encoder_attention import (
-        MMEncoderAttention,
-    )
-    from vllm.utils.math_utils import round_up
-
-    attn = None
-    with contextlib.suppress(ValueError, ImportError):
-        attn = MMEncoderAttention(
-            num_heads=num_heads,
-            head_size=head_dim,
-            prefix="visual.blocks.0.attn",
-        ).to("cuda")
-
-    if attn is None or not attn.fp8_enabled:
-        pytest.skip("FP8 MMEncoderAttention not available")
-    assert attn is not None  # mypy narrowing
-
-    # FP8 always needs fp8_padded_hidden_size for correct cu_seqlens
-    fp8_padded_hidden_size = num_heads * round_up(head_dim, 16)
-
-    cu_seqlens, max_seqlen, sequence_lengths = _build_cu_seqlens_and_meta(
-        seq_len, num_heads, head_dim, fp8_padded_hidden_size=fp8_padded_hidden_size
-    )
-
-    q = torch.randn(
-        seq_len,
-        num_heads,
-        head_dim,
-        device="cuda",
-        dtype=torch.bfloat16,
-    )
-    k = torch.randn_like(q)
-    v = torch.randn_like(q)
-
-    output = attn._forward_flashinfer(q, k, v, cu_seqlens, max_seqlen, sequence_lengths)
-
-    # Output should have original head_dim (un-padded)
-    assert output.shape[-1] == head_dim
-    assert output.dtype == torch.bfloat16
-
-
-@pytest.mark.skipif(
-    not (HAS_TRITON and _has_flashinfer_cudnn()),
-    reason="Triton and FlashInfer cuDNN required",
-)
-@pytest.mark.parametrize("head_dim", HEAD_DIMS)
-@pytest.mark.parametrize("seq_len", SEQ_LENS)
-@pytest.mark.parametrize("num_heads", NUM_HEADS)
-def test_fp8_vs_bf16_close(
-    head_dim: int, seq_len: int, num_heads: int, _fp8_attention
-) -> None:
-    """FP8 attention output should be reasonably close to BF16 baseline."""
-    from vllm.model_executor.layers.attention.mm_encoder_attention import (
-        MMEncoderAttention,
-    )
-    from vllm.utils.math_utils import round_up
-
-    torch.manual_seed(42)
-    q = torch.randn(
-        1,
-        seq_len,
-        num_heads,
-        head_dim,
-        device="cuda",
-        dtype=torch.bfloat16,
-    )
-    k = torch.randn_like(q)
-    v = torch.randn_like(q)
-
-    # FP8 path
-    attn_fp8 = None
-    with contextlib.suppress(ValueError, ImportError):
-        attn_fp8 = MMEncoderAttention(
-            num_heads=num_heads,
-            head_size=head_dim,
-            prefix="visual.blocks.0.attn",
-        ).to("cuda")
-
-    if attn_fp8 is None or not attn_fp8.fp8_enabled:
-        pytest.skip("FP8 MMEncoderAttention not available")
-    assert attn_fp8 is not None  # mypy narrowing
-
-    fp8_padded_hidden_size = num_heads * round_up(head_dim, 16)
-    cu_seqlens, max_seqlen, seq_lengths = _build_cu_seqlens_and_meta(
-        seq_len,
-        num_heads,
-        head_dim,
-        fp8_padded_hidden_size=fp8_padded_hidden_size,
-    )
-
-    out_fp8 = attn_fp8._forward_flashinfer(
-        q.clone(),
-        k.clone(),
-        v.clone(),
-        cu_seqlens,
-        max_seqlen,
-        seq_lengths,
-    )
-
-    # BF16 baseline (create non-FP8 attention by using scale=attn_fp8.scale
-    # and calling the wrapper directly without FP8 quantization)
-    from vllm.model_executor.layers.attention.mm_encoder_attention import (
-        _get_flashinfer_workspace_buffer,
-    )
-    from vllm.v1.attention.ops.vit_attn_wrappers import (
-        vit_flashinfer_wrapper,
-    )
-
-    out_bf16 = vit_flashinfer_wrapper(
-        q=q.clone(),
-        k=k.clone(),
-        v=v.clone(),
-        scale=attn_fp8.scale,
-        workspace_buffer=_get_flashinfer_workspace_buffer(),
-        cu_seqlens=cu_seqlens,
-        max_seqlen=max_seqlen,
-        sequence_lengths=seq_lengths,
-    )
-
-    out_fp8_f = out_fp8.float()
-    out_bf16_f = out_bf16.float()
-
-    abs_diff = (out_fp8_f - out_bf16_f).abs()
-    abs_diff_flat = abs_diff.flatten()
-
-    # Relative diff (avoid division by zero)
-    denom = out_bf16_f.abs().clamp(min=1e-6)
-    rel_diff_flat = (abs_diff / denom).flatten()
-
-    cosine_sim = torch.nn.functional.cosine_similarity(
-        out_fp8_f.flatten().unsqueeze(0),
-        out_bf16_f.flatten().unsqueeze(0),
-    ).item()
-
-    pcts = [50, 90, 95, 99, 99.9]
-    abs_pct = {p: torch.quantile(abs_diff_flat, p / 100).item() for p in pcts}
-    rel_pct = {p: torch.quantile(rel_diff_flat, p / 100).item() for p in pcts}
-
-    print(f"\nFP8 vs BF16 (head_dim={head_dim}, seq_len={seq_len}):")
-    print(f"  cosine_sim={cosine_sim:.6f}")
-    print(
-        f"  abs_diff: max={abs_diff_flat.max().item():.6f}, "
-        f"mean={abs_diff_flat.mean().item():.6f}, "
-        + ", ".join(f"p{p}={abs_pct[p]:.6f}" for p in pcts)
-    )
-    print(
-        f"  rel_diff: max={rel_diff_flat.max().item():.6f}, "
-        f"mean={rel_diff_flat.mean().item():.6f}, "
-        + ", ".join(f"p{p}={rel_pct[p]:.6f}" for p in pcts)
-    )
-
-    assert abs_diff_flat.max().item() < 0.3, (
-        f"FP8 vs BF16 max abs diff too large: {abs_diff_flat.max().item()}"
-    )
-    assert abs_diff_flat.mean().item() < 0.03, (
-        f"FP8 vs BF16 mean abs diff too large: {abs_diff_flat.mean().item()}"
-    )
-    assert cosine_sim > 0.99, f"Cosine similarity too low: {cosine_sim:.6f}"

tests/kernels/core/test_vit_fp8_quant.py

@@ -1,124 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Tests for the stride-aware FP8 quantization kernel with head_dim padding."""
-
-import pytest
-import torch
-
-from vllm.platforms import current_platform
-from vllm.triton_utils import HAS_TRITON
-
-if HAS_TRITON:
-    from vllm.kernels.triton.qkv_padded_fp8_quant import (
-        quantize_fp8_pad_head_dim_triton,
-    )
-
-HEAD_DIMS = [72, 80, 128]
-SEQ_LENS = [64, 256]
-NUM_HEADS = [16]
-SCALES = [0.01, 0.1, 1.0]
-
-
-def _naive_fp8_quantize(
-    tensor: torch.Tensor, scale: torch.Tensor, skip_scale: bool
-) -> torch.Tensor:
-    """Reference FP8 quantization in PyTorch."""
-    fp8_dtype = current_platform.fp8_dtype()
-    fp8_max = torch.finfo(fp8_dtype).max
-    fp8_min = -fp8_max
-
-    x = tensor.float()
-    if not skip_scale:
-        x = x / scale.item()
-    x = x.clamp(fp8_min, fp8_max)
-    return x.to(fp8_dtype)
-
-
-@pytest.mark.skipif(not HAS_TRITON, reason="Triton not available")
-@pytest.mark.parametrize("head_dim", HEAD_DIMS)
-@pytest.mark.parametrize("seq_len", SEQ_LENS)
-@pytest.mark.parametrize("num_heads", NUM_HEADS)
-@pytest.mark.parametrize("scale_val", SCALES)
-def test_quantize_contiguous(
-    head_dim: int, seq_len: int, num_heads: int, scale_val: float
-) -> None:
-    """Test quantization of contiguous 3D tensors."""
-    torch.manual_seed(42)
-    tensor = torch.randn(
-        seq_len, num_heads, head_dim, device="cuda", dtype=torch.bfloat16
-    )
-    scale = torch.tensor([scale_val], dtype=torch.float32, device="cuda").view(
-        1, 1, 1, 1
-    )
-
-    result = quantize_fp8_pad_head_dim_triton(tensor, scale)
-
-    padded_dim = (head_dim + 15) // 16 * 16
-    assert result.shape == (seq_len, num_heads, padded_dim)
-    assert result.is_contiguous()
-    assert result.dtype == current_platform.fp8_dtype()
-
-    # Compare unpadded portion against reference
-    ref = _naive_fp8_quantize(tensor, scale, skip_scale=False)
-    torch.testing.assert_close(result[:, :, :head_dim].float(), ref.float())
-
-    # Padded region should be zero
-    if padded_dim > head_dim:
-        assert (result[:, :, head_dim:].float() == 0).all()
-
-
-@pytest.mark.skipif(not HAS_TRITON, reason="Triton not available")
-@pytest.mark.parametrize("head_dim", [72, 80])
-def test_quantize_non_contiguous(head_dim: int) -> None:
-    """Test quantization from non-contiguous QKV views (interleaved buffer)."""
-    seq_len, num_heads = 64, 16
-    # Simulate interleaved QKV buffer: shape (seq_len, 3 * num_heads, head_dim)
-    qkv = torch.randn(
-        seq_len, 3 * num_heads, head_dim, device="cuda", dtype=torch.bfloat16
-    )
-    # Q is every 3rd head slice - non-contiguous view
-    q = qkv[:, 0::3, :]
-    assert not q.is_contiguous()
-
-    scale = torch.tensor([0.1], dtype=torch.float32, device="cuda").view(1, 1, 1, 1)
-    result = quantize_fp8_pad_head_dim_triton(q, scale)
-
-    padded_dim = (head_dim + 15) // 16 * 16
-    assert result.shape == (seq_len, num_heads, padded_dim)
-    assert result.is_contiguous()
-
-    # Compare against contiguous reference
-    ref = _naive_fp8_quantize(q.contiguous(), scale, skip_scale=False)
-    torch.testing.assert_close(result[:, :, :head_dim].float(), ref.float())
-
-
-@pytest.mark.skipif(not HAS_TRITON, reason="Triton not available")
-def test_skip_scale() -> None:
-    """Test skip_scale=True produces cast-only output (no division)."""
-    seq_len, num_heads, head_dim = 32, 8, 80
-    tensor = torch.randn(
-        seq_len, num_heads, head_dim, device="cuda", dtype=torch.bfloat16
-    )
-    scale = torch.tensor([0.5], dtype=torch.float32, device="cuda").view(1, 1, 1, 1)
-
-    result_skip = quantize_fp8_pad_head_dim_triton(tensor, scale, skip_scale=True)
-    result_noskip = quantize_fp8_pad_head_dim_triton(tensor, scale, skip_scale=False)
-
-    # skip_scale should just cast, not divide
-    ref_cast = _naive_fp8_quantize(tensor, scale, skip_scale=True)
-    torch.testing.assert_close(result_skip[:, :, :head_dim].float(), ref_cast.float())
-
-    # With scale != 1.0, skip and no-skip should differ
-    assert not torch.equal(result_skip.float(), result_noskip.float())
-
-
-@pytest.mark.skipif(not HAS_TRITON, reason="Triton not available")
-def test_4d_input() -> None:
-    """Test that 4D input (B, S, H, D) is handled correctly."""
-    B, S, H, D = 2, 32, 8, 72
-    tensor = torch.randn(B, S, H, D, device="cuda", dtype=torch.bfloat16)
-    scale = torch.tensor([0.1], dtype=torch.float32, device="cuda").view(1, 1, 1, 1)
-
-    result = quantize_fp8_pad_head_dim_triton(tensor, scale)
-    padded_dim = (D + 15) // 16 * 16
-    assert result.shape == (B, S, H, padded_dim)

tests/kernels/core/test_vit_fp8_scaling.py

@@ -1,251 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Tests for FP8 scaling (dynamic and static) in MMEncoderAttention."""
-
-import contextlib
-import json
-from types import SimpleNamespace
-from unittest.mock import patch
-
-import pytest
-import torch
-
-from vllm.model_executor.layers.attention.mm_encoder_attention import (
-    _FP8_AMAX_HISTORY_LEN,
-    _FP8_MAX,
-)
-from vllm.utils.flashinfer import (
-    is_flashinfer_cudnn_fp8_prefill_attn_supported,
-)
-
-LAYER_0 = "visual.blocks.0.attn.attn"
-LAYER_1 = "visual.blocks.1.attn.attn"
-NUM_HEADS = 16
-HEAD_DIM = 72
-
-
-@contextlib.contextmanager
-def _build_attention(mm_config):
-    """Yield an MMEncoderAttention with the given multimodal config.
-
-    The VllmConfig context stays active while the test runs so that
-    ``get_multimodal_config()`` calls during the forward path resolve. Also
-    invokes ``process_weights_after_loading`` to simulate the model loader's
-    auto-scan. Yields ``None`` if FlashInfer cuDNN is not available.
-    """
-    from vllm.config import VllmConfig, set_current_vllm_config
-    from vllm.model_executor.layers.attention.mm_encoder_attention import (
-        MMEncoderAttention,
-    )
-    from vllm.v1.attention.backends.registry import AttentionBackendEnum
-
-    if not is_flashinfer_cudnn_fp8_prefill_attn_supported():
-        yield None
-        return
-
-    vllm_config = VllmConfig()
-    vllm_config.model_config = SimpleNamespace(multimodal_config=mm_config)
-
-    with (
-        set_current_vllm_config(vllm_config),
-        patch(
-            "vllm.model_executor.layers.attention.mm_encoder_attention"
-            ".get_vit_attn_backend",
-            return_value=AttentionBackendEnum.FLASHINFER,
-        ),
-    ):
-        attn = MMEncoderAttention(
-            num_heads=NUM_HEADS,
-            head_size=HEAD_DIM,
-            prefix=LAYER_0,
-        )
-        attn.process_weights_after_loading(torch.bfloat16)
-        yield attn
-
-
-@pytest.fixture
-def _make_attention():
-    """Create an MMEncoderAttention with dynamic FP8 scaling."""
-    from vllm.config.multimodal import MultiModalConfig
-
-    with _build_attention(MultiModalConfig(mm_encoder_attn_dtype="fp8")) as attn:
-        yield attn
-
-
-@pytest.fixture
-def _make_static_attention(tmp_path):
-    """Create an MMEncoderAttention with static FP8 scales from a file."""
-    from vllm.config.multimodal import MultiModalConfig
-
-    scale_file = tmp_path / "scales.json"
-    scale_file.write_text(
-        json.dumps(
-            {
-                LAYER_0: {"q": 224.0, "k": 198.0, "v": 210.0},
-                LAYER_1: {"q": 100.0, "k": 110.0, "v": 120.0},
-            }
-        )
-    )
-    with _build_attention(
-        MultiModalConfig(
-            mm_encoder_attn_dtype="fp8",
-            mm_encoder_fp8_scale_path=str(scale_file),
-        )
-    ) as attn:
-        yield attn
-
-
-def test_dynamic_scaling_updates_scales(_make_attention) -> None:
-    """Verify that _record_amax_and_update_scales updates scale buffers."""
-    attn = _make_attention
-    if attn is None or not attn.fp8_enabled:
-        pytest.skip("FP8 attention not available (FlashInfer backend required)")
-
-    attn = attn.to("cuda")
-
-    S, H, D = 32, NUM_HEADS, HEAD_DIM
-    q = torch.full((S, H, D), 2.0, device="cuda", dtype=torch.bfloat16)
-    k = torch.full((S, H, D), 3.0, device="cuda", dtype=torch.bfloat16)
-    v = torch.full((S, H, D), 4.0, device="cuda", dtype=torch.bfloat16)
-
-    attn._record_amax_and_update_scales(q, k, v)
-
-    expected_q_scale = 2.0 / _FP8_MAX
-    expected_k_scale = 3.0 / _FP8_MAX
-    expected_v_scale = 4.0 / _FP8_MAX
-
-    torch.testing.assert_close(attn._fp8_q_scale.item(), expected_q_scale)
-    torch.testing.assert_close(attn._fp8_k_scale.item(), expected_k_scale)
-    torch.testing.assert_close(attn._fp8_v_scale.item(), expected_v_scale)
-
-
-def test_circular_buffer_wraps(_make_attention) -> None:
-    """Verify the amax circular buffer wraps at HISTORY_LEN."""
-    attn = _make_attention
-    if attn is None or not attn.fp8_enabled:
-        pytest.skip("FP8 attention not available (FlashInfer backend required)")
-
-    attn = attn.to("cuda")
-    S, H, D = 16, NUM_HEADS, HEAD_DIM
-
-    for i in range(_FP8_AMAX_HISTORY_LEN + 2):
-        mag = float(i + 1)
-        q = torch.full((S, H, D), mag, device="cuda", dtype=torch.bfloat16)
-        k = torch.full((S, H, D), mag, device="cuda", dtype=torch.bfloat16)
-        v = torch.full((S, H, D), mag, device="cuda", dtype=torch.bfloat16)
-        attn._record_amax_and_update_scales(q, k, v)
-
-    assert attn._fp8_amax_pos == 2
-
-    expected_max = float(_FP8_AMAX_HISTORY_LEN + 2)
-    expected_scale = expected_max / _FP8_MAX
-    torch.testing.assert_close(attn._fp8_q_scale.item(), expected_scale)
-
-
-def test_static_scales_loaded(_make_static_attention) -> None:
-    """Verify static scales are loaded from the JSON file."""
-    attn = _make_static_attention
-    if attn is None or not attn.fp8_enabled:
-        pytest.skip("FP8 attention not available (FlashInfer backend required)")
-
-    assert attn.fp8_enabled
-    assert not attn._fp8_dynamic_scale
-
-    # Layer 0 scales (the layer this attention was created with).
-    assert attn._fp8_q_scale.item() == 224.0
-    assert attn._fp8_k_scale.item() == 198.0
-    assert attn._fp8_v_scale.item() == 210.0
-
-    assert not attn.skip_scale_q
-    assert not attn.skip_scale_k
-    assert not attn.skip_scale_v
-
-    # No amax history buffers for static scaling.
-    assert not hasattr(attn, "_fp8_q_amax")
-
-
-def test_static_scales_missing_layer(tmp_path) -> None:
-    """Verify error when requested layer is not in the scale file."""
-    from vllm.config import VllmConfig, set_current_vllm_config
-    from vllm.config.multimodal import MultiModalConfig
-    from vllm.v1.attention.backends.registry import AttentionBackendEnum
-
-    if not is_flashinfer_cudnn_fp8_prefill_attn_supported():
-        pytest.skip("FlashInfer cuDNN not available")
-
-    scale_file = tmp_path / "wrong_layer.json"
-    scale_file.write_text(
-        json.dumps({"visual.blocks.99.attn": {"q": 1.0, "k": 1.0, "v": 1.0}})
-    )
-    mm_config = MultiModalConfig(
-        mm_encoder_attn_dtype="fp8",
-        mm_encoder_fp8_scale_path=str(scale_file),
-    )
-    vllm_config = VllmConfig()
-    vllm_config.model_config = SimpleNamespace(multimodal_config=mm_config)
-
-    from vllm.model_executor.layers.attention.mm_encoder_attention import (
-        MMEncoderAttention,
-    )
-
-    with (
-        set_current_vllm_config(vllm_config),
-        patch(
-            "vllm.model_executor.layers.attention.mm_encoder_attention"
-            ".get_vit_attn_backend",
-            return_value=AttentionBackendEnum.FLASHINFER,
-        ),
-    ):
-        attn = MMEncoderAttention(
-            num_heads=NUM_HEADS,
-            head_size=HEAD_DIM,
-            prefix=LAYER_0,
-        )
-        with pytest.raises(ValueError, match="scales not found for layer"):
-            attn.process_weights_after_loading(torch.bfloat16)
-
-
-def test_dynamic_scales_auto_save(tmp_path) -> None:
-    """Verify scales are saved to disk after the amax buffer fills."""
-    import vllm.model_executor.layers.attention.mm_encoder_attention as _mod
-    from vllm.config.multimodal import MultiModalConfig
-
-    if not is_flashinfer_cudnn_fp8_prefill_attn_supported():
-        pytest.skip("FlashInfer cuDNN not available")
-
-    # Reset module-level state between runs (other tests may have left
-    # state behind after triggering a save).
-    _mod._fp8_scale_save_path = None
-    _mod._fp8_saved_scale_refs.clear()
-
-    save_file = tmp_path / "auto_scales.json"
-    with _build_attention(
-        MultiModalConfig(
-            mm_encoder_attn_dtype="fp8",
-            mm_encoder_fp8_scale_save_path=str(save_file),
-        )
-    ) as attn:
-        if attn is None or not attn.fp8_enabled:
-            pytest.skip("FP8 attention not available")
-
-        attn = attn.to("cuda")
-        S, H, D = 16, NUM_HEADS, HEAD_DIM
-
-        # Run exactly _FP8_AMAX_HISTORY_LEN forward passes.
-        for i in range(_FP8_AMAX_HISTORY_LEN):
-            mag = float(i + 1)
-            q = torch.full((S, H, D), mag, device="cuda", dtype=torch.bfloat16)
-            k = torch.full((S, H, D), mag * 0.5, device="cuda", dtype=torch.bfloat16)
-            v = torch.full((S, H, D), mag * 0.3, device="cuda", dtype=torch.bfloat16)
-            attn._record_amax_and_update_scales(q, k, v)
-
-    # File should have been written on the 16th call (buffer wrap).
-    assert save_file.is_file(), "Scale file was not saved"
-    scales = json.loads(save_file.read_text())
-    assert LAYER_0 in scales
-    assert set(scales[LAYER_0].keys()) == {"q", "k", "v"}
-    for val in scales[LAYER_0].values():
-        assert isinstance(val, float) and val > 0
-
-    # Path is cleared after the one-shot save fires.
-    assert _mod._fp8_scale_save_path is None

tests/kernels/moe/test_deepgemm.py

@@ -1,7 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """
-Unit-test DeepGEMM FP8 and FP4 kernels (no DeepEP).
+Unit-test DeepGEMM FP8 kernels (no DeepEP).
 Compare DeepGEMM path against the Triton fallback inside vLLM's fused_experts.
 """
 
@@ -21,8 +21,6 @@ from vllm.model_executor.layers.fused_moe.all2all_utils import (
     maybe_make_prepare_finalize,
 )
 from vllm.model_executor.layers.fused_moe.config import (
-    FusedMoEQuantConfig,
-    FusedMoEQuantDesc,
     fp8_w8a8_moe_quant_config,
 )
 from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts
@@ -206,195 +204,3 @@ def test_deepgemm_vs_triton(m, n, k, topk, num_experts, monkeypatch, workspace_i
             f"DeepGEMM path was not executed during the test. "
             f"Call counter: {call_counter['cnt']}"
         )
-
-
-# ---------------------------------------------------------------------------
-# FP4 weight tests (DeepGEMM m_grouped_fp8_fp4_gemm_nt_contiguous)
-# ---------------------------------------------------------------------------
-
-
-def make_mxfp4_weights(
-    e: int,
-    n: int,
-    k: int,
-):
-    """
-    Generate (w1, w2) expert weights in MXFP4 packed format with float32 scales,
-    plus BF16 reference weights for validation.
-
-      w1 shape: (E, 2N, K//2) uint8    — packed FP4
-      w2 shape: (E, K, N//2)  uint8    — packed FP4
-      w1_s shape: (E, 2N, K//32) float32  — per-row block-32 scales
-      w2_s shape: (E, K, N//32)  float32  — per-row block-32 scales
-      w1_bf16: (E, 2N, K)   — original BF16 for reference
-      w2_bf16: (E, K, N)    — original BF16 for reference
-    """
-    from deep_gemm.utils.math import per_token_cast_to_fp4
-
-    dtype = torch.bfloat16
-    gran_k = 32  # MXFP4 block size
-
-    # bf16 reference weights — scale by 1/sqrt(dim) for numerical stability
-    w1_bf16 = torch.randn(e, 2 * n, k, device="cuda", dtype=dtype) * (k**-0.5)
-    w2_bf16 = torch.randn(e, k, n, device="cuda", dtype=dtype) * (n**-0.5)
-
-    # Quantize per-expert to FP4
-    w1 = torch.empty(e, 2 * n, k // 2, device="cuda", dtype=torch.uint8)
-    w2 = torch.empty(e, k, n // 2, device="cuda", dtype=torch.uint8)
-    w1_s = torch.empty(
-        e, 2 * n, math.ceil(k / gran_k), device="cuda", dtype=torch.float32
-    )
-    w2_s = torch.empty(e, k, math.ceil(n / gran_k), device="cuda", dtype=torch.float32)
-
-    for i in range(e):
-        w1[i], w1_s[i] = per_token_cast_to_fp4(
-            w1_bf16[i].float(), use_ue8m0=True, gran_k=gran_k
-        )
-        w2[i], w2_s[i] = per_token_cast_to_fp4(
-            w2_bf16[i].float(), use_ue8m0=True, gran_k=gran_k
-        )
-
-    return w1, w2, w1_s, w2_s, w1_bf16, w2_bf16
-
-
-def _bf16_moe_reference(x, w1, w2, topk_weights, topk_ids):
-    """BF16 token-loop MoE reference for correctness testing."""
-    import torch.nn.functional as F
-
-    num_tokens, hidden_size = x.shape
-    intermediate = w1.shape[1] // 2
-    top_k = topk_ids.shape[1]
-
-    output = torch.zeros(num_tokens, hidden_size, dtype=torch.float32, device=x.device)
-    for t in range(num_tokens):
-        for kk in range(top_k):
-            e = topk_ids[t, kk].item()
-            w = topk_weights[t, kk].item()
-            fc1 = x[t : t + 1].float() @ w1[e].float().T
-            linear = fc1[:, :intermediate]
-            gate = fc1[:, intermediate:]
-            act = F.silu(gate) * linear
-            fc2 = act @ w2[e].float().T
-            output[t] += w * fc2[0]
-    return output.to(torch.bfloat16)
-
-
-def run_single_fp4_case(m, n, k, topk, num_experts):
-    """
-    Run one (M,N,K) configuration with FP4 weights on DeepGEMM and assert
-    DeepGEMM FP4 == BF16 reference within tolerance.
-    """
-    tokens_bf16 = torch.randn(m, k, device="cuda", dtype=torch.bfloat16) * (k**-0.5)
-
-    # FP4 expert weight tensors + BF16 originals for reference
-    w1, w2, w1_s, w2_s, w1_bf16, w2_bf16 = make_mxfp4_weights(num_experts, n, k)
-
-    router_logits = torch.randn(m, num_experts, device="cuda", dtype=torch.float32)
-    topk_weights, topk_ids = torch.topk(router_logits, k=topk, dim=-1)
-    topk_weights = torch.nn.functional.softmax(topk_weights, dim=-1)
-
-    from vllm.model_executor.layers.quantization.utils.quant_utils import (
-        GroupShape,
-    )
-    from vllm.platforms import current_platform
-
-    _fp8_dtype = current_platform.fp8_dtype()
-    _block_shape = GroupShape(128, 128)
-    quant_config = FusedMoEQuantConfig(
-        _a1=FusedMoEQuantDesc(_fp8_dtype, _block_shape, None, None, None, None),
-        _a2=FusedMoEQuantDesc(_fp8_dtype, _block_shape, None, None, None, None),
-        _w1=FusedMoEQuantDesc("mxfp4", None, w1_s, None, None, None),
-        _w2=FusedMoEQuantDesc("mxfp4", None, w2_s, None, None, None),
-    )
-    moe_config = make_dummy_moe_config()
-
-    from vllm.model_executor.layers.fused_moe.experts.deep_gemm_moe import (
-        DeepGemmFP4Experts,
-    )
-
-    deep_gemm_fp4_experts = mk.FusedMoEKernel(
-        prepare_finalize=maybe_make_prepare_finalize(
-            moe=moe_config,
-            quant_config=quant_config,
-            allow_new_interface=True,
-            use_monolithic=False,
-        ),
-        fused_experts=DeepGemmFP4Experts(
-            moe_config=moe_config,
-            quant_config=quant_config,
-        ),
-        inplace=False,
-    )
-
-    # DeepGEMM FP4 path
-    out_deepgemm_fp4 = deep_gemm_fp4_experts.apply(
-        hidden_states=tokens_bf16,
-        w1=w1,
-        w2=w2,
-        topk_weights=topk_weights,
-        topk_ids=topk_ids,
-        global_num_experts=num_experts,
-        activation=MoEActivation.SILU,
-        apply_router_weight_on_input=False,
-        expert_map=None,
-    )
-
-    # BF16 reference using the same original weights
-    out_ref = _bf16_moe_reference(tokens_bf16, w1_bf16, w2_bf16, topk_weights, topk_ids)
-
-    # FP4 vs BF16 reference: quantization error from FP4 weights + FP8 activations
-    diff = calc_diff(out_deepgemm_fp4, out_ref)
-    assert diff < 0.05, f"FP4 diff exceeded 5%: {diff}"
-
-
-# DeepSeek V4 dims: H=4096, I=2048, so N=2*I=4096, K=H=4096.
-# FP4 quantization with block_k=32 needs large K for good accuracy.
-FP4_MNKs = [
-    (128, 4096, 4096),  # DeepSeek V4 shape
-    (256, 2048, 2048),  # Half-size variant
-]
-
-FP4_TOPKS = [2]
-FP4_NUM_EXPERTS = [8]
-
-
-@pytest.mark.parametrize(("m", "n", "k"), FP4_MNKs)
-@pytest.mark.parametrize("topk", FP4_TOPKS)
-@pytest.mark.parametrize("num_experts", FP4_NUM_EXPERTS)
-@pytest.mark.skipif(not is_deep_gemm_supported(), reason="Requires deep_gemm kernels")
-def test_deepgemm_fp4_vs_triton(
-    m, n, k, topk, num_experts, monkeypatch, workspace_init
-):
-    pytest.importorskip("deep_gemm.utils.math")
-    with monkeypatch.context() as mp:
-        mp.setenv("VLLM_USE_DEEP_GEMM", "1")
-
-        _DeepGemmFP4Experts = importlib.import_module(
-            "vllm.model_executor.layers.fused_moe.experts.deep_gemm_moe"
-        ).DeepGemmFP4Experts
-
-        call_counter = {"cnt": 0}
-
-        orig_fn = _DeepGemmFP4Experts.apply
-
-        def _spy_apply(*args, **kwargs):
-            call_counter["cnt"] += 1
-            return orig_fn(*args, **kwargs)
-
-        monkeypatch.setattr(_DeepGemmFP4Experts, "apply", _spy_apply)
-        if topk > num_experts:
-            pytest.skip(f"topk={topk} > num_experts={num_experts}")
-
-        run_single_fp4_case(
-            m=m,
-            n=n,
-            k=k,
-            topk=topk,
-            num_experts=num_experts,
-        )
-
-        # ensure that the DeepGEMM FP4 path was indeed taken.
-        assert call_counter["cnt"] == 1, (
-            f"DeepGEMM FP4 path was not executed during the test. "
-            f"Call counter: {call_counter['cnt']}"
-        )

tests/kernels/moe/test_moe_layer.py

@@ -1632,7 +1632,6 @@ def _parallel_worker(
                 if all2all_manager is not None:
                     all2all_manager.destroy()
             total = total + 1
-            torch.distributed.barrier()
 
     skipped = total - (passed + failed)
 

tests/kernels/moe/test_topk_softplus_sqrt.py

@@ -1,186 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-
-import pytest
-import torch
-import torch.nn.functional as F
-
-from vllm.model_executor.layers.fused_moe.config import (
-    RoutingMethodType,
-    get_routing_method_type,
-)
-from vllm.model_executor.layers.fused_moe.router.fused_topk_bias_router import (
-    fused_topk_bias,
-)
-from vllm.platforms import current_platform
-
-
-def _torch_topk_softplus_sqrt(
-    gating_output: torch.Tensor,
-    topk: int,
-    renormalize: bool,
-    routed_scaling_factor: float,
-    e_score_correction_bias: torch.Tensor | None = None,
-    input_ids: torch.Tensor | None = None,
-    hash_indices_table: torch.Tensor | None = None,
-):
-    scores = F.softplus(gating_output.float()).sqrt()
-    original_scores = scores
-    if e_score_correction_bias is not None:
-        scores_for_choice = scores + e_score_correction_bias.unsqueeze(0)
-    else:
-        scores_for_choice = scores
-
-    if hash_indices_table is not None:
-        assert input_ids is not None
-        topk_ids = hash_indices_table[input_ids.long()]
-    else:
-        topk_ids = torch.topk(scores_for_choice, k=topk, dim=-1, sorted=True)[1]
-
-    topk_weights = original_scores.gather(1, topk_ids.long())
-    if renormalize:
-        topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
-    if routed_scaling_factor != 1.0:
-        topk_weights = topk_weights * routed_scaling_factor
-    return topk_weights.to(torch.float32), topk_ids.to(torch.int32)
-
-
-def test_sqrtsoftplus_bias_uses_deepseek_v4_routing_method():
-    assert (
-        get_routing_method_type(
-            scoring_func="sqrtsoftplus",
-            top_k=8,
-            renormalize=True,
-            num_expert_group=None,
-            has_e_score_bias=True,
-        )
-        == RoutingMethodType.DeepseekV4
-    )
-    assert (
-        get_routing_method_type(
-            scoring_func="sqrtsoftplus",
-            top_k=8,
-            renormalize=False,
-            num_expert_group=None,
-            has_e_score_bias=True,
-        )
-        == RoutingMethodType.Unspecified
-    )
-
-
-@pytest.mark.skipif(
-    not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform."
-)
-@pytest.mark.parametrize("num_tokens", [1, 33, 128])
-@pytest.mark.parametrize("hidden_size", [1024, 2048])
-@pytest.mark.parametrize("num_experts", [128, 256, 384, 512])
-@pytest.mark.parametrize("topk", [6, 8, 16])
-@pytest.mark.parametrize("renormalize", [True, False])
-@pytest.mark.parametrize("routed_scaling_factor", [1.0, 1.5])
-@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.half, torch.float32])
-def test_fused_topk_softplus_sqrt(
-    num_tokens: int,
-    hidden_size: int,
-    num_experts: int,
-    topk: int,
-    renormalize: bool,
-    routed_scaling_factor: float,
-    dtype: torch.dtype,
-):
-    torch.manual_seed(0)
-    hidden_states = torch.randn((num_tokens, hidden_size), dtype=dtype, device="cuda")
-    gating_output = torch.randn((num_tokens, num_experts), dtype=dtype, device="cuda")
-    e_score_correction_bias = torch.randn(
-        (num_experts,), dtype=torch.float32, device="cuda"
-    )
-
-    topk_weights_ref, topk_ids_ref = _torch_topk_softplus_sqrt(
-        gating_output=gating_output,
-        topk=topk,
-        renormalize=renormalize,
-        routed_scaling_factor=routed_scaling_factor,
-        e_score_correction_bias=e_score_correction_bias,
-    )
-
-    topk_weights, topk_ids = fused_topk_bias(
-        hidden_states=hidden_states,
-        gating_output=gating_output,
-        scoring_func="sqrtsoftplus",
-        e_score_correction_bias=e_score_correction_bias,
-        topk=topk,
-        renormalize=renormalize,
-        routed_scaling_factor=routed_scaling_factor,
-    )
-
-    # Different kernels may return the topk experts in different orders when
-    # scores tie; sort by expert id before comparing.
-    sorted_ref_ids, idx_ref = topk_ids_ref.sort(dim=-1)
-    sorted_ids, idx_ops = topk_ids.sort(dim=-1)
-    torch.testing.assert_close(sorted_ref_ids, sorted_ids, atol=0, rtol=0)
-
-    sorted_w_ref = topk_weights_ref.gather(1, idx_ref)
-    sorted_w = topk_weights.gather(1, idx_ops)
-    torch.testing.assert_close(sorted_w_ref, sorted_w, atol=2e-2, rtol=1e-2)
-
-
-@pytest.mark.skipif(
-    not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform."
-)
-@pytest.mark.parametrize("num_tokens", [1, 33, 128])
-@pytest.mark.parametrize("hidden_size", [1024, 2048])
-@pytest.mark.parametrize("num_experts", [256, 384, 512])
-@pytest.mark.parametrize("topk", [6, 8, 16])
-@pytest.mark.parametrize("renormalize", [True, False])
-@pytest.mark.parametrize("routed_scaling_factor", [1.0, 2.5])
-@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.half, torch.float32])
-def test_fused_topk_softplus_sqrt_hash(
-    num_tokens: int,
-    hidden_size: int,
-    num_experts: int,
-    topk: int,
-    renormalize: bool,
-    routed_scaling_factor: float,
-    dtype: torch.dtype,
-):
-    torch.manual_seed(0)
-    vocab_size = 1024
-    hidden_states = torch.randn((num_tokens, hidden_size), dtype=dtype, device="cuda")
-    gating_output = torch.randn((num_tokens, num_experts), dtype=dtype, device="cuda")
-    # Per-token fixed expert selection: for each vocab id pick `topk` distinct
-    # experts.
-    hash_indices_table = torch.stack(
-        [torch.randperm(num_experts)[:topk] for _ in range(vocab_size)]
-    ).to(device="cuda", dtype=torch.int32)
-    input_ids = torch.randint(
-        0, vocab_size, (num_tokens,), dtype=torch.int32, device="cuda"
-    )
-
-    topk_weights_ref, topk_ids_ref = _torch_topk_softplus_sqrt(
-        gating_output=gating_output,
-        topk=topk,
-        renormalize=renormalize,
-        routed_scaling_factor=routed_scaling_factor,
-        input_ids=input_ids,
-        hash_indices_table=hash_indices_table,
-    )
-
-    topk_weights, topk_ids = fused_topk_bias(
-        hidden_states=hidden_states,
-        gating_output=gating_output,
-        scoring_func="sqrtsoftplus",
-        e_score_correction_bias=None,
-        topk=topk,
-        renormalize=renormalize,
-        input_tokens=input_ids,
-        hash_indices_table=hash_indices_table,
-        routed_scaling_factor=routed_scaling_factor,
-    )
-
-    sorted_ref_ids, idx_ref = topk_ids_ref.sort(dim=-1)
-    sorted_ids, idx_ops = topk_ids.sort(dim=-1)
-    torch.testing.assert_close(sorted_ref_ids, sorted_ids, atol=0, rtol=0)
-
-    sorted_w_ref = topk_weights_ref.gather(1, idx_ref)
-    sorted_w = topk_weights.gather(1, idx_ops)
-    torch.testing.assert_close(sorted_w_ref, sorted_w, atol=2e-2, rtol=1e-2)

tests/kernels/test_compressor_kv_cache.py

@@ -1,311 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Round-trip tests for compressor → FP8 quant + KV cache insert → gather + dequant.
-
-Two paths tested:
-  A) DeepseekV4 Attention: head_dim=512 (448 FP8 nope + 64 bf16 rope), quant_block=64
-  B) Indexer:       head_dim=128 (all FP8), quant_block=128
-
-These serve as golden references for validating the future fused
-compressor+quant+cache kernel.
-"""
-
-import math
-
-import pytest
-import torch
-
-from vllm import _custom_ops as ops
-from vllm.v1.attention.ops.deepseek_v4_ops import (
-    dequantize_and_gather_k_cache,
-    quantize_and_insert_k_cache,
-)
-
-
-def _ue8m0_reference(x: torch.Tensor, block_size: int, fp8_max: float):
-    """PyTorch reference for UE8M0 FP8 quantization (per-block, power-of-2 scale).
-
-    Returns (x_fp8, scales) where x_fp8 is float8_e4m3fn and scales are float32.
-    """
-    assert x.dim() == 1
-    n = x.numel()
-    n_blocks = math.ceil(n / block_size)
-    x_fp8 = torch.zeros(n, dtype=torch.float8_e4m3fn, device=x.device)
-    scales = torch.zeros(n_blocks, dtype=torch.float32, device=x.device)
-
-    for i in range(n_blocks):
-        start = i * block_size
-        end = min(start + block_size, n)
-        block = x[start:end].float()
-        amax = block.abs().max().clamp(min=1e-4)
-        raw_scale = amax / fp8_max
-        exponent = math.ceil(math.log2(raw_scale.item()))
-        scale = 2.0**exponent
-        scales[i] = scale
-        quantized = (block / scale).clamp(-fp8_max, fp8_max)
-        x_fp8[start:end] = quantized.to(torch.float8_e4m3fn)
-
-    return x_fp8, scales
-
-
-# ── Test A: DeepseekV4 Attention path ──────────────────────────────────────────────
-
-
-@pytest.mark.parametrize("num_tokens", [1, 4, 8, 17])
-@pytest.mark.parametrize("block_size", [16, 64])
-def test_deepseek_v4_attention_quant_cache_roundtrip(num_tokens: int, block_size: int):
-    """compressed_kv → quantize_and_insert_k_cache → dequantize_and_gather_k_cache
-    → compare against original."""
-
-    HEAD_DIM = 512
-    NOPE_DIM = 448
-    HEAD_BYTES = 584  # 448 fp8 + 128 bf16 + 8 uint8 scale
-    FP8_MAX = 448.0
-    QUANT_BLOCK = 64
-
-    num_blocks = (num_tokens + block_size - 1) // block_size + 1
-    device = "cuda"
-
-    # Random compressed_kv (simulates compressor output)
-    compressed_kv = torch.randn(
-        num_tokens, HEAD_DIM, dtype=torch.bfloat16, device=device
-    )
-
-    # ── Quant + insert ──────────────────────────────────────────────────
-    k_cache = torch.zeros(
-        num_blocks, block_size, HEAD_BYTES, dtype=torch.uint8, device=device
-    )
-    k_cache_2d = k_cache.view(num_blocks, -1)
-
-    # Sequential slot mapping: token i → slot i
-    slot_mapping = torch.arange(num_tokens, dtype=torch.int64, device=device)
-
-    quantize_and_insert_k_cache(
-        compressed_kv, k_cache_2d, slot_mapping, block_size=block_size
-    )
-
-    # ── Gather + dequant ────────────────────────────────────────────────
-    num_reqs = 1
-    max_blocks_per_seq = num_blocks
-    out = torch.zeros(
-        num_reqs, num_tokens, HEAD_DIM, dtype=torch.bfloat16, device=device
-    )
-    seq_lens = torch.tensor([num_tokens], dtype=torch.int32, device=device)
-    # block_table: request 0 uses physical blocks 0, 1, ...
-    block_table = torch.arange(
-        max_blocks_per_seq, dtype=torch.int32, device=device
-    ).unsqueeze(0)
-
-    dequantize_and_gather_k_cache(
-        out, k_cache, seq_lens, None, block_table, block_size, offset=0
-    )
-
-    recovered = out[0, :num_tokens]
-
-    # ── NoPE portion (first 448): FP8 quantized, expect UE8M0 error ──
-    nope_orig = compressed_kv[:, :NOPE_DIM].float()
-    nope_recv = recovered[:, :NOPE_DIM].float()
-    nope_diff = (nope_recv - nope_orig).abs()
-
-    # Per-token check: FP8 e4m3 (3-bit mantissa) worst-case error is
-    # half-ULP at the largest representable value.  At y ≈ 448 (max),
-    # ULP = 2^(8-3) = 32, so error ≤ 16 * scale.
-    for t in range(num_tokens):
-        _, scales = _ue8m0_reference(
-            compressed_kv[t, :NOPE_DIM].float(), QUANT_BLOCK, FP8_MAX
-        )
-        max_allowed = 16.0 * scales.max().item()
-        token_diff = nope_diff[t].max().item()
-        assert token_diff <= max_allowed, (
-            f"Token {t} nope diff {token_diff} exceeds max_allowed "
-            f"{max_allowed} (scale={scales.max().item()})"
-        )
-
-    # ── RoPE portion (last 64): stored as bf16, should be exact ─────
-    rope_diff = (recovered[:, NOPE_DIM:] - compressed_kv[:, NOPE_DIM:]).abs()
-    assert rope_diff.max().item() == 0.0, (
-        f"RoPE portion should be exact but got max diff {rope_diff.max().item()}"
-    )
-
-
-# ── Test B: Indexer path ────────────────────────────────────────────────────
-
-
-@pytest.mark.parametrize("num_tokens", [1, 4, 8, 17])
-@pytest.mark.parametrize("block_size", [16, 64])
-def test_indexer_quant_cache_roundtrip(num_tokens: int, block_size: int):
-    """k → indexer_k_quant_and_cache → cp_gather_indexer_k_quant_cache
-    → manual dequant → compare against original."""
-
-    HEAD_DIM = 128
-    QUANT_BLOCK_SIZE = 128
-    # cache_stride = head_dim + (head_dim * 4 / quant_block_size) = 128 + 4 = 132
-    CACHE_STRIDE = HEAD_DIM + HEAD_DIM * 4 // QUANT_BLOCK_SIZE
-
-    num_blocks = (num_tokens + block_size - 1) // block_size + 1
-    device = "cuda"
-
-    # Random K (simulates compressor output for indexer)
-    k = torch.randn(num_tokens, HEAD_DIM, dtype=torch.bfloat16, device=device)
-
-    # ── Quant + insert ──────────────────────────────────────────────────
-    kv_cache = torch.zeros(
-        num_blocks, block_size, CACHE_STRIDE, dtype=torch.uint8, device=device
-    )
-    slot_mapping = torch.arange(num_tokens, dtype=torch.int64, device=device)
-
-    ops.indexer_k_quant_and_cache(k, kv_cache, slot_mapping, QUANT_BLOCK_SIZE, "ue8m0")
-
-    # ── Gather ──────────────────────────────────────────────────────────
-    max_blocks_per_seq = num_blocks
-    block_table = torch.arange(
-        max_blocks_per_seq, dtype=torch.int32, device=device
-    ).unsqueeze(0)
-    cu_seq_lens = torch.tensor([0, num_tokens], dtype=torch.int32, device=device)
-
-    # dst_k: [total_seq_len, head_dim] as uint8 (raw FP8 bytes)
-    dst_k = torch.zeros(num_tokens, HEAD_DIM, dtype=torch.uint8, device=device)
-    # dst_scale: [total_seq_len, head_dim/quant_block*4] as uint8 (raw float32 bytes)
-    num_scale_bytes = HEAD_DIM * 4 // QUANT_BLOCK_SIZE  # 4
-    dst_scale = torch.zeros(
-        num_tokens, num_scale_bytes, dtype=torch.uint8, device=device
-    )
-
-    ops.cp_gather_indexer_k_quant_cache(
-        kv_cache, dst_k, dst_scale, block_table, cu_seq_lens
-    )
-
-    # ── Manual dequant ──────────────────────────────────────────────────
-    k_fp8 = dst_k.view(torch.float8_e4m3fn).float()  # [num_tokens, 128]
-    scale = dst_scale.view(torch.float32)  # [num_tokens, 1]
-    k_recovered = k_fp8 * scale  # [num_tokens, 128]
-
-    # ── Compare ─────────────────────────────────────────────────────────
-    diff = (k_recovered - k.float()).abs()
-    k_abs = k.float().abs()
-
-    for t in range(num_tokens):
-        amax = k_abs[t].max().clamp(min=1e-4).item()
-        # UE8M0: scale = 2^ceil(log2(amax / 448))
-        exponent = math.ceil(math.log2(amax / 448.0))
-        ue8m0_scale = 2.0**exponent
-        # FP8 e4m3 (3-bit mantissa): worst-case error = 16 * scale
-        max_allowed = 16.0 * ue8m0_scale
-        token_diff = diff[t].max().item()
-        assert token_diff <= max_allowed, (
-            f"Token {t} diff {token_diff} exceeds max_allowed "
-            f"{max_allowed} (scale={ue8m0_scale})"
-        )
-
-
-def test_indexer_gather_accepts_upper_bound_output():
-    """Gather only exact cu_seq_lens even when dst is over-allocated."""
-
-    head_dim = 128
-    quant_block_size = 128
-    cache_stride = head_dim + head_dim * 4 // quant_block_size
-    valid_tokens = 9
-    upper_bound_tokens = 13
-    block_size = 16
-    num_blocks = 2
-    sentinel = 123
-    device = "cuda"
-
-    k = torch.randn(valid_tokens, head_dim, dtype=torch.bfloat16, device=device)
-    kv_cache = torch.zeros(
-        num_blocks, block_size, cache_stride, dtype=torch.uint8, device=device
-    )
-    slot_mapping = torch.arange(valid_tokens, dtype=torch.int64, device=device)
-    ops.indexer_k_quant_and_cache(k, kv_cache, slot_mapping, quant_block_size, "ue8m0")
-
-    block_table = torch.arange(num_blocks, dtype=torch.int32, device=device).unsqueeze(
-        0
-    )
-    cu_seq_lens = torch.tensor([0, valid_tokens], dtype=torch.int32, device=device)
-    dst_k = torch.full(
-        (upper_bound_tokens, head_dim), sentinel, dtype=torch.uint8, device=device
-    )
-    num_scale_bytes = head_dim * 4 // quant_block_size
-    dst_scale = torch.full(
-        (upper_bound_tokens, num_scale_bytes),
-        sentinel,
-        dtype=torch.uint8,
-        device=device,
-    )
-
-    ops.cp_gather_indexer_k_quant_cache(
-        kv_cache, dst_k, dst_scale, block_table, cu_seq_lens
-    )
-    torch.accelerator.synchronize()
-
-    k_recovered = dst_k[:valid_tokens].view(torch.float8_e4m3fn).float() * dst_scale[
-        :valid_tokens
-    ].view(torch.float32)
-    diff = (k_recovered - k.float()).abs()
-    max_allowed = (16.0 * dst_scale[:valid_tokens].view(torch.float32).max()).item()
-    assert diff.max().item() <= max_allowed
-    assert torch.all(dst_k[valid_tokens:] == sentinel)
-    assert torch.all(dst_scale[valid_tokens:] == sentinel)
-
-
-# ── Test C: DeepseekV4 attention with values at different magnitudes ───────────
-
-
-def test_deepseek_v4_quant_magnitude_range():
-    """Test that quantization handles a range of magnitudes correctly."""
-
-    HEAD_DIM = 512
-    NOPE_DIM = 448
-    HEAD_BYTES = 584
-    block_size = 16
-    num_tokens = 4
-    num_blocks = 2
-    device = "cuda"
-
-    # Create inputs with varying magnitudes: small, medium, large
-    compressed_kv = torch.zeros(
-        num_tokens, HEAD_DIM, dtype=torch.bfloat16, device=device
-    )
-    compressed_kv[0] = 0.001  # very small
-    compressed_kv[1] = 1.0  # unit scale
-    compressed_kv[2] = 100.0  # large
-    compressed_kv[3] = torch.randn(HEAD_DIM, dtype=torch.bfloat16, device=device)
-
-    k_cache = torch.zeros(
-        num_blocks, block_size, HEAD_BYTES, dtype=torch.uint8, device=device
-    )
-    slot_mapping = torch.arange(num_tokens, dtype=torch.int64, device=device)
-
-    quantize_and_insert_k_cache(
-        compressed_kv, k_cache.view(num_blocks, -1), slot_mapping, block_size
-    )
-
-    out = torch.zeros(1, num_tokens, HEAD_DIM, dtype=torch.bfloat16, device=device)
-    seq_lens = torch.tensor([num_tokens], dtype=torch.int32, device=device)
-    block_table = torch.arange(num_blocks, dtype=torch.int32, device=device).unsqueeze(
-        0
-    )
-
-    dequantize_and_gather_k_cache(
-        out, k_cache, seq_lens, None, block_table, block_size, offset=0
-    )
-
-    recovered = out[0, :num_tokens]
-
-    # RoPE portion must be exact
-    rope_diff = (recovered[:, NOPE_DIM:] - compressed_kv[:, NOPE_DIM:]).abs().max()
-    assert rope_diff.item() == 0.0, f"RoPE diff {rope_diff.item()}"
-
-    # NoPE: relative error should be reasonable
-    for t in range(num_tokens):
-        orig = compressed_kv[t, :NOPE_DIM].float()
-        recv = recovered[t, :NOPE_DIM].float()
-        abs_diff = (recv - orig).abs().max().item()
-        magnitude = orig.abs().max().item()
-        if magnitude > 0.01:
-            rel_err = abs_diff / magnitude
-            assert rel_err < 0.15, (
-                f"Token {t}: rel_err={rel_err:.4f}, abs_diff={abs_diff:.6f}, "
-                f"magnitude={magnitude:.4f}"
-            )

tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py

@@ -1,359 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Standalone unit test for the horizontally-fused DeepseekV4-MLA kernel:
-
-  fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert
-    - Q side:  per-head RMSNorm (no weight) + GPT-J RoPE on last 64 dims
-    - KV side: GPT-J RoPE on last 64 + UE8M0 FP8 quant + paged cache insert
-
-We compare against:
-  - PyTorch reference for RMSNorm + GPT-J RoPE on Q
-  - Existing Triton `quantize_and_insert_k_cache` + round-trip via
-    `dequantize_and_gather_k_cache` for KV
-
-The kernel is imported via
-`torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert`.
-"""
-
-import pytest
-import torch
-
-from vllm.v1.attention.ops.deepseek_v4_ops import (
-    dequantize_and_gather_k_cache,
-    quantize_and_insert_k_cache,
-)
-
-# ── Constants matching the kernel ────────────────────────────────────────────
-HEAD_DIM = 512
-ROPE_DIM = 64
-NOPE_DIM = HEAD_DIM - ROPE_DIM  # 448
-QUANT_BLOCK = 64
-FP8_MAX = 448.0
-HEAD_BYTES = NOPE_DIM + ROPE_DIM * 2 + 8  # 448 + 128 + 8 = 584
-
-
-# ── PyTorch reference implementations ────────────────────────────────────────
-
-
-def make_cos_sin_cache(max_pos: int, rope_dim: int, dtype, device):
-    """Build a cos||sin cache matching DeepseekV4ScalingRotaryEmbedding layout.
-    cos_sin_cache[pos, :rope_dim/2] = cos(theta), [rope_dim/2:] = sin(theta).
-    """
-    base = 10000.0
-    inv_freq = 1.0 / (
-        base
-        ** (torch.arange(0, rope_dim, 2, dtype=torch.float32, device=device) / rope_dim)
-    )
-    t = torch.arange(max_pos, dtype=torch.float32, device=device)
-    freqs = torch.einsum("i,j -> ij", t, inv_freq)  # [max_pos, rope_dim/2]
-    cache = torch.cat((freqs.cos(), freqs.sin()), dim=-1)  # [max_pos, rope_dim]
-    return cache.to(dtype)
-
-
-def apply_rope_gptj_last_k(
-    x: torch.Tensor, positions: torch.Tensor, cos_sin_cache: torch.Tensor
-) -> torch.Tensor:
-    """GPT-J-style (interleaved-pair) RoPE on the LAST rope_dim elements.
-
-    x: [..., head_dim] float32
-    positions: [num_tokens] int64 (positions[i] corresponds to x[i, ...])
-    cos_sin_cache: [max_pos, rope_dim] float (cos|sin layout)
-
-    Returns rotated x (same shape/dtype).
-    """
-    rope_dim = cos_sin_cache.shape[-1]
-    half = rope_dim // 2
-    head_dim = x.shape[-1]
-    nope_dim = head_dim - rope_dim
-
-    # Gather cos/sin for each token position: [num_tokens, rope_dim]
-    cs = cos_sin_cache[positions].to(torch.float32)  # [N, rope_dim]
-    cos = cs[..., :half]  # [N, half]
-    sin = cs[..., half:]  # [N, half]
-
-    # Reshape leading dims so we can broadcast: x shape [..., head_dim].
-    # Bring token dim to front; assume x is [num_tokens, ..., head_dim].
-    # We rely on positions being per-token and all other dims sharing the same pos.
-    rope = x[..., nope_dim:].float()  # [..., rope_dim]
-    # Make rope pairs: reshape last dim to [half, 2]
-    shape = rope.shape
-    rope = rope.reshape(*shape[:-1], half, 2)
-    even = rope[..., 0]  # [..., half]
-    odd = rope[..., 1]
-
-    # Broadcast cos/sin over any heads dim in between.  cos/sin are [N, half].
-    # Add singleton dims for intermediate axes.
-    for _ in range(rope.ndim - 3):
-        cos = cos.unsqueeze(1)
-        sin = sin.unsqueeze(1)
-
-    new_even = even * cos - odd * sin
-    new_odd = even * sin + odd * cos
-    rope_rotated = torch.stack((new_even, new_odd), dim=-1).reshape(shape)
-
-    out = x.clone().float()
-    out[..., nope_dim:] = rope_rotated
-    return out.to(x.dtype)
-
-
-def rmsnorm_no_weight(x: torch.Tensor, eps: float) -> torch.Tensor:
-    """RMSNorm with no learnable weight, matching
-    `RMSNorm(head_dim, has_weight=False)`."""
-    orig_dtype = x.dtype
-    xf = x.float()
-    variance = xf.pow(2).mean(dim=-1, keepdim=True)
-    return (xf * torch.rsqrt(variance + eps)).to(orig_dtype)
-
-
-# ── Dispatch to the CUDA op (skip test cleanly if it isn't built in) ─────────
-
-
-def _op_available() -> bool:
-    return hasattr(torch.ops._C, "fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert")
-
-
-pytestmark = pytest.mark.skipif(
-    not torch.cuda.is_available() or not _op_available(),
-    reason="CUDA not available or fused DeepseekV4 op not built in",
-)
-
-
-def _call_fused(q, kv, k_cache, slot_mapping, positions, cos_sin_cache, eps, bs):
-    torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
-        q, kv, k_cache, slot_mapping, positions, cos_sin_cache, eps, bs
-    )
-
-
-# ── Test 1: Q path numerical parity ──────────────────────────────────────────
-
-
-@pytest.mark.parametrize("num_tokens", [1, 4, 17, 64])
-@pytest.mark.parametrize("n_heads", [8, 64])
-def test_q_path_matches_reference(num_tokens: int, n_heads: int):
-    torch.manual_seed(0)
-    device = "cuda"
-    dtype = torch.bfloat16
-    eps = 1e-6
-    max_pos = 4096
-
-    q = torch.randn(num_tokens, n_heads, HEAD_DIM, dtype=dtype, device=device)
-    positions = torch.arange(num_tokens, dtype=torch.int64, device=device)
-    cos_sin_cache = make_cos_sin_cache(max_pos, ROPE_DIM, torch.float32, device)
-
-    # Reference: RMSNorm (no weight) per head, then GPT-J RoPE on last 64.
-    q_ref = rmsnorm_no_weight(q, eps)
-    q_ref = apply_rope_gptj_last_k(q_ref, positions, cos_sin_cache)
-
-    # Fused call with dummy KV tensors (KV branch will write slot_mapping=-1 → noop).
-    num_blocks = 2
-    bs = 16
-    kv = torch.zeros(num_tokens, HEAD_DIM, dtype=dtype, device=device)
-    k_cache = torch.zeros(
-        num_blocks, bs, HEAD_BYTES, dtype=torch.uint8, device=device
-    ).view(num_blocks, -1)
-    slot_mapping = torch.full((num_tokens,), -1, dtype=torch.int64, device=device)
-    q_fused = q.clone()
-    _call_fused(q_fused, kv, k_cache, slot_mapping, positions, cos_sin_cache, eps, bs)
-
-    torch.testing.assert_close(q_fused, q_ref, rtol=1e-2, atol=1e-2)
-
-
-# ── Test 2: KV path round-trip byte/value parity ─────────────────────────────
-
-
-def _ue8m0_per_block_scales(kv_roped_nope_f32: torch.Tensor, qblock: int):
-    """Return per-token per-block max scale (used to bound FP8 error)."""
-    n_tok, nope = kv_roped_nope_f32.shape
-    n_blocks = nope // qblock
-    blocks = kv_roped_nope_f32.view(n_tok, n_blocks, qblock)
-    absmax = blocks.abs().amax(dim=-1).clamp(min=1e-4)
-    raw = absmax / FP8_MAX
-    exponent = torch.ceil(torch.log2(raw))
-    return torch.pow(2.0, exponent)  # [n_tok, n_blocks]
-
-
-@pytest.mark.parametrize("num_tokens", [1, 4, 17, 64])
-@pytest.mark.parametrize("block_size", [16, 64])
-def test_kv_path_matches_reference(num_tokens: int, block_size: int):
-    torch.manual_seed(1)
-    device = "cuda"
-    dtype = torch.bfloat16
-    eps = 1e-6
-    max_pos = 4096
-
-    kv = torch.randn(num_tokens, HEAD_DIM, dtype=dtype, device=device)
-    positions = torch.arange(num_tokens, dtype=torch.int64, device=device)
-    cos_sin_cache = make_cos_sin_cache(max_pos, ROPE_DIM, torch.float32, device)
-
-    num_blocks = (num_tokens + block_size - 1) // block_size + 1
-    slot_mapping = torch.arange(num_tokens, dtype=torch.int64, device=device)
-
-    # ── Reference path: RoPE on kv, then existing Triton quant+insert ──────
-    kv_ref = apply_rope_gptj_last_k(kv, positions, cos_sin_cache)
-    k_cache_ref = torch.zeros(
-        num_blocks, block_size * HEAD_BYTES, dtype=torch.uint8, device=device
-    )
-    quantize_and_insert_k_cache(
-        kv_ref, k_cache_ref, slot_mapping, block_size=block_size
-    )
-
-    # ── Fused path (dummy q, single head) ──────────────────────────────────
-    k_cache_fused = torch.zeros_like(k_cache_ref)
-    q_dummy = torch.zeros(num_tokens, 1, HEAD_DIM, dtype=dtype, device=device)
-    _call_fused(
-        q_dummy,
-        kv,
-        k_cache_fused,
-        slot_mapping,
-        positions,
-        cos_sin_cache,
-        eps,
-        block_size,
-    )
-
-    # ── Round-trip compare via dequant+gather ──────────────────────────────
-    def _dequant(k_cache_2d):
-        num_reqs = 1
-        max_blocks = num_blocks
-        out = torch.zeros(
-            num_reqs, num_tokens, HEAD_DIM, dtype=torch.bfloat16, device=device
-        )
-        seq_lens = torch.tensor([num_tokens], dtype=torch.int32, device=device)
-        block_table = torch.arange(
-            max_blocks, dtype=torch.int32, device=device
-        ).unsqueeze(0)
-        # gather_lens arg is None (use seq_lens)
-        k_cache_3d = k_cache_2d.view(num_blocks, block_size, HEAD_BYTES)
-        dequantize_and_gather_k_cache(
-            out, k_cache_3d, seq_lens, None, block_table, block_size, offset=0
-        )
-        return out[0, :num_tokens]
-
-    recovered_ref = _dequant(k_cache_ref)
-    recovered_fused = _dequant(k_cache_fused)
-
-    # NoPE: per-block UE8M0 FP8 error bound (half-ULP at max = 16 * scale).
-    scales = _ue8m0_per_block_scales(kv_ref[:, :NOPE_DIM].float(), QUANT_BLOCK)
-    for t in range(num_tokens):
-        max_allowed = 16.0 * scales[t].max().item()
-        diff_ref = (
-            (recovered_ref[t, :NOPE_DIM] - kv_ref[t, :NOPE_DIM]).abs().max().item()
-        )
-        diff_fused = (
-            (recovered_fused[t, :NOPE_DIM] - kv_ref[t, :NOPE_DIM]).abs().max().item()
-        )
-        assert diff_ref <= max_allowed, (
-            f"ref NoPE token {t} diff {diff_ref} > {max_allowed}"
-        )
-        assert diff_fused <= max_allowed, (
-            f"fused NoPE token {t} diff {diff_fused} > {max_allowed}"
-        )
-
-    # RoPE region: bf16 stored exactly → zero diff.
-    rope_diff = (recovered_fused[:, NOPE_DIM:] - kv_ref[:, NOPE_DIM:]).abs().max()
-    assert rope_diff.item() == 0.0, f"RoPE portion not exact: {rope_diff.item()}"
-
-    # Exact byte equality of the two cache buffers — strong parity.
-    torch.testing.assert_close(k_cache_fused, k_cache_ref, rtol=0, atol=0)
-
-
-# ── Test 2b: DP padding (slot_mapping shorter than q/kv) ─────────────────────
-
-
-@pytest.mark.parametrize("num_tokens", [4, 17])
-@pytest.mark.parametrize("pad", [1, 5])
-@pytest.mark.parametrize("block_size", [16, 64])
-def test_kv_path_with_dp_padding(num_tokens: int, pad: int, block_size: int):
-    """slot_mapping.size(0) < q.size(0): the kernel must skip padded
-    tokens in the KV branch while still running Q-norm+RoPE on all rows."""
-    torch.manual_seed(3)
-    device = "cuda"
-    dtype = torch.bfloat16
-    eps = 1e-6
-    max_pos = 4096
-    total = num_tokens + pad
-
-    kv = torch.randn(total, HEAD_DIM, dtype=dtype, device=device)
-    positions = torch.arange(total, dtype=torch.int64, device=device)
-    cos_sin_cache = make_cos_sin_cache(max_pos, ROPE_DIM, torch.float32, device)
-
-    num_blocks = (num_tokens + block_size - 1) // block_size + 1
-    slot_mapping = torch.arange(num_tokens, dtype=torch.int64, device=device)
-
-    # Reference: only the first num_tokens kv rows get inserted.
-    kv_ref = apply_rope_gptj_last_k(
-        kv[:num_tokens], positions[:num_tokens], cos_sin_cache
-    )
-    k_cache_ref = torch.zeros(
-        num_blocks, block_size * HEAD_BYTES, dtype=torch.uint8, device=device
-    )
-    quantize_and_insert_k_cache(
-        kv_ref, k_cache_ref, slot_mapping, block_size=block_size
-    )
-
-    # Fused: pass full-sized q/kv/positions, shorter slot_mapping.
-    q_dummy = torch.zeros(total, 1, HEAD_DIM, dtype=dtype, device=device)
-    k_cache_fused = torch.zeros_like(k_cache_ref)
-    _call_fused(
-        q_dummy,
-        kv,
-        k_cache_fused,
-        slot_mapping,
-        positions,
-        cos_sin_cache,
-        eps,
-        block_size,
-    )
-
-    torch.testing.assert_close(k_cache_fused, k_cache_ref, rtol=0, atol=0)
-
-
-# ── Test 3: combined single-call Q + KV parity ───────────────────────────────
-
-
-@pytest.mark.parametrize("num_tokens", [1, 4, 17])
-@pytest.mark.parametrize("n_heads", [8, 64])
-@pytest.mark.parametrize("block_size", [16, 64])
-def test_combined_q_and_kv(num_tokens: int, n_heads: int, block_size: int):
-    torch.manual_seed(2)
-    device = "cuda"
-    dtype = torch.bfloat16
-    eps = 1e-6
-    max_pos = 4096
-
-    q = torch.randn(num_tokens, n_heads, HEAD_DIM, dtype=dtype, device=device)
-    kv = torch.randn(num_tokens, HEAD_DIM, dtype=dtype, device=device)
-    positions = torch.arange(num_tokens, dtype=torch.int64, device=device)
-    cos_sin_cache = make_cos_sin_cache(max_pos, ROPE_DIM, torch.float32, device)
-
-    num_blocks = (num_tokens + block_size - 1) // block_size + 1
-    slot_mapping = torch.arange(num_tokens, dtype=torch.int64, device=device)
-
-    # Reference.
-    q_ref = rmsnorm_no_weight(q, eps)
-    q_ref = apply_rope_gptj_last_k(q_ref, positions, cos_sin_cache)
-    kv_ref = apply_rope_gptj_last_k(kv, positions, cos_sin_cache)
-    k_cache_ref = torch.zeros(
-        num_blocks, block_size * HEAD_BYTES, dtype=torch.uint8, device=device
-    )
-    quantize_and_insert_k_cache(
-        kv_ref, k_cache_ref, slot_mapping, block_size=block_size
-    )
-
-    # Fused single call.
-    q_fused = q.clone()
-    k_cache_fused = torch.zeros_like(k_cache_ref)
-    _call_fused(
-        q_fused,
-        kv,
-        k_cache_fused,
-        slot_mapping,
-        positions,
-        cos_sin_cache,
-        eps,
-        block_size,
-    )
-
-    torch.testing.assert_close(q_fused, q_ref, rtol=1e-2, atol=1e-2)
-    torch.testing.assert_close(k_cache_fused, k_cache_ref, rtol=0, atol=0)

tests/kernels/test_fused_indexer_q_rope_quant.py

@@ -1,98 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Unit test for fused_indexer_q_rope_quant.
-
-Compares the fused Triton kernel against the unfused reference flow used by
-the DeepseekV4 indexer in model_tracking:
-    q_rot = ops.rotary_embedding(positions, q, None, head_dim, cos_sin_cache,
-                                 is_neox_style=False,
-                                 rope_dim_offset=head_dim - rope_dim)
-    q_fp8, q_scale = per_token_group_quant_fp8(q_rot, head_dim, use_ue8m0=True)
-    weights_out = weights * q_scale * softmax_scale * head_scale
-
-Expects bit-exact equality on both q_fp8 and weights_out.
-"""
-
-import pytest
-import torch
-
-from vllm import _custom_ops as ops
-from vllm.model_executor.layers.quantization.utils.fp8_utils import (
-    per_token_group_quant_fp8,
-)
-from vllm.v1.attention.ops.deepseek_v4_ops.fused_indexer_q import (
-    fused_indexer_q_rope_quant,
-)
-
-HEAD_DIM = 128
-ROPE_DIM = 64
-N_HEAD = 64
-MAX_POS = 4096
-
-
-def _reference(
-    positions: torch.Tensor,
-    q: torch.Tensor,
-    cos_sin_cache: torch.Tensor,
-    weights: torch.Tensor,
-    softmax_scale: float,
-    head_scale: float,
-) -> tuple[torch.Tensor, torch.Tensor]:
-    q_rot = q.clone()
-    ops.rotary_embedding(
-        positions,
-        q_rot,
-        None,
-        HEAD_DIM,
-        cos_sin_cache,
-        False,  # is_neox_style=False → GPT-J interleaved
-        HEAD_DIM - ROPE_DIM,  # rope_dim_offset → rotate the tail
-        False,
-    )
-    q_fp8, q_scale = per_token_group_quant_fp8(
-        q_rot.view(-1, HEAD_DIM).contiguous(),
-        HEAD_DIM,
-        use_ue8m0=True,
-    )
-    q_fp8 = q_fp8.view(-1, N_HEAD, HEAD_DIM)
-    q_scale = q_scale.view(-1, N_HEAD)
-
-    weights_out = weights.to(torch.float32) * q_scale * softmax_scale * head_scale
-    return q_fp8, weights_out
-
-
-@pytest.mark.parametrize("num_tokens", [1, 7, 32, 257])
-@pytest.mark.parametrize("cache_dtype", [torch.float32, torch.bfloat16])
-@torch.inference_mode()
-def test_fused_indexer_q_rope_quant_matches_unfused(num_tokens, cache_dtype):
-    device = "cuda"
-    torch.manual_seed(0)
-
-    q = torch.randn(num_tokens, N_HEAD, HEAD_DIM, dtype=torch.bfloat16, device=device)
-    positions = torch.randint(
-        0, MAX_POS, (num_tokens,), dtype=torch.int64, device=device
-    )
-    cos_sin_cache = torch.randn(MAX_POS, ROPE_DIM, dtype=cache_dtype, device=device)
-    weights = torch.randn(num_tokens, N_HEAD, dtype=torch.bfloat16, device=device)
-    softmax_scale = HEAD_DIM**-0.5
-    head_scale = N_HEAD**-0.5
-
-    q_fp8_ref, weights_ref = _reference(
-        positions, q, cos_sin_cache, weights, softmax_scale, head_scale
-    )
-    q_fp8_fused, weights_fused = fused_indexer_q_rope_quant(
-        positions, q.clone(), cos_sin_cache, weights, softmax_scale, head_scale
-    )
-
-    # fp8 tensors aren't directly comparable via torch.equal — reinterpret as int8.
-    ref_bits = q_fp8_ref.view(torch.int8)
-    fused_bits = q_fp8_fused.view(torch.int8)
-    assert torch.equal(ref_bits, fused_bits), (
-        f"q_fp8 mismatch: "
-        f"{(ref_bits != fused_bits).sum().item()} / {ref_bits.numel()} bytes differ"
-    )
-
-    assert torch.equal(weights_ref, weights_fused), (
-        f"weights mismatch: max abs diff "
-        f"{(weights_ref - weights_fused).abs().max().item()}"
-    )

tests/kernels/test_fused_inv_rope_fp8_quant.py

@@ -1,908 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Unit tests for the fused inverse RoPE + block-scaled FP8 quantization kernel.
-
-Tests compare the fused kernel against a reference implementation built from
-the existing separate operations (inverse RoPE via rotate_neox + FP8 quant
-via per_token_group_quant_fp8).
-
-The reference faithfully reproduces the exact flow in deepseek_v4_attention.py:295-310:
-  1. Apply inverse RoPE (NeoX style, last rope_dim=64 dims of each head)
-  2. Reshape [T, H, head_dim] -> [T, G, D]
-  3. Transpose+flatten to [G*T, D], quantize, reshape back
-  4. Return o_fp8 and o_scale with strides (D, T*D, 1) and (S, T*S, 1)
-     (non-contiguous [T, G, ...] view backed by contiguous [G, T, ...] memory)
-
-Usage:
-    pytest tests/kernels/test_fused_inv_rope_fp8_quant.py -v
-"""
-
-import pytest
-import torch
-
-from vllm.v1.attention.ops.deepseek_v4_ops import fused_inv_rope_fp8_quant
-
-# -- Default dimensions matching DeepSeek V3/V4 --------------------------
-HEAD_DIM = 512
-NOPE_DIM = 448
-ROPE_DIM = 64
-QUANT_GROUP_SIZE = 128
-FP8_MAX = 448.0  # torch.finfo(torch.float8_e4m3fn).max
-FP8_DTYPE = torch.float8_e4m3fn
-EPS = 1e-10
-
-
-# =========================================================================
-# Helpers
-# =========================================================================
-
-
-def assert_dequant_close(
-    fp8_a: torch.Tensor,
-    scale_a: torch.Tensor,
-    fp8_b: torch.Tensor,
-    scale_b: torch.Tensor,
-    msg: str = "",
-):
-    """Compare two FP8-quantized tensors via their dequantized values.
-
-    Uses cosine-similarity-based diff (same as deep_gemm calc_diff).
-    Both fused and reference paths rotate in fp32 using an fp32
-    cos_sin_cache, so differences are only fp32 ordering ULPs that can
-    occasionally shift FP8 values at quantization boundaries.
-    """
-    S = scale_a.shape[-1]
-    shape = fp8_a.shape
-
-    dq_a = fp8_a.float() * scale_a.unsqueeze(-1).expand(
-        *shape[:-1], S, QUANT_GROUP_SIZE
-    ).reshape(shape)
-    dq_b = fp8_b.float() * scale_b.unsqueeze(-1).expand(
-        *shape[:-1], S, QUANT_GROUP_SIZE
-    ).reshape(shape)
-
-    # Cosine diff: 1 - cos_sim (0 = identical, higher = worse)
-    dq_a_flat = dq_a.flatten().float()
-    dq_b_flat = dq_b.flatten().float()
-    cos_sim = torch.nn.functional.cosine_similarity(
-        dq_a_flat.unsqueeze(0), dq_b_flat.unsqueeze(0)
-    ).item()
-    diff = 1.0 - cos_sim
-
-    assert diff < 1e-4, f"Dequant diff too large: {diff:.8f} (expected < 1e-4). {msg}"
-
-
-def rotate_gptj(x: torch.Tensor) -> torch.Tensor:
-    """GPT-J style rotation: interleaved pairs, negate-swap.
-
-    Matches vllm/model_executor/layers/rotary_embedding/common.py:23-27.
-    DeepseekV4 uses is_neox_style=False, so this is the correct rotation.
-    """
-    x1 = x[..., ::2]
-    x2 = x[..., 1::2]
-    x = torch.stack((-x2, x1), dim=-1)
-    return x.flatten(-2)
-
-
-def make_cos_sin_cache(
-    max_pos: int,
-    rope_dim: int = ROPE_DIM,
-    dtype: torch.dtype = torch.float32,
-    device: str = "cuda",
-) -> torch.Tensor:
-    """Create a synthetic cos_sin_cache matching the layout used by
-    DeepseekV4ScalingRotaryEmbedding._compute_cos_sin_cache.
-
-    Shape: [max_pos, rope_dim] where first half is cos, second half is sin.
-    The fused kernel requires fp32; callers can override dtype if passing
-    the cache into the bf16-only paths.
-    """
-    half = rope_dim // 2
-    # Use random but bounded frequencies so cos/sin are well-behaved
-    inv_freq = 1.0 / (
-        10000.0 ** (torch.arange(0, half, device=device, dtype=torch.float32) / half)
-    )
-    t = torch.arange(max_pos, device=device, dtype=torch.float32)
-    freqs = torch.outer(t, inv_freq)  # [max_pos, half]
-    cos = freqs.cos()
-    sin = freqs.sin()
-    cache = torch.cat((cos, sin), dim=-1)  # [max_pos, rope_dim]
-    return cache.to(dtype)
-
-
-def reference_inv_rope(
-    o: torch.Tensor,
-    positions: torch.Tensor,
-    cos_sin_cache: torch.Tensor,
-    nope_dim: int = NOPE_DIM,
-    rope_dim: int = ROPE_DIM,
-) -> torch.Tensor:
-    """Apply inverse RoPE to the last rope_dim dimensions of each head.
-
-    Matches the GPT-J inverse rotation in pos_encoding_kernels.cu, which
-    promotes the cache to fp32 and performs the rotation in fp32. The
-    result is cast back to the input dtype.
-
-    Args:
-        o: [T, H, head_dim] bf16
-        positions: [T] int64
-        cos_sin_cache: [max_pos, rope_dim] fp32
-
-    Returns:
-        o with inverse RoPE applied on the rope portion (bf16).
-    """
-    assert cos_sin_cache.dtype == torch.float32
-    cos_sin = cos_sin_cache[positions]  # [T, rope_dim] fp32
-    half = rope_dim // 2
-    cos = cos_sin[:, :half]
-    sin = cos_sin[:, half:]
-
-    # GPT-J style: repeat_interleave (not repeat) to match interleaved pairs
-    cos = cos.repeat_interleave(2, dim=-1).unsqueeze(1)
-    sin = sin.repeat_interleave(2, dim=-1).unsqueeze(1)
-    sin = -sin  # inverse
-
-    o_pass = o[..., :nope_dim]
-    o_rot_f32 = o[..., nope_dim:].float()
-    o_rot_f32 = o_rot_f32 * cos + rotate_gptj(o_rot_f32) * sin
-    o_rot = o_rot_f32.to(o.dtype)
-
-    return torch.cat([o_pass, o_rot], dim=-1)
-
-
-def _ref_ue8m0_quant_block(x_f32: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
-    """Per-block UE8M0 FP8 quantization in pure float32.
-
-    Matches the Triton kernel logic exactly:
-      absmax -> 2^ceil(log2(absmax / fp8_max)) -> clamp(x / scale) -> fp8
-
-    Args:
-        x_f32: [..., quant_group_size] float32 — one or more 128-element blocks.
-
-    Returns:
-        x_fp8: same shape, float8_e4m3fn
-        scales: [...] float32, one scale per block
-    """
-    assert x_f32.shape[-1] == QUANT_GROUP_SIZE
-    absmax = x_f32.abs().amax(dim=-1, keepdim=True).clamp(min=EPS)
-    scale_raw = absmax * (1.0 / FP8_MAX)
-    scale = torch.exp2(torch.ceil(torch.log2(scale_raw)))
-    x_scaled = (x_f32 / scale).clamp(-FP8_MAX, FP8_MAX)
-    x_fp8 = x_scaled.to(FP8_DTYPE)
-    return x_fp8, scale.squeeze(-1)
-
-
-def reference_inv_rope_fp8_quant(
-    o: torch.Tensor,
-    positions: torch.Tensor,
-    cos_sin_cache: torch.Tensor,
-    n_groups: int,
-    heads_per_group: int,
-    nope_dim: int = NOPE_DIM,
-    rope_dim: int = ROPE_DIM,
-    quant_group_size: int = QUANT_GROUP_SIZE,
-) -> tuple[torch.Tensor, torch.Tensor]:
-    """Full reference: inverse RoPE in fp32 + UE8M0 FP8 quant in fp32.
-
-    Mimics the Triton kernel's precision path exactly:
-      Load bf16 -> cast to fp32 -> apply inverse RoPE with fp32 cos/sin ->
-      UE8M0 quant in fp32 -> write fp8 + scale
-
-    Returns:
-        o_fp8: [T, G, D] FP8 with strides (D, T*D, 1)
-        o_scale: [T, G, S] FP32 with strides (S, T*S, 1)
-    """
-    assert cos_sin_cache.dtype == torch.float32
-    T, _H, head_dim = o.shape
-    d = heads_per_group * head_dim
-    S = d // quant_group_size
-    half_rope = rope_dim // 2
-    chunks_per_head = head_dim // quant_group_size
-
-    # Reshape [T, H, head_dim] -> [T, G, heads_per_group, head_dim]
-    o_4d = o.view(T, n_groups, heads_per_group, head_dim)
-
-    # Lookup cos/sin directly in fp32
-    cos_sin = cos_sin_cache[positions]  # [T, rope_dim] fp32
-    cos = cos_sin[:, :half_rope]  # [T, half_rope] fp32
-    sin = cos_sin[:, half_rope:]  # [T, half_rope] fp32
-
-    # Allocate outputs in [G, T, ...] contiguous layout
-    fp8_buf = torch.empty(n_groups, T, d, dtype=FP8_DTYPE, device=o.device)
-    scale_buf = torch.empty(n_groups, T, S, dtype=torch.float32, device=o.device)
-
-    # Process each quant block, matching the Triton kernel's per-program logic
-    for g in range(n_groups):
-        for qb in range(S):
-            head_in_group = qb // chunks_per_head
-            chunk_in_head = qb % chunks_per_head
-            offset = chunk_in_head * quant_group_size
-
-            # Load 128 bf16 elements and promote to fp32 for rotation+quant
-            block = o_4d[:, g, head_in_group, offset : offset + quant_group_size]
-            x = block.float()
-
-            # Apply inverse RoPE in fp32 if this is the last chunk
-            # GPT-J style: interleaved pairs (even=x, odd=y)
-            if chunk_in_head == chunks_per_head - 1:
-                rope_start = nope_dim % quant_group_size  # 64
-                rope_region = x[:, rope_start:].clone()
-                x_vals = rope_region[:, ::2]
-                y_vals = rope_region[:, 1::2]
-                x_new = x_vals * cos + y_vals * sin
-                y_new = y_vals * cos - x_vals * sin
-                x = x.clone()
-                x[:, rope_start::2] = x_new
-                x[:, rope_start + 1 :: 2] = y_new
-
-            # UE8M0 quant in fp32
-            x_fp8, scale = _ref_ue8m0_quant_block(x)
-
-            # Write to [G, T, D] contiguous memory
-            fp8_buf[g, :, qb * quant_group_size : (qb + 1) * quant_group_size] = x_fp8
-            scale_buf[g, :, qb] = scale
-
-    # Return transposed views
-    return fp8_buf.transpose(0, 1), scale_buf.transpose(0, 1)
-
-
-# =========================================================================
-# Tests
-# =========================================================================
-
-
-@pytest.mark.parametrize("num_tokens", [1, 7, 32, 128])
-@pytest.mark.parametrize(
-    "num_heads,n_groups",
-    [(64, 8), (32, 4), (128, 8)],
-    ids=["H64_G8", "H32_G4", "H128_G8"],
-)
-@pytest.mark.parametrize("seed", [0, 42])
-@torch.inference_mode()
-def test_correctness(num_tokens, num_heads, n_groups, seed):
-    """Compare fused kernel against reference for FP8 values and scales."""
-    torch.manual_seed(seed)
-
-    heads_per_group = num_heads // n_groups
-    max_pos = 4096
-    device = "cuda"
-
-    # Create inputs
-    o = torch.randn(
-        num_tokens, num_heads, HEAD_DIM, device=device, dtype=torch.bfloat16
-    )
-    positions = torch.randint(
-        0, max_pos, (num_tokens,), device=device, dtype=torch.long
-    )
-    cos_sin_cache = make_cos_sin_cache(
-        max_pos, ROPE_DIM, dtype=torch.float32, device=device
-    )
-
-    # Reference
-    ref_fp8, ref_scale = reference_inv_rope_fp8_quant(
-        o.clone(),
-        positions,
-        cos_sin_cache,
-        n_groups,
-        heads_per_group,
-    )
-
-    # Fused kernel
-    fused_fp8, fused_scale = fused_inv_rope_fp8_quant(
-        o.clone(),
-        positions,
-        cos_sin_cache,
-        n_groups,
-        heads_per_group,
-    )
-
-    # Check shapes
-    d = heads_per_group * HEAD_DIM
-    S = d // QUANT_GROUP_SIZE
-    assert ref_fp8.shape == (num_tokens, n_groups, d)
-    assert fused_fp8.shape == (num_tokens, n_groups, d)
-    assert ref_scale.shape == (num_tokens, n_groups, S)
-    assert fused_scale.shape == (num_tokens, n_groups, S)
-
-    # Scales: exact match (both use identical UE8M0 algorithm)
-    # Scales may differ by one UE8M0 step (factor of 2) if fp32 rotation
-    # ordering shifts absmax across a power-of-2 boundary. Check ratio is
-    # close to 1.
-    scale_ratio = fused_scale / ref_scale.clamp(min=1e-30)
-    assert scale_ratio.max() <= 2.0 and scale_ratio.min() >= 0.5, (
-        f"Scale ratio out of [0.5, 2]: min={scale_ratio.min():.4f} "
-        f"max={scale_ratio.max():.4f}"
-    )
-
-    # Compare via dequant (Triton vs PyTorch fp32 may differ by ULPs)
-    assert_dequant_close(ref_fp8, ref_scale, fused_fp8, fused_scale)
-
-
-@pytest.mark.parametrize("num_tokens", [1, 7, 32, 128])
-@pytest.mark.parametrize(
-    "num_heads,n_groups",
-    [(64, 8), (128, 8)],
-    ids=["H64_G8", "H128_G8"],
-)
-@torch.inference_mode()
-def test_output_strides(num_tokens, num_heads, n_groups):
-    """Verify fused output layout:
-    - FP8: logical [T, G, D] backed by contiguous [G, T, D].
-    - Scale: MN-major TMA-aligned (column-major: T-stride=1).
-    """
-
-    heads_per_group = num_heads // n_groups
-    max_pos = 4096
-    device = "cuda"
-
-    o = torch.randn(
-        num_tokens, num_heads, HEAD_DIM, device=device, dtype=torch.bfloat16
-    )
-    positions = torch.randint(
-        0, max_pos, (num_tokens,), device=device, dtype=torch.long
-    )
-    cos_sin_cache = make_cos_sin_cache(max_pos, device=device)
-
-    fused_fp8, fused_scale = fused_inv_rope_fp8_quant(
-        o.clone(),
-        positions,
-        cos_sin_cache,
-        n_groups,
-        heads_per_group,
-    )
-
-    # FP8: logical [T, G, D] backed by [G, T, D] row-major
-    d = heads_per_group * HEAD_DIM
-    expected_fp8_stride = (d, num_tokens * d, 1)
-    assert fused_fp8.stride() == expected_fp8_stride, (
-        f"FP8 stride mismatch: got {fused_fp8.stride()}, expected {expected_fp8_stride}"
-    )
-
-    # Scale: MN-major TMA-aligned layout. After fp8_einsum permutes
-    # [T,G,S] -> [G,T,S], T-dim should have stride 1.
-    # Our output is [T,G,S] = transpose of [G,T,S].
-    # So fused_scale.permute(1,0,2) should have T-stride=1.
-    perm = fused_scale.permute(1, 0, 2)  # [G, T, S]
-    assert perm.stride(1) == 1 or num_tokens == 1, (
-        f"Scale T-stride (after permute to [G,T,S]) should be 1, got {perm.stride(1)}"
-    )
-
-
-@pytest.mark.parametrize("num_tokens", [1, 7, 32, 128])
-@torch.inference_mode()
-def test_per_group_contiguity(num_tokens):
-    """FP8 per-group slices must be contiguous. Scale per-group slices
-    are column-major (T-stride=1) — not row-major contiguous, which is
-    correct for TMA loads."""
-    num_heads, n_groups = 64, 8
-    heads_per_group = num_heads // n_groups
-    max_pos = 4096
-    device = "cuda"
-
-    o = torch.randn(
-        num_tokens, num_heads, HEAD_DIM, device=device, dtype=torch.bfloat16
-    )
-    positions = torch.randint(
-        0, max_pos, (num_tokens,), device=device, dtype=torch.long
-    )
-    cos_sin_cache = make_cos_sin_cache(max_pos, device=device)
-
-    fused_fp8, fused_scale = fused_inv_rope_fp8_quant(
-        o.clone(),
-        positions,
-        cos_sin_cache,
-        n_groups,
-        heads_per_group,
-    )
-
-    for g in range(n_groups):
-        fp8_slice = fused_fp8[:, g, :]
-        assert fp8_slice.is_contiguous(), (
-            f"o_fp8[:, {g}, :] is not contiguous: "
-            f"shape={list(fp8_slice.shape)}, stride={list(fp8_slice.stride())}"
-        )
-
-
-@torch.inference_mode()
-def test_scales_are_power_of_two():
-    """Verify all scales are exact powers of 2 (UE8M0 property)."""
-    num_tokens, num_heads, n_groups = 32, 64, 8
-    heads_per_group = num_heads // n_groups
-    max_pos = 4096
-    device = "cuda"
-
-    o = torch.randn(
-        num_tokens, num_heads, HEAD_DIM, device=device, dtype=torch.bfloat16
-    )
-    positions = torch.randint(
-        0, max_pos, (num_tokens,), device=device, dtype=torch.long
-    )
-    cos_sin_cache = make_cos_sin_cache(max_pos, device=device)
-
-    _, fused_scale = fused_inv_rope_fp8_quant(
-        o.clone(),
-        positions,
-        cos_sin_cache,
-        n_groups,
-        heads_per_group,
-    )
-
-    # log2 of a power-of-two is an exact integer
-    log2_scales = torch.log2(fused_scale)
-    residual = (log2_scales - log2_scales.round()).abs()
-    assert residual.max() < 1e-5, (
-        f"Not all scales are powers of 2: max log2 residual = {residual.max().item()}"
-    )
-
-
-@torch.inference_mode()
-def test_nope_dims_unchanged():
-    """Nope dimensions (first 448 per head) should only be quantized,
-    not rotated. Verify by dequantizing and comparing against
-    quantize-only reference (no RoPE)."""
-    num_tokens, num_heads, n_groups = 16, 64, 8
-    heads_per_group = num_heads // n_groups
-    max_pos = 4096
-    device = "cuda"
-    torch.manual_seed(0)
-
-    o = torch.randn(
-        num_tokens, num_heads, HEAD_DIM, device=device, dtype=torch.bfloat16
-    )
-    positions = torch.randint(
-        0, max_pos, (num_tokens,), device=device, dtype=torch.long
-    )
-    cos_sin_cache = make_cos_sin_cache(max_pos, device=device)
-
-    # Fused kernel result
-    fused_fp8, fused_scale = fused_inv_rope_fp8_quant(
-        o.clone(),
-        positions,
-        cos_sin_cache,
-        n_groups,
-        heads_per_group,
-    )
-
-    # Reference: quantize without RoPE (identity rotation)
-    # Create a zero-sin cache so RoPE is identity
-    zero_cache = torch.zeros_like(cos_sin_cache)
-    half = ROPE_DIM // 2
-    zero_cache[:, :half] = 1.0  # cos = 1
-    # sin = 0 (already zero)
-
-    norope_fp8, norope_scale = fused_inv_rope_fp8_quant(
-        o.clone(),
-        positions,
-        zero_cache,
-        n_groups,
-        heads_per_group,
-    )
-
-    # Extract nope quant blocks only (first 3 of every 4 blocks per head)
-    chunks_per_head = HEAD_DIM // QUANT_GROUP_SIZE  # 4
-
-    for h in range(heads_per_group):
-        for c in range(chunks_per_head - 1):  # skip last chunk (has rope)
-            qb = h * chunks_per_head + c
-            start = qb * QUANT_GROUP_SIZE
-            end = start + QUANT_GROUP_SIZE
-
-            fused_nope = fused_fp8[:, :, start:end].view(torch.uint8)
-            norope_nope = norope_fp8[:, :, start:end].view(torch.uint8)
-            assert torch.equal(fused_nope, norope_nope), (
-                f"Nope block (head={h}, chunk={c}) differs between "
-                f"fused and no-rope reference"
-            )
-
-            fused_s = fused_scale[:, :, qb]
-            norope_s = norope_scale[:, :, qb]
-            assert torch.equal(fused_s, norope_s), (
-                f"Nope scale (head={h}, chunk={c}) differs"
-            )
-
-
-@torch.inference_mode()
-def test_single_token():
-    """Edge case: single token."""
-    num_tokens, num_heads, n_groups = 1, 64, 8
-    heads_per_group = num_heads // n_groups
-    max_pos = 4096
-    device = "cuda"
-
-    o = torch.randn(
-        num_tokens, num_heads, HEAD_DIM, device=device, dtype=torch.bfloat16
-    )
-    positions = torch.tensor([42], device=device, dtype=torch.long)
-    cos_sin_cache = make_cos_sin_cache(max_pos, device=device)
-
-    ref_fp8, ref_scale = reference_inv_rope_fp8_quant(
-        o.clone(),
-        positions,
-        cos_sin_cache,
-        n_groups,
-        heads_per_group,
-    )
-    fused_fp8, fused_scale = fused_inv_rope_fp8_quant(
-        o.clone(),
-        positions,
-        cos_sin_cache,
-        n_groups,
-        heads_per_group,
-    )
-
-    assert_dequant_close(ref_fp8, ref_scale, fused_fp8, fused_scale)
-
-
-@torch.inference_mode()
-def test_zero_positions():
-    """Edge case: all positions are 0."""
-    num_tokens, num_heads, n_groups = 16, 64, 8
-    heads_per_group = num_heads // n_groups
-    max_pos = 4096
-    device = "cuda"
-
-    o = torch.randn(
-        num_tokens, num_heads, HEAD_DIM, device=device, dtype=torch.bfloat16
-    )
-    positions = torch.zeros(num_tokens, device=device, dtype=torch.long)
-    cos_sin_cache = make_cos_sin_cache(max_pos, device=device)
-
-    ref_fp8, ref_scale = reference_inv_rope_fp8_quant(
-        o.clone(),
-        positions,
-        cos_sin_cache,
-        n_groups,
-        heads_per_group,
-    )
-    fused_fp8, fused_scale = fused_inv_rope_fp8_quant(
-        o.clone(),
-        positions,
-        cos_sin_cache,
-        n_groups,
-        heads_per_group,
-    )
-
-    assert_dequant_close(ref_fp8, ref_scale, fused_fp8, fused_scale)
-
-
-@torch.inference_mode()
-def test_large_values():
-    """Edge case: values near FP8 saturation to test clamping."""
-    num_tokens, num_heads, n_groups = 8, 64, 8
-    heads_per_group = num_heads // n_groups
-    max_pos = 4096
-    device = "cuda"
-
-    # Create inputs with large values that will saturate FP8
-    o = torch.randn(
-        num_tokens, num_heads, HEAD_DIM, device=device, dtype=torch.bfloat16
-    )
-    o = o * 1000.0  # scale up to force saturation
-    positions = torch.randint(
-        0, max_pos, (num_tokens,), device=device, dtype=torch.long
-    )
-    cos_sin_cache = make_cos_sin_cache(max_pos, device=device)
-
-    ref_fp8, ref_scale = reference_inv_rope_fp8_quant(
-        o.clone(),
-        positions,
-        cos_sin_cache,
-        n_groups,
-        heads_per_group,
-    )
-    fused_fp8, fused_scale = fused_inv_rope_fp8_quant(
-        o.clone(),
-        positions,
-        cos_sin_cache,
-        n_groups,
-        heads_per_group,
-    )
-
-    assert_dequant_close(ref_fp8, ref_scale, fused_fp8, fused_scale)
-
-
-@torch.inference_mode()
-def test_dequant_numerical_accuracy():
-    """Verify dequantized values are close to the original (after inv RoPE)."""
-    num_tokens, num_heads, n_groups = 32, 64, 8
-    heads_per_group = num_heads // n_groups
-    max_pos = 4096
-    device = "cuda"
-    torch.manual_seed(0)
-
-    o = torch.randn(
-        num_tokens, num_heads, HEAD_DIM, device=device, dtype=torch.bfloat16
-    )
-    positions = torch.randint(
-        0, max_pos, (num_tokens,), device=device, dtype=torch.long
-    )
-    cos_sin_cache = make_cos_sin_cache(max_pos, device=device)
-
-    # Get the post-inv-RoPE values (ground truth before quantization)
-    o_after_rope = reference_inv_rope(o.clone(), positions, cos_sin_cache)
-    d = heads_per_group * HEAD_DIM
-    o_after_rope = o_after_rope.view(num_tokens, n_groups, d)
-
-    # Get fused quantized output
-    fused_fp8, fused_scale = fused_inv_rope_fp8_quant(
-        o.clone(),
-        positions,
-        cos_sin_cache,
-        n_groups,
-        heads_per_group,
-    )
-
-    # Dequantize: broadcast scale [T, G, S] to [T, G, D] via repeat
-    S = d // QUANT_GROUP_SIZE
-    scale_expanded = (
-        fused_scale.unsqueeze(-1)
-        .expand(num_tokens, n_groups, S, QUANT_GROUP_SIZE)
-        .reshape(num_tokens, n_groups, d)
-    )
-    dequant = fused_fp8.float() * scale_expanded
-
-    # Check relative error.
-    # FP8 e4m3 with UE8M0 (power-of-two scales that round UP) quantizes more
-    # coarsely than optimal scaling. Both paths rotate in fp32, so the bulk
-    # of the error comes from UE8M0 quantization itself (~10-12% typical).
-    o_gt = o_after_rope.transpose(0, 1).contiguous().transpose(0, 1)
-    dequant_contig = dequant.transpose(0, 1).contiguous().transpose(0, 1)
-
-    abs_err = (dequant_contig.float() - o_gt.float()).abs()
-    rel_err = abs_err / (o_gt.float().abs().clamp(min=1e-6))
-    mean_rel_err = rel_err.mean().item()
-
-    assert mean_rel_err < 0.15, (
-        f"Mean relative error too high: {mean_rel_err:.4f} (expected < 0.15)"
-    )
-
-
-def _unfused_inv_rope_fp8_quant(
-    o: torch.Tensor,
-    positions: torch.Tensor,
-    cos_sin_cache: torch.Tensor,
-    n_groups: int,
-    heads_per_group: int,
-    nope_dim: int = NOPE_DIM,
-    rope_dim: int = ROPE_DIM,
-) -> tuple[torch.Tensor, torch.Tensor]:
-    """Unfused path matching deepseek_v4_attention.py:295-310.
-
-    Uses the production CUDA RoPE kernel + per_token_group_quant_fp8.
-    """
-    from vllm import _custom_ops as ops
-    from vllm.model_executor.layers.quantization.utils.fp8_utils import (
-        per_token_group_quant_fp8,
-    )
-
-    head_dim = o.shape[-1]
-    rope_dim_offset = head_dim - rope_dim
-
-    # Step 1: In-place CUDA RoPE (same as production)
-    ops.rotary_embedding(
-        positions,
-        o,
-        None,
-        head_dim,
-        cos_sin_cache,
-        False,  # is_neox=False for DeepseekV4 (GPT-J style)
-        rope_dim_offset=rope_dim_offset,
-        inverse=True,
-    )
-
-    # Step 2: Reshape + quant + reshape (same as production)
-    T = o.shape[0]
-    d = heads_per_group * head_dim
-    o = o.view(T, n_groups, -1)
-    o_flat = o.transpose(0, 1).contiguous().reshape(-1, d)
-    o_fp8, o_scale = per_token_group_quant_fp8(
-        o_flat,
-        group_size=QUANT_GROUP_SIZE,
-        use_ue8m0=True,
-    )
-    o_fp8 = o_fp8.view(n_groups, T, d).transpose(0, 1)
-    o_scale = o_scale.view(n_groups, T, -1).transpose(0, 1)
-    return o_fp8, o_scale
-
-
-# =========================================================================
-# End-to-end test including fp8_einsum
-# =========================================================================
-
-
-@pytest.mark.parametrize("num_tokens", [1, 7, 32, 128, 1024])
-@pytest.mark.parametrize(
-    "num_heads,n_groups",
-    [(64, 8)],
-    ids=["H64_G8"],
-)
-@torch.inference_mode()
-def test_einsum_end_to_end(num_tokens, num_heads, n_groups):
-    """End-to-end: fused inv_rope+quant → fp8_einsum must match
-    unfused CUDA_rope+quant → fp8_einsum bitwise.
-
-    This catches stride/layout bugs that only manifest when the einsum
-    kernel actually consumes the quantized activations.
-    """
-    from deep_gemm.utils.math import ceil_div
-
-    from vllm.utils.deep_gemm import (
-        fp8_einsum,
-        per_block_cast_to_fp8,
-        transform_sf_into_required_layout,
-    )
-
-    heads_per_group = num_heads // n_groups
-    d = heads_per_group * HEAD_DIM
-    o_lora_rank = 1024
-    max_pos = 4096
-    device = "cuda"
-    torch.manual_seed(0)
-
-    o = torch.randn(
-        num_tokens, num_heads, HEAD_DIM, device=device, dtype=torch.bfloat16
-    )
-    positions = torch.randint(
-        0, max_pos, (num_tokens,), device=device, dtype=torch.long
-    )
-    cos_sin_cache = make_cos_sin_cache(max_pos, device=device)
-
-    # -- Weight quantization (shared between both paths) --
-    w = torch.randn(n_groups, o_lora_rank, d, device=device, dtype=torch.bfloat16)
-    w_fp8 = torch.empty_like(w, dtype=torch.float8_e4m3fn)
-    w_scale = torch.empty(
-        n_groups,
-        ceil_div(o_lora_rank, 128),
-        ceil_div(d, 128),
-        device=device,
-        dtype=torch.float32,
-    )
-    for g in range(n_groups):
-        w_fp8[g], w_scale[g] = per_block_cast_to_fp8(w[g], use_ue8m0=True)
-
-    recipe = (1, 1, 128)
-    w_scale_t = transform_sf_into_required_layout(
-        sf=w_scale,
-        mn=o_lora_rank,
-        k=d,
-        recipe=(1, 128, 128),
-        num_groups=n_groups,
-        is_sfa=False,
-    )
-
-    # -- UNFUSED path --
-    ref_fp8, ref_scale = _unfused_inv_rope_fp8_quant(
-        o.clone(),
-        positions,
-        cos_sin_cache,
-        n_groups,
-        heads_per_group,
-    )
-    z_ref = torch.empty(
-        num_tokens, n_groups, o_lora_rank, device=device, dtype=torch.bfloat16
-    )
-    fp8_einsum(
-        "bhr,hdr->bhd", (ref_fp8, ref_scale), (w_fp8, w_scale_t), z_ref, recipe=recipe
-    )
-
-    # -- FUSED path --
-    fused_fp8, fused_scale = fused_inv_rope_fp8_quant(
-        o.clone(),
-        positions,
-        cos_sin_cache,
-        n_groups,
-        heads_per_group,
-    )
-    z_fused = torch.empty(
-        num_tokens, n_groups, o_lora_rank, device=device, dtype=torch.bfloat16
-    )
-    fp8_einsum(
-        "bhr,hdr->bhd",
-        (fused_fp8, fused_scale),
-        (w_fp8, w_scale_t),
-        z_fused,
-        recipe=recipe,
-    )
-
-    # -- Checks --
-    # Einsum output: Triton and CUDA both rotate in fp32 now, so diffs
-    # come from fp32 ordering and UE8M0 boundary shifts only.
-    # Use relative diff (same metric as test_fp8_einsum.py).
-    from deep_gemm.testing import calc_diff
-
-    z_diff = calc_diff(z_fused, z_ref)
-    assert z_diff < 0.01, (
-        f"Einsum output diff too large: {z_diff:.6f} (expected < 0.01)"
-    )
-
-
-@pytest.mark.parametrize("num_tokens", [1, 32, 256])
-@torch.inference_mode()
-def test_with_real_deepseek_v4_rope(num_tokens, default_vllm_config):
-    """Test with real DeepseekV4ScalingRotaryEmbedding (GPT-J style,
-    mscale=0, YaRN scaling) matching the production config."""
-
-    num_heads = 64
-    n_groups = 8
-    heads_per_group = num_heads // n_groups
-    device = "cuda"
-    torch.manual_seed(0)
-
-    # Build YaRN-scaled cos_sin_cache matching real DeepSeek V3/V4 config
-    # (mscale=0 → mscale=1.0, so no magnitude scaling)
-    from vllm.model_executor.layers.rotary_embedding.common import (
-        yarn_find_correction_range,
-        yarn_linear_ramp_mask,
-    )
-
-    scaling_factor = 16
-    base = 10000.0
-    max_pos = 65536
-    beta_fast, beta_slow = 32, 1
-
-    pos_freqs = base ** (
-        torch.arange(0, ROPE_DIM, 2, dtype=torch.float32, device=device) / ROPE_DIM
-    )
-    inv_freq_extra = 1.0 / pos_freqs
-    inv_freq_interp = 1.0 / (scaling_factor * pos_freqs)
-    low, high = yarn_find_correction_range(
-        beta_fast, beta_slow, ROPE_DIM, base, max_pos
-    )
-    mask = 1 - yarn_linear_ramp_mask(low, high, ROPE_DIM // 2, dtype=torch.float32).to(
-        device
-    )
-    inv_freq = inv_freq_interp * (1 - mask) + inv_freq_extra * mask
-    t = torch.arange(max_pos * scaling_factor, device=device, dtype=torch.float32)
-    freqs = torch.outer(t, inv_freq)
-    # mscale=0 → yarn_get_mscale returns 1.0
-    cos_sin_cache = torch.cat([freqs.cos(), freqs.sin()], dim=-1)  # fp32
-
-    o = torch.randn(
-        num_tokens, num_heads, HEAD_DIM, device=device, dtype=torch.bfloat16
-    )
-    positions = torch.randint(0, 4096, (num_tokens,), device=device, dtype=torch.long)
-
-    # UNFUSED: CUDA RoPE with is_neox=False (GPT-J)
-    from vllm import _custom_ops as ops
-    from vllm.model_executor.layers.quantization.utils.fp8_utils import (
-        per_token_group_quant_fp8,
-    )
-
-    o_unfused = o.clone()
-    ops.rotary_embedding(
-        positions,
-        o_unfused,
-        None,
-        HEAD_DIM,
-        cos_sin_cache,
-        False,  # is_neox=False (GPT-J style)
-        rope_dim_offset=NOPE_DIM,
-        inverse=True,
-    )
-    d = heads_per_group * HEAD_DIM
-    T = num_tokens
-    o_unfused = o_unfused.view(T, n_groups, d)
-    o_flat = o_unfused.transpose(0, 1).contiguous().reshape(-1, d)
-    ref_fp8, ref_scale = per_token_group_quant_fp8(
-        o_flat,
-        group_size=QUANT_GROUP_SIZE,
-        use_ue8m0=True,
-    )
-    ref_fp8 = ref_fp8.view(n_groups, T, d).transpose(0, 1)
-    ref_scale = ref_scale.view(n_groups, T, -1).transpose(0, 1)
-
-    # FUSED: use the real YaRN-scaled cos_sin_cache
-    fused_fp8, fused_scale = fused_inv_rope_fp8_quant(
-        o.clone(),
-        positions,
-        cos_sin_cache,
-        n_groups,
-        heads_per_group,
-    )
-
-    # Scales must match exactly (same UE8M0 algorithm)
-    # Compare via dequant (Triton bf16 rotation may differ from CUDA by 1 ULP)
-    assert_dequant_close(
-        ref_fp8, ref_scale, fused_fp8, fused_scale, msg="Real DeepSeek V4 rope"
-    )

tests/kernels/test_top_k_per_row.py

@@ -718,6 +718,7 @@ def test_persistent_topk_stress() -> None:
         pytest.param(
             {
                 "seq_lens": [2000, 6000, 30000, 80000],
+                "top_k": 2048,
                 "data_type": "random",
             },
             id="mixed_all_paths",
@@ -726,6 +727,7 @@ def test_persistent_topk_stress() -> None:
         pytest.param(
             {
                 "seq_lens": [2048, 4096, 8192, 16000],
+                "top_k": 2048,
                 "data_type": "random",
             },
             id="all_decode_medium",
@@ -734,6 +736,7 @@ def test_persistent_topk_stress() -> None:
         pytest.param(
             {
                 "seq_lens": [70000, 100000, 163840],
+                "top_k": 2048,
                 "data_type": "random",
             },
             id="all_large",
@@ -742,6 +745,7 @@ def test_persistent_topk_stress() -> None:
         pytest.param(
             {
                 "seq_lens": [32767, 32768, 32769, 32772],
+                "top_k": 2048,
                 "data_type": "random",
             },
             id="large_threshold_boundary",
@@ -750,6 +754,7 @@ def test_persistent_topk_stress() -> None:
         pytest.param(
             {
                 "seq_lens": [5000],
+                "top_k": 2048,
                 "data_type": "random",
             },
             id="single_row_medium",
@@ -767,15 +772,15 @@ def test_persistent_topk_stress() -> None:
         pytest.param(
             {
                 "seq_lens": [100, 2048, 10000, 80000],
+                "top_k": 2048,
                 "data_type": "random",
             },
             id="trivial_medium_large_mix",
         ),
     ],
 )
-@pytest.mark.parametrize("top_k", [512, 2048])
 @torch.inference_mode()
-def test_persistent_topk(test_config: dict, top_k: int) -> None:
+def test_persistent_topk(test_config: dict) -> None:
     """
     Tests specific to the persistent_topk kernel:
     - Mixed medium/large rows in the same batch (dynamic per-row dispatch)
@@ -785,15 +790,14 @@ def test_persistent_topk(test_config: dict, top_k: int) -> None:
     run_large_context_topk_test(
         batch_size=len(test_config["seq_lens"]),
         seq_lens=test_config["seq_lens"],
-        top_k=top_k,
+        top_k=test_config["top_k"],
         data_type=test_config.get("data_type", "random"),
     )
 
 
 @pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA")
-@pytest.mark.parametrize("top_k", [512, 2048])
 @torch.inference_mode()
-def test_persistent_topk_padded_stride(top_k: int) -> None:
+def test_persistent_topk_padded_stride() -> None:
     """
     Test persistent_topk with padded logits (large stride, small seq_len)
     to simulate the e2e CUDAGraph scenario where fp8_paged_mqa_logits
@@ -802,6 +806,7 @@ def test_persistent_topk_padded_stride(top_k: int) -> None:
     set_random_seed(42)
     torch.set_default_device("cuda:0")
 
+    top_k = 2048
     batch_size = 4
     padded_stride = 163840  # DeepSeek-V3.2 max_model_len
     actual_seq_lens = [3000, 5000, 8000, 12000]

tests/lora/test_gptoss_tp.py

@@ -125,7 +125,7 @@ def test_gpt_oss_lora_tp2(
             enable_lora=True,
             max_loras=2,
             max_num_seqs=2,
-            max_num_batched_tokens=2048,
+            max_num_batched_tokens=1024,
             tensor_parallel_size=2,
             gpu_memory_utilization=0.8,
             fully_sharded_loras=fully_sharded_loras,

tests/model_executor/test_routed_experts_capture.py

@@ -41,9 +41,7 @@ class DummyRouter(BaseRouter):
     def routing_method_type(self) -> RoutingMethodType:
         return RoutingMethodType.FUSED_TOPK
 
-    def _compute_routing(
-        self, hidden_states, router_logits, indices_type, *, input_ids=None
-    ):
+    def _compute_routing(self, hidden_states, router_logits, indices_type):
         topk_ids = torch.tensor([[1, 2], [3, 4]], dtype=torch.int64)
         topk_weights = torch.ones_like(topk_ids, dtype=torch.float32)
         return topk_weights, topk_ids

tests/models/registry.py

@@ -260,9 +260,6 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
         trust_remote_code=True,
     ),
     "DeepseekV32ForCausalLM": _HfExamplesInfo("deepseek-ai/DeepSeek-V3.2-Exp"),
-    "DeepseekV4ForCausalLM": _HfExamplesInfo(
-        "deepseek-ai/DeepSeek-V4-Flash", is_available_online=False
-    ),
     "Ernie4_5ForCausalLM": _HfExamplesInfo("baidu/ERNIE-4.5-0.3B-PT"),
     "Ernie4_5_MoeForCausalLM": _HfExamplesInfo("baidu/ERNIE-4.5-21B-A3B-PT"),
     "ExaoneForCausalLM": _HfExamplesInfo(
@@ -594,9 +591,6 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
     "MiMoV2FlashForCausalLM": _HfExamplesInfo(
         "XiaomiMiMo/MiMo-V2-Flash", trust_remote_code=True
     ),
-    "MiMoV2ProForCausalLM": _HfExamplesInfo(
-        "XiaomiMiMo/MiMo-V2.5-Pro", trust_remote_code=True, is_available_online=False
-    ),
     "Dots1ForCausalLM": _HfExamplesInfo("rednote-hilab/dots.llm1.inst"),
 }
 
@@ -964,18 +958,6 @@ _MULTIMODAL_EXAMPLE_MODELS = {
         "PerceptronAI/Isaac-0.1",
         trust_remote_code=True,
         extras={"0.2-2B-Preview": "PerceptronAI/Isaac-0.2-2B-Preview"},
-        max_transformers_version="4.57",
-        transformers_version_reason={
-            "vllm": (
-                "Custom Isaac code is not compatible with Transformers v5. "
-                "The model should be upstreamed to Transformers for "
-                "long-term support."
-            ),
-            "hf": (
-                "Isaac's remote model and processor code import or configure "
-                "APIs that changed in Transformers v5."
-            ),
-        },
     ),
     "InternS1ForConditionalGeneration": _HfExamplesInfo(
         "internlm/Intern-S1",
@@ -1072,9 +1054,6 @@ _MULTIMODAL_EXAMPLE_MODELS = {
     "MiDashengLMModel": _HfExamplesInfo(
         "mispeech/midashenglm-7b", trust_remote_code=True
     ),
-    "MiMoV2OmniForCausalLM": _HfExamplesInfo(
-        "XiaomiMiMo/MiMo-V2.5-Omni", trust_remote_code=True, is_available_online=False
-    ),
     "MiniCPMO": _HfExamplesInfo(
         "openbmb/MiniCPM-o-2_6",
         trust_remote_code=True,
@@ -1503,12 +1482,6 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {
         speculative_model="luccafong/deepseek_mtp_draft_random",
         trust_remote_code=True,
     ),
-    "DeepSeekV4MTPModel": _HfExamplesInfo(
-        "deepseek-ai/DeepSeek-V4-Flash",
-        speculative_model="deepseek-ai/DeepSeek-V4-Flash",
-        trust_remote_code=True,
-        is_available_online=False,
-    ),
     "ErnieMTPModel": _HfExamplesInfo(
         "baidu/ERNIE-4.5-21B-A3B-PT",
         trust_remote_code=True,
@@ -1558,18 +1531,6 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {
         trust_remote_code=True,
         speculative_model="XiaomiMiMo/MiMo-7B-RL",
     ),
-    "MiMoV2MTPModel": _HfExamplesInfo(
-        "XiaomiMiMo/MiMo-V2.5-Pro",
-        trust_remote_code=True,
-        speculative_model="XiaomiMiMo/MiMo-V2.5-Pro",
-        is_available_online=False,
-    ),
-    "MiMoV2OmniMTPModel": _HfExamplesInfo(
-        "XiaomiMiMo/MiMo-V2.5-Omni",
-        trust_remote_code=True,
-        speculative_model="XiaomiMiMo/MiMo-V2.5-Omni",
-        is_available_online=False,
-    ),
     "NemotronHMTPModel": _HfExamplesInfo(
         "nvidia/Nemotron-Super-Placeholder",
         speculative_model="nvidia/Nemotron-Super-Placeholder",

tests/models/test_deepseek_v4_mega_moe.py

@@ -1,184 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-from types import SimpleNamespace
-
-import pytest
-import torch
-
-from vllm.model_executor.models.deepseek_v4 import (
-    DeepseekV4MegaMoEExperts,
-    _stage_deepseek_v4_mega_moe_inputs,
-    make_deepseek_v4_expert_params_mapping,
-)
-from vllm.platforms import current_platform
-
-pytestmark = pytest.mark.skipif(
-    not current_platform.is_cuda(),
-    reason="DeepSeek V4 MegaMoE requires CUDA",
-)
-
-
-def test_deepseek_v4_mega_moe_expert_mapping():
-    mapping = make_deepseek_v4_expert_params_mapping(2)
-
-    assert mapping == [
-        ("experts.w13_", "experts.0.w1.", 0, "w1"),
-        ("experts.w2_", "experts.0.w2.", 0, "w2"),
-        ("experts.w13_", "experts.0.w3.", 0, "w3"),
-        ("experts.w13_", "experts.1.w1.", 1, "w1"),
-        ("experts.w2_", "experts.1.w2.", 1, "w2"),
-        ("experts.w13_", "experts.1.w3.", 1, "w3"),
-    ]
-
-
-def test_deepseek_v4_mega_moe_ue8m0_uint8_to_float():
-    raw = torch.tensor([0, 126, 127, 128], dtype=torch.uint8)
-
-    decoded = DeepseekV4MegaMoEExperts._ue8m0_uint8_to_float(raw)
-
-    assert torch.equal(decoded.view(torch.int32), raw.to(torch.int32) << 23)
-    assert decoded[0].item() == 0.0
-    assert decoded[1].item() == 0.5
-    assert decoded[2].item() == 1.0
-    assert decoded[3].item() == 2.0
-
-
-def test_deepseek_v4_mega_moe_weight_loader_uses_ep_expert_ownership():
-    vllm_config = SimpleNamespace(
-        scheduler_config=SimpleNamespace(max_num_batched_tokens=4)
-    )
-    experts = DeepseekV4MegaMoEExperts(
-        vllm_config,
-        num_experts=4,
-        num_local_experts=2,
-        experts_start_idx=2,
-        top_k=2,
-        hidden_size=128,
-        intermediate_size=128,
-    )
-
-    nonlocal_weight = torch.ones(128, 64, dtype=torch.uint8)
-    assert (
-        experts.weight_loader(
-            experts.w13_weight,
-            nonlocal_weight,
-            "experts.w13_weight",
-            shard_id="w1",
-            expert_id=1,
-            return_success=True,
-        )
-        is False
-    )
-
-    w1 = torch.full((128, 64), 3, dtype=torch.uint8)
-    w3 = torch.full((128, 64), 7, dtype=torch.uint8)
-    w2 = torch.full((128, 64), 11, dtype=torch.uint8)
-
-    assert experts.weight_loader(
-        experts.w13_weight,
-        w1,
-        "experts.w13_weight",
-        shard_id="w1",
-        expert_id=2,
-        return_success=True,
-    )
-    assert experts.weight_loader(
-        experts.w13_weight,
-        w3,
-        "experts.w13_weight",
-        shard_id="w3",
-        expert_id=2,
-        return_success=True,
-    )
-    assert experts.weight_loader(
-        experts.w2_weight,
-        w2,
-        "experts.w2_weight",
-        shard_id="w2",
-        expert_id=2,
-        return_success=True,
-    )
-
-    assert torch.equal(experts.w13_weight[0, :128], w1)
-    assert torch.equal(experts.w13_weight[0, 128:], w3)
-    assert torch.equal(experts.w2_weight[0], w2)
-    assert torch.count_nonzero(experts.w13_weight[1]) == 0
-
-
-@pytest.mark.skipif(
-    not torch.cuda.is_available(),
-    reason="DeepSeek V4 MegaMoE fused input staging requires CUDA.",
-)
-def test_deepseek_v4_mega_moe_fused_input_staging_is_bitwise_exact():
-    from vllm.third_party.deep_gemm.utils import per_token_cast_to_fp8
-
-    device = torch.device("cuda")
-    num_tokens = 7
-    hidden_size = 256
-    top_k = 8
-
-    generator = torch.Generator(device=device)
-    generator.manual_seed(0)
-    hidden_states = (
-        torch.randn(
-            num_tokens,
-            hidden_size,
-            device=device,
-            dtype=torch.float32,
-            generator=generator,
-        )
-        * 17.0
-    ).to(torch.bfloat16)
-    hidden_states[0, :32] = 0
-    hidden_states[1, 32:64] = 1.0e-6
-    hidden_states[2, 64:96] = -1.0e-6
-
-    topk_ids = torch.randint(
-        0,
-        256,
-        (num_tokens, top_k),
-        device=device,
-        dtype=torch.int32,
-        generator=generator,
-    )
-    topk_weights = torch.randn(
-        num_tokens,
-        top_k,
-        device=device,
-        dtype=torch.float32,
-        generator=generator,
-    )
-
-    ref_x, ref_x_sf = per_token_cast_to_fp8(
-        hidden_states,
-        use_ue8m0=True,
-        gran_k=32,
-        use_packed_ue8m0=True,
-    )
-    ref_topk_idx = topk_ids.to(torch.int64)
-    ref_topk_weights = topk_weights.clone()
-
-    fused_x = torch.empty_like(ref_x)
-    fused_x_sf = torch.empty_like(ref_x_sf)
-    fused_topk_idx = torch.empty_like(ref_topk_idx)
-    fused_topk_weights = torch.empty_like(ref_topk_weights)
-
-    _stage_deepseek_v4_mega_moe_inputs(
-        hidden_states,
-        topk_weights,
-        topk_ids,
-        fused_x,
-        fused_x_sf,
-        fused_topk_idx,
-        fused_topk_weights,
-    )
-    torch.accelerator.synchronize()
-
-    assert torch.equal(fused_x.view(torch.uint8), ref_x.view(torch.uint8))
-    assert torch.equal(fused_x_sf, ref_x_sf)
-    assert torch.equal(fused_topk_idx, ref_topk_idx)
-    assert torch.equal(
-        fused_topk_weights.view(torch.uint8),
-        ref_topk_weights.view(torch.uint8),
-    )

tests/multimodal/test_registry.py

@@ -5,8 +5,6 @@ Unit tests for MultiModalRegistry.supports_multimodal_inputs and
 Qwen2.5-VL visual component loading behavior.
 """
 
-from types import SimpleNamespace
-
 import pytest
 
 from vllm.multimodal import MULTIMODAL_REGISTRY
@@ -34,17 +32,3 @@ def test_supports_multimodal_inputs(model_id, limit_mm_per_prompt, expected):
         limit_mm_per_prompt=limit_mm_per_prompt,
     )
     assert MULTIMODAL_REGISTRY.supports_multimodal_inputs(ctx.model_config) is expected
-
-
-def test_create_processor_error_uses_served_model_name():
-    model_config = SimpleNamespace(
-        is_multimodal_model=False,
-        model="/path/to/model/weights",
-        served_model_name="friendly-model-name",
-    )
-
-    with pytest.raises(
-        ValueError,
-        match="friendly-model-name is not a multimodal model",
-    ):
-        MULTIMODAL_REGISTRY.create_processor(model_config)

tests/quantization/test_cutlass_w4a16.py

@@ -1,185 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Tests for Cutlass W4A16 (Machete) kernel on Hopper.
-
-Verifies that W4A16 quantized models loaded through vllm select the
-MacheteLinearKernel on sm_90 GPUs, that weights are correctly repacked,
-and that inference produces valid output.
-
-Run `pytest tests/quantization/test_cutlass_w4a16.py`.
-"""
-
-import pytest
-import torch
-
-from vllm.platforms import current_platform
-
-if not current_platform.has_device_capability(90):
-    pytest.skip(
-        "Machete W4A16 requires Hopper (sm_90).",
-        allow_module_level=True,
-    )
-
-from vllm.model_executor.kernels.linear import (
-    MPLinearLayerConfig,
-    choose_mp_linear_kernel,
-)
-from vllm.model_executor.kernels.linear.mixed_precision import (
-    MacheteLinearKernel,
-)
-from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors import (  # noqa: E501
-    CompressedTensorsLinearMethod,
-    CompressedTensorsWNA16,
-)
-from vllm.scalar_type import scalar_types
-
-
-@pytest.fixture(scope="function", autouse=True)
-def enable_pickle(monkeypatch):
-    """`LLM.apply_model` requires pickling a function."""
-    monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
-
-
-@pytest.mark.parametrize(
-    "act_type,weight_type,group_size,zero_points",
-    [
-        (torch.float16, scalar_types.uint4b8, 128, False),
-        (torch.bfloat16, scalar_types.uint4b8, 128, False),
-        (torch.float16, scalar_types.uint4, 128, True),
-        (torch.float16, scalar_types.uint4b8, -1, False),
-    ],
-    ids=[
-        "fp16-gptq-g128",
-        "bf16-gptq-g128",
-        "fp16-awq-g128",
-        "fp16-channelwise",
-    ],
-)
-def test_machete_kernel_selected(act_type, weight_type, group_size, zero_points):
-    """Verify choose_mp_linear_kernel picks MacheteLinearKernel."""
-    config = MPLinearLayerConfig(
-        full_weight_shape=(4096, 4096),
-        partition_weight_shape=(4096, 4096),
-        act_type=act_type,
-        weight_type=weight_type,
-        group_size=group_size,
-        zero_points=zero_points,
-        has_g_idx=False,
-    )
-    kernel = choose_mp_linear_kernel(config)
-    assert kernel is MacheteLinearKernel, (
-        f"Expected MacheteLinearKernel, got {kernel.__name__}"
-    )
-
-
-@pytest.mark.parametrize(
-    "full_shape,part_shape,weight_type,group_size,has_g_idx,expected_reason",
-    [
-        ((4096, 4096), (2048, 4096), scalar_types.uint4b8, 128, True, "Act reordering"),
-        (
-            (4096, 4096),
-            (4096, 4096),
-            scalar_types.float6_e3m2f,
-            128,
-            False,
-            "Quant type",
-        ),
-        ((4096, 4096), (4096, 4096), scalar_types.uint4b8, 32, False, "Group size"),
-    ],
-    ids=["partitioned-g_idx", "unsupported-quant-type", "unsupported-group-size"],
-)
-def test_machete_rejects_invalid_config(
-    full_shape, part_shape, weight_type, group_size, has_g_idx, expected_reason
-):
-    """Verify Machete rejects unsupported configurations."""
-    config = MPLinearLayerConfig(
-        full_weight_shape=full_shape,
-        partition_weight_shape=part_shape,
-        act_type=torch.float16,
-        weight_type=weight_type,
-        group_size=group_size,
-        zero_points=False,
-        has_g_idx=has_g_idx,
-    )
-    can_impl, reason = MacheteLinearKernel.can_implement(config)
-    assert not can_impl
-    assert expected_reason in reason
-
-
-def test_kernel_selection_with_disabled_machete(monkeypatch):
-    """Verify kernel selection falls back when Machete is disabled."""
-    monkeypatch.setattr("vllm.envs.VLLM_DISABLED_KERNELS", ["MacheteLinearKernel"])
-
-    config = MPLinearLayerConfig(
-        full_weight_shape=(4096, 4096),
-        partition_weight_shape=(4096, 4096),
-        act_type=torch.float16,
-        weight_type=scalar_types.uint4b8,
-        group_size=128,
-        zero_points=False,
-        has_g_idx=False,
-    )
-    kernel = choose_mp_linear_kernel(config)
-    assert kernel is not MacheteLinearKernel, "MacheteLinearKernel should be disabled"
-
-
-@pytest.mark.parametrize(
-    "model_name",
-    [
-        "nm-testing/tinyllama-oneshot-w4a16-channel-v2",
-        "nm-testing/TinyLlama-1.1B-Chat-v1.0-W4A16-G128-Asym-Updated-ActOrder",
-    ],
-)
-def test_w4a16_machete_e2e(vllm_runner, model_name):
-    """Load a W4A16 model, verify Machete kernel is used, and generate."""
-    with vllm_runner(model_name, enforce_eager=True, gpu_memory_utilization=0.5) as llm:
-
-        def check_model(model):
-            layer = model.model.layers[0]
-            qkv_proj = layer.self_attn.qkv_proj
-
-            assert isinstance(qkv_proj.quant_method, CompressedTensorsLinearMethod)
-            assert isinstance(qkv_proj.scheme, CompressedTensorsWNA16)
-            assert isinstance(qkv_proj.scheme.kernel, MacheteLinearKernel), (
-                f"Expected MacheteLinearKernel on Hopper, "
-                f"got {type(qkv_proj.scheme.kernel).__name__}"
-            )
-
-            assert hasattr(qkv_proj, "weight_packed")
-            assert hasattr(qkv_proj, "weight_scale")
-            assert qkv_proj.weight_packed.dtype == torch.int32
-
-        llm.apply_model(check_model)
-
-        output = llm.generate_greedy("Hello my name is", max_tokens=10)
-        assert output
-        assert len(output[0][1]) > 0
-
-
-def test_w4a16_machete_bfloat16_deterministic(vllm_runner):
-    """Verify Machete works with bf16 activations and is deterministic."""
-    model_name = "nm-testing/tinyllama-oneshot-w4a16-channel-v2"
-    prompt = "The capital of France is"
-
-    with vllm_runner(
-        model_name,
-        enforce_eager=True,
-        dtype="bfloat16",
-        gpu_memory_utilization=0.5,
-    ) as llm:
-
-        def check_kernel_type(model):
-            layer = model.model.layers[0]
-            scheme = layer.self_attn.qkv_proj.scheme
-            assert isinstance(scheme.kernel, MacheteLinearKernel), (
-                f"Expected MacheteLinearKernel with bf16, "
-                f"got {type(scheme.kernel).__name__}"
-            )
-
-        llm.apply_model(check_kernel_type)
-
-        out1 = llm.generate_greedy(prompt, max_tokens=10)
-        out2 = llm.generate_greedy(prompt, max_tokens=10)
-        assert out1[0][1] == out2[0][1], (
-            f"Non-deterministic: '{out1[0][1]}' vs '{out2[0][1]}'"
-        )

tests/reasoning/test_deepseekv3_reasoning_parser.py

@@ -6,7 +6,6 @@ from transformers import AutoTokenizer
 
 from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
 from vllm.entrypoints.openai.engine.protocol import DeltaMessage
-from vllm.reasoning import ReasoningParserManager
 from vllm.reasoning.deepseek_r1_reasoning_parser import DeepSeekR1ReasoningParser
 from vllm.reasoning.deepseek_v3_reasoning_parser import DeepSeekV3ReasoningParser
 from vllm.reasoning.identity_reasoning_parser import IdentityReasoningParser
@@ -34,12 +33,6 @@ def test_parser_selection(tokenizer, thinking, expected_parser_type):
     assert isinstance(parser._parser, expected_parser_type)
 
 
-def test_deepseek_v4_reasoning_parser_alias():
-    parser_cls = ReasoningParserManager.get_reasoning_parser("deepseek_v4")
-
-    assert parser_cls is DeepSeekV3ReasoningParser
-
-
 def test_identity_reasoning_parser_basic(tokenizer):
     parser = IdentityReasoningParser(tokenizer)
 

tests/tokenizers_/fixtures/deepseek_v4/test_input_1.json

@@ -1,81 +0,0 @@
-{
-    "tools": [
-        {
-            "type": "function",
-            "function": {
-                "name": "get_weather",
-                "description": "Get the weather for a specific location",
-                "parameters": {
-                    "type": "object",
-                    "properties": {
-                        "location": {
-                            "type": "string",
-                            "description": "The city name"
-                        },
-                        "unit": {
-                            "type": "string",
-                            "enum": ["celsius", "fahrenheit"],
-                            "description": "Temperature unit"
-                        }
-                    },
-                    "required": ["location"]
-                }
-            }
-        },
-        {
-            "type": "function",
-            "function": {
-                "name": "search",
-                "description": "Search the web for information",
-                "parameters": {
-                    "type": "object",
-                    "properties": {
-                        "query": {
-                            "type": "string",
-                            "description": "Search query"
-                        },
-                        "num_results": {
-                            "type": "integer",
-                            "description": "Number of results to return"
-                        }
-                    },
-                    "required": ["query"]
-                }
-            }
-        }
-    ],
-    "messages": [
-        {
-            "role": "system",
-            "content": "You are a helpful assistant."
-        },
-        {
-            "role": "user",
-            "content": "What's the weather in Beijing?"
-        },
-        {
-            "role": "assistant",
-            "reasoning": "The user wants to know the weather in Beijing. I should use the get_weather tool.",
-            "tool_calls": [
-                {
-                    "id": "call_001",
-                    "type": "function",
-                    "function": {
-                        "name": "get_weather",
-                        "arguments": "{\"location\": \"Beijing\", \"unit\": \"celsius\"}"
-                    }
-                }
-            ]
-        },
-        {
-            "role": "tool",
-            "tool_call_id": "call_001",
-            "content": "{\"temperature\": 22, \"condition\": \"sunny\", \"humidity\": 45}"
-        },
-        {
-            "role": "assistant",
-            "reasoning": "Got the weather data. Let me format a nice response.",
-            "content": "The weather in Beijing is currently sunny with a temperature of 22°C and 45% humidity."
-        }
-    ]
-}

tests/tokenizers_/fixtures/deepseek_v4/test_input_2.json

@@ -1,24 +0,0 @@
-[
-  {
-    "role": "system",
-    "content": "You are a helpful assistant."
-  },
-  {
-    "role": "user",
-    "content": "Hello"
-  },
-  {
-    "role": "assistant",
-    "reasoning": "The user said hello, I should greet back.",
-    "content": "Hi there! How can I help you?"
-  },
-  {
-    "role": "user",
-    "content": "What is the capital of France?"
-  },
-  {
-    "role": "assistant",
-    "reasoning": "The user asks about the capital of France. It is Paris.",
-    "content": "The capital of France is Paris."
-  }
-]

tests/tokenizers_/fixtures/deepseek_v4/test_input_3.json

@@ -1,159 +0,0 @@
-[
-  {
-    "role": "system",
-    "content": "该助手为DeepSeek，由深度求索公司创造。"
-  },
-  {
-    "role": "latest_reminder",
-    "content": "2026-02-21,星期六,广州,App,中文"
-  },
-  {
-    "role": "developer",
-    "content": "小柴胡冲剂和布洛芬能一起吃吗？\n\nCITATION FORMAT: 【{cursor_id}†L{start_line_id}(-L{end_line_id})?】",
-    "tools": [
-      {
-        "type": "function",
-        "function": {
-          "name": "search",
-          "description": "Web search. Split multiple queries with '||'.",
-          "parameters": {
-            "type": "object",
-            "properties": {
-              "queries": {
-                "type": "string",
-                "description": "query1||query2"
-              }
-            },
-            "required": [
-              "queries"
-            ],
-            "additionalProperties": false,
-            "$schema": "http://json-schema.org/draft-07/schema#"
-          }
-        }
-      },
-      {
-        "type": "function",
-        "function": {
-          "name": "open",
-          "description": "Batch open IDs (format 【{id}†...】) or URLs.",
-          "parameters": {
-            "type": "object",
-            "properties": {
-              "open_list": {
-                "type": "array",
-                "items": {
-                  "type": "object",
-                  "properties": {
-                    "id": {
-                      "description": "ID or URL",
-                      "anyOf": [
-                        {
-                          "type": "integer"
-                        },
-                        {
-                          "type": "string"
-                        }
-                      ],
-                      "default": -1
-                    },
-                    "cursor": {
-                      "type": "integer",
-                      "description": "",
-                      "default": -1
-                    },
-                    "loc": {
-                      "type": "integer",
-                      "description": "Start line",
-                      "default": -1
-                    },
-                    "num_lines": {
-                      "type": "integer",
-                      "description": "",
-                      "default": -1
-                    },
-                    "view_source": {
-                      "type": "boolean",
-                      "description": "",
-                      "default": false
-                    }
-                  },
-                  "additionalProperties": false
-                },
-                "description": ""
-              }
-            },
-            "required": [
-              "open_list"
-            ],
-            "additionalProperties": false,
-            "$schema": "http://json-schema.org/draft-07/schema#"
-          }
-        }
-      },
-      {
-        "type": "function",
-        "function": {
-          "name": "find",
-          "description": "Find exact text pattern in pages.",
-          "parameters": {
-            "type": "object",
-            "properties": {
-              "find_list": {
-                "type": "array",
-                "items": {
-                  "type": "object",
-                  "properties": {
-                    "pattern": {
-                      "type": "string",
-                      "description": ""
-                    },
-                    "cursor": {
-                      "type": "integer",
-                      "description": "",
-                      "default": -1
-                    }
-                  },
-                  "required": [
-                    "pattern"
-                  ],
-                  "additionalProperties": false
-                },
-                "description": ""
-              }
-            },
-            "required": [
-              "find_list"
-            ],
-            "additionalProperties": false,
-            "$schema": "http://json-schema.org/draft-07/schema#"
-          }
-        }
-      }
-    ]
-  },
-  {
-    "role": "assistant",
-    "content": "",
-    "reasoning": "用户想知道小柴胡冲剂和布洛芬能否一起服用。",
-    "tool_calls": [
-      {
-        "type": "function",
-        "function": {
-          "name": "search",
-          "arguments": "{\"queries\": \"小柴胡冲剂 布洛芬 相互作用 一起吃\"}"
-        }
-      }
-    ]
-  },
-  {
-    "role": "tool",
-    "content": "[0]"
-  },
-  {
-    "role": "assistant",
-    "content": "请及时就医。",
-    "reasoning": "现在开始组织回答。",
-    "tool_calls": []
-  }
-]

tests/tokenizers_/fixtures/deepseek_v4/test_input_4.json

@@ -1,28 +0,0 @@
-[
-  {
-    "role": "system",
-    "content": "该助手为DeepSeek-V3，由深度求索公司创造。\n今天是2025年10月17日，星期五。"
-  },
-  {
-    "role": "latest_reminder",
-    "content": "2024-11-15,上海市,App,中文"
-  },
-  {
-    "role": "user",
-    "content": "热海大滚锅是世界著名温泉吗"
-  },
-  {
-    "role": "assistant",
-    "content": "关于热海大滚锅是否是世界著名温泉，可以这样总结：它在**中国乃至全球的地热奇观中占有重要地位**，但“世界著名”的称号更侧重于它作为**独特的地质现象和旅游景点**，而非普遍的温泉疗养体验。\n\n为了让你快速了解，我整理了一个简要的表格：\n\n| 维度 | 热海大滚锅的地位与特点 |\n| :--- | :--- |\n| **地质奇观** | **世界罕见**的高温地热系统，泉眼多、水温高、形态多样。 |\n| **国内声誉** | **中国三大地热区之一**，**国家5A级旅游景区**，自古闻名（徐霞客曾游历并记载）。 |\n| **国际知名度** | 在一些旅游平台被提及为“世界六大温泉”之一，但此说法流传不广，其国际声誉更多建立在地质独特性上。 |\n| **核心体验** | **观赏地热奇观**（如97℃沸腾的“大滚锅”）、**体验温泉煮鸡蛋**。 |\n\n### 💡 游玩攻略与温馨提示\n\n如果你计划前往热海大滚锅，这里有一些实用信息供你参考：\n\n- **门票与开放时间**：\n    - **门票**：景区门票约为**50元/人**。如果选择包含温泉沐浴的套餐，价格会更高，例如约**288元**。\n    - **开放时间**：景区一般**08:00-18:00**开放，但具体时间可能变动，建议提前核实。\n\n- **特色体验**：\n    - **温泉煮鸡蛋**：这几乎是必试项目。可以在景区门口购买用草绳串起的生鸡蛋（约5-8元/串），然后到“大滚锅”旁的指定区域蒸煮，几分钟便可熟食，趣味十足。\n    - **金汤足浴**：可以直接用从“大滚锅”流出的温泉水泡脚，缓解旅途疲劳。\n\n- **注意事项**：\n    - **安全第一**：“大滚锅”水温极高，务必遵守游览规则，在指定区域内观赏，切勿随意触碰泉水。\n    - **规划行程**：建议为热海景区预留**3-4小时**的游览时间。景区内步道不走回头路，出入口有观光车接送。\n\n希望这些信息能帮助你更好地了解热海大滚锅。如果你对腾冲的其他景点或者行程规划有更多疑问，我很乐意提供进一步的信息。",
-    "mask": 1
-  },
-  {
-    "role": "user",
-    "content": "世界著名温泉有哪些",
-    "task": "action"
-  },
-  {
-    "role": "assistant",
-    "content": "Search"
-  }
-]

tests/tokenizers_/fixtures/deepseek_v4/test_output_1.txt

@@ -1,36 +0,0 @@
-<｜begin▁of▁sentence｜>
-
-## Tools
-
-You have access to a set of tools to help answer the user's question. You can invoke tools by writing a "<｜DSML｜tool_calls>" block like the following:
-
-<｜DSML｜tool_calls>
-<｜DSML｜invoke name="$TOOL_NAME">
-<｜DSML｜parameter name="$PARAMETER_NAME" string="true|false">$PARAMETER_VALUE</｜DSML｜parameter>
-...
-</｜DSML｜invoke>
-<｜DSML｜invoke name="$TOOL_NAME2">
-...
-</｜DSML｜invoke>
-</｜DSML｜tool_calls>
-
-String parameters should be specified as is and set `string="true"`. For all other types (numbers, booleans, arrays, objects), pass the value in JSON format and set `string="false"`.
-
-If thinking_mode is enabled (triggered by <think>), you MUST output your complete reasoning inside <think>...</think> BEFORE any tool calls or final response.
-
-Otherwise, output directly after </think> with tool calls or final response.
-
-### Available Tool Schemas
-
-{"name": "get_weather", "description": "Get the weather for a specific location", "parameters": {"type": "object", "properties": {"location": {"type": "string", "description": "The city name"}, "unit": {"type": "string", "enum": ["celsius", "fahrenheit"], "description": "Temperature unit"}}, "required": ["location"]}}
-{"name": "search", "description": "Search the web for information", "parameters": {"type": "object", "properties": {"query": {"type": "string", "description": "Search query"}, "num_results": {"type": "integer", "description": "Number of results to return"}}, "required": ["query"]}}
-
-You MUST strictly follow the above defined tool name and parameter schemas to invoke tool calls.
-You are a helpful assistant.<｜User｜>What's the weather in Beijing?<｜Assistant｜><think>The user wants to know the weather in Beijing. I should use the get_weather tool.</think>
-
-<｜DSML｜tool_calls>
-<｜DSML｜invoke name="get_weather">
-<｜DSML｜parameter name="location" string="true">Beijing</｜DSML｜parameter>
-<｜DSML｜parameter name="unit" string="true">celsius</｜DSML｜parameter>
-</｜DSML｜invoke>
-</｜DSML｜tool_calls><｜end▁of▁sentence｜><｜User｜><tool_result>{"temperature": 22, "condition": "sunny", "humidity": 45}</tool_result><｜Assistant｜><think>Got the weather data. Let me format a nice response.</think>The weather in Beijing is currently sunny with a temperature of 22°C and 45% humidity.<｜end▁of▁sentence｜>

tests/tokenizers_/fixtures/deepseek_v4/test_output_2.txt

@@ -1 +0,0 @@
-<｜begin▁of▁sentence｜>You are a helpful assistant.<｜User｜>Hello<｜Assistant｜></think>Hi there! How can I help you?<｜end▁of▁sentence｜><｜User｜>What is the capital of France?<｜Assistant｜><think>The user asks about the capital of France. It is Paris.</think>The capital of France is Paris.<｜end▁of▁sentence｜>

tests/tokenizers_/fixtures/deepseek_v4/test_output_3.txt

@@ -1,38 +0,0 @@
-<｜begin▁of▁sentence｜>该助手为DeepSeek，由深度求索公司创造。<｜latest_reminder｜>2026-02-21,星期六,广州,App,中文<｜User｜>小柴胡冲剂和布洛芬能一起吃吗？
-
-CITATION FORMAT: 【{cursor_id}†L{start_line_id}(-L{end_line_id})?】
-
-## Tools
-
-You have access to a set of tools to help answer the user's question. You can invoke tools by writing a "<｜DSML｜tool_calls>" block like the following:
-
-<｜DSML｜tool_calls>
-<｜DSML｜invoke name="$TOOL_NAME">
-<｜DSML｜parameter name="$PARAMETER_NAME" string="true|false">$PARAMETER_VALUE</｜DSML｜parameter>
-...
-</｜DSML｜invoke>
-<｜DSML｜invoke name="$TOOL_NAME2">
-...
-</｜DSML｜invoke>
-</｜DSML｜tool_calls>
-
-String parameters should be specified as is and set `string="true"`. For all other types (numbers, booleans, arrays, objects), pass the value in JSON format and set `string="false"`.
-
-If thinking_mode is enabled (triggered by <think>), you MUST output your complete reasoning inside <think>...</think> BEFORE any tool calls or final response.
-
-Otherwise, output directly after </think> with tool calls or final response.
-
-### Available Tool Schemas
-
-{"name": "search", "description": "Web search. Split multiple queries with '||'.", "parameters": {"type": "object", "properties": {"queries": {"type": "string", "description": "query1||query2"}}, "required": ["queries"], "additionalProperties": false, "$schema": "http://json-schema.org/draft-07/schema#"}}
-{"name": "open", "description": "Batch open IDs (format 【{id}†...】) or URLs.", "parameters": {"type": "object", "properties": {"open_list": {"type": "array", "items": {"type": "object", "properties": {"id": {"description": "ID or URL", "anyOf": [{"type": "integer"}, {"type": "string"}], "default": -1}, "cursor": {"type": "integer", "description": "", "default": -1}, "loc": {"type": "integer", "description": "Start line", "default": -1}, "num_lines": {"type": "integer", "description": "", "default": -1}, "view_source": {"type": "boolean", "description": "", "default": false}}, "additionalProperties": false}, "description": ""}}, "required": ["open_list"], "additionalProperties": false, "$schema": "http://json-schema.org/draft-07/schema#"}}
-{"name": "find", "description": "Find exact text pattern in pages.", "parameters": {"type": "object", "properties": {"find_list": {"type": "array", "items": {"type": "object", "properties": {"pattern": {"type": "string", "description": ""}, "cursor": {"type": "integer", "description": "", "default": -1}}, "required": ["pattern"], "additionalProperties": false}, "description": ""}}, "required": ["find_list"], "additionalProperties": false, "$schema": "http://json-schema.org/draft-07/schema#"}}
-
-You MUST strictly follow the above defined tool name and parameter schemas to invoke tool calls.
-<｜Assistant｜><think>用户想知道小柴胡冲剂和布洛芬能否一起服用。</think>
-
-<｜DSML｜tool_calls>
-<｜DSML｜invoke name="search">
-<｜DSML｜parameter name="queries" string="true">小柴胡冲剂 布洛芬 相互作用 一起吃</｜DSML｜parameter>
-</｜DSML｜invoke>
-</｜DSML｜tool_calls><｜end▁of▁sentence｜><｜User｜><tool_result>[0]</tool_result><｜Assistant｜><think>现在开始组织回答。</think>请及时就医。<｜end▁of▁sentence｜>

tests/tokenizers_/fixtures/deepseek_v4/test_output_4.txt

@@ -1,29 +0,0 @@
-<｜begin▁of▁sentence｜>该助手为DeepSeek-V3，由深度求索公司创造。
-今天是2025年10月17日，星期五。<｜latest_reminder｜>2024-11-15,上海市,App,中文<｜User｜>热海大滚锅是世界著名温泉吗<｜Assistant｜></think>关于热海大滚锅是否是世界著名温泉，可以这样总结：它在**中国乃至全球的地热奇观中占有重要地位**，但“世界著名”的称号更侧重于它作为**独特的地质现象和旅游景点**，而非普遍的温泉疗养体验。
-
-为了让你快速了解，我整理了一个简要的表格：
-
-| 维度 | 热海大滚锅的地位与特点 |
-| :--- | :--- |
-| **地质奇观** | **世界罕见**的高温地热系统，泉眼多、水温高、形态多样。 |
-| **国内声誉** | **中国三大地热区之一**，**国家5A级旅游景区**，自古闻名（徐霞客曾游历并记载）。 |
-| **国际知名度** | 在一些旅游平台被提及为“世界六大温泉”之一，但此说法流传不广，其国际声誉更多建立在地质独特性上。 |
-| **核心体验** | **观赏地热奇观**（如97℃沸腾的“大滚锅”）、**体验温泉煮鸡蛋**。 |
-
-### 💡 游玩攻略与温馨提示
-
-如果你计划前往热海大滚锅，这里有一些实用信息供你参考：
-
-- **门票与开放时间**：
-    - **门票**：景区门票约为**50元/人**。如果选择包含温泉沐浴的套餐，价格会更高，例如约**288元**。
-    - **开放时间**：景区一般**08:00-18:00**开放，但具体时间可能变动，建议提前核实。
-
-- **特色体验**：
-    - **温泉煮鸡蛋**：这几乎是必试项目。可以在景区门口购买用草绳串起的生鸡蛋（约5-8元/串），然后到“大滚锅”旁的指定区域蒸煮，几分钟便可熟食，趣味十足。
-    - **金汤足浴**：可以直接用从“大滚锅”流出的温泉水泡脚，缓解旅途疲劳。
-
-- **注意事项**：
-    - **安全第一**：“大滚锅”水温极高，务必遵守游览规则，在指定区域内观赏，切勿随意触碰泉水。
-    - **规划行程**：建议为热海景区预留**3-4小时**的游览时间。景区内步道不走回头路，出入口有观光车接送。
-
-希望这些信息能帮助你更好地了解热海大滚锅。如果你对腾冲的其他景点或者行程规划有更多疑问，我很乐意提供进一步的信息。<｜end▁of▁sentence｜><｜User｜>世界著名温泉有哪些<｜Assistant｜></think><｜action｜>Search<｜end▁of▁sentence｜>

tests/tokenizers_/test_deepseek_v4.py

@@ -1,224 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import json
-from pathlib import Path
-from types import SimpleNamespace
-
-import pytest
-
-from vllm.entrypoints.chat_utils import parse_chat_messages
-from vllm.renderers.registry import RENDERER_REGISTRY
-from vllm.tokenizers.deepseek_v4 import get_deepseek_v4_tokenizer
-from vllm.tokenizers.registry import TokenizerRegistry
-
-FIXTURES_DIR = Path(__file__).parent / "fixtures" / "deepseek_v4"
-
-
-class FakeHfTokenizer:
-    vocab_size = 100
-
-    def get_added_vocab(self) -> dict[str, int]:
-        return {"</think>": 100}
-
-    def encode(
-        self,
-        text: str,
-        add_special_tokens: bool = False,
-        **kwargs,
-    ) -> list[int]:
-        self.last_encode = (text, add_special_tokens, kwargs)
-        return [len(text)]
-
-
-def _tokenizer():
-    return get_deepseek_v4_tokenizer(FakeHfTokenizer())
-
-
-def _model_config():
-    return SimpleNamespace(
-        multimodal_config=None,
-        allowed_local_media_path="",
-        allowed_media_domains=None,
-    )
-
-
-def _load_reference_case(case_id: int):
-    data = json.loads((FIXTURES_DIR / f"test_input_{case_id}.json").read_text())
-    if isinstance(data, dict):
-        return data["messages"], data.get("tools")
-    return data, None
-
-
-def _render_reference_case(case_id: int, **kwargs):
-    messages, tools = _load_reference_case(case_id)
-    conversation, _, _ = parse_chat_messages(
-        messages,
-        _model_config(),
-        content_format="string",
-    )
-    return _tokenizer().apply_chat_template(
-        conversation=conversation,
-        messages=messages,
-        tools=tools,
-        tokenize=False,
-        **kwargs,
-    )
-
-
-def test_deepseek_v4_tokenizer_registered():
-    assert TokenizerRegistry.load_tokenizer_cls("deepseek_v4").__name__ == (
-        "DeepseekV4Tokenizer"
-    )
-    assert RENDERER_REGISTRY.load_renderer_cls("deepseek_v4").__name__ == (
-        "DeepseekV4Renderer"
-    )
-
-
-def test_deepseek_v4_defaults_to_chat_mode():
-    prompt = _tokenizer().apply_chat_template(
-        [{"role": "user", "content": "Hello"}],
-        tokenize=False,
-    )
-
-    assert prompt == ("<｜begin▁of▁sentence｜><｜User｜>Hello<｜Assistant｜></think>")
-
-
-@pytest.mark.parametrize("kwargs", [{"thinking": True}, {"enable_thinking": True}])
-def test_deepseek_v4_enables_thinking_with_compatible_kwargs(kwargs):
-    prompt = _tokenizer().apply_chat_template(
-        [{"role": "user", "content": "Hello"}],
-        tokenize=False,
-        **kwargs,
-    )
-
-    assert prompt == ("<｜begin▁of▁sentence｜><｜User｜>Hello<｜Assistant｜><think>")
-
-
-def test_deepseek_v4_uses_v4_tool_prompt_from_request_tools():
-    tools = [
-        {
-            "type": "function",
-            "function": {
-                "name": "get_weather",
-                "description": "Get weather for a city",
-                "parameters": {
-                    "type": "object",
-                    "properties": {"city": {"type": "string"}},
-                    "required": ["city"],
-                },
-            },
-        }
-    ]
-
-    prompt = _tokenizer().apply_chat_template(
-        [{"role": "user", "content": "Weather?"}],
-        tools=tools,
-        tokenize=False,
-    )
-
-    assert "## Tools" in prompt
-    assert "<｜DSML｜tool_calls>" in prompt
-    assert "</｜DSML｜tool_calls>" in prompt
-    assert "function_calls" not in prompt
-    assert '"name": "get_weather"' in prompt
-    assert prompt.endswith("<｜User｜>Weather?<｜Assistant｜></think>")
-
-
-def test_deepseek_v4_renders_parsed_history_tool_arguments():
-    messages = [
-        {"role": "user", "content": "List the repo"},
-        {
-            "role": "assistant",
-            "tool_calls": [
-                {
-                    "id": "call_1",
-                    "type": "function",
-                    "function": {
-                        "name": "str_replace_editor",
-                        "arguments": '{"command": "view", "path": "/testbed"}',
-                    },
-                }
-            ],
-        },
-        {
-            "role": "tool",
-            "tool_call_id": "call_1",
-            "content": "file list",
-        },
-    ]
-    tools = [
-        {
-            "type": "function",
-            "function": {
-                "name": "str_replace_editor",
-                "description": "Edit files",
-                "parameters": {
-                    "type": "object",
-                    "properties": {
-                        "command": {"type": "string"},
-                        "path": {"type": "string"},
-                    },
-                    "required": ["command", "path"],
-                },
-            },
-        }
-    ]
-    conversation, _, _ = parse_chat_messages(
-        messages,
-        _model_config(),
-        content_format="string",
-    )
-
-    prompt = _tokenizer().apply_chat_template(
-        conversation=conversation,
-        messages=messages,
-        tools=tools,
-        tokenize=False,
-    )
-
-    assert '<｜DSML｜parameter name="command" string="true">view' in prompt
-    assert '<｜DSML｜parameter name="path" string="true">/testbed' in prompt
-    assert 'parameter name="arguments"' not in prompt
-
-
-@pytest.mark.parametrize("reasoning_effort", ["none", "low", "medium", "high"])
-def test_deepseek_v4_accepts_openai_reasoning_effort_values(reasoning_effort):
-    prompt = _tokenizer().apply_chat_template(
-        [{"role": "user", "content": "Hello"}],
-        tokenize=False,
-        enable_thinking=True,
-        reasoning_effort=reasoning_effort,
-    )
-
-    assert prompt.endswith("<｜Assistant｜><think>")
-    assert "Reasoning Effort: Absolute maximum" not in prompt
-
-
-def test_deepseek_v4_preserves_reference_max_reasoning_effort():
-    prompt = _tokenizer().apply_chat_template(
-        [{"role": "user", "content": "Hello"}],
-        tokenize=False,
-        enable_thinking=True,
-        reasoning_effort="max",
-    )
-
-    assert prompt.startswith(
-        "<｜begin▁of▁sentence｜>Reasoning Effort: Absolute maximum"
-    )
-
-
-@pytest.mark.parametrize(
-    ("case_id", "kwargs"),
-    [
-        (1, {"thinking": True}),
-        (2, {"thinking": True}),
-        (3, {"thinking": True}),
-        (4, {}),
-    ],
-)
-def test_deepseek_v4_matches_reference_golden_fixtures(case_id, kwargs):
-    prompt = _render_reference_case(case_id, **kwargs)
-
-    expected = (FIXTURES_DIR / f"test_output_{case_id}.txt").read_text()
-    assert prompt == expected

tests/tool_parsers/test_deepseekv32_tool_parser.py

@@ -484,58 +484,6 @@ class TestExtractToolCallsStreaming:
         # Should have no tool call deltas yet
         assert all(not d.tool_calls for d in deltas)
 
-    def test_no_marker_leak_chunked(self, parser):
-        """Chunked streaming must NOT leak DSML start-marker fragments
-        as content (GitHub #40801)."""
-        full_text = build_tool_call("fn", {"k": "v"})
-        deltas = self._stream_chunked(parser, full_text, chunk_size=5)
-        content = "".join(d.content for d in deltas if d.content is not None)
-        assert content == ""
-        args_str = self._reconstruct_args(deltas)
-        assert json.loads(args_str) == {"k": "v"}
-
-    def test_no_marker_leak_with_prefix_chunked(self, parser):
-        """Content before a tool call must not include start-marker
-        fragments when chunked (GitHub #40801)."""
-        full_text = "Hello!" + build_tool_call("fn", {"a": "b"})
-        deltas = self._stream_chunked(parser, full_text, chunk_size=5)
-        content = "".join(d.content for d in deltas if d.content is not None)
-        assert content == "Hello!"
-        assert "DSML" not in content
-        assert "<｜" not in content
-        args_str = self._reconstruct_args(deltas)
-        assert json.loads(args_str) == {"a": "b"}
-
-    def test_no_marker_leak_char_by_char(self, parser):
-        """Character-by-character streaming must not leak marker
-        fragments (GitHub #40801)."""
-        full_text = build_tool_call("fn", {"k": "v"})
-        deltas = self._stream_chunked(parser, full_text, chunk_size=1)
-        content = "".join(d.content for d in deltas if d.content is not None)
-        assert content == ""
-        args_str = self._reconstruct_args(deltas)
-        assert json.loads(args_str) == {"k": "v"}
-
-    def test_no_marker_leak_all_split_points(self, parser):
-        """Start token split at every possible boundary must not
-        leak (GitHub #40801)."""
-        for chunk_size in range(1, len(FC_START) + 2):
-            p = make_parser()
-            full_text = build_tool_call("fn", {"k": "v"})
-            deltas = self._stream_chunked(p, full_text, chunk_size=chunk_size)
-            content = "".join(d.content for d in deltas if d.content is not None)
-            assert content == "", (
-                f"Leaked content {content!r} at chunk_size={chunk_size}"
-            )
-
-    def test_false_partial_marker_emitted(self, parser):
-        """Text ending with a prefix of the start token that turns out
-        NOT to be a marker must still be emitted as content."""
-        full_text = "<｜DSM some regular text"
-        deltas = self._stream_chunked(parser, full_text, chunk_size=3)
-        content = "".join(d.content for d in deltas if d.content is not None)
-        assert content == full_text
-
 
 class TestDelimiterPreservation:
     """Regression: fast detokenization skipping DSML delimiters (PR #33964)."""