From 1b3732a6ed87ea3c537dc073a9fc812d7e46d237 Mon Sep 17 00:00:00 2001 From: George Hotz Date: Mon, 10 Aug 2026 01:15:26 -0700 Subject: [PATCH] speed up Kimi K3 TP8 loading on MI350X --- docs/kimi_k3_mi350.md | 37 +++++++++ extra/benchmark_kimi_k3.py | 82 +++++++++++++++++++ test/unit/test_attention.py | 17 ++++ test/unit/test_llm_k3.py | 24 +++++- tinygrad/llm/kimi_k3.py | 154 ++++++++++++++++++++++++++---------- tinygrad/llm/model.py | 9 ++- 6 files changed, 274 insertions(+), 49 deletions(-) create mode 100644 extra/benchmark_kimi_k3.py diff --git a/docs/kimi_k3_mi350.md b/docs/kimi_k3_mi350.md index e5e41ded74..c9ea1cfede 100644 --- a/docs/kimi_k3_mi350.md +++ b/docs/kimi_k3_mi350.md @@ -76,6 +76,43 @@ DEV=AMD DEBUG=1 python -m tinygrad.llm.cli --model /raid/weights/kimi-k3 \ --devices 8 --max_context 4096 --warmup --benchmark 20 ``` +## MI350X validation results (2026-08-10) + +The official directory was audited in place: 96 shards, 497,220 indexed tensors, 497,052 language tensors, and 1,560,860,324,864 total bytes. All eight devices reported `gfx950`. No checkpoint file was converted, copied, or modified, and every model run used a single process. The actual text tower is 1,559,965,606,912 bytes; its checked TP8 layout is 196,784,397,312 bytes per GPU. + +The preserved first full-checkpoint error was an `A_log` shape mismatch, `(128,) -> (96, 1)`. K3 stores one decay value per 128-wide KDA channel, not one per head. The loader now keeps this field replicated and applies the official channel-wise broadcast. A numerical unit test covers the distinction from the older head-wise Kimi Linear behavior. + +Load speed was fixed before generation. The original loader opened thousands of individual expert tensors and independently realized eight strided TP slices. The MI350 path now does the following without changing the checkpoint: + +- copies contiguous axis-zero shards and replicas directly into their final device buffers; +- stages an inner-axis tensor once and schedules all eight TP slices together; +- reads each layer's contiguous 15.72 GB expert region once, reorders its lexicographically stored expert records on GPU 0, and realizes all six packed/scale destinations together; +- retains only final MultiBuffer identities, drops the reorder graph, and flushes the 15.72 GB staging allocation before the next layer. + +One real expert layer leaves exactly 1,965,293,568 bytes resident on each GPU and zero bytes in the GPU-0 allocator cache. Complete context-128 loads measured 527.20 seconds before the final staging cleanup and 490.05/489.59 seconds afterward. Peak host RSS for the unprofiled correctness run was 2.11 GiB with zero swap. RAID variability produced later loads from 489.06 to 532.85 seconds. + +The fixed XTML prompt `Reply with exactly: OK` encodes to 93 tokens. After excluding the cold JIT capture from replay comparison, two greedy runs produced the identical eight-token sequence: + +```text +[9545, 59991, 10580, 14404, 9545, 59991, 9545, 59991] +``` + +At context 128, steady prefill was 14.32 seconds (6.49 tok/s) and eight-token decode was 2.27 seconds (3.53 tok/s, 283.3 ms/token). The same first tokens remained stable at every admitted context. These rates are much lower than the planning estimates below and should be treated as the current measured baseline. + +| Maximum context | Load | Short-prompt replay | Result | +|---:|---:|---:|---| +| 128 | 489.59s | 14.32s | stable 8-token replay | +| 4,096 | 489.06s | 14.32s | stable replay, zero swap | +| 32,768 | 532.85s | 14.33s | stable replay, zero swap | +| 131,072 | 520.91s | 14.37s | stable first token, zero swap | +| 262,144 | 497.34s | 14.41s | stable first token, zero swap | + +These are maximum-context/cache admission tests with the same 93-token prompt, not full-length 32K/131K/262K prefills. The full cache allocation path was exercised, but filling those contexts remains a separate long-running throughput test. + +Runtime profiling bracketed four steady decode tokens. It recorded 6,304 kernel events and 477.97 ms of summed GPU work across the eight devices inside a 1,573.99 ms profiled wall interval. The packed `mxfp4_expert_linear_wave64` kernels accounted for only 22.47 ms summed; the largest families were small 1,792-wide reductions. This identifies launch/synchronization granularity as the immediate MI350 bottleneck rather than packed-weight bandwidth. `JIT_BATCH_SIZE=64` produced the same 3.53 tok/s as 32. A gfx950 fused MXFP8 QDQ experiment was bit-exact but slower on the real device (about 95 microseconds versus 57--64 microseconds), so it was rejected. The next retained performance work should fuse the reduction/collective boundaries called out below. + +The checkpoint's bundled Transformers code was used as the architectural reference for channel decay and tensor mapping. A full independent Transformers/vLLM token comparison was not run on this host because the required `compressed_tensors`/serving backend is not installed; deterministic tinygrad replay and the numerical KDA, loader-layout, NULL gfx950 compile, and real TP8 smoke tests are the completed correctness gates. + ## Known hardware-only gate The correctness path now consumes packed MXFP4 expert weights directly on gfx950 with a wave64 software-decode kernel, so it does not create selected-expert BF16 weight expansions. MXFP8 activation quantization is still emulated. tinygrad has gfx950/CDNA4 BF16 and FP8 matrix-core support, but this branch does not yet have a hardware-validated native MXFP4×MXFP8 expert GEMM. Expect the first run to be a correctness bring-up, not production throughput. Capture profiles on MI350X before changing the representation: native FP4 work cannot be validated faithfully on the available gfx1100 cards. diff --git a/extra/benchmark_kimi_k3.py b/extra/benchmark_kimi_k3.py new file mode 100644 index 0000000000..f812ed5da1 --- /dev/null +++ b/extra/benchmark_kimi_k3.py @@ -0,0 +1,82 @@ +#!/usr/bin/env python3 +"""Bounded correctness and load/prefill/decode benchmark for the official TP8 Kimi K3 checkpoint.""" +import argparse, resource, time + +from tinygrad import Device +from tinygrad.helpers import profile_marker +from tinygrad.llm.cli import KimiK3Template, SimpleTokenizer +from tinygrad.llm.kimi_k3 import load_kimi_k3, load_kimi_tokenizer_data + +def sync(devices:int) -> None: + for i in range(devices): Device[f"AMD:{i}"].synchronize() + +def fresh_generate(model, prompt:list[int], chunk_size:int): + # Never reuse a prefix or recurrent state across correctness/benchmark trials. + model._cached_tokens = [-1] * len(prompt) + return model.generate(prompt.copy(), chunk_size=chunk_size, temperature=0.0) + +def timed_next(gen, devices:int) -> tuple[int, float]: + begin = time.perf_counter() + token = next(gen) + sync(devices) + return token, time.perf_counter()-begin + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("model", help="official unmodified Kimi K3 checkpoint directory") + parser.add_argument("--devices", type=int, default=8) + parser.add_argument("--max-context", type=int, default=128) + parser.add_argument("--prompt", default="Reply with exactly: OK") + parser.add_argument("--stable-tokens", type=int, default=8) + parser.add_argument("--decode-tokens", type=int, default=8) + parser.add_argument("--chunk-size", type=int, default=8) + args = parser.parse_args() + + begin = time.perf_counter() + model = load_kimi_k3(args.model, max_context=args.max_context, devices=args.devices) + sync(args.devices) + load_time = time.perf_counter()-begin + print(f"load: {load_time:.3f}s", flush=True) + + normal, special, bos, eos = load_kimi_tokenizer_data(args.model) + tok = SimpleTokenizer(normal, special, "kimi-k2", bos_id=bos, eos_id=eos, eot_id=eos) + rendered = KimiK3Template().render(messages=[{"role":"user", "content":args.prompt}], add_generation_prompt=True) + prompt = tok.encode(rendered) + needed = len(prompt) + max(args.stable_tokens, args.decode_tokens+3) + if needed > args.max_context: raise ValueError(f"prompt and output need {needed} tokens but max context is {args.max_context}") + print(f"prompt: {len(prompt)} tokens, chunk={args.chunk_size}", flush=True) + + sequences:list[list[int]] = [] + # The first execution captures the prefill and rollout JITs. Correctness comparisons must use + # identical replay paths, rather than comparing compilation/capture numerics to replay numerics. + for trial in range(3): + gen = fresh_generate(model, prompt, args.chunk_size) + sequence:list[int] = [] + prefill = 0.0 + for step in range(args.stable_tokens): + token, elapsed = timed_next(gen, args.devices) + sequence.append(token) + if step == 0: prefill = elapsed + if trial: sequences.append(sequence) + print(f"{'capture warmup' if trial == 0 else f'stable trial {trial}'}: prefill={prefill:.3f}s " + f"({len(prompt)/prefill:.3f} tok/s), tokens={sequence}", flush=True) + if sequences[0] != sequences[1]: raise RuntimeError(f"greedy output is not repeatable: {sequences}") + print(f"stable text: {tok.decode(sequences[0])!r}", flush=True) + + gen = fresh_generate(model, prompt, args.chunk_size) + profile_marker("kimi k3 steady prefill start") + first, prefill = timed_next(gen, args.devices) + profile_marker("kimi k3 steady prefill end") + warmup = [timed_next(gen, args.devices)[0] for _ in range(2)] + profile_marker("kimi k3 steady decode start") + begin = time.perf_counter() + output = [next(gen) for _ in range(args.decode_tokens)] + sync(args.devices) + decode = time.perf_counter()-begin + profile_marker("kimi k3 steady decode end") + print(f"prefill replay: {prefill:.3f}s ({len(prompt)/prefill:.3f} tok/s), token={first}", flush=True) + print(f"decode after warmup {warmup}: {decode:.3f}s ({args.decode_tokens/decode:.3f} tok/s, " + f"{decode/args.decode_tokens*1e3:.3f} ms/tok), output={output}", flush=True) + print(f"peak RSS: {resource.getrusage(resource.RUSAGE_SELF).ru_maxrss/1024:.1f} MiB", flush=True) + +if __name__ == "__main__": main() diff --git a/test/unit/test_attention.py b/test/unit/test_attention.py index 1fe564a81e..da8f63f0e3 100644 --- a/test/unit/test_attention.py +++ b/test/unit/test_attention.py @@ -214,6 +214,23 @@ class TestGatedDeltaNetBlock(unittest.TestCase): alpha = np.exp(-5.0 / (1.0 + np.exp(-(exp_a * gate_logits)))).reshape(1, 2, 1, 2) np.testing.assert_allclose(block.recurrent_state.numpy(), initial_state.numpy() * alpha, rtol=2e-5, atol=2e-5) + def test_kda_per_channel_a(self): + config = self._make_config(n_heads=2, kda_full_rank_gate=True, kda_gate_lower_bound=-5.0, + ssm=SSMConfig(conv_kernel=2, state_size=2, group_count=2, time_step_rank=2, inner_size=4, kda=True, channel_decay=True)) + block, x = GatedDeltaNetBlock(config, config.ssm), Tensor([[[1., 2., 0., 0.]]]) + block.ssm_f_a.weight = Tensor([[1., 0., 0., 0.], [0., 1., 0., 0.]]) + block.ssm_f_b.weight = Tensor([[1., 0.], [0., 1.], [1., 1.], [2., 1.]]) + block.ssm_dt["bias"] = Tensor.zeros(4) + block.ssm_a = Tensor([[-2.], [-3.]]) + block._init_state(x) + initial_state = Tensor.arange(8, dtype=dtypes.float32).reshape(1, 2, 2, 2) + block.recurrent_state.assign(initial_state).realize() + block._attention(x, 0).realize() + gate_logits = np.arange(1, 5, dtype=np.float32).reshape(1, 2, 2) + exp_a = np.array([2., 3.], dtype=np.float32).reshape(1, 1, 2) + alpha = np.exp(-5.0 / (1.0 + np.exp(-(exp_a * gate_logits)))).reshape(1, 2, 1, 2) + np.testing.assert_allclose(block.recurrent_state.numpy(), initial_state.numpy() * alpha, rtol=2e-5, atol=2e-5) + def test_kda_chunked_prefill_matches_decode(self): config = self._make_config(max_context=4, n_heads=2, ssm=SSMConfig(conv_kernel=2, state_size=2, group_count=2, time_step_rank=2, inner_size=4, kda=True), kda_split_qkv=True) diff --git a/test/unit/test_llm_k3.py b/test/unit/test_llm_k3.py index b0ef3c68e6..0f826627f1 100644 --- a/test/unit/test_llm_k3.py +++ b/test/unit/test_llm_k3.py @@ -3,12 +3,17 @@ from dataclasses import replace import numpy as np from tinygrad import Tensor, dtypes, nn from tinygrad.llm.kimi_k3 import KIMI_K3_FULL_ATTN_LAYERS, KIMI_K3_SSM_LAYERS, KIMI_K3_TEXT_SIZE, KIMI_K3_TP8_BYTES_PER_GPU, \ - _layer_sources, _load_stacked_experts, _shard_kimi_k3, _validate_config, kimi_k3_config, kimi_k3_smoke_config + _layer_sources, _load_stacked_experts, _replace, _shard_kimi_k3, _validate_config, kimi_k3_config, kimi_k3_smoke_config from tinygrad.llm.model import FFNBlock, Transformer def small_k3_config(max_context:int=4): return replace(kimi_k3_smoke_config(max_context), num_experts=8) class TestKimiK3(unittest.TestCase): + def test_smoke_config_preserves_gfx950_expert_alignment(self): + c = kimi_k3_smoke_config() + self.assertEqual(c.routed_expert_dim % 64, 0) + self.assertEqual((c.hidden_dim // 8) % 64, 0) + def test_official_config(self): c = kimi_k3_config(1_048_576) self.assertEqual((c.num_blocks, c.dim, c.n_heads, c.num_experts, c.num_experts_per_tok), (93, 7168, 96, 896, 16)) @@ -16,6 +21,7 @@ class TestKimiK3(unittest.TestCase): self.assertEqual(KIMI_K3_FULL_ATTN_LAYERS, (*range(3, 93, 4), 92)) self.assertEqual((c.routed_expert_dim, c.hidden_dim, c.shared_expert_dim), (3584, 3072, 6144)) self.assertTrue(c.route_weights_uncorrected and c.kda_full_rank_gate and c.attn_output_gate) + self.assertTrue(c.ssm is not None and c.ssm.channel_decay) self.assertEqual((c.activation_situ_beta, c.activation_situ_linear_beta, c.kda_gate_lower_bound), (4.0, 25.0, -5.0)) def test_config_rejects_wrong_checkpoint(self): @@ -34,11 +40,13 @@ class TestKimiK3(unittest.TestCase): self.assertEqual(targets, set(state)) self.assertEqual(state["blk.1.ffn_gate_exps.weight"].shape, (896, 3072, 1792)) self.assertEqual(state["blk.1.ffn_gate_exps.weight_scale"].shape, (896, 3072, 112)) + self.assertEqual(state["blk.0.ssm_a"].shape, (128, 1)) _shard_kimi_k3(model, tuple(f"NULL:{i}" for i in range(8))) total, per_gpu = 0, 0 for name,value in state.items(): - dtype = dtypes.uint8 if name.endswith(("weight_scale", "_exps.weight")) else \ - dtypes.float32 if name.endswith(("ssm_a", "ssm_dt.bias")) else dtypes.bfloat16 + dtype = dtypes.uint8 if name.endswith(("weight_scale", "_exps.weight")) else dtypes.float32 if name.endswith( + ("exp_probs_b.bias", "ssm_q_conv1d.weight", "ssm_k_conv1d.weight", "ssm_v_conv1d.weight", "ssm_norm.weight", "ssm_a", "ssm_dt.bias")) \ + else dtypes.bfloat16 size = value.numel() * dtype.itemsize total += size per_gpu += size if value.uop.axis is None else size//8 @@ -75,6 +83,7 @@ class TestKimiK3(unittest.TestCase): self.assertEqual(state[name].uop.axis, axis, name) self.assertIsNone(state["blk.1.attn_res_norm.weight"].uop.axis) self.assertIsNone(state["blk.1.ffn_routed_norm.weight"].uop.axis) + self.assertIsNone(state["blk.0.ssm_a"].uop.axis) def test_direct_expert_staging(self): devices = tuple(f"PYTHON:{i}" for i in range(4)) @@ -86,6 +95,15 @@ class TestKimiK3(unittest.TestCase): _load_stacked_experts(dst, sources) np.testing.assert_equal(dst.numpy(), expected) + def test_direct_tp_replacement(self): + devices = tuple(f"PYTHON:{i}" for i in range(4)) + source = Tensor.arange(64, dtype=dtypes.float32).reshape(8, 8).realize() + expected = source.numpy() + for axis in (None, 0, 1): + dst = Tensor.zeros(8, 8, device="PYTHON").shard(devices, axis=axis) + _replace(dst, source) + np.testing.assert_equal(dst.numpy(), expected) + def test_chunked_recurrent_generate(self): model = Transformer(small_k3_config(max_context=8)) for name,value in nn.state.get_state_dict(model).items(): diff --git a/tinygrad/llm/kimi_k3.py b/tinygrad/llm/kimi_k3.py index 944c724b10..455cd96f6b 100644 --- a/tinygrad/llm/kimi_k3.py +++ b/tinygrad/llm/kimi_k3.py @@ -1,17 +1,18 @@ from __future__ import annotations -import gc, json, pathlib +import gc, json, math, pathlib from dataclasses import replace from collections import defaultdict from typing import Callable, cast -from tinygrad import Tensor, Device, nn +from tinygrad import Tensor, Device, dtypes, nn from tinygrad.device import Buffer +from tinygrad.uop.ops import UOp from tinygrad.nn.state import safe_load from tinygrad.llm.kimi import load_kimi_tokenizer_data from tinygrad.llm.model import SSMConfig, Transformer, TransformerConfig KIMI_K3_TOTAL_SIZE = 1_560_860_324_864 -KIMI_K3_TEXT_SIZE = 1_559_945_066_624 -KIMI_K3_TP8_BYTES_PER_GPU = 196_781_639_152 +KIMI_K3_TEXT_SIZE = 1_559_965_606_912 +KIMI_K3_TP8_BYTES_PER_GPU = 196_784_397_312 KIMI_K3_SHARDS = 96 KIMI_K3_EXPERTS = 896 KIMI_K3_LAYERS = 93 @@ -25,7 +26,7 @@ def kimi_k3_config(max_context:int) -> TransformerConfig: q_lora_rank=1536, kv_lora_rank=512, num_experts=896, num_experts_per_tok=16, norm_topk_prob=True, shared_expert_dim=6144, leading_dense_blocks=1, dense_hidden_dim=33792, routed_scaling_factor=1.0, expert_bias=True, expert_mxfp4=True, bf16_activations=True, kda_split_qkv=True, - ssm=SSMConfig(conv_kernel=4, state_size=128, group_count=96, time_step_rank=96, inner_size=12288, kda=True), + ssm=SSMConfig(conv_kernel=4, state_size=128, group_count=96, time_step_rank=96, inner_size=12288, kda=True, channel_decay=True), ssm_layers=KIMI_K3_SSM_LAYERS, shared_expert_gate=False, attn_output_gate=True, activation_situ_beta=4.0, activation_situ_linear_beta=25.0, routed_expert_dim=3584, latent_moe_norm=True, route_weights_uncorrected=True, attn_res_block_size=12, kda_full_rank_gate=True, kda_gate_lower_bound=-5.0, @@ -33,10 +34,12 @@ def kimi_k3_config(max_context:int) -> TransformerConfig: def kimi_k3_smoke_config(max_context:int=4) -> TransformerConfig: """Reduced K3 with every architectural feature retained for cheap compile/hardware admission tests.""" - return replace(kimi_k3_config(max_context), num_blocks=2, dim=32, hidden_dim=256, n_heads=8, n_kv_heads=8, + # Keep both the routed latent and the TP8-local expert hidden dimension wave64 aligned. The real + # gfx950 packed-expert kernel requires this, so the hardware smoke test must preserve the constraint. + return replace(kimi_k3_config(max_context), num_blocks=2, dim=32, hidden_dim=512, n_heads=8, n_kv_heads=8, vocab_size=64, head_dim=8, rope_dim=4, v_head_dim=4, q_lora_rank=16, kv_lora_rank=8, num_experts=512, - num_experts_per_tok=2, shared_expert_dim=32, dense_hidden_dim=64, routed_expert_dim=32, - ssm=SSMConfig(4, 4, 8, 8, 32, True), ssm_layers=(True, False), attn_res_block_size=1) + num_experts_per_tok=2, shared_expert_dim=32, dense_hidden_dim=64, routed_expert_dim=64, + ssm=SSMConfig(4, 4, 8, 8, 32, True, True), ssm_layers=(True, False), attn_res_block_size=1) def _shard_kimi_k3(model:Transformer, devices:tuple[str, ...]) -> None: """Tensor parallel layout for K3. The official dimensions are divisible by TP8.""" @@ -52,7 +55,7 @@ def _shard_kimi_k3(model:Transformer, devices:tuple[str, ...]) -> None: ".attn_output.weight", ".ssm_out.weight")): axis = 1 elif name.endswith((".attn_q_b.weight", ".attn_k_b.weight", ".attn_v_b.weight", ".attn_gate.weight", ".attn_q.weight", ".attn_k.weight", ".attn_v.weight", ".ssm_f_b.weight", ".ssm_g_full.weight", ".ssm_beta.weight")): axis = 0 - elif name.endswith((".ssm_q_conv1d.weight", ".ssm_k_conv1d.weight", ".ssm_v_conv1d.weight", ".ssm_a", ".ssm_dt.bias")): axis = 0 + elif name.endswith((".ssm_q_conv1d.weight", ".ssm_k_conv1d.weight", ".ssm_v_conv1d.weight", ".ssm_dt.bias")): axis = 0 value.shard_(devices, axis=axis) def _validate_config(config:dict) -> None: @@ -125,8 +128,40 @@ def _layer_sources(i:int, is_kda:bool) -> dict[str, str]: def _replace(dst:Tensor, src:Tensor) -> None: if dst.shape != src.shape: raise ValueError(f"shape mismatch: expected {dst.shape}, got {src.shape}") - dst.replace(src if isinstance(src.device, tuple) else src.shard(dst.device, dst.uop.axis) if isinstance(dst.device, tuple) else src.to(dst.device)) - dst.realize() + if not isinstance(dst.device, tuple): + dst.replace(src.to(dst.device)).realize() + return + if isinstance(src.device, tuple): + dst.replace(src.shard_like(dst)).realize() + return + + # Build the final MultiBuffer directly. The generic shard().realize() path schedules several + # kernels per tensor and recompiles them for every DISK: device. Axis-0 shards and + # replicas are contiguous, so copy those bytes straight into their final device buffers. + devices, axis, shape = dst.device, dst.uop.axis, tuple(int(x) for x in dst.shape) + try: src_buffer = cast(Buffer, src.uop.buffer) + except (AssertionError, RuntimeError): + src = src.clone().realize() + src_buffer = cast(Buffer, src.uop.buffer) + + if axis is None or axis == 0: + part_shape = shape if axis is None else (shape[0]//len(devices), *shape[1:]) + part_numel = math.prod(part_shape) + parts:list[Tensor] = [] + for i,device in enumerate(devices): + part = Tensor.empty(*part_shape, dtype=src.dtype, device=device).realize() + source = src_buffer if axis is None else src_buffer.view(part_numel, src.dtype, i*part_numel*src.dtype.itemsize) + cast(Buffer, part.uop.buffer).ensure_allocated().copy_from(source.ensure_allocated()) + parts.append(part) + else: + # Inner-axis TP slices are strided in row-major safetensors. Stage one complete tensor on + # GPU 0, then schedule all slice kernels and peer copies as one multi-device graph. + staging = Tensor.empty(*shape, dtype=src.dtype, device=devices[0]).realize() + cast(Buffer, staging.uop.buffer).ensure_allocated().copy_from(src_buffer.ensure_allocated()) + dst.replace(staging.shard(devices, axis=axis)).realize() + return + dst.replace(Tensor(parts[0].uop.mstack(*(x.uop for x in parts[1:])).unshard(axis)) if axis is not None else + Tensor(parts[0].uop.mstack(*(x.uop for x in parts[1:])))) def _load_stacked_experts(dst:Tensor, sources:list[Tensor]) -> None: """Read expert tensors once into a transient GPU staging buffer, then redistribute TP slices over the GPU fabric.""" @@ -148,23 +183,10 @@ def _load_stacked_experts(dst:Tensor, sources:list[Tensor]) -> None: staging_buffer.view(cast(int, source.numel()), source.dtype, offset).ensure_allocated().copy_from(source_buffer.ensure_allocated()) offset += source.nbytes() - # The full projection is at most the packed down-expert matrix. Materialize one TP slice at a time, - # transfer it peer-to-peer, and immediately release the GPU-0 gather temporary. - parts:list[Tensor] = [] - for i,device in enumerate(devices): - bounds = [(0, int(s)) for s in shape] - shard_size = int(shape[axis])//len(devices) - bounds[axis] = (i*shard_size, (i+1)*shard_size) - local_part = staging.shrink(tuple(bounds)).contiguous().realize() - part = local_part if device == devices[0] else local_part.to(device).realize() - parts.append(part) - if part is not local_part: - del local_part - gc.collect() - free_staging_cache() - - dst.replace(Tensor(parts[0].uop.mstack(*(x.uop for x in parts[1:])).unshard(axis))) - del staging, parts + # Schedule all TP slices and peer copies together. This avoids eight independent realization + # passes and lets the runtime overlap the multi-device transfer graph. + dst.replace(staging.shard(devices, axis=axis)).realize() + del staging gc.collect() free_staging_cache() @@ -184,10 +206,14 @@ def _load_nonexperts(root:pathlib.Path, weight_map:dict[str, str], model:Transfo shard = safe_load(root / filename) for source in sources: value, targets = shard[source], mappings[source].split("|") - if source.endswith("A_log"): value = -value.float().exp().reshape(96, 1) + # A_log is the only checkpoint tensor requiring arithmetic during load. Realize its 128 + # channel values on CPU so replicating it does not try to render the disk/PYTHON graph. + if source.endswith("A_log"): value = (-value.to("CPU").float().exp()).reshape(model_state[targets[0]].shape).realize() if source.endswith("conv1d.weight"): value = value.squeeze(1) if source.endswith("kv_b_proj.weight"): - value = value.reshape(96, 256, 512) + # Splitting K/V includes a transpose, which cannot be rendered against a disk buffer. + # Materialize only this one 25 MiB projection on CPU, then release it with the shard. + value = value.to("CPU").realize().reshape(96, 256, 512) values:tuple[Tensor, ...] = (value[:, :128].transpose(1, 2), value[:, 128:]) else: values = (value,) for target,tensor in zip(targets, values): _replace(model_state[target], tensor) @@ -197,20 +223,62 @@ def _load_nonexperts(root:pathlib.Path, weight_map:dict[str, str], model:Transfo return consumed def _load_experts(root:pathlib.Path, weight_map:dict[str, str], model:Transformer, progress:Callable[[str], None]) -> set[str]: - model_state, consumed = nn.state.get_state_dict(model), set() + model_state, consumed = nn.state.get_state_dict(model), set[str]() for i in range(1, KIMI_K3_LAYERS): - for wid,dst_name in (("w1","ffn_gate_exps"),("w3","ffn_up_exps"),("w2","ffn_down_exps")): - base = f"language_model.model.layers.{i}.block_sparse_moe.experts" - packed_keys = [f"{base}.{e}.{wid}.weight_packed" for e in range(KIMI_K3_EXPERTS)] - scale_keys = [f"{base}.{e}.{wid}.weight_scale" for e in range(KIMI_K3_EXPERTS)] - files = sorted({weight_map[k] for k in packed_keys+scale_keys}) - progress(f"loading layer {i}/92 {wid} routed experts from {', '.join(files)}") - shards = {fn:safe_load(root / fn) for fn in files} - _load_stacked_experts(model_state[f"blk.{i}.{dst_name}.weight"], [shards[weight_map[k]][k] for k in packed_keys]) - _load_stacked_experts(model_state[f"blk.{i}.{dst_name}.weight_scale"], [shards[weight_map[k]][k] for k in scale_keys]) - consumed.update(packed_keys+scale_keys) - del shards - gc.collect() + base = f"language_model.model.layers.{i}.block_sparse_moe.experts" + fields = tuple((wid, suffix, dst_name) for wid,dst_name in (("w1","ffn_gate_exps"),("w2","ffn_down_exps"),("w3","ffn_up_exps")) + for suffix in ("weight_packed", "weight_scale")) + keys = {(e,wid,suffix):f"{base}.{e}.{wid}.{suffix}" for e in range(KIMI_K3_EXPERTS) for wid,suffix,_ in fields} + files = sorted({weight_map[k] for k in keys.values()}) + progress(f"loading layer {i}/92 routed experts from {', '.join(files)}") + shards = {fn:safe_load(root / fn) for fn in files} + + # Official K3 stores all six tensors for an expert contiguously and all 896 experts for a + # layer in one contiguous shard region, ordered lexicographically by expert name. Read that + # region once, then reorder/split on GPU 0 directly into the six TP8 destinations. + blocks:list[tuple[int, int, int, list[Buffer]]] = [] + for e in range(KIMI_K3_EXPERTS): + bufs = [cast(Buffer, shards[weight_map[keys[e,wid,suffix]]][keys[e,wid,suffix]].uop.buffer) for wid,suffix,_ in fields] + if not all(bufs[j].device == bufs[0].device and bufs[j].offset+bufs[j].nbytes == bufs[j+1].offset for j in range(len(bufs)-1)): + raise ValueError(f"layer {i} expert {e} tensors are not contiguous in the official shard") + blocks.append((bufs[0].offset, bufs[-1].offset+bufs[-1].nbytes, e, bufs)) + blocks.sort() + if len(files) != 1 or not all(blocks[j][1] == blocks[j+1][0] for j in range(len(blocks)-1)): + raise ValueError(f"layer {i} routed experts are not one contiguous official-shard region") + row_bytes = blocks[0][1]-blocks[0][0] + if any(end-start != row_bytes for start,end,_,_ in blocks): raise ValueError(f"layer {i} expert records have inconsistent sizes") + + devices = cast(tuple[str, ...], model_state[f"blk.{i}.ffn_gate_exps.weight"].device) + raw = Tensor.empty(KIMI_K3_EXPERTS, row_bytes, dtype=dtypes.uint8, device=devices[0]).realize() + raw_buffer, first_buffer = cast(Buffer, raw.uop.buffer), blocks[0][3][0] + raw_buffer.ensure_allocated().copy_from(first_buffer.base.view(KIMI_K3_EXPERTS*row_bytes, dtypes.uint8, blocks[0][0]).ensure_allocated()) + lexpos = {expert:pos for pos,(_,_,expert,_) in enumerate(blocks)} + permutation = Tensor([lexpos[e] for e in range(KIMI_K3_EXPERTS)], device=devices[0]) + field_offset, outputs = 0, [] + for field_idx,(wid,suffix,dst_name) in enumerate(fields): + field_bytes = blocks[0][3][field_idx].nbytes + dst = model_state[f"blk.{i}.{dst_name}.weight" + ("_scale" if suffix == "weight_scale" else "")] + axis = dst.uop.axis + if axis is None: raise ValueError(f"layer {i} expert destination {dst_name} is not TP-sharded") + value = raw[:, field_offset:field_offset+field_bytes][permutation].reshape(dst.shape).shard(devices, axis=axis) + # Realize into a buffer-identity tensor, then retain only that identity in the model. Keeping + # value's arithmetic UOp would also keep the 15.7 GB raw staging tensor and its reorder graph + # alive for every loaded weight, wasting about 44 GB on GPU 0 after the load completes. + shard_shape = tuple(int(x) for x in value.uop.shard_shape) + storage = UOp.new_buffer(devices, math.prod(shard_shape), dst.dtype).reshape(shard_shape).unshard(axis) + final = Tensor(storage) + final.assign(value) + outputs.append((dst, final, storage)) + field_offset += field_bytes + if field_offset != row_bytes: raise ValueError(f"layer {i} expert field sizes do not cover the contiguous record") + outputs[0][1].realize(*(value for _,value,_ in outputs[1:])) + for dst,_,storage in outputs: dst.replace(Tensor(storage)) + consumed.update(keys.values()) + # Drop the realized assignment graphs before flushing the allocator cache. Their final storage + # UOps remain in model_state, while the graphs themselves still reference raw and permutation. + del shards, raw, raw_buffer, permutation, outputs, value, final, storage, dst + gc.collect() + if (free_cache:=getattr(Device[devices[0]].allocator, "free_cache", None)) is not None: free_cache() return consumed def load_kimi_k3(model_dir:str|pathlib.Path, max_context:int=4096, devices:int=8, diff --git a/tinygrad/llm/model.py b/tinygrad/llm/model.py index 9b3a109582..779dd13a92 100644 --- a/tinygrad/llm/model.py +++ b/tinygrad/llm/model.py @@ -83,6 +83,7 @@ class SSMConfig: time_step_rank: int inner_size: int kda: bool = False + channel_decay: bool = False @dataclass(frozen=True) class TransformerConfig: @@ -413,7 +414,8 @@ class GatedDeltaNetBlock(FFNBlock): self.ssm_alpha = Linear(config.dim, self.num_v_heads, bias=False) self.ssm_beta = Linear(config.dim, self.num_v_heads, bias=False) self.ssm_dt = {"bias": Tensor.zeros(ssm.inner_size if ssm.kda else self.num_v_heads)} - self.ssm_a = Tensor.zeros(self.num_v_heads, 1) if ssm.kda else Tensor.zeros(self.num_v_heads) + self.ssm_a = Tensor.zeros(self.head_v_dim if ssm.channel_decay else self.num_v_heads, 1) if ssm.kda else Tensor.zeros(self.num_v_heads) + self.kda_channel_decay = ssm.channel_decay self.ssm_norm, self.ssm_out = nn.RMSNorm(self.head_v_dim, config.norm_eps), Linear(ssm.inner_size, config.dim, bias=False) def _attention(self, x:Tensor, start_pos:int|UOp) -> Tensor: @@ -466,9 +468,10 @@ class GatedDeltaNetBlock(FFNBlock): v = v.reshape(B, T, self.num_v_heads, self.head_v_dim).transpose(1, 2).float() beta = (beta_logits.float() if self.config.ssm and self.config.ssm.kda else beta_logits).sigmoid().transpose(1, 2) gate_logits = (alpha_logits.float() + self.ssm_dt["bias"]).reshape(B, T, self.num_v_heads, -1) + a_shape = (1, 1, 1, self.head_v_dim) if self.kda_channel_decay else (1, 1, self.num_v_heads, 1) if self.config.kda_gate_lower_bound: - log_alpha = self.config.kda_gate_lower_bound * ((-self.ssm_a).reshape(1, 1, self.num_v_heads, -1) * gate_logits).sigmoid() - else: log_alpha = gate_logits.softplus() * self.ssm_a.reshape(1, 1, self.num_v_heads, -1) + log_alpha = self.config.kda_gate_lower_bound * ((-self.ssm_a).reshape(a_shape) * gate_logits).sigmoid() + else: log_alpha = gate_logits.softplus() * self.ssm_a.reshape(a_shape) alpha = log_alpha.squeeze(-1).transpose(1, 2).exp() if log_alpha.shape[-1] == 1 else log_alpha.permute(0, 2, 1, 3).exp() if T == 1: # Keep decode on the small elementwise graph. The fused prefill kernel writes a temporary