diff --git a/csrc/activation/silu_and_mul.cu b/csrc/activation/silu_and_mul.cu index faadcda1..08c05c7a 100644 --- a/csrc/activation/silu_and_mul.cu +++ b/csrc/activation/silu_and_mul.cu @@ -2,6 +2,7 @@ #include "../ops.h" #include +#include #include #include #include @@ -77,28 +78,94 @@ __device__ __forceinline__ T compute_scalar(T gate, T up) { return Math::mul(activated, up); } -template +template +__device__ __forceinline__ Vector load_vector(const T* input, int64_t index, + int64_t count) { + constexpr int elements = 16 / sizeof(T); + const int offset_bytes = aligned ? 0 : reinterpret_cast(input + index) & 15; + Vector result; + if (offset_bytes == 0) { + *reinterpret_cast(&result) = + *reinterpret_cast(input + index); + return result; + } + + const int offset = offset_bytes / sizeof(T); + const int64_t lower = index - offset; + Vector head{}, tail{}; + // Never read before or after the input view, even if its storage is larger. + if (lower >= 0) { + *reinterpret_cast(&head) = *reinterpret_cast(input + lower); + } else { +#pragma unroll + for (int j = 0; j < elements; ++j) { + if (j >= offset) reinterpret_cast(&head)[j] = input[lower + j]; + } + } + if (lower + 2 * elements <= count) { + *reinterpret_cast(&tail) = + *reinterpret_cast(input + lower + elements); + } else { +#pragma unroll + for (int j = 0; j < elements; ++j) { + if (j < offset) reinterpret_cast(&tail)[j] = input[lower + elements + j]; + } + } + + // Reassemble the requested logical vector with register moves and bit shifts. + uint4 lo = *reinterpret_cast(&head); + uint4 hi = *reinterpret_cast(&tail); + if (offset_bytes & 8) { + lo = make_uint4(lo.z, lo.w, hi.x, hi.y); + hi = make_uint4(hi.z, hi.w, 0, 0); + } + if (offset_bytes & 4) { + lo = make_uint4(lo.y, lo.z, lo.w, hi.x); + hi.x = hi.y; + } + if constexpr (sizeof(T) == 2) { + if (offset_bytes & 2) { + lo = make_uint4(__funnelshift_r(lo.x, lo.y, 16), __funnelshift_r(lo.y, lo.z, 16), + __funnelshift_r(lo.z, lo.w, 16), __funnelshift_r(lo.w, hi.x, 16)); + } + } + *reinterpret_cast(&result) = lo; + return result; +} + +enum class AccessMode { Aligned, Misaligned, Scalar }; + +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) { + if constexpr (mode != AccessMode::Scalar) { constexpr int elements = 16 / sizeof(T); constexpr int pairs = 16 / sizeof(typename Math::Pair); - const int64_t vectors = count / elements; + int64_t prefix = 0; + if constexpr (mode == AccessMode::Misaligned) { + prefix = min(count, static_cast( + ((16 - (reinterpret_cast(out) & 15)) & 15) / sizeof(T))); + } + for (int64_t i = thread; i < prefix; i += stride) { + out[i] = compute_scalar(gate[i], up[i]); + } + const int64_t vectors = (count - prefix) / 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); + const int64_t index = prefix + i * elements; + Vector g = load_vector(gate, index, count); + Vector u = load_vector(up, index, count); + Vector result; #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); + *reinterpret_cast(out + index) = *reinterpret_cast(&result); } - tail_start = vectors * elements; + tail_start = prefix + vectors * elements; } - // Only valid scalar addresses are touched, including for count < one vector. + // Scalar mode starts at zero; vector modes handle only their remaining tail. for (int64_t i = tail_start + thread; i < count; i += stride) { out[i] = compute_scalar(gate[i], up[i]); } @@ -114,8 +181,19 @@ void launch(const at::Tensor& gate, const at::Tensor& up, at::Tensor& out) { 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; + bool scalar = false; + if (!vectorized) { + // Use 160 reference vector blocks per SM to determine the element cutoff. + // The actual launch grid is calculated separately and may be capped. + constexpr int reference_blocks_per_sm = 160; + constexpr int elements_per_vector = 16 / sizeof(T); + const int sms = at::cuda::getDeviceProperties(gate.get_device())->multiProcessorCount; + const int64_t threshold = + static_cast(sms) * reference_blocks_per_sm * threads * elements_per_vector; + scalar = count < threshold; + } + const int64_t work = scalar ? count : std::max(count / (16 / sizeof(T)), 1); int max_blocks = 4096; if constexpr (std::is_same_v) { // More blocks improve large FP32 streaming workloads. Require enough work @@ -128,9 +206,11 @@ void launch(const at::Tensor& gate, const at::Tensor& up, at::Tensor& out) { 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); + silu_and_mul_kernel<<>>(g, u, result, count); + } else if (scalar) { + silu_and_mul_kernel<<>>(g, u, result, count); } else { - silu_and_mul_kernel<<>>(g, u, result, count); + silu_and_mul_kernel<<>>(g, u, result, count); } C10_CUDA_KERNEL_LAUNCH_CHECK(); } diff --git a/docs/kernels.md b/docs/kernels.md index a1fda4e1..9f01fc48 100644 --- a/docs/kernels.md +++ b/docs/kernels.md @@ -38,10 +38,31 @@ nor copies incompatible inputs. 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. +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 may have independently misaligned storage offsets. For these +inputs, the scalar/vector cutoff scales with the input GPU's actual SM count: +`SMs * 160 * 256 * (16 / sizeof(dtype))` total elements. The reference is 160 +vector blocks per SM, with 256 threads and 16 bytes per vector: eight FP16/BF16 +elements or four FP32 elements. On a 148-SM B300, the corresponding cutoffs are +48,496,640 FP16/BF16 elements and 24,248,320 FP32 elements. The reference block +count specifies the workload cutoff; the actual launch grid is computed below +and may be capped. This threshold is a heuristic, not a measured optimum. + +Below the cutoff, misaligned inputs use scalar computation with one work item +per element. At or above it, the kernel loads each input's adjacent aligned +vectors and reassembles the required logical vector with register moves and bit +shifts. Scalar loads handle input prefixes and suffixes where a full aligned +read would extend outside the view. Aligned and misaligned vector paths share the +same calculation and output loop; output prefix and tail use scalar +computation. No input copy or extra kernel launch is needed. Fully aligned inputs +retain their direct vector loads even for small inputs. Empty inputs return an +empty output without launching. On B300, correctness was tested for all +independent 16-byte input offsets just above the cutoff and for every vector +remainder with selected offsets. The cutoff applies to total element count, +independent of tensor rank or model width. Raising it to 160 reference blocks +per SM also sends many previously vectorized inputs through the slower scalar +path; the performance impact depends on the workload. SiLU uses FP32 arithmetic and rounds to the input dtype before multiplying by `up`, preserving the reference's intermediate rounding. FP16/BF16 vector paths @@ -55,10 +76,11 @@ 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; +Each vector work item processes four FP32 elements or eight FP16/BF16 elements; +the scalar path for small misaligned inputs processes one element per work item. +The streaming threshold requires +enough work to fill the larger grid. Aligned inputs retain their existing launch +geometry. 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 diff --git a/tests/test_activation.py b/tests/test_activation.py index dbedb464..0b877a0b 100644 --- a/tests/test_activation.py +++ b/tests/test_activation.py @@ -1,10 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 """Numerics, tails, layouts and stream contracts for SiLU-and-Mul.""" +import json import math import subprocess import sys import textwrap +from itertools import product import pytest import torch @@ -160,7 +162,115 @@ def test_silu_and_mul_cuda_vectors_tails_and_storage_offsets(dtype, shape, offse @pytest.mark.gpu @pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) -def test_silu_and_mul_cuda_special_values_and_underflow(dtype): +def test_silu_and_mul_cuda_all_alignment_pairs_match_aligned_path(dtype): + # The two inputs can have different offsets. Cover every 16-byte alignment + # pair in the scalar fallback, including zero and short input lengths. + elements = 16 // torch.empty((), dtype=dtype).element_size() + counts = [*range(2 * elements + 2), 17 * elements - 1, 17 * elements + 1] + integer = torch.int32 if dtype == torch.float32 else torch.int16 + generator = torch.Generator(device="cuda").manual_seed(42) + for count in counts: + for gate_offset in range(elements): + for up_offset in range(elements): + gate = torch.randn( + count + gate_offset, device="cuda", dtype=dtype, generator=generator + )[gate_offset:] + up = torch.randn( + count + up_offset, device="cuda", dtype=dtype, generator=generator + )[up_offset:] + expected = silu_and_mul(gate.clone(), up.clone()) + actual = silu_and_mul(gate, up) + assert torch.equal(actual.view(integer), expected.view(integer)), ( + dtype, + count, + gate_offset, + up_offset, + ) + + +@pytest.mark.gpu +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("delta", range(-1, 8)) +def test_silu_and_mul_cuda_threshold_and_vector_offsets_match_aligned_path(dtype, delta): + # Short-offset tests now take the scalar fallback. Exercise all independent + # input alignments above the threshold and each possible scalar tail. + sms = torch.cuda.get_device_properties(torch.cuda.current_device()).multi_processor_count + elements = 16 // torch.empty((), dtype=dtype).element_size() + count = sms * 160 * 256 * elements + delta + integer = torch.int32 if dtype == torch.float32 else torch.int16 + generator = torch.Generator(device="cuda").manual_seed(43) + gate_storage = torch.randn( + count + elements - 1, device="cuda", dtype=dtype, generator=generator + ) + up_storage = torch.randn(count + elements - 1, device="cuda", dtype=dtype, generator=generator) + offsets = product(range(elements), repeat=2) if delta == 1 else [(1, 3), (elements - 1, 1)] + for go, uo in offsets: + gate = gate_storage[go : go + count] + up = up_storage[uo : uo + count] + expected = silu_and_mul(gate.clone(), up.clone()) + actual = silu_and_mul(gate, up) + assert torch.equal(actual.view(integer), expected.view(integer)), (dtype, count, go, uo) + + +@pytest.mark.gpu +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_silu_and_mul_cuda_uses_device_scaled_threshold(dtype, tmp_path): + # Both paths can have the same capped grid, so also inspect the kernel mode. + sms = torch.cuda.get_device_properties(torch.cuda.current_device()).multi_processor_count + elements = 16 // torch.empty((), dtype=dtype).element_size() + threshold = sms * 160 * 256 * elements + cases = [(threshold - 1, 1, 3), (threshold, 1, 3), (threshold, 0, 0)] + tensors = [ + ( + torch.randn(count + go, device="cuda", dtype=dtype)[go:], + torch.randn(count + uo, device="cuda", dtype=dtype)[uo:], + ) + for count, go, uo in cases + ] + for gate, up in tensors: + silu_and_mul(gate, up) + torch.cuda.synchronize() + with torch.profiler.profile( + activities=[torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA] + ) as prof: + for gate, up in tensors: + silu_and_mul(gate, up) + torch.cuda.synchronize() + trace = tmp_path / "silu-dispatch.json" + prof.export_chrome_trace(str(trace)) + events = json.loads(trace.read_text())["traceEvents"] + kernels = sorted( + ( + event + for event in events + if event.get("cat") == "kernel" and "silu_and_mul_kernel" in event.get("name", "") + ), + key=lambda event: event["ts"], + ) + assert len(kernels) == len(cases) + + def blocks_for(work): + cap = 16384 if dtype == torch.float32 and work >= 16384 * 256 else 4096 + return min(math.ceil(work / 256), cap) + + scalar_blocks = blocks_for(threshold - 1) + vector_blocks = blocks_for(threshold // elements) + assert [event["args"]["grid"][0] for event in kernels] == [ + scalar_blocks, + vector_blocks, + vector_blocks, + ] + assert all(event["args"]["block"] == [256, 1, 1] for event in kernels) + assert all( + f"AccessMode){mode}>" in event["name"] + for event, mode in zip(kernels, (2, 1, 0), strict=True) + ) + + +@pytest.mark.gpu +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("offsets", [(0, 0), (1, 3)]) +def test_silu_and_mul_cuda_special_values_and_underflow(dtype, offsets): info = torch.finfo(dtype) values = [ float("nan"), @@ -184,6 +294,13 @@ def test_silu_and_mul_cuda_special_values_and_underflow(dtype): indexing="ij", ) gate, up = gate.contiguous(), up.contiguous() + if offsets != (0, 0): + count = gate.numel() + shifted_gate = torch.empty(count + offsets[0], device="cuda", dtype=dtype)[offsets[0] :] + shifted_up = torch.empty(count + offsets[1], device="cuda", dtype=dtype)[offsets[1] :] + shifted_gate.copy_(gate.flatten()) + shifted_up.copy_(up.flatten()) + gate, up = shifted_gate.view(gate.shape), shifted_up.view(up.shape) torch.testing.assert_close( silu_and_mul(gate, up), silu_and_mul_reference(gate, up), equal_nan=True ) @@ -210,11 +327,17 @@ def test_silu_and_mul_cuda_rejects_autograd_but_allows_inference(): @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): +@pytest.mark.parametrize("offsets", [(0, 0), (1, 3)]) +@pytest.mark.parametrize("large", [False, True]) +def test_silu_and_mul_cuda_current_stream_and_graph_replay(dtype, offsets, large): + sms = torch.cuda.get_device_properties(torch.cuda.current_device()).multi_processor_count + elements = 16 // torch.empty((), dtype=dtype).element_size() + threshold = sms * 160 * 256 * elements + count = threshold + 1 if large else min(4097, threshold - 1) stream = torch.cuda.Stream() with torch.cuda.stream(stream): - gate = torch.randn(4097, device="cuda", dtype=dtype) - up = torch.randn_like(gate) + gate = torch.randn(count + offsets[0], device="cuda", dtype=dtype)[offsets[0] :] + up = torch.randn(count + offsets[1], device="cuda", dtype=dtype)[offsets[1] :] for _ in range(3): silu_and_mul(gate, up) stream.synchronize() @@ -227,6 +350,25 @@ def test_silu_and_mul_cuda_current_stream_and_graph_replay(dtype): torch.testing.assert_close(out, silu_and_mul_reference(gate, up)) +@pytest.mark.gpu +def test_silu_and_mul_first_misaligned_call_during_graph_capture(): + # The device-property lookup used by dispatch must also work on first use. + script = textwrap.dedent(""" + import torch + from vllm_rlt.kernels.activation import silu_and_mul, silu_and_mul_reference + n = 4097 + gate = torch.randn(n + 1, device='cuda', dtype=torch.bfloat16)[1:] + up = torch.randn(n + 3, device='cuda', dtype=torch.bfloat16)[3:] + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + output = silu_and_mul(gate, up) + graph.replay() + torch.testing.assert_close(output, silu_and_mul_reference(gate, up), rtol=0, atol=0) + """) + subprocess.run([sys.executable, "-c", script], check=True, capture_output=True, timeout=60) + + @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): @@ -238,9 +380,10 @@ def test_silu_and_mul_cuda_fp32_streaming_grid_boundary(count): @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. +@pytest.mark.parametrize("offsets", [(0, 0), (1, 0), (1, 3), (7, 5)]) +def test_silu_and_mul_cuda_low_precision_storage_patterns(dtype, offsets): + # Exercise all gate bit patterns through aligned and independently shifted + # vector loads, including subnormal products at the scalar tail. 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) @@ -265,16 +408,23 @@ def test_silu_and_mul_cuda_low_precision_storage_patterns(dtype, offset): 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. + pattern_count = values.numel() * factors.numel() + sms = torch.cuda.get_device_properties(torch.cuda.current_device()).multi_processor_count + threshold = sms * 160 * 256 * (16 // values.element_size()) + count = max(pattern_count + 1, threshold + 1) if offsets != (0, 0) else pattern_count + 1 + gate = torch.empty(count + offsets[0], device="cuda", dtype=dtype)[offsets[0] :] + up = torch.empty(count + offsets[1], device="cuda", dtype=dtype)[offsets[1] :] + gate.zero_() + up.zero_() + gate[:pattern_count].copy_(values.repeat(factors.numel())) + up[:pattern_count].copy_(factors.repeat_interleave(values.numel())) + # A scalar tail multiplication produces 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) + expected = silu_and_mul_reference(gate[:pattern_count], up[:pattern_count]) + actual_patterns = actual[:pattern_count] + torch.testing.assert_close(actual_patterns, 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)) + assert torch.equal(actual_patterns[valid].view(torch.int16), expected[valid].view(torch.int16)) + expected_tail = silu_and_mul_reference(gate[-1:], up[-1:]) + torch.testing.assert_close(actual[-1:], expected_tail, rtol=0, atol=0)