wuyang
86 lines
2.8 KiB
Markdown
86 lines
2.8 KiB
Markdown
# 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.
|