take 2

Signed-off-by: Tyler Michael Smith <tysmith@redhat.com>
2025-07-11 14:42:44 +00:00
154 changed files with 1516 additions and 10782 deletions
@@ -563,8 +563,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
  cuda_archs_loose_intersection(MLA_ARCHS "10.0a" "${CUDA_ARCHS}")
  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND MLA_ARCHS)
    set(SRCS
-      "csrc/attention/mla/cutlass_mla_kernels.cu"
-      "csrc/attention/mla/sm100_cutlass_mla_kernel.cu")
+      "csrc/attention/mla/cutlass_mla_kernels.cu")
    set_gencode_flags_for_srcs(
      SRCS "${SRCS}"
      CUDA_ARCHS "${MLA_ARCHS}")
@@ -69,7 +69,7 @@ vLLM is flexible and easy to use with:

 - Seamless integration with popular Hugging Face models
 - High-throughput serving with various decoding algorithms, including *parallel sampling*, *beam search*, and more
- Tensor, pipeline, data and expert parallelism support for distributed inference
+- Tensor parallelism and pipeline parallelism support for distributed inference
 - Streaming outputs
 - OpenAI-compatible API server
 - Support NVIDIA GPUs, AMD CPUs and GPUs, Intel CPUs and GPUs, PowerPC CPUs, TPU, and AWS Neuron
@@ -324,9 +324,6 @@ class RandomDataset(BenchmarkDataset):
        input_low = int(real_input_len * (1 - range_ratio))
        input_high = int(real_input_len * (1 + range_ratio))
        output_low = int(output_len * (1 - range_ratio))
-        # Ensure the lower bound for output length is at least 1 to prevent
-        # sampling 0 tokens, which can cause request failures.
-        output_low = max(output_low, 1)
        output_high = int(output_len * (1 + range_ratio))

        # Add logging for debugging
@@ -1,98 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-import itertools
-from typing import Callable
-
-import torch
-
-from vllm import _custom_ops as ops
-from vllm.config import CompilationConfig, VllmConfig, set_current_vllm_config
-from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8
-from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape
-from vllm.triton_utils import triton
-
-
-# TODO(luka): use standalone_compile utility
-def with_dyn_arg(fn: Callable, arg_index: int, dim_index: int):
-    def inner(*args):
-        torch._dynamo.mark_dynamic(args[arg_index], dim_index)
-        return fn(*args)
-
-    return inner
-
-
-torch._dynamo.config.recompile_limit = 8888
-compilation_config = CompilationConfig(custom_ops=["none"])
-with set_current_vllm_config(VllmConfig(compilation_config=compilation_config)):
-    torch_per_token_quant_fp8 = torch.compile(
-        QuantFP8(False, GroupShape.PER_TOKEN),
-        fullgraph=True,
-        dynamic=False,  # recompile for different shapes
-    )
-
-    # First dim is explicitly dynamic to simulate vLLM usage
-    torch_per_token_quant_fp8 = with_dyn_arg(torch_per_token_quant_fp8, 0, 0)
-
-
-def cuda_per_token_quant_fp8(
-    input: torch.Tensor,
-) -> tuple[torch.Tensor, torch.Tensor]:
-    return ops.scaled_fp8_quant(input)
-
-
-def calculate_diff(batch_size: int, seq_len: int):
-    """Calculate difference between Triton and CUDA implementations."""
-    device = torch.device("cuda")
-    x = torch.rand((batch_size * seq_len, 4096), dtype=torch.float16, device=device)
-
-    torch_out, torch_scale = torch_per_token_quant_fp8(x)
-    cuda_out, cuda_scale = cuda_per_token_quant_fp8(x)
-
-    if torch.allclose(
-        cuda_out.to(torch.float32), torch_out.to(torch.float32), rtol=1e-3, atol=1e-5
-    ) and torch.allclose(cuda_scale, torch_scale, rtol=1e-3, atol=1e-5):
-        print("✅ All implementations match")
-    else:
-        print("❌ Implementations differ")
-
-
-batch_size_range = [1, 16, 32, 64, 128]
-seq_len_range = [1, 16, 64, 128, 256, 512, 1024, 2048, 4096]
-
-configs = list(itertools.product(batch_size_range, seq_len_range))
-
-
-@triton.testing.perf_report(
-    triton.testing.Benchmark(
-        x_names=["batch_size", "seq_len"],
-        x_vals=configs,
-        line_arg="provider",
-        line_vals=["torch", "cuda"],
-        line_names=["Torch", "CUDA"],
-        styles=[("blue", "-"), ("green", "-")],
-        ylabel="us",
-        plot_name="per-token-dynamic-quant-fp8-performance",
-        args={},
-    )
-)
-def benchmark_quantization(batch_size, seq_len, provider):
-    dtype = torch.float16
-    device = torch.device("cuda")
-
-    x = torch.randn(batch_size * seq_len, 4096, device=device, dtype=dtype)
-
-    quantiles = [0.5, 0.2, 0.8]
-
-    if provider == "torch":
-        fn = lambda: torch_per_token_quant_fp8(x.clone())
-    elif provider == "cuda":
-        fn = lambda: cuda_per_token_quant_fp8(x.clone())
-
-    ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(fn, quantiles=quantiles)
-
-    return 1000 * ms, 1000 * max_ms, 1000 * min_ms
-
-
-if __name__ == "__main__":
-    calculate_diff(batch_size=4, seq_len=4096)
-    benchmark_quantization.run(print_data=True)
@@ -86,9 +86,6 @@ def benchmark_config(
            (num_experts, 2 * shard_intermediate_size), dtype=torch.float32
        )
        w2_scale = torch.randn((hidden_size, num_experts), dtype=torch.float32)
-    if use_deep_gemm:
-        # we use the default block shape for deepgemm
-        block_quant_shape = [128, 128]
    if use_fp8_w8a8:
        if block_quant_shape:
            block_n, block_k = block_quant_shape[0], block_quant_shape[1]
@@ -1,240 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import csv
-import os
-import random
-from datetime import datetime
-
-import flashinfer
-import torch
-
-FLOAT32_BYTES = torch.finfo(torch.float).bits // 8
-
-# KV Cache Layout for TRT-LLM
-# kv_cache_shape = (num_blocks, 2, num_kv_heads, page_size, head_dim)
-
-
-def to_float8(x, dtype=torch.float8_e4m3fn):
-    finfo = torch.finfo(dtype)
-    min_val, max_val = x.aminmax()
-    amax = torch.maximum(min_val.abs(), max_val.abs()).clamp(min=1e-12)
-    scale = finfo.max / amax * 0.1
-    x_scl_sat = (x * scale).clamp(min=finfo.min, max=finfo.max)
-    return x_scl_sat.to(dtype), scale.float().reciprocal()
-
-
-@torch.no_grad()
-def benchmark_decode(
-    num_seqs,
-    max_seq_len,
-    page_size=16,
-    dtype=torch.bfloat16,
-    kv_layout="HND",
-    num_kv_heads=8,
-    kv_cache_dtype="auto",
-    head_dim=128,
-    warmup=10,
-    trials=20,
-):
-    torch.set_default_device("cuda")
-    device = "cuda"
-    torch.manual_seed(0)
-
-    # Currently only HEAD_GRP_SIZE == 8 is supported
-    HEAD_GRP_SIZE = 8
-    MAX_SEQ_LEN = max_seq_len
-
-    # large number to reduce kv_cache reuse
-    NUM_BLOCKS = int(256000 / page_size)
-
-    workspace_buffer = torch.empty(1024 * 1024 * 1024, dtype=torch.int8, device=device)
-
-    # For decode, batch_size is num_decode_token
-    num_qo_heads = num_kv_heads * HEAD_GRP_SIZE
-    sm_scale = float(1.0 / (head_dim**0.5))
-    q = torch.randn(num_seqs, num_qo_heads, head_dim, device=device, dtype=dtype)
-    kv_lens = [random.randint(1, MAX_SEQ_LEN) for _ in range(num_seqs)]
-
-    max_kv_len = max(kv_lens)
-    kv_lens_tensor = torch.tensor(kv_lens, dtype=torch.int, device=device)
-    max_num_blocks_per_seq = (max_kv_len + page_size - 1) // page_size
-
-    block_tables = torch.randint(
-        0, NUM_BLOCKS, (num_seqs, max_num_blocks_per_seq), dtype=torch.int32
-    )
-
-    kv_cache_shape = (NUM_BLOCKS, 2, num_kv_heads, page_size, head_dim)
-    kv_cache = torch.randn(size=kv_cache_shape, device=device, dtype=dtype)
-    k_scale = v_scale = 1.0
-
-    if kv_cache_dtype.startswith("fp8"):
-        kv_cache, _ = to_float8(kv_cache)
-
-    # Benchmark TRT decode
-    def trt_decode():
-        return flashinfer.decode.trtllm_batch_decode_with_kv_cache(
-            q,
-            kv_cache,
-            workspace_buffer,
-            num_qo_heads,
-            num_kv_heads,
-            sm_scale,
-            block_tables,
-            kv_lens_tensor,
-            page_size,
-            max_kv_len,
-            kv_cache_dtype,
-            k_scale,
-            v_scale,
-        )
-
-    def time_fn(fn, warmup=10, trials=20):
-        torch.cuda.synchronize()
-        start = torch.cuda.Event(enable_timing=True)
-        end = torch.cuda.Event(enable_timing=True)
-        times = []
-        for i in range(warmup):
-            fn()
-        for i in range(trials):
-            start.record()
-            fn()
-            end.record()
-            torch.cuda.synchronize()
-            times.append(start.elapsed_time(end))  # ms
-        return sum(times) / len(times), torch.std(torch.tensor(times))
-
-    # TRT Decode
-    trt_mean, trt_std = time_fn(trt_decode)
-
-    kv_indptr = [0]
-    kv_indices = []
-    kv_last_page_lens = []
-    for i in range(num_seqs):
-        seq_len = kv_lens[i]
-        assert seq_len > 0
-        num_blocks = (seq_len + page_size - 1) // page_size
-        kv_indices.extend(block_tables[i, :num_blocks])
-        kv_indptr.append(kv_indptr[-1] + num_blocks)
-        kv_last_page_len = seq_len % page_size
-        if kv_last_page_len == 0:
-            kv_last_page_len = page_size
-        kv_last_page_lens.append(kv_last_page_len)
-
-    kv_indptr = torch.tensor(kv_indptr, dtype=torch.int32)
-    kv_indices = torch.tensor(kv_indices, dtype=torch.int32)
-    kv_last_page_lens = torch.tensor(kv_last_page_lens, dtype=torch.int32)
-
-    wrapper = flashinfer.BatchDecodeWithPagedKVCacheWrapper(
-        workspace_buffer,
-        kv_layout,
-        use_tensor_cores=((num_qo_heads // num_kv_heads) > 4),
-    )
-
-    wrapper.plan(
-        kv_indptr,
-        kv_indices,
-        kv_last_page_lens,
-        num_qo_heads,
-        num_kv_heads,
-        head_dim,
-        page_size,
-        "NONE",
-        q_data_type=dtype,
-        kv_data_type=torch.float8_e4m3fn if kv_cache_dtype.startswith("fp8") else dtype,
-    )
-
-    def baseline_decode():
-        return wrapper.run(q, kv_cache, sm_scale, k_scale, v_scale)
-
-    baseline_mean, baseline_std = time_fn(baseline_decode)
-
-    # Calculate percentage speedup (positive means TRT is faster)
-    speedup_percent = (baseline_mean - trt_mean) / baseline_mean
-
-    print(
-        f"\t{num_seqs}\t{max_seq_len}\t{trt_mean:.3f}\t{trt_std.item():.3f}"
-        f"\t{baseline_mean:.3f}\t{baseline_std.item():.3f}\t{speedup_percent:.3f}"
-    )
-
-    # Return results for CSV writing
-    return {
-        "num_seqs": num_seqs,
-        "trt_mean": trt_mean,
-        "trt_std": trt_std.item(),
-        "baseline_mean": baseline_mean,
-        "baseline_std": baseline_std.item(),
-        "speedup_percent": speedup_percent,
-        "q_dtype": str(dtype),
-        "kv_cache_dtype": kv_cache_dtype,
-        "page_size": page_size,
-        "num_kv_heads": num_kv_heads,
-        "head_dim": head_dim,
-        "max_seq_len": max_seq_len,
-    }
-
-
-def write_results_to_csv(results, filename=None):
-    """Write benchmark results to CSV file."""
-    if filename is None:
-        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
-        filename = f"flashinfer_trtllm_benchmark_{timestamp}.csv"
-
-    fieldnames = [
-        "num_seqs",
-        "trt_mean",
-        "trt_std",
-        "baseline_mean",
-        "baseline_std",
-        "speedup_percent",
-        "q_dtype",
-        "kv_cache_dtype",
-        "page_size",
-        "num_kv_heads",
-        "head_dim",
-        "max_seq_len",
-    ]
-
-    file_exists = os.path.exists(filename)
-
-    with open(filename, "a", newline="") as csvfile:
-        writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
-
-        if not file_exists:
-            writer.writeheader()
-
-        for result in results:
-            writer.writerow(result)
-
-    print(f"Results written to {filename}")
-
-
-if __name__ == "__main__":
-    num_seqs = [1, 4, 8, 16, 32, 64, 128, 256]
-    max_seq_lens = [1024, 2048, 4096, 8192, 16384, 32768, 65536, 131072]
-    all_results = []
-
-    print("Running benchmark for kv_cache_dtype: bfloat16")
-    print(
-        "\tnum_seqs\tmax_seq_len\ttrt_mean\ttrt_std\tbaseline_mean\tbaseline_std\tspeedup_percent"
-    )
-    for max_seq_len in max_seq_lens:
-        for bs in num_seqs:
-            result = benchmark_decode(
-                bs, max_seq_len, dtype=torch.bfloat16, kv_cache_dtype="auto"
-            )
-            all_results.append(result)
-
-    print("Running benchmark for q_dtype = bfloat16, kv_cache_dtype: fp8")
-    print(
-        "\tnum_seqs\tmax_seq_len\ttrt_mean\ttrt_std\tbaseline_mean\tbaseline_std\tspeedup_percent"
-    )
-    for max_seq_len in max_seq_lens:
-        for bs in num_seqs:
-            result = benchmark_decode(
-                bs, max_seq_len, dtype=torch.bfloat16, kv_cache_dtype="fp8"
-            )
-            all_results.append(result)
-
-    # Write all results to CSV
-    write_results_to_csv(all_results)
@@ -1,364 +0,0 @@
-/***************************************************************************************************
- * Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
- * SPDX-License-Identifier: BSD-3-Clause
- *
- * Redistribution and use in source and binary forms, with or without
- * modification, are permitted provided that the following conditions are met:
- *
- * 1. Redistributions of source code must retain the above copyright notice,
- *this list of conditions and the following disclaimer.
- *
- * 2. Redistributions in binary form must reproduce the above copyright notice,
- * this list of conditions and the following disclaimer in the documentation
- * and/or other materials provided with the distribution.
- *
- * 3. Neither the name of the copyright holder nor the names of its
- * contributors may be used to endorse or promote products derived from
- * this software without specific prior written permission.
- *
- * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
- * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
- * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
- *ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
- *LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
- *CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
- *SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
- *INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
- *CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
- *ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
- *POSSIBILITY OF SUCH DAMAGE.
- *
- **************************************************************************************************/
-/*
- * Taken from SGLANG PR https://github.com/sgl-project/sglang/pull/6929
- * by Alcanderian JieXin Liang
- */
-
-/*!
- \file
- \brief An universal device layer for cutlass 3.x-style kernels.
-*/
-
-// clang-format off
-#pragma once
-
-// common
-#include "cutlass/cutlass.h"
-#include "cutlass/device_kernel.h"
-
-#if !defined(__CUDACC_RTC__)
-#include "cutlass/cluster_launch.hpp"
-#include "cutlass/trace.h"
-#endif // !defined(__CUDACC_RTC__)
-
-#include "../kernel/sm100_fmha_mla_tma_warpspecialized.hpp"
-#include "../kernel/sm100_fmha_mla_reduction.hpp"
-
-////////////////////////////////////////////////////////////////////////////////
-
-namespace cutlass::fmha::device {
-
-using namespace cute;
-using namespace cutlass::fmha::kernel;
-
-
-////////////////////////////////////////////////////////////////////////////////
-////////////////////////////// CUTLASS 3.x API /////////////////////////////////
-////////////////////////////////////////////////////////////////////////////////
-
-template<
-    class Kernel_
->
-class MLA {
-public:
-
-  using Kernel = Kernel_;
-
-  using ReductionKernel = cutlass::fmha::kernel::Sm100FmhaMlaReductionKernel<
-      typename Kernel::ElementOut,
-      typename Kernel::ElementAcc,
-      typename Kernel::ElementAcc,
-      Kernel::TileShapeH::value,
-      Kernel::TileShapeL::value,
-      256 /*Max split*/
-  >;
-
-  /// Argument structure: User API
-  using KernelArguments = typename Kernel::Arguments;
-  using ReductionArguments = typename ReductionKernel::Arguments;
-
-  using Arguments = KernelArguments;
-
-  /// Argument structure: Kernel API
-  using KernelParams = typename Kernel::Params;
-  using ReductionParams = typename ReductionKernel::Params;
-  struct Params {
-    KernelParams fmha_params;
-    ReductionParams reduction_params;
-  };
-
-private:
-
-  /// Kernel API parameters object
-  Params params_;
-
-  bool is_initialized(bool set = false) {
-    static bool initialized = false;
-    if (set) initialized = true;
-    return initialized;
-  }
-
-  static ReductionArguments to_reduction_args(Arguments const& args) {
-    auto [H, K, D, B] = args.problem_shape;
-    return ReductionArguments{
-      nullptr, args.epilogue.ptr_o, nullptr, args.epilogue.ptr_lse,
-      args.mainloop.softmax_scale, B, args.split_kv, K, args.mainloop.ptr_seq,
-      args.ptr_split_kv, Kernel::TileShapeS::value
-    };
-  }
-
-public:
-
-  /// Access the Params structure
-  Params const& params() const {
-    return params_;
-  }
-
-  static void set_split_kv (KernelArguments& args) {
-    if (args.split_kv >= 1) return;
-    auto [H, K, D, B] = args.problem_shape;
-    int sm_count = args.hw_info.sm_count;
-    int max_splits = ceil_div(K, 128);
-    int sms_per_batch = max(1, sm_count / B);
-    int split_heur = min(max_splits, sms_per_batch);
-    int waves = ceil_div(B * split_heur, sm_count);
-    int k_waves = ceil_div(max_splits, split_heur);
-    int split_wave_aware = ceil_div(max_splits, k_waves);
-    args.split_kv = split_wave_aware;
-  }
-
-  /// Determines whether the GEMM can execute the given problem.
-  static Status
-  can_implement(Arguments const& args) {
-    if (! Kernel::can_implement(args)) {
-      return Status::kInvalid;
-    }
-    if (! ReductionKernel::can_implement(to_reduction_args(args))) {
-      return Status::kInvalid;
-    }
-    return Status::kSuccess;
-  }
-
-  /// Gets the workspace size
-  static size_t
-  get_workspace_size(Arguments const& args) {
-    size_t workspace_bytes = 0;
-    workspace_bytes += Kernel::get_workspace_size(args);
-    workspace_bytes += ReductionKernel::get_workspace_size(to_reduction_args(args));
-    return workspace_bytes;
-  }
-
-  /// Computes the maximum number of active blocks per multiprocessor
-  static int maximum_active_blocks(int /* smem_capacity */ = -1) {
-    CUTLASS_TRACE_HOST("MLA::maximum_active_blocks()");
-    int max_active_blocks = -1;
-    int smem_size = Kernel::SharedStorageSize;
-
-    // first, account for dynamic smem capacity if needed
-    cudaError_t result;
-    if (smem_size >= (48 << 10)) {
-      CUTLASS_TRACE_HOST("  Setting smem size to " << smem_size);
-      result = cudaFuncSetAttribute(
-          device_kernel<Kernel>,
-          cudaFuncAttributeMaxDynamicSharedMemorySize,
-          smem_size);
-      if (cudaSuccess != result) {
-        result = cudaGetLastError(); // to clear the error bit
-        CUTLASS_TRACE_HOST(
-          "  cudaFuncSetAttribute() returned error: "
-          << cudaGetErrorString(result));
-        return -1;
-      }
-    }
-
-    // query occupancy after setting smem size
-    result = cudaOccupancyMaxActiveBlocksPerMultiprocessor(
-        &max_active_blocks,
-        device_kernel<Kernel>,
-        Kernel::MaxThreadsPerBlock,
-        smem_size);
-
-    if (cudaSuccess != result) {
-      result = cudaGetLastError(); // to clear the error bit
-      CUTLASS_TRACE_HOST(
-        "  cudaOccupancyMaxActiveBlocksPerMultiprocessor() returned error: "
-        << cudaGetErrorString(result));
-      return -1;
-    }
-
-    CUTLASS_TRACE_HOST("  max_active_blocks: " << max_active_blocks);
-    return max_active_blocks;
-  }
-
-  /// Initializes GEMM state from arguments.
-  Status
-  initialize(Arguments const& args, void* workspace = nullptr, cudaStream_t stream = nullptr) {
-    CUTLASS_TRACE_HOST("MLA::initialize() - workspace "
-      << workspace << ", stream: " << (stream ? "non-null" : "null"));
-
-    // Initialize the workspace
-    Status status = Kernel::initialize_workspace(args, workspace, stream);
-    if (status != Status::kSuccess) {
-      return status;
-    }
-    status = ReductionKernel::initialize_workspace(to_reduction_args(args), workspace, stream);
-    if (status != Status::kSuccess) {
-      return status;
-    }
-    KernelParams kernel_params = Kernel::to_underlying_arguments(args, workspace);
-
-    ReductionArguments reduction_args = to_reduction_args(args);
-    if (reduction_args.split_kv > 1) {
-      reduction_args.ptr_oaccum   = kernel_params.epilogue.ptr_o_acc;
-      reduction_args.ptr_lseaccum = kernel_params.epilogue.ptr_lse_acc;
-    }
-    ReductionParams reduction_params = ReductionKernel::to_underlying_arguments(reduction_args, workspace);
-    // Initialize the Params structure
-    params_ = Params {kernel_params, reduction_params};
-
-    if (is_initialized()) return Status::kSuccess;
-
-    // account for dynamic smem capacity if needed
-    // no dynamic smem is needed for reduction kernel
-    int smem_size = Kernel::SharedStorageSize;
-    if (smem_size >= (48 << 10)) {
-      CUTLASS_TRACE_HOST("  Setting smem size to " << smem_size);
-      cudaError_t result = cudaFuncSetAttribute(
-          device_kernel<Kernel>,
-          cudaFuncAttributeMaxDynamicSharedMemorySize,
-          smem_size);
-      if (cudaSuccess != result) {
-        result = cudaGetLastError(); // to clear the error bit
-        CUTLASS_TRACE_HOST("  cudaFuncSetAttribute() returned error: " << cudaGetErrorString(result));
-        return Status::kErrorInternal;
-      }
-    }
-
-    is_initialized(true);
-
-    return Status::kSuccess;
-  }
-
-  /// Update API is preserved in 3.0, but does not guarantee a lightweight update of params.
-  Status
-  update(Arguments const& args, void* workspace = nullptr) {
-    CUTLASS_TRACE_HOST("MLA()::update() - workspace: " << workspace);
-
-    size_t workspace_bytes = get_workspace_size(args);
-    if (workspace_bytes > 0 && nullptr == workspace) {
-      return Status::kErrorWorkspaceNull;
-    }
-
-    auto fmha_params = Kernel::to_underlying_arguments(args, workspace);
-
-    ReductionArguments reduction_args = to_reduction_args(args);
-    if (reduction_args.split_kv > 1) {
-      reduction_args.ptr_oaccum   = fmha_params.epilogue.ptr_o_acc;
-      reduction_args.ptr_lseaccum = fmha_params.epilogue.ptr_lse_acc;
-    }
-    ReductionParams reduction_params = ReductionKernel::to_underlying_arguments(reduction_args, workspace);
-    // Initialize the Params structure
-    params_ = Params {fmha_params, reduction_params};
-
-    return Status::kSuccess;
-  }
-
-  /// Primary run() entry point API that is static allowing users to create and manage their own params.
-  /// Supplied params struct must be construct by calling Kernel::to_underling_arguments()
-  static Status
-  run(Params& params, cudaStream_t stream = nullptr) {
-    CUTLASS_TRACE_HOST("MLA::run()");
-    dim3 const block = Kernel::get_block_shape();
-    dim3 const grid = Kernel::get_grid_shape(params.fmha_params);
-
-    // configure smem size and carveout
-    int smem_size = Kernel::SharedStorageSize;
-
-    Status launch_result;
-    // Use extended launch API only for mainloops that use it
-    if constexpr(Kernel::ArchTag::kMinComputeCapability >= 90) {
-      dim3 cluster(cute::size<0>(typename Kernel::ClusterShape{}),
-                   cute::size<1>(typename Kernel::ClusterShape{}),
-                   cute::size<2>(typename Kernel::ClusterShape{}));
-      void const* kernel = (void const*) device_kernel<Kernel>;
-      void* kernel_params[] = {&params.fmha_params};
-      launch_result = ClusterLauncher::launch(grid, cluster, block, smem_size, stream, kernel, kernel_params);
-    }
-    else {
-      launch_result = Status::kSuccess;
-      device_kernel<Kernel><<<grid, block, smem_size, stream>>>(params.fmha_params);
-    }
-
-    cudaError_t result = cudaGetLastError();
-    if (cudaSuccess != result or Status::kSuccess != launch_result) {
-      //return Status::kSuccess;
-      CUTLASS_TRACE_HOST("  Kernel launch failed. Reason: " << result);
-      return Status::kErrorInternal;
-    }
-    if (params.reduction_params.split_kv > 1) {
-      // launch reduction kernel
-      dim3 const block = ReductionKernel::get_block_shape();
-      dim3 const grid  = ReductionKernel::get_grid_shape(params.reduction_params);
-      device_kernel<ReductionKernel><<<grid, block, 0, stream>>>(params.reduction_params);
-      cudaError_t result = cudaGetLastError();
-      if (cudaSuccess == result) {
-        return Status::kSuccess;
-      }
-      else {
-        CUTLASS_TRACE_HOST("  Kernel launch failed. Reason: " << result);
-        return Status::kErrorInternal;
-      }
-    }
-    else {
-      return Status::kSuccess;
-    }
-  }
-
-  //
-  // Non-static launch overloads that first create and set the internal params struct of this kernel handle.
-  //
-
-  /// Launches the kernel after first constructing Params internal state from supplied arguments.
-  Status
-  run(Arguments const& args, void* workspace = nullptr, cudaStream_t stream = nullptr) {
-    Status status = initialize(args, workspace, stream);
-    if (Status::kSuccess == status) {
-      status = run(params_, stream);
-    }
-    return status;
-  }
-
-  /// Launches the kernel after first constructing Params internal state from supplied arguments.
-  Status
-  operator()(Arguments const& args, void* workspace = nullptr, cudaStream_t stream = nullptr) {
-    return run(args, workspace, stream);
-  }
-
-  /// Overload that allows a user to re-launch the same kernel without updating internal params struct.
-  Status
-  run(cudaStream_t stream = nullptr) {
-    return run(params_, stream);
-  }
-
-  /// Overload that allows a user to re-launch the same kernel without updating internal params struct.
-  Status
-  operator()(cudaStream_t stream = nullptr) {
-    return run(params_, stream);
-  }
-};
-
-////////////////////////////////////////////////////////////////////////////////
-
-} // namespace cutlass::fmha::device
-
-////////////////////////////////////////////////////////////////////////////////
@@ -1,203 +0,0 @@
-/***************************************************************************************************
- * Copyright (c) 2024 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights
- *reserved. SPDX-License-Identifier: BSD-3-Clause
- *
- * Redistribution and use in source and binary forms, with or without
- * modification, are permitted provided that the following conditions are met:
- *
- * 1. Redistributions of source code must retain the above copyright notice,
- *this list of conditions and the following disclaimer.
- *
- * 2. Redistributions in binary form must reproduce the above copyright notice,
- * this list of conditions and the following disclaimer in the documentation
- * and/or other materials provided with the distribution.
- *
- * 3. Neither the name of the copyright holder nor the names of its
- * contributors may be used to endorse or promote products derived from
- * this software without specific prior written permission.
- *
- * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
- * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
- * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
- *ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
- *LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
- *CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
- *SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
- *INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
- *CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
- *ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
- *POSSIBILITY OF SUCH DAMAGE.
- *
- **************************************************************************************************/
-/*
- * Taken from SGLANG PR https://github.com/sgl-project/sglang/pull/6929
- * by Alcanderian JieXin Liang
- */
-
-// clang-format off
-#pragma once
-
-#include "cutlass/cutlass.h"
-#include "cutlass/arch/arch.h"
-#include "cute/tensor.hpp"
-
-namespace cutlass::fmha::kernel {
-
-using namespace cute;
-template<
-    class ElementOut,
-    class ElementAcc,
-    class ElementScale,
-    size_t kNumHeads,
-    size_t kHeadDimLatent,
-    int kMaxSplits
->
-struct Sm100FmhaMlaReductionKernel {
-
-  static const int SharedStorageSize = 0;
-  static const int MaxThreadsPerBlock = 128;
-  static const int MinBlocksPerMultiprocessor = 1;
-
-  using ArchTag = cutlass::arch::Sm100;
-
-  static_assert(kHeadDimLatent % MaxThreadsPerBlock == 0);
-  struct Arguments {
-    ElementAcc* ptr_oaccum = nullptr;
-    ElementOut* ptr_o = nullptr;
-    ElementAcc* ptr_lseaccum = nullptr;
-    ElementAcc* ptr_lse = nullptr;
-    ElementScale scale = 1.f;
-    int num_batches = 0;
-    int split_kv = -1;
-    int dim_k = -1;
-    int* ptr_seq = nullptr;
-    int* ptr_split_kv = nullptr;
-    int tile_shape_s = 128;
-  };
-  using Params = Arguments;
-
-  static Params to_underlying_arguments(Arguments const& args, void* workspace) {
-    return {args.ptr_oaccum, args.ptr_o, args.ptr_lseaccum, args.ptr_lse,
-	    args.scale, args.num_batches, args.split_kv, args.dim_k, args.ptr_seq,
-	    args.ptr_split_kv, args.tile_shape_s};
-  }
-
-  static size_t get_workspace_size(Arguments const& /*args*/) {
-    return 0;
-  }
-
-  static Status initialize_workspace(
-      Arguments const& /*args*/, void* /*ws*/, cudaStream_t /*stream*/) {
-    return Status::kSuccess;
-  }
-
-  static dim3 get_grid_shape(Params const& params) {
-    return dim3(kNumHeads, 1, params.num_batches);
-  }
-
-  static dim3 get_block_shape() {
-    return dim3(MaxThreadsPerBlock, 1, 1);
-  }
-
-  static bool can_implement(Arguments const& args) {
-    if (args.num_batches <= 0) return false;
-    if (args.split_kv <= 0) return false;
-    return true;
-  }
-
-  CUTLASS_DEVICE void operator() (Params const& params, char* smem_raw) {
-    if (params.split_kv <= 1) return;
-    auto blk_coord = make_coord(blockIdx.x, _0{}, blockIdx.z);
-
-    __shared__ ElementAcc sLseScale[kMaxSplits];
-    const size_t offset_lseaccum = get<0>(blk_coord) + kNumHeads * params.split_kv * get<2>(blk_coord);
-    const size_t offset_lse = get<0>(blk_coord) + kNumHeads * get<2>(blk_coord);
-
-    Tensor gLSEaccum = make_tensor(make_gmem_ptr(params.ptr_lseaccum + offset_lseaccum),
-                                   make_shape(params.split_kv), Stride<Int<kNumHeads>>{});
-
-    Tensor gLSE = make_tensor(make_gmem_ptr(params.ptr_lse + offset_lse),
-                              Shape<_1>{}, Stride<_1>{});
-
-    auto dim_k = params.ptr_seq == nullptr ?  params.dim_k : params.ptr_seq[get<2>(blk_coord)];
-    auto local_split_kv = params.ptr_split_kv == nullptr ? params.split_kv : params.ptr_split_kv[get<2>(blk_coord)];
-    auto k_tile_total = ceil_div(dim_k, params.tile_shape_s);
-    auto k_tile_per_cta = ceil_div(k_tile_total, local_split_kv);
-    local_split_kv = ceil_div(k_tile_total, k_tile_per_cta);
-
-    int warp_idx = cutlass::canonical_warp_idx_sync();
-    if (warp_idx == 0) {
-      constexpr int kNLsePerThread = cute::ceil_div(kMaxSplits, 32);
-
-      ElementAcc local_lse[kNLsePerThread];
-
-      CUTLASS_PRAGMA_UNROLL
-      for (int i = 0; i < kNLsePerThread; ++i) {
-        const int split = i * 32 + threadIdx.x;
-        local_lse[i] = split < local_split_kv ? gLSEaccum(split) : -std::numeric_limits<ElementAcc>::infinity();
-      }
-
-      ElementAcc lse_max = -std::numeric_limits<ElementAcc>::infinity();
-      CUTLASS_PRAGMA_UNROLL
-      for (int i = 0; i < kNLsePerThread; ++i) {
-        lse_max = max(lse_max, local_lse[i]);
-      }
-      CUTLASS_PRAGMA_UNROLL
-      for (int offset = 16; offset >= 1; offset /= 2) {
-        lse_max = max(lse_max, __shfl_xor_sync(0xffffffff, lse_max, offset));
-      }
-      lse_max = lse_max == -std::numeric_limits<ElementAcc>::infinity() ? 0.0f : lse_max;  // In case all local LSEs are -inf
-      lse_max = __shfl_sync(0xffffffff, lse_max, 0);
-
-      ElementAcc sum_lse = 0;
-      CUTLASS_PRAGMA_UNROLL
-      for (int i = 0; i < kNLsePerThread; ++i) {
-        sum_lse = sum_lse + expf(local_lse[i] - lse_max);
-      }
-
-      CUTLASS_PRAGMA_UNROLL
-      for (int offset = 16; offset >= 1; offset /= 2) {
-        sum_lse = sum_lse + __shfl_xor_sync(0xffffffff, sum_lse, offset);
-      }
-
-      sum_lse = __shfl_sync(0xffffffff, sum_lse, 0);
-
-      ElementAcc global_lse = (sum_lse == 0.f || sum_lse != sum_lse) ? std::numeric_limits<ElementAcc>::infinity() : logf(sum_lse) + lse_max;
-      if (threadIdx.x == 0 and params.ptr_lse != nullptr) {
-        gLSE(0) = global_lse;
-      }
-
-      CUTLASS_PRAGMA_UNROLL
-      for (int i = 0; i < kNLsePerThread; ++i) {
-        const int split = i * 32 + threadIdx.x;
-        if (split < local_split_kv) {
-          sLseScale[split] = expf(local_lse[i] - global_lse);
-        }
-      }
-    }
-    __syncthreads();
-
-    constexpr int Elements = kHeadDimLatent / MaxThreadsPerBlock;
-    const size_t offset_oaccum = kHeadDimLatent * params.split_kv * (get<0>(blk_coord) + kNumHeads * get<2>(blk_coord));
-    Tensor gOaccum = make_tensor(make_gmem_ptr(params.ptr_oaccum + offset_oaccum),
-                               Shape<Int<kHeadDimLatent>>{}, Stride<_1>{});
-    ElementAcc local_val[Elements] = {0};
-    for (int split = 0; split < local_split_kv; ++split) {
-      ElementAcc lse_scale = sLseScale[split];
-      CUTLASS_PRAGMA_UNROLL
-      for(int i = 0; i < Elements; ++i) {
-        local_val[i] += lse_scale * gOaccum(threadIdx.x + MaxThreadsPerBlock * i);
-      }
-      gOaccum.data() = gOaccum.data() + kHeadDimLatent;
-    }
-    auto ptr_o_local = params.ptr_o + (get<0>(blk_coord) + get<2>(blk_coord) * kNumHeads) * kHeadDimLatent;
-    Tensor gO = make_tensor(make_gmem_ptr(ptr_o_local), Shape<Int<kHeadDimLatent>>{}, Stride<_1>{});
-
-    CUTLASS_PRAGMA_UNROLL
-    for(int i = 0; i < Elements; ++i) {
-      gO(threadIdx.x + MaxThreadsPerBlock * i) = static_cast<ElementOut>(local_val[i]);
-    }
-  }
-};
-
-}  // namespace cutlass::fmha::kernel
@@ -1,165 +0,0 @@
-/***************************************************************************************************
- * Copyright (c) 2024 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights
- *reserved. SPDX-License-Identifier: BSD-3-Clause
- *
- * Redistribution and use in source and binary forms, with or without
- * modification, are permitted provided that the following conditions are met:
- *
- * 1. Redistributions of source code must retain the above copyright notice,
- *this list of conditions and the following disclaimer.
- *
- * 2. Redistributions in binary form must reproduce the above copyright notice,
- * this list of conditions and the following disclaimer in the documentation
- * and/or other materials provided with the distribution.
- *
- * 3. Neither the name of the copyright holder nor the names of its
- * contributors may be used to endorse or promote products derived from
- * this software without specific prior written permission.
- *
- * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
- * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
- * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
- *ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
- *LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
- *CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
- *SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
- *INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
- *CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
- *ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
- *POSSIBILITY OF SUCH DAMAGE.
- *
- **************************************************************************************************/
-/*
- * Taken from SGLANG PR https://github.com/sgl-project/sglang/pull/6929
- * by Alcanderian JieXin Liang
- */
-
-// clang-format off
-#pragma once
-
-#include "cutlass/cutlass.h"
-#include "cutlass/fast_math.h"
-#include "cutlass/kernel_hardware_info.h"
-
-namespace cutlass::fmha::kernel {
-
-////////////////////////////////////////////////////////////////////////////////
-
-struct Sm100MlaIndividualTileScheduler {
-
-  struct Params {
-    dim3 grid;
-  };
-
-  bool valid_ = true;
-
-  CUTLASS_DEVICE
-  Sm100MlaIndividualTileScheduler(Params const&) {}
-
-  template<class ProblemShape, class ClusterShape>
-  static Params to_underlying_arguments(
-      ProblemShape const& problem_shape, KernelHardwareInfo hw_info,
-      ClusterShape const& cluster_shape, int const& split_kv) {
-    using namespace cute;
-    dim3 grid(get<0>(cluster_shape), get<3>(problem_shape) /* Batch */, split_kv /*Maximum Split KV*/);
-    return Params{ grid };
-  }
-
-  static dim3 get_grid_shape(Params const& params) {
-    return params.grid;
-  }
-
-  CUTLASS_DEVICE
-  bool is_valid() {
-    return valid_;
-  }
-
-  CUTLASS_DEVICE
-  auto get_block_coord() {
-    using namespace cute;
-    return make_coord(blockIdx.x, _0{}, blockIdx.y, blockIdx.z);
-  }
-
-  CUTLASS_DEVICE
-  Sm100MlaIndividualTileScheduler& operator++() {
-    valid_ = false;
-    return *this;
-  }
-};
-
-////////////////////////////////////////////////////////////////////////////////
-
-struct Sm100MlaPersistentTileScheduler {
-
-  struct Params {
-    int num_blocks;
-    FastDivmod divmod_m_block;
-    FastDivmod divmod_b;
-    FastDivmod divmod_split_kv;
-    KernelHardwareInfo hw_info;
-  };
-
-  int block_idx = 0;
-  Params params;
-
-  CUTLASS_DEVICE
-  Sm100MlaPersistentTileScheduler(Params const& params) : block_idx(blockIdx.x), params(params) {}
-
-  template<class ProblemShape, class ClusterShape>
-  static Params to_underlying_arguments(
-      ProblemShape const& problem_shape, KernelHardwareInfo hw_info,
-      ClusterShape const& cluster_shape, int const& split_kv) {
-    using namespace cute;
-    // Get SM count if needed, otherwise use user supplied SM count
-    int sm_count = hw_info.sm_count;
-    if (sm_count <= 1 || sm_count % size<0>(cluster_shape) != 0) {
-      CUTLASS_TRACE_HOST("  WARNING: Arguments do not include a valid SM count.\n"
-          "  For optimal performance, populate the arguments KernelHardwareInfo struct with the SM count.");
-      sm_count = KernelHardwareInfo::query_device_multiprocessor_count(hw_info.device_id);
-    }
-
-    CUTLASS_TRACE_HOST("to_underlying_arguments(): Setting persistent grid SM count to " << sm_count);
-    hw_info.sm_count = sm_count;
-
-    int num_m_blocks = size<0>(cluster_shape);
-    int num_blocks = num_m_blocks * get<3>(problem_shape)  /* Batch */;
-    num_blocks *= split_kv; /* Maximum Split KV*/
-
-    return Params {
-      num_blocks,
-      { num_m_blocks}, { get<3>(problem_shape) }, {split_kv},
-      hw_info
-    };
-  }
-
-  static dim3 get_grid_shape(Params const& params) {
-    dim3 grid(std::min(params.num_blocks, params.hw_info.sm_count), 1, 1);
-    return grid;
-  }
-
-  CUTLASS_DEVICE
-  bool is_valid() {
-    return block_idx < params.num_blocks;
-  }
-
-  CUTLASS_DEVICE
-  auto get_block_coord() {
-    using namespace cute;
-    int block_decode = block_idx;
-    int m_block, bidb, n_split_kv;
-    params.divmod_m_block(block_decode, m_block, block_decode);
-    params.divmod_b(block_decode, bidb, block_decode);
-    params.divmod_split_kv(block_decode, n_split_kv, block_decode);
-    return make_coord(m_block, _0{}, bidb, n_split_kv);
-  }
-
-  CUTLASS_DEVICE
-  Sm100MlaPersistentTileScheduler& operator++() {
-    block_idx += gridDim.x;
-    return *this;
-  }
-};
-
-////////////////////////////////////////////////////////////////////////////////
-
-} // namespace cutlass::fmha::kernel
@@ -1,273 +0,0 @@
-/*
-Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
-Copyright 2025 SGLang Team. All Rights Reserved.
-
-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.
-==============================================================================*/
-/*
- * Taken from SGLANG PR https://github.com/sgl-project/sglang/pull/6929
- * by Alcanderian JieXin Liang
- */
-
-#include <ATen/cuda/CUDAContext.h>
-#include <c10/cuda/CUDAGuard.h>
-#include <cutlass/cutlass.h>
-#include <cutlass/kernel_hardware_info.h>
-#include <torch/all.h>
-
-#include <cute/tensor.hpp>
-#include <iostream>
-
-#include "cutlass_sm100_mla/device/sm100_mla.hpp"
-#include "cutlass_sm100_mla/kernel/sm100_mla_tile_scheduler.hpp"
-
-// clang-format off
-#if !defined(CUDA_VERSION) || CUDA_VERSION < 12040
-void sm100_cutlass_mla_decode(
-    torch::Tensor const& out,
-    torch::Tensor const& q_nope,
-    torch::Tensor const& q_pe,
-    torch::Tensor const& kv_c_and_k_pe_cache,
-    torch::Tensor const& seq_lens,
-    torch::Tensor const& page_table,
-    torch::Tensor const& workspace,
-    int64_t num_kv_splits) {
-  TORCH_CHECK(false, "CUDA version must be >= 12.4 for cutlass_mla_decode");
-}
-int64_t sm100_cutlass_mla_get_workspace_size(int64_t max_seq_len, int64_t num_batches, int64_t sm_count, int64_t num_kv_splits) {
-  TORCH_CHECK(false, "CUDA version must be >= 12.4 for cutlass_mla_get_workspace_size");
-}
-#else
-
-#define CUTLASS_CHECK(status)                                                       \
-  {                                                                                 \
-    cutlass::Status error = status;                                                 \
-    TORCH_CHECK(error == cutlass::Status::kSuccess, cutlassGetStatusString(error)); \
-  }
-
-using namespace cute;
-using namespace cutlass::fmha::kernel;
-
-template <bool v>
-struct IsPersistent {
-  static const bool value = v;
-};
-
-template <typename T, bool IsPaged128, typename PersistenceOption = IsPersistent<true>>
-struct MlaSm100 {
-  using Element = T;
-  using ElementAcc = float;
-  using ElementOut = T;
-
-  using TileShape = Shape<_128, _128, Shape<_512, _64>>;
-  using TileShapeH = cute::tuple_element_t<0, TileShape>;
-  using TileShapeD = cute::tuple_element_t<2, TileShape>;
-
-  // H K (D_latent D_rope) B
-  using ProblemShape = cute::tuple<TileShapeH, int, TileShapeD, int>;
-
-  using StrideQ = cute::tuple<int64_t, _1, int64_t>;  // H D B
-  using StrideK = cute::tuple<int64_t, _1, int64_t>;  // K D B
-  using StrideO = StrideK;                            // H D B
-  using StrideLSE = cute::tuple<_1, int>;             // H B
-
-  using TileScheduler =
-      std::conditional_t<PersistenceOption::value, Sm100MlaPersistentTileScheduler, Sm100MlaIndividualTileScheduler>;
-
-  using FmhaKernel = cutlass::fmha::kernel::Sm100FmhaMlaKernelTmaWarpspecialized<
-      TileShape,
-      Element,
-      ElementAcc,
-      ElementOut,
-      ElementAcc,
-      TileScheduler,
-      /*kIsCpAsync=*/!IsPaged128>;
-  using Fmha = cutlass::fmha::device::MLA<FmhaKernel>;
-};
-
-template <typename T>
-typename T::Fmha::Arguments args_from_options(
-    at::Tensor const& out,
-    at::Tensor const& q_nope,
-    at::Tensor const& q_pe,
-    at::Tensor const& kv_c_and_k_pe_cache,
-    at::Tensor const& seq_lens,
-    at::Tensor const& page_table,
-    double sm_scale,
-    int64_t num_kv_splits) {
-  cutlass::KernelHardwareInfo hw_info;
-  hw_info.device_id = q_nope.device().index();
-  hw_info.sm_count = cutlass::KernelHardwareInfo::query_device_multiprocessor_count(hw_info.device_id);
-
-  int batches = q_nope.sizes()[0];
-  int page_count_per_seq = page_table.sizes()[1];
-  int page_count_total = kv_c_and_k_pe_cache.sizes()[0];
-  int page_size = kv_c_and_k_pe_cache.sizes()[1];
-  int max_seq_len = page_size * page_count_per_seq;
-  using TileShapeH = typename T::TileShapeH;
-  using TileShapeD = typename T::TileShapeD;
-  auto problem_shape = cute::make_tuple(TileShapeH{}, max_seq_len, TileShapeD{}, batches);
-
-  auto [H, K, D, B] = problem_shape;
-  auto [D_latent, D_rope] = D;
-
-  float scale = float(sm_scale);
-
-  using StrideQ = typename T::StrideQ;
-  using StrideK = typename T::StrideK;
-  using StrideO = typename T::StrideO;
-  using StrideLSE = typename T::StrideLSE;
-
-  StrideQ stride_Q_nope = cute::make_tuple(
-      static_cast<int64_t>(q_nope.stride(1)), _1{}, static_cast<int64_t>(q_nope.stride(0)));
-  StrideQ stride_Q_pe = cute::make_tuple(
-      static_cast<int64_t>(q_pe.stride(1)), _1{}, static_cast<int64_t>(q_pe.stride(0)));
-
-  StrideK stride_C = cute::make_tuple(
-      static_cast<int64_t>(0 + D_latent + D_rope), _1{}, static_cast<int64_t>(page_size * (D_latent + D_rope)));
-  StrideLSE stride_PT = cute::make_stride(_1{}, page_count_per_seq);
-  StrideLSE stride_LSE = cute::make_tuple(_1{}, 0 + H);
-  StrideO stride_O = cute::make_tuple(static_cast<int64_t>(0 + D_latent), _1{}, static_cast<int64_t>(0 + H * D_latent));
-
-  using Element = typename T::Element;
-  using ElementOut = typename T::ElementOut;
-  using ElementAcc = typename T::ElementAcc;
-  auto Q_nope_ptr = static_cast<Element*>(q_nope.data_ptr());
-  auto Q_pe_ptr = static_cast<Element*>(q_pe.data_ptr());
-  auto C_ptr = static_cast<Element*>(kv_c_and_k_pe_cache.data_ptr());
-  typename T::Fmha::Arguments arguments{
-      problem_shape,
-      {scale,
-       Q_nope_ptr,
-       stride_Q_nope,
-       Q_pe_ptr,
-       stride_Q_pe,
-       C_ptr,
-       stride_C,
-       C_ptr + D_latent,
-       stride_C,
-       static_cast<int*>(seq_lens.data_ptr()),
-       static_cast<int*>(page_table.data_ptr()),
-       stride_PT,
-       page_count_total,
-       page_size},
-      {static_cast<ElementOut*>(out.data_ptr()), stride_O, static_cast<ElementAcc*>(nullptr), stride_LSE},
-      hw_info,
-      // TODO(trevor-m): Change split_kv back to -1 when
-      // https://github.com/NVIDIA/cutlass/issues/2274 is fixed. Split_kv=1 will
-      // perform worse with larger context length and smaller batch sizes.
-      num_kv_splits, // split_kv
-      nullptr,       // is_var_split_kv
-  };
-  // TODO(kaixih@nvidia): When split_kv=-1 and is_var_split_kv=false, we compute
-  // split_kv automatically based on batch size and sequence length to balance
-  // workload across available SMs. Consider using var_split_kv for manual
-  // control if needed.
-  T::Fmha::set_split_kv(arguments);
-  return arguments;
-}
-
-template <typename Element, bool IsPaged128, typename PersistenceOption>
-void runMla(
-    at::Tensor const& out,
-    at::Tensor const& q_nope,
-    at::Tensor const& q_pe,
-    at::Tensor const& kv_c_and_k_pe_cache,
-    at::Tensor const& seq_lens,
-    at::Tensor const& page_table,
-    at::Tensor const& workspace,
-    double sm_scale,
-    int64_t num_kv_splits,
-    cudaStream_t stream) {
-  using MlaSm100Type = MlaSm100<Element, IsPaged128, PersistenceOption>;
-  typename MlaSm100Type::Fmha fmha;
-  auto arguments = args_from_options<MlaSm100Type>(out, q_nope, q_pe, kv_c_and_k_pe_cache, seq_lens, page_table, sm_scale, num_kv_splits);
-
-  CUTLASS_CHECK(fmha.can_implement(arguments));
-
-  CUTLASS_CHECK(fmha.initialize(arguments, workspace.data_ptr(), stream));
-
-  CUTLASS_CHECK(fmha.run(arguments, workspace.data_ptr(), stream));
-}
-
-#define DISPATCH_BOOL(expr, const_expr, ...) \
-  [&]() -> bool {                            \
-    if (expr) {                              \
-      constexpr bool const_expr = true;      \
-      return __VA_ARGS__();                  \
-    } else {                                 \
-      constexpr bool const_expr = false;     \
-      return __VA_ARGS__();                  \
-    }                                        \
-  }()
-
-void sm100_cutlass_mla_decode(
-    torch::Tensor const& out,
-    torch::Tensor const& q_nope,
-    torch::Tensor const& q_pe,
-    torch::Tensor const& kv_c_and_k_pe_cache,
-    torch::Tensor const& seq_lens,
-    torch::Tensor const& page_table,
-    torch::Tensor const& workspace,
-    double sm_scale,
-    int64_t num_kv_splits) {
-  auto in_dtype = q_nope.dtype();
-  at::cuda::CUDAGuard device_guard{(char)q_nope.get_device()};
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream(q_nope.get_device());
-  const int page_size = kv_c_and_k_pe_cache.sizes()[1];
-  
-  // NOTE(alcanderian): IsPersistent has bug with manual split_kv.
-  // Kernel will hang if batch is too large with large num_kv_splits. (for example bs=8, num_kv_splits=8)
-  // Maybe per batch split kv will fix this.
-  DISPATCH_BOOL(page_size == 128, IsPaged128, [&] {
-    DISPATCH_BOOL(num_kv_splits <= 1, NotManualSplitKV, [&] {
-      if (in_dtype == at::ScalarType::Half) {
-        runMla<cutlass::half_t, IsPaged128, IsPersistent<NotManualSplitKV>>(
-          out, q_nope, q_pe, kv_c_and_k_pe_cache, seq_lens, page_table, workspace, sm_scale, num_kv_splits, stream);
-      } else if (in_dtype == at::ScalarType::BFloat16) {
-        runMla<cutlass::bfloat16_t, IsPaged128, IsPersistent<NotManualSplitKV>>(
-          out, q_nope, q_pe, kv_c_and_k_pe_cache, seq_lens, page_table, workspace, sm_scale, num_kv_splits, stream);
-      } else if (in_dtype == at::ScalarType::Float8_e4m3fn) {
-        runMla<cutlass::float_e4m3_t, IsPaged128, IsPersistent<NotManualSplitKV>>(
-          out, q_nope, q_pe, kv_c_and_k_pe_cache, seq_lens, page_table, workspace, sm_scale, num_kv_splits, stream);
-      } else {
-        TORCH_CHECK(false, "Unsupported input data type of MLA");
-      }
-      return true;
-    });
-    return true;
-  });
-}
-
-int64_t sm100_cutlass_mla_get_workspace_size(int64_t max_seq_len, int64_t num_batches, int64_t sm_count, int64_t num_kv_splits) {
-  // Workspace size depends on ElementAcc and ElementLSE (same as ElementAcc)
-  // which are float, so Element type here doesn't matter.
-  using MlaSm100Type = MlaSm100<cutlass::half_t, true>;
-
-  // Get split kv. Requires problem shape and sm_count only.
-  typename MlaSm100Type::Fmha::Arguments arguments;
-  using TileShapeH = typename MlaSm100Type::TileShapeH;
-  using TileShapeD = typename MlaSm100Type::TileShapeD;
-  arguments.problem_shape =
-      cute::make_tuple(TileShapeH{}, static_cast<int>(max_seq_len), TileShapeD{}, static_cast<int>(num_batches));
-  // Assumes device 0 when getting sm_count.
-  arguments.hw_info.sm_count =
-      sm_count <= 0 ? cutlass::KernelHardwareInfo::query_device_multiprocessor_count(/*device_id=*/0) : sm_count;
-  arguments.split_kv = num_kv_splits;
-  MlaSm100Type::Fmha::set_split_kv(arguments);
-
-  return MlaSm100Type::Fmha::get_workspace_size(arguments);
-}
-
-#endif
-// clang-format on
@@ -167,19 +167,6 @@ void cutlass_mla_decode(torch::Tensor const& out, torch::Tensor const& q_nope,
                        torch::Tensor const& seq_lens,
                        torch::Tensor const& page_table, double scale);

-void sm100_cutlass_mla_decode(
-    torch::Tensor const& out, torch::Tensor const& q_nope,
-    torch::Tensor const& q_pe, torch::Tensor const& kv_c_and_k_pe_cache,
-    torch::Tensor const& seq_lens, torch::Tensor const& page_table,
-    torch::Tensor const& workspace, double sm_scale,
-    int64_t num_kv_splits =
-        1 /* Set to 1 to avoid cuda_graph issue by default. */);
-
-int64_t sm100_cutlass_mla_get_workspace_size(
-    int64_t max_seq_len, int64_t num_batches, int64_t sm_count = 0,
-    int64_t num_kv_splits =
-        1 /* Set to 1 to avoid cuda_graph issue by default. */);
-
 torch::Tensor get_cuda_view_from_cpu_tensor(torch::Tensor& cpu_tensor);

 #ifndef USE_ROCM
@@ -201,10 +201,11 @@ void run_blockwise_scaled_group_mm(
      reinterpret_cast<typename ScheduleConfig::LayoutSFB*>(
          layout_sfb.data_ptr())};

-  int device_id = a_ptrs.device().index();
-  static const cutlass::KernelHardwareInfo hw_info{
-      device_id, cutlass::KernelHardwareInfo::query_device_multiprocessor_count(
-                     device_id)};
+  cutlass::KernelHardwareInfo hw_info;
+  hw_info.device_id = a_ptrs.get_device();
+  hw_info.sm_count =
+      cutlass::KernelHardwareInfo::query_device_multiprocessor_count(
+          hw_info.device_id);

  // Epilogue Arguments
  typename GemmKernel::EpilogueArguments epilogue_args{
@@ -30,40 +30,35 @@

 #include "cutlass/util/packed_stride.hpp"

-#include "core/math.hpp"
-
 using namespace cute;

 #if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED)
+// Kernel Perf config
+template <typename T>
+struct KernelTraits;

-// Configuration for M in (256, inf)
-struct sm100_fp4_config_default {
-  using KernelSchedule = cutlass::gemm::collective::KernelScheduleAuto;
-  using EpilogueSchedule = cutlass::epilogue::collective::EpilogueScheduleAuto;
-  using TileShape = Shape<_256, _256, _256>;
-  using ClusterShape = Shape<_2, _1, _1>;
-  using PerSmTileShape_MNK = Shape<_128, _256, _256>;
-};
-
-// Configuration for M in (16, 256]
-struct sm100_fp4_config_M256 {
-  using KernelSchedule = cutlass::gemm::collective::KernelScheduleAuto;
-  using EpilogueSchedule = cutlass::epilogue::collective::EpilogueScheduleAuto;
-  using TileShape = Shape<_256, _128, _256>;
-  using ClusterShape = Shape<_2, _1, _1>;
-  using PerSmTileShape_MNK = Shape<_128, _128, _256>;
-};
-
-// Configuration for M in [1, 16]
-struct sm100_fp4_config_M16 {
-  using KernelSchedule = cutlass::gemm::collective::KernelScheduleAuto;
-  using EpilogueSchedule = cutlass::epilogue::collective::EpilogueScheduleAuto;
-  using TileShape = Shape<_128, _128, _256>;
+template <>
+struct KernelTraits<float> {
+  using MmaTileShape = Shape<_128, _128, _256>;
  using ClusterShape = Shape<_1, _1, _1>;
  using PerSmTileShape_MNK = Shape<_128, _128, _256>;
 };

-template <typename Config, typename OutType>
+template <>
+struct KernelTraits<cutlass::half_t> {
+  using MmaTileShape = Shape<_256, _256, _256>;
+  using ClusterShape = Shape<_4, _4, _1>;
+  using PerSmTileShape_MNK = Shape<_128, _256, _256>;
+};
+
+template <>
+struct KernelTraits<cutlass::bfloat16_t> {
+  using MmaTileShape = Shape<_256, _256, _256>;
+  using ClusterShape = Shape<_4, _4, _1>;
+  using PerSmTileShape_MNK = Shape<_128, _256, _256>;
+};
+
+template <typename T>
 struct Fp4GemmSm100 {
  // A matrix configuration
  using ElementA = cutlass::nv_float4_t<cutlass::float_e2m1_t>;
@@ -76,22 +71,21 @@ struct Fp4GemmSm100 {
  static constexpr int AlignmentB = 32;

  // C/D matrix configuration
-  using ElementD = OutType;
-  using ElementC = OutType;
+  using ElementD = T;
+  using ElementC = T;
  using LayoutCTag = cutlass::layout::RowMajor;
  using LayoutDTag = cutlass::layout::RowMajor;
  static constexpr int AlignmentD = 128 / cutlass::sizeof_bits<ElementD>::value;
  static constexpr int AlignmentC = 128 / cutlass::sizeof_bits<ElementC>::value;
-
  // Kernel functional config
  using ElementAccumulator = float;
  using ArchTag = cutlass::arch::Sm100;
  using OperatorClass = cutlass::arch::OpClassBlockScaledTensorOp;

-  // Use config's tile shapes
-  using MmaTileShape = typename Config::TileShape;
-  using ClusterShape = typename Config::ClusterShape;
-  using PerSmTileShape_MNK = typename Config::PerSmTileShape_MNK;
+  // Kernel Perf config
+  using MmaTileShape = typename KernelTraits<T>::MmaTileShape;
+  using ClusterShape = typename KernelTraits<T>::ClusterShape;
+  using PerSmTileShape_MNK = typename KernelTraits<T>::PerSmTileShape_MNK;

  using CollectiveEpilogue =
      typename cutlass::epilogue::collective::CollectiveBuilder<
@@ -125,22 +119,22 @@ struct Fp4GemmSm100 {
  using LayoutD = decltype(cute::make_layout(make_shape(0, 0, 0), StrideD{}));
 };

-template <typename Config>
-typename Config::Gemm::Arguments args_from_options(
+template <typename T>
+typename T::Gemm::Arguments args_from_options(
    at::Tensor& D, at::Tensor const& A, at::Tensor const& B,
    at::Tensor const& A_sf, at::Tensor const& B_sf, at::Tensor const& alpha,
    int64_t M, int64_t N, int64_t K) {
-  using ElementA = typename Config::Gemm::ElementA;
-  using ElementB = typename Config::Gemm::ElementB;
+  using ElementA = typename T::Gemm::ElementA;
+  using ElementB = typename T::Gemm::ElementB;
  using ElementSFA = cutlass::float_ue4m3_t;
  using ElementSFB = cutlass::float_ue4m3_t;
-  using ElementD = typename Config::Gemm::ElementD;
+  using ElementD = typename T::Gemm::ElementD;
  using ElementCompute = float;
-  using StrideA = typename Config::StrideA;
-  using StrideB = typename Config::StrideB;
-  using StrideD = typename Config::StrideD;
-  using Sm100BlkScaledConfig = typename Config::Gemm::GemmKernel::
-      CollectiveMainloop::Sm1xxBlkScaledConfig;
+  using StrideA = typename T::StrideA;
+  using StrideB = typename T::StrideB;
+  using StrideD = typename T::StrideD;
+  using Sm100BlkScaledConfig =
+      typename T::Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig;

  int m = static_cast<int>(M);
  int n = static_cast<int>(N);
@@ -154,7 +148,7 @@ typename Config::Gemm::Arguments args_from_options(
  auto layout_SFB = Sm100BlkScaledConfig::tile_atom_to_shape_SFB(
      cute::make_shape(m, n, k, 1));

-  typename Config::Gemm::Arguments arguments{
+  typename T::Gemm::Arguments arguments{
      cutlass::gemm::GemmUniversalMode::kGemm,
      {m, n, k, 1},
      {// Mainloop arguments
@@ -173,17 +167,17 @@ typename Config::Gemm::Arguments args_from_options(
  return arguments;
 }

-template <typename Config>
+template <typename T>
 void runGemm(at::Tensor& D, at::Tensor const& A, at::Tensor const& B,
             at::Tensor const& A_sf, at::Tensor const& B_sf,
             at::Tensor const& alpha, int64_t m, int64_t n, int64_t k,
             cudaStream_t stream) {
-  typename Config::Gemm gemm;
+  typename Fp4GemmSm100<T>::Gemm gemm;

  auto arguments =
-      args_from_options<Config>(D, A, B, A_sf, B_sf, alpha, m, n, k);
+      args_from_options<Fp4GemmSm100<T>>(D, A, B, A_sf, B_sf, alpha, m, n, k);

-  size_t workspace_size = Config::Gemm::get_workspace_size(arguments);
+  size_t workspace_size = Fp4GemmSm100<T>::Gemm::get_workspace_size(arguments);
  auto const workspace_options =
      torch::TensorOptions().dtype(torch::kUInt8).device(A.device());
  auto workspace = torch::empty(workspace_size, workspace_options);
@@ -194,40 +188,12 @@ void runGemm(at::Tensor& D, at::Tensor const& A, at::Tensor const& B,

  CUTLASS_CHECK(gemm.run(arguments, workspace.data_ptr(), stream));
 }
-
-// Dispatch function to select appropriate config based on M
-template <typename OutType>
-void cutlass_fp4_gemm_dispatch(torch::Tensor& D, torch::Tensor const& A,
-                               torch::Tensor const& B,
-                               torch::Tensor const& A_sf,
-                               torch::Tensor const& B_sf,
-                               torch::Tensor const& alpha, int64_t m, int64_t n,
-                               int64_t k, cudaStream_t stream) {
-  uint32_t const mp2 = std::max(static_cast<uint32_t>(16), next_pow_2(m));
-
-  if (mp2 <= 16) {
-    // m in [1, 16]
-    runGemm<Fp4GemmSm100<sm100_fp4_config_M16, OutType>>(
-        D, A, B, A_sf, B_sf, alpha, m, n, k, stream);
-  } else if (mp2 <= 256) {
-    // m in (16, 256]
-    runGemm<Fp4GemmSm100<sm100_fp4_config_M256, OutType>>(
-        D, A, B, A_sf, B_sf, alpha, m, n, k, stream);
-  } else {
-    // m in (256, inf)
-    runGemm<Fp4GemmSm100<sm100_fp4_config_default, OutType>>(
-        D, A, B, A_sf, B_sf, alpha, m, n, k, stream);
-  }
-}
-
 #else
-template <typename OutType>
-void cutlass_fp4_gemm_dispatch(torch::Tensor& D, torch::Tensor const& A,
-                               torch::Tensor const& B,
-                               torch::Tensor const& A_sf,
-                               torch::Tensor const& B_sf,
-                               torch::Tensor const& alpha, int64_t m, int64_t n,
-                               int64_t k, cudaStream_t stream) {
+template <typename T>
+void runGemm(at::Tensor& D, at::Tensor const& A, at::Tensor const& B,
+             at::Tensor const& A_sf, at::Tensor const& B_sf,
+             at::Tensor const& alpha, int64_t m, int64_t n, int64_t k,
+             cudaStream_t stream) {
  TORCH_CHECK(false,
              "Unsupported CUTLASS version. Set VLLM_CUTLASS_SRC_DIR to "
              "a CUTLASS 3.8 source directory to enable support.");
@@ -305,13 +271,12 @@ void cutlass_scaled_fp4_mm_sm100a(torch::Tensor& D, torch::Tensor const& A,
  const cudaStream_t stream = at::cuda::getCurrentCUDAStream(A.get_device());

  if (out_dtype == at::ScalarType::Half) {
-    cutlass_fp4_gemm_dispatch<cutlass::half_t>(D, A, B, A_sf, B_sf, alpha, m, n,
-                                               k, stream);
+    runGemm<cutlass::half_t>(D, A, B, A_sf, B_sf, alpha, m, n, k, stream);
  } else if (out_dtype == at::ScalarType::BFloat16) {
-    cutlass_fp4_gemm_dispatch<cutlass::bfloat16_t>(D, A, B, A_sf, B_sf, alpha,
-                                                   m, n, k, stream);
+    runGemm<cutlass::bfloat16_t>(D, A, B, A_sf, B_sf, alpha, m, n, k, stream);
+  } else if (out_dtype == at::ScalarType::Float) {
+    runGemm<float>(D, A, B, A_sf, B_sf, alpha, m, n, k, stream);
  } else {
-    TORCH_CHECK(false, "Unsupported output data type of nvfp4 mm (", out_dtype,
-                ")");
+    TORCH_CHECK(false, "Unsupported output data type of nvfp4 mm");
  }
 }
@@ -514,23 +514,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
      "                   Tensor page_table, float scale) -> ()");
  ops.impl("cutlass_mla_decode", torch::kCUDA, &cutlass_mla_decode);

-  // SM100 CUTLASS MLA decode
-  ops.def(
-      "sm100_cutlass_mla_decode(Tensor! out, Tensor q_nope, Tensor q_pe,"
-      "                         Tensor kv_c_and_k_pe_cache, Tensor seq_lens,"
-      "                         Tensor page_table, Tensor workspace, float "
-      "scale,"
-      "                         int num_kv_splits) -> ()");
-  ops.impl("sm100_cutlass_mla_decode", torch::kCUDA, &sm100_cutlass_mla_decode);
-
-  // SM100 CUTLASS MLA workspace
-  ops.def(
-      "sm100_cutlass_mla_get_workspace_size(int max_seq_len, int num_batches,"
-      "                                     int sm_count, int num_kv_splits) "
-      "-> int");
-  ops.impl("sm100_cutlass_mla_get_workspace_size",
-           &sm100_cutlass_mla_get_workspace_size);
-
  // Compute NVFP4 block quantized tensor.
  ops.def(
      "scaled_fp4_quant(Tensor! output, Tensor input,"
@@ -408,8 +408,7 @@ RUN --mount=type=cache,target=/root/.cache/uv bash - <<'BASH'

          # Needed to build AOT kernels
          pushd flashinfer
-            TORCH_CUDA_ARCH_LIST="${FI_TORCH_CUDA_ARCH_LIST}" \
-              python3 -m flashinfer.aot
+            python3 -m flashinfer.aot
            TORCH_CUDA_ARCH_LIST="${FI_TORCH_CUDA_ARCH_LIST}" \
              uv pip install --system --no-build-isolation .
          popd
@@ -36,7 +36,7 @@ vLLM is flexible and easy to use with:

 - Seamless integration with popular HuggingFace models
 - High-throughput serving with various decoding algorithms, including *parallel sampling*, *beam search*, and more
- Tensor, pipeline, data and expert parallelism support for distributed inference
+- Tensor parallelism and pipeline parallelism support for distributed inference
 - Streaming outputs
 - OpenAI-compatible API server
 - Support NVIDIA GPUs, AMD CPUs and GPUs, Intel CPUs, Gaudi® accelerators and GPUs, IBM Power CPUs, TPU, and AWS Trainium and Inferentia Accelerators.
@@ -272,80 +272,3 @@ The new format of `--lora-modules` is mainly to support the display of parent mo
        ]
    }
    ```
-
-## Default LoRA Models For Multimodal Models
-
-Some models, e.g., [Granite Speech](https://huggingface.co/ibm-granite/granite-speech-3.3-8b) and [Phi-4-multimodal-instruct](https://huggingface.co/microsoft/Phi-4-multimodal-instruct) multimodal, contain LoRA adapter(s) that are expected to always be applied when a given modality is present. This can be a bit tedious to manage with the above approaches, as it requires the user to send the `LoRARequest` (offline) or to filter requests between the base model and LoRA model (server) depending on the content of the request's multimodal data.
-
-To this end, we allow registration of default multimodal LoRAs to handle this automatically, where users can map each modality to a LoRA adapter to automatically apply it when the corresponding inputs are present. Note that currently, we only allow one LoRA per prompt; if several modalities are provided, each of which are registered to a given modality, none of them will be applied.
-
-??? code "Example usage for offline inference"
-
-    ```python
-    from transformers import AutoTokenizer
-    from vllm import LLM, SamplingParams
-    from vllm.assets.audio import AudioAsset
-
-    model_id = "ibm-granite/granite-speech-3.3-2b"
-    tokenizer = AutoTokenizer.from_pretrained(model_id)
-
-    def get_prompt(question: str, has_audio: bool):
-        """Build the input prompt to send to vLLM."""
-        if has_audio:
-            question = f"<|audio|>{question}"
-        chat = [
-            {
-                "role": "user",
-                "content": question
-            }
-        ]
-        return tokenizer.apply_chat_template(chat, tokenize=False)
-
-
-    model = LLM(
-        model=model_id,
-        enable_lora=True,
-        max_lora_rank=64,
-        max_model_len=2048,
-        limit_mm_per_prompt={"audio": 1},
-        # Will always pass a `LoRARequest` with the `model_id`
-        # whenever audio is contained in the request data.
-        default_mm_loras = {"audio": model_id},
-        enforce_eager=True,
-    )
-
-    question = "can you transcribe the speech into a written format?"
-    prompt_with_audio = get_prompt(
-        question=question,
-        has_audio=True,
-    )
-    audio = AudioAsset("mary_had_lamb").audio_and_sample_rate
-
-    inputs = {
-        "prompt": prompt_with_audio,
-        "multi_modal_data": {
-            "audio": audio,
-        }
-    }
-
-
-    outputs = model.generate(
-        inputs,
-        sampling_params=SamplingParams(
-            temperature=0.2,
-            max_tokens=64,
-        ),
-    )
-    ```
-
-You can also pass a json dictionary of `--default-mm-loras` mapping modalities to LoRA model IDs. For example, when starting the server:
-
-```bash
-vllm serve ibm-granite/granite-speech-3.3-2b \
-    --max-model-len 2048 \
-    --enable-lora \
-    --default-mm-loras '{"audio":"ibm-granite/granite-speech-3.3-2b"}' \
-    --max-lora-rank 64
-```
-
-Note: Default multimodal LoRAs are currently only available for `.generate` and chat completions.
@@ -256,12 +256,12 @@ speculative decoding, breaking down the guarantees into three key areas:
 2. **Algorithmic Losslessness**
   \- vLLM’s implementation of speculative decoding is algorithmically validated to be lossless. Key validation tests include:

-    > - **Rejection Sampler Convergence**: Ensures that samples from vLLM’s rejection sampler align with the target
-    >   distribution. [View Test Code](https://github.com/vllm-project/vllm/blob/47b65a550866c7ffbd076ecb74106714838ce7da/tests/samplers/test_rejection_sampler.py#L252)
-    > - **Greedy Sampling Equality**: Confirms that greedy sampling with speculative decoding matches greedy sampling
-    >   without it. This verifies that vLLM's speculative decoding framework, when integrated with the vLLM forward pass and the vLLM rejection sampler,
-    >   provides a lossless guarantee. Almost all of the tests in <gh-dir:tests/spec_decode/e2e>.
-    >   verify this property using [this assertion implementation](https://github.com/vllm-project/vllm/blob/b67ae00cdbbe1a58ffc8ff170f0c8d79044a684a/tests/spec_decode/e2e/conftest.py#L291)
+   > - **Rejection Sampler Convergence**: Ensures that samples from vLLM’s rejection sampler align with the target
+   >   distribution. [View Test Code](https://github.com/vllm-project/vllm/blob/47b65a550866c7ffbd076ecb74106714838ce7da/tests/samplers/test_rejection_sampler.py#L252)
+   > - **Greedy Sampling Equality**: Confirms that greedy sampling with speculative decoding matches greedy sampling
+   >   without it. This verifies that vLLM's speculative decoding framework, when integrated with the vLLM forward pass and the vLLM rejection sampler,
+   >   provides a lossless guarantee. Almost all of the tests in <gh-dir:tests/spec_decode/e2e>.
+   >   verify this property using [this assertion implementation](https://github.com/vllm-project/vllm/blob/b67ae00cdbbe1a58ffc8ff170f0c8d79044a684a/tests/spec_decode/e2e/conftest.py#L291)

 3. **vLLM Logprob Stability**
   \- vLLM does not currently guarantee stable token log probabilities (logprobs). This can result in different outputs for the
@@ -1,112 +0,0 @@
-# Data Parallel Deployment
-
-vLLM supports Data Parallel deployment, where model weights are replicated across separate instances/GPUs to process independent batches of requests.
-
-This will work with both dense and MoE models.
-
-For MoE models, particularly those like DeepSeek that employ MLA (Multi-head Latent Attention), it can be advantageous to use data parallel for the attention layers and expert or tensor parallel (EP or TP) for the expert layers.
-
-In these cases, the data parallel ranks are not completely independent. Forward passes must be aligned, and expert layers across all ranks are required to synchronize during every forward pass, even when there are fewer requests to be processed than DP ranks.
-
-The expert layers will by default form a (DP x TP) sized tensor parallel group. To enable expert parallelism, include the `--enable-expert-parallel` CLI arg (on all nodes in the multi-node case).
-
-In vLLM, each DP rank is deployed as a separate "core engine" process that communicates with front-end process(es) via ZMQ sockets. Data Parallel attention can be combined with Tensor Parallel attention, in which case each DP engine owns a number of per-GPU worker processes equal to the configured TP size.
-
-For MoE models, when any requests are in progress in any rank, we must ensure that empty "dummy" forward passes are performed in all ranks that don't currently have any requests scheduled. This is handled via a separate DP Coordinator process that communicates with all ranks, and a collective operation performed every N steps to determine when all ranks become idle and can be paused. When TP is used in conjunction with DP, expert layers form an EP or TP group of size (DP x TP).
-
-In all cases, it is beneficial to load-balance requests between DP ranks. For online deployments, this balancing can be optimized by taking into account the state of each DP engine - in particular its currently scheduled and waiting (queued) requests, and KV cache state. Each DP engine has an independent KV cache, and the benefit of prefix caching can be maximized by directing prompts intelligently.
-
-This document focuses on online deployments (with the API server). DP + EP is also supported for offline usage (via the LLM class), for an example see <gh-file:examples/offline_inference/data_parallel.py>.
-
-There are two distinct modes supported for online deployments - self-contained with internal load balancing, or externally per-rank process deployment and load balancing.
-
-## Internal Load Balancing
-
-vLLM supports "self-contained" data parallel deployments that expose a single API endpoint.
-
-It can be configured by simply including e.g. `--data-parallel-size=4` in the vllm serve command line arguments. This will require 4 GPUs. It can be combined with tensor parallel, for example `--data-parallel-size=4 --tensor-parallel-size=2`, which would require 8 GPUs.
-
-Running a single data parallel deployment across multiple nodes requires a different `vllm serve` to be run on each node, specifying which DP ranks should run on that node. In this case, there will still be a single HTTP entrypoint - the API server(s) will run only on one node, but it doesn't necessarily need to be co-located with the DP ranks.
-
-This will run DP=4, TP=2 on a single 8-GPU node:
-
-```bash
-vllm serve $MODEL --data-parallel-size 4 --tensor-parallel-size 2
-```
-
-This will run DP=4 with DP ranks 0 and 1 on the head node and ranks 2 and 3 on the second node:
-
-```bash
-# Node 0  (with ip address 10.99.48.128)
-vllm serve $MODEL --data-parallel-size 4 --data-parallel-size-local 2 \
-                  --data-parallel-address 10.99.48.128 --data-parallel-rpc-port 13345
-# Node 1
-vllm serve $MODEL --headless --data-parallel-size 4 --data-parallel-size-local 2 \
-                  --data-parallel-start-rank 2 \
-                  --data-parallel-address 10.99.48.128 --data-parallel-rpc-port 13345
-```
-
-This will run DP=4 with only the API server on the first node and all engines on the second node:
-
-```bash
-# Node 0  (with ip address 10.99.48.128)
-vllm serve $MODEL --data-parallel-size 4 --data-parallel-size-local 0 \
-                  --data-parallel-address 10.99.48.128 --data-parallel-rpc-port 13345
-# Node 1
-vllm serve $MODEL --headless --data-parallel-size 4 --data-parallel-size-local 4 \
-                  --data-parallel-address 10.99.48.128 --data-parallel-rpc-port 13345
-```
-
-This DP mode can also be used with Ray, in which case only a single launch command is needed irrespective of the number of nodes:
-
-```bash
-vllm serve $MODEL --data-parallel-size 16 --tensor-parallel-size 2 --data-parallel-backend=ray
-```
-
-Currently, the internal DP load balancing is done within the API server process(es) and is based on the running and waiting queues in each of the engines. This could be made more sophisticated in future by incorporating KV cache aware logic.
-
-When deploying large DP sizes using this method, the API server process can become a bottleneck. In this case, the orthogonal `--api-server-count` command line option can be used to scale this out (for example `--api-server-count=4`). This is transparent to users - a single HTTP endpoint / port is still exposed. Note that this API server scale-out is "internal" and still confined to the "head" node.
-
-<figure markdown="1">
-![DP Internal LB Diagram](../assets/deployment/dp_internal_lb.png)
-</figure>
-
-## External Load Balancing
-
-For larger scale deployments especially, it can make sense to handle the orchestration and load balancing of data parallel ranks externally.
-
-In this case, it's more convenient to treat each DP rank like a separate vLLM deployment, with its own endpoint, and have an external router balance HTTP requests between them, making use of appropriate real-time telemetry from each server for routing decisions.
-
-This can already be done trivially for non-MoE models, since each deployed server is fully independent. No data parallel CLI options need to be used for this.
-
-We support an equivalent topology for MoE DP+EP which can be configured via the following CLI arguments.
-
-If DP ranks are co-located (same node / ip address), a default RPC port is used, but a different HTTP server port must be specified for each rank:
-
-```bash
-# Rank 0
-CUDA_VISIBLE_DEVICES=0 vllm serve $MODEL --data-parallel-size 2 --data-parallel-rank 0 \
-                                         --port 8000
-# Rank 1
-CUDA_VISIBLE_DEVICES=1 vllm serve $MODEL --data-parallel-size 2 --data-parallel-rank 1 \
-                                         --port 8001
-```
-
-For multi-node cases, the address/port of rank 0 must also be specified:
-
-```bash
-# Rank 0  (with ip address 10.99.48.128)
-vllm serve $MODEL --data-parallel-size 2 --data-parallel-rank 0 \
-                  --data-parallel-address 10.99.48.128 --data-parallel-rpc-port 13345
-# Rank 1
-vllm serve $MODEL --data-parallel-size 2 --data-parallel-rank 1 \
-                  --data-parallel-address 10.99.48.128 --data-parallel-rpc-port 13345
-```
-
-The coordinator process also runs in this scenario, co-located with the DP rank 0 engine.
-
-<figure markdown="1">
-![DP External LB Diagram](../assets/deployment/dp_external_lb.png)
-</figure>
-
-In the above diagram, each of the dotted boxes corresponds to a separate launch of `vllm serve` - these could be separate Kubernetes pods, for example.
@@ -15,10 +15,6 @@ After adding enough GPUs and nodes to hold the model, you can run vLLM first, wh
 !!! note
    There is one edge case: if the model fits in a single node with multiple GPUs, but the number of GPUs cannot divide the model size evenly, you can use pipeline parallelism, which splits the model along layers and supports uneven splits. In this case, the tensor parallel size should be 1 and the pipeline parallel size should be the number of GPUs.

-### Distributed serving of MoE (Mixture of Experts) models
-
-It is often advantageous to exploit the inherent parallelism of experts by using a separate parallelism strategy for the expert layers. vLLM supports large-scale deployment combining Data Parallel attention with Expert or Tensor Parallel MoE layers. See the page on [Data Parallel Deployment](data_parallel_deployment.md) for more information.
-
 ## Running vLLM on a single node

 vLLM supports distributed tensor-parallel and pipeline-parallel inference and serving. Currently, we support [Megatron-LM's tensor parallel algorithm](https://arxiv.org/pdf/1909.08053.pdf). We manage the distributed runtime with either [Ray](https://github.com/ray-project/ray) or python native multiprocessing. Multiprocessing can be used when deploying on a single node, multi-node inference currently requires Ray.
@@ -21,9 +21,7 @@ prometheus-fastapi-instrumentator >= 7.0.0
 tiktoken >= 0.6.0  # Required for DBRX tokenizer
 lm-format-enforcer >= 0.10.11, < 0.11
 llguidance >= 0.7.11, < 0.8.0; platform_machine == "x86_64" or platform_machine == "arm64" or platform_machine == "aarch64"
-outlines_core == 0.2.10
-# required for outlines backend disk cache
-diskcache == 5.6.3
+outlines == 0.1.11
 lark == 1.2.2
 xgrammar == 0.1.19; platform_machine == "x86_64" or platform_machine == "aarch64" or platform_machine == "arm64"
 typing_extensions >= 4.10
@@ -40,7 +38,7 @@ six>=1.16.0; python_version > '3.11' # transitive dependency of pandas that need
 setuptools>=77.0.3,<80; python_version > '3.11' # Setuptools is used by triton, we need to ensure a modern version is installed for 3.12+ so that it does not try to import distutils, which was removed in 3.12
 einops # Required for Qwen2-VL.
 compressed-tensors == 0.10.2 # required for compressed-tensors
-depyf==0.19.0 # required for profiling and debugging with compilation config
+depyf==0.18.0 # required for profiling and debugging with compilation config
 cloudpickle # allows pickling lambda functions in model_executor/models/registry.py
 watchfiles # required for http server to monitor the updates of TLS files
 python-json-logger # Used by logging as per examples/others/logging_configuration.md
@@ -3,7 +3,6 @@

 from __future__ import annotations

-import tempfile
 from typing import Any, Optional, Union

 import pytest
@@ -112,11 +111,6 @@ def test_full_graph(
                           pass_config=PassConfig(enable_fusion=True,
                                                  enable_noop=True)), model)
        for model in models_list(keywords=["FP8-dynamic", "quantized.w8a8"])
-    ] + [
-        # Test depyf integration works
-        (CompilationConfig(level=CompilationLevel.PIECEWISE,
-                           debug_dump_path=tempfile.gettempdir()),
-         ("facebook/opt-125m", {})),
    ])
 # only test some of the models
@create_new_process_for_each_test()
@@ -44,9 +44,7 @@ class TestModel(torch.nn.Module):
        ]
        self.fp8_linear = Fp8LinearOp(
            cutlass_fp8_supported=cutlass_fp8_enabled,
-            act_quant_static=static,
-            act_quant_group_shape=group_shape,
-        )
+            use_per_token_if_dynamic=True)

    def forward(self, x):
        resid = torch.sqrt(x)
@@ -93,10 +91,9 @@ def test_fusion_rmsnorm_quant(dtype, hidden_size, num_tokens, eps, static,
    maybe_create_device_identity()  # needed for certain non-cutlass fp8 paths

    vllm_config = VllmConfig(compilation_config=CompilationConfig(
-        level=CompilationLevel.PIECEWISE,
-        custom_ops=["+rms_norm", "+quant_fp8"],
-        pass_config=PassConfig(enable_fusion=True, enable_noop=True),
-    ))
+        level=CompilationLevel.PIECEWISE, custom_ops=["+rms_norm"]))
+    vllm_config.compilation_config.pass_config = \
+        PassConfig(enable_fusion=True, enable_noop=True)
    with vllm.config.set_current_vllm_config(vllm_config):
        # Reshape pass is needed for the fusion pass to work
        noop_pass = NoOpEliminationPass(vllm_config)
@@ -50,7 +50,6 @@ def test_attention_fusion(example_prompts, monkeypatch, model: str,
        # DYNAMO_ONCE does not properly propagate shapes.
        level=CompilationLevel.DYNAMO_AS_IS,
        backend="tests.compile.test_fusion_attn.backend_unfused",
-        custom_ops=["+quant_fp8"],
    )
    vllm_config = VllmConfig(compilation_config=compile_config)
    backend_unfused = TestBackend(NoOpEliminationPass(vllm_config))
@@ -74,7 +73,6 @@ def test_attention_fusion(example_prompts, monkeypatch, model: str,
        # DYNAMO_ONCE does not properly propagate shapes.
        level=CompilationLevel.DYNAMO_AS_IS,
        backend="tests.compile.test_fusion_attn.backend",
-        custom_ops=["+quant_fp8"],
    )
    vllm_config = VllmConfig(compilation_config=compile_config)

@@ -4,56 +4,33 @@ import pytest
 import torch

 import vllm.envs as envs
+from vllm._custom_ops import scaled_fp8_quant
 from vllm.compilation.activation_quant_fusion import ActivationQuantFusionPass
 from vllm.compilation.fx_utils import find_auto_fn, find_auto_fn_maybe
-from vllm.compilation.noop_elimination import NoOpEliminationPass
 from vllm.config import CompilationConfig, PassConfig, VllmConfig
 from vllm.model_executor.layers.activation import SiluAndMul
-from vllm.model_executor.layers.quantization.utils.quant_utils import (
-    GroupShape)
-from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
-    CUTLASS_FP8_SUPPORTED, Fp8LinearOp)
-from vllm.platforms import current_platform

 from .backend import TestBackend


 class TestModel(torch.nn.Module):

-    def __init__(self, hidden_size: int, cutlass_fp8_enabled: bool, *args,
-                 **kwargs):
+    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.silu_and_mul = SiluAndMul()
-        self.wscale = torch.rand(1, dtype=torch.float32)
        self.scale = torch.rand(1, dtype=torch.float32)

-        self.w = (torch.rand(
-            hidden_size,
-            hidden_size).to(dtype=current_platform.fp8_dtype()).t())
-
-        self.fp8_linear = Fp8LinearOp(
-            cutlass_fp8_supported=cutlass_fp8_enabled,
-            act_quant_static=True,
-            act_quant_group_shape=GroupShape.PER_TENSOR,
-        )
-
    def forward(self, x):
        y = self.silu_and_mul(x)
-        x2 = self.fp8_linear.apply(y,
-                                   self.w,
-                                   self.wscale,
-                                   input_scale=self.wscale)
+        x2 = scaled_fp8_quant(y, self.scale)
        return x2


@pytest.mark.parametrize("num_tokens", [256])
@pytest.mark.parametrize("hidden_size", [64])
-@pytest.mark.parametrize("cutlass_fp8_enabled",
-                         [True, False] if CUTLASS_FP8_SUPPORTED else [False])
@pytest.mark.skipif(envs.VLLM_TARGET_DEVICE not in ["cuda", "rocm"],
                    reason="Only test on CUDA and ROCm")
-def test_fusion_silu_and_mul_quant(num_tokens, hidden_size,
-                                   cutlass_fp8_enabled):
+def test_fusion_silu_and_mul_quant(num_tokens, hidden_size):
    torch.set_default_device("cuda")
    torch.set_default_dtype(torch.float16)

@@ -63,11 +40,11 @@ def test_fusion_silu_and_mul_quant(num_tokens, hidden_size,
        pass_config=PassConfig(enable_fusion=True, enable_noop=True))
    fusion_pass = ActivationQuantFusionPass(config)

-    backend = TestBackend(NoOpEliminationPass(config), fusion_pass)
-    model = TestModel(hidden_size, cutlass_fp8_enabled)
+    backend = TestBackend(fusion_pass)
+    model = TestModel()

    # First dimension dynamic
-    x = torch.rand(num_tokens, hidden_size * 2)
+    x = torch.rand(num_tokens, hidden_size)
    torch._dynamo.mark_dynamic(x, 0)

    result = model(x)
@@ -69,12 +69,6 @@ def test_lm_eval_accuracy_v1_engine(model, monkeypatch: pytest.MonkeyPatch):
        more_args = None
        if current_platform.is_tpu():
            # Limit compilation time for TPU V1
-
-            if model == "google/gemma-3-1b-it":
-                pytest.skip(
-                    "Temporarily disabled due to test failures"
-                    "(timeout or accuracy mismatch). Re-enable once fixed.")
-
            more_args = "max_model_len=2048,max_num_seqs=64"

            # Add TP test (if provided)
@@ -16,18 +16,14 @@ from vllm.outputs import RequestOutput
 from vllm.sampling_params import GuidedDecodingParams, SamplingParams

 MODEL_NAME = "Qwen/Qwen2.5-1.5B-Instruct"
-
-# Separate backends which support grammars vs ones
-# which only support regex based constraints in tests.
-GRAMMAR_DECODING_BACKENDS = [
+GUIDED_DECODING_BACKENDS = [
    # (backend, disable_any_whitespace),
+    ("outlines", False),
    ("lm-format-enforcer", False),
    ("xgrammar", True),
    ("guidance", True),
 ]

-ALL_DECODING_BACKENDS = ([("outlines", False)] + GRAMMAR_DECODING_BACKENDS)
-

@pytest.fixture(scope="module")
 def llm():
@@ -43,7 +39,7 @@ def llm():

@pytest.mark.skip_global_cleanup
@pytest.mark.parametrize("guided_decoding_backend,disable_any_whitespace",
-                         ALL_DECODING_BACKENDS)
+                         GUIDED_DECODING_BACKENDS)
 def test_guided_regex(sample_regex, llm, guided_decoding_backend: str,
                      disable_any_whitespace: bool):
    sampling_params = SamplingParams(
@@ -53,7 +49,6 @@ def test_guided_regex(sample_regex, llm, guided_decoding_backend: str,
            regex=sample_regex,
            backend=guided_decoding_backend,
            disable_any_whitespace=disable_any_whitespace))
-
    outputs = llm.generate(prompts=[
        f"Give an example IPv4 address with this regex: {sample_regex}"
    ] * 2,
@@ -74,7 +69,7 @@ def test_guided_regex(sample_regex, llm, guided_decoding_backend: str,

@pytest.mark.skip_global_cleanup
@pytest.mark.parametrize("guided_decoding_backend,disable_any_whitespace",
-                         ALL_DECODING_BACKENDS)
+                         GUIDED_DECODING_BACKENDS)
 def test_guided_json_completion(sample_json_schema, llm,
                                guided_decoding_backend: str,
                                disable_any_whitespace: bool):
@@ -108,7 +103,7 @@ def test_guided_json_completion(sample_json_schema, llm,

@pytest.mark.skip_global_cleanup
@pytest.mark.parametrize("guided_decoding_backend,disable_any_whitespace",
-                         ALL_DECODING_BACKENDS)
+                         GUIDED_DECODING_BACKENDS)
 def test_guided_complex_json_completion(sample_complex_json_schema, llm,
                                        guided_decoding_backend: str,
                                        disable_any_whitespace: bool):
@@ -143,7 +138,7 @@ def test_guided_complex_json_completion(sample_complex_json_schema, llm,

@pytest.mark.skip_global_cleanup
@pytest.mark.parametrize("guided_decoding_backend,disable_any_whitespace",
-                         ALL_DECODING_BACKENDS)
+                         GUIDED_DECODING_BACKENDS)
 def test_guided_definition_json_completion(sample_definition_json_schema, llm,
                                           guided_decoding_backend: str,
                                           disable_any_whitespace: bool):
@@ -178,7 +173,7 @@ def test_guided_definition_json_completion(sample_definition_json_schema, llm,

@pytest.mark.skip_global_cleanup
@pytest.mark.parametrize("guided_decoding_backend,disable_any_whitespace",
-                         ALL_DECODING_BACKENDS)
+                         GUIDED_DECODING_BACKENDS)
 def test_guided_enum_json_completion(sample_enum_json_schema, llm,
                                     guided_decoding_backend: str,
                                     disable_any_whitespace: bool):
@@ -223,7 +218,7 @@ def test_guided_enum_json_completion(sample_enum_json_schema, llm,

@pytest.mark.skip_global_cleanup
@pytest.mark.parametrize("guided_decoding_backend,disable_any_whitespace",
-                         ALL_DECODING_BACKENDS)
+                         GUIDED_DECODING_BACKENDS)
 def test_guided_choice_completion(sample_guided_choice, llm,
                                  guided_decoding_backend: str,
                                  disable_any_whitespace: bool):
@@ -253,7 +248,7 @@ def test_guided_choice_completion(sample_guided_choice, llm,

@pytest.mark.skip_global_cleanup
@pytest.mark.parametrize("guided_decoding_backend,disable_any_whitespace",
-                         GRAMMAR_DECODING_BACKENDS)
+                         GUIDED_DECODING_BACKENDS)
 def test_guided_grammar(sample_sql_statements, llm,
                        guided_decoding_backend: str,
                        disable_any_whitespace: bool):
@@ -349,7 +344,7 @@ def test_disable_guided_decoding_fallback(sample_regex, llm):

@pytest.mark.skip_global_cleanup
@pytest.mark.parametrize("guided_decoding_backend,disable_any_whitespace",
-                         GRAMMAR_DECODING_BACKENDS)
+                         GUIDED_DECODING_BACKENDS)
 def test_guided_json_object(llm, guided_decoding_backend: str,
                            disable_any_whitespace: bool):
    sampling_params = SamplingParams(
@@ -382,9 +377,7 @@ def test_guided_json_object(llm, guided_decoding_backend: str,

            # Parse to verify it is valid JSON
            parsed_json = json.loads(generated_text)
-            # A list is not what was intended, but is still valid
-            # json.
-            assert isinstance(parsed_json, (dict, list))
+            assert isinstance(parsed_json, dict)


 class CarType(str, Enum):
@@ -402,7 +395,7 @@ class CarDescription(BaseModel):

@pytest.mark.skip_global_cleanup
@pytest.mark.parametrize("guided_decoding_backend,disable_any_whitespace",
-                         ALL_DECODING_BACKENDS)
+                         GUIDED_DECODING_BACKENDS)
 def test_guided_json_completion_with_enum(llm, guided_decoding_backend: str,
                                          disable_any_whitespace: bool):
    json_schema = CarDescription.model_json_schema()
@@ -434,7 +427,7 @@ def test_guided_json_completion_with_enum(llm, guided_decoding_backend: str,

@pytest.mark.skip_global_cleanup
@pytest.mark.parametrize("guided_decoding_backend,disable_any_whitespace",
-                         ALL_DECODING_BACKENDS)
+                         GUIDED_DECODING_BACKENDS)
 def test_guided_number_range_json_completion(llm, guided_decoding_backend: str,
                                             disable_any_whitespace: bool):
    sample_output_schema = {
@@ -1113,7 +1113,10 @@ async def test_http_chat_no_model_name_with_curl(server: RemoteOpenAIServer):


@pytest.mark.asyncio
-async def test_http_chat_no_model_name_with_openai(server: RemoteOpenAIServer):
+@pytest.mark.parametrize("model_name", [MODEL_NAME, ""])
+async def test_http_chat_no_model_name_with_openai(server: RemoteOpenAIServer,
+                                                   model_name: str):
+
    openai_api_key = "EMPTY"
    openai_api_base = f"http://localhost:{server.port}/v1"

@@ -1132,35 +1135,3 @@ async def test_http_chat_no_model_name_with_openai(server: RemoteOpenAIServer):
        messages=messages,
    )
    assert response.model == MODEL_NAME
-
-
-@pytest.mark.asyncio
-async def test_invocations(server: RemoteOpenAIServer,
-                           client: openai.AsyncOpenAI):
-    messages = [{
-        "role": "system",
-        "content": "you are a helpful assistant"
-    }, {
-        "role": "user",
-        "content": "what is 1+1?"
-    }]
-
-    request_args = {
-        "model": MODEL_NAME,
-        "messages": messages,
-        "max_completion_tokens": 5,
-        "temperature": 0.0,
-        "logprobs": False,
-    }
-
-    chat_completion = await client.chat.completions.create(**request_args)
-
-    invocation_response = requests.post(server.url_for("invocations"),
-                                        json=request_args)
-    invocation_response.raise_for_status()
-
-    chat_output = chat_completion.model_dump()
-    invocation_output = invocation_response.json()
-
-    assert chat_output.keys() == invocation_output.keys()
-    assert chat_output["choices"] == invocation_output["choices"]
@@ -155,25 +155,3 @@ def test_batch_classification_empty_list(server: RemoteOpenAIServer,
    assert output.object == "list"
    assert isinstance(output.data, list)
    assert len(output.data) == 0
-
-
-@pytest.mark.asyncio
-async def test_invocations(server: RemoteOpenAIServer):
-    request_args = {
-        "model": MODEL_NAME,
-        "input": "This product was excellent and exceeded my expectations"
-    }
-
-    classification_response = requests.post(server.url_for("classify"),
-                                            json=request_args)
-    classification_response.raise_for_status()
-
-    invocation_response = requests.post(server.url_for("invocations"),
-                                        json=request_args)
-    invocation_response.raise_for_status()
-
-    classification_output = classification_response.json()
-    invocation_output = invocation_response.json()
-
-    assert classification_output.keys() == invocation_output.keys()
-    assert classification_output["data"] == invocation_output["data"]
@@ -11,7 +11,6 @@ import openai  # use the official client for correctness check
 import pytest
 import pytest_asyncio
 import regex as re
-import requests
 # downloading lora to test lora requests
 from huggingface_hub import snapshot_download
 from openai import BadRequestError
@@ -834,27 +833,3 @@ async def test_echo_stream_completion(client: openai.AsyncOpenAI,
            assert content is not None and saying in content
        else:
            assert content is not None and saying not in content
-
-
-@pytest.mark.asyncio
-async def test_invocations(server: RemoteOpenAIServer,
-                           client: openai.AsyncOpenAI):
-    request_args = {
-        "model": MODEL_NAME,
-        "prompt": "Hello, my name is",
-        "max_tokens": 5,
-        "temperature": 0.0,
-        "logprobs": None,
-    }
-
-    completion = await client.completions.create(**request_args)
-
-    invocation_response = requests.post(server.url_for("invocations"),
-                                        json=request_args)
-    invocation_response.raise_for_status()
-
-    completion_output = completion.model_dump()
-    invocation_output = invocation_response.json()
-
-    assert completion_output.keys() == invocation_output.keys()
-    assert completion_output["choices"] == invocation_output["choices"]
@@ -1,107 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import os
-
-import openai  # use the official client for correctness check
-import pytest
-import pytest_asyncio
-from huggingface_hub import snapshot_download
-
-from ...conftest import AudioTestAssets
-from ...utils import RemoteOpenAIServer
-
-# NOTE - the tests in this module are currently analogous to test_chat, but are
-# separated to avoid OOM killing due to module-scoped servers, since we
-# need a multimodal model for these tests.
-
-# Contains a modality specific lora alongside the base model
-MULTIMODAL_MODEL_NAME = snapshot_download(
-    "microsoft/Phi-4-multimodal-instruct")
-AUDIO_LORA_PATH = os.path.join(MULTIMODAL_MODEL_NAME, "speech-lora")
-
-ACTIVE_MM_LORA_RESPONSE = "Spoken text: The first words I spoke in the original chronograph, a little piece of practical poetry. Mary had a little lamb, it slept with quite a snow, and everywhere that Mary went, the lamb was sure to go."  # noqa: E501
-
-
-@pytest.fixture(scope="module")
-def monkeypatch_module():
-    from _pytest.monkeypatch import MonkeyPatch
-    mpatch = MonkeyPatch()
-    yield mpatch
-    mpatch.undo()
-
-
-@pytest.fixture(scope="module", params=[False, True])
-def multimodal_server(request, monkeypatch_module):  # noqa: F811
-
-    use_v1 = request.param
-    monkeypatch_module.setenv('VLLM_USE_V1', '1' if use_v1 else '0')
-
-    args = [
-        # use half precision for speed and memory savings in CI environment
-        "--dtype",
-        "half",
-        "--max-model-len",
-        "12800",
-        "--enforce-eager",
-        # lora config below
-        "--enable-lora",
-        "--lora-modules",
-        f"speech={AUDIO_LORA_PATH}",
-        "--max-lora-rank",
-        "320",
-        "--max-num-seqs",
-        "2",
-        "--trust-remote-code",
-        "--gpu-memory-utilization",
-        "0.8",
-        "--default-mm-loras",
-        f"{{\"audio\": \"{AUDIO_LORA_PATH}\"}}",
-    ]
-
-    with RemoteOpenAIServer(MULTIMODAL_MODEL_NAME, args) as remote_server:
-        yield remote_server
-
-
-@pytest_asyncio.fixture
-async def multi_modal_client(multimodal_server):
-    async with multimodal_server.get_async_client() as async_client:
-        yield async_client
-
-
-@pytest.mark.asyncio
-@pytest.mark.parametrize(
-    # base model with default lora should give the same response as lora model
-    "model_name",
-    [MULTIMODAL_MODEL_NAME, "speech"],
-)
-async def test_default_mm_lora_chat_completions(
-    model_name: str,
-    multi_modal_client: openai.AsyncOpenAI,
-    audio_assets: AudioTestAssets,
-):
-    messages = [{
-        "role":
-        "user",
-        "content": [{
-            "type": "text",
-            "text": "Can you transcribe this audio?",
-        }, {
-            "type": "audio_url",
-            "audio_url": {
-                "url": audio_assets[0].url
-            },
-        }]
-    }]
-
-    chat_completion = await multi_modal_client.chat.completions.create(
-        model=model_name,
-        messages=messages,
-        max_completion_tokens=128,
-        temperature=0.0)
-
-    assert len(chat_completion.choices) > 0
-
-    message = chat_completion.choices[0].message
-    assert message.content is not None and len(message.content) >= 0
-    assert message.content == ACTIVE_MM_LORA_RESPONSE
@@ -296,63 +296,3 @@ async def test_single_embedding_truncation_invalid(client: openai.AsyncOpenAI,
        assert "error" in response.object
        assert "truncate_prompt_tokens value is greater than max_model_len. "\
               "Please, select a smaller truncation size." in response.message
-
-
-@pytest.mark.asyncio
-async def test_invocations(server: RemoteOpenAIServer,
-                           client: openai.AsyncOpenAI):
-    input_texts = [
-        "The chef prepared a delicious meal.",
-    ]
-
-    request_args = {
-        "model": MODEL_NAME,
-        "input": input_texts,
-        "encoding_format": "float",
-    }
-
-    completion_response = await client.embeddings.create(**request_args)
-
-    invocation_response = requests.post(server.url_for("invocations"),
-                                        json=request_args)
-    invocation_response.raise_for_status()
-
-    completion_output = completion_response.model_dump()
-    invocation_output = invocation_response.json()
-
-    assert completion_output.keys() == invocation_output.keys()
-    assert completion_output["data"] == invocation_output["data"]
-
-
-@pytest.mark.asyncio
-async def test_invocations_conversation(server: RemoteOpenAIServer):
-    messages = [{
-        "role": "user",
-        "content": "The cat sat on the mat.",
-    }, {
-        "role": "assistant",
-        "content": "A feline was resting on a rug.",
-    }, {
-        "role": "user",
-        "content": "Stars twinkle brightly in the night sky.",
-    }]
-
-    request_args = {
-        "model": MODEL_NAME,
-        "messages": messages,
-        "encoding_format": "float",
-    }
-
-    chat_response = requests.post(server.url_for("v1/embeddings"),
-                                  json=request_args)
-    chat_response.raise_for_status()
-
-    invocation_response = requests.post(server.url_for("invocations"),
-                                        json=request_args)
-    invocation_response.raise_for_status()
-
-    chat_output = chat_response.json()
-    invocation_output = invocation_response.json()
-
-    assert chat_output.keys() == invocation_output.keys()
-    assert chat_output["data"] == invocation_output["data"]
@@ -13,7 +13,7 @@ from vllm.transformers_utils.tokenizer import get_tokenizer

 from ...utils import RemoteOpenAIServer

-MODEL_NAME = "internlm/internlm2-1_8b-reward"
+MODEL_NAME = "jason9693/Qwen2.5-1.5B-apeach"
 DUMMY_CHAT_TEMPLATE = """{% for message in messages %}{{message['role'] + ': ' + message['content'] + '\\n'}}{% endfor %}"""  # noqa: E501


@@ -21,16 +21,15 @@ DUMMY_CHAT_TEMPLATE = """{% for message in messages %}{{message['role'] + ': ' +
 def server():
    args = [
        "--task",
-        "reward",
+        "classify",
        # use half precision for speed and memory savings in CI environment
        "--dtype",
        "bfloat16",
        "--enforce-eager",
        "--max-model-len",
-        "512",
+        "8192",
        "--chat-template",
        DUMMY_CHAT_TEMPLATE,
-        "--trust-remote-code",
    ]

    with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
@@ -58,10 +57,10 @@ async def test_single_pooling(server: RemoteOpenAIServer, model_name: str):

    assert poolings.id is not None
    assert len(poolings.data) == 1
-    assert len(poolings.data[0].data) == 8
+    assert len(poolings.data[0].data) == 2
    assert poolings.usage.completion_tokens == 0
-    assert poolings.usage.prompt_tokens == 8
-    assert poolings.usage.total_tokens == 8
+    assert poolings.usage.prompt_tokens == 7
+    assert poolings.usage.total_tokens == 7

    # test using token IDs
    input_tokens = [1, 1, 1, 1, 1]
@@ -78,7 +77,7 @@ async def test_single_pooling(server: RemoteOpenAIServer, model_name: str):

    assert poolings.id is not None
    assert len(poolings.data) == 1
-    assert len(poolings.data[0].data) == 5
+    assert len(poolings.data[0].data) == 2
    assert poolings.usage.completion_tokens == 0
    assert poolings.usage.prompt_tokens == 5
    assert poolings.usage.total_tokens == 5
@@ -105,10 +104,10 @@ async def test_batch_pooling(server: RemoteOpenAIServer, model_name: str):

    assert poolings.id is not None
    assert len(poolings.data) == 3
-    assert len(poolings.data[0].data) == 8
+    assert len(poolings.data[0].data) == 2
    assert poolings.usage.completion_tokens == 0
-    assert poolings.usage.prompt_tokens == 29
-    assert poolings.usage.total_tokens == 29
+    assert poolings.usage.prompt_tokens == 25
+    assert poolings.usage.total_tokens == 25

    # test list[list[int]]
    input_tokens = [[4, 5, 7, 9, 20], [15, 29, 499], [24, 24, 24, 24, 24],
@@ -126,7 +125,7 @@ async def test_batch_pooling(server: RemoteOpenAIServer, model_name: str):

    assert poolings.id is not None
    assert len(poolings.data) == 4
-    assert len(poolings.data[0].data) == 5
+    assert len(poolings.data[0].data) == 2
    assert poolings.usage.completion_tokens == 0
    assert poolings.usage.prompt_tokens == 17
    assert poolings.usage.total_tokens == 17
@@ -158,11 +157,7 @@ async def test_conversation_pooling(server: RemoteOpenAIServer,
    chat_response.raise_for_status()
    chat_poolings = PoolingResponse.model_validate(chat_response.json())

-    tokenizer = get_tokenizer(
-        tokenizer_name=model_name,
-        tokenizer_mode="fast",
-        trust_remote_code=True,
-    )
+    tokenizer = get_tokenizer(tokenizer_name=model_name, tokenizer_mode="fast")
    prompt = tokenizer.apply_chat_template(
        messages,
        chat_template=DUMMY_CHAT_TEMPLATE,
@@ -211,9 +206,6 @@ async def test_batch_base64_pooling(server: RemoteOpenAIServer,
    )
    float_response.raise_for_status()
    responses_float = PoolingResponse.model_validate(float_response.json())
-    float_data = [
-        np.array(d.data).squeeze(-1).tolist() for d in responses_float.data
-    ]

    base64_response = requests.post(
        server.url_for("pooling"),
@@ -232,10 +224,11 @@ async def test_batch_base64_pooling(server: RemoteOpenAIServer,
            np.frombuffer(base64.b64decode(data.data),
                          dtype="float32").tolist())

-    check_embeddings_close(embeddings_0_lst=float_data,
-                           embeddings_1_lst=decoded_responses_base64_data,
-                           name_0="float32",
-                           name_1="base64")
+    check_embeddings_close(
+        embeddings_0_lst=[d.data for d in responses_float.data],
+        embeddings_1_lst=decoded_responses_base64_data,
+        name_0="float32",
+        name_1="base64")

    # Default response is float32 decoded from base64 by OpenAI Client
    default_response = requests.post(
@@ -247,71 +240,9 @@ async def test_batch_base64_pooling(server: RemoteOpenAIServer,
    )
    default_response.raise_for_status()
    responses_default = PoolingResponse.model_validate(default_response.json())
-    default_data = [
-        np.array(d.data).squeeze(-1).tolist() for d in responses_default.data
-    ]

-    check_embeddings_close(embeddings_0_lst=float_data,
-                           embeddings_1_lst=default_data,
-                           name_0="float32",
-                           name_1="default")
-
-
-@pytest.mark.asyncio
-async def test_invocations(server: RemoteOpenAIServer):
-    input_texts = [
-        "The chef prepared a delicious meal.",
-    ]
-
-    request_args = {
-        "model": MODEL_NAME,
-        "input": input_texts,
-        "encoding_format": "float",
-    }
-
-    completion_response = requests.post(server.url_for("pooling"),
-                                        json=request_args)
-    completion_response.raise_for_status()
-
-    invocation_response = requests.post(server.url_for("invocations"),
-                                        json=request_args)
-    invocation_response.raise_for_status()
-
-    completion_output = completion_response.json()
-    invocation_output = invocation_response.json()
-
-    assert completion_output.keys() == invocation_output.keys()
-    assert completion_output["data"] == invocation_output["data"]
-
-
-@pytest.mark.asyncio
-async def test_invocations_conversation(server: RemoteOpenAIServer):
-    messages = [{
-        "role": "user",
-        "content": "The cat sat on the mat.",
-    }, {
-        "role": "assistant",
-        "content": "A feline was resting on a rug.",
-    }, {
-        "role": "user",
-        "content": "Stars twinkle brightly in the night sky.",
-    }]
-
-    request_args = {
-        "model": MODEL_NAME,
-        "messages": messages,
-        "encoding_format": "float",
-    }
-
-    chat_response = requests.post(server.url_for("pooling"), json=request_args)
-    chat_response.raise_for_status()
-
-    invocation_response = requests.post(server.url_for("invocations"),
-                                        json=request_args)
-    invocation_response.raise_for_status()
-
-    chat_output = chat_response.json()
-    invocation_output = invocation_response.json()
-
-    assert chat_output.keys() == invocation_output.keys()
-    assert chat_output["data"] == invocation_output["data"]
+    check_embeddings_close(
+        embeddings_0_lst=[d.data for d in responses_default.data],
+        embeddings_1_lst=[d.data for d in responses_default.data],
+        name_0="float32",
+        name_1="base64")
@@ -94,30 +94,3 @@ def test_rerank_max_model_len(server: RemoteOpenAIServer, model_name: str):
    # Assert just a small fragments of the response
    assert "Please reduce the length of the input." in \
        rerank_response.text
-
-
-def test_invocations(server: RemoteOpenAIServer):
-    query = "What is the capital of France?"
-    documents = [
-        "The capital of Brazil is Brasilia.", "The capital of France is Paris."
-    ]
-
-    request_args = {
-        "model": MODEL_NAME,
-        "query": query,
-        "documents": documents,
-    }
-
-    rerank_response = requests.post(server.url_for("rerank"),
-                                    json=request_args)
-    rerank_response.raise_for_status()
-
-    invocation_response = requests.post(server.url_for("invocations"),
-                                        json=request_args)
-    invocation_response.raise_for_status()
-
-    rerank_output = rerank_response.json()
-    invocation_output = invocation_response.json()
-
-    assert rerank_output.keys() == invocation_output.keys()
-    assert rerank_output["results"] == invocation_output["results"]
@@ -191,28 +191,3 @@ class TestModel:
        assert score_response.status_code == 400
        assert "Please, select a smaller truncation size." in \
            score_response.text
-
-    def test_invocations(self, server: RemoteOpenAIServer, model: dict[str,
-                                                                       Any]):
-        text_1 = "What is the capital of France?"
-        text_2 = "The capital of France is Paris."
-
-        request_args = {
-            "model": model["name"],
-            "text_1": text_1,
-            "text_2": text_2,
-        }
-
-        score_response = requests.post(server.url_for("score"),
-                                       json=request_args)
-        score_response.raise_for_status()
-
-        invocation_response = requests.post(server.url_for("invocations"),
-                                            json=request_args)
-        invocation_response.raise_for_status()
-
-        score_output = score_response.json()
-        invocation_output = invocation_response.json()
-
-        assert score_output.keys() == invocation_output.keys()
-        assert score_output["data"] == invocation_output["data"]
@@ -1,140 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-from typing import Optional
-
-import flashinfer
-import pytest
-import torch
-
-from vllm.platforms import current_platform
-
-if not current_platform.is_device_capability(100):
-    pytest.skip("This TRTLLM kernel requires NVIDIA Blackwell.",
-                allow_module_level=True)
-
-FLOAT32_BYTES = torch.finfo(torch.float).bits // 8
-
-# KV Cache Layout for TRT-LLM
-# kv_cache_shape = (num_blocks, 2, num_kv_heads, page_size, head_dim)
-
-NUM_HEADS = [(64, 8), (16, 16), (40, 8), (32, 8)]
-HEAD_SIZES = [128]
-BLOCK_SIZES = [16, 32]
-DTYPES = [torch.float16, torch.bfloat16]
-NUM_BLOCKS = 32768  # Large enough to test overflow in index calculation.
-SOFT_CAPS = [None, 30.0, 50.0]
-
-
-def to_float8(x, dtype=torch.float8_e4m3fn):
-    finfo = torch.finfo(dtype)
-    min_val, max_val = x.aminmax()
-    amax = torch.maximum(min_val.abs(), max_val.abs()).clamp(min=1e-12)
-    scale = finfo.max / amax * 0.1
-    x_scl_sat = (x * scale).clamp(min=finfo.min, max=finfo.max)
-    return x_scl_sat.to(dtype), scale.float().reciprocal()
-
-
-@pytest.mark.parametrize("kv_lens", [[1328, 18, 463], [1, 54, 293, 70]])
-@pytest.mark.parametrize("num_heads", NUM_HEADS)
-@pytest.mark.parametrize("head_size", HEAD_SIZES)
-@pytest.mark.parametrize("block_size", BLOCK_SIZES)
-@pytest.mark.parametrize("kv_layout", ["HND"])
-@pytest.mark.parametrize("dtype", DTYPES)
-@pytest.mark.parametrize("soft_cap", SOFT_CAPS)
-@torch.inference_mode
-def test_flashinfer_trtllm_decode_with_baseline(
-    kv_lens: list[int],
-    num_heads: tuple[int, int],
-    head_size: int,
-    dtype: torch.dtype,
-    block_size: int,
-    soft_cap: Optional[float],
-    kv_layout: str,
-) -> None:
-    torch.set_default_device("cuda")
-    current_platform.seed_everything(0)
-    num_seqs = len(kv_lens)
-    num_query_heads = num_heads[0]
-    num_kv_heads = num_heads[1]
-
-    assert num_query_heads % num_kv_heads == 0
-    max_kv_len = max(kv_lens)
-    scale = head_size**-0.5
-
-    query = torch.randn(num_seqs, num_query_heads, head_size, dtype=dtype)
-    kv_cache_shape = None
-    if kv_layout == "NHD":
-        kv_cache_shape = (NUM_BLOCKS, 2, block_size, num_kv_heads, head_size)
-    elif kv_layout == "HND":
-        kv_cache_shape = (NUM_BLOCKS, 2, num_kv_heads, block_size, head_size)
-    else:
-        raise ValueError(f"Invalid kv_layout: {kv_layout}")
-    key_value_cache = torch.randn(kv_cache_shape, dtype=dtype)
-
-    max_num_blocks_per_seq = (max_kv_len + block_size - 1) // block_size
-    block_tables = torch.randint(0,
-                                 NUM_BLOCKS,
-                                 (num_seqs, max_num_blocks_per_seq),
-                                 dtype=torch.int32)
-    k_scale = v_scale = 1.0
-    kv_indptr = [0]
-    kv_indices = []
-    kv_last_page_lens = []
-    for i in range(num_seqs):
-        seq_len = kv_lens[i]
-        assert seq_len > 0
-        num_blocks = (seq_len + block_size - 1) // block_size
-        kv_indices.extend(block_tables[i, :num_blocks])
-        kv_indptr.append(kv_indptr[-1] + num_blocks)
-        kv_last_page_len = seq_len % block_size
-        if kv_last_page_len == 0:
-            kv_last_page_len = block_size
-        kv_last_page_lens.append(kv_last_page_len)
-
-    kv_indptr = torch.tensor(kv_indptr, dtype=torch.int32)
-    kv_indices = torch.tensor(kv_indices, dtype=torch.int32)
-    kv_last_page_lens = torch.tensor(kv_last_page_lens, dtype=torch.int32)
-
-    workspace_buffer = torch.empty(128 * 1024 * 1024, dtype=torch.int8)
-    wrapper = flashinfer.\
-        BatchDecodeWithPagedKVCacheWrapper(workspace_buffer, kv_layout,
-                use_tensor_cores=(
-                    (num_query_heads//num_kv_heads) > 4)
-                )
-    wrapper.plan(kv_indptr,
-                 kv_indices,
-                 kv_last_page_lens,
-                 num_query_heads,
-                 num_kv_heads,
-                 head_size,
-                 block_size,
-                 "NONE",
-                 q_data_type=dtype,
-                 kv_data_type=dtype,
-                 logits_soft_cap=soft_cap)
-
-    output = wrapper.run(query, key_value_cache, scale)
-
-    # TRTLLM Decode
-    max_kv_len = max(kv_lens)
-    kv_lens_tensor = torch.tensor(kv_lens,
-                                  dtype=torch.int,
-                                  device=query.device)
-    output_trtllm = flashinfer.decode.trtllm_batch_decode_with_kv_cache(
-        query.contiguous(),
-        key_value_cache,
-        workspace_buffer,
-        num_query_heads,
-        num_kv_heads,
-        scale,
-        block_tables,
-        kv_lens_tensor,
-        block_size,
-        max_kv_len,
-        "auto",
-        k_scale,
-        v_scale,
-    )
-
-    torch.testing.assert_close(output, output_trtllm, atol=1e-2, rtol=1e-2), \
-        f"{torch.max(torch.abs(output - output_trtllm))}"
@@ -1,160 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import argparse
-
-import torch
-
-import vllm.model_executor.layers.fused_moe.modular_kernel as mk
-from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
-
-from .common import Config
-from .mk_objects import (MK_ALL_PREPARE_FINALIZE_TYPES, MK_FUSED_EXPERT_TYPES,
-                         MK_SINGLE_GPU_PREPARE_FINALIZE_TYPES)
-
-
-def make_config_arg_parser(description: str):
-
-    def to_pf_class_type(s: str) -> mk.FusedMoEPrepareAndFinalize:
-        for pf in MK_ALL_PREPARE_FINALIZE_TYPES:
-            if pf.__name__ == s:
-                return pf
-        raise ValueError(
-            f"Cannot find a PrepareFinalize type that matches {s}")
-
-    def to_experts_class_type(s: str) -> mk.FusedMoEPermuteExpertsUnpermute:
-        for fe in MK_FUSED_EXPERT_TYPES:
-            if fe.__name__ == s:
-                return fe
-        raise ValueError(f"Cannot find a FusedExperts type that matches {s}")
-
-    def to_quant_torch_dtype(s: str) -> torch.dtype:
-        if s == "torch.float8_e4m3fn":
-            return torch.float8_e4m3fn
-        raise ValueError(f"Unsupported quant type {s}")
-
-    parser = argparse.ArgumentParser(description=description)
-
-    parser.add_argument(
-        "--world-size",
-        type=int,
-        default=2,
-        help="Number of ranks that participate in all2all",
-    )
-    parser.add_argument(
-        "--pf-type",
-        type=to_pf_class_type,
-        required=True,
-        help=("Choose a PrepareFinalize Type : "
-              f"{[x.__name__ for x in MK_ALL_PREPARE_FINALIZE_TYPES]}"),
-    )
-    parser.add_argument(
-        "--experts-type",
-        type=to_experts_class_type,
-        required=True,
-        help=(f"Choose a FusedExpert type : "
-              f"{[x.__name__ for x in MK_FUSED_EXPERT_TYPES]}"),
-    )
-    parser.add_argument(
-        "-m",
-        nargs="+",
-        type=int,
-        default=[64],
-        help="num tokens per rank",
-    )
-    parser.add_argument(
-        "-k",
-        type=int,
-        default=7168,
-        help="hidden-size",
-    )
-    parser.add_argument(
-        "-n",
-        type=int,
-        default=1024,
-        help="N dimension of the first fused-moe matmul",
-    )
-    parser.add_argument("--num-experts",
-                        type=int,
-                        default=32,
-                        help="Global num experts")
-    parser.add_argument("--topk",
-                        nargs="+",
-                        type=int,
-                        default=[4, 1],
-                        help="num topk")
-    parser.add_argument(
-        "--fused-moe-chunk-size",
-        nargs="+",
-        type=int,
-        help="Fused moe chunk size used for the non-batched fused experts impl."
-    )
-
-    # Quant args
-    parser.add_argument("--quant-dtype",
-                        type=to_quant_torch_dtype,
-                        help="Quant datatype")
-    parser.add_argument("--per-token-quantized-activations",
-                        action='store_true',
-                        help=("The input activations must be per-token "
-                              "quantized"))
-    parser.add_argument("--per-channel-quantized-weights",
-                        action="store_true",
-                        help="The weights must be per-channel quantized.")
-    parser.add_argument("--block-shape",
-                        nargs="+",
-                        type=int,
-                        help="Quantization block shape")
-
-    # Torch trace profile generation args
-    parser.add_argument("--torch-trace-dir-path",
-                        type=str,
-                        default=None,
-                        help="Get torch trace for single execution")
-
-    return parser
-
-
-def _validate_args(args: argparse.Namespace):
-
-    if args.quant_dtype is not None:
-        assert args.quant_dtype == torch.float8_e4m3fn
-        if args.block_shape is not None:
-            assert len(args.block_shape) == 2, (
-                f"block shape must have 2 elements. got {args.block_shape}")
-
-    if args.experts_type in MK_SINGLE_GPU_PREPARE_FINALIZE_TYPES:
-        assert args.world_size == 1, (
-            "Single GPU objects need world size set to 1")
-
-    if args.torch_trace_dir_path is not None:
-        from pathlib import Path
-        assert Path(args.torch_trace_dir_path).is_dir(), (
-            f"Please create {args.torch_trace_dir_path}")
-
-
-def make_config(args: argparse.Namespace) -> Config:
-
-    _validate_args(args)
-
-    quant_config = None
-    if args.quant_dtype is not None:
-        quant_config = FusedMoEQuantConfig(
-            quant_dtype=args.quant_dtype,
-            per_act_token_quant=args.per_token_quantized_activations,
-            per_out_ch_quant=args.per_channel_quantized_weights,
-            block_shape=args.block_shape)
-
-    return Config(
-        Ms=args.m,
-        K=args.k,
-        N=args.n,
-        E=args.num_experts,
-        topks=args.topk,
-        dtype=torch.bfloat16,  # hard-code
-        quant_config=quant_config,
-        prepare_finalize_type=args.pf_type,
-        fused_experts_type=args.experts_type,
-        fused_moe_chunk_size=args.fused_moe_chunk_size,
-        world_size=args.world_size,
-        torch_trace_dir_path=args.torch_trace_dir_path)
@@ -1,641 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-from dataclasses import dataclass
-from typing import Any, Optional, Union
-
-import torch
-
-import vllm._custom_ops as ops
-import vllm.model_executor.layers.fused_moe.modular_kernel as mk
-from tests.kernels.utils import torch_experts
-from vllm.config import VllmConfig
-from vllm.distributed import get_dp_group, get_tensor_model_parallel_world_size
-# Fused experts and PrepareFinalize imports
-from vllm.model_executor.layers.fused_moe.batched_deep_gemm_moe import (
-    BatchedDeepGemmExperts)
-from vllm.model_executor.layers.fused_moe.batched_triton_or_deep_gemm_moe import (  # noqa: E501
-    BatchedTritonOrDeepGemmExperts)
-from vllm.model_executor.layers.fused_moe.config import (
-    FusedMoEConfig, FusedMoEParallelConfig, FusedMoEQuantConfig)
-from vllm.model_executor.layers.fused_moe.cutlass_moe import CutlassExpertsFp8
-from vllm.model_executor.layers.fused_moe.deep_gemm_moe import DeepGemmExperts
-from vllm.model_executor.layers.fused_moe.fused_batched_moe import (
-    BatchedTritonExperts, NaiveBatchedExperts)
-from vllm.model_executor.layers.fused_moe.fused_moe import fused_topk
-from vllm.model_executor.layers.fused_moe.layer import (FusedMoEMethodBase,
-                                                        TritonExperts)
-from vllm.model_executor.layers.fused_moe.prepare_finalize import (
-    MoEPrepareAndFinalizeNoEP)
-from vllm.model_executor.layers.fused_moe.triton_deep_gemm_moe import (
-    TritonOrDeepGemmExperts)
-from vllm.utils import has_deep_ep, has_deep_gemm, has_pplx
-
-from .parallel_utils import ProcessGroupInfo
-from .utils import (make_block_quant_fp8_weights, make_non_quant_weights,
-                    make_quant_fp8_weights, per_token_cast_to_fp8)
-
-if has_pplx():
-    from vllm.model_executor.layers.fused_moe.pplx_prepare_finalize import (
-        PplxPrepareAndFinalize)
-if has_deep_ep():
-    from vllm.model_executor.layers.fused_moe.deepep_ht_prepare_finalize import (  # noqa: E501
-        DeepEPHTPrepareAndFinalize)
-    from vllm.model_executor.layers.fused_moe.deepep_ll_prepare_finalize import (  # noqa: E501
-        DeepEPLLPrepareAndFinalize)
-
-
-def _describe_tensor(t: Optional[torch.Tensor], name: str) -> str:
-    if t is None:
-        return f"{name} : None"
-    else:
-        return f"{name} : {t.shape} {t.dtype} {t.device}"
-
-
-@dataclass
-class Config:
-    Ms: Union[list[int], int]
-    K: int
-    N: int
-    E: int
-    topks: Union[list[int], int]
-    dtype: torch.dtype
-    quant_config: Optional[FusedMoEQuantConfig]
-
-    prepare_finalize_type: mk.FusedMoEPrepareAndFinalize
-    fused_experts_type: mk.FusedMoEPermuteExpertsUnpermute
-
-    fused_moe_chunk_size: Optional[int]
-    world_size: int
-
-    torch_trace_dir_path: Optional[str] = None
-
-    def describe(self) -> str:
-        s = ""
-        s += "== Config: \n"
-        s += f" world_size={self.world_size} \n"
-        s += f" PF={self.prepare_finalize_type.__name__} \n"
-        s += f" FE={self.fused_experts_type.__name__} \n"
-        s += f" topk={self.topks} \n"
-        s += f" dtype={self.dtype} \n"
-        s += f" fused_moe_chunk_size={self.fused_moe_chunk_size} \n"
-        s += " Quant: \n"
-        s += f" fused_moe_chunk_size={self.fused_moe_chunk_size} \n "
-        if self.quant_config is not None:
-            s += f"     q_dtype={self.quant_dtype} \n"
-            s += f"     q_block_shape={self.quant_block_shape} \n"
-            s += f"     q_per_out_ch_quant={self.is_per_out_ch_quant} \n"
-            s += f"     q_per_act_token={self.is_per_act_token_quant} \n"
-        else:
-            s += "     quant=None \n"
-        return s
-
-    @property
-    def M(self) -> int:
-        assert isinstance(self.Ms, int)
-        return self.Ms
-
-    @property
-    def quant_dtype(self) -> Optional[torch.dtype]:
-        if self.quant_config is None:
-            return None
-        return self.quant_config.quant_dtype
-
-    @property
-    def is_per_act_token_quant(self) -> bool:
-        if self.quant_config is None:
-            return False
-        return self.quant_config.per_act_token_quant
-
-    @property
-    def is_per_tensor_act_quant(self) -> bool:
-        if self.quant_config is None:
-            return False
-        return (not self.is_per_act_token_quant
-                and self.quant_block_shape is None)
-
-    @property
-    def is_per_out_ch_quant(self) -> bool:
-        if self.quant_config is None:
-            return False
-        return self.quant_config.per_out_ch_quant
-
-    @property
-    def quant_block_shape(self) -> Optional[list[int]]:
-        if self.quant_config is None:
-            return None
-        return self.quant_config.block_shape
-
-    @property
-    def topk(self) -> int:
-        assert isinstance(self.topks, int)
-        return self.topks
-
-    @property
-    def topk_ids_dtype(self) -> Optional[torch.dtype]:
-        topk_ids_dtype = None
-        if self.prepare_finalize_type == PplxPrepareAndFinalize:
-            topk_ids_dtype = torch.uint32
-        elif self.prepare_finalize_type in [
-                DeepEPHTPrepareAndFinalize, DeepEPLLPrepareAndFinalize
-        ]:
-            topk_ids_dtype = torch.int64
-        return topk_ids_dtype
-
-    @property
-    def num_local_experts(self) -> int:
-        return self.E // self.world_size
-
-    def make_env_data(self) -> tuple[VllmConfig, dict[Any, Any]]:
-        """
-        make env data for vllm launch. 
-        """
-        vllm_config = VllmConfig()
-        vllm_config.parallel_config.data_parallel_size = self.world_size
-        vllm_config.parallel_config.enable_expert_parallel = True
-
-        env_dict = {
-            "VLLM_ALL2ALL_BACKEND": self.all2all_backend(),
-            "VLLM_USE_DEEP_GEMM": str(int(self.needs_deep_gemm())),
-        }
-        if self.fused_moe_chunk_size is not None:
-            env_dict.update(
-                {"VLLM_FUSED_MOE_CHUNK_SIZE": str(self.fused_moe_chunk_size)})
-        return vllm_config, env_dict
-
-    def is_fp8_block_quantized(self):
-        return (self.quant_dtype == torch.float8_e4m3fn
-                and self.quant_block_shape is not None)
-
-    def is_batched_prepare_finalize(self):
-        return self.prepare_finalize_type in [
-            PplxPrepareAndFinalize, DeepEPLLPrepareAndFinalize
-        ]
-
-    def is_batched_fused_experts(self):
-        return self.fused_experts_type in [
-            CutlassExpertsFp8, BatchedDeepGemmExperts, BatchedTritonExperts,
-            NaiveBatchedExperts, BatchedTritonOrDeepGemmExperts
-        ]
-
-    def is_standard_fused_experts(self):
-        return self.fused_experts_type in [
-            CutlassExpertsFp8, DeepGemmExperts, TritonOrDeepGemmExperts,
-            TritonExperts
-        ]
-
-    def is_fe_16bit_supported(self):
-        return self.fused_experts_type in [
-            BatchedTritonExperts, BatchedTritonOrDeepGemmExperts,
-            NaiveBatchedExperts, TritonExperts
-        ]
-
-    def is_fe_fp8_supported(self):
-        return self.fused_experts_type in [
-            BatchedDeepGemmExperts,
-            BatchedTritonExperts,
-            BatchedTritonOrDeepGemmExperts,
-            CutlassExpertsFp8,
-            DeepGemmExperts,
-            TritonExperts,
-            TritonOrDeepGemmExperts,
-            NaiveBatchedExperts,
-        ]
-
-    def is_fe_block_fp8_supported(self):
-        return self.fused_experts_type in [
-            BatchedDeepGemmExperts,
-            BatchedTritonOrDeepGemmExperts,
-            DeepGemmExperts,
-            TritonExperts,
-            TritonOrDeepGemmExperts,
-            BatchedTritonExperts,
-            NaiveBatchedExperts,
-        ]
-
-    def is_fe_supports_chunking(self):
-        return self.fused_experts_type in [
-            CutlassExpertsFp8, DeepGemmExperts, TritonOrDeepGemmExperts,
-            TritonExperts
-        ]
-
-    def needs_deep_gemm(self):
-        return self.fused_experts_type in [
-            BatchedDeepGemmExperts,
-            DeepGemmExperts,
-        ]
-
-    def needs_pplx(self):
-        return self.prepare_finalize_type in [PplxPrepareAndFinalize]
-
-    def needs_deep_ep(self):
-        return self.prepare_finalize_type in [
-            DeepEPHTPrepareAndFinalize, DeepEPLLPrepareAndFinalize
-        ]
-
-    def all2all_backend(self):
-        if self.needs_pplx():
-            return "pplx"
-        if self.prepare_finalize_type == DeepEPHTPrepareAndFinalize:
-            return "deepep_high_throughput"
-        if self.prepare_finalize_type == DeepEPLLPrepareAndFinalize:
-            return "deepep_low_latency"
-        return "naive"
-
-    def needs_all2all(self):
-        return self.prepare_finalize_type in [
-            PplxPrepareAndFinalize, DeepEPHTPrepareAndFinalize,
-            DeepEPLLPrepareAndFinalize
-        ]
-
-    def is_valid(self):
-        # Check prepare-finalize and fused-experts compatibility
-        if self.is_batched_prepare_finalize():
-            if not self.is_batched_fused_experts():
-                return False
-        else:
-            if not self.is_standard_fused_experts():
-                return False
-
-        use_chunking = self.fused_moe_chunk_size is not None
-        if use_chunking and not self.is_fe_supports_chunking():
-            return False
-
-        # Check quantization sanity
-        if (int(self.is_per_act_token_quant) +
-                int(self.is_per_tensor_act_quant) +
-                int(self.quant_block_shape is not None)) > 1:
-            # invalid quant config
-            return False
-
-        # check bf16 / fp16 support
-        is_16bit = (self.dtype.itemsize == 2 and self.quant_dtype is None)
-        if is_16bit and not self.is_fe_16bit_supported():
-            return False
-
-        # Check fp8 support
-        is_fp8 = self.quant_dtype == torch.float8_e4m3fn
-        if is_fp8 and not self.is_fe_fp8_supported():
-            return False
-
-        # Check fp8 block quanization support
-        is_block_quatized = self.quant_block_shape is not None
-        if is_block_quatized and not is_fp8:
-            return False
-        if is_block_quatized and not self.is_fe_block_fp8_supported():
-            return False
-
-        # deep_gemm only works with block-quantized
-        if self.needs_deep_gemm() and not is_block_quatized:
-            return False
-
-        # Check dependencies
-        if self.needs_deep_ep() and not has_deep_ep():
-            return False
-        if self.needs_deep_gemm() and not has_deep_gemm():
-            return False
-        if self.needs_pplx() and not has_pplx():  # noqa: SIM103
-            return False
-
-        return True
-
-
-@dataclass
-class WeightTensors:
-    w1: torch.Tensor
-    w2: torch.Tensor
-    w1_scale: Optional[torch.Tensor]
-    w2_scale: Optional[torch.Tensor]
-
-    def describe(self):
-        s = ""
-        s += "== Weight Tensors: \n"
-        s += f' - {_describe_tensor(self.w1, "w1")} \n'
-        s += f' - {_describe_tensor(self.w2, "w2")} \n'
-        s += f' - {_describe_tensor(self.w1_scale, "w1_scale")} \n'
-        s += f' - {_describe_tensor(self.w2_scale, "w2_scale")} \n'
-        return s
-
-    def to_current_device(self):
-        self.w1 = self.w1.to(device=torch.cuda.current_device())
-        self.w2 = self.w2.to(device=torch.cuda.current_device())
-        is_quantized = self.w1.dtype == torch.float8_e4m3fn
-        if is_quantized:
-            assert self.w1_scale is not None
-            assert self.w2_scale is not None
-            self.w1_scale = self.w1_scale.to(
-                device=torch.cuda.current_device())
-            self.w2_scale = self.w2_scale.to(
-                device=torch.cuda.current_device())
-
-    def slice_weights(self, rank: int,
-                      num_local_experts: int) -> "WeightTensors":
-        s = rank * num_local_experts
-        e = s + num_local_experts
-        w1 = self.w1[s:e, :, :]
-        w2 = self.w2[s:e, :, :]
-        is_quantized = self.w1.dtype == torch.float8_e4m3fn
-        w1_scale, w2_scale = (None, None)
-        if is_quantized:
-            assert self.w1_scale is not None
-            assert self.w2_scale is not None
-            w1_scale = self.w1_scale[s:e, :, :]
-            w2_scale = self.w2_scale[s:e, :, :]
-        return WeightTensors(w1, w2, w1_scale, w2_scale)
-
-    @staticmethod
-    def make(config: Config) -> "WeightTensors":
-
-        if config.quant_dtype is None:
-            # just make normal dtype weights
-            w1, w2 = make_non_quant_weights(e=config.E,
-                                            n=config.N,
-                                            k=config.K,
-                                            dtype=config.dtype)
-            return WeightTensors(w1=w1, w2=w2, w1_scale=None, w2_scale=None)
-
-        assert config.quant_dtype == torch.float8_e4m3fn
-        if not config.is_fp8_block_quantized():
-            w1, w2, w1_scale, w2_scale = make_quant_fp8_weights(
-                e=config.E,
-                n=config.N,
-                k=config.K,
-                per_out_channel_quant=config.is_per_out_ch_quant,
-            )
-            return WeightTensors(w1=w1,
-                                 w2=w2,
-                                 w1_scale=w1_scale,
-                                 w2_scale=w2_scale)
-
-        assert config.quant_block_shape is not None
-        w1, w2, w1_scale, w2_scale = make_block_quant_fp8_weights(
-            e=config.E,
-            n=config.N,
-            k=config.K,
-            block_size=config.quant_block_shape,
-        )
-        return WeightTensors(w1=w1,
-                             w2=w2,
-                             w1_scale=w1_scale,
-                             w2_scale=w2_scale)
-
-
-@dataclass
-class RankTensors:
-    hidden_states: torch.Tensor
-    hidden_states_scale: Optional[torch.Tensor]
-
-    topk_weights: torch.Tensor
-    topk_ids: torch.Tensor
-    expert_map: Optional[torch.Tensor]
-
-    quant_config: Optional[FusedMoEQuantConfig]
-
-    def describe(self):
-        s = ""
-        s += "== Rank Tensors: \n"
-        s += f' - {_describe_tensor(self.hidden_states, "HS")} \n'
-        s += f' - {_describe_tensor(self.hidden_states_scale, "HS_scale")} \n'
-        s += f' - {_describe_tensor(self.topk_weights, "topk_weights")} \n'
-        s += f' - {_describe_tensor(self.topk_ids, "topk_ids")} \n'
-        s += f' - {_describe_tensor(self.expert_map, "expert_map")} \n'
-        return s
-
-    @staticmethod
-    def make_hidden_states(
-            config: Config) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
-        """
-        Return hidden_states
-        """
-        m, k, dtype = (config.M, config.K, config.dtype)
-        a = (torch.randn(
-            (m, k), device=torch.cuda.current_device(), dtype=dtype) / 15.0)
-
-        if config.quant_dtype is None:
-            return a, None
-
-        # We dequant and use that as hidden_states so the tests are stable.
-        # quantizing and dequantizing yield slightly different results
-        # depending on the hardware. Here we, quantize and dequantize
-        # first - so further quantize and dequantize will yeild the same
-        # values.
-        if config.is_per_tensor_act_quant:
-            a_q, a_scales = ops.scaled_fp8_quant(
-                a, use_per_token_if_dynamic=False)
-            return a_q.float().mul(a_scales).to(dtype), a_scales
-
-        if config.is_per_act_token_quant:
-            a_q, a_scales = ops.scaled_fp8_quant(a,
-                                                 use_per_token_if_dynamic=True)
-            return a_q.float().mul(a_scales).to(dtype), None
-
-        assert config.quant_block_shape is not None
-        block_k = config.quant_block_shape[1]
-        a_q, a_scales = per_token_cast_to_fp8(a, block_size=block_k)
-        return a_q.float().view(
-            (-1, block_k)).mul(a_scales.view(-1, 1)).view(m, k).to(dtype), None
-
-    @staticmethod
-    def make(config: Config, pgi: ProcessGroupInfo):
-
-        dtype = config.dtype
-        topk, m, _ = (config.topk, config.M, config.K)
-        hidden_states, hidden_states_scale = RankTensors.make_hidden_states(
-            config)
-
-        num_local_experts, global_num_experts = (config.num_local_experts,
-                                                 config.E)
-        score = torch.randn((m, global_num_experts),
-                            device="cuda",
-                            dtype=dtype)
-        topk_weights, topk_ids, _ = fused_topk(hidden_states, score, topk,
-                                               False)
-        topk_ids = topk_ids.to(config.topk_ids_dtype)
-
-        # distribute topk_ids evenly
-        for mi in range(m):
-            topk_ids[mi] = torch.randperm(config.E)[:topk]
-        topk_ids = topk_ids.to(device=torch.cuda.current_device())
-
-        expert_map = None
-        if config.world_size > 1:
-            expert_map = torch.full((global_num_experts, ),
-                                    fill_value=-1,
-                                    dtype=torch.int32)
-            s = pgi.rank * num_local_experts
-            e = s + num_local_experts
-            expert_map[s:e] = torch.tensor(list(range(num_local_experts)))
-            expert_map = expert_map.to(device=torch.cuda.current_device(),
-                                       dtype=torch.int32)
-
-        return RankTensors(
-            hidden_states=hidden_states,
-            hidden_states_scale=hidden_states_scale,
-            topk_weights=topk_weights,
-            topk_ids=topk_ids,
-            expert_map=expert_map,
-            quant_config=config.quant_config,
-        )
-
-
-def reference_moe_impl(config: Config, weights: WeightTensors,
-                       rank_tensors: RankTensors) -> torch.Tensor:
-
-    return torch_experts(a=rank_tensors.hidden_states,
-                         w1=weights.w1,
-                         w2=weights.w2,
-                         topk_weight=rank_tensors.topk_weights,
-                         topk_ids=rank_tensors.topk_ids,
-                         global_num_experts=config.E,
-                         expert_map=None,
-                         w1_scale=weights.w1_scale,
-                         w2_scale=weights.w2_scale,
-                         a1_scale=rank_tensors.hidden_states_scale,
-                         quant_dtype=config.quant_dtype,
-                         per_act_token_quant=config.is_per_act_token_quant,
-                         block_shape=config.quant_block_shape,
-                         apply_router_weights_on_input=config.topk == 1)
-
-
-def make_fused_experts(
-        config: Config, moe: FusedMoEConfig,
-        num_dispatchers: int) -> mk.FusedMoEPermuteExpertsUnpermute:
-
-    use_fp8 = config.quant_dtype == torch.float8_e4m3fn
-    batch_kwargs = {
-        "max_num_tokens": moe.max_num_tokens,
-        "num_dispatchers": num_dispatchers,
-    }
-    quant_kwargs = {
-        "use_fp8_w8a8": use_fp8,
-        "use_int8_w8a8": False,
-        "use_int8_w8a16": False,
-        "use_int4_w4a16": False,
-        "block_shape": config.quant_block_shape,
-        "per_act_token_quant": config.is_per_act_token_quant,
-    }
-    deepgemm_kwargs = {"allow_deep_gemm": has_deep_gemm()}
-
-    if config.fused_experts_type == BatchedDeepGemmExperts:
-        kwargs = batch_kwargs | {
-            "block_shape": config.quant_block_shape,
-            "per_act_token_quant": config.is_per_act_token_quant,
-        }
-        print(f"Making BatchedDeepGemmExperts {kwargs} ...")
-        experts = BatchedDeepGemmExperts(**kwargs)
-    elif config.fused_experts_type == BatchedTritonExperts:
-        kwargs = batch_kwargs | quant_kwargs
-        print(f"Making BatchedTritonExperts {kwargs} ...")
-        experts = BatchedTritonExperts(**kwargs)
-    elif config.fused_experts_type == BatchedTritonOrDeepGemmExperts:
-        kwargs = batch_kwargs | quant_kwargs | deepgemm_kwargs
-        print(f"Making BatchedTritonOrDeepGemmExperts {kwargs} ...")
-        experts = BatchedTritonOrDeepGemmExperts(**kwargs)
-    elif config.fused_experts_type == DeepGemmExperts:
-        print("Making DeepGemmExperts () ...")
-        experts = DeepGemmExperts()
-    elif config.fused_experts_type == TritonExperts:
-        kwargs = quant_kwargs
-        print(f"Making TritonExperts {kwargs} ...")
-        experts = TritonExperts(**kwargs)
-    elif config.fused_experts_type == TritonOrDeepGemmExperts:
-        kwargs = quant_kwargs | deepgemm_kwargs
-        print(f"Making TritonOrDeepGemmExperts {kwargs} ...")
-        experts = TritonOrDeepGemmExperts(**kwargs)
-    elif config.fused_experts_type == NaiveBatchedExperts:
-        kwargs = batch_kwargs | quant_kwargs
-        print(f"Making NaiveBatchedExperts {kwargs} ...")
-        experts = NaiveBatchedExperts(**kwargs)
-    elif config.fused_experts_type == CutlassExpertsFp8:
-        use_batched_format = config.is_batched_prepare_finalize()
-        num_experts = (moe.num_local_experts
-                       if use_batched_format else moe.num_experts)
-        kwargs = {
-            "max_experts_per_worker": num_experts,
-            "out_dtype": moe.in_dtype,
-            "per_act_token_quant": config.is_per_act_token_quant,
-            "per_out_ch_quant": config.is_per_out_ch_quant,
-            "block_shape": config.quant_block_shape,
-            "num_dispatchers": num_dispatchers,
-            "use_batched_format": use_batched_format
-        }
-        print(f"Making CutlassExpertsFp8 {kwargs} ...")
-        experts = CutlassExpertsFp8(**kwargs)
-
-    return experts
-
-
-def make_modular_kernel(config: Config,
-                        vllm_config: VllmConfig) -> mk.FusedMoEModularKernel:
-
-    def next_power_of_2(x):
-        import math
-        if x == 0:
-            return 1
-        return 2**math.ceil(math.log2(x))
-
-    # make moe config
-    moe_parallel_config: FusedMoEParallelConfig = FusedMoEParallelConfig.make(
-        tp_size_=get_tensor_model_parallel_world_size(),
-        dp_size_=get_dp_group().world_size,
-        vllm_parallel_config=vllm_config.parallel_config,
-    )
-    moe = FusedMoEConfig(
-        num_experts=config.E,
-        experts_per_token=config.topk,
-        hidden_dim=config.K,
-        num_local_experts=config.num_local_experts,
-        moe_parallel_config=moe_parallel_config,
-        in_dtype=config.dtype,
-        quant_config=config.quant_config,
-        max_num_tokens=next_power_of_2(config.M),
-    )
-
-    # make modular kernel
-    prepare_finalize = None
-    if config.needs_all2all():
-        prepare_finalize = FusedMoEMethodBase.maybe_make_prepare_finalize(moe)
-        assert prepare_finalize is not None
-    else:
-        prepare_finalize = MoEPrepareAndFinalizeNoEP()
-
-    fused_experts = make_fused_experts(config, moe,
-                                       prepare_finalize.num_dispatchers())
-
-    modular_kernel = mk.FusedMoEModularKernel(
-        prepare_finalize=prepare_finalize, fused_experts=fused_experts)
-
-    return modular_kernel
-
-
-def run_modular_kernel(
-    pgi: ProcessGroupInfo,
-    vllm_config: VllmConfig,
-    config: Config,
-    weights: WeightTensors,
-    rank_tensors: RankTensors,
-) -> torch.Tensor:
-    assert isinstance(config.Ms, int)
-    assert isinstance(config.topks, int)
-
-    # weights for rank
-    rank_weights = weights.slice_weights(pgi.rank, config.num_local_experts)
-
-    mk = make_modular_kernel(config, vllm_config)
-
-    mk_kwargs = {
-        "hidden_states": rank_tensors.hidden_states.clone(
-        ),  # impls might update the tensor in place
-        "w1": rank_weights.w1,
-        "w2": rank_weights.w2,
-        "topk_weights": rank_tensors.topk_weights,
-        "topk_ids": rank_tensors.topk_ids,
-        "expert_map": rank_tensors.expert_map,
-        "w1_scale": rank_weights.w1_scale,
-        "w2_scale": rank_weights.w2_scale,
-        "a1_scale": rank_tensors.hidden_states_scale,
-        "global_num_experts": config.E,
-        "apply_router_weight_on_input": config.topk == 1,
-    }
-    out = mk.forward(**mk_kwargs)
-
-    return out
@@ -1,173 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import copy
-from enum import Enum
-from itertools import product
-from typing import Optional
-
-import torch
-from tqdm import tqdm
-
-from vllm.config import VllmConfig, set_current_vllm_config
-from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
-from vllm.platforms import current_platform
-
-from .common import (Config, RankTensors, WeightTensors, reference_moe_impl,
-                     run_modular_kernel)
-from .mk_objects import (MK_FUSED_EXPERT_TYPES,
-                         MK_MULTI_GPU_PREPARE_FINALIZE_TYPES, MK_QUANT_CONFIGS)
-from .parallel_utils import ProcessGroupInfo, parallel_launch_with_config
-
-
-class Result(Enum):
-    PASS = 1
-    FAIL = 2
-    SKIP = 3
-
-
-def rank_worker(
-    pgi: ProcessGroupInfo,
-    vllm_config: VllmConfig,
-    cpu_group,
-    config: Config,
-    weights: WeightTensors,
-):
-    current_platform.seed_everything(pgi.rank)
-
-    # sanity check
-    from vllm import envs
-    if config.fused_moe_chunk_size is not None:
-        assert (config.fused_moe_chunk_size == envs.VLLM_FUSED_MOE_CHUNK_SIZE)
-
-    # get weights to this device
-    weights.to_current_device()
-
-    Ms = config.Ms
-    assert isinstance(Ms, list)
-    TOPKs = config.topks
-    assert isinstance(TOPKs, list)
-
-    for m, topk in product(Ms, TOPKs):
-        print(f"Running m={m}, topk={topk} ...")
-        # override m and topk
-        cfgx = copy.deepcopy(config)
-        cfgx.Ms = m
-        cfgx.topks = topk
-
-        # inputs for rank
-        rank_tensors = RankTensors.make(cfgx, pgi)
-
-        # modular kernel out
-        mk_out = run_modular_kernel(pgi, vllm_config, cfgx, weights,
-                                    rank_tensors)
-
-        with set_current_vllm_config(vllm_config):
-            ref_out = reference_moe_impl(cfgx, weights, rank_tensors)
-
-        torch.testing.assert_close(ref_out, mk_out, atol=3e-2, rtol=3e-2)
-
-
-def make_feature_matrix(csv_file_path: str):
-
-    from dataclasses import asdict
-
-    import pandas as pd
-
-    def add_to_results(config: Config,
-                       success: Result,
-                       results_df: Optional[pd.DataFrame] = None):
-        config_dict = asdict(config)
-        config_dict['prepare_finalize_type'] = config_dict[
-            'prepare_finalize_type'].__name__
-        config_dict['fused_experts_type'] = config_dict[
-            'fused_experts_type'].__name__
-        config_dict['per_tensor_act_quant'] = config.is_per_tensor_act_quant
-        quant_config_dict = config_dict['quant_config']
-        del config_dict['quant_config']
-        if quant_config_dict is None:
-            quant_config = FusedMoEQuantConfig(None)
-            quant_config_dict = asdict(quant_config)
-
-        config_dict |= quant_config_dict
-        result_dict = config_dict | {'success': success.name}
-
-        result_df = pd.DataFrame([result_dict])
-        if results_df is None:
-            results_df = result_df
-        else:
-            results_df = pd.concat([results_df, result_df], ignore_index=True)
-
-        return results_df
-
-    Ms = [64]
-    Ks = [7168]  # hidden sizes
-    Ns = [2048]
-    TOPKs = [[4, 1]]
-    Es = [32]
-    DTYPEs = [torch.bfloat16]
-    PF_TYPES = MK_MULTI_GPU_PREPARE_FINALIZE_TYPES
-    FE_TYPES = MK_FUSED_EXPERT_TYPES
-    Q_TYPES = MK_QUANT_CONFIGS
-
-    combinations = list(
-        product(Ms, Ks, Ns, Es, TOPKs, DTYPEs, PF_TYPES, FE_TYPES, Q_TYPES))
-
-    results_df: Optional[pd.DataFrame] = None
-    for m, k, n, e, topks, dtype, pf_type, experts_type, quant_config in tqdm(
-            combinations):  #noqa: E501
-        config = Config(Ms=[m],
-                        K=k,
-                        N=n,
-                        E=e,
-                        topks=topks,
-                        dtype=dtype,
-                        prepare_finalize_type=pf_type,
-                        fused_experts_type=experts_type,
-                        quant_config=quant_config,
-                        world_size=2,
-                        fused_moe_chunk_size=None)
-
-        success = None
-        if config.is_valid():
-            print(f"Running config : {config.describe()} ...")
-            try:
-                weights: WeightTensors = WeightTensors.make(config)
-                vllm_config, env_dict = config.make_env_data()
-                parallel_launch_with_config(config.world_size, rank_worker,
-                                            vllm_config, env_dict, config,
-                                            weights)
-                success = Result.PASS
-            except Exception as _:
-                success = Result.FAIL
-        else:
-            success = Result.SKIP
-
-        results_df = add_to_results(config, success, results_df)
-
-    if results_df is not None:
-        results_df.to_csv(f"{csv_file_path}")
-
-
-if __name__ == '__main__':
-    import argparse
-    from pathlib import Path
-    parser = argparse.ArgumentParser(description=(
-        "Make ModularKernel feature matrix \n"
-        "Example : python3 -m tests.kernels.moe.modular_kernel_tools.make_feature_matrix "  #noqa: E501
-        "-f ./feature_matrices/feature_matrix.csv"))
-
-    parser.add_argument("-f",
-                        "--feature-matrix-csv-file-path",
-                        type=str,
-                        required=True,
-                        help="File name to Generate a .csv file")
-    args = parser.parse_args()
-
-    csv_path = args.feature_matrix_csv_file_path
-    assert csv_path.endswith(
-        'csv'), f"Need a file path ending with .csv, got {csv_path}"
-    assert Path(csv_path).parent.is_dir(
-    ), f"Cannot find parent directory for {Path(csv_path).parent}"
-
-    make_feature_matrix(args.feature_matrix_csv_file_path)
@@ -1,87 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import torch
-
-# Fused experts and PrepareFinalize imports
-from vllm.model_executor.layers.fused_moe.batched_deep_gemm_moe import (
-    BatchedDeepGemmExperts)
-from vllm.model_executor.layers.fused_moe.batched_triton_or_deep_gemm_moe import (  # noqa: E501
-    BatchedTritonOrDeepGemmExperts)
-from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
-from vllm.model_executor.layers.fused_moe.cutlass_moe import CutlassExpertsFp8
-from vllm.model_executor.layers.fused_moe.deep_gemm_moe import DeepGemmExperts
-from vllm.model_executor.layers.fused_moe.fused_batched_moe import (
-    BatchedTritonExperts, NaiveBatchedExperts)
-from vllm.model_executor.layers.fused_moe.layer import TritonExperts
-from vllm.model_executor.layers.fused_moe.prepare_finalize import (
-    MoEPrepareAndFinalizeNoEP)
-from vllm.model_executor.layers.fused_moe.triton_deep_gemm_moe import (
-    TritonOrDeepGemmExperts)
-from vllm.utils import has_deep_ep, has_pplx
-
-if has_deep_ep():
-    from vllm.model_executor.layers.fused_moe.deepep_ht_prepare_finalize import (  # noqa: E501
-        DeepEPHTPrepareAndFinalize)
-    from vllm.model_executor.layers.fused_moe.deepep_ll_prepare_finalize import (  # noqa: E501
-        DeepEPLLPrepareAndFinalize)
-
-if has_pplx():
-    from vllm.model_executor.layers.fused_moe.pplx_prepare_finalize import (
-        PplxPrepareAndFinalize)
-
-MK_MULTI_GPU_PREPARE_FINALIZE_TYPES = []
-if has_pplx():
-    MK_MULTI_GPU_PREPARE_FINALIZE_TYPES += [PplxPrepareAndFinalize]
-if has_deep_ep():
-    MK_MULTI_GPU_PREPARE_FINALIZE_TYPES += [
-        DeepEPHTPrepareAndFinalize, DeepEPLLPrepareAndFinalize
-    ]
-
-MK_SINGLE_GPU_PREPARE_FINALIZE_TYPES = [MoEPrepareAndFinalizeNoEP]
-
-MK_ALL_PREPARE_FINALIZE_TYPES = (MK_MULTI_GPU_PREPARE_FINALIZE_TYPES +
-                                 MK_SINGLE_GPU_PREPARE_FINALIZE_TYPES)
-
-MK_FUSED_EXPERT_TYPES = [
-    BatchedDeepGemmExperts,
-    BatchedTritonExperts,
-    NaiveBatchedExperts,
-    BatchedTritonOrDeepGemmExperts,
-    CutlassExpertsFp8,
-    DeepGemmExperts,
-    TritonOrDeepGemmExperts,
-    TritonExperts,
-]
-
-MK_QUANT_CONFIGS = [
-    None,
-    # per-channel / per-column weights and per-tensor activations
-    FusedMoEQuantConfig(quant_dtype=torch.float8_e4m3fn,
-                        per_out_ch_quant=True,
-                        per_act_token_quant=False,
-                        block_shape=None),
-    # per-channel / per-column weights and per-token activations
-    FusedMoEQuantConfig(quant_dtype=torch.float8_e4m3fn,
-                        per_out_ch_quant=True,
-                        per_act_token_quant=True,
-                        block_shape=None),
-    # per-tensor weights and per-tensor activations
-    FusedMoEQuantConfig(quant_dtype=torch.float8_e4m3fn,
-                        per_out_ch_quant=False,
-                        per_act_token_quant=False,
-                        block_shape=None),
-    # per-tensor weights and per-token activations
-    FusedMoEQuantConfig(quant_dtype=torch.float8_e4m3fn,
-                        per_out_ch_quant=False,
-                        per_act_token_quant=True,
-                        block_shape=None),
-    # block-quantized weights and 128 block per-token activations
-    FusedMoEQuantConfig(quant_dtype=torch.float8_e4m3fn,
-                        per_out_ch_quant=False,
-                        per_act_token_quant=False,
-                        block_shape=[128, 128]),
-    # TODO (varun) : Should we test the following combinations ?
-    # block-quantized weights and per-token activations
-    # block-quantized weights and per-tensor activations
-]
@@ -1,138 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-import dataclasses
-import os
-import traceback
-from typing import Any, Callable, Optional
-
-import torch
-from torch.multiprocessing import (
-    spawn)  # pyright: ignore[reportPrivateImportUsage]
-from typing_extensions import Concatenate, ParamSpec
-
-from vllm.config import VllmConfig, set_current_vllm_config
-from vllm.distributed import (init_distributed_environment,
-                              initialize_model_parallel)
-from vllm.utils import get_open_port
-
-## Parallel Processes Utils
-
-P = ParamSpec("P")
-
-
-@dataclasses.dataclass
-class ProcessGroupInfo:
-    world_size: int
-    world_local_size: int
-    rank: int
-    node_rank: int
-    local_rank: int
-    device: torch.device
-
-
-def _set_vllm_config(vllm_config: VllmConfig, world_size: int, rank: int,
-                     local_rank: int):
-
-    import tempfile
-    temp_file = tempfile.mkstemp()[1]
-
-    set_current_vllm_config(vllm_config)
-    with set_current_vllm_config(vllm_config):
-        init_distributed_environment(
-            world_size=world_size,
-            rank=rank,
-            distributed_init_method=f"file://{temp_file}",
-            local_rank=local_rank,
-            backend="nccl",
-        )
-
-        initialize_model_parallel(
-            tensor_model_parallel_size=vllm_config.parallel_config.
-            tensor_parallel_size,
-            pipeline_model_parallel_size=vllm_config.parallel_config.
-            pipeline_parallel_size,
-        )
-        cpu_group = torch.distributed.new_group(list(range(world_size)),
-                                                backend="gloo")
-    return cpu_group
-
-
-def _worker_parallel_launch(
-    local_rank: int,
-    world_size: int,
-    world_local_size: int,
-    node_rank: int,
-    init_method: str,
-    worker: Callable[Concatenate[ProcessGroupInfo, Optional[VllmConfig], Any,
-                                 P], None],
-    vllm_config: Optional[VllmConfig],
-    env_dict: Optional[dict],
-    *args: P.args,
-    **kwargs: P.kwargs,
-) -> None:
-    rank = node_rank * world_local_size + local_rank
-    torch.cuda.set_device(local_rank)
-    device = torch.device("cuda", local_rank)
-    torch.distributed.init_process_group(
-        backend="cpu:gloo,cuda:nccl",
-        init_method=init_method,
-        rank=rank,
-        world_size=world_size,
-        device_id=device,
-    )
-    barrier = torch.tensor([rank], device=device)
-    torch.distributed.all_reduce(barrier)
-
-    if env_dict is not None:
-        os.environ.update(env_dict)
-
-    cpu_group = None
-    if vllm_config is not None:
-        cpu_group = _set_vllm_config(vllm_config, world_size, rank, local_rank)
-
-    try:
-        worker(
-            ProcessGroupInfo(
-                world_size=world_size,
-                world_local_size=world_local_size,
-                rank=rank,
-                node_rank=node_rank,
-                local_rank=local_rank,
-                device=device,
-            ),
-            vllm_config,
-            cpu_group,
-            *args,
-            **kwargs,
-        )
-    except Exception as ex:
-        print(ex)
-        traceback.print_exc()
-        raise
-    finally:
-        torch.distributed.destroy_process_group()
-
-
-def parallel_launch_with_config(
-    world_size: int,
-    worker: Callable[Concatenate[ProcessGroupInfo, VllmConfig, Any, P], None],
-    vllm_config: VllmConfig,
-    env_dict: dict[Any, Any],
-    *args: P.args,
-    **kwargs: P.kwargs,
-) -> None:
-    assert not kwargs
-    spawn(
-        _worker_parallel_launch,
-        args=(
-            world_size,
-            world_size,
-            0,
-            f"tcp://{os.getenv('LOCALHOST', 'localhost')}:{get_open_port()}",
-            worker,
-            vllm_config,
-            env_dict,
-        ) + args,
-        nprocs=world_size,
-        join=True,
-    )
@@ -1,127 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import copy
-from itertools import product
-from typing import Any, Callable
-
-import torch
-
-from vllm.config import VllmConfig
-from vllm.platforms import current_platform
-
-from .common import Config, RankTensors, WeightTensors, make_modular_kernel
-from .parallel_utils import ProcessGroupInfo, parallel_launch_with_config
-
-
-def do_profile(fn: Callable,
-               fn_kwargs: dict[Any, Any],
-               pgi: ProcessGroupInfo,
-               config: Config,
-               num_warmups: int = 5):
-    for _ in range(num_warmups):
-        fn(**fn_kwargs)
-
-    with torch.profiler.profile(
-            activities=[
-                torch.profiler.ProfilerActivity.CPU,
-                torch.profiler.ProfilerActivity.CUDA,
-            ],
-            with_stack=True,
-            record_shapes=True,
-    ) as tprof:
-        fn(**fn_kwargs)
-        torch.cuda.synchronize(torch.cuda.current_device())
-
-    # TODO (varun): Add a descriptive trace file name
-    tprof.export_chrome_trace(
-        f"{config.torch_trace_dir_path}/m{config.M}_{pgi.rank}_trace.json")
-
-
-def profile_modular_kernel(
-    pgi: ProcessGroupInfo,
-    vllm_config: VllmConfig,
-    config: Config,
-    weights: WeightTensors,
-    rank_tensors: RankTensors,
-) -> None:
-    assert isinstance(config.Ms, int)
-    assert isinstance(config.topks, int)
-
-    # weights for rank
-    rank_weights = weights.slice_weights(pgi.rank, config.num_local_experts)
-
-    # make modular kernel
-    mk = make_modular_kernel(config, vllm_config)
-
-    mk_kwargs = {
-        "hidden_states": rank_tensors.hidden_states,
-        "w1": rank_weights.w1,
-        "w2": rank_weights.w2,
-        "topk_weights": rank_tensors.topk_weights,
-        "topk_ids": rank_tensors.topk_ids,
-        "expert_map": rank_tensors.expert_map,
-        "w1_scale": rank_weights.w1_scale,
-        "w2_scale": rank_weights.w2_scale,
-        "a1_scale": rank_tensors.hidden_states_scale,
-        "global_num_experts": config.E,
-        "apply_router_weight_on_input": config.topk == 1,
-    }
-
-    do_profile(mk.forward, mk_kwargs, pgi, config)
-
-
-def rank_worker(
-    pgi: ProcessGroupInfo,
-    vllm_config: VllmConfig,
-    cpu_group,
-    config: Config,
-    weights: WeightTensors,
-):
-    current_platform.seed_everything(pgi.rank)
-
-    # sanity check
-    from vllm import envs
-    if config.fused_moe_chunk_size is not None:
-        assert (config.fused_moe_chunk_size == envs.VLLM_FUSED_MOE_CHUNK_SIZE)
-
-    # get weights to this device
-    weights.to_current_device()
-
-    Ms = config.Ms
-    assert isinstance(Ms, list)
-    TOPKs = config.topks
-    assert isinstance(TOPKs, list)
-
-    for m, topk in product(Ms, TOPKs):
-        print(f"Running m={m}, topk={topk} ...")
-        # override m and topk
-        cfgx = copy.deepcopy(config)
-        cfgx.Ms = m
-        cfgx.topks = topk
-
-        # inputs for rank
-        rank_tensors = RankTensors.make(cfgx, pgi)
-        profile_modular_kernel(pgi, vllm_config, cfgx, weights, rank_tensors)
-
-
-def run(config: Config):
-    weights: WeightTensors = WeightTensors.make(config)
-    vllm_config, env_dict = config.make_env_data()
-    parallel_launch_with_config(config.world_size, rank_worker, vllm_config,
-                                env_dict, config, weights)
-
-
-if __name__ == '__main__':
-    from .cli_args import make_config, make_config_arg_parser
-    parser = make_config_arg_parser(description=(
-        "Run single prepare-finalize & fused-experts combination test"
-        "Example : python3 -m tests.kernels.moe.modular_kernel_tools.profile_modular_kernel "  #noqa: E501
-        "--pf-type PplxPrepareAndFinalize --experts-type BatchedTritonExperts"
-    ))
-    args = parser.parse_args()
-    assert args.torch_trace_dir_path is not None, (
-        "Please pass in a directory to store torch traces")
-    config = make_config(args)
-
-    run(config)
@@ -1,142 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-import math
-
-import torch
-
-import vllm._custom_ops as ops
-
-
-def per_token_cast_to_fp8(
-        x: torch.Tensor, block_size: int) -> tuple[torch.Tensor, torch.Tensor]:
-    assert x.dim() == 2
-    m, n = x.shape
-    pad_size = (block_size - (n % block_size)) % block_size
-    x = torch.nn.functional.pad(x,
-                                (0, pad_size), value=0) if pad_size > 0 else x
-    x_view = x.view(m, -1, block_size)
-    x_amax = x_view.abs().float().amax(dim=2).view(m, -1).clamp(1e-4)
-    fp8_data = (x_view * (448.0 / x_amax.unsqueeze(2))).to(torch.float8_e4m3fn)
-    return fp8_data.view(m, n + pad_size)[:, :n], (x_amax / 448.0).view(m, -1)
-
-
-def per_block_cast_to_fp8(
-        x: torch.Tensor, block_size_k: int,
-        block_size_n: int) -> tuple[torch.Tensor, torch.Tensor]:
-    assert x.dim() == 2
-    m, n = x.shape
-    x_padded = torch.zeros(
-        (
-            int(math.ceil(m / block_size_k)) * block_size_k,
-            int(math.ceil(n / block_size_n)) * block_size_n,
-        ),
-        dtype=x.dtype,
-        device=x.device,
-    )
-    x_padded[:m, :n] = x
-    x_view = x_padded.view(-1, block_size_k,
-                           x_padded.size(1) // block_size_k, block_size_n)
-    x_amax = x_view.abs().float().amax(dim=(1, 3), keepdim=True).clamp(1e-4)
-    x_scaled = (x_view * (448.0 / x_amax)).to(torch.float8_e4m3fn)
-    x_scaled_sub = x_scaled.view_as(x_padded)[:m, :n].contiguous()
-    scales = (x_amax / 448.0).view(x_view.size(0), x_view.size(2))
-    return x_scaled_sub, scales
-
-
-def make_non_quant_weights(
-    e: int,
-    n: int,
-    k: int,
-    dtype: torch.dtype,
-) -> tuple[torch.Tensor, torch.Tensor]:
-    """
-    Return weights w1, w2
-    """
-    device = torch.cuda.current_device()
-    w1 = torch.randn((e, 2 * n, k), device=device, dtype=dtype) / 15
-    w2 = torch.randn((e, k, n), device=device, dtype=dtype) / 15
-    return w1, w2
-
-
-def make_block_quant_fp8_weights(
-    e: int,
-    n: int,
-    k: int,
-    block_size: list[int],
-) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
-    """
-    Return weights w1, w2, w1_scale, w2_scale
-    """
-    dtype = torch.bfloat16
-    device = torch.cuda.current_device()
-
-    fp8_info = torch.finfo(torch.float8_e4m3fn)
-    fp8_max, fp8_min = fp8_info.max, fp8_info.min
-
-    w1_bf16, w2_bf16 = make_non_quant_weights(e, n, k, dtype)
-    w1_bf16 = w1_bf16.clamp(min=fp8_min, max=fp8_max).to(dtype=dtype)
-    w2_bf16 = w2_bf16.clamp(min=fp8_min, max=fp8_max).to(dtype=dtype)
-
-    block_n, block_k = block_size[0], block_size[1]
-    n_tiles_w1 = ((2 * n) + block_n - 1) // block_n
-    k_tiles_w1 = (k + block_k - 1) // block_k
-    n_tiles_w2 = (k + block_n - 1) // block_n
-    k_tiles_w2 = (n + block_k - 1) // block_k
-
-    w1 = torch.empty_like(w1_bf16, dtype=torch.float8_e4m3fn, device=device)
-    w2 = torch.empty_like(w2_bf16, dtype=torch.float8_e4m3fn, device=device)
-
-    w1_s = torch.empty((e, n_tiles_w1, k_tiles_w1),
-                       device=device,
-                       dtype=torch.float32)
-    w2_s = torch.empty((e, n_tiles_w2, k_tiles_w2),
-                       device=device,
-                       dtype=torch.float32)
-
-    assert w1_s.shape == (e, (2 * n + (block_n - 1)) // block_n,
-                          (k + (block_k - 1)) // block_k)
-    assert (w2.shape[-2] + block_n - 1) // block_n == w2_s.shape[-2]
-
-    for i in range(e):
-        w1[i], w1_s[i] = per_block_cast_to_fp8(w1_bf16[i],
-                                               block_size_k=block_k,
-                                               block_size_n=block_n)
-        w2[i], w2_s[i] = per_block_cast_to_fp8(w2_bf16[i],
-                                               block_size_k=block_k,
-                                               block_size_n=block_n)
-
-    return w1, w2, w1_s, w2_s
-
-
-def make_quant_fp8_weights(
-    e: int,
-    n: int,
-    k: int,
-    per_out_channel_quant: bool,
-) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
-    """
-    Return w1, w2, w1_scale, w2_scale
-    """
-    q_dtype = torch.float8_e4m3fn
-
-    w1, w2 = make_non_quant_weights(e, n, k, dtype=torch.bfloat16)
-
-    # w1 -> w1_q, w2 -> w2_q
-    w1_q = torch.empty((e, 2 * n, k), device="cuda", dtype=q_dtype)
-    w2_q = torch.empty((e, k, n), device="cuda", dtype=q_dtype)
-
-    n_b_scales = 2 * n if per_out_channel_quant else 1
-    k_b_scales = k if per_out_channel_quant else 1
-    w1_scale = torch.empty((e, n_b_scales, 1),
-                           device="cuda",
-                           dtype=torch.float32)
-    w2_scale = torch.empty((e, k_b_scales, 1),
-                           device="cuda",
-                           dtype=torch.float32)
-
-    for expert in range(e):
-        w1_q[expert], w1_scale[expert] = ops.scaled_fp8_quant(
-            w1[expert], use_per_token_if_dynamic=per_out_channel_quant)
-        w2_q[expert], w2_scale[expert] = ops.scaled_fp8_quant(
-            w2[expert], use_per_token_if_dynamic=per_out_channel_quant)
-    return w1_q, w2_q, w1_scale, w2_scale
@@ -4,6 +4,7 @@
 DeepEP test utilities
 """
 import dataclasses
+import importlib
 import os
 import traceback
 from typing import Callable, Optional
@@ -14,9 +15,10 @@ from torch.multiprocessing import (
    spawn)  # pyright: ignore[reportPrivateImportUsage]
 from typing_extensions import Concatenate, ParamSpec

-from vllm.utils import get_open_port, has_deep_ep
+from vllm.utils import get_open_port

-if has_deep_ep():
+has_deep_ep = importlib.util.find_spec("deep_ep") is not None
+if has_deep_ep:
    from vllm.model_executor.layers.fused_moe.deepep_ht_prepare_finalize import (  # noqa: E501
        DeepEPHTPrepareAndFinalize)
    from vllm.model_executor.layers.fused_moe.deepep_ll_prepare_finalize import (  # noqa: E501
@@ -15,13 +15,13 @@ from vllm.model_executor.layers.fused_moe.deep_gemm_moe import (
 from vllm.model_executor.layers.fused_moe.fused_moe import (
    fused_topk, modular_triton_fused_moe)
 from vllm.platforms import current_platform
-from vllm.utils import has_deep_gemm
-from vllm.utils.deep_gemm import is_blackwell_deep_gemm_used

-dg_available = has_deep_gemm()
-
-if dg_available:
-    from deep_gemm import get_m_alignment_for_contiguous_layout
+dg_available = False
+try:
+    import deep_gemm
+    dg_available = True
+except ImportError:
+    pass

 if current_platform.get_device_capability() < (9, 0):
    pytest.skip("FP8 Triton requires CUDA 9.0 or higher",
@@ -224,7 +224,6 @@ def test_w8a8_block_fp8_fused_moe(M, N, K, E, topk, block_size, dtype, seed,
@pytest.mark.parametrize("topk", TOP_KS)
@pytest.mark.parametrize("seed", SEEDS)
@pytest.mark.skipif(not dg_available, reason="DeepGemm kernels not available.")
-@pytest.mark.skipif(is_blackwell_deep_gemm_used(), reason="Not E8M0 scale MOE")
@torch.inference_mode()
 def test_w8a8_block_fp8_deep_gemm_fused_moe(M, N, K, E, topk, seed,
                                            monkeypatch):
@@ -239,7 +238,8 @@ def test_w8a8_block_fp8_deep_gemm_fused_moe(M, N, K, E, topk, seed,
    torch.manual_seed(seed)

    monkeypatch.setenv("VLLM_FUSED_MOE_CHUNK_SIZE", str(chunk_size))
-    block_m = get_m_alignment_for_contiguous_layout()
+
+    block_m = deep_gemm.get_m_alignment_for_contiguous_layout()
    block_size = [block_m, block_m]
    dtype = torch.bfloat16

@@ -1,140 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Tests compute_expert_num_tokens kernels
-"""
-
-import dataclasses
-from typing import Optional
-
-import pytest
-import torch
-
-from vllm.model_executor.layers.fused_moe.utils import count_expert_num_tokens
-
-
-@dataclasses.dataclass
-class TestTensors:
-
-    topk_ids: torch.Tensor
-    expert_map: Optional[torch.Tensor] = None
-
-    def to_device(self, device: str):
-        self.topk_ids = self.topk_ids.to(device=device)
-        if self.expert_map is not None:
-            self.expert_map = self.expert_map.to(device=device)
-
-    @staticmethod
-    def make(num_tokens: int, num_topk: int, num_experts: int, device: str,
-             topk_ids_dtype: torch.dtype) -> "TestTensors":
-
-        # make topk ids
-        topk_ids = torch.empty((num_tokens, num_topk),
-                               device=device,
-                               dtype=torch.int64)
-        for x in range(num_tokens):
-            topk_ids[x] = torch.randperm(num_experts)[:num_topk]
-        topk_ids = topk_ids.to(dtype=torch.int64)
-        return TestTensors(topk_ids=topk_ids)
-
-    def with_ep_rank(self, ep_rank: int, num_global_experts: int,
-                     num_local_experts: int, device: str):
-        # make an expert map
-        expert_map = torch.empty((num_global_experts),
-                                 device=device,
-                                 dtype=torch.int32)
-        expert_map.fill_(-1)
-        s = ep_rank * num_local_experts
-        e = s + num_local_experts
-        expert_map[s:e] = torch.tensor(list(range(num_local_experts)),
-                                       device=device)
-
-        return TestTensors(topk_ids=self.topk_ids.clone(),
-                           expert_map=expert_map)
-
-
-def ref_impl(tt: TestTensors, expert_num_tokens: torch.Tensor):
-    # do the reference in cpu
-    tt.to_device("cpu")
-    expert_ids, counts = tt.topk_ids.unique(return_counts=True)
-
-    for eid, count in zip(expert_ids, counts):
-        if eid != -1 and tt.expert_map is not None:
-            eid = tt.expert_map[eid]
-
-        if eid == -1:
-            continue
-
-        expert_num_tokens[eid] += count
-
-
-def do_test_compute_expert_num_tokens(num_tokens: int, num_topk: int,
-                                      num_experts: int, ep_size: int,
-                                      topk_ids_dtype: torch.dtype):
-
-    assert num_topk <= num_experts
-
-    tt = TestTensors.make(num_tokens,
-                          num_topk,
-                          num_experts,
-                          topk_ids_dtype=topk_ids_dtype,
-                          device="cpu")
-
-    num_global_experts = num_experts
-    assert num_global_experts % ep_size == 0
-    num_local_experts = num_global_experts // ep_size
-    for ep_rank in range(ep_size):
-        tt_rank = tt.with_ep_rank(ep_rank, num_global_experts,
-                                  num_local_experts, "cpu")
-
-        ref_expert_num_tokens = torch.zeros((num_local_experts),
-                                            device="cpu",
-                                            dtype=torch.int32)
-        ref_impl(tt_rank, ref_expert_num_tokens)
-        ref_expert_num_tokens = ref_expert_num_tokens.to("cuda")
-
-        tt_rank.to_device("cuda")
-        # Test with expert_map
-        triton_expert_num_tokens_w_emap = count_expert_num_tokens(
-            tt_rank.topk_ids, num_local_experts, tt_rank.expert_map)
-
-        # Test without expert map
-        topk_ids = tt_rank.expert_map[tt_rank.topk_ids].to(topk_ids_dtype)
-        triton_expert_num_tokens_wo_emap = count_expert_num_tokens(
-            topk_ids, num_local_experts, expert_map=None)
-
-        torch.testing.assert_close(ref_expert_num_tokens,
-                                   triton_expert_num_tokens_w_emap,
-                                   atol=0,
-                                   rtol=0)
-        torch.testing.assert_close(ref_expert_num_tokens,
-                                   triton_expert_num_tokens_wo_emap,
-                                   atol=0,
-                                   rtol=0)
-
-
-@pytest.mark.parametrize(
-    "num_tokens", [1, 4, 8, 11, 19, 128, 127, 405, 1024, 3333, 6666, 7317])
-@pytest.mark.parametrize("num_topk", [2, 6, 8])
-@pytest.mark.parametrize("num_experts", [64])
-@pytest.mark.parametrize("ep_size", [1, 2, 4])
-@pytest.mark.parametrize("topk_ids_dtype", [torch.int64])
-def test_compute_expert_num_tokens(num_tokens: int, num_topk: int,
-                                   num_experts: int, ep_size: int,
-                                   topk_ids_dtype: torch.dtype):
-    do_test_compute_expert_num_tokens(num_tokens, num_topk, num_experts,
-                                      ep_size, topk_ids_dtype)
-
-
-@pytest.mark.parametrize("numel", list(range(1, 8192, 11)))
-@pytest.mark.parametrize("num_experts", [32])
-@pytest.mark.parametrize("ep_size", [2])
-@pytest.mark.parametrize("topk_ids_dtype", [torch.int64])
-def test_compute_expert_num_tokens_from_numel(numel: int, num_experts: int,
-                                              ep_size: int,
-                                              topk_ids_dtype: torch.dtype):
-    do_test_compute_expert_num_tokens(num_tokens=numel,
-                                      num_topk=1,
-                                      num_experts=num_experts,
-                                      ep_size=ep_size,
-                                      topk_ids_dtype=topk_ids_dtype)
@@ -20,7 +20,6 @@ from vllm.model_executor.layers.fused_moe.modular_kernel import (
    FusedMoEModularKernel)
 from vllm.platforms import current_platform
 from vllm.utils import has_deep_ep, has_deep_gemm
-from vllm.utils.deep_gemm import is_blackwell_deep_gemm_used

 from .parallel_utils import ProcessGroupInfo, parallel_launch
 from .utils import make_test_weights
@@ -369,8 +368,6 @@ NUM_EXPERTS = [32]
@pytest.mark.parametrize("world_dp_size", [(2, 1)])
@requires_deep_ep
@requires_deep_gemm
-@pytest.mark.skipif(is_blackwell_deep_gemm_used(),
-                    reason="Skipping test for Blackwell DeepGEMM")
 def test_ht_deepep_deepgemm_moe(mnk: tuple[int, int, int], num_experts: int,
                                topk: int, world_dp_size: tuple[int, int]):
    """
@@ -426,8 +423,6 @@ USE_FP8_DISPATCH = [False]
@pytest.mark.parametrize("world_dp_size", [(2, 1)])
@requires_deep_ep
@requires_deep_gemm
-@pytest.mark.skipif(is_blackwell_deep_gemm_used(),
-                    reason="Skipping test for Blackwell DeepGEMM")
 def test_ll_deepep_deepgemm_moe(
    mnk: tuple[int, int, int],
    num_experts: int,
@@ -13,18 +13,48 @@ import torch

 # vLLM fused-expert reference (Triton fallback + DeepGEMM option)
 from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts
-from vllm.utils import has_deep_gemm
-from vllm.utils.deep_gemm import (calc_diff, per_block_cast_to_fp8,
-                                  per_token_group_cast_to_fp8)
+from vllm.model_executor.layers.quantization.utils.fp8_utils import (
+    per_token_group_quant_fp8)
+from vllm.utils import cdiv

-BLOCK_SIZE = [128, 128]
+has_deep_gemm = importlib.util.find_spec("deep_gemm") is not None
+
+if has_deep_gemm:
+    import deep_gemm
+    BLOCK_M = deep_gemm.get_m_alignment_for_contiguous_layout()
+    BLOCK_SIZE = [BLOCK_M, BLOCK_M]

 requires_deep_gemm = pytest.mark.skipif(
-    not has_deep_gemm(),
+    not has_deep_gemm,
    reason="Requires deep_gemm kernels",
 )


+def calc_diff(x: torch.Tensor, y: torch.Tensor):
+    x, y = x.double(), y.double()
+    denominator = (x * x + y * y).sum()
+    sim = 2 * (x * y).sum() / denominator
+    return 1 - sim
+
+
+def per_block_cast_to_fp8(
+        x: torch.Tensor,
+        block_size_n: int = 128) -> tuple[torch.Tensor, torch.Tensor]:
+    assert x.dim() == 2
+    m, n = x.shape
+    x_padded = torch.zeros(
+        (cdiv(m, 128) * 128, cdiv(n, block_size_n) * block_size_n),
+        dtype=x.dtype,
+        device=x.device)
+    x_padded[:m, :n] = x
+    x_view = x_padded.view(-1, 128, x_padded.size(1) // 128, block_size_n)
+    x_amax = x_view.abs().float().amax(dim=(1, 3), keepdim=True).clamp(1e-4)
+    x_scaled = (x_view * (448.0 / x_amax)).to(torch.float8_e4m3fn)
+    x_scaled_sub = x_scaled.view_as(x_padded)[:m, :n].contiguous()
+    scales = (x_amax / 448.0).view(x_view.size(0), x_view.size(2))
+    return x_scaled_sub, scales
+
+
 def make_block_quant_fp8_weights(
    e: int,
    n: int,
@@ -81,7 +111,7 @@ def run_single_case(m, n, k, topk, num_experts, block_size):
    """
    tokens_bf16 = torch.randn(
        m, k, device="cuda", dtype=torch.bfloat16).clamp_min_(-1).clamp_max_(1)
-    _, a1_scale = per_token_group_cast_to_fp8(tokens_bf16, block_size[1])
+    _, a1_scale = per_token_group_quant_fp8(tokens_bf16, block_size[1])

    # expert weight tensors
    w1, w2, w1_s, w2_s = make_block_quant_fp8_weights(num_experts, n, k,
@@ -125,8 +155,17 @@ def run_single_case(m, n, k, topk, num_experts, block_size):
        block_shape=block_size,
        allow_deep_gemm=True,
    )
-    diff = calc_diff(out_deepgemm, out_triton)
-    assert diff < 0.001, f"Diff exceeded 1%: {diff}"
+
+    base = out_triton.abs().mean()
+    atol = 0.1 * base.clamp(min=1e-2)  # 10% of mean, but not lower than 1e-3
+    rtol = 0.05
+    # ----- Compare -----
+    torch.testing.assert_close(
+        out_deepgemm.to(torch.float32),
+        out_triton.to(torch.float32),
+        rtol=rtol,
+        atol=float(atol),
+    )


 # Note: W1 has shape (E, 2N, K), so N = 512
@@ -1,214 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import copy
-from itertools import product
-from typing import Optional
-
-import pytest
-import torch
-
-import vllm.model_executor.layers.fused_moe.modular_kernel as mk
-from vllm.config import VllmConfig, current_platform, set_current_vllm_config
-from vllm.model_executor.layers.fused_moe.batched_triton_or_deep_gemm_moe import (  # noqa: E501
-    BatchedTritonOrDeepGemmExperts)
-from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
-from vllm.model_executor.layers.fused_moe.cutlass_moe import CutlassExpertsFp8
-from vllm.model_executor.layers.fused_moe.fused_batched_moe import (
-    BatchedTritonExperts)
-from vllm.model_executor.layers.fused_moe.layer import TritonExperts
-from vllm.model_executor.layers.fused_moe.triton_deep_gemm_moe import (
-    TritonOrDeepGemmExperts)
-from vllm.utils import has_deep_ep, has_deep_gemm, has_pplx
-
-from .modular_kernel_tools.common import (Config, RankTensors, WeightTensors,
-                                          reference_moe_impl,
-                                          run_modular_kernel)
-from .modular_kernel_tools.mk_objects import (
-    MK_FUSED_EXPERT_TYPES, MK_MULTI_GPU_PREPARE_FINALIZE_TYPES,
-    MK_QUANT_CONFIGS, MK_SINGLE_GPU_PREPARE_FINALIZE_TYPES)
-from .modular_kernel_tools.parallel_utils import (ProcessGroupInfo,
-                                                  parallel_launch_with_config)
-
-# TODO (varun): These requirements are very strict and could be relaxed.
-has_all_packages = (has_deep_ep() and has_deep_gemm() and has_pplx())
-
-meets_package_requirements = pytest.mark.skipif(
-    not has_all_packages,
-    reason="Requires deep_ep & deep_gemm & pplx packages",
-)
-
-
-def rank_worker(
-    pgi: ProcessGroupInfo,
-    vllm_config: VllmConfig,
-    cpu_group,
-    config: Config,
-    weights: WeightTensors,
-):
-    current_platform.seed_everything(pgi.rank)
-
-    # sanity check
-    from vllm import envs
-    if config.fused_moe_chunk_size is not None:
-        assert (config.fused_moe_chunk_size == envs.VLLM_FUSED_MOE_CHUNK_SIZE)
-
-    # get weights to this device
-    weights.to_current_device()
-
-    Ms = config.Ms
-    assert isinstance(Ms, list)
-    TOPKs = config.topks
-    assert isinstance(TOPKs, list)
-
-    for m, topk in product(Ms, TOPKs):
-        print(f"Running m={m}, topk={topk} ...")
-        # override m and topk
-        cfgx = copy.deepcopy(config)
-        cfgx.Ms = m
-        cfgx.topks = topk
-
-        # inputs for rank
-        rank_tensors = RankTensors.make(cfgx, pgi)
-
-        # modular kernel out
-        mk_out = run_modular_kernel(pgi, vllm_config, cfgx, weights,
-                                    rank_tensors)
-
-        with set_current_vllm_config(vllm_config):
-            ref_out = reference_moe_impl(cfgx, weights, rank_tensors)
-
-        torch.testing.assert_close(ref_out, mk_out, atol=3e-2, rtol=3e-2)
-
-
-def run(config: Config):
-    assert config.is_valid()
-    print(f"Testing config \n{config.describe()} ...")
-
-    weights: WeightTensors = WeightTensors.make(config)
-
-    vllm_config, env_dict = config.make_env_data()
-    parallel_launch_with_config(config.world_size, rank_worker, vllm_config,
-                                env_dict, config, weights)
-
-
-Ms = [32, 64]
-Ks = [7168]  # hidden sizes
-Ns = [2048]
-TOPKs = [4, 1]
-Es = [32]
-DTYPEs = [torch.bfloat16]
-FUSED_MOE_CHUNK_SIZEs = [None, 16]
-
-
-def is_nyi_config(config: Config) -> bool:
-    # We know these configs to be legitimate. but still fail.
-
-    if (config.fused_experts_type in [
-            BatchedTritonExperts, BatchedTritonOrDeepGemmExperts,
-            TritonExperts, TritonOrDeepGemmExperts
-    ]):
-        # The triton kernels expect both per-act-token-quant and
-        # per-out-ch-quant or neither.
-        unsupported_quant_config = ((config.is_per_act_token_quant +
-                                     config.is_per_out_ch_quant) == 1)
-        return unsupported_quant_config
-
-    # cutlass kernels dont support expert_maps yet.
-    return config.fused_experts_type == CutlassExpertsFp8
-
-
-@pytest.mark.parametrize("k", Ks)
-@pytest.mark.parametrize("n", Ns)
-@pytest.mark.parametrize("e", Es)
-@pytest.mark.parametrize("dtype", DTYPEs)
-@pytest.mark.parametrize("quant_config", MK_QUANT_CONFIGS)
-@pytest.mark.parametrize(
-    "combination",
-    product(MK_MULTI_GPU_PREPARE_FINALIZE_TYPES, MK_FUSED_EXPERT_TYPES))
-@pytest.mark.parametrize("fused_moe_chunk_size", FUSED_MOE_CHUNK_SIZEs)
-@pytest.mark.parametrize("world_size", [2])
-@meets_package_requirements
-def test_modular_kernel_combinations_multigpu(
-        k: int, n: int, e: int, dtype: torch.dtype,
-        quant_config: FusedMoEQuantConfig,
-        combination: tuple[mk.FusedMoEPrepareAndFinalize,
-                           mk.FusedMoEPermuteExpertsUnpermute],
-        fused_moe_chunk_size: Optional[int], world_size: int):
-
-    config = Config(
-        Ms=Ms,
-        K=k,
-        N=n,
-        E=e,
-        topks=TOPKs,
-        dtype=dtype,
-        quant_config=quant_config,
-        prepare_finalize_type=combination[0],
-        fused_experts_type=combination[1],
-        fused_moe_chunk_size=fused_moe_chunk_size,
-        world_size=world_size,
-    )
-    if not config.is_valid():
-        pytest.skip(f"Tests config {config} is not valid. Skipping ...")
-
-    if is_nyi_config(config):
-        pytest.skip(f"Tests config {config} is nyi. Skipping ...")
-
-    print(f"{config.describe()}")
-    run(config)
-
-
-@pytest.mark.parametrize("k", Ks)
-@pytest.mark.parametrize("n", Ns)
-@pytest.mark.parametrize("e", Es)
-@pytest.mark.parametrize("dtype", DTYPEs)
-@pytest.mark.parametrize("quant_config", MK_QUANT_CONFIGS)
-@pytest.mark.parametrize(
-    "combination",
-    product(MK_SINGLE_GPU_PREPARE_FINALIZE_TYPES, MK_FUSED_EXPERT_TYPES))
-@pytest.mark.parametrize("fused_moe_chunk_size", FUSED_MOE_CHUNK_SIZEs)
-@pytest.mark.parametrize("world_size", [1])
-@meets_package_requirements
-def test_modular_kernel_combinations_singlegpu(
-        k: int, n: int, e: int, dtype: torch.dtype,
-        quant_config: FusedMoEQuantConfig,
-        combination: tuple[mk.FusedMoEPrepareAndFinalize,
-                           mk.FusedMoEPermuteExpertsUnpermute],
-        fused_moe_chunk_size: Optional[int], world_size: int):
-    config = Config(
-        Ms=Ms,
-        K=k,
-        N=n,
-        E=e,
-        topks=TOPKs,
-        dtype=dtype,
-        quant_config=quant_config,
-        prepare_finalize_type=combination[0],
-        fused_experts_type=combination[1],
-        fused_moe_chunk_size=fused_moe_chunk_size,
-        world_size=world_size,
-    )
-
-    if not config.is_valid():
-        pytest.skip(f"Tests config {config} is not valid. Skipping ...")
-
-    if is_nyi_config(config):
-        pytest.skip(f"Tests config {config} is nyi. Skipping ...")
-
-    run(config)
-
-
-if __name__ == '__main__':
-    # Ability to test individual PrepareAndFinalize and FusedExperts combination
-    from .modular_kernel_tools.cli_args import (make_config,
-                                                make_config_arg_parser)
-    parser = make_config_arg_parser(description=(
-        "Run single prepare-finalize & fused-experts combination test"
-        "Example : python3 -m tests.kernels.moe.test_modular_kernel_combinations "  #noqa: E501
-        "--pf-type PplxPrepareAndFinalize --experts-type BatchedTritonExperts"
-    ))
-    args = parser.parse_args()
-    config = make_config(args)
-
-    run(config)
@@ -32,8 +32,6 @@ from vllm.model_executor.layers.fused_moe.fused_batched_moe import (
 from vllm.model_executor.layers.fused_moe.fused_moe import get_default_config
 from vllm.model_executor.layers.fused_moe.modular_kernel import (
    FusedMoEModularKernel)
-from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
-    TopKWeightAndReduceDelegate)
 from vllm.platforms import current_platform
 from vllm.utils import round_up

@@ -373,7 +371,6 @@ def pplx_prepare_finalize(
        chunk_topk_weight,
        chunk_topk_ids,
        False,
-        weight_and_reduce_impl=TopKWeightAndReduceDelegate(),
    )

    torch.cuda.synchronize()
@@ -8,15 +8,19 @@ import pytest
 import torch

 from tests.kernels.quant_utils import (native_per_token_group_quant_fp8,
-                                       native_w8a8_block_matmul)
+                                       native_w8a8_block_matmul,
+                                       per_block_cast_to_fp8)
 from vllm.config import VllmConfig
 from vllm.model_executor.layers.quantization.utils.fp8_utils import (
-    get_col_major_tma_aligned_tensor, per_token_group_quant_fp8,
-    w8a8_block_fp8_matmul)
+    per_token_group_quant_fp8, w8a8_block_fp8_matmul)
 from vllm.platforms import current_platform
-from vllm.utils import has_deep_gemm
-from vllm.utils.deep_gemm import (fp8_gemm_nt, per_block_cast_to_fp8,
-                                  per_token_group_cast_to_fp8)
+
+dg_available = False
+try:
+    import deep_gemm
+    dg_available = True
+except ImportError:
+    pass

 if current_platform.get_device_capability() < (9, 0):
    pytest.skip("FP8 Triton requires CUDA 9.0 or higher",
@@ -102,8 +106,7 @@ def test_w8a8_block_fp8_matmul(M, N, K, block_size, out_dtype, seed):
@pytest.mark.parametrize(
    "M,N,K,block_size,out_dtype,seed",
    itertools.product(M, N, K, BLOCK_SIZE, OUT_DTYPES, SEEDS))
-@pytest.mark.skipif(not has_deep_gemm(),
-                    reason="DeepGemm kernels not available.")
+@pytest.mark.skipif(not dg_available, reason="DeepGemm kernels not available.")
@torch.inference_mode()
 def test_w8a8_block_fp8_deep_gemm_matmul(M, N, K, block_size, out_dtype, seed):
    # only aligned sizes
@@ -117,7 +120,9 @@ def test_w8a8_block_fp8_deep_gemm_matmul(M, N, K, block_size, out_dtype, seed):
    A_fp32 = (torch.rand(M, K, dtype=torch.float32) - 0.5) * 2 * fp8_max
    B_fp32 = (torch.rand(N, K, dtype=torch.float32) - 0.5) * 2 * fp8_max

-    A_fp8, As_fp8 = per_token_group_cast_to_fp8(A_fp32, block_size[1])
+    _, block_k = block_size[0], block_size[1]
+
+    A_fp8, As_fp8 = per_token_group_quant_fp8(A_fp32, block_k)
    B_fp8, Bs_fp8 = per_block_cast_to_fp8(B_fp32)

    As = As_fp8.to(torch.float32)
@@ -127,14 +132,14 @@ def test_w8a8_block_fp8_deep_gemm_matmul(M, N, K, block_size, out_dtype, seed):
                                       out_dtype)

    # Transpose earlier so that the testing will not trigger transposing kernels
-    As_fp8 = get_col_major_tma_aligned_tensor(As_fp8)
+    As_fp8 = deep_gemm.get_col_major_tma_aligned_tensor(As_fp8)

    out = torch.zeros((M, N), device='cuda', dtype=out_dtype)

    assert As_fp8.shape == (M, (K + 127) //
                            128), f"{As_fp8.shape} != {(M, (K + 127) // 128)}"

-    fp8_gemm_nt((A_fp8, As_fp8), (B_fp8, Bs_fp8), out)
+    deep_gemm.gemm_fp8_fp8_bf16_nt((A_fp8, As_fp8), (B_fp8, Bs_fp8), out)

    rel_diff = (torch.mean(
        torch.abs(out.to(torch.float32) - ref_out.to(torch.float32))) /
@@ -1072,7 +1072,6 @@ def torch_experts(
    quant_dtype: Optional[torch.dtype] = None,
    per_act_token_quant=False,
    block_shape: Optional[list[int]] = None,
-    apply_router_weights_on_input: bool = False,
 ) -> torch.Tensor:
    assert (global_num_experts == -1
            or (global_num_experts == w1.shape[0] and expert_map is None)
@@ -1082,17 +1081,11 @@ def torch_experts(
    M, K = a.shape
    topk = topk_ids.shape[1]

-    if apply_router_weights_on_input:
-        assert topk == 1
-        a = a * topk_weight.to(a.dtype)
-
    a = a.view(M, -1, K).repeat(1, topk, 1).reshape(-1, K)

    out = torch.zeros(M * topk, w2.shape[1], dtype=a.dtype, device=a.device)

-    if a1_scale:
-        assert not per_act_token_quant and block_shape is None
-    a, a_scale = moe_kernel_quantize_input(a, a1_scale, quant_dtype,
+    a, a_scale = moe_kernel_quantize_input(a, None, quant_dtype,
                                           per_act_token_quant, block_shape)

    num_experts = w1.shape[0]
@@ -1111,7 +1104,6 @@ def torch_experts(
                tmp2 = SiluAndMul()(tmp1)
                out[mask] = tmp2 @ w2[i].transpose(0, 1)
            elif block_shape is not None:
-                # block quantized
                assert (a_scale is not None and w1_scale is not None
                        and w2_scale is not None)
                tmp1 = native_w8a8_block_matmul(a[mask], w1[i], a_scale[mask],
@@ -1129,27 +1121,15 @@ def torch_experts(
                assert (a_scale is not None and w1_scale is not None
                        and w2_scale is not None)
                scales = a_scale if a_scale.numel() == 1 else a_scale[mask]
-
                tmp1 = a[mask].to(f32) * scales
                w1_dq = (w1[i].to(f32) * w1_scale[i]).transpose(0, 1)
-                tmp1 = (tmp1 @ w1_dq).to(out.dtype)
-
-                tmp2 = SiluAndMul()(tmp1).to(out.dtype)
-
-                tmp2, b_scale = moe_kernel_quantize_input(
-                    tmp2, a2_scale, quant_dtype, per_act_token_quant,
-                    block_shape)
-                assert b_scale is not None
-
-                tmp2 = tmp2.to(f32) * b_scale
+                tmp1 = tmp1 @ w1_dq
+                tmp2 = SiluAndMul()(tmp1)
                w2_dq = (w2[i].to(f32) * w2_scale[i]).transpose(0, 1)
                out[mask] = (tmp2 @ w2_dq).to(out.dtype)

-    if apply_router_weights_on_input:
-        return out
-    else:
-        return (out.view(M, -1, w2.shape[1]).to(f32) *
-                topk_weight.view(M, -1, 1)).sum(dim=1).to(out.dtype)
+    return (out.view(M, -1, w2.shape[1]).to(f32) *
+            topk_weight.view(M, -1, 1)).sum(dim=1).to(out.dtype)


 def torch_moe(a: torch.Tensor,
@@ -1,118 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Tests for applying default registered multimodal loras.
-"""
-
-import os
-
-from huggingface_hub import snapshot_download
-
-from vllm.lora.request import LoRARequest
-
-from ..conftest import AudioTestAssets, VllmRunner
-
-MODEL_PATH = snapshot_download("microsoft/Phi-4-multimodal-instruct")
-AUDIO_LORA_PATH = os.path.join(MODEL_PATH, "speech-lora")
-IMAGE_LORA_PATH = os.path.join(MODEL_PATH, "vision-lora")
-
-AUDIO_PROMPT = "<|user|><|audio_1|>Can you transcribe this audio?<|end|><|assistant|>"  # noqa: E501
-
-# Responses are greedy decoded; we just check the end of
-# the generated text. If the lora is inactive, this model
-# generates commentary on the transcription.
-RESPONSE_SUFFIX_WITH_LORA = "Spoken text: The first words I spoke in the original chronograph, a little piece of practical poetry. Mary had a little lamb, it slept with quite a snow, and everywhere that Mary went, the lamb was sure to go."  # noqa: E501
-RESPONSE_SUFFIX_WITHOUT_LORA = "Certainly! Here is the transcription of the audio you provided:\n\nThe first words I spoke in the original phonograph record: A little piece of practical poetry. Mary had a little lamb; its fleece was white as snow, and everywhere that Mary went, the lamb was sure to go."  # noqa: E501
-
-VLLM_RUNNER_BASE_KWARGS = {
-    "model_name": MODEL_PATH,
-    "dtype": "half",
-    "enable_lora": "True",
-    "max_num_seqs": 2,
-    "max_lora_rank": 320,
-    "max_model_len": 12800,
-    "gpu_memory_utilization": 0.8,
-    "limit_mm_per_prompt": {
-        "audio": 1
-    },
-    "enforce_eager": True,
-}
-
-
-def run_test(vllm_runner, audio_assets, lora_request, expected_suffix,
-             **kwargs):
-    inputs = [([AUDIO_PROMPT], [audio_assets[0].audio_and_sample_rate[0]])]
-
-    # Apply any additional kwargs as overrides to the base kwargs
-    vllm_runner_kwargs = {**VLLM_RUNNER_BASE_KWARGS, **kwargs}
-
-    with vllm_runner(**vllm_runner_kwargs) as vllm_model:
-        vllm_outputs_with_default_lora = [
-            vllm_model.generate_greedy(
-                prompts,
-                max_tokens=128,
-                audios=audios,
-                lora_request=lora_request,
-            ) for prompts, audios in inputs
-        ]
-
-        assert vllm_outputs_with_default_lora[-1][-1][-1].endswith(
-            expected_suffix)
-
-
-def test_active_default_mm_lora(
-    vllm_runner: type[VllmRunner],
-    audio_assets: AudioTestAssets,
-):
-    """Ensure that we can use the default audio lora."""
-    run_test(
-        vllm_runner,
-        audio_assets,
-        lora_request=None,
-        default_mm_loras={"audio": AUDIO_LORA_PATH},
-        expected_suffix=RESPONSE_SUFFIX_WITH_LORA,
-    )
-
-
-def test_inactive_default_mm_lora(
-    vllm_runner: type[VllmRunner],
-    audio_assets: AudioTestAssets,
-):
-    """Ensure that modalities are filtered properly."""
-    # Default image lora won't be active since we only pass audio
-    run_test(
-        vllm_runner,
-        audio_assets,
-        lora_request=None,
-        default_mm_loras={"image": IMAGE_LORA_PATH},
-        expected_suffix=RESPONSE_SUFFIX_WITHOUT_LORA,
-    )
-
-
-def test_default_mm_lora_succeeds_with_redundant_lora_request(
-    vllm_runner: type[VllmRunner],
-    audio_assets: AudioTestAssets,
-):
-    """Ensure that redundantly providing the lora works."""
-    run_test(
-        vllm_runner,
-        audio_assets,
-        lora_request=LoRARequest("audio", 1, AUDIO_LORA_PATH),
-        default_mm_loras={"audio": AUDIO_LORA_PATH},
-        expected_suffix=RESPONSE_SUFFIX_WITH_LORA,
-    )
-
-
-def test_default_mm_lora_fails_with_overridden_lora_request(
-    vllm_runner: type[VllmRunner],
-    audio_assets: AudioTestAssets,
-):
-    """Ensure that if the lora_request conflicts with default_mm_loras,
-    we use the lora_request."""
-    run_test(
-        vllm_runner,
-        audio_assets,
-        lora_request=LoRARequest("speech", 2, AUDIO_LORA_PATH),
-        default_mm_loras={"audio": IMAGE_LORA_PATH},
-        expected_suffix=RESPONSE_SUFFIX_WITH_LORA,
-    )
@@ -46,15 +46,20 @@ def test_guided_logits_processors(zephyr_7B_tokenzer, sample_regex,
                                  whitespace_pattern=None,
                                  reasoner=None)

+    token_ids = zephyr_7B_tokenzer.encode(
+        f"Give an example IPv4 address with this regex: {sample_regex}")
    tensor = torch.rand(32000)
    original_tensor = torch.clone(tensor)
-    tensor = regex_LP([], tensor)
+    regex_LP(token_ids, tensor)
    assert tensor.shape == original_tensor.shape
    assert not torch.allclose(tensor, original_tensor)

+    token_ids = zephyr_7B_tokenzer.encode(
+        f"Give an employee profile that fits this schema: {sample_json_schema}"
+    )
    tensor = torch.rand(32000)
    original_tensor = torch.clone(tensor)
-    tensor = json_LP([], tensor)
+    json_LP(token_ids, tensor)
    assert tensor.shape == original_tensor.shape
    assert not torch.allclose(tensor, original_tensor)

@@ -76,6 +81,8 @@ async def test_guided_logits_processor_black_box(backend: str, is_local: bool,
        seed=0,
        dtype="bfloat16",
    )
+    token_ids = zephyr_7B_tokenzer.encode(
+        f"Give an example IPv4 address with this regex: {sample_regex}")
    regex_request = GuidedDecodingParams(regex=sample_regex, backend=backend)

    regex_lp = get_local_guided_decoding_logits_processor(
@@ -85,11 +92,13 @@ async def test_guided_logits_processor_black_box(backend: str, is_local: bool,
    assert regex_lp is not None
    tensor = torch.rand(32000)
    original_tensor = torch.clone(tensor)
-    # allowed tokens at state 0
-    tensor = regex_lp([], tensor)
+    tensor = regex_lp(token_ids, tensor)
    assert tensor.shape == original_tensor.shape
    assert not torch.allclose(tensor, original_tensor)

+    token_ids = zephyr_7B_tokenzer.encode(
+        f"Give an employee profile that fits this schema: {sample_json_schema}"
+    )
    json_request = GuidedDecodingParams(json=sample_json_schema,
                                        backend=backend)
    json_lp = await get_guided_decoding_logits_processor(
@@ -97,7 +106,7 @@ async def test_guided_logits_processor_black_box(backend: str, is_local: bool,
    assert json_lp is not None
    tensor = torch.rand(32000)
    original_tensor = torch.clone(tensor)
-    tensor = json_lp([], tensor)
+    tensor = json_lp(token_ids, tensor)
    assert tensor.shape == original_tensor.shape
    assert not torch.allclose(tensor, original_tensor)

@@ -121,6 +130,7 @@ async def test_guided_logits_processor_with_reasoning(
        dtype="bfloat16",
    )
    token_ids = deepseek_r1_qwen_tokenizer.encode(
+        f"Give an example IPv4 address with this regex: {sample_regex}."
        "<think>here is the thinking process")
    regex_request = GuidedDecodingParams(regex=sample_regex, backend=backend)

@@ -131,13 +141,14 @@ async def test_guided_logits_processor_with_reasoning(
                    regex_request, deepseek_r1_qwen_tokenizer, config,
                    reasoning_backend)
    assert regex_lp is not None
-    tensor = torch.rand(151664)
+    tensor = torch.rand(32000)
    original_tensor = torch.clone(tensor)
    tensor = regex_lp(token_ids, tensor)
    assert tensor.shape == original_tensor.shape
    assert torch.allclose(tensor, original_tensor)

    token_ids = deepseek_r1_qwen_tokenizer.encode(
+        f"Give an employee profile that fits this schema: {sample_json_schema}."
        "<think>here is the thinking process")
    json_request = GuidedDecodingParams(json=sample_json_schema,
                                        backend=backend)
@@ -147,7 +158,7 @@ async def test_guided_logits_processor_with_reasoning(
        await get_guided_decoding_logits_processor(
            json_request, deepseek_r1_qwen_tokenizer, config, reasoning_backend)
    assert json_lp is not None
-    tensor = torch.rand(151664)
+    tensor = torch.rand(32000)
    original_tensor = torch.clone(tensor)
    tensor = json_lp(token_ids, tensor)
    assert tensor.shape == original_tensor.shape
@@ -155,7 +166,8 @@ async def test_guided_logits_processor_with_reasoning(

    # Thinking is over, so the tensor should change.
    token_ids = deepseek_r1_qwen_tokenizer.encode(
-        "<think>here is the thinking process</think>")
+        f"Give an employee profile that fits this schema: {sample_json_schema}."
+        "<think>here is the thinking process</think> Then")
    json_request = GuidedDecodingParams(json=sample_json_schema,
                                        backend=backend)
    json_lp = get_local_guided_decoding_logits_processor(
@@ -164,7 +176,7 @@ async def test_guided_logits_processor_with_reasoning(
        await get_guided_decoding_logits_processor(
            json_request, deepseek_r1_qwen_tokenizer, config, reasoning_backend)
    assert json_lp is not None
-    tensor = torch.rand(151664)
+    tensor = torch.rand(32000)
    original_tensor = torch.clone(tensor)
    tensor = json_lp(token_ids, tensor)
    assert tensor.shape == original_tensor.shape
@@ -14,7 +14,7 @@ from transformers import BitsAndBytesConfig
 from tests.quantization.utils import is_quant_method_supported

 from ...utils import compare_two_settings, multi_gpu_test
-from ..utils import check_embeddings_close, check_logprobs_close
+from ..utils import check_embeddings_close

 models_4bit_to_test = [
    ("facebook/opt-125m", "quantize opt model inflight"),
@@ -26,10 +26,6 @@ models_4bit_to_embedding_test = [
    ("intfloat/e5-mistral-7b-instruct", "quantize embedding model inflight"),
 ]

-models_4bit_to_moe_test = [
-    ("allenai/OLMoE-1B-7B-0125-Instruct", "quantize moe model inflight"),
-]
-
 models_pre_qaunt_4bit_to_test = [
    ('PrunaAI/Einstein-v6.1-Llama3-8B-bnb-4bit-smashed',
     'read pre-quantized 4-bit FP4 model'),
@@ -119,35 +115,6 @@ def test_load_pp_4bit_bnb_model(model_name, description) -> None:
    compare_two_settings(model_name, common_args, pp_args)


-@pytest.mark.skipif(not is_quant_method_supported("bitsandbytes"),
-                    reason='bitsandbytes is not supported on this GPU type.')
-@pytest.mark.parametrize("model_name, description", models_4bit_to_moe_test)
-def test_4bit_bnb_moe_model(hf_runner, vllm_runner, example_prompts,
-                            model_name, description) -> None:
-
-    hf_model_kwargs = dict(quantization_config=BitsAndBytesConfig(
-        load_in_4bit=True,
-        bnb_4bit_quant_type="nf4",
-        bnb_4bit_use_double_quant=True,
-    ))
-    with vllm_runner(model_name,
-                     quantization='bitsandbytes',
-                     enforce_eager=False) as llm:
-        vllm_outputs = llm.generate_greedy_logprobs(example_prompts,
-                                                    max_tokens=32,
-                                                    num_logprobs=5)
-
-    with hf_runner(model_name, model_kwargs=hf_model_kwargs) as llm:
-        transformers_outputs = llm.generate_greedy_logprobs_limit(
-            example_prompts, max_tokens=32, num_logprobs=5)
-    check_logprobs_close(
-        outputs_0_lst=transformers_outputs,
-        outputs_1_lst=vllm_outputs,
-        name_0="transformers",
-        name_1="vllm",
-    )
-
-
@pytest.mark.skipif(not is_quant_method_supported("bitsandbytes"),
                    reason='bitsandbytes is not supported on this GPU type.')
@pytest.mark.parametrize("model_name, description",
@@ -215,8 +182,7 @@ def validate_generated_texts(hf_runner,
                             model_name,
                             pre_quant=False,
                             hf_model_kwargs=None,
-                             vllm_tp_size=1,
-                             max_tokens=8):
+                             vllm_tp_size=1):

    # NOTE: run vLLM first, as it requires a clean process
    # when using distributed inference
@@ -224,8 +190,7 @@ def validate_generated_texts(hf_runner,
                     quantization=None if pre_quant else 'bitsandbytes',
                     tensor_parallel_size=vllm_tp_size,
                     enforce_eager=False) as llm:
-
-        vllm_outputs = llm.generate_greedy(prompts, max_tokens)
+        vllm_outputs = llm.generate_greedy(prompts, 8)
        vllm_logs = log_generated_texts(prompts, vllm_outputs, "VllmRunner")

    # Clean up the GPU memory for the next test
@@ -237,17 +202,19 @@ def validate_generated_texts(hf_runner,

    # Run with HF runner
    with hf_runner(model_name, model_kwargs=hf_model_kwargs) as llm:
-        hf_outputs = llm.generate_greedy(prompts, max_tokens)
+        hf_outputs = llm.generate_greedy(prompts, 8)
        hf_logs = log_generated_texts(prompts, hf_outputs, "HfRunner")

    # Clean up the GPU memory for the next test
    gc.collect()
    torch.cuda.empty_cache()
+
    # Compare the generated strings
    for hf_log, vllm_log in zip(hf_logs, vllm_logs):
        hf_str = hf_log["generated_text"]
        vllm_str = vllm_log["generated_text"]
        prompt = hf_log["prompt"]
+
        assert hf_str == vllm_str, (f"Model: {model_name}"
                                    f"Mismatch between HF and vLLM outputs:\n"
                                    f"Prompt: {prompt}\n"
@@ -218,8 +218,6 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
                                         trust_remote_code=True),
    "MiniCPM3ForCausalLM": _HfExamplesInfo("openbmb/MiniCPM3-4B",
                                         trust_remote_code=True),
-    "MiniMaxForCausalLM": _HfExamplesInfo("MiniMaxAI/MiniMax-Text-01-hf",
-                                          min_transformers_version="4.53"),
    "MiniMaxText01ForCausalLM": _HfExamplesInfo("MiniMaxAI/MiniMax-Text-01",
                                                trust_remote_code=True,
                                                revision="a59aa9cbc53b9fb8742ca4e9e1531b9802b6fdc3"),  # noqa: E501
@@ -1,195 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-# ruff: noqa: E501
-
-import json
-
-import pytest
-
-from vllm.entrypoints.openai.protocol import FunctionCall, ToolCall
-from vllm.entrypoints.openai.tool_parsers import KimiK2ToolParser
-from vllm.transformers_utils.tokenizer import get_tokenizer
-
-pytest.skip("skip kimi_k2 parser test", allow_module_level=True)
-
-# Use a common model that is likely to be available
-MODEL = "moonshotai/Kimi-K2-Instruct"
-
-
-@pytest.fixture(scope="module")
-def kimi_k2_tokenizer():
-    return get_tokenizer(tokenizer_name=MODEL, trust_remote_code=True)
-
-
-@pytest.fixture
-def kimi_k2_tool_parser(kimi_k2_tokenizer):
-    return KimiK2ToolParser(kimi_k2_tokenizer)
-
-
-def assert_tool_calls(actual_tool_calls: list[ToolCall],
-                      expected_tool_calls: list[ToolCall]):
-    assert len(actual_tool_calls) == len(expected_tool_calls)
-
-    for actual_tool_call, expected_tool_call in zip(actual_tool_calls,
-                                                    expected_tool_calls):
-
-        assert actual_tool_call.type == "function"
-        assert actual_tool_call.function == expected_tool_call.function
-
-        # assert tool call id format
-        assert actual_tool_call.id.startswith("functions.")
-        assert actual_tool_call.id.split(':')[-1].isdigit()
-        assert actual_tool_call.id.split('.')[1].split(
-            ':')[0] == expected_tool_call.function.name
-
-
-def test_extract_tool_calls_no_tools(kimi_k2_tool_parser):
-    model_output = "This is a test"
-    extracted_tool_calls = kimi_k2_tool_parser.extract_tool_calls(
-        model_output, request=None)  # type: ignore[arg-type]
-    assert not extracted_tool_calls.tools_called
-    assert extracted_tool_calls.tool_calls == []
-    assert extracted_tool_calls.content == model_output
-
-
-@pytest.mark.parametrize(
-    ids=[
-        "tool_call_with_content_before",
-        "multi_tool_call_with_content_before",
-    ],
-    argnames=["model_output", "expected_tool_calls", "expected_content"],
-    argvalues=[
-        (
-            """I'll help you check the weather. <|tool_calls_section_begin|> <|tool_call_begin|>
-functions.get_weather:0 <|tool_call_argument_begin|> {"city": "Beijing"} <|tool_call_end|> <|tool_calls_section_end|>""",
-            [
-                ToolCall(id='functions.get_weather:0',
-                         function=FunctionCall(
-                             name="get_weather",
-                             arguments=json.dumps({
-                                 "city": "Beijing",
-                             }, ),
-                         ),
-                         type='function')
-            ],
-            "I'll help you check the weather. ",
-        ),
-        (
-            """I'll help you check the weather. <|tool_calls_section_begin|> <|tool_call_begin|>
-functions.get_weather:0 <|tool_call_argument_begin|> {"city": "Beijing"} <|tool_call_end|> <|tool_call_begin|>
-functions.get_weather:1 <|tool_call_argument_begin|> {"city": "Shanghai"} <|tool_call_end|> <|tool_calls_section_end|>""",
-            [
-                ToolCall(id='functions.get_weather:0',
-                         function=FunctionCall(
-                             name="get_weather",
-                             arguments=json.dumps({
-                                 "city": "Beijing",
-                             }, ),
-                         ),
-                         type='function'),
-                ToolCall(id='functions.get_weather:1',
-                         function=FunctionCall(
-                             name="get_weather",
-                             arguments=json.dumps({
-                                 "city": "Shanghai",
-                             }, ),
-                         ),
-                         type='function')
-            ],
-            "I'll help you check the weather. ",
-        ),
-    ],
-)
-def test_extract_tool_calls(kimi_k2_tool_parser, model_output,
-                            expected_tool_calls, expected_content):
-    extracted_tool_calls = kimi_k2_tool_parser.extract_tool_calls(
-        model_output, request=None)  # type: ignore[arg-type]
-    assert extracted_tool_calls.tools_called
-
-    assert_tool_calls(extracted_tool_calls.tool_calls, expected_tool_calls)
-
-    assert extracted_tool_calls.content == expected_content
-
-
-def test_extract_tool_calls_invalid_json(kimi_k2_tool_parser):
-    """we'll return every funcall result"""
-    model_output = """I'll help you check the weather. <|tool_calls_section_begin|> <|tool_call_begin|>
-functions.invalid_get_weather:0 <|tool_call_argument_begin|> {"city": "Beijing" <|tool_call_end|> <|tool_call_begin|>
-functions.valid_get_weather:1 <|tool_call_argument_begin|> {"city": "Shanghai"} <|tool_call_end|> <|tool_calls_section_end|>"""
-
-    extracted_tool_calls = kimi_k2_tool_parser.extract_tool_calls(
-        model_output, request=None)  # type: ignore[arg-type]
-
-    assert extracted_tool_calls.tools_called
-    # Should extract only the valid JSON tool calls
-    assert len(extracted_tool_calls.tool_calls) == 2
-    assert extracted_tool_calls.tool_calls[
-        0].function.name == "invalid_get_weather"
-    assert extracted_tool_calls.tool_calls[
-        1].function.name == "valid_get_weather"
-
-
-def test_extract_tool_calls_invalid_funcall(kimi_k2_tool_parser):
-    """we'll return every funcall result"""
-    model_output = """I'll help you check the weather. <|tool_calls_section_begin|> <|tool_call_begin|>
-functions.invalid_get_weather.0 <|tool_call_argument_begin|> {"city": "Beijing"} <|tool_call_end|> <|tool_call_begin|>
-functions.valid_get_weather:1 <|tool_call_argument_begin|> {"city": "Shanghai"} <|tool_call_end|> <|tool_calls_section_end|>"""
-
-    extracted_tool_calls = kimi_k2_tool_parser.extract_tool_calls(
-        model_output, request=None)  # type: ignore[arg-type]
-
-    assert extracted_tool_calls.tools_called
-    # Should extract only the valid JSON tool calls
-    assert len(extracted_tool_calls.tool_calls) == 1
-    assert extracted_tool_calls.tool_calls[
-        0].function.name == "valid_get_weather"
-
-
-def test_streaming_basic_functionality(kimi_k2_tool_parser):
-    """Test basic streaming functionality."""
-    # Reset streaming state
-    kimi_k2_tool_parser.current_tool_name_sent = False
-    kimi_k2_tool_parser.prev_tool_call_arr = []
-    kimi_k2_tool_parser.current_tool_id = -1
-    kimi_k2_tool_parser.streamed_args_for_tool = []
-
-    # Test with a simple tool call
-    current_text = """ check the weather. <|tool_calls_section_begin|> <|tool_call_begin|>
-functions.get_weather:0 <|tool_call_argument_begin|> {"city": "Beijing"} <|tool_call_end|> <|tool_calls_section_end|>"""
-
-    # First call should handle the initial setup
-    result = kimi_k2_tool_parser.extract_tool_calls_streaming(
-        previous_text="I'll help you",
-        current_text=current_text,
-        delta_text="<|tool_calls_section_end|>",
-        previous_token_ids=[],
-        current_token_ids=[],
-        delta_token_ids=[],
-        request=None,
-    )
-
-    # The result might be None or contain tool call information
-    # This depends on the internal state management
-    if result is not None and hasattr(result,
-                                      'tool_calls') and result.tool_calls:
-        assert len(result.tool_calls) >= 0
-
-
-def test_streaming_no_tool_calls(kimi_k2_tool_parser):
-    """Test streaming when there are no tool calls."""
-    current_text = "This is just regular text without any tool calls."
-
-    result = kimi_k2_tool_parser.extract_tool_calls_streaming(
-        previous_text="This is just regular text",
-        current_text=current_text,
-        delta_text=" without any tool calls.",
-        previous_token_ids=[],
-        current_token_ids=[],
-        delta_token_ids=[],
-        request=None,
-    )
-
-    # Should return the delta text as content
-    assert result is not None
-    assert hasattr(result, 'content')
-    assert result.content == " without any tool calls."
@@ -72,7 +72,7 @@ def _compile_and_check(tools: list[ChatCompletionToolsParam], sample_output,
    assert isinstance(schema, dict)

    # use build_regex_from_schema used in JSONLogitsProcessor to create Guide
-    from outlines_core.json_schema import build_regex_from_schema
+    from outlines_core.fsm.json_schema import build_regex_from_schema
    regex = build_regex_from_schema(json.dumps(schema))
    compiled = re.compile(regex)
    matches = compiled.fullmatch(json.dumps(sample_output)) is not None
@@ -41,10 +41,6 @@ PARAMS_MODELS_BACKENDS_TOKENIZER_MODE = [
    ("mistralai/Ministral-8B-Instruct-2410", "guidance", "auto", None),
    ("mistralai/Ministral-8B-Instruct-2410", "xgrammar", "mistral", None),
    ("Qwen/Qwen2.5-1.5B-Instruct", "xgrammar", "auto", None),
-    ("mistralai/Ministral-8B-Instruct-2410", "outlines", "auto", None),
-    ("mistralai/Ministral-8B-Instruct-2410", "outlines", "mistral", None),
-    ("mistralai/Ministral-8B-Instruct-2410", "outlines", "auto",
-     NGRAM_SPEC_CONFIG),
    #FIXME: This test is flaky on CI thus disabled
    #("Qwen/Qwen2.5-1.5B-Instruct", "guidance", "auto"),
    ("mistralai/Ministral-8B-Instruct-2410", "guidance", "auto",
@@ -110,15 +106,13 @@ def test_structured_output(
    enforce_eager = bool(not current_platform.is_tpu())
    # Use a single LLM instance for several scenarios to
    # speed up the test suite.
-    llm = LLM(
-        model=model_name,
-        enforce_eager=enforce_eager,
-        max_model_len=1024,
-        guided_decoding_backend=guided_decoding_backend,
-        guided_decoding_disable_any_whitespace=(guided_decoding_backend
-                                                in {"xgrammar", "guidance"}),
-        tokenizer_mode=tokenizer_mode,
-        speculative_config=speculative_config)
+    llm = LLM(model=model_name,
+              enforce_eager=enforce_eager,
+              max_model_len=1024,
+              guided_decoding_backend=guided_decoding_backend,
+              guided_decoding_disable_any_whitespace=True,
+              tokenizer_mode=tokenizer_mode,
+              speculative_config=speculative_config)

    #
    # Test 1: Generate JSON output based on a provided schema
@@ -152,33 +146,32 @@ def test_structured_output(
    #
    # Test 2: Generate JSON object without a schema
    #
-    if guided_decoding_backend != "outlines":
-        sampling_params = SamplingParams(
-            temperature=1.0,
-            max_tokens=4096,
-            n=2,
-            guided_decoding=GuidedDecodingParams(json_object=True))
+    sampling_params = SamplingParams(
+        temperature=1.0,
+        max_tokens=4096,
+        n=2,
+        guided_decoding=GuidedDecodingParams(json_object=True))

-        outputs = llm.generate(prompts=(
-            "Generate a JSON object with curly braces for a person with "
-            "name and age fields for John Smith who is 31 years old. "
-            "Make the response as short as possible."),
-                               sampling_params=sampling_params,
-                               use_tqdm=True)
+    outputs = llm.generate(
+        prompts=("Generate a JSON object with curly braces for a person with "
+                 "name and age fields for John Smith who is 31 years old. "
+                 "Make the response as short as possible."),
+        sampling_params=sampling_params,
+        use_tqdm=True)

-        assert outputs is not None
-        for output in outputs:
-            assert output is not None
-            assert isinstance(output, RequestOutput)
+    assert outputs is not None
+    for output in outputs:
+        assert output is not None
+        assert isinstance(output, RequestOutput)

-            for i in range(2):
-                generated_text = output.outputs[i].text
-                print(generated_text)
-                assert generated_text is not None
+        for i in range(2):
+            generated_text = output.outputs[i].text
+            print(generated_text)
+            assert generated_text is not None

-                # Parse to verify it is a valid JSON object
-                parsed_json = json.loads(generated_text)
-                assert isinstance(parsed_json, dict)
+            # Parse to verify it is a valid JSON object
+            parsed_json = json.loads(generated_text)
+            assert isinstance(parsed_json, dict)

    #
    # Test 3: test a jsonschema incompatible with xgrammar
@@ -217,98 +210,97 @@ def test_structured_output(
            parsed_json = json.loads(generated_text)
            assert isinstance(parsed_json, dict)

-    if guided_decoding_backend != "outlines":
-        #
-        # Test 4: Generate SQL statement using EBNF grammar
-        #
-        sampling_params = SamplingParams(
-            temperature=0.8,
-            top_p=0.95,
-            max_tokens=1000,
-            guided_decoding=GuidedDecodingParams(grammar=sample_sql_ebnf))
-        outputs = llm.generate(
+    #
+    # Test 4: Generate SQL statement using EBNF grammar
+    #
+    sampling_params = SamplingParams(
+        temperature=0.8,
+        top_p=0.95,
+        max_tokens=1000,
+        guided_decoding=GuidedDecodingParams(grammar=sample_sql_ebnf))
+    outputs = llm.generate(
+        prompts=(
+            "Generate a sql statement that selects col_1 from "
+            "table_1 where it is equal to 1. Make the response as short as "
+            "possible."),
+        sampling_params=sampling_params,
+        use_tqdm=True,
+    )
+
+    assert outputs is not None
+    for output in outputs:
+        assert output is not None
+        assert isinstance(output, RequestOutput)
+        prompt = output.prompt
+
+        generated_text = output.outputs[0].text
+        assert generated_text is not None
+
+        # remove spaces for comparison b/c we removed them in the grammar
+        ground_truth = "SELECT col_1 from table_1 where col_1 = 1".replace(
+            " ", "")
+
+        assert generated_text.strip() == ground_truth
+
+        print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
+
+    #
+    # Test 5: Generate SQL statement using Lark grammar
+    #
+    sampling_params = SamplingParams(
+        temperature=0.8,
+        top_p=0.95,
+        max_tokens=1000,
+        guided_decoding=GuidedDecodingParams(grammar=sample_sql_lark))
+    outputs = llm.generate(
+        prompts=(
+            "Generate a sql statement that selects col_1 from "
+            "table_1 where it is equal to 1. Make the response as short as "
+            "possible."),
+        sampling_params=sampling_params,
+        use_tqdm=True,
+    )
+
+    assert outputs is not None
+    for output in outputs:
+        assert output is not None
+        assert isinstance(output, RequestOutput)
+        prompt = output.prompt
+
+        generated_text = output.outputs[0].text
+        assert generated_text is not None
+
+        # use Lark to parse the output, and make sure it's a valid parse tree
+        from lark import Lark
+        parser = Lark(sample_sql_lark)
+        parser.parse(generated_text)
+
+        # remove spaces for comparison b/c we removed them in the grammar
+        ground_truth = "SELECT col_1 from table_1 where col_1 = 1".replace(
+            " ", "")
+
+        assert generated_text.strip() == ground_truth
+
+        print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
+
+    #
+    # Test 6: Test invalid grammar input
+    #
+    sampling_params = SamplingParams(
+        temperature=0.8,
+        top_p=0.95,
+        max_tokens=1000,
+        guided_decoding=GuidedDecodingParams(grammar="not a grammar"))
+    with pytest.raises(ValueError, match="Failed to convert the grammar "):
+        llm.generate(
            prompts=(
                "Generate a sql statement that selects col_1 from "
-                "table_1 where it is equal to 1. Make the response as short as "
-                "possible."),
+                "table_1 where it is equal to 1. Make the response as short "
+                "as possible."),
            sampling_params=sampling_params,
            use_tqdm=True,
        )

-        assert outputs is not None
-        for output in outputs:
-            assert output is not None
-            assert isinstance(output, RequestOutput)
-            prompt = output.prompt
-
-            generated_text = output.outputs[0].text
-            assert generated_text is not None
-
-            # remove spaces for comparison b/c we removed them in the grammar
-            ground_truth = "SELECT col_1 from table_1 where col_1 = 1".replace(
-                " ", "")
-
-            assert generated_text.strip() == ground_truth
-
-            print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
-
-        #
-        # Test 5: Generate SQL statement using Lark grammar
-        #
-        sampling_params = SamplingParams(
-            temperature=0.8,
-            top_p=0.95,
-            max_tokens=1000,
-            guided_decoding=GuidedDecodingParams(grammar=sample_sql_lark))
-        outputs = llm.generate(
-            prompts=(
-                "Generate a sql statement that selects col_1 from "
-                "table_1 where it is equal to 1. Make the response as short as "
-                "possible."),
-            sampling_params=sampling_params,
-            use_tqdm=True,
-        )
-
-        assert outputs is not None
-        for output in outputs:
-            assert output is not None
-            assert isinstance(output, RequestOutput)
-            prompt = output.prompt
-
-            generated_text = output.outputs[0].text
-            assert generated_text is not None
-
-            # use Lark to parse the output, and make sure it's a valid parse tree
-            from lark import Lark
-            parser = Lark(sample_sql_lark)
-            parser.parse(generated_text)
-
-            # remove spaces for comparison b/c we removed them in the grammar
-            ground_truth = "SELECT col_1 from table_1 where col_1 = 1".replace(
-                " ", "")
-
-            assert generated_text.strip() == ground_truth
-
-            print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
-
-        #
-        # Test 6: Test invalid grammar input
-        #
-        sampling_params = SamplingParams(
-            temperature=0.8,
-            top_p=0.95,
-            max_tokens=1000,
-            guided_decoding=GuidedDecodingParams(grammar="not a grammar"))
-        with pytest.raises(ValueError, match="Failed to convert the grammar "):
-            llm.generate(
-                prompts=
-                ("Generate a sql statement that selects col_1 from "
-                 "table_1 where it is equal to 1. Make the response as short "
-                 "as possible."),
-                sampling_params=sampling_params,
-                use_tqdm=True,
-            )
-
    #
    # Test 7: Generate text based on a regex pattern
    #
@@ -429,36 +421,35 @@ def test_structured_output(
        output_json = json.loads(generated_text)
        jsonschema.validate(instance=output_json, schema=json_schema)

-    if guided_decoding_backend != "outlines":
-        #
-        # Test 11: Generate structured output using structural_tag format
-        #
-        structural_tag_config = {
-            "type":
-            "structural_tag",
-            "structures": [{
-                "begin": "<function=get_weather>",
-                "schema": {
-                    "type": "object",
-                    "properties": {
-                        "city": {
-                            "type": "string"
-                        }
-                    },
-                    "additionalProperties": False
+    #
+    # Test 11: Generate structured output using structural_tag format
+    #
+    structural_tag_config = {
+        "type":
+        "structural_tag",
+        "structures": [{
+            "begin": "<function=get_weather>",
+            "schema": {
+                "type": "object",
+                "properties": {
+                    "city": {
+                        "type": "string"
+                    }
                },
-                "end": "</function>"
-            }],
-            "triggers": ["<function="]
-        }
+                "additionalProperties": False
+            },
+            "end": "</function>"
+        }],
+        "triggers": ["<function="]
+    }

-        sampling_params = SamplingParams(
-            temperature=0.0,
-            max_tokens=4096,
-            guided_decoding=GuidedDecodingParams(
-                structural_tag=json.dumps(structural_tag_config)))
+    sampling_params = SamplingParams(
+        temperature=0.0,
+        max_tokens=4096,
+        guided_decoding=GuidedDecodingParams(
+            structural_tag=json.dumps(structural_tag_config)))

-        prompt = """
+    prompt = """
 You have access to the following function to retrieve the weather in a city:

    {
@@ -478,7 +469,7 @@ where

 start_tag => `<function`
 parameters => a JSON dict with the function argument name
-            as key and function argument value as value.
+              as key and function argument value as value.
 end_tag => `</function>`

 Here is an example,
@@ -497,37 +488,37 @@ Given the previous instructions, what is the weather in New York City? \
 Make the response as short as possible.
 """

-        # Change this once other backends support structural_tag
-        outputs = llm.generate(prompts=prompt,
-                               sampling_params=sampling_params,
-                               use_tqdm=True)
-        assert outputs is not None
+    # Change this once other backends support structural_tag
+    outputs = llm.generate(prompts=prompt,
+                           sampling_params=sampling_params,
+                           use_tqdm=True)
+    assert outputs is not None

-        for output in outputs:
-            assert output is not None
-            assert isinstance(output, RequestOutput)
-            generated_text = output.outputs[0].text
-            assert generated_text is not None
+    for output in outputs:
+        assert output is not None
+        assert isinstance(output, RequestOutput)
+        generated_text = output.outputs[0].text
+        assert generated_text is not None

-            # Search for function call pattern in the response
-            function_call_pattern = r'<function=get_weather>(.*?)</function>'
-            matches = re.findall(function_call_pattern, generated_text)
+        # Search for function call pattern in the response
+        function_call_pattern = r'<function=get_weather>(.*?)</function>'
+        matches = re.findall(function_call_pattern, generated_text)

-            if not matches:
-                print(f"Warning: No function calls found in response: "
-                      f"{generated_text!r}")
-                continue
+        if not matches:
+            print(f"Warning: No function calls found in response: "
+                  f"{generated_text!r}")
+            continue

-            # Take the first function call if multiple are found
-            json_str = matches[0]
-            try:
-                json_content = json.loads(json_str)
-                assert "city" in json_content
-                assert isinstance(json_content["city"], str)
-                print(f"Found valid function call: {generated_text!r}")
-            except (json.JSONDecodeError, AssertionError) as e:
-                pytest.fail("Invalid function call format: "
-                            f"{generated_text!r}\nError: {str(e)}")
+        # Take the first function call if multiple are found
+        json_str = matches[0]
+        try:
+            json_content = json.loads(json_str)
+            assert "city" in json_content
+            assert isinstance(json_content["city"], str)
+            print(f"Found valid function call: {generated_text!r}")
+        except (json.JSONDecodeError, AssertionError) as e:
+            pytest.fail("Invalid function call format: "
+                        f"{generated_text!r}\nError: {str(e)}")


@pytest.mark.skip_global_cleanup
@@ -3,10 +3,16 @@
 import filecmp
 import shutil
 import tempfile
+from collections import defaultdict
 from pathlib import Path

 from vllm import LLM, SamplingParams
-from vllm.config import KVTransferConfig
+from vllm.config import KVTransferConfig, VllmConfig
+from vllm.distributed.kv_transfer.kv_connector.factory import (
+    KVConnectorFactory)
+from vllm.distributed.kv_transfer.kv_connector.v1.shared_storage_connector import (  # noqa
+    SharedStorageConnector)
+from vllm.v1.core.kv_cache_manager import KVCacheBlocks

 MODEL_NAME = "meta-llama/Llama-3.2-1B-Instruct"

@@ -19,6 +25,65 @@ PROMPTS = [
 SAMPLING_PARAMS = SamplingParams(temperature=0, max_tokens=20)


+class TestSharedStorageConnector(SharedStorageConnector):
+
+    def __init__(self, config: VllmConfig, role):
+        self.name = config.kv_transfer_config.kv_connector_extra_config["name"]
+        self._connector = SharedStorageConnector(config, role)
+        self.call_record: dict[str, int] = defaultdict(int)
+        # Use a unique temp file per connector
+        self._event_file = tempfile.gettempdir(
+        ) + f"/connector_{self.name}-{self.role.name}_events.log"
+        # Start with an empty file
+        with open(self._event_file, "w") as _:
+            pass
+
+    def __getattribute__(self, name):
+        if name in ("_connector", "call_record", "name", "_event_file",
+                    "__class__", "__dict__", "__getattribute__",
+                    "__init__"):  # avoid recursion
+            return object.__getattribute__(self, name)
+        if not hasattr(self._connector, name):
+            return object.__getattribute__(self, name)
+        attr = getattr(self._connector, name)
+
+        # Intercept calls to the connector interface and write an event
+        # for each one to a file, which can be read back in the main test proc.
+        if callable(attr):
+
+            def wrapper(*args, **kwargs):
+                self.call_record[name] += 1
+
+                # Include args that we're interested in
+                to_log = [name]
+                for arg in args:
+                    if isinstance(arg, int):
+                        to_log.append(str(arg))
+                    elif isinstance(arg, KVCacheBlocks):
+                        to_log.append(
+                            f"num_blocks={[len(b) for b in arg.blocks]}")
+
+                # Log the event as a line to the file
+                try:
+                    with open(self._event_file, "a") as f:
+                        f.write(' '.join(to_log) + "\n")
+                except Exception as e:
+                    print(f"[ERROR] Could not log event {name} "
+                          f"for {self.name}: {e}")
+                return attr(*args, **kwargs)
+
+            return wrapper
+        return attr
+
+
+# This relies on "fork" multiprocessing method being used.
+# It's the default but vLLM may fall back to spawn if for example CUDA
+# is already initialized.
+KVConnectorFactory.register_connector("TestSharedStorageConnector",
+                                      TestSharedStorageConnector.__module__,
+                                      TestSharedStorageConnector.__name__)
+
+
 # Helper function to compare directories recursively
 def _compare_directories(dir1: Path, dir2: Path) -> bool:
    """Compares two directories recursively for identical content."""
@@ -53,27 +118,19 @@ def test_multi_shared_storage_connector_consistency():
        kv_role="kv_both",
        kv_connector_extra_config={
            "connectors": [{
-                "kv_connector":
-                "TestSharedStorageConnector",
-                "kv_role":
-                "kv_both",
+                "kv_connector": "TestSharedStorageConnector",
+                "kv_role": "kv_both",
                "kv_connector_extra_config": {
                    "shared_storage_path": str(storage_1_path),
                    "name": "storage1",
-                },
-                "kv_connector_module_path":
-                "tests.v1.kv_connector.unit.utils",
+                }
            }, {
-                "kv_connector":
-                "TestSharedStorageConnector",
-                "kv_role":
-                "kv_both",
+                "kv_connector": "TestSharedStorageConnector",
+                "kv_role": "kv_both",
                "kv_connector_extra_config": {
                    "shared_storage_path": str(storage_2_path),
                    "name": "storage2",
-                },
-                "kv_connector_module_path":
-                "tests.v1.kv_connector.unit.utils",
+                }
            }]
        },
    )
@@ -1,7 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-import tempfile
-from collections import defaultdict
 from typing import Any, Optional

 import torch
@@ -9,11 +7,6 @@ import torch
 from vllm import SamplingParams
 from vllm.config import (CacheConfig, DeviceConfig, KVTransferConfig,
                         ModelConfig, SchedulerConfig, VllmConfig)
-from vllm.distributed.kv_transfer.kv_connector.factory import (
-    KVConnectorFactory)
-from vllm.distributed.kv_transfer.kv_connector.v1.shared_storage_connector import (  # noqa
-    SharedStorageConnector)
-from vllm.v1.core.kv_cache_manager import KVCacheBlocks
 from vllm.v1.core.sched.scheduler import Scheduler
 from vllm.v1.kv_cache_interface import (FullAttentionSpec, KVCacheConfig,
                                        KVCacheGroupSpec)
@@ -194,58 +187,3 @@ def create_model_runner_output(
        finished_sending=finished_sending,
        finished_recving=finished_recving,
    )
-
-
-class TestSharedStorageConnector(SharedStorageConnector):
-
-    def __init__(self, config: VllmConfig, role):
-        self.name = config.kv_transfer_config.kv_connector_extra_config["name"]
-        self._connector = SharedStorageConnector(config, role)
-        self.call_record: dict[str, int] = defaultdict(int)
-        # Use a unique temp file per connector
-        self._event_file = tempfile.gettempdir(
-        ) + f"/connector_{self.name}-{self.role.name}_events.log"
-        # Start with an empty file
-        with open(self._event_file, "w") as _:
-            pass
-
-    def __getattribute__(self, name):
-        if name in ("_connector", "call_record", "name", "_event_file",
-                    "__class__", "__dict__", "__getattribute__",
-                    "__init__"):  # avoid recursion
-            return object.__getattribute__(self, name)
-        if not hasattr(self._connector, name):
-            return object.__getattribute__(self, name)
-        attr = getattr(self._connector, name)
-
-        # Intercept calls to the connector interface and write an event
-        # for each one to a file, which can be read back in the main test proc.
-        if callable(attr):
-
-            def wrapper(*args, **kwargs):
-                self.call_record[name] += 1
-
-                # Include args that we're interested in
-                to_log = [name]
-                for arg in args:
-                    if isinstance(arg, int):
-                        to_log.append(str(arg))
-                    elif isinstance(arg, KVCacheBlocks):
-                        to_log.append(
-                            f"num_blocks={[len(b) for b in arg.blocks]}")
-
-                # Log the event as a line to the file
-                try:
-                    with open(self._event_file, "a") as f:
-                        f.write(' '.join(to_log) + "\n")
-                except Exception as e:
-                    print(f"[ERROR] Could not log event {name} "
-                          f"for {self.name}: {e}")
-                return attr(*args, **kwargs)
-
-            return wrapper
-        return attr
-
-
-KVConnectorFactory.register_connector("TestSharedStorageConnector", __name__,
-                                      TestSharedStorageConnector.__name__)
@@ -1843,26 +1843,6 @@ def cutlass_mla_decode(out: torch.Tensor, q_nope: torch.Tensor,
    return out


-def sm100_cutlass_mla_decode(out: torch.Tensor, q_nope: torch.Tensor,
-                             q_pe: torch.Tensor,
-                             kv_c_and_k_pe_cache: torch.Tensor,
-                             seq_lens: torch.Tensor, page_table: torch.Tensor,
-                             workspace: torch.Tensor, scale: float,
-                             num_kv_splits: int) -> torch.Tensor:
-    torch.ops._C.sm100_cutlass_mla_decode(out, q_nope, q_pe,
-                                          kv_c_and_k_pe_cache, seq_lens,
-                                          page_table, workspace, scale,
-                                          num_kv_splits)
-    return out
-
-
-def sm100_cutlass_mla_get_workspace_size(max_seq_len: int, num_batches: int,
-                                         sm_count: int,
-                                         num_kv_splits: int) -> int:
-    return torch.ops._C.sm100_cutlass_mla_get_workspace_size(
-        max_seq_len, num_batches, sm_count, num_kv_splits)
-
-
 if hasattr(torch.ops._C, "weight_packed_linear"):

    @register_fake("_C::weight_packed_linear")
@@ -9,8 +9,6 @@ from typing import (TYPE_CHECKING, Any, Dict, Generic, List, Optional,

 import torch

-from vllm.model_executor.layers.quantization.utils.quant_utils import (
-    GroupShape)
 from vllm.multimodal import MultiModalPlaceholderMap

 if TYPE_CHECKING:
@@ -291,7 +289,7 @@ class AttentionImpl(ABC, Generic[T]):
        raise NotImplementedError

    def fused_output_quant_supported(self, dtype: torch.dtype, static: bool,
-                                     group_shape: GroupShape):
+                                     group_shape: tuple[int, int]):
        """
        Does this attention implementation support fused output quantization.
        This is used by the AttnFusionPass to only fuse output quantization
@@ -300,7 +298,7 @@ class AttentionImpl(ABC, Generic[T]):
        TODO(luka) merge parameters into QuantDescriptor
        :param dtype: quantized dtype
        :param static: static or dynamic quantization
-        :param group_shape: quant group shape.
+        :param group_shape: quant group shape. (-1, -1) for per-tensor.
        :return: is fusion supported for this type of quantization
        """
        return False
@@ -11,8 +11,7 @@ from vllm.multimodal import MultiModalPlaceholderMap

 try:
    from flashinfer import BatchDecodeWithPagedKVCacheWrapper
-    from flashinfer.decode import (CUDAGraphBatchDecodeWithPagedKVCacheWrapper,
-                                   trtllm_batch_decode_with_kv_cache)
+    from flashinfer.decode import CUDAGraphBatchDecodeWithPagedKVCacheWrapper
    from flashinfer.prefill import BatchPrefillWithPagedKVCacheWrapper

    from vllm.vllm_flash_attn import flash_attn_varlen_func
@@ -23,10 +22,7 @@ except ImportError:
        BatchDecodeWithPagedKVCacheWrapper = None
        CUDAGraphBatchDecodeWithPagedKVCacheWrapper = None
        BatchPrefillWithPagedKVCacheWrapper = None
-        trtllm_batch_decode_with_kv_cache = None
    FLASHINFER_WORKSPACE_BUFFER_SIZE = 0
-    raise ImportError("FlashInfer is not installed. Please install it from "
-                      "https://github.com/flashinfer-ai/flashinfer") from None

 import torch

@@ -44,7 +40,6 @@ from vllm.attention.layer import Attention
 from vllm.attention.ops.paged_attn import PagedAttention
 from vllm.config import VllmConfig, get_layers_from_vllm_config
 from vllm.logger import init_logger
-from vllm.platforms import current_platform
 from vllm.utils import (async_tensor_h2d, get_kv_cache_torch_dtype,
                        make_tensor_with_pad)

@@ -54,9 +49,10 @@ if TYPE_CHECKING:
    from vllm.worker.model_runner import (ModelInputForGPUBuilder,
                                          ModelInputForGPUWithSamplingMetadata)

+FLASHINFER_KV_CACHE_LAYOUT: str = envs.VLLM_KV_CACHE_LAYOUT or "NHD"
+

 class FlashInferBackend(AttentionBackend):
-    cached_sm100a_supported: Optional[bool] = None

    @staticmethod
    def get_name() -> str:
@@ -89,7 +85,7 @@ class FlashInferBackend(AttentionBackend):

    @staticmethod
    def get_kv_cache_stride_order() -> Tuple[int, ...]:
-        cache_layout = FlashInferState.get_kv_cache_layout()
+        cache_layout = FLASHINFER_KV_CACHE_LAYOUT
        assert (cache_layout in ("NHD", "HND"))
        stride_order = (0, 1, 2, 3, 4) if cache_layout == "NHD" else (0, 1, 3,
                                                                      2, 4)
@@ -123,47 +119,6 @@ class FlashInferBackend(AttentionBackend):
        else:
            raise ValueError(f"Unrecognized FP8 dtype: {kv_cache_dtype}")

-    @staticmethod
-    def use_trtllm_decode_attention(
-        batch_size: int,
-        max_seq_len: int,
-        kv_cache_dtype: str,
-        num_qo_heads: Optional[int],
-        num_kv_heads: Optional[int],
-        attn_head_size: Optional[int],
-    ) -> bool:
-        if FlashInferBackend.cached_sm100a_supported is None:
-            FlashInferBackend.cached_sm100a_supported = (
-                current_platform.has_device_capability(100))
-        if not FlashInferBackend.cached_sm100a_supported:
-            return False
-        # Check if the dimensions are supported by TRTLLM decode attention
-        if (attn_head_size is None or num_qo_heads is None
-                or num_kv_heads is None or num_qo_heads // num_kv_heads > 8
-                or num_qo_heads % num_kv_heads != 0 or attn_head_size != 128):
-            return False
-        env_value = envs.VLLM_USE_TRTLLM_DECODE_ATTENTION
-        if env_value is not None:
-            logger.info_once("VLLM_USE_TRTLLM_DECODE_ATTENTION is set to %s",
-                             env_value)
-            # Environment variable is set - respect it
-            # Making the conditional check for zero because
-            # the path is automatically enabled if the batch size condition
-            # is satisfied.
-            no_use_trtllm = (env_value == "0")
-            if not no_use_trtllm:
-                logger.info_once("Using TRTLLM decode attention.")
-            return not no_use_trtllm
-        else:
-            # Environment variable not set - use auto-detection
-            use_trtllm = (FlashInferBackend.cached_sm100a_supported
-                          and batch_size <= 256 and max_seq_len < 131072
-                          and kv_cache_dtype == "auto")
-            if use_trtllm:
-                logger.warning_once(
-                    "Using TRTLLM decode attention (auto-detected).")
-        return use_trtllm
-

@dataclass
 class PerLayerParameters:
@@ -252,19 +207,10 @@ class FlashInferState(AttentionState):
                device=self.runner.device)
        return self._workspace_buffer

-    @staticmethod
-    def get_kv_cache_layout():
-        from vllm.v1.attention.backends.utils import _KV_CACHE_LAYOUT_OVERRIDE
-        if _KV_CACHE_LAYOUT_OVERRIDE is not None:
-            logger.info_once("Using KV cache layout %s",
-                             _KV_CACHE_LAYOUT_OVERRIDE)
-            return _KV_CACHE_LAYOUT_OVERRIDE
-        cache_layout = envs.VLLM_KV_CACHE_LAYOUT
-        if cache_layout is None:
-            logger.info_once("Using default KV cache layout NHD")
-            return "NHD"
-        logger.info_once("Using KV cache layout %s", cache_layout)
-        return cache_layout
+    def get_kv_cache_layout(self):
+        if self._kv_cache_layout is None:
+            self._kv_cache_layout = FLASHINFER_KV_CACHE_LAYOUT
+        return self._kv_cache_layout

    def _get_prefill_wrapper(self):
        if self._prefill_wrapper is None:
@@ -377,8 +323,6 @@ class FlashInferState(AttentionState):
            num_prefill_tokens=0,
            num_decode_tokens=batch_size,
            max_prefill_seq_len=0,
-            max_decode_seq_len=0,
-            seq_lens_tensor=self._graph_seq_lens,
            block_tables=self._graph_block_tables,
            paged_kv_indptr=paged_kv_indptr_tensor_host,
            paged_kv_indices=paged_kv_indices_tensor_host,
@@ -404,8 +348,6 @@ class FlashInferState(AttentionState):
                                attn_metadata,
                                is_encoder_decoder_model: bool = False):
        return {
-            "block_tables": attn_metadata.block_tables,
-            "seq_lens_tensor": attn_metadata.seq_lens_tensor,
            "slot_mapping": attn_metadata.slot_mapping,
        }

@@ -413,13 +355,7 @@ class FlashInferState(AttentionState):
                                    input_buffers,
                                    attn_metadata,
                                    is_encoder_decoder_model: bool = False):
-        # FlashInfer-specific logic: copy additional tensors
-        num_total_blocks = attn_metadata.decode_metadata.seq_lens_tensor.shape[
-            0]
-        input_buffers["seq_lens_tensor"][:num_total_blocks].copy_(
-            attn_metadata.seq_lens_tensor, non_blocking=True)
-        input_buffers["block_tables"][:num_total_blocks].copy_(
-            attn_metadata.block_tables, non_blocking=True)
+        return

    def begin_forward(self, model_input):
        assert not self._is_graph_capturing
@@ -452,8 +388,6 @@ class FlashInferMetadata(AttentionMetadata):
    # Maximum sequence length among prefill batch. 0 if there are decoding
    # requests only.
    max_prefill_seq_len: int
-    max_decode_seq_len: int
-
    # Number of query tokens for each request in the batch.
    # Currently, we require that all requests have the same number of query
    # tokens during the decoding phase. When speculavie decoding is enabled,
@@ -856,7 +790,6 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
        use_captured_graph = cuda_graph_pad_size != -1

        max_prefill_seq_len = max(self.prefill_seq_lens, default=0)
-        max_decode_seq_len = max(self.curr_seq_lens, default=0)
        num_decode_tokens = self.num_decode_tokens
        decode_query_len = max(query_lens[self.num_prefills:], default=1)

@@ -962,7 +895,6 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
            num_prefill_tokens=self.num_prefill_tokens,
            num_decode_tokens=num_decode_tokens,
            max_prefill_seq_len=max_prefill_seq_len,
-            max_decode_seq_len=max_decode_seq_len,
            block_tables=block_tables,
            paged_kv_indptr=paged_kv_indptr_tensor,
            paged_kv_indices=paged_kv_indices_tensor,
@@ -1149,36 +1081,13 @@ class FlashInferImpl(AttentionImpl):
            assert decode_meta.decode_wrapper._logits_soft_cap == (
                logits_soft_cap or 0.0)
            assert decode_meta.decode_wrapper._sm_scale == softmax_scale
-            # TODO: @pavanimajety Remove this once the switch happens
-            # inside flashinfer.
-            if not FlashInferBackend.use_trtllm_decode_attention(
-                    num_decode_tokens, attn_metadata.max_decode_seq_len,
-                    kv_cache_dtype, attn_metadata.num_qo_heads,
-                    attn_metadata.num_kv_heads, attn_metadata.head_dim):
-                decode_output = decode_meta.decode_wrapper.run(
-                    decode_query,
-                    kv_cache.permute(*stride_order),
-                    k_scale=layer._k_scale_float,
-                    v_scale=layer._v_scale_float,
-                )
-            else:
-                workspace_buffer = (
-                    decode_meta.decode_wrapper._int_workspace_buffer)
-                assert FlashInferState.get_kv_cache_layout() == "HND"
-                decode_output = trtllm_batch_decode_with_kv_cache(
-                    query=decode_query,
-                    kv_cache=kv_cache.permute(*stride_order),
-                    workspace_buffer=workspace_buffer,
-                    num_heads=num_heads,
-                    num_kv_heads=num_kv_heads,
-                    scale=softmax_scale,
-                    block_tables=attn_metadata.block_tables,
-                    seq_lens=decode_meta.seq_lens_tensor,
-                    block_size=attn_metadata.page_size,
-                    max_seq_len=attn_metadata.max_decode_seq_len,
-                    kv_cache_dtype=kv_cache_dtype,
-                    k_scale=layer._k_scale_float,
-                    v_scale=layer._v_scale_float)
+
+            decode_output = decode_meta.decode_wrapper.run(
+                decode_query,
+                kv_cache.permute(*stride_order),
+                k_scale=layer._k_scale_float,
+                v_scale=layer._v_scale_float,
+            )

        if prefill_output is None and decode_output is not None:
            # Decode only batch.
@@ -19,8 +19,6 @@ from vllm.attention.ops.paged_attn import (PagedAttention,
                                           PagedAttentionMetadata)
 from vllm.config import get_current_vllm_config
 from vllm.logger import init_logger
-from vllm.model_executor.layers.quantization.utils.quant_utils import (
-    GroupShape)
 from vllm.platforms import current_platform
 from vllm.platforms.rocm import use_rocm_custom_paged_attention

@@ -600,10 +598,10 @@ class ROCmFlashAttentionImpl(AttentionImpl):
                                                    head_dim))

    def fused_output_quant_supported(self, dtype: torch.dtype, static: bool,
-                                     group_shape: GroupShape):
+                                     group_shape: tuple[int, int]):
        if self.use_triton_flash_attn:
            return dtype == current_platform.fp8_dtype(
-            ) and static and group_shape == GroupShape.PER_TENSOR
+            ) and static and group_shape == (-1, -1)  # per-tensor

        # Only supported in the Triton backend
        return False
@@ -10,7 +10,6 @@ import torch.nn.functional as F
 import vllm.envs as envs
 from vllm.attention import AttentionType
 from vllm.attention.selector import backend_name_to_enum, get_attn_backend
-from vllm.attention.utils.kv_sharing_utils import validate_kv_sharing_target
 from vllm.config import CacheConfig, get_current_vllm_config
 from vllm.distributed.kv_transfer import (get_kv_transfer_group,
                                          has_kv_transfer_group,
@@ -22,6 +21,7 @@ from vllm.model_executor.layers.quantization.base_config import (
 from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod
 from vllm.platforms import _Backend, current_platform
 from vllm.utils import direct_register_custom_op
+from vllm.v1.attention.backends.utils import validate_kv_sharing_target


 class Attention(nn.Module):
@@ -1,33 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-def validate_kv_sharing_target(current_layer_name, target_layer_name,
-                               static_forward_context):
-    error_msg = (f"Specified KV sharing target layer for {current_layer_name} "
-                 f"is not valid: target layer {target_layer_name} ")
-
-    if current_layer_name == target_layer_name:
-        raise ValueError(error_msg +
-                         "cannot be the same as the current layer.")
-
-    if target_layer_name not in static_forward_context:
-        from vllm.model_executor.models.utils import extract_layer_index
-
-        # If target layer name is not in the static fwd context, it means either
-        # a) the target layer does not come BEFORE the current layer, or
-        # b) the target layer is not an Attention layer that exists in the model
-        current_layer_idx = extract_layer_index(current_layer_name)
-        target_layer_idx = extract_layer_index(target_layer_name)
-        if current_layer_idx <= target_layer_idx:
-            raise ValueError(error_msg + "must come before the current layer.")
-        else:
-            raise ValueError(error_msg +
-                             "is not a valid Attention layer in the model.")
-
-    # Currently KV sharing is only supported between layers of the same type
-    target_layer_attn_type = static_forward_context[
-        target_layer_name].attn_type
-    expected = static_forward_context[current_layer_name].attn_type
-    if target_layer_attn_type != expected:
-        raise ValueError(
-            error_msg +
-            f"must be the same type as the current layer ({expected}).")
@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-from typing import Callable, NamedTuple, Optional
+from typing import Callable, ClassVar, NamedTuple, Optional

 import torch
 import torch._inductor.pattern_matcher as pm
@@ -11,8 +11,6 @@ from torch._ops import OpOverload

 from vllm.config import VllmConfig
 from vllm.logger import init_logger
-from vllm.model_executor.layers.quantization.utils.quant_utils import (
-    GroupShape)
 from vllm.platforms import current_platform

 from .fx_utils import find_getitem_maybe
@@ -35,6 +33,27 @@ RMS_OP = torch.ops._C.rms_norm.default
 RMS_ADD_OP = torch.ops._C.fused_add_rms_norm.default


+# Use proxy as NamedTuple direct subclasses cannot have static members
+class _GroupShape(NamedTuple):
+    row: int
+    col: int
+
+
+class GroupShape(_GroupShape):
+    """
+    This class describes the quantization group shape.
+    It includes static members for common shapes (per-tensor, per-token).
+    """
+
+    # Aliases for common quantization group shapes
+    PER_TENSOR: ClassVar['GroupShape']
+    PER_TOKEN: ClassVar['GroupShape']
+
+
+GroupShape.PER_TENSOR = GroupShape(-1, -1)
+GroupShape.PER_TOKEN = GroupShape(1, -1)
+
+
 class QuantKey(NamedTuple):
    """
    Named tuple for identifying the type of quantization.
@@ -6,7 +6,13 @@ import time
 import torch
 from torch._dynamo.utils import lazy_format_graph_code

-from vllm.config import VllmConfig
+from vllm.config import PassConfig, VllmConfig
+# yapf: disable
+from vllm.distributed import get_tensor_model_parallel_rank as get_tp_rank
+from vllm.distributed import (
+    get_tensor_model_parallel_world_size as get_tp_world_size)
+from vllm.distributed import model_parallel_is_initialized as p_is_init
+# yapf: enable
 from vllm.logger import init_logger

 from .inductor_pass import InductorPass
@@ -28,9 +34,22 @@ class VllmInductorPass(InductorPass):
            else None
        self.pass_name = self.__class__.__name__

-    def dump_graph(self, graph: torch.fx.Graph, stage: str):
+    def dump_graph(self, graph: torch.fx.Graph, stage: str, always=False):
        lazy_format_graph_code(stage, graph.owning_module)

+        if stage in self.pass_config.dump_graph_stages or always:
+            # Make sure filename includes rank in the distributed setting
+            parallel = p_is_init() and get_tp_world_size() > 1
+            rank = f"-{get_tp_rank()}" if parallel else ""
+            filepath = self.pass_config.dump_graph_dir / f"{stage}{rank}.py"
+
+            logger.info("%s printing graph to %s", self.pass_name, filepath)
+            with open(filepath, "w") as f:
+                src = graph.python_code(root_module="self", verbose=True).src
+                # Add imports so it's not full of errors
+                print("import torch; from torch import device", file=f)
+                print(src, file=f)
+
    def begin(self):
        self._start_time = time.perf_counter_ns()

@@ -42,9 +61,10 @@ class VllmInductorPass(InductorPass):

 class PrinterInductorPass(VllmInductorPass):

-    def __init__(self, name: str, config: VllmConfig):
+    def __init__(self, name: str, config: PassConfig, always=False):
        super().__init__(config)
        self.name = name
+        self.always = always

    def __call__(self, graph: torch.fx.Graph):
-        self.dump_graph(graph, self.name)
+        self.dump_graph(graph, self.name, always=self.always)
@@ -16,6 +16,7 @@ from dataclasses import (MISSING, Field, asdict, field, fields, is_dataclass,
                         replace)
 from functools import cached_property
 from importlib.util import find_spec
+from pathlib import Path
 from typing import (TYPE_CHECKING, Any, Callable, ClassVar, Literal, Optional,
                    Protocol, TypeVar, Union, cast, get_args)

@@ -32,7 +33,6 @@ import vllm.envs as envs
 from vllm import version
 from vllm.compilation.inductor_pass import CallableInductorPass, InductorPass
 from vllm.logger import init_logger
-from vllm.model_executor.layers.quantization import QuantizationMethods
 from vllm.platforms import current_platform
 from vllm.transformers_utils.config import (
    ConfigFormat, get_config, get_hf_image_processor_config,
@@ -1142,7 +1142,7 @@ class ModelConfig:
        if not hasattr(self.hf_text_config, "model_type"):
            return False
        elif self.hf_text_config.model_type in \
-            ('deepseek_v2', 'deepseek_v3', 'deepseek_mtp', 'kimi_k2'):
+            ('deepseek_v2', 'deepseek_v3', 'deepseek_mtp'):
            return self.hf_text_config.kv_lora_rank is not None
        elif self.hf_text_config.model_type == 'eagle':
            # if the model is an EAGLE module, check for the
@@ -1330,17 +1330,6 @@ class ModelConfig:

            return sum(t == 1 for t in attn_type_list[start:end])

-    def get_mamba_chunk_size(self) -> Optional[int]:
-        """
-        Returns the mamba chunk size if it exists
-        """
-        # used by e.g. Bamba, FalconH1, Granite, PLaMo2
-        chunk_size = getattr(self.hf_text_config, "mamba_chunk_size", None)
-        if chunk_size is None:
-            # used by e.g. Mamba2, NemotronH, Zamba
-            chunk_size = getattr(self.hf_text_config, "chunk_size", None)
-        return chunk_size
-
    def get_multimodal_config(self) -> "MultiModalConfig":
        """
        Get the multimodal configuration of the model.
@@ -3000,16 +2989,6 @@ class LoRAConfig:
    trained with those scaling factors to be used at the same time. If not
    specified, only adapters trained with the base model scaling factor are
    allowed."""
-    default_mm_loras: Optional[dict[str, str]] = None
-    """Dictionary mapping specific modalities to LoRA model paths; this field
-    is only applicable to multimodal models and should be leveraged when a
-    model always expects a LoRA to be active when a given modality is present.
-    Note that currently, if a request provides multiple additional
-    modalities, each of which have their own LoRA, we do NOT apply
-    default_mm_loras because we currently only support one lora adapter
-    per prompt. When run in offline mode, the lora IDs for n modalities
-    will be automatically assigned to 1-n with the names of the modalities
-    in alphabetic order."""
    bias_enabled: bool = False
    """Enable bias for LoRA adapters."""

@@ -3601,8 +3580,7 @@ def get_served_model_name(model: str,

 GuidedDecodingBackendV0 = Literal["auto", "outlines", "lm-format-enforcer",
                                  "xgrammar", "guidance"]
-
-GuidedDecodingBackendV1 = Literal["auto", "xgrammar", "guidance", "outlines"]
+GuidedDecodingBackendV1 = Literal["auto", "xgrammar", "guidance"]
 GuidedDecodingBackend = Literal[GuidedDecodingBackendV0,
                                GuidedDecodingBackendV1]

@@ -3952,6 +3930,11 @@ class PassConfig:
    don't all have access to full configuration - that would create a cycle as
    the `PassManager` is set as a property of config."""

+    dump_graph_stages: list[str] = field(default_factory=list)
+    """List of stages for which we want to dump the graph. Each pass defines
+    its own stages (before, after, maybe in-between)."""
+    dump_graph_dir: Path = Path(".")
+    """Directory to dump the graphs."""
    enable_fusion: bool = field(default_factory=lambda: not envs.VLLM_USE_V1)
    """Whether to enable the custom fusion (RMSNorm/SiluMul+quant) pass."""
    enable_attn_fusion: bool = False
@@ -3969,9 +3952,12 @@ class PassConfig:
        """
        Produces a hash unique to the pass configuration.
        Any new fields that affect compilation should be added to the hash.
-        Any future fields that don't affect compilation should be excluded.
+        Do not include dump_graph_* in the hash - they don't affect
+        compilation.
        """
-        return InductorPass.hash_dict(asdict(self))
+        exclude = {"dump_graph_stages", "dump_graph_dir"}
+        dict_ = {k: v for k, v in asdict(self).items() if k not in exclude}
+        return InductorPass.hash_dict(dict_)

    def __post_init__(self) -> None:
        if not self.enable_noop:
@@ -240,7 +240,8 @@ class DeviceCommunicatorBase:
            if module.__class__.__name__ == "FusedMoE"
        ]
        for module in moe_modules:
-            module.quant_method.init_prepare_finalize(module.moe_config)
+            module.quant_method.init_prepare_finalize(module.moe_config,
+                                                      module.quant_config)

    def dispatch(
            self, hidden_states: torch.Tensor,
@@ -57,7 +57,7 @@ class KVConnectorRole(enum.Enum):
    WORKER = 1


-class KVConnectorMetadata(ABC):  # noqa: B024
+class KVConnectorMetadata:
    """
    Abstract Metadata used to communicate between the
    Scheduler KVConnector and Worker KVConnector.
@@ -71,7 +71,7 @@ class KVConnectorBase_V1(ABC):
        logger.warning(
            "Initializing KVConnectorBase_V1. This API is experimental and "
            "subject to change in the future as we iterate the design.")
-        self._connector_metadata: Optional[KVConnectorMetadata] = None
+        self._connector_metadata = KVConnectorMetadata()
        self._vllm_config = vllm_config
        self._role = role

@@ -102,7 +102,7 @@ class KVConnectorBase_V1(ABC):
        This function should be called by the model runner every time 
        after the model execution.
        """
-        self._connector_metadata = None
+        self._connector_metadata = KVConnectorMetadata()

    def _get_connector_metadata(self) -> KVConnectorMetadata:
        """Get the connector metadata.
@@ -112,9 +112,6 @@ class KVConnectorBase_V1(ABC):
        Returns:
            ConnectorMetadata: the connector metadata.
        """
-
-        # Should only be called while set to valid metadata.
-        assert self._connector_metadata is not None
        return self._connector_metadata

    def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
@@ -395,8 +395,6 @@ class EngineArgs:
    enable_lora_bias: bool = LoRAConfig.bias_enabled
    max_loras: int = LoRAConfig.max_loras
    max_lora_rank: int = LoRAConfig.max_lora_rank
-    default_mm_loras: Optional[Dict[str, str]] = \
-        LoRAConfig.default_mm_loras
    fully_sharded_loras: bool = LoRAConfig.fully_sharded_loras
    max_cpu_loras: Optional[int] = LoRAConfig.max_cpu_loras
    lora_dtype: Optional[Union[str, torch.dtype]] = LoRAConfig.lora_dtype
@@ -809,8 +807,6 @@ class EngineArgs:
                                **lora_kwargs["max_cpu_loras"])
        lora_group.add_argument("--fully-sharded-loras",
                                **lora_kwargs["fully_sharded_loras"])
-        lora_group.add_argument("--default-mm-loras",
-                                **lora_kwargs["default_mm_loras"])

        # PromptAdapter related configs
        prompt_adapter_kwargs = get_kwargs(PromptAdapterConfig)
@@ -1288,16 +1284,10 @@ class EngineArgs:
            disable_hybrid_kv_cache_manager,
        )

-        if not model_config.is_multimodal_model and self.default_mm_loras:
-            raise ValueError(
-                "Default modality-specific LoRA(s) were provided for a "
-                "non multimodal model")
-
        lora_config = LoRAConfig(
            bias_enabled=self.enable_lora_bias,
            max_lora_rank=self.max_lora_rank,
            max_loras=self.max_loras,
-            default_mm_loras=self.default_mm_loras,
            fully_sharded_loras=self.fully_sharded_loras,
            lora_extra_vocab_size=self.lora_extra_vocab_size,
            long_lora_scaling_factors=self.long_lora_scaling_factors,
@@ -1424,8 +1414,6 @@ class EngineArgs:
                from vllm.attention.utils.fa_utils import (
                    flash_attn_supports_fp8)
                supported = flash_attn_supports_fp8()
-            elif envs.VLLM_USE_TRTLLM_DECODE_ATTENTION:
-                supported = True
            if not supported:
                _raise_or_fallback(feature_name="--kv-cache-dtype",
                                   recommend_to_remove=False)
@@ -1502,7 +1490,6 @@ class EngineArgs:
            "TRITON_ATTN_VLLM_V1",
            "TRITON_MLA",
            "CUTLASS_MLA_VLLM_V1",
-            "SM100_CUTLASS_MLA_VLLM_V1",
            "FLASHMLA",
            "FLASHINFER",
            "FLASHINFER_VLLM_V1",
@@ -499,10 +499,6 @@ class LLM:
        _validate_truncation_size(self.llm_engine.model_config.max_model_len,
                                  truncate_prompt_tokens, tokenization_kwargs)

-        # Add any modality specific loras to the corresponding prompts
-        lora_request = self._get_modality_specific_lora_reqs(
-            parsed_prompts, lora_request)
-
        self._validate_and_add_requests(
            prompts=parsed_prompts,
            params=sampling_params,
@@ -517,83 +513,6 @@ class LLM:
        outputs = self._run_engine(use_tqdm=use_tqdm)
        return self.engine_class.validate_outputs(outputs, RequestOutput)

-    def _get_modality_specific_lora_reqs(
-            self, parsed_prompts: Union[PromptType, Sequence[PromptType]],
-            lora_request: Optional[Union[list[LoRARequest], LoRARequest]]):
-        # Grab the lora config off the vllm config on the engine,
-        # since this is the same for both v0 & v1.
-        lora_config = self.llm_engine.vllm_config.lora_config
-
-        # If there's no lora config / default_mm_loras, or the model
-        # isn't multimodal, leave the lora as is.
-        if (lora_config is None
-                or not self.llm_engine.model_config.is_multimodal_model
-                or (lora_config and lora_config.default_mm_loras is None)):
-            return lora_request
-
-        if not isinstance(parsed_prompts, Sequence):
-            parsed_prompts = [parsed_prompts]
-
-        optional_loras = ([lora_request] * len(parsed_prompts)
-                          if not isinstance(lora_request, Sequence) else
-                          lora_request)
-
-        return [
-            self._resolve_single_prompt_mm_lora(
-                parsed_prompt,
-                opt_lora_req,
-                lora_config.default_mm_loras,
-            ) for parsed_prompt, opt_lora_req in zip(parsed_prompts,
-                                                     optional_loras)
-        ]
-
-    def _resolve_single_prompt_mm_lora(self, parsed_prompt: PromptType,
-                                       lora_request: Optional[LoRARequest],
-                                       default_mm_loras: Optional[dict[str,
-                                                                       str]]):
-        if (not default_mm_loras or not isinstance(parsed_prompt, dict)
-                or "multi_modal_data" not in parsed_prompt):
-            return lora_request
-
-        parsed_prompt = cast(Union[TextPrompt, TokensPrompt], parsed_prompt)
-
-        intersection = set(
-            parsed_prompt["multi_modal_data"].keys()).intersection(
-                default_mm_loras.keys())
-        if not intersection:
-            return lora_request
-        if len(intersection) > 1:
-            # TODO: Would be nice to be able to have multiple loras per prompt
-            logger.warning(
-                "Multiple modality specific loras were registered and would be"
-                " used by a single prompt consuming several modalities; "
-                " currently we only support one lora per request; as such,"
-                " lora(s) registered with modalities: %s"
-                " will be skipped", intersection)
-            return lora_request
-
-        # Build the LoRA request; the ID of the default mm lora is the
-        # index of the modality name sorted alphabetically + 1.
-        modality_name = intersection.pop()
-        modality_lora_path = default_mm_loras[modality_name]
-        modality_lora_id = sorted(default_mm_loras).index(modality_name) + 1
-
-        # If we have a collision, warn if there is a collision,
-        # but always send the explicitly provided request.
-        if lora_request:
-            if lora_request.lora_int_id != modality_lora_id:
-                logger.warning(
-                    "A modality with a registered lora and a lora_request "
-                    "with a different ID were provided; falling back to the "
-                    "lora_request as we only apply one LoRARequest per prompt")
-            return lora_request
-
-        return LoRARequest(
-            modality_name,
-            modality_lora_id,
-            modality_lora_path,
-        )
-
    def collective_rpc(self,
                       method: Union[str, Callable[..., _R]],
                       timeout: Optional[float] = None,
@@ -18,7 +18,7 @@ from collections.abc import AsyncIterator, Awaitable
 from contextlib import asynccontextmanager
 from functools import partial
 from http import HTTPStatus
-from typing import Annotated, Any, Callable, Optional
+from typing import Annotated, Any, Optional

 import prometheus_client
 import pydantic
@@ -61,9 +61,13 @@ from vllm.entrypoints.openai.protocol import (ChatCompletionRequest,
                                              CompletionResponse,
                                              DetokenizeRequest,
                                              DetokenizeResponse,
+                                              EmbeddingChatRequest,
+                                              EmbeddingCompletionRequest,
                                              EmbeddingRequest,
                                              EmbeddingResponse, ErrorResponse,
                                              LoadLoRAAdapterRequest,
+                                              PoolingChatRequest,
+                                              PoolingCompletionRequest,
                                              PoolingRequest, PoolingResponse,
                                              RerankRequest, RerankResponse,
                                              ResponsesRequest,
@@ -83,7 +87,6 @@ from vllm.entrypoints.openai.serving_completion import OpenAIServingCompletion
 from vllm.entrypoints.openai.serving_embedding import OpenAIServingEmbedding
 from vllm.entrypoints.openai.serving_engine import OpenAIServing
 from vllm.entrypoints.openai.serving_models import (BaseModelPath,
-                                                    LoRAModulePath,
                                                    OpenAIServingModels)
 from vllm.entrypoints.openai.serving_pooling import OpenAIServingPooling
 from vllm.entrypoints.openai.serving_responses import OpenAIServingResponses
@@ -430,7 +433,6 @@ async def get_server_load_metrics(request: Request):
    # - /v1/chat/completions
    # - /v1/completions
    # - /v1/audio/transcriptions
-    # - /v1/audio/translations
    # - /v1/embeddings
    # - /pooling
    # - /classify
@@ -954,6 +956,31 @@ async def do_rerank_v2(request: RerankRequest, raw_request: Request):
    return await do_rerank(request, raw_request)


+TASK_HANDLERS: dict[str, dict[str, tuple]] = {
+    "generate": {
+        "messages": (ChatCompletionRequest, create_chat_completion),
+        "default": (CompletionRequest, create_completion),
+    },
+    "embed": {
+        "messages": (EmbeddingChatRequest, create_embedding),
+        "default": (EmbeddingCompletionRequest, create_embedding),
+    },
+    "score": {
+        "default": (RerankRequest, do_rerank)
+    },
+    "rerank": {
+        "default": (RerankRequest, do_rerank)
+    },
+    "reward": {
+        "messages": (PoolingChatRequest, create_pooling),
+        "default": (PoolingCompletionRequest, create_pooling),
+    },
+    "classify": {
+        "messages": (PoolingChatRequest, create_pooling),
+        "default": (PoolingCompletionRequest, create_pooling),
+    },
+}
+
 if envs.VLLM_SERVER_DEV_MODE:
    logger.warning("SECURITY WARNING: Development endpoints are enabled! "
                   "This should NOT be used in production!")
@@ -1005,30 +1032,6 @@ if envs.VLLM_SERVER_DEV_MODE:
        return JSONResponse(content={"is_sleeping": is_sleeping})


-# TODO: RequestType = TypeForm[BaseModel] when recognized by type checkers
-# (requires typing_extensions >= 4.13)
-RequestType = Any
-GetHandlerFn = Callable[[Request], Optional[OpenAIServing]]
-EndpointFn = Callable[[RequestType, Request], Awaitable[Any]]
-
-# NOTE: Items defined earlier take higher priority
-INVOCATION_TYPES: list[tuple[RequestType, tuple[GetHandlerFn, EndpointFn]]] = [
-    (ChatCompletionRequest, (chat, create_chat_completion)),
-    (CompletionRequest, (completion, create_completion)),
-    (EmbeddingRequest, (embedding, create_embedding)),
-    (ClassificationRequest, (classify, create_classify)),
-    (ScoreRequest, (score, create_score)),
-    (RerankRequest, (rerank, do_rerank)),
-    (PoolingRequest, (pooling, create_pooling)),
-]
-
-# NOTE: Construct the TypeAdapters only once
-INVOCATION_VALIDATORS = [
-    (pydantic.TypeAdapter(request_type), (get_handler, endpoint))
-    for request_type, (get_handler, endpoint) in INVOCATION_TYPES
-]
-
-
@router.post("/invocations",
             dependencies=[Depends(validate_json_request)],
             responses={
@@ -1043,34 +1046,32 @@ INVOCATION_VALIDATORS = [
                 },
             })
 async def invocations(raw_request: Request):
-    """For SageMaker, routes requests based on the request type."""
+    """
+    For SageMaker, routes requests to other handlers based on model `task`.
+    """
    try:
        body = await raw_request.json()
    except json.JSONDecodeError as e:
        raise HTTPException(status_code=HTTPStatus.BAD_REQUEST.value,
                            detail=f"JSON decode error: {e}") from e

-    valid_endpoints = [(validator, endpoint)
-                       for validator, (get_handler,
-                                       endpoint) in INVOCATION_VALIDATORS
-                       if get_handler(raw_request) is not None]
+    task = raw_request.app.state.task

-    for request_validator, endpoint in valid_endpoints:
-        try:
-            request = request_validator.validate_python(body)
-        except pydantic.ValidationError:
-            continue
+    if task not in TASK_HANDLERS:
+        raise HTTPException(
+            status_code=400,
+            detail=f"Unsupported task: '{task}' for '/invocations'. "
+            f"Expected one of {set(TASK_HANDLERS.keys())}")

-        return await endpoint(request, raw_request)
+    handler_config = TASK_HANDLERS[task]
+    if "messages" in body:
+        request_model, handler = handler_config["messages"]
+    else:
+        request_model, handler = handler_config["default"]

-    type_names = [
-        t.__name__ if isinstance(t := validator._type, type) else str(t)
-        for validator, _ in valid_endpoints
-    ]
-    msg = ("Cannot find suitable handler for request. "
-           f"Expected one of: {type_names}")
-    res = base(raw_request).create_error_response(message=msg)
-    return JSONResponse(content=res.model_dump(), status_code=res.code)
+    # this is required since we lose the FastAPI automatic casting
+    request = request_model.model_validate(body)
+    return await handler(request, raw_request)


 if envs.VLLM_TORCH_PROFILER_DIR:
@@ -1480,28 +1481,11 @@ async def init_app_state(
                    "This discrepancy may lead to performance degradation.",
                    resolved_chat_template, args.model)

-    # Merge default_mm_loras into the static lora_modules
-    default_mm_loras = (vllm_config.lora_config.default_mm_loras
-                        if vllm_config.lora_config is not None else {})
-
-    lora_modules = args.lora_modules
-    if default_mm_loras:
-        default_mm_lora_paths = [
-            LoRAModulePath(
-                name=modality,
-                path=lora_path,
-            ) for modality, lora_path in default_mm_loras.items()
-        ]
-        if args.lora_modules is None:
-            lora_modules = default_mm_lora_paths
-        else:
-            lora_modules += default_mm_lora_paths
-
    state.openai_serving_models = OpenAIServingModels(
        engine_client=engine_client,
        model_config=model_config,
        base_model_paths=base_model_paths,
-        lora_modules=lora_modules,
+        lora_modules=args.lora_modules,
        prompt_adapters=args.prompt_adapters,
    )
    await state.openai_serving_models.init_static_loras()
@@ -11,12 +11,6 @@ from typing import Annotated, Any, ClassVar, Literal, Optional, Union
 import regex as re
 import torch
 from fastapi import HTTPException, UploadFile
-# yapf: disable
-from openai.types.chat.chat_completion_audio import (
-    ChatCompletionAudio as OpenAIChatCompletionAudio)
-from openai.types.chat.chat_completion_message import (
-    Annotation as OpenAIAnnotation)
-# yapf: enable
 from openai.types.responses import (ResponseInputParam, ResponseOutputItem,
                                    ResponseOutputMessage, ResponsePrompt,
                                    ResponseStatus, ResponseTextConfig)
@@ -1399,16 +1393,11 @@ class CompletionResponseChoice(OpenAIBaseModel):

 class CompletionResponse(OpenAIBaseModel):
    id: str = Field(default_factory=lambda: f"cmpl-{random_uuid()}")
-    object: Literal["text_completion"] = "text_completion"
+    object: str = "text_completion"
    created: int = Field(default_factory=lambda: int(time.time()))
    model: str
    choices: list[CompletionResponseChoice]
-    service_tier: Optional[Literal["auto", "default", "flex", "scale",
-                                   "priority"]] = None
-    system_fingerprint: Optional[str] = None
    usage: UsageInfo
-
-    # vLLM-specific fields that are not in OpenAI spec
    kv_transfer_params: Optional[dict[str, Any]] = Field(
        default=None, description="KVTransfer parameters.")

@@ -1560,15 +1549,9 @@ class ExtractedToolCallInformation(BaseModel):

 class ChatMessage(OpenAIBaseModel):
    role: str
-    content: Optional[str] = None
-    refusal: Optional[str] = None
-    annotations: Optional[OpenAIAnnotation] = None
-    audio: Optional[OpenAIChatCompletionAudio] = None
-    function_call: Optional[FunctionCall] = None
-    tool_calls: list[ToolCall] = Field(default_factory=list)
-
-    # vLLM-specific fields that are not in OpenAI spec
    reasoning_content: Optional[str] = None
+    content: Optional[str] = None
+    tool_calls: list[ToolCall] = Field(default_factory=list)


 class ChatCompletionLogProb(OpenAIBaseModel):
@@ -1604,12 +1587,7 @@ class ChatCompletionResponse(OpenAIBaseModel):
    created: int = Field(default_factory=lambda: int(time.time()))
    model: str
    choices: list[ChatCompletionResponseChoice]
-    service_tier: Optional[Literal["auto", "default", "flex", "scale",
-                                   "priority"]] = None
-    system_fingerprint: Optional[str] = None
    usage: UsageInfo
-
-    # vLLM-specific fields that are not in OpenAI spec
    prompt_logprobs: Optional[list[Optional[dict[int, Logprob]]]] = None
    kv_transfer_params: Optional[dict[str, Any]] = Field(
        default=None, description="KVTransfer parameters.")
@@ -153,8 +153,7 @@ class OpenAIServingChat(OpenAIServing):
            (
                lora_request,
                prompt_adapter_request,
-            ) = self._maybe_get_adapters(request,
-                                         supports_default_mm_loras=True)
+            ) = self._maybe_get_adapters(request)

            model_name = self._get_model_name(request.model, lora_request)

@@ -458,74 +458,20 @@ class OpenAIServing:
            err_type="NotFoundError",
            status_code=HTTPStatus.NOT_FOUND)

-    def _get_active_default_mm_loras(
-            self, request: AnyRequest) -> Optional[LoRARequest]:
-        """Determine if there are any active default multimodal loras."""
-        # TODO: Currently this is only enabled for chat completions
-        # to be better aligned with only being enabled for .generate
-        # when run offline. It would be nice to support additional
-        # tasks types in the future.
-        message_types = self._get_message_types(request)
-        default_mm_loras = set()
-
-        for lora in self.models.lora_requests.values():
-            # Best effort match for default multimodal lora adapters;
-            # There is probably a better way to do this, but currently
-            # this matches against the set of 'types' in any content lists
-            # up until '_', e.g., to match audio_url -> audio
-            if lora.lora_name in message_types:
-                default_mm_loras.add(lora)
-
-        # Currently only support default modality specific loras if
-        # we have exactly one lora matched on the request.
-        if len(default_mm_loras) == 1:
-            return default_mm_loras.pop()
-        return None
-
    def _maybe_get_adapters(
-        self,
-        request: AnyRequest,
-        supports_default_mm_loras: bool = False,
+        self, request: AnyRequest
    ) -> Union[tuple[None, None], tuple[LoRARequest, None], tuple[
            None, PromptAdapterRequest]]:
-
-        if request.model in self.models.lora_requests:
-            return self.models.lora_requests[request.model], None
-
-        # Currently only support default modality specific loras
-        # if we have exactly one lora matched on the request.
-        if supports_default_mm_loras:
-            default_mm_lora = self._get_active_default_mm_loras(request)
-            if default_mm_lora is not None:
-                return default_mm_lora, None
-
        if self._is_model_supported(request.model):
            return None, None
-
+        if request.model in self.models.lora_requests:
+            return self.models.lora_requests[request.model], None
        for prompt_adapter in self.models.prompt_adapter_requests:
            if request.model == prompt_adapter.prompt_adapter_name:
                return None, prompt_adapter
        # if _check_model has been called earlier, this will be unreachable
        raise ValueError(f"The model `{request.model}` does not exist.")

-    def _get_message_types(self, request: AnyRequest) -> set[str]:
-        """Retrieve the set of types from message content dicts up
-        until `_`; we use this to match potential multimodal data
-        with default per modality loras.
-        """
-        message_types: set[str] = set()
-
-        if not hasattr(request, "messages"):
-            return message_types
-
-        for message in request.messages:
-            if (isinstance(message, dict) and "content" in message
-                    and isinstance(message["content"], list)):
-                for content_dict in message["content"]:
-                    if "type" in content_dict:
-                        message_types.add(content_dict["type"].split("_")[0])
-        return message_types
-
    async def _normalize_prompt_text_to_input(
        self,
        request: AnyRequest,
@@ -216,8 +216,8 @@ class ServingScores(OpenAIServing):
                # cross_encoder models defaults to using pad_token.
                tokenized_prompts = await asyncio.gather(*(
                    tokenize_async(
-                        text=t1,  # type: ignore[arg-type]
-                        text_pair=t2,  # type: ignore[arg-type]
+                        text=t1, # type: ignore[arg-type]
+                        text_pair=t2, # type: ignore[arg-type]
                        **tokenization_kwargs) for t1, t2 in input_pairs))
            else:
                # `llm as reranker` models defaults to not using pad_token.
@@ -8,7 +8,6 @@ from .granite_tool_parser import GraniteToolParser
 from .hermes_tool_parser import Hermes2ProToolParser
 from .internlm2_tool_parser import Internlm2ToolParser
 from .jamba_tool_parser import JambaToolParser
-from .kimi_k2_tool_parser import KimiK2ToolParser
 from .llama4_pythonic_tool_parser import Llama4PythonicToolParser
 from .llama_tool_parser import Llama3JsonToolParser
 from .minimax_tool_parser import MinimaxToolParser
@@ -22,6 +21,5 @@ __all__ = [
    "GraniteToolParser", "Hermes2ProToolParser", "MistralToolParser",
    "Internlm2ToolParser", "Llama3JsonToolParser", "JambaToolParser",
    "Llama4PythonicToolParser", "PythonicToolParser", "Phi4MiniJsonToolParser",
-    "DeepSeekV3ToolParser", "xLAMToolParser", "MinimaxToolParser",
-    "KimiK2ToolParser"
+    "DeepSeekV3ToolParser", "xLAMToolParser", "MinimaxToolParser"
 ]
@@ -1,377 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-# code modified from deepseekv3_tool_parser.py
-
-from collections.abc import Sequence
-from typing import Union
-
-import regex as re
-
-from vllm.entrypoints.openai.protocol import (ChatCompletionRequest,
-                                              DeltaFunctionCall, DeltaMessage,
-                                              DeltaToolCall,
-                                              ExtractedToolCallInformation,
-                                              FunctionCall, ToolCall)
-from vllm.entrypoints.openai.tool_parsers.abstract_tool_parser import (
-    ToolParser, ToolParserManager)
-from vllm.logger import init_logger
-from vllm.transformers_utils.tokenizer import AnyTokenizer
-
-logger = init_logger(__name__)
-
-
-@ToolParserManager.register_module(["kimi_k2"])
-class KimiK2ToolParser(ToolParser):
-
-    def __init__(self, tokenizer: AnyTokenizer):
-        super().__init__(tokenizer)
-        self.current_tool_name_sent: bool = False
-        self.prev_tool_call_arr: list[dict] = []
-        self.current_tool_id: int = -1
-        self.streamed_args_for_tool: list[str] = (
-            [])  # map what has been streamed for each tool so far to a list
-
-        self.tool_calls_start_token: str = "<|tool_calls_section_begin|>"
-        self.tool_calls_end_token: str = "<|tool_calls_section_end|>"
-
-        self.tool_call_start_token: str = "<|tool_call_begin|>"
-        self.tool_call_end_token: str = "<|tool_call_end|>"
-
-        self.tool_call_regex = re.compile(
-            r"<\|tool_call_begin\|>\s*(?P<tool_call_id>[\w\.]+:\d+)\s*<\|tool_call_argument_begin\|>\s*(?P<function_arguments>.*?)\s*<\|tool_call_end\|>"
-        )
-
-        self.stream_tool_call_portion_regex = re.compile(
-            r"(?P<tool_call_id>[\w\.]+:\d+)\s*<\|tool_call_argument_begin\|>\s*(?P<function_arguments>.*)"
-        )
-
-        self.stream_tool_call_name_regex = re.compile(
-            r"(?P<tool_call_id>[\w\.]+:\d+)\s*")
-
-        if not self.model_tokenizer:
-            raise ValueError(
-                "The model tokenizer must be passed to the ToolParser "
-                "constructor during construction.")
-        self.tool_calls_start_token_id = self.vocab.get(
-            self.tool_calls_start_token)
-        self.tool_calls_end_token_id = self.vocab.get(
-            self.tool_calls_end_token)
-
-        self.tool_call_start_token_id = self.vocab.get(
-            self.tool_call_start_token)
-        self.tool_call_end_token_id = self.vocab.get(self.tool_call_end_token)
-
-        if (self.tool_calls_start_token_id is None
-                or self.tool_calls_end_token_id is None):
-            raise RuntimeError(
-                "Kimi-K2 Tool parser could not locate tool call start/end "
-                "tokens in the tokenizer!")
-
-    def extract_tool_calls(
-        self,
-        model_output: str,
-        request: ChatCompletionRequest,
-    ) -> ExtractedToolCallInformation:
-
-        # sanity check; avoid unnecessary processing
-        if self.tool_calls_start_token not in model_output:
-            return ExtractedToolCallInformation(tools_called=False,
-                                                tool_calls=[],
-                                                content=model_output)
-
-        else:
-            try:
-                # there are two possible captures - between tags, or between a
-                # tag and end-of-string so the result of
-                # findall is an array of tuples where one is a function call and
-                # the other is None
-                function_call_tuples = self.tool_call_regex.findall(
-                    model_output)
-
-                logger.debug("function_call_tuples: %s", function_call_tuples)
-
-                tool_calls = []
-                for match in function_call_tuples:
-                    function_id, function_args = match
-                    # function_id: functions.get_weather:0
-                    function_name = function_id.split('.')[1].split(':')[0]
-                    tool_calls.append(
-                        ToolCall(
-                            id=function_id,
-                            type='function',
-                            function=FunctionCall(name=function_name,
-                                                  arguments=function_args),
-                        ))
-
-                content = model_output[:model_output.
-                                       find(self.tool_calls_start_token)]
-                return ExtractedToolCallInformation(
-                    tools_called=True,
-                    tool_calls=tool_calls,
-                    content=content if content else None,
-                )
-
-            except Exception:
-                logger.exception(
-                    "Error in extracting tool call from response.")
-                return ExtractedToolCallInformation(tools_called=False,
-                                                    tool_calls=[],
-                                                    content=model_output)
-
-    def extract_tool_calls_streaming(
-        self,
-        previous_text: str,
-        current_text: str,
-        delta_text: str,
-        previous_token_ids: Sequence[int],
-        current_token_ids: Sequence[int],
-        delta_token_ids: Sequence[int],
-        request: ChatCompletionRequest,
-    ) -> Union[DeltaMessage, None]:
-
-        logger.debug("delta_text: %s", delta_text)
-        logger.debug("delta_token_ids: %s", delta_token_ids)
-        # check to see if we should be streaming a tool call - is there a
-        if self.tool_calls_start_token_id not in current_token_ids:
-            logger.debug("No tool call tokens found!")
-            return DeltaMessage(content=delta_text)
-        delta_text = delta_text.replace(self.tool_calls_start_token,
-                                        "").replace(self.tool_calls_end_token,
-                                                    "")
-        try:
-
-            # figure out where we are in the parsing by counting tool call
-            # start & end tags
-            prev_tool_start_count = previous_token_ids.count(
-                self.tool_call_start_token_id)
-            prev_tool_end_count = previous_token_ids.count(
-                self.tool_call_end_token_id)
-            cur_tool_start_count = current_token_ids.count(
-                self.tool_call_start_token_id)
-            cur_tool_end_count = current_token_ids.count(
-                self.tool_call_end_token_id)
-            tool_call_portion = None
-            text_portion = None
-
-            # case: if we're generating text, OR rounding out a tool call
-            if (cur_tool_start_count == cur_tool_end_count
-                    and prev_tool_end_count == cur_tool_end_count
-                    and self.tool_call_end_token not in delta_text):
-                logger.debug("Generating text content! skipping tool parsing.")
-                return DeltaMessage(content=delta_text)
-
-            if self.tool_call_end_token in delta_text:
-                logger.debug("tool_call_end_token in delta_text")
-                full_text = current_text + delta_text
-                tool_call_portion = full_text.split(
-                    self.tool_call_start_token)[-1].split(
-                        self.tool_call_end_token)[0].rstrip()
-                delta_text = delta_text.split(
-                    self.tool_call_end_token)[0].rstrip()
-                text_portion = delta_text.split(
-                    self.tool_call_end_token)[-1].lstrip()
-
-            # case -- we're starting a new tool call
-            if (cur_tool_start_count > cur_tool_end_count
-                    and cur_tool_start_count > prev_tool_start_count):
-                if len(delta_token_ids) > 1:
-                    tool_call_portion = current_text.split(
-                        self.tool_call_start_token)[-1]
-                else:
-                    tool_call_portion = None
-                    delta = None
-
-                text_portion = None
-
-                # set cursors and state appropriately
-                self.current_tool_id += 1
-                self.current_tool_name_sent = False
-                self.streamed_args_for_tool.append("")
-                logger.debug("Starting on a new tool %s", self.current_tool_id)
-
-            # case -- we're updating an existing tool call
-            elif (cur_tool_start_count > cur_tool_end_count
-                  and cur_tool_start_count == prev_tool_start_count):
-
-                # get the portion of the text that's the tool call
-                tool_call_portion = current_text.split(
-                    self.tool_call_start_token)[-1]
-                text_portion = None
-
-            # case -- the current tool call is being closed.
-            elif (cur_tool_start_count == cur_tool_end_count
-                  and cur_tool_end_count >= prev_tool_end_count):
-                if self.prev_tool_call_arr is None or len(
-                        self.prev_tool_call_arr) == 0:
-                    logger.debug(
-                        "attempting to close tool call, but no tool call")
-                    return None
-                diff = self.prev_tool_call_arr[self.current_tool_id].get(
-                    "arguments")
-                if diff:
-                    diff = (diff.encode("utf-8").decode("unicode_escape")
-                            if diff is str else diff)
-                    if '"}' not in delta_text:
-                        return None
-                    end_loc = delta_text.rindex('"}')
-                    diff = delta_text[:end_loc] + '"}'
-                    logger.debug(
-                        "Finishing tool and found diff that had not "
-                        "been streamed yet: %s",
-                        diff,
-                    )
-                    self.streamed_args_for_tool[self.current_tool_id] += diff
-                    return DeltaMessage(tool_calls=[
-                        DeltaToolCall(
-                            index=self.current_tool_id,
-                            function=DeltaFunctionCall(
-                                arguments=diff).model_dump(exclude_none=True),
-                        )
-                    ])
-
-            # case -- otherwise we're just generating text
-            else:
-                text = delta_text.replace(self.tool_call_start_token, "")
-                text = text.replace(self.tool_call_end_token, "")
-                delta = DeltaMessage(tool_calls=[], content=text)
-                return delta
-
-            current_tool_call = dict()
-            if tool_call_portion:
-                current_tool_call_matches = (
-                    self.stream_tool_call_portion_regex.match(
-                        tool_call_portion))
-                if current_tool_call_matches:
-                    tool_id, tool_args = (current_tool_call_matches.groups())
-                    tool_name = tool_id.split('.')[1].split(':')[0]
-                    current_tool_call['id'] = tool_id
-                    current_tool_call["name"] = tool_name
-                    current_tool_call["arguments"] = tool_args
-                else:
-                    current_tool_call_name_matches = (
-                        self.stream_tool_call_name_regex.match(
-                            tool_call_portion))
-                    if current_tool_call_name_matches:
-                        tool_id_str, = current_tool_call_name_matches.groups()
-                        tool_name = tool_id_str.split('.')[1].split(':')[0]
-                        current_tool_call['id'] = tool_id_str
-                        current_tool_call["name"] = tool_name
-                        current_tool_call["arguments"] = ""
-                    else:
-                        logger.debug("Not enough token")
-                        return None
-
-            # case - we haven't sent the tool name yet. If it's available, send
-            #   it. otherwise, wait until it's available.
-            if not self.current_tool_name_sent:
-                if current_tool_call is None:
-                    return None
-                function_name: Union[str, None] = current_tool_call.get("name")
-                tool_id = current_tool_call.get("id")
-                if function_name:
-                    self.current_tool_name_sent = True
-                    return DeltaMessage(tool_calls=[
-                        DeltaToolCall(
-                            index=self.current_tool_id,
-                            type="function",
-                            id=tool_id,
-                            function=DeltaFunctionCall(
-                                name=function_name).model_dump(
-                                    exclude_none=True),
-                        )
-                    ])
-                else:
-                    return None
-
-            # case -- otherwise, send the tool call delta
-
-            # if the tool call portion is None, send the delta as text
-            if tool_call_portion is None:
-                # if there's text but not tool calls, send that -
-                # otherwise None to skip chunk
-                delta = (DeltaMessage(
-                    content=delta_text) if text_portion is not None else None)
-                return delta
-
-            # now, the nitty-gritty of tool calls
-            # now we have the portion to parse as tool call.
-
-            logger.debug("Trying to parse current tool call with ID %s",
-                         self.current_tool_id)
-
-            # if we're starting a new tool call, push an empty object in as
-            #   a placeholder for the arguments
-            if len(self.prev_tool_call_arr) <= self.current_tool_id:
-                self.prev_tool_call_arr.append({})
-
-            # main logic for tool parsing here - compare prev. partially-parsed
-            #   JSON to the current partially-parsed JSON
-            prev_arguments = self.prev_tool_call_arr[self.current_tool_id].get(
-                "arguments")
-            cur_arguments = current_tool_call.get("arguments")
-
-            logger.debug("diffing old arguments: %s", prev_arguments)
-            logger.debug("against new ones: %s", cur_arguments)
-
-            # case -- no arguments have been created yet. skip sending a delta.
-            if not cur_arguments and not prev_arguments:
-                logger.debug("Skipping text %s - no arguments", delta_text)
-                delta = None
-
-            # case -- prev arguments are defined, but non are now.
-            #   probably impossible, but not a fatal error - just keep going
-            elif not cur_arguments and prev_arguments:
-                logger.error("should be impossible to have arguments reset "
-                             "mid-call. skipping streaming anything.")
-                delta = None
-
-            # case -- we now have the first info about arguments available from
-            #   autocompleting the JSON
-            elif cur_arguments and not prev_arguments:
-
-                delta = DeltaMessage(tool_calls=[
-                    DeltaToolCall(
-                        index=self.current_tool_id,
-                        function=DeltaFunctionCall(
-                            arguments=cur_arguments).model_dump(
-                                exclude_none=True),
-                    )
-                ])
-                self.streamed_args_for_tool[
-                    self.current_tool_id] = cur_arguments
-
-            # last case -- we have an update to existing arguments.
-            elif cur_arguments and prev_arguments:
-                if (isinstance(delta_text, str)
-                        and cur_arguments != prev_arguments
-                        and len(cur_arguments) > len(prev_arguments)
-                        and cur_arguments.startswith(prev_arguments)):
-                    delta_arguments = cur_arguments[len(prev_arguments):]
-                    logger.debug("got diff %s", delta_text)
-
-                    delta = DeltaMessage(tool_calls=[
-                        DeltaToolCall(
-                            index=self.current_tool_id,
-                            function=DeltaFunctionCall(
-                                arguments=delta_arguments).model_dump(
-                                    exclude_none=True),
-                        )
-                    ])
-                    self.streamed_args_for_tool[
-                        self.current_tool_id] = cur_arguments
-                else:
-                    delta = None
-
-            # handle saving the state for the current tool into
-            # the "prev" list for use in diffing for the next iteration
-            if self.current_tool_id == len(self.prev_tool_call_arr) - 1:
-                self.prev_tool_call_arr[
-                    self.current_tool_id] = current_tool_call
-            else:
-                self.prev_tool_call_arr.append(current_tool_call)
-
-            return delta
-
-        except Exception:
-            logger.exception("Error trying to handle streaming tool call.")
-            return None  # do not stream a delta. skip this token ID.
@@ -11,6 +11,7 @@ from vllm.entrypoints.chat_utils import (
    ChatCompletionContentPartImageParam, ChatCompletionContentPartTextParam,
    MultiModalItemTracker, _ContentPart, _parse_chat_message_content_part)
 from vllm.inputs import TokensPrompt
+from vllm.model_executor.model_loader import get_model_cls
 from vllm.model_executor.models.interfaces import supports_score_template
 from vllm.multimodal.inputs import MultiModalDataDict
 from vllm.outputs import PoolingRequestOutput
@@ -139,8 +140,6 @@ def apply_score_template(
    prompt_1: str,
    prompt_2: str,
 ) -> str:
-    # NOTE(Simon): lazy import to avoid bring in all dependencies (e.g. gguf)
-    from vllm.model_executor.model_loader import get_model_cls

    model = get_model_cls(model_config)
    if supports_score_template(model):
@@ -163,9 +162,6 @@ def post_process_tokens(
    Note:
        This is an in-place operation.
    """
-    # NOTE(Simon): lazy import to avoid bring in all dependencies (e.g. gguf)
-    from vllm.model_executor.model_loader import get_model_cls
-
    model = get_model_cls(model_config)
    if supports_score_template(model):
        model.post_process_tokens(prompt)
@@ -117,7 +117,6 @@ if TYPE_CHECKING:
    VLLM_RANDOMIZE_DP_DUMMY_INPUTS: bool = False
    VLLM_MARLIN_USE_ATOMIC_ADD: bool = False
    VLLM_V0_USE_OUTLINES_CACHE: bool = False
-    VLLM_V1_USE_OUTLINES_CACHE: bool = False
    VLLM_TPU_BUCKET_PADDING_GAP: int = 0
    VLLM_TPU_MOST_MODEL_LEN: Optional[int] = None
    VLLM_USE_DEEP_GEMM: bool = False
@@ -339,6 +338,10 @@ environment_variables: dict[str, Callable[[], Any]] = {
    "CUDA_VISIBLE_DEVICES":
    lambda: os.environ.get("CUDA_VISIBLE_DEVICES", None),

+    # used to control the visible devices in the distributed setting
+    "VLLM_VISIBLE_DEVICES":
+    lambda: os.environ.get("VLLM_VISIBLE_DEVICES", None),
+
    # timeout for each iteration in the engine
    "VLLM_ENGINE_ITERATION_TIMEOUT_S":
    lambda: int(os.environ.get("VLLM_ENGINE_ITERATION_TIMEOUT_S", "60")),
@@ -848,12 +851,6 @@ environment_variables: dict[str, Callable[[], Any]] = {
    "VLLM_V0_USE_OUTLINES_CACHE":
    lambda: os.environ.get("VLLM_V0_USE_OUTLINES_CACHE", "0") == "1",

-    # Whether to turn on the outlines cache for V1
-    # This cache is unbounded and on disk, so it's not safe to use in
-    # an environment with potentially malicious users.
-    "VLLM_V1_USE_OUTLINES_CACHE":
-    lambda: os.environ.get("VLLM_V1_USE_OUTLINES_CACHE", "0") == "1",
-
    # Gap between padding buckets for the forward pass. So we have
    # 8, we will run forward pass with [16, 24, 32, ...].
    "VLLM_TPU_BUCKET_PADDING_GAP":
@@ -959,11 +956,7 @@ environment_variables: dict[str, Callable[[], Any]] = {
    # consumer. This is only applicable when using NixlConnector in a
    # disaggregated decode-prefill setup.
    "VLLM_NIXL_ABORT_REQUEST_TIMEOUT":
-    lambda: int(os.getenv("VLLM_NIXL_ABORT_REQUEST_TIMEOUT", "120")),
-
-    # If set to 1, use the TRTLLM Decode Attention backend in flashinfer.
-    "VLLM_USE_TRTLLM_DECODE_ATTENTION":
-    lambda: os.getenv("VLLM_USE_TRTLLM_DECODE_ATTENTION", None),
+    lambda: int(os.getenv("VLLM_NIXL_ABORT_REQUEST_TIMEOUT", "120"))
 }

 # --8<-- [end:env-vars-definition]
@@ -53,12 +53,6 @@ DEFAULT_LOGGING_CONFIG = {
 }


-@lru_cache
-def _print_debug_once(logger: Logger, msg: str, *args: Hashable) -> None:
-    # Set the stacklevel to 2 to print the original caller's line info
-    logger.debug(msg, *args, stacklevel=2)
-
-
@lru_cache
 def _print_info_once(logger: Logger, msg: str, *args: Hashable) -> None:
    # Set the stacklevel to 2 to print the original caller's line info
@@ -80,13 +74,6 @@ class _VllmLogger(Logger):
        `intel_extension_for_pytorch.utils._logger`.
    """

-    def debug_once(self, msg: str, *args: Hashable) -> None:
-        """
-        As [`debug`][logging.Logger.debug], but subsequent calls with
-        the same message are silently dropped.
-        """
-        _print_debug_once(self, msg, *args)
-
    def info_once(self, msg: str, *args: Hashable) -> None:
        """
        As [`info`][logging.Logger.info], but subsequent calls with
@@ -145,7 +132,6 @@ def init_logger(name: str) -> _VllmLogger:
    logger = logging.getLogger(name)

    methods_to_patch = {
-        "debug_once": _print_debug_once,
        "info_once": _print_info_once,
        "warning_once": _print_warning_once,
    }
@@ -79,33 +79,20 @@ def maybe_backend_fallback(
                fallback_or_error(
                    guided_params,
                    "xgrammar does not support Lark grammars and the "
-                    "grammar failed to convert to GBNF.", "guidance")
+                    "grammar failed to convert to GBNF.", "outlines")

        # If the xgrammar module cannot be imported successfully,
        # we should still allow users to use guided decoding with a fallback.
        elif not xgr_installed:
            fallback_or_error(
                guided_params,
-                "xgrammar module cannot be imported successfully.", "guidance")
+                "xgrammar module cannot be imported successfully.", "outlines")

-    if guided_params.backend == "outlines":
-        if guided_params.json_object is not None:
-            # outlines doesn't support json_object, fallback to guidance
-            fallback_or_error(guided_params,
-                              "outlines does not support json_object.",
-                              "guidance")
-        elif guided_params.grammar is not None:
-            # outlines grammar support has been removed, fallback to guidance
-            # if it is a lark-based grammar and xgrammar otherwise
-            if grammar_is_likely_lark(guided_params.grammar):
-                fallback_or_error(guided_params,
-                                  "outlines no longer supports grammars.",
-                                  "guidance")
-            else:
-                # The grammar is likely already GBNF format.
-                fallback_or_error(guided_params,
-                                  "outlines no longer supports grammars.",
-                                  "xgrammar")
+    if (guided_params.backend == "outlines"
+            and guided_params.json_object is not None):
+        # outlines doesn't support json_object, fallback to guidance
+        fallback_or_error(guided_params,
+                          "outlines does not support json_object.", "guidance")

    return guided_params

@@ -124,6 +111,7 @@ async def get_guided_decoding_logits_processor(

    guided_params = maybe_backend_fallback(guided_params)

+    # CFG grammar not supported by LMFE, so we use outlines instead
    if guided_params.backend == 'outlines':
        # NOTE: lazy import outlines to avoid https://github.com/vllm-project/vllm/issues/4193
        from vllm.model_executor.guided_decoding.outlines_decoding import (  # noqa
@@ -164,6 +152,7 @@ def get_local_guided_decoding_logits_processor(
            reasoning_backend)
        reasoner = reasoner_class(tokenizer)

+    # CFG grammar not supported by LMFE, so we use outlines instead
    if guided_params.backend == 'outlines':
        # NOTE: lazy import outlines to avoid https://github.com/vllm-project/vllm/issues/4193
        from vllm.model_executor.guided_decoding.outlines_decoding import (  # noqa
@@ -12,7 +12,7 @@ from regex import escape as regex_escape
 from transformers import PreTrainedTokenizerBase

 from vllm.model_executor.guided_decoding.outlines_logits_processors import (
-    JSONLogitsProcessor, RegexLogitsProcessor)
+    CFGLogitsProcessor, JSONLogitsProcessor, RegexLogitsProcessor)
 from vllm.reasoning import ReasoningParser
 from vllm.sampling_params import GuidedDecodingParams

@@ -21,8 +21,36 @@ class GuidedDecodingMode(Enum):
    JSON = "json"
    REGEX = "regex"
    CHOICE = "choice"
+    GRAMMAR = "grammar"


+# https://github.com/outlines-dev/outlines/blob/main/outlines/grammars/json.lark
+# the main difference is that we changed the start: value to
+# start: object | array, so we are denying scalar values as the root of the
+# JSON. Starting with scalars as the root seems to cause llama to generate
+# without stop.
+JSON_GRAMMAR = r"""
+?start: object | array
+
+?value: object
+| array
+| UNESCAPED_STRING
+| SIGNED_NUMBER      -> number
+| "true"             -> true
+| "false"            -> false
+| "null"             -> null
+
+array  : "[" [value ("," value)*] "]"
+object : "{" [pair ("," pair)*] "}"
+pair   : UNESCAPED_STRING ":" value
+
+%import common.UNESCAPED_STRING
+%import common.SIGNED_NUMBER
+%import common.WS
+
+%ignore WS
+"""
+
 global_thread_pool = None  # used for generating logits processor fsm

 # It's not yet clear that using more provides a benefit, and it could
@@ -32,12 +60,16 @@ _MAX_THREADPOOL_WORKERS = 16


 async def get_outlines_guided_decoding_logits_processor(
-    guided_params: GuidedDecodingParams, tokenizer: PreTrainedTokenizerBase,
-    reasoner: Optional[ReasoningParser]
-) -> Union[JSONLogitsProcessor, RegexLogitsProcessor, None]:
+    guided_params: GuidedDecodingParams,
+    tokenizer: PreTrainedTokenizerBase,
+    reasoner: Optional[ReasoningParser],
+) -> Union[JSONLogitsProcessor, RegexLogitsProcessor, CFGLogitsProcessor,
+           None]:
    """
    Given an OpenAI-compatible request, check for guided decoding parameters
    and get the necessary logits processor for the given guide.
+    We cache logit processors by (guide, tokenizer), and on cache hit
+    we make a shallow copy to reuse the same underlying FSM.
    """
    global global_thread_pool
    guide, mode = _get_guide_and_mode(guided_params)
@@ -51,6 +83,7 @@ async def get_outlines_guided_decoding_logits_processor(
        global_thread_pool = concurrent.futures.ThreadPoolExecutor(
            max_workers=max_workers)
    loop = asyncio.get_running_loop()
+
    return await loop.run_in_executor(global_thread_pool,
                                      _get_logits_processor, guide, tokenizer,
                                      mode, guided_params.whitespace_pattern,
@@ -58,12 +91,16 @@ async def get_outlines_guided_decoding_logits_processor(


 def get_local_outlines_guided_decoding_logits_processor(
-    guided_params: GuidedDecodingParams, tokenizer: PreTrainedTokenizerBase,
-    reasoner: Optional[ReasoningParser]
-) -> Union[JSONLogitsProcessor, RegexLogitsProcessor, None]:
+    guided_params: GuidedDecodingParams,
+    tokenizer: PreTrainedTokenizerBase,
+    reasoner: Optional[ReasoningParser],
+) -> Union[JSONLogitsProcessor, RegexLogitsProcessor, CFGLogitsProcessor,
+           None]:
    """
    Given an OpenAI-compatible request, check for guided decoding parameters
    and get the necessary logits processor for the given guide.
+    We cache logit processors by (guide, tokenizer), and on cache hit
+    we make a shallow copy to reuse the same underlying FSM.
    """
    guide, mode = _get_guide_and_mode(guided_params)
    if not guide or not mode:
@@ -93,10 +130,9 @@ def _get_guide_and_mode(
        choices_regex = "(" + "|".join(choices) + ")"
        return choices_regex, GuidedDecodingMode.CHOICE
    elif guided_params.grammar:
-        raise ValueError(
-            "The `outlines` guided decoding backend no longer supports grammar "
-            "guided generation. Please use either the `xgrammar` or `guidance` "
-            "backend")
+        return guided_params.grammar, GuidedDecodingMode.GRAMMAR
+    elif guided_params.json_object:
+        return JSON_GRAMMAR, GuidedDecodingMode.GRAMMAR
    else:
        return None, None

@@ -107,11 +143,13 @@ def _get_logits_processor(
    mode: GuidedDecodingMode,
    whitespace_pattern: Union[str, None],
    reasoner: Optional[ReasoningParser],
-) -> Union[JSONLogitsProcessor, RegexLogitsProcessor]:
+) -> Union[JSONLogitsProcessor, RegexLogitsProcessor, CFGLogitsProcessor]:
    if mode == GuidedDecodingMode.JSON:
        return JSONLogitsProcessor(guide, tokenizer, whitespace_pattern,
                                   reasoner)
    elif mode == GuidedDecodingMode.REGEX or mode == GuidedDecodingMode.CHOICE:
        return RegexLogitsProcessor(guide, tokenizer, reasoner)
+    elif mode == GuidedDecodingMode.GRAMMAR:
+        return CFGLogitsProcessor(guide, tokenizer, reasoner)
    else:
        raise ValueError(f"Unknown guided decoding mode {mode}")
@@ -1,124 +1,168 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-# SPDX-FileCopyrightText: Copyright 2024-present the Outlines developers
-from __future__ import annotations

+# Copyright 2024- the Outlines developers
+# This file is adapted from
+# https://github.com/outlines-dev/outlines/blob/main/outlines/serve/vllm.py
+#
+# 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.
 import copy
-import hashlib
-import importlib.metadata
 import json
-import os
-from typing import Optional, Union
+from collections import defaultdict
+from functools import lru_cache
+from typing import Callable, Optional, Union

-import regex as re
+import numpy as np
 import torch
-from cachetools import LRUCache
-from diskcache import Cache
-from outlines_core import Guide, Index, Vocabulary
-from outlines_core.json_schema import build_regex_from_schema
-from outlines_core.kernels.torch import (_apply_token_bitmask_inplace_kernel,
-                                         allocate_token_bitmask)
+from outlines import grammars
+from outlines.caching import cache, disable_cache
+from outlines.fsm.guide import (CFGGuide, CFGState, Generate, Guide,
+                                RegexGuide, Write)
+from outlines.fsm.parsing import PartialLark
+from outlines_core.fsm.json_schema import build_regex_from_schema
 from pydantic import BaseModel
 from transformers import PreTrainedTokenizerBase
-from transformers.file_utils import SPIECE_UNDERLINE
-from transformers.models.gpt2.tokenization_gpt2 import bytes_to_unicode

 import vllm.envs as envs
 from vllm.logger import init_logger
+from vllm.platforms import current_platform
 from vllm.reasoning import ReasoningParser
-from vllm.transformers_utils.tokenizer import AnyTokenizer

 logger = init_logger(__name__)

-CACHE = None
+if envs.VLLM_V0_USE_OUTLINES_CACHE:
+    logger.warning("Enabling outlines cache. This is an unbounded on-disk "
+                   "cache. It may consume a lot of disk space and should "
+                   "not be used with untrusted clients.")
+else:
+    disable_cache()


 class BaseLogitsProcessor:

-    def __init__(self, guide: Guide, eos_token_id: int,
-                 reasoner: Optional[ReasoningParser]) -> None:
+    def __init__(self, guide: Guide, reasoner: Optional[ReasoningParser]):
        self._guide: Guide = guide
-        self._eos_token_id: int = eos_token_id
        self._reasoner: Optional[ReasoningParser] = reasoner
-        self._mask: Optional[torch.Tensor] = None
+        # CFGState is used for the FSM state for CFGGuide
+        self._fsm_state: defaultdict[int, Union[int,
+                                                CFGState]] = defaultdict(int)
+
+    def clone(self) -> "BaseLogitsProcessor":
+        cloned = copy.copy(self)
+        cloned._guide = self._guide.copy()
+        cloned._fsm_state = copy.deepcopy(self._fsm_state)
+        return cloned

    def __call__(self, input_ids: list[int],
                 scores: torch.Tensor) -> torch.Tensor:
-        if self._mask is None:
-            self._mask = allocate_token_bitmask(scores.size(-1))
+        """Use the FSM to bias the logits before sampling the next token."""

        # Skip the structured logits processing if reasoning is not finished.
        # reasoner is not None only when `--reasoning-parser` is set.
-        if self._reasoner is not None and not self._reasoner.is_reasoning_end(
-                input_ids):
-            return scores
+        if self._reasoner is not None:
+            if not self._reasoner.is_reasoning_end(input_ids):
+                return scores
+            else:
+                # Remove the reasoning tokens from the input_ids
+                # We need this because our implementation relies on the
+                # hash of the input_ids to store the FSM state.
+                input_ids = self._reasoner.extract_content_ids(input_ids)

-        # Remove the reasoning tokens from the input_ids
-        # We need this because our implementation relies on the
-        # input_ids sequence to store the FSM state.
-        input_ids = (self._reasoner.extract_content_ids(input_ids)
-                     if self._reasoner is not None else input_ids)
+        seq_id = hash(tuple(input_ids))

-        # Vllm V0 engine has a weird bug where we have to repeat
-        # the eos token id twice for generation to stop, or at least
-        # that is what we have to do from here in any case.
-        # This is a patch until a better solution can be pushed
-        # to outlines_core
-        if input_ids and input_ids[-1] != self._eos_token_id:
-            self._guide.advance(token_id=input_ids[-1], return_tokens=False)
+        if len(input_ids) > 0:
+            last_token = input_ids[-1]
+            last_seq_id = hash(tuple(input_ids[:-1]))
+            self._fsm_state[seq_id] = self._guide.get_next_state(
+                state=self._fsm_state[last_seq_id], token_id=last_token)
+        else:
+            # Note: this is a hack.
+            # Lark pickling does not work properly (silent failure),
+            # which breaks the RPC (which uses python pickleing).
+            # We need to find a better solution.
+            # On the first time this is called, we simply re-create
+            # the Lark object.
+            if isinstance(self._guide, CFGGuide):
+                self._guide.parser = PartialLark(
+                    self._guide.cfg_string,
+                    parser="lalr",
+                    import_paths=[grammars.GRAMMAR_PATH],
+                )
+                self._fsm_state[seq_id] = CFGState(
+                    parser_state=self._guide.parser.parse(""), prev_token=None)

-        self._guide.write_mask_into(
-            data_ptr=self._mask.data_ptr(),
-            numel=self._mask.numel(),
-            element_size=self._mask.element_size(),
-        )
+        instruction = self._guide.get_next_instruction(
+            state=self._fsm_state[seq_id])

-        # Any allowed tokens beyond the length of the scores will
-        # be ignored by the kernel, taking care of the issue with
-        # models such as Llama 3.2 Vision with an `<|image|>` token
-        # with id 128256, but scores.shape == torch.Size([128256])
-        _apply_token_bitmask_inplace_kernel(
-            logits=scores.unsqueeze(dim=0),
-            # mask must be on same device
-            mask=self._mask.to(scores.device, non_blocking=True))
-        self._mask.to("cpu", non_blocking=True)
+        if type(instruction) == Generate:  # noqa: E721
+            allowed_tokens = instruction.tokens
+        elif type(instruction) == Write:  # noqa: E721
+            # TODO: support fast forward tokens
+            allowed_tokens = [instruction.tokens[0]]
+        else:
+            raise TypeError(
+                f"Unsupported instruction type {type(instruction)}")

+        mask = torch.full((scores.shape[-1], ),
+                          -torch.inf,
+                          device=scores.device)
+        # The tokenizer may support more token ids than the model can generate,
+        # eg. Llama 3.2 Vision models have an `<|image|>` token with id 128256
+        # but scores.shape == torch.Size([128256])
+        # Using NumPy is faster for filtering token ids
+        allowed_tokens = np.array(allowed_tokens, dtype=np.int64)
+        allowed_tokens = torch.tensor(allowed_tokens, device=scores.device)
+        allowed_tokens = allowed_tokens.masked_select(
+            allowed_tokens < scores.shape[-1])
+        mask.index_fill_(0, allowed_tokens, 0)
+        if current_platform.is_hpu():
+            # Workaround for HPU bug where add_() raise RuntimeError:
+            # synNodeCreateWithId failed for node: strided_insert
+            # with synStatus 1 [Invalid argument], hopefully it will
+            # be fixed in the future releases of the HPU runtime.
+            scores = scores.add(mask)
+        else:
+            scores.add_(mask)
        return scores

-    def clone(self) -> BaseLogitsProcessor:
-        guide = copy.deepcopy(self._guide)
-        guide.reset()
-        return BaseLogitsProcessor(guide=guide,
-                                   eos_token_id=self._eos_token_id,
-                                   reasoner=self._reasoner)
-

 class RegexLogitsProcessor(BaseLogitsProcessor):

    @classmethod
+    @cache()
    def _get_guide(cls, regex_string: str,
                   tokenizer: PreTrainedTokenizerBase) -> Guide:
-        global CACHE
-        if CACHE is None:
-            CACHE = get_cache()
-        vocabulary = get_vocabulary(tokenizer)  # type: ignore[arg-type]
-        cache_key = f"{vocabulary._hash}_{regex_string}"
-        if CACHE is not None and cache_key in CACHE:
-            return Guide(CACHE[cache_key])
+        tokenizer = _adapt_tokenizer(tokenizer)
+        return RegexGuide.from_regex(regex_string, tokenizer)

-        index = Index(regex_string, vocabulary.inner)
+    def __init__(
+        self,
+        regex_string: str,
+        tokenizer: PreTrainedTokenizerBase,
+        reasoner: Optional[ReasoningParser],
+    ):
+        """Compile the FSM that drives the regex-structured generation.

-        if CACHE is not None:
-            CACHE[cache_key] = index
+        Parameters
+        ----------
+        regex_string
+            A string that represents a regular expression
+        tokenizer
+            The model's tokenizer

-        return Guide(index)
-
-    def __init__(self, regex_string: str, tokenizer: PreTrainedTokenizerBase,
-                 reasoner: Optional[ReasoningParser]) -> None:
+        """
        super().__init__(
-            guide=RegexLogitsProcessor._get_guide(regex_string, tokenizer),
-            eos_token_id=tokenizer.eos_token_id,  # type: ignore
-            reasoner=reasoner)
+            RegexLogitsProcessor._get_guide(regex_string, tokenizer), reasoner)


 class JSONLogitsProcessor(RegexLogitsProcessor):
@@ -126,8 +170,22 @@ class JSONLogitsProcessor(RegexLogitsProcessor):
    def __init__(self, schema: Union[str, dict, BaseModel],
                 tokenizer: PreTrainedTokenizerBase,
                 whitespace_pattern: Union[str, None],
-                 reasoner: Optional[ReasoningParser]) -> None:
+                 reasoner: Optional[ReasoningParser]):
+        """Compile the FSM that drives the JSON-guided generation.

+        Parameters
+        ----------
+        schema
+            A JSON schema that encodes the structure we want the model to
+            generate
+        tokenizer
+            The model's tokenizer
+        whitespace_pattern
+            Pattern to use for JSON syntactic whitespace (doesn't impact
+            string literals)
+            Example: allow only a single space or newline with
+            `whitespace_pattern=r"[\n ]?"`
+        """
        if isinstance(schema, type(BaseModel)):
            schema_str = json.dumps(schema.model_json_schema())
        elif isinstance(schema, dict):
@@ -139,42 +197,63 @@ class JSONLogitsProcessor(RegexLogitsProcessor):
                f"Cannot parse schema {schema}. The schema must be either "
                f"a Pydantic object, a dictionary or a string that contains "
                f"the JSON Schema specification")
-
        regex_string = build_regex_from_schema(schema_str, whitespace_pattern)
        super().__init__(regex_string, tokenizer, reasoner)


-class OutlinesVocabulary:
-    """
-    Wrapper class for `outlines_core.Vocabulary`,
-    which allows us to store a hash with the vocabulary
+class CFGLogitsProcessor(BaseLogitsProcessor):
+
+    @classmethod
+    @cache()
+    def _get_guide(cls, cfg: str, tokenizer: PreTrainedTokenizerBase) -> Guide:
+        tokenizer = _adapt_tokenizer(tokenizer)
+        return CFGGuide(cfg, tokenizer)
+
+    def __init__(self, cfg: str, tokenizer: PreTrainedTokenizerBase,
+                 reasoner: Optional[ReasoningParser]):
+        """Compile the FSM that drives the context free grammar generation.
+
+        Parameters
+        ----------
+        cfg
+            A string that represents a context-free grammar
+        tokenizer
+            The model's tokenizer
+
+        """
+        super().__init__(CFGLogitsProcessor._get_guide(cfg, tokenizer),
+                         reasoner)
+        self._guide = self._guide.copy()
+
+    def clone(self) -> "CFGLogitsProcessor":
+        cloned = copy.copy(self)
+        cloned._fsm_state = copy.deepcopy(self._fsm_state)
+        cloned._guide = self._guide.copy()
+        return cloned
+
+
+@lru_cache(maxsize=32)
+def _adapt_tokenizer(tokenizer: PreTrainedTokenizerBase):
+    """Adapt vLLM's tokenizer to use to compile the FSM.
+
+    The API of Outlines tokenizers is slightly different to that of
+    `transformers`. The decoder of outlines, returns a list whereas
+    the decode of vLLM returns an str. To sync the vLLM decoder with
+    outlines internal api, the decoder should be adapted. In addition
+    we need to handle the missing spaces to Llama's tokenizer to be
+    able to compile FSMs for this model.
+
    """
+    if getattr(tokenizer, "_outlines_adapted", False):
+        return tokenizer

-    def __init__(self, vocabulary: Vocabulary) -> None:
-        # Actual vocabulary object
-        self.inner = vocabulary
-        # Have to do abs(hash()) because python hashes can
-        # be negative, and we are using hash as a cache key.
-        hex_str = hashlib.sha256(
-            vocabulary.__repr__().encode('utf-8')).hexdigest()
-        hash_int = int(hex_str, 16)
-        self._hash = hash_int
+    tokenizer = copy.deepcopy(tokenizer)

-
-re_llama_byte_token = re.compile(r"^<0x[0-9A-F]{2}>$")
-re_replacement_seq = re.compile(r"^.{0,6}�+.{0,6}$")
-
-
-def _reduced_vocabulary(tokenizer: AnyTokenizer,
-                        eos_token_id: int) -> dict[bytes, list[int]]:
-    """Create a map from vocabulary tokens to lists of equivalent token ids.
-    
-    Returns:
-        A Dict of token string -> equivalent token ids
-    """
-    unicode_to_bytes = {v: k for k, v in bytes_to_unicode().items()}
+    tokenizer.vocabulary = tokenizer.get_vocab()
+    tokenizer.special_tokens = set(tokenizer.all_special_tokens)

    def convert_token_to_string(token: str) -> str:
+        from transformers.file_utils import SPIECE_UNDERLINE

        string = tokenizer.convert_tokens_to_string([token])

@@ -185,123 +264,21 @@ def _reduced_vocabulary(tokenizer: AnyTokenizer,

        return string

-    vocabulary: dict[bytes, list[int]] = {}
-    empty_token_ids: list[int] = []
-    for token, token_idx in tokenizer.get_vocab().items():
-        if token in tokenizer.all_special_tokens:  # type: ignore
-            continue
+    def change_decoder(
+        decoder: Callable[[list[int]],
+                          str]) -> Callable[[list[int]], list[str]]:
+        """Sync vLLM's decoder with the outlines by returning list."""

-        token_str = convert_token_to_string(token)
-        if token_str:
-            if isinstance(token, (bytes, bytearray)):
-                # For BPE tokenizers where tokens are stored as bytes.
+        def new_decoder(inp_tokens: list[int]) -> list[str]:
+            if (isinstance(inp_tokens, list) and len(inp_tokens) == 1
+                    and isinstance(inp_tokens[0], list)):
+                inp_tokens = inp_tokens[0]
+            return [decoder(inp_tokens)]

-                # safe to ignore since token_str is of type (bytearray, bytes)
-                # by this point.
-                token_bytes = bytes(token_str)  # type: ignore[arg-type]
+        return new_decoder

-            elif "\ufffd" in token_str and not re_replacement_seq.match(
-                    token_str):
-                # Handle tokens with invalid UTF-8 sequences.
-                if re_llama_byte_token.match(token):
-                    # Llama-like tokenizers use <0xXX> for incomplete sequences.
-                    token_bytes = bytes([int(token[3:5], 16)])
-                else:
-                    # GPT2 tokenizers: map each byte back using unicode_to_bytes
-                    byte_vals = [unicode_to_bytes.get(c) for c in token]
-                    if None in byte_vals:
-                        raise RuntimeError(
-                            f"Cannot convert token `{token}`"
-                            f" ({token_idx}) to bytes: {token_str}")
-                    # safe to ignore, since if None in byte_vals,
-                    # an error is thrown.
-                    token_bytes = bytes(byte_vals)  # type: ignore[arg-type]
-            else:
-                token_bytes = token_str.encode('utf-8')
+    tokenizer.convert_token_to_string = convert_token_to_string
+    tokenizer.decode = change_decoder(tokenizer.decode)
+    setattr(tokenizer, "_outlines_adapted", True)  # noqa: B010

-            if token_idx != eos_token_id:
-                vocabulary.setdefault(token_bytes, []).append(token_idx)
-        else:
-            empty_token_ids.append(token_idx)
-
-    return vocabulary
-
-
-def get_vocabulary(tokenizer: AnyTokenizer) -> Vocabulary:
-    """Get the `Vocabulary` object for a given tokenizer.
-    """
-    if hasattr(tokenizer, "_outlines_vocabulary"):
-        return tokenizer._outlines_vocabulary  # type: ignore
-
-    try:
-        if hasattr(
-                tokenizer,
-                "eos_token_id",
-        ) and tokenizer.eos_token_id is not None:
-            eos_token_id = tokenizer.eos_token_id
-        else:
-            raise ValueError(
-                f"Error during guided decoding setup: Tokenizer"
-                f" ({type(tokenizer)}) has no `eos_token_id` property, "
-                "but `eos_token_id` is required for guided decoding"
-                " to work properly.")
-
-        reduced_vocab = _reduced_vocabulary(
-            tokenizer,
-            eos_token_id  #type: ignore
-        )
-        vocabulary = OutlinesVocabulary(Vocabulary(eos_token_id,
-                                                   reduced_vocab))
-        tokenizer._outlines_vocabulary = vocabulary  # type: ignore
-
-        return vocabulary
-    except AttributeError as e:
-        raise ValueError(f"Cannot get the vocabulary of the tokenizer "
-                         f"({type(tokenizer)}). The tokenizer should have a "
-                         "get_vocab method.") from e
-
-
-def get_cache_path() -> str:
-    """Get the context object that contains previously-computed return values"""
-    outlines_cache_dir = os.getenv("OUTLINES_CACHE_DIR")
-    xdg_cache_home = os.getenv("XDG_CACHE_HOME")
-    home_dir = os.path.expanduser("~")
-
-    if outlines_cache_dir:
-        # OUTLINES_CACHE_DIR takes precedence
-        return outlines_cache_dir
-    elif xdg_cache_home:
-        return os.path.join(xdg_cache_home, ".cache", "outlines")
-    # If homedir is "/", we may be inside a container, and thus writing to
-    # root would be problematic, so we fallback to using a tempfile.
-    # Also validate the path exists, since os.path.expanduser does
-    # not garuntee existence.
-    elif os.path.isdir(home_dir) and home_dir != "/":
-        # Default Unix fallback: ~/.cache/outlines
-        return os.path.join(home_dir, ".cache", "outlines")
-    else:
-        import tempfile
-
-        # home_dir may be / inside a docker container without existing user
-        tempdir = tempfile.gettempdir()
-        return os.path.join(tempdir, ".cache", "outlines")
-
-
-def get_cache():
-    """Get the Cache instance to be used for index caching"""
-
-    cache_dir = get_cache_path()
-    if envs.VLLM_V0_USE_OUTLINES_CACHE:
-        logger.warning("Enabling outlines cache. This is an unbounded on-disk "
-                       "cache. It may consume a lot of disk space and should "
-                       "not be used with untrusted clients.")
-        cache = Cache(cache_dir, eviction_policy="none", cull_limit=0)
-        outlines_version = importlib.metadata.version("outlines_core")
-
-        cached_version = cache.get('__version__', None)
-        if cached_version != outlines_version:
-            cache.clear()
-        cache.set('__version__', outlines_version)
-        return cache
-    else:
-        return LRUCache(maxsize=128)
+    return tokenizer
@@ -7,11 +7,8 @@ import torch
 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
 from vllm.logger import init_logger
 from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
-from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
-    TopKWeightAndReduceDelegate)
 from vllm.model_executor.layers.fused_moe.utils import _resize_cache
 from vllm.triton_utils import tl, triton
-from vllm.utils.deep_gemm import fp8_m_grouped_gemm_nt_masked

 logger = init_logger(__name__)

@@ -220,10 +217,6 @@ class BatchedDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
    def supports_expert_map(self) -> bool:
        return False

-    def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
-        # Let PrepareAndFinalize::finalize() decide the impl.
-        return TopKWeightAndReduceDelegate()
-
    def workspace_shapes(
        self,
        a: torch.Tensor,
@@ -272,6 +265,7 @@ class BatchedDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
        assert expert_tokens_meta is not None
        expert_num_tokens = expert_tokens_meta.expert_num_tokens

+        import deep_gemm as dg
        assert hidden_states.ndim == 3
        assert self.block_shape is not None

@@ -289,15 +283,18 @@ class BatchedDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
        # for the M expectation of each batch, correctly setting this value
        # may lead to better performance.
        expected_m = max_num_tokens
-        fp8_m_grouped_gemm_nt_masked((a1q, a1q_scale), (w1, w1_scale),
-                                     out=workspace1,
-                                     masked_m=expert_num_tokens,
-                                     expected_m=expected_m)
+
+        dg.m_grouped_gemm_fp8_fp8_bf16_nt_masked((a1q, a1q_scale),
+                                                 (w1, w1_scale),
+                                                 out=workspace1,
+                                                 masked_m=expert_num_tokens,
+                                                 expected_m=expected_m)

        a2q, a2q_scale = silu_mul_fp8_quant_deep_gemm(workspace1,
                                                      expert_num_tokens)

-        fp8_m_grouped_gemm_nt_masked((a2q, a2q_scale), (w2, w2_scale),
-                                     out=output,
-                                     masked_m=expert_num_tokens,
-                                     expected_m=expected_m)
+        dg.m_grouped_gemm_fp8_fp8_bf16_nt_masked((a2q, a2q_scale),
+                                                 (w2, w2_scale),
+                                                 out=output,
+                                                 masked_m=expert_num_tokens,
+                                                 expected_m=expected_m)
@@ -37,6 +37,7 @@ class BatchedTritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
                block_shape=block_shape,
                per_act_token_quant=per_act_token_quant,
            ))
+        self.allow_deep_gemm = allow_deep_gemm

        self.batched_triton_experts = BatchedTritonExperts(
            max_num_tokens=max_num_tokens,
@@ -87,25 +88,6 @@ class BatchedTritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
        return ((bdge is None or bdge.supports_expert_map())
                and (bte is None or bte.supports_expert_map()))

-    def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
-        bdge = self.batched_deep_gemm_experts
-        bte = self.batched_triton_experts
-        bdge_war = bdge.finalize_weight_and_reduce_impl() if bdge else None
-        bte_war = bte.finalize_weight_and_reduce_impl() if bte else None
-        is_bdge_war = bdge_war is not None
-        is_bte_war = bte_war is not None
-
-        if is_bdge_war and is_bte_war:
-            assert bdge_war == bte_war, (
-                "Both implementations should agree on WeightAndReduce impls. "
-                f"Got bdge_war: {bdge_war}, and bte_war: {bte_war}")
-
-        if bdge_war is not None:
-            return bdge_war
-
-        assert bte_war is not None
-        return bte_war
-
    def workspace_shapes(
        self,
        a: torch.Tensor,
@@ -11,8 +11,6 @@ from vllm.logger import init_logger
 from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
 from vllm.model_executor.layers.fused_moe.prepare_finalize import (
    MoEPrepareAndFinalizeNoEP)
-from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
-    TopKWeightAndReduceDelegate)
 from vllm.model_executor.layers.fused_moe.utils import (_fp8_perm,
                                                        _fp8_quantize,
                                                        _resize_cache)
@@ -257,10 +255,6 @@ class CutlassExpertsFp8(mk.FusedMoEPermuteExpertsUnpermute):
    def supports_expert_map(self) -> bool:
        return not self.use_batched_format

-    def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
-        # Let PrepareAndFinalize::finalize() decide the impl.
-        return TopKWeightAndReduceDelegate()
-
    def workspace_shapes(
        self,
        a: torch.Tensor,
@@ -553,10 +547,8 @@ def cutlass_moe_fp4(a: torch.Tensor,
    return out.to(dtype=out_dtype)


-def _valid_cutlass_block_scaled_grouped_gemm(
-        w1: torch.Tensor, w2: torch.Tensor, inplace: bool, activation: str,
-        apply_router_weight_on_input: bool,
-        expert_map: Optional[torch.Tensor]) -> bool:
+def _valid_cutlass_block_scaled_grouped_gemm(w1: torch.Tensor,
+                                             w2: torch.Tensor) -> bool:

    def _valid_cutlass_block_scaled_grouped_gemm_shape(N: int, K: int):
        return N % 128 == 0 and K % 128 == 0
@@ -572,29 +564,6 @@ def _valid_cutlass_block_scaled_grouped_gemm(
            "CutlassBlockScaledGroupedGemm disabled: invalid weight dtype(s).")
        return False

-    if expert_map is not None:
-        logger.debug(
-            "CutlassBlockScaledGroupedGemm disabled: expert_parallel is"
-            " not supported.")
-        return False
-
-    if activation != "silu":
-        logger.debug(
-            "CutlassBlockScaledGroupedGemm disabled: only activation silu is"
-            " supported.")
-        return False
-
-    if apply_router_weight_on_input:
-        logger.debug("CutlassBlockScaledGroupedGemm disabled:"
-                     " apply_router_weight_on_input is not supported.")
-        return False
-
-    if inplace:
-        logger.debug(
-            "CutlassBlockScaledGroupedGemm disabled: inplace is not supported."
-        )
-        return False
-
    return True


@@ -12,12 +12,9 @@ from vllm.model_executor.layers.fused_moe.moe_permute_unpermute import (
    _moe_permute)
 from vllm.model_executor.layers.fused_moe.prepare_finalize import (
    MoEPrepareAndFinalizeNoEP)
-from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
-    TopKWeightAndReduceDelegate)
-from vllm.model_executor.layers.fused_moe.utils import _resize_cache
+from vllm.model_executor.layers.fused_moe.utils import (
+    _resize_cache, per_token_group_quant_fp8)
 from vllm.utils import has_deep_gemm, round_up
-from vllm.utils.deep_gemm import (m_grouped_fp8_gemm_nt_contiguous,
-                                  per_token_group_cast_to_fp8)

 logger = init_logger(__name__)

@@ -88,10 +85,6 @@ class DeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
    def supports_expert_map(self) -> bool:
        return True

-    def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
-        # Let PrepareAndFinalize::finalize() decide the impl.
-        return TopKWeightAndReduceDelegate()
-
    def workspace_shapes(
        self, a: torch.Tensor, aq: torch.Tensor, M: int, N: int, K: int,
        topk: int, global_num_experts: int, local_num_experts: int
@@ -105,7 +98,7 @@ class DeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
        M_sum = round_up(M_sum, block_m)
        workspace1 = (M_sum, max(N * 2, K))
        workspace2 = (M_sum, max(N, K))
-        output = (M, topk, K)
+        output = (M * topk, K)
        return (workspace1, workspace2, output, a.dtype)

    def apply(
@@ -128,6 +121,7 @@ class DeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
        workspace2: torch.Tensor,
        expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
    ):
+        import deep_gemm as dg
        assert self.block_shape is not None

        a1q = hidden_states
@@ -164,21 +158,21 @@ class DeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
                                  (M_sum, N // 2))
        mm2_out = _resize_cache(workspace2, (M_sum, K))

-        m_grouped_fp8_gemm_nt_contiguous((a1q, a1q_scale), (w1, w1_scale),
-                                         mm1_out, expert_ids)
+        dg.m_grouped_gemm_fp8_fp8_bf16_nt_contiguous(
+            (a1q, a1q_scale), (w1, w1_scale), mm1_out, expert_ids)

        self.activation(activation, act_out, mm1_out.view(-1, N))

        a2q_scale: Optional[torch.Tensor] = None
-        a2q, a2q_scale = per_token_group_cast_to_fp8(act_out,
-                                                     self.block_shape[1],
-                                                     column_major_scales=True,
-                                                     out_q=quant_out)
+        a2q, a2q_scale = per_token_group_quant_fp8(act_out,
+                                                   self.block_shape[1],
+                                                   column_major_scales=True,
+                                                   out_q=quant_out)

-        m_grouped_fp8_gemm_nt_contiguous((a2q, a2q_scale), (w2, w2_scale),
-                                         mm2_out, expert_ids)
+        dg.m_grouped_gemm_fp8_fp8_bf16_nt_contiguous(
+            (a2q, a2q_scale), (w2, w2_scale), mm2_out, expert_ids)

-        torch.index_select(mm2_out, 0, inv_perm, out=output.view((-1, K)))
+        torch.index_select(mm2_out, 0, inv_perm, out=output)


 def deep_gemm_moe_fp8(
@@ -6,9 +6,8 @@ import deep_ep
 import torch

 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
+from vllm import _custom_ops as ops
 from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
-from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
-    TopKWeightAndReduceContiguous, TopKWeightAndReduceDelegate)
 from vllm.model_executor.layers.fused_moe.utils import (
    moe_kernel_quantize_input)

@@ -188,25 +187,45 @@ class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
        return (expert_x, expert_x_scale, expert_tokens_meta, expert_topk_ids,
                expert_topk_weights)

+    def _apply_weights_and_reduce(self, num_tokens: int,
+                                  fused_expert_output: torch.Tensor,
+                                  topk_weights: torch.Tensor,
+                                  apply_router_weight_on_input: bool,
+                                  output_dtype: torch.dtype):
+
+        hidden_dim = fused_expert_output.size(-1)
+        if fused_expert_output.ndim == 2:
+            fused_expert_output = fused_expert_output.view(
+                num_tokens, -1, hidden_dim)
+
+        if not apply_router_weight_on_input:
+            # The DeepEP combine kernels don't do the topk weight
+            # multiplication. We multiply the weights locally.
+            m_x_topk = fused_expert_output.size(0)
+            fused_expert_output.mul_(topk_weights.view(m_x_topk, -1, 1))
+
+        out = torch.empty((num_tokens, hidden_dim),
+                          device=fused_expert_output.device,
+                          dtype=output_dtype)
+        ops.moe_sum(fused_expert_output, out)
+
+        return out
+
    def finalize(self, output: torch.Tensor, fused_expert_output: torch.Tensor,
                 topk_weights: torch.Tensor, topk_ids: torch.Tensor,
-                 apply_router_weight_on_input: bool,
-                 weight_and_reduce_impl: mk.TopKWeightAndReduce) -> None:
+                 apply_router_weight_on_input: bool) -> None:

        assert self.handle is not None

        # fused_expert_output can have 0 tokens - This happens when none of the
        # tokens from the all2all reach this EP rank.
        if fused_expert_output.numel() != 0:
-            if isinstance(weight_and_reduce_impl, TopKWeightAndReduceDelegate):
-                weight_and_reduce_impl = TopKWeightAndReduceContiguous()
-            fused_expert_output = weight_and_reduce_impl.apply(
-                output=None,
+            fused_expert_output = self._apply_weights_and_reduce(
+                num_tokens=topk_ids.size(0),
                fused_expert_output=fused_expert_output,
                topk_weights=topk_weights,
-                topk_ids=topk_ids,
                apply_router_weight_on_input=apply_router_weight_on_input,
-            )
+                output_dtype=output.dtype)

        combined_x, _, event = self.buffer.combine(
            x=fused_expert_output,
@@ -7,8 +7,6 @@ import torch

 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
 from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
-from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
-    TopKWeightAndReduceDelegate)
 from vllm.model_executor.layers.fused_moe.utils import (
    moe_kernel_quantize_input, normalize_batched_scales_shape)

@@ -168,11 +166,8 @@ class DeepEPLLPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

    def finalize(self, output: torch.Tensor, fused_expert_output: torch.Tensor,
                 topk_weights: torch.Tensor, topk_ids: torch.Tensor,
-                 apply_router_weight_on_input: bool,
-                 weight_and_reduce_impl: mk.TopKWeightAndReduce) -> None:
-        assert isinstance(
-            weight_and_reduce_impl, TopKWeightAndReduceDelegate
-        ), ("Weight application and reduction happens in the combine kernel.")
+                 apply_router_weight_on_input: bool) -> None:
+
        assert self.handle is not None

        combine_topk_weights = topk_weights
@@ -11,8 +11,6 @@ import vllm.model_executor.layers.fused_moe.modular_kernel as mk
 from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
 from vllm.model_executor.layers.fused_moe.fused_moe import (
    get_config_dtype_str, try_get_optimal_moe_config)
-from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
-    TopKWeightAndReduceDelegate, TopKWeightAndReduceNaiveBatched)
 from vllm.model_executor.layers.fused_moe.utils import (
    _resize_cache, moe_kernel_quantize_input, normalize_batched_scales_shape,
    normalize_scales_shape)
@@ -602,17 +600,25 @@ class BatchedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
        topk_weights: torch.Tensor,
        topk_ids: torch.Tensor,
        apply_router_weight_on_input: bool,
-        weight_and_reduce_impl: mk.TopKWeightAndReduce,
    ) -> None:
-        if isinstance(weight_and_reduce_impl, TopKWeightAndReduceDelegate):
-            weight_and_reduce_impl = TopKWeightAndReduceNaiveBatched(self.rank)
-        weight_and_reduce_impl.apply(
-            output=output,
-            fused_expert_output=fused_expert_output,
-            topk_weights=topk_weights,
-            topk_ids=topk_ids,
-            apply_router_weight_on_input=apply_router_weight_on_input,
-        )
+        num_tokens = topk_ids.size(0)
+        num_local_experts = fused_expert_output.size(0)
+        K = fused_expert_output.size(-1)
+        assert output.size(0) == num_tokens and output.size(1) == K
+
+        output.fill_(0)
+
+        first_expert = num_local_experts * self.rank
+        last_expert = first_expert + num_local_experts
+
+        for expert_id in range(first_expert, last_expert):
+            matching_tokens = topk_ids == expert_id
+            topks = torch.any(matching_tokens, dim=1).flatten()
+            rows = torch.count_nonzero(topks)
+            rhs = fused_expert_output[expert_id - first_expert, :rows, :]
+            if not apply_router_weight_on_input:
+                rhs.mul_(topk_weights[matching_tokens].view(rhs.size(0), 1))
+            output[topks] = output[topks] + rhs


 class NaiveBatchedExperts(mk.FusedMoEPermuteExpertsUnpermute):
@@ -664,10 +670,6 @@ class NaiveBatchedExperts(mk.FusedMoEPermuteExpertsUnpermute):
    def supports_expert_map(self) -> bool:
        return False

-    def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
-        # Let PrepareAndFinalize::finalize() decide the impl.
-        return TopKWeightAndReduceDelegate()
-
    def workspace_shapes(
        self,
        a: torch.Tensor,
@@ -875,10 +877,6 @@ class BatchedTritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
    def supports_expert_map(self) -> bool:
        return False

-    def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
-        # Let PrepareAndFinalize::finalize() decide the impl.
-        return TopKWeightAndReduceDelegate()
-
    def workspace_shapes(
        self,
        a: torch.Tensor,
@@ -25,8 +25,6 @@ from vllm.model_executor.layers.fused_moe.moe_align_block_size import (
    moe_align_block_size)
 from vllm.model_executor.layers.fused_moe.prepare_finalize import (
    MoEPrepareAndFinalizeNoEP)
-from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
-    TopKWeightAndReduceDelegate)
 from vllm.model_executor.layers.fused_moe.utils import (
    _resize_cache, moe_kernel_quantize_input)
 from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (
@@ -34,7 +32,6 @@ from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (
 from vllm.platforms import current_platform
 from vllm.triton_utils import tl, triton
 from vllm.utils import direct_register_custom_op
-from vllm.utils.deep_gemm import is_blackwell_deep_gemm_used

 from .rocm_aiter_fused_moe import is_rocm_aiter_moe_enabled

@@ -1172,15 +1169,9 @@ def fused_experts(
        allow_cutlass_block_scaled_grouped_gemm: bool = False) -> torch.Tensor:
    # For now, disable DeepGemm for small N (<= 512) until better
    # permute/unpermute ops are available.
-    # However, on B200, we use DeepGemm for all cases becuase they only support
-    # E8M0 scale, which means we requantize the weight and input to the specific
-    # scale. Fallen back to cutlass or triton for some cases would cause
-    # accuracy issue.
    N = w1.size(1)
-    should_use_deep_gemm = ((N > 512
-                             and _valid_deep_gemm(hidden_states, w1, w2))
-                            or is_blackwell_deep_gemm_used())
-    if (allow_deep_gemm and use_fp8_w8a8 and should_use_deep_gemm):
+    if (allow_deep_gemm and use_fp8_w8a8 and N > 512
+            and _valid_deep_gemm(hidden_states, w1, w2)):
        assert apply_router_weight_on_input is False
        return deep_gemm_moe_fp8(
            hidden_states=hidden_states,
@@ -1199,9 +1190,8 @@ def fused_experts(
            apply_router_weight_on_input=apply_router_weight_on_input,
        )
    elif (allow_cutlass_block_scaled_grouped_gemm and use_fp8_w8a8
-          and _valid_cutlass_block_scaled_grouped_gemm(
-              w1, w2, inplace, activation, apply_router_weight_on_input,
-              expert_map)):
+          and _valid_cutlass_block_scaled_grouped_gemm(w1, w2)):
+        assert apply_router_weight_on_input is False
        return run_cutlass_block_scaled_fused_experts(
            a=hidden_states,
            w1=w1,
@@ -1370,6 +1360,7 @@ def fused_experts_impl(

        curr_topk_ids = topk_ids[begin_chunk_idx:end_chunk_idx]
        curr_topk_weights = topk_weights[begin_chunk_idx:end_chunk_idx]
+
        qcurr_hidden_states, a1q_scale = moe_kernel_quantize_input(
            A=curr_hidden_states,
            A_scale=a1_scale,
@@ -1605,10 +1596,6 @@ class TritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
    def supports_expert_map(self) -> bool:
        return True

-    def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
-        # Let PrepareAndFinalize::finalize() decide the impl.
-        return TopKWeightAndReduceDelegate()
-
    def workspace_shapes(
        self,
        a: torch.Tensor,
--- a/Show More
+++ b/Show More