wuyang/llm-atlas
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: execute FlashKDA on RTX 5090

Browse Source
This commit is contained in:
wuyang
2026-07-29 13:36:49 +08:00
parent 1146208b5e
commit 2ef846f751
16 changed files with 1189 additions and 51 deletions
Download Patch File Download Diff File Expand all files Collapse all files
experiments/k3/flashkda/Dockerfile
+30
View File
@@ -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
experiments/k3/flashkda/README.md
+85
View File
@@ -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.
experiments/k3/flashkda/build_reference_helper.py
+50
View File
@@ -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 <torch/extension.h>
#include <cuda_runtime.h>
__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<float>(), output.data_ptr<float>(), 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()
experiments/k3/flashkda/run_probe.py
+454
View File
@@ -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()