diff --git a/benchmarks/benchmark_gguf_dense_projection.py b/benchmarks/benchmark_gguf_dense_projection.py new file mode 100644 index 000000000..38eaddc47 --- /dev/null +++ b/benchmarks/benchmark_gguf_dense_projection.py @@ -0,0 +1,197 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Native SM70 dense microbenchmark on all Flash-Next TP4 weight shapes. + +Run control and candidate with their complete source runtimes and identical +arguments under the GPU lock. Model loading, attention and MTP are not run. +""" + +import argparse +import hashlib +import json +import statistics +from pathlib import Path + +import numpy as np +import torch +import vllm._C + +from vllm.model_executor.layers.quantization.gguf_dense_hmma_formats import decode, pack +from vllm.transformers_utils.gguf_tensor_reader import GGUFReader, dequantize + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--source-sha", required=True) + parser.add_argument("--samples", type=int, default=12) + args = parser.parse_args() + torch.set_num_threads(1) + torch.backends.cuda.matmul.allow_tf32 = False + rng = np.random.default_rng(20261010) + readers = [GGUFReader(f) for f in sorted(args.model.glob("*.gguf"))] + tensors = {t.name: t for reader in readers for t in reader.tensors} + groups = [] + for layer in range(48): + prefix = f"blk.{layer}." + if prefix + "attn_qkv.weight" in tensors: + members = [ + ("gdn_in", [prefix + "attn_qkv.weight", prefix + "attn_gate.weight"]), + ("gdn_out", [prefix + "ssm_out.weight"]), + ] + else: + members = [ + ("attn_in", [prefix + f"attn_{v}.weight" for v in ("q", "k", "v")]), + ("attn_out", [prefix + "attn_output.weight"]), + ] + for kind, names in members: + codes, high, stats, fmts, ns, refs, types = [], [], [], [], [], [], [] + x = None + for name in names: + tensor = tensors[name] + raw = tensor.data + if kind.endswith("_out"): + raw = np.ascontiguousarray(raw[:, : raw.shape[1] // 4]) + else: + # Two 256-dimensional KV heads are replicated over TP4. + n = ( + 256 + if name.endswith(("attn_k.weight", "attn_v.weight")) + else raw.shape[0] // 4 + ) + raw = np.ascontiguousarray(raw[:n]) + fmt, q, s, minimum, group = decode(raw, int(tensor.tensor_type)) + k = q.shape[1] + if x is None: + x = torch.from_numpy( + rng.normal(0, 0.5, (20, k)).astype(np.float16) + ).cuda() + c, h, sc = [ + torch.from_numpy(p).cuda() for p in pack(fmt, q, s, minimum, group) + ] + codes.append(c) + high.append(h) + stats.append(sc) + fmts.append(fmt) + ns.append(q.shape[0]) + types.append(int(tensor.tensor_type)) + w = torch.from_numpy( + dequantize(raw, int(tensor.tensor_type)).astype(np.float32) + ).cuda() + refs.append(x.float() @ w.T) + tiles = sum((n + 31) // 32 for n in ns) + groups.append( + dict( + kind=kind, + layer=layer, + x=x, + k=k, + codes=codes, + high=high, + stats=stats, + fmts=fmts, + ns=ns, + types=types, + out=torch.empty(20, sum(ns), device="cuda", dtype=torch.float16), + ref=torch.cat(refs, dim=1), + ws=torch.empty(tiles * 3 * 256, device="cuda"), + cnt=torch.zeros(tiles * 3, device="cuda", dtype=torch.int32), + ) + ) + + def run(g, m): + torch.ops._C.gguf_dense_segments_sm70_out( + g["x"][:m], + g["codes"], + g["high"], + g["stats"], + list(g["out"][:m].split(g["ns"], dim=1)), + g["fmts"], + g["ns"], + g["k"], + 1, + 8 if g["k"] == 1536 else 4, + g["ws"], + g["cnt"], + None, + ) + + report = dict( + scope="Native operator microbenchmark; no end-to-end claim", + source=args.source_sha, + core=str(vllm._C.__file__), + core_sha256=hashlib.sha256(Path(vllm._C.__file__).read_bytes()).hexdigest(), + torch=torch.__version__, + cuda=torch.version.cuda, + gpu=torch.cuda.get_device_name(), + measurements=[], + checks=[], + ) + for m in (5, 20): + for g in groups: + run(g, m) + torch.accelerator.synchronize() + y, ref = g["out"][:m].float(), g["ref"][:m] + error = float((y - ref).norm() / ref.norm()) + assert error < 0.003 + assert not bool(g["cnt"].any()) + report["checks"].append( + dict( + M=m, + kind=g["kind"], + layer=g["layer"], + official_rel_l2=error, + output_sha256=hashlib.sha256( + g["out"][:m].cpu().numpy().tobytes() + ).hexdigest(), + ) + ) + for kind in ("gdn_in", "gdn_out", "attn_in", "attn_out"): + subset = [g for g in groups if g["kind"] == kind] + stream = torch.cuda.Stream() + with torch.cuda.stream(stream): + for g in subset: + run(g, m) + torch.accelerator.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + for g in subset: + run(g, m) + observations = [] + for _ in range(args.samples): + for _ in range(3): + graph.replay() + start, stop = ( + torch.cuda.Event(enable_timing=True), + torch.cuda.Event(enable_timing=True), + ) + start.record() + for _ in range(8): + graph.replay() + stop.record() + stop.synchronize() + observations.append(start.elapsed_time(stop) * 1000 / 8) + record = dict( + M=m, + kind=kind, + layers=len(subset), + K=subset[0]["k"], + N=subset[0]["ns"], + types=sorted({tuple(g["types"]) for g in subset}), + weight_bytes=sum( + t.numel() + for g in subset + for name in ("codes", "high", "stats") + for t in g[name] + ), + median_us=statistics.median(observations), + samples_us=observations, + ) + report["measurements"].append(record) + print(json.dumps(record), flush=True) + args.output.write_text(json.dumps(report, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/csrc/sm70_turbomind/ops/gguf_dense_segments_sm70.cuh b/csrc/sm70_turbomind/ops/gguf_dense_segments_sm70.cuh index c8dcfe760..e8c9cf405 100644 --- a/csrc/sm70_turbomind/ops/gguf_dense_segments_sm70.cuh +++ b/csrc/sm70_turbomind/ops/gguf_dense_segments_sm70.cuh @@ -289,6 +289,8 @@ __device__ __forceinline__ void write_out(const Seg& sg, int t, int v, } } +#include "gguf_dense_weight_major_sm70.cuh" + template __global__ void __launch_bounds__(32 * W) dense_mv(Segs segs, const half* __restrict__ x, int ldx, int M, int K, @@ -367,6 +369,21 @@ __global__ void __launch_bounds__(32 * W) template void launch(const Segs& segs, const half* x, int ldx, int M, int K, int S, int G, int split, float* ws, int* cnt, int tiles, cudaStream_t st) { + // Small-M verification needs three token tiles, but each uses the same + // weights. Keep the original path for other shapes and split-K schedules. + int total_n = 0; + for (int i = 0; i < segs.nseg; ++i) total_n += segs.s[i].n; + if constexpr (W == 4 || W == 8) { + if (M == 20 && split == 1 && + ((K == 2560 && (total_n == 4096 || total_n == 3584)) || + (K == 1536 && total_n == 2560))) { + dense_weight_major + <<>>( + segs, x, ldx, M, K, S, G, split, ws, cnt); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return; + } + } const size_t xs_bytes = static_cast(W) * 256 * 16 + 1024; const size_t red_bytes = static_cast(W) * 256 * 4; const size_t smem = xs_bytes > red_bytes ? xs_bytes : red_bytes; diff --git a/csrc/sm70_turbomind/ops/gguf_dense_weight_major_sm70.cuh b/csrc/sm70_turbomind/ops/gguf_dense_weight_major_sm70.cuh new file mode 100644 index 000000000..bcb4529c3 --- /dev/null +++ b/csrc/sm70_turbomind/ops/gguf_dense_weight_major_sm70.cuh @@ -0,0 +1,176 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +#pragma once + +// Included after the segment decoder and output helpers. All twenty +// verification rows share one weight load and decode; their independent +// eight-row MMA accumulators retain the original K order and FP32 reduction +// tree. +template +__device__ __forceinline__ void decode_weight_major(const Ld& L, int st, + uint32_t (&hw)[16]) { + if constexpr (FMT != LUT4) { + decode(L, st, hw, nullptr); + } else { + const half2 s2 = lo2(word(L.sc, st)); + const uint32_t cr[4] = {L.c[st][0].x, L.c[st][0].y, L.c[st][0].z, + L.c[st][0].w}; +#pragma unroll + for (int c = 0; c < 4; ++c) { + // kIQ4 + 128 as bytes; nibble bit 3 picks the upper table half. + const uint32_t T0 = 0x3F2D1801u, T1 = 0x766A5D4Fu, T2 = 0xA6998D81u, + T3 = 0xF1D9C5B5u; + const half2 b1152 = h2(0x64806480u); + const uint32_t q = cr[c]; + const uint32_t qm = q & 0x77777777u; + uint32_t sel; + asm("lop3.b32 %0, %1, %2, %3, 0xEA;" + : "=r"(sel) + : "r"(q >> 1), "r"(0x44444444u), "r"(0x32103210u)); +#pragma unroll + for (int i = 0; i < 2; ++i) { + uint32_t lo, hi, v, a, b; + asm("prmt.b32 %0, %1, %2, %3;" + : "=r"(lo) + : "r"(T0), "r"(T1), "r"(qm >> (16 * i))); + asm("prmt.b32 %0, %1, %2, %3;" + : "=r"(hi) + : "r"(T2), "r"(T3), "r"(qm >> (16 * i))); + asm("prmt.b32 %0, %1, %2, %3;" + : "=r"(v) + : "r"(lo), "r"(hi), "r"(sel >> (16 * i))); + asm("prmt.b32 %0, %1, %2, 0x4140;" + : "=r"(a) + : "r"(v), "r"(0x64646464u)); + asm("prmt.b32 %0, %1, %2, 0x4342;" + : "=r"(b) + : "r"(v), "r"(0x64646464u)); + hw[4 * c + 2 * i] = u32(__hmul2(__hsub2(h2(a), b1152), s2)); + hw[4 * c + 2 * i + 1] = u32(__hmul2(__hsub2(h2(b), b1152), s2)); + } + } + } +} + +template +__device__ __forceinline__ void body_weight_major(const Seg& sg, int t, int g0, + int g1, int S, int G, + const half* x, int ldx, int M, + int K, uint4* xs, half2* lut, + float (&acc)[TM][8]) { + const int warp = threadIdx.x / 32, lane = threadIdx.x % 32; + const int r = (lane & 3) + ((lane & 16) ? 4 : 0); + int g = g0 + warp; + Ld A; + if (g < g1) load(A, sg, t, g, S, G, lane); + while (g < g1) { + Ld B; + const int gn = g + W; + if (gn < g1) load(B, sg, t, gn, S, G, lane); +#pragma unroll + for (int st = 0; st < 4; ++st) { + if (g * 4 + st < S) { + uint32_t hw[16]; + decode_weight_major(A, st, hw); +#pragma unroll + for (int b = 0; b < TM; ++b) { +#pragma unroll + for (int j = 0; j < 4; ++j) { + uint4 xa = make_uint4(0, 0, 0, 0); + if (b * 8 + r < M) + xa = __ldg(reinterpret_cast( + x + (b * 8 + r) * ldx + g * 128 + st * 32 + j * 8)); + mma(acc[b], xa.x, xa.y, hw[4 * j], hw[4 * j + 1]); + mma(acc[b], xa.z, xa.w, hw[4 * j + 2], hw[4 * j + 3]); + } + } + } + } + if (gn >= g1) break; + A = B; + g = gn; + } +} + +template +__global__ void __launch_bounds__(32 * W) + dense_weight_major(Segs segs, const half* x, int ldx, int M, int K, int S, + int G, int split, float* ws, int* cnt) { + extern __shared__ uint4 smem[]; + __shared__ int last; + const int tile = blockIdx.x, sp = blockIdx.y; + int si = 0; +#pragma unroll + for (int i = 1; i < MAXSEG; ++i) + if (i < segs.nseg && tile >= segs.s[i].tile0) si = i; + const Seg& sg = segs.s[si]; + const int t = tile - sg.tile0; + const int gps = (G + split - 1) / split; + const int g0 = min(G, sp * gps), g1 = min(G, g0 + gps); + float acc[TM][8] = {}; + half2* lut = reinterpret_cast(smem); + switch (sg.fmt) { + case Q4K: + body_weight_major(sg, t, g0, g1, S, G, x, ldx, M, K, smem, + lut, acc); + break; + case Q5K: + body_weight_major(sg, t, g0, g1, S, G, x, ldx, M, K, smem, + lut, acc); + break; + case Q6K: + body_weight_major(sg, t, g0, g1, S, G, x, ldx, M, K, smem, + lut, acc); + break; + case LUT4: + body_weight_major(sg, t, g0, g1, S, G, x, ldx, M, K, smem, + lut, acc); + break; + default: + body_weight_major(sg, t, g0, g1, S, G, x, ldx, M, K, smem, lut, + acc); + break; + } + float* red = reinterpret_cast(smem); + const int warp = threadIdx.x / 32, lane = threadIdx.x % 32; +#pragma unroll + for (int b = 0; b < TM; ++b) +#pragma unroll + for (int i = 0; i < 8; ++i) + red[(b * W + warp) * 256 + lane * 8 + i] = acc[b][i]; + __syncthreads(); +#pragma unroll + for (int b = 0; b < TM; ++b) { + for (int v = threadIdx.x; v < 256; v += 32 * W) { + float sum = 0.f; +#pragma unroll + for (int w = 0; w < W; ++w) sum += red[(b * W + w) * 256 + v]; + if (split == 1) + write_out(sg, t, v, sum, min(8, M - b * 8), segs.sgate, b * 8); + else + ws[((static_cast(b) * gridDim.x + tile) * split + sp) * 256 + + v] = sum; + } + } + if (split == 1) return; + __threadfence(); + __syncthreads(); + if (threadIdx.x == 0) { + last = atomicAdd(cnt + tile, 1) == split - 1; + if (last) __threadfence(); + } + __syncthreads(); + if (!last) return; +#pragma unroll + for (int b = 0; b < TM; ++b) + for (int v = threadIdx.x; v < 256; v += 32 * W) { + float sum = 0.f; + for (int sp2 = 0; sp2 < split; ++sp2) + sum += __ldcg( + ws + + ((static_cast(b) * gridDim.x + tile) * split + sp2) * 256 + + v); + write_out(sg, t, v, sum, min(8, M - b * 8), segs.sgate, b * 8); + } + if (threadIdx.x == 0) cnt[tile] = 0; +} diff --git a/docs/design/gguf_dense_weight_reuse.md b/docs/design/gguf_dense_weight_reuse.md new file mode 100644 index 000000000..21c9031f5 --- /dev/null +++ b/docs/design/gguf_dense_weight_reuse.md @@ -0,0 +1,114 @@ +# SM70 dense GGUF weight reuse + +The M20 projection previously launched three independent eight-row token +tiles. Each tile loaded and decoded the complete weight matrix. The new +weight-major kernel loads each group once and feeds three independent Volta +MMA accumulators. FP16 weight reconstruction, FP16 activations, FP32 K order, +warp reduction and final output rounding are preserved. + +Admission is limited to the measured M20, split1, W4/W8 shapes: + +| Projection | Local K | Local output widths | Layers | +| --- | ---: | --- | ---: | +| GDN qkv + z | 2560 | 2560 + 1536 | 36 | +| Attention q/k/v | 2560 | 3072 + 256 + 256 | 12 | +| GDN output | 1536 | 2560 | 36 | +| Attention output | 1536 | 2560 | 12 | + +The two 256-dimensional attention KV heads are replicated over TP4; cutting +each GGUF K/V tensor into four 128-row shards would benchmark the wrong +shape. M5, other M, split-K and other dimensions retain the old kernel. +No weight expansion, second resident bank or additional graph allocation is +introduced. Existing segment dispatch and scratch ownership are retained. + +## Community references and structural screens + +[Marlin](https://github.com/IST-DASLab/marlin) discusses register reuse, +dequantization/MMA scheduling and offline layouts. Its Ampere asynchronous +loads are not used here. [FLUTE](https://github.com/HanGuo97/flute) studies +lookup layout and shared-memory lookup costs. +[GemLite](https://github.com/dropbox/gemlite) separates small-batch algorithms +and packing widths. These are design references; no external implementation +or license-incompatible source is incorporated. The decoder derives from +the existing Apache-2.0 SM70 projection implementation. + +The tested alternatives preserve the current arithmetic, including +FP16 coefficient rounding. They are research screens, not model results: + +| Screen | M5 total ms | M20 total ms | Decision | +| --- | ---: | ---: | --- | +| Matched installed control | 1.2466 | 3.1011 | Reference | +| Register IQ4 lookup | 1.2579 | 3.0651 | No useful overall gain | +| Register activation exchange | 1.3301 | 3.1930 | Reject | +| Lossless Q6/IQ4 byte expansion | 1.3313 | 3.1878 | Reject; more storage and traffic | +| Three token groups sharing weights | 1.3309 | 2.4392 | Select M20 only | +| Format-combination specialization | 1.3210 | 3.0068 | Reject as an overall policy | + +Format specialization reduces registers for LUT-only kernels from about 144 +to 95, but this does not improve the complete mixed-format M5 workload. +Another screen interleaves four decoded half2 operands with each pair of +MMA instructions; its M5 improvement is negligible and it is not selected. +Lower register counts and fewer lookup instructions alone are not latency +evidence. + +An initial CUDA 12.0 screen is retained separately. Its control was faster +than the installed CUDA 12.8 kernel, so those cross-compiler differences are +not counted as optimization gains. All results in the table use CUDA 12.8. + +## Microbenchmark contract + +Four V100-SXM2-32GB GPUs with full NV2 connectivity; one locked GPU is used +for these independent TP4 shard projections. Torch 2.10.0+cu128 and NVCC +12.8.93; SM70, FP16 operands and FP32 accumulation. No fast-math flag is +introduced. The weights come from all 48 layers of Flash-Next GSQ-RCO +IQ3_S, retaining their actual mixed Q4_K, Q5_K, Q6_K and IQ4_XS types. + +Each graph rotates through 12 or 36 distinct layer banks, substantially more +than V100's 6 MB L2 capacity. Graphs are warmed before CUDA-event timing. +Research arms are interleaved in forward/reverse order, ten observations per +arm, with eight graph replays per observation. Warp counts and split counts +are fixed to production choices, not selected from a tuning sweep. + +| M20 role | Control us/layer | Weight-major us/layer | +| --- | ---: | ---: | +| GDN input | 43.68 | 36.97 | +| GDN output | 22.74 | 15.26 | +| Attention input | 38.71 | 32.80 | +| Attention output | 20.47 | 13.77 | + +The weighted total improves 3.1011 to 2.4392 ms, about 21.3%. This is the sum +of 96 independently supplied dense projections, not end-to-end C4 latency. +All banks contain 540.34 MB of packed weight planes per rank. M20 previously +issued approximately 1.62 GB of weight loads across its three token tiles; +the new schedule issues approximately 540 MB. Cache hits can change measured +DRAM traffic; these are logical issued bytes, not Nsight DRAM counters. +Unique-weight bytes per service time improve approximately 174 to 222 GB/s. + +The M5 installed control measures 1.2466 ms, about 433 GB/s of logical +weight traffic. The historical 2.3-ms dense service sum comes from a different +profile and execution context; it cannot be replaced with this isolated +microbenchmark total or subtracted from the 17.4-ms round baseline. + +## Correctness and reproducibility + +The initial six-screen CUDA 12.8 comparison passes 2,304 real-shard output +bit comparisons, across M5/M20 and two different activation inputs. Maximum +relative L2 error against official FP32 GGUF dequantization is 0.0008975; +candidate outputs are bit-identical to the installed control. Only the M20 +weight-major schedule is incorporated into native dispatch. + +The additional GPU suite covers all six dense source types, real admitted +dimensions, mixed segments, strided output, shared-gate rounding, poisoned +output/scratch, and changed-input CUDA graph replay. The reproducible native +benchmark uses only the normal CUDA extension: + +```bash +flock /tmp/gpu0-3.lock python benchmarks/benchmark_gguf_dense_projection.py \ + --model "$GGUF_MODEL_DIR" --source-sha "$SOURCE_SHA" --output result.json +``` + +Run the same command with complete control and candidate source runtimes. +The JSON includes the native library hash, output fingerprints, official +dequantization error, exact N/K/type lists, logical weight bytes and every +timing observation. No whole model, prefill, attention, MTP sampling or +acceptance benchmark is run by this command. diff --git a/tests/kernels/quantization/test_gguf_dense_weight_major.py b/tests/kernels/quantization/test_gguf_dense_weight_major.py new file mode 100644 index 000000000..5f7bb5038 --- /dev/null +++ b/tests/kernels/quantization/test_gguf_dense_weight_major.py @@ -0,0 +1,117 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""M20 weight reuse: official references, strided output and graph rewrites.""" + +import gguf +import numpy as np +import pytest +import torch + +from vllm.model_executor.layers.quantization.gguf_dense_hmma_formats import decode, pack + + +def _raw(source, n, k): + block, size = gguf.GGML_QUANT_SIZES[gguf.GGMLQuantizationType(source)] + rng = np.random.default_rng(2300 + source) + data = rng.integers(0, 256, (n, k // block, size), dtype=np.uint8) + offset = 208 if source == 14 else 0 + data[..., offset : offset + 2] = np.array([0.0007], np.float16).view(np.uint8) + if source in (12, 13): + data[..., 2:4] = np.array([0.0003], np.float16).view(np.uint8) + return data.reshape(n, -1) + + +@pytest.mark.parametrize( + "sources", [(8,), (12,), (13,), (14,), (20,), (23,), (12, 14), (23, 14, 13)] +) +@pytest.mark.parametrize("m", [5, 20]) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_actual_k_changed_graph(sources, m): + if torch.cuda.get_device_capability() != (7, 0): + pytest.skip("SM70 required") + import vllm._C # noqa: F401 + + k = 2560 + n_shards = {1: [4096], 2: [2560, 1536], 3: [3072, 256, 256]}[len(sources)] + widths, fmts, codes, high, scales, official = [], [], [], [], [], [] + for source, n in zip(sources, n_shards): + raw = _raw(source, n, k) + fmt, q, s, minimum, group = decode(raw, source) + payload = [torch.from_numpy(p).cuda() for p in pack(fmt, q, s, minimum, group)] + widths.append(n) + fmts.append(fmt) + codes.append(payload[0]) + high.append(payload[1]) + scales.append(payload[2]) + official.append( + torch.from_numpy( + gguf.quants.dequantize(raw, gguf.GGMLQuantizationType(source)) + ).cuda() + ) + x = torch.randn(m, k, device="cuda", dtype=torch.float16) + storage = torch.empty(m, sum(widths) + 7, device="cuda", dtype=torch.float16) + out = storage[:, 3 : 3 + sum(widths)] + views = list(out.split(widths, dim=1)) + tiles = sum((v + 31) // 32 for v in widths) + ws = torch.full((tiles * ((m + 7) // 8) * 256,), float("nan"), device="cuda") + counter = torch.zeros(tiles * ((m + 7) // 8), device="cuda", dtype=torch.int32) + gate = torch.randn(m, device="cuda", dtype=torch.float16) + + def run(): + torch.ops._C.gguf_dense_segments_sm70_out( + x, codes, high, scales, views, fmts, widths, k, 1, 4, ws, counter, gate + ) + + run() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run() + for trial in range(4): + x.normal_(mean=trial / 20) + out.fill_(float("nan")) + graph.replay() + reference = torch.cat([x.float() @ w.T for w in official], dim=1).half() + reference = (reference.float() * torch.sigmoid(gate.float())[:, None]).half() + torch.accelerator.synchronize() + assert torch.isfinite(out).all() + assert ( + float((out.float() - reference.float()).norm() / reference.float().norm()) + < 0.003 + ) + assert not counter.any() + + +@pytest.mark.parametrize("source", [12, 13, 14]) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_output_projection_k1536(source): + if torch.cuda.get_device_capability() != (7, 0): + pytest.skip("SM70 required") + import vllm._C # noqa: F401 + + raw = _raw(source, 2560, 1536) + fmt, q, s, minimum, group = decode(raw, source) + payload = [torch.from_numpy(p).cuda() for p in pack(fmt, q, s, minimum, group)] + x = torch.randn(20, 1536, device="cuda", dtype=torch.float16) + y = torch.empty(20, 2560, device="cuda", dtype=torch.float16) + ws = torch.empty(80 * 3 * 256, device="cuda") + counter = torch.zeros(80 * 3, device="cuda", dtype=torch.int32) + torch.ops._C.gguf_dense_segments_sm70_out( + x, + [payload[0]], + [payload[1]], + [payload[2]], + [y], + [fmt], + [2560], + 1536, + 1, + 8, + ws, + counter, + None, + ) + w = torch.from_numpy( + gguf.quants.dequantize(raw, gguf.GGMLQuantizationType(source)) + ).cuda() + reference = x.float() @ w.T + assert float((y.float() - reference).norm() / reference.norm()) < 0.003