From 416f4d0fb8ecd5d9b601e496803c346c7c31f40b Mon Sep 17 00:00:00 2001 From: 0z5a Date: Wed, 2 Sep 2026 18:47:45 +0800 Subject: [PATCH 1/2] perf: tune bf16 gemv for common transformer shapes --- astrai/extension/backend/linear.py | 39 +- benchmarks/infraswe/README.md | 36 ++ .../infraswe/astrai-gemv-common-draft.json | 103 +++++ .../gemv_common_l20_sm89_infraswe_score.json | 122 ++++++ csrc/kernels/gemv/bf16_gemv.cu | 91 ++++ .../gemv_common_chain_l20_sm89.json | 138 ++++++ .../gemv_common_kernel_l20_sm89.json | 348 +++++++++++++++ .../gemv_common_main_baseline_l20_sm89.csv | 24 ++ .../gemv_common_warp_tiled_l20_sm89.csv | 24 ++ docs/developer/cuda_kernels.md | 65 ++- docs/developer/decode_linear_benchmark.md | 17 + scripts/tools/benchmark_gemv_common.py | 396 ++++++++++++++++++ tests/extension/test_gemv.py | 28 ++ tests/extension/test_linear_dispatch.py | 51 ++- 14 files changed, 1456 insertions(+), 26 deletions(-) create mode 100644 benchmarks/infraswe/README.md create mode 100644 benchmarks/infraswe/astrai-gemv-common-draft.json create mode 100644 benchmarks/results/gemv_common_l20_sm89_infraswe_score.json create mode 100644 docs/benchmarks/gemv_common_chain_l20_sm89.json create mode 100644 docs/benchmarks/gemv_common_kernel_l20_sm89.json create mode 100644 docs/benchmarks/gemv_common_main_baseline_l20_sm89.csv create mode 100644 docs/benchmarks/gemv_common_warp_tiled_l20_sm89.csv create mode 100644 scripts/tools/benchmark_gemv_common.py diff --git a/astrai/extension/backend/linear.py b/astrai/extension/backend/linear.py index 4c85c82a..4b1ea979 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,41 @@ 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 and M=8 remain empty +# on SM89; their isolated winners do not consistently survive the chain gate. +_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), +} + _AUTO_GEMV_SHAPES: dict[tuple[int, int], dict[int, frozenset[tuple[int, int]]]] = { (8, 9): { - 2: frozenset( + 2: _COMMON_TRANSFORMER_SM89_SHAPES + | frozenset( { (256, 1536), (1536, 1536), (100000, 1536), } ), - 4: frozenset({(256, 1536), (1536, 1536)}), + 4: _COMMON_TRANSFORMER_SM89_M4_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..be156488 --- /dev/null +++ b/benchmarks/infraswe/README.md @@ -0,0 +1,36 @@ +# 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. Neither 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 **91.68/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..1e7b81ef --- /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": 1, + "state": "D3-contract-proposed", + "created_by": "0z5a" + }, + "target": { + "mode": "repository", + "repository": "https://github.com/ViperEkura/AstrAI", + "revision": "sha256:671ae3f0331ace4f60b4331e963a61987662a7e3234596d0f057834c307acb74", + "project_profile_sha256": "sha256:953067c47f5298f06846c8fe873db55afae49b6588a81926e49d9ef6b0994e08" + }, + "candidate": { + "kind": "git-diff", + "revision": "sha256:51e8f7630ff67a623cbb41f055648cd7f5fc1c9650a7da5f9ecbb817d48f88d5", + "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:671ae3f0331ace4f60b4331e963a61987662a7e3234596d0f057834c307acb74" + }, + "deployment": { + "workload_portfolio": { + "id": "astrai-common-shape-gemv-chains-v1", + "sha256": "sha256:b5752a012369bd216ac45611c91e6ba140f36113be703c44e02240988e6ec089", + "path": "docs/benchmarks/gemv_common_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:113db75a413d558d7cab3de71c72b56ca1932f0a1f5bb790585b07c08a31f42d", + "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:253af102dda5a2673efaf9bc4d35491d46a52e3fee6baa606a80bc54e0c77cbc", + "probe_set_sha256": "sha256:835ff89893c9dc5aefe6b0f3e0c5fddba452b4467745308ac5fe98a5b8238a4d", + "hidden_probe_policy_sha256": "sha256:e181380d1b2435c18569629904153feb91bf0494d12d05e353e50c71ebab90bd" + }, + "project_objectives": { + "edge_ecosystem_policy": "experimental", + "profile_set_sha256": "sha256:b2312952e3b6cf61a527a61806eb39278a1490758e5356436b3b2c36ddffae67" + }, + "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..a15636ec --- /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": 91.68455150382596, + "component_values": { + "evolutionary_maintainability": 0.8833373508742498, + "project_contract_fit": 1.0, + "performance_reuse_utilization": 0.9160131161571949, + "operational_fit": 0.8215838362577492 + }, + "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.75, + "locality": 0.95, + "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.75, + "retention": 1.0, + "family": 0.9, + "compile": 1.0 + }, + "operational_fit": { + "replay": 0.9, + "load": 0.75, + "resource": 0.9, + "coldsteady": 0.75 + } + }, + "input_rationale": { + "evolution": "The candidate builds on the merged AstrAI GEMV stack and records a rejected M=2 warp-tiling experiment, but has no upstream maintenance history yet.", + "locality": "The CUDA change is isolated to one exact-shape M=4 launcher branch; dispatcher, benchmark, tests, and documentation are separately scoped.", + "tests": "The complete CPU-capable AstrAI suite passed (613 passed, 141 skipped) and the built L20 extension passed 96 GPU tests.", + "failure": "Unaligned K, unsupported M, training/autograd, unmeasured shapes, other architectures, and the sub-threshold LLaMA 2 7B M=4 chain fail closed to the existing path.", + "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 without spills in the added kernel, 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.", + "attainment": "Every enabled common-shape chain exceeds the 3% gate; the only sub-threshold M=4 family is explicitly excluded.", + "coverage": "The required NVIDIA L20/SM89 cell was tested; optional A100/SM80 and H100/SM90 cells remain untested and fall back in auto mode.", + "retention": "The exact-shape launcher avoids AstrAI's pre-existing shape bands; their direct benchmark and the regression suite passed.", + "family": "Eleven LLaMA 2/3 and GPT-NeoX matrix shapes plus four dependent synthetic chains cover M=2 and M=4, 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 include Python dispatch; a real LLaMA/GPT 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:671ae3f0331ace4f60b4331e963a61987662a7e3234596d0f057834c307acb74", + "change_intent": "add-fastpath", + "semantic_contract_sha256": "sha256:113db75a413d558d7cab3de71c72b56ca1932f0a1f5bb790585b07c08a31f42d", + "acceptance_contract_sha256": "sha256:253af102dda5a2673efaf9bc4d35491d46a52e3fee6baa606a80bc54e0c77cbc", + "probe_set_sha256": "sha256:835ff89893c9dc5aefe6b0f3e0c5fddba452b4467745308ac5fe98a5b8238a4d", + "workload_portfolio_sha256": "sha256:b5752a012369bd216ac45611c91e6ba140f36113be703c44e02240988e6ec089", + "performance_target_sha256": "sha256:b2312952e3b6cf61a527a61806eb39278a1490758e5356436b3b2c36ddffae67", + "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": "869fa79df766d8499ae3a6c1e4dacfe9f354f6444f4dc8df7dc7684cc177bdf3", + "infraswe_draft_engine_tests": "53 passed", + "astrai_tests": "613 passed, 141 skipped", + "astrai_l20_gpu_tests": "96 passed", + "astrai_lint": "ruff format and import-order checks passed", + "sm89_build": "bf16_gemv target built successfully; added M=4 kernel 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..5bbfe827 100644 --- a/csrc/kernels/gemv/bf16_gemv.cu +++ b/csrc/kernels/gemv/bf16_gemv.cu @@ -13,6 +13,7 @@ namespace { constexpr int kThreads = 256; constexpr int kWarpSize = 32; +constexpr int kWarpTiledThreads = 128; __device__ __forceinline__ float warp_sum(float value) { #pragma unroll @@ -158,6 +159,82 @@ __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 void launch_bf16_gemv( const __nv_bfloat16* x, @@ -168,6 +245,20 @@ void launch_bf16_gemv( int k, cudaStream_t stream ) { + 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; + } + } 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_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..76a9a3fa 100644 --- a/docs/developer/cuda_kernels.md +++ b/docs/developer/cuda_kernels.md @@ -21,15 +21,19 @@ 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 @@ -38,12 +42,45 @@ 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: -| 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, +| 2 | AstrAI `(256,1536)`, `(1536,1536)`, `(100000,1536)` plus all common shapes below | +14.0% on AstrAI 1B; +5.66% to +8.50% common chains | +| 4 | AstrAI `(256,1536)`, `(1536,1536)` plus gated common shapes below | +11.8% on AstrAI 1B; +5.67% to +6.44% common chains | + +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. + +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 | + +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`. + +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..17313f0f 100644 --- a/docs/developer/decode_linear_benchmark.md +++ b/docs/developer/decode_linear_benchmark.md @@ -20,3 +20,20 @@ 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 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`. +The synthetic-chain suite alternates `ASTRAI_GEMV=0` and `auto`, includes +dependent MLP work and Python dispatch, and rotates through distinct weights. +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..7a623880 --- /dev/null +++ b/scripts/tools/benchmark_gemv_common.py @@ -0,0 +1,396 @@ +"""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 + + +@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), +) + +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), +) + + +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 + + 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), + "gate": _weight(spec.intermediate, spec.hidden, device, std), + } + ) + return weights + + +def _chain_fn( + x: torch.Tensor, weights: dict[str, torch.Tensor], fused_qkv: bool +) -> Callable[[], torch.Tensor]: + def run() -> torch.Tensor: + output_projection = linear(x, weights["o"]) + up = linear(x, weights["up"]) + if 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"]) + gate = linear(x, weights["gate"]) + hidden = F.silu(gate) * 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]] = [] + 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.fused_qkv) + with torch.inference_mode(): + _set_mode("0") + reference = run() + _set_mode("auto") + 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("auto"), + ) + 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("--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("--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": { + "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..c3b91eab 100644 --- a/tests/extension/test_gemv.py +++ b/tests/extension/test_gemv.py @@ -43,6 +43,34 @@ 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), + ], +) +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 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..9d066113 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,27 @@ def test_linear_backend_is_public(): assert linear is public_linear +def test_sm89_common_shape_policy_keeps_subthreshold_m4_family_disabled(): + 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)} + policy = _AUTO_GEMV_SHAPES[(8, 9)] + assert common <= policy[2] + assert common - subthreshold_m4 <= policy[4] + assert subthreshold_m4.isdisjoint(policy[4]) + + def test_model_linear_routes_through_backend(monkeypatch): sentinel = torch.randn(2, 4) @@ -109,10 +131,27 @@ 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), + ], +) +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 +172,12 @@ 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), ], ) def test_auto_rejects_measured_small_batch_losers(monkeypatch, m, n, k): From 8545006c6abacbaaab3a6b822704097c2567b953 Mon Sep 17 00:00:00 2001 From: 0z5a Date: Wed, 2 Sep 2026 23:33:02 +0800 Subject: [PATCH 2/2] perf: deepen common-shape bf16 gemv tuning --- astrai/extension/backend/linear.py | 37 +- benchmarks/infraswe/README.md | 7 +- .../infraswe/astrai-gemv-common-draft.json | 22 +- .../gemv_common_l20_sm89_infraswe_score.json | 60 +- csrc/kernels/gemv/bf16_gemv.cu | 55 +- .../gemv_common_extended_chain_l20_sm89.json | 213 +++ .../gemv_common_extended_kernel_l20_sm89.json | 701 ++++++++++ ...emv_common_halfcta_vs_256cta_l20_sm89.json | 1168 +++++++++++++++++ docs/developer/cuda_kernels.md | 38 +- docs/developer/decode_linear_benchmark.md | 5 +- scripts/tools/benchmark_gemv_common.py | 72 +- tests/extension/test_gemv.py | 32 + tests/extension/test_linear_dispatch.py | 39 +- 13 files changed, 2381 insertions(+), 68 deletions(-) create mode 100644 docs/benchmarks/gemv_common_extended_chain_l20_sm89.json create mode 100644 docs/benchmarks/gemv_common_extended_kernel_l20_sm89.json create mode 100644 docs/benchmarks/gemv_common_halfcta_vs_256cta_l20_sm89.json diff --git a/astrai/extension/backend/linear.py b/astrai/extension/backend/linear.py index 4b1ea979..38fd275c 100644 --- a/astrai/extension/backend/linear.py +++ b/astrai/extension/backend/linear.py @@ -31,8 +31,9 @@ # 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 projection-chain/engine -# >=3% gates pass and output argmax remains stable. M=1 and M=8 remain empty -# on SM89; their isolated winners do not consistently survive the chain gate. +# >=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 @@ -53,10 +54,37 @@ (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): { + 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), @@ -64,7 +92,10 @@ (100000, 1536), } ), - 4: _COMMON_TRANSFORMER_SM89_M4_SHAPES | 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 index be156488..fd0be307 100644 --- a/benchmarks/infraswe/README.md +++ b/benchmarks/infraswe/README.md @@ -25,10 +25,13 @@ 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. Neither is presented as whole-model throughput. +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 **91.68/100** and BenchmarkTrust of +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, diff --git a/benchmarks/infraswe/astrai-gemv-common-draft.json b/benchmarks/infraswe/astrai-gemv-common-draft.json index 1e7b81ef..d44e9d65 100644 --- a/benchmarks/infraswe/astrai-gemv-common-draft.json +++ b/benchmarks/infraswe/astrai-gemv-common-draft.json @@ -2,19 +2,19 @@ "schema_version": "0.5", "draft": { "id": "astrai-common-shape-bf16-gemv-v1", - "revision": 1, + "revision": 2, "state": "D3-contract-proposed", "created_by": "0z5a" }, "target": { "mode": "repository", "repository": "https://github.com/ViperEkura/AstrAI", - "revision": "sha256:671ae3f0331ace4f60b4331e963a61987662a7e3234596d0f057834c307acb74", + "revision": "sha256:59c45de7cb0356e9661cfb73d38bd11c888a8b8d564e9123a90c9fa3910cc498", "project_profile_sha256": "sha256:953067c47f5298f06846c8fe873db55afae49b6588a81926e49d9ef6b0994e08" }, "candidate": { "kind": "git-diff", - "revision": "sha256:51e8f7630ff67a623cbb41f055648cd7f5fc1c9650a7da5f9ecbb817d48f88d5", + "revision": "sha256:d9ae3993caf937517b121d4fe855f2c5172cf38f69aa4bda10bbb8daad52a730", "intent": "add-fastpath", "implementation_kind": "cuda-native", "entrypoints": [ @@ -29,13 +29,13 @@ }, "baseline": { "mode": "target-head", - "revision": "sha256:671ae3f0331ace4f60b4331e963a61987662a7e3234596d0f057834c307acb74" + "revision": "sha256:59c45de7cb0356e9661cfb73d38bd11c888a8b8d564e9123a90c9fa3910cc498" }, "deployment": { "workload_portfolio": { "id": "astrai-common-shape-gemv-chains-v1", - "sha256": "sha256:b5752a012369bd216ac45611c91e6ba140f36113be703c44e02240988e6ec089", - "path": "docs/benchmarks/gemv_common_chain_l20_sm89.json" + "sha256": "sha256:c3cf734200844d7b8308a4803f456cd9a3f0d7e19839e7d3e74778f6d1307e4d", + "path": "docs/benchmarks/gemv_common_extended_chain_l20_sm89.json" }, "required_cells": [ "cuda-sm89-nvidia-l20" @@ -46,7 +46,7 @@ ], "request_or_step_protocol": { "id": "astrai-gemv-common-benchmark-v1", - "sha256": "sha256:113db75a413d558d7cab3de71c72b56ca1932f0a1f5bb790585b07c08a31f42d", + "sha256": "sha256:74808a77e551ffa5f77e3256930f2647ebd080828cfdbdcc2770f9d1fa298ff6", "path": "docs/developer/decode_linear_benchmark.md" } }, @@ -65,13 +65,13 @@ "acceptance_contract": { "status": "proposed", "path": "tests/extension/test_gemv.py,tests/extension/test_linear_dispatch.py", - "sha256": "sha256:253af102dda5a2673efaf9bc4d35491d46a52e3fee6baa606a80bc54e0c77cbc", - "probe_set_sha256": "sha256:835ff89893c9dc5aefe6b0f3e0c5fddba452b4467745308ac5fe98a5b8238a4d", - "hidden_probe_policy_sha256": "sha256:e181380d1b2435c18569629904153feb91bf0494d12d05e353e50c71ebab90bd" + "sha256": "sha256:15ffa9e1378d7cb1b06f018f9e1aac16bf56f14446271d51129f86c13f3827bb", + "probe_set_sha256": "sha256:beafc95a2a4aaa587079f1938db8c0fb115d63c2b0c3b185debc9a6132397653", + "hidden_probe_policy_sha256": "sha256:e6bdf6b5a307c9ebdaf643d18ae88b51c6e554d7d9df4427c2f7ba0eb5cc8e93" }, "project_objectives": { "edge_ecosystem_policy": "experimental", - "profile_set_sha256": "sha256:b2312952e3b6cf61a527a61806eb39278a1490758e5356436b3b2c36ddffae67" + "profile_set_sha256": "sha256:22da954f941ba912f60d79b7a58121436d4ba97e64f0182f66b8274080c3a75a" }, "benchmark_loop": { "fast_stage_max_official_fraction": 0.05, diff --git a/benchmarks/results/gemv_common_l20_sm89_infraswe_score.json b/benchmarks/results/gemv_common_l20_sm89_infraswe_score.json index a15636ec..b326c35c 100644 --- a/benchmarks/results/gemv_common_l20_sm89_infraswe_score.json +++ b/benchmarks/results/gemv_common_l20_sm89_infraswe_score.json @@ -5,12 +5,12 @@ "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": 91.68455150382596, + "diagnostic_project_fit_100": 92.6415162827573, "component_values": { - "evolutionary_maintainability": 0.8833373508742498, + "evolutionary_maintainability": 0.8835232637338853, "project_contract_fit": 1.0, - "performance_reuse_utilization": 0.9160131161571949, - "operational_fit": 0.8215838362577492 + "performance_reuse_utilization": 0.9515978740065268, + "operational_fit": 0.8438741102892859 }, "component_floors": { "evolutionary_maintainability": 0.6, @@ -20,8 +20,8 @@ }, "subcomponent_inputs": { "evolutionary_maintainability": { - "evolution": 0.75, - "locality": 0.95, + "evolution": 0.78, + "locality": 0.9, "tests": 1.0, "failure": 1.0, "contract": 0.9 @@ -35,36 +35,36 @@ }, "performance_reuse_utilization": { "attainment": 1.0, - "coverage": 0.75, + "coverage": 0.82, "retention": 1.0, - "family": 0.9, + "family": 1.0, "compile": 1.0 }, "operational_fit": { "replay": 0.9, - "load": 0.75, + "load": 0.82, "resource": 0.9, "coldsteady": 0.75 } }, "input_rationale": { - "evolution": "The candidate builds on the merged AstrAI GEMV stack and records a rejected M=2 warp-tiling experiment, but has no upstream maintenance history yet.", - "locality": "The CUDA change is isolated to one exact-shape M=4 launcher branch; dispatcher, benchmark, tests, and documentation are separately scoped.", - "tests": "The complete CPU-capable AstrAI suite passed (613 passed, 141 skipped) and the built L20 extension passed 96 GPU tests.", - "failure": "Unaligned K, unsupported M, training/autograd, unmeasured shapes, other architectures, and the sub-threshold LLaMA 2 7B M=4 chain fail closed to the existing path.", + "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 without spills in the added kernel, 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.", - "attainment": "Every enabled common-shape chain exceeds the 3% gate; the only sub-threshold M=4 family is explicitly excluded.", - "coverage": "The required NVIDIA L20/SM89 cell was tested; optional A100/SM80 and H100/SM90 cells remain untested and fall back in auto mode.", - "retention": "The exact-shape launcher avoids AstrAI's pre-existing shape bands; their direct benchmark and the regression suite passed.", - "family": "Eleven LLaMA 2/3 and GPT-NeoX matrix shapes plus four dependent synthetic chains cover M=2 and M=4, without claiming a whole-model benchmark.", + "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 include Python dispatch; a real LLaMA/GPT service workload was not run.", + "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." }, @@ -97,13 +97,13 @@ }, "comparison_cell": { "target_project_profile_sha256": "sha256:953067c47f5298f06846c8fe873db55afae49b6588a81926e49d9ef6b0994e08", - "target_repository_or_baseline_sha256": "sha256:671ae3f0331ace4f60b4331e963a61987662a7e3234596d0f057834c307acb74", + "target_repository_or_baseline_sha256": "sha256:59c45de7cb0356e9661cfb73d38bd11c888a8b8d564e9123a90c9fa3910cc498", "change_intent": "add-fastpath", - "semantic_contract_sha256": "sha256:113db75a413d558d7cab3de71c72b56ca1932f0a1f5bb790585b07c08a31f42d", - "acceptance_contract_sha256": "sha256:253af102dda5a2673efaf9bc4d35491d46a52e3fee6baa606a80bc54e0c77cbc", - "probe_set_sha256": "sha256:835ff89893c9dc5aefe6b0f3e0c5fddba452b4467745308ac5fe98a5b8238a4d", - "workload_portfolio_sha256": "sha256:b5752a012369bd216ac45611c91e6ba140f36113be703c44e02240988e6ec089", - "performance_target_sha256": "sha256:b2312952e3b6cf61a527a61806eb39278a1490758e5356436b3b2c36ddffae67", + "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", @@ -112,11 +112,11 @@ }, "execution": { "infraswe_commit": "811bc775ed5b3a6ec853219245f3469f78818020", - "draft_resolution_sha256": "869fa79df766d8499ae3a6c1e4dacfe9f354f6444f4dc8df7dc7684cc177bdf3", + "draft_resolution_sha256": "952ee658e669126582b8eb9aeedede9f70dae6348ea0f1119d64a36d5da100fe", "infraswe_draft_engine_tests": "53 passed", - "astrai_tests": "613 passed, 141 skipped", - "astrai_l20_gpu_tests": "96 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; added M=4 kernel used 42 registers, zero barriers, and zero spills" + "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 5bbfe827..58278a82 100644 --- a/csrc/kernels/gemv/bf16_gemv.cu +++ b/csrc/kernels/gemv/bf16_gemv.cu @@ -12,6 +12,7 @@ namespace { constexpr int kThreads = 256; +constexpr int kHalfCtaThreads = 128; constexpr int kWarpSize = 32; constexpr int kWarpTiledThreads = 128; @@ -23,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, @@ -37,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 @@ -130,7 +131,6 @@ __global__ void bf16_gemv_kernel( } } - #pragma unroll for (int row = 0; row < Rows; ++row) { sums[row] = warp_sum(sums[row]); @@ -147,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) { @@ -235,6 +235,44 @@ constexpr bool use_warp_tiled_kernel(int n, int k) { 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, @@ -259,7 +297,14 @@ void launch_bf16_gemv( return; } } - bf16_gemv_kernel<<>>( + 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_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/developer/cuda_kernels.md b/docs/developer/cuda_kernels.md index 76a9a3fa..48a1fc80 100644 --- a/docs/developer/cuda_kernels.md +++ b/docs/developer/cuda_kernels.md @@ -39,13 +39,14 @@ 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 | Validated gain | |---:|---|---:| -| 2 | AstrAI `(256,1536)`, `(1536,1536)`, `(100000,1536)` plus all common shapes below | +14.0% on AstrAI 1B; +5.66% to +8.50% common chains | -| 4 | AstrAI `(256,1536)`, `(1536,1536)` plus gated common shapes below | +11.8% on AstrAI 1B; +5.67% to +6.44% common chains | +| 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 @@ -56,6 +57,24 @@ up/down. In `(N,K)` form it is `(1024,4096)`, `(4096,4096)`, 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 @@ -74,11 +93,20 @@ measured: | 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`. +`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 diff --git a/docs/developer/decode_linear_benchmark.md b/docs/developer/decode_linear_benchmark.md index 17313f0f..92494eec 100644 --- a/docs/developer/decode_linear_benchmark.md +++ b/docs/developer/decode_linear_benchmark.md @@ -25,14 +25,17 @@ 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 python scripts/tools/benchmark_gemv_common.py \ +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 diff --git a/scripts/tools/benchmark_gemv_common.py b/scripts/tools/benchmark_gemv_common.py index 7a623880..33b8f0c9 100644 --- a/scripts/tools/benchmark_gemv_common.py +++ b/scripts/tools/benchmark_gemv_common.py @@ -36,6 +36,7 @@ class Chain: kv: int intermediate: int fused_qkv: bool = False + gated_mlp: bool = True @dataclass(frozen=True) @@ -64,6 +65,17 @@ class Timing: 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 = ( @@ -71,6 +83,9 @@ class Timing: 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), ) @@ -197,6 +212,12 @@ def benchmark_kernels( 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: @@ -254,27 +275,31 @@ def _chain_weights( "q": _weight(spec.hidden, spec.hidden, device, std), "k": _weight(spec.kv, spec.hidden, device, std), "v": _weight(spec.kv, spec.hidden, device, std), - "gate": _weight(spec.intermediate, 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], fused_qkv: bool + 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 fused_qkv: + 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"]) - gate = linear(x, weights["gate"]) - hidden = F.silu(gate) * up + 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 @@ -285,15 +310,22 @@ def benchmark_chains( args: argparse.Namespace, device: torch.device ) -> list[dict[str, object]]: results: list[dict[str, object]] = [] - for spec in CHAINS: + 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.fused_qkv) + run = _chain_fn(x, weights, spec) with torch.inference_mode(): _set_mode("0") reference = run() - _set_mode("auto") + _set_mode(args.candidate_mode) actual = run() baseline, candidate = _measure_pair( run, @@ -302,7 +334,7 @@ def benchmark_chains( samples=args.samples, inner=args.chain_inner, prepare_baseline=lambda: _set_mode("0"), - prepare_candidate=lambda: _set_mode("auto"), + prepare_candidate=lambda: _set_mode(args.candidate_mode), ) results.append( _print_result( @@ -333,11 +365,27 @@ def parse_args() -> argparse.Namespace: 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( @@ -379,6 +427,12 @@ def main() -> None: "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, diff --git a/tests/extension/test_gemv.py b/tests/extension/test_gemv.py index c3b91eab..1ebecc41 100644 --- a/tests/extension/test_gemv.py +++ b/tests/extension/test_gemv.py @@ -59,6 +59,17 @@ def test_bf16_gemv_matches_small_decode_batches(m, n, k): (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): @@ -71,6 +82,27 @@ def test_bf16_gemv_matches_common_transformer_shapes(m, n, k): 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 9d066113..d23fa927 100644 --- a/tests/extension/test_linear_dispatch.py +++ b/tests/extension/test_linear_dispatch.py @@ -24,7 +24,7 @@ def test_linear_backend_is_public(): assert linear is public_linear -def test_sm89_common_shape_policy_keeps_subthreshold_m4_family_disabled(): +def test_sm89_common_shape_policy_keeps_only_validated_families_enabled(): common = { (1024, 4096), (4096, 4096), @@ -39,10 +39,28 @@ def test_sm89_common_shape_policy_keeps_subthreshold_m4_family_disabled(): (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): @@ -115,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( @@ -145,6 +163,17 @@ def test_auto_m1_falls_back_until_end_to_end_gate_passes(monkeypatch): (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): @@ -178,6 +207,12 @@ def test_auto_selects_measured_sm89_small_batch_winner(monkeypatch, m, n, k): (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):