p

Signed-off-by: khluu <khluu000@gmail.com>
2026-04-02 14:51:37 -07:00
133 changed files with 1190 additions and 4564 deletions
@@ -56,7 +56,7 @@ steps:
        'cd tests &&
        pytest -v -s v1/core --ignore=v1/core/test_reset_prefix_cache_e2e.py --ignore=v1/core/test_scheduler_e2e.py &&
        pytest -v -s v1/engine --ignore=v1/engine/test_output_processor.py &&
-        pytest -v -s v1/sample --ignore=v1/sample/test_logprobs.py --ignore=v1/sample/test_logprobs_e2e.py -k "not test_topk_only and not test_topp_only and not test_topk_and_topp" &&
+        pytest -v -s v1/sample --ignore=v1/sample/test_logprobs.py --ignore=v1/sample/test_logprobs_e2e.py &&
        pytest -v -s v1/worker --ignore=v1/worker/test_gpu_model_runner.py --ignore=v1/worker/test_worker_memory_snapshot.py &&
        pytest -v -s v1/structured_output &&
        pytest -v -s v1/test_serial_utils.py &&
@@ -23,22 +23,22 @@ if [ "$failed_req" -ne 0 ]; then
  exit 1
 fi

-#echo "--- DP+TP"
-#vllm serve meta-llama/Llama-3.2-3B-Instruct -tp=2 -dp=2 --max-model-len=4096 &
-#server_pid=$!
-#timeout 600 bash -c "until curl localhost:8000/v1/models > /dev/null 2>&1; do sleep 1; done" || exit 1
-#vllm bench serve \
-#    --backend vllm \
-#    --dataset-name random \
-#    --model meta-llama/Llama-3.2-3B-Instruct \
-#    --num-prompts 20 \
-#    --result-dir ./test_results \
-#    --result-filename dp_pp.json \
-#    --save-result \
-#    --endpoint /v1/completions
-#kill -s SIGTERM $server_pid; wait $server_pid || true
-#failed_req=$(jq '.failed' ./test_results/dp_pp.json)
-#if [ "$failed_req" -ne 0 ]; then
-#  echo "Some requests were failed!"
-#  exit 1
-#fi
+echo "--- DP+TP"
+vllm serve meta-llama/Llama-3.2-3B-Instruct -tp=2 -dp=2 --max-model-len=4096 &
+server_pid=$!
+timeout 600 bash -c "until curl localhost:8000/v1/models > /dev/null 2>&1; do sleep 1; done" || exit 1
+vllm bench serve \
+    --backend vllm \
+    --dataset-name random \
+    --model meta-llama/Llama-3.2-3B-Instruct \
+    --num-prompts 20 \
+    --result-dir ./test_results \
+    --result-filename dp_pp.json \
+    --save-result \
+    --endpoint /v1/completions
+kill -s SIGTERM $server_pid; wait $server_pid || true
+failed_req=$(jq '.failed' ./test_results/dp_pp.json)
+if [ "$failed_req" -ne 0 ]; then
+  echo "Some requests were failed!"
+  exit 1
+fi
@@ -72,7 +72,6 @@ steps:
  - vllm/v1/attention/backends/flashinfer.py
  - vllm/compilation/ # TODO(luka) limit to vllm/compilation/passes
  - tests/compile/passes/test_fusion_attn.py
-  - tests/compile/passes/test_mla_attn_quant_fusion.py
  - tests/compile/passes/test_silu_mul_quant_fusion.py
  - tests/compile/passes/distributed/test_fusion_all_reduce.py
  - tests/compile/fullgraph/test_full_graph.py
@@ -80,7 +79,6 @@ steps:
    # b200 runners are limited, so we limit the tests to the minimum set only supported on Blackwell
    - nvidia-smi
    - pytest -v -s tests/compile/passes/test_fusion_attn.py -k FLASHINFER
-    - pytest -v -s tests/compile/passes/test_mla_attn_quant_fusion.py
    - pytest -v -s tests/compile/passes/test_silu_mul_quant_fusion.py
    # this runner has 2 GPUs available even though num_devices=2 is not set
    - pytest -v -s tests/compile/passes/distributed/test_fusion_all_reduce.py
@@ -18,6 +18,5 @@ steps:
  # Avoid importing model tests that cause CUDA reinitialization error
  - pytest models/test_transformers.py -v -s -m 'distributed(num_gpus=2)'
  - pytest models/language -v -s -m 'distributed(num_gpus=2)'
-  - pytest models/multimodal/generation/test_phi4siglip.py -v -s -m 'distributed(num_gpus=2)'
-  - pytest models/multimodal -v -s -m 'distributed(num_gpus=2)' --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_phi4siglip.py
+  - pytest models/multimodal -v -s -m 'distributed(num_gpus=2)' --ignore models/multimodal/generation/test_whisper.py
  - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest models/multimodal/generation/test_whisper.py -v -s -m 'distributed(num_gpus=2)'
