[Kernel]Add norm fused kernel for Step-5 Preview and Muse Spark Open (#62132)
Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai>
This commit is contained in:
Jee Jee Li
@@ -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)
|
||||
|
||||
@@ -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()
|
||||
@@ -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 <algorithm>
|
||||
#include <cmath>
|
||||
#include <cuda_runtime.h>
|
||||
#include <type_traits>
|
||||
|
||||
#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 <cuda_fp8.h>
|
||||
#include "quantization/w8a8/fp8/nvidia/quant_utils.cuh"
|
||||
#else
|
||||
#include <hip/hip_fp8.h>
|
||||
#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<uint32_t>(lg);
|
||||
uint64_t const m = ((1ull << p) + static_cast<uint64_t>(denom) - 1) /
|
||||
static_cast<uint64_t>(denom);
|
||||
return FastDiv{static_cast<uint32_t>(m), p - 32, denom};
|
||||
}
|
||||
|
||||
__device__ __forceinline__ void fastDivmod(int& quo, int& rem, int src,
|
||||
FastDiv const fd) {
|
||||
int const q =
|
||||
static_cast<int>(__umulhi(static_cast<uint32_t>(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 <int LANES>
|
||||
__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 <typename scalar_t, int kElems>
|
||||
__device__ __forceinline__ uint4 loadElems(scalar_t const* __restrict__ src,
|
||||
float (&elems)[kElems]) {
|
||||
using Converter = vllm::_typeConvert<scalar_t>;
|
||||
using T2 = typename Converter::packed_hip_type;
|
||||
uint4 raw = *reinterpret_cast<uint4 const*>(src);
|
||||
auto const* p = reinterpret_cast<T2 const*>(&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 <typename rope_t, int kPairs>
|
||||
__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<rope_t, float>) {
|
||||
float4 const cv = VLLM_LDG(reinterpret_cast<float4 const*>(cos_ptr));
|
||||
float4 const sv = VLLM_LDG(reinterpret_cast<float4 const*>(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<rope_t>;
|
||||
using T2 = typename Converter::packed_hip_type;
|
||||
uint2 const cv = VLLM_LDG(reinterpret_cast<uint2 const*>(cos_ptr));
|
||||
uint2 const sv = VLLM_LDG(reinterpret_cast<uint2 const*>(sin_ptr));
|
||||
auto const* cp = reinterpret_cast<T2 const*>(&cv);
|
||||
auto const* sp = reinterpret_cast<T2 const*>(&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 <typename scalar_t, int kElems>
|
||||
__device__ __forceinline__ uint4 packElems(float const (&elems)[kElems]) {
|
||||
using Converter = vllm::_typeConvert<scalar_t>;
|
||||
using T2 = typename Converter::packed_hip_type;
|
||||
uint4 raw;
|
||||
auto* p = reinterpret_cast<T2*>(&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 <typename scalar_t, typename cache_t, Fp8KVCacheDataType kv_dt,
|
||||
int kElems>
|
||||
__device__ __forceinline__ void storeCache(cache_t* __restrict__ dst,
|
||||
uint4 const& packed,
|
||||
float const inv_scale) {
|
||||
if constexpr (kv_dt == Fp8KVCacheDataType::kAuto) {
|
||||
*reinterpret_cast<uint4*>(dst) = packed;
|
||||
} else {
|
||||
using Converter = vllm::_typeConvert<scalar_t>;
|
||||
using HipT = typename Converter::hip_type;
|
||||
auto const* rounded = reinterpret_cast<HipT const*>(&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<uint2*>(dst) = *reinterpret_cast<uint2 const*>(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<cache_t, float, kv_dt>(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 <typename scalar_t, typename cache_t, Fp8KVCacheDataType kv_dt,
|
||||
typename rope_t, int HEAD_DIM, int HEADS_PER_GROUP, bool kDoRope,
|
||||
bool kInsertKV>
|
||||
__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<BFloat16> 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<scalar_t, c10::BFloat16>) {
|
||||
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<WarpMask>(0xffffffffu)
|
||||
: ((WarpMask{1} << kLanes) - WarpMask{1});
|
||||
unsigned const group_shift =
|
||||
(threadIdx.x & (kWarpSize - 1)) & ~static_cast<unsigned>(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<int64_t>(tokenIdx) * qkv_stride +
|
||||
static_cast<int64_t>(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<uint4 const*>(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<rope_t, kPairs>(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<int>(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<scalar_t>;
|
||||
using T2 = typename Converter::packed_hip_type;
|
||||
auto const* p = reinterpret_cast<T2 const*>(&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<kLanes>(sumsq, group_mask);
|
||||
float const rms_rcp =
|
||||
rsqrtf(sumsq / static_cast<float>(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<scalar_t, kElems>(elems);
|
||||
if (q_fp8_out && slot < nq) {
|
||||
int64_t const q_off =
|
||||
(static_cast<int64_t>(tokenIdx) * nq + slot) * HEAD_DIM + dimBase;
|
||||
storeCache<scalar_t, uint8_t, Fp8KVCacheDataType::kFp8E4M3, kElems>(
|
||||
q_fp8_out + q_off, packed, q_inv);
|
||||
} else {
|
||||
*reinterpret_cast<uint4*>(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<int64_t>(cache_block) * kv_s_block +
|
||||
head * kv_s_head +
|
||||
static_cast<int64_t>(cache_tok) * kv_s_token +
|
||||
(isK ? 0 : HEAD_DIM) + dimBase;
|
||||
storeCache<scalar_t, cache_t, kv_dt, kElems>(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 <typename scalar_t, typename cache_t, Fp8KVCacheDataType kv_dt,
|
||||
typename rope_t, int HEAD_DIM, int HEADS_PER_GROUP>
|
||||
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<int64_t>(num_tokens) * chunks_per_token * kLanes;
|
||||
if (total_lanes == 0) return;
|
||||
|
||||
constexpr int kBlockSize = 256;
|
||||
int const grid =
|
||||
static_cast<int>((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><<<grid, kBlockSize, 0, stream>>>( \
|
||||
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 <typename scalar_t, typename cache_t, Fp8KVCacheDataType kv_dt>
|
||||
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<int64_t>(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<int64_t>(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<int64_t>(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<scalar_t, cache_t, kv_dt, RT, HD, HPG>( \
|
||||
qkv, reinterpret_cast<RT const*>(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<st, CACHE_T, KV_DTYPE>( \
|
||||
reinterpret_cast<st*>(qkv.data_ptr()), \
|
||||
do_rope ? cos_sin_cache->data_ptr() : nullptr, rope_is_fp32, \
|
||||
reinterpret_cast<int64_t const*>(positions.data_ptr()), \
|
||||
insert_kv ? reinterpret_cast<int64_t const*>(slot_mapping->data_ptr()) \
|
||||
: nullptr, \
|
||||
insert_kv ? reinterpret_cast<CACHE_T*>(kv_cache->data_ptr()) : nullptr, \
|
||||
needs_scales ? reinterpret_cast<float const*>(k_scale->data_ptr()) \
|
||||
: nullptr, \
|
||||
needs_scales ? reinterpret_cast<float const*>(v_scale->data_ptr()) \
|
||||
: nullptr, \
|
||||
q_fp8_out ? reinterpret_cast<uint8_t*>(q_fp8_out->data_ptr()) : nullptr, \
|
||||
q_scale ? reinterpret_cast<float const*>(q_scale->data_ptr()) : nullptr, \
|
||||
static_cast<float>(eps), static_cast<int>(rotary_dim), num_tokens, nq, \
|
||||
nkv, hd, qkv_stride, static_cast<int>(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<torch::stable::Tensor> cos_sin_cache, // None => NoPE layer
|
||||
int64_t rotary_dim,
|
||||
std::optional<torch::stable::Tensor> slot_mapping, // [N] i64
|
||||
std::optional<torch::stable::Tensor> kv_cache, // [nb, nkv, bs, 2*hd]
|
||||
int64_t block_size, const std::string& kv_cache_dtype,
|
||||
std::optional<torch::stable::Tensor> k_scale, // [1] fp32, fp8 cache only
|
||||
std::optional<torch::stable::Tensor> v_scale,
|
||||
std::optional<torch::stable::Tensor> q_fp8_out,
|
||||
std::optional<torch::stable::Tensor> 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<int>(qkv.size(0));
|
||||
int const nq = static_cast<int>(num_heads);
|
||||
int const nkv = static_cast<int>(num_kv_heads);
|
||||
int const hd = static_cast<int>(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<int64_t>(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<int64_t>(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
|
||||
@@ -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 <torch/csrc/stable/macros.h>
|
||||
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
#include <algorithm>
|
||||
#include <type_traits>
|
||||
|
||||
namespace vllm {
|
||||
|
||||
namespace {
|
||||
|
||||
constexpr int kWarpSize = 32;
|
||||
constexpr int kThreadsPerBlock = 128;
|
||||
|
||||
template <typename scalar_t, typename weight_t, typename rope_t, int HEAD_DIM,
|
||||
bool WIDE_192_ROTARY = false, bool FP8_CACHE = false>
|
||||
__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<scalar_t, torch::headeronly::BFloat16> ||
|
||||
std::is_same_v<weight_t, torch::headeronly::BFloat16> ||
|
||||
std::is_same_v<rope_t, torch::headeronly::BFloat16>) {
|
||||
return;
|
||||
} else {
|
||||
#endif
|
||||
using ScalarConverter = _typeConvert<scalar_t>;
|
||||
using WeightConverter = _typeConvert<weight_t>;
|
||||
using RopeConverter = _typeConvert<rope_t>;
|
||||
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<const Scalar*>(qkv_void);
|
||||
auto* q_out = static_cast<Scalar*>(q_out_void);
|
||||
auto* k_out = static_cast<Scalar*>(k_out_void);
|
||||
auto* v_out = static_cast<Scalar*>(v_out_void);
|
||||
const auto* q_weight = static_cast<const Weight*>(q_weight_void);
|
||||
const auto* k_weight = static_cast<const Weight*>(k_weight_void);
|
||||
const auto* cos = static_cast<const Rope*>(cos_void);
|
||||
const auto* sin = static_cast<const Rope*>(sin_void);
|
||||
auto* k_cache = static_cast<Scalar*>(k_cache_void);
|
||||
auto* v_cache = static_cast<Scalar*>(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<int64_t>(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<int>(global_head -
|
||||
static_cast<int64_t>(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<int64_t>(head) * HEAD_DIM;
|
||||
output = q_out;
|
||||
output_offset =
|
||||
token * q_token_stride + static_cast<int64_t>(head) * HEAD_DIM;
|
||||
} else if (is_k) {
|
||||
input_offset += static_cast<int64_t>(num_q_heads + head) * HEAD_DIM;
|
||||
output = k_out;
|
||||
output_offset =
|
||||
token * k_token_stride + static_cast<int64_t>(head) * HEAD_DIM;
|
||||
} else {
|
||||
input_offset +=
|
||||
static_cast<int64_t>(num_q_heads + num_kv_heads + head) * HEAD_DIM;
|
||||
output = v_out;
|
||||
output_offset =
|
||||
token * v_token_stride + static_cast<int64_t>(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<int64_t>(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<uint8_t*>(v_cache_void)[cache_offset + dim] =
|
||||
fp8::scaled_convert<uint8_t, float,
|
||||
Fp8KVCacheDataType::kFp8E4M3>(
|
||||
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<float>(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<uint8_t*>(k_cache_void)[cache_offset + dim] =
|
||||
fp8::scaled_convert<uint8_t, float,
|
||||
Fp8KVCacheDataType::kFp8E4M3>(
|
||||
ScalarConverter::convert(value), *k_scale);
|
||||
} else {
|
||||
k_cache[cache_offset + dim] = value;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#if !defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
template <typename scalar_t, typename weight_t, typename rope_t, bool FP8_CACHE>
|
||||
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<int>(k_cache.size(1));
|
||||
const int block_size_shift =
|
||||
(block_size & (block_size - 1)) == 0
|
||||
? __builtin_ctz(static_cast<unsigned int>(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<scalar_t, weight_t, rope_t, HEAD_DIM, \
|
||||
WIDE_192_ROTARY, FP8_CACHE> \
|
||||
<<<grid, block, 0, stream>>>( \
|
||||
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<const int64_t*>(positions.data_ptr()), \
|
||||
reinterpret_cast<const int64_t*>(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<int>(qkv.size(0)), \
|
||||
static_cast<int>(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<torch::stable::Tensor> k_scale,
|
||||
std::optional<torch::stable::Tensor> 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<input_t, norm_weight_t,
|
||||
rope_cache_t,
|
||||
decltype(fp8_tag)::value>(
|
||||
qkv, q_out, k_out, v_out, q_weight, k_weight, cos,
|
||||
sin, positions, slot_mapping, k_cache, v_cache,
|
||||
static_cast<int>(head_dim),
|
||||
static_cast<int>(num_q_heads),
|
||||
static_cast<int>(num_kv_heads),
|
||||
static_cast<int>(rotary_pairs),
|
||||
static_cast<float>(epsilon),
|
||||
static_cast<float>(norm_weight_bias), stream,
|
||||
fp8_cache
|
||||
? static_cast<const float*>(k_scale->data_ptr())
|
||||
: nullptr,
|
||||
fp8_cache
|
||||
? static_cast<const float*>(v_scale->data_ptr())
|
||||
: nullptr);
|
||||
};
|
||||
if (fp8_cache) {
|
||||
launch(std::true_type{});
|
||||
} else {
|
||||
launch(std::false_type{});
|
||||
}
|
||||
});
|
||||
});
|
||||
});
|
||||
STD_CUDA_KERNEL_LAUNCH_CHECK();
|
||||
}
|
||||
|
||||
} // namespace vllm
|
||||
@@ -0,0 +1,100 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
#include <cuda_runtime.h>
|
||||
#include <torch/csrc/stable/accelerator.h>
|
||||
#include <torch/csrc/stable/tensor.h>
|
||||
#include <torch/headeronly/core/ScalarType.h>
|
||||
|
||||
#include <algorithm>
|
||||
|
||||
#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 <typename scalar_t>
|
||||
__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<int64_t>(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<kMuseSoftCapVecSize>(
|
||||
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<float>(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<int>(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<int64_t>(
|
||||
kTargetBlocks / num_tokens, 1, std::max<int64_t>(blocks_to_cover, 1));
|
||||
|
||||
dim3 grid(static_cast<unsigned>(num_tokens),
|
||||
static_cast<unsigned>(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<scalar_t>
|
||||
<<<grid, block, 0, stream>>>(out.mutable_data_ptr<float>(),
|
||||
input.const_data_ptr<scalar_t>(), d,
|
||||
static_cast<float>(multiplier),
|
||||
static_cast<float>(soft_cap_temp));
|
||||
});
|
||||
}
|
||||
+33
@@ -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<torch::stable::Tensor> k_scale,
|
||||
std::optional<torch::stable::Tensor> 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<torch::stable::Tensor> cos_sin_cache, int64_t rotary_dim,
|
||||
std::optional<torch::stable::Tensor> slot_mapping,
|
||||
std::optional<torch::stable::Tensor> kv_cache, int64_t block_size,
|
||||
const std::string& kv_cache_dtype,
|
||||
std::optional<torch::stable::Tensor> k_scale,
|
||||
std::optional<torch::stable::Tensor> v_scale,
|
||||
std::optional<torch::stable::Tensor> q_fp8_out,
|
||||
std::optional<torch::stable::Tensor> 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<int64_t, torch::stable::Tensor> 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,
|
||||
|
||||
@@ -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));
|
||||
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user