K2-Horizon-7B-MXFP8

MXFP8-quantized version of IFM/K2-Horizon-7B, targeting NVIDIA Blackwell GPUs.

This repository provides an MXFP8 conversion of K2-Horizon-7B together with the custom runtime used for optimized batch-1 inference.

Important — download or clone the complete repository to use the optimized inference path. Do not download only the .safetensors checkpoint. inference.py also requires the included configuration, tokenizer, model architecture, chat template, MXFP8 GEMV, fusion, and FlashAttention helper files to remain beside it.

Quantization

The model weights were converted to MXFP8 using:

• E4M3 FP8 values
• E8M0 block scaling
• Blackwell-compatible block-scaled matrix multiplication
• SafeTensors checkpoint format

The converted model contains 254 MXFP8 modules. The optimized runtime replaces 253 linear modules with its batch-1 M=1 MXFP8 GEMV path and loads the embedding separately.

Base Model

Base model: IFM/K2-Horizon-7B

Please refer to the original model card for architecture details, training information, intended use, and limitations.

Optimized Inference

The primary portable entry point is:

inference.py

It provides BF16 KV cache, FlashAttention 2, custom M=1 MXFP8 GEMV, CUDA-graph decode, GPU top-p sampling, normal EOS-aware generation, and readable assistant output.

Interactive mode:

python inference.py --context 16384

One prompt:

python inference.py --context 16384 --max-new-tokens 4096 --prompt "Your prompt here"

An experimental PyTorch 2.14 multi-turn runner is also included:

python inference_multiturn_gpu_while.py

This version captures decode in a CUDA conditional WHILE graph. EOS evaluation and loop termination remain on the GPU, with one host synchronization after each assistant turn. It keeps conversation history, uses a 16,384-token context, and supports /clear and /exit.

PyTorch 2.14 currently rejects RNG operations inside CUDA conditional nodes. The multi-turn runner therefore generates uniform samples on the GPU before graph replay and consumes them with an inverse-CDF top-p sampler inside the loop. It preserves the requested categorical/top-p distribution, but seeded output is not token-for-token identical to torch.multinomial.

The conditional-WHILE runner imports a private PyTorch 2.14 control-flow API. Treat it as version-pinned experimental code; inference.py remains the recommended general entry point.

Installation

Validated requirements:

• NVIDIA Blackwell GPU supported by the runtime, currently SM120 or newer
• Python 3.13
• PyTorch 2.14.0+cu130
• CUDA 13.0
• Transformers 5.15.0
• Comfy-Kitchen 0.2.35
• Triton 3.8
• FlashAttention 2.8.3 with flash_attn_with_kvcache
• accelerate and safetensors

Install the CUDA 13 PyTorch build appropriate for your platform first, then install the remaining packages:

python -m pip install -r requirements.txt

PyTorch and FlashAttention wheels are platform-, Python-, GPU-, and CUDA-specific. On Windows, a compatible FlashAttention build may need to be installed manually if pip cannot provide a suitable wheel.

All runtime paths are resolved relative to the downloaded repository. There are no user-directory or drive-letter dependencies. The checkpoint is discovered automatically when exactly one .safetensors file is present. K2_MODEL_DIR and K2_CHECKPOINT may be used as explicit overrides.

Required Files

Keep the following files together in the downloaded model directory:

• K2-Horizon-7B-MXFP8.safetensors or one compatible .safetensors checkpoint
• config.json
• generation_config.json
• configuration_k2_horizon.py
• modeling_k2_horizon.py
• tokenizer.json
• tokenizer_config.json
• chat_template.jinja
• inference.py
• k2_mxfp8_engine_portable.py
• k2_fa2_bf16_backend_portable.py
• k2_mxfp8_fusions_portable.py
• mxfp8_gemv_portable.py

For GPU-controlled multi-turn inference, also keep inference_multiturn_gpu_while.py.

Inference Performance

Real-generation throughput was measured over a roughly 16K-token run using identical sampling settings.

Decode KV / Attention Backend	Beginning (~79–207 ctx)	Midpoint (~8.1K ctx)	Near 16K	Full-run Avg	KV Storage	Peak PyTorch VRAM
FP8 KV + Triton	52.38 tok/s	40.94 tok/s	32.21 tok/s	39.88 tok/s	1.143 GiB	9.951 GiB
BF16 KV + Triton	50.14 tok/s	41.85 tok/s	37.85 tok/s	42.72 tok/s	2.250 GiB	11.090 GiB
BF16 KV + FlashAttention 2	50.24 tok/s	44.35 tok/s	39.99 tok/s	43.26 tok/s	2.250 GiB	11.090 GiB

BF16 KV + FlashAttention 2 provides the best sustained throughput up to approximately 16K context.

FP8 KV uses roughly half the KV-cache memory and is preferable when VRAM is limited or when targeting longer contexts such as 24K–32K.

Asynchronous No-Sync Benchmark

A fixed-length benchmark measured the same 7B MXFP8/BF16-KV/FA2 path without any CPU synchronization inside the decode loop. Generated IDs were copied asynchronously into pinned system RAM and consumed only after the final synchronization.

Measured result:

Context	Prompt	Generated	Start	Average	Median	Last	Decode-loop CPU synchronizations
16,384	55 tokens	16,329 tokens	58.03 tok/s	48.39 tok/s	47.59 tok/s	43.90 tok/s	0

Additional measurements:

• TTFT: 49.74 ms
• Timed decode tokens: 16,328
• GPU decode time: 337.445 seconds
• Wall-clock decode average: 48.39 tok/s
• Timing window: 128 tokens
• Token storage: pinned CPU RAM using asynchronous device-to-host copies

This is a fixed-length throughput benchmark: it intentionally does not perform per-token EOS checks. It should not be confused with normal interactive generation. The host still submits one-token graph replays; it simply never waits for an individual token during the measured loop.

The multi-turn CUDA-WHILE runner is architecturally different: the CPU submits the conditional decode graph once per assistant turn and the GPU itself stops on EOS. Its TTFT includes prefill of the retained conversation history, so later turns can have substantially higher TTFT even when decode throughput remains similar.

Test Stack

• GPU: NVIDIA GeForce RTX 5070 Ti Laptop GPU
• Architecture: Blackwell / SM120
• VRAM: 12 GB (11.94 GiB)
• Python: 3.13
• PyTorch: 2.14.0+cu130
• CUDA: 13.0
• Transformers: 5.15.0
• Comfy-Kitchen: 0.2.35
• Triton-Windows: 3.8.0.post28
• FlashAttention: 2.8.3

Inference uses:

• MXFP8 weights with E4M3 data and E8M0 block scales
• Custom M=1 MXFP8 GEMV
• CUDA Graph decode
• GPU top-p sampling
• Static KV-cache allocation
• Temperature 1.0
• Top-p 0.95
• High reasoning effort
• Batch size 1

Measurement Notes

The backend comparison used real autoregressive generation from a 79-token prompt with 16,000 generated tokens and a 16,384-token KV-cache capacity.

The beginning, midpoint, and near-16K results are 128-token rolling windows. Full-run average measures steady-state sampled decode throughput.

FP8 KV stores E4M3FN K/V values with BF16 scaling, while the BF16 backends store K/V directly in BF16.

All backend-comparison runs used the same generation settings and seed, although sampled token trajectories are not bit-identical across attention backends because of numerical differences in attention and KV precision.

Acknowledgements

All model architecture, training, and base-model credit belongs to the authors of K2-Horizon-7B.

This repository provides an MXFP8 conversion and custom inference runtime for the original model.