diff --git a/CMakeLists.txt b/CMakeLists.txt index d59f56c00a..e67221978c 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -415,6 +415,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") "csrc/cuda_view.cu" "csrc/cuda_utils_kernels.cu" "csrc/activation_kernels.cu" + "csrc/muse_spark_mul_tanh_soft_cap_kernel.cu" "csrc/ngram_embedding_kernels.cu" "csrc/vocab_parallel_embedding_kernels.cu" "csrc/quantization/activation_kernels.cu" @@ -423,6 +424,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") "csrc/quantization/w8a8/fp8/per_token_group_quant.cu" "csrc/quantization/gptq/q_gemm.cu" "csrc/pos_encoding_kernels.cu" + "csrc/fused_muse_spark_qknorm_rope_kv_insert_kernel.cu" "csrc/fused_qknorm_rope_kernel.cu" "csrc/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu" "csrc/layernorm_kernels.cu" @@ -441,6 +443,11 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") "csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu" "csrc/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu") + if(VLLM_GPU_LANG STREQUAL "CUDA") + list(APPEND VLLM_STABLE_EXT_SRC + "csrc/fused_qknorm_rope_cache_kernel.cu") + endif() + if(VLLM_GPU_LANG STREQUAL "CUDA" AND DEFINED CMAKE_CUDA_COMPILER_VERSION AND CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 12.0) diff --git a/benchmarks/kernels/benchmark_muse_spark_qknorm_rope_kv_insert.py b/benchmarks/kernels/benchmark_muse_spark_qknorm_rope_kv_insert.py new file mode 100644 index 0000000000..115459d415 --- /dev/null +++ b/benchmarks/kernels/benchmark_muse_spark_qknorm_rope_kv_insert.py @@ -0,0 +1,326 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Benchmark the fused Muse-Spark QK-norm + RoPE + KV-insert kernel vs unfused. + +``tuned_us`` forces the cp.async schedule ``fused_qk_norm_rope`` only auto-selects +on SM 9.0; ``unfused_us`` is what production runs. +""" + +import argparse +import itertools + +import torch + +import vllm._custom_ops as ops + +HEAD_DIM = 64 +DEFAULT_TOKENS = [1, 8, 64, 256, 1024, 4096, 8192, 16384, 32768, 65536] + + +def make_cos_sin_cache(max_pos, rotary_dim, base, device): + inv_freq = 1.0 / ( + base + ** ( + torch.arange(0, rotary_dim, 2, dtype=torch.float32, device=device) + / rotary_dim + ) + ) + t = torch.arange(max_pos, dtype=torch.float32, device=device) + freqs = torch.einsum("i,j->ij", t, inv_freq) + return torch.cat((freqs.cos(), freqs.sin()), dim=-1) + + +class Case: + """Everything a single (num_tokens, config) measurement needs.""" + + def __init__(self, num_tokens, nq, nkv, dtype, kv_cache_dtype, device): + self.num_tokens = num_tokens + self.nq, self.nkv = nq, nkv + self.dtype = dtype + self.kv_cache_dtype = kv_cache_dtype + self.eps = 1e-5 + self.block_size = 128 + self.rotary_dim = HEAD_DIM + + self.qsz, self.kvsz = nq * HEAD_DIM, nkv * HEAD_DIM + self.row = self.qsz + 2 * self.kvsz + self.qsz # [q | k | v | o_gate] + self.qkv = torch.randn(num_tokens, self.row, dtype=dtype, device=device) + self.pristine = self.qkv.clone() + + max_pos = 32768 + self.cos_sin = make_cos_sin_cache(max_pos, self.rotary_dim, 500_000.0, device) + self.positions = torch.randint( + 0, max_pos, (num_tokens,), dtype=torch.int64, device=device + ) + + num_blocks = (num_tokens + self.block_size - 1) // self.block_size + 1 + storage = torch.uint8 if kv_cache_dtype != "auto" else dtype + self.kv_cache = torch.zeros( + num_blocks, + nkv, + self.block_size, + 2 * HEAD_DIM, + dtype=storage, + device=device, + ) + self.k_cache, self.v_cache = self.kv_cache.transpose(1, 2).split( + HEAD_DIM, dim=-1 + ) + self.slot_mapping = torch.randperm( + num_blocks * self.block_size, dtype=torch.int64, device=device + )[:num_tokens] + self.k_scale = torch.tensor(1.0, dtype=torch.float32, device=device) + self.v_scale = torch.tensor(1.0, dtype=torch.float32, device=device) + self.ones = torch.ones(HEAD_DIM, dtype=dtype, device=device) + self.q_fp8 = torch.empty( + num_tokens, self.qsz, dtype=torch.float8_e4m3fn, device=device + ) + self.q_scale = torch.tensor(0.37, dtype=torch.float32, device=device) + + def views(self): + q, k, v, _ = self.qkv.split([self.qsz, self.kvsz, self.kvsz, self.qsz], dim=-1) + return q, k, v + + def fused(self, use_rope=True, quant_query=False): + ops.fused_muse_spark_qknorm_rope_kv_insert( + self.qkv, + self.positions, + self.nq, + self.nkv, + HEAD_DIM, + self.eps, + cos_sin_cache=self.cos_sin if use_rope else None, + rotary_dim=self.rotary_dim if use_rope else 0, + slot_mapping=self.slot_mapping, + kv_cache=self.kv_cache, + block_size=self.block_size, + kv_cache_dtype=self.kv_cache_dtype, + k_scale=self.k_scale if self.kv_cache_dtype != "auto" else None, + v_scale=self.v_scale if self.kv_cache_dtype != "auto" else None, + q_fp8_out=self.q_fp8 if quant_query else None, + q_scale=self.q_scale if quant_query else None, + ) + + def separate_query_quant(self, use_rope=True): + self.fused(use_rope) + ops.scaled_fp8_quant(self.qkv[:, : self.qsz], scale=self.q_scale) + + def _insert(self): + _, k, v = self.views() + ops.reshape_and_cache_flash( + k.view(-1, self.nkv, HEAD_DIM), + v.view(-1, self.nkv, HEAD_DIM), + self.k_cache, + self.v_cache, + self.slot_mapping, + self.kv_cache_dtype, + self.k_scale, + self.v_scale, + ) + + def unfused_rope(self, token_heads_per_warp: int = -1): + ops.fused_qk_norm_rope( + self.qkv, + self.nq, + self.nkv, + self.nkv + self.nq, # v heads plus the o-gate + HEAD_DIM, + self.eps, + self.ones, + self.ones, + self.cos_sin, + False, + self.positions, + token_heads_per_warp, + ) + self._insert() + + def fused_bytes(self): + el = self.qkv.element_size() + cache_el = self.kv_cache.element_size() + read = (self.nq + 2 * self.nkv) * HEAD_DIM * el + write = (self.nq + self.nkv) * HEAD_DIM * el + cache = 2 * self.nkv * HEAD_DIM * cache_el + return self.num_tokens * (read + write + cache) + + +def _l2_bytes(device) -> int: + return torch.cuda.get_device_properties(device).L2_cache_size + + +def _peak_bandwidth(device) -> float: + """Measured large-copy bandwidth, in GB/s -- the practical ceiling.""" + n = 1 << 28 # 256 Mi elements = 512 MiB each way at 2 bytes + a = torch.empty(n, dtype=torch.bfloat16, device=device) + b = torch.empty_like(a) + for _ in range(3): + b.copy_(a) + torch.accelerator.synchronize() + start, end = torch.cuda.Event(True), torch.cuda.Event(True) + start.record() + iters = 20 + for _ in range(iters): + b.copy_(a) + end.record() + torch.accelerator.synchronize() + ms = start.elapsed_time(end) / iters + return (2 * a.numel() * a.element_size()) / (ms * 1e-3) / 1e9 + + +def time_eager(fn, warmup=20, iters=100) -> float: + for _ in range(warmup): + fn() + torch.accelerator.synchronize() + start, end = torch.cuda.Event(True), torch.cuda.Event(True) + start.record() + for _ in range(iters): + fn() + end.record() + torch.accelerator.synchronize() + return start.elapsed_time(end) / iters * 1e3 # us + + +def time_graph(fn, warmup=10, iters=50, inner=32) -> float: + """GPU time per call; ``inner`` calls per graph amortize the ~2 us replay.""" + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + for _ in range(warmup): + fn() + torch.cuda.current_stream().wait_stream(side) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + for _ in range(inner): + fn() + for _ in range(3): + graph.replay() + torch.accelerator.synchronize() + start, end = torch.cuda.Event(True), torch.cuda.Event(True) + start.record() + for _ in range(iters): + graph.replay() + end.record() + torch.accelerator.synchronize() + return start.elapsed_time(end) / (iters * inner) * 1e3 # us + + +def main(): + p = argparse.ArgumentParser() + p.add_argument("--nq", type=int, default=32, help="query heads per rank") + p.add_argument("--nkv", type=int, default=4, help="kv heads per rank") + p.add_argument("--kv-cache-dtype", default="auto", choices=["auto", "fp8"]) + p.add_argument("--dtype", default="bfloat16", choices=["bfloat16", "float16"]) + p.add_argument("--tokens", type=int, nargs="*", default=DEFAULT_TOKENS) + p.add_argument("--csv", default=None) + p.add_argument("--no-graph", action="store_true") + p.add_argument( + "--query-fp8", action="store_true", help="compare fused Q quantization" + ) + args = p.parse_args() + + device = "cuda" + dtype = getattr(torch, args.dtype) + if args.query_fp8: + if args.kv_cache_dtype != "fp8": + p.error("--query-fp8 requires --kv-cache-dtype fp8") + timer = time_eager if args.no_graph else time_graph + print(f"device: {torch.cuda.get_device_name()}; nq={args.nq}, nkv={args.nkv}") + print("layer,tokens,separate_us,fused_us,speedup") + rows = [] + for use_rope, nt in itertools.product([True, False], args.tokens): + case = Case(nt, args.nq, args.nkv, dtype, args.kv_cache_dtype, device) + separate_us = timer(lambda c=case, r=use_rope: c.separate_query_quant(r)) + case.qkv.copy_(case.pristine) + fused_us = timer(lambda c=case, r=use_rope: c.fused(r, quant_query=True)) + row = ( + "rope" if use_rope else "nope", + nt, + separate_us, + fused_us, + separate_us / fused_us, + ) + rows.append(row) + print(f"{row[0]},{nt},{separate_us:.3f},{fused_us:.3f},{row[4]:.3f}") + if args.csv: + import csv + + with open(args.csv, "w", newline="") as f: + writer = csv.writer(f) + writer.writerow( + ["layer", "tokens", "separate_us", "fused_us", "speedup"] + ) + writer.writerows(rows) + return + peak = _peak_bandwidth(device) + print( + f"device: {torch.cuda.get_device_name()} measured copy peak: {peak:.0f} GB/s" + ) + print( + f"config: nq={args.nq} nkv={args.nkv} head_dim={HEAD_DIM} " + f"dtype={args.dtype} kv_cache={args.kv_cache_dtype}\n" + ) + + modes = ["rope"] + l2 = _l2_bytes(device) + print( + f"rows marked * have a working set inside the {l2 / 2**20:.0f} MiB L2, " + "so their bandwidth is not an HBM number\n" + ) + header = ( + f"{'layer':>5} {'tokens':>7} {'fused_us':>9} {'unfused_us':>11} " + f"{'tuned_us':>9} {'speedup':>8} {'GB/s':>8} {'SOL%':>6}" + ) + rows = [] + for mode, use_graph in itertools.product(modes, [False, True]): + if use_graph and args.no_graph: + continue + label = "graph" if use_graph else "eager" + print(f"── {mode} layer, {label} ──") + print(header) + timer = time_graph if use_graph else time_eager + for nt in args.tokens: + case = Case(nt, args.nq, args.nkv, dtype, args.kv_cache_dtype, device) + use_rope = mode == "rope" + fused_us = timer(lambda c=case, r=use_rope: c.fused(r)) + baseline = case.unfused_rope + unfused_us = timer(baseline) + tuned_us = ( + min(timer(lambda c=case, t=t: c.unfused_rope(t)) for t in (4, 8)) + if use_rope + else unfused_us + ) + nbytes = case.fused_bytes() + gbps = nbytes / (fused_us * 1e-6) / 1e9 + sol = 100.0 * gbps / peak + resident = "*" if nbytes < l2 else " " + print( + f"{mode:>5} {nt:>7} {fused_us:>9.2f} {unfused_us:>11.2f} " + f"{tuned_us:>9.2f} {tuned_us / fused_us:>7.2f}x {gbps:>8.0f} " + f"{sol:>5.1f}%{resident}" + ) + rows.append((mode, label, nt, fused_us, unfused_us, tuned_us, gbps, sol)) + print() + + if args.csv: + import csv + + with open(args.csv, "w", newline="") as f: + w = csv.writer(f) + w.writerow( + [ + "layer", + "timing", + "tokens", + "fused_us", + "unfused_us", + "tuned_us", + "gbps", + "sol", + ] + ) + w.writerows(rows) + print(f"wrote {args.csv}") + + +if __name__ == "__main__": + main() diff --git a/csrc/fused_muse_spark_qknorm_rope_kv_insert_kernel.cu b/csrc/fused_muse_spark_qknorm_rope_kv_insert_kernel.cu new file mode 100644 index 0000000000..c50bdbafbc --- /dev/null +++ b/csrc/fused_muse_spark_qknorm_rope_kv_insert_kernel.cu @@ -0,0 +1,752 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * SPDX-FileCopyrightText: Copyright contributors to the vLLM project + * + * Fused Muse Spark attention pre-processing: weightless QK RMSNorm, GPT-J + * interleaved RoPE (skipped on NoPE layers, ``cos_sin_cache == None``) and + * the K/V insert of ``reshape_and_cache_flash``, over ``[q | k | v | o_gate]`` + * rows. + * + * A group of ``head_dim * sizeof(scalar_t) / 16`` lanes owns one head (8 lanes + * for bf16 head_dim=64), so each lane moves one 16-byte vector, the rotated + * pair (2i, 2i+1) is lane-local, and the norm reduction is 3 shuffles. + * + * Slot enumeration per token: + * [0, nq) Q heads -> norm (+RoPE), in place or to + * ``q_fp8_out`` + * [nq, nq+nkv) K heads -> norm (+RoPE), in place and into + * the K cache + * [nq+nkv, nq+2*nkv) V heads -> copied into the V cache (no norm) + * V slots exist only when a KV cache is bound. + * + * Norm and rotation run in fp32 and round once at the store. Both fp8 paths + * quantize after rounding to the model dtype, the same boundary + * ``reshape_and_cache_flash`` and ``scaled_fp8_quant`` use. + */ + +#include +#include +#include +#include + +#include "torch_utils.h" + +#include "core/cuda_compat.h" +#include "type_convert.cuh" +#include "core/vec_dtypes/dtype_fp8.cuh" +#include "dispatch_utils.h" + +#ifndef USE_ROCM + #include + #include "quantization/w8a8/fp8/nvidia/quant_utils.cuh" +#else + #include + #include "quantization/w8a8/fp8/amd/quant_utils.cuh" +#endif + +#ifndef FINAL_MASK + #ifdef USE_ROCM + #define FINAL_MASK 0xffffffffffffffffULL + #else + #define FINAL_MASK 0xffffffffu + #endif +#endif + +namespace vllm { +namespace muse_spark_fused_ops { + +namespace { +inline int getSMVersion() { + auto* props = get_device_prop(); + return props->major * 10 + props->minor; +} +inline int getSMCount() { + auto* props = get_device_prop(); + return props->multiProcessorCount; +} +} // namespace + +// 16 bytes per lane: the widest vector access, and the reason a head is owned +// by a group of lanes rather than a whole warp. +constexpr int kBytesPerLane = 16; + +#ifdef USE_ROCM +using WarpMask = uint64_t; +constexpr unsigned kWarpSize = 64; +#else +using WarpMask = unsigned; +constexpr unsigned kWarpSize = 32; +#endif + +// Slot -> (token, head) needs a division by a value only known at launch. The +// kernel is issue-bound, and a 32-bit hardware division is ~20 instructions of +// it, so the reciprocal is computed once on the host instead. +struct FastDiv { + uint32_t mul; + uint32_t shr; + int div; +}; + +inline FastDiv makeFastDiv(int denom) { + // The reciprocal form needs denom >= 2; denom == 1 is carried as a flag the + // device side selects on. Dividends stay below 2^31, which covers every slot + // index this kernel can form. + if (denom <= 1) return FastDiv{0u, 0u, 1}; + int lg = 31 - __builtin_clz(denom); + if ((1 << lg) < denom) lg++; + uint32_t const p = 31 + static_cast(lg); + uint64_t const m = ((1ull << p) + static_cast(denom) - 1) / + static_cast(denom); + return FastDiv{static_cast(m), p - 32, denom}; +} + +__device__ __forceinline__ void fastDivmod(int& quo, int& rem, int src, + FastDiv const fd) { + int const q = + static_cast(__umulhi(static_cast(src), fd.mul) >> fd.shr); + // A select, not a branch: fd.div is uniform across the launch. + quo = (fd.div == 1) ? src : q; + rem = src - quo * fd.div; +} + +// Sum across the LANES contiguous lanes that share one head. The group is +// LANES-aligned inside the warp, so an XOR butterfly over masks < LANES never +// leaves it. `group_mask` names exactly the group: whole groups are uniformly +// active or inactive, so a warp whose other groups bailed out still syncs. +template +__device__ __forceinline__ float groupReduceSum(float val, + WarpMask const group_mask) { +#pragma unroll + for (int mask = LANES >> 1; mask > 0; mask >>= 1) { + val += __shfl_xor_sync(group_mask, val, mask, 32); + } + return val; +} + +// 16 contiguous bytes -> kElems fp32 registers, keeping the raw vector so a +// pass-through slot (V) can be re-stored without a conversion round trip. +template +__device__ __forceinline__ uint4 loadElems(scalar_t const* __restrict__ src, + float (&elems)[kElems]) { + using Converter = vllm::_typeConvert; + using T2 = typename Converter::packed_hip_type; + uint4 raw = *reinterpret_cast(src); + auto const* p = reinterpret_cast(&raw); +#pragma unroll + for (int i = 0; i < kElems / 2; i++) { + float2 f2 = Converter::convert(p[i]); + elems[2 * i] = f2.x; + elems[2 * i + 1] = f2.y; + } + return raw; +} + +// The lane's slice of the rope table: kPairs cos and kPairs sin values, each +// run contiguous, so each is one vector load (16 bytes for the fp32 table, 8 +// for a model-dtype one) instead of kPairs scalar ones. rotary_dim being a +// multiple of 16 is what keeps both runs naturally aligned. +template +__device__ __forceinline__ void loadRopeVec(rope_t const* __restrict__ cos_ptr, + rope_t const* __restrict__ sin_ptr, + float (&c)[kPairs], + float (&s)[kPairs]) { + static_assert(kPairs == 4, "rope vector load assumes 4 pairs per lane"); + if constexpr (std::is_same_v) { + float4 const cv = VLLM_LDG(reinterpret_cast(cos_ptr)); + float4 const sv = VLLM_LDG(reinterpret_cast(sin_ptr)); + c[0] = cv.x, c[1] = cv.y, c[2] = cv.z, c[3] = cv.w; + s[0] = sv.x, s[1] = sv.y, s[2] = sv.z, s[3] = sv.w; + } else { + using Converter = vllm::_typeConvert; + using T2 = typename Converter::packed_hip_type; + uint2 const cv = VLLM_LDG(reinterpret_cast(cos_ptr)); + uint2 const sv = VLLM_LDG(reinterpret_cast(sin_ptr)); + auto const* cp = reinterpret_cast(&cv); + auto const* sp = reinterpret_cast(&sv); +#pragma unroll + for (int i = 0; i < kPairs / 2; i++) { + float2 const cf = Converter::convert(cp[i]); + float2 const sf = Converter::convert(sp[i]); + c[2 * i] = cf.x, c[2 * i + 1] = cf.y; + s[2 * i] = sf.x, s[2 * i + 1] = sf.y; + } + } +} + +// kElems fp32 registers -> the 16-byte scalar_t vector they round to. Returned +// rather than stored so the fp8 cache path can reuse the *rounded* values, +// which is what reshape_and_cache_flash quantizes from. +template +__device__ __forceinline__ uint4 packElems(float const (&elems)[kElems]) { + using Converter = vllm::_typeConvert; + using T2 = typename Converter::packed_hip_type; + uint4 raw; + auto* p = reinterpret_cast(&raw); +#pragma unroll + for (int i = 0; i < kElems / 2; i++) { + p[i] = Converter::convert(make_float2(elems[2 * i], elems[2 * i + 1])); + } + return raw; +} + +// Scatter one lane's 16 bytes into the KV cache. kAuto stores the model dtype +// verbatim; the fp8 dtypes quantize the value *after* it has been rounded to +// scalar_t, which is the boundary reshape_and_cache_flash uses. +// +// `inv_scale` is taken once per thread: per-element fp32 divisions would cost +// more than the rest of the kernel. The extra rounding is 2^-24 relative. +template +__device__ __forceinline__ void storeCache(cache_t* __restrict__ dst, + uint4 const& packed, + float const inv_scale) { + if constexpr (kv_dt == Fp8KVCacheDataType::kAuto) { + *reinterpret_cast(dst) = packed; + } else { + using Converter = vllm::_typeConvert; + using HipT = typename Converter::hip_type; + auto const* rounded = reinterpret_cast(&packed); +#ifndef USE_ROCM + constexpr __nv_fp8_interpretation_t kInterp = + (kv_dt == Fp8KVCacheDataType::kFp8E4M3) ? __NV_E4M3 : __NV_E5M2; + __nv_fp8x2_storage_t out[kElems / 2]; + #pragma unroll + for (int i = 0; i < kElems / 2; i++) { + float2 v = + make_float2(Converter::convert(rounded[2 * i]) * inv_scale, + Converter::convert(rounded[2 * i + 1]) * inv_scale); + out[i] = __nv_cvt_float2_to_fp8x2(v, __NV_SATFINITE, kInterp); + } + // kElems fp8 bytes: 8 for the 2-byte dtypes this kernel serves. + static_assert(kElems == 8, "fp8 store assumes 8 elements per lane"); + *reinterpret_cast(dst) = *reinterpret_cast(out); +#else + #pragma unroll + for (int i = 0; i < kElems; i++) { + float const v = Converter::convert(rounded[i]) * inv_scale; + dst[i] = fp8::scaled_convert(v, 1.0f); + } +#endif + } +} + +// Grid is 1D over (token, head-chunk) pairs. A chunk is HEADS_PER_GROUP +// consecutive head slots of one token, served by kLanes contiguous lanes; the +// whole chunk shares the token's position and its slice of the rope table, so +// the slot division, the position load and the cos/sin load are paid once per +// chunk. `slots_per_token % HEADS_PER_GROUP == 0` keeps a chunk inside one +// token. `rope_t` is the cos/sin cache dtype, independent of the qkv dtype. +template +__global__ +__launch_bounds__(256, 6) void fusedMuseSparkQKNormRopeKVInsertKernel( + scalar_t* __restrict__ qkv, // [N, qkv_stride] in/out + rope_t const* __restrict__ cos_sin_cache, // [max_pos, rotary_dim] + int64_t const* __restrict__ positions, // [N] i64 + int64_t const* __restrict__ slot_mapping, // [N] i64 or nullptr + cache_t* __restrict__ kv_cache, // [nb, nkv, bs, 2*HEAD_DIM] + float const* __restrict__ k_scale, // [1] or nullptr + float const* __restrict__ v_scale, // [1] or nullptr + uint8_t* __restrict__ q_fp8_out, // [N, nq * HEAD_DIM] or nullptr + float const* __restrict__ q_scale, // [1] or nullptr + float const eps, int const rotary_dim, int const num_tokens, int const nq, + int const nkv, FastDiv const chunks_div, int64_t const qkv_stride, + FastDiv const block_div, int64_t const kv_s_block, int64_t const kv_s_head, + int64_t const kv_s_token) { +#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM) + // _typeConvert is unavailable on pre-Ampere; this kernel only ever + // runs with bf16/fp16 inputs. Discard the bf16 body there. + if constexpr (std::is_same_v) { + return; + } else { +#endif + constexpr int kElems = kBytesPerLane / sizeof(scalar_t); + constexpr int kLanes = HEAD_DIM / kElems; + constexpr int kPairs = kElems / 2; + constexpr int kH = HEADS_PER_GROUP; + static_assert(kLanes > 0 && kLanes <= 32 && (32 % kLanes) == 0, + "head_dim must give a power-of-two lane group <= 32"); + + int const tid = blockIdx.x * blockDim.x + threadIdx.x; + int const groupIdx = tid / kLanes; + int const dimBase = (tid % kLanes) * kElems; + + int tokenIdx, chunk; + fastDivmod(tokenIdx, chunk, groupIdx, chunks_div); + int const firstSlot = chunk * kH; + + // Out-of-range groups must not exit before the reductions below: the XOR + // butterfly needs every lane the mask names, and the mask covers the whole + // group. They ride along and skip every memory access instead. + bool const active = tokenIdx < num_tokens; + + constexpr WarpMask kGroupBits = + (kLanes == 32) ? static_cast(0xffffffffu) + : ((WarpMask{1} << kLanes) - WarpMask{1}); + unsigned const group_shift = + (threadIdx.x & (kWarpSize - 1)) & ~static_cast(kLanes - 1); + WarpMask const group_mask = kGroupBits << group_shift; + + // PDL: wait for the qkv-projection GEMM that produces `qkv`. No-op when + // PDL is not enabled on the launch. +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + cudaGridDependencySynchronize(); +#endif + + // Q, K and V are contiguous head slots at the head of the row, so the slot + // index doubles as the head offset. The trailing o-gate is skipped by the + // row stride alone. + scalar_t* const base = qkv + static_cast(tokenIdx) * qkv_stride + + static_cast(firstSlot) * HEAD_DIM + dimBase; + + // Issue every load in the chunk before touching any of it, so the group has + // kH loads in flight instead of one. An inactive group leaves these unset; + // nothing downstream reads them, since every store is gated on `active`. + uint4 raw[kH]; + if (active) { +#pragma unroll + for (int h = 0; h < kH; h++) { + raw[h] = *reinterpret_cast(base + h * HEAD_DIM); + } + } + + float cosv[kPairs], sinv[kPairs]; + bool const laneRotates = kDoRope && dimBase < rotary_dim; + if constexpr (kDoRope) { + if (active && laneRotates) { + int64_t const pos = positions[tokenIdx]; + rope_t const* cache_ptr = cos_sin_cache + pos * rotary_dim; + loadRopeVec(cache_ptr + dimBase / 2, + cache_ptr + rotary_dim / 2 + dimBase / 2, + cosv, sinv); + } + } + + int64_t sm = -1; + if constexpr (kInsertKV) { + if (active) sm = slot_mapping[tokenIdx]; + } + int cache_block = 0, cache_tok = 0; + float k_inv = 1.0f, v_inv = 1.0f; + float const q_inv = q_fp8_out ? 1.0f / *q_scale : 1.0f; + if constexpr (kInsertKV) { + if (sm >= 0) { + fastDivmod(cache_block, cache_tok, static_cast(sm), block_div); + } + if constexpr (kv_dt != Fp8KVCacheDataType::kAuto) { + k_inv = 1.0f / *k_scale; + v_inv = 1.0f / *v_scale; + } + } + +#pragma unroll + for (int h = 0; h < kH; h++) { + int const slot = firstSlot + h; + bool const isV = kInsertKV && slot >= nq + nkv; + bool const isK = slot >= nq && slot < nq + nkv; + bool const isQK = active && !isV; + + float elems[kElems]; + uint4 packed = raw[h]; + if (isQK) { + using Converter = vllm::_typeConvert; + using T2 = typename Converter::packed_hip_type; + auto const* p = reinterpret_cast(&raw[h]); +#pragma unroll + for (int i = 0; i < kElems / 2; i++) { + float2 const f2 = Converter::convert(p[i]); + elems[2 * i] = f2.x; + elems[2 * i + 1] = f2.y; + } + + float sumsq = 0.0f; +#pragma unroll + for (int i = 0; i < kElems; i++) sumsq += elems[i] * elems[i]; + sumsq = groupReduceSum(sumsq, group_mask); + float const rms_rcp = + rsqrtf(sumsq / static_cast(HEAD_DIM) + eps); +#pragma unroll + for (int i = 0; i < kElems; i++) elems[i] *= rms_rcp; + + if constexpr (kDoRope) { + if (laneRotates) { +#pragma unroll + for (int i = 0; i < kPairs; i++) { + float const v0 = elems[2 * i]; + float const v1 = elems[2 * i + 1]; + elems[2 * i] = v0 * cosv[i] - v1 * sinv[i]; + elems[2 * i + 1] = v0 * sinv[i] + v1 * cosv[i]; + } + } + } + // Rounding to scalar_t happens exactly once, here, and both the + // Q output and the fp8 cache quantization consume the same + // rounded value. V passes through with the vector it was loaded with. + packed = packElems(elems); + if (q_fp8_out && slot < nq) { + int64_t const q_off = + (static_cast(tokenIdx) * nq + slot) * HEAD_DIM + dimBase; + storeCache( + q_fp8_out + q_off, packed, q_inv); + } else { + *reinterpret_cast(base + h * HEAD_DIM) = packed; + } + } + + if constexpr (kInsertKV) { + if (sm >= 0 && (isK || isV)) { // sm < 0 marks a padded token + int const head = isK ? (slot - nq) : (slot - nq - nkv); + int64_t const off = static_cast(cache_block) * kv_s_block + + head * kv_s_head + + static_cast(cache_tok) * kv_s_token + + (isK ? 0 : HEAD_DIM) + dimBase; + storeCache(kv_cache + off, packed, + isK ? k_inv : v_inv); + } + } + } + + // PDL: let a dependent successor launch early. +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + cudaTriggerProgrammaticLaunchCompletion(); +#endif +#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM) + } +#endif +} + +template +void launchFusedMuseSparkTyped( + scalar_t* qkv, rope_t const* cos_sin_cache, int64_t const* positions, + int64_t const* slot_mapping, cache_t* kv_cache, float const* k_scale, + float const* v_scale, uint8_t* q_fp8_out, float const* q_scale, + float const eps, int const rotary_dim, int const num_tokens, int const nq, + int const nkv, int64_t const qkv_stride, int const block_size, + int64_t const kv_s_block, int64_t const kv_s_head, int64_t const kv_s_token, + bool const do_rope, bool const insert_kv, cudaStream_t stream) { + FastDiv const block_div = makeFastDiv(block_size); + constexpr int kElems = kBytesPerLane / sizeof(scalar_t); + constexpr int kLanes = HEAD_DIM / kElems; + + int const slots_per_token = nq + nkv + (insert_kv ? nkv : 0); + int const chunks_per_token = slots_per_token / HEADS_PER_GROUP; + FastDiv const chunks_div = makeFastDiv(chunks_per_token); + int64_t const total_lanes = + static_cast(num_tokens) * chunks_per_token * kLanes; + if (total_lanes == 0) return; + + constexpr int kBlockSize = 256; + int const grid = + static_cast((total_lanes + kBlockSize - 1) / kBlockSize); + +#ifndef USE_ROCM + // PDL wherever the hardware supports it (SM90+); a plain launch elsewhere. + static int const sm_version = getSMVersion(); + cudaLaunchConfig_t config; + config.gridDim = dim3(grid); + config.blockDim = dim3(kBlockSize); + config.dynamicSmemBytes = 0; + config.stream = stream; + cudaLaunchAttribute attrs[1]; + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = 1; + config.attrs = attrs; + config.numAttrs = (sm_version >= 90) ? 1 : 0; + + #define LAUNCH_MUSE(DO_ROPE, INSERT) \ + cudaLaunchKernelEx(&config, \ + fusedMuseSparkQKNormRopeKVInsertKernel< \ + scalar_t, cache_t, kv_dt, rope_t, HEAD_DIM, \ + HEADS_PER_GROUP, DO_ROPE, INSERT>, \ + qkv, cos_sin_cache, positions, slot_mapping, kv_cache, \ + k_scale, v_scale, q_fp8_out, q_scale, eps, rotary_dim, \ + num_tokens, nq, nkv, chunks_div, qkv_stride, block_div, \ + kv_s_block, kv_s_head, kv_s_token) +#else + // clang-format off + #define LAUNCH_MUSE(DO_ROPE, INSERT) \ + fusedMuseSparkQKNormRopeKVInsertKernel< \ + scalar_t, cache_t, kv_dt, rope_t, HEAD_DIM, HEADS_PER_GROUP, DO_ROPE, \ + INSERT><<>>( \ + qkv, cos_sin_cache, positions, slot_mapping, kv_cache, k_scale, \ + v_scale, q_fp8_out, q_scale, eps, rotary_dim, num_tokens, nq, nkv, \ + chunks_div, qkv_stride, block_div, kv_s_block, kv_s_head, kv_s_token) + // clang-format on +#endif + + if (do_rope) { + if (insert_kv) { + LAUNCH_MUSE(true, true); + } else { + LAUNCH_MUSE(true, false); + } + } else { + if (insert_kv) { + LAUNCH_MUSE(false, true); + } else { + LAUNCH_MUSE(false, false); + } + } +#undef LAUNCH_MUSE +} + +template +void launchFusedMuseSpark( + scalar_t* qkv, void const* cos_sin_cache, bool const rope_is_fp32, + int64_t const* positions, int64_t const* slot_mapping, cache_t* kv_cache, + float const* k_scale, float const* v_scale, uint8_t* q_fp8_out, + float const* q_scale, float const eps, int const rotary_dim, + int const num_tokens, int const nq, int const nkv, int const head_dim, + int64_t const qkv_stride, int const block_size, int64_t const kv_s_block, + int64_t const kv_s_head, int64_t const kv_s_token, bool const do_rope, + bool const insert_kv, cudaStream_t stream) { + // The group width must divide the slot count. Widen only while the grid + // still covers a full wave; decode-sized batches keep the narrow schedule. + int const slots_per_token = nq + nkv + (insert_kv ? nkv : 0); + constexpr int kElemsSel = kBytesPerLane / sizeof(scalar_t); + constexpr int kBlockSel = 256; + int const lanes_per_head = head_dim / kElemsSel; + int64_t const full_wave = + static_cast(getSMCount()) * (2048 / kBlockSel); + int heads_per_group = 1; + for (int h : {4, 2}) { + if (slots_per_token % h != 0) continue; + int64_t const blocks = (static_cast(num_tokens) * + (slots_per_token / h) * lanes_per_head + + kBlockSel - 1) / + kBlockSel; + if (blocks >= full_wave) { + heads_per_group = h; + break; + } + } + + // TP4/TP8 amortize head addressing before filling a full CTA wave. RoPE + // benefits earlier; FP8 conversion keeps wider groups expensive for longer. + if (insert_kv && head_dim == 64 && + ((nq == 16 && nkv == 2) || (nq == 32 && nkv == 4))) { + int64_t const query_heads = static_cast(num_tokens) * nq; + if (do_rope && query_heads >= 24576) { + heads_per_group = std::max(heads_per_group, 2); + if constexpr (kv_dt == Fp8KVCacheDataType::kAuto) { + if (query_heads >= 49152) heads_per_group = 4; + } + } else if (!do_rope && query_heads >= 32768) { + if constexpr (kv_dt == Fp8KVCacheDataType::kAuto) { + heads_per_group = std::max(heads_per_group, 2); + } + } + } + +#define DISPATCH_ARGS(RT, HD, HPG) \ + launchFusedMuseSparkTyped( \ + qkv, reinterpret_cast(cos_sin_cache), positions, \ + slot_mapping, kv_cache, k_scale, v_scale, q_fp8_out, q_scale, eps, \ + rotary_dim, num_tokens, nq, nkv, qkv_stride, block_size, kv_s_block, \ + kv_s_head, kv_s_token, do_rope, insert_kv, stream) + +#define DISPATCH_HPG(RT, HD) \ + switch (heads_per_group) { \ + case 4: \ + DISPATCH_ARGS(RT, HD, 4); \ + break; \ + case 2: \ + DISPATCH_ARGS(RT, HD, 2); \ + break; \ + default: \ + DISPATCH_ARGS(RT, HD, 1); \ + break; \ + } + + // A NoPE launch never reads the table, so it can share the fp32 branch. +#define DISPATCH_ROPE_T(HD) \ + if (!do_rope || rope_is_fp32) { \ + DISPATCH_HPG(float, HD); \ + } else { \ + DISPATCH_HPG(scalar_t, HD); \ + } + + switch (head_dim) { + case 64: + DISPATCH_ROPE_T(64); + break; + case 128: + DISPATCH_ROPE_T(128); + break; + case 256: + DISPATCH_ROPE_T(256); + break; + default: + STD_TORCH_CHECK(false, "unsupported head_dim: ", head_dim); + } +#undef DISPATCH_ROPE_T +#undef DISPATCH_HPG +#undef DISPATCH_ARGS +} + +} // namespace muse_spark_fused_ops +} // namespace vllm + +// clang-format off +#define CALL_FUSED_MUSE_SPARK(_RAW_T, CACHE_T, KV_DTYPE) \ + vllm::muse_spark_fused_ops::launchFusedMuseSpark( \ + reinterpret_cast(qkv.data_ptr()), \ + do_rope ? cos_sin_cache->data_ptr() : nullptr, rope_is_fp32, \ + reinterpret_cast(positions.data_ptr()), \ + insert_kv ? reinterpret_cast(slot_mapping->data_ptr()) \ + : nullptr, \ + insert_kv ? reinterpret_cast(kv_cache->data_ptr()) : nullptr, \ + needs_scales ? reinterpret_cast(k_scale->data_ptr()) \ + : nullptr, \ + needs_scales ? reinterpret_cast(v_scale->data_ptr()) \ + : nullptr, \ + q_fp8_out ? reinterpret_cast(q_fp8_out->data_ptr()) : nullptr, \ + q_scale ? reinterpret_cast(q_scale->data_ptr()) : nullptr, \ + static_cast(eps), static_cast(rotary_dim), num_tokens, nq, \ + nkv, hd, qkv_stride, static_cast(block_size), kv_s_block, \ + kv_s_head, kv_s_token, do_rope, insert_kv, stream) +// clang-format on + +void fused_muse_spark_qknorm_rope_kv_insert( + torch::stable::Tensor& qkv, // [N, (nq + 2*nkv + gate)*hd] + torch::stable::Tensor const& positions, // [N] i64 + int64_t num_heads, int64_t num_kv_heads, int64_t head_dim, double eps, + std::optional cos_sin_cache, // None => NoPE layer + int64_t rotary_dim, + std::optional slot_mapping, // [N] i64 + std::optional kv_cache, // [nb, nkv, bs, 2*hd] + int64_t block_size, const std::string& kv_cache_dtype, + std::optional k_scale, // [1] fp32, fp8 cache only + std::optional v_scale, + std::optional q_fp8_out, + std::optional q_scale) { + STD_TORCH_CHECK(qkv.is_cuda() && qkv.dim() == 2 && qkv.stride(1) == 1, + "qkv must be a 2-D CUDA tensor with a contiguous row"); + STD_TORCH_CHECK( + qkv.scalar_type() == torch::headeronly::ScalarType::Half || + qkv.scalar_type() == torch::headeronly::ScalarType::BFloat16, + "qkv must be float16 or bfloat16"); + STD_TORCH_CHECK( + positions.is_cuda() && + positions.scalar_type() == torch::headeronly::ScalarType::Long, + "positions must be int64 CUDA"); + + int const num_tokens = static_cast(qkv.size(0)); + int const nq = static_cast(num_heads); + int const nkv = static_cast(num_kv_heads); + int const hd = static_cast(head_dim); + int64_t const qkv_stride = qkv.stride(0); + + STD_TORCH_CHECK(q_fp8_out.has_value() == q_scale.has_value(), + "q_fp8_out and q_scale must be provided together"); + if (q_fp8_out) { + STD_TORCH_CHECK( + q_fp8_out->is_cuda() && + q_fp8_out->get_device_index() == qkv.get_device_index() && + q_fp8_out->scalar_type() == + torch::headeronly::ScalarType::Float8_e4m3fn && + q_fp8_out->dim() == 2 && q_fp8_out->is_contiguous() && + q_fp8_out->size(0) == num_tokens && + q_fp8_out->size(1) == static_cast(nq) * hd, + "q_fp8_out must be contiguous E4M3 [num_tokens, num_heads * head_dim] " + "on the qkv device"); + STD_TORCH_CHECK( + q_scale->is_cuda() && + q_scale->get_device_index() == qkv.get_device_index() && + q_scale->scalar_type() == torch::headeronly::ScalarType::Float && + q_scale->numel() == 1, + "q_scale must be a single fp32 value on the qkv device"); + } + + STD_TORCH_CHECK(hd == 64 || hd == 128 || hd == 256, + "head_dim must be 64, 128 or 256"); + STD_TORCH_CHECK(qkv.size(1) >= static_cast(nq + 2 * nkv) * hd, + "qkv row must hold at least [q|k|v]"); + // A lane moves 16 bytes, so every row start and every head start must land + // on a 16-byte boundary. Every supported head_dim is a multiple of 8 at 2 + // bytes, so heads are fine; the row stride is the one that can break it. + STD_TORCH_CHECK(qkv_stride % 8 == 0, + "qkv row stride must be a multiple of 8 elements, got ", + qkv_stride); + + bool const do_rope = cos_sin_cache.has_value(); + bool rope_is_fp32 = true; + if (do_rope) { + STD_TORCH_CHECK(cos_sin_cache->is_cuda() && cos_sin_cache->dim() == 2 && + cos_sin_cache->is_contiguous(), + "cos_sin_cache must be a contiguous 2-D CUDA tensor"); + STD_TORCH_CHECK(cos_sin_cache->size(1) == rotary_dim, + "cos_sin_cache shape must be [max_pos, rotary_dim]"); + STD_TORCH_CHECK( + cos_sin_cache->scalar_type() == torch::headeronly::ScalarType::Float || + cos_sin_cache->scalar_type() == qkv.scalar_type(), + "cos_sin_cache must be float32 or match the qkv dtype"); + rope_is_fp32 = + cos_sin_cache->scalar_type() == torch::headeronly::ScalarType::Float; + // A lane owns 16 bytes = 8 elements of the head and must be wholly inside + // or wholly outside the rotated prefix. + STD_TORCH_CHECK(rotary_dim > 0 && rotary_dim <= hd && rotary_dim % 16 == 0, + "rotary_dim must be a positive multiple of 16 and <= " + "head_dim"); + } + + vllm::Fp8KVCacheDataType const kv_dt = + vllm::get_fp8_kv_cache_data_type(kv_cache_dtype); + bool const insert_kv = kv_cache.has_value() && kv_cache->numel() > 0; + bool const needs_scales = + insert_kv && kv_dt != vllm::Fp8KVCacheDataType::kAuto; + + int64_t kv_s_block = 0, kv_s_head = 0, kv_s_token = 0; + if (insert_kv) { + STD_TORCH_CHECK( + slot_mapping.has_value() && slot_mapping->is_cuda() && + slot_mapping->scalar_type() == torch::headeronly::ScalarType::Long, + "insert mode requires an int64 CUDA slot_mapping"); + STD_TORCH_CHECK(slot_mapping->numel() >= num_tokens, + "slot_mapping must cover every token"); + STD_TORCH_CHECK(kv_cache->dim() == 4 && kv_cache->stride(3) == 1, + "kv_cache must be [nb, nkv, bs, 2*head_dim] with a " + "contiguous content dim"); + STD_TORCH_CHECK(kv_cache->size(3) == 2 * hd, + "kv_cache content dim must be 2*head_dim"); + STD_TORCH_CHECK(kv_cache->size(1) == nkv, + "kv_cache head dim must equal num_kv_heads"); + STD_TORCH_CHECK(kv_cache->size(2) == block_size, + "kv_cache token dim must equal block_size"); + if (kv_dt == vllm::Fp8KVCacheDataType::kAuto) { + STD_TORCH_CHECK(kv_cache->scalar_type() == qkv.scalar_type(), + "auto kv_cache dtype must match qkv"); + } else { + STD_TORCH_CHECK( + kv_cache->scalar_type() == torch::headeronly::ScalarType::Byte || + kv_cache->scalar_type() == + torch::headeronly::ScalarType::Float8_e4m3fn, + "fp8 kv_cache must use uint8 or float8_e4m3fn storage"); + STD_TORCH_CHECK(k_scale.has_value() && v_scale.has_value(), + "fp8 kv cache requires k_scale and v_scale"); + STD_TORCH_CHECK( + k_scale->numel() == 1 && v_scale->numel() == 1 && + k_scale->scalar_type() == torch::headeronly::ScalarType::Float && + v_scale->scalar_type() == torch::headeronly::ScalarType::Float, + "k_scale and v_scale must be single fp32 values"); + } + kv_s_block = kv_cache->stride(0); + kv_s_head = kv_cache->stride(1); + kv_s_token = kv_cache->stride(2); + } + + const torch::stable::accelerator::DeviceGuard device_guard( + qkv.get_device_index()); + auto stream = get_current_cuda_stream(qkv.get_device_index()); + + VLLM_STABLE_DISPATCH_HALF_TYPES( + qkv.scalar_type(), "fused_muse_spark_qknorm_rope_kv_insert", [&] { + using st = scalar_t; + DISPATCH_BY_KV_CACHE_DTYPE(qkv.scalar_type(), kv_cache_dtype, + CALL_FUSED_MUSE_SPARK); + }); +} + +#undef CALL_FUSED_MUSE_SPARK diff --git a/csrc/fused_qknorm_rope_cache_kernel.cu b/csrc/fused_qknorm_rope_cache_kernel.cu new file mode 100644 index 0000000000..a5ad757f50 --- /dev/null +++ b/csrc/fused_qknorm_rope_cache_kernel.cu @@ -0,0 +1,567 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// Step5 fused QK norm, RoPE, and KV-cache update kernel. + +// This kernel is built into `_C_stable_libtorch` and uses the libtorch stable +// ABI (torch::stable::Tensor / STD_TORCH_CHECK) instead of ATen and c10. +#include "torch_utils.h" +#include "dispatch_utils.h" +#include "type_convert.cuh" +#include "core/vec_dtypes/dtype_fp8.cuh" +#include "quantization/w8a8/fp8/nvidia/quant_utils.cuh" + +#include + +#include + +#include +#include + +namespace vllm { + +namespace { + +constexpr int kWarpSize = 32; +constexpr int kThreadsPerBlock = 128; + +template +__global__ void qknorm_rope_cache_kernel( + const void* __restrict__ qkv_void, void* __restrict__ q_out_void, + void* __restrict__ k_out_void, void* __restrict__ v_out_void, + const void* __restrict__ q_weight_void, + const void* __restrict__ k_weight_void, const void* __restrict__ cos_void, + const void* __restrict__ sin_void, const int64_t* __restrict__ positions, + const int64_t* __restrict__ slot_mapping, void* __restrict__ k_cache_void, + void* __restrict__ v_cache_void, int64_t qkv_token_stride, + int64_t q_token_stride, int64_t k_token_stride, int64_t v_token_stride, + int64_t cos_token_stride, int64_t sin_token_stride, + int64_t cache_block_stride, int64_t cache_page_stride, + int64_t cache_head_stride, int num_tokens, int num_slots, int num_q_heads, + int num_kv_heads, int rotary_pairs, int block_size, float epsilon, + int block_size_shift, float norm_weight_bias, const float* k_scale, + const float* v_scale) { +#if !defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800 + if constexpr (std::is_same_v || + std::is_same_v || + std::is_same_v) { + return; + } else { +#endif + using ScalarConverter = _typeConvert; + using WeightConverter = _typeConvert; + using RopeConverter = _typeConvert; + using Scalar = typename ScalarConverter::hip_type; + using Weight = typename WeightConverter::hip_type; + using Rope = typename RopeConverter::hip_type; + + // Preserve the original per-head reduction and activation rounding. + static_assert(!WIDE_192_ROTARY || HEAD_DIM == 192); + constexpr int kThreadsPerHead = + WIDE_192_ROTARY ? 32 + : (HEAD_DIM <= 64 ? 8 : (HEAD_DIM <= 192 ? 16 : 32)); + constexpr int kValuesPerVector = + WIDE_192_ROTARY ? 8 : (HEAD_DIM == 192 ? 4 : 8); + constexpr int kValuesPerTile = kThreadsPerHead * kValuesPerVector; + constexpr int kNumVectors = + (HEAD_DIM + kValuesPerTile - 1) / kValuesPerTile; + constexpr int kValuesPerThread = kValuesPerVector * kNumVectors; + constexpr int kTileSize = kValuesPerTile * kNumVectors; + constexpr int kHeadsPerBlock = kThreadsPerBlock / kThreadsPerHead; + + const auto* qkv = static_cast(qkv_void); + auto* q_out = static_cast(q_out_void); + auto* k_out = static_cast(k_out_void); + auto* v_out = static_cast(v_out_void); + const auto* q_weight = static_cast(q_weight_void); + const auto* k_weight = static_cast(k_weight_void); + const auto* cos = static_cast(cos_void); + const auto* sin = static_cast(sin_void); + auto* k_cache = static_cast(k_cache_void); + auto* v_cache = static_cast(v_cache_void); + + const int head_in_block = threadIdx.x / kThreadsPerHead; + const int lane = threadIdx.x % kThreadsPerHead; + const int total_heads = num_q_heads + 2 * num_kv_heads; + int token; + int local_head; + const int64_t global_head = + static_cast(blockIdx.x) * kHeadsPerBlock + head_in_block; + token = global_head / total_heads; + // Reuse the quotient instead of spelling this as global_head % + // total_heads. nvcc otherwise emits a second runtime integer-division + // sequence for the flattened token/head mapping. + local_head = static_cast(global_head - + static_cast(token) * total_heads); + + if (token >= num_tokens) return; + + const bool is_q = local_head < num_q_heads; + const bool is_k = !is_q && local_head < num_q_heads + num_kv_heads; + const int head = is_q ? local_head + : (is_k ? local_head - num_q_heads + : local_head - num_q_heads - num_kv_heads); + + int64_t input_offset = token * qkv_token_stride; + Scalar* output; + int64_t output_offset; + if (is_q) { + input_offset += static_cast(head) * HEAD_DIM; + output = q_out; + output_offset = + token * q_token_stride + static_cast(head) * HEAD_DIM; + } else if (is_k) { + input_offset += static_cast(num_q_heads + head) * HEAD_DIM; + output = k_out; + output_offset = + token * k_token_stride + static_cast(head) * HEAD_DIM; + } else { + input_offset += + static_cast(num_q_heads + num_kv_heads + head) * HEAD_DIM; + output = v_out; + output_offset = + token * v_token_stride + static_cast(head) * HEAD_DIM; + } + + bool write_cache = false; + int64_t cache_offset = 0; + if (!is_q && token < num_slots) { + const int64_t slot = slot_mapping[token]; + if (slot >= 0) { + int64_t block; + int64_t page; + if (block_size_shift >= 0) { + // vLLM production cache blocks are powers of two. Avoid the costly + // runtime int64 division on that path while retaining generic block + // sizes through the fallback below. + block = slot >> block_size_shift; + page = slot & (block_size - 1); + } else { + block = slot / block_size; + page = slot - block * block_size; + } + cache_offset = block * cache_block_stride + page * cache_page_stride + + static_cast(head) * cache_head_stride; + write_cache = true; + } + } + + if (!is_q && !is_k) { +#pragma unroll + for (int block = 0; block < kNumVectors; ++block) { +#pragma unroll + for (int element = 0; element < kValuesPerVector; ++element) { + const int dim = block * kThreadsPerHead * kValuesPerVector + + lane * kValuesPerVector + element; + if constexpr (kTileSize > HEAD_DIM) { + if (dim >= HEAD_DIM) continue; + } + const Scalar value = qkv[input_offset + dim]; + output[output_offset + dim] = value; + if (write_cache) { + if constexpr (FP8_CACHE) { + reinterpret_cast(v_cache_void)[cache_offset + dim] = + fp8::scaled_convert( + ScalarConverter::convert(value), *v_scale); + } else { + v_cache[cache_offset + dim] = value; + } + } + } + } + return; + } + + float values[kValuesPerThread]; + float sum_squares = 0.0f; +#pragma unroll + for (int block = 0; block < kNumVectors; ++block) { +#pragma unroll + for (int element = 0; element < kValuesPerVector; ++element) { + const int value_index = block * kValuesPerVector + element; + const int dim = block * kThreadsPerHead * kValuesPerVector + + lane * kValuesPerVector + element; + float value = 0.0f; + if constexpr (kTileSize > HEAD_DIM) { + if (dim < HEAD_DIM) { + value = ScalarConverter::convert(qkv[input_offset + dim]); + } + } else { + value = ScalarConverter::convert(qkv[input_offset + dim]); + } + values[value_index] = value; + // CuTeDSL lowers its fragment reduction to a sequential fma.rn + // chain. Keep this explicit: changing it changes RMSNorm bits. + sum_squares = __fmaf_rn(value, value, sum_squares); + } + } + + const int warp_lane = threadIdx.x % kWarpSize; + const int subwarp_base = (warp_lane / kThreadsPerHead) * kThreadsPerHead; + constexpr uint32_t kSubwarpBits = + kThreadsPerHead == kWarpSize ? 0xffffffffu + : ((uint32_t{1} << kThreadsPerHead) - 1); + const uint32_t subwarp_mask = kSubwarpBits << subwarp_base; +#pragma unroll + for (int offset = 1; offset < kThreadsPerHead; offset <<= 1) { + // CuTeDSL's warp_reduce uses ascending XOR offsets (1, 2, 4, ...). + sum_squares += + __shfl_xor_sync(subwarp_mask, sum_squares, offset, kWarpSize); + } + const float mean_square = + __fdiv_rn(sum_squares, static_cast(HEAD_DIM)); + const float inverse_rms = rsqrtf(mean_square + epsilon); + const Weight* weight = is_q ? q_weight : k_weight; + + // CuTeDSL writes RMSNorm results to activation-typed shared memory and + // reloads them for RoPE. This intermediate FP16/BF16 rounding boundary + // is intentional and is necessary for bitwise output parity. + __shared__ Scalar normalized[kHeadsPerBlock * HEAD_DIM]; + Scalar* normalized_head = normalized + head_in_block * HEAD_DIM; +#pragma unroll + for (int block = 0; block < kNumVectors; ++block) { +#pragma unroll + for (int element = 0; element < kValuesPerVector; ++element) { + const int value_index = block * kValuesPerVector + element; + const int dim = block * kThreadsPerHead * kValuesPerVector + + lane * kValuesPerVector + element; + if constexpr (kTileSize > HEAD_DIM) { + if (dim >= HEAD_DIM) continue; + } + const float x_hat = values[value_index] * inverse_rms; + const float norm_weight = + WeightConverter::convert(weight[dim]) + norm_weight_bias; + normalized_head[dim] = ScalarConverter::convert(x_hat * norm_weight); + } + } + + __syncwarp(subwarp_mask); + + // Apply NeoX-style RoPE. CuTeDSL assigns four rotary pairs to each + // participating lane and reloads activation-typed normalized values. + constexpr int kRotaryPairsPerThread = 4; + if (lane < rotary_pairs / kRotaryPairsPerThread) { + const int64_t position = positions[token]; + const Rope* token_cos = cos + position * cos_token_stride; + const Rope* token_sin = sin + position * sin_token_stride; +#pragma unroll + for (int element = 0; element < kRotaryPairsPerThread; ++element) { + const int pair = lane * kRotaryPairsPerThread + element; + const int dim0 = pair; + const int dim1 = pair + rotary_pairs; + const float value0 = ScalarConverter::convert(normalized_head[dim0]); + const float value1 = ScalarConverter::convert(normalized_head[dim1]); + const float cos_value = RopeConverter::convert(token_cos[pair]); + const float sin_value = RopeConverter::convert(token_sin[pair]); + + // Match CuTeDSL PTX exactly. The first half is mul + mul + sub; the + // second is one mul followed by one fused multiply-add. + const float rotated0 = __fsub_rn(__fmul_rn(value0, cos_value), + __fmul_rn(value1, sin_value)); + const float rotated1 = + __fmaf_rn(value0, sin_value, __fmul_rn(value1, cos_value)); + normalized_head[dim0] = ScalarConverter::convert(rotated0); + normalized_head[dim1] = ScalarConverter::convert(rotated1); + } + } + + __syncwarp(subwarp_mask); + +#pragma unroll + for (int block = 0; block < kNumVectors; ++block) { +#pragma unroll + for (int element = 0; element < kValuesPerVector; ++element) { + const int dim = block * kThreadsPerHead * kValuesPerVector + + lane * kValuesPerVector + element; + if constexpr (kTileSize > HEAD_DIM) { + if (dim >= HEAD_DIM) continue; + } + const Scalar value = normalized_head[dim]; + output[output_offset + dim] = value; + if (write_cache) { + if constexpr (FP8_CACHE) { + reinterpret_cast(k_cache_void)[cache_offset + dim] = + fp8::scaled_convert( + ScalarConverter::convert(value), *k_scale); + } else { + k_cache[cache_offset + dim] = value; + } + } + } + } + +#if !defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800 + } +#endif +} + +template +void launch_qknorm_rope_cache( + const torch::stable::Tensor& qkv, torch::stable::Tensor& q_out, + torch::stable::Tensor& k_out, torch::stable::Tensor& v_out, + const torch::stable::Tensor& q_weight, + const torch::stable::Tensor& k_weight, const torch::stable::Tensor& cos, + const torch::stable::Tensor& sin, const torch::stable::Tensor& positions, + const torch::stable::Tensor& slot_mapping, torch::stable::Tensor& k_cache, + torch::stable::Tensor& v_cache, int head_dim, int num_q_heads, + int num_kv_heads, int rotary_pairs, float epsilon, float norm_weight_bias, + cudaStream_t stream, const float* k_scale, const float* v_scale) { + const int total_heads = num_q_heads + 2 * num_kv_heads; + const int64_t total_work = qkv.size(0) * total_heads; + const int block_size = static_cast(k_cache.size(1)); + const int block_size_shift = + (block_size & (block_size - 1)) == 0 + ? __builtin_ctz(static_cast(block_size)) + : -1; + +#define LAUNCH_QKNORM_ROPE_CACHE(HEAD_DIM, WIDE_192_ROTARY) \ + do { \ + constexpr int threads_per_head = \ + WIDE_192_ROTARY ? 32 \ + : (HEAD_DIM <= 64 ? 8 : (HEAD_DIM <= 192 ? 16 : 32)); \ + constexpr int heads_per_block = kThreadsPerBlock / threads_per_head; \ + const dim3 grid((total_work + heads_per_block - 1) / heads_per_block); \ + const dim3 block(kThreadsPerBlock); \ + qknorm_rope_cache_kernel \ + <<>>( \ + qkv.data_ptr(), q_out.data_ptr(), k_out.data_ptr(), \ + v_out.data_ptr(), q_weight.data_ptr(), k_weight.data_ptr(), \ + cos.data_ptr(), sin.data_ptr(), \ + reinterpret_cast(positions.data_ptr()), \ + reinterpret_cast(slot_mapping.data_ptr()), \ + k_cache.data_ptr(), v_cache.data_ptr(), qkv.stride(0), \ + q_out.stride(0), k_out.stride(0), v_out.stride(0), cos.stride(0), \ + sin.stride(0), k_cache.stride(0), k_cache.stride(1), \ + k_cache.stride(2), static_cast(qkv.size(0)), \ + static_cast(slot_mapping.numel()), num_q_heads, num_kv_heads, \ + rotary_pairs, block_size, epsilon, block_size_shift, \ + norm_weight_bias, k_scale, v_scale); \ + } while (false) + + switch (head_dim) { + case 64: + LAUNCH_QKNORM_ROPE_CACHE(64, false); + break; + case 128: + LAUNCH_QKNORM_ROPE_CACHE(128, false); + break; + case 192: + if (rotary_pairs > 64) { + LAUNCH_QKNORM_ROPE_CACHE(192, true); + } else { + LAUNCH_QKNORM_ROPE_CACHE(192, false); + } + break; + case 256: + LAUNCH_QKNORM_ROPE_CACHE(256, false); + break; + default: + STD_TORCH_CHECK(false, "Unsupported head_dim: ", head_dim); + } +#undef LAUNCH_QKNORM_ROPE_CACHE +} + +void check_cuda_tensor(const torch::stable::Tensor& tensor, const char* name) { + STD_TORCH_CHECK(tensor.is_cuda(), name, " must be a CUDA tensor"); +} + +void check_same_device(const torch::stable::Tensor& reference, + const torch::stable::Tensor& tensor, const char* name) { + STD_TORCH_CHECK(tensor.get_device_index() == reference.get_device_index(), + name, " must be on the same CUDA device as qkv"); +} + +} // namespace + +void fused_qknorm_rope_cache( + torch::stable::Tensor& q_out, torch::stable::Tensor& k_out, + torch::stable::Tensor& v_out, torch::stable::Tensor& qkv, + torch::stable::Tensor& q_weight, torch::stable::Tensor& k_weight, + torch::stable::Tensor& cos, torch::stable::Tensor& sin, + torch::stable::Tensor& positions, torch::stable::Tensor& slot_mapping, + torch::stable::Tensor& k_cache, torch::stable::Tensor& v_cache, + int64_t head_dim, int64_t num_q_heads, int64_t num_kv_heads, + int64_t rotary_pairs, double epsilon, double norm_weight_bias, + const std::string& kv_cache_dtype, + std::optional k_scale, + std::optional v_scale) { + check_cuda_tensor(qkv, "qkv"); + check_cuda_tensor(q_out, "q_out"); + check_cuda_tensor(k_out, "k_out"); + check_cuda_tensor(v_out, "v_out"); + check_cuda_tensor(q_weight, "q_weight"); + check_cuda_tensor(k_weight, "k_weight"); + check_cuda_tensor(cos, "cos"); + check_cuda_tensor(sin, "sin"); + check_cuda_tensor(positions, "positions"); + check_cuda_tensor(slot_mapping, "slot_mapping"); + check_cuda_tensor(k_cache, "k_cache"); + check_cuda_tensor(v_cache, "v_cache"); + check_same_device(qkv, q_out, "q_out"); + check_same_device(qkv, k_out, "k_out"); + check_same_device(qkv, v_out, "v_out"); + check_same_device(qkv, q_weight, "q_weight"); + check_same_device(qkv, k_weight, "k_weight"); + check_same_device(qkv, cos, "cos"); + check_same_device(qkv, sin, "sin"); + check_same_device(qkv, positions, "positions"); + check_same_device(qkv, slot_mapping, "slot_mapping"); + check_same_device(qkv, k_cache, "k_cache"); + check_same_device(qkv, v_cache, "v_cache"); + + STD_TORCH_CHECK(qkv.dim() == 2 && qkv.stride(1) == 1, + "qkv must be a 2D tensor contiguous in its last dimension"); + STD_TORCH_CHECK(q_out.dim() == 2 && q_out.is_contiguous(), + "q_out must be contiguous and 2D"); + STD_TORCH_CHECK(k_out.dim() == 2 && k_out.is_contiguous(), + "k_out must be contiguous and 2D"); + STD_TORCH_CHECK(v_out.dim() == 2 && v_out.is_contiguous(), + "v_out must be contiguous and 2D"); + STD_TORCH_CHECK(q_weight.dim() == 1 && q_weight.is_contiguous(), + "q_weight must be contiguous and 1D"); + STD_TORCH_CHECK(k_weight.dim() == 1 && k_weight.is_contiguous(), + "k_weight must be contiguous and 1D"); + STD_TORCH_CHECK(cos.dim() == 2 && cos.stride(1) == 1, + "cos must be 2D and contiguous in its last dimension"); + STD_TORCH_CHECK(sin.dim() == 2 && sin.stride(1) == 1, + "sin must be 2D and contiguous in its last dimension"); + STD_TORCH_CHECK( + positions.scalar_type() == torch::headeronly::ScalarType::Long, + "positions must be an int64 tensor"); + STD_TORCH_CHECK(positions.is_contiguous(), "positions must be contiguous"); + STD_TORCH_CHECK( + slot_mapping.scalar_type() == torch::headeronly::ScalarType::Long, + "slot_mapping must be an int64 tensor"); + STD_TORCH_CHECK(slot_mapping.is_contiguous(), + "slot_mapping must be contiguous"); + STD_TORCH_CHECK(k_cache.dim() == 4 && v_cache.dim() == 4, + "K/V cache must have shape [blocks, block_size, heads, dim]"); + for (int64_t d = 0; d < 4; ++d) { + STD_TORCH_CHECK(k_cache.size(d) == v_cache.size(d) && + k_cache.stride(d) == v_cache.stride(d), + "K/V cache must have identical shapes and strides"); + } + STD_TORCH_CHECK(k_cache.stride(3) == 1, + "K/V cache head dimension must be contiguous"); + + STD_TORCH_CHECK(num_q_heads > 0 && num_kv_heads > 0, + "Q/KV head counts must be positive"); + STD_TORCH_CHECK( + head_dim == 64 || head_dim == 128 || head_dim == 192 || head_dim == 256, + "Unsupported head_dim"); + STD_TORCH_CHECK(cos.size(0) == sin.size(0), + "cos/sin cache must have the same number of positions"); + const int64_t num_tokens = qkv.size(0); + const int64_t packed_width = (num_q_heads + 2 * num_kv_heads) * head_dim; + STD_TORCH_CHECK(qkv.size(1) == packed_width, + "qkv width must equal (num_q_heads + 2 * num_kv_heads) * " + "head_dim"); + STD_TORCH_CHECK( + q_out.size(0) == num_tokens && q_out.size(1) == num_q_heads * head_dim, + "q_out has an invalid shape"); + STD_TORCH_CHECK( + k_out.size(0) == num_tokens && k_out.size(1) == num_kv_heads * head_dim, + "k_out has an invalid shape"); + STD_TORCH_CHECK( + v_out.size(0) == k_out.size(0) && v_out.size(1) == k_out.size(1), + "v_out has an invalid shape"); + STD_TORCH_CHECK(q_weight.numel() == head_dim && k_weight.numel() == head_dim, + "Q/K norm weights must have head_dim elements"); + STD_TORCH_CHECK(positions.numel() == num_tokens, + "positions must have one entry per qkv token"); + STD_TORCH_CHECK(rotary_pairs >= 0 && rotary_pairs * 2 <= head_dim && + rotary_pairs % 4 == 0, + "rotary_pairs must be a multiple of 4 in [0, head_dim / 2]"); + STD_TORCH_CHECK(cos.size(1) >= rotary_pairs && sin.size(1) >= rotary_pairs, + "cos/sin cache is smaller than rotary_pairs"); + STD_TORCH_CHECK( + k_cache.size(2) == num_kv_heads && k_cache.size(3) == head_dim, + "K/V cache head shape does not match the kernel arguments"); + STD_TORCH_CHECK(k_cache.size(1) > 0, "K/V cache block_size must be positive"); + STD_TORCH_CHECK(slot_mapping.numel() <= num_tokens, + "slot_mapping cannot contain more entries than qkv tokens"); + STD_TORCH_CHECK(qkv.scalar_type() == q_out.scalar_type() && + qkv.scalar_type() == k_out.scalar_type() && + qkv.scalar_type() == v_out.scalar_type(), + "qkv and Q/K/V outputs must have the same dtype"); + const bool fp8_cache = + kv_cache_dtype == "fp8" || kv_cache_dtype == "fp8_e4m3"; + STD_TORCH_CHECK(fp8_cache || kv_cache_dtype == "auto" || + kv_cache_dtype == "bfloat16" || + kv_cache_dtype == "float16", + "Unsupported KV cache dtype: ", kv_cache_dtype); + STD_TORCH_CHECK(k_cache.scalar_type() == v_cache.scalar_type(), + "K/V caches must have the same dtype"); + if (fp8_cache) { + STD_TORCH_CHECK( + k_cache.scalar_type() == torch::headeronly::ScalarType::Byte || + k_cache.scalar_type() == + torch::headeronly::ScalarType::Float8_e4m3fn, + "FP8 cache requires uint8 or float8_e4m3fn storage"); + STD_TORCH_CHECK(k_scale.has_value() && v_scale.has_value(), + "FP8 cache requires K/V scales"); + for (const auto* scale : {&*k_scale, &*v_scale}) { + check_cuda_tensor(*scale, "KV scale"); + check_same_device(qkv, *scale, "KV scale"); + STD_TORCH_CHECK( + scale->numel() == 1 && + scale->scalar_type() == torch::headeronly::ScalarType::Float, + "KV scales must be single FP32 values"); + } + } else { + STD_TORCH_CHECK(qkv.scalar_type() == k_cache.scalar_type(), + "Unquantized cache must match QKV dtype"); + } + STD_TORCH_CHECK(q_weight.scalar_type() == k_weight.scalar_type(), + "Q/K norm weights must have the same dtype"); + STD_TORCH_CHECK(cos.scalar_type() == sin.scalar_type(), + "cos and sin must have the same dtype"); + + if (num_tokens == 0) return; + + const torch::stable::accelerator::DeviceGuard device_guard( + qkv.get_device_index()); + const cudaStream_t stream = get_current_cuda_stream(qkv.get_device_index()); + VLLM_STABLE_DISPATCH_HALF_TYPES( + qkv.scalar_type(), "qknorm_rope_cache_input", [&] { + using input_t = scalar_t; + VLLM_STABLE_DISPATCH_FLOATING_TYPES( + q_weight.scalar_type(), "qknorm_rope_cache_weight", [&] { + using norm_weight_t = scalar_t; + VLLM_STABLE_DISPATCH_FLOATING_TYPES( + cos.scalar_type(), "qknorm_rope_cache_rope", [&] { + using rope_cache_t = scalar_t; + auto launch = [&](auto fp8_tag) { + launch_qknorm_rope_cache( + qkv, q_out, k_out, v_out, q_weight, k_weight, cos, + sin, positions, slot_mapping, k_cache, v_cache, + static_cast(head_dim), + static_cast(num_q_heads), + static_cast(num_kv_heads), + static_cast(rotary_pairs), + static_cast(epsilon), + static_cast(norm_weight_bias), stream, + fp8_cache + ? static_cast(k_scale->data_ptr()) + : nullptr, + fp8_cache + ? static_cast(v_scale->data_ptr()) + : nullptr); + }; + if (fp8_cache) { + launch(std::true_type{}); + } else { + launch(std::false_type{}); + } + }); + }); + }); + STD_CUDA_KERNEL_LAUNCH_CHECK(); +} + +} // namespace vllm diff --git a/csrc/muse_spark_mul_tanh_soft_cap_kernel.cu b/csrc/muse_spark_mul_tanh_soft_cap_kernel.cu new file mode 100644 index 0000000000..e3eefae163 --- /dev/null +++ b/csrc/muse_spark_mul_tanh_soft_cap_kernel.cu @@ -0,0 +1,100 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +#include +#include +#include +#include + +#include + +#include "dispatch_utils.h" +#include "quantization/vectorization_utils.cuh" +#include "torch_utils.h" + +namespace vllm { + +// 16 B of output per thread; low-precision inputs are widened in registers. +static constexpr int kMuseSoftCapVecSize = 4; + +// blockIdx.x walks rows, blockIdx.y splits a row. Rows are the sampled tokens +// (often one, at decode) while a row is the whole vocab, so splitting the row +// is what keeps the GPU busy; a row-per-block grid would leave a single-token +// decode on one SM. +template +__global__ void muse_spark_mul_tanh_soft_cap_kernel( + float* __restrict__ out, const scalar_t* __restrict__ input, int d, + float multiplier, float soft_cap_temp) { + const int64_t row_offset = static_cast(blockIdx.x) * d; + const scalar_t* __restrict__ in_row = input + row_offset; + float* __restrict__ out_row = out + row_offset; + + const int tid = blockIdx.y * blockDim.x + threadIdx.x; + const int stride = gridDim.y * blockDim.x; + + vectorize_with_alignment( + in_row, out_row, d, tid, stride, + [=] __device__(float& dst, const scalar_t& src) { + // Multiply then divide, as MulTanhSoftCap.forward does, so fp32 + // rounding matches the eager path. + dst = soft_cap_temp * + tanhf(static_cast(src) * multiplier / soft_cap_temp); + }); +} + +} // namespace vllm + +// out = soft_cap_temp * tanh(input * multiplier / soft_cap_temp), fp32. +void muse_spark_mul_tanh_soft_cap(torch::stable::Tensor& out, + torch::stable::Tensor& input, + double multiplier, double soft_cap_temp) { + STD_TORCH_CHECK( + input.scalar_type() == torch::headeronly::ScalarType::Float || + input.scalar_type() == torch::headeronly::ScalarType::Half || + input.scalar_type() == torch::headeronly::ScalarType::BFloat16, + "muse_spark_mul_tanh_soft_cap: input must be float32, float16 or " + "bfloat16"); + STD_TORCH_CHECK(out.scalar_type() == torch::headeronly::ScalarType::Float, + "muse_spark_mul_tanh_soft_cap: out must be float32"); + STD_TORCH_CHECK(input.is_contiguous(), + "muse_spark_mul_tanh_soft_cap: input must be contiguous"); + STD_TORCH_CHECK(out.is_contiguous(), + "muse_spark_mul_tanh_soft_cap: out must be contiguous"); + STD_TORCH_CHECK(out.numel() == input.numel(), + "muse_spark_mul_tanh_soft_cap: out and input must have the " + "same number of elements"); + STD_TORCH_CHECK( + soft_cap_temp != 0.0, + "muse_spark_mul_tanh_soft_cap: soft_cap_temp must be nonzero"); + + if (input.numel() == 0) { + return; + } + + const int d = static_cast(input.size(input.dim() - 1)); + const int64_t num_tokens = input.numel() / d; + + constexpr int kThreads = 128; + const int64_t vecs_per_row = + (d + vllm::kMuseSoftCapVecSize - 1) / vllm::kMuseSoftCapVecSize; + const int64_t blocks_to_cover = (vecs_per_row + kThreads - 1) / kThreads; + constexpr int64_t kTargetBlocks = 32768; + const int64_t blocks_per_row = std::clamp( + kTargetBlocks / num_tokens, 1, std::max(blocks_to_cover, 1)); + + dim3 grid(static_cast(num_tokens), + static_cast(blocks_per_row)); + dim3 block(kThreads); + + const torch::stable::accelerator::DeviceGuard device_guard( + input.get_device_index()); + const cudaStream_t stream = get_current_cuda_stream(); + VLLM_STABLE_DISPATCH_FLOATING_TYPES( + input.scalar_type(), "muse_spark_mul_tanh_soft_cap", [&] { + vllm::muse_spark_mul_tanh_soft_cap_kernel + <<>>(out.mutable_data_ptr(), + input.const_data_ptr(), d, + static_cast(multiplier), + static_cast(soft_cap_temp)); + }); +} diff --git a/csrc/ops.h b/csrc/ops.h index 6e40dc7e54..8304b1a312 100644 --- a/csrc/ops.h +++ b/csrc/ops.h @@ -363,6 +363,35 @@ minimax_allreduce_rms_qk(torch::stable::Tensor qkv, int64_t const nranks, double const eps); #endif +#ifndef USE_ROCM +namespace vllm { +void fused_qknorm_rope_cache( + torch::stable::Tensor& q_out, torch::stable::Tensor& k_out, + torch::stable::Tensor& v_out, torch::stable::Tensor& qkv, + torch::stable::Tensor& q_weight, torch::stable::Tensor& k_weight, + torch::stable::Tensor& cos, torch::stable::Tensor& sin, + torch::stable::Tensor& positions, torch::stable::Tensor& slot_mapping, + torch::stable::Tensor& k_cache, torch::stable::Tensor& v_cache, + int64_t head_dim, int64_t num_q_heads, int64_t num_kv_heads, + int64_t rotary_pairs, double epsilon, double norm_weight_bias, + const std::string& kv_cache_dtype, + std::optional k_scale, + std::optional v_scale); +} // namespace vllm +#endif + +void fused_muse_spark_qknorm_rope_kv_insert( + torch::stable::Tensor& qkv, torch::stable::Tensor const& positions, + int64_t num_heads, int64_t num_kv_heads, int64_t head_dim, double eps, + std::optional cos_sin_cache, int64_t rotary_dim, + std::optional slot_mapping, + std::optional kv_cache, int64_t block_size, + const std::string& kv_cache_dtype, + std::optional k_scale, + std::optional v_scale, + std::optional q_fp8_out, + std::optional q_scale); + // Horizontally-fused MiniMax-M3 QK-norm + partial NeoX RoPE (+ optional KV / // index-cache insert). Dense layer: norm+RoPE only; sparse layer: also packs // the index branch and scatters k/v/index_k into their paged caches. @@ -546,6 +575,10 @@ std::tuple allocate_shared_buffer_and_handle( int64_t open_mem_handle(torch::stable::Tensor& mem_handle); void free_shared_buffer(int64_t buffer); +void muse_spark_mul_tanh_soft_cap(torch::stable::Tensor& out, + torch::stable::Tensor& input, + double multiplier, double soft_cap_temp); + // Activation kernels (shared CUDA/ROCm) void silu_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input); void silu_and_mul_clamp(torch::stable::Tensor& out, diff --git a/csrc/torch_bindings.cpp b/csrc/torch_bindings.cpp index 223cf5e101..f5eba99b86 100644 --- a/csrc/torch_bindings.cpp +++ b/csrc/torch_bindings.cpp @@ -501,6 +501,27 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { "float eps) -> (Tensor, Tensor)"); #endif +#ifndef USE_ROCM + ops.def( + "fused_qknorm_rope_cache(" + "Tensor! q_out, Tensor! k_out, Tensor! v_out, Tensor qkv, " + "Tensor q_weight, Tensor k_weight, Tensor cos, Tensor sin, " + "Tensor positions, Tensor slot_mapping, Tensor! k_cache, " + "Tensor! v_cache, int head_dim, int num_q_heads, int num_kv_heads, " + "int rotary_pairs, float epsilon, float norm_weight_bias, " + "str kv_cache_dtype=\"auto\", Tensor? k_scale=None, Tensor? " + "v_scale=None) -> ()"); +#endif + + ops.def( + "fused_muse_spark_qknorm_rope_kv_insert(" + "Tensor! qkv, Tensor positions, int num_heads, int num_kv_heads, " + "int head_dim, float eps, Tensor? cos_sin_cache, int rotary_dim, " + "Tensor? slot_mapping, Tensor!? kv_cache, int block_size, " + "str kv_cache_dtype, Tensor? k_scale=None, Tensor? v_scale=None, " + "Tensor!? q_fp8_out=None, Tensor? q_scale=None) " + "-> ()"); + // Horizontally-fused MiniMax-M3 QK-norm + partial NeoX RoPE + KV-insert. ops.def( "fused_minimax_m3_qknorm_rope_kv_insert(" @@ -598,6 +619,10 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { "y_q, Tensor! y_s, bool use_ue8m0) -> ()"); ops.def("weak_ref_tensor(Tensor input) -> Tensor"); + ops.def( + "muse_spark_mul_tanh_soft_cap(Tensor! out, Tensor input, " + "float multiplier, float soft_cap_temp) -> ()"); + // Activation function used in SwiGLU. ops.def("silu_and_mul(Tensor! result, Tensor input) -> ()"); @@ -828,7 +853,11 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) { TORCH_BOX(&fused_kimi_k3_mla_decode_q_concat_ds_mla_insert)); #ifndef USE_ROCM ops.impl("minimax_allreduce_rms_qk", TORCH_BOX(&minimax_allreduce_rms_qk)); + ops.impl("fused_qknorm_rope_cache", + TORCH_BOX(&vllm::fused_qknorm_rope_cache)); #endif + ops.impl("fused_muse_spark_qknorm_rope_kv_insert", + TORCH_BOX(&fused_muse_spark_qknorm_rope_kv_insert)); ops.impl("fused_minimax_m3_qknorm_rope_kv_insert", TORCH_BOX(&fused_minimax_m3_qknorm_rope_kv_insert)); #ifdef VLLM_ENABLE_FUSED_KDA_DECODE @@ -864,6 +893,8 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) { TORCH_BOX(&persistent_masked_m_silu_mul_quant)); ops.impl("weak_ref_tensor", TORCH_BOX(&weak_ref_tensor)); ops.impl("silu_and_mul_quant", TORCH_BOX(&silu_and_mul_quant)); + ops.impl("muse_spark_mul_tanh_soft_cap", + TORCH_BOX(&muse_spark_mul_tanh_soft_cap)); ops.impl("silu_and_mul", TORCH_BOX(&silu_and_mul)); ops.impl("mul_and_silu", TORCH_BOX(&mul_and_silu)); ops.impl("gelu_and_mul", TORCH_BOX(&gelu_and_mul)); diff --git a/tests/kernels/test_fused_muse_spark_qknorm_rope_kv_insert.py b/tests/kernels/test_fused_muse_spark_qknorm_rope_kv_insert.py new file mode 100644 index 0000000000..99012d4943 --- /dev/null +++ b/tests/kernels/test_fused_muse_spark_qknorm_rope_kv_insert.py @@ -0,0 +1,340 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch + +import vllm._custom_ops as ops +from vllm.platforms import current_platform + +HEAD_DIM = 64 + + +pytestmark = pytest.mark.skipif( + not current_platform.is_cuda_alike() or not torch.cuda.is_available(), + reason="Muse Spark kernel requires a CUDA or ROCm GPU", +) + + +def make_cos_sin_cache(max_pos, rotary_dim, base, device, dtype=torch.float32): + """The table ``RotaryEmbedding`` builds: cos||sin, fp32 for Muse-Spark.""" + inv_freq = 1.0 / ( + base + ** ( + torch.arange(0, rotary_dim, 2, dtype=torch.float32, device=device) + / rotary_dim + ) + ) + t = torch.arange(max_pos, dtype=torch.float32, device=device) + freqs = torch.einsum("i,j->ij", t, inv_freq) + return torch.cat((freqs.cos(), freqs.sin()), dim=-1).to(dtype) + + +def qk_norm_rope_ref(x, positions, cos_sin_cache, rotary_dim, eps): + """fp32 reference for one q or k branch, rounded once; no table means NoPE.""" + out = x.float() + out = out * torch.rsqrt(out.pow(2).mean(dim=-1, keepdim=True) + eps) + if cos_sin_cache is not None: + half = rotary_dim // 2 + cs = cos_sin_cache[positions].float() + cos = cs[..., :half].unsqueeze(1) + sin = cs[..., half:].unsqueeze(1) + # GPT-J interleaved: the rotated pair is (2i, 2i+1). + x0 = out[..., :rotary_dim:2] + x1 = out[..., 1:rotary_dim:2] + out = out.clone() + out[..., :rotary_dim:2] = x0 * cos - x1 * sin + out[..., 1:rotary_dim:2] = x0 * sin + x1 * cos + return out.to(x.dtype) + + +def assert_fp8_cache_close(got, expected, exact_fraction=0.999): + """Compare raw cache codes in the platform's FP8 encoding.""" + diff = (got.int() - expected.int()).abs() + within = diff <= 1 + assert bool(within.all()), ( + f"fp8 cache differs by more than one code in {int((~within).sum())} " + f"of {diff.numel()} elements" + ) + exact = float((diff == 0).float().mean()) + assert exact >= exact_fraction, ( + f"only {exact:.5f} of the fp8 cache is bit-identical to " + f"reshape_and_cache_flash, expected at least {exact_fraction}" + ) + + +def make_case(num_tokens, nq, nkv, dtype, device, with_gate=True, seed=0): + torch.manual_seed(seed) + qsz, kvsz = nq * HEAD_DIM, nkv * HEAD_DIM + gate = qsz if with_gate else 0 + qkv = torch.randn(num_tokens, qsz + 2 * kvsz + gate, dtype=dtype, device=device) + return qkv, qsz, kvsz, gate + + +def split_qkv(qkv, qsz, kvsz, gate): + return qkv.split([qsz, kvsz, kvsz, gate], dim=-1) + + +def check_qk_norm_rope( + num_tokens, + nq, + nkv, + dtype, + use_rope, + rotary_dim=HEAD_DIM, + with_gate=True, + rope_dtype=torch.float32, +): + device, eps, max_pos = "cuda", 1e-5, 4096 + + qkv, qsz, kvsz, gate = make_case(num_tokens, nq, nkv, dtype, device, with_gate) + qkv_orig = qkv.clone() + cos_sin = ( + make_cos_sin_cache(max_pos, rotary_dim, 500_000.0, device, rope_dtype) + if use_rope + else None + ) + positions = torch.randint( + 0, max_pos, (num_tokens,), dtype=torch.int64, device=device + ) + + ops.fused_muse_spark_qknorm_rope_kv_insert( + qkv, + positions, + nq, + nkv, + HEAD_DIM, + eps, + cos_sin_cache=cos_sin, + rotary_dim=rotary_dim if use_rope else 0, + ) + + q_out, k_out, v_out, og_out = split_qkv(qkv, qsz, kvsz, gate) + q_in, k_in, v_in, og_in = split_qkv(qkv_orig, qsz, kvsz, gate) + + q_ref = qk_norm_rope_ref( + q_in.view(num_tokens, nq, HEAD_DIM), positions, cos_sin, rotary_dim, eps + ).view(num_tokens, qsz) + k_ref = qk_norm_rope_ref( + k_in.view(num_tokens, nkv, HEAD_DIM), positions, cos_sin, rotary_dim, eps + ).view(num_tokens, kvsz) + + # rsqrtf and lane summation order cost a couple of fp32 ulp, well inside + # one ulp of the stored dtype. + torch.testing.assert_close(q_out, q_ref, rtol=1e-2, atol=1e-2) + torch.testing.assert_close(k_out, k_ref, rtol=1e-2, atol=1e-2) + torch.testing.assert_close(v_out, v_in, rtol=0, atol=0) + torch.testing.assert_close(og_out, og_in, rtol=0, atol=0) + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +@pytest.mark.parametrize("use_rope", [True, False]) +def test_qk_norm_rope(dtype, use_rope): + check_qk_norm_rope(17, 32, 4, dtype, use_rope) + + +@pytest.mark.parametrize( + "num_tokens,variant", + [ + pytest.param(64, {"rotary_dim": 32}, id="partial_rotary_dim"), + # A [q | k | v] row (no o-gate) is addressed by its own stride. + pytest.param(33, {"with_gate": False}, id="no_gate_row"), + pytest.param(128, {"rope_dtype": torch.bfloat16}, id="bf16_cos_sin_cache"), + ], +) +def test_qk_norm_rope_layout_variants(num_tokens, variant): + check_qk_norm_rope(num_tokens, 32, 4, torch.bfloat16, True, **variant) + + +# Cover narrow/wide head groups for TP4/TP8 and the one-chunk-per-token case. +@pytest.mark.parametrize( + "num_tokens,nq,nkv,kv_cache_dtype,use_rope", + [ + pytest.param(1, 32, 4, "auto", True, id="tp4-decode"), + pytest.param(1, 16, 2, "fp8", False, id="tp8-decode-fp8-nope"), + pytest.param(1025, 32, 4, "auto", True, id="tp4-group2"), + pytest.param(1025, 32, 4, "fp8", True, id="tp4-group2-fp8"), + pytest.param(1537, 32, 4, "auto", True, id="tp4-group4"), + pytest.param(1537, 32, 4, "auto", False, id="tp4-nope"), + pytest.param(3073, 16, 2, "auto", True, id="tp8-group4"), + pytest.param(3073, 16, 2, "fp8", False, id="tp8-fp8-nope"), + pytest.param(40000, 2, 1, "auto", True, id="single-chunk-per-token"), + ], +) +def test_kv_insert(num_tokens, nq, nkv, kv_cache_dtype, use_rope): + device, dtype, eps = "cuda", torch.bfloat16, 1e-5 + rotary_dim, max_pos, block_size = HEAD_DIM, 4096, 16 + + qkv, qsz, kvsz, gate = make_case(num_tokens, nq, nkv, dtype, device, seed=1) + qkv_orig = qkv.clone() + cos_sin = ( + make_cos_sin_cache(max_pos, rotary_dim, 500_000.0, device) if use_rope else None + ) + positions = torch.randint( + 0, max_pos, (num_tokens,), dtype=torch.int64, device=device + ) + + num_blocks = (num_tokens + block_size - 1) // block_size + 1 + storage = torch.uint8 if kv_cache_dtype == "fp8" else dtype + kv_cache = torch.zeros( + num_blocks, nkv, block_size, 2 * HEAD_DIM, dtype=storage, device=device + ) + slot_mapping = torch.randperm( + num_blocks * block_size, dtype=torch.int64, device=device + )[:num_tokens] + if num_tokens > 1: + # -1 marks a token the scheduler did not place; the kernel must skip it + # and leave the cache alone. + slot_mapping[::3] = -1 + + # Non-unit scales, so a kernel that ignored them would fail loudly. + k_scale = torch.tensor(0.37, dtype=torch.float32, device=device) + v_scale = torch.tensor(1.83, dtype=torch.float32, device=device) + + ops.fused_muse_spark_qknorm_rope_kv_insert( + qkv, + positions, + nq, + nkv, + HEAD_DIM, + eps, + cos_sin_cache=cos_sin, + rotary_dim=rotary_dim if use_rope else 0, + slot_mapping=slot_mapping, + kv_cache=kv_cache, + block_size=block_size, + kv_cache_dtype=kv_cache_dtype, + k_scale=k_scale if kv_cache_dtype != "auto" else None, + v_scale=v_scale if kv_cache_dtype != "auto" else None, + ) + + q_out, k_out, v_out, og_out = split_qkv(qkv, qsz, kvsz, gate) + q_in, k_in, v_in, og_in = split_qkv(qkv_orig, qsz, kvsz, gate) + + q_ref = qk_norm_rope_ref( + q_in.view(num_tokens, nq, HEAD_DIM), positions, cos_sin, rotary_dim, eps + ).view(num_tokens, qsz) + k_ref = qk_norm_rope_ref( + k_in.view(num_tokens, nkv, HEAD_DIM), positions, cos_sin, rotary_dim, eps + ).view(num_tokens, kvsz) + torch.testing.assert_close(q_out, q_ref, rtol=1e-2, atol=1e-2) + torch.testing.assert_close(k_out, k_ref, rtol=1e-2, atol=1e-2) + torch.testing.assert_close(v_out, v_in, rtol=0, atol=0) + torch.testing.assert_close(og_out, og_in, rtol=0, atol=0) + + expected = torch.zeros_like(kv_cache) + exp_k, exp_v = expected.transpose(1, 2).split(HEAD_DIM, dim=-1) + ops.reshape_and_cache_flash( + k_out.view(num_tokens, nkv, HEAD_DIM), + v_out.view(num_tokens, nkv, HEAD_DIM), + exp_k, + exp_v, + slot_mapping, + kv_cache_dtype, + k_scale, + v_scale, + ) + if kv_cache_dtype == "auto": + torch.testing.assert_close(kv_cache, expected, rtol=0, atol=0) + else: + assert_fp8_cache_close(kv_cache, expected) + + +# The fused Q output requires E4M3FN; scaled_fp8_quant uses the platform dtype. +@pytest.mark.skipif( + current_platform.is_rocm() and current_platform.fp8_dtype() != torch.float8_e4m3fn, + reason="Fused E4M3FN Q output cannot be compared with ROCm FNUZ quantization", +) +@pytest.mark.parametrize( + "head_dim,num_tokens,rotary_fraction,use_cuda_graph", + [ + pytest.param(64, 17, 0, False, id="head64-nope"), + pytest.param(128, 17, 0.5, False, id="head128-partial-rope"), + pytest.param(256, 17, 0, False, id="head256-nope"), + pytest.param(256, 3073, 0.5, False, id="head256-partial-rope-large"), + pytest.param(64, 3073, 0.5, True, id="graph-rope"), + pytest.param(64, 3073, 0, True, id="graph-nope"), + ], +) +def test_fused_query_quant_matches_separate_quant( + head_dim, num_tokens, rotary_fraction, use_cuda_graph +): + """Preserve the rounding boundary, scales, padded Q, and KV/gate writes.""" + torch.manual_seed(42) + nq, nkv, block_size = 16, 2, 16 + qsz, kvsz = nq * head_dim, nkv * head_dim + qkv = torch.randn(num_tokens, 2 * (qsz + kvsz), device="cuda", dtype=torch.bfloat16) + original = qkv.clone() + reference = qkv.clone() + q_fp8 = torch.empty(num_tokens, qsz, device="cuda", dtype=torch.float8_e4m3fn) + positions = torch.arange(num_tokens, device="cuda", dtype=torch.int64) + rotary_dim = int(head_dim * rotary_fraction) + cos_sin = ( + make_cos_sin_cache(num_tokens, rotary_dim, 500_000.0, "cuda") + if rotary_dim + else None + ) + cache = torch.zeros( + (num_tokens + block_size - 1) // block_size, + nkv, + block_size, + 2 * head_dim, + device="cuda", + dtype=torch.uint8, + ) + expected_cache = torch.zeros_like(cache) + slots = positions.clone() + slots[::3] = -1 + # Small enough that part of Q saturates the e4m3 range. + scale = torch.tensor(0.001, device="cuda", dtype=torch.float32) + kwargs = dict( + cos_sin_cache=cos_sin, + rotary_dim=rotary_dim, + slot_mapping=slots, + block_size=block_size, + kv_cache_dtype="fp8", + k_scale=torch.tensor(0.37, device="cuda"), + v_scale=torch.tensor(1.83, device="cuda"), + ) + ops.fused_muse_spark_qknorm_rope_kv_insert( + reference, + positions, + nq, + nkv, + head_dim, + 1e-5, + kv_cache=expected_cache, + **kwargs, + ) + expected_q, _ = ops.scaled_fp8_quant(reference[:, :qsz], scale=scale) + + def run(): + qkv.copy_(original) + ops.fused_muse_spark_qknorm_rope_kv_insert( + qkv, + positions, + nq, + nkv, + head_dim, + 1e-5, + kv_cache=cache, + q_fp8_out=q_fp8, + q_scale=scale, + **kwargs, + ) + + run() + if use_cuda_graph: + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run() + qkv.zero_() + cache.zero_() + q_fp8.zero_() + graph.replay() + torch.testing.assert_close( + q_fp8.view(torch.uint8), expected_q.view(torch.uint8), rtol=0, atol=0 + ) + torch.testing.assert_close(cache, expected_cache, rtol=0, atol=0) + torch.testing.assert_close(qkv[:, :qsz], original[:, :qsz], rtol=0, atol=0) + torch.testing.assert_close(qkv[:, qsz:], reference[:, qsz:], rtol=0, atol=0) diff --git a/tests/kernels/test_fused_qknorm_rope_cache.py b/tests/kernels/test_fused_qknorm_rope_cache.py new file mode 100644 index 0000000000..425a623ab5 --- /dev/null +++ b/tests/kernels/test_fused_qknorm_rope_cache.py @@ -0,0 +1,195 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Weighted RMSNorm + partial NeoX RoPE with separate paged K/V caches.""" + +import pytest +import torch + +import vllm._custom_ops as ops +from vllm.platforms import current_platform + +pytestmark = pytest.mark.skipif( + not current_platform.is_cuda() or not torch.cuda.is_available(), + reason="Step-5 kernel is only built for CUDA", +) + + +def norm_rope_reference(x, weight, cos_sin, positions, rotary_dim, eps, bias): + xf = x.float() + normed = xf * torch.rsqrt(xf.square().mean(-1, keepdim=True) + eps) + out = (normed * (weight.float() + bias)).to(x.dtype) + half = rotary_dim // 2 + cos, sin = cos_sin[positions].float().unsqueeze(1).chunk(2, dim=-1) + x0, x1 = out[..., :half].float(), out[..., half:rotary_dim].float() + out[..., :half] = x0 * cos - x1 * sin + out[..., half:rotary_dim] = x0 * sin + x1 * cos + return out + + +def check_norm_rope_and_cache( + dtype=torch.bfloat16, + kv_cache_dtype="auto", + head_dim=128, + rotary_dim=64, + num_tokens=19, + block_size=16, + fp32_weights=False, + bias=0.0, + strided_cache=False, + use_cuda_graph=False, +): + """Cover rounding, 192-wide RoPE, padded rows and short/negative slots.""" + torch.manual_seed(0) + nq, nkv, max_pos, eps = 5, 2, 64, 1e-5 + width = (nq + 2 * nkv) * head_dim + storage = torch.randn(num_tokens, width + 8, device="cuda", dtype=dtype) + qkv = storage[:, :width] + original = storage.clone() + weight_dtype = torch.float32 if fp32_weights else dtype + qw = torch.randn(head_dim, device="cuda", dtype=weight_dtype) + kw = torch.randn_like(qw) + angles = torch.randn(max_pos, rotary_dim // 2, device="cuda") + rope_dtype = torch.float32 if fp32_weights else dtype + cos_sin = torch.cat((angles.cos(), angles.sin()), dim=-1).to(rope_dtype) + positions = torch.randint(max_pos, (num_tokens,), device="cuda") + num_blocks = (num_tokens + block_size - 1) // block_size + 1 + cache_dtype = torch.uint8 if kv_cache_dtype == "fp8" else dtype + shape = (num_blocks, block_size, nkv, head_dim) + if strided_cache: + shape = (num_blocks, nkv, block_size, head_dim) + key_cache = torch.zeros(shape, device="cuda", dtype=cache_dtype) + value_cache = torch.zeros_like(key_cache) + if strided_cache: + key_cache = key_cache.transpose(1, 2) + value_cache = value_cache.transpose(1, 2) + # A short mapping models graph padding; negative slots must stay unwritten. + num_slots = max(0, num_tokens - 3) if num_tokens > 1 else num_tokens + slots = torch.randperm(num_blocks * block_size, device="cuda")[:num_slots] + slots[1::3] = -1 + k_scale = torch.tensor(0.37, device="cuda") + v_scale = torch.tensor(1.83, device="cuda") + + def run(): + return ops.fused_qknorm_rope_cache( + qkv, + qw, + kw, + cos_sin, + positions, + slots, + key_cache, + value_cache, + nq, + nkv, + head_dim, + rotary_dim, + eps, + norm_weight_bias=bias, + kv_cache_dtype=kv_cache_dtype, + k_scale=k_scale, + v_scale=v_scale, + ) + + q, k, v = run() + if use_cuda_graph: + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + q, k, v = run() + key_cache.zero_() + value_cache.zero_() + graph.replay() + + q_in, k_in, v_in = qkv.split( + (nq * head_dim, nkv * head_dim, nkv * head_dim), dim=-1 + ) + tol = 1e-2 if dtype == torch.bfloat16 else 2e-3 + for actual, src, weight, heads in ((q, q_in, qw, nq), (k, k_in, kw, nkv)): + expected = norm_rope_reference( + src.reshape(num_tokens, heads, head_dim), + weight, + cos_sin, + positions, + rotary_dim, + eps, + bias, + ).reshape_as(actual) + torch.testing.assert_close(actual, expected, rtol=tol, atol=tol) + torch.testing.assert_close(v, v_in, rtol=0, atol=0) + torch.testing.assert_close(storage, original, rtol=0, atol=0) + + expected_k, expected_v = torch.zeros_like(key_cache), torch.zeros_like(value_cache) + if num_slots: + ops.reshape_and_cache_flash( + k.reshape(num_tokens, nkv, head_dim), + v.reshape(num_tokens, nkv, head_dim), + expected_k, + expected_v, + slots, + kv_cache_dtype, + k_scale, + v_scale, + ) + for actual, expected in ((key_cache, expected_k), (value_cache, expected_v)): + if kv_cache_dtype == "fp8": + # Reciprocal multiplication versus division may differ by one code. + assert (actual.int() - expected.int()).abs().max() <= 1 + else: + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +@pytest.mark.parametrize( + "case", + [ + pytest.param(dict(head_dim=64, rotary_dim=32, num_tokens=1), id="decode"), + pytest.param( + dict(block_size=7, fp32_weights=True, bias=1.0, strided_cache=True), + id="strided-cache-fp32-weights", + ), + pytest.param( + dict(head_dim=192, rotary_dim=128, fp32_weights=True), + id="head192-partial-rope", + ), + pytest.param( + dict( + head_dim=192, + rotary_dim=192, + block_size=7, + bias=1.0, + strided_cache=True, + kv_cache_dtype="fp8", + ), + id="head192-full-rope-fp8", + ), + pytest.param( + dict( + head_dim=256, rotary_dim=256, strided_cache=True, kv_cache_dtype="fp8" + ), + id="head256-full-rope-fp8", + ), + pytest.param( + dict(head_dim=192, rotary_dim=0, fp32_weights=True, bias=1.0), + id="no-rope", + ), + pytest.param( + dict(head_dim=64, rotary_dim=8, fp32_weights=True), + id="minimum-rotary-dim", + ), + ], +) +def test_norm_rope_and_cache(dtype, case): + check_norm_rope_and_cache(dtype=dtype, **case) + + +@pytest.mark.parametrize( + "case", + [ + pytest.param(dict(num_tokens=0), id="empty-input"), + pytest.param(dict(use_cuda_graph=True), id="graph-auto-cache"), + pytest.param( + dict(use_cuda_graph=True, kv_cache_dtype="fp8"), id="graph-fp8-cache" + ), + ], +) +def test_empty_input_and_graph_replay(case): + check_norm_rope_and_cache(**case) diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index d74bb67816..586e4d19aa 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -339,6 +339,17 @@ def fused_add_rms_norm( torch.ops._C.fused_add_rms_norm(input, residual, weight, epsilon) +def muse_spark_mul_tanh_soft_cap( + x: torch.Tensor, + multiplier: float, + soft_cap_temp: float, +) -> torch.Tensor: + """Cast to FP32 and apply ``temp * tanh(x * multiplier / temp)`` in one pass.""" + out = torch.empty_like(x, dtype=torch.float32) + torch.ops._C.muse_spark_mul_tanh_soft_cap(out, x, multiplier, soft_cap_temp) + return out + + def fused_qk_norm_rope( qkv: torch.Tensor, num_heads_q: int, @@ -2598,6 +2609,105 @@ def reshape_and_cache_flash( ) +def fused_qknorm_rope_cache( + qkv: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + cos_sin_cache: torch.Tensor, + positions: torch.Tensor, + slot_mapping: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + num_q_heads: int, + num_kv_heads: int, + head_dim: int, + rotary_dim: int, + eps: float, + norm_weight_bias: float = 0.0, + kv_cache_dtype: str = "auto", + k_scale: torch.Tensor | None = None, + v_scale: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Fuse per-head RMSNorm, partial NeoX RoPE and KV insertion. + + Norm weights may be FP32 independently of QKV. A zero rotary dimension + skips RoPE; norm outputs are rounded to QKV dtype before rotation. + """ + if rotary_dim < 0 or rotary_dim > head_dim or rotary_dim % 8: + raise ValueError("rotary_dim must be a multiple of 8 in [0, head_dim]") + q = qkv.new_empty((qkv.shape[0], num_q_heads * head_dim)) + k = qkv.new_empty((qkv.shape[0], num_kv_heads * head_dim)) + v = torch.empty_like(k) + cos, sin = cos_sin_cache.chunk(2, dim=-1) + torch.ops._C.fused_qknorm_rope_cache( + q, + k, + v, + qkv, + q_weight, + k_weight, + cos, + sin, + positions, + slot_mapping, + key_cache, + value_cache, + head_dim, + num_q_heads, + num_kv_heads, + rotary_dim // 2, + eps, + norm_weight_bias, + kv_cache_dtype, + k_scale, + v_scale, + ) + return q, k, v + + +def fused_muse_spark_qknorm_rope_kv_insert( + qkv: torch.Tensor, + positions: torch.Tensor, + num_heads: int, + num_kv_heads: int, + head_dim: int, + eps: float, + cos_sin_cache: torch.Tensor | None = None, + rotary_dim: int = 0, + slot_mapping: torch.Tensor | None = None, + kv_cache: torch.Tensor | None = None, + block_size: int = 0, + kv_cache_dtype: str = "auto", + k_scale: torch.Tensor | None = None, + v_scale: torch.Tensor | None = None, + q_fp8_out: torch.Tensor | None = None, + q_scale: torch.Tensor | None = None, +) -> None: + """In-place weightless QK RMSNorm + GPT-J RoPE, optionally writing k/v to cache. + + When ``kv_cache`` is given, the caller must skip the attention layer's own KV write. + With ``q_fp8_out``, write scaled E4M3 Q there instead of rewriting Q in ``qkv``. + """ + torch.ops._C.fused_muse_spark_qknorm_rope_kv_insert( + qkv, + positions, + num_heads, + num_kv_heads, + head_dim, + eps, + cos_sin_cache, + rotary_dim, + slot_mapping, + kv_cache, + block_size, + kv_cache_dtype, + k_scale, + v_scale, + q_fp8_out, + q_scale, + ) + + def fused_minimax_m3_qknorm_rope_kv_insert( qkv: torch.Tensor, q_norm_weight: torch.Tensor,