tokenspeed-situ 0.1.0.post20260727

AOT TensorRT-LLM-Gen MXFP4+SiTU MoE kernel bundle

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:

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:

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:

import tokenspeed_situ as situ

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

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

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:

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