Merge remote-tracking branch 'origin/main' into sync/lesj/qwen4exp-bugfixes

Signed-off-by: lesj0610 <lesj0610@users.noreply.github.com>

.buildkite/intel_jobs/test-intel.yaml

@@ -98,7 +98,7 @@ steps:
       - >-
         bash .buildkite/scripts/hardware_ci/run-intel-test.sh
         'cd tests && pip install dpctl &&
-        pytest -v -s v1/core --ignore=v1/core/test_reset_prefix_cache_e2e.py --ignore=v1/core/test_scheduler_e2e.py --ignore=v1/core/prefix_cache/test_hybrid_prefix_cache_hit_rate.py &&
+        pytest -v -s v1/core --ignore=v1/core/test_reset_prefix_cache_e2e.py --ignore=v1/core/test_scheduler_e2e.py &&
         pytest -v -s v1/engine --ignore=v1/engine/test_output_processor.py &&
         pytest -v -s v1/sample --ignore=v1/sample/test_logprobs.py --ignore=v1/sample/test_logprobs_e2e.py -k "not test_topk_only and not test_topp_only and not test_topk_and_topp" &&
         pytest -v -s v1/worker --ignore=v1/worker/test_gpu_model_runner.py --ignore=v1/worker/test_worker_memory_snapshot.py &&

csrc/cpu/cpu_attn_impl.hpp

@@ -956,6 +956,28 @@ void print_logits(const char* name, T* ptr, int32_t row, int32_t col,
   std::printf("%s", ss.str().c_str());
 }
 
+// This allows us to store probabilities in a packed format
+// for GEMM implementations that need packed A.
+template <typename attention_impl_t, typename = void>
+struct ProbabilityTokenStore {
+  using prob_buffer_t = typename attention_impl_t::prob_buffer_t;
+  static constexpr int32_t TokenStride = 1;
+
+  FORCE_INLINE static void store_probabilities(
+      prob_buffer_t* __restrict__ probability, const vec_op::FP32Vec16& values,
+      const int32_t, const int64_t) {
+    using prob_buffer_vec_t = typename VecTypeTrait<prob_buffer_t>::vec_t;
+    prob_buffer_vec_t output_vec(values);
+    output_vec.save(probability);
+  }
+};
+
+template <typename attention_impl_t>
+struct ProbabilityTokenStore<
+    attention_impl_t,
+    std::void_t<typename attention_impl_t::ProbabilityTokenStore>>
+    : attention_impl_t::ProbabilityTokenStore {};
+
 template <typename attention_impl_t>
 class AttentionMainLoop {
  public:
@@ -966,6 +988,7 @@ class AttentionMainLoop {
   using partial_output_buffer_t =
       typename attention_impl_t::partial_output_buffer_t;
   using prob_buffer_t = typename attention_impl_t::prob_buffer_t;
+  using probability_token_store_t = ProbabilityTokenStore<attention_impl_t>;
 
   static constexpr int64_t max_q_head_num_per_iter =
       attention_impl_t::MaxQHeadNumPerIteration;
@@ -974,6 +997,8 @@ class AttentionMainLoop {
   static constexpr int64_t headdim_alignment =
       attention_impl_t::HeadDimAlignment;
   static constexpr int64_t head_dim = attention_impl_t::HeadDim;
+  static constexpr int32_t probability_token_stride =
+      probability_token_store_t::TokenStride;
   static constexpr ISA ISAType = attention_impl_t::ISAType;
   static constexpr bool scale_on_logits =
       attention_impl_t::scale_on_logits;  // apply scale on logits, otherwise
@@ -1152,8 +1177,6 @@ class AttentionMainLoop {
                 ? kv_tile_token_num
                 : kv_tile_token_num *
                       (sizeof(logits_buffer_t) / sizeof(prob_buffer_t));
-        constexpr int64_t prob_buffer_elem_size =
-            prequantize_probabilities ? sizeof(uint8_t) : sizeof(prob_buffer_t);
         partial_output_buffer_t* curr_partial_q_buffer = partial_q_buffer;
         bool accum_c = !is_first_iter;
         for (int32_t block_idx = start_block_idx; block_idx < end_block_idx;
@@ -1187,9 +1210,7 @@ class AttentionMainLoop {
           remaining_group_num -= curr_group_num_in_block;
           curr_group_offset = 0;
           curr_group_num_in_block = token_group_num_per_block;
-          curr_prob_buffer = reinterpret_cast<pv_prob_buffer_t*>(
-              reinterpret_cast<uint8_t*>(curr_prob_buffer) +
-              curr_token_num * prob_buffer_elem_size);
+          curr_prob_buffer += curr_token_num * probability_token_stride;
           curr_partial_q_buffer = partial_q_buffer;
           accum_c = true;
         }
@@ -1262,7 +1283,6 @@ class AttentionMainLoop {
           std::is_same_v<query_t, c10::Half>;
 #endif
 
-      using prob_buffer_vec_t = typename VecTypeTrait<prob_buffer_t>::vec_t;
       constexpr bool prequantize_probabilities = []() {
         if constexpr (requires { tile_gemm_t::prequantize_probabilities; }) {
           return tile_gemm_t::prequantize_probabilities;
@@ -1331,12 +1351,6 @@ class AttentionMainLoop {
               vec = fast_exp(vec);
             }
 
-            if constexpr (prequantize_probabilities) {
-              vec.save(curr_logits_buffer_iter);
-            } else {
-              prob_buffer_vec_t output_vec(vec);
-              output_vec.save(curr_prob_buffer_iter);
-            }
 #else
             vec.save(curr_logits_buffer_iter);
             for (int32_t k = 0; k < 16; ++k) {
@@ -1345,10 +1359,19 @@ class AttentionMainLoop {
             vec = vec_op::FP32Vec16(curr_logits_buffer_iter);
 #endif
 
+            if constexpr (prequantize_probabilities) {
+              vec.save(curr_logits_buffer_iter);
+            } else {
+              probability_token_store_t::store_probabilities(
+                  curr_prob_buffer_iter, vec, i,
+                  logits_buffer_stride * sizeof(logits_buffer_t) /
+                      sizeof(prob_buffer_t));
+            }
+
             sum_vec = sum_vec + vec;
 
             curr_logits_buffer_iter += 16;
-            curr_prob_buffer_iter += 16;
+            curr_prob_buffer_iter += 16 * probability_token_stride;
           }
           if constexpr (prequantize_probabilities) {
             tile_gemm_t::quantize_probability_row(

csrc/cpu/cpu_attn_neon.hpp

@@ -11,7 +11,7 @@ namespace cpu_attention {
 
 namespace {
 
-#define BLOCK_SIZE_ALIGNMENT 32
+#define BLOCK_SIZE_ALIGNMENT 16
 #define HEAD_SIZE_ALIGNMENT 32
 #define MAX_Q_HEAD_NUM_PER_ITER 16
 
@@ -387,7 +387,7 @@ class AttentionImpl<ISA::NEON, scalar_t, head_dim, kv_cache_scalar_t> {
 };
 
 #ifdef ARM_BF16_SUPPORT
-// For BF16 on Arm, reuse the BFMMLA kernels with 32-token alignment.
+// For BF16 on Arm, reuse the BFMMLA kernels with the same token alignment.
 template <int64_t head_dim>
 class AttentionImpl<ISA::NEON, c10::BFloat16, head_dim, c10::BFloat16>
     : public AttentionImplNEONBFMMLA<BLOCK_SIZE_ALIGNMENT, ISA::NEON,

csrc/cpu/cpu_attn_neon_bfmmla.hpp

@@ -5,432 +5,275 @@
 
 #include "cpu_attn_impl.hpp"
 
+#include <arm_bf16.h>
 #include <arm_neon.h>
+#include <c10/util/BFloat16.h>
 
+#include <algorithm>
 #include <cstdint>
-#include <vector>
 
 namespace cpu_attention {
 
+class BfmmlaGemm {
+ public:
+  static constexpr int32_t KTile = 4;
+  static constexpr int32_t NTile = 8;
+  static constexpr int32_t MaxRows = 8;
+
+  FORCE_INLINE static void gemm(const c10::BFloat16* __restrict__ a,
+                                const c10::BFloat16* __restrict__ b,
+                                float* __restrict__ c, const int32_t m,
+                                const int32_t n, const int32_t k,
+                                const int64_t a_pair_stride,
+                                const int64_t b_n_group_stride,
+                                const int64_t b_k_group_stride,
+                                const int64_t ldc, const bool accumulate) {
+    const auto* a_ptr = reinterpret_cast<const bfloat16_t*>(a);
+    const auto* b_ptr = reinterpret_cast<const bfloat16_t*>(b);
+
+    for (int32_t n_idx = 0; n_idx < n; n_idx += 16) {
+      const auto* b_panel = b_ptr + (n_idx / NTile) * b_n_group_stride;
+      float* c_panel = c + n_idx;
+
+      // Preserve this range so the inactive row pair is optimized away.
+      if (m <= 2) {
+        gemm_4x16(a_ptr, b_panel, c_panel, m, k, a_pair_stride,
+                  b_n_group_stride, b_k_group_stride, ldc, accumulate);
+      } else if (m <= 4) {
+        gemm_4x16(a_ptr, b_panel, c_panel, m, k, a_pair_stride,
+                  b_n_group_stride, b_k_group_stride, ldc, accumulate);
+      } else {
+        gemm_8x8(a_ptr, b_panel, c_panel, m, k, a_pair_stride, b_k_group_stride,
+                 ldc, accumulate);
+        gemm_8x8(a_ptr, b_panel + b_n_group_stride, c_panel + NTile, m, k,
+                 a_pair_stride, b_k_group_stride, ldc, accumulate);
+      }
+    }
+  }
+
+ private:
+  FORCE_INLINE static float32x4_t zip_low_pairs(const float32x4_t a,
+                                                const float32x4_t b) {
+    return vreinterpretq_f32_f64(
+        vzip1q_f64(vreinterpretq_f64_f32(a), vreinterpretq_f64_f32(b)));
+  }
+
+  FORCE_INLINE static float32x4_t zip_high_pairs(const float32x4_t a,
+                                                 const float32x4_t b) {
+    return vreinterpretq_f32_f64(
+        vzip2q_f64(vreinterpretq_f64_f32(a), vreinterpretq_f64_f32(b)));
+  }
+
+  FORCE_INLINE static void init_accumulators(
+      float32x4_t& acc01, float32x4_t& acc23, float32x4_t& acc45,
+      float32x4_t& acc67, const float* __restrict__ c, const int64_t ldc,
+      const int32_t rows, const bool accumulate) {
+    if (!accumulate || rows == 0) {
+      acc01 = vdupq_n_f32(0.0f);
+      acc23 = vdupq_n_f32(0.0f);
+      acc45 = vdupq_n_f32(0.0f);
+      acc67 = vdupq_n_f32(0.0f);
+      return;
+    }
+
+    const float32x4_t row0_0123 = vld1q_f32(c);
+    const float32x4_t row0_4567 = vld1q_f32(c + 4);
+    const float32x4_t row1_0123 =
+        (rows == 2) ? vld1q_f32(c + ldc) : vdupq_n_f32(0.0f);
+    const float32x4_t row1_4567 =
+        (rows == 2) ? vld1q_f32(c + ldc + 4) : vdupq_n_f32(0.0f);
+
+    acc01 = zip_low_pairs(row0_0123, row1_0123);
+    acc23 = zip_high_pairs(row0_0123, row1_0123);
+    acc45 = zip_low_pairs(row0_4567, row1_4567);
+    acc67 = zip_high_pairs(row0_4567, row1_4567);
+  }
+
+  FORCE_INLINE static void store_accumulators(
+      const float32x4_t acc01, const float32x4_t acc23, const float32x4_t acc45,
+      const float32x4_t acc67, float* __restrict__ c, const int64_t ldc,
+      const int32_t rows) {
+    if (rows == 0) {
+      return;
+    }
+
+    vst1q_f32(c, zip_low_pairs(acc01, acc23));
+    vst1q_f32(c + 4, zip_low_pairs(acc45, acc67));
+    if (rows == 2) {
+      vst1q_f32(c + ldc, zip_high_pairs(acc01, acc23));
+      vst1q_f32(c + ldc + 4, zip_high_pairs(acc45, acc67));
+    }
+  }
+
+  FORCE_INLINE static bfloat16x8_t load_a_pair(const bfloat16_t* __restrict__ a,
+                                               const int32_t rows) {
+    if (rows == 0) {
+      return vdupq_n_bf16(bfloat16_t{});
+    }
+    // Packed A reserves both rows for an M tail.
+    return vld1q_bf16(a);
+  }
+
+  FORCE_INLINE static void gemm_4x16(const bfloat16_t* __restrict__ a,
+                                     const bfloat16_t* __restrict__ b,
+                                     float* __restrict__ c, const int32_t m,
+                                     const int32_t k,
+                                     const int64_t a_pair_stride,
+                                     const int64_t b_n_group_stride,
+                                     const int64_t b_k_group_stride,
+                                     const int64_t ldc, const bool accumulate) {
+    const int32_t rows01 = std::min(2, std::max(0, m));
+    const int32_t rows23 = std::min(2, std::max(0, m - 2));
+    float32x4_t acc0101, acc0123, acc0145, acc0167;
+    float32x4_t acc2301, acc2323, acc2345, acc2367;
+    float32x4_t acc0189, acc011011, acc011213, acc011415;
+    float32x4_t acc2389, acc231011, acc231213, acc231415;
+    init_accumulators(acc0101, acc0123, acc0145, acc0167, c, ldc, rows01,
+                      accumulate);
+    init_accumulators(acc2301, acc2323, acc2345, acc2367, c + 2 * ldc, ldc,
+                      rows23, accumulate);
+    init_accumulators(acc0189, acc011011, acc011213, acc011415, c + 8, ldc,
+                      rows01, accumulate);
+    init_accumulators(acc2389, acc231011, acc231213, acc231415, c + 2 * ldc + 8,
+                      ldc, rows23, accumulate);
+
+    const bfloat16_t* a01 = a;
+    const bfloat16_t* a23 = a + a_pair_stride;
+    const bfloat16_t* b0 = b;
+    const bfloat16_t* b1 = b + b_n_group_stride;
+
+#pragma GCC unroll 4
+    for (int32_t k_idx = 0; k_idx < k; k_idx += KTile) {
+      const bfloat16x8_t av01 = load_a_pair(a01, rows01);
+      const bfloat16x8_t av23 = load_a_pair(a23, rows23);
+      const bfloat16x8_t b01 = vld1q_bf16(b0);
+      const bfloat16x8_t b23 = vld1q_bf16(b0 + NTile);
+      const bfloat16x8_t b45 = vld1q_bf16(b0 + 2 * NTile);
+      const bfloat16x8_t b67 = vld1q_bf16(b0 + 3 * NTile);
+      const bfloat16x8_t b89 = vld1q_bf16(b1);
+      const bfloat16x8_t b1011 = vld1q_bf16(b1 + NTile);
+      const bfloat16x8_t b1213 = vld1q_bf16(b1 + 2 * NTile);
+      const bfloat16x8_t b1415 = vld1q_bf16(b1 + 3 * NTile);
+
+      acc0101 = vbfmmlaq_f32(acc0101, av01, b01);
+      acc2301 = vbfmmlaq_f32(acc2301, av23, b01);
+      acc0123 = vbfmmlaq_f32(acc0123, av01, b23);
+      acc2323 = vbfmmlaq_f32(acc2323, av23, b23);
+      acc0145 = vbfmmlaq_f32(acc0145, av01, b45);
+      acc2345 = vbfmmlaq_f32(acc2345, av23, b45);
+      acc0167 = vbfmmlaq_f32(acc0167, av01, b67);
+      acc2367 = vbfmmlaq_f32(acc2367, av23, b67);
+      acc0189 = vbfmmlaq_f32(acc0189, av01, b89);
+      acc2389 = vbfmmlaq_f32(acc2389, av23, b89);
+      acc011011 = vbfmmlaq_f32(acc011011, av01, b1011);
+      acc231011 = vbfmmlaq_f32(acc231011, av23, b1011);
+      acc011213 = vbfmmlaq_f32(acc011213, av01, b1213);
+      acc231213 = vbfmmlaq_f32(acc231213, av23, b1213);
+      acc011415 = vbfmmlaq_f32(acc011415, av01, b1415);
+      acc231415 = vbfmmlaq_f32(acc231415, av23, b1415);
+
+      a01 += 2 * KTile;
+      a23 += 2 * KTile;
+      b0 += b_k_group_stride;
+      b1 += b_k_group_stride;
+    }
+
+    store_accumulators(acc0101, acc0123, acc0145, acc0167, c, ldc, rows01);
+    store_accumulators(acc2301, acc2323, acc2345, acc2367, c + 2 * ldc, ldc,
+                       rows23);
+    store_accumulators(acc0189, acc011011, acc011213, acc011415, c + 8, ldc,
+                       rows01);
+    store_accumulators(acc2389, acc231011, acc231213, acc231415,
+                       c + 2 * ldc + 8, ldc, rows23);
+  }
+
+  FORCE_INLINE static void gemm_8x8(const bfloat16_t* __restrict__ a,
+                                    const bfloat16_t* __restrict__ b,
+                                    float* __restrict__ c, const int32_t m,
+                                    const int32_t k,
+                                    const int64_t a_pair_stride,
+                                    const int64_t b_k_group_stride,
+                                    const int64_t ldc, const bool accumulate) {
+    const int32_t rows01 = std::min(2, std::max(0, m));
+    const int32_t rows23 = std::min(2, std::max(0, m - 2));
+    const int32_t rows45 = std::min(2, std::max(0, m - 4));
+    const int32_t rows67 = std::min(2, std::max(0, m - 6));
+    float32x4_t acc0101, acc0123, acc0145, acc0167;
+    float32x4_t acc2301, acc2323, acc2345, acc2367;
+    float32x4_t acc4501, acc4523, acc4545, acc4567;
+    float32x4_t acc6701, acc6723, acc6745, acc6767;
+    init_accumulators(acc0101, acc0123, acc0145, acc0167, c, ldc, rows01,
+                      accumulate);
+    init_accumulators(acc2301, acc2323, acc2345, acc2367, c + 2 * ldc, ldc,
+                      rows23, accumulate);
+    init_accumulators(acc4501, acc4523, acc4545, acc4567, c + 4 * ldc, ldc,
+                      rows45, accumulate);
+    init_accumulators(acc6701, acc6723, acc6745, acc6767, c + 6 * ldc, ldc,
+                      rows67, accumulate);
+
+    const bfloat16_t* a01 = a;
+    const bfloat16_t* a23 = a + a_pair_stride;
+    const bfloat16_t* a45 = a + 2 * a_pair_stride;
+    const bfloat16_t* a67 = a + 3 * a_pair_stride;
+    const bfloat16_t* b_ptr = b;
+
+#pragma GCC unroll 4
+    for (int32_t k_idx = 0; k_idx < k; k_idx += KTile) {
+      const bfloat16x8_t av01 = load_a_pair(a01, rows01);
+      const bfloat16x8_t av23 = load_a_pair(a23, rows23);
+      const bfloat16x8_t av45 = load_a_pair(a45, rows45);
+      const bfloat16x8_t av67 = load_a_pair(a67, rows67);
+      const bfloat16x8_t b01 = vld1q_bf16(b_ptr);
+      const bfloat16x8_t b23 = vld1q_bf16(b_ptr + NTile);
+      const bfloat16x8_t b45 = vld1q_bf16(b_ptr + 2 * NTile);
+      const bfloat16x8_t b67 = vld1q_bf16(b_ptr + 3 * NTile);
+
+      acc0101 = vbfmmlaq_f32(acc0101, av01, b01);
+      acc2301 = vbfmmlaq_f32(acc2301, av23, b01);
+      acc4501 = vbfmmlaq_f32(acc4501, av45, b01);
+      acc6701 = vbfmmlaq_f32(acc6701, av67, b01);
+      acc0123 = vbfmmlaq_f32(acc0123, av01, b23);
+      acc2323 = vbfmmlaq_f32(acc2323, av23, b23);
+      acc4523 = vbfmmlaq_f32(acc4523, av45, b23);
+      acc6723 = vbfmmlaq_f32(acc6723, av67, b23);
+      acc0145 = vbfmmlaq_f32(acc0145, av01, b45);
+      acc2345 = vbfmmlaq_f32(acc2345, av23, b45);
+      acc4545 = vbfmmlaq_f32(acc4545, av45, b45);
+      acc6745 = vbfmmlaq_f32(acc6745, av67, b45);
+      acc0167 = vbfmmlaq_f32(acc0167, av01, b67);
+      acc2367 = vbfmmlaq_f32(acc2367, av23, b67);
+      acc4567 = vbfmmlaq_f32(acc4567, av45, b67);
+      acc6767 = vbfmmlaq_f32(acc6767, av67, b67);
+
+      a01 += 2 * KTile;
+      a23 += 2 * KTile;
+      a45 += 2 * KTile;
+      a67 += 2 * KTile;
+      b_ptr += b_k_group_stride;
+    }
+
+    store_accumulators(acc0101, acc0123, acc0145, acc0167, c, ldc, rows01);
+    store_accumulators(acc2301, acc2323, acc2345, acc2367, c + 2 * ldc, ldc,
+                       rows23);
+    store_accumulators(acc4501, acc4523, acc4545, acc4567, c + 4 * ldc, ldc,
+                       rows45);
+    store_accumulators(acc6701, acc6723, acc6745, acc6767, c + 6 * ldc, ldc,
+                       rows67);
+  }
+};
+
 namespace {
 
-// BFMMLA tile dimensions
-constexpr int32_t TILE_ROWS = 2;  // M dimension
-constexpr int32_t TILE_K = 4;     // K reduction
-constexpr int32_t TILE_COLS = 2;  // N dimension (column-pair)
-
-// Derived constants
-constexpr int32_t OUTPUT_COLS_PER_BLOCK = 8;   // 4 column-pairs
-constexpr int32_t K_TOKENS_PER_GROUP = 8;      // Tokens grouped in K cache
-constexpr int32_t V_TOKENS_PER_ROW_BLOCK = 4;  // Tokens per V cache row block
-constexpr int32_t K_INNER_STRIDE = K_TOKENS_PER_GROUP * TILE_K;
-constexpr int32_t V_INNER_STRIDE = V_TOKENS_PER_ROW_BLOCK * TILE_COLS;
-constexpr int32_t PACK_ELEMENTS_PER_K_CHUNK = TILE_ROWS * TILE_K;  // A packing
-
-// Matrix Packing and Accumulator
-// Reshape two rows of Q into BFMMLA-friendly interleaved
-// Input:  row0 = [a0,a1,a2,a3], row1 = [b0,b1,b2,b3]
-// Output: [a0,a1,a2,a3,b0,b1,b2,b3, a4,a5,a6,a7,b4,b5,b6,b7]
-// For K tail (K % TILE_K != 0): pads with zeros to complete the final chunk
-FORCE_INLINE void reshape_Q_2xK_for_bfmmla(const c10::BFloat16* __restrict r0,
-                                           const c10::BFloat16* __restrict r1,
-                                           c10::BFloat16* __restrict dst,
-                                           int32_t K) {
-  const uint16_t* s0 = reinterpret_cast<const uint16_t*>(r0);
-  const uint16_t* s1 = reinterpret_cast<const uint16_t*>(r1);
-  uint16_t* d = reinterpret_cast<uint16_t*>(dst);
-
-  // Process TILE_K elements at a time (PACK_ELEMENTS_PER_K_CHUNK output)
-  int32_t k = 0;
-  for (; k + TILE_K <= K; k += TILE_K, d += PACK_ELEMENTS_PER_K_CHUNK) {
-    vst1q_u16(d, vcombine_u16(vld1_u16(s0 + k), vld1_u16(s1 + k)));
-  }
-
-  // Handle K tail: pack remaining elements with zero-padding
-  const int32_t tail = K - k;
-  if (tail > 0) {
-    // Pack remaining tail elements: [r0[k..k+tail-1], pad, r1[k..k+tail-1],
-    // pad]
-    for (int32_t t = 0; t < tail; ++t) {
-      d[t] = s0[k + t];
-      d[t + TILE_K] = s1[k + t];
-    }
-    // Zero-pad the rest
-    for (int32_t t = tail; t < TILE_K; ++t) {
-      d[t] = 0;
-      d[t + TILE_K] = 0;
-    }
-  }
-}
-
-// 2x2 accumulator load/store with compile-time row count
-template <int32_t m_rows>
-FORCE_INLINE float32x4_t load_acc_2x2(float* base, int64_t ldc, int col_off) {
-  static_assert(m_rows == 1 || m_rows == 2);
-  float32x2_t row0 = vld1_f32(base + col_off);
-  float32x2_t row1 =
-      (m_rows == 2) ? vld1_f32(base + ldc + col_off) : vdup_n_f32(0.f);
-  return vcombine_f32(row0, row1);
-}
-
-template <int32_t m_rows>
-FORCE_INLINE void store_acc_2x2(float32x4_t acc, float* base, int64_t ldc,
-                                int col_off) {
-  static_assert(m_rows == 1 || m_rows == 2);
-  vst1_f32(base + col_off, vget_low_f32(acc));
-  if constexpr (m_rows == 2) {
-    vst1_f32(base + ldc + col_off, vget_high_f32(acc));
-  }
-}
-
-// Initialize 4 column-pair accumulators for 2 rows (8 columns total)
-#define INIT_ACC_ROWPAIR_4(a0, a1, a2, a3, Crow, ldc, m_rows, accum) \
-  do {                                                               \
-    if (accum) {                                                     \
-      if (m_rows == 2) {                                             \
-        a0 = load_acc_2x2<2>(Crow, ldc, 0);                          \
-        a1 = load_acc_2x2<2>(Crow, ldc, 2);                          \
-        a2 = load_acc_2x2<2>(Crow, ldc, 4);                          \
-        a3 = load_acc_2x2<2>(Crow, ldc, 6);                          \
-      } else {                                                       \
-        a0 = load_acc_2x2<1>(Crow, ldc, 0);                          \
-        a1 = load_acc_2x2<1>(Crow, ldc, 2);                          \
-        a2 = load_acc_2x2<1>(Crow, ldc, 4);                          \
-        a3 = load_acc_2x2<1>(Crow, ldc, 6);                          \
-      }                                                              \
-    } else {                                                         \
-      a0 = a1 = a2 = a3 = vdupq_n_f32(0.f);                          \
-    }                                                                \
-  } while (0)
-
-// Store 4 column-pair accumulators back to C matrix
-#define STORE_ACC_ROWPAIR_4(a0, a1, a2, a3, Crow, ldc, m_rows) \
-  do {                                                         \
-    if (m_rows == 2) {                                         \
-      store_acc_2x2<2>(a0, Crow, ldc, 0);                      \
-      store_acc_2x2<2>(a1, Crow, ldc, 2);                      \
-      store_acc_2x2<2>(a2, Crow, ldc, 4);                      \
-      store_acc_2x2<2>(a3, Crow, ldc, 6);                      \
-    } else {                                                   \
-      store_acc_2x2<1>(a0, Crow, ldc, 0);                      \
-      store_acc_2x2<1>(a1, Crow, ldc, 2);                      \
-      store_acc_2x2<1>(a2, Crow, ldc, 4);                      \
-      store_acc_2x2<1>(a3, Crow, ldc, 6);                      \
-    }                                                          \
-  } while (0)
-
-// Perform 4 BFMMLA operations: acc += A @ B for 4 column-pairs
-#define BFMMLA_COMPUTE_4(r0, r1, r2, r3, a, b0, b1, b2, b3) \
-  do {                                                      \
-    r0 = vbfmmlaq_f32(r0, a, b0);                           \
-    r1 = vbfmmlaq_f32(r1, a, b1);                           \
-    r2 = vbfmmlaq_f32(r2, a, b2);                           \
-    r3 = vbfmmlaq_f32(r3, a, b3);                           \
-  } while (0)
-
-// Micro-kernel: updates a small fixed tile using BFMMLA.
-// RP = number of row-pairs (1,2,4)
-// Computes C[TILE_ROWS*RP, OUTPUT_COLS_PER_BLOCK] += A_packed @ B.
-// A_packed interleaves RP row-pairs; B layout is driven by the attention phase:
-// - AttentionGemmPhase::QK -> token-column layout (Q @ K^T)
-// - AttentionGemmPhase::PV -> token-row layout (P @ V)
-// K_static < 0 enables runtime K (PV only)
-template <int32_t RP, int32_t K_static, AttentionGemmPhase phase>
-FORCE_INLINE void gemm_rowpairs_x8_bfmmla_neon(
-    const bfloat16_t* const* __restrict A_packed_rp,
-    const int32_t* __restrict m_rows_rp, const bfloat16_t* __restrict B_blk,
-    float* __restrict C, int64_t ldc, bool accumulate, int64_t b_stride,
-    int32_t K_runtime = 0) {
-  static_assert(RP == 1 || RP == 2 || RP == 4, "RP must be 1,2,4");
-  static_assert(K_static < 0 || K_static % TILE_K == 0,
-                "K must be divisible by TILE_K");
-  static_assert(K_static >= 0 || phase == AttentionGemmPhase::PV,
-                "Runtime K only supported for PV");
-
-  constexpr bool runtime_k = (K_static < 0);
-  const int32_t K_iters =
-      runtime_k ? (K_runtime / TILE_K) : (K_static / TILE_K);
-  const int32_t K_tail = runtime_k ? (K_runtime % TILE_K) : 0;
-
-  if (!runtime_k) {
-    // Help the compiler fold away unused K_runtime when K is compile-time
-    (void)K_runtime;
-  }
-
-  auto* C_al = C;
-  const auto* B_al = B_blk;
-
-  // Setup A pointers
-  const bfloat16_t* a_ptr[4] = {
-      A_packed_rp[0],
-      (RP >= 2) ? A_packed_rp[1] : nullptr,
-      (RP >= 4) ? A_packed_rp[2] : nullptr,
-      (RP >= 4) ? A_packed_rp[3] : nullptr,
-  };
-
-  // Setup B pointers based on layout
-  const bfloat16_t* b_ptr[4];
-  if constexpr (phase == AttentionGemmPhase::PV) {
-    b_ptr[0] = B_blk + 0 * b_stride;
-    b_ptr[1] = B_blk + 1 * b_stride;
-    b_ptr[2] = B_blk + 2 * b_stride;
-    b_ptr[3] = B_blk + 3 * b_stride;
-  }
-
-  float32x4_t acc[4][4];
-
-// Initialize accumulators
-#define INIT_RP(rp)                                                            \
-  if constexpr (RP > rp) {                                                     \
-    INIT_ACC_ROWPAIR_4(acc[rp][0], acc[rp][1], acc[rp][2], acc[rp][3],         \
-                       C_al + (rp * 2) * ldc, ldc, m_rows_rp[rp], accumulate); \
-  }
-  INIT_RP(0);
-  INIT_RP(1);
-  INIT_RP(2);
-  INIT_RP(3);
-#undef INIT_RP
-
-  // Main compute loop
-  for (int32_t ki = 0; ki < K_iters; ++ki) {
-    bfloat16x8_t b0, b1, b2, b3;
-    if constexpr (phase == AttentionGemmPhase::PV) {
-      b0 = vld1q_bf16(b_ptr[0] + ki * V_INNER_STRIDE);
-      b1 = vld1q_bf16(b_ptr[1] + ki * V_INNER_STRIDE);
-      b2 = vld1q_bf16(b_ptr[2] + ki * V_INNER_STRIDE);
-      b3 = vld1q_bf16(b_ptr[3] + ki * V_INNER_STRIDE);
-    } else {
-      const bfloat16_t* b_base = B_al + ki * b_stride;
-      b0 = vld1q_bf16(b_base + 0 * V_INNER_STRIDE);
-      b1 = vld1q_bf16(b_base + 1 * V_INNER_STRIDE);
-      b2 = vld1q_bf16(b_base + 2 * V_INNER_STRIDE);
-      b3 = vld1q_bf16(b_base + 3 * V_INNER_STRIDE);
-    }
-
-#define COMPUTE_RP(rp)                                                       \
-  if constexpr (RP > rp) {                                                   \
-    bfloat16x8_t a = vld1q_bf16(a_ptr[rp] + ki * PACK_ELEMENTS_PER_K_CHUNK); \
-    BFMMLA_COMPUTE_4(acc[rp][0], acc[rp][1], acc[rp][2], acc[rp][3], a, b0,  \
-                     b1, b2, b3);                                            \
-  }
-    COMPUTE_RP(0);
-    COMPUTE_RP(1);
-    COMPUTE_RP(2);
-    COMPUTE_RP(3);
-#undef COMPUTE_RP
-  }
-
-  // K tail for runtime PV: fallback path
-  if constexpr (runtime_k) {
-    if (K_tail > 0) {
-      const int32_t tail_offset = K_iters * V_INNER_STRIDE;
-      const int32_t a_tail_offset = K_iters * PACK_ELEMENTS_PER_K_CHUNK;
-      for (int32_t kt = 0; kt < K_tail; ++kt) {
-        float32x4_t b_vecs[4];
-        for (int32_t p = 0; p < 4; ++p) {
-          const bfloat16_t* bp = b_ptr[p] + tail_offset + kt * TILE_COLS;
-          const float b0 = vcvtah_f32_bf16(bp[0]);
-          const float b1 = vcvtah_f32_bf16(bp[1]);
-          const float32x2_t b_pair = vset_lane_f32(b1, vdup_n_f32(b0), 1);
-          b_vecs[p] = vcombine_f32(b_pair, b_pair);
-        }
-
-#define TAIL_RP(rp)                                                     \
-  if constexpr (RP > rp) {                                              \
-    const bfloat16_t* ap = A_packed_rp[rp] + a_tail_offset;             \
-    float a_row0 = vcvtah_f32_bf16(ap[kt]);                             \
-    float a_row1 =                                                      \
-        (m_rows_rp[rp] == 2) ? vcvtah_f32_bf16(ap[kt + TILE_K]) : 0.0f; \
-    const float32x4_t a_vec =                                           \
-        vcombine_f32(vdup_n_f32(a_row0), vdup_n_f32(a_row1));           \
-    for (int32_t p = 0; p < 4; ++p) {                                   \
-      acc[rp][p] = vmlaq_f32(acc[rp][p], a_vec, b_vecs[p]);             \
-    }                                                                   \
-  }
-        TAIL_RP(0);
-        TAIL_RP(1);
-        TAIL_RP(2);
-        TAIL_RP(3);
-#undef TAIL_RP
-      }
-    }
-  }
-
-  // Store results
-#define STORE_RP(rp)                                                    \
-  if constexpr (RP > rp) {                                              \
-    STORE_ACC_ROWPAIR_4(acc[rp][0], acc[rp][1], acc[rp][2], acc[rp][3], \
-                        C_al + (rp * 2) * ldc, ldc, m_rows_rp[rp]);     \
-  }
-  STORE_RP(0);
-  STORE_RP(1);
-  STORE_RP(2);
-  STORE_RP(3);
-#undef STORE_RP
-}
-
-// Meso-kernel: packs a small MBxK slice of A, then tiles over N and calls the
-// micro-kernel for each OUTPUT_COLS_PER_BLOCK chunk. K_static < 0 enables
-// runtime K (PV only).
-template <int32_t MB, int32_t N, int32_t K_static, AttentionGemmPhase phase>
-FORCE_INLINE void gemm_packA_compute_MB_xN(
-    const c10::BFloat16* __restrict A, const c10::BFloat16* __restrict B,
-    float* __restrict C, int32_t K_runtime, int64_t lda, int64_t ldc,
-    int64_t b_layout_stride, int64_t b_reduction_stride, bool accumulate) {
-  static_assert(MB >= 1 && MB <= 8, "MB must be in [1,8]");
-  static_assert(N % OUTPUT_COLS_PER_BLOCK == 0,
-                "N must be a multiple of OUTPUT_COLS_PER_BLOCK");
-  static_assert(K_static < 0 || K_static % TILE_K == 0,
-                "K must be divisible by TILE_K");
-  static_assert(K_static >= 0 || phase == AttentionGemmPhase::PV,
-                "Runtime K only supported for PV");
-
-  constexpr bool runtime_k = (K_static < 0);
-  const int32_t K_val = runtime_k ? K_runtime : K_static;
-
-  // Keep small packs on-stack to avoid heap churn
-  constexpr int32_t STACK_PACK_STRIDE =
-      (1024 / TILE_K) * PACK_ELEMENTS_PER_K_CHUNK;
-
-  constexpr int32_t ROW_PAIRS = (MB + 1) / TILE_ROWS;
-  const int32_t pack_stride =
-      runtime_k ? ((K_val + TILE_K - 1) / TILE_K) * PACK_ELEMENTS_PER_K_CHUNK
-                : (K_static / TILE_K) * PACK_ELEMENTS_PER_K_CHUNK;
-
-  alignas(64) c10::BFloat16 A_packed_stack[ROW_PAIRS * STACK_PACK_STRIDE];
-  std::vector<c10::BFloat16> A_packed_heap;
-  c10::BFloat16* A_packed =
-      (pack_stride <= STACK_PACK_STRIDE)
-          ? A_packed_stack
-          : (A_packed_heap.resize(ROW_PAIRS * pack_stride),
-             A_packed_heap.data());
-
-  for (int32_t rp = 0; rp < ROW_PAIRS; ++rp) {
-    const int32_t m = rp * TILE_ROWS;
-    const int32_t m_rows = (m + 1 < MB) ? TILE_ROWS : 1;
-    const c10::BFloat16* A0 = A + m * lda;
-    const c10::BFloat16* A1 = (m_rows == TILE_ROWS) ? (A + (m + 1) * lda) : A0;
-    reshape_Q_2xK_for_bfmmla(A0, A1, A_packed + rp * pack_stride, K_val);
-  }
-
-  for (int32_t n = 0; n < N; n += OUTPUT_COLS_PER_BLOCK) {
-    const c10::BFloat16* B_blk_c10 =
-        (phase == AttentionGemmPhase::PV)
-            ? (B + (n / TILE_COLS) * b_layout_stride)
-            : (B + (n / OUTPUT_COLS_PER_BLOCK) * b_layout_stride);
-    const bfloat16_t* B_blk = reinterpret_cast<const bfloat16_t*>(B_blk_c10);
-
-    // Process row-pairs in groups of 4, 2, then 1
-    int32_t row_pair_idx = 0;
-
-#define PROCESS_RP_GROUP(group_size)                                       \
-  for (; row_pair_idx + (group_size - 1) < ROW_PAIRS;                      \
-       row_pair_idx += group_size) {                                       \
-    const bfloat16_t* Ap[group_size];                                      \
-    int32_t mr[group_size];                                                \
-    for (int32_t i = 0; i < group_size; ++i) {                             \
-      Ap[i] = reinterpret_cast<const bfloat16_t*>(                         \
-          A_packed + (row_pair_idx + i) * pack_stride);                    \
-      mr[i] = (((row_pair_idx + i) * TILE_ROWS + 1) < MB) ? TILE_ROWS : 1; \
-    }                                                                      \
-    float* C_blk = C + (row_pair_idx * TILE_ROWS) * ldc + n;               \
-    if constexpr (runtime_k) {                                             \
-      gemm_rowpairs_x8_bfmmla_neon<group_size, -1, phase>(                 \
-          Ap, mr, B_blk, C_blk, ldc, accumulate, b_layout_stride, K_val);  \
-    } else {                                                               \
-      gemm_rowpairs_x8_bfmmla_neon<group_size, K_static, phase>(           \
-          Ap, mr, B_blk, C_blk, ldc, accumulate,                           \
-          (phase == AttentionGemmPhase::PV) ? b_layout_stride              \
-                                            : b_reduction_stride);         \
-    }                                                                      \
-  }
-
-    PROCESS_RP_GROUP(4);
-    PROCESS_RP_GROUP(2);
-    PROCESS_RP_GROUP(1);
-#undef PROCESS_RP_GROUP
-  }
-}
-
-// Macro-kernel: iterates over M in MB={8,4,2,1} chunks.
-// Supports compile-time K specialization when K >= 0; otherwise uses runtime K
-// (runtime K path is only supported for PV).
-template <AttentionGemmPhase phase, int32_t N, int32_t K = -1>
-FORCE_INLINE void gemm_macro_neon_bfmmla(
-    const c10::BFloat16* __restrict A, const c10::BFloat16* __restrict B,
-    float* __restrict C, int32_t M, int32_t K_runtime, int64_t lda, int64_t ldc,
-    int64_t b_layout_stride, int64_t b_reduction_stride, bool accumulate) {
-  static_assert(N % OUTPUT_COLS_PER_BLOCK == 0,
-                "N must be a multiple of OUTPUT_COLS_PER_BLOCK");
-
-  if constexpr (K >= 0) {
-    static_assert(K % TILE_K == 0, "K must be divisible by TILE_K");
-    for (int32_t m = 0; m < M;) {
-      const int32_t rem = M - m;
-      const c10::BFloat16* A_blk = A + m * lda;
-      float* C_blk = C + m * ldc;
-
-#define DISPATCH_MB(mb)                                                   \
-  gemm_packA_compute_MB_xN<mb, N, K, phase>(A_blk, B, C_blk, 0, lda, ldc, \
-                                            b_layout_stride,              \
-                                            b_reduction_stride, accumulate)
-
-      if (rem >= 8) {
-        DISPATCH_MB(8);
-        m += 8;
-      } else if (rem >= 4) {
-        DISPATCH_MB(4);
-        m += 4;
-      } else if (rem >= 2) {
-        DISPATCH_MB(2);
-        m += 2;
-      } else {
-        DISPATCH_MB(1);
-        m += 1;
-      }
-#undef DISPATCH_MB
-    }
-  } else {
-    static_assert(phase == AttentionGemmPhase::PV,
-                  "Runtime K specialization only supported for PV.");
-    const int32_t K_val = K_runtime;
-
-    for (int32_t m = 0; m < M;) {
-      const int32_t rem = M - m;
-      const c10::BFloat16* A_blk = A + m * lda;
-      float* C_blk = C + m * ldc;
-
-#define DISPATCH_MB_RUNTIME(mb)                                                \
-  gemm_packA_compute_MB_xN<mb, N, -1, phase>(A_blk, B, C_blk, K_val, lda, ldc, \
-                                             b_layout_stride,                  \
-                                             b_reduction_stride, accumulate)
-
-      if (rem >= 8) {
-        DISPATCH_MB_RUNTIME(8);
-        m += 8;
-      } else if (rem >= 4) {
-        DISPATCH_MB_RUNTIME(4);
-        m += 4;
-      } else if (rem >= 2) {
-        DISPATCH_MB_RUNTIME(2);
-        m += 2;
-      } else {
-        DISPATCH_MB_RUNTIME(1);
-        m += 1;
-      }
-#undef DISPATCH_MB_RUNTIME
-    }
-  }
-}
-
-#undef INIT_ACC_ROWPAIR_4
-#undef STORE_ACC_ROWPAIR_4
-#undef BFMMLA_COMPUTE_4
+constexpr int32_t TILE_K = BfmmlaGemm::KTile;
+constexpr int32_t TILE_COLS = 2;
+constexpr int32_t OUTPUT_COLS_PER_BLOCK = BfmmlaGemm::NTile;
+constexpr int32_t K_TOKENS_PER_GROUP = 8;
+constexpr int32_t V_TOKENS_PER_ROW_BLOCK = 4;
+constexpr int32_t K_CACHE_K_GROUP_STRIDE = K_TOKENS_PER_GROUP * TILE_K;
+constexpr int32_t B_COL_PAIR_STRIDE = V_TOKENS_PER_ROW_BLOCK * TILE_COLS;
 
 }  // namespace
 
-// TileGemm Adapter for Attention
-
 template <typename kv_cache_t, int32_t BlockTokens, int32_t HeadDim>
 class TileGemmNEONBFMMLA {
  public:
@@ -443,49 +286,31 @@ class TileGemmNEONBFMMLA {
                                 [[maybe_unused]] const int32_t block_size,
                                 [[maybe_unused]] const int32_t dynamic_k_size,
                                 const bool accum_c) {
-    static_assert(BlockTokens % OUTPUT_COLS_PER_BLOCK == 0);
-    // BFMMLA kernels require compile-time head_dim; keep head_dim_ct only for
-    // API parity with other tile_gemm implementations.
+    static_assert(BlockTokens % 16 == 0);
     if constexpr (head_dim_ct >= 0) {
-      static_assert(head_dim_ct == HeadDim,
-                    "BFMMLA expects head_dim_ct to match HeadDim; PV passes "
-                    "-1 for API parity.");
+      static_assert(head_dim_ct == HeadDim);
     }
 
+    const auto* a = reinterpret_cast<const c10::BFloat16*>(a_tile);
     if constexpr (phase == AttentionGemmPhase::QK) {
-      const int64_t b_reduction_stride = K_INNER_STRIDE;
-      const int64_t b_token_block_stride = (HeadDim / TILE_K) * K_INNER_STRIDE;
+      constexpr int64_t b_n_group_stride =
+          (HeadDim / BfmmlaGemm::KTile) * K_CACHE_K_GROUP_STRIDE;
 
-      gemm_macro_neon_bfmmla<AttentionGemmPhase::QK, BlockTokens, HeadDim>(
-          reinterpret_cast<const c10::BFloat16*>(a_tile), b_tile, c_tile,
-          m_size, 0, lda, ldc, b_token_block_stride, b_reduction_stride,
-          accum_c);
+      for (int32_t row = 0; row < m_size; row += BfmmlaGemm::MaxRows) {
+        const int32_t panel_m = std::min(BfmmlaGemm::MaxRows, m_size - row);
+        BfmmlaGemm::gemm(a + row * HeadDim, b_tile, c_tile + row * ldc, panel_m,
+                         BlockTokens, HeadDim, 2 * HeadDim, b_n_group_stride,
+                         K_CACHE_K_GROUP_STRIDE, ldc, accum_c);
+      }
     } else {
-      const int64_t b_pair_stride =
-          (block_size / V_TOKENS_PER_ROW_BLOCK) * V_INNER_STRIDE;
+      const int64_t b_n_group_stride =
+          (block_size / V_TOKENS_PER_ROW_BLOCK) * K_CACHE_K_GROUP_STRIDE;
 
-      // PV gemm with runtime K specialization
-      switch (dynamic_k_size) {
-        case 32:
-          gemm_macro_neon_bfmmla<AttentionGemmPhase::PV, HeadDim, 32>(
-              reinterpret_cast<const c10::BFloat16*>(a_tile), b_tile, c_tile,
-              m_size, 32, lda, ldc, b_pair_stride, 0, accum_c);
-          break;
-        case 128:
-          gemm_macro_neon_bfmmla<AttentionGemmPhase::PV, HeadDim, 128>(
-              reinterpret_cast<const c10::BFloat16*>(a_tile), b_tile, c_tile,
-              m_size, 128, lda, ldc, b_pair_stride, 0, accum_c);
-          break;
-        case 256:
-          gemm_macro_neon_bfmmla<AttentionGemmPhase::PV, HeadDim, 256>(
-              reinterpret_cast<const c10::BFloat16*>(a_tile), b_tile, c_tile,
-              m_size, 256, lda, ldc, b_pair_stride, 0, accum_c);
-          break;
-        default:
-          gemm_macro_neon_bfmmla<AttentionGemmPhase::PV, HeadDim>(
-              reinterpret_cast<const c10::BFloat16*>(a_tile), b_tile, c_tile,
-              m_size, dynamic_k_size, lda, ldc, b_pair_stride, 0, accum_c);
-          break;
+      for (int32_t row = 0; row < m_size; row += BfmmlaGemm::MaxRows) {
+        const int32_t panel_m = std::min(BfmmlaGemm::MaxRows, m_size - row);
+        BfmmlaGemm::gemm(a + row * lda, b_tile, c_tile + row * ldc, panel_m,
+                         HeadDim, dynamic_k_size, 2 * lda, b_n_group_stride,
+                         K_CACHE_K_GROUP_STRIDE, ldc, accum_c);
       }
     }
   }
@@ -512,9 +337,11 @@ class AttentionImplNEONBFMMLA {
   static constexpr int64_t HeadDim = head_dim;
   static constexpr ISA ISAType = isa_type;
   static constexpr bool scale_on_logits = false;
+  static constexpr int64_t VCacheNGroup = OUTPUT_COLS_PER_BLOCK;
+  static constexpr int64_t VCacheKGroupStride = VCacheNGroup * TILE_K;
 
-  static_assert(HeadDim % OUTPUT_COLS_PER_BLOCK == 0);
-  static_assert(BlockSizeAlignment % OUTPUT_COLS_PER_BLOCK == 0);
+  static_assert(HeadDim % (2 * OUTPUT_COLS_PER_BLOCK) == 0);
+  static_assert(BlockSizeAlignment % K_TOKENS_PER_GROUP == 0);
   static_assert(HeadDim % TILE_K == 0, "HeadDim must be a multiple of TILE_K");
 
  public:
@@ -527,19 +354,36 @@ class AttentionImplNEONBFMMLA {
     attention_iteration(CPU_ATTENTION_PARAMS);
   }
 
+  struct ProbabilityTokenStore {
+    static constexpr int32_t TokenStride = 2;
+
+    FORCE_INLINE static void store_probabilities(
+        c10::BFloat16* __restrict__ probability,
+        const vec_op::FP32Vec16& values, const int32_t row,
+        const int64_t row_stride) {
+      const int32_t row_in_pair = row & 1;
+      auto* dst = reinterpret_cast<bfloat16_t*>(
+          probability + row_in_pair * (BfmmlaGemm::KTile - row_stride));
+      vst1_bf16(dst, vcvt_bf16_f32(values.reg.val[0]));
+      vst1_bf16(dst + 2 * BfmmlaGemm::KTile, vcvt_bf16_f32(values.reg.val[1]));
+      vst1_bf16(dst + 4 * BfmmlaGemm::KTile, vcvt_bf16_f32(values.reg.val[2]));
+      vst1_bf16(dst + 6 * BfmmlaGemm::KTile, vcvt_bf16_f32(values.reg.val[3]));
+    }
+  };
+
   // Key cache stride per token group (TokenColumn layout; QK)
   static constexpr int64_t k_cache_token_group_stride(
       [[maybe_unused]] const int32_t block_size) {
     static_assert(BlockSizeAlignment % K_TOKENS_PER_GROUP == 0);
     return (BlockSizeAlignment / K_TOKENS_PER_GROUP) *
-           ((head_dim / TILE_K) * K_INNER_STRIDE);
+           ((head_dim / TILE_K) * K_CACHE_K_GROUP_STRIDE);
   }
 
   // Value cache stride per token group (TokenRow layout; PV)
   static constexpr int64_t v_cache_token_group_stride(
       [[maybe_unused]] const int32_t block_size) {
     static_assert(BlockSizeAlignment % V_TOKENS_PER_ROW_BLOCK == 0);
-    return (BlockSizeAlignment / V_TOKENS_PER_ROW_BLOCK) * V_INNER_STRIDE;
+    return (BlockSizeAlignment / V_TOKENS_PER_ROW_BLOCK) * VCacheKGroupStride;
   }
 
   // The stride to move to the "next" head_dim group
@@ -550,7 +394,7 @@ class AttentionImplNEONBFMMLA {
     return head_dim * block_size;
   }
 
-  // Convert Q heads to BF16 and apply scale factor using native BF16 intrinsics
+  // Scale Q and write row pairs in BFMMLA reduction order.
   static void copy_q_heads_tile(c10::BFloat16* __restrict__ src,
                                 c10::BFloat16* __restrict__ q_buffer,
                                 const int32_t q_num,
@@ -559,25 +403,42 @@ class AttentionImplNEONBFMMLA {
                                 const int64_t q_head_stride, float scale) {
     constexpr int32_t dim = static_cast<int32_t>(head_dim);
     const float32x4_t scale_vec = vdupq_n_f32(scale);
+    const bfloat16x4_t zero = vdup_n_bf16(bfloat16_t{});
+    const int32_t row_num = q_num * q_heads_per_kv;
 
-    for (int32_t qi = 0; qi < q_num; ++qi) {
-      for (int32_t hi = 0; hi < q_heads_per_kv; ++hi) {
-        c10::BFloat16* __restrict__ curr_q =
-            src + qi * q_num_stride + hi * q_head_stride;
-        c10::BFloat16* __restrict__ dst =
-            q_buffer + qi * q_heads_per_kv * head_dim + hi * head_dim;
+    for (int32_t row = 0; row < row_num; row += 2) {
+      const int32_t q0 = row / q_heads_per_kv;
+      const int32_t h0 = row % q_heads_per_kv;
+      const auto* row0 = reinterpret_cast<const bfloat16_t*>(
+          src + q0 * q_num_stride + h0 * q_head_stride);
+      const bool has_row1 = row + 1 < row_num;
+      const int32_t q1 = (row + 1) / q_heads_per_kv;
+      const int32_t h1 = (row + 1) % q_heads_per_kv;
+      const auto* row1 = has_row1
+                             ? reinterpret_cast<const bfloat16_t*>(
+                                   src + q1 * q_num_stride + h1 * q_head_stride)
+                             : nullptr;
+      auto* dst = reinterpret_cast<bfloat16_t*>(q_buffer + row * head_dim);
 
-        for (int32_t i = 0; i < dim; i += OUTPUT_COLS_PER_BLOCK) {
-          bfloat16x8_t in8 =
-              vld1q_bf16(reinterpret_cast<const bfloat16_t*>(curr_q + i));
-          float32x4_t lo = vmulq_f32(vcvtq_low_f32_bf16(in8), scale_vec);
-          float32x4_t hi = vmulq_f32(vcvtq_high_f32_bf16(in8), scale_vec);
+      for (int32_t k = 0; k < dim; k += OUTPUT_COLS_PER_BLOCK) {
+        const bfloat16x8_t in0 = vld1q_bf16(row0 + k);
+        const bfloat16x4_t out0_lo =
+            vcvt_bf16_f32(vmulq_f32(vcvtq_low_f32_bf16(in0), scale_vec));
+        const bfloat16x4_t out0_hi =
+            vcvt_bf16_f32(vmulq_f32(vcvtq_high_f32_bf16(in0), scale_vec));
 
-          bfloat16x4_t lo_b = vcvt_bf16_f32(lo);
-          bfloat16x4_t hi_b = vcvt_bf16_f32(hi);
-          bfloat16x8_t out = vcombine_bf16(lo_b, hi_b);
-          vst1q_bf16(reinterpret_cast<bfloat16_t*>(dst + i), out);
+        bfloat16x4_t out1_lo = zero;
+        bfloat16x4_t out1_hi = zero;
+        if (has_row1) {
+          const bfloat16x8_t in1 = vld1q_bf16(row1 + k);
+          out1_lo =
+              vcvt_bf16_f32(vmulq_f32(vcvtq_low_f32_bf16(in1), scale_vec));
+          out1_hi =
+              vcvt_bf16_f32(vmulq_f32(vcvtq_high_f32_bf16(in1), scale_vec));
         }
+
+        vst1q_bf16(dst + 2 * k, vcombine_bf16(out0_lo, out1_lo));
+        vst1q_bf16(dst + 2 * k + 8, vcombine_bf16(out0_hi, out1_hi));
       }
     }
   }
@@ -585,11 +446,12 @@ class AttentionImplNEONBFMMLA {
  public:
   // Reshape and cache K/V into BFMMLA-optimized layouts
   // K cache:
-  // [block_size/K_TOKENS_PER_GROUP][head_dim/TILE_K][K_INNER_STRIDE]
+  // [block_size/K_TOKENS_PER_GROUP][head_dim/TILE_K]
+  // [K_CACHE_K_GROUP_STRIDE]
   // - TokenColumn
   // V cache:
-  // [head_dim/TILE_COLS][block_size/V_TOKENS_PER_ROW_BLOCK][V_INNER_STRIDE]
-  // - TokenRows
+  // [head_dim/VCacheNGroup][block_size/V_TOKENS_PER_ROW_BLOCK]
+  // [VCacheKGroupStride]
   static void reshape_and_cache(
       const c10::BFloat16* __restrict__ key,
       const c10::BFloat16* __restrict__ value,
@@ -604,13 +466,13 @@ class AttentionImplNEONBFMMLA {
       const int64_t block_size,
       [[maybe_unused]] const int64_t block_size_stride,
       const float /*k_inv*/ = 0.0f, const float /*v_inv*/ = 0.0f) {
-    const int64_t k_block_stride = (head_dim / TILE_K) * K_INNER_STRIDE;
-    const int64_t v_pair_stride =
-        (block_size / V_TOKENS_PER_ROW_BLOCK) * V_INNER_STRIDE;
+    const int64_t k_block_stride = (head_dim / TILE_K) * K_CACHE_K_GROUP_STRIDE;
+    const int64_t v_n_group_stride =
+        (block_size / V_TOKENS_PER_ROW_BLOCK) * VCacheKGroupStride;
 
-#pragma omp parallel for
-    for (int64_t head_idx = 0; head_idx < head_num; ++head_idx) {
-      for (int64_t token_idx = 0; token_idx < token_num; ++token_idx) {
+#pragma omp parallel for collapse(2)
+    for (int64_t token_idx = 0; token_idx < token_num; ++token_idx) {
+      for (int64_t head_idx = 0; head_idx < head_num; ++head_idx) {
         const int64_t pos = slot_mapping[token_idx];
         if (pos < 0) continue;
 
@@ -636,12 +498,12 @@ class AttentionImplNEONBFMMLA {
               key_base + block_in_block * k_block_stride;
 
           for (int64_t hd4 = 0; hd4 < head_dim / TILE_K; ++hd4) {
-            uint16_t* dst_u16 = reinterpret_cast<uint16_t*>(
-                block_base + hd4 * K_INNER_STRIDE +
-                pair_in_block * V_INNER_STRIDE + lane_base);
-            const uint16_t* src_u16 =
-                reinterpret_cast<const uint16_t*>(key_src + hd4 * TILE_K);
-            vst1_u16(dst_u16, vld1_u16(src_u16));
+            bfloat16_t* dst = reinterpret_cast<bfloat16_t*>(
+                block_base + hd4 * K_CACHE_K_GROUP_STRIDE +
+                pair_in_block * B_COL_PAIR_STRIDE + lane_base);
+            const bfloat16_t* src =
+                reinterpret_cast<const bfloat16_t*>(key_src + hd4 * TILE_K);
+            vst1_bf16(dst, vld1_bf16(src));
           }
         }
 
@@ -655,22 +517,33 @@ class AttentionImplNEONBFMMLA {
               value_cache + block_idx * num_blocks_stride +
               head_idx * cache_head_num_stride;
 
-          const int64_t row_block = block_offset / V_TOKENS_PER_ROW_BLOCK;
+          const int64_t token_group = block_offset / V_TOKENS_PER_ROW_BLOCK;
           const int64_t lane = block_offset & (V_TOKENS_PER_ROW_BLOCK - 1);
 
-          c10::BFloat16* __restrict row_block_base =
-              value_base + row_block * V_INNER_STRIDE;
+          const auto* src = reinterpret_cast<const bfloat16_t*>(value_src);
+          auto* dst = reinterpret_cast<bfloat16_t*>(value_base);
+          for (int64_t hd8 = 0; hd8 < head_dim / OUTPUT_COLS_PER_BLOCK; ++hd8) {
+            const bfloat16x8_t values =
+                vld1q_bf16(src + hd8 * OUTPUT_COLS_PER_BLOCK);
+            const bfloat16x4_t low = vget_low_bf16(values);
+            const bfloat16x4_t high = vget_high_bf16(values);
+            bfloat16_t* group =
+                dst + hd8 * v_n_group_stride + token_group * VCacheKGroupStride;
 
-          for (int64_t hd2 = 0; hd2 < head_dim / TILE_COLS; ++hd2) {
-            c10::BFloat16* __restrict dst_val =
-                row_block_base + hd2 * v_pair_stride;
-
-            const uint16_t* src_u16 =
-                reinterpret_cast<const uint16_t*>(value_src);
-            uint16_t* dst_u16 = reinterpret_cast<uint16_t*>(dst_val);
-            dst_u16[lane] = src_u16[hd2 * TILE_COLS + 0];
-            dst_u16[lane + V_TOKENS_PER_ROW_BLOCK] =
-                src_u16[hd2 * TILE_COLS + 1];
+            vst1_lane_bf16(group + lane, low, 0);
+            vst1_lane_bf16(group + V_TOKENS_PER_ROW_BLOCK + lane, low, 1);
+            vst1_lane_bf16(group + B_COL_PAIR_STRIDE + lane, low, 2);
+            vst1_lane_bf16(
+                group + B_COL_PAIR_STRIDE + V_TOKENS_PER_ROW_BLOCK + lane, low,
+                3);
+            vst1_lane_bf16(group + 2 * B_COL_PAIR_STRIDE + lane, high, 0);
+            vst1_lane_bf16(
+                group + 2 * B_COL_PAIR_STRIDE + V_TOKENS_PER_ROW_BLOCK + lane,
+                high, 1);
+            vst1_lane_bf16(group + 3 * B_COL_PAIR_STRIDE + lane, high, 2);
+            vst1_lane_bf16(
+                group + 3 * B_COL_PAIR_STRIDE + V_TOKENS_PER_ROW_BLOCK + lane,
+                high, 3);
           }
         }
       }
@@ -680,4 +553,4 @@ class AttentionImplNEONBFMMLA {
 
 }  // namespace cpu_attention
 
-#endif  // CPU_ATTN_ASIMD_BFMMLA_HPP
+#endif  // CPU_ATTN_NEON_BFMMLA_HPP

examples/disaggregated/disaggregated_encoder/README.md

@@ -106,6 +106,82 @@ If you enable prefill instance (`--prefill-servers-urls` not disabled), you will
 | `--decode-servers-urls` | Comma-separated list of decode endpoints. Non-stream and stream paths both round-robin over this list. |
 | `--host`, `--port` | Bind address for the proxy itself (defaults: `0.0.0.0:8000`). |
 
+### Dynamic registration
+
+Alternatively, let the external launcher register ready instances over HTTP.
+No vLLM configuration changes or worker registration threads are required.
+Set `ADMIN_API_KEY` on the proxy and supply it as `X-API-Key` for registration
+and removal. These are trusted control-plane APIs; do not expose them publicly.
+
+```bash
+export ADMIN_API_KEY="your-admin-key"
+python disagg_epd_proxy.py --port 8000 --dynamic-registration
+```
+
+Omit the static server URL flags and register every stage. The roles determine
+the topology: `encode` + `prefill_decode` gives E+PD; `encode` + `prefill` +
+`decode` gives E+P+D. Standalone `decode` always requires an available `prefill`,
+even when D registers first or all P instances go offline; otherwise requests
+return `503`. The proxy rejects mixing combined PD with standalone P/D,
+including unhealthy instances still in the registry. Explicitly remove the
+old topology's P/D or PD registrations before switching topologies.
+
+After each instance is ready, the launcher registers its reachable HTTP URL:
+
+```bash
+curl --fail-with-body http://proxy-host:8000/instances \
+    -H "X-API-Key: $ADMIN_API_KEY" -H 'Content-Type: application/json' \
+    -d '{"role":"encode","url":"http://e-host:8001"}'
+curl --fail-with-body http://proxy-host:8000/instances \
+    -H "X-API-Key: $ADMIN_API_KEY" -H 'Content-Type: application/json' \
+    -d '{"role":"prefill_decode","url":"http://pd-host:8002"}'
+```
+
+For E+P+D, register P with `role: "prefill"` and D with `role: "decode"`.
+Example and NIXL EC connectors need no additional registration fields.
+For Mooncake, register `ec_zmq_addrs` on the **EC consumer** (`prefill_decode`
+or `prefill`), using its configured `ec_ip` and `ec_port`. Keep the fixed-port
+layout: supply the TP-rank-0 address for each DP replica, in DP-rank order;
+replica `r` uses `ec_port + r * tensor_parallel_size`. The connector discovers
+the remaining TP ranks itself. For example, DP=2, TP=2, `ec_port=19019`:
+
+```json
+{
+  "role": "prefill_decode",
+  "url": "http://pd-host:8002",
+  "dp_size": 2,
+  "ec_zmq_addrs": ["tcp://pd-host:19019", "tcp://pd-host:19021"]
+}
+```
+
+The proxy selects one consumer replica and uses it for both the encoder push
+and the consumer HTTP request. Standalone D does not need EC control addresses.
+Port allocation and avoiding collisions remain the launcher's responsibility.
+
+Inspect or remove instances without restarting the proxy:
+
+```bash
+curl http://proxy-host:8000/instances
+curl --fail-with-body -X DELETE \
+    'http://proxy-host:8000/instances?url=http://e-host:8001' \
+    -H "X-API-Key: $ADMIN_API_KEY"
+```
+
+Registration is idempotent. The proxy probes registered instances every
+`--probe-interval` seconds (default 5), with `--probe-timeout` seconds per probe
+(default 2). After `--fail-threshold` consecutive failures (default 3), it stops
+sending new requests to that instance. Healthy instances rejoin automatically;
+unreachable ones are forgotten after `--evicted-ttl` seconds (default 900;
+`0` retains them indefinitely). Removal stops new routing; already routed
+requests retain their selected endpoints, so drain requests before stopping
+the instance. The launcher must re-register instances after a proxy restart.
+
+The example launch scripts also support this flow: export `ADMIN_API_KEY` and
+set `DYNAMIC_REGISTRATION=1` when running `disagg_1e1pd_example.sh` or
+`disagg_1e1p1d_example.sh`. Their default remains static routing.
+
+### Static configuration
+
 The proxy batches images from the same user request assigned to the same encoder.
 Set the proxy environment variable `ENCODER_MAX_BATCH_SIZE` to limit the number
 of images per encoder subrequest. It defaults to `0` (unlimited); `1` sends each
@@ -137,6 +213,6 @@ For E + P + D setup:
 ```bash
 $ python disagg_encoder_proxy.py \
       --encode-servers-urls "http://e1:8001,http://e2:8001" \
-      --prefill-servers-urls "http://p1:8003,http://p2:8004" \ 
+      --prefill-servers-urls "http://p1:8003,http://p2:8004" \
       --decode-servers-urls "http://d1:8005,http://d2:8006"
 ```

examples/disaggregated/disaggregated_encoder/disagg_1e1p1d_example.sh

@@ -189,17 +189,34 @@ wait_for_server "$DECODE_PORT"
 ###############################################################################
 # Proxy
 ###############################################################################
+PROXY_ARGS=()
+if [[ "${DYNAMIC_REGISTRATION:-0}" == "1" ]]; then
+    : "${ADMIN_API_KEY:?Set ADMIN_API_KEY for dynamic registration}"
+    export ADMIN_API_KEY
+    PROXY_ARGS=(--dynamic-registration)
+else
+    PROXY_ARGS=(
+        --encode-servers-urls "http://localhost:$ENCODE_PORT"
+        --prefill-servers-urls "http://localhost:$PREFILL_PORT"
+        --decode-servers-urls "http://localhost:$DECODE_PORT"
+    )
+fi
 python disagg_epd_proxy.py \
     --host "0.0.0.0" \
     --port "$PROXY_PORT" \
-    --encode-servers-urls "http://localhost:$ENCODE_PORT" \
-    --prefill-servers-urls "http://localhost:$PREFILL_PORT" \
-    --decode-servers-urls "http://localhost:$DECODE_PORT" \
+    "${PROXY_ARGS[@]}" \
     >"${PROXY_LOG}" 2>&1 &
 
 PIDS+=($!)
 
 wait_for_server "$PROXY_PORT"
+if [[ "${DYNAMIC_REGISTRATION:-0}" == "1" ]]; then
+    for endpoint in "encode:$ENCODE_PORT" "prefill:$PREFILL_PORT" "decode:$DECODE_PORT"; do
+        curl --fail-with-body "http://127.0.0.1:$PROXY_PORT/instances" \
+            -H "X-API-Key: $ADMIN_API_KEY" -H 'Content-Type: application/json' \
+            -d "{\"role\":\"${endpoint%%:*}\",\"url\":\"http://127.0.0.1:${endpoint#*:}\"}"
+    done
+fi
 echo "All services are up!"
 
 ###############################################################################
@@ -238,4 +255,4 @@ curl http://127.0.0.1:"${PROXY_PORT}"/v1/chat/completions \
 
 # cleanup
 echo "cleanup..."
-cleanup
+cleanup

examples/disaggregated/disaggregated_encoder/disagg_1e1pd_example.sh

@@ -152,17 +152,33 @@ wait_for_server "$PREFILL_DECODE_PORT"
 ###############################################################################
 # Proxy
 ###############################################################################
+PROXY_ARGS=()
+if [[ "${DYNAMIC_REGISTRATION:-0}" == "1" ]]; then
+    : "${ADMIN_API_KEY:?Set ADMIN_API_KEY for dynamic registration}"
+    export ADMIN_API_KEY
+    PROXY_ARGS=(--dynamic-registration)
+else
+    PROXY_ARGS=(
+        --encode-servers-urls "http://localhost:$ENCODE_PORT"
+        --decode-servers-urls "http://localhost:$PREFILL_DECODE_PORT"
+    )
+fi
 python disagg_epd_proxy.py \
     --host "0.0.0.0" \
     --port "$PROXY_PORT" \
-    --encode-servers-urls "http://localhost:$ENCODE_PORT" \
-    --prefill-servers-urls "disable" \
-    --decode-servers-urls "http://localhost:$PREFILL_DECODE_PORT" \
+    "${PROXY_ARGS[@]}" \
     >"${PROXY_LOG}" 2>&1 &
 
 PIDS+=($!)
 
 wait_for_server "$PROXY_PORT"
+if [[ "${DYNAMIC_REGISTRATION:-0}" == "1" ]]; then
+    for endpoint in "encode:$ENCODE_PORT" "prefill_decode:$PREFILL_DECODE_PORT"; do
+        curl --fail-with-body "http://127.0.0.1:$PROXY_PORT/instances" \
+            -H "X-API-Key: $ADMIN_API_KEY" -H 'Content-Type: application/json' \
+            -d "{\"role\":\"${endpoint%%:*}\",\"url\":\"http://127.0.0.1:${endpoint#*:}\"}"
+    done
+fi
 echo "All services are up!"
 
 ###############################################################################
@@ -201,4 +217,4 @@ curl http://127.0.0.1:"${PROXY_PORT}"/v1/chat/completions \
 
 # cleanup
 echo "cleanup..."
-cleanup
+cleanup

examples/disaggregated/disaggregated_encoder/disagg_epd_proxy.py

@@ -1,12 +1,13 @@
 #!/usr/bin/env python3
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""disagg_encoder_proxy.py
+"""OpenAI-compatible proxy for E+PD and E+P+D disaggregation.
 
-Proxy that routes OpenAI-compatible “/v1/chat/completions” requests to two
-clusters:
-  • encode  (multimodal feature extraction)
-  • decode  (language-model inference)
+Use static server URLs, or --dynamic-registration for launcher-managed HTTP
+registration of E, P/PD and D instances. Both modes route OpenAI-compatible
+``/v1/chat/completions`` requests through encoder and inference clusters and
+share encoder fan-out, metadata-only rewriting, prefill, and retry/streaming
+forwarding.
 
 For MM input we:
     1. Extract *every* image/audio/video item.
@@ -20,22 +21,27 @@ from __future__ import annotations
 
 import argparse
 import asyncio
+import enum
 import hashlib
 import itertools
 import json
 import logging
 import os
 import random
+import secrets
 import time
 import uuid
 from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager, suppress
+from dataclasses import dataclass, field
 from typing import Any
 
 import aiohttp
 import msgspec
 import uvicorn
-from fastapi import FastAPI, HTTPException, Request
+from fastapi import Depends, FastAPI, Header, HTTPException, Request
 from fastapi.responses import JSONResponse, Response, StreamingResponse
+from pydantic import AnyHttpUrl, BaseModel, Field, model_validator
 
 ###############################################################################
 # FastAPI app & global state
@@ -44,6 +50,340 @@ from fastapi.responses import JSONResponse, Response, StreamingResponse
 logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s: %(message)s")
 logger = logging.getLogger("proxy")
 
+DEFAULT_PROBE_INTERVAL = 5.0
+DEFAULT_PROBE_TIMEOUT = 2.0
+DEFAULT_FAIL_THRESHOLD = 3
+# Stop probing an instance that has been down this long. 0 probes forever,
+# which is what a cluster that restarts instances in place wants.
+DEFAULT_EVICTED_TTL = 900.0
+
+
+class InstanceRole(str, enum.Enum):
+    ENCODE = "encode"
+    PREFILL = "prefill"
+    DECODE = "decode"
+    PREFILL_DECODE = "prefill_decode"
+
+
+@dataclass
+class InstanceRecord:
+    """One registered instance.
+
+    Attributes:
+        role: Which stage this instance serves.
+        url: Base OpenAI-compatible URL, e.g. ``http://host:8000``.
+        ec_zmq_addrs: Mooncake TP-rank-0 control addresses, one per DP replica,
+            supplied by the launcher from the consumer's fixed port config.
+        dp_size: Data-parallel replicas behind `url`, so the proxy can pick
+            a replica and name the same one to both halves of a request.
+
+    """
+
+    role: InstanceRole
+    url: str
+    ec_zmq_addrs: list[str] = field(default_factory=list)
+    dp_size: int = 1
+
+
+class InstanceRegistry:
+    def __init__(
+        self,
+        probe_interval: float = DEFAULT_PROBE_INTERVAL,
+        probe_timeout: float = DEFAULT_PROBE_TIMEOUT,
+        fail_threshold: int = DEFAULT_FAIL_THRESHOLD,
+        evicted_ttl: float = DEFAULT_EVICTED_TTL,
+    ):
+        self._probe_interval = probe_interval
+        self._probe_timeout = probe_timeout
+        self._fail_threshold = fail_threshold
+        self._evicted_ttl = evicted_ttl
+
+        self._live: dict[str, InstanceRecord] = {}
+        self._evicted: dict[str, InstanceRecord] = {}
+        self._evicted_since: dict[str, float] = {}
+        self._fail_counts: dict[str, int] = {}
+        # One cursor per role, only ever incremented. Rebuilding it whenever
+        # the roster changes -- what an `itertools.cycle` over a mutable list
+        # forces -- restarts every fan-out at the first instance and hot-spots
+        # it after each registration.
+        self._cursors: dict[InstanceRole, int] = {role: 0 for role in InstanceRole}
+        self._replica_cursors: dict[str, int] = {}
+        self._probe_task: asyncio.Task | None = None
+
+    def register(self, record: InstanceRecord) -> bool:
+        """Add or refresh an instance without overriding its health status."""
+        roles = {
+            other.role
+            for other in itertools.chain(self._live.values(), self._evicted.values())
+        }
+        split_roles = {InstanceRole.PREFILL, InstanceRole.DECODE}
+        if (record.role is InstanceRole.PREFILL_DECODE and roles & split_roles) or (
+            record.role in split_roles and InstanceRole.PREFILL_DECODE in roles
+        ):
+            raise ValueError("Cannot mix prefill_decode with standalone prefill/decode")
+        key = record.url
+        previous = self._live.get(key) or self._evicted.get(key)
+        if previous is not None:
+            if previous == record:
+                return False
+            if previous.role is not record.role:
+                raise ValueError("Unregister the instance before changing its role")
+            target = self._live if key in self._live else self._evicted
+            target[key] = record
+            return False
+        self._live[key] = record
+        logger.info("Registered instance %s: %s", record.role.value, record.url)
+        return True
+
+    def unregister(self, url: str) -> bool:
+        """Drop an instance for good, so a probe cannot bring it back."""
+        key = url
+        if key not in self._live and key not in self._evicted:
+            return False
+        self._live.pop(key, None)
+        self._evicted.pop(key, None)
+        self._evicted_since.pop(key, None)
+        self._fail_counts.pop(key, None)
+        self._replica_cursors.pop(key, None)
+        logger.info("Unregistered instance: %s", url)
+        return True
+
+    def instances(self, role: InstanceRole) -> list[InstanceRecord]:
+        return [record for record in self._live.values() if record.role is role]
+
+    def urls(self, role: InstanceRole) -> list[str]:
+        return [record.url for record in self.instances(role)]
+
+    def pick(self, role: InstanceRole) -> InstanceRecord | None:
+        """Take the next instance of `role` in round-robin order."""
+        picked = self.pick_many(role, 1)
+        return picked[0] if picked else None
+
+    def pick_many(self, role: InstanceRole, count: int) -> list[InstanceRecord]:
+        """Take `count` instances, continuing the rotation across calls.
+
+        A multimodal request fans out one encoder request per item, so the
+        assignment has to be contiguous with the previous request's rather
+        than restart at the first instance every time.
+        """
+        alive = self.instances(role)
+        if not alive or count <= 0:
+            return []
+        start = self._cursors[role]
+        self._cursors[role] = start + count
+        return [alive[(start + offset) % len(alive)] for offset in range(count)]
+
+    def next_replica(self, record: InstanceRecord) -> int:
+        """Take the next data-parallel replica of `record`, round-robin.
+
+        The encoder pushes to one replica's receive channel, so the request
+        has to run on that same replica; the caller names the rank to both
+        halves.
+        """
+        if record.dp_size <= 1:
+            return 0
+        cursor = self._replica_cursors.get(record.url, 0)
+        self._replica_cursors[record.url] = cursor + 1
+        return cursor % record.dp_size
+
+    def status(self) -> dict[str, Any]:
+        return {
+            role.value: {
+                "live": [record.url for record in self.instances(role)],
+                "evicted": [
+                    record.url
+                    for record in self._evicted.values()
+                    if record.role is role
+                ],
+            }
+            for role in InstanceRole
+        }
+
+    def start_probing(self) -> None:
+        if self._probe_task is None and self._probe_interval > 0:
+            self._probe_task = asyncio.create_task(self._probe_loop())
+
+    async def stop_probing(self) -> None:
+        if self._probe_task is None:
+            return
+        self._probe_task.cancel()
+        with suppress(asyncio.CancelledError):
+            await self._probe_task
+        self._probe_task = None
+
+    async def _probe_loop(self) -> None:
+        timeout = aiohttp.ClientTimeout(total=self._probe_timeout)
+        async with aiohttp.ClientSession(timeout=timeout) as session:
+            while True:
+                await asyncio.sleep(self._probe_interval)
+                try:
+                    await self._probe_once(session)
+                except Exception:
+                    logger.exception("EPD registry probe round failed")
+
+    async def _probe_once(self, session: aiohttp.ClientSession) -> None:
+        targets = list(self._live.values()) + list(self._evicted.values())
+        if not targets:
+            return
+        results = await asyncio.gather(
+            *(self._probe(session, record.url) for record in targets),
+            return_exceptions=True,
+        )
+        now = time.monotonic()
+        for record, healthy in zip(targets, results):
+            current = self._live.get(record.url) or self._evicted.get(record.url)
+            if current is not record:
+                continue
+            if healthy is True:
+                self._on_probe_success(record)
+            else:
+                self._on_probe_failure(record, now)
+        self._drop_expired(now)
+
+    async def _probe(self, session: aiohttp.ClientSession, url: str) -> bool:
+        async with session.get(f"{url}/health") as resp:
+            return resp.status == 200
+
+    def _on_probe_success(self, record: InstanceRecord) -> None:
+        key = record.url
+        self._fail_counts.pop(key, None)
+        if key in self._evicted:
+            self._evicted.pop(key, None)
+            self._evicted_since.pop(key, None)
+            self._live[key] = record
+            logger.info(
+                "Instance %s (%s) is healthy again; routing resumed",
+                record.url,
+                record.role.value,
+            )
+
+    def _on_probe_failure(self, record: InstanceRecord, now: float) -> None:
+        key = record.url
+        if key not in self._live:
+            # Either already evicted, or unregistered while this probe was in
+            # flight. A round snapshots its targets and then awaits, so a
+            # removal inside that window would otherwise be undone by a result
+            # describing a registry that no longer exists.
+            return
+        failures = self._fail_counts.get(key, 0) + 1
+        self._fail_counts[key] = failures
+        if failures < self._fail_threshold:
+            return
+        self._live.pop(key, None)
+        self._evicted[key] = record
+        self._evicted_since[key] = now
+        logger.warning(
+            "Instance %s (%s) failed %d consecutive probes; stopped routing "
+            "to it. It rejoins on its own once it responds again.",
+            record.url,
+            record.role.value,
+            failures,
+        )
+
+    def _drop_expired(self, now: float) -> None:
+        if self._evicted_ttl <= 0:
+            return
+        for key, since in list(self._evicted_since.items()):
+            if now - since < self._evicted_ttl:
+                continue
+            record = self._evicted.pop(key, None)
+            self._evicted_since.pop(key, None)
+            self._fail_counts.pop(key, None)
+            self._replica_cursors.pop(key, None)
+            if record is not None:
+                logger.warning("Instance %s stayed down; forgetting it", record.url)
+
+
+class InstanceRegistration(BaseModel):
+    role: InstanceRole
+    url: AnyHttpUrl
+    ec_zmq_addrs: list[str] = Field(default_factory=list)
+    dp_size: int = Field(default=1, ge=1)
+
+    @model_validator(mode="after")
+    def validate_consumer_addresses(self):
+        if self.ec_zmq_addrs and len(self.ec_zmq_addrs) != self.dp_size:
+            raise ValueError("Provide one Mooncake control address per DP replica")
+        return self
+
+
+def require_admin_key(x_api_key: str = Header(default="")) -> None:
+    expected = os.getenv("ADMIN_API_KEY", "")
+    if not expected or not secrets.compare_digest(x_api_key, expected):
+        raise HTTPException(403, "Invalid admin API key")
+
+
+@dataclass
+class EPDProxyConfig:
+    probe_interval: float = 5.0
+    probe_timeout: float = 2.0
+    fail_threshold: int = 3
+    evicted_ttl: float = 900.0
+
+
+@dataclass
+class _Route:
+    """The instances one request was assigned to.
+
+    Attributes:
+        consumer: The stage that receives the embedding, if any stage does.
+        consumer_zmq: The receive address named to the encoders.
+        dp_rank: Which replica of `consumer` was named, so the request can be
+            pinned to it.
+
+    """
+
+    encoder_urls: list[str]
+    prefill: InstanceRecord | None
+    decode: InstanceRecord
+    consumer: InstanceRecord | None = None
+    consumer_zmq: str | None = None
+    dp_rank: int | None = None
+
+
+class EPDProxy:
+    def __init__(self, registry: InstanceRegistry):
+        self.registry = registry
+
+    # ---------------------------------------------------------------- #
+    # Routing                                                          #
+    # ---------------------------------------------------------------- #
+    def route(self, num_items: int) -> _Route:
+        decode = self.registry.pick(InstanceRole.PREFILL_DECODE) or self.registry.pick(
+            InstanceRole.DECODE
+        )
+        if decode is None:
+            raise HTTPException(
+                status_code=503, detail="No decode instance is registered"
+            )
+        prefill = (
+            self.registry.pick(InstanceRole.PREFILL)
+            if decode.role is InstanceRole.DECODE
+            else None
+        )
+        if decode.role is InstanceRole.DECODE and prefill is None:
+            raise HTTPException(
+                status_code=503, detail="No prefill instance is registered"
+            )
+        encoders = self.registry.urls(InstanceRole.ENCODE)
+        if num_items and not encoders:
+            raise HTTPException(
+                status_code=503, detail="No encode instance is registered"
+            )
+        route = _Route(encoder_urls=encoders, prefill=prefill, decode=decode)
+        self._name_consumer(route)
+        return route
+
+    def _name_consumer(self, route: _Route) -> None:
+        """Pin the EC consumer replica and, for push connectors, its endpoint."""
+        candidate = route.prefill or route.decode
+        route.consumer = candidate
+        rank = self.registry.next_replica(candidate)
+        route.dp_rank = rank if candidate.dp_size > 1 else None
+        if candidate.ec_zmq_addrs:
+            route.consumer_zmq = candidate.ec_zmq_addrs[rank]
+
+
 app = FastAPI()
 encode_session: aiohttp.ClientSession | None = None
 prefill_session: aiohttp.ClientSession | None = None
@@ -448,6 +788,7 @@ async def maybe_prefill(
     req_data: dict,
     p_url: str,
     req_id: str,
+    dp_rank: int | None = None,
 ) -> dict:
     """- Do prefill-only task if p_url exist;
     - Return a new body carrying kv transfer params (for nixl connector)
@@ -459,7 +800,7 @@ async def maybe_prefill(
     if p_url:
         logger.info("[%s] Processing through prefill: %s", req_id, p_url)
 
-        prefill_response = await process_prefill_stage(req_data, p_url, req_id)
+        prefill_response = await process_prefill_stage(req_data, p_url, req_id, dp_rank)
         # for nixl connector to facilitate kv transfer...
         prefill_response_json = msgspec.json.decode(await prefill_response.read())
         kv_transfer_params = prefill_response_json.get("kv_transfer_params", {})
@@ -473,6 +814,7 @@ async def process_prefill_stage(
     req_data: dict,
     p_url: str,
     req_id: str,
+    dp_rank: int | None = None,
 ) -> dict:
     """Process request through Prefill stage and return kv_transfer_params"""
     logger.info("[%s] Sending prefill request to: %s", req_id, p_url)
@@ -494,6 +836,8 @@ async def process_prefill_stage(
         del prefill_request["stream_options"]
 
     headers = {"x-request-id": req_id, "Content-Type": "application/json"}
+    if dp_rank is not None:
+        headers["X-data-parallel-rank"] = str(dp_rank)
     try:
         prefill_response = await prefill_session.post(
             f"{p_url}/v1/chat/completions",
@@ -594,9 +938,7 @@ async def on_startup() -> None:
         ),
     )
     encode_session = aiohttp.ClientSession(timeout=timeout, connector=connector)
-    if app.state.p_urls:
-        # only setup if prefill instance(s) exist
-        prefill_session = aiohttp.ClientSession(timeout=timeout, connector=connector)
+    prefill_session = aiohttp.ClientSession(timeout=timeout, connector=connector)
     decode_session = aiohttp.ClientSession(timeout=timeout, connector=connector)
 
 
@@ -622,6 +964,7 @@ async def prepare_for_decode(
     e_urls: list[str],
     p_url: str,
     consumer_zmq: str | None,
+    prefill_dp_rank: int | None = None,
 ) -> tuple[dict, float, float]:
     """Encode, rewrite and prefill, returning the body to send to decode.
 
@@ -642,7 +985,7 @@ async def prepare_for_decode(
         handles.update(ec_params)
         prepared = {**prepared, "ec_transfer_params": handles}
     _t2 = time.perf_counter()
-    prepared = await maybe_prefill(prepared, p_url, req_id)
+    prepared = await maybe_prefill(prepared, p_url, req_id, prefill_dp_rank)
     return prepared, _t1 - _t0, _t2 - _t1
 
 
@@ -654,12 +997,13 @@ async def forward_non_stream(
     d_url: str,
     consumer_zmq: str | None,
     dp_rank: int | None = None,
+    prefill_dp_rank: int | None = None,
 ) -> Response:
     try:
         for attempt in range(DECODE_RETRIES + 1):
             _t0 = time.perf_counter()
             prepared, encode_s, rewrite_s = await prepare_for_decode(
-                req_data, req_id, e_urls, p_url, consumer_zmq
+                req_data, req_id, e_urls, p_url, consumer_zmq, prefill_dp_rank
             )
             _t2 = time.perf_counter()
 
@@ -733,12 +1077,13 @@ async def forward_stream(
     d_url: str,
     consumer_zmq: str | None,
     dp_rank: int | None = None,
+    prefill_dp_rank: int | None = None,
 ) -> AsyncIterator[bytes]:
     try:
         for attempt in range(DECODE_RETRIES + 1):
             _t0 = time.perf_counter()
             prepared, encode_s, rewrite_s = await prepare_for_decode(
-                req_data, req_id, e_urls, p_url, consumer_zmq
+                req_data, req_id, e_urls, p_url, consumer_zmq, prefill_dp_rank
             )
             _t2 = time.perf_counter()
 
@@ -976,10 +1321,147 @@ async def stop_profile(request: Request):
     return await _profile_cmd("stop", body, e_url, p_url, d_url)
 
 
+def build_app(config: EPDProxyConfig | None = None) -> FastAPI:
+    config = config or EPDProxyConfig()
+    registry = InstanceRegistry(
+        probe_interval=config.probe_interval,
+        probe_timeout=config.probe_timeout,
+        fail_threshold=config.fail_threshold,
+        evicted_ttl=config.evicted_ttl,
+    )
+    proxy = EPDProxy(registry)
+
+    @asynccontextmanager
+    async def lifespan(app: FastAPI):
+        await on_startup()
+        try:
+            registry.start_probing()
+            yield
+        finally:
+            await registry.stop_probing()
+            await on_shutdown()
+
+    app = FastAPI(lifespan=lifespan)
+    app.state.proxy = proxy
+    app.state.registry = registry
+
+    @app.get("/instances")
+    async def list_instances():
+        return registry.status()
+
+    @app.post("/instances", dependencies=[Depends(require_admin_key)])
+    async def register_instance(body: InstanceRegistration):
+        if body.ec_zmq_addrs and body.role not in (
+            InstanceRole.PREFILL,
+            InstanceRole.PREFILL_DECODE,
+        ):
+            raise HTTPException(400, "Only EC consumers accept Mooncake addresses")
+        try:
+            created = registry.register(
+                InstanceRecord(
+                    body.role,
+                    str(body.url).rstrip("/"),
+                    body.ec_zmq_addrs,
+                    body.dp_size,
+                )
+            )
+        except ValueError as error:
+            raise HTTPException(409, str(error)) from error
+        return {"registered": created}
+
+    @app.delete("/instances", dependencies=[Depends(require_admin_key)])
+    async def unregister_instance(url: AnyHttpUrl):
+        return {"removed": registry.unregister(str(url).rstrip("/"))}
+
+    @app.post("/v1/chat/completions")
+    async def chat_completions(request: Request):
+        req_data = msgspec.json.decode(await request.body())
+        req_id = request.headers.get("x-request-id", str(uuid.uuid4()))
+        route = proxy.route(len(extract_mm_items(req_data)))
+        args = (
+            req_data,
+            req_id,
+            route.encoder_urls,
+            route.prefill.url if route.prefill else None,
+            route.decode.url,
+            route.consumer_zmq,
+        )
+        ranks = {
+            "dp_rank": route.dp_rank if route.prefill is None else None,
+            "prefill_dp_rank": route.dp_rank if route.prefill is not None else None,
+        }
+        if req_data.get("stream", False):
+            return StreamingResponse(
+                forward_stream(*args, **ranks), media_type="text/event-stream"
+            )
+        return await forward_non_stream(*args, **ranks)
+
+    @app.get("/v1/models")
+    async def list_models():
+        decode = registry.pick(InstanceRole.PREFILL_DECODE) or registry.pick(
+            InstanceRole.DECODE
+        )
+        if decode is None:
+            raise HTTPException(
+                status_code=503, detail="No decode instance is registered"
+            )
+        async with decode_session.get(f"{decode.url}/v1/models") as resp:
+            resp.raise_for_status()
+            return await resp.json()
+
+    @app.get("/health")
+    async def health():
+        status = registry.status()
+        # An empty roster is not unhealthy: the proxy is meant to come up
+        # before anything registers with it.
+        return JSONResponse({"proxy": "healthy", "instances": status})
+
+    @app.post("/start_profile")
+    async def start_profile(request: Request):
+        return await _profile(registry, "start", await request.json())
+
+    @app.post("/stop_profile")
+    async def stop_profile(request: Request):
+        return await _profile(registry, "stop", await request.json())
+
+    return app
+
+
+async def _profile(registry: InstanceRegistry, cmd: str, payload: dict) -> dict:
+    headers = {"Authorization": f"Bearer {os.getenv('OPENAI_API_KEY', '')}"}
+    targets = {role.value: registry.pick(role) for role in InstanceRole}
+    results = await asyncio.gather(
+        *(
+            _post_if_available(
+                decode_session, f"{record.url}/{cmd}_profile", payload, headers
+            )
+            for record in targets.values()
+            if record is not None
+        )
+    )
+    reachable = [result for result in results if result is not None]
+    if not reachable:
+        raise HTTPException(
+            status_code=503,
+            detail="Profiling endpoints are disabled on every instance",
+        )
+    live = [name for name, record in targets.items() if record is not None]
+    return dict(zip(live, results))
+
+
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
     parser.add_argument("--host", default="0.0.0.0")
     parser.add_argument("--port", type=int, default=8000)
+    parser.add_argument(
+        "--dynamic-registration",
+        action="store_true",
+        help="Enable launcher-managed HTTP registration for E+PD or E+P+D.",
+    )
+    parser.add_argument("--probe-interval", type=float, default=DEFAULT_PROBE_INTERVAL)
+    parser.add_argument("--probe-timeout", type=float, default=DEFAULT_PROBE_TIMEOUT)
+    parser.add_argument("--fail-threshold", type=int, default=DEFAULT_FAIL_THRESHOLD)
+    parser.add_argument("--evicted-ttl", type=float, default=DEFAULT_EVICTED_TTL)
     parser.add_argument(
         "--log-requests",
         action="store_true",
@@ -995,12 +1477,12 @@ if __name__ == "__main__":
     )
     parser.add_argument(
         "--encode-servers-urls",
-        required=True,
+        default="",
         help='Comma-separated encode URLs ("http://e1:8001,http://e2:8001")',
     )
     parser.add_argument(
         "--prefill-servers-urls",
-        required=True,
+        default="none",
         help=(
             'Comma-separated prefill URLs ("http://p1:8003,http://p2:8004") '
             'to enable E->P->D, set "disable" or "none" to enable E->PD'
@@ -1008,7 +1490,7 @@ if __name__ == "__main__":
     )
     parser.add_argument(
         "--decode-servers-urls",
-        required=True,
+        default="",
         help='Comma-separated decode URLs ("http://d1:8005,http://d2:8006")',
     )
     parser.add_argument(
@@ -1043,58 +1525,87 @@ if __name__ == "__main__":
     )
 
     args = parser.parse_args()
+    if args.fail_threshold < 1:
+        parser.error("--fail-threshold must be at least 1")
     if args.log_requests:
         logging.getLogger().setLevel(logging.DEBUG)
         app.middleware("http")(log_requests)
     NO_REWRITE = args.no_rewrite
     DECODE_RETRIES = max(0, args.decode_retries)
-    app.state.e_urls = [
-        u.strip() for u in args.encode_servers_urls.split(",") if u.strip()
-    ]
-    app.state.d_urls = [
-        u.strip() for u in args.decode_servers_urls.split(",") if u.strip()
-    ]
-    app.state.d_ec_urls = [
-        u.strip() for u in args.ec_consumer_zmq_addrs.split(",") if u.strip()
-    ]
-    if args.ec_consumer_dp_size < 1:
-        parser.error("--ec-consumer-dp-size must be at least 1")
-    app.state.ec_consumer_dp_size = args.ec_consumer_dp_size
-    app.state.replica_counter = itertools.count()
-    expected = len(app.state.d_urls) * args.ec_consumer_dp_size
-    if app.state.d_ec_urls and len(app.state.d_ec_urls) != expected:
-        parser.error(
-            "--ec-consumer-zmq-addrs must contain one address per consumer "
-            f"replica: expected {expected} "
-            f"({len(app.state.d_urls)} servers x {args.ec_consumer_dp_size} replicas), "
-            f"got {len(app.state.d_ec_urls)}"
-        )
-    # handle prefill instances
-    if args.prefill_servers_urls.lower() in ("disable", "none", ""):
-        app.state.p_urls = []
-        logger.info(
-            "Disaggregated prefill phase explicitly disabled by user. Running E + PD..."
+    if args.dynamic_registration:
+        if not os.getenv("ADMIN_API_KEY"):
+            parser.error("--dynamic-registration requires ADMIN_API_KEY")
+        if (
+            args.encode_servers_urls
+            or args.prefill_servers_urls.lower() not in ("disable", "none", "")
+            or args.decode_servers_urls
+            or args.ec_consumer_zmq_addrs
+        ):
+            parser.error(
+                "With --dynamic-registration, instances register through /instances"
+            )
+        app = build_app(
+            EPDProxyConfig(
+                probe_interval=args.probe_interval,
+                probe_timeout=args.probe_timeout,
+                fail_threshold=args.fail_threshold,
+                evicted_ttl=args.evicted_ttl,
+            )
         )
+        if args.log_requests:
+            app.middleware("http")(log_requests)
     else:
-        app.state.p_urls = [
-            u.strip() for u in args.prefill_servers_urls.split(",") if u.strip()
+        if not args.encode_servers_urls or not args.decode_servers_urls:
+            parser.error(
+                "Static routing requires --encode-servers-urls "
+                "and --decode-servers-urls"
+            )
+        app.state.e_urls = [
+            u.strip() for u in args.encode_servers_urls.split(",") if u.strip()
         ]
-        logger.info("Disaggregated prefill phase is enabled. Running E + P + D...")
-    try:
-        validate_ec_consumer_routing(app.state.p_urls, app.state.d_ec_urls)
-    except ValueError as exc:
-        parser.error(str(exc))
+        app.state.d_urls = [
+            u.strip() for u in args.decode_servers_urls.split(",") if u.strip()
+        ]
+        app.state.d_ec_urls = [
+            u.strip() for u in args.ec_consumer_zmq_addrs.split(",") if u.strip()
+        ]
+        if args.ec_consumer_dp_size < 1:
+            parser.error("--ec-consumer-dp-size must be at least 1")
+        app.state.ec_consumer_dp_size = args.ec_consumer_dp_size
+        app.state.replica_counter = itertools.count()
+        expected = len(app.state.d_urls) * args.ec_consumer_dp_size
+        if app.state.d_ec_urls and len(app.state.d_ec_urls) != expected:
+            parser.error(
+                "--ec-consumer-zmq-addrs must contain one address per consumer "
+                f"replica: expected {expected} "
+                f"({len(app.state.d_urls)} servers x "
+                f"{args.ec_consumer_dp_size} replicas), "
+                f"got {len(app.state.d_ec_urls)}"
+            )
+        # handle prefill instances
+        if args.prefill_servers_urls.lower() in ("disable", "none", ""):
+            app.state.p_urls = []
+            logger.info("Disaggregated prefill disabled. Running E + PD...")
+        else:
+            app.state.p_urls = [
+                u.strip() for u in args.prefill_servers_urls.split(",") if u.strip()
+            ]
+            logger.info("Disaggregated prefill phase is enabled. Running E + P + D...")
+        try:
+            validate_ec_consumer_routing(app.state.p_urls, app.state.d_ec_urls)
+        except ValueError as exc:
+            parser.error(str(exc))
 
-    logger.info("Proxy listening on %s:%s", args.host, args.port)
-    logger.info("Encode servers: %s", app.state.e_urls)
-    logger.info("Prefill instances %s", app.state.p_urls)
-    logger.info("Decode servers: %s", app.state.d_urls)
-    if app.state.ec_consumer_dp_size > 1:
-        logger.info(
-            "EC consumer replicas per server: %d (control addresses: %s)",
-            app.state.ec_consumer_dp_size,
-            app.state.d_ec_urls,
-        )
+        logger.info("Proxy listening on %s:%s", args.host, args.port)
+        logger.info("Encode servers: %s", app.state.e_urls)
+        logger.info("Prefill instances %s", app.state.p_urls)
+        logger.info("Decode servers: %s", app.state.d_urls)
+        if app.state.ec_consumer_dp_size > 1:
+            logger.info(
+                "EC consumer replicas per server: %d (control addresses: %s)",
+                app.state.ec_consumer_dp_size,
+                app.state.d_ec_urls,
+            )
 
     uvicorn.run(
         app,

rust/src/engine-core-client/src/protocol/output.rs

@@ -130,10 +130,36 @@ pub struct EngineCoreOutput {
     #[serde(default)]
     pub new_sampling_mask: Option<MaybeWireSamplingMask>,
     /// Per-request speculative-decoding acceptance metrics, set on the final
-    /// output when `--per-request-spec-decode-metrics` is enabled. Opaque here;
-    /// the Rust frontend does not yet surface it in responses.
+    /// output when `--per-request-spec-decode-metrics` is enabled.
     #[serde(default)]
-    pub spec_decode_metrics: Option<OpaqueValue>,
+    pub spec_decode_metrics: Option<RequestSpecDecodeMetrics>,
+}
+
+/// Raw per-sequence speculative-decoding accumulator.
+///
+/// Python models this as a plain `@dataclass`, so it is serialized by msgspec
+/// as a map (named fields). Derived values such as acceptance rates are
+/// computed by the frontend, not carried on the wire.
+///
+/// Original Python definition: `RequestSpecDecodeMetrics` in
+/// `vllm/v1/metrics/stats.py`.
+#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
+pub struct RequestSpecDecodeMetrics {
+    /// Configured `num_speculative_tokens`.
+    #[serde(default)]
+    pub num_spec_tokens: u64,
+    /// Verify-step counts indexed by accepted draft-token count.
+    #[serde(default)]
+    pub histogram: Vec<u64>,
+    /// Total proposed draft tokens.
+    #[serde(default)]
+    pub num_draft_tokens: u64,
+    /// Accepted draft count per verify step; empty unless `detailed`.
+    #[serde(default)]
+    pub per_step_accepted: Vec<u64>,
+    /// Proposed draft count per verify step; empty unless `detailed`.
+    #[serde(default)]
+    pub per_step_drafted: Vec<u64>,
 }
 
 impl EngineCoreOutput {

rust/src/llm/src/output.rs

@@ -11,7 +11,9 @@ use futures::stream::FusedStream;
 use futures::{Stream, StreamExt as _, pin_mut};
 use serde::{Deserialize, Serialize};
 use vllm_engine_core_client::protocol::logprobs::Logprobs;
-use vllm_engine_core_client::protocol::output::{EngineCoreFinishReason, StopReason};
+use vllm_engine_core_client::protocol::output::{
+    EngineCoreFinishReason, RequestSpecDecodeMetrics, StopReason,
+};
 use vllm_engine_core_client::protocol::sampling_mask::SamplingMask;
 use vllm_engine_core_client::{AbortCause, EngineCoreOutputStream};
 
@@ -47,6 +49,8 @@ pub struct CollectedGenerateOutput {
     pub ec_transfer_params: Option<serde_json::Value>,
     /// Sampling support sets aligned one-to-one with generated token positions.
     pub sampling_mask: Option<SamplingMask>,
+    /// Per-request speculative-decoding metrics from the terminal output.
+    pub spec_decode_metrics: Option<RequestSpecDecodeMetrics>,
 }
 
 /// Prompt-scoped metadata emitted only once on the first [`GenerateOutput`] for
@@ -160,6 +164,8 @@ pub struct GenerateOutput {
     pub ec_transfer_params: Option<serde_json::Value>,
     /// Sampling support sets aligned one-to-one with `token_ids`.
     pub sampling_mask: Option<SamplingMask>,
+    /// Per-request speculative-decoding metrics, present on terminal outputs.
+    pub spec_decode_metrics: Option<RequestSpecDecodeMetrics>,
 }
 
 impl GenerateOutput {
@@ -208,6 +214,7 @@ impl GenerateOutput {
             kv_transfer_params: None,
             ec_transfer_params: None,
             sampling_mask: None,
+            spec_decode_metrics: None,
         }
     }
 }
@@ -308,6 +315,7 @@ impl Stream for GenerateOutputStream {
             kv_transfer_params: raw.kv_transfer_params,
             ec_transfer_params: raw.ec_transfer_params,
             sampling_mask,
+            spec_decode_metrics: raw.spec_decode_metrics,
         };
 
         Poll::Ready(Some(Ok(output)))
@@ -396,6 +404,7 @@ impl<T: Stream<Item = Result<GenerateOutput>> + Send> T {
                         kv_transfer_params: None,
                         ec_transfer_params: None,
                         sampling_mask,
+                        spec_decode_metrics: None,
                     });
                 }
 
@@ -409,6 +418,7 @@ impl<T: Stream<Item = Result<GenerateOutput>> + Send> T {
                     };
                     collected.kv_transfer_params = output.kv_transfer_params;
                     collected.ec_transfer_params = output.ec_transfer_params;
+                    collected.spec_decode_metrics = output.spec_decode_metrics;
                     if let Some(mask) = collected.sampling_mask.as_ref()
                         && mask.rows.len() != collected.token_ids.len()
                     {

rust/src/llm/tests/generate.rs

@@ -413,6 +413,7 @@ async fn collect_output_rejects_partial_sampling_mask() {
         sampling_mask: Some(SamplingMask {
             rows: vec![vec![1, 33, 99]],
         }),
+        spec_decode_metrics: None,
     };
 
     let error = futures::stream::iter([Ok(output)]).collect_output().await.unwrap_err();

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

@@ -29,7 +29,8 @@ use vllm_llm::{
 use self::convert::{ResponseOptions, prepare_generate_request};
 use self::types::{
     GenerateLogprob, GenerateResponse, GenerateResponseChoice, GenerateResponseStreamChoice,
-    GenerateStreamResponse, MultiModalPlaceholders, PlaceholderRangeInfo,
+    GenerateStreamResponse, MultiModalPlaceholders, PerRequestMetrics, PlaceholderRangeInfo,
+    SpeculativeDecodingMetrics, StreamingSpeculativeDecodingMetrics,
 };
 pub(crate) use self::types::{GenerateRequest, GenerateSamplingParams};
 pub(crate) use self::validate::validate_request_compat;
@@ -160,10 +161,14 @@ async fn generate_chunk_stream(
     let mut prompt_tokens = None;
     let mut prompt_token_ids = None;
     let mut usage = TokenUsage::default();
+    let mut spec_decode_metrics = None;
 
     while let Some(next) = stream.next().await {
         match next {
             Ok(output) => {
+                if let Some(metrics) = output.spec_decode_metrics {
+                    spec_decode_metrics = Some(SpeculativeDecodingMetrics::from(metrics));
+                }
                 if prompt_tokens.is_none()
                     && let Some(info) = output.prompt_info.as_ref()
                 {
@@ -223,6 +228,7 @@ async fn generate_chunk_stream(
                         .then(|| Usage::from_token_usage(usage, enable_prompt_tokens_details)),
                     mm_placeholders: prompt_token_ids.as_ref().and_then(|_| mm_placeholders.take()),
                     prompt_token_ids,
+                    metrics: None,
                 })
                 .await;
             }
@@ -243,6 +249,11 @@ async fn generate_chunk_stream(
             usage: Some(Usage::from_token_usage(usage, enable_prompt_tokens_details)),
             prompt_token_ids: None,
             mm_placeholders: None,
+            metrics: spec_decode_metrics.map(|speculative_decoding| PerRequestMetrics {
+                speculative_decoding: StreamingSpeculativeDecodingMetrics::from(
+                    speculative_decoding,
+                ),
+            }),
         })
         .await;
     }
@@ -318,6 +329,9 @@ fn collect_generate(
         mm_placeholders: return_token_ids.then_some(mm_placeholders).flatten(),
         kv_transfer_params: collected.kv_transfer_params,
         ec_transfer_params: collected.ec_transfer_params,
+        metrics: collected.spec_decode_metrics.map(|metrics| PerRequestMetrics {
+            speculative_decoding: SpeculativeDecodingMetrics::from(metrics),
+        }),
     })
 }
 
@@ -465,6 +479,7 @@ mod tests {
 
     use futures::{TryStreamExt as _, stream};
     use vllm_engine_core_client::protocol::multimodal::{MmModality, PlaceholderRange};
+    use vllm_engine_core_client::protocol::output::RequestSpecDecodeMetrics;
     use vllm_llm::GeneratePromptInfo;
 
     use super::*;
@@ -482,6 +497,7 @@ mod tests {
                 kv_transfer_params: None,
                 ec_transfer_params: None,
                 sampling_mask: None,
+                spec_decode_metrics: None,
             }),
             Ok(GenerateOutput {
                 request_id: String::new(),
@@ -496,6 +512,7 @@ mod tests {
                 kv_transfer_params: None,
                 ec_transfer_params: None,
                 sampling_mask: None,
+                spec_decode_metrics: None,
             }),
         ]);
 
@@ -578,6 +595,7 @@ mod tests {
             kv_transfer_params: None,
             ec_transfer_params: None,
             sampling_mask: None,
+            spec_decode_metrics: None,
         }
     }
 
@@ -645,6 +663,80 @@ mod tests {
         assert!(chunks[1].mm_placeholders.is_none());
     }
 
+    #[tokio::test]
+    async fn generate_chunk_stream_returns_terminal_spec_decode_metrics() {
+        let metrics = RequestSpecDecodeMetrics {
+            num_spec_tokens: 3,
+            histogram: vec![0, 1, 1, 0],
+            num_draft_tokens: 6,
+            ..Default::default()
+        };
+        let mut output = stream_output(None, Vec::new(), Some(FinishReason::stop_eos()));
+        output.spec_decode_metrics = Some(metrics.clone());
+
+        for include_usage in [false, true] {
+            let chunks: Vec<_> = generate_chunk_stream(
+                stream::iter(vec![Ok(output.clone())]),
+                "raw-stream".to_string(),
+                ApiServerOptions::default(),
+                ResponseOptions {
+                    include_usage,
+                    ..Default::default()
+                },
+                None,
+            )
+            .try_collect()
+            .await
+            .expect("collect chunks");
+            assert_eq!(chunks.len(), 1 + usize::from(include_usage));
+            assert!(chunks[0].metrics.is_none());
+            if include_usage {
+                assert!(chunks[1].choices.is_empty());
+                assert_eq!(
+                    chunks[1].metrics,
+                    Some(PerRequestMetrics {
+                        speculative_decoding: StreamingSpeculativeDecodingMetrics::from(
+                            SpeculativeDecodingMetrics::from(metrics.clone())
+                        )
+                    })
+                );
+                let json = serde_json::to_value(&chunks[1]).expect("serialize chunk");
+                let stats = &json["metrics"]["speculative_decoding"];
+                assert_eq!(stats.as_object().expect("metrics object").len(), 7);
+                assert_eq!(stats["num_draft_tokens"], 6);
+                assert_eq!(stats["num_accepted_draft_tokens"], 3);
+                assert_eq!(stats["mean_acceptance_length"], 2.5);
+                assert_eq!(
+                    stats["acceptance_histogram"],
+                    serde_json::json!([0, 1, 1, 0])
+                );
+                assert!(stats.get("per_step_accepted").is_none());
+            }
+        }
+    }
+
+    #[test]
+    fn spec_decode_metrics_zero_steps_and_detailed() {
+        let empty = SpeculativeDecodingMetrics::from(RequestSpecDecodeMetrics {
+            num_spec_tokens: 2,
+            histogram: vec![0, 0, 0],
+            ..Default::default()
+        });
+        assert_eq!(empty.mean_acceptance_length, 1.0);
+        assert_eq!(empty.draft_acceptance_rate, 0.0);
+        assert_eq!(empty.per_step_accepted, None);
+
+        let detailed = SpeculativeDecodingMetrics::from(RequestSpecDecodeMetrics {
+            num_spec_tokens: 2,
+            histogram: vec![0, 1, 1],
+            num_draft_tokens: 4,
+            per_step_accepted: vec![1, 2],
+            per_step_drafted: vec![2, 2],
+        });
+        assert_eq!(detailed.per_step_accepted, Some(vec![1, 2]));
+        assert_eq!(detailed.per_step_drafted, Some(vec![2, 2]));
+    }
+
     #[test]
     fn collect_generate_omits_prompt_metadata_by_default() {
         let output = CollectedGenerateOutput {
@@ -662,6 +754,7 @@ mod tests {
             ec_transfer_params: None,
             prompt_token_ids: vec![10, 20],
             sampling_mask: None,
+            spec_decode_metrics: None,
         };
 
         let response = collect_generate(
@@ -683,6 +776,50 @@ mod tests {
         assert!(response.mm_placeholders.is_none());
     }
 
+    #[test]
+    fn collect_generate_returns_spec_decode_metrics() {
+        let metrics = RequestSpecDecodeMetrics {
+            num_spec_tokens: 3,
+            histogram: vec![0, 1, 1, 0],
+            num_draft_tokens: 6,
+            ..Default::default()
+        };
+        let output = CollectedGenerateOutput {
+            request_id: "raw-metrics".to_string(),
+            prompt_logprobs: None,
+            token_ids: vec![30],
+            logprobs: None,
+            finish_reason: FinishReason::stop_eos(),
+            usage: TokenUsage::default(),
+            kv_transfer_params: None,
+            ec_transfer_params: None,
+            prompt_token_ids: vec![10],
+            sampling_mask: None,
+            spec_decode_metrics: Some(metrics.clone()),
+        };
+
+        let response = collect_generate(
+            output,
+            "raw-metrics".to_string(),
+            ApiServerOptions::default(),
+            ResponseOptions::default(),
+            None,
+        )
+        .expect("response");
+
+        assert_eq!(
+            response.metrics,
+            Some(PerRequestMetrics {
+                speculative_decoding: SpeculativeDecodingMetrics::from(metrics)
+            })
+        );
+        let json = serde_json::to_value(&response).expect("serialize response");
+        let stats = &json["metrics"]["speculative_decoding"];
+        assert_eq!(stats.as_object().expect("metrics object").len(), 9);
+        assert!(stats["per_step_accepted"].is_null());
+        assert!(stats["per_step_drafted"].is_null());
+    }
+
     #[test]
     fn extract_mm_placeholders_returns_none_without_features() {
         assert!(extract_mm_placeholders(None).is_none());
@@ -749,6 +886,7 @@ mod tests {
             ec_transfer_params: None,
             prompt_token_ids: vec![10, 20],
             sampling_mask: None,
+            spec_decode_metrics: None,
         };
 
         let response = collect_generate(
@@ -782,6 +920,7 @@ mod tests {
             kv_transfer_params: None,
             ec_transfer_params: None,
             sampling_mask: None,
+            spec_decode_metrics: None,
             prompt_token_ids,
         };
 

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

@@ -7,6 +7,7 @@ use llm_multimodal::MediaContentPart;
 use serde::{Deserialize, Serialize};
 use serde_json::{Map, Value};
 use validator::Validate;
+use vllm_engine_core_client::protocol::output::RequestSpecDecodeMetrics;
 use vllm_text::SamplingParams;
 
 use crate::routes::openai::utils::types::{ChatLogProbs, Normalizable, StreamOptions, Usage};
@@ -84,6 +85,7 @@ pub(super) struct GenerateStreamResponse {
     pub usage: Option<Usage>,
     pub prompt_token_ids: Option<Vec<u32>>,
     pub mm_placeholders: Option<MultiModalPlaceholders>,
+    pub metrics: Option<PerRequestMetrics<StreamingSpeculativeDecodingMetrics>>,
 }
 
 /// Mirrors the Python vLLM `GenerateResponse` class.
@@ -96,6 +98,91 @@ pub(super) struct GenerateResponse {
     pub mm_placeholders: Option<MultiModalPlaceholders>,
     pub kv_transfer_params: Option<Value>,
     pub ec_transfer_params: Option<Value>,
+    pub metrics: Option<PerRequestMetrics<SpeculativeDecodingMetrics>>,
+}
+
+#[derive(Debug, Clone, PartialEq, Serialize)]
+pub(super) struct PerRequestMetrics<T> {
+    pub speculative_decoding: T,
+}
+
+/// Mirrors the Python vLLM `SpeculativeDecodingMetrics` class.
+///
+/// Derived from the raw engine accumulator the same way as Python
+/// `RequestSpecDecodeMetrics.to_dict`.
+#[derive(Debug, Clone, PartialEq, Serialize)]
+pub(super) struct SpeculativeDecodingMetrics {
+    pub mean_acceptance_length: f64,
+    pub draft_acceptance_rate: f64,
+    pub acceptance_histogram: Vec<u64>,
+    pub num_spec_steps: u64,
+    pub num_accepted_draft_tokens: u64,
+    pub num_draft_tokens: u64,
+    pub num_spec_tokens: u64,
+    pub per_step_accepted: Option<Vec<u64>>,
+    pub per_step_drafted: Option<Vec<u64>>,
+}
+
+impl From<RequestSpecDecodeMetrics> for SpeculativeDecodingMetrics {
+    fn from(raw: RequestSpecDecodeMetrics) -> Self {
+        let num_spec_steps: u64 = raw.histogram.iter().sum();
+        let num_accepted_draft_tokens: u64 =
+            (0u64..).zip(&raw.histogram).map(|(accepted, count)| accepted * count).sum();
+        let ratio = |num: u64, den: u64| {
+            if den == 0 {
+                0.0
+            } else {
+                num as f64 / den as f64
+            }
+        };
+        let detailed = !raw.per_step_accepted.is_empty();
+        Self {
+            mean_acceptance_length: if num_spec_steps == 0 {
+                1.0
+            } else {
+                1.0 + ratio(num_accepted_draft_tokens, num_spec_steps)
+            },
+            draft_acceptance_rate: ratio(num_accepted_draft_tokens, raw.num_draft_tokens),
+            acceptance_histogram: raw.histogram,
+            num_spec_steps,
+            num_accepted_draft_tokens,
+            num_draft_tokens: raw.num_draft_tokens,
+            num_spec_tokens: raw.num_spec_tokens,
+            per_step_accepted: detailed.then_some(raw.per_step_accepted),
+            per_step_drafted: detailed.then_some(raw.per_step_drafted),
+        }
+    }
+}
+
+/// Streaming form omits detailed fields when summary metrics are requested.
+#[serde_with::skip_serializing_none]
+#[derive(Debug, Clone, PartialEq, Serialize)]
+pub(super) struct StreamingSpeculativeDecodingMetrics {
+    pub mean_acceptance_length: f64,
+    pub draft_acceptance_rate: f64,
+    pub acceptance_histogram: Vec<u64>,
+    pub num_spec_steps: u64,
+    pub num_accepted_draft_tokens: u64,
+    pub num_draft_tokens: u64,
+    pub num_spec_tokens: u64,
+    pub per_step_accepted: Option<Vec<u64>>,
+    pub per_step_drafted: Option<Vec<u64>>,
+}
+
+impl From<SpeculativeDecodingMetrics> for StreamingSpeculativeDecodingMetrics {
+    fn from(metrics: SpeculativeDecodingMetrics) -> Self {
+        Self {
+            mean_acceptance_length: metrics.mean_acceptance_length,
+            draft_acceptance_rate: metrics.draft_acceptance_rate,
+            acceptance_histogram: metrics.acceptance_histogram,
+            num_spec_steps: metrics.num_spec_steps,
+            num_accepted_draft_tokens: metrics.num_accepted_draft_tokens,
+            num_draft_tokens: metrics.num_draft_tokens,
+            num_spec_tokens: metrics.num_spec_tokens,
+            per_step_accepted: metrics.per_step_accepted,
+            per_step_drafted: metrics.per_step_drafted,
+        }
+    }
 }
 
 pub(super) type MultiModalPlaceholders = HashMap<String, Vec<PlaceholderRangeInfo>>;

tests/distributed/test_mnnvl_alltoall.py

@@ -14,6 +14,7 @@ import torch
 import torch.multiprocessing as mp
 
 from vllm.distributed import get_ep_group
+from vllm.platforms import current_platform
 from vllm.utils.flashinfer import (
     has_flashinfer_nvlink_one_sided,
     has_flashinfer_nvlink_two_sided,
@@ -23,6 +24,8 @@ from vllm.utils.network_utils import get_open_port
 
 from ..utils import init_test_distributed_environment
 
+DEVICE = current_platform.device_type
+
 # ---------------------------------------------------------------------------
 # Helpers
 # ---------------------------------------------------------------------------
@@ -483,7 +486,7 @@ def _args_dispatch_combine_worker(rank, world_size):
     from vllm.forward_context import get_forward_context
 
     cpu_group = get_ep_group().cpu_group
-    device = torch.device(f"cuda:{rank}")
+    device = torch.device(f"{DEVICE}:{rank}")
 
     hidden_size = 64
     tokens_per_rank = 16
@@ -624,7 +627,7 @@ def _two_sided_data_worker(rank, world_size):
     # Use DP group because MnnvlMoe workspace allocation calls get_dp_group()
     # internally and requires dp_size == ep_size.
     cpu_group = get_dp_group().cpu_group
-    device = torch.device(f"cuda:{rank}")
+    device = torch.device(f"{DEVICE}:{rank}")
     num_gpus = torch.accelerator.device_count()
 
     hidden_size = 128
@@ -769,7 +772,7 @@ def _one_sided_data_worker(rank, world_size):
     from vllm.forward_context import get_forward_context
 
     cpu_group = get_dp_group().cpu_group
-    device = torch.device(f"cuda:{rank}")
+    device = torch.device(f"{DEVICE}:{rank}")
 
     hidden_size = 256
     tokens_per_rank = 32

tests/distributed/test_multiproc_executor.py

@@ -15,6 +15,7 @@ import pytest
 from tests.utils import multi_gpu_test
 from vllm.config import VllmConfig
 from vllm.engine.arg_utils import EngineArgs
+from vllm.platforms import current_platform
 from vllm.v1.core.sched.output import SchedulerOutput
 from vllm.v1.executor import multiproc_executor
 from vllm.v1.executor.multiproc_executor import MultiprocExecutor
@@ -378,8 +379,8 @@ def test_multiproc_executor_properties():
 def test_multiproc_executor_multi_node():
     """Test MultiprocExecutor with multi-node configuration.
     This simulates 2 nodes with TP=4:
-    - Node 0 (rank 0): Uses GPUs 0,1 (CUDA_VISIBLE_DEVICES=0,1) with TP=2
-    - Node 1 (rank 1): Uses GPUs 2,3 (CUDA_VISIBLE_DEVICES=2,3) with TP=2
+    - Node 0 (rank 0): Uses GPUs 0,1 (device_control_env_var=0,1) with TP=2
+    - Node 1 (rank 1): Uses GPUs 2,3 (device_control_env_var=2,3) with TP=2
     Total world_size = 4, nnodes = 2
     """
     # Python 3.14+ changed default multiprocessing start method to 'forkserver'
@@ -392,15 +393,20 @@ def test_multiproc_executor_multi_node():
     # symm_mem does not work for simulating multi instance in single node
     os.environ["VLLM_ALLREDUCE_USE_SYMM_MEM"] = "0"
 
+    # Use the platform-specific device control env var (e.g.
+    # CUDA_VISIBLE_DEVICES on CUDA, ZE_AFFINITY_MASK on XPU) so that each
+    # node only sees its own devices on different platform.
+    device_control_env_var = current_platform.device_control_env_var
+
     def run_node(node_rank: int, result_queue: multiprocessing.Queue, port: int):
         """Run a single node's executor."""
         executor = None
         try:
-            # Set CUDA_VISIBLE_DEVICES for this node
+            # Restrict visible devices for this node
             if node_rank == 0:
-                os.environ["CUDA_VISIBLE_DEVICES"] = "0,1"
+                os.environ[device_control_env_var] = "0,1"
             else:
-                os.environ["CUDA_VISIBLE_DEVICES"] = "2,3"
+                os.environ[device_control_env_var] = "2,3"
 
             # Create config for this node
             vllm_config = create_vllm_config(

tests/entrypoints/scale_out/derender/test_derender.py

@@ -52,6 +52,7 @@ def _make_generate_response(
     logprobs: dict | None = None,
     prompt_logprobs: list | None = None,
     kv_transfer_params: dict | None = None,
+    metrics: dict | None = None,
 ) -> dict:
     choice: dict = {
         "index": 0,
@@ -64,6 +65,7 @@ def _make_generate_response(
         "choices": [choice],
         "prompt_logprobs": prompt_logprobs,
         "kv_transfer_params": kv_transfer_params,
+        "metrics": metrics,
     }
 
 
@@ -266,6 +268,49 @@ async def test_derender_chat_kv_transfer_params_passthrough(client):
     assert response.json()["kv_transfer_params"] == kv
 
 
+@pytest.mark.asyncio
+async def test_derender_chat_metrics_passthrough(client):
+    gen_req = await _render_chat(client)
+    metrics = {
+        "speculative_decoding": {
+            "mean_acceptance_length": 2.0,
+            "draft_acceptance_rate": 0.5,
+            "acceptance_histogram": [0, 1],
+            "num_spec_steps": 1,
+            "num_accepted_draft_tokens": 1,
+            "num_draft_tokens": 2,
+            "num_spec_tokens": 1,
+        }
+    }
+    response = await client.post(
+        "/v1/chat/completions/derender",
+        json={
+            "model": MODEL_NAME,
+            "generate_response": _make_generate_response(
+                gen_req["token_ids"][:3], metrics=metrics
+            ),
+        },
+    )
+    assert response.status_code == 200
+    actual = response.json()["metrics"]
+    assert set(actual) == {
+        "time_to_first_token_ms",
+        "generation_time_ms",
+        "queue_time_ms",
+        "mean_itl_ms",
+        "tokens_per_second",
+        "speculative_decoding",
+    }
+    assert all(
+        actual[name] is None for name in actual if name != "speculative_decoding"
+    )
+    assert actual["speculative_decoding"] == {
+        **metrics["speculative_decoding"],
+        "per_step_accepted": None,
+        "per_step_drafted": None,
+    }
+
+
 @pytest.mark.asyncio
 async def test_derender_chat_empty_token_ids(client):
     """Empty token_ids list returns 400."""

tests/entrypoints/scale_out/derender/test_derender_stream.py

@@ -27,6 +27,8 @@ from vllm.entrypoints.generate.base.protocol import (
     DeltaFunctionCall,
     DeltaMessage,
     DeltaToolCall,
+    PerRequestMetrics,
+    SpeculativeDecodingMetrics,
 )
 from vllm.entrypoints.scale_out.token_in_token_out.protocol import (
     DerenderStreamState,
@@ -164,6 +166,7 @@ def _make_stream_chunk(
     finish_reason: str | None = None,
     request_id: str = "test-req",
     usage: dict | None = None,
+    metrics: PerRequestMetrics | None = None,
 ) -> GenerateStreamResponse:
     """Build a GenerateStreamResponse SSE chunk."""
     return GenerateStreamResponse(
@@ -176,6 +179,7 @@ def _make_stream_chunk(
             )
         ],
         usage=UsageInfo(**usage) if usage else None,
+        metrics=metrics,
     )
 
 
@@ -512,6 +516,26 @@ class TestDerenderChatStream:
         assert chunk1.choices[0].delta.role == "assistant"
         assert chunk2.choices[0].delta.role is None
 
+    @pytest.mark.asyncio
+    async def test_metrics_passthrough(self, derenderer):
+        metrics = PerRequestMetrics(
+            speculative_decoding=SpeculativeDecodingMetrics(
+                mean_acceptance_length=2.0,
+                draft_acceptance_rate=0.5,
+                acceptance_histogram=[0, 1],
+                num_spec_steps=1,
+                num_accepted_draft_tokens=1,
+                num_draft_tokens=2,
+                num_spec_tokens=1,
+            )
+        )
+        chunk, _ = await derenderer.derender_chat_stream(
+            model=MODEL_NAME,
+            generate_chunk=_make_stream_chunk([], metrics=metrics),
+            state=DerenderStreamState(),
+        )
+        assert chunk.metrics == metrics
+
     @pytest.mark.asyncio
     async def test_chunked_equals_oneshot(self, derenderer, tokenizer):
         """Sum of streaming content deltas == one shot decode."""

tests/entrypoints/scale_out/token_in_token_out/test_generate_stream.py

@@ -28,6 +28,7 @@ from vllm.renderers import renderer_from_config
 from vllm.renderers.online_renderer import OnlineRenderer
 from vllm.sampling_params import SamplingParams
 from vllm.v1.engine.async_llm import AsyncLLM
+from vllm.v1.metrics.stats import RequestSpecDecodeMetrics
 
 MODEL_NAME = "openai-community/gpt2"
 BASE_MODEL_PATHS = [
@@ -130,6 +131,7 @@ def _make_request_output(
     logprobs: list[dict[int, Any] | None] | None = None,
     num_cached_tokens: int | None = None,
     index: int = 0,
+    spec_decode_metrics: RequestSpecDecodeMetrics | None = None,
 ) -> RequestOutput:
     return RequestOutput(
         request_id=request_id,
@@ -144,6 +146,7 @@ def _make_request_output(
                 cumulative_logprob=None,
                 logprobs=logprobs,
                 finish_reason=finish_reason,
+                spec_decode_metrics=spec_decode_metrics,
             )
         ],
         finished=finished,
@@ -321,6 +324,145 @@ async def test_stream_basic():
     assert data_chunks[2]["choices"][0]["finish_reason"] == "stop"
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("stream", [False, True])
+async def test_serve_tokens_returns_spec_decode_metrics(stream: bool):
+    engine = _mock_engine()
+    metrics = RequestSpecDecodeMetrics.new(num_spec_tokens=3)
+    metrics.observe(num_draft_tokens=3, num_accepted=2)
+
+    async def mock_generate(*args, **kwargs):
+        yield _make_request_output(
+            "req-1",
+            token_ids=[10],
+            finish_reason="stop",
+            finished=True,
+            spec_decode_metrics=metrics,
+        )
+
+    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),
+        model=MODEL_NAME,
+        stream=stream,
+        stream_options=StreamOptions(include_usage=True) if stream else None,
+    )
+
+    response = await serving.serve_tokens(request)
+    if stream:
+        chunks = [chunk async for chunk in response]
+        payload = next(
+            chunk
+            for chunk in _parse_sse_chunks(chunks)
+            if isinstance(chunk, dict) and chunk.get("metrics")
+        )["metrics"]["speculative_decoding"]
+    else:
+        assert isinstance(response, GenerateResponse)
+        assert response.metrics is not None
+        assert response.metrics.speculative_decoding is not None
+        payload = response.metrics.speculative_decoding.model_dump()
+
+    assert payload["num_draft_tokens"] == 3
+    assert payload["num_accepted_draft_tokens"] == 2
+    assert payload["num_spec_steps"] == 1
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("include_usage", [False, True])
+async def test_stream_returns_spec_decode_metrics_on_empty_terminal_output(
+    include_usage,
+):
+    engine = _mock_engine()
+    metrics = RequestSpecDecodeMetrics.new(num_spec_tokens=3)
+    metrics.observe(num_draft_tokens=3, num_accepted=2)
+
+    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,
+            spec_decode_metrics=metrics,
+        )
+
+    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),
+        model=MODEL_NAME,
+        stream=True,
+        stream_options=StreamOptions(include_usage=include_usage),
+    )
+
+    chunks = [chunk async for chunk in await serving.serve_tokens(request)]
+    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)
+    if include_usage:
+        assert data_chunks[1]["choices"] == []
+        assert (
+            data_chunks[1]["metrics"]["speculative_decoding"]["num_draft_tokens"] == 3
+        )
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("stream", [False, True])
+async def test_serve_tokens_omits_spec_decode_metrics_for_parallel_sampling(
+    stream: bool,
+):
+    engine = _mock_engine()
+    metrics = RequestSpecDecodeMetrics.new(num_spec_tokens=3)
+    metrics.observe(num_draft_tokens=3, num_accepted=2)
+    result = _make_request_output(
+        "req-1",
+        token_ids=[10],
+        finish_reason="stop",
+        finished=True,
+        spec_decode_metrics=metrics,
+    )
+    result.outputs.append(
+        CompletionOutput(
+            index=1,
+            text="",
+            token_ids=[20],
+            cumulative_logprob=None,
+            logprobs=None,
+            finish_reason="stop",
+            spec_decode_metrics=metrics,
+        )
+    )
+
+    async def mock_generate(*args, **kwargs):
+        yield result
+
+    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, n=2),
+        model=MODEL_NAME,
+        stream=stream,
+        stream_options=StreamOptions(include_usage=True) if stream else None,
+    )
+
+    response = await serving.serve_tokens(request)
+    if stream:
+        chunks = [chunk async for chunk in response]
+        assert all(
+            "metrics" not in chunk
+            for chunk in _parse_sse_chunks(chunks)
+            if isinstance(chunk, dict)
+        )
+    else:
+        assert isinstance(response, GenerateResponse)
+        assert response.metrics is None
+
+
 @pytest.mark.asyncio
 async def test_non_stream_returns_final_prompt_token_ids_when_requested():
     engine = _mock_engine()

tests/kernels/moe/test_moe.py

@@ -1384,24 +1384,33 @@ def test_humming_global_valid_shape_m(
     assert result == expected
 
 
-def test_humming_permute_scratch_is_keyed_by_runtime_topk(
+def test_humming_permute_scratch_is_shared_by_config(
     monkeypatch: pytest.MonkeyPatch,
 ):
     from types import SimpleNamespace
     from unittest.mock import Mock
 
     import vllm.model_executor.layers.fused_moe.experts.fused_humming_moe as humming
+    import vllm.model_executor.layers.fused_moe.moe_permute_unpermute as permute
+    import vllm.v1.worker.workspace as workspace
+
+    manager = workspace.WorkspaceManager(torch.device("cpu"))
+    monkeypatch.setattr(workspace, "_manager", manager)
 
     scratch_topk6 = Mock()
     scratch_topk1 = Mock()
     scratch_type = Mock(side_effect=[scratch_topk6, scratch_topk1])
     monkeypatch.setattr(humming, "moe_permute_unpermute_supported", lambda: True)
-    monkeypatch.setattr(humming, "MoEPermuteScratch", scratch_type)
+    monkeypatch.setattr(permute, "MoEPermuteScratch", scratch_type)
     moe_config = make_dummy_moe_config(max_num_tokens=512, experts_per_token=6)
     moe_config.moe_parallel_config.dp_size = 2
-    experts = SimpleNamespace(_permute_scratch={}, moe_config=moe_config)
+    experts = SimpleNamespace(moe_config=moe_config)
+    other_layer = SimpleNamespace(moe_config=moe_config)
 
     assert humming.HummingExpertsBase._get_permute_scratch(experts, 6) is scratch_topk6
+    assert (
+        humming.HummingExpertsBase._get_permute_scratch(other_layer, 6) is scratch_topk6
+    )
     assert humming.HummingExpertsBase._get_permute_scratch(experts, 1) is scratch_topk1
     assert humming.HummingExpertsBase._get_permute_scratch(experts, 6) is scratch_topk6
 
@@ -1412,6 +1421,11 @@ def test_humming_permute_scratch_is_keyed_by_runtime_topk(
     assert second_call.kwargs["max_num_tokens"] == 6144
     assert second_call.kwargs["topk"] == 1
 
+    manager.lock()
+    assert humming.HummingExpertsBase._get_permute_scratch(experts, 6) is scratch_topk6
+    with pytest.raises(AssertionError, match="was not allocated during warmup"):
+        humming.HummingExpertsBase._get_permute_scratch(experts, 2)
+
 
 def test_humming_delegates_to_instance_activation():
     from types import SimpleNamespace

tests/kernels/moe/test_moe_permute_unpermute.py

@@ -15,12 +15,14 @@ from vllm.model_executor.layers.fused_moe.expert_map_manager import (
 )
 from vllm.model_executor.layers.fused_moe.moe_permute_unpermute import (
     MoEPermuteScratch,
+    get_moe_permute_scratch,
     moe_permute,
     moe_permute_unpermute_supported,
     moe_unpermute,
 )
 from vllm.platforms import current_platform
 from vllm.utils.torch_utils import set_random_seed
+from vllm.v1.worker.workspace import current_workspace_manager
 
 NUM_EXPERTS = [16, 64, 256]
 TOP_KS = [2, 6, 8]
@@ -120,6 +122,7 @@ def torch_unpermute(
 @pytest.mark.parametrize("topk", TOP_KS)
 @pytest.mark.parametrize("dtype", [torch.bfloat16])
 @pytest.mark.parametrize("ep_size", EP_SIZE)
+@pytest.mark.parametrize("use_scratch", [False, True])
 def test_moe_permute_unpermute(
     n_token: int,
     n_hidden: int,
@@ -127,6 +130,8 @@ def test_moe_permute_unpermute(
     n_expert: int,
     ep_size: int,
     dtype: torch.dtype,
+    use_scratch: bool,
+    workspace_init,
 ):
     if not moe_permute_unpermute_supported():
         pytest.skip("moe_permute_unpermute is not supported on this platform.")
@@ -160,6 +165,18 @@ def test_moe_permute_unpermute(
         expert_map=expert_map,
     )
 
+    scratch = None
+    if use_scratch:
+        scratch = get_moe_permute_scratch(
+            max_num_tokens=n_token,
+            topk=topk,
+            num_experts=n_expert,
+            num_local_experts=n_local_expert,
+            device=hidden_states.device,
+            hidden_size=n_hidden,
+            hidden_dtype=dtype,
+        )
+
     (
         permuted_hidden_states,
         _,
@@ -173,6 +190,7 @@ def test_moe_permute_unpermute(
         n_expert=n_expert,
         n_local_expert=n_local_expert,
         expert_map=expert_map,
+        scratch=scratch,
     )
 
     # check expert_first_token_offset
@@ -212,20 +230,22 @@ def test_moe_permute_unpermute(
     torch.testing.assert_close(result4, gold4, atol=2e-2, rtol=0)
 
 
-@pytest.mark.parametrize("dtype", [torch.bfloat16])
-def test_moe_permute_reuses_scratch_buffers(dtype: torch.dtype):
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
+@pytest.mark.parametrize("n_token", [1, 33, 128])
+@pytest.mark.parametrize("topk", [1, 6, 8])
+def test_moe_permute_reuses_scratch_buffers(
+    dtype: torch.dtype, n_token: int, topk: int, workspace_init
+):
     if not moe_permute_unpermute_supported():
         pytest.skip("moe_permute_unpermute is not supported on this platform.")
 
-    n_token = 64
     n_hidden = 2048
     n_expert = 16
-    topk = 4
     hidden_states = torch.randn((n_token, n_hidden), device="cuda").to(dtype)
     gating_output = torch.randn((n_token, n_expert), device="cuda").to(dtype)
     _, topk_ids, _ = fused_topk(hidden_states, gating_output, topk, False)
 
-    scratch = MoEPermuteScratch(
+    scratch_config = dict(
         max_num_tokens=n_token,
         topk=topk,
         num_experts=n_expert,
@@ -234,6 +254,7 @@ def test_moe_permute_reuses_scratch_buffers(dtype: torch.dtype):
         hidden_size=n_hidden,
         hidden_dtype=hidden_states.dtype,
     )
+    scratch = get_moe_permute_scratch(**scratch_config)
 
     first = moe_permute(
         hidden_states=hidden_states,
@@ -242,13 +263,26 @@ def test_moe_permute_reuses_scratch_buffers(dtype: torch.dtype):
         n_expert=n_expert,
         scratch=scratch,
     )
-    second = moe_permute(
-        hidden_states=hidden_states,
-        a1q_scale=None,
-        topk_ids=topk_ids,
-        n_expert=n_expert,
-        scratch=scratch,
-    )
+    current_workspace_manager().lock()
+    assert get_moe_permute_scratch(**scratch_config) is scratch
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        second = moe_permute(
+            hidden_states=hidden_states,
+            a1q_scale=None,
+            topk_ids=topk_ids,
+            n_expert=n_expert,
+            scratch=get_moe_permute_scratch(**scratch_config),
+        )
+
+    for _ in range(2):
+        hidden_states.add_(1)
+        topk_ids.copy_(topk_ids.roll(1, 0))
+        graph.replay()
+        expected = moe_permute(hidden_states, None, topk_ids, n_expert)
+        for actual, reference in zip(second, expected):
+            if actual is not None:
+                torch.testing.assert_close(actual, reference)
 
     (
         permuted_hidden_states_1,
@@ -288,6 +322,123 @@ def test_moe_permute_reuses_scratch_buffers(dtype: torch.dtype):
     )
 
 
+def test_moe_permute_scratch_reused_across_graph_sizes(workspace_init) -> None:
+    """Switching graph sizes must not expose metadata left by a larger batch."""
+    if not moe_permute_unpermute_supported():
+        pytest.skip("moe_permute_unpermute is not supported on this platform.")
+
+    device = torch.device("cuda:0")
+    scratch = get_moe_permute_scratch(
+        max_num_tokens=128,
+        topk=6,
+        num_experts=16,
+        num_local_experts=16,
+        device=device,
+        hidden_size=2048,
+        hidden_dtype=torch.bfloat16,
+    )
+    current_workspace_manager().lock()
+    runs = []
+    for n_token in (128, 1, 33):
+        hidden = torch.randn((n_token, 2048), dtype=torch.bfloat16, device=device)
+        topk_ids = torch.randint(16, (n_token, 6), device=device)
+        moe_permute(hidden, None, topk_ids, 16, scratch=scratch)
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            output = moe_permute(hidden, None, topk_ids, 16, scratch=scratch)
+        runs.append((graph, hidden, topk_ids, output))
+
+    for graph, hidden, topk_ids, output in runs + runs[::-1]:
+        hidden.add_(1)
+        topk_ids.copy_((topk_ids + 1) % 16)
+        graph.replay()
+        expected = moe_permute(hidden, None, topk_ids, 16)
+        for actual, reference in zip(output, expected):
+            if actual is not None:
+                torch.testing.assert_close(actual, reference)
+
+
+def test_moe_permute_scratch_isolated_across_execution_slots(monkeypatch) -> None:
+    """Concurrent graph replays in different ubatches/lanes cannot share scratch."""
+    if not moe_permute_unpermute_supported():
+        pytest.skip("moe_permute_unpermute is not supported on this platform.")
+
+    from vllm.v1.worker import workspace
+
+    device = torch.device("cuda:0")
+    manager = workspace.WorkspaceManager(device, num_ubatches=2, num_lanes=2)
+    monkeypatch.setattr(workspace, "_manager", manager)
+    ubatch = 0
+    monkeypatch.setattr(workspace, "dbo_current_ubatch_id", lambda: ubatch)
+    scratch_config = dict(
+        max_num_tokens=64,
+        topk=6,
+        num_experts=16,
+        num_local_experts=16,
+        device=device,
+        hidden_size=2048,
+        hidden_dtype=torch.bfloat16,
+    )
+    runs = []
+    for i, n_token in enumerate((1, 7, 33, 64)):
+        ubatch = i // 2
+        stream = torch.cuda.Stream()
+        stream.wait_stream(torch.cuda.current_stream())
+        with workspace.use_workspace_lane(i % 2), torch.cuda.stream(stream):
+            hidden = torch.randn((n_token, 2048), dtype=torch.bfloat16, device=device)
+            topk_ids = torch.randint(16, (n_token, 6), device=device)
+            scratch = get_moe_permute_scratch(**scratch_config)
+            moe_permute(hidden, None, topk_ids, 16, scratch=scratch)
+            graph = torch.cuda.CUDAGraph()
+            with torch.cuda.graph(graph, stream=stream):
+                output = moe_permute(hidden, None, topk_ids, 16, scratch=scratch)
+            runs.append((stream, graph, hidden, topk_ids, output, scratch))
+        torch.cuda.current_stream().wait_stream(stream)
+
+    assert len({run[-1].permuted_hidden_states.data_ptr() for run in runs}) == 4
+    manager.lock()
+    for _ in range(3):
+        for i, (stream, graph, hidden, topk_ids, _, scratch) in enumerate(runs):
+            ubatch = i // 2
+            with workspace.use_workspace_lane(i % 2), torch.cuda.stream(stream):
+                assert get_moe_permute_scratch(**scratch_config) is scratch
+                hidden.add_(1)
+                topk_ids.copy_((topk_ids + 1) % 16)
+                graph.replay()
+        for stream, _, hidden, topk_ids, output, _ in runs:
+            torch.cuda.current_stream().wait_stream(stream)
+            expected = moe_permute(hidden, None, topk_ids, 16)
+            for actual, reference in zip(output, expected):
+                if actual is not None:
+                    torch.testing.assert_close(actual, reference)
+
+
+def test_moe_permute_scratch_without_manager(monkeypatch) -> None:
+    """Standalone calls get independent scratch with initialized row indices."""
+    if not moe_permute_unpermute_supported():
+        pytest.skip("moe_permute_unpermute is not supported on this platform.")
+
+    from vllm.v1.worker import workspace
+
+    monkeypatch.setattr(workspace, "_manager", None)
+    config = dict(
+        max_num_tokens=4,
+        topk=2,
+        num_experts=4,
+        num_local_experts=4,
+        device=torch.device("cuda"),
+    )
+    first = get_moe_permute_scratch(**config)
+    second = get_moe_permute_scratch(**config)
+
+    assert (
+        first.token_expert_indices.data_ptr() != second.token_expert_indices.data_ptr()
+    )
+    torch.testing.assert_close(
+        first.token_expert_indices, torch.arange(8, dtype=torch.int32, device="cuda")
+    )
+
+
 def test_moe_permute_ignores_invalid_expert_ids_with_scratch() -> None:
     if not moe_permute_unpermute_supported():
         pytest.skip("moe_permute_unpermute is not supported on this platform.")

tests/kernels/quantization/test_marlin_gemm.py

@@ -17,6 +17,7 @@ from vllm.model_executor.layers.quantization.utils.int8_utils import (
     per_token_quant_int8,
 )
 from vllm.model_executor.layers.quantization.utils.marlin_utils import (
+    get_marlin_workspace,
     marlin_make_empty,
     marlin_make_workspace_new,
     marlin_permute_bias,
@@ -43,6 +44,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
 )
 from vllm.platforms import current_platform
 from vllm.scalar_type import scalar_types
+from vllm.v1.worker.workspace import current_workspace_manager
 
 if current_platform.is_rocm():
     pytest.skip(
@@ -477,6 +479,62 @@ def test_marlin_gemm(
     assert max_diff < 0.04
 
 
+def test_marlin_persistent_locks_cuda_graph_streams(workspace_init):
+    """Compiled GEMMs select stream-local locks at runtime and survive replay."""
+    size_m, size_k, size_n = 32, 1024, 2048
+    inputs = [rand_data((size_m, size_k)) for _ in range(2)]
+    weight = rand_data((size_k, size_n))
+    quant_type = scalar_types.uint4b8
+    w_ref, qweight, scales = marlin_quantize(weight, quant_type, 128)
+    zeros = marlin_make_empty(scales.device)
+    streams = [torch.cuda.Stream() for _ in inputs]
+    graphs = [torch.cuda.CUDAGraph() for _ in inputs]
+    locks = []
+    outputs = []
+
+    @torch.compile(backend="eager", fullgraph=True)
+    def run(x):
+        return torch.ops.vllm.marlin_gemm(
+            x,
+            qweight,
+            None,
+            scales,
+            None,
+            None,
+            zeros,
+            None,
+            quant_type.id,
+            size_m,
+            size_n,
+            size_k,
+            use_atomic_add=False,
+            use_fp32_reduce=True,
+        )
+
+    for stream, graph, x in zip(streams, graphs, inputs):
+        stream.wait_stream(torch.cuda.current_stream())
+        with torch.cuda.stream(stream):
+            locks.append(get_marlin_workspace(x.device))
+            run(x)
+            with torch.cuda.graph(graph, stream=stream):
+                outputs.append(run(x))
+        torch.cuda.current_stream().wait_stream(stream)
+
+    assert locks[0].data_ptr() != locks[1].data_ptr()
+    current_workspace_manager().lock()
+    for _ in range(3):
+        for stream, graph, x, lock in zip(streams, graphs, inputs, locks):
+            with torch.cuda.stream(stream):
+                assert get_marlin_workspace(x.device) is lock
+                x.add_(0.1)
+                graph.replay()
+        for stream in streams:
+            torch.cuda.current_stream().wait_stream(stream)
+        for x, output, lock in zip(inputs, outputs, locks):
+            assert compute_max_diff(output, x @ w_ref) < 0.04
+            assert torch.count_nonzero(lock) == 0
+
+
 def test_marlin_gemm_subset_input():
     quant_type = scalar_types.uint4b8
     group_size = 128

tests/kernels/quantization/test_marlin_tile_padding.py

@@ -14,6 +14,7 @@ from vllm import _custom_ops as ops
 from vllm.model_executor.layers.quantization.utils.marlin_utils import (
     GPTQ_MARLIN_TILE,
     apply_gptq_marlin_linear,
+    get_marlin_workspace,
     marlin_make_empty,
     marlin_make_workspace_new,
     marlin_moe_padded_intermediate,
@@ -225,7 +226,7 @@ def test_fp8_marlin_padded_round_trip(shape, use_bias):
         input=x,
         weight=layer.weight,
         weight_scale=layer.weight_scale,
-        workspace=layer.workspace,
+        workspace=get_marlin_workspace(layer.weight.device),
         size_n=size_n,
         size_k=size_k,
         bias=layer.bias if use_bias else None,
@@ -288,7 +289,7 @@ def test_nvfp4_marlin_padded_round_trip(shape):
         weight=layer.weight,
         weight_scale=layer.weight_scale,
         weight_global_scale=layer.weight_global_scale,
-        workspace=layer.workspace,
+        workspace=get_marlin_workspace(layer.weight.device),
         size_n=size_n,
         size_k=size_k,
     )
@@ -392,7 +393,7 @@ def test_fp8_block_marlin_padded_round_trip(shape):
         input=x,
         weight=layer.weight,
         weight_scale=layer.weight_scale_inv,
-        workspace=layer.workspace,
+        workspace=get_marlin_workspace(layer.weight.device),
         size_n=size_n,
         size_k=size_k,
         bias=None,
@@ -440,7 +441,7 @@ def test_mxfp8_marlin_padded_round_trip(shape):
         input=x,
         weight=layer.weight,
         weight_scale=layer.weight_scale,
-        workspace=layer.workspace,
+        workspace=get_marlin_workspace(layer.weight.device),
         size_n=size_n,
         size_k=size_k,
     )
@@ -756,7 +757,7 @@ def test_fp8_marlin_moe_padded_round_trip(shape, quant):
         topk_ids,
         quant_type_id=scalar_types.float8_e4m3fn.id,
         global_num_experts=e,
-        workspace=layer.workspace,
+        workspace=get_marlin_workspace(device),
     )
     with set_current_vllm_config(VllmConfig()):
         ref = torch_experts(
@@ -833,7 +834,7 @@ def test_mxfp8_marlin_moe_padded_round_trip(shape):
         topk_ids,
         quant_type_id=scalar_types.float8_e4m3fn.id,
         global_num_experts=e,
-        workspace=layer.workspace,
+        workspace=get_marlin_workspace(device),
     )
     with set_current_vllm_config(VllmConfig()):
         ref = torch_experts(

tests/kernels/test_kpool_decode_update_batched.py

@@ -384,9 +384,16 @@ def test_leading_invalid_tail_slot():
     _assert_eq(r_ref, r_kern)
 
 
-@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm required")
-def test_amd_prefill_seed_honors_padded_tail_block_stride():
-    """The tail shares a padded indexer allocation in production."""
+def test_prefill_seed_honors_padded_tail_block_stride():
+    """The tail shares a padded indexer allocation in production.
+
+    ``get_kv_cache_config_from_groups`` aliases each tail tensor onto its
+    indexer tensor with the indexer's block stride (38016 B for GLM-5.3-Flash
+    vs a dense 2048 B tail block), so a seed kernel that addresses blocks
+    densely writes into an unrelated indexer block and leaves the request's
+    tail block untouched. Runs on every platform; the NVIDIA kernel had this
+    bug while the AMD kernel did not.
+    """
     kpool = 4
     num_blocks = 6
     logical_block_elems = 2 * kpool * HEAD_DIM

tests/kernels/test_top_k_per_row.py

@@ -1871,17 +1871,32 @@ def test_sparse_indexer_topk_backend_resolution() -> None:
     has_coop = _has_cooperative_topk()
     has_fi = _has_flashinfer_topk()
 
-    # "auto" is exactly the pre-existing chain: cooperative -> persistent ->
-    # per_row. It must never select the opt-in backends by itself.
+    # "auto" is the chain cooperative -> persistent -> per_row, with the
+    # cooperative/persistent switch at cooperative_topk's cluster-wave row
+    # limit. It must never select the opt-in backends by itself.
     if has_coop:
-        assert resolve("auto") == "cooperative"
+        from vllm.model_executor.layers.indexer_topk import (
+            AUTO_COOPERATIVE_MAX_ROWS,
+        )
+
+        # cooperative for every batch within its cluster-wave row limit...
+        assert (
+            resolve("auto", k=512, num_rows=AUTO_COOPERATIVE_MAX_ROWS) == "cooperative"
+        )
         assert resolve("auto", k=512, num_rows=8) == "cooperative"
-        # Past cooperative's 64-row limit -> persistent, even where
-        # DeepSelect would be applicable.
+        # ...persistent past the row limit, even where DeepSelect would be
+        # applicable.
+        assert (
+            resolve("auto", k=512, num_rows=AUTO_COOPERATIVE_MAX_ROWS + 1)
+            == "persistent"
+        )
         assert resolve("auto", k=512, num_rows=128) == "persistent"
-        # 16B-aligned (cooperative TMA ok) but not DeepSelect-aligned:
-        # still cooperative at <= 64 rows...
-        assert resolve("auto", unaligned_logits, num_rows=64) == "cooperative"
+        # 16B-aligned (cooperative TMA ok) but not DeepSelect-aligned: the
+        # same row-limit logic applies.
+        assert (
+            resolve("auto", unaligned_logits, num_rows=AUTO_COOPERATIVE_MAX_ROWS)
+            == "cooperative"
+        )
     # ...and persistent past the row limit.
     assert resolve("auto", unaligned_logits, num_rows=128) == "persistent"
     # Unsupported topk (outside {512, 1024, 2048} for the workspace kernels)

tests/model_executor/model_loader/test_reload.py

@@ -330,8 +330,8 @@ def _load_marlin_checkpoint_format_weights(layer):
     )
 
 
-def test_marlin_post_load_preserves_runtime_tensor_addresses(monkeypatch, dist_init):
-    """Marlin must reuse the workspace storage captured by CUDA graphs."""
+def test_marlin_post_load_does_not_own_workspace(monkeypatch, dist_init):
+    """Weight reload must not create layer-owned Marlin lock storage."""
     _stub_marlin_ops(monkeypatch)
     kernel = _make_marlin_kernel()
 
@@ -339,19 +339,17 @@ def test_marlin_post_load_preserves_runtime_tensor_addresses(monkeypatch, dist_i
     _load_marlin_checkpoint_format_weights(layer)
     kernel.process_weights_after_loading(layer)
 
-    workspace_ptr = kernel.workspace.data_ptr()
+    assert not hasattr(kernel, "workspace")
 
     _load_marlin_checkpoint_format_weights(layer)
     kernel.process_weights_after_loading(layer)
 
-    assert kernel.workspace.data_ptr() == workspace_ptr
-    assert torch.all(kernel.workspace == 0)
+    assert not hasattr(kernel, "workspace")
 
 
 @pytest.mark.parametrize("variant", ["fp8", "mxfp8", "nvfp4"])
-def test_marlin_prepare_layer_preserves_workspace_address(monkeypatch, variant):
-    """The Marlin fallback prepare_* functions rerun on weight reload and must
-    reuse the workspace storage whose address captured CUDA graphs hold."""
+def test_marlin_prepare_layer_does_not_own_workspace(monkeypatch, variant):
+    """Weight preparation must not attach runtime workspace to model layers."""
     from vllm import _custom_ops as ops
     from vllm.model_executor.layers.quantization.utils import (
         marlin_utils,
@@ -417,34 +415,57 @@ def test_marlin_prepare_layer_preserves_workspace_address(monkeypatch, variant):
 
     load_checkpoint_format_weights()
     prepare(layer)
-    workspace_ptr = layer.workspace.data_ptr()
+    assert not hasattr(layer, "workspace")
 
     # Reload: fresh checkpoint-format tensors, prepare runs again
     load_checkpoint_format_weights()
     prepare(layer)
 
-    assert layer.workspace.data_ptr() == workspace_ptr
-    assert torch.all(layer.workspace == 0)
+    assert not hasattr(layer, "workspace")
 
 
-def test_marlin_make_workspace_new_rejects_incompatible_existing(monkeypatch):
-    """An incompatible existing workspace means the address captured by CUDA
-    graphs is already unusable; allocating a replacement would hide that."""
+def test_marlin_workspace_uses_persistent_workspace_manager(monkeypatch):
+    """Calls reuse initialized locks; independent streams get separate storage."""
     from vllm.model_executor.layers.quantization.utils import marlin_utils
+    from vllm.utils import torch_utils
+    from vllm.v1.worker import workspace as workspace_module
 
     monkeypatch.setattr(marlin_utils, "num_compute_units", lambda _: 4)
     device = torch.device("cpu")
+    manager = workspace_module.WorkspaceManager(device)
+    monkeypatch.setattr(workspace_module, "_manager", manager)
+    stream = "main"
+    monkeypatch.setattr(torch_utils, "current_stream", lambda: stream)
 
-    workspace = marlin_utils.marlin_make_workspace_new(device)
-    reused = marlin_utils.marlin_make_workspace_new(device, existing=workspace)
-    assert reused is workspace
+    workspace = marlin_utils.get_marlin_workspace(device)
+    assert workspace.shape == (4 * marlin_utils.MARLIN_MAX_BLOCKS_PER_SM,)
+    assert workspace.dtype == torch.int32
+    assert torch.count_nonzero(workspace) == 0
 
-    with pytest.raises(ValueError, match="incompatible"):
-        marlin_utils.marlin_make_workspace_new(device, 4, existing=workspace)
-    with pytest.raises(ValueError, match="incompatible"):
-        marlin_utils.marlin_make_workspace_new(
-            device, existing=workspace.to(torch.int64)
-        )
+    workspace.fill_(1)
+    assert marlin_utils.get_marlin_workspace(device) is workspace
+    assert torch.all(workspace == 1)
+
+    stream = "aux"
+    aux_workspace = marlin_utils.get_marlin_workspace(device)
+    assert aux_workspace.data_ptr() != workspace.data_ptr()
+    assert torch.count_nonzero(aux_workspace) == 0
+
+    manager.lock()
+    assert marlin_utils.get_marlin_workspace(device) is aux_workspace
+
+
+def test_marlin_workspace_without_manager(monkeypatch):
+    from vllm.model_executor.layers.quantization.utils import marlin_utils
+    from vllm.v1.worker import workspace as workspace_module
+
+    monkeypatch.setattr(workspace_module, "_manager", None)
+    monkeypatch.setattr(marlin_utils, "num_compute_units", lambda _: 4)
+    first = marlin_utils.get_marlin_workspace(torch.device("cpu"))
+    second = marlin_utils.get_marlin_workspace(torch.device("cpu"))
+    assert first.data_ptr() != second.data_ptr()
+    assert first.shape == (4 * marlin_utils.MARLIN_MAX_BLOCKS_PER_SM,)
+    assert torch.count_nonzero(first) == 0
 
 
 def test_model_cleanup(dist_init, default_vllm_config):

tests/models/glm5next/test_sparse_indexer_topk_dispatch.py

@@ -1,11 +1,11 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Focused unit tests for the GLM decode top-k dispatch."""
+"""Tests for the GLM kpool indexer's top-k backend wiring."""
 
 import pytest
 import torch
 
-from vllm.models.glm5next.nvidia.sparse_indexer import _use_cooperative_topk
+from vllm.config import VllmConfig, set_current_vllm_config
 from vllm.platforms import current_platform
 
 pytestmark = pytest.mark.skipif(
@@ -13,29 +13,31 @@ pytestmark = pytest.mark.skipif(
 )
 
 
-@pytest.mark.parametrize("num_rows", [1, 64], ids=["1row", "64rows"])
-@torch.inference_mode()
-def test_cooperative_topk_is_selected_for_supported_batch(num_rows: int):
-    logits = torch.empty(num_rows, 8192, dtype=torch.float32, device="cuda")
-    assert _use_cooperative_topk(logits, 512, num_rows)
+def _require_deep_gemm() -> None:
+    from vllm.utils.deep_gemm import has_deep_gemm
+
+    if not has_deep_gemm():
+        pytest.skip("kpool indexer requires DeepGEMM")
 
 
-@pytest.mark.parametrize(
-    ("num_rows", "select_k", "stride_padding"),
-    [
-        (65, 512, 0),
-        (1, 256, 0),
-        (1, 512, 1),
-    ],
-    ids=["too_many_rows", "unsupported_k", "unaligned_stride"],
-)
-@torch.inference_mode()
-def test_cooperative_topk_falls_back_on_unsupported_batch(
-    num_rows: int, select_k: int, stride_padding: int
-):
-    logits = torch.empty(
-        num_rows, 8192 + stride_padding, dtype=torch.float32, device="cuda"
-    )
-    if stride_padding:
-        logits = logits[:, :8192]
-    assert not _use_cooperative_topk(logits, select_k, num_rows)
+@pytest.mark.parametrize("backend", ["auto", "persistent", "cooperative", "torch"])
+def test_kpool_indexer_dispatches_through_shared_topk_backend(backend: str) -> None:
+    """The kpool indexer must read kernel_config.sparse_indexer_topk_backend
+    and hand it to the shared SparseIndexerTopk dispatcher, rather than
+    hard-coding its own cooperative/persistent/per_row choice."""
+    _require_deep_gemm()
+    from vllm.models.glm5next.nvidia.sparse_indexer import SparseAttnIndexerKpool
+
+    cfg = VllmConfig(kernel_config={"sparse_indexer_topk_backend": backend})
+    with set_current_vllm_config(cfg):
+        op = SparseAttnIndexerKpool(
+            k_cache=None,
+            quant_block_size=128,
+            scale_fmt="ue8m0",
+            topk_tokens=2048,
+            head_dim=128,
+            max_pool_len=4096,
+            max_total_seq_len=8192,
+            topk_indices_buffer=torch.empty(8, 2176, dtype=torch.int32, device="cuda"),
+        )
+    assert op.topk_backend == backend

tests/models/multimodal/generation/test_vit_cudagraph.py

@@ -55,6 +55,10 @@ def kimi_vl_chat_template(content: str) -> str:
     )
 
 
+def deepseek_v41_chat_template(content: str) -> str:
+    return f"<｜begin▁of▁sentence｜><｜User｜>{content}<｜Assistant｜></think>"
+
+
 def step3_vl_chat_template(content: str) -> str:
     return (
         "<｜begin▁of▁sentence｜> You are a helpful assistant.<|BOT|>user\n "
@@ -221,6 +225,25 @@ MODEL_CONFIGS: dict[str, VitCudagraphTestConfig] = {
         vllm_runner_kwargs={"gpu_memory_utilization": 0.80},
         marks=[pytest.mark.core_model],
     ),
+    "deepseek_v41": VitCudagraphTestConfig(
+        model="deepseek-ai/DeepSeek-V4.1-Flash",
+        modalities=["image"],
+        image_prompt=deepseek_v41_chat_template(
+            "<｜deepseek_image｜>\n\nWhat is in this image?"
+        ),
+        compilation_config_overrides={
+            "encoder_cudagraph_token_budgets": [1024],
+        },
+        vllm_runner_kwargs={
+            "load_format": "dummy",
+            "attention_backend": "FLASHMLA_SPARSE_DSV41",
+            "hf_overrides": partial(
+                dummy_hf_overrides,
+                model_arch="DeepseekV41ForCausalLM",
+                exist_overrides={"vision_n_layers": 1},
+            ),
+        },
+    ),
     "step3_vl": VitCudagraphTestConfig(
         model="stepfun-ai/Step3-VL-10B",
         modalities=["image"],

tests/models/multimodal/processing/test_tensor_schema.py

@@ -218,8 +218,14 @@ def test_model_tensor_schema(model_id: str):
 
     factories = model_cls._processor_factory
 
+    # Capture helpers return capture-buffer containers (e.g.
+    # EncoderCudaGraphCaptureInputs), not TensorSchema mm inputs.
+    capture_helpers = {"prepare_encoder_cudagraph_capture_inputs"}
+
     inputs_parse_methods = []
     for attr_name in dir(model_cls):
+        if attr_name in capture_helpers:
+            continue
         attr = getattr(model_cls, attr_name)
         if hasattr(attr, "__annotations__"):
             return_type = attr.__annotations__.get("return", None)

tests/models/test_deepseek_v4_rocm_wo_a.py

@@ -0,0 +1,44 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+import torch
+
+from vllm.models.deepseek_v4.amd.rocm import _wo_a_block_scale_to_e8m0
+
+
+def test_wo_a_block_scale_to_e8m0_from_float():
+    scale = torch.tensor([[0.5, 1.0, 2.0, 4.0]], dtype=torch.float32)
+
+    encoded = _wo_a_block_scale_to_e8m0(scale)
+
+    assert encoded is not None
+    torch.testing.assert_close(
+        encoded,
+        torch.tensor([[126, 127, 128, 129]], dtype=torch.uint8),
+    )
+    assert encoded.is_contiguous()
+
+
+def test_wo_a_block_scale_to_e8m0_preserves_encoded_scales():
+    raw = torch.tensor([[125, 127, 131]], dtype=torch.uint8)
+    encoded = raw.view(torch.float8_e8m0fnu)
+
+    converted = _wo_a_block_scale_to_e8m0(encoded)
+
+    assert converted is not None
+    torch.testing.assert_close(converted, raw)
+
+
+@pytest.mark.parametrize(
+    "scale",
+    [
+        torch.tensor([[0.0, 1.0]]),
+        torch.tensor([[-1.0, 1.0]]),
+        torch.tensor([[0.75, 1.0]]),
+        torch.tensor([[float("inf"), 1.0]]),
+        torch.ones(1, dtype=torch.int32),
+    ],
+)
+def test_wo_a_block_scale_to_e8m0_rejects_invalid_scales(scale: torch.Tensor):
+    assert _wo_a_block_scale_to_e8m0(scale) is None

tests/renderers/test_truncate_after_mm_expansion.py

@@ -0,0 +1,166 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Truncation must bound the prompt *after* multimodal placeholders expand.
+
+``TokenizeParams.apply_post_tokenization`` truncates the unexpanded prompt, so
+one placeholder token can expand into many and push the result back over
+``truncate_prompt_tokens`` (vllm-project/vllm#57038).
+"""
+
+import numpy as np
+import pytest
+import torch
+
+from vllm.config import ModelConfig, VllmConfig
+from vllm.exceptions import VLLMValidationError
+from vllm.multimodal.inputs import PlaceholderRange
+from vllm.renderers import renderer_from_config
+from vllm.renderers.base import BaseRenderer
+from vllm.renderers.hf import _ensure_prompt_embeds_placeholder_token
+from vllm.renderers.params import TokenizeParams
+
+
+def _renderer():
+    """A BaseRenderer stub: only the tokenizer attribute is needed here."""
+
+    class _StubRenderer(BaseRenderer):
+        def __init__(self):
+            self.tokenizer = None
+
+        def render_messages(self, *args, **kwargs):
+            raise NotImplementedError
+
+    return _StubRenderer()
+
+
+def _mm_input(num_tokens: int, offset: int, length: int, **extra):
+    return {
+        "type": "multimodal",
+        "prompt_token_ids": list(range(num_tokens)),
+        "mm_kwargs": {},
+        "mm_hashes": {},
+        "mm_placeholders": {"audio": [PlaceholderRange(offset=offset, length=length)]},
+        **extra,
+    }
+
+
+def _params(limit: int | None, side: str = "right"):
+    return TokenizeParams(
+        max_total_tokens=4096,
+        truncate_prompt_tokens=limit,
+        truncation_side=side,
+    )
+
+
+def test_expanded_prompt_is_bounded():
+    """The reported case: audio fits, trailing text is cut to the limit."""
+    engine_input = _mm_input(num_tokens=147, offset=1, length=52)
+
+    _renderer()._truncate_expanded_prompt(engine_input, _params(96))
+
+    assert engine_input["prompt_token_ids"] == list(range(96))
+    assert engine_input["mm_placeholders"]["audio"][0] == PlaceholderRange(
+        offset=1, length=52
+    )
+
+
+def test_prompt_within_limit_is_untouched():
+    engine_input = _mm_input(num_tokens=64, offset=1, length=52)
+
+    _renderer()._truncate_expanded_prompt(engine_input, _params(96))
+
+    assert engine_input["prompt_token_ids"] == list(range(64))
+
+
+def test_split_multimodal_span_is_rejected():
+    """Cutting mid-span would desync features from token ids, so reject."""
+    engine_input = _mm_input(num_tokens=147, offset=80, length=50)
+
+    with pytest.raises(VLLMValidationError, match="truncate_prompt_tokens"):
+        _renderer()._truncate_expanded_prompt(engine_input, _params(96))
+
+
+def test_dropped_multimodal_span_is_rejected():
+    """A span entirely past the cut would be discarded silently."""
+    engine_input = _mm_input(num_tokens=147, offset=100, length=20)
+
+    with pytest.raises(VLLMValidationError, match="truncate_prompt_tokens"):
+        _renderer()._truncate_expanded_prompt(engine_input, _params(96))
+
+
+def test_left_truncation_shifts_placeholder_offsets():
+    engine_input = _mm_input(num_tokens=147, offset=100, length=20)
+
+    _renderer()._truncate_expanded_prompt(engine_input, _params(96, side="left"))
+
+    assert engine_input["prompt_token_ids"] == list(range(51, 147))
+    # The span started at 100 and 51 leading tokens were dropped.
+    assert engine_input["mm_placeholders"]["audio"][0] == PlaceholderRange(
+        offset=49, length=20
+    )
+
+
+def test_text_only_input_is_untouched():
+    """Only multimodal prompts grow after tokenization, so nothing else runs."""
+    engine_input = {
+        "type": "token",
+        "prompt_token_ids": list(range(147)),
+        "prompt_token_offsets": [(i, i + 1) for i in range(147)],
+    }
+
+    _renderer()._truncate_expanded_prompt(engine_input, _params(96))
+
+    assert engine_input["prompt_token_ids"] == list(range(147))
+    assert len(engine_input["prompt_token_offsets"]) == 147
+
+
+def test_no_truncation_requested_is_noop():
+    engine_input = _mm_input(num_tokens=147, offset=1, length=52)
+
+    _renderer()._truncate_expanded_prompt(engine_input, _params(None))
+
+    assert engine_input["prompt_token_ids"] == list(range(147))
+
+
+def test_mixed_prompt_embeds_truncation_covers_embeds_spans():
+    """Bound in _process_singleton, where prompt_embeds spans are already attached."""
+    model_config = ModelConfig(
+        "Qwen/Qwen3-ASR-0.6B-hf",
+        revision="7f1569a48a89f3e3f4dc3a5c9d28bddd903bc76c",
+        dtype="float32",
+        max_model_len=4096,
+        mm_processor_cache_gb=0,
+        enable_prompt_embeds=True,
+    )
+    renderer = renderer_from_config(VllmConfig(model_config=model_config))
+    tok = renderer.tokenizer
+    pe_id = _ensure_prompt_embeds_placeholder_token(tok)
+    hidden = model_config.get_hidden_size()
+    audio = np.zeros(4 * 16000, dtype=np.float32)
+
+    pe_len = 20
+    tensors = [torch.zeros(pe_len, hidden, dtype=torch.float32)]
+
+    def _build_prompt():
+        text = "<|audio_start|><|audio_pad|><|audio_end|>"
+        ids = tok.encode(text) + [pe_id] + tok.encode(" tail" * 40)
+        return {
+            "prompt_token_ids": ids,
+            "multi_modal_data": {"audio": (audio, 16000)},
+            "_prompt_embeds": (tensors, pe_id),
+        }
+
+    # A limit past every span (audio ends at 52, prompt_embeds at 74):
+    # truncation should succeed and bound the prompt.
+    out = renderer._process_singleton(
+        _build_prompt(),
+        tok_params=_params(79),
+    )
+    assert len(out["prompt_token_ids"]) == 79
+
+    # A limit inside the prompt_embeds span (54..74): must be rejected.
+    with pytest.raises(VLLMValidationError, match="truncate_prompt_tokens"):
+        renderer._process_singleton(
+            _build_prompt(),
+            tok_params=_params(64),
+        )

tests/utils.py

@@ -1427,8 +1427,10 @@ def init_test_distributed_environment(
         get_current_vllm_config_or_none,
         set_current_vllm_config,
     )
+    from vllm.platforms import current_platform
 
     distributed_init_method = f"tcp://localhost:{distributed_init_port}"
+    backend = current_platform.dist_backend
 
     if data_parallel_size > 1:
         # For DP we need to set a common DP master port
@@ -1449,6 +1451,7 @@ def init_test_distributed_environment(
                 rank=rank % tp_pp_world,
                 distributed_init_method=distributed_init_method,
                 local_rank=local_rank if local_rank >= 0 else rank,
+                backend=backend,
             )
             ensure_model_parallel_initialized(tp_size, pp_size)
         return
@@ -1460,6 +1463,7 @@ def init_test_distributed_environment(
             rank=rank,
             distributed_init_method=distributed_init_method,
             local_rank=local_rank,
+            backend=backend,
         )
         ensure_model_parallel_initialized(tp_size, pp_size)
     else:
@@ -1470,6 +1474,7 @@ def init_test_distributed_environment(
                 rank=rank,
                 distributed_init_method=distributed_init_method,
                 local_rank=local_rank,
+                backend=backend,
             )
             ensure_model_parallel_initialized(tp_size, pp_size)
 

tests/v1/attention/test_kpool_tail_slot_mapping.py

@@ -391,3 +391,66 @@ def test_interleaved_decode_pollution_legacy_vs_circular():
 
     # The circular mapping keeps the rings isolated under interleaving.
     torch.testing.assert_close(circular, ground_truth)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device")
+@pytest.mark.parametrize(
+    "per_req,num_actual,padded_len",
+    [
+        ([list(range(10)), list(range(12))], 22, 22),
+        ([list(range(10)), list(range(12))], 22, 30),
+        ([[3, 4], [0], [7, 8, 9]], 6, 8),
+        ([[5]], 1, 1),
+    ],
+)
+def test_triton_mapping_matches_cpu(per_req, num_actual, padded_len):
+    """The CUDA (Triton) path must match the CPU torch reference, including
+    tokens between the last request boundary and num_actual_tokens (mapped to
+    the last request) and untouched padding beyond num_actual."""
+    positions, qsl, slot_mapping, _, num_reqs = make_batch(
+        per_req, padded_len=padded_len
+    )
+    # Replace the all--1 placeholder slots with sentinel values to check the
+    # padding range is copied through untouched.
+    slot_mapping = torch.arange(padded_len, dtype=torch.int64) + 1000
+    bt = make_tail_block_table(list(range(5, 5 + num_reqs)))
+
+    ref = circular_tail_slots(slot_mapping, bt, qsl, positions, num_actual, num_reqs)
+    got = circular_tail_slots(
+        slot_mapping.cuda(),
+        bt.cuda(),
+        qsl.cuda().to(torch.int32),
+        positions.cuda(),
+        num_actual,
+        num_reqs,
+    )
+    torch.testing.assert_close(got.cpu(), ref)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device")
+def test_triton_mapping_reads_strided_block_table():
+    """The kernel must address the block table through its real stride(0),
+    not a dense assumption (the #57477 class of bug): a tail block table that
+    is a column view of a wider allocation must still read column 0."""
+    per_req = [list(range(10)), list(range(12))]
+    positions, qsl, slot_mapping, num_actual, num_reqs = make_batch(per_req)
+    own_blocks = [11, 22]
+    wide = torch.zeros(num_reqs, 8, dtype=torch.int32)
+    bt = wide[:, 1:4]  # non-contiguous view, stride(0) == 8
+    bt[:, 0] = torch.tensor(own_blocks, dtype=torch.int32)
+
+    ref = circular_tail_slots(
+        slot_mapping, bt.contiguous(), qsl, positions, num_actual, num_reqs
+    )
+    got = circular_tail_slots(
+        slot_mapping.cuda(),
+        bt.cuda(),
+        qsl.cuda().to(torch.int32),
+        positions.cuda(),
+        num_actual,
+        num_reqs,
+    )
+    torch.testing.assert_close(got.cpu(), ref)
+    for req, blk in enumerate(own_blocks):
+        start, end = int(qsl[req]), int(qsl[req + 1])
+        assert (got[start:end] // KPOOL == blk).all()

tests/v1/core/prefix_cache/test_hybrid_prefix_cache_hit_rate.py

@@ -4,6 +4,7 @@
 
 import asyncio
 import json
+from math import lcm
 from uuid import uuid4
 
 import httpx
@@ -12,6 +13,7 @@ from openai.types.chat import ChatCompletionMessageParam
 from prometheus_client.parser import text_string_to_metric_families
 
 from tests.utils import RemoteOpenAIServer
+from vllm.utils.math_utils import round_down
 
 MODEL = "Qwen/Qwen3.5-0.8B"
 NUM_CONVERSATIONS = 8
@@ -34,6 +36,29 @@ def _user_turns(conversation: int) -> list[str]:
     ]
 
 
+async def _scheduler_block_size(client: httpx.AsyncClient) -> int:
+    """Return the coarsest boundary an expected cache hit can land on.
+
+    Mirrors ``scheduler_block_size`` in ``resolve_kv_cache_block_sizes``, which
+    the engine keeps internal: the LCM of the resolved group block sizes. A
+    finer ``prefix_match_unit`` only adds hits, so this stays conservative.
+    """
+    response = await client.get("/metrics")
+    response.raise_for_status()
+    for family in text_string_to_metric_families(response.text):
+        for sample in family.samples:
+            if sample.name != "vllm:cache_config_info":
+                continue
+            sizes = [
+                int(value)
+                for name in ("block_size", "mamba_block_size")
+                if (value := sample.labels.get(name, "None")) != "None"
+            ]
+            assert sizes, f"no block size in {sample.labels}"
+            return lcm(*sizes)
+    raise AssertionError("missing vllm:cache_config_info metric")
+
+
 async def _prefix_cache_counters(
     client: httpx.AsyncClient,
 ) -> tuple[float, float]:
@@ -83,13 +108,13 @@ def server(request):
 
 @pytest.mark.asyncio
 @pytest.mark.parametrize(
-    "server, min_hit_rate",
-    # MTP re-prefills an extra block; both floors reject one more missed block.
-    [pytest.param(False, 0.90, id="base"), pytest.param(True, 0.75, id="mtp")],
+    "server, dropped_blocks",
+    # MTP's prefill lookahead pollutes the tail block, so EAGLE drops one.
+    [pytest.param(False, 0, id="base"), pytest.param(True, 1, id="mtp")],
     indirect=["server"],
 )
 async def test_prefix_cache_hit_rate(
-    server: RemoteOpenAIServer, min_hit_rate: float
+    server: RemoteOpenAIServer, dropped_blocks: int
 ) -> None:
     # Isolate sessions so another conversation cannot hide a cache miss.
     salts = [uuid4().hex for _ in range(NUM_CONVERSATIONS)]
@@ -102,6 +127,8 @@ async def test_prefix_cache_hit_rate(
         server.get_async_client() as client,
         httpx.AsyncClient(base_url=server.url_root, timeout=METRICS_TIMEOUT) as metrics,
     ):
+        block_size = await _scheduler_block_size(metrics)
+        print(f"{block_size=}")
 
         async def run_conversation(session: int, *, replay: bool) -> tuple[int, int]:
             history = histories[session]
@@ -126,15 +153,30 @@ async def test_prefix_cache_hit_rate(
                 assert cached is not None, context
                 assert 0 <= cached < prompt, context
                 assert prompt >= MIN_PROMPT_TOKENS, context
-                if not replay and turn == 0:
-                    assert cached == 0, context
-                else:
-                    assert cached / prompt >= min_hit_rate, (
-                        f"{context}, {min_hit_rate=}"
-                    )
                 if replay:
                     assert prompt == prompt_lengths[session][turn], context
+                    # The identical request computed this whole prompt, but at
+                    # least one token is always recomputed.
+                    computed = prompt - 1
+                elif turn:
+                    computed = prompt_lengths[session][turn - 1]
                 else:
+                    assert cached == 0, context
+                    computed = 0
+                # Every token below the last reachable boundary must be reused.
+                floor = max(
+                    0, round_down(computed, block_size) - dropped_blocks * block_size
+                )
+                print(
+                    f"{session=} {turn=} {replay=} {prompt=} {cached=} {floor=} "
+                    f"hit_rate={cached / prompt:.4f}"
+                )
+                assert cached >= floor, (
+                    f"{context}, {computed=}, {block_size=}, "
+                    f"hit_rate={cached / prompt:.4f}, "
+                    f"expected_hit_rate={floor / prompt:.4f}"
+                )
+                if not replay:
                     prompt_lengths[session].append(prompt)
                     history.append(
                         {

tests/v1/ec_connector/unit/test_epd_proxy.py

@@ -0,0 +1,279 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Routing and registration behaviour of the EPD proxy."""
+
+import asyncio
+from typing import Any
+from unittest.mock import AsyncMock, Mock
+
+import msgspec
+import pytest
+from fastapi import HTTPException
+from fastapi.testclient import TestClient
+
+from examples.disaggregated.disaggregated_encoder import (
+    disagg_epd_proxy as proxy_module,
+)
+from examples.disaggregated.disaggregated_encoder.disagg_epd_proxy import (
+    EPDProxy,
+    EPDProxyConfig,
+    InstanceRecord,
+    InstanceRegistry,
+    InstanceRole,
+    build_app,
+    extract_mm_items,
+)
+
+ENCODE = InstanceRole.ENCODE
+PREFILL = InstanceRole.PREFILL
+DECODE = InstanceRole.DECODE
+PD = InstanceRole.PREFILL_DECODE
+
+IMAGE_ITEM = {"type": "image_url", "image_url": {"url": "http://img/0.png"}}
+
+
+@pytest.fixture
+def proxy():
+    registry = InstanceRegistry(probe_interval=0)
+    return EPDProxy(registry)
+
+
+class TestRouting:
+    def test_no_decode_instance_is_service_unavailable(self, proxy):
+        """The proxy comes up before anything registers, so this is normal."""
+        with pytest.raises(HTTPException) as excinfo:
+            proxy.route(num_items=0)
+        assert excinfo.value.status_code == 503
+
+    def test_media_without_an_encoder_is_service_unavailable(self, proxy):
+        proxy.registry.register(InstanceRecord(PD, "http://d0:8000"))
+        assert proxy.route(num_items=0).decode.url == "http://d0:8000"
+        with pytest.raises(HTTPException) as excinfo:
+            proxy.route(num_items=1)
+        assert excinfo.value.status_code == 503
+
+    def test_prefill_is_optional(self, proxy):
+        """An E+PD deployment registers no prefill instance at all."""
+        proxy.registry.register(InstanceRecord(PD, "http://d0:8000"))
+        assert proxy.route(num_items=0).prefill is None
+
+    def test_encoder_roster_tracks_new_instances(self, proxy):
+        proxy.registry.register(InstanceRecord(PD, "http://d0:8000"))
+        for index in range(2):
+            proxy.registry.register(InstanceRecord(ENCODE, f"http://e{index}:8000"))
+        assert proxy.route(num_items=3).encoder_urls == [
+            "http://e0:8000",
+            "http://e1:8000",
+        ]
+        proxy.registry.unregister("http://e0:8000")
+        assert proxy.route(num_items=3).encoder_urls == ["http://e1:8000"]
+
+    @pytest.mark.parametrize("unregister", [False, True])
+    def test_standalone_decode_always_requires_healthy_prefill(self, unregister):
+        app = build_app()
+        proxy = app.state.proxy
+        proxy.registry.register(InstanceRecord(DECODE, "http://d0:8000"))
+        with pytest.raises(HTTPException) as excinfo:
+            proxy.route(num_items=0)
+        assert excinfo.value.status_code == 503
+        proxy.registry.register(InstanceRecord(PREFILL, "http://p0:8000"))
+        route = proxy.route(num_items=0)
+        assert route.decode.url == "http://d0:8000"
+        assert route.consumer.url == "http://p0:8000"
+        if unregister:
+            proxy.registry.unregister("http://p0:8000")
+        else:
+            for _ in range(proxy.registry._fail_threshold):
+                proxy.registry._on_probe_failure(route.prefill, 0)
+        with pytest.raises(HTTPException) as excinfo:
+            proxy.route(num_items=0)
+        assert excinfo.value.status_code == 503
+
+
+class TestConsumerAddress:
+    """Which stage receives the embedding depends on the topology."""
+
+    def test_shared_storage_connectors_name_no_target(self, proxy):
+        """Rotate replicas without requiring a push endpoint."""
+        proxy.registry.register(InstanceRecord(PD, "http://d0:8000", dp_size=2))
+        route = proxy.route(num_items=0)
+        assert route.consumer_zmq is None
+        assert route.dp_rank == 0
+        assert proxy.route(num_items=0).dp_rank == 1
+
+    def test_decode_is_the_consumer_when_prefill_is_not_split_out(self, proxy):
+        proxy.registry.register(
+            InstanceRecord(PD, "http://d0:8000", ec_zmq_addrs=["tcp://d0:20001"])
+        )
+        assert proxy.route(num_items=0).consumer_zmq == "tcp://d0:20001"
+
+    def test_prefill_is_the_consumer_when_it_is_split_out(self, proxy):
+        """In E+P+D the prefill instance consumes the embedding, not decode.
+
+        Reading the address out of whichever list it came from -- rather than
+        from the record that reported one -- sends the push to the wrong
+        instance in this topology.
+        """
+        proxy.registry.register(
+            InstanceRecord(PREFILL, "http://p0:8000", ec_zmq_addrs=["tcp://p0:20001"])
+        )
+        proxy.registry.register(InstanceRecord(DECODE, "http://d0:8000"))
+        assert proxy.route(num_items=0).consumer_zmq == "tcp://p0:20001"
+
+
+class TestProxyApi:
+    @pytest.fixture
+    def client(self, monkeypatch):
+        monkeypatch.setenv("ADMIN_API_KEY", "test-key")
+        config = EPDProxyConfig(probe_interval=0)
+        with TestClient(build_app(config)) as client:
+            client.headers["X-API-Key"] = "test-key"
+            yield client
+
+    def test_proxy_serves_before_anything_registers(self, client):
+        assert client.get("/health").status_code == 200
+        assert client.get("/instances").json()["encode"]["live"] == []
+
+    def test_requests_are_refused_until_a_decoder_registers(self, client):
+        response = client.post(
+            "/v1/chat/completions", json={"model": "m", "messages": []}
+        )
+        assert response.status_code == 503
+        assert client.get("/v1/models").status_code == 503
+
+    def test_http_registration_keeps_consumer_replica_and_address_together(
+        self, client
+    ):
+        payload: dict[str, Any] = {
+            "role": "prefill_decode",
+            "url": "http://d0:8000/",
+            "dp_size": 2,
+            "ec_zmq_addrs": ["tcp://d0:20001", "tcp://d0:20003"],
+        }
+        assert client.post("/instances", json=payload).json() == {"registered": True}
+        assert client.post("/instances", json=payload).json() == {"registered": False}
+        proxy = client.app.state.proxy
+        for rank in [0, 1, 0]:
+            route = proxy.route(0)
+            assert route.dp_rank == rank
+            assert route.consumer_zmq == payload["ec_zmq_addrs"][rank]
+        assert client.delete("/instances", params={"url": payload["url"]}).json() == {
+            "removed": True
+        }
+        assert proxy.registry.pick(PD) is None
+
+    def test_registration_rejects_invalid_topology(self, client):
+        cases: list[tuple[dict[str, Any], int]] = [
+            ({"dp_size": 0}, 422),
+            ({"dp_size": 2, "ec_zmq_addrs": ["tcp://d0:20001"]}, 422),
+            ({"role": "decode", "ec_zmq_addrs": ["tcp://d0:20001"]}, 400),
+            ({"role": "encode", "ec_zmq_addrs": ["tcp://d0:20001"]}, 400),
+        ]
+        for fields, status in cases:
+            response = client.post(
+                "/instances",
+                json={"url": "http://d0:8000", "role": "prefill_decode", **fields},
+            )
+            assert response.status_code == status
+
+    def test_registration_requires_admin_key(self, client):
+        client.headers.pop("X-API-Key")
+        assert (
+            client.post(
+                "/instances", json={"role": "encode", "url": "http://e0:8000"}
+            ).status_code
+            == 403
+        )
+        assert (
+            client.delete("/instances", params={"url": "http://e0:8000"}).status_code
+            == 403
+        )
+
+    @pytest.mark.parametrize("split_role", ["prefill", "decode"])
+    @pytest.mark.parametrize("combined_first", [False, True])
+    def test_registration_rejects_mixed_topologies(
+        self, client, split_role, combined_first
+    ):
+        roles = ["prefill_decode", split_role]
+        if not combined_first:
+            roles.reverse()
+        first = {"role": roles[0], "url": "http://first:8000"}
+        second = {"role": roles[1], "url": "http://second:8000"}
+        assert client.post("/instances", json=first).status_code == 200
+        assert client.post("/instances", json=second).status_code == 409
+        # Health eviction must not silently permit a topology switch.
+        registry = client.app.state.registry
+        record = registry._live[first["url"]]
+        for _ in range(registry._fail_threshold):
+            registry._on_probe_failure(record, 0)
+        assert client.post("/instances", json=second).status_code == 409
+        assert client.delete("/instances", params={"url": first["url"]}).json() == {
+            "removed": True
+        }
+        assert client.post("/instances", json=second).status_code == 200
+
+
+def test_extract_mm_items_finds_media_across_messages():
+    req = {
+        "messages": [
+            {"role": "user", "content": "plain text"},
+            {"role": "user", "content": [IMAGE_ITEM, {"type": "text", "text": "hi"}]},
+            {"role": "user", "content": [IMAGE_ITEM]},
+        ]
+    }
+    assert extract_mm_items(req) == [IMAGE_ITEM, IMAGE_ITEM]
+
+
+@pytest.mark.parametrize("push", [False, True])
+def test_encoder_handles_and_json_metadata_survive_rewrite(proxy, push, monkeypatch):
+    """Keep main's hash-keyed handles and per-item push IDs after the proxy move."""
+    proxy.registry.register(InstanceRecord(ENCODE, "http://encoder"))
+    proxy.registry.register(
+        InstanceRecord(
+            PD,
+            "http://decode",
+            ec_zmq_addrs=["tcp://decode:14579"] if push else [],
+        )
+    )
+    handle = {
+        "metadata": {"image_grid_thw": [[1, 2, 3]]},
+        "item_indices": [0, 1],
+        "peer_port": 5601,
+    }
+    response = Mock(
+        status=200,
+        read=AsyncMock(
+            return_value=msgspec.json.encode(
+                {"ec_transfer_params": {"engine-hash": handle}}
+            )
+        ),
+    )
+    session = Mock(post=AsyncMock(return_value=response))
+    monkeypatch.setattr(proxy_module, "encode_session", session)
+    original = {"messages": [{"role": "user", "content": [IMAGE_ITEM, IMAGE_ITEM]}]}
+    route = proxy.route(2)
+    prepared, _, _ = asyncio.run(
+        proxy_module.prepare_for_decode(
+            original, "request", route.encoder_urls, None, route.consumer_zmq
+        )
+    )
+    assert "ec_transfer_params" not in original
+    assert prepared["ec_transfer_params"]["engine-hash"] == handle
+    items = prepared["ec_transfer_params"]["ec_items"]
+    assert [item["mm_hash"] for item in items] == ["engine-hash", "engine-hash"]
+    assert items[0]["transfer_id"] != items[1]["transfer_id"]
+    assert session.post.call_count == 1
+    sent = msgspec.json.decode(session.post.call_args.kwargs["data"]).get(
+        "ec_transfer_params"
+    )
+    if push:
+        assert sent == {
+            "consumer_zmq": "tcp://decode:14579",
+            "ec_items": [{"transfer_id": item["transfer_id"]} for item in items],
+        }
+    else:
+        assert sent is None
+    for index in range(2):
+        assert prepared["messages"][0]["content"][index]["image_embeds"] == {
+            "image_grid_thw": [1, 2, 3]
+        }

tests/v1/ec_connector/unit/test_epd_proxy_retry.py

@@ -8,6 +8,7 @@ change to them is what these tests catch.
 
 import asyncio
 import importlib.util
+import sys
 from pathlib import Path
 from types import SimpleNamespace
 from typing import Any
@@ -27,6 +28,7 @@ def proxy():
     spec = importlib.util.spec_from_file_location("disagg_epd_proxy_retry", path)
     assert spec is not None and spec.loader is not None
     module = importlib.util.module_from_spec(spec)
+    sys.modules[spec.name] = module
     spec.loader.exec_module(module)
     return module
 
@@ -48,7 +50,7 @@ def test_maybe_prefill_leaves_the_caller_body_untouched(proxy, monkeypatch):
     """
     served = [{"remote_block_ids": [1, 2]}, {}]
 
-    async def _stage(req_data, p_url, req_id):
+    async def _stage(req_data, p_url, req_id, dp_rank=None):
         assert "kv_transfer_params" not in req_data
         return _Response(served.pop(0))
 
@@ -222,7 +224,7 @@ def test_a_decode_retry_does_not_inherit_the_previous_handles(proxy, monkeypatch
         _EncoderSession([{"encoder-side-hash": handle}, {}]),
     )
 
-    async def _no_prefill(req_data, p_url, req_id):
+    async def _no_prefill(req_data, p_url, req_id, dp_rank=None):
         return req_data
 
     monkeypatch.setattr(proxy, "maybe_prefill", _no_prefill)
@@ -280,6 +282,7 @@ def test_raw_media_keeps_encoder_transfer_identity(proxy, monkeypatch):
 @pytest.mark.asyncio
 @pytest.mark.parametrize("stream", [False, True])
 @pytest.mark.parametrize("prefill", [False, True])
+@pytest.mark.parametrize("dynamic", [False, True])
 @pytest.mark.parametrize(
     "request_options",
     [
@@ -293,7 +296,7 @@ def test_raw_media_keeps_encoder_transfer_identity(proxy, monkeypatch):
     ],
 )
 async def test_http_roundtrip_preserves_payload_and_response_bytes(
-    proxy, monkeypatch, stream, prefill, request_options
+    proxy, monkeypatch, stream, prefill, dynamic, request_options
 ):
     """Preserve media and scheduling context across HTTP hops and retries."""
     seen: dict[str, list[dict]] = {"encode": [], "prefill": [], "decode": []}
@@ -307,6 +310,10 @@ async def test_http_roundtrip_preserves_payload_and_response_bytes(
 
     async def backend(request):
         stage = request.match_info["stage"]
+        if dynamic:
+            consumer = "prefill" if prefill else "decode"
+            expected_rank = "0" if stage == consumer else None
+            assert request.headers.get("X-data-parallel-rank") == expected_rank
         assert request.content_type == "application/json"
         body = await request.json()
         seen[stage].append(body)
@@ -354,6 +361,10 @@ async def test_http_roundtrip_preserves_payload_and_response_bytes(
         monkeypatch.setattr(proxy.app.state, "ec_consumer_dp_size", 1, raising=False)
         monkeypatch.setattr(proxy, "DECODE_RETRIES", 1)
         monkeypatch.setattr(proxy, "NO_REWRITE", False)
+        target_app = proxy.app
+        if dynamic:
+            monkeypatch.setenv("ADMIN_API_KEY", "test-key")
+            target_app = proxy.build_app(proxy.EPDProxyConfig(probe_interval=0))
         item = {"type": "image_url", "image_url": {"url": "data:image/png;base64,YWJj"}}
         body = {
             "model": "test",
@@ -377,9 +388,37 @@ async def test_http_roundtrip_preserves_payload_and_response_bytes(
         await proxy.on_startup()
         try:
             async with httpx.AsyncClient(
-                transport=httpx.ASGITransport(app=proxy.app), base_url="http://proxy"
+                transport=httpx.ASGITransport(app=target_app), base_url="http://proxy"
             ) as client:
+                if dynamic:
+                    registrations = [("encode", "encode")]
+                    registrations += (
+                        [("prefill", "prefill"), ("decode", "decode")]
+                        if prefill
+                        else [("prefill_decode", "decode")]
+                    )
+                    for role, stage in registrations:
+                        consumer = "prefill" if prefill else "decode"
+                        registered = await client.post(
+                            "/instances",
+                            headers={"X-API-Key": "test-key"},
+                            json={
+                                "role": role,
+                                "url": str(server.make_url(f"/{stage}")),
+                                "dp_size": 2 if stage == consumer else 1,
+                            },
+                        )
+                        assert registered.status_code == 200
                 response = await client.post("/v1/chat/completions", json=body)
+                if dynamic:
+                    removed = await client.delete(
+                        "/instances",
+                        headers={"X-API-Key": "test-key"},
+                        params={"url": str(server.make_url("/decode"))},
+                    )
+                    assert removed.json() == {"removed": True}
+                    unavailable = await client.post("/v1/chat/completions", json=body)
+                    assert unavailable.status_code == 503
             assert response.status_code == 200
             assert response.content == expected
             assert response.headers["content-type"].startswith(

tests/v1/ec_connector/unit/test_epd_proxy_round_robin.py

@@ -14,6 +14,7 @@ across calls so consecutive requests keep advancing through every URL.
 """
 
 import importlib.util
+import sys
 from collections import Counter
 from pathlib import Path
 
@@ -27,6 +28,7 @@ def _load_proxy_module():
     spec = importlib.util.spec_from_file_location("disagg_epd_proxy_under_test", path)
     assert spec is not None and spec.loader is not None
     module = importlib.util.module_from_spec(spec)
+    sys.modules[spec.name] = module
     spec.loader.exec_module(module)
     return module
 

tests/v1/ec_connector/unit/test_epd_registry.py

@@ -0,0 +1,156 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Dynamic registration and liveness for the EPD proxy's instance registry."""
+
+from unittest.mock import patch
+
+import pytest
+
+from examples.disaggregated.disaggregated_encoder.disagg_epd_proxy import (
+    InstanceRecord,
+    InstanceRegistry,
+    InstanceRole,
+)
+
+ENCODE = InstanceRole.ENCODE
+DECODE = InstanceRole.DECODE
+
+
+@pytest.fixture
+def registry():
+    return InstanceRegistry(probe_interval=0, fail_threshold=3, evicted_ttl=60.0)
+
+
+def _fill(registry, role, count, prefix="http://h"):
+    for index in range(count):
+        registry.register(InstanceRecord(role, f"{prefix}{index}:8000"))
+
+
+async def _probe_round(registry, healthy: set[str], now: float = 0.0):
+    """Run one probe pass where only `healthy` URLs answer."""
+
+    async def fake_probe(self, session, url):
+        return url in healthy
+
+    with (
+        patch.object(InstanceRegistry, "_probe", fake_probe),
+        patch(
+            "examples.disaggregated.disaggregated_encoder.disagg_epd_proxy.time.monotonic",
+            return_value=now,
+        ),
+    ):
+        await registry._probe_once(session=None)
+
+
+class TestRegistration:
+    def test_registering_is_idempotent_and_reports_novelty(self, registry):
+        record = InstanceRecord(ENCODE, "http://e0:8000")
+        assert registry.register(record) is True
+        assert registry.register(record) is False
+        assert registry.urls(ENCODE) == ["http://e0:8000"]
+
+    def test_transfer_addresses_stay_with_the_registered_consumer(self, registry):
+        """Roster updates must not shift a consumer's transfer addresses."""
+        registry.register(InstanceRecord(DECODE, "http://d0:8000"))
+        registry.register(
+            InstanceRecord(
+                DECODE,
+                "http://d0:8000",
+                ec_zmq_addrs=["tcp://d0:20001", "tcp://d0:20002"],
+                dp_size=2,
+            )
+        )
+        picked = registry.pick(DECODE)
+        assert picked.ec_zmq_addrs == ["tcp://d0:20001", "tcp://d0:20002"]
+        assert picked.dp_size == 2
+
+    def test_unregister_is_terminal(self, registry):
+        registry.register(InstanceRecord(ENCODE, "http://e0:8000"))
+        assert registry.unregister("http://e0:8000") is True
+        assert registry.unregister("http://e0:8000") is False
+        assert registry.pick(ENCODE) is None
+
+    def test_round_robin_survives_registration(self, registry):
+        _fill(registry, ENCODE, 2)
+        assert [record.url for record in registry.pick_many(ENCODE, 2)] == [
+            "http://h0:8000",
+            "http://h1:8000",
+        ]
+        registry.register(InstanceRecord(ENCODE, "http://h2:8000"))
+        assert registry.pick(ENCODE).url == "http://h2:8000"
+
+
+class TestLiveness:
+    @pytest.mark.asyncio
+    async def test_identical_reregistration_preserves_inflight_probe_results(
+        self, registry
+    ):
+        """Repeated registrations must not suppress eviction or recovery."""
+        url = "http://e0:8000"
+        registry.register(InstanceRecord(ENCODE, url))
+        healthy = False
+
+        async def probe(self, session, url):
+            registry.register(InstanceRecord(ENCODE, url))
+            return healthy
+
+        with patch.object(InstanceRegistry, "_probe", probe):
+            for _ in range(registry._fail_threshold):
+                await registry._probe_once(None)
+            assert registry.status()["encode"] == {"live": [], "evicted": [url]}
+
+            healthy = True
+            await registry._probe_once(None)
+            assert registry.status()["encode"] == {"live": [url], "evicted": []}
+
+    @pytest.mark.asyncio
+    async def test_consecutive_failures_stop_routing(self, registry):
+        registry.register(InstanceRecord(ENCODE, "http://e0:8000"))
+        registry.register(InstanceRecord(ENCODE, "http://e1:8000"))
+        for _ in range(registry._fail_threshold):
+            await _probe_round(registry, healthy={"http://e1:8000"})
+        assert registry.urls(ENCODE) == ["http://e1:8000"]
+        assert registry.status()["encode"]["evicted"] == ["http://e0:8000"]
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize("role", [ENCODE, DECODE])
+    async def test_a_recovered_instance_rejoins_without_re_registering(
+        self, registry, role
+    ):
+        """Restarting every encoder to recover from a blip is not acceptable."""
+        registry.register(InstanceRecord(role, "http://e0:8000"))
+        for _ in range(registry._fail_threshold):
+            await _probe_round(registry, healthy=set())
+        assert registry.urls(role) == []
+
+        await _probe_round(registry, healthy={"http://e0:8000"}, now=1)
+        assert registry.urls(role) == ["http://e0:8000"]
+
+    @pytest.mark.asyncio
+    async def test_re_registration_does_not_override_failed_health_checks(
+        self, registry
+    ):
+        record = InstanceRecord(ENCODE, "http://e0:8000")
+        registry.register(record)
+        for _ in range(registry._fail_threshold):
+            await _probe_round(registry, healthy=set())
+        registry.register(record)
+        assert registry.urls(ENCODE) == []
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize("healthy", [False, True])
+    async def test_inflight_probe_cannot_restore_unregistered_instance(
+        self, registry, healthy
+    ):
+        record = InstanceRecord(ENCODE, "http://e0:8000")
+        registry.register(record)
+        for _ in range(registry._fail_threshold):
+            await _probe_round(registry, healthy=set())
+
+        async def probe(self, session, url):
+            registry.unregister(url)
+            return healthy
+
+        with patch.object(InstanceRegistry, "_probe", probe):
+            await registry._probe_once(None)
+        assert registry.status()["encode"] == {"live": [], "evicted": []}

tests/v1/simple_kv_offload/test_kv_events.py

@@ -19,8 +19,10 @@ from tests.v1.simple_kv_offload.test_scheduler import (
     SchedulerFixture,
     _alloc_and_register,
     _allocate_cp_gpu_blocks,
+    _allocate_gpu_blocks,
     _make_cp_request,
     _make_kv_cache_config,
+    _make_scratch_kv_cache_config,
     _make_vllm_config,
     make_request,
     make_scheduler_output,
@@ -757,3 +759,37 @@ def test_mamba_align_skips_positional_event_metadata() -> None:
             get_block_hash(make_block_hash_with_group_id(req.block_hashes[idx], 0))
         )
         assert ev.block_hashes == [expected_hash]
+
+
+def test_scratch_group_store_emits_events_for_cacheable_groups_only() -> None:
+    """Event metadata is resolved only for prefix-cacheable groups.
+
+    A scratch group (GLM-5.3-Flash kpool tail) has a block size that does not
+    divide the hash block size, so viewing the request hashes at its block
+    size is impossible; the store must skip it and still emit BlockStored
+    for the attention blocks.
+    """
+    fix = make_events_scheduler(kv_cache_config=_make_scratch_kv_cache_config(16))
+    sched = fix.scheduler
+    gpu_pool = fix.gpu_block_pool
+    num_blocks = 2
+
+    req = make_request(num_blocks=num_blocks)
+    fa_blocks = _allocate_gpu_blocks(gpu_pool, req, num_blocks, group_id=0)
+    scratch_block = gpu_pool.get_new_blocks(1)
+    kv_blocks = KVCacheBlocks(blocks=(fa_blocks, scratch_block))
+    req.num_computed_tokens = num_blocks * BLOCK_SIZE
+    sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0)
+    sched_out = make_scheduler_output(
+        {req.request_id: num_blocks * BLOCK_SIZE},
+        new_reqs={req.request_id: kv_blocks.get_block_ids()},
+    )
+    meta = sched.build_connector_meta(sched_out)
+    assert sorted(meta.store_gpu_blocks) == sorted(b.block_id for b in fa_blocks)
+    simulate_store_completion(sched, meta.store_event)
+
+    events = list(sched.take_events())
+    assert len(events) == num_blocks
+    assert all(isinstance(e, BlockStored) for e in events)
+    assert {e.group_idx for e in events} == {0}
+    assert all(len(e.token_ids) == BLOCK_SIZE for e in events)

tests/v1/simple_kv_offload/test_scheduler.py

@@ -41,7 +41,9 @@ from vllm.v1.core.single_type_kv_cache_manager import (
     register_all_kvcache_specs,
 )
 from vllm.v1.kv_cache_interface import (
+    CircularBufferSpec,
     FullAttentionSpec,
+    KpoolTailSpec,
     KVCacheConfig,
     KVCacheGroupSpec,
     KVCacheTensor,
@@ -125,6 +127,45 @@ def _make_kv_cache_config(
     )
 
 
+def _make_scratch_kv_cache_config(
+    num_blocks: int, scratch_block_size: int = 4
+) -> KVCacheConfig:
+    """FullAttention group plus a non-prefix-cacheable scratch group.
+
+    Mirrors GLM-5.3-Flash, whose kpool-tail group holds one per-request
+    scratch block of ``index_kpool`` tokens: it is excluded from prefix
+    caching, so its block size neither divides nor is divided by the
+    hash block size in general.
+    """
+    register_all_kvcache_specs(vllm_config=None)
+    fa_config = _make_kv_cache_config(num_blocks, num_groups=1)
+    scratch_layers = ["layer_scratch"]
+    scratch_spec = KpoolTailSpec(
+        block_size=scratch_block_size,
+        num_kv_heads=2,
+        head_size=HEAD_SIZE,
+        head_size_v=0,
+        dtype=DTYPE,
+        sliding_window=scratch_block_size,
+    )
+    assert not scratch_spec.prefix_cacheable
+    scratch_bytes = scratch_spec.page_size_bytes * num_blocks
+    return KVCacheConfig(
+        num_blocks=num_blocks,
+        kv_cache_tensors=fa_config.kv_cache_tensors
+        + [
+            KVCacheTensor(
+                size=scratch_bytes,
+                layers=scratch_layers,
+                layer_stride=scratch_bytes,
+                block_stride=scratch_spec.page_size_bytes,
+            )
+        ],
+        kv_cache_groups=fa_config.kv_cache_groups
+        + [KVCacheGroupSpec(scratch_layers, scratch_spec)],
+    )
+
+
 def _make_vllm_config(block_size: int = BLOCK_SIZE) -> VllmConfig:
     """Minimal VllmConfig for scheduler tests (no GPU)."""
     model_config = ModelConfig(
@@ -174,9 +215,13 @@ def make_scheduler(
     num_gpu_blocks: int = 16,
     num_groups: int = 1,
     lazy: bool = False,
+    kv_cache_config: KVCacheConfig | None = None,
 ) -> SchedulerFixture:
     """Build a SimpleCPUOffloadScheduler with small block pools."""
-    kv_cache_config = _make_kv_cache_config(num_gpu_blocks, num_groups)
+    if kv_cache_config is None:
+        kv_cache_config = _make_kv_cache_config(num_gpu_blocks, num_groups)
+    else:
+        num_groups = len(kv_cache_config.kv_cache_groups)
     vllm_config = _make_vllm_config()
     cpu_capacity_bytes = _BYTES_PER_BLOCK * num_cpu_blocks * num_groups
 
@@ -1392,6 +1437,63 @@ def test_multi_group_null_blocks_skipped() -> None:
     )
 
 
+# ---------------------------------------------------------------------------
+# Test 8b: Non-prefix-cacheable scratch groups take no part in store or load
+# ---------------------------------------------------------------------------
+def test_scratch_group_excluded_from_store_and_load() -> None:
+    """A scratch group (GLM-5.3-Flash kpool tail) is never stored or loaded.
+
+    Its single per-request block carries no hash, so the store must offload
+    only the attention blocks and the load must pair only attention blocks.
+    """
+    fix = make_scheduler(
+        num_cpu_blocks=8,
+        num_gpu_blocks=16,
+        kv_cache_config=_make_scratch_kv_cache_config(16),
+    )
+    sched = fix.scheduler
+    gpu_pool = fix.gpu_block_pool
+    num_blocks = 2
+
+    req = make_request(num_blocks=num_blocks)
+    fa_blocks = _allocate_gpu_blocks(gpu_pool, req, num_blocks, group_id=0)
+    scratch_block = gpu_pool.get_new_blocks(1)
+    kv_blocks = KVCacheBlocks(blocks=(fa_blocks, scratch_block))
+    req.num_computed_tokens = num_blocks * BLOCK_SIZE
+    sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0)
+    sched_out = make_scheduler_output(
+        {req.request_id: num_blocks * BLOCK_SIZE},
+        new_reqs={req.request_id: kv_blocks.get_block_ids()},
+    )
+    meta = sched.build_connector_meta(sched_out)
+    assert sorted(meta.store_gpu_blocks) == sorted(b.block_id for b in fa_blocks)
+    simulate_store_completion(sched, meta.store_event)
+
+    req2 = Request(
+        request_id="req-scratch-load",
+        prompt_token_ids=req.prompt_token_ids,
+        sampling_params=req.sampling_params,
+        pooling_params=None,
+        mm_features=None,
+        block_hasher=req._block_hasher,
+    )
+    hit_tokens, is_async = sched.get_num_new_matched_tokens(req2, num_computed_tokens=0)
+    assert hit_tokens == num_blocks * BLOCK_SIZE
+    assert is_async is True
+
+    fa_blocks2 = gpu_pool.get_new_blocks(num_blocks)
+    scratch_block2 = gpu_pool.get_new_blocks(1)
+    kv_blocks2 = KVCacheBlocks(blocks=(fa_blocks2, scratch_block2))
+    sched.update_state_after_alloc(req2, kv_blocks2, num_external_tokens=hit_tokens)
+    sched_out2 = make_scheduler_output(
+        {req2.request_id: 1},
+        new_reqs={req2.request_id: kv_blocks2.get_block_ids()},
+    )
+    meta2 = sched.build_connector_meta(sched_out2)
+    assert sorted(meta2.load_gpu_blocks) == sorted(b.block_id for b in fa_blocks2)
+    assert len(meta2.load_cpu_blocks) == num_blocks
+
+
 # ---------------------------------------------------------------------------
 # Test 9: Chunked prefill accumulates block_ids across steps
 # ---------------------------------------------------------------------------
@@ -2844,3 +2946,122 @@ def test_finished_eager_store_recovers_nulled_window_tail() -> None:
     hit_tokens, is_async = sched.get_num_new_matched_tokens(req2, num_computed_tokens=0)
     assert hit_tokens == num_blocks * BLOCK_SIZE
     assert is_async is True
+
+
+QSA_RING_BLOCK_SIZE = 4
+
+
+def _make_qsa_hybrid_kv_cache_config(num_blocks: int) -> KVCacheConfig:
+    register_all_kvcache_specs(vllm_config=None)
+    full_spec = FullAttentionSpec(
+        block_size=BLOCK_SIZE,
+        num_kv_heads=NUM_KV_HEADS,
+        head_size=HEAD_SIZE,
+        dtype=DTYPE,
+    )
+    ring_spec = CircularBufferSpec(
+        block_size=QSA_RING_BLOCK_SIZE,
+        num_kv_heads=NUM_KV_HEADS,
+        head_size=HEAD_SIZE,
+        head_size_v=0,
+        dtype=DTYPE,
+    )
+    groups = [
+        KVCacheGroupSpec(["layer_0"], full_spec),
+        KVCacheGroupSpec(["qsa_ring"], ring_spec),
+    ]
+    tensors = [
+        KVCacheTensor(
+            size=group.kv_cache_spec.page_size_bytes * num_blocks,
+            layers=list(group.layer_names),
+            layer_stride=group.kv_cache_spec.page_size_bytes * num_blocks,
+            block_stride=group.kv_cache_spec.page_size_bytes,
+        )
+        for group in groups
+    ]
+    return KVCacheConfig(
+        num_blocks=num_blocks,
+        kv_cache_tensors=tensors,
+        kv_cache_groups=groups,
+    )
+
+
+def _make_qsa_scheduler(
+    lazy: bool = False,
+) -> tuple[SimpleCPUOffloadScheduler, BlockPool]:
+    kv_cache_config = _make_qsa_hybrid_kv_cache_config(num_blocks=16)
+    sched = SimpleCPUOffloadScheduler(
+        vllm_config=_make_vllm_config(),
+        kv_cache_config=kv_cache_config,
+        cpu_capacity_bytes=_BYTES_PER_BLOCK * 8 * 2,
+        scheduler_block_size=BLOCK_SIZE,
+        hash_block_size=BLOCK_SIZE,
+        lazy_offload=lazy,
+    )
+    gpu_pool = BlockPool(
+        num_gpu_blocks=16, enable_caching=True, hash_block_size=BLOCK_SIZE
+    )
+    sched.bind_gpu_block_pool(gpu_pool)
+    return sched, gpu_pool
+
+
+def test_qsa_ring_group_is_never_stored_or_loaded() -> None:
+    sched, gpu_pool = _make_qsa_scheduler()
+    assert sched.prefix_cacheable_group_ids == (0,)
+
+    num_blocks = 2
+    req = make_request(num_blocks=num_blocks)
+    fa_blocks = _allocate_gpu_blocks(gpu_pool, req, num_blocks, group_id=0)
+    ring_block = gpu_pool.get_new_blocks(1)[0]
+    kv_blocks = KVCacheBlocks(blocks=(fa_blocks, [ring_block]))
+    req.num_computed_tokens = num_blocks * BLOCK_SIZE
+    sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0)
+
+    sched_out = make_scheduler_output(
+        {req.request_id: num_blocks * BLOCK_SIZE},
+        new_reqs={req.request_id: kv_blocks.get_block_ids()},
+    )
+    meta = sched.build_connector_meta(sched_out)
+    assert set(meta.store_gpu_blocks) == {block.block_id for block in fa_blocks}
+    assert ring_block.block_id not in meta.store_gpu_blocks
+    assert meta.store_event >= 0
+    simulate_store_completion(sched, meta.store_event)
+
+    req2 = Request(
+        request_id="req-qsa-load",
+        prompt_token_ids=req.prompt_token_ids,
+        sampling_params=req.sampling_params,
+        pooling_params=None,
+        mm_features=None,
+        block_hasher=req._block_hasher,
+    )
+    hit_tokens, is_async = sched.get_num_new_matched_tokens(req2, num_computed_tokens=0)
+    assert hit_tokens == num_blocks * BLOCK_SIZE
+    assert is_async is True
+
+    fa_blocks2 = gpu_pool.get_new_blocks(num_blocks)
+    ring_block2 = gpu_pool.get_new_blocks(1)[0]
+    kv_blocks2 = KVCacheBlocks(blocks=(fa_blocks2, [ring_block2]))
+    sched.update_state_after_alloc(req2, kv_blocks2, num_external_tokens=hit_tokens)
+
+    meta2 = sched.build_connector_meta(make_scheduler_output({req2.request_id: 1}))
+    assert set(meta2.load_gpu_blocks) == {block.block_id for block in fa_blocks2}
+    assert ring_block2.block_id not in meta2.load_gpu_blocks
+    assert len(meta2.load_cpu_blocks) == num_blocks
+
+
+def test_lazy_target_blocks_ignore_non_prefix_cacheable_groups() -> None:
+    with_ring = _make_qsa_hybrid_kv_cache_config(num_blocks=16)
+    without_ring = KVCacheConfig(
+        num_blocks=16,
+        kv_cache_tensors=with_ring.kv_cache_tensors[:1],
+        kv_cache_groups=with_ring.kv_cache_groups[:1],
+    )
+    max_batched = 64
+    target_with = SimpleCPUOffloadScheduler._estimate_lazy_target_blocks(
+        with_ring, max_batched
+    )
+    target_without = SimpleCPUOffloadScheduler._estimate_lazy_target_blocks(
+        without_ring, max_batched
+    )
+    assert target_with == target_without

tests/v1/worker/test_workspace.py

@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 from typing import cast
+from unittest.mock import Mock
 
 import pytest
 import torch
@@ -124,3 +125,65 @@ def test_workspace_lane_validation(monkeypatch) -> None:
 
     with pytest.raises(ValueError, match="at least one"):
         workspace.WorkspaceManager(torch.device("cpu"), num_lanes=0)
+
+
+def test_persistent_resources_are_lazy_and_isolated(monkeypatch) -> None:
+    """Only the requesting slot allocates; each ubatch/lane owns its resource."""
+    active_ubatch = 0
+    monkeypatch.setattr(workspace, "dbo_current_ubatch_id", lambda: active_ubatch)
+    manager = workspace.WorkspaceManager(
+        torch.device("cpu"), num_ubatches=2, num_lanes=2
+    )
+    factory = Mock(side_effect=lambda: {"buffer": torch.zeros(8)})
+    resources = []
+    for ubatch in range(2):
+        active_ubatch = ubatch
+        for lane in range(2):
+            with workspace.use_workspace_lane(lane):
+                resource = manager.get_persistent_resource("scratch", factory)
+                resources.append(resource)
+                resource["buffer"].fill_(len(resources))
+                assert manager.get_persistent_resource("scratch", factory) is resource
+                assert factory.call_count == len(resources)
+
+    assert len({r["buffer"].data_ptr() for r in resources}) == 4
+    for i, resource in enumerate(resources, 1):
+        assert torch.all(resource["buffer"] == i)
+
+
+def test_persistent_resource_lock_applies_to_each_slot(monkeypatch) -> None:
+    monkeypatch.setattr(workspace, "dbo_current_ubatch_id", lambda: 0)
+    manager = workspace.WorkspaceManager(torch.device("cpu"), num_lanes=2)
+    factory = Mock(side_effect=object)
+    first = manager.get_persistent_resource("scratch", factory)
+    manager.lock()
+    assert manager.get_persistent_resource("scratch", factory) is first
+    with pytest.raises(AssertionError, match="was not allocated during warmup"):
+        manager.get_persistent_resource("new", factory)
+    with workspace.use_workspace_lane(1):
+        with pytest.raises(AssertionError, match="was not allocated during warmup"):
+            manager.get_persistent_resource("scratch", factory)
+        assert factory.call_count == 1
+        manager.unlock()
+        assert manager.get_persistent_resource("scratch", factory) is not first
+    assert manager.get_persistent_resource("scratch", factory) is first
+
+
+def test_persistent_tensor_preserves_contents_and_rejects_changes() -> None:
+    """Initialization happens once and transient scratch cannot overwrite locks."""
+    manager = workspace.WorkspaceManager(torch.device("cpu"))
+    first = manager.get_persistent("locks", (8,), torch.int32, zero_init=True)
+    assert torch.count_nonzero(first) == 0
+    first.fill_(1)
+    (scratch,) = manager.get_simultaneous(((1024,), torch.uint8))
+    scratch.zero_()
+    assert manager.get_persistent("locks", (8,), torch.int32, zero_init=True) is first
+    assert torch.all(first == 1)
+    with pytest.raises(ValueError, match="requested shape"):
+        manager.get_persistent("locks", (16,), torch.int32)
+    with pytest.raises(ValueError, match="requested shape"):
+        manager.get_persistent("locks", (8,), torch.int64)
+    manager.lock()
+    assert manager.get_persistent("locks", (8,), torch.int32) is first
+    with pytest.raises(AssertionError, match="was not allocated during warmup"):
+        manager.get_persistent("new", (8,), torch.int32)

tools/recipes/README.md

@@ -21,6 +21,21 @@ flowchart LR
 The recipe is the baseline. Hardware and workload information can optionally
 refine the initial configuration. Sweep tuning is an optional validation step.
 
+## Getting Started
+
+Use `serve_with_recipe.sh` to generate the recipe configuration and start vLLM
+in one step:
+
+```bash
+tools/recipes/serve_with_recipe.sh \
+  --model meta-llama/Llama-3.1-8B-Instruct \
+  --hardware xeon6
+```
+
+For `xeon6`, the script enables hardware detection automatically. Generated
+`config.yml` and `env.sh` files remain in the directory where the script was
+invoked.
+
 ## 1. vLLM Recipes Only
 
 Use the converter directly when the recipe already contains the deployment

tools/recipes/serve_with_recipe.sh

@@ -0,0 +1,50 @@
+#!/usr/bin/env bash
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+set -euo pipefail
+
+SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
+CONVERTER="${SCRIPT_DIR}/recipe_json_to_vllm_config.py"
+OUTPUT_DIR="$(pwd -P)"
+CONFIG_PATH="${OUTPUT_DIR}/config.yml"
+ENV_PATH="${OUTPUT_DIR}/env.sh"
+
+args=("$@")
+hardware=""
+
+for ((i = 0; i < ${#args[@]}; i++)); do
+    case "${args[$i]}" in
+        --hardware)
+            if ((i + 1 < ${#args[@]})); then
+                hardware="${args[$((i + 1))]}"
+            fi
+            ;;
+        --hardware=*)
+            hardware="${args[$i]#*=}"
+            ;;
+    esac
+done
+
+if [[ "${hardware,,}" == "xeon6" ]]; then
+    args+=(--detect-hardware)
+fi
+
+echo "==> Generating vLLM configuration"
+python3 "${CONVERTER}" "${args[@]}"
+
+echo
+echo "==> Loading recipe environment"
+# shellcheck disable=SC1090
+source "${ENV_PATH}"
+
+# Released vLLM images keep recipe-relative assets such as examples/*.jinja
+# under /vllm-workspace. Launch from there when available while keeping the
+# generated config and environment files in the caller's output directory.
+if [[ -d /vllm-workspace/examples ]]; then
+    cd /vllm-workspace
+fi
+
+echo
+echo "==> Starting vLLM"
+exec vllm serve --config "${CONFIG_PATH}"

vllm/entrypoints/scale_out/derender/serving.py

@@ -179,6 +179,7 @@ class ServingDerender(BaseServing):
             usage=usage,
             prompt_logprobs=gen.prompt_logprobs,
             kv_transfer_params=gen.kv_transfer_params,
+            metrics=gen.metrics,
         )
 
     async def derender_completion_response(
@@ -245,6 +246,9 @@ class ServingDerender(BaseServing):
             choices=choices,
             usage=usage,
             kv_transfer_params=kv_params,
+            # Metrics describe one prompt. Parallel samples are already
+            # suppressed by /generate; multi-prompt responses cannot be merged.
+            metrics=first.metrics if len(request.generate_responses) == 1 else None,
         )
 
     async def derender_chat_stream_response(

vllm/entrypoints/scale_out/token_in_token_out/protocol.py

@@ -11,7 +11,11 @@ from pydantic import (
 )
 
 from vllm.config import ModelConfig
-from vllm.entrypoints.generate.base.protocol import StreamOptions, validate_cache_salt
+from vllm.entrypoints.generate.base.protocol import (
+    PerRequestMetrics,
+    StreamOptions,
+    validate_cache_salt,
+)
 from vllm.entrypoints.openai.chat_completion.protocol import (
     ChatCompletionLogProbs,
     ChatCompletionRequest,
@@ -360,6 +364,7 @@ class GenerateStreamResponse(BaseModel):
     usage: UsageInfo | None = Field(default=None)
     prompt_token_ids: list[int] | None = None
     mm_placeholders: dict[str, list[PlaceholderRangeInfo]] | None = None
+    metrics: PerRequestMetrics | None = None
 
 
 class GenerateResponse(BaseModel):
@@ -378,6 +383,7 @@ class GenerateResponse(BaseModel):
     prompt_logprobs: list[dict[int, Logprob] | None] | None = None
     prompt_token_ids: list[int] | None = None
     mm_placeholders: dict[str, list[PlaceholderRangeInfo]] | None = None
+    metrics: PerRequestMetrics | None = None
 
     kv_transfer_params: dict[str, Any] | None = Field(
         default=None,

vllm/entrypoints/scale_out/token_in_token_out/serving.py

@@ -12,9 +12,13 @@ from fastapi import Request
 
 from vllm.engine.protocol import EngineClient
 from vllm.entrypoints.chat_utils import AsyncMultiModalItemTracker
-from vllm.entrypoints.generate.base.protocol import RequestResponseMetadata
+from vllm.entrypoints.generate.base.protocol import (
+    PerRequestMetrics,
+    RequestResponseMetadata,
+)
 from vllm.entrypoints.generate.base.serving import (
     GenerateBaseServing,
+    build_spec_decoding_metrics,
     clamp_prompt_logprobs,
 )
 from vllm.entrypoints.openai.chat_completion.protocol import (
@@ -373,6 +377,11 @@ class ServingTokens(GenerateBaseServing):
 
         request_metadata.final_usage_info = usage
 
+        per_request_metrics = None
+        if request.sampling_params.n == 1:
+            spec_stats = build_spec_decoding_metrics(final_res)
+            if spec_stats is not None:
+                per_request_metrics = PerRequestMetrics(speculative_decoding=spec_stats)
         response = GenerateResponse(
             request_id=request_id,
             created=created_time,
@@ -384,6 +393,7 @@ class ServingTokens(GenerateBaseServing):
                 final_res.prompt_token_ids if request.return_token_ids else None
             ),
             mm_placeholders=request._response_mm_placeholders,
+            metrics=per_request_metrics,
             kv_transfer_params=final_res.kv_transfer_params,
             ec_transfer_params=final_res.ec_transfer_params,
         )
@@ -423,6 +433,7 @@ class ServingTokens(GenerateBaseServing):
         prompt_token_ids: list[int] | None = None
         num_cached_tokens = None
         sampling_params: SamplingParams = request.sampling_params
+        last_res: RequestOutput | None = None
 
         include_usage, include_continuous_usage = should_include_usage(
             request.stream_options, False
@@ -430,6 +441,7 @@ class ServingTokens(GenerateBaseServing):
 
         try:
             async for res in result_generator:
+                last_res = res
                 if first_iteration:
                     if res.prompt_token_ids is not None:
                         num_prompt_tokens = len(res.prompt_token_ids)
@@ -496,10 +508,10 @@ class ServingTokens(GenerateBaseServing):
                             total_tokens=(num_prompt_tokens + num_generated_tokens[i]),
                         )
 
-                    # Omit absent prompt metadata (Rust skips None fields).
+                    # Omit fields that are absent from token-bearing chunks.
                     exclude = {
                         name
-                        for name in ("prompt_token_ids", "mm_placeholders")
+                        for name in ("prompt_token_ids", "mm_placeholders", "metrics")
                         if getattr(chunk, name) is None
                     }
                     yield f"data: {chunk.model_dump_json(exclude=exclude)}\n\n"
@@ -517,10 +529,18 @@ class ServingTokens(GenerateBaseServing):
                 )
 
             if include_usage:
+                per_request_metrics = None
+                if sampling_params.n == 1:
+                    spec_stats = build_spec_decoding_metrics(last_res)
+                    if spec_stats is not None:
+                        per_request_metrics = PerRequestMetrics(
+                            speculative_decoding=spec_stats
+                        )
                 final_chunk = GenerateStreamResponse(
                     request_id=request_id,
                     choices=[],
                     usage=final_usage_info,
+                    metrics=per_request_metrics,
                 )
                 yield f"data: {final_chunk.model_dump_json(exclude_none=True)}\n\n"
 

vllm/model_executor/kernels/linear/mixed_precision/marlin.py

@@ -10,7 +10,6 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils import (
     apply_gptq_marlin_linear,
     marlin_act_int8_process_scales,
     marlin_make_empty,
-    marlin_make_workspace_new,
     marlin_pad_dim,
     marlin_pad_qweight,
     marlin_pad_scales,
@@ -91,11 +90,6 @@ class MarlinLinearKernel(MPLinearKernel):
         size_k, size_n = c.partition_weight_shape
         padded_n, padded_k = marlin_padded_nk(size_n, size_k, c.group_size)
 
-        # Allocate marlin workspace, reusing existing storage on reload.
-        self.workspace = marlin_make_workspace_new(
-            device, existing=getattr(self, "workspace", None)
-        )
-
         # Default name since marlin requires empty parameter for zp,
         # TODO: remove this requirement from marlin (allow optional tensors)
         if self.w_zp_name is None:
@@ -197,7 +191,7 @@ class MarlinLinearKernel(MPLinearKernel):
             weight=w_q,
             weight_scale=w_s,
             weight_zp=w_zp,  # type: ignore
-            workspace=self.workspace,
+            workspace=None,
             wtype=c.weight_type,
             input_size_per_partition=c.partition_weight_shape[0],
             output_size_per_partition=c.partition_weight_shape[1],

vllm/model_executor/kernels/linear/mxfp4/marlin.py

@@ -58,7 +58,7 @@ class MarlinMxFp4LinearKernel(MxFp4LinearKernel):
             weight=layer.weight,
             weight_scale=layer.weight_scale,
             weight_global_scale=None,
-            workspace=layer.workspace,
+            workspace=None,
             size_n=layer.output_size_per_partition,
             size_k=layer.input_size_per_partition,
             bias=bias,

vllm/model_executor/kernels/linear/mxfp8/marlin.py

@@ -46,7 +46,7 @@ class MarlinMxfp8LinearKernel(Mxfp8LinearKernel):
             input=x,
             weight=layer.weight,
             weight_scale=layer.weight_scale,
-            workspace=layer.workspace,
+            workspace=None,
             size_n=layer.output_size_per_partition,
             size_k=layer.input_size_per_partition,
             bias=bias,

vllm/model_executor/kernels/linear/nvfp4/marlin.py

@@ -50,7 +50,7 @@ class MarlinNvFp4LinearKernel(NvFp4LinearKernel):
             weight=layer.weight,
             weight_scale=layer.weight_scale,
             weight_global_scale=layer.weight_global_scale,
-            workspace=layer.workspace,
+            workspace=None,
             size_n=layer.output_size_per_partition,
             size_k=layer.input_size_per_partition,
             bias=bias,

vllm/model_executor/kernels/linear/scaled_mm/marlin.py

@@ -93,7 +93,7 @@ class MarlinFP8ScaledMMLinearKernel(FP8ScaledMMLinearKernel):
             input=x,
             weight=layer.weight,
             weight_scale=weight_scale,
-            workspace=layer.workspace,
+            workspace=None,
             size_n=layer.output_size_per_partition,
             size_k=layer.input_size_per_partition,
             input_dtype=self.marlin_input_dtype,

vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py

@@ -27,6 +27,7 @@ from vllm.model_executor.layers.fused_moe.moe_align_block_size import (
 from vllm.model_executor.layers.fused_moe.moe_fused_mul_sum import moe_fused_mul_sum
 from vllm.model_executor.layers.fused_moe.moe_permute_unpermute import (
     MoEPermuteScratch,
+    get_moe_permute_scratch,
     moe_permute,
     moe_permute_unpermute_supported,
     moe_unpermute,
@@ -176,7 +177,6 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular):
             max_num_tokens=max_num_tokens,
             num_dispatchers=num_dispatchers,
         )
-        self._permute_scratch: dict[int, MoEPermuteScratch] = {}
 
     def init_humming_moe(self):
         from vllm.utils.humming import get_heuristics_config
@@ -259,24 +259,20 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular):
         if not moe_permute_unpermute_supported():
             return None
 
-        scratch = self._permute_scratch.get(topk)
-        if scratch is None:
-            max_expanded_rows = (
-                self.moe_config.max_num_tokens
-                * self.moe_config.dp_size
-                * self.moe_config.experts_per_token
-            )
-            scratch = MoEPermuteScratch(
-                max_num_tokens=math.ceil(max_expanded_rows / topk),
-                topk=topk,
-                num_experts=self.moe_config.num_experts,
-                num_local_experts=self.moe_config.num_local_experts,
-                device=torch.device(self.moe_config.device),
-                hidden_size=self.moe_config.hidden_dim,
-                hidden_dtype=self.moe_config.in_dtype,
-            )
-            self._permute_scratch[topk] = scratch
-        return scratch
+        max_expanded_rows = (
+            self.moe_config.max_num_tokens
+            * self.moe_config.dp_size
+            * self.moe_config.experts_per_token
+        )
+        return get_moe_permute_scratch(
+            max_num_tokens=math.ceil(max_expanded_rows / topk),
+            topk=topk,
+            num_experts=self.moe_config.num_experts,
+            num_local_experts=self.moe_config.num_local_experts,
+            device=torch.device(self.moe_config.device),
+            hidden_size=self.moe_config.hidden_dim,
+            hidden_dtype=self.moe_config.in_dtype,
+        )
 
     def get_global_valid_shape_m(self, topk_ids: torch.Tensor):
         ctx = get_forward_context()

vllm/model_executor/layers/fused_moe/experts/marlin_moe.py

@@ -35,7 +35,7 @@ from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
 from vllm.model_executor.layers.fused_moe.utils import _resize_cache
 from vllm.model_executor.layers.quantization.utils.marlin_utils import (
     get_marlin_input_dtype,
-    marlin_make_workspace_new,
+    get_marlin_workspace,
     marlin_moe_intermediate_size,
     marlin_quant_input,
 )
@@ -95,7 +95,7 @@ def _fused_marlin_moe(
     N = marlin_moe_intermediate_size(w1, w2)
     w13_num_shards = 2 if activation.is_gated else 1
     if workspace is None:
-        workspace = marlin_make_workspace_new(hidden_states.device, 4)
+        workspace = get_marlin_workspace(hidden_states.device)
 
     if intermediate_cache13 is None:
         intermediate_cache13 = torch.empty(

vllm/model_executor/layers/fused_moe/moe_permute_unpermute.py

@@ -5,6 +5,11 @@ from dataclasses import dataclass, field
 
 import torch
 
+from vllm.v1.worker.workspace import (
+    current_workspace_manager,
+    is_workspace_manager_initialized,
+)
+
 
 @dataclass
 class MoEPermuteScratch:
@@ -102,6 +107,48 @@ class MoEPermuteScratch:
         return topk_ids_int32
 
 
+def get_moe_permute_scratch(
+    *,
+    max_num_tokens: int,
+    topk: int,
+    num_experts: int,
+    num_local_experts: int,
+    device: torch.device,
+    hidden_size: int | None = None,
+    hidden_dtype: torch.dtype | None = None,
+) -> MoEPermuteScratch:
+    """Share scratch across sequential layers in the current ubatch and lane.
+
+    Without a workspace manager, allocate new scratch for each call.
+    """
+
+    def create_scratch() -> MoEPermuteScratch:
+        return MoEPermuteScratch(
+            max_num_tokens=max_num_tokens,
+            topk=topk,
+            num_experts=num_experts,
+            num_local_experts=num_local_experts,
+            device=device,
+            hidden_size=hidden_size,
+            hidden_dtype=hidden_dtype,
+        )
+
+    if not is_workspace_manager_initialized():
+        return create_scratch()
+
+    key = (
+        MoEPermuteScratch,
+        max_num_tokens,
+        topk,
+        num_experts,
+        num_local_experts,
+        device,
+        hidden_size,
+        hidden_dtype,
+    )
+    return current_workspace_manager().get_persistent_resource(key, create_scratch)
+
+
 def moe_permute(
     hidden_states: torch.Tensor,
     a1q_scale: torch.Tensor | None,

vllm/model_executor/layers/indexer_topk.py

@@ -14,6 +14,12 @@ from vllm.v1.worker.workspace import current_workspace_manager
 
 RADIX_TOPK_WORKSPACE_SIZE = 1024 * 1024
 
+# cooperative_topk's hard row limit: one cluster wave covers at most this many
+# rows. "auto" uses cooperative_topk for every batch within the limit and falls
+# back to persistent_topk past it; explicit ``cooperative`` requests are
+# validated against the same bound.
+AUTO_COOPERATIVE_MAX_ROWS = 64
+
 # ---------------------------------------------------------------------------
 # DeepSelect (vllm._deepselect_C)
 # ---------------------------------------------------------------------------
@@ -209,8 +215,13 @@ class SparseIndexerTopk(torch.nn.Module):
     def _resolve_auto(
         self, logits: torch.Tensor, topk_tokens: int, num_rows: int
     ) -> str:
-        """The pre-existing priority chain: cooperative -> persistent ->
-        per_row. deep_select/flashinfer/torch are opt-in only."""
+        """The priority chain: cooperative -> persistent -> per_row.
+        deep_select/flashinfer/torch are opt-in only.
+
+        cooperative_topk is preferred whenever it is applicable, i.e. within
+        its AUTO_COOPERATIVE_MAX_ROWS row limit; larger batches go to
+        persistent_topk.
+        """
         if not self._cooperative_constraints(logits, topk_tokens, num_rows):
             return "cooperative"
         if self._is_cuda and topk_tokens in (512, 1024, 2048):
@@ -228,8 +239,10 @@ class SparseIndexerTopk(torch.nn.Module):
             failures.append(
                 f"topk_tokens must be in (512, 1024, 2048), got {topk_tokens}"
             )
-        if num_rows > 64:
-            failures.append(f"num_rows must be <= 64, got {num_rows}")
+        if num_rows > AUTO_COOPERATIVE_MAX_ROWS:
+            failures.append(
+                f"num_rows must be <= {AUTO_COOPERATIVE_MAX_ROWS}, got {num_rows}"
+            )
         if logits.stride(0) % 4 != 0:
             failures.append(
                 f"logits.stride(0) must be divisible by 4, got {logits.stride(0)}"
@@ -293,13 +306,18 @@ class SparseIndexerTopk(torch.nn.Module):
             (topk_workspace,) = current_workspace_manager().get_simultaneous(
                 ((RADIX_TOPK_WORKSPACE_SIZE,), torch.uint8),
             )
+            # persistent_topk's last argument is the host-side "max seq_len
+            # across rows" bound (it gates the sampled_topk path and clamps
+            # per-row lengths), not the padded row width: logits is sized to
+            # max_model_len, so passing its width would incorrectly enable the
+            # long-row kernels on short-context batches.
             torch.ops._C.persistent_topk(
                 logits,
                 seq_lens,
                 topk_indices,
                 topk_workspace,
                 topk_tokens,
-                logits.shape[1],
+                max_seq_len,
             )
         elif backend == "flashinfer":
             # Deferred: importing flashinfer initializes CUDA at import time.

vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py

@@ -167,6 +167,8 @@ def _log_gdn_backend_decision(
             head_k_dim,
         )
         return
+    elif current_platform.is_xpu():
+        return
 
     chosen = {
         "flashinfer": "FlashInfer",
@@ -524,6 +526,8 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):
                 self.gdn_decode_kernel = "triton"
         elif current_platform.is_cpu():
             self.gdn_decode_kernel = "CPU"
+        elif current_platform.is_xpu():
+            self.gdn_decode_kernel = "XPU"
 
         self.enable_fused_gdn_decode = self.gdn_decode_kernel == "cuda"
         logger.info_once("GDN decode kernel: %s", self.gdn_decode_kernel)

vllm/model_executor/layers/quantization/auto_awq.py

@@ -48,7 +48,6 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils import (
     check_marlin_supports_layer,
     check_moe_marlin_supports_layer,
     get_marlin_input_dtype,
-    marlin_make_workspace_new,
     verify_marlin_supported,
 )
 from vllm.model_executor.layers.quantization.utils.quant_utils import (
@@ -642,9 +641,6 @@ class AutoAWQMoEMethod(FusedMoEMethodBase):
         layer.register_parameter("w2_qzeros", w2_qzeros)
         set_weight_attrs(w2_qzeros, extra_weight_attrs)
 
-        device = layer.w13_qweight.device
-        layer.workspace = marlin_make_workspace_new(device, 4)
-
     def process_weights_after_loading(self, layer: RoutedExperts) -> None:
         converted = convert_to_wna16_moe_kernel_format(
             backend=self.wna16_moe_backend,

vllm/model_executor/layers/quantization/auto_gptq.py

@@ -16,7 +16,6 @@ from vllm.model_executor.kernels.linear import (
 )
 from vllm.model_executor.layers.fused_moe import (
     FusedMoEConfig,
-    FusedMoEExpertsModular,
     FusedMoEMethodBase,
     FusedMoEQuantConfig,
     FusedMoeWeightScaleSupported,
@@ -46,7 +45,6 @@ from vllm.model_executor.layers.quantization.utils.gptq_utils import (
 from vllm.model_executor.layers.quantization.utils.marlin_utils import (
     check_moe_marlin_supports_layer,
     get_marlin_input_dtype,
-    marlin_make_workspace_new,
     marlin_repeat_scales_on_all_ranks,
     verify_marlin_supported,
 )
@@ -601,12 +599,6 @@ class AutoGPTQMoEMethod(FusedMoEMethodBase):
         layer.register_parameter("w2_bias", w2_bias)
         set_weight_attrs(w2_bias, extra_weight_attrs)
 
-        if self.experts_cls is not None and issubclass(
-            self.experts_cls, FusedMoEExpertsModular
-        ):
-            device = layer.w13_qweight.device
-            layer.workspace = marlin_make_workspace_new(device, 4)
-
     def process_weights_after_loading(self, layer: RoutedExperts) -> None:
         def replace_or_register(name: str, val: torch.Tensor | None):
             if val is None:

vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py

@@ -11,7 +11,6 @@ from compressed_tensors.quantization import (
 
 from vllm.logger import init_logger
 from vllm.model_executor.layers.fused_moe import (
-    FusedMoEExpertsModular,
     RoutedExperts,
     SharedExperts,
 )
@@ -36,7 +35,6 @@ from vllm.model_executor.layers.quantization.compressed_tensors.schemes.compress
 from vllm.model_executor.layers.quantization.utils.marlin_utils import (
     check_moe_marlin_supports_config,
     get_marlin_input_dtype,
-    marlin_make_workspace_new,
 )
 from vllm.model_executor.layers.quantization.utils.quant_utils import (
     GroupShape,
@@ -466,18 +464,6 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):
                     torch.nn.Parameter(w2_input_global_scale, requires_grad=False),
                 )
 
-            # Marlin workspace — only needed for Marlin-family backends, not emulation.
-            if (
-                self.experts_cls is not None
-                and issubclass(self.experts_cls, FusedMoEExpertsModular)
-                and self.wna16_backend != WNA16MoEBackend.EMULATION
-            ):
-                layer.workspace = marlin_make_workspace_new(
-                    layer.w13_weight_packed.device,
-                    4,
-                    existing=getattr(layer, "workspace", None),
-                )
-
         # Alias packed weights to w13_weight/w2_weight for the modular kernel interface
         layer.w13_weight = layer.w13_weight_packed
         layer.w2_weight = layer.w2_weight_packed

vllm/model_executor/layers/quantization/utils/marlin_utils.py

@@ -8,7 +8,6 @@ import numpy
 import torch
 
 import vllm.envs as envs
-from vllm import _custom_ops as ops
 from vllm.distributed.utils import verify_group_size_divides_partition
 from vllm.logger import init_logger
 from vllm.model_executor.layers.fused_moe import RoutedExperts
@@ -23,6 +22,7 @@ from vllm.platforms import current_platform
 from vllm.scalar_type import ScalarType, scalar_types
 from vllm.utils.math_utils import round_up
 from vllm.utils.platform_utils import num_compute_units
+from vllm.utils.torch_utils import direct_register_custom_op
 
 from .quant_utils import pack_cols, unpack_cols
 
@@ -32,6 +32,7 @@ GPTQ_MARLIN_TILE = 16
 GPTQ_MARLIN_MIN_THREAD_N = 64
 GPTQ_MARLIN_MIN_THREAD_K = 128
 GPTQ_MARLIN_MAX_PARALLEL = 16
+MARLIN_MAX_BLOCKS_PER_SM = 4
 
 MARLIN_SUPPORTED_GROUP_SIZES = [-1, 32, 64, 128]
 
@@ -407,29 +408,96 @@ def marlin_moe_intermediate_size(w1_packed: torch.Tensor, w2_packed: torch.Tenso
 def marlin_make_workspace_new(
     device: torch.device,
     max_blocks_per_sm: int = 1,
-    existing: torch.Tensor | None = None,
 ) -> torch.Tensor:
     # In the new marlin kernel, we use the num of threadblocks as workspace
     # size. The num of threadblocks is sms_count * max_blocks_per_sm.
     sms = num_compute_units(device.index)
-    size = sms * max_blocks_per_sm
-    # On weight reload, reuse the existing storage so the workspace address
-    # captured by CUDA graphs stays valid.
-    if existing is not None:
-        if (
-            existing.device != device
-            or existing.dtype != torch.int
-            or existing.numel() != size
-        ):
-            raise ValueError(
-                f"Existing Marlin workspace is incompatible "
-                f"(device={existing.device}, dtype={existing.dtype}, "
-                f"numel={existing.numel()}; expected device={device}, "
-                f"dtype={torch.int}, numel={size}). Reload must reuse the "
-                f"workspace storage captured by CUDA graphs."
-            )
-        return existing.zero_()
-    return torch.zeros(size, dtype=torch.int, device=device, requires_grad=False)
+    return torch.zeros(sms * max_blocks_per_sm, dtype=torch.int, device=device)
+
+
+def get_marlin_workspace(device: torch.device) -> torch.Tensor:
+    """Get zero-initialized locks for the current stream, ubatch, and lane.
+
+    Marlin resets its locks after each call, so ordered calls can reuse them.
+    """
+    from vllm.utils.torch_utils import current_stream
+    from vllm.v1.worker.workspace import (
+        current_workspace_manager,
+        is_workspace_manager_initialized,
+    )
+
+    if not is_workspace_manager_initialized():
+        return marlin_make_workspace_new(device, MARLIN_MAX_BLOCKS_PER_SM)
+
+    return current_workspace_manager().get_persistent(
+        ("marlin_locks", current_stream()),
+        (num_compute_units(device.index) * MARLIN_MAX_BLOCKS_PER_SM,),
+        torch.int,
+        zero_init=True,
+    )
+
+
+def _marlin_gemm(
+    a: torch.Tensor,
+    b_q_weight: torch.Tensor,
+    b_bias: torch.Tensor | None,
+    b_scales: torch.Tensor,
+    a_scales: torch.Tensor | None,
+    global_scale: torch.Tensor | None,
+    b_zeros: torch.Tensor | None,
+    workspace: torch.Tensor | None,
+    b_q_type_id: int,
+    size_m: int,
+    size_n: int,
+    size_k: int,
+    use_atomic_add: bool = False,
+    use_fp32_reduce: bool = False,
+    is_zp_float: bool = False,
+) -> torch.Tensor:
+    if workspace is None:
+        workspace = get_marlin_workspace(a.device)
+    return torch.ops._C.marlin_gemm(
+        a,
+        None,
+        b_q_weight,
+        b_bias,
+        b_scales,
+        a_scales,
+        global_scale,
+        b_zeros,
+        workspace,
+        b_q_type_id,
+        size_m,
+        size_n,
+        size_k,
+        use_atomic_add,
+        use_fp32_reduce,
+        is_zp_float,
+    )
+
+
+def _marlin_gemm_fake(
+    a: torch.Tensor,
+    b_q_weight: torch.Tensor,
+    b_bias: torch.Tensor | None,
+    b_scales: torch.Tensor,
+    a_scales: torch.Tensor | None,
+    global_scale: torch.Tensor | None,
+    b_zeros: torch.Tensor | None,
+    workspace: torch.Tensor | None,
+    b_q_type_id: int,
+    size_m: int,
+    size_n: int,
+    size_k: int,
+    use_atomic_add: bool = False,
+    use_fp32_reduce: bool = False,
+    is_zp_float: bool = False,
+) -> torch.Tensor:
+    dtype = a.dtype if a.dtype in (torch.float16, torch.bfloat16) else b_scales.dtype
+    return torch.empty((size_m, size_n), device=a.device, dtype=dtype)
+
+
+direct_register_custom_op("marlin_gemm", _marlin_gemm, fake_impl=_marlin_gemm_fake)
 
 
 def marlin_repeat_scales_on_all_ranks(group_size: int, is_row_parallel: bool) -> bool:
@@ -686,7 +754,7 @@ def apply_gptq_marlin_linear(
     weight: torch.Tensor,
     weight_scale: torch.Tensor,
     weight_zp: torch.Tensor,
-    workspace: torch.Tensor,
+    workspace: torch.Tensor | None,
     wtype: ScalarType,
     output_size_per_partition: int,
     input_size_per_partition: int,
@@ -723,9 +791,8 @@ def apply_gptq_marlin_linear(
 
         reshaped_x, a_scales = marlin_quant_input(reshaped_x, input_dtype)
 
-    output = ops.marlin_gemm(
+    output = torch.ops.vllm.marlin_gemm(
         reshaped_x,
-        None,
         weight,
         bias,
         weight_scale,
@@ -733,7 +800,7 @@ def apply_gptq_marlin_linear(
         None,
         weight_zp,
         workspace,
-        wtype,
+        wtype.id,
         size_m=reshaped_x.shape[0],
         size_n=padded_n,
         size_k=padded_k,

vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py

@@ -10,7 +10,6 @@ from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts
 from vllm.model_executor.layers.quantization.utils.marlin_utils import (
     USE_FP32_REDUCE_DEFAULT,
     get_marlin_input_dtype,
-    marlin_make_workspace_new,
     marlin_pad_dim,
     marlin_pad_qweight,
     marlin_pad_scales,
@@ -161,7 +160,7 @@ def apply_fp4_marlin_linear(
     weight: torch.Tensor,
     weight_scale: torch.Tensor,
     weight_global_scale: torch.Tensor | None,
-    workspace: torch.Tensor,
+    workspace: torch.Tensor | None,
     size_n: int,
     size_k: int,
     bias: torch.Tensor | None = None,
@@ -196,9 +195,8 @@ def apply_fp4_marlin_linear(
 
         inputs, a_scales = marlin_quant_input(inputs, torch.float8_e4m3fn)
 
-    output = ops.marlin_gemm(
+    output = torch.ops.vllm.marlin_gemm(
         a=inputs,
-        c=None,
         b_q_weight=weight,
         b_bias=bias,
         b_scales=weight_scale,
@@ -206,7 +204,7 @@ def apply_fp4_marlin_linear(
         global_scale=weight_global_scale,
         b_zeros=None,
         workspace=workspace,
-        b_q_type=scalar_types.float4_e2m1f,
+        b_q_type_id=scalar_types.float4_e2m1f.id,
         size_m=reshaped_x.size(0),
         size_n=padded_n,
         size_k=padded_k,
@@ -237,13 +235,6 @@ def prepare_fp4_layer_for_marlin(
 
     assert layer.weight.shape == (part_size_n, part_size_k // 2)
 
-    device = layer.weight.device
-
-    # WORKSPACE
-    layer.workspace = marlin_make_workspace_new(
-        device, existing=getattr(layer, "workspace", None)
-    )
-
     # WEIGHT
     # Repack weights to marlin format
     qweight = layer.weight.view(torch.int32).T.contiguous()
@@ -389,15 +380,9 @@ def prepare_nvfp4_moe_layer_for_marlin(
             return x
         return torch.nn.functional.pad(x, (0, (padded_N - N) // packing))
 
-    device = w13.device
     param_dtype = layer.params_dtype
     is_a_8bit = input_dtype is not None and input_dtype.itemsize == 1
 
-    # WORKSPACE
-    layer.workspace = marlin_make_workspace_new(
-        device, 4, existing=getattr(layer, "workspace", None)
-    )
-
     # WEIGHT
     # Repack weights to marlin format
     def repack_weight(weight: torch.Tensor, name: str) -> torch.Tensor:
@@ -481,12 +466,7 @@ def prepare_moe_fp4_layer_for_marlin(
     k = layer.moe_config.hidden_dim
     n = layer.moe_config.intermediate_size_per_partition
 
-    # WORKSPACE
-    device = layer.w13_weight.device
     param_dtype = layer.params_dtype
-    layer.workspace = marlin_make_workspace_new(
-        device, 4, existing=getattr(layer, "workspace", None)
-    )
     is_a_8bit = input_dtype is not None and input_dtype.itemsize == 1
 
     # WEIGHT

vllm/model_executor/layers/quantization/utils/marlin_utils_fp8.py

@@ -9,7 +9,6 @@ from vllm.logger import init_logger
 from vllm.model_executor.layers.quantization.utils.marlin_utils import (
     USE_FP32_REDUCE_DEFAULT,
     get_marlin_input_dtype,
-    marlin_make_workspace_new,
     marlin_moe_padded_intermediate,
     marlin_pad_dim,
     marlin_pad_qweight,
@@ -50,7 +49,7 @@ def apply_fp8_marlin_linear(
     input: torch.Tensor,
     weight: torch.Tensor,
     weight_scale: torch.Tensor,
-    workspace: torch.Tensor,
+    workspace: torch.Tensor | None,
     size_n: int,
     size_k: int,
     bias: torch.Tensor | None,
@@ -80,9 +79,8 @@ def apply_fp8_marlin_linear(
         # inputs, a_scales = marlin_quant_input(inputs, torch.float8_e4m3fn)
         raise RuntimeError("Marlin W8A8 is not supported.")
 
-    output = ops.marlin_gemm(
+    output = torch.ops.vllm.marlin_gemm(
         a=inputs,
-        c=None,
         b_q_weight=weight,
         b_bias=bias,
         b_scales=weight_scale,
@@ -90,7 +88,7 @@ def apply_fp8_marlin_linear(
         global_scale=None,
         b_zeros=None,
         workspace=workspace,
-        b_q_type=scalar_types.float8_e4m3fn,
+        b_q_type_id=scalar_types.float8_e4m3fn.id,
         size_m=reshaped_x.size(0),
         size_n=padded_n,
         size_k=padded_k,
@@ -127,13 +125,6 @@ def prepare_fp8_layer_for_marlin(
     else:
         assert layer.weight.shape == (part_size_n, part_size_k)
 
-    device = layer.weight.device
-
-    # WORKSPACE
-    layer.workspace = marlin_make_workspace_new(
-        device, existing=getattr(layer, "workspace", None)
-    )
-
     # WEIGHT
     # Repack weights to marlin format
     qweight = pack_fp8_to_int32(layer.weight, size_k_first)
@@ -241,12 +232,7 @@ def prepare_fp8_moe_layer_for_marlin(
     w13_weight_scale: torch.Tensor,
     w2_weight_scale: torch.Tensor,
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
-    """Shuffle weights and scales into marlin format.
-
-    Note that this function has the side effect of adding a `workspace`
-    attribute to the layer. This `workspace` does not need to be
-    registered as a Parameter as it is not used during weight reloading.
-    """
+    """Shuffle weights and scales into marlin format."""
     logger.warning_once(
         "Your GPU does not have native support for FP8 computation but "
         "FP8 quantization is being used. Weight-only FP8 compression will "
@@ -272,12 +258,6 @@ def prepare_fp8_moe_layer_for_marlin(
         w13_weight = _moe_pad_shard_rows(w13_weight, n, padded_n)
         w2_weight = _moe_pad_last(w2_weight, n, padded_n)
 
-    # WORKSPACE
-    device = layer.w13_weight.device
-    layer.workspace = marlin_make_workspace_new(
-        device, 4, existing=getattr(layer, "workspace", None)
-    )
-
     # WEIGHT
     # Repack weights to marlin format
     def repack_weight(name: str, weight: torch.Tensor) -> torch.Tensor:
@@ -396,7 +376,7 @@ def apply_mxfp8_marlin_linear(
     input: torch.Tensor,
     weight: torch.Tensor,
     weight_scale: torch.Tensor,
-    workspace: torch.Tensor,
+    workspace: torch.Tensor | None,
     size_n: int,
     size_k: int,
     bias: torch.Tensor | None = None,
@@ -416,9 +396,8 @@ def apply_mxfp8_marlin_linear(
         dtype=input.dtype,
     )
 
-    output = ops.marlin_gemm(
+    output = torch.ops.vllm.marlin_gemm(
         a=reshaped_x,
-        c=None,
         b_q_weight=weight,
         b_bias=bias,
         b_scales=weight_scale,
@@ -426,7 +405,7 @@ def apply_mxfp8_marlin_linear(
         global_scale=None,
         b_zeros=None,
         workspace=workspace,
-        b_q_type=scalar_types.float8_e4m3fn,
+        b_q_type_id=scalar_types.float8_e4m3fn.id,
         size_m=reshaped_x.size(0),
         size_n=padded_n,
         size_k=padded_k,
@@ -451,13 +430,6 @@ def prepare_mxfp8_layer_for_marlin(layer: torch.nn.Module) -> None:
     group_size = 32  # MX standard block size
     padded_n, padded_k = marlin_padded_nk(part_size_n, part_size_k, group_size)
 
-    device = layer.weight.device
-
-    # WORKSPACE
-    layer.workspace = marlin_make_workspace_new(
-        device, existing=getattr(layer, "workspace", None)
-    )
-
     # WEIGHT - repack FP8 weights to Marlin format
     qweight = pack_fp8_to_int32(layer.weight, size_k_first=False)
     qweight = qweight.T.contiguous()
@@ -512,7 +484,7 @@ def prepare_mxfp8_moe_layer_for_marlin(
     """Repack MXFP8 MoE weights and scales into Marlin kernel format.
 
     Args:
-        layer: MoE layer (used to read params_dtype and attach workspace).
+        layer: MoE layer.
         w13: [E, 2*N, K] float8_e4m3fn weights.
         w2:  [E, K, N] float8_e4m3fn weights.
         w13_scale: [E, 2*N, K//32] uint8 e8m0 scales.
@@ -537,13 +509,8 @@ def prepare_mxfp8_moe_layer_for_marlin(
         n = padded_n
     w13_n = w13.shape[1]
 
-    device = w13.device
     param_dtype = torch.get_default_dtype()
 
-    layer.workspace = marlin_make_workspace_new(
-        device, 4, existing=getattr(layer, "workspace", None)
-    )
-
     def repack_weight(weight: torch.Tensor, name: str) -> torch.Tensor:
         if "w13" in name:
             size_n, size_k = w13_n, k

vllm/model_executor/models/diffusion_gemma.py

@@ -543,7 +543,7 @@ def _compiled_sample_step(
     if tp_size > 1:
         soft_embeds = torch.ops.vllm.all_reduce(soft_embeds, group_name=tp_group_name)
     soft_embeds = soft_embeds * normalizer
-    sc_embeds[decode_slots] = soft_embeds * sc_keep
+    sc_embeds[decode_slots] = (soft_embeds * sc_keep).to(sc_embeds.dtype)
 
     # Overwrite canvas with argmax for newly converged denoise requests
     newly_converged = (converged & is_denoise).unsqueeze(1)

vllm/model_executor/models/mimo_v2.py

@@ -27,6 +27,7 @@ from vllm.model_executor.layers.fused_moe import (
     FusedMoEFactory,
     fused_moe_make_expert_params_mapping,
 )
+from vllm.model_executor.layers.fused_moe.router.gate_linear import GateLinear
 from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.linear import (
     MergedColumnParallelLinear,
@@ -156,14 +157,16 @@ class MiMoV2MoE(nn.Module):
 
         dtype = getattr(config, "moe_router_dtype", "float32")
         self.gate_dtype = str_dtype_to_torch_dtype(dtype)
-        self.gate = nn.Linear(
+        self.gate = GateLinear(
             config.hidden_size,
             config.n_routed_experts,
             bias=False,
-            dtype=self.gate_dtype,
+            params_dtype=self.gate_dtype,
+            out_dtype=torch.float32,
+            prefix=f"{prefix}.gate",
         )
         self.gate.e_score_correction_bias = nn.Parameter(
-            torch.empty(config.n_routed_experts, dtype=self.gate_dtype)
+            torch.empty(config.n_routed_experts, dtype=self.gate.out_dtype)
         )
 
         self.experts = FusedMoEFactory(
@@ -182,7 +185,7 @@ class MiMoV2MoE(nn.Module):
             num_expert_group=config.n_group,
             topk_group=config.topk_group,
             scoring_func="sigmoid",
-            router_logits_dtype=self.gate_dtype,
+            router_logits_dtype=self.gate.out_dtype,
         )
 
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
@@ -198,7 +201,7 @@ class MiMoV2MoE(nn.Module):
             gate_input = hidden_states.to(self.gate_dtype)
         else:
             gate_input = hidden_states
-        router_logits = self.gate(gate_input)
+        router_logits, _ = self.gate(gate_input)
         final_hidden_states = self.experts(
             hidden_states=hidden_states, router_logits=router_logits
         )

vllm/model_executor/models/mimo_v2_omni.py

@@ -52,6 +52,7 @@ from vllm.transformers_utils.processors.mimo_v2_omni import (
 
 from .interfaces import (
     MultiModalEmbeddings,
+    SupportsEagle3,
     SupportsMultiModal,
     SupportsPP,
     SupportsQuant,
@@ -1209,7 +1210,9 @@ class MiMoV2OmniDummyInputsBuilder(BaseDummyInputsBuilder[MiMoV2OmniProcessingIn
     info=MiMoV2OmniProcessingInfo,
     dummy_inputs=MiMoV2OmniDummyInputsBuilder,
 )
-class MiMoV2OmniForCausalLM(nn.Module, SupportsMultiModal, SupportsPP, SupportsQuant):
+class MiMoV2OmniForCausalLM(
+    nn.Module, SupportsMultiModal, SupportsPP, SupportsQuant, SupportsEagle3
+):
     # To ensure correct weight loading and mapping.
     hf_to_vllm_mapper = WeightsMapper(
         orig_to_new_prefix={

vllm/model_executor/models/qwen3_dflash.py

@@ -168,6 +168,7 @@ class DFlashQwen3Attention(nn.Module):
         rms_norm_eps: float = 1e-06,
         attention_bias: bool = False,
         add_swa_attention_sink_bias: bool = False,
+        v_scale: float | None = None,
         sliding_window: int | None = None,
         causal: bool = False,
         is_neox_style: bool = True,
@@ -193,6 +194,7 @@ class DFlashQwen3Attention(nn.Module):
         self.q_size = self.num_heads * self.head_dim
         self.kv_size = self.num_kv_heads * self.head_dim
         self.scaling = self.head_dim**-0.5
+        self.v_scale = v_scale
 
         self.qkv_proj = QKVParallelLinear(
             hidden_size,
@@ -264,6 +266,8 @@ class DFlashQwen3Attention(nn.Module):
 
         q, k = self.rotary_emb(positions, q, k)
 
+        if self.v_scale is not None:
+            v = v * self.v_scale
         attn_output = self.attn(q, k, v)
         output, _ = self.o_proj(attn_output)
         return output
@@ -312,6 +316,7 @@ class DFlashQwen3DecoderLayer(nn.Module):
             rms_norm_eps=config.rms_norm_eps,
             attention_bias=getattr(config, "attention_bias", False),
             add_swa_attention_sink_bias=add_swa_attention_sink_bias,
+            v_scale=dflash_config.get("attention_value_scale"),
             sliding_window=sliding_window,
             causal=causal,
             is_neox_style=is_neox_style,
@@ -633,6 +638,10 @@ class DFlashQwen3Model(nn.Module):
         if context_slot_mapping is None:
             return
 
+        v_scale = getattr(self.layers[0].self_attn, "v_scale", None)
+        if v_scale is not None:
+            all_v.mul_(v_scale)
+
         # --- Per-layer cache insert ---
         all_k_final = all_k_flat.view(L, num_ctx, nkv, hd)
         per_layer = isinstance(context_slot_mapping, (list, tuple))

vllm/model_executor/models/voyage.py

@@ -8,6 +8,7 @@ import torch
 import torch.nn as nn
 
 from vllm.config import VllmConfig
+from vllm.model_executor.models.interfaces import SupportsLoRA
 from vllm.model_executor.models.qwen3 import Qwen3Model
 from vllm.model_executor.models.utils import AutoWeightsLoader, maybe_prefix
 from vllm.sequence import IntermediateTensors
@@ -15,7 +16,7 @@ from vllm.sequence import IntermediateTensors
 WeightItem = tuple[str, torch.Tensor]
 
 
-class VoyageQwen3BidirectionalEmbedModel(nn.Module):
+class VoyageQwen3BidirectionalEmbedModel(nn.Module, SupportsLoRA):
     """Qwen3Model + Voyage embedding head + bidirectional attention.
 
     Checkpoint conventions (HF):
@@ -30,6 +31,11 @@ class VoyageQwen3BidirectionalEmbedModel(nn.Module):
       - No "model." prefix
     """
 
+    packed_modules_mapping = {
+        "qkv_proj": ["q_proj", "k_proj", "v_proj"],
+        "gate_up_proj": ["gate_proj", "up_proj"],
+    }
+
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         super().__init__()
         self.config = vllm_config.model_config.hf_config

vllm/models/deepseek_v4/amd/rocm.py

@@ -7,6 +7,7 @@ from typing import cast
 
 import torch
 
+from vllm import envs
 from vllm.distributed import (
     get_tensor_model_parallel_world_size,
     tensor_model_parallel_all_reduce,
@@ -42,6 +43,49 @@ from vllm.v1.worker.workspace import current_workspace_manager
 logger = init_logger(__name__)
 
 
+def _wo_a_block_scale_to_e8m0(scale: torch.Tensor) -> torch.Tensor | None:
+    """Normalize checkpoint WO_A scales to raw OCP MX E8M0 bytes.
+
+    E8M0 is an unsigned exponent-only scale format with bias 127. A finite
+    encoded byte ``b`` in ``[0, 254]`` represents ``2 ** (b - 127)``;
+    ``0xFF`` is reserved for NaN. This is the vendor-neutral OCP encoding used
+    by AMD AITER/OPUS, not an NVIDIA-specific convention. Reference: OCP
+    Microscaling Formats (MX) Specification, section 5.4.1:
+    https://www.opencompute.org/documents/ocp-microscaling-formats-mx-v1-0-spec-final-pdf
+
+    Loaders may preserve the encoded byte as ``float8_e8m0fnu``/``uint8``,
+    or decode it to a floating-point power of two. Floating-point inputs are
+    accepted only when the original E8M0 byte can be recovered losslessly.
+    This function does not quantize or round arbitrary scales.
+    """
+    if scale.dtype == torch.float8_e8m0fnu:
+        # Reinterpret the native E8M0 storage. ``to(uint8)`` would perform a
+        # numeric conversion instead of preserving the encoded exponent byte.
+        return scale.view(torch.uint8).contiguous()
+    if scale.dtype == torch.uint8:
+        # The checkpoint loader already exposed the E8M0 wire representation.
+        return scale.contiguous()
+    if not scale.dtype.is_floating_point:
+        return None
+
+    scale_f32 = scale.detach().float()
+    if not bool(torch.isfinite(scale_f32).all()) or bool((scale_f32 <= 0).any()):
+        return None
+
+    # With no mantissa, E8M0 can represent only exact powers of two. Rebuild
+    # the value before encoding so this adapter never silently quantizes a
+    # general floating-point checkpoint scale.
+    exponent = torch.round(torch.log2(scale_f32))
+    if not torch.equal(torch.exp2(exponent), scale_f32):
+        return None
+
+    encoded = exponent.to(torch.int32) + 127
+    # 0xFF is NaN in OCP E8M0, not a finite exponent.
+    if int(encoded.min()) < 0 or int(encoded.max()) > 254:
+        return None
+    return encoded.to(torch.uint8).contiguous()
+
+
 def _trust_dsv4_extra_cache_nan_free(
     kv_cache_dtype: str,
     has_kv_transfer: bool,
@@ -555,6 +599,10 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention):
         # Block scale for the preshuffled weight; None = not preshuffled.
         self._wqa_wkv_scale: torch.Tensor | None = None
         self._wo_b_scale: torch.Tensor | None = None
+        self._wo_a_fp8_weight: torch.Tensor | None = None
+        self._wo_a_e8m0_scale: torch.Tensor | None = None
+        self._wo_a_cos_cache: torch.Tensor | None = None
+        self._wo_a_sin_cache: torch.Tensor | None = None
         self._fused_compressor_weight: torch.Tensor | None
         self.register_buffer("_fused_compressor_weight", None, persistent=False)
         self._fused_compressor_split_sizes: tuple[int, int] | None = None
@@ -815,6 +863,93 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention):
 
         self._wqa_wkv_scale = _prep(self.fused_wqa_wkv)
         self._wo_b_scale = _prep(self.wo_b)
+        if _ON_GFX950 and envs.VLLM_ROCM_USE_AITER_FP8BMM:
+            self._prepare_fp8_wo_a()
+
+    def _prepare_fp8_wo_a(self) -> None:
+        try:
+            from aiter.ops.batched_gemm_op_a8w8 import (
+                batched_gemm_a8w8_mxscale as mxscale_op,
+            )
+            from aiter.ops.inverse_rope_group_quant import (
+                inverse_rope_group_quant as inverse_quant_op,
+            )
+        except ImportError:
+            logger.warning_once(
+                "The DeepSeek V4 FP8 WO_A path requires AITER >= 0.1.20; "
+                "falling back to BF16 WO_A."
+            )
+            return
+        del mxscale_op, inverse_quant_op
+
+        from vllm.model_executor.layers.quantization.utils.fp8_utils import (
+            get_fp8_block_weight_scale,
+            is_fp8,
+        )
+
+        weight = getattr(self.wo_a, "weight", None)
+        scale = get_fp8_block_weight_scale(self.wo_a)
+        if scale is None:
+            # ModelOpt MXFP8 stores the multiplicative E8M0 scale without the
+            # historical ``_inv`` suffix.
+            scale = getattr(self.wo_a, "weight_scale", None)
+        if weight is None or scale is None:
+            logger.warning_once(
+                "DeepSeek V4 FP8 WO_A needs a block-scaled FP8 wo_a weight; "
+                "the layer exposes no weight/weight scale. Falling back to "
+                "BF16 WO_A."
+            )
+            return
+        if weight.dim() != 2 or scale.dim() != 2 or not is_fp8(weight.dtype):
+            logger.warning_once(
+                "DeepSeek V4 FP8 WO_A needs a 2-D FP8 wo_a weight with a 2-D "
+                "block scale, got weight %s%s and scale %s. Falling back to "
+                "BF16 WO_A.",
+                weight.dtype,
+                tuple(weight.shape),
+                tuple(scale.shape),
+            )
+            return
+
+        groups = self.n_local_groups
+        out_per_group = self.o_lora_rank
+        out_features, in_features = weight.shape
+        if (
+            out_features != groups * out_per_group
+            or out_per_group % 128 != 0
+            or in_features % 128 != 0
+            or scale.shape != (out_features // 128, in_features // 128)
+        ):
+            logger.warning_once(
+                "DeepSeek V4 FP8 WO_A needs group-128 blocks for %d groups of "
+                "%d outputs, got weight %s and scale %s. Falling back to BF16 "
+                "WO_A.",
+                groups,
+                out_per_group,
+                tuple(weight.shape),
+                tuple(scale.shape),
+            )
+            return
+
+        e8m0_scale = _wo_a_block_scale_to_e8m0(scale)
+        if e8m0_scale is None:
+            logger.warning_once(
+                "DeepSeek V4 FP8 WO_A could not losslessly encode the %s wo_a "
+                "block scale as OCP E8M0. Falling back to BF16 WO_A.",
+                scale.dtype,
+            )
+            return
+
+        self._wo_a_fp8_weight = weight.view(groups, out_per_group, in_features)
+        self._wo_a_e8m0_scale = e8m0_scale.view(
+            groups, out_per_group // 128, in_features // 128
+        )
+        cache = getattr(self.rotary_emb, "cos_sin_cache_bf16", None)
+        if cache is None:
+            cache = self.rotary_emb.cos_sin_cache.to(dtype=torch.bfloat16)
+        cos_cache, sin_cache = cache.chunk(2, dim=-1)
+        self._wo_a_cos_cache = cos_cache.contiguous()
+        self._wo_a_sin_cache = sin_cache.contiguous()
 
     def prepare_compressor_gemm_fusion(self) -> bool:
         if self._fused_compressor_weight is not None:
@@ -973,17 +1108,44 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention):
         )
 
     def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
-        # ROCm BF16 reference wo_a path (inverse RoPE + einsum) + wo_b.
-        z = rocm_inv_rope_einsum(
-            self.rotary_emb,
-            o,
-            positions,
-            self.rope_head_dim,
-            self.n_local_groups,
-            self.o_lora_rank,
-            self.wo_a,
-        )
-        zf = z.flatten(1)
+        if self._wo_a_fp8_weight is not None:
+            from aiter.ops.batched_gemm_op_a8w8 import (
+                batched_gemm_a8w8_mxscale,
+            )
+            from aiter.ops.inverse_rope_group_quant import (
+                inverse_rope_group_quant,
+            )
+
+            assert self._wo_a_cos_cache is not None
+            assert self._wo_a_sin_cache is not None
+            o_fp8, o_scale = inverse_rope_group_quant(
+                o.view(o.shape[0], self.n_local_heads, self.head_dim),
+                positions.to(torch.int64),
+                self._wo_a_cos_cache,
+                self._wo_a_sin_cache,
+                num_groups=self.n_local_groups,
+                quant_group_size=128,
+            )
+            assert self._wo_a_e8m0_scale is not None
+            zf = batched_gemm_a8w8_mxscale(
+                o_fp8,
+                self._wo_a_fp8_weight,
+                o_scale,
+                self._wo_a_e8m0_scale,
+                dtype=o.dtype,
+            ).flatten(1)
+        else:
+            # ROCm BF16 reference wo_a path (inverse RoPE + einsum) + wo_b.
+            z = rocm_inv_rope_einsum(
+                self.rotary_emb,
+                o,
+                positions,
+                self.rope_head_dim,
+                self.n_local_groups,
+                self.o_lora_rank,
+                self.wo_a,
+            )
+            zf = z.flatten(1)
         if self._wo_b_scale is not None and zf.dim() == 2:
             return self._bpre_attn_gemm(self.wo_b.weight, self._wo_b_scale, zf, True)
         return self.wo_b(zf)

vllm/models/deepseek_v4/common/vision.py

@@ -40,8 +40,7 @@ from vllm.model_executor.models.vision import (
 )
 
 
-@lru_cache(8)
-def get_vision_cos_sin(
+def _compute_vision_cos_sin(
     n_h: int, n_w: int, dim: int, theta: float
 ) -> tuple[torch.Tensor, torch.Tensor]:
     inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim))
@@ -52,6 +51,13 @@ def get_vision_cos_sin(
     return freqs.cos().unsqueeze(1), freqs.sin().unsqueeze(1)
 
 
+@lru_cache(8)
+def get_vision_cos_sin(
+    n_h: int, n_w: int, dim: int, theta: float
+) -> tuple[torch.Tensor, torch.Tensor]:
+    return _compute_vision_cos_sin(n_h, n_w, dim, theta)
+
+
 def apply_rotary(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
     dtype = x.dtype
     x1, x2 = x.float().chunk(2, dim=-1)
@@ -123,15 +129,23 @@ class DeepseekV4VisionAttention(nn.Module):
         )
 
     def forward(
-        self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
+        self,
+        x: torch.Tensor,
+        cos: torch.Tensor,
+        sin: torch.Tensor,
+        cu_seqlens: torch.Tensor | None = None,
+        max_seqlen: torch.Tensor | None = None,
     ) -> torch.Tensor:
         n = x.size(0)
         qkv, _ = self.wqkv(x)
         q, k, v = (t.view(n, self.n_heads, self.head_dim) for t in qkv.chunk(3, -1))
         q = apply_rotary(q, cos, sin).unsqueeze(0)  # (b=1, n, h, d)
         k = apply_rotary(k, cos, sin).unsqueeze(0)
-        # One image per call: a dense batch, no varlen packing metadata.
-        o = self.attn(q, k, v.unsqueeze(0))
+        # Dense (cu_seqlens=None): one image per call. Varlen (cu_seqlens
+        # set): packed multi-image batch for encoder CUDA graph replay.
+        o = self.attn(
+            q, k, v.unsqueeze(0), cu_seqlens=cu_seqlens, max_seqlen=max_seqlen
+        )
         out, _ = self.wo(o.reshape(n, -1))
         return out
 
@@ -173,9 +187,14 @@ class DeepseekV4VisionBlock(nn.Module):
         self.mlp = DeepseekV4VisionMLP(config, prefix=f"{prefix}.mlp")
 
     def forward(
-        self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
+        self,
+        x: torch.Tensor,
+        cos: torch.Tensor,
+        sin: torch.Tensor,
+        cu_seqlens: torch.Tensor | None = None,
+        max_seqlen: torch.Tensor | None = None,
     ) -> torch.Tensor:
-        x = x + self.attn(self.norm1(x), cos, sin)
+        x = x + self.attn(self.norm1(x), cos, sin, cu_seqlens, max_seqlen)
         return x + self.mlp(self.norm2(x))
 
 
@@ -206,6 +225,115 @@ class DeepseekV4ViT(nn.Module):
             x = block(x, cos, sin)
         return self.norm(x)
 
+    def forward_packed(
+        self,
+        patches: torch.Tensor,
+        cos: torch.Tensor,
+        sin: torch.Tensor,
+        cu_seqlens: torch.Tensor,
+        max_seqlen: torch.Tensor,
+    ) -> torch.Tensor:
+        """Varlen path for encoder CUDA graphs: multiple images packed along
+        rows, per-image attention via ``cu_seqlens``, RoPE tables precomputed
+        on device. Pure-tensor (no host reads, no H2D), safe to capture."""
+        x = self.patch_embed(patches)
+        for block in self.blocks:
+            x = block(x, cos, sin, cu_seqlens, max_seqlen)
+        return self.norm(x)
+
+
+def build_packed_vit_metadata(
+    grids: list[list[int]] | list[tuple[int, int]],
+    *,
+    rope_dim: int,
+    rope_theta: float,
+    device: torch.device,
+    max_seqlen_override: int | None = None,
+    cached: bool = True,
+) -> dict[str, torch.Tensor]:
+    """Precompute packed-batch ViT metadata for ``DeepseekV4ViT.forward_packed``.
+
+    Args:
+        grids: ``[n_vit_h, n_vit_w]`` per image, in packing order.
+        rope_dim: RoPE embedding dimension of the ViT.
+        rope_theta: RoPE base frequency of the ViT.
+        device: Device for the returned tensors (except ``max_seqlen``,
+            which stays on the host).
+        max_seqlen_override: Worst-case value baked in at CUDA graph capture
+            (the attention wrapper reads ``max_seqlen`` on the host, so the
+            capture-time value becomes a graph constant).
+        cached: Use the shared ``lru_cache`` for RoPE tables. Capture-time
+            dummy grids pass ``False`` to avoid evicting real entries.
+
+    """
+    cos_sin = get_vision_cos_sin if cached else _compute_vision_cos_sin
+    cos_list: list[torch.Tensor] = []
+    sin_list: list[torch.Tensor] = []
+    cu_seqlens = [0]
+    max_seqlen = 0
+    for n_h, n_w in grids:
+        cos, sin = cos_sin(n_h, n_w, rope_dim, rope_theta)
+        cos_list.append(cos)
+        sin_list.append(sin)
+        n = n_h * n_w
+        cu_seqlens.append(cu_seqlens[-1] + n)
+        max_seqlen = max(max_seqlen, n)
+    if cos_list:
+        cos = torch.cat(cos_list).to(device)
+        sin = torch.cat(sin_list).to(device)
+    else:
+        cos = torch.zeros((0, 1, rope_dim), dtype=torch.float32, device=device)
+        sin = torch.zeros((0, 1, rope_dim), dtype=torch.float32, device=device)
+    if max_seqlen_override is not None:
+        max_seqlen = max_seqlen_override
+    return {
+        "vit_cos": cos,
+        "vit_sin": sin,
+        "cu_seqlens": torch.tensor(cu_seqlens, dtype=torch.int32, device=device),
+        # Read on the host by the attention wrapper; keep on CPU.
+        "max_seqlen": torch.tensor(max_seqlen, dtype=torch.int32),
+    }
+
+
+def build_packed_merge_metadata(
+    grids: list[list[int]] | list[tuple[int, int]],
+    downsample_ratio: int,
+    *,
+    device: torch.device,
+    dtype: torch.dtype,
+) -> dict[str, torch.Tensor]:
+    """Gather indices/mask for ``DeepseekV4Aligner.forward_packed``.
+
+    For each aligner output row, the ``r x r`` source patch rows in the
+    packed ViT output; positions past the image edge (the eager path pads
+    them with zeros) get mask 0 and a clamped in-range index.
+    """
+    r = downsample_ratio
+    idx_list: list[torch.Tensor] = []
+    mask_list: list[torch.Tensor] = []
+    offset = 0
+    for n_h, n_w in grids:
+        rows = -(-n_h // r)
+        cols = -(-n_w // r)
+        bi = torch.arange(rows).repeat_interleave(cols)
+        bj = torch.arange(cols).repeat(rows)
+        di = torch.arange(r).repeat_interleave(r)
+        dj = torch.arange(r).repeat(r)
+        hi = bi[:, None] * r + di[None, :]
+        wj = bj[:, None] * r + dj[None, :]
+        valid = (hi < n_h) & (wj < n_w)
+        idx = offset + hi.clamp(max=n_h - 1) * n_w + wj.clamp(max=n_w - 1)
+        idx_list.append(idx)
+        mask_list.append(valid)
+        offset += n_h * n_w
+    if idx_list:
+        merge_idx = torch.cat(idx_list).to(device)
+        merge_mask = torch.cat(mask_list).unsqueeze(-1).to(device=device, dtype=dtype)
+    else:
+        merge_idx = torch.zeros((0, r * r), dtype=torch.int64, device=device)
+        merge_mask = torch.zeros((0, r * r, 1), dtype=dtype, device=device)
+    return {"merge_idx": merge_idx, "merge_mask": merge_mask}
+
 
 class DeepseekV4Aligner(nn.Module):
     """Spatial merge (downsample_ratio x downsample_ratio) + MLP projector."""
@@ -240,6 +368,27 @@ class DeepseekV4Aligner(nn.Module):
         out, _ = self.w2(F.gelu(hidden))
         return out
 
+    def forward_packed(
+        self,
+        x: torch.Tensor,
+        merge_idx: torch.Tensor,
+        merge_mask: torch.Tensor,
+    ) -> torch.Tensor:
+        """Packed multi-image path for encoder CUDA graphs.
+
+        ``merge_idx``/``merge_mask`` come from ``build_packed_merge_metadata``;
+        gather+mask reproduces the eager ``F.pad`` + ``F.unfold`` merge
+        exactly (unfold rows are channel-major, so the gathered patch rows
+        are transposed before flattening).
+        """
+        r2 = self.downsample_ratio**2
+        m = merge_idx.shape[0]
+        gathered = (x[merge_idx] * merge_mask).view(m, r2, -1)
+        gathered = gathered.permute(0, 2, 1).reshape(m, -1)
+        hidden, _ = self.w1(gathered)
+        out, _ = self.w2(F.gelu(hidden))
+        return out
+
 
 def run_dp_sharded_vision_tower(
     vision_model: DeepseekV4ViT,

vllm/models/deepseek_v41/amd/rocm.py

@@ -13,7 +13,10 @@ from vllm.distributed import (
 )
 from vllm.forward_context import get_forward_context
 from vllm.logger import init_logger
-from vllm.models.deepseek_v41.attention import DeepseekV4Attention
+from vllm.models.deepseek_v41.attention import (
+    DeepseekV4Attention,
+    _replace_layer_index,
+)
 from vllm.models.deepseek_v41.common.ops import dequantize_and_gather_k_cache
 from vllm.models.deepseek_v41.sparse_mla import (
     DeepseekV4FlashMLAMetadata,
@@ -391,6 +394,10 @@ def _copy_ragged_to_graph_buffers(
     return ragged_out, indptr_out
 
 
+# (ragged_indices, ragged_indptr, lens) as consumed by rocm_sparse_attn_decode.
+_TopkRagged = tuple[torch.Tensor, torch.Tensor, torch.Tensor]
+
+
 @dataclass
 class DeepseekV4ROCMAiterSparseSWAMetadata(DeepseekSparseSWAMetadata):
     decode_swa_ragged_indices: torch.Tensor | None = None
@@ -500,6 +507,23 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention):
         self._fused_compressor_weight: torch.Tensor | None
         self.register_buffer("_fused_compressor_weight", None, persistent=False)
         self._fused_compressor_split_sizes: tuple[int, int] | None = None
+        # Decode ragged topk metadata is a pure function of the indices its
+        # index source published, so every consumer of a source rebuilds the
+        # same thing. The source owns one cache per compress ratio, refreshed
+        # when it runs; consumers below it read through.
+        self._topk_ragged_cache: dict[int, _TopkRagged] = {}
+        self._index_source_prefix: str | None = None
+        if self.compress_ratio > 0:
+            assert self.index_source_layer_id is not None
+            self._index_source_prefix = _replace_layer_index(
+                self.prefix, self.index_source_layer_id
+            )
+            if self._index_source_prefix not in self._static_forward_context:
+                raise NotImplementedError(
+                    f"Index source {self._index_source_prefix} not found on "
+                    "this rank; PP splits inside a v4.1 index-sharing group "
+                    "are not supported."
+                )
 
     @classmethod
     def get_padded_num_q_heads(cls, num_heads: int) -> int:
@@ -733,6 +757,42 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention):
                 output=output[:num_decode_tokens],
             )
 
+    def _decode_topk_ragged(
+        self,
+        swa_metadata: DeepseekV4ROCMAiterSparseSWAMetadata,
+        attn_metadata: DeepseekV4FlashMLAMetadata | None,
+        num_decodes: int,
+        num_decode_tokens: int,
+    ) -> _TopkRagged:
+        """Ragged form of the topk indices this layer's index source published.
+
+        The packing depends only on the shared ``topk_indices_buffer`` and
+        per-step metadata, apart from the layer's own compress ratio, so the
+        source memoizes one result per ratio for all the layers below it.
+        """
+        assert attn_metadata is not None
+        assert swa_metadata.is_valid_token is not None
+        assert self.topk_indices_buffer is not None
+        assert self._index_source_prefix is not None
+
+        source = self._static_forward_context[self._index_source_prefix]
+        if source is self:
+            # Fresh indices as of this layer; drop what the last step cached.
+            source._topk_ragged_cache = {}
+        cached = source._topk_ragged_cache.get(self.compress_ratio)
+        if cached is not None:
+            return cached
+
+        built = compute_global_topk_ragged_indices_and_indptr(
+            self.topk_indices_buffer[:num_decode_tokens],
+            swa_metadata.token_to_req_indices,
+            attn_metadata.block_table[:num_decodes],
+            attn_metadata.block_size // self.compress_ratio,
+            swa_metadata.is_valid_token[:num_decode_tokens],
+        )
+        source._topk_ragged_cache[self.compress_ratio] = built
+        return built
+
     def _forward_decode(
         self,
         q: torch.Tensor,
@@ -749,22 +809,15 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention):
         topk_ragged_indices = None
         topk_ragged_indptr = None
         if not swa_only:
-            # Local indices filled by the index-source layer's indexer.
-            assert attn_metadata is not None
-            assert swa_metadata.is_valid_token is not None
-            assert self.topk_indices_buffer is not None
-            block_size = attn_metadata.block_size // self.compress_ratio
-            is_valid = swa_metadata.is_valid_token[:num_decode_tokens]
             (
                 topk_ragged_indices,
                 topk_ragged_indptr,
                 topk_lens,
-            ) = compute_global_topk_ragged_indices_and_indptr(
-                self.topk_indices_buffer[:num_decode_tokens],
-                swa_metadata.token_to_req_indices,
-                attn_metadata.block_table[:num_decodes],
-                block_size,
-                is_valid,
+            ) = self._decode_topk_ragged(
+                swa_metadata=swa_metadata,
+                attn_metadata=attn_metadata,
+                num_decodes=num_decodes,
+                num_decode_tokens=num_decode_tokens,
             )
 
         rocm_sparse_attn_decode(

vllm/models/deepseek_v41/amd/vl_model.py

@@ -26,6 +26,7 @@ from vllm.distributed import get_tensor_model_parallel_world_size
 from vllm.model_executor.models.interfaces import (
     MultiModalEmbeddings,
     SupportsEagle3,
+    SupportsEncoderCudaGraph,
     SupportsMultiModal,
     SupportsPP,
 )
@@ -44,15 +45,12 @@ from vllm.multimodal import MULTIMODAL_REGISTRY
 from vllm.utils.tensor_schema import TensorSchema, TensorShape
 
 from ..common.mm_preprocess import (
-    IMAGE,
-    IMAGE_END,
-    IMAGE_NEW_LINE,
     IMAGE_PLACEHOLDER,
-    IMAGE_START,
     DeepseekV4VLDummyInputsBuilder,
     DeepseekV4VLMultiModalProcessor,
     DeepseekV4VLProcessingInfo,
 )
+from ..common.vl_cudagraph import DeepseekV4VLEncoderCudaGraphMixin
 from .model import (
     DeepseekV41LLMForCausalLM,
     _linear_scale_param_name,
@@ -115,11 +113,20 @@ def _make_deepseek_v4_vl_weights_mapper(
     info=DeepseekV4VLProcessingInfo,
     dummy_inputs=DeepseekV4VLDummyInputsBuilder,
 )
-class DeepseekV41ForCausalLM(nn.Module, SupportsMultiModal, SupportsPP, SupportsEagle3):
+class DeepseekV41ForCausalLM(
+    nn.Module,
+    DeepseekV4VLEncoderCudaGraphMixin,
+    SupportsMultiModal,
+    SupportsEncoderCudaGraph,
+    SupportsPP,
+    SupportsEagle3,
+):
     """Multimodal entry point for DeepSeek-V4.1 checkpoints with a vision tower.
 
     ``SupportsEagle3`` (aux hidden-state plumbing for MTP/DSpark drafters)
     delegates through ``language_model`` via the protocol defaults.
+    ``SupportsEncoderCudaGraph`` (ViT CUDA graph capture/replay) is
+    implemented by ``DeepseekV4VLEncoderCudaGraphMixin``.
     """
 
     supports_encoder_tp_data = True
@@ -196,29 +203,6 @@ class DeepseekV41ForCausalLM(nn.Module, SupportsMultiModal, SupportsPP, Supports
             resolve_bindings={"p": self.config.vision_patch_size},
         )
 
-    def _encode_image(
-        self,
-        patches: torch.Tensor,
-        n_vit_h: int,
-        n_vit_w: int,
-    ) -> torch.Tensor:
-        # Aligner rows in reading order, one per IMAGE slot.
-        return self.aligner(self.vision(patches, n_vit_h, n_vit_w), n_vit_h, n_vit_w)
-
-    def _build_image_span(
-        self, image_embeds: torch.Tensor, types: torch.Tensor
-    ) -> torch.Tensor:
-        """Full image span: aligner rows at IMAGE slots, the learned
-        delimiter vectors at IMAGE_START/IMAGE_NEW_LINE/IMAGE_END."""
-        types = types.to(image_embeds.device)
-        span = image_embeds.new_empty(types.numel(), image_embeds.shape[-1])
-        dtype = image_embeds.dtype
-        span[types == IMAGE_START] = self.image_start.to(dtype)
-        span[types == IMAGE_END] = self.image_end.to(dtype)
-        span[types == IMAGE_NEW_LINE] = self.image_newline.to(dtype)
-        span[types == IMAGE] = image_embeds
-        return span
-
     def _process_image_input(
         self,
         image_input: DeepseekV4VLImagePixelInputs,

vllm/models/deepseek_v41/common/vl_cudagraph.py

@@ -0,0 +1,390 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Encoder CUDA graph support for the DeepSeek-V4.1 vision tower.
+
+Mixin implementing the ``SupportsEncoderCudaGraph`` protocol
+(``vllm/model_executor/models/interfaces.py``) on top of the shared
+``DeepseekV4ViT``/``DeepseekV4Aligner``. The captured graph packs a whole
+batch of images into one varlen run: ViT blocks attend per image via
+``cu_seqlens``, and the aligner's spatial merge becomes a gather+mask
+against precomputed indices. Span assembly (IMAGE_START/IMAGE_NEW_LINE/
+IMAGE_END delimiters) stays eager in ``postprocess_encoder_output`` but is
+batched over all items of the packed batch.
+
+Enabled via ``--compilation-config '{"cudagraph_mm_encoder": true}'``; the
+manager (``vllm/v1/worker/encoder_cudagraph.py``) handles budget packing, DP
+sharding and eager fallback.
+"""
+
+from collections.abc import Hashable
+from typing import TYPE_CHECKING, Any
+
+import torch
+
+from vllm.models.deepseek_v4.common.vision import (
+    build_packed_merge_metadata,
+    build_packed_vit_metadata,
+)
+from vllm.v1.worker.encoder_cudagraph_defs import (
+    EncoderCudaGraphCaptureInputs,
+    EncoderCudaGraphConfig,
+    EncoderCudaGraphReplayBuffers,
+    EncoderItemSpec,
+)
+
+from .mm_preprocess import IMAGE, IMAGE_END, IMAGE_NEW_LINE, IMAGE_START
+
+if TYPE_CHECKING:
+    from vllm.config import VllmConfig
+    from vllm.config.multimodal import MultiModalConfig
+    from vllm.models.deepseek_v4.common.vision import (
+        DeepseekV4Aligner,
+        DeepseekV4ViT,
+    )
+    from vllm.transformers_utils.configs.deepseek_v41 import DeepseekV41Config
+
+
+class DeepseekV4VLEncoderCudaGraphMixin:
+    """``SupportsEncoderCudaGraph`` for ``DeepseekV41ForCausalLM``.
+
+    Expects ``self.vision``/``self.aligner``/``self.config``/
+    ``self.multimodal_config`` and the ``image_start``/``image_newline``/
+    ``image_end`` parameters from the host model.
+    """
+
+    # Provided by the host model; declared here so mixin methods type-check.
+    vision: "DeepseekV4ViT"
+    aligner: "DeepseekV4Aligner"
+    config: "DeepseekV41Config"
+    multimodal_config: "MultiModalConfig | None"
+    image_start: torch.nn.Parameter
+    image_end: torch.nn.Parameter
+    image_newline: torch.nn.Parameter
+
+    # -- helpers shared with the eager multimodal path --
+
+    def _encode_image(
+        self,
+        patches: torch.Tensor,
+        n_vit_h: int,
+        n_vit_w: int,
+    ) -> torch.Tensor:
+        # Aligner rows in reading order, one per IMAGE slot.
+        return self.aligner(self.vision(patches, n_vit_h, n_vit_w), n_vit_h, n_vit_w)
+
+    def _build_image_span(
+        self, image_embeds: torch.Tensor, types: torch.Tensor
+    ) -> torch.Tensor:
+        """Full image span: aligner rows at IMAGE slots, the learned
+        delimiter vectors at IMAGE_START/IMAGE_NEW_LINE/IMAGE_END."""
+        types = types.to(image_embeds.device)
+        span = image_embeds.new_empty(types.numel(), image_embeds.shape[-1])
+        dtype = image_embeds.dtype
+        span[types == IMAGE_START] = self.image_start.to(dtype)
+        span[types == IMAGE_END] = self.image_end.to(dtype)
+        span[types == IMAGE_NEW_LINE] = self.image_newline.to(dtype)
+        span[types == IMAGE] = image_embeds
+        return span
+
+    @staticmethod
+    def _get_grid_list(mm_kwargs: dict[str, Any], key: str) -> list[list[int]]:
+        grid = mm_kwargs[key]
+        if isinstance(grid, torch.Tensor):
+            grid = grid.tolist()
+        return [[int(v) for v in row] for row in grid]
+
+    _encoder_cg_pad_totals: dict[int, int] = {}
+
+    # -- SupportsEncoderCudaGraph protocol --
+
+    def get_encoder_cudagraph_config(self) -> EncoderCudaGraphConfig:
+        mm_cfg = self.multimodal_config
+        if (
+            mm_cfg is not None
+            and mm_cfg.mm_encoder_attn_dtype == "fp8"
+            and mm_cfg.mm_encoder_fp8_scale_path is None
+        ):
+            raise ValueError(
+                "cudagraph_mm_encoder is incompatible with dynamic FP8 "
+                "scaling for the vision encoder (the amax history is "
+                "host-side state that cannot be captured). Provide a static "
+                "scale file via --mm-encoder-fp8-scale-path or disable "
+                "cudagraph_mm_encoder."
+            )
+
+        from vllm.v1.attention.backends.registry import AttentionBackendEnum
+
+        backend = self.vision.blocks[0].attn.attn.attn_backend
+        # FlashInfer reads max_seqlen on the host and TORCH_SDPA calls
+        # .tolist() on CUDA cu_seqlens; neither can be captured.
+        if backend in (
+            AttentionBackendEnum.FLASHINFER,
+            AttentionBackendEnum.TORCH_SDPA,
+        ):
+            raise ValueError(
+                f"cudagraph_mm_encoder is not supported with the {backend.name} "
+                "ViT attention backend for DeepSeek-V4.1. Set "
+                "--mm-encoder-attn-backend FLASH_ATTN or disable "
+                "cudagraph_mm_encoder."
+            )
+
+        pad_totals = self._encoder_cg_pad_totals
+
+        def pad_cu_seqlens(dst: torch.Tensor, src: torch.Tensor) -> None:
+            # Varlen attention requires cu_seqlens[-1] to equal the number of
+            # rows actually passed in. The captured buffers are sized for the
+            # full token budget, so a smaller real batch is completed with one
+            # trailing padding sequence; declaring fewer rows than the buffer
+            # holds is undefined behaviour and returns NaN on FlashAttn.
+            total = pad_totals.get(dst.data_ptr())
+            if total is None:
+                raise RuntimeError(
+                    "cu_seqlens replay buffer was not registered at capture "
+                    "time; the manager must replay into the capture-time "
+                    "buffers."
+                )
+            n = min(src.shape[0], dst.shape[0])
+            dst[:n].copy_(src[:n])
+            dst[n:] = total
+
+        return EncoderCudaGraphConfig(
+            modalities=["image"],
+            buffer_keys=[
+                "patches",
+                "vit_cos",
+                "vit_sin",
+                "cu_seqlens",
+                "max_seqlen",
+                "merge_idx",
+                "merge_mask",
+            ],
+            out_hidden_size=self.config.hidden_size,
+            padding_logics={"cu_seqlens": pad_cu_seqlens},
+        )
+
+    def get_encoder_cudagraph_budget_range(
+        self,
+        vllm_config: "VllmConfig",
+    ) -> tuple[int, int]:
+        # Min: comfortably above the smallest possible image span (4 tokens
+        # for a 1x1 grid); also sets max_batch_size = 64 via the manager's
+        # auto-inference (max_batch_size <= min_budget invariant).
+        min_budget = 64
+        max_budget = min(
+            vllm_config.scheduler_config.max_num_batched_tokens,
+            vllm_config.model_config.max_model_len,
+        )
+        return (min_budget, max_budget)
+
+    def get_encoder_cudagraph_item_specs(
+        self,
+        mm_kwargs: dict[str, Any],
+    ) -> list[EncoderItemSpec]:
+        vit_grid = self._get_grid_list(mm_kwargs, "vit_grid")
+        llm_grid = self._get_grid_list(mm_kwargs, "llm_grid")
+        return [
+            EncoderItemSpec(
+                input_size=h * w,
+                output_tokens=lh * (lw + 1) + 2,
+            )
+            for (h, w), (lh, lw) in zip(vit_grid, llm_grid, strict=True)
+        ]
+
+    def select_encoder_cudagraph_items(
+        self,
+        mm_kwargs: dict[str, Any],
+        indices: list[int],
+    ) -> dict[str, Any]:
+        vit_grid = self._get_grid_list(mm_kwargs, "vit_grid")
+        llm_grid = self._get_grid_list(mm_kwargs, "llm_grid")
+        patches = mm_kwargs["patches"]
+        types = mm_kwargs["types"]
+
+        if len(indices) == 0:
+            return {
+                "patches": patches[:0],
+                "vit_grid": torch.zeros((0, 2), dtype=torch.int64),
+                "llm_grid": torch.zeros((0, 2), dtype=torch.int64),
+                "types": types[:0],
+            }
+
+        cum_patches = [0]
+        for h, w in vit_grid:
+            cum_patches.append(cum_patches[-1] + h * w)
+        cum_spans = [0]
+        for lh, lw in llm_grid:
+            cum_spans.append(cum_spans[-1] + lh * (lw + 1) + 2)
+
+        return {
+            "patches": torch.cat(
+                [patches[cum_patches[i] : cum_patches[i + 1]] for i in indices]
+            ),
+            "vit_grid": torch.tensor([vit_grid[i] for i in indices], dtype=torch.int64),
+            "llm_grid": torch.tensor([llm_grid[i] for i in indices], dtype=torch.int64),
+            "types": torch.cat(
+                [types[cum_spans[i] : cum_spans[i + 1]] for i in indices]
+            ),
+        }
+
+    def prepare_encoder_cudagraph_capture_inputs(
+        self,
+        token_budget: int,
+        max_batch_size: int,
+        max_frames_per_batch: int,
+        device: torch.device,
+        dtype: torch.dtype,
+        path: str = "default",
+        axis_keys: tuple[Hashable, ...] | None = None,
+    ) -> EncoderCudaGraphCaptureInputs:
+        if (
+            self.multimodal_config is not None
+            and self.multimodal_config.get_limit_per_prompt("image") == 0
+        ):
+            raise RuntimeError(
+                "cudagraph_mm_encoder requires the vision tower, but it was "
+                "stubbed out because the image limit is 0. Disable "
+                "cudagraph_mm_encoder."
+            )
+
+        config = self.config
+        r = config.vision_downsample_ratio
+        # Span rows per dummy item (ceiling so the buffers tile the budget).
+        # Each dummy image is a single aligner row of per_item_out blocks:
+        # r x (r * per_item_out) patches. One real image using the whole
+        # budget needs at most r^2 * token_budget patches, which this fits.
+        per_item_out = -(-token_budget // max_batch_size)
+        grids = [[r, r * per_item_out]] * max_batch_size
+        total_patches = r * r * per_item_out * max_batch_size
+
+        p = config.vision_patch_size
+        patches = torch.zeros(total_patches, 3, p, p, device=device, dtype=dtype)
+
+        metadata = build_packed_vit_metadata(
+            grids,
+            rope_dim=self.vision.rope_dim,
+            rope_theta=self.vision.rope_theta,
+            device=device,
+            max_seqlen_override=total_patches,
+            cached=False,
+        )
+        # Spare cu_seqlens slots let replay append a padding sequence that
+        # covers the rows a smaller real batch does not fill.
+        real_cu = metadata.pop("cu_seqlens")
+        cu_seqlens = torch.full(
+            (max_batch_size + 2,), total_patches, dtype=torch.int32, device=device
+        )
+        cu_seqlens[: real_cu.numel()] = real_cu
+
+        merge = build_packed_merge_metadata(grids, r, device=device, dtype=dtype)
+
+        self._encoder_cg_pad_totals[cu_seqlens.data_ptr()] = total_patches
+
+        return EncoderCudaGraphCaptureInputs(
+            values={
+                "patches": patches,
+                **metadata,
+                "cu_seqlens": cu_seqlens,
+                **merge,
+            }
+        )
+
+    def prepare_encoder_cudagraph_replay_buffers(
+        self,
+        mm_kwargs: dict[str, Any],
+        max_batch_size: int,
+        max_frames_per_batch: int,
+        path: str = "default",
+    ) -> EncoderCudaGraphReplayBuffers:
+        vit_grid = self._get_grid_list(mm_kwargs, "vit_grid")
+        patches = mm_kwargs["patches"]
+        dtype = self.aligner.w1.weight.dtype
+        if patches.dtype != dtype:
+            patches = patches.to(dtype)
+
+        # Unpadded: the manager zero-pads patches/cos/sin/merge buffers, and
+        # pad_cu_seqlens appends the padding sequence covering the tail rows.
+        metadata = build_packed_vit_metadata(
+            vit_grid,
+            rope_dim=self.vision.rope_dim,
+            rope_theta=self.vision.rope_theta,
+            device=patches.device,
+        )
+        merge = build_packed_merge_metadata(
+            vit_grid,
+            self.config.vision_downsample_ratio,
+            device=patches.device,
+            dtype=dtype,
+        )
+        return EncoderCudaGraphReplayBuffers(
+            values={"patches": patches, **metadata, **merge}
+        )
+
+    def encoder_cudagraph_forward(
+        self,
+        inputs: dict[str, torch.Tensor],
+        path: str = "default",
+    ) -> torch.Tensor:
+        vit_out = self.vision.forward_packed(
+            inputs["patches"],
+            inputs["vit_cos"],
+            inputs["vit_sin"],
+            inputs["cu_seqlens"],
+            inputs["max_seqlen"],
+        )
+        return self.aligner.forward_packed(
+            vit_out, inputs["merge_idx"], inputs["merge_mask"]
+        )
+
+    def encoder_eager_forward(
+        self,
+        mm_kwargs: dict[str, Any],
+        path: str = "default",
+    ) -> torch.Tensor:
+        patches = mm_kwargs["patches"].to(self.aligner.w1.weight.dtype)
+        vit_grid = self._get_grid_list(mm_kwargs, "vit_grid")
+        outs: list[torch.Tensor] = []
+        offset = 0
+        for n_vit_h, n_vit_w in vit_grid:
+            n_vit = n_vit_h * n_vit_w
+            outs.append(
+                self._encode_image(patches[offset : offset + n_vit], n_vit_h, n_vit_w)
+            )
+            offset += n_vit
+        if not outs:
+            return patches.new_zeros((0, self.config.hidden_size))
+        return torch.cat(outs, dim=0)
+
+    def postprocess_encoder_output(
+        self,
+        outputs: dict[str, torch.Tensor],
+        indices: list[int],
+        per_item_out_tokens: list[int],
+        dest: dict[int, torch.Tensor] | list[torch.Tensor | None],
+        clone: bool = False,
+        batch_mm_kwargs: dict[str, Any] | None = None,
+    ) -> None:
+        if not indices:
+            return
+        assert batch_mm_kwargs is not None
+        aligner_out = outputs["default"]
+        r = self.config.vision_downsample_ratio
+        vit_grid = self._get_grid_list(batch_mm_kwargs, "vit_grid")
+        llm_grid = self._get_grid_list(batch_mm_kwargs, "llm_grid")
+        types = batch_mm_kwargs["types"].to(aligner_out.device)
+
+        n_rows = sum(-(-h // r) * (-(-w // r)) for h, w in vit_grid)
+        span_lens = [lh * (lw + 1) + 2 for lh, lw in llm_grid]
+
+        # Batched span assembly: one masked fill per role across all items.
+        dtype = aligner_out.dtype
+        span = aligner_out.new_empty(sum(span_lens), self.config.hidden_size)
+        span[types == IMAGE] = aligner_out[:n_rows]
+        span[types == IMAGE_START] = self.image_start.to(dtype)
+        span[types == IMAGE_END] = self.image_end.to(dtype)
+        span[types == IMAGE_NEW_LINE] = self.image_newline.to(dtype)
+
+        # Freshly allocated, so later replays cannot clobber the results.
+        offset = 0
+        for idx, span_len in zip(indices, span_lens):
+            dest[idx] = span[offset : offset + span_len]
+            offset += span_len

vllm/models/deepseek_v41/nvidia/vl_model.py

@@ -26,6 +26,7 @@ from vllm.distributed import get_tensor_model_parallel_world_size
 from vllm.model_executor.models.interfaces import (
     MultiModalEmbeddings,
     SupportsEagle3,
+    SupportsEncoderCudaGraph,
     SupportsMultiModal,
     SupportsPP,
 )
@@ -44,15 +45,12 @@ from vllm.multimodal import MULTIMODAL_REGISTRY
 from vllm.utils.tensor_schema import TensorSchema, TensorShape
 
 from ..common.mm_preprocess import (
-    IMAGE,
-    IMAGE_END,
-    IMAGE_NEW_LINE,
     IMAGE_PLACEHOLDER,
-    IMAGE_START,
     DeepseekV4VLDummyInputsBuilder,
     DeepseekV4VLMultiModalProcessor,
     DeepseekV4VLProcessingInfo,
 )
+from ..common.vl_cudagraph import DeepseekV4VLEncoderCudaGraphMixin
 from .model import (
     DeepseekV41LLMForCausalLM,
     _linear_scale_param_name,
@@ -115,11 +113,20 @@ def _make_deepseek_v4_vl_weights_mapper(
     info=DeepseekV4VLProcessingInfo,
     dummy_inputs=DeepseekV4VLDummyInputsBuilder,
 )
-class DeepseekV41ForCausalLM(nn.Module, SupportsMultiModal, SupportsPP, SupportsEagle3):
+class DeepseekV41ForCausalLM(
+    nn.Module,
+    DeepseekV4VLEncoderCudaGraphMixin,
+    SupportsMultiModal,
+    SupportsEncoderCudaGraph,
+    SupportsPP,
+    SupportsEagle3,
+):
     """Multimodal entry point for DeepSeek-V4.1 checkpoints with a vision tower.
 
     ``SupportsEagle3`` (aux hidden-state plumbing for MTP/DSpark drafters)
     delegates through ``language_model`` via the protocol defaults.
+    ``SupportsEncoderCudaGraph`` (ViT CUDA graph capture/replay) is
+    implemented by ``DeepseekV4VLEncoderCudaGraphMixin``.
     """
 
     supports_encoder_tp_data = True
@@ -196,29 +203,6 @@ class DeepseekV41ForCausalLM(nn.Module, SupportsMultiModal, SupportsPP, Supports
             resolve_bindings={"p": self.config.vision_patch_size},
         )
 
-    def _encode_image(
-        self,
-        patches: torch.Tensor,
-        n_vit_h: int,
-        n_vit_w: int,
-    ) -> torch.Tensor:
-        # Aligner rows in reading order, one per IMAGE slot.
-        return self.aligner(self.vision(patches, n_vit_h, n_vit_w), n_vit_h, n_vit_w)
-
-    def _build_image_span(
-        self, image_embeds: torch.Tensor, types: torch.Tensor
-    ) -> torch.Tensor:
-        """Full image span: aligner rows at IMAGE slots, the learned
-        delimiter vectors at IMAGE_START/IMAGE_NEW_LINE/IMAGE_END."""
-        types = types.to(image_embeds.device)
-        span = image_embeds.new_empty(types.numel(), image_embeds.shape[-1])
-        dtype = image_embeds.dtype
-        span[types == IMAGE_START] = self.image_start.to(dtype)
-        span[types == IMAGE_END] = self.image_end.to(dtype)
-        span[types == IMAGE_NEW_LINE] = self.image_newline.to(dtype)
-        span[types == IMAGE] = image_embeds
-        return span
-
     def _process_image_input(
         self,
         image_input: DeepseekV4VLImagePixelInputs,

vllm/models/glm5next/nvidia/ops/kpool_compress.py

@@ -383,6 +383,8 @@ def _kpool_tail_seed_kernel(
     tslot_ptr,
     tail_ptr,
     n_tokens,
+    TAIL_BLOCK_ELEMS: tl.constexpr,
+    KPOOL_HEAD: tl.constexpr,
     HEAD_DIM: tl.constexpr,
     KPOOL: tl.constexpr,
     BLOCK_D: tl.constexpr,
@@ -393,6 +395,11 @@ def _kpool_tail_seed_kernel(
     ahead belongs to a different tail block (or is past the batch / padding,
     slot < 0). ``tslot = block * KPOOL + pos % KPOOL``; the destination is
     ``tail[block, {0:K, 1:score}, pos % KPOOL, :]``.
+
+    The tail cache aliases the indexer cache with the indexer's (padded) block
+    stride, so blocks are addressed through ``TAIL_BLOCK_ELEMS`` /
+    ``KPOOL_HEAD`` (``tail.stride(0)`` / ``tail.stride(1)``), never as a dense
+    ``[num_blocks, 2, KPOOL, HEAD_DIM]`` array.
     """
     i = tl.program_id(0)
     t = tl.load(tslot_ptr + i).to(tl.int64)
@@ -409,11 +416,11 @@ def _kpool_tail_seed_kernel(
         return
     offs = tl.arange(0, BLOCK_D)
     m = offs < HEAD_DIM
-    base = (blk * 2 * KPOOL + t % KPOOL) * HEAD_DIM
+    base = blk * TAIL_BLOCK_ELEMS + (t % KPOOL) * HEAD_DIM
     k = tl.load(key_ptr + i * HEAD_DIM + offs, mask=m)
     s = tl.load(score_ptr + i * HEAD_DIM + offs, mask=m)
     tl.store(tail_ptr + base + offs, k, mask=m)
-    tl.store(tail_ptr + base + KPOOL * HEAD_DIM + offs, s, mask=m)
+    tl.store(tail_ptr + base + KPOOL_HEAD + offs, s, mask=m)
 
 
 def kpool_seed_tail_cache(
@@ -426,6 +433,8 @@ def kpool_seed_tail_cache(
 ) -> None:
     """Seed the paged tail cache from a prefill batch (see the kernel)."""
     assert tail_kv_cache.dtype == torch.bfloat16
+    assert tail_kv_cache.ndim == 4 and tail_kv_cache.shape[1] == 2
+    assert tail_kv_cache.stride(3) == 1 and tail_kv_cache.stride(2) == head_dim
     assert key.dtype == torch.bfloat16
     n = tslot.shape[0]
     if n == 0:
@@ -436,6 +445,8 @@ def kpool_seed_tail_cache(
         tslot,
         tail_kv_cache,
         n,
+        TAIL_BLOCK_ELEMS=tail_kv_cache.stride(0),
+        KPOOL_HEAD=tail_kv_cache.stride(1),
         HEAD_DIM=head_dim,
         KPOOL=kpool,
         BLOCK_D=triton.next_power_of_2(head_dim),

vllm/models/glm5next/nvidia/sparse_indexer.py

@@ -10,6 +10,7 @@ from vllm.config import get_current_vllm_config_or_none
 from vllm.forward_context import get_forward_context
 from vllm.logger import init_logger
 from vllm.model_executor.custom_op import CustomOp
+from vllm.model_executor.layers.indexer_topk import get_indexer_topk
 from vllm.models.glm5next.common.sparse_indexer import (
     RADIX_TOPK_WORKSPACE_SIZE,
     _build_decode_scatter_indices,
@@ -39,21 +40,6 @@ if current_platform.is_cuda_alike():
 logger = init_logger(__name__)
 
 
-def _use_cooperative_topk(
-    logits: torch.Tensor,
-    select_k: int,
-    num_rows: int,
-) -> bool:
-    return (
-        current_platform.is_cuda()
-        and select_k in (512, 1024, 2048)
-        and num_rows <= 64
-        and logits.stride(0) % 4 == 0
-        and current_platform.has_device_capability(90)
-        and not current_platform.is_device_capability_family(120)
-    )
-
-
 # kpool write helper: form pools from the current token batch and compress them
 # into the index K cache via the fused Triton kernel.
 
@@ -134,6 +120,7 @@ def sparse_attn_indexer_kpool(
     # path and when the tail cache is disabled.
     tail_kv_cache: torch.Tensor | None = None,
     tail_prefix: str | None = None,
+    topk_backend: str = "auto",
 ) -> torch.Tensor:
     # careful! this will be None in dummy run
     attn_metadata = get_forward_context().attn_metadata
@@ -584,35 +571,18 @@ def sparse_attn_indexer_kpool(
         else:
             topk_dst = topk_indices_buffer[:num_padded_tokens, :topk_tokens]
 
-        if current_platform.is_cuda() and select_k in (512, 1024, 2048):
-            workspace_manager = current_workspace_manager()
-            (topk_workspace,) = workspace_manager.get_simultaneous(
-                ((RADIX_TOPK_WORKSPACE_SIZE,), torch.uint8),
-            )
-            topk_op = (
-                torch.ops._C.cooperative_topk
-                if _use_cooperative_topk(logits, select_k, num_rows)
-                else torch.ops._C.persistent_topk
-            )
-            topk_op(
-                logits,
-                seq_lens,
-                topk_dst,
-                topk_workspace,
-                select_k,
-                attn_metadata_narrowed.max_seq_len,
-            )
-        else:
-            torch.ops._C.top_k_per_row_decode(
-                logits,
-                next_n,
-                seq_lens,
-                topk_dst,
-                num_rows,
-                logits.stride(0),
-                logits.stride(1),
-                select_k,
-            )
+        # Shared dispatcher with the DSA sparse indexer, so the kpool select
+        # honors kernel_config.sparse_indexer_topk_backend and picks between
+        # cooperative/persistent by batch size instead of always taking the
+        # same kernel.
+        get_indexer_topk(topk_backend)(
+            logits,
+            seq_lens,
+            next_n,
+            topk_dst,
+            select_k,
+            attn_metadata_narrowed.max_seq_len,
+        )
 
         # Resolve to token-level indices in the output buffer.
         if index_kpool > 1:
@@ -692,6 +662,11 @@ class SparseAttnIndexerKpool(CustomOp):
                 "Sparse Attention Indexer CUDA op requires DeepGEMM to be installed."
             )
         _cfg = get_current_vllm_config_or_none()
+        self.topk_backend = (
+            _cfg.kernel_config.sparse_indexer_topk_backend
+            if _cfg is not None
+            else "auto"
+        )
         _parallel = _cfg.parallel_config if _cfg is not None else None
         if (
             _parallel is not None
@@ -766,4 +741,5 @@ class SparseAttnIndexerKpool(CustomOp):
             positions,
             self.tail_cache.kv_cache if self.tail_cache is not None else None,
             self.tail_cache.prefix if self.tail_cache is not None else None,
+            self.topk_backend,
         )

vllm/renderers/base.py

@@ -7,11 +7,13 @@ from abc import ABC, abstractmethod
 from collections.abc import Mapping, Sequence
 from concurrent.futures import Future, ThreadPoolExecutor
 from contextlib import ExitStack
+from dataclasses import replace
 from functools import cached_property
 from typing import TYPE_CHECKING, Any, Generic, overload
 
 from typing_extensions import TypeVar
 
+from vllm.exceptions import VLLMValidationError
 from vllm.inputs import (
     EmbedsInput,
     EmbedsPrompt,
@@ -962,27 +964,103 @@ class BaseRenderer(ABC, Generic[_T]):
 
         return engine_input
 
+    def _truncate_expanded_prompt(
+        self,
+        engine_input: TokensInput | MultiModalInput,
+        params: TokenizeParams | None,
+    ) -> None:
+        """Re-apply `truncate_prompt_tokens` once placeholders are expanded.
+
+        `TokenizeParams.apply_post_tokenization` truncates the prompt before
+        multimodal expansion, so a single placeholder token can expand into
+        many and push the result back over the requested bound.
+
+        Multimodal spans are never cut: a request whose spans would not all
+        survive intact is rejected, matching the post-expansion rejection in
+        `vllm.v1.engine.input_processor._validate_model_input`.
+
+        Only multimodal inputs can grow after `apply_post_tokenization`, so
+        everything else returns immediately.
+        """
+        if engine_input["type"] != "multimodal":
+            return
+
+        if params is None or params.truncate_prompt_tokens is None:
+            return
+
+        prompt_token_ids = engine_input.get("prompt_token_ids")
+        if prompt_token_ids is None:
+            return
+
+        num_tokens = len(prompt_token_ids)
+        truncation = params._truncation_slice(self.tokenizer, num_tokens)
+        if truncation is None:
+            return
+
+        start, stop, _ = truncation.indices(num_tokens)
+
+        mm_placeholders = engine_input["mm_placeholders"]
+        for modality, positions in mm_placeholders.items():
+            for position in positions:
+                end = position.offset + position.length
+                if position.offset >= start and end <= stop:
+                    continue
+                raise VLLMValidationError(
+                    f"truncate_prompt_tokens={params.truncate_prompt_tokens} "
+                    f"would split or drop a(n) {modality} item, which spans "
+                    f"tokens {position.offset} to {end} of the "
+                    f"{num_tokens}-token prompt. Multimodal placeholders "
+                    f"expand after truncation is applied, so the limit must "
+                    f"leave every multimodal item intact. Please raise "
+                    f"truncate_prompt_tokens or shorten the multimodal input.",
+                    parameter="truncate_prompt_tokens",
+                    value=params.truncate_prompt_tokens,
+                )
+
+        if start:
+            engine_input["mm_placeholders"] = {
+                modality: [
+                    replace(position, offset=position.offset - start)
+                    for position in positions
+                ]
+                for modality, positions in mm_placeholders.items()
+            }
+
+        engine_input["prompt_token_ids"] = prompt_token_ids[truncation]
+
     def _process_singleton(
         self,
         prompt: SingletonTokPrompt,
         *,
         skip_mm_cache: bool = False,
+        tok_params: TokenizeParams | None = None,
     ) -> SingletonInput:
         if "prompt_embeds" in prompt:
             return self._process_embeds(prompt)  # type: ignore[arg-type]
 
-        return self._process_tokens(prompt, skip_mm_cache=skip_mm_cache)  # type: ignore[arg-type]
+        engine_input = self._process_tokens(  # type: ignore[arg-type]
+            prompt,  # type: ignore[arg-type]
+            skip_mm_cache=skip_mm_cache,
+        )
+        self._truncate_expanded_prompt(engine_input, tok_params)
+        return engine_input
 
     async def _process_singleton_async(
         self,
         prompt: SingletonTokPrompt,
         *,
         skip_mm_cache: bool = False,
+        tok_params: TokenizeParams | None = None,
     ) -> SingletonInput:
         if "prompt_embeds" in prompt:
             return self._process_embeds(prompt)  # type: ignore[arg-type]
 
-        return await self._process_tokens_async(prompt, skip_mm_cache=skip_mm_cache)  # type: ignore[arg-type]
+        engine_input = await self._process_tokens_async(  # type: ignore[arg-type]
+            prompt,  # type: ignore[arg-type]
+            skip_mm_cache=skip_mm_cache,
+        )
+        self._truncate_expanded_prompt(engine_input, tok_params)
+        return engine_input
 
     def _get_skip_decoder_start_token(self) -> bool:
         """Whether the multimodal processor supplies a complete decoder prefix."""
@@ -1048,12 +1126,15 @@ class BaseRenderer(ABC, Generic[_T]):
         arrival_time: float,
         *,
         skip_mm_cache: bool = False,
+        tok_params: TokenizeParams | None = None,
     ) -> EngineInput:
         engine_input: EngineInput
         if "encoder_prompt" in prompt:
             engine_input = self._process_enc_dec(prompt, skip_mm_cache=skip_mm_cache)  # type: ignore[arg-type]
         else:
-            engine_input = self._process_singleton(prompt, skip_mm_cache=skip_mm_cache)
+            engine_input = self._process_singleton(
+                prompt, skip_mm_cache=skip_mm_cache, tok_params=tok_params
+            )
 
         engine_input["arrival_time"] = arrival_time
 
@@ -1065,6 +1146,7 @@ class BaseRenderer(ABC, Generic[_T]):
         arrival_time: float,
         *,
         skip_mm_cache: bool = False,
+        tok_params: TokenizeParams | None = None,
     ) -> EngineInput:
         engine_input: EngineInput
         if "encoder_prompt" in prompt:
@@ -1074,7 +1156,7 @@ class BaseRenderer(ABC, Generic[_T]):
             )
         else:
             engine_input = await self._process_singleton_async(
-                prompt, skip_mm_cache=skip_mm_cache
+                prompt, skip_mm_cache=skip_mm_cache, tok_params=tok_params
             )
 
         engine_input["arrival_time"] = arrival_time
@@ -1101,7 +1183,9 @@ class BaseRenderer(ABC, Generic[_T]):
         self._apply_prompt_extras(tok_prompts, prompt_extras)
 
         return [
-            self.process_for_engine(prompt, arrival_time, skip_mm_cache=skip_mm_cache)
+            self.process_for_engine(
+                prompt, arrival_time, skip_mm_cache=skip_mm_cache, tok_params=tok_params
+            )
             for prompt in tok_prompts
         ]
 
@@ -1126,7 +1210,7 @@ class BaseRenderer(ABC, Generic[_T]):
         return await asyncio.gather(
             *(
                 self.process_for_engine_async(
-                    p, arrival_time, skip_mm_cache=skip_mm_cache
+                    p, arrival_time, skip_mm_cache=skip_mm_cache, tok_params=tok_params
                 )
                 for p in tok_prompts
             )
@@ -1164,7 +1248,9 @@ class BaseRenderer(ABC, Generic[_T]):
         self._apply_prompt_extras(tok_prompts, prompt_extras)
 
         eng_prompts = [
-            self.process_for_engine(prompt, arrival_time, skip_mm_cache=skip_mm_cache)
+            self.process_for_engine(
+                prompt, arrival_time, skip_mm_cache=skip_mm_cache, tok_params=tok_params
+            )
             for prompt in tok_prompts
         ]
 
@@ -1204,7 +1290,7 @@ class BaseRenderer(ABC, Generic[_T]):
         eng_prompts = await asyncio.gather(
             *(
                 self.process_for_engine_async(
-                    p, arrival_time, skip_mm_cache=skip_mm_cache
+                    p, arrival_time, skip_mm_cache=skip_mm_cache, tok_params=tok_params
                 )
                 for p in tok_prompts
             )

vllm/renderers/hf.py

@@ -1194,22 +1194,20 @@ class HfRenderer(BaseRenderer[HfTokenizer]):
         coordinate space, no offset shifting needed afterwards.
         """
         prompt_embeds_info = cast(dict, prompt).pop("_prompt_embeds", None)
-        if prompt_embeds_info is not None:
-            tensors, placeholder_token_id = prompt_embeds_info
-            mm_updates = _build_prompt_embeds_updates(tensors, placeholder_token_id)
-            cast(dict, prompt)["prompt_token_ids"] = _expand_prompt_embeds_placeholders(
-                list(prompt["prompt_token_ids"]), mm_updates
-            )
+        if prompt_embeds_info is None:
+            return super()._process_tokens(prompt, skip_mm_cache=skip_mm_cache)
 
+        tensors, placeholder_token_id = prompt_embeds_info
+        mm_updates = _build_prompt_embeds_updates(tensors, placeholder_token_id)
+        cast(dict, prompt)["prompt_token_ids"] = _expand_prompt_embeds_placeholders(
+            list(prompt["prompt_token_ids"]), mm_updates
+        )
         engine_input = super()._process_tokens(prompt, skip_mm_cache=skip_mm_cache)
-
-        if prompt_embeds_info is not None:
-            tensors, _ = prompt_embeds_info
-            self._apply_prompt_embeds_to_engine_input(
-                cast(MultiModalInput, engine_input),
-                tensors,
-                mm_updates,
-            )
+        self._apply_prompt_embeds_to_engine_input(
+            cast(MultiModalInput, engine_input),
+            tensors,
+            mm_updates,
+        )
 
         return engine_input
 
@@ -1222,24 +1220,24 @@ class HfRenderer(BaseRenderer[HfTokenizer]):
     ) -> TokensInput | MultiModalInput:
         """Async equivalent of `_process_tokens`."""
         prompt_embeds_info = cast(dict, prompt).pop("_prompt_embeds", None)
-        if prompt_embeds_info is not None:
-            tensors, placeholder_token_id = prompt_embeds_info
-            mm_updates = _build_prompt_embeds_updates(tensors, placeholder_token_id)
-            cast(dict, prompt)["prompt_token_ids"] = _expand_prompt_embeds_placeholders(
-                list(prompt["prompt_token_ids"]), mm_updates
+        if prompt_embeds_info is None:
+            return await super()._process_tokens_async(
+                prompt, skip_mm_cache=skip_mm_cache
             )
 
+        tensors, placeholder_token_id = prompt_embeds_info
+        mm_updates = _build_prompt_embeds_updates(tensors, placeholder_token_id)
+        cast(dict, prompt)["prompt_token_ids"] = _expand_prompt_embeds_placeholders(
+            list(prompt["prompt_token_ids"]), mm_updates
+        )
         engine_input = await super()._process_tokens_async(
             prompt, skip_mm_cache=skip_mm_cache
         )
-
-        if prompt_embeds_info is not None:
-            tensors, _ = prompt_embeds_info
-            self._apply_prompt_embeds_to_engine_input(
-                cast(MultiModalInput, engine_input),
-                tensors,
-                mm_updates,
-            )
+        self._apply_prompt_embeds_to_engine_input(
+            cast(MultiModalInput, engine_input),
+            tensors,
+            mm_updates,
+        )
 
         return engine_input
 

vllm/renderers/online_derenderer.py

@@ -398,6 +398,7 @@ class OnlineDerenderer:
             model=model,
             choices=stream_choices,
             usage=usage,
+            metrics=generate_chunk.metrics,
         )
         return chunk, updated_state
 
@@ -595,6 +596,7 @@ class OnlineDerenderer:
             model=model,
             choices=stream_choices,
             usage=usage,
+            metrics=generate_chunk.metrics,
         )
         return chunk, updated_state
 
@@ -741,6 +743,7 @@ class OnlineDerenderer:
             model=model,
             choices=stream_choices,
             usage=usage,
+            metrics=generate_chunk.metrics,
         )
         return chunk, updated_state
 

vllm/v1/attention/backends/cpu_attn.py

@@ -512,11 +512,10 @@ def _get_attn_isa(
         )
     if supports_amx and dtype in (torch.bfloat16,) and block_size % 32 == 0:
         return "amx"
+    elif supports_arm:
+        return "neon"
     elif block_size % 32 == 0:
-        if supports_arm:
-            # support ARM NEON FMLA and BFMMLA (bf16) for block size 32
-            return "neon"
-        elif supports_riscv and _riscv_supports_rvv():
+        if supports_riscv and _riscv_supports_rvv():
             return "rvv"
         elif supports_vxe:
             return "vxe"

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

@@ -634,6 +634,38 @@ class DeepseekV32IndexerMetadata:
     prefill: DeepseekV32IndexerPrefillMetadata | None = None
 
 
+@triton.jit(do_not_specialize=["num_reqs", "num_actual_tokens", "num_tokens"])
+def _kpool_tail_slot_mapping_kernel(
+    slot_mapping_ptr,
+    block_table_ptr,
+    block_table_stride,
+    query_start_loc_ptr,
+    positions_ptr,
+    out_ptr,
+    num_reqs,
+    num_actual_tokens,
+    num_tokens,
+    kpool,
+    BLOCK: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    if pid < num_reqs:
+        start = tl.load(query_start_loc_ptr + pid)
+        # Tokens past the last request's boundary (if any) also map to it.
+        end = tl.load(query_start_loc_ptr + pid + 1)
+        end = tl.where(pid == num_reqs - 1, num_actual_tokens, end)
+        own_block = tl.load(block_table_ptr + pid * block_table_stride).to(tl.int64)
+        for i in range(start, end, BLOCK):
+            offs = i + tl.arange(0, BLOCK)
+            mask = offs < end
+            pos = tl.load(positions_ptr + offs, mask=mask, other=0).to(tl.int64)
+            tl.store(out_ptr + offs, own_block * kpool + pos % kpool, mask=mask)
+    else:
+        offs = num_actual_tokens + (pid - num_reqs) * BLOCK + tl.arange(0, BLOCK)
+        mask = offs < num_tokens
+        tl.store(out_ptr + offs, tl.load(slot_mapping_ptr + offs, mask=mask), mask=mask)
+
+
 def compute_kpool_tail_slot_mapping(
     slot_mapping: torch.Tensor,
     block_table: torch.Tensor,
@@ -646,10 +678,33 @@ def compute_kpool_tail_slot_mapping(
 ) -> torch.Tensor:
     """Map every token to its request's one circular tail block."""
     if out is None:
-        out = slot_mapping.clone()
+        out = torch.empty_like(slot_mapping)
     else:
         assert out.shape == slot_mapping.shape
-        out.copy_(slot_mapping)
+    if slot_mapping.is_cuda and slot_mapping.dim() == 1 and num_reqs > 0:
+        block = 256
+        num_tokens = slot_mapping.shape[0]
+        num_actual_tokens = min(num_actual_tokens, num_tokens)
+        grid = (num_reqs + triton.cdiv(num_tokens - num_actual_tokens, block),)
+        _kpool_tail_slot_mapping_kernel[grid](
+            slot_mapping,
+            block_table,
+            block_table.stride(0),
+            query_start_loc,
+            positions,
+            out,
+            num_reqs,
+            num_actual_tokens,
+            num_tokens,
+            kpool,
+            BLOCK=block,
+            num_warps=4,
+        )
+        return out
+    # Torch fallback: CPU tensors (the CPU unit tests), non-1D or empty inputs.
+    # Production always takes the Triton path above — spec-decode tokens arrive
+    # flattened token-major, so slot_mapping is 1D there too.
+    out.copy_(slot_mapping)
     if num_actual_tokens == 0:
         return out
     tokens = torch.arange(num_actual_tokens, device=slot_mapping.device)

vllm/v1/simple_kv_offload/manager.py

@@ -165,6 +165,9 @@ class SimpleCPUOffloadScheduler:
             kv_cache_config, offload_capacity
         )
         self.num_cpu_blocks = self.cpu_kv_cache_config.num_blocks
+        self.prefix_cacheable_group_ids = (
+            self.cpu_kv_cache_config.prefix_cacheable_group_ids
+        )
         self.kv_event_medium = MEDIUM_STORAGE if disk_capacity_bytes > 0 else MEDIUM_CPU
         # Find the full attention kv group for prefix cache matching.
         self.fa_gidx = -1
@@ -297,7 +300,7 @@ class SimpleCPUOffloadScheduler:
         """GPU blocks to keep available (free/offloaded) per step in lazy mode."""
         WATERMARK_RATIO = 1.0  # Reserve larger space to avoid running out of GPU blocks
         target = 0
-        for g in kv_cache_config.kv_cache_groups:
+        for g in kv_cache_config.prefix_cacheable_groups:
             spec = g.kv_cache_spec
             # Only full attention is sharded across DCP ranks; replicated specs
             # (mamba, sliding window, chunked-local) keep their own block size.
@@ -442,6 +445,9 @@ class SimpleCPUOffloadScheduler:
         # the rest will be released along with the temp pin below.
         cpu_hit_blocks: list[list[KVCacheBlock]] = []
         for g in range(num_groups):
+            if g not in self.prefix_cacheable_group_ids:
+                cpu_hit_blocks.append([])
+                continue
             g_block_size = self.group_block_sizes[g]
             n_take_g = cdiv(num_external_tokens, g_block_size)
             cpu_hit_blocks.append(cpu_hit_blocks_full[g][:n_take_g])
@@ -828,6 +834,8 @@ class SimpleCPUOffloadScheduler:
         num_free = self.cpu_block_pool.get_num_free_blocks()
 
         for g, group_gpu_ids in enumerate(block_ids_by_group):
+            if g not in self.prefix_cacheable_group_ids:
+                continue
             if len(gpu_block_ids) >= num_free:
                 break
             group_manager = self.cpu_coordinator.single_type_managers[g]

vllm/v1/worker/gpu/cudagraph_utils.py

@@ -662,6 +662,7 @@ class ModelCudaGraphManager(CudaGraphManager):
             model_inputs = {
                 "input_ids": input_buffers.input_ids[:num_tokens],
                 "positions": input_buffers.positions[:num_tokens],
+                "intermediate_tensors": None,
                 **model_state.prepare_dummy_inputs(num_reqs, num_tokens),
             }
             if not self.is_first_pp_rank:

vllm/v1/worker/gpu/model_runner.py

@@ -1291,7 +1291,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
         # Get the number of draft tokens for each request.
         draft_tokens = scheduler_output.scheduled_spec_decode_tokens
         num_draft_tokens_per_req = None
-        if not draft_tokens:
+        if not draft_tokens and self.model_state.num_new_sampled_tokens_per_step == 1:
             # No draft token scheduled (common case).
             total_num_draft_tokens = 0
             total_num_logits = num_reqs
@@ -1322,9 +1322,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
             np.cumsum(num_logits, out=cu_num_logits_np[1:])
             cu_num_logits = async_tensor_h2d(cu_num_logits_np, device=self.device)
 
-        adaptive_verification = (
-            self.adaptive_verification if num_draft_tokens_per_req is not None else None
-        )
+        # The general branch also serves a no-draft batch with k != 1, and
+        # compact_batch's budget is only primed when drafts are scheduled.
+        adaptive_verification = self.adaptive_verification if draft_tokens else None
         num_scheduled_tokens_upper_bound = num_scheduled_tokens_np
         if adaptive_verification is not None:
             # num_scheduled_tokens represents the draft budget evenly distributed across
@@ -1354,7 +1354,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
                 adaptive_verification.reallocate_drafts(req_ids, idx_mapping)
             )
             total_num_logits = num_reqs * num_bonus_tokens + total_num_draft_tokens
-        if draft_tokens:
+        if num_draft_tokens_per_req is not None:
             expanded_idx_mapping, expanded_local_pos = expand_idx_mapping(
                 idx_mapping, total_num_logits, cu_num_logits, self.decode_query_len
             )
@@ -1534,7 +1534,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
             sample_hidden_states = hidden_states[input_batch.logits_indices]
             logits = self.model.compute_logits(sample_hidden_states)
 
-        if grammar_output is not None:
+        # A diffusion prefill has no logit rows even when a bitmask row
+        # arrived for it.
+        if grammar_output is not None and logits.shape[0] > 0:
             # Apply grammar bitmask to the logits in-place.
             assert self.structured_outputs_worker is not None
             self.structured_outputs_worker.apply_grammar_bitmask(

vllm/v1/worker/workspace.py

@@ -3,11 +3,12 @@
 
 import inspect
 import os
-from collections.abc import Iterator
+from collections.abc import Callable, Hashable, Iterator
 from contextlib import contextmanager
 from contextvars import ContextVar
 from itertools import accumulate
 from math import prod
+from typing import TypeVar, cast
 
 import torch
 
@@ -17,6 +18,7 @@ from vllm.utils.math_utils import round_up
 from vllm.v1.worker.ubatching import dbo_current_ubatch_id
 
 logger = init_logger(__name__)
+T = TypeVar("T")
 
 
 def _compute_bytes(shape: tuple[int, ...], dtype: torch.dtype) -> int:
@@ -47,8 +49,8 @@ def use_workspace_lane(lane: int) -> Iterator[None]:
 class WorkspaceManager:
     """Manager for workspace allocation.
 
-    Manages one workspace buffer per active ``(ubatch, lane)`` slot.
-    Can be locked to prevent further growth during execution.
+    Manages shared scratch and keyed persistent resources per ``(ubatch, lane)``.
+    Can be locked to prevent growth or new resources during execution.
     """
 
     def __init__(
@@ -66,6 +68,9 @@ class WorkspaceManager:
         self._current_workspaces: list[torch.Tensor | None] = [None] * (
             self._num_ubatches * self._num_lanes
         )
+        self._persistent_resources: list[dict[Hashable, object]] = [
+            {} for _ in self._current_workspaces
+        ]
         self._locked: bool = False
 
     @staticmethod
@@ -76,10 +81,10 @@ class WorkspaceManager:
         return workspace.numel() * workspace.element_size()
 
     def lock(self) -> None:
-        """Lock the workspace to prevent further growth.
+        """Lock the workspace to prevent growth or new persistent resources.
 
-        After locking, any attempt to allocate a larger workspace will raise
-        an assertion error. This ensures workspace size is fixed during execution.
+        Larger scratch requests and new persistent resources raise an assertion
+        error. Existing resources retain their storage throughout execution.
         """
         self._locked = True
         if envs.VLLM_DEBUG_WORKSPACE:
@@ -113,6 +118,63 @@ class WorkspaceManager:
         """Check if workspace is locked."""
         return self._locked
 
+    def _get_workspace_id(self) -> int:
+        lane = _workspace_lane.get()
+        if lane >= self._num_lanes:
+            raise RuntimeError(
+                f"Workspace lane {lane} is not configured; manager has "
+                f"{self._num_lanes} lane(s)."
+            )
+        return dbo_current_ubatch_id() * self._num_lanes + lane
+
+    def get_persistent_resource(self, key: Hashable, factory: Callable[[], T]) -> T:
+        """Return the resource cached for ``key`` in the current ubatch and lane.
+
+        The first request calls ``factory`` and stores its result for the lifetime
+        of this manager. Later requests return the same object without calling
+        the factory again. Each ubatch/lane pair has a separate cache. Once the
+        manager is locked, requesting a key missing from that cache raises an
+        AssertionError.
+
+        Keys must identify the resource type and configuration. Callers must not
+        resize or replace tensor storage held by a cached resource. Uses of the
+        same resource must not overlap; include the CUDA stream in the key when
+        streams can execute concurrently.
+
+        With torch.compile, call this inside a runtime custom op so resource
+        selection happens during execution.
+        """
+        resources = self._persistent_resources[self._get_workspace_id()]
+        if key not in resources:
+            if self._locked:
+                raise AssertionError(
+                    f"Workspace is locked but {key!r} was not allocated during "
+                    "warmup for the current ubatch and lane."
+                )
+            resources[key] = factory()
+        return cast(T, resources[key])
+
+    def get_persistent(
+        self,
+        key: Hashable,
+        shape: tuple[int, ...],
+        dtype: torch.dtype,
+        *,
+        zero_init: bool = False,
+    ) -> torch.Tensor:
+        """Get a fixed-shape tensor, optionally zeroed on first allocation only."""
+        factory = torch.zeros if zero_init else torch.empty
+        workspace = self.get_persistent_resource(
+            (torch.Tensor, key),
+            lambda: factory(shape, dtype=dtype, device=self._device),
+        )
+        if workspace.shape != shape or workspace.dtype != dtype:
+            raise ValueError(
+                f"Persistent workspace {key!r} has shape={workspace.shape}, "
+                f"dtype={workspace.dtype}; requested shape={shape}, dtype={dtype}."
+            )
+        return workspace
+
     def get_simultaneous(
         self, *shapes_and_dtypes: tuple[tuple[int, ...], torch.dtype]
     ) -> list[torch.Tensor]:
@@ -151,14 +213,7 @@ class WorkspaceManager:
             The current workspace tensor.
 
         """
-        ubatch_id = dbo_current_ubatch_id()
-        lane = _workspace_lane.get()
-        if lane >= self._num_lanes:
-            raise RuntimeError(
-                f"Workspace lane {lane} is not configured; manager has "
-                f"{self._num_lanes} lane(s)."
-            )
-        workspace_id = ubatch_id * self._num_lanes + lane
+        workspace_id = self._get_workspace_id()
         current_workspace = self._current_workspaces[workspace_id]
         current_size = self._workspace_size_bytes(current_workspace)
 
@@ -218,8 +273,8 @@ class WorkspaceManager:
                     get_caller_info(),
                     current_size / _MB,
                     required_bytes / _MB,
-                    ubatch_id,
-                    lane,
+                    workspace_id // self._num_lanes,
+                    workspace_id % self._num_lanes,
                 )
 
         return current_workspace
@@ -277,12 +332,11 @@ def init_workspace_manager(
 
 
 def lock_workspace() -> None:
-    """Lock the workspace to prevent further growth.
+    """Lock the workspace to prevent growth or new persistent resources.
 
-    After calling this function, any attempt to allocate a workspace larger
-    than the current size will raise an AssertionError. This ensures that
-    workspace size is fixed during execution and prevents unexpected memory
-    allocations in the hot path.
+    Larger scratch requests and new persistent resources raise an AssertionError.
+    All required resources must be initialized in each execution slot during
+    warmup, so their storage remains fixed during execution.
 
     Example:
         # During initialization