Skip to content

[Kernels] Native CUDA SiLU-and-Mul with unified dispatch and automatic builds - #160

Merged
bjf-frz merged 5 commits into
InferMatrixLab:mainfrom
bjf-frz:native-cuda-silu-and-mul
Oct 31, 2026
Merged

bjf-frz merged 5 commits into
InferMatrixLab:mainfrom
bjf-frz:native-cuda-silu-and-mul

Conversation

@bjf-frz

@bjf-frz bjf-frz commented Oct 31, 2026 •

Copy link
Copy Markdown
Collaborator

Required human contributor answers

These record the author's stated development requirements, summarized from the author-directed implementation work.

1. What outcome are you trying to achieve?

Add the first native CUDA SiLU-and-Mul implementation and a complete, reusable structure for developing further CUDA operators in vllm-rlt, including build, tests, comparisons, and model integration.

2. What is the highest-leverage variable for that outcome?

Activation memory movement, kernel launches, and efficient parallel computation, while preserving the existing numerical contract.

3. How will this work move that variable rather than create busywork?

Fuse SiLU and multiplication, vectorize the aligned body while handling the scalar tail in the same kernel, and use packed FP16/BF16 multiplication. Verify numerical behavior and measure against Torch, vLLM, and real-model inference rather than relying on intermediate metrics.

4. Who is the customer or user, and what do they actually need?

The author is developing native CUDA operators from vllm-rlt and needs a complete implementation/build/validation path that also supports adding more operators later.

PR necessity and bugfix evidence

Yes: the author requested a PR for the completed first operator and framework. This requires repository changes to native compilation, operator registration, and the Ouro MLP call site.

A separate prerequisite commit fixes an observed benchmark import failure with the configured lm-eval==0.4.9.2 and existing Transformers 5.14.1. On clean base 9ad03fce, pytest tests/test_gsm8k_accuracy.py::test_run_refuses_the_stored_baseline_at_another_depth fails with AttributeError: module transformers has no attribute AutoModelForVision2Seq: importing stopping helpers initializes lm-eval's unrelated model registry. The fix loads the installed standalone helper module without that registry, retaining the harness's exact stopping/postprocessing functions, and performs protocol validation before loading generation backends. A fresh-process regression checks this import path. No package pins or stored accuracy baseline are changed; this does not qualify the official HF release model's full generation path on Transformers 5.

Purpose

Ouro's CUDA MLP replaces separate SiLU and multiplication with one native operator on its separate gate/up projection outputs. Model code is simply:

self.down_proj(silu_and_mul(self.gate_proj(hidden), self.up_proj(hidden)))
  • Load the optional extension once at kernel-module initialization; reuse PyTorch dispatch for native CUDA versus Torch fallback. No per-forward availability/device branch, backend registry, or runtime compilation. Importing the package does not create a CUDA context. Broken ABI/dependencies and native execution errors propagate.
  • FP32/FP16/BF16 contiguous CUDA inputs; 16-byte aligned body plus scalar tail, scalar path for misaligned contiguous views; current device/stream and empty inputs handled. Preserve SiLU intermediate rounding before multiplication; FP16/BF16 use packed products, FP32 uses scalar products. Inference only; no torch.compile fake registration.
  • Automatically build through setuptools CUDAExtension when CUDA PyTorch, the toolkit's nvcc, and a visible GPU are available. PyTorch locates the toolkit and selects target architectures; no project build switch or manual architecture setting is required. Without the requirements, produce a Python-only distribution with an explicit skip message. Compilation failures propagate after a native build is selected. One shared extension supports future sources; CMake is not required. The first-run user guide includes automatic build/install commands and runtime selection. See docs/kernels.md.
  • Existing [Kernels] Bitwise-exact fused RMSNorm, RoPE and SiLU*up for CUDA graph decode #146 overlaps on SiLU fusion but implements Triton kernels with graph-scoped selection. This PR supplies a native CUDA implementation/build path; it does not incorporate that PR's RMSNorm/RoPE work or claim its measured speedups. Only Ouro is integrated here.

Test Plan

vllm-rlt Version / Commit: base 9ad03fce4aa92839a5fcc9a5bf4ae7f9c36d101d; tested working-tree contents match this PR's commit trees.

  • Environment: Python 3.12.3, torch 2.13.0+cu130, Transformers 5.14.1, lm-eval 0.4.9.2, CUDA toolkit 13.4, driver 610.43.02; NVIDIA B300 SXM6 AC. Triton attention for engine/model checks.
  • pre-commit run --all-files: Ruff lint/format and full CPU regression suite passed after the import fix. GPU skips in this hook are not GPU verification.
  • pytest tests/test_activation.py tests/test_ouro.py tests/test_cuda_graph.py --run-gpu -k 'not flash' -q: 218 passed, 12 FlashAttention combinations deselected. Includes real CUDA execution, FP16/BF16 exhaustive gate storage patterns with multiple up factors, tails/misalignment, numerical boundaries, streams, model dispatch/fallback, graph capture/replay, cancellation, KV release, and request ID reuse. Native tests require the built extension so fallback cannot hide a missing kernel.
  • Accuracy-helper tests: 110 passed, 1 skipped. Fresh-process checks retain the pinned helper source and avoid unrelated model registrations.
  • Automatic-build tests: 8 passed, covering missing PyTorch, CPU PyTorch, missing toolkit/nvcc, no visible GPU, automatic visible-device builds, and PyTorch's optional architecture override for cross-compilation. Actual B300 wheel compilation with no project build switch or architecture variable generated SM103 and compute_103 PTX. An actual no-visible-GPU build produced a pure Python wheel. Both installed wheels passed three-dtype/odd-width numerics and full MLP Graph replay checks; the full pre-commit suite passed.
  • Built CUDA and Python-only wheels via pip; installed each into a separate target directory, then exercised three dtypes, odd widths, full MLP graph replay with updated input contents, and correct native/fallback selection.
  • Final microbenchmark: 18 cases (3 dtypes × rows 1/128/4096 × widths 5632/5633), three repeats, eager and graph timing; matched original vLLM source at 2b678ac739166bd0356883048673bbd1c34a3b0e compiled as a standalone baseline. All tested native outputs match eager Torch exactly. A staged FP64 oracle is also recorded.

Test Result

Kernel comparison

Median CUDA Graph device latency in microseconds, shape [4096,5632]. vLLM receives packed input and a preallocated output prepared before timing; the native operator receives separate inputs and allocates output. These are kernel comparisons, not E2E speedups.

dtype Torch Native CUDA vLLM prepacked
bfloat16 41.492 23.134 26.307
float16 41.012 22.941 25.327
float32 69.808 40.636 41.217

Real Ouro-1.4B BF16 inference

One B300, public LLM worker, greedy, fixed 4 loops and ignore-EOS; three unprofiled interleaved A/B repeats after workload-specific warmup. Both backends use the unified API; the baseline uses its Torch implementation with native discovery disabled before worker imports. Scheduler max sequences=16, token budget=4096, prefill chunk=1024; last_exited KV, 4416 blocks × 16 tokens; synchronous scheduling, Graph max rows=16. Input IDs derive from 32 GSM8K questions repeated/cropped to exact lengths. Timing includes local IPC and generation; these are short closed-loop bursts, not saturation measurements. Physical GPU 5 was monitored for foreign processes throughout this final run.

Mode Concurrency / input / output / requests Torch tok/s Native tok/s Change
eager 1 / 128 / 32 / 8 29.38 31.25 +6.37%
eager 16 / 1024 / 32 / 32 119.05 119.67 +0.52%
graph 1 / 128 / 32 / 8 66.37 65.72 -0.99%
graph 16 / 1024 / 32 / 32 131.69 132.43 +0.56%

Model-level benefit is limited: eager single-request improves, while concurrent and Graph results are around the run-to-run variation; this is not a general throughput-speedup claim. TTFT/E2E/TPOT samples and per-run distributions were retained locally. All 240 fixed-length requests / 7680 output tokens per backend match in generated token IDs and exit depths. An additional 8 natural GSM8K questions, up to 64 output tokens with adaptive exits, match on both execution modes; this is output consistency, not a newly scored GSM8K task result. All requests and KV blocks return to zero after each batch.

Attribution and limits

Four separate complete profiles were captured after timing: Torch/native × ops-only/full (CPU+CUDA; full adds shapes, memory, stacks). For a 16-token prefill plus one decode, 24 layers × 4 loops × 2 submissions gives 192 MLP activations. Separate 192 SiLU + 192 Mul kernels become 192 fused kernels; all device kernels fall from 11288 to 11096. Full traces contain [16,5632]/[1,5632] shapes and OuroMLP stacks. Total temporary allocation/free volume falls by exactly 18,382,848 bytes, matching the removed SiLU intermediate; this is cumulative allocation volume, not peak-memory savings. Launch counts are not a latency breakdown.

Final traces, raw runs, wheel checks, source fingerprints and environment records are retained locally under ignored artifacts/silu-and-mul-pr-20261031/; large artifacts are not committed or publicly uploaded. Only B300 was exercised; FlashAttention, other GPUs, other model integrations, autograd and torch.compile are not qualified here.

Code-quality review reused Torch dispatch, kept process-wide registration/availability in the kernel modules, and avoided per-layer flags, speculative adapters and unrelated runtime changes.

Build and reproduce

Use an existing compatible CUDA PyTorch and toolkit; keep the allocated GPU visible. --no-build-isolation exposes that environment's PyTorch to the build. An isolated build without PyTorch cannot compile the extension and produces Python-only output. Without a GPU, an optional standard TORCH_CUDA_ARCH_LIST override enables cross-compilation; ordinary GPU installs detect architectures automatically.

python -m pip install "torch>=2.5" setuptools wheel ninja
python -m pip install --no-build-isolation -e .
CUDA_VISIBLE_DEVICES=<allocated_gpu> python -m pytest \
  tests/test_activation.py tests/test_ouro.py tests/test_cuda_graph.py \
  --run-gpu -k 'not flash' -q
CUDA_VISIBLE_DEVICES=<allocated_gpu> python -m benchmarks.kernels.silu_and_mul \
  --output artifacts/silu-and-mul/results.json \
  --profile-dir artifacts/silu-and-mul/profiles

@bjf-frz bjf-frz changed the title [Kernels] Add native CUDA SiLU-and-Mul with unified dispatch and optional pip build [Kernels] Native CUDA SiLU-and-Mul with unified dispatch and automatic builds Oct 31, 2026
@bjf-frz
bjf-frz marked this pull request as ready for review October 31, 2026 09:43
@bjf-frz
bjf-frz merged commit 88c8f96 into InferMatrixLab:main Oct 31, 2026
1 check passed
kunkunblueberry added a commit to kunkunblueberry/vllm-rlt that referenced this pull request Oct 31, 2026
Brings in InferMatrixLab#159 (scheduler wavefront selection / prefill progress) and InferMatrixLab#160
(native CUDA SiLU-and-Mul) so the branch is measured against current main.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant