From 2ef846f7512253999bd4833ed2d8a4f65673c6e7 Mon Sep 17 00:00:00 2001 From: wuyang <5700876+banisherwy@user.noreply.gitee.com> Date: Wed, 29 Jul 2026 13:36:49 +0800 Subject: [PATCH] feat: execute FlashKDA on RTX 5090 --- PROGRESS.md | 10 +- README.md | 7 +- experiments/k3/checkpoint_probe.py | 6 +- experiments/k3/flashkda/Dockerfile | 30 ++ experiments/k3/flashkda/README.md | 85 ++++ .../k3/flashkda/build_reference_helper.py | 50 ++ experiments/k3/flashkda/run_probe.py | 454 ++++++++++++++++++ research/K3_ARTIFACT_AUDIT.md | 116 ++++- scripts/check-k3-browser.mjs | 15 +- src/components/K3ArtifactLab.astro | 113 ++++- src/data/k3-artifact-snapshot.json | 6 +- src/data/k3-flashkda-runtime.json | 328 +++++++++++++ src/data/k3Artifacts.ts | 4 +- src/pages/index.astro | 5 +- src/pages/k3/index.astro | 6 +- src/pages/progress/index.astro | 5 +- 16 files changed, 1189 insertions(+), 51 deletions(-) create mode 100644 experiments/k3/flashkda/Dockerfile create mode 100644 experiments/k3/flashkda/README.md create mode 100644 experiments/k3/flashkda/build_reference_helper.py create mode 100644 experiments/k3/flashkda/run_probe.py create mode 100644 src/data/k3-flashkda-runtime.json diff --git a/PROGRESS.md b/PROGRESS.md index a2b0368..a58a1f4 100644 --- a/PROGRESS.md +++ b/PROGRESS.md @@ -8,7 +8,7 @@ |---|---:|---:|---| | 研究框架与规范 | 进行中 | 83% | Scaling Laws 二轮拟合复现与逐图精读 | | 网站设计系统 | 进行中 | 89% | 打印样式与更多通用可视化组件 | -| Kimi K3 深读 | 三轮实证进行中 | 92% | 匹配 CUDA 12.9+ 执行 FlashKDA，并接入真实 hidden-state / expert-load traces | +| Kimi K3 深读 | 三轮实证进行中 | 94% | 接入真实 hidden-state / expert-load / cache traces，并重绘报告数值图 | | 语言模型前史 | 完成首版 | 78% | Kneser–Ney、LSTM、Bahdanau 逐图精读与真实小语料复现 | | Transformer 基础 | 完成首版 | 79% | 多头电路、归一化 traces 与真实 kernel / KV 配置 | | 表示、位置与残差高速公路 | 完成首版 | 81% | 真实 hidden-state / norm traces、长上下文位置外推与深层稳定性消融 | @@ -164,10 +164,14 @@ - [x] 第三轮证据快照、可复现探针脚本与正式审计账本已进入开源树；原始权重字节不提交，真实观测 O、推导 D、执行 X、合成 S 与未决 U 分开标记。 - [x] K3 新四视图本地真实 Chrome 回归通过：93 层条带、tensor group、参数分布、benchmark/router 切换、键盘 tabs、桌面与 390px 移动端均无异常。 - [x] K3 三轮开放工件里程碑以源提交 `be2b291`、不可变镜像 `20260729T044605Z-be2b291` 发布；OCI digest `sha256:99aa953e…d00cdf`，NAS、VPS/Tailscale、NPM、DNS、HTTPS、证书、门户与十六套生产 Chrome 回归全链路通过；保留 `20260729T040336Z-b669615` 回滚。 +- [x] FlashKDA 隔离构建闭环：主机 CUDA 13.1 / glibc 2.43 仍在 `rsqrt` declarations 冲突；改用 CUDA 13.0.2 / Ubuntu 24.04 / glibc 2.39 后成功产出 CPython 3.12、`sm_120a` wheel，SHA-256 `14687b6d…2158d`。 +- [x] RTX 5090 正确性闸门通过：one chunk、partial tail、multi-chunk、96 heads 与 varlen 共 6 组，FlashKDA output / final state 和官方 `torch_ref.py` 逐元素完全相等，max absolute error 0。 +- [x] RTX 5090 K3 形状受控计时完成：fixed / varlen、BF16 / no-state / FP32 state 共 1,800 个 CUDA Event samples；fixed BF16 mean 2.6210 ms、P95 2.6437 ms，varlen BF16 mean 2.3335 ms、P95 2.3574 ms。 +- [x] 可复现 Dockerfile、运行脚本与机器可读 JSON 接入开源树；网站第四视图把作者 H20/GB200 表、本机 RTX 5090 值、exact suite、构建链、synthetic router 和 `A_log` 未决冲突分开显示。 ## 正在进行 -- [ ] K3 三轮下一闸门：在匹配 CUDA 12.9+ 环境执行 FlashKDA correctness / benchmark，获得真实 token hidden states、expert load 与 cache traces，再做逐图数值重绘和独立小模型复现。 +- [ ] K3 三轮下一闸门：获得真实 token hidden states、expert load 与 cache traces，解释或修订 `A_log [128]` 工件冲突，再做 Figure 3/4/5 数值重绘和独立小模型复现。 - [ ] DeepSeek 三轮：真实专家负载、MLA kernel、FP8 / pipeline 与 R1-like RL traces，外加独立小模型复现。 - [ ] 表示、位置与残差二轮：真实 hidden-state / norm traces、长上下文位置外推复现与 mHC / AttnRes 深层稳定性消融。 - [ ] 评测安全二轮：真实 cross-harness / pass@k 复跑、Judge 元评测、动态污染与过拒案例。 @@ -188,6 +192,8 @@ | 日期 | 决策/发现 | 影响 | |---|---|---| +| 2026-07-29 | K3 FlashKDA 在隔离 CUDA 13.0 / glibc 2.39 构建并回到 RTX 5090 执行 | 主机头文件 ABI 问题与 GPU 架构支持分离；wheel checksum、Dockerfile 与执行 JSON 可审计 | +| 2026-07-29 | FlashKDA 本机实测永久与作者 H20 / GB200 表分账 | 6/6 exact 与 1,800 个 latency samples 可称 X；未跑本机 FLA 就不计算本机 speedup | | 2026-07-28 | 网站命名为 **LLM Atlas / 大模型技术全景** | 既能容纳 K3 深读，也能承载完整 LLM 课程 | | 2026-07-28 | K3 作为“汇流点”，不是课程起点 | 初学者可以先学基础，高阶读者可以从 K3 反向跳转 | | 2026-07-28 | 优先重绘论文图并标明“简化/改绘” | 图可缩放、可交互，也减少脱离上下文复制论文图片 | diff --git a/README.md b/README.md index 6403cdd..bd7b15a 100644 --- a/README.md +++ b/README.md @@ -21,11 +21,12 @@ 语言模型前史、Transformer 基础、表示/位置/残差、DeepSeek 技术谱系、Scaling Laws、数据工程、长上下文、MoE、指令微调与人类偏好、推理、工具使用与长程 Agent、原生多模态、训练系统、推理服务、数值优化，以及评测与安全深度专题， 以及 71 个覆盖核心机制的原创交互视图。K3 二轮导读以 32 张问题账、16 图 / 5 表审计、 8 个交互实验和 100 个一手/官方节点，完整覆盖架构、预训练、后训练、系统、评测、案例与附录。 -第三轮已完成首个开放工件里程碑：固定官方模型与 FlashKDA revisions，审计 96 个 checkpoint shards、 +第三轮已完成开放工件与首个真实 kernel 里程碑：固定官方模型与 FlashKDA revisions，审计 96 个 checkpoint shards、 497,220 个 tensor entries、真实 KDA / MLA / MoE / MoonViT shapes 与小范围参数统计，并用 4 个新视图 -明确区分官方观测、确定性推导、本机执行、合成探针和未决矛盾。详见 +明确区分官方观测、确定性推导、本机执行、合成探针和未决矛盾；同时用隔离 CUDA 13.0 环境编译 +`sm_120a` wheel，在 RTX 5090 上完成 6/6 官方参考 exact-match 和 K3 fixed / varlen 形状计时。详见 [K3_ARTIFACT_AUDIT.md](./research/K3_ARTIFACT_AUDIT.md) 与 -[checkpoint_probe.py](./experiments/k3/checkpoint_probe.py)。 +[checkpoint_probe.py](./experiments/k3/checkpoint_probe.py)、[FlashKDA probe](./experiments/k3/flashkda/)。 DeepSeek 二轮专题以 24 张问题账、10 次技术转向、 4 个交互实验和 60 个一手/官方节点，串起 Dense、MoE、MLA、V3 协同、R1 与 V4。 其余专题按进度账本持续扩建。 diff --git a/experiments/k3/checkpoint_probe.py b/experiments/k3/checkpoint_probe.py index 52de01b..3661ed3 100644 --- a/experiments/k3/checkpoint_probe.py +++ b/experiments/k3/checkpoint_probe.py @@ -370,11 +370,11 @@ def main() -> None: else None, "libc": list(platform.libc_ver()), }, - "local_build": { - "status": "blocked_before_kernel execution", + "baseline_host_build": { + "status": "blocked in the default CUDA 12.8 environment; superseded by the separate CUDA 13 runtime probe", "attempt_1": "system g++ 15 exceeds CUDA 12.8 host compiler range", "attempt_2": "temporary g++ 13 reaches nvcc, then CUDA 12.8 headers conflict with current glibc math declarations", - "interpretation": "GPU architecture is listed by the repository, but the local CUDA 12.8 stack is below the official CUDA 12.9 requirement", + "interpretation": "this snapshot records the first host path only; see src/data/k3-flashkda-runtime.json for the successful isolated build and RTX 5090 execution", }, }, } diff --git a/experiments/k3/flashkda/Dockerfile b/experiments/k3/flashkda/Dockerfile new file mode 100644 index 0000000..28553f4 --- /dev/null +++ b/experiments/k3/flashkda/Dockerfile @@ -0,0 +1,30 @@ +FROM nvidia/cuda:13.0.2-devel-ubuntu24.04 + +ENV DEBIAN_FRONTEND=noninteractive + +RUN apt-get update \ + && apt-get install -y --no-install-recommends \ + build-essential \ + git \ + ninja-build \ + python3 \ + python3-dev \ + python3-pip \ + python3-venv \ + && rm -rf /var/lib/apt/lists/* + +RUN python3 -m venv /opt/flashkda \ + && /opt/flashkda/bin/python -m pip install --upgrade \ + pip \ + setuptools \ + wheel \ + ninja \ + numpy \ + pytest \ + --index-url https://pypi.tuna.tsinghua.edu.cn/simple \ + && /opt/flashkda/bin/python -m pip install \ + torch==2.11.0 \ + --index-url https://download.pytorch.org/whl/cu130 + +ENV PATH=/opt/flashkda/bin:${PATH} +ENV CUDA_HOME=/usr/local/cuda diff --git a/experiments/k3/flashkda/README.md b/experiments/k3/flashkda/README.md new file mode 100644 index 0000000..8781c09 --- /dev/null +++ b/experiments/k3/flashkda/README.md @@ -0,0 +1,85 @@ +# FlashKDA RTX 5090 execution probe + +This directory reproduces the local execution evidence shown in the K3 +artifact lab. It does **not** download or run the 1.56 TB Kimi K3 checkpoint. +All inputs are deterministic synthetic tensors with shapes allowed by the +official FlashKDA API. + +Pinned upstream revision: + +```text +MoonshotAI/FlashKDA@1ce47ea3bb22c84eb9cc665028399cf35e8ffb0b +``` + +## Why the build is isolated + +The workstation uses glibc 2.43. CUDA 13.1's published Linux support matrix +currently lists distributions up to glibc 2.41, and compiling this extension +directly on the host reaches an `rsqrt` / `rsqrtf` exception-specification +conflict in the CUDA and system math headers. The container fixes the build +ABI at Ubuntu 24.04 / glibc 2.39 while still targeting `sm_120a`. + +## Build the wheel + +Clone FlashKDA with submodules and verify the revision before building: + +```bash +git clone --recursive https://github.com/MoonshotAI/FlashKDA.git /tmp/FlashKDA +git -C /tmp/FlashKDA checkout 1ce47ea3bb22c84eb9cc665028399cf35e8ffb0b +git -C /tmp/FlashKDA submodule update --init --recursive + +docker build -t llm-atlas-flashkda-cu130 experiments/k3/flashkda +mkdir -p /tmp/flashkda-wheelhouse +docker run --rm \ + -e FLASH_KDA_ARCHS=120a \ + -e MAX_JOBS=12 \ + -v /tmp/FlashKDA:/src:ro \ + -v /tmp/flashkda-wheelhouse:/wheelhouse \ + llm-atlas-flashkda-cu130 \ + python -m pip wheel /src --no-build-isolation --no-deps -w /wheelhouse +``` + +The audited wheel was built for CPython 3.12 and has SHA-256: + +```text +14687b6d84a256d4552f0c73ccf93a601be582aeabcdf49ae3a409266872158d +``` + +## Run the probe + +Use CPython 3.12 with PyTorch 2.11.0+cu130 and install the wheel. Prebuild the +small CUDA helper loaded by upstream `tests/torch_ref.py` without needing a +Docker GPU runtime: + +```bash +mkdir -p /tmp/k3-torch-extensions +docker run --rm \ + --user "$(id -u):$(id -g)" \ + -e HOME=/tmp \ + -e TORCH_EXTENSIONS_DIR=/cache \ + -e TORCH_CUDA_ARCH_LIST=12.0a \ + -e MAX_JOBS=12 \ + -v "$PWD:/atlas:ro" \ + -v /tmp/k3-torch-extensions:/cache \ + llm-atlas-flashkda-cu130 \ + python /atlas/experiments/k3/flashkda/build_reference_helper.py +``` + +Then run the GPU probe on the host: + +```bash +export TORCH_EXTENSIONS_DIR=/tmp/k3-torch-extensions +export TORCH_CUDA_ARCH_LIST=12.0a +export CUDA_HOME=/usr/local/cuda + +python experiments/k3/flashkda/run_probe.py \ + --flashkda-source /tmp/FlashKDA \ + --wheel /tmp/flashkda-wheelhouse/flash_kda-0.0.1+1ce47ea-cp312-cp312-linux_x86_64.whl \ + --output src/data/k3-flashkda-runtime.json +``` + +The correctness suite imports the upstream reference implementation instead +of copying it into this repository. It checks exact BF16 equality at one +chunk, a partial tail chunk, multiple chunks, 96 heads, and a variable-length +batch. The performance cases are local kernel timings, not a comparison with +the authors' H20 or GB200 tables. diff --git a/experiments/k3/flashkda/build_reference_helper.py b/experiments/k3/flashkda/build_reference_helper.py new file mode 100644 index 0000000..6128835 --- /dev/null +++ b/experiments/k3/flashkda/build_reference_helper.py @@ -0,0 +1,50 @@ +#!/usr/bin/env python3 +"""Prebuild the CUDA helper used verbatim by upstream tests/torch_ref.py. + +This runs in the CUDA build container without a GPU. Keeping the extension +name, sources, generated function, and compiler flags identical lets PyTorch +reuse the cache when the official reference module is imported on the host. +""" + +from torch.utils.cpp_extension import load_inline + + +CUDA_SOURCE = r""" +#include +#include + +__global__ void sigmoid_tanh_fp32_kernel(const float* __restrict__ input, + float* __restrict__ output, int n) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx < n) { + float xh = input[idx] * 0.5f; + float th; + asm("tanh.approx.f32 %0, %1;" : "=f"(th) : "f"(xh)); + output[idx] = th * 0.5f + 0.5f; + } +} + +torch::Tensor sigmoid_tanh_fp32(torch::Tensor input) { + auto output = torch::empty_like(input); + int n = input.numel(); + sigmoid_tanh_fp32_kernel<<<(n + 255) / 256, 256>>>( + input.data_ptr(), output.data_ptr(), n); + return output; +} +""" + + +def main() -> None: + module = load_inline( + name="sigmoid_ext", + cpp_sources="torch::Tensor sigmoid_tanh_fp32(torch::Tensor input);", + cuda_sources=CUDA_SOURCE, + functions=["sigmoid_tanh_fp32"], + extra_cuda_cflags=["-O2"], + verbose=True, + ) + print(module.__file__) + + +if __name__ == "__main__": + main() diff --git a/experiments/k3/flashkda/run_probe.py b/experiments/k3/flashkda/run_probe.py new file mode 100644 index 0000000..3554133 --- /dev/null +++ b/experiments/k3/flashkda/run_probe.py @@ -0,0 +1,454 @@ +#!/usr/bin/env python3 +"""Run auditable FlashKDA correctness and latency probes on one CUDA GPU. + +The upstream torch reference is imported from a pinned FlashKDA checkout. The +script deliberately uses valid synthetic A_log[H] tensors: it does not resolve +the public Kimi K3 checkpoint's A_log[128] versus H=96 inconsistency. +""" + +from __future__ import annotations + +import argparse +import hashlib +import importlib.metadata +import json +import math +import platform +import statistics +import subprocess +import sys +import time +from datetime import datetime, timezone +from pathlib import Path +from typing import Any, Callable + +import torch +import torch.nn.functional as F + +import flash_kda +import flash_kda_C + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--flashkda-source", type=Path, required=True) + parser.add_argument("--wheel", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--warmup", type=int, default=20) + parser.add_argument("--iters", type=int, default=100) + parser.add_argument("--repeats", type=int, default=3) + parser.add_argument("--seed", type=int, default=20260729) + parser.add_argument("--captured-at", default=None) + return parser.parse_args() + + +def sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + for block in iter(lambda: handle.read(1024 * 1024), b""): + digest.update(block) + return digest.hexdigest() + + +def percentile(sorted_values: list[float], fraction: float) -> float: + if not sorted_values: + return float("nan") + index = fraction * (len(sorted_values) - 1) + lower = math.floor(index) + upper = math.ceil(index) + if lower == upper: + return sorted_values[lower] + weight = index - lower + return sorted_values[lower] * (1 - weight) + sorted_values[upper] * weight + + +def nvidia_smi() -> dict[str, Any]: + fields = [ + "name", + "driver_version", + "memory.total", + "power.limit", + "clocks.max.sm", + "clocks.max.memory", + ] + output = subprocess.check_output( + [ + "nvidia-smi", + f"--query-gpu={','.join(fields)}", + "--format=csv,noheader,nounits", + ], + text=True, + ).strip() + values = [value.strip() for value in output.split(",")] + return dict(zip(fields, values, strict=True)) + + +def make_inputs( + sequence_lengths: list[int], + heads: int, + seed: int, +) -> dict[str, torch.Tensor | float | None]: + batch = 1 + dimension = 128 + total_tokens = sum(sequence_lengths) + sequences = len(sequence_lengths) + generator = torch.Generator(device="cuda").manual_seed(seed) + + q = F.normalize( + torch.randn( + (batch, total_tokens, heads, dimension), + dtype=torch.float32, + device="cuda", + generator=generator, + ), + p=2, + dim=-1, + ).to(torch.bfloat16) + k = F.normalize( + torch.randn( + (batch, total_tokens, heads, dimension), + dtype=torch.float32, + device="cuda", + generator=generator, + ), + p=2, + dim=-1, + ).to(torch.bfloat16) + v = torch.randn( + (batch, total_tokens, heads, dimension), + dtype=torch.bfloat16, + device="cuda", + generator=generator, + ) + g = torch.randn( + (batch, total_tokens, heads, dimension), + dtype=torch.bfloat16, + device="cuda", + generator=generator, + ) + beta = torch.randn( + (batch, total_tokens, heads), + dtype=torch.bfloat16, + device="cuda", + generator=generator, + ) + a_log = torch.rand( + heads, dtype=torch.float32, device="cuda", generator=generator + ) + dt_bias = torch.rand( + (heads, dimension), + dtype=torch.float32, + device="cuda", + generator=generator, + ) + initial_state = torch.randn( + (sequences, heads, dimension, dimension), + dtype=torch.bfloat16, + device="cuda", + generator=generator, + ) + cu_seqlens = None + if len(sequence_lengths) > 1: + offsets = [0] + for length in sequence_lengths: + offsets.append(offsets[-1] + length) + cu_seqlens = torch.tensor(offsets, dtype=torch.long, device="cuda") + + return { + "q": q, + "k": k, + "v": v, + "g": g, + "beta": beta, + "A_log": a_log, + "dt_bias": dt_bias, + "initial_state": initial_state, + "cu_seqlens": cu_seqlens, + "scale": 1 / math.sqrt(dimension), + } + + +def run_correctness_case( + torch_ref: Callable[..., None], + name: str, + sequence_lengths: list[int], + heads: int, + seed: int, +) -> dict[str, Any]: + inputs = make_inputs(sequence_lengths, heads, seed) + q = inputs["q"] + assert isinstance(q, torch.Tensor) + initial_state = inputs["initial_state"] + assert isinstance(initial_state, torch.Tensor) + + out_kernel = torch.zeros_like(q) + out_reference = torch.zeros_like(q) + state_kernel = torch.zeros_like(initial_state) + state_reference = torch.zeros_like(initial_state) + common = { + "A_log": inputs["A_log"], + "dt_bias": inputs["dt_bias"], + "lower_bound": -5.0, + "initial_state": initial_state.clone(), + "cu_seqlens": inputs["cu_seqlens"], + } + + start = time.perf_counter() + flash_kda.fwd( + inputs["q"], + inputs["k"], + inputs["v"], + inputs["g"], + inputs["beta"], + inputs["scale"], + out_kernel, + final_state=state_kernel, + **common, + ) + torch.cuda.synchronize() + kernel_ms = (time.perf_counter() - start) * 1000 + + start = time.perf_counter() + torch_ref( + inputs["q"], + inputs["k"], + inputs["v"], + inputs["g"], + inputs["beta"], + inputs["scale"], + out_reference, + final_state=state_reference, + **common, + ) + torch.cuda.synchronize() + reference_ms = (time.perf_counter() - start) * 1000 + + output_diff = (out_kernel.float() - out_reference.float()).abs() + state_diff = (state_kernel.float() - state_reference.float()).abs() + result = { + "name": name, + "sequence_lengths": sequence_lengths, + "heads": heads, + "dimension": 128, + "kernel_ms_first_measured": kernel_ms, + "reference_ms": reference_ms, + "output_exact": bool(torch.equal(out_kernel, out_reference)), + "state_exact": bool(torch.equal(state_kernel, state_reference)), + "output_max_abs_diff": output_diff.max().item(), + "output_mean_abs_diff": output_diff.mean().item(), + "state_max_abs_diff": state_diff.max().item(), + "state_mean_abs_diff": state_diff.mean().item(), + } + if not result["output_exact"] or not result["state_exact"]: + raise AssertionError(f"exact correctness failed: {result}") + return result + + +def bench_events( + function: Callable[[], None], + warmup: int, + iters: int, + repeats: int, +) -> list[float]: + for _ in range(max(warmup, 1)): + function() + torch.cuda.synchronize() + + elapsed: list[float] = [] + for _ in range(repeats): + starts = [torch.cuda.Event(enable_timing=True) for _ in range(iters)] + ends = [torch.cuda.Event(enable_timing=True) for _ in range(iters)] + torch.cuda.synchronize() + for index in range(iters): + starts[index].record() + function() + ends[index].record() + torch.cuda.synchronize() + elapsed.extend( + start.elapsed_time(end) for start, end in zip(starts, ends, strict=True) + ) + return elapsed + + +def summarize_latency( + values: list[float], + total_tokens: int, +) -> dict[str, Any]: + ordered = sorted(float(value) for value in values) + mean = statistics.fmean(ordered) + return { + "samples": len(ordered), + "mean_ms": mean, + "min_ms": ordered[0], + "p50_ms": percentile(ordered, 0.50), + "p95_ms": percentile(ordered, 0.95), + "max_ms": ordered[-1], + "sequence_tokens_per_second": total_tokens / (mean / 1000), + } + + +def run_benchmark_case( + name: str, + sequence_lengths: list[int], + heads: int, + seed: int, + warmup: int, + iters: int, + repeats: int, +) -> dict[str, Any]: + torch.cuda.empty_cache() + torch.cuda.reset_peak_memory_stats() + inputs = make_inputs(sequence_lengths, heads, seed) + q = inputs["q"] + initial_state = inputs["initial_state"] + assert isinstance(q, torch.Tensor) + assert isinstance(initial_state, torch.Tensor) + out = torch.zeros_like(q) + state_bf16 = torch.zeros_like(initial_state) + initial_fp32 = initial_state.float() + state_fp32 = torch.zeros_like(initial_fp32) + common = { + "A_log": inputs["A_log"], + "dt_bias": inputs["dt_bias"], + "lower_bound": -5.0, + "cu_seqlens": inputs["cu_seqlens"], + } + + def invoke(initial: torch.Tensor | None, final: torch.Tensor | None) -> None: + flash_kda.fwd( + inputs["q"], + inputs["k"], + inputs["v"], + inputs["g"], + inputs["beta"], + inputs["scale"], + out, + initial_state=initial, + final_state=final, + **common, + ) + + variants = { + "bf16_state": lambda: invoke(initial_state, state_bf16), + "no_state": lambda: invoke(None, None), + "fp32_state": lambda: invoke(initial_fp32, state_fp32), + } + timings = { + variant: summarize_latency( + bench_events(function, warmup, iters, repeats), + sum(sequence_lengths), + ) + for variant, function in variants.items() + } + torch.cuda.synchronize() + return { + "name": name, + "sequence_lengths": sequence_lengths, + "total_tokens": sum(sequence_lengths), + "heads": heads, + "dimension": 128, + "warmup": warmup, + "iters": iters, + "repeats": repeats, + "timings": timings, + "output_abs_mean_after_last_run": out.float().abs().mean().item(), + "peak_allocated_mib": torch.cuda.max_memory_allocated() / 2**20, + "peak_reserved_mib": torch.cuda.max_memory_reserved() / 2**20, + } + + +def main() -> None: + args = parse_args() + if not torch.cuda.is_available(): + raise RuntimeError("CUDA GPU is required") + if torch.cuda.get_device_capability(0) < (9, 0): + raise RuntimeError("FlashKDA requires SM90 or newer") + tests_dir = args.flashkda_source / "tests" + sys.path.insert(0, str(tests_dir)) + from torch_ref import torch_ref + + revision = subprocess.check_output( + ["git", "-C", str(args.flashkda_source), "rev-parse", "HEAD"], + text=True, + ).strip() + torch.manual_seed(args.seed) + torch.cuda.manual_seed_all(args.seed) + + correctness_specs = [ + ("one_chunk", [16], 1), + ("partial_tail", [17], 1), + ("two_chunks", [32], 1), + ("multi_chunk_multi_head", [65], 2), + ("k3_head_count", [17], 96), + ("varlen_partial_chunks", [17, 31], 2), + ] + correctness = [ + run_correctness_case(torch_ref, name, lengths, heads, args.seed + index) + for index, (name, lengths, heads) in enumerate(correctness_specs) + ] + + benchmark_specs = [ + ("teaching_scale", [512], 8), + ("intermediate", [2048], 32), + ("k3_fixed_shape", [8192], 96), + ("k3_varlen_shape", [1300, 547, 2048, 963, 271, 3063], 96), + ] + benchmarks = [ + run_benchmark_case( + name, + lengths, + heads, + args.seed + 100 + index, + args.warmup, + args.iters, + args.repeats, + ) + for index, (name, lengths, heads) in enumerate(benchmark_specs) + ] + + captured_at = args.captured_at or datetime.now(timezone.utc).isoformat() + result = { + "schema_version": 1, + "captured_at": captured_at, + "evidence_identity": "X / local execution on deterministic synthetic tensors", + "boundary": { + "k3_checkpoint_loaded": False, + "real_token_hidden_states": False, + "a_log_shape_conflict_resolved": False, + "benchmark_comparison": "local FlashKDA timings only; author H20/GB200 tables remain separate", + }, + "provenance": { + "flashkda_revision": revision, + "flashkda_package": importlib.metadata.version("flash-kda"), + "wheel_filename": args.wheel.name, + "wheel_sha256": sha256(args.wheel), + "runner": str(Path(__file__).relative_to(Path.cwd())), + }, + "environment": { + "python": platform.python_version(), + "platform": platform.platform(), + "libc": list(platform.libc_ver()), + "torch": torch.__version__, + "torch_cuda": torch.version.cuda, + "gpu": torch.cuda.get_device_name(0), + "capability": list(torch.cuda.get_device_capability(0)), + "nvidia_smi": nvidia_smi(), + "flash_kda_module": flash_kda.__file__, + "flash_kda_extension": flash_kda_C.__file__, + }, + "correctness": { + "all_exact": all( + case["output_exact"] and case["state_exact"] + for case in correctness + ), + "cases": correctness, + }, + "benchmarks": benchmarks, + } + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(result, indent=2) + "\n") + print(json.dumps(result, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/research/K3_ARTIFACT_AUDIT.md b/research/K3_ARTIFACT_AUDIT.md index b55e3bd..a3b2c4a 100644 --- a/research/K3_ARTIFACT_AUDIT.md +++ b/research/K3_ARTIFACT_AUDIT.md @@ -22,7 +22,7 @@ - 没有在单张 RTX 5090 上加载 1.56 TB checkpoint； - 没有获得真实 token hidden states、线上 expert load 或生产 cache trace； - 没有把随机向量上的 router 行为写成真实数据分布； -- 没有把未成功执行的 FlashKDA kernel 写成“本机 benchmark”； +- 没有把合法 synthetic shape 上的 FlashKDA 实测写成“K3 checkpoint 已跑通”； - 没有替官方解释下面发现的 `A_log` 形状不一致。 ## 2. 一手工件与校验 @@ -268,7 +268,7 @@ RMS=1 的各向同性随机 hidden vectors，再比较 top-16： ### 7.1 官方仓库事实 -FlashKDA `1ce47ea3`： +[FlashKDA 官方仓库](https://github.com/MoonshotAI/FlashKDA) `1ce47ea3`： - CUTLASS kernels； - 支持 `90a / 100a / 103a / 120a`； @@ -288,51 +288,127 @@ FlashKDA `1ce47ea3`： 这些是作者仓库 benchmark，不是本站复跑值。 -### 7.2 本机真实构建边界 +### 7.2 从主机构建失败到隔离 wheel -本机： +GPU 主机： - RTX 5090，compute capability `12.0`； -- PyTorch `2.11.0+cu128`； -- PyTorch CUDA `12.8`； +- Ubuntu 26.04，glibc `2.43`； +- 常用环境为 PyTorch `2.11.0+cu128` / CUDA `12.8`； - 官方源码明确包含 `sm_120a`，所以不是 GPU architecture 缺失。 -两次可复现构建： +主机路径依次暴露了两个不同问题： 1. 系统 `g++ 15.2`：PyTorch extension 在编译前拒绝，CUDA 12.8 要求 host compiler `<14`； 2. 临时解包 `g++ 13.4`：成功进入 nvcc，但 CUDA 12.8 headers 与当前 glibc math declarations - 在 `cospi / sinpi / rsqrt` exception specification 处冲突。 + 冲突； +3. 新建 PyTorch `2.11.0+cu130` 环境并改用 CUDA 13.1 后，版本主线已经匹配，但仍在 + `rsqrt / rsqrtf` exception specification 处失败。 -结论： +第三步说明“升级到 CUDA 13”还不够。[CUDA 13.1 官方 Linux support matrix](https://docs.nvidia.com/cuda/archive/13.1.0/cuda-installation-guide-linux/index.html) 当前列出的 +glibc 上限是 2.41，而主机是 2.43；NVIDIA 官方论坛也记录了同类 CUDA 13.1 / 新 glibc +数学声明[冲突案例](https://forums.developer.nvidia.com/t/fedora-43-and-nvcc-cuda13-1-error-exception-specification-is-incompatible-rsqrt-rsqrtf/354510)。这里能写的是“构建失败与超出验证矩阵一致”，不能写成 CUDA 对未来 glibc +永久不兼容。 -- kernel 尚未在本站机器执行； -- 失败与 README 的 CUDA 12.9+ 要求一致； -- 不能把 `sm_120a` 支持写成本机已经跑通； -- 下一次应使用匹配 PyTorch 的 CUDA 12.9+ toolchain 或官方容器后再复跑 correctness + benchmark。 +最终构建固定为： + +| layer | pinned value | +|---|---| +| base image | `nvidia/cuda:13.0.2-devel-ubuntu24.04` | +| container libc | glibc `2.39` | +| Python | `3.12` | +| PyTorch | `2.11.0+cu130` | +| target | `sm_120a` | +| FlashKDA | `1ce47ea3bb22c84eb9cc665028399cf35e8ffb0b` | +| wheel | `flash_kda-0.0.1+1ce47ea-cp312-cp312-linux_x86_64.whl` | +| wheel SHA-256 | `14687b6d84a256d4552f0c73ccf93a601be582aeabcdf49ae3a409266872158d` | + +容器没有 NVIDIA runtime，也不需要 GPU 才能为固定 architecture 编译。CUDA 编译器实际产出 +`sm_120a` 对象，最终 wheel 为 3,806,174 bytes。wheel 随后装入主机独立 CPython 3.12.11 / +PyTorch 2.11.0+cu130 环境；该运行环境能看到 RTX 5090、capability 12.0，并成功导入 +`flash_kda` 与 `flash_kda_C`。 + +### 7.3 exact correctness：输出与 recurrent state 都必须相等 + +正确性对照直接调用官方仓库的 `tests/torch_ref.py`，而不是本站重写一份近似参考。 +测试输入满足公开 API 的 `A_log[H]`；它们故意不使用 checkpoint 的 `A_log[128]`，因此不会 +偷偷替未决形状冲突下结论。 + +| case | sequence lengths | H | output exact | final state exact | max abs error | +|---|---|---:|---:|---:|---:| +| one chunk | `[16]` | 1 | yes | yes | 0 | +| partial tail | `[17]` | 1 | yes | yes | 0 | +| two chunks | `[32]` | 1 | yes | yes | 0 | +| multi-chunk / multi-head | `[65]` | 2 | yes | yes | 0 | +| K3 head count | `[17]` | 96 | yes | yes | 0 | +| varlen partial chunks | `[17,31]` | 2 | yes | yes | 0 | + +六组用 BF16 q/k/v/g、BF16 beta logits、F32 `A_log/dt_bias` 和 BF16 recurrent state； +output 与 final state 全部逐元素相等，所有 max / mean absolute difference 都是 0。 +这把证据从“源码声称支持 120a”推进成了“本站 RTX 5090 实际执行并通过官方参考”。 + +### 7.4 本机受控计时：只报本站 FlashKDA，不伪造本机 FLA 对照 + +协议： + +- deterministic synthetic tensors； +- `D=128`，K3 shape 使用 `H=96, T_total=8192`； +- warmup 20； +- 每轮 100 次、3 repeats，即每个 state mode 300 个 CUDA Event samples； +- output buffer 重用； +- 分开测 BF16 state、无 state I/O、FP32 state； +- throughput 是 sequence tokens/s，不是 head-tokens/s； +- 没有安装并复跑 FLA，所以不计算本站 speedup。 + +| RTX 5090 case | state mode | mean | P95 | sequence tokens/s | +|---|---|---:|---:|---:| +| fixed `[8192]` | BF16 | 2.6210 ms | 2.6437 ms | 3.126 M/s | +| fixed `[8192]` | none | 2.6154 ms | 2.6362 ms | 3.132 M/s | +| fixed `[8192]` | FP32 | 2.6267 ms | 2.6500 ms | 3.119 M/s | +| varlen `[1300,547,2048,963,271,3063]` | BF16 | 2.3335 ms | 2.3574 ms | 3.511 M/s | +| same varlen | none | 2.3104 ms | 2.3254 ms | 3.546 M/s | +| same varlen | FP32 | 2.3596 ms | 2.3845 ms | 3.472 M/s | + +fixed case 三种 mode 合并的 peak allocated 为 1,755.7 MiB；varlen 为 1,845.7 MiB。 +作者 H20 fixed 的 `2.6220 ms` 与本站 BF16 fixed 的 `2.6210 ms` 数值接近只是一次观测， +不能据此宣布 RTX 5090 与 H20 等价：软件栈、频率、功耗、输入、测量时段和对照 kernel +都不足以支持这种外推。 ## 8. 可复现实验入口 -脚本： +checkpoint 工件脚本： ```text experiments/k3/checkpoint_probe.py ``` +FlashKDA 构建与执行： + +```text +experiments/k3/flashkda/Dockerfile +experiments/k3/flashkda/build_reference_helper.py +experiments/k3/flashkda/run_probe.py +experiments/k3/flashkda/README.md +``` + 提交的数据快照： ```text src/data/k3-artifact-snapshot.json +src/data/k3-flashkda-runtime.json ``` -脚本会： +两条实验路径共同完成： 1. 解析 config、index 与 selected headers； 2. 校验小范围字节长度和 SHA-256； 3. 统计真实 KDA / router 参数； 4. 运行明确标注的 synthetic router counterexample； 5. 解析 FlashKDA 官方 H20 / GB200 benchmark； -6. 输出本机环境与构建边界； -7. 不提交原始权重。 +6. 在隔离 CUDA 13.0 / glibc 2.39 环境编译 `sm_120a` wheel； +7. 用官方 torch reference 跑六组 exact-match； +8. 在 RTX 5090 跑 fixed / varlen、三种 state mode 的受控计时； +9. 不提交原始权重或二进制 wheel，只提交 checksum、环境、代码和 JSON 结果。 ## 9. 网站实现合同 @@ -341,7 +417,7 @@ src/data/k3-artifact-snapshot.json 1. **Layer map**：93 层真实 config 条带；显示 KDA/MLA、dense/MoE、AttnRes block。 2. **Tensor anatomy**：checkpoint / expert / vision / MLA tensor shape 与数量。 3. **Parameter audit**：真实 Range statistics，并把 `A_log` mismatch 放在主视区。 -4. **Reproduction boundary**：官方 benchmark、本站构建失败点、synthetic router counterexample。 +4. **Reproduction boundary**：作者 benchmark、本站 exact / latency、构建链、未决形状冲突与 synthetic router counterexample。 每个视图必须显示证据类型： @@ -359,7 +435,9 @@ src/data/k3-artifact-snapshot.json - [x] selected open-weight ranges 做真实参数统计； - [x] 发现并限定 `A_log` shape inconsistency； - [x] FlashKDA RTX 5090 构建尝试留下可复现边界； -- [ ] 使用 CUDA 12.9+ 匹配环境跑 FlashKDA exact correctness； +- [x] 使用 CUDA 13.0 / glibc 2.39 隔离环境产出可校验 `sm_120a` wheel； +- [x] RTX 5090 上六组 FlashKDA / official torch reference exact correctness； +- [x] K3 fixed / varlen shape、三种 state mode 的 1,800 个 CUDA Event samples； - [ ] 取得真实 hidden-state / router load trace； - [ ] 取得可加载的 reduced checkpoint、官方 trace 或多机资源； - [ ] 对 Figure 3 / 4 / 5 做真实数值重绘； diff --git a/scripts/check-k3-browser.mjs b/scripts/check-k3-browser.mjs index 54ec668..765c04a 100644 --- a/scripts/check-k3-browser.mjs +++ b/scripts/check-k3-browser.mjs @@ -221,15 +221,24 @@ const artifacts = await evaluate(`(() => { flash: text("[data-benchmark-flash]"), fla: text("[data-benchmark-fla]"), speedup: text("[data-benchmark-speedup]"), + localMean: text("[data-local-mean]"), + localP95: text("[data-local-p95]"), + localThroughput: text("[data-local-throughput]"), + exactSuite: root.textContent.includes("6 / 6 PASS") && root.textContent.includes("MAX ABS ERROR"), cv: text("[data-router-cv]"), zero: text("[data-router-zero]"), }; input("[data-benchmark-device]", "gb200"); input("[data-benchmark-case]", "Varlen, \\\`seq_lens\\\`=\\\`1024 x 8\\\`"); + input("[data-local-case]", "k3_varlen_shape"); + input("[data-local-state]", "fp32_state"); input("[data-router-mode]", "bias"); const reproductionChanged = { flash: text("[data-benchmark-flash]"), speedup: text("[data-benchmark-speedup]"), + localMean: text("[data-local-mean]"), + localP95: text("[data-local-p95]"), + localMode: text("[data-local-mode-copy]"), cv: text("[data-router-cv]"), zero: text("[data-router-zero]"), }; @@ -253,6 +262,7 @@ await pause(180); await screenshot("/tmp/llm-atlas-k3-lab-desktop.png"); await evaluate(`(() => { document.querySelector("[data-k3-artifact-lab]").scrollIntoView({ block: "start", behavior: "instant" }); + document.querySelector('[data-artifact-tab="reproduction"]')?.click(); window.scrollBy(0, -82); })()`); await pause(180); @@ -297,6 +307,7 @@ await pause(180); await screenshot("/tmp/llm-atlas-k3-mobile.png"); await evaluate(`(() => { document.querySelector("[data-k3-artifact-lab]").scrollIntoView({ block: "start", behavior: "instant" }); + document.querySelector('[data-artifact-tab="reproduction"]')?.click(); window.scrollBy(0, -64); })()`); await pause(180); @@ -334,8 +345,8 @@ if (artifacts.tensors.panel !== "tensors" || artifacts.tensors.groups !== 3 || a if (!artifacts.mla.visible || artifacts.mla.rows !== 5 || !artifacts.mla.has576) failures.push("MLA header shape 视图异常"); if (artifacts.parameterInitial.panel !== "parameters" || artifacts.parameterInitial.shape !== "[128] F32" || !artifacts.parameterInitial.conflict) failures.push("A_log 工件冲突审计异常"); if (artifacts.parameterChanged.shape !== "[96,128] F32" || !artifacts.parameterChanged.count.includes("12,288")) failures.push("真实 dt_bias 参数切换异常"); -if (artifacts.reproductionInitial.panel !== "reproduction" || numeric(artifacts.reproductionInitial.speedup) !== 1.85 || numeric(artifacts.reproductionInitial.cv) < 2) failures.push("FlashKDA H20 或 router 初始探针异常"); -if (numeric(artifacts.reproductionChanged.speedup) !== 3.27 || numeric(artifacts.reproductionChanged.flash) !== 0.7064 || numeric(artifacts.reproductionChanged.cv) <= numeric(artifacts.reproductionInitial.cv) || numeric(artifacts.reproductionChanged.zero) <= numeric(artifacts.reproductionInitial.zero)) failures.push("GB200 benchmark 或 synthetic router counterexample 未更新"); +if (artifacts.reproductionInitial.panel !== "reproduction" || numeric(artifacts.reproductionInitial.speedup) !== 1.85 || numeric(artifacts.reproductionInitial.localMean) < 2.6 || !artifacts.reproductionInitial.exactSuite || numeric(artifacts.reproductionInitial.cv) < 2) failures.push("FlashKDA H20、本机 exact suite 或 router 初始探针异常"); +if (numeric(artifacts.reproductionChanged.speedup) !== 3.27 || numeric(artifacts.reproductionChanged.flash) !== 0.7064 || numeric(artifacts.reproductionChanged.localMean) >= numeric(artifacts.reproductionInitial.localMean) || !artifacts.reproductionChanged.localMode.includes("FP32 state") || numeric(artifacts.reproductionChanged.cv) <= numeric(artifacts.reproductionInitial.cv) || numeric(artifacts.reproductionChanged.zero) <= numeric(artifacts.reproductionInitial.zero)) failures.push("GB200 benchmark、本机 varlen/state 或 synthetic router counterexample 未更新"); if (artifacts.keyboardSelected !== "tensors" || artifacts.keyboardVisible !== "tensors") failures.push("开放工件键盘 tab 导航异常"); if (!mobile.menuVisible || mobile.menuOpen !== "true" || mobile.tabs !== 8 || mobile.artifactTabs !== 4 || mobile.artifactLayers !== 93) failures.push("移动端导航或实验异常"); if (mobile.offenders.length) failures.push(`移动端越界元素：${JSON.stringify(mobile.offenders)}`); diff --git a/src/components/K3ArtifactLab.astro b/src/components/K3ArtifactLab.astro index 3465d2d..38e5c2b 100644 --- a/src/components/K3ArtifactLab.astro +++ b/src/components/K3ArtifactLab.astro @@ -3,6 +3,7 @@ import { k3ArtifactEvidence, k3ArtifactLayers, k3ArtifactSnapshot as snapshot, + k3FlashKdaRuntime as runtime, k3ArtifactViews, } from "@/data/k3Artifacts"; @@ -10,6 +11,9 @@ const checkpoint = snapshot.checkpoint; const audit = snapshot.parameter_audit; const probe = snapshot.router_stress_probe; const flash = snapshot.flashkda; +const localBenchmarks = runtime.benchmarks.filter((row) => row.name.startsWith("k3_")); +const localInitial = localBenchmarks[0]; +const localInitialTiming = localInitial.timings.bf16_state; const bytes = (value: number) => { if (value >= 2 ** 40) return `${(value / 2 ** 40).toFixed(3)} TiB`; @@ -263,8 +267,8 @@ const benchmarkDevices = [