Merge remote-tracking branch 'upstream/main' into patches

.buildkite/test_areas/misc.yaml

@@ -655,12 +655,14 @@ steps:
   source_file_dependencies:
     - vllm/v1/attention
     - vllm/model_executor/layers
+    - vllm/lora/ops/triton_ops
     - tests/v1/determinism/
   commands:
     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
     - pip install pytest-timeout pytest-forked
     - pytest -v -s v1/determinism/test_batch_invariance.py
     - pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py
+    - pytest -v -s v1/determinism/test_lora_shrink_batch_invariant.py
     - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_MLA
     - VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k FLASH_ATTN
   mirror:
@@ -690,12 +692,14 @@ steps:
   source_file_dependencies:
     - vllm/v1/attention
     - vllm/model_executor/layers
+    - vllm/lora/ops/triton_ops
     - tests/v1/determinism/
   commands:
     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
     - pip install pytest-timeout pytest-forked
     - pytest -v -s v1/determinism/test_batch_invariance.py
     - pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py
+    - pytest -v -s v1/determinism/test_lora_shrink_batch_invariant.py
     - pytest -v -s v1/determinism/test_batch_invariance_vlm.py
     - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_MLA
     - VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k FLASH_ATTN

.github/mergify.yml

@@ -112,6 +112,32 @@ pull_request_rules:
       add:
         - frontend
 
+- name: label-pooling
+  description: Automatically apply pooling label
+  conditions:
+    - -closed
+    - or:
+      - files~=^vllm/entrypoints/pooling/
+      - files~=^vllm/model_executor/layers/pooler/
+      - files=vllm/model_executor/models/transformers/pooling.py
+      - files=vllm/model_executor/models/adapters.py
+      - files=vllm/config/pooler.py
+      - files=vllm/pooling_params.py
+      - files~=^vllm/v1/pool/
+      - files~=^vllm/v1/worker/gpu/pool/
+      - files~=^tests/entrypoints/pooling/
+      - files~=^tests/models/(language|multimodal)/pooling(_mteb_test)?/
+      - files~=^tests/model_executor/layers/test_pooler_.*\.py$
+      - files=tests/test_pooling_params.py
+      - files~=^tests/v1/e2e/general/test_pooling_.*\.py$
+      - files~=^docs/models/pooling_models/
+      - files~=^examples/pooling/
+      - title~=(?i)\bpooling\b
+  actions:
+    label:
+      add:
+        - pooling
+
 - name: label-rust
   description: Automatically apply rust label
   conditions:

.github/workflows/issue_autolabel.yml

@@ -255,6 +255,27 @@ jobs:
               // phrases go in `substrings`: they are specific enough on their
               // own, and a boundary match would miss ordinary variations --
               // "structured outputs", "tool parsers", "prefix caching".
