llm: accelerate Kimi serving on gfx11

docs/kimi_k3_mi350.md

@@ -81,7 +81,18 @@ DEV=AMD DEBUG=1 python -m tinygrad.llm.cli --model /models/Kimi-K3 \
 
 The correctness path expands only selected MXFP4 expert weights and emulates MXFP8 activation quantization. tinygrad has gfx950/CDNA4 BF16 and FP8 matrix-core support, but this branch does not yet have a hardware-validated fused 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.
 
-Recurrent prefill is fused and defaults to eight-token chunks. The portable custom kernel compiles for gfx950, but the wave-parallel version is deliberately restricted to gfx11 because it uses RDNA3 wave32 swizzles. A gfx950-tuned wave64/MFMA recurrent kernel remains a performance task for the rented machine; do not enable the gfx11 kernel on CDNA without rewriting its lane reduction and validating every recurrent-state transition.
+Recurrent prefill is fused and defaults to 32-token chunks. The portable custom kernel compiles for gfx950, but the wave-parallel version is deliberately restricted to gfx11 because it uses RDNA3 wave32 swizzles. A gfx950-tuned wave64/MFMA recurrent kernel remains a performance task for the rented machine; do not enable the gfx11 kernel on CDNA without rewriting its lane reduction and validating every recurrent-state transition.
+
+The following serving changes are portable and therefore apply to the official K3 path: recurrent-state reset graph capture, correct fresh-prompt benchmark resets, 32-token recurrent chunks, materialized gate/up boundaries that prevent pathological fused reduction kernels, separate greedy decode JITs, and K3's uncorrected routed probability semantics. The new packed expert, software MXFP8, fused KDA Q/K/V, combined TP down-partial, and exact greedy output-head kernels are intentionally gated to gfx11. They improve the local 7900 XTX bring-up but will fall back to the generic implementation on gfx950.
+
+After hardware admission on MI350X, profile before porting those kernels. The likely implementation order is:
+
+1. A native packed MXFP4×MXFP8 grouped expert GEMM using CDNA4 matrix instructions.
+2. A wave64/MFMA KDA Q/K/V decode projection.
+3. Combined routed/shared down-projection TP partials so each layer performs one XGMI all-reduce.
+4. A CDNA4 output-head matvec and router matvec if they remain visible in the profile.
+
+Every port needs a direct numerical comparison with the generic graph and an end-to-end greedy-token comparison before performance measurements. None of the gfx11 custom kernels should be enabled on gfx950 by changing only the architecture guard.
 
 The official checkpoint also contains MoonViT-V2 and multimodal projector weights. They are skipped by the text loader. Image input remains a separate implementation and validation task.
 
@@ -90,20 +101,23 @@ The official checkpoint also contains MoonViT-V2 and multimodal projector weight
 The pre-rental benchmark uses the converted `Kimi-Linear-48B-A3B-Instruct-MXFP4-v2` checkpoint on four gfx1100 GPUs. It is a useful regression test for the KDA/MLA/MoE text path, not a projection of K3 throughput on MI350X.
 
 ```sh
-DEV=AMD python extra/benchmark_kimi.py \
-  /home/tiny/models/Kimi-Linear-48B-A3B-Instruct-MXFP4-v2 \
-  --devices 4 --max-context 128 --prompt-tokens 32 --decode-tokens 8 --chunk-size 8
+DEV=AMD JIT_BATCH_SIZE=64 python extra/benchmark_kimi.py \
+  /raid/models/Kimi-Linear-48B-A3B-Instruct-MXFP4-v2 \
+  --devices 4 --max-context 128 --prompt-tokens 32 --decode-tokens 8 --chunk-size 32
 ```
 
 Results from 2026-08-10:
 
-- load: 310.280s for 29.27 GB
-- cold prefill including compilation: 47.622s
-- captured prefill: 1.020s, 31.38 tok/s
-- steady prefill replay: 0.898s, 35.64 tok/s
-- steady decode replay: 23.03 tok/s, 43.43 ms/token
-- peak host RSS: 901,192 KiB; swap: 0
+- load from RAID: 43.9–44.4s for the 29.27 GB checkpoint
+- first 32-token prefill includes roughly 10s of compilation/capture
+- steady fresh-prompt prefill replay: 0.120s, 267.18 tok/s
+- steady decode replay: 66.19 tok/s, 15.11 ms/token
+- peak host RSS: 799.9 MiB; swap was not used
 
-The same resident model measured 19.60 tok/s with tokenwise prefill, so the selected eight-token chunk is 1.82× faster. A 32-token chunk fell to 3.03 tok/s because the current selected-expert path expands packed weights per token; larger batches multiply that temporary dequantization work. Chunk 8 is therefore a conservative default until native grouped MXFP4×MXFP8 expert GEMM exists.
+The load and prefill targets of less than 60 seconds and more than 200 tok/s are met on this host. Decode improved from 23.03 tok/s to 66.19 tok/s but remains below the 100 tok/s target. The remaining local profile is dominated by many small router/shared projections, collective and graph-launch overhead, and the unavoidable active expert traffic; this result must not be reported as reaching the decode goal.
+
+Four 7900 XTX cards provide 96 GB aggregate VRAM and about 3.84 TB/s aggregate physical memory bandwidth. Their nominal aggregate vector FP16 rate is about 245.6 TFLOP/s, or about 492 TFLOP/s through matrix instructions. Kimi Linear activates roughly 3.107B parameters per token; a simple active-weight accounting gives approximately 4.05 GB/token and an optimistic bandwidth-only ceiling near 948 tok/s. The measured decode rate is much lower because this MoE decode workload is a collection of small matrix-vector operations plus PCIe collectives, not one ideal streaming kernel.
+
+The generic loader currently rereads logical TP shards and accounts for roughly 227 GB of disk traffic for a TP4 load. RAID bandwidth hides that inefficiency locally, but a direct one-pass shard loader remains worthwhile before slow remote storage is used. It was not retained here because the attempted direct-shard graph exposed an unresolved scheduler/renderer edge; correctness and bounded memory take priority over avoiding the redundant reads.
 
 Different chunk sizes can choose a different final token because their matrix kernels use different floating-point reduction orders. Each measured shape was repeatable between cold and captured execution. For official K3 validation, compare logits/tokens against the reference at one fixed chunk size and greedy settings rather than requiring bitwise agreement between performance shapes.

extra/benchmark_kimi.py

@@ -13,6 +13,12 @@ def timed_next(gen, devices:int) -> tuple[int, float]:
   sync(devices)
   return token, time.perf_counter()-begin
 
+def fresh_generate(model, prompt:list[int], chunk_size:int):
+  # Force recurrent/KV state reset so repeated runs and chunk sweeps measure the entire prompt,
+  # rather than silently reusing the prefix cached by the previous measurement.
+  model._cached_tokens = [-1] * len(prompt)
+  return model.generate(prompt.copy(), chunk_size=chunk_size)
+
 def main() -> None:
   parser = argparse.ArgumentParser()
   parser.add_argument("model", help="converted Kimi-Linear-48B-A3B MXFP4-v2 directory")
@@ -40,10 +46,10 @@ def main() -> None:
     # Recurrent prefill has a static token dimension. Give each swept shape its own capture;
     # the rollout JIT remains shared and independently benchmarks chunk 1/decode.
     if chunk != 1: model.prefill_jit = TinyJit(model.forward)
-    cold = model.generate(prompt.copy(), chunk_size=chunk)
+    cold = fresh_generate(model, prompt, chunk)
     first, cold_prefill = timed_next(cold, args.devices)
     print(f"chunk {chunk}: cold prefill {cold_prefill:.3f}s, token={first}", flush=True)
-    warm = model.generate(prompt.copy(), chunk_size=chunk)
+    warm = fresh_generate(model, prompt, chunk)
     warm_first, prefill = timed_next(warm, args.devices)
     if first != warm_first: raise RuntimeError(f"chunk {chunk} is not repeatable: cold={first}, warm={warm_first}")
     timings.append((prefill, chunk))
@@ -52,7 +58,7 @@ def main() -> None:
 
   prefill, best_chunk = min(timings)
   if best_chunk != 1: model.prefill_jit = prefill_jits[best_chunk]
-  warm = model.generate(prompt.copy(), chunk_size=best_chunk)
+  warm = fresh_generate(model, prompt, best_chunk)
   first, replay_prefill = timed_next(warm, args.devices)
   _, cold_decode = timed_next(warm, args.devices)
   _, capture_decode = timed_next(warm, args.devices)

tinygrad/llm/kernels/__init__.py

@@ -11,6 +11,83 @@ def amd_custom_kernels_supported(device:str|tuple[str, ...]|None) -> bool:
   with Context(ALLOW_DEVICE_USAGE=1):
     return (target:=getattr(Device[device], "target", None)) is not None and target[0] == 11
 
+def mxfp4_expert_linear(sel:Tensor, x:Tensor, weight:Tensor, scale:Tensor, partial:bool=False) -> Tensor:
+  """Run a TP routed projection without materializing selected BF16 weights."""
+  from tinygrad.llm.kernels.amd import _mxfp4_expert_linear_kernel
+  batch, tokens, topk = sel.shape
+  out_features = weight.shape[1]
+  weight_axis = weight.uop.axis
+  if isinstance(weight.device, tuple):
+    devices = weight.device
+    # Gate/up shard their output dimension. Down shards its reduction dimension;
+    # represent each GPU's partial as a size-one device axis, then all-reduce it.
+    axis = 3 if weight_axis == 1 else 4
+    shard_shape: tuple[int|UOp, ...]
+    if weight_axis == 1:
+      if out_features % len(devices): raise ValueError(f"expert output {out_features} is not divisible by {len(devices)} devices")
+      shard_shape = (batch, tokens, topk, out_features//len(devices))
+    elif weight_axis == 2:
+      shard_shape = (batch, tokens, topk, out_features, 1)
+    else: raise ValueError(f"unsupported expert TP axis {weight_axis}")
+    partial_dtype = dtypes.float32 if weight_axis == 2 else dtypes.bfloat16
+    parts = [Tensor.empty(*shard_shape, dtype=partial_dtype, device=device).uop for device in devices]
+    out = Tensor(parts[0].mstack(*parts[1:]).unshard(axis))
+  else:
+    out = Tensor.empty(batch, tokens, topk, out_features, dtype=dtypes.bfloat16, device=weight.device)
+  out = Tensor.custom_kernel(out, sel.contiguous(), x.contiguous(), weight, scale, fxn=_mxfp4_expert_linear_kernel)[0]
+  return out if weight_axis == 2 and partial else out.sum(4).cast(dtypes.bfloat16) if weight_axis == 2 else out
+
+def bf16_partial_linear(x:Tensor, weight:Tensor) -> Tensor:
+  """Return output-shaped FP32 TP partials with a final device axis, without all-reduce."""
+  from tinygrad.llm.kernels.amd import _bf16_partial_linear_kernel
+  if not isinstance(weight.device, tuple) or weight.uop.axis != 1: raise ValueError("partial linear expects input-sharded TP weight")
+  batch, tokens, _ = x.shape
+  devices, out_features = weight.device, weight.shape[0]
+  shard_shape = (batch, tokens, out_features, 1)
+  parts = [Tensor.empty(*shard_shape, dtype=dtypes.float32, device=device).uop for device in devices]
+  out = Tensor(parts[0].mstack(*parts[1:]).unshard(3))
+  return Tensor.custom_kernel(out, x.contiguous(), weight, fxn=_bf16_partial_linear_kernel)[0]
+
+def bf16_matvec(x:Tensor, weight:Tensor) -> Tensor:
+  from tinygrad.llm.kernels.amd import _bf16_matvec_kernel
+  batch, tokens, _ = x.shape
+  out_features = weight.shape[0]
+  if isinstance(weight.device, tuple):
+    if weight.uop.axis != 0: raise ValueError("bf16_matvec expects output-sharded TP weight")
+    devices = weight.device
+    shard_shape = (batch, tokens, out_features//len(devices))
+    parts = [Tensor.empty(*shard_shape, dtype=dtypes.bfloat16, device=device).uop for device in devices]
+    out = Tensor(parts[0].mstack(*parts[1:]).unshard(2))
+  else: out = Tensor.empty(batch, tokens, out_features, dtype=dtypes.bfloat16, device=weight.device)
+  return Tensor.custom_kernel(out, x.contiguous(), weight, fxn=_bf16_matvec_kernel)[0]
+
+def mxfp8_quantize_dequantize(x:Tensor) -> Tensor:
+  """gfx11 software MXFP8 round trip without a multi-kernel reduction graph."""
+  from tinygrad.llm.kernels.amd import _mxfp8_qdq_kernel
+  out = Tensor.empty_like(x, dtype=dtypes.bfloat16)
+  return Tensor.custom_kernel(out, x.contiguous(), fxn=_mxfp8_qdq_kernel)[0]
+
+def kda_qkv_linear(x:Tensor, qw:Tensor, kw:Tensor, vw:Tensor) -> tuple[Tensor, Tensor, Tensor]:
+  """Fuse equal-sized output-sharded KDA Q/K/V decode projections."""
+  from tinygrad.llm.kernels.amd import _kda_qkv_kernel
+  batch, tokens, _ = x.shape
+  out_features = qw.shape[0]
+  if not (qw.shape == kw.shape == vw.shape): raise ValueError("fused KDA Q/K/V weights must have equal shapes")
+  if isinstance(qw.device, tuple):
+    devices = qw.device
+    if qw.uop.axis != 0 or out_features % len(devices): raise ValueError("fused KDA Q/K/V expects output-sharded weights")
+    shard_shape = (batch, tokens, out_features//len(devices))
+    def make_out() -> Tensor:
+      parts = [Tensor.empty(*shard_shape, dtype=dtypes.bfloat16, device=device).uop for device in devices]
+      return Tensor(parts[0].mstack(*parts[1:]).unshard(2))
+    outs: tuple[Tensor, Tensor, Tensor] = (make_out(), make_out(), make_out())
+  else:
+    outs = (Tensor.empty(batch, tokens, out_features, dtype=dtypes.bfloat16, device=qw.device),
+            Tensor.empty(batch, tokens, out_features, dtype=dtypes.bfloat16, device=qw.device),
+            Tensor.empty(batch, tokens, out_features, dtype=dtypes.bfloat16, device=qw.device))
+  ret = Tensor.custom_kernel(*outs, x.contiguous(), qw, kw, vw, fxn=_kda_qkv_kernel)
+  return ret[0], ret[1], ret[2]
+
 @functools.cache
 def _gated_delta_prefill_kernel(core:UOp, next_state:UOp, q:UOp, k:UOp, v:UOp, beta:UOp, alpha:UOp, state:UOp, kq:UOp) -> UOp:
   batch, heads, tokens, value_dim = cast(tuple[int, int, int, int], core.shape)

tinygrad/llm/kernels/amd.py

@@ -5,14 +5,136 @@ from tinygrad import UOp
 from tinygrad.uop.ops import AxisType, KernelInfo, Ops
 from tinygrad.dtype import AddrSpace, dtypes
 
-def warp_reduce(val:UOp, full_wave:bool=False) -> UOp:
+def warp_reduce(val:UOp, full_wave:bool=False, maximum:bool=False) -> UOp:
   for offset in ((16, 8, 4, 2, 1) if full_wave else (8, 4, 2, 1)):
     if val.op is Ops.INDEX and val.addrspace == AddrSpace.REG: val = val.load()
     other = UOp(Ops.CUSTOM, dtypes.float, (val,), arg=
       f"__builtin_bit_cast(float, __builtin_amdgcn_ds_swizzle(__builtin_bit_cast(int, {{0}}), {0x1f | offset<<10}))")
-    val = val + other
+    val = val.maximum(other) if maximum else val + other
   return val
 
+@functools.cache
+def _mxfp8_qdq_kernel(out:UOp, x:UOp) -> UOp:
+  """Software OCP E4M3/E8M0 round trip, one wave per 32-value MX block."""
+  groups = cast(int, x.shape[-1])//32
+  outer = x.numel()//cast(int, x.shape[-1])
+  block, lane = UOp.range(outer*groups, 0), UOp.range(32, 1, axis_type=AxisType.LOCAL)
+  value = x.reshape(outer, groups, 32)[block//groups, block%groups, lane].float()
+  amax = warp_reduce(value.abs(), full_wave=True, maximum=True)
+  exponent = (amax.maximum(1e-38)/448.0).log2().round().maximum(-127.0).minimum(127.0)
+  block_scale = amax.eq(0).where(1.0, exponent.exp2())
+  normalized = value/block_scale
+  magnitude = normalized.abs().minimum(448.0)
+  elem_exp = magnitude.maximum(2**-9).log2().floor().maximum(-6.0).minimum(8.0)
+  quantum = (elem_exp-3.0).exp2()
+  quantized = (magnitude/quantum).round()*quantum
+  quantized = (normalized < 0).where(-quantized, quantized).maximum(-448.0).minimum(448.0)
+  store = out.reshape(outer, groups, 32)[block//groups, block%groups, lane].store((quantized*block_scale).cast(out.dtype))
+  return store.end(lane, block).sink(arg=KernelInfo(name="mxfp8_qdq", opts_to_apply=()))
+
+def _mxfp4_value(code:UOp) -> UOp:
+  """Decode one OCP E2M1 nibble without a lookup-table memory access."""
+  magnitude = code & 7
+  value = magnitude.eq(7).where(6.0, magnitude.eq(6).where(4.0, magnitude.eq(5).where(3.0, magnitude.float()*0.5)))
+  return (code & 8).ne(0).where(-value, value)
+
+@functools.cache
+def _kda_qkv_kernel(qout:UOp, kout:UOp, vout:UOp, x:UOp, qw:UOp, kw:UOp, vw:UOp) -> UOp:
+  """Fused BF16 decode projection for equal-sized KDA Q/K/V tensors."""
+  batch, tokens, out_features = cast(tuple[int, int, int], qout.shape)
+  in_features, output_tile = cast(int, x.shape[-1]), 1
+  assert qout.shape == kout.shape == vout.shape and out_features % output_tile == 0 and in_features % 32 == 0
+  row, lane = UOp.range(batch*tokens*(out_features//output_tile), 0), UOp.range(32, 1, axis_type=AxisType.LOCAL)
+  token, output_block = row // (out_features//output_tile), row % (out_features//output_tile)
+  outputs = tuple(output_block*output_tile+i for i in range(output_tile))
+  acc = UOp.placeholder((3, output_tile), dtypes.float32, slot=0, addrspace=AddrSpace.REG)
+  acc = acc.after(acc.store(acc.const_like(0.0)))
+  group = UOp.range(in_features//32, 2, AxisType.REDUCE)
+  activation = x.reshape(batch*tokens, in_features)[token, group*32+lane].float()
+  updates = [acc.after(group)[p, i].load()+activation*w[output, group*32+lane].float()
+             for p,w in enumerate((qw, kw, vw)) for i,output in enumerate(outputs)]
+  update = acc.store(UOp.stack(*updates).reshape(3, output_tile)).end(group)
+  outs = (qout, kout, vout)
+  stores = (outs[p].reshape(batch*tokens, out_features)[token, output.valid(lane.eq(0))].store(
+    warp_reduce(acc.after(update)[p, i], full_wave=True).cast(outs[p].dtype))
+    for p in range(3) for i,output in enumerate(outputs))
+  return UOp.group(*stores).end(lane, row).sink(arg=KernelInfo(name="kda_qkv", opts_to_apply=()))
+
+@functools.cache
+def _mxfp4_expert_linear_kernel(out:UOp, sel:UOp, x:UOp, weight:UOp, scale:UOp) -> UOp:
+  """Wave32 decode GEMM which consumes selected experts directly from packed MXFP4 storage."""
+  batch, tokens, topk, out_features = cast(tuple[int, int, int, int], out.shape[:4])
+  partials = cast(int, out.shape[4]) if len(out.shape) == 5 else 1
+  output_tile = 1
+  assert out_features % output_tile == 0
+  in_features = cast(int, weight.shape[-1])*2
+  assert in_features % 32 == 0 and x.shape[-1] == in_features and sel.shape == (batch, tokens, topk)
+  xchoices = cast(int, x.shape[-2])
+  assert xchoices in (1, topk)
+  row, lane = UOp.range(batch*tokens*topk*(out_features//output_tile)*partials, 0), UOp.range(32, 1, axis_type=AxisType.LOCAL)
+  partial, output_block, route = row % partials, (row//partials) % (out_features//output_tile), \
+    row // ((out_features//output_tile)*partials)
+  outputs = tuple(output_block*output_tile+i for i in range(output_tile))
+  token, choice = route // topk, route % topk
+  expert = sel.reshape(batch*tokens, topk)[token, choice]
+  xv = x.reshape(batch*tokens, xchoices, in_features)
+  acc = UOp.placeholder((output_tile,), dtypes.float32, slot=0, addrspace=AddrSpace.REG)
+  acc = acc.after(acc.store(acc.const_like(0.0)))
+  group = UOp.range(in_features//32, 2, AxisType.REDUCE)
+  activation = xv[token, 0 if xchoices == 1 else choice, group*32+lane].float()
+  updates = []
+  for i,output in enumerate(outputs):
+    packed = weight[expert, output, group*16 + lane//2]
+    code = (packed >> ((lane&1)*4).cast(dtypes.uint8)) & 15
+    w = _mxfp4_value(code) * (scale[expert, output, group].float()-127.0).exp2()
+    updates.append(acc.after(group)[i].load()+activation*w)
+  update = acc.store(UOp.stack(*updates)).end(group)
+  out = out.reshape(batch*tokens, topk, out_features, partials)
+  stores = (out[token, choice, output, partial.valid(lane.eq(0))].store(warp_reduce(acc.after(update)[i], full_wave=True).cast(out.dtype))
+            for i,output in enumerate(outputs))
+  return UOp.group(*stores).end(lane, row).sink(arg=KernelInfo(name="mxfp4_expert_linear", opts_to_apply=()))
+
+@functools.cache
+def _bf16_partial_linear_kernel(out:UOp, x:UOp, weight:UOp) -> UOp:
+  """Per-device BF16 down projection; its dummy final axis is reduced after combining TP partials."""
+  batch, tokens, out_features, partials = cast(tuple[int, int, int, int], out.shape)
+  in_features, output_tile = cast(int, x.shape[-1]), 1
+  assert out_features % output_tile == 0 and in_features % 16 == 0
+  row, lane = UOp.range(batch*tokens*(out_features//output_tile)*partials, 0), UOp.range(16, 1, axis_type=AxisType.LOCAL)
+  partial, output_block, token = row%partials, (row//partials)%(out_features//output_tile), row//(partials*(out_features//output_tile))
+  outputs = tuple(output_block*output_tile+i for i in range(output_tile))
+  acc = UOp.placeholder((output_tile,), dtypes.float32, slot=0, addrspace=AddrSpace.REG)
+  acc = acc.after(acc.store(acc.const_like(0.0)))
+  chunk, group = in_features//16, UOp.range(in_features//16, 2, AxisType.REDUCE)
+  input_idx = lane*chunk+group
+  activation = x.reshape(batch*tokens, in_features)[token, input_idx].float()
+  update = acc.store(UOp.stack(*(acc.after(group)[i].load()+(activation*weight[output, input_idx].float()).cast(dtypes.bfloat16).float()
+    for i,output in enumerate(outputs)))).end(group)
+  local = UOp.placeholder((output_tile, 16), dtypes.float32, slot=1, addrspace=AddrSpace.LOCAL)
+  barrier = UOp.group(*(local[i, lane].store(acc.after(update)[i]) for i in range(output_tile))).barrier()
+  stores = (out.reshape(batch*tokens, out_features, partials)[token, output, partial.valid(lane.eq(0))].store(
+    sum((local.after(barrier)[i, j] for j in range(16)), UOp.const(0, dtypes.float32))) for i,output in enumerate(outputs))
+  return UOp.group(*stores).end(lane, row).sink(arg=KernelInfo(name="bf16_partial_linear", opts_to_apply=()))
+
+@functools.cache
+def _bf16_matvec_kernel(out:UOp, x:UOp, weight:UOp) -> UOp:
+  batch, tokens, out_features = cast(tuple[int, int, int], out.shape)
+  in_features = cast(int, x.shape[-1])
+  assert in_features % 16 == 0
+  row, lane = UOp.range(batch*tokens*out_features, 0), UOp.range(16, 1, axis_type=AxisType.LOCAL)
+  token, output = row//out_features, row%out_features
+  acc = UOp.placeholder((), dtypes.float32, slot=0, addrspace=AddrSpace.REG)
+  acc = acc.after(acc.store(0.0))
+  chunk, group = in_features//16, UOp.range(in_features//16, 2, AxisType.REDUCE)
+  input_idx = lane*chunk+group
+  product = (x.reshape(batch*tokens, in_features)[token, input_idx].float()*weight[output, input_idx].float()).cast(dtypes.bfloat16).float()
+  update = acc.store(acc.after(group)+product).end(group)
+  local = UOp.placeholder((16,), dtypes.float32, slot=1, addrspace=AddrSpace.LOCAL)
+  barrier = local[lane].store(acc.after(update)).barrier()
+  total = sum((local.after(barrier)[j] for j in range(16)), UOp.const(0, dtypes.float32))
+  return out.reshape(batch*tokens, out_features)[token, output.valid(lane.eq(0))].store(total.cast(out.dtype)).end(lane, row).sink(
+    arg=KernelInfo(name="bf16_matvec", opts_to_apply=()))
+
 @functools.cache
 def _gated_delta_prefill_kernel(core:UOp, next_state:UOp, q:UOp, k:UOp, v:UOp, beta:UOp, alpha:UOp, state:UOp, kq:UOp) -> UOp:
   batch, heads, tokens, value_dim, row_tile = *core.shape, 4

tinygrad/llm/kimi.py

@@ -15,7 +15,7 @@ def kimi_config(max_context:int, expert_mxfp4:bool=True) -> TransformerConfig:
     num_experts=256, num_experts_per_tok=8, norm_topk_prob=True, shared_expert_dim=1024, leading_dense_blocks=1,
     dense_hidden_dim=9216, routed_scaling_factor=2.446, expert_bias=True, max_context=max_context, expert_mxfp4=expert_mxfp4,
     shared_expert_gate=False, bf16_activations=True, kda_split_qkv=True,
-    recurrent_prefill_chunked=True, recurrent_prefill_chunk_size=8,
+    recurrent_prefill_chunked=True, recurrent_prefill_chunk_size=32,
     ssm=SSMConfig(conv_kernel=4, state_size=128, group_count=32, time_step_rank=32, inner_size=4096, kda=True),
     ssm_layers=KIMI_SSM_LAYERS)
 

tinygrad/llm/model.py

@@ -1,11 +1,13 @@
 from __future__ import annotations
 import functools, itertools, pathlib
 from dataclasses import dataclass, replace
+from typing import cast
 from tinygrad import Tensor, nn, UOp, TinyJit, getenv, function, dtypes
 from tinygrad.nn import Linear
 from tinygrad.llm.gguf import gguf_load
 from tinygrad.llm.quant import dequantize_mxfp4, quantize_dequantize_mxfp8
-from tinygrad.llm.kernels import gated_delta_prefill
+from tinygrad.llm.kernels import amd_custom_kernels_supported, bf16_matvec, bf16_partial_linear, gated_delta_prefill, kda_qkv_linear, \
+  mxfp4_expert_linear, mxfp8_quantize_dequantize
 from tinygrad.uop.ops import resolve
 
 @functools.cache
@@ -29,11 +31,17 @@ class MXFP4ExpertWeights:
     self.in_features, self.out_features = in_features, out_features
     self.weight = Tensor.zeros(num_experts, out_features, in_features//2, dtype=dtypes.uint8)
     self.weight_scale = Tensor.full((num_experts, out_features, in_features//32), 127, dtype=dtypes.uint8)
-  def __call__(self, sel:Tensor, x:Tensor) -> Tensor:
+  def __call__(self, sel:Tensor, x:Tensor, quantized:bool=False, partial:bool=False) -> Tensor:
     # Only selected weights are expanded, so packed storage remains resident during generation.
     if isinstance(self.weight.device, tuple) and not isinstance(sel.device, tuple): sel = sel.shard(self.weight.device, axis=None)
+    if not quantized:
+      x = mxfp8_quantize_dequantize(x.cast(dtypes.bfloat16)) if amd_custom_kernels_supported(x.device) else \
+        quantize_dequantize_mxfp8(x.cast(dtypes.bfloat16))
+    # gfx11 has no native FP4 instructions, but decoding nibbles inside the dot product still avoids
+    # the much larger selected-expert BF16 temporary. Gate/up weights are output-sharded in TP.
+    if amd_custom_kernels_supported(self.weight.device):
+      return mxfp4_expert_linear(sel, x, self.weight, self.weight_scale, partial=partial)
     weight = dequantize_mxfp4(self.weight[sel], self.weight_scale[sel], dtype=dtypes.bfloat16)
-    x = quantize_dequantize_mxfp8(x.cast(dtypes.bfloat16))
     return (x.unsqueeze(-2) @ weight.transpose(-1, -2)).contiguous().squeeze(-2)
 
 def apply_rope(x:Tensor, freqs_cis:Tensor) -> Tensor:
@@ -168,29 +176,46 @@ class FFNBlock:
       logits = x.float().linear(self.ffn_gate_inp.weight.float().transpose()) if self.config.bf16_activations else self.ffn_gate_inp(x)
       if hasattr(self, 'exp_probs_b'):
         scores = logits.sigmoid()
+        adjusted_scores = scores + self.exp_probs_b["bias"]
         topk = iterative_topk if self.config.num_experts >= 512 else pairwise_topk
-        _, sel = topk(scores + self.exp_probs_b["bias"], self.config.num_experts_per_tok)
+        _, sel = topk(adjusted_scores, self.config.num_experts_per_tok)
+        probs = (scores if self.config.route_weights_uncorrected else adjusted_scores).gather(-1, sel)
         # Kimi-Linear-48B's older reference weights corrected scores. K3 selects with the correction
         # but gathers the uncorrected sigmoid scores, so keep this an explicit compatibility switch.
-        probs = (scores if self.config.route_weights_uncorrected else scores + self.exp_probs_b["bias"]).gather(-1, sel)
         if self.config.norm_topk_prob: probs = probs / (probs.sum(axis=-1, keepdim=True) + 1e-20)
       else:
         vals, sel = pairwise_topk(logits, self.config.num_experts_per_tok)
         probs = vals.softmax(-1) if self.config.norm_topk_prob else logits.softmax(-1).gather(-1, sel)
       probs = probs * self.config.routed_scaling_factor
       if hasattr(self, 'ffn_routed_down'): h = self.ffn_routed_down(x).unsqueeze(2)
-      x_down = self.ffn_down_exps(sel, self._activation(self.ffn_gate_exps(sel, h), self.ffn_up_exps(sel, h)).contiguous())
-      out = (x_down * probs.unsqueeze(-1)).sum(axis=2).cast(x_down.dtype)  # (B, T, D)
+      if isinstance(self.ffn_gate_exps, MXFP4ExpertWeights) and amd_custom_kernels_supported(h.device):
+        hq = mxfp8_quantize_dequantize(h.cast(dtypes.bfloat16))
+        gate = self.ffn_gate_exps(sel, hq, quantized=True)
+        up = cast(MXFP4ExpertWeights, self.ffn_up_exps)(sel, hq, quantized=True)
+      else: gate, up = self.ffn_gate_exps(sel, h), self.ffn_up_exps(sel, h)
+      routed_activation = self._activation(gate, up).contiguous()
+      combine_down = resolve(x.shape[1] == 1) and isinstance(self.ffn_down_exps, MXFP4ExpertWeights) and \
+        hasattr(self, 'ffn_gate_shexp') and not hasattr(self, 'ffn_routed_up') and amd_custom_kernels_supported(x.device)
+      x_down = cast(MXFP4ExpertWeights, self.ffn_down_exps)(sel, routed_activation, partial=True) if combine_down else \
+        self.ffn_down_exps(sel, routed_activation)
+      out = (x_down * probs.unsqueeze(-1).unsqueeze(-1)).sum(axis=2) if combine_down else \
+        (x_down * probs.unsqueeze(-1)).sum(axis=2).cast(x_down.dtype)  # (B, T, D[, devices])
       if hasattr(self, 'ffn_routed_up'):
         if hasattr(self, 'ffn_routed_norm'): out = self.ffn_routed_norm(out)
         out = self.ffn_routed_up(out)
       if hasattr(self, 'ffn_gate_shexp'):
-        shexp = self.ffn_down_shexp(self._activation(self.ffn_gate_shexp(x), self.ffn_up_shexp(x)).contiguous())
-        if hasattr(self, 'ffn_gate_inp_shexp'): shexp = shexp * (x * self.ffn_gate_inp_shexp["weight"]).sum(axis=-1, keepdim=True).sigmoid()
-        out = out + shexp
+        shared_gate, shared_up = self.ffn_gate_shexp(x).contiguous(), self.ffn_up_shexp(x).contiguous()
+        shared_activation = self._activation(shared_gate, shared_up).contiguous()
+        if combine_down:
+          out = (out + bf16_partial_linear(shared_activation, self.ffn_down_shexp.weight)).sum(3).cast(dtypes.bfloat16)
+        else:
+          shexp = self.ffn_down_shexp(shared_activation)
+          if hasattr(self, 'ffn_gate_inp_shexp'):
+            shexp = shexp * (x * self.ffn_gate_inp_shexp["weight"]).sum(axis=-1, keepdim=True).sigmoid()
+          out = out + shexp
       return out
     # TODO: remove the need for this contiguous
-    return self.ffn_down(self._activation(self.ffn_gate(x), self.ffn_up(x)).contiguous())
+    return self.ffn_down(self._activation(self.ffn_gate(x).contiguous(), self.ffn_up(x).contiguous()).contiguous())
 
   # given the token-prefix match, return how much cached state this block can still reuse
   def _reusable_prefix_len(self, prefix_len:int, cached_len:int) -> int: return prefix_len
@@ -388,7 +413,10 @@ class GatedDeltaNetBlock(FFNBlock):
     # update is fused into one kernel so prefill doesn't build a Python-unrolled graph.
     split_qkv = hasattr(self, "attn_q")
     if split_qkv:
-      projected_q, projected_k, projected_v = self.attn_q(x), self.attn_k(x), self.attn_v(x)
+      if resolve(T == 1) and amd_custom_kernels_supported(x.device) and \
+         self.attn_q.weight.shape == self.attn_k.weight.shape == self.attn_v.weight.shape:
+        projected_q, projected_k, projected_v = kda_qkv_linear(x, self.attn_q.weight, self.attn_k.weight, self.attn_v.weight)
+      else: projected_q, projected_k, projected_v = self.attn_q(x), self.attn_k(x), self.attn_v(x)
       # Snapshot mutable caches before constructing the recurrence. Otherwise the final store can
       # overwrite their buffers before earlier outputs in a multi-token lazy graph consume them.
       conv_state_q, conv_state_k, conv_state_v = self.conv_state_q.clone(), self.conv_state_k.clone(), self.conv_state_v.clone()
@@ -488,8 +516,14 @@ class Transformer:
     # we specialize the JIT for prefill and rollout
     self.prefill_jit = TinyJit(self.forward)
     self.rollout_jit = TinyJit(self.forward)
+    self.greedy_prefill_jit = TinyJit(self.forward)
+    self.greedy_rollout_jit = TinyJit(self.forward)
+    self.reset_jit = TinyJit(self._reset_state)
 
-  def forward(self, tokens:Tensor, start_pos:int|UOp, temperature:Tensor) -> Tensor:
+  def _reset_state(self) -> None:
+    if resets := [r for b in self.blk for r in b._state_reset_ops()]: Tensor.realize(*resets)
+
+  def forward(self, tokens:Tensor, start_pos:int|UOp, temperature:Tensor|None) -> Tensor:
     if len(tokens.shape) == 1: tokens = tokens.reshape(1, -1)
     x = self.token_embd(tokens).cast(dtypes.bfloat16) if self.config.bf16_activations else self.token_embd(tokens).float()
     block_residual = Tensor.zeros(x.shape[0]*x.shape[1], 0, x.shape[2], device=x.device, dtype=x.dtype) \
@@ -505,12 +539,18 @@ class Transformer:
     if block_residual is not None:
       x = FFNBlock._apply_attn_res(x.reshape(-1, x.shape[-1]), block_residual,
                                    self.output_attn_res_proj, self.output_attn_res_norm).reshape(x.shape)
-    logits = self.output(self.output_norm(x))[:, -1, :]
+    final_x = self.output_norm(x)
+    if temperature is None and resolve(tokens.numel() == 1) and amd_custom_kernels_supported(x.device):
+      return bf16_matvec(final_x, self.output.weight).argmax(-1, keepdim=True)
+    logits = self.output(final_x)[:, -1, :]
+    if temperature is None: return logits.argmax(-1, keepdim=True)
     # Gumbel-max trick: argmax(logits/temp - log(-log(uniform))) is equivalent to sampling from softmax(logits/temp)
     return (logits / temperature.maximum(1e-12) - (Tensor.rand_like(logits).maximum(1e-12).log().neg()).log()).argmax(-1, keepdim=True)
 
-  def __call__(self, tokens:Tensor, start_pos:int|UOp, temperature:Tensor) -> Tensor:
+  def __call__(self, tokens:Tensor, start_pos:int|UOp, temperature:Tensor|None) -> Tensor:
     token_count = tokens.numel()
+    if temperature is None:
+      return (self.greedy_prefill_jit if resolve(token_count != 1) else self.greedy_rollout_jit)(tokens.flatten().contiguous(), start_pos, None)
     return (self.prefill_jit if resolve(token_count != 1) else self.rollout_jit)(tokens.flatten().contiguous(), start_pos, temperature)
 
   @staticmethod
@@ -612,12 +652,12 @@ class Transformer:
     v_toks = UOp.variable("toks", 1, chunk_size)
     # TODO: use UOp.variable for temperature once float variables are supported
     model_device = self.token_embd.weight.device
-    temp = Tensor([temperature], device=model_device)
+    temp = None if temperature == 0.0 else Tensor([temperature], device=model_device)
     # assign all input tokens once, then slice from start_pos for the model call
     t = Tensor(tokens + [0] * (self.max_context - len(tokens)), dtype="int32", device=model_device).reshape(1, self.max_context)
     # recompute start_pos from what's currently valid in the caches
     start_pos = self.get_start_pos(tokens)
-    if start_pos < len(self._cached_tokens) and (resets := [r for b in self.blk for r in b._state_reset_ops()]): Tensor.realize(*resets)
+    if start_pos < len(self._cached_tokens) and self.has_recurrent_block: self.reset_jit()
     out, prompt_len = None, len(tokens)
     while len(tokens) < self.max_context:
       n_toks = min(chunk_size, len(tokens) - start_pos)