@@ -1,264 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Benchmark: Fused FP8 output quantization in merge_attn_states
-
-Compares fused vs unfused approaches for producing FP8-quantized merged
-attention output:
-  1. Fused CUDA     -- single CUDA kernel (merge + FP8 quant)
-  2. Fused Triton   -- single Triton kernel (merge + FP8 quant)
-  3. Unfused CUDA   -- CUDA merge + torch.compiled FP8 quant
-  4. Unfused Triton  -- Triton merge + torch.compiled FP8 quant
-
-Usage:
-    python benchmarks/fused_kernels/merge_attn_states_benchmarks.py
-    python benchmarks/fused_kernels/merge_attn_states_benchmarks.py --tp 1 4 8
-    python benchmarks/fused_kernels/merge_attn_states_benchmarks.py --dtype bfloat16
-"""
-
-import argparse
-import itertools
-
-import torch
-
-from vllm._custom_ops import merge_attn_states as merge_attn_states_cuda
-from vllm.benchmarks.lib.utils import default_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.platforms import current_platform
-from vllm.triton_utils import triton
-from vllm.v1.attention.ops.triton_merge_attn_states import (
-    merge_attn_states as merge_attn_states_triton,
-)
-
-# ---------------------------------------------------------------------------
-# Configuration defaults
-# ---------------------------------------------------------------------------
-
-NUM_TOKENS_LIST = [1, 16, 64, 256, 1024, 4096]
-
-# (label, num_heads, head_size) — num_heads is for TP=1
-HEAD_CONFIGS = [
-    ("DeepSeek-V3 MLA", 128, 128),
-    ("Llama-70B", 64, 128),
-    ("Llama-8B", 32, 128),
-]
-
-TP_SIZES = [1, 2, 4, 8]
-
-INPUT_DTYPES = [torch.float32, torch.float16, torch.bfloat16]
-
-QUANTILES = [0.5, 0.2, 0.8]
-
-
-# ---------------------------------------------------------------------------
-# Helpers
-# ---------------------------------------------------------------------------
-
-
-def short_dtype(dtype: torch.dtype) -> str:
-    return str(dtype).removeprefix("torch.")
-
-
-def make_inputs(
-    num_tokens: int,
-    num_heads: int,
-    head_size: int,
-    dtype: torch.dtype,
-):
-    """Create random prefix/suffix outputs and LSEs."""
-    prefix_output = torch.randn(
-        (num_tokens, num_heads, head_size), dtype=dtype, device="cuda"
-    )
-    suffix_output = torch.randn(
-        (num_tokens, num_heads, head_size), dtype=dtype, device="cuda"
-    )
-    prefix_lse = torch.randn(num_heads, num_tokens, dtype=torch.float32, device="cuda")
-    suffix_lse = torch.randn(num_heads, num_tokens, dtype=torch.float32, device="cuda")
-    # Sprinkle some inf values to exercise edge-case paths
-    mask = torch.rand(num_heads, num_tokens, device="cuda") < 0.05
-    prefix_lse[mask] = float("inf")
-    mask2 = torch.rand(num_heads, num_tokens, device="cuda") < 0.05
-    suffix_lse[mask2] = float("inf")
-    return prefix_output, suffix_output, prefix_lse, suffix_lse
-
-
-def build_configs(head_configs, num_tokens_list, input_dtypes, tp_sizes):
-    """Build (num_tokens, num_heads, head_size, dtype_str) config tuples,
-    applying TP division to num_heads and skipping invalid combos."""
-    configs = []
-    for (_, nh, hs), nt, dtype, tp in itertools.product(
-        head_configs, num_tokens_list, input_dtypes, tp_sizes
-    ):
-        nh_tp = nh // tp
-        if nh_tp >= 1:
-            configs.append((nt, nh_tp, hs, short_dtype(dtype)))
-    return configs
-
-
-def parse_args():
-    parser = argparse.ArgumentParser(
-        description="Benchmark merge_attn_states fused FP8 quantization"
-    )
-    parser.add_argument(
-        "--num-tokens",
-        type=int,
-        nargs="+",
-        default=None,
-        help=f"Override token counts (default: {NUM_TOKENS_LIST})",
-    )
-    parser.add_argument(
-        "--tp",
-        type=int,
-        nargs="+",
-        default=None,
-        help=f"TP sizes to simulate (divides num_heads) (default: {TP_SIZES})",
-    )
-    parser.add_argument(
-        "--dtype",
-        type=str,
-        nargs="+",
-        default=None,
-        help="Input dtypes (e.g. bfloat16 float16 float32). "
-        f"Default: {[short_dtype(d) for d in INPUT_DTYPES]}",
-    )
-    return parser.parse_args()
-
-
-# ---------------------------------------------------------------------------
-# Parse args and build configs before decorators
-# ---------------------------------------------------------------------------
-
-args = parse_args()
-
-num_tokens_list = args.num_tokens if args.num_tokens else NUM_TOKENS_LIST
-tp_sizes = args.tp if args.tp else TP_SIZES
-
-if args.dtype:
-    from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
-
-    input_dtypes = [STR_DTYPE_TO_TORCH_DTYPE[d] for d in args.dtype]
-else:
-    input_dtypes = INPUT_DTYPES
-
-configs = build_configs(HEAD_CONFIGS, num_tokens_list, input_dtypes, tp_sizes)
-
-torch._dynamo.config.recompile_limit = 8888
-
-
-# ---------------------------------------------------------------------------
-# Benchmark function
-# ---------------------------------------------------------------------------
-
-
-@triton.testing.perf_report(
-    triton.testing.Benchmark(
-        x_names=["num_tokens", "num_heads", "head_size", "dtype_str"],
-        x_vals=configs,
-        line_arg="provider",
-        line_vals=["fused_cuda", "fused_triton", "unfused_cuda", "unfused_triton"],
-        line_names=["Fused CUDA", "Fused Triton", "Unfused CUDA", "Unfused Triton"],
-        styles=[("blue", "-"), ("green", "-"), ("blue", "--"), ("green", "--")],
-        ylabel="us",
-        plot_name="merge_attn_states FP8 (fused vs unfused)",
-        args={},
-    )
-)
-@default_vllm_config()
-def benchmark(num_tokens, num_heads, head_size, dtype_str, provider):
-    input_dtype = getattr(torch, dtype_str)
-    fp8_dtype = current_platform.fp8_dtype()
-    prefix_out, suffix_out, prefix_lse, suffix_lse = make_inputs(
-        num_tokens, num_heads, head_size, input_dtype
-    )
-    output_scale = torch.tensor([0.1], dtype=torch.float32, device="cuda")
-
-    if provider == "fused_cuda":
-        output = torch.empty(
-            (num_tokens, num_heads, head_size), dtype=fp8_dtype, device="cuda"
-        )
-        fn = lambda: merge_attn_states_cuda(
-            output,
-            prefix_out,
-            prefix_lse,
-            suffix_out,
-            suffix_lse,
-            output_scale=output_scale,
-        )
-    elif provider == "fused_triton":
-        output = torch.empty(
-            (num_tokens, num_heads, head_size), dtype=fp8_dtype, device="cuda"
-        )
-        fn = lambda: merge_attn_states_triton(
-            output,
-            prefix_out,
-            prefix_lse,
-            suffix_out,
-            suffix_lse,
-            output_scale=output_scale,
-        )
-    elif provider == "unfused_cuda":
-        merge_buf = torch.empty(
-            (num_tokens, num_heads, head_size), dtype=input_dtype, device="cuda"
-        )
-        quant_fp8 = QuantFP8(
-            static=True,
-            group_shape=GroupShape.PER_TENSOR,
-            column_major_scales=False,
-        )
-        quant_input = merge_buf.view(-1, head_size)
-        compiled_quant = torch.compile(
-            quant_fp8.forward_native, fullgraph=True, dynamic=False
-        )
-
-        def unfused_fn():
-            merge_attn_states_cuda(
-                merge_buf, prefix_out, prefix_lse, suffix_out, suffix_lse
-            )
-            compiled_quant(quant_input, output_scale)
-
-        fn = unfused_fn
-    else:  # unfused_triton
-        merge_buf = torch.empty(
-            (num_tokens, num_heads, head_size), dtype=input_dtype, device="cuda"
-        )
-        quant_fp8 = QuantFP8(
-            static=True,
-            group_shape=GroupShape.PER_TENSOR,
-            column_major_scales=False,
-        )
-        quant_input = merge_buf.view(-1, head_size)
-        compiled_quant = torch.compile(
-            quant_fp8.forward_native, fullgraph=True, dynamic=False
-        )
-
-        def unfused_fn():
-            merge_attn_states_triton(
-                merge_buf, prefix_out, prefix_lse, suffix_out, suffix_lse
-            )
-            compiled_quant(quant_input, output_scale)
-
-        fn = unfused_fn
-
-    ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(fn, quantiles=QUANTILES)
-    return 1000 * ms, 1000 * max_ms, 1000 * min_ms  # us
-
-
-# ---------------------------------------------------------------------------
-# Main
-# ---------------------------------------------------------------------------
-
-
-def main():
-    device_name = current_platform.get_device_name()
-    print(f"Device: {device_name}")
-    print(f"Token counts: {num_tokens_list}")
-    print(f"TP sizes: {tp_sizes}")
-    print(f"Input dtypes: {[short_dtype(d) for d in input_dtypes]}")
-    print(f"Head configs: {[(c[0], c[1], c[2]) for c in HEAD_CONFIGS]}")
-    benchmark.run(print_data=True)
-
-
-if __name__ == "__main__":
-    with torch.inference_mode():
-        main()
@@ -7,29 +7,19 @@

 #include "attention_dtypes.h"
 #include "attention_utils.cuh"
-#include "../quantization/w8a8/fp8/common.cuh"
-#include "../dispatch_utils.h"

 namespace vllm {

 // Implements section 2.2 of https://www.arxiv.org/pdf/2501.01005
 // can be used to combine partial attention results (in the split-KV case)
-template <typename scalar_t, typename output_t, const uint NUM_THREADS,
-          bool USE_FP8_OUTPUT>
+template <typename scalar_t, const uint NUM_THREADS>
 __global__ void merge_attn_states_kernel(
-    output_t* output, float* output_lse, const scalar_t* prefix_output,
+    scalar_t* output, float* output_lse, const scalar_t* prefix_output,
    const float* prefix_lse, const scalar_t* suffix_output,
    const float* suffix_lse, const uint num_tokens, const uint num_heads,
    const uint head_size, const uint prefix_head_stride,
-    const uint output_head_stride, const uint prefix_num_tokens,
-    const float* output_scale) {
-  // Inputs always load 128-bit packs (pack_size elements of scalar_t).
-  // Outputs store pack_size elements of output_t, which is smaller for FP8.
-  using input_pack_t = uint4;
-  using output_pack_t =
-      std::conditional_t<USE_FP8_OUTPUT,
-                         std::conditional_t<sizeof(scalar_t) == 4, uint, uint2>,
-                         uint4>;
+    const uint output_head_stride, const uint prefix_num_tokens) {
+  using pack_128b_t = uint4;
  const uint pack_size = 16 / sizeof(scalar_t);
  const uint threads_per_head = head_size / pack_size;

@@ -52,36 +42,15 @@ __global__ void merge_attn_states_kernel(
                               head_idx * output_head_stride;
  const scalar_t* prefix_head_ptr = prefix_output + src_head_offset;
  const scalar_t* suffix_head_ptr = suffix_output + src_head_offset;
-  output_t* output_head_ptr = output + dst_head_offset;
-
-  // Pre-invert scale: multiplication is faster than division
-  float fp8_scale_inv = 1.0f;
-  if constexpr (USE_FP8_OUTPUT) {
-    fp8_scale_inv = 1.0f / *output_scale;
-  }
+  scalar_t* output_head_ptr = output + dst_head_offset;

  // If token_idx >= prefix_num_tokens, just copy from suffix
  if (token_idx >= prefix_num_tokens) {
    if (pack_offset < head_size) {
-      input_pack_t s_out_pack = reinterpret_cast<const input_pack_t*>(
+      pack_128b_t s_out_pack = reinterpret_cast<const pack_128b_t*>(
          suffix_head_ptr)[pack_offset / pack_size];
-
-      if constexpr (USE_FP8_OUTPUT) {
-        output_t o_out_pack[pack_size];
-#pragma unroll
-        for (uint i = 0; i < pack_size; ++i) {
-          const float val =
-              vllm::to_float(reinterpret_cast<const scalar_t*>(&s_out_pack)[i]);
-          o_out_pack[i] =
-              vllm::scaled_fp8_conversion<true, output_t>(val, fp8_scale_inv);
-        }
-        reinterpret_cast<output_pack_t*>(
-            output_head_ptr)[pack_offset / pack_size] =
-            *reinterpret_cast<output_pack_t*>(o_out_pack);
-      } else {
-        reinterpret_cast<output_pack_t*>(
-            output_head_ptr)[pack_offset / pack_size] = s_out_pack;
-      }
+      reinterpret_cast<pack_128b_t*>(output_head_ptr)[pack_offset / pack_size] =
+          s_out_pack;
    }
    if (output_lse != nullptr && pack_idx == 0) {
      float s_lse = suffix_lse[head_idx * num_tokens + token_idx];
@@ -101,34 +70,20 @@ __global__ void merge_attn_states_kernel(
  /* In certain edge cases, MLA can produce p_lse = s_lse = -inf;
     continuing the pipeline then yields NaN. Root cause: with chunked prefill
     a batch may be split into two chunks; if a request in that batch has no
-     prefix hit, every LSE entry for that request's position is -inf, and at
+     prefix hit, every LSE entry for that request’s position is -inf, and at
     this moment we merge cross-attention at first. For now we simply emit
     prefix_output (expected to be all zeros) and prefix_lse (-inf) to fix
     this problem.
  */
  if (std::isinf(max_lse)) {
    if (pack_offset < head_size) {
-      input_pack_t p_out_pack = reinterpret_cast<const input_pack_t*>(
+      // Pack 128b load
+      pack_128b_t p_out_pack = reinterpret_cast<const pack_128b_t*>(
          prefix_head_ptr)[pack_offset / pack_size];

-      if constexpr (USE_FP8_OUTPUT) {
-        // Convert prefix values to FP8 (since -inf means no data,
-        // prefix_output is expected to be zeros)
-        output_t o_out_pack[pack_size];
-#pragma unroll
-        for (uint i = 0; i < pack_size; ++i) {
-          const float val =
-              vllm::to_float(reinterpret_cast<const scalar_t*>(&p_out_pack)[i]);
-          o_out_pack[i] =
-              vllm::scaled_fp8_conversion<true, output_t>(val, fp8_scale_inv);
-        }
-        reinterpret_cast<output_pack_t*>(
-            output_head_ptr)[pack_offset / pack_size] =
-            *reinterpret_cast<output_pack_t*>(o_out_pack);
-      } else {
-        reinterpret_cast<output_pack_t*>(
-            output_head_ptr)[pack_offset / pack_size] = p_out_pack;
-      }
+      // Pack 128b storage
+      reinterpret_cast<pack_128b_t*>(output_head_ptr)[pack_offset / pack_size] =
+          p_out_pack;
    }
    // We only need to write to output_lse once per head.
    if (output_lse != nullptr && pack_idx == 0) {
@@ -146,43 +101,30 @@ __global__ void merge_attn_states_kernel(
  const float s_scale = s_se / out_se;

  if (pack_offset < head_size) {
-    input_pack_t p_out_pack = reinterpret_cast<const input_pack_t*>(
+    // Pack 128b load
+    pack_128b_t p_out_pack = reinterpret_cast<const pack_128b_t*>(
        prefix_head_ptr)[pack_offset / pack_size];
-    input_pack_t s_out_pack = reinterpret_cast<const input_pack_t*>(
+    pack_128b_t s_out_pack = reinterpret_cast<const pack_128b_t*>(
        suffix_head_ptr)[pack_offset / pack_size];
+    pack_128b_t o_out_pack;

-    // Compute merged values in float32
-    float o_out_f[pack_size];
 #pragma unroll
    for (uint i = 0; i < pack_size; ++i) {
+      // Always use float for FMA to keep high precision.
+      // half(uint16_t), bfloat16, float -> float.
      const float p_out_f =
          vllm::to_float(reinterpret_cast<const scalar_t*>(&p_out_pack)[i]);
      const float s_out_f =
          vllm::to_float(reinterpret_cast<const scalar_t*>(&s_out_pack)[i]);
-      o_out_f[i] = p_out_f * p_scale + (s_out_f * s_scale);
+      // fma: a * b + c = p_out_f * p_scale + (s_out_f * s_scale)
+      const float o_out_f = p_out_f * p_scale + (s_out_f * s_scale);
+      // float -> half(uint16_t), bfloat16, float.
+      vllm::from_float(reinterpret_cast<scalar_t*>(&o_out_pack)[i], o_out_f);
    }

-    // Convert and store
-    if constexpr (USE_FP8_OUTPUT) {
-      output_t o_out_pack[pack_size];
-#pragma unroll
-      for (uint i = 0; i < pack_size; ++i) {
-        o_out_pack[i] = vllm::scaled_fp8_conversion<true, output_t>(
-            o_out_f[i], fp8_scale_inv);
-      }
-      reinterpret_cast<output_pack_t*>(
-          output_head_ptr)[pack_offset / pack_size] =
-          *reinterpret_cast<output_pack_t*>(o_out_pack);
-    } else {
-      output_pack_t o_out_pack;
-#pragma unroll
-      for (uint i = 0; i < pack_size; ++i) {
-        vllm::from_float(reinterpret_cast<scalar_t*>(&o_out_pack)[i],
-                         o_out_f[i]);
-      }
-      reinterpret_cast<output_pack_t*>(
-          output_head_ptr)[pack_offset / pack_size] = o_out_pack;
-    }
+    // Pack 128b storage
+    reinterpret_cast<pack_128b_t*>(output_head_ptr)[pack_offset / pack_size] =
+        o_out_pack;
  }
  // We only need to write to output_lse once per head.
  if (output_lse != nullptr && pack_idx == 0) {
@@ -209,26 +151,24 @@ __global__ void merge_attn_states_kernel(
    }                                                                   \
  }

-#define LAUNCH_MERGE_ATTN_STATES(scalar_t, output_t, NUM_THREADS,           \
-                                 USE_FP8_OUTPUT)                            \
+#define LAUNCH_MERGE_ATTN_STATES(scalar_t, NUM_THREADS)                     \
  {                                                                         \
-    vllm::merge_attn_states_kernel<scalar_t, output_t, NUM_THREADS,         \
-                                   USE_FP8_OUTPUT>                          \
+    vllm::merge_attn_states_kernel<scalar_t, NUM_THREADS>                   \
        <<<grid, block, 0, stream>>>(                                       \
-            reinterpret_cast<output_t*>(output.data_ptr()), output_lse_ptr, \
+            reinterpret_cast<scalar_t*>(output.data_ptr()), output_lse_ptr, \
            reinterpret_cast<scalar_t*>(prefix_output.data_ptr()),          \
            reinterpret_cast<float*>(prefix_lse.data_ptr()),                \
            reinterpret_cast<scalar_t*>(suffix_output.data_ptr()),          \
            reinterpret_cast<float*>(suffix_lse.data_ptr()), num_tokens,    \
            num_heads, head_size, prefix_head_stride, output_head_stride,   \
-            prefix_num_tokens, output_scale_ptr);                           \
+            prefix_num_tokens);                                             \
  }

 /*@brief Merges the attention states from prefix and suffix
 * into the output tensor. NUM_TOKENS: n, NUM_HEADS: h, HEAD_SIZE: d
 *
 * @param output [n,h,d] The output tensor to store the merged attention states.
- * @param output_lse [h,n] Optional tensor to store the log-sum-exp values.
+ * @param output_lse [h,d] Optional tensor to store the log-sum-exp values.
 * @param prefix_output [n,h,d] The prefix attention states.
 * @param prefix_lse [h,n] The log-sum-exp values for the prefix attention
 * states.
@@ -240,23 +180,19 @@ __global__ void merge_attn_states_kernel(
 * is computed by merging prefix_output and suffix_output. For remaining tokens
 * (prefill_tokens_with_context <= token_idx < n), output is copied directly
 * from suffix_output.
- * @param output_scale Optional scalar tensor for FP8 static quantization.
- * When provided, output must be FP8 dtype.
 */
 template <typename scalar_t>
 void merge_attn_states_launcher(
    torch::Tensor& output, std::optional<torch::Tensor> output_lse,
    const torch::Tensor& prefix_output, const torch::Tensor& prefix_lse,
    const torch::Tensor& suffix_output, const torch::Tensor& suffix_lse,
-    const std::optional<int64_t> prefill_tokens_with_context,
-    const std::optional<torch::Tensor>& output_scale) {
+    const std::optional<int64_t> prefill_tokens_with_context) {
  constexpr uint NUM_THREADS = 128;
  const uint num_tokens = output.size(0);
  const uint num_heads = output.size(1);
  const uint head_size = output.size(2);
  const uint prefix_head_stride = prefix_output.stride(1);
  const uint output_head_stride = output.stride(1);
-  // Thread mapping is based on input BF16 pack_size
  const uint pack_size = 16 / sizeof(scalar_t);
  TORCH_CHECK(head_size % pack_size == 0,
              "headsize must be multiple of pack_size:", pack_size);
@@ -272,10 +208,6 @@ void merge_attn_states_launcher(
  if (output_lse.has_value()) {
    output_lse_ptr = output_lse.value().data_ptr<float>();
  }
-  float* output_scale_ptr = nullptr;
-  if (output_scale.has_value()) {
-    output_scale_ptr = output_scale.value().data_ptr<float>();
-  }
  // Process one pack elements per thread. for float, the
  // pack_size is 4 for half/bf16, the pack_size is 8.
  const uint threads_per_head = head_size / pack_size;
@@ -287,44 +219,20 @@ void merge_attn_states_launcher(
  const c10::cuda::OptionalCUDAGuard device_guard(prefix_output.device());
  auto stream = at::cuda::getCurrentCUDAStream();

-  if (output_scale.has_value()) {
-    // FP8 output path - dispatch on output FP8 type
-    VLLM_DISPATCH_FP8_TYPES(output.scalar_type(), "merge_attn_states_fp8", [&] {
-      LAUNCH_MERGE_ATTN_STATES(scalar_t, fp8_t, NUM_THREADS, true);
-    });
-  } else {
-    // Original BF16/FP16/FP32 output path
-    LAUNCH_MERGE_ATTN_STATES(scalar_t, scalar_t, NUM_THREADS, false);
-  }
+  LAUNCH_MERGE_ATTN_STATES(scalar_t, NUM_THREADS);
 }

 #define CALL_MERGE_ATTN_STATES_LAUNCHER(scalar_t)                     \
  {                                                                   \
    merge_attn_states_launcher<scalar_t>(                             \
        output, output_lse, prefix_output, prefix_lse, suffix_output, \
-        suffix_lse, prefill_tokens_with_context, output_scale);       \
+        suffix_lse, prefill_tokens_with_context);                     \
  }

-void merge_attn_states(torch::Tensor& output,
-                       std::optional<torch::Tensor> output_lse,
-                       const torch::Tensor& prefix_output,
-                       const torch::Tensor& prefix_lse,
-                       const torch::Tensor& suffix_output,
-                       const torch::Tensor& suffix_lse,
-                       std::optional<int64_t> prefill_tokens_with_context,
-                       const std::optional<torch::Tensor>& output_scale) {
-  if (output_scale.has_value()) {
-    TORCH_CHECK(output.scalar_type() == at::ScalarType::Float8_e4m3fn ||
-                    output.scalar_type() == at::ScalarType::Float8_e4m3fnuz,
-                "output must be FP8 when output_scale is provided, got: ",
-                output.scalar_type());
-  } else {
-    TORCH_CHECK(output.scalar_type() == prefix_output.scalar_type(),
-                "output dtype (", output.scalar_type(),
-                ") must match prefix_output dtype (",
-                prefix_output.scalar_type(), ") when output_scale is not set");
-  }
-  // Always dispatch on prefix_output (input) dtype
-  DISPATCH_BY_SCALAR_DTYPE(prefix_output.dtype(),
-                           CALL_MERGE_ATTN_STATES_LAUNCHER);
+void merge_attn_states(
+    torch::Tensor& output, std::optional<torch::Tensor> output_lse,
+    const torch::Tensor& prefix_output, const torch::Tensor& prefix_lse,
+    const torch::Tensor& suffix_output, const torch::Tensor& suffix_lse,
+    std::optional<int64_t> prefill_tokens_with_context = std::nullopt) {
+  DISPATCH_BY_SCALAR_DTYPE(output.dtype(), CALL_MERGE_ATTN_STATES_LAUNCHER);
 }
@@ -10,10 +10,6 @@ void swap_blocks(torch::Tensor& src, torch::Tensor& dst,
                 int64_t block_size_in_bytes,
                 const torch::Tensor& block_mapping);

-void swap_blocks_batch(const torch::Tensor& src_ptrs,
-                       const torch::Tensor& dst_ptrs,
-                       const torch::Tensor& sizes);
-
 void reshape_and_cache(torch::Tensor& key, torch::Tensor& value,
                       torch::Tensor& key_cache, torch::Tensor& value_cache,
                       torch::Tensor& slot_mapping,
@@ -24,8 +24,6 @@
 #ifdef USE_ROCM
  #include <hip/hip_bf16.h>
 typedef __hip_bfloat16 __nv_bfloat16;
-#else
-  #include <cuda.h>
 #endif

 #if defined(__gfx942__)
@@ -75,59 +73,6 @@ void swap_blocks(torch::Tensor& src, torch::Tensor& dst,
  }
 }

-void swap_blocks_batch(const torch::Tensor& src_ptrs,
-                       const torch::Tensor& dst_ptrs,
-                       const torch::Tensor& sizes) {
-  TORCH_CHECK(src_ptrs.device().is_cpu(), "src_ptrs must be on CPU");
-  TORCH_CHECK(dst_ptrs.device().is_cpu(), "dst_ptrs must be on CPU");
-  TORCH_CHECK(sizes.device().is_cpu(), "sizes must be on CPU");
-  TORCH_CHECK(src_ptrs.dtype() == torch::kInt64, "src_ptrs must be int64");
-  TORCH_CHECK(dst_ptrs.dtype() == torch::kInt64, "dst_ptrs must be int64");
-  TORCH_CHECK(sizes.dtype() == torch::kInt64, "sizes must be int64");
-
-  const int64_t n = src_ptrs.size(0);
-  TORCH_CHECK(dst_ptrs.size(0) == n, "dst_ptrs length must match src_ptrs");
-  TORCH_CHECK(sizes.size(0) == n, "sizes length must match src_ptrs");
-
-  if (n == 0) return;
-
-  const int64_t* src_data = src_ptrs.data_ptr<int64_t>();
-  const int64_t* dst_data = dst_ptrs.data_ptr<int64_t>();
-  const int64_t* size_data = sizes.data_ptr<int64_t>();
-
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
-
-  // Use cuMemcpyBatchAsync (CUDA 12.8+) to submit all copies in a single
-  // driver call, amortizing per-copy submission overhead.
-  // int64_t and CUdeviceptr/size_t are both 8 bytes on 64-bit platforms,
-  // so we reinterpret_cast the tensor data directly to avoid copies.
-  static_assert(sizeof(CUdeviceptr) == sizeof(int64_t));
-  static_assert(sizeof(size_t) == sizeof(int64_t));
-#if !defined(USE_ROCM) && defined(CUDA_VERSION) && CUDA_VERSION >= 12080
-  CUmemcpyAttributes attr = {};
-  attr.srcAccessOrder = CU_MEMCPY_SRC_ACCESS_ORDER_STREAM;
-  size_t attrs_idx = 0;
-  size_t fail_idx = 0;
-  CUresult result = cuMemcpyBatchAsync(
-      reinterpret_cast<CUdeviceptr*>(const_cast<int64_t*>(dst_data)),
-      reinterpret_cast<CUdeviceptr*>(const_cast<int64_t*>(src_data)),
-      reinterpret_cast<size_t*>(const_cast<int64_t*>(size_data)),
-      static_cast<size_t>(n), &attr, &attrs_idx, 1, &fail_idx,
-      static_cast<CUstream>(stream));
-  TORCH_CHECK(result == CUDA_SUCCESS, "cuMemcpyBatchAsync failed at index ",
-              fail_idx, " with error ", result);
-#else
-  // Fallback for CUDA < 12.8 and ROCm: individual async copies.
-  // cudaMemcpyDefault lets the driver infer direction from pointer types.
-  for (int64_t i = 0; i < n; i++) {
-    cudaMemcpyAsync(reinterpret_cast<void*>(dst_data[i]),
-                    reinterpret_cast<void*>(src_data[i]),
-                    static_cast<size_t>(size_data[i]), cudaMemcpyDefault,
-                    stream);
-  }
-#endif
-}
-
 namespace vllm {

 // Grid: (num_layers, num_pairs)
@@ -8,6 +8,8 @@
 // libraries use different ISAs.
 #define TORCH_EXTENSION_NAME _C

+std::string init_cpu_threads_env(const std::string& cpu_ids);
+
 void release_dnnl_matmul_handler(int64_t handler);

 int64_t create_onednn_scaled_mm_handler(const torch::Tensor& b,
@@ -352,6 +354,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
      "str act, str isa) -> ()");
  ops.impl("cpu_fused_moe", torch::kCPU, &cpu_fused_moe);
 #endif
+  ops.def("init_cpu_threads_env(str cpu_ids) -> str", &init_cpu_threads_env);
  ops.def(
      "mla_decode_kvcache("
      "   Tensor! out, Tensor query, Tensor kv_cache,"
@@ -21,6 +21,150 @@ std::string init_cpu_threads_env(const std::string& cpu_ids) {

 #endif

+#ifndef VLLM_NUMA_DISABLED
+std::string init_cpu_threads_env(const std::string& cpu_ids) {
+  bitmask* omp_cpu_mask = numa_parse_cpustring_all(cpu_ids.c_str());
+  TORCH_CHECK(omp_cpu_mask != nullptr,
+              "Failed to parse CPU string: " + cpu_ids);
+  TORCH_CHECK(omp_cpu_mask->size > 0);
+  std::vector<int> omp_cpu_ids;
+  omp_cpu_ids.reserve(omp_cpu_mask->size);
+
+  constexpr int group_size = 8 * sizeof(*omp_cpu_mask->maskp);
+
+  for (int offset = 0; offset < omp_cpu_mask->size; offset += group_size) {
+    unsigned long group_mask = omp_cpu_mask->maskp[offset / group_size];
+    int i = 0;
+    while (group_mask) {
+      if (group_mask & 1) {
+        omp_cpu_ids.emplace_back(offset + i);
+      }
+      ++i;
+      group_mask >>= 1;
+    }
+  }
+
+  // Memory node binding
+  if (numa_available() != -1) {
+    std::set<int> node_ids;
+    for (const auto& cpu_id : omp_cpu_ids) {
+      int node_id = numa_node_of_cpu(cpu_id);
+      if (node_id != -1) {
+        node_ids.insert(node_id);
+      }
+    }
+    // Concatenate all node_ids into a single comma-separated string
+    if (!node_ids.empty()) {
+      std::string node_ids_str;
+      for (const int node_id : node_ids) {
+        if (!node_ids_str.empty()) {
+          node_ids_str += ",";
+        }
+        node_ids_str += std::to_string(node_id);
+      }
+
+      bitmask* mask = numa_parse_nodestring(node_ids_str.c_str());
+      bitmask* src_mask = numa_get_mems_allowed();
+
+      int pid = getpid();
+
+      if (mask && src_mask) {
+        // move all existing pages to the specified numa node.
+        *(src_mask->maskp) = *(src_mask->maskp) ^ *(mask->maskp);
+        int page_num = numa_migrate_pages(pid, src_mask, mask);
+        if (page_num == -1) {
+          TORCH_WARN("numa_migrate_pages failed. errno: " +
+                     std::to_string(errno));
+        }
+
+        // Restrict memory allocation to the selected NUMA node(s).
+        // Enhances memory locality for the threads bound to those NUMA CPUs.
+        if (node_ids.size() > 1) {
+          errno = 0;
+          numa_set_interleave_mask(mask);
+          if (errno != 0) {
+            TORCH_WARN("numa_set_interleave_mask failed. errno: " +
+                       std::to_string(errno));
+          } else {
+            TORCH_WARN(
+                "NUMA binding: Using INTERLEAVE policy for memory "
+                "allocation across multiple NUMA nodes (nodes: " +
+                node_ids_str +
+                "). Memory allocations will be "
+                "interleaved across the specified NUMA nodes.");
+          }
+        } else {
+          errno = 0;
+          numa_set_membind(mask);
+          if (errno != 0) {
+            TORCH_WARN("numa_set_membind failed. errno: " +
+                       std::to_string(errno));
+          } else {
+            TORCH_WARN(
+                "NUMA binding: Using MEMBIND policy for memory "
+                "allocation on the NUMA nodes (" +
+                node_ids_str +
+                "). Memory allocations will be "
+                "strictly bound to these NUMA nodes.");
+          }
+        }
+
+        numa_set_strict(1);
+
+        numa_free_nodemask(mask);
+        numa_free_nodemask(src_mask);
+      } else {
+        TORCH_WARN(
+            "numa_parse_nodestring or numa_get_run_node_mask failed. errno: " +
+            std::to_string(errno));
+      }
+    }
+  }
+
+  // OMP threads binding
+  omp_set_num_threads((int)omp_cpu_ids.size());
+  torch::set_num_threads((int)omp_cpu_ids.size());
+  TORCH_CHECK_EQ(omp_cpu_ids.size(), torch::get_num_threads());
+  TORCH_CHECK_EQ(omp_cpu_ids.size(), omp_get_max_threads());
+
+  std::vector<std::pair<int, int>> thread_core_mapping;
+  thread_core_mapping.reserve(omp_cpu_ids.size());
+  omp_lock_t writelock;
+  omp_init_lock(&writelock);
+
+  #pragma omp parallel for schedule(static, 1)
+  for (size_t i = 0; i < omp_cpu_ids.size(); ++i) {
+    cpu_set_t mask;
+    CPU_ZERO(&mask);
+    CPU_SET(omp_cpu_ids[i], &mask);
+    int ret = sched_setaffinity(0, sizeof(cpu_set_t), &mask);
+    if (ret == -1) {
+      TORCH_CHECK(false,
+                  "sched_setaffinity failed. errno: " + std::to_string(errno));
+    }
+
+    omp_set_lock(&writelock);
+    thread_core_mapping.emplace_back(gettid(), omp_cpu_ids[i]);
+    omp_unset_lock(&writelock);
+  }
+
+  omp_destroy_lock(&writelock);
+
+  numa_free_nodemask(omp_cpu_mask);
+
+  std::stringstream ss;
+  ss << "OMP threads binding of Process " << getpid() << ":\n";
+  std::sort(thread_core_mapping.begin(), thread_core_mapping.end(),
+            [](auto&& a, auto&& b) { return a.second < b.second; });
+  for (auto&& item : thread_core_mapping) {
+    ss << "\t"
+       << "OMP tid: " << item.first << ", core " << item.second << "\n";
+  }
+
+  return ss.str();
+}
+#endif  // VLLM_NUMA_DISABLED
+
 namespace cpu_utils {
 ScratchPadManager::ScratchPadManager() : size_(0), ptr_(nullptr) {
  this->realloc(allocation_unit * 128);
@@ -26,10 +26,8 @@ using namespace cute;
 template <class OutType, int ScaleGranularityM,
          int ScaleGranularityN, int ScaleGranularityK,
          class MmaTileShape, class ClusterShape,
-          class EpilogueScheduler, class MainloopScheduler,
-          bool swap_ab_ = false>
+          class EpilogueScheduler, class MainloopScheduler>
 struct cutlass_3x_gemm_fp8_blockwise {
-  static constexpr bool swap_ab = swap_ab_;
  using ElementAB = cutlass::float_e4m3_t;

  using ElementA = ElementAB;
@@ -57,13 +55,9 @@ struct cutlass_3x_gemm_fp8_blockwise {
  using ElementCompute = float;
  using ElementBlockScale = float; 

-  using ScaleConfig = conditional_t<swap_ab,
-      cutlass::detail::Sm120BlockwiseScaleConfig<
+  using ScaleConfig = cutlass::detail::Sm120BlockwiseScaleConfig<
        ScaleGranularityM, ScaleGranularityN, ScaleGranularityK,
-        cute::UMMA::Major::K, cute::UMMA::Major::MN>,
-      cutlass::detail::Sm120BlockwiseScaleConfig<
-        ScaleGranularityM, ScaleGranularityN, ScaleGranularityK,
-        cute::UMMA::Major::MN, cute::UMMA::Major::K>>;
+        cute::UMMA::Major::MN, cute::UMMA::Major::K>;

  // layout_SFA and layout_SFB cannot be swapped since they are deduced.
  using LayoutSFA = decltype(ScaleConfig::deduce_layoutSFA());
@@ -84,32 +78,17 @@ struct cutlass_3x_gemm_fp8_blockwise {
      ElementAccumulator,
      ElementCompute,
      ElementC,
-      conditional_t<swap_ab, LayoutC_Transpose, LayoutC>,
+      LayoutC,
      AlignmentC,
      ElementD,
-      conditional_t<swap_ab, LayoutD_Transpose, LayoutD>,
+      LayoutD,
      AlignmentD,
      EpilogueScheduler,
      DefaultOperation
  >::CollectiveOp;
 
  using StageCountType = cutlass::gemm::collective::StageCountAuto; 
-  using CollectiveMainloop = conditional_t<swap_ab,
-      typename cutlass::gemm::collective::CollectiveBuilder<
-          ArchTag,
-          OperatorClass,
-          ElementB,
-          cute::tuple<LayoutB_Transpose, LayoutSFA>,
-          AlignmentB,
-          ElementA,
-          cute::tuple<LayoutA_Transpose, LayoutSFB>,
-          AlignmentA,
-          ElementAccumulator,
-          MmaTileShape,
-          ClusterShape,
-          cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
-          MainloopScheduler
-      >::CollectiveOp,
+  using CollectiveMainloop = 
      typename cutlass::gemm::collective::CollectiveBuilder<
          ArchTag,
          OperatorClass,
@@ -124,7 +103,7 @@ struct cutlass_3x_gemm_fp8_blockwise {
          ClusterShape,
          cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
          MainloopScheduler
-      >::CollectiveOp>;
+      >::CollectiveOp;

  // SM12x family to support both SM120 (RTX 5090) and SM121 (DGX Spark)
  using KernelType = enable_sm120_family<cutlass::gemm::kernel::GemmUniversal<
@@ -136,7 +115,7 @@ struct cutlass_3x_gemm_fp8_blockwise {
 // Tile configurations for different M ranges
 template <typename OutType>
 struct sm120_blockwise_fp8_config_default {
-  // use 128x128x128 tile with Cooperative (Auto) schedule
+  // M > 256: use 128x128x128 tile with Cooperative (Auto) schedule
  using KernelSchedule = cutlass::gemm::collective::KernelScheduleAuto;
  using EpilogueSchedule = cutlass::epilogue::collective::EpilogueScheduleAuto;
  using TileShape = Shape<_128, _128, _128>;
@@ -148,8 +127,8 @@ struct sm120_blockwise_fp8_config_default {
 };

 template <typename OutType>
-struct sm120_blockwise_fp8_config_pingpong {
-  // use 64x128x128 tile with Pingpong schedule
+struct sm120_blockwise_fp8_config_M64 {
+  // M in [1, 256]: use 64x128x128 tile with Pingpong schedule
  using KernelSchedule = cutlass::gemm::KernelTmaWarpSpecializedBlockwisePingpongSm120;
  using EpilogueSchedule = cutlass::epilogue::collective::EpilogueScheduleAuto;
  using TileShape = Shape<_64, _128, _128>;
@@ -160,24 +139,11 @@ struct sm120_blockwise_fp8_config_pingpong {
      EpilogueSchedule, KernelSchedule>;
 };

-template <typename OutType>
-struct sm120_blockwise_fp8_config_swapab {
-  // use 128x32x128 tile with Cooperative schedule
-  using KernelSchedule = cutlass::gemm::KernelTmaWarpSpecializedBlockwiseCooperativeSm120;
-  using EpilogueSchedule = cutlass::epilogue::collective::EpilogueScheduleAuto;
-  using TileShape = Shape<_128, _32, _128>;
-  using ClusterShape = Shape<_1, _1, _1>;
-  using Gemm = cutlass_3x_gemm_fp8_blockwise<
-      OutType, 128, 1, 128, TileShape, ClusterShape,
-      EpilogueSchedule, KernelSchedule, true>;
-};
-
 template <typename Gemm>
 void cutlass_gemm_caller_blockwise(torch::stable::Tensor& out, torch::stable::Tensor const& a,
                                   torch::stable::Tensor const& b,
                                   torch::stable::Tensor const& a_scales,
                                   torch::stable::Tensor const& b_scales) {
-  static constexpr bool swap_ab = Gemm::swap_ab;
  using GemmKernel = typename Gemm::GemmKernel;
  using StrideA = typename Gemm::GemmKernel::StrideA;
  using StrideB = typename Gemm::GemmKernel::StrideB;
@@ -201,13 +167,11 @@ void cutlass_gemm_caller_blockwise(torch::stable::Tensor& out, torch::stable::Te
  b_stride =
      cutlass::make_cute_packed_stride(StrideB{}, cute::make_shape(n, k, 1));
  c_stride =
-      cutlass::make_cute_packed_stride(StrideC{}, swap_ab ? cute::make_shape(n, m, 1) : cute::make_shape(m, n, 1));
+      cutlass::make_cute_packed_stride(StrideC{}, cute::make_shape(m, n, 1));

-  LayoutSFA layout_SFA = swap_ab ?
-      ScaleConfig::tile_atom_to_shape_SFA(make_shape(n, m, k, 1)) :
+  LayoutSFA layout_SFA = 
      ScaleConfig::tile_atom_to_shape_SFA(make_shape(m, n, k, 1));
-  LayoutSFB layout_SFB = swap_ab ?
-      ScaleConfig::tile_atom_to_shape_SFB(make_shape(n, m, k, 1)) :
+  LayoutSFB layout_SFB = 
      ScaleConfig::tile_atom_to_shape_SFB(make_shape(m, n, k, 1));

  auto a_ptr = static_cast<ElementAB const*>(a.data_ptr());
@@ -216,24 +180,15 @@ void cutlass_gemm_caller_blockwise(torch::stable::Tensor& out, torch::stable::Te
  auto b_scales_ptr = static_cast<ElementBlockScale const*>(b_scales.data_ptr());

  typename GemmKernel::MainloopArguments mainloop_args{};
+  mainloop_args.ptr_A = a_ptr;
+  mainloop_args.dA = a_stride;
+  mainloop_args.ptr_B = b_ptr;
+  mainloop_args.dB = b_stride;
+  mainloop_args.ptr_SFA = a_scales_ptr;
  mainloop_args.layout_SFA = layout_SFA;
+  mainloop_args.ptr_SFB = b_scales_ptr;
  mainloop_args.layout_SFB = layout_SFB;
-  if (swap_ab) {
-    mainloop_args.ptr_A = b_ptr;
-    mainloop_args.dA = b_stride;
-    mainloop_args.ptr_B = a_ptr;
-    mainloop_args.dB = a_stride;
-    mainloop_args.ptr_SFA = b_scales_ptr;
-    mainloop_args.ptr_SFB = a_scales_ptr;
-  } else {
-    mainloop_args.ptr_A = a_ptr;
-    mainloop_args.dA = a_stride;
-    mainloop_args.ptr_B = b_ptr;
-    mainloop_args.dB = b_stride;
-    mainloop_args.ptr_SFA = a_scales_ptr;
-    mainloop_args.ptr_SFB = b_scales_ptr;
-  }
-  auto prob_shape = swap_ab ? cute::make_shape(n, m, k, 1) : cute::make_shape(m, n, k, 1);
+  auto prob_shape = cute::make_shape(m, n, k, 1);

  auto c_ptr = static_cast<ElementD*>(out.data_ptr());
  typename GemmKernel::EpilogueArguments epilogue_args{
@@ -249,26 +204,15 @@ void cutlass_gemm_blockwise_sm120_fp8_dispatch(torch::stable::Tensor& out,
                                               torch::stable::Tensor const& a_scales,
                                               torch::stable::Tensor const& b_scales) {
  int M = a.size(0);
-  // more heuristic tuning can be done here by checking N/K dimensions as well
-  bool swap_ab = (M <= 64) || (M % 4 != 0);
-
-  if (!swap_ab) {
-    if (M <= 256) {
-      using Gemm = typename sm120_blockwise_fp8_config_pingpong<OutType>::Gemm;
-      return cutlass_gemm_caller_blockwise<Gemm>(
-          out, a, b, a_scales, b_scales);
-    }
-    // M > 256: use default 128x128x128 config with Cooperative (Auto) schedule
-    using Gemm = typename sm120_blockwise_fp8_config_default<OutType>::Gemm;
-    return cutlass_gemm_caller_blockwise<Gemm>(
-        out, a, b, a_scales, b_scales);
-  } else {
-    // Swap A/B for small M to improve performance
-    // Use TILE_N=32 as the minimum compatible tile size.
-    using Gemm = typename sm120_blockwise_fp8_config_swapab<OutType>::Gemm;
+  if (M <= 256) {
+    using Gemm = typename sm120_blockwise_fp8_config_M64<OutType>::Gemm;
    return cutlass_gemm_caller_blockwise<Gemm>(
        out, a, b, a_scales, b_scales);
  }
+  // M > 256: use default 128x128x128 config with Cooperative (Auto) schedule
+  using Gemm = typename sm120_blockwise_fp8_config_default<OutType>::Gemm;
+  return cutlass_gemm_caller_blockwise<Gemm>(
+      out, a, b, a_scales, b_scales);
 }

 }  // namespace vllm
@@ -57,8 +57,7 @@ void merge_attn_states(
    torch::Tensor& output, std::optional<torch::Tensor> output_lse,
    const torch::Tensor& prefix_output, const torch::Tensor& prefix_lse,
    const torch::Tensor& suffix_output, const torch::Tensor& suffix_lse,
-    const std::optional<int64_t> prefill_tokens_with_context,
-    const std::optional<torch::Tensor>& output_scale = std::nullopt);
+    const std::optional<int64_t> prefill_tokens_with_context);
 #ifndef USE_ROCM
 void convert_vertical_slash_indexes(
    torch::Tensor& block_count,      // [BATCH, N_HEADS, NUM_ROWS]
@@ -73,8 +73,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
      "    Tensor prefix_lse,"
      "    Tensor suffix_output,"
      "    Tensor suffix_lse,"
-      "    int!? prefill_tokens_with_context,"
-      "    Tensor? output_scale=None) -> ()");
+      "    int!? prefill_tokens_with_context) -> ()");
  ops.impl("merge_attn_states", torch::kCUDA, &merge_attn_states);
 #ifndef USE_ROCM
  ops.def(
@@ -508,12 +507,6 @@ TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _cache_ops), cache_ops) {
      "            int block_size_in_bytes, Tensor block_mapping) -> ()");
  cache_ops.impl("swap_blocks", torch::kCUDA, &swap_blocks);

-  // Batch swap: submit all block copies in a single driver call.
-  cache_ops.def(
-      "swap_blocks_batch(Tensor src_ptrs, Tensor dst_ptrs,"
-      "                  Tensor sizes) -> ()");
-  cache_ops.impl("swap_blocks_batch", torch::kCPU, &swap_blocks_batch);
-
  // Reshape the key and value tensors and cache them.
  cache_ops.def(
      "reshape_and_cache(Tensor key, Tensor value,"
@@ -689,6 +689,10 @@ RUN --mount=type=cache,target=/root/.cache/uv \
 . /etc/environment && \
 uv pip list

+# Pin transformers to 5.5.0, overwriting the version from dependencies
+RUN --mount=type=cache,target=/root/.cache/uv \
+    uv pip install --system "transformers==5.5.0"
+
 # Install deepgemm wheel that has been built in the `build` stage
 RUN --mount=type=cache,target=/root/.cache/uv \
    --mount=type=bind,from=build,source=/tmp/deepgemm/dist,target=/tmp/deepgemm/dist,ro \
@@ -203,8 +203,7 @@ WORKDIR /vllm-workspace
 RUN --mount=type=cache,target=/root/.cache/uv \
    --mount=type=cache,target=/root/.cache/ccache \
    --mount=type=bind,from=vllm-build,src=/vllm-workspace/dist,target=dist \
-    uv pip install dist/*.whl && \
-    uv pip install "vllm[audio]"
+    uv pip install dist/*.whl

 # Add labels to document build configuration
 LABEL org.opencontainers.image.title="vLLM CPU"
@@ -390,20 +390,7 @@ ENV MIOPEN_DEBUG_CONV_GEMM=0
 RUN mkdir src && mv vllm src/vllm

 # This is a workaround to ensure pytest exits with the correct status code in CI tests.
-RUN cat << 'EOF' > /vllm-workspace/conftest.py
-import os
-
-_exit_code = 1
-
-def pytest_sessionfinish(session, exitstatus):
-    global _exit_code
-    _exit_code = int(exitstatus)
-
-def pytest_unconfigure(config):
-    sys.stdout.flush()
-    sys.stderr.flush()
-    os._exit(_exit_code)
-EOF
+RUN echo "import os\n\ndef pytest_sessionfinish(session, exitstatus):\n    os._exit(int(exitstatus))" > /vllm-workspace/conftest.py

 # -----------------------
 # Final vLLM image
@@ -22,7 +22,6 @@ or just on the low or high end.
 | ------------------------------------------------------------------------------ | ---------------------------- | ---------------------------------------------- | ------------------------------ | ------------------ | --------- | ------------ |
 | [AllReduce + RMSNorm](#allreduce--rmsnorm-fuse_allreduce_rms)                  | `fuse_allreduce_rms`         | All-reduce → RMSNorm (+residual_add) (→ quant) | O2 (Hopper/Blackwell + TP > 1) | 5-20%              | No        | Low          |
 | [Attention + Quant](#attention--quantization-fuse_attn_quant)                  | `fuse_attn_quant`            | Attention output → FP8/NVFP4 quant             | Off by default                 | 3-7%               | Yes       | Always       |
-| [MLA Attention + Quant](#attention--quantization-fuse_attn_quant)              | `fuse_attn_quant`            | MLA Attention output → FP8/NVFP4 quant         | Off by default                 | TBD                | Yes       | Always       |
 | [RoPE + KV-Cache Update](#rope--kv-cache-update-fuse_rope_kvcache)             | `fuse_rope_kvcache`          | Rotary embedding → KV cache write              | O2 (ROCm/AITER only)           | 2-4%               | No        | Low          |
 | [QK Norm + RoPE](#qk-norm--rope-enable_qk_norm_rope_fusion)                    | `enable_qk_norm_rope_fusion` | Q/K RMSNorm → rotary embedding                 | Off by default                 | 2-3%               | No        | Low          |
 | [Sequence Parallelism](#sequence-parallelism-enable_sp)                        | `enable_sp`                  | AllReduce → ReduceScatter + AllGather          | Off by default                 | Prereq for AsyncTP | Yes       | High         |
@@ -41,7 +40,6 @@ The table below lists the quantization schemes supported by each fusion on each
 | ---------------------------- | ---------------------------------------- | ---------------------------------------- | ---------------------------------------- | ------------- | ---------------------------------------- |
 | `fuse_allreduce_rms`         | FP16/BF16, FP8 static, NVFP4             | FP16/BF16, FP8 static                    | —                                        | —             | —                                        |
 | `fuse_attn_quant`\*          | FP8 static\*, NVFP4\*                    | FP8 static\*                             | FP8 static\*                             | —             | FP8 static\*                             |
-| `fuse_attn_quant` (MLA)\*    | FP8 static\*, NVFP4\*                    | FP8 static\*                             | FP8 static\*                             | —             | FP8 static(untested)\*                   |
 | `fuse_rope_kvcache`          | —                                        | —                                        | —                                        | —             | FP16/BF16                                |
 | `enable_qk_norm_rope_fusion` | FP16/BF16                                | FP16/BF16                                | FP16/BF16†                               | FP16/BF16†    | —                                        |
 | `enable_sp`                  | FP16/BF16, FP8 static†                   | FP16/BF16, FP8 static                    | FP16/BF16†                               | FP16/BF16†    | —                                        |
@@ -131,8 +129,7 @@ on SM90/SM100) and configurable via `PassConfig.fi_allreduce_fusion_max_size_mb`
    explicitly. It requires the full model graph to be visible (Inductor partition or `splitting_ops=[]`).

 **What it fuses.** Fuses the attention output quantization directly after the attention computation,
-eliminating a full-precision memory round-trip of the attention output. This fusion supports both
-standard `Attention` and `MLAAttention` (used by DeepSeek-V2/V3/R1 models). Patterns covered:
+eliminating a full-precision memory round-trip of the attention output. Patterns covered:

 `Attention → FP8 static quant`:

@@ -145,24 +142,11 @@ standard `Attention` and `MLAAttention` (used by DeepSeek-V2/V3/R1 models). Patt

 - `FLASHINFER`: CUDA sm100+ with FlashInfer installed

-`MLAAttention → FP8 static quant` / `MLAAttention → NVFP4 dynamic quant`:
-
-The MLA fusion operates at the graph level on the `unified_mla_attention_with_output` op and works
-with all MLA decode and prefill backend combinations. Unlike standard `Attention` backends (where
-the kernel writes FP8 output directly), no MLA prefill or decode backend currently supports direct
-FP8/FP4 output. The fusion writes to an intermediate buffer and quantizes in a separate step, so
-there is no memory round-trip elimination yet.
-
-!!! info
-    The MLA attention fusion is not expected to yield a measurable speedup yet.
-    This will improve once MLA prefill/decode kernels support direct FP8/FP4 output.
-
 Other attention backends do not support fused output quantization yet.

 **Code locations.**

- Pass (Attention): [`vllm/compilation/passes/fusion/attn_quant_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/attn_quant_fusion.py)
- Pass (MLAAttention): [`vllm/compilation/passes/fusion/mla_attn_quant_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/mla_attn_quant_fusion.py)
+- Pass: [`vllm/compilation/passes/fusion/attn_quant_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/attn_quant_fusion.py)
 - Attention backends: [`vllm/v1/attention/backends/`](https://github.com/vllm-project/vllm/blob/main/vllm/v1/attention/backends/)

 ### RoPE + KV-Cache Update (`fuse_rope_kvcache`)
@@ -481,7 +481,6 @@ th {
 | `Step3p5ForCausalLM` | Step-3.5-flash | `stepfun-ai/Step-3.5-Flash`, etc. | | ✅︎ |
 | `TeleChatForCausalLM` | TeleChat | `chuhac/TeleChat2-35B`, etc. | ✅︎ | ✅︎ |
 | `TeleChat2ForCausalLM` | TeleChat2 | `Tele-AI/TeleChat2-3B`, `Tele-AI/TeleChat2-7B`, `Tele-AI/TeleChat2-35B`, etc. | ✅︎ | ✅︎ |
-| `TeleChat3ForCausalLM` | TeleChat3 | `Tele-AI/TeleChat3-36B-Thinking`, `Tele-AI/TeleChat3-Coder-36B-Thinking`, etc. | ✅︎ | ✅︎ |
 | `TeleFLMForCausalLM` | TeleFLM | `CofeAI/FLM-2-52B-Instruct-2407`, `CofeAI/Tele-FLM`, etc. | ✅︎ | ✅︎ |
 | `XverseForCausalLM` | XVERSE | `xverse/XVERSE-7B-Chat`, `xverse/XVERSE-13B-Chat`, `xverse/XVERSE-65B-Chat`, etc. | ✅︎ | ✅︎ |
 | `MiniMaxM1ForCausalLM` | MiniMax-Text | `MiniMaxAI/MiniMax-M1-40k`, `MiniMaxAI/MiniMax-M1-80k`, etc. | | |
@@ -15,4 +15,4 @@ torch==2.10.0+xpu
 torchaudio
 torchvision

-vllm_xpu_kernels @ https://github.com/vllm-project/vllm-xpu-kernels/releases/download/v0.1.5/vllm_xpu_kernels-0.1.5-cp38-abi3-manylinux_2_28_x86_64.whl
+vllm_xpu_kernels @ https://github.com/vllm-project/vllm-xpu-kernels/releases/download/v0.1.4/vllm_xpu_kernels-0.1.4-cp38-abi3-manylinux_2_28_x86_64.whl
@@ -170,3 +170,14 @@ class TestFullCUDAGraph:
                piecewise_res.outputs[0].text.lower()
                == full_res.outputs[0].text.lower()
            )
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="Skip if not cuda")
+def test_full_cudagraph_with_invalid_backend():
+    # Flex_Attention is not supported with full cuda graph
+    with pytest.raises(RuntimeError):
+        LLM(
+            model="Qwen/Qwen2-1.5B-Instruct",
+            compilation_config=CompilationConfig(cudagraph_mode="FULL"),
+            attention_config={"backend": "FLEX_ATTENTION"},
+        )
@@ -84,14 +84,10 @@ def run_e2e_fusion_test(monkeypatch, caplog_mp_spawn):
        rocm_aiter_ops.refresh_env_variables()

        # Filter here to reduce code duplication
-        backend_name = attn_backend.backend.name.lower()
        requires_mla = "deepseek" in model_name.lower()
-        is_mla = "mla" in backend_name
-        # DeepSeek V3.2 uses sparse MLA
-        requires_sparse = "v3.2" in model_name.lower()
-        is_sparse = "sparse" in backend_name
+        is_mla = "mla" in attn_backend.backend.name.lower()

-        if requires_mla != is_mla or requires_sparse != is_sparse:
+        if requires_mla != is_mla:
            pytest.skip(
                f"Incompatible model '{model_name}' and "
                f"attention backend '{attn_backend.backend.name}'"
@@ -235,9 +231,7 @@ def run_e2e_fusion_test(monkeypatch, caplog_mp_spawn):
                )

            elif match_name == "attn_quant_fusion":
-                actual_match = match_table.get(
-                    "attn_quant_fusion", 0
-                ) + match_table.get("mla_attn_quant_fusion", 0)
+                actual_match = match_table.get(match_name, 0)
                assert actual_match == expected_matches * n_expected, (
                    f"Could not find {expected_matches * n_expected} "
                    f"{match_name} (found {actual_match})."
@@ -58,15 +58,6 @@ TRITON_MLA_ATTN = pytest.param(
    id="TRITON_MLA",
 )

-FLASHMLA_SPARSE_ATTN = pytest.param(
-    AttentionBackendCase(backend=AttentionBackendEnum.FLASHMLA_SPARSE),
-    id="FLASHMLA_SPARSE",
-    marks=pytest.mark.skipif(
-        not is_blackwell(),
-        reason="FlashMLA Sparse requires Blackwell",
-    ),
-)
-
 # Models
 llama3_8b = ModelFusionInfo(
    model_name="meta-llama/Llama-3.1-8B-Instruct",
@@ -150,18 +141,6 @@ qwen3_a3b_fp8 = ModelFusionInfo(
    ),
 )

-deepseek_coder_v2_lite_fp8 = ModelFusionInfo(
-    model_name="RedHatAI/DeepSeek-Coder-V2-Lite-Instruct-FP8",
-    matches=lambda n_layers: Matches(
-        # first_k_dense_replace=1; MoE hides most rms+quant sites
-        rms_quant_fusion=1,
-        act_quant_fusion=min(1, n_layers),  # dense layers only
-        # MLA attn + static FP8 quant
-        attn_quant_fusion=n_layers,
-        ar_rms_fusion=n_layers * 2 + 1,
-    ),
-)
-
 deepseek_v3_fp8 = ModelFusionInfo(
    model_name="deepseek-ai/DeepSeek-V3",
    matches=lambda n_layers: Matches(
@@ -173,7 +152,7 @@ deepseek_v3_fp8 = ModelFusionInfo(
        rms_quant_fusion=n_layers * 2 + min(3, n_layers),  # add for 3 dense layers
        # silu+block quant
        act_quant_fusion=min(3, n_layers),  # dense layers only
-        # MLA attn + per-group FP8 quant not supported yet:
+        # MLA attn + quant not supported yet:
        # https://github.com/vllm-project/vllm/issues/35792
        attn_quant_fusion=0,
        ar_rms_fusion=n_layers * 2 + 1,
@@ -183,16 +162,6 @@ deepseek_v3_fp8 = ModelFusionInfo(
    ),
 )

-deepseek_v32_fp4 = ModelFusionInfo(
-    model_name="nvidia/DeepSeek-V3.2-NVFP4",
-    matches=lambda n_layers: Matches(
-        rms_quant_fusion=0,
-        act_quant_fusion=0,
-        attn_quant_fusion=n_layers,
-        ar_rms_fusion=n_layers * 2 + 1,
-    ),
-)
-
 gpt_oss_20b = ModelFusionInfo(
    model_name="openai/gpt-oss-20b",
    matches=lambda n_layers: Matches(
@@ -18,14 +18,11 @@ from .common import (
 from .models import (
    FLASHINFER_ATTN,
    FLASHINFER_MLA_ATTN,
-    FLASHMLA_SPARSE_ATTN,
    ROCM_AITER_UNIFIED_ATTN,
    ROCM_ATTN,
    TRITON_ATTN,
    TRITON_MLA_ATTN,
-    deepseek_coder_v2_lite_fp8,
    deepseek_v3_fp8,
-    deepseek_v32_fp4,
    llama3_8b_fp4,
    llama3_8b_fp8,
    llama4_scout_fp4,
@@ -40,7 +37,6 @@ from .models import (
        (*llama3_8b_fp8, False),
        (*qwen3_a3b_fp8, False),
        (*qwen3_a3b_fp8, True),
-        (*deepseek_coder_v2_lite_fp8, False),
        (*deepseek_v3_fp8, False),
        (*deepseek_v3_fp8, True),
        pytest.param(
@@ -148,12 +144,9 @@ def test_tp1_fp8_fusions(

@pytest.mark.parametrize(
    "model_name, matches_fn, model_kwargs, hf_overrides",
-    [llama3_8b_fp4, llama4_scout_fp4, deepseek_v32_fp4],
-)
-@pytest.mark.parametrize(
-    "attn_backend",
-    [FLASHINFER_ATTN, FLASHMLA_SPARSE_ATTN],
+    [llama3_8b_fp4, llama4_scout_fp4],
 )
+@pytest.mark.parametrize("attn_backend", [FLASHINFER_ATTN])
@pytest.mark.parametrize("n_layers", [6])
@pytest.mark.parametrize("custom_ops", custom_ops_combos("rms_norm"))
@pytest.mark.parametrize("inductor_graph_partition", INDUCTOR_GRAPH_PARTITION)
@@ -18,11 +18,8 @@ from .common import (
 from .models import (
    FLASHINFER_ATTN,
    FLASHINFER_MLA_ATTN,
-    FLASHMLA_SPARSE_ATTN,
    TRITON_ATTN,
-    deepseek_coder_v2_lite_fp8,
    deepseek_v3_fp8,
-    deepseek_v32_fp4,
    gpt_oss_20b,
    llama3_8b,
    llama3_8b_fp4,
@@ -40,13 +37,7 @@ pytestmark = pytest.mark.skipif(not current_platform.is_cuda(), reason="Only tes
@pytest.mark.parametrize(
    "model_name, matches_fn, model_kwargs, hf_overrides",
    # qwen3 & dsv3 should still fuse AR+rms even though group quant is not yet supported
-    [
-        llama3_8b_fp8,
-        llama4_scout_fp8,
-        qwen3_a3b_fp8,
-        deepseek_coder_v2_lite_fp8,
-        deepseek_v3_fp8,
-    ],
+    [llama3_8b_fp8, llama4_scout_fp8, qwen3_a3b_fp8, deepseek_v3_fp8],
 )
@pytest.mark.parametrize(
    "attn_backend", [TRITON_ATTN, FLASHINFER_ATTN, FLASHINFER_MLA_ATTN]
@@ -113,12 +104,9 @@ def test_tp2_ar_rms_fp8_fusions(
@multi_gpu_test(num_gpus=2)
@pytest.mark.parametrize(
    "model_name, matches_fn, model_kwargs, hf_overrides",
-    [llama3_8b_fp4, llama4_scout_fp4, deepseek_v32_fp4],
-)
-@pytest.mark.parametrize(
-    "attn_backend",
-    [FLASHINFER_ATTN, FLASHMLA_SPARSE_ATTN],
+    [llama3_8b_fp4, llama4_scout_fp4],
 )
+@pytest.mark.parametrize("attn_backend", [FLASHINFER_ATTN])
@pytest.mark.parametrize("n_layers", [4])
@pytest.mark.parametrize("custom_ops", custom_ops_combos("rms_norm"))
@pytest.mark.parametrize("inductor_graph_partition", INDUCTOR_GRAPH_PARTITION)
@@ -1,508 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-import copy
-
-import pytest
-import torch._dynamo
-
-from tests.compile.backend import LazyInitPass, TestBackend
-from tests.utils import TestFP8Layer, flat_product
-from tests.v1.attention.utils import BatchSpec, create_common_attn_metadata
-from vllm._custom_ops import cutlass_scaled_fp4_mm, scaled_fp4_quant
-from vllm.compilation.passes.fusion.matcher_utils import QUANT_OPS
-from vllm.compilation.passes.fusion.mla_attn_quant_fusion import (
-    MLA_ATTN_OP,
-    MLAAttnQuantFusionPass,
-)
-from vllm.compilation.passes.fx_utils import find_op_nodes
-from vllm.compilation.passes.utility.noop_elimination import NoOpEliminationPass
-from vllm.compilation.passes.utility.post_cleanup import PostCleanupPass
-from vllm.config import (
-    AttentionConfig,
-    CacheConfig,
-    CompilationConfig,
-    CompilationMode,
-    ModelConfig,
-    PassConfig,
-    SchedulerConfig,
-    VllmConfig,
-    set_current_vllm_config,
-)
-from vllm.forward_context import get_forward_context, set_forward_context
-from vllm.model_executor.layers.attention import MLAAttention
-from vllm.model_executor.layers.linear import ColumnParallelLinear
-from vllm.model_executor.layers.quantization.fp8 import Fp8Config
-from vllm.model_executor.layers.quantization.modelopt import ModelOptNvFp4Config
-from vllm.model_executor.layers.quantization.utils.quant_utils import (
-    QuantKey,
-    kFp8StaticTensorSym,
-    kNvfp4Dynamic,
-)
-from vllm.platforms import current_platform
-from vllm.v1.attention.backend import AttentionMetadata
-from vllm.v1.attention.backends.registry import AttentionBackendEnum
-from vllm.v1.kv_cache_interface import MLAAttentionSpec
-
-FP8_DTYPE = current_platform.fp8_dtype()
-FP4_DTYPE = torch.uint8
-
-
-class MLAAttentionQuantPatternModel(torch.nn.Module):
-    """Base model for MLA AttentionQuantPattern fusion."""
-
-    def __init__(
-        self,
-        num_heads: int,
-        qk_nope_head_dim: int,
-        qk_rope_head_dim: int,
-        v_head_dim: int,
-        kv_lora_rank: int,
-        kv_cache_dtype: torch.dtype,
-        device: torch.device,
-        vllm_config: VllmConfig,
-        **kwargs,
-    ):
-        super().__init__()
-        self.num_heads = num_heads
-        self.qk_nope_head_dim = qk_nope_head_dim
-        self.qk_rope_head_dim = qk_rope_head_dim
-        self.qk_head_dim = qk_nope_head_dim + qk_rope_head_dim
-        self.v_head_dim = v_head_dim
-        self.kv_lora_rank = kv_lora_rank
-        self.output_dim = num_heads * v_head_dim
-        self.head_size = kv_lora_rank + qk_rope_head_dim
-        self.kv_cache_dtype = kv_cache_dtype
-        self.device = device
-        self.vllm_config = vllm_config
-
-        # Create kv_b_proj (ColumnParallelLinear) on device.
-        # Reuse weights from prior model instance when available, because
-        # ColumnParallelLinear may get NaN from recycled CUDA memory after
-        # torch.compile runs in the same process.
-        kv_b_proj = ColumnParallelLinear(
-            input_size=kv_lora_rank,
-            output_size=num_heads * (qk_nope_head_dim + v_head_dim),
-            bias=False,
-            prefix="model.layers.0.self_attn.kv_b_proj",
-        ).to(device)
-        kv_b_proj_weight = kwargs.get("kv_b_proj_weight")
-        if kv_b_proj_weight is not None:
-            kv_b_proj.weight.data.copy_(kv_b_proj_weight)
-        elif kv_b_proj.weight.data.isnan().any():
-            # Sanitize NaN from recycled CUDA memory
-            kv_b_proj.weight.data.normal_()
-
-        # Create MLAAttention
-        self.mla_attn = MLAAttention(
-            num_heads=num_heads,
-            scale=1.0 / (self.qk_head_dim**0.5),
-            qk_nope_head_dim=qk_nope_head_dim,
-            qk_rope_head_dim=qk_rope_head_dim,
-            v_head_dim=v_head_dim,
-            q_lora_rank=None,
-            kv_lora_rank=kv_lora_rank,
-            kv_b_proj=kv_b_proj,
-            cache_config=vllm_config.cache_config,
-            quant_config=self.quant_config,
-            prefix="model.layers.0.self_attn.attn",
-        )
-        self.mla_attn._k_scale = self.mla_attn._k_scale.to(device)
-        self.mla_attn._v_scale = self.mla_attn._v_scale.to(device)
-
-        # Initialize W_UK_T and W_UV from kv_b_proj weights
-        self.mla_attn.process_weights_after_loading(torch.get_default_dtype())
-        self.kv_b_proj_weight = kv_b_proj.weight.data.clone()
-
-        self.block_size = 16
-
-        # Initialize MLA MetadataBuilder
-        self.builder = self.mla_attn.attn_backend.get_builder_cls()(
-            kv_cache_spec=MLAAttentionSpec(
-                block_size=self.block_size,
-                num_kv_heads=1,
-                head_size=self.head_size,
-                dtype=self.kv_cache_dtype,
-            ),
-            layer_names=[self.mla_attn.layer_name],
-            vllm_config=self.vllm_config,
-            device=self.device,
-        )
-
-    def build_attn_metadata(self, batch_size: int) -> AttentionMetadata:
-        """Initialize MLA attention metadata.
-
-        NOTE: Uses decode-only batch (query_len=1 per request). The prefill
-        (forward_mha) path is not separately tested here because it requires
-        FlashAttention availability and different input tensor shapes. The
-        quant logic in forward_impl is identical for both paths — it quantizes
-        the full output[:num_actual_toks] buffer after both forward_mha and
-        forward_mqa have written their results.
-        """
-
-        batch_spec = BatchSpec(seq_lens=[1] * batch_size, query_lens=[1] * batch_size)
-        common_attn_metadata = create_common_attn_metadata(
-            batch_spec, self.block_size, self.device, arange_block_indices=True
-        )
-
-        max_blocks = (max(batch_spec.seq_lens) + self.block_size - 1) // self.block_size
-        num_blocks = batch_size * max_blocks
-
-        # MLA KV cache is 3D: (num_blocks, block_size, head_size)
-        attn_backend = self.mla_attn.attn_backend
-        kv_cache_shape = attn_backend.get_kv_cache_shape(
-            num_blocks, self.block_size, 1, self.head_size
-        )
-        try:
-            kv_cache_stride_order = attn_backend.get_kv_cache_stride_order()
-        except (AttributeError, NotImplementedError):
-            kv_cache_stride_order = tuple(range(len(kv_cache_shape)))
-
-        ordered_shape = tuple(kv_cache_shape[i] for i in kv_cache_stride_order)
-        inv_order = [
-            kv_cache_stride_order.index(i) for i in range(len(kv_cache_stride_order))
-        ]
-
-        raw_tensor = torch.zeros(
-            ordered_shape, dtype=self.kv_cache_dtype, device=self.device
-        )
-        kv_cache = raw_tensor.permute(*inv_order)
-
-        self.mla_attn.kv_cache = kv_cache
-
-        self.attn_metadata = self.builder.build(
-            common_prefix_len=0, common_attn_metadata=common_attn_metadata
-        )
-
-        return self.attn_metadata
-
-
-class TestMLAAttentionFp8StaticQuantPatternModel(MLAAttentionQuantPatternModel):
-    """Test model for MLA Attention + FP8 static quant fusion."""
-
-    quant_key = kFp8StaticTensorSym
-    quant_config = Fp8Config()
-
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-
-        self.fp8_linear = TestFP8Layer(
-            weight_shape=(self.output_dim, self.output_dim),
-            activation_quant_key=self.quant_key,
-            weight_quant_key=self.quant_key,
-            device=self.device,
-        )
-
-        w = kwargs.get("w")
-        if w is not None:
-            self.fp8_linear.weight = w["weight"]
-            self.fp8_linear.weight_scale = w["wscale"]
-            self.fp8_linear.input_scale = w["scale"]
-
-        self.w = {
-            "weight": self.fp8_linear.weight,
-            "wscale": self.fp8_linear.weight_scale,
-            "scale": self.fp8_linear.input_scale,
-        }
-
-    def forward(
-        self,
-        q: torch.Tensor,
-        kv_c_normed: torch.Tensor,
-        k_pe: torch.Tensor,
-    ):
-        """Forward pass that creates the MLA attention + FP8 quant pattern."""
-        attn_output = self.mla_attn(
-            q,
-            kv_c_normed,
-            k_pe,
-            output_shape=(q.shape[0], self.output_dim),
-        )
-        return self.fp8_linear(attn_output)
-
-
-class TestMLAAttentionNvfp4QuantPatternModel(MLAAttentionQuantPatternModel):
-    """Test model for MLA Attention + NVFP4 quant fusion."""
-
-    quant_key = kNvfp4Dynamic
-    quant_config = ModelOptNvFp4Config(
-        is_checkpoint_nvfp4_serialized=False,
-        kv_cache_quant_algo=None,
-        exclude_modules=[],
-    )
-
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-
-        self.w = kwargs.get(
-            "w",
-            {
-                "weight": torch.randint(
-                    256,
-                    (self.output_dim, self.output_dim // 2),
-                    dtype=FP4_DTYPE,
-                    device=self.device,
-                ),
-                "wscale_swizzled": torch.randn(
-                    self.output_dim, self.output_dim // 16
-                ).to(dtype=FP8_DTYPE, device=self.device),
-                "wscale": torch.tensor([500], dtype=torch.float32, device=self.device),
-                "scale": torch.tensor([0.002], dtype=torch.float32, device=self.device),
-            },
-        )
-
-    def forward(
-        self,
-        q: torch.Tensor,
-        kv_c_normed: torch.Tensor,
-        k_pe: torch.Tensor,
-    ):
-        """Forward pass that creates the MLA attention + NVFP4 quant pattern."""
-        attn_output = self.mla_attn(
-            q,
-            kv_c_normed,
-            k_pe,
-            output_shape=(q.shape[0], self.output_dim),
-        )
-        quant_output, output_block_scale = scaled_fp4_quant(
-            attn_output, 1 / self.w["scale"]
-        )
-        return cutlass_scaled_fp4_mm(
-            a=quant_output,
-            b=self.w["weight"],
-            block_scale_a=output_block_scale,
-            block_scale_b=self.w["wscale_swizzled"],
-            alpha=self.w["scale"] * self.w["wscale"],
-            out_dtype=attn_output.dtype,
-        )
-
-
-def is_nvfp4_supported():
-    return current_platform.has_device_capability(100)
-
-
-# MLA test configuration
-MLA_DIMS: list[tuple[int, int, int, int, int]] = []
-PATTERN_TEST_MODELS_MLA_FP8: list[tuple[str, type]] = []
-PATTERN_TEST_MODELS_MLA_FP4: list[tuple[str, type]] = []
-BACKENDS_MLA_FP8: list[AttentionBackendEnum] = []
-BACKENDS_MLA_FP4: list[AttentionBackendEnum] = []
-
-if current_platform.is_cuda():
-    # (num_heads, qk_nope_head_dim, qk_rope_head_dim, v_head_dim, kv_lora_rank)
-    MLA_DIMS = [(16, 128, 64, 128, 512)]
-    PATTERN_TEST_MODELS_MLA_FP8 = [
-        (
-            "deepseek-ai/DeepSeek-V2-Lite",
-            TestMLAAttentionFp8StaticQuantPatternModel,
-        )
-    ]
-    PATTERN_TEST_MODELS_MLA_FP4 = [
-        (
-            "deepseek-ai/DeepSeek-V2-Lite",
-            TestMLAAttentionNvfp4QuantPatternModel,
-        )
-    ]
-    BACKENDS_MLA_FP8 = [AttentionBackendEnum.TRITON_MLA]
-    BACKENDS_MLA_FP4 = [AttentionBackendEnum.TRITON_MLA]
-
-
-@pytest.mark.parametrize(
-    "num_heads, qk_nope_head_dim, qk_rope_head_dim, v_head_dim, kv_lora_rank",
-    MLA_DIMS,
-)
-@pytest.mark.parametrize("batch_size", [7, 256] if current_platform.is_cuda() else [8])
-@pytest.mark.parametrize("dtype", [torch.bfloat16])
-@pytest.mark.parametrize(
-    "backend, model_name, model_class, custom_ops",
-    list(
-        flat_product(
-            BACKENDS_MLA_FP8,
-            PATTERN_TEST_MODELS_MLA_FP8,
-            ["+quant_fp8", "-quant_fp8"],
-        )
-    )
-    + list(flat_product(BACKENDS_MLA_FP4, PATTERN_TEST_MODELS_MLA_FP4, [""])),
-)
-@pytest.mark.skipif(
-    not current_platform.is_cuda_alike(), reason="Only test ROCm or CUDA"
-)
-@pytest.mark.skipif(not current_platform.supports_fp8(), reason="Need FP8")
-def test_mla_attention_quant_pattern(
-    num_heads: int,
-    qk_nope_head_dim: int,
-    qk_rope_head_dim: int,
-    v_head_dim: int,
-    kv_lora_rank: int,
-    batch_size: int,
-    dtype: torch.dtype,
-    custom_ops: str,
-    model_name: str,
-    model_class: type[MLAAttentionQuantPatternModel],
-    backend: AttentionBackendEnum,
-    dist_init,
-    monkeypatch,
-    use_fresh_inductor_cache,
-):
-    """Test MLA AttentionQuantPattern fusion pass"""
-    if (
-        model_class is TestMLAAttentionNvfp4QuantPatternModel
-        and not is_nvfp4_supported()
-    ):
-        pytest.skip("NVFP4 is not supported on this GPU (requires SM 100+).")
-
-    monkeypatch.setenv("VLLM_DISABLE_COMPILE_CACHE", "1")
-
-    custom_ops_list = custom_ops.split(",") if custom_ops else []
-
-    device = torch.device("cuda:0")
-    torch.set_default_dtype(dtype)
-    torch.manual_seed(42)
-
-    model_config = ModelConfig(
-        model=model_name,
-        max_model_len=2048,
-        dtype=dtype,
-    )
-    vllm_config = VllmConfig(
-        model_config=model_config,
-        scheduler_config=SchedulerConfig(
-            max_num_seqs=1024,
-            max_model_len=model_config.max_model_len,
-            is_encoder_decoder=model_config.is_encoder_decoder,
-        ),
-        compilation_config=CompilationConfig(
-            mode=CompilationMode.VLLM_COMPILE,
-            custom_ops=custom_ops_list,
-        ),
-        cache_config=CacheConfig(cache_dtype="auto"),
-        attention_config=AttentionConfig(backend=backend),
-    )
-
-    # MLA inputs: q(B, N, qk_head_dim), kv_c_normed(B, L), k_pe(B, 1, R)
-    qk_head_dim = qk_nope_head_dim + qk_rope_head_dim
-    q = torch.randn(batch_size, num_heads, qk_head_dim, dtype=dtype, device=device)
-    kv_c_normed = torch.randn(batch_size, kv_lora_rank, dtype=dtype, device=device)
-    k_pe = torch.randn(batch_size, 1, qk_rope_head_dim, dtype=dtype, device=device)
-
-    # Mark first dimension as dynamic
-    torch._dynamo.mark_dynamic(q, 0)
-    torch._dynamo.mark_dynamic(kv_c_normed, 0)
-    torch._dynamo.mark_dynamic(k_pe, 0)
-
-    # Run model without fusion
-    vllm_config_unfused = copy.deepcopy(vllm_config)
-    with (
-        set_current_vllm_config(vllm_config_unfused),
-        set_forward_context(attn_metadata=None, vllm_config=vllm_config_unfused),
-    ):
-        model_unfused = model_class(
-            num_heads=num_heads,
-            qk_nope_head_dim=qk_nope_head_dim,
-            qk_rope_head_dim=qk_rope_head_dim,
-            v_head_dim=v_head_dim,
-            kv_lora_rank=kv_lora_rank,
-            kv_cache_dtype=dtype,
-            device=device,
-            vllm_config=vllm_config_unfused,
-        )
-        model_unfused = model_unfused.to(device)
-        # HACK: See #131044
-        result_unfused_0 = model_unfused(q, kv_c_normed, k_pe)  # noqa: F841
-
-        forward_ctx = get_forward_context()
-        forward_ctx.attn_metadata = model_unfused.build_attn_metadata(batch_size)
-
-        compiled_unfused = torch.compile(model_unfused, fullgraph=True)
-        result_unfused = compiled_unfused(q, kv_c_normed, k_pe)
-
-    # Run model with attn fusion enabled
-    vllm_config.compilation_config.pass_config = PassConfig(
-        fuse_attn_quant=True, eliminate_noops=True
-    )
-    with (
-        set_current_vllm_config(vllm_config),
-        set_forward_context(attn_metadata=None, vllm_config=vllm_config),
-    ):
-        model_fused = model_class(
-            num_heads=num_heads,
-            qk_nope_head_dim=qk_nope_head_dim,
-            qk_rope_head_dim=qk_rope_head_dim,
-            v_head_dim=v_head_dim,
-            kv_lora_rank=kv_lora_rank,
-            kv_cache_dtype=dtype,
-            device=device,
-            vllm_config=vllm_config,
-            w=model_unfused.w,
-            kv_b_proj_weight=model_unfused.kv_b_proj_weight,
-        )
-        model_fused = model_fused.to(device)
-
-        forward_ctx = get_forward_context()
-        forward_ctx.attn_metadata = model_fused.build_attn_metadata(batch_size)
-
-        # Create test backend with fusion passes
-        noop_pass = NoOpEliminationPass(vllm_config)
-        attn_pass = LazyInitPass(MLAAttnQuantFusionPass, vllm_config)
-        cleanup_pass = PostCleanupPass(vllm_config)
-
-        test_backend = TestBackend(noop_pass, attn_pass, cleanup_pass)
-        # HACK: See https://github.com/vllm-project/vllm/issues/31044
-        result_fused_0 = model_fused(q, kv_c_normed, k_pe)  # noqa: F841
-
-        compiled_fused = torch.compile(
-            model_fused, backend=test_backend, fullgraph=True
-        )
-
-        result_fused = compiled_fused(q, kv_c_normed, k_pe)
-
-    # Check attn fusion support
-    quant_key: QuantKey = model_class.quant_key
-    attn_fusion_supported = [
-        layer.impl.fused_output_quant_supported(quant_key)
-        for key, layer in vllm_config.compilation_config.static_forward_context.items()
-        if isinstance(layer, MLAAttention)
-    ]
-    assert sum(attn_fusion_supported) == len(attn_fusion_supported), (
-        "All MLA layers should support attention fusion"
-    )
-
-    # Check quantization ops in the graph
-    quant_op = (
-        torch.ops.aten.reciprocal
-        if "-quant_fp8" in custom_ops_list
-        else QUANT_OPS[quant_key]
-    )
-    test_backend.check_before_ops([quant_op], fully_replaced=quant_key is kNvfp4Dynamic)
-
-    assert attn_pass.pass_.matched_count == sum(attn_fusion_supported)
-
-    # Check MLA attention ops in the graph
-    attn_nodes_pre = list(find_op_nodes(MLA_ATTN_OP, test_backend.graph_pre_pass))
-    attn_nodes_post = list(find_op_nodes(MLA_ATTN_OP, test_backend.graph_post_pass))
-
-    assert len(attn_nodes_pre) > 0, "Should have MLA attention nodes before fusion"
-    assert len(attn_nodes_pre) == len(attn_nodes_post), (
-        "Should have same number of MLA attention nodes before and after fusion"
-    )
-    assert attn_nodes_pre[0].kwargs.get("output_scale") is None, (
-        "MLA attention should not have output_scale before fusion"
-    )
-    assert attn_nodes_post[0].kwargs.get("output_scale") is not None, (
-        "MLA attention should have output_scale after fusion"
-    )
-
-    assert attn_nodes_pre[0].kwargs.get("output_block_scale") is None, (
-        "MLA attention should not have output_block_scale before fusion"
-    )
-
-    if quant_key.dtype == FP8_DTYPE:
-        assert attn_nodes_post[0].kwargs.get("output_block_scale") is None, (
-            "MLA attention should not have output_block_scale after FP8 fusion"
-        )
-    elif quant_key.dtype == FP4_DTYPE:
-        assert attn_nodes_post[0].kwargs.get("output_block_scale") is not None, (
-            "MLA attention should have output_block_scale after FP4 fusion"
-        )
-
-    # Check numerical correctness
-    torch.testing.assert_close(result_unfused, result_fused, atol=1e-2, rtol=1e-2)
@@ -1,474 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import json
-from dataclasses import dataclass, field
-from typing import Any
-from unittest.mock import AsyncMock, MagicMock
-
-import pytest
-
-from vllm.config.multimodal import MultiModalConfig
-from vllm.entrypoints.openai.engine.protocol import StreamOptions
-from vllm.entrypoints.openai.models.protocol import BaseModelPath
-from vllm.entrypoints.openai.models.serving import OpenAIServingModels
-from vllm.entrypoints.serve.disagg.protocol import GenerateRequest
-from vllm.entrypoints.serve.disagg.serving import ServingTokens
-from vllm.entrypoints.serve.render.serving import OpenAIServingRender
-from vllm.logprobs import Logprob
-from vllm.outputs import CompletionOutput, RequestOutput
-from vllm.renderers import renderer_from_config
-from vllm.sampling_params import SamplingParams
-from vllm.v1.engine.async_llm import AsyncLLM
-
-MODEL_NAME = "openai-community/gpt2"
-BASE_MODEL_PATHS = [
-    BaseModelPath(name=MODEL_NAME, model_path=MODEL_NAME),
-]
-
-
-@dataclass
-class MockHFConfig:
-    model_type: str = "any"
-
-
-@dataclass
-class MockModelConfig:
-    task = "generate"
-    runner_type = "generate"
-    model = MODEL_NAME
-    tokenizer = MODEL_NAME
-    trust_remote_code = False
-    tokenizer_mode = "auto"
-    max_model_len = 100
-    tokenizer_revision = None
-    multimodal_config = MultiModalConfig()
-    hf_config = MockHFConfig()
-    hf_text_config = MockHFConfig()
-    logits_processors: list[str] | None = None
-    diff_sampling_param: dict | None = None
-    allowed_local_media_path: str = ""
-    allowed_media_domains: list[str] | None = None
-    encoder_config = None
-    generation_config: str = "auto"
-    media_io_kwargs: dict[str, dict[str, Any]] = field(default_factory=dict)
-    skip_tokenizer_init = False
-    is_encoder_decoder: bool = False
-    is_multimodal_model: bool = False
-    renderer_num_workers: int = 1
-
-    def get_diff_sampling_param(self):
-        return self.diff_sampling_param or {}
-
-
-@dataclass
-class MockParallelConfig:
-    _api_process_rank: int = 0
-
-
-@dataclass
-class MockVllmConfig:
-    model_config: MockModelConfig
-    parallel_config: MockParallelConfig
-
-
-def _build_renderer(model_config: MockModelConfig):
-    return renderer_from_config(
-        MockVllmConfig(model_config, parallel_config=MockParallelConfig()),
-    )
-
-
-def _build_serving_tokens(engine: AsyncLLM, **kwargs) -> ServingTokens:
-    models = OpenAIServingModels(
-        engine_client=engine,
-        base_model_paths=BASE_MODEL_PATHS,
-    )
-    serving_render = OpenAIServingRender(
-        model_config=engine.model_config,
-        renderer=engine.renderer,
-        io_processor=engine.io_processor,
-        model_registry=models.registry,
-        request_logger=None,
-        chat_template=None,
-        chat_template_content_format="auto",
-    )
-    serving = ServingTokens(
-        engine,
-        models,
-        openai_serving_render=serving_render,
-        request_logger=None,
-        **kwargs,
-    )
-
-    async def _fake_preprocess(*args, **kwargs):
-        return [{"prompt_token_ids": [1, 2, 3]}]
-
-    serving.openai_serving_render.preprocess_completion = AsyncMock(
-        side_effect=_fake_preprocess
-    )
-    return serving
-
-
-def _make_request_output(
-    request_id: str,
-    token_ids: list[int],
-    finish_reason: str | None = None,
-    finished: bool = False,
-    prompt_token_ids: list[int] | None = None,
-    logprobs: list[dict[int, Any] | None] | None = None,
-    num_cached_tokens: int | None = None,
-    index: int = 0,
-) -> RequestOutput:
-    return RequestOutput(
-        request_id=request_id,
-        prompt=None,
-        prompt_token_ids=prompt_token_ids or [1, 2, 3],
-        prompt_logprobs=None,
-        outputs=[
-            CompletionOutput(
-                index=index,
-                text="",
-                token_ids=token_ids,
-                cumulative_logprob=None,
-                logprobs=logprobs,
-                finish_reason=finish_reason,
-            )
-        ],
-        finished=finished,
-        metrics=None,
-        lora_request=None,
-        encoder_prompt=None,
-        encoder_prompt_token_ids=None,
-        num_cached_tokens=num_cached_tokens,
-    )
-
-
-def _mock_engine() -> MagicMock:
-    engine = MagicMock(spec=AsyncLLM)
-    engine.errored = False
-    engine.model_config = MockModelConfig()
-    engine.input_processor = MagicMock()
-    engine.io_processor = MagicMock()
-    engine.renderer = _build_renderer(engine.model_config)
-    return engine
-
-
-def _parse_sse_chunks(chunks: list[str]) -> list[Any]:
-    """Parse SSE chunks into dicts (JSON) or raw strings ([DONE])."""
-    parsed: list[Any] = []
-    for chunk in chunks:
-        assert chunk.startswith("data: ") and chunk.endswith("\n\n")
-        payload = chunk[len("data: ") : -len("\n\n")]
-        if payload == "[DONE]":
-            parsed.append("[DONE]")
-        else:
-            parsed.append(json.loads(payload))
-    return parsed
-
-
-@pytest.mark.asyncio
-async def test_stream_basic():
-    """Streaming returns SSE chunks with correct token_ids and ends with [DONE]."""
-    engine = _mock_engine()
-
-    async def mock_generate(*args, **kwargs):
-        yield _make_request_output("req-1", token_ids=[10])
-        yield _make_request_output("req-1", token_ids=[20, 30])
-        yield _make_request_output(
-            "req-1", token_ids=[40], finish_reason="stop", finished=True
-        )
-
-    engine.generate = MagicMock(side_effect=mock_generate)
-    serving = _build_serving_tokens(engine)
-
-    request = GenerateRequest(
-        token_ids=[1, 2, 3],
-        sampling_params=SamplingParams(max_tokens=10),
-        model=MODEL_NAME,
-        stream=True,
-    )
-
-    response = await serving.serve_tokens(request)
-    chunks = []
-    async for chunk in response:
-        chunks.append(chunk)
-
-    parsed = _parse_sse_chunks(chunks)
-
-    # 3 data chunks + [DONE]
-    assert parsed[-1] == "[DONE]"
-    data_chunks = [c for c in parsed if c != "[DONE]"]
-    assert len(data_chunks) == 3
-
-    assert data_chunks[0]["choices"][0]["token_ids"] == [10]
-    assert data_chunks[1]["choices"][0]["token_ids"] == [20, 30]
-    assert data_chunks[2]["choices"][0]["token_ids"] == [40]
-    assert data_chunks[2]["choices"][0]["finish_reason"] == "stop"
-
-
-@pytest.mark.asyncio
-async def test_stream_error_mid_generation():
-    """finish_reason='error' mid-stream yields error chunk then [DONE]."""
-    engine = _mock_engine()
-
-    async def mock_generate(*args, **kwargs):
-        yield _make_request_output("req-1", token_ids=[10])
-        yield _make_request_output(
-            "req-1", token_ids=[20], finish_reason="error", finished=True
-        )
-
-    engine.generate = MagicMock(side_effect=mock_generate)
-    serving = _build_serving_tokens(engine)
-
-    request = GenerateRequest(
-        token_ids=[1, 2, 3],
-        sampling_params=SamplingParams(max_tokens=10),
-        model=MODEL_NAME,
-        stream=True,
-    )
-
-    response = await serving.serve_tokens(request)
-    chunks = []
-    async for chunk in response:
-        chunks.append(chunk)
-
-    assert len(chunks) >= 2
-    assert any("Internal server error" in chunk for chunk in chunks), (
-        f"Expected error message in chunks: {chunks}"
-    )
-    assert chunks[-1] == "data: [DONE]\n\n"
-
-
-@pytest.mark.asyncio
-async def test_stream_error_with_empty_delta():
-    """finish_reason='error' with empty delta_token_ids still raises."""
-    engine = _mock_engine()
-
-    async def mock_generate(*args, **kwargs):
-        yield _make_request_output("req-1", token_ids=[10])
-        yield _make_request_output(
-            "req-1", token_ids=[], finish_reason="error", finished=True
-        )
-
-    engine.generate = MagicMock(side_effect=mock_generate)
-    serving = _build_serving_tokens(engine)
-
-    request = GenerateRequest(
-        token_ids=[1, 2, 3],
-        sampling_params=SamplingParams(max_tokens=10),
-        model=MODEL_NAME,
-        stream=True,
-    )
-
-    response = await serving.serve_tokens(request)
-    chunks = []
-    async for chunk in response:
-        chunks.append(chunk)
-
-    assert any("Internal server error" in chunk for chunk in chunks), (
-        f"Expected error message in chunks: {chunks}"
-    )
-    assert chunks[-1] == "data: [DONE]\n\n"
-
-
-@pytest.mark.asyncio
-async def test_stream_skips_empty_token_output():
-    """Outputs with empty token_ids are skipped (no chunk emitted)."""
-    engine = _mock_engine()
-
-    async def mock_generate(*args, **kwargs):
-        yield _make_request_output("req-1", token_ids=[10])
-        yield _make_request_output("req-1", token_ids=[])
-        yield _make_request_output(
-            "req-1", token_ids=[20], finish_reason="stop", finished=True
-        )
-
-    engine.generate = MagicMock(side_effect=mock_generate)
-    serving = _build_serving_tokens(engine)
-
-    request = GenerateRequest(
-        token_ids=[1, 2, 3],
-        sampling_params=SamplingParams(max_tokens=10),
-        model=MODEL_NAME,
-        stream=True,
-    )
-
-    response = await serving.serve_tokens(request)
-    chunks = []
-    async for chunk in response:
-        chunks.append(chunk)
-
-    parsed = _parse_sse_chunks(chunks)
-    assert parsed[-1] == "[DONE]"
-    data_chunks = [c for c in parsed if c != "[DONE]"]
-
-    # Only 2 data chunks — the empty one is skipped
-    assert len(data_chunks) == 2
-    assert data_chunks[0]["choices"][0]["token_ids"] == [10]
-    assert data_chunks[1]["choices"][0]["token_ids"] == [20]
-
-
-@pytest.mark.asyncio
-async def test_stream_include_usage():
-    """stream_options.include_usage emits a final usage-only chunk."""
-    engine = _mock_engine()
-
-    async def mock_generate(*args, **kwargs):
-        yield _make_request_output("req-1", token_ids=[10])
-        yield _make_request_output(
-            "req-1", token_ids=[20], finish_reason="stop", finished=True
-        )
-
-    engine.generate = MagicMock(side_effect=mock_generate)
-    serving = _build_serving_tokens(engine)
-
-    request = GenerateRequest(
-        token_ids=[1, 2, 3],
-        sampling_params=SamplingParams(max_tokens=10),
-        model=MODEL_NAME,
-        stream=True,
-        stream_options=StreamOptions(include_usage=True),
-    )
-
-    response = await serving.serve_tokens(request)
-    chunks = []
-    async for chunk in response:
-        chunks.append(chunk)
-
-    parsed = _parse_sse_chunks(chunks)
-    assert parsed[-1] == "[DONE]"
-
-    # The chunk before [DONE] should be the usage-only chunk
-    usage_chunk = parsed[-2]
-    assert usage_chunk["choices"] == []
-    assert usage_chunk["usage"]["prompt_tokens"] == 3
-    assert usage_chunk["usage"]["completion_tokens"] == 2
-    assert usage_chunk["usage"]["total_tokens"] == 5
-
-
-@pytest.mark.asyncio
-async def test_stream_continuous_usage():
-    """continuous_usage_stats adds usage to every data chunk."""
-    engine = _mock_engine()
-
-    async def mock_generate(*args, **kwargs):
-        yield _make_request_output("req-1", token_ids=[10])
-        yield _make_request_output(
-            "req-1", token_ids=[20], finish_reason="stop", finished=True
-        )
-
-    engine.generate = MagicMock(side_effect=mock_generate)
-    serving = _build_serving_tokens(engine)
-
-    request = GenerateRequest(
-        token_ids=[1, 2, 3],
-        sampling_params=SamplingParams(max_tokens=10),
-        model=MODEL_NAME,
-        stream=True,
-        stream_options=StreamOptions(
-            include_usage=True,
-            continuous_usage_stats=True,
-        ),
-    )
-
-    response = await serving.serve_tokens(request)
-    chunks = []
-    async for chunk in response:
-        chunks.append(chunk)
-
-    parsed = _parse_sse_chunks(chunks)
-    data_chunks = [c for c in parsed if isinstance(c, dict) and c.get("choices")]
-
-    # Every data chunk should have usage
-    for i, dc in enumerate(data_chunks):
-        assert dc["usage"] is not None, f"chunk {i} missing usage"
-        assert dc["usage"]["prompt_tokens"] == 3
-
-    # First chunk: 1 completion token
-    assert data_chunks[0]["usage"]["completion_tokens"] == 1
-    assert data_chunks[0]["usage"]["total_tokens"] == 4
-
-    # Second chunk: 2 completion tokens (cumulative)
-    assert data_chunks[1]["usage"]["completion_tokens"] == 2
-    assert data_chunks[1]["usage"]["total_tokens"] == 5
-
-
-@pytest.mark.asyncio
-async def test_stream_with_logprobs():
-    """Streaming with logprobs includes logprob data in each chunk."""
-    engine = _mock_engine()
-
-    async def mock_generate(*args, **kwargs):
-        yield _make_request_output(
-            "req-1",
-            token_ids=[10],
-            logprobs=[{10: Logprob(logprob=-0.5)}],
-        )
-        yield _make_request_output(
-            "req-1",
-            token_ids=[20],
-            logprobs=[{20: Logprob(logprob=-1.0)}],
-            finish_reason="stop",
-            finished=True,
-        )
-
-    engine.generate = MagicMock(side_effect=mock_generate)
-    serving = _build_serving_tokens(engine)
-
-    request = GenerateRequest(
-        token_ids=[1, 2, 3],
-        sampling_params=SamplingParams(max_tokens=10, logprobs=1),
-        model=MODEL_NAME,
-        stream=True,
-    )
-
-    response = await serving.serve_tokens(request)
-    chunks = []
-    async for chunk in response:
-        chunks.append(chunk)
-
-    parsed = _parse_sse_chunks(chunks)
-    data_chunks = [c for c in parsed if isinstance(c, dict) and c.get("choices")]
-
-    for dc in data_chunks:
-        lp = dc["choices"][0]["logprobs"]
-        assert lp is not None
-        assert len(lp["content"]) == 1
-        assert lp["content"][0]["token"].startswith("token_id:")
-
-
-@pytest.mark.asyncio
-async def test_stream_prompt_tokens_details():
-    """enable_prompt_tokens_details includes cached_tokens in final usage."""
-    engine = _mock_engine()
-
-    async def mock_generate(*args, **kwargs):
-        yield _make_request_output(
-            "req-1",
-            token_ids=[10],
-            finish_reason="stop",
-            finished=True,
-            num_cached_tokens=2,
-        )
-
-    engine.generate = MagicMock(side_effect=mock_generate)
-    serving = _build_serving_tokens(engine, enable_prompt_tokens_details=True)
-
-    request = GenerateRequest(
-        token_ids=[1, 2, 3],
-        sampling_params=SamplingParams(max_tokens=10),
-        model=MODEL_NAME,
-        stream=True,
-        stream_options=StreamOptions(include_usage=True),
-    )
-
-    response = await serving.serve_tokens(request)
-    chunks = []
-    async for chunk in response:
-        chunks.append(chunk)
-
-    parsed = _parse_sse_chunks(chunks)
-    # Usage-only chunk (before [DONE])
-    usage_chunk = parsed[-2]
-    assert usage_chunk["choices"] == []
-    assert usage_chunk["usage"]["prompt_tokens_details"]["cached_tokens"] == 2
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

-import json
 import os

 import httpx
@@ -114,54 +113,6 @@ async def test_generate_endpoint(client):
    assert "choices" in data


-@pytest.mark.asyncio
-async def test_generate_stream(client):
-    payload = {
-        "model": MODEL_NAME,
-        "token_ids": [1, 2, 3],
-        "sampling_params": {"max_tokens": 5},
-        "stream": True,
-    }
-    async with client.stream("POST", GEN_ENDPOINT, json=payload) as resp:
-        resp.raise_for_status()
-        chunks = []
-        async for line in resp.aiter_lines():
-            if not line.startswith("data: "):
-                continue
-            payload_str = line[len("data: ") :]
-            if payload_str == "[DONE]":
-                break
-            chunks.append(json.loads(payload_str))
-
-    assert len(chunks) > 0
-    # Every chunk has choices with token_ids
-    all_token_ids = []
-    for chunk in chunks:
-        assert "choices" in chunk
-        assert len(chunk["choices"]) == 1
-        choice = chunk["choices"][0]
-        assert "token_ids" in choice
-        assert len(choice["token_ids"]) > 0
-        all_token_ids.extend(choice["token_ids"])
-
-    # Last chunk should have a finish_reason
-    assert chunks[-1]["choices"][0]["finish_reason"] is not None
-
-    # Streaming should produce the same tokens as non-streaming
-    non_stream_resp = await client.post(
-        GEN_ENDPOINT,
-        json={
-            "model": MODEL_NAME,
-            "token_ids": [1, 2, 3],
-            "sampling_params": {"max_tokens": 5, "temperature": 0.0},
-            "stream": False,
-        },
-    )
-    non_stream_data = non_stream_resp.json()
-    # Just verify we got the right number of tokens
-    assert len(all_token_ids) == len(non_stream_data["choices"][0]["token_ids"])
-
-
@pytest.mark.asyncio
@pytest.mark.parametrize("logprobs_value", [0, 1, 5])
 async def test_generate_logprobs(client, logprobs_value):
@@ -1,8 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-model_name: amd/gpt-oss-20b-w-mxfp4-a-bf16
-metric_threshold: 0.568
-reasoning_effort: low
-server_args: "--attention-backend ROCM_AITER_UNIFIED_ATTN --moe-backend aiter"
-env:
-  VLLM_ROCM_USE_AITER: "1"
@@ -1,6 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-model_name: amd/gpt-oss-20b-w-mxfp4-a-bf16
-metric_threshold: 0.568
-reasoning_effort: low
-server_args: "--attention-backend ROCM_AITER_UNIFIED_ATTN --moe-backend triton"
@@ -1,6 +1,4 @@
 # GFX950 model configurations for GPQA evaluation
 # Tests different environment variable combinations
 gpt-oss-20b-rocm-baseline.yaml
-gpt-oss-20b-rocm-quark-mxfp4-bf16-aiter.yaml
-gpt-oss-20b-rocm-quark-mxfp4-bf16-triton.yaml
-gpt-oss-20b-rocm-quark-mxfp4-fp8-triton.yaml
+gpt-oss-20b-rocm-mxfp4-fp8.yaml
@@ -1,8 +0,0 @@
-model_name: "amd/Qwen3.5-35B-A3B-MXFP4"
-accuracy_threshold: 0.82
-tolerance: 0.03
-num_questions: 1319
-num_fewshot: 5
-server_args: >-
-  --max-model-len 4096
-  --tensor-parallel-size 2
@@ -1,2 +1 @@
 Qwen3.5-35B-A3B-DEP2.yaml
-Qwen3.5-35B-A3B-MXFP4-TP2.yaml
@@ -4,12 +4,7 @@
 import pytest
 import torch

-from vllm._custom_ops import (
-    merge_attn_states as merge_attn_states_cuda,
-)
-from vllm._custom_ops import (
-    scaled_fp8_quant,
-)
+from vllm._custom_ops import merge_attn_states as merge_attn_states_cuda
 from vllm.platforms import current_platform
 from vllm.v1.attention.ops.triton_merge_attn_states import (
    merge_attn_states as merge_attn_states_triton,
@@ -26,7 +21,6 @@ def merge_attn_states_torch(
    suffix_lse: torch.Tensor,  # [NUM_HEADS, NUM_TOKENS]
    output_lse: torch.Tensor | None = None,  # [NUM_HEADS, NUM_TOKENS]
    prefill_tokens_with_context: int | None = None,
-    output_scale: torch.Tensor | None = None,  # scalar, per-tensor FP8 scale
 ):
    # Apply prefill_tokens_with_context mask if needed
    if prefill_tokens_with_context is None:
@@ -55,13 +49,9 @@ def merge_attn_states_torch(
    s_scale = s_lse_exp / out_se  # [NUM_HEADS, NUM_TOKENS]
    p_scale = torch.transpose(p_scale, 0, 1).unsqueeze(2)  # [NUM_TOKENS, NUM_HEADS, 1]
    s_scale = torch.transpose(s_scale, 0, 1).unsqueeze(2)  # [NUM_TOKENS, NUM_HEADS, 1]
-    output = prefix_output * p_scale * mask + suffix_output * (
-        s_scale * mask + (1 - mask)
+    output.copy_(
+        prefix_output * p_scale * mask + suffix_output * (s_scale * mask + (1 - mask))
    )
-    if output_scale is not None:
-        shape = output.shape
-        output, _ = scaled_fp8_quant(output.float().view(-1, shape[-1]), output_scale)
-        output = output.view(shape)
    return output, output_lse


@@ -112,20 +102,18 @@ def generate_markdown_table():
        )


-@pytest.mark.parametrize("use_fp8", [False, True])
@pytest.mark.parametrize("prefill_tokens_with_context", [None, 128])
@pytest.mark.parametrize("num_tokens", NUM_BATCH_TOKENS)
@pytest.mark.parametrize("num_query_heads", NUM_QUERY_HEADS)
@pytest.mark.parametrize("head_size", HEAD_SIZES)
-@pytest.mark.parametrize("input_dtype", DTYPES)
+@pytest.mark.parametrize("output_dtype", DTYPES)
@torch.inference_mode()
 def test_merge_attn_states(
    prefill_tokens_with_context: int | None,
    num_tokens: int,
    num_query_heads: int,
    head_size: int,
-    input_dtype: torch.dtype,
-    use_fp8: bool,
+    output_dtype: torch.dtype,
 ):
    if not current_platform.is_cuda():
        pytest.skip(
@@ -137,18 +125,9 @@ def test_merge_attn_states(
    NUM_HEADS = num_query_heads
    HEAD_SIZE = head_size

-    # When use_fp8 is set, inputs stay as input_dtype (bf16/fp16/fp32)
-    # and output becomes FP8.
-    output_dtype = input_dtype
-    output_scale = None
-    if use_fp8:
-        output_dtype = current_platform.fp8_dtype()
-        output_scale = torch.tensor([0.05], dtype=torch.float32, device="cuda")
-
    print(
        f"\nNUM_TOKENS:{NUM_TOKENS}, NUM_HEADS:{NUM_HEADS}, "
-        f"HEAD_SIZE:{HEAD_SIZE}, input_dtype: {input_dtype}, "
-        f"output_dtype: {output_dtype}, use_fp8: {use_fp8}, "
+        f"HEAD_SIZE:{HEAD_SIZE}, DTYPE: {output_dtype}, "
        f"prefill_tokens_with_context: {prefill_tokens_with_context}, "
        f"Device: {current_platform.get_device_name()}"
    )
@@ -177,10 +156,10 @@ def test_merge_attn_states(
        (NUM_HEADS, NUM_TOKENS), dtype=torch.float32, device="cuda"
    )
    prefix_output = torch.randn(
-        (NUM_TOKENS, NUM_HEADS, HEAD_SIZE), dtype=input_dtype, device="cuda"
+        (NUM_TOKENS, NUM_HEADS, HEAD_SIZE), dtype=output_dtype, device="cuda"
    )
    suffix_output = torch.randn(
-        (NUM_TOKENS, NUM_HEADS, HEAD_SIZE), dtype=input_dtype, device="cuda"
+        (NUM_TOKENS, NUM_HEADS, HEAD_SIZE), dtype=output_dtype, device="cuda"
    )

    warmup_times = 2
@@ -204,7 +183,6 @@ def test_merge_attn_states(
            suffix_lse_torch,
            output_lse_torch,
            prefill_tokens_with_context,
-            output_scale,
        )
    torch.accelerator.synchronize()

@@ -218,7 +196,6 @@ def test_merge_attn_states(
            suffix_lse_torch,
            output_lse_torch,
            prefill_tokens_with_context,
-            output_scale,
        )
        end.record()
        torch.accelerator.synchronize()
@@ -243,7 +220,6 @@ def test_merge_attn_states(
            suffix_lse,
            output_lse_ref_triton,
            prefill_tokens_with_context,
-            output_scale,
        )
    torch.accelerator.synchronize()

@@ -257,7 +233,6 @@ def test_merge_attn_states(
            suffix_lse,
            output_lse_ref_triton,
            prefill_tokens_with_context,
-            output_scale,
        )
        end.record()
        torch.accelerator.synchronize()
@@ -279,7 +254,6 @@ def test_merge_attn_states(
            suffix_lse,
            output_lse_cuda,
            prefill_tokens_with_context,
-            output_scale,
        )
    torch.accelerator.synchronize()

@@ -293,7 +267,6 @@ def test_merge_attn_states(
            suffix_lse,
            output_lse_cuda,
            prefill_tokens_with_context,
-            output_scale,
        )
        end.record()
        torch.accelerator.synchronize()
@@ -315,19 +288,7 @@ def test_merge_attn_states(
    # Liger Kernel: Efficient Triton Kernels for LLM Training
    # https://arxiv.org/pdf/2410.10989, 3.3 Correctness
    # use rtol = 1e-2 for bfloat16.
-    if use_fp8:
-        # Compare in dequantized space (multiply back by scale) so that
-        # absolute differences reflect real precision, not amplified FP8
-        # quantization steps.
-        atol, rtol = 1e-1, 1e-1
-        assert output_scale is not None
-        scale = output_scale.item()
-    elif output_dtype == torch.bfloat16:
-        atol, rtol = 1e-3, 1e-2
-        scale = 1.0
-    else:
-        atol, rtol = 1e-3, 1e-3
-        scale = 1.0
+    rtol = 1e-2 if output_dtype == torch.bfloat16 else 1e-3

    def diff(a: torch.Tensor, b: torch.Tensor):
        max_diff = torch.max(torch.abs(a.float() - b.float()))
@@ -339,26 +300,16 @@ def test_merge_attn_states(
    output_ref = output_ref_triton
    output_lse_ref = output_lse_ref_triton
    torch.testing.assert_close(
-        output_cuda.float() * scale,
-        output_ref.float() * scale,
-        atol=atol,
-        rtol=rtol,
+        output_cuda.float(), output_ref.float(), atol=1e-3, rtol=rtol
    )
-    print(
-        "Output all match, max abs diff (dequantized):"
-        if use_fp8
-        else "Output all match, max abs diff:"
-    )
-    _diff = diff(output_ref.float() * scale, output_torch.float() * scale)
-    print(f"(Triton vs Torch) : {_diff}")
-    _diff = diff(output_torch.float() * scale, output_cuda.float() * scale)
-    print(f"  (CUDA vs Torch) : {_diff}")
-    _diff = diff(output_ref.float() * scale, output_cuda.float() * scale)
-    print(f"  (CUDA vs Triton): {_diff}")
+    print("Output all match, max abs diff:")
+    print(f"(Triton vs Torch) : {diff(output_torch, output_ref)}")
+    print(f"  (CUDA vs Torch) : {diff(output_torch, output_cuda)}")
+    print(f"  (CUDA vs Triton): {diff(output_ref, output_cuda)}")
    print("-" * 100)

    torch.testing.assert_close(
-        output_lse_cuda.float(), output_lse_ref.float(), atol=atol, rtol=rtol
+        output_lse_cuda.float(), output_lse_ref.float(), atol=1e-3, rtol=rtol
    )
    print("Output LSE all match, max abs diff:")
    print(f"(Triton vs Torch) : {diff(output_lse_torch, output_lse_ref)}")
@@ -26,59 +26,6 @@ MINIMUM_TORCH_VERSION = version.parse("2.7.0")
 DIRECT_BUILD_VERSION = version.parse("2.9.dev0")


-@pytest.mark.skipif(
-    not torch.cuda.is_available() or TORCH_VERSION < MINIMUM_TORCH_VERSION,
-    reason="CUDA not available or PyTorch version < 2.7",
-)
-def test_flex_attention_full_cudagraphs(vllm_runner):
-    """Test the numerics for flex attention full cudagraphs support."""
-    model_name = "Qwen/Qwen2.5-1.5B-Instruct"
-    seed = 42
-    max_tokens = 24
-    num_logprobs = 5
-    prompts = [
-        "Hello, my name is",
-        "The president of the United States is",
-        "The capital of France is",
-    ]
-
-    # Run with flex attention eager
-    set_random_seed(seed)
-    with vllm_runner(
-        model_name,
-        runner="generate",
-        tensor_parallel_size=1,
-        num_gpu_blocks_override=128,
-        enforce_eager=True,
-        attention_config={"backend": "FLEX_ATTENTION"},
-    ) as llm_flex:
-        output_eager = llm_flex.generate_greedy_logprobs(
-            prompts, max_tokens, num_logprobs
-        )
-
-    # Run with flex attention compiled
-    set_random_seed(seed)
-    with vllm_runner(
-        model_name,
-        runner="generate",
-        tensor_parallel_size=1,
-        num_gpu_blocks_override=128,
-        enforce_eager=False,
-        gpu_memory_utilization=0.85,
-        attention_config={"backend": "FLEX_ATTENTION"},
-    ) as llm_default:
-        output_compile = llm_default.generate_greedy_logprobs(
-            prompts, max_tokens, num_logprobs
-        )
-
-    check_logprobs_close(
-        outputs_0_lst=output_eager,
-        outputs_1_lst=output_compile,
-        name_0="eager",
-        name_1="compile",
-    )
-
-
@pytest.mark.skipif(
    not torch.cuda.is_available() or TORCH_VERSION < MINIMUM_TORCH_VERSION,
    reason="CUDA not available or PyTorch version < 2.7",
@@ -637,7 +637,7 @@ def use_fused_moe_lora_kernel_tensor_parallel(

    set_random_seed(seed)

-    device = torch.device(f"{DEVICE_TYPE}:{local_rank}")
+    device = torch.device(f"cuda:{local_rank}")
    torch.accelerator.set_device_index(device)
    torch.set_default_device(device)
    torch.set_default_dtype(dtype)
@@ -60,12 +60,8 @@ pytestmark = pytest.mark.skipif(
    reason="Backend not supported",
 )

-DEVICE_TYPE = current_platform.device_type
 DEVICES = (
-    [
-        f"{DEVICE_TYPE}:{i}"
-        for i in range(1 if torch.accelerator.device_count() == 1 else 2)
-    ]
+    [f"cuda:{i}" for i in range(1 if torch.accelerator.device_count() == 1 else 2)]
    if current_platform.is_cuda_alike()
    else ["cpu"]
 )
@@ -200,7 +196,7 @@ def create_random_inputs(
    input_size: tuple[int, ...],
    input_range: tuple[float, float],
    input_type: torch.dtype = torch.int,
-    device: torch.device = DEVICE_TYPE,
+    device: torch.device = "cuda",
 ) -> tuple[list[torch.Tensor], list[int], list[int]]:
    """Creates random inputs.

