Metadata-Version: 2.4
Name: tokenspeed-situ
Version: 0.1.0.post20260727
Summary: AOT TensorRT-LLM-Gen MXFP4+SiTU MoE kernel bundle
License-Expression: LicenseRef-Proprietary
Project-URL: Homepage, https://github.com/lightseekorg/tokenspeed
Classifier: Development Status :: 2 - Pre-Alpha
Classifier: Operating System :: POSIX :: Linux
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3 :: Only
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
License-File: THIRD_PARTY_NOTICES.md
Requires-Dist: apache-tvm-ffi<0.2,>=0.1.12
Dynamic: license-file

# TokenSpeed SiTU

`tokenspeed-situ` is the sidecar distribution for the
TensorRT-LLM-Gen MXFP4 + SiTU MoE path used by Kimi K3 on NVIDIA B300
(`sm103a`, CUDA 13). Its Python import package is `tokenspeed_situ`.
It deliberately does **not** install, replace, or pin the public
`flashinfer-python` package.

The bundled cubins are not covered by the wrapper's MIT license. Do not
publish or redistribute the wheel until the required NVIDIA and model-owner
permissions and export-control review are complete; see
`THIRD_PARTY_NOTICES.md`.
Source-tree identity and validation status are recorded in `PROVENANCE.md`.

The wheel contains:

- one release-mode TVM-FFI runner shared library;
- **267 cubins**: the bf16-activation (w4a16) chain — 30 FC1 MXFP4 + SiTU +
  30 matching FC2, `sm103a` — plus the **w4a8** (MXFP8-activation) chain —
  the full dtype family of 106 FC1 (`MxE4m3_MxE2m1MxE4m3` siTuGlu) + 101 FC2
  (`Bfloat16_MxE2m1MxE4m3`), `sm100f` (B300-loadable);
- a manifest with SHA-256 checksums and source/build provenance;
- a lazy loader and the low-level FFI entry points.

The current `0.1.0.post20260727` artifact uses a runner compiled from the
complete 1,696-entry metadata table. The w4a8
chain ships the **whole** dtype family (not a per-shape-captured subset) so any
parallel config (EP8, TP8, other intermediate sizes) is covered with no runtime
`BundleError` on an untested shape — the runner only ever selects tiles within
the family. A release candidate should be GPU-tested on B300 before publication.

## Install

Install the current bundle from PyPI:

```bash
python -m pip install "tokenspeed-situ==0.1.0.post20260727"
python -c \
  'import tokenspeed_situ as s; print(s.verify_bundle())'   # cubin_count 267
```

## Activation precision (w4a16 vs w4a8)

The wheel serves both MoE activation precisions; the runtime selects via
`--moe-activation-dtype`:

- `bf16` (default): bf16 activation × MXFP4 weight (w4a16) — today's path.
- `mxfp8`: MXFP8-quantized activation × MXFP4 weight (w4a8) — a run-wide
  deployment choice (stable for the whole process, not per batch) that speeds
  up large-batch prefill on the MoE at a small decode cost. The runner-selected
  chain is FC1-w4a8 + FC2-w4a16 (the FC1 fp8 output is dequantized to bf16
  before FC2). Weights stay MXFP4 either way.

The current TokenSpeed monorepo has not yet switched to this standalone import
name, so do not use this wheel as a drop-in replacement until the integration
changes listed below land. After that switch, install this wheel before
importing `tokenspeed_kernel` or launching a Kimi-K3 process. The wheel only
adds the K3 SiTU runner; it does not install, replace, or pin FlashInfer, so
TokenSpeed can keep using its public
`flashinfer-python==0.6.15` dependency.

A clean checkout includes the reviewed runner, manifest, license files, and
all 267 cubins. Maintainers can build and verify the wheel directly:

```bash
python -m pip install auditwheel==6.7.0 build patchelf==0.17.2.4
python -m build --wheel --outdir wheelhouse/raw
python -m auditwheel repair \
  --only-plat \
  --exclude 'libcuda.so.*' \
  --exclude 'libcudart.so.*' \
  --plat manylinux_2_39_x86_64 \
  --wheel-dir dist \
  wheelhouse/raw/*.whl
python tools/verify_wheel.py dist/*.whl
```

## Supported contract

- activation: BF16 for w4a16, or MXFP8 for w4a8 with a per-token
  `hidden_states_scale` block-scale tensor;
- weights: MXFP4 (`MxE2m1`);
- activation function: SiTU (`9`);
- experts: 896 total, with balanced contiguous local expert shards;
- hidden size: 3584;
- MoE TP8/EP1: 896 local experts, physical intermediate size 384 (native, 128-aligned)
  (logical shard size 384);
- MoE TP1/EP8: 112 local experts per rank, physical intermediate size 3072,
  with `local_expert_offset = ep_rank * 112`;
- top-k: 16;
- routing: precomputed packed or unpacked top-k;
- architecture/toolkit: `sm103a`, CUDA 13.

Top-k expert IDs remain global (`0..895`). The native routing kernel filters
out experts outside the current rank's contiguous shard and produces a partial
output. The caller must all-reduce those partial outputs across the MoE TP/EP
group. The bundle does not own or initialize distributed communication, and
mixed MoE tensor parallelism plus expert parallelism is not currently a
supported TokenSpeed topology.

## Runtime API

Importing the package does not load CUDA, the runner, or FlashInfer:

```python
import tokenspeed_situ as situ

print(situ.bundle_info())
situ.verify_bundle()
```

The native module is loaded only when an operation is requested:

```python
from tokenspeed_situ import get_valid_k3_situ_configs

tactics = get_valid_k3_situ_configs(
    num_tokens=128,
    intermediate_size=3072,
    local_num_experts=112,
)
```

`mxfp4_situ_routed_moe` provides a narrow keyword-only wrapper for K3 TP and
contiguous EP shards. Callers pass the physical `intermediate_size`, global
expert count, local expert offset, and local expert count. The default tactic
`(-1, -1)` asks the native runner for its first valid configuration for the
runtime shape; it performs no Python benchmarking or autotuning.

`trtllm_fp4_block_scale_moe_raw` also exposes the complete low-level TVM-FFI ABI.
The existing TokenSpeed adapter owns tensor preparation and the output
all-reduce; this package owns only the runner and cubin callback.

## TokenSpeed integration

The source monorepo still uses its original in-tree sidecar name. When it opts
into this standalone distribution, update the import and installation guidance
in `tokenspeed-kernel/.../trtllm_mxfp4.py`,
`python/tokenspeed/runtime/models/kimi_k3.py`,
`test/runtime/test_kimi_k3_moe_backend.py`, `docs/recipes/models.md`, and the
pre-commit large-runner exclusion together.

## Rebuild the wheel (maintainers only)

To refresh the reviewed binary assets, stage a new payload and then build:

```bash
python tools/stage_assets.py \
  --snapshot-root /path/to/flashinfer-v0613rc1-situ \
  --pool-root /path/to/20260617_v0613rc1_situ_v0611_barrierfix \
  --module /path/to/fused_moe_trtllm_sm100_private_hidden_ffi_lean_a61443e1.so
python -m build --wheel --outdir wheelhouse/raw
python -m auditwheel repair \
  --only-plat \
  --exclude 'libcuda.so.*' \
  --exclude 'libcudart.so.*' \
  --plat manylinux_2_39_x86_64 \
  --wheel-dir dist \
  wheelhouse/raw/*.whl
python tools/verify_wheel.py dist/*.whl
```
