From e8cd8845643487dac396344f74b1750f92d2501e Mon Sep 17 00:00:00 2001 From: 0z5a Date: Sun, 25 Oct 2026 17:17:58 +0100 Subject: [PATCH] [Kernels] Bitwise-exact fused RMSNorm, RoPE and SiLU*up for CUDA graph decode At batch 1 an Ouro-1.4B token launches about 5,700 GPU kernels, most of them small elementwise and reduction kernels from RMSNorm (8 per norm, 4 norms per layer), RoPE and SiLU*up, repeated for 24 layers and 4 loops. CUDA graphs remove their launch cost but not their minimum run time. This adds fused Triton kernels for the three that are bitwise identical to the PyTorch reference: - RMSNorm: torch still computes the mean (so the reduction order is the reference's); one Triton kernel squares, one computes rsqrt(mean + eps), the scale, the cast and the weight multiply. 3 launches instead of 8; with a residual, torch adds it (4 instead of 9). - RoPE and SiLU*up: one kernel each. - Rounding points are forced with inline PTX (cvt.rn.*.f32, mul.rn.f32; kernels/rounding.py). Writing the same casts is not enough: Triton/LLVM keep an fp32 product live across .to(bf16).to(f32) and contract a*b+c into an FMA, so a fused residual epilogue computed round(r + w*n) where torch computes round(r + round(w*n)); 18.4% of 75,776 bf16 elements differed by one ulp. Fused kernels run only while a CUDA graph is being captured (VLLM_RLT_FUSED_SCOPE=graph, default): a Triton launch from Python costs more host time than the torch ops it replaces, so eager paths (prefill, sync eager decode) stay on torch. With fusion in every scope, sync eager decode on a GH200 ran at 0.40x and prefill steps took 33-46% longer. VLLM_RLT_FUSED_NORM/ROPE/ACT = exact (default) | approx | 0. The approximate one-kernel versions are opt-in: they are faster but change greedy outputs. GH200, Ouro-1.4B, 512-token prompts, 128 output tokens, async scheduling + CUDA graphs, one run per level, all arms in one job: kernels per token (batch 1): 5689 -> 2885; GPU time 13.05 -> 8.91 ms throughput vs unfused: 1.44x / 1.44x / 1.23x / 1.14x at 1 / 8 / 32 / 128 token identity vs unfused: 169/169 requests GSM8K (87-question subset, greedy, serial): 62/87 both, 0 sequences differ. SiLU*up: 0 mismatching elements in 8.3e8 (bf16/fp16/fp32, 3 gate scales). Approximate set, for comparison: 1.70x at batch 1, but -2.30 and -3.45 points on the same GSM8K subset. Signed-off-by: 0z5a Co-Authored-By: Claude Opus 5.5 --- tests/test_normalization.py | 111 +++++++++++++++ tests/test_ouro.py | 38 +++++ tests/test_rotary_embedding.py | 64 +++++++++ vllm_rlt/kernels/activation.py | 68 +++++++++ vllm_rlt/kernels/dispatch.py | 34 +++++ vllm_rlt/kernels/norm.py | 209 ++++++++++++++++++++++++++++ vllm_rlt/kernels/rope.py | 123 ++++++++++++++++ vllm_rlt/kernels/rounding.py | 53 +++++++ vllm_rlt/layers/normalization.py | 30 +++- vllm_rlt/layers/rotary_embedding.py | 21 +++ vllm_rlt/models/ouro.py | 25 +++- 11 files changed, 771 insertions(+), 5 deletions(-) create mode 100644 vllm_rlt/kernels/activation.py create mode 100644 vllm_rlt/kernels/dispatch.py create mode 100644 vllm_rlt/kernels/norm.py create mode 100644 vllm_rlt/kernels/rope.py create mode 100644 vllm_rlt/kernels/rounding.py 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 b99e3788..d96fcb21 100644 --- a/tests/test_ouro.py +++ b/tests/test_ouro.py @@ -227,3 +227,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 new file mode 100644 index 00000000..1a404290 --- /dev/null +++ b/vllm_rlt/kernels/activation.py @@ -0,0 +1,68 @@ +"""Fused SiLU(gate) * up (Triton). + +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 + +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 eae29724..570da135 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 torch.nn import functional as F +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(F.silu(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(F.silu(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 )