+              pooling: {
+                keywords: [
+                  { term: "pooling",   searchIn: "title" },
+                  { term: "pooler",    searchIn: "title" },
+                  { term: "rerank",    searchIn: "title" },
+                  { term: "reranker",  searchIn: "title" },
+                  { term: "reranking", searchIn: "title" },
+                ],
+                substrings: [
+                  { term: "embedding model",         searchIn: "title" },
+                  { term: "embedding API",           searchIn: "title" },
+                  { term: "embeddings API",          searchIn: "title" },
+                  { term: "sentence embedding",      searchIn: "title" },
+                  { term: "sequence classification", searchIn: "title" },
+                  { term: "token classification",    searchIn: "title" },
+                  { term: "cross-encoder",           searchIn: "title" },
+                  { term: "late interaction",        searchIn: "title" },
+                  { term: "late-interaction",        searchIn: "title" },
+                  { term: "reward model",            searchIn: "title" },
+                ],
+              },
               "speculative-decoding": {
                 keywords: [
                   { term: "eagle",  searchIn: "title" },

csrc/rocm/moe_q_gemm_rdna3.cu

@@ -57,16 +57,6 @@ __forceinline__ __device__ bf16_t tzero<bf16_t>() {
   return __float2bfloat16(0.0f);
 }
 
-__forceinline__ __device__ float dot22_8_f(half2 (&dq)[4], const half* a_ptr) {
-  float result = 0.0f;
-  const half2* a2_ptr = (const half2*)a_ptr;
-  #pragma unroll
-  for (int i = 0; i < 4; i++) {
-    result = __builtin_amdgcn_fdot2(dq[i], *a2_ptr++, result, /*clamp=*/false);
-  }
-  return result;
-}
-
 __forceinline__ __device__ void atomic_add_pk4_f16(half* addr, half2 v01,
                                                    half2 v23) {
   unsigned long long* addr_u = reinterpret_cast<unsigned long long*>(addr);
@@ -177,7 +167,8 @@ __global__ void moe_gemm_q4_kernel_rdna3(
                 "BLOCK_KN_SIZE must equal THREADS_X");
 
   // For bf16 M=1, we can skip LDS and read A from global (same as dense).
-  // fp16 always needs LDS due to the dot22_8_f indexing pattern.
+  // fp16 always needs LDS: the exact-factored inner loop indexes
+  // block_a[m][a_off] directly.
   constexpr bool USE_LDS_A = (BLOCK_SIZE_M > 1) || std::is_same<T, half>::value;
 
   const int offset_m_base = token_block * BLOCK_SIZE_M;
@@ -211,8 +202,10 @@ __global__ void moe_gemm_q4_kernel_rdna3(
   int qk = offset_k / 8;
   const uint32_t* b_ptr = expert_weights + qk * size_n + n;
 
-  // Per-column dequant constants (4 columns per thread)
+  // Per-column dequant constants (4 columns per thread). fp16 uses the
-  half2 z1z16_h[4][2], y1y16_h[4][2];
+  // exact factored form (fp32 y/z pairs, no baked bias — see
+  // qdq_4_rdna3.cuh); bf16 keeps fp32 scalars as before.
+  float yf_h[4][2], zf_h[4][2];
   float z_b_f[4], y_b_f[4];
 
   // GPTQv1: zero_offset = 1
@@ -228,8 +221,8 @@ __global__ void moe_gemm_q4_kernel_rdna3(
     if constexpr (std::is_same<T, half>::value) {
   #pragma unroll
       for (int i = 0; i < 4; ++i) {
-        gptq_rdna3::prep_zero_scale_fp16((uint32_t)(zeros[i] + zero_offset),
+        gptq_rdna3::prep_zero_scale_fp16_f32((uint32_t)(zeros[i] + zero_offset),
-                                         scales[i], z1z16_h[i], y1y16_h[i]);
+                                             scales[i], yf_h[i], zf_h[i]);
       }
     } else {
   #pragma unroll
@@ -266,29 +259,42 @@ __global__ void moe_gemm_q4_kernel_rdna3(
     }
     b_ptr += 4 * size_n;
 
+    // fp16 exact factored form: activation sums per row, split by nibble
+    // slot, accumulated across the 32-K round; the z correction is applied
+    // once after the j loop (see qdq_4_rdna3.cuh).
+    float sum_lo[BLOCK_SIZE_M] = {}, sum_hi[BLOCK_SIZE_M] = {};
+
   #pragma unroll
     for (int j = 0; j < 4; ++j) {
       const int a_off = (k - offset_k) + 8 * j;
 
       if constexpr (std::is_same<T, half>::value) {
-        // fp16 path: dequant via bit-trick, dot via v_dot2_f32_f16
+        // fp16 path: magic values without bias, dots via v_dot2_f32_f16,
-        half2 dq[4][4];
+        // the (y, z) correction applied in fp32 (qdq_4_rdna3.cuh).
-        gptq_rdna3::dequant_4bit_8_fp16((uint32_t)b_w[j].x, dq[0], z1z16_h[0],
+        half2 qm[4][4];
-                                        y1y16_h[0]);
+        gptq_rdna3::magic_4bit_8_fp16((uint32_t)b_w[j].x, qm[0]);
-        gptq_rdna3::dequant_4bit_8_fp16((uint32_t)b_w[j].y, dq[1], z1z16_h[1],
+        gptq_rdna3::magic_4bit_8_fp16((uint32_t)b_w[j].y, qm[1]);
-                                        y1y16_h[1]);
+        gptq_rdna3::magic_4bit_8_fp16((uint32_t)b_w[j].z, qm[2]);
-        gptq_rdna3::dequant_4bit_8_fp16((uint32_t)b_w[j].z, dq[2], z1z16_h[2],
+        gptq_rdna3::magic_4bit_8_fp16((uint32_t)b_w[j].w, qm[3]);
-                                        y1y16_h[2]);
+
-        gptq_rdna3::dequant_4bit_8_fp16((uint32_t)b_w[j].w, dq[3], z1z16_h[3],
+        const half2 ones = gptq_rdna3::ones_half2_fp16();
-                                        y1y16_h[3]);
 
   #pragma unroll
         for (int m = 0; m < BLOCK_SIZE_M; ++m) {
-          const half* a_ptr = reinterpret_cast<const half*>(&block_a[m][a_off]);
+          const half2* a2 = reinterpret_cast<const half2*>(&block_a[m][a_off]);
-          block_c[m][0] += dot22_8_f(dq[0], a_ptr);
+          sum_lo[m] = __builtin_amdgcn_fdot2(a2[0], ones, sum_lo[m], false);
-          block_c[m][1] += dot22_8_f(dq[1], a_ptr);
+          sum_hi[m] = __builtin_amdgcn_fdot2(a2[1], ones, sum_hi[m], false);
-          block_c[m][2] += dot22_8_f(dq[2], a_ptr);
+          sum_lo[m] = __builtin_amdgcn_fdot2(a2[2], ones, sum_lo[m], false);
-          block_c[m][3] += dot22_8_f(dq[3], a_ptr);
+          sum_hi[m] = __builtin_amdgcn_fdot2(a2[3], ones, sum_hi[m], false);
+  #pragma unroll
+          for (int c = 0; c < 4; ++c) {
+            float pl = 0.0f, ph = 0.0f;
+            pl = __builtin_amdgcn_fdot2(qm[c][0], a2[0], pl, /*clamp=*/false);
+            ph = __builtin_amdgcn_fdot2(qm[c][1], a2[1], ph, /*clamp=*/false);
+            pl = __builtin_amdgcn_fdot2(qm[c][2], a2[2], pl, /*clamp=*/false);
+            ph = __builtin_amdgcn_fdot2(qm[c][3], a2[3], ph, /*clamp=*/false);
+            block_c[m][c] += yf_h[c][0] * pl + yf_h[c][1] * ph;
+          }
         }
       } else if constexpr (BLOCK_SIZE_M == 1) {
         // bf16 M=1: v_dot2_f32_bf16 with InstCombine-defeating opacity
@@ -415,6 +421,18 @@ __global__ void moe_gemm_q4_kernel_rdna3(
         }
       }
     }
+
+    if constexpr (std::is_same<T, half>::value) {
+      // z is a group constant and sum_a is complete for this round: apply
+      // the fp16 z correction once per round (see qdq_4_rdna3.cuh).
+  #pragma unroll
+      for (int m = 0; m < BLOCK_SIZE_M; ++m) {
+  #pragma unroll
+        for (int c = 0; c < 4; ++c) {
+          block_c[m][c] += zf_h[c][0] * sum_lo[m] + zf_h[c][1] * sum_hi[m];
+        }
+      }
+    }
     k += 32;
   }
 
@@ -587,6 +605,10 @@ void moe_gptq_gemm_rdna3(torch::Tensor a, torch::Tensor c,
   int size_k = (int)a.size(1);
   int size_n = (int)b_q_weight.size(2);
   int groups = (int)b_scales.size(1);
+  TORCH_CHECK(size_k % (groups * 32) == 0,
+              "group size (K/groups = ", size_k / groups,
+              ") must be a multiple of 32: the kernel checks group "
+              "transitions at 32-K granularity");
 
   // Per-expert strides
   int expert_weight_stride = (int)(b_q_weight.size(1) * b_q_weight.size(2));

csrc/rocm/q_gemm_rdna3.cu

@@ -5,31 +5,42 @@
 // activation dtype (half or __hip_bfloat16). Adapted from exllamav2's 4-bit
 // kernel (csrc/quantization/gptq/q_gemm.cu) with the following changes:
 //
-//   1. Direct write to the T-typed output via packed CAS-loop on a 64-bit
+//   1. Deterministic split-K epilogue. K is split across gridDim.z blocks;
-//      word (atomic_add_pk4_{f16,bf16}). gfx11 has no native
+//      each block stores its FP32 block partial to a scratch tensor, and the
-//      v_global_atomic_pk_add_{f16,bf16}, so the kernel emulates one with
+//      last z block of each (x, y) tile reduces the z-slices in fixed
-//      global_atomic_cmpswap_b64. This avoids the M*N*4-byte FP32 scratch
+//      ascending order with a single final cast to T — the reduction runs
-//      buffer + memset + cast-pass that an fp32-accumulator design would
+//      in the GEMM kernel itself, gated by a per-call ticket counter (no
-//      need; the caller passes a zero-initialised T-typed output tensor
+//      cross-launch state; see the epilogue). When gridDim.z == 1 each
-//      and every block atomically adds its partial sum into it.
+//      output element has exactly one writer and the kernel stores the
+//      rounded accumulator directly — no scratch, no reduce, no atomics.
+//      (Legacy design, pre-#54706: a packed CAS-loop atomic add on a 64-bit
+//      word emulating v_global_atomic_pk_add_{f16,bf16}, which gfx11 lacks.
+//      It narrowed every split partial to bf16/fp16 BEFORE accumulation, so
+//      the execution-dependent CAS completion order changed the rounded
+//      result for identical inputs. The CAS epilogue is kept below the
+//      partials branch for A/B comparison only; it requires a pre-zeroed
+//      output tensor.)
 //
 //   2. The bf16 path uses a dedicated bit-trick that avoids the fp16-only
 //      "upper nibble * 16" trick, which would overflow the 7-bit bf16
-//      mantissa. See qdq_4_rdna3.cuh for details.
+//      mantissa. See qdq_4_rdna3.cuh for details. The fp16 path uses the
+//      exact factored dequant (also in qdq_4_rdna3.cuh): magic values in
+//      fp16, the (y, z) bias correction in fp32.
 //
 //   3. Wave32 geometry sized for high CU saturation: THREADS_X=256
 //      (8 waves per block) and BLOCK_KN_SIZE=256, with each thread
 //      computing 4 N output columns. gridDim.z = K / BLOCK_KN_SIZE
-//      splits K and the output is atomically accumulated. fp16 uses
+//      splits K; the epilogue stores FP32 partials for the fixed-order
+//      reduce (see item 1). fp16 uses
 //      v_dot2_f32_f16 (__builtin_amdgcn_fdot2) for the inner dot;
 //      bf16 widens to fp32 (no v_pk_fma_bf16 on gfx11) and accumulates
 //      with v_fma_f32. M_COUNT ∈ {1,2,4,8} is selected at launch
 //      based on size_m.
 //
-//   4. The bf16 dispatch with M >= 16 forwards to the WMMA kernel in
+//   4. The dispatch with M >= 12 forwards to the WMMA kernel in
 //      q_gemm_rdna3_wmma.cu (separate translation unit) where
-//      v_wmma_f32_16x16x16_bf16_w32 wins. The fp16 path always stays
+//      v_wmma_f32_16x16x16_bf16_w32 wins. Below that the scalar path wins
-//      scalar (the bit-trick dequant beats WMMA below M=64).
+//      and is the more accurate of the two.
 
 #include <cstdint>
 #include <cstdio>
@@ -92,40 +103,16 @@ __forceinline__ __device__ bf16_t tzero<bf16_t>() {
   return __float2bfloat16(0.0f);
 }
 
-__forceinline__ __device__ float dot22_8_f(half2 (&dq)[4], const half* a_ptr) {
-  // RDNA3 has v_dot2_f32_f16 (`__builtin_amdgcn_fdot2`) which computes
-  // fp32 += a.x*b.x + a.y*b.y in a single instruction with the accumulator
-  // staying in fp32 throughout. hipcc 7.2 does NOT peephole the obvious
-  // `__hfma2 + cast + add` pattern into v_dot2 (verified by ISA
-  // disassembly: 0 v_dot2_f32_f16 vs 256 v_cvt_f32_f16 + 218 v_add_f32 in
-  // the M_COUNT=8 kernel before this change), so we issue the builtin
-  // explicitly. Saves the trailing 2× v_cvt_f32_f16 + v_add_f32 (3 ops)
-  // per dot22_8_f call vs the half2-accumulator form. With 128 calls per
-  // K=32 step that's ~384 ops/K-step less issue pressure on the VALU.
-  //
-  // Numerical bonus: accumulator stays fp32 throughout the dot. The old
-  // form accumulated 8 muladds in fp16 (10-bit mantissa) before casting,
-  // which could lose ~3 bits of precision on borderline magnitudes.
-  float result = 0.0f;
-  const half2* a2_ptr = (const half2*)a_ptr;
-  #pragma unroll
-  for (int i = 0; i < 4; i++) {
-    result = __builtin_amdgcn_fdot2(dq[i], *a2_ptr++, result, /*clamp=*/false);
-  }
-  return result;
-}
-
 // ---------------------------------------------------------------------------
 // Packed atomic-add via CAS-loop on a 64-bit word (4 fp16/bf16 lanes per CAS).
 // RDNA3 (gfx11) does NOT have native v_global_atomic_pk_add_f16 / _bf16 (those
 // landed on gfx940 / gfx1250 respectively), so this lowers to
-// global_atomic_cmpswap_b64 plus retry. We use this in the kernel epilogue to
+// global_atomic_cmpswap_b64 plus retry. This is the LEGACY (pre-#54706)
-// write 4 output columns per row in a single atomic operation — half the
+// epilogue: the caller must pre-zero the output, and every split block
-// atomic instruction count and half the contention vs two 32-bit CAS calls.
+// atomically adds its low-precision partial into it. Because the addition
-//
+// happens in bf16/fp16 AFTER narrowing, the result depends on CAS completion
-// Writing directly to fp16/bf16 (instead of through an FP32 scratch buffer +
+// order — the nondeterminism PR #54706 fixes. Kept for A/B comparison with
-// cast pass) saves M*N*4 bytes of allocation, the memset, and the epilogue
+// the deterministic path; the shipped dispatch never selects it.
-// cast pass that an fp32-accumulator design would need.
 //
 // 64-bit alignment: the kernel writes at `out + n` where n = offset_n + t*4
 // (always multiple of 4), and partition_weight_shape[1] is required to be a
@@ -198,13 +185,26 @@ __forceinline__ __device__ void load4_scales(const T* scales_row, int n,
 // ---------------------------------------------------------------------------
 
 template <typename T, int M_COUNT>
-__global__ void gemm_q4_kernel_rdna3(const T* __restrict__ a,
+__global__ void gemm_q4_kernel_rdna3(
-                                     const uint32_t* __restrict__ b_q_weight,
+    const T* __restrict__ a, const uint32_t* __restrict__ b_q_weight,
-                                     const uint32_t* __restrict__ b_qzeros,
+    const uint32_t* __restrict__ b_qzeros, const T* __restrict__ b_scales,
-                                     const T* __restrict__ b_scales,
+    T* __restrict__ c, const int size_m, const int size_n, const int size_k,
-                                     T* __restrict__ c, const int size_m,
+    const int groups, const int zero_offset,
-                                     const int size_n, const int size_k,
+    // Deterministic split-K epilogue: when
-                                     const int groups, const int zero_offset) {
+    // non-null, each split block stores its
+    // FP32 partial to
+    // partials[(z*size_m + m)*size_n + n]
+    // instead of CAS-atomically accumulating
+    // a low-precision partial into c. The
+    // last z block of each (x, y) tile then
+    // performs the fixed ascending-z
+    // reduction in-kernel (see the epilogue
+    // below). `tickets` (same length as the
+    // number of tiles) is a per-call
+    // zero-filled arrival counter array —
+    // one ticket per tile — so concurrent
+    // invocations never share state.
+    float* __restrict__ partials, int* __restrict__ tickets) {
   const int t = threadIdx.x;
   const int offset_n = blockIdx.x * BLOCK_KN_SIZE * 4;
   const int offset_m = blockIdx.y * M_COUNT;
@@ -270,27 +270,24 @@ __global__ void gemm_q4_kernel_rdna3(const T* __restrict__ a,
   int qk = offset_k / 8;
   const uint32_t* b_ptr = b_q_weight + qk * size_n + n;
 
-  // Per-column dequant constants. We hold one set of (z, y) pairs per column.
+  // Per-column dequant constants, fp32 for both dtypes. fp16 needs two (y, z)
-  // fp16 uses the exllama (z1z16, y1y16) double-pair to enable the upper-
+  // pairs because the upper-nibble trick carries a factor of 16; bf16 needs one
-  // nibble-*16 trick. bf16 uses fp32 scalars (z, y) for the inline v_dot2
+  // because its dequant produces fp32 directly — see prep_zero_scale_bf16_f32 /
-  // bias correction.
+  // the FMA bypass for the missing v_pk_fma_bf16 on gfx11. Neither narrows the
-  half2 z1z16_h[4][2], y1y16_h[4][2];
+  // bias to the activation dtype: see the exact factored form in
+  // qdq_4_rdna3.cuh for why that matters.
   float z_b_f[4], y_b_f[4];
+  float yf_h[4][2], zf_h[4][2];
 
+  // bf16-only group refresh (fp16 preps from prefetched raw words below).
   auto refresh_group = [&](int g) {
-    const uint32_t* qz_row = b_qzeros + g * (size_n / 8);
+    if constexpr (!std::is_same<T, half>::value) {
-    const T* sc_row = b_scales + g * size_n;
+      const uint32_t* qz_row = b_qzeros + g * (size_n / 8);
-    int zeros[4];
+      const T* sc_row = b_scales + g * size_n;
-    T scales[4];
+      int zeros[4];
-    load4_zeros(qz_row, n, zeros);
+      T scales[4];
-    load4_scales<T>(sc_row, n, scales);
+      load4_zeros(qz_row, n, zeros);
-    if constexpr (std::is_same<T, half>::value) {
+      load4_scales<T>(sc_row, n, scales);
-  #pragma unroll
-      for (int i = 0; i < 4; ++i) {
-        prep_zero_scale_fp16((uint32_t)(zeros[i] + zero_offset), scales[i],
-                             z1z16_h[i], y1y16_h[i]);
-      }
-    } else {
   #pragma unroll
       for (int i = 0; i < 4; ++i) {
         prep_zero_scale_bf16_f32((uint32_t)(zeros[i] + zero_offset), scales[i],
@@ -299,7 +296,32 @@ __global__ void gemm_q4_kernel_rdna3(const T* __restrict__ a,
     }
   };
 
-  refresh_group(group);
+  // fp16: the raw zero/scale words of the NEXT group are loaded during the
+  // round before the boundary, so a group change no longer stalls on a full
+  // memory round trip before the weights of the round are even issued.
+  uint32_t rz = 0;
+  T rs[4];
+  auto load_raw_group = [&](int g) {
+    rz = b_qzeros[g * (size_n / 8) + n / 8] >> ((n & 7) * 4);
+  #pragma unroll
+    for (int i = 0; i < 4; ++i) rs[i] = b_scales[g * size_n + n + i];
+  };
+  auto prep_raw_fp16 = [&]() {
+    if constexpr (std::is_same<T, half>::value) {
+  #pragma unroll
+      for (int i = 0; i < 4; ++i) {
+        prep_zero_scale_fp16_f32(((rz >> (4 * i)) & 0xF) + zero_offset, rs[i],
+                                 yf_h[i], zf_h[i]);
+      }
+    }
+  };
+
+  if constexpr (std::is_same<T, half>::value) {
+    load_raw_group(group);
+    prep_raw_fp16();
+  } else {
+    refresh_group(group);
+  }
 
   float block_c[M_COUNT][4];
   #pragma unroll
@@ -309,11 +331,12 @@ __global__ void gemm_q4_kernel_rdna3(const T* __restrict__ a,
   }
 
   // Note on group-transition granularity: we check `k == nextgroup` at the
-  // start of each outer iteration (which advances K by 32). This is correct
+  // start of each outer iteration (which advances K by 32), so a group
-  // when group_size >= 32 OR group_size divides 32 evenly (groupsize is one
+  // boundary must never fall strictly inside a 32-K round: groupsize must
-  // of {1,2,4,8,16,32,64,128,...}). For group_size in {16, 8, 4, ...} the
+  // be a multiple of 32 (all real GPTQ group sizes 32/64/128/256, or
-  // inner loop would cross a group boundary between j-iterations; we require
+  // groups == 1). groupsize in {16, 8, ...} or non-multiples like 48 would
-  // group_size >= 32 here, mirroring exllama's assumption.
+  // skip the boundary check and silently use stale constants; the host-side
+  // TORCH_CHECK below enforces it. Same assumption as exllama.
   //
   // Software pipelining: we issue all 4 vectorized weight loads up front
   // before any dequant/FMA depends on them. This gives the AMDGPU backend
@@ -325,7 +348,11 @@ __global__ void gemm_q4_kernel_rdna3(const T* __restrict__ a,
     if (k == nextgroup) {
       group++;
       nextgroup += groupsize;
-      refresh_group(group);
+      if constexpr (std::is_same<T, half>::value) {
+        prep_raw_fp16();  // from the words prefetched in the previous round
+      } else {
+        refresh_group(group);
+      }
     }
 
     // Prefetch all four j-iterations' weight words. The compiler emits 4
@@ -337,25 +364,55 @@ __global__ void gemm_q4_kernel_rdna3(const T* __restrict__ a,
       b_w[j] = *(const int4*)(b_ptr + j * size_n);
     }
     b_ptr += 4 * size_n;
+    if constexpr (std::is_same<T, half>::value) {
+      if (k + 32 == nextgroup && k + 32 < end_k) load_raw_group(group + 1);
+    }
+
+    // fp16 exact factored form: raw dot products and activation sums, both
+    // split by nibble slot, accumulated across the whole 32-K round. Neither
+    // depends on the output column, and the (y, z) constants are group
+    // constants, so both corrections are applied once after the j loop
+    // instead of four times inside it — see qdq_4_rdna3.cuh.
+    float sum_lo[M_COUNT] = {}, sum_hi[M_COUNT] = {};
+    float acc_lo[M_COUNT][4] = {}, acc_hi[M_COUNT][4] = {};
 
   #pragma unroll
     for (int j = 0; j < 4; ++j) {
       const int a_off = (k - offset_k) + 8 * j;
 
       if constexpr (std::is_same<T, half>::value) {
-        half2 dq[4][4];
+        // Magic values with no bias folded in; the (y, z) correction is
-        dequant_4bit_8_fp16((uint32_t)b_w[j].x, dq[0], z1z16_h[0], y1y16_h[0]);
+        // applied in fp32 below — see the exact factored form in
-        dequant_4bit_8_fp16((uint32_t)b_w[j].y, dq[1], z1z16_h[1], y1y16_h[1]);
+        // qdq_4_rdna3.cuh.
-        dequant_4bit_8_fp16((uint32_t)b_w[j].z, dq[2], z1z16_h[2], y1y16_h[2]);
+        half2 qm[4][4];
-        dequant_4bit_8_fp16((uint32_t)b_w[j].w, dq[3], z1z16_h[3], y1y16_h[3]);
+        magic_4bit_8_fp16((uint32_t)b_w[j].x, qm[0]);
+        magic_4bit_8_fp16((uint32_t)b_w[j].y, qm[1]);
+        magic_4bit_8_fp16((uint32_t)b_w[j].z, qm[2]);
+        magic_4bit_8_fp16((uint32_t)b_w[j].w, qm[3]);
+
+        const half2 ones = ones_half2_fp16();
 
   #pragma unroll
         for (int m = 0; m < M_COUNT; ++m) {
-          const half* a_ptr = reinterpret_cast<const half*>(&block_a[m][a_off]);
+          const half2* a2 = reinterpret_cast<const half2*>(&block_a[m][a_off]);
-          block_c[m][0] += dot22_8_f(dq[0], a_ptr);
+          // sum_a split by nibble slot (the high pairs carry a factor of
-          block_c[m][1] += dot22_8_f(dq[1], a_ptr);
+          // 16). It does not depend on the output column, so one v_dot2 per
-          block_c[m][2] += dot22_8_f(dq[2], a_ptr);
+          // pair is shared by the four columns this thread owns.
-          block_c[m][3] += dot22_8_f(dq[3], a_ptr);
+          sum_lo[m] = __builtin_amdgcn_fdot2(a2[0], ones, sum_lo[m], false);
+          sum_hi[m] = __builtin_amdgcn_fdot2(a2[1], ones, sum_hi[m], false);
+          sum_lo[m] = __builtin_amdgcn_fdot2(a2[2], ones, sum_lo[m], false);
+          sum_hi[m] = __builtin_amdgcn_fdot2(a2[3], ones, sum_hi[m], false);
+  #pragma unroll
+          for (int c = 0; c < 4; ++c) {
+            acc_lo[m][c] =
+                __builtin_amdgcn_fdot2(qm[c][0], a2[0], acc_lo[m][c], false);
+            acc_hi[m][c] =
+                __builtin_amdgcn_fdot2(qm[c][1], a2[1], acc_hi[m][c], false);
+            acc_lo[m][c] =
+                __builtin_amdgcn_fdot2(qm[c][2], a2[2], acc_lo[m][c], false);
+            acc_hi[m][c] =
+                __builtin_amdgcn_fdot2(qm[c][3], a2[3], acc_hi[m][c], false);
+          }
         }
       } else if constexpr (M_COUNT == 1) {
         // bf16 decode (M=1), v_dot2_f32_bf16 path. Mirrors the data-flow of
@@ -520,18 +577,84 @@ __global__ void gemm_q4_kernel_rdna3(const T* __restrict__ a,
         }
       }
     }
+
+    if constexpr (std::is_same<T, half>::value) {
+      // y and z are group constants and this round's dots and activation
+      // sums are complete: one pair of FMAs per (row, column) per round
+      // instead of per j-iteration. Per-round z applications sum exactly to
+      // the per-group correction because z does not change inside a group.
+  #pragma unroll
+      for (int m = 0; m < M_COUNT; ++m) {
+  #pragma unroll
+        for (int c = 0; c < 4; ++c) {
+          block_c[m][c] += yf_h[c][0] * acc_lo[m][c] +
+                           yf_h[c][1] * acc_hi[m][c] + zf_h[c][0] * sum_lo[m] +
+                           zf_h[c][1] * sum_hi[m];
+        }
+      }
+    }
     k += 32;  // 4 weight words * 8 nibbles = 32 K elements
   }
 
-  // Pack the 4 FP32 partial sums into 2 packed pairs and atomically add all
+  // ---- Epilogue: three store modes, selected by launch shape ----
-  // four lanes in a single 64-bit CAS write directly to the T-typed output
+  //
-  // (caller pre-zeros it). On gfx11 the packed atomic is a CAS-loop, but with
+  // partials != nullptr: deterministic split-K. Plain FP32 stores to
-  // a single b64 op we halve the atomic instruction count vs two b32 CAS
+  //   partials[(blockIdx.z * size_m + m) * size_n + n]
-  // calls, AND save the FP32 buffer + memset + cast pass entirely.
+  // (one writer per in-range slot; threads with n >= size_n returned before
+  // the epilogue and rows past size_m are skipped — exactly the slots the
+  // in-kernel fixed-order reduce below never reads). n is a multiple of 4
+  // and size_n % 8 == 0, so the 4-lane store never crosses the right edge.
+  // After the stores, the block publishes its arrival on the tile's ticket
+  // (a per-call zero-filled counter); the LAST z block of the tile then
+  // sums the z slices in fixed ascending order with a single final cast —
+  // the reduction order is a pure function of the launch shape, so the
+  // result stays bit-reproducible for identical inputs. Tickets are
+  // per-invocation state (see launch_gemm_q4_deterministic); nothing is
+  // shared across kernel launches, so concurrent invocations on other
+  // streams cannot interfere.
+  //
+  // partials == nullptr && gridDim.z == 1: single split block per output
+  //   element — direct store of the rounded accumulator. No scratch, no
+  //   reduce, no atomics; c may be left uninitialized (torch::empty).
+  //
+  // partials == nullptr && gridDim.z > 1: LEGACY pre-#54706 CAS epilogue
+  //   (A/B control only; never selected by launch_gemm_q4_deterministic).
+  //   Adds the narrowed 4-lane partial into c with one 64-bit CAS per 4
+  //   columns; REQUIRES a pre-zeroed output because the CAS adds into
+  //   whatever is already there.
   #pragma unroll
   for (int m = 0; m < M_COUNT; ++m) {
     if (offset_m + m >= size_m) continue;  // skip padding rows past size_m
+    if (partials != nullptr) {
+      float* p =
+          partials + ((long)blockIdx.z * size_m + (offset_m + m)) * size_n + n;
+      p[0] = block_c[m][0];
+      p[1] = block_c[m][1];
+      p[2] = block_c[m][2];
+      p[3] = block_c[m][3];
+      continue;
+    }
     T* out = c + (offset_m + m) * size_n + n;
+    if (gridDim.z == 1) {
+      // Single writer per element: round once and store the packed 4 lanes
+      // (8-byte aligned, see the note above) in one go.
+      if constexpr (std::is_same<T, half>::value) {
+        half2 packed[2] = {__halves2half2(__float2half_rn(block_c[m][0]),
+                                          __float2half_rn(block_c[m][1])),
+                           __halves2half2(__float2half_rn(block_c[m][2]),
+                                          __float2half_rn(block_c[m][3]))};
+        __builtin_memcpy(out, packed, sizeof(packed));
+      } else {
+        bf162_t packed[2];
+        packed[0].x = __float2bfloat16(block_c[m][0]);
+        packed[0].y = __float2bfloat16(block_c[m][1]);
+        packed[1].x = __float2bfloat16(block_c[m][2]);
+        packed[1].y = __float2bfloat16(block_c[m][3]);
+        __builtin_memcpy(out, packed, sizeof(packed));
+      }
+      continue;
+    }
+    // Legacy CAS-atomic epilogue: see the store-mode comment above.
     if constexpr (std::is_same<T, half>::value) {
       half2 r01 = __halves2half2(__float2half_rn(block_c[m][0]),
                                  __float2half_rn(block_c[m][1]));
@@ -548,6 +671,67 @@ __global__ void gemm_q4_kernel_rdna3(const T* __restrict__ a,
       atomic_add_pk4_bf16(out, r01, r23);
     }
   }
+
+  // Deterministic in-kernel split-K reduction (fp16 only — on gfx1100 the
+  // bf16 scalar kernel measured neutral at M=1 and up to 10% slower at
+  // M=4-8 with the reduction fused, so bf16 keeps the separate reducer;
+  // see launch_gemm_q4_deterministic). Every (x, y) tile has gridDim.z
+  // producer blocks; the last one to finish sums the tile's z slices in
+  // fixed ascending order and writes the output with a single final cast.
+  // Memory ordering: each thread's partials stores precede its
+  // __threadfence() (release, device scope); the fence precedes the block
+  // barrier, so thread 0's ticket increment happens after every store of
+  // the block is device-visible. The last block's second __threadfence()
+  // pairs with the writers' fences (acquire side), making the other blocks'
+  // partials visible before they are read. The ticket is re-armed after the
+  // reduction so the same per-call array serves the next row tile
+  // (launches on a stream are ordered, so the re-arm is complete before any
+  // later launch reads the ticket).
+  if constexpr (std::is_same<T, half>::value) {
+    if (partials != nullptr) {
+      __shared__ int s_last;
+      __threadfence();
+      __syncthreads();
+      if (t == 0) {
+        s_last = atomicAdd(tickets + blockIdx.y * gridDim.x + blockIdx.x, 1) ==
+                 (int)gridDim.z - 1;
+      }
+      __syncthreads();
+      if (!s_last) return;
+      __threadfence();
+
+      typedef float f4v __attribute__((ext_vector_type(4)));
+      const long zs = ((long)size_m * size_n) / 4;  // z-slice stride, f4 units
+      const int zc = (int)gridDim.z;
+  #pragma unroll
+      for (int m = 0; m < M_COUNT; ++m) {
+        if (offset_m + m >= size_m) continue;  // skip padding rows past size_m
+        const f4v* p =
+            (const f4v*)(partials + (long)(offset_m + m) * size_n + n);
+        f4v acc = {};
+        int z = 0;
+        // Ascending-z accumulation with batched loads (one L2 round trip per
+        // eight slices); the addition order is fixed, so the result is
+        // bit-reproducible.
+        for (; z + 8 <= zc; z += 8) {
+          f4v v[8];
+  #pragma unroll
+          for (int u = 0; u < 8; ++u)
+            v[u] = __builtin_nontemporal_load(p + (long)(z + u) * zs);
+  #pragma unroll
+          for (int u = 0; u < 8; ++u) acc += v[u];
+        }
+        for (; z < zc; ++z) acc += __builtin_nontemporal_load(p + (long)z * zs);
+
+        T* out = c + (long)(offset_m + m) * size_n + n;
+        half2 packed[2] = {
+            __halves2half2(__float2half_rn(acc[0]), __float2half_rn(acc[1])),
+            __halves2half2(__float2half_rn(acc[2]), __float2half_rn(acc[3]))};
+        __builtin_memcpy(out, packed, sizeof(packed));
+      }
+      if (t == 0) *(tickets + blockIdx.y * gridDim.x + blockIdx.x) = 0;
+    }
+  }
 }
 
 #else  // non-RDNA3 device pass: empty __global__ for symbol parity.
@@ -555,7 +739,8 @@ __global__ void gemm_q4_kernel_rdna3(const T* __restrict__ a,
 template <typename T, int M_COUNT>
 __global__ void gemm_q4_kernel_rdna3(const T*, const uint32_t*, const uint32_t*,
                                      const T*, T*, const int, const int,
-                                     const int, const int, const int) {}
+                                     const int, const int, const int, float*,
+                                     int*) {}
 
 #endif  // __HIP__RDNA3__ || !__HIP_DEVICE_COMPILE__
 
@@ -567,20 +752,28 @@ template <typename T, int M_COUNT>
 void launch_gemm_q4_for_mcount(const T* a, const uint32_t* b_q_weight,
                                const uint32_t* b_qzeros, const T* b_scales,
                                T* c, int size_m, int size_n, int size_k,
-                               int groups, int zero_offset,
+                               int groups, int zero_offset, float* partials,
-                               cudaStream_t stream) {
+                               int* tickets, cudaStream_t stream) {
   dim3 block(THREADS_X);
   dim3 grid((size_n + BLOCK_KN_SIZE * 4 - 1) / (BLOCK_KN_SIZE * 4),
             (size_m + M_COUNT - 1) / M_COUNT,
             (size_k + BLOCK_KN_SIZE - 1) / BLOCK_KN_SIZE);
 
-  gemm_q4_kernel_rdna3<T, M_COUNT>
+  gemm_q4_kernel_rdna3<T, M_COUNT><<<grid, block, 0, stream>>>(
-      <<<grid, block, 0, stream>>>(a, b_q_weight, b_qzeros, b_scales, c, size_m,
+      a, b_q_weight, b_qzeros, b_scales, c, size_m, size_n, size_k, groups,
-                                   size_n, size_k, groups, zero_offset);
+      zero_offset, partials, tickets);
+}
+
+// The M_COUNT template the launcher below picks for a given row count.
+static inline int mcount_for_rows(int rows) {
+  if (rows == 1) return 1;
+  if (rows <= 3) return 2;
+  if (rows <= 7) return 4;
+  return 8;
 }
 
 // Dispatch to the largest M_COUNT template that doesn't waste more than
-// half a tile. Caps at 8: above that, the WMMA-prefill kernel (M >= 16) is
+// half a tile. Caps at 8: above that, the WMMA-prefill kernel (M >= 12) is
 // the right tool, not bigger M_COUNT in the scalar dot-product path.
 //
 // Tile-waste table:
@@ -594,28 +787,130 @@ template <typename T>
 void launch_gemm_q4(const T* a, const uint32_t* b_q_weight,
                     const uint32_t* b_qzeros, const T* b_scales, T* c,
                     int size_m, int size_n, int size_k, int groups,
-                    bool use_v2_format, cudaStream_t stream) {
+                    bool use_v2_format, float* partials, int* tickets,
+                    cudaStream_t stream) {
   const int zero_offset = use_v2_format ? 0 : 1;
+  switch (mcount_for_rows(size_m)) {
+    case 1:
+      launch_gemm_q4_for_mcount<T, 1>(a, b_q_weight, b_qzeros, b_scales, c,
+                                      size_m, size_n, size_k, groups,
+                                      zero_offset, partials, tickets, stream);
+      break;
+    case 2:
+      launch_gemm_q4_for_mcount<T, 2>(a, b_q_weight, b_qzeros, b_scales, c,
+                                      size_m, size_n, size_k, groups,
+                                      zero_offset, partials, tickets, stream);
+      break;
+    case 4:
+      launch_gemm_q4_for_mcount<T, 4>(a, b_q_weight, b_qzeros, b_scales, c,
+                                      size_m, size_n, size_k, groups,
+                                      zero_offset, partials, tickets, stream);
+      break;
+    default:
+      // M_COUNT=8 covers the whole scalar domain (M <= 11 in practice: the
+      // public dispatch routes M >= 12 to WMMA); this branch is only reached
+      // when the caller falls through, and still produces correct output.
+      launch_gemm_q4_for_mcount<T, 8>(a, b_q_weight, b_qzeros, b_scales, c,
+                                      size_m, size_n, size_k, groups,
+                                      zero_offset, partials, tickets, stream);
+      break;
+  }
+}
 
-  if (size_m == 1) {
+// Deterministic split-K reduction for the bf16 launches: one thread per
-    launch_gemm_q4_for_mcount<T, 1>(a, b_q_weight, b_qzeros, b_scales, c,
+// output element sums the grid.z FP32 partial slices in fixed ascending-z
-                                    size_m, size_n, size_k, groups, zero_offset,
+// order and rounds to the output dtype exactly once. The order is a pure
-                                    stream);
+// function of the launch shape, so the result is bit-reproducible.
-  } else if (size_m <= 3) {
+template <typename T>
-    launch_gemm_q4_for_mcount<T, 2>(a, b_q_weight, b_qzeros, b_scales, c,
+__global__ void reduce_partials_rdna3(const float* __restrict__ partials,
-                                    size_m, size_n, size_k, groups, zero_offset,
+                                      T* __restrict__ c, const int z_count,
-                                    stream);
+                                      const int size_m, const int size_n) {
-  } else if (size_m <= 7) {
+  const long idx = (long)blockIdx.x * blockDim.x + threadIdx.x;
-    launch_gemm_q4_for_mcount<T, 4>(a, b_q_weight, b_qzeros, b_scales, c,
+  if (idx >= (long)size_m * size_n) return;
-                                    size_m, size_n, size_k, groups, zero_offset,
+  const int m = (int)(idx / size_n);
-                                    stream);
+  const int n = (int)(idx % size_n);
+  float acc = 0.0f;
+  for (int z = 0; z < z_count; ++z)
+    acc += partials[((long)z * size_m + m) * size_n + n];
+  if constexpr (std::is_same<T, half>::value) {
+    c[idx] = __float2half_rn(acc);
   } else {
-    // M_COUNT=8 covers M up to 15 here; M >= 16 should ideally take the
+    c[idx] = __float2bfloat16(acc);
-    // WMMA path, but if it falls through we still produce correct output —
+  }
-    // just leaving 3-5× of throughput on the table for prefill workloads.
+}
-    launch_gemm_q4_for_mcount<T, 8>(a, b_q_weight, b_qzeros, b_scales, c,
+
-                                    size_m, size_n, size_k, groups, zero_offset,
+// Deterministic scalar GEMM: split-K blocks write FP32 partials to scratch
-                                    stream);
+// (no atomics, no intermediate low-precision rounding), and the last z
+// block of each (x, y) tile performs the fixed ascending-z reduction
+// in-kernel with a single final cast (see the kernel epilogue). The scalar
+// kernel writes every in-range (z, m, n) partial exactly once, so a plain
+// at::empty scratch suffices; row-tile bound:
+//   scratch_bytes = z_count * TILE_M * size_n * 4
+// is independent of the caller's M (the scalar domain is M < 64, so a single
+// tile covers it). The PyTorch caching allocator (including its CUDA-graph
+// capture pool) owns scratch reuse and lifetime.
+//
+// The arrival tickets are per-call at::zeros state: every launch of this
+// call counts on tickets that start at zero (the kernel re-arms them for
+// the next row tile), and concurrent invocations get distinct allocations
+// from the stream-aware caching allocator, so no synchronization state is
+// ever shared between invocations. When z_count == 1 the kernel's
+// direct-store epilogue (gridDim.z == 1) is already deterministic, so the
+// scratch and ticket allocations are skipped entirely.
+template <typename T>
+void launch_gemm_q4_deterministic(const T* a, const uint32_t* b_q_weight,
+                                  const uint32_t* b_qzeros, const T* b_scales,
+                                  T* c, int size_m, int size_n, int size_k,
+                                  int groups, bool use_v2_format,
+                                  cudaStream_t stream) {
+  constexpr int TILE_M = 64;  // single tile covers the scalar domain
+  const int z_count = (size_k + BLOCK_KN_SIZE - 1) / BLOCK_KN_SIZE;
+  if (z_count == 1) {
+    launch_gemm_q4(a, b_q_weight, b_qzeros, b_scales, c, size_m, size_n, size_k,
+                   groups, use_v2_format,
+                   /*partials=*/nullptr, /*tickets=*/nullptr, stream);
+    return;
+  }
+  at::TensorOptions dev_opts = at::TensorOptions().device(
+      at::Device(at::kCUDA, c10::cuda::current_device()));
+  at::Tensor partials = at::empty({z_count, std::min(TILE_M, size_m), size_n},
+                                  dev_opts.dtype(at::kFloat));
+  float* partials_ptr = partials.data_ptr<float>();
+
+  int* tickets_ptr = nullptr;
+  if constexpr (std::is_same<T, half>::value) {
+    // fp16: the reduction runs in the kernel (see the epilogue). One
+    // ticket per (x, y) tile, sized for the widest row-tile launch.
+    const int gx = (size_n + BLOCK_KN_SIZE * 4 - 1) / (BLOCK_KN_SIZE * 4);
+    int max_gy = 1;
+    for (int row0 = 0; row0 < size_m; row0 += TILE_M) {
+      const int rows = std::min(TILE_M, size_m - row0);
+      const int mc = mcount_for_rows(rows);
+      max_gy = std::max(max_gy, (rows + mc - 1) / mc);
+    }
+    at::Tensor tickets =
+        at::zeros({(long)gx * max_gy}, dev_opts.dtype(at::kInt));
+    for (int row0 = 0; row0 < size_m; row0 += TILE_M) {
+      const int rows = std::min(TILE_M, size_m - row0);
+      launch_gemm_q4(a + (long)row0 * size_k, b_q_weight, b_qzeros, b_scales,
+                     c + (long)row0 * size_n, rows, size_n, size_k, groups,
+                     use_v2_format, partials_ptr, tickets.data_ptr<int>(),
+                     stream);
+    }
+  } else {
+    // bf16: separate fixed-order reduce (measured faster than the fused
+    // form on this path).
+    for (int row0 = 0; row0 < size_m; row0 += TILE_M) {
+      const int rows = std::min(TILE_M, size_m - row0);
+      launch_gemm_q4(a + (long)row0 * size_k, b_q_weight, b_qzeros, b_scales,
+                     c + (long)row0 * size_n, rows, size_n, size_k, groups,
+                     use_v2_format, partials_ptr, tickets_ptr, stream);
+      const long total = (long)rows * size_n;
+      const int threads = 256;
+      const int blocks = (int)((total + threads - 1) / threads);
+      reduce_partials_rdna3<T><<<blocks, threads, 0, stream>>>(
+          partials_ptr, c + (long)row0 * size_n, z_count, rows, size_n);
+    }
   }
 }
 
@@ -640,13 +935,40 @@ torch::Tensor gptq_gemm_rdna3_wmma(torch::Tensor a, torch::Tensor b_q_weight,
                                    torch::Tensor b_qzeros,
                                    torch::Tensor b_scales, bool use_v2_format);
 
+// WMMA dispatch threshold, from the measured scalar/WMMA crossover on
+// gfx1100 (ratio = scalar_us / wmma_us, >1 means WMMA faster; median of
+// 300 CUDA-event timings per arm, identical inputs, both ops called
+// directly):
+//
+//   K / N (dtype)     M=1   M=4   M=8  M=11  M=12  M=16  M=24  M=32
+//   1536/5120 fp16   0.69  0.97  1.36  1.60  1.59  1.60  1.22  1.62
+//   4096/4096 fp16   0.46  0.69  1.01  1.30  1.30  1.32  1.05  1.40
+//   4352/5120 fp16   0.42  0.63  0.93  1.37  1.37  1.37  1.19  1.74
+//   5120/2048 fp16   0.44  0.62  0.87  1.00  1.00  1.01  1.15  1.34
+//   5120/8704 fp16   0.30  0.53  0.84  1.39  1.42  1.43  1.24  1.58
+//   1536/5120 bf16   1.07  1.23  1.50  1.71  1.69  1.00  1.00  1.00
+//   4096/4096 bf16   0.70  0.82  1.08  1.33  1.34  1.00  1.00  1.00
+//   4352/5120 bf16   0.62  0.74  0.98  1.40  1.40  1.00  1.00  1.00
+//   5120/2048 bf16   0.64  0.73  0.89  1.00  1.00  1.00  1.01  1.00
+//   5120/8704 bf16   0.45  0.65  0.92  1.46  1.49  1.00  1.00  1.00
+//
+// At M >= 12 WMMA wins or ties on every shape for both dtypes (exact 1.00
+// cells and the one sub-1.0 sample, bf16 4096/4096 at M=24 = 0.996, are
+// within measurement noise); below
+// that the scalar path wins on most shapes (and is the more accurate of
+// the two since the exact fp16 factored dequant, so it keeps the band).
+// fp16 used to stay scalar until M=64 because the old baked-dequant
+// scalar kernel was faster there — that speed came with ~100x more
+// dequant error, and with the exact dequant the crossover moved down to
+// the same 12 as bf16. A shape-aware (N/K) dispatch would do better
+// still; that is tuning work beyond this fix.
+constexpr int64_t WMMA_MIN_M = 12;
+
 torch::Tensor gptq_gemm_rdna3(torch::Tensor a, torch::Tensor b_q_weight,
                               torch::Tensor b_qzeros, torch::Tensor b_scales,
                               bool use_v2_format) {
   if (a.dim() == 2 && b_q_weight.dim() == 2 && a.size(1) % 16 == 0 &&
-      b_q_weight.size(1) % 16 == 0 &&
+      b_q_weight.size(1) % 16 == 0 && a.size(0) >= WMMA_MIN_M) {
-      ((a.scalar_type() == torch::kBFloat16 && a.size(0) >= 16) ||
-       (a.scalar_type() == torch::kHalf && a.size(0) >= 64))) {
     return gptq_gemm_rdna3_wmma(a, b_q_weight, b_qzeros, b_scales,
                                 use_v2_format);
   }
@@ -673,22 +995,38 @@ torch::Tensor gptq_gemm_rdna3(torch::Tensor a, torch::Tensor b_q_weight,
 
   TORCH_CHECK(b_q_weight.size(0) * 8 == size_k,
               "b_q_weight first dim must be K/8");
+  TORCH_CHECK(size_k % (groups * 32) == 0,
+              "group size (K/groups = ", size_k / groups,
+              ") must be a multiple of 32: the kernel checks group "
+              "transitions at 32-K granularity");
   TORCH_CHECK(b_scales.size(0) == groups,
               "b_scales must have same group count as qzeros");
   TORCH_CHECK(b_scales.size(1) == size_n, "b_scales last dim must be N");
-  TORCH_CHECK(size_n % 8 == 0, "N must be a multiple of 8 (64-bit atomic CAS)");
+  TORCH_CHECK(size_n % 8 == 0,
+              "N must be a multiple of 8 (packed qzeros layout: 8 4-bit zero "
+              "points per uint32 along N; also keeps the 4-column epilogue "
+              "store in bounds)");
 
   auto opts = torch::TensorOptions().dtype(a.dtype()).device(a.device());
-  at::Tensor c = torch::zeros({size_m, size_n}, opts);
+  // The deterministic epilogue writes every output element exactly once
+  // (reduce pass below — or the direct store when z_count == 1), so c needs
+  // no zero-initialization. This is a one-way door: the legacy CAS epilogue
+  // (see the kernel's store-mode comment) ADDS into c and would require
+  // torch::zeros here if it were ever re-selected.
+  at::Tensor c = torch::empty({size_m, size_n}, opts);
 
+  // Deterministic split-K: FP32 partials + fixed-order reduction (see
+  // launch_gemm_q4_deterministic). The CAS-atomic low-precision epilogue
+  // this replaces was order-dependent and produced different results for
+  // identical inputs on gfx11 once more than a few split blocks contended.
   if (a.scalar_type() == torch::kHalf) {
-    vllm::gptq_rdna3::launch_gemm_q4<half>(
+    vllm::gptq_rdna3::launch_gemm_q4_deterministic<half>(
         (const half*)a.data_ptr(), (const uint32_t*)b_q_weight.data_ptr(),
         (const uint32_t*)b_qzeros.data_ptr(), (const half*)b_scales.data_ptr(),
         (half*)c.data_ptr(), size_m, size_n, size_k, groups, use_v2_format,
         stream);
   } else {
-    vllm::gptq_rdna3::launch_gemm_q4<vllm::gptq_rdna3::bf16_t>(
+    vllm::gptq_rdna3::launch_gemm_q4_deterministic<vllm::gptq_rdna3::bf16_t>(
         (const vllm::gptq_rdna3::bf16_t*)a.data_ptr(),
         (const uint32_t*)b_q_weight.data_ptr(),
         (const uint32_t*)b_qzeros.data_ptr(),

csrc/rocm/q_gemm_rdna3_wmma.cu

@@ -24,13 +24,15 @@
 //     of the 16x16 output, with 8 elements alternating M rows by lane_hi
 //     (lanes 0..15 = even rows, lanes 16..31 = odd rows). See the layout
 //     diagram on `gemm_q4_wmma_kernel_16x16_1w` below for the full mapping.
-//   * No native v_global_atomic_pk_add_{f16,bf16} on gfx11; the K-split
+//   * No native v_global_atomic_pk_add_{f16,bf16} on gfx11. The split-K
-//     epilogue (gridDim.z > 1) emulates packed atomic add via a CAS-32
+//     epilogue therefore stores FP32 per-split partials to scratch and a
-//     retry loop on a uint32 word covering 2 fp16/bf16 lanes. Within a
+//     separate pass reduces them in fixed ascending-z order with a single
-//     block we shuffle adjacent lanes via shfl_xor first so each pair of
+//     final low-precision rounding — deterministic by construction
-//     output cols goes through a single atomic — no intra-block
+//     (PR #54706). K_SPLIT == 1 keeps the original direct-write path with
-//     contention. K_SPLIT == 1 keeps the original direct-write path with
+//     each block owning its 16M × 16N output tile. The legacy pre-#54706
-//     each block owning its 16M × 16N output tile.
+//     design (a CAS-32 retry loop accumulating narrowed partials directly
+//     into c, order-dependent) is kept in the store path below the
+//     partials branch for A/B comparison only.
 
 #include <cstdint>
 
@@ -182,16 +184,19 @@ __forceinline__ __device__ void dequant_4bit_8_bf16_to_bf16(uint32_t qa,
 }
 
 // ---------------------------------------------------------------------------
-// Packed atomic-add helpers used by the K-split epilogue.
+// Packed atomic-add helpers — LEGACY (pre-#54706) K-split epilogue.
 //
-// When the kernel is launched with gridDim.z > 1, multiple K-segments
+// When the kernel is launched with partials == nullptr and gridDim.z > 1,
-// accumulate into the same 16x16 output tile and need atomic write-back.
+// multiple K-segments accumulate into the same 16x16 output tile and the
-// gfx11 has no native v_global_atomic_pk_add_{f16,bf16}, so we issue a
+// legacy path atomically adds narrowed partials into a pre-zeroed c via a
 // CAS-loop on a 32-bit word covering 2 packed fp16/bf16 lanes. Within a
 // block the kernel pairs adjacent lanes via shfl_xor first, so each pair
 // of cols (n=lane_lo even, lane_lo+1) goes through a SINGLE atomic — no
 // intra-block contention on the same uint32 target. Inter-block
 // contention from gridDim.z (4-way at K_SPLIT=4) is the residual cost.
+// The addition happens in bf16/fp16 AFTER narrowing, so the result depends
+// on CAS completion order — the nondeterminism PR #54706 replaces. Kept
+// for A/B comparison only; the shipped launchers never select it.
 // ---------------------------------------------------------------------------
 
 __forceinline__ __device__ void atomic_add_pk_f16(half2* addr, half2 val) {
@@ -230,12 +235,15 @@ __forceinline__ __device__ void atomic_add_pk_bf16(bf162_t* addr, bf162_t val) {
 // K-split factor heuristic. Returns the gridDim.z to use for a given K.
 // Aim: each block does at least ~16 K-tiles (= K=256) so the per-block
 // constant overhead (LDS init, kernel prologue) is amortised. Upper
-// bound K_SPLIT=4 to cap inter-block atomic contention to 4-way.
+// bound K_SPLIT=4 — historically to cap inter-block atomic contention to
+// 4-way under the legacy CAS epilogue; now it caps the FP32 scratch and
+// reduce-pass traffic the deterministic epilogue adds.
 //
 // For typical Qwen-class shapes K ∈ {4096, 5120, 11008}, all return 4.
-// Smaller K (e.g., embedding lookups) fall back to 1 (no split, no
+// Smaller K (e.g., embedding lookups) fall back to 1: no split, no
-// atomic). K must be divisible by (K_SPLIT × 16) for the split to be
+// scratch, direct store. K must be divisible by (K_SPLIT × 16) for the
-// valid; the heuristic checks divisibility before raising the factor.
+// split to be valid; the heuristic checks divisibility before raising
+// the factor.
 __host__ __device__ static inline int compute_wmma_k_split(int size_k) {
   if (size_k >= 1024 && size_k % 64 == 0) return 4;
   if (size_k >= 512 && size_k % 32 == 0) return 2;
@@ -245,17 +253,18 @@ __host__ __device__ static inline int compute_wmma_k_split(int size_k) {
 // M-and-N-aware K-split heuristic for the v3/v4/v5 launchers.
 //
 // The original `compute_wmma_k_split` was K-only and always returns 4 for
-// Qwen-class K, which over-subscribes wave slots and pays the atomic CAS
+// Qwen-class K, which over-subscribes wave slots and pays the split-K
-// epilogue once-per-K-segment per output cell. With v3/v4/v5's larger
+// epilogue (historically the atomic CAS; now FP32 scratch + reduce) once
+// per K-segment per output cell. With v3/v4/v5's larger
 // tiles (64M × 16/32/64N) and 4 resident waves per block, the no-split
 // grid is often already well-saturated on gfx1100's 96 CUs / 3072 wave
-// slots — adding gridDim.z just adds atomic overhead.
+// slots — adding gridDim.z just adds epilogue overhead.
 //
 // Heuristic: compute the no-split block count gridDim.x × gridDim.y, then
 // pick the smallest K_SPLIT that brings total waves to at least
 // ~2× over-subscription (~6000 waves for our 3072 slots, i.e. 1500 blocks
-// at 4 waves/block). Above that threshold, K_SPLIT=1 — direct write, no
+// at 4 waves/block). Above that threshold, K_SPLIT=1: no scratch, no
-// atomic.
+// reduce pass, direct store.
 //
 // Args:
 //   size_m, size_n, size_k     — GEMM dims
@@ -373,7 +382,11 @@ __global__ void gemm_q4_wmma_kernel_16x16_1w(
     const T* __restrict__ a, const uint32_t* __restrict__ b_q,
     const uint32_t* __restrict__ b_qzeros, const T* __restrict__ b_scales,
     T* __restrict__ c, const int size_m, const int size_n, const int size_k,
-    const int groups, const int zero_offset) {
+    const int groups, const int zero_offset,
+    // Deterministic split-K: non-null => each block stores its FP32 partial
+    // to partials[(z*size_m + m)*size_n + n]; a separate pass then reduces
+    // the z-slices in fixed order with a single low-precision rounding.
+    float* __restrict__ partials) {
   using E = typename WmmaNative<T>::elem;
   using V16 = typename WmmaNative<T>::v16;
 
@@ -525,7 +538,21 @@ __global__ void gemm_q4_wmma_kernel_16x16_1w(
   //   lane_hi == 0  →  rows 0, 2, 4, ..., 14   (even rows)
   //   lane_hi == 1  →  rows 1, 3, 5, ..., 15   (odd rows)
   // c_acc[i] corresponds to actual row m = 2*i + lane_hi at column lane_lo.
-  if (gridDim.z > 1) {
+  if (partials != nullptr) {
+    // Deterministic split-K: plain FP32 store per (z, m, n); the reduce
+    // pass fixes the accumulation order.
+    const int out_n = n_tile + lane_lo;
+    if (out_n < size_n) {
+  #pragma unroll
+      for (int i = 0; i < 8; i++) {
+        const int out_m = m_tile + 2 * i + lane_hi;
+        if (out_m < size_m) {
+          partials[((long)blockIdx.z * size_m + out_m) * size_n + out_n] =
+              c_acc[i];
+        }
+      }
+    }
+  } else if (gridDim.z > 1) {
     // K-split path: 4 (or whatever the split factor is) K-segments per
     // output cell contend → atomic accumulation. Caller has zero-init'd c.
     //
@@ -588,9 +615,78 @@ template <typename T>
 __global__ void gemm_q4_wmma_kernel_16x16_1w(const T*, const uint32_t*,
                                              const uint32_t*, const T*, T*,
                                              const int, const int, const int,
-                                             const int, const int) {}
+                                             const int, const int, float*) {}
 #endif
 
+// Deterministic split-K reduction for the WMMA path: one thread per output
+// element sums the grid.z FP32 partial slices in fixed ascending-z order and
+// rounds to the output dtype exactly once. Order is a pure function of the
+// launch shape, so the result is bit-reproducible for identical inputs.
+// (Kernel bodies live in q_gemm_rdna3_wmma.cu alongside their users; HIP
+// kernel templates are not shared across translation units here.)
+template <typename T>
+__global__ void reduce_partials_wmma(const float* __restrict__ partials,
+                                     T* __restrict__ c, const int z_count,
+                                     const int size_m, const int size_n) {
+  const long idx = (long)blockIdx.x * blockDim.x + threadIdx.x;
+  if (idx >= (long)size_m * size_n) return;
+  const int m = (int)(idx / size_n);
+  const int n = (int)(idx % size_n);
+  float acc = 0.0f;
+  for (int z = 0; z < z_count; ++z)
+    acc += partials[((long)z * size_m + m) * size_n + n];
+  if constexpr (std::is_same<T, half>::value) {
+    c[idx] = __float2half_rn(acc);
+  } else {
+    c[idx] = __float2bfloat16(acc);
+  }
+}
+
+// FP32 split-partial scratch for one deterministic WMMA launch.
+//
+// Coverage invariant — why at::empty is safe: every scratch slot that the
+// fixed-order reducer reads, i.e. every (z, m, n) with z < k_split,
+// m < size_m, n < size_n, has EXACTLY ONE writer:
+//   * z: grid.z == k_split and each z-block stores slice blockIdx.z;
+//   * m: grid.y tiles M by the kernel's M_TILE; within a tile, wave w owns
+//     rows [16w, 16w+16) and (i, lane_hi) with out_m = m_tile + 16w + 2i +
+//     lane_hi enumerates the 16 rows of the wave uniquely (i in [0,8),
+//     lane_hi in {0,1});
+//   * n: grid.x tiles N by N_TILE; each 16-wide accumulator slice stores
+//     columns n_base + lane_lo with lane_lo in [0,16) unique per column.
+// The out_m >= size_m / out_n >= size_n guards in the store paths skip
+// only slots the reducer never reads (it iterates m < size_m,
+// n < size_n), and the store index ((z * size_m + m) * size_n + n) stays
+// inside [0, k_split * size_m * size_n) — no out-of-bounds writes either.
+// This holds for every WMMA kernel variant in this TU (16x16_1w, 32x16_2w,
+// 64x16_4w, 64x32_4w, 64x64_4w, 128x64_k16, 128x64_k32) because they all
+// share the wave/lane output mapping above; the tiled launchers pass
+// size_m = rows of the current row tile so the same argument applies per
+// tile. A stale at::empty scratch therefore never reaches the reducer.
+//
+// One-way door: the legacy CAS epilogue (partials == nullptr &&
+// gridDim.z > 1, kept for A/B comparison) accumulates into the OUTPUT
+// instead and would require a zero-initialized c if re-selected.
+static inline at::Tensor alloc_wmma_partials(int k_split, int size_m,
+                                             int size_n) {
+  return at::empty(
+      {k_split, size_m, size_n},
+      at::TensorOptions()
+          .dtype(at::kFloat)
+          .device(at::Device(at::kCUDA, c10::cuda::current_device())));
+}
+
+template <typename T>
+static inline void launch_wmma_reduce(const at::Tensor& partials, T* c,
+                                      int k_split, int size_m, int size_n,
+                                      cudaStream_t stream) {
+  const long total = (long)size_m * size_n;
+  const int threads = 256;
+  const int blocks = (int)((total + threads - 1) / threads);
+  reduce_partials_wmma<T><<<blocks, threads, 0, stream>>>(
+      partials.data_ptr<float>(), c, k_split, size_m, size_n);
+}
+
 template <typename T>
 void launch_gemm_q4_wmma_16x16_1w(const T* a, const uint32_t* b_q_weight,
                                   const uint32_t* b_qzeros, const T* b_scales,
@@ -599,13 +695,24 @@ void launch_gemm_q4_wmma_16x16_1w(const T* a, const uint32_t* b_q_weight,
                                   cudaStream_t stream) {
   // 1 wave per block (32 lanes), 16x16 C tile per block. gridDim.z splits
   // K so that more blocks (and therefore more waves) are in flight; with
-  // K_SPLIT > 1 the kernel switches to atomic write-back at the epilogue.
+  // K_SPLIT > 1 the kernel stores FP32 partials and a reduce pass fixes
+  // the accumulation order.
   const int k_split = compute_wmma_k_split(size_k);
   dim3 block(32);
   dim3 grid((size_n + 15) / 16, (size_m + 15) / 16, k_split);
-  gemm_q4_wmma_kernel_16x16_1w<T>
+  // k_split == 1: single writer per cell, the kernel's direct-store path is
-      <<<grid, block, 0, stream>>>(a, b_q_weight, b_qzeros, b_scales, c, size_m,
+  // already deterministic. k_split > 1: FP32 partials + fixed-order reduce.
-                                   size_n, size_k, groups, zero_offset);
+  at::Tensor partials;
+  float* partials_ptr = nullptr;
+  if (k_split > 1) {
+    partials = alloc_wmma_partials(k_split, size_m, size_n);
+    partials_ptr = partials.data_ptr<float>();
+  }
+  gemm_q4_wmma_kernel_16x16_1w<T><<<grid, block, 0, stream>>>(
+      a, b_q_weight, b_qzeros, b_scales, c, size_m, size_n, size_k, groups,
+      zero_offset, partials_ptr);
+  if (k_split > 1)
+    launch_wmma_reduce<T>(partials, c, k_split, size_m, size_n, stream);
 }
 
 #if defined(__HIP__RDNA3__) || !defined(__HIP_DEVICE_COMPILE__)
@@ -647,7 +754,11 @@ __global__ void gemm_q4_wmma_kernel_32x16_2w(
     const T* __restrict__ a, const uint32_t* __restrict__ b_q,
     const uint32_t* __restrict__ b_qzeros, const T* __restrict__ b_scales,
     T* __restrict__ c, const int size_m, const int size_n, const int size_k,
-    const int groups, const int zero_offset) {
+    const int groups, const int zero_offset,
+    // Deterministic split-K: non-null => each block stores its FP32 partial
+    // to partials[(z*size_m + m)*size_n + n]; a separate pass then reduces
+    // the z-slices in fixed order with a single low-precision rounding.
+    float* __restrict__ partials) {
   using E = typename WmmaNative<T>::elem;
   using V16 = typename WmmaNative<T>::v16;
 
@@ -775,7 +886,21 @@ __global__ void gemm_q4_wmma_kernel_32x16_2w(
   // Each wave owns rows [m_tile + wave_id*16 .. + 16) of the output tile.
   const int m_tile_wave = m_tile + wave_id * 16;
 
-  if (gridDim.z > 1) {
+  if (partials != nullptr) {
+    // Deterministic split-K: plain FP32 store per (z, m, n); the reduce
+    // pass fixes the accumulation order.
+    const int out_n = n_tile + lane_lo;
+    if (out_n < size_n) {
+  #pragma unroll
+      for (int i = 0; i < 8; i++) {
+        const int out_m = m_tile_wave + 2 * i + lane_hi;
+        if (out_m < size_m) {
+          partials[((long)blockIdx.z * size_m + out_m) * size_n + out_n] =
+              c_acc[i];
+        }
+      }
+    }
+  } else if (gridDim.z > 1) {
     // K-split atomic path. Pair-shuffle within wave to halve atomic count.
     // shfl_xor here is wave-local (wave32 semantics) so each wave does its
     // own pairing — the two waves don't interact during the store.
@@ -826,7 +951,7 @@ template <typename T>
 __global__ void gemm_q4_wmma_kernel_32x16_2w(const T*, const uint32_t*,
                                              const uint32_t*, const T*, T*,
                                              const int, const int, const int,
-                                             const int, const int) {}
+                                             const int, const int, float*) {}
 #endif
 
 template <typename T>
@@ -856,9 +981,17 @@ void launch_gemm_q4_wmma_32x16_2w(const T* a, const uint32_t* b_q_weight,
   const int k_split = compute_wmma_k_split(size_k);
   dim3 block(64);
   dim3 grid((size_n + 15) / 16, (size_m + 31) / 32, k_split);
-  gemm_q4_wmma_kernel_32x16_2w<T>
+  at::Tensor partials;
-      <<<grid, block, 0, stream>>>(a, b_q_weight, b_qzeros, b_scales, c, size_m,
+  float* partials_ptr = nullptr;
-                                   size_n, size_k, groups, zero_offset);
+  if (k_split > 1) {
+    partials = alloc_wmma_partials(k_split, size_m, size_n);
+    partials_ptr = partials.data_ptr<float>();
+  }
+  gemm_q4_wmma_kernel_32x16_2w<T><<<grid, block, 0, stream>>>(
+      a, b_q_weight, b_qzeros, b_scales, c, size_m, size_n, size_k, groups,
+      zero_offset, partials_ptr);
+  if (k_split > 1)
+    launch_wmma_reduce<T>(partials, c, k_split, size_m, size_n, stream);
 }
 
 #if defined(__HIP__RDNA3__) || !defined(__HIP_DEVICE_COMPILE__)
@@ -898,7 +1031,11 @@ __global__ void gemm_q4_wmma_kernel_64x16_4w(
     const T* __restrict__ a, const uint32_t* __restrict__ b_q,
     const uint32_t* __restrict__ b_qzeros, const T* __restrict__ b_scales,
     T* __restrict__ c, const int size_m, const int size_n, const int size_k,
-    const int groups, const int zero_offset) {
+    const int groups, const int zero_offset,
+    // Deterministic split-K: non-null => each block stores its FP32 partial
+    // to partials[(z*size_m + m)*size_n + n]; a separate pass then reduces
+    // the z-slices in fixed order with a single low-precision rounding.
+    float* __restrict__ partials) {
   using E = typename WmmaNative<T>::elem;
   using V16 = typename WmmaNative<T>::v16;
 
@@ -1024,7 +1161,21 @@ __global__ void gemm_q4_wmma_kernel_64x16_4w(
   // ---- Store C ---- Each wave owns 16 M-rows of the output tile.
   const int m_tile_wave = m_tile + wave_id * 16;
 
-  if (gridDim.z > 1) {
+  if (partials != nullptr) {
+    // Deterministic split-K: plain FP32 store per (z, m, n); the reduce
+    // pass fixes the accumulation order.
+    const int out_n = n_tile + lane_lo;
+    if (out_n < size_n) {
+  #pragma unroll
+      for (int i = 0; i < 8; i++) {
+        const int out_m = m_tile_wave + 2 * i + lane_hi;
+        if (out_m < size_m) {
+          partials[((long)blockIdx.z * size_m + out_m) * size_n + out_n] =
+              c_acc[i];
+        }
+      }
+    }
+  } else if (gridDim.z > 1) {
     // K-split atomic path. Pair-shuffle within wave to halve atomic count.
     const bool is_even_lane = (lane_lo & 1) == 0;
     const int out_n_pair = n_tile + lane_lo;
@@ -1073,7 +1224,7 @@ template <typename T>
 __global__ void gemm_q4_wmma_kernel_64x16_4w(const T*, const uint32_t*,
                                              const uint32_t*, const T*, T*,
                                              const int, const int, const int,
-                                             const int, const int) {}
+                                             const int, const int, float*) {}
 #endif
 
 template <typename T>
@@ -1094,9 +1245,17 @@ void launch_gemm_q4_wmma_64x16_4w(const T* a, const uint32_t* b_q_weight,
   const int k_split = compute_wmma_k_split_mn(size_m, size_n, size_k, 64, 16);
   dim3 block(128);
   dim3 grid((size_n + 15) / 16, (size_m + 63) / 64, k_split);
-  gemm_q4_wmma_kernel_64x16_4w<T>
+  at::Tensor partials;
-      <<<grid, block, 0, stream>>>(a, b_q_weight, b_qzeros, b_scales, c, size_m,
+  float* partials_ptr = nullptr;
-                                   size_n, size_k, groups, zero_offset);
+  if (k_split > 1) {
+    partials = alloc_wmma_partials(k_split, size_m, size_n);
+    partials_ptr = partials.data_ptr<float>();
+  }
+  gemm_q4_wmma_kernel_64x16_4w<T><<<grid, block, 0, stream>>>(
+      a, b_q_weight, b_qzeros, b_scales, c, size_m, size_n, size_k, groups,
+      zero_offset, partials_ptr);
+  if (k_split > 1)
+    launch_wmma_reduce<T>(partials, c, k_split, size_m, size_n, stream);
 }
 
 #if defined(__HIP__RDNA3__) || !defined(__HIP_DEVICE_COMPILE__)
@@ -1130,7 +1289,11 @@ __global__ void gemm_q4_wmma_kernel_64x32_4w(
     const T* __restrict__ a, const uint32_t* __restrict__ b_q,
     const uint32_t* __restrict__ b_qzeros, const T* __restrict__ b_scales,
     T* __restrict__ c, const int size_m, const int size_n, const int size_k,
-    const int groups, const int zero_offset) {
+    const int groups, const int zero_offset,
+    // Deterministic split-K: non-null => each block stores its FP32 partial
+    // to partials[(z*size_m + m)*size_n + n]; a separate pass then reduces
+    // the z-slices in fixed order with a single low-precision rounding.
+    float* __restrict__ partials) {
   using E = typename WmmaNative<T>::elem;
   using V16 = typename WmmaNative<T>::v16;
 
@@ -1266,6 +1429,22 @@ __global__ void gemm_q4_wmma_kernel_64x32_4w(
   // Helper: store one v8fp32 accumulator's 8 outputs (covers 16 N-cols at
   // n_base via lane_lo + interleaved M rows m_tile_wave + 2i + lane_hi).
   auto store_acc = [&](const v8fp32& acc, int n_base) {
+    if (partials != nullptr) {
+      // Deterministic split-K: plain FP32 store per (z, m, n). Each element
+      // has one writer per z-slice and the reduce pass fixes the
+      // accumulation order, so the result is bit-reproducible.
+      const int out_n = n_base + lane_lo;
+      if (out_n >= size_n) return;
+  #pragma unroll
+      for (int i = 0; i < 8; i++) {
+        const int out_m = m_tile_wave + 2 * i + lane_hi;
+        if (out_m < size_m) {
+          partials[((long)blockIdx.z * size_m + out_m) * size_n + out_n] =
+              acc[i];
+        }
+      }
+      return;
+    }
     if (gridDim.z > 1) {
       const bool is_even_lane = (lane_lo & 1) == 0;
       const int out_n_pair = n_base + lane_lo;
@@ -1316,7 +1495,7 @@ template <typename T>
 __global__ void gemm_q4_wmma_kernel_64x32_4w(const T*, const uint32_t*,
                                              const uint32_t*, const T*, T*,
                                              const int, const int, const int,
-                                             const int, const int) {}
+                                             const int, const int, float*) {}
 #endif
 
 template <typename T>
@@ -1339,9 +1518,17 @@ void launch_gemm_q4_wmma_64x32_4w(const T* a, const uint32_t* b_q_weight,
   const int k_split = compute_wmma_k_split_mn(size_m, size_n, size_k, 64, 32);
   dim3 block(128);
   dim3 grid((size_n + 31) / 32, (size_m + 63) / 64, k_split);
-  gemm_q4_wmma_kernel_64x32_4w<T>
+  at::Tensor partials;
-      <<<grid, block, 0, stream>>>(a, b_q_weight, b_qzeros, b_scales, c, size_m,
+  float* partials_ptr = nullptr;
-                                   size_n, size_k, groups, zero_offset);
+  if (k_split > 1) {
+    partials = alloc_wmma_partials(k_split, size_m, size_n);
+    partials_ptr = partials.data_ptr<float>();
+  }
+  gemm_q4_wmma_kernel_64x32_4w<T><<<grid, block, 0, stream>>>(
+      a, b_q_weight, b_qzeros, b_scales, c, size_m, size_n, size_k, groups,
+      zero_offset, partials_ptr);
+  if (k_split > 1)
+    launch_wmma_reduce<T>(partials, c, k_split, size_m, size_n, stream);
 }
 
 #if defined(__HIP__RDNA3__) || !defined(__HIP_DEVICE_COMPILE__)
@@ -1375,7 +1562,11 @@ __global__ void gemm_q4_wmma_kernel_64x64_4w(
     const T* __restrict__ a, const uint32_t* __restrict__ b_q,
     const uint32_t* __restrict__ b_qzeros, const T* __restrict__ b_scales,
     T* __restrict__ c, const int size_m, const int size_n, const int size_k,
-    const int groups, const int zero_offset) {
+    const int groups, const int zero_offset,
+    // Deterministic split-K: non-null => each block stores its FP32 partial
+    // to partials[(z*size_m + m)*size_n + n]; a separate pass then reduces
+    // the z-slices in fixed order with a single low-precision rounding.
+    float* __restrict__ partials) {
   using E = typename WmmaNative<T>::elem;
   using V16 = typename WmmaNative<T>::v16;
 
@@ -1502,6 +1693,22 @@ __global__ void gemm_q4_wmma_kernel_64x64_4w(
   // ---- Store C ---- Each wave owns 16M × 64N. Helper writes one acc slice.
   const int m_tile_wave = m_tile + wave_id * 16;
   auto store_acc = [&](const v8fp32& acc, int n_base) {
+    if (partials != nullptr) {
+      // Deterministic split-K: plain FP32 store per (z, m, n). Each element
+      // has one writer per z-slice and the reduce pass fixes the
+      // accumulation order, so the result is bit-reproducible.
+      const int out_n = n_base + lane_lo;
+      if (out_n >= size_n) return;
+  #pragma unroll
+      for (int i = 0; i < 8; i++) {
+        const int out_m = m_tile_wave + 2 * i + lane_hi;
+        if (out_m < size_m) {
+          partials[((long)blockIdx.z * size_m + out_m) * size_n + out_n] =
+              acc[i];
+        }
+      }
+      return;
+    }
     if (gridDim.z > 1) {
       const bool is_even_lane = (lane_lo & 1) == 0;
       const int out_n_pair = n_base + lane_lo;
@@ -1566,7 +1773,11 @@ __global__ void gemm_q4_wmma_kernel_128x64_k16(
     const T* __restrict__ a, const uint32_t* __restrict__ b_q,
     const uint32_t* __restrict__ b_qzeros, const T* __restrict__ b_scales,
     T* __restrict__ c, const int size_m, const int size_n, const int size_k,
-    const int groups, const int zero_offset) {
+    const int groups, const int zero_offset,
+    // Deterministic split-K: non-null => each block stores its FP32 partial
+    // to partials[(z*size_m + m)*size_n + n]; a separate pass then reduces
+    // the z-slices in fixed order with a single low-precision rounding.
+    float* __restrict__ partials) {
   using E = typename WmmaNative<T>::elem;
   using V16 = typename WmmaNative<T>::v16;
 
@@ -1695,6 +1906,22 @@ __global__ void gemm_q4_wmma_kernel_128x64_k16(
   // ---- Store C ---- Each wave owns 16M × 64N.
   const int m_tile_wave = m_tile + wave_id * 16;
   auto store_acc = [&](const v8fp32& acc, int n_base) {
+    if (partials != nullptr) {
+      // Deterministic split-K: plain FP32 store per (z, m, n). Each element
+      // has one writer per z-slice and the reduce pass fixes the
+      // accumulation order, so the result is bit-reproducible.
+      const int out_n = n_base + lane_lo;
+      if (out_n >= size_n) return;
+  #pragma unroll
+      for (int i = 0; i < 8; i++) {
+        const int out_m = m_tile_wave + 2 * i + lane_hi;
+        if (out_m < size_m) {
+          partials[((long)blockIdx.z * size_m + out_m) * size_n + out_n] =
+              acc[i];
+        }
+      }
+      return;
+    }
     if (gridDim.z > 1) {
       const bool is_even_lane = (lane_lo & 1) == 0;
       const int out_n_pair = n_base + lane_lo;
@@ -1764,7 +1991,11 @@ __global__ void gemm_q4_wmma_kernel_128x64_k32(
     const T* __restrict__ a, const uint32_t* __restrict__ b_q,
     const uint32_t* __restrict__ b_qzeros, const T* __restrict__ b_scales,
     T* __restrict__ c, const int size_m, const int size_n, const int size_k,
-    const int groups, const int zero_offset) {
+    const int groups, const int zero_offset,
+    // Deterministic split-K: non-null => each block stores its FP32 partial
+    // to partials[(z*size_m + m)*size_n + n]; a separate pass then reduces
+    // the z-slices in fixed order with a single low-precision rounding.
+    float* __restrict__ partials) {
   using E = typename WmmaNative<T>::elem;
   using V16 = typename WmmaNative<T>::v16;
 
@@ -1907,6 +2138,22 @@ __global__ void gemm_q4_wmma_kernel_128x64_k32(
   // ---- Store C ---- Same as V7.
   const int m_tile_wave = m_tile + wave_id * 16;
   auto store_acc = [&](const v8fp32& acc, int n_base) {
+    if (partials != nullptr) {
+      // Deterministic split-K: plain FP32 store per (z, m, n). Each element
+      // has one writer per z-slice and the reduce pass fixes the
+      // accumulation order, so the result is bit-reproducible.
+      const int out_n = n_base + lane_lo;
+      if (out_n >= size_n) return;
+  #pragma unroll
+      for (int i = 0; i < 8; i++) {
+        const int out_m = m_tile_wave + 2 * i + lane_hi;
+        if (out_m < size_m) {
+          partials[((long)blockIdx.z * size_m + out_m) * size_n + out_n] =
+              acc[i];
+        }
+      }
+      return;
+    }
     if (gridDim.z > 1) {
       const bool is_even_lane = (lane_lo & 1) == 0;
       const int out_n_pair = n_base + lane_lo;
@@ -1960,17 +2207,17 @@ template <typename T>
 __global__ void gemm_q4_wmma_kernel_64x64_4w(const T*, const uint32_t*,
                                              const uint32_t*, const T*, T*,
                                              const int, const int, const int,
-                                             const int, const int) {}
+                                             const int, const int, float*) {}
 template <typename T>
 __global__ void gemm_q4_wmma_kernel_128x64_k16(const T*, const uint32_t*,
                                                const uint32_t*, const T*, T*,
                                                const int, const int, const int,
-                                               const int, const int) {}
+                                               const int, const int, float*) {}
 template <typename T>
 __global__ void gemm_q4_wmma_kernel_128x64_k32(const T*, const uint32_t*,
                                                const uint32_t*, const T*, T*,
                                                const int, const int, const int,
-                                               const int, const int) {}
+                                               const int, const int, float*) {}
 #endif
 
 template <typename T>
@@ -1994,15 +2241,46 @@ void launch_gemm_q4_wmma_64x64_4w(const T* a, const uint32_t* b_q_weight,
         compute_wmma_k_split_mn(size_m, size_n, size_k, 128, 64);
     const int groupsize = size_k / groups;
     dim3 block(256);
-    dim3 grid((size_n + 63) / 64, (size_m + 127) / 128, k_split);
+    if (k_split == 1) {
-    if (size_k % 32 == 0 && groupsize >= 32 && (size_k / k_split) % 32 == 0) {
+      // Single writer per cell: the kernel's direct-store path is already
-      gemm_q4_wmma_kernel_128x64_k32<T><<<grid, block, 0, stream>>>(
+      // deterministic — no scratch, no reduce pass (same invariant as the
-          a, b_q_weight, b_qzeros, b_scales, c, size_m, size_n, size_k, groups,
+      // 64x64 path below). No row tiling either: tiling exists only to
-          zero_offset);
+      // bound the FP32 scratch, which is not allocated here; grid.y tiles
-    } else {
+      // M by 128 internally for any size_m.
-      gemm_q4_wmma_kernel_128x64_k16<T><<<grid, block, 0, stream>>>(
+      dim3 grid((size_n + 63) / 64, (size_m + 127) / 128, 1);
-          a, b_q_weight, b_qzeros, b_scales, c, size_m, size_n, size_k, groups,
+      if (size_k % 32 == 0 && groupsize >= 32) {
-          zero_offset);
+        gemm_q4_wmma_kernel_128x64_k32<T><<<grid, block, 0, stream>>>(
+            a, b_q_weight, b_qzeros, b_scales, c, size_m, size_n, size_k,
+            groups, zero_offset, /*partials=*/nullptr);
+      } else {
+        gemm_q4_wmma_kernel_128x64_k16<T><<<grid, block, 0, stream>>>(
+            a, b_q_weight, b_qzeros, b_scales, c, size_m, size_n, size_k,
+            groups, zero_offset, /*partials=*/nullptr);
+      }
+      return;
+    }
+    // Deterministic split-K, row-tiled so the FP32 scratch bound
+    //   scratch_bytes = k_split * TILE_M * size_n * 4
+    // stays independent of the caller's M (k_split <= 4 here).
+    constexpr int TILE_M = 512;
+    at::Tensor partials =
+        alloc_wmma_partials(k_split, std::min(TILE_M, size_m), size_n);
+    float* partials_ptr = partials.data_ptr<float>();
+    for (int row0 = 0; row0 < size_m; row0 += TILE_M) {
+      const int rows = std::min(TILE_M, size_m - row0);
+      const T* a_t = a + (long)row0 * size_k;
+      T* c_t = c + (long)row0 * size_n;
+      dim3 grid((size_n + 63) / 64, (rows + 127) / 128, k_split);
+      if (size_k % 32 == 0 && groupsize >= 32 && (size_k / k_split) % 32 == 0) {
+        gemm_q4_wmma_kernel_128x64_k32<T><<<grid, block, 0, stream>>>(
+            a_t, b_q_weight, b_qzeros, b_scales, c_t, rows, size_n, size_k,
+            groups, zero_offset, partials_ptr);
+      } else {
+        gemm_q4_wmma_kernel_128x64_k16<T><<<grid, block, 0, stream>>>(
+            a_t, b_q_weight, b_qzeros, b_scales, c_t, rows, size_n, size_k,
+            groups, zero_offset, partials_ptr);
+      }
+      launch_wmma_reduce<T>(partials, c_t, k_split, rows, size_n, stream);
     }
     return;
   }
@@ -2010,10 +2288,29 @@ void launch_gemm_q4_wmma_64x64_4w(const T* a, const uint32_t* b_q_weight,
   // 4 waves per block (128 threads), 64M × 64N tile per block.
   const int k_split = compute_wmma_k_split_mn(size_m, size_n, size_k, 64, 64);
   dim3 block(128);
-  dim3 grid((size_n + 63) / 64, (size_m + 63) / 64, k_split);
+  if (k_split == 1) {
-  gemm_q4_wmma_kernel_64x64_4w<T>
+    // Single writer per cell: the kernel's direct-store path is already
-      <<<grid, block, 0, stream>>>(a, b_q_weight, b_qzeros, b_scales, c, size_m,
+    // deterministic — no scratch, no reduce pass.
-                                   size_n, size_k, groups, zero_offset);
+    dim3 grid((size_n + 63) / 64, (size_m + 63) / 64, 1);
+    gemm_q4_wmma_kernel_64x64_4w<T><<<grid, block, 0, stream>>>(
+        a, b_q_weight, b_qzeros, b_scales, c, size_m, size_n, size_k, groups,
+        zero_offset, /*partials=*/nullptr);
+    return;
+  }
+  constexpr int TILE_M = 512;
+  at::Tensor partials =
+      alloc_wmma_partials(k_split, std::min(TILE_M, size_m), size_n);
+  float* partials_ptr = partials.data_ptr<float>();
+  for (int row0 = 0; row0 < size_m; row0 += TILE_M) {
+    const int rows = std::min(TILE_M, size_m - row0);
+    const T* a_t = a + (long)row0 * size_k;
+    T* c_t = c + (long)row0 * size_n;
+    dim3 grid((size_n + 63) / 64, (rows + 63) / 64, k_split);
+    gemm_q4_wmma_kernel_64x64_4w<T><<<grid, block, 0, stream>>>(
+        a_t, b_q_weight, b_qzeros, b_scales, c_t, rows, size_n, size_k, groups,
+        zero_offset, partials_ptr);
+    launch_wmma_reduce<T>(partials, c_t, k_split, rows, size_n, stream);
+  }
 }
 
 }  // namespace gptq_rdna3_wmma
@@ -2070,11 +2367,13 @@ torch::Tensor gptq_gemm_rdna3_wmma(torch::Tensor a, torch::Tensor b_q_weight,
   TORCH_CHECK(size_k % 16 == 0, "WMMA path requires K % 16 == 0");
 
   auto opts = torch::TensorOptions().dtype(a.dtype()).device(a.device());
-  // Always zero-init the output: some V3-V8 boundary threads may exit
+  // Every path writes each output element exactly once: direct store when
-  // without writing their output cell (e.g. out_m >= size_m), leaving
+  // k_split == 1, otherwise the fixed-order FP32 reduce over per-split
-  // uninitialized garbage when torch::empty is used.  The cost is
+  // partials (whose scratch coverage invariant is documented at
-  // negligible (< 1.5% of prefill time on gfx1100).
+  // alloc_wmma_partials). c therefore needs no zero-initialization.
-  at::Tensor c = torch::zeros({size_m, size_n}, opts);
+  // One-way door: the legacy CAS epilogue ADDS into c and would require
+  // torch::zeros here if it were ever re-selected.
+  at::Tensor c = torch::empty({size_m, size_n}, opts);
 
   const int zero_offset = use_v2_format ? 0 : 1;
 

csrc/rocm/qdq_4_rdna3.cuh

@@ -52,61 +52,75 @@ __forceinline__ __device__ void shuffle_4bit_8(uint32_t* q) {
 }
 
 // ---------------------------------------------------------------------------
-// fp16 path
+// fp16 path: exact factored form.
+// ---------------------------------------------------------------------------
+// The pre-#54706-v2 kernels baked the zero/scale bias into every weight,
+// which forced z = scale * (-1024 - zero) to be *stored* as fp16. At
+// scale ~0.02 that costs ~0.008 abs per (group, column) constant and the
+// error then accumulates along the K/groups axis (measured max abs
+// error ~1.0 vs an FP32 dequantized reference at K=4096). Computing the
+// product in fp32 does not help: the error is in the fp16 *representation*
+// of z, not in the multiply.
+//
+// Instead use the identity, exact for any grouping:
+//
+//     sum_k a_k * w_k  =  y * sum_k a_k * (1024 + q_k)  +  z * sum_k a_k
+//
+// with one (y, z) pair per nibble slot (the high pairs carry a factor of
+// 16): y_lo = s, z_lo = s*(-1024 - zero), y_hi = s/16, z_hi = s*(-64 -
+// zero). The y/z constants live in fp32 and are never narrowed; the magic
+// values (1024 + q) / (1024 + 16q) stay in fp16 where they are exact. The
+// z * sum_a correction is applied once per 32-K round in the kernel (z is
+// constant within a group, so per-round applications sum to the per-group
+// application). sum_k a_k does not depend on the output column, so the
+// extra v_dot2 against half2(1,1) is shared by all four columns a thread
+// owns. This is the same structure the bf16 branches of q_gemm_rdna3.cu
+// already use; fp16 needs two (y, z) pairs and two activation sums because
+// its magic values carry the upper-nibble *16 trick.
 // ---------------------------------------------------------------------------
 
-// Precompute scale-baked constants for a single zero/scale pair.
+// fp32 (y, z) pairs for the exact factored form above.
-//   z1z16[0] = scale * (-1024 - zero)            (used for "low" pairs)
+__forceinline__ __device__ void prep_zero_scale_fp16_f32(uint32_t zero,
-//   z1z16[1] = scale * (-64   - zero)            (used for "high" pairs)
+                                                         half scale,
-//   y1y16[0] = scale * 1                          (low pairs are q + 1024)
+                                                         float (&y)[2],
-//   y1y16[1] = scale * (1/16)                     (high pairs are q*16 + 1024)
+                                                         float (&z)[2]) {
-__forceinline__ __device__ void prep_zero_scale_fp16(uint32_t zero, half scale,
+  const float s = __half2float(scale);
-                                                     half2 (&z1z16)[2],
+  const float zf = (float)(int)zero;
-                                                     half2 (&y1y16)[2]) {
+  y[0] = s;  // low pairs:  (1024 + q) * s
-  // half(-1024 - zero) via the exllamav2 bit-trick:
+  z[0] = s * (-1024.0f - zf);
-  //   half bits 0xE400 == -1024.0 ; ORing the zero into mantissa subtracts it.
+  y[1] = s * (1.0f / 16.0f);  // high pairs: (1024 + 16q) * s/16
-  union {
+  z[1] = s * (-64.0f - zf);
-    uint16_t u;
-    half h;
-  } z1u;
-  z1u.u = (uint16_t)(0xE400 | zero);
-  half z1 = z1u.h;
-  half z16 = __hsub(__int2half_rn(-64), __int2half_rn((int)zero));
-
-  half2 scale2 = __half2half2(scale);
-  z1z16[0] = __hmul2(scale2, __half2half2(z1));
-  z1z16[1] = __hmul2(scale2, __half2half2(z16));
-
-  half y1 = __float2half_rn(1.0f);
-  half y16 = __float2half_rn(1.0f / 16.0f);
-  y1y16[0] = __hmul2(scale2, __half2half2(y1));
-  y1y16[1] = __hmul2(scale2, __half2half2(y16));
 }
 
-// Dequantize one int32 (8 shuffled 4-bit weights) into 4 half2 pairs:
+// Magic values only, with no y/z folded in: low pairs land in [0] and [2],
-//   dq[0] = (q[0], q[1]) * scale - zero*scale
+// high pairs in [1] and [3]. The lane mapping assumes the host-side
-//   dq[1] = (q[2], q[3]) * scale - zero*scale
+// gptq_shuffle interleave (see shuffle_4bit_8 above): even elements sit in
-//   dq[2] = (q[4], q[5]) * scale - zero*scale
+// the low nibbles, odd elements in the high nibbles.
-//   dq[3] = (q[6], q[7]) * scale - zero*scale
+__forceinline__ __device__ void magic_4bit_8_fp16(uint32_t qa, half2 (&qm)[4]) {
-__forceinline__ __device__ void dequant_4bit_8_fp16(uint32_t qa, half2 (&dq)[4],
-                                                    half2 (&z1z16)[2],
-                                                    half2 (&y1y16)[2]) {
   const uint32_t c0 = 0x64006400;
-
   union {
     uint32_t u;
     half2 h2;
-  } q0, q1, q2, q3;
+  } t;
-  q0.u = (qa & 0x000F000F) | c0;  // half2(q[0]+1024, q[1]+1024)
+  t.u = (qa & 0x000F000F) | c0;  // half2(1024 + q[0], 1024 + q[1])
-  q1.u = (qa & 0x00F000F0) | c0;  // half2(q[2]*16+1024, q[3]*16+1024)
+  qm[0] = t.h2;
-  uint32_t qa_hi = qa >> 8;
+  t.u = (qa & 0x00F000F0) | c0;  // half2(1024 + q[2]*16, 1024 + q[3]*16)
-  q2.u = (qa_hi & 0x000F000F) | c0;  // half2(q[4]+1024, q[5]+1024)
+  qm[1] = t.h2;
-  q3.u = (qa_hi & 0x00F000F0) | c0;  // half2(q[6]*16+1024, q[7]*16+1024)
+  const uint32_t qa_hi = qa >> 8;
+  t.u = (qa_hi & 0x000F000F) | c0;  // half2(1024 + q[4], 1024 + q[5])
+  qm[2] = t.h2;
+  t.u = (qa_hi & 0x00F000F0) | c0;  // half2(1024 + q[6]*16, 1024 + q[7]*16)
+  qm[3] = t.h2;
+}
 
-  dq[0] = __hfma2(q0.h2, y1y16[0], z1z16[0]);
+// half2(1.0, 1.0) for the activation-sum v_dot2 in the exact factored form.
-  dq[1] = __hfma2(q1.h2, y1y16[1], z1z16[1]);
+__forceinline__ __device__ half2 ones_half2_fp16() {
-  dq[2] = __hfma2(q2.h2, y1y16[0], z1z16[0]);
+  union {
-  dq[3] = __hfma2(q3.h2, y1y16[1], z1z16[1]);
+    uint32_t u;
+    half2 h2;
+  } t;
+  t.u = 0x3C003C00u;
+  return t.h2;
 }
 
 // ---------------------------------------------------------------------------

docs/features/batch_invariance.md

@@ -130,7 +130,8 @@ Batch invariance has been tested and verified on the following models:
 - **Llama 3**: Llama3.1 and 3.2 series, `meta-llama/Llama-3.2-1B-Instruct`, `meta-llama/Llama-3.2-3B-Instruct` for example
 - **GPT-OSS**: `openai/gpt-oss-20b`, `openai/gpt-oss-120b`
 - **Mistral**: `mistralai/Mistral-7B-v0.3`
-- **Phi series**: `microsoft/Phi-3.5-mini-instruct`
+- **GLM-4 series**: `zai-org/glm-4-9b-chat`
+- **Phi series**: `microsoft/Phi-3.5-mini-instruct`, `microsoft/phi-4`
 - **Granite 3.1 (MoE)**: `ibm-granite/granite-3.1-1b-a400m-instruct`, `ibm-granite/granite-3.1-3b-a800m-instruct`
 - **Granite 3.1 (Dense)**: `ibm-granite/granite-3.1-2b-instruct`, `ibm-granite/granite-3.1-8b-instruct`
 - **EXAONE 4.0 series**: `LGAI-EXAONE/EXAONE-4.0-1.2B`, `LGAI-EXAONE/EXAONE-4.0.1-32B`, `LGAI-EXAONE/EXAONE-4.0-32B`
@@ -152,6 +153,7 @@ When batch invariance is enabled, vLLM:
 5. On CUDA devices with tuned matmul table entries for the model's bf16 unquantized forward linear layers (Ada, Hopper, Blackwell),
    runs without `torch.compile` using breakable CUDA graphs so tile configs follow the runtime batch size; set `VLLM_USE_BREAKABLE_CUDAGRAPH=0` to opt out
    (not applied when sequence parallelism / async TP are enabled, since those are `torch.compile` passes).
+6. Splits the LoRA shrink reduction into 8 fixed K chunks that are summed in a fixed order instead of with atomics
 
 !!! warning
     Batch invariance under tensor parallelism is not yet supported for all-reduces whose size isn't a multiple of 16 bytes, for example a hidden size that isn't a multiple of 8 in fp16/bf16. Such tensors can switch between NCCL and custom all-reduce depending on batch size. All validated models meet this requirement.

docs/features/sleep_mode.md

@@ -101,6 +101,41 @@ llm.release_kv_cache_memory()
 llm.wake_up(tags=["kv_cache"])  # Reallocate KV cache and resume scheduling.
 ```
 
+#### Offloading CUDA graph memory
+
+By default, CUDA graph memory stays on the GPU while asleep. With
+`sleep_mode_offload_cudagraph=True` (off by default), CUDA graphs are captured
+into a cuMem pool that sleep backs up to CPU memory at both levels and any wake
+restores in place, so graphs are replayed, not recaptured. It needs the default
+`cumem` backend, CUDA and CUDA graphs; otherwise it has no effect.
+
+```python
+llm = LLM("Qwen/Qwen3-8B", enable_sleep_mode=True, sleep_mode_offload_cudagraph=True)
+```
+
+or `vllm serve <model> --enable-sleep-mode --sleep-mode-offload-cudagraph`.
+
+The cost is pinned host memory for the pool's backup, also at level 2, and one
+extra copy per captured custom allreduce (about 3% decode latency at batch
+size 1). `NCCL_GRAPH_REGISTER` defaults to `0`, since NCCL graph registration
+would pin the pool; an explicit value is kept with a warning. Graph executables
+outside PyTorch pools stay resident.
+
+#### Releasing NCCL communicator memory
+
+By default, NCCL communicators keep their GPU buffers while asleep. With
+`enable_nccl_comm_suspend=True` (off by default, experimental), sleep releases
+them with `ncclCommSuspend` and wake restores them with `ncclCommResume`. The
+communicators keep their topology, so they are not re-created. It needs NCCL
+2.29.7 or newer; with an older library, a warning is logged and the memory stays
+on the GPU.
+
+```python
+llm = LLM("Qwen/Qwen3-8B", enable_sleep_mode=True, enable_nccl_comm_suspend=True)
+```
+
+or `vllm serve <model> --enable-sleep-mode --enable-nccl-comm-suspend`.
+
 ### Online Serving
 
 To enable sleep mode in a vLLM server you need to initialize it with the flag `VLLM_SERVER_DEV_MODE=1` and pass `--enable-sleep-mode` to the vLLM server.

requirements/common.txt

@@ -7,7 +7,7 @@ requests >= 2.33.0
 tqdm
 blake3
 py-cpuinfo
-transformers >= 5.16.1, < 5.18.0
+transformers >= 5.16.1, < 5.19.0
 huggingface_hub >= 1.31.0  # Required for the `huggingface_hub.utils.httpx` re-export
 tokenizers >= 0.23.2  # Required for fast incremental detokenization.
 safetensors >= 0.6.2  # MXFP4/MXFP6 dtype support (F8_E8M0, F4) added in 0.6.0: https://github.com/huggingface/safetensors/pull/611

requirements/cuda.txt

@@ -30,7 +30,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==0.1.8; platform_system == "Linux"
+tokenspeed-mla==0.2.15; platform_system == "Linux"
 
 # Humming kernels for quantization gemm
 humming-kernels[cu13]==0.1.16

requirements/test/cpu.txt

@@ -713,7 +713,7 @@ pathvalidate==3.2.1
     # via pytablewriter
 patsy==1.0.1
     # via statsmodels
-peft==0.19.1
+peft==0.21.1
     # via -r requirements/test/cuda.in
 perceptron==0.1.4
     # via -r requirements/test/cuda.in
@@ -1228,7 +1228,7 @@ traitlets==5.16.1
     # via
     #   jupyter-client
     #   jupyter-core
-transformers==5.17.0
+transformers==5.18.0
     # via
     #   -r requirements/test/../common.txt
     #   -r requirements/test/cuda.in

requirements/test/cuda.in

@@ -20,7 +20,7 @@ httpx
 librosa # required for audio tests
 vector_quantize_pytorch # required for minicpmo_26 test
 vocos # required for minicpmo_26 test
-peft>=0.19.1 # required for phi-4-mm 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)
 sentence-transformers>=5.2.0 # required for embedding tests
@@ -40,7 +40,7 @@ datamodel_code_generator>=0.64.0 # required for minicpm3 test
 nltk>=3.10.3 # pin transitive rouge-score/lm-eval dep for known GHSAs
 lm-eval[api]>=0.4.12 # required for model evaluation test
 mteb[bm25s]>=2, <3 # required for mteb test
-transformers==5.17.0
+transformers==5.18.0
 schemathesis>=4.0.0 # Required for openai schema test.
 # quantization
 buildkite-test-collector==0.1.9

requirements/test/cuda.txt

@@ -801,7 +801,7 @@ pathvalidate==3.2.1
     # via pytablewriter
 patsy==1.0.1
     # via statsmodels
-peft==0.19.1
+peft==0.21.1
     # via -r requirements/test/cuda.in
 perceptron==0.1.4
     # via -r requirements/test/cuda.in
@@ -1335,7 +1335,7 @@ traitlets==5.16.1
     # via
     #   jupyter-client
     #   jupyter-core
-transformers==5.17.0
+transformers==5.18.0
     # via
     #   -c requirements/common.txt
     #   -r requirements/test/../common.txt

requirements/test/nightly-torch.txt

@@ -29,7 +29,7 @@ opencv-python-headless >= 4.13.0 # required for video test
 datamodel_code_generator>=0.64.0 # required for minicpm3 test
 lm-eval[api]>=0.4.12 # required for model evaluation test
 mteb[bm25s]>=2, <3 # required for mteb test
-transformers==5.17.0
+transformers==5.18.0
 schemathesis>=4.0.0 # Required for openai schema test.
 # quantization
 buildkite-test-collector==0.1.9

requirements/test/rocm.in

@@ -19,7 +19,7 @@ httpx
 librosa # required for audio tests
 vector_quantize_pytorch # required for minicpmo_26 test
 vocos # required for minicpmo_26 test
-peft>=0.19.1 # required for phi-4-mm 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
 sentence-transformers>=5.2.0 # required for embedding tests
@@ -36,7 +36,7 @@ datamodel_code_generator>=0.64.0 # required for minicpm3 test
 nltk>=3.10.3 # pin transitive rouge-score/lm-eval dep for known GHSAs
 lm-eval[api]>=0.4.12 # required for model evaluation test
 mteb[bm25s]>=2, <3 # required for mteb test
-transformers==5.17.0
+transformers==5.18.0
 schemathesis>=4.0.0 # Required for openai schema test
 # quantization
 bitsandbytes==0.50.0

requirements/test/rocm.txt

@@ -773,7 +773,7 @@ pathvalidate==3.3.1
     # via pytablewriter
 patsy==1.0.2
     # via statsmodels
-peft==0.19.1
+peft==0.21.1
     # via -r requirements/test/rocm.in
 perceptron==0.1.4
     # via -r requirements/test/rocm.in
@@ -1282,7 +1282,7 @@ traitlets==5.16.1
     # via
     #   jupyter-client
     #   jupyter-core
-transformers==5.17.0
+transformers==5.18.0
     # via
     #   -c requirements/common.txt
     #   -r requirements/test/../common.txt

requirements/test/xpu.in

@@ -20,7 +20,7 @@ nltk>=3.10.3 # pin transitive rouge-score/lm-eval dep for known GHSAs
 modelscope<1.38
 runai-model-streamer[s3,gcs,azure]>=0.15.7
 tensorizer==2.10.1
-transformers==5.17.0
+transformers==5.18.0
 
 # --- Audio Processing ---
 av

requirements/test/xpu.txt

@@ -1134,7 +1134,7 @@ traitlets==5.16.1
     # via
     #   jupyter-client
     #   jupyter-core
-transformers==5.17.0
+transformers==5.18.0
     # via
     #   -c requirements/common.txt
     #   -r requirements/test/../common.txt

tests/basic_correctness/memory/cumem/test_cumem.py

@@ -423,3 +423,58 @@ def test_cumem_with_cudagraph():
 
     # output content is as expected
     assert torch.allclose(y, x + 1)
+
+
+@pytest.mark.parametrize("level", [1, 2], ids=["sleep-1", "sleep-2"])
+@create_new_process_for_each_test("fork" if current_platform.is_cuda() else "spawn")
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="cuMem CUDA graph pool")
+def test_cudagraph_pool_sleep(level):
+    """Routing, and the graph pool backed up at both sleep levels and restored
+    in place by any wake."""
+    from contextlib import nullcontext
+    from types import SimpleNamespace
+
+    import vllm.distributed.device_communicators.pynccl_allocator as nccl_alloc
+    from vllm.compilation.cudagraph_pool import (
+        capture_outside_cumem_pool,
+        capture_pool,
+    )
+    from vllm.device_allocator.sleep_mode_backend import CuMemBackend
+
+    allocator = get_mem_allocator_instance()
+    on, off = (SimpleNamespace(use_cumem_cudagraph_pool=v) for v in (True, False))
+    handle = current_platform.graph_pool_handle()
+    for cfg, ctx in ((off, nullcontext()), (on, capture_outside_cumem_pool())):
+        with ctx, capture_pool(handle, cfg) as used:
+            assert used == handle and allocator.current_tag != "cudagraph"
+    with allocator.use_memory_pool("weights"):
+        weight = torch.full((1 << 20,), 2.0, device=DEVICE_TYPE)
+    x = torch.ones_like(weight)
+
+    def capture() -> list:
+        graph = torch.cuda.CUDAGraph()
+        stream = torch.cuda.Stream()
+        with (
+            capture_pool(handle, on) as pool,
+            torch.cuda.graph(graph, pool=pool, stream=stream),
+        ):
+            assert nccl_alloc._graph_pool_id == pool != handle
+            const = torch.empty_like(x)  # Never written by replay: needs a backup.
+            y = x * weight + const
+        return [graph, y, const.fill_(3.0)]
+
+    def graph_pool() -> dict[int, int]:
+        data = allocator.pointer_to_data.items()
+        return {
+            p: d.handle[1] for p, d in data if d.tag == "cudagraph" and not d.is_asleep
+        }
+
+    held, backend = capture(), CuMemBackend()
+    mapped = graph_pool()
+    backend.suspend(level=level)
+    backend.resume(tags=["kv_cache"])
+    assert graph_pool() == mapped
+    backend.resume(tags=["weights"])
+    weight.fill_(2.0)  # Level 2 discards weights; emulate the reload.
+    held[0].replay()
+    assert torch.equal(held[1], torch.full_like(x, 5.0))

tests/basic_correctness/memory/sleep_mode/test_sleep_mode.py

@@ -3,6 +3,7 @@
 
 import asyncio
 import math
+import os
 
 import pytest
 import torch
@@ -14,7 +15,12 @@ from vllm.model_executor.model_loader import get_model_loader
 from vllm.platforms import current_platform
 from vllm.utils.mem_constants import GiB_bytes
 
-from ....utils import create_new_process_for_each_test, multi_gpu_test, requires_fp8
+from ....utils import (
+    create_new_process_for_each_test,
+    multi_gpu_marks,
+    multi_gpu_test,
+    requires_fp8,
+)
 
 
 @create_new_process_for_each_test("fork" if current_platform.is_cuda() else "spawn")
@@ -195,6 +201,73 @@ def test_deep_sleep_lora_tp2(monkeypatch):
     assert output[0].outputs[0].text == output2[0].outputs[0].text
 
 
+def _cudagraph_bytes(worker) -> int:
+    data = cumem.CuMemAllocator.get_instance().pointer_to_data.values()
+    return sum(d.handle[1] for d in data if d.tag == "cudagraph")
+
+
+def _custom_ar_active(worker) -> bool:
+    from vllm.distributed.parallel_state import get_tp_group
+
+    ca_comm = get_tp_group().device_communicator.ca_comm
+    return ca_comm is not None and not ca_comm.disabled
+
+
+@pytest.mark.parametrize(
+    ("mode", "breakable", "tp", "offload"),
+    [
+        ("FULL", False, 1, True),
+        ("PIECEWISE", False, 1, True),
+        ("PIECEWISE", True, 1, True),
+        pytest.param(
+            "FULL_AND_PIECEWISE", False, 2, True, marks=multi_gpu_marks(num_gpus=2)
+        ),
+        ("FULL_AND_PIECEWISE", False, 1, False),
+    ],
+    ids=["full", "piecewise", "breakable", "custom-ar-tp2", "off-by-default"],
+)
+@create_new_process_for_each_test()
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="cuMem CUDA graph pool")
+def test_sleep_cudagraph_pool(monkeypatch, mode, breakable, tp, offload):
+    """Each capture site uses the pool, exact across sleeps; TP=2 needs
+    unregistered custom AR; off changes nothing."""
+    for name, value in [
+        ("VLLM_USE_V2_MODEL_RUNNER", "1"),
+        ("VLLM_USE_BREAKABLE_CUDAGRAPH", "1" if breakable else "0"),
+        ("VLLM_ALLOW_INSECURE_SERIALIZATION", "1"),
+        ("VLLM_ALLREDUCE_USE_FLASHINFER", "0"),
+        ("VLLM_ALLREDUCE_USE_SYMM_MEM", "0"),
+    ]:
+        monkeypatch.setenv(name, value)
+    monkeypatch.delenv("NCCL_GRAPH_REGISTER", raising=False)
+    llm = LLM(
+        "hmellor/tiny-random-LlamaForCausalLM",
+        enable_sleep_mode=True,
+        sleep_mode_offload_cudagraph=offload,
+        tensor_parallel_size=tp,
+        compilation_config={
+            "cudagraph_mode": mode,
+            "pass_config": {"fuse_allreduce_rms": False},
+        },
+    )
+    graph_bytes = llm.collective_rpc(_cudagraph_bytes)
+    assert os.environ.get("NCCL_GRAPH_REGISTER") == ("0" if offload else None)
+    assert all(graph_bytes) if offload else not any(graph_bytes)
+    if not offload:
+        return
+    if tp > 1 and not all(llm.collective_rpc(_custom_ar_active)):
+        pytest.skip("Custom allreduce is unavailable on these GPUs")
+    prompt, params = "How are you?", SamplingParams(temperature=0, max_tokens=10)
+    expected = llm.generate(prompt, params)[0].outputs[0].text
+    for level in (1, 2):
+        llm.sleep(level=level)
+        llm.wake_up(tags=["weights"])
+        if level == 2:
+            llm.collective_rpc("reload_weights")
+        llm.wake_up(tags=["kv_cache"])
+        assert llm.generate(prompt, params)[0].outputs[0].text == expected
+
+
 @create_new_process_for_each_test()
 def test_deep_sleep_async(gpu_memory_cleared):
     async def test():

tests/distributed/test_comm_ops.py

@@ -376,7 +376,6 @@ def test_flashinfer_standalone_workspace_size(
         "_resolve_fi_ar_backend",
         Mock(return_value=("mnnvl", False)),
     )
-    monkeypatch.setattr(flashinfer_all_reduce, "get_node_count", lambda: 2)
     monkeypatch.setattr(
         flashinfer_all_reduce,
         "_get_tuned_standalone_max_size",
@@ -389,6 +388,34 @@ def test_flashinfer_standalone_workspace_size(
     assert create_workspace.call_args.args[3] == expected
 
 
+@pytest.mark.parametrize(
+    ("node_count", "is_blackwell", "expected"),
+    [
+        (1, False, ("trtllm", False)),
+        (1, True, ("trtllm", False)),
+        (2, False, (None, False)),
+        (2, True, ("mnnvl", False)),
+    ],
+)
+def test_flashinfer_auto_backend_uses_group_topology(
+    monkeypatch: pytest.MonkeyPatch,
+    node_count: int,
+    is_blackwell: bool,
+    expected: tuple[str | None, bool],
+) -> None:
+    monkeypatch.setattr(
+        flashinfer_all_reduce.envs, "VLLM_FLASHINFER_ALLREDUCE_BACKEND", "auto"
+    )
+    monkeypatch.setattr(flashinfer_all_reduce, "_node_count", lambda _: node_count)
+    monkeypatch.setattr(
+        flashinfer_all_reduce.current_platform,
+        "has_device_capability",
+        lambda capability: is_blackwell,
+    )
+
+    assert flashinfer_all_reduce._resolve_fi_ar_backend(Mock()) == expected
+
+
 def test_flashinfer_workspace_failure_is_not_retried(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
@@ -402,7 +429,6 @@ def test_flashinfer_workspace_failure_is_not_retried(
         "_resolve_fi_ar_backend",
         Mock(return_value=("mnnvl", True)),
     )
-    monkeypatch.setattr(flashinfer_all_reduce, "get_node_count", lambda: 1)
     monkeypatch.setattr(
         flashinfer_all_reduce, "_get_tuned_standalone_max_size", Mock(return_value=None)
     )

tests/distributed/test_engram_dp_shard.py

@@ -177,7 +177,9 @@ def _worker(rank: int, tp_size: int, port: int) -> None:
     # Supply the resolved Engram settings without constructing a full model.
     vllm_config.engram_config = EngramConfig()
     vllm_config.model_config = SimpleNamespace(
-        architecture="DeepseekV41ForCausalLM", is_moe=True
+        architecture="DeepseekV41ForCausalLM",
+        is_moe=True,
+        sleep_mode_offload_cudagraph=False,
     )
     try:
         init_distributed_environment()

tests/entrypoints/openai/chat_completion/test_chat_error.py

@@ -248,7 +248,8 @@ async def test_renderer_only_chat_request_skips_mm_cache():
 
 
 @pytest.mark.asyncio
-async def test_chat_error_stream():
+@pytest.mark.parametrize("include_usage", [False, True])
+async def test_chat_error_stream(include_usage: bool):
     """Test finish_reason='error' returns 500 InternalServerError (streaming)."""
     mock_engine = MagicMock(spec=AsyncLLM)
     mock_engine.errored = False
@@ -313,6 +314,67 @@ async def test_chat_error_stream():
         messages=[{"role": "user", "content": "Test prompt"}],
         max_tokens=10,
         stream=True,
+        stream_options={"include_usage": include_usage},
+    )
+
+    response = await serving_chat.create_chat_completion(request)
+
+    chunks = []
+    async for chunk in response:
+        chunks.append(chunk)
+
+    assert len(chunks) >= 2
+    assert any("Internal server error" in chunk for chunk in chunks), (
+        f"Expected error message in chunks: {chunks}"
+    )
+    assert chunks[-1] == "data: [DONE]\n\n"
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("include_usage", [False, True])
+async def test_chat_error_stream_first_chunk(include_usage: bool):
+    """An error before the first token must reach the streaming client."""
+    mock_engine = MagicMock(spec=AsyncLLM)
+    mock_engine.errored = False
+    mock_engine.model_config = MockModelConfig()
+    mock_engine.input_processor = MagicMock()
+    mock_engine.renderer = _build_renderer(mock_engine.model_config)
+
+    serving_chat = _build_serving_chat(mock_engine)
+
+    completion_output = CompletionOutput(
+        index=0,
+        text="",
+        token_ids=[],
+        cumulative_logprob=None,
+        logprobs=None,
+        finish_reason="error",
+    )
+
+    request_output = RequestOutput(
+        request_id="test-id",
+        prompt="Test prompt",
+        prompt_token_ids=[1, 2, 3],
+        prompt_logprobs=None,
+        outputs=[completion_output],
+        finished=True,
+        metrics=None,
+        lora_request=None,
+        encoder_prompt=None,
+        encoder_prompt_token_ids=None,
+    )
+
+    async def mock_generate(*args, **kwargs):
+        yield request_output
+
+    mock_engine.generate = MagicMock(side_effect=mock_generate)
+
+    request = ChatCompletionRequest(
+        model=MODEL_NAME,
+        messages=[{"role": "user", "content": "Test prompt"}],
+        max_tokens=10,
+        stream=True,
+        stream_options={"include_usage": include_usage},
     )
 
     response = await serving_chat.create_chat_completion(request)

tests/entrypoints/openai/completion/test_completion_error.py

@@ -383,10 +383,14 @@ async def test_completion_suffix_supports_text_prompt_list():
 
     assert result == [{"ok": True}]
     assert online_renderer.preprocess_completion.call_args.kwargs["prompt_input"] == [
-        f"{DEEPSEEK_V4_FIM_BEGIN}prefix A"
+        (
-        f"{DEEPSEEK_V4_FIM_HOLE} suffix{DEEPSEEK_V4_FIM_END}",
+            f"{DEEPSEEK_V4_FIM_BEGIN}prefix A"
-        f"{DEEPSEEK_V4_FIM_BEGIN}prefix B"
+            f"{DEEPSEEK_V4_FIM_HOLE} suffix{DEEPSEEK_V4_FIM_END}"
-        f"{DEEPSEEK_V4_FIM_HOLE} suffix{DEEPSEEK_V4_FIM_END}",
+        ),
+        (
+            f"{DEEPSEEK_V4_FIM_BEGIN}prefix B"
+            f"{DEEPSEEK_V4_FIM_HOLE} suffix{DEEPSEEK_V4_FIM_END}"
+        ),
     ]
 
 
@@ -619,7 +623,8 @@ async def test_renderer_only_completion_request_skips_mm_cache():
 
 
 @pytest.mark.asyncio
-async def test_completion_error_stream():
+@pytest.mark.parametrize("include_usage", [False, True])
+async def test_completion_error_stream(include_usage: bool):
     """Test finish_reason='error' returns 500 InternalServerError (streaming)."""
     mock_engine = MagicMock(spec=AsyncLLM)
     mock_engine.errored = False
@@ -684,6 +689,67 @@ async def test_completion_error_stream():
         prompt="Test prompt",
         max_tokens=10,
         stream=True,
+        stream_options={"include_usage": include_usage},
+    )
+
+    response = await serving_completion.create_completion(request)
+
+    chunks = []
+    async for chunk in response:
+        chunks.append(chunk)
+
+    assert len(chunks) >= 2
+    assert any("Internal server error" in chunk for chunk in chunks), (
+        f"Expected error message in chunks: {chunks}"
+    )
+    assert chunks[-1] == "data: [DONE]\n\n"
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("include_usage", [False, True])
+async def test_completion_error_stream_first_chunk(include_usage: bool):
+    """An error before the first token must reach the streaming client."""
+    mock_engine = MagicMock(spec=AsyncLLM)
+    mock_engine.errored = False
+    mock_engine.model_config = MockModelConfig()
+    mock_engine.input_processor = MagicMock()
+    mock_engine.renderer = _build_renderer(mock_engine.model_config)
+
+    serving_completion = _build_serving_completion(mock_engine)
+
+    completion_output = CompletionOutput(
+        index=0,
+        text="",
+        token_ids=[],
+        cumulative_logprob=None,
+        logprobs=None,
+        finish_reason="error",
+    )
+
+    request_output = RequestOutput(
+        request_id="test-id",
+        prompt="Test prompt",
+        prompt_token_ids=[1, 2, 3],
+        prompt_logprobs=None,
+        outputs=[completion_output],
+        finished=True,
+        metrics=None,
+        lora_request=None,
+        encoder_prompt=None,
+        encoder_prompt_token_ids=None,
+    )
+
+    async def mock_generate(*args, **kwargs):
+        yield request_output
+
+    mock_engine.generate = MagicMock(side_effect=mock_generate)
+
+    request = CompletionRequest(
+        model=MODEL_NAME,
+        prompt="Test prompt",
+        max_tokens=10,
+        stream=True,
+        stream_options={"include_usage": include_usage},
     )
 
     response = await serving_completion.create_completion(request)

tests/entrypoints/openai/responses/test_serving_responses.py

@@ -67,7 +67,7 @@ from vllm.entrypoints.serve.engine.protocol import ErrorResponse
 from vllm.inputs import tokens_input
 from vllm.logprobs import Logprob as SampleLogprob
 from vllm.outputs import CompletionOutput, RequestOutput
-from vllm.parser.harmony import Segment
+from vllm.parser.harmony import ChunkResult, HarmonyParser, Segment
 from vllm.renderers import TokenizeParams
 from vllm.renderers.online_renderer import (
     OnlineRenderer,
@@ -2129,6 +2129,87 @@ async def test_stream_completed_response_reuses_streamed_items(monkeypatch):
     assert function_call.call_id == "chatcmpl-tool-parser-id"
 
 
+def _harmony_msg(channel: str, text: str, recipient: str | None = None):
+    msg = OpenAIHarmonyMessage.from_role_and_content(Role.ASSISTANT, text)
+    msg = msg.with_channel(channel)
+    return msg.with_recipient(recipient) if recipient else msg
+
+
+def _harmony_segments(msg, *, streamed: bool = True) -> list[Segment]:
+    """Segments the parser yields for one message: a content delta (absent
+    for a zero-delta message), then the completed message."""
+    delta = [Segment(msg.channel, msg.recipient, msg.content[0].text)]
+    completed = Segment(msg.channel, msg.recipient, "", completed_message=msg)
+    return (delta if streamed else []) + [completed]
+
+
+async def _harmony_stream_events(segment_chunks: list[list[Segment]]):
+    serving = _make_serving_instance()
+    serving.use_harmony = True
+    parser = MagicMock(spec=HarmonyParser)
+    parser.process_chunk.side_effect = [
+        ChunkResult(segments=segments, reasoning_token_count=0)
+        for segments in segment_chunks
+    ]
+    context = HarmonyContext([], [], frozenset({"get_weather"}), parser)
+
+    async def result_generator():
+        for token_id in range(len(segment_chunks)):
+            context.append_output(_make_request_output("", [token_id]))
+            yield context
+
+    return [
+        event
+        async for event in serving.responses_stream_generator(
+            request=ResponsesRequest(input="hi", stream=True, store=False),
+            sampling_params=SamplingParams(max_tokens=16),
+            result_generator=result_generator(),
+            context=context,
+            model_name="test-model",
+            tokenizer=MagicMock(),
+            request_metadata=RequestResponseMetadata(request_id="req"),
+        )
+    ]
+
+
+@pytest.mark.asyncio
+async def test_harmony_stream_completed_response_reuses_streamed_ids():
+    """response.completed must carry the id and call_id each item was streamed
+    with, not ones minted when the harmony messages are converted again."""
+    messages = [
+        _harmony_msg("analysis", "think"),
+        _harmony_msg("commentary", '{"city": "Paris"}', "functions.get_weather"),
+        _harmony_msg("final", "Sunny"),
+    ]
+
+    events = await _harmony_stream_events([_harmony_segments(msg) for msg in messages])
+
+    streamed = [e.item for e in events if e.type == "response.output_item.done"]
+    output = events[-1].response.output
+    assert [item.type for item in output] == ["reasoning", "function_call", "message"]
+    assert [item.id for item in output] == [item.id for item in streamed]
+    assert output[1].call_id == streamed[1].call_id
+    assert output[2].content[0].text == "Sunny"
+
+
+@pytest.mark.asyncio
+async def test_harmony_stream_unstreamed_item_keeps_own_id():
+    """An item with no done event (zero-delta) must not take the streamed id
+    of a later item of the same type."""
+    unstreamed = _harmony_msg("analysis", "skipped")
+    reasoning = _harmony_msg("analysis", "think")
+
+    events = await _harmony_stream_events(
+        [_harmony_segments(unstreamed, streamed=False), _harmony_segments(reasoning)]
+    )
+
+    (streamed,) = [e.item for e in events if e.type == "response.output_item.done"]
+    first, second = events[-1].response.output
+    assert first.content[0].text == "skipped"
+    assert first.id != streamed.id
+    assert second.id == streamed.id
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize(("parallel_tool_calls", "num_calls"), [(False, 1), (None, 3)])
 async def test_streaming_parallel_tool_calls(

tests/entrypoints/scale_out/token_in_token_out/test_generate_stream.py

@@ -456,11 +456,13 @@ async def test_stream_returns_spec_decode_metrics_on_empty_terminal_output(
     data_chunks = [chunk for chunk in _parse_sse_chunks(chunks) if chunk != "[DONE]"]
 
     assert "metrics" not in data_chunks[0]
-    assert len(data_chunks) == 1 + int(include_usage)
+    assert data_chunks[1]["choices"][0]["finish_reason"] == "stop"
+    assert "metrics" not in data_chunks[1]
+    assert len(data_chunks) == 2 + int(include_usage)
     if include_usage:
-        assert data_chunks[1]["choices"] == []
+        assert data_chunks[2]["choices"] == []
         assert (
-            data_chunks[1]["metrics"]["speculative_decoding"]["num_draft_tokens"] == 3
+            data_chunks[2]["metrics"]["speculative_decoding"]["num_draft_tokens"] == 3
         )
 
 
@@ -823,30 +825,76 @@ async def test_stream_zero_token_completion_still_delivers_prompt_metadata():
 
 
 @pytest.mark.asyncio
-async def test_stream_zero_token_completion_emits_no_chunk_by_default():
+@pytest.mark.parametrize("finish_reason", ["stop", "length", "abort"])
+async def test_stream_emits_terminal_output_without_new_tokens(finish_reason):
+    """The final output may carry no new tokens (always for an abort); the
+    client must still see why the choice finished."""
     engine = _mock_engine()
 
     async def mock_generate(*args, **kwargs):
+        yield _make_request_output("req-1", token_ids=[10])
         yield _make_request_output(
-            "req-1", token_ids=[], finish_reason="stop", finished=True
+            "req-1", token_ids=[], finish_reason=finish_reason, finished=True
         )
 
     engine.generate = MagicMock(side_effect=mock_generate)
     serving = _build_serving_tokens(engine)
     request = GenerateRequest(
         token_ids=[1, 2, 3],
-        sampling_params=SamplingParams(max_tokens=1),
+        sampling_params=SamplingParams(max_tokens=10),
         model=MODEL_NAME,
         stream=True,
     )
 
     response = await serving.serve_tokens(request)
     parsed = _parse_sse_chunks([chunk async for chunk in response])
+    choices = [c["choices"][0] for c in parsed[:-1]]
 
-    assert not [c for c in parsed if isinstance(c, dict) and c.get("choices")]
+    assert [c["token_ids"] for c in choices] == [[10], []]
+    assert [c["finish_reason"] for c in choices] == [None, finish_reason]
     assert parsed[-1] == "[DONE]"
 
 
+@pytest.mark.asyncio
+async def test_stream_parallel_sampling_first_output_missing_a_choice():
+    """With n=2 the first output may carry only one choice; the other must
+    still stream, and an abort finishes both."""
+    engine = _mock_engine()
+
+    async def mock_generate(*args, **kwargs):
+        yield _make_request_output("req-1", token_ids=[10], index=0)
+        yield _make_request_output("req-1", token_ids=[11], index=1)
+        aborted = _make_request_output(
+            "req-1", token_ids=[], finish_reason="abort", finished=True
+        )
+        aborted.outputs += _make_request_output(
+            "req-1", token_ids=[], finish_reason="abort", index=1
+        ).outputs
+        yield aborted
+
+    engine.generate = MagicMock(side_effect=mock_generate)
+    serving = _build_serving_tokens(engine)
+    request = GenerateRequest(
+        token_ids=[1, 2, 3],
+        sampling_params=SamplingParams(max_tokens=10, n=2),
+        model=MODEL_NAME,
+        stream=True,
+        stream_options=StreamOptions(include_usage=True),
+    )
+
+    response = await serving.serve_tokens(request)
+    parsed = _parse_sse_chunks([chunk async for chunk in response])
+    choices = [c["choices"][0] for c in parsed[:-2]]
+
+    assert [(c["index"], c["token_ids"], c["finish_reason"]) for c in choices] == [
+        (0, [10], None),
+        (1, [11], None),
+        (0, [], "abort"),
+        (1, [], "abort"),
+    ]
+    assert parsed[-2]["usage"]["completion_tokens"] == 2
+
+
 @pytest.mark.asyncio
 async def test_stream_error_mid_generation():
     """finish_reason='error' mid-stream yields error chunk then [DONE]."""

tests/entrypoints/scale_out/token_in_token_out/test_serving_tokens.py

@@ -710,21 +710,23 @@ async def test_text_mode_rejected_when_tokens_only(client):
 @pytest.mark.asyncio
 @pytest.mark.skipif(
     envs.VLLM_USE_RUST_FRONTEND,
-    reason="output_mode is not supported by the Rust frontend",
+    reason="the Rust frontend does not serve /inference/v1/abort_requests",
 )
-async def test_text_mode_stream_delivers_abort_finish_chunk(
+@pytest.mark.parametrize("output_mode", ["tokens", "text"])
-    client, tokenizer, messages
+async def test_stream_delivers_abort_finish_chunk(
+    client, tokenizer, messages, output_mode
 ):
-    """The final output after an abort has no new token IDs, so it
+    """The final output after an abort has no new token IDs, but the
-    reaches a text stream only because text mode emits finish-only chunks."""
+    stream must still deliver its finish reason."""
     payload = {
         **_text_mode_payload(
             _chat_prompt_token_ids(tokenizer, messages),
             max_tokens=900,
             ignore_eos=True,
         ),
+        "output_mode": output_mode,
         "stream": True,
-        "request_id": "text-abort-e2e",
+        "request_id": f"{output_mode}-abort-e2e",
     }
 
     chunks = []
@@ -736,7 +738,7 @@ async def test_text_mode_stream_delivers_abort_finish_chunk(
                 continue
             chunk = json.loads(line[len("data: ") :])
             chunks.append(chunk)
-            if not aborted and chunk["choices"] and chunk["choices"][0]["text"]:
+            if not aborted and chunk["choices"] and chunk["choices"][0]["token_ids"]:
                 abort = await client.post(
                     ABORT_ENDPOINT,
                     json={"request_ids": [chunk["request_id"]]},
@@ -745,7 +747,7 @@ async def test_text_mode_stream_delivers_abort_finish_chunk(
                 aborted = True
 
     assert aborted
-    assert {chunk["output_mode"] for chunk in chunks} == {"text"}
+    assert {chunk["output_mode"] for chunk in chunks} == {output_mode}
     final_choice = chunks[-1]["choices"][0]
     assert final_choice["finish_reason"] == "abort"
     generated = sum(len(chunk["choices"][0]["token_ids"] or []) for chunk in chunks)

tests/entrypoints/serve/dev/rlhf/test_weight_operation_metrics.py

@@ -2,11 +2,8 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Weight-transfer metrics scraped from a real two-API-server deployment."""
 
-from pathlib import Path
-
 import requests
 import torch
-from prometheus_client.multiprocess import MultiProcessCollector
 from prometheus_client.parser import text_string_to_metric_families
 from transformers import AutoModelForCausalLM
 
@@ -21,10 +18,7 @@ from vllm.distributed.weight_transfer.ipc_engine import IPCTrainerInitInfo
 
 DURATION = "vllm:rl_weight_update_operation_duration_seconds"
 IN_FLIGHT = "vllm:rl_weight_update_operations_in_flight"
-# The API servers accept from one shared listening socket, so the kernel may
+SYNCS = 8
-# hand one process most connections. Sync until both have recorded, within a cap.
-MIN_SYNCS = 8
-MAX_SYNCS = 64
 
 
 def scrape(url: str) -> dict[tuple[str, str], float]:
@@ -37,17 +31,6 @@ def scrape(url: str) -> dict[tuple[str, str], float]:
     }
 
 
-def pids_that_recorded(multiproc_dir: Path) -> set[str]:
-    return {
-        path.stem.rsplit("_", 1)[1]
-        for path in multiproc_dir.glob("histogram_*.db")
-        if any(
-            metric.name == DURATION
-            for metric in MultiProcessCollector.merge([str(path)])
-        )
-    }
-
-
 def test_weight_sync_metrics_aggregate_across_api_servers(tmp_path):
     args = [
         "--load-format",
@@ -77,17 +60,11 @@ def test_weight_sync_metrics_aggregate_across_api_servers(tmp_path):
             client=HTTPVLLMWeightSyncClient(url),
             source=ModuleSource(trainer),
         )
-        syncs = 0
+        for _ in range(SYNCS):
-        while syncs < MIN_SYNCS or (
-            len(pids_that_recorded(tmp_path)) < 2 and syncs < MAX_SYNCS
-        ):
             engine.send_weights()
-            syncs += 1
         metrics = scrape(url)
 
-    calls = {"init": 1, "start": syncs, "update": syncs, "finish": syncs}
+    calls = {"init": 1, "start": SYNCS, "update": SYNCS, "finish": SYNCS}
     for operation, count in calls.items():
         assert metrics[(f"{DURATION}_count", operation)] == count, operation
         assert metrics[(IN_FLIGHT, operation)] == 0, operation
-    # The counts above are sums over processes only if both servers recorded.
-    assert len(pids_that_recorded(tmp_path)) == 2

tests/kernels/attention/test_prefix_prefill.py

@@ -5,6 +5,7 @@ import math
 import random
 import time
 from collections.abc import Callable
+from types import SimpleNamespace
 
 import pytest
 import torch
@@ -323,6 +324,63 @@ def test_contexted_kv_attention(
     torch.testing.assert_close(output, output_ref, atol=atol, rtol=0)
 
 
+@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm backend")
+@pytest.mark.parametrize("sliding_window", [None, 1, 4])
+@pytest.mark.parametrize("query_len", [1, 3])
+@torch.inference_mode()
+def test_rocm_backend_sliding_window_includes_boundary_token(
+    sliding_window: int | None, query_len: int
+) -> None:
+    """A window of W must include the current token and W - 1 preceding tokens."""
+    from vllm.v1.attention.backends.rocm_attn import RocmAttentionImpl
+    from vllm.v1.attention.ops.paged_attn import PagedAttention
+
+    device, dtype = "cuda:0", torch.float16
+    num_heads, head_size, seq_len = 4, 128, 8
+    query = torch.zeros(query_len, num_heads, head_size, dtype=dtype, device=device)
+    key = torch.zeros(seq_len, 1, head_size, dtype=dtype, device=device)
+    values = torch.arange(1, seq_len + 1, dtype=dtype, device=device)
+    value = values[:, None, None].expand_as(key)
+    kv_cache = torch.zeros(2, 1, 16, head_size, dtype=dtype, device=device)
+    _, value_cache = PagedAttention.split_kv_cache(kv_cache, 1, head_size)
+    value_cache[..., :seq_len] = value.permute(1, 2, 0)
+    scale = torch.tensor(1.0, device=device)
+    metadata = SimpleNamespace(
+        use_cascade=False,
+        num_actual_tokens=query_len,
+        query_start_loc=torch.tensor([0, query_len], dtype=torch.int32, device=device),
+        seq_lens=torch.tensor([seq_len], dtype=torch.int32, device=device),
+        max_query_len=query_len,
+        max_seq_len=seq_len,
+        block_table=torch.zeros(1, 1, dtype=torch.int32, device=device),
+        causal=True,
+    )
+    impl = RocmAttentionImpl(
+        num_heads=num_heads,
+        head_size=head_size,
+        scale=head_size**-0.5,
+        num_kv_heads=1,
+        alibi_slopes=None,
+        sliding_window=sliding_window,
+        kv_cache_dtype="auto",
+    )
+    output = torch.empty_like(query)
+    impl.forward(
+        SimpleNamespace(_k_scale=scale, _v_scale=scale),
+        query,
+        key[-query_len:],
+        value[-query_len:],
+        kv_cache.transpose(0, 1),
+        metadata,
+        output,
+    )
+    mask = create_causal_attention_mask_for_sdpa(
+        [query_len], [seq_len], sliding_window or 0, device=device, dtype=dtype
+    )
+    expected = (mask.float().softmax(dim=-1) @ values.float())[:, None, None]
+    torch.testing.assert_close(output, expected.to(dtype).expand_as(output))
+
+
 @pytest.mark.parametrize("num_heads", NUM_HEADS)
 @pytest.mark.parametrize("num_queries_per_kv", NUM_QUERIES_PER_KV)
 @pytest.mark.parametrize("head_size", HEAD_SIZES)

tests/kernels/quantization/test_rdna3_w4a16_determinism.py

@@ -0,0 +1,585 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Determinism + FP32-reference correctness tests for the RDNA3 W4A16 GEMM.
+
+The split-K epilogues of ``gptq_gemm_rdna3`` (scalar path) and the WMMA
+path previously accumulated per-split low-precision partials directly into
+the output via CAS atomics. With multiple writers per output element the
+accumulation order — and therefore the rounded result — varied per call:
+repeated fixed-input calls produced a different output almost every time on
+gfx1100 once more than a few split blocks contended.
+
+Both paths now store FP32 split partials and reduce them in a fixed order
+with a single final low-precision rounding, which is deterministic by
+construction. When the split count is 1 every path stores directly
+(single writer per output element).
+
+These tests assert, through the public op path:
+
+A. Repeatability — ``torch.equal`` across repeated identical calls.
+B. Numerical correctness — against an FP32 dequantized reference.
+
+(B) exists because repeatability alone cannot see the failure mode the
+deterministic epilogue introduces: the FP32 partials scratch is
+``at::empty`` (its coverage invariant is documented at
+``alloc_wmma_partials`` in ``csrc/rocm/q_gemm_rdna3_wmma.cu``), and across
+repeated calls the caching allocator hands back the same block — a slot
+that ever went unwritten would be bitwise-repeatable *and* wrong.
+
+Reference semantics (uint4b8, GPTQv1, synthesized zero points):
+    W[k, n] = (q[k, n] - 8) * float(round_to_dtype(scale[g, n])),  g = k//G
+    ref[m, n] = x[m, :] @ W[:, n]  (FP32 matmul)
+
+Tolerances are derived from the kernel's rounding structure, not tuned to
+pass (measured max errors on W7900/gfx1100 in parentheses):
+
+* scalar bf16  — dequant keeps FP32 precision end-to-end (magic-value
+  v_dot2 path); the only rounding is the final cast. Measured max ~0.06
+  on the tested shapes (|ref| up to ~30); tolerance 0.10 with headroom
+  for larger |ref| at the final-cast ulp.
+* scalar fp16  — the dequant uses the exact factored form (magic values in
+  fp16, (y, z) correction in fp32; see qdq_4_rdna3.cuh), so the only fp16
+  roundings are the final output cast and the B-tile-free dot inputs:
+  max_abs <= 0.5 * ulp_fp16(max|ref|) * 1.5 + 1e-3   (measured ~0.015 at
+  |ref|~37; the pre-fix baked-constant form measured ~1.0-1.4 on the same
+  shapes — 94-200x worse).
+* WMMA bf16    — B is narrowed to bf16 per cell (one rounding,
+  <= 0.5*ulp_bf16(|B|)), accumulated in FP32; measured max ≈ 0.09 at
+  K=4096 — tolerance 0.25.
+* WMMA fp16    — B narrowed to fp16 via the precise sub-then-mul variant
+  (one rounding, ~1e-4); measured max ≈ 0.015 — tolerance 0.05.
+
+All four are orders of magnitude below what an unwritten scratch slot
+(caching-allocator garbage / NaN) or an indexing bug (errors of the same
+scale as |ref| ≈ 7) would produce.
+
+Test hygiene: the op path needs a single-process distributed group and a
+current VLLM config; tests use the repo-standard ``dist_init`` fixture from
+``tests/conftest.py`` (real context manager, temp-file rendezvous so no
+hardcoded MASTER_PORT can collide in parallel CI, and the suite's standard
+post-test cleanup).
+"""
+
+import pytest
+import torch
+
+from vllm.platforms import current_platform
+
+if not current_platform.is_rocm():
+    pytest.skip("RDNA3 W4A16 kernel is ROCm-only", allow_module_level=True)
+
+from vllm.model_executor.kernels.linear.mixed_precision.MPLinearKernel import (  # noqa: E402
+    MPLinearLayerConfig,
+)
+from vllm.model_executor.kernels.linear.mixed_precision.rdna3_w4a16 import (  # noqa: E402
+    RDNA3W4A16LinearKernel,
+)
+from vllm.model_executor.layers.quantization.utils.quant_utils import (  # noqa: E402
+    pack_quantized_values_into_int32,
+)
+from vllm.model_executor.parameter import (  # noqa: E402
+    GroupQuantScaleParameter,
+    PackedvLLMParameter,
+)
+from vllm.platforms.rocm import on_gfx1100  # noqa: E402
+from vllm.scalar_type import scalar_types  # noqa: E402
+
+device = "cuda"
+WEIGHT_TYPE = scalar_types.uint4b8
+
+gfx1100_only = pytest.mark.skipif(
+    not (
+        on_gfx1100()
+        and hasattr(torch.ops, "_rocm_C")
+        and hasattr(torch.ops._rocm_C, "gptq_gemm_rdna3")
+    ),
+    reason="requires gfx1100 with the _rocm_C.gptq_gemm_rdna3 op built in",
+)
+
+REPEATS = 20
+GROUP = 128
+DTYPES = [torch.bfloat16, torch.float16]
+
+# max_abs_error tolerances vs the FP32 reference, per (path, dtype).
+# Derivation in the module docstring; ~2-3x headroom over measured maxima.
+TOL = {
+    ("scalar", torch.bfloat16): 0.10,
+    ("scalar", torch.float16): 0.025,
+    ("wmma", torch.bfloat16): 0.25,
+    ("wmma", torch.float16): 0.05,
+}
+
+
+def _build_layer(k, n, seed, dtype, group=GROUP):
+    torch.manual_seed(seed)
+    q_int4_kn = torch.randint(0, 16, (k, n), dtype=torch.int32)
+    scales_gn = (torch.randn(k // group, n) * 0.01 + 0.02).to(dtype)
+    qweight = pack_quantized_values_into_int32(q_int4_kn, WEIGHT_TYPE, packed_dim=0)
+    no_loader = lambda *a, **kw: None  # noqa: E731
+
+    class DummyLayer(torch.nn.Module):
+        pass
+
+    layer = DummyLayer()
+    layer.register_parameter(
+        "qweight",
+        PackedvLLMParameter(
+            data=qweight,
+            weight_loader=no_loader,
+            input_dim=0,
+            output_dim=1,
+            packed_dim=0,
+            packed_factor=8,
+        ),
+    )
+    layer.register_parameter(
+        "scales",
+        GroupQuantScaleParameter(
+            data=scales_gn, weight_loader=no_loader, input_dim=0, output_dim=1
+        ),
+    )
+    layer.to(device)
+    return layer, q_int4_kn, scales_gn
+
+
+def _prepare(layer, dtype, k, n, group=GROUP):
+    cfg = MPLinearLayerConfig(
+        full_weight_shape=(k, n),
+        partition_weight_shape=(k, n),
+        weight_type=WEIGHT_TYPE,
+        act_type=dtype,
+        group_size=group,
+        zero_points=False,
+    )
+    kernel = RDNA3W4A16LinearKernel(
+        cfg, w_q_param_name="qweight", w_s_param_name="scales", w_zp_param_name=None
+    )
+    kernel.process_weights_after_loading(layer)
+    return kernel._get_weight_params(layer)
+
+
+def _reference(m, k, dtype, seed, q_int4_kn, scales_gn, group=GROUP):
+    torch.manual_seed(seed + 1)
+    x = torch.randn(m, k, device=device, dtype=dtype)
+    # FP32 reference; see module docstring for the dequant semantics.
+    w_f32 = (q_int4_kn.to(device).float() - 8.0) * scales_gn.to(
+        device
+    ).repeat_interleave(group, dim=0).float()
+    return x, x.float() @ w_f32
+
+
+def _run_op(x, w_q, w_zp, w_s):
+    return torch.ops._rocm_C.gptq_gemm_rdna3(x, w_q, w_zp, w_s, False)
+
+
+def _outputs_and_ref(M, K, N, seed, dtype, repeats=1, group=GROUP):
+    layer, q_int4_kn, scales_gn = _build_layer(K, N, seed, dtype, group)
+    w_q, w_s, w_zp = _prepare(layer, dtype, K, N, group)
+    x, ref = _reference(M, K, dtype, seed, q_int4_kn, scales_gn, group)
+    outs = [_run_op(x, w_q, w_zp, w_s) for _ in range(repeats)]
+    return outs, ref
+
+
+def _path_for(dtype, m):
+    """Public dispatch: both dtypes reach WMMA at M>=12 (see the
+    WMMA_MIN_M crossover table in csrc/rocm/q_gemm_rdna3.cu)."""
+    if m >= 12:
+        return "wmma"
+    return "scalar"
+
+
+def _assert_repeatable(outs):
+    for o in outs[1:]:
+        assert torch.equal(o, outs[0]), (
+            "RDNA3 W4A16 split-K produced differing outputs for identical inputs"
+        )
+
+
+def _assert_close_to_ref(out, ref, path, dtype):
+    err = (out.float() - ref).abs()
+    max_abs = err.max().item()
+    mean_abs = err.mean().item()
+    tol = TOL[(path, dtype)]
+    assert torch.isfinite(out.float()).all(), "non-finite output"
+    assert max_abs <= tol, (
+        f"RDNA3 W4A16 {path} output deviates from the FP32 dequantized "
+        f"reference: max_abs={max_abs:.4f} > tol={tol} (mean_abs="
+        f"{mean_abs:.4f}). This catches values that are bitwise-repeatable "
+        "but wrong (e.g. an unwritten FP32 scratch slot)."
+    )
+
+
+# ---------------------------------------------------------------------------
+# A. Repeatability (bit-exact across identical calls)
+# ---------------------------------------------------------------------------
+
+
+@gfx1100_only
+@pytest.mark.parametrize("dtype", DTYPES)
+def test_scalar_splitk_bit_repeatable(dist_init, dtype):
+    """Scalar path at production-like contention: K=4096 splits into
+    ceil(4096/256) = 16 concurrent writers per output element (the old CAS
+    bug's onset was between 2 and 4 writers, so this sits well inside the
+    failure regime rather than at the empirical threshold)."""
+    outs, _ = _outputs_and_ref(
+        M=1, K=4096, N=4096, seed=1234, dtype=dtype, repeats=REPEATS
+    )
+    _assert_repeatable(outs)
+
+
+@gfx1100_only
+@pytest.mark.parametrize("dtype", DTYPES)
+def test_scalar_small_m_bit_repeatable(dist_init, dtype):
+    """Scalar M>1 tiles share the same deterministic epilogue."""
+    outs, _ = _outputs_and_ref(
+        M=8, K=4096, N=4096, seed=1235, dtype=dtype, repeats=REPEATS
+    )
+    _assert_repeatable(outs)
+
+
+@gfx1100_only
+@pytest.mark.parametrize("dtype", DTYPES)
+def test_scalar_single_split_bit_repeatable(dist_init, dtype):
+    """z_count == 1 (K=256 < BLOCK_KN_SIZE): direct-store epilogue, no
+    scratch, no reduce pass."""
+    outs, _ = _outputs_and_ref(
+        M=1, K=256, N=4096, seed=1238, dtype=dtype, repeats=REPEATS
+    )
+    _assert_repeatable(outs)
+
+
+@gfx1100_only
+@pytest.mark.parametrize("dtype", DTYPES)
+def test_wmma_splitk_bit_repeatable(dist_init, dtype):
+    """WMMA path (both dtypes M >= 12), split-K active.
+
+    K=6656 gives K_SPLIT=4 under the upstream heuristic; the deterministic
+    epilogue must be bit-repeatable regardless of the split count.
+    """
+    m = 12  # smallest M on the WMMA side of the dispatch boundary
+    outs, _ = _outputs_and_ref(
+        M=m, K=6656, N=4096, seed=1236, dtype=dtype, repeats=REPEATS
+    )
+    _assert_repeatable(outs)
+
+
+@gfx1100_only
+@pytest.mark.parametrize("dtype", DTYPES)
+def test_wmma_large_m_bit_repeatable(dist_init, dtype):
+    """Large-M WMMA tiles (128x64 kernels) with the deterministic epilogue."""
+    outs, _ = _outputs_and_ref(
+        M=128, K=6656, N=4096, seed=1237, dtype=dtype, repeats=REPEATS
+    )
+    _assert_repeatable(outs)
+
+
+# ---------------------------------------------------------------------------
+# B. Numerical correctness against the FP32 dequantized reference
+# ---------------------------------------------------------------------------
+
+
+@gfx1100_only
+@pytest.mark.parametrize("dtype", DTYPES)
+def test_scalar_matches_fp32_reference(dist_init, dtype):
+    """Scalar split-K (K=4096, 16 writers) vs the FP32 reference."""
+    m = 1
+    outs, ref = _outputs_and_ref(M=m, K=4096, N=512, seed=1240, dtype=dtype)
+    _assert_close_to_ref(outs[0], ref, _path_for(dtype, m), dtype)
+
+
+@gfx1100_only
+@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize("m", [2, 3, 4, 5, 8])
+def test_scalar_multi_m_matches_fp32_reference(dist_init, m, dtype):
+    """The exact factored dequant across every M_COUNT instantiation
+    (M=2,3 -> M_COUNT=2; M=4,5 -> M_COUNT=4 with a padded row; M=8 -> 8):
+    the (y, z) correction and the nibble-slot activation sums must hold
+    for every tile width, not just the M=1 fast path."""
+    outs, ref = _outputs_and_ref(M=m, K=4096, N=512, seed=1247, dtype=dtype)
+    _assert_repeatable(outs)
+    _assert_close_to_ref(outs[0], ref, _path_for(dtype, m), dtype)
+
+
+@gfx1100_only
+@pytest.mark.parametrize("dtype", DTYPES)
+def test_scalar_single_split_matches_fp32_reference(dist_init, dtype):
+    """Scalar z_count == 1 direct-store path vs the FP32 reference."""
+    m = 1
+    outs, ref = _outputs_and_ref(M=m, K=256, N=512, seed=1241, dtype=dtype)
+    _assert_close_to_ref(outs[0], ref, _path_for(dtype, m), dtype)
+
+
+@gfx1100_only
+@pytest.mark.parametrize("dtype", DTYPES)
+def test_wmma_matches_fp32_reference(dist_init, dtype):
+    """WMMA split-K path (16x16_1w at the M=12 boundary) vs FP32 reference."""
+    m = 12  # smallest M on the WMMA side of the dispatch boundary
+    outs, ref = _outputs_and_ref(M=m, K=4096, N=512, seed=1242, dtype=dtype)
+    _assert_close_to_ref(outs[0], ref, _path_for(dtype, m), dtype)
+
+
+@gfx1100_only
+@pytest.mark.parametrize("dtype", DTYPES)
+def test_wmma_single_split_matches_fp32_reference(dist_init, dtype):
+    """WMMA k_split == 1 direct-store path (K=256 -> compute_wmma_k_split
+    returns 1 for every dispatch level) vs the FP32 reference."""
+    m = 12  # smallest M on the WMMA side of the dispatch boundary
+    outs, ref = _outputs_and_ref(M=m, K=256, N=512, seed=1243, dtype=dtype)
+    _assert_close_to_ref(outs[0], ref, _path_for(dtype, m), dtype)
+
+
+@gfx1100_only
+@pytest.mark.parametrize("dtype", DTYPES)
+def test_wmma_large_m_matches_fp32_reference(dist_init, dtype):
+    """Large-M WMMA (128x64_k32, row-tiled scratch + reduce) vs FP32."""
+    m = 128
+    outs, ref = _outputs_and_ref(M=m, K=4096, N=512, seed=1244, dtype=dtype)
+    _assert_close_to_ref(outs[0], ref, _path_for(dtype, m), dtype)
+
+
+# ---------------------------------------------------------------------------
+# C. V7/V8 (128x64) no-split path: k_split == 1 -> direct store, no scratch,
+#    no reduce pass (partials == nullptr).
+# ---------------------------------------------------------------------------
+
+# The measured production-like shape from the W7900 benchmark campaign.
+# Routing arithmetic (see _assert_v7v8_no_split_routing): with (128, 64)
+# tiles, blocks_xy = ceil(25600/64) * ceil(512/128) = 400 * 4 = 1600 >=
+# 1500, so compute_wmma_k_split_mn returns 1. Smaller no-split shapes are
+# not practical: N would have to grow past ~24k at M=512 (or M past ~1k at
+# N=12k) to reach the 1500-block threshold, so this is effectively the
+# smallest production-realistic V7/V8 no-split shape.
+V7V8_NOSPLIT = (512, 25600, 6656)  # (M, N, K)
+
+
+def _compute_wmma_k_split_mn(m, n, k, m_tile, n_tile):
+    """Faithful replica of compute_wmma_k_split_mn (C++ heuristic)."""
+    blocks_xy = ((n + n_tile - 1) // n_tile) * ((m + m_tile - 1) // m_tile)
+    if blocks_xy >= 1500:
+        return 1
+    if blocks_xy * 2 >= 1500 and k >= 512 and k % 32 == 0:
+        return 2
+    if blocks_xy * 4 >= 1500 and k >= 1024 and k % 64 == 0:
+        return 4
+    if k >= 1024 and k % 64 == 0:
+        return 4
+    if k >= 512 and k % 32 == 0:
+        return 2
+    return 1
+
+
+def _assert_v7v8_no_split_routing(dtype, m, n, k):
+    """Assert the shape deterministically routes to the V7/V8 128x64
+    kernel with k_split == 1: WMMA dispatch (both dtypes M >= 12), the
+    M >= 128 branch, and the no-split threshold."""
+    assert m >= 12, "not the WMMA path"
+    assert m >= 128, "below the V7/V8 (128x64) branch"
+    assert _compute_wmma_k_split_mn(m, n, k, 128, 64) == 1, (
+        "shape does not route to the k_split == 1 no-split path"
+    )
+
+
+@gfx1100_only
+@pytest.mark.parametrize("dtype", DTYPES)
+def test_wmma_v7v8_no_split_bit_repeatable(dist_init, dtype):
+    """V7/V8 128x64_k32 with k_split == 1: single direct writer per output
+    cell — no FP32 scratch, no reduce pass."""
+    m, n, k = V7V8_NOSPLIT
+    _assert_v7v8_no_split_routing(dtype, m, n, k)
+    outs, _ = _outputs_and_ref(M=m, K=k, N=n, seed=1245, dtype=dtype, repeats=REPEATS)
+    _assert_repeatable(outs)
+
+
+@gfx1100_only
+@pytest.mark.parametrize("dtype", DTYPES)
+def test_wmma_v7v8_no_split_matches_fp32_reference(dist_init, dtype):
+    """V7/V8 128x64_k32 no-split direct-store path vs the FP32 reference."""
+    m, n, k = V7V8_NOSPLIT
+    _assert_v7v8_no_split_routing(dtype, m, n, k)
+    outs, ref = _outputs_and_ref(M=m, K=k, N=n, seed=1246, dtype=dtype)
+    _assert_close_to_ref(outs[0], ref, _path_for(dtype, m), dtype)
+
+
+# ---------------------------------------------------------------------------
+# D. Dispatch boundary: M=11 must stay scalar, M=12 must route to WMMA
+#    (WMMA_MIN_M in csrc/rocm/q_gemm_rdna3.cu). Observed via the CUDA
+#    profiler: the scalar and WMMA kernels have disjoint names.
+# ---------------------------------------------------------------------------
+
+
+def _launched_kernel_names(fn) -> set:
+    from torch.profiler import ProfilerActivity, profile
+
+    fn()
+    torch.accelerator.synchronize()
+    with profile(activities=[ProfilerActivity.CUDA]) as prof:
+        fn()
+        torch.accelerator.synchronize()
+    return {
+        e.name
+        for e in prof.events()
+        if "gemm_q4" in e.name or "reduce_partials" in e.name
+    }
+
+
+@pytest.mark.parametrize("dtype", DTYPES)
+@gfx1100_only
+def test_dispatch_boundary(dist_init, dtype):
+    """The scalar/WMMA crossover sits between M=11 and M=12 for both
+    dtypes; the boundary must be exact (see the WMMA_MIN_M table)."""
+
+    def run(m):
+        outs, _ = _outputs_and_ref(M=m, K=4096, N=512, seed=1248, dtype=dtype)
+        return outs[0]
+
+    m11 = run(11)
+    names11 = _launched_kernel_names(lambda: run(11))
+    assert any("gemm_q4_kernel_rdna3" in n for n in names11), names11
+    assert not any("wmma" in n for n in names11), names11
+
+    m12 = run(12)
+    names12 = _launched_kernel_names(lambda: run(12))
+    assert any("wmma" in n for n in names12), names12
+    assert not any("gemm_q4_kernel_rdna3" in n for n in names12), names12
+
+    # Both sides of the boundary stay correct against the reference.
+    _, ref11 = _outputs_and_ref(M=11, K=4096, N=512, seed=1248, dtype=dtype)
+    _, ref12 = _outputs_and_ref(M=12, K=4096, N=512, seed=1248, dtype=dtype)
+    _assert_close_to_ref(m11, ref11, "scalar", dtype)
+    _assert_close_to_ref(m12, ref12, "wmma", dtype)
+
+
+# ---------------------------------------------------------------------------
+# E. In-kernel split-K reduction: stress the per-call ticket epilogue.
+#    The fused reduction must stay deterministic and correct under
+#    back-to-back unsynchronized calls, concurrent invocations on two
+#    streams (per-call tickets must never share state), and CUDA-graph
+#    capture/replay.
+# ---------------------------------------------------------------------------
+
+
+@gfx1100_only
+@pytest.mark.parametrize("dtype", DTYPES)
+def test_scalar_splitk_100x_bit_repeatable(dist_init, dtype):
+    """100 back-to-back calls (no intermediate sync) on the split-K path."""
+    outs, _ = _outputs_and_ref(M=1, K=4096, N=512, seed=1249, dtype=dtype, repeats=100)
+    _assert_repeatable(outs)
+
+
+@gfx1100_only
+@pytest.mark.parametrize("dtype", DTYPES)
+def test_scalar_multistream_concurrent(dist_init, dtype):
+    """Two concurrent invocations on separate streams, distinct weights and
+    activations: outputs must be bit-identical to the serial invocation of
+    each (the arrival tickets are per-call, so cross-stream interference
+    would show up here as a corrupted or unreduced output)."""
+    layer_a, q_a, s_a = _build_layer(4096, 512, 4321, dtype)
+    layer_b, q_b, s_b = _build_layer(4096, 512, 4322, dtype)
+    wq_a, ws_a, wz_a = _prepare(layer_a, dtype, 4096, 512)
+    wq_b, ws_b, wz_b = _prepare(layer_b, dtype, 4096, 512)
+    torch.manual_seed(4323)
+    x_a = torch.randn(1, 4096, device=device, dtype=dtype)
+    x_b = torch.randn(1, 4096, device=device, dtype=dtype)
+
+    serial_a = _run_op(x_a, wq_a, wz_a, ws_a)
+    serial_b = _run_op(x_b, wq_b, wz_b, ws_b)
+    torch.accelerator.synchronize()
+
+    s1, s2 = torch.cuda.Stream(), torch.cuda.Stream()
+    for _ in range(25):
+        with torch.cuda.stream(s1):
+            out_a = _run_op(x_a, wq_a, wz_a, ws_a)
+        with torch.cuda.stream(s2):
+            out_b = _run_op(x_b, wq_b, wz_b, ws_b)
+    torch.accelerator.synchronize()
+    assert torch.equal(out_a, serial_a), "concurrent stream A corrupted"
+    assert torch.equal(out_b, serial_b), "concurrent stream B corrupted"
+
+
+@gfx1100_only
+@pytest.mark.parametrize("dtype", DTYPES)
+def test_scalar_graph_capture_replay(dist_init, dtype):
+    """The op must survive CUDA-graph capture: the replayed output must be
+    bit-identical to the eager output and stable across repeated replays."""
+    layer, q_int4, scales_gn = _build_layer(4096, 512, 4324, dtype)
+    wq, ws, wz = _prepare(layer, dtype, 4096, 512)
+    torch.manual_seed(4325)
+    x = torch.randn(1, 4096, device=device, dtype=dtype)
+
+    eager = _run_op(x, wq, wz, ws)
+    torch.accelerator.synchronize()
+
+    graph = torch.cuda.CUDAGraph()
+    _run_op(x, wq, wz, ws)  # warm-up side streams/allocator on capture path
+    torch.accelerator.synchronize()
+    with torch.cuda.graph(graph):
+        captured = _run_op(x, wq, wz, ws)
+    graph.replay()
+    torch.accelerator.synchronize()
+    assert torch.equal(captured, eager), "graph replay differs from eager"
+    graph.replay()
+    graph.replay()
+    torch.accelerator.synchronize()
+    assert torch.equal(captured, eager), "repeated replay not repeatable"
+
+
+@gfx1100_only
+@pytest.mark.parametrize("dtype", DTYPES)
+def test_scalar_multi_row_tile_bit_repeatable(dist_init, dtype):
+    """Row tiling with ticket reuse: N % 16 == 8 bypasses the WMMA
+    dispatch (N % 16 != 0), so M=100 runs scalar across two 64-row tiles
+    (the second with gy > 1 y-blocks). Every tile's tickets must be armed,
+    consumed, and re-armed correctly for all outputs to be written exactly
+    once — and the result must be bit-repeatable."""
+    outs, ref = _outputs_and_ref(
+        M=100, K=4096, N=2056, seed=1250, dtype=dtype, repeats=20
+    )
+    _assert_repeatable(outs)
+    _assert_close_to_ref(outs[0], ref, "scalar", dtype)
+
+
+# ---------------------------------------------------------------------------
+# F. Split-K geometry corners: clipped non-first z blocks, small split
+#    counts, and non-128 group sizes.
+# ---------------------------------------------------------------------------
+
+
+@gfx1100_only
+@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize("k,group", [(4224, 128), (2112, 32)])
+def test_scalar_clipped_tail_splitk(dist_init, dtype, k, group):
+    """K % 256 != 0 with z_count > 1: the LAST z block is clipped by end_k,
+    its group prefetch guard (`k + 32 < end_k`) is what keeps the raw
+    zero/scale loads in bounds, and the reduce tail loop runs (z=17 -> one
+    remainder slice; z=9 -> one batch + one). K=2112/G=32 also puts two
+    group switches inside the clipped tail block."""
+    outs, ref = _outputs_and_ref(M=1, K=k, N=512, seed=1251, dtype=dtype, group=group)
+    _assert_repeatable(outs)
+    _assert_close_to_ref(outs[0], ref, "scalar", dtype)
+
+
+@gfx1100_only
+@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize("group", [32, 64])
+def test_scalar_small_group_matches_fp32_reference(dist_init, dtype, group):
+    """Group sizes below 128 make the group constants change every 32-K
+    round (G=32) or three times per z block (G=64); the factored-dequant
+    correction must track every change. G=32 is also the production MoE
+    group size."""
+    outs, ref = _outputs_and_ref(
+        M=1, K=2048, N=512, seed=1252, dtype=dtype, group=group
+    )
+    _assert_repeatable(outs)
+    _assert_close_to_ref(outs[0], ref, "scalar", dtype)
+
+
+@gfx1100_only
+@pytest.mark.parametrize("dtype", DTYPES)
+def test_scalar_two_way_split_bit_repeatable(dist_init, dtype):
+    """z_count == 2 (K=512): the minimal split-K configuration — the old
+    CAS bug's onset regime — and the pure-remainder reduce tail."""
+    outs, _ = _outputs_and_ref(
+        M=1, K=512, N=512, seed=1253, dtype=dtype, repeats=REPEATS
+    )
+    _assert_repeatable(outs)

tests/kernels/test_compressor_kv_cache.py

@@ -711,7 +711,8 @@ def test_v41_rope_insert_plain_row(compress_ratio: int, store_fp8: bool):
 def test_v41_compressor_metadata_maps_tokens_to_their_ring():
     """The ring group's generic slot mapping is disabled (all PAD), so the
     builder must map every real token to ``ring_block * capacity + pos %
-    capacity`` and keep padding tokens at PAD."""
+    capacity``. Padding tokens and requests on the null block (block 0, as in
+    dummy batches) must stay at PAD."""
     from unittest.mock import MagicMock
 
     from vllm.models.deepseek_v41.compressor import CompressorMetadataBuilder
@@ -731,29 +732,30 @@ def test_v41_compressor_metadata_maps_tokens_to_their_ring():
     device = torch.device("cuda")
     builder = CompressorMetadataBuilder(spec, ["state"], vllm_config, device)
 
-    # Two requests: 3 tokens at positions 13..15 on ring block 5, then 2
+    # Three requests: 3 tokens at positions 13..15 on ring block 5, 2 tokens
-    # tokens at positions 7..8 on ring block 2; three padding tokens.
+    # at positions 7..8 on ring block 2, then 2 tokens on the null block;
-    query_start_loc = torch.tensor([0, 3, 5], dtype=torch.int32, device=device)
+    # three padding tokens.
-    positions = torch.tensor([13, 14, 15, 7, 8, 0, 0, 0], device=device)
+    query_start_loc = torch.tensor([0, 3, 5, 7], dtype=torch.int32, device=device)
-    block_table = torch.tensor([[5], [2]], dtype=torch.int32, device=device)
+    positions = torch.tensor([13, 14, 15, 7, 8, 0, 1, 0, 0, 0], device=device)
+    block_table = torch.tensor([[5], [2], [0]], dtype=torch.int32, device=device)
     common = CommonAttentionMetadata(
         query_start_loc=query_start_loc,
         query_start_loc_cpu=query_start_loc.cpu(),
-        seq_lens=torch.tensor([16, 9], dtype=torch.int32, device=device),
+        seq_lens=torch.tensor([16, 9, 2], dtype=torch.int32, device=device),
-        num_reqs=2,
+        num_reqs=3,
-        num_actual_tokens=5,
+        num_actual_tokens=7,
         max_query_len=3,
         max_seq_len=16,
         block_table_tensor=block_table,
-        slot_mapping=torch.full((8,), -1, dtype=torch.int64, device=device),
+        slot_mapping=torch.full((10,), -1, dtype=torch.int64, device=device),
         positions=positions,
     )
     metadata = builder.build(0, common)
 
-    expected = [5 * 8 + 5, 5 * 8 + 6, 5 * 8 + 7, 2 * 8 + 7, 2 * 8 + 0, -1, -1, -1]
+    ring = [5 * 8 + 5, 5 * 8 + 6, 5 * 8 + 7, 2 * 8 + 7, 2 * 8 + 0]
-    assert metadata.slot_mapping.tolist() == expected
+    assert metadata.slot_mapping.tolist() == ring + [-1] * 5
     assert metadata.query_start_loc is query_start_loc
-    assert metadata.token_to_req_indices.tolist() == [0, 0, 0, 1, 1]
+    assert metadata.token_to_req_indices.tolist() == [0, 0, 0, 1, 1, 2, 2]
 
 
 @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA stream coverage")

tests/models/qwen4_exp/test_ple.py

@@ -26,6 +26,7 @@ from vllm.config import SpeculativeConfig, VllmConfig
 from vllm.config.compilation import CUDAGraphMode
 from vllm.config.quantization import QuantizationConfigArgs
 from vllm.model_executor.layers.quantization.fp8 import Fp8Config
+from vllm.model_executor.layers.quantization.inc import INCConfig
 from vllm.model_executor.layers.quantization.modelopt import (
     ModelOptMixedPrecisionConfig,
     ModelOptNvFp4Config,
@@ -573,6 +574,24 @@ def test_ple_embedding_respects_modelopt_exclusion() -> None:
     )
 
 
+def test_ple_embedding_respects_inc_layer_config() -> None:
+    prefix = "model.layers.1.ple.ple_embedding.ngram_embedding"
+    quant_config = INCConfig(
+        weight_bits=4,
+        group_size=128,
+        block_name_to_quantize="model.layers",
+        extra_config={".*ple.*": {"bits": 16, "data_type": "float"}},
+    )
+    assert isinstance(
+        Qwen4ExpPLEEmbeddingMethod.from_quant_config(quant_config, prefix),
+        Qwen4ExpPLEUnquantizedEmbeddingMethod,
+    )
+
+    quant_config.extra_config = None
+    with pytest.raises(NotImplementedError, match="INCConfig"):
+        Qwen4ExpPLEEmbeddingMethod.from_quant_config(quant_config, prefix)
+
+
 def test_ple_embedding_dtype_overrides_modelopt_exclusion() -> None:
     prefix = "model.layers.1.ple.ple_embedding.ngram_embedding"
     quant_config = ModelOptNvFp4Config(exclude_modules=[prefix])

tests/models/registry.py

@@ -380,12 +380,9 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
     "Llama4ForCausalLM": _HfExamplesInfo(
         "meta-llama/Llama-4-Scout-17B-16E-Instruct",
     ),
-    "LongcatFlashForCausalLM": _HfExamplesInfo(
+    "LongcatFlashForCausalLM": _HfExamplesInfo("meituan-longcat/LongCat-Flash-Chat"),
-        "meituan-longcat/LongCat-Flash-Chat", trust_remote_code=True
-    ),
     "LongcatFlashNgramForCausalLM": _HfExamplesInfo(
         "meituan-longcat/LongCat-Flash-Lite",
-        trust_remote_code=True,
         # Shrink the ~62GB n-gram tables (ngram_vocab_size_ratio * vocab_size)
         # so the dummy-weight init test fits in CI memory.
         hf_overrides={"ngram_vocab_size_ratio": 1},
@@ -1752,7 +1749,6 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {
     ),
     "LongCatFlashMTPModel": _HfExamplesInfo(
         "meituan-longcat/LongCat-Flash-Chat",
-        trust_remote_code=True,
         speculative_model="meituan-longcat/LongCat-Flash-Chat",
     ),
     "MiMoMTPModel": _HfExamplesInfo(

tests/multimodal/test_audio.py

@@ -216,17 +216,19 @@ class TestNormalizeAudio:
         assert result.ndim == 1
         torch.testing.assert_close(result, torch.tensor([0.0, 1.0]))
 
-    def test_mono_passthrough_for_1d_numpy(self):
+    @pytest.mark.parametrize("shape", [(3,), (1, 3), (3, 1), (1, 1)])
-        """1D numpy array should pass through unchanged with mono spec."""
+    def test_mono_passthrough_for_numpy(self, shape):
-        mono = np.array([1.0, 2.0, 3.0], dtype=np.float32)
+        """Mono output is 1D, including audio with an explicit channel axis."""
-        result = normalize_audio(mono, MONO_AUDIO_SPEC)
+        mono = np.arange(max(shape), dtype=np.float32)
+        result = normalize_audio(mono.reshape(shape), MONO_AUDIO_SPEC)
         assert result.ndim == 1
         np.testing.assert_array_equal(result, mono)
 
-    def test_mono_passthrough_for_1d_torch(self):
+    @pytest.mark.parametrize("shape", [(3,), (1, 3), (3, 1), (1, 1)])
-        """1D torch tensor should pass through unchanged with mono spec."""
+    def test_mono_passthrough_for_torch(self, shape):
-        mono = torch.tensor([1.0, 2.0, 3.0])
+        """Mono output is 1D, including audio with an explicit channel axis."""
-        result = normalize_audio(mono, MONO_AUDIO_SPEC)
+        mono = torch.arange(max(shape), dtype=torch.float32)
+        result = normalize_audio(mono.reshape(shape), MONO_AUDIO_SPEC)
         assert result.ndim == 1
         torch.testing.assert_close(result, mono)
 
@@ -410,8 +412,9 @@ class TestMultiModalDataParserChannelNormalization:
         # When target_channels=None, stereo audio should be preserved
         assert audio_item.ndim == 2, f"Expected 2D stereo audio, got {audio_item.ndim}D"
 
-    def test_parser_mono_passthrough_when_target_channels_1(self):
+    @pytest.mark.parametrize("shape", [(16000,), (1, 16000), (16000, 1)])
-        """Parser should pass through mono audio unchanged when target_channels=1."""
+    def test_parser_mono_passthrough_when_target_channels_1(self, shape):
+        """Mono layouts must preserve samples and report the waveform length."""
         from vllm.multimodal.parse import MultiModalDataParser
 
         # Create parser with mono normalization enabled
@@ -420,17 +423,20 @@ class TestMultiModalDataParserChannelNormalization:
             target_channels=1,
         )
 
-        # Create mono audio (already 1D)
+        mono_audio = np.linspace(-1.0, 1.0, 16000, dtype=np.float32)
-        mono_audio = np.random.randn(16000).astype(np.float32)
 
         # Parse audio data
-        result = parser._parse_audio_data((mono_audio, 16000))
+        result = parser.parse_mm_data({"audio": (mono_audio.reshape(shape), 16000)})[
+            "audio"
+        ]
 
         # Check that result is still mono (1D)
         assert result is not None
         audio_item = result.get(0)
         assert audio_item.ndim == 1
         assert audio_item.shape == (16000,)
+        assert len(audio_item) == 16000
+        np.testing.assert_array_equal(audio_item, mono_audio)
 
     def test_parser_with_target_channels_2(self):
         """Parser should reduce 6-channel to 2-channel when target_channels=2."""

tests/v1/attention/test_mla_backends.py

@@ -1098,8 +1098,14 @@ def test_flashinfer_mla_dspark_dcp_supports_target_and_draft(monkeypatch):
         pytest.param(False, 3, 1, 0, id="noncausal-multi-token"),
     ],
 )
+@pytest.mark.parametrize("cp_interleave_size", [1, 16, 896])
 def test_tokenspeed_mla_decode_contract(
-    monkeypatch, causal, tokens_per_decode, dcp_world_size, dcp_rank
+    monkeypatch,
+    causal,
+    tokens_per_decode,
+    dcp_world_size,
+    dcp_rank,
+    cp_interleave_size,
 ):
     decode_call = None
     num_decodes = 2
@@ -1135,7 +1141,7 @@ def test_tokenspeed_mla_decode_contract(
     impl = object.__new__(tokenspeed_mla_module.TokenspeedMLAImpl)
     impl.dcp_world_size = dcp_world_size
     impl.dcp_rank = dcp_rank
-    impl.cp_kv_cache_interleave_size = 1
+    impl._parallel_config = SimpleNamespace(cp_kv_cache_interleave_size=1)
     impl.need_to_return_lse_for_decode = True
     impl.kv_lora_rank = kv_lora_rank
     impl.qk_rope_head_dim = qk_rope_head_dim
@@ -1164,6 +1170,8 @@ def test_tokenspeed_mla_decode_contract(
         dtype=torch.float8_e4m3fn,
     )
 
+    # NIXL may resolve the interleave size after the implementation is created.
+    impl._parallel_config.cp_kv_cache_interleave_size = cp_interleave_size
     out, lse = impl.forward_mqa(
         q,
         kv_cache,
@@ -1194,6 +1202,142 @@ def test_tokenspeed_mla_decode_contract(
     assert decode_call["return_lse"] is True
     assert decode_call["cp_world"] == dcp_world_size
     assert decode_call["cp_rank"] == dcp_rank
+    assert decode_call.get("cp_interleave_size", 1) == cp_interleave_size
+
+
+@pytest.mark.skipif(
+    AttentionBackendEnum.TOKENSPEED_MLA not in BACKENDS_TO_TEST,
+    reason="TokenSpeed MLA requires an SM100-family GPU and tokenspeed-mla",
+)
+@pytest.mark.parametrize(
+    ("cp_interleave_size", "seq_len"),
+    [
+        pytest.param(1, 5, id="token-interleave-empty-local-ranks"),
+        pytest.param(1, 8 + 17, id="token-interleave-full-cycle"),
+        pytest.param(896, 5, id="block-interleave-empty-local-ranks"),
+        pytest.param(896, 8 * 896 + 17, id="block-interleave-full-cycle"),
+    ],
+)
+def test_tokenspeed_mla_dcp_matches_unsharded_decode(cp_interleave_size, seq_len):
+    from tokenspeed_mla import tokenspeed_mla_decode
+
+    from vllm.v1.attention.ops.dcp import (
+        _lse_weighted_combine,
+        mask_dcp_empty_shards_,
+    )
+
+    torch.manual_seed(7)
+    device = torch.device("cuda")
+    dcp_world_size = 8
+    kernel_block_size = 64
+    num_heads = 128
+    kv_lora_rank = 512
+    qk_rope_head_dim = 64
+    head_size = kv_lora_rank + qk_rope_head_dim
+
+    def make_paged_cache(tokens: torch.Tensor):
+        # Empty DCP ranks still have the null block in vLLM's physical cache.
+        num_pages = max(1, cdiv(tokens.shape[0], kernel_block_size))
+        cache = torch.zeros(
+            num_pages,
+            kernel_block_size,
+            head_size,
+            dtype=tokens.dtype,
+            device=device,
+        )
+        cache.view(-1, head_size)[: tokens.shape[0]].copy_(tokens)
+        block_table = torch.arange(
+            num_pages, dtype=torch.int32, device=device
+        ).unsqueeze(0)
+        return cache, block_table
+
+    # Four query positions exercise the MTP/DSpark causal-mask path.
+    query = (torch.randn(1, 4, num_heads, head_size, device=device) * 0.1).to(
+        torch.float8_e4m3fn
+    )
+    global_tokens = (torch.randn(seq_len, head_size, device=device) * 0.1).to(
+        torch.float8_e4m3fn
+    )
+    workspace = tokenspeed_mla_module._get_workspace(device, num_heads, kv_lora_rank)
+    global_cache, global_block_table = make_paged_cache(global_tokens)
+    global_seq_len = torch.tensor([seq_len], dtype=torch.int32, device=device)
+    reference, _ = tokenspeed_mla_decode(
+        query=query,
+        kv_cache=global_cache,
+        workspace_buffer=workspace,
+        kv_lora_rank=kv_lora_rank,
+        qk_rope_head_dim=qk_rope_head_dim,
+        block_tables=global_block_table,
+        seq_lens=global_seq_len,
+        max_seq_len=seq_len,
+        softmax_scale=head_size**-0.5,
+        return_lse=True,
+    )
+
+    positions = torch.arange(seq_len, device=device)
+    owners = positions.div(cp_interleave_size, rounding_mode="floor").remainder(
+        dcp_world_size
+    )
+    max_local_seq_len = (
+        cdiv(seq_len, dcp_world_size * cp_interleave_size) * cp_interleave_size
+    )
+    partial_outputs = []
+    partial_lses = []
+    for dcp_rank in range(dcp_world_size):
+        local_tokens = global_tokens[owners == dcp_rank]
+        local_cache, local_block_table = make_paged_cache(local_tokens)
+        local_seq_len = torch.tensor(
+            [local_tokens.shape[0]], dtype=torch.int32, device=device
+        )
+        impl = object.__new__(tokenspeed_mla_module.TokenspeedMLAImpl)
+        impl.dcp_world_size = dcp_world_size
+        impl.dcp_rank = dcp_rank
+        impl._parallel_config = SimpleNamespace(
+            cp_kv_cache_interleave_size=cp_interleave_size
+        )
+        impl.need_to_return_lse_for_decode = True
+        impl.kv_lora_rank = kv_lora_rank
+        impl.qk_rope_head_dim = qk_rope_head_dim
+        impl.num_heads = num_heads
+        impl.scale = head_size**-0.5
+        impl.softmax_scale = None
+        impl.output_scale = None
+        impl._workspace_buffer = workspace
+        metadata = SimpleNamespace(
+            num_decodes=1,
+            num_decode_tokens=query.shape[1],
+            max_seq_len=max_local_seq_len,
+            causal=True,
+            decode=SimpleNamespace(
+                block_table=local_block_table,
+                seq_lens=local_seq_len,
+                dcp_tot_seq_lens=global_seq_len,
+            ),
+        )
+        output, lse = impl.forward_mqa(
+            query.view(-1, num_heads, head_size),
+            local_cache,
+            metadata,
+            SimpleNamespace(_q_scale_float=1.0, _k_scale_float=1.0),
+        )
+        # Empty-shard normalization belongs to the downstream DCP combine path.
+        mask_dcp_empty_shards_(
+            lse,
+            local_seq_len,
+            torch.tensor([0, query.shape[1]], dtype=torch.int32, device=device),
+        )
+        partial_outputs.append(
+            output.view(query.shape[1], num_heads, kv_lora_rank).float()
+        )
+        partial_lses.append(lse.view(query.shape[1], num_heads).float())
+
+    # TokenSpeed returns base-2 LSE. Merge the per-rank partial attention states
+    # the same way MLADCPManager does before comparing to the unsharded result.
+    merged = _lse_weighted_combine(
+        torch.stack(partial_outputs), torch.stack(partial_lses), is_lse_base_on_e=False
+    ).unsqueeze(0)
+    # FP8 attention rounds differently when KV is partitioned across DCP ranks.
+    torch.testing.assert_close(merged, reference.float(), atol=1e-2, rtol=1e-2)
 
 
 @pytest.mark.parametrize("is_fp8_kvcache", [False, True], ids=["bf16", "fp8"])

tests/v1/cudagraph/test_cudagraph_manager.py

@@ -50,6 +50,7 @@ def _create_vllm_config() -> MagicMock:
     vllm_config.parallel_config = ParallelConfig()
     vllm_config.speculative_config = None
     vllm_config.num_speculative_tokens = 0
+    vllm_config.use_cumem_cudagraph_pool = False
     return vllm_config
 
 

tests/v1/determinism/test_batch_invariance.py

@@ -18,6 +18,7 @@ from utils import (
 
 import vllm.envs as envs
 from vllm import LLM, SamplingParams
+from vllm.inputs import TokensPrompt
 
 
 @skip_unsupported
@@ -778,32 +779,16 @@ def test_decode_logprobs_match_prefill_logprobs(
         print(f"[Prompt {prompt_idx}] Generated {len(token_ids)} tokens: {token_ids}")
         print(f"[Prompt {prompt_idx}] Decode logprobs: {decode_logprobs.tolist()}")
 
-        # Step 2: For each token position, run prefill and compare
+        # Step 2: For each token position, run prefill and compare.
+        prompt_token_ids = list(decode_output.prompt_token_ids or [])
+
         print(f"\n[Prompt {prompt_idx}] Verifying each token via prefill...")
 
         for token_idx in range(len(token_ids)):
-            # Construct the prefix up to (but not including) this token
+            # Construct the prefix up to (but not including) this token using
+            # token ids directly.
             current_token = token_ids[token_idx]
-
+            prefix_token_ids = prompt_token_ids + list(token_ids[:token_idx])
-            # We need to detokenize to get the text prefix
-            # For this, we'll use the tokenizer from the LLM
-            # However, the LLM API doesn't expose tokenizer easily, so we'll
-            # construct the prefix by decoding from the original prompt
-
-            # Get text up to this point by using the output text
-            # This is approximate but should work for verification
-            if token_idx == 0:
-                prefix_prompt = prompt
-            else:
-                # No per-token text, so regenerate with max_tokens = token_idx
-                # to get the prefix text
-                prefix_sp = SamplingParams(
-                    temperature=0.0,
-                    max_tokens=token_idx,
-                    logprobs=1,
-                )
-                prefix_output = llm.generate([prompt], prefix_sp, use_tqdm=False)[0]
-                prefix_prompt = prompt + prefix_output.outputs[0].text
 
             # Now run prefill with max_tokens=1 to get the logprob of the next token
             prefill_sp = SamplingParams(
@@ -814,19 +799,23 @@ def test_decode_logprobs_match_prefill_logprobs(
 
             print(
                 f"  [Token {token_idx}] Running prefill for prefix "
-                f"(len={len(prefix_prompt)})..."
+                f"(num_tokens={len(prefix_token_ids)})..."
+            )
+            prefill_output = llm.generate(
+                [TokensPrompt(prompt_token_ids=prefix_token_ids)],
+                prefill_sp,
+                use_tqdm=False,
+            )[0]
+            prefill_logprobs, prefill_token_ids_out = _extract_step_logprobs(
+                prefill_output
             )
-            prefill_output = llm.generate([prefix_prompt], prefill_sp, use_tqdm=False)[
-                0
-            ]
-            prefill_logprobs, prefill_token_ids = _extract_step_logprobs(prefill_output)
 
             if prefill_logprobs is None:
                 print(f"  [Token {token_idx}] Warning: No prefill logprobs available")
                 continue
 
             # The first token from prefill should match the current token
-            prefill_token = prefill_token_ids[0]
+            prefill_token = prefill_token_ids_out[0]
             prefill_logprob = prefill_logprobs[0].item()
             decode_logprob = decode_logprobs[token_idx].item()
 
@@ -851,7 +840,7 @@ def test_decode_logprobs_match_prefill_logprobs(
                         "decode_logprob": decode_logprob,
                         "prefill_logprob": prefill_logprob,
                         "prompt_text": prompt[:100],
-                        "prefix_text": prefix_prompt[:100],
+                        "prefix_token_ids": prefix_token_ids[:50],
                     }
                 )
                 print(f"  [Token {token_idx}] ✗ TOKEN MISMATCH!")
@@ -871,7 +860,7 @@ def test_decode_logprobs_match_prefill_logprobs(
                         "prefill_logprob": prefill_logprob,
                         "diff": diff,
                         "prompt_text": prompt[:100],
-                        "prefix_text": prefix_prompt[:100],
+                        "prefix_token_ids": prefix_token_ids[:50],
                         "decode_all_tokens": token_ids,
                         "decode_all_logprobs": decode_logprobs.tolist(),
                     }
@@ -908,7 +897,7 @@ def test_decode_logprobs_match_prefill_logprobs(
             for i, fail in enumerate(failures[:5]):  # Show first 5 failures per prompt
                 print(f"\n  [Failure {i + 1}] Token position {fail['token_idx']}:")
                 print(f"    Reason: {fail['reason']}")
-                print(f"    Prefix text: '{fail['prefix_text']}...'")
+                print(f"    Prefix token ids: {fail['prefix_token_ids']}...")
                 print(
                     f"    Decode:  token={fail['decode_token']}, "
                     f"logprob={fail['decode_logprob']:.10f}"

tests/v1/determinism/test_lora_shrink_batch_invariant.py

@@ -0,0 +1,99 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Batch-invariant LoRA shrink: fixed split-K partials, fixed-order reduction."""
+
+import pytest
+import torch
+from utils import skip_if_not_cuda
+
+import vllm.lora.ops.triton_ops.lora_shrink_op as shrink_op
+import vllm.lora.ops.triton_ops.utils as lora_utils
+from vllm.lora.ops.triton_ops import LoRAKernelMeta, lora_shrink
+from vllm.lora.ops.triton_ops.utils import _LORA_A_PTR_DICT
+from vllm.utils.torch_utils import set_random_seed
+
+DEVICE = "cuda"
+
+
+@pytest.fixture(autouse=True)
+def batch_invariant(monkeypatch):
+    monkeypatch.setattr(lora_utils, "is_batch_invariant", True)
+    monkeypatch.setattr(shrink_op, "is_batch_invariant", True)
+    lora_utils.get_lora_op_configs.cache_clear()
+    yield
+    lora_utils.get_lora_op_configs.cache_clear()
+
+
+def _shrink(inputs, weights, mapping, scaling):
+    num_tokens = inputs.size(0)
+    meta = LoRAKernelMeta.make(weights[0].size(0), num_tokens, device=DEVICE)
+    meta.prepare_tensors(mapping)
+    out = torch.zeros(
+        (len(weights), num_tokens, weights[0].size(1)),
+        dtype=torch.float32,
+        device=DEVICE,
+    )
+    _LORA_A_PTR_DICT.clear()
+    args = meta.meta_args(token_nums=num_tokens, specialize_active_lora=False)
+    lora_shrink(inputs, weights, out, *args, scaling)
+    return out
+
+
+@skip_if_not_cuda
+@pytest.mark.parametrize(
+    "hidden_size,rank,nslices,dtype",
+    [
+        (2560, 8, 3, torch.bfloat16),
+        (9728, 16, 1, torch.float16),
+        (4097, 64, 1, torch.bfloat16),  # K tail block
+        (1024, 8, 1, torch.float16),  # K < 8 * BLOCK_K: some splits are empty
+    ],
+)
+def test_lora_shrink_matches_reference(hidden_size, rank, nslices, dtype):
+    set_random_seed(0)
+    inputs = torch.randn((33, hidden_size), dtype=dtype, device=DEVICE)
+    weights = [
+        torch.randn((2, rank, hidden_size), dtype=dtype, device=DEVICE)
+        for _ in range(nslices)
+    ]
+    mapping = torch.randint(-1, 2, (33,), dtype=torch.int32, device=DEVICE)
+    out = _shrink(inputs, weights, mapping, 0.5).double().cpu()
+
+    ids, x = mapping.cpu(), inputs.double().cpu()
+    ref = torch.zeros_like(out)
+    for s, weight in enumerate(weights):
+        for lora_id in (0, 1):
+            rows = ids == lora_id
+            ref[s, rows] = 0.5 * x[rows] @ weight[lora_id].double().cpu().T
+    torch.testing.assert_close(out, ref, rtol=5e-3, atol=5e-3)
+
+
+@skip_if_not_cuda
+@pytest.mark.parametrize("nslices", [1, 3])
+def test_lora_shrink_invariant_across_batch_sizes(monkeypatch, nslices):
+    reduce_kernel, reduce_grids = shrink_op._lora_shrink_reduce_kernel, []
+
+    class RecordReduce:
+        def __getitem__(self, grid):
+            reduce_grids.append(grid)
+            return reduce_kernel[grid]
+
+    monkeypatch.setattr(shrink_op, "_lora_shrink_reduce_kernel", RecordReduce())
+    set_random_seed(0)
+    # K = 2560 is ten BLOCK_K=256 blocks, so all 8 splits contribute.
+    dtype, hidden_size = torch.bfloat16, 2560
+    weights = [
+        torch.randn((3, 16, hidden_size), dtype=dtype, device=DEVICE)
+        for _ in range(nslices)
+    ]
+    target = torch.randn((1, hidden_size), dtype=dtype, device=DEVICE)
+    target_id = torch.zeros(1, dtype=torch.int32, device=DEVICE)
+    expected = _shrink(target, weights, target_id, 0.7)[:, 0]
+    for num_tokens in (31, 32, 33, 127, 128, 129):
+        inputs = torch.randn((num_tokens, hidden_size), dtype=dtype, device=DEVICE)
+        mapping = torch.randint(-1, 3, (num_tokens,), dtype=torch.int32, device=DEVICE)
+        pos = num_tokens // 2
+        inputs[pos], mapping[pos] = target[0], 0
+        out = _shrink(inputs, weights, mapping, 0.7)
+        assert torch.equal(out[:, pos], expected), num_tokens
+    assert len(reduce_grids) == 7

tests/v1/simple_kv_offload/test_scheduler.py

@@ -1864,6 +1864,31 @@ def test_reset_pending_lazy_stores() -> None:
     assert hit_tokens == 0, "CPU cache should be empty after reset"
 
 
+@pytest.mark.parametrize("num_groups", [1, 2])
+@pytest.mark.parametrize("lazy", [False, True])
+def test_reset_releases_pending_cpu_hits(num_groups: int, lazy: bool) -> None:
+    """A lookup awaiting GPU allocation must not prevent a cache reset."""
+    fix = make_scheduler(num_groups=num_groups, lazy=lazy)
+    sched = fix.scheduler
+    req = make_request(num_blocks=2)
+    cpu_blocks = []
+    for group_id in range(num_groups):
+        cpu_blocks.extend(
+            _allocate_gpu_blocks(sched.cpu_block_pool, req, 2, group_id=group_id)
+        )
+    sched.cpu_block_pool.free_blocks(cpu_blocks)
+
+    assert sched.get_num_new_matched_tokens(req, 0) == (2 * BLOCK_SIZE, True)
+    assert all(block.ref_cnt == 1 for block in cpu_blocks)
+    assert not sched._reqs_to_load
+
+    assert sched.reset()
+    assert all(block.ref_cnt == 0 for block in cpu_blocks)
+    assert not sched._pending_cpu_hits
+    assert sched.get_num_new_matched_tokens(req, 0) == (0, False)
+    assert sched.reset()
+
+
 # ---------------------------------------------------------------------------
 # Test 14: Reset with pending loads waits for completion
 # ---------------------------------------------------------------------------

tests/v1/worker/test_gpu_model_runner.py

@@ -486,6 +486,67 @@ def test_preferred_block_size_rejects_backends_with_no_common_size():
         Platform._preferred_block_size_for_backends(classes, 16, None)
 
 
+@pytest.mark.parametrize("cpu", [False, True])
+def test_alignment_rejected_by_a_sibling_backend_raises(monkeypatch, cpu):
+    # Alignment that lands on a size an exact-size sibling rejects must fail
+    # here, not later in select_common_block_size().
+    backends = [_mock_backend([MultipleOf(16)]), _mock_backend([16], exact=True)]
+    monkeypatch.setattr(
+        Platform, "_find_non_ssm_backends", classmethod(lambda cls, c: backends)
+    )
+
+    def bump(cls, vllm_config, backend_cls):
+        vllm_config.cache_config.block_size = 32
+
+    monkeypatch.setattr(Platform, "_align_hybrid_block_size", classmethod(bump))
+    cache_config = SimpleNamespace(
+        block_size=16, user_specified_block_size=True, kv_cache_dtype_skip_layers=None
+    )
+    vllm_config = SimpleNamespace(
+        cache_config=cache_config, model_config=SimpleNamespace(is_hybrid=True)
+    )
+    from vllm.platforms.cpu import CpuPlatform  # overrides the method; must check too
+
+    platform = CpuPlatform if cpu else Platform
+    with pytest.raises(ValueError, match="MOCK_EXACT"):
+        platform.update_block_size_for_backend(vllm_config)
+
+
+def test_xpu_gdn_rounding_rejected_by_a_sibling_backend_raises(monkeypatch):
+    # XPU's GDN rounding runs after super()'s check, so it must check again.
+    try:
+        from vllm.platforms.xpu import XPUPlatform
+    except ImportError:
+        pytest.skip("vllm_xpu_kernels not importable outside an XPU stack")
+
+    backends = [_mock_backend([MultipleOf(16)]), _mock_backend([16], exact=True)]
+    monkeypatch.setattr(
+        Platform, "update_block_size_for_backend", classmethod(lambda cls, c: None)
+    )
+    monkeypatch.setattr(
+        Platform, "_find_non_ssm_backends", classmethod(lambda cls, c: backends)
+    )
+    gdn_layer = SimpleNamespace(
+        get_attn_backend=lambda: SimpleNamespace(get_name=lambda: "GDN_ATTN")
+    )
+    monkeypatch.setattr(
+        "vllm.config.vllm.get_layers_from_vllm_config",
+        lambda vllm_config, layer_type: {"gdn": gdn_layer},
+    )
+    cache_config = SimpleNamespace(
+        block_size=16,
+        mamba_cache_mode="none",
+        mamba_page_size_padded=None,
+        user_specified_block_size=True,
+        kv_cache_dtype_skip_layers=None,
+    )
+    vllm_config = SimpleNamespace(
+        cache_config=cache_config, model_config=SimpleNamespace(is_hybrid=True)
+    )
+    with pytest.raises(ValueError, match="MOCK_EXACT"):
+        XPUPlatform.update_block_size_for_backend(vllm_config)
+
+
 def test_set_active_mm_loras_builds_tower_and_connector_mappings():
     model = Mock()
     model.get_mm_lora_token_counts.side_effect = (

vllm/compilation/breakable_cudagraph.py

@@ -34,9 +34,9 @@ from typing import Any, ClassVar, TypeVar
 import torch
 
 import vllm.envs as envs
+from vllm.compilation.cudagraph_pool import capture_pool
 from vllm.compilation.monitor import validate_cudagraph_capturing_enabled
 from vllm.config import CUDAGraphMode, VllmConfig
-from vllm.distributed.device_communicators.pynccl_allocator import set_graph_pool_id
 from vllm.forward_context import (
     BatchDescriptor,
     get_forward_context,
@@ -370,11 +370,6 @@ class BreakableCUDAGraphWrapper:
 
         entry.input_addresses = self._collect_tensor_addresses(args, kwargs)
 
-        if self.graph_pool is not None:
-            set_graph_pool_id(self.graph_pool)
-        else:
-            set_graph_pool_id(current_platform.graph_pool_handle())
-
         # Match torch.cuda.graph()'s pre-capture cleanup, which we bypass.
         # Skip it when gc is disabled: bulk capture runs under
         # freeze_gc_for_cudagraph_capture, which already did this cleanup,
@@ -386,8 +381,10 @@ class BreakableCUDAGraphWrapper:
         # pre-capture prefetches are complete and don't leak into the graph.
         get_offloader().sync_prev_onload()
 
-        capture = BreakableCUDAGraphCapture(pool=self.graph_pool)
+        with (
-        with capture:
+            capture_pool(self.graph_pool, self.vllm_config) as pool,
+            BreakableCUDAGraphCapture(pool=pool) as capture,
+        ):
             output = self.runnable(*args, **kwargs)
             # Join the offloader's copy stream while we still hold the last
             # segment open, so the join is captured into the graph (otherwise

vllm/compilation/cuda_graph.py

@@ -13,9 +13,9 @@ from unittest.mock import patch
 import torch
 
 from vllm.compilation.counter import compilation_counter
+from vllm.compilation.cudagraph_pool import capture_pool
 from vllm.compilation.monitor import validate_cudagraph_capturing_enabled
 from vllm.config import CUDAGraphMode, VllmConfig
-from vllm.distributed.device_communicators.pynccl_allocator import set_graph_pool_id
 from vllm.forward_context import (
     BatchDescriptor,
     get_forward_context,
@@ -302,10 +302,9 @@ class CUDAGraphWrapper:
                         )
                     )
 
-                if self.graph_pool is not None:
+                graph_pool = stack.enter_context(
-                    set_graph_pool_id(self.graph_pool)
+                    capture_pool(self.graph_pool, self.vllm_config)
-                else:
+                )
-                    set_graph_pool_id(current_platform.graph_pool_handle())
 
                 # Sync offloader's copy stream before capture.
                 # Ensure any pre-capture prefetches from offloader are complete.
@@ -314,7 +313,7 @@ class CUDAGraphWrapper:
                 # mind-exploding: carefully manage the reference and memory.
                 with torch.cuda.graph(
                     cudagraph,
-                    pool=self.graph_pool,
+                    pool=graph_pool,
                     stream=current_stream(),
                 ):
                     # `output` is managed by pytorch's cudagraph pool

vllm/compilation/cudagraph_pool.py

@@ -0,0 +1,45 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Graph pool routing for CUDA graph capture under sleep mode."""
+
+from collections.abc import Iterator
+from contextlib import contextmanager
+from contextvars import ContextVar
+from typing import TYPE_CHECKING, cast
+
+from vllm.config import VllmConfig
+from vllm.device_allocator import get_mem_allocator_instance
+from vllm.distributed.device_communicators.pynccl_allocator import set_graph_pool_id
+from vllm.platforms import current_platform
+
+if TYPE_CHECKING:
+    from vllm.device_allocator.cumem import CuMemAllocator
+
+_outside_cumem: ContextVar[bool] = ContextVar("capture_outside_cumem", default=False)
+
+
+@contextmanager
+def capture_outside_cumem_pool() -> Iterator[None]:
+    """CUDA graphs captured here stay out of cuMem, so that memory profiling can
+    free its throwaway graphs normally."""
+    token = _outside_cumem.set(True)
+    try:
+        yield
+    finally:
+        _outside_cumem.reset(token)
+
+
+@contextmanager
+def capture_pool(
+    pool: tuple[int, int] | None, vllm_config: VllmConfig
+) -> Iterator[tuple[int, int] | None]:
+    """Yield the pool to capture into, the cuMem graph pool when sleep offloads
+    graph memory, and point NCCL's graph allocator at it."""
+    if vllm_config.use_cumem_cudagraph_pool and not _outside_cumem.get():
+        allocator = cast("CuMemAllocator", get_mem_allocator_instance())
+        with allocator.cudagraph_pool() as cumem_pool:
+            set_graph_pool_id(cumem_pool)
+            yield cumem_pool
+    else:
+        set_graph_pool_id(pool or current_platform.graph_pool_handle())
+        yield pool

vllm/config/model.py

@@ -111,6 +111,8 @@ LayerBlockType = Literal["attention", "linear_attention", "mamba"]
 _ATTENTION_LAYER_TYPES = frozenset(
     {
         "full_attention",
+        "indexed_attention",
+        # TODO: Delete below once Transformers 5.18.0 is the minimum required version.
         "deepseek_sparse_attention",
         "qwen_sparse_attention",
     }
@@ -376,6 +378,9 @@ class ModelConfig:
     (default) uses the built-in ``CuMemAllocator`` and is behavior-compatible
     with prior releases. Additional backends (CUDA checkpoint, CRIU, durable
     snapshot) may be registered in-tree or by plugins (RFC #34303)."""
+    sleep_mode_offload_cudagraph: bool = False
+    """Back up CUDA graph memory to CPU during sleep, restored in place on wake.
+    Takes effect with enable_sleep_mode, the cumem backend and CUDA graphs."""
     enable_nccl_comm_suspend: bool = False
     """Enable releasing NCCL communicator memory during sleep mode
     (``ncclCommSuspend``/``ncclCommResume``). Experimental; when disabled

vllm/config/speculative.py

@@ -960,7 +960,7 @@ class SpeculativeConfig:
                     "architectures": [architecture],
                 }
             )
-        if hf_config.model_type in ("longcat_flash", "longcat_flash_ngram"):
+        if hf_config.model_type == "longcat_flash":
             hf_config.model_type = "longcat_flash_mtp"
             n_predict = getattr(hf_config, "num_nextn_predict_layers", 1)
             hf_config.update(

vllm/config/vllm.py

@@ -704,6 +704,21 @@ class VllmConfig:
         """
         return 1 + self.num_speculative_tokens
 
+    @property
+    def use_cumem_cudagraph_pool(self) -> bool:
+        """Whether CUDA graphs go to the cuMem pool that sleep offloads."""
+        from vllm.platforms import current_platform
+
+        model_config = self.model_config
+        return (
+            model_config is not None
+            and model_config.sleep_mode_offload_cudagraph
+            and model_config.enable_sleep_mode
+            and model_config.sleep_mode_backend == "cumem"
+            and self.compilation_config.cudagraph_mode != CUDAGraphMode.NONE
+            and current_platform.is_cuda()
+        )
+
     @property
     def use_v2_model_runner(self) -> bool:
         if self.attention_config.hisparse_config is not None:
@@ -2328,6 +2343,14 @@ class VllmConfig:
                 custom_ops.append("+quant_fp8")
 
         self._verify_kv_transfer_compat()
+        if self.use_cumem_cudagraph_pool:
+            # NCCL graph registration pins the offloaded pool; workers inherit this.
+            value = os.environ.setdefault("NCCL_GRAPH_REGISTER", "0")
+            if value != "0":
+                logger.warning(
+                    "NCCL_GRAPH_REGISTER=%s pins the CUDA graph pool during sleep.",
+                    value,
+                )
         # Log the custom passes that are enabled
         self.compilation_config.pass_config.log_enabled_passes()
 

vllm/device_allocator/cumem.py

@@ -430,6 +430,23 @@ class CuMemAllocator:
             if allocation["allocated_size"] == 0 and data and not data.is_asleep:
                 unmap_and_release(self._python_free_callback(allocation["address"]))
 
+    @contextmanager
+    def cudagraph_pool(self) -> Iterator[tuple[int, int]]:
+        """Tag CUDA graph capture allocations and yield the graph pool id."""
+        if "cudagraph" not in self.allocator_and_pools:
+            allocator = get_pluggable_allocator(
+                self.python_malloc_callback, self.python_free_callback
+            )
+            mem_pool = torch.cuda.memory.MemPool(allocator._allocator)
+            self.allocator_and_pools["cudagraph"] = [(mem_pool, allocator)]
+        old_tag = self.current_tag
+        self.current_tag = "cudagraph"
+        try:
+            # capture_begin routes allocations to this pool; no use_mem_pool.
+            yield self.allocator_and_pools["cudagraph"][0][0].id
+        finally:
+            self.current_tag = old_tag
+
     def get_current_usage(self) -> int:
         """Get the total number of bytes allocated in the memory pool."""
         sum_bytes: int = 0

vllm/device_allocator/sleep_mode_backend.py

@@ -129,8 +129,11 @@ class CuMemBackend(SleepModeBackend):
 
         self._state = "SUSPENDED"
         allocator = get_mem_allocator_instance()
-        # Runtime state is kept at every level; weights only at level 1.
+        # Runtime state and graph pools are kept at every level; weights only
-        offload_tags = ("weights", "runtime") if level == 1 else ("runtime",)
+        # at level 1.
+        offload_tags: tuple[str, ...] = ("runtime", "cudagraph")
+        if level == 1:
+            offload_tags = ("weights", *offload_tags)
         allocator.sleep(offload_tags=offload_tags)
 
     def resume(self, tags: list[str] | None = None) -> None:

vllm/distributed/device_communicators/cuda_communicator.py

@@ -7,6 +7,7 @@ from torch.distributed import ProcessGroup
 
 import vllm.envs as envs
 from vllm._aiter_ops import rocm_aiter_ops
+from vllm.config import get_current_vllm_config_or_none
 from vllm.distributed.device_communicators.all_reduce_utils import (
     NCCL_SYMM_MEM_ALL_REDUCE_CONFIG,
     should_nccl_symm_mem_ag_rs,
@@ -114,6 +115,10 @@ class CudaCommunicator(DeviceCommunicatorBase):
         self.aiter_ar_comm: AiterCustomAllreduce | None = None
         self.use_aiter_ag_rs: bool = False
 
+        # cuMem graph buffers cannot be IPC-registered; capture copies them instead.
+        config = get_current_vllm_config_or_none()
+        register = config is None or not config.use_cumem_cudagraph_pool
+
         if use_torch_symm_mem and current_platform.is_cuda():
             self.symm_mem_comm = SymmMemCommunicator(
                 group=self.cpu_group,
@@ -151,6 +156,7 @@ class CudaCommunicator(DeviceCommunicatorBase):
                 symm_mem_enabled=(
                     self.symm_mem_comm is not None and not self.symm_mem_comm.disabled
                 ),
+                register_graph_buffers=register,
             )
 
         # AITER custom all-gather/reduce-scatter DP-attention dispatch/combine

vllm/distributed/device_communicators/custom_all_reduce.py

@@ -147,18 +147,22 @@ class CustomAllreduce:
             _DEFAULT_MNNVL_MULTIMEM_REDUCE_SCATTER_MAX_SIZE
         ),
         symm_mem_enabled=False,
+        *,
+        register_graph_buffers: bool = True,
     ) -> None:
         """Args:
             group: the process group to work on. If None, it will use the
                 default process group.
             device: the device to bind the CustomAllreduce to. If None,
                 it will be bound to f"cuda:{local_rank}".
+            register_graph_buffers: whether graph capture IPC-registers its buffers.
         It is the caller's responsibility to make sure each communicator
         is bind to a unique device, and all communicators in this group
         are in the same node.
 
         """
         self._IS_CAPTURING = False
+        self._capture_registered = register_graph_buffers
         self._ptr = 0
         self.disabled = True
         self.mnnvl_buffer = None
@@ -549,7 +553,7 @@ class CustomAllreduce:
             return None
         if self._IS_CAPTURING:
             if torch.cuda.is_current_stream_capturing():
-                return self.all_reduce(input, registered=True)
+                return self.all_reduce(input, registered=self._capture_registered)
             else:
                 # If warm up, mimic the allocation pattern since custom
                 # allreduce is out-of-place.

vllm/distributed/device_communicators/flashinfer_all_reduce.py

@@ -17,7 +17,7 @@ from vllm.config.compilation import PassConfig
 from vllm.distributed.device_communicators.all_reduce_utils import (
     FI_MNNVL_ALLREDUCE_MAX_SIZE_MB,
 )
-from vllm.distributed.parallel_state import _node_count, get_node_count
+from vllm.distributed.parallel_state import _node_count
 from vllm.logger import init_logger
 from vllm.platforms import current_platform
 from vllm.utils.gpu_sync_debug import gpu_sync_allowed
@@ -140,14 +140,13 @@ def _create_workspace(
     return workspace
 
 
-def _resolve_fi_ar_backend() -> tuple[str, bool]:
+def _resolve_fi_ar_backend(group: ProcessGroup) -> tuple[str | None, bool]:
     """Resolve the flashinfer allreduce backend for the current setup.
 
     Returns:
-        A ``(backend, allow_trtllm_fallback)`` tuple. ``allow_trtllm_fallback``
+        A ``(backend, allow_trtllm_fallback)`` tuple. ``backend`` is ``None``
-        is True only when ``auto`` selects mnnvl for a single node, so that
+        when automatic selection has no supported backend for the process
-        workspace creation can fall back to trtllm on single-node topologies
+        group.
-        without NVSwitch multicast support (where mnnvl is unavailable).
 
     """
     backend = envs.VLLM_FLASHINFER_ALLREDUCE_BACKEND
@@ -155,15 +154,23 @@ def _resolve_fi_ar_backend() -> tuple[str, bool]:
         logger.debug_once("Using flashinfer allreduce backend: %s", backend)
         return backend, False
 
-    # Default to mnnvl for both single- and multi-node setups. The mnnvl
+    node_count = _node_count(group)
-    # cudagraph hang that previously forced single-node to trtllm
+    if node_count == 1:
-    # (https://github.com/vllm-project/vllm/issues/35772) was fixed upstream in
+        # TRTLLM supports single-node allreduce on Hopper and Blackwell. Do not
-    # FlashInfer (>= 0.6.12, vLLM pins 0.6.15), so mnnvl is safe here. trtllm
+        # create an MNNVL workspace when the process group is node-local.
-    # does not support multi-node allreduce, so mnnvl is required there anyway.
+        backend = "trtllm"
-    # mnnvl needs NVSwitch multicast; on single-node topologies without it,
+        allow_trtllm_fallback = False
-    # fall back to trtllm so fused allreduce stays enabled.
+    elif current_platform.has_device_capability(100):
-    backend = "mnnvl"
+        # MNNVL requires Blackwell-class GPUs and is the only supported
-    allow_trtllm_fallback = get_node_count() == 1
+        # FlashInfer backend for a multi-node allreduce process group.
+        backend = "mnnvl"
+        allow_trtllm_fallback = False
+    else:
+        logger.debug_once(
+            "FlashInfer allreduce fusion is disabled: multi-node process groups "
+            "require an MNNVL-capable Blackwell system."
+        )
+        return None, False
 
     logger.debug_once("Auto-selected flashinfer allreduce backend: %s", backend)
     return backend, allow_trtllm_fallback
@@ -187,9 +194,12 @@ def get_fi_ar_workspace(
     if _fi_ar_workspace is not None or _fi_ar_workspace_failed:
         return _fi_ar_workspace
 
-    backend, allow_trtllm_fallback = _resolve_fi_ar_backend()
+    backend, allow_trtllm_fallback = _resolve_fi_ar_backend(group)
 
-    if get_node_count() > 1 and backend == "trtllm":
+    if backend is None:
+        return None
+
+    if backend == "trtllm" and _node_count(group) > 1:
         raise ValueError(
             "Flashinfer allreduce is not supported for multi-node allreduce with "
             "'trtllm' backend. Please use 'mnnvl' backend instead."
@@ -253,9 +263,12 @@ def get_fi_ar_quant_workspace(
     if _fi_ar_quant_workspace is not None or _fi_ar_quant_workspace_failed:
         return _fi_ar_quant_workspace
 
-    backend, allow_trtllm_fallback = _resolve_fi_ar_backend()
+    backend, allow_trtllm_fallback = _resolve_fi_ar_backend(group)
 
-    if get_node_count() > 1 and backend == "trtllm":
+    if backend is None:
+        return None
+
+    if backend == "trtllm" and _node_count(group) > 1:
         raise ValueError(
             "Flashinfer allreduce quantization fusion is not supported for "
             "multi-node allreduce with 'trtllm' backend. Please use 'mnnvl' "
@@ -397,7 +410,13 @@ class FlashInferAllReduce:
                 self.world_size,
             )
             return
-        backend, _ = _resolve_fi_ar_backend()
+        backend, _ = _resolve_fi_ar_backend(self.group)
+        if backend is None:
+            logger.info(
+                "FlashInfer All Reduce is disabled because no supported backend "
+                "is available for this process group."
+            )
+            return
         tuned_max_size = _get_tuned_standalone_max_size(
             self.world_size,
             backend,

vllm/engine/arg_utils.py

@@ -770,6 +770,7 @@ class EngineArgs:
 
     generation_config: str = ModelConfig.generation_config
     enable_sleep_mode: bool = ModelConfig.enable_sleep_mode
+    sleep_mode_offload_cudagraph: bool = ModelConfig.sleep_mode_offload_cudagraph
     sleep_preserve_parameter_names: list[str] = get_field(
         ModelConfig, "sleep_preserve_parameter_names"
     )
@@ -1007,6 +1008,10 @@ class EngineArgs:
         model_group.add_argument(
             "--enable-sleep-mode", **model_kwargs["enable_sleep_mode"]
         )
+        model_group.add_argument(
+            "--sleep-mode-offload-cudagraph",
+            **model_kwargs["sleep_mode_offload_cudagraph"],
+        )
         model_group.add_argument(
             "--sleep-preserve-parameter-names",
             **model_kwargs["sleep_preserve_parameter_names"],
@@ -1984,6 +1989,7 @@ class EngineArgs:
             generation_config=self.generation_config,
             override_generation_config=self.override_generation_config,
             enable_sleep_mode=self.enable_sleep_mode,
+            sleep_mode_offload_cudagraph=self.sleep_mode_offload_cudagraph,
             sleep_preserve_parameter_names=self.sleep_preserve_parameter_names,
             enable_cumem_allocator=self.enable_cumem_allocator,
             enable_nccl_comm_suspend=self.enable_nccl_comm_suspend,

vllm/entrypoints/openai/chat_completion/serving.py

@@ -631,6 +631,8 @@ class OpenAIServingChat(GenerateBaseServing):
                     if finish_reason_sent[i]:
                         continue
 
+                    self._raise_if_error(output.finish_reason, request_id)
+
                     if request.logprobs and (
                         request.top_logprobs is not None or request.logprob_token_ids
                     ):
@@ -758,10 +760,6 @@ class OpenAIServingChat(GenerateBaseServing):
 
                     # if the model is finished generating
                     else:
-                        # check for error finish reason and abort streaming
-                        # finish_reason='error' indicates a retryable error
-                        self._raise_if_error(output.finish_reason, request_id)
-
                         # Send the finish response for each request.n only once
                         # In OpenAI's API, when a tool is called, the
                         # finish_reason is:

vllm/entrypoints/openai/completion/serving.py

@@ -326,6 +326,8 @@ class OpenAIServingCompletion(GenerateBaseServing):
 
                 for output in res.outputs:
                     i = output.index + prompt_idx * num_choices
+                    finish_reason = output.finish_reason
+                    self._raise_if_error(finish_reason, request_id)
 
                     # Useful when request.return_token_ids is True
                     # Returning prompt token IDs shares the same logic
@@ -392,11 +394,8 @@ class OpenAIServingCompletion(GenerateBaseServing):
 
                     previous_text_lens[i] += len(delta_text)
                     previous_num_tokens[i] += len(output.token_ids)
-                    finish_reason = output.finish_reason
                     stop_reason = output.stop_reason
 
-                    self._raise_if_error(finish_reason, request_id)
-
                     chunk = CompletionStreamResponse(
                         id=request_id,
                         object="text_completion",

vllm/entrypoints/openai/responses/context.py

@@ -656,6 +656,11 @@ class HarmonyContext(ConversationContext):
 
         self.last_append_segments: list[Segment] = []
         self.last_append_flush_status: bool = False
+        # Each completed message with the items sent for it in
+        # output_item.done events, so the final response can reuse their ids.
+        self.streamed_items_by_message: list[
+            tuple[Message, list[ResponseOutputItem]]
+        ] = []
 
         # Turn tracking - replaces multiple individual tracking variables
         self.current_turn_metrics = TurnMetrics()

vllm/entrypoints/openai/responses/serving.py

@@ -12,6 +12,7 @@ from typing import Any, Final, cast
 from fastapi import Request
 from openai.types.responses import (
     ResponseOutputItem,
+    ResponseOutputItemDoneEvent,
     ResponseOutputMessage,
     ResponseOutputText,
     ResponseStatus,
@@ -68,6 +69,7 @@ from vllm.entrypoints.openai.responses.utils import (
     build_response_output_items,
     extract_function_tool_names,
     extract_tool_types,
+    reuse_streamed_item_ids,
 )
 from vllm.entrypoints.serve.engine.protocol import ErrorResponse
 from vllm.entrypoints.serve.utils.api_utils import get_max_tokens
@@ -781,17 +783,21 @@ class OpenAIServingResponses(GenerateBaseServing):
             assert isinstance(context, HarmonyContext)
             output = []
             harmony_msgs = context.messages[context.num_init_messages :]
-            if harmony_msgs:
+            # Streamed messages are a subsequence of harmony_msgs, in order.
-                fn_names = context.function_tool_names
+            streamed = iter(context.streamed_items_by_message)
-                for msg in harmony_msgs[:-1]:
+            next_streamed = next(streamed, None)
-                    output.extend(harmony_to_response_output(msg, fn_names))
+            for i, msg in enumerate(harmony_msgs):
-                output.extend(
+                items = harmony_to_response_output(
-                    harmony_to_response_output(
+                    msg,
-                        harmony_msgs[-1],
+                    context.function_tool_names,
-                        fn_names,
+                    incomplete=(
-                        incomplete=context.last_append_flush_status,
+                        i == len(harmony_msgs) - 1 and context.last_append_flush_status
-                    )
+                    ),
                 )
+                if next_streamed is not None and next_streamed[0] is msg:
+                    reuse_streamed_item_ids(items, next_streamed[1])
+                    next_streamed = next(streamed, None)
+                output.extend(items)
 
             if request.enable_response_messages:
                 input_messages = context.messages[: context.num_init_messages]
@@ -1291,16 +1297,24 @@ class OpenAIServingResponses(GenerateBaseServing):
                         yield _increment_sequence_number_and_return(event)
 
                 elif completed_message := segment.completed_message:
+                    done_items: list[ResponseOutputItem] = []
                     # TODO: Fix browser emitted as MCP calls
                     for event in emit_previous_item_done_events(
                         completed_message, state, ctx.function_tool_names
                     ):
+                        if isinstance(event, ResponseOutputItemDoneEvent):
+                            done_items.append(event.item)
                         yield _increment_sequence_number_and_return(event)
 
                     for event in emit_tool_action_events(
                         completed_message, state, self.tool_server
                     ):
+                        if isinstance(event, ResponseOutputItemDoneEvent):
+                            done_items.append(event.item)
                         yield _increment_sequence_number_and_return(event)
+                    ctx.streamed_items_by_message.append(
+                        (completed_message, done_items)
+                    )
                     state.reset_for_new_item()
 
     async def responses_stream_generator(

vllm/entrypoints/openai/responses/streaming_events.py

@@ -202,7 +202,7 @@ def emit_reasoning_delta_events(
     events: list[StreamingResponsesResponse] = []
     if not state.sent_output_item_added:
         state.sent_output_item_added = True
-        state.current_item_id = f"msg_{random_uuid()}"
+        state.current_item_id = f"rs_{random_uuid()}"
         events.append(
             ResponseOutputItemAddedEvent(
                 type="response.output_item.added",

vllm/entrypoints/openai/responses/utils.py

@@ -112,6 +112,28 @@ def build_response_output_items(
     return outputs
 
 
+def reuse_streamed_item_ids(
+    items: list[ResponseOutputItem],
+    streamed_items: list[ResponseOutputItem],
+) -> None:
+    """Give rebuilt output items the ids already streamed for the same message.
+
+    Items are paired by type in order; an item without a streamed counterpart
+    keeps its own id.
+    """
+    remaining = list(streamed_items)
+    for item in items:
+        for i, streamed in enumerate(remaining):
+            if streamed.type != item.type:
+                continue
+            item.id = streamed.id
+            if isinstance(item, ResponseFunctionToolCall):
+                assert isinstance(streamed, ResponseFunctionToolCall)
+                item.call_id = streamed.call_id
+            del remaining[i]
+            break
+
+
 def should_continue_final_message(
     request_input: str | list[ResponseInputOutputItem],
 ) -> bool:

vllm/entrypoints/scale_out/token_in_token_out/serving.py

@@ -516,11 +516,12 @@ class ServingTokens(GenerateBaseServing):
         request_metadata: RequestResponseMetadata,
     ) -> AsyncGenerator[str, None]:
         num_prompt_tokens = 0
-        num_generated_tokens: list[int] = []
         first_iteration = True
         prompt_token_ids: list[int] | None = None
         num_cached_tokens = None
         sampling_params: SamplingParams = request.sampling_params
+        # With n > 1, the first output may not include every choice yet.
+        num_generated_tokens = [0] * sampling_params.n
         last_res: RequestOutput | None = None
         text_mode = request.output_mode == "text"
         tokenizer = self._logprobs_tokenizer(text_mode)
@@ -540,7 +541,6 @@ class ServingTokens(GenerateBaseServing):
                     if res.encoder_prompt_token_ids is not None:
                         num_prompt_tokens += len(res.encoder_prompt_token_ids)
                     num_cached_tokens = res.num_cached_tokens
-                    num_generated_tokens = [0] * len(res.outputs)
                     first_iteration = False
 
                 for output in res.outputs:
@@ -551,21 +551,15 @@ class ServingTokens(GenerateBaseServing):
                     finish_reason = output.finish_reason
                     self._raise_if_error(finish_reason, request_id)
 
-                    if text_mode:
+                    # Terminal outputs are always emitted so the client sees
-                        # Text held back for stop string matching is flushed
+                    # the finish reason, e.g. an abort, which has no new
-                        # with the final output, and the abort output has no
+                    # tokens. Text mode also emits text held back for stop
-                        # new token IDs, so text or a finish reason is enough.
+                    # string matching, which can arrive without token IDs.
-                        emit = bool(
+                    if not (
-                            delta_token_ids or output.text or finish_reason is not None
+                        delta_token_ids
-                        )
+                        or finish_reason is not None
-                    else:
+                        or (text_mode and output.text)
-                        # Still emit a terminal empty chunk while prompt
+                    ):
-                        # metadata is pending, so zero-token completions
-                        # deliver it.
-                        emit = bool(delta_token_ids) or (
-                            finish_reason is not None and prompt_token_ids is not None
-                        )
-                    if not emit:
                         continue
 
                     if sampling_params.logprobs is not None:

vllm/entrypoints/serve/engine/serving.py

@@ -63,8 +63,12 @@ class BaseServing:
             ):
                 error_response = load_result
 
+        served_names = ", ".join(model.name for model in self.models.base_model_paths)
         return error_response or self.create_error_response(
-            message=f"The model `{request.model}` does not exist.",
+            message=(
+                f"The model `{request.model}` does not exist. "
+                f"Valid aliases: {served_names}."
+            ),
             err_type="NotFoundError",
             status_code=HTTPStatus.NOT_FOUND,
             param="model",

vllm/lora/ops/triton_ops/kernel_utils.py

@@ -257,6 +257,7 @@ def do_shrink_kernel(
     output_d0_stride,
     output_d1_stride,
     output_d2_stride,
+    output_split_stride,
     scaling,
     BLOCK_M: tl.constexpr,
     BLOCK_N: tl.constexpr,
@@ -265,6 +266,7 @@ def do_shrink_kernel(
     SPLIT_K: tl.constexpr,
     SLICE_NUM: tl.constexpr,
     USE_GDC: tl.constexpr,
+    STORE_PARTIALS: tl.constexpr,
 ):
     """Given an array of integers that identifies the rows of A, ram,
     a lora index that identifies which LoRA to use from lora_ptr, lora_index,
@@ -324,17 +326,31 @@ def do_shrink_kernel(
     # Identify the C output pointers to store the results of the accumulator.
     offset_cn = tl.arange(0, BLOCK_N) + pid_n * BLOCK_N
     offset_cm = tl.arange(0, BLOCK_M)
-    cur_out_ptr = out_ptr if SLICE_NUM == 1 else out_ptr + slice_id * output_d0_stride
+    if STORE_PARTIALS:
+        cur_out_ptr = (
+            out_ptr
+            + slice_id.to(tl.int64) * output_d0_stride
+            + pid_sk.to(tl.int64) * output_split_stride
+        )
+    else:
+        cur_out_ptr = (
+            out_ptr if SLICE_NUM == 1 else out_ptr + slice_id * output_d0_stride
+        )
     c_ptr = (
         cur_out_ptr
         + ram[:, None] * output_d1_stride
         + offset_cn[None, :] * output_d2_stride
     )
     c_mask = (offset_cm[:, None] < M_LEN) & (offset_cn[None, :] < N)
-    accumulator *= scaling
 
-    # handles write-back with reduction-splitting
+    if STORE_PARTIALS:
-    if SPLIT_K == 1:
+        # Unscaled FP32 partial per K split; the reducer adds the splits in a
+        # fixed order and applies scaling once.
         tl.store(c_ptr, accumulator, mask=c_mask)
     else:
-        tl.atomic_add(c_ptr, accumulator, mask=c_mask, sem="relaxed")
+        accumulator *= scaling
+        # handles write-back with reduction-splitting
+        if SPLIT_K == 1:
+            tl.store(c_ptr, accumulator, mask=c_mask)
+        else:
+            tl.atomic_add(c_ptr, accumulator, mask=c_mask, sem="relaxed")

vllm/lora/ops/triton_ops/lora_shrink_op.py

@@ -13,11 +13,90 @@ from vllm.lora.ops.triton_ops.kernel_utils import do_shrink_kernel
 from vllm.lora.ops.triton_ops.utils import (
     _get_lora_a_ptr,
     get_lora_op_configs,
+    is_batch_invariant,
     supports_pdl,
 )
 from vllm.triton_utils import tl, triton
 from vllm.utils.torch_utils import direct_register_custom_op
 
+# Batch-invariant shrink splits K into this many fixed partials (no atomics).
+_BI_SPLIT_K = 8
+
+
+@triton.jit
+def _lora_shrink_reduce_kernel(
+    partial_ptr,
+    out_ptr,
+    M,
+    N: tl.constexpr,
+    token_indices_sorted_by_lora_ids,
+    num_tokens_per_lora,
+    lora_token_start_loc,
+    lora_ids,
+    scaling,
+    partial_d0_stride,
+    partial_d1_stride,
+    partial_d2_stride,
+    partial_d3_stride,
+    output_d0_stride,
+    output_d1_stride,
+    output_d2_stride,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    SPLIT_K: tl.constexpr,
+    GROUP_SIZE_M: tl.constexpr,
+):
+    cta_n_num = tl.cdiv(N, BLOCK_N)
+    cta_m_num = tl.cdiv(M, BLOCK_M)
+    pid_m_n = tl.program_id(axis=0)
+    num_pid_in_group = GROUP_SIZE_M * cta_n_num
+    group_id = pid_m_n // num_pid_in_group
+    first_pid_m = group_id * GROUP_SIZE_M
+    group_size_m = min(cta_m_num - first_pid_m, GROUP_SIZE_M)
+    pid_m = first_pid_m + ((pid_m_n % num_pid_in_group) % group_size_m)
+    pid_n = (pid_m_n % num_pid_in_group) // group_size_m
+
+    slice_id = tl.program_id(axis=1)
+    lora_idx = tl.program_id(axis=2)
+    lora_id = tl.load(lora_ids + lora_idx)
+    if lora_id == -1:
+        return
+
+    lora_m_size = tl.load(num_tokens_per_lora + lora_idx)
+    cta_m_offset = pid_m * BLOCK_M
+    if cta_m_offset >= lora_m_size:
+        return
+
+    cta_m_len = min(BLOCK_M, lora_m_size - cta_m_offset)
+    lora_m_indices_start = tl.load(lora_token_start_loc + lora_idx)
+    cta_lora_seq_indices = (
+        token_indices_sorted_by_lora_ids + lora_m_indices_start + cta_m_offset
+    )
+    offset_m = tl.arange(0, BLOCK_M) % cta_m_len
+    ram = tl.load(cta_lora_seq_indices + offset_m)
+    offset_n = tl.arange(0, BLOCK_N) + pid_n * BLOCK_N
+    offset_cm = tl.arange(0, BLOCK_M)
+    mask = (offset_cm[:, None] < cta_m_len) & (offset_n[None, :] < N)
+
+    accumulator = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+    for split_id in range(SPLIT_K):
+        partial_in = (
+            partial_ptr
+            + slice_id.to(tl.int64) * partial_d0_stride
+            + tl.cast(split_id, tl.int64) * partial_d1_stride
+            + ram[:, None] * partial_d2_stride
+            + offset_n[None, :] * partial_d3_stride
+        )
+        accumulator += tl.load(partial_in, mask=mask, other=0.0)
+
+    output = (
+        out_ptr
+        + slice_id * output_d0_stride
+        + ram[:, None] * output_d1_stride
+        + offset_n[None, :] * output_d2_stride
+    )
+    tl.store(output, accumulator * scaling, mask=mask)
+
 
 @triton.jit
 def _lora_shrink_kernel(
@@ -49,7 +128,15 @@ def _lora_shrink_kernel(
     SLICE_NUM: tl.constexpr,
     USE_GDC: tl.constexpr,
     launch_pdl: tl.constexpr,
+    STORE_PARTIALS: tl.constexpr = False,
+    output_split_stride=0,
+    PARTIAL_N: tl.constexpr = 0,
+    PARTIAL_K: tl.constexpr = 0,
 ):
+    if STORE_PARTIALS:
+        # The split-K partial stage specializes on N and K.
+        N = PARTIAL_N
+        K = PARTIAL_K
     cta_n_num = tl.cdiv(N, BLOCK_N)
     cta_m_num = tl.cdiv(M, BLOCK_M)
 
@@ -116,6 +203,7 @@ def _lora_shrink_kernel(
         output_d0_stride,
         output_d1_stride,
         output_d2_stride,
+        output_split_stride,
         scaling,
         BLOCK_M,
         BLOCK_N,
@@ -124,6 +212,7 @@ def _lora_shrink_kernel(
         SPLIT_K,
         SLICE_NUM,
         USE_GDC,
+        STORE_PARTIALS,
     )
 
 
@@ -167,6 +256,23 @@ def _lora_shrink(
 
     """
     assert no_lora_flag_cpu.numel() == 1
+    if inputs.size(0) == 0:
+        # An empty first call has no work to do.
+        return
+    partials: torch.Tensor | None = None
+    if is_batch_invariant:
+        # Allocated before the no-LoRA return so profiling sees the scratch.
+        # Per-call allocation keeps each CUDA graph's scratch in its own pool.
+        partials = torch.empty(
+            (
+                output_tensor.size(0),
+                _BI_SPLIT_K,
+                inputs.size(0),
+                output_tensor.size(-1),
+            ),
+            dtype=torch.float32,
+            device=output_tensor.device,
+        )
     if no_lora_flag_cpu.item():
         # None of the inputs require LoRA.
         return
@@ -213,6 +319,11 @@ def _lora_shrink(
     NUM_STAGES = kernel_config["num_stages"]
     NUM_CTAS = kernel_config["num_ctas"]
     GROUP_SIZE_M = kernel_config.get("group_size_m", 8)
+
+    # Every M must keep the same K partition for batch invariance.
+    if partials is not None:
+        BLOCK_K = 256
+        SPLIT_K = _BI_SPLIT_K
     EVEN_K = K % (BLOCK_K * SPLIT_K) == 0  # type: ignore
 
     # TODO (varun): This grid formulation maximizes parallelization at the
@@ -225,11 +336,24 @@ def _lora_shrink(
     )
 
     # PDL only works when dual-stream is being used.
-    use_gdc = supports_pdl(inputs.device) and envs.VLLM_LORA_ENABLE_DUAL_STREAM
+    use_gdc = (
+        supports_pdl(inputs.device)
+        and envs.VLLM_LORA_ENABLE_DUAL_STREAM
+        and partials is None
+    )
+    first_out = output_tensor if partials is None else partials
+    partial_args = {}
+    if partials is not None:
+        partial_args = dict(
+            STORE_PARTIALS=True,
+            output_split_stride=partials.stride(1),
+            PARTIAL_N=N,
+            PARTIAL_K=K,
+        )
     _lora_shrink_kernel[grid](
         inputs,
         lora_ptr_tensor,
-        output_tensor,
+        first_out,
         M,
         N,
         K,
@@ -243,9 +367,9 @@ def _lora_shrink(
         lora_strides_d0,
         lora_strides_d1,
         lora_strides_d2,
-        output_tensor.stride(0),
+        first_out.stride(0),
-        output_tensor.stride(1),
+        first_out.stride(-2),
-        output_tensor.stride(2),
+        first_out.stride(-1),
         BLOCK_M,
         BLOCK_N,
         BLOCK_K,
@@ -258,7 +382,28 @@ def _lora_shrink(
         num_ctas=NUM_CTAS,
         num_stages=NUM_STAGES,
         launch_pdl=use_gdc,
+        **partial_args,
     )
+    if partials is not None:
+        _lora_shrink_reduce_kernel[(grid[0] // SPLIT_K, *grid[1:])](
+            partials,
+            output_tensor,
+            M,
+            N,
+            token_indices_sorted_by_lora_ids,
+            num_tokens_per_lora,
+            lora_token_start_loc,
+            lora_ids,
+            scaling,
+            *partials.stride(),
+            *output_tensor.stride(),
+            BLOCK_M,
+            BLOCK_N,
+            SPLIT_K,
+            GROUP_SIZE_M,
+            num_warps=NUM_WARPS,
+            num_stages=NUM_STAGES,
+        )
 
     return
 

vllm/model_executor/determinism/batch_invariant.py

@@ -918,7 +918,7 @@ def mean_batch_invariant(input, dim, keepdim=False, dtype: torch.dtype | None =
         for d in sorted_dims:
             result = result.squeeze(d)
 
-    return result
+    return result if dtype is not None else result.to(input.dtype)
 
 
 @triton.jit

vllm/model_executor/models/lightonocr.py

@@ -50,6 +50,19 @@ class LightOnOCRMultiModalProcessor(BaseMultiModalProcessor[LightOnOCRProcessing
     def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str:
         return self.dummy_inputs.get_dummy_text(mm_counts)
 
+    def _call_hf_processor(
+        self,
+        hf_data: Mapping[str, object],
+        hf_kwargs: Mapping[str, object],
+    ) -> BatchFeature:
+        # The checkpoint sets `do_pad: null`, which disables padding since
+        # Transformers v5.18, so images of different sizes cannot be stacked.
+        return self.info.ctx.call_hf_processor(
+            self.info.get_hf_processor(**hf_kwargs),
+            hf_data,
+            {**hf_kwargs, "do_pad": True},
+        )
+
     def _postprocess_hf_mm_data(
         self,
         mm_data: Mapping[str, object],

vllm/model_executor/models/longcat_flash.py

@@ -39,7 +39,7 @@ from itertools import islice
 
 import torch
 from torch import nn
-from transformers import PreTrainedConfig
+from transformers import LongcatFlashConfig
 
 from vllm.compilation.decorators import support_torch_compile
 from vllm.config import CacheConfig, VllmConfig
@@ -80,121 +80,6 @@ from .utils import (
 logger = init_logger(__name__)
 
 
-class FlashConfig(PreTrainedConfig):
-    """Flash model configuration."""
-
-    moe_intermediate_size: int
-
-    model_type = "longcat_flash"
-    keys_to_ignore_at_inference = ["past_key_values"]
-
-    def __init__(
-        self,
-        vocab_size=131072,
-        hidden_size=4096,
-        intermediate_size=8192,
-        num_layers=28,
-        num_hidden_layers=None,
-        num_attention_heads=96,
-        num_key_value_heads=128,
-        ep_size=1,
-        kv_lora_rank=512,
-        q_lora_rank=1536,
-        qk_rope_head_dim=64,
-        v_head_dim=128,
-        qk_nope_head_dim=128,
-        num_experts_per_tok=None,
-        norm_topk_prob=False,
-        max_position_embeddings=8192,
-        initializer_range=0.02,
-        rms_norm_eps=1e-05,
-        use_cache=True,
-        pad_token_id=None,
-        bos_token_id=100000,
-        eos_token_id=100001,
-        pretraining_tp=1,
-        tie_word_embeddings=False,
-        rope_parameters=None,
-        attention_bias=False,
-        attention_dropout=0.0,
-        mla_scale_q_lora=False,
-        mla_scale_kv_lora=False,
-        dtype="bfloat16",
-        params_dtype="bfloat16",
-        router_dtype="float32",
-        router_bias=False,
-        topk_method=None,
-        routed_scaling_factor=1.0,
-        zero_expert_num=0,
-        zero_expert_type=None,
-        nextn_use_scmoe=False,
-        **kwargs,
-    ):
-        super().__init__(
-            pad_token_id=pad_token_id,
-            bos_token_id=bos_token_id,
-            eos_token_id=eos_token_id,
-            tie_word_embeddings=tie_word_embeddings,
-            dtype=dtype,
-            params_dtype=params_dtype,
-            router_dtype=router_dtype,
-            topk_method=topk_method,
-            router_bias=router_bias,
-            nextn_use_scmoe=nextn_use_scmoe,
-            **kwargs,
-        )
-        self.vocab_size = vocab_size
-        self.max_position_embeddings = max_position_embeddings
-        self.hidden_size = hidden_size
-        self.num_hidden_layers = (
-            num_hidden_layers if num_hidden_layers is not None else num_layers
-        )
-        self.num_attention_heads = num_attention_heads
-        self.ep_size = ep_size
-        self.kv_lora_rank = kv_lora_rank
-        self.q_lora_rank = q_lora_rank
-        self.qk_rope_head_dim = qk_rope_head_dim
-        self.v_head_dim = v_head_dim
-        self.qk_nope_head_dim = qk_nope_head_dim
-        self.num_experts_per_tok = num_experts_per_tok
-        self.norm_topk_prob = norm_topk_prob
-        # for backward compatibility
-        if num_key_value_heads is None:
-            num_key_value_heads = num_attention_heads
-
-        self.num_key_value_heads = num_key_value_heads
-        self.initializer_range = initializer_range
-        self.rms_norm_eps = rms_norm_eps
-        self.pretraining_tp = pretraining_tp
-        self.use_cache = use_cache
-        # Try to set `rope_scaling` if available, otherwise use `rope_parameters`
-        rope_scaling = kwargs.pop("rope_scaling", None)
-        rope_parameters = rope_scaling or rope_parameters or {"rope_type": "default"}
-        rope_theta = kwargs.pop("rope_theta", 1000000.0)
-        if "rope_theta" not in rope_parameters:
-            rope_parameters["rope_theta"] = rope_theta
-        self.rope_parameters = rope_parameters
-        self.attention_bias = attention_bias
-        self.attention_dropout = attention_dropout
-        self.mla_scale_q_lora = mla_scale_q_lora
-        self.mla_scale_kv_lora = mla_scale_kv_lora
-        self.zero_expert_num = zero_expert_num
-        self.zero_expert_type = zero_expert_type
-        self.routed_scaling_factor = routed_scaling_factor
-        self.hidden_act = "silu"
-        self.intermediate_size = (
-            self.ffn_hidden_size
-            if hasattr(self, "ffn_hidden_size")
-            else intermediate_size
-        )
-        if hasattr(self, "moe_intermediate_size"):
-            self.moe_intermediate_size = self.moe_intermediate_size
-        elif hasattr(self, "expert_ffn_hidden_size"):
-            self.moe_intermediate_size = self.expert_ffn_hidden_size
-        else:
-            self.moe_intermediate_size = self.intermediate_size
-
-
 class FlashMLP(nn.Module):
     """Flash MLP layer."""
 
@@ -242,22 +127,17 @@ class FlashMLP(nn.Module):
 class LongcatRouter(nn.Module):
     def __init__(
         self,
-        config: FlashConfig,
+        config: LongcatFlashConfig,
         zero_expert_num: int,
         router_params_dtype: torch.dtype,
         prefix: str = "",
     ):
         super().__init__()
-        self.n_routed_experts = (
+        self.n_routed_experts = config.n_routed_experts + zero_expert_num
-            config.n_routed_experts
-            if hasattr(config, "n_routed_experts")
-            else config.num_experts[0]
-        )
-        self.n_routed_experts = self.n_routed_experts + zero_expert_num
         self.classifier = GateLinear(
             config.hidden_size,
             self.n_routed_experts,
-            bias=config.router_bias,
+            bias=False,
             out_dtype=router_params_dtype,
             params_dtype=router_params_dtype,
             prefix=f"{prefix}.classifier",
@@ -274,7 +154,7 @@ class LongcatRouter(nn.Module):
 class LongcatMoe(nn.Module):
     def __init__(
         self,
-        config: FlashConfig,
+        config: LongcatFlashConfig,
         num_experts: int,
         top_k: int,
         hidden_size: int,
@@ -286,10 +166,7 @@ class LongcatMoe(nn.Module):
     ):
         super().__init__()
         self.hidden_size = hidden_size
-        # Gate always runs at half / full precision for now.
+        self.router_params_dtype = torch.float32
-        self.router_params_dtype = params_dtype
-        if config.router_dtype == "float32":
-            self.router_params_dtype = torch.float32
 
         self.router = LongcatRouter(
             config=config,
@@ -354,7 +231,7 @@ class FlashDecoderLayer(nn.Module):
     def __init__(
         self,
         vllm_config: VllmConfig,
-        config: FlashConfig,
+        config: LongcatFlashConfig,
         cache_config: CacheConfig | None = None,
         quant_config: QuantizationConfig | None = None,
         prefix: str = "",
@@ -376,9 +253,7 @@ class FlashDecoderLayer(nn.Module):
                     qk_nope_head_dim=config.qk_nope_head_dim,
                     qk_rope_head_dim=config.qk_rope_head_dim,
                     v_head_dim=config.v_head_dim,
-                    q_lora_rank=(
+                    q_lora_rank=config.q_lora_rank,
-                        config.q_lora_rank if hasattr(config, "q_lora_rank") else None
-                    ),
                     kv_lora_rank=config.kv_lora_rank,
                     max_position_embeddings=max_position_embeddings,
                     cache_config=cache_config,
@@ -415,12 +290,8 @@ class FlashDecoderLayer(nn.Module):
 
         self.mlp = LongcatMoe(
             config=config,
-            num_experts=config.n_routed_experts
+            num_experts=config.n_routed_experts,
-            if hasattr(config, "n_routed_experts")
+            top_k=config.moe_topk,
-            else config.num_experts[self.layer_idx],
-            top_k=config.moe_topk
-            if hasattr(config, "moe_topk")
-            else config.num_experts_per_tok,
             hidden_size=config.hidden_size,
             intermediate_size=config.moe_intermediate_size,
             quant_config=quant_config,
@@ -482,7 +353,7 @@ class FlashModel(nn.Module):
 
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         super().__init__()
-        config = FlashConfig(**vllm_config.model_config.hf_config.__dict__)
+        config = vllm_config.model_config.hf_config
         cache_config = vllm_config.cache_config
         quant_config = vllm_config.quant_config
         self.config = config
@@ -499,7 +370,7 @@ class FlashModel(nn.Module):
         else:
             self.embed_tokens = PPMissingLayer()
         self.start_layer, self.end_layer, self.layers = make_layers(
-            config.num_hidden_layers,
+            config.num_layers,
             lambda prefix: FlashDecoderLayer(
                 vllm_config,
                 config,
@@ -561,9 +432,7 @@ class FlashModel(nn.Module):
             ckpt_gate_proj_name="gate_proj",
             ckpt_down_proj_name="down_proj",
             ckpt_up_proj_name="up_proj",
-            num_experts=self.config.n_routed_experts
+            num_experts=self.config.n_routed_experts,
-            if hasattr(self.config, "n_routed_experts")
-            else self.config.num_experts[0],
         )
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
@@ -659,7 +528,7 @@ class FlashModel(nn.Module):
                     )
                     weight_loader(param, loaded_weight)
             loaded_params.add(name)
-        for layer_id in range(self.config.num_hidden_layers):
+        for layer_id in range(self.config.num_layers):
             for i in range(2):
                 if isinstance(self.layers[layer_id], PPMissingLayer):
                     continue
@@ -727,15 +596,10 @@ class LongcatFlashForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
 
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         super().__init__()
-        config = FlashConfig(**vllm_config.model_config.hf_config.__dict__)
+        config = vllm_config.model_config.hf_config
         quant_config = vllm_config.quant_config
 
         self.config = config
-        config.intermediate_size = (
-            config.ffn_hidden_size
-            if hasattr(config, "ffn_hidden_size")
-            else config.intermediate_size
-        )
 
         self.quant_config = quant_config
 

vllm/model_executor/models/longcat_flash_mtp.py

@@ -20,7 +20,6 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
     VocabParallelEmbedding,
 )
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
-from vllm.model_executor.models.longcat_flash import FlashConfig
 from vllm.sequence import IntermediateTensors
 
 from .deepseek_v2 import DeepseekV2DecoderLayer
@@ -80,9 +79,9 @@ class LongCatMultiTokenPredictor(nn.Module):
         prefix: str = "",
     ):
         super().__init__()
-        config = FlashConfig(**vllm_config.model_config.hf_config.__dict__)
+        config = vllm_config.model_config.hf_config
-        vllm_config.model_config.hf_config.intermediate_size = config.intermediate_size
+        # TODO: Use `num_hidden_layers` once Transformers 5.18.0 is the minimum
-        self.mtp_start_layer_idx = config.num_hidden_layers * 2
+        self.mtp_start_layer_idx = config.num_layers * 2
         self.num_mtp_layers = 1
         self.layers = torch.nn.ModuleDict(
             {
@@ -130,7 +129,7 @@ class LongCatFlashMTP(nn.Module):
         # builds a dense MLP. object.__setattr__ bypasses the ngram remote
         # config's strict validation (it rejects setting the int field to None).
         object.__setattr__(vllm_config.model_config.hf_config, "n_routed_experts", None)
-        self.config = FlashConfig(**vllm_config.model_config.hf_config.__dict__)
+        self.config = vllm_config.model_config.hf_config
         self.quant_config = (
             None
             if "mtp" in getattr(self.config, "disable_quant_module", [])
@@ -264,7 +263,7 @@ class LongCatFlashMTP(nn.Module):
                 weight_loader = getattr(param, "weight_loader", default_weight_loader)
                 weight_loader(param, loaded_weight)
             loaded_params.add(name)
-        spec_layer_id = self.config.num_hidden_layers * 2
+        spec_layer_id = self.model.mtp_start_layer_idx
         self_attn = self.model.layers[str(spec_layer_id)].mtp_block.self_attn
         if (
             self.quant_config is not None
@@ -356,5 +355,5 @@ class LongCatFlashMTP(nn.Module):
         self, config: PreTrainedConfig, weight_name: str
     ) -> int | None:
         if "model.mtp" in weight_name:
-            return config.num_hidden_layers * 2
+            return self.model.mtp_start_layer_idx
         return None

vllm/model_executor/models/longcat_flash_ngram.py

@@ -15,6 +15,7 @@ from typing import Any
 
 import torch
 from torch import nn
+from transformers import LongcatFlashConfig
 
 from vllm import _custom_ops as ops
 from vllm.config import VllmConfig
@@ -30,15 +31,15 @@ from vllm.v1.worker.gpu.model_states.default import DefaultModelState
 from vllm.v1.worker.gpu.states import RequestState
 
 from .interfaces import SupportsLoRA, SupportsPP
-from .longcat_flash import FlashConfig, FlashModel
+from .longcat_flash import FlashModel
 from .utils import AutoWeightsLoader, PPMissingLayer, WeightsMapper, maybe_prefix
 
 
-def uses_ngram_embedding(config: FlashConfig) -> bool:
+def uses_ngram_embedding(config: LongcatFlashConfig) -> bool:
     return getattr(config, "ngram_vocab_size_ratio", None) is not None
 
 
-def _config_dtype(config: FlashConfig) -> torch.dtype:
+def _config_dtype(config: LongcatFlashConfig) -> torch.dtype:
     dt = getattr(config, "torch_dtype", None) or getattr(config, "dtype", None)
     if isinstance(dt, torch.dtype):
         return dt
@@ -54,7 +55,7 @@ class NgramEmbedding(nn.Module):
     single ``bmm``. Hashing math is ported from the HF reference.
     """
 
-    def __init__(self, config: FlashConfig, base_embeddings: nn.Module) -> None:
+    def __init__(self, config: LongcatFlashConfig, base_embeddings: nn.Module) -> None:
         super().__init__()
         self.config = config
         self.word_embeddings = base_embeddings
@@ -170,14 +171,6 @@ class FlashNgramModel(FlashModel):
     """FlashModel whose input embedding is an :class:`NgramEmbedding`."""
 
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
-        # Each FlashDecoderLayer is a *dual* layer (2 attentions), so the number
-        # of decoder layers is ``num_layers``. The ngram HF config sets
-        # ``num_hidden_layers`` to a multiple of that (attention-module count),
-        # which FlashModel would otherwise build as too many (dead) layers.
-        hf = vllm_config.model_config.hf_config
-        num_layers = getattr(hf, "num_layers", None)
-        if num_layers is not None and hf.num_hidden_layers != num_layers:
-            hf.num_hidden_layers = num_layers
         super().__init__(vllm_config=vllm_config, prefix=prefix)
         self.ngram_embeddings: NgramEmbedding | None
         if get_pp_group().is_first_rank and uses_ngram_embedding(self.config):
@@ -222,10 +215,7 @@ class LongcatFlashNgramForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
                 "V2 model runner for its n-gram embedding state; it is selected "
                 "automatically unless VLLM_USE_V2_MODEL_RUNNER=0 is set."
             )
-        config = FlashConfig(**vllm_config.model_config.hf_config.__dict__)
+        config = vllm_config.model_config.hf_config
-        config.intermediate_size = getattr(
-            config, "ffn_hidden_size", config.intermediate_size
-        )
         self.config = config
         self.quant_config = vllm_config.quant_config
 

vllm/models/deepseek_v4/xpu/xpu_sparse_decode_fp8.py

@@ -166,8 +166,7 @@ def xpu_sparse_decode_fp8(
     num_tokens = q.shape[0]
     device = q.device
 
-    # Determine max topk and swa widths
+    if not swa_only and topk_indices is not None and kv_cache is not None:
-    if not swa_only and topk_indices is not None:
         topk_idx_2d = (
             topk_indices.squeeze(1) if topk_indices.dim() == 3 else topk_indices
         )
@@ -176,109 +175,57 @@ def xpu_sparse_decode_fp8(
         topk_idx_2d = None
         max_topk = 0
 
+    use_topk = topk_idx_2d is not None
     swa_idx_2d = swa_indices.squeeze(1) if swa_indices.dim() == 3 else swa_indices
     max_swa = swa_idx_2d.shape[1]
-
     K_total = max_topk + max_swa
 
-    # Allocate flat workspace: [num_tokens * K_total, 512] bf16
+    # Topk rows precede SWA rows in the flat workspace.
+    n_topk_rows = num_tokens * max_topk
     workspace = torch.empty(
         (num_tokens * K_total, OUTPUT_DIM), dtype=torch.bfloat16, device=device
     )
-    ws_3d = workspace.view(num_tokens, K_total, OUTPUT_DIM)
 
-    # Dequant+gather topk slots from compressed cache
+    if topk_idx_2d is not None and kv_cache is not None and n_topk_rows > 0:
-    if not swa_only and topk_idx_2d is not None and kv_cache is not None:
+        dequant_gather_slots(
-        topk_flat = topk_idx_2d.reshape(-1).to(torch.int32)
+            workspace[:n_topk_rows],
-        topk_buf = torch.empty(
+            kv_cache,
-            (num_tokens * max_topk, OUTPUT_DIM), dtype=torch.bfloat16, device=device
+            topk_idx_2d.reshape(-1).to(torch.int32),
+            kv_cache.shape[1],
         )
-        compressed_block_size = kv_cache.shape[1]
-        dequant_gather_slots(topk_buf, kv_cache, topk_flat, compressed_block_size)
-        ws_3d[:, :max_topk, :] = topk_buf.view(num_tokens, max_topk, OUTPUT_DIM)
 
-    # Dequant+gather SWA slots
+    dequant_gather_slots(
-    swa_flat = swa_idx_2d.reshape(-1).to(torch.int32)
+        workspace[n_topk_rows:],
-    swa_buf = torch.empty(
+        swa_kv_cache,
-        (num_tokens * max_swa, OUTPUT_DIM), dtype=torch.bfloat16, device=device
+        swa_idx_2d.reshape(-1).to(torch.int32),
+        swa_kv_cache.shape[1],
     )
-    swa_block_size = swa_kv_cache.shape[1]
-    dequant_gather_slots(swa_buf, swa_kv_cache, swa_flat, swa_block_size)
 
-    ws_3d[:, max_topk:, :] = swa_buf.view(num_tokens, max_swa, OUTPUT_DIM)
+    # Use a shape-derived width so replay accepts changing lens without host sync.
-
-    # Build combined indices into the flat workspace and combined lengths.
-    # Workspace layout per token t: [topk_0..topk_{max_topk-1}, swa_0..swa_{max_swa-1}]
-    # Flat index for token t, position p = t * K_total + p
-    #
-    # IMPORTANT: The attention kernel uses combined_lens as a position cutoff —
-    # it only reads indices[0..combined_lens-1]. So indices must be PACKED
-    # contiguously: [valid_topk_indices..., valid_swa_indices..., -1 padding...]
-    if not swa_only and topk_lens is not None:
-        combined_lens = (topk_lens + swa_lens).to(torch.int32)
-    else:
-        combined_lens = swa_lens.to(torch.int32)
-
-    max_combined = int(combined_lens.max().item()) if combined_lens.numel() > 0 else 0
-    # Round up to BLOCK_N=16 alignment for kernel efficiency
     _BLOCK_N = 16
-    max_combined_padded = ((max_combined + _BLOCK_N - 1) // _BLOCK_N) * _BLOCK_N
+    width = max(_BLOCK_N, ((K_total + _BLOCK_N - 1) // _BLOCK_N) * _BLOCK_N)
 
-    # Build packed index table: [num_tokens, max_combined_padded]
+    pos = torch.arange(width, device=device, dtype=torch.int32).unsqueeze(0)
-    # Each token t: [topk_0..topk_{tlen-1}, swa_0..swa_{slen-1}, -1 padding]
+    tok = torch.arange(num_tokens, device=device, dtype=torch.int32).unsqueeze(1)
-    # Vectorized: for each token, topk indices are t*K_total + 0..tlen-1,
+
-    #             swa indices are t*K_total + max_topk + 0..slen-1
+    slen = swa_lens.reshape(-1, 1).to(torch.int32).clamp(min=0, max=max_swa)
-    combined_indices = torch.full(
+    if use_topk and topk_lens is not None:
-        (num_tokens, max_combined_padded),
+        tlen = topk_lens.reshape(-1, 1).to(torch.int32).clamp(min=0, max=max_topk)
-        fill_value=-1,
+    else:
-        dtype=torch.int32,
+        tlen = torch.zeros_like(slen)
-        device=device,
+
+    in_topk = pos < tlen
+    in_swa = (pos >= tlen) & (pos < tlen + slen)
+
+    topk_ws = tok * max_topk + pos
+    swa_ws = n_topk_rows + tok * max_swa + (pos - tlen)
+
+    combined_indices = torch.where(
+        in_topk,
+        topk_ws,
+        torch.where(in_swa, swa_ws, torch.full_like(topk_ws, -1)),
     )
 
-    token_offsets = (
-        torch.arange(num_tokens, device=device, dtype=torch.int32) * K_total
-    )  # [B]
-
-    if not swa_only and topk_lens is not None:
-        # Pack topk: for each token, write t*K_total + 0..tlen-1 at positions 0..tlen-1
-        max_tlen = int(topk_lens.max().item())
-        topk_range = torch.arange(max_tlen, device=device, dtype=torch.int32).unsqueeze(
-            0
-        )
-        topk_valid = topk_range < topk_lens.unsqueeze(1)
-        topk_ws_indices = token_offsets.unsqueeze(1) + topk_range
-        combined_indices[:, :max_tlen] = torch.where(
-            topk_valid,
-            topk_ws_indices,
-            torch.tensor(-1, dtype=torch.int32, device=device),
-        )
-        # Pack swa after topk: positions tlen..tlen+slen-1
-        # Since tlen varies per token, we need per-token offset
-        swa_range = torch.arange(max_swa, device=device, dtype=torch.int32).unsqueeze(0)
-        swa_valid = swa_range < swa_lens.unsqueeze(1)
-        swa_ws_indices = token_offsets.unsqueeze(1) + max_topk + swa_range
-        # Write at position topk_lens[t] + swa_pos for each token
-        for t_idx in range(num_tokens):
-            tlen = int(topk_lens[t_idx].item())
-            slen = int(swa_lens[t_idx].item())
-            combined_indices[t_idx, tlen : tlen + slen] = swa_ws_indices[t_idx, :slen]
-    else:
-        # SWA-only: pack swa indices at positions 0..slen-1
-        # Use min(max_swa, max_combined_padded) because combined_indices only
-        # has max_combined_padded columns, and all valid entries fit within it.
-        effective_swa = min(max_swa, max_combined_padded)
-        swa_range = torch.arange(
-            effective_swa, device=device, dtype=torch.int32
-        ).unsqueeze(0)
-        swa_valid = swa_range < swa_lens.unsqueeze(1)
-        swa_ws_indices = token_offsets.unsqueeze(1) + swa_range  # max_topk=0
-        combined_indices[:, :effective_swa] = torch.where(
-            swa_valid,
-            swa_ws_indices,
-            torch.tensor(-1, dtype=torch.int32, device=device),
-        )
-
-    # Call BF16 sparse MLA kernel
     out_attn, _, _ = triton_bf16_mla_sparse_interface(
         q=q,
         kv=workspace.unsqueeze(1),

vllm/models/deepseek_v41/compressor.py

@@ -75,8 +75,12 @@ def _ring_slot_mapping_kernel(
     block = tl.load(block_table_ptr + req * block_table_stride, mask=valid, other=0)
     pos = tl.load(positions_ptr + offsets, mask=valid, other=0)
     slot = block.to(tl.int64) * CAPACITY + pos % CAPACITY
+    # Block 0 is the null block. A request without a ring block (dummy or padding
+    # rows, which the runners fill with the null block) must not write.
     tl.store(
-        slot_mapping_ptr + offsets, tl.where(valid, slot, -1), mask=offsets < num_tokens
+        slot_mapping_ptr + offsets,
+        tl.where(valid & (block != 0), slot, -1),
+        mask=offsets < num_tokens,
     )
 
 

vllm/models/hy_v4/nvidia/attention.py

@@ -47,7 +47,13 @@ from vllm.v1.attention.selector import get_attn_backend
 
 logger = init_logger(__name__)
 
-_SPARSE_LAYER_TYPES = ("sparse_attention", "sparse", "deepseek_sparse_attention")
+_SPARSE_LAYER_TYPES = (
+    "sparse_attention",
+    "sparse",
+    "indexed_attention",
+    # TODO: Delete below once Transformers 5.18.0 is the minimum required version.
+    "deepseek_sparse_attention",
+)
 _WEIGHT_LAYER_INDEX_RE = re.compile(r"(?:^|\.)layers\.(\d+)(?:\.|$)")
 
 

vllm/models/qwen4_exp/amd/model.py

@@ -80,6 +80,10 @@ from .low_latency_gemm import enable_qwen4_exp_low_latency_gemm
 from .ple_layer import Qwen4ExpPLELayer
 from .qsa import Qwen4ExpQSAAttention
 
+# Transformers v5.18 renamed `qwen_sparse_attention` to `indexed_attention`
+# TODO: Delete qwen_... once Transformers 5.18.0 is the minimum required version.
+_QSA_LAYER_TYPES = ("qwen_sparse_attention", "indexed_attention")
+
 
 def without_modelopt_fp4(
     quant_config: QuantizationConfig | None,
@@ -208,7 +212,7 @@ class Qwen4ExpDecoderLayer(nn.Module):
                 prefix=f"{prefix}.linear_attn",
                 gqa_interleaved_layout=False,
             )
-        elif layer_type == "qwen_sparse_attention":
+        elif layer_type in _QSA_LAYER_TYPES:
             use_qsa = getattr(config, "indexer_n_heads", None) is not None
             if not use_qsa:
                 self.self_attn = Qwen3NextAttention(
@@ -290,7 +294,7 @@ class Qwen4ExpDecoderLayer(nn.Module):
 
         if self.layer_type == "linear_attention":
             attn_out = self.linear_attn(hidden_states=block_input)
-        elif self.layer_type == "qwen_sparse_attention":
+        elif self.layer_type in _QSA_LAYER_TYPES:
             attn_out = self.self_attn(
                 hidden_states=block_input,
                 positions=positions,
@@ -381,7 +385,7 @@ class Qwen4ExpModel(nn.Module):
         self._qsa_layer_ids = frozenset(
             layer_idx
             for layer_idx, layer_type in enumerate(config.layer_types)
-            if layer_type == "qwen_sparse_attention"
+            if layer_type in _QSA_LAYER_TYPES
             and getattr(config, "indexer_n_heads", None) is not None
         )
         self.embed_tokens = VocabParallelEmbedding(self.vocab_size, config.hidden_size)

vllm/models/qwen4_exp/common/ngram_embedding.py

@@ -27,6 +27,7 @@ from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tenso
     CompressedTensorsConfig,
 )
 from vllm.model_executor.layers.quantization.fp8 import Fp8Config
+from vllm.model_executor.layers.quantization.inc import INCConfig
 from vllm.model_executor.layers.quantization.modelopt import (
     ModelOptMixedPrecisionConfig,
     ModelOptQuantConfigBase,
@@ -190,6 +191,11 @@ class Qwen4ExpPLEEmbeddingMethod(QuantizeMethodBase):
             and quant_config.get_scheme_dict(None, layer_name=prefix) is None
         ):
             return Qwen4ExpPLEUnquantizedEmbeddingMethod()
+        if (
+            isinstance(quant_config, INCConfig)
+            and not quant_config.config_parser.resolve(None, prefix).quantized
+        ):
+            return Qwen4ExpPLEUnquantizedEmbeddingMethod()
         if not isinstance(quant_config, Fp8Config):
             raise NotImplementedError(
                 "Qwen4Exp PLE embedding does not support quantization config "

vllm/models/qwen4_exp/nvidia/model.py

@@ -82,6 +82,10 @@ from .ops.cute_dsl.hc_down_silu import request_hc_down_silu_warmup
 from .ple_layer import Qwen4ExpPLELayer
 from .qsa import Qwen4ExpQSAAttention
 
+# Transformers v5.18 renamed `qwen_sparse_attention` to `indexed_attention`
+# TODO: Delete qwen_... once Transformers 5.18.0 is the minimum required version.
+_QSA_LAYER_TYPES = ("qwen_sparse_attention", "indexed_attention")
+
 
 def without_modelopt_fp4(
     quant_config: QuantizationConfig | None,
@@ -212,7 +216,7 @@ class Qwen4ExpDecoderLayer(nn.Module):
                 prefix=f"{prefix}.linear_attn",
                 gqa_interleaved_layout=False,
             )
-        elif layer_type == "qwen_sparse_attention":
+        elif layer_type in _QSA_LAYER_TYPES:
             use_qsa = getattr(config, "indexer_n_heads", None) is not None
             if not use_qsa:
                 self.self_attn = Qwen3NextAttention(
@@ -296,7 +300,7 @@ class Qwen4ExpDecoderLayer(nn.Module):
 
         if self.layer_type == "linear_attention":
             attn_out = self.linear_attn(hidden_states=block_input)
-        elif self.layer_type == "qwen_sparse_attention":
+        elif self.layer_type in _QSA_LAYER_TYPES:
             attn_out = self.self_attn(
                 hidden_states=block_input,
                 positions=positions,
@@ -376,7 +380,7 @@ class Qwen4ExpModel(nn.Module):
         self._qsa_layer_ids = frozenset(
             layer_idx
             for layer_idx, layer_type in enumerate(config.layer_types)
-            if layer_type == "qwen_sparse_attention"
+            if layer_type in _QSA_LAYER_TYPES
             and getattr(config, "indexer_n_heads", None) is not None
         )
         self.embed_tokens = VocabParallelEmbedding(self.vocab_size, config.hidden_size)

vllm/multimodal/audio.py

@@ -146,7 +146,7 @@ def normalize_audio(
 
     # No reduction needed if already at target
     if num_channels == spec.target_channels:
-        return audio
+        return audio[0] if num_channels == 1 else audio
 
     # Cannot expand channels
     if num_channels < spec.target_channels:

vllm/platforms/cpu.py

@@ -523,7 +523,9 @@ class CpuPlatform(Platform):
         if not backend_classes:
             return
 
+        pre_block_size = vllm_config.cache_config.block_size
         cls._align_hybrid_block_size(vllm_config, backend_classes[0])
+        cls._check_aligned_block_size(vllm_config, backend_classes, pre_block_size)
 
     @classmethod
     def discover_numa_topology(cls) -> list[list[int]]:

vllm/platforms/interface.py

@@ -698,6 +698,7 @@ class Platform:
 
         # Phase 2: Align block/mamba sizes for hybrid models
         # (may override user settings).
+        pre_block_size = cache_config.block_size
         if model_config.is_hybrid:
             cls._align_hybrid_block_size(vllm_config, backend_classes[0])
 
@@ -706,6 +707,36 @@ class Platform:
         # May override the user's --block-size.
         if cache_config.kv_cache_dtype_skip_layers:
             cls._align_heterogeneous_kv_block_size(vllm_config, backend_classes[0])
+        cls._check_aligned_block_size(vllm_config, backend_classes, pre_block_size)
+
+    @classmethod
+    def _check_aligned_block_size(
+        cls,
+        vllm_config: "VllmConfig",
+        backend_classes: "list[type[AttentionBackend]]",
+        pre_block_size: int,
+    ) -> None:
+        """Fail now, not at worker setup, if alignment left a rejected size.
+
+        Alignment rounds by the first backend's kernel block only, so it can
+        land on a size a sibling backend rejects.
+        """
+        from vllm.config.vllm import set_current_vllm_config
+
+        block_size = vllm_config.cache_config.block_size
+        if block_size == pre_block_size:
+            return
+        with set_current_vllm_config(vllm_config):
+            rejecting = [
+                b.get_name()
+                for b in backend_classes
+                if not b.supports_block_size(block_size)
+            ]
+        if rejecting:
+            raise ValueError(
+                f"Aligned KV cache block size {block_size} is not supported by "
+                f"attention backend(s): {', '.join(rejecting)}."
+            )
 
     @classmethod
     def _align_heterogeneous_kv_block_size(

vllm/platforms/xpu.py

@@ -452,6 +452,7 @@ class XPUPlatform(Platform):
         if new_block_size == cache_config.block_size:
             return
 
+        pre_block_size = cache_config.block_size
         if cache_config.mamba_cache_mode == "align":
             cache_config.mamba_block_size = new_block_size
         original_mamba_page_size_padded = cache_config.mamba_page_size_padded
@@ -474,6 +475,10 @@ class XPUPlatform(Platform):
                 original_mamba_page_size_padded,
                 cache_config.mamba_page_size_padded,
             )
+        # This rounding runs after super()'s check, so check its result too.
+        cls._check_aligned_block_size(
+            vllm_config, cls._find_non_ssm_backends(vllm_config), pre_block_size
+        )
 
     @classmethod
     def support_hybrid_kv_cache(cls) -> bool:

vllm/transformers_utils/config.py

@@ -147,6 +147,12 @@ _CONFIG_REGISTRY: dict[str, type[PreTrainedConfig]] = LazyConfigDict(
 
 _SPECULATIVE_DECODING_CONFIGS: set[str] = {"eagle", "speculators", "medusa"}
 
+# Checkpoints whose config.json has no model_type, only remote code in auto_map
+_ARCHITECTURE_TO_MODEL_TYPE: dict[str, str] = {
+    "LongcatFlashForCausalLM": "longcat_flash",
+    "LongcatFlashNgramForCausalLM": "longcat_flash",
+}
+
 _PATCH_HF_VALIDATE_ROPE: set[str] = {"sarvam_mla"}
 
 # Model types whose checkpoints carry shared RoPE parameters alongside the
@@ -352,6 +358,8 @@ class HFConfigParser(ConfigParserBase):
                 if config_dict.get("speculators_config") is not None
                 else model_type
             )
+        if model_type is None and (architectures := config_dict.get("architectures")):
+            model_type = _ARCHITECTURE_TO_MODEL_TYPE.get(architectures[0])
         # Allow hf_overrides to override model_type before checking _CONFIG_REGISTRY
         if (hf_overrides := kwargs.pop("hf_overrides", None)) is not None:
             if isinstance(hf_overrides, dict) and "model_type" in hf_overrides:
@@ -374,6 +382,15 @@ class HFConfigParser(ConfigParserBase):
         if extra_layer_types := _PATCH_HF_ALLOWED_LAYER_TYPES.get(model_type):
             _patch_hf_transformers_allowed_layer_types(extra_layer_types)
 
+        if model_type == "vlm":
+            # HyperCLOVAX remote code registers this alias in a bare
+            # `try/except` that fails silently since transformers 5.18
+            from transformers import CONFIG_MAPPING, Qwen2_5_VLVisionConfig
+
+            CONFIG_MAPPING.register(
+                "qwen2_5_vl_visual", Qwen2_5_VLVisionConfig, exist_ok=True
+            )
+
         rope_parameters = config_dict.get("rope_parameters") or {}
         if model_type == "gemma4_text" and "full_attention" in rope_parameters:
             from transformers import Gemma4TextConfig
@@ -387,6 +404,11 @@ class HFConfigParser(ConfigParserBase):
             }
             kwargs.setdefault("name_or_path", str(model))
             config = Gemma4TextConfig.from_dict(config_dict, **kwargs)
+        elif model_type in _ARCHITECTURE_TO_MODEL_TYPE.values():
+            from transformers import CONFIG_MAPPING
+
+            kwargs.setdefault("name_or_path", str(model))
+            config = CONFIG_MAPPING[model_type].from_dict(config_dict, **kwargs)
         elif model_type in _SPECULATIVE_DECODING_CONFIGS:
             config_class = _CONFIG_REGISTRY[model_type]
             config = config_class.from_pretrained(

vllm/transformers_utils/model_arch_config_convertor.py

@@ -329,7 +329,6 @@ class ModelArchConfigConvertorBase:
             "kimi_k2",
             "kimi_linear",
             "longcat_flash",
-            "longcat_flash_ngram",
             "pangu_ultra_moe",
             "pangu_ultra_moe_mtp",
             "bailing_hybrid",
@@ -714,6 +713,12 @@ class PanguUltraMoeMTPModelArchConfigConvertor(ModelArchConfigConvertorBase):
         return getattr(self.hf_text_config, "num_nextn_predict_layers", 0)
 
 
+class LongCatFlashModelArchConfigConvertor(ModelArchConfigConvertorBase):
+    def get_num_hidden_layers(self) -> int:
+        # `num_hidden_layers` counts attention sublayers, two per decoder layer
+        return self.hf_text_config.num_layers
+
+
 class LongCatFlashMTPModelArchConfigConvertor(ModelArchConfigConvertorBase):
     def get_num_hidden_layers(self) -> int:
         return getattr(self.hf_text_config, "num_nextn_predict_layers", 1)
@@ -825,6 +830,7 @@ MODEL_ARCH_CONFIG_CONVERTORS = {
     "gemma4_unified_text": Gemma4ModelArchConfigConvertor,
     "glm4_moe_mtp": GLM4MoeMTPModelArchConfigConvertor,
     "glm_ocr_mtp": GLM4MoeMTPModelArchConfigConvertor,
+    "longcat_flash": LongCatFlashModelArchConfigConvertor,
     "longcat_flash_mtp": LongCatFlashMTPModelArchConfigConvertor,
     "mamba": MambaModelArchConfigConvertor,
     "medusa": MedusaModelArchConfigConvertor,

vllm/v1/attention/backends/mla/tokenspeed_mla.py

@@ -6,6 +6,7 @@ from typing import TYPE_CHECKING, ClassVar
 
 import torch
 
+from vllm.config import get_current_vllm_config
 from vllm.config.cache import CacheDType
 from vllm.logger import init_logger
 from vllm.model_executor.layers.attention.mla_attention import (
@@ -155,6 +156,7 @@ class TokenspeedMLABackend(MLACommonBackend):
 class TokenspeedMLAImpl(MLACommonImpl[MLACommonMetadata]):
     can_return_lse_for_decode: bool = True
     supports_dcp: bool = True
+    supports_mtp_with_cp_non_trivial_interleave_size: bool = True
     # tokenspeed_mla_decode returns LSE in log2 units; its own DCP test merges
     # partial outputs with exp2(lse).
     lse_base_on_e: bool = False
@@ -216,6 +218,9 @@ class TokenspeedMLAImpl(MLACommonImpl[MLACommonMetadata]):
         self._workspace_buffer: torch.Tensor | None = None
         self.softmax_scale: float | None = None
         self.output_scale: float | None = None
+        # NIXL resolves interleaving after model construction; retain the config
+        # rather than caching its initial interleave size.
+        self._parallel_config = get_current_vllm_config().parallel_config
 
         # Pre-JIT BF16 and FP8 prefill kernels here too — decode impl always
         # runs when tokenspeed is selected, prefill backend may not (user can
@@ -311,9 +316,13 @@ class TokenspeedMLAImpl(MLACommonImpl[MLACommonMetadata]):
             causal_seqs=causal_seqs if self.dcp_world_size > 1 else None,
             cp_world=self.dcp_world_size,
             cp_rank=self.dcp_rank,
+            cp_interleave_size=self._parallel_config.cp_kv_cache_interleave_size,
         )
         if return_lse:
             o, lse = kernel_out
+            # DCP ranks with no local KV may return undefined outputs (e.g.,
+            # NaNs). The downstream DCP combine path must mask empty
+            # shards and ignore their outputs when merging partial attention.
             lse = lse.view(-1, lse.shape[-1])
         else:
             o, lse = kernel_out, None

vllm/v1/attention/backends/rocm_attn.py

@@ -481,7 +481,9 @@ class RocmAttentionImpl(AttentionImpl):
             k_scale=layer._k_scale,
             v_scale=layer._v_scale,
             alibi_slopes=self.alibi_slopes,
-            sliding_window=self.sliding_window[0],
+            # self.sliding_window[0] is the FlashAttention-style left span (W - 1).
+            # chunked_prefill_paged_decode expects the full window length W.
+            sliding_window=1 + self.sliding_window[0],
             sm_scale=self.scale,
             output_scale=output_scale,
             sinks=self.sinks,

vllm/v1/simple_kv_offload/manager.py

@@ -1395,6 +1395,10 @@ class SimpleCPUOffloadScheduler:
         the transfer finished, then release refs without caching abandoned
         store results.
         """
+        for pending in self._pending_cpu_hits.values():
+            self._free_pending_cpu_hit(pending)
+        self._pending_cpu_hits.clear()
+
         self._abandoned_store_event_to_blocks.update(self._store_event_to_blocks)
         for transfer in self._pending_finished_stores:
             self._release_transfer_refs(transfer)

vllm/v1/worker/gpu/attn_utils.py

@@ -398,10 +398,7 @@ def init_kv_cache(
     # Dual-attention models (e.g. LongCat-Flash) put two Attention modules per
     # decoder layer, so a layer name carries two integers (layer + module index).
     num_attn_module = (
-        2
+        2 if vllm_config.model_config.hf_config.model_type == "longcat_flash" else 1
-        if vllm_config.model_config.hf_config.model_type
-        in ("longcat_flash", "longcat_flash_ngram")
-        else 1
     )
     bindable_caches = {
         name: cache for name, cache in kv_caches.items() if name in forward_context

vllm/v1/worker/gpu/cudagraph_utils.py

@@ -19,10 +19,13 @@ from vllm.compilation.breakable_cudagraph import (
 )
 from vllm.compilation.counter import compilation_counter
 from vllm.compilation.cuda_graph import CUDAGraphStat, CUDAGraphWrapper
+from vllm.compilation.cudagraph_pool import (
+    capture_outside_cumem_pool,
+    capture_pool,
+)
 from vllm.compilation.wrapper import TorchCompileWithNoGuardsWrapper
 from vllm.config import VllmConfig, set_current_vllm_config
 from vllm.config.compilation import CUDAGraphMode
-from vllm.distributed.device_communicators.pynccl_allocator import set_graph_pool_id
 from vllm.distributed.parallel_state import (
     get_pp_group,
     graph_capture,
@@ -481,15 +484,14 @@ class CudaGraphManager:
                         # Sync offloader's copy stream before capture.
                         # Ensure any pre-capture prefetches from offloader are complete.
                         get_offloader().sync_prev_onload()
-                        if self.pool is not None:
-                            set_graph_pool_id(self.pool)
-                        else:
-                            set_graph_pool_id(current_platform.graph_pool_handle())
                         if self._capture_mem_samples is not None:
                             torch.accelerator.synchronize()
                             free_before = torch.accelerator.get_memory_info()[0]
-                        with torch.cuda.graph(
+                        with (
-                            graph, self.pool, stream=self._capture_stream(desc)
+                            capture_pool(self.pool, self.vllm_config) as pool,
+                            torch.cuda.graph(
+                                graph, pool, stream=self._capture_stream(desc)
+                            ),
                         ):
                             forward_fn(CUDAGraphMode.NONE)
                             # Join offloader's copy stream after forward to avoid
@@ -941,7 +943,8 @@ def profile_cudagraph_memory(runner: "GPUModelRunner") -> int:
             mem_samples: list[int] = []
             manager._capture_mem_samples = mem_samples
 
-            measured = int(runner.capture_model(profile_only=True))
+            with capture_outside_cumem_pool():
+                measured = int(runner.capture_model(profile_only=True))
 
             # The measured delta covers PIECEWISE, encoder and speculator graphs
             # plus the sampled FULL graphs; swap the sampled FULL cost for the