@@ -35,9 +35,9 @@ EMBEDDING_MODULES = {
    "lm_head": "output_embeddings",
 }

-DEVICE_TYPE = current_platform.device_type
+
 DEVICES = (
-    [f"{DEVICE_TYPE}:{i}" for i in range(min(torch.accelerator.device_count(), 2))]
+    [f"cuda:{i}" for i in range(1 if torch.accelerator.device_count() == 1 else 2)]
    if current_platform.is_cuda_alike()
    else ["cpu"]
 )
@@ -6,9 +6,6 @@ import pytest
 import torch

 from vllm import _custom_ops as ops
-from vllm.platforms import current_platform
-
-DEVICE_TYPE = current_platform.device_type


 def round_up(x, base):
@@ -30,7 +27,7 @@ def sample_data(num_experts, max_loras, num_tokens, topk_num):
            topk_ids[i, j] = pool[j]
        token_lora_mapping[i] = random.randint(0, max_loras - 1)

-    return topk_ids.to(DEVICE_TYPE), token_lora_mapping.to(DEVICE_TYPE)
+    return topk_ids.to("cuda"), token_lora_mapping.to("cuda")


@pytest.mark.parametrize("num_tokens", [100, 200, 1024, 4096])  # 81920
@@ -59,21 +56,14 @@ def test_moe_lora_align_block_size(
        (max_loras * max_num_tokens_padded,),
        topk_ids.numel(),
        dtype=torch.int32,
-        device=DEVICE_TYPE,
+        device="cuda",
    )
    expert_ids = torch.full(
-        (max_loras * max_num_m_blocks,),
-        num_experts,
-        dtype=torch.int32,
-        device=DEVICE_TYPE,
+        (max_loras * max_num_m_blocks,), num_experts, dtype=torch.int32, device="cuda"
    )
-    num_tokens_post_pad = torch.zeros(
-        (max_loras,), dtype=torch.int32, device=DEVICE_TYPE
-    )
-    adapter_enabled = torch.ones(
-        (max_loras + 1,), dtype=torch.int32, device=DEVICE_TYPE
-    )
-    lora_ids = torch.arange(max_loras + 2, dtype=torch.int32, device=DEVICE_TYPE)
+    num_tokens_post_pad = torch.zeros((max_loras,), dtype=torch.int32, device="cuda")
+    adapter_enabled = torch.ones((max_loras + 1,), dtype=torch.int32, device="cuda")
+    lora_ids = torch.arange(max_loras + 2, dtype=torch.int32, device="cuda")

    # call kernel
    ops.moe_lora_align_block_size(
@@ -9,13 +9,10 @@ import vllm.lora.ops.torch_ops as torch_ops
 import vllm.lora.ops.triton_ops as triton_ops
 from vllm.lora.ops.triton_ops import LoRAKernelMeta
 from vllm.lora.ops.triton_ops.utils import _LORA_A_PTR_DICT, _LORA_B_PTR_DICT
-from vllm.platforms import current_platform
 from vllm.utils.torch_utils import set_random_seed

 from .utils import PunicaTensors, assert_close, generate_data_for_nslices

-DEVICE_TYPE = current_platform.device_type
-

@pytest.fixture(autouse=True)
 def reset_device(reset_default_device):
@@ -149,9 +146,7 @@ def check_lora_shrink_kernel(

    # Setup metadata information for the LoRA kernel.
    lora_meta = LoRAKernelMeta.make(
-        max_loras=num_loras,
-        max_num_tokens=token_nums,
-        device=DEVICE_TYPE,
+        max_loras=num_loras, max_num_tokens=token_nums, device="cuda"
    )
    lora_meta.prepare_tensors(data.token_lora_mapping)

@@ -224,9 +219,7 @@ def check_lora_expand_kernel(

    # Setup metadata information for the LoRA kernel.
    lora_meta = LoRAKernelMeta.make(
-        max_loras=num_loras,
-        max_num_tokens=token_nums,
-        device=DEVICE_TYPE,
+        max_loras=num_loras, max_num_tokens=token_nums, device="cuda"
    )
    lora_meta.prepare_tensors(data.token_lora_mapping)

@@ -374,7 +367,7 @@ test_params = {
 }

 DTYPES = [torch.float16, torch.bfloat16]
-DEVICES = [f"{DEVICE_TYPE}:{0}"]
+DEVICES = [f"cuda:{0}"]
 SEED = [0]


@@ -28,11 +28,9 @@ from vllm.lora.ops.triton_ops.lora_shrink_fp8_op import (
    _SHRINK_LORA_SCALE_PTR_DICT,
 )
 from vllm.lora.ops.triton_ops.utils import _LORA_A_PTR_DICT, _LORA_B_PTR_DICT
-from vllm.platforms import current_platform
 from vllm.utils.torch_utils import set_random_seed

-DEVICE_TYPE = current_platform.device_type
-DEVICES = [f"{DEVICE_TYPE}:{0}"]
+DEVICES = [f"cuda:{0}"]
 SEED = [0]

 _dict_lock = Lock()
@@ -19,14 +19,11 @@ from vllm.config.load import LoadConfig
 from vllm.config.lora import LoRAConfig
 from vllm.lora.model_manager import LoRAMapping
 from vllm.lora.request import LoRARequest
-from vllm.platforms import current_platform
 from vllm.v1.worker.gpu_worker import Worker

 MODEL_PATH = "Qwen/Qwen3-0.6B"
 NUM_LORAS = 16

-DEVICE_TYPE = current_platform.device_type
-

@patch.dict(os.environ, {"RANK": "0"})
 def test_worker_apply_lora(qwen3_lora_files):
@@ -64,7 +61,7 @@ def test_worker_apply_lora(qwen3_lora_files):
            max_num_seqs=32,
            max_num_partial_prefills=32,
        ),
-        device_config=DeviceConfig(DEVICE_TYPE),
+        device_config=DeviceConfig("cuda"),
        cache_config=CacheConfig(
            block_size=16,
            cache_dtype="auto",
@@ -9,13 +9,10 @@ import torch
 from safetensors.torch import save_file

 from vllm.lora.lora_weights import LoRALayerWeights, PackedLoRALayerWeights
-from vllm.platforms import current_platform
-
-DEVICE_TYPE = current_platform.device_type


 class DummyLoRAManager:
-    def __init__(self, device: torch.device = f"{DEVICE_TYPE}:0"):
+    def __init__(self, device: torch.device = "cuda:0"):
        super().__init__()
        self._loras: dict[str, LoRALayerWeights] = {}
        self._device = device
@@ -60,8 +57,8 @@ class DummyLoRAManager:
            module_name,
            rank=rank,
            lora_alpha=1,
-            lora_a=torch.rand([rank, input_dim], device=DEVICE_TYPE),
-            lora_b=torch.rand([output_dim, input_dim], device=DEVICE_TYPE),
+            lora_a=torch.rand([rank, input_dim], device="cuda"),
+            lora_b=torch.rand([output_dim, input_dim], device="cuda"),
            embeddings_tensor=embeddings_tensor,
        )
        self.set_module_lora(module_name, lora)
@@ -60,14 +60,6 @@ MAX_NUM_SEQS = 4
 ATTN_BACKEND = "TRITON_ATTN" if current_platform.is_rocm() else "auto"


-def _set_conv_state_layout(monkeypatch, layout: str) -> None:
-    """Set conv state layout env var and clear cache to pick up new value."""
-    from vllm.model_executor.layers.mamba import mamba_utils
-
-    monkeypatch.setenv("VLLM_SSM_CONV_STATE_LAYOUT", layout)
-    mamba_utils.get_conv_state_layout.cache_clear()
-
-
@pytest.mark.parametrize("model", SSM_MODELS + HYBRID_MODELS)
@pytest.mark.parametrize("max_tokens", [64])
@pytest.mark.parametrize("num_logprobs", [5])
@@ -110,15 +102,12 @@ def test_models(
@pytest.mark.parametrize("model", [SSM_MODELS[0], HYBRID_MODELS[0]])
@pytest.mark.parametrize("max_tokens", [64])
@pytest.mark.parametrize("num_logprobs", [5])
-@pytest.mark.parametrize("conv_state_layout", ["SD", "DS"])
 def test_batching(
    vllm_runner,
    example_prompts,
-    monkeypatch,
    model: str,
    max_tokens: int,
    num_logprobs: int,
-    conv_state_layout: str,
 ) -> None:
    try:
        model_info = HF_EXAMPLE_MODELS.find_hf_info(model)
@@ -127,8 +116,6 @@ def test_batching(
    except ValueError:
        pass

-    _set_conv_state_layout(monkeypatch, conv_state_layout)
-
    for_loop_outputs = []
    with vllm_runner(model, max_num_seqs=MAX_NUM_SEQS) as vllm_model:
        for prompt in example_prompts:
@@ -151,14 +138,11 @@ def test_batching(

@pytest.mark.parametrize("model", [SSM_MODELS[0], HYBRID_MODELS[0]])
@pytest.mark.parametrize("max_tokens", [10])
-@pytest.mark.parametrize("conv_state_layout", ["SD", "DS"])
 def test_chunked_prefill_with_parallel_sampling(
    vllm_runner,
    example_prompts,
-    monkeypatch,
    model: str,
    max_tokens: int,
-    conv_state_layout: str,
 ) -> None:
    """
    Tests chunked prefill in conjunction with n > 1.
@@ -170,8 +154,6 @@ def test_chunked_prefill_with_parallel_sampling(
    decoding steps inside a chunked prefill forward pass
    (where we have both prefill and decode together)
    """
-    _set_conv_state_layout(monkeypatch, conv_state_layout)
-
    sampling_params = SamplingParams(n=3, temperature=1, seed=0, max_tokens=max_tokens)
    with vllm_runner(
        model,
@@ -186,22 +168,17 @@ def test_chunked_prefill_with_parallel_sampling(

@pytest.mark.parametrize("model", [SSM_MODELS[0], HYBRID_MODELS[0]])
@pytest.mark.parametrize("max_tokens", [20])
-@pytest.mark.parametrize("conv_state_layout", ["SD", "DS"])
 def test_mamba_cache_cg_padding(
    vllm_runner,
    example_prompts,
-    monkeypatch,
    model: str,
    max_tokens: int,
-    conv_state_layout: str,
 ) -> None:
    """
    This test is for verifying that mamba cache is padded to CG captured
    batch size. If it's not, a torch RuntimeError will be raised because
    tensor dimensions aren't compatible.
    """
-    _set_conv_state_layout(monkeypatch, conv_state_layout)
-
    vllm_config = EngineArgs(model=model, trust_remote_code=True).create_engine_config()
    cudagraph_dispatcher = CudagraphDispatcher(vllm_config)
    cudagraph_dispatcher.initialize_cudagraph_keys(
@@ -1,187 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-from collections.abc import Sequence
-
-import pytest
-import regex as re
-from transformers import AutoModelForCausalLM, AutoTokenizer
-
-from vllm.logprobs import SampleLogprobs
-from vllm.multimodal.image import rescale_image_size
-
-from ....conftest import (
-    IMAGE_ASSETS,
-    HfRunner,
-    PromptImageInput,
-    VllmRunner,
-)
-from ....utils import multi_gpu_test
-from ...utils import check_logprobs_close
-
-MODEL_ID = "microsoft/Phi-4-reasoning-vision-15B"
-
-HF_IMAGE_PROMPTS = IMAGE_ASSETS.prompts(
-    {
-        "stop_sign": "<|user|>\n<image>\nWhat's the content of the image?<|end|>\n<|assistant|>\n",  # noqa: E501
-        "cherry_blossom": "<|user|>\n<image>\nPlease infer the season with reason in details.<|end|>\n<|assistant|>\n",  # noqa: E501
-    }
-)
-HF_MULTIIMAGE_IMAGE_PROMPT = (
-    "<|user|>\n<image>\n<image>\nDescribe these images.<|end|>\n<|assistant|>\n"  # noqa: E501
-)
-
-DTYPE = "half"
-MAX_TOKENS = 128
-NUM_LOGPROBS = 10
-
-
-def vllm_to_hf_output(
-    vllm_output: tuple[list[int], str, SampleLogprobs | None], model: str
-):
-    """Sanitize vllm output to be comparable with hf output."""
-    _, output_str, out_logprobs = vllm_output
-
-    output_str_without_image = re.sub(r"(<image>)+", "", output_str)
-    if output_str_without_image and output_str_without_image[0] == " ":
-        output_str_without_image = output_str_without_image[1:]
-
-    hf_output_str = output_str_without_image + "<|end|><|endoftext|>"
-
-    tokenizer = AutoTokenizer.from_pretrained(model, trust_remote_code=True)
-    hf_output_ids = tokenizer.encode(output_str_without_image)
-    if hf_output_ids and hf_output_ids[0] == tokenizer.bos_token_id:
-        hf_output_ids = hf_output_ids[1:]
-
-    return hf_output_ids, hf_output_str, out_logprobs
-
-
-def _build_single_image_inputs(
-    image_assets,
-) -> list[tuple[list[str], PromptImageInput]]:
-    """Build single-image inputs for all size_factors at once."""
-    images = [asset.pil_image for asset in image_assets]
-    all_inputs: list[tuple[list[str], PromptImageInput]] = []
-    for size_factors in [[1.0], [0.25, 0.5, 1.0]]:
-        for image, prompt in zip(images, HF_IMAGE_PROMPTS):
-            all_inputs.append(
-                (
-                    [prompt for _ in size_factors],
-                    [rescale_image_size(image, f) for f in size_factors],
-                )
-            )
-    return all_inputs
-
-
-def _build_multi_image_inputs(
-    image_assets,
-) -> list[tuple[list[str], PromptImageInput]]:
-    """Build multi-image inputs for all size_factors at once."""
-    images = [asset.pil_image for asset in image_assets]
-    all_inputs: list[tuple[list[str], PromptImageInput]] = []
-    for size_factors in [[0.5], [0.15, 0.30]]:
-        all_inputs.append(
-            (
-                [HF_MULTIIMAGE_IMAGE_PROMPT for _ in size_factors],
-                [
-                    [rescale_image_size(image, factor) for image in images]
-                    for factor in size_factors
-                ],
-            )
-        )
-    return all_inputs
-
-
-def _run_and_compare(
-    hf_runner: type[HfRunner],
-    vllm_runner: type[VllmRunner],
-    all_inputs: Sequence[tuple[list[str], PromptImageInput]],
-    model: str,
-    max_model_len: int,
-    max_num_seqs: int,
-    mm_limit: int,
-    gpu_memory_utilization: float,
-):
-    """Load each runner once, run all inputs, then compare."""
-    # NOTE: run vLLM first, then HF.  vLLM needs a fresh process without
-    # cuda initialization; running HF first would break the multiprocessing
-    # backend with fork method.
-    with vllm_runner(
-        model,
-        runner="generate",
-        max_model_len=max_model_len,
-        max_num_seqs=max_num_seqs,
-        gpu_memory_utilization=gpu_memory_utilization,
-        dtype=DTYPE,
-        limit_mm_per_prompt={"image": mm_limit},
-        tensor_parallel_size=2,
-        trust_remote_code=True,
-        enforce_eager=True,
-    ) as vllm_model:
-        vllm_outputs_per_case = [
-            vllm_model.generate_greedy_logprobs(
-                prompts,
-                MAX_TOKENS,
-                num_logprobs=NUM_LOGPROBS,
-                images=images,
-            )
-            for prompts, images in all_inputs
-        ]
-
-    hf_model_kwargs = {"_attn_implementation": "sdpa", "device_map": "auto"}
-    with hf_runner(
-        model,
-        dtype=DTYPE,
-        model_kwargs=hf_model_kwargs,
-        auto_cls=AutoModelForCausalLM,
-        trust_remote_code=True,
-    ) as hf_model:
-        hf_outputs_per_case = [
-            hf_model.generate_greedy_logprobs_limit(
-                prompts,
-                MAX_TOKENS,
-                num_logprobs=NUM_LOGPROBS,
-                images=images,
-            )
-            for prompts, images in all_inputs
-        ]
-
-    for hf_outputs, vllm_outputs in zip(hf_outputs_per_case, vllm_outputs_per_case):
-        check_logprobs_close(
-            outputs_0_lst=hf_outputs,
-            outputs_1_lst=vllm_outputs,
-            name_0="hf",
-            name_1="vllm",
-        )
-
-
-@multi_gpu_test(num_gpus=2)
-@pytest.mark.parametrize("model", [MODEL_ID])
-def test_models(hf_runner, vllm_runner, image_assets, model) -> None:
-    all_inputs = _build_single_image_inputs(image_assets)
-    _run_and_compare(
-        hf_runner,
-        vllm_runner,
-        all_inputs,
-        model,
-        max_model_len=8192,
-        max_num_seqs=2,
-        mm_limit=1,
-        gpu_memory_utilization=0.80,
-    )
-
-
-@multi_gpu_test(num_gpus=2)
-@pytest.mark.parametrize("model", [MODEL_ID])
-def test_multi_images_models(hf_runner, vllm_runner, image_assets, model) -> None:
-    all_inputs = _build_multi_image_inputs(image_assets)
-    _run_and_compare(
-        hf_runner,
-        vllm_runner,
-        all_inputs,
-        model,
-        max_model_len=8192,
-        max_num_seqs=2,
-        mm_limit=2,
-        gpu_memory_utilization=0.80,
-    )
@@ -7,7 +7,6 @@ from typing import Any, Literal

 import pytest
 from packaging.version import Version
-from transformers import PretrainedConfig
 from transformers import __version__ as TRANSFORMERS_VERSION

 from vllm.config.model import ModelDType, TokenizerMode
@@ -538,9 +537,6 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
    "TeleChat2ForCausalLM": _HfExamplesInfo(
        "Tele-AI/TeleChat2-3B", trust_remote_code=True
    ),
-    "TeleChat3ForCausalLM": _HfExamplesInfo(
-        "Tele-AI/TeleChat3-36B-Thinking", trust_remote_code=True
-    ),
    "TeleFLMForCausalLM": _HfExamplesInfo(
        "CofeAI/FLM-2-52B-Instruct-2407", trust_remote_code=True
    ),
@@ -1005,26 +1001,7 @@ _MULTIMODAL_EXAMPLE_MODELS = {
        trust_remote_code=True,
    ),
    "NemotronH_Nano_VL_V2": _HfExamplesInfo(
-        "nvidia/NVIDIA-Nemotron-Nano-12B-v2-VL-BF16",
-        max_model_len=4096,
-        # NemotronH layers are constructed via `hybrid_override_pattern`:
-        use_original_num_layers=True,
-        hf_overrides={
-            "vision_config": PretrainedConfig(
-                args={
-                    "min_num_patches": 1,  # Trigger image dynamic res
-                    "max_num_patches": 12,
-                    "model": "vit_huge_patch16_224",
-                },
-                # Trigger conv3d:
-                video_temporal_patch_size=2,
-            ),
-            "text_config": {
-                "num_hidden_layers": 2,
-                "hybrid_override_pattern": "M*",
-            },
-        },
-        trust_remote_code=True,
+        "nano_vl_dummy", is_available_online=False, trust_remote_code=True
    ),
    "OpenCUAForConditionalGeneration": _HfExamplesInfo(
        "xlangai/OpenCUA-7B", trust_remote_code=True
@@ -1069,9 +1046,6 @@ _MULTIMODAL_EXAMPLE_MODELS = {
        },  # noqa: E501
        extras={"phi3.5": "microsoft/Phi-3.5-vision-instruct"},
    ),
-    "Phi4ForCausalLMV": _HfExamplesInfo(
-        "microsoft/Phi-4-reasoning-vision-15B", trust_remote_code=True
-    ),
    "Phi4MMForCausalLM": _HfExamplesInfo(
        "microsoft/Phi-4-multimodal-instruct", trust_remote_code=True
    ),
@@ -447,16 +447,9 @@ def dummy_hf_overrides(
    Dummy HF overrides function used to create dummy model
    with only minimum nums of layer.
    """
-    # Copy because this helper is called more than once
-    # while loading config, and we `.pop()`
-    exist_overrides = (exist_overrides or {}).copy()
-    text_config_override = exist_overrides.pop("text_config", None)
-    hf_config.update(exist_overrides)
+    hf_config.update(exist_overrides or {})

    text_config = hf_config.get_text_config()
-    if text_config_override is not None:
-        # multimodal test models may override *some* text-model fields
-        text_config.update(text_config_override)

    # Ensure at least 2 expert per group
    # Since `grouped_topk` assumes top-2
@@ -40,8 +40,6 @@ BACKENDS_TO_TEST = [
    "FLEX_ATTENTION_SLOW",
 ]

-DEVICE_TYPE = current_platform.device_type
-
 # Remove flashinfer from the list if it's not available
 try:
    import flashinfer  # noqa: F401
@@ -368,7 +366,7 @@ def _test_backend_correctness(
        num_gpu_blocks=8192,
        hf_config_override=hf_config_override,
    )
-    device = torch.device(f"{DEVICE_TYPE}:0")
+    device = torch.device("cuda:0")

    kv_cache_spec = create_standard_kv_cache_spec(vllm_config)

@@ -7,7 +7,6 @@ import pytest
 import torch

 from tests.v1.attention.utils import BatchSpec, create_common_attn_metadata
-from vllm.platforms import current_platform
 from vllm.v1.attention.backends.utils import make_local_attention_virtual_batches


@@ -23,8 +22,6 @@ class LocalAttentionTestData:
    expected_local_block_table: list[list[int]]


-DEVICE_TYPE = current_platform.device_type
-
 test_data_list = [
    # Same as example in docstring of make_local_attention_virtual_batches
    # except block table has 9 columns instead of 10
@@ -154,7 +151,7 @@ test_data_list = [

@pytest.mark.parametrize("test_data", test_data_list)
 def test_local_attention_virtual_batches(test_data: LocalAttentionTestData):
-    device = torch.device(f"{DEVICE_TYPE}:0")
+    device = torch.device("cuda:0")
    batch_spec = test_data.batch_spec
    attn_chunk_size = test_data.attn_chunk_size
    block_size = test_data.block_size
@@ -42,8 +42,6 @@ BACKENDS_TO_TEST = [
    AttentionBackendEnum.TRITON_MLA,
 ]

-DEVICE_TYPE = current_platform.device_type
-
 # Remove sm100 backends from the list if not using sm100
 if not torch.cuda.is_available() or torch.cuda.get_device_properties(0).major < 10:
    BACKENDS_TO_TEST.remove(AttentionBackendEnum.CUTLASS_MLA)
@@ -765,7 +763,7 @@ def test_backend_correctness(
            method="ngram", num_speculative_tokens=query_len - 1
        )

-    device = torch.device(f"{DEVICE_TYPE}:0")
+    device = torch.device("cuda:0")

    # 1. Setup
    batch_size = batch_spec.batch_size
@@ -64,8 +64,6 @@ SPARSE_BACKEND_BATCH_SPECS["large_q_pure_prefill"] = BatchSpec(
    seq_lens=[256] * 2, query_lens=[256] * 2
 )

-DEVICE_TYPE = current_platform.device_type
-

 def _float_to_e8m0_truncate(f: float) -> float:
    """Simulate SM100's float -> e8m0 -> bf16 scale conversion.
@@ -224,7 +222,7 @@ def test_sparse_backend_decode_correctness(
    batch_spec = SPARSE_BACKEND_BATCH_SPECS[batch_name]
    use_fp8_ds_mla_quantization = kv_cache_dtype == "fp8_ds_mla"

-    device = torch.device(DEVICE_TYPE)
+    device = torch.device("cuda")
    dtype = torch.bfloat16

    # Model hyper-parameters (kept intentionally small for the unit test)
@@ -588,7 +586,7 @@ def _triton_convert_reference_impl(
 def test_triton_convert_req_index_to_global_index_decode_only(
    block_size, num_topk_tokens
 ):
-    device = torch.device(DEVICE_TYPE)
+    device = torch.device("cuda")
    num_tokens = 8
    num_requests = 4
    max_blocks_per_req = 10
@@ -641,7 +639,7 @@ def test_triton_convert_req_index_to_global_index_decode_only(
    reason="FlashMLASparseBackend requires CUDA 9.0 or higher",
 )
 def test_triton_convert_req_index_to_global_index_with_prefill_workspace(block_size):
-    device = torch.device(DEVICE_TYPE)
+    device = torch.device("cuda")
    num_requests = 4
    max_blocks_per_req = 8
    num_topk_tokens = 128
@@ -796,7 +794,7 @@ def test_split_indexer_prefill_chunks_single_request_overflow():

 def test_triton_convert_returns_valid_counts():
    """Test that return_valid_counts correctly counts non-negative indices."""
-    device = torch.device(DEVICE_TYPE)
+    device = torch.device("cuda")
    num_tokens = 8
    num_requests = 2
    max_blocks_per_req = 10
@@ -55,7 +55,6 @@ class MockAttentionLayer:
 MODEL = "Qwen/Qwen2.5-0.5B"
 BLOCK_SIZE = 16
 NUM_GPU_BLOCKS = 8192
-DEVICE_TYPE = current_platform.device_type

 BATCH_SPECS = {
    "decode_only": BatchSpec(
@@ -173,7 +172,7 @@ def _run_trtllm_integration(batch_spec):
    """Run TRTLLM attention through the full FlashInfer pipeline
    and compare against an SDPA reference."""
    set_random_seed(42)
-    device = torch.device(f"{DEVICE_TYPE}:0")
+    device = torch.device("cuda:0")

    vllm_config = create_vllm_config(
        model_name=MODEL,
@@ -23,8 +23,6 @@ from vllm.forward_context import BatchDescriptor, set_forward_context
 from vllm.platforms import current_platform
 from vllm.v1.cudagraph_dispatcher import CudagraphDispatcher

-DEVICE_TYPE = current_platform.device_type
-

 # Helper MLP for testing
 class SimpleMLP(nn.Module):
@@ -271,9 +269,9 @@ class TestCudagraphDispatcher:
 class TestCUDAGraphWrapper:
    def setup_method(self):
        self.vllm_config = _create_vllm_config(CompilationConfig())
-        self.model = SimpleMLP().to(DEVICE_TYPE)
-        self.persistent_input_buffer = torch.zeros(1, 10, device=DEVICE_TYPE)
-        self.input_tensor = torch.randn(1, 10, device=DEVICE_TYPE)
+        self.model = SimpleMLP().to("cuda")
+        self.persistent_input_buffer = torch.zeros(1, 10, device="cuda")
+        self.input_tensor = torch.randn(1, 10, device="cuda")

    def test_capture_and_replay(self):
        wrapper = CUDAGraphWrapper(
@@ -430,10 +428,10 @@ class TestCudagraphIntegration:

    @create_new_process_for_each_test("spawn")
    def test_capture_replay_bypass_logic(self):
-        model = SimpleMLP().to(DEVICE_TYPE)
+        model = SimpleMLP().to("cuda")
        full_wrapper = CUDAGraphWrapper(model, self.vllm_config, CUDAGraphMode.FULL)
        max_bs = 16
-        persistent_input_buffer = torch.zeros(max_bs, 10, device=DEVICE_TYPE)
+        persistent_input_buffer = torch.zeros(max_bs, 10, device="cuda")
        input_1 = persistent_input_buffer[:1]
        input_2 = persistent_input_buffer[:2]
        input_3 = persistent_input_buffer[:3]
@@ -488,17 +486,17 @@ class TestCudagraphIntegration:
    @create_new_process_for_each_test("spawn")
    def test_nested_wrappers(self):
        """Tests a scenario with a PIECEWISE wrapper inside a FULL one."""
-        model = SimpleMLP().to(DEVICE_TYPE)
+        model = SimpleMLP().to("cuda")
        full_wrapper = CUDAGraphWrapper(model, self.vllm_config, CUDAGraphMode.FULL)
-        input_1 = torch.randn(1, 10, device=DEVICE_TYPE)
+        input_1 = torch.randn(1, 10, device="cuda")

        # Setup: Inner model is wrapped with PIECEWISE, outer with FULL
-        inner_model = SimpleMLP().to(DEVICE_TYPE)
+        inner_model = SimpleMLP().to("cuda")
        piecewise_wrapper = CUDAGraphWrapper(
            inner_model, self.vllm_config, CUDAGraphMode.PIECEWISE
        )
        inner_model.forward = MagicMock(wraps=inner_model.forward)
-        outer_model = SimpleMLP().to(DEVICE_TYPE)
+        outer_model = SimpleMLP().to("cuda")
        # When outer model is called, it calls the piecewise_wrapper
        outer_model.forward = MagicMock(
            wraps=outer_model.forward, side_effect=piecewise_wrapper
@@ -187,7 +187,7 @@ def test_logprobs_bitwise_batch_invariance_bs1_vs_bsN(
        tensor_parallel_size=tp_size,
        max_num_seqs=128,
        max_model_len=8192,
-        dtype="auto",  # not everything is supported
+        dtype="bfloat16",  # not everything is supported
        gpu_memory_utilization=0.9,
        enforce_eager=IS_DEVICE_CAPABILITY_BELOW_90,
        attention_config={"backend": backend},
@@ -400,7 +400,7 @@ def test_simple_generation(backend):
        tensor_parallel_size=int(os.getenv("VLLM_TP_SIZE", "1")),
        gpu_memory_utilization=0.9,
        max_model_len=2048,
-        dtype="auto",
+        dtype="bfloat16",
        enable_prefix_caching=False,
        enforce_eager=IS_DEVICE_CAPABILITY_BELOW_90,
        attention_config={"backend": backend},
@@ -466,7 +466,7 @@ def test_logprobs_without_batch_invariance_should_fail(
        tensor_parallel_size=tp_size,
        max_num_seqs=32,
        max_model_len=8192,
-        dtype="auto",
+        dtype="bfloat16",
        enforce_eager=IS_DEVICE_CAPABILITY_BELOW_90,
        attention_config={"backend": backend},
    )
@@ -686,7 +686,7 @@ def test_decode_logprobs_match_prefill_logprobs(
        tensor_parallel_size=tp_size,
        max_num_seqs=32,
        max_model_len=8192,
-        dtype="auto",
+        dtype="bfloat16",
        enforce_eager=IS_DEVICE_CAPABILITY_BELOW_90,
        attention_config={"backend": backend},
    )
@@ -931,7 +931,7 @@ def LLM_with_max_seqs(
        max_num_seqs=max_num_seqs,
        gpu_memory_utilization=gpu_memory_utilization,
        max_model_len=max_model_len,
-        dtype="auto",
+        dtype="bfloat16",
        tensor_parallel_size=int(os.getenv("VLLM_TP_SIZE", "1")),
        enable_prefix_caching=False,
        enforce_eager=IS_DEVICE_CAPABILITY_BELOW_90,
@@ -13,9 +13,6 @@ from utils import skip_unsupported

 from vllm.model_executor.layers.batch_invariant import rms_norm as triton_rms_norm
 from vllm.model_executor.layers.layernorm import RMSNorm
-from vllm.platforms import current_platform
-
-DEVICE_TYPE = current_platform.device_type


@skip_unsupported
@@ -37,7 +34,7 @@ def test_rms_norm_batch_invariant_vs_standard(
    equivalent results to the standard CUDA implementation across various
    configurations.
    """
-    device = torch.device(DEVICE_TYPE)
+    device = torch.device("cuda")

    # Create test input and weight
    torch.manual_seed(42)
@@ -84,7 +81,7 @@ def test_rms_norm_3d_input(
    Ensures that the batch-invariant RMS norm correctly handles multi-dimensional
    inputs that are common in transformer models.
    """
-    device = torch.device(DEVICE_TYPE)
+    device = torch.device("cuda")
    dtype = torch.bfloat16
    eps = 1e-6

@@ -123,7 +120,7 @@ def test_rms_norm_numerical_stability(default_vllm_config):
    Ensures that both implementations handle edge cases like very small or large
    values without producing NaN or Inf.
    """
-    device = torch.device(DEVICE_TYPE)
+    device = torch.device("cuda")
    dtype = torch.float16
    eps = 1e-6
    hidden_size = 2048
@@ -182,7 +179,7 @@ def test_rms_norm_formula(default_vllm_config):

    Verifies: output = input / sqrt(mean(input^2) + eps) * weight
    """
-    device = torch.device(DEVICE_TYPE)
+    device = torch.device("cuda")
    dtype = torch.float32  # Use float32 for higher precision in formula check
    eps = 1e-6
    hidden_size = 1024
@@ -217,7 +214,7 @@ def test_rms_norm_different_hidden_sizes(default_vllm_config, hidden_size: int):
    The Triton kernel uses a fixed BLOCK_SIZE=1024, so this tests that it
    correctly handles hidden sizes both smaller and larger than the block size.
    """
-    device = torch.device(DEVICE_TYPE)
+    device = torch.device("cuda")
    dtype = torch.bfloat16
    eps = 1e-6
    batch_size = 16
@@ -254,7 +251,7 @@ def test_rms_norm_determinism(default_vllm_config):
    Runs the same input through the kernel multiple times and verifies
    identical outputs.
    """
-    device = torch.device(DEVICE_TYPE)
+    device = torch.device("cuda")
    dtype = torch.bfloat16
    eps = 1e-6
    hidden_size = 4096
@@ -286,7 +283,7 @@ if __name__ == "__main__":
    # Run a quick smoke test
    print("Running quick smoke test of RMS norm implementations...")

-    device = torch.device(DEVICE_TYPE)
+    device = torch.device("cuda")
    batch_size = 8
    hidden_size = 4096
    dtype = torch.bfloat16
@@ -16,7 +16,6 @@ from vllm import LLM, SamplingParams, TokensPrompt
 from vllm.config import CacheConfig
 from vllm.distributed import cleanup_dist_env_and_memory
 from vllm.model_executor.layers.mamba.mamba_utils import MambaStateCopyFunc
-from vllm.platforms import current_platform
 from vllm.sequence import IntermediateTensors
 from vllm.v1.attention.backends.utils import CommonAttentionMetadata
 from vllm.v1.core.kv_cache_manager import KVCacheBlocks, KVCacheManager
@@ -49,7 +48,6 @@ num_accepted_tokens = 1
 prompt_token_ids: list[int] = []
 MODEL = "Qwen/Qwen3-Next-80B-A3B-Instruct-FP8"
 BLOCK_SIZE = 560
-DEVICE_TYPE = current_platform.device_type
 NUM_HIDDEN_LAYERS = 1
 cur_step_action_idx = 0
 cur_step_action: StepAction | None = None
@@ -73,7 +71,7 @@ def get_fake_sample_fn() -> SamplerOutput:
            return SamplerOutput(
                sampled_token_ids=torch.tensor(
                    [[prompt_token_ids[first_token_id_index]]],
-                    device=DEVICE_TYPE,
+                    device="cuda",
                    dtype=torch.int32,
                ),
                logprobs_tensors=None,
@@ -85,9 +83,7 @@ def get_fake_sample_fn() -> SamplerOutput:
        sampled_token_ids = accepted_tokens
        return SamplerOutput(
            sampled_token_ids=torch.tensor(
-                [sampled_token_ids],
-                device=DEVICE_TYPE,
-                dtype=torch.int32,
+                [sampled_token_ids], device="cuda", dtype=torch.int32
            ),
            logprobs_tensors=None,
        )
@@ -132,23 +128,17 @@ def get_fake_propose_draft_token_ids_fn():
                - 1
                + num_accepted_tokens
            ],
-            device=DEVICE_TYPE,
+            device="cuda",
            dtype=torch.int32,
        )

        valid_sampled_tokens_count = torch.tensor(
-            [num_accepted_tokens],
-            device=DEVICE_TYPE,
-            dtype=torch.int32,
+            [num_accepted_tokens], device="cuda", dtype=torch.int32
        )

        self._copy_valid_sampled_token_count(next_token_ids, valid_sampled_tokens_count)

-        return torch.tensor(
-            proposed_draft_token_ids,
-            device=DEVICE_TYPE,
-            dtype=torch.int32,
-        )
+        return torch.tensor(proposed_draft_token_ids, device="cuda", dtype=torch.int32)

    return fake_propose_draft_token_ids_fn

@@ -10,7 +10,6 @@ from vllm.distributed.kv_transfer.kv_transfer_state import (
    get_kv_transfer_group,
 )
 from vllm.v1.core.sched.output import CachedRequestData, SchedulerOutput
-from vllm.v1.worker.gpu.kv_connector import ActiveKVConnector
 from vllm.v1.worker.kv_connector_model_runner_mixin import KVConnectorModelRunnerMixin

 # Importing utils registers TestExampleConnector with the factory
@@ -60,29 +59,3 @@ def test_kv_connector_mixin_clears_metadata():
    finally:
        # Ensure we clean up the global connector between tests
        ensure_kv_transfer_shutdown()
-
-
-def test_active_kv_connector_runs_lifecycle_hooks_for_empty_metadata():
-    vllm_config = create_vllm_config()
-    vllm_config.kv_transfer_config.kv_connector = "TestExampleConnector"
-    vllm_config.kv_transfer_config.kv_role = "kv_both"
-    vllm_config.kv_transfer_config.kv_connector_extra_config["name"] = "empty"
-
-    ensure_kv_transfer_initialized(vllm_config)
-
-    try:
-        wrapped = get_kv_transfer_group()
-        connector = ActiveKVConnector(vllm_config, {})
-        scheduler_output = _make_empty_scheduler_output()
-
-        connector.pre_forward(scheduler_output)
-        connector.post_forward(scheduler_output)
-
-        assert wrapped.call_record.get("bind_connector_metadata", 0) == 1
-        assert wrapped.call_record.get("handle_preemptions", 0) == 1
-        assert wrapped.call_record.get("start_load_kv", 0) == 1
-        assert wrapped.call_record.get("wait_for_save", 0) == 1
-        assert wrapped.call_record.get("get_finished", 0) == 1
-        assert wrapped.call_record.get("clear_connector_metadata", 0) == 1
-    finally:
-        ensure_kv_transfer_shutdown()
@@ -6,7 +6,6 @@ import time
 import pytest
 import torch

-from vllm.platforms import current_platform
 from vllm.utils.torch_utils import set_random_seed
 from vllm.v1.kv_offload.mediums import CPULoadStoreSpec, GPULoadStoreSpec
 from vllm.v1.kv_offload.spec import (
@@ -22,8 +21,7 @@ GPU_PAGE_SIZES = [512, 1024]
 BLOCK_SIZE_FACTORS = [1, 3]
 NUM_TENSORS = [4]
 SEEDS = [0]
-DEVICE_TYPE = current_platform.device_type
-DEVICES = [f"{DEVICE_TYPE}:0"]
+CUDA_DEVICES = ["cuda:0"]
 NUM_MAPPINGS = [3]


@@ -35,7 +33,7 @@ NUM_MAPPINGS = [3]
@pytest.mark.parametrize("num_cpu_blocks", NUM_CPU_BLOCKS)
@pytest.mark.parametrize("num_tensors", NUM_TENSORS)
@pytest.mark.parametrize("seed", SEEDS)
-@pytest.mark.parametrize("device", DEVICES)
+@pytest.mark.parametrize("device", CUDA_DEVICES)
@torch.inference_mode()
 def test_transfer(
    default_vllm_config,
@@ -39,9 +39,8 @@ PIN_MEMORY_AVAILABLE = is_pin_memory_available()
 MAX_NUM_REQS = 256
 VOCAB_SIZE = 1024
 NUM_OUTPUT_TOKENS = 20
-DEVICE_TYPE = current_platform.device_type
-DEVICES = [
-    f"{DEVICE_TYPE}:{i}"
+CUDA_DEVICES = [
+    f"{current_platform.device_type}:{i}"
    for i in range(1 if current_platform.device_count() == 1 else 2)
 ]
 MAX_NUM_PROMPT_TOKENS = 64
@@ -802,7 +801,7 @@ def _assert_valid(


@create_new_process_for_each_test()
-@pytest.mark.parametrize("device", DEVICES)
+@pytest.mark.parametrize("device", CUDA_DEVICES)
@pytest.mark.parametrize("reqs_per_logitproc", [REQS_PER_LOGITPROC])
@pytest.mark.parametrize("logitsprocs_under_test", _get_test_cases())
 def test_logitsprocs(
@@ -19,7 +19,7 @@ from vllm.v1.sample.rejection_sampler import (
 from vllm.v1.sample.sampler import Sampler, SamplerOutput
 from vllm.v1.spec_decode.metadata import SpecDecodeMetadata

-DEVICE_TYPE = current_platform.device_type
+DEVICE = current_platform.device_type


@pytest.fixture
@@ -57,7 +57,7 @@ def create_logits_tensor(
    will produce desired token ids on argmax"""
    token_ids = [tokens[:-1] for tokens in output_token_ids]
    num_total_tokens = sum(len(tokens) for tokens in token_ids)
-    logits = torch.full((num_total_tokens, vocab_size), -100.0, device=DEVICE_TYPE)
+    logits = torch.full((num_total_tokens, vocab_size), -100.0, device=DEVICE)
    start_loc = 0
    for tokens in token_ids:
        for j, token_id in enumerate(tokens):
@@ -99,9 +99,9 @@ def create_sampling_metadata(
        assert output_token_ids
        assert len(output_token_ids) > 0

-        frequency_penalties = torch.tensor(frequency_penalties, device=DEVICE_TYPE)
-        presence_penalties = torch.tensor(presence_penalties, device=DEVICE_TYPE)
-        repetition_penalties = torch.tensor(repetition_penalties, device=DEVICE_TYPE)
+        frequency_penalties = torch.tensor(frequency_penalties, device=DEVICE)
+        presence_penalties = torch.tensor(presence_penalties, device=DEVICE)
+        repetition_penalties = torch.tensor(repetition_penalties, device=DEVICE)
    else:
        no_penalties = True
        frequency_penalties = torch.tensor([])
@@ -320,27 +320,14 @@ def test_deterministic_when_seeded(
    n_rep: int,
 ):
    num_tokens = batch_size * k
-    draft_probs = torch.rand(
-        num_tokens,
-        vocab_size,
-        dtype=torch.float32,
-        device=DEVICE_TYPE,
-    )
+    draft_probs = torch.rand(num_tokens, vocab_size, dtype=torch.float32, device=DEVICE)
    draft_probs = F.softmax(draft_probs, dim=-1)
    target_logits = torch.rand_like(draft_probs)
    bonus_token_ids = torch.randint(
-        low=0,
-        high=vocab_size,
-        size=(batch_size, 1),
-        dtype=torch.int64,
-        device=DEVICE_TYPE,
+        low=0, high=vocab_size, size=(batch_size, 1), dtype=torch.int64, device=DEVICE
    )
    draft_token_ids = torch.randint(
-        low=0,
-        high=vocab_size,
-        size=(batch_size, k),
-        dtype=torch.int64,
-        device=DEVICE_TYPE,
+        low=0, high=vocab_size, size=(batch_size, k), dtype=torch.int64, device=DEVICE
    )

    seeded_mask = torch.rand(batch_size, dtype=torch.float32) <= frac_seeded
@@ -348,12 +335,12 @@ def test_deterministic_when_seeded(
    results = []
    for _ in range(n_rep):
        seeded_seqs = {
-            i: torch.Generator(device=DEVICE_TYPE).manual_seed(i)
+            i: torch.Generator(device=DEVICE).manual_seed(i)
            for i in range(batch_size)
            if seeded_mask[i]
        }

-        temperature = torch.ones(batch_size, dtype=torch.float32, device=DEVICE_TYPE)
+        temperature = torch.ones(batch_size, dtype=torch.float32, device=DEVICE)
        sampling_metadata = create_sampling_metadata(
            all_greedy=False, temperature=temperature, generators=seeded_seqs
        )
@@ -400,7 +387,7 @@ def test_rejection_sampling_approximates_target_distribution():
    much more than the distance improvement between the observed
    distribution and the random distribution.
    """
-    torch.set_default_device(DEVICE_TYPE)
+    torch.set_default_device(DEVICE)
    vocab_size = 10
    k = 2
    num_reference_probs = 100
@@ -423,7 +410,7 @@ def test_rejection_sampling_approximates_target_distribution():
        rej_sample_probs = estimate_rejection_sampling_pdf(
            draft_probs, target_logits, k, vocab_size, num_samples
        )
-        rej_sample_probs = rej_sample_probs.to(DEVICE_TYPE)
+        rej_sample_probs = rej_sample_probs.to(DEVICE)

        # Average distance from reference probs.
        reference_vs_rejsample_dist = (
@@ -504,11 +491,11 @@ def estimate_rejection_sampling_pdf(
    draft_probs = draft_probs.view(num_tokens, vocab_size)

    # Bonus tokens not used but required.
-    bonus_token_ids = torch.zeros((1, 1), dtype=torch.int64, device=DEVICE_TYPE).repeat(
+    bonus_token_ids = torch.zeros((1, 1), dtype=torch.int64, device=DEVICE).repeat(
        num_samples, 1
    )

-    temperature = torch.ones(num_samples, dtype=torch.float32, device=DEVICE_TYPE)
+    temperature = torch.ones(num_samples, dtype=torch.float32, device=DEVICE)
    sampling_metadata = create_sampling_metadata(
        all_greedy=False, temperature=temperature
    )
@@ -613,7 +600,7 @@ def _test_masked_logits(

    # Create random draft probabilities.
    draft_probs = torch.rand(
-        (num_tokens, vocab_size), dtype=torch.float32, device=DEVICE_TYPE
+        (num_tokens, vocab_size), dtype=torch.float32, device=DEVICE
    )
    draft_probs = F.softmax(draft_probs, dim=-1)

@@ -623,11 +610,7 @@ def _test_masked_logits(
    draft_token_ids = draft_token_ids.tolist()

    # Bonus tokens not used but required
-    bonus_token_ids = torch.zeros(
-        (batch_size, 1),
-        dtype=torch.int64,
-        device=DEVICE_TYPE,
-    )
+    bonus_token_ids = torch.zeros((batch_size, 1), dtype=torch.int64, device=DEVICE)

    # Create spec decode metadata
    spec_decode_metadata = create_spec_decode_metadata(draft_token_ids, target_logits)
@@ -662,13 +645,12 @@ def test_top_k(rejection_sampler, top_k):

    # Randomly create top-k indices.
    top_k_indices = [
-        torch.randperm(vocab_size, device=DEVICE_TYPE)[:top_k]
-        for _ in range(num_tokens)
+        torch.randperm(vocab_size, device=DEVICE)[:top_k] for _ in range(num_tokens)
    ]
    top_k_indices = torch.stack(top_k_indices)

    # Create logits with the uniform distribution.
-    target_logits = torch.zeros((num_tokens, vocab_size), device=DEVICE_TYPE)
+    target_logits = torch.zeros((num_tokens, vocab_size), device=DEVICE)

    # Increment the logits for top-k indices, a little bit more than the other
    # ones. If the masking is effective, the non-topk indices will never be
@@ -677,11 +659,11 @@ def test_top_k(rejection_sampler, top_k):
        target_logits[i, top_k_indices[i]] += 0.1

    # Create sampling metadata
-    temperature = torch.ones(batch_size, dtype=torch.float32, device=DEVICE_TYPE)
+    temperature = torch.ones(batch_size, dtype=torch.float32, device=DEVICE)
    sampling_metadata = create_sampling_metadata(
        all_greedy=False,
        temperature=temperature,
-        top_k=torch.tensor([top_k] * batch_size, device=DEVICE_TYPE, dtype=torch.int64),
+        top_k=torch.tensor([top_k] * batch_size, device=DEVICE, dtype=torch.int64),
    )

    _test_masked_logits(
@@ -704,8 +686,8 @@ def test_top_p(rejection_sampler, top_p):
    num_tokens = batch_size * num_draft_tokens

    # Create logits with the uniform distribution.
-    target_logits = torch.randn((num_tokens, vocab_size), device=DEVICE_TYPE)
-    temperature = torch.ones(batch_size, dtype=torch.float32, device=DEVICE_TYPE)
+    target_logits = torch.randn((num_tokens, vocab_size), device=DEVICE)
+    temperature = torch.ones(batch_size, dtype=torch.float32, device=DEVICE)
    rescaled_logits = target_logits / temperature

    logits_sort, logits_idx = rescaled_logits.sort(dim=-1, descending=False)
@@ -724,11 +706,7 @@ def test_top_p(rejection_sampler, top_p):
    sampling_metadata = create_sampling_metadata(
        all_greedy=False,
        temperature=temperature,
-        top_p=torch.tensor(
-            [top_p] * batch_size,
-            device=DEVICE_TYPE,
-            dtype=torch.float32,
-        ),
+        top_p=torch.tensor([top_p] * batch_size, device=DEVICE, dtype=torch.float32),
    )

    _test_masked_logits(
@@ -754,10 +732,7 @@ def test_frequency_penalties(rejection_sampler):
        all_greedy=True,
        output_token_ids=[[2], [3], [4]],
        spec_token_ids=spec_tokens,
-        prompt_token_ids=torch.tensor(
-            [[5, 6, 7], [6, 7, 8], [7, 8, 9]],
-            device=DEVICE_TYPE,
-        ),
+        prompt_token_ids=torch.tensor([[5, 6, 7], [6, 7, 8], [7, 8, 9]], device=DEVICE),
        frequency_penalties=[1.5, 1.5, 0.7],
        presence_penalties=[0.0] * num_requests,
        repetition_penalties=[1.0] * num_requests,
@@ -883,26 +858,21 @@ def test_sample_recovered_tokens(
    num_tokens = batch_size * max_spec_len

    # Create random draft probabilities.
-    draft_probs = torch.rand(
-        num_tokens,
-        vocab_size,
-        dtype=torch.float32,
-        device=DEVICE_TYPE,
-    )
+    draft_probs = torch.rand(num_tokens, vocab_size, dtype=torch.float32, device=DEVICE)
    draft_probs = F.softmax(draft_probs, dim=-1)

    # Create random target probabilities.
    target_logits = torch.rand(
-        num_tokens, vocab_size, dtype=torch.float32, device=DEVICE_TYPE
+        num_tokens, vocab_size, dtype=torch.float32, device=DEVICE
    )
    target_probs = F.softmax(target_logits, dim=-1)

    # Randomly sample draft token ids from draft probs
    draft_token_ids = torch.multinomial(draft_probs, num_samples=1).to(torch.int32)

-    temperature = torch.ones(batch_size, dtype=torch.float32, device=DEVICE_TYPE)
+    temperature = torch.ones(batch_size, dtype=torch.float32, device=DEVICE)
    generators = {
-        i: torch.Generator(device=DEVICE_TYPE).manual_seed(i) for i in range(batch_size)
+        i: torch.Generator(device=DEVICE).manual_seed(i) for i in range(batch_size)
    }
    sampling_metadata = create_sampling_metadata(
        all_greedy=False, temperature=temperature, generators=generators
@@ -920,7 +890,7 @@ def test_sample_recovered_tokens(
        None if no_draft_probs else draft_probs,
        target_probs,
        sampling_metadata,
-        device=DEVICE_TYPE,
+        device=DEVICE,
    )
    recovered_token_ids = sample_recovered_tokens(
        max_spec_len,
@@ -930,6 +900,6 @@ def test_sample_recovered_tokens(
        None if no_draft_probs else draft_probs,
        target_probs,
        sampling_metadata,
-        device=DEVICE_TYPE,
+        device=DEVICE,
    )
    assert torch.equal(recovered_token_ids, ref_recovered_token_ids)
@@ -17,9 +17,8 @@ PIN_MEMORY_AVAILABLE = is_pin_memory_available()
 MAX_NUM_REQS = 256
 VOCAB_SIZE = 1024
 NUM_OUTPUT_TOKENS = 20
-DEVICE_TYPE = current_platform.device_type
-DEVICES = [
-    f"{DEVICE_TYPE}:{i}"
+CUDA_DEVICES = [
+    f"{current_platform.device_type}:{i}"
    for i in range(1 if current_platform.device_count() == 1 else 2)
 ]
 MAX_NUM_PROMPT_TOKENS = 64
@@ -200,7 +199,7 @@ def _create_weighted_output_token_list(
    return output_token_ids, sorted_token_ids_in_output


-@pytest.mark.parametrize("device", DEVICES)
+@pytest.mark.parametrize("device", CUDA_DEVICES)
@pytest.mark.parametrize("batch_size", [1, 2, 32])
@pytest.mark.parametrize("presence_penalty", [-2.0, 2.0])
 def test_sampler_presence_penalty(
@@ -250,7 +249,7 @@ def test_sampler_presence_penalty(
            assert penalized_token_id not in output_token_ids[batch_idx]


-@pytest.mark.parametrize("device", DEVICES)
+@pytest.mark.parametrize("device", CUDA_DEVICES)
@pytest.mark.parametrize("batch_size", [1, 2, 32])
@pytest.mark.parametrize("frequency_penalty", [-2.0, 2.0])
 def test_sampler_frequency_penalty(
@@ -306,7 +305,7 @@ def test_sampler_frequency_penalty(
            assert penalized_token_id not in distinct_sorted_token_ids_in_output


-@pytest.mark.parametrize("device", DEVICES)
+@pytest.mark.parametrize("device", CUDA_DEVICES)
@pytest.mark.parametrize("batch_size", [1, 2, 32])
@pytest.mark.parametrize("repetition_penalty", [0.1, 1.9])
 def test_sampler_repetition_penalty(
@@ -364,7 +363,7 @@ def test_sampler_repetition_penalty(
            )


-@pytest.mark.parametrize("device", DEVICES)
+@pytest.mark.parametrize("device", CUDA_DEVICES)
@pytest.mark.parametrize("batch_size", [1, 2, 32])
@pytest.mark.parametrize("num_allowed_token_ids", [0, 1, 2])
 def test_sampler_allowed_token_ids(
@@ -410,7 +409,7 @@ def test_sampler_allowed_token_ids(
                assert logits_for_req[token_id] != -float("inf")


-@pytest.mark.parametrize("device", DEVICES)
+@pytest.mark.parametrize("device", CUDA_DEVICES)
@pytest.mark.parametrize("batch_size", [1, 2, 32])
@pytest.mark.parametrize("bad_words_lengths", [(1,), (1, 3), (2, 2)])
 def test_sampler_bad_words(
@@ -7,7 +7,8 @@ from torch import Generator
 from vllm.platforms import current_platform
 from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p_pytorch

-DEVICE_TYPE = current_platform.device_type
+CUDA_DEVICE = "cuda" if current_platform.is_cuda() else None
+DEVICE = current_platform.device_type

 BATCH_SIZE = 1024
 VOCAB_SIZE = 128 * 1024
@@ -25,8 +26,8 @@ def reset_default_device():


 def test_topk_impl_equivalence():
-    torch.set_default_device(DEVICE_TYPE)
-    generator = Generator(device=DEVICE_TYPE).manual_seed(33)
+    torch.set_default_device(DEVICE)
+    generator = Generator(device=DEVICE).manual_seed(33)

    logits = torch.rand((BATCH_SIZE, VOCAB_SIZE), generator=generator)

@@ -75,8 +76,8 @@ def test_flashinfer_sampler():
    if not FLASHINFER_ENABLED:
        pytest.skip("FlashInfer not installed or not available on this platform.")

-    torch.set_default_device(DEVICE_TYPE)
-    generator = Generator(device=DEVICE_TYPE).manual_seed(42)
+    torch.set_default_device(DEVICE)
+    generator = Generator(device=DEVICE).manual_seed(42)

    # Generate random logits
    logits = torch.rand((BATCH_SIZE, VOCAB_SIZE), generator=generator)
@@ -127,15 +128,15 @@ def test_flashinfer_sampler():
 # =============================================================================


-@pytest.mark.skipif("CPU" in DEVICE_TYPE, reason="CUDA/XPU not available")
+@pytest.mark.skipif(CUDA_DEVICE is None, reason="CUDA not available")
 class TestTritonTopkTopp:
    """Tests for the Triton top-k/top-p kernel."""

    @pytest.fixture(autouse=True)
    def setup(self):
        """Set up test fixtures."""
-        torch.set_default_device(DEVICE_TYPE)
-        self.generator = Generator(device=DEVICE_TYPE).manual_seed(42)
+        torch.set_default_device(CUDA_DEVICE)
+        self.generator = Generator(device=CUDA_DEVICE).manual_seed(42)

    def _compare_results(
        self,
@@ -42,7 +42,6 @@ dflash_target_dir = "Qwen/Qwen3-8B"
 dflash_dir = "z-lab/Qwen3-8B-DFlash-b16"

 BLOCK_SIZE = 16
-DEVICE_TYPE = current_platform.device_type


 def _create_proposer(
@@ -93,7 +92,7 @@ def _create_proposer(
        # Overwrite pard_token to avoid crash during init
        speculative_config.draft_model_config.hf_config.pard_token = 0

-    device = DEVICE_TYPE
+    device = current_platform.device_type
    vllm_config = VllmConfig(
        model_config=model_config,
        cache_config=CacheConfig(block_size=16),
@@ -125,7 +124,7 @@ def test_prepare_next_token_ids():
    either the GPU tensor of sampled_token_ids with -1 for rejected tokens,
    or the CPU python list[list[int]] with the rejected tokens removed.
    """
-    device = torch.device(DEVICE_TYPE)
+    device = torch.device(current_platform.device_type)

    num_requests = 4
    num_speculative_tokens = 4
@@ -208,7 +207,7 @@ def test_prepare_inputs():
                    a, a + 1, ..., a + b - n2 - 1,
                    a + b, a + b + 1, ..., a + b + c - n3 - 1]
    """
-    device = torch.device(DEVICE_TYPE)
+    device = torch.device(current_platform.device_type)

    # q1 = 4, q2 = 7, q3 = 5
    # n1 = 1, n2 = 3, n3 = 2
@@ -301,7 +300,7 @@ def test_prepare_inputs_padded():
            from the original indices to sample from.
    """

-    device = torch.device(DEVICE_TYPE)
+    device = torch.device(current_platform.device_type)

    expected_token_indices_to_sample = torch.tensor(
        [1, 5, 6], dtype=torch.int32, device=device
@@ -371,7 +370,7 @@ def test_set_inputs_first_pass_default_eagle():
    - After inserting next_tokens [100, 200, 300]:
        [a2, a3, 100, b2, 200, c2, c3, c4, 300]
    """
-    device = torch.device(DEVICE_TYPE)
+    device = torch.device(current_platform.device_type)

    num_speculative_tokens = 3
    proposer = _create_proposer("eagle", num_speculative_tokens)
@@ -472,7 +471,7 @@ def test_set_inputs_first_pass_draft_model():
      - idx 5: token 21, pos 1
      - idx 6: token 200, pos 2 (bonus token)
    """
-    device = torch.device(DEVICE_TYPE)
+    device = torch.device(current_platform.device_type)

    num_speculative_tokens = 2
    block_size = BLOCK_SIZE
@@ -610,7 +609,7 @@ def test_set_inputs_first_pass_parallel_drafting():
      - idx 9: bonus token 200
      - idx 10-11: parallel_drafting_tokens, is_masked=True
    """
-    device = torch.device(DEVICE_TYPE)
+    device = torch.device(current_platform.device_type)

    num_speculative_tokens = 3
    block_size = BLOCK_SIZE
@@ -860,7 +859,7 @@ def test_propose(method, attn_backend, num_speculative_tokens, monkeypatch):
        monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1")

    # Use GPU device
-    device = torch.device(DEVICE_TYPE)
+    device = torch.device(current_platform.device_type)

    # Setup test parameters
    batch_size = 2
@@ -1031,7 +1030,7 @@ def test_propose(method, attn_backend, num_speculative_tokens, monkeypatch):
 )
 def test_propose_tree(spec_token_tree):
    # Get GPU device.
-    device = torch.device(DEVICE_TYPE)
+    device = torch.device(current_platform.device_type)

    # Setup test parameters.
    batch_size = 2
@@ -5,14 +5,11 @@
 import pytest
 import torch

-from vllm.platforms import current_platform
 from vllm.v1.spec_decode.utils import (
    PADDING_SLOT_ID,
    eagle_step_update_slot_mapping_and_metadata,
 )

-DEVICE_TYPE = current_platform.device_type
-
 # Skip if no CUDA - Triton kernel requires GPU
 pytest.importorskip("triton")
 if not torch.cuda.is_available():
@@ -50,7 +47,7 @@ def _reference_eagle_step_slot_mapping(

 def test_eagle_step_slot_mapping_kernel():
    """Test fused kernel matches Python reference for slot mapping and metadata."""
-    device = torch.device(DEVICE_TYPE)
+    device = torch.device("cuda")
    batch_size = 32
    block_size = 16
    max_model_len = 4096
@@ -96,7 +93,7 @@ def test_eagle_step_slot_mapping_kernel():

 def test_eagle_step_slot_mapping_kernel_exceeds_max():
    """Test fused kernel when position exceeds max_model_len."""
-    device = torch.device(DEVICE_TYPE)
+    device = torch.device("cuda")
    batch_size = 4
    block_size = 16
    max_model_len = 100
@@ -133,7 +130,7 @@ def test_eagle_step_slot_mapping_kernel_exceeds_max():
 def test_eagle_step_slot_mapping_kernel_cudagraph_padding():
    """Test that padding threads write PADDING_SLOT_ID when
    input_batch_size > batch_size (cudagraph padding)."""
-    device = torch.device(DEVICE_TYPE)
+    device = torch.device("cuda")
    batch_size = 4
    input_batch_size = 8
    block_size = 16
@@ -27,7 +27,6 @@ from vllm.v1.spec_decode.extract_hidden_states import ExtractHiddenStatesPropose
 from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch

 model_dir = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
-DEVICE_TYPE = current_platform.device_type


 def _create_proposer(
@@ -52,7 +51,7 @@ def _create_proposer(
        },
    )

-    device = DEVICE_TYPE
+    device = current_platform.device_type
    vllm_config = VllmConfig(
        model_config=model_config,
        cache_config=CacheConfig(),
@@ -102,7 +101,7 @@ def test_proposer_initialization_missing_layer_ids():
        },
    )

-    device = DEVICE_TYPE
+    device = current_platform.device_type
    vllm_config = VllmConfig(
        model_config=model_config,
        cache_config=CacheConfig(),
@@ -131,7 +130,7 @@ def test_prepare_next_token_ids_padded():
    For each request we either use the sampled token (if valid and not discarded)
    or a backup token from the request state.
    """
-    device = torch.device(DEVICE_TYPE)
+    device = torch.device(current_platform.device_type)

    num_requests = 4
    req_ids = [f"req_{i + 1}" for i in range(num_requests)]
@@ -198,7 +197,7 @@ def test_propose():
    2. Return the sampled tokens as "draft" tokens (shape [batch_size, 1])
    3. Cache the hidden states in the model's KV cache
    """
-    device = torch.device(DEVICE_TYPE)
+    device = torch.device(current_platform.device_type)

    # Setup test parameters
    batch_size = 2
@@ -274,7 +273,7 @@ def test_propose():
@pytest.mark.parametrize("num_hidden_layers", [1, 4, 8])
 def test_propose_different_layer_counts(num_hidden_layers):
    """Test that propose works correctly with different numbers of hidden layers."""
-    device = torch.device(DEVICE_TYPE)
+    device = torch.device(current_platform.device_type)

    batch_size = 2
    num_tokens = 5
@@ -28,7 +28,6 @@ from vllm.v1.attention.backends.registry import AttentionBackendEnum
 from vllm.v1.spec_decode.eagle import EagleProposer

 mimo_7b_dir = "XiaomiMiMo/MiMo-7B-Base"
-DEVICE_TYPE = current_platform.device_type


 def _create_mtp_proposer(num_speculative_tokens: int) -> EagleProposer:
@@ -49,7 +48,7 @@ def _create_mtp_proposer(num_speculative_tokens: int) -> EagleProposer:
        model_config=model_config,
        cache_config=CacheConfig(),
        speculative_config=speculative_config,
-        device_config=DeviceConfig(device=DEVICE_TYPE),
+        device_config=DeviceConfig(device=current_platform.device_type),
        parallel_config=ParallelConfig(),
        load_config=LoadConfig(),
        scheduler_config=SchedulerConfig(
@@ -58,7 +57,7 @@ def _create_mtp_proposer(num_speculative_tokens: int) -> EagleProposer:
        ),
    )

-    return EagleProposer(vllm_config=vllm_config, device=DEVICE_TYPE)
+    return EagleProposer(vllm_config=vllm_config, device=current_platform.device_type)


@mock.patch("vllm.v1.spec_decode.eagle.get_pp_group")
@@ -119,7 +118,7 @@ def test_mtp_load_model_unified(mock_get_model, mock_get_layers, mock_get_pp_gro
 def test_mtp_propose(num_speculative_tokens, monkeypatch):
    """Test that MTP's forward method returns hidden states directly"""

-    device = torch.device(DEVICE_TYPE)
+    device = torch.device(current_platform.device_type)
    batch_size = 2
    seq_lens = [5, 3]
    total_tokens = sum(seq_lens)
@@ -18,8 +18,6 @@ from vllm.v1.attention.backend import CommonAttentionMetadata
 from vllm.v1.attention.backends.fa_utils import is_flash_attn_varlen_func_available
 from vllm.v1.attention.backends.registry import AttentionBackendEnum

-DEVICE_TYPE = current_platform.device_type
-
 if not is_flash_attn_varlen_func_available():
    pytest.skip(
        "This test requires flash_attn_varlen_func, but it's not available.",
@@ -172,9 +170,9 @@ def _get_available_reference_backends() -> list[AttentionBackendEnum]:


 class MockAttentionLayer(torch.nn.Module):
-    _q_scale = torch.tensor(1.0, dtype=torch.float32, device=DEVICE_TYPE)
-    _k_scale = torch.tensor(1.0, dtype=torch.float32, device=DEVICE_TYPE)
-    _v_scale = torch.tensor(1.0, dtype=torch.float32, device=DEVICE_TYPE)
+    _q_scale = torch.tensor(1.0, dtype=torch.float32, device="cuda")
+    _k_scale = torch.tensor(1.0, dtype=torch.float32, device="cuda")
+    _v_scale = torch.tensor(1.0, dtype=torch.float32, device="cuda")
    layer_name = "mock_layer"

    def __init__(self):
@@ -22,8 +22,10 @@ from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch
 VOCAB_SIZE = 1024
 NUM_OUTPUT_TOKENS = 20
 MAX_PROMPT_SIZE = 100
-DEVICE_TYPE = current_platform.device_type
-DEVICES = [f"{DEVICE_TYPE}:{i}" for i in range(min(current_platform.device_count(), 2))]
+CUDA_DEVICES = [
+    f"{current_platform.device_type}:{i}"
+    for i in range(min(current_platform.device_count(), 2))
+]
 MAX_NUM_PROMPT_TOKENS = 64


@@ -217,7 +219,7 @@ def _construct_cached_request_state(req_id_suffix: int):
    )


-@pytest.mark.parametrize("device", DEVICES)
+@pytest.mark.parametrize("device", CUDA_DEVICES)
@pytest.mark.parametrize("batch_size", [1, 2, 32, 64])
 def test_sampling_metadata_in_input_batch(device: str, batch_size: int):
    """
@@ -311,7 +313,7 @@ def test_sampling_metadata_in_input_batch(device: str, batch_size: int):
    )


-@pytest.mark.parametrize("device", DEVICES)
+@pytest.mark.parametrize("device", CUDA_DEVICES)
@pytest.mark.parametrize("batch_size", [32])
@pytest.mark.parametrize("swap_list", [((0, 1),)])
 def test_swap_states_in_input_batch(device: str, batch_size: int, swap_list: list):
@@ -398,7 +400,7 @@ def _construct_pooling_request(req_id_suffix: int, pooling_params=None):
    )


-@pytest.mark.parametrize("device", DEVICES)
+@pytest.mark.parametrize("device", CUDA_DEVICES)
 def test_pooling_prompt_lens_not_aliased(device: str):
    """Verify that prompt_lens in PoolingMetadata does not share memory
    with the internal num_prompt_tokens pinned buffer. Guards against possible
@@ -45,7 +45,7 @@ from vllm.v1.worker.utils import AttentionGroup, select_common_block_size

 BLOCK_SIZE = 16
 NUM_BLOCKS = 10
-DEVICE_TYPE = current_platform.device_type
+DEVICE = current_platform.device_type


 def initialize_kv_cache(runner: GPUModelRunner):
@@ -121,7 +121,7 @@ def model_runner():
        vllm_config.compilation_config.static_forward_context["layer.0"] = Attention(
            num_heads, head_size, 0.1
        )
-        runner = GPUModelRunner(vllm_config, DEVICE_TYPE)
+        runner = GPUModelRunner(vllm_config, DEVICE)
        initialize_kv_cache(runner)
        yield runner

@@ -340,7 +340,7 @@ def test_get_nans_in_logits(model_runner, dist_init):
            [1.0, 2.0, 3.0],
            [3.0, 2.0, 1.0],
        ],
-        device=DEVICE_TYPE,
+        device=DEVICE,
    )
    result = model_runner._get_nans_in_logits(logits)
    assert result == {"req_0": 0, "req_1": 0}
@@ -350,7 +350,7 @@ def test_get_nans_in_logits(model_runner, dist_init):
            [1.0, float("nan"), 3.0],
            [4.0, float("nan"), float("nan")],
        ],
-        device=DEVICE_TYPE,
+        device=DEVICE,
    )
    result = model_runner._get_nans_in_logits(logits)
    assert result == {"req_0": 1, "req_1": 2}
@@ -360,7 +360,7 @@ def test_get_nans_in_logits(model_runner, dist_init):
            [1.0, 2.0, 3.0],
            [4.0, float("nan"), float("nan")],
        ],
-        device=DEVICE_TYPE,
+        device=DEVICE,
    )
    result = model_runner._get_nans_in_logits(logits)
    assert result == {"req_0": 0, "req_1": 2}
@@ -372,7 +372,7 @@ def test_get_nans_in_logits(model_runner, dist_init):
        [
            [1.0, float("nan"), 3.0],
        ],
-        device=DEVICE_TYPE,
+        device=DEVICE,
    )
    result = model_runner._get_nans_in_logits(logits)
    assert result == {"req_0": 1, "req_1": 0}
@@ -383,7 +383,7 @@ def test_get_nans_in_logits(model_runner, dist_init):
            [1.0, 2.0, 3.0],
            [float("nan"), 2.0, 3.0],
        ],
-        device=DEVICE_TYPE,
+        device=DEVICE,
    )
    result = model_runner._get_nans_in_logits(logits)
    assert result == {"req_0": 2, "req_1": 0}
@@ -643,7 +643,7 @@ def test_init_kv_cache_without_kv_sharing(default_vllm_config):
    # Set high context length to test max context length estimation
    vllm_config.model_config.max_model_len = 3_000_000
    vllm_ctx = vllm_config.compilation_config.static_forward_context
-    runner = GPUModelRunner(vllm_config, DEVICE_TYPE)
+    runner = GPUModelRunner(vllm_config, DEVICE)
    kv_cache_spec = runner.get_kv_cache_spec()
    assert len(kv_cache_spec) == 2
    assert len(runner.shared_kv_cache_layers) == 0
@@ -711,7 +711,7 @@ def test_init_kv_cache_with_kv_sharing_valid(default_vllm_config):
    # Set high context length to test max context length estimation
    vllm_config.model_config.max_model_len = 3_000_000
    vllm_ctx = vllm_config.compilation_config.static_forward_context
-    runner = GPUModelRunner(vllm_config, DEVICE_TYPE)
+    runner = GPUModelRunner(vllm_config, DEVICE)
    kv_cache_spec = runner.get_kv_cache_spec()
    assert len(kv_cache_spec) == 1
    assert layer_0 in kv_cache_spec
@@ -850,7 +850,7 @@ def test_hybrid_attention_mamba_tensor_shapes():
        assert fwd_context is not None
        vllm_ctx = vllm_config.compilation_config.static_forward_context

-        runner = GPUModelRunner(vllm_config, DEVICE_TYPE)
+        runner = GPUModelRunner(vllm_config, DEVICE)
        current_platform.update_block_size_for_backend(vllm_config)
        kv_cache_spec = runner.get_kv_cache_spec()

@@ -896,13 +896,13 @@ def test_hybrid_attention_mamba_tensor_shapes():
    ssm_constant_shape = ssm_shape[1:]

    attn_blocks_constant = torch.full(
-        (test_block_size, *attn_constant_shape), device=DEVICE_TYPE, fill_value=3.33
+        (test_block_size, *attn_constant_shape), device=DEVICE, fill_value=3.33
    )
    conv_blocks_constant = torch.full(
-        (test_block_size, *conv_constant_shape), device=DEVICE_TYPE, fill_value=6.66
+        (test_block_size, *conv_constant_shape), device=DEVICE, fill_value=6.66
    )
    ssm_blocks_constant = torch.full(
-        (test_block_size, *ssm_constant_shape), device=DEVICE_TYPE, fill_value=9.99
+        (test_block_size, *ssm_constant_shape), device=DEVICE, fill_value=9.99
    )

    # Fill attention blocks with constants using kv block indices
@@ -997,7 +997,7 @@ def test_hybrid_block_table_initialization():
        max_num_blocks_per_req=max_num_blocks_per_req,
        max_num_batched_tokens=max_num_batched_tokens,
        pin_memory=False,
-        device=torch.device(DEVICE_TYPE),
+        device=torch.device(DEVICE),
        kernel_block_size=kernel_block_sizes[0],
        cp_kv_cache_interleave_size=cp_kv_cache_interleave_size,
    )
@@ -1036,7 +1036,7 @@ def test_input_batch_with_kernel_block_sizes():
    max_num_reqs = 10
    max_model_len = 512
    max_num_batched_tokens = 512
-    device = torch.device(DEVICE_TYPE)
+    device = torch.device(DEVICE)
    pin_memory = False
    vocab_size = 50272

@@ -1083,7 +1083,7 @@ def test_hybrid_cache_integration(default_vllm_config, dist_init):
        num_heads, head_size, 0.1
    )

-    runner = GPUModelRunner(vllm_config, DEVICE_TYPE)
+    runner = GPUModelRunner(vllm_config, DEVICE)

    # Initialize KV cache with configuration
    attn_spec = FullAttentionSpec(
@@ -1306,7 +1306,7 @@ def test_mamba_cache_raises_when_max_num_seqs_exceeds_blocks():
            )
        assert fwd_context is not None

-        runner = GPUModelRunner(vllm_config, DEVICE_TYPE)
+        runner = GPUModelRunner(vllm_config, DEVICE)
        current_platform.update_block_size_for_backend(vllm_config)
        kv_cache_spec = runner.get_kv_cache_spec()

@@ -265,7 +265,6 @@ def merge_attn_states(
    suffix_lse: torch.Tensor,
    output_lse: torch.Tensor | None = None,
    prefill_tokens_with_context: int | None = None,
-    output_scale: torch.Tensor | None = None,
 ) -> None:
    torch.ops._C.merge_attn_states(
        output,
@@ -275,7 +274,6 @@ def merge_attn_states(
        suffix_output,
        suffix_lse,
        prefill_tokens_with_context,
-        output_scale,
    )


@@ -2641,22 +2639,6 @@ def swap_blocks(
    torch.ops._C_cache_ops.swap_blocks(src, dst, block_size_in_bytes, block_mapping)


-def swap_blocks_batch(
-    src_ptrs: torch.Tensor,
-    dst_ptrs: torch.Tensor,
-    sizes: torch.Tensor,
-) -> None:
-    """
-    Batch version of swap_blocks: submit all copies in a single driver call.
-
-    Each entry specifies a raw pointer copy: src_ptrs[i] -> dst_ptrs[i]
-    of sizes[i] bytes. All three tensors must be int64 CPU tensors.
-    On CUDA 12.8+ this uses cuMemcpyBatchAsync for minimal submission
-    overhead; on older CUDA it falls back to a loop of cudaMemcpyAsync.
-    """
-    torch.ops._C_cache_ops.swap_blocks_batch(src_ptrs, dst_ptrs, sizes)
-
-
 def convert_fp8(
    output: torch.Tensor, input: torch.Tensor, scale: float = 1.0, kv_dtype: str = "fp8"
 ) -> None:
@@ -1,262 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-from collections.abc import Callable
-
-import torch
-from torch._higher_order_ops.auto_functionalize import auto_functionalized
-
-from vllm._custom_ops import create_fp4_output_tensors
-from vllm.config import VllmConfig, get_layers_from_vllm_config
-from vllm.logger import init_logger
-from vllm.model_executor.layers.attention.mla_attention import MLAAttention
-from vllm.model_executor.layers.quantization.utils.quant_utils import (
-    kFp8StaticTensorSym,
-    kNvfp4Dynamic,
-)
-from vllm.platforms import current_platform
-
-from ..vllm_inductor_pass import VllmFusionPatternMatcherPass, VllmPatternReplacement
-from .matcher_utils import MatcherQuantFP8
-from .rms_quant_fusion import QUANT_OPS
-
-logger = init_logger(__name__)
-
-FP8_DTYPE = current_platform.fp8_dtype()
-FP4_DTYPE = torch.uint8
-
-MLA_ATTN_OP = torch.ops.vllm.unified_mla_attention_with_output.default
-
-
-class MLAAttnFp8StaticQuantPattern(VllmPatternReplacement[..., torch.Tensor]):
-    """
-    Fusion for MLA Attention+Fp8StaticQuant.
-
-    Matches the pattern: MLA attention -> static FP8 quant, and replaces
-    it with MLA attention(output_scale=scale, output=fp8_buffer).
-    """
-
-    def __init__(self, layer: MLAAttention, dtype: torch.dtype) -> None:
-        self._layer_name = layer.layer_name
-        self._num_heads = layer.num_heads
-        self._v_head_dim = layer.v_head_dim
-        self._kv_lora_rank = layer.kv_lora_rank
-        self._qk_rope_head_dim = layer.qk_rope_head_dim
-        self._qk_head_dim = layer.qk_nope_head_dim + layer.qk_rope_head_dim
-        self._output_dim = layer.num_heads * layer.v_head_dim
-        self._dtype = dtype
-        self._quant_matcher = MatcherQuantFP8(kFp8StaticTensorSym)
-
-    @property
-    def pattern(self) -> Callable[..., torch.Tensor]:
-        def _pattern(
-            q: torch.Tensor,
-            kv_c_normed: torch.Tensor,
-            k_pe: torch.Tensor,
-            output_attn: torch.Tensor,
-            scale: torch.Tensor,
-            kv_cache_dummy_dep: torch.Tensor,
-        ) -> torch.Tensor:
-            at1 = auto_functionalized(
-                MLA_ATTN_OP,
-                q=q,
-                kv_c_normed=kv_c_normed,
-                k_pe=k_pe,
-                output=output_attn,
-                layer_name=self._layer_name,
-                output_scale=None,
-                output_block_scale=None,
-                kv_cache_dummy_dep=kv_cache_dummy_dep,
-            )
-            # MLA output is already 2D (T, N*V), no reshape needed
-            return self._quant_matcher(at1[1], scale)[0]
-
-        return _pattern
-
-    @property
-    def replacement(self) -> Callable[..., torch.Tensor]:
-        def _replacement(
-            q: torch.Tensor,
-            kv_c_normed: torch.Tensor,
-            k_pe: torch.Tensor,
-            output_attn: torch.Tensor,
-            scale: torch.Tensor,
-            kv_cache_dummy_dep: torch.Tensor,
-        ) -> torch.Tensor:
-            # MLA output in quant_dtype
-            output_attn = torch.empty(
-                [q.shape[0], self._output_dim],
-                dtype=FP8_DTYPE,
-                device=q.device,
-            )
-            at1 = auto_functionalized(
-                MLA_ATTN_OP,
-                q=q,
-                kv_c_normed=kv_c_normed,
-                k_pe=k_pe,
-                output=output_attn,
-                layer_name=self._layer_name,
-                output_scale=scale,
-                output_block_scale=None,
-                kv_cache_dummy_dep=kv_cache_dummy_dep,
-            )
-            return at1[1]
-
-        return _replacement
-
-    def get_inputs(self) -> list[torch.Tensor]:
-        return [
-            self.empty(5, self._num_heads, self._qk_head_dim, dtype=self._dtype),
-            self.empty(5, self._kv_lora_rank, dtype=self._dtype),
-            self.empty(5, 1, self._qk_rope_head_dim, dtype=self._dtype),
-            self.empty(5, self._output_dim, dtype=self._dtype),
-            self.empty_fp32(1, 1),
-            self.empty(0, dtype=self._dtype),
-        ]
-
-
-class MLAAttnNvfp4QuantPattern(
-    VllmPatternReplacement[..., tuple[torch.Tensor, torch.Tensor]]
-):
-    """
-    Fusion for MLA Attention+Nvfp4Quant.
-
-    Matches the pattern: MLA attention -> NVFP4 quant, and replaces
-    it with MLA attention(output_scale=scale, output_block_scale=block_scale,
-    output=fp4_buffer).
-    """
-
-    def __init__(self, layer: MLAAttention, dtype: torch.dtype) -> None:
-        self._layer_name = layer.layer_name
-        self._num_heads = layer.num_heads
-        self._v_head_dim = layer.v_head_dim
-        self._kv_lora_rank = layer.kv_lora_rank
-        self._qk_rope_head_dim = layer.qk_rope_head_dim
-        self._qk_head_dim = layer.qk_nope_head_dim + layer.qk_rope_head_dim
-        self._output_dim = layer.num_heads * layer.v_head_dim
-        self._dtype = dtype
-        self._QUANT_OP = QUANT_OPS[kNvfp4Dynamic]
-
-    @property
-    def pattern(
-        self,
-    ) -> Callable[..., tuple[torch.Tensor, torch.Tensor]]:
-        def _pattern(
-            q: torch.Tensor,
-            kv_c_normed: torch.Tensor,
-            k_pe: torch.Tensor,
-            output_attn: torch.Tensor,
-            input_scale: torch.Tensor,
-            kv_cache_dummy_dep: torch.Tensor,
-        ) -> tuple[torch.Tensor, torch.Tensor]:
-            at1 = auto_functionalized(
-                MLA_ATTN_OP,
-                q=q,
-                kv_c_normed=kv_c_normed,
-                k_pe=k_pe,
-                output=output_attn,
-                layer_name=self._layer_name,
-                output_scale=None,
-                output_block_scale=None,
-                kv_cache_dummy_dep=kv_cache_dummy_dep,
-            )
-            # Replicate what scaled_fp4_quant() does: allocate output
-            # tensors inline then call the .out variant.
-            output_quant, output_scale = create_fp4_output_tensors(
-                at1[1].shape[0], at1[1].shape[1], at1[1].device, True
-            )
-            at2 = auto_functionalized(
-                self._QUANT_OP,
-                input=at1[1],
-                input_scale=input_scale,
-                is_sf_swizzled_layout=True,
-                output=output_quant,
-                output_scale=output_scale,
-            )
-            output_scale_view = torch.ops.aten.view.dtype(at2[2], FP8_DTYPE)
-            return at2[1], output_scale_view
-
-        return _pattern
-
-    @property
-    def replacement(
-        self,
-    ) -> Callable[..., tuple[torch.Tensor, torch.Tensor]]:
-        def _replacement(
-            q: torch.Tensor,
-            kv_c_normed: torch.Tensor,
-            k_pe: torch.Tensor,
-            output_attn: torch.Tensor,
-            input_scale: torch.Tensor,
-            kv_cache_dummy_dep: torch.Tensor,
-        ) -> tuple[torch.Tensor, torch.Tensor]:
-            # MLA output in quant_dtype (FP4 packed as uint8)
-            output_attn = torch.empty(
-                [q.shape[0], self._output_dim // 2],
-                dtype=FP4_DTYPE,
-                device=q.device,
-            )
-            # attention output block scale
-            output_scale = create_fp4_output_tensors(
-                q.shape[0], self._output_dim, q.device, True
-            )[1]
-            output_scale_view = torch.ops.aten.view.dtype(output_scale, FP8_DTYPE)
-            at2 = auto_functionalized(
-                MLA_ATTN_OP,
-                q=q,
-                kv_c_normed=kv_c_normed,
-                k_pe=k_pe,
-                output=output_attn,
-                layer_name=self._layer_name,
-                output_scale=input_scale,
-                output_block_scale=output_scale_view,
-                kv_cache_dummy_dep=kv_cache_dummy_dep,
-            )
-            return at2[1], at2[2]
-
-        return _replacement
-
-    def get_inputs(self) -> list[torch.Tensor]:
-        return [
-            self.empty(5, self._num_heads, self._qk_head_dim, dtype=self._dtype),
-            self.empty(5, self._kv_lora_rank, dtype=self._dtype),
-            self.empty(5, 1, self._qk_rope_head_dim, dtype=self._dtype),
-            self.empty(5, self._output_dim, dtype=self._dtype),
-            self.empty_fp32(1, 1),
-            self.empty(0, dtype=self._dtype),
-        ]
-
-
-class MLAAttnQuantFusionPass(VllmFusionPatternMatcherPass):
-    """
-    This pass fuses post-attention quantization onto MLA attention if supported.
-
-    It uses the pattern matcher and matches each MLA layer manually, as strings
-    cannot be wildcarded. This also lets us check support on attention layers
-    upon registration instead of during pattern matching.
-    """
-
-    def __init__(self, config: VllmConfig) -> None:
-        super().__init__(config, "mla_attn_quant_fusion")
-
-        dtype = config.model_config.dtype
-        layers = list(get_layers_from_vllm_config(config, MLAAttention).values())
-
-        if len(layers) == 0:
-            logger.warning(
-                "MLA attention + quant fusion is enabled, but no MLA "
-                "attention layers were found in "
-                "CompilationConfig.static_forward_context "
-                "so no fusion patterns were registered."
-            )
-
-        for layer in layers:
-            if layer.impl.fused_output_quant_supported(kFp8StaticTensorSym):
-                self.register(MLAAttnFp8StaticQuantPattern(layer, dtype))
-
-        if current_platform.is_cuda() and hasattr(torch.ops._C, "scaled_fp4_quant"):
-            for layer in layers:
-                if layer.impl.fused_output_quant_supported(kNvfp4Dynamic):
-                    self.register(MLAAttnNvfp4QuantPattern(layer, dtype))
-
-        self.dump_patterns(config, self.pm_pass)
@@ -27,7 +27,6 @@ if rocm_aiter_ops.is_enabled():
 if current_platform.is_cuda_alike():
    from .fusion.act_quant_fusion import ActivationQuantFusionPass
    from .fusion.attn_quant_fusion import AttnQuantFusionPass
-    from .fusion.mla_attn_quant_fusion import MLAAttnQuantFusionPass
    from .fusion.qk_norm_rope_fusion import QKNormRoPEFusionPass
    from .fusion.rms_quant_fusion import RMSNormQuantFusionPass
    from .fusion.rope_kvcache_fusion import RopeKVCacheFusionPass
@@ -158,7 +157,6 @@ class PostGradPassManager(CustomGraphPass):  # type: ignore[misc]

            if self.pass_config.fuse_attn_quant:
                self.passes += [AttnQuantFusionPass(config)]
-                self.passes += [MLAAttnQuantFusionPass(config)]

            if self.pass_config.enable_qk_norm_rope_fusion:
                self.passes += [SplitCoalescingPass(config)]
@@ -121,7 +121,7 @@ class PassConfig:
    fuse_act_quant: bool = None  # type: ignore[assignment]
    """Fuse the custom SiluMul + quant ops."""
    fuse_attn_quant: bool = None  # type: ignore[assignment]
-    """Fuse the custom Attention and MLAAttention + quant ops."""
+    """Fuse the custom attention + quant ops."""
    eliminate_noops: bool = Field(default=True)
    """Eliminate no-op ops."""
    enable_sp: bool = None  # type: ignore[assignment]
@@ -521,11 +521,7 @@ class NixlConnector(KVConnectorBase_V1, SupportsHMA):
    def wait_for_save(self):
        assert self.connector_worker is not None
        assert isinstance(self._connector_metadata, NixlConnectorMetadata)
-        if (
-            self.connector_worker.use_host_buffer
-            and self.connector_worker.copy_blocks
-            and self._connector_metadata.reqs_to_save
-        ):
+        if self.connector_worker.use_host_buffer and self.connector_worker.copy_blocks:
            self.connector_worker.save_kv_to_host(self._connector_metadata)

    def shutdown(self):
@@ -2470,16 +2466,6 @@ class NixlConnectorWorker:
        Start loading by triggering non-blocking nixl_xfer.
        We check for these trnxs to complete in each step().
        """
-        # skip the empty path
-        if (
-            not metadata.reqs_to_recv
-            and not metadata.reqs_to_send
-            and not metadata.reqs_in_batch
-            and not metadata.reqs_not_processed
-            and self._ready_requests.empty()
-        ):
-            return
-
        for req_id, meta in metadata.reqs_to_recv.items():
            meta.local_physical_block_ids = self._logical_to_kernel_block_ids(
                meta.local_block_ids
@@ -6,7 +6,7 @@ from pydantic import BaseModel, Field, field_validator

 from vllm.config import ModelConfig
 from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionLogProbs
-from vllm.entrypoints.openai.engine.protocol import StreamOptions, UsageInfo
+from vllm.entrypoints.openai.engine.protocol import StreamOptions
 from vllm.logprobs import Logprob
 from vllm.renderers import TokenizeParams
 from vllm.sampling_params import SamplingParams
@@ -122,26 +122,6 @@ class GenerateResponseChoice(BaseModel):
    token_ids: list[int] | None = None


-class GenerateResponseStreamChoice(BaseModel):
-    index: int
-    logprobs: ChatCompletionLogProbs | None = None
-    finish_reason: str | None = None
-    token_ids: list[int] | None = None
-
-
-class GenerateStreamResponse(BaseModel):
-    request_id: str = Field(
-        default_factory=lambda: f"{random_uuid()}",
-        description=(
-            "The request_id related to this request. If the caller does "
-            "not set it, a random_uuid will be generated. This id is used "
-            "through out the inference process and return in response."
-        ),
-    )
-    choices: list[GenerateResponseStreamChoice]
-    usage: UsageInfo | None = Field(default=None)
-
-
 class GenerateResponse(BaseModel):
    request_id: str = Field(
        default_factory=lambda: f"{random_uuid()}",
@@ -18,7 +18,6 @@ from vllm.entrypoints.openai.chat_completion.protocol import (
 )
 from vllm.entrypoints.openai.engine.protocol import (
    ErrorResponse,
-    GenerationError,
    PromptTokenUsageInfo,
    RequestResponseMetadata,
    UsageInfo,
@@ -29,15 +28,12 @@ from vllm.entrypoints.serve.disagg.protocol import (
    GenerateRequest,
    GenerateResponse,
    GenerateResponseChoice,
-    GenerateResponseStreamChoice,
-    GenerateStreamResponse,
 )
 from vllm.entrypoints.serve.render.serving import OpenAIServingRender
-from vllm.entrypoints.utils import should_include_usage
 from vllm.logger import init_logger
 from vllm.logprobs import Logprob
 from vllm.outputs import RequestOutput
-from vllm.sampling_params import RequestOutputKind, SamplingParams
+from vllm.sampling_params import SamplingParams
 from vllm.utils.collection_utils import as_list

 logger = init_logger(__name__)
@@ -78,7 +74,7 @@ class ServingTokens(OpenAIServing):
        self,
        request: GenerateRequest,
        raw_request: Request | None = None,
-    ) -> GenerateResponse | ErrorResponse | AsyncGenerator[str, None]:
+    ) -> GenerateResponse | ErrorResponse:
        error_check_ret = await self._check_model(request)
        if error_check_ret is not None:
            logger.error("Error with model %s", error_check_ret)
@@ -114,8 +110,6 @@ class ServingTokens(OpenAIServing):
        sampling_params = request.sampling_params
        if self.force_no_detokenize:
            sampling_params.detokenize = False
-        if request.stream:
-            sampling_params.output_kind = RequestOutputKind.DELTA

        self._log_inputs(
            request_id,
@@ -139,17 +133,9 @@ class ServingTokens(OpenAIServing):
            priority=request.priority,
        )

+        # TODO(NickLucche): Implement streaming response
+
        assert result_generator is not None
-
-        if request.stream:
-            return self.serve_tokens_stream_generator(
-                request,
-                result_generator,
-                request_id,
-                model_name,
-                request_metadata,
-            )
-
        return await self.serve_tokens_full_generator(
            request, result_generator, request_id, model_name, request_metadata
        )
@@ -250,109 +236,6 @@ class ServingTokens(OpenAIServing):

        return response

-    async def serve_tokens_stream_generator(
-        self,
-        request: GenerateRequest,
-        result_generator: AsyncGenerator[RequestOutput, None],
-        request_id: str,
-        model_name: str,
-        request_metadata: RequestResponseMetadata,
-    ) -> AsyncGenerator[str, None]:
-        num_prompt_tokens = 0
-        num_generated_tokens: list[int] = []
-        first_iteration = True
-        num_cached_tokens = None
-        sampling_params: SamplingParams = request.sampling_params
-
-        include_usage, include_continuous_usage = should_include_usage(
-            request.stream_options, False
-        )
-
-        try:
-            async for res in result_generator:
-                if first_iteration:
-                    if res.prompt_token_ids is not None:
-                        num_prompt_tokens = len(res.prompt_token_ids)
-                    if res.encoder_prompt_token_ids is not None:
-                        num_prompt_tokens += len(res.encoder_prompt_token_ids)
-                    num_cached_tokens = res.num_cached_tokens
-                    num_generated_tokens = [0] * len(res.outputs)
-                    first_iteration = False
-
-                for output in res.outputs:
-                    i = output.index
-                    delta_token_ids = output.token_ids
-                    num_generated_tokens[i] += len(delta_token_ids)
-
-                    finish_reason = output.finish_reason
-                    self._raise_if_error(finish_reason, request_id)
-
-                    if not delta_token_ids:
-                        continue
-
-                    if sampling_params.logprobs is not None:
-                        out_logprobs = output.logprobs
-                        assert out_logprobs is not None, "Did not output logprobs"
-                        logprobs = self._create_tokens_logprobs(
-                            token_ids=delta_token_ids,
-                            top_logprobs=out_logprobs,
-                            num_output_top_logprobs=sampling_params.logprobs,
-                        )
-                    else:
-                        logprobs = None
-
-                    chunk = GenerateStreamResponse(
-                        request_id=request_id,
-                        choices=[
-                            GenerateResponseStreamChoice(
-                                index=i,
-                                logprobs=logprobs,
-                                finish_reason=finish_reason,
-                                token_ids=as_list(delta_token_ids),
-                            )
-                        ],
-                    )
-                    if include_continuous_usage:
-                        chunk.usage = UsageInfo(
-                            prompt_tokens=num_prompt_tokens,
-                            completion_tokens=num_generated_tokens[i],
-                            total_tokens=(num_prompt_tokens + num_generated_tokens[i]),
-                        )
-
-                    yield f"data: {chunk.model_dump_json()}\n\n"
-
-            total_completion_tokens = sum(num_generated_tokens)
-            final_usage_info = UsageInfo(
-                prompt_tokens=num_prompt_tokens,
-                completion_tokens=total_completion_tokens,
-                total_tokens=num_prompt_tokens + total_completion_tokens,
-            )
-
-            if self.enable_prompt_tokens_details and num_cached_tokens:
-                final_usage_info.prompt_tokens_details = PromptTokenUsageInfo(
-                    cached_tokens=num_cached_tokens
-                )
-
-            if include_usage:
-                final_chunk = GenerateStreamResponse(
-                    request_id=request_id,
-                    choices=[],
-                    usage=final_usage_info,
-                )
-                yield f"data: {final_chunk.model_dump_json(exclude_none=True)}\n\n"
-
-            request_metadata.final_usage_info = final_usage_info
-
-        except GenerationError as e:
-            yield (
-                f"data: {self._convert_generation_error_to_streaming_response(e)}\n\n"
-            )
-        except Exception as e:
-            logger.exception("Error in token generation stream.")
-            data = self.create_streaming_error_response(e)
-            yield f"data: {data}\n\n"
-        yield "data: [DONE]\n\n"
-
    def _create_tokens_logprobs(
        self,
        token_ids: GenericSequence[int],
@@ -191,7 +191,6 @@ if TYPE_CHECKING:
    VLLM_MQ_MAX_CHUNK_BYTES_MB: int = 16
    VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS: int = 300
    VLLM_KV_CACHE_LAYOUT: Literal["NHD", "HND"] | None = None
-    VLLM_SSM_CONV_STATE_LAYOUT: Literal["SD", "DS"] | None = None
    VLLM_COMPUTE_NANS_IN_LOGITS: bool = False
    VLLM_USE_NVFP4_CT_EMULATIONS: bool = False
    VLLM_ROCM_QUICK_REDUCE_QUANTIZATION: Literal[
@@ -1410,13 +1409,6 @@ environment_variables: dict[str, Callable[[], Any]] = {
    "VLLM_KV_CACHE_LAYOUT": env_with_choices(
        "VLLM_KV_CACHE_LAYOUT", None, ["NHD", "HND"]
    ),
-    # SSM conv state layout used for Mamba models.
-    # - SD: (state_len, dim) — dim contiguous (default)
-    # - DS: (dim, state_len) — TP-sharded dim on dim1,
-    #   consistent with SSM temporal state and HND KV cache layout.
-    "VLLM_SSM_CONV_STATE_LAYOUT": env_with_choices(
-        "VLLM_SSM_CONV_STATE_LAYOUT", None, ["SD", "DS"]
-    ),
    # Enable checking whether the generated logits contain NaNs,
    # indicating corrupted output. Useful for debugging low level bugs
    # or bad hardware but it may add compute overhead.
@@ -449,11 +449,6 @@ class MLAAttention(nn.Module, AttentionLayerBase):
            group_shape=GroupShape.PER_TENSOR,
            compile_native=True,
        )
-        self._quant_fp8_op = QuantFP8(
-            static=True,
-            group_shape=GroupShape.PER_TENSOR,
-            compile_native=True,
-        )

    @property
    def chunked_prefill_workspace_size(self) -> int:
@@ -550,19 +545,9 @@ class MLAAttention(nn.Module, AttentionLayerBase):
    ) -> torch.Tensor:
        assert output is not None, "Output tensor must be provided."

-        use_quant = output_scale is not None or output_block_scale is not None
-        if use_quant:
-            # The fusion pass has allocated output with quantized dtype
-            # (FP8 or uint8 for FP4). We can't write into it directly,
-            # so we swap in a temp buffer for computation, then quantize
-            # into the real output at the end.
-            # NOTE(carlyou): this is temporary until kernels support fp8 output
-            quant_output = output
-            output = torch.empty(
-                output.shape[0],
-                self.num_heads * self.v_head_dim,
-                dtype=q.dtype,
-                device=output.device,
+        if output_scale is not None or output_block_scale is not None:
+            raise NotImplementedError(
+                "fused output quantization is not yet supported for MLA"
            )

        if attn_metadata is None:
@@ -582,8 +567,6 @@ class MLAAttention(nn.Module, AttentionLayerBase):
            # The zero fill is required when used with DP + EP
            # to ensure all ranks within a DP group compute the
            # same expert outputs.
-            if use_quant:
-                return quant_output.fill_(0)
            return output.fill_(0)

        if self.impl.dcp_world_size == -1:
@@ -723,21 +706,6 @@ class MLAAttention(nn.Module, AttentionLayerBase):

            # v_up projection
            self._v_up_proj(attn_out, out=mqa_output_slice)
-
-        if use_quant:
-            # Quantize the BF16 computation result into the quantized output
-            actual = output[:num_actual_toks]
-            if output_block_scale is not None:
-                # NVFP4: two FP4 values packed into one uint8
-                fp4_data, fp4_scales = ops.scaled_fp4_quant(actual, output_scale)
-                quant_output[:num_actual_toks].copy_(fp4_data)
-                output_block_scale.copy_(fp4_scales)
-            else:
-                # Static FP8 quantization
-                fp8_data, _ = self._quant_fp8_op(actual, output_scale)
-                quant_output[:num_actual_toks].copy_(fp8_data)
-            return quant_output
-
        return output_padded

    def process_weights_after_loading(self, act_dtype: torch.dtype):
@@ -2101,14 +2069,6 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
    understand this class
    """

-    def fused_output_quant_supported(self, quant_key):
-        from vllm.model_executor.layers.quantization.utils.quant_utils import (
-            kFp8StaticTensorSym,
-            kNvfp4Dynamic,
-        )
-
-        return quant_key in (kFp8StaticTensorSym, kNvfp4Dynamic)
-
    def __init__(
        self,
        num_heads: int,
@@ -2553,12 +2513,8 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
                if hasattr(self.kv_b_proj, "weight")
                else self.kv_b_proj.params_dtype
            )
-            # For NVFP4, weights are packed uint8 — keep input in model dtype
-            # since the NVFP4 linear layer quantizes internally.
-            if (
-                use_fp8_prefill or _kv_b_proj_w_dtype != current_platform.fp8_dtype()
-            ) and _kv_b_proj_w_dtype != torch.uint8:
-                kv_c_normed = kv_c_normed.to(self.kv_b_proj.weight.dtype)
+            if use_fp8_prefill or _kv_b_proj_w_dtype != current_platform.fp8_dtype():
+                kv_c_normed = kv_c_normed.to(_kv_b_proj_w_dtype)

            k_pe = workspace[:toks][..., self.kv_lora_rank :].unsqueeze(1)
            kv_nope = self.kv_b_proj(kv_c_normed)[0].view(
@@ -10,7 +10,6 @@ import vllm.envs as envs
 from vllm.logger import init_logger
 from vllm.platforms import current_platform
 from vllm.triton_utils import tl, triton
-from vllm.utils.mem_utils import get_max_shared_memory_bytes
 from vllm.utils.platform_utils import num_compute_units
 from vllm.utils.torch_utils import is_torch_equal_or_newer
 from vllm.v1.attention.backends.registry import AttentionBackendEnum
@@ -178,7 +177,7 @@ def matmul_persistent(
        },
        torch.float16: {
            "BLOCK_SIZE_M": 128,
-            "BLOCK_SIZE_N": _fp16_block_size_n,
+            "BLOCK_SIZE_N": 256,
            "BLOCK_SIZE_K": 64,
            "GROUP_SIZE_M": 8,
            "num_stages": 3,
@@ -701,7 +700,7 @@ def bmm_batch_invariant(a, b, *, out=None):
        },
        torch.float16: {
            "BLOCK_SIZE_M": 128,
-            "BLOCK_SIZE_N": _fp16_block_size_n,
+            "BLOCK_SIZE_N": 256,
            "BLOCK_SIZE_K": 64,
            "num_stages": 3,
            "num_warps": 8,
@@ -753,8 +752,7 @@ def addmm_batch_invariant(bias, a, b):


 def _log_softmax_batch_invariant(input, dim, _half_to_float):
-    if _half_to_float:
-        return log_softmax(input.float(), dim=dim)
+    assert not _half_to_float, "not implemented"
    return log_softmax(input, dim=dim)


@@ -925,15 +923,12 @@ _original_fp16_reduction_precision = None
 _original_bf16_reduction_precision = None
 _original_cublas_workspace_cfg = None
 _original_cublaslt_workspace_size = None
-_fp16_block_size_n = 256


 def enable_batch_invariant_mode():
    global _batch_invariant_MODE, _batch_invariant_LIB, _original_torch_bmm
    global _original_fp16_reduction_precision, _original_bf16_reduction_precision
    global _original_cublas_workspace_cfg, _original_cublaslt_workspace_size
-    global _fp16_block_size_n
-
    if _batch_invariant_MODE:
        return

@@ -949,10 +944,6 @@ def enable_batch_invariant_mode():
        _batch_invariant_LIB.impl("aten::addmm", addmm_batch_invariant, "CUDA")
        _batch_invariant_LIB.impl("aten::matmul", matmul_batch_invariant, "CUDA")
        _batch_invariant_LIB.impl("aten::linear", linear_batch_invariant, "CUDA")
-
-        # Query the shared memory size and set block size
-        # accordingly to avoid triton OutOfResources
-        _fp16_block_size_n = 256 if get_max_shared_memory_bytes() > 106496 else 128
    else:
        # Only source of batch invariance for Hopper is split-k, can disable through
        # cuBLAS workspace config
@@ -16,7 +16,7 @@ from .chunk_scaled_dot_kkt import chunk_scaled_dot_kkt_fwd
 from .cumsum import chunk_local_cumsum
 from .l2norm import l2norm_fwd
 from .solve_tril import solve_tril
-from .utils import FLA_CHUNK_SIZE, SUPPRESS_LEVEL, input_guard
+from .utils import SUPPRESS_LEVEL, input_guard
 from .wy_fast import recompute_w_u_fwd


@@ -30,24 +30,13 @@ def chunk_gated_delta_rule_fwd(
    initial_state: torch.Tensor,
    output_final_state: bool,
    cu_seqlens: torch.Tensor | None = None,
-    chunk_indices: torch.Tensor | None = None,
-    chunk_offsets: torch.Tensor | None = None,
 ):
-    g = chunk_local_cumsum(
-        g, chunk_size=FLA_CHUNK_SIZE, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices
-    )
+    g = chunk_local_cumsum(g, chunk_size=64, cu_seqlens=cu_seqlens)
    # obtain WY representation. u is actually the new v.
    A = chunk_scaled_dot_kkt_fwd(
-        k=k,
-        beta=beta,
-        g=g,
-        cu_seqlens=cu_seqlens,
-        chunk_indices=chunk_indices,
-        output_dtype=torch.float32,
-    )
-    A = solve_tril(
-        A=A, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices, output_dtype=k.dtype
+        k=k, beta=beta, g=g, cu_seqlens=cu_seqlens, output_dtype=torch.float32
    )
+    A = solve_tril(A=A, cu_seqlens=cu_seqlens, output_dtype=k.dtype)
    w, u = recompute_w_u_fwd(
        k=k,
        v=v,
@@ -55,7 +44,6 @@ def chunk_gated_delta_rule_fwd(
        A=A,
        g_cumsum=g,
        cu_seqlens=cu_seqlens,
-        chunk_indices=chunk_indices,
    )
    h, v_new, final_state = chunk_gated_delta_rule_fwd_h(
        k=k,
@@ -65,8 +53,6 @@ def chunk_gated_delta_rule_fwd(
        initial_state=initial_state,
        output_final_state=output_final_state,
        cu_seqlens=cu_seqlens,
-        chunk_indices=chunk_indices,
-        chunk_offsets=chunk_offsets,
    )
    o = chunk_fwd_o(
        q=q,
@@ -76,7 +62,6 @@ def chunk_gated_delta_rule_fwd(
        g=g,
        scale=scale,
        cu_seqlens=cu_seqlens,
-        chunk_indices=chunk_indices,
    )
    if SUPPRESS_LEVEL < 3:
        return g, o, A, final_state, None, None, None
@@ -99,8 +84,6 @@ class ChunkGatedDeltaRuleFunction(torch.autograd.Function):
        initial_state: torch.Tensor,
        output_final_state: bool,
        cu_seqlens: torch.Tensor | None = None,
-        chunk_indices: torch.Tensor | None = None,
-        chunk_offsets: torch.Tensor | None = None,
        use_qk_l2norm_in_kernel: bool = False,
    ):
        if use_qk_l2norm_in_kernel:
@@ -117,8 +100,6 @@ class ChunkGatedDeltaRuleFunction(torch.autograd.Function):
            initial_state=initial_state,
            output_final_state=output_final_state,
            cu_seqlens=cu_seqlens,
-            chunk_indices=chunk_indices,
-            chunk_offsets=chunk_offsets,
        )
        ctx.scale = scale
        ctx.use_qk_l2norm_in_kernel = use_qk_l2norm_in_kernel
@@ -136,8 +117,6 @@ def chunk_gated_delta_rule(
    initial_state: torch.Tensor = None,
    output_final_state: bool = False,
    cu_seqlens: torch.Tensor | None = None,
-    chunk_indices: torch.Tensor | None = None,
-    chunk_offsets: torch.Tensor | None = None,
    use_qk_l2norm_in_kernel: bool = False,
 ):
    r"""
@@ -227,8 +206,6 @@ def chunk_gated_delta_rule(
        initial_state,
        output_final_state,
        cu_seqlens,
-        chunk_indices,
-        chunk_offsets,
        use_qk_l2norm_in_kernel,
    )
    return o, final_state
@@ -14,7 +14,7 @@ from vllm.triton_utils import tl, triton

 from .index import prepare_chunk_indices, prepare_chunk_offsets
 from .op import exp
-from .utils import FLA_CHUNK_SIZE, use_cuda_graph
+from .utils import use_cuda_graph

 NUM_WARPS = [2, 4, 8, 16]

@@ -286,11 +286,9 @@ def chunk_gated_delta_rule_fwd_h(
    gk: torch.Tensor | None = None,
    initial_state: torch.Tensor | None = None,
    output_final_state: bool = False,
-    chunk_size: int = FLA_CHUNK_SIZE,
+    chunk_size: int = 64,  # SY: remove this argument and force chunk size 64?
    save_new_value: bool = True,
    cu_seqlens: torch.Tensor | None = None,
-    chunk_indices: torch.Tensor | None = None,
-    chunk_offsets: torch.Tensor | None = None,
 ) -> tuple[torch.Tensor, torch.Tensor]:
    # This kernel is slightly different from fla to support Q/K with different head numbers.
    # In fla, Q/K always have the same head number, so Hg is always equal to H.
@@ -298,15 +296,20 @@ def chunk_gated_delta_rule_fwd_h(
    H = u.shape[-2]
    BT = chunk_size

-    if chunk_indices is None and cu_seqlens is not None:
-        chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size)
+    chunk_indices = (
+        prepare_chunk_indices(cu_seqlens, chunk_size)
+        if cu_seqlens is not None
+        else None
+    )
    # N: the actual number of sequences in the batch with either equal or variable lengths
    if cu_seqlens is None:
        N, NT, chunk_offsets = B, triton.cdiv(T, BT), None
    else:
-        N, NT = len(cu_seqlens) - 1, len(chunk_indices)
-        if chunk_offsets is None:
-            chunk_offsets = prepare_chunk_offsets(cu_seqlens, BT)
+        N, NT, chunk_offsets = (
+            len(cu_seqlens) - 1,
+            len(chunk_indices),
+            prepare_chunk_offsets(cu_seqlens, BT),
+        )
    assert K <= 256, "current kernel does not support head dimension larger than 256."

    h = k.new_empty(B, NT, H, V, K)
@@ -146,14 +146,14 @@ def chunk_fwd_o(
    g: torch.Tensor | None = None,  # cumsum of log decay
    scale: float | None = None,
    cu_seqlens: torch.Tensor | None = None,
-    chunk_indices: torch.Tensor | None = None,
    chunk_size: int = FLA_CHUNK_SIZE,
 ) -> torch.Tensor:
    B, T, Hg, K, V = *q.shape, v.shape[-1]
    H = v.shape[-2]
    BT = chunk_size
-    if chunk_indices is None and cu_seqlens is not None:
-        chunk_indices = prepare_chunk_indices(cu_seqlens, BT)
+    chunk_indices = (
+        prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None
+    )
    NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)
    if scale is None:
        scale = k.shape[-1] ** -0.5
@@ -14,7 +14,6 @@ from vllm.triton_utils import tl, triton

 from .index import prepare_chunk_indices
 from .op import exp
-from .utils import FLA_CHUNK_SIZE


@triton.heuristics(
@@ -104,8 +103,7 @@ def chunk_scaled_dot_kkt_fwd(
    g: torch.Tensor | None = None,
    beta: torch.Tensor | None = None,
    cu_seqlens: torch.Tensor | None = None,
-    chunk_indices: torch.Tensor | None = None,
-    chunk_size: int = FLA_CHUNK_SIZE,
+    chunk_size: int = 64,
    output_dtype: torch.dtype = torch.float32,
 ) -> torch.Tensor:
    r"""
@@ -121,9 +119,6 @@ def chunk_scaled_dot_kkt_fwd(
        cu_seqlens (torch.Tensor):
            The cumulative sequence lengths of the input tensor.
            Default: None
-        chunk_indices (torch.Tensor):
-            Pre-computed chunk indices. If None and cu_seqlens is provided,
-            computed internally. Default: None
        chunk_size (int):
            The chunk size. Default: 64.
        output_dtype (torch.dtype):
@@ -137,8 +132,9 @@ def chunk_scaled_dot_kkt_fwd(
    B, T, Hg, K = k.shape
    H = beta.shape[-1]
    BT = chunk_size
-    if chunk_indices is None and cu_seqlens is not None:
-        chunk_indices = prepare_chunk_indices(cu_seqlens, BT)
+    chunk_indices = (
+        prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None
+    )
    NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)

    A = torch.empty(B, T, H, BT, device=k.device, dtype=output_dtype)
@@ -162,7 +162,6 @@ def chunk_local_cumsum_scalar(
    chunk_size: int,
    reverse: bool = False,
    cu_seqlens: torch.Tensor | None = None,
-    chunk_indices: torch.Tensor | None = None,
    head_first: bool = False,
    output_dtype: torch.dtype | None = torch.float,
 ) -> torch.Tensor:
@@ -173,9 +172,10 @@ def chunk_local_cumsum_scalar(
    assert chunk_size == 2 ** (chunk_size.bit_length() - 1), (
        "chunk_size must be a power of 2"
    )
-    if chunk_indices is None and cu_seqlens is not None:
-        chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size)
    BT = chunk_size
+    chunk_indices = (
+        prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None
+    )
    NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)
    g_org, g = g, torch.empty_like(g, dtype=output_dtype or g.dtype)
    grid = (NT, B * H)
@@ -199,7 +199,6 @@ def chunk_local_cumsum_vector(
    chunk_size: int,
    reverse: bool = False,
    cu_seqlens: torch.Tensor | None = None,
-    chunk_indices: torch.Tensor | None = None,
    head_first: bool = False,
    output_dtype: torch.dtype | None = torch.float,
 ) -> torch.Tensor:
@@ -207,13 +206,16 @@ def chunk_local_cumsum_vector(
        B, H, T, S = g.shape
    else:
        B, T, H, S = g.shape
+    BT = chunk_size
+    chunk_indices = (
+        prepare_chunk_indices(cu_seqlens, chunk_size)
+        if cu_seqlens is not None
+        else None
+    )
+    NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)
    assert chunk_size == 2 ** (chunk_size.bit_length() - 1), (
        "chunk_size must be a power of 2"
    )
-    if chunk_indices is None and cu_seqlens is not None:
-        chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size)
-    BT = chunk_size
-    NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)

    g_org, g = g, torch.empty_like(g, dtype=output_dtype or g.dtype)

@@ -245,7 +247,6 @@ def chunk_local_cumsum(
    chunk_size: int,
    reverse: bool = False,
    cu_seqlens: torch.Tensor | None = None,
-    chunk_indices: torch.Tensor | None = None,
    head_first: bool = False,
    output_dtype: torch.dtype | None = torch.float,
    **kwargs,
@@ -256,23 +257,11 @@ def chunk_local_cumsum(
        )
    if len(g.shape) == 3:
        return chunk_local_cumsum_scalar(
-            g,
-            chunk_size,
-            reverse,
-            cu_seqlens,
-            chunk_indices,
-            head_first,
-            output_dtype,
+            g, chunk_size, reverse, cu_seqlens, head_first, output_dtype
        )
    elif len(g.shape) == 4:
        return chunk_local_cumsum_vector(
-            g,
-            chunk_size,
-            reverse,
-            cu_seqlens,
-            chunk_indices,
-            head_first,
-            output_dtype,
+            g, chunk_size, reverse, cu_seqlens, head_first, output_dtype
        )
    else:
        raise ValueError(
@@ -23,7 +23,7 @@ from .index import prepare_chunk_indices
 from .l2norm import l2norm_fwd
 from .op import exp, log
 from .solve_tril import solve_tril
-from .utils import FLA_CHUNK_SIZE, is_amd
+from .utils import is_amd

 BT_LIST_AUTOTUNE = [32, 64, 128]
 NUM_WARPS_AUTOTUNE = [2, 4, 8, 16] if is_amd else [4, 8, 16, 32]
@@ -721,7 +721,7 @@ def chunk_kda_scaled_dot_kkt_fwd(
    beta: torch.Tensor | None = None,
    scale: float | None = None,
    cu_seqlens: torch.Tensor | None = None,
-    chunk_size: int = FLA_CHUNK_SIZE,
+    chunk_size: int = 64,
    output_dtype: torch.dtype = torch.float32,
 ) -> tuple[torch.Tensor, torch.Tensor]:
    r"""
@@ -1178,7 +1178,7 @@ def chunk_kda_fwd(
    output_final_state: bool,
    cu_seqlens: torch.Tensor | None = None,
 ):
-    chunk_size = FLA_CHUNK_SIZE
+    chunk_size = 64
    g = chunk_local_cumsum(g, chunk_size=chunk_size, cu_seqlens=cu_seqlens)
    # the intra Aqk is kept in fp32
    # the computation has very marginal effect on the entire throughput
@@ -1189,7 +1189,6 @@ def chunk_kda_fwd(
        beta=beta,
        scale=scale,
        cu_seqlens=cu_seqlens,
-        chunk_size=chunk_size,
        output_dtype=torch.float32,
    )
    A = solve_tril(A=A, cu_seqlens=cu_seqlens, output_dtype=k.dtype)
@@ -507,7 +507,6 @@ def merge_16x16_to_64x64_inverse_kernel(
 def solve_tril(
    A: torch.Tensor,
    cu_seqlens: torch.Tensor | None = None,
-    chunk_indices: torch.Tensor | None = None,
    output_dtype: torch.dtype = torch.float,
 ) -> torch.Tensor:
    """
@@ -519,8 +518,6 @@ def solve_tril(
            [B, T, H, BT], where BT should only be 16, 32, or 64.
        cu_seqlens (torch.Tensor):
            The cumulative sequence lengths of the input tensor. Default: `None`.
-        chunk_indices (torch.Tensor):
-            Pre-computed chunk indices. Default: `None`.
        output_dtype (torch.dtype):
            The dtype of the output tensor. Default: `torch.float`.
            If `None`, the output dtype will be the same as the input dtype.
@@ -532,8 +529,9 @@ def solve_tril(
    output_dtype = A.dtype if output_dtype is None else output_dtype

    B, T, H, BT = A.shape
-    if chunk_indices is None and cu_seqlens is not None:
-        chunk_indices = prepare_chunk_indices(cu_seqlens, BT)
+    chunk_indices = (
+        prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None
+    )
    NT = len(chunk_indices) if cu_seqlens is not None else triton.cdiv(T, BT)

    Ai = torch.zeros_like(A, dtype=output_dtype)
@@ -123,14 +123,14 @@ def recompute_w_u_fwd(
    g_cumsum: torch.Tensor,
    A: torch.Tensor,
    cu_seqlens: torch.Tensor | None,
-    chunk_indices: torch.Tensor | None = None,
 ) -> tuple[torch.Tensor, torch.Tensor]:
    B, T, Hg, K, V = *k.shape, v.shape[-1]
    H = v.shape[-2]
    BT = A.shape[-1]

-    if chunk_indices is None and cu_seqlens is not None:
-        chunk_indices = prepare_chunk_indices(cu_seqlens, BT)
+    chunk_indices = (
+        prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None
+    )
    NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)
    BK = 64
    BV = 64
@@ -54,8 +54,8 @@ class Mxfp4MoeBackend(Enum):
    # Marlin
    BATCHED_MARLIN = "BATCHED_MARLIN"
    MARLIN = "MARLIN"
-    # ROCm AITER
-    AITER = "AITER"
+    # ROCm AITER (CK)
+    CK = "CK"
    # Triton
    TRITON = "TRITON"
    TRITON_UNFUSED = "TRITON_UNFUSED"
@@ -130,7 +130,7 @@ def backend_to_kernel_cls(

        return [BatchedMarlinExperts]

-    elif backend == Mxfp4MoeBackend.AITER:
+    elif backend == Mxfp4MoeBackend.CK:
        from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
            AiterExperts,
        )
@@ -155,7 +155,7 @@ def map_mxfp4_backend(runner_backend: str) -> Mxfp4MoeBackend:
        "flashinfer_cutlass_afp8": Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_MXFP8,
        "triton": Mxfp4MoeBackend.TRITON,
        "marlin": Mxfp4MoeBackend.MARLIN,
-        "aiter": Mxfp4MoeBackend.AITER,
+        "ck": Mxfp4MoeBackend.CK,
        "xpu": Mxfp4MoeBackend.XPU,
    }
    if backend := mapping.get(runner_backend):
@@ -173,7 +173,7 @@ def _get_priority_backends() -> list[Mxfp4MoeBackend]:
    """
    _AVAILABLE_BACKENDS = [
        Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_BF16,
-        Mxfp4MoeBackend.AITER,
+        Mxfp4MoeBackend.CK,
        Mxfp4MoeBackend.TRITON,
        Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_BF16,
        Mxfp4MoeBackend.TRITON_UNFUSED,
@@ -656,7 +656,7 @@ def convert_to_mxfp4_moe_kernel_format(
                w2_bias,
            )

-    elif mxfp4_backend == Mxfp4MoeBackend.AITER:
+    elif mxfp4_backend == Mxfp4MoeBackend.CK:
        from vllm._aiter_ops import rocm_aiter_ops

        if w13_bias is not None:
@@ -794,7 +794,7 @@ def make_mxfp4_moe_quant_config(
        Mxfp4MoeBackend.TRITON_UNFUSED,
        Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_BF16,
        Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_BF16,
-        Mxfp4MoeBackend.AITER,
+        Mxfp4MoeBackend.CK,
    ):
        return mxfp4_w4a16_moe_quant_config(
            w1_bias=w1_bias,
@@ -222,18 +222,6 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
                    self.cpu_fused_moe = cpu_fused_moe.CPUFusedMOE(layer)
            else:
                self.cpu_fused_moe = cpu_fused_moe.CPUFusedMOE(layer)
-        elif current_platform.is_xpu():
-            w13 = layer.w13_weight
-            w2 = layer.w2_weight
-
-            w13.data = w13.transpose(-1, -2).contiguous()
-            w2.data = w2.transpose(-1, -2).contiguous()
-
-            self._setup_kernel(
-                layer=layer,
-                w13=w13,
-                w2=w2,
-            )
        else:
            self._setup_kernel(
                layer=layer,
@@ -31,11 +31,7 @@ from .linear import (
    RowParallelLinear,
 )
 from .mamba.abstract import MambaBase
-from .mamba.mamba_utils import (
-    MambaStateDtypeCalculator,
-    MambaStateShapeCalculator,
-    is_conv_state_dim_first,
-)
+from .mamba.mamba_utils import MambaStateDtypeCalculator, MambaStateShapeCalculator
 from .mamba.ops.causal_conv1d import causal_conv1d_fn, causal_conv1d_update
 from .quantization.base_config import QuantizationConfig

@@ -319,12 +315,10 @@ class KimiDeltaAttention(nn.Module, MambaBase):
        beta = beta[:num_actual_tokens]

        (conv_state_q, conv_state_k, conv_state_v, recurrent_state) = constant_caches
-        # conv_state must be (..., dim, width-1) for the conv kernels.
-        # DS layout stores it that way directly; SD layout needs a transpose.
-        if not is_conv_state_dim_first():
-            conv_state_q = conv_state_q.transpose(-1, -2)
-            conv_state_k = conv_state_k.transpose(-1, -2)
-            conv_state_v = conv_state_v.transpose(-1, -2)
+        # deal with strides
+        conv_state_q = conv_state_q.transpose(-1, -2)
+        conv_state_k = conv_state_k.transpose(-1, -2)
+        conv_state_v = conv_state_v.transpose(-1, -2)

        q_conv_weights = self.q_conv1d.weight.view(
            self.q_conv1d.weight.size(0), self.q_conv1d.weight.size(2)
@@ -560,11 +560,6 @@ class RMSNormGated(CustomOp):
            activation=self.activation,
        )

-    def forward_xpu(
-        self, x: torch.Tensor, z: torch.Tensor | None = None
-    ) -> torch.Tensor:
-        return self.forward_cuda(x, z)
-

 class LayerNorm(nn.Module):
    """
@@ -41,7 +41,6 @@ from vllm.model_executor.layers.mamba.mamba_mixer2 import mamba_v2_sharded_weigh
 from vllm.model_executor.layers.mamba.mamba_utils import (
    MambaStateDtypeCalculator,
    MambaStateShapeCalculator,
-    is_conv_state_dim_first,
 )
 from vllm.model_executor.layers.mamba.ops.causal_conv1d import (
    causal_conv1d_fn,
@@ -163,8 +162,6 @@ class ChunkGatedDeltaRule(CustomOp):
        initial_state: torch.Tensor,
        output_final_state: bool,
        cu_seqlens: torch.Tensor | None = None,
-        chunk_indices: torch.Tensor | None = None,
-        chunk_offsets: torch.Tensor | None = None,
        use_qk_l2norm_in_kernel: bool = True,
    ):
        return fi_chunk_gated_delta_rule(
@@ -189,8 +186,6 @@ class ChunkGatedDeltaRule(CustomOp):
        initial_state: torch.Tensor,
        output_final_state: bool,
        cu_seqlens: torch.Tensor | None = None,
-        chunk_indices: torch.Tensor | None = None,
-        chunk_offsets: torch.Tensor | None = None,
        use_qk_l2norm_in_kernel: bool = True,
    ):
        return fla_chunk_gated_delta_rule(
@@ -202,8 +197,6 @@ class ChunkGatedDeltaRule(CustomOp):
            initial_state=initial_state,
            output_final_state=output_final_state,
            cu_seqlens=cu_seqlens,
-            chunk_indices=chunk_indices,
-            chunk_offsets=chunk_offsets,
            use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
        )

@@ -268,9 +261,6 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase):
            else 0
        )
        self.gqa_interleaved_layout = gqa_interleaved_layout
-        self._forward_method = (
-            self.forward_xpu if current_platform.is_xpu() else self.forward_cuda
-        )

        # QKV
        self.conv_dim = self.key_dim * 2 + self.value_dim
@@ -502,13 +492,6 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase):
        self,
        hidden_states: torch.Tensor,
        output: torch.Tensor,
-    ):
-        self._forward_method(hidden_states, output)
-
-    def forward_cuda(
-        self,
-        hidden_states: torch.Tensor,
-        output: torch.Tensor,
    ):
        """
        Forward pass with three parts:
@@ -583,90 +566,6 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase):
        core_attn_out = rearrange(core_attn_out, "... h d -> ... (h d)")
        output[:num_tokens], _ = self.out_proj(core_attn_out)

-    def forward_xpu(
-        self,
-        hidden_states: torch.Tensor,
-        output: torch.Tensor,
-    ):
-        """
-        Forward pass with three parts:
-        1. Input projection
-        2. Core attention (custom op)
-        3. Output projection
-        """
-        num_tokens = hidden_states.size(0)
-
-        assert not hasattr(self, "in_proj_qkv"), "lora isn't supported on XPU."
-
-        # ============================================================
-        # Part 1: Input Projection
-        # ============================================================
-        projected_states_qkvz, _ = self.in_proj_qkvz(hidden_states)
-        projected_states_ba, _ = self.in_proj_ba(hidden_states)
-
-        # ============================================================
-        # Part 2: Core Attention
-        # ============================================================
-        forward_context = get_forward_context()
-        attn_metadata: AttentionMetadata = forward_context.attn_metadata
-        core_attn_out = torch.zeros(
-            (num_tokens, self.num_v_heads // self.tp_size, self.head_v_dim),
-            dtype=hidden_states.dtype,
-            device=hidden_states.device,
-        )
-        z = torch.empty_like(core_attn_out)
-        if attn_metadata is not None:
-            attn_metadata = attn_metadata[self.prefix]
-
-            # TODO: xpu does not support this param yet
-            spec_sequence_masks = attn_metadata.spec_sequence_masks
-            assert spec_sequence_masks is None
-
-            conv_weights = self.conv1d.weight.view(
-                self.conv1d.weight.size(0), self.conv1d.weight.size(2)
-            )
-
-            conv_state = self.kv_cache[0]
-            ssm_state = self.kv_cache[1]
-
-            torch.ops._xpu_C.gdn_attention(
-                core_attn_out,
-                z,
-                projected_states_qkvz,
-                projected_states_ba,
-                self.num_k_heads,
-                self.num_v_heads,
-                self.head_k_dim,
-                self.head_v_dim,
-                conv_state=conv_state,
-                ssm_state=ssm_state,
-                conv_weights=conv_weights,
-                conv_bias=self.conv1d.bias,
-                activation=self.activation,
-                A_log=self.A_log,
-                dt_bias=self.dt_bias,
-                num_prefills=attn_metadata.num_prefills,
-                num_decodes=attn_metadata.num_decodes,
-                has_initial_state=attn_metadata.has_initial_state,
-                non_spec_query_start_loc=attn_metadata.non_spec_query_start_loc,
-                non_spec_state_indices_tensor=attn_metadata.non_spec_state_indices_tensor,
-                num_actual_tokens=attn_metadata.num_actual_tokens,
-                tp_size=self.tp_size,
-                reorder_input=not self.gqa_interleaved_layout,
-            )
-
-        # ============================================================
-        # Part 3: Output Projection
-        # ============================================================
-        z_shape_og = z.shape
-        # Reshape input data into 2D tensor
-        core_attn_out = core_attn_out.reshape(-1, core_attn_out.shape[-1])
-        z = z.reshape(-1, z.shape[-1])
-        core_attn_out = self.norm(core_attn_out, z)
-        core_attn_out = core_attn_out.reshape(z_shape_og)
-        core_attn_out = rearrange(core_attn_out, "... h d -> ... (h d)")
-        output[:num_tokens], _ = self.out_proj(core_attn_out)
-
    def _warmup_prefill_kernels(self, mixed_qkv: torch.Tensor) -> None:
        """Warm up GDN prefill kernels during V1 profiling.

@@ -800,13 +699,7 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase):
        spec_state_indices_tensor = attn_metadata.spec_state_indices_tensor  # noqa: E501
        non_spec_state_indices_tensor = attn_metadata.non_spec_state_indices_tensor  # noqa: E501
        self_kv_cache = self.kv_cache
-        # conv_state must be (..., dim, width-1) for the conv kernels.
-        # DS layout stores it that way directly; SD layout needs a transpose.
-        conv_state = (
-            self_kv_cache[0]
-            if is_conv_state_dim_first()
-            else self_kv_cache[0].transpose(-1, -2)
-        )
+        conv_state = self_kv_cache[0].transpose(-1, -2)
        ssm_state = self_kv_cache[1]
        num_actual_tokens = attn_metadata.num_actual_tokens
        num_accepted_tokens = attn_metadata.num_accepted_tokens
@@ -965,8 +858,6 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase):
                initial_state=initial_state,
                output_final_state=True,
                cu_seqlens=non_spec_query_start_loc,
-                chunk_indices=attn_metadata.chunk_indices,
-                chunk_offsets=attn_metadata.chunk_offsets,
                use_qk_l2norm_in_kernel=False,
            )
            # Init cache
@@ -1023,13 +914,7 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase):
        """
        non_spec_state_indices_tensor = attn_metadata.non_spec_state_indices_tensor  # noqa: E501
        self_kv_cache = self.kv_cache
-        # conv_state must be (..., dim, width-1) for the conv kernels.
-        # DS layout stores it that way directly; SD layout needs a transpose.
-        conv_state = (
-            self_kv_cache[0]
-            if is_conv_state_dim_first()
-            else self_kv_cache[0].transpose(-1, -2)
-        )
+        conv_state = self_kv_cache[0].transpose(-1, -2)
        ssm_state = self_kv_cache[1]
        num_actual_tokens = attn_metadata.num_actual_tokens

@@ -24,7 +24,6 @@ from vllm.model_executor.layers.mamba.abstract import MambaBase
 from vllm.model_executor.layers.mamba.mamba_utils import (
    MambaStateDtypeCalculator,
    MambaStateShapeCalculator,
-    is_conv_state_dim_first,
 )
 from vllm.model_executor.layers.mamba.ops.causal_conv1d import (
    causal_conv1d_fn,
@@ -268,12 +267,9 @@ class MambaMixer(MambaBase, PluggableLayer):
            query_start_loc_p = attn_metadata.query_start_loc_p
            state_indices_tensor_p = attn_metadata.state_indices_tensor_p
            state_indices_tensor_d = attn_metadata.state_indices_tensor_d
-            conv_state = (
-                self.kv_cache[0]
-                if is_conv_state_dim_first()
-                else self.kv_cache[0].transpose(-1, -2)
-            )
-            ssm_state = self.kv_cache[1]
+            self_kv_cache = self.kv_cache
+            conv_state = self_kv_cache[0].transpose(-1, -2)
+            ssm_state = self_kv_cache[1]
            has_initial_states_p = attn_metadata.has_initial_states_p
            cu_chunk_seqlen_p = attn_metadata.cu_chunk_seqlen_p
            last_chunk_indices_p = attn_metadata.last_chunk_indices_p
@@ -24,7 +24,6 @@ from vllm.model_executor.layers.mamba.abstract import MambaBase
 from vllm.model_executor.layers.mamba.mamba_utils import (
    MambaStateDtypeCalculator,
    MambaStateShapeCalculator,
-    is_conv_state_dim_first,
 )
 from vllm.model_executor.layers.mamba.ops.causal_conv1d import (
    causal_conv1d_fn,
@@ -576,15 +575,10 @@ class MambaMixer2(MambaBase, PluggableLayer):
            assert isinstance(attn_metadata, dict)
            attn_metadata = attn_metadata[self.prefix]
            assert isinstance(attn_metadata, Mamba2AttentionMetadata)
-            # conv_state must be (..., dim, width-1) for the conv kernels.
-            # DS layout stores it that way directly; SD layout needs a
-            # transpose (which keeps dim contiguous via stride tricks).
-            conv_state = (
-                self.kv_cache[0]
-                if is_conv_state_dim_first()
-                else self.kv_cache[0].transpose(-1, -2)
-            )
-            ssm_state = self.kv_cache[1]
+            self_kv_cache = self.kv_cache
+            # conv_state = (..., dim, width-1) yet contiguous along 'dim'
+            conv_state = self_kv_cache[0].transpose(-1, -2)
+            ssm_state = self_kv_cache[1]
            has_initial_states_p = attn_metadata.has_initial_states_p
            prep_initial_states = attn_metadata.prep_initial_states
            chunk_size = attn_metadata.chunk_size
@@ -1,52 +1,20 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

-import functools
 from collections.abc import Callable
 from dataclasses import dataclass
-from typing import Literal, TypeAlias
+from typing import TypeAlias

 import torch

-import vllm.envs as envs
 from vllm.config.cache import MambaDType
 from vllm.config.model import ModelDType
 from vllm.distributed import divide
-from vllm.logger import init_logger
 from vllm.utils.torch_utils import (
    STR_DTYPE_TO_TORCH_DTYPE,
    get_kv_cache_torch_dtype,
 )

-logger = init_logger(__name__)
-
-ConvStateLayoutType = Literal["SD", "DS"]
-
-
-@functools.lru_cache
-def get_conv_state_layout() -> ConvStateLayoutType:
-    """Return the SSM conv state layout.
-
-    SD = (state_len, dim) — dim is the innermost contiguous dimension.
-    DS = (dim, state_len) — TP-sharded dim is on dim-1 (like HND for KV
-         cache), consistent with SSM temporal state layout.
-    """
-    layout: ConvStateLayoutType | None = envs.VLLM_SSM_CONV_STATE_LAYOUT
-    if layout is not None:
-        logger.info_once(
-            "VLLM_SSM_CONV_STATE_LAYOUT env detected. "
-            "Setting SSM conv state layout to %s.",
-            layout,
-        )
-        return layout
-
-    return "SD"
-
-
-def is_conv_state_dim_first() -> bool:
-    """True when the conv state is stored as (dim, state_len) per block."""
-    return get_conv_state_layout() == "DS"
-

 class MambaStateDtypeCalculator:
    @classmethod
@@ -139,13 +107,6 @@ class MambaStateShapeCalculator:
        state_shape = (num_heads // tp_size, head_dim, head_dim)
        return (state_shape,)

-    @staticmethod
-    def _orient_conv_shape(dim: int, state_len: int) -> tuple[int, int]:
-        """Return (dim, state_len) for DS layout, (state_len, dim) for SD."""
-        if is_conv_state_dim_first():
-            return (dim, state_len)
-        return (state_len, dim)
-
    @classmethod
    def mamba1_state_shape(
        cls,
@@ -154,11 +115,12 @@ class MambaStateShapeCalculator:
        state_size: int,
        conv_kernel: int,
    ) -> tuple[tuple[int, int], tuple[int, int]]:
-        conv_dim = divide(intermediate_size, tp_world_size)
-        conv_state_shape = cls._orient_conv_shape(conv_dim, conv_kernel - 1)
+        conv_state_shape = (divide(intermediate_size, tp_world_size), conv_kernel - 1)

        temporal_state_shape = (divide(intermediate_size, tp_world_size), state_size)

+        conv_state_shape = conv_state_shape[1], conv_state_shape[0]
+
        return conv_state_shape, temporal_state_shape

    @classmethod
@@ -179,9 +141,8 @@ class MambaStateShapeCalculator:
        # heads and n_groups are TP-ed
        conv_dim = intermediate_size + 2 * n_groups * state_size

-        conv_state_shape = cls._orient_conv_shape(
-            divide(conv_dim, tp_world_size), conv_kernel - 1 + num_spec
-        )
+        # contiguous along 'dim' axis
+        conv_state_shape = (conv_kernel - 1 + num_spec, divide(conv_dim, tp_world_size))

        # These are not TP-ed as they depend on A, dt_bias, D
        # - they are typically small
@@ -197,7 +158,7 @@ class MambaStateShapeCalculator:
        conv_kernel: int,
    ) -> tuple[tuple[int, int]]:
        conv_dim = divide(intermediate_size, tp_world_size)
-        conv_state_shape = cls._orient_conv_shape(conv_dim, conv_kernel - 1)
+        conv_state_shape = (conv_kernel - 1, conv_dim)
        return (conv_state_shape,)

    @classmethod
@@ -224,11 +185,13 @@ class MambaStateShapeCalculator:
        num_spec: int = 0,
    ):
        conv_dim = head_k_dim * num_k_heads * 2 + head_v_dim * num_v_heads
-        conv_state_shape = cls._orient_conv_shape(
+        conv_state_shape = (
            divide(conv_dim, tp_world_size),
            conv_kernel_size - 1 + num_spec,
        )

+        conv_state_shape = conv_state_shape[1], conv_state_shape[0]
+
        temporal_state_shape = (
            divide(num_v_heads, tp_world_size),
            head_v_dim,
@@ -255,13 +218,12 @@ class MambaStateShapeCalculator:
        proj_size = num_heads * head_dim
        proj_k_size = num_k_heads * head_k_dim

-        conv_state_shape = cls._orient_conv_shape(
-            divide(proj_size, tp_world_size), conv_kernel_size - 1
-        )
-        conv_state_k_shape = cls._orient_conv_shape(
-            divide(proj_k_size, tp_world_size), conv_kernel_size - 1
-        )
+        conv_state_shape = (divide(proj_size, tp_world_size), conv_kernel_size - 1)
+        conv_state_k_shape = (divide(proj_k_size, tp_world_size), conv_kernel_size - 1)
        recurrent_state_shape = (divide(num_heads, tp_world_size), head_dim, head_dim)
+
+        conv_state_shape = conv_state_shape[1], conv_state_shape[0]
+        conv_state_k_shape = conv_state_k_shape[1], conv_state_k_shape[0]
        return (
            conv_state_shape,
            conv_state_k_shape,
@@ -305,27 +267,9 @@ def get_conv_copy_spec(
    cur_block_idx: int,
    num_accepted_tokens: int,
 ) -> MambaCopySpec:
-    """Return a MambaCopySpec for copying a convolutional state slice.
-
-    Works for both SD layout ``(num_blocks, state_len, dim)`` and
-    DS layout ``(num_blocks, dim, state_len)``.
-    """
+    """Return a MambaCopySpec for copying a convolutional state slice."""
    src_block_id = block_ids[cur_block_idx]
-    offset = num_accepted_tokens - 1
-    if is_conv_state_dim_first():
-        # DS layout: (num_blocks, dim, state_len) — state_len is last.
-        if offset > 0:
-            # Slicing along the last dim yields a non-contiguous view
-            # because features (dim) are strided by state_len.
-            raise NotImplementedError(
-                "DS conv state layout does not yet support speculative "
-                "decoding with mamba_cache_mode='align' "
-                "(num_accepted_tokens > 1)."
-            )
-        src_state = state[src_block_id]
-    else:
-        # SD layout: (num_blocks, state_len, dim) — dim contiguous.
-        src_state = state[src_block_id, offset:]
+    src_state = state[src_block_id, num_accepted_tokens - 1 :]
    return MambaCopySpec(
        start_addr=src_state.data_ptr(), num_elements=src_state.numel()
    )
@@ -592,6 +592,7 @@ def causal_conv1d_fn(
        stride_istate_seq = conv_states.stride(0)
        stride_istate_dim = conv_states.stride(1)
        stride_istate_token = conv_states.stride(2)
+        assert stride_istate_dim == 1
    if out.dim() == 2:
        stride_o_dim = out.stride(0)
        stride_o_token = out.stride(1)
@@ -1148,6 +1149,9 @@ def causal_conv1d_update(

    if validate_data:
        assert dim == weight.size(0)
+        assert conv_state.stride(-2) == 1, (
+            f"ERROR: expect contiguous along feat-dim of conv_state (currently stride={conv_state.stride()})"
+        )
        assert state_len >= width - 1
        # when above happens, we don't shift-left to keep any records in conv_state
        assert dim == conv_state.size(1)
@@ -17,7 +17,6 @@ from vllm.model_executor.layers.mamba.abstract import MambaBase
 from vllm.model_executor.layers.mamba.mamba_utils import (
    MambaStateDtypeCalculator,
    MambaStateShapeCalculator,
-    is_conv_state_dim_first,
 )
 from vllm.model_executor.layers.mamba.ops.causal_conv1d import (
    causal_conv1d_fn,
@@ -118,11 +117,8 @@ class ShortConv(MambaBase, CustomOp):
            assert isinstance(attn_metadata, dict)
            attn_metadata = attn_metadata[self.prefix]
            assert isinstance(attn_metadata, ShortConvAttentionMetadata)
-            conv_state = (
-                self.kv_cache[0]
-                if is_conv_state_dim_first()
-                else self.kv_cache[0].transpose(-1, -2)
-            )
+            self_kv_cache = self.kv_cache
+            conv_state = self_kv_cache[0].transpose(-1, -2)
            state_indices_tensor_p = attn_metadata.state_indices_tensor_p
            state_indices_tensor_d = attn_metadata.state_indices_tensor_d
            has_initial_states_p = attn_metadata.has_initial_states_p
@@ -8,7 +8,6 @@ from safetensors.torch import _TYPES as _SAFETENSORS_TO_TORCH_DTYPE
 from transformers import PretrainedConfig

 from vllm import _custom_ops as ops
-from vllm import envs
 from vllm.logger import init_logger
 from vllm.model_executor.layers.fused_moe.layer import FusedMoE
 from vllm.model_executor.layers.linear import (
@@ -274,9 +273,8 @@ class AWQLinearMethod(LinearMethodBase):

        # num_tokens >= threshold
        FP16_MATMUL_HEURISTIC_CONDITION = x.shape[:-1].numel() >= 256
-        # Batch invariant mode requires torch.matmul path
-        # for Triton override
-        if FP16_MATMUL_HEURISTIC_CONDITION or envs.VLLM_BATCH_INVARIANT:
+
+        if FP16_MATMUL_HEURISTIC_CONDITION:
            out = ops.awq_dequantize(qweight, scales, qzeros, 0, 0, 0)
            out = torch.matmul(reshaped_x, out)
        else:
@@ -10,7 +10,6 @@ from transformers import PretrainedConfig

 import vllm.model_executor.layers.fused_moe  # noqa
 from vllm import _custom_ops as ops
-from vllm import envs
 from vllm.logger import init_logger
 from vllm.model_executor.kernels.linear import (
    MPLinearLayerConfig,
@@ -234,11 +233,6 @@ class AWQMarlinConfig(QuantizationConfig):
    def override_quantization_method(
        cls, hf_quant_cfg, user_quant
    ) -> "QuantizationMethods | None":
-        # Skip override to marlin kernels, as they are not
-        # batch invariant
-        if envs.VLLM_BATCH_INVARIANT:
-            return None
-
        can_convert = cls.is_awq_marlin_compatible(hf_quant_cfg)
        is_valid_user_quant = (
            user_quant is None or user_quant == "marlin" or user_quant == "awq_marlin"
@@ -1028,10 +1028,6 @@ class Fp8OnlineMoEMethod(Fp8MoEMethod):
                layer.w2_weight[expert, :, :]
            )

-        if current_platform.is_xpu():
-            w13.data = w13.transpose(-1, -2).contiguous()
-            w2.data = w2.transpose(-1, -2).contiguous()
-
        # Shuffle weights to runtime format and setup kernel.
        self._setup_kernel(
            layer,
--- a/Show More
+++ b/Show More