diff --git a/MANIFEST.in b/MANIFEST.in new file mode 100644 index 00000000..2a4ebb76 --- /dev/null +++ b/MANIFEST.in @@ -0,0 +1 @@ +recursive-include csrc *.h *.cuh *.cpp *.cu diff --git a/benchmarks/gsm8k.py b/benchmarks/gsm8k.py index b753a55d..d1ecc7f0 100644 --- a/benchmarks/gsm8k.py +++ b/benchmarks/gsm8k.py @@ -559,8 +559,6 @@ def run(args): import torch from transformers import AutoTokenizer - from benchmarks.gsm8k_backends import Generator - protocol = json.loads(Path(args.protocol).read_text()) if ( "baseline" in protocol @@ -579,6 +577,8 @@ def run(args): prompt_mode, scorer = validate_prompt_recipe(protocol) if protocol["source"]["packages"] != source()["packages"]: raise ValueError("Evaluation package versions changed after preparation") + from benchmarks.gsm8k_backends import Generator + model = Path(protocol["model"]) for name, expected in protocol["model_files"].items(): if file_digest(model / name) != expected: diff --git a/benchmarks/gsm8k_backends.py b/benchmarks/gsm8k_backends.py index 2931f9d5..15d57165 100644 --- a/benchmarks/gsm8k_backends.py +++ b/benchmarks/gsm8k_backends.py @@ -1,7 +1,10 @@ """Single-request BF16 generation with identical token inputs and stop criteria.""" +import importlib.util +from pathlib import Path + +import lm_eval import torch -from lm_eval.models.utils import postprocess_generated_text, stop_sequences_criteria from transformers import AutoModelForCausalLM, DynamicCache, LogitsProcessor from benchmarks.gsm8k import is_fixed @@ -16,6 +19,23 @@ from vllm_rlt.models import AutoModelForCausalLM as NativeModel +def _load_lm_eval_utils(): + # Importing lm_eval.models in the pinned harness eagerly imports every model + # backend, including multimodal classes removed from Transformers 5. Load + # its standalone helpers without importing the unrelated model registry; + # retain the pinned harness's exact stopping and postprocessing behavior. + path = Path(lm_eval.__file__).parent / "models" / "utils.py" + spec = importlib.util.spec_from_file_location("_gsm8k_lm_eval_utils", path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +_lm_eval_utils = _load_lm_eval_utils() +postprocess_generated_text = _lm_eval_utils.postprocess_generated_text +stop_sequences_criteria = _lm_eval_utils.stop_sequences_criteria + + def check_logits(logits): if not torch.isfinite(logits).all().item(): raise FloatingPointError("Non-finite generation logits") diff --git a/benchmarks/kernels/__init__.py b/benchmarks/kernels/__init__.py new file mode 100644 index 00000000..63a79e82 --- /dev/null +++ b/benchmarks/kernels/__init__.py @@ -0,0 +1 @@ +"""Reproducible standalone kernel comparisons.""" diff --git a/benchmarks/kernels/silu_and_mul.py b/benchmarks/kernels/silu_and_mul.py new file mode 100644 index 00000000..eaf79cfb --- /dev/null +++ b/benchmarks/kernels/silu_and_mul.py @@ -0,0 +1,216 @@ +"""SiLU-and-Mul accuracy and latency against eager Torch and vLLM. + +CUDA_VISIBLE_DEVICES= python -m benchmarks.kernels.silu_and_mul \ + --output artifacts/silu-and-mul/results.json +""" + +import argparse +import hashlib +import importlib +import importlib.metadata +import json +import random +import statistics +import subprocess +import time +from pathlib import Path + +import torch +import torch.nn.functional as F + +from vllm_rlt.kernels._extension import load_extension +from vllm_rlt.kernels.activation import silu_and_mul + + +def errors(actual, expected): + finite = torch.isfinite(actual) & torch.isfinite(expected) + delta = (actual[finite].double() - expected[finite].double()).abs() + integer = torch.int32 if actual.dtype == torch.float32 else torch.int16 + sign = 2 ** (31 if integer == torch.int32 else 15) + + def ordered(x): + bits = x[finite].view(integer).long() + return torch.where(bits < 0, -sign - bits, bits) + + nonfinite_match = (actual == expected) | (torch.isnan(actual) & torch.isnan(expected)) + return { + "max_abs": delta.max().item() if delta.numel() else 0, + "rmse": delta.square().mean().sqrt().item() if delta.numel() else 0, + "max_ulp": (ordered(actual) - ordered(expected)).abs().max().item() if delta.numel() else 0, + "unequal_fraction": ((actual != expected) & ~(torch.isnan(actual) & torch.isnan(expected))) + .float() + .mean() + .item(), + "nonfinite_mismatches": ((~finite) & ~nonfinite_match).sum().item(), + } + + +def make_functions(gate, up, vllm_op): + # vLLM receives prepacked input and a preallocated output outside timing. + packed = torch.cat((gate, up), dim=-1) + output = torch.empty_like(gate) + + def vllm_prepacked(): + vllm_op(output, packed) + return output + + return { + "torch_eager": lambda: F.silu(gate) * up, + "rlt": lambda: silu_and_mul(gate, up), + "vllm_prepacked": vllm_prepacked, + } + + +def measure(functions, repeats, iterations): + graphs = {} + unroll = 32 + for name, function in functions.items(): + for _ in range(10): + function() + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + for _ in range(unroll): + function() + graphs[name] = graph + + result = {name: {"graph_us": [], "eager_us": []} for name in functions} + rng = random.Random(42) + for _ in range(repeats): + names = list(functions) + rng.shuffle(names) + for name in names: + graph = graphs[name] + graph.replay() + start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(iterations): + graph.replay() + end.record() + end.synchronize() + result[name]["graph_us"].append(start.elapsed_time(end) * 1000 / (iterations * unroll)) + + torch.cuda.synchronize() + before = time.perf_counter() + for _ in range(iterations): + functions[name]() + torch.cuda.synchronize() + result[name]["eager_us"].append((time.perf_counter() - before) * 1e6 / iterations) + for measurements in result.values(): + for mode in ("graph", "eager"): + measurements[f"{mode}_median_us"] = statistics.median(measurements[f"{mode}_us"]) + return result + + +def profile(functions, output): + output.mkdir(parents=True, exist_ok=True) + for full in (False, True): + mode = "full" if full else "ops-only" + for function in functions.values(): + function() + torch.cuda.synchronize() + with torch.profiler.profile( + activities=[torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA], + record_shapes=full, + profile_memory=full, + with_stack=full, + experimental_config=torch.profiler._ExperimentalConfig(verbose=full), + ) as prof: + for name, function in functions.items(): + with torch.profiler.record_function(name): + function() + torch.cuda.synchronize() + prof.export_chrome_trace(str(output / f"{mode}.json")) + (output / f"{mode}.txt").write_text(prof.key_averages().table(row_limit=40)) + if full: + prof.export_stacks(str(output / "stacks.txt"), "self_cuda_time_total") + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--vllm-library", type=Path) + parser.add_argument("--vllm-namespace", default="_C") + parser.add_argument("--vllm-source", type=Path) + parser.add_argument("--tokens", default="1,16,128,1024,4096") + parser.add_argument("--widths", default="5632,5633") + parser.add_argument("--dtypes", default="bfloat16,float16,float32") + parser.add_argument("--repeats", type=int, default=3) + parser.add_argument("--iterations", type=int, default=50) + parser.add_argument("--profile-dir", type=Path) + args = parser.parse_args() + load_extension() # Do not benchmark the Torch fallback as a native kernel. + torch.set_num_threads(1) + torch.manual_seed(42) + if args.vllm_library: + torch.ops.load_library(str(args.vllm_library.resolve())) + baseline = {"library": str(args.vllm_library.resolve())} + else: + extension = importlib.import_module("vllm._C_stable_libtorch") + baseline = {"version": importlib.metadata.version("vllm"), "library": extension.__file__} + baseline["sha256"] = hashlib.sha256(Path(baseline["library"]).read_bytes()).hexdigest() + if args.vllm_source: + baseline["source_commit"] = subprocess.check_output( + ["git", "-C", str(args.vllm_source), "rev-parse", "HEAD"], text=True + ).strip() + vllm_op = getattr(getattr(torch.ops, args.vllm_namespace), "silu_and_mul") + patch = subprocess.check_output(["git", "diff", "--binary"]) + report = { + "torch": torch.__version__, + "cuda_runtime": torch.version.cuda, + "gpu": torch.cuda.get_device_name(), + "capability": torch.cuda.get_device_capability(), + "git_head": subprocess.check_output(["git", "rev-parse", "HEAD"], text=True).strip(), + "diff_sha256": hashlib.sha256(patch).hexdigest(), + "vllm": baseline, + "seed": 42, + "gate_std": 3, + "repeats": args.repeats, + "iterations": args.iterations, + "graph_unroll": 32, + "rows": [], + } + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.with_suffix(".patch").write_bytes(patch) + with torch.inference_mode(): + for dtype_name in args.dtypes.split(","): + dtype = getattr(torch, dtype_name) + for tokens in map(int, args.tokens.split(",")): + for width in map(int, args.widths.split(",")): + gate = torch.randn(tokens, width, device="cuda", dtype=dtype) * 3 + up = torch.randn_like(gate) + functions = make_functions(gate, up, vllm_op) + expected = F.silu(gate) * up + g64 = gate.double() + activated = (g64 / (1 + torch.exp(-g64))).to(dtype) + oracle = (activated.double() * up.double()).to(dtype) + accuracy = {} + for name, function in functions.items(): + actual = function() + torch.testing.assert_close(actual, expected) + accuracy[name] = { + "vs_torch": errors(actual, expected), + "vs_staged_fp64": errors(actual, oracle), + } + timing = measure(functions, args.repeats, args.iterations) + row = { + "dtype": dtype_name, + "tokens": tokens, + "width": width, + "accuracy": accuracy, + "timing": timing, + } + report["rows"].append(row) + args.output.write_text(json.dumps(report, indent=2) + "\n") + medians = { + name: round(value["graph_median_us"], 3) for name, value in timing.items() + } + print(dtype_name, tokens, width, medians, flush=True) + if args.profile_dir: + gate = torch.randn(128, 5633, device="cuda", dtype=torch.bfloat16) + up = torch.randn_like(gate) + profile(make_functions(gate, up, vllm_op), args.profile_dir) + + +if __name__ == "__main__": + main() diff --git a/csrc/activation/silu_and_mul.cu b/csrc/activation/silu_and_mul.cu new file mode 100644 index 00000000..faadcda1 --- /dev/null +++ b/csrc/activation/silu_and_mul.cu @@ -0,0 +1,169 @@ +// SPDX-License-Identifier: Apache-2.0 +#include "../ops.h" + +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include + +namespace vllm_rlt { +namespace { + +__device__ __forceinline__ float silu(float x) { + return x / (1.0f + expf(-x)); +} + +template +struct Math; + +template <> +struct Math<__half> { + using Pair = __half2; + __device__ static float to_float(__half x) { return __half2float(x); } + __device__ static __half from_float(float x) { return __float2half_rn(x); } + __device__ static __half mul(__half a, __half b) { return __hmul(a, b); } + __device__ static Pair compute(Pair gate, Pair up) { + float2 g = __half22float2(gate); + Pair activated = __floats2half2_rn(silu(g.x), silu(g.y)); + return __hmul2(activated, up); + } +}; + +template <> +struct Math<__nv_bfloat16> { + using Pair = __nv_bfloat162; + __device__ static float to_float(__nv_bfloat16 x) { return __bfloat162float(x); } + __device__ static __nv_bfloat16 from_float(float x) { return __float2bfloat16_rn(x); } + __device__ static __nv_bfloat16 mul(__nv_bfloat16 a, __nv_bfloat16 b) { + return __hmul(a, b); + } + __device__ static Pair compute(Pair gate, Pair up) { + float2 g = __bfloat1622float2(gate); + Pair activated = __floats2bfloat162_rn(silu(g.x), silu(g.y)); + return __hmul2(activated, up); + } +}; + +template <> +struct Math { + using Pair = float2; + __device__ static float to_float(float x) { return x; } + __device__ static float from_float(float x) { return x; } + __device__ static float mul(float a, float b) { return __fmul_rn(a, b); } + __device__ static Pair compute(Pair gate, Pair up) { + Pair activated = make_float2(silu(gate.x), silu(gate.y)); + // Packed FP32 products increased register allocation and reduced occupancy + // in the SM103 benchmarks, unlike the FP16/BF16 packed products above. + return make_float2(mul(activated.x, up.x), mul(activated.y, up.y)); + } +}; + +template +struct alignas(16) Vector { + typename Math::Pair pairs[16 / sizeof(typename Math::Pair)]; +}; + +template +__device__ __forceinline__ T compute_scalar(T gate, T up) { + // Match F.silu(gate) * up: round the activation before multiplying. + T activated = Math::from_float(silu(Math::to_float(gate))); + return Math::mul(activated, up); +} + +template +__global__ void silu_and_mul_kernel(const T* gate, const T* up, T* out, int64_t count) { + const int64_t thread = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const int64_t stride = static_cast(gridDim.x) * blockDim.x; + int64_t tail_start = 0; + if constexpr (vectorized) { + constexpr int elements = 16 / sizeof(T); + constexpr int pairs = 16 / sizeof(typename Math::Pair); + const int64_t vectors = count / elements; + for (int64_t i = thread; i < vectors; i += stride) { + Vector g, u, result; + *reinterpret_cast(&g) = *reinterpret_cast(gate + i * elements); + *reinterpret_cast(&u) = *reinterpret_cast(up + i * elements); +#pragma unroll + for (int j = 0; j < pairs; ++j) { + result.pairs[j] = Math::compute(g.pairs[j], u.pairs[j]); + } + *reinterpret_cast(out + i * elements) = *reinterpret_cast(&result); + } + tail_start = vectors * elements; + } + // Only valid scalar addresses are touched, including for count < one vector. + for (int64_t i = tail_start + thread; i < count; i += stride) { + out[i] = compute_scalar(gate[i], up[i]); + } +} + +template +void launch(const at::Tensor& gate, const at::Tensor& up, at::Tensor& out) { + const T* g = reinterpret_cast(gate.const_data_ptr()); + const T* u = reinterpret_cast(up.const_data_ptr()); + T* result = reinterpret_cast(out.mutable_data_ptr()); + // A contiguous storage-offset view is not necessarily 16-byte aligned. + const bool vectorized = ((reinterpret_cast(g) | + reinterpret_cast(u) | + reinterpret_cast(result)) & 15) == 0; + const int64_t count = gate.numel(); + const int64_t work = vectorized ? std::max(count / (16 / sizeof(T)), 1) : count; + constexpr int threads = 256; + int max_blocks = 4096; + if constexpr (std::is_same_v) { + // More blocks improve large FP32 streaming workloads. Require enough work + // to fill the enlarged grid, preserving smaller shapes' launch geometry. + constexpr int streaming_blocks = 16384; + if (work >= static_cast(streaming_blocks) * threads) { + max_blocks = streaming_blocks; + } + } + const int blocks = static_cast(std::min((work + threads - 1) / threads, max_blocks)); + auto stream = c10::cuda::getCurrentCUDAStream(gate.get_device()); + if (vectorized) { + silu_and_mul_kernel<<>>(g, u, result, count); + } else { + silu_and_mul_kernel<<>>(g, u, result, count); + } + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +} // namespace + +at::Tensor silu_and_mul_cuda(const at::Tensor& gate, const at::Tensor& up) { + TORCH_CHECK(gate.is_cuda() && up.is_cuda(), "silu_and_mul expects CUDA inputs"); + TORCH_CHECK(gate.device() == up.device(), "silu_and_mul expects inputs on the same device"); + TORCH_CHECK(gate.sizes() == up.sizes(), "silu_and_mul expects inputs with the same shape"); + TORCH_CHECK(gate.scalar_type() == up.scalar_type(), + "silu_and_mul expects inputs with the same dtype"); + TORCH_CHECK(gate.scalar_type() == at::kFloat || gate.scalar_type() == at::kHalf || + gate.scalar_type() == at::kBFloat16, + "silu_and_mul supports float32, float16 and bfloat16"); + + TORCH_CHECK(gate.is_contiguous() && up.is_contiguous(), + "silu_and_mul expects contiguous inputs"); + TORCH_CHECK(!at::GradMode::is_enabled() || (!gate.requires_grad() && !up.requires_grad()), + "silu_and_mul is inference-only; use torch.no_grad() or silu_and_mul_reference"); + + const c10::cuda::CUDAGuard guard(gate.device()); + auto out = at::empty(gate.sizes(), gate.options()); + if (gate.numel() == 0) { + return out; + } + switch (gate.scalar_type()) { + case at::kHalf: launch<__half>(gate, up, out); break; + case at::kBFloat16: launch<__nv_bfloat16>(gate, up, out); break; + case at::kFloat: launch(gate, up, out); break; + default: TORCH_CHECK(false, "unsupported dtype"); + } + return out; +} + +} // namespace vllm_rlt diff --git a/csrc/ops.h b/csrc/ops.h new file mode 100644 index 00000000..e00c054e --- /dev/null +++ b/csrc/ops.h @@ -0,0 +1,10 @@ +// SPDX-License-Identifier: Apache-2.0 +#pragma once + +#include + +namespace vllm_rlt { + +at::Tensor silu_and_mul_cuda(const at::Tensor& gate, const at::Tensor& up); + +} // namespace vllm_rlt diff --git a/csrc/torch_bindings.cpp b/csrc/torch_bindings.cpp new file mode 100644 index 00000000..afd5cf41 --- /dev/null +++ b/csrc/torch_bindings.cpp @@ -0,0 +1,15 @@ +// SPDX-License-Identifier: Apache-2.0 +#include +#include + +#include "ops.h" + +TORCH_LIBRARY(vllm_rlt, ops) { + ops.def("silu_and_mul(Tensor gate, Tensor up) -> Tensor"); +} + +TORCH_LIBRARY_IMPL(vllm_rlt, CUDA, ops) { + ops.impl("silu_and_mul", &vllm_rlt::silu_and_mul_cuda); +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) {} diff --git a/docs/README.md b/docs/README.md index d83666cb..421e0f08 100644 --- a/docs/README.md +++ b/docs/README.md @@ -7,6 +7,7 @@ an environment to receiving the first generated response. The | Guide | Topics | | --- | --- | | [Developer Must-Read](developer-guide.md) | Requirement understanding, correctness, performance analysis, refactoring rules, and completion criteria | +| [Native kernel development](kernels.md) | Optional CUDA extension, operator layout, and SiLU-and-Mul reference contract | | [First-run user guide](launching.md) | Environment setup, model download, server startup, first request, CLI, Python API, and troubleshooting | | [Unified architecture](unified_execution.md) | Stage process boundaries, component responsibilities, shared resources and loop-level execution | | [Model and KV semantics](design.md) | PRELUDE/RECURRENT/CODA, request progress and depth-specific history | diff --git a/docs/kernels.md b/docs/kernels.md new file mode 100644 index 00000000..a1fda4e1 --- /dev/null +++ b/docs/kernels.md @@ -0,0 +1,147 @@ +# Native kernel development + +Native kernels are optional and live alongside the existing Triton kernels. +Importing `vllm_rlt.kernels` loads the built library once and registers the Torch +fallbacks. It neither compiles code nor creates a CUDA context. + +## Layout + +| Location | Responsibility | +| --- | --- | +| `vllm_rlt/kernels/activation.py` | Python activation entry and Torch reference | +| `vllm_rlt/kernels/_extension.py` | One-time loading and missing-library detection | +| `csrc/ops.h` | Native operator declarations | +| `csrc/torch_bindings.cpp` | Operator schemas and CUDA dispatch registration | +| `csrc/activation/silu_and_mul.cu` | SiLU-and-Mul CUDA validation, launch and computation | +| `setup.py` | Optional build of `vllm_rlt._C` | +| `tests/test_activation.py` | Reference, CUDA numerics, layouts, tails and graph replay | +| `benchmarks/kernels/silu_and_mul.py` | Matched Torch/native/vLLM accuracy and latency comparison | + +Add future kernels to the same extension and group sources by operation family. +No backend registry or base class is required. Put future kernel performance +benchmarks under `benchmarks/kernels/`. + +## SiLU-and-Mul contract and current status + +`silu_and_mul_reference(gate, up)` computes `F.silu(gate) * up`. The inputs have +identical shape, dtype and device; broadcasting is rejected. Supported dtypes +are FP32, FP16 and BF16. The output owns new storage and preserves input shape, +dtype and device. Neither input is modified. The Torch reference preserves +the intermediate SiLU rounding of the existing MLP expression. + +Separate inputs avoid concatenating Ouro and Nanbeige's separate projection +outputs. The reference also accepts Huginn's noncontiguous split-projection +views. The native entry requires both inputs to be contiguous; it neither pads +nor copies incompatible inputs. + +`silu_and_mul(gate, up)` uses PyTorch dispatch: CUDA inputs use the native kernel +when the extension is installed; CPU inputs and installations without the +extension use the Torch reference. Import it from `vllm_rlt.kernels`. The model +call site does not check devices or extension availability. The native kernel +flattens contiguous inputs, using 16-byte +loads/stores for the aligned vector body and scalar accesses for the remaining +elements in the same kernel. Contiguous views with misaligned storage offsets +use scalar accesses. Empty inputs return an empty output without launching. + +SiLU uses FP32 arithmetic and rounds to the input dtype before multiplying by +`up`, preserving the reference's intermediate rounding. FP16/BF16 vector paths +use `__hmul2`; FP32 uses separate round-to-nearest products. On SM103, forcing +`__fmul2_rn` increased the 16-byte vector kernel's registers/thread from 28 to +34 and reduced resident warps, hurting memory throughput. The scalar path uses +`__hmul` for FP16/BF16 and `__fmul_rn` for FP32, retaining the SiLU rounding +boundary without explicitly converting the rounded activation back to FP32. +The device guard follows the input device, and launches use that device's current +stream. + +Blocks have 256 threads. The grid is capped at 4096 blocks, except for FP32 +workloads with at least `16384 * 256` work items: these use up to 16384 blocks. +Each aligned work item processes four FP32 elements; an unaligned work item +processes one scalar. The threshold requires enough work to fill the larger +grid. This keeps small and medium shapes' launch geometry unchanged. Grid sizes, +16/32-byte transfers and scalar/packed FP32 products were compared on B300; +the selected configuration is not a universal hardware optimum. + +Ouro's inference MLP calls this registered operator directly. Availability is +resolved once during module initialization; extension ABI or dependency errors +propagate rather than silently selecting another path. The native CUDA dispatch +reuses its C++ input checks; a missing extension does not mask native execution +errors when an installed kernel is selected. +Nanbeige and Huginn do not use this operator yet. + +Autograd is unsupported: gradient-enabled inputs are rejected unless gradients +are disabled, consistent with Ouro's inference-only model. `torch.compile` has +no fake kernel registration. Tests cover the operator and Ouro MLP, including +CUDA Graph capture/replay. + +## Build and reference checks + +The extension builds automatically when the build environment has CUDA-enabled +PyTorch, a CUDA toolkit with `nvcc`, and a visible GPU. PyTorch locates the toolkit +and detects the visible GPUs' target architectures. No project build switch or +manual architecture setting is needed. A compatible C++ compiler is also required; +a driver alone is insufficient. Run these commands with the Python interpreter +used for inference, keeping your allocated GPU visible: + +```bash +python -m pip install "torch>=2.5" setuptools wheel ninja +python -m pip install --no-build-isolation -e . +python -m pytest tests/test_activation.py -q +CUDA_VISIBLE_DEVICES= python -m pytest \ + tests/test_activation.py --run-gpu -q +``` + +`pip` invokes setuptools, whose PyTorch `CUDAExtension`/`BuildExtension` compiles +the C++ bindings and CUDA sources with the host compiler and `nvcc`, then installs +`vllm_rlt._C`. `--no-build-isolation` uses the environment's existing PyTorch +headers and libraries; it avoids a separate build environment with a potentially +different Torch version. An isolated build without PyTorch cannot build the +extension and produces a Python-only package. Missing CUDA PyTorch, toolkit, +`nvcc`, or a GPU target also produces a Python-only package with a skip message. +Once native compilation is selected, compilation errors fail the installation. +No runtime JIT compilation or CMake is required. + +For a wheel targeting the visible GPUs: + +```bash +python -m pip wheel --no-build-isolation --no-deps . -w dist +python -m pip install dist/vllm_rlt-*.whl +``` + +Without a visible GPU, the deployment architecture cannot be detected. For +cross-compilation, PyTorch's standard `TORCH_CUDA_ARCH_LIST` override remains +available, for example `TORCH_CUDA_ARCH_LIST="8.0;10.3"` for SM80 and SM103. +Choose targets supported by the deployment hardware and toolkit. An installed +binary without compatible device code raises an error; there is no fallback +after a native launch failure. GPU kernel tests and the benchmark require the +native extension and fail if it is absent. + +Restart Python after rebuilding the library. A successful build or passing CPU +reference tests establishes no CUDA compute correctness or speedup. + +## Accuracy and performance comparison + +With vLLM installed in the same environment: + +```bash +CUDA_VISIBLE_DEVICES= .venv/bin/python -m benchmarks.kernels.silu_and_mul \ + --output artifacts/silu-and-mul/results.json \ + --profile-dir artifacts/silu-and-mul/profiles +``` + +The benchmark checks matched random inputs against eager Torch and a staged +FP64 oracle that rounds SiLU before multiplication. It retains maximum absolute +and ULP errors, RMSE, mismatch fractions, and nonfinite mismatches. The default +matrix covers FP32/FP16/BF16, 1/16/128/1024/4096 rows, and widths 5632/5633. + +Latency is measured without profiling, with three repetitions and shuffled +backend order. CUDA Graph timing amortizes replay over 32 calls; eager timing +includes Python, allocation, launch and device completion. `vllm_prepacked` +receives its packed input and preallocated output before timing. These operator +comparisons do not establish full-model throughput or task quality. + +For a source-built standalone vLLM baseline, pass `--vllm-library ` and +`--vllm-namespace `, plus `--vllm-source ` to record its +commit. The report retains library hashes, the local diff and raw measurements. +Optional ops-only/full profiles are captured separately after timing; full +capture enables shapes, memory and stacks. Large traces stay under ignored +`artifacts/`. diff --git a/docs/launching.md b/docs/launching.md index 6fbc0de5..ca8e9350 100644 --- a/docs/launching.md +++ b/docs/launching.md @@ -34,8 +34,19 @@ cd vllm-rlt ## 4. Install the project +Select your allocated GPU before installation so automatic compilation targets +the GPU you will use: + ```bash -uv pip install -e . +export CUDA_VISIBLE_DEVICES=0 +``` + +Use your allocated GPU ID in place of `0`. If your job scheduler already sets +`CUDA_VISIBLE_DEVICES`, keep its value and skip the export. + +```bash +uv pip install "torch>=2.5" setuptools wheel ninja +uv pip install --no-build-isolation -e . python -m vllm_rlt.entrypoints.serve --help ``` @@ -52,15 +63,35 @@ A100, and H100. There is no CUDA version to select based on the GPU model; your NVIDIA driver must support the CUDA runtime supplied with PyTorch. FlashAttention and NIXL are optional and are not needed for this guide. -Select your GPU and confirm that PyTorch can use it: +Confirm that PyTorch can use your GPU: ```bash -export CUDA_VISIBLE_DEVICES=0 python -c "import torch; print('PyTorch:', torch.__version__, 'CUDA:', torch.version.cuda); assert torch.cuda.is_available(), 'CUDA unavailable'; assert torch.cuda.is_bf16_supported(), 'BF16 required'" ``` -Use your allocated GPU ID in place of `0`. If your job scheduler already sets -`CUDA_VISIBLE_DEVICES`, keep its value and skip the export. +### Automatic native CUDA build + +Installation automatically builds Ouro's fused SiLU-and-Mul implementation when +it finds CUDA-enabled PyTorch, a CUDA toolkit with `nvcc`, and a visible GPU. +PyTorch finds the toolkit and detects the visible GPUs' target architectures; +no build switch or architecture setting is required. The NVIDIA driver and +PyTorch CUDA runtime alone do not provide the compiler; a compatible toolkit +and C++ compiler must also be installed. + +If the requirements are missing, installation reports why the extension was +skipped and retains the Torch implementation. To add the extension later, install +the toolkit, keep your allocated GPU visible, and rerun: + +```bash +uv pip install --no-build-isolation -e . +``` + +This compiles and installs `vllm_rlt._C` using the environment's existing PyTorch. +Ouro selects the native kernel automatically for CUDA inputs; no server flag is +needed. Without the extension, it retains the Torch implementation. After changing +CUDA or C++ sources, rerun the build command and restart the server or Python +process. See [native kernel development](kernels.md#build-and-reference-checks) +for wheel builds, target GPU architectures and GPU validation commands. ## 5. Download the model and tokenizer diff --git a/setup.py b/setup.py new file mode 100644 index 00000000..b83b309d --- /dev/null +++ b/setup.py @@ -0,0 +1,48 @@ +"""Build native kernels automatically when a CUDA build environment is available.""" + +import os +from pathlib import Path + +from setuptools import setup + + +def native_build_options(): + try: + import torch + except ModuleNotFoundError as exc: + if exc.name != "torch": + raise + print("CUDA extension skipped: install PyTorch and use --no-build-isolation to build it.") + return {} + + from torch.utils.cpp_extension import CUDA_HOME, BuildExtension, CUDAExtension + + if torch.version.cuda is None or CUDA_HOME is None: + print("CUDA extension skipped: CUDA-enabled PyTorch and the CUDA toolkit are required.") + return {} + if not (Path(CUDA_HOME) / "bin" / "nvcc").is_file(): + print("CUDA extension skipped: nvcc was not found in the CUDA toolkit.") + return {} + targets = os.environ.get("TORCH_CUDA_ARCH_LIST", "").strip() + if targets in ("", "native") and not torch.cuda.is_available(): + print("CUDA extension skipped: no visible GPU to detect a target architecture.") + return {} + + # CUDA_HOME and visible-device architecture detection are owned by PyTorch. + # Its standard architecture override also supports cross-compiling without a GPU. + return { + "ext_modules": [ + CUDAExtension( + "vllm_rlt._C", + sources=[ + "csrc/torch_bindings.cpp", + "csrc/activation/silu_and_mul.cu", + ], + extra_compile_args={"cxx": ["-O3"], "nvcc": ["-O3"]}, + ) + ], + "cmdclass": {"build_ext": BuildExtension}, + } + + +setup(**native_build_options()) diff --git a/tests/test_activation.py b/tests/test_activation.py new file mode 100644 index 00000000..dbedb464 --- /dev/null +++ b/tests/test_activation.py @@ -0,0 +1,280 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Numerics, tails, layouts and stream contracts for SiLU-and-Mul.""" + +import math +import subprocess +import sys +import textwrap + +import pytest +import torch + +from vllm_rlt.kernels.activation import silu_and_mul, silu_and_mul_reference + + +def test_missing_native_extension_allows_optional_selection_but_not_explicit_load(monkeypatch): + from vllm_rlt.kernels import _extension + + def missing(name): + raise ModuleNotFoundError(name=name) + + _extension.try_load_extension.cache_clear() + monkeypatch.setattr(_extension, "import_module", missing) + try: + assert not _extension.try_load_extension() + with pytest.raises(RuntimeError, match="Native CUDA kernels are not built"): + _extension.load_extension() + finally: + _extension.try_load_extension.cache_clear() + + +def test_broken_native_extension_dependency_is_not_hidden(monkeypatch): + from vllm_rlt.kernels import _extension + + def broken(name): + raise ModuleNotFoundError(name="missing_internal_dependency") + + _extension.try_load_extension.cache_clear() + monkeypatch.setattr(_extension, "import_module", broken) + try: + with pytest.raises(ModuleNotFoundError) as exc: + _extension.try_load_extension() + assert exc.value.name == "missing_internal_dependency" + finally: + _extension.try_load_extension.cache_clear() + + +def test_silu_and_mul_known_values(): + gate = torch.tensor([0.0, 1.0, -1.0]) + up = torch.tensor([2.0, 3.0, 4.0]) + expected = torch.tensor([0.0, 2.1931758, -1.0757657]) + torch.testing.assert_close(silu_and_mul_reference(gate, up), expected) + + +@pytest.mark.parametrize("shape", [(7,), (2, 7), (2, 3, 9), (0, 7), (2, 0)]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +def test_silu_and_mul_reference_contract(shape, dtype): + generator = torch.Generator().manual_seed(42) + gate = torch.randn(shape, dtype=dtype, generator=generator) + up = torch.randn(shape, dtype=dtype, generator=generator) + gate_before, up_before = gate.clone(), up.clone() + + out = silu_and_mul_reference(gate, up) + # Preserve the intermediate SiLU rounding used by the existing MLPs. + activated = (gate.double() * torch.sigmoid(gate.double())).to(dtype) + expected = (activated.double() * up.double()).to(dtype) + torch.testing.assert_close(out, expected) + assert out.shape == gate.shape + assert out.dtype == gate.dtype + assert out.device == gate.device + torch.testing.assert_close(gate, gate_before, rtol=0, atol=0) + torch.testing.assert_close(up, up_before, rtol=0, atol=0) + if out.numel(): + assert out.data_ptr() not in (gate.data_ptr(), up.data_ptr()) + + +def test_silu_and_mul_reference_accepts_split_projection_views(): + packed = torch.arange(20, dtype=torch.float32).reshape(2, 10) / 10 + gate, up = packed.chunk(2, dim=-1) + assert not gate.is_contiguous() + out = silu_and_mul_reference(gate, up) + expected = silu_and_mul_reference(gate.contiguous(), up.contiguous()) + torch.testing.assert_close(out, expected) + + +@pytest.mark.parametrize("operator", [silu_and_mul_reference, silu_and_mul]) +def test_silu_and_mul_rejects_broadcasting(operator): + with pytest.raises(ValueError, match="same shape"): + operator(torch.ones(2, 7), torch.ones(1, 7)) + + +@pytest.mark.parametrize("operator", [silu_and_mul_reference, silu_and_mul]) +def test_silu_and_mul_rejects_mixed_dtypes(operator): + with pytest.raises(TypeError, match="same dtype"): + operator(torch.ones(7), torch.ones(7, dtype=torch.float16)) + + +@pytest.mark.parametrize("operator", [silu_and_mul_reference, silu_and_mul]) +@pytest.mark.parametrize("dtype", [torch.int32, torch.float64]) +def test_silu_and_mul_rejects_unsupported_dtypes(operator, dtype): + with pytest.raises(TypeError, match="supports float32"): + operator(torch.ones(7, dtype=dtype), torch.ones(7, dtype=dtype)) + + +def test_silu_and_mul_dispatches_cpu_to_torch_reference(): + gate, up = torch.randn(3, 7), torch.randn(3, 7) + torch.testing.assert_close(silu_and_mul(gate, up), silu_and_mul_reference(gate, up)) + + +@pytest.mark.parametrize("missing_extension", [False, True]) +def test_kernel_initialization_and_cpu_fallback_do_not_initialize_cuda(missing_extension): + script = textwrap.dedent(f""" + import importlib + import torch + original_import = importlib.import_module + attempts = [] + def import_module(name, *args, **kwargs): + if name == 'vllm_rlt._C': + attempts.append(name) + if {missing_extension}: + raise ModuleNotFoundError(name=name) + return original_import(name, *args, **kwargs) + importlib.import_module = import_module + from vllm_rlt.kernels import silu_and_mul + for _ in range(3): + gate, up = torch.randn(3, 7), torch.randn(3, 7) + torch.testing.assert_close(silu_and_mul(gate, up), torch.nn.functional.silu(gate)*up) + assert len(attempts) == 1, attempts + assert not torch.cuda.is_initialized() + """) + subprocess.run([sys.executable, "-c", script], check=True, capture_output=True, timeout=60) + + +@pytest.fixture(autouse=True) +def require_native_extension_for_gpu_tests(request): + if request.node.get_closest_marker("gpu"): + from vllm_rlt.kernels._extension import load_extension + + load_extension() + + +@pytest.mark.gpu +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +@pytest.mark.parametrize( + "shape", [(), (0, 7), (2, 0), (1,), (7,), (8,), (9,), (3, 7), (2, 3, 9), (3, 5633), (128, 5632)] +) +@pytest.mark.parametrize("offsets", [(0, 0), (1, 0), (0, 1)]) +def test_silu_and_mul_cuda_vectors_tails_and_storage_offsets(dtype, shape, offsets): + count = math.prod(shape) + gate = torch.randn(count + offsets[0], device="cuda", dtype=dtype)[offsets[0] :].view(shape) + up = torch.randn(count + offsets[1], device="cuda", dtype=dtype)[offsets[1] :].view(shape) + gate_before, up_before = gate.clone(), up.clone() + out = silu_and_mul(gate, up) + torch.testing.assert_close(out, silu_and_mul_reference(gate, up)) + assert out.shape == gate.shape and out.dtype == dtype and out.device == gate.device + torch.testing.assert_close(gate, gate_before, rtol=0, atol=0) + torch.testing.assert_close(up, up_before, rtol=0, atol=0) + if count: + assert out.data_ptr() not in (gate.data_ptr(), up.data_ptr()) + + +@pytest.mark.gpu +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +def test_silu_and_mul_cuda_special_values_and_underflow(dtype): + info = torch.finfo(dtype) + values = [ + float("nan"), + float("inf"), + -float("inf"), + -100, + -12, + -1, + -0.0, + 0.0, + 1, + 12, + 100, + info.tiny, + info.tiny / 128, + info.max, + ] + gate, up = torch.meshgrid( + torch.tensor(values, device="cuda", dtype=dtype), + torch.tensor(values, device="cuda", dtype=dtype), + indexing="ij", + ) + gate, up = gate.contiguous(), up.contiguous() + torch.testing.assert_close( + silu_and_mul(gate, up), silu_and_mul_reference(gate, up), equal_nan=True + ) + + +@pytest.mark.gpu +def test_silu_and_mul_cuda_rejects_noncontiguous_input_without_copying(): + gate, up = torch.randn(2, 14, device="cuda").chunk(2, dim=-1) + with pytest.raises(RuntimeError, match="contiguous inputs"): + silu_and_mul(gate, up) + + +@pytest.mark.gpu +def test_silu_and_mul_cuda_rejects_autograd_but_allows_inference(): + gate = torch.randn(16, device="cuda", requires_grad=True) + up = torch.randn_like(gate) + with pytest.raises(RuntimeError, match="inference-only"): + silu_and_mul(gate, up) + with torch.no_grad(): + out = silu_and_mul(gate, up) + torch.testing.assert_close(out, silu_and_mul_reference(gate, up)) + assert not out.requires_grad + + +@pytest.mark.gpu +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +def test_silu_and_mul_cuda_current_stream_and_graph_replay(dtype): + stream = torch.cuda.Stream() + with torch.cuda.stream(stream): + gate = torch.randn(4097, device="cuda", dtype=dtype) + up = torch.randn_like(gate) + for _ in range(3): + silu_and_mul(gate, up) + stream.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + out = silu_and_mul(gate, up) + gate.fill_(3) + up.fill_(2) + graph.replay() + torch.testing.assert_close(out, silu_and_mul_reference(gate, up)) + + +@pytest.mark.gpu +@pytest.mark.parametrize("count", [(1 << 24) - 1, 1 << 24, (1 << 24) + 1]) +def test_silu_and_mul_cuda_fp32_streaming_grid_boundary(count): + # Four FP32 elements/vector: this boundary fills 16384 blocks of 256 threads. + gate = torch.randn(count, device="cuda") + up = torch.randn_like(gate) + torch.testing.assert_close(silu_and_mul(gate, up), silu_and_mul_reference(gate, up)) + + +@pytest.mark.gpu +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("offset", [0, 1]) +def test_silu_and_mul_cuda_low_precision_storage_patterns(dtype, offset): + # Exercise all gate bit patterns with both vector/tail and scalar-only access. + values = torch.arange(65536, device="cuda", dtype=torch.int32).to(torch.int16).view(dtype) + info = torch.finfo(dtype) + subnormal = info.tiny / (1024 if dtype == torch.float16 else 128) + factors = torch.tensor( + [ + 0.0, + -0.0, + 1.0, + -1.0, + 0.5, + -0.5, + 2.0, + -2.0, + info.tiny, + subnormal, + info.max, + -info.max, + float("inf"), + -float("inf"), + float("nan"), + ], + device="cuda", + dtype=dtype, + ) + gate = torch.empty(values.numel() * factors.numel() + 1 + offset, device="cuda", dtype=dtype)[ + offset: + ] + up = torch.empty_like(gate) + gate[:-1].copy_(values.repeat(factors.numel())) + up[:-1].copy_(factors.repeat_interleave(values.numel())) + # The aligned path ends with one scalar multiplication yielding a subnormal. + gate[-1], up[-1] = 1.0, subnormal + actual = silu_and_mul(gate, up) + expected = silu_and_mul_reference(gate, up) + torch.testing.assert_close(actual, expected, rtol=0, atol=0, equal_nan=True) + valid = ~torch.isnan(expected) + assert torch.equal(actual[valid].view(torch.int16), expected[valid].view(torch.int16)) diff --git a/tests/test_gsm8k_backends.py b/tests/test_gsm8k_backends.py index 59adce7d..29be03e7 100644 --- a/tests/test_gsm8k_backends.py +++ b/tests/test_gsm8k_backends.py @@ -1,5 +1,7 @@ """CPU tests for native generation stopping and cleanup.""" +import subprocess +import sys from types import SimpleNamespace import pytest @@ -16,6 +18,21 @@ } +def test_backend_helpers_import_without_registering_unrelated_lm_eval_models(): + pytest.importorskip("lm_eval") + script = """ +import sys +from pathlib import Path +import lm_eval +from benchmarks import gsm8k_backends as backends +assert 'lm_eval.models' not in sys.modules +expected = Path(lm_eval.__file__).parent / 'models' / 'utils.py' +for helper in (backends.stop_sequences_criteria, backends.postprocess_generated_text): + assert Path(helper.__code__.co_filename) == expected +""" + subprocess.run([sys.executable, "-c", script], check=True, capture_output=True, timeout=60) + + class Engine: def __init__(self, failure=None, depths=(4, 4)): self.failure = failure diff --git a/tests/test_native_build.py b/tests/test_native_build.py new file mode 100644 index 00000000..2a4bf8f1 --- /dev/null +++ b/tests/test_native_build.py @@ -0,0 +1,68 @@ +"""Build selection without requiring a local GPU or compiler in CPU CI.""" + +import runpy +import sys +from pathlib import Path +from types import ModuleType, SimpleNamespace + +import pytest + + +@pytest.fixture +def setup_calls(monkeypatch): + calls = [] + setuptools = ModuleType("setuptools") + setuptools.setup = lambda **options: calls.append(options) + monkeypatch.setitem(sys.modules, "setuptools", setuptools) + return calls + + +@pytest.mark.parametrize( + "cuda_version,toolkit,nvcc,visible,architecture,expected", + [ + (None, True, True, True, None, False), + ("13.0", False, False, True, None, False), + ("13.0", True, False, True, None, False), + ("13.0", True, True, False, None, False), + ("13.0", True, True, False, "native", False), + ("13.0", True, True, True, None, True), + ("13.0", True, True, False, "10.3", True), + ], +) +def test_native_build_detects_toolchain_and_target( + monkeypatch, tmp_path, setup_calls, cuda_version, toolkit, nvcc, visible, architecture, expected +): + home = tmp_path / "cuda" + if nvcc: + (home / "bin").mkdir(parents=True) + (home / "bin" / "nvcc").touch() + torch = ModuleType("torch") + torch.version = SimpleNamespace(cuda=cuda_version) + torch.cuda = SimpleNamespace(is_available=lambda: visible) + extension = ModuleType("torch.utils.cpp_extension") + extension.CUDA_HOME = str(home) if toolkit else None + extension.CUDAExtension = lambda name, **options: {"name": name, **options} + extension.BuildExtension = object() + for name, module in { + "torch": torch, + "torch.utils": ModuleType("torch.utils"), + "torch.utils.cpp_extension": extension, + }.items(): + monkeypatch.setitem(sys.modules, name, module) + if architecture is None: + monkeypatch.delenv("TORCH_CUDA_ARCH_LIST", raising=False) + else: + monkeypatch.setenv("TORCH_CUDA_ARCH_LIST", architecture) + runpy.run_path(str(Path(__file__).resolve().parents[1] / "setup.py")) + assert len(setup_calls) == 1 + assert bool(setup_calls[0].get("ext_modules")) == expected + if expected: + assert setup_calls[0]["ext_modules"][0]["name"] == "vllm_rlt._C" + # Architecture selection is delegated to PyTorch, not hardcoded here. + assert "-gencode" not in repr(setup_calls[0]) + + +def test_native_build_without_torch_produces_python_package(monkeypatch, setup_calls): + monkeypatch.setitem(sys.modules, "torch", None) + runpy.run_path(str(Path(__file__).resolve().parents[1] / "setup.py")) + assert setup_calls == [{}] diff --git a/tests/test_ouro.py b/tests/test_ouro.py index b99e3788..e7f29cd3 100644 --- a/tests/test_ouro.py +++ b/tests/test_ouro.py @@ -1,6 +1,9 @@ """CPU numerical checks against an independent, dense Ouro implementation.""" import json +import subprocess +import sys +import textwrap import pytest import torch @@ -9,6 +12,7 @@ from tests.helpers import make_cpu_cache, tiny_ouro_config from tests.reference import dense_reference from vllm_rlt.models import OuroConfig, OuroForCausalLM +from vllm_rlt.models.ouro import OuroMLP @pytest.fixture @@ -17,6 +21,74 @@ def model(): return OuroForCausalLM(tiny_ouro_config()) +@pytest.mark.gpu +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32]) +@pytest.mark.parametrize("rows", [1, 7, 32]) +def test_cuda_mlp_matches_torch_expression(dtype, rows): + torch.manual_seed(72) + layer = OuroMLP(tiny_ouro_config(intermediate_size=17)).to(device="cuda", dtype=dtype) + layer.requires_grad_(False) + hidden = torch.randn(rows, layer.gate_proj.in_features, device="cuda", dtype=dtype) + expected = layer.down_proj( + torch.nn.functional.silu(layer.gate_proj(hidden)) * layer.up_proj(hidden) + ) + torch.testing.assert_close(layer(hidden), expected, rtol=0, atol=0) + + +@pytest.mark.gpu +@pytest.mark.parametrize("native_available", [True, False]) +def test_cuda_mlp_selects_native_or_missing_extension_fallback(native_available): + # Dispatcher registrations last for the process lifetime. Exercise an absent + # extension in a fresh process rather than changing an already loaded backend. + script = textwrap.dedent(f""" + import importlib + import torch + original_import = importlib.import_module + def import_module(name, *args, **kwargs): + if name == 'vllm_rlt._C' and not {native_available}: + raise ModuleNotFoundError(name=name) + return original_import(name, *args, **kwargs) + importlib.import_module = import_module + from tests.helpers import tiny_ouro_config + from vllm_rlt.models.ouro import OuroMLP + assert not torch.cuda.is_initialized() + assert torch._C._dispatch_has_kernel_for_dispatch_key( + 'vllm_rlt::silu_and_mul', 'CUDA') == {native_available} + layer = OuroMLP(tiny_ouro_config()).cuda().bfloat16().requires_grad_(False) + hidden = torch.randn(3, layer.gate_proj.in_features, + device='cuda', dtype=torch.bfloat16) + expected = layer.down_proj( + torch.nn.functional.silu(layer.gate_proj(hidden)) * layer.up_proj(hidden)) + with torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CPU]) as prof: + actual = layer(hidden) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + names = {{event.key for event in prof.key_averages()}} + assert 'vllm_rlt::silu_and_mul' in names + assert ('aten::silu' in names) == (not {native_available}) + """) + subprocess.run([sys.executable, "-c", script], check=True, capture_output=True, timeout=60) + + +@pytest.mark.gpu +def test_cuda_mlp_graph_replays_updated_hidden(): + layer = OuroMLP(tiny_ouro_config(intermediate_size=17)).cuda().bfloat16().requires_grad_(False) + stream = torch.cuda.Stream() + with torch.cuda.stream(stream): + hidden = torch.randn(3, layer.gate_proj.in_features, device="cuda", dtype=torch.bfloat16) + for _ in range(3): + layer(hidden) + stream.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + actual = layer(hidden) + hidden.fill_(0.5) + graph.replay() + expected = layer.down_proj( + torch.nn.functional.silu(layer.gate_proj(hidden)) * layer.up_proj(hidden) + ) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + @pytest.mark.parametrize("kv_heads", [1, 2, 4]) def test_packed_prefill_matches_dense_at_every_depth(kv_heads): torch.manual_seed(8) diff --git a/vllm_rlt/kernels/__init__.py b/vllm_rlt/kernels/__init__.py index 64b75c22..7cf7031f 100644 --- a/vllm_rlt/kernels/__init__.py +++ b/vllm_rlt/kernels/__init__.py @@ -1 +1,5 @@ -"""Attention backends; importing this package does not initialize a GPU.""" +"""Register native and Torch kernels once without initializing a GPU.""" + +from .activation import silu_and_mul + +__all__ = ["silu_and_mul"] diff --git a/vllm_rlt/kernels/_extension.py b/vllm_rlt/kernels/_extension.py new file mode 100644 index 00000000..93e0199a --- /dev/null +++ b/vllm_rlt/kernels/_extension.py @@ -0,0 +1,23 @@ +"""Load the optional native library once, outside kernel implementation code.""" + +from functools import cache +from importlib import import_module + + +@cache +def try_load_extension() -> bool: + """Load once; a missing optional extension is cached as unavailable.""" + try: + import_module("vllm_rlt._C") + except ModuleNotFoundError as exc: + if exc.name != "vllm_rlt._C": + raise + return False + return True + + +def load_extension() -> None: + if not try_load_extension(): + raise RuntimeError( + "Native CUDA kernels are not built. See docs/kernels.md for build instructions." + ) diff --git a/vllm_rlt/kernels/activation.py b/vllm_rlt/kernels/activation.py new file mode 100644 index 00000000..4123006c --- /dev/null +++ b/vllm_rlt/kernels/activation.py @@ -0,0 +1,39 @@ +"""MLP activation operators and their unfused correctness reference.""" + +import torch +import torch.nn.functional as F + +from vllm_rlt.kernels._extension import try_load_extension + + +def _validate_inputs(gate: torch.Tensor, up: torch.Tensor) -> None: + if gate.shape != up.shape: + raise ValueError("silu_and_mul expects inputs with the same shape") + if gate.device != up.device: + raise ValueError("silu_and_mul expects inputs on the same device") + if gate.dtype != up.dtype: + raise TypeError("silu_and_mul expects inputs with the same dtype") + if gate.dtype not in (torch.float32, torch.float16, torch.bfloat16): + raise TypeError("silu_and_mul supports float32, float16 and bfloat16") + + +def silu_and_mul_reference(gate: torch.Tensor, up: torch.Tensor) -> torch.Tensor: + """Compute ``F.silu(gate) * up`` without broadcasting or modifying either input. + + Inputs have equal shape, dtype and device. The result has the same shape, + dtype and device and owns new storage. Noncontiguous inputs are supported + by this reference; the native entry requires contiguous inputs. + """ + _validate_inputs(gate, up) + return F.silu(gate) * up + + +# Keep the library alive: Python registrations are removed when it is destroyed. +# Native CUDA takes precedence over the composite Torch implementation; CPU and +# installations without the optional extension use the same schema and fallback. +_library = torch.library.Library("vllm_rlt", "FRAGMENT") +if not try_load_extension(): + _library.define("silu_and_mul(Tensor gate, Tensor up) -> Tensor") +_library.impl("silu_and_mul", silu_and_mul_reference, "CompositeExplicitAutograd") + +silu_and_mul = torch.ops.vllm_rlt.silu_and_mul diff --git a/vllm_rlt/models/ouro.py b/vllm_rlt/models/ouro.py index 4af3ff44..31b37f9c 100644 --- a/vllm_rlt/models/ouro.py +++ b/vllm_rlt/models/ouro.py @@ -18,8 +18,8 @@ import torch from torch import nn -from torch.nn import functional as F +from vllm_rlt.kernels import silu_and_mul from vllm_rlt.layers import ( RMSNorm, RotaryEmbedding, @@ -194,7 +194,7 @@ 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)) + return self.down_proj(silu_and_mul(self.gate_proj(hidden), self.up_proj(hidden))) class OuroDecoderLayer(nn.Module):