From 53b6d1110cef80de7c42dbd5d30dbedf8313190d Mon Sep 17 00:00:00 2001 From: alexbi29 Date: Thu, 2 Jul 2026 09:54:44 +0000 Subject: [PATCH] fix(dspark): reserve V2 padded Q scratch --- ...deepseek_v4_qnorm_rope_kv_insert_kernel.cu | 23 ++--- csrc/libtorch_stable/ops.h | 9 +- csrc/libtorch_stable/torch_bindings.cpp | 4 +- ..._fused_deepseek_v4_qnorm_rope_kv_insert.py | 77 +++++++++++++++- .../passes/utility/fix_functionalization.py | 12 ++- vllm/models/deepseek_v4/attention.py | 91 ++++++++++++++++++- vllm/v1/worker/gpu/model_runner.py | 11 +++ 7 files changed, 202 insertions(+), 25 deletions(-) diff --git a/csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu b/csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu index 7bc435b8e0da..a59cd1668846 100644 --- a/csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu +++ b/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,7 +998,7 @@ 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(q_in.size(1)); - int const num_heads_q_padded = static_cast(q_head_padded); + int const num_heads_q_padded = static_cast(q_out.size(1)); int const cache_block_size_i = static_cast(cache_block_size); int const kv_block_stride = static_cast(k_cache.stride(0)); @@ -999,11 +1006,6 @@ torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( 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; } // ──────────────────────────────────────────────────────────────────────────── diff --git a/csrc/libtorch_stable/ops.h b/csrc/libtorch_stable/ops.h index 7cf34d8b03ab..88f50cd0e139 100644 --- a/csrc/libtorch_stable/ops.h +++ b/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, diff --git a/csrc/libtorch_stable/torch_bindings.cpp b/csrc/libtorch_stable/torch_bindings.cpp index 158999a6633b..61510a057076 100644 --- a/csrc/libtorch_stable/torch_bindings.cpp +++ b/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. diff --git a/tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py b/tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py index ed163a0472ad..bf2e9ceaee31 100644 --- a/tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py +++ b/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 ───────────────────────────── diff --git a/vllm/compilation/passes/utility/fix_functionalization.py b/vllm/compilation/passes/utility/fix_functionalization.py index 75994ba58039..6988b30d9fac 100644 --- a/vllm/compilation/passes/utility/fix_functionalization.py +++ b/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") diff --git a/vllm/models/deepseek_v4/attention.py b/vllm/models/deepseek_v4/attention.py index 519f5f9a1449..b94177298d97 100644 --- a/vllm/models/deepseek_v4/attention.py +++ b/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; diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 5b067c87e0d4..679e5ab0711b 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/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 )