Repository navigation
[Kernels] Vectorize misaligned SiLU inputs with a per-SM cutoff - #163
Merged
bjf-frz merged 6 commits intoNov 4, 2026
Merged
Conversation
bjf-frz
marked this pull request as ready for review
November 4, 2026 07:20
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Main processes an entire SiLU-and-Mul input scalarly when any input or output address is misaligned. This PR adds a vector path for sufficiently large misaligned contiguous inputs while retaining the original arithmetic and a small-input scalar path. On B300, BF16
[8611, 5632]with gate/up offsets(1, 3)measured 115.832 µs on main and 64.196 µs on this PR; details and scope are below.Required human contributor answers
The author's public answers in #160, “Required human contributor answers” remain the intent for this follow-up; this PR improves that merged operator rather than introducing a new framework or changing its arithmetic contract.
1. What outcome are you trying to achieve?
Reuse the author's answer in #160, question 1: develop native CUDA operators with a complete, reusable build, integration and validation path. This follow-up improves the first operator's supported misaligned-input path.
2. What is the highest-leverage variable for that outcome?
Reuse #160, question 2: activation memory movement, kernel launches and efficient parallel computation while preserving the numerical contract. Here the changed variable is device latency for independently shifted contiguous inputs.
3. How will this work move that variable rather than create busywork?
Reuse #160, question 3: vectorized fused computation, preserved rounding, numerical checks and matched performance measurements. The concrete extension is safe aligned reads for misaligned inputs, with a measured scalar fallback where reconstruction costs more.
4. Who is the customer or user, and what do they actually need?
Reuse #160, question 4: the author developing native CUDA operators needs a complete implementation/build/validation path that can support subsequent operators.
PR necessity and bugfix evidence
The author explicitly requested a PR for this follow-up. It requires native load/dispatch changes; an existing configuration cannot enable vector processing of these inputs. This is a measured performance optimization, not a claim that the previous scalar implementation returned incorrect results.
Purpose
This PR improves the existing SiLU-and-Mul CUDA operator for contiguous inputs whose starting addresses are not 16-byte aligned. Main currently falls back to scalar processing for the entire input if gate, up or output is misaligned. The PR adds a vector path for sufficiently large misaligned inputs while retaining the scalar path below a size cutoff. Aligned inputs continue to use vector loads and stores. The operator keeps its existing output and arithmetic contract.
Inputs and relationship to vLLM
Ouro has separate
gate_projandup_projlinear layers. For example, Ouro-1.4B transformshiddenwith shape[N, 2048]into two separate contiguous tensorsgateandup, each with shape[N, 5632]. They need not be adjacent in memory.silu_and_mul(gate, up)directly consumes those two tensors and returns an output of the same shape:The original vLLM operator used in the comparison accepts a single packed tensor. Adapting Ouro's separate tensors requires
cat((gate, up), dim=-1), producing[N, 11264], with each row laid out as[gate_row | up_row]. If an upstream projection already produces that packed layout, the original vLLM kernel can consume it directly and does not need this concatenation. The performance tables below therefore report both vLLM modes.Dispatch and launch strategy
The CUDA entry point validates the same shape, device and dtype, requires contiguous inputs, allocates a fresh output, and uses the current CUDA stream. Empty inputs return without a launch. Supported dtypes are BF16, FP16 and FP32; the native operator is inference-only.
launch<T>selects one of three paths:Alignednumel < thresholdScalarnumel >= thresholdMisalignedThe number of elements is the total across the contiguous input, rather than a separate dispatch for each row. The cutoff uses the input GPU's actual SM count:
For the 148-SM B300 used below, the cutoff is 48,496,640 BF16/FP16 elements or 24,248,320 FP32 elements. The
160is a reference workload for the heuristic, not a number of blocks physically assigned to each SM. The real launch uses 256 threads per block and computes its grid separately, capped at 4096 blocks, or 16384 for sufficiently large FP32 work. This cutoff follows the requested B300-based heuristic; it has not been tuned or validated as optimal for other GPUs.Misaligned vector processing
silu_and_mul_kernelhandles a scalar output prefix when required to make the first vector store 16-byte aligned. A 16-byte-aligned output has prefix length zero. The middle handles one 16-byte logical vector per thread iteration; the remaining incomplete output tail is scalar, in the same kernel. These regions partition the work; there is no device-wide barrier between them.For each middle vector,
load_vectorhandles gate and up independently:uint4directly.__funnelshift_r.For example, with BF16 storage
AandBboth 16-byte aligned, supposegate = A[2:]andup = B[6:], and the output is aligned. For logical elements 8–15, the vector body processes:Here indices are BF16 elements and the example assumes the displayed aligned chunks lie inside the input views. Gate and up can have different offsets; each is reconstructed separately. The selection is register work, not a separate global-memory copy or kernel launch. A 16-byte input chunk may overlap the chunk read for an adjacent logical vector.
Calculation and rounding
The vector body is shared by aligned and misaligned modes. A 16-byte vector contains 8 BF16/FP16 or 4 FP32 elements. For each element, SiLU is evaluated in FP32, rounded to the input dtype, and then multiplied by
up, preserving the existing staged arithmetic. BF16/FP16 multiply the rounded activations andupin pairs using__hmul2. FP32 still uses two scalar FP32 multiplications perfloat2; its memory loads and stores remain 16-byte vectors. Scalar prefix, tail and fallback use the corresponding scalar calculation.Measured examples
BF16 on one B300, CUDA Graph median latency in µs. The first two rows come from the November 4 model-width sweep; the last comes from the November 3 cutoff sweep. Offsets are gate/up storage offsets in elements.
[128, 5632](0, 0)[128, 5632](1, 3)[8611, 5632](1, 3)For the common Ouro examples, the PR and main have nearly the same latency; the smaller misaligned case remains scalar. In the 126-case November 4 sweep, the PR was faster than vLLM + cat in every measured case, but vLLM's already-packed input can be faster, as shown above. These are separate-input operator timings, not end-to-end model-throughput results. The complete test method and both 126-case matrices follow.
Test Plan
vllm-rlt Version / Commit: base
989921ac1d3897c5c43e27ec8cddbc91fa953194, tested PR head21c706e646ae350ab20844ec24bd6dc63ed6e27b.pre-commit run --all-filespassed Ruff lint, Ruff formatting and CPU regression tests.torch.equal; signed zero is not distinguished). This is not an exhaustive numerical proof.fe71d62; Compute Sanitizer memcheck passed 24 cases with 0 errors. The subsequent source difference to the current head is the enum name and documentation only; current-head targeted tests are reported separately above.Test Result: performance
Measured on one idle B300, with no other compute process detected on that device during the run. The other GPUs were in use. These are operator-level measurements, not model throughput. Each point reports the median of 5 shuffled rounds of unprofiled CUDA Graph timing, 32 calls/graph, 30 graph replays per round (8 at the threshold); each function has 10 warmups. Clocks were not locked. Inputs use fixed seed 42, gate sampled from normal(std=3), up from normal(std=1). Gate/up are separate contiguous allocations; offsets in the table are in elements. The GPU results are limited to this hardware and workload.
The three code versions are:
21c706e, the restored output-prefix implementation.989921a, compiled from the same SiLU source in a separate operator namespace; misaligned inputs use scalar access.2b678ac739166bd0356883048673bbd1c34a3b0e, compiled directly from its originalcsrc/activation_kernels.cu. Its interface consumes one packed[gate | up]tensor, so it is shown in two timing modes:torch.cat((gate, up), dim=-1), output allocation, and the vLLM kernel are inside the timed graph. CUDA Graph timings measure device work; host-side allocation overhead is not included. This column reflects adapting separate inputs to the vLLM interface.Representative latency, in microseconds (lower is better):
All rows above are BF16. For the 120 common-shape measurements (BF16/FP16/FP32, 1–1024 rows, widths 5632/5633, aligned and independently shifted inputs), the PR/main latency ratio has a median of approximately 1.00; the misaligned inputs are below the cutoff and use the scalar path. For the 60 aligned cases, the observed main/PR speedup ranges from 0.960× to 1.015×. Small differences at these microsecond times are workload-specific; no blanket aligned-path speedup is claimed.
The cutoff is easier to see with fixed width 5632 and neighboring row counts. In each dtype the first row is just below the cutoff and the second is just above it:
Immediately above the cutoff, the PR is 1.80× faster for BF16, 1.87× for FP16, and 1.26× for FP32 than main on these shapes. Below it, the PR and main remain on the scalar path and have nearly identical latencies. For the above-cutoff examples, the PR's five-round ranges were 64.171–64.220 µs (BF16), 63.244–63.292 µs (FP16), and 53.564–53.592 µs (FP32); main's corresponding ranges were 115.800–115.840, 118.436–118.484, and 67.688–67.720 µs. These are repetitions of one benchmark run, not independent multi-run confidence intervals.
The vLLM packed column can be faster: at BF16
[8611, 5632]it takes 51.240 µs versus the PR's 64.196 µs, while adapting separate inputs withcattakes 218.988 µs. The packed layout is different, so that comparison does not isolate the cost of independent misalignment. For width 5633 the local vLLM dispatch falls back to scalar because the hidden dimension does not divide its vector width; the RLT operator handles the separately allocated tensors as flat contiguous arrays. Short ops-only and full profiles confirm the scalar and vector branches and show a separatecatkernel for the final column; profiled timings are not used as latency evidence.All 126 measured cases (CUDA Graph median, µs)
BF16
FP16
FP32
The benchmark reuses the repository's
benchmarks/kernels/silu_and_mul.pytiming and profiling helpers. The standalone comparison script, all five timing samples including eager timings, native build metadata, profiles, and GPU process-monitor log are retained in the local ignoredartifacts/silu-unified-20261103/directory; they are not committed. The displayed table includes all measured shapes.Additional model-width measurements (2026-11-04)
The operator was rerun at FFN intermediate widths read from the local model configurations: Ouro-1.4B 5632, Nanbeige4.2-3B 10752, and Qwen3.8-27B text 17408. These are synthetic operator inputs using real configuration widths, not runs of those models; using a model's width does not mean this PR integrates the operator into that model. In this PR's Ouro code,
gate_projandup_projare independent linear layers and return separate tensors beforesilu_and_mul.On an otherwise idle B300 with no other compute process detected on that GPU, 126 cases were measured: BF16 rows 1/4/16/64/128/256/512/1024/2048/4096/8192; FP16 and FP32 rows 1/16/128/1024/4096; each at all three widths with separately allocated aligned
(0, 0)and intentionally shifted(1, 3)input offsets. Offsets are in elements. The same PR, main and original vLLM commits listed above were used. Five shuffled CUDA Graph rounds, 32 calls per graph, 30/10/5 replays depending on size; 10 warmups per function. Results are medians in µs. As above, vLLM packed excludes packing and uses a packed input; vLLM + cat includes the GPU work to pack the separate inputs. Output allocation happens during graph capture, so Python/host allocator time is not measured.The PR was faster than vLLM + cat in all 126/126 cases here, with a 1.41–4.29× ratio of median latencies. That is restricted to this GPU, matrix and separate-input contract; a model that already produces packed gate/up can use the vLLM packed path without
cat. In the shifted subset, the PR was faster than vLLM packed in 25/63 cases, emphasizing that the two input layouts must not be conflated. All 126 cases matched elementwise across backends usingtorch.equalon finite random input (signed zeros not distinguished).Representative BF16 aligned inputs:
Representative BF16 shifted inputs:
For BF16 Ouro width 5632, even
[8192, 5632]is below the B300 cutoff of 48,496,640 elements; the shifted input therefore uses the scalar path in both PR and main. Its advantage over vLLM + cat is not a speedup introduced by this PR. At BF16[8192, 10752]with offsets(1, 3), the PR enters the misaligned vector path and measures 113.325 µs versus main's 209.165 µs. Five-round ranges for that case are 113.286–113.363 µs and 209.165–209.204 µs, respectively. Representative ops-only/full profiles confirm the extracatkernel and the selected PR/main paths; profiling was separate from timing.All 126 additional model-width cases (CUDA Graph median, µs)
Ouro-1.4B width 5632
Nanbeige4.2-3B width 10752
Qwen3.8-27B text width 17408
The exact comparison script, per-round CUDA Graph and eager timings, model-config hashes, profiles, and GPU process-monitor log are retained locally under ignored
artifacts/silu-unified-20261103/realistic-shapes-*files; the complete median matrix is included above.AI assistance was used for implementation, measurements and author self-review.