[GLM5] Implement op fusion for GLM5/DSV3.2 (#46876)
This commit is contained in:
Woosuk Kwon
@@ -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")
|
||||
@@ -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.
|
||||
|
||||
@@ -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
|
||||
):
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
@@ -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
|
||||
@@ -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
|
||||
|
||||
@@ -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.
|
||||
|
||||
Reference in New Issue
Block a user