K2-Horizon-3.7B-MXFP8

MXFP8 conversion of IFM/K2-Horizon-3.7B, packaged with an optimized batch-1 inference runtime for NVIDIA Blackwell GPUs.

Clone or download the complete repository. The checkpoint is not a standalone Transformers weight file. The included model code and custom MXFP8 runtime are required for inference.

Quantization

• E4M3 FP8 weights
• E8M0 block scales stored as uint8
• Group size: 32
• Comfy-Kitchen / TensorCore MXFP8 blocked layout
• 253 quantized modules
  • 252 transformer linear layers
  • 1 embedding layer
• BF16 output head and remaining non-quantized parameters
• SafeTensors checkpoint

A generic transformers.pipeline() cannot directly execute the MXFP8 side tensors.

---

Benchmarks

Test system:

• NVIDIA GeForce RTX 5070 Ti Laptop GPU
• Blackwell SM120
• 12 GB VRAM (11.94 GiB)
• Python 3.13
• PyTorch 2.14.0+cu130
• CUDA 13.0
• FlashAttention 2.8.3
• batch size 1
• one-token CUDA Graph decode

0–32K Decode

Backend	Start tok/s	~16K tok/s	~32K tok/s	Average tok/s	KV GiB	Peak GiB
FP8 KV + Triton	85.17	38.75	25.05	38.72	2.285	7.958
BF16 KV + PyTorch SDPA	10.82	10.84	10.84	10.84	4.500	10.173
BF16 KV + cuDNN SDPA	51.80	48.19	48.55	49.23	4.500	10.173
BF16 KV + FlashAttention 2	84.17	57.29	46.57	54.99	4.500	10.173

FlashAttention 2 gives the best overall throughput across low and moderate context lengths.

cuDNN SDPA scales significantly better at very long context and overtakes FlashAttention 2 near the upper end of the tested range.

FP8 KV cuts cache memory roughly in half, but the experimental Triton FP8 attention backend currently loses substantial throughput as context grows.

32K–40K Long-Context Decode

Backend	Prefill tok/s	Start @32K	Mid	Last @40K	Average tok/s	Median tok/s	KV GiB
BF16 KV + FlashAttention 2	135.14	46.46	44.94	43.85	45.03	44.91	5.625
BF16 KV + cuDNN SDPA	133.20	46.67	45.90	46.00	46.14	46.15	5.625

The benchmark suggests two useful operating regimes:

• FlashAttention 2: best for low-to-moderate context and highest overall decode throughput.
• cuDNN SDPA: stronger throughput stability at very long context.

The public inference runner currently uses BF16 KV + FlashAttention 2.

---

Inference

inference.py is the recommended entry point.

It uses:

• MXFP8 model weights
• custom Triton M=1 MXFP8 GEMV for autoregressive decode
• Comfy-Kitchen MXFP8 kernels for larger matrix shapes
• BF16 static KV cache
• FlashAttention 2 flash_attn_with_kvcache
• one-token CUDA Graph replay
• GPU temperature/top-p sampling
• EOS-aware generation

Interactive mode

python inference.py --context 16384

Each prompt is independent. Previous prompts are not retained.

Single request

python inference.py --context 16384 --max-new-tokens 4096 --prompt "Explain MXFP8."

Default generation settings:

• temperature=1.0
• top_p=0.95
• reasoning_effort=high

The response is streamed during generation and decode-speed statistics are printed afterward.

---

Multi-turn Inference

inference_multiturn_gpu_while.py maintains conversation history and uses a PyTorch 2.14 CUDA conditional-WHILE graph.

python inference_multiturn_gpu_while.py --max-new-tokens 4096

Commands:

• /clear — clear conversation history
• /exit — exit

The multi-turn runner uses a fixed 16,384-token context.

When history becomes too large, the oldest complete user/assistant pairs are removed. The retained conversation is re-prefilled on every turn; KV state is not reused between turns.

Decode, sampling, EOS detection, token writes, and loop control execute on the GPU, with no per-token host synchronization.

Unlike inference.py, the completed assistant response is displayed after the GPU loop finishes rather than streamed token-by-token.

This runner depends on an experimental PyTorch 2.14 control-flow API. Use inference.py for the more conventional inference path.

---

Installation

Validated environment:

• NVIDIA Blackwell GPU
• Windows
• 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

Install an appropriate CUDA 13 PyTorch build first.

PyTorch is intentionally not pinned in requirements.txt because the correct package depends on the operating system and CUDA installation.

Then install the remaining dependencies:

python -m pip install -r requirements.txt

On Windows, FlashAttention may require a compatible prebuilt wheel matching Python, PyTorch, CUDA, and GPU architecture.

---

Repository Files

model.safetensors
config.json
generation_config.json
configuration_k2_horizon.py
modeling_k2_horizon.py

tokenizer.json
tokenizer_config.json
chat_template.jinja

inference.py
inference_multiturn_gpu_while.py

k2_mxfp8_engine_portable.py
k2_fa2_bf16_backend_portable.py
k2_mxfp8_fusions_portable.py
mxfp8_gemv_portable.py

requirements.txt
README.md

Keep these files together in the repository/model directory.

---

Runtime Design

The optimized inference path works as follows:

1. The K2 architecture is created on the meta device.
2. MXFP8 tensors are loaded module-by-module from the SafeTensors checkpoint.
3. Standard linears and the embedding layer are replaced by custom MXFP8 modules.
4. Batch-1 M=1 linear operations use the specialized Triton MXFP8 GEMV kernel.
5. Larger matrix shapes use Comfy-Kitchen MXFP8 kernels.
6. A BF16 KV cache is allocated for the requested context length.
7. Attention modules are patched to use the FlashAttention 2 KV-cache backend.
8. The prompt is prefetched.
9. A one-token CUDA Graph is captured and replayed during autoregressive decode.

KV-cache memory is approximately:

• 16K BF16: 2.25 GiB
• 32K BF16: 4.50 GiB
• 40K BF16: 5.625 GiB

Long prefills can require substantially more temporary memory than steady-state decode.

---

Acknowledgements

The original K2-Horizon architecture, training, model behavior, intended-use guidance, and base-model limitations belong to the K2-Horizon authors.

This repository provides the MXFP8 conversion and optimized inference runtime.