llm-atlas/experiments/k3/flashkda

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:

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:

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:

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:

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:

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.