diff --git a/tests/test_normalization.py b/tests/test_normalization.py index f7111244..c3240ae5 100644 --- a/tests/test_normalization.py +++ b/tests/test_normalization.py @@ -69,3 +69,114 @@ def test_rmsnorm_explicit_epsilon(): # Variance of ones is 1.0; 1.0 / sqrt(1.0 + 1.0) = 1.0 / sqrt(2) ≈ 0.7071 expected_large = 1.0 / (2.0**0.5) assert torch.allclose(out_large, torch.full_like(x, expected_large), atol=1e-4) + + +@pytest.mark.gpu +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32]) +@pytest.mark.parametrize("width", [2048, 5280, 10560]) +def test_fused_rmsnorm_matches_reference_within_rounding(dtype, width): + from vllm_rlt.kernels.norm import fused_rmsnorm, fused_rmsnorm_supported + + torch.manual_seed(0) + norm = RMSNorm(width).to(device="cuda", dtype=dtype) + with torch.no_grad(): + norm.weight.copy_(torch.randn(width, device="cuda") * 0.5 + 1.0) + hidden = (torch.randn(37, width, device="cuda") * 3).to(dtype) + assert fused_rmsnorm_supported(hidden) + fused = fused_rmsnorm(hidden, norm.weight, norm.variance_epsilon) + reference = norm.reference(hidden) + assert fused.dtype == dtype and fused.shape == reference.shape + # Same arithmetic and rounding points; only the FP32 reduction order differs. A + # flipped normalized value can move the weighted output by up to two ulps + # (two chained roundings); the vast majority of elements are identical. + if dtype == torch.float32: + assert torch.allclose(fused, reference, rtol=1e-5, atol=1e-6) + else: + diff = (fused.float() - reference.float()).abs() + ulp = torch.finfo(dtype).eps * reference.float().abs().clamp_min(torch.finfo(dtype).tiny) + assert (diff <= 2.01 * ulp).all() + assert (diff == 0).float().mean().item() > 0.99 + + +@pytest.mark.gpu +def test_fused_rmsnorm_residual_epilogue_rounds_once(): + """The approx kernel's residual epilogue computes round(residual + w * normalized). + + torch's separate add rounds the weight product first. The two oracles differ + in ~18% of elements, by more than one ulp of the result under cancellation, so + the check is membership: every element equals the single-rounded oracle (or, + in rare fp16 cases, the double-rounded one). + """ + from vllm_rlt.kernels.norm import fused_rmsnorm + + torch.manual_seed(1) + norm = RMSNorm(2048).to(device="cuda", dtype=torch.bfloat16) + with torch.no_grad(): + # Unit weights would make the product exact and hide the epilogue's rounding. + norm.weight.copy_(torch.randn(2048, device="cuda") * 0.5 + 1.0) + hidden = (torch.randn(9, 2048, device="cuda") * 3).to(torch.bfloat16) + residual = (torch.randn(9, 2048, device="cuda") * 3).to(torch.bfloat16) + ones = torch.ones_like(norm.weight) + normalized = fused_rmsnorm(hidden, ones, norm.variance_epsilon) # round(x * inv), exact weight + plain = fused_rmsnorm(hidden, norm.weight, norm.variance_epsilon) + fused = fused_rmsnorm(hidden, norm.weight, norm.variance_epsilon, residual) + product = norm.weight.float() * normalized.float() + single = (residual.float() + product).to(torch.bfloat16) + double = residual + plain + assert torch.equal(plain, product.to(torch.bfloat16)) + assert ((fused == single) | (fused == double)).all() + assert (fused == single).float().mean().item() > 0.99 + assert not torch.equal(single, double) + + +@pytest.mark.gpu +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32]) +@pytest.mark.parametrize("width", [2048, 5280, 10560]) +@pytest.mark.parametrize("with_residual", [False, True]) +@pytest.mark.parametrize("fuse_inverse", [False, True]) +def test_fused_rmsnorm_exact_is_bitwise_identical_to_reference( + dtype, width, with_residual, fuse_inverse +): + from vllm_rlt.kernels.norm import fused_rmsnorm_exact + + torch.manual_seed(2) + norm = RMSNorm(width).to(device="cuda", dtype=dtype) + with torch.no_grad(): + norm.weight.copy_(torch.randn(width, device="cuda") * 0.5 + 1.0) + hidden = (torch.randn(37, width, device="cuda") * 3).to(dtype) + residual = (torch.randn(37, width, device="cuda") * 3).to(dtype) if with_residual else None + reference = norm.reference(hidden) + if residual is not None: + reference = residual + reference + fused = fused_rmsnorm_exact( + hidden, norm.weight, norm.variance_epsilon, residual, fuse_inverse=fuse_inverse + ) + assert fused.dtype == dtype and fused.shape == reference.shape + assert torch.equal(fused, reference) + # The row statistic is torch's own, so the exact mode stays bitwise under any width. + assert torch.equal( + fused_rmsnorm_exact( + hidden[:1], norm.weight, norm.variance_epsilon, fuse_inverse=fuse_inverse + ), + norm.reference(hidden[:1]), + ) + + +def test_fused_scope_graph_uses_reference_outside_capture(monkeypatch): + """Default scope: fused kernels only during CUDA graph capture; eager uses the reference.""" + import vllm_rlt.kernels.dispatch as dispatch + import vllm_rlt.layers.normalization as normalization + + calls = [] + monkeypatch.setattr(normalization, "FUSED_NORM", "exact") + monkeypatch.setattr( + normalization.RMSNorm, "reference", lambda self, h: calls.append("ref") or h + ) + norm = RMSNorm(8) + hidden = torch.ones(2, 8) + monkeypatch.setattr(dispatch, "FUSED_SCOPE", "graph") + assert dispatch.fused_allowed() is False # no capture on CPU + norm(hidden) + assert calls == ["ref"] + monkeypatch.setattr(dispatch, "FUSED_SCOPE", "always") + assert dispatch.fused_allowed() is True diff --git a/tests/test_ouro.py b/tests/test_ouro.py index e7f29cd3..52814237 100644 --- a/tests/test_ouro.py +++ b/tests/test_ouro.py @@ -299,3 +299,41 @@ def test_dtype_conversion_preserves_float32_rotary_frequencies(model): angles = torch.cat([angles, angles], dim=-1) torch.testing.assert_close(cos[:, 0], angles.cos().bfloat16(), atol=0, rtol=0) torch.testing.assert_close(sin[:, 0], angles.sin().bfloat16(), atol=0, rtol=0) + + +@pytest.mark.gpu +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +def test_fused_silu_mul_tracks_reference(dtype): + import torch.nn.functional as F + + from vllm_rlt.kernels.activation import fused_silu_mul, fused_silu_mul_supported + + torch.manual_seed(0) + gate = (torch.randn(64, 5632, device="cuda") * 4).to(dtype) + up = (torch.randn(64, 5632, device="cuda") * 2).to(dtype) + assert fused_silu_mul_supported(gate, up) + fused = fused_silu_mul(gate, up) + reference = F.silu(gate) * up + diff = (fused.float() - reference.float()).abs() + ulp = torch.finfo(dtype).eps * reference.float().abs().clamp_min(torch.finfo(dtype).tiny) + assert (diff <= 1.01 * ulp).all() + assert (diff == 0).float().mean().item() > 0.999 + + +@pytest.mark.gpu +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32]) +def test_fused_silu_mul_exact_matches_reference(dtype): + """Forced rounding points: the only possible difference left is the transcendental.""" + import torch.nn.functional as F + + from vllm_rlt.kernels.activation import fused_silu_mul + + torch.manual_seed(0) + gate = (torch.randn(64, 5632, device="cuda") * 4).to(dtype) + up = (torch.randn(64, 5632, device="cuda") * 2).to(dtype) + fused = fused_silu_mul(gate, up, exact=True) + reference = F.silu(gate) * up + mismatches = int((fused != reference).sum()) + # libdevice exp is the CUDA math library PyTorch links, so this is expected to be + # zero; the bound keeps the test informative if a platform's exp differs. + assert mismatches <= 2, mismatches diff --git a/tests/test_rotary_embedding.py b/tests/test_rotary_embedding.py index 9dc4f3ac..8e52365f 100644 --- a/tests/test_rotary_embedding.py +++ b/tests/test_rotary_embedding.py @@ -90,3 +90,67 @@ def test_rotary_embedding_forward_precision_and_shapes(dtype): assert sin.shape == (len(positions), 1, head_dim) assert cos.dtype == dtype assert sin.dtype == dtype + + +@pytest.mark.gpu +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +@pytest.mark.parametrize("heads,dim", [(16, 128), (55, 96)]) +def test_fused_rope_approx_rounds_the_sum_once(dtype, heads, dim): + from vllm_rlt.kernels.rope import fused_rope_, fused_rope_supported + from vllm_rlt.layers.rotary_embedding import apply_rotary_pos_emb_reference + + if (dim // 2) & (dim // 2 - 1): + pytest.skip("fused RoPE needs a power-of-two half dimension") + torch.manual_seed(0) + rope = RotaryEmbedding(head_dim=dim).to("cuda") + rows = 37 + hidden = torch.zeros(rows, 8, device="cuda", dtype=dtype) + positions = torch.randint(0, 60000, (rows,), device="cuda") + cos, sin = rope(hidden, positions) + q = (torch.randn(rows, heads, dim, device="cuda") * 2).to(dtype) + k = (torch.randn(rows, heads, dim, device="cuda") * 2).to(dtype) + assert fused_rope_supported(q, cos) + expected_q, expected_k = apply_rotary_pos_emb_reference(q, k, cos, sin) + fused_q, fused_k = fused_rope_(q.clone(), cos, sin), fused_rope_(k.clone(), cos, sin) + # The compiler is free to keep either product unrounded (FP32, or fused into an FMA + # with the other term), so each element is one of four roundings of + # q*cos + rotate_half(q)*sin: both products exact, either one rounded first, or + # both rounded (the reference). GH200 showed ~9% of elements outside the first and + # last, i.e. a mixed form. Under cancellation these differ by more than any fixed + # number of result ulps, so the contract is membership in that set. + from vllm_rlt.layers.rotary_embedding import rotate_half + + for value, fused, expected in ((q, fused_q, expected_q), (k, fused_k, expected_k)): + p1 = value.float() * cos.float() + p2 = rotate_half(value).float() * sin.float() + r1, r2 = p1.to(dtype).float(), p2.to(dtype).float() + candidates = [(p1 + p2), (r1 + p2), (p1 + r2)] + member = fused == expected + for c in candidates: + member |= fused == c.to(dtype) + assert member.float().mean().item() > 0.999 + assert (fused == expected).float().mean().item() > 0.7 + + +@pytest.mark.gpu +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32]) +@pytest.mark.parametrize("heads,dim", [(16, 128), (4, 64)]) +def test_fused_rope_exact_is_bitwise_identical_to_reference(dtype, heads, dim): + from vllm_rlt.kernels.rope import fused_rope_, fused_rope_supported + from vllm_rlt.layers.rotary_embedding import apply_rotary_pos_emb_reference + + torch.manual_seed(0) + rope = RotaryEmbedding(head_dim=dim).to("cuda") + rows = 37 + hidden = torch.zeros(rows, 8, device="cuda", dtype=dtype) + positions = torch.randint(0, 60000, (rows,), device="cuda") + cos, sin = rope(hidden, positions) + q = (torch.randn(rows, heads, dim, device="cuda") * 2).to(dtype) + k = (torch.randn(rows, heads, dim, device="cuda") * 2).to(dtype) + assert fused_rope_supported(q, cos) + expected_q, expected_k = apply_rotary_pos_emb_reference(q, k, cos, sin) + fused_q = fused_rope_(q.clone(), cos, sin, exact=True) + fused_k = fused_rope_(k.clone(), cos, sin, exact=True) + # Forced rounding points reproduce the reference's three roundings exactly. + assert torch.equal(fused_q, expected_q) + assert torch.equal(fused_k, expected_k) diff --git a/vllm_rlt/kernels/activation.py b/vllm_rlt/kernels/activation.py index 4123006c..ef498fa0 100644 --- a/vllm_rlt/kernels/activation.py +++ b/vllm_rlt/kernels/activation.py @@ -1,4 +1,17 @@ -"""MLP activation operators and their unfused correctness reference.""" +"""MLP activation operators and their unfused correctness reference. + +``silu_and_mul`` dispatches to the optional native CUDA kernel or the Torch +reference. ``fused_silu_mul`` is the fused Triton SiLU(gate) * up used while +fusion is allowed (see ``vllm_rlt.kernels.dispatch``). + +The reference ``F.silu(gate) * up`` is two kernels, each rounding to the activation dtype. +This kernel computes ``silu`` in FP32 as ``x / (1 + exp(-x))`` with libdevice ``exp`` and +IEEE division (the same formula PyTorch's CUDA kernel uses), rounds, multiplies by ``up`` +and rounds again. The approximate kernel (VLLM_RLT_FUSED_ACT=1) lets the compiler elide the +intermediate rounding; the exact kernel (VLLM_RLT_FUSED_ACT=exact) forces it through inline +PTX so the only remaining difference could be the transcendental's last fp32 ulp (libdevice +``exp`` is the CUDA math library PyTorch links). The GPU test reports the identical fraction. +""" import torch import torch.nn.functional as F @@ -37,3 +50,59 @@ def silu_and_mul_reference(gate: torch.Tensor, up: torch.Tensor) -> torch.Tensor _library.impl("silu_and_mul", silu_and_mul_reference, "CompositeExplicitAutograd") silu_and_mul = torch.ops.vllm_rlt.silu_and_mul + +try: + import triton + import triton.language as tl + from triton.language.extra import libdevice + + from vllm_rlt.kernels.rounding import mul_round, round_to +except ImportError: # pragma: no cover - CPU-only installs + triton = None + + +if triton is not None: + + @triton.jit + def _silu_mul_kernel(gate_ptr, up_ptr, out_ptr, numel, BLOCK: tl.constexpr): + offsets = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + mask = offsets < numel + dtype = out_ptr.dtype.element_ty + g = tl.load(gate_ptr + offsets, mask=mask, other=0.0).to(tl.float32) + u = tl.load(up_ptr + offsets, mask=mask, other=0.0).to(tl.float32) + act = tl.math.div_rn(g, 1.0 + libdevice.exp(-g)).to(dtype).to(tl.float32) + tl.store(out_ptr + offsets, (act * u).to(dtype), mask=mask) + + @triton.jit + def _silu_mul_exact_kernel(gate_ptr, up_ptr, out_ptr, numel, BLOCK: tl.constexpr): + offsets = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + mask = offsets < numel + dtype = out_ptr.dtype.element_ty + g = tl.load(gate_ptr + offsets, mask=mask, other=0.0).to(tl.float32) + u = tl.load(up_ptr + offsets, mask=mask, other=0.0).to(tl.float32) + # silu rounded to dtype (PyTorch's F.silu output), then the dtype product. + act = round_to(tl.math.div_rn(g, 1.0 + libdevice.exp(-g)), dtype) + tl.store(out_ptr + offsets, mul_round(act, u, dtype).to(dtype), mask=mask) + + +def fused_silu_mul_supported(gate: torch.Tensor, up: torch.Tensor) -> bool: + return ( + triton is not None + and gate.is_cuda + and gate.dtype in (torch.bfloat16, torch.float16, torch.float32) + and gate.dtype == up.dtype + and gate.shape == up.shape + and gate.is_contiguous() + and up.is_contiguous() + ) + + +def fused_silu_mul(gate: torch.Tensor, up: torch.Tensor, *, exact: bool = False) -> torch.Tensor: + out = torch.empty_like(gate) + numel = gate.numel() + if numel == 0: + return out + block = 1024 + kernel = _silu_mul_exact_kernel if exact else _silu_mul_kernel + kernel[((numel + block - 1) // block,)](gate, up, out, numel, BLOCK=block, num_warps=4) + return out diff --git a/vllm_rlt/kernels/dispatch.py b/vllm_rlt/kernels/dispatch.py new file mode 100644 index 00000000..2ec5bb7a --- /dev/null +++ b/vllm_rlt/kernels/dispatch.py @@ -0,0 +1,34 @@ +"""Where fused elementwise kernels run: inside CUDA graph capture only, or everywhere. + +A fused Triton kernel replaces several PyTorch kernels, but each Triton launch from Python +costs more host time than a PyTorch op. Inside a captured CUDA graph that cost is paid +once at capture; in eager execution (prefill, Ouro prelude/coda, sync eager decode) it is +paid every call, and on GH200's Grace CPU it outweighs the saved launches: sync eager +Ouro-1.4B decode ran at 0.40x and prefill steps took 33-46% longer with the fused set on. + +VLLM_RLT_FUSED_SCOPE=graph (default) uses the fused kernels only while a CUDA graph is +being captured; =always also uses them eagerly. The exact kernels are bitwise identical +to the eager reference, so the scope never changes numerics in exact mode. +""" + +import os + +import torch + +FUSED_SCOPE = os.environ.get("VLLM_RLT_FUSED_SCOPE", "graph") +if FUSED_SCOPE not in ("graph", "always"): + raise ValueError("VLLM_RLT_FUSED_SCOPE must be 'graph' or 'always'") + + +_CUDA = None + + +def fused_allowed() -> bool: + """Whether a fused kernel may replace the eager reference at this call site.""" + global _CUDA + if FUSED_SCOPE == "always": + return True + if _CUDA is None: + # CPU-only builds have no capture query at all; nothing is ever captured there. + _CUDA = torch.cuda.is_available() + return _CUDA and torch.cuda.is_current_stream_capturing() diff --git a/vllm_rlt/kernels/norm.py b/vllm_rlt/kernels/norm.py new file mode 100644 index 00000000..f7c3ed05 --- /dev/null +++ b/vllm_rlt/kernels/norm.py @@ -0,0 +1,209 @@ +"""Fused RMSNorm (Triton) for CUDA rows. + +The reference ``RMSNorm.forward`` is eight small kernels per call (cast, square, mean, +add, rsqrt, mul, cast, weight mul). ``fused_rmsnorm`` does the same arithmetic in one +launch: FP32 reduction, ``value * rsqrt(mean + eps)`` in FP32, round to the activation +dtype, then the weight multiply. Its FP32 reduction order differs (a block tree sum +instead of torch's reduction) and the compiler may merge roundings, so outputs can differ +in the last ulp; it is opt-in (VLLM_RLT_FUSED_NORM=approx). + +``fused_rmsnorm_exact`` (the default) keeps torch's own ``mean`` for the statistic, so the +reduction order is the reference's, and fuses the elementwise parts around it: cast+square +before, ``rsqrt(mean + eps)``/scale/cast/weight after. Three launches instead of eight (four +instead of nine with a residual, which torch adds). Every remaining operation is one +correctly rounded FP32 operation followed by the same cast, and the casts are forced +rounding points (``kernels/rounding.py``) that the compiler cannot fold, so the result is +bitwise identical to the reference. +""" + +import torch + +try: + import triton + import triton.language as tl + + from vllm_rlt.kernels.rounding import mul_round, round_to +except ImportError: # pragma: no cover - CPU-only installs + triton = None + + +if triton is not None: + + @triton.jit + def _rmsnorm_kernel( + x_ptr, + w_ptr, + res_ptr, + out_ptr, + hidden, + eps, + x_stride, + res_stride, + out_stride, + HAS_RESIDUAL: tl.constexpr, + BLOCK: tl.constexpr, + ): + row = tl.program_id(0) + offsets = tl.arange(0, BLOCK) + mask = offsets < hidden + x = tl.load(x_ptr + row * x_stride + offsets, mask=mask, other=0.0).to(tl.float32) + variance = tl.sum(x * x, axis=0) / hidden + inv = 1.0 / tl.sqrt(variance + eps) + normalized = (x * inv).to(out_ptr.dtype.element_ty) + w = tl.load(w_ptr + offsets, mask=mask, other=0.0) + out = (w.to(tl.float32) * normalized.to(tl.float32)).to(out_ptr.dtype.element_ty) + if HAS_RESIDUAL: + # residual + norm(x). The compiler keeps the weight product in fp32 across + # the cast (bf16/fp16) or contracts it into an FMA (fp32), so this is a + # single rounding of residual + w * normalized, not torch's two roundings: + # about 18% of elements differ from `residual + reference` by one ulp + # (measured on GH200). Exact mode adds the residual with torch instead. + res = tl.load(res_ptr + row * res_stride + offsets, mask=mask, other=0.0) + out = (res.to(tl.float32) + out.to(tl.float32)).to(out_ptr.dtype.element_ty) + tl.store(out_ptr + row * out_stride + offsets, out, mask=mask) + + @triton.jit + def _square_kernel(x_ptr, out_ptr, hidden, x_stride, out_stride, BLOCK: tl.constexpr): + row = tl.program_id(0) + offsets = tl.arange(0, BLOCK) + mask = offsets < hidden + x = tl.load(x_ptr + row * x_stride + offsets, mask=mask, other=0.0).to(tl.float32) + tl.store(out_ptr + row * out_stride + offsets, x * x, mask=mask) + + @triton.jit + def _rmsnorm_tail_kernel( + x_ptr, + inv_ptr, + w_ptr, + res_ptr, + out_ptr, + hidden, + eps, + x_stride, + res_stride, + out_stride, + HAS_RESIDUAL: tl.constexpr, + FUSE_INVERSE: tl.constexpr, + BLOCK: tl.constexpr, + ): + # Same rounding points as the reference: fp32 scale -> dtype, fp32 weight + # product -> dtype. Products of two half-precision values are exact in fp32, + # so each step rounds once. With FUSE_INVERSE the row statistic arrives as the + # mean of squares and rsqrt(mean + eps) is computed here: the add is one IEEE + # rounding and rsqrt lowers to the same approximate instruction torch uses. + row = tl.program_id(0) + offsets = tl.arange(0, BLOCK) + mask = offsets < hidden + dtype = out_ptr.dtype.element_ty + x = tl.load(x_ptr + row * x_stride + offsets, mask=mask, other=0.0).to(tl.float32) + inv = tl.load(inv_ptr + row) + if FUSE_INVERSE: + inv = tl.math.rsqrt(inv + eps) + # Forced rounding points (inline PTX): the compiler may otherwise keep the + # fp32 products across the casts, as it does in the approx residual epilogue. + normalized = round_to(x * inv, dtype) + w = tl.load(w_ptr + offsets, mask=mask, other=0.0) + out = mul_round(w.to(tl.float32), normalized, dtype).to(dtype) + if HAS_RESIDUAL: + res = tl.load(res_ptr + row * res_stride + offsets, mask=mask, other=0.0) + out = (res.to(tl.float32) + out.to(tl.float32)).to(dtype) + tl.store(out_ptr + row * out_stride + offsets, out, mask=mask) + + +def fused_rmsnorm_supported(hidden: torch.Tensor) -> bool: + return ( + triton is not None + and hidden.is_cuda + and hidden.ndim == 2 + and hidden.dtype in (torch.bfloat16, torch.float16, torch.float32) + and hidden.stride(1) == 1 + and hidden.shape[1] <= 16384 + ) + + +def fused_rmsnorm( + hidden: torch.Tensor, weight: torch.Tensor, eps: float, residual: torch.Tensor | None = None +) -> torch.Tensor: + """``weight * rmsnorm(hidden)`` or ``residual + weight * rmsnorm(hidden)`` in one launch.""" + rows, width = hidden.shape + out = torch.empty_like(hidden) + if rows == 0: + return out + if residual is not None and (residual.shape != hidden.shape or residual.stride(1) != 1): + raise ValueError("residual must match hidden's shape with unit inner stride") + block = triton.next_power_of_2(width) + warps = 4 if block <= 2048 else (8 if block <= 8192 else 16) + _rmsnorm_kernel[(rows,)]( + hidden, + weight, + residual if residual is not None else hidden, + out, + width, + eps, + hidden.stride(0), + residual.stride(0) if residual is not None else 0, + out.stride(0), + HAS_RESIDUAL=residual is not None, + BLOCK=block, + num_warps=warps, + ) + return out + + +# rsqrt(mean + eps) inside the tail kernel saves two launches per norm. It is +# bitwise with torch when Triton's rsqrt lowers to the same instruction as +# torch's; the GPU test checks both settings. +FUSE_INVERSE = True + + +def fused_rmsnorm_exact( + hidden: torch.Tensor, + weight: torch.Tensor, + eps: float, + residual: torch.Tensor | None = None, + *, + fuse_inverse: bool | None = None, +) -> torch.Tensor: + """Bitwise ``RMSNorm.reference`` (plus optional residual) in 3-5 (4-6) launches, not 8 (9). + + torch computes the row mean of squares exactly as the reference does; Triton + fuses the elementwise work around it (and, by default, ``rsqrt(mean + eps)``). + """ + if fuse_inverse is None: + fuse_inverse = FUSE_INVERSE + rows, width = hidden.shape + out = torch.empty_like(hidden) + if rows == 0: + return out + if residual is not None and (residual.shape != hidden.shape or residual.stride(1) != 1): + raise ValueError("residual must match hidden's shape with unit inner stride") + block = triton.next_power_of_2(width) + warps = 4 if block <= 2048 else (8 if block <= 8192 else 16) + squares = torch.empty((rows, width), dtype=torch.float32, device=hidden.device) + _square_kernel[(rows,)]( + hidden, squares, width, hidden.stride(0), squares.stride(0), BLOCK=block, num_warps=warps + ) + statistic = squares.mean(-1) + if not fuse_inverse: + statistic = torch.rsqrt(statistic + eps) + _rmsnorm_tail_kernel[(rows,)]( + hidden, + statistic, + weight, + hidden, + out, + width, + eps, + hidden.stride(0), + 0, + out.stride(0), + HAS_RESIDUAL=False, + FUSE_INVERSE=fuse_inverse, + BLOCK=block, + num_warps=warps, + ) + # An in-kernel residual epilogue rounds residual + w * normalized once (the + # compiler keeps the product unrounded across the cast, or forms an FMA in + # fp32), which differs from torch's `residual + out` in ~18% of elements by + # one ulp. torch's own add keeps this mode exact for one more launch (6 vs 9). + return out if residual is None else residual + out diff --git a/vllm_rlt/kernels/rope.py b/vllm_rlt/kernels/rope.py new file mode 100644 index 00000000..0144ae98 --- /dev/null +++ b/vllm_rlt/kernels/rope.py @@ -0,0 +1,123 @@ +"""Fused rotary embedding (Triton) for ``apply_rotary_pos_emb``. + +The reference computes ``q * cos + rotate_half(q) * sin`` as separate half-precision ops, +rounding each product and the sum. The approximate kernel (VLLM_RLT_FUSED_ROPE=1) writes +the same casts, but the compiler keeps the products in FP32 before the add (measured on an +RTX 5090 and GH200), so outputs agree with the reference to within two ulps rather than +bitwise. The exact kernel (VLLM_RLT_FUSED_ROPE=exact) forces each product's rounding +through inline PTX and is bitwise with the reference. One launch per tensor instead of +about five either way. +""" + +import torch + +try: + import triton + import triton.language as tl + + from vllm_rlt.kernels.rounding import mul_round +except ImportError: # pragma: no cover - CPU-only installs + triton = None + + +if triton is not None: + + @triton.jit + def _rope_kernel( + x_ptr, + cos_ptr, + sin_ptr, + heads, + x_stride_row, + x_stride_head, + cs_stride_row, + HALF: tl.constexpr, + ): + row = tl.program_id(0) + head = tl.program_id(1) + offsets = tl.arange(0, HALF) + base = x_ptr + row * x_stride_row + head * x_stride_head + dtype = x_ptr.dtype.element_ty + first = tl.load(base + offsets).to(tl.float32) + second = tl.load(base + HALF + offsets).to(tl.float32) + cos_a = tl.load(cos_ptr + row * cs_stride_row + offsets).to(tl.float32) + cos_b = tl.load(cos_ptr + row * cs_stride_row + HALF + offsets).to(tl.float32) + sin_a = tl.load(sin_ptr + row * cs_stride_row + offsets).to(tl.float32) + sin_b = tl.load(sin_ptr + row * cs_stride_row + HALF + offsets).to(tl.float32) + # rotate_half(x) = cat(-second, first) + prod_a = (first * cos_a).to(dtype).to(tl.float32) + rot_a = ((-second) * sin_a).to(dtype).to(tl.float32) + prod_b = (second * cos_b).to(dtype).to(tl.float32) + rot_b = (first * sin_b).to(dtype).to(tl.float32) + tl.store(base + offsets, (prod_a + rot_a).to(dtype)) + tl.store(base + HALF + offsets, (prod_b + rot_b).to(dtype)) + + @triton.jit + def _rope_exact_kernel( + x_ptr, + cos_ptr, + sin_ptr, + heads, + x_stride_row, + x_stride_head, + cs_stride_row, + HALF: tl.constexpr, + ): + row = tl.program_id(0) + head = tl.program_id(1) + offsets = tl.arange(0, HALF) + base = x_ptr + row * x_stride_row + head * x_stride_head + dtype = x_ptr.dtype.element_ty + first = tl.load(base + offsets).to(tl.float32) + second = tl.load(base + HALF + offsets).to(tl.float32) + cos_a = tl.load(cos_ptr + row * cs_stride_row + offsets).to(tl.float32) + cos_b = tl.load(cos_ptr + row * cs_stride_row + HALF + offsets).to(tl.float32) + sin_a = tl.load(sin_ptr + row * cs_stride_row + offsets).to(tl.float32) + sin_b = tl.load(sin_ptr + row * cs_stride_row + HALF + offsets).to(tl.float32) + # Each product rounded to dtype exactly as the reference's separate multiplies; + # the fp32 sum of two rounded terms is then rounded once by the store cast. + prod_a = mul_round(first, cos_a, dtype) + rot_a = mul_round(-second, sin_a, dtype) + prod_b = mul_round(second, cos_b, dtype) + rot_b = mul_round(first, sin_b, dtype) + tl.store(base + offsets, (prod_a + rot_a).to(dtype)) + tl.store(base + HALF + offsets, (prod_b + rot_b).to(dtype)) + + +def fused_rope_supported(x: torch.Tensor, cos: torch.Tensor) -> bool: + return ( + triton is not None + and x.is_cuda + and x.ndim == 3 + and x.stride(2) == 1 + and x.dtype in (torch.bfloat16, torch.float16, torch.float32) + and cos.dtype == x.dtype + and cos.ndim == 3 + and cos.shape[1] == 1 + and cos.shape[0] == x.shape[0] + and cos.stride(2) == 1 + and x.shape[2] % 2 == 0 + and (x.shape[2] // 2) & (x.shape[2] // 2 - 1) == 0 + ) + + +def fused_rope_( + x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, *, exact: bool = False +) -> torch.Tensor: + """Apply rotary embedding to ``x`` [rows, heads, dim] in place; returns ``x``.""" + rows, heads, dim = x.shape + if rows == 0: + return x + kernel = _rope_exact_kernel if exact else _rope_kernel + kernel[(rows, heads)]( + x, + cos, + sin, + heads, + x.stride(0), + x.stride(1), + cos.stride(0), + HALF=dim // 2, + num_warps=1 if dim <= 128 else 2, + ) + return x diff --git a/vllm_rlt/kernels/rounding.py b/vllm_rlt/kernels/rounding.py new file mode 100644 index 00000000..195af22b --- /dev/null +++ b/vllm_rlt/kernels/rounding.py @@ -0,0 +1,53 @@ +"""Rounding points the compiler cannot elide (Triton device helpers). + +Triton/LLVM keeps ``(a * b).to(bf16).to(f32)`` in fp32 and contracts fp32 ``a * b + c`` +into an FMA, so a kernel that mirrors a PyTorch op sequence with explicit casts is not +automatically bitwise with it (measured on GH200 for RMSNorm's residual epilogue and for +RoPE). These helpers go through inline PTX, which is opaque to those passes: ``round_to`` +rounds an fp32 value to ``dtype`` and widens it back, ``mul_rn`` is a plain IEEE fp32 +multiply that cannot be fused into a following add. +""" + +try: + import triton + import triton.language as tl +except ImportError: # pragma: no cover - CPU-only installs + triton = None + + +if triton is not None: + + @triton.jit + def round_to(x, dtype: tl.constexpr): + """fp32 -> ``dtype`` -> fp32 with the rounding guaranteed to happen.""" + if dtype == tl.bfloat16: + y = tl.inline_asm_elementwise( + "cvt.rn.bf16.f32 $0, $1;", "=h,r", [x], dtype=tl.bfloat16, is_pure=True, pack=1 + ) + return y.to(tl.float32) + elif dtype == tl.float16: + y = tl.inline_asm_elementwise( + "cvt.rn.f16.f32 $0, $1;", "=h,r", [x], dtype=tl.float16, is_pure=True, pack=1 + ) + return y.to(tl.float32) + else: + return x + + @triton.jit + def mul_rn(a, b): + """IEEE fp32 multiply that the compiler cannot contract into an FMA.""" + return tl.inline_asm_elementwise( + "mul.rn.f32 $0, $1, $2;", "=r,r,r", [a, b], dtype=tl.float32, is_pure=True, pack=1 + ) + + @triton.jit + def mul_round(a, b, dtype: tl.constexpr): + """``(a * b)`` rounded to ``dtype`` exactly once, as PyTorch rounds a ``dtype`` product. + + For bf16/fp16 operands the fp32 product is exact, so rounding it to ``dtype`` is + the one rounding PyTorch applies; for fp32 the multiply itself is the rounding. + """ + if dtype == tl.float32: + return mul_rn(a, b) + else: + return round_to(a * b, dtype) diff --git a/vllm_rlt/layers/normalization.py b/vllm_rlt/layers/normalization.py index 9d32247d..43c43d5a 100644 --- a/vllm_rlt/layers/normalization.py +++ b/vllm_rlt/layers/normalization.py @@ -1,9 +1,20 @@ # SPDX-License-Identifier: Apache-2.0 """Shared normalization layers across models.""" +import os + import torch from torch import nn +from vllm_rlt.kernels.dispatch import fused_allowed + +# Fused Triton RMSNorm on CUDA (see vllm_rlt/kernels/norm.py): +# "exact" (default): torch computes the mean, Triton the elementwise parts; bitwise identical +# to ``reference`` with 3 launches per norm instead of 8. "1"/"approx": one launch with its +# own fp32 reduction order (last-ulp differences, changes greedy outputs). "0": reference. +FUSED_NORM = os.environ.get("VLLM_RLT_FUSED_NORM", "exact") +FUSED_NORM = {"1": "approx", "approx": "approx", "exact": "exact"}.get(FUSED_NORM) + class RMSNorm(nn.Module): """Root Mean Square Layer Normalization. @@ -17,7 +28,24 @@ def __init__(self, hidden_size: int, eps: float = 1e-6) -> None: self.weight = nn.Parameter(torch.ones(hidden_size)) self.variance_epsilon = eps - def forward(self, hidden: torch.Tensor) -> torch.Tensor: + def forward(self, hidden: torch.Tensor, residual: torch.Tensor | None = None) -> torch.Tensor: + """``weight * norm(hidden)``, plus ``residual`` when given (``residual + norm``).""" + if FUSED_NORM and self.weight.dtype == hidden.dtype and fused_allowed(): + from vllm_rlt.kernels.norm import ( + fused_rmsnorm, + fused_rmsnorm_exact, + fused_rmsnorm_supported, + ) + + if fused_rmsnorm_supported(hidden) and ( + residual is None or residual.dtype == hidden.dtype + ): + kernel = fused_rmsnorm_exact if FUSED_NORM == "exact" else fused_rmsnorm + return kernel(hidden, self.weight, self.variance_epsilon, residual) + out = self.reference(hidden) + return out if residual is None else residual + out + + def reference(self, hidden: torch.Tensor) -> torch.Tensor: value = hidden.float() value = value * torch.rsqrt(value.square().mean(-1, keepdim=True) + self.variance_epsilon) return self.weight * value.to(hidden.dtype) diff --git a/vllm_rlt/layers/rotary_embedding.py b/vllm_rlt/layers/rotary_embedding.py index 297e98d6..265ccfb1 100644 --- a/vllm_rlt/layers/rotary_embedding.py +++ b/vllm_rlt/layers/rotary_embedding.py @@ -1,9 +1,19 @@ # SPDX-License-Identifier: Apache-2.0 """Shared rotary position embedding (RoPE) layers across models.""" +import os + import torch from torch import nn +from vllm_rlt.kernels.dispatch import fused_allowed + +# Fused Triton RoPE on CUDA: "exact" (default) is bitwise with the reference (forced +# rounding points), "1"/"approx" lets the compiler keep products in fp32, "0" is eager. +FUSED_ROPE = {"1": "approx", "approx": "approx", "exact": "exact"}.get( + os.environ.get("VLLM_RLT_FUSED_ROPE", "exact") +) + def rotate_half(value: torch.Tensor) -> torch.Tensor: """Rotate half the hidden dimensions of the input tensor.""" @@ -18,6 +28,17 @@ def apply_rotary_pos_emb( sin: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: """Applies Rotary Position Embedding (RoPE) to Query and Key states.""" + if FUSED_ROPE and fused_allowed(): + from vllm_rlt.kernels.rope import fused_rope_, fused_rope_supported + + if fused_rope_supported(q, cos) and fused_rope_supported(k, cos): + # q and k are fresh projection outputs here, so in-place is safe. + exact = FUSED_ROPE == "exact" + return fused_rope_(q, cos, sin, exact=exact), fused_rope_(k, cos, sin, exact=exact) + return apply_rotary_pos_emb_reference(q, k, cos, sin) + + +def apply_rotary_pos_emb_reference(q, k, cos, sin): q_embed = q * cos + rotate_half(q) * sin k_embed = k * cos + rotate_half(k) * sin return q_embed, k_embed diff --git a/vllm_rlt/models/ouro.py b/vllm_rlt/models/ouro.py index 31b37f9c..9f247e6a 100644 --- a/vllm_rlt/models/ouro.py +++ b/vllm_rlt/models/ouro.py @@ -11,6 +11,7 @@ import json import math +import os from collections.abc import Sequence from dataclasses import asdict, dataclass, fields from pathlib import Path @@ -20,6 +21,7 @@ from torch import nn from vllm_rlt.kernels import silu_and_mul +from vllm_rlt.kernels.dispatch import fused_allowed from vllm_rlt.layers import ( RMSNorm, RotaryEmbedding, @@ -150,6 +152,15 @@ def to_dict(self) -> dict[str, Any]: _rotate_half = rotate_half +# Opt-in fused SiLU(gate) * up on CUDA (see vllm_rlt.kernels.activation). +# "exact" (default): fused silu(gate) * up with forced rounding points; bitwise with the +# eager reference on GH200 over 8.3e8 random elements (bf16/fp16/fp32, three scales). +# "1"/"approx": compiler-chosen rounding, "0": eager. +FUSED_ACT = {"1": "approx", "approx": "approx", "exact": "exact"}.get( + os.environ.get("VLLM_RLT_FUSED_ACT", "exact") +) + + class OuroAttention(nn.Module): def __init__(self, config: OuroConfig, layer_idx: int) -> None: super().__init__() @@ -194,7 +205,13 @@ def __init__(self, config: OuroConfig) -> None: self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False) def forward(self, hidden: torch.Tensor) -> torch.Tensor: - return self.down_proj(silu_and_mul(self.gate_proj(hidden), self.up_proj(hidden))) + gate, up = self.gate_proj(hidden), self.up_proj(hidden) + if FUSED_ACT and fused_allowed(): + from vllm_rlt.kernels.activation import fused_silu_mul, fused_silu_mul_supported + + if fused_silu_mul_supported(gate, up): + return self.down_proj(fused_silu_mul(gate, up, exact=FUSED_ACT == "exact")) + return self.down_proj(silu_and_mul(gate, up)) class OuroDecoderLayer(nn.Module): @@ -215,9 +232,9 @@ def forward( cache: "KVCacheManager", ) -> torch.Tensor: attention = self.self_attn(self.input_layernorm(hidden), position_embeddings, batch, cache) - hidden = hidden + self.input_layernorm_2(attention) - return hidden + self.post_attention_layernorm_2( - self.mlp(self.post_attention_layernorm(hidden)) + hidden = self.input_layernorm_2(attention, residual=hidden) + return self.post_attention_layernorm_2( + self.mlp(self.post_attention_layernorm(hidden)), residual=hidden )