diff --git a/astrai/extension/backend/linear.py b/astrai/extension/backend/linear.py index 4c85c82a..38fd275c 100644 --- a/astrai/extension/backend/linear.py +++ b/astrai/extension/backend/linear.py @@ -1,9 +1,9 @@ """Inference-only dispatch for AstrAI linear layers. The CUDA GEMV path is deliberately narrow: automatic selection is enabled -only for single-row BF16 shapes measured to beat ``F.linear`` on a supported -architecture. Every training, prefill, unsupported-layout, and unmeasured -call falls back to PyTorch. +only for small decode batches and BF16 shapes measured to beat ``F.linear`` +on a supported architecture. Every training, prefill, unsupported-layout, +and unmeasured call falls back to PyTorch. """ import logging @@ -30,20 +30,72 @@ logger = logging.getLogger(__name__) # Shape keys are (N, K) for Y[M, N] = X[M, K] @ W[N, K].T. A band is -# automatic only after both the per-shape >=5% and end-to-end decode >=3% -# gates pass and checkpoint greedy output remains stable. M=1 and M=8 remain -# empty on SM89; the safe M=2/4 bands improve real-engine throughput by -# 11.8-14.0%. +# automatic only after both the per-shape >=5% and projection-chain/engine +# >=3% gates pass and output argmax remains stable. M=1 is limited to OPT 1.3B; +# M=8 remains empty because at least one projection in each measured family +# misses the per-shape gate even when its aggregate chain result is positive. +_COMMON_TRANSFORMER_SM89_SHAPES = frozenset( + { + (1024, 4096), # LLaMA 3 8B K/V + (4096, 4096), # LLaMA 2/3 7B/8B Q/O + (11008, 4096), # LLaMA 2 7B gate/up + (4096, 11008), # LLaMA 2 7B down + (14336, 4096), # LLaMA 3 8B gate/up + (4096, 14336), # LLaMA 3 8B down + (5120, 5120), # LLaMA 2 13B Q/K/V/O + (13824, 5120), # LLaMA 2 13B gate/up + (5120, 13824), # LLaMA 2 13B down + (16384, 4096), # GPT-NeoX MLP up + (4096, 16384), # GPT-NeoX MLP down + } +) +_COMMON_TRANSFORMER_SM89_M4_SHAPES = _COMMON_TRANSFORMER_SM89_SHAPES - { + (4096, 4096), + (11008, 4096), + (4096, 11008), +} +_QWEN2_7B_SM89_SHAPES = frozenset( + { + (512, 3584), # K/V + (3584, 3584), # Q/O + (18944, 3584), # gate/up + (3584, 18944), # down + } +) +_LLAMA3_70B_SM89_SHAPES = frozenset( + { + (1024, 8192), # K/V + (8192, 8192), # Q/O + (28672, 8192), # gate/up + (8192, 28672), # down + } +) +_OPT_1_3B_SM89_SHAPES = frozenset( + { + (2048, 2048), # Q/K/V/O + (8192, 2048), # MLP up + (2048, 8192), # MLP down + } +) + _AUTO_GEMV_SHAPES: dict[tuple[int, int], dict[int, frozenset[tuple[int, int]]]] = { (8, 9): { - 2: frozenset( + 1: _OPT_1_3B_SM89_SHAPES, + 2: _COMMON_TRANSFORMER_SM89_SHAPES + | _QWEN2_7B_SM89_SHAPES + | _LLAMA3_70B_SM89_SHAPES + | _OPT_1_3B_SM89_SHAPES + | frozenset( { (256, 1536), (1536, 1536), (100000, 1536), } ), - 4: frozenset({(256, 1536), (1536, 1536)}), + 4: _COMMON_TRANSFORMER_SM89_M4_SHAPES + | _QWEN2_7B_SM89_SHAPES + | _LLAMA3_70B_SM89_SHAPES + | frozenset({(256, 1536), (1536, 1536)}), } } _AUTO_GEMV_M = frozenset( diff --git a/benchmarks/infraswe/README.md b/benchmarks/infraswe/README.md new file mode 100644 index 00000000..fd0be307 --- /dev/null +++ b/benchmarks/infraswe/README.md @@ -0,0 +1,39 @@ +# InfraSWE Draft: common-shape BF16 GEMV + +This directory binds the common LLaMA/GPT-NeoX GEMV benchmark to AstrAI as an +explicit repository target. AstrAI is not one of InfraSWE v0.5's built-in +projects, so selecting a built-in default would score the change against the +wrong host contract. The Draft remains `D3-contract-proposed`; it does not +claim maintainer review, sealing, hidden-probe completion, or an official +ProjectFit. + +The Draft was validated and resolved before opening the PR with InfraSWE commit +`811bc775ed5b3a6ec853219245f3469f78818020`: + +```bash +PYTHONPATH=src .venv/bin/infraswe draft validate \ + /path/to/AstrAI/benchmarks/infraswe/astrai-gemv-common-draft.json + +PYTHONPATH=src .venv/bin/infraswe draft resolve \ + --local-draft \ + /path/to/AstrAI/benchmarks/infraswe/astrai-gemv-common-draft.json \ + --output /tmp/astrai-gemv-common-draft-resolution.json +``` + +The candidate and contract digests bind the ordered source, tests, +documentation, and raw benchmark evidence. The required comparison cell is one +NVIDIA L20 (`sm_89`); optional A100 and H100 cells are explicitly untested. +Compilation happens before timed cases. The checked-in kernel suite uses 21 +paired/interleaved samples, while the distinct-weight synthetic chains include +the guarded Python dispatcher. The extended portfolio adds Qwen2-7B, LLaMA 3 +70B, and OPT-1.3B shapes at M=1/2/4/8, plus an interleaved +128-thread-versus-256-thread CTA comparison. Neither suite is presented as +whole-model throughput. + +Applying InfraSWE's frozen `project-fit-kernel-v0.5` formula to the visible +evidence yields a diagnostic ProjectFit of **92.64/100** and BenchmarkTrust of +**95.87/100**. The machine-readable rationale is +`benchmarks/results/gemv_common_l20_sm89_infraswe_score.json`. Both numbers are +non-official: official ProjectFit remains unresolved until the Draft is sealed, +five or more fresh-process replays and system traces exist, hidden probes are +complete, and the evidence manifest is verified. diff --git a/benchmarks/infraswe/astrai-gemv-common-draft.json b/benchmarks/infraswe/astrai-gemv-common-draft.json new file mode 100644 index 00000000..d44e9d65 --- /dev/null +++ b/benchmarks/infraswe/astrai-gemv-common-draft.json @@ -0,0 +1,103 @@ +{ + "schema_version": "0.5", + "draft": { + "id": "astrai-common-shape-bf16-gemv-v1", + "revision": 2, + "state": "D3-contract-proposed", + "created_by": "0z5a" + }, + "target": { + "mode": "repository", + "repository": "https://github.com/ViperEkura/AstrAI", + "revision": "sha256:59c45de7cb0356e9661cfb73d38bd11c888a8b8d564e9123a90c9fa3910cc498", + "project_profile_sha256": "sha256:953067c47f5298f06846c8fe873db55afae49b6588a81926e49d9ef6b0994e08" + }, + "candidate": { + "kind": "git-diff", + "revision": "sha256:d9ae3993caf937517b121d4fe855f2c5172cf38f69aa4bda10bbb8daad52a730", + "intent": "add-fastpath", + "implementation_kind": "cuda-native", + "entrypoints": [ + "astrai.extension.bf16_gemv", + "astrai.extension.backend.linear", + "scripts/tools/benchmark_gemv_common.py" + ], + "operator_family": "dense-gemm", + "phase": "inference", + "backend": "cuda", + "primary_host_candidate": "astrai" + }, + "baseline": { + "mode": "target-head", + "revision": "sha256:59c45de7cb0356e9661cfb73d38bd11c888a8b8d564e9123a90c9fa3910cc498" + }, + "deployment": { + "workload_portfolio": { + "id": "astrai-common-shape-gemv-chains-v1", + "sha256": "sha256:c3cf734200844d7b8308a4803f456cd9a3f0d7e19839e7d3e74778f6d1307e4d", + "path": "docs/benchmarks/gemv_common_extended_chain_l20_sm89.json" + }, + "required_cells": [ + "cuda-sm89-nvidia-l20" + ], + "optional_cells": [ + "cuda-sm80-a100", + "cuda-sm90-h100" + ], + "request_or_step_protocol": { + "id": "astrai-gemv-common-benchmark-v1", + "sha256": "sha256:74808a77e551ffa5f77e3256930f2647ebd080828cfdbdcc2770f9d1fa298ff6", + "path": "docs/developer/decode_linear_benchmark.md" + } + }, + "retrieval": { + "enabled": true, + "corpus_cutoff": "2026-09-02T09:45:00Z", + "sources": [ + "target-code", + "merged-prs", + "rejected-prs", + "review-comments", + "ci-failures" + ], + "precedent_set_sha256": "sha256:9034271948fd791fb53d5099c2bc246acfd4e27e45926857220b48fd841389b9" + }, + "acceptance_contract": { + "status": "proposed", + "path": "tests/extension/test_gemv.py,tests/extension/test_linear_dispatch.py", + "sha256": "sha256:15ffa9e1378d7cb1b06f018f9e1aac16bf56f14446271d51129f86c13f3827bb", + "probe_set_sha256": "sha256:beafc95a2a4aaa587079f1938db8c0fb115d63c2b0c3b185debc9a6132397653", + "hidden_probe_policy_sha256": "sha256:e6bdf6b5a307c9ebdaf643d18ae88b51c6e554d7d9df4427c2f7ba0eb5cc8e93" + }, + "project_objectives": { + "edge_ecosystem_policy": "experimental", + "profile_set_sha256": "sha256:22da954f941ba912f60d79b7a58121436d4ba97e64f0182f66b8274080c3a75a" + }, + "benchmark_loop": { + "fast_stage_max_official_fraction": 0.05, + "affected_stage_max_official_fraction": 0.2, + "official_replays": 7, + "early_exit_on_hard_gate": true, + "affected_case_selection": "required", + "benchmark_budget_policy_id": "draft-staged-budget-v0.5", + "evidence_policy_id": "v0.4-evidence-ladder-plus-seal-v0.5", + "precompile": { + "mode": "auto", + "trigger": "when-compilation-required", + "cache_policy": "content-addressed-evidence-identity", + "cache_miss_action": "precompile-before-timed-cases", + "timing_phases": [ + "precompile", + "cold-start", + "steady-state" + ], + "steady_state_compile_allowed": false + } + }, + "scoring": { + "formula_template_id": "project-fit-kernel-v0.5", + "provisional_scoring_allowed": true, + "official_scoring_requires_seal": true, + "project_season": "astrai-2026q3" + } +} diff --git a/benchmarks/results/gemv_common_l20_sm89_infraswe_score.json b/benchmarks/results/gemv_common_l20_sm89_infraswe_score.json new file mode 100644 index 00000000..b326c35c --- /dev/null +++ b/benchmarks/results/gemv_common_l20_sm89_infraswe_score.json @@ -0,0 +1,122 @@ +{ + "schema_version": "0.5", + "score_kind": "diagnostic-project-fit", + "score_is_official": false, + "draft_id": "astrai-common-shape-bf16-gemv-v1", + "draft_state": "D3-contract-proposed", + "formula_template_id": "project-fit-kernel-v0.5", + "diagnostic_project_fit_100": 92.6415162827573, + "component_values": { + "evolutionary_maintainability": 0.8835232637338853, + "project_contract_fit": 1.0, + "performance_reuse_utilization": 0.9515978740065268, + "operational_fit": 0.8438741102892859 + }, + "component_floors": { + "evolutionary_maintainability": 0.6, + "project_contract_fit": 0.6, + "performance_reuse_utilization": 0.4, + "operational_fit": 0.6 + }, + "subcomponent_inputs": { + "evolutionary_maintainability": { + "evolution": 0.78, + "locality": 0.9, + "tests": 1.0, + "failure": 1.0, + "contract": 0.9 + }, + "project_contract_fit": { + "integration": 1.0, + "interface": 1.0, + "lifecycle": 1.0, + "buildtest": 1.0, + "policy": 1.0 + }, + "performance_reuse_utilization": { + "attainment": 1.0, + "coverage": 0.82, + "retention": 1.0, + "family": 1.0, + "compile": 1.0 + }, + "operational_fit": { + "replay": 0.9, + "load": 0.82, + "resource": 0.9, + "coldsteady": 0.75 + } + }, + "input_rationale": { + "evolution": "The candidate extends the merged AstrAI GEMV stack with a rejected broad half-CTA experiment and an evidence-gated exact selector, but has no upstream maintenance history yet.", + "locality": "The CUDA change remains inside the GEMV launcher and one cooperative kernel template; dispatcher, benchmark, tests, and documentation are separately scoped.", + "tests": "The complete CPU-capable AstrAI suite passed (625 passed, 185 skipped) and the built L20 extension passed 140 focused GPU tests.", + "failure": "Unaligned K, unsupported M, training/autograd, unmeasured shapes, other architectures, losing half-CTA cells, and sub-threshold automatic bands fail closed to the existing kernel or PyTorch.", + "contract": "The benchmark and tests are digest-bound, but the D3 acceptance contract has not received maintainer review or a seal.", + "integration": "The fast path remains behind AstrAI's existing bf16_gemv primitive and linear backend.", + "interface": "No public signature or environment-variable contract changes; ASTRAI_GEMV modes remain compatible.", + "lifecycle": "The path remains inference-only, uses the current CUDA stream, fuses optional bias, and is CUDA Graph safe.", + "buildtest": "The SM89 extension built all 128/256-thread row specializations without spills, Ruff format/import checks passed, and both local and L20 tests passed.", + "policy": "The candidate adds no dependency and enables automatic dispatch only for measured architecture/M/(N,K) cells that clear per-shape and chain gates.", + "attainment": "Every newly enabled Qwen2, LLaMA 3 70B, and OPT chain exceeds the 3% gate; losing M/family combinations are explicitly excluded.", + "coverage": "Twenty-two traditional shapes and seven model families were measured in the required NVIDIA L20/SM89 cell; optional A100/SM80 and H100/SM90 cells remain untested and fall back in auto mode.", + "retention": "The exact row/shape half-CTA selector preserves the 256-thread fallback and existing warp-tiled bands; direct A/B evidence and the regression suite passed.", + "family": "Twenty-two LLaMA, Qwen2, OPT, and GPT-NeoX matrix shapes plus seven dependent synthetic chains cover M=1/2/4/8, without claiming a whole-model benchmark.", + "compile": "AOT CUDA compilation completed before timed cases; no compilation is permitted in steady-state timing.", + "replay": "Paired order alternation, warmup, 21 samples, median and p90 are captured, but not seven fresh processes.", + "load": "Distinct-weight dependent chains cover seven model families and include Python dispatch; a real end-to-end service workload was not run.", + "resource": "GPU5 memory and utilization were checked before and after; the co-resident idle AstrAI service remained allocated and is disclosed.", + "coldsteady": "Precompile and steady-state phases are separated, but cold-start latency is not part of the checked-in result." + }, + "benchmark_trust": { + "formula_version": "benchmark-trust-v0.5", + "status": "scored", + "score_100": 95.87315155141826, + "components": { + "reproducibility": 1.0, + "evidence": 1.0, + "statistics": 0.9, + "environment": 0.9 + }, + "failure_codes": [ + "DRAFT_UNSEALED", + "FRESH_PROCESS_REPLAY_INCOMPLETE", + "OPTIONAL_CELLS_UNTESTED" + ] + }, + "official_project_fit": { + "status": "unresolved", + "score_100": null, + "failure_codes": [ + "DRAFT_SEAL_MISSING", + "FRESH_PROCESS_REPLAYS_BELOW_MINIMUM", + "SYSTEM_TRACE_EVIDENCE_MISSING", + "HIDDEN_PROBES_INCOMPLETE", + "EVIDENCE_MANIFEST_UNVERIFIED" + ] + }, + "comparison_cell": { + "target_project_profile_sha256": "sha256:953067c47f5298f06846c8fe873db55afae49b6588a81926e49d9ef6b0994e08", + "target_repository_or_baseline_sha256": "sha256:59c45de7cb0356e9661cfb73d38bd11c888a8b8d564e9123a90c9fa3910cc498", + "change_intent": "add-fastpath", + "semantic_contract_sha256": "sha256:74808a77e551ffa5f77e3256930f2647ebd080828cfdbdcc2770f9d1fa298ff6", + "acceptance_contract_sha256": "sha256:15ffa9e1378d7cb1b06f018f9e1aac16bf56f14446271d51129f86c13f3827bb", + "probe_set_sha256": "sha256:beafc95a2a4aaa587079f1938db8c0fb115d63c2b0c3b185debc9a6132397653", + "workload_portfolio_sha256": "sha256:c3cf734200844d7b8308a4803f456cd9a3f0d7e19839e7d3e74778f6d1307e4d", + "performance_target_sha256": "sha256:22da954f941ba912f60d79b7a58121436d4ba97e64f0182f66b8274080c3a75a", + "required_deployment_cell_set_sha256": "sha256:6c00c84e931a5ad6fdaa5b4f7d4497c872530c75f8510431aff883006a4a2779", + "formula_template_id": "project-fit-kernel-v0.5", + "evidence_policy_id": "v0.4-evidence-ladder-plus-seal-v0.5", + "project_season": "astrai-2026q3", + "cross_project_ranking_allowed": false + }, + "execution": { + "infraswe_commit": "811bc775ed5b3a6ec853219245f3469f78818020", + "draft_resolution_sha256": "952ee658e669126582b8eb9aeedede9f70dae6348ea0f1119d64a36d5da100fe", + "infraswe_draft_engine_tests": "53 passed", + "astrai_tests": "625 passed, 185 skipped", + "astrai_l20_gpu_tests": "140 passed", + "astrai_lint": "ruff format and import-order checks passed", + "sm89_build": "bf16_gemv target built successfully; 128/256-thread cooperative specializations used 26-52 registers, one barrier, and zero spills; M=4 warp tiling used 42 registers, zero barriers, and zero spills" + } +} diff --git a/csrc/kernels/gemv/bf16_gemv.cu b/csrc/kernels/gemv/bf16_gemv.cu index 70cce50b..58278a82 100644 --- a/csrc/kernels/gemv/bf16_gemv.cu +++ b/csrc/kernels/gemv/bf16_gemv.cu @@ -12,7 +12,9 @@ namespace { constexpr int kThreads = 256; +constexpr int kHalfCtaThreads = 128; constexpr int kWarpSize = 32; +constexpr int kWarpTiledThreads = 128; __device__ __forceinline__ float warp_sum(float value) { #pragma unroll @@ -22,7 +24,7 @@ __device__ __forceinline__ float warp_sum(float value) { return value; } -template +template __global__ void bf16_gemv_kernel( const __nv_bfloat16* __restrict__ x, const __nv_bfloat16* __restrict__ weight, @@ -36,7 +38,7 @@ __global__ void bf16_gemv_kernel( const int warp = threadIdx.x / kWarpSize; float sums[Rows] = {}; - __shared__ float warp_sums[Rows][kThreads / kWarpSize]; + __shared__ float warp_sums[Rows][Threads / kWarpSize]; // Weight row: scalar head/tail around a 16-byte-aligned uint4 middle so // any K is accepted while keeping 128-bit weight loads, which dominate // bandwidth on decode shapes. x pairs with scalar loads: it is a tiny @@ -129,7 +131,6 @@ __global__ void bf16_gemv_kernel( } } - #pragma unroll for (int row = 0; row < Rows; ++row) { sums[row] = warp_sum(sums[row]); @@ -146,7 +147,7 @@ __global__ void bf16_gemv_kernel( #pragma unroll for (int row = 0; row < Rows; ++row) { float sum = - lane < (kThreads / kWarpSize) ? warp_sums[row][lane] : 0.0f; + lane < (Threads / kWarpSize) ? warp_sums[row][lane] : 0.0f; sum = warp_sum(sum); if (lane == 0) { if (bias != nullptr) { @@ -158,6 +159,120 @@ __global__ void bf16_gemv_kernel( } } +template +__global__ void bf16_gemv_aligned_warp_tiled_kernel( + const __nv_bfloat16* __restrict__ x, + const __nv_bfloat16* __restrict__ weight, + const __nv_bfloat16* __restrict__ bias, + __nv_bfloat16* __restrict__ output, + int n, + int k +) { + constexpr int kWarpsPerBlock = kWarpTiledThreads / kWarpSize; + const int lane = threadIdx.x & (kWarpSize - 1); + const int warp = threadIdx.x / kWarpSize; + const int output_index = blockIdx.x * kWarpsPerBlock + warp; + if (output_index >= n) { + return; + } + + // The launcher selects this path only when each row is 16-byte aligned. + // Four independent output rows per CTA remove the block-wide reduction + // barrier and improve occupancy for the medium LLaMA projection bands. + const int vectors = k / 8; + const auto* x4 = reinterpret_cast(x); + const auto* w4 = reinterpret_cast(weight) + + static_cast(output_index) * vectors; + float sums[Rows] = {}; + for (int vector = lane; vector < vectors; vector += kWarpSize) { + const uint4 wv_raw = w4[vector]; + const auto* wv = reinterpret_cast(&wv_raw); +#pragma unroll + for (int row = 0; row < Rows; ++row) { + const uint4 xv_raw = + x4[static_cast(row) * vectors + vector]; + const auto* xv = reinterpret_cast(&xv_raw); +#pragma unroll + for (int pair = 0; pair < 4; ++pair) { + sums[row] = fmaf( + __bfloat162float(__low2bfloat16(xv[pair])), + __bfloat162float(__low2bfloat16(wv[pair])), + sums[row] + ); + sums[row] = fmaf( + __bfloat162float(__high2bfloat16(xv[pair])), + __bfloat162float(__high2bfloat16(wv[pair])), + sums[row] + ); + } + } + } + +#pragma unroll + for (int row = 0; row < Rows; ++row) { + sums[row] = warp_sum(sums[row]); + if (lane == 0) { + if (bias != nullptr) { + sums[row] += __bfloat162float(bias[output_index]); + } + output[row * n + output_index] = __float2bfloat16_rn(sums[row]); + } + } +} + +template +constexpr bool use_warp_tiled_kernel(int n, int k) { + // These bands are intentionally narrow and are validated by the common + // transformer benchmark. The 256-thread cooperative kernel remains the + // fallback for arbitrary K, larger projections, and M=2 (where the + // single-warp reduction regresses the current vectorized kernel). + if constexpr (Rows == 4) { + return (n == 1024 && k == 4096) || + (n == 4096 && k == 4096) || + (n == 11008 && k == 4096) || + (n == 4096 && k == 11008); + } + return false; +} + +template +constexpr bool use_half_cta_kernel(int n, int k) { + // A 128-thread CTA reduces synchronization and scheduling overhead for + // selected medium decode projections. Keep the selector exact: long-K + // and bandwidth-saturated shapes regress, and the winning bands differ + // materially with the number of reused input rows. + if constexpr (Rows == 1) { + return n == 8192 && k == 2048; + } + if constexpr (Rows == 2) { + return (n == 4096 && k == 4096) || + (n == 11008 && k == 4096) || + (n == 3584 && k == 3584) || + (n == 2048 && k == 2048) || + (n == 8192 && k == 2048); + } + if constexpr (Rows == 4) { + return (n == 5120 && k == 5120) || + (n == 3584 && k == 3584) || + (n == 2048 && k == 2048) || + (n == 8192 && k == 2048); + } + if constexpr (Rows == 8) { + return (n == 4096 && k == 4096) || + (n == 11008 && k == 4096) || + (n == 4096 && k == 11008) || + (n == 1024 && k == 4096) || + (n == 5120 && k == 5120) || + (n == 512 && k == 3584) || + (n == 3584 && k == 3584) || + (n == 1024 && k == 8192) || + (n == 2048 && k == 2048) || + (n == 8192 && k == 2048) || + (n == 2048 && k == 8192); + } + return false; +} + template void launch_bf16_gemv( const __nv_bfloat16* x, @@ -168,7 +283,28 @@ void launch_bf16_gemv( int k, cudaStream_t stream ) { - bf16_gemv_kernel<<>>( + const bool aligned_rows = k % 8 == 0 && + (reinterpret_cast(x) & 15u) == 0u && + (reinterpret_cast(weight) & 15u) == 0u; + if constexpr (Rows == 4) { + if (aligned_rows && use_warp_tiled_kernel(n, k)) { + constexpr int kWarpsPerBlock = kWarpTiledThreads / kWarpSize; + const int blocks = (n + kWarpsPerBlock - 1) / kWarpsPerBlock; + bf16_gemv_aligned_warp_tiled_kernel + <<>>( + x, weight, bias, output, n, k + ); + return; + } + } + if (aligned_rows && use_half_cta_kernel(n, k)) { + bf16_gemv_kernel + <<>>( + x, weight, bias, output, n, k + ); + return; + } + bf16_gemv_kernel<<>>( x, weight, bias, output, n, k ); } diff --git a/docs/benchmarks/gemv_common_chain_l20_sm89.json b/docs/benchmarks/gemv_common_chain_l20_sm89.json new file mode 100644 index 00000000..27b27cb4 --- /dev/null +++ b/docs/benchmarks/gemv_common_chain_l20_sm89.json @@ -0,0 +1,138 @@ +{ + "environment": { + "device": "NVIDIA L20", + "capability": "8.9", + "torch": "2.11.0+cu128", + "cuda": "12.8" + }, + "parameters": { + "seed": 20260902, + "weight_std": 0.02, + "warmup": 40, + "samples": 21, + "inner": 100, + "chain_inner": 50 + }, + "results": [ + { + "suite": "synthetic_chain", + "label": "llama2_7b", + "m": 2, + "n": 4096, + "k": 11008, + "torch_median_ms": 0.5955379104614258, + "torch_p90_ms": 0.5965619277954102, + "candidate_median_ms": 0.5489459228515625, + "candidate_p90_ms": 0.5493555068969727, + "speedup_pct": 8.487536872090406, + "max_abs": 0.0625, + "relative_l2": 0.0010812067720616032, + "argmax_equal": true + }, + { + "suite": "synthetic_chain", + "label": "llama2_7b", + "m": 4, + "n": 4096, + "k": 11008, + "torch_median_ms": 0.5960294342041016, + "torch_p90_ms": 0.5965619277954102, + "candidate_median_ms": 0.5963980865478515, + "candidate_p90_ms": 0.5966643142700195, + "speedup_pct": -0.06181313321842463, + "max_abs": 0.0, + "relative_l2": 0.0, + "argmax_equal": true + }, + { + "suite": "synthetic_chain", + "label": "llama3_8b", + "m": 2, + "n": 4096, + "k": 14336, + "torch_median_ms": 0.6417817687988281, + "torch_p90_ms": 0.6427635192871094, + "candidate_median_ms": 0.5915033721923828, + "candidate_p90_ms": 0.5918515014648438, + "speedup_pct": 8.500103121997515, + "max_abs": 0.03125, + "relative_l2": 0.0007109941736025917, + "argmax_equal": true + }, + { + "suite": "synthetic_chain", + "label": "llama3_8b", + "m": 4, + "n": 4096, + "k": 14336, + "torch_median_ms": 0.6444236755371093, + "torch_p90_ms": 0.6447309112548828, + "candidate_median_ms": 0.6054105758666992, + "candidate_p90_ms": 0.6057369613647461, + "speedup_pct": 6.44407303499106, + "max_abs": 0.0625, + "relative_l2": 0.0009759692858008782, + "argmax_equal": true + }, + { + "suite": "synthetic_chain", + "label": "llama2_13b", + "m": 2, + "n": 5120, + "k": 13824, + "torch_median_ms": 0.9096806335449219, + "torch_p90_ms": 0.910376968383789, + "candidate_median_ms": 0.8609587097167969, + "candidate_p90_ms": 0.8614297485351563, + "speedup_pct": 5.659031412104709, + "max_abs": 0.0625, + "relative_l2": 0.0010442947319278225, + "argmax_equal": true + }, + { + "suite": "synthetic_chain", + "label": "llama2_13b", + "m": 4, + "n": 5120, + "k": 13824, + "torch_median_ms": 0.9117491149902344, + "torch_p90_ms": 0.9125888061523437, + "candidate_median_ms": 0.8628633880615234, + "candidate_p90_ms": 0.8632115173339844, + "speedup_pct": 5.665523373118875, + "max_abs": 0.125, + "relative_l2": 0.0010491116899757787, + "argmax_equal": true + }, + { + "suite": "synthetic_chain", + "label": "gpt_neox_20b", + "m": 2, + "n": 4096, + "k": 16384, + "torch_median_ms": 0.5899878311157226, + "torch_p90_ms": 0.5925696182250977, + "candidate_median_ms": 0.5516704177856445, + "candidate_p90_ms": 0.5521408081054687, + "speedup_pct": 6.945707454077521, + "max_abs": 0.03125, + "relative_l2": 0.0007355517078123456, + "argmax_equal": true + }, + { + "suite": "synthetic_chain", + "label": "gpt_neox_20b", + "m": 4, + "n": 4096, + "k": 16384, + "torch_median_ms": 0.5916057586669922, + "torch_p90_ms": 0.5919334411621093, + "candidate_median_ms": 0.5585100936889649, + "candidate_p90_ms": 0.558837776184082, + "speedup_pct": 5.9257057933241875, + "max_abs": 0.0625, + "relative_l2": 0.0007807545463126493, + "argmax_equal": true + } + ] +} diff --git a/docs/benchmarks/gemv_common_extended_chain_l20_sm89.json b/docs/benchmarks/gemv_common_extended_chain_l20_sm89.json new file mode 100644 index 00000000..10554130 --- /dev/null +++ b/docs/benchmarks/gemv_common_extended_chain_l20_sm89.json @@ -0,0 +1,213 @@ +{ + "environment": { + "device": "NVIDIA L20", + "capability": "8.9", + "torch": "2.11.0+cu128", + "cuda": "12.8" + }, + "parameters": { + "suite": "chain", + "family": "traditional", + "m": [ + 1, + 2, + 4, + 8 + ], + "shape_labels": null, + "chain_labels": [ + "qwen2_7b", + "llama3_70b", + "opt_1_3b" + ], + "candidate_mode": "auto", + "seed": 20260902, + "weight_std": 0.02, + "warmup": 30, + "samples": 15, + "inner": 100, + "chain_inner": 20 + }, + "results": [ + { + "suite": "synthetic_chain", + "label": "qwen2_7b", + "m": 1, + "n": 3584, + "k": 18944, + "torch_median_ms": 0.6211584091186524, + "torch_p90_ms": 0.6384640216827393, + "candidate_median_ms": 0.6222847938537598, + "candidate_p90_ms": 0.6400512218475342, + "speedup_pct": -0.18100791570556662, + "max_abs": 0.0, + "relative_l2": 0.0, + "argmax_equal": true + }, + { + "suite": "synthetic_chain", + "label": "qwen2_7b", + "m": 2, + "n": 3584, + "k": 18944, + "torch_median_ms": 0.6860288143157959, + "torch_p90_ms": 0.6872064113616944, + "candidate_median_ms": 0.6382592201232911, + "candidate_p90_ms": 0.6390768051147461, + "speedup_pct": 7.484356306404361, + "max_abs": 0.0625, + "relative_l2": 0.0006854576325413914, + "argmax_equal": true + }, + { + "suite": "synthetic_chain", + "label": "qwen2_7b", + "m": 4, + "n": 3584, + "k": 18944, + "torch_median_ms": 0.6879744052886962, + "torch_p90_ms": 0.6889472007751465, + "candidate_median_ms": 0.6401023864746094, + "candidate_p90_ms": 0.6407167911529541, + "speedup_pct": 7.478806488715661, + "max_abs": 0.046875, + "relative_l2": 0.0008338340873294167, + "argmax_equal": true + }, + { + "suite": "synthetic_chain", + "label": "qwen2_7b", + "m": 8, + "n": 3584, + "k": 18944, + "torch_median_ms": 0.6909440040588379, + "torch_p90_ms": 0.6916096210479736, + "candidate_median_ms": 0.6909952163696289, + "candidate_p90_ms": 0.6916096210479736, + "speedup_pct": -0.007411384272681953, + "max_abs": 0.0, + "relative_l2": 0.0, + "argmax_equal": true + }, + { + "suite": "synthetic_chain", + "label": "llama3_70b", + "m": 1, + "n": 8192, + "k": 28672, + "torch_median_ms": 2.273740768432617, + "torch_p90_ms": 2.2749696731567384, + "candidate_median_ms": 2.2747135162353516, + "candidate_p90_ms": 2.2809600830078125, + "speedup_pct": -0.042763530255196525, + "max_abs": 0.0, + "relative_l2": 0.0, + "argmax_equal": true + }, + { + "suite": "synthetic_chain", + "label": "llama3_70b", + "m": 2, + "n": 8192, + "k": 28672, + "torch_median_ms": 2.453555107116699, + "torch_p90_ms": 2.4563199996948244, + "candidate_median_ms": 2.2766592025756838, + "candidate_p90_ms": 2.2790143966674803, + "speedup_pct": 7.769977357212077, + "max_abs": 0.25, + "relative_l2": 0.0027889899820575568, + "argmax_equal": true + }, + { + "suite": "synthetic_chain", + "label": "llama3_70b", + "m": 4, + "n": 8192, + "k": 28672, + "torch_median_ms": 2.463795280456543, + "torch_p90_ms": 2.465740776062012, + "candidate_median_ms": 2.2878719329833985, + "candidate_p90_ms": 2.2896127700805664, + "speedup_pct": 7.689387895228017, + "max_abs": 0.25, + "relative_l2": 0.0027640196353711528, + "argmax_equal": true + }, + { + "suite": "synthetic_chain", + "label": "llama3_70b", + "m": 8, + "n": 8192, + "k": 28672, + "torch_median_ms": 2.468351936340332, + "torch_p90_ms": 2.4706560134887696, + "candidate_median_ms": 2.4688608169555666, + "candidate_p90_ms": 2.4712703704833983, + "speedup_pct": -0.020611960453165157, + "max_abs": 0.0, + "relative_l2": 0.0, + "argmax_equal": true + }, + { + "suite": "synthetic_chain", + "label": "opt_1_3b", + "m": 1, + "n": 2048, + "k": 8192, + "torch_median_ms": 0.107315194606781, + "torch_p90_ms": 0.10864640474319458, + "candidate_median_ms": 0.10265599489212036, + "candidate_p90_ms": 0.1047551989555359, + "speedup_pct": 4.538653314457597, + "max_abs": 0.0, + "relative_l2": 0.0, + "argmax_equal": true + }, + { + "suite": "synthetic_chain", + "label": "opt_1_3b", + "m": 2, + "n": 2048, + "k": 8192, + "torch_median_ms": 0.1276927947998047, + "torch_p90_ms": 0.12861440181732178, + "candidate_median_ms": 0.10199040174484253, + "candidate_p90_ms": 0.10362880229949951, + "speedup_pct": 25.20079597221696, + "max_abs": 0.03125, + "relative_l2": 0.0007647863858926024, + "argmax_equal": true + }, + { + "suite": "synthetic_chain", + "label": "opt_1_3b", + "m": 4, + "n": 2048, + "k": 8192, + "torch_median_ms": 0.12702720165252684, + "torch_p90_ms": 0.12830719947814942, + "candidate_median_ms": 0.14146560430526733, + "candidate_p90_ms": 0.14325759410858155, + "speedup_pct": -10.20629906728705, + "max_abs": 0.0, + "relative_l2": 0.0, + "argmax_equal": true + }, + { + "suite": "synthetic_chain", + "label": "opt_1_3b", + "m": 8, + "n": 2048, + "k": 8192, + "torch_median_ms": 0.12810239791870118, + "torch_p90_ms": 0.12943359613418579, + "candidate_median_ms": 0.12851200103759766, + "candidate_p90_ms": 0.13015040159225463, + "speedup_pct": -0.318727523958362, + "max_abs": 0.0, + "relative_l2": 0.0, + "argmax_equal": true + } + ] +} diff --git a/docs/benchmarks/gemv_common_extended_kernel_l20_sm89.json b/docs/benchmarks/gemv_common_extended_kernel_l20_sm89.json new file mode 100644 index 00000000..9953ff2d --- /dev/null +++ b/docs/benchmarks/gemv_common_extended_kernel_l20_sm89.json @@ -0,0 +1,701 @@ +{ + "environment": { + "device": "NVIDIA L20", + "capability": "8.9", + "torch": "2.11.0+cu128", + "cuda": "12.8" + }, + "parameters": { + "suite": "kernel", + "family": "traditional", + "m": [ + 1, + 2, + 4, + 8 + ], + "shape_labels": [ + "qwen2_7b_kv", + "qwen2_7b_qo", + "qwen2_7b_up_gate", + "qwen2_7b_down", + "llama3_70b_kv", + "llama3_70b_qo", + "llama3_70b_up_gate", + "llama3_70b_down", + "opt_1_3b_qkvo", + "opt_1_3b_up", + "opt_1_3b_down" + ], + "chain_labels": null, + "candidate_mode": "auto", + "seed": 20260902, + "weight_std": 0.02, + "warmup": 40, + "samples": 21, + "inner": 200, + "chain_inner": 20 + }, + "results": [ + { + "suite": "kernel", + "label": "qwen2_7b_kv", + "m": 1, + "n": 512, + "k": 3584, + "torch_median_ms": 0.01112064003944397, + "torch_p90_ms": 0.011166720390319825, + "candidate_median_ms": 0.006773759722709656, + "candidate_p90_ms": 0.006819999814033508, + "speedup_pct": 64.1723429037643, + "max_abs": 0.0, + "relative_l2": 0.0, + "argmax_equal": true + }, + { + "suite": "kernel", + "label": "qwen2_7b_kv", + "m": 2, + "n": 512, + "k": 3584, + "torch_median_ms": 0.01459712028503418, + "torch_p90_ms": 0.014704639911651612, + "candidate_median_ms": 0.006702079772949219, + "candidate_p90_ms": 0.00675328016281128, + "speedup_pct": 117.79985884308246, + "max_abs": 0.015625, + "relative_l2": 0.0026976845668545417, + "argmax_equal": true + }, + { + "suite": "kernel", + "label": "qwen2_7b_kv", + "m": 4, + "n": 512, + "k": 3584, + "torch_median_ms": 0.014638079404830933, + "torch_p90_ms": 0.014750720262527465, + "candidate_median_ms": 0.006640639901161194, + "candidate_p90_ms": 0.0067123198509216305, + "speedup_pct": 120.43175993131766, + "max_abs": 0.015625, + "relative_l2": 0.0026301932766628084, + "argmax_equal": true + }, + { + "suite": "kernel", + "label": "qwen2_7b_kv", + "m": 8, + "n": 512, + "k": 3584, + "torch_median_ms": 0.014668799638748169, + "torch_p90_ms": 0.01480191946029663, + "candidate_median_ms": 0.0066969597339630125, + "candidate_p90_ms": 0.006768640279769898, + "speedup_pct": 119.03670055468166, + "max_abs": 0.015625, + "relative_l2": 0.0026439480090191202, + "argmax_equal": true + }, + { + "suite": "kernel", + "label": "qwen2_7b_qo", + "m": 1, + "n": 3584, + "k": 3584, + "torch_median_ms": 0.014392319917678833, + "torch_p90_ms": 0.014417920112609863, + "candidate_median_ms": 0.00956928014755249, + "candidate_p90_ms": 0.009579520225524902, + "speedup_pct": 50.401280929787795, + "max_abs": 0.0, + "relative_l2": 0.0, + "argmax_equal": true + }, + { + "suite": "kernel", + "label": "qwen2_7b_qo", + "m": 2, + "n": 3584, + "k": 3584, + "torch_median_ms": 0.01904639959335327, + "torch_p90_ms": 0.019082239866256713, + "candidate_median_ms": 0.009958400130271912, + "candidate_p90_ms": 0.009978880286216735, + "speedup_pct": 91.25963351738923, + "max_abs": 0.015625, + "relative_l2": 0.0001875343957076546, + "argmax_equal": true + }, + { + "suite": "kernel", + "label": "qwen2_7b_qo", + "m": 4, + "n": 3584, + "k": 3584, + "torch_median_ms": 0.019118080139160155, + "torch_p90_ms": 0.019148800373077392, + "candidate_median_ms": 0.01168895959854126, + "candidate_p90_ms": 0.01171455979347229, + "speedup_pct": 63.556730417188035, + "max_abs": 0.015625, + "relative_l2": 0.00014083795183464072, + "argmax_equal": true + }, + { + "suite": "kernel", + "label": "qwen2_7b_qo", + "m": 8, + "n": 3584, + "k": 3584, + "torch_median_ms": 0.019215359687805175, + "torch_p90_ms": 0.019266560077667236, + "candidate_median_ms": 0.019537919759750368, + "candidate_p90_ms": 0.019578880071640013, + "speedup_pct": -1.6509437847610164, + "max_abs": 0.0078125, + "relative_l2": 0.00010323309107383507, + "argmax_equal": true + }, + { + "suite": "kernel", + "label": "qwen2_7b_up_gate", + "m": 1, + "n": 18944, + "k": 3584, + "torch_median_ms": 0.1831167984008789, + "torch_p90_ms": 0.18320383071899415, + "candidate_median_ms": 0.17959936141967772, + "candidate_p90_ms": 0.17967615127563477, + "speedup_pct": 1.9584908060902517, + "max_abs": 0.00390625, + "relative_l2": 2.3701111300909968e-05, + "argmax_equal": true + }, + { + "suite": "kernel", + "label": "qwen2_7b_up_gate", + "m": 2, + "n": 18944, + "k": 3584, + "torch_median_ms": 0.19084287643432618, + "torch_p90_ms": 0.19095552444458008, + "candidate_median_ms": 0.17985023498535158, + "candidate_p90_ms": 0.1798860740661621, + "speedup_pct": 6.112108471734801, + "max_abs": 0.015625, + "relative_l2": 0.00012973139947699189, + "argmax_equal": true + }, + { + "suite": "kernel", + "label": "qwen2_7b_up_gate", + "m": 4, + "n": 18944, + "k": 3584, + "torch_median_ms": 0.19150848388671876, + "torch_p90_ms": 0.19164159774780273, + "candidate_median_ms": 0.18037248611450196, + "candidate_p90_ms": 0.18044416427612306, + "speedup_pct": 6.173889384186659, + "max_abs": 0.015625, + "relative_l2": 0.00012086182936299403, + "argmax_equal": true + }, + { + "suite": "kernel", + "label": "qwen2_7b_up_gate", + "m": 8, + "n": 18944, + "k": 3584, + "torch_median_ms": 0.19242992401123046, + "torch_p90_ms": 0.19248128890991212, + "candidate_median_ms": 0.18143232345581053, + "candidate_p90_ms": 0.18149375915527344, + "speedup_pct": 6.061544241921424, + "max_abs": 0.015625, + "relative_l2": 0.00015294284982227395, + "argmax_equal": true + }, + { + "suite": "kernel", + "label": "qwen2_7b_down", + "m": 1, + "n": 3584, + "k": 18944, + "torch_median_ms": 0.1842892837524414, + "torch_p90_ms": 0.18454511642456053, + "candidate_median_ms": 0.18156032562255858, + "candidate_p90_ms": 0.18160640716552734, + "speedup_pct": 1.503058622816078, + "max_abs": 0.0, + "relative_l2": 0.0, + "argmax_equal": true + }, + { + "suite": "kernel", + "label": "qwen2_7b_down", + "m": 2, + "n": 3584, + "k": 18944, + "torch_median_ms": 0.19400192260742188, + "torch_p90_ms": 0.19409919738769532, + "candidate_median_ms": 0.18149375915527344, + "candidate_p90_ms": 0.18160640716552734, + "speedup_pct": 6.891787084230994, + "max_abs": 0.03125, + "relative_l2": 0.00044727133148071604, + "argmax_equal": true + }, + { + "suite": "kernel", + "label": "qwen2_7b_down", + "m": 4, + "n": 3584, + "k": 18944, + "torch_median_ms": 0.1941043281555176, + "torch_p90_ms": 0.1942425537109375, + "candidate_median_ms": 0.18183679580688478, + "candidate_p90_ms": 0.18190336227416992, + "speedup_pct": 6.746452110639489, + "max_abs": 0.03125, + "relative_l2": 0.00040355112653001876, + "argmax_equal": true + }, + { + "suite": "kernel", + "label": "qwen2_7b_down", + "m": 8, + "n": 3584, + "k": 18944, + "torch_median_ms": 0.19467775344848634, + "torch_p90_ms": 0.19479551315307617, + "candidate_median_ms": 0.18391551971435546, + "candidate_p90_ms": 0.18406400680541993, + "speedup_pct": 5.851726787845868, + "max_abs": 0.0625, + "relative_l2": 0.00042564765557899596, + "argmax_equal": true + }, + { + "suite": "kernel", + "label": "llama3_70b_kv", + "m": 1, + "n": 1024, + "k": 8192, + "torch_median_ms": 0.014883840084075927, + "torch_p90_ms": 0.014996639490127563, + "candidate_median_ms": 0.007413759827613831, + "candidate_p90_ms": 0.007434239983558655, + "speedup_pct": 100.7596743104422, + "max_abs": 0.03125, + "relative_l2": 0.0027323557431972703, + "argmax_equal": true + }, + { + "suite": "kernel", + "label": "llama3_70b_kv", + "m": 2, + "n": 1024, + "k": 8192, + "torch_median_ms": 0.01474560022354126, + "torch_p90_ms": 0.014914560317993163, + "candidate_median_ms": 0.008038399815559387, + "candidate_p90_ms": 0.008079360127449036, + "speedup_pct": 83.43949743578563, + "max_abs": 0.03125, + "relative_l2": 0.002428879255111259, + "argmax_equal": true + }, + { + "suite": "kernel", + "label": "llama3_70b_kv", + "m": 4, + "n": 1024, + "k": 8192, + "torch_median_ms": 0.014730240106582642, + "torch_p90_ms": 0.01792512059211731, + "candidate_median_ms": 0.009937919974327088, + "candidate_p90_ms": 0.009968640208244324, + "speedup_pct": 48.222567143181784, + "max_abs": 0.03125, + "relative_l2": 0.0026361971223301064, + "argmax_equal": true + }, + { + "suite": "kernel", + "label": "llama3_70b_kv", + "m": 8, + "n": 1024, + "k": 8192, + "torch_median_ms": 0.014663679599761963, + "torch_p90_ms": 0.014975680112838745, + "candidate_median_ms": 0.016624640226364135, + "candidate_p90_ms": 0.016655360460281373, + "speedup_pct": -11.795507150238294, + "max_abs": 0.03125, + "relative_l2": 0.002637206430000059, + "argmax_equal": true + }, + { + "suite": "kernel", + "label": "llama3_70b_qo", + "m": 1, + "n": 8192, + "k": 8192, + "torch_median_ms": 0.18086399078369142, + "torch_p90_ms": 0.18099199295043944, + "candidate_median_ms": 0.17788415908813476, + "candidate_p90_ms": 0.17794559478759767, + "speedup_pct": 1.6751529258320685, + "max_abs": 0.00390625, + "relative_l2": 2.7130160561412286e-05, + "argmax_equal": true + }, + { + "suite": "kernel", + "label": "llama3_70b_qo", + "m": 2, + "n": 8192, + "k": 8192, + "torch_median_ms": 0.19515392303466797, + "torch_p90_ms": 0.1953228759765625, + "candidate_median_ms": 0.17703935623168945, + "candidate_p90_ms": 0.17726463317871094, + "speedup_pct": 10.231943443847701, + "max_abs": 0.015625, + "relative_l2": 0.0001814438414194483, + "argmax_equal": true + }, + { + "suite": "kernel", + "label": "llama3_70b_qo", + "m": 4, + "n": 8192, + "k": 8192, + "torch_median_ms": 0.19595264434814452, + "torch_p90_ms": 0.19606016159057618, + "candidate_median_ms": 0.1782067108154297, + "candidate_p90_ms": 0.17833471298217773, + "speedup_pct": 9.958061316273593, + "max_abs": 0.03125, + "relative_l2": 0.00027871101042350544, + "argmax_equal": true + }, + { + "suite": "kernel", + "label": "llama3_70b_qo", + "m": 8, + "n": 8192, + "k": 8192, + "torch_median_ms": 0.19649023056030274, + "torch_p90_ms": 0.19668479919433593, + "candidate_median_ms": 0.17957887649536133, + "candidate_p90_ms": 0.17977855682373048, + "speedup_pct": 9.417229016564344, + "max_abs": 0.03125, + "relative_l2": 0.00020848420253535642, + "argmax_equal": true + }, + { + "suite": "kernel", + "label": "llama3_70b_up_gate", + "m": 1, + "n": 28672, + "k": 8192, + "torch_median_ms": 0.6258892822265625, + "torch_p90_ms": 0.6261350250244141, + "candidate_median_ms": 0.618419189453125, + "candidate_p90_ms": 0.6187059020996094, + "speedup_pct": 1.2079335345404285, + "max_abs": 0.0078125, + "relative_l2": 2.642172027691496e-05, + "argmax_equal": true + }, + { + "suite": "kernel", + "label": "llama3_70b_up_gate", + "m": 2, + "n": 28672, + "k": 8192, + "torch_median_ms": 0.6681497955322265, + "torch_p90_ms": 0.6683853149414063, + "candidate_median_ms": 0.620323829650879, + "candidate_p90_ms": 0.6205542373657227, + "speedup_pct": 7.709838570648531, + "max_abs": 0.03125, + "relative_l2": 0.0002490715541388539, + "argmax_equal": true + }, + { + "suite": "kernel", + "label": "llama3_70b_up_gate", + "m": 4, + "n": 28672, + "k": 8192, + "torch_median_ms": 0.6681446075439453, + "torch_p90_ms": 0.6685491180419922, + "candidate_median_ms": 0.6240716934204101, + "candidate_p90_ms": 0.6242047882080078, + "speedup_pct": 7.062155612599663, + "max_abs": 0.03125, + "relative_l2": 0.00024391069344504231, + "argmax_equal": true + }, + { + "suite": "kernel", + "label": "llama3_70b_up_gate", + "m": 8, + "n": 28672, + "k": 8192, + "torch_median_ms": 0.6714214324951172, + "torch_p90_ms": 0.6716159820556641, + "candidate_median_ms": 0.6219520187377929, + "candidate_p90_ms": 0.6220800018310547, + "speedup_pct": 7.953895520384169, + "max_abs": 0.03125, + "relative_l2": 0.00024773541899195583, + "argmax_equal": true + }, + { + "suite": "kernel", + "label": "llama3_70b_down", + "m": 1, + "n": 8192, + "k": 28672, + "torch_median_ms": 0.628326416015625, + "torch_p90_ms": 0.628592643737793, + "candidate_median_ms": 0.6206105422973632, + "candidate_p90_ms": 0.6207587051391602, + "speedup_pct": 1.2432714548643142, + "max_abs": 6.103515625e-05, + "relative_l2": 1.98126826296999e-07, + "argmax_equal": true + }, + { + "suite": "kernel", + "label": "llama3_70b_down", + "m": 2, + "n": 8192, + "k": 28672, + "torch_median_ms": 0.6553958129882812, + "torch_p90_ms": 0.655467529296875, + "candidate_median_ms": 0.6208819198608398, + "candidate_p90_ms": 0.6209382247924805, + "speedup_pct": 5.558849762476092, + "max_abs": 0.0625, + "relative_l2": 0.0025582788887146407, + "argmax_equal": true + }, + { + "suite": "kernel", + "label": "llama3_70b_down", + "m": 4, + "n": 8192, + "k": 28672, + "torch_median_ms": 0.6599577331542968, + "torch_p90_ms": 0.6600806427001953, + "candidate_median_ms": 0.6216960144042969, + "candidate_p90_ms": 0.6218291091918945, + "speedup_pct": 6.154409528692573, + "max_abs": 0.0625, + "relative_l2": 0.00257769657804324, + "argmax_equal": true + }, + { + "suite": "kernel", + "label": "llama3_70b_down", + "m": 8, + "n": 8192, + "k": 28672, + "torch_median_ms": 0.6522316741943359, + "torch_p90_ms": 0.6523804473876953, + "candidate_median_ms": 0.6264627075195313, + "candidate_p90_ms": 0.6266726303100586, + "speedup_pct": 4.113407927638102, + "max_abs": 0.0625, + "relative_l2": 0.00048372845611638154, + "argmax_equal": true + }, + { + "suite": "kernel", + "label": "opt_1_3b_qkvo", + "m": 1, + "n": 2048, + "k": 2048, + "torch_median_ms": 0.010961920022964478, + "torch_p90_ms": 0.011033600568771363, + "candidate_median_ms": 0.006758400201797485, + "candidate_p90_ms": 0.006799200177192688, + "speedup_pct": 62.19696519376008, + "max_abs": 0.0, + "relative_l2": 0.0, + "argmax_equal": true + }, + { + "suite": "kernel", + "label": "opt_1_3b_qkvo", + "m": 2, + "n": 2048, + "k": 2048, + "torch_median_ms": 0.011939840316772461, + "torch_p90_ms": 0.012006399631500244, + "candidate_median_ms": 0.006599680185317993, + "candidate_p90_ms": 0.006732800006866455, + "speedup_pct": 80.91543804402035, + "max_abs": 0.0078125, + "relative_l2": 0.00016236410125643511, + "argmax_equal": true + }, + { + "suite": "kernel", + "label": "opt_1_3b_qkvo", + "m": 4, + "n": 2048, + "k": 2048, + "torch_median_ms": 0.011934720277786255, + "torch_p90_ms": 0.011991039514541627, + "candidate_median_ms": 0.006630399823188781, + "candidate_p90_ms": 0.006722559928894043, + "speedup_pct": 80.0000089896004, + "max_abs": 0.00390625, + "relative_l2": 5.538061428058506e-05, + "argmax_equal": true + }, + { + "suite": "kernel", + "label": "opt_1_3b_qkvo", + "m": 8, + "n": 2048, + "k": 2048, + "torch_median_ms": 0.011975680589675903, + "torch_p90_ms": 0.012098560333251953, + "candidate_median_ms": 0.009728000164031983, + "candidate_p90_ms": 0.00975871980190277, + "speedup_pct": 23.105267143748897, + "max_abs": 0.001953125, + "relative_l2": 1.959762156592698e-05, + "argmax_equal": true + }, + { + "suite": "kernel", + "label": "opt_1_3b_up", + "m": 1, + "n": 8192, + "k": 2048, + "torch_median_ms": 0.0156876802444458, + "torch_p90_ms": 0.015733760595321656, + "candidate_median_ms": 0.011207679510116577, + "candidate_p90_ms": 0.011223039627075194, + "speedup_pct": 39.97259852305166, + "max_abs": 7.62939453125e-06, + "relative_l2": 9.629044896962479e-08, + "argmax_equal": true + }, + { + "suite": "kernel", + "label": "opt_1_3b_up", + "m": 2, + "n": 8192, + "k": 2048, + "torch_median_ms": 0.01619968056678772, + "torch_p90_ms": 0.01625599980354309, + "candidate_median_ms": 0.011699199676513672, + "candidate_p90_ms": 0.011755520105361938, + "speedup_pct": 38.46828000815161, + "max_abs": 0.0078125, + "relative_l2": 0.0001241812396716967, + "argmax_equal": true + }, + { + "suite": "kernel", + "label": "opt_1_3b_up", + "m": 4, + "n": 8192, + "k": 2048, + "torch_median_ms": 0.01646080017089844, + "torch_p90_ms": 0.016604160070419312, + "candidate_median_ms": 0.01696768045425415, + "candidate_p90_ms": 0.017142080068588257, + "speedup_pct": -2.987328083660523, + "max_abs": 0.015625, + "relative_l2": 0.0001591981462703407, + "argmax_equal": true + }, + { + "suite": "kernel", + "label": "opt_1_3b_up", + "m": 8, + "n": 8192, + "k": 2048, + "torch_median_ms": 0.016860159635543822, + "torch_p90_ms": 0.017141760587692262, + "candidate_median_ms": 0.029204480648040772, + "candidate_p90_ms": 0.029726719856262206, + "speedup_pct": -42.268585979203465, + "max_abs": 0.0078125, + "relative_l2": 8.000112697917659e-05, + "argmax_equal": true + }, + { + "suite": "kernel", + "label": "opt_1_3b_down", + "m": 1, + "n": 2048, + "k": 8192, + "torch_median_ms": 0.014259040355682373, + "torch_p90_ms": 0.014305280447006226, + "candidate_median_ms": 0.011499520540237427, + "candidate_p90_ms": 0.011525119543075562, + "speedup_pct": 23.996824961434182, + "max_abs": 0.0, + "relative_l2": 0.0, + "argmax_equal": true + }, + { + "suite": "kernel", + "label": "opt_1_3b_down", + "m": 2, + "n": 2048, + "k": 8192, + "torch_median_ms": 0.036275200843811035, + "torch_p90_ms": 0.036413440704345705, + "candidate_median_ms": 0.012313599586486817, + "candidate_p90_ms": 0.012349439859390259, + "speedup_pct": 194.59461134029522, + "max_abs": 0.015625, + "relative_l2": 0.0002446471270340609, + "argmax_equal": true + }, + { + "suite": "kernel", + "label": "opt_1_3b_down", + "m": 4, + "n": 2048, + "k": 8192, + "torch_median_ms": 0.03610111951828003, + "torch_p90_ms": 0.03620863914489746, + "candidate_median_ms": 0.01481727957725525, + "candidate_p90_ms": 0.014847999811172486, + "speedup_pct": 143.64202166837563, + "max_abs": 0.03125, + "relative_l2": 0.0003093974547856938, + "argmax_equal": true + }, + { + "suite": "kernel", + "label": "opt_1_3b_down", + "m": 8, + "n": 2048, + "k": 8192, + "torch_median_ms": 0.036090879440307616, + "torch_p90_ms": 0.036362240314483645, + "candidate_median_ms": 0.0255948805809021, + "candidate_p90_ms": 0.02564608097076416, + "speedup_pct": 41.00819625326644, + "max_abs": 0.015625, + "relative_l2": 0.000243702763099258, + "argmax_equal": true + } + ] +} diff --git a/docs/benchmarks/gemv_common_halfcta_vs_256cta_l20_sm89.json b/docs/benchmarks/gemv_common_halfcta_vs_256cta_l20_sm89.json new file mode 100644 index 00000000..e062239c --- /dev/null +++ b/docs/benchmarks/gemv_common_halfcta_vs_256cta_l20_sm89.json @@ -0,0 +1,1168 @@ +{ + "environment": { + "device": "NVIDIA L20", + "torch": "2.11.0+cu128", + "cuda": "12.8" + }, + "parameters": { + "baseline_module": "/home/kxqandccx/lxt/AstrAI-bench-gemv-baseline-416f4d0/astrai/extension/lib/bf16_gemv.cpython-312-x86_64-linux-gnu.so", + "candidate_module": "/home/kxqandccx/lxt/AstrAI-bench-gemv-halfcta-416f4d0/astrai/extension/lib/bf16_gemv_halfcta.cpython-312-x86_64-linux-gnu.so", + "m": [ + 1, + 2, + 4, + 8 + ], + "warmup": 40, + "samples": 21, + "inner": 200, + "seed": 20260902, + "output": "docs/benchmarks/gemv_common_halfcta_vs_256cta_l20_sm89.json" + }, + "results": [ + { + "label": "llama2_7b_qo", + "m": 1, + "n": 4096, + "k": 4096, + "baseline_median_ms": 0.011248639822006225, + "baseline_p90_ms": 0.011263999938964844, + "candidate_median_ms": 0.011304960250854493, + "candidate_p90_ms": 0.01132032036781311, + "delta_pct": -0.49819218819465716, + "max_abs": 0.0, + "argmax_equal": true + }, + { + "label": "llama2_7b_qo", + "m": 2, + "n": 4096, + "k": 4096, + "baseline_median_ms": 0.012354559898376465, + "baseline_p90_ms": 0.012375040054321289, + "candidate_median_ms": 0.011724799871444702, + "candidate_p90_ms": 0.011755520105361938, + "delta_pct": 5.3711793278921505, + "max_abs": 9.5367431640625e-07, + "argmax_equal": true + }, + { + "label": "llama2_7b_qo", + "m": 4, + "n": 4096, + "k": 4096, + "baseline_median_ms": 0.01630720019340515, + "baseline_p90_ms": 0.01633280038833618, + "candidate_median_ms": 0.016302080154418946, + "candidate_p90_ms": 0.016312320232391358, + "delta_pct": 0.031407274027017706, + "max_abs": 0.0, + "argmax_equal": true + }, + { + "label": "llama2_7b_qo", + "m": 8, + "n": 4096, + "k": 4096, + "baseline_median_ms": 0.030100479125976562, + "baseline_p90_ms": 0.030197761058807372, + "candidate_median_ms": 0.02437648057937622, + "candidate_p90_ms": 0.024483840465545654, + "delta_pct": 23.481644644974487, + "max_abs": 0.00390625, + "argmax_equal": true + }, + { + "label": "llama2_7b_up_gate", + "m": 1, + "n": 11008, + "k": 4096, + "baseline_median_ms": 0.025338881015777588, + "baseline_p90_ms": 0.025569279193878174, + "candidate_median_ms": 0.025318400859832765, + "candidate_p90_ms": 0.025518081188201903, + "delta_pct": 0.08089040085195176, + "max_abs": 0.00048828125, + "argmax_equal": true + }, + { + "label": "llama2_7b_up_gate", + "m": 2, + "n": 11008, + "k": 4096, + "baseline_median_ms": 0.02895872116088867, + "baseline_p90_ms": 0.029035520553588868, + "candidate_median_ms": 0.027351040840148926, + "candidate_p90_ms": 0.027448320388793947, + "delta_pct": 5.877949326081255, + "max_abs": 3.814697265625e-06, + "argmax_equal": true + }, + { + "label": "llama2_7b_up_gate", + "m": 4, + "n": 11008, + "k": 4096, + "baseline_median_ms": 0.031027359962463377, + "baseline_p90_ms": 0.03107327938079834, + "candidate_median_ms": 0.03102207899093628, + "candidate_p90_ms": 0.031047680377960206, + "delta_pct": 0.0170232676173665, + "max_abs": 0.0, + "argmax_equal": true + }, + { + "label": "llama2_7b_up_gate", + "m": 8, + "n": 11008, + "k": 4096, + "baseline_median_ms": 0.07837183952331543, + "baseline_p90_ms": 0.07922175884246827, + "candidate_median_ms": 0.06047743797302246, + "candidate_p90_ms": 0.06155776023864746, + "delta_pct": 29.588557567989614, + "max_abs": 0.0078125, + "argmax_equal": true + }, + { + "label": "llama2_7b_down", + "m": 1, + "n": 4096, + "k": 11008, + "baseline_median_ms": 0.026224639415740967, + "baseline_p90_ms": 0.026275839805603027, + "candidate_median_ms": 0.026844160556793215, + "candidate_p90_ms": 0.02694144010543823, + "delta_pct": -2.3078432262448656, + "max_abs": 0.015625, + "argmax_equal": true + }, + { + "label": "llama2_7b_down", + "m": 2, + "n": 4096, + "k": 11008, + "baseline_median_ms": 0.02950144052505493, + "baseline_p90_ms": 0.029521920680999757, + "candidate_median_ms": 0.030914559364318847, + "candidate_p90_ms": 0.030970880985260008, + "delta_pct": -4.571046355895724, + "max_abs": 0.0, + "argmax_equal": true + }, + { + "label": "llama2_7b_down", + "m": 4, + "n": 4096, + "k": 11008, + "baseline_median_ms": 0.040422401428222655, + "baseline_p90_ms": 0.040698881149291995, + "candidate_median_ms": 0.04043263912200928, + "candidate_p90_ms": 0.040632319450378415, + "delta_pct": -0.025320369901471462, + "max_abs": 0.0, + "argmax_equal": true + }, + { + "label": "llama2_7b_down", + "m": 8, + "n": 4096, + "k": 11008, + "baseline_median_ms": 0.08580096244812012, + "baseline_p90_ms": 0.08653823852539062, + "candidate_median_ms": 0.06162432193756104, + "candidate_p90_ms": 0.06214144229888916, + "delta_pct": 39.232302685707964, + "max_abs": 0.015625, + "argmax_equal": true + }, + { + "label": "llama3_8b_kv", + "m": 1, + "n": 1024, + "k": 4096, + "baseline_median_ms": 0.005329920053482056, + "baseline_p90_ms": 0.005375999808311462, + "candidate_median_ms": 0.005314559936523438, + "candidate_p90_ms": 0.00535040020942688, + "delta_pct": 0.28901954521309836, + "max_abs": 0.0, + "argmax_equal": true + }, + { + "label": "llama3_8b_kv", + "m": 2, + "n": 1024, + "k": 4096, + "baseline_median_ms": 0.005570560097694397, + "baseline_p90_ms": 0.005575680136680603, + "candidate_median_ms": 0.005258240103721619, + "candidate_p90_ms": 0.005278720259666443, + "delta_pct": 5.939629758476195, + "max_abs": 0.0, + "argmax_equal": true + }, + { + "label": "llama3_8b_kv", + "m": 4, + "n": 1024, + "k": 4096, + "baseline_median_ms": 0.006369280219078064, + "baseline_p90_ms": 0.006487039923667908, + "candidate_median_ms": 0.006364160180091858, + "candidate_p90_ms": 0.006389600038528442, + "delta_pct": 0.08045113324177677, + "max_abs": 0.0, + "argmax_equal": true + }, + { + "label": "llama3_8b_kv", + "m": 8, + "n": 1024, + "k": 4096, + "baseline_median_ms": 0.010199040174484253, + "baseline_p90_ms": 0.010245120525360108, + "candidate_median_ms": 0.009169920086860657, + "candidate_p90_ms": 0.009195520281791688, + "delta_pct": 11.222781418762807, + "max_abs": 0.00048828125, + "argmax_equal": true + }, + { + "label": "llama3_8b_up_gate", + "m": 1, + "n": 14336, + "k": 4096, + "baseline_median_ms": 0.1537228775024414, + "baseline_p90_ms": 0.15385600090026855, + "candidate_median_ms": 0.15530495643615722, + "candidate_p90_ms": 0.15557632446289063, + "delta_pct": -1.0186918499063946, + "max_abs": 0.0, + "argmax_equal": true + }, + { + "label": "llama3_8b_up_gate", + "m": 2, + "n": 14336, + "k": 4096, + "baseline_median_ms": 0.1540454387664795, + "baseline_p90_ms": 0.1541119956970215, + "candidate_median_ms": 0.15562239646911621, + "candidate_p90_ms": 0.15565312385559082, + "delta_pct": -1.0133231067095627, + "max_abs": 0.000244140625, + "argmax_equal": true + }, + { + "label": "llama3_8b_up_gate", + "m": 4, + "n": 14336, + "k": 4096, + "baseline_median_ms": 0.15432703971862793, + "baseline_p90_ms": 0.15435263633728027, + "candidate_median_ms": 0.15609343528747557, + "candidate_p90_ms": 0.15613439559936523, + "delta_pct": -1.1316270704110565, + "max_abs": 0.0078125, + "argmax_equal": true + }, + { + "label": "llama3_8b_up_gate", + "m": 8, + "n": 14336, + "k": 4096, + "baseline_median_ms": 0.15612416267395018, + "baseline_p90_ms": 0.15629311561584472, + "candidate_median_ms": 0.1584179210662842, + "candidate_p90_ms": 0.15860223770141602, + "delta_pct": -1.447915978757397, + "max_abs": 0.0078125, + "argmax_equal": true + }, + { + "label": "llama3_8b_down", + "m": 1, + "n": 4096, + "k": 14336, + "baseline_median_ms": 0.15725055694580078, + "baseline_p90_ms": 0.15730688095092774, + "candidate_median_ms": 0.15714816093444825, + "candidate_p90_ms": 0.15718912124633788, + "delta_pct": 0.06515889892930726, + "max_abs": 0.0, + "argmax_equal": true + }, + { + "label": "llama3_8b_down", + "m": 2, + "n": 4096, + "k": 14336, + "baseline_median_ms": 0.15744511604309083, + "baseline_p90_ms": 0.15746560096740722, + "candidate_median_ms": 0.15731712341308593, + "candidate_p90_ms": 0.15737343788146974, + "delta_pct": 0.08135963029836102, + "max_abs": 0.0, + "argmax_equal": true + }, + { + "label": "llama3_8b_down", + "m": 4, + "n": 4096, + "k": 14336, + "baseline_median_ms": 0.15765503883361817, + "baseline_p90_ms": 0.1577011203765869, + "candidate_median_ms": 0.15737855911254883, + "candidate_p90_ms": 0.15746047973632812, + "delta_pct": 0.17567813724335135, + "max_abs": 1.52587890625e-05, + "argmax_equal": true + }, + { + "label": "llama3_8b_down", + "m": 8, + "n": 4096, + "k": 14336, + "baseline_median_ms": 0.16097791671752928, + "baseline_p90_ms": 0.16116735458374024, + "candidate_median_ms": 0.15961600303649903, + "candidate_p90_ms": 0.1597439956665039, + "delta_pct": 0.8532438196180214, + "max_abs": 0.0078125, + "argmax_equal": true + }, + { + "label": "llama2_13b_qo", + "m": 1, + "n": 5120, + "k": 5120, + "baseline_median_ms": 0.016808960437774658, + "baseline_p90_ms": 0.01682944059371948, + "candidate_median_ms": 0.01627135992050171, + "candidate_p90_ms": 0.01628159999847412, + "delta_pct": 3.3039679528911403, + "max_abs": 1.52587890625e-05, + "argmax_equal": true + }, + { + "label": "llama2_13b_qo", + "m": 2, + "n": 5120, + "k": 5120, + "baseline_median_ms": 0.017546240091323852, + "baseline_p90_ms": 0.01784319996833801, + "candidate_median_ms": 0.016921600103378297, + "candidate_p90_ms": 0.017151999473571777, + "delta_pct": 3.691376608178154, + "max_abs": 0.0, + "argmax_equal": true + }, + { + "label": "llama2_13b_qo", + "m": 4, + "n": 5120, + "k": 5120, + "baseline_median_ms": 0.0235263991355896, + "baseline_p90_ms": 0.023557119369506836, + "candidate_median_ms": 0.021186399459838866, + "candidate_p90_ms": 0.02120192050933838, + "delta_pct": 11.044819957192153, + "max_abs": 0.000244140625, + "argmax_equal": true + }, + { + "label": "llama2_13b_qo", + "m": 8, + "n": 5120, + "k": 5120, + "baseline_median_ms": 0.04256256103515625, + "baseline_p90_ms": 0.042777600288391116, + "candidate_median_ms": 0.03810303926467896, + "candidate_p90_ms": 0.038251519203186035, + "delta_pct": 11.703847925357525, + "max_abs": 0.0078125, + "argmax_equal": true + }, + { + "label": "llama2_13b_up_gate", + "m": 1, + "n": 13824, + "k": 5120, + "baseline_median_ms": 0.18877952575683593, + "baseline_p90_ms": 0.1890764808654785, + "candidate_median_ms": 0.18909696578979493, + "candidate_p90_ms": 0.1893427276611328, + "delta_pct": -0.16787156347705112, + "max_abs": 0.0078125, + "argmax_equal": true + }, + { + "label": "llama2_13b_up_gate", + "m": 2, + "n": 13824, + "k": 5120, + "baseline_median_ms": 0.1888870429992676, + "baseline_p90_ms": 0.18892799377441405, + "candidate_median_ms": 0.1893427276611328, + "candidate_p90_ms": 0.1894553565979004, + "delta_pct": -0.24066657721375506, + "max_abs": 1.52587890625e-05, + "argmax_equal": true + }, + { + "label": "llama2_13b_up_gate", + "m": 4, + "n": 13824, + "k": 5120, + "baseline_median_ms": 0.18933759689331053, + "baseline_p90_ms": 0.1893734359741211, + "candidate_median_ms": 0.1897011184692383, + "candidate_p90_ms": 0.18977792739868163, + "delta_pct": -0.19162858862463716, + "max_abs": 0.001953125, + "argmax_equal": true + }, + { + "label": "llama2_13b_up_gate", + "m": 8, + "n": 13824, + "k": 5120, + "baseline_median_ms": 0.19117055892944335, + "baseline_p90_ms": 0.19130863189697267, + "candidate_median_ms": 0.19213823318481446, + "candidate_p90_ms": 0.19232767105102538, + "delta_pct": -0.5036344091081113, + "max_abs": 0.0078125, + "argmax_equal": true + }, + { + "label": "llama2_13b_down", + "m": 1, + "n": 5120, + "k": 13824, + "baseline_median_ms": 0.1890559959411621, + "baseline_p90_ms": 0.18912256240844727, + "candidate_median_ms": 0.18902528762817383, + "candidate_p90_ms": 0.18909664154052735, + "delta_pct": 0.01624561103628075, + "max_abs": 0.001953125, + "argmax_equal": true + }, + { + "label": "llama2_13b_down", + "m": 2, + "n": 5120, + "k": 13824, + "baseline_median_ms": 0.1892915153503418, + "baseline_p90_ms": 0.18942464828491212, + "candidate_median_ms": 0.18920448303222656, + "candidate_p90_ms": 0.1892608070373535, + "delta_pct": 0.04599907820388438, + "max_abs": 9.5367431640625e-07, + "argmax_equal": true + }, + { + "label": "llama2_13b_down", + "m": 4, + "n": 5120, + "k": 13824, + "baseline_median_ms": 0.1894963264465332, + "baseline_p90_ms": 0.1895167922973633, + "candidate_median_ms": 0.1893836784362793, + "candidate_p90_ms": 0.18942464828491212, + "delta_pct": 0.05948137198730219, + "max_abs": 0.001953125, + "argmax_equal": true + }, + { + "label": "llama2_13b_down", + "m": 8, + "n": 5120, + "k": 13824, + "baseline_median_ms": 0.19152383804321288, + "baseline_p90_ms": 0.1916262435913086, + "candidate_median_ms": 0.1977190399169922, + "candidate_p90_ms": 0.19799039840698243, + "delta_pct": -3.133336008702159, + "max_abs": 0.00390625, + "argmax_equal": true + }, + { + "label": "gpt_neox_up", + "m": 1, + "n": 16384, + "k": 4096, + "baseline_median_ms": 0.17534975051879884, + "baseline_p90_ms": 0.17538560867309572, + "candidate_median_ms": 0.17821184158325196, + "candidate_p90_ms": 0.17826303482055664, + "delta_pct": -1.6060049876742255, + "max_abs": 0.0009765625, + "argmax_equal": true + }, + { + "label": "gpt_neox_up", + "m": 2, + "n": 16384, + "k": 4096, + "baseline_median_ms": 0.1755392074584961, + "baseline_p90_ms": 0.17556991577148437, + "candidate_median_ms": 0.17774591445922852, + "candidate_p90_ms": 0.17777151107788086, + "delta_pct": -1.2414952025457615, + "max_abs": 0.000244140625, + "argmax_equal": true + }, + { + "label": "gpt_neox_up", + "m": 4, + "n": 16384, + "k": 4096, + "baseline_median_ms": 0.1760767936706543, + "baseline_p90_ms": 0.1761177635192871, + "candidate_median_ms": 0.1783603286743164, + "candidate_p90_ms": 0.17839103698730469, + "delta_pct": -1.2802931126191242, + "max_abs": 0.0078125, + "argmax_equal": true + }, + { + "label": "gpt_neox_up", + "m": 8, + "n": 16384, + "k": 4096, + "baseline_median_ms": 0.1780121612548828, + "baseline_p90_ms": 0.1781964874267578, + "candidate_median_ms": 0.1804902458190918, + "candidate_p90_ms": 0.18060800552368164, + "delta_pct": -1.3729742308029258, + "max_abs": 0.001953125, + "argmax_equal": true + }, + { + "label": "gpt_neox_down", + "m": 1, + "n": 4096, + "k": 16384, + "baseline_median_ms": 0.17917440414428712, + "baseline_p90_ms": 0.17921024322509765, + "candidate_median_ms": 0.17931264877319336, + "candidate_p90_ms": 0.17934335708618165, + "delta_pct": -0.0770969755073514, + "max_abs": 0.015625, + "argmax_equal": true + }, + { + "label": "gpt_neox_down", + "m": 2, + "n": 4096, + "k": 16384, + "baseline_median_ms": 0.17925615310668946, + "baseline_p90_ms": 0.1792972755432129, + "candidate_median_ms": 0.17945087432861329, + "candidate_p90_ms": 0.1794715118408203, + "delta_pct": -0.10850948631615509, + "max_abs": 0.000244140625, + "argmax_equal": true + }, + { + "label": "gpt_neox_down", + "m": 4, + "n": 4096, + "k": 16384, + "baseline_median_ms": 0.17968128204345704, + "baseline_p90_ms": 0.17971200942993165, + "candidate_median_ms": 0.1798246383666992, + "candidate_p90_ms": 0.17986047744750977, + "delta_pct": -0.0797200675859755, + "max_abs": 0.0078125, + "argmax_equal": true + }, + { + "label": "gpt_neox_down", + "m": 8, + "n": 4096, + "k": 16384, + "baseline_median_ms": 0.18399744033813475, + "baseline_p90_ms": 0.1840537643432617, + "candidate_median_ms": 0.1815398406982422, + "candidate_p90_ms": 0.18169343948364258, + "delta_pct": 1.353752228954308, + "max_abs": 0.0078125, + "argmax_equal": true + }, + { + "label": "qwen2_7b_kv", + "m": 1, + "n": 512, + "k": 3584, + "baseline_median_ms": 0.0053708797693252566, + "baseline_p90_ms": 0.005442559719085693, + "candidate_median_ms": 0.0053862398862838745, + "candidate_p90_ms": 0.005463039875030518, + "delta_pct": -0.28517328011573584, + "max_abs": 0.0, + "argmax_equal": true + }, + { + "label": "qwen2_7b_kv", + "m": 2, + "n": 512, + "k": 3584, + "baseline_median_ms": 0.0053350400924682614, + "baseline_p90_ms": 0.005391200184822083, + "candidate_median_ms": 0.005237759947776794, + "candidate_p90_ms": 0.00529919981956482, + "delta_pct": 1.8572852834303433, + "max_abs": 0.0, + "argmax_equal": true + }, + { + "label": "qwen2_7b_kv", + "m": 4, + "n": 512, + "k": 3584, + "baseline_median_ms": 0.0053355199098587035, + "baseline_p90_ms": 0.0053708797693252566, + "candidate_median_ms": 0.005283839702606201, + "candidate_p90_ms": 0.00532480001449585, + "delta_pct": 0.9780805278216897, + "max_abs": 0.0, + "argmax_equal": true + }, + { + "label": "qwen2_7b_kv", + "m": 8, + "n": 512, + "k": 3584, + "baseline_median_ms": 0.007388160228729248, + "baseline_p90_ms": 0.007393280267715454, + "candidate_median_ms": 0.0062412798404693605, + "candidate_p90_ms": 0.0062566399574279785, + "delta_pct": 18.37572449200802, + "max_abs": 1.9073486328125e-06, + "argmax_equal": true + }, + { + "label": "qwen2_7b_qo", + "m": 1, + "n": 3584, + "k": 3584, + "baseline_median_ms": 0.009589920043945313, + "baseline_p90_ms": 0.009605119824409486, + "candidate_median_ms": 0.009338880181312561, + "candidate_p90_ms": 0.009349120259284973, + "delta_pct": 2.688115253208756, + "max_abs": 3.0517578125e-05, + "argmax_equal": true + }, + { + "label": "qwen2_7b_qo", + "m": 2, + "n": 3584, + "k": 3584, + "baseline_median_ms": 0.010562560558319091, + "baseline_p90_ms": 0.010577600002288818, + "candidate_median_ms": 0.009988800287246705, + "candidate_p90_ms": 0.009994239807128906, + "delta_pct": 5.744035865898134, + "max_abs": 0.001953125, + "argmax_equal": true + }, + { + "label": "qwen2_7b_qo", + "m": 4, + "n": 3584, + "k": 3584, + "baseline_median_ms": 0.014279680252075195, + "baseline_p90_ms": 0.014371839761734008, + "candidate_median_ms": 0.01225216031074524, + "candidate_p90_ms": 0.012328959703445434, + "delta_pct": 16.548264876617758, + "max_abs": 0.00048828125, + "argmax_equal": true + }, + { + "label": "qwen2_7b_qo", + "m": 8, + "n": 3584, + "k": 3584, + "baseline_median_ms": 0.025246880054473876, + "baseline_p90_ms": 0.0253439998626709, + "candidate_median_ms": 0.020423519611358642, + "candidate_p90_ms": 0.020500481128692627, + "delta_pct": 23.616695529955066, + "max_abs": 0.0078125, + "argmax_equal": true + }, + { + "label": "qwen2_7b_up_gate", + "m": 1, + "n": 18944, + "k": 3584, + "baseline_median_ms": 0.17954303741455077, + "baseline_p90_ms": 0.17956335067749024, + "candidate_median_ms": 0.18173952102661134, + "candidate_p90_ms": 0.18176511764526368, + "delta_pct": -1.2085888636951703, + "max_abs": 0.000244140625, + "argmax_equal": true + }, + { + "label": "qwen2_7b_up_gate", + "m": 2, + "n": 18944, + "k": 3584, + "baseline_median_ms": 0.17977855682373048, + "baseline_p90_ms": 0.1798092842102051, + "candidate_median_ms": 0.18187776565551758, + "candidate_p90_ms": 0.18189823150634765, + "delta_pct": -1.1541866177106397, + "max_abs": 0.00048828125, + "argmax_equal": true + }, + { + "label": "qwen2_7b_up_gate", + "m": 4, + "n": 18944, + "k": 3584, + "baseline_median_ms": 0.18035200119018555, + "baseline_p90_ms": 0.18039295196533203, + "candidate_median_ms": 0.1824665641784668, + "candidate_p90_ms": 0.18250240325927736, + "delta_pct": -1.158876968940481, + "max_abs": 0.001953125, + "argmax_equal": true + }, + { + "label": "qwen2_7b_up_gate", + "m": 8, + "n": 18944, + "k": 3584, + "baseline_median_ms": 0.18180608749389648, + "baseline_p90_ms": 0.18197504043579102, + "candidate_median_ms": 0.1844068717956543, + "candidate_p90_ms": 0.18456031799316405, + "delta_pct": -1.4103510766343863, + "max_abs": 0.0078125, + "argmax_equal": true + }, + { + "label": "qwen2_7b_down", + "m": 1, + "n": 3584, + "k": 18944, + "baseline_median_ms": 0.1816268730163574, + "baseline_p90_ms": 0.1816422462463379, + "candidate_median_ms": 0.18170879364013673, + "candidate_p90_ms": 0.1817241668701172, + "delta_pct": -0.04508346686927789, + "max_abs": 0.0009765625, + "argmax_equal": true + }, + { + "label": "qwen2_7b_down", + "m": 2, + "n": 3584, + "k": 18944, + "baseline_median_ms": 0.18167295455932617, + "baseline_p90_ms": 0.18171392440795897, + "candidate_median_ms": 0.18178560256958007, + "candidate_p90_ms": 0.18180095672607421, + "delta_pct": -0.06196750934155393, + "max_abs": 0.0078125, + "argmax_equal": true + }, + { + "label": "qwen2_7b_down", + "m": 4, + "n": 3584, + "k": 18944, + "baseline_median_ms": 0.18192895889282226, + "baseline_p90_ms": 0.18195968627929687, + "candidate_median_ms": 0.1820672035217285, + "candidate_p90_ms": 0.18208255767822265, + "delta_pct": -0.07593054994649018, + "max_abs": 0.001953125, + "argmax_equal": true + }, + { + "label": "qwen2_7b_down", + "m": 8, + "n": 3584, + "k": 18944, + "baseline_median_ms": 0.18429439544677734, + "baseline_p90_ms": 0.1844326400756836, + "candidate_median_ms": 0.18709503173828124, + "candidate_p90_ms": 0.18755071640014648, + "delta_pct": -1.4969057518435802, + "max_abs": 0.015625, + "argmax_equal": true + }, + { + "label": "llama3_70b_kv", + "m": 1, + "n": 1024, + "k": 8192, + "baseline_median_ms": 0.00736240029335022, + "baseline_p90_ms": 0.007367680072784424, + "candidate_median_ms": 0.0073471999168395995, + "candidate_p90_ms": 0.007352319955825806, + "delta_pct": 0.2068866599884034, + "max_abs": 1.52587890625e-05, + "argmax_equal": true + }, + { + "label": "llama3_70b_kv", + "m": 2, + "n": 1024, + "k": 8192, + "baseline_median_ms": 0.00796176016330719, + "baseline_p90_ms": 0.007976959943771362, + "candidate_median_ms": 0.008002560138702392, + "candidate_p90_ms": 0.008012800216674805, + "delta_pct": -0.5098365359090806, + "max_abs": 0.0, + "argmax_equal": true + }, + { + "label": "llama3_70b_kv", + "m": 4, + "n": 1024, + "k": 8192, + "baseline_median_ms": 0.009861119985580445, + "baseline_p90_ms": 0.009876319766044616, + "candidate_median_ms": 0.009456639885902405, + "candidate_p90_ms": 0.009466879963874818, + "delta_pct": 4.277207386114212, + "max_abs": 0.00048828125, + "argmax_equal": true + }, + { + "label": "llama3_70b_kv", + "m": 8, + "n": 1024, + "k": 8192, + "baseline_median_ms": 0.01817088007926941, + "baseline_p90_ms": 0.018201600313186645, + "candidate_median_ms": 0.016721919775009156, + "candidate_p90_ms": 0.01675264000892639, + "delta_pct": 8.665035616459061, + "max_abs": 0.0, + "argmax_equal": true + }, + { + "label": "llama3_70b_qo", + "m": 1, + "n": 8192, + "k": 8192, + "baseline_median_ms": 0.17747455596923828, + "baseline_p90_ms": 0.17754112243652342, + "candidate_median_ms": 0.17870847702026368, + "candidate_p90_ms": 0.1787392044067383, + "delta_pct": -0.6904658758215998, + "max_abs": 0.0078125, + "argmax_equal": true + }, + { + "label": "llama3_70b_qo", + "m": 2, + "n": 8192, + "k": 8192, + "baseline_median_ms": 0.17808895111083983, + "baseline_p90_ms": 0.17811456680297852, + "candidate_median_ms": 0.17896448135375975, + "candidate_p90_ms": 0.1789798355102539, + "delta_pct": -0.48922011580010416, + "max_abs": 3.0517578125e-05, + "argmax_equal": true + }, + { + "label": "llama3_70b_qo", + "m": 4, + "n": 8192, + "k": 8192, + "baseline_median_ms": 0.17745920181274413, + "baseline_p90_ms": 0.17748992919921874, + "candidate_median_ms": 0.17948671340942382, + "candidate_p90_ms": 0.17957376480102538, + "delta_pct": -1.1296165371610378, + "max_abs": 0.001953125, + "argmax_equal": true + }, + { + "label": "llama3_70b_qo", + "m": 8, + "n": 8192, + "k": 8192, + "baseline_median_ms": 0.17998336791992187, + "baseline_p90_ms": 0.18009599685668945, + "candidate_median_ms": 0.18143232345581053, + "candidate_p90_ms": 0.1816422462463379, + "delta_pct": -0.7986203937037528, + "max_abs": 0.0078125, + "argmax_equal": true + }, + { + "label": "llama3_70b_up_gate", + "m": 1, + "n": 28672, + "k": 8192, + "baseline_median_ms": 0.6200934219360351, + "baseline_p90_ms": 0.6202521514892578, + "candidate_median_ms": 0.6198988723754882, + "candidate_p90_ms": 0.6200627136230469, + "delta_pct": 0.03138408040674889, + "max_abs": 0.015625, + "argmax_equal": true + }, + { + "label": "llama3_70b_up_gate", + "m": 2, + "n": 28672, + "k": 8192, + "baseline_median_ms": 0.6201087951660156, + "baseline_p90_ms": 0.6202214431762695, + "candidate_median_ms": 0.6210969543457031, + "candidate_p90_ms": 0.6212556838989258, + "delta_pct": -0.15909902194393144, + "max_abs": 0.001953125, + "argmax_equal": true + }, + { + "label": "llama3_70b_up_gate", + "m": 4, + "n": 28672, + "k": 8192, + "baseline_median_ms": 0.6240409469604492, + "baseline_p90_ms": 0.6241126251220703, + "candidate_median_ms": 0.6228633499145508, + "candidate_p90_ms": 0.6229504013061523, + "delta_pct": 0.1890618618128448, + "max_abs": 0.015625, + "argmax_equal": true + }, + { + "label": "llama3_70b_up_gate", + "m": 8, + "n": 28672, + "k": 8192, + "baseline_median_ms": 0.6230116653442382, + "baseline_p90_ms": 0.6231449508666992, + "candidate_median_ms": 0.6242457580566406, + "candidate_p90_ms": 0.6246092987060546, + "delta_pct": -0.1976934078405712, + "max_abs": 0.03125, + "argmax_equal": true + }, + { + "label": "llama3_70b_down", + "m": 1, + "n": 8192, + "k": 28672, + "baseline_median_ms": 0.6207334518432617, + "baseline_p90_ms": 0.6208256149291992, + "candidate_median_ms": 0.6206054306030273, + "candidate_p90_ms": 0.62076416015625, + "delta_pct": 0.020628443439485444, + "max_abs": 0.0, + "argmax_equal": true + }, + { + "label": "llama3_70b_down", + "m": 2, + "n": 8192, + "k": 28672, + "baseline_median_ms": 0.6209843063354492, + "baseline_p90_ms": 0.6210867309570313, + "candidate_median_ms": 0.6208051300048828, + "candidate_p90_ms": 0.6209587097167969, + "delta_pct": 0.028861928148860017, + "max_abs": 0.00048828125, + "argmax_equal": true + }, + { + "label": "llama3_70b_down", + "m": 4, + "n": 8192, + "k": 28672, + "baseline_median_ms": 0.6219929504394531, + "baseline_p90_ms": 0.6221004867553711, + "candidate_median_ms": 0.6217216110229492, + "candidate_p90_ms": 0.6217984008789063, + "delta_pct": 0.043643233835388706, + "max_abs": 0.00390625, + "argmax_equal": true + }, + { + "label": "llama3_70b_down", + "m": 8, + "n": 8192, + "k": 28672, + "baseline_median_ms": 0.6267750549316407, + "baseline_p90_ms": 0.6273331069946289, + "candidate_median_ms": 0.6232320022583008, + "candidate_p90_ms": 0.6234316635131836, + "delta_pct": 0.5684965888307358, + "max_abs": 0.015625, + "argmax_equal": true + }, + { + "label": "opt_1_3b_qkvo", + "m": 1, + "n": 2048, + "k": 2048, + "baseline_median_ms": 0.005340160131454468, + "baseline_p90_ms": 0.005437440276145935, + "candidate_median_ms": 0.00526848018169403, + "candidate_p90_ms": 0.005432320237159729, + "delta_pct": 1.3605432171786225, + "max_abs": 0.0, + "argmax_equal": true + }, + { + "label": "opt_1_3b_qkvo", + "m": 2, + "n": 2048, + "k": 2048, + "baseline_median_ms": 0.0057190400362014775, + "baseline_p90_ms": 0.005729280114173889, + "candidate_median_ms": 0.005242879986763, + "candidate_p90_ms": 0.005253120064735412, + "delta_pct": 9.082032215894053, + "max_abs": 7.62939453125e-06, + "argmax_equal": true + }, + { + "label": "opt_1_3b_qkvo", + "m": 4, + "n": 2048, + "k": 2048, + "baseline_median_ms": 0.007782400250434875, + "baseline_p90_ms": 0.00779263973236084, + "candidate_median_ms": 0.006568959951400757, + "candidate_p90_ms": 0.006579200029373169, + "delta_pct": 18.47233516434159, + "max_abs": 0.0, + "argmax_equal": true + }, + { + "label": "opt_1_3b_qkvo", + "m": 8, + "n": 2048, + "k": 2048, + "baseline_median_ms": 0.013061120510101318, + "baseline_p90_ms": 0.013071360588073731, + "candidate_median_ms": 0.009707520008087158, + "candidate_p90_ms": 0.009712640047073364, + "delta_pct": 34.546418644724255, + "max_abs": 4.76837158203125e-07, + "argmax_equal": true + }, + { + "label": "opt_1_3b_up", + "m": 1, + "n": 8192, + "k": 2048, + "baseline_median_ms": 0.013004800081253052, + "baseline_p90_ms": 0.013015040159225465, + "candidate_median_ms": 0.011146080493927003, + "candidate_p90_ms": 0.011166720390319825, + "delta_pct": 16.675992859900667, + "max_abs": 0.0, + "argmax_equal": true + }, + { + "label": "opt_1_3b_up", + "m": 2, + "n": 8192, + "k": 2048, + "baseline_median_ms": 0.016112639904022216, + "baseline_p90_ms": 0.01628159999847412, + "candidate_median_ms": 0.0123852801322937, + "candidate_p90_ms": 0.012462079524993896, + "delta_pct": 30.09507844727468, + "max_abs": 0.001953125, + "argmax_equal": true + }, + { + "label": "opt_1_3b_up", + "m": 4, + "n": 8192, + "k": 2048, + "baseline_median_ms": 0.024191999435424806, + "baseline_p90_ms": 0.024232959747314452, + "candidate_median_ms": 0.017679359912872315, + "candidate_p90_ms": 0.017740800380706787, + "delta_pct": 36.83753006131545, + "max_abs": 6.103515625e-05, + "argmax_equal": true + }, + { + "label": "opt_1_3b_up", + "m": 8, + "n": 8192, + "k": 2048, + "baseline_median_ms": 0.04454400062561035, + "baseline_p90_ms": 0.04510704040527344, + "candidate_median_ms": 0.02998784065246582, + "candidate_p90_ms": 0.030392320156097413, + "delta_pct": 48.54020715208654, + "max_abs": 0.0001220703125, + "argmax_equal": true + }, + { + "label": "opt_1_3b_down", + "m": 1, + "n": 2048, + "k": 8192, + "baseline_median_ms": 0.011473920345306397, + "baseline_p90_ms": 0.011586560010910034, + "candidate_median_ms": 0.011806720495223999, + "candidate_p90_ms": 0.011934720277786255, + "delta_pct": -2.8187348896099085, + "max_abs": 0.0, + "argmax_equal": true + }, + { + "label": "opt_1_3b_down", + "m": 2, + "n": 2048, + "k": 8192, + "baseline_median_ms": 0.012344319820404053, + "baseline_p90_ms": 0.01240064024925232, + "candidate_median_ms": 0.01260543942451477, + "candidate_p90_ms": 0.012687360048294067, + "delta_pct": -2.07148355021165, + "max_abs": 0.0, + "argmax_equal": true + }, + { + "label": "opt_1_3b_down", + "m": 4, + "n": 2048, + "k": 8192, + "baseline_median_ms": 0.015621119737625122, + "baseline_p90_ms": 0.015646719932556154, + "candidate_median_ms": 0.01602560043334961, + "candidate_p90_ms": 0.016076799631118775, + "delta_pct": -2.5239659344230003, + "max_abs": 0.0, + "argmax_equal": true + }, + { + "label": "opt_1_3b_down", + "m": 8, + "n": 2048, + "k": 8192, + "baseline_median_ms": 0.033950560092926026, + "baseline_p90_ms": 0.03397119998931885, + "candidate_median_ms": 0.026997759342193603, + "candidate_p90_ms": 0.02703360080718994, + "delta_pct": 25.75325108504911, + "max_abs": 0.00390625, + "argmax_equal": true + } + ] +} diff --git a/docs/benchmarks/gemv_common_kernel_l20_sm89.json b/docs/benchmarks/gemv_common_kernel_l20_sm89.json new file mode 100644 index 00000000..4875724e --- /dev/null +++ b/docs/benchmarks/gemv_common_kernel_l20_sm89.json @@ -0,0 +1,348 @@ +{ + "environment": { + "device": "NVIDIA L20", + "capability": "8.9", + "torch": "2.11.0+cu128", + "cuda": "12.8" + }, + "parameters": { + "seed": 20260902, + "weight_std": 0.02, + "warmup": 40, + "samples": 21, + "inner": 200, + "chain_inner": 20 + }, + "results": [ + { + "suite": "kernel", + "label": "llama2_7b_qo", + "m": 2, + "n": 4096, + "k": 4096, + "torch_median_ms": 0.020986878871917726, + "torch_p90_ms": 0.021032960414886476, + "candidate_median_ms": 0.0122980797290802, + "candidate_p90_ms": 0.012313599586486817, + "speedup_pct": 70.65167354779686, + "max_abs": 0.015625, + "relative_l2": 0.00024786205047956324, + "argmax_equal": true + }, + { + "suite": "kernel", + "label": "llama2_7b_qo", + "m": 4, + "n": 4096, + "k": 4096, + "torch_median_ms": 0.02096127986907959, + "torch_p90_ms": 0.020986878871917726, + "candidate_median_ms": 0.016061439514160156, + "candidate_p90_ms": 0.016102399826049805, + "speedup_pct": 30.506856814419493, + "max_abs": 0.0078125, + "relative_l2": 7.349992864901514e-05, + "argmax_equal": true + }, + { + "suite": "kernel", + "label": "llama2_7b_up_gate", + "m": 2, + "n": 11008, + "k": 4096, + "torch_median_ms": 0.039777278900146484, + "torch_p90_ms": 0.040135679244995115, + "candidate_median_ms": 0.02741247892379761, + "candidate_p90_ms": 0.02762239933013916, + "speedup_pct": 45.106464142557414, + "max_abs": 0.0078125, + "relative_l2": 9.189538707426296e-05, + "argmax_equal": true + }, + { + "suite": "kernel", + "label": "llama2_7b_up_gate", + "m": 4, + "n": 11008, + "k": 4096, + "torch_median_ms": 0.04090367794036865, + "torch_p90_ms": 0.04160511970520019, + "candidate_median_ms": 0.029184160232543947, + "candidate_p90_ms": 0.029762558937072754, + "speedup_pct": 40.15711815739002, + "max_abs": 0.015625, + "relative_l2": 0.00018090539722389557, + "argmax_equal": true + }, + { + "suite": "kernel", + "label": "llama2_7b_down", + "m": 2, + "n": 4096, + "k": 11008, + "torch_median_ms": 0.0516710376739502, + "torch_p90_ms": 0.051742720603942874, + "candidate_median_ms": 0.02998784065246582, + "candidate_p90_ms": 0.030115840435028077, + "speedup_pct": 72.30663011976965, + "max_abs": 0.015625, + "relative_l2": 0.00017434120320033395, + "argmax_equal": true + }, + { + "suite": "kernel", + "label": "llama2_7b_down", + "m": 4, + "n": 4096, + "k": 11008, + "torch_median_ms": 0.05163519859313965, + "torch_p90_ms": 0.05170703887939453, + "candidate_median_ms": 0.039761760234832765, + "candidate_p90_ms": 0.03986943960189819, + "speedup_pct": 29.861450519751685, + "max_abs": 0.03125, + "relative_l2": 0.00034510913677468784, + "argmax_equal": true + }, + { + "suite": "kernel", + "label": "llama3_8b_kv", + "m": 2, + "n": 1024, + "k": 4096, + "torch_median_ms": 0.014115840196609497, + "torch_p90_ms": 0.014320640563964844, + "candidate_median_ms": 0.006584320068359375, + "candidate_p90_ms": 0.006676480174064636, + "speedup_pct": 114.38569282867141, + "max_abs": 0.015625, + "relative_l2": 0.002633088689341241, + "argmax_equal": true + }, + { + "suite": "kernel", + "label": "llama3_8b_kv", + "m": 4, + "n": 1024, + "k": 4096, + "torch_median_ms": 0.014085119962692261, + "torch_p90_ms": 0.014146560430526733, + "candidate_median_ms": 0.006543359756469727, + "candidate_p90_ms": 0.006615359783172608, + "speedup_pct": 115.25822340374367, + "max_abs": 0.03125, + "relative_l2": 0.0025596959926785638, + "argmax_equal": true + }, + { + "suite": "kernel", + "label": "llama3_8b_up_gate", + "m": 2, + "n": 14336, + "k": 4096, + "torch_median_ms": 0.16598016738891602, + "torch_p90_ms": 0.16606719970703124, + "candidate_median_ms": 0.1538969612121582, + "candidate_p90_ms": 0.15395839691162108, + "speedup_pct": 7.851491076617312, + "max_abs": 0.015625, + "relative_l2": 0.0001283098848252402, + "argmax_equal": true + }, + { + "suite": "kernel", + "label": "llama3_8b_up_gate", + "m": 4, + "n": 14336, + "k": 4096, + "torch_median_ms": 0.16636911392211914, + "torch_p90_ms": 0.16643583297729492, + "candidate_median_ms": 0.15441408157348632, + "candidate_p90_ms": 0.15444992065429688, + "speedup_pct": 7.742190496365686, + "max_abs": 0.015625, + "relative_l2": 0.00016471143444015762, + "argmax_equal": true + }, + { + "suite": "kernel", + "label": "llama3_8b_down", + "m": 2, + "n": 4096, + "k": 14336, + "torch_median_ms": 0.16633344650268556, + "torch_p90_ms": 0.16658416748046875, + "candidate_median_ms": 0.1570251178741455, + "candidate_p90_ms": 0.1573529624938965, + "speedup_pct": 5.927923350454423, + "max_abs": 0.03125, + "relative_l2": 0.00027259474851604274, + "argmax_equal": true + }, + { + "suite": "kernel", + "label": "llama3_8b_down", + "m": 4, + "n": 4096, + "k": 14336, + "torch_median_ms": 0.16647680282592772, + "torch_p90_ms": 0.16680959701538087, + "candidate_median_ms": 0.15723520278930664, + "candidate_p90_ms": 0.15757328033447265, + "speedup_pct": 5.8775642303236175, + "max_abs": 0.03125, + "relative_l2": 0.0003724078071877287, + "argmax_equal": true + }, + { + "suite": "kernel", + "label": "llama2_13b_qo", + "m": 2, + "n": 5120, + "k": 5120, + "torch_median_ms": 0.026772480010986328, + "torch_p90_ms": 0.0268339204788208, + "candidate_median_ms": 0.017315839529037477, + "candidate_p90_ms": 0.017351679801940918, + "speedup_pct": 54.612659502247716, + "max_abs": 0.015625, + "relative_l2": 0.00014752121382419835, + "argmax_equal": true + }, + { + "suite": "kernel", + "label": "llama2_13b_qo", + "m": 4, + "n": 5120, + "k": 5120, + "torch_median_ms": 0.026910719871520997, + "torch_p90_ms": 0.026956799030303954, + "candidate_median_ms": 0.021724159717559813, + "candidate_p90_ms": 0.021790719032287596, + "speedup_pct": 23.874618035370297, + "max_abs": 0.015625, + "relative_l2": 0.0001289835715547304, + "argmax_equal": true + }, + { + "suite": "kernel", + "label": "llama2_13b_up_gate", + "m": 2, + "n": 13824, + "k": 5120, + "torch_median_ms": 0.19898368835449218, + "torch_p90_ms": 0.199050235748291, + "candidate_median_ms": 0.18885120391845703, + "candidate_p90_ms": 0.18894336700439454, + "speedup_pct": 5.3653268953531175, + "max_abs": 0.015625, + "relative_l2": 0.00012896713728075684, + "argmax_equal": true + }, + { + "suite": "kernel", + "label": "llama2_13b_up_gate", + "m": 4, + "n": 13824, + "k": 5120, + "torch_median_ms": 0.19945472717285156, + "torch_p90_ms": 0.1995110321044922, + "candidate_median_ms": 0.18927616119384766, + "candidate_p90_ms": 0.18933759689331053, + "speedup_pct": 5.37762701589215, + "max_abs": 0.03125, + "relative_l2": 0.00022244987307607752, + "argmax_equal": true + }, + { + "suite": "kernel", + "label": "llama2_13b_down", + "m": 2, + "n": 5120, + "k": 13824, + "torch_median_ms": 0.20000255584716797, + "torch_p90_ms": 0.20012544631958007, + "candidate_median_ms": 0.1890764808654785, + "candidate_p90_ms": 0.1892198371887207, + "speedup_pct": 5.778653659977406, + "max_abs": 0.03125, + "relative_l2": 0.00033799466525286537, + "argmax_equal": true + }, + { + "suite": "kernel", + "label": "llama2_13b_down", + "m": 4, + "n": 5120, + "k": 13824, + "torch_median_ms": 0.2002841567993164, + "torch_p90_ms": 0.2004377555847168, + "candidate_median_ms": 0.1893939208984375, + "candidate_p90_ms": 0.18960895538330078, + "speedup_pct": 5.750045117191904, + "max_abs": 0.03125, + "relative_l2": 0.0002938090515862735, + "argmax_equal": true + }, + { + "suite": "kernel", + "label": "gpt_neox_up", + "m": 2, + "n": 16384, + "k": 4096, + "torch_median_ms": 0.19273727416992187, + "torch_p90_ms": 0.1927987289428711, + "candidate_median_ms": 0.17540096282958983, + "candidate_p90_ms": 0.17543167114257813, + "speedup_pct": 9.883817660211513, + "max_abs": 0.015625, + "relative_l2": 0.0001748977553606047, + "argmax_equal": true + }, + { + "suite": "kernel", + "label": "gpt_neox_up", + "m": 4, + "n": 16384, + "k": 4096, + "torch_median_ms": 0.193121280670166, + "torch_p90_ms": 0.1931929588317871, + "candidate_median_ms": 0.17605119705200195, + "candidate_p90_ms": 0.17608192443847656, + "speedup_pct": 9.69609062818353, + "max_abs": 0.015625, + "relative_l2": 0.00016855482573421412, + "argmax_equal": true + }, + { + "suite": "kernel", + "label": "gpt_neox_down", + "m": 2, + "n": 4096, + "k": 16384, + "torch_median_ms": 0.19429887771606447, + "torch_p90_ms": 0.1946112060546875, + "candidate_median_ms": 0.1790924835205078, + "candidate_p90_ms": 0.1792153549194336, + "speedup_pct": 8.49080536303768, + "max_abs": 0.03125, + "relative_l2": 0.00038089508822365355, + "argmax_equal": true + }, + { + "suite": "kernel", + "label": "gpt_neox_down", + "m": 4, + "n": 4096, + "k": 16384, + "torch_median_ms": 0.1944780731201172, + "torch_p90_ms": 0.19456480026245118, + "candidate_median_ms": 0.17938943862915038, + "candidate_p90_ms": 0.17957376480102538, + "speedup_pct": 8.411105250270268, + "max_abs": 0.03125, + "relative_l2": 0.00031749888615750923, + "argmax_equal": true + } + ] +} diff --git a/docs/benchmarks/gemv_common_main_baseline_l20_sm89.csv b/docs/benchmarks/gemv_common_main_baseline_l20_sm89.csv new file mode 100644 index 00000000..7475deb8 --- /dev/null +++ b/docs/benchmarks/gemv_common_main_baseline_l20_sm89.csv @@ -0,0 +1,24 @@ +# device=NVIDIA L20, capability=8.9, seed=20260902, weight_std=0.02 +suite,label,m,n,k,torch_median_ms,torch_p90_ms,candidate_median_ms,candidate_p90_ms,speedup_pct,max_abs,relative_l2,argmax_equal +kernel,llama2_7b_qo,2,4096,4096,0.021094,0.021391,0.012349,0.012442,+70.81,0.015625,0.00024786,true +kernel,llama2_7b_qo,4,4096,4096,0.021084,0.021146,0.016292,0.016323,+29.42,0.007812,0.00007721,true +kernel,llama2_7b_up_gate,2,11008,4096,0.039813,0.039875,0.027453,0.027525,+45.02,0.007812,0.00009190,true +kernel,llama2_7b_up_gate,4,11008,4096,0.040294,0.041339,0.037939,0.038881,+6.21,0.015625,0.00018091,true +kernel,llama2_7b_down,2,4096,11008,0.051825,0.053258,0.030085,0.030454,+72.26,0.015625,0.00017434,true +kernel,llama2_7b_down,4,4096,11008,0.051999,0.052183,0.043407,0.043674,+19.79,0.031250,0.00034518,true +kernel,llama3_8b_kv,2,1024,4096,0.014305,0.014449,0.006492,0.006554,+120.35,0.015625,0.00263309,true +kernel,llama3_8b_kv,4,1024,4096,0.014377,0.014541,0.006697,0.006820,+114.68,0.031250,0.00256016,true +kernel,llama3_8b_up_gate,2,14336,4096,0.165478,0.166062,0.153508,0.153979,+7.80,0.015625,0.00012831,true +kernel,llama3_8b_up_gate,4,14336,4096,0.166369,0.166584,0.154132,0.154491,+7.94,0.015625,0.00016471,true +kernel,llama3_8b_down,2,4096,14336,0.165939,0.166687,0.156600,0.157420,+5.96,0.031250,0.00027259,true +kernel,llama3_8b_down,4,4096,14336,0.165949,0.166820,0.156897,0.157614,+5.77,0.031250,0.00037241,true +kernel,llama2_13b_qo,2,5120,5120,0.026778,0.026900,0.017306,0.017357,+54.73,0.015625,0.00014752,true +kernel,llama2_13b_qo,4,5120,5120,0.026972,0.027034,0.021821,0.021924,+23.60,0.015625,0.00012898,true +kernel,llama2_13b_up_gate,2,13824,5120,0.198984,0.199086,0.188867,0.188989,+5.36,0.015625,0.00012897,true +kernel,llama2_13b_up_gate,4,13824,5120,0.199526,0.201226,0.189317,0.189737,+5.39,0.031250,0.00022245,true +kernel,llama2_13b_down,2,5120,13824,0.199967,0.200264,0.188887,0.189327,+5.87,0.031250,0.00033799,true +kernel,llama2_13b_down,4,5120,13824,0.200346,0.202015,0.189563,0.190177,+5.69,0.031250,0.00029381,true +kernel,gpt_neox_up,2,16384,4096,0.192952,0.193126,0.175462,0.175524,+9.97,0.015625,0.00017490,true +kernel,gpt_neox_up,4,16384,4096,0.193167,0.194109,0.176077,0.176252,+9.71,0.015625,0.00016855,true +kernel,gpt_neox_down,2,4096,16384,0.194191,0.194468,0.178821,0.179374,+8.60,0.031250,0.00038090,true +kernel,gpt_neox_down,4,4096,16384,0.194662,0.196485,0.179374,0.179692,+8.52,0.031250,0.00031750,true diff --git a/docs/benchmarks/gemv_common_warp_tiled_l20_sm89.csv b/docs/benchmarks/gemv_common_warp_tiled_l20_sm89.csv new file mode 100644 index 00000000..37d1e6ac --- /dev/null +++ b/docs/benchmarks/gemv_common_warp_tiled_l20_sm89.csv @@ -0,0 +1,24 @@ +# device=NVIDIA L20, capability=8.9, seed=20260902, weight_std=0.02 +suite,label,m,n,k,torch_median_ms,torch_p90_ms,candidate_median_ms,candidate_p90_ms,speedup_pct,max_abs,relative_l2,argmax_equal +kernel,llama2_7b_qo,2,4096,4096,0.021064,0.021187,0.012349,0.012411,+70.56,0.015625,0.00024786,true +kernel,llama2_7b_qo,4,4096,4096,0.021043,0.021146,0.016108,0.016138,+30.64,0.007812,0.00007350,true +kernel,llama2_7b_up_gate,2,11008,4096,0.039782,0.039885,0.027423,0.027474,+45.07,0.007812,0.00009190,true +kernel,llama2_7b_up_gate,4,11008,4096,0.040141,0.040376,0.028539,0.028723,+40.65,0.015625,0.00018091,true +kernel,llama2_7b_down,2,4096,11008,0.052019,0.052142,0.032614,0.032840,+59.50,0.015625,0.00017442,true +kernel,llama2_7b_down,4,4096,11008,0.052163,0.052224,0.039803,0.039997,+31.05,0.031250,0.00034511,true +kernel,llama3_8b_kv,2,1024,4096,0.014100,0.014172,0.006390,0.006513,+120.67,0.015625,0.00263309,true +kernel,llama3_8b_kv,4,1024,4096,0.014141,0.014510,0.006390,0.006431,+121.31,0.031250,0.00255970,true +kernel,llama3_8b_up_gate,2,14336,4096,0.165550,0.166083,0.153477,0.153989,+7.87,0.015625,0.00012831,true +kernel,llama3_8b_up_gate,4,14336,4096,0.166523,0.168100,0.154450,0.154767,+7.82,0.015625,0.00016471,true +kernel,llama3_8b_down,2,4096,14336,0.166205,0.166799,0.156641,0.157399,+6.11,0.031250,0.00027259,true +kernel,llama3_8b_down,4,4096,14336,0.167199,0.167393,0.157102,0.158034,+6.43,0.031250,0.00037241,true +kernel,llama2_13b_qo,2,5120,5120,0.026829,0.026972,0.017357,0.017459,+54.57,0.015625,0.00014752,true +kernel,llama2_13b_qo,4,5120,5120,0.026952,0.027023,0.021637,0.021668,+24.56,0.015625,0.00012898,true +kernel,llama2_13b_up_gate,2,13824,5120,0.199014,0.199209,0.188856,0.188979,+5.38,0.015625,0.00012897,true +kernel,llama2_13b_up_gate,4,13824,5120,0.199506,0.199823,0.189307,0.189358,+5.39,0.031250,0.00022245,true +kernel,llama2_13b_down,2,5120,13824,0.199844,0.200274,0.188826,0.189358,+5.84,0.031250,0.00033799,true +kernel,llama2_13b_down,4,5120,13824,0.200397,0.202629,0.189307,0.190259,+5.86,0.031250,0.00029381,true +kernel,gpt_neox_up,2,16384,4096,0.192819,0.193208,0.175401,0.175514,+9.93,0.015625,0.00017490,true +kernel,gpt_neox_up,4,16384,4096,0.193259,0.193352,0.176108,0.176220,+9.74,0.015625,0.00016855,true +kernel,gpt_neox_down,2,4096,16384,0.194540,0.194714,0.179026,0.179292,+8.67,0.031250,0.00038090,true +kernel,gpt_neox_down,4,4096,16384,0.194253,0.194867,0.179108,0.179661,+8.46,0.031250,0.00031750,true diff --git a/docs/developer/cuda_kernels.md b/docs/developer/cuda_kernels.md index b2dbd20c..48a1fc80 100644 --- a/docs/developer/cuda_kernels.md +++ b/docs/developer/cuda_kernels.md @@ -21,29 +21,94 @@ model linear dispatcher described below. `astrai.extension.bf16_gemv(x, weight, bias=None)` accepts a contiguous BF16 input shaped `[K]` or `[M, K]`, with `M` in `[1, 8]` and any positive `K`, and -row-major weights `[N, K]`. One CTA reduces each output row and computes all M -results together, reusing the weight row across tokens. The weight stream uses -128-bit vectorized loads anchored at each row's first 16-byte-aligned address -with scalar head/tail sweeps for unaligned remainders, so arbitrary `K` and -storage offsets stay correct; x loads are vectorized when every row base is -16-byte aligned (always true for K % 8 == 0 with allocator-aligned tensors) -and scalar otherwise. Accumulation is FP32; the optional BF16 bias is fused -before the BF16 store. The launcher uses the current CUDA stream, is CUDA -Graph capture-safe, and requires sm_80 or newer. +row-major weights `[N, K]`. The general path assigns one 256-thread CTA to an +output row and computes all M results together, reusing the weight row across +tokens. For measured aligned M=4 medium projections, a 128-thread CTA instead +assigns one output to each of four warps. That removes the CTA-wide reduction +barrier and exposes four neighboring outputs without changing accumulation. + +The weight stream uses 128-bit vectorized loads anchored at each row's first +16-byte-aligned address with scalar head/tail sweeps for unaligned remainders, +so arbitrary `K` and storage offsets stay correct. The warp-tiled path is used +only when both tensors and every row are 16-byte aligned; all other calls keep +the general arbitrary-K path. Accumulation is FP32; the optional BF16 bias is +fused before the BF16 store. The launcher uses the current CUDA stream, is +CUDA Graph capture-safe, and requires sm_80 or newer. Model `Linear` calls route through the lightweight linear backend. Set `ASTRAI_GEMV=0` for an unconditional `F.linear` fallback, `1` to force the kernel for any supported M in [1, 8], or `auto` (the default) to select only architecture/shape bands that pass both the per-shape and end-to-end gates. -M=1 has no automatic SM89 band because isolated winners did not reach the 3% -whole-graph gate. Measured SM89 small-M bands are enabled as follows: +Measured SM89 small-M bands are enabled as follows: -| M | Automatic `(N, K)` bands | Engine throughput | +| M | Automatic `(N, K)` bands | Validated gain | |---:|---|---:| -| 2 | `(256,1536)`, `(1536,1536)`, `(100000,1536)` | +14.0% | -| 4 | `(256,1536)`, `(1536,1536)` | +11.8% | - -These A→B→B→A results use the real `InferenceEngine`, including scheduler, +| 1 | OPT-1.3B Q/K/V/O and MLP | +4.54% OPT projection chain | +| 2 | AstrAI `(256,1536)`, `(1536,1536)`, `(100000,1536)` plus all common shapes below | +14.0% on AstrAI 1B; +5.66% to +25.20% common chains | +| 4 | AstrAI `(256,1536)`, `(1536,1536)` plus gated common shapes below | +11.8% on AstrAI 1B; +5.67% to +7.71% common chains | +| 8 | none | at least one projection in every measured family missed the per-shape gate | + +The common set covers LLaMA 2 7B Q/O, gate/up, and down; LLaMA 3 8B K/V, +gate/up, and down; LLaMA 2 13B Q/K/V/O, gate/up, and down; and GPT-NeoX MLP +up/down. In `(N,K)` form it is `(1024,4096)`, `(4096,4096)`, +`(11008,4096)`, `(4096,11008)`, `(14336,4096)`, `(4096,14336)`, +`(5120,5120)`, `(13824,5120)`, `(5120,13824)`, `(16384,4096)`, and +`(4096,16384)`. M=2 enables all eleven. M=4 excludes the three LLaMA 2 7B +bands `(4096,4096)`, `(11008,4096)`, and `(4096,11008)` because their combined +projection chain reached only +1.89%, below the 3% automatic-dispatch gate. + +The extended common set adds Qwen2-7B `(512,3584)`, `(3584,3584)`, +`(18944,3584)`, and `(3584,18944)`; LLaMA 3 70B `(1024,8192)`, +`(8192,8192)`, `(28672,8192)`, and `(8192,28672)`; and OPT-1.3B +`(2048,2048)`, `(8192,2048)`, and `(2048,8192)`. Qwen2 and LLaMA 3 70B are +enabled at M=2/4. OPT-1.3B is enabled at M=1/2. Other rows retain their +previous policy or fall back to PyTorch. + +Inside the primitive, a templated cooperative kernel uses either 256 threads +or a shape-gated 128-thread CTA. The smaller CTA is enabled only where an +interleaved direct-module comparison against the original 256-thread kernel +cleared 5%: OPT up at M=1; selected LLaMA 2 7B, Qwen2, and OPT projections at +M=2; LLaMA 2 13B Q/O, Qwen2 Q/O, and selected OPT projections at M=4; and +selected LLaMA 2, Qwen2, LLaMA 3 KV, and OPT projections at M=8. Confirmed +direct-kernel gains range from +5.37% to +48.54%. Long-K and saturated shapes +keep the 256-thread fallback. This internal selector is separate from model +automatic dispatch, whose Python/wrapper overhead is included in the gates +above. + +On NVIDIA L20 (SM89), the common-shape microbenchmark reports +5.37% to ++114.39% for M=2 and +5.38% to +115.26% for M=4 versus `F.linear`. The paired +main-versus-warp-tiling run used identical interleaved settings; for the four +M=4 selected shapes, candidate latency changed from 0.016292 to 0.016108 ms +for `(4096,4096)`, 0.037939 to 0.028539 ms for `(11008,4096)`, 0.043407 to +0.039803 ms for `(4096,11008)`, and 0.006697 to 0.006390 ms for +`(1024,4096)`. + +The dependent projection-chain gate, which includes Python dispatch and +rotates through distinct weights instead of repeatedly warming one matrix, +measured: + +| Synthetic chain | M=2 | M=4 | Row argmax parity | +|---|---:|---:|---| +| LLaMA 2 7B | +8.49% | fallback (M=4 bands excluded) | exact | +| LLaMA 3 8B | +8.50% | +6.44% | exact | +| LLaMA 2 13B | +5.66% | +5.67% | exact | +| GPT-NeoX 20B | +6.95% | +5.93% | exact | +| Qwen2 7B | +7.48% | +7.48% | exact | +| LLaMA 3 70B | +7.77% | +7.69% | exact | +| OPT 1.3B | +25.20% | fallback (M=4 up projection regresses) | exact | + +OPT 1.3B M=1 is +4.54%. Qwen2 and LLaMA 3 70B M=1, and all three new +families at M=8, remain exact PyTorch fallbacks. + +These are synthetic projection-chain measurements, not whole-model throughput +claims. Reproduce them with `scripts/tools/benchmark_gemv_common.py`; the raw +L20 environment, parameters, timings, and numerical diagnostics live in +`docs/benchmarks/gemv_common_{kernel,chain}_l20_sm89.json`. Extended-shape and +CTA-selection evidence is in +`gemv_common_extended_{kernel,chain}_l20_sm89.json` and +`gemv_common_halfcta_vs_256cta_l20_sm89.json` in the same directory. + +The AstrAI 1B A→B→B→A results use the real `InferenceEngine`, including scheduler, sampling, and CUDA Graph. M=8 stays on PyTorch because its remaining greedy-stable winners missed the 3% end-to-end gate. Long-K MLP-down bands are also excluded because their valid BF16 error changed a checkpoint greedy diff --git a/docs/developer/decode_linear_benchmark.md b/docs/developer/decode_linear_benchmark.md index 02bf113c..92494eec 100644 --- a/docs/developer/decode_linear_benchmark.md +++ b/docs/developer/decode_linear_benchmark.md @@ -20,3 +20,23 @@ change the decode batch sizes. Compare each GPU architecture only with its own baseline; do not use absolute A100-versus-L20 numbers as a dispatch criterion. Keep the raw JSON as the source of truth and generate tables with `--markdown-output` rather than transcribing measurements by hand. + +For direct A/B coverage of the custom kernel and guarded dispatcher across +traditional LLaMA and GPT-NeoX decode shapes, use: + +```bash +CUDA_VISIBLE_DEVICES=0 PYTHONPATH=. python scripts/tools/benchmark_gemv_common.py \ + --suite all --family traditional --m 2 4 \ + --output results/gemv_common.json +``` + +The kernel suite compares the directly callable primitive with `F.linear`. +Use repeatable `--shape-label` and `--chain-label` filters for a focused run. +The synthetic-chain suite alternates `ASTRAI_GEMV=0` and `auto`, includes +dependent MLP work and Python dispatch, and rotates through distinct weights. +Pass `--candidate-mode 1` to characterize a family before adding it to the +automatic shape table; the checked-in final evidence always uses `auto`. +It is deliberately not labeled a whole-model throughput benchmark. Both +suites report median/p90 CUDA-event latency plus maximum absolute error, +relative L2 error, and row-wise argmax parity. The checked-in L20 evidence is +under `docs/benchmarks/`. diff --git a/scripts/tools/benchmark_gemv_common.py b/scripts/tools/benchmark_gemv_common.py new file mode 100644 index 00000000..33b8f0c9 --- /dev/null +++ b/scripts/tools/benchmark_gemv_common.py @@ -0,0 +1,450 @@ +"""Benchmark the BF16 GEMV primitive and guarded linear dispatcher. + +The kernel suite covers AstrAI's native projections plus common LLaMA and +GPT-NeoX matrix shapes. The chain suite is a synthetic projection/MLP chain; +it measures dispatcher overhead and dependent MLP work, but is deliberately +not presented as a whole-model throughput benchmark. +""" + +import argparse +import gc +import json +import math +import os +import statistics +from collections.abc import Callable +from dataclasses import dataclass +from pathlib import Path + +import torch +import torch.nn.functional as F + +from astrai.extension import bf16_gemv, is_available, linear + + +@dataclass(frozen=True) +class Shape: + label: str + n: int + k: int + + +@dataclass(frozen=True) +class Chain: + label: str + hidden: int + kv: int + intermediate: int + fused_qkv: bool = False + gated_mlp: bool = True + + +@dataclass(frozen=True) +class Timing: + median_ms: float + p90_ms: float + + +ASTRAI_SHAPES = ( + Shape("astrai_qkv", 256, 1536), + Shape("astrai_square", 1536, 1536), + Shape("astrai_up_gate", 6912, 1536), + Shape("astrai_down", 1536, 6912), + Shape("astrai_lm_head", 100000, 1536), +) + +TRADITIONAL_SHAPES = ( + Shape("llama2_7b_qo", 4096, 4096), + Shape("llama2_7b_up_gate", 11008, 4096), + Shape("llama2_7b_down", 4096, 11008), + Shape("llama3_8b_kv", 1024, 4096), + Shape("llama3_8b_up_gate", 14336, 4096), + Shape("llama3_8b_down", 4096, 14336), + Shape("llama2_13b_qo", 5120, 5120), + Shape("llama2_13b_up_gate", 13824, 5120), + Shape("llama2_13b_down", 5120, 13824), + Shape("gpt_neox_up", 16384, 4096), + Shape("gpt_neox_down", 4096, 16384), + Shape("qwen2_7b_kv", 512, 3584), + Shape("qwen2_7b_qo", 3584, 3584), + Shape("qwen2_7b_up_gate", 18944, 3584), + Shape("qwen2_7b_down", 3584, 18944), + Shape("llama3_70b_kv", 1024, 8192), + Shape("llama3_70b_qo", 8192, 8192), + Shape("llama3_70b_up_gate", 28672, 8192), + Shape("llama3_70b_down", 8192, 28672), + Shape("opt_1_3b_qkvo", 2048, 2048), + Shape("opt_1_3b_up", 8192, 2048), + Shape("opt_1_3b_down", 2048, 8192), +) + +CHAINS = ( + Chain("llama2_7b", 4096, 4096, 11008), + Chain("llama3_8b", 4096, 1024, 14336), + Chain("llama2_13b", 5120, 5120, 13824), + Chain("gpt_neox_20b", 4096, 4096, 16384, fused_qkv=True), + Chain("qwen2_7b", 3584, 512, 18944), + Chain("llama3_70b", 8192, 1024, 28672), + Chain("opt_1_3b", 2048, 2048, 8192, gated_mlp=False), +) + + +def _elapsed_ms(fn: Callable[[], torch.Tensor], inner: int) -> float: + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(inner): + fn() + end.record() + end.synchronize() + return start.elapsed_time(end) / inner + + +def _timing(values: list[float]) -> Timing: + ordered = sorted(values) + p90_index = max(0, math.ceil(0.9 * len(ordered)) - 1) + return Timing(statistics.median(ordered), ordered[p90_index]) + + +def _measure_pair( + baseline: Callable[[], torch.Tensor], + candidate: Callable[[], torch.Tensor], + *, + warmup: int, + samples: int, + inner: int, + prepare_baseline: Callable[[], None] = lambda: None, + prepare_candidate: Callable[[], None] = lambda: None, +) -> tuple[Timing, Timing]: + cases = ( + ("baseline", prepare_baseline, baseline), + ("candidate", prepare_candidate, candidate), + ) + for iteration in range(warmup): + _, prepare, fn = cases[iteration % 2] + prepare() + fn() + torch.cuda.synchronize() + + values: dict[str, list[float]] = {"baseline": [], "candidate": []} + for sample in range(samples): + order = cases if sample % 2 == 0 else tuple(reversed(cases)) + for label, prepare, fn in order: + prepare() + values[label].append(_elapsed_ms(fn, inner)) + return _timing(values["baseline"]), _timing(values["candidate"]) + + +def _print_header() -> None: + print( + "suite,label,m,n,k,torch_median_ms,torch_p90_ms," + "candidate_median_ms,candidate_p90_ms,speedup_pct," + "max_abs,relative_l2,argmax_equal" + ) + + +def _print_result( + suite: str, + label: str, + m: int, + n: int, + k: int, + baseline: Timing, + candidate: Timing, + reference: torch.Tensor, + actual: torch.Tensor, +) -> dict[str, object]: + difference = actual.float() - reference.float() + max_abs = difference.abs().max().item() + relative_l2 = difference.norm().item() / max(reference.float().norm().item(), 1e-12) + argmax_equal = torch.equal(actual.argmax(dim=-1), reference.argmax(dim=-1)) + speedup = (baseline.median_ms / candidate.median_ms - 1.0) * 100.0 + result: dict[str, object] = { + "suite": suite, + "label": label, + "m": m, + "n": n, + "k": k, + "torch_median_ms": baseline.median_ms, + "torch_p90_ms": baseline.p90_ms, + "candidate_median_ms": candidate.median_ms, + "candidate_p90_ms": candidate.p90_ms, + "speedup_pct": speedup, + "max_abs": max_abs, + "relative_l2": relative_l2, + "argmax_equal": argmax_equal, + } + print( + f"{suite},{label},{m},{n},{k}," + f"{baseline.median_ms:.6f},{baseline.p90_ms:.6f}," + f"{candidate.median_ms:.6f},{candidate.p90_ms:.6f}," + f"{speedup:+.2f},{max_abs:.6f},{relative_l2:.8f}," + f"{str(argmax_equal).lower()}", + flush=True, + ) + return result + + +def _weight(n: int, k: int, device: torch.device, std: float) -> torch.Tensor: + weight = torch.empty((n, k), device=device, dtype=torch.bfloat16) + weight.normal_(mean=0.0, std=std) + return weight.requires_grad_(True) + + +def _kernel_functions( + x: torch.Tensor, weight: torch.Tensor +) -> tuple[Callable[[], torch.Tensor], Callable[[], torch.Tensor]]: + def baseline() -> torch.Tensor: + return F.linear(x, weight) + + def candidate() -> torch.Tensor: + return bf16_gemv(x, weight.detach()) + + return baseline, candidate + + +def benchmark_kernels( + args: argparse.Namespace, device: torch.device +) -> list[dict[str, object]]: + if args.family == "astrai": + shapes = ASTRAI_SHAPES + elif args.family == "traditional": + shapes = TRADITIONAL_SHAPES + else: + shapes = ASTRAI_SHAPES + TRADITIONAL_SHAPES + if args.shape_label: + requested = set(args.shape_label) + shapes = tuple(shape for shape in shapes if shape.label in requested) + missing = requested - {shape.label for shape in shapes} + if missing: + raise ValueError(f"unknown shape labels: {', '.join(sorted(missing))}") + + results: list[dict[str, object]] = [] + for shape in shapes: + weight = _weight(shape.n, shape.k, device, args.weight_std) + for m in args.m: + x = torch.randn((m, shape.k), device=device, dtype=torch.bfloat16) + baseline_fn, candidate_fn = _kernel_functions(x, weight) + with torch.inference_mode(): + reference = baseline_fn() + actual = candidate_fn() + baseline, candidate = _measure_pair( + baseline_fn, + candidate_fn, + warmup=args.warmup, + samples=args.samples, + inner=args.inner, + ) + results.append( + _print_result( + "kernel", + shape.label, + m, + shape.n, + shape.k, + baseline, + candidate, + reference, + actual, + ) + ) + del baseline_fn, candidate_fn, x, reference, actual + del weight + gc.collect() + torch.cuda.empty_cache() + return results + + +def _set_mode(mode: str) -> None: + os.environ["ASTRAI_GEMV"] = mode + + +def _chain_weights( + spec: Chain, device: torch.device, std: float +) -> dict[str, torch.Tensor]: + weights = { + "o": _weight(spec.hidden, spec.hidden, device, std), + "up": _weight(spec.intermediate, spec.hidden, device, std), + "down": _weight(spec.hidden, spec.intermediate, device, std), + } + if spec.fused_qkv: + weights["qkv"] = _weight(3 * spec.hidden, spec.hidden, device, std) + else: + weights.update( + { + "q": _weight(spec.hidden, spec.hidden, device, std), + "k": _weight(spec.kv, spec.hidden, device, std), + "v": _weight(spec.kv, spec.hidden, device, std), + } + ) + if spec.gated_mlp: + weights["gate"] = _weight(spec.intermediate, spec.hidden, device, std) + return weights + + +def _chain_fn( + x: torch.Tensor, weights: dict[str, torch.Tensor], spec: Chain +) -> Callable[[], torch.Tensor]: + def run() -> torch.Tensor: + output_projection = linear(x, weights["o"]) + up = linear(x, weights["up"]) + if spec.fused_qkv: + attention_projection = linear(x, weights["qkv"])[..., : x.shape[-1]] + hidden = F.gelu(up) + else: + attention_projection = linear(x, weights["q"]) + linear(x, weights["k"]) + linear(x, weights["v"]) + if spec.gated_mlp: + gate = linear(x, weights["gate"]) + hidden = F.silu(gate) * up + else: + hidden = F.gelu(up) + down = linear(hidden, weights["down"]) + return attention_projection + output_projection + down + + return run + + +def benchmark_chains( + args: argparse.Namespace, device: torch.device +) -> list[dict[str, object]]: + results: list[dict[str, object]] = [] + chains = CHAINS + if args.chain_label: + requested = set(args.chain_label) + chains = tuple(chain for chain in chains if chain.label in requested) + missing = requested - {chain.label for chain in chains} + if missing: + raise ValueError(f"unknown chain labels: {', '.join(sorted(missing))}") + for spec in chains: + weights = _chain_weights(spec, device, args.weight_std) + for m in args.m: + x = torch.randn((m, spec.hidden), device=device, dtype=torch.bfloat16) + run = _chain_fn(x, weights, spec) + with torch.inference_mode(): + _set_mode("0") + reference = run() + _set_mode(args.candidate_mode) + actual = run() + baseline, candidate = _measure_pair( + run, + run, + warmup=args.warmup, + samples=args.samples, + inner=args.chain_inner, + prepare_baseline=lambda: _set_mode("0"), + prepare_candidate=lambda: _set_mode(args.candidate_mode), + ) + results.append( + _print_result( + "synthetic_chain", + spec.label, + m, + spec.hidden, + spec.intermediate, + baseline, + candidate, + reference, + actual, + ) + ) + del x, reference, actual + del weights + gc.collect() + torch.cuda.empty_cache() + return results + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--suite", choices=("kernel", "chain", "all"), default="all") + parser.add_argument( + "--family", choices=("astrai", "traditional", "all"), default="all" + ) + parser.add_argument( + "--m", type=int, nargs="+", choices=(1, 2, 4, 8), default=(1, 2, 4, 8) + ) + parser.add_argument( + "--shape-label", + action="append", + help="limit the kernel suite to one or more named shape labels", + ) + parser.add_argument( + "--chain-label", + action="append", + help="limit the chain suite to one or more named model families", + ) + parser.add_argument("--device", type=int, default=0) + parser.add_argument("--warmup", type=int, default=20) + parser.add_argument("--samples", type=int, default=9) + parser.add_argument("--inner", type=int, default=100) + parser.add_argument("--chain-inner", type=int, default=20) + parser.add_argument( + "--candidate-mode", + choices=("auto", "1"), + default="auto", + help="dispatcher mode for the candidate side of the chain suite", + ) + parser.add_argument("--weight-std", type=float, default=0.02) + parser.add_argument("--seed", type=int, default=20260902) + parser.add_argument( + "--output", + type=Path, + help="optional JSON output; stdout always retains the compact CSV table", + ) + return parser.parse_args() + + +def main() -> None: + args = parse_args() + if not torch.cuda.is_available() or not is_available("bf16_gemv"): + raise RuntimeError("benchmark requires CUDA and the built bf16_gemv extension") + if args.warmup < 0 or args.samples < 1 or args.inner < 1 or args.chain_inner < 1: + raise ValueError("warmup must be non-negative and sample/inner counts positive") + + torch.cuda.set_device(args.device) + device = torch.device("cuda", args.device) + torch.manual_seed(args.seed) + torch.cuda.manual_seed_all(args.seed) + properties = torch.cuda.get_device_properties(device) + print( + f"# device={properties.name}, capability={properties.major}.{properties.minor}, " + f"seed={args.seed}, weight_std={args.weight_std}" + ) + _print_header() + results: list[dict[str, object]] = [] + if args.suite in ("kernel", "all"): + results.extend(benchmark_kernels(args, device)) + if args.suite in ("chain", "all"): + results.extend(benchmark_chains(args, device)) + if args.output is not None: + payload = { + "environment": { + "device": properties.name, + "capability": f"{properties.major}.{properties.minor}", + "torch": torch.__version__, + "cuda": torch.version.cuda, + }, + "parameters": { + "suite": args.suite, + "family": args.family, + "m": args.m, + "shape_labels": args.shape_label, + "chain_labels": args.chain_label, + "candidate_mode": args.candidate_mode, + "seed": args.seed, + "weight_std": args.weight_std, + "warmup": args.warmup, + "samples": args.samples, + "inner": args.inner, + "chain_inner": args.chain_inner, + }, + "results": results, + } + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(payload, indent=2) + "\n") + + +if __name__ == "__main__": + main() diff --git a/tests/extension/test_gemv.py b/tests/extension/test_gemv.py index c49bf67e..1ebecc41 100644 --- a/tests/extension/test_gemv.py +++ b/tests/extension/test_gemv.py @@ -43,6 +43,66 @@ def test_bf16_gemv_matches_small_decode_batches(m, n, k): torch.testing.assert_close(actual, expected, rtol=0.02, atol=0.5) +@skip_no_gemv +@pytest.mark.parametrize("m", [2, 4]) +@pytest.mark.parametrize( + "n,k", + [ + (1024, 4096), + (4096, 4096), + (11008, 4096), + (4096, 11008), + (14336, 4096), + (4096, 14336), + (5120, 5120), + (13824, 5120), + (5120, 13824), + (16384, 4096), + (4096, 16384), + (512, 3584), + (3584, 3584), + (18944, 3584), + (3584, 18944), + (1024, 8192), + (8192, 8192), + (28672, 8192), + (8192, 28672), + (2048, 2048), + (8192, 2048), + (2048, 8192), + ], +) +def test_bf16_gemv_matches_common_transformer_shapes(m, n, k): + torch.manual_seed(2026 + m + n + k) + x = torch.randn(m, k, device="cuda", dtype=torch.bfloat16) + weight = torch.empty(n, k, device="cuda", dtype=torch.bfloat16) + weight.normal_(mean=0.0, std=0.02) + actual = bf16_gemv(x, weight) + expected = F.linear(x, weight) + torch.testing.assert_close(actual, expected, rtol=0.02, atol=0.25) + + +@skip_no_gemv +@pytest.mark.parametrize( + "m,n,k", + [ + (1, 8192, 2048), + (8, 4096, 11008), + (8, 512, 3584), + (8, 1024, 8192), + (8, 2048, 8192), + ], +) +def test_bf16_gemv_matches_half_cta_edge_bands(m, n, k): + torch.manual_seed(2026 + m + n + k) + x = torch.randn(m, k, device="cuda", dtype=torch.bfloat16) + weight = torch.empty(n, k, device="cuda", dtype=torch.bfloat16) + weight.normal_(mean=0.0, std=0.02) + actual = bf16_gemv(x, weight) + expected = F.linear(x, weight) + torch.testing.assert_close(actual, expected, rtol=0.02, atol=0.25) + + @skip_no_gemv def test_bf16_gemv_preserves_singleton_batch_and_fuses_bias(): torch.manual_seed(23) diff --git a/tests/extension/test_linear_dispatch.py b/tests/extension/test_linear_dispatch.py index ae9cf017..d23fa927 100644 --- a/tests/extension/test_linear_dispatch.py +++ b/tests/extension/test_linear_dispatch.py @@ -6,6 +6,7 @@ from astrai.extension import explain, is_available, linear, op_backend from astrai.extension.backend import linear as public_linear +from astrai.extension.backend.linear import _AUTO_GEMV_SHAPES from astrai.model.components.linear import Linear GEMV_AVAILABLE = ( @@ -23,6 +24,45 @@ def test_linear_backend_is_public(): assert linear is public_linear +def test_sm89_common_shape_policy_keeps_only_validated_families_enabled(): + common = { + (1024, 4096), + (4096, 4096), + (11008, 4096), + (4096, 11008), + (14336, 4096), + (4096, 14336), + (5120, 5120), + (13824, 5120), + (5120, 13824), + (16384, 4096), + (4096, 16384), + } + subthreshold_m4 = {(4096, 4096), (11008, 4096), (4096, 11008)} + qwen2_7b = { + (512, 3584), + (3584, 3584), + (18944, 3584), + (3584, 18944), + } + llama3_70b = { + (1024, 8192), + (8192, 8192), + (28672, 8192), + (8192, 28672), + } + opt_1_3b = {(2048, 2048), (8192, 2048), (2048, 8192)} + policy = _AUTO_GEMV_SHAPES[(8, 9)] + assert policy[1] == opt_1_3b + assert common <= policy[2] + assert qwen2_7b | llama3_70b | opt_1_3b <= policy[2] + assert common - subthreshold_m4 <= policy[4] + assert qwen2_7b | llama3_70b <= policy[4] + assert subthreshold_m4.isdisjoint(policy[4]) + assert opt_1_3b.isdisjoint(policy[4]) + assert 8 not in policy + + def test_model_linear_routes_through_backend(monkeypatch): sentinel = torch.randn(2, 4) @@ -93,7 +133,7 @@ def test_mode_one_dispatches_supported_small_batches(monkeypatch, m): @skip_no_gemv -def test_auto_m1_falls_back_until_end_to_end_gate_passes(monkeypatch): +def test_auto_unmeasured_m1_falls_back(monkeypatch): monkeypatch.setenv("ASTRAI_GEMV", "auto") x = torch.randn(1, 1536, device="cuda", dtype=torch.bfloat16) winning = torch.randn( @@ -109,10 +149,38 @@ def test_auto_m1_falls_back_until_end_to_end_gate_passes(monkeypatch): @skip_no_gemv -def test_auto_selects_measured_sm89_small_batch_winner(monkeypatch): +@pytest.mark.parametrize( + "m,n,k", + [ + (4, 256, 1536), + (2, 1024, 4096), + (2, 11008, 4096), + (2, 4096, 11008), + (2, 14336, 4096), + (4, 4096, 14336), + (2, 5120, 5120), + (4, 13824, 5120), + (2, 5120, 13824), + (4, 16384, 4096), + (2, 4096, 16384), + (2, 512, 3584), + (4, 3584, 3584), + (2, 18944, 3584), + (4, 3584, 18944), + (2, 1024, 8192), + (4, 8192, 8192), + (2, 28672, 8192), + (4, 8192, 28672), + (1, 2048, 2048), + (2, 8192, 2048), + (1, 2048, 8192), + ], +) +def test_auto_selects_measured_sm89_small_batch_winner(monkeypatch, m, n, k): monkeypatch.setenv("ASTRAI_GEMV", "auto") - x = torch.randn(4, 1536, device="cuda", dtype=torch.bfloat16) - weight = torch.randn(256, 1536, device="cuda", dtype=torch.bfloat16) + x = torch.randn(m, k, device="cuda", dtype=torch.bfloat16) + weight = torch.empty(n, k, device="cuda", dtype=torch.bfloat16) + weight.normal_(mean=0.0, std=0.02) with torch.no_grad(): trace = explain("linear", x, weight) if torch.cuda.get_device_capability() == (8, 9): @@ -133,6 +201,18 @@ def test_auto_selects_measured_sm89_small_batch_winner(monkeypatch): (4, 100000, 1536), # LM head misses the 5% M=4 gate (4, 1536, 6912), # long-K accumulation changed checkpoint greedy output (8, 256, 1536), # remaining M=8 winners miss the 3% end-to-end gate + (1, 4096, 4096), # isolated M=1 winner misses the projection-chain gate + (8, 1024, 4096), # isolated M=8 winner misses the projection-chain gate + (4, 12288, 4096), # GPT-NeoX fused QKV was not measured as a winner + (4, 4096, 4096), # LLaMA 2 7B M=4 chain misses the 3% gate + (4, 11008, 4096), + (4, 4096, 11008), + (1, 3584, 3584), # Qwen2 M=1 chain misses the 3% gate + (8, 3584, 3584), # Qwen2 Q/O misses the M=8 per-shape gate + (1, 8192, 8192), # LLaMA 3 70B M=1 projections miss the per-shape gate + (8, 1024, 8192), # LLaMA 3 70B K/V loses at wrapper level for M=8 + (4, 8192, 2048), # OPT up loses at wrapper level for M=4 + (8, 2048, 2048), # OPT M=8 chain and Q/K/V/O both regress ], ) def test_auto_rejects_measured_small_batch_losers(monkeypatch, m, n, k):