fix(dspark): reserve V2 padded Q scratch

csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu

@@ -942,19 +942,21 @@ static void launchFullCacheKernel(
 // ────────────────────────────────────────────────────────────────────────────
 // Torch op wrapper
 // ────────────────────────────────────────────────────────────────────────────
-torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
+void fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
     torch::stable::Tensor const& q_in,           // [N, num_heads_q, 512] bf16
     torch::stable::Tensor const& kv,             // [N, 512] bf16 (read-only)
+    torch::stable::Tensor& q_out,                // [N, q_head_padded, 512] bf16
     torch::stable::Tensor& k_cache,              // [num_blocks, block_bytes] uint8
     torch::stable::Tensor const& slot_mapping,   // [N] int64
     torch::stable::Tensor const& position_ids,   // [N] int64
     torch::stable::Tensor const& cos_sin_cache,  // [max_pos, rope_dim] bf16
-    int64_t q_head_padded,                       // padded Q head count for output
     double eps, int64_t cache_block_size) {
   STD_TORCH_CHECK(q_in.device().is_cuda() && q_in.is_contiguous(),
                   "q_in must be contiguous CUDA");
   STD_TORCH_CHECK(kv.device().is_cuda() && kv.is_contiguous(),
                   "kv must be contiguous CUDA");
+  STD_TORCH_CHECK(q_out.device().is_cuda() && q_out.is_contiguous(),
+                  "q_out must be contiguous CUDA");
   STD_TORCH_CHECK(k_cache.device().is_cuda(), "k_cache must be CUDA");
   STD_TORCH_CHECK(slot_mapping.device().is_cuda() &&
                       slot_mapping.scalar_type() ==
@@ -970,8 +972,13 @@ torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
   STD_TORCH_CHECK(kv.dim() == 2 && kv.size(1) == 512, "kv shape [N, 512]");
   STD_TORCH_CHECK(q_in.scalar_type() == kv.scalar_type(),
                   "q_in and kv dtype must match");
-  STD_TORCH_CHECK(q_head_padded >= q_in.size(1),
-                  "q_head_padded must be >= q_in.size(1) (num_heads_q)");
+  STD_TORCH_CHECK(q_out.scalar_type() == q_in.scalar_type(),
+                  "q_out dtype must match q_in");
+  STD_TORCH_CHECK(q_out.dim() == 3 && q_out.size(0) == q_in.size(0) &&
+                      q_out.size(1) >= q_in.size(1) &&
+                      q_out.size(2) == q_in.size(2),
+                  "q_out shape [N, q_head_padded, 512] with "
+                  "q_head_padded >= num_heads_q");
   STD_TORCH_CHECK(k_cache.scalar_type() == torch::headeronly::ScalarType::Byte,
                   "k_cache must be uint8");
   STD_TORCH_CHECK(cos_sin_cache.dim() == 2 && cos_sin_cache.size(1) == 64,
@@ -991,19 +998,14 @@ torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
   STD_TORCH_CHECK(num_tokens_insert <= num_tokens_full,
                   "slot_mapping must not exceed q row count");
   int const num_heads_q = static_cast<int>(q_in.size(1));
-  int const num_heads_q_padded = static_cast<int>(q_head_padded);
+  int const num_heads_q_padded = static_cast<int>(q_out.size(1));
   int const cache_block_size_i = static_cast<int>(cache_block_size);
   int const kv_block_stride = static_cast<int>(k_cache.stride(0));
 
   const torch::stable::accelerator::DeviceGuard device_guard(
       q_in.get_device_index());
   const cudaStream_t stream = get_current_cuda_stream(q_in.get_device_index());
 
-  // Allocate the padded q output.  The kernel writes every element (live
-  // region gets RMSNorm+RoPE; pad region gets zeros), so `empty` is safe.
-  auto q_out = torch::stable::new_empty(
-      q_in, {q_in.size(0), q_head_padded, q_in.size(2)}, q_in.scalar_type());
-
   VLLM_STABLE_DISPATCH_HALF_TYPES(
       q_in.scalar_type(), "fused_deepseek_v4_qnorm_rope_kv_insert", [&] {
         using qkv_scalar_t = scalar_t;
@@ -1020,7 +1022,6 @@ torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
                 num_heads_q_padded, cache_block_size_i, kv_block_stride,
                 stream);
       });
-  return q_out;
 }
 
 // ────────────────────────────────────────────────────────────────────────────

csrc/libtorch_stable/ops.h

@@ -263,12 +263,13 @@ void fused_qk_norm_rope(torch::stable::Tensor& qkv, int64_t num_heads_q,
                         torch::stable::Tensor& position_ids,
                         int64_t forced_token_heads_per_warp);
 
-torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
+void fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
     torch::stable::Tensor const& q_in, torch::stable::Tensor const& kv,
-    torch::stable::Tensor& k_cache, torch::stable::Tensor const& slot_mapping,
+    torch::stable::Tensor& q_out, torch::stable::Tensor& k_cache,
+    torch::stable::Tensor const& slot_mapping,
     torch::stable::Tensor const& position_ids,
-    torch::stable::Tensor const& cos_sin_cache, int64_t q_head_padded,
-    double eps, int64_t cache_block_size);
+    torch::stable::Tensor const& cos_sin_cache, double eps,
+    int64_t cache_block_size);
 
 void fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert(
     torch::stable::Tensor& q, torch::stable::Tensor const& kv,

csrc/libtorch_stable/torch_bindings.cpp

@@ -430,9 +430,9 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
 
   ops.def(
       "fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert("
-      "Tensor q_in, Tensor kv, Tensor! k_cache, "
+      "Tensor q_in, Tensor kv, Tensor! q_out, Tensor! k_cache, "
       "Tensor slot_mapping, Tensor position_ids, Tensor cos_sin_cache, "
-      "int q_head_padded, float eps, int cache_block_size) -> Tensor");
+      "float eps, int cache_block_size) -> ()");
 
   // FlashInfer V4 full-cache variants: write Q in place (bf16) or to a separate
   // FP8 tensor, and KV into a contiguous 512-wide token-strided cache.

tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py

@@ -13,7 +13,8 @@
     `dequantize_and_gather_k_cache` for KV
 
 The kernel is imported via
-`torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert`.
+`torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert` and writes Q
+into a caller-owned output tensor.
 """
 
 import pytest
@@ -148,17 +149,50 @@ def _full_cache_bf16_op_available() -> bool:
 def _call_fused(
     q_in, q_head_padded, kv, k_cache, slot_mapping, positions, cos_sin_cache, eps, bs
 ):
-    return torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
+    q_out = torch.empty(
+        q_in.shape[0],
+        q_head_padded,
+        q_in.shape[2],
+        dtype=q_in.dtype,
+        device=q_in.device,
+    )
+    torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
         q_in,
         kv,
+        q_out,
         k_cache,
         slot_mapping,
         positions,
         cos_sin_cache,
-        q_head_padded,
         eps,
         bs,
     )
+    return q_out
+
+
+def _call_fused_out(
+    q_in,
+    q_out,
+    kv,
+    k_cache,
+    slot_mapping,
+    positions,
+    cos_sin_cache,
+    eps,
+    bs,
+):
+    torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
+        q_in,
+        kv,
+        q_out,
+        k_cache,
+        slot_mapping,
+        positions,
+        cos_sin_cache,
+        eps,
+        bs,
+    )
+    return q_out
 
 
 def _as_stored_fp8(t: torch.Tensor) -> torch.Tensor:
@@ -269,6 +303,43 @@ def test_q_path_matches_reference(num_tokens: int, n_heads: int, padded_heads: i
         )
 
 
+def test_quant_insert_writes_caller_owned_q_out():
+    torch.manual_seed(4)
+    device = "cuda"
+    dtype = torch.bfloat16
+    eps = 1e-6
+    num_tokens = 17
+    n_heads = 8
+    padded_heads = 16
+    block_size = 16
+
+    q = torch.randn(num_tokens, n_heads, HEAD_DIM, dtype=dtype, device=device)
+    kv = torch.randn(num_tokens, HEAD_DIM, dtype=dtype, device=device)
+    positions = torch.arange(num_tokens, dtype=torch.int64, device=device)
+    cos_sin_cache = make_cos_sin_cache(4096, ROPE_DIM, torch.float32, device)
+    k_cache = torch.zeros(
+        2, block_size * HEAD_BYTES, dtype=torch.uint8, device=device
+    )
+    slot_mapping = torch.full((num_tokens,), -1, dtype=torch.int64, device=device)
+    q_out = torch.full(
+        (num_tokens, padded_heads, HEAD_DIM),
+        -123.0,
+        dtype=dtype,
+        device=device,
+    )
+
+    returned = _call_fused_out(
+        q, q_out, kv, k_cache, slot_mapping, positions, cos_sin_cache, eps, block_size
+    )
+
+    assert returned.data_ptr() == q_out.data_ptr()
+    q_ref = apply_rope_gptj_last_k(
+        rmsnorm_no_weight(q, eps), positions, cos_sin_cache
+    ).to(dtype)
+    torch.testing.assert_close(q_out[:, :n_heads], q_ref, rtol=1e-2, atol=1e-2)
+    assert q_out[:, n_heads:padded_heads].abs().max().item() == 0.0
+
+
 # ── Test 2: KV path round-trip byte/value parity ─────────────────────────────
 
 

vllm/compilation/passes/utility/fix_functionalization.py

@@ -191,7 +191,17 @@ def __call__(self, graph: torch.fx.Graph) -> None:
                 }
                 self.defunctionalize(graph, node, mutated_args=mutated_args)
             elif at_target in fused_deepseek_v4_mla_targets:
-                mutated_args = {1: "q", 2: "k_cache"}
+                if (
+                    hasattr(
+                        torch.ops._C,
+                        "fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert",
+                    )
+                    and at_target
+                    == torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert.default  # noqa: E501
+                ):
+                    mutated_args = {1: "q_out", 2: "k_cache"}
+                else:
+                    mutated_args = {1: "q", 2: "k_cache"}
                 self.defunctionalize(graph, node, mutated_args)
             elif (
                 hasattr(torch.ops.vllm, "fused_rope_unified_mla_kv_cache_update")

vllm/models/deepseek_v4/attention.py

@@ -57,6 +57,7 @@
 )
 from vllm.v1.attention.backends.mla.sparse_swa import DeepseekV4SWACache
 from vllm.v1.kv_cache_interface import KVCacheSpec, MLAAttentionSpec
+from vllm.v1.worker.ubatching import dbo_current_ubatch_id
 
 logger = init_logger(__name__)
 
@@ -117,6 +118,9 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
     # workspace allocated in _forward_prefill and is also read by the dummy-run
     # path to pre-reserve that workspace.
     PREFILL_CHUNK_SIZE: ClassVar[int] = 4
+    _q_padded_scratch_by_key: ClassVar[
+        dict[tuple[str, int, int, torch.dtype, int, int], torch.Tensor]
+    ] = {}
 
     @classmethod
     @abstractmethod
@@ -280,6 +284,10 @@ def __init__(
             vllm_config.scheduler_config.max_num_batched_tokens
         )
         self.max_model_len = vllm_config.model_config.max_model_len
+        self._q_padded_scratch_num_ubatches = (
+            2 if vllm_config.parallel_config.enable_dbo else 1
+        )
+        self._q_padded_scratch_dtype = vllm_config.model_config.dtype
 
         # Resolve the kv-cache dtype from this backend's block format. The same
         # resolution drives the SWA cache tensor dtype below.
@@ -317,6 +325,78 @@ def __init__(
                 k_cache_prefix=self.prefix,
             )
 
+    @staticmethod
+    def _q_padded_scratch_device_index(device: torch.device) -> int:
+        if device.index is not None:
+            return int(device.index)
+        if device.type == "cuda":
+            return int(torch.cuda.current_device())
+        return -1
+
+    @classmethod
+    def _reserve_q_padded_scratch_buffer(
+        cls,
+        num_tokens: int,
+        padded_heads: int,
+        head_dim: int,
+        dtype: torch.dtype,
+        device: torch.device,
+        ubatch_id: int,
+    ) -> torch.Tensor:
+        key = (
+            device.type,
+            cls._q_padded_scratch_device_index(device),
+            ubatch_id,
+            dtype,
+            padded_heads,
+            head_dim,
+        )
+        scratch = cls._q_padded_scratch_by_key.get(key)
+        if scratch is not None and scratch.shape[0] >= num_tokens:
+            return scratch
+
+        old_scratch = cls._q_padded_scratch_by_key.pop(key, None)
+        if old_scratch is not None:
+            del old_scratch
+            torch.accelerator.empty_cache()
+
+        scratch = torch.empty(
+            (num_tokens, padded_heads, head_dim),
+            dtype=dtype,
+            device=device,
+        )
+        cls._q_padded_scratch_by_key[key] = scratch
+        return scratch
+
+    def _get_q_padded_scratch(self, q: torch.Tensor) -> torch.Tensor:
+        num_tokens = q.shape[0]
+        reserved_tokens = max(num_tokens, int(self.max_num_batched_tokens))
+        scratch = self._reserve_q_padded_scratch_buffer(
+            reserved_tokens,
+            int(self.padded_heads),
+            int(self.head_dim),
+            q.dtype,
+            q.device,
+            dbo_current_ubatch_id(),
+        )
+        return scratch[:num_tokens]
+
+    def reserve_profile_scratch(self) -> None:
+        if self.kv_cache_torch_dtype != torch.uint8:
+            return
+        device = self.q_norm.weight.device
+        if device.type not in ("cuda", "xpu"):
+            return
+        for ubatch_id in range(self._q_padded_scratch_num_ubatches):
+            self._reserve_q_padded_scratch_buffer(
+                max(1, int(self.max_num_batched_tokens)),
+                int(self.padded_heads),
+                int(self.head_dim),
+                self._q_padded_scratch_dtype,
+                device,
+                ubatch_id,
+            )
+
     def forward(
         self,
         positions: torch.Tensor,
@@ -518,6 +598,8 @@ def _fused_qnorm_rope_kv_insert(
         if not isinstance(attn_metadata, dict):
             # Profile run: kernel doesn't fire; produce a padded tensor so
             # downstream FlashMLA gets the right shape.
+            if self.kv_cache_torch_dtype == torch.uint8:
+                return self._get_q_padded_scratch(q)
             if self.n_local_heads < self.padded_heads:
                 return F.pad(
                     q,
@@ -543,21 +625,22 @@ def _fused_qnorm_rope_kv_insert(
         if cache_dtype == torch.uint8:
             # fp8_ds_mla UE8M0 paged path. Horizontally fused:
             #   Q side:  per-head RMSNorm (no weight) + GPT-J RoPE, zero-filling
-            #            the padding head slots; the kernel allocates and returns
-            #            the padded q tensor.
+            #            the padding head slots into a reusable q buffer.
             #   KV side: GPT-J RoPE + UE8M0 FP8 quant + paged cache insert.
             swa_kv_cache_2d = swa_kv_cache.view(swa_kv_cache.shape[0], -1)
-            return torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
+            q_out = self._get_q_padded_scratch(q)
+            torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
                 q,
                 kv,
+                q_out,
                 swa_kv_cache_2d,
                 swa_metadata.slot_mapping,
                 positions,
                 cos_sin_cache,
-                self.padded_heads,
                 self.eps,
                 swa_metadata.block_size,
             )
+            return q_out
 
         # Plain-row path: the [num_blocks, block_size, 512] cache stores the KV
         # row in its element dtype (no Q padding). bf16 rewrites q in place;

vllm/v1/worker/gpu/model_runner.py

@@ -643,8 +643,19 @@ def _dummy_pooler_run(self, hidden_states: torch.Tensor) -> None:
         assert self.pooling_runner is not None
         self.pooling_runner.dummy_pooler_run(hidden_states)
 
+    def _reserve_profile_scratch(self) -> None:
+        seen: set[int] = set()
+        for module in self.compilation_config.static_forward_context.values():
+            if id(module) in seen:
+                continue
+            seen.add(id(module))
+            reserve = getattr(module, "reserve_profile_scratch", None)
+            if reserve is not None:
+                reserve()
+
     @torch.inference_mode()
     def profile_run(self) -> None:
+        self._reserve_profile_scratch()
         hidden_states, sample_hidden_states = self._dummy_run(
             self.max_num_tokens, skip_attn=True, is_profile=True
         )