From 89876b0c548afdd932d41d5ef81b14347edb9ff7 Mon Sep 17 00:00:00 2001 From: Woosuk Kwon Date: Sun, 28 Jun 2026 08:17:39 -0700 Subject: [PATCH] [GLM5] Implement op fusion for GLM5/DSV3.2 (#46876) --- .../test_fused_deepseek_v32_norm_rope.py | 423 +++++++++ .../model_executor/layers/fused_moe/config.py | 6 + .../layers/fused_moe/runner/moe_runner.py | 2 + .../layers/sparse_attn_indexer.py | 10 +- vllm/models/deepseek_v32/nvidia/attention.py | 251 ++++-- vllm/models/deepseek_v32/nvidia/fused_ops.py | 63 ++ vllm/models/deepseek_v32/nvidia/kernels.py | 823 ++++++++++++++++++ vllm/models/deepseek_v32/nvidia/model.py | 24 +- vllm/models/deepseek_v32/nvidia/mtp.py | 21 +- 9 files changed, 1531 insertions(+), 92 deletions(-) create mode 100644 tests/kernels/test_fused_deepseek_v32_norm_rope.py create mode 100644 vllm/models/deepseek_v32/nvidia/fused_ops.py create mode 100644 vllm/models/deepseek_v32/nvidia/kernels.py diff --git a/tests/kernels/test_fused_deepseek_v32_norm_rope.py b/tests/kernels/test_fused_deepseek_v32_norm_rope.py new file mode 100644 index 0000000000..a6f6d71b48 --- /dev/null +++ b/tests/kernels/test_fused_deepseek_v32_norm_rope.py @@ -0,0 +1,423 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for the horizontally-fused deepseek_v32 (NVIDIA SM100) Triton +kernels used by the specialized DSA model: + + fused_norm_rope + - q : q_lora RMSNorm + - kv : kv_lora RMSNorm + (interleaved) RoPE on k_pe + MLA cache insert + (bf16 or per-tensor fp8) + - idx: indexer-K LayerNorm + RoPE (interleaved or NeoX) + UE8M0 fp8 quant + + packed indexer cache insert; plus the top-k buffer (-1) fill + fused_q + - mqa: ql_nope + (interleaved) RoPE'd q_pe, concat-quantized to the fp8 MQA + query + - idx: indexer-Q RoPE (interleaved or NeoX) + UE8M0 fp8 quant + folded + index weights + fused_eh_norm (MTP): zero-at-pos-0 + enorm RMSNorm(embeds) + hnorm + RMSNorm(prev), concatenated side-by-side + +Each kernel is compared against a PyTorch reference. The kernel keeps the whole +pipeline in fp32 and rounds once, so it can land on the opposite side of a +round-to-nearest tie from the reference for a few elements: deterministic fp8 +outputs are checked within 1 representable-step (ULP); bf16 norm/RoPE outputs use +rtol/atol=1e-2 (the tolerance the sibling deepseek_v4 fused-kernel test uses). +""" + +import pytest +import torch + +from vllm.models.deepseek_v32.nvidia import kernels as K +from vllm.platforms import current_platform + +FP8 = torch.float8_e4m3fn +FP8_MAX = 448.0 + +# GLM-5.2 / DeepSeek-V3.2 shapes (TP8 local heads). +Q_LORA = 2048 +KV_LORA = 512 +ROPE_DIM = 64 +NUM_HEADS = 8 +INDEX_HEADS = 32 +INDEX_HEAD_DIM = 128 +HIDDEN = 6144 +EPS = 1e-6 + +pytestmark = pytest.mark.skipif( + not current_platform.is_cuda() or not current_platform.has_device_capability(89), + reason="deepseek_v32 fused kernels require CUDA with fp8 (SM89+)", +) + + +# ── reference helpers ──────────────────────────────────────────────────────── + + +def make_cos_sin(max_pos: int, rot_dim: int, device) -> torch.Tensor: + """cos||sin cache: row[pos] = [cos(theta)(rot/2), sin(theta)(rot/2)].""" + half = rot_dim // 2 + inv_freq = 1.0 / ( + 10000.0 ** (torch.arange(0, half, dtype=torch.float32, device=device) / half) + ) + t = torch.arange(max_pos, dtype=torch.float32, device=device) + freqs = torch.einsum("i,j->ij", t, inv_freq) + return torch.cat([freqs.cos(), freqs.sin()], dim=-1) + + +def rms_norm(x: torch.Tensor, w: torch.Tensor) -> torch.Tensor: + """RMSNorm matching kernels._rms_norm (fp32, eps inside rsqrt). Returns fp32.""" + xf = x.float() + ms = xf.pow(2).mean(dim=-1, keepdim=True) + return xf * torch.rsqrt(ms + EPS) * w.float() + + +def layer_norm(x: torch.Tensor, w: torch.Tensor, b: torch.Tensor) -> torch.Tensor: + xf = x.float() + mean = xf.mean(dim=-1, keepdim=True) + var = (xf - mean).pow(2).mean(dim=-1, keepdim=True) + return (xf - mean) * torch.rsqrt(var + EPS) * w.float() + b.float() + + +def rope( + x: torch.Tensor, pos: torch.Tensor, cos_sin: torch.Tensor, interleave: bool +) -> torch.Tensor: + """Apply RoPE to the first ``rot_dim`` elements of x's last dim. + + x: [..., head_dim] fp32. ``cos_sin`` is [max_pos, rot_dim]. ``interleave`` + selects adjacent-pair (GLM) vs split-half NeoX (DeepSeek-V3.2) layout. + """ + rot = cos_sin.shape[-1] + half = rot // 2 + cs = cos_sin[pos.long()] + cos, sin = cs[..., :half], cs[..., half:] + out = x.float().clone() + r = out[..., :rot] + if interleave: + x1, x2 = r[..., 0::2].clone(), r[..., 1::2].clone() + r[..., 0::2] = x1 * cos - x2 * sin + r[..., 1::2] = x2 * cos + x1 * sin + else: + x1, x2 = r[..., :half].clone(), r[..., half:].clone() + r[..., :half] = x1 * cos - x2 * sin + r[..., half:] = x2 * cos + x1 * sin + return out + + +def ue8m0_quant(vals: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Per-row (last dim) UE8M0 fp8 quant matching kernels._fp8_ue8m0_quantize.""" + amax = vals.float().abs().amax(dim=-1, keepdim=True) + scale = torch.clamp(amax, min=1e-4) / FP8_MAX + scale = torch.exp2(torch.ceil(torch.log2(scale))) + q = (vals.float() / scale).to(FP8) + return q, scale.squeeze(-1) + + +def _bf16_ulp(a: torch.Tensor, b: torch.Tensor) -> int: + def key(t): + u = t.contiguous().view(torch.int16).to(torch.int64) & 0xFFFF + return torch.where(u >= 0x8000, 0xFFFF - u, u + 0x8000) + + return int((key(a) - key(b)).abs().max().item()) + + +def _fp8_ulp(a: torch.Tensor, b: torch.Tensor) -> int: + def key(t): + u = t.contiguous().view(torch.uint8).to(torch.int64) + return torch.where(u >= 0x80, 0xFF - u, u + 0x80) + + return int((key(a) - key(b)).abs().max().item()) + + +def assert_bf16(got: torch.Tensor, ref_fp32: torch.Tensor, msg: str): + # Kernel keeps RMSNorm/RoPE in fp32 and rounds to bf16 once; the fp32 + # reduction/FMA order differs from torch, so a few elements land on the + # opposite side of a round-to-nearest tie. Use the same tolerance the + # sibling deepseek_v4 fused-kernel test uses for this bf16 norm+rope class. + torch.testing.assert_close( + got.float(), ref_fp32.float(), rtol=1e-2, atol=1e-2, msg=lambda m: f"{msg}: {m}" + ) + + +def assert_fp8(got: torch.Tensor, ref: torch.Tensor, msg: str): + assert _fp8_ulp(got, ref) <= 1, f"{msg}: >1 fp8 ULP" + + +# ── fused_norm_rope ────────────────────────────────────────────────────────── + + +@pytest.mark.parametrize("num_tokens", [1, 4, 17, 512, 4096]) +@pytest.mark.parametrize("index_interleave", [True, False]) +@pytest.mark.parametrize("mla_fp8", [False, True]) +def test_fused_norm_rope(num_tokens: int, index_interleave: bool, mla_fp8: bool): + torch.manual_seed(0) + dev = "cuda" + max_pos = 8192 + pos = torch.arange(num_tokens, device=dev, dtype=torch.int64) % max_pos + + q_c = torch.randn(num_tokens, Q_LORA, device=dev, dtype=torch.bfloat16) + kv_c = torch.randn(num_tokens, KV_LORA, device=dev, dtype=torch.bfloat16) + k_pe = torch.randn(num_tokens, ROPE_DIM, device=dev, dtype=torch.bfloat16) + qw = torch.randn(Q_LORA, device=dev, dtype=torch.bfloat16) + kvw = torch.randn(KV_LORA, device=dev, dtype=torch.bfloat16) + ik = torch.randn(num_tokens, INDEX_HEAD_DIM, device=dev, dtype=torch.bfloat16) + ikw = torch.randn(INDEX_HEAD_DIM, device=dev, dtype=torch.float32) + ikb = torch.randn(INDEX_HEAD_DIM, device=dev, dtype=torch.float32) + + mla_cos_sin = make_cos_sin(max_pos, ROPE_DIM, dev) # MLA k_pe: interleaved + idx_cos_sin = make_cos_sin(max_pos, ROPE_DIM, dev) + + bs = max_pos # single block covering all tokens + mla_dim = KV_LORA + ROPE_DIM + if mla_fp8: + mla_cache = torch.zeros(1, bs, mla_dim, device=dev, dtype=torch.uint8) + mla_dtype = "fp8" + mla_k_scale = torch.tensor([0.3], device=dev, dtype=torch.float32) + else: + mla_cache = torch.zeros(1, bs, mla_dim, device=dev, dtype=torch.bfloat16) + mla_dtype = "auto" + mla_k_scale = None + idx_row = INDEX_HEAD_DIM + INDEX_HEAD_DIM // 128 * 4 # 132 + idx_cache = torch.zeros(1, bs, idx_row, device=dev, dtype=torch.uint8) + slot = torch.arange(num_tokens, device=dev, dtype=torch.int64) + topk = torch.full((num_tokens, 2048), 7, device=dev, dtype=torch.int32) + + q_out = K.fused_norm_rope( + pos, + q_c, + qw, + EPS, + kv_c, + kvw, + EPS, + k_pe, + mla_cos_sin, + ik, + ikw, + ikb, + EPS, + idx_cos_sin, + topk, + slot_mapping=slot, + indexer_k_cache=idx_cache, + mla_kv_cache=mla_cache, + mla_kv_cache_dtype=mla_dtype, + mla_k_scale=mla_k_scale, + has_indexer=True, + index_rope_interleave=index_interleave, + ) + + # q_lora RMSNorm + assert_bf16(q_out, rms_norm(q_c, qw), "q_c rmsnorm") + + # MLA cache: [kv_c_normed | k_pe_roped(interleaved)] + kv_ref = rms_norm(kv_c, kvw) + kpe_ref = rope(k_pe.float(), pos, mla_cos_sin, interleave=True) + if mla_fp8: + cache = mla_cache.view(FP8)[0, :num_tokens] + s = mla_k_scale.item() + assert_fp8(cache[:, :KV_LORA], (kv_ref / s).to(FP8), "MLA kv fp8") + assert_fp8(cache[:, KV_LORA:], (kpe_ref / s).to(FP8), "MLA k_pe fp8") + else: + cache = mla_cache[0, :num_tokens] + assert_bf16(cache[:, :KV_LORA], kv_ref, "MLA kv bf16") + assert_bf16(cache[:, KV_LORA:], kpe_ref, "MLA k_pe bf16") + + # Indexer-K cache (packed [bs*head_dim fp8 | bs*4 fp32 scale]). + ik_ref = layer_norm(ik, ikw, ikb) + ik_ref = rope(ik_ref, pos, idx_cos_sin, interleave=index_interleave) + q_ref, s_ref = ue8m0_quant(ik_ref) + flat = idx_cache[0].reshape(-1) + vals = flat[: bs * INDEX_HEAD_DIM].view(FP8).reshape(bs, INDEX_HEAD_DIM) + scales = flat[bs * INDEX_HEAD_DIM :].view(torch.float32) + assert_fp8(vals[:num_tokens], q_ref, "indexer-K fp8") + torch.testing.assert_close(scales[:num_tokens], s_ref, rtol=0, atol=0) + + # Top-k buffer cleared to -1 on indexer layers. + assert (topk == -1).all(), "topk buffer not cleared on indexer layer" + + +@pytest.mark.parametrize("num_tokens", [1, 17, 512]) +def test_fused_norm_rope_no_indexer(num_tokens: int): + """Shared (no-indexer) layer: q + kv/MLA only; top-k buffer untouched.""" + torch.manual_seed(1) + dev = "cuda" + max_pos = 8192 + pos = torch.arange(num_tokens, device=dev, dtype=torch.int64) + + q_c = torch.randn(num_tokens, Q_LORA, device=dev, dtype=torch.bfloat16) + kv_c = torch.randn(num_tokens, KV_LORA, device=dev, dtype=torch.bfloat16) + k_pe = torch.randn(num_tokens, ROPE_DIM, device=dev, dtype=torch.bfloat16) + qw = torch.randn(Q_LORA, device=dev, dtype=torch.bfloat16) + kvw = torch.randn(KV_LORA, device=dev, dtype=torch.bfloat16) + mla_cos_sin = make_cos_sin(max_pos, ROPE_DIM, dev) + + bs = max_pos + mla_cache = torch.zeros(1, bs, KV_LORA + ROPE_DIM, device=dev, dtype=torch.bfloat16) + slot = torch.arange(num_tokens, device=dev, dtype=torch.int64) + topk = torch.full((num_tokens, 2048), 7, device=dev, dtype=torch.int32) + + q_out = K.fused_norm_rope( + pos, + q_c, + qw, + EPS, + kv_c, + kvw, + EPS, + k_pe, + mla_cos_sin, + None, + None, + None, + EPS, + None, + topk, + slot_mapping=slot, + indexer_k_cache=None, + mla_kv_cache=mla_cache, + mla_kv_cache_dtype="auto", + mla_k_scale=None, + has_indexer=False, + index_rope_interleave=False, + ) + + assert_bf16(q_out, rms_norm(q_c, qw), "q_c rmsnorm (no-indexer)") + cache = mla_cache[0, :num_tokens] + assert_bf16(cache[:, :KV_LORA], rms_norm(kv_c, kvw), "MLA kv (no-indexer)") + assert_bf16( + cache[:, KV_LORA:], + rope(k_pe.float(), pos, mla_cos_sin, interleave=True), + "MLA k_pe (no-indexer)", + ) + # Shared layers reuse the previous indexer's top-k: buffer must be untouched. + assert (topk == 7).all(), "topk buffer should be untouched on shared layer" + + +# ── fused_q ────────────────────────────────────────────────────────────────── + + +@pytest.mark.parametrize("num_tokens", [1, 4, 17, 512, 4096]) +@pytest.mark.parametrize("index_interleave", [True, False]) +def test_fused_q(num_tokens: int, index_interleave: bool): + torch.manual_seed(2) + dev = "cuda" + max_pos = 8192 + pos = torch.arange(num_tokens, device=dev, dtype=torch.int64) % max_pos + + q_pe = torch.randn( + num_tokens, NUM_HEADS, ROPE_DIM, device=dev, dtype=torch.bfloat16 + ) + ql_nope = torch.randn( + num_tokens, NUM_HEADS, KV_LORA, device=dev, dtype=torch.bfloat16 + ) + index_q = torch.randn( + num_tokens, INDEX_HEADS, INDEX_HEAD_DIM, device=dev, dtype=torch.bfloat16 + ) + index_w = torch.randn(num_tokens, INDEX_HEADS, device=dev, dtype=torch.float32) + q_scale = torch.tensor([0.37], device=dev, dtype=torch.float32) + softmax_scale = INDEX_HEAD_DIM**-0.5 + head_scale = INDEX_HEADS**-0.5 + q_cos_sin = make_cos_sin(max_pos, ROPE_DIM, dev) # q_pe: interleaved + idx_cos_sin = make_cos_sin(max_pos, ROPE_DIM, dev) + + iq_fp8, iw_out, mqa = K.fused_q( + pos, + q_pe, + q_cos_sin, + index_q, + idx_cos_sin, + ql_nope, + q_scale, + index_w, + softmax_scale, + head_scale, + has_indexer=True, + index_rope_interleave=index_interleave, + ) + + s = q_scale.item() + # MQA query: [ql_nope | q_pe RoPE'd (interleaved)], per-tensor fp8. + mqa_nope_ref = (ql_nope.float() / s).to(FP8) + qpe_ref = rope( + q_pe.float(), + pos.unsqueeze(-1).expand(num_tokens, NUM_HEADS), + q_cos_sin, + interleave=True, + ) + mqa_pe_ref = (qpe_ref / s).to(FP8) + assert_fp8(mqa[:, :, :KV_LORA], mqa_nope_ref, "mqa ql_nope") + assert_fp8(mqa[:, :, KV_LORA:], mqa_pe_ref, "mqa q_pe") + + # Indexer-Q: RoPE + UE8M0 fp8 quant; index weights fold in q-scale. + iq_ref = rope( + index_q.float(), + pos.unsqueeze(-1).expand(num_tokens, INDEX_HEADS), + idx_cos_sin, + interleave=index_interleave, + ) + q_ref, scale_ref = ue8m0_quant(iq_ref) + assert_fp8(iq_fp8, q_ref, "indexer-Q fp8") + iw_ref = index_w * scale_ref * softmax_scale * head_scale + torch.testing.assert_close(iw_out, iw_ref, rtol=1e-3, atol=1e-3) + + +@pytest.mark.parametrize("num_tokens", [1, 17, 512]) +def test_fused_q_no_indexer(num_tokens: int): + torch.manual_seed(3) + dev = "cuda" + max_pos = 8192 + pos = torch.arange(num_tokens, device=dev, dtype=torch.int64) + q_pe = torch.randn( + num_tokens, NUM_HEADS, ROPE_DIM, device=dev, dtype=torch.bfloat16 + ) + ql_nope = torch.randn( + num_tokens, NUM_HEADS, KV_LORA, device=dev, dtype=torch.bfloat16 + ) + q_scale = torch.tensor([0.5], device=dev, dtype=torch.float32) + q_cos_sin = make_cos_sin(max_pos, ROPE_DIM, dev) + + _, _, mqa = K.fused_q( + pos, + q_pe, + q_cos_sin, + None, + None, + ql_nope, + q_scale, + None, + 0.0, + 0.0, + has_indexer=False, + index_rope_interleave=False, + ) + s = q_scale.item() + assert_fp8(mqa[:, :, :KV_LORA], (ql_nope.float() / s).to(FP8), "mqa ql_nope") + qpe_ref = rope( + q_pe.float(), + pos.unsqueeze(-1).expand(num_tokens, NUM_HEADS), + q_cos_sin, + interleave=True, + ) + assert_fp8(mqa[:, :, KV_LORA:], (qpe_ref / s).to(FP8), "mqa q_pe") + + +# ── fused_eh_norm (MTP) ────────────────────────────────────────────────────── + + +@pytest.mark.parametrize("num_tokens", [1, 4, 17, 512]) +def test_fused_eh_norm(num_tokens: int): + torch.manual_seed(4) + dev = "cuda" + # Mix in a position-0 token to exercise the embeds-zeroing branch. + pos = torch.arange(num_tokens, device=dev, dtype=torch.int64) + pos[0] = 0 + embeds = torch.randn(num_tokens, HIDDEN, device=dev, dtype=torch.bfloat16) + prev = torch.randn(num_tokens, HIDDEN, device=dev, dtype=torch.bfloat16) + ew = torch.randn(HIDDEN, device=dev, dtype=torch.bfloat16) + hw = torch.randn(HIDDEN, device=dev, dtype=torch.bfloat16) + + out = K.fused_eh_norm(pos, embeds, prev, ew, hw, EPS) + + masked = torch.where(pos.unsqueeze(-1) == 0, torch.zeros_like(embeds), embeds) + ref = torch.cat([rms_norm(masked, ew), rms_norm(prev, hw)], dim=-1) + assert out.shape == (num_tokens, 2 * HIDDEN) + assert_bf16(out, ref, "eh_norm") diff --git a/vllm/model_executor/layers/fused_moe/config.py b/vllm/model_executor/layers/fused_moe/config.py index b065d2142c..55c7238f64 100644 --- a/vllm/model_executor/layers/fused_moe/config.py +++ b/vllm/model_executor/layers/fused_moe/config.py @@ -1287,6 +1287,12 @@ class FusedMoEConfig: has_bias: bool = False is_lora_enabled: bool = False + # When True, the MoE skips its final cross-rank all-reduce (and the separate + # shared-expert reduce), returning the partial per-rank sum. The caller is + # then responsible for the reduction (e.g. fusing it into the next RMSNorm). + # Only honored on the non-reduced (late-AR) TP path. Default False. + skip_final_all_reduce: bool = False + # SwiGLU clamp limit. When set, backends that do not implement the clamp # are filtered out by `FusedMoEExperts.is_supported_config` so the oracle # cannot silently select one and drop the clamp. diff --git a/vllm/model_executor/layers/fused_moe/runner/moe_runner.py b/vllm/model_executor/layers/fused_moe/runner/moe_runner.py index b638db13fd..140466c7f4 100644 --- a/vllm/model_executor/layers/fused_moe/runner/moe_runner.py +++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner.py @@ -423,6 +423,7 @@ class MoERunner(MoERunnerInterface): if ( shared_output is not None and not self.moe_config.is_sequence_parallel + and not self.moe_config.skip_final_all_reduce and self._fused_output_is_reduced ): shared_output = tensor_model_parallel_all_reduce(shared_output) @@ -445,6 +446,7 @@ class MoERunner(MoERunnerInterface): # - The MK already reduced the fused output itself. if ( not self.moe_config.is_sequence_parallel + and not self.moe_config.skip_final_all_reduce and (self.moe_config.tp_size > 1 or self.moe_config.ep_size > 1) and not self._fused_output_is_reduced ): diff --git a/vllm/model_executor/layers/sparse_attn_indexer.py b/vllm/model_executor/layers/sparse_attn_indexer.py index c1bc731ee6..80c0c4ec36 100644 --- a/vllm/model_executor/layers/sparse_attn_indexer.py +++ b/vllm/model_executor/layers/sparse_attn_indexer.py @@ -227,6 +227,7 @@ def sparse_attn_indexer( topk_indices_buffer: torch.Tensor, skip_k_cache_insert: bool, use_fp4_cache: bool = False, + skip_topk_buffer_clear: bool = False, ) -> torch.Tensor: # careful! this will be None in dummy run attn_metadata = get_forward_context().attn_metadata @@ -303,7 +304,13 @@ def sparse_attn_indexer( scale_fmt, ) - topk_indices_buffer[: hidden_states.shape[0]] = -1 + # The buffer must be pre-filled with -1 (the "no token" sentinel) before the + # top-k kernels scatter valid indices into it. On the fused deepseek_v32 + # nvidia path, _fused_norm_rope_kernel already cleared the same + # [:num_tokens, :topk] region earlier in this forward, so skip the redundant + # fill. + if not skip_topk_buffer_clear: + topk_indices_buffer[: hidden_states.shape[0]] = -1 if has_prefill: prefill_metadata = attn_metadata_narrowed.prefill assert prefill_metadata is not None @@ -546,6 +553,7 @@ def sparse_attn_indexer_fake( topk_indices_buffer: torch.Tensor | None, skip_k_cache_insert: bool, use_fp4_cache: bool = False, + skip_topk_buffer_clear: bool = False, ) -> torch.Tensor: return topk_indices_buffer diff --git a/vllm/models/deepseek_v32/nvidia/attention.py b/vllm/models/deepseek_v32/nvidia/attention.py index 21b0c2c441..420e4e9378 100644 --- a/vllm/models/deepseek_v32/nvidia/attention.py +++ b/vllm/models/deepseek_v32/nvidia/attention.py @@ -1,7 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from typing import TYPE_CHECKING - import torch import torch.nn as nn from transformers import DeepseekV2Config, DeepseekV3Config @@ -23,7 +21,10 @@ from vllm.model_executor.layers.quantization.utils.fp8_utils import ( per_token_group_quant_fp8, ) from vllm.model_executor.layers.rotary_embedding import get_rope -from vllm.model_executor.layers.sparse_attn_indexer import SparseAttnIndexer +from vllm.model_executor.layers.sparse_attn_indexer import ( + SparseAttnIndexer, + sparse_attn_indexer, +) from vllm.model_executor.models.deepseek_v2 import ( DeepSeekV2FusedQkvAProjLinear, DeepseekV32IndexerCache, @@ -32,10 +33,7 @@ from vllm.model_executor.models.deepseek_v2 import ( from vllm.model_executor.models.utils import extract_layer_index from vllm.utils.torch_utils import is_quantized_kv_cache -if TYPE_CHECKING: - from vllm.model_executor.layers.attention.mla_attention import ( - MLACommonMetadata, - ) +from .kernels import fused_norm_rope, fused_q class DeepseekV32Indexer(nn.Module): @@ -160,6 +158,10 @@ class DeepseekV32Indexer(nn.Module): class DeepseekV32Attention(MLAAttention): + # Narrow the base's broadly-typed `indexer` to the concrete type so the + # `if self.indexer is not None` guards below type-check its attributes. + indexer: "DeepseekV32Indexer | None" + def __init__( self, vllm_config: VllmConfig, @@ -263,21 +265,27 @@ class DeepseekV32Attention(MLAAttention): self.num_local_heads = num_local_heads self.qk_head_dim = qk_head_dim self.indexer = indexer + self.topk_indices_buffer = topk_indices_buffer # Runtime toggle for index_share_for_mtp_iteration: MTP draft step 0 # computes the top-k, steps 1+ set this True to reuse it. self.skip_topk = False - # Whether the paged KV cache must be viewed as fp8 before the attention - # (per-tensor fp8; the fp8_ds_mla layout is read as uint8). - self._fp8_kv_needs_view = ( - is_quantized_kv_cache(self.kv_cache_dtype) - and self.kv_cache_dtype != "fp8_ds_mla" - ) - # Whether the backend takes an fp8-quantized query (FlashInfer sparse) - # vs the (ql_nope, q_pe) tuple (FlashMLA sparse). - self._use_concat_quant = ( + # Single fused fp8 path: Triton fused norm/rope/cache + fused-q write a + # single fp8 MQA query and the contiguous [kv_c; k_pe] MLA cache layout. + # This requires an fp8 KV cache and a sparse MLA backend that accepts a + # quantized query (FlashInfer sparse on SM100). + assert ( is_quantized_kv_cache(self.kv_cache_dtype) and self.impl.supports_quant_query_input + ), ( + "deepseek_v32 (nvidia) requires an fp8 KV cache served by the " + "FlashInfer sparse MLA backend (which accepts a quantized query). " + "Launch with --kv-cache-dtype fp8." ) + # The paged KV cache is stored as uint8 and viewed as fp8 for the decode + # (per-tensor fp8; never the fp8_ds_mla layout on this path). + self._fp8_kv_needs_view = self.kv_cache_dtype != "fp8_ds_mla" + # GLM-5.2 uses interleaved indexer RoPE; DeepSeek-V3.2 uses NeoX. + self._index_rope_interleave = getattr(config, "indexer_rope_interleave", False) # Remaining MLA projections (registered on this module). self.fused_qkv_a_proj = DeepSeekV2FusedQkvAProjLinear( @@ -295,10 +303,14 @@ class DeepseekV32Attention(MLAAttention): prefix=f"{prefix}.q_b_proj", ) self.kv_a_layernorm = RMSNorm(kv_lora_rank, eps=config.rms_norm_eps) + # reduce_results=False: the attention all-reduce is fused with the + # following post_attention_layernorm in the decoder layer via + # fused_allreduce_rms_norm. self.o_proj = RowParallelLinear( num_heads * v_head_dim, hidden_size, bias=False, + reduce_results=False, quant_config=quant_config, prefix=f"{prefix}.o_proj", ) @@ -322,102 +334,177 @@ class DeepseekV32Attention(MLAAttention): positions: torch.Tensor, hidden_states: torch.Tensor, ) -> torch.Tensor: + # Captured: A-projections (+ indexer A-GEMM on indexer layers). qkv_lora = self.fused_qkv_a_proj(hidden_states)[0] - q_c, kv_lora = qkv_lora.split( - [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim], dim=-1 + q_c, kv_c, k_pe = qkv_lora.split( + [self.q_lora_rank, self.kv_lora_rank, self.qk_rope_head_dim], dim=-1 ) - q_c = self.q_a_layernorm(q_c) - q = self.q_b_proj(q_c)[0] - kv_c, k_pe = kv_lora.split([self.kv_lora_rank, self.qk_rope_head_dim], dim=-1) - kv_c_normed = self.kv_a_layernorm(kv_c) - - q = q.view(-1, self.num_local_heads, self.qk_head_dim) - k_pe = k_pe.unsqueeze(1) - q[..., self.qk_nope_head_dim :], k_pe = self.rotary_emb( - positions, q[..., self.qk_nope_head_dim :], k_pe - ) + if self.indexer is not None and not self.skip_topk: + kw = self.indexer.wk_weights_proj(hidden_states)[0] + index_k = kw[:, : self.indexer.head_dim] + index_weights = kw[:, self.indexer.head_dim :] + else: + index_k = None + index_weights = None num_tokens = hidden_states.shape[0] - q_nope, q_pe = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) - q_nope = q_nope.transpose(0, 1) # (N, B, P) - ql_nope = torch.bmm(q_nope, self.W_UK_T).transpose(0, 1) # (B, N, L) - - # Lightning indexer writes the top-k indices into the shared buffer. - # "Shared" layers (indexer is None) reuse the top-k from the previous - # indexer layer already sitting in the buffer. - if self.indexer is not None and not self.skip_topk: - self.indexer(hidden_states, q_c, positions, self.indexer_rope_emb) # type: ignore[operator] - - attn_latent = torch.empty( - (num_tokens, self.num_local_heads, self.kv_lora_rank), - dtype=q.dtype, - device=q.device, - ) - self._sparse_attention(kv_c_normed, k_pe, ql_nope, q_pe, attn_latent) - - # V up-projection + output projection are metadata-independent GEMMs and - # stay captured. output = torch.empty( (num_tokens, self.num_local_heads * self.v_head_dim), - dtype=q.dtype, - device=q.device, + dtype=hidden_states.dtype, + device=hidden_states.device, + ) + self._fused_attention( + positions, q_c, kv_c, k_pe, index_k, index_weights, output ) - self._v_up_proj(attn_latent, out=output) return self.o_proj(output)[0] @eager_break_during_capture - def _sparse_attention( + def _fused_attention( self, - kv_c_normed: torch.Tensor, + positions: torch.Tensor, + q_c: torch.Tensor, + kv_c: torch.Tensor, k_pe: torch.Tensor, - ql_nope: torch.Tensor, - q_pe: torch.Tensor, - attn_latent: torch.Tensor, + index_k: torch.Tensor | None, + index_weights: torch.Tensor | None, + output: torch.Tensor, ) -> None: + # One eager break for the whole attention. In FULL cudagraph mode (pure + # decode) this decorator is a no-op, so everything here is captured; in + # PIECEWISE (prefill) it runs eagerly. The cache writes, sparse indexer, + # and forward_mqa all depend on per-step metadata and must not be split + # out (PIECEWISE capture would otherwise miss them). forward_context = get_forward_context() attn_metadata_raw = forward_context.attn_metadata - attn_metadata: MLACommonMetadata | None if isinstance(attn_metadata_raw, dict): - attn_metadata = attn_metadata_raw[self.layer_name] # type: ignore[assignment] + attn_metadata = attn_metadata_raw.get(self.layer_name) elif isinstance(attn_metadata_raw, list): - # Speculative decoding: [0] is the base-model metadata dict. - attn_metadata = attn_metadata_raw[0][self.layer_name] # type: ignore[assignment] + attn_metadata = attn_metadata_raw[0].get(self.layer_name) else: attn_metadata = attn_metadata_raw slot_mapping = forward_context.slot_mapping assert isinstance(slot_mapping, dict) - self.impl.do_kv_cache_update( # type: ignore[attr-defined] - kv_c_normed, + mla_slot = slot_mapping.get(self.layer_name) + + if self.indexer is not None: + has_indexer = True + indexer_k_norm_w = self.indexer.k_norm.weight + indexer_k_norm_bias = self.indexer.k_norm.bias + indexer_k_norm_eps = self.indexer.k_norm.variance_epsilon + indexer_k_rope_cos_sin_cache = self.indexer_rope_emb.cos_sin_cache + indexer_k_cache = self.indexer.k_cache.kv_cache + indexer_softmax_scale = self.indexer.softmax_scale + indexer_n_head_scale = self.indexer.n_head**-0.5 + else: + has_indexer = False + indexer_k_norm_w = None + indexer_k_norm_bias = None + indexer_k_norm_eps = 1e-6 + indexer_k_rope_cos_sin_cache = None + indexer_k_cache = None + indexer_softmax_scale = 0.0 + indexer_n_head_scale = 0.0 + + if attn_metadata is None: + mla_kv_cache = None + mla_k_scale = None + indexer_k_cache = None + mla_slot = None + else: + mla_kv_cache = self.kv_cache + mla_k_scale = self._k_scale + + q_c = fused_norm_rope( + positions, + q_c, + self.q_a_layernorm.weight, + self.q_a_layernorm.variance_epsilon, + kv_c, + self.kv_a_layernorm.weight, + self.kv_a_layernorm.variance_epsilon, k_pe, - self.kv_cache, - slot_mapping.get(self.layer_name), - self.kv_cache_dtype, - self._k_scale, + self.rotary_emb.cos_sin_cache, + index_k, + indexer_k_norm_w, + indexer_k_norm_bias, + indexer_k_norm_eps, + indexer_k_rope_cos_sin_cache, + self.topk_indices_buffer, + slot_mapping=mla_slot, + indexer_k_cache=indexer_k_cache, + mla_kv_cache=mla_kv_cache, + mla_kv_cache_dtype=self.kv_cache_dtype, + mla_k_scale=mla_k_scale, + has_indexer=has_indexer, + index_rope_interleave=self._index_rope_interleave, + ) + + q = self.q_b_proj(q_c)[0].view(-1, self.num_local_heads, self.qk_head_dim) + q_nope, q_pe = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) + q_nope = q_nope.transpose(0, 1) + ql_nope = torch.bmm(q_nope, self.W_UK_T).transpose(0, 1) + + if self.indexer is not None: + index_q = self.indexer.wq_b(q_c)[0] + index_q = index_q.view(-1, self.indexer.n_head, self.indexer.head_dim) + else: + index_q = None + + index_q_fp8, index_weights_out, mqa_q = fused_q( + positions, + q_pe, + self.rotary_emb.cos_sin_cache, + index_q, + self.indexer_rope_emb.cos_sin_cache if has_indexer else None, + ql_nope, + self._q_scale, + index_weights, + indexer_softmax_scale, + indexer_n_head_scale, + has_indexer=has_indexer, + index_rope_interleave=self._index_rope_interleave, ) if attn_metadata is None: - # Profile / warmup: zero-fill for DP+EP determinism. - attn_latent.zero_() + output.zero_() return - num_actual = attn_metadata.num_actual_tokens + if self.indexer is not None: + sparse_attn_indexer( + q_c, + self.indexer.k_cache.prefix, + self.indexer.k_cache.kv_cache, + index_q_fp8, + None, # q_scale folded into weights on the fp8 path + None, # k unused when skip_k_cache_insert=True + index_weights_out, + self.indexer.quant_block_size, + self.indexer.scale_fmt, + self.indexer.topk_tokens, + self.indexer.head_dim, + self.indexer.max_model_len, + self.indexer.max_total_seq_len, + self.topk_indices_buffer, + True, # skip_k_cache_insert + False, # use_fp4_cache + True, # skip_topk_buffer_clear (fused_norm_rope already did it) + ) + + num_actual = attn_metadata.num_actual_tokens # type: ignore[attr-defined] kv_cache = self.kv_cache if self._fp8_kv_needs_view: kv_cache = kv_cache.view(torch.float8_e4m3fn) - - ql_nope = ql_nope[:num_actual] - q_pe = q_pe[:num_actual] - # FlashInfer sparse takes a single fp8-quantized query; FlashMLA sparse - # takes the (ql_nope, q_pe) tuple and concatenates internally. - mqa_q: torch.Tensor | tuple[torch.Tensor, torch.Tensor] - if self._use_concat_quant: - mqa_q = self._decode_concat_quant_fp8_op(ql_nope, q_pe, self._q_scale) - else: - mqa_q = (ql_nope, q_pe) - - attn_out, _ = self.impl.forward_mqa(mqa_q, kv_cache, attn_metadata, self) # type: ignore[attr-defined] - attn_latent[:num_actual] = attn_out.view( - num_actual, self.num_local_heads, self.kv_lora_rank + attn_out, _ = self.impl.forward_mqa( # type: ignore[attr-defined] + mqa_q[:num_actual], kv_cache, attn_metadata, self ) + x = attn_out.view( + num_actual, self.num_local_heads, self.kv_lora_rank + ).transpose(0, 1) + out = ( + output[:num_actual] + .view(num_actual, self.num_local_heads, self.v_head_dim) + .transpose(0, 1) + ) + torch.bmm(x, self.W_UV, out=out) diff --git a/vllm/models/deepseek_v32/nvidia/fused_ops.py b/vllm/models/deepseek_v32/nvidia/fused_ops.py new file mode 100644 index 0000000000..6a795e2a15 --- /dev/null +++ b/vllm/models/deepseek_v32/nvidia/fused_ops.py @@ -0,0 +1,63 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Fused ops for deepseek_v32 (eager / breakable-cudagraph path). + +These recover fusions that vLLM's torch.compile passes would normally do but +that don't fire when running eager under the breakable CUDA graph. +""" + +import torch + +from vllm.distributed import ( + get_tensor_model_parallel_world_size, + tensor_model_parallel_all_reduce, +) +from vllm.model_executor.layers.fused_allreduce_gemma_rms_norm import ( + _AR_RESIDUAL_RMS_NORM, + _can_use_flashinfer, + flashinfer_trtllm_fused_allreduce_norm, +) +from vllm.model_executor.layers.layernorm import RMSNorm + + +def fused_allreduce_rms_norm( + hidden_states: torch.Tensor, + residual: torch.Tensor, + norm: RMSNorm, +) -> tuple[torch.Tensor, torch.Tensor]: + """All-reduce + add residual + (standard) RMSNorm, fused via flashinfer. + + ``hidden_states`` is the per-rank *partial* output of a row-parallel linear + run with ``reduce_results=False``; ``norm`` is the RMSNorm applied right + after. Returns ``(normed_output, new_residual)``, equivalent to + ``norm(all_reduce(hidden_states), residual)``. Falls back to an explicit + all-reduce + RMSNorm when the flashinfer fast path is unavailable. + """ + tp_size = get_tensor_model_parallel_world_size() + if tp_size == 1: + return norm(hidden_states, residual) + + if flashinfer_trtllm_fused_allreduce_norm is not None: + ok, max_token_num = _can_use_flashinfer(hidden_states, tp_size) + if ok: + norm_out = torch.empty_like(hidden_states) + # With norm_out provided, the kernel writes the new residual + # (all_reduce(hidden_states) + residual) into the hidden_states + # buffer and the normalized result into norm_out. + flashinfer_trtllm_fused_allreduce_norm( + allreduce_in=hidden_states, + residual=residual, + rms_gamma=norm.weight, + rms_eps=norm.variance_epsilon, + world_size=tp_size, + weight_bias=0.0, # standard RMSNorm (Gemma would use 1.0) + launch_with_pdl=True, + fp32_acc=True, + max_token_num=max_token_num, + pattern_code=_AR_RESIDUAL_RMS_NORM, + norm_out=norm_out, + ) + return norm_out, hidden_states + + reduced = tensor_model_parallel_all_reduce(hidden_states) + return norm(reduced, residual) diff --git a/vllm/models/deepseek_v32/nvidia/kernels.py b/vllm/models/deepseek_v32/nvidia/kernels.py new file mode 100644 index 0000000000..86419b5060 --- /dev/null +++ b/vllm/models/deepseek_v32/nvidia/kernels.py @@ -0,0 +1,823 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import torch + +from vllm.triton_utils import tl, triton + +# Cache of tiny 1-element dummy tensors (per device, dtype) reused by the +# has_indexer=False path so the indexer args don't allocate every call. +_DUMMY_CACHE: dict[tuple, torch.Tensor] = {} + + +def _dummy(shape: tuple, dtype: torch.dtype, device: torch.device) -> torch.Tensor: + key = (shape, dtype, device) + t = _DUMMY_CACHE.get(key) + if t is None: + t = torch.empty(shape, dtype=dtype, device=device) + _DUMMY_CACHE[key] = t + return t + + +@triton.jit +def _rms_norm(x, w, eps, HIDDEN_SIZE: tl.constexpr): + x = x.to(tl.float32) + mean_sq = tl.sum(x * x, axis=0) / HIDDEN_SIZE + rrms = tl.rsqrt(mean_sq + eps) + w = w.to(tl.float32) + return (x * rrms) * w + + +@triton.jit +def _get_cos_sin( + cos_sin_cache_ptr, + cos_sin_cache_stride, + pos, + HALF_ROT_DIM: tl.constexpr, +): + block = tl.arange(0, HALF_ROT_DIM) + cos = tl.load(cos_sin_cache_ptr + pos * cos_sin_cache_stride + block) + cos = cos.to(tl.float32) + sin = tl.load(cos_sin_cache_ptr + pos * cos_sin_cache_stride + block + HALF_ROT_DIM) + sin = sin.to(tl.float32) + return cos, sin + + +@triton.jit +def _fp8_ue8m0_quantize(vals): + """Quantize float32 values to FP8 E4M3 with a ue8m0 (power-of-2) scale. + + Returns (fp8_vals, scale) so the caller can store them or reuse the scale. + """ + vals = vals.to(tl.float32) + amax = tl.max(tl.abs(vals)) + scale = tl.div_rn(tl.maximum(amax, 1e-4), 448.0) + scale = tl.math.exp2(tl.math.ceil(tl.math.log2(scale))) + fp8_vals = tl.div_rn(vals, scale).to(tl.float8e4nv) + return fp8_vals, scale + + +@triton.jit +def _fp8_quant_and_cache_write( + vals, + mask, + slot_idx, + kv_cache_ptr, + kv_cache_scale_ptr, + cache_block_size, + cache_stride, + offsets, + HEAD_DIM: tl.constexpr, +): + k_fp8, scale = _fp8_ue8m0_quantize(vals) + + block_idx = slot_idx // cache_block_size + block_offset = slot_idx % cache_block_size + block_start = block_idx * cache_block_size * cache_stride + + tl.store( + kv_cache_ptr + block_start + block_offset * HEAD_DIM + offsets, + k_fp8, + mask=mask, + ) + scale_byte_off = block_start + cache_block_size * HEAD_DIM + block_offset * 4 + tl.store(kv_cache_scale_ptr + scale_byte_off // 4, scale) + + +@triton.jit +def _fused_norm_rope_kernel( + pos_ptr, + # Q RMS norm + q_c_ptr, + q_c_stride, + q_rms_norm_w_ptr, + q_rms_eps, + q_c_out_ptr, + q_c_out_stride, + Q_DIM: tl.constexpr, + Q_BLOCK_SIZE: tl.constexpr, + # KV RMS norm + kv_ptr, + kv_stride, + kv_rms_norm_w_ptr, + kv_rms_eps, + KV_DIM: tl.constexpr, + # KV RoPE + kpe_ptr, + kpe_stride, + kpe_rope_cos_sin_cache_ptr, + kpe_rope_cos_sin_cache_stride, + KPE_HALF_ROT_DIM: tl.constexpr, + # Index K layer norm + index_k_ptr, + index_k_stride, + index_k_layer_norm_w_ptr, + index_k_layer_norm_bias_ptr, + index_k_layer_norm_eps, + INDEX_K_DIM: tl.constexpr, + INDEX_K_BLOCK_SIZE: tl.constexpr, + # Index K RoPE + index_k_rope_cos_sin_cache_ptr, + index_k_rope_cos_sin_cache_stride, + INDEX_K_HALF_ROT_DIM: tl.constexpr, + # Cache params (shared by indexer K and MLA) + slot_mapping_ptr, + # Index K FP8 cache + indexer_cache_ptr, + indexer_cache_scale_ptr, + indexer_cache_block_size, + indexer_cache_stride, + # MLA KV cache (concat kv_c_normed + k_pe_roped, uses slot_mapping_ptr) + mla_cache_ptr, + mla_cache_block_stride, + mla_cache_entry_stride, + MLA_CACHE_FP8: tl.constexpr, + mla_cache_scale_ptr, + # Top k indices + topk_indices_ptr, + topk_indices_stride, + TOPK: tl.constexpr, + TOPK_BLOCK_SIZE: tl.constexpr, + HAS_INDEXER: tl.constexpr, + INDEX_ROPE_INTERLEAVE: tl.constexpr, +): + pid = tl.program_id(0) + tok_idx = tl.program_id(1) + if pid == 3: + if not HAS_INDEXER: + # Shared layer: reuse the previous indexer layer's top-k; do not + # clear the buffer. + return + # Fill top k indices buffer with -1 + for i in range(0, TOPK, TOPK_BLOCK_SIZE): + offset = i + tl.arange(0, TOPK_BLOCK_SIZE) + mask = offset < TOPK + tl.store( + topk_indices_ptr + tok_idx * topk_indices_stride + offset, + -1, + mask=mask, + ) + return + + if slot_mapping_ptr is None: + # Memory profiling run. + return + slot_idx = tl.load(slot_mapping_ptr + tok_idx) + if slot_idx < 0: + # Padding + return + + if pid == 2: + # Q RMS norm + q_block = tl.arange(0, Q_BLOCK_SIZE) + q_mask = q_block < Q_DIM + q_c = tl.load(q_c_ptr + tok_idx * q_c_stride + q_block, mask=q_mask, other=0.0) + q_c_rms_w = tl.load(q_rms_norm_w_ptr + q_block, mask=q_mask) + q_c = _rms_norm(q_c, q_c_rms_w, q_rms_eps, Q_DIM) + tl.store(q_c_out_ptr + tok_idx * q_c_out_stride + q_block, q_c, mask=q_mask) + elif pid == 1: + # KV RMS Norm + KV RoPE + MLA concat_and_cache. + # Merged so the normed kv_c and RoPE'd k_pe can be written + # to the MLA KV cache directly without a separate kernel. + + # KV RMS Norm (result stays in registers for MLA cache write) + kv_block = tl.arange(0, KV_DIM) + kv_c = tl.load(kv_ptr + tok_idx * kv_stride + kv_block) + kv_c_rms_w = tl.load(kv_rms_norm_w_ptr + kv_block) + kv_c = _rms_norm(kv_c, kv_c_rms_w, kv_rms_eps, KV_DIM) + + # KV RoPE (interleaved) on k_pe — in registers only. + # k_pe is not needed after the cache write (MLA decode reads + # from kv_cache), so we skip writing back to kpe_ptr. + pos = tl.load(pos_ptr + tok_idx) + cos, sin = _get_cos_sin( + kpe_rope_cos_sin_cache_ptr, + kpe_rope_cos_sin_cache_stride, + pos, + KPE_HALF_ROT_DIM, + ) + dim_off = tl.arange(0, KPE_HALF_ROT_DIM) + kpe_base = kpe_ptr + tok_idx * kpe_stride + x1 = tl.load(kpe_base + dim_off * 2).to(tl.float32) + x2 = tl.load(kpe_base + dim_off * 2 + 1).to(tl.float32) + r1 = x1 * cos - x2 * sin + r2 = x2 * cos + x1 * sin + + # MLA concat_and_cache: write [kv_c_normed, k_pe_roped] to cache. + if mla_cache_entry_stride == 0: + return + + mla_block_size = mla_cache_block_stride // mla_cache_entry_stride + mla_block_idx = slot_idx // mla_block_size + mla_block_off = slot_idx % mla_block_size + dst = ( + mla_cache_ptr + + mla_block_idx * mla_cache_block_stride + + mla_block_off * mla_cache_entry_stride + ) + # kv_c_normed (KV_DIM elements) + if MLA_CACHE_FP8: + scale = tl.load(mla_cache_scale_ptr) + kv_c_fp8 = (kv_c.to(tl.float32) / scale).to(tl.float8e4nv) + tl.store(dst + kv_block, kv_c_fp8) + else: + tl.store(dst + kv_block, kv_c) + # k_pe_roped (from registers, interleaved layout) + if MLA_CACHE_FP8: + tl.store(dst + KV_DIM + dim_off * 2, (r1 / scale).to(tl.float8e4nv)) + tl.store(dst + KV_DIM + dim_off * 2 + 1, (r2 / scale).to(tl.float8e4nv)) + else: + tl.store(dst + KV_DIM + dim_off * 2, r1) + tl.store(dst + KV_DIM + dim_off * 2 + 1, r2) + elif pid == 0: + if not HAS_INDEXER: + # Shared layer: no indexer K to process. + return + # Fused: Index K LayerNorm + RoPE + FP8 quant + cache write. + # Eliminates the separate indexer_k_quant_and_cache kernel launch. + + index_k_block = tl.arange(0, INDEX_K_BLOCK_SIZE) + index_k_mask = index_k_block < INDEX_K_DIM + index_k = tl.load( + index_k_ptr + tok_idx * index_k_stride + index_k_block, + mask=index_k_mask, + other=0.0, + ).to(tl.float32) + index_k_w = tl.load( + index_k_layer_norm_w_ptr + index_k_block, mask=index_k_mask + ).to(tl.float32) + index_k_b = tl.load( + index_k_layer_norm_bias_ptr + index_k_block, mask=index_k_mask + ).to(tl.float32) + + # 1. LayerNorm. Keep (mean, rstd) so the RoPE rotation partner can be + # re-normalized in registers below, avoiding a global scratch buffer. + mean = tl.sum(index_k, axis=0) / INDEX_K_DIM + diff = tl.where(index_k_mask, index_k - mean, 0.0) + var = tl.sum(diff * diff, axis=0) / INDEX_K_DIM + rstd = tl.rsqrt(var + index_k_layer_norm_eps) + normed = (index_k - mean) * rstd * index_k_w + index_k_b + + # 2. RoPE on the rotation region. Supports both interleaved (adjacent + # pairs, e.g. GLM-5.2) and NeoX (split-half, e.g. DeepSeek-V3.2). The + # rotation partner is gathered from the read-only inputs and + # re-normalized with the same (mean, rstd) — no scratch, no atomics. + pos = tl.load(pos_ptr + tok_idx) + in_rope = index_k_block < 2 * INDEX_K_HALF_ROT_DIM + if INDEX_ROPE_INTERLEAVE: + # pair i = block // 2; partner = block ^ 1; even -> -sin, odd -> +sin. + cos_idx = index_k_block // 2 + partner_offs = tl.where(in_rope, index_k_block ^ 1, index_k_block) + sign = tl.where(index_k_block % 2 == 0, -1.0, 1.0) + else: + # NeoX: pair across halves; partner = block ^ HALF. + cos_idx = index_k_block % INDEX_K_HALF_ROT_DIM + partner_offs = tl.where( + in_rope, index_k_block ^ INDEX_K_HALF_ROT_DIM, index_k_block + ) + sign = tl.where(index_k_block < INDEX_K_HALF_ROT_DIM, -1.0, 1.0) + cos_full = tl.load( + index_k_rope_cos_sin_cache_ptr + + pos * index_k_rope_cos_sin_cache_stride + + cos_idx, + mask=in_rope, + other=1.0, + ).to(tl.float32) + sin_full = tl.load( + index_k_rope_cos_sin_cache_ptr + + pos * index_k_rope_cos_sin_cache_stride + + INDEX_K_HALF_ROT_DIM + + cos_idx, + mask=in_rope, + other=0.0, + ).to(tl.float32) + # normed[partner_offs] == (raw_partner - mean) * rstd * w_partner + + # b_partner: gather the raw partner and its norm affine (read-only + # loads), then apply the same per-token mean/rstd. + raw_partner = tl.load( + index_k_ptr + tok_idx * index_k_stride + partner_offs, + mask=index_k_mask, + other=0.0, + ).to(tl.float32) + w_partner = tl.load( + index_k_layer_norm_w_ptr + partner_offs, mask=index_k_mask + ).to(tl.float32) + b_partner = tl.load( + index_k_layer_norm_bias_ptr + partner_offs, mask=index_k_mask + ).to(tl.float32) + normed_partner = (raw_partner - mean) * rstd * w_partner + b_partner + roped = normed * cos_full + sign * normed_partner * sin_full + result = tl.where(in_rope, roped, normed) + + # 3. FP8 quantize + cache write from registers. + # No need to write back to index_k_ptr — the only consumer + # (sparse_attn_indexer) reads from the cache, not index_k. + _fp8_quant_and_cache_write( + result, + index_k_mask, + slot_idx, + indexer_cache_ptr, + indexer_cache_scale_ptr, + indexer_cache_block_size, + indexer_cache_stride, + index_k_block, + INDEX_K_DIM, + ) + + +def fused_norm_rope( + positions: torch.Tensor, + q_c: torch.Tensor, + q_rms_norm_w: torch.Tensor, + q_rms_eps: float, + kv_c: torch.Tensor, + kv_rms_norm_w: torch.Tensor, + kv_rms_eps: float, + k_pe: torch.Tensor, + k_rope_cos_sin_cache: torch.Tensor, + index_k: torch.Tensor | None, + index_k_layer_norm_w: torch.Tensor | None, + index_k_layer_norm_bias: torch.Tensor | None, + index_k_layer_norm_eps: float, + index_k_rope_cos_sin_cache: torch.Tensor | None, + topk_indices_buffer: torch.Tensor, + # Cache params for fused writes (single slot_mapping for both caches) + slot_mapping: torch.Tensor | None = None, + indexer_k_cache: torch.Tensor | None = None, + mla_kv_cache: torch.Tensor | None = None, + mla_kv_cache_dtype: str = "auto", + mla_k_scale: torch.Tensor | None = None, + has_indexer: bool = True, + index_rope_interleave: bool = False, + q_c_out: torch.Tensor | None = None, +) -> torch.Tensor: + assert positions.ndim == 1 + assert q_c.ndim == 2 + assert kv_c.ndim == 2 + assert k_pe.ndim == 2 + assert topk_indices_buffer.ndim == 2 + + num_tokens = positions.shape[0] + q_dim = q_c.shape[-1] + kv_dim = kv_c.shape[-1] + device = positions.device + + # Shared (no-indexer) layers: substitute cached 1-element dummies so the + # kernel launches cleanly; pid 0/3 (indexer + topk fill) skipped by + # HAS_INDEXER and never dereference them. + if not has_indexer: + indexer_k_cache = None + index_k = _dummy((1, 1), q_c.dtype, device) + index_k_layer_norm_w = _dummy((1,), torch.float32, device) + index_k_layer_norm_bias = _dummy((1,), torch.float32, device) + index_k_rope_cos_sin_cache = k_rope_cos_sin_cache + assert index_k is not None + assert index_k_rope_cos_sin_cache is not None + index_k_dim = index_k.shape[-1] + topk = topk_indices_buffer.shape[-1] + + # --- Indexer K cache setup --- + if indexer_k_cache is not None: + assert slot_mapping is not None + idx_cache_scale_view = indexer_k_cache.view(torch.uint8).view(torch.float32) + idx_cache_block_size = indexer_k_cache.shape[1] + idx_cache_stride = indexer_k_cache.shape[2] + if indexer_k_cache.dtype == torch.uint8: + indexer_k_cache = indexer_k_cache.view(torch.float8_e4m3fn) + else: + # No indexer cache (shared layer / MLA-only fusion). Use dummies but + # KEEP the caller's slot_mapping so the MLA write (pid 1) still runs. + idx_cache_scale_view = torch.empty(0, dtype=torch.float32, device=device) + indexer_k_cache = torch.empty(0, dtype=torch.float8_e4m3fn, device=device) + idx_cache_block_size = 1 + idx_cache_stride = 1 + if mla_kv_cache is None: + # Pure profiling run (no caches at all): skip all per-token writes. + slot_mapping = torch.full( + (num_tokens,), -1, dtype=torch.int64, device=device + ) + + # --- MLA KV cache setup --- + mla_cache_fp8 = mla_kv_cache_dtype != "auto" + if mla_kv_cache is not None: + mla_block_stride = mla_kv_cache.stride(0) + mla_entry_stride = mla_kv_cache.stride(1) + if mla_cache_fp8 and mla_kv_cache.dtype == torch.uint8: + mla_kv_cache = mla_kv_cache.view(torch.float8_e4m3fn) + if mla_k_scale is None: + mla_k_scale = torch.ones(1, dtype=torch.float32, device=device) + else: + # Dummy values — pid 2 will skip the MLA cache write because + # slot_mapping is all -1. + mla_kv_cache = torch.empty(0, dtype=torch.bfloat16, device=device) + mla_block_stride = 0 + mla_entry_stride = 0 + mla_k_scale = torch.ones(1, dtype=torch.float32, device=device) + + if q_c_out is None: + q_c_out = torch.empty_like(q_c) + _fused_norm_rope_kernel[(4, num_tokens)]( + positions, + # Q RMS norm + q_c, + q_c.stride(0), + q_rms_norm_w, + q_rms_eps, + q_c_out, + q_c_out.stride(0), + q_dim, + triton.next_power_of_2(q_dim), + # KV RMS norm + kv_c, + kv_c.stride(0), + kv_rms_norm_w, + kv_rms_eps, + kv_dim, + # KV RoPE + k_pe, + k_pe.stride(0), + k_rope_cos_sin_cache, + k_rope_cos_sin_cache.stride(0), + k_rope_cos_sin_cache.shape[-1] // 2, + # Index K layer norm + RoPE + FP8 quant + index_k, + index_k.stride(0), + index_k_layer_norm_w, + index_k_layer_norm_bias, + index_k_layer_norm_eps, + index_k_dim, + triton.next_power_of_2(index_k_dim), + index_k_rope_cos_sin_cache, + index_k_rope_cos_sin_cache.stride(0), + index_k_rope_cos_sin_cache.shape[-1] // 2, + # Cache params + slot_mapping, + indexer_k_cache, + idx_cache_scale_view, + idx_cache_block_size, + idx_cache_stride, + # MLA KV cache (uses same slot_mapping) + mla_kv_cache, + mla_block_stride, + mla_entry_stride, + mla_cache_fp8, + mla_k_scale, + # Top k indices buffer + topk_indices_buffer, + topk_indices_buffer.stride(0), + topk, + TOPK_BLOCK_SIZE=1024, + HAS_INDEXER=has_indexer, + INDEX_ROPE_INTERLEAVE=index_rope_interleave, + ) + return q_c_out + + +@triton.jit +def _fused_q_kernel( + pos_ptr, + # MQA query PE: RoPE + FP8 pack into output tail + q_pe_ptr, + q_pe_stride0, + q_pe_stride1, + NUM_Q_HEADS: tl.constexpr, + q_pe_cos_sin_ptr, + q_pe_cos_sin_stride, + Q_PE_HALF_ROT_DIM: tl.constexpr, + # Index Q RoPE + index_q_ptr, + index_q_stride0, + index_q_stride1, + NUM_INDEX_Q_HEADS: tl.constexpr, + index_q_cos_sin_ptr, + index_q_cos_sin_stride, + INDEX_Q_HALF_ROT_DIM: tl.constexpr, + # Index Q Quantize + index_q_fp8_ptr, + index_q_fp8_stride0, + index_q_fp8_stride1, + INDEX_Q_HEAD_DIM: tl.constexpr, + # MQA query pack: quantize ql_nope and RoPE+quantize q_pe into mqa_q_fp8 + ql_nope_ptr, + ql_nope_stride0, + ql_nope_stride1, + mqa_q_fp8_ptr, + mqa_q_fp8_stride0, + mqa_q_fp8_stride1, + q_scale_ptr, + QL_NOPE_DIM: tl.constexpr, + QL_NOPE_BLOCK: tl.constexpr, + # Index weights + index_weights_ptr, + index_weights_stride, + index_weights_softmax_scale, + index_weights_head_scale, + index_weights_out_ptr, + index_weights_out_stride, + HAS_INDEXER: tl.constexpr, + INDEX_ROPE_INTERLEAVE: tl.constexpr, +): + pid = tl.program_id(0) + tok_idx = tl.program_id(1) + head_idx = tl.program_id(2) + + if pid == 2: + # ql_nope quantize + pack into the front of mqa_q_fp8. + if 2 * head_idx >= NUM_Q_HEADS: + return + + scale = tl.load(q_scale_ptr) + for local_head in range(2): + q_head_idx = head_idx * 2 + local_head + if q_head_idx < NUM_Q_HEADS: + ql_nope_off = tl.arange(0, QL_NOPE_BLOCK) + ql_nope_mask = ql_nope_off < QL_NOPE_DIM + ql_nope = tl.load( + ql_nope_ptr + + tok_idx * ql_nope_stride0 + + q_head_idx * ql_nope_stride1 + + ql_nope_off, + mask=ql_nope_mask, + ).to(tl.float32) + ql_nope_fp8 = (ql_nope / scale).to(tl.float8e4nv) + tl.store( + mqa_q_fp8_ptr + + tok_idx * mqa_q_fp8_stride0 + + q_head_idx * mqa_q_fp8_stride1 + + ql_nope_off, + ql_nope_fp8, + mask=ql_nope_mask, + ) + return + elif pid == 0: + # q_pe RoPE + quantize + pack into the tail of mqa_q_fp8. + if 2 * head_idx >= NUM_Q_HEADS: + return + + pos = tl.load(pos_ptr + tok_idx) + cos, sin = _get_cos_sin( + q_pe_cos_sin_ptr, + q_pe_cos_sin_stride, + pos, + Q_PE_HALF_ROT_DIM, + ) + + scale = tl.load(q_scale_ptr) + for local_head in range(2): + q_head_idx = head_idx * 2 + local_head + if q_head_idx < NUM_Q_HEADS: + rot_off = tl.arange(0, Q_PE_HALF_ROT_DIM) + x1 = tl.load( + q_pe_ptr + + tok_idx * q_pe_stride0 + + q_head_idx * q_pe_stride1 + + rot_off * 2, + ).to(tl.float32) + x2 = tl.load( + q_pe_ptr + + tok_idx * q_pe_stride0 + + q_head_idx * q_pe_stride1 + + rot_off * 2 + + 1 + ).to(tl.float32) + r1 = x1 * cos - x2 * sin + r2 = x2 * cos + x1 * sin + tl.store( + mqa_q_fp8_ptr + + tok_idx * mqa_q_fp8_stride0 + + q_head_idx * mqa_q_fp8_stride1 + + QL_NOPE_DIM + + rot_off * 2, + (r1 / scale).to(tl.float8e4nv), + ) + tl.store( + mqa_q_fp8_ptr + + tok_idx * mqa_q_fp8_stride0 + + q_head_idx * mqa_q_fp8_stride1 + + QL_NOPE_DIM + + rot_off * 2 + + 1, + (r2 / scale).to(tl.float8e4nv), + ) + return + elif pid == 1: + # Index Q RoPE + fp8 quant, all in registers. The roped bf16 index_q is + # never consumed (only the fp8 below is), so we avoid an in-place + # store-then-reload round-trip. + if not HAS_INDEXER: + return + if head_idx >= NUM_INDEX_Q_HEADS: + return + + pos = tl.load(pos_ptr + tok_idx) + index_q_block = tl.arange(0, INDEX_Q_HEAD_DIM) + iq_base = index_q_ptr + tok_idx * index_q_stride0 + head_idx * index_q_stride1 + index_q = tl.load(iq_base + index_q_block).to(tl.float32) + + # RoPE in registers (interleaved for GLM-5.2, NeoX for DeepSeek-V3.2), + # gathering the rotation partner from the read-only input. + in_rope = index_q_block < 2 * INDEX_Q_HALF_ROT_DIM + if INDEX_ROPE_INTERLEAVE: + cos_idx = index_q_block // 2 + partner_offs = tl.where(in_rope, index_q_block ^ 1, index_q_block) + sign = tl.where(index_q_block % 2 == 0, -1.0, 1.0) + else: + cos_idx = index_q_block % INDEX_Q_HALF_ROT_DIM + partner_offs = tl.where( + in_rope, index_q_block ^ INDEX_Q_HALF_ROT_DIM, index_q_block + ) + sign = tl.where(index_q_block < INDEX_Q_HALF_ROT_DIM, -1.0, 1.0) + cos_full = tl.load( + index_q_cos_sin_ptr + pos * index_q_cos_sin_stride + cos_idx, + mask=in_rope, + other=1.0, + ).to(tl.float32) + sin_full = tl.load( + index_q_cos_sin_ptr + + pos * index_q_cos_sin_stride + + INDEX_Q_HALF_ROT_DIM + + cos_idx, + mask=in_rope, + other=0.0, + ).to(tl.float32) + partner = tl.load(iq_base + partner_offs).to(tl.float32) + roped = index_q * cos_full + sign * partner * sin_full + index_q = tl.where(in_rope, roped, index_q) + + # Index Q Quantize (from registers) + index_q_fp8, index_q_scale = _fp8_ue8m0_quantize(index_q) + tl.store( + index_q_fp8_ptr + + tok_idx * index_q_fp8_stride0 + + head_idx * index_q_fp8_stride1 + + index_q_block, + index_q_fp8, + ) + + # Index weights update + index_weights = tl.load( + index_weights_ptr + tok_idx * index_weights_stride + head_idx + ) + index_weights = index_weights.to(tl.float32) + index_weights *= index_q_scale + index_weights *= index_weights_softmax_scale + index_weights *= index_weights_head_scale + tl.store( + index_weights_out_ptr + tok_idx * index_weights_out_stride + head_idx, + index_weights, + ) + + +def fused_q( + positions: torch.Tensor, + q_pe: torch.Tensor, + q_pe_cos_sin_cache: torch.Tensor, + index_q: torch.Tensor | None, + index_q_cos_sin_cache: torch.Tensor | None, + ql_nope: torch.Tensor, + q_scale: torch.Tensor, + # Index weights + index_weights: torch.Tensor | None, + index_weights_softmax_scale: float, + index_weights_head_scale: float, + has_indexer: bool = True, + index_rope_interleave: bool = False, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + assert positions.ndim == 1 + assert q_pe.ndim == 3 + assert q_pe_cos_sin_cache.ndim == 2 + assert ql_nope.ndim == 3 + assert ql_nope.shape[:2] == q_pe.shape[:2] + + num_tokens = positions.shape[0] + num_q_heads = q_pe.shape[1] + # Grid's 3rd dim must cover the MQA-pack heads (pid 0/2 iterate 2 heads + # each) and, when present, the indexer heads (pid 1). + mqa_grid_heads = (num_q_heads + 1) // 2 + if not has_indexer: + # Shared layer: cached 1-element dummies; pid 1 skipped by HAS_INDEXER + # and never dereferences them. + index_q = _dummy((1, 1, 1), q_pe.dtype, q_pe.device) + index_q_cos_sin_cache = q_pe_cos_sin_cache + index_weights = _dummy((1, 1), torch.float32, q_pe.device) + assert index_q is not None and index_q.ndim == 3 + assert index_q_cos_sin_cache is not None + assert index_weights is not None + num_index_q_heads = index_q.shape[1] + index_q_head_dim = index_q.shape[2] + grid_heads = max(mqa_grid_heads, num_index_q_heads) + mqa_q_fp8 = torch.empty( + q_pe.shape[0], + q_pe.shape[1], + ql_nope.shape[2] + q_pe.shape[2], + dtype=torch.float8_e4m3fn, + device=q_pe.device, + ) + + index_q_fp8 = torch.empty_like(index_q, dtype=torch.float8_e4m3fn) + index_weights_out = torch.empty_like(index_weights, dtype=torch.float32) + _fused_q_kernel[(3, num_tokens, grid_heads)]( + positions, + q_pe, + q_pe.stride(0), + q_pe.stride(1), + num_q_heads, + q_pe_cos_sin_cache, + q_pe_cos_sin_cache.stride(0), + q_pe_cos_sin_cache.shape[-1] // 2, + index_q, + index_q.stride(0), + index_q.stride(1), + num_index_q_heads, + index_q_cos_sin_cache, + index_q_cos_sin_cache.stride(0), + index_q_cos_sin_cache.shape[-1] // 2, + index_q_fp8, + index_q_fp8.stride(0), + index_q_fp8.stride(1), + index_q_head_dim, + ql_nope, + ql_nope.stride(0), + ql_nope.stride(1), + mqa_q_fp8, + mqa_q_fp8.stride(0), + mqa_q_fp8.stride(1), + q_scale, + ql_nope.shape[2], + triton.next_power_of_2(ql_nope.shape[2]), + index_weights, + index_weights.stride(0), + index_weights_softmax_scale, + index_weights_head_scale, + index_weights_out, + index_weights_out.stride(0), + HAS_INDEXER=has_indexer, + INDEX_ROPE_INTERLEAVE=index_rope_interleave, + # num_warps=1 is optimal here: each program is a single 128-element + # rope+quant, so the kernel is program-count/occupancy bound, not + # per-program compute bound (swept 1/2/4/8 — 1 wins or ties everywhere). + num_warps=1, + ) + return index_q_fp8, index_weights_out, mqa_q_fp8 + + +@triton.jit +def _fused_eh_norm_kernel( + pos_ptr, + embeds_ptr, + embeds_stride, + prev_ptr, + prev_stride, + enorm_w_ptr, + hnorm_w_ptr, + eps, + out_ptr, + out_stride, + H: tl.constexpr, + BLOCK: tl.constexpr, +): + """MTP input fusion: zero embeds at position 0, RMSNorm(embeds) with enorm + and RMSNorm(prev_hidden) with hnorm, written side-by-side into ``out`` + ([N, 2H]) ready for the eh_proj GEMM. Replaces where + 2x RMSNorm + cat.""" + tok = tl.program_id(0) + off = tl.arange(0, BLOCK) + mask = off < H + + pos = tl.load(pos_ptr + tok) + e = tl.load(embeds_ptr + tok * embeds_stride + off, mask=mask, other=0.0) + e = tl.where(pos == 0, 0.0, e.to(tl.float32)) + ew = tl.load(enorm_w_ptr + off, mask=mask) + e_normed = _rms_norm(e, ew, eps, H) + tl.store(out_ptr + tok * out_stride + off, e_normed, mask=mask) + + p = tl.load(prev_ptr + tok * prev_stride + off, mask=mask, other=0.0) + hw = tl.load(hnorm_w_ptr + off, mask=mask) + p_normed = _rms_norm(p, hw, eps, H) + tl.store(out_ptr + tok * out_stride + H + off, p_normed, mask=mask) + + +def fused_eh_norm( + positions: torch.Tensor, + inputs_embeds: torch.Tensor, + previous_hidden: torch.Tensor, + enorm_w: torch.Tensor, + hnorm_w: torch.Tensor, + eps: float, +) -> torch.Tensor: + """Returns cat([enorm(masked embeds), hnorm(prev_hidden)]) -> [N, 2H].""" + n, h = inputs_embeds.shape + out = torch.empty(n, 2 * h, dtype=inputs_embeds.dtype, device=inputs_embeds.device) + _fused_eh_norm_kernel[(n,)]( + positions, + inputs_embeds, + inputs_embeds.stride(0), + previous_hidden, + previous_hidden.stride(0), + enorm_w, + hnorm_w, + eps, + out, + out.stride(0), + h, + triton.next_power_of_2(h), + ) + return out diff --git a/vllm/models/deepseek_v32/nvidia/model.py b/vllm/models/deepseek_v32/nvidia/model.py index dd9e1d65ea..0cd7fb02ed 100644 --- a/vllm/models/deepseek_v32/nvidia/model.py +++ b/vllm/models/deepseek_v32/nvidia/model.py @@ -36,6 +36,7 @@ from vllm.model_executor.models.utils import ( from vllm.sequence import IntermediateTensors from .attention import DeepseekV32Attention +from .fused_ops import fused_allreduce_rms_norm class DeepseekV32DecoderLayer(torch.nn.Module): @@ -77,6 +78,10 @@ class DeepseekV32DecoderLayer(torch.nn.Module): quant_config=quant_config, prefix=f"{prefix}.mlp", ) + # Defer the MoE cross-rank all-reduce; it is fused into the next + # layer's input_layernorm (or the final norm) via + # fused_allreduce_rms_norm. self.mlp.experts is the MoERunner. + self.mlp.experts.moe_config.skip_final_all_reduce = True else: self.mlp = DeepseekV2MLP( hidden_size=config.hidden_size, @@ -84,6 +89,7 @@ class DeepseekV32DecoderLayer(torch.nn.Module): hidden_act=config.hidden_act, quant_config=quant_config, prefix=f"{prefix}.mlp", + reduce_results=False, ) self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.post_attention_layernorm = RMSNorm( @@ -98,12 +104,23 @@ class DeepseekV32DecoderLayer(torch.nn.Module): residual: torch.Tensor | None, ) -> tuple[torch.Tensor, torch.Tensor]: if residual is None: + # First layer: hidden_states is the (already reduced) embedding. residual = hidden_states hidden_states = self.input_layernorm(hidden_states) else: - hidden_states, residual = self.input_layernorm(hidden_states, residual) + # The previous layer's MLP/MoE output is left un-reduced; fuse its + # all-reduce into this input_layernorm. + hidden_states, residual = fused_allreduce_rms_norm( + hidden_states, residual, self.input_layernorm + ) + # self_attn's o_proj runs reduce_results=False; fuse its all-reduce with + # the post-attention RMSNorm. hidden_states = self.self_attn(positions=positions, hidden_states=hidden_states) - hidden_states, residual = self.post_attention_layernorm(hidden_states, residual) + hidden_states, residual = fused_allreduce_rms_norm( + hidden_states, residual, self.post_attention_layernorm + ) + # MLP/MoE runs un-reduced; its all-reduce is fused into the next layer's + # input_layernorm (or the model's final norm). hidden_states = self.mlp(hidden_states) return hidden_states, residual @@ -201,7 +218,8 @@ class DeepseekV32Model(torch.nn.Module): {"hidden_states": hidden_states, "residual": residual} ) - hidden_states, _ = self.norm(hidden_states, residual) + # Last layer's MoE output is un-reduced; fuse its all-reduce into norm. + hidden_states, _ = fused_allreduce_rms_norm(hidden_states, residual, self.norm) if len(aux_hidden_states) > 0: return hidden_states, aux_hidden_states return hidden_states diff --git a/vllm/models/deepseek_v32/nvidia/mtp.py b/vllm/models/deepseek_v32/nvidia/mtp.py index 482ebecc52..7a04d4f3d3 100644 --- a/vllm/models/deepseek_v32/nvidia/mtp.py +++ b/vllm/models/deepseek_v32/nvidia/mtp.py @@ -1,6 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project - import typing from collections.abc import Callable, Iterable @@ -9,6 +8,7 @@ import torch.nn as nn from vllm._aiter_ops import rocm_aiter_ops from vllm.config import VllmConfig +from vllm.distributed import tensor_model_parallel_all_reduce from vllm.model_executor.layers.fused_moe import ( fused_moe_make_expert_params_mapping, ) @@ -35,6 +35,7 @@ from vllm.model_executor.models.utils import ( from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors +from .kernels import fused_eh_norm from .model import DeepseekV32DecoderLayer @@ -75,15 +76,23 @@ class DeepseekV32MultiTokenPredictorLayer(nn.Module): spec_step_index: int = 0, ) -> torch.Tensor: assert inputs_embeds is not None - inputs_embeds = torch.where(positions.unsqueeze(-1) == 0, 0, inputs_embeds) - inputs_embeds = self.enorm(inputs_embeds) - previous_hidden_states = self.hnorm(previous_hidden_states) - hidden_states = self.eh_proj( - torch.cat([inputs_embeds, previous_hidden_states], dim=-1) + # Fused: zero pos-0 embeds + enorm(embeds) + hnorm(prev) + cat -> [N, 2H]. + eh_input = fused_eh_norm( + positions, + inputs_embeds, + previous_hidden_states, + self.enorm.weight, + self.hnorm.weight, + self.enorm.variance_epsilon, ) + hidden_states = self.eh_proj(eh_input) hidden_states, residual = self.mtp_block( positions=positions, hidden_states=hidden_states, residual=None ) + # mtp_block's MoE output is left un-reduced (skip_final_all_reduce); the + # main model fuses that all-reduce into the next norm, but here the + # recycle hidden is consumed directly, so reduce it now. + hidden_states = tensor_model_parallel_all_reduce(hidden_states) # Return the pre-final-norm recycle hidden (re-fed as the next spec # step's previous_hidden_states); shared_head norm is applied in # compute_logits. Matches the V2-runner / deepseek_v4 MTP contract.