# 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.