Merge branch 'main' into patches

Signed-off-by: MOVZX <movzx@yahoo.com>

.buildkite/check-torch-abi.py

@@ -17,7 +17,6 @@ ALLOWED_UNSTABLE_LIBRARIES: tuple[str, ...] = (
     # _C_stable_libtorch, but these two modules still bind libtorch directly.
     "_C.abi3.so",
     "_rocm_C.abi3.so",
-    "third_party/deep_gemm/_C*.so",
 )
 
 

.buildkite/intel_jobs/model_executor_intel.yaml

@@ -28,7 +28,7 @@ steps:
       export VLLM_WORKER_MULTIPROC_METHOD=spawn &&
       export PYTHONFAULTHANDLER=1 &&
       cd tests &&
-      pytest -v -s model_executor -m "not slow_test" --ignore="model_executor/layers/test_rocm_unquantized_gemm.py" --ignore="model_executor/model_loader/tensorizer_loader/test_tensorizer.py" --ignore="model_executor/model_loader/test_reload.py"'
+      pytest -v -s model_executor -m "not slow_test" --ignore="model_executor/layers/test_rocm_unquantized_gemm.py" --ignore="model_executor/model_loader/tensorizer_loader/test_tensorizer.py" --ignore="model_executor/model_loader/test_reload.py" --deselect="tests/model_executor/test_enabled_custom_ops.py::test_topk_gating_enabled_by_cdna"'
 
 - label: ":intel: (BMG) Tensorizer Model Loader Tests (XPU)"
   key: tensorizer-model-loader-tests-xpu

.buildkite/test-amd.yaml

@@ -3439,6 +3439,26 @@ steps:
   commands:
   - pytest -v -s tests/kernels/test_rocm_gemm_rs_ar.py
 
+- label: ":amd: (MI355) Qwen4 Exp SP"
+  timeout_in_minutes: 30
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
+  agent_pool: mi355_4
+  num_gpus: 4
+  optional: true
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - vllm/config/parallel.py
+  - vllm/models/common/ops/sequence_parallel.py
+  - vllm/models/qwen4_exp/
+  - vllm/model_executor/layers/vocab_parallel_embedding.py
+  - vllm/model_executor/models/qwen3_next.py
+  - vllm/platforms/rocm.py
+  - tests/models/qwen4_exp/test_ple.py
+  - tests/models/qwen4_exp/test_tpsp.py
+  commands:
+  - pytest -v -s models/qwen4_exp/test_tpsp.py
+
 - label: ":amd: (MI355) Distributed DP Basic" # TBD
   timeout_in_minutes: 45
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
@@ -4684,16 +4704,32 @@ steps:
       kernels/test_fused_inv_rope_fp8_quant.py
       kernels/test_fused_minimax_m3_qknorm_rope_kv_insert.py
       kernels/test_fused_quant_activation.py
-      kernels/test_gate_linear_rocm_dispatch.py
       kernels/test_gemm_rs_ar.py
       kernels/test_kpool_decode_update_batched.py
       kernels/test_minimax_m3_amd_ops.py
       kernels/test_minimax_m3_sparse_attn_fp8_scale.py
       kernels/test_onednn.py
       kernels/test_relu2_fp8_quant.py
-      kernels/test_rocm_fp32_router_gemm.py
       --deselect=tests/kernels/test_bf16_skinny_gemm.py::test_every_dsv3_routed_shape_is_instantiated
 
+- label: ":amd: (MI355) Gate Linear Router Kernels"
+  timeout_in_minutes: 20
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
+  agent_pool: mi355_1
+  optional: true
+  working_dir: "/vllm-workspace/"
+  source_file_dependencies:
+  - vllm/model_executor/layers/fused_moe/router/gate_linear.py
+  - vllm/model_executor/layers/fused_moe/router/rocm_fp32_router_gemm.py
+  - vllm/model_executor/layers/fused_moe/router/bf16x3_router_gemm_rocm.py
+  - vllm/platforms/rocm.py
+  - tests/kernels/test_rocm_fp32_router_gemm.py
+  - tests/kernels/test_gate_linear_rocm_dispatch.py
+  - tests/kernels/test_bf16x3_router_gemm_rocm.py
+  commands:
+  - pytest -v -s tests/kernels/test_rocm_fp32_router_gemm.py tests/kernels/test_gate_linear_rocm_dispatch.py tests/kernels/test_bf16x3_router_gemm_rocm.py
+
 - label: ":amd: (MI355 DPX) Attention DiffKV Kernels"
   timeout_in_minutes: 35
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]

.buildkite/test_areas/distributed.yaml

@@ -14,6 +14,7 @@ steps:
   - tests/distributed
   commands:
   - pytest -v -s distributed/test_comm_ops.py
+  - pytest -v -s distributed/test_cpu_all_reduce.py
   - pytest -v -s distributed/test_shm_broadcast.py
   - pytest -v -s distributed/test_shm_buffer.py
   - pytest -v -s distributed/test_shm_storage.py
@@ -109,6 +110,7 @@ steps:
     amd:
       label: ":amd: (MI355) Distributed DP Basic"
       dind: false
+      automatic_shard: true
       device: mi355_2
       timeout_in_minutes: 45
       depends_on:
@@ -331,6 +333,7 @@ steps:
     amd:
       label: ':amd: (MI355) Distributed DP Extended'
       dind: false
+      automatic_shard: true
       device: mi355_4
       timeout_in_minutes: 60
       working_dir: /vllm-workspace/tests
@@ -637,7 +640,7 @@ steps:
   - tests/distributed/
   - examples/features/data_parallel/data_parallel_offline.py
   commands:
-    - ./.buildkite/scripts/run-multi-node-test.sh /vllm-workspace/tests 2 2 $IMAGE_TAG "VLLM_TEST_SAME_HOST=0 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_same_node.py | grep 'Same node test passed' && NUM_NODES=2 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_node_count.py | grep 'Node count test passed' && python3 ../examples/features/data_parallel/data_parallel_offline.py -dp=2 -tp=1 --dp-num-nodes=2 --dp-node-rank=0 --dp-master-addr=192.168.10.10 --dp-master-port=12345 --enforce-eager --trust-remote-code && VLLM_MULTI_NODE=1 pytest -v -s distributed/test_multi_node_assignment.py && VLLM_MULTI_NODE=1 pytest -v -s distributed/test_pipeline_parallel.py" "VLLM_TEST_SAME_HOST=0 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_same_node.py | grep 'Same node test passed' && NUM_NODES=2 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_node_count.py | grep 'Node count test passed' && python3 ../examples/features/data_parallel/data_parallel_offline.py -dp=2 -tp=1 --dp-num-nodes=2 --dp-node-rank=1 --dp-master-addr=192.168.10.10 --dp-master-port=12345 --enforce-eager --trust-remote-code"
+    - ./.buildkite/scripts/run-multi-node-test.sh /vllm-workspace/tests 2 2 $IMAGE_TAG "VLLM_TEST_SAME_HOST=0 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_same_node.py | grep 'Same node test passed' && NUM_NODES=2 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_node_count.py | grep 'Node count test passed' && python3 ../examples/features/data_parallel/data_parallel_offline.py -dp=2 -tp=1 --dp-num-nodes=2 --dp-node-rank=0 --dp-master-addr=192.168.10.10 --dp-master-port=12345 --enforce-eager --trust-remote-code && VLLM_MULTI_NODE=1 pytest -v -s distributed/test_pipeline_parallel.py" "VLLM_TEST_SAME_HOST=0 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_same_node.py | grep 'Same node test passed' && NUM_NODES=2 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_node_count.py | grep 'Node count test passed' && python3 ../examples/features/data_parallel/data_parallel_offline.py -dp=2 -tp=1 --dp-num-nodes=2 --dp-node-rank=1 --dp-master-addr=192.168.10.10 --dp-master-port=12345 --enforce-eager --trust-remote-code"
 
 - label: ":nvidia: (L4) Pipeline + Context Parallelism"
   key: pipeline-context-parallelism-4-gpus
@@ -706,6 +709,7 @@ steps:
     amd:
       label: ':amd: (MI355) RayExecutorV2'
       dind: false
+      automatic_shard: true
       device: mi355_4
       timeout_in_minutes: 75
       working_dir: /vllm-workspace/tests

.buildkite/test_areas/entrypoints.yaml

@@ -41,6 +41,7 @@ steps:
     amd:
       label: ":amd: (MI355 DPX) Entrypoints Unit"
       dind: false
+      automatic_shard: true
       device: mi355_dpx
       timeout_in_minutes: 50
       depends_on:
@@ -207,6 +208,7 @@ steps:
     amd:
       label: ":amd: (MI355 DPX) Entrypoints Integration (Speech to Text)"
       dind: false
+      automatic_shard: true
       device: mi355_dpx
       timeout_in_minutes: 60
       working_dir: "/vllm-workspace/tests"
@@ -252,6 +254,7 @@ steps:
     amd:
       label: ":amd: (MI355 DPX) Entrypoints Integration (Pooling)"
       dind: false
+      automatic_shard: true
       device: mi355_dpx
       timeout_in_minutes: 65
       working_dir: "/vllm-workspace/tests"

.buildkite/test_areas/kernels.yaml

@@ -138,11 +138,15 @@ steps:
   - csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu
   - vllm/models/deepseek_v4/common/ops/
   - vllm/models/deepseek_v4/nvidia/
+  - vllm/model_executor/layers/fused_moe/deep_gemm_mega_moe.py
+  - vllm/model_executor/layers/fused_moe/prepare_megamoe.py
   - vllm/utils/deep_gemm.py
   - tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py
+  - tests/kernels/moe/test_deep_gemm_mega_moe.py
   - tests/models/test_deepseek_v4_mega_moe.py
   commands:
     - pytest -v -s kernels/test_fused_deepseek_v4_*.py
+    - pytest -v -s kernels/moe/test_deep_gemm_mega_moe.py
     - pytest -v -s models/test_deepseek_v4_mega_moe.py
   mirror:
     amd:
@@ -286,7 +290,6 @@ steps:
           kernels/test_apply_repetition_penalties.py
           kernels/test_awq_int4_to_int8.py
           kernels/test_bf16_skinny_gemm.py
-          kernels/test_bf16x3_router_gemm_rocm.py
           kernels/test_cache_kernels.py
           kernels/test_compressor_kv_cache.py
           kernels/test_cp_gather_fp8.py
@@ -301,14 +304,12 @@ steps:
           kernels/test_fused_inv_rope_fp8_quant.py
           kernels/test_fused_minimax_m3_qknorm_rope_kv_insert.py
           kernels/test_fused_quant_activation.py
-          kernels/test_gate_linear_rocm_dispatch.py
           kernels/test_gemm_rs_ar.py
           kernels/test_kpool_decode_update_batched.py
           kernels/test_minimax_m3_amd_ops.py
           kernels/test_minimax_m3_sparse_attn_fp8_scale.py
           kernels/test_onednn.py
           kernels/test_relu2_fp8_quant.py
-          kernels/test_rocm_fp32_router_gemm.py
           --deselect=tests/kernels/test_bf16_skinny_gemm.py::test_every_dsv3_routed_shape_is_instantiated
 
 - label: ":nvidia: (L4) Attention Kernels Shard %N"
@@ -528,6 +529,7 @@ steps:
     amd:
       label: ":amd: (MI250) Mamba Kernels"
       dind: false
+      automatic_shard: true
       device: mi250_1
       timeout_in_minutes: 55
       depends_on:
@@ -720,6 +722,24 @@ steps:
     - nvidia-smi
     - pytest -v -s tests/kernels/test_ll_bf16_gemm.py
     - pytest -v -s tests/kernels/test_ll_fp32w_gemm.py
+  mirror:
+    amd:
+      label: ":amd: (MI355) Gate Linear Router Kernels"
+      dind: false
+      device: mi355_1
+      timeout_in_minutes: 20
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - vllm/model_executor/layers/fused_moe/router/gate_linear.py
+      - vllm/model_executor/layers/fused_moe/router/rocm_fp32_router_gemm.py
+      - vllm/model_executor/layers/fused_moe/router/bf16x3_router_gemm_rocm.py
+      - vllm/platforms/rocm.py
+      - tests/kernels/test_rocm_fp32_router_gemm.py
+      - tests/kernels/test_gate_linear_rocm_dispatch.py
+      - tests/kernels/test_bf16x3_router_gemm_rocm.py
+      commands:
+      - pytest -v -s tests/kernels/test_rocm_fp32_router_gemm.py tests/kernels/test_gate_linear_rocm_dispatch.py tests/kernels/test_bf16x3_router_gemm_rocm.py
 
 - label: ":nvidia: (DGX) Spark b12x Linear Kernels Nightly"
   key: b12x-linear-kernels-dgx-spark-nightly

.buildkite/test_areas/misc.yaml

@@ -32,6 +32,7 @@ steps:
     amd:
       label: ":amd: (MI355 DPX) V1 Sample"
       dind: false
+      automatic_shard: true
       device: mi355_dpx
       timeout_in_minutes: 70
       depends_on:
@@ -219,6 +220,7 @@ steps:
     - pytest -v -s -m 'cpu_test' v1/core
     - pytest -v -s v1/structured_output
     - pytest -v -s v1/test_serial_utils.py
+    - pytest -v -s v1/worker/test_cpu_worker_kv_budget.py
     - pytest -v -s v1/test_kv_cache_spec_registry.py
     - pytest -v -s v1/cudagraph/test_cudagraph_manager.py
     - pytest -v -s -m 'cpu_test' v1/kv_connector/unit
@@ -707,6 +709,7 @@ steps:
     amd:
       label: ":amd: (MI300) Batch Invariance"
       dind: false
+      automatic_shard: true
       device: mi300_1
       timeout_in_minutes: 60
       depends_on:

.buildkite/test_areas/models_basic.yaml

@@ -69,10 +69,11 @@ steps:
   - tests/models/test_deepseek_v4_vl_rocm.py
   - tests/models/test_deepseek_v41_decoder_replay_layers.py
   - tests/models/test_deepseek_v41_replay_batch.py
+  - tests/models/test_deepseek_v41_pcp.py
   - tests/models/transformers/test_layer_registry_e2e.py
   - tests/models/test_deepseek_v41_replay_start.py
   commands:
-    - pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py models/test_deepseek_v4_vl_rocm.py models/transformers/test_layer_registry_e2e.py models/test_deepseek_v41_replay_start.py models/test_deepseek_v41_decoder_replay_layers.py models/test_deepseek_v41_replay_batch.py
+    - pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py models/test_deepseek_v4_vl_rocm.py models/transformers/test_layer_registry_e2e.py models/test_deepseek_v41_replay_start.py models/test_deepseek_v41_decoder_replay_layers.py models/test_deepseek_v41_replay_batch.py models/test_deepseek_v41_pcp.py
   mirror:
     amd:
       label: ":amd: (MI355) Basic Models (Other)"

.buildkite/test_areas/models_distributed.yaml

@@ -18,6 +18,22 @@ steps:
   - tests/models/qwen4_exp/test_tpsp.py
   commands:
   - pytest -v -s models/qwen4_exp/test_tpsp.py
+  mirror:
+    amd:
+      label: ":amd: (MI355) Qwen4 Exp SP"
+      dind: false
+      device: mi355_4
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - vllm/config/parallel.py
+      - vllm/models/common/ops/sequence_parallel.py
+      - vllm/models/qwen4_exp/
+      - vllm/model_executor/layers/vocab_parallel_embedding.py
+      - vllm/model_executor/models/qwen3_next.py
+      - vllm/platforms/rocm.py
+      - tests/models/qwen4_exp/test_ple.py
+      - tests/models/qwen4_exp/test_tpsp.py
 
 - label: ":nvidia: (L4) Distributed Models Shard %N"
   key: distributed-model-tests-2-gpus

.buildkite/test_areas/pytorch.yaml

@@ -160,6 +160,7 @@ steps:
     amd:
       label: ':amd: (MI300) PyTorch Compilation Passes'
       dind: false
+      automatic_shard: true
       device: mi300_1
       timeout_in_minutes: 80
       working_dir: /vllm-workspace/tests

.buildkite/test_areas/quantization.yaml

@@ -11,6 +11,7 @@ steps:
   - csrc/
   - vllm/model_executor/layers/quantization
   - vllm/models/deepseek_v41/quant_config.py
+  - vllm/utils/humming.py
   - tests/quantization
   commands:
   # Pin torchao to the version compatible with the CI PyTorch and CUDA stack.
@@ -136,6 +137,7 @@ steps:
     amd:
       label: ":amd: (MI355 DPX) Quantized Models"
       dind: false
+      automatic_shard: true
       device: mi355_dpx
       timeout_in_minutes: 80
       working_dir: "/vllm-workspace/tests"

.buildkite/test_areas/spec_decode.yaml

@@ -28,6 +28,7 @@ steps:
     amd:
       label: ":amd: (MI300) V1 Spec Decode"
       dind: false
+      automatic_shard: true
       device: mi300_1
       timeout_in_minutes: 50
       depends_on:

benchmarks/kernels/benchmark_fused_topk.py

@@ -11,7 +11,7 @@ from vllm.utils.argparse_utils import FlexibleArgumentParser
 
 num_tokens_range = [2**i for i in range(0, 8, 2)]
 num_experts_range = [16, 32, 64, 128, 256, 512]
-topk_range = [3, 4]
+topk_range = [3, 4, 8, 10]
 configs = list(itertools.product(num_tokens_range, num_experts_range, topk_range))
 
 
@@ -39,9 +39,9 @@ def get_benchmark(scoring_func):
             x_names=["num_tokens", "num_experts", "topk"],
             x_vals=[list(_) for _ in configs],
             line_arg="provider",
-            line_vals=["torch", "vllm"],
-            line_names=["Torch", "vLLM"],
-            styles=[("blue", "-"), ("red", "-")],
+            line_vals=["torch", "vllm", "aiter_gating"],
+            line_names=["Torch", "vLLM", "AITER topk_gating"],
+            styles=[("blue", "-"), ("red", "-"), ("green", "-")],
             ylabel="us",
             plot_name=f"fused-topk-perf-{scoring_func}",
             args={},
@@ -70,6 +70,28 @@ def get_benchmark(scoring_func):
                 ),
                 quantiles=quantiles,
             )
+        elif provider == "aiter_gating":
+            from vllm._aiter_ops import rocm_aiter_ops
+
+            if scoring_func != "softmax" or not rocm_aiter_ops.is_topk_gating_enabled():
+                return float("nan"), float("nan"), float("nan")
+            topk_weights = torch.empty(
+                (num_tokens, topk), dtype=torch.float32, device="cuda"
+            )
+            topk_ids = torch.empty((num_tokens, topk), dtype=torch.int32, device="cuda")
+            token_expert_indices = torch.empty(
+                (num_tokens, topk), dtype=torch.int32, device="cuda"
+            )
+            ms, min_ms, max_ms = triton.testing.do_bench(
+                lambda: rocm_aiter_ops.topk_gating(
+                    topk_weights,
+                    topk_ids,
+                    token_expert_indices,
+                    gating_output,
+                    renormalize,
+                ),
+                quantiles=quantiles,
+            )
         else:
             ms, min_ms, max_ms = triton.testing.do_bench(
                 lambda: fused_topk(

cmake/external_projects/deepgemm.cmake

@@ -32,8 +32,9 @@ else()
   # Pinned to the fork's dev branch: upstream's 26/09/30 release (2.8.1)
   # plus the SM120 port, the SM90 paged-MQA kv_block=32/next_n=4 port,
   # configurable SwiGLU alpha/beta, SiTU for FP8/FP4 Mega MoE, and NVFP4
-  # Mega MoE. TORCH_LIBRARY build with a CPython 3.10 stable-ABI floor.
-  set(_DEEPGEMM_UPSTREAM_TAG "6ce1c5ec6627a6bf11dbcc698b0b1e19c9248432")
+  # Mega MoE, plus swiglustep for BF16/FP8/NVFP4 Mega MoE.
+  # Uses the PyTorch 2.11 stable ABI and CPython 3.10 limited API.
+  set(_DEEPGEMM_UPSTREAM_TAG "a7a927731b109fdfec770df7c2a5be5d7e2216b5")
 
   set(_deepgemm_fc_root "${FETCHCONTENT_BASE_DIR}")
   if(NOT _deepgemm_fc_root)

cmake/external_projects/flashmla.cmake

@@ -19,7 +19,7 @@ else()
   FetchContent_Declare(
         flashmla
         GIT_REPOSITORY https://github.com/vllm-project/FlashMLA
-        GIT_TAG 0eee43b12f034b657133cf2afca6a72ebb6efccf
+        GIT_TAG 5dceadbd9eabc9acfc547ef18ec8f7b14b980682
         GIT_PROGRESS TRUE
         CONFIGURE_COMMAND ""
         BUILD_COMMAND ""
@@ -81,72 +81,83 @@ if(FLASH_MLA_ARCHS)
         ${flashmla_SOURCE_DIR}/csrc/api/dense_decode.cpp
         ${flashmla_SOURCE_DIR}/csrc/api/fused_norm_rope_attn_rope_cast_fwd.cpp
 
-        # Misc kernels for decoding
-        ${flashmla_SOURCE_DIR}/csrc/kernels/smxx/decode/get_decoding_sched_meta/get_decoding_sched_meta.cu
-        ${flashmla_SOURCE_DIR}/csrc/kernels/smxx/decode/combine/combine.cu
-
-        # sm90 dense decode
-        ${flashmla_SOURCE_DIR}/csrc/kernels/sm90/decode/dense/instantiations/fp16.cu
-        ${flashmla_SOURCE_DIR}/csrc/kernels/sm90/decode/dense/instantiations/bf16.cu
-
-        # sm90 sparse decode
-        ${flashmla_SOURCE_DIR}/csrc/kernels/sm90/decode/sparse/instantiations/v4_persistent_h64.cu
-        ${flashmla_SOURCE_DIR}/csrc/kernels/sm90/decode/sparse/instantiations/v4_persistent_h128.cu
-        ${flashmla_SOURCE_DIR}/csrc/kernels/sm90/decode/sparse/instantiations/v32_persistent_h64.cu
-        ${flashmla_SOURCE_DIR}/csrc/kernels/sm90/decode/sparse/instantiations/v32_persistent_h128.cu
-
-        # sm90 sparse prefill
-        ${flashmla_SOURCE_DIR}/csrc/kernels/sm90/prefill/sparse/instantiations/phase1_k512.cu
-        ${flashmla_SOURCE_DIR}/csrc/kernels/sm90/prefill/sparse/instantiations/phase1_k512_topklen.cu
-        ${flashmla_SOURCE_DIR}/csrc/kernels/sm90/prefill/sparse/instantiations/phase1_k576.cu
-        ${flashmla_SOURCE_DIR}/csrc/kernels/sm90/prefill/sparse/instantiations/phase1_k576_topklen.cu
+        # Current kernels (DeepSeek V4.1, token-interleaved KV cache scales)
+        ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/smxx/decode/get_decoding_sched_meta/get_decoding_sched_meta.cu
+        ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/smxx/decode/combine/combine.cu
 
         # sm100 dense prefill (inference-only: dense backward stays out)
-        ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/dense/fmha_cutlass_fwd_sm100.cu
+        ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/sm100/prefill/dense/fmha_cutlass_fwd_sm100.cu
 
         # sm100 sparse prefill
-        ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd/head64/instantiations/phase1_h64_k512.cu
-        ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd/head64/instantiations/phase1_h64_k576.cu
-        ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd/head128/instantiations/phase1_k512.cu
-        ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd/head128/instantiations/phase1_k576.cu
-        ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_k512.cu
-
-        # sm100 fused norm + rope + attn + rope + cast
-        ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h64_prefill_norm.cu
-        ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h64_prefill_nonorm.cu
-        ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h128_prefill_norm.cu
-        ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h128_prefill_nonorm.cu
-        ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h64_decode_norm.cu
-        ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h64_decode_nonorm.cu
-        ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h128_decode_norm.cu
-        ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h128_decode_nonorm.cu
-        ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41_h64_decode_norm.cu
-        ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41_h64_decode_nonorm.cu
-        ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41_h128_decode_norm.cu
-        ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41_h128_decode_nonorm.cu
-        ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41fp4_h64_decode_norm.cu
-        ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41fp4_h64_decode_nonorm.cu
-        ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41fp4_h128_decode_norm.cu
-        ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41fp4_h128_decode_nonorm.cu
-        ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/permute_q_b_proj/kernel.cu
-        ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/permute_wv_proj/kernel.cu
+        ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/sm100/prefill/sparse/fwd/head64/instantiations/phase1_h64_k512.cu
+        ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/sm100/prefill/sparse/fwd/head128/instantiations/phase1_k512.cu
+        ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_k512.cu
 
         # sm100 sparse decode
-        ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/decode/sparse/head64/instantiations/v32_h64.cu
-        ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/decode/sparse/head64/instantiations/v32_h64_no_split.cu
-        ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/decode/sparse/head64/instantiations/v4_h64.cu
-        ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/decode/sparse/head64/instantiations/v4_h64_no_split.cu
-        ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/decode/sparse/head64/instantiations/v41_h64.cu
-        ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/decode/sparse/head64/instantiations/v41_h64_no_split.cu
-        ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/decode/sparse/head64/instantiations/v41fp4_h64.cu
-        ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/decode/sparse/head64/instantiations/v41fp4_h64_no_split.cu
-        ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/decode/sparse/nvfp4_head64/instantiations/v32_nvfp4_fp8rope.cu
-        ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512.cu
-        ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512_splitkv.cu
-        ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512_v41.cu
-        ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512_v41_splitkv.cu
-        ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512_v41fp4.cu
-        ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512_v41fp4_splitkv.cu
+        ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/sm100/decode/sparse/head64/instantiations/v41_h64.cu
+        ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/sm100/decode/sparse/head64/instantiations/v41_h64_no_split.cu
+        ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/sm100/decode/sparse/head64/instantiations/v41fp4_h64.cu
+        ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/sm100/decode/sparse/head64/instantiations/v41fp4_h64_no_split.cu
+        ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512_v41.cu
+        ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512_v41_splitkv.cu
+        ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512_v41fp4.cu
+        ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512_v41fp4_splitkv.cu
+
+        # sm100 fused norm + rope + attn + rope + cast
+        ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41_h64_prefill_norm.cu
+        ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41_h64_prefill_nonorm.cu
+        ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41_h64_decode_norm.cu
+        ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41_h64_decode_nonorm.cu
+        ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41fp4_h64_decode_norm.cu
+        ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41fp4_h64_decode_nonorm.cu
+        ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41_h128_prefill_norm.cu
+        ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41_h128_prefill_nonorm.cu
+        ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41_h128_decode_norm.cu
+        ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41_h128_decode_nonorm.cu
+        ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41fp4_h128_decode_norm.cu
+        ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41fp4_h128_decode_nonorm.cu
+        ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/permute_q_b_proj/kernel.cu
+        ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/permute_wv_proj/kernel.cu
+
+        # Legacy kernels (upstream keeps them under cuda_kernels/legacy/):
+        # V4, V3.2 and V3.2-NVFP4, and everything on Hopper
+        ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/legacy/smxx/decode/get_decoding_sched_meta/get_decoding_sched_meta.cu
+        ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/legacy/smxx/decode/combine/combine.cu
+
+        # sm90 dense decode
+        ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/legacy/sm90/decode/dense/instantiations/fp16.cu
+        ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/legacy/sm90/decode/dense/instantiations/bf16.cu
+
+        # sm90 sparse decode
+        ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/legacy/sm90/decode/sparse/instantiations/v4_persistent_h64.cu
+        ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/legacy/sm90/decode/sparse/instantiations/v4_persistent_h128.cu
+        ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/legacy/sm90/decode/sparse/instantiations/v32_persistent_h64.cu
+        ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/legacy/sm90/decode/sparse/instantiations/v32_persistent_h128.cu
+
+        # sm90 sparse prefill
+        ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/legacy/sm90/prefill/sparse/instantiations/phase1_k512.cu
+        ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/legacy/sm90/prefill/sparse/instantiations/phase1_k512_topklen.cu
+        ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/legacy/sm90/prefill/sparse/instantiations/phase1_k576.cu
+        ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/legacy/sm90/prefill/sparse/instantiations/phase1_k576_topklen.cu
+
+        # sm100 sparse prefill, d_qk = 576
+        ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/legacy/sm100/prefill/sparse/fwd/head64/instantiations/phase1_h64_k576.cu
+        ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/legacy/sm100/prefill/sparse/fwd/head128/instantiations/phase1_k576.cu
+
+        # sm100 sparse decode
+        ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/legacy/sm100/decode/sparse/head64/instantiations/v32_h64.cu
+        ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/legacy/sm100/decode/sparse/head64/instantiations/v32_h64_no_split.cu
+        ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/legacy/sm100/decode/sparse/head64/instantiations/v4_h64.cu
+        ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/legacy/sm100/decode/sparse/head64/instantiations/v4_h64_no_split.cu
+        ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/legacy/sm100/decode/sparse/nvfp4_head64/instantiations/v32_nvfp4_fp8rope.cu
+        ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/legacy/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512.cu
+        ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/legacy/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512_splitkv.cu
+
+        # sm100 fused norm + rope + attn + rope + cast, V4 decode
+        ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/legacy/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h64_decode_norm.cu
+        ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/legacy/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h64_decode_nonorm.cu
+        ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/legacy/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h128_decode_norm.cu
+        ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/legacy/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h128_decode_nonorm.cu
     )
 
     set(FlashMLA_Extension_SOURCES
@@ -158,17 +169,17 @@ if(FLASH_MLA_ARCHS)
 
     set(FlashMLA_INCLUDES
         ${flashmla_SOURCE_DIR}/csrc
-        ${flashmla_SOURCE_DIR}/csrc/kerutils/include
-        ${flashmla_SOURCE_DIR}/csrc/cutlass/include
-        ${flashmla_SOURCE_DIR}/csrc/cutlass/tools/util/include
+        ${flashmla_SOURCE_DIR}/csrc/3rdparty/kerutils/include
+        ${flashmla_SOURCE_DIR}/csrc/3rdparty/cutlass/include
+        ${flashmla_SOURCE_DIR}/csrc/3rdparty/cutlass/tools/util/include
     )
 
     set(FlashMLA_Extension_INCLUDES
         ${flashmla_SOURCE_DIR}/csrc
-        ${flashmla_SOURCE_DIR}/csrc/kerutils/include
+        ${flashmla_SOURCE_DIR}/csrc/3rdparty/kerutils/include
         ${flashmla_SOURCE_DIR}/csrc/extension/sm90/dense_fp8/
-        ${flashmla_SOURCE_DIR}/csrc/cutlass/include
-        ${flashmla_SOURCE_DIR}/csrc/cutlass/tools/util/include
+        ${flashmla_SOURCE_DIR}/csrc/3rdparty/cutlass/include
+        ${flashmla_SOURCE_DIR}/csrc/3rdparty/cutlass/tools/util/include
     )
 
     set_gencode_flags_for_srcs(
@@ -202,6 +213,12 @@ if(FLASH_MLA_ARCHS)
         target_compile_definitions(_flashmla_C PRIVATE USE_CUDA)
     endif()
 
+    # Select the CUDA branch of FlashMLA's params.h and kerutils; kerutils
+    # detects CUDA by itself in .cu files, so only host files need the flag.
+    target_compile_definitions(_flashmla_C PRIVATE
+        FLASH_MLA_IS_BUILD_ON_CUDA
+        $<$<COMPILE_LANGUAGE:CXX>:KERUTILS_IS_BUILD_ON_CUDA>)
+
     define_extension_target(
         _flashmla_extension_C
         DESTINATION vllm

csrc/cache_kernels.cu

@@ -273,7 +273,7 @@ __global__ void reshape_and_cache_kernel(
     const int64_t* __restrict__ slot_mapping,  // [num_tokens]
     const int key_stride, const int value_stride, const int num_heads,
     const int head_size, const int block_size, const int x,
-    const float* k_scale, const float* v_scale) {
+    const float* k_scale, const float* v_scale, const bool interleaved_v) {
   const int64_t token_idx = blockIdx.x;
   const int64_t slot_idx = slot_mapping[token_idx];
   if (slot_idx < 0) {
@@ -301,10 +301,20 @@ __global__ void reshape_and_cache_kernel(
       key_cache + block_idx * num_heads * h_block_count * block_size * x +
       head_idx * h_block_count * block_size * x + h_block * block_size * x +
       block_offset * x;
-  const int64_t tgt_value_start =
-      block_idx * num_heads * h_block_count * x * block_size +
-      head_idx * h_block_count * x * block_size + h_block * x * block_size +
-      block_offset;
+  // V cache target: standard [.., head_size, block_size], or the interleaved
+  // [.., block_size/x, head_size, x] the fused AITER writer/reader use.
+  const int64_t v_head_base = block_idx * num_heads * head_size * block_size +
+                              head_idx * head_size * block_size;
+  int64_t tgt_value_start;
+  int v_dim_stride;
+  if (interleaved_v) {
+    tgt_value_start = v_head_base + (block_offset / x) * (head_size * x) +
+                      (h_block * x) * x + (block_offset % x);
+    v_dim_stride = x;
+  } else {
+    tgt_value_start = v_head_base + h_block * x * block_size + block_offset;
+    v_dim_stride = block_size;
+  }
 
   constexpr int VEC_SIZE = (sizeof(scalar_t) == 2) ? 8 : 4;
   float k_scale_val = (kv_dt == Fp8KVCacheDataType::kAuto) ? 0.f : *k_scale;
@@ -318,7 +328,7 @@ __global__ void reshape_and_cache_kernel(
   cache_t* __restrict__ value_dst = value_cache + tgt_value_start;
 #pragma unroll
   for (int i = 0; i < x; i++) {
-    v_op(value_dst[i * block_size], value_src[i]);
+    v_op(value_dst[i * v_dim_stride], value_src[i]);
   }
 }
 
@@ -775,7 +785,7 @@ __global__ void cp_gather_indexer_k_quant_cache_kernel(
           slot_mapping.const_data_ptr<int64_t>(), key_stride, value_stride, \
           num_heads, head_size, block_size, x,                              \
           reinterpret_cast<const float*>(k_scale.data_ptr()),               \
-          reinterpret_cast<const float*>(v_scale.data_ptr()));
+          reinterpret_cast<const float*>(v_scale.data_ptr()), interleaved_v);
 
 void reshape_and_cache(
     torch::stable::Tensor& key,    // [num_tokens, num_heads, head_size]
@@ -783,10 +793,12 @@ void reshape_and_cache(
     torch::stable::Tensor&
         key_cache,  // [num_blocks, num_heads, head_size/x, block_size, x]
     torch::stable::Tensor&
-        value_cache,  // [num_blocks, num_heads, head_size, block_size]
+        value_cache,  // [num_blocks, num_heads, head_size, block_size], or the
+                      // interleaved [.., block_size/x, head_size, x] when
+                      // interleaved_v is set
     torch::stable::Tensor& slot_mapping,  // [num_tokens]
     const std::string& kv_cache_dtype, torch::stable::Tensor& k_scale,
-    torch::stable::Tensor& v_scale) {
+    torch::stable::Tensor& v_scale, bool interleaved_v) {
   int num_tokens = slot_mapping.size(0);
   int num_heads = key.size(1);
   int head_size = key.size(2);

csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu

@@ -32,9 +32,10 @@
  *     [0,      bs*576):          token data, 448 fp8 + 128 bf16 each
  *     [bs*576, bs*576 + bs*8):   UE8M0 scales, 7 real + 1 pad per token
  *   V4.1 (528 B/token, kv_mxfp8=true) -- RoPE is quantized as well, so the
- *   whole row is MXFP8 (fp8 e4m3 with one UE8M0 scale per 32 dims)
- *     [0,      bs*512):          token data, 512 fp8 each
- *     [bs*512, bs*512 + bs*16):  UE8M0 scales, 16 per token
+ *   whole row is MXFP8 (fp8 e4m3 with one UE8M0 scale per 32 dims), and each
+ *   token's scales follow its own data
+ *     [i*528, i*528 + 512):      token i data, 512 fp8
+ *     [i*528 + 512, (i+1)*528):  token i UE8M0 scales, 16
  */
 
 #include "torch_utils.h"
@@ -104,6 +105,8 @@ constexpr int kTokenDataBytes = kNopeDim + kRopeDim * 2;  // 448 + 128 = 576
 constexpr int kMxfp8QuantBlock = 32;
 constexpr int kMxfp8NumQuantBlocks = kHeadDim / kMxfp8QuantBlock;  // 16
 constexpr int kMxfp8TokenDataBytes = kHeadDim;                     // 512
+constexpr int kMxfp8BytesPerToken =
+    kMxfp8TokenDataBytes + kMxfp8NumQuantBlocks;  // 528
 // FNUZ on gfx942 / OCP elsewhere. FNUZ uses 224.0 (not the dtype's raw
 // 240.0) to match the rest of vLLM's FNUZ pipeline.
 #if defined(USE_ROCM) && defined(__gfx942__)
@@ -356,20 +359,20 @@ __device__ __forceinline__ void processDeepseekV4Slot(
     // V4.1 (kv_mxfp8) quantizes them too.
     int64_t const slot_id = slot_mapping[tokenIdx];
     if (slot_id >= 0) {
-      int const token_data_bytes =
-          kv_mxfp8 ? kMxfp8TokenDataBytes : kTokenDataBytes;
-      int const scale_bytes =
-          kv_mxfp8 ? kMxfp8NumQuantBlocks : kScaleBytesPerToken;
       int64_t const block_idx = slot_id / cache_block_size;
       int64_t const pos_in_block = slot_id % cache_block_size;
       uint8_t* block_base =
           k_cache + block_idx * static_cast<int64_t>(kv_block_stride);
-      uint8_t* token_fp8_ptr = block_base + pos_in_block * token_data_bytes;
+      uint8_t* token_fp8_ptr =
+          block_base +
+          pos_in_block * (kv_mxfp8 ? kMxfp8BytesPerToken : kTokenDataBytes);
       uint8_t* token_bf16_ptr = token_fp8_ptr + kNopeDim;
       uint8_t* token_scale_ptr =
-          block_base +
-          static_cast<int64_t>(cache_block_size) * token_data_bytes +
-          pos_in_block * scale_bytes;
+          kv_mxfp8
+              ? token_fp8_ptr + kMxfp8TokenDataBytes
+              : block_base +
+                    static_cast<int64_t>(cache_block_size) * kTokenDataBytes +
+                    pos_in_block * kScaleBytesPerToken;
 
 #pragma unroll
       for (int i = 0; i < kElemsPerLane; i++) {
@@ -790,7 +793,7 @@ void launchFusedDeepseekV4QNormRopeKVRopeQuantInsert(
 // backend.  Differences from the legacy path:
 //   * No Q head padding — output Q layout matches the input num_heads_q.
 //   * KV is written as a *contiguous* 512-wide row per token (token-strided),
-//     not the legacy UE8M0 paged layout with a separate scale tail.
+//     not the V4 paged record with a separate scale tail.
 //   * Q/KV are stored either as bf16 or as per-tensor E4M3 FP8 (one global
 //     scale), selected by the STORE_Q_FP8 / STORE_KV_FP8 template flags.
 //

csrc/ops.h

@@ -650,7 +650,7 @@ void reshape_and_cache(torch::stable::Tensor& key, torch::stable::Tensor& value,
                        torch::stable::Tensor& slot_mapping,
                        const std::string& kv_cache_dtype,
                        torch::stable::Tensor& k_scale,
-                       torch::stable::Tensor& v_scale);
+                       torch::stable::Tensor& v_scale, bool interleaved_v);
 
 void reshape_and_cache_flash(
     torch::stable::Tensor& key, torch::stable::Tensor& value,

csrc/quantization/activation_kernels.cu

@@ -4,6 +4,7 @@
 
 #include "core/utils.hpp"
 #include "core/cuda_compat.h"
+#include "cuda_vec_utils.cuh"
 #include "dispatch_utils.h"
 #include "quantization/w8a8/fp8/common.cuh"
 
@@ -43,7 +44,7 @@ template <typename scalar_t, scalar_t (*ACT_FN)(const scalar_t&),
 __global__ void act_and_mul_quant_kernel(
     fp8_type* __restrict__ out,          // [..., d]
     const scalar_t* __restrict__ input,  // [..., 2, d]
-    const float* scale, const int d) {
+    const float* scale, const int d, const bool vectorize) {
   const int32_t blocks_per_token = gridDim.y;
 
   const int32_t elems_per_128bit_load = (128 / 8) / sizeof(scalar_t);
@@ -64,9 +65,12 @@ __global__ void act_and_mul_quant_kernel(
   const scalar_t* __restrict__ y_ptr = input + token_idx * 2 * d + d;
   fp8_type* __restrict__ out_ptr = out + token_idx * d;
 
-  // 128-bit vectorized code
-  const int32_t vec_loop_end =
-      round_to_previous_multiple_of(elems_per_128bit_load, block_end);
+  // 128-bit vectorized code, when d and the data pointers allow 16-byte
+  // loads and 8-byte stores on every row; otherwise the cleanup loop below
+  // handles the whole block.
+  const int32_t vec_loop_end = vectorize ? round_to_previous_multiple_of(
+                                               block_end, elems_per_128bit_load)
+                                         : block_start;
   const int32_t vec_end_idx = vec_loop_end / elems_per_128bit_load;
   const int32_t vec_start_idx = block_start / elems_per_128bit_load;
 
@@ -559,26 +563,30 @@ __global__ void silu_mul_fp8_quant_deep_gemm_kernel(
 }  // namespace vllm
 
 // Launch activation, gating, and quantize kernel.
-#define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL)                             \
-  int d = input.size(-1) / 2;                                             \
-  int64_t num_tokens = input.numel() / input.size(-1);                    \
-  dim3 grid(num_tokens, num_tokens > 16 ? num_tokens > 32 ? 1 : 2 : 4);   \
-  dim3 block(std::min(d, 512));                                           \
-  const torch::stable::accelerator::DeviceGuard device_guard(             \
-      input.get_device_index());                                          \
-  const cudaStream_t stream =                                             \
-      get_current_cuda_stream(input.get_device_index());                  \
-  VLLM_STABLE_DISPATCH_FLOATING_TYPES(                                    \
-      input.scalar_type(), "act_and_mul_kernel", [&] {                    \
-        VLLM_STABLE_DISPATCH_FP8_TYPES(                                   \
-            out.scalar_type(), "act_and_mul_quant_kernel_fp8_type", [&] { \
-              vllm::act_and_mul_quant_kernel<scalar_t, KERNEL<scalar_t>,  \
-                                             fp8_t>                       \
-                  <<<grid, block, 0, stream>>>(                           \
-                      out.mutable_data_ptr<fp8_t>(),                      \
-                      input.const_data_ptr<scalar_t>(),                   \
-                      scale.const_data_ptr<float>(), d);                  \
-            });                                                           \
+#define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL)                                 \
+  int d = input.size(-1) / 2;                                                 \
+  int64_t num_tokens = input.numel() / input.size(-1);                        \
+  dim3 grid(num_tokens, num_tokens > 16 ? num_tokens > 32 ? 1 : 2 : 4);       \
+  dim3 block(std::min(d, 512));                                               \
+  const torch::stable::accelerator::DeviceGuard device_guard(                 \
+      input.get_device_index());                                              \
+  const cudaStream_t stream =                                                 \
+      get_current_cuda_stream(input.get_device_index());                      \
+  VLLM_STABLE_DISPATCH_FLOATING_TYPES(                                        \
+      input.scalar_type(), "act_and_mul_kernel", [&] {                        \
+        VLLM_STABLE_DISPATCH_FP8_TYPES(                                       \
+            out.scalar_type(), "act_and_mul_quant_kernel_fp8_type", [&] {     \
+              const scalar_t* in_ptr = input.const_data_ptr<scalar_t>();      \
+              fp8_t* out_ptr = out.mutable_data_ptr<fp8_t>();                 \
+              const bool vectorize = d % (16 / sizeof(scalar_t)) == 0 &&      \
+                                     vllm::is_16byte_aligned(in_ptr) &&       \
+                                     vllm::is_16byte_aligned(out_ptr);        \
+              vllm::act_and_mul_quant_kernel<scalar_t, KERNEL<scalar_t>,      \
+                                             fp8_t>                           \
+                  <<<grid, block, 0, stream>>>(out_ptr, in_ptr,               \
+                                               scale.const_data_ptr<float>(), \
+                                               d, vectorize);                 \
+            });                                                               \
       });
 
 void silu_and_mul_quant(torch::stable::Tensor& out,    // [..., d]

csrc/rocm/attention.cu

@@ -355,7 +355,8 @@ __device__ float warpReduceMax(float val) {
 // clang-format off
 template <typename scalar_t, typename cache_t,
           vllm::Fp8KVCacheDataType KV_DTYPE, typename OUTT, int BLOCK_SIZE,
-          int HEAD_SIZE, int NUM_THREADS, bool ALIBI_ENABLED, int GQA_RATIO, MFMAType MFMA_TYPE>
+          int HEAD_SIZE, int NUM_THREADS, bool ALIBI_ENABLED, int GQA_RATIO, MFMAType MFMA_TYPE,
+          bool USE_INTERLEAVED_V_CACHE>
 __global__
 __launch_bounds__(NUM_THREADS, 4) void paged_attention_ll4mi_QKV_mfma16_kernel(
     const scalar_t* __restrict__ q,         // [num_seqs, num_heads, head_size]
@@ -593,26 +594,59 @@ __launch_bounds__(NUM_THREADS, 4) void paged_attention_ll4mi_QKV_mfma16_kernel(
 
   _B16x8 Vlocal[VTLOOP][VHELOOP][VTLANELOOP];  // this could be B8x16 too
 
-  const cache_t* v_ptr = v_cache + wg_start_kv_head_idx * kv_head_stride +
-                         ((rowid * VTOKENS_PER_LANE) % BLOCK_SIZE);
+  // KX_V = number of contiguous token elements per interleaved chunk
+  // (e.g. 8 for bf16, 16 for fp8). This is also the 'x' dimension size.
+  constexpr int KX_V = 16 / sizeof(cache_t);
 
-  // v fetches are 16head elems across lanes x 16 tokens per lane
-  for (int vhe_depth = 0; vhe_depth < VHELOOP; vhe_depth++) {
-    const int vhead_elem = vhe_depth * NWARPS * 16 + warpid * 16 + lane16id;
-    const cache_t* v_ptr2 = v_ptr + vhead_elem * BLOCK_SIZE;
-
-    for (int vtoken_depth = 0; vtoken_depth < VTLOOP; vtoken_depth++) {
-      for (int vfetch_depth = 0; vfetch_depth < VTLANELOOP; vfetch_depth++) {
+  if constexpr (USE_INTERLEAVED_V_CACHE) {
+    // Interleaved V-cache layout: [num_blocks, num_heads, block_size/x,
+    // head_dim, x] Address for v[block][head][token t][dim d]:
+    //   chunk = t / x,  sub = t % x
+    //   addr = base + chunk * head_dim * x + d * x + sub
+    // Each 16B fetch reads KX_V contiguous elements within a single chunk, so
+    // a lane's token offset must land on a chunk boundary. This holds only
+    // when BLOCK_SIZE is a multiple of KX_V.
+    static_assert(BLOCK_SIZE % KX_V == 0,
+                  "Interleaved V-cache read requires BLOCK_SIZE % KX_V == 0");
+    const cache_t* v_ptr_head = v_cache + wg_start_kv_head_idx * kv_head_stride;
+    for (int vhe_depth = 0; vhe_depth < VHELOOP; vhe_depth++) {
+      const int vhead_elem = vhe_depth * NWARPS * 16 + warpid * 16 + lane16id;
+      for (int vtoken_depth = 0; vtoken_depth < VTLOOP; vtoken_depth++) {
         const int vblock_depth = 0;
         const int64_t vblock_number = static_cast<int64_t>(
             vphysical_block_number[vtoken_depth][vblock_depth]);
-        const cache_t* v_ptr3 = v_ptr2 + (vblock_number * kv_block_stride);
-
-        const cache_t* v_fetch_ptr =
-            v_ptr3 + vfetch_depth * CONTIGUOUS_KV_ELEMS_16B_LOAD;
-        const _B16x8* v_fetch_ptr_16B =
-            reinterpret_cast<const _B16x8*>(v_fetch_ptr);
-        Vlocal[vtoken_depth][vhe_depth][vfetch_depth] = *v_fetch_ptr_16B;
+        const cache_t* v_block_ptr =
+            v_ptr_head + vblock_number * kv_block_stride;
+        for (int vfetch_depth = 0; vfetch_depth < VTLANELOOP; vfetch_depth++) {
+          const int bo = ((rowid * VTOKENS_PER_LANE) % BLOCK_SIZE) +
+                         vfetch_depth * CONTIGUOUS_KV_ELEMS_16B_LOAD;
+          const int bo_chunk = bo / KX_V;  // which chunk of x tokens
+          const cache_t* v_fetch_ptr =
+              v_block_ptr + (bo_chunk * HEAD_SIZE + vhead_elem) * KX_V +
+              (bo % KX_V);
+          Vlocal[vtoken_depth][vhe_depth][vfetch_depth] =
+              *reinterpret_cast<const _B16x8*>(v_fetch_ptr);
+        }
+      }
+    }
+  } else {
+    const cache_t* v_ptr = v_cache + wg_start_kv_head_idx * kv_head_stride +
+                           ((rowid * VTOKENS_PER_LANE) % BLOCK_SIZE);
+    for (int vhe_depth = 0; vhe_depth < VHELOOP; vhe_depth++) {
+      const int vhead_elem = vhe_depth * NWARPS * 16 + warpid * 16 + lane16id;
+      const cache_t* v_ptr2 = v_ptr + vhead_elem * BLOCK_SIZE;
+      for (int vtoken_depth = 0; vtoken_depth < VTLOOP; vtoken_depth++) {
+        for (int vfetch_depth = 0; vfetch_depth < VTLANELOOP; vfetch_depth++) {
+          const int vblock_depth = 0;
+          const int64_t vblock_number = static_cast<int64_t>(
+              vphysical_block_number[vtoken_depth][vblock_depth]);
+          const cache_t* v_ptr3 = v_ptr2 + (vblock_number * kv_block_stride);
+          const cache_t* v_fetch_ptr =
+              v_ptr3 + vfetch_depth * CONTIGUOUS_KV_ELEMS_16B_LOAD;
+          const _B16x8* v_fetch_ptr_16B =
+              reinterpret_cast<const _B16x8*>(v_fetch_ptr);
+          Vlocal[vtoken_depth][vhe_depth][vfetch_depth] = *v_fetch_ptr_16B;
+        }
       }
     }
   }
@@ -996,7 +1030,8 @@ __launch_bounds__(NUM_THREADS, 4) void paged_attention_ll4mi_QKV_mfma16_kernel(
 template <typename scalar_t, typename cache_t,
           vllm::Fp8KVCacheDataType KV_DTYPE, typename OUTT, int BLOCK_SIZE,
           int HEAD_SIZE, int NUM_THREADS, bool ALIBI_ENABLED,
-          int GQA_RATIO>
+          int GQA_RATIO,
+          bool USE_INTERLEAVED_V_CACHE>
 __global__
 __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_mfma4_kernel(
     const scalar_t* __restrict__ q,         // [num_seqs, num_heads, head_size]
@@ -1177,25 +1212,41 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_mfma4_kernel(
       }
     }
 
+    // KX_V = number of contiguous token elements per interleaved chunk
+    constexpr int KX_V = 16 / sizeof(cache_t);
     const cache_t* v_ptr = v_cache + wg_start_kv_head_idx * kv_head_stride;
-    // fetch vcache in kv cache auto case
+    // fetch vcache in kv cache auto case (bf16/fp16)
     if constexpr (KV_DTYPE == vllm::Fp8KVCacheDataType::kAuto) {
       const _B16x8* v_ptrh8 = reinterpret_cast<const _B16x8*>(v_ptr);
-      // iterate over each v block
       for (int b = 0; b < VBLOCKS; b++) {
-        // int32 physical_block_number leads to overflow when multiplied with
-        // kv_block_stride
         const int64_t vphysical_block_number =
             static_cast<int64_t>(vphysical_blocks[b]);
         const _B16x8* v_ptrh8b =
             v_ptrh8 + (vphysical_block_number * kv_block_stride) / 8;
-        // iterate over each head elem (within head_size)
         for (int h = 0; h < VHELOOP; h++) {
           const int head_size_elem = h * WARP_SIZE + laneid;
-          const _B16x8* v_ptrh8be = v_ptrh8b + head_size_elem * BLOCK_SIZE / 8;
-          // iterate over all velems within block
-          for (int d = 0; d < BLOCK_SIZE / 8; d++) {
-            Vlocal[h][b * BLOCK_SIZE / 8 + d] = v_ptrh8be[d];
+          if constexpr (USE_INTERLEAVED_V_CACHE) {
+            // Interleaved layout: [num_blocks, num_heads, block_size/x,
+            // head_dim, x] Iterate over chunks of x tokens; within each chunk,
+            // x values for the same head_size_elem are contiguous (16-byte
+            // aligned).
+            for (int bo_chunk = 0; bo_chunk < BLOCK_SIZE / KX_V; bo_chunk++) {
+              const _B16x8* v_chunk =
+                  v_ptrh8b + (bo_chunk * HEAD_SIZE + head_size_elem) * KX_V / 8;
+              for (int xd = 0; xd < KX_V / 8; xd++) {
+                Vlocal[h][b * BLOCK_SIZE / 8 + bo_chunk * (KX_V / 8) + xd] =
+                    v_chunk[xd];
+              }
+            }
+          } else {
+            // Standard PagedAttention layout: [num_blocks, num_heads, head_dim,
+            // block_size] All block_size tokens for one head_size_elem are
+            // contiguous.
+            const _B16x8* v_ptrh8be =
+                v_ptrh8b + head_size_elem * BLOCK_SIZE / 8;
+            for (int d = 0; d < BLOCK_SIZE / 8; d++) {
+              Vlocal[h][b * BLOCK_SIZE / 8 + d] = v_ptrh8be[d];
+            }
           }
         }
       }
@@ -1203,21 +1254,29 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_mfma4_kernel(
     // fetch vcache in fp8 case
     else {  // if constexpr (KV_DTYPE != vllm::Fp8KVCacheDataType::kAuto)
       const _B8x8* v_ptrh8 = reinterpret_cast<const _B8x8*>(v_ptr);
-      // iterate over each v block
       for (int b = 0; b < VBLOCKS; b++) {
-        // int32 physical_block_number leads to overflow when multiplied with
-        // kv_block_stride
         const int64_t vphysical_block_number =
             static_cast<int64_t>(vphysical_blocks[b]);
         const _B8x8* v_ptrh8b =
             v_ptrh8 + (vphysical_block_number * kv_block_stride) / 8;
-        // iterate over each head elem (within head_size)
         for (int h = 0; h < VHELOOP; h++) {
           const int head_size_elem = h * WARP_SIZE + laneid;
-          const _B8x8* v_ptrh8be = v_ptrh8b + head_size_elem * BLOCK_SIZE / 8;
-          // iterate over all velems within block
-          for (int d = 0; d < BLOCK_SIZE / 8; d++) {
-            Vlocalb8[h][b * BLOCK_SIZE / 8 + d] = v_ptrh8be[d];
+          if constexpr (USE_INTERLEAVED_V_CACHE) {
+            // Same interleaved addressing as bf16 above, but with fp8
+            // elements (KX_V=16, using _B8x8 loads instead of _B16x8).
+            for (int bo_chunk = 0; bo_chunk < BLOCK_SIZE / KX_V; bo_chunk++) {
+              const _B8x8* v_chunk =
+                  v_ptrh8b + (bo_chunk * HEAD_SIZE + head_size_elem) * KX_V / 8;
+              for (int xd = 0; xd < KX_V / 8; xd++) {
+                Vlocalb8[h][b * BLOCK_SIZE / 8 + bo_chunk * (KX_V / 8) + xd] =
+                    v_chunk[xd];
+              }
+            }
+          } else {
+            const _B8x8* v_ptrh8be = v_ptrh8b + head_size_elem * BLOCK_SIZE / 8;
+            for (int d = 0; d < BLOCK_SIZE / 8; d++) {
+              Vlocalb8[h][b * BLOCK_SIZE / 8 + d] = v_ptrh8be[d];
+            }
           }
         }
       }
@@ -1831,7 +1890,8 @@ __device__ __forceinline__ _B16x8 from_floatx8(const floatx8& inp) {
 template <typename scalar_t, typename cache_t,
           vllm::Fp8KVCacheDataType KV_DTYPE, typename OUTT, int BLOCK_SIZE,
           int HEAD_SIZE, int NUM_THREADS, bool ALIBI_ENABLED, int GQA_RATIO,
-          MFMAType MFMA_TYPE>
+          MFMAType MFMA_TYPE,
+          bool USE_INTERLEAVED_V_CACHE>
 __global__
 __launch_bounds__(NUM_THREADS, 3) void paged_attention_ll4mi_QKV_mfma16_kernel(
     const scalar_t* __restrict__ q,       // [num_seqs, num_heads, head_size]
@@ -2285,8 +2345,8 @@ __launch_bounds__(NUM_THREADS, 3) void paged_attention_ll4mi_QKV_mfma16_kernel(
 
 template <typename scalar_t, typename cache_t,
           vllm::Fp8KVCacheDataType KV_DTYPE, typename OUTT, int BLOCK_SIZE,
-          int HEAD_SIZE, int NUM_THREADS, bool ALIBI_ENABLED,
-          int GQA_RATIO>
+          int HEAD_SIZE, int NUM_THREADS, bool ALIBI_ENABLED, int GQA_RATIO,
+          bool USE_INTERLEAVED_V_CACHE>
 __global__
 __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_mfma4_kernel(
     const scalar_t* __restrict__ q,       // [num_seqs, num_heads, head_size]
@@ -2611,7 +2671,8 @@ __device__ __forceinline__ _B16x8 from_floatx8(const floatx8& inp) {
 template <typename scalar_t, typename cache_t,
           vllm::Fp8KVCacheDataType KV_DTYPE, typename OUTT, int BLOCK_SIZE,
           int HEAD_SIZE, int NUM_THREADS, bool ALIBI_ENABLED, int GQA_RATIO,
-          MFMAType MFMA_TYPE>
+          MFMAType MFMA_TYPE,
+          bool USE_INTERLEAVED_V_CACHE>
 __global__
 __launch_bounds__(NUM_THREADS, 3) void paged_attention_ll4mi_QKV_mfma16_kernel(
     const scalar_t* __restrict__ q,       // [num_seqs, num_heads, head_size]
@@ -3030,8 +3091,8 @@ __launch_bounds__(NUM_THREADS, 3) void paged_attention_ll4mi_QKV_mfma16_kernel(
 
 template <typename scalar_t, typename cache_t,
           vllm::Fp8KVCacheDataType KV_DTYPE, typename OUTT, int BLOCK_SIZE,
-          int HEAD_SIZE, int NUM_THREADS, bool ALIBI_ENABLED,
-          int GQA_RATIO>
+          int HEAD_SIZE, int NUM_THREADS, bool ALIBI_ENABLED, int GQA_RATIO,
+          bool USE_INTERLEAVED_V_CACHE>
 __global__
 __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_mfma4_kernel(
     const scalar_t* __restrict__ q,       // [num_seqs, num_heads, head_size]
@@ -3261,7 +3322,8 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_reduce_kernel(
 template <typename scalar_t, typename cache_t,
           vllm::Fp8KVCacheDataType KV_DTYPE, typename OUTT, int BLOCK_SIZE,
           int HEAD_SIZE, int NUM_THREADS, bool ALIBI_ENABLED,
-          int GQA_RATIO, MFMAType MFMA_TYPE>
+          int GQA_RATIO, MFMAType MFMA_TYPE,
+          bool USE_INTERLEAVED_V_CACHE>
 __global__
 __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_mfma16_kernel(
     const scalar_t* __restrict__ q,         // [num_seqs, num_heads, head_size]
@@ -3288,7 +3350,8 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_mfma16_kernel(
 template <typename scalar_t, typename cache_t,
           vllm::Fp8KVCacheDataType KV_DTYPE, typename OUTT, int BLOCK_SIZE,
           int HEAD_SIZE, int NUM_THREADS, bool ALIBI_ENABLED,
-          int GQA_RATIO>
+          int GQA_RATIO,
+          bool USE_INTERLEAVED_V_CACHE>
 __global__
 __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_mfma4_kernel(
     const scalar_t* __restrict__ q,          // [num_seqs, num_heads, head_size]
@@ -3330,10 +3393,11 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_reduce_kernel(
 
 #endif
 
+// clang-format off
 #define LAUNCH_CUSTOM_ATTENTION_MFMA16(GQA_RATIO)                              \
-  paged_attention_ll4mi_QKV_mfma16_kernel<T, KVT, KV_DTYPE, OUTT, BLOCK_SIZE,  \
-                                          HEAD_SIZE, NTHR, ALIBI_ENABLED,      \
-                                          GQA_RATIO, MFMA_TYPE>                \
+  paged_attention_ll4mi_QKV_mfma16_kernel<                                     \
+      T, KVT, KV_DTYPE, OUTT, BLOCK_SIZE, HEAD_SIZE, NTHR, ALIBI_ENABLED,      \
+      GQA_RATIO, MFMA_TYPE, USE_INTERLEAVED_V_CACHE>                           \
       <<<grid, block, 0, stream>>>(                                            \
           query_ptr, key_cache_ptr, value_cache_ptr, num_kv_heads, scale,      \
           block_tables_ptr, seq_lens_ptr, query_start_loc_ptr,                 \
@@ -3342,15 +3406,16 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_reduce_kernel(
           max_ctx_blocks, k_scale_ptr, v_scale_ptr);
 
 #define LAUNCH_CUSTOM_ATTENTION_MFMA4(GQA_RATIO)                               \
-  paged_attention_ll4mi_QKV_mfma4_kernel<T, KVT, KV_DTYPE, OUTT, BLOCK_SIZE,   \
-                                         HEAD_SIZE, NTHR, ALIBI_ENABLED,       \
-                                         GQA_RATIO>                            \
+  paged_attention_ll4mi_QKV_mfma4_kernel<                                      \
+      T, KVT, KV_DTYPE, OUTT, BLOCK_SIZE, HEAD_SIZE, NTHR, ALIBI_ENABLED,      \
+      GQA_RATIO, USE_INTERLEAVED_V_CACHE>                                      \
       <<<grid, block, 0, stream>>>(                                            \
           query_ptr, key_cache_ptr, value_cache_ptr, num_kv_heads, scale,      \
           block_tables_ptr, seq_lens_ptr, query_start_loc_ptr,                 \
           max_num_blocks_per_seq, alibi_slopes_ptr, q_stride, kv_block_stride, \
           kv_head_stride, exp_sums_ptr, max_logits_ptr, tmp_out_ptr, out_ptr,  \
           max_ctx_blocks, k_scale_ptr, v_scale_ptr);
+// clang-format on
 
 #define LAUNCH_CUSTOM_REDUCTION(NPAR_LOOPS)                                 \
   paged_attention_ll4mi_reduce_kernel<T, OUTT, HEAD_SIZE, HEAD_SIZE,        \
@@ -3361,7 +3426,7 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_reduce_kernel(
 
 template <typename T, typename KVT, vllm::Fp8KVCacheDataType KV_DTYPE,
           int BLOCK_SIZE, int HEAD_SIZE, typename OUTT, int PARTITION_SIZE_OLD,
-          bool ALIBI_ENABLED, MFMAType MFMA_TYPE>
+          bool ALIBI_ENABLED, MFMAType MFMA_TYPE, bool USE_INTERLEAVED_V_CACHE>
 void paged_attention_custom_launcher(
     torch::Tensor& out, torch::Tensor& exp_sums, torch::Tensor& max_logits,
     torch::Tensor& tmp_out, torch::Tensor& query, torch::Tensor& key_cache,
@@ -3517,7 +3582,7 @@ void paged_attention_custom_launcher(
 
 template <typename T, typename KVT, vllm::Fp8KVCacheDataType KV_DTYPE,
           int BLOCK_SIZE, int HEAD_SIZE, typename OUTT, int PARTITION_SIZE_OLD,
-          bool ALIBI_ENABLED, MFMAType MFMA_TYPE>
+          bool ALIBI_ENABLED, MFMAType MFMA_TYPE, bool USE_INTERLEAVED_V_CACHE>
 void paged_attention_custom_launcher_navi(
     torch::Tensor& out, torch::Tensor& exp_sums, torch::Tensor& max_logits,
     torch::Tensor& tmp_out, torch::Tensor& query, torch::Tensor& key_cache,
@@ -3688,21 +3753,42 @@ void paged_attention_custom_launcher_navi(
   }
 }
 
-#define CALL_CUSTOM_LAUNCHER(T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE, OUTT,   \
-                             PSIZE, ALIBI_ENABLED, MFMA_TYPE)               \
-  if (!is_navi) {                                                           \
-    paged_attention_custom_launcher<T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE,  \
-                                    OUTT, PSIZE, ALIBI_ENABLED, MFMA_TYPE>( \
-        out, exp_sums, max_logits, tmp_out, query, key_cache, value_cache,  \
-        num_kv_heads, scale, block_tables, seq_lens, query_start_loc,       \
-        max_seq_len, alibi_slopes, k_scale, v_scale, fp8_out_scale);        \
-  } else {                                                                  \
-    paged_attention_custom_launcher_navi<T, KVT, KV_DTYPE, BLK_SIZE,        \
-                                         HEAD_SIZE, OUTT, PSIZE,            \
-                                         ALIBI_ENABLED, MFMA_TYPE>(         \
-        out, exp_sums, max_logits, tmp_out, query, key_cache, value_cache,  \
-        num_kv_heads, scale, block_tables, seq_lens, query_start_loc,       \
-        max_seq_len, alibi_slopes, k_scale, v_scale);                       \
+// USE_INTERLEAVED_V_CACHE gates the V-cache read path at compile time.
+// Currently only the GFX9 kernels (MI300x/MI325x/MI350x/MI355x) have the
+// actual interleaved V-fetch logic implemented.  The GFX11 (RDNA 3) and
+// GFX12 (RDNA 4) mfma16 kernels do NOT support interleaved V-cache addressing.
+// The template parameter is carried on all variants for compilation
+// compatibility; a TORCH_CHECK prevents dispatch to NAVI with interleaved V.
+// TODO: implement interleaved V-cache support in the NAVI (GFX11/GFX12)
+// kernels if shuffle KV cache layout is needed on RDNA GPUs.
+#define CALL_CUSTOM_LAUNCHER(T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE, OUTT,    \
+                             PSIZE, ALIBI_ENABLED, MFMA_TYPE)                \
+  if (use_interleaved_v_cache) {                                             \
+    TORCH_CHECK(!is_navi,                                                    \
+                "Reading interleaved V-cache layout from paged decode "      \
+                "attention not supported on NAVI GPUs");                     \
+    paged_attention_custom_launcher<T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE,   \
+                                    OUTT, PSIZE, ALIBI_ENABLED, MFMA_TYPE,   \
+                                    true>(                                   \
+        out, exp_sums, max_logits, tmp_out, query, key_cache, value_cache,   \
+        num_kv_heads, scale, block_tables, seq_lens, query_start_loc,        \
+        max_seq_len, alibi_slopes, k_scale, v_scale, fp8_out_scale);         \
+  } else {                                                                   \
+    if (!is_navi) {                                                          \
+      paged_attention_custom_launcher<T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE, \
+                                      OUTT, PSIZE, ALIBI_ENABLED, MFMA_TYPE, \
+                                      false>(                                \
+          out, exp_sums, max_logits, tmp_out, query, key_cache, value_cache, \
+          num_kv_heads, scale, block_tables, seq_lens, query_start_loc,      \
+          max_seq_len, alibi_slopes, k_scale, v_scale, fp8_out_scale);       \
+    } else {                                                                 \
+      paged_attention_custom_launcher_navi<T, KVT, KV_DTYPE, BLK_SIZE,       \
+                                           HEAD_SIZE, OUTT, PSIZE,           \
+                                           ALIBI_ENABLED, MFMA_TYPE, false>( \
+          out, exp_sums, max_logits, tmp_out, query, key_cache, value_cache, \
+          num_kv_heads, scale, block_tables, seq_lens, query_start_loc,      \
+          max_seq_len, alibi_slopes, k_scale, v_scale);                      \
+    }                                                                        \
   }
 
 #define CALL_CUSTOM_LAUNCHER_ALIBI(T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE,    \
@@ -3799,7 +3885,8 @@ void paged_attention(
     const std::string& kv_cache_dtype, torch::Tensor& k_scale,
     torch::Tensor& v_scale,
     const std::optional<torch::Tensor>& fp8_out_scale,
-    const std::string& mfma_type) {
+    const std::string& mfma_type,
+    bool use_interleaved_v_cache) {
   // clang-format on
   bool is_navi = is_navi_gpu();
   const int head_size = query.size(2);

csrc/rocm/ops.h

@@ -50,4 +50,4 @@ void paged_attention(
     int64_t max_seq_len, const std::optional<torch::Tensor>& alibi_slopes,
     const std::string& kv_cache_dtype, torch::Tensor& k_scale,
     torch::Tensor& v_scale, const std::optional<torch::Tensor>& fp8_out_scale,
-    const std::string& mfma_type);
+    const std::string& mfma_type, bool use_interleaved_v_cache);

csrc/rocm/torch_bindings.cpp

@@ -92,7 +92,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, rocm_ops) {
       "                str kv_cache_dtype,"
       "                Tensor k_scale, Tensor v_scale,"
       "                Tensor? fp8_out_scale,"
-      "                str mfma_type) -> ()");
+      "                str mfma_type,"
+      "                bool use_interleaved_v_cache=False) -> ()");
   rocm_ops.impl("paged_attention", torch::kCUDA, &paged_attention);
 }
 

csrc/torch_bindings.cpp

@@ -955,7 +955,8 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C_cache_ops, ops) {
       "                  Tensor! key_cache, Tensor! value_cache,"
       "                  Tensor slot_mapping,"
       "                  str kv_cache_dtype,"
-      "                  Tensor k_scale, Tensor v_scale) -> ()");
+      "                  Tensor k_scale, Tensor v_scale,"
+      "                  bool interleaved_v=False) -> ()");
 
   // Reshape the key and value tensors and cache them.
   ops.def(

docker/Dockerfile.rocm_base_gfx1250

@@ -1,16 +1,16 @@
 ARG BASE_IMAGE=ubuntu:24.04
-ARG ROCM_WHEEL_INDEX=https://rocm.devreleases.amd.com/whl-multi-arch/
-ARG ROCM_SDK_VERSION=7.14.0a20260623
-ARG TORCH_VERSION=2.11.0+rocm7.14.0a20260623
-ARG TORCHVISION_VERSION=0.26.0+rocm7.14.0a20260623
-ARG TORCHAUDIO_VERSION=2.11.0+rocm7.14.0a20260623
-ARG TRITON_VERSION=3.7.1+git110cd8e2.rocm7.14.0a20260623
-ARG APEX_VERSION=1.11.0+rocm7.14.0a20260623
+ARG ROCM_WHEEL_INDEX=https://rocm.nightlies.amd.com/whl-multi-arch/
+ARG ROCM_SDK_VERSION=10.1.0a20260822
+ARG TORCH_VERSION=2.11.0+rocm10.1.0a20260822
+ARG TORCHVISION_VERSION=0.26.0+rocm10.1.0a20260822
+ARG TORCHAUDIO_VERSION=2.11.0+rocm10.1.0a20260822
+ARG TRITON_VERSION=3.8.0+git4cff872c.rocm10.1.0a20260822
+ARG APEX_VERSION=1.11.0+rocm10.1.0a20260822
 
 
 ARG FA_BRANCH="jpvillam/gfx1250_wip"
 ARG FA_REPO="https://github.com/jpvillam-amd/flash-attention.git"
-ARG AITER_BRANCH="main"
+ARG AITER_BRANCH="f85b24b"
 ARG AITER_REPO="https://github.com/ROCm/aiter.git"
 ARG MORI_BRANCH="v1.1.0"
 ARG MORI_REPO="https://github.com/ROCm/mori.git"
@@ -119,10 +119,10 @@ ENV PYTHONPATH=${SITE_PACKAGES}/_rocm_sdk_core/share/amd_smi
 # which keeps the set consistent and avoids pip backtracking.
 RUN pip install --pre --index-url "${ROCM_WHEEL_INDEX}" \
     --extra-index-url https://pypi.org/simple \
-    "torch[device-all]==${TORCH_VERSION}" \
+    "torch[device-gfx1250]==${TORCH_VERSION}" \
     "torchvision==${TORCHVISION_VERSION}" \
     "torchaudio==${TORCHAUDIO_VERSION}" \
-    "rocm[libraries,devel,device-all]==${ROCM_SDK_VERSION}" && \
+    "rocm[libraries,devel,device-gfx1250]==${ROCM_SDK_VERSION}" && \
     rocm-sdk init
  
 # Torch runtime deps that may not be published on the ROCm wheel index;
@@ -200,15 +200,6 @@ ENV SCCACHE_S3_NO_CREDENTIALS=${USE_SCCACHE:+${SCCACHE_S3_NO_CREDENTIALS}}
 ENV SCCACHE_IDLE_TIMEOUT=${USE_SCCACHE:+0}
 
 
-###
-### AMD SMI Build
-###
-FROM base AS build_amdsmi
-RUN cd ${SITE_PACKAGES}/_rocm_sdk_core/share/amd_smi \
-    && pip wheel . --wheel-dir=dist
-RUN mkdir -p /app/install && cp ${SITE_PACKAGES}/_rocm_sdk_core/share/amd_smi/dist/*.whl /app/install
-
-
 ###
 ### MORI Build TODO(Build needs fixing)
 ###
@@ -253,7 +244,10 @@ FROM base AS build_aiter
 ARG AITER_BRANCH
 ARG AITER_REPO
 ARG USE_SCCACHE
-RUN git clone --recursive --branch ${AITER_BRANCH} ${AITER_REPO}
+
+RUN git clone ${AITER_REPO} \
+    && cd aiter \
+    && git checkout --detach ${AITER_BRANCH}
 RUN cd aiter \
     && git submodule update --init --recursive \
     && pip install -r requirements.txt
@@ -278,8 +272,6 @@ FROM base AS debs_wheel_release
 RUN mkdir /app/debs
 RUN --mount=type=bind,from=build_fa,src=/app/install/,target=/install \
     if ls /install/*.whl >/dev/null 2>&1; then cp /install/*.whl /app/debs; fi
-RUN --mount=type=bind,from=build_amdsmi,src=/app/install/,target=/install \
-    cp /install/*.whl /app/debs
 RUN --mount=type=bind,from=build_aiter,src=/app/install/,target=/install \
     cp /install/*.whl /app/debs
 
@@ -288,8 +280,6 @@ FROM base AS debs
 RUN mkdir /app/debs
 RUN --mount=type=bind,from=build_fa,src=/app/install/,target=/install \
     if ls /install/*.whl >/dev/null 2>&1; then cp /install/*.whl /app/debs; fi
-RUN --mount=type=bind,from=build_amdsmi,src=/app/install/,target=/install \
-    cp /install/*.whl /app/debs
 RUN --mount=type=bind,from=build_aiter,src=/app/install/,target=/install \
     cp /install/*.whl /app/debs
 RUN --mount=type=bind,from=build_mori,src=/app/install/,target=/install \

docker/Dockerfile.rocm_gfx1250

@@ -15,6 +15,7 @@ ARG NIC_BACKEND=all
 # Users can specify a custom version compatible with their host drivers.
 # The default version has been tested with ionic-dkms=26.03.3.001
 # (libionic1=54.0-187-1), required for MoRI-EP/NIXL on Pollara/AINIC.
+# TODO (jmathias) : Check if we can revert this Pinned AINIC version
 ARG AINIC_VERSION=1.117.5-a-77
 ARG UBUNTU_CODENAME=jammy
 
@@ -115,8 +116,8 @@ WORKDIR ${COMMON_WORKDIR}
 FROM base AS fetch_vllm_0
 ONBUILD COPY ./ vllm/
 FROM base AS fetch_vllm_1
-ARG VLLM_REPO="https://github.com/ROCm/vllm.git"
-ARG VLLM_BRANCH="455_wip"
+ARG VLLM_REPO="https://github.com/vllm-project/vllm.git"
+ARG VLLM_BRANCH="main"
 ENV VLLM_REPO=${VLLM_REPO}
 ENV VLLM_BRANCH=${VLLM_BRANCH}
 ONBUILD RUN git clone ${VLLM_REPO} \
@@ -726,13 +727,14 @@ RUN echo "VLLM_BASE_IMAGE=${BASE_IMAGE}" >> ${COMMON_WORKDIR}/versions.txt \
 
 # Download bench scripts for gfx1250
 RUN curl -k -O https://raw.githubusercontent.com/ROCm/vllm/refs/heads/gfx1250_bench/vllm_smoketest.sh
+RUN curl -k -O https://raw.githubusercontent.com/ROCm/vllm/refs/heads/gfx1250_bench/vllm_benchserve.sh
 
 
 ### Install triton from upstream for AITER Deps TODO: (JPVILLAM) If possible to get this on whls it would be better
 RUN pip3 uninstall -y triton && \
     git clone https://github.com/triton-lang/triton.git && \
     cd triton && \
-    git checkout c517f38c && \
+    git checkout aedd1f6 && \
     TRITON_APPEND_CMAKE_ARGS="-DCMAKE_FIND_USE_CMAKE_ENVIRONMENT_PATH=FALSE" pip3 install .
 
 CMD ["/bin/bash"]

docker/Dockerfile.xpu

@@ -170,8 +170,8 @@ FROM vllm-base AS ucx-nixl-build
 # Keep UCX_VERSION aligned with the UCX that NIXL itself builds against for this
 # NIXL_VERSION (nixl contrib/Dockerfile.manylinux: ARG UCX_REF), so XPU runs the
 # same UCX commit as the released nixl-cu* wheels rather than an untested pairing.
-ARG UCX_VERSION=v1.21.x
-ARG NIXL_VERSION=v1.3.2
+ARG UCX_VERSION=v1.23.x
+ARG NIXL_VERSION=v1.5.0
 
 # Build-time only: compiler, autotools, and verbs dev headers
 RUN apt-get update -y && apt-get install -y --no-install-recommends \
@@ -209,7 +209,7 @@ ENV UV_COMPILE_BYTECODE=1
 
 # Must match the ucx-nixl-build stage above: ARG does not cross stages, so a stale
 # value here installs a meta package that disagrees with the wheel built there.
-ARG NIXL_VERSION=v1.3.2
+ARG NIXL_VERSION=v1.5.0
 
 # Copy compiled UCX runtime libraries and the pre-built NIXL wheel.
 # No compiler or autotools are installed in this stage.

docs/design/cuda_graphs_multimodal.md

@@ -12,6 +12,7 @@ For two-tower vision encoders (e.g., DeepSeek-OCR's SAM + CLIP with dynamic tili
 | Architecture | Models | CG for Image | CG for Video | Multi-Path Graph |
 | ------------ | ------ | ------------ | ------------ | --------------- |
 | `DeepseekOCRForCausalLM` | `DeepSeek-OCR` | ✅︎ | ❌︎ | ✅︎ |
+| `DeepseekVLV2ForCausalLM` | `DeepSeek-VL2` | ✅︎ | ❌︎ | ❌︎ |
 | `DeepseekV41ForCausalLM` | `DeepSeek-V4.1-Flash` | ✅︎ | ❌︎ | ❌︎ |
 | `Ernie4_5_VLMoeForConditionalGeneration` | `ERNIE-4.5-VL` | ✅︎ | ❌︎ | ❌︎ |
 | `Gemma3ForConditionalGeneration` | `Gemma3` | ✅︎ | ❌︎ | ❌︎ |

docs/features/speculative_decoding/mtp.md

@@ -74,3 +74,12 @@ vllm serve XiaomiMiMo/MiMo-7B-Base \
   is a good default to start with.
 - If your model does not support MTP, use another method such as EAGLE or draft
   model speculation.
+- When the MTP head shares the target's unquantized lm_head (most MTP models),
+  `"draft_lm_head_quantization": "fp8"` or `"nvfp4"` in `--speculative-config`
+  gives the drafter a quantized copy of that head (Model Runner V2). The target
+  still verifies with its own head, so outputs are unchanged; only the
+  acceptance length can move. It helps most on bandwidth-bound GPUs with a
+  large vocabulary. The copy costs about one byte (FP8) or 0.56 bytes (NVFP4)
+  per lm_head weight, e.g. 0.6 or 0.35 GiB for a 248k x 2560 head; it is
+  counted in the model weights, so with `--kv-cache-memory-bytes` leave room
+  for it.

docs/features/watermarking.md

@@ -125,8 +125,10 @@ generation.
 ## Speculative decoding
 
 Watermarking requires speculative decoding to use probabilistic draft sampling,
-standard rejection sampling, and an autoregressive model-based method (`dspark`,
-`eagle`, `eagle3`, or `mtp`). Parallel drafting is supported only by `dspark`.
+standard rejection sampling, and a model-based method (`dflash`, `dspark`,
+`eagle`, `eagle3`, or `mtp`). Parallel drafting is supported only by `dflash`
+and `dspark`. DFlash drafters with a candidate head (`DFlash2DraftModel`,
+`LiLiCorrDraftModel`) are not supported.
 
 A watermarking algorithm without native speculative-decoding support is
 rejected before model loading. Set

docs/models/structured_decisions_models.md

@@ -304,12 +304,16 @@ vllm serve Cloudflare/clef-omni
 
 ### Decision 3.0
 
-The Decision 3.0 models ([`vllm-sr/d3`](https://huggingface.co/vllm-sr/d3) and
-its smaller variants such as `vllm-sr/d3-mini` and `vllm-sr/d3-edge`) are a
-Qwen3.5 backbone with an answer-code readout: one forward pass per question,
-and the readout scores the question's options at the last token. vLLM detects
-the checkpoint's `decision_config.json` and serves it as a classification
-model.
+The Decision 3.0 models are a Qwen3.5 backbone with an answer-code readout:
+one forward pass per question, and the readout scores the question's options
+at the last token. vLLM detects the checkpoint's `decision_config.json` and
+serves it as a classification model. All six sizes take the same command:
+[`vllm-sr/d3`](https://huggingface.co/vllm-sr/d3) (27B),
+[`vllm-sr/d3-flash`](https://huggingface.co/vllm-sr/d3-flash) (9B),
+[`vllm-sr/d3-mini`](https://huggingface.co/vllm-sr/d3-mini) (4B),
+[`vllm-sr/d3-nano`](https://huggingface.co/vllm-sr/d3-nano) (2B),
+[`vllm-sr/d3-lite`](https://huggingface.co/vllm-sr/d3-lite) (0.8B) and
+[`vllm-sr/d3-edge`](https://huggingface.co/vllm-sr/d3-edge) (0.6B).
 
 ```bash
 vllm serve vllm-sr/d3

docs/serving/online_serving/derenderer.md

@@ -89,7 +89,7 @@ The derenderer builds its tool and reasoning parsers from its own server flags p
 
 A mismatch doesn't fail. The parser just splits `reasoning`, `content` and `tool_calls` differently from what `vllm serve` would return.
 
-`/v1/completions/derender` only detokenizes. It never runs tool or reasoning parsers, the same as `/v1/completions` on `vllm serve`. It only reads `skip_special_tokens` from `completion_request`.
+`/v1/completions/derender` only detokenizes. It never runs tool or reasoning parsers, the same as `/v1/completions` on `vllm serve`. From `completion_request` it reads `skip_special_tokens`, and `logprobs` and `logprob_token_ids` to cut `top_logprobs` the way `/v1/completions` does.
 
 ## Streaming
 
@@ -105,7 +105,7 @@ The server keeps no state between calls. Everything the next call needs is in `s
 - **Send `prompt_token_ids` on the first chunk, even without a parser.** It seeds detokenization so the first token's leading space matches the coupled endpoint (see [Request format](#request-format)). Later chunks already carry that context in `stream_state`.
 - **Don't forward `[DONE]`.** It marks the end of the generate stream and isn't a chunk.
 
-`/inference/v1/generate` returns `GenerateLogProbs` with integer `token_id`s (the generate server has no tokenizer), and derender turns those into `ChatCompletionLogProbs` / `CompletionLogProbs`, filling `token` and `bytes` from the tokenizer with the usual U+FFFD byte-fallback correction. Streaming derender does the same per chunk on the plain detokenization path (see [Streaming state and logprobs](#streaming-state-and-logprobs)). The parser path doesn't resolve them yet and drops them.
+`/inference/v1/generate` returns `GenerateLogProbs` with integer `token_id`s (the generate server has no tokenizer), and derender turns those into `ChatCompletionLogProbs` / `CompletionLogProbs`, filling `token` and `bytes` from the tokenizer with the usual U+FFFD byte-fallback correction. Generate returns every top-k candidate the engine produced; derender keeps the ones the coupled endpoint would, with the same helpers: the first `top_logprobs` on chat (all of them for `top_logprobs=-1`) and the first `logprobs + 1` on completions (all of them for `logprobs=-1`). It reads those values from `chat_request` / `completion_request`, so without the original request, or with one that does not ask for logprobs, it returns every candidate. Requests that use `logprob_token_ids` get no output logprobs from generate (their `sampling_params.logprobs` is unset), so derender has none to return for them either. Streaming derender does the same per chunk on the plain detokenization path (see [Streaming state and logprobs](#streaming-state-and-logprobs)). The parser path doesn't resolve them yet and drops them.
 
 ## Streaming cost
 

docs/serving/online_serving/renderer.md

@@ -108,10 +108,13 @@ choice carries a `GenerateLogProbs`; `output_mode: "text"` returns decoded
 - `content` has one entry per generated token, in generation order.
 - `top_logprobs` is a list, not a dict: JSON turns dict keys into strings and
   the ordering would be implicit. It follows the engine's order: the sampled
-  token first, then the remaining candidates in rank order. With non-greedy
-  sampling the sampled token can sit outside the top k (for example ranks
-  `[5, 1, 2]` at `logprobs=2`); it then takes one of the `logprobs` slots and
-  the rank-k candidate is left out, as on the OpenAI endpoints.
+  token first, then the remaining candidates in rank order. It holds every
+  candidate the engine returned and is not cut to `logprobs`: k entries when
+  the sampled token is in the top k, k + 1 when non-greedy sampling picked a
+  token outside it (for example ranks `[5, 1, 2]` at `logprobs=2`). The OpenAI
+  endpoints cut this differently (`/v1/chat/completions` keeps the first k,
+  `/v1/completions` the first k + 1), so [derender](derenderer.md) applies the
+  cut of the endpoint it renders for.
 - `rank` is the token's rank in the vocabulary distribution (1 = most likely)
   on every entry, the sampled one included; a top-k candidate's rank is its
   top-k position. The list is not sorted by it, so sort by `rank` if you need
@@ -132,6 +135,16 @@ choice carries a `GenerateLogProbs`; `output_mode: "text"` returns decoded
     `return_tokens_as_token_ids` on `/v1/chat/completions` and `/v1/completions`
     is unchanged: it is a user-facing OpenAI option and still uses the
     `token_id:N` format.
+    `top_logprobs` is no longer cut to `max(logprobs, 1)` entries, in either
+    `output_mode`: it now holds every candidate the engine returned, which is
+    `logprobs + 1` entries when the sampled token is outside the top
+    `logprobs`. Clients that assume `len(top_logprobs) == logprobs` should
+    take the first `logprobs` entries after sorting by `rank`.
+    Derender applies each endpoint's cut from the original request, so send
+    `chat_request` / `completion_request` to derender whenever logprobs are
+    requested: without it derender returns every candidate, which on chat is
+    `top_logprobs + 1` entries where `/v1/chat/completions` returns
+    `top_logprobs`.
 
 ## Multimodal Render Features
 

examples/generate/multimodal/vision_language_offline.py

@@ -2274,6 +2274,7 @@ MODELS_SUPPORT_VIT_CUDA_GRAPH = [
     "stepvl",
     "glm4_1v",
     "deepseek_ocr",
+    "deepseek_vl_v2",
     "ernie45_vl",
     "minicpmv2_5_vl",
     "minicpmv2_6_vl",

requirements/cuda.txt

@@ -10,7 +10,7 @@ torchaudio==2.11.0
 torchvision==0.29.0 # Required for phi3v processor. See https://github.com/pytorch/vision?tab=readme-ov-file#installation for corresponding version
 torchcodec >= 0.14
 PyNvVideoCodec==2.0.4
-flashinfer-python==0.7.1
+flashinfer-python==0.7.2
 apache-tvm-ffi==0.1.11
 tilelang==0.1.12
 nvidia-cudnn-frontend>=1.19.1
@@ -25,7 +25,7 @@ nvidia-cutlass-dsl[cu13]==4.8.0
 quack-kernels==0.6.5  # Required by MSA
 
 # Tokenspeed_MLA for faster mla with spec decode
-tokenspeed-mla==2.0.2; platform_system == "Linux"
+tokenspeed-mla==2.0.3; platform_system == "Linux"
 
 # Humming kernels for quantization gemm
 humming-kernels[cu13]==0.1.17

requirements/test/cpu.txt

@@ -157,8 +157,6 @@ cryptography==50.0.1
     #   azure-storage-blob
     #   msal
     #   pyjwt
-cupy-cuda12x==13.6.0
-    # via ray
 cycler==0.12.1
     # via matplotlib
 datamodel-code-generator==0.79.0
@@ -239,8 +237,6 @@ fastavro==1.12.2
     # via cohere
 fastparquet==2024.11.0
     # via genai-perf
-fastrlock==0.8.2
-    # via cupy-cuda12x
 fastsafetensors==0.3.3
     # via -r requirements/test/cuda.in
 filelock==4.0.7
@@ -569,7 +565,6 @@ numpy==2.2.6
     #   albumentations
     #   bm25s
     #   contourpy
-    #   cupy-cuda12x
     #   datasets
     #   decord
     #   einx

requirements/test/cuda.in

@@ -22,7 +22,7 @@ vector_quantize_pytorch # required for minicpmo_26 test
 vocos # required for minicpmo_26 test
 peft>=0.20.0 # required for phi-4-mm test
 pqdm
-ray[cgraph,default]>=2.56.0,<2.57 # Ray Compiled Graph, required by pipeline parallelism tests; sharded-RDT weight transfer needs Ray Direct Transport (>=2.56)
+ray[default]>=2.56.0,<2.57 # sharded-RDT weight transfer needs Ray Direct Transport (>=2.56)
 sentence-transformers>=5.2.0 # required for embedding tests
 soundfile # required for audio tests
 jiwer # required for audio tests

requirements/test/cuda.txt

@@ -170,8 +170,6 @@ cuda-pathfinder==1.3.3
     # via cuda-bindings
 cuda-toolkit==13.0.3.0
     # via torch
-cupy-cuda12x==13.6.0
-    # via ray
 cycler==0.12.1
     # via matplotlib
 datamodel-code-generator==0.79.0
@@ -255,8 +253,6 @@ fastavro==1.12.2
     # via cohere
 fastparquet==2024.11.0
     # via genai-perf
-fastrlock==0.8.2
-    # via cupy-cuda12x
 fastsafetensors==0.3.3
     # via
     #   -c requirements/cuda.txt
@@ -604,7 +600,6 @@ numpy==2.2.6
     #   albumentations
     #   bm25s
     #   contourpy
-    #   cupy-cuda12x
     #   datasets
     #   decord
     #   einx

requirements/test/nightly-torch.txt

@@ -16,7 +16,7 @@ librosa # required for audio tests
 vocos # required for minicpmo_26 test
 peft
 pqdm
-ray[cgraph,default]>=2.56.0,<2.57 # Ray Compiled Graph, required by pipeline parallelism tests; sharded-RDT weight transfer needs Ray Direct Transport (>=2.56)
+ray[default]>=2.56.0,<2.57 # sharded-RDT weight transfer needs Ray Direct Transport (>=2.56)
 sentence-transformers>=5.2.0 # required for embedding tests
 soundfile # required for audio tests
 jiwer # required for audio tests

requirements/test/rocm.in

@@ -21,7 +21,7 @@ vector_quantize_pytorch # required for minicpmo_26 test
 vocos # required for minicpmo_26 test
 peft>=0.20.0 # required for phi-4-mm test
 pqdm
-ray[cgraph,default]==2.55.1 # Includes worker startup getenv/setenv race fix
+ray[default]==2.55.1 # Includes worker startup getenv/setenv race fix
 sentence-transformers>=5.2.0 # required for embedding tests
 soundfile # required for audio tests
 jiwer # required for audio tests

requirements/test/rocm.txt

@@ -606,7 +606,6 @@ numpy==2.2.6
     #   bitsandbytes
     #   bm25s
     #   contourpy
-    #   cupy-cuda12x
     #   datasets
     #   decord
     #   einx
@@ -1439,7 +1438,6 @@ zipp==3.23.0
 # cuda-bindings
 # cuda-pathfinder
 # cuda-toolkit
-# cupy-cuda12x
 # nvidia-cublas
 # nvidia-cuda-cupti
 # nvidia-cuda-nvrtc

rust/src/server/src/routes/inference/generate.rs

@@ -5,7 +5,7 @@ mod convert;
 mod types;
 mod validate;
 
-use std::collections::{HashMap, HashSet};
+use std::collections::HashMap;
 use std::convert::Infallible;
 use std::result::Result;
 use std::sync::Arc;
@@ -158,7 +158,7 @@ async fn generate_chunk_stream(
     ResponseOptions {
         include_usage,
         include_continuous_usage,
-        logprobs: output_logprobs,
+        include_logprobs,
         // Ignored: raw generate streaming has no prompt-logprobs wire shape.
         include_prompt_logprobs: _,
         return_token_ids,
@@ -213,15 +213,13 @@ async fn generate_chunk_stream(
                     continue;
                 }
 
-                let logprobs = if let Some(requested) = output_logprobs
-                    && !token_ids.is_empty()
-                {
+                let logprobs = if include_logprobs && !token_ids.is_empty() {
                     let logprobs = output.logprobs.as_ref().ok_or_else(|| {
                         server_error!(
                             "raw generate stream requested logprobs but generation returned none"
                         )
                     })?;
-                    Some(raw_logprobs_to_generate(logprobs, requested)?)
+                    Some(raw_logprobs_to_generate(logprobs)?)
                 } else {
                     None
                 };
@@ -289,19 +287,19 @@ fn collect_generate(
         include_usage: _,
         // Ignored: continuous usage is a streaming-only option.
         include_continuous_usage: _,
-        logprobs: output_logprobs,
+        include_logprobs,
         include_prompt_logprobs,
         return_token_ids,
     }: ResponseOptions,
     mm_placeholders: Option<MultiModalPlaceholders>,
 ) -> Result<GenerateResponse, ApiError> {
-    let logprobs = if let Some(requested) = output_logprobs {
+    let logprobs = if include_logprobs {
         let logprobs = collected.logprobs.as_ref().ok_or_else(|| {
             ApiError::server_error(
                 "raw generate response requested logprobs but generation returned none".to_string(),
             )
         })?;
-        Some(raw_logprobs_to_generate(logprobs, requested)?)
+        Some(raw_logprobs_to_generate(logprobs)?)
     } else {
         None
     };
@@ -381,14 +379,11 @@ fn extract_mm_placeholders(features: Option<&[MmFeatureSpec]>) -> Option<MultiMo
     Some(placeholders)
 }
 
-fn raw_logprobs_to_generate(
-    logprobs: &Logprobs,
-    requested: i32,
-) -> Result<GenerateLogProbs, ApiError> {
+fn raw_logprobs_to_generate(logprobs: &Logprobs) -> Result<GenerateLogProbs, ApiError> {
     let content = logprobs
         .positions
         .iter()
-        .map(|position| position_to_generate_logprobs_content(position, requested))
+        .map(position_to_generate_logprobs_content)
         .collect::<Result<Vec<_>, _>>()?;
 
     Ok(GenerateLogProbs {
@@ -409,12 +404,12 @@ fn raw_prompt_logprobs_to_maps(
         .collect()
 }
 
-/// `top_logprobs` matches the Python generate server: the engine's entries (the
-/// sampled token first, then ranks 1..k) with repeated token ids dropped, cut to
-/// `max(requested, 1)` entries, or all of them for `-1`.
+/// `top_logprobs` matches the Python generate server: every candidate the
+/// engine returned (the sampled token first, then ranks 1..k) with the repeated
+/// sampled entry dropped, and no cut. Derender applies each OpenAI endpoint's
+/// cut.
 fn position_to_generate_logprobs_content(
     position: &PositionLogprobs,
-    requested: i32,
 ) -> Result<GenerateLogProbsContent, ApiError> {
     let chosen = position.entries.first().ok_or_else(|| {
         ApiError::server_error(
@@ -422,9 +417,9 @@ fn position_to_generate_logprobs_content(
         )
     })?;
 
-    // One pass: the first entry of each token id wins, as in Python's dict. The
-    // sampled token is the only one the engine can list twice.
-    let mut seen = HashSet::with_capacity(position.entries.len());
+    // The engine row is the sampled token followed by the top k, which are
+    // distinct; only the sampled token can repeat, at its own rank (anywhere in
+    // 1..=k). Skip that repeat, as Python's dict does: linear, no allocation.
     Ok(GenerateLogProbsContent {
         token_id: chosen.token_id,
         logprob: clamp_logprob(chosen.logprob),
@@ -432,9 +427,9 @@ fn position_to_generate_logprobs_content(
         top_logprobs: position
             .entries
             .iter()
-            .filter(|entry| seen.insert(entry.token_id))
-            .take(usize::try_from(requested).map_or(usize::MAX, |n| n.max(1)))
-            .map(|entry| GenerateLogProb {
+            .enumerate()
+            .filter(|(i, entry)| *i == 0 || entry.token_id != chosen.token_id)
+            .map(|(_, entry)| GenerateLogProb {
                 token_id: entry.token_id,
                 logprob: clamp_logprob(entry.logprob),
                 rank: wire_rank(entry.rank),
@@ -563,8 +558,8 @@ mod tests {
         }
     }
 
-    fn top(position: &PositionLogprobs, requested: i32) -> Vec<(u32, Option<u32>)> {
-        position_to_generate_logprobs_content(position, requested)
+    fn top(position: &PositionLogprobs) -> Vec<(u32, Option<u32>)> {
+        position_to_generate_logprobs_content(position)
             .unwrap()
             .top_logprobs
             .iter()
@@ -576,24 +571,31 @@ mod tests {
     fn generate_top_logprobs_drop_the_repeated_sampled_token() {
         // Sampled token 7 is also top-1: the engine row is [7, 7, 8].
         let pos = position(&[(7, -0.1, 1), (7, -0.1, 1), (8, -2.0, 2)]);
-        assert_eq!(top(&pos, 2), vec![(7, Some(1)), (8, Some(2))]);
+        assert_eq!(top(&pos), vec![(7, Some(1)), (8, Some(2))]);
     }
 
     #[test]
-    fn generate_top_logprobs_keep_sampled_first_when_outside_top_k() {
-        // Sampled token 50 has vocab rank 5; it takes a slot and rank 2 drops
-        // out, as on the Python generate server and the OpenAI endpoints.
+    fn generate_top_logprobs_keep_every_candidate_when_sampled_is_outside_top_k() {
+        // Sampled token 50 has vocab rank 5 at logprobs=2: sampled first, then
+        // ranks 1 and 2, all kept (derender cuts per endpoint).
         let pos = position(&[(50, -3.0, 5), (10, -0.2, 1), (20, -1.0, 2)]);
-        assert_eq!(top(&pos, 2), vec![(50, Some(5)), (10, Some(1))]);
+        assert_eq!(top(&pos), vec![(50, Some(5)), (10, Some(1)), (20, Some(2))]);
     }
 
     #[test]
-    fn generate_top_logprobs_respect_zero_and_all() {
-        let pos = position(&[(7, -0.1, 1), (7, -0.1, 1), (8, -2.0, 2), (9, -3.0, 3)]);
-        assert_eq!(top(&pos, 0), vec![(7, Some(1))]);
+    fn generate_top_logprobs_drop_the_sampled_repeat_at_its_own_rank() {
+        // Sampled token 30 is rank 3: the engine row is [30, 10, 20, 30, 40],
+        // so the repeat is at entries[3], not entries[1].
+        let pos = position(&[
+            (30, -1.0, 3),
+            (10, -0.2, 1),
+            (20, -0.5, 2),
+            (30, -1.0, 3),
+            (40, -2.0, 4),
+        ]);
         assert_eq!(
-            top(&pos, -1),
-            vec![(7, Some(1)), (8, Some(2)), (9, Some(3))]
+            top(&pos),
+            vec![(30, Some(3)), (10, Some(1)), (20, Some(2)), (40, Some(4))]
         );
     }
 
@@ -606,7 +608,7 @@ mod tests {
         entries.extend((0..VOCAB).map(|t| (t, -1.0 - t as f32, t + 1)));
         let pos = position(&entries);
         let start = std::time::Instant::now();
-        let got = top(&pos, -1);
+        let got = top(&pos);
         assert!(start.elapsed() < std::time::Duration::from_secs(5));
         assert_eq!(got.len(), VOCAB as usize);
         assert_eq!(got[0], (5, Some(1)));
@@ -616,10 +618,10 @@ mod tests {
     fn generate_rank_zero_from_a_nan_logprob_is_none() {
         // The engine reports rank 0 for a sampled token whose logprob is NaN.
         let pos = position(&[(7, f32::NAN, 0), (8, -0.2, 1)]);
-        let content = position_to_generate_logprobs_content(&pos, 1).unwrap();
+        let content = position_to_generate_logprobs_content(&pos).unwrap();
         assert_eq!(content.rank, None);
         assert_eq!(content.logprob, -9999.0);
-        assert_eq!(top(&pos, 1), vec![(7, None)]);
+        assert_eq!(top(&pos), vec![(7, None), (8, Some(1))]);
     }
 
     #[tokio::test]

rust/src/server/src/routes/inference/generate/convert.rs

@@ -26,10 +26,8 @@ pub(super) struct ResponseOptions {
     pub include_usage: bool,
     /// Whether the caller asked for usage on every streamed chunk.
     pub include_continuous_usage: bool,
-    /// Requested output logprobs (`sampling_params.logprobs`): `None` omits
-    /// them, `Some(n)` returns `max(n, 1)` candidates per position, `Some(-1)`
-    /// all of them.
-    pub logprobs: Option<i32>,
+    /// Whether the caller requested output logprobs on generate choices.
+    pub include_logprobs: bool,
     /// Whether the caller requested top-level prompt logprobs.
     pub include_prompt_logprobs: bool,
     /// Whether the caller requested final prompt token metadata.
@@ -58,7 +56,7 @@ pub(super) fn prepare_generate_request(
             .as_ref()
             .and_then(|options| options.continuous_usage_stats)
             .unwrap_or(false);
-    let logprobs = request.sampling_params.inner.logprobs;
+    let include_logprobs = request.sampling_params.inner.logprobs.is_some();
     let include_prompt_logprobs = request.sampling_params.inner.prompt_logprobs.is_some();
     let return_token_ids = request.return_token_ids.unwrap_or(false);
     let mut sampling_params = request.sampling_params.inner;
@@ -98,7 +96,7 @@ pub(super) fn prepare_generate_request(
         options: ResponseOptions {
             include_usage,
             include_continuous_usage,
-            logprobs,
+            include_logprobs,
             include_prompt_logprobs,
             return_token_ids,
         },

tests/basic_correctness/models/test_basic_correctness.py

@@ -216,15 +216,6 @@ def test_models_distributed(
         pytest.skip(f"Skip test for {TARGET_TEST_SUITE}")
 
     with monkeypatch.context() as monkeypatch_context:
-        if (
-            model == "meta-llama/Llama-3.2-1B-Instruct"
-            and distributed_executor_backend == "ray"
-            and attention_backend == ""
-            and TARGET_TEST_SUITE == "L4"
-            and enable_prompt_embeds
-        ):  # noqa
-            pytest.skip("enable_prompt_embeds does not work with ray compiled dag.")
-
         for k, v in extra_env.items():
             monkeypatch_context.setenv(k, v)
 

tests/compile/passes/test_rocm_aiter_qk_norm_rope_kvcache_fusion.py

@@ -157,10 +157,17 @@ class QKNormRoPEKVCacheTestModel(torch.nn.Module):
         )
 
         if self.kv_cache_dtype != self.dtype:
-            self.attn._k_scale = torch.tensor(1.0, dtype=torch.float32, device=device)
-            self.attn._v_scale = torch.tensor(1.0, dtype=torch.float32, device=device)
-            self.attn._k_scale_float = 1.0
-            self.attn._v_scale_float = 1.0
+            # Non-unit scales so the fused and unfused writers must agree on them.
+            self.attn._k_scale = torch.tensor(0.5, dtype=torch.float32, device=device)
+            self.attn._v_scale = torch.tensor(0.25, dtype=torch.float32, device=device)
+            self.attn._k_scale_cpu = torch.tensor(
+                0.5, dtype=torch.float32, device="cpu"
+            )
+            self.attn._v_scale_cpu = torch.tensor(
+                0.25, dtype=torch.float32, device="cpu"
+            )
+            self.attn._k_scale_float = 0.5
+            self.attn._v_scale_float = 0.25
         else:
             self.attn._k_scale = self.attn._k_scale.to(device)
             self.attn._v_scale = self.attn._v_scale.to(device)
@@ -527,16 +534,6 @@ def _run_qk_norm_rope_kvcache_fusion_test(
         else:
             torch.testing.assert_close(q_unfused, q_fused, atol=ATOL, rtol=RTOL)
 
-        if not expect_fusion or (not is_fp8_cache and mrope_section is None):
-            # The AITER PTS kernel populates k_out only for non-FP8 caches.
-            # With FP8, the kernel writes quantized K directly to the cache
-            # and may leave k_out uninitialised.  In production this is fine
-            # because downstream attention reads K from the cache.
-            torch.testing.assert_close(k_unfused, k_fused, atol=ATOL, rtol=RTOL)
-
-        # Should be bit exact since no processing had been done on v for both paths
-        torch.testing.assert_close(v_unfused, v_fused, atol=0.0, rtol=0.0)
-
         # Fused and unfused arithmetic can straddle an FP8 rounding boundary.
         # Allow one E4M3 quantization step while still comparing every block.
         if is_fp8_cache:
@@ -544,6 +541,18 @@ def _run_qk_norm_rope_kvcache_fusion_test(
         else:
             cache_atol, cache_rtol = ATOL, RTOL
 
+        # The fused k_out is the dequantized cache write, so it carries the same
+        # rounding; FA and ROCM_ATTN prefill consume it directly. Unified reads K
+        # from the cache and asks the kernel not to write k_out at all.
+        if attn_backend != AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN:
+            torch.testing.assert_close(
+                k_unfused, k_fused, atol=cache_atol, rtol=cache_rtol
+            )
+
+        # Should be bit exact since no processing had been done on v for both paths
+        torch.testing.assert_close(v_unfused, v_fused, atol=0.0, rtol=0.0)
+
+        # Whole cache, so the interleaved V write (ROCM_ATTN kv_cache[1]) is checked.
         torch.testing.assert_close(
             kv_cache_unfused.float(),
             kv_cache_fused.float(),
@@ -570,20 +579,52 @@ _FUSION_CONFIGS = [
     _FUSION_CONFIGS,
 )
 @pytest.mark.parametrize(
-    "attn_backend",
+    "attn_backend, use_shuffle_kv_layout, kv_layout",
     [
-        AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN,
-        AttentionBackendEnum.ROCM_AITER_FA,
+        # ROCM_ATTN always writes the shuffled, interleaved-V layout and needs each
+        # K/V side contiguous per block, so head-major only.
+        pytest.param(
+            AttentionBackendEnum.ROCM_ATTN,
+            "0",
+            KVCacheLayout.LBHNC,
+            id="rocm_attn",
+        ),
+        # Unified never reads the shuffle env; its strided views take any layout.
+        pytest.param(
+            AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN,
+            "0",
+            KVCacheLayout.LBHNC,
+            id="unified-head_major",
+        ),
+        pytest.param(
+            AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN,
+            "0",
+            KVCacheLayout.LBNHC,
+            id="unified-token_major",
+        ),
+        # FA takes any layout unshuffled; its shuffle writer needs each K/V side
+        # contiguous per block, so head-major only.
+        pytest.param(
+            AttentionBackendEnum.ROCM_AITER_FA,
+            "0",
+            KVCacheLayout.LBHNC,
+            id="fa-head_major",
+        ),
+        pytest.param(
+            AttentionBackendEnum.ROCM_AITER_FA,
+            "0",
+            KVCacheLayout.LBNHC,
+            id="fa-token_major",
+        ),
+        pytest.param(
+            AttentionBackendEnum.ROCM_AITER_FA,
+            "1",
+            KVCacheLayout.LBHNC,
+            id="fa-shuffle",
+        ),
     ],
 )
 @pytest.mark.parametrize("num_tokens", [5, 2048])
-@pytest.mark.parametrize(
-    "kv_layout",
-    [
-        pytest.param(KVCacheLayout.LBHNC, id="head_major"),
-        pytest.param(KVCacheLayout.LBNHC, id="token_major"),
-    ],
-)
 @pytest.mark.parametrize("enable_aiter_triton_rope", [True, False])
 @pytest.mark.parametrize("block_size", [16, 32, 64])
 @pytest.mark.parametrize("dtype", [torch.bfloat16])
@@ -615,6 +656,7 @@ def test_qk_norm_rope_kvcache_fusion(
     kv_cache_dtype: CacheDType,
     rms_norm_eps: float,
     custom_op: str,
+    use_shuffle_kv_layout: str,
     monkeypatch: pytest.MonkeyPatch,
 ):
     _run_qk_norm_rope_kvcache_fusion_test(
@@ -627,7 +669,7 @@ def test_qk_norm_rope_kvcache_fusion(
         rotary_dim=rotary_dim,
         block_size=block_size,
         is_neox=is_neox,
-        use_shuffle_kv_layout="0",
+        use_shuffle_kv_layout=use_shuffle_kv_layout,
         kv_layout=kv_layout,
         dtype=dtype,
         kv_cache_dtype=kv_cache_dtype,

tests/distributed/test_cpu_all_reduce.py

@@ -0,0 +1,84 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for cpu_all_reduce_ (gloo, CPU only)."""
+
+import pytest
+import torch
+import torch.distributed as dist
+import torch.multiprocessing as mp
+
+from vllm.distributed.parallel_state import GroupCoordinator
+from vllm.distributed.utils import (
+    cpu_all_reduce_,
+    stateless_init_torch_distributed_process_group,
+)
+from vllm.utils.network_utils import get_open_port
+
+pytestmark = pytest.mark.cpu_test
+
+
+def _make_input(rank: int, world_size: int, step: int) -> torch.Tensor:
+    # Same shape as the V2 DP metadata tensor; values differ per rank/step.
+    t = torch.zeros(6, world_size, dtype=torch.int32)
+    t[:, rank] = torch.arange(6, dtype=torch.int32) * (rank + 1) + step
+    return t
+
+
+def _worker(rank: int, world_size: int, port: int, stateless_port: int):
+    dist.init_process_group(
+        "gloo",
+        init_method=f"tcp://127.0.0.1:{port}",
+        rank=rank,
+        world_size=world_size,
+    )
+    coord = GroupCoordinator(
+        group_ranks=[list(range(world_size))],
+        local_rank=0,
+        torch_distributed_backend="gloo",
+        use_device_communicator=False,
+        group_name="cpu_all_reduce_test",
+    )
+
+    # Several steps to exercise tag reuse across consecutive calls.
+    for step in range(5):
+        expected = _make_input(rank, world_size, step)
+        dist.all_reduce(expected, group=coord.cpu_group)
+        t = _make_input(rank, world_size, step)
+        out = coord.cpu_all_reduce_(t)
+        assert out is t
+        torch.testing.assert_close(t, expected, rtol=0, atol=0)
+
+    # Float sums must be bitwise identical on every rank.
+    t = torch.randn(6, world_size, generator=torch.Generator().manual_seed(rank))
+    expected = t.clone()
+    dist.all_reduce(expected, group=coord.cpu_group)
+    coord.cpu_all_reduce_(t)
+    torch.testing.assert_close(t, expected)
+    gathered = [torch.empty_like(t) for _ in range(world_size)]
+    dist.all_gather(gathered, t, group=coord.cpu_group)
+    for other in gathered:
+        assert torch.equal(other, t)
+
+    # Unregistered (stateless) process group, as used by the engine-core DP
+    # group and elastic EP.
+    stateless_pg = stateless_init_torch_distributed_process_group(
+        "127.0.0.1", stateless_port, rank, world_size, backend="gloo"
+    )
+    expected = _make_input(rank, world_size, 0)
+    dist.all_reduce(expected)
+    t = _make_input(rank, world_size, 0)
+    cpu_all_reduce_(t, stateless_pg)
+    torch.testing.assert_close(t, expected, rtol=0, atol=0)
+
+    dist.destroy_process_group()
+
+
+# Non-power-of-two sizes exercise the cross-block exchange.
+@pytest.mark.parametrize("world_size", [2, 3, 4, 6, 8, 9, 12])
+def test_cpu_all_reduce(world_size: int):
+    mp.spawn(
+        _worker,
+        args=(world_size, get_open_port(), get_open_port()),
+        nprocs=world_size,
+        join=True,
+    )

tests/distributed/test_distributed_oot.py

@@ -1,6 +1,10 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import json
+
+import pytest
+
 from tests.plugins_tests.test_oot_registration_online import (
     run_and_test_dummy_opt_api_server,
 )
@@ -8,3 +12,17 @@ from tests.plugins_tests.test_oot_registration_online import (
 
 def test_distributed_oot(dummy_opt_path: str):
     run_and_test_dummy_opt_api_server(dummy_opt_path, tp=2)
+
+
+@pytest.mark.parametrize("enforce_eager", [True, False])
+def test_distributed_oot_pp_raw_input_tokens(dummy_opt_path: str, enforce_eager: bool):
+    # The model asserts that input_ids reaches every pipeline stage.
+    overrides = {"architectures": ["MyOPTRawInputTokensForCausalLM"]}
+    extra_args = ["-pp", "2", "--hf-overrides", json.dumps(overrides)]
+    if enforce_eager:
+        extra_args.append("--enforce-eager")
+    run_and_test_dummy_opt_api_server(
+        dummy_opt_path,
+        extra_args=extra_args,
+        env_dict={"VLLM_USE_V2_MODEL_RUNNER": "1"},
+    )

tests/distributed/test_elastic_ep.py

@@ -51,7 +51,7 @@ def _traffic_loop(
     stop: threading.Event,
     finished: threading.Event,
     is_probe: bool = False,
-) -> list[tuple[float, float, int | None]]:
+) -> list[tuple[float, float, int | None, str | None]]:
     url = server.url_for("is_scaling_elastic_ep" if is_probe else "v1/completions")
     payload = {"model": MODEL_NAME, "prompt": "Hello", "max_tokens": 4}
     headers = None if dp_rank is None else {"X-data-parallel-rank": str(dp_rank)}
@@ -60,15 +60,18 @@ def _traffic_loop(
     is_ready = False
     while not stop.is_set():
         request_start = time.perf_counter()
+        phase = None
         try:
             response = requests.post(
                 url, json=request_payload, headers=headers, timeout=120
             )
             status_code = response.status_code
+            if is_probe and status_code == 200:
+                phase = response.json()["phase"]
         except requests.exceptions.RequestException:
             status_code = None
-        responses.append((request_start, time.perf_counter(), status_code))
-        if status_code == 200:
+        responses.append((request_start, time.perf_counter(), status_code, phase))
+        if status_code == 200 and (not is_probe or phase in ("idle", "completed")):
             if not is_ready:
                 ready.wait(timeout=120)
                 is_ready = True
@@ -78,14 +81,16 @@ def _traffic_loop(
     return responses
 
 
-def _downtime(responses: list[tuple[float, float, int | None]]) -> float:
-    rejected = [end for _, end, status in responses if status == 503]
-    if not rejected:
+def _downtime(responses: list[tuple[float, float, int | None, str | None]]) -> float:
+    committing = [end for _, end, _, phase in responses if phase == "committing"]
+    if not committing:
         return 0
     recovered = next(
-        end for _, end, status in responses if status == 200 and end > rejected[-1]
+        end
+        for _, end, status, phase in responses
+        if status == 200 and phase in ("idle", "completed") and end > committing[-1]
     )
-    return recovered - rejected[0]
+    return recovered - committing[0]
 
 
 def _scale_with_traffic(
@@ -124,16 +129,20 @@ def _scale_with_traffic(
     bad_statuses = {
         status
         for responses in [probe_result, *results]
-        for _, _, status in responses
+        for _, _, status, _ in responses
         if status not in (200, 503)
     }
     assert not bad_statuses, f"traffic got unexpected statuses {bad_statuses}"
-    probe_503 = [start for start, _, status in probe_result if status == 503]
-    assert probe_503, "Scaling probe did not observe commit"
+    probe_committing = [
+        start for start, _, _, phase in probe_result if phase == "committing"
+    ]
+    assert probe_committing, "Scaling probe did not observe commit"
     assert not results or any(
-        status == 200 and start_time <= request_start and request_end < probe_503[0]
+        status == 200
+        and start_time <= request_start
+        and request_end < probe_committing[0]
         for responses in results
-        for request_start, request_end, status in responses
+        for request_start, request_end, status, _ in responses
     ), "No request completed successfully during preparation"
 
     print(

tests/distributed/test_engram_dp_shard.py

@@ -585,6 +585,7 @@ def _check_dummy_hash_model_forward(
 
     model = model_ops.DeepseekV4Model.__new__(model_ops.DeepseekV4Model)
     torch.nn.Module.__init__(model)
+    model.parallel_config = vllm_config.parallel_config
     model.use_native_mega_moe = False
     model.use_sequence_parallel = False
     model.fuse_mhc_all_reduce = False

tests/distributed/test_multi_node_assignment.py

@@ -1,79 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Make sure ray assigns GPU workers to the correct node.
-
-Run:
-```sh
-cd $VLLM_PATH/tests
-
-pytest distributed/test_multi_node_assignment.py
-```
-"""
-
-import os
-import time
-
-import pytest
-import ray
-from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
-
-from vllm import initialize_ray_cluster
-from vllm.config import ParallelConfig
-from vllm.utils.network_utils import get_ip
-from vllm.v1.executor.ray_utils import available_resources_per_node
-
-VLLM_MULTI_NODE = os.getenv("VLLM_MULTI_NODE", "0") == "1"
-
-
-@pytest.mark.skipif(
-    not VLLM_MULTI_NODE, reason="Need at least 2 nodes to run the test."
-)
-def test_multi_node_assignment() -> None:
-    # NOTE: important to keep this class definition here
-    # to let ray use cloudpickle to serialize it.
-    class Actor:
-        def get_ip(self):
-            return get_ip()
-
-    for _ in range(10):
-        config = ParallelConfig(1, 2)
-        initialize_ray_cluster(config)
-
-        current_ip = get_ip()
-        workers = []
-        for bundle_id, bundle in enumerate(config.placement_group.bundle_specs):
-            if not bundle.get("GPU", 0):
-                continue
-            scheduling_strategy = PlacementGroupSchedulingStrategy(
-                placement_group=config.placement_group,
-                placement_group_capture_child_tasks=True,
-                placement_group_bundle_index=bundle_id,
-            )
-
-            worker = ray.remote(
-                num_cpus=0,
-                num_gpus=1,
-                scheduling_strategy=scheduling_strategy,
-            )(Actor).remote()
-            worker_ip = ray.get(worker.get_ip.remote())
-            assert worker_ip == current_ip
-            workers.append(worker)
-
-        for worker in workers:
-            ray.kill(worker)
-
-        ray.util.remove_placement_group(config.placement_group)
-
-        gpus_needed = len(workers)
-        deadline = time.monotonic() + 60
-        node_id = ray.get_runtime_context().get_node_id()
-        while True:
-            free_gpus = available_resources_per_node().get(node_id, {}).get("GPU", 0)
-            if free_gpus >= gpus_needed:
-                break
-
-            remaining = deadline - time.monotonic()
-            if remaining <= 0:
-                raise TimeoutError("placement group resources were not released")
-
-            time.sleep(min(1, remaining))

tests/entrypoints/cohere/test_api_router.py

@@ -598,6 +598,21 @@ class TestCohereErrorEnvelope:
         assert "message" in body
         assert isinstance(body["message"], str) and body["message"]
 
+    @pytest.mark.parametrize("path", ["/cohere/v2/chat", "/cohere/v2/chat/render"])
+    def test_required_tool_choice_without_tools_returns_400(self, path):
+        """Reject an impossible tool choice before reaching either handler."""
+        app = _build_app_with_vllm_handlers(handler=None)
+        with TestClient(app) as client:
+            response = client.post(
+                path,
+                json={**_minimal_request_body(), "tool_choice": "REQUIRED"},
+            )
+        assert response.status_code == HTTPStatus.BAD_REQUEST
+        body = response.json()
+        assert "error" not in body
+        assert "REQUIRED" in body["message"]
+        assert "tools" in body["message"]
+
     def test_http_error_body_is_cohere_shaped(self):
         # A raised ``HTTPException`` from anywhere in the request cycle
         # is routed through ``http_exception_handler`` (producing the

tests/entrypoints/cohere/test_protocol.py

@@ -95,6 +95,27 @@ class TestCohereChatV2Request:
                 tool_choice="ANY",  # not REQUIRED/NONE
             )
 
+    @pytest.mark.parametrize("tool_fields", [{}, {"tools": None}, {"tools": []}])
+    def test_required_tool_choice_requires_tools(self, tool_fields):
+        with pytest.raises(VLLMValidationError, match="REQUIRED.*tools") as exc_info:
+            CohereChatV2Request(
+                model="m",
+                messages=[{"role": "user", "content": "hi"}],
+                tool_choice="REQUIRED",
+                **tool_fields,
+            )
+        assert exc_info.value.parameter == "tool_choice"
+
+    @pytest.mark.parametrize("tool_choice", [None, "NONE"])
+    def test_nonrequired_tool_choice_allows_empty_tools(self, tool_choice):
+        request = CohereChatV2Request(
+            model="m",
+            messages=[{"role": "user", "content": "hi"}],
+            tool_choice=tool_choice,
+            tools=[],
+        )
+        assert request.tool_choice == tool_choice
+
     def test_stop_sequences_bounded_by_env_limit(self):
         messages = [{"role": "user", "content": "hi"}]
         limit = envs.VLLM_MAX_STOP_STRINGS

tests/entrypoints/openai/completion/test_prompt_validation.py

@@ -64,6 +64,7 @@ def test_load_prompt_embeds(
     model_config = Mock(spec=ModelConfig)
     model_config.enable_prompt_embeds = True
     model_config.get_hidden_size.return_value = hidden_size
+    model_config.hf_text_config = object()
     model_config.dtype = dtype
 
     # construct arbitrary tensors of various dtypes, layouts, and sizes.

tests/entrypoints/scale_out/derender/test_derender_stream.py

@@ -348,6 +348,85 @@ def test_non_streaming_derender_initializes_parser_prefix(
     assert choices[0].finish_reason == "length"
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    ("stop", "include_stop_str_in_output", "dropped"),
+    [
+        ("eos", False, True),
+        ("stop_token_id", False, True),
+        ("stop_string", False, False),
+        ("eos", True, False),
+    ],
+)
+async def test_parsed_derender_leaves_out_stop_token_like_engine(
+    tokenizer, stop, include_stop_str_in_output, dropped
+):
+    """A parser keeps special tokens in the decoded text, so the token that
+    stopped generation must be dropped the way the engine drops it."""
+    eos = tokenizer.eos_token_id
+    answer_ids = tokenizer.encode("answer", add_special_tokens=False)
+    bang = tokenizer.encode("!", add_special_tokens=False)[-1]
+    last_id, stop_reason = {
+        "eos": (eos, None),
+        "stop_token_id": (bang, bang),
+        "stop_string": (bang, "!"),
+    }[stop]
+
+    parser = MagicMock()
+    parser.parse.side_effect = lambda text, *args, **kwargs: (None, text, [])
+    parser.parse_delta.side_effect = lambda text, *args, **kwargs: DeltaMessage(
+        content=text
+    )
+    request = _chat_request(include_stop_str_in_output=include_stop_str_in_output)
+
+    with ThreadPoolExecutor(max_workers=1) as executor:
+        renderer = MagicMock()
+        renderer.get_tokenizer.return_value = tokenizer
+        renderer.get_eos_token_id.return_value = eos
+        renderer._executor = executor
+        model_config = MagicMock()
+        model_config.hf_config.model_type = "llama"
+        model_config.model = MODEL_NAME
+        model_config.try_get_generation_config.return_value = {}
+        derenderer = OnlineDerenderer(
+            model_config=model_config,
+            renderer=renderer,
+            request_logger=None,
+            chat_template=None,
+            chat_template_content_format="string",
+        )
+        derenderer.parser = MagicMock(return_value=parser)
+
+        def batch_content(token_ids: list[int], finish_reason: str) -> str:
+            choice = GenerateTokensChoice(
+                index=0,
+                token_ids=token_ids,
+                finish_reason=finish_reason,
+                stop_reason=stop_reason if finish_reason == "stop" else None,
+            )
+            response = GenerateTokensResponse(
+                request_id="test", prompt_token_ids=[11, 12], choices=[choice]
+            )
+            return derenderer._derender_chat(response, request)[0].message.content
+
+        # Same decode window and seeding; "length" never drops a token.
+        expected = batch_content(
+            answer_ids if dropped else answer_ids + [last_id], "length"
+        )
+        batch = batch_content(answer_ids + [last_id], "stop")
+        stream, _ = await derenderer.derender_chat_stream(
+            model=MODEL_NAME,
+            generate_chunk=_make_stream_chunk(
+                answer_ids + [last_id], finish_reason="stop", stop_reason=stop_reason
+            ),
+            chat_request=request,
+            prompt_token_ids=[11, 12],
+        )
+
+    assert batch == expected
+    assert stream.choices[0].delta.content == expected
+
+
 class TestDetokenizeDelta:
     """_detokenize_delta: chunked decode must equal one shot decode."""
 

tests/entrypoints/scale_out/derender/test_generate_logprobs_conversion.py

@@ -139,7 +139,7 @@ def test_sentencepiece_leading_space_matches_engine():
     assert [t.token for t in entry.top_logprobs] == convert_ids_list_to_tokens(
         tokenizer, [spaced, bare]
     )
-    completion = _convert_chat_logprobs_to_completion_logprobs(resolved)
+    completion = _convert_chat_logprobs_to_completion_logprobs(resolved, logprobs)
     assert completion.top_logprobs == [{" true": -0.5, "true": -1.5}]
 
 

tests/entrypoints/scale_out/derender/test_logprobs_endpoint_parity.py

@@ -0,0 +1,212 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""generate + derender must return the same `top_logprobs` as the coupled
+`/v1/chat/completions` and `/v1/completions` endpoints (#59513).
+
+Both sides start from the same per-position engine dicts, built with the
+engine's own `append_logprobs_for_next_position` and detokenized the way the
+engine does, and go through the real coupled serving objects on one side and
+`/inference/v1/generate` + derender on the other. The cases include a sampled
+token outside the top k, where chat keeps k entries and completions k + 1; the
+greedy end-to-end parity tests never produce that case.
+"""
+
+import pytest
+
+from tests.entrypoints.openai.chat_completion.test_serving_chat import (
+    _build_mock_engine,
+    _build_serving_chat,
+)
+from tests.entrypoints.openai.completion.test_completion_error import (
+    _build_serving_completion,
+)
+from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
+from vllm.entrypoints.openai.completion.protocol import CompletionRequest
+from vllm.entrypoints.scale_out.token_in_token_out.protocol import GenerateLogProbs
+from vllm.entrypoints.scale_out.token_in_token_out.serving import ServingTokens
+from vllm.logprobs import append_logprobs_for_next_position
+from vllm.renderers.online_derenderer import (
+    _chat_top_logprobs_limit,
+    _completion_top_logprobs_limit,
+    _convert_chat_logprobs_to_completion_logprobs,
+    _resolve_logprobs,
+)
+from vllm.tokenizers import get_tokenizer
+from vllm.tokenizers.detokenizer_utils import convert_ids_list_to_tokens
+
+MODEL_NAME = "hmellor/tiny-random-LlamaForCausalLM"
+
+
+@pytest.fixture(scope="module")
+def tokenizer():
+    return get_tokenizer(MODEL_NAME)
+
+
+@pytest.fixture(scope="module")
+def serving_chat():
+    return _build_serving_chat(_build_mock_engine())
+
+
+@pytest.fixture(scope="module")
+def serving_completion():
+    return _build_serving_completion(_build_mock_engine())
+
+
+@pytest.fixture(scope="module")
+def vocab_ids(tokenizer):
+    """Ordinary word pieces (no byte-fallback tokens, which need the U+FFFD
+    context repair that is compared elsewhere)."""
+    ids = tokenizer.encode(
+        "the quick brown fox jumps over a lazy dog while seven cats watch "
+        "from an old wooden fence near the river bank",
+        add_special_tokens=False,
+    )
+    pieces = tokenizer.convert_ids_to_tokens(ids)
+    return list(dict.fromkeys(i for i, p in zip(ids, pieces) if "<0x" not in p))
+
+
+def _positions(tokenizer, vocab_ids, k, sampled_ranks):
+    """Engine output for one completion: per position the sampled id and its
+    dict (sampled first, then ranks 1..k), with decoded tokens as the engine
+    fills them. k = -1 returns every candidate, as for the whole vocabulary."""
+    token_ids, dicts = [], []
+    for pos, sampled_rank in enumerate(sampled_ranks):
+        n_top = max(k, sampled_rank) if k != -1 else len(vocab_ids) - 1
+        ranked = [vocab_ids[(pos + j) % len(vocab_ids)] for j in range(n_top + 1)]
+        sampled = ranked[sampled_rank - 1]
+        top = ranked[: len(ranked) - 1] if k == -1 else ranked[:k]
+        ids = [sampled, *top]
+        logprobs = [-0.25 * sampled_rank] + [-0.25 * (r + 1) for r in range(len(top))]
+        per_position: list = []
+        append_logprobs_for_next_position(
+            per_position,
+            ids,
+            logprobs,
+            convert_ids_list_to_tokens(tokenizer, ids),
+            sampled_rank,
+            len(top),
+        )
+        token_ids.append(sampled)
+        dicts.append(per_position[0])
+    return token_ids, dicts
+
+
+def _generate(token_ids, dicts) -> GenerateLogProbs:
+    out = ServingTokens._create_tokens_logprobs(
+        None, token_ids=token_ids, top_logprobs=dicts
+    )
+    # Over the wire, as a derender client would send it back.
+    return GenerateLogProbs.model_validate_json(out.model_dump_json())
+
+
+# Sampled token rank per position: top-1, inside the top k, outside it.
+RANKS = [1, 2, 1, 7, 3, 9, 1]
+
+
+@pytest.mark.parametrize("k", [0, 1, 2, 5, -1])
+def test_chat_derender_matches_coupled_chat(tokenizer, vocab_ids, serving_chat, k):
+    token_ids, dicts = _positions(tokenizer, vocab_ids, k, RANKS)
+    coupled = serving_chat._create_chat_logprobs(
+        token_ids=token_ids,
+        top_logprobs=dicts,
+        tokenizer=tokenizer,
+        num_output_top_logprobs=k,
+    )
+    request = ChatCompletionRequest(
+        model=MODEL_NAME,
+        messages=[{"role": "user", "content": "hi"}],
+        logprobs=True,
+        top_logprobs=k,
+    )
+    derendered = _resolve_logprobs(
+        _generate(token_ids, dicts),
+        tokenizer,
+        top_limit=_chat_top_logprobs_limit(request),
+    )
+    assert derendered.model_dump() == coupled.model_dump()
+
+
+@pytest.mark.parametrize("k", [0, 1, 2, 5, -1])
+def test_completion_derender_matches_coupled_completion(
+    tokenizer, vocab_ids, serving_completion, k
+):
+    token_ids, dicts = _positions(tokenizer, vocab_ids, k, RANKS)
+    coupled = serving_completion._create_completion_logprobs(
+        token_ids=token_ids,
+        top_logprobs=dicts,
+        num_output_top_logprobs=k,
+        tokenizer=tokenizer,
+    )
+    request = CompletionRequest(model=MODEL_NAME, prompt="hi", logprobs=k)
+    generated = _generate(token_ids, dicts)
+    derendered = _convert_chat_logprobs_to_completion_logprobs(
+        _resolve_logprobs(
+            generated,
+            tokenizer,
+            top_limit=_completion_top_logprobs_limit(request),
+        ),
+        generated,
+    )
+    assert derendered.model_dump() == coupled.model_dump()
+    if k > 0:
+        # The case #59513 is about: completions keeps k + 1 entries when the
+        # sampled token is outside the top k.
+        outside = [i for i, r in enumerate(RANKS) if r > k]
+        assert all(len(derendered.top_logprobs[i]) == k + 1 for i in outside)
+    if k == -1:
+        # /v1/completions returns every candidate for logprobs=-1 (#55717).
+        assert [len(top) for top in derendered.top_logprobs] == [len(d) for d in dicts]
+
+
+@pytest.mark.parametrize(
+    "top_limit",
+    [
+        _chat_top_logprobs_limit(None),
+        # A request that does not ask for logprobs says nothing about the cut.
+        _chat_top_logprobs_limit(
+            ChatCompletionRequest(
+                model=MODEL_NAME, messages=[{"role": "user", "content": "hi"}]
+            )
+        ),
+        _completion_top_logprobs_limit(None),
+        _completion_top_logprobs_limit(
+            CompletionRequest(model=MODEL_NAME, prompt="hi")
+        ),
+    ],
+)
+def test_without_a_logprobs_request_derender_keeps_every_candidate(
+    tokenizer, vocab_ids, top_limit
+):
+    token_ids, dicts = _positions(tokenizer, vocab_ids, 2, RANKS)
+    derendered = _resolve_logprobs(
+        _generate(token_ids, dicts), tokenizer, top_limit=top_limit
+    )
+    assert [len(e.top_logprobs) for e in derendered.content] == [len(d) for d in dicts]
+
+
+@pytest.mark.parametrize("k", [-1, 0, 2])
+def test_completion_position_without_logprobs_stays_none(
+    tokenizer, vocab_ids, serving_completion, k
+):
+    """A position the engine returned no logprobs for is `None` in
+    `/v1/completions`, not `{}`; derender keeps it that way."""
+    token_ids, dicts = _positions(tokenizer, vocab_ids, k, RANKS)
+    dicts[3] = None
+    coupled = serving_completion._create_completion_logprobs(
+        token_ids=token_ids,
+        top_logprobs=dicts,
+        num_output_top_logprobs=k,
+        tokenizer=tokenizer,
+    )
+    request = CompletionRequest(model=MODEL_NAME, prompt="hi", logprobs=k)
+    generated = _generate(token_ids, dicts)
+    derendered = _convert_chat_logprobs_to_completion_logprobs(
+        _resolve_logprobs(
+            generated,
+            tokenizer,
+            top_limit=_completion_top_logprobs_limit(request),
+        ),
+        generated,
+    )
+    assert derendered.top_logprobs == coupled.top_logprobs
+    assert derendered.top_logprobs[3] is None

tests/entrypoints/scale_out/token_in_token_out/test_tokens_logprobs.py

@@ -14,11 +14,11 @@ from vllm.logprobs import Logprob
 
 def test_top_logprobs_alternatives_have_own_token_ids():
     """Each top_logprobs alternative must carry its own integer token id."""
+    # Engine dict at logprobs=2 with the sampled token 262 also top-1.
     result = ServingTokens._create_tokens_logprobs(
         None,
         token_ids=[262],
-        top_logprobs=[{262: Logprob(-0.1), 257: Logprob(-1.2), 428: Logprob(-2.3)}],
-        num_output_top_logprobs=2,
+        top_logprobs=[{262: Logprob(-0.1), 257: Logprob(-1.2)}],
     )
     token_ids = {e.token_id for e in result.content[0].top_logprobs}
     assert token_ids == {262, 257}, f"got {token_ids}"
@@ -32,7 +32,6 @@ def test_sampled_entry_carries_token_id_and_rank():
         top_logprobs=[
             {262: Logprob(-0.1, rank=1), 257: Logprob(-1.2, rank=2)},
         ],
-        num_output_top_logprobs=2,
     )
     entry = result.content[0]
     assert entry.token_id == 262
@@ -50,7 +49,6 @@ def test_sampled_token_absent_from_topk_uses_sentinel():
         None,
         token_ids=[5],
         top_logprobs=[{7: Logprob(-0.9)}],
-        num_output_top_logprobs=1,
     )
     entry = result.content[0]
     assert entry.token_id == 5
@@ -60,14 +58,13 @@ def test_sampled_token_absent_from_topk_uses_sentinel():
 
 
 def test_logprobs_zero_emits_sampled_token():
-    """logprobs=0 must still emit 1 entry (the sampled token)."""
+    """logprobs=0: the engine returns only the sampled token, and it is kept."""
     result = ServingTokens._create_tokens_logprobs(
         None,
         token_ids=[7],
-        top_logprobs=[{7: Logprob(-0.9), 8: Logprob(-1.1)}],
-        num_output_top_logprobs=0,
+        top_logprobs=[{7: Logprob(-0.9)}],
     )
-    assert len(result.content[0].top_logprobs) == 1
+    assert [t.token_id for t in result.content[0].top_logprobs] == [7]
 
 
 def test_logprobs_minus_one_emits_all_tokens():
@@ -75,7 +72,6 @@ def test_logprobs_minus_one_emits_all_tokens():
         None,
         token_ids=[7],
         top_logprobs=[{7: Logprob(-0.9), 8: Logprob(-1.1)}],
-        num_output_top_logprobs=-1,
     )
     assert len(result.content[0].top_logprobs) == 2
 
@@ -86,7 +82,6 @@ def test_text_logprobs_without_tokenizer_are_placeholders_without_bytes():
         None,
         token_ids=[7],
         top_logprobs=[{7: Logprob(-0.9), 8: Logprob(-1.1)}],
-        num_output_top_logprobs=2,
     )
     entry = result.content[0]
     assert entry.token == "token_id:7"
@@ -102,7 +97,6 @@ def test_resolved_logprobs_use_engine_decoded_tokens_and_bytes():
         top_logprobs=[
             {7: Logprob(-0.9, decoded_token="é"), 8: Logprob(-1.1, decoded_token="e")}
         ],
-        num_output_top_logprobs=2,
         tokenizer=tokenizer,
     )
     entry = result.content[0]
@@ -121,7 +115,6 @@ def test_resolved_logprobs_decode_tokens_the_engine_did_not():
         None,
         token_ids=[7, 9],
         top_logprobs=[{7: Logprob(-0.9)}, None],
-        num_output_top_logprobs=1,
         tokenizer=tokenizer,
     )
     assert [e.token for e in result.content] == ["<7>", "<9>"]
@@ -130,9 +123,9 @@ def test_resolved_logprobs_decode_tokens_the_engine_did_not():
 
 
 def test_sampled_token_outside_topk_comes_first():
-    """The engine puts the sampled token first, then ranks 1..k. When the
-    sampled token is outside the top k it takes one of the k slots, so rank k
-    is left out (same as the OpenAI endpoints)."""
+    """The engine puts the sampled token first, then ranks 1..k. Generate keeps
+    all k + 1 candidates when the sampled token is outside the top k; derender
+    applies each endpoint's cut (#59513)."""
     result = ServingTokens._create_tokens_logprobs(
         None,
         token_ids=[50],
@@ -143,9 +136,8 @@ def test_sampled_token_outside_topk_comes_first():
                 20: Logprob(-1.0, rank=2),
             }
         ],
-        num_output_top_logprobs=2,
     )
-    assert [t.rank for t in result.content[0].top_logprobs] == [5, 1]
+    assert [t.rank for t in result.content[0].top_logprobs] == [5, 1, 2]
 
 
 def test_logprob_is_required_on_the_wire():
@@ -160,7 +152,6 @@ def test_rank_zero_from_a_nan_logprob_is_none():
         None,
         token_ids=[7],
         top_logprobs=[{7: Logprob(float("nan"), rank=0), 8: Logprob(-0.2, rank=1)}],
-        num_output_top_logprobs=1,
     )
     entry = result.content[0]
     assert entry.rank is None
@@ -176,7 +167,6 @@ def test_text_logprobs_clamp_a_nan_logprob():
         None,
         token_ids=[7],
         top_logprobs=[{7: Logprob(float("nan"), rank=0, decoded_token="a")}],
-        num_output_top_logprobs=1,
         tokenizer=MagicMock(),
     )
     entry = result.content[0]

tests/kernels/attention/test_attention.py

@@ -122,6 +122,7 @@ def ref_single_query_cached_kv_attention(
 @pytest.mark.parametrize("kv_cache_dtype", KV_CACHE_DTYPE)
 @pytest.mark.parametrize("seed", SEEDS)
 @pytest.mark.parametrize("device", CUDA_DEVICES)
+@pytest.mark.parametrize("use_interleaved_v_cache", [False, True])
 def test_paged_attention(
     kv_cache_factory,
     num_seqs: int,
@@ -133,9 +134,14 @@ def test_paged_attention(
     kv_cache_dtype: str,
     seed: int,
     device: str,
+    use_interleaved_v_cache: bool,
 ) -> None:
     if current_platform.is_navi() and (
-        kv_cache_dtype == "fp8" or head_size != 128 or block_size != 16 or use_alibi
+        kv_cache_dtype == "fp8"
+        or head_size != 128
+        or block_size != 16
+        or use_alibi
+        or use_interleaved_v_cache
     ):
         pytest.skip()
 
@@ -201,6 +207,20 @@ def test_paged_attention(
             last_block_idx = block_tables_lst[seq_idx][seq_len // block_size]
             value_cache[last_block_idx, :, :, padding_start:] = padding_nan
 
+    standard_value_cache = value_cache
+    if use_interleaved_v_cache:
+        # Hand the kernel V in the interleaved layout (same 4D shape); the
+        # reference below gets the standard layout back.
+        x_v = 16 // value_cache.element_size()
+        value_cache = (
+            value_cache.view(
+                NUM_BLOCKS, num_kv_heads, head_size, block_size // x_v, x_v
+            )
+            .permute(0, 1, 3, 2, 4)
+            .contiguous()
+            .view(NUM_BLOCKS, num_kv_heads, head_size, block_size)
+        )
+
     # Call the paged attention kernel.
     output = torch.empty_like(query)
     num_partitions = (max_seq_len + PARTITION_SIZE_ROCM - 1) // PARTITION_SIZE_ROCM
@@ -234,6 +254,7 @@ def test_paged_attention(
         kv_cache_dtype,
         k_scale,
         v_scale,
+        use_interleaved_v_cache=use_interleaved_v_cache,
     )
 
     opcheck(
@@ -259,10 +280,13 @@ def test_paged_attention(
             v_scale,
             None,
             "f16",
+            use_interleaved_v_cache,
         ),
         cond=(head_size == 64 and block_size == BLOCK_SIZES[0]),
     )
 
+    value_cache = standard_value_cache
+
     # Run the reference implementation.
     if kv_cache_dtype == "fp8":
         # Convert cache data back to dtype.

tests/kernels/attention/test_prefix_prefill.py

@@ -109,6 +109,7 @@ def create_alibi_causal_mask(
 @pytest.mark.parametrize("device", CUDA_DEVICES)
 @pytest.mark.parametrize("sliding_window", SLIDING_WINDOW)
 @pytest.mark.parametrize("op", OPS)
+@pytest.mark.parametrize("interleave_v", [False, True])
 @torch.inference_mode()
 def test_contexted_kv_attention(
     num_heads: int,
@@ -119,6 +120,7 @@ def test_contexted_kv_attention(
     kv_cache_dtype: str,
     device: str,
     op: Callable,
+    interleave_v: bool,
     block_size: int = 32,
 ) -> None:
     if "fp8" in kv_cache_dtype and not current_platform.has_device_capability(89):
@@ -225,6 +227,23 @@ def test_contexted_kv_attention(
     )
     k_scale = v_scale = torch.tensor(1.0, dtype=torch.float32, device=device)
 
+    extra_op_kwargs = {}
+    if interleave_v:
+        # Rewrite the V cache into the interleaved layout (same 4D shape). Both
+        # context_attention_fwd (prefill) and kernel_paged_attention_2d (decode)
+        # read it back interleaved.
+        x_v = 16 // v_cache.element_size()
+        v_cache = (
+            v_cache.view(cache_size, num_kv_heads, head_size, block_size // x_v, x_v)
+            .permute(0, 1, 3, 2, 4)
+            .contiguous()
+            .view(cache_size, num_kv_heads, head_size, block_size)
+        )
+        if op is context_attention_fwd:
+            extra_op_kwargs["interleaved_v_pack_factor"] = x_v
+        else:
+            extra_op_kwargs["use_interleaved_v_cache"] = True
+
     # Warm up the Triton kernel by calling it once before actually measuring
     # generation time
     op(
@@ -243,6 +262,7 @@ def test_contexted_kv_attention(
         k_scale,
         v_scale,
         sliding_window=sliding_window,
+        **extra_op_kwargs,
     )
     torch.accelerator.synchronize()
     start_time = time.time()
@@ -262,6 +282,7 @@ def test_contexted_kv_attention(
         k_scale,
         v_scale,
         sliding_window=sliding_window,
+        **extra_op_kwargs,
     )
     torch.accelerator.synchronize()
     end_time = time.time()
@@ -386,6 +407,7 @@ def test_rocm_backend_sliding_window_includes_boundary_token(
 @pytest.mark.parametrize("dtype", DTYPES)
 @pytest.mark.parametrize("kv_cache_dtype", KV_CACHE_DTYPES)
 @pytest.mark.parametrize("device", CUDA_DEVICES)
+@pytest.mark.parametrize("interleave_v", [False, True])
 @torch.inference_mode()
 def test_contexted_kv_attention_cached_kv(
     num_heads: int,
@@ -394,6 +416,7 @@ def test_contexted_kv_attention_cached_kv(
     dtype: torch.dtype,
     kv_cache_dtype: str,
     device: str,
+    interleave_v: bool,
     block_size: int = 32,
 ) -> None:
     # Exercises the KV_FROM_CACHE path of context_attention_fwd: the current
@@ -487,6 +510,18 @@ def test_contexted_kv_attention_cached_kv(
     )
     k_scale = v_scale = torch.tensor(1.0, dtype=torch.float32, device=device)
 
+    interleaved_v_pack_factor = 0
+    if interleave_v:
+        # Rewrite the V cache into the interleaved layout (same 4D shape).
+        interleaved_v_pack_factor = 16 // v_cache.element_size()
+        x_v = interleaved_v_pack_factor
+        v_cache = (
+            v_cache.view(cache_size, num_kv_heads, head_size, block_size // x_v, x_v)
+            .permute(0, 1, 3, 2, 4)
+            .contiguous()
+            .view(cache_size, num_kv_heads, head_size, block_size)
+        )
+
     # Cached-K/V path: current-chunk k and v are None.
     context_attention_fwd(
         query,
@@ -504,6 +539,7 @@ def test_contexted_kv_attention_cached_kv(
         k_scale,
         v_scale,
         sliding_window=0,
+        interleaved_v_pack_factor=interleaved_v_pack_factor,
     )
     torch.accelerator.synchronize()
 
@@ -544,8 +580,11 @@ def test_contexted_kv_attention_cached_kv(
 
 
 @pytest.mark.parametrize("device", CUDA_DEVICES)
+@pytest.mark.parametrize("interleave_v", [False, True])
 @torch.inference_mode()
-def test_contexted_kv_attention_cached_kv_block_table_boundary(device: str) -> None:
+def test_contexted_kv_attention_cached_kv_block_table_boundary(
+    device: str, interleave_v: bool
+) -> None:
     # Boundary guard for the KV_FROM_CACHE block-table load. With an
     # exact-sized block table (row length == number of blocks for the
     # sequence) and a query that ends the sequence, the last K/V tile has
@@ -611,6 +650,18 @@ def test_contexted_kv_attention_cached_kv_block_table_boundary(device: str) -> N
     )
     k_scale = v_scale = torch.tensor(1.0, dtype=torch.float32, device=device)
 
+    interleaved_v_pack_factor = 0
+    if interleave_v:
+        # Rewrite the V cache into the interleaved layout (same 4D shape).
+        interleaved_v_pack_factor = 16 // v_cache.element_size()
+        x_v = interleaved_v_pack_factor
+        v_cache = (
+            v_cache.view(num_blocks, num_kv_heads, head_size, block_size // x_v, x_v)
+            .permute(0, 1, 3, 2, 4)
+            .contiguous()
+            .view(num_blocks, num_kv_heads, head_size, block_size)
+        )
+
     # Cached-K/V path: current-chunk k and v are None.
     context_attention_fwd(
         query,
@@ -628,6 +679,7 @@ def test_contexted_kv_attention_cached_kv_block_table_boundary(device: str) -> N
         k_scale,
         v_scale,
         sliding_window=0,
+        interleaved_v_pack_factor=interleaved_v_pack_factor,
     )
     torch.accelerator.synchronize()
 
@@ -723,6 +775,7 @@ def test_contexted_kv_attention_cached_kv_alibi_unsupported(device: str) -> None
 @pytest.mark.parametrize("kv_cache_dtype", KV_CACHE_DTYPES)
 @pytest.mark.parametrize("device", CUDA_DEVICES)
 @pytest.mark.parametrize("op", OPS)
+@pytest.mark.parametrize("interleave_v", [False, True])
 @torch.inference_mode()
 def test_contexted_kv_attention_alibi(
     num_heads: int,
@@ -732,6 +785,7 @@ def test_contexted_kv_attention_alibi(
     kv_cache_dtype: str,
     device: str,
     op: Callable,
+    interleave_v: bool,
     block_size: int = 32,
 ) -> None:
     if "fp8" in kv_cache_dtype and not current_platform.has_device_capability(89):
@@ -860,6 +914,23 @@ def test_contexted_kv_attention_alibi(
     )
     k_scale = v_scale = torch.tensor(1.0, dtype=torch.float32, device=device)
 
+    extra_op_kwargs = {}
+    if interleave_v:
+        # Rewrite the V cache into the interleaved layout (same 4D shape). Both
+        # context_attention_fwd (prefill) and kernel_paged_attention_2d (decode)
+        # read it back interleaved.
+        x_v = 16 // v_cache.element_size()
+        v_cache = (
+            v_cache.view(cache_size, num_kv_heads, head_size, block_size // x_v, x_v)
+            .permute(0, 1, 3, 2, 4)
+            .contiguous()
+            .view(cache_size, num_kv_heads, head_size, block_size)
+        )
+        if op is context_attention_fwd:
+            extra_op_kwargs["interleaved_v_pack_factor"] = x_v
+        else:
+            extra_op_kwargs["use_interleaved_v_cache"] = True
+
     # Warm up the Triton kernel by calling it once before actually measuring
     # generation time
     op(
@@ -878,6 +949,7 @@ def test_contexted_kv_attention_alibi(
         k_scale,
         v_scale,
         alibi_slopes=alibi_slopes,
+        **extra_op_kwargs,
     )
     torch.accelerator.synchronize()
     start_time = time.time()
@@ -897,6 +969,7 @@ def test_contexted_kv_attention_alibi(
         k_scale,
         v_scale,
         alibi_slopes=alibi_slopes,
+        **extra_op_kwargs,
     )
     torch.accelerator.synchronize()
     end_time = time.time()
@@ -992,6 +1065,7 @@ def test_contexted_kv_attention_alibi(
 @pytest.mark.parametrize("device", CUDA_DEVICES)
 @pytest.mark.parametrize("sliding_window", SLIDING_WINDOW)
 @pytest.mark.parametrize("op", OPS)
+@pytest.mark.parametrize("interleave_v", [False, True])
 @torch.inference_mode()
 def test_contexted_kv_attention_f32(
     num_heads: int,
@@ -1002,6 +1076,7 @@ def test_contexted_kv_attention_f32(
     kv_cache_dtype: str,
     device: str,
     op: Callable,
+    interleave_v: bool,
 ) -> None:
     test_contexted_kv_attention(
         num_heads,
@@ -1012,6 +1087,7 @@ def test_contexted_kv_attention_f32(
         kv_cache_dtype,
         device,
         op,
+        interleave_v=interleave_v,
     )
 
 
@@ -1023,6 +1099,7 @@ def test_contexted_kv_attention_f32(
 @pytest.mark.parametrize("kv_cache_dtype", KV_CACHE_DTYPES)
 @pytest.mark.parametrize("device", CUDA_DEVICES)
 @pytest.mark.parametrize("op", OPS)
+@pytest.mark.parametrize("interleave_v", [False, True])
 @torch.inference_mode()
 def test_contexted_kv_attention_alibi_f32(
     num_heads: int,
@@ -1032,9 +1109,17 @@ def test_contexted_kv_attention_alibi_f32(
     kv_cache_dtype: str,
     device: str,
     op: Callable,
+    interleave_v: bool,
 ) -> None:
     test_contexted_kv_attention_alibi(
-        num_heads, num_queries_per_kv, head_size, dtype, kv_cache_dtype, device, op
+        num_heads,
+        num_queries_per_kv,
+        head_size,
+        dtype,
+        kv_cache_dtype,
+        device,
+        op,
+        interleave_v=interleave_v,
     )
 
 
@@ -1052,6 +1137,7 @@ NONSTANDARD_BLOCK_SIZE_SHAPES = [
 @pytest.mark.parametrize("dtype", DTYPES)
 @pytest.mark.parametrize("device", CUDA_DEVICES)
 @pytest.mark.parametrize("op", OPS)
+@pytest.mark.parametrize("interleave_v", [False, True])
 @torch.inference_mode()
 def test_qwen3_nonstandard_block_size(
     num_heads: int,
@@ -1061,6 +1147,7 @@ def test_qwen3_nonstandard_block_size(
     dtype: torch.dtype,
     device: str,
     op: Callable,
+    interleave_v: bool,
 ) -> None:
     """Non-power-of-2 pages must match, even when a tile straddles a page."""
     if not current_platform.is_rocm():
@@ -1076,4 +1163,5 @@ def test_qwen3_nonstandard_block_size(
         kv_cache_dtype="auto",
         device=device,
         op=op,
+        interleave_v=interleave_v,
     )

tests/kernels/attention/test_rocm_aiter_fa.py

@@ -1061,3 +1061,25 @@ def test_aiter_mha_varlen_fp8_kv(dtype):
         rtol=rtol,
     )
     torch.testing.assert_close(output, ref, atol=atol, rtol=rtol)
+
+
+def test_asm_kv_scale_tables_shared_across_builders():
+    """Builders share one AsmKvScaleTables; equal (shape, scale) requests must
+    return the same dense table, a different scale a different one."""
+    from vllm.v1.attention.backends.rocm_aiter_fa import (
+        AiterFlashAttentionMetadataBuilder,
+        AsmKvScaleTables,
+    )
+
+    tables = AsmKvScaleTables()
+    device = torch.device("cuda")
+    table = tables.get(device, 8, 2, 16, 0.5)
+    assert tables.get(device, 8, 2, 16, 0.5) is table
+    assert tables.get(device, 8, 2, 16, 1.0) is not table
+    assert table.shape == (8, 2, 16)
+    assert table.dtype == torch.float32
+    assert torch.all(table == 0.5)
+    # The builder class holds the single process-wide instance.
+    assert isinstance(
+        AiterFlashAttentionMetadataBuilder.asm_kv_scale_tables, AsmKvScaleTables
+    )

tests/kernels/attention/test_rocm_triton_attn_dsv4.py

@@ -1989,8 +1989,8 @@ def test_dsv41_adaptive_mla_swa_metadata_graph_replay(monkeypatch) -> None:
     """Replay V4.1 ROCm decode after device-only query reallocation."""
     from tests.v1.attention.utils import create_vllm_config
     from vllm.models.deepseek_v41.amd.rocm import (
-        DeepseekV4ROCMAiterSparseSWAMetadataBuilder,
         DeepseekV41ROCMAiterMLASparseMetadataBuilder,
+        DeepseekV41ROCMAiterSparseSWAMetadataBuilder,
         compute_global_topk_ragged_indices_and_indptr,
     )
     from vllm.v1.attention.backend import AttentionCGSupport, CommonAttentionMetadata
@@ -2044,7 +2044,7 @@ def test_dsv41_adaptive_mla_swa_metadata_graph_replay(monkeypatch) -> None:
     mla_builder = DeepseekV41ROCMAiterMLASparseMetadataBuilder(
         mla_spec, ["c2a"], vllm_config, device
     )
-    swa_builder = DeepseekV4ROCMAiterSparseSWAMetadataBuilder(
+    swa_builder = DeepseekV41ROCMAiterSparseSWAMetadataBuilder(
         swa_spec, ["c2a"], vllm_config, device
     )
     assert (

tests/kernels/moe/test_flashinfer.py

@@ -52,6 +52,7 @@ from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
 )
 from vllm.model_executor.layers.quantization.utils.fp8_utils import input_to_float8
 from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
+    dequant_mxfp8_to_bf16,
     mxfp8_e4m3_quantize,
 )
 from vllm.model_executor.layers.quantization.utils.quant_utils import (
@@ -65,7 +66,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
 )
 from vllm.model_executor.models.llama4 import Llama4MoE
 from vllm.platforms import current_platform
-from vllm.utils.math_utils import next_power_of_2
+from vllm.utils.math_utils import next_power_of_2, round_up
 from vllm.utils.torch_utils import set_random_seed
 
 try:
@@ -788,7 +789,7 @@ def test_trtllm_fp8_swiglu_clamp_support(
 
 
 def _make_mxfp8_moe_weights(
-    e: int, n: int, k: int
+    e: int, n: int, k: int, is_gated: bool = True
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
     def quantize(w: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
         q, s = zip(
@@ -796,8 +797,9 @@ def _make_mxfp8_moe_weights(
         )
         return torch.stack(q), torch.stack(s)
 
+    w1_rows = 2 * n if is_gated else n
     w1, w1_scale = quantize(
-        torch.randn(e, 2 * n, k, device="cuda", dtype=torch.bfloat16) / 10
+        torch.randn(e, w1_rows, k, device="cuda", dtype=torch.bfloat16) / 10
     )
     w2, w2_scale = quantize(
         torch.randn(e, k, n, device="cuda", dtype=torch.bfloat16) / 10
@@ -806,16 +808,20 @@ def _make_mxfp8_moe_weights(
 
 
 @pytest.mark.parametrize("block_shape", [[128, 128], [1, 32]], ids=["block", "mxfp8"])
+@pytest.mark.parametrize("transposed", [True, False], ids=["transposed", "row-major"])
 @pytest.mark.parametrize("m", [1, 16])
 @pytest.mark.skipif(
     not current_platform.is_device_capability_family(100),
     reason="Requires TRTLLM-Gen FP8 MoE (SM100)",
 )
 def test_trtllm_fp8_block_moe_deferred_finalize(
-    m: int, block_shape: list[int], workspace_init
+    m: int, block_shape: list[int], transposed: bool, workspace_init
 ):
     """TRTLLM-Gen block-FP8 and MXFP8 modular experts can leave the top-k
-    finalize to the caller."""
+    finalize to the caller. For fp8-block, ``transposed`` toggles the oracle's
+    is_scale_transposed column-major activation-scale layout."""
+    if block_shape != [128, 128] and transposed:
+        pytest.skip("column-major scales apply to fp8 [128, 128] block only")
     e, topk, n, k = 32, 4, 1024, 1024
     set_random_seed(7)
     with set_current_vllm_config(vllm_config):
@@ -849,6 +855,9 @@ def test_trtllm_fp8_block_moe_deferred_finalize(
             a2_scale=None,
             block_shape=block_shape,
         )
+        is_block = block_shape == [128, 128]
+        if is_block and not transposed:
+            quant_config = replace(quant_config, is_scale_transposed=False)
         # One rank of a TP group, which deferral needs.
         moe_config = FusedMoEConfig(
             num_experts=e,
@@ -866,19 +875,27 @@ def test_trtllm_fp8_block_moe_deferred_finalize(
             routing_method=RoutingMethodType.TopK,
             max_num_tokens=next_power_of_2(m),
         )
-        kernel = mk.FusedMoEKernel(
-            maybe_make_prepare_finalize(
-                moe=moe_config, quant_config=quant_config, allow_new_interface=True
-            ),
-            TrtLlmFp8ExpertsModular(moe_config=moe_config, quant_config=quant_config),
-        )
+
+        def make_kernel(quant_config: FusedMoEQuantConfig) -> mk.FusedMoEKernel:
+            return mk.FusedMoEKernel(
+                maybe_make_prepare_finalize(
+                    moe=moe_config,
+                    quant_config=quant_config,
+                    allow_new_interface=True,
+                ),
+                TrtLlmFp8ExpertsModular(
+                    moe_config=moe_config, quant_config=quant_config
+                ),
+            )
+
+        kernel = make_kernel(quant_config)
 
         a = torch.randn((m, k), device="cuda", dtype=torch.bfloat16) / 10
         score = torch.randn((m, e), device="cuda", dtype=torch.bfloat16)
         topk_weights, topk_ids, _ = fused_topk(a, score, topk, renormalize=False)
-        check_deferred_moe_finalize(
-            moe_config,
-            lambda: kernel.apply(
+
+        def apply(k_: mk.FusedMoEKernel) -> torch.Tensor:
+            return k_.apply(
                 hidden_states=a,
                 w1=w1,
                 w2=w2,
@@ -888,7 +905,34 @@ def test_trtllm_fp8_block_moe_deferred_finalize(
                 global_num_experts=e,
                 expert_map=None,
                 apply_router_weight_on_input=False,
-            ),
+            )
+
+        if is_block and transposed:
+            # The oracle gates column-major scales on TRTLLM fp8-block.
+            assert quant_config.is_scale_transposed
+            assert not make_fp8_moe_quant_config(
+                Fp8MoeBackend.DEEPGEMM,
+                w1_scale,
+                w2_scale,
+                a1_scale=None,
+                a2_scale=None,
+                block_shape=block_shape,
+            ).is_scale_transposed
+            # no_dp_ep emits column-major [M, G] scales (physical [G, M]), so
+            # the experts' a1q_scale.t().contiguous() is a no-op.
+            _, a1q_scale, _, _, _ = kernel.prepare_finalize.prepare(
+                a, topk_weights, topk_ids, e, None, False, quant_config
+            )
+            assert a1q_scale.t().is_contiguous()
+            # Bitwise-equal to running with row-major scales.
+            row_major = replace(quant_config, is_scale_transposed=False)
+            torch.testing.assert_close(
+                apply(kernel), apply(make_kernel(row_major)), atol=0, rtol=0
+            )
+
+        check_deferred_moe_finalize(
+            moe_config,
+            lambda: apply(kernel),
             router_weights=topk_weights,
         )
 
@@ -986,3 +1030,123 @@ def test_trtllm_mxfp8_minimax2_routing_applies_routed_scale(m: int, workspace_in
         )
 
     torch.testing.assert_close(fused, reference, atol=2e-2, rtol=2e-2)
+
+
+@pytest.mark.parametrize(
+    ("activation", "n", "k"),
+    [
+        (MoEActivation.SILU, 1024, 1024),
+        (MoEActivation.SILU, 928, 1024),
+        (MoEActivation.SILU, 1024, 2688),
+        (MoEActivation.RELU2_NO_MUL, 1024, 1024),
+        # Nemotron-3.5 Lightning's non-gated MoE: intermediate 1856, hidden 2688.
+        (MoEActivation.RELU2_NO_MUL, 1856, 2688),
+    ],
+)
+@pytest.mark.parametrize("m", [1, 32])
+@pytest.mark.skipif(
+    not current_platform.is_device_capability_family(100),
+    reason="Requires TRTLLM-Gen FP8 MoE (SM100)",
+)
+def test_trtllm_mxfp8_moe_pads_intermediate(
+    m: int, activation: MoEActivation, n: int, k: int, workspace_init
+):
+    """TRTLLM-Gen MXFP8 experts need the intermediate size to be a multiple
+    of 128 and the hidden size a multiple of 256. Misaligned experts, gated or
+    not, are padded at load and, with the activations padded and the output
+    cropped as the MoE runner does, must still match the unquantized reference
+    on the dequantized weights."""
+    e, topk = 8, 2
+    is_gated = activation.is_gated
+    block_shape = [1, 32]
+    set_random_seed(7)
+    with set_current_vllm_config(vllm_config):
+        w1, w2, w1_scale, w2_scale = _make_mxfp8_moe_weights(e, n, k, is_gated)
+        w1_ref = dequant_mxfp8_to_bf16(w1, w1_scale)
+        w2_ref = dequant_mxfp8_to_bf16(w2, w2_scale)
+
+        layer = SimpleNamespace(
+            weight_block_size=block_shape,
+            moe_config=SimpleNamespace(
+                is_act_and_mul=is_gated, intermediate_size_per_partition=n
+            ),
+            activation=activation,
+        )
+        w1, w2, w1_scale, w2_scale = convert_to_fp8_moe_kernel_format(
+            Fp8MoeBackend.FLASHINFER_TRTLLM,
+            layer,
+            w1,
+            w2,
+            w1_scale,
+            w2_scale,
+            w13_input_scale=None,
+            w2_input_scale=None,
+        )
+        padded_n = layer.moe_config.intermediate_size_per_partition
+        padded_k = layer.moe_config.hidden_dim
+        assert padded_n == round_up(n, 128)
+        assert padded_k == round_up(k, 256)
+
+        quant_config = make_fp8_moe_quant_config(
+            Fp8MoeBackend.FLASHINFER_TRTLLM,
+            w1_scale,
+            w2_scale,
+            a1_scale=None,
+            a2_scale=None,
+            block_shape=block_shape,
+        )
+        moe_config = FusedMoEConfig(
+            num_experts=e,
+            experts_per_token=topk,
+            hidden_dim=padded_k,
+            intermediate_size=padded_n,
+            num_local_experts=e,
+            num_logical_experts=e,
+            activation=activation,
+            device="cuda",
+            moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),
+            in_dtype=torch.bfloat16,
+            routing_method=RoutingMethodType.Renormalize,
+            max_num_tokens=next_power_of_2(m),
+        )
+        # The monolithic experts pass the activation type to the kernel; the
+        # modular ones always run it gated.
+        kernel = mk.FusedMoEKernel(
+            maybe_make_prepare_finalize(
+                moe=moe_config,
+                quant_config=quant_config,
+                allow_new_interface=True,
+                use_monolithic=True,
+            ),
+            TrtLlmFp8ExpertsMonolithic(
+                moe_config=moe_config, quant_config=quant_config
+            ),
+        )
+
+        a = torch.randn((m, k), device="cuda", dtype=torch.bfloat16) / 10
+        logits = torch.randn((m, e), device="cuda", dtype=torch.float32)
+        output = kernel.apply_monolithic(
+            hidden_states=torch.nn.functional.pad(a, (0, padded_k - k)),
+            w1=w1,
+            w2=w2,
+            router_logits=logits,
+            activation=activation,
+            global_num_experts=e,
+            expert_map=None,
+            apply_router_weight_on_input=False,
+            routed_scaling_factor=1.0,
+            routing_sink=None,
+        )[:, :k]
+        topk_weights, topk_ids, _ = fused_topk(a, logits, topk, renormalize=True)
+        reference = fused_experts(
+            a,
+            w1_ref,
+            w2_ref,
+            topk_weights=topk_weights,
+            topk_ids=topk_ids,
+            activation=activation,
+            global_num_experts=e,
+            expert_map=None,
+        )
+
+    check_accuracy(ref_output=reference, actual_output=output)

tests/kernels/moe/test_ocp_mx_moe.py

@@ -42,13 +42,15 @@ ROCM_AVAILABLE = current_platform.is_rocm()
 ROCM_TRITON_KERNELS_AVAILABLE = False
 ROCM_AITER_AVAILABLE = is_aiter_found()
 ROCM_GFX950 = False
+ROCM_GFX1250 = False
 
 if ROCM_AVAILABLE:
-    from vllm.platforms.rocm import on_gfx950
+    from vllm.platforms.rocm import on_gfx950, on_gfx1250
     from vllm.utils.import_utils import has_triton_kernels
 
     ROCM_TRITON_KERNELS_AVAILABLE = has_triton_kernels()
     ROCM_GFX950 = on_gfx950()
+    ROCM_GFX1250 = on_gfx1250()
 
     if ROCM_AITER_AVAILABLE:
         from aiter.ops.triton.quant import dynamic_mxfp4_quant
@@ -1568,6 +1570,15 @@ ROCM_BACKEND_CONFIGS = {
         "requires_aiter": True,
         "requires_gfx950": True,
     },
+    "AITER_TRITON_MXFP4_MXFP4": {
+        "activation": "SILU",
+        "act_type": "mxfp4",
+        "rtol": 1.0,
+        "percent": 0.8,
+        "requires_aiter": True,
+        "requires_gfx950": False,
+        "requires_gfx1250": True,
+    },
 }
 
 
@@ -1607,6 +1618,8 @@ def test_rocm_mxfp4_moe_oracle(
         pytest.skip(f"Backend {backend_name} requires AITER")
     if config["requires_gfx950"] and not ROCM_GFX950:
         pytest.skip(f"Backend {backend_name} requires GFX950")
+    if config.get("requires_gfx1250") and not ROCM_GFX1250:
+        pytest.skip(f"Backend {backend_name} requires GFX1250")
 
     import vllm.distributed.parallel_state as ps
     from vllm.config import VllmConfig, set_current_vllm_config

tests/kernels/moe/test_rocm_aiter_topk.py

@@ -219,3 +219,169 @@ def test_rocm_aiter_grouped_topk_torch_compile_compatibility():
         topk_weights_original, topk_weights_compiled, rtol=1e-2, atol=1e-2
     )
     assert torch.allclose(topk_ids_original, topk_ids_compiled)
+
+
+def test_rocm_aiter_topk_gating_custom_op_registration():
+    assert callable(torch.ops.vllm.rocm_aiter_topk_gating)
+
+
+def _sort_routing(weights: torch.Tensor, ids: torch.Tensor):
+    ids_sorted, order = torch.sort(ids, dim=-1)
+    weights_sorted = torch.gather(weights, 1, order)
+    return weights_sorted, ids_sorted
+
+
+@pytest.mark.parametrize("num_tokens", [1, 4, 16, 64])
+@pytest.mark.parametrize("num_experts", [256, 512])
+@pytest.mark.parametrize("topk", [8, 10])
+@pytest.mark.parametrize("renormalize", [True, False])
+def test_rocm_aiter_topk_gating_matches_softmax_reference(
+    num_tokens: int,
+    num_experts: int,
+    topk: int,
+    renormalize: bool,
+):
+    """Numerical check vs softmax + topk, including Qwen3.8 (E=512, k=10)."""
+    torch.manual_seed(0)
+    gating_output = torch.randn(
+        (num_tokens, num_experts), dtype=torch.bfloat16, device="cuda"
+    )
+    scores = torch.softmax(gating_output.float(), dim=-1)
+    cutoff = torch.topk(scores, k=topk, dim=-1).values[:, -1:]
+
+    topk_weights = torch.empty((num_tokens, topk), dtype=torch.float32, device="cuda")
+    topk_ids = torch.empty((num_tokens, topk), dtype=torch.int32, device="cuda")
+    token_expert_indices = torch.empty(
+        (num_tokens, topk), dtype=torch.int32, device="cuda"
+    )
+    torch.ops.vllm.rocm_aiter_topk_gating(
+        topk_weights,
+        topk_ids,
+        token_expert_indices,
+        gating_output,
+        renormalize,
+        0,
+        "",
+    )
+
+    # BF16 logits can tie at the top-k boundary, so a different tied expert is
+    # valid. Check selection against the cutoff and weights against the scores
+    # of the experts the kernel actually selected.
+    assert ((topk_ids >= 0) & (topk_ids < num_experts)).all()
+    sorted_ids = torch.sort(topk_ids, dim=-1).values
+    assert (sorted_ids[:, 1:] != sorted_ids[:, :-1]).all()
+    selected_scores = torch.gather(scores, 1, topk_ids.long())
+    assert (selected_scores >= cutoff).all()
+    expected_weights = selected_scores
+    if renormalize:
+        expected_weights = expected_weights / expected_weights.sum(dim=-1, keepdim=True)
+    torch.testing.assert_close(topk_weights, expected_weights, atol=2e-2, rtol=2e-2)
+
+
+def test_rocm_aiter_topk_gating_shared_expert_suffix():
+    """Fused shared expert: softmax top-k on the prefix, sigmoid appended.
+
+    Matches the Qwen3.8 FSE launch: 512 routed experts, top-10, one shared
+    column, weights wide enough for the suffix, and the shared id prefilled.
+    """
+    torch.manual_seed(0)
+    num_tokens, num_experts, topk, num_shared = 4, 512, 10, 1
+    routed = torch.randn((num_tokens, num_experts), dtype=torch.bfloat16, device="cuda")
+    shared = torch.randn((num_tokens, num_shared), dtype=torch.bfloat16, device="cuda")
+    shared[0, 0] = 40
+    gating_output = torch.cat([routed, shared], dim=-1).contiguous()
+    width = topk + num_shared
+    shared_ids = torch.arange(
+        num_experts, num_experts + num_shared, device="cuda", dtype=torch.int32
+    )
+    topk_ids = torch.full((num_tokens, width), -1, dtype=torch.int32, device="cuda")
+    topk_ids[:, topk:] = shared_ids
+    topk_weights = torch.empty((num_tokens, width), dtype=torch.float32, device="cuda")
+    token_expert_indices = torch.empty(
+        (num_tokens, topk), dtype=torch.int32, device="cuda"
+    )
+
+    torch.ops.vllm.rocm_aiter_topk_gating(
+        topk_weights,
+        topk_ids[:, :topk],
+        token_expert_indices,
+        gating_output,
+        True,
+        num_shared,
+        "sigmoid",
+    )
+
+    scores = torch.softmax(routed.float(), dim=-1)
+    cutoff = torch.topk(scores, k=topk, dim=-1).values[:, -1:]
+    routed_ids = topk_ids[:, :topk]
+    assert ((routed_ids >= 0) & (routed_ids < num_experts)).all()
+    sorted_ids = torch.sort(routed_ids, dim=-1).values
+    assert (sorted_ids[:, 1:] != sorted_ids[:, :-1]).all()
+    selected = torch.gather(scores, 1, routed_ids.long())
+    assert (selected >= cutoff).all()
+    expected = selected / selected.sum(dim=-1, keepdim=True)
+    torch.testing.assert_close(topk_weights[:, :topk], expected, atol=2e-2, rtol=2e-2)
+    torch.testing.assert_close(
+        topk_weights[:, topk:],
+        torch.sigmoid(shared.float()),
+        atol=2e-2,
+        rtol=2e-2,
+    )
+    assert torch.equal(topk_ids[:, topk:], shared_ids.expand_as(topk_ids[:, topk:]))
+
+
+def test_rocm_aiter_topk_gating_torch_compile_compatibility():
+    token, expert, topk, renormalize = 4, 512, 10, True
+    gating_output = torch.randn((token, expert), dtype=torch.bfloat16, device="cuda")
+    device = gating_output.device
+
+    def gating_fn(gating_output, topk_weights, topk_ids, token_expert_indices):
+        return torch.ops.vllm.rocm_aiter_topk_gating(
+            topk_weights,
+            topk_ids,
+            token_expert_indices,
+            gating_output,
+            renormalize,
+            0,
+            "",
+        )
+
+    torch.library.opcheck(
+        torch.ops.vllm.rocm_aiter_topk_gating,
+        (
+            torch.empty((token, topk), dtype=torch.float32, device=device),
+            torch.empty((token, topk), dtype=torch.int32, device=device),
+            torch.empty((token, topk), dtype=torch.int32, device=device),
+            gating_output,
+        ),
+        kwargs={
+            "renormalize": renormalize,
+            "num_shared_experts": 0,
+            "shared_expert_scoring_func": "",
+        },
+        test_utils=("test_faketensor"),
+    )
+
+    compiled_fn = torch.compile(
+        gating_fn,
+        fullgraph=True,
+        backend="inductor",
+        mode="reduce-overhead",
+        dynamic=False,
+    )
+
+    def alloc():
+        return (
+            torch.empty((token, topk), dtype=torch.float32, device=device),
+            torch.empty((token, topk), dtype=torch.int32, device=device),
+            torch.empty((token, topk), dtype=torch.int32, device=device),
+        )
+
+    w0, i0, t0 = alloc()
+    w1, i1, t1 = alloc()
+    gating_fn(gating_output, w0, i0, t0)
+    compiled_fn(gating_output, w1, i1, t1)
+    w0, i0 = _sort_routing(w0, i0)
+    w1, i1 = _sort_routing(w1, i1)
+    assert torch.allclose(w0, w1, rtol=1e-2, atol=1e-2)
+    assert torch.equal(i0, i1)

tests/kernels/moe/test_topk_softplus_sqrt.py

@@ -81,32 +81,11 @@ def _torch_topk_softplus_sqrt(
     return topk_weights.to(torch.float32), topk_ids.to(torch.int32)
 
 
-@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm backend selection")
-@pytest.mark.parametrize("has_hash_routing", [False, True])
-@pytest.mark.parametrize(
-    "backend,expected_experts",
-    [
-        ("aiter_triton_mxfp4_bf16", "AiterW4A16ExpertsMonolithic"),
-        ("aiter", "AiterMxfp4Experts"),
-        ("triton_unfused", "UnfusedOAITritonExperts"),
-    ],
-)
-def test_hash_routing_backend_selection(
-    dist_init,
-    default_vllm_config,
-    monkeypatch,
-    has_hash_routing,
-    backend,
-    expected_experts,
-):
-    """Hash tables must reach backend selection; modular AITER remains eligible."""
+def _dsv4_moe_config(monkeypatch, default_vllm_config, has_hash_routing, backend):
     from dataclasses import replace
 
     from vllm._aiter_ops import rocm_aiter_ops
     from vllm.model_executor.layers.fused_moe.layer import FusedMoEFactory
-    from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import (
-        select_deepseek_v4_mxfp4_moe_backend,
-    )
     from vllm.platforms import rocm
 
     # Exercise both native MX and monolithic AITER selection on any ROCm device.
@@ -130,14 +109,63 @@ def test_hash_routing_backend_selection(
         ),
         prefix="hash_routing_selection",
     )
-    config = replace(layer.routed_experts.moe_config, moe_backend=backend)
+    return layer, replace(layer.routed_experts.moe_config, moe_backend=backend)
 
-    if has_hash_routing and backend == "aiter_triton_mxfp4_bf16":
-        with pytest.raises(ValueError, match="hash routing"):
-            select_deepseek_v4_mxfp4_moe_backend(config)
-    else:
-        _, experts_cls = select_deepseek_v4_mxfp4_moe_backend(config)
-        assert experts_cls.__name__ == expected_experts
+
+@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm backend selection")
+@pytest.mark.parametrize("has_hash_routing", [False, True])
+@pytest.mark.parametrize(
+    "backend,expected_experts",
+    [
+        ("aiter_triton_mxfp4_bf16", "AiterW4A16ExpertsMonolithic"),
+        ("aiter", "AiterMxfp4Experts"),
+        ("triton_unfused", "UnfusedOAITritonExperts"),
+    ],
+)
+def test_hash_routing_backend_selection(
+    dist_init,
+    default_vllm_config,
+    monkeypatch,
+    has_hash_routing,
+    backend,
+    expected_experts,
+):
+    """Hash tables reach backend selection and hash-capable kernels."""
+    from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import (
+        select_deepseek_v4_mxfp4_moe_backend,
+    )
+
+    layer, config = _dsv4_moe_config(
+        monkeypatch, default_vllm_config, has_hash_routing, backend
+    )
+    _, experts_cls = select_deepseek_v4_mxfp4_moe_backend(config)
+    assert experts_cls.__name__ == expected_experts
+    # Monolithic kernels route internally, so they read the table from the layer.
+    assert (layer.routed_experts.hash_indices_table is not None) == has_hash_routing
+
+
+@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm backend selection")
+def test_hash_routing_rejects_monolithic_without_support(
+    dist_init, default_vllm_config, monkeypatch
+):
+    """A monolithic kernel that does not implement hash routing is rejected."""
+    from vllm.model_executor.layers.fused_moe.experts.aiter_mxfp4_w4a16_moe import (
+        AiterW4A16ExpertsMonolithic,
+    )
+    from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import (
+        select_deepseek_v4_mxfp4_moe_backend,
+    )
+
+    monkeypatch.setattr(
+        AiterW4A16ExpertsMonolithic,
+        "_supports_hash_routing",
+        staticmethod(lambda: False),
+    )
+    _, config = _dsv4_moe_config(
+        monkeypatch, default_vllm_config, True, "aiter_triton_mxfp4_bf16"
+    )
+    with pytest.raises(ValueError, match="hash routing"):
+        select_deepseek_v4_mxfp4_moe_backend(config)
 
 
 def test_torch_topk_softplus_sqrt_breaks_ties_by_expert_id():

tests/kernels/quantization/test_rocm_mxfp4.py

@@ -50,7 +50,7 @@ SKINNY_GEMM_PASS_RATES = {
 }
 PRESHUFFLED_SHAPES = [
     (64, 4096, 8192),
-    (32, 8192, 16384),
+    (32, 16384, 16384),
 ]
 
 
@@ -421,8 +421,12 @@ def test_aiter_fp4_gemm_preshuffled_tuned_shapes(shape):
     M, K, N = shape
 
     assert M <= 64
-    # The predicate takes K as packed bytes, matching `weight.shape[1]`.
+    # The predicate takes K as packed bytes, matching `weight.shape[1]`, and
+    # must agree with the config the kernel itself resolves.
     assert rocm_aiter_ops.is_triton_gemm_afp4wfp4_presh_ws_tuned(N, K // 2)
+    from aiter.ops.triton.utils.gemm_config_utils import get_gemm_config
+
+    assert get_gemm_config("GEMM-AFP4WFP4_PRESHUFFLED", M, N, K)[1]
 
     A = torch.randn(M, K, dtype=torch.bfloat16)
     B = torch.randn(N, K, dtype=torch.bfloat16)
@@ -464,6 +468,18 @@ def test_aiter_fp4_gemm_preshuffled_tuned_shapes(shape):
     _assert_deterministic(run_preshuffled, n_runs=3)
 
 
+@pytest.mark.skipif(not on_gfx950(), reason="gfx950 ROCm only")
+def test_aiter_fp4_preshuffled_tuned_gate_uses_logical_k():
+    """AITER ships N=16384-K=16384 but not N=16384-K=8192, so a K=8192 layer
+    has no tuned preshuffled config even though twice its K does."""
+    _assert_aiter_supported()
+    from vllm._aiter_ops import rocm_aiter_ops
+
+    rocm_aiter_ops.is_triton_gemm_afp4wfp4_presh_ws_tuned.cache_clear()
+    assert rocm_aiter_ops.is_triton_gemm_afp4wfp4_presh_ws_tuned(16384, 16384 // 2)
+    assert not rocm_aiter_ops.is_triton_gemm_afp4wfp4_presh_ws_tuned(16384, 8192 // 2)
+
+
 @pytest.mark.skipif(not on_gfx950(), reason="gfx950 ROCm only")
 def test_aiter_fp4_gemm_a4w4_determinism():
     """AITER FP4 A4W4 GEMM is bitwise deterministic across repeated runs.

tests/kernels/test_compressor_kv_cache.py

@@ -14,6 +14,7 @@ These tests cover:
 import math
 from types import SimpleNamespace
 
+import numpy as np
 import pytest
 import torch
 
@@ -157,6 +158,104 @@ def test_dspark_context_kv_matches_query_insert(
         torch.testing.assert_close(kv, before, rtol=0, atol=0)
 
 
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA only")
+@pytest.mark.parametrize("kv_mxfp8", [False, True])
+def test_v41_pcp_swa_cache_matches_non_pcp_writer(kv_mxfp8, monkeypatch):
+    """PCP SWA writes must preserve the established CUDA cache bytes."""
+    from vllm.models.deepseek_v41 import attention as attention_module
+    from vllm.models.deepseek_v41.attention import DeepseekV4Attention
+
+    torch.manual_seed(43)
+    block_size = 64
+    num_local_tokens = 5
+    num_tokens = 2 * num_local_tokens
+    row_size = 528 if kv_mxfp8 else 584
+    num_blocks = 2
+    page_stride = math.ceil(block_size * row_size / 576) * 576 + 576
+    backing = torch.full(
+        (num_blocks, page_stride), 165, dtype=torch.uint8, device="cuda"
+    )
+    cache = backing.as_strided(
+        (num_blocks, block_size, row_size), (page_stride, row_size, 1)
+    )
+    reference_backing = backing.clone()
+    reference = reference_backing.as_strided(cache.shape, cache.stride())
+
+    local_kv = torch.randn(num_local_tokens, 512, dtype=torch.bfloat16, device="cuda")
+    remote_kv = torch.randn_like(local_kv)
+    gathered_kv = torch.cat((local_kv, remote_kv))
+    local_positions = torch.arange(num_local_tokens, device="cuda") + 3
+    cache_positions = torch.cat((local_positions, local_positions + 11))
+    cache_slots = torch.randperm(
+        num_blocks * block_size, dtype=torch.int64, device="cuda"
+    )[:num_tokens]
+    cache_slots[2] = -1
+    angles = torch.randn(32, 32, device="cuda")
+    cos_sin = torch.cat((angles.cos(), angles.sin()), dim=-1)
+    num_heads = 8
+    q = torch.randn(
+        num_local_tokens,
+        num_heads,
+        512,
+        dtype=torch.bfloat16,
+        device="cuda",
+    )
+
+    torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
+        torch.zeros(num_tokens, num_heads, 512, dtype=torch.bfloat16, device="cuda"),
+        gathered_kv,
+        reference.view(num_blocks, -1),
+        cache_slots,
+        cache_positions,
+        cos_sin,
+        num_heads,
+        1e-20,
+        block_size,
+        False,
+        kv_mxfp8,
+    )
+
+    monkeypatch.setattr(
+        attention_module,
+        "get_pcp_group",
+        lambda: SimpleNamespace(
+            all_gather=lambda tensor, dim: torch.cat((tensor, remote_kv), dim=dim)
+        ),
+    )
+    metadata = SimpleNamespace(
+        slot_mapping=torch.full(
+            (num_local_tokens,), -1, dtype=torch.int64, device="cuda"
+        ),
+        cache_slot_mapping=cache_slots,
+        cache_positions=cache_positions,
+        block_size=block_size,
+    )
+    attention = SimpleNamespace(
+        use_pcp=True,
+        swa_cache_layer=SimpleNamespace(
+            prefix="swa", kv_cache=cache, block_size=block_size
+        ),
+        rotary_emb=SimpleNamespace(cos_sin_cache=cos_sin),
+        _flashinfer_fp8_kv_scale=None,
+        kv_mxfp8=kv_mxfp8,
+        accepts_unnormed_unroped_query=False,
+        n_local_heads=num_heads,
+        padded_heads=num_heads,
+        eps=1e-20,
+        head_dim=512,
+    )
+    DeepseekV4Attention._fused_qnorm_rope_kv_insert(
+        attention, q, local_kv, local_positions, {"swa": metadata}
+    )
+
+    torch.testing.assert_close(
+        backing.view(torch.uint8),
+        reference_backing.view(torch.uint8),
+        rtol=0,
+        atol=0,
+    )
+
+
 @pytest.mark.skipif(
     not current_platform.is_cuda_alike(), reason="graph capture coverage"
 )
@@ -360,7 +459,7 @@ def _mxfp8_record_reference(
 @pytest.mark.parametrize("compress_ratio", [1, 2])
 def test_v41_rope_insert_mxfp8_record(compress_ratio: int):
     """The 528-byte V4.1 record quantizes the RoPE dims instead of keeping them
-    in bf16, so a page is [block_size x 512 fp8][block_size x 16 UE8M0].
+    in bf16, so a token is [512 fp8][16 UE8M0].
 
     Byte-exact against the reference quantizer, including the rows and page
     padding the kernel must leave alone.
@@ -398,9 +497,8 @@ def test_v41_rope_insert_mxfp8_record(compress_ratio: int):
         values, scales = _mxfp8_record_reference(
             latent[t], cos_sin[pos // compress_ratio * compress_ratio]
         )
-        expected[page, row * 512 : (row + 1) * 512] = values
-        scale_offset = cache_block * 512 + row * 16
-        expected[page, scale_offset : scale_offset + 16] = scales
+        expected[page, row * 528 : row * 528 + 512] = values
+        expected[page, row * 528 + 512 : (row + 1) * 528] = scales
 
     torch.testing.assert_close(cache_backing, expected, rtol=0, atol=0)
 
@@ -510,14 +608,13 @@ def test_v41_mxfp8_cache_round_trip():
         values, scales = _ue8m0_reference(k[t].float(), 32, 448.0)
         page, row = divmod(t, block_size)
         torch.testing.assert_close(
-            cache[page, row * 512 : (row + 1) * 512],
+            cache[page, row * 528 : row * 528 + 512],
             values.view(torch.uint8),
             rtol=0,
             atol=0,
         )
-        scale_offset = block_size * 512 + row * 16
         torch.testing.assert_close(
-            cache[page, scale_offset : scale_offset + 16],
+            cache[page, row * 528 + 512 : (row + 1) * 528],
             (scales.log2() + 127).to(torch.uint8),
             rtol=0,
             atol=0,
@@ -564,7 +661,8 @@ def _decode_nvfp4_row(packed: torch.Tensor, scales: torch.Tensor) -> torch.Tenso
 @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA only")
 @pytest.mark.parametrize("compress_ratio", [1, 2])
 def test_v41_rope_insert_nvfp4_record(compress_ratio: int):
-    """The 288-byte V4.1 NVFP4 record: 256 B of e2m1 pairs then 32 e4m3 scales.
+    """The 288-byte V4.1 NVFP4 record: 256 B of e2m1 pairs then 32 e4m3 scales
+    per token.
 
     Scales are byte-exact against the reference (``amax / 6`` clamped to the
     e4m3 range); values are checked as round-to-nearest on the e2m1 grid, whose
@@ -604,17 +702,13 @@ def test_v41_rope_insert_nvfp4_record(compress_ratio: int):
         amax = rotated.view(32, 16).abs().amax(-1)
         scale = (amax / 6.0).clamp(2.0**-9, 448.0).to(torch.float8_e4m3fn)
 
-        got_scales = cache_backing[page, cache_block * 256 + row * 32 :][:32]
+        record = cache_backing[page, row * 288 : (row + 1) * 288]
+        got_scales = record[256:]
         torch.testing.assert_close(got_scales, scale.view(torch.uint8), rtol=0, atol=0)
-        decoded = _decode_nvfp4_row(
-            cache_backing[page, row * 256 : (row + 1) * 256], got_scales
-        )
+        decoded = _decode_nvfp4_row(record[:256], got_scales)
         step = scale.float().repeat_interleave(16)
         assert ((decoded - rotated).abs() <= step).all()
-        untouched[page, row * 256 : (row + 1) * 256] = cache_backing[
-            page, row * 256 : (row + 1) * 256
-        ]
-        untouched[page, cache_block * 256 + row * 32 :][:32] = got_scales
+        untouched[page, row * 288 : (row + 1) * 288] = record
     # Rows the kernel must not have touched, page padding included.
     torch.testing.assert_close(cache_backing, untouched, rtol=0, atol=0)
 
@@ -670,10 +764,8 @@ def test_v41_nvfp4_gather_matches_insert(num_reqs, num_tokens, partial):
             offset=offset,
         )
 
-    expected = _decode_nvfp4_row(
-        backing[:, : block_size * 256].reshape(-1, 256),
-        backing[:, block_size * 256 : block_size * 288].reshape(-1, 32),
-    )
+    records = backing[:, : block_size * 288].reshape(-1, 288)
+    expected = _decode_nvfp4_row(records[:, :256], records[:, 256:])
 
     def check(shorter_by):
         assert (storage[:, num_tokens + offset :] == -123).all()
@@ -702,6 +794,72 @@ def test_v41_nvfp4_gather_matches_insert(num_reqs, num_tokens, partial):
     check(1)
 
 
+@pytest.mark.skipif(
+    not current_platform.is_device_capability_family(100),
+    reason="The V4.1 records are SM100-only",
+)
+@pytest.mark.parametrize("record", ["swa_fp8", "compressed_fp8", "compressed_fp4"])
+def test_v41_cache_writers_match_flashmla_decode(record: str):
+    """FlashMLA and vLLM's gathers read back the token a V4.1 writer put mid
+    page in a padded block, so all agree on the interleaved record layout."""
+    import vllm.v1.attention.ops.flashmla as fm
+    from vllm.models.deepseek_v41.common.ops.cache_utils import (
+        dequantize_and_gather_k_cache,
+        dequantize_and_gather_k_cache_triton,
+    )
+    from vllm.models.deepseek_v41.common.ops.fused_compress_quant_cache import (
+        rope_quant_insert,
+    )
+
+    device, block_size, topk = "cuda", 64, 64
+    slot = block_size + 37
+    bytes_per_token = 288 if record == "compressed_fp4" else 528
+    torch.manual_seed(0)
+    kv = torch.randn(1, 512, dtype=torch.bfloat16, device=device)
+    positions = torch.tensor([slot], device=device)
+    # cos = 1, sin = 0: RoPE is the identity, so the cache should hold `kv`.
+    cos_sin = torch.zeros(slot + 1, 64, device=device)
+    cos_sin[:, :32] = 1
+    # Padded like a packed block: a multiple of the record and of 512 B.
+    cache = torch.zeros(2, block_size + 32, bytes_per_token, dtype=torch.uint8)
+    cache = cache.to(device)[:, :block_size]
+    if record == "swa_fp8":
+        torch.ops._C.fused_deepseek_v4_kv_rope_insert(
+            kv, cache, positions, positions, cos_sin, block_size, None, True
+        )
+    else:
+        rope_quant_insert(kv, positions, cos_sin, cache, positions, 1)
+
+    indices = torch.full((1, 1, topk), -1, dtype=torch.int32, device=device)
+    indices[..., 0] = slot
+    extra = None
+    if record == "compressed_fp4":
+        cache, extra = torch.zeros(2, block_size, 528).to(cache), cache
+    out, _ = fm.flash_mla_with_kvcache(
+        q=torch.randn(1, 1, 64, 512, dtype=torch.bfloat16, device=device),
+        k_cache=cache.unsqueeze(2),
+        block_table=None,
+        cache_seqlens=None,
+        head_dim_v=512,
+        tile_scheduler_metadata=fm.get_mla_metadata()[0],
+        is_fp8_kvcache=True,
+        indices=indices if extra is None else torch.full_like(indices, -1),
+        extra_k_cache=None if extra is None else extra.unsqueeze(2),
+        extra_indices_in_kvcache=None if extra is None else indices,
+    )
+    for gather in (dequantize_and_gather_k_cache, dequantize_and_gather_k_cache_triton):
+        rows = torch.zeros(1, slot + 1, 512, dtype=torch.bfloat16, device=device)
+        block_table = torch.arange(2, dtype=torch.int32, device=device)[None]
+        seq_lens = torch.tensor([slot + 1], dtype=torch.int32, device=device)
+        k_cache = cache if extra is None else extra
+        gather(rows, k_cache, seq_lens, None, block_table, block_size, offset=0)
+        torch.testing.assert_close(
+            out[0, 0], rows[0, slot].expand(64, -1), rtol=0, atol=0
+        )
+    cos = torch.nn.functional.cosine_similarity(rows[0, slot].float(), kv[0].float(), 0)
+    assert cos > (0.98 if extra is not None else 0.999)
+
+
 @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA only")
 @pytest.mark.parametrize("compress_ratio", [1, 2])
 @pytest.mark.parametrize("store_fp8", [False, True])
@@ -796,6 +954,7 @@ def test_v41_compressor_metadata_maps_tokens_to_their_ring():
     capacity = 8
     vllm_config = MagicMock()
     vllm_config.scheduler_config.max_num_batched_tokens = 16
+    vllm_config.parallel_config.prefill_context_parallel_size = 1
     spec = CircularBufferSpec(
         block_size=capacity,
         num_kv_heads=1,
@@ -832,6 +991,67 @@ def test_v41_compressor_metadata_maps_tokens_to_their_ring():
     assert metadata.token_to_req_indices.tolist() == [0, 0, 0, 1, 1, 2, 2]
 
 
+def test_v41_pcp_compressor_latent_is_gathered_for_cache_writes(monkeypatch):
+    """PCP compresses locally and publishes the gathered latent to both caches."""
+    from vllm.forward_context import ForwardContext, override_forward_context
+    from vllm.models.deepseek_v41 import attention as attention_module
+    from vllm.models.deepseek_v41.attention import DeepseekV4Attention
+
+    local_latent = torch.arange(8, dtype=torch.float32).view(2, 4)
+    cache_positions = torch.tensor([0, 1, 10, 11])
+    observed = {}
+
+    class Compressor:
+        def __call__(self, kv_score, positions):
+            return local_latent
+
+        def insert_cache(self, latent, positions, rotary_emb):
+            observed["insert"] = (latent, positions)
+
+    def prepare_indexer(qr, latent, weights, positions, rotary, qr_scale, **kwargs):
+        observed["indexer"] = (latent, kwargs)
+        return None, None, None
+
+    monkeypatch.setattr(
+        attention_module,
+        "get_pcp_group",
+        lambda: SimpleNamespace(
+            all_gather=lambda tensor, dim: torch.cat((tensor, tensor + 10), dim=dim)
+        ),
+    )
+    metadata = {"swa": SimpleNamespace(cache_positions=cache_positions)}
+    attention = SimpleNamespace(
+        use_pcp=True,
+        compressor=Compressor(),
+        indexer=prepare_indexer,
+        aux_stream_list=None,
+        ln_events=[None, None],
+        swa_cache_layer=SimpleNamespace(prefix="swa"),
+        rotary_emb=SimpleNamespace(),
+        indexer_rotary_emb=SimpleNamespace(),
+        n_local_heads=1,
+        head_dim=4,
+        _wq_b_proj=lambda qr, scale: qr,
+        _fused_qnorm_rope_kv_insert=lambda q, kv, pos, meta: q,
+        _sparse_indexer_and_attn=lambda *args: None,
+    )
+    rows = torch.zeros(2, 4)
+    positions = torch.arange(2)
+    with override_forward_context(ForwardContext({}, metadata, {})):
+        DeepseekV4Attention._prepare_and_attn(
+            attention, rows, rows, rows, None, rows, rows, positions, rows
+        )
+
+    expected_global = torch.cat((local_latent, local_latent + 10))
+    insert_latent, insert_positions = observed["insert"]
+    torch.testing.assert_close(insert_latent, expected_global)
+    torch.testing.assert_close(insert_positions, cache_positions)
+    indexer_latent, indexer_kwargs = observed["indexer"]
+    torch.testing.assert_close(indexer_latent, local_latent)
+    torch.testing.assert_close(indexer_kwargs["cache_latent"], expected_global)
+    torch.testing.assert_close(indexer_kwargs["cache_positions"], cache_positions)
+
+
 @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA stream coverage")
 @pytest.mark.parametrize(
     "use_aux,use_graph", [(False, False), (True, False), (True, True)]
@@ -871,6 +1091,7 @@ def test_v41_attention_joins_cache_writes_before_consumption(use_aux, use_graph)
     context = ForwardContext({}, metadata, {})
     compressor = DeepseekCompressor.__new__(DeepseekCompressor)
     torch.nn.Module.__init__(compressor)
+    compressor.pcp_size = 1
     compressor.head_dim, compressor.rope_head_dim, compressor.compress_ratio = (
         512,
         64,
@@ -896,7 +1117,7 @@ def test_v41_attention_joins_cache_writes_before_consumption(use_aux, use_graph)
         use_fp4_kv=False,
     )
 
-    def prepare_indexer(qr, latent, weights, positions, rotary, qr_scale):
+    def prepare_indexer(qr, latent, weights, positions, rotary, qr_scale, **kwargs):
         DeepseekV4Indexer._produce_k(indexer, latent, positions, rotary)
         return None, None, None
 
@@ -906,6 +1127,7 @@ def test_v41_attention_joins_cache_writes_before_consumption(use_aux, use_graph)
 
     auxiliary = [torch.cuda.Stream()] if use_aux else None
     attention = SimpleNamespace(
+        use_pcp=False,
         compressor=compressor,
         indexer=prepare_indexer,
         aux_stream_list=auxiliary,
@@ -963,6 +1185,291 @@ def test_v41_attention_joins_cache_writes_before_consumption(use_aux, use_graph)
         previous = expected
 
 
+@pytest.mark.parametrize("pcp_size", [2, 4])
+@pytest.mark.parametrize(
+    "lengths,prefilling,starts",
+    [
+        ([9, 1], [True, False], [1, 4]),
+        ([1, 1], [False, False], [16, 3]),
+        ([3, 1], [True, False], [1, 0]),
+    ],
+)
+def test_v41_pcp_ring_links_match_request_positions(
+    pcp_size, lengths, prefilling, starts, monkeypatch
+):
+    """Cross-chunk predecessors and open ring writes must belong to the request."""
+    from vllm.models.deepseek_v41 import compressor as compressor_module
+    from vllm.models.deepseek_v41.compressor import (
+        CompressorMetadata,
+        CompressorMetadataBuilder,
+    )
+    from vllm.v1.worker.gpu import pcp_manager as pcp_manager_module
+    from vllm.v1.worker.gpu.pcp_manager import PCPManager
+
+    monkeypatch.setattr(
+        pcp_manager_module,
+        "async_tensor_h2d",
+        lambda array, **kwargs: torch.from_numpy(array),
+    )
+    manager = PCPManager(pcp_size, 0, torch.device("cpu"), shard_decode_requests=True)
+    query_start = np.array([0, *np.cumsum(lengths)], dtype=np.int32)
+    segments, counts = manager._build_batch_layout(
+        np.array(lengths), np.array(starts), np.array(prefilling), query_start
+    )
+    width = max(counts)
+    capacity = 8
+    blocks = [1, (1 << 20) + 2]
+    num_rows = min(width, 2 * len(lengths))
+    metadata_by_rank = []
+    positions_by_rank = []
+    gathered_keys = []
+    for local_segments in segments:
+        positions = torch.zeros(width, dtype=torch.int64)
+        slots = torch.full((width,), -1, dtype=torch.int64)
+        local_query_start = [0]
+        keys = [-1] * num_rows
+        for row, segment in enumerate(local_segments):
+            req = segment.global_batch_req_idx
+            begin = starts[req] + segment.global_batch_slice.start - query_start[req]
+            local = segment.rank_local_batch_slice
+            positions[local] = torch.arange(begin, begin + segment.num_tokens)
+            slots[local] = blocks[req] * capacity + positions[local] % capacity
+            local_query_start.append(local.stop)
+            keys[row] = (blocks[req] << 32) | int(positions[local.stop - 1])
+        if not local_segments:
+            local_query_start.append(0)
+        metadata_by_rank.append(
+            CompressorMetadata(
+                slot_mapping=slots,
+                query_start_loc=torch.tensor(local_query_start, dtype=torch.int32),
+                token_to_req_indices=torch.empty(0, dtype=torch.int32),
+            )
+        )
+        positions_by_rank.append(positions)
+        gathered_keys.extend(keys)
+
+    keys = torch.tensor(gathered_keys, dtype=torch.int64)
+    contributed_keys = []
+
+    def all_gather(tensor, dim):
+        contributed_keys.append(tensor.clone())
+        return keys
+
+    monkeypatch.setattr(
+        compressor_module,
+        "get_pcp_group",
+        lambda: SimpleNamespace(all_gather=all_gather),
+    )
+    builder = SimpleNamespace(
+        capacity=capacity,
+        device=torch.device("cpu"),
+        vllm_config=SimpleNamespace(
+            scheduler_config=SimpleNamespace(max_num_seqs=len(lengths))
+        ),
+    )
+    expected_rows = {
+        block: max(key for key in gathered_keys if key >> 32 == block)
+        for block in blocks
+    }
+    expected_slots = [
+        block * capacity + (key & 0xFFFFFFFF) % capacity
+        for block, key in expected_rows.items()
+        if key % 2 == 0
+    ]
+    for metadata, positions in zip(metadata_by_rank, positions_by_rank):
+        num_reqs = metadata.query_start_loc.numel() - 1
+        CompressorMetadataBuilder._link_pcp_chunks(
+            builder, metadata, num_reqs, positions
+        )
+        assert metadata.prev_row_indices is not None
+        assert metadata.gather_tokens is not None
+        assert metadata.ring_write_slots is not None
+        for row, token in enumerate(metadata.gather_tokens.tolist()):
+            key = int(contributed_keys[-1][row])
+            if key >= 0:
+                assert int(metadata.slot_mapping[token]) // capacity == key >> 32
+                assert int(positions[token]) == key & 0xFFFFFFFF
+        for start, prev in zip(
+            metadata.query_start_loc[:-1].tolist(), metadata.prev_row_indices.tolist()
+        ):
+            position = int(positions[start])
+            block = int(metadata.slot_mapping[start]) // capacity
+            wanted = (block << 32) | (position - 1)
+            if position % 2 == 1 and wanted in gathered_keys:
+                assert prev >= 0 and gathered_keys[prev] == wanted
+            else:
+                assert prev == -1
+        writes = metadata.ring_write_slots.tolist()
+        assert sorted(slot for slot in writes if slot >= 0) == sorted(expected_slots)
+        for row, slot in enumerate(writes):
+            if slot >= 0:
+                assert gathered_keys[row] == expected_rows[slot // capacity]
+                assert (gathered_keys[row] & 0xFFFFFFFF) % capacity == slot % capacity
+    torch.testing.assert_close(torch.cat(contributed_keys), keys)
+
+
+def _build_pcp_ring_metadata(ranks, all_gather, spec, config, device, monkeypatch):
+    from vllm.models.deepseek_v41 import compressor as compressor_module
+    from vllm.models.deepseek_v41.compressor import CompressorMetadataBuilder
+
+    monkeypatch.setattr(
+        compressor_module,
+        "get_pcp_group",
+        lambda: SimpleNamespace(all_gather=all_gather),
+    )
+    for local in ranks:
+        builder = CompressorMetadataBuilder(spec, ["ring"], config, device)
+        local.metadata = builder.build(0, local.common)
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="Triton kernel")
+@pytest.mark.parametrize("pcp_size", [2, 4])
+def test_v41_pcp_compressor_preserves_pairs_and_decode_state(pcp_size, monkeypatch):
+    """Pairs split across PCP chunks and decode owner changes keep exact latents."""
+    from vllm.forward_context import ForwardContext, override_forward_context
+    from vllm.models.deepseek_v41 import compressor as compressor_module
+    from vllm.models.deepseek_v41.common.ops.fused_compress_quant_cache import (
+        fused_save_compress_norm,
+    )
+    from vllm.models.deepseek_v41.compressor import DeepseekCompressor
+    from vllm.v1.attention.backend import CommonAttentionMetadata
+    from vllm.v1.kv_cache_interface import CircularBufferSpec
+    from vllm.v1.worker.gpu.pcp_manager import PCPManager
+
+    torch.manual_seed(42)
+    device = torch.device("cuda")
+    norm = torch.randn(512, dtype=torch.bfloat16, device=device)
+    # Block 0 is reserved as the null block; requests use physical blocks 1/2.
+    reference_state = torch.zeros(3, 8, 1024, device=device)
+    rank_states = [reference_state.clone() for _ in range(pcp_size)]
+    config = SimpleNamespace(
+        scheduler_config=SimpleNamespace(max_num_batched_tokens=32, max_num_seqs=4),
+        parallel_config=SimpleNamespace(prefill_context_parallel_size=pcp_size),
+        speculative_config=None,
+    )
+    spec = CircularBufferSpec(
+        block_size=8, num_kv_heads=1, head_size=1024, head_size_v=0, dtype=torch.float32
+    )
+    starts = np.zeros(2, dtype=np.int32)
+    for lengths, prefilling in [
+        ([7, 3], [True, True]),
+        ([9, 1], [True, False]),
+        ([1, 1], [False, False]),
+        ([1, 1], [False, False]),
+    ]:
+        lengths = np.array(lengths, dtype=np.int32)
+        query_start = np.concatenate(([0], np.cumsum(lengths))).astype(np.int32)
+        qsl = torch.from_numpy(query_start).to(device)
+        positions = torch.cat(
+            [torch.arange(s, s + n, device=device) for s, n in zip(starts, lengths)]
+        )
+        req_ids = torch.repeat_interleave(
+            torch.arange(2, device=device, dtype=torch.int32),
+            torch.from_numpy(lengths).to(device),
+        )
+        slots = (req_ids.long() + 1) * 8 + positions % 8
+        raw = torch.randn(sum(lengths), 1024, device=device)
+        expected = torch.empty(sum(lengths), 512, dtype=torch.bfloat16, device=device)
+        fused_save_compress_norm(
+            raw,
+            positions,
+            reference_state,
+            slots,
+            qsl,
+            req_ids,
+            norm,
+            1e-20,
+            2,
+            expected,
+        )
+
+        ranks = []
+        for rank in range(pcp_size):
+            manager = PCPManager(pcp_size, rank, device, shard_decode_requests=True)
+            segments, counts = manager._build_batch_layout(
+                lengths, starts, np.array(prefilling), query_start
+            )
+            width = max(counts)
+            local_indices = manager._padded_gather_idx[
+                rank * width : (rank + 1) * width
+            ]
+            local_lengths = [s.num_tokens for s in segments[rank]] or [0]
+            local_qsl = torch.tensor([0, *np.cumsum(local_lengths)], dtype=torch.int32)
+            local_reqs = [s.global_batch_req_idx for s in segments[rank]] or [0]
+            common = CommonAttentionMetadata(
+                query_start_loc=local_qsl.to(device),
+                query_start_loc_cpu=local_qsl,
+                seq_lens=torch.zeros(len(local_reqs), dtype=torch.int32, device=device),
+                num_reqs=len(local_reqs),
+                num_actual_tokens=counts[rank],
+                max_query_len=max(local_lengths),
+                max_seq_len=int(max(starts + lengths)),
+                block_table_tensor=torch.tensor(
+                    np.array(local_reqs) + 1, dtype=torch.int32, device=device
+                ).unsqueeze(1),
+                slot_mapping=torch.full((width * pcp_size,), -1, device=device),
+                positions=positions[local_indices],
+            )
+            ranks.append(
+                SimpleNamespace(
+                    indices=local_indices[: counts[rank]],
+                    raw=raw[local_indices],
+                    positions=common.positions,
+                    common=common,
+                )
+            )
+
+        # Each rank's builder contributes its last-row candidates to a gather.
+        candidates: list[torch.Tensor] = []
+        _build_pcp_ring_metadata(
+            ranks,
+            lambda t, dim, out=candidates: out.append(t.clone()) or t.repeat(pcp_size),
+            spec,
+            config,
+            device,
+            monkeypatch,
+        )
+        _build_pcp_ring_metadata(
+            ranks,
+            lambda t, dim, out=candidates: torch.cat(out),
+            spec,
+            config,
+            device,
+            monkeypatch,
+        )
+        gathered_rows = torch.cat([r.raw[r.metadata.gather_tokens] for r in ranks])
+        monkeypatch.setattr(
+            compressor_module,
+            "get_pcp_group",
+            lambda rows=gathered_rows: SimpleNamespace(all_gather=lambda t, dim: rows),
+        )
+        for rank, local in enumerate(ranks):
+            compressor = SimpleNamespace(
+                pcp_size=pcp_size,
+                head_dim=512,
+                compress_ratio=2,
+                rms_norm_eps=1e-20,
+                norm=SimpleNamespace(weight=norm),
+                state_cache=SimpleNamespace(prefix="ring", kv_cache=rank_states[rank]),
+            )
+            context = ForwardContext({}, {"ring": local.metadata}, {})
+            with override_forward_context(context):
+                actual = DeepseekCompressor.forward(
+                    compressor, local.raw, local.positions
+                )
+            assert isinstance(actual, torch.Tensor)
+            valid = (positions[local.indices] + 1) % 2 == 0
+            torch.testing.assert_close(
+                actual[: local.indices.numel()][valid],
+                expected[local.indices][valid],
+                rtol=0,
+                atol=0,
+            )
+        for state in rank_states[1:]:
+            torch.testing.assert_close(state, rank_states[0], rtol=0, atol=0)
+        starts += lengths
+
+
 def _on_gfx950() -> bool:
     if not current_platform.is_rocm():
         return False

tests/kernels/test_engram.py

@@ -502,11 +502,12 @@ def test_lookback_window_reproduces_single_instance(runner, capture):
 
 
 @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required")
-def test_v2_model_state_gathers_lookback_window():
+@pytest.mark.parametrize("pcp", [False, True])
+def test_v2_model_state_gathers_lookback_window(pcp):
     """The window is gathered on device from the runner's token history."""
     from vllm.models.deepseek_v41.nvidia.model_state import DeepseekV41ModelState
 
-    depth, max_num_reqs, max_model_len = 3, 4, 16
+    depth, max_num_reqs, max_model_len = 3, 6, 16
     state = DeepseekV41ModelState.__new__(DeepseekV41ModelState)
     state.rope_state = None
     state.supports_mm_inputs = False
@@ -527,15 +528,23 @@ def test_v2_model_state_gathers_lookback_window():
     # Batch rows -> request state rows: a request two decode steps in, a fresh
     # request, and one at the end of a 5-token prompt.
     input_batch = SimpleNamespace(
-        idx_mapping=torch.tensor([2, 0, 3], dtype=torch.int64, device="cuda")
+        idx_mapping=torch.tensor(
+            [2, 0, 3, 2] if pcp else [2, 0, 3], dtype=torch.int64, device="cuda"
+        ),
+        query_start_loc=torch.arange(5 if pcp else 4, dtype=torch.int32, device="cuda"),
+        positions=torch.tensor([8, 0, 6, 5] if pcp else [8, 0, 6], device="cuda"),
     )
     window = state.prepare_inputs(input_batch, req_states)["lookback_token_ids"]
-    assert window.cpu().tolist() == [
+    expected = [
         [22, 21, 16],
         [-1, -1, -1],
         [0, 0, 0],
-        [-1, -1, -1],
     ]
+    if pcp:
+        # The second PCP chunk of request 2 starts at 5, not at its global 8.
+        expected.append([15, 14, 13])
+    expected.extend([[-1] * depth] * (max_num_reqs - len(expected)))
+    assert window.cpu().tolist() == expected
     # Graph capture runs on the dummy inputs and replays read the same buffer.
     dummy = state.prepare_dummy_inputs(num_reqs=3, num_tokens=3)["lookback_token_ids"]
     assert dummy is window and torch.all(dummy == -1)

tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py

@@ -525,9 +525,8 @@ def test_kv_path_mxfp8_record_matches_reference(num_tokens: int, block_size: int
     step = 32.0 * tile_scale.repeat_interleave(32, dim=-1)
     assert ((rec_fused - rec_ref).abs() <= step).all()
     # The NoPE dims see no rotation at all, so they must be byte-identical.
-    nope_bytes = slice(0, block_size * HEAD_DIM)
-    ref_rows = k_cache_ref[:, nope_bytes].view(num_blocks, block_size, HEAD_DIM)
-    fused_rows = k_cache_fused[:, nope_bytes].view(num_blocks, block_size, HEAD_DIM)
+    ref_rows = k_cache_ref.view(num_blocks, block_size, V41_HEAD_BYTES)
+    fused_rows = k_cache_fused.view(num_blocks, block_size, V41_HEAD_BYTES)
     torch.testing.assert_close(
         fused_rows[..., :NOPE_DIM], ref_rows[..., :NOPE_DIM], rtol=0, atol=0
     )

tests/kernels/test_gate_linear_rocm_dispatch.py

@@ -109,7 +109,7 @@ def test_rocm_set_out_dtype_respects_bias_guard(monkeypatch):
 
 @pytest.mark.parametrize(
     ("input_size", "output_size"),
-    [(3072, 256), (4096, 8), (4096, 192), (6144, 128), (6144, 256)],
+    [(3072, 256), (4096, 8), (6144, 128), (2048, 64), (2560, 512), (7168, 384)],
 )
 def test_rocm_gfx950_enables_fp32_router_gemm(
     monkeypatch, input_size: int, output_size: int
@@ -147,7 +147,6 @@ def test_rocm_fp32_router_gemm_is_replicated_across_ep_sizes(
     ("input_size", "output_size", "params_dtype", "bias", "on_gfx950"),
     [
         pytest.param(6144, 128, torch.float32, False, False, id="gfx942"),
-        pytest.param(2048, 64, torch.float32, False, True, id="shape"),
         pytest.param(6144, 128, torch.bfloat16, False, True, id="bf16-weight"),
         pytest.param(6144, 128, torch.float32, True, True, id="bias"),
     ],

tests/kernels/test_rocm_fp32_router_gemm.py

@@ -5,6 +5,9 @@
 import pytest
 import torch
 
+from vllm.model_executor.layers.fused_moe.router import (
+    rocm_fp32_router_gemm as rocm_router,
+)
 from vllm.model_executor.layers.fused_moe.router.gate_linear import (
     fp32_router_gemm_dispatch_impl,
 )
@@ -19,6 +22,11 @@ SHAPES = [
     (4096, 192),
     (6144, 128),
     (6144, 256),
+    (2560, 512),
+    (2880, 128),
+    (5120, 16),
+    (7168, 384),
+    (1000, 100),
 ]
 MAX_TOKENS = 32
 ATOL = 5e-4
@@ -80,7 +88,7 @@ def test_rocm_fp32_router_gemm_preserves_topk(
     x = torch.randn(num_tokens, hidden_size, dtype=torch.bfloat16, device=device)
     weight = torch.randn(num_experts, hidden_size, dtype=torch.float32, device=device)
 
-    top_k = 2 if (hidden_size, num_experts) == (4096, 8) else 8
+    top_k = 2 if num_experts <= 16 else 8
     actual = rocm_fp32_router_gemm(x, weight).topk(top_k, dim=-1).indices
     reference = _reference(x, weight)
     reference_values, expected = reference.topk(top_k, dim=-1)
@@ -113,15 +121,22 @@ def test_rocm_fp32_router_gemm_rejects_invalid_inputs() -> None:
             weight,
         )
 
-    with pytest.raises(ValueError, match="shape"):
-        rocm_fp32_router_gemm(x[:, :2048].contiguous(), weight[:, :2048].contiguous())
+    with pytest.raises(ValueError, match="hidden sizes"):
+        rocm_fp32_router_gemm(x, weight[:, :2048].contiguous())
+
+    with pytest.raises(ValueError, match="num_experts"):
+        big_weight = torch.empty(1024, 7168, dtype=torch.float32, device=device)
+        rocm_fp32_router_gemm(
+            torch.empty(MAX_TOKENS, 7168, dtype=torch.bfloat16, device=device),
+            big_weight,
+        )
 
     with pytest.raises(ValueError, match="contiguous"):
         wide_x = torch.randn(4, hidden_size * 2, dtype=torch.bfloat16, device=device)
         rocm_fp32_router_gemm(wide_x[:, ::2], weight)
 
-    with pytest.raises(ValueError, match="float32"):
-        rocm_fp32_router_gemm(x, weight.to(torch.bfloat16))
+    with pytest.raises(ValueError, match="float32 or bfloat16"):
+        rocm_fp32_router_gemm(x, weight.to(torch.float16))
 
     with pytest.raises(ValueError, match="bfloat16 or float32"):
         rocm_fp32_router_gemm(x.to(torch.float16), weight)
@@ -226,3 +241,154 @@ def test_rocm_fp32_router_gemm_dynamic_compile_dispatch() -> None:
         )
         output = compiled_dispatch(x)
         torch.testing.assert_close(output, _reference(x, weight), atol=ATOL, rtol=RTOL)
+
+
+@pytest.mark.parametrize("hidden_size", [64, 100, 513, 1000, 2049, 3000])
+@pytest.mark.parametrize("num_tokens", [1, 5, 32])
+@torch.inference_mode()
+def test_rocm_fp32_router_gemm_any_hidden_size(num_tokens: int, hidden_size: int):
+    """Hidden sizes that do not divide the K tile exercise the masked tail."""
+    torch.manual_seed(6100 + num_tokens + hidden_size)
+    x = torch.randn(num_tokens, hidden_size, dtype=torch.bfloat16, device="cuda")
+    weight = torch.randn(64, hidden_size, dtype=torch.float32, device="cuda")
+    torch.testing.assert_close(
+        rocm_fp32_router_gemm(x, weight), _reference(x, weight), atol=ATOL, rtol=RTOL
+    )
+
+
+@pytest.mark.parametrize("case", ["large", "near_zero", "zero_input", "ones"])
+@torch.inference_mode()
+def test_rocm_fp32_router_gemm_special_values(case: str) -> None:
+    torch.manual_seed(7100)
+    x = torch.randn(4, 5120, dtype=torch.bfloat16, device="cuda")
+    weight = torch.randn(128, 5120, dtype=torch.float32, device="cuda")
+    if case in ("large", "near_zero"):
+        scale = 100.0 if case == "large" else 1e-4
+        x = (x * scale).to(torch.bfloat16)
+        weight *= scale
+    elif case == "zero_input":
+        x.zero_()
+    else:
+        x.fill_(1.0)
+    output = rocm_fp32_router_gemm(x, weight)
+    expected = _reference(x, weight)
+    assert torch.isfinite(output).all()
+    if case == "zero_input":
+        assert torch.equal(output, torch.zeros_like(output))
+    else:
+        assert ((output - expected).norm() / expected.norm()).item() < 1e-5
+
+
+@pytest.mark.parametrize("num_tokens", [1, 4, 16, 32])
+@torch.inference_mode()
+def test_rocm_fp32_router_gemm_deterministic(num_tokens: int) -> None:
+    torch.manual_seed(8100 + num_tokens)
+    x = torch.randn(num_tokens, 7168, dtype=torch.bfloat16, device="cuda")
+    weight = torch.randn(384, 7168, dtype=torch.float32, device="cuda")
+    first = rocm_fp32_router_gemm(x, weight)
+    assert torch.equal(first, rocm_fp32_router_gemm(x, weight))
+
+
+@pytest.mark.parametrize("num_tokens", [1, 4, 16, 32])
+@pytest.mark.parametrize(("hidden_size", "num_experts"), SHAPES)
+@torch.inference_mode()
+def test_rocm_fp32_router_gemm_bf16_weight(
+    num_tokens: int, hidden_size: int, num_experts: int
+) -> None:
+    torch.manual_seed(9100 + num_tokens + hidden_size + num_experts)
+    x = torch.randn(num_tokens, hidden_size, dtype=torch.bfloat16, device="cuda")
+    weight = torch.randn(num_experts, hidden_size, dtype=torch.bfloat16, device="cuda")
+    # Past this size hipBLASLt beats the kernel for BF16 weights on MI355.
+    eligible = num_tokens * num_experts <= 8192
+    assert rocm_router.can_use_rocm_fp32_router_gemm(x, weight) == eligible
+    if eligible:
+        torch.testing.assert_close(
+            rocm_fp32_router_gemm(x, weight),
+            _reference(x.float(), weight.float()),
+            atol=ATOL,
+            rtol=RTOL,
+        )
+
+
+def _make_gate_linear(
+    monkeypatch: pytest.MonkeyPatch,
+    input_size,
+    output_size,
+    params_dtype=torch.float32,
+):
+    for module in ("layers.linear", "parameter"):
+        monkeypatch.setattr(
+            f"vllm.model_executor.{module}.get_tensor_model_parallel_rank", lambda: 0
+        )
+        monkeypatch.setattr(
+            f"vllm.model_executor.{module}.get_tensor_model_parallel_world_size",
+            lambda: 1,
+        )
+    from vllm.model_executor.layers.fused_moe.router.gate_linear import GateLinear
+
+    gate = GateLinear(
+        input_size=input_size,
+        output_size=output_size,
+        bias=False,
+        out_dtype=torch.float32,
+        params_dtype=params_dtype,
+    ).cuda()
+    with torch.no_grad():
+        gate.weight.normal_()
+    gate.quant_method.process_weights_after_loading(gate)
+    return gate
+
+
+@pytest.mark.parametrize("x_dtype", [torch.bfloat16, torch.float32])
+@pytest.mark.parametrize("num_tokens", [1, 32, 33])
+@torch.inference_mode()
+def test_gate_linear_uses_rocm_fp32_router_gemm(
+    monkeypatch: pytest.MonkeyPatch, num_tokens: int, x_dtype: torch.dtype
+) -> None:
+    """A router shape outside the old tuned list takes the fast path for M<=32."""
+    gate = _make_gate_linear(monkeypatch, 2560, 512)
+    assert gate.allow_fp32_router_gemm
+    calls = []
+    kernel = rocm_router.rocm_fp32_router_gemm
+
+    def spy(hidden_states, router_weight):
+        calls.append(hidden_states.shape[0])
+        return kernel(hidden_states, router_weight)
+
+    monkeypatch.setattr(rocm_router, "rocm_fp32_router_gemm", spy)
+    x = torch.randn(num_tokens, 2560, dtype=x_dtype, device="cuda")
+    output, bias = gate(x)
+
+    assert bias is None
+    assert calls == ([num_tokens] if num_tokens <= MAX_TOKENS else [])
+    torch.testing.assert_close(output, _reference(x, gate.weight), atol=ATOL, rtol=RTOL)
+
+
+@pytest.mark.parametrize(
+    ("num_tokens", "num_experts", "uses_kernel"),
+    [(1, 256, True), (32, 256, True), (33, 256, False), (32, 512, False)],
+)
+@torch.inference_mode()
+def test_gate_linear_bf16_weight_uses_rocm_router_gemm(
+    monkeypatch: pytest.MonkeyPatch, num_tokens, num_experts, uses_kernel
+) -> None:
+    """BF16-weight gates with FP32 logits use the kernel only where it beats
+    hipBLASLt; everything else keeps the torch.mm path."""
+    gate = _make_gate_linear(monkeypatch, 7168, num_experts, torch.bfloat16)
+    assert gate.allow_cublas_router_gemm
+    calls = []
+    kernel = rocm_router.rocm_fp32_router_gemm
+
+    def spy(hidden_states, router_weight):
+        calls.append(hidden_states.shape[0])
+        return kernel(hidden_states, router_weight)
+
+    monkeypatch.setattr(rocm_router, "rocm_fp32_router_gemm", spy)
+    x = torch.randn(num_tokens, 7168, dtype=torch.bfloat16, device="cuda")
+    output, _ = gate(x)
+
+    assert output.dtype == torch.float32
+    assert calls == ([num_tokens] if uses_kernel else [])
+    torch.testing.assert_close(
+        output, _reference(x.float(), gate.weight.float()), atol=ATOL, rtol=RTOL
+    )

tests/model_executor/layers/test_mla_short_prefill_indexer.py

@@ -830,15 +830,19 @@ def test_sparse_mqa_logits_prefill_matches_dense_masked(capture: bool):
     _assert_same_selection(dense, out, ref)
 
 
-def _build_fp4_paged_cache(req_lens: torch.Tensor, page_kv: int, head_dim: int):
+def _build_fp4_paged_cache(
+    req_lens: torch.Tensor, page_kv: int, head_dim: int, page_pad: int = 0
+):
     """MXFP4 paged cache in DeepGEMM's fused layout: per page, all tokens'
-    fp4 data first, then their packed UE8M0 scales."""
+    fp4 data first, then their packed UE8M0 scales; ``page_pad`` bytes pad
+    each page like the other layers' pages in a packed block."""
     device = req_lens.device
     pages_per = (req_lens + page_kv - 1) // page_kv
     num_pages, max_pages = int(pages_per.sum()), int(pages_per.max())
+    head_bytes = head_dim // 2 + 4
     kv_cache = torch.zeros(
-        num_pages, page_kv, 1, head_dim // 2 + 4, dtype=torch.uint8, device=device
-    )
+        num_pages, page_kv * head_bytes + page_pad, dtype=torch.uint8, device=device
+    )[:, : page_kv * head_bytes].view(num_pages, page_kv, 1, head_bytes)
     block_table = torch.zeros(
         len(req_lens), max_pages, dtype=torch.int32, device=device
     )
@@ -864,7 +868,10 @@ def _build_fp4_paged_cache(req_lens: torch.Tensor, page_kv: int, head_dim: int):
 
 @pytest.mark.parametrize("next_n", [1, 2])
 @pytest.mark.parametrize("capture", [False, True])
-def test_sparse_mqa_logits_paged_matches_dense_masked(next_n: int, capture: bool):
+@pytest.mark.parametrize("page_pad", [0, 256])
+def test_sparse_mqa_logits_paged_matches_dense_masked(
+    next_n: int, capture: bool, page_pad: int
+):
     """Paged sparse-MQA decode selects exactly the values ``torch.topk`` picks
     from the dense bf16 masked logits, incl. MTP rows (next_n > 1 flattened to
     one query row each), CUDA graph capture/replay and schedule reuse by a
@@ -886,8 +893,10 @@ def test_sparse_mqa_logits_paged_matches_dense_masked(next_n: int, capture: bool
         .clamp(min=0)
         .to(torch.int32)
     )
-    kv_cache, block_table = _build_fp4_paged_cache(req_lens, page_kv, head_dim)
-    assert kv_cache.stride(0) % 512 == 0  # the paged sparse kernel requires it
+    kv_cache, block_table = _build_fp4_paged_cache(
+        req_lens, page_kv, head_dim, page_pad
+    )
+    assert kv_cache.stride(0) % 512 == page_pad  # padded pages are not 512B
 
     q_fp, q_sf = _quant_fp4(
         torch.randn(rows * num_heads, head_dim, device="cuda", dtype=torch.bfloat16)

tests/model_executor/layers/test_rocm_unquantized_gemm.py

@@ -17,6 +17,15 @@ if current_platform.is_cuda():
 from vllm.model_executor.layers import utils
 
 
+def _on_gfx11() -> bool:
+    """gfx11 check that is safe to evaluate at collection time on any platform."""
+    if not current_platform.is_rocm():
+        return False
+    from vllm.platforms.rocm import on_gfx11
+
+    return on_gfx11()
+
+
 @pytest.mark.parametrize("m", [1, 128])
 def test_rocm_unquantized_gemm_gfx1x_wvsplitk_path(monkeypatch, m):
     x = torch.randn(1, 64, dtype=torch.float16)
@@ -145,6 +154,101 @@ def test_rocm_unquantized_gemm_rejects_unsupported_skinny_layouts(
     assert torch.allclose(out, ref, atol=1e-3, rtol=1e-3)
 
 
+def _row_stride_padded(weight: torch.Tensor, pad: int) -> torch.Tensor:
+    """A weight whose row stride is offset by ``pad`` elements, as produced by
+    ``maybe_pad_weight_avoid_cache_aliasing``."""
+    padded = torch.nn.functional.pad(weight, (0, pad))[..., :-pad]
+    assert padded.shape == weight.shape
+    assert padded.stride() == (weight.shape[1] + pad, 1)
+    return padded
+
+
+def test_rocm_unquantized_gemm_wvsplitk_accepts_row_stride_padded_weight(monkeypatch):
+    # wvSplitK reads the weight row stride, so a row-stride-padded weight must
+    # still reach it instead of falling back to torch.nn.functional.linear.
+    x = torch.randn(4, 1024, dtype=torch.float16)
+    weight = _row_stride_padded(torch.randn(128, 1024, dtype=torch.float16), 64)
+
+    monkeypatch.setattr(utils, "use_aiter_triton_gemm", lambda *args: False)
+    monkeypatch.setattr(utils.rocm_aiter_ops, "is_tgemm_enabled", lambda: False)
+    monkeypatch.setattr(utils.envs, "VLLM_ROCM_USE_SKINNY_GEMM", True)
+    monkeypatch.setattr("vllm.platforms.rocm.on_gfx11", lambda: True)
+    monkeypatch.setattr("vllm.platforms.rocm.on_gfx1x", lambda: True)
+    monkeypatch.setattr("vllm.platforms.rocm.on_gfx9", lambda: False)
+    monkeypatch.setattr("vllm.platforms.rocm.on_gfx950", lambda: False)
+    monkeypatch.setattr("vllm.platforms.rocm.on_gfx1250", lambda: False)
+    monkeypatch.setattr(utils, "num_compute_units", lambda: 120)
+
+    wvsplitk_mock = MagicMock(side_effect=lambda w, x_view, _, __: x_view @ w.t())
+    monkeypatch.setattr(utils.ops, "wvSplitK", wvsplitk_mock)
+
+    out = utils.rocm_unquantized_gemm_impl(x, weight, None)
+    ref = torch.nn.functional.linear(x, weight, None)
+
+    wvsplitk_mock.assert_called_once()
+    assert wvsplitk_mock.call_args.args[0].stride() == weight.stride()
+    assert torch.allclose(out, ref, atol=1e-3, rtol=1e-3)
+
+
+@pytest.mark.parametrize(
+    "k,element_size,expected",
+    [
+        (4096, 2, 64),  # 8192 B row stride, aliases
+        (12288, 2, 64),  # 24576 B row stride, aliases
+        (1024, 2, 64),  # smallest aliasing bf16/fp16 K
+        (2560, 2, 0),  # 5120 B, no aliasing
+        (11008, 2, 0),  # 22016 B, no aliasing
+        (2048, 1, 128),  # fp8: 2048 B, aliases, twice the pad
+        (512, 4, 32),  # fp32: 2048 B, aliases, half the pad
+    ],
+)
+def test_cache_aliasing_pad_elems(k, element_size, expected):
+    assert utils.cache_aliasing_pad_elems(k, element_size) == expected
+
+
+def test_maybe_pad_weight_avoid_cache_aliasing_skips_non_gfx11(monkeypatch):
+    # Padding is gfx11-only, even for a K whose row stride would alias there.
+    monkeypatch.setattr(utils.current_platform, "is_rocm", lambda: True)
+    monkeypatch.setattr("vllm.platforms.rocm.on_gfx11", lambda: False)
+
+    weight = torch.randn(128, 4096, dtype=torch.bfloat16)
+    assert utils.maybe_pad_weight_avoid_cache_aliasing(weight) is weight
+
+
+def test_maybe_pad_weight_avoid_cache_aliasing(monkeypatch):
+    monkeypatch.setattr(utils.current_platform, "is_rocm", lambda: True)
+    monkeypatch.setattr("vllm.platforms.rocm.on_gfx11", lambda: True)
+
+    weight = torch.randn(128, 4096, dtype=torch.bfloat16)
+    padded = utils.maybe_pad_weight_avoid_cache_aliasing(weight)
+
+    assert padded.shape == weight.shape
+    assert padded.stride() == (4096 + 64, 1)
+    assert torch.equal(padded, weight)
+    # Idempotent: the padded stride no longer aliases, so a second call is a no-op
+    # and preserves the pointer captured by any graph.
+    repadded = utils.maybe_pad_weight_avoid_cache_aliasing(padded)
+    assert repadded.data_ptr() == padded.data_ptr()
+
+
+@pytest.mark.parametrize(
+    "weight",
+    [
+        torch.randn(128, 2560, dtype=torch.bfloat16),  # no aliasing
+        torch.randn(2560, 128, dtype=torch.bfloat16).t(),  # not contiguous
+        torch.randn(2, 128, 4096, dtype=torch.bfloat16),  # not 2D
+    ],
+)
+def test_maybe_pad_weight_avoid_cache_aliasing_skips(monkeypatch, weight):
+    monkeypatch.setattr(utils.current_platform, "is_rocm", lambda: True)
+    monkeypatch.setattr("vllm.platforms.rocm.on_gfx11", lambda: True)
+
+    assert (
+        utils.maybe_pad_weight_avoid_cache_aliasing(weight).data_ptr()
+        == weight.data_ptr()
+    )
+
+
 @pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm-only kernel test")
 @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
 def test_rocm_unquantized_gemm_noncontiguous_activation_real_kernel(monkeypatch, dtype):
@@ -164,6 +268,26 @@ def test_rocm_unquantized_gemm_noncontiguous_activation_real_kernel(monkeypatch,
     torch.testing.assert_close(out, ref, atol=1e-2, rtol=1e-2)
 
 
+@pytest.mark.skipif(
+    not _on_gfx11(), reason="cache-aliasing padding only applies on gfx11"
+)
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
+def test_rocm_unquantized_gemm_padded_weight_real_kernel(monkeypatch, dtype):
+    x = torch.randn(4, 4096, device="cuda", dtype=dtype)
+    weight = _row_stride_padded(torch.randn(128, 4096, device="cuda", dtype=dtype), 64)
+
+    monkeypatch.setattr(utils.envs, "VLLM_ROCM_USE_SKINNY_GEMM", True)
+    original_wvsplitk = utils.ops.wvSplitK
+    wvsplitk_mock = MagicMock(side_effect=original_wvsplitk)
+    monkeypatch.setattr(utils.ops, "wvSplitK", wvsplitk_mock)
+
+    out = utils.rocm_unquantized_gemm_impl(x, weight, None)
+    ref = torch.nn.functional.linear(x, weight, None)
+
+    wvsplitk_mock.assert_called_once()
+    torch.testing.assert_close(out, ref, atol=1e-2, rtol=1e-2)
+
+
 @pytest.mark.parametrize("m", [1, 128])
 def test_rocm_unquantized_gemm_gfx1x_n_gt_5_falls_back(monkeypatch, m):
     # wvSplitK skinny GEMM handles n in [1, 5] (see PR #40687); n > 5 must

tests/model_executor/test_deep_gemm_warmup.py

@@ -123,19 +123,12 @@ def test_mega_moe_m_values_reach_mega_gate(monkeypatch) -> None:
     threshold (1 for 128 experts), so the forward actually reaches the kernel."""
     dg = SimpleNamespace(
         get_bf16_mega_gate_config=lambda m, *_: {"num_sms": -(-m // 16)},
-        get_block_m_for_mega_moe=lambda *_: 16,
     )
     monkeypatch.setattr(deep_gemm_warmup, "_import_deep_gemm", lambda: dg)
     monkeypatch.setattr(deep_gemm_warmup, "_mega_moe_uses_mega_gate", lambda _: True)
-    monkeypatch.setattr(
-        deep_gemm_warmup, "get_ep_group", lambda: SimpleNamespace(world_size=8)
-    )
-    buffer = SimpleNamespace(num_max_tokens_per_rank=64)
     module = SimpleNamespace(
         gate=SimpleNamespace(input_size=4096, output_size=128),
-        experts=SimpleNamespace(
-            top_k=6, num_experts=128, get_symm_buffer=lambda: buffer
-        ),
+        experts=SimpleNamespace(top_k=6, get_block_m=lambda num_tokens: 16),
         use_sequence_parallel=False,
     )
     assert deep_gemm_warmup._get_fp8_fp4_mega_moe_m_values(module, 40) == [16, 32, 40]

tests/model_executor/test_enabled_custom_ops.py

@@ -139,6 +139,122 @@ def test_topk_softmax_dispatch(use_rocm_aiter: bool):
         assert topk_func == vllm_topk_softmax
 
 
+def _topk_gating_launch(
+    num_tokens: int = 4,
+    num_experts: int = 512,
+    topk: int = 10,
+    num_shared_experts: int = 0,
+    weights_width: int | None = None,
+    ids_width: int | None = None,
+    contiguous: bool = True,
+) -> dict:
+    """Tensors for one softmax top-k launch, as dispatch_topk_softmax_func sees them."""
+    width = topk + num_shared_experts
+    weights = torch.empty(num_tokens, weights_width or width)
+    ids = torch.empty(num_tokens, ids_width or width, dtype=torch.int32)
+    ids = ids[:, :topk]
+    if contiguous:
+        gating = torch.empty(num_tokens, num_experts + num_shared_experts)
+    else:
+        gating = torch.empty(num_experts + num_shared_experts + 3, num_tokens).t()
+    return dict(topk_weights=weights, topk_indices=ids, gating_output=gating)
+
+
+@pytest.mark.parametrize(
+    "launch, num_shared_experts, scoring_func, expect_gating",
+    [
+        pytest.param({}, 0, "", True, id="plain-softmax"),
+        pytest.param({"num_tokens": 4096}, 0, "", True, id="max-tokens"),
+        pytest.param({"num_tokens": 4097}, 0, "", False, id="too-many-tokens"),
+        pytest.param({"contiguous": False}, 0, "", False, id="non-contiguous"),
+        pytest.param({}, 0, "sigmoid", False, id="scoring-without-shared"),
+        pytest.param({"num_shared_experts": 1}, 1, "sigmoid", True, id="one-shared"),
+        pytest.param({"num_shared_experts": 8}, 8, "sigmoid", True, id="eight-shared"),
+        pytest.param(
+            {"num_shared_experts": 3}, 3, "sigmoid", False, id="unsupported-count"
+        ),
+        pytest.param({"num_shared_experts": 1}, 1, "", False, id="no-shared-scoring"),
+        pytest.param(
+            {"num_shared_experts": 1}, 1, "softmax", False, id="wrong-shared-scoring"
+        ),
+        pytest.param(
+            {"num_shared_experts": 1, "weights_width": 10},
+            1,
+            "sigmoid",
+            False,
+            id="weights-too-narrow",
+        ),
+        pytest.param(
+            {"num_shared_experts": 1, "ids_width": 12},
+            1,
+            "sigmoid",
+            False,
+            id="row-stride-mismatch",
+        ),
+        pytest.param(
+            {"num_tokens": 4097, "num_shared_experts": 1},
+            1,
+            "sigmoid",
+            False,
+            id="shared-too-many-tokens",
+        ),
+    ],
+)
+@pytest.mark.parametrize("gating_enabled", [True, False])
+def test_topk_softmax_dispatch_aiter_topk_gating(
+    monkeypatch: pytest.MonkeyPatch,
+    launch: dict,
+    num_shared_experts: int,
+    scoring_func: str,
+    expect_gating: bool,
+    gating_enabled: bool,
+):
+    """AITER launches use topk_gating only where it is enabled, and only when it
+    supports them.
+
+    Every other AITER launch must reach the legacy topk_softmax, so this is the
+    one place the choice is made.
+    """
+    monkeypatch.setattr(
+        rocm_aiter_ops, "is_topk_gating_enabled", lambda: gating_enabled
+    )
+    topk_func = dispatch_topk_softmax_func(
+        True,
+        **_topk_gating_launch(**launch),
+        num_shared_experts=num_shared_experts,
+        shared_expert_scoring_func=scoring_func,
+    )
+    if expect_gating and gating_enabled:
+        assert topk_func == rocm_aiter_ops.topk_gating
+    else:
+        assert topk_func == rocm_aiter_ops.topk_softmax
+
+
+@pytest.mark.parametrize(
+    "cdna, fused_moe_enabled, expected",
+    [
+        pytest.param(3, True, True, id="cdna3"),
+        pytest.param(4, True, True, id="cdna4"),
+        pytest.param(2, True, False, id="cdna2"),
+        pytest.param(3, False, False, id="cdna3-aiter-moe-off"),
+        pytest.param(4, False, False, id="cdna4-aiter-moe-off"),
+    ],
+)
+def test_topk_gating_enabled_by_cdna(
+    monkeypatch: pytest.MonkeyPatch,
+    cdna: int,
+    fused_moe_enabled: bool,
+    expected: bool,
+):
+    """topk_gating is on for CDNA3 and CDNA4 with AITER MoE enabled, nowhere else."""
+    monkeypatch.setattr("vllm._aiter_ops.is_aiter_found_and_supported", lambda: True)
+    monkeypatch.setattr(
+        rocm_aiter_ops, "is_fused_moe_enabled", lambda: fused_moe_enabled
+    )
+    monkeypatch.setattr("vllm.platforms.rocm.get_cdna_version", lambda: cdna)
+    assert bool(rocm_aiter_ops.is_topk_gating_enabled()) is expected
+
+
 @pytest.mark.parametrize(
     "use_rocm_aiter", [True, False] if current_platform.is_rocm() else [False]
 )

tests/model_executor/test_routed_experts_capture.py

@@ -7,9 +7,16 @@ from unittest.mock import Mock, patch
 import pytest
 import torch
 
-from vllm.config import VllmConfig
+from vllm.config import CUDAGraphMode, VllmConfig
 from vllm.config.compilation import CompilationMode
 from vllm.distributed.eplb.eplb_state import EplbLayerState
+from vllm.forward_context import (
+    BatchDescriptor,
+    DPMetadata,
+    ForwardContext,
+    get_forward_context,
+    set_forward_context,
+)
 from vllm.model_executor.layers.fused_moe.config import RoutingMethodType
 from vllm.model_executor.layers.fused_moe.experts.cpu_int4_moe import CPUExpertsInt4
 from vllm.model_executor.layers.fused_moe.modular_kernel import (
@@ -36,6 +43,12 @@ pytestmark = pytest.mark.cpu_test
 _REC_MODULE = "vllm.model_executor.layers.fused_moe.routed_experts_capturer"
 
 
+def _forward_context(**kwargs):
+    return ForwardContext(
+        no_compile_layers={}, attn_metadata={}, slot_mapping={}, **kwargs
+    )
+
+
 def _capturer_with_buffer(
     *,
     max_tokens: int = 8,
@@ -319,13 +332,120 @@ def test_routed_experts_capturer_single_dp_no_metadata():
     """dp_metadata is None: capture writes the full topk_ids rows."""
     capturer = _capturer_with_buffer(dp_rank=0)
     topk = torch.tensor([[1, 2], [3, 4], [5, 6]], dtype=torch.int32)
-    ctx = SimpleNamespace(dp_metadata=None)
+    ctx = _forward_context()
     with patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx):
         capturer.capture(layer_id=0, topk_ids=topk)
     assert torch.equal(capturer.device_buffer[:3, 0, :], topk)
     assert capturer.device_buffer[3, 0, 0].item() == -1
 
 
+@pytest.mark.parametrize("tp_size, num_tokens", [(2, 6), (2, 3), (2, 1), (4, 2)])
+@pytest.mark.parametrize(
+    "cg_mode", [CUDAGraphMode.NONE, CUDAGraphMode.FULL, CUDAGraphMode.PIECEWISE]
+)
+def test_single_dp_ep_capture_restores_tp_shards(tp_size, num_tokens, cg_mode):
+    """Capture all current rows after model-level SP, including buffer reuse."""
+    capturer = _capturer_with_buffer(tp_size=tp_size)
+    config = SimpleNamespace(
+        parallel_config=SimpleNamespace(
+            data_parallel_size=1,
+            tensor_parallel_size=tp_size,
+            enable_expert_parallel=True,
+            use_sequence_parallel_moe=False,
+            is_moe_model=True,
+        ),
+        compilation_config=SimpleNamespace(
+            fast_moe_cold_start=False, static_forward_context={}
+        ),
+    )
+    tp_group = Mock()
+    shard_n = (num_tokens + tp_size - 1) // tp_size
+    with (
+        set_forward_context(
+            None,
+            config,
+            num_tokens=num_tokens,
+            cudagraph_runtime_mode=cg_mode,
+        ),
+        patch(f"{_REC_MODULE}.get_tp_group", return_value=tp_group),
+    ):
+        assert get_forward_context().dp_metadata is None
+        for offset in (20, 1):
+            expected = torch.arange(num_tokens * 2, dtype=torch.int32).reshape(-1, 2)
+            expected += offset
+            padded = torch.full((shard_n * tp_size, 2), 99, dtype=torch.int32)
+            padded[:num_tokens] = expected
+            tp_group.all_gather.return_value = padded
+            capturer.capture(0, padded[:shard_n])
+            assert torch.equal(
+                capturer.snapshot_routing_data(num_tokens)[:, 0].int(), expected
+            )
+            assert (capturer.device_buffer[num_tokens:, 0] == -1).all()
+    if num_tokens == 1:
+        tp_group.all_gather.assert_not_called()
+    else:
+        assert tp_group.all_gather.call_count == 2
+
+
+@pytest.mark.parametrize("num_tokens, num_rows", [(6, 6), (3, 4), (0, 0)])
+def test_single_dp_tp_capture_does_not_regather_full_batch(num_tokens, num_rows):
+    capturer = _capturer_with_buffer(tp_size=2)
+    topk = torch.arange(num_rows * 2, dtype=torch.int32).reshape(-1, 2)
+    ctx = _forward_context(batch_descriptor=BatchDescriptor(num_tokens=num_tokens))
+    with (
+        patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx),
+        patch(f"{_REC_MODULE}.get_tp_group") as get_tp_group,
+    ):
+        capturer.capture(0, topk)
+    get_tp_group.assert_not_called()
+    assert torch.equal(capturer.device_buffer[:num_tokens, 0], topk[:num_tokens])
+    assert (capturer.device_buffer[num_tokens:, 0] == -1).all()
+
+
+@pytest.mark.parametrize("num_tokens, num_rows", [(6, 2), (0, 1)])
+def test_single_dp_tp_capture_rejects_unexpected_batch_without_writing(
+    num_tokens, num_rows
+):
+    capturer = _capturer_with_buffer(tp_size=2)
+    topk = torch.ones(num_rows, 2, dtype=torch.int32)
+    ctx = _forward_context(batch_descriptor=BatchDescriptor(num_tokens=num_tokens))
+    with (
+        patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx),
+        pytest.raises(AssertionError, match="unexpected topk_ids batch dim"),
+    ):
+        capturer.capture(0, topk)
+    assert (capturer.device_buffer == -1).all()
+
+
+@pytest.mark.parametrize("tp_rank", [0, 1])
+def test_capture_uses_pre_sp_size_after_model_shards_padding_mask(tp_rank):
+    """Model-level SP shortens the mask; capture still needs all three rows."""
+    from vllm.models.common.ops import sequence_parallel as sp
+
+    capturer = _capturer_with_buffer(tp_size=2)
+    expected = torch.arange(8, dtype=torch.int32).reshape(4, 2)
+    ctx = _forward_context(
+        is_padding=torch.zeros(3, dtype=torch.bool),
+        batch_descriptor=BatchDescriptor(num_tokens=3),
+    )
+    with (
+        patch.object(sp, "get_tensor_model_parallel_world_size", return_value=2),
+        patch.object(sp, "get_tensor_model_parallel_rank", return_value=tp_rank),
+    ):
+        ctx.is_padding = sp.sp_padding_mask(ctx.is_padding, torch.empty(3, 1))
+    assert len(ctx.is_padding) == 2
+    tp_group = Mock()
+    tp_group.all_gather.return_value = expected
+    with (
+        patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx),
+        patch(f"{_REC_MODULE}.get_tp_group", return_value=tp_group),
+    ):
+        capturer.capture(0, expected[2 * tp_rank : 2 * (tp_rank + 1)])
+    tp_group.all_gather.assert_called_once()
+    torch.testing.assert_close(capturer.device_buffer[:3, 0], expected[:3])
+    assert (capturer.device_buffer[3:, 0] == -1).all()
+
+
 @pytest.mark.parametrize(
     ("num_experts", "dtype_name", "torch_dtype"),
     [(256, "uint8", torch.uint8), (257, "uint16", torch.uint16)],
@@ -359,7 +479,7 @@ def test_routed_experts_capturer_narrows_snapshot(output_dtype):
     capturer = _capturer_with_buffer(dtype=torch.int32)
     capturer.output_dtype = output_dtype
     topk = torch.tensor([[1, 2], [254, 255]], dtype=torch.int64)
-    ctx = SimpleNamespace(dp_metadata=None)
+    ctx = _forward_context()
     with patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx):
         capturer.capture(layer_id=0, topk_ids=topk)
 
@@ -375,8 +495,8 @@ def test_routed_experts_capturer_dp_naive_concatenated_all_ranks():
     """N == sum(num_tokens_dp): slice this rank's segment from concatenated topk."""
     capturer = _capturer_with_buffer(dp_rank=1)
     num_tokens_dp = torch.tensor([2, 3], dtype=torch.int32)
-    ctx = SimpleNamespace(
-        dp_metadata=SimpleNamespace(num_tokens_across_dp_cpu=num_tokens_dp)
+    ctx = _forward_context(
+        dp_metadata=DPMetadata(num_tokens_across_dp_cpu=num_tokens_dp)
     )
     # Concatenated order: rank0 rows then rank1 rows.
     topk = torch.tensor(
@@ -392,8 +512,8 @@ def test_routed_experts_capturer_dp_modular_local_tokens():
     """N == token_num_per_dp: topk is already local to this DP rank."""
     capturer = _capturer_with_buffer(dp_rank=1)
     num_tokens_dp = torch.tensor([2, 3], dtype=torch.int32)
-    ctx = SimpleNamespace(
-        dp_metadata=SimpleNamespace(num_tokens_across_dp_cpu=num_tokens_dp)
+    ctx = _forward_context(
+        dp_metadata=DPMetadata(num_tokens_across_dp_cpu=num_tokens_dp)
     )
     topk = torch.tensor([[10, 11], [12, 13], [14, 15]], dtype=torch.int32)
     with patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx):
@@ -404,8 +524,8 @@ def test_routed_experts_capturer_dp_modular_local_tokens():
 def test_routed_experts_capturer_dp_ep_gathered_shards():
     capturer = _capturer_with_buffer(dp_rank=1, tp_size=2)
     num_tokens_dp = torch.tensor([3, 2], dtype=torch.int32)
-    ctx = SimpleNamespace(
-        dp_metadata=SimpleNamespace(
+    ctx = _forward_context(
+        dp_metadata=DPMetadata(
             num_tokens_across_dp_cpu=num_tokens_dp,
             local_sizes=[2, 2, 1, 1],
         )
@@ -421,13 +541,17 @@ def test_routed_experts_capturer_dp_ep_gathered_shards():
     assert torch.equal(capturer.device_buffer[:2, 0], topk[4:])
 
 
-def test_routed_experts_capturer_sp_modular_gathers_tp_shards():
-    capturer = _capturer_with_buffer(dp_rank=1, tp_size=2)
+@pytest.mark.parametrize("single_dp", [False, True])
+def test_routed_experts_capturer_sp_modular_gathers_tp_shards(single_dp):
+    capturer = _capturer_with_buffer(dp_rank=0 if single_dp else 1, tp_size=2)
     num_tokens_dp = torch.tensor([2, 3], dtype=torch.int32)
-    ctx = SimpleNamespace(
-        dp_metadata=SimpleNamespace(num_tokens_across_dp_cpu=num_tokens_dp)
+    ctx = _forward_context(
+        dp_metadata=(
+            None if single_dp else DPMetadata(num_tokens_across_dp_cpu=num_tokens_dp)
+        ),
+        batch_descriptor=BatchDescriptor(num_tokens=3),
     )
-    local_shard = torch.tensor([[10, 11], [12, 13]], dtype=torch.int32)
+    local_shard = torch.tensor([[10, 12], [11, 13]], dtype=torch.int32).t()
     gathered = torch.tensor([[10, 11], [12, 13], [14, 15], [-1, -1]], dtype=torch.int32)
     tp_group = Mock()
     tp_group.all_gather.return_value = gathered
@@ -438,16 +562,83 @@ def test_routed_experts_capturer_sp_modular_gathers_tp_shards():
     ):
         capturer.capture(layer_id=0, topk_ids=local_shard)
 
-    tp_group.all_gather.assert_called_once_with(local_shard, dim=0)
+    tp_group.all_gather.assert_called_once()
+    gather_input = tp_group.all_gather.call_args.args[0]
+    assert gather_input.is_contiguous()
+    torch.testing.assert_close(gather_input, local_shard)
+    assert tp_group.all_gather.call_args.kwargs == {"dim": 0}
     assert torch.equal(capturer.device_buffer[:3, 0], gathered[:3])
 
 
+def test_routed_experts_capturer_rejects_incomplete_dp_shard_groups():
+    """Invalid gather metadata must not silently select another rank's rows."""
+    capturer = _capturer_with_buffer(dp_rank=1, tp_size=2)
+    ctx = _forward_context(
+        dp_metadata=DPMetadata(
+            num_tokens_across_dp_cpu=torch.tensor([3, 2]),
+            local_sizes=[2, 2, 2],
+        )
+    )
+    with (
+        patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx),
+        pytest.raises(ValueError, match="not divisible by DP rank count"),
+    ):
+        capturer.capture(0, torch.ones(6, 2, dtype=torch.int32))
+    assert (capturer.device_buffer == -1).all()
+
+
+@pytest.mark.parametrize("layer_id", [-1, 4])
+def test_routed_experts_capturer_rejects_layer_before_restoring_rows(layer_id):
+    capturer = _capturer_with_buffer(tp_size=2)
+    ctx = _forward_context(batch_descriptor=BatchDescriptor(num_tokens=3))
+    with (
+        patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx),
+        patch(f"{_REC_MODULE}.get_tp_group") as get_tp_group,
+        pytest.raises(IndexError, match="exceeds capture buffer"),
+    ):
+        capturer.capture(layer_id, torch.ones(2, 2, dtype=torch.int32))
+    get_tp_group.assert_not_called()
+    assert (capturer.device_buffer == -1).all()
+
+
+@pytest.mark.parametrize(
+    ("counts", "local_sizes", "start", "stop"),
+    [
+        ([3, 0], [2, 2, 0, 0], 3, 3),
+        ([0, 3], None, 0, 3),
+        ([0, 0], [0, 0, 0, 0], 0, 0),
+        ([2, 2], [1, 1, 1, 1], 2, 4),
+    ],
+    ids=["idle-rank", "only-active-rank", "all-idle", "unpadded-gather-tie"],
+)
+def test_routed_experts_capturer_dp_preserves_empty_and_unpadded_spans(
+    counts, local_sizes, start, stop
+):
+    capturer = _capturer_with_buffer(dp_rank=1, tp_size=2)
+    topk = torch.arange(sum(counts) * 2, dtype=torch.int32).reshape(-1, 2)
+    ctx = _forward_context(
+        dp_metadata=DPMetadata(
+            num_tokens_across_dp_cpu=torch.tensor(counts), local_sizes=local_sizes
+        )
+    )
+    with (
+        patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx),
+        patch(f"{_REC_MODULE}.get_tp_group") as get_tp_group,
+    ):
+        capturer.capture(0, topk)
+    get_tp_group.assert_not_called()
+    torch.testing.assert_close(
+        capturer.device_buffer[: stop - start, 0], topk[start:stop]
+    )
+    assert (capturer.device_buffer[stop - start :, 0] == -1).all()
+
+
 def test_routed_experts_capturer_dp_unexpected_batch_raises():
     """Mismatch between topk batch dim and DP layout: fail fast."""
     capturer = _capturer_with_buffer(dp_rank=0)
     num_tokens_dp = torch.tensor([2, 3], dtype=torch.int32)
-    ctx = SimpleNamespace(
-        dp_metadata=SimpleNamespace(num_tokens_across_dp_cpu=num_tokens_dp)
+    ctx = _forward_context(
+        dp_metadata=DPMetadata(num_tokens_across_dp_cpu=num_tokens_dp)
     )
     # total=5, local=2: n=1 matches neither naive (5) nor modular (2).
     topk = torch.tensor([[1, 2]], dtype=torch.int32)

tests/models/kimi_k3/test_amd_kda_decode.py

@@ -1,11 +1,17 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""The ROCm fused KDA decode kernel must match the Triton chain it replaces.
+"""The ROCm fused KDA decode kernels must match the Triton chain they replace.
 
-The fused kernel folds the packed causal conv1d update, the gated delta-rule
-recurrence and the gated output RMSNorm into one launch, and updates both the
-conv state and the recurrent state in place. Every one of those outputs is
-compared against the three-kernel path the AMD layer falls back to.
+Two fused paths fold the packed causal conv1d update, the gated delta-rule
+recurrence and the gated output RMSNorm into one launch, updating the conv
+state and the recurrent state in place:
+
+* the vLLM HIP kernel ``ops.fused_kda_decode`` (gfx942 / gfx950), and
+* AITER's Gluon kernels behind ``rocm_aiter_ops`` (CDNA4+): the packed
+  non-spec decode and the speculative-decode ``fused_recurrent_kda``.
+
+Every output of each is compared against the unfused path the AMD layer falls
+back to, with the argument layout the layer itself uses.
 """
 
 import pytest
@@ -14,7 +20,7 @@ import torch
 from vllm.platforms import current_platform
 
 
-def _on_supported_arch() -> bool:
+def _hip_kernel_arch() -> bool:
     if not current_platform.is_rocm():
         return False
     from vllm.platforms.rocm import on_gfx942, on_gfx950
@@ -22,10 +28,26 @@ def _on_supported_arch() -> bool:
     return on_gfx950() or on_gfx942()
 
 
+def _aiter_gluon_kda_available() -> bool:
+    if not current_platform.is_rocm():
+        return False
+    from vllm._aiter_ops import rocm_aiter_ops
+    from vllm.platforms.rocm import get_cdna_version
+
+    return get_cdna_version() >= 4 and rocm_aiter_ops._kda_gluon_decode_importable()
+
+
 pytestmark = pytest.mark.skipif(
-    not _on_supported_arch(),
+    not current_platform.is_rocm(), reason="ROCm-only KDA decode kernels"
+)
+requires_hip_kernel_arch = pytest.mark.skipif(
+    not _hip_kernel_arch(),
     reason="The fused KDA decode kernel is only built for gfx942 / gfx950",
 )
+requires_aiter_gluon_kda = pytest.mark.skipif(
+    not _aiter_gluon_kda_available(),
+    reason="AITER Gluon KDA decode needs CDNA4+ and an AITER build that ships it",
+)
 
 # Kimi-K3 KDA: 96 heads x 128, conv width 4, gate_lower_bound -5.0.
 HEAD_DIM = 128
@@ -41,7 +63,11 @@ def _requires_kernel() -> None:
 
 
 class KdaDecodeInputs:
-    """One decode step of a Kimi-K3 KDA layer, in the layer's own layout."""
+    """One decode step of a Kimi-K3 KDA layer, in the layer's own layout.
+
+    ``num_spec`` > 0 sizes the conv state for speculative decoding
+    (``conv_width - 1 + num_spec`` taps, as ``kda_state_shape`` does).
+    """
 
     def __init__(
         self,
@@ -49,6 +75,7 @@ class KdaDecodeInputs:
         num_heads: int,
         num_slots: int,
         seed: int = 0,
+        num_spec: int = 0,
     ) -> None:
         torch.manual_seed(seed)
         device = "cuda"
@@ -70,9 +97,15 @@ class KdaDecodeInputs:
                 for i in range(3)
             ]
         )
-        # SD cache layout: [slots, width - 1, 3 * dim]; the layer transposes it.
+        # SD cache layout: [slots, taps, 3 * dim]; the layer transposes it.
         self.conv_state = (
-            torch.randn(num_slots, CONV_WIDTH - 1, 3 * dim, device=device, dtype=DTYPE)
+            torch.randn(
+                num_slots,
+                CONV_WIDTH - 1 + num_spec,
+                3 * dim,
+                device=device,
+                dtype=DTYPE,
+            )
             * 0.5
         )
         self.recurrent_state = (
@@ -112,6 +145,12 @@ class KdaDecodeInputs:
     def conv_state_view(self, state: torch.Tensor) -> torch.Tensor:
         return state.transpose(-1, -2)
 
+    def empty_out(self) -> torch.Tensor:
+        num_tokens = self.mixed_qkv.shape[0]
+        return torch.empty(
+            1, num_tokens, self.num_heads, HEAD_DIM, device="cuda", dtype=DTYPE
+        )
+
 
 def _gated_rmsnorm(
     x: torch.Tensor, gate: torch.Tensor, weight: torch.Tensor, eps: float
@@ -196,6 +235,7 @@ def _run_fused(
     return out, conv_state, recurrent_state
 
 
+@requires_hip_kernel_arch
 @pytest.mark.parametrize("num_heads", [12, 24, 96])
 @pytest.mark.parametrize("num_tokens", [1, 7, 128])
 @torch.inference_mode()
@@ -214,6 +254,7 @@ def test_fused_kda_decode_matches_triton_chain(num_heads: int, num_tokens: int)
     torch.testing.assert_close(actual_state, expected_state, atol=2e-3, rtol=2e-3)
 
 
+@requires_hip_kernel_arch
 @torch.inference_mode()
 def test_fused_kda_decode_leaves_untouched_slots_alone() -> None:
     """Only the slots named by state_indices may be written."""
@@ -234,6 +275,7 @@ def test_fused_kda_decode_leaves_untouched_slots_alone() -> None:
     )
 
 
+@requires_hip_kernel_arch
 @torch.inference_mode()
 def test_fused_kda_decode_without_output_norm() -> None:
     """Omitting the gate/norm pair returns the raw recurrent output."""
@@ -265,6 +307,7 @@ def test_fused_kda_decode_without_output_norm() -> None:
     assert torch.isfinite(out).all()
 
 
+@requires_hip_kernel_arch
 @torch.inference_mode()
 def test_fused_kda_decode_skips_null_block_padding() -> None:
     """A CUDA-graph decode batch is padded with NULL_BLOCK_ID (0).
@@ -291,6 +334,7 @@ def test_fused_kda_decode_skips_null_block_padding() -> None:
     torch.testing.assert_close(actual_state, expected_state, atol=2e-3, rtol=2e-3)
 
 
+@requires_hip_kernel_arch
 def test_decode_weights_staged_after_loading(dist_init) -> None:
     """The fused-decode weight copies must not depend on the weight loaders.
 
@@ -323,11 +367,15 @@ def test_decode_weights_staged_after_loading(dist_init) -> None:
             config, vllm_config, prefix="model.layers.0.self_attn"
         )
     assert layer.decode_conv1d_weight is not None
+    # The AITER Gluon decode path takes o_norm.weight directly, so only the HIP
+    # kernel gets an fp32 norm copy.
+    assert (layer.decode_norm_weight is None) == layer.use_aiter_decode
     with torch.no_grad():
         for param in layer.parameters():
             param.normal_()
         layer.decode_conv1d_weight.fill_(float("nan"))
-        layer.decode_norm_weight.fill_(float("nan"))
+        if layer.decode_norm_weight is not None:
+            layer.decode_norm_weight.fill_(float("nan"))
 
     process_weights_after_loading(
         torch.nn.ModuleList([layer]), vllm_config.model_config, torch.device("cuda")
@@ -341,7 +389,271 @@ def test_decode_weights_staged_after_loading(dist_init) -> None:
     torch.testing.assert_close(
         layer.decode_conv1d_weight, expected_conv, atol=0, rtol=0
     )
+    if layer.decode_norm_weight is not None:
+        torch.testing.assert_close(
+            layer.decode_norm_weight, layer.o_norm.weight.float(), atol=0, rtol=0
+        )
+        assert not layer.decode_norm_weight.requires_grad
+    assert not layer.decode_conv1d_weight.requires_grad
+
+
+# --------------------------------------------------------------------------
+# non-spec decode: fused_recurrent_kda_packed_decode
+# --------------------------------------------------------------------------
+
+
+def _run_aiter_packed_decode(
+    inp: KdaDecodeInputs,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Exactly the call KimiK3DeltaAttention._forward makes for decode batches."""
+    from vllm._aiter_ops import rocm_aiter_ops
+
+    conv_state = inp.conv_state.clone()
+    recurrent_state = inp.recurrent_state.clone()
+    num_tokens = inp.mixed_qkv.shape[0]
+    core_attn_out = inp.empty_out()
+    rocm_aiter_ops.kda_fused_recurrent_packed_decode(
+        mixed_qkv=inp.mixed_qkv,
+        g=inp.g1[0],
+        beta=inp.beta[0],
+        A_log=inp.A_log,
+        dt_bias=inp.dt_bias,
+        lower_bound=GATE_LOWER_BOUND,
+        initial_state=recurrent_state,
+        ssm_state_indices=inp.state_indices[:num_tokens],
+        conv_state=inp.conv_state_view(conv_state),
+        conv_weight=inp.decode_conv1d_weight,
+        out_gate=inp.g2[:num_tokens],
+        norm_weight=inp.norm_weight_bf16,
+        norm_eps=NORM_EPS,
+        out=core_attn_out[0, :num_tokens],
+    )
+    return core_attn_out, conv_state, recurrent_state
+
+
+@requires_aiter_gluon_kda
+@pytest.mark.parametrize("num_heads", [12, 24])
+@pytest.mark.parametrize("num_tokens", [1, 7, 128])
+@torch.inference_mode()
+def test_aiter_packed_decode_matches_triton_chain(
+    num_heads: int, num_tokens: int
+) -> None:
+    inp = KdaDecodeInputs(num_tokens, num_heads, num_slots=max(num_tokens, 4) + 3)
+
+    expected_out, expected_conv, expected_state = _run_triton_chain(inp)
+    actual_out, actual_conv, actual_state = _run_aiter_packed_decode(inp)
+
+    # The fused kernel keeps the recurrent output in fp32 through the norm,
+    # while the Triton chain rounds it to BF16 in between.
+    torch.testing.assert_close(actual_out, expected_out, atol=3e-2, rtol=3e-2)
+    # The conv state is a pure shift-and-append: it must be bit-exact.
+    torch.testing.assert_close(actual_conv, expected_conv, atol=0, rtol=0)
+    torch.testing.assert_close(actual_state, expected_state, atol=2e-3, rtol=2e-3)
+
+
+@requires_aiter_gluon_kda
+@torch.inference_mode()
+def test_aiter_packed_decode_leaves_untouched_slots_alone() -> None:
+    """Only the slots named by the state indices may be written."""
+    num_tokens, num_heads, num_slots = 4, 12, 9
+    inp = KdaDecodeInputs(num_tokens, num_heads, num_slots, seed=3)
+
+    _, actual_conv, actual_state = _run_aiter_packed_decode(inp)
+
+    touched = set(inp.state_indices.tolist())
+    untouched = [slot for slot in range(num_slots) if slot not in touched]
+    assert untouched, "test needs at least one unused slot"
     torch.testing.assert_close(
-        layer.decode_norm_weight, layer.o_norm.weight.float(), atol=0, rtol=0
+        actual_conv[untouched], inp.conv_state[untouched], atol=0, rtol=0
+    )
+    torch.testing.assert_close(
+        actual_state[untouched], inp.recurrent_state[untouched], atol=0, rtol=0
+    )
+
+
+@requires_aiter_gluon_kda
+@torch.inference_mode()
+def test_aiter_packed_decode_skips_null_block_padding() -> None:
+    """A CUDA-graph decode batch is padded with NULL_BLOCK_ID (0).
+
+    gdn_attn.py fills the tail of ``non_spec_state_indices_tensor`` with 0 and
+    the Triton chain leaves slot 0 alone; the fused kernel's pad_slot_guard
+    must do the same rather than read-modify-write slot 0 once per padded row.
+    """
+    num_real, num_padded, num_heads = 3, 5, 12
+    inp = KdaDecodeInputs(num_real + num_padded, num_heads, num_slots=12, seed=11)
+    inp.state_indices[num_real:] = 0
+
+    expected_out, expected_conv, expected_state = _run_triton_chain(inp)
+    actual_out, actual_conv, actual_state = _run_aiter_packed_decode(inp)
+
+    torch.testing.assert_close(
+        actual_out[0, :num_real], expected_out[0, :num_real], atol=3e-2, rtol=3e-2
+    )
+    assert torch.isfinite(actual_out).all()
+    # Slot 0 must be untouched by both paths.
+    torch.testing.assert_close(actual_conv[0], inp.conv_state[0], atol=0, rtol=0)
+    torch.testing.assert_close(actual_state[0], inp.recurrent_state[0], atol=0, rtol=0)
+    torch.testing.assert_close(actual_conv, expected_conv, atol=0, rtol=0)
+    torch.testing.assert_close(actual_state, expected_state, atol=2e-3, rtol=2e-3)
+
+
+# --------------------------------------------------------------------------
+# spec-only decode: fused_recurrent_kda with conv / gate / norm fused in
+# --------------------------------------------------------------------------
+
+
+class SpecBatch:
+    """A spec-only KDA batch: ``num_spec + 1`` query tokens per request at most.
+
+    ``seq_lens`` are the per-request query lengths, ``num_accepted`` the tokens
+    verified in the previous step; ``state_indices`` is the [N, num_spec + 1]
+    per-token slot table the GDN metadata hands the layer.
+    """
+
+    def __init__(
+        self,
+        seq_lens: list[int],
+        num_accepted: list[int],
+        num_heads: int,
+        num_spec: int,
+        seed: int = 0,
+    ) -> None:
+        assert len(seq_lens) == len(num_accepted)
+        assert max(seq_lens) <= num_spec + 1
+        n_seqs = len(seq_lens)
+        slots_per_seq = num_spec + 1
+        num_slots = n_seqs * slots_per_seq + 2
+        self.inp = KdaDecodeInputs(
+            sum(seq_lens), num_heads, num_slots, seed=seed, num_spec=num_spec
+        )
+        # Distinct slots per (request, token), shuffled, none at slot 0.
+        perm = torch.randperm(num_slots - 1, device="cuda")[: n_seqs * slots_per_seq]
+        self.state_indices = (perm + 1).to(torch.int32).view(n_seqs, slots_per_seq)
+        self.query_start_loc = torch.tensor(
+            [0] + torch.tensor(seq_lens).cumsum(0).tolist(),
+            dtype=torch.int32,
+            device="cuda",
+        )
+        self.num_accepted_tokens = torch.tensor(
+            num_accepted, dtype=torch.int32, device="cuda"
+        )
+        self.max_query_len = slots_per_seq
+
+
+def _run_spec_triton_chain(
+    sb: SpecBatch,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """The layer's unfused spec path: conv update -> recurrent KDA -> norm."""
+    from vllm.model_executor.layers.mamba.ops.causal_conv1d import (
+        causal_conv1d_update,
+    )
+    from vllm.models.kimi_k3.amd.ops.third_party.kda import fused_recurrent_kda
+
+    inp = sb.inp
+    conv_state = inp.conv_state.clone()
+    recurrent_state = inp.recurrent_state.clone()
+    conv_out = causal_conv1d_update(
+        inp.mixed_qkv,
+        inp.conv_state_view(conv_state),
+        inp.conv_weights,
+        None,
+        activation="silu",
+        conv_state_indices=sb.state_indices[:, 0],
+        num_accepted_tokens=sb.num_accepted_tokens,
+        query_start_loc=sb.query_start_loc,
+        max_query_len=sb.max_query_len,
+        validate_data=False,
+        out=torch.empty_like(inp.mixed_qkv),
+    )
+    q, k, v = (
+        x.view(1, -1, inp.num_heads, HEAD_DIM) for x in conv_out.split(inp.dim, dim=-1)
+    )
+    core_attn_out, _ = fused_recurrent_kda(
+        q=q,
+        k=k,
+        v=v,
+        raw_g=inp.g1,
+        raw_beta=inp.beta,
+        A_log=inp.A_log,
+        dt_bias=inp.dt_bias,
+        lower_bound=GATE_LOWER_BOUND,
+        initial_state=recurrent_state,
+        cu_seqlens=sb.query_start_loc,
+        ssm_state_indices=sb.state_indices,
+        num_accepted_tokens=sb.num_accepted_tokens,
+    )
+    out = _gated_rmsnorm(core_attn_out, inp.g2, inp.norm_weight_bf16, NORM_EPS)
+    return out, conv_state, recurrent_state
+
+
+def _run_aiter_spec(
+    sb: SpecBatch,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Exactly the call KimiK3DeltaAttention._forward makes for spec-only batches."""
+    from vllm._aiter_ops import rocm_aiter_ops
+
+    inp = sb.inp
+    conv_state = inp.conv_state.clone()
+    recurrent_state = inp.recurrent_state.clone()
+    q, k, v = (
+        x.view(1, -1, inp.num_heads, HEAD_DIM)
+        for x in inp.mixed_qkv.split(inp.dim, dim=-1)
+    )
+    n_spec = q.shape[1]
+    core_attn_out = inp.empty_out()
+    rocm_aiter_ops.kda_fused_recurrent_spec_decode(
+        q=q,
+        k=k,
+        v=v,
+        g=inp.g1,
+        beta=inp.beta,
+        A_log=inp.A_log,
+        dt_bias=inp.dt_bias,
+        lower_bound=GATE_LOWER_BOUND,
+        initial_state=recurrent_state,
+        cu_seqlens=sb.query_start_loc,
+        ssm_state_indices=sb.state_indices,
+        num_accepted_tokens=sb.num_accepted_tokens,
+        conv_state=inp.conv_state_view(conv_state),
+        conv_weight=inp.decode_conv1d_weight,
+        out_gate=inp.g2[:n_spec].unsqueeze(0),
+        norm_weight=inp.norm_weight_bf16,
+        norm_eps=NORM_EPS,
+        out=core_attn_out[:, :n_spec],
+    )
+    return core_attn_out, conv_state, recurrent_state
+
+
+@requires_aiter_gluon_kda
+@pytest.mark.parametrize("num_heads", [12, 24])
+@pytest.mark.parametrize(
+    "num_spec, seq_lens, num_accepted",
+    [
+        (3, [4, 4, 4], [1, 2, 4]),  # Kimi-K3 recipe: 3 draft tokens, uniform
+        (3, [4, 2, 1, 4], [4, 1, 1, 3]),  # ragged after partial acceptance
+        (6, [7, 3], [1, 4]),  # longer window, two requests
+    ],
+)
+@torch.inference_mode()
+def test_aiter_spec_decode_matches_triton_chain(
+    num_heads: int, num_spec: int, seq_lens: list[int], num_accepted: list[int]
+) -> None:
+    sb = SpecBatch(seq_lens, num_accepted, num_heads, num_spec, seed=5)
+
+    expected_out, expected_conv, expected_state = _run_spec_triton_chain(sb)
+    actual_out, actual_conv, actual_state = _run_aiter_spec(sb)
+
+    torch.testing.assert_close(actual_out, expected_out, atol=3e-2, rtol=3e-2)
+    torch.testing.assert_close(actual_conv, expected_conv, atol=0, rtol=0)
+    torch.testing.assert_close(actual_state, expected_state, atol=2e-3, rtol=3e-2)
+    # Nothing outside the slot table may change.
+    touched = set(sb.state_indices.flatten().tolist())
+    untouched = [s for s in range(sb.inp.conv_state.shape[0]) if s not in touched]
+    assert untouched
+    torch.testing.assert_close(
+        actual_conv[untouched], sb.inp.conv_state[untouched], atol=0, rtol=0
+    )
+    torch.testing.assert_close(
+        actual_state[untouched], sb.inp.recurrent_state[untouched], atol=0, rtol=0
     )
-    assert not layer.decode_norm_weight.requires_grad

tests/models/language/generation/test_common.py

@@ -75,6 +75,9 @@ AITER_MODEL_LIST = [
             "facebook/opt-125m",  # opt
             marks=[pytest.mark.core_model, pytest.mark.cpu_model],
         ),
+        pytest.param(
+            "facebook/opt-350m",  # opt with projected token embeddings
+        ),
         pytest.param(
             "microsoft/phi-2",  # phi
             marks=[pytest.mark.core_model, pytest.mark.slow_test],

tests/models/multimodal/generation/test_vit_cudagraph.py

@@ -70,6 +70,10 @@ def gemma3_chat_template(content: str) -> str:
     return f"<bos><start_of_turn>user\n{content}<end_of_turn>\n<start_of_turn>model\n"
 
 
+def deepseek_vl2_chat_template(content: str) -> str:
+    return f"<|User|>: {content}\n\n<|Assistant|>:"
+
+
 def ernie45_vl_chat_template(content: str) -> str:
     return (
         f"<|begin_of_sentence|>User: {content}"
@@ -255,7 +259,7 @@ MODEL_CONFIGS: dict[str, VitCudagraphTestConfig] = {
         vllm_runner_kwargs={
             "load_format": "dummy",
             "attention_backend": (
-                "ROCM_FLASHMLA_SPARSE_DSV4"
+                "ROCM_FLASHMLA_SPARSE_DSV41"
                 if current_platform.is_rocm()
                 else "FLASHMLA_SPARSE_DSV41"
             ),
@@ -356,6 +360,15 @@ MODEL_CONFIGS: dict[str, VitCudagraphTestConfig] = {
         },
         skip=True,  # TODO: Re-enable this once OOM issues are resolved on CI.
     ),
+    "deepseek_vl2": VitCudagraphTestConfig(
+        model="deepseek-ai/deepseek-vl2-tiny",
+        modalities=["image"],
+        image_prompt=deepseek_vl2_chat_template("<image>\nWhat is in this image?"),
+        vllm_runner_kwargs={
+            "hf_overrides": {"architectures": ["DeepseekVLV2ForCausalLM"]},
+        },
+        marks=[pytest.mark.core_model],
+    ),
     "gemma4": VitCudagraphTestConfig(
         model="google/gemma-4-E2B-it",
         image_prompt=(

tests/models/multimodal/pooling/test_d3.py

@@ -1,10 +1,12 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""d3-edge against the checkpoint's own runtime (`trust_remote_code`).
+"""Decision 3.0 against the checkpoint's own runtime (`trust_remote_code`).
 
 The runtime reads a question's options from the 255 code logits of its readout
 at the last prompt token. vLLM's classify head must give the same code logits
 for the same prompt tokens, and for an image prompt that vLLM processes itself.
+d3-edge has an irregular layer layout; d3-nano ships its backbone as a single
+weight file without an index, next to the readout file.
 """
 
 import pytest
@@ -14,7 +16,7 @@ from transformers import AutoModel
 from vllm.assets.image import ImageAsset
 from vllm.distributed import cleanup_dist_env_and_memory
 
-MODEL = "vllm-sr/d3-edge"
+MODELS = ["vllm-sr/d3-edge", "vllm-sr/d3-nano"]
 STATE = (
     "The order arrived damaged yesterday. The customer has a receipt and asks "
     "for a replacement today."
@@ -47,7 +49,7 @@ def option_probs(logits: torch.Tensor, count: int) -> torch.Tensor:
     return torch.softmax(logits[:count].float(), dim=-1)
 
 
-@pytest.mark.parametrize("model", [MODEL])
+@pytest.mark.parametrize("model", MODELS)
 def test_d3_matches_the_checkpoint_runtime(vllm_runner, model: str) -> None:
     reference = AutoModel.from_pretrained(model, trust_remote_code=True)
     runtime = reference.runtime

tests/models/qwen4_exp/test_config.py

@@ -70,20 +70,21 @@ def test_moe_sp_rejects_pipeline_parallel() -> None:
         ),
     )
     with (
-        patch("vllm.platforms.current_platform.is_cuda", return_value=True),
+        patch("vllm.platforms.current_platform.is_cuda_alike", return_value=True),
         pytest.raises(ValueError, match="MoE SP requires PP=1"),
     ):
         Qwen4ExpForConditionalGenerationConfig.verify_and_update_config(config)
 
 
+@spawn_new_process_for_each_test
+@pytest.mark.parametrize("backend", ["amd", "nvidia"])
 @pytest.mark.parametrize("use_sp", [False, True])
-def test_qwen4_exp_mtp_returns_sample_and_multi_streams(use_sp: bool) -> None:
+def test_qwen4_exp_mtp_returns_sample_and_multi_streams(
+    backend: str, use_sp: bool
+) -> None:
     """MTP returns complete, contiguous outputs after gathering padded SP shards."""
-    from vllm.models.qwen4_exp.nvidia.mtp import (
-        Qwen4ExpMultiTokenPredictor,
-    )
-
-    model = object.__new__(Qwen4ExpMultiTokenPredictor)
+    mtp_module = f"vllm.models.qwen4_exp.{backend}.mtp"
+    model = object.__new__(import_module(mtp_module).Qwen4ExpMultiTokenPredictor)
     model.use_hc_sequence_parallel = use_sp
     model.use_sequence_parallel = False
     torch.nn.Module.__init__(model)
@@ -98,7 +99,9 @@ def test_qwen4_exp_mtp_returns_sample_and_multi_streams(use_sp: bool) -> None:
     def decoder(**kwargs):
         """Check that the embedding residual shares the hidden state's token slice."""
         hidden = kwargs["hidden_states"]
-        assert kwargs["prev_block_output"].shape == (hidden.shape[0], model.hidden_size)
+        if backend == "nvidia":
+            prev_block_output = kwargs["prev_block_output"]
+            assert prev_block_output.shape == (hidden.shape[0], model.hidden_size)
         return hidden, hidden, torch.zeros(hidden.shape[0], model.hc_count)
 
     model.layers = [decoder]
@@ -120,9 +123,9 @@ def test_qwen4_exp_mtp_returns_sample_and_multi_streams(use_sp: bool) -> None:
         return torch.cat([local, remote, torch.full_like(remote, float("nan"))])
 
     with (
-        patch("vllm.models.qwen4_exp.nvidia.mtp.get_pp_group", return_value=pp_group),
-        patch("vllm.models.qwen4_exp.nvidia.mtp.sp_shard", side_effect=lambda x: x[:2]),
-        patch("vllm.models.qwen4_exp.nvidia.mtp.sp_all_gather", side_effect=gather),
+        patch(f"{mtp_module}.get_pp_group", return_value=pp_group),
+        patch(f"{mtp_module}.sp_shard", side_effect=lambda x: x[:2]),
+        patch(f"{mtp_module}.sp_all_gather", side_effect=gather),
     ):
         sample_hidden, returned_multi_hidden = model.forward(
             input_ids=None,

tests/models/qwen4_exp/test_tpsp.py

@@ -14,8 +14,8 @@ from transformers import Qwen4ExpTextConfig
 
 from vllm.platforms import current_platform
 
-if not current_platform.is_cuda():
-    pytest.skip("Qwen4Exp SP requires CUDA", allow_module_level=True)
+if not current_platform.is_cuda_alike():
+    pytest.skip("Qwen4Exp SP requires CUDA or ROCm", allow_module_level=True)
 
 from tests.utils import multi_gpu_test
 from vllm.config import (
@@ -38,14 +38,19 @@ from vllm.models.common.ops.sequence_parallel import (
     sp_padding_mask,
     sp_shard,
 )
-from vllm.models.qwen4_exp.nvidia import model as qwen4_model
-from vllm.models.qwen4_exp.nvidia.model import Qwen4ExpDecoderLayer
 from vllm.utils.network_utils import get_open_port
 from vllm.v1.attention.backends.short_conv_attn import PleShortConvAttentionMetadata
 from vllm.v1.worker.workspace import init_workspace_manager, reset_workspace_manager
 
 from .test_ple import _ConvBatchCase, _make_conv_metadata
 
+if current_platform.is_rocm():
+    from vllm.models.qwen4_exp.amd import model as qwen4_model
+    from vllm.models.qwen4_exp.amd.model import Qwen4ExpDecoderLayer
+else:
+    from vllm.models.qwen4_exp.nvidia import model as qwen4_model
+    from vllm.models.qwen4_exp.nvidia.model import Qwen4ExpDecoderLayer
+
 
 class _TPAttentionProjection(nn.Module):
     """Keep real TP communication without setting up an attention KV cache."""
@@ -168,7 +173,8 @@ def _check_decoder_ple(rank: int) -> None:
                 set_current_vllm_config(config),
                 set_forward_context({layer.ple.prefix: metadata}, config),
             ):
-                layer.ple.start_prefetch(ids, qsl, context)
+                if hasattr(layer.ple, "start_prefetch"):
+                    layer.ple.start_prefetch(ids, qsl, context)
                 output = layer(
                     *inputs,
                     positions,
@@ -177,7 +183,10 @@ def _check_decoder_ple(rank: int) -> None:
                     ngram_context=context,
                 )
                 if config.parallel_config.enable_hc_sp:
-                    output = tuple(sp_all_gather(x)[:num_tokens] for x in output)
+                    # ROCm all-gather requires contiguous inputs.
+                    output = tuple(
+                        sp_all_gather(x.contiguous())[:num_tokens] for x in output
+                    )
                 # A later MoE call may reuse workspace-backed output storage.
                 outputs.append(tuple(x.clone() for x in output))
         torch.testing.assert_close(outputs[1], outputs[0], atol=0.002, rtol=0.02)
@@ -223,11 +232,38 @@ def _check_moe_sp(rank: int, config: VllmConfig) -> None:
             query_start_loc=None,
             ngram_context=None,
         )
-        actual = tuple(sp_all_gather(x)[:num_tokens] for x in output)
+        actual = tuple(sp_all_gather(x.contiguous())[:num_tokens] for x in output)
     torch.testing.assert_close(actual, expected, atol=0.002, rtol=0.02)
 
 
-def _run_sp(rank: int, port: int, moe_sp: bool) -> None:
+def _check_compiled_sp_ops() -> None:
+    """Compiled SP collectives stay correct for token counts unseen at trace time."""
+    from vllm.models.qwen4_exp.amd.ops.sequence_parallel import (
+        sp_reduce_scatter,
+        sp_shard,
+    )
+
+    def collectives(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        gathered = sp_all_gather(sp_shard(x) * 2)[: x.shape[0]]
+        reduced = sp_all_gather(sp_reduce_scatter(x))[: x.shape[0]]
+        return gathered, reduced
+
+    # vLLM drops Dynamo guards, so one trace must serve every token count.
+    compiled = torch.compile(
+        collectives,
+        fullgraph=True,
+        options={"guard_filter_fn": torch.compiler.skip_all_guards_unsafe},
+    )
+    traced = torch.randn(8, 16)
+    torch._dynamo.mark_dynamic(traced, 0)
+    compiled(traced)
+    for num_tokens in (7, 3, 1):
+        torch.manual_seed(num_tokens)
+        x = torch.randn(num_tokens, 16)
+        torch.testing.assert_close(compiled(x), collectives(x), atol=0, rtol=0)
+
+
+def _run_sp(rank: int, port: int, moe_sp: bool, compiled_ops: bool = False) -> None:
     """Initialize rank-local runtime state and run one equivalence check."""
     torch.accelerator.set_device_index(rank)
     torch.set_num_threads(1)
@@ -247,6 +283,8 @@ def _run_sp(rank: int, port: int, moe_sp: bool) -> None:
                 torch.set_default_dtype(torch.bfloat16)
                 if moe_sp:
                     _check_moe_sp(rank, config)
+                elif compiled_ops:
+                    _check_compiled_sp_ops()
                 else:
                     _check_decoder_ple(rank)
     finally:
@@ -265,3 +303,10 @@ def test_moe_sp_equivalence(monkeypatch) -> None:
     """Native MoE SP preserves outputs with unequal DP batches and padding."""
     monkeypatch.setenv("VLLM_MOE_SKIP_PADDING", "1")
     mp.spawn(_run_sp, args=(get_open_port(), True), nprocs=4)
+
+
+@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm-only SP ops")
+@multi_gpu_test(num_gpus=2)
+def test_compiled_sp_ops_handle_untraced_token_counts() -> None:
+    """Padding in compiled SP collectives must not specialize on the traced size."""
+    mp.spawn(_run_sp, args=(get_open_port(), False, True), nprocs=2)

tests/models/registry.py

@@ -462,7 +462,8 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
     "OlmoHybridForCausalLM": _HfExamplesInfo("allenai/Olmo-Hybrid-7B"),
     "OlmoeForCausalLM": _HfExamplesInfo("allenai/OLMoE-1B-7B-0924-Instruct"),
     "OPTForCausalLM": _HfExamplesInfo(
-        "facebook/opt-125m", {"1b": "facebook/opt-iml-max-1.3b"}
+        "facebook/opt-125m",
+        {"350m": "facebook/opt-350m", "1b": "facebook/opt-iml-max-1.3b"},
     ),
     "OrionForCausalLM": _HfExamplesInfo(
         "OrionStarAI/Orion-14B-Chat", trust_remote_code=True
@@ -731,7 +732,13 @@ _SEQUENCE_CLASSIFICATION_EXAMPLE_MODELS = {
     ),
     "D3ForSequenceClassification": _HfExamplesInfo(
         "vllm-sr/d3-edge",
-        extras={"d3": "vllm-sr/d3", "d3-mini": "vllm-sr/d3-mini"},
+        extras={
+            "d3": "vllm-sr/d3",
+            "d3-flash": "vllm-sr/d3-flash",
+            "d3-mini": "vllm-sr/d3-mini",
+            "d3-nano": "vllm-sr/d3-nano",
+            "d3-lite": "vllm-sr/d3-lite",
+        },
         max_model_len=4096,
     ),
     "GPT2ForSequenceClassification": _HfExamplesInfo(

tests/models/test_deepseek_v41_pcp.py

@@ -0,0 +1,66 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+import pytest
+import torch
+
+from vllm.models.deepseek_v41.nvidia import model as model_module
+
+
+@pytest.mark.parametrize(
+    "pcp_size,use_ep", [(1, False), (2, True), (2, False), (4, False)]
+)
+@pytest.mark.parametrize("use_sp", [False, True])
+def test_pcp_gathers_routing_ids_once_after_local_input_preparation(
+    pcp_size, use_ep, use_sp, monkeypatch
+):
+    """Embedding keeps local IDs; all decoder layers receive aligned routing IDs."""
+    gather = Mock(side_effect=lambda ids, dim: torch.cat([ids] * pcp_size, dim=dim))
+    monkeypatch.setattr(
+        model_module, "get_pcp_group", lambda: SimpleNamespace(all_gather=gather)
+    )
+    monkeypatch.setattr(
+        model_module,
+        "get_pp_group",
+        lambda: SimpleNamespace(is_first_rank=True, is_last_rank=True),
+    )
+    monkeypatch.setattr(model_module, "is_forward_context_available", lambda: False)
+    monkeypatch.setattr(model_module, "sp_shard", lambda tensor: tensor[::2])
+    monkeypatch.setattr(
+        model_module, "sp_all_gather", lambda tensor: tensor.repeat_interleave(2, dim=0)
+    )
+    monkeypatch.setattr(model_module, "hc_collapse_triton", lambda hidden, mix: hidden)
+    embed = Mock(side_effect=lambda ids: ids[:, None].float())
+    run_layers = Mock(
+        side_effect=lambda layer_ids, hidden, *args: (hidden, None, None, None, hidden)
+    )
+    model = SimpleNamespace(
+        parallel_config=SimpleNamespace(
+            prefill_context_parallel_size=pcp_size, enable_expert_parallel=use_ep
+        ),
+        use_native_mega_moe=False,
+        use_sequence_parallel=use_sp,
+        engram_hash=None,
+        embed_input_ids=embed,
+        start_layer=0,
+        decoder_replay_start=2,
+        decoder_replay_layers=None,
+        aux_hidden_state_layers=(),
+        _mtp_hidden_buffer=None,
+        _run_layers=run_layers,
+        _collapse=lambda hidden, *args: hidden,
+        norm=lambda hidden: hidden,
+    )
+    for offset in (0, 10):
+        local_ids = torch.tensor([3, 7, 11, 15]) + offset
+        model_module.DeepseekV4Model.forward(model, local_ids, torch.arange(4), None)
+        assert embed.call_args.args[0] is local_ids
+        routing_ids = local_ids[::2] if use_sp else local_ids
+        needs_gather = pcp_size > 1 and not use_ep
+        if needs_gather:
+            torch.testing.assert_close(gather.call_args.args[0], routing_ids)
+            routing_ids = routing_ids.repeat(pcp_size)
+        torch.testing.assert_close(run_layers.call_args.args[3], routing_ids)
+    assert gather.call_count == 2 * int(needs_gather)

tests/models/test_deepseek_v41_replay_batch.py

@@ -4,7 +4,7 @@
 
 from dataclasses import replace
 from types import SimpleNamespace
-from unittest.mock import MagicMock
+from unittest.mock import MagicMock, patch
 
 import numpy as np
 import pytest
@@ -16,6 +16,9 @@ from vllm.forward_context import (
     ForwardContext,
     override_forward_context,
 )
+from vllm.model_executor.layers.fused_moe.routed_experts_capturer import (
+    RoutedExpertsCapturer,
+)
 from vllm.models.deepseek_v41.decoder_replay_layers import DecoderReplayLayers
 from vllm.models.deepseek_v41.nvidia.model_state import DeepseekV41ModelState
 from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor
@@ -63,6 +66,7 @@ def state(monkeypatch):
     cfg.scheduler_config.max_num_seqs = 16
     cfg.scheduler_config.max_num_batched_tokens = 1024
     cfg.parallel_config.data_parallel_size = 1
+    cfg.aux_output_config.enable_return_routed_experts = False
     cfg.compilation_config.fast_moe_cold_start = False
     cfg.compilation_config.cudagraph_mode = CUDAGraphMode.NONE
     layers = DecoderReplayLayers(WINDOW, MagicMock(), MagicMock(), set(), "swa_first")
@@ -172,6 +176,31 @@ def test_replay_batch_keeps_each_request_window(state):
     assert context.dp_metadata is None and context.is_padding is None
 
 
+def test_replay_context_restores_tp_routing_rows(state):
+    """Capture uses the compact replay batch's size, not the outer batch's."""
+    batch = _input_batch(QUERY_LENS, SEQ_LENS, PREFILLING)
+    replay, _ = _prepare(state, batch, CUDAGraphMode.NONE)
+    context = replay.forward_context
+    assert context.is_padding is None and context.dp_metadata is None
+    count = len(REPLAY_ROWS)
+    assert context.batch_descriptor.num_tokens == count
+    capturer = RoutedExpertsCapturer.__new__(RoutedExpertsCapturer)
+    capturer.dp_rank, capturer.tp_size = 0, 2
+    capturer.device_buffer = torch.full((401, 1, 2), -1, device=DEVICE)
+    expected = torch.tensor(REPLAY_ROWS, device=DEVICE)[:, None].expand(-1, 2)
+    gathered = torch.cat((expected, torch.full((1, 2), -1, device=DEVICE)))
+    with (
+        override_forward_context(context),
+        patch(
+            "vllm.model_executor.layers.fused_moe.routed_experts_capturer.get_tp_group"
+        ) as group,
+    ):
+        group.return_value.all_gather.return_value = gathered
+        capturer.capture(0, gathered[: (count + 1) // 2])
+    torch.testing.assert_close(capturer.device_buffer[:count, 0], expected)
+    assert (capturer.device_buffer[count:] == -1).all()
+
+
 def test_replay_batch_keeps_device_decode_boundaries(state):
     """Adaptive verification resizes decodes on the GPU (CPU [2, 2], device
     [1, 3]): their rows stay and the boundaries shift only by trimmed rows."""
@@ -195,6 +224,27 @@ def test_prompt_logprobs_requests_keep_their_rows(state):
     assert build.batch.max_query_len == 300
 
 
+@pytest.mark.parametrize("cg_mode", [CUDAGraphMode.NONE, CUDAGraphMode.PIECEWISE])
+def test_aux_output_requests_keep_every_routing_row(state, cg_mode):
+    """R3 needs every prefill row even when final logits only need the window."""
+    state.vllm_config.aux_output_config.enable_return_routed_experts = True
+    for index in range(3):
+        state.add_request(
+            index,
+            SimpleNamespace(
+                replay_start=0,
+                sampling_params=SimpleNamespace(prompt_logprobs=None),
+            ),
+        )
+    batch = _input_batch(QUERY_LENS, SEQ_LENS, PREFILLING)
+    state.decoder_replay_layers.trim_threshold = WINDOW + 1
+    replay, _ = _prepare(state, batch, cg_mode)
+    if cg_mode == CUDAGraphMode.NONE:
+        assert replay is None
+    else:
+        torch.testing.assert_close(replay.rows, torch.arange(401, device=DEVICE))
+
+
 def test_replay_batch_keeps_adaptive_verification_query_bound(state):
     """Adaptive verification bounds the decodes' device-side query lengths
     above their CPU lengths; the replay batch keeps that bound."""

tests/models/test_deepseek_v4_mega_moe.py

@@ -8,10 +8,8 @@ import torch
 from transformers import DeepseekV4Config
 
 from vllm.model_executor.layers.fused_moe.deep_gemm_mega_moe import (
-    DeepGemmSm100MegaMoEBackend,
-)
-from vllm.model_executor.layers.fused_moe.routed_experts_capturer import (
-    bind_routed_experts_capturer,
+    DeepGemmMegaMoEExperts,
+    MegaMoEActivation,
 )
 from vllm.model_executor.models.interfaces import (
     get_mixture_of_experts_model,
@@ -20,17 +18,14 @@ from vllm.model_executor.models.interfaces import (
 from vllm.models.deepseek_v4.nvidia.dspark import DSparkDeepseekV4ForCausalLM
 from vllm.models.deepseek_v4.nvidia.model import (
     DeepseekV4ForCausalLM,
-    DeepseekV4MegaMoEExperts,
     DeepseekV4MoE,
     make_deepseek_v4_expert_params_mapping,
     prepare_mega_gate_routing_metadata,
 )
-from vllm.models.deepseek_v4.nvidia.ops.prepare_megamoe import prepare_megamoe_inputs
 from vllm.models.deepseek_v41.common.mm_preprocess import IMAGE_SENTINEL_BASE_ID
 from vllm.models.deepseek_v41.nvidia.model import DeepseekV4MoE as DeepseekV41MoE
 from vllm.platforms import current_platform
 from vllm.transformers_utils.configs.deepseek_v41 import DeepseekV41Config
-from vllm.utils.deep_gemm import is_deep_gemm_supported
 from vllm.utils.torch_utils import set_default_torch_dtype
 
 pytestmark = pytest.mark.skipif(
@@ -40,11 +35,12 @@ pytestmark = pytest.mark.skipif(
 
 
 @pytest.fixture(autouse=True)
-def sm100_mega_moe_backend(monkeypatch):
+def sm100_mega_moe_layout(monkeypatch):
     """Build the SM100 kernel layout on any GPU; kernels are gated per test."""
     monkeypatch.setattr(
-        "vllm.models.deepseek_v4.nvidia.model.get_deep_gemm_mega_moe_backend",
-        lambda *args: DeepGemmSm100MegaMoEBackend(),
+        "vllm.model_executor.layers.fused_moe.deep_gemm_mega_moe."
+        "check_deep_gemm_mega_moe_supported",
+        lambda *args: None,
     )
 
 
@@ -125,7 +121,8 @@ def test_deepseek_v41_moe_routes_without_hash_table(
         )
     assert moe.gate.tid2eid is None
     assert moe.gate.weight.shape == (num_experts, config.hidden_size)
-    assert isinstance(moe.experts, DeepseekV4MegaMoEExperts)
+    assert isinstance(moe.experts, DeepGemmMegaMoEExperts)
+    assert moe.experts.activation == MegaMoEActivation("swiglu", clamp=10.0)
     assert moe.experts.w13_weight.shape == (num_experts, 256, 128)
     assert moe.experts.top_k == top_k
     assert (config.n_routed_experts, config.num_experts_per_tok) == (384, 6)
@@ -151,10 +148,9 @@ def test_deepseek_v41_moe_routes_without_hash_table(
 
     routed = {}
 
-    def check_routing(x, weights, ids, *, activation_clamp):
+    def check_routing(x, weights, ids, shared_experts):
         routed["ids"], routed["weights"] = ids, weights
-        assert activation_clamp == config.swiglu_limit
-        return x.clone()
+        return x + shared_experts(x)
 
     monkeypatch.setattr(moe.experts, "forward", check_routing)
     monkeypatch.setattr(moe.shared_experts, "forward", lambda x: 2 * x)
@@ -280,10 +276,10 @@ def test_deepseek_v4_mega_gate_hash_routing_correctness(
         dim=-1, keepdim=True
     )
 
-    def check_routing(x, weights, ids, *, activation_clamp):
+    def check_routing(x, weights, ids, shared_experts):
         torch.testing.assert_close(ids, expected_ids)
         torch.testing.assert_close(weights, expected_weights, rtol=1e-3, atol=1e-4)
-        return x.clone()
+        return x + shared_experts(x)
 
     monkeypatch.setattr(moe.experts, "forward", check_routing)
     monkeypatch.setattr(moe.shared_experts, "forward", torch.zeros_like)
@@ -308,634 +304,6 @@ def test_deepseek_v4_mega_moe_expert_mapping():
     ]
 
 
-@pytest.mark.parametrize("use_kimi", [False, True])
-def test_deep_gemm_mega_moe_capture_precedes_eplb(monkeypatch, use_kimi):
-    experts_cls = DeepseekV4MegaMoEExperts
-    if use_kimi:
-        from vllm.models.kimi_k3.nvidia.model import KimiK3MegaMoEExperts
-
-        experts_cls = KimiK3MegaMoEExperts
-
-    experts = experts_cls.__new__(experts_cls)
-    torch.nn.Module.__init__(experts)
-    if use_kimi:
-        experts.synchronize_first_launch = lambda: None
-    experts.prefix = "model.layers.3.ffn.experts"
-    experts.max_num_tokens = 4
-    experts.capture_fn = None
-    experts.get_symm_buffer = lambda: object()
-    experts.eplb_state = SimpleNamespace(
-        logical_to_physical_map=torch.empty(1),
-        expert_load_view=torch.empty(1),
-        logical_replica_count=torch.empty(1),
-        should_record_tensor=torch.empty(1),
-        num_unpadded_tokens_tensors=None,
-    )
-
-    topk_ids = torch.tensor([[1, 2], [3, 4]])
-    captured: list[tuple[int, torch.Tensor]] = []
-    bind_routed_experts_capturer(
-        SimpleNamespace(modules=lambda: [experts]),
-        SimpleNamespace(capture=lambda layer_id, ids: captured.append((layer_id, ids))),
-    )
-
-    class MappingReached(Exception):
-        pass
-
-    def map_ids(**kwargs):
-        assert captured == [(3, topk_ids)]
-        raise MappingReached
-
-    monkeypatch.setattr(
-        f"{experts_cls.__module__}.eplb_map_to_physical_and_record",
-        map_ids,
-    )
-    monkeypatch.setattr(
-        "vllm.utils.deep_gemm._import_deep_gemm", lambda: SimpleNamespace()
-    )
-
-    with pytest.raises(MappingReached):
-        experts(
-            torch.empty(2, 8),
-            torch.empty(2, 2),
-            topk_ids,
-            activation_clamp=None,
-        )
-
-
-def test_deepseek_v4_mega_moe_weight_loader_uses_ep_expert_ownership():
-    vllm_config = SimpleNamespace(
-        scheduler_config=SimpleNamespace(max_num_batched_tokens=4),
-        compilation_config=SimpleNamespace(static_forward_context={}),
-    )
-    experts = DeepseekV4MegaMoEExperts(
-        vllm_config,
-        num_experts=4,
-        num_local_experts=2,
-        experts_start_idx=2,
-        top_k=2,
-        hidden_size=128,
-        intermediate_size=128,
-    )
-
-    nonlocal_weight = torch.ones(128, 64, dtype=torch.uint8)
-    assert (
-        experts.weight_loader(
-            experts.w13_weight,
-            nonlocal_weight,
-            "experts.w13_weight",
-            shard_id="w1",
-            expert_id=1,
-            return_success=True,
-        )
-        is False
-    )
-
-    w1 = torch.full((128, 64), 3, dtype=torch.uint8)
-    w3 = torch.full((128, 64), 7, dtype=torch.uint8)
-    w2 = torch.full((128, 64), 11, dtype=torch.uint8)
-
-    assert experts.weight_loader(
-        experts.w13_weight,
-        w1,
-        "experts.w13_weight",
-        shard_id="w1",
-        expert_id=2,
-        return_success=True,
-    )
-    assert experts.weight_loader(
-        experts.w13_weight,
-        w3,
-        "experts.w13_weight",
-        shard_id="w3",
-        expert_id=2,
-        return_success=True,
-    )
-    assert experts.weight_loader(
-        experts.w2_weight,
-        w2,
-        "experts.w2_weight",
-        shard_id="w2",
-        expert_id=2,
-        return_success=True,
-    )
-
-    interleaved = experts.w13_weight[0].view(torch.uint8).view(16, 2, 8, 64)
-    assert torch.equal(interleaved[:, 0], w1.view(16, 8, 64))
-    assert torch.equal(interleaved[:, 1], w3.view(16, 8, 64))
-    assert torch.equal(experts.w2_weight[0].view(torch.uint8), w2)
-    assert torch.count_nonzero(experts.w13_weight[1]) == 0
-
-
-def _reference_mega_moe_layout(w13, s13, w2, s2):
-    """CPU copy of DeepGEMM's transform_weights_for_mega_moe on packed scales."""
-
-    def interleave(t):
-        g, n, k = t.shape
-        halves = t.view(g, 2, n // 16, 8, k)
-        return torch.stack([halves[:, 0], halves[:, 1]], dim=2).reshape(g, n, k)
-
-    def utccp(sf):
-        g, mn, k = sf.shape
-        return sf.reshape(g, -1, 4, 32, k).transpose(2, 3).reshape(g, mn, k)
-
-    return (
-        interleave(w13),
-        utccp(interleave(s13.view(torch.int32))),
-        w2,
-        utccp(s2.view(torch.int32)),
-    )
-
-
-def _make_mega_moe_experts(num_local_experts, num_logical_experts):
-    return DeepseekV4MegaMoEExperts(
-        SimpleNamespace(
-            scheduler_config=SimpleNamespace(max_num_batched_tokens=4),
-            compilation_config=SimpleNamespace(static_forward_context={}),
-        ),
-        num_experts=num_local_experts,
-        num_local_experts=num_local_experts,
-        experts_start_idx=0,
-        top_k=1,
-        hidden_size=256,
-        intermediate_size=384,
-        num_logical_experts=num_logical_experts,
-    )
-
-
-def _load_mega_moe_checkpoint_shuffled(experts, num_logical_experts, seed):
-    hidden, inter = experts.hidden_size, experts.intermediate_size
-    shapes = {
-        "w13_weight": (num_logical_experts, 2 * inter, hidden // 2),
-        "w13_weight_scale": (num_logical_experts, 2 * inter, hidden // 32),
-        "w2_weight": (num_logical_experts, hidden, inter // 2),
-        "w2_weight_scale": (num_logical_experts, hidden, inter // 32),
-    }
-    generator = torch.Generator().manual_seed(seed)
-    ckpt = {
-        name: torch.randint(0, 256, shape, dtype=torch.uint8, generator=generator)
-        for name, shape in shapes.items()
-    }
-    shards = []
-    for expert_id in range(num_logical_experts):
-        for name, tensor in ckpt.items():
-            if name.startswith("w2_"):
-                shards.append((name, tensor[expert_id], "w2", expert_id))
-            else:
-                shards.append((name, tensor[expert_id, :inter], "w1", expert_id))
-                shards.append((name, tensor[expert_id, inter:], "w3", expert_id))
-    for i in torch.randperm(len(shards), generator=generator).tolist():
-        name, loaded_weight, shard_id, expert_id = shards[i]
-        if name.endswith("scale"):
-            loaded_weight = loaded_weight.view(torch.float8_e8m0fnu)
-        assert experts.weight_loader(
-            getattr(experts, name),
-            loaded_weight,
-            f"experts.{name}",
-            shard_id=shard_id,
-            expert_id=expert_id,
-            return_success=True,
-        )
-    return ckpt
-
-
-def _expected_mega_moe_layout(ckpt, num_local_experts, num_logical_experts):
-    physical = torch.arange(num_local_experts) % num_logical_experts
-    ckpt = {name: tensor[physical] for name, tensor in ckpt.items()}
-    expected = _reference_mega_moe_layout(*ckpt.values())
-    if is_deep_gemm_supported():
-        from vllm.utils.deep_gemm import _import_deep_gemm
-
-        deep_gemm = _import_deep_gemm()
-        w13, s13, w2, s2 = (t.cuda() for t in ckpt.values())
-        hidden, inter = w2.shape[1], w2.shape[2] * 2
-        scales = [
-            deep_gemm.transform_sf_into_required_layout(
-                (sf.to(torch.int32) << 23).view(torch.float32),
-                mn,
-                k,
-                (1, 32),
-                num_local_experts,
-            )
-            for sf, mn, k in ((s13, 2 * inter, hidden), (s2, hidden, inter))
-        ]
-        (l1, l1_sf), (l2, l2_sf) = deep_gemm.transform_weights_for_mega_moe(
-            (w13.view(torch.int8), scales[0]), (w2.view(torch.int8), scales[1])
-        )
-        for actual, reference in zip((l1, l1_sf, l2, l2_sf), expected):
-            assert torch.equal(actual.cpu().view(reference.dtype), reference)
-    return expected
-
-
-def _assert_mega_moe_layout(experts, expected):
-    actual = (*experts._transformed_l1_weights, *experts._transformed_l2_weights)
-    for tensor, reference in zip(actual, expected):
-        assert torch.equal(tensor.view(reference.dtype), reference)
-
-
-def test_deepseek_v4_mega_moe_loads_shards_into_kernel_layout():
-    """Shards in any order, including EPLB replicas, land in DeepGEMM's layout."""
-    experts = _make_mega_moe_experts(num_local_experts=4, num_logical_experts=3)
-    ckpt = _load_mega_moe_checkpoint_shuffled(experts, 3, seed=0)
-    _assert_mega_moe_layout(experts, _expected_mega_moe_layout(ckpt, 4, 3))
-    # DeepGEMM's MN-major scale strides.
-    assert experts.w13_weight_scale.stride() == (768 * 2, 1, 768)
-    assert experts.w2_weight_scale.stride() == (256 * 3, 1, 256)
-
-
-def test_deepseek_v4_mega_moe_reload_overwrites_kernel_layout_in_place():
-    """A second load must update the live kernel tensors without reallocating."""
-    experts = _make_mega_moe_experts(num_local_experts=2, num_logical_experts=2)
-    _load_mega_moe_checkpoint_shuffled(experts, 2, seed=0)
-    data_ptrs = [param.data_ptr() for param in experts.parameters()]
-    ckpt = _load_mega_moe_checkpoint_shuffled(experts, 2, seed=1)
-    assert [param.data_ptr() for param in experts.parameters()] == data_ptrs
-    _assert_mega_moe_layout(experts, _expected_mega_moe_layout(ckpt, 2, 2))
-
-
-@pytest.mark.parametrize("intermediate_size", [512, 2304])
-def test_deepseek_v4_mega_moe_preserves_checkpoint_dimensions(intermediate_size):
-    """Keep native widths so V4.1 routed and shared experts can fuse."""
-    experts = DeepseekV4MegaMoEExperts(
-        SimpleNamespace(
-            scheduler_config=SimpleNamespace(max_num_batched_tokens=4),
-            compilation_config=SimpleNamespace(static_forward_context={}),
-        ),
-        num_experts=2,
-        num_local_experts=2,
-        experts_start_idx=0,
-        top_k=1,
-        hidden_size=128,
-        intermediate_size=intermediate_size,
-    )
-
-    assert experts.intermediate_size == intermediate_size
-    assert experts.w13_weight.shape == (2, 2 * intermediate_size, 64)
-    assert experts.w13_weight_scale.shape == (2, 2 * intermediate_size, 1)
-    assert experts.w2_weight.shape == (2, 128, intermediate_size // 2)
-    assert experts.w2_weight_scale.shape == (2, 128, intermediate_size // 128)
-    transformed = (*experts._transformed_l1_weights, *experts._transformed_l2_weights)
-    for actual, param in zip(transformed, experts.parameters()):
-        assert actual.data_ptr() == param.data_ptr()
-
-
-def test_deepseek_v4_mega_moe_loads_shared_experts_into_kernel_layout(monkeypatch):
-    """Shared expert shards land directly in the MegaMoE kernel layout."""
-    monkeypatch.setattr(
-        DeepGemmSm100MegaMoEBackend, "supports_shared_experts", lambda *args: True
-    )
-    experts = DeepseekV4MegaMoEExperts(
-        SimpleNamespace(
-            scheduler_config=SimpleNamespace(max_num_batched_tokens=4),
-            compilation_config=SimpleNamespace(static_forward_context={}),
-        ),
-        num_experts=2,
-        num_local_experts=2,
-        experts_start_idx=0,
-        top_k=1,
-        hidden_size=256,
-        intermediate_size=384,
-        num_shared_experts=1,
-    )
-    assert experts.has_fused_shared_experts
-    hidden, inter = 256, 384
-    generator = torch.Generator().manual_seed(42)
-
-    def rand(*shape):
-        return torch.randint(0, 256, shape, dtype=torch.uint8, generator=generator)
-
-    w1, w3, w2 = rand(inter, hidden), rand(inter, hidden), rand(hidden, inter)
-    s1, s3, s2 = (
-        rand(inter, hidden // 32),
-        rand(inter, hidden // 32),
-        rand(hidden, inter // 32),
-    )
-    fp8, e8m0 = torch.float8_e4m3fn, torch.float8_e8m0fnu
-    # The down shard arrives as `.w2.` (MTP, DSpark) or `.down_proj.` (mapped).
-    for name, data in [
-        ("shared_experts.w1.weight", w1.view(fp8)),
-        ("shared_experts.w3.weight", w3.view(fp8)),
-        ("shared_experts.down_proj.weight", w2.view(fp8)),
-        ("shared_experts.w1.weight_scale", s1.view(e8m0)),
-        ("shared_experts.w3.weight_scale", s3.view(e8m0)),
-        ("shared_experts.w2.weight_scale", s2.view(e8m0)),
-    ]:
-        loaded = experts.load_shared_expert_weight(f"ffn.{name}", data)
-        # Weight tracking checks this name against the registered parameters.
-        assert loaded in {f"ffn.experts.{n}" for n, _ in experts.named_parameters()}
-
-    def l1_halves():
-        return experts.shared_l1_weight.view(torch.uint8).view(-1, 2, 8, hidden)
-
-    # Every tensor matches DeepGEMM's batched transform of the checkpoint.
-    expected = _reference_mega_moe_layout(
-        *(t.unsqueeze(0) for t in (torch.cat([w1, w3]), torch.cat([s1, s3]), w2, s2))
-    )
-    actual = (
-        experts.shared_l1_weight,
-        experts.shared_l1_weight_scale,
-        experts.shared_l2_weight,
-        experts.shared_l2_weight_scale,
-    )
-    for tensor, reference in zip(actual, expected):
-        assert torch.equal(tensor.view(reference.dtype), reference[0])
-    # A reload of one shard overwrites in place and leaves the other half alone.
-    ptr = experts.shared_l1_weight.data_ptr()
-    w1b = rand(inter, hidden)
-    experts.load_shared_expert_weight("shared_experts.w1.weight", w1b.view(fp8))
-    assert experts.shared_l1_weight.data_ptr() == ptr
-    assert torch.equal(l1_halves()[:, 0], w1b.view(-1, 8, hidden))
-    assert torch.equal(l1_halves()[:, 1], w3.view(-1, 8, hidden))
-    with pytest.raises(ValueError):
-        experts.load_shared_expert_weight("shared_experts.gate_up_proj.weight", w1)
-
-
-def _fake_deep_gemm(with_shared: bool):
-    def old_buffer(group): ...
-
-    def old_kernel(y): ...
-
-    def new_buffer(group, num_shared_experts=0): ...
-
-    def new_kernel(y, shared_l1_weights=None, shared_l2_weights=None): ...
-
-    return SimpleNamespace(
-        get_symm_buffer_for_mega_moe=new_buffer if with_shared else old_buffer,
-        fp8_fp4_mega_moe=new_kernel if with_shared else old_kernel,
-        get_block_m_for_mega_moe=lambda *args: 0,
-    )
-
-
-@pytest.mark.parametrize(
-    "weight_dtype,scale_dtype,scale_shape,with_shared,fused",
-    [
-        (torch.float8_e4m3fn, torch.float8_e8m0fnu, (2, 2), True, True),
-        (torch.float8_e4m3fn, torch.uint8, (256, 8), True, True),
-        (torch.float8_e4m3fn, torch.float8_e8m0fnu, (2, 2), False, False),
-        (torch.bfloat16, None, None, True, False),
-        (torch.float8_e4m3fn, torch.float32, (2, 2), True, False),
-        (torch.float8_e4m3fn, torch.uint8, (256, 16), True, False),
-    ],
-)
-def test_deepseek_v4_mega_moe_decides_shared_fusion_at_construction(
-    monkeypatch, weight_dtype, scale_dtype, scale_shape, with_shared, fused
-):
-    """Fuse only shared experts the installed kernel and the 1x32 layout fit."""
-    monkeypatch.setattr(
-        "vllm.utils.deep_gemm._import_deep_gemm", lambda: _fake_deep_gemm(with_shared)
-    )
-
-    def linear():
-        scale = (
-            None if scale_dtype is None else torch.empty(scale_shape).to(scale_dtype)
-        )
-        return SimpleNamespace(
-            weight=torch.empty(256, 256, dtype=weight_dtype), weight_scale_inv=scale
-        )
-
-    shared = SimpleNamespace(gate_up_proj=linear(), down_proj=linear())
-    experts = DeepseekV4MegaMoEExperts(
-        SimpleNamespace(
-            scheduler_config=SimpleNamespace(max_num_batched_tokens=4),
-            compilation_config=SimpleNamespace(static_forward_context={}),
-        ),
-        num_experts=2,
-        num_local_experts=2,
-        experts_start_idx=0,
-        top_k=1,
-        hidden_size=256,
-        intermediate_size=128,
-        num_shared_experts=1,
-        shared_experts=shared,
-    )
-
-    assert experts.has_fused_shared_experts == fused
-    assert hasattr(experts, "shared_l1_weight") == fused
-
-
-@pytest.mark.parametrize("fused", [False, True])
-def test_deepseek_v4_mega_moe_does_not_double_add_fused_shared_expert(
-    monkeypatch, fused
-):
-    class FakeGate(torch.nn.Module):
-        weight = torch.empty(2, 128)
-        tid2eid = None
-        e_score_correction_bias = None
-
-        def forward(self, hidden_states):
-            return torch.empty(hidden_states.shape[0], 2), None
-
-    class FakeExperts(torch.nn.Module):
-        has_fused_shared_experts = fused
-        backend = DeepGemmSm100MegaMoEBackend()
-
-        def forward(self, hidden_states, *args, **kwargs):
-            return torch.ones_like(hidden_states)
-
-    class FakeSharedExperts(torch.nn.Module):
-        calls = 0
-
-        def forward(self, hidden_states):
-            self.calls += 1
-            return torch.full_like(hidden_states, 2)
-
-    moe = DeepseekV4MoE.__new__(DeepseekV4MoE)
-    torch.nn.Module.__init__(moe)
-    moe.use_native_mega_moe = True
-    moe.gate = FakeGate()
-    moe.experts = FakeExperts()
-    moe.shared_experts = None if fused else FakeSharedExperts()
-    moe.scoring_func = "sqrtsoftplus"
-    moe.n_activated_experts = 1
-    moe.renormalize = True
-    moe.hash_indices_dtype = torch.int64
-    moe.routed_scaling_factor = 1.0
-    moe.swiglu_limit = 10.0
-    monkeypatch.setattr(
-        "vllm.models.deepseek_v4.nvidia.model.fused_topk_bias",
-        lambda **kwargs: (
-            torch.ones(kwargs["hidden_states"].shape[0], 1),
-            torch.zeros(kwargs["hidden_states"].shape[0], 1, dtype=torch.int64),
-        ),
-    )
-
-    output = moe(torch.zeros(2, 128))
-
-    expected = 1 if fused else 3
-    assert torch.all(output == expected)
-    if not fused:
-        assert moe.shared_experts.calls == 1
-
-
-@pytest.mark.skipif(
-    not torch.cuda.is_available(),
-    reason="DeepSeek V4 MegaMoE fused input staging requires CUDA.",
-)
-@pytest.mark.parametrize("num_tokens", [7, 63, 64, 65, 127, 128, 135, 434, 16384])
-@pytest.mark.parametrize("hidden_size", [256, 5120])
-def test_deepseek_v4_mega_moe_fused_input_staging_is_bitwise_exact(
-    num_tokens, hidden_size
-):
-    from vllm.third_party.deep_gemm.utils import per_token_cast_to_fp8
-
-    device = torch.device("cuda")
-    top_k = 8
-
-    generator = torch.Generator(device=device)
-    generator.manual_seed(0)
-    hidden_states = (
-        torch.randn(
-            num_tokens,
-            hidden_size,
-            device=device,
-            dtype=torch.float32,
-            generator=generator,
-        )
-        * 17.0
-    ).to(torch.bfloat16)
-    hidden_states[0, :32] = 0
-    hidden_states[1, 32:64] = 1.0e-6
-    hidden_states[2, 64:96] = -1.0e-6
-
-    topk_ids = torch.randint(
-        0,
-        256,
-        (num_tokens, top_k),
-        device=device,
-        dtype=torch.int32,
-        generator=generator,
-    )
-    topk_weights = torch.randn(
-        num_tokens,
-        top_k,
-        device=device,
-        dtype=torch.float32,
-        generator=generator,
-    )
-
-    ref_x, ref_x_sf = per_token_cast_to_fp8(
-        hidden_states,
-        use_ue8m0=True,
-        gran_k=32,
-        use_packed_ue8m0=True,
-    )
-    ref_topk_idx = topk_ids.to(torch.int64)
-    ref_topk_weights = topk_weights.clone()
-
-    fused_x = torch.empty_like(ref_x)
-    fused_x_sf = torch.empty_like(ref_x_sf)
-    fused_topk_idx = torch.empty_like(ref_topk_idx)
-    fused_topk_weights = torch.empty_like(ref_topk_weights)
-
-    prepare_megamoe_inputs(
-        hidden_states,
-        topk_weights,
-        topk_ids,
-        fused_x,
-        fused_x_sf,
-        fused_topk_idx,
-        fused_topk_weights,
-    )
-    torch.accelerator.synchronize()
-
-    assert torch.equal(fused_x.view(torch.uint8), ref_x.view(torch.uint8))
-    assert torch.equal(fused_x_sf, ref_x_sf)
-    assert torch.equal(fused_topk_idx, ref_topk_idx)
-    assert torch.equal(
-        fused_topk_weights.view(torch.uint8),
-        ref_topk_weights.view(torch.uint8),
-    )
-
-
-@pytest.mark.parametrize("shared_block_m", [8, 32, 96, 128, 192])
-def test_deepseek_v4_mega_moe_stages_shared_scale_tma_layout(shared_block_m):
-    from vllm.third_party.deep_gemm.utils import per_token_cast_to_fp8
-
-    device = torch.device("cuda")
-    num_tokens = shared_block_m + 7
-    hidden_size = 256
-    top_k = 8
-    generator = torch.Generator(device=device)
-    generator.manual_seed(shared_block_m)
-    hidden_states = torch.randn(
-        num_tokens,
-        hidden_size,
-        device=device,
-        dtype=torch.bfloat16,
-        generator=generator,
-    )
-    topk_ids = torch.randint(
-        0,
-        256,
-        (num_tokens, top_k),
-        device=device,
-        dtype=torch.int32,
-        generator=generator,
-    )
-    topk_weights = torch.randn(
-        num_tokens,
-        top_k,
-        device=device,
-        dtype=torch.float32,
-        generator=generator,
-    )
-
-    ref_x, ref_x_sf = per_token_cast_to_fp8(
-        hidden_states,
-        use_ue8m0=True,
-        gran_k=32,
-        use_packed_ue8m0=True,
-    )
-    aligned_block_m = ((shared_block_m + 127) // 128) * 128
-    num_shared_rows = ((num_tokens + shared_block_m - 1) // shared_block_m) * (
-        aligned_block_m
-    )
-    ref_shared_x_sf = torch.zeros(
-        num_shared_rows,
-        hidden_size // 128,
-        dtype=torch.int32,
-        device=device,
-    )
-    for token_id in range(num_tokens):
-        m_in_block = token_id % shared_block_m
-        transposed_m = (
-            (m_in_block // 128) * 128 + (m_in_block % 32) * 4 + (m_in_block % 128) // 32
-        )
-        shared_row = token_id // shared_block_m * aligned_block_m + transposed_m
-        ref_shared_x_sf[shared_row].copy_(ref_x_sf[token_id])
-
-    fused_x = torch.empty_like(ref_x)
-    fused_x_sf = torch.empty_like(ref_x_sf)
-    fused_shared_storage = torch.full(
-        (hidden_size // 128, num_shared_rows),
-        -1,
-        dtype=torch.int32,
-        device=device,
-    )
-    fused_shared_x_sf = fused_shared_storage.t()
-    fused_topk_idx = torch.empty_like(topk_ids, dtype=torch.int64)
-    fused_topk_weights = torch.empty_like(topk_weights)
-
-    prepare_megamoe_inputs(
-        hidden_states,
-        topk_weights,
-        topk_ids,
-        fused_x,
-        fused_x_sf,
-        fused_topk_idx,
-        fused_topk_weights,
-        shared_x_sf=fused_shared_x_sf,
-        shared_block_m=shared_block_m,
-    )
-    torch.accelerator.synchronize()
-
-    populated = ref_shared_x_sf != 0
-    assert torch.equal(fused_x.view(torch.uint8), ref_x.view(torch.uint8))
-    assert torch.equal(fused_x_sf, ref_x_sf)
-    assert torch.equal(fused_shared_x_sf[populated], ref_shared_x_sf[populated])
-    assert torch.all(fused_shared_x_sf[~populated] == -1)
-
-
 def test_deepseek_v4_pwal_hook_finalizes_mhc_broadcast():
     """The loader invokes the model-level PWAL hook for every load format,
     so it must finalize mhc broadcast weights to cover dummy load, which
@@ -1040,85 +408,3 @@ def test_dspark_draft_registers_mixture_of_experts(
     assert len(draft.moe_layers) == draft.model.num_dspark_layers == 2
     assert draft.num_routed_experts == hf_config.n_routed_experts
     assert draft.num_redundant_experts == 4
-
-
-@pytest.mark.skipif(
-    not torch.cuda.is_available(),
-    reason="DeepSeek V4 MegaMoE fused input staging requires CUDA.",
-)
-@pytest.mark.parametrize("nonfinite_padding", [False, True])
-@pytest.mark.parametrize("num_tokens", [7, 65, 135])
-def test_deepseek_v4_mega_moe_fused_input_staging_masks_padding(
-    nonfinite_padding, num_tokens
-):
-    from vllm.third_party.deep_gemm.utils import per_token_cast_to_fp8
-
-    device = torch.device("cuda")
-    hidden_size = 256
-    top_k = 8
-
-    generator = torch.Generator(device=device)
-    generator.manual_seed(1)
-    hidden_states = torch.randn(
-        num_tokens,
-        hidden_size,
-        device=device,
-        dtype=torch.bfloat16,
-        generator=generator,
-    )
-    topk_ids = torch.randint(
-        0,
-        256,
-        (num_tokens, top_k),
-        device=device,
-        dtype=torch.int32,
-        generator=generator,
-    )
-    topk_weights = torch.randn(
-        num_tokens,
-        top_k,
-        device=device,
-        dtype=torch.float32,
-        generator=generator,
-    )
-    is_padding = torch.zeros(num_tokens, device=device, dtype=torch.bool)
-    is_padding[1::3] = True
-    is_padding[-1] = True
-    if nonfinite_padding:
-        topk_weights[is_padding] = float("nan")
-
-    ref_x, ref_x_sf = per_token_cast_to_fp8(
-        hidden_states,
-        use_ue8m0=True,
-        gran_k=32,
-        use_packed_ue8m0=True,
-    )
-    ref_topk_idx = topk_ids.to(torch.int64)
-    ref_topk_idx[is_padding] = -1
-    ref_topk_weights = topk_weights.clone()
-    ref_topk_weights[is_padding] = 0.0
-
-    fused_x = torch.empty_like(ref_x)
-    fused_x_sf = torch.empty_like(ref_x_sf)
-    fused_topk_idx = torch.empty_like(ref_topk_idx)
-    fused_topk_weights = torch.empty_like(ref_topk_weights)
-
-    prepare_megamoe_inputs(
-        hidden_states,
-        topk_weights,
-        topk_ids,
-        fused_x,
-        fused_x_sf,
-        fused_topk_idx,
-        fused_topk_weights,
-        is_padding=is_padding,
-    )
-    torch.accelerator.synchronize()
-
-    assert torch.equal(fused_x.view(torch.uint8), ref_x.view(torch.uint8))
-    assert torch.equal(fused_x_sf, ref_x_sf)
-    assert torch.equal(fused_topk_idx, ref_topk_idx)
-    assert torch.equal(
-        fused_topk_weights.view(torch.uint8),
-        ref_topk_weights.view(torch.uint8),
-    )

tests/models/test_registry.py

@@ -127,9 +127,9 @@ def test_registry_is_pp(model_arch, is_pp):
 @pytest.mark.parametrize(
     "model_arch,supported",
     [
-        ("Qwen4ExpForCausalLM", not current_platform.is_rocm()),
-        ("Qwen4ExpForConditionalGeneration", not current_platform.is_rocm()),
-        ("Qwen4ExpMTP", not current_platform.is_rocm()),
+        ("Qwen4ExpForCausalLM", True),
+        ("Qwen4ExpForConditionalGeneration", True),
+        ("Qwen4ExpMTP", True),
         ("LlamaForCausalLM", False),
     ],
 )

tests/parser/engine/test_parser_engine.py

@@ -161,12 +161,15 @@ class TestReasoningEndTokenIds:
         )
         engine = _make_engine(cfg)
         assert engine.reasoning_end_token_ids == {201, 202}
+        assert engine.implicit_reasoning_end_token_ids == {202}
         assert engine.find_reasoning_end([5, 201]) == ReasoningEnd(1, False)
         assert engine.find_reasoning_end([5, 202]) == ReasoningEnd(1, True)
 
     def test_transition_staying_in_reasoning_is_ignored(self):
         cfg = _with_reasoning_exits(("THINK_START", ParserState.REASONING, ()))
-        assert _make_engine(cfg).reasoning_end_token_ids == {201}
+        engine = _make_engine(cfg)
+        assert engine.reasoning_end_token_ids == {201}
+        assert engine.implicit_reasoning_end_token_ids == frozenset()
 
     def test_unreported_exit_fails_closed(self):
         cfg = _with_reasoning_exits(

tests/plugins/vllm_add_dummy_model/vllm_add_dummy_model/__init__.py

@@ -6,11 +6,16 @@ from vllm import ModelRegistry
 
 def register():
     # Test directly passing the model
-    from .my_opt import MyOPTForCausalLM
+    from .my_opt import MyOPTForCausalLM, MyOPTRawInputTokensForCausalLM
 
     if "MyOPTForCausalLM" not in ModelRegistry.get_supported_archs():
         ModelRegistry.register_model("MyOPTForCausalLM", MyOPTForCausalLM)
 
+    if "MyOPTRawInputTokensForCausalLM" not in ModelRegistry.get_supported_archs():
+        ModelRegistry.register_model(
+            "MyOPTRawInputTokensForCausalLM", MyOPTRawInputTokensForCausalLM
+        )
+
     # Test passing lazy model
     if "MyGemma2Embedding" not in ModelRegistry.get_supported_archs():
         ModelRegistry.register_model(

tests/plugins/vllm_add_dummy_model/vllm_add_dummy_model/my_opt.py

@@ -5,6 +5,7 @@
 import torch
 
 from vllm.model_executor.models.opt import OPTForCausalLM
+from vllm.sequence import IntermediateTensors
 
 
 class MyOPTForCausalLM(OPTForCausalLM):
@@ -15,3 +16,21 @@ class MyOPTForCausalLM(OPTForCausalLM):
             logits.zero_()
             logits[:, 0] += 1.0
         return logits
+
+
+class MyOPTRawInputTokensForCausalLM(MyOPTForCausalLM):
+    # Like models whose decoder layers read token ids, this model needs
+    # input_ids on every pipeline-parallel stage, not only the first.
+    requires_raw_input_tokens = True
+
+    def forward(
+        self,
+        input_ids: torch.Tensor | None,
+        positions: torch.Tensor,
+        intermediate_tensors: IntermediateTensors | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+    ) -> torch.Tensor | IntermediateTensors:
+        assert input_ids is not None, "requires_raw_input_tokens model got no input_ids"
+        return super().forward(
+            input_ids, positions, intermediate_tensors, inputs_embeds
+        )

tests/plugins_tests/test_oot_registration_online.py

@@ -7,7 +7,7 @@ chatml_jinja_path = VLLM_PATH / "examples/template_chatml.jinja"
 assert chatml_jinja_path.exists()
 
 
-def run_and_test_dummy_opt_api_server(model, tp=1):
+def run_and_test_dummy_opt_api_server(model, tp=1, extra_args=(), env_dict=None):
     # the model is registered through the plugin
     server_args = [
         "--gpu-memory-utilization",
@@ -20,8 +20,9 @@ def run_and_test_dummy_opt_api_server(model, tp=1):
         "dummy",
         "-tp",
         f"{tp}",
+        *extra_args,
     ]
-    with RemoteOpenAIServer(model, server_args) as server:
+    with RemoteOpenAIServer(model, server_args, env_dict=env_dict) as server:
         client = server.get_client()
         completion = client.chat.completions.create(
             model=model,

tests/quantization/test_humming_compile.py

@@ -0,0 +1,43 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import pytest
+import torch
+
+from vllm.platforms import current_platform
+from vllm.utils.import_utils import has_humming
+
+pytestmark = pytest.mark.skipif(
+    not (has_humming() and current_platform.is_cuda()),
+    reason="requires humming and CUDA",
+)
+
+
+def test_humming_input_quant_compiles_fullgraph():
+    """Humming's forward helpers look up ``DeviceInfo`` on every call; vLLM
+    traces them with ``fullgraph=True`` (vllm-project/humming#147)."""
+    import vllm.utils.humming as hm
+
+    weight_schema = hm.HummingWeightSchema(b_dtype="int4", weight_scale_group_size=128)
+    input_schema = hm.HummingInputSchema(a_dtype="int8")
+    from humming.transform import prepare_layer_config
+
+    config = prepare_layer_config(
+        shape_n=256,
+        shape_k=256,
+        weight_schema=weight_schema,
+        input_schema=input_schema,
+        torch_dtype=torch.bfloat16,
+        device=torch.device("cuda"),
+    )
+    x = torch.randn(16, 256, dtype=torch.bfloat16, device="cuda")
+
+    def quant(t: torch.Tensor):
+        return hm.may_process_input(config, t)
+
+    expected = quant(x)
+    actual = torch.compile(quant, fullgraph=True)(x)
+    for a, e in zip(actual, expected):
+        if e is None:
+            assert a is None
+        else:
+            torch.testing.assert_close(a, e)

tests/renderers/chat_utils/test_prompt_embeds.py

@@ -105,6 +105,7 @@ def _make_mock_model_config(*, enable_prompt_embeds: bool = True) -> mock.MagicM
     # `safe_load_prompt_embeds` pins each tensor to the model's hidden_size
     # and dtype, so the mock must return concrete values.
     mc.get_hidden_size.return_value = _MOCK_HIDDEN_SIZE
+    mc.hf_text_config = object()
     mc.dtype = _MOCK_DTYPE
     return mc
 

tests/renderers/test_completions.py

@@ -34,6 +34,7 @@ class MockModelConfig:
     tokenizer_revision = None
     tokenizer_mode = "auto"
     hf_config = MockHFConfig()
+    hf_text_config = hf_config
     encoder_config: dict[str, Any] | None = None
     enable_prompt_embeds: bool = True
     skip_tokenizer_init: bool = False

tests/renderers/test_prompt_embeds_payload_validation.py

@@ -65,6 +65,40 @@ def _saved_tensor(tensor: torch.Tensor) -> bytes:
     return buffer.getvalue()
 
 
+@pytest.mark.parametrize(
+    "model, hidden_size, input_size, buffer_size",
+    [
+        ("facebook/opt-350m", 1024, 512, 512),
+        ("facebook/opt-125m", 768, 768, 768),
+        ("openai/clip-vit-base-patch32", 512, 512, 1024),
+        ("google/siglip-base-patch16-224", 768, 768, 1024),
+    ],
+)
+def test_prompt_embeds_input_width(model, hidden_size, input_size, buffer_size):
+    """OPT uses its token width; dual encoders retain their text input width."""
+    from vllm.config import ModelConfig
+
+    config = ModelConfig(
+        model=model,
+        dtype="float32",
+        enable_prompt_embeds=True,
+    )
+    if "opt-" not in model:
+        config.hf_text_config.projection_dim = buffer_size
+    assert config.get_hidden_size() == hidden_size
+    assert config.get_inputs_embeds_size() == buffer_size
+
+    tensor = torch.randn(4, input_size)
+    loaded = safe_load_prompt_embeds(config, _encode(_saved_tensor(tensor)))
+    torch.testing.assert_close(loaded, tensor)
+
+    wrong_size = hidden_size if hidden_size != input_size else input_size + 1
+    with pytest.raises(VLLMValidationError, match="does not match"):
+        safe_load_prompt_embeds(
+            config, _encode(_saved_tensor(torch.randn(4, wrong_size)))
+        )
+
+
 UNPARSABLE_PAYLOADS = [
     pytest.param(b"", id="empty"),
     pytest.param(b"\x80", id="single-byte"),

tests/test_config.py

@@ -2132,8 +2132,8 @@ def test_draft_parallel_config_preserves_ep_without_model(
     [
         pytest.param("eagle3", False, 0, id="eagle3"),
         pytest.param("eagle3", True, 7, id="p-eagle"),
-        pytest.param("dflash", True, 8, id="dflash"),
-        pytest.param("dspark", True, 7, id="dspark"),
+        pytest.param("dflash", True, 0, id="dflash"),
+        pytest.param("dspark", True, 0, id="dspark"),
         pytest.param("mtp", False, 0, id="mtp"),
         pytest.param("ngram", False, 0, id="ngram"),
         pytest.param("draft_model", False, 1, id="draft-model"),
@@ -3769,34 +3769,48 @@ def test_dual_key_gumbel_requires_probabilistic_drafting():
         config._check_supports_watermarking()
 
 
-@pytest.mark.parametrize("method", ["eagle", "eagle3", "mtp"])
-def test_dual_key_gumbel_supports_probabilistic_speculative_decoding(method):
+@pytest.mark.parametrize(
+    ("method", "parallel_drafting"),
+    [
+        ("eagle", False),
+        ("eagle3", False),
+        ("mtp", False),
+        ("dspark", True),
+        ("dflash", True),
+    ],
+)
+def test_dual_key_gumbel_supports_probabilistic_speculative_decoding(
+    method, parallel_drafting
+):
     config = _watermarked_vllm_config()
     config.watermark_config = WatermarkConfig(algorithm="dual_key_gumbel", key=42)
     config.speculative_config = SimpleNamespace(
         method=method,
         draft_sample_method="probabilistic",
         rejection_sample_method="standard",
-        parallel_drafting=False,
+        parallel_drafting=parallel_drafting,
     )
 
     config._check_supports_watermarking()
 
 
-def test_dual_key_gumbel_supports_dspark():
+@pytest.mark.parametrize("architecture", ["DFlash2DraftModel", "LiLiCorrDraftModel"])
+def test_dual_key_gumbel_rejects_dflash_candidate_drafts(architecture):
     config = _watermarked_vllm_config()
     config.watermark_config = WatermarkConfig(algorithm="dual_key_gumbel", key=42)
     config.speculative_config = SimpleNamespace(
-        method="dspark",
+        method="dflash",
         draft_sample_method="probabilistic",
         rejection_sample_method="standard",
         parallel_drafting=True,
+        draft_model_config=SimpleNamespace(architectures=[architecture]),
     )
 
-    config._check_supports_watermarking()
+    with pytest.raises(ValueError, match="candidate-head drafters"):
+        config._check_supports_watermarking()
 
 
-def test_dual_key_gumbel_rejects_non_autoregressive_speculation():
+def test_dual_key_gumbel_rejects_unsupported_speculative_method():
     config = _watermarked_vllm_config()
     config.watermark_config = WatermarkConfig(algorithm="dual_key_gumbel", key=42)
     config.speculative_config = SimpleNamespace(
@@ -3806,7 +3820,7 @@ def test_dual_key_gumbel_rejects_non_autoregressive_speculation():
         parallel_drafting=False,
     )
 
-    with pytest.raises(ValueError, match="autoregressive model-based"):
+    with pytest.raises(ValueError, match="Watermarking supports only"):
         config._check_supports_watermarking()
 
 
@@ -4167,7 +4181,7 @@ def test_revision_resolved_for_model(mock_resolve):
 
 
 @pytest.mark.parametrize(
-    "enabled,supported,is_cuda,tp_size,pp_size,error",
+    "enabled,supported,cuda_alike,tp_size,pp_size,error",
     [
         (True, True, True, 2, 1, None),
         (True, False, True, 2, 1, "not supported for this model"),
@@ -4178,7 +4192,7 @@ def test_revision_resolved_for_model(mock_resolve):
         (False, False, False, 2, 1, None),
     ],
 )
-def test_hc_sp_model_support(enabled, supported, is_cuda, tp_size, pp_size, error):
+def test_hc_sp_model_support(enabled, supported, cuda_alike, tp_size, pp_size, error):
     """Validate active HC SP configurations and ignore the flag with TP=1."""
     from vllm.model_executor.models.registry import ModelRegistry
 
@@ -4196,7 +4210,10 @@ def test_hc_sp_model_support(enabled, supported, is_cuda, tp_size, pp_size, erro
         distributed_executor_backend="mp",
     )
     with (
-        patch("vllm.config.model.current_platform.is_cuda", return_value=is_cuda),
+        patch(
+            "vllm.config.model.current_platform.is_cuda_alike",
+            return_value=cuda_alike,
+        ),
         patch.object(ModelRegistry, "is_hc_sp_supported_model", return_value=supported),
         patch.object(ModelRegistry, "is_pp_supported_model", return_value=True),
     ):

tests/tokenizers_/test_basic.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from pathlib import Path
+from types import SimpleNamespace
 from typing import _get_protocol_attrs  # type: ignore
 from unittest.mock import patch
 
@@ -15,7 +16,11 @@ from transformers import (
 from vllm.tokenizers import TokenizerLike, get_tokenizer
 from vllm.tokenizers.hf import HfTokenizer
 from vllm.tokenizers.mistral import MistralTokenizer
-from vllm.tokenizers.registry import resolve_tokenizer_args
+from vllm.tokenizers.registry import (
+    cached_tokenizer_from_config,
+    resolve_tokenizer_args,
+    tokenizer_args_from_config,
+)
 
 
 def _get_missing_attrs(obj: object, target: type):
@@ -65,7 +70,11 @@ def test_tokenizer_revision(tokenizer_name: str):
 @pytest.mark.parametrize(
     "code_revision, error_match",
     [
-        pytest.param("never", "not a valid git identifier", id="unknown-revision"),
+        pytest.param(
+            "never",
+            "not a valid git identifier|couldn't find them in the cached files",
+            id="unknown-revision",
+        ),
         pytest.param(
             "73293ac504e43f4a895688b57c65cebe24f01a38",
             r"tokenization_fast\.py",
@@ -143,6 +152,141 @@ def test_tokenizer_code_revision_default_or_matching(code_revision: str | None):
     _assert_tokenizer_like(tokenizer)
 
 
+@pytest.mark.skip_global_cleanup
+@pytest.mark.parametrize(
+    "code_revision",
+    [
+        None,
+        "c8ae8a2e5845fc3ad4e5237e4346e5a9ba1e87d2",
+        "2efddce40dddaccd37bae208c3c7ca66dbedf68a",
+    ],
+)
+def test_tokenizer_config_code_revision(code_revision):
+    """Load real remote config code at its pin, not the data revision."""
+    revision = "c8ae8a2e5845fc3ad4e5237e4346e5a9ba1e87d2"
+
+    with patch(
+        "vllm.tokenizers.registry.CachedHfTokenizer.from_pretrained",
+        return_value=SimpleNamespace(is_fast=True),
+    ) as load_tokenizer:
+        get_tokenizer(
+            "hf-internal-testing/test_dynamic_model",
+            tokenizer_mode="hf",
+            trust_remote_code=True,
+            revision=revision,
+            code_revision=code_revision,
+        )
+
+    config = load_tokenizer.call_args.kwargs["config"]
+    assert (code_revision or revision) in type(config).__module__.split(".")
+    assert config.new_hidden_size == 12
+
+
+@pytest.mark.skip_global_cleanup
+@pytest.mark.parametrize(
+    "loader", (cached_tokenizer_from_config, tokenizer_args_from_config)
+)
+@pytest.mark.parametrize(
+    "model, code_revision, kwargs, error_match",
+    [
+        pytest.param(
+            None,
+            "never",
+            {},
+            "not a valid git identifier|couldn't find them in the cached files",
+            id="same-repo",
+        ),
+        pytest.param(
+            None,
+            "73293ac504e43f4a895688b57c65cebe24f01a38",
+            {},
+            r"tokenization_fast\.py",
+            id="same-repo-missing-code",
+        ),
+        pytest.param(
+            "facebook/opt-125m",
+            "27dcfa74d334bc871f3234de431e71c6eeba5dd6",
+            {},
+            None,
+            id="separate-repo",
+        ),
+        pytest.param(
+            None,
+            "never",
+            {"code_revision": "08a35780b7a78794436d8654aafa26ee0165c149"},
+            None,
+            id="explicit-revision",
+        ),
+        pytest.param(None, "never", {"code_revision": None}, None, id="explicit-none"),
+    ],
+)
+def test_config_tokenizer_code_revision(
+    loader, model, code_revision, kwargs, error_match
+):
+    """Only same-repo tokenizers inherit model code revisions; overrides win."""
+    revision = "08a35780b7a78794436d8654aafa26ee0165c149"
+    repo = "hf-internal-testing/test_dynamic_tokenizer"
+    config = SimpleNamespace(
+        model=model or repo,
+        tokenizer=repo,
+        runner_type="generate",
+        tokenizer_mode="hf",
+        tokenizer_revision=revision,
+        code_revision=code_revision,
+        trust_remote_code=True,
+        skip_tokenizer_init=False,
+    )
+
+    def load():
+        result = loader(config, **kwargs)
+        if loader is tokenizer_args_from_config:
+            mode, name, args, resolved_kwargs = result
+            return get_tokenizer(name, *args, tokenizer_mode=mode, **resolved_kwargs)
+        return result
+
+    if error_match is not None:
+        with pytest.raises(OSError, match=error_match):
+            load()
+    else:
+        tokenizer = load()
+        reference = AutoTokenizer.from_pretrained(
+            repo,
+            revision=revision,
+            code_revision=kwargs.get("code_revision"),
+            trust_remote_code=True,
+        )
+        assert tokenizer.encode("hello world") == reference.encode("hello world")
+        _assert_tokenizer_like(tokenizer)
+
+
+@pytest.mark.parametrize(
+    "loader", (cached_tokenizer_from_config, tokenizer_args_from_config)
+)
+@pytest.mark.parametrize("tokenizer_mode", ["mistral", "auto"])
+@pytest.mark.parametrize("code_revision", [None, "never"])
+def test_config_mistral_code_revision(loader, tokenizer_mode, code_revision):
+    """Config loading must not pass HF-only revisions to the Mistral backend."""
+    repo = "mistralai/Mistral-7B-Instruct-v0.3"
+    config = SimpleNamespace(
+        model=repo,
+        tokenizer=repo,
+        runner_type="generate",
+        tokenizer_mode=tokenizer_mode,
+        tokenizer_revision=None,
+        code_revision=code_revision,
+        trust_remote_code=False,
+        skip_tokenizer_init=False,
+    )
+    tokenizer = loader(config)
+    if loader is tokenizer_args_from_config:
+        mode, name, args, kwargs = tokenizer
+        tokenizer = get_tokenizer(name, *args, tokenizer_mode=mode, **kwargs)
+
+    reference = get_tokenizer(repo, tokenizer_mode="mistral")
+    assert isinstance(tokenizer, MistralTokenizer)
+    assert tokenizer.encode("hello world") == reference.encode("hello world")
+
+
 def test_hub_tokenizer_loads_from_local_snapshot():
     """Transformers makes a Hub request per file when given a repo ID."""
     with patch.object(

tests/tokenizers_/test_detokenize.py

@@ -15,6 +15,7 @@ from vllm.tokenizers.detokenizer_utils import (
 )
 from vllm.tokenizers.mistral import MistralTokenizer
 from vllm.tokenizers.protocol import TokenizerLike
+from vllm.transformers_utils.repo_utils import with_retry
 from vllm.v1.engine import EngineCoreRequest
 from vllm.v1.engine.detokenizer import (
     FastIncrementalDetokenizer,
@@ -95,10 +96,13 @@ def _run_incremental_decode(
 
 @pytest.fixture
 def tokenizer(tokenizer_name):
-    return (
-        MistralTokenizer.from_pretrained(tokenizer_name)
-        if "mistral" in tokenizer_name
-        else AutoTokenizer.from_pretrained(tokenizer_name)
+    return with_retry(
+        lambda: (
+            MistralTokenizer.from_pretrained(tokenizer_name)
+            if "mistral" in tokenizer_name
+            else AutoTokenizer.from_pretrained(tokenizer_name)
+        ),
+        log_msg=f"Error loading tokenizer {tokenizer_name}",
     )
 
 

tests/tokenizers_/test_registry.py

@@ -9,7 +9,7 @@ import pytest
 from transformers import AutoConfig, Qwen3_5MoeConfig
 from transformers.models.auto.configuration_auto import CONFIG_MAPPING
 
-from vllm.tokenizers import TokenizerLike
+from vllm.tokenizers import TokenizerLike, registry
 from vllm.tokenizers.registry import (
     TokenizerRegistry,
     cached_get_tokenizer,
@@ -17,6 +17,7 @@ from vllm.tokenizers.registry import (
     cached_tokenizer_from_config,
     get_tokenizer,
     resolve_tokenizer_args,
+    tokenizer_args_from_config,
 )
 
 
@@ -116,11 +117,13 @@ def test_cached_tokenizer_from_config_registers_local_config(tmp_path: Path):
     )
 
     model_config = SimpleNamespace(
+        model=str(tmp_path),
         skip_tokenizer_init=False,
         tokenizer=str(tmp_path),
         runner_type="generate",
         tokenizer_mode="hf",
         tokenizer_revision=None,
+        code_revision=None,
         trust_remote_code=True,
         hf_config=Qwen3_5MoeConfig(),
     )
@@ -164,3 +167,117 @@ def test_cached_tokenizer_from_config_registers_local_config(tmp_path: Path):
         CONFIG_MAPPING._extra_content.pop("qwen3_5_moe", None)
         if registered_config is not None:
             CONFIG_MAPPING._extra_content["qwen3_5_moe"] = registered_config
+
+
+@pytest.fixture
+def config_tokenizer_backend(monkeypatch):
+    """Keep config helpers and dispatch real; replace Hub/backend loading."""
+    calls = []
+
+    class Backend:
+        @classmethod
+        def from_pretrained(cls, name, *args, **kwargs):
+            calls.append(kwargs)
+            return SimpleNamespace(is_fast=True)
+
+    monkeypatch.setattr(registry, "get_config", lambda *a, **kw: None)
+    monkeypatch.setattr(registry, "_maybe_register_hf_config", lambda config: None)
+    monkeypatch.setattr(registry, "is_mistral_model_repo", lambda **kw: True)
+    monkeypatch.setattr(registry, "any_pattern_in_repo_files", lambda **kw: True)
+    monkeypatch.setattr(
+        registry, "get_non_weight_snapshot_path", lambda name, **kw: name
+    )
+    monkeypatch.setattr(TokenizerRegistry, "load_tokenizer_cls", lambda mode: Backend)
+    cached_get_tokenizer.cache_clear()
+    cached_resolve_tokenizer_args.cache_clear()
+    yield calls
+    cached_get_tokenizer.cache_clear()
+    cached_resolve_tokenizer_args.cache_clear()
+
+
+def _load_config_tokenizer(loader, config, **kwargs):
+    result = loader(config, **kwargs)
+    if loader is tokenizer_args_from_config:
+        mode, name, args, resolved_kwargs = result
+        return get_tokenizer(name, *args, tokenizer_mode=mode, **resolved_kwargs)
+    return result
+
+
+@pytest.mark.parametrize(
+    "loader", (tokenizer_args_from_config, cached_tokenizer_from_config)
+)
+@pytest.mark.parametrize("mode", ["mistral", "auto"])
+@pytest.mark.parametrize("code_revision", [None, "model-code"])
+@pytest.mark.parametrize("overrides", [{}, {"code_revision": None}])
+def test_config_mistral_does_not_forward_code_revision(
+    loader, mode, code_revision, overrides, config_tokenizer_backend
+):
+    """HF code revisions must not reach MistralCommonBackend's strict kwargs."""
+    config = SimpleNamespace(
+        model="test-model",
+        tokenizer="test-model",
+        runner_type="generate",
+        tokenizer_mode=mode,
+        tokenizer_revision=None,
+        code_revision=code_revision,
+        trust_remote_code=False,
+        skip_tokenizer_init=False,
+    )
+    _load_config_tokenizer(loader, config, **overrides)
+    assert len(config_tokenizer_backend) == 1
+    assert "code_revision" not in config_tokenizer_backend[0]
+
+
+@pytest.mark.parametrize(
+    "loader", (tokenizer_args_from_config, cached_tokenizer_from_config)
+)
+@pytest.mark.parametrize("missing", [False, True])
+def test_config_tokenizer_without_code_revision(
+    loader, missing, config_tokenizer_backend
+):
+    """Default or minimal configs must not introduce an HF-only plugin kwarg."""
+    config = SimpleNamespace(
+        model="test-model",
+        tokenizer="test-model",
+        runner_type="generate",
+        tokenizer_mode="test_tokenizer",
+        tokenizer_revision=None,
+        trust_remote_code=False,
+        skip_tokenizer_init=False,
+    )
+    if not missing:
+        config.code_revision = None
+    _load_config_tokenizer(loader, config)
+    assert len(config_tokenizer_backend) == 1
+    assert "code_revision" not in config_tokenizer_backend[0]
+
+
+@pytest.mark.parametrize(
+    "loader", (tokenizer_args_from_config, cached_tokenizer_from_config)
+)
+@pytest.mark.parametrize("mode", ["hf", "cohere", "kimi_k3", "inkling", "deepseek_v4"])
+@pytest.mark.parametrize(
+    "overrides, expected",
+    [
+        ({}, "model-code"),
+        ({"code_revision": None}, None),
+        ({"code_revision": "override"}, "override"),
+    ],
+)
+def test_config_hf_backed_tokenizer_preserves_code_revision(
+    loader, mode, overrides, expected, config_tokenizer_backend
+):
+    """Filtering Mistral kwargs must preserve HF wrappers and caller overrides."""
+    config = SimpleNamespace(
+        model="test-model",
+        tokenizer="test-model",
+        runner_type="generate",
+        tokenizer_mode=mode,
+        tokenizer_revision=None,
+        code_revision="model-code",
+        trust_remote_code=True,
+        skip_tokenizer_init=False,
+    )
+    _load_config_tokenizer(loader, config, **overrides)
+    assert len(config_tokenizer_backend) == 1
+    assert config_tokenizer_backend[0]["code_revision"] == expected

tests/tool_parsers/test_apertus_tool_parser.py

@@ -244,6 +244,22 @@ class TestStreamingExtraction:
         assert tcs[0]["name"] == "get_weather"
         assert json.loads(tcs[0]["arguments"]) == {"location": "Paris, France"}
 
+    @pytest.mark.parametrize(
+        "arguments",
+        ['{"value": -2e-3}', '{"value": 1.5e3}', '{"values": [1.5e3, -2e-3]}'],
+    )
+    def test_streaming_scientific_notation(self, parser, mock_request, arguments):
+        model_output = (
+            f'{TOOL_CALLS_PREFIX}[{{"calculate": {arguments}}}]{TOOL_CALLS_SUFFIX}'
+        )
+
+        results = self._simulate_streaming(parser, mock_request, list(model_output))
+        tool_calls = self._collect_tool_calls(results)
+
+        assert len(tool_calls) == 1
+        assert tool_calls[0]["name"] == "calculate"
+        assert json.loads(tool_calls[0]["arguments"]) == json.loads(arguments)
+
     def test_streaming_missing_tool_suffix(self, parser, mock_request):
         chunks = [
             "<|tools_prefix|>",