From 725041aabc9935e4b06bde3d86acfde54c66e0e0 Mon Sep 17 00:00:00 2001 From: BlivionIaG <15323295+BlivionIaG@users.noreply.github.com> Date: Tue, 27 Oct 2026 05:27:38 +0100 Subject: [PATCH 1/3] EXL3 Flash Next on CUDA: one rank, experts resident or offloaded to pinned host RAM, sm_120 gate --- docs/kernel-contract.md | 2 + tools/exl3_len_abba.py | 97 +++ tools/exl3_longctx_client.py | 131 ++++ tools/exl3_pack_sizes.py | 26 + tools/exl3_serve.sh | 13 + tools/exl3_serve_gate.py | 56 ++ tools/exl3_single_request_bench.py | 48 ++ tools/exl3_snapshot_gate.py | 91 +++ tools/exl3_text_abba.py | 45 ++ tools/tests/test_exl3_snapshot_gate.py | 86 +++ zig/build/cuda.zig | 3 +- zig/build/cuda_kernels.zig | 1 + .../cuda/flashnext_exl3/flashnext_exl3.cu | 15 + .../core/contract/flashnext_exl3_pending.zig | 33 + zig/src/core/contract/pending_132.zig | 4 + zig/src/core/contract/registry.zig | 3 +- .../core/contract/registry_pending_test.zig | 17 + zig/src/core/contract/registry_sets_test.zig | 3 +- zig/src/cuda/exl3_dma.zig | 38 ++ zig/src/cuda/exl3_dma_test.zig | 216 +++++++ zig/src/cuda/exl3_pool.zig | 331 ++++++++++ zig/src/cuda/kernels.zig | 1 + zig/src/cuda/root.zig | 2 + zig/src/families/flashnext/cuda_batch.zig | 35 +- .../families/flashnext/cuda_batch_gather.zig | 40 ++ zig/src/families/flashnext/cuda_batch_ple.zig | 32 + .../families/flashnext/cuda_batch_state.zig | 238 ++++--- zig/src/families/flashnext/cuda_config.zig | 9 +- zig/src/families/flashnext/cuda_embed.zig | 22 + zig/src/families/flashnext/cuda_lanes.zig | 67 +- .../families/flashnext/cuda_lanes_bytes.zig | 3 + .../families/flashnext/cuda_lanes_load.zig | 38 +- .../families/flashnext/cuda_lanes_snap.zig | 51 ++ zig/src/families/flashnext/cuda_launch.zig | 10 + .../flashnext/cuda_paging_batch_test.zig | 2 + zig/src/families/flashnext/cuda_ple.zig | 13 +- .../families/flashnext/cuda_projection.zig | 3 + .../flashnext/cuda_projection_load.zig | 3 + zig/src/families/flashnext/cuda_prompt.zig | 102 +-- .../flashnext/cuda_prompt_lifecycle.zig | 14 + .../families/flashnext/cuda_prompt_load.zig | 16 +- .../flashnext/cuda_prompt_telemetry.zig | 233 +++++++ .../families/flashnext/cuda_prompt_test.zig | 53 ++ zig/src/families/flashnext/cuda_prompt_wt.zig | 29 +- .../flashnext/cuda_scratch_layout_test.zig | 41 ++ zig/src/families/flashnext/cuda_snap.zig | 41 +- .../flashnext/exl3/kv/cuda_exl3_kv.zig | 83 +++ .../flashnext/exl3/kv/cuda_exl3_kv_test.zig | 20 + .../families/flashnext/exl3/kv/exl3_kv.zig | 98 +++ .../flashnext/exl3/model/cuda_exl3.zig | 31 + .../flashnext/exl3/model/cuda_exl3_arch.zig | 32 + .../exl3/model/cuda_exl3_dispatch_test.zig | 74 +++ .../exl3/model/cuda_exl3_fake_driver_test.zig | 23 + .../flashnext/exl3/model/cuda_exl3_pass.zig | 374 +++++++++++ .../flashnext/exl3/model/cuda_exl3_ple.zig | 543 ++++++++++++++++ .../exl3/model/cuda_exl3_telemetry.zig | 85 +++ .../exl3/model/cuda_exl3_telemetry_test.zig | 38 ++ .../families/flashnext/exl3/model/exl3.zig | 33 + .../flashnext/exl3/model/exl3_attn_cpu.zig | 441 +++++++++++++ .../flashnext/exl3/model/exl3_cache.zig | 168 +++++ .../flashnext/exl3/model/exl3_cpu.zig | 380 +++++++++++ .../flashnext/exl3/model/exl3_fwd_cpu.zig | 414 ++++++++++++ .../flashnext/exl3/model/exl3_gdn_cpu.zig | 399 ++++++++++++ .../families/flashnext/exl3/model/exl3_hc.zig | 253 ++++++++ .../flashnext/exl3/model/exl3_moe.zig | 147 +++++ .../flashnext/exl3/model/exl3_norm.zig | 90 +++ .../flashnext/exl3/model/exl3_round.zig | 599 ++++++++++++++++++ .../flashnext/exl3/model/exl3_route.zig | 94 +++ .../flashnext/exl3/model/exl3_router.zig | 252 ++++++++ .../flashnext/exl3/model/exl3_telemetry.zig | 226 +++++++ .../exl3/model/exl3_telemetry_test.zig | 96 +++ .../exl3/offload/cuda_exl3_frequency.zig | 122 ++++ .../offload/cuda_exl3_frequency_state.zig | 444 +++++++++++++ .../exl3/offload/cuda_exl3_frequency_test.zig | 299 +++++++++ .../flashnext/exl3/offload/cuda_exl3_pool.zig | 345 ++++++++++ .../exl3/offload/exl3_frequency_seed.zig | 139 ++++ .../exl3/offload/exl3_line_deltas.zig | 120 ++++ .../flashnext/exl3/offload/exl3_offload.zig | 52 ++ .../exl3/offload/exl3_offload_counts.zig | 53 ++ .../exl3/offload/exl3_offload_fill.zig | 145 +++++ .../exl3/offload/exl3_offload_fill_test.zig | 207 ++++++ .../exl3/offload/exl3_offload_policy.zig | 191 ++++++ .../exl3/offload/exl3_offload_policy_test.zig | 359 +++++++++++ .../exl3/offload/exl3_offload_used.zig | 146 +++++ .../exl3/offload/exl3_offload_used_test.zig | 276 ++++++++ .../exl3/offload/exl3_offload_vram.zig | 216 +++++++ .../exl3/offload/exl3_offload_vram_test.zig | 246 +++++++ .../exl3/offload/exl3_profile_line_budget.zig | 468 ++++++++++++++ .../exl3/offload/exl3_recent_demand.zig | 212 +++++++ .../flashnext/exl3/pack/cuda_exl3_config.zig | 113 ++++ .../flashnext/exl3/pack/cuda_exl3_load.zig | 527 +++++++++++++++ .../flashnext/exl3/pack/cuda_exl3_memory.zig | 134 ++++ .../flashnext/exl3/pack/cuda_exl3_store.zig | 121 ++++ .../flashnext/exl3/pack/exl3_format.zig | 1 + .../flashnext/exl3/pack/exl3_place.zig | 309 +++++++++ .../flashnext/exl3/pack/exl3_qmat.zig | 229 +++++++ .../flashnext/exl3/pack/exl3_scan.zig | 206 ++++++ .../flashnext/exl3/pack/exl3_shape.zig | 66 ++ .../exl3/pack/exl3_startup_memory.zig | 72 +++ .../exl3/pack/fixtures/exl3_config.json | 168 +++++ .../flashnext/exl3/tp/cuda_exl3_fork.zig | 126 ++++ .../flashnext/exl3/tp/cuda_exl3_fork_test.zig | 105 +++ zig/src/native/cuda.zig | 5 +- zig/src/native/cuda_snaps.zig | 44 +- zig/src/native/cuda_startup_test.zig | 14 +- zig/tests/cuda/contract_limits.zig | 2 + zig/tests/cuda/flashnext_exl3/kv/exl3_kv.zig | 1 + .../flashnext_exl3/model/exl3_attn_cpu.zig | 1 + .../cuda/flashnext_exl3/model/exl3_cache.zig | 1 + .../cuda/flashnext_exl3/model/exl3_cpu.zig | 1 + .../flashnext_exl3/model/exl3_fwd_cpu.zig | 1 + .../flashnext_exl3/model/exl3_gdn_cpu.zig | 1 + .../cuda/flashnext_exl3/model/exl3_hc.zig | 1 + .../cuda/flashnext_exl3/model/exl3_moe.zig | 1 + .../flashnext_exl3/model/exl3_ple2_tests.zig | 598 +++++++++++++++++ .../cuda/flashnext_exl3/model/exl3_route.zig | 1 + .../cuda/flashnext_exl3/model/exl3_router.zig | 1 + .../model/flashnext_exl3_engine_tests.zig | 536 ++++++++++++++++ .../model/flashnext_exl3_layer_tests.zig | 386 +++++++++++ .../model/flashnext_exl3_names_tests.zig | 374 +++++++++++ .../model/flashnext_exl3_proj_tests.zig | 269 ++++++++ .../model/flashnext_exl3_rig.zig | 251 ++++++++ .../model/flashnext_exl3_tests.zig | 50 ++ .../offload/exl3_pool_bench_tests.zig | 255 ++++++++ .../flashnext_exl3/offload/exl3_pool_data.zig | 284 +++++++++ .../offload/exl3_pool_equal_tests.zig | 151 +++++ .../offload/exl3_pool_hits_tests.zig | 166 +++++ .../offload/exl3_pool_tests.zig | 57 ++ .../cuda/flashnext_exl3/pack/exl3_format.zig | 1 + .../flashnext_exl3/pack/exl3_pack_tests.zig | 522 +++++++++++++++ .../cuda/flashnext_exl3/pack/exl3_place.zig | 1 + .../cuda/flashnext_exl3/pack/exl3_qmat.zig | 1 + .../cuda/flashnext_exl3/pack/exl3_scan.zig | 1 + .../cuda/flashnext_exl3/pack/exl3_shape.zig | 1 + .../prefill/exl3_prefill_copy.zig | 157 +++++ .../prefill/exl3_prefill_dense.zig | 171 +++++ .../prefill/exl3_prefill_experts.zig | 118 ++++ .../prefill/exl3_prefill_pack.zig | 218 +++++++ .../prefill/exl3_prefill_tests.zig | 37 ++ zig/tests/cuda/main.zig | 18 + 140 files changed, 17967 insertions(+), 298 deletions(-) create mode 100644 tools/exl3_len_abba.py create mode 100644 tools/exl3_longctx_client.py create mode 100644 tools/exl3_pack_sizes.py create mode 100755 tools/exl3_serve.sh create mode 100755 tools/exl3_serve_gate.py create mode 100644 tools/exl3_single_request_bench.py create mode 100644 tools/exl3_snapshot_gate.py create mode 100644 tools/exl3_text_abba.py create mode 100644 tools/tests/test_exl3_snapshot_gate.py create mode 100644 zig/kernels/cuda/flashnext_exl3/flashnext_exl3.cu create mode 100644 zig/src/core/contract/flashnext_exl3_pending.zig create mode 100644 zig/src/cuda/exl3_dma.zig create mode 100644 zig/src/cuda/exl3_dma_test.zig create mode 100644 zig/src/cuda/exl3_pool.zig create mode 100644 zig/src/families/flashnext/cuda_batch_gather.zig create mode 100644 zig/src/families/flashnext/cuda_batch_ple.zig create mode 100644 zig/src/families/flashnext/cuda_lanes_snap.zig create mode 100644 zig/src/families/flashnext/cuda_prompt_lifecycle.zig create mode 100644 zig/src/families/flashnext/cuda_prompt_telemetry.zig create mode 100644 zig/src/families/flashnext/cuda_prompt_test.zig create mode 100644 zig/src/families/flashnext/cuda_scratch_layout_test.zig create mode 100644 zig/src/families/flashnext/exl3/kv/cuda_exl3_kv.zig create mode 100644 zig/src/families/flashnext/exl3/kv/cuda_exl3_kv_test.zig create mode 100644 zig/src/families/flashnext/exl3/kv/exl3_kv.zig create mode 100644 zig/src/families/flashnext/exl3/model/cuda_exl3.zig create mode 100644 zig/src/families/flashnext/exl3/model/cuda_exl3_arch.zig create mode 100644 zig/src/families/flashnext/exl3/model/cuda_exl3_dispatch_test.zig create mode 100644 zig/src/families/flashnext/exl3/model/cuda_exl3_fake_driver_test.zig create mode 100644 zig/src/families/flashnext/exl3/model/cuda_exl3_pass.zig create mode 100644 zig/src/families/flashnext/exl3/model/cuda_exl3_ple.zig create mode 100644 zig/src/families/flashnext/exl3/model/cuda_exl3_telemetry.zig create mode 100644 zig/src/families/flashnext/exl3/model/cuda_exl3_telemetry_test.zig create mode 100644 zig/src/families/flashnext/exl3/model/exl3.zig create mode 100644 zig/src/families/flashnext/exl3/model/exl3_attn_cpu.zig create mode 100644 zig/src/families/flashnext/exl3/model/exl3_cache.zig create mode 100644 zig/src/families/flashnext/exl3/model/exl3_cpu.zig create mode 100644 zig/src/families/flashnext/exl3/model/exl3_fwd_cpu.zig create mode 100644 zig/src/families/flashnext/exl3/model/exl3_gdn_cpu.zig create mode 100644 zig/src/families/flashnext/exl3/model/exl3_hc.zig create mode 100644 zig/src/families/flashnext/exl3/model/exl3_moe.zig create mode 100644 zig/src/families/flashnext/exl3/model/exl3_norm.zig create mode 100644 zig/src/families/flashnext/exl3/model/exl3_round.zig create mode 100644 zig/src/families/flashnext/exl3/model/exl3_route.zig create mode 100644 zig/src/families/flashnext/exl3/model/exl3_router.zig create mode 100644 zig/src/families/flashnext/exl3/model/exl3_telemetry.zig create mode 100644 zig/src/families/flashnext/exl3/model/exl3_telemetry_test.zig create mode 100644 zig/src/families/flashnext/exl3/offload/cuda_exl3_frequency.zig create mode 100644 zig/src/families/flashnext/exl3/offload/cuda_exl3_frequency_state.zig create mode 100644 zig/src/families/flashnext/exl3/offload/cuda_exl3_frequency_test.zig create mode 100644 zig/src/families/flashnext/exl3/offload/cuda_exl3_pool.zig create mode 100644 zig/src/families/flashnext/exl3/offload/exl3_frequency_seed.zig create mode 100644 zig/src/families/flashnext/exl3/offload/exl3_line_deltas.zig create mode 100644 zig/src/families/flashnext/exl3/offload/exl3_offload.zig create mode 100644 zig/src/families/flashnext/exl3/offload/exl3_offload_counts.zig create mode 100644 zig/src/families/flashnext/exl3/offload/exl3_offload_fill.zig create mode 100644 zig/src/families/flashnext/exl3/offload/exl3_offload_fill_test.zig create mode 100644 zig/src/families/flashnext/exl3/offload/exl3_offload_policy.zig create mode 100644 zig/src/families/flashnext/exl3/offload/exl3_offload_policy_test.zig create mode 100644 zig/src/families/flashnext/exl3/offload/exl3_offload_used.zig create mode 100644 zig/src/families/flashnext/exl3/offload/exl3_offload_used_test.zig create mode 100644 zig/src/families/flashnext/exl3/offload/exl3_offload_vram.zig create mode 100644 zig/src/families/flashnext/exl3/offload/exl3_offload_vram_test.zig create mode 100644 zig/src/families/flashnext/exl3/offload/exl3_profile_line_budget.zig create mode 100644 zig/src/families/flashnext/exl3/offload/exl3_recent_demand.zig create mode 100644 zig/src/families/flashnext/exl3/pack/cuda_exl3_config.zig create mode 100644 zig/src/families/flashnext/exl3/pack/cuda_exl3_load.zig create mode 100644 zig/src/families/flashnext/exl3/pack/cuda_exl3_memory.zig create mode 100644 zig/src/families/flashnext/exl3/pack/cuda_exl3_store.zig create mode 120000 zig/src/families/flashnext/exl3/pack/exl3_format.zig create mode 100644 zig/src/families/flashnext/exl3/pack/exl3_place.zig create mode 100644 zig/src/families/flashnext/exl3/pack/exl3_qmat.zig create mode 100644 zig/src/families/flashnext/exl3/pack/exl3_scan.zig create mode 100644 zig/src/families/flashnext/exl3/pack/exl3_shape.zig create mode 100644 zig/src/families/flashnext/exl3/pack/exl3_startup_memory.zig create mode 100644 zig/src/families/flashnext/exl3/pack/fixtures/exl3_config.json create mode 100644 zig/src/families/flashnext/exl3/tp/cuda_exl3_fork.zig create mode 100644 zig/src/families/flashnext/exl3/tp/cuda_exl3_fork_test.zig create mode 120000 zig/tests/cuda/flashnext_exl3/kv/exl3_kv.zig create mode 120000 zig/tests/cuda/flashnext_exl3/model/exl3_attn_cpu.zig create mode 120000 zig/tests/cuda/flashnext_exl3/model/exl3_cache.zig create mode 120000 zig/tests/cuda/flashnext_exl3/model/exl3_cpu.zig create mode 120000 zig/tests/cuda/flashnext_exl3/model/exl3_fwd_cpu.zig create mode 120000 zig/tests/cuda/flashnext_exl3/model/exl3_gdn_cpu.zig create mode 120000 zig/tests/cuda/flashnext_exl3/model/exl3_hc.zig create mode 120000 zig/tests/cuda/flashnext_exl3/model/exl3_moe.zig create mode 100644 zig/tests/cuda/flashnext_exl3/model/exl3_ple2_tests.zig create mode 120000 zig/tests/cuda/flashnext_exl3/model/exl3_route.zig create mode 120000 zig/tests/cuda/flashnext_exl3/model/exl3_router.zig create mode 100644 zig/tests/cuda/flashnext_exl3/model/flashnext_exl3_engine_tests.zig create mode 100644 zig/tests/cuda/flashnext_exl3/model/flashnext_exl3_layer_tests.zig create mode 100644 zig/tests/cuda/flashnext_exl3/model/flashnext_exl3_names_tests.zig create mode 100644 zig/tests/cuda/flashnext_exl3/model/flashnext_exl3_proj_tests.zig create mode 100644 zig/tests/cuda/flashnext_exl3/model/flashnext_exl3_rig.zig create mode 100644 zig/tests/cuda/flashnext_exl3/model/flashnext_exl3_tests.zig create mode 100644 zig/tests/cuda/flashnext_exl3/offload/exl3_pool_bench_tests.zig create mode 100644 zig/tests/cuda/flashnext_exl3/offload/exl3_pool_data.zig create mode 100644 zig/tests/cuda/flashnext_exl3/offload/exl3_pool_equal_tests.zig create mode 100644 zig/tests/cuda/flashnext_exl3/offload/exl3_pool_hits_tests.zig create mode 100644 zig/tests/cuda/flashnext_exl3/offload/exl3_pool_tests.zig create mode 120000 zig/tests/cuda/flashnext_exl3/pack/exl3_format.zig create mode 100644 zig/tests/cuda/flashnext_exl3/pack/exl3_pack_tests.zig create mode 120000 zig/tests/cuda/flashnext_exl3/pack/exl3_place.zig create mode 120000 zig/tests/cuda/flashnext_exl3/pack/exl3_qmat.zig create mode 120000 zig/tests/cuda/flashnext_exl3/pack/exl3_scan.zig create mode 120000 zig/tests/cuda/flashnext_exl3/pack/exl3_shape.zig create mode 100644 zig/tests/cuda/flashnext_exl3/prefill/exl3_prefill_copy.zig create mode 100644 zig/tests/cuda/flashnext_exl3/prefill/exl3_prefill_dense.zig create mode 100644 zig/tests/cuda/flashnext_exl3/prefill/exl3_prefill_experts.zig create mode 100644 zig/tests/cuda/flashnext_exl3/prefill/exl3_prefill_pack.zig create mode 100644 zig/tests/cuda/flashnext_exl3/prefill/exl3_prefill_tests.zig diff --git a/docs/kernel-contract.md b/docs/kernel-contract.md index 1ed43ac146..8f3a676fc4 100644 --- a/docs/kernel-contract.md +++ b/docs/kernel-contract.md @@ -215,6 +215,8 @@ shared memory, the device's limits and the planted controls are checked the same A restored 13.0 receipt can arrive after the later compiler sets. Such records mark `later_pending` so 13.0 and 13.1 use their base usage, while 13.2 and later keep their independently measured or pending selection. A missing later receipt remains pending; a 13.0 measurement never substitutes for it. +Flash Next's EXL3 launch (`flashnext_exl3_pending.zig`) carries only an nvcc 13.4 receipt: it stays pending under +13.0 to 13.3 and takes its 13.4 usage from `cuda_134`, which `pending_132.lookup` serves for that release. Gemma 4 CUDA tower QKV norms and nine-term pooling register FP32 sums and BF16 outputs. Their contract keeps the exact FP32 mirror and adds independent FP64 ranges at the BF16 nodes. diff --git a/tools/exl3_len_abba.py b/tools/exl3_len_abba.py new file mode 100644 index 0000000000..fe4cafce50 --- /dev/null +++ b/tools/exl3_len_abba.py @@ -0,0 +1,97 @@ +"""Concurrent two-request HTTP bench (repeated-ID prompts) used for the EXL3 ABBA runs: PORT IDS.json OUT.json LENGTH DRAFT.""" +import concurrent.futures +import json +import statistics +import sys +import time +import urllib.request + + +def one(port, prompt, tokens=128, draft=False): + body = { + "prompt": prompt, + "max_tokens": tokens, + "temperature": 0, + "ignore_eos": True, + "stream": True, + "stream_options": {"include_usage": True}, + "draft": draft, + } + request = urllib.request.Request( + f"http://127.0.0.1:{port}/v1/completions", + json.dumps(body).encode(), + {"Content-Type": "application/json"}, + ) + started = time.perf_counter() + first = last = None + usage = {} + runtime = {} + text_events = 0 + complete = False + with urllib.request.urlopen(request, timeout=1800) as response: + for raw in response: + line = raw.decode().strip() + if line == "data: [DONE]": + complete = True + break + if not line.startswith("data:"): + continue + event = json.loads(line[5:]) + usage = event.get("usage") or usage + runtime = event.get("tensorfold") or runtime + for choice in event.get("choices", []): + piece = choice.get("text") or (choice.get("delta") or {}).get("content") or "" + if piece: + now = time.perf_counter() + first = first or now + last = now + text_events += 1 + generated = int(usage.get("completion_tokens", 0)) + span = (last - first) if first is not None and last is not None else 0.0 + return { + "complete": complete, + "prompt_tokens": usage.get("prompt_tokens"), + "tokens": generated, + "ttft_s": (first - started) if first is not None else None, + "first_token_at": first, + "last_token_at": last, + "prefill_s": runtime.get("prefill_seconds"), + "decode_tps_server": runtime.get("tokens_per_second"), + "decode_tps_client": ((generated - 1) / span) if generated > 1 and span > 0 else None, + "itl_ms_client": (1000 * span / (generated - 1)) if generated > 1 and span > 0 else None, + "token_sha256": runtime.get("token_sha256"), + "text_events": text_events, + } + + +def batch(port, prompt, tokens, draft=False): + start = time.perf_counter() + with concurrent.futures.ThreadPoolExecutor(max_workers=2) as pool: + results = list(pool.map(lambda _: one(port, prompt, tokens, draft), range(2))) + end = time.perf_counter() + total = sum(row["tokens"] for row in results) + first_times = [row["first_token_at"] for row in results if row["first_token_at"] is not None] + last_times = [row["last_token_at"] for row in results if row["last_token_at"] is not None] + decode_span = max(last_times) - min(first_times) if first_times and last_times else 0.0 + return { + "requests": results, + "batch_wall_s": end - start, + "aggregate_decode_tps": total / decode_span if decode_span > 0 else None, + "end_to_end_aggregate_tps": total / (end - start) if end > start else None, + "mean_prefill_s": statistics.mean(row["prefill_s"] for row in results if row["prefill_s"] is not None), + } + + +if __name__ == "__main__": + port = int(sys.argv[1]) + with open(sys.argv[2], encoding="utf-8") as source: + base = json.load(source) + length = int(sys.argv[4]); draft = sys.argv[5] == "1" + ids = (base * (length // len(base) + 1))[:length] + report = {"prompt_tokens": len(ids), "tokens": 128, "drafts": draft, "batches": []} + for phase in ("warmup-1", "measure-1", "measure-2"): + result = batch(port, ids, int(__import__("os").environ.get("TOK", "128")), draft) + report["batches"].append({"phase": phase, **result}) + print(json.dumps(report["batches"][-1]), flush=True) + with open(sys.argv[3], "w", encoding="utf-8") as target: + json.dump(report, target, indent=2) diff --git a/tools/exl3_longctx_client.py b/tools/exl3_longctx_client.py new file mode 100644 index 0000000000..8ef9a845bb --- /dev/null +++ b/tools/exl3_longctx_client.py @@ -0,0 +1,131 @@ +#!/usr/bin/env python3 +"""Long-context checks against tensorfold-native serve (EXL3 Flash Next, docs/plans/exl3-kv-memory.md). + +usage: + exl3_longctx_client.py PORT needle CTX[,CTX..] CONC TEXT + A haystack of real text (the file TEXT) cut to about CTX prompt tokens (counted by the + server's /tokenize on the rendered chat), one needle sentence at a depth, a question at the end; greedy chat + reply (thinking off). CONC requests at once, each with its own needle, depth and text offset. Every request + runs drafts off, then all again drafts on: the answer must hold the needle's value and the drafted reply must + equal the undrafted one (token_sha256). + exl3_longctx_client.py PORT cont N[,N..] GEN IDS + The first N real-text ids (the little-endian u32 file IDS), GEN greedy tokens, drafts off: prints the reply's + sha and ids (for comparing two servers' continuations).""" +import hashlib, json, os, struct, sys, threading, time, urllib.request + +port, mode = int(sys.argv[1]), sys.argv[2] +base = f"http://127.0.0.1:{port}" + + +def post(path, body, timeout=7200): + req = urllib.request.Request(base + path, json.dumps(body).encode(), {"Content-Type": "application/json"}) + return json.load(urllib.request.urlopen(req, timeout=timeout)) + + +NEEDLES = [ + ("The secret code for the blue vault is 7391-4826.", "What is the secret code for the blue vault?", "7391-4826"), + ("The password for the red door is amber-heron-55.", "What is the password for the red door?", "amber-heron-55"), + ("The launch date written on the green folder is 14 March 2031.", "What launch date is written on the green folder?", "14 March 2031"), + ("The name of the night watchman at the north gate is Oskar Lindqvist.", "What is the name of the night watchman at the north gate?", "Lindqvist"), +] +DEPTHS = [0.5, 0.25, 0.75, 0.1] + + +def messages(hay, needle, question): + content = ("Read the following document carefully. Somewhere in it is one sentence that does not belong to it.\n\n" + + hay + "\n\nQuestion: " + question + " Answer with the exact value only.") + return [{"role": "user", "content": content}] + + +def build(text, ctx, i): + needle, question, answer = NEEDLES[i % len(NEEDLES)] + depth = DEPTHS[i % len(DEPTHS)] + off = (i * 97_000) % max(1, len(text) // 2) + chars = int(ctx * 3.3) + for _ in range(4): + hay = text[off:off + chars] + if len(hay) < chars: + hay = (text[off:] + "\n" + text)[:chars] + cut = int(len(hay) * depth) + cut = hay.rfind("\n", 0, cut) + 1 or cut + doc = hay[:cut] + needle + "\n" + hay[cut:] + msgs = messages(doc, needle, question) + n = post("/v1/tokenize", {"messages": msgs, "chat_template_kwargs": {"enable_thinking": False}})["count"] + if abs(n - ctx) <= max(64, ctx // 200) and n <= ctx: + break + chars = int(chars * (ctx - 48) / n) + return msgs, answer, n + + +def chat(msgs, draft, max_tokens=32): + body = {"messages": msgs, "max_tokens": max_tokens, "temperature": 0, "chat_template_kwargs": {"enable_thinking": False}} + if not draft: + body["draft"] = False + t0 = time.time() + r = post("/v1/chat/completions", body) + r["_wall"] = time.time() - t0 + return r + + +def run_all(reqs, draft): + out = [None] * len(reqs) + + def one(k): + try: + out[k] = chat(reqs[k][0], draft) + except Exception as e: # noqa: BLE001 - a failed request is reported, not raised + out[k] = {"error": repr(e)} + + ts = [threading.Thread(target=one, args=(k,)) for k in range(len(reqs))] + t0 = time.time() + for t in ts: + t.start() + for t in ts: + t.join() + return out, time.time() - t0 + + +if mode == "needle": + ctxs = [int(x) for x in sys.argv[3].split(",")] + conc = int(sys.argv[4]) + text = open(sys.argv[5], encoding="utf-8", errors="replace").read() + all_ok = True + for ctx in ctxs: + reqs = [build(text, ctx, i) for i in range(conc)] + res = {} + for draft in (False, True): + out, wall = run_all(reqs, draft) + res[draft] = out + for k, r in enumerate(out): + if "error" in r: + print(f"NEEDLE ctx={ctx} req={k} drafts={'on' if draft else 'off'}: ERROR {r['error']}", flush=True) + all_ok = False + continue + tf, sp = r.get("tensorfold", {}), r.get("speculative", {}) + ans = r["choices"][0]["message"]["content"] + ok = reqs[k][1].lower() in ans.lower() + all_ok &= ok + print(f"NEEDLE ctx={ctx} req={k} prompt_tokens={r['usage']['prompt_tokens']} drafts={'on' if draft else 'off'}: " + f"{'CORRECT' if ok else 'WRONG'} answer={ans.strip()[:80]!r} want={reqs[k][1]!r} prefill {tf.get('prefill_seconds', 0):.1f} s " + f"decode {tf.get('tokens_per_second', 0):.1f} tok/s accepted {sp.get('accepted')}/{sp.get('drafted')} " + f"sha {tf.get('token_sha256', '')[:12]} wall {r['_wall']:.1f} s", flush=True) + print(f"NEEDLE ctx={ctx} conc={conc} drafts={'on' if draft else 'off'}: all requests done in {wall:.1f} s", flush=True) + for k in range(conc): + a, b = res[False][k], res[True][k] + same = "error" not in a and "error" not in b and a["tensorfold"]["token_sha256"] == b["tensorfold"]["token_sha256"] + all_ok &= same + print(f"NEEDLE ctx={ctx} req={k}: drafted {'==' if same else '!='} undrafted", flush=True) + print("NEEDLE ALL OK" if all_ok else "NEEDLE NOT ALL OK", flush=True) +elif mode == "cont": + ns = [int(x) for x in sys.argv[3].split(",")] + gen = int(sys.argv[4]) + raw = open(sys.argv[5], "rb").read() + ids = list(struct.unpack("<%dI" % (len(raw) // 4), raw)) + for n in ns: + body = {"prompt": ids[:n], "max_tokens": gen, "temperature": 0, "ignore_eos": True, "draft": False} + r = post("/v1/completions", body) + tf = r["tensorfold"] + print(f"CONT n={n} gen={gen}: sha {tf['token_sha256'][:16]} prefill {tf['prefill_seconds']:.1f} s decode {tf['tokens_per_second']:.1f} tok/s " + f"text {r['choices'][0]['text'][:100]!r}", flush=True) +else: + sys.exit(__doc__) diff --git a/tools/exl3_pack_sizes.py b/tools/exl3_pack_sizes.py new file mode 100644 index 0000000000..ccaadbccb0 --- /dev/null +++ b/tools/exl3_pack_sizes.py @@ -0,0 +1,26 @@ +"""Group an EXL3 pack's safetensors tensors by role and print their sizes (reads headers only). usage: exl3_pack_sizes.py PACK_DIR.""" +import json, struct, glob, collections, sys +P = sys.argv[1] # the pack directory +g = collections.Counter() +def cat(n): + for k in ("ngram_embedding", "ple_", ".ple."): + if k in n and "ngram" in n: return "ple n-gram embedding shards" + if ".ple." in n: return "ple other (projections)" + if "mtp." in n: return "mtp (head)" + if "experts." in n and "shared" not in n: return "routed experts" + if "shared_expert" in n: return "shared expert" + if "hyper_connection" in n: return "hyper-connection mixes" + if "linear_attn" in n: return "linear (DeltaNet) attention" + if "self_attn" in n: return "full attention" + if "embed_tokens" in n: return "embed_tokens" + if "lm_head" in n: return "lm_head" + if "router" in n or "gate." in n: return "router" + return "other: " + n.split(".")[-2] if "." in n else "other" +for f in sorted(glob.glob(P + "/*.safetensors")): + with open(f, "rb") as h: + n = struct.unpack(" 1 else 8080 +TOK = int(os.environ.get("TOK", "192")) +root = os.environ.get("TF_SRC") or os.path.dirname(os.path.dirname(os.path.abspath(__file__))) # the checkout whose sources are the corpus +a = t.corpus([root + "/zig/src/families/flashnext/*.zig", root + "/zig/src/core/**/*.zig"], 50000) +b = t.corpus([root + "/docs/**/*.md", root + "/*.md", root + "/zig/tests/**/*.zig"], 50000) + + +def one(prompt): + body = { + "model": os.environ.get("MODEL", "qwen3.8-flash-next-iq3_xxs"), + "messages": [{"role": "user", "content": ("%s\n%s" % (uuid.uuid4().hex if not os.environ.get("NOUUID") else "fixed", prompt))}], + "max_tokens": TOK, "temperature": 0, "ignore_eos": True, "stream": True, + "stream_options": {"include_usage": True}, + "chat_template_kwargs": {"enable_thinking": False}, + } + req = urllib.request.Request("http://127.0.0.1:%d/v1/chat/completions" % PORT, json.dumps(body).encode(), {"Content-Type": "application/json"}) + start = time.perf_counter(); first = last = None; usage = {}; pieces = 0; digest = hashlib.sha256() + with urllib.request.urlopen(req, timeout=1800) as r: + for raw in r: + line = raw.decode().strip() + if line == "data: [DONE]": + break + if not line.startswith("data:"): + continue + ev = json.loads(line[5:]) + usage = ev.get("usage") or usage + for c in ev.get("choices", []): + d = c.get("delta") or {} + if d.get("content") or d.get("reasoning_content"): + now = time.perf_counter(); first = first or now; last = now; pieces += 1; digest.update((d.get("content") or d.get("reasoning_content")).encode()) + n = int(usage.get("completion_tokens", 0)) + span = (last - first) if first and last else 0 + return {"prompt_tokens": usage.get("prompt_tokens"), "completion_tokens": n, + "ttft_s": round(first - start, 2) if first else None, + "decode_tps": round((n - 1) / span, 2) if n > 1 and span > 0 else None, + "text_sha": digest.hexdigest()[:12]} + + +short = "Write a long, detailed story about a lighthouse keeper who finds a message in a bottle. Keep writing until told to stop." +for name, prompt in (("short", short), ("zig", a), ("docs", b)): + for rep in range(3): + r = one(prompt); r["prompt"] = name; r["rep"] = rep + print(json.dumps(r), flush=True) diff --git a/tools/exl3_snapshot_gate.py b/tools/exl3_snapshot_gate.py new file mode 100644 index 0000000000..eb1ef83481 --- /dev/null +++ b/tools/exl3_snapshot_gate.py @@ -0,0 +1,91 @@ +#!/usr/bin/env python3 +"""Strict two-lane serving gate for BF16 paged prompt-cache restore.""" + +import argparse +import concurrent.futures +import json +import threading +import time +import urllib.request +from pathlib import Path + + +def main(argv=None): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("url") + parser.add_argument("--tokens", type=int, default=4099) + parser.add_argument("--decode", type=int, default=32) + parser.add_argument("--drafts", action="store_true") + parser.add_argument("--reference", type=Path) + parser.add_argument("--record-reference", action="store_true", help="record a cold baseline; use a fresh server with prompt cache disabled") + parser.add_argument("--output", type=Path) + args = parser.parse_args(argv) + if args.tokens < 32 or args.decode < 1: + parser.error("require at least 32 prompt tokens and one decode token") + if args.record_reference and not args.output: + parser.error("--record-reference requires --output") + if args.record_reference and args.reference: + parser.error("--record-reference and --reference are mutually exclusive") + shared = " cache" * max(args.tokens, 1) + history = [{"role": "user", "content": shared}, {"role": "assistant", "content": "I have reviewed the shared document."}] + prompts = [history + [{"role": "user", "content": f"Original ending for lane {lane}."}] for lane in range(2)] + reference = json.loads(args.reference.read_text()) if args.reference else None + report = {"prompt_tokens": args.tokens, "decode_tokens": args.decode, "drafts": args.drafts, "phases": {}} + if reference: + for field in ("prompt_tokens", "decode_tokens", "drafts"): + if reference[field] != report[field]: + raise AssertionError(f"reference configuration mismatch: {field}") + + def pair(request_prompts): + barrier = threading.Barrier(2) + + def ask(ids): + started = time.perf_counter() + body = {"messages": ids, "max_tokens": args.decode, "temperature": 0, "ignore_eos": True, "draft": args.drafts} + request = urllib.request.Request(args.url.rstrip("/") + "/v1/chat/completions", json.dumps(body).encode(), {"Content-Type": "application/json"}) + barrier.wait(timeout=10) + with urllib.request.urlopen(request, timeout=900) as reply: + response = json.load(reply) + usage = response["usage"] + if usage["prompt_tokens"] < args.tokens or usage["completion_tokens"] != args.decode: + raise AssertionError(f"inexact token counts: {usage}") + runtime = response["tensorfold"] + digest = runtime.get("token_sha256") or runtime["token_sha"] + if len(digest) < 12: + raise AssertionError("missing output token digest") + return { + "sha": digest, + "cached_tokens": usage["prompt_tokens_details"]["cached_tokens"], + "prompt_tokens": usage["prompt_tokens"], + "elapsed_seconds": time.perf_counter() - started, + "prefill_seconds": runtime.get("prefill_seconds"), + "tokens_per_second": runtime.get("tokens_per_second"), + } + + with concurrent.futures.ThreadPoolExecutor(max_workers=2) as pool: + return list(pool.map(ask, request_prompts)) + + branches = [history + [{"role": "user", "content": f"Edited ending for branch {lane}."}] for lane in range(2)] + for phase, ids in (("cold", prompts), ("warm", prompts), ("branch", branches), ("rewind", prompts)): + results = pair(ids) + for lane, result in enumerate(results): + if args.record_reference and result["cached_tokens"] != 0: + raise AssertionError(f"reference phase {phase} lane {lane} was not cold: {result}") + if not args.record_reference and phase != "cold" and not 0 < result["cached_tokens"] < result["prompt_tokens"]: + raise AssertionError(f"{phase} lane {lane} did not restore a cached prefix: {result}") + if phase in ("warm", "rewind") and result["sha"] != report["phases"]["cold"][lane]["sha"]: + raise AssertionError(f"{phase} lane {lane} differs from its cold continuation") + if reference and result["sha"] != reference["phases"][phase][lane]["sha"]: + raise AssertionError(f"{phase} lane {lane} differs from reference continuation") + report["phases"][phase] = results + print(json.dumps({"phase": phase, "lanes": results}), flush=True) + if args.output: + args.output.write_text(json.dumps(report, indent=2) + "\n") + if args.record_reference: + print("PASS: recorded no-cache cold continuations for cold, warm, branch and rewind inputs") + else: + print("PASS: exact prompt/reply lengths, two submitted lanes, warm/branch/rewind cache hits and cold-baseline parity") + + +if __name__ == "__main__": + main() diff --git a/tools/exl3_text_abba.py b/tools/exl3_text_abba.py new file mode 100644 index 0000000000..ff773c5405 --- /dev/null +++ b/tools/exl3_text_abba.py @@ -0,0 +1,45 @@ +"""Concurrent two-request bench with two different real-text prompts built from a checkout: PORT OUT.json CHARS.""" +import concurrent.futures, glob, json, os, statistics, sys, time, urllib.request +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +import exl3_len_abba as base + + +def corpus(patterns, chars): + out, total = [], 0 + for pattern in patterns: + for path in sorted(glob.glob(os.path.expanduser(pattern), recursive=True)): + try: + text = open(path, encoding="utf-8").read() + except Exception: + continue + out.append("// file: %s\n%s\n" % (os.path.basename(path), text)) + total += len(out[-1]) + if total >= chars: + return "".join(out)[:chars] + return "".join(out)[:chars] + + +def post(port, prompt, tokens): + return base.one(port, prompt, tokens, False) + + +if __name__ == "__main__": + port, out_path, chars, tokens = int(sys.argv[1]), sys.argv[2], int(sys.argv[3]), int(os.environ.get("TOK", "128")) + root = os.environ.get("TF_SRC") or os.path.dirname(os.path.dirname(os.path.abspath(__file__))) # the checkout whose sources are the corpus + a = corpus([root + "/zig/src/families/flashnext/*.zig", root + "/zig/src/core/**/*.zig"], chars) + b = corpus([root + "/docs/**/*.md", root + "/*.md", root + "/zig/tests/**/*.zig"], chars) + prompts = [a, b] + report = {"chars": [len(a), len(b)], "batches": []} + for phase in ("warmup-1", "measure-1", "measure-2"): + start = time.perf_counter() + with concurrent.futures.ThreadPoolExecutor(max_workers=2) as pool: + rows = list(pool.map(lambda p: post(port, p, tokens), prompts)) + batch = { + "phase": phase, + "requests": rows, + "mean_prefill_s": statistics.mean(r["prefill_s"] for r in rows), + "batch_wall_s": time.perf_counter() - start, + } + report["batches"].append(batch) + print(json.dumps(batch), flush=True) + json.dump(report, open(out_path, "w"), indent=2) diff --git a/tools/tests/test_exl3_snapshot_gate.py b/tools/tests/test_exl3_snapshot_gate.py new file mode 100644 index 0000000000..7363c7c345 --- /dev/null +++ b/tools/tests/test_exl3_snapshot_gate.py @@ -0,0 +1,86 @@ +"""Model-free checks for the strict chat-history snapshot gate.""" + +from contextlib import redirect_stderr, redirect_stdout +import hashlib +import importlib.util +import io +import json +from pathlib import Path +import tempfile +import unittest +from unittest.mock import patch + + +SCRIPT = Path(__file__).resolve().parents[1] / "exl3_snapshot_gate.py" +SPEC = importlib.util.spec_from_file_location("exl3_snapshot_gate", SCRIPT) +gate = importlib.util.module_from_spec(SPEC) +SPEC.loader.exec_module(gate) + + +class Response: + def __init__(self, payload): + self.payload = payload + + def __enter__(self): + return self + + def __exit__(self, *args): + return False + + def read(self): + return json.dumps(self.payload).encode() + + +class SnapshotGateTests(unittest.TestCase): + def invoke(self, args, cached): + requests = [] + + def open_request(request, timeout): + body = json.loads(request.data) + requests.append((request.full_url, body)) + digest = hashlib.sha256(json.dumps(body["messages"]).encode()).hexdigest() + return Response({ + "usage": {"prompt_tokens": 80, "completion_tokens": 2, + "prompt_tokens_details": {"cached_tokens": cached}}, + "tensorfold": {"token_sha256": digest}, + }) + + output, errors = io.StringIO(), io.StringIO() + with patch.object(gate.urllib.request, "urlopen", open_request), redirect_stdout(output), redirect_stderr(errors): + try: + gate.main(["http://example.invalid", "--tokens", "40", "--decode", "2", *args]) + code = 0 + except AssertionError as error: + errors.write(str(error)) + code = 1 + return code, output.getvalue(), errors.getvalue(), requests + + def test_records_cold_reference_then_requires_exact_chat_cache_reuse(self): + with tempfile.TemporaryDirectory() as directory: + reference = Path(directory) / "reference.json" + receipt = Path(directory) / "receipt.json" + code, _, _, cold_requests = self.invoke(["--record-reference", "--output", str(reference)], 0) + self.assertEqual(code, 0) + self.assertEqual(len(cold_requests), 8) + self.assertTrue(all(url.endswith("/v1/chat/completions") for url, _ in cold_requests)) + self.assertTrue(all(len(body["messages"]) == 3 for _, body in cold_requests)) + code, output, _, warm_requests = self.invoke(["--reference", str(reference), "--output", str(receipt)], 40) + self.assertEqual(code, 0) + self.assertEqual(len(warm_requests), 8) + self.assertIn("PASS:", output) + saved = json.loads(receipt.read_text()) + for phase in ("warm", "branch", "rewind"): + self.assertTrue(all(0 < lane["cached_tokens"] < lane["prompt_tokens"] + for lane in saved["phases"][phase])) + + def test_missing_warm_cache_hit_fails(self): + with tempfile.TemporaryDirectory() as directory: + reference = Path(directory) / "reference.json" + self.invoke(["--record-reference", "--output", str(reference)], 0) + code, output, errors, _ = self.invoke(["--reference", str(reference)], 0) + self.assertEqual(code, 1) + self.assertIn("did not restore a cached prefix", errors) + + +if __name__ == "__main__": + unittest.main() diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index 054d85922b..e53e74b715 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -433,7 +433,8 @@ pub fn hostTests(b: *std.Build, draft_ids: *std.Build.Module, build_options: *st golden_build.step(b, host, server_tests, native_modules.api, mods.tokenizer, template, native_modules.engines); @import("cortheon_fuzz.zig").add(b, host, native_modules.api, mods.tokenizer, template); } - // the tokenizer's own tests (memo, steps, regex, unicode) run here: no other test binary reaches them + // the tokenizer's own tests run here: no other test binary reaches them; the EXL3 hooks alone run as "exl3: ..." + b.step("test-flashnext-exl3-host", "Flash Next EXL3 admission, store, sizing and loader host tests").dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = flashnextModule(b, host, .debug, cuda, mods.core, mods.lanes, mods.qwen_vision, mods.api), .filters = &.{"exl3: "} })).step); for ([_]*std.Build.Module{ cuda, mods.core, mods.lanes, mods.nemotron, mods.qwen35, mods.qwen_vision, flashnextModule(b, host, .debug, cuda, mods.core, mods.lanes, mods.qwen_vision, mods.api), lfm2Module(b, host, .debug, cuda, mods.core), mods.radio, mods.heat, mods.tokenizer, native_modules.api, native, stagger(b, host, .debug) }) |m| step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = m })).step); const selected = b.option([]const u8, "modelopt-host-filter", "One portable ModelOpt regression") orelse "modelopt host:"; const run = b.addRunArtifact(b.addTest(.{ .root_module = mods.nemotron, .filters = &.{selected} })); diff --git a/zig/build/cuda_kernels.zig b/zig/build/cuda_kernels.zig index 887b2d84ba..3dc4c2a816 100644 --- a/zig/build/cuda_kernels.zig +++ b/zig/build/cuda_kernels.zig @@ -120,6 +120,7 @@ pub const kernels = [_]Kernel{ .{ .name = "drama3", .flags = &.{"-O3"} }, .{ .name = "drama3_fp8", .flags = &.{"-O3"} }, .{ .name = "drama3_codec", .flags = &.{"-O3"} }, + .{ .name = "flashnext_exl3", .src = "flashnext_exl3/flashnext_exl3", .flags = glue }, // ours: Flash Next's EXL3 glue (an fp16 row split exactly into two bf16) }; /// torch.utils.cpp_extension's own nvcc flags (torch 2.13): C++20 and which half/bf16 operators the headers define. diff --git a/zig/kernels/cuda/flashnext_exl3/flashnext_exl3.cu b/zig/kernels/cuda/flashnext_exl3/flashnext_exl3.cu new file mode 100644 index 0000000000..ae58b291de --- /dev/null +++ b/zig/kernels/cuda/flashnext_exl3/flashnext_exl3.cu @@ -0,0 +1,15 @@ +// The Flash Next family's EXL3 glue (cuda_exl3_pass.zig): an fp16 row as two bf16 rows whose sum is it exactly. + +#include +#include + +// hi = bf16(x), lo = bf16(x - hi): an fp16 has 11 significant bits, hi keeps 8 and the remainder has at most 3. +extern "C" __global__ void fn_exl3_split_f16(const __half* __restrict__ x, __nv_bfloat16* __restrict__ hi, + __nv_bfloat16* __restrict__ lo, long long n) { + const long long i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i >= n) return; + const float v = __half2float(x[i]); + const __nv_bfloat16 h = __float2bfloat16_rn(v); + hi[i] = h; + lo[i] = __float2bfloat16_rn(v - __bfloat162float(h)); +} diff --git a/zig/src/core/contract/flashnext_exl3_pending.zig b/zig/src/core/contract/flashnext_exl3_pending.zig new file mode 100644 index 0000000000..e8d0f82fec --- /dev/null +++ b/zig/src/core/contract/flashnext_exl3_pending.zig @@ -0,0 +1,33 @@ +//! Flash Next EXL3 launches (the fp16 row split); nvcc 13.4 measured, older pending. +const std = @import("std"); +const limits = @import("limits.zig"); +const registry = @import("registry.zig"); + +/// Launches with no nvcc 13.0 to 13.3 receipt yet: the geometry is recorded and usage stays pending until measured. +pub const records = [_]registry.CudaRecord{ + .{ .launch = .{ .name = "flashnext_exl3 fn_exl3_split_f16", .symbol = "fn_exl3_split_f16", .image = "flashnext_exl3", .threads = 256 }, .usage = registry.none }, +}; + +/// nvcc 13.4's measured registers, static shared bytes and launch bounds for sm_80, sm_86, sm_89, sm_120, sm_121. +pub const cuda_134 = [_]registry.Override{ + .{ .name = "flashnext_exl3 fn_exl3_split_f16", .usage = .{ .{ .regs = 10, .static = 0, .bound = 0 }, .{ .regs = 10, .static = 0, .bound = 0 }, .{ .regs = 10, .static = 0, .bound = 0 }, .{ .regs = 12, .static = 0, .bound = 0 }, .{ .regs = 12, .static = 0, .bound = 0 } } }, +}; + +test "Flash Next EXL3 launches: pending below nvcc 13.4, measured 13.4 resources fit every SM launch limit" { + const t = std.testing; + try t.expectEqual(records.len, cuda_134.len); + for (records, cuda_134) |r, o| { + try t.expectEqualStrings(r.launch.name, o.name); + try t.expectEqual(limits.sms.len, try registry.checkCuda(t.allocator, &.{r}, true)); + const measured = registry.usageFor(r, .{ .major = 13, .minor = 4 }); + try t.expect(std.meta.eql(measured, o.usage)); + var current = r; + current.usage = measured; + try t.expectEqual(@as(usize, 0), try registry.checkCuda(t.allocator, &.{current}, true)); + for (measured) |u| { + try t.expect(u.regs > 0); + try t.expect(u.bound == 0 or u.bound == r.launch.threads); + } + for (registry.cuda) |known| try t.expect(!std.mem.eql(u8, r.launch.symbol, known.launch.symbol)); + } +} diff --git a/zig/src/core/contract/pending_132.zig b/zig/src/core/contract/pending_132.zig index 09f0d0023a..11e92b7807 100644 --- a/zig/src/core/contract/pending_132.zig +++ b/zig/src/core/contract/pending_132.zig @@ -8,6 +8,10 @@ pub fn lookup(name: []const u8, release: limits.Nvcc) ?[limits.sms.len]limits.Us for (@import("measurements/qwen_lanes_133.zig").records ++ @import("measurements/pending_133.zig").records) |r| if (std.mem.eql(u8, r.name, name)) return r.usage; } + if (release.major == 13 and release.minor == 4) { + for (@import("flashnext_exl3_pending.zig").cuda_134) |r| + if (std.mem.eql(u8, r.name, name)) return r.usage; + } if (release.major < 13 or (release.major == 13 and release.minor < 2)) return null; for (records) |r| if (std.mem.eql(u8, r.name, name)) return r.usage; return null; diff --git a/zig/src/core/contract/registry.zig b/zig/src/core/contract/registry.zig index 7d83d71249..493c9d5a02 100644 --- a/zig/src/core/contract/registry.zig +++ b/zig/src/core/contract/registry.zig @@ -324,6 +324,7 @@ pub const cuda_sets = [_]CudaSet{ .{ .name = "router tf_nemo_topk", .usage = .{ .{ .regs = 94, .static = 0, .bound = 32 }, .{ .regs = 95, .static = 0, .bound = 32 }, .{ .regs = 95, .static = 0, .bound = 32 }, .{ .regs = 96, .static = 0, .bound = 32 }, .{ .regs = 96, .static = 0, .bound = 32 } } }, .{ .name = "flashnext_batch fn_qsa_select", .usage = .{ .{ .regs = 32, .static = 1168, .bound = 1024 }, .{ .regs = 50, .static = 1168, .bound = 1024 }, .{ .regs = 50, .static = 1168, .bound = 1024 }, .{ .regs = 44, .static = 1168, .bound = 1024 }, .{ .regs = 44, .static = 1168, .bound = 1024 } } }, } ++ @import("flashnext_tile_records.zig").cuda_133 ++ @import("glm_cuda_records.zig").cuda_133 ++ @import("nvcc_133_records.zig").cuda_133) }, + .{ .release = .{ .major = 13, .minor = 4 }, .usage = &.{.{ .name = "exl3_candidate tf_exl3_prompt_2", .usage = .{ .{ .regs = 80, .static = 32768, .bound = 256 }, .{ .regs = 80, .static = 32768, .bound = 256 }, .{ .regs = 80, .static = 32768, .bound = 256 }, .{ .regs = 77, .static = 32768, .bound = 256 }, .{ .regs = 77, .static = 32768, .bound = 256 } } }} }, }; pub fn newer(x: limits.Nvcc, y: limits.Nvcc) bool { return x.major > y.major or (x.major == y.major and x.minor > y.minor); @@ -578,7 +579,7 @@ const learning = @import("registry_learning.zig"); pub const gradient_pending = learning.gradient_pending; pub const qwen_learning_pending = learning.qwen_learning_pending; pub const glm_pending = @import("glm_sampling_pending.zig"); -pub const landed_cuda = cuda ++ @import("dory_cuda_records.zig").cuda ++ gradient_pending ++ qwen_learning_pending ++ paging.extra ++ @import("qwen_served_records.zig").records ++ cuda_pending ++ @import("qwen_decode_pending.zig").records ++ glm_pending.sampling ++ @import("glm_kv_pending.zig").records ++ @import("proposal_sampling_records.zig").records ++ @import("flashnext_pages_pending.zig").records ++ @import("glm_f32w_pending.zig").records ++ @import("penalty_sampling_records.zig").records; +pub const landed_cuda = cuda ++ @import("dory_cuda_records.zig").cuda ++ gradient_pending ++ qwen_learning_pending ++ paging.extra ++ @import("qwen_served_records.zig").records ++ cuda_pending ++ @import("qwen_decode_pending.zig").records ++ glm_pending.sampling ++ @import("glm_kv_pending.zig").records ++ @import("proposal_sampling_records.zig").records ++ @import("flashnext_pages_pending.zig").records ++ @import("flashnext_exl3_pending.zig").records ++ @import("glm_f32w_pending.zig").records ++ @import("penalty_sampling_records.zig").records; test { _ = learning; } diff --git a/zig/src/core/contract/registry_pending_test.zig b/zig/src/core/contract/registry_pending_test.zig index c708e15bbc..2453c70030 100644 --- a/zig/src/core/contract/registry_pending_test.zig +++ b/zig/src/core/contract/registry_pending_test.zig @@ -9,6 +9,23 @@ const cuda = registry.cuda; const checkCuda = registry.checkCuda; const targetDecision = registry.targetDecision; const TargetDecision = registry.TargetDecision; + +test { + _ = @import("flashnext_exl3_pending.zig"); +} +test "EXL3 prompt resource records select measured CUDA 13.4 usage without changing older compiler records" { + const record = for (cuda) |candidate| { + if (std.mem.eql(u8, candidate.launch.name, "exl3_candidate tf_exl3_prompt_2")) break candidate; + } else return error.MissingExl3PromptRecord; + const older = registry.usageFor(record, .{ .major = 13, .minor = 3 }); + const current = registry.usageFor(record, .{ .major = 13, .minor = 4 }); + for (current, 0..) |usage, index| { + try std.testing.expectEqual(@as(u32, if (index < 3) 80 else 77), usage.regs); + try std.testing.expectEqual(@as(u32, if (index < 3) 80 else 79), older[index].regs); + try std.testing.expectEqual(@as(u32, 32768), usage.static); + try std.testing.expectEqual(@as(u32, 256), usage.bound); + } +} test "CUDA lane resources distinguish recorded launches from missing receipts" { for (cuda_pending) |r| { try std.testing.expect(r.launch.threads > 0); diff --git a/zig/src/core/contract/registry_sets_test.zig b/zig/src/core/contract/registry_sets_test.zig index a494f045cb..d964985766 100644 --- a/zig/src/core/contract/registry_sets_test.zig +++ b/zig/src/core/contract/registry_sets_test.zig @@ -39,7 +39,8 @@ test "a compiler's records: its own set, else the newest below it, else 13.0's" .{ .{ .major = 13, .minor = 1 }, .{ .major = 13, .minor = 0 }, 91, 45 }, .{ .{ .major = 13, .minor = 2 }, .{ .major = 13, .minor = 2 }, 91, 43 }, .{ .{ .major = 13, .minor = 3 }, .{ .major = 13, .minor = 3 }, 96, 43 }, - .{ .{ .major = 14, .minor = 0 }, .{ .major = 13, .minor = 3 }, 96, 43 }, + .{ .{ .major = 13, .minor = 4 }, .{ .major = 13, .minor = 4 }, 96, 43 }, + .{ .{ .major = 14, .minor = 0 }, .{ .major = 13, .minor = 4 }, 96, 43 }, }; for (cases) |c| { try std.testing.expectEqual(c[1], setFor(c[0])); diff --git a/zig/src/cuda/exl3_dma.zig b/zig/src/cuda/exl3_dma.zig new file mode 100644 index 0000000000..f2b78ffb6c --- /dev/null +++ b/zig/src/cuda/exl3_dma.zig @@ -0,0 +1,38 @@ +const std = @import("std"); +const Stream = @import("stream.zig").Stream; + +pub fn directLines(stream: Stream, dst: u64, capacity: usize, offsets: [3]usize, src: [3][]const u8, pinned_bank: ?[]const u8) !usize { + const Span = struct { dst: u64, bytes: []const u8, banked: bool }; + var spans: [3]Span = undefined; + var count: usize = 0; + var moved: usize = 0; + _ = std.math.add(u64, dst, capacity) catch return error.Invalid; + const bank_start = if (pinned_bank) |bank| @intFromPtr(bank.ptr) else 0; + const bank_end = if (pinned_bank) |bank| std.math.add(usize, bank_start, bank.len) catch return error.Invalid else 0; + for (src, offsets) |bytes, offset| { + if (offset > capacity or bytes.len > capacity - offset) return error.Invalid; + const start = @intFromPtr(bytes.ptr); + const end = std.math.add(usize, start, bytes.len) catch return error.Invalid; + moved = std.math.add(usize, moved, bytes.len) catch return error.Invalid; + const address = dst + offset; + const banked = pinned_bank != null and start >= bank_start and end <= bank_end; + if (count > 0 and bytes.len > 0 and banked) { + const previous = &spans[count - 1]; + if (previous.banked and previous.bytes.len > 0 and + @intFromPtr(previous.bytes.ptr) + previous.bytes.len == start and + previous.dst + previous.bytes.len == address) + { + previous.bytes = pinned_bank.?[@intFromPtr(previous.bytes.ptr) - bank_start ..][0 .. previous.bytes.len + bytes.len]; + continue; + } + } + spans[count] = .{ .dst = address, .bytes = bytes, .banked = banked }; + count += 1; + } + for (spans[0..count]) |span| try stream.d.check(stream.d.api.cuMemcpyHtoDAsync_v2(span.dst, span.bytes.ptr, span.bytes.len, stream.handle), "cuMemcpyHtoDAsync"); + return moved; +} + +test { + _ = @import("exl3_dma_test.zig"); +} diff --git a/zig/src/cuda/exl3_dma_test.zig b/zig/src/cuda/exl3_dma_test.zig new file mode 100644 index 0000000000..8cbafd61cb --- /dev/null +++ b/zig/src/cuda/exl3_dma_test.zig @@ -0,0 +1,216 @@ +const std = @import("std"); +const dma = @import("exl3_dma.zig"); +const pool = @import("exl3_pool.zig"); +const abi = @import("abi.zig"); +const Driver = @import("driver.zig").Driver; +const Stream = @import("stream.zig").Stream; +const Event = @import("stream.zig").Event; + +const Fake = struct { + const Call = enum { query, wait, copy, table, event }; + const Copy = struct { dst: u64, src: usize, len: usize, stream: abi.Stream }; + const Word = struct { dst: u64, value: u32 }; + var calls: [32]Call = undefined; + var count: usize = 0; + var copies: [3]Copy = undefined; + var copied: usize = 0; + var fail_copy: usize = 0; + var words: [18]Word = undefined; + var written: usize = 0; + + fn note(call: Call) void { + calls[count] = call; + count += 1; + } + + fn copy(dst: u64, src: ?*const anyopaque, len: usize, stream: abi.Stream) callconv(.c) abi.Result { + note(.copy); + copies[copied] = .{ .dst = dst, .src = @intFromPtr(src), .len = len, .stream = stream }; + copied += 1; + return if (copied == fail_copy) abi.error_not_ready else abi.success; + } + + fn table(dst: u64, value: c_uint, len: usize, stream: abi.Stream) callconv(.c) abi.Result { + std.debug.assert(len == 1 and stream == @as(abi.Stream, @ptrFromInt(7))); + note(.table); + words[written] = .{ .dst = dst, .value = value }; + written += 1; + return abi.success; + } + + fn event(_: abi.Event, _: abi.Stream) callconv(.c) abi.Result { + note(.event); + return abi.success; + } + + fn query(_: abi.Event) callconv(.c) abi.Result { + note(.query); + return abi.error_not_ready; + } + + fn wait(_: abi.Stream, _: abi.Event, _: c_uint) callconv(.c) abi.Result { + note(.wait); + return abi.success; + } + + fn driver() Driver { + count = 0; + copied = 0; + fail_copy = 0; + written = 0; + var result: Driver = undefined; + result.api.cuMemcpyHtoDAsync_v2 = copy; + result.api.cuMemsetD32Async = table; + result.api.cuEventRecord = event; + result.api.cuEventQuery = query; + result.api.cuStreamWaitEvent = wait; + return result; + } +}; + +fn run(src: [3][]const u8, offsets: [3]usize, bank: ?[]const u8, lengths: []const usize) !void { + var driver = Fake.driver(); + const stream: Stream = .{ .d = &driver, .handle = @ptrFromInt(7) }; + const moved = try dma.directLines(stream, 4096, 32, offsets, src, bank); + try std.testing.expectEqual(src[0].len + src[1].len + src[2].len, moved); + try std.testing.expectEqual(lengths.len, Fake.copied); + var projection: usize = 0; + for (Fake.copies[0..Fake.copied], lengths) |copy, len| { + try std.testing.expectEqual(len, copy.len); + try std.testing.expectEqual(4096 + offsets[projection], copy.dst); + try std.testing.expectEqual(@intFromPtr(src[projection].ptr), copy.src); + try std.testing.expectEqual(stream.handle, copy.stream); + var consumed: usize = 0; + while (projection < 3) { + const bytes = src[projection]; + const uploaded: [*]const u8 = @ptrFromInt(copy.src); + try std.testing.expectEqualSlices(u8, bytes, uploaded[consumed..][0..bytes.len]); + consumed += bytes.len; + projection += 1; + if (consumed == len) break; + } + try std.testing.expectEqual(len, consumed); + } + try std.testing.expectEqual(@as(usize, 3), projection); +} + +test "DMA coalescing: exact unequal widths merge three projections or either pair" { + const bank = [_]u8{ 1, 2, 3, 4, 5, 6, 7, 8, 9, 10 }; + const src: [3][]const u8 = .{ bank[0..2], bank[2..5], bank[5..9] }; + try run(src, .{ 0, 2, 5 }, &bank, &.{9}); + try run(src, .{ 0, 2, 8 }, &bank, &.{ 5, 4 }); + try run(src, .{ 0, 4, 7 }, &bank, &.{ 2, 7 }); + try run(.{ bank[0..2], bank[2..5], bank[6..10] }, .{ 0, 2, 5 }, &bank, &.{ 5, 4 }); + try run(.{ bank[0..2], bank[3..6], bank[6..10] }, .{ 0, 2, 5 }, &bank, &.{ 2, 7 }); +} + +test "DMA coalescing: absent provenance and source or destination gaps retain exact fallback copies" { + const bank: [32]u8 = @splat(1); + const src: [3][]const u8 = .{ bank[0..2], bank[2..5], bank[5..9] }; + try run(src, .{ 0, 2, 5 }, null, &.{ 2, 3, 4 }); + try run(src, .{ 0, 4, 9 }, &bank, &.{ 2, 3, 4 }); + try run(.{ bank[0..2], bank[4..7], bank[9..13] }, .{ 0, 2, 5 }, &bank, &.{ 2, 3, 4 }); + try run(src, .{ 0, 2, 5 }, bank[10..20], &.{ 2, 3, 4 }); + try run(src, .{ 0, 2, 5 }, bank[0..4], &.{ 2, 3, 4 }); + try run(src, .{ 0, 2, 5 }, bank[2..9], &.{ 2, 7 }); + try run(src, .{ 0, 2, 5 }, bank[0..0], &.{ 2, 3, 4 }); + const foreign: [9]u8 = @splat(3); + try run(.{ foreign[0..2], foreign[2..5], foreign[5..9] }, .{ 0, 2, 5 }, &bank, &.{ 2, 3, 4 }); +} + +test "DMA coalescing: zero lengths preserve fallback and never bridge a gap" { + const bank: [9]u8 = @splat(2); + try run(.{ bank[0..0], bank[0..3], bank[3..9] }, .{ 0, 0, 3 }, &bank, &.{ 0, 9 }); + try run(.{ bank[0..2], bank[2..2], bank[2..9] }, .{ 0, 2, 2 }, &bank, &.{ 2, 0, 7 }); + try run(.{ bank[0..2], bank[2..2], bank[3..9] }, .{ 0, 2, 3 }, &bank, &.{ 2, 0, 6 }); + try run(.{ bank[0..0], bank[0..0], bank[0..0] }, .{ 0, 0, 0 }, &bank, &.{ 0, 0, 0 }); +} + +test "DMA coalescing: bounds and overflowing source bank or destination fail before copying" { + var driver = Fake.driver(); + const stream: Stream = .{ .d = &driver, .handle = null }; + const bank: [9]u8 = @splat(1); + const src: [3][]const u8 = .{ bank[0..2], bank[2..5], bank[5..9] }; + try std.testing.expectError(error.Invalid, dma.directLines(stream, 4096, 8, .{ 0, 2, 5 }, src, &bank)); + try std.testing.expectError(error.Invalid, dma.directLines(stream, 4096, 9, .{ 0, 10, 5 }, src, &bank)); + try std.testing.expectError(error.Invalid, dma.directLines(stream, 4096, 9, .{ 0, 2, std.math.maxInt(usize) }, src, &bank)); + try std.testing.expectError(error.Invalid, dma.directLines(stream, std.math.maxInt(u64) - 8, 9, .{ 0, 2, 5 }, src, &bank)); + const high: [*]const u8 = @ptrFromInt(std.math.maxInt(usize) - 1); + try std.testing.expectError(error.Invalid, dma.directLines(stream, 4096, 9, .{ 0, 2, 5 }, .{ src[0], src[1], high[0..4] }, null)); + try std.testing.expectError(error.Invalid, dma.directLines(stream, 4096, 9, .{ 0, 2, 5 }, src, high[0..4])); + try std.testing.expectEqual(@as(usize, 0), Fake.count); +} + +test "DMA coalescing: production three-bit sizes issue one exact transfer" { + const width = 614400; + const bank = try std.testing.allocator.alloc(u8, 3 * width); + defer std.testing.allocator.free(bank); + @memset(bank, 1); + var driver = Fake.driver(); + const stream: Stream = .{ .d = &driver, .handle = null }; + const layout = pool.Layout.init(.{ width, width, width }); + const src: [3][]const u8 = .{ bank[0..width], bank[width .. 2 * width], bank[2 * width ..] }; + try std.testing.expectEqual(@as(usize, 3 * width), try dma.directLines(stream, 4096, layout.line_bytes, layout.off, src, bank)); + try std.testing.expectEqual(@as(usize, 1), Fake.copied); + try std.testing.expectEqual(@as(usize, 3 * width), Fake.copies[0].len); +} + +test "DMA coalescing: pool preserves reader waits table writes event and expert byte statistics" { + const bank: [1536]u8 = @splat(1); + const src: [3][]const u8 = .{ bank[0..256], bank[256..768], bank[768..1536] }; + for ([_]?[]const u8{ null, &bank, bank[256..] }, [_]usize{ 3, 1, 2 }) |provenance, copies| { + for (0..copies + 1) |failure| { + var driver = Fake.driver(); + Fake.fail_copy = failure; + var epochs = [_]u64{1}; + var pending = [_]bool{false}; + var events = [_]Event{.{ .d = &driver, .handle = @ptrFromInt(2) }}; + var lines: pool.Lines = undefined; + lines.d = &driver; + lines.spec = .{ .experts = 2, .bytes = .{ 256, 512, 768 }, .lines = 1, .staging = 0 }; + lines.layout = pool.Layout.init(lines.spec.bytes); + lines.tlay = pool.Tables.init(2); + lines.mem = .{ .d = &driver, .ptr = 4096, .len = 1536 }; + lines.tables = .{ .d = &driver, .ptr = 8192, .len = lines.tlay.bytes }; + lines.staging = null; + lines.pinned_bank = provenance; + lines.copy = .{ .d = &driver, .handle = @ptrFromInt(7) }; + lines.line_ev = &events; + lines.line_pending = &pending; + lines.line_epoch = &epochs; + lines.rel_ev[1] = .{ .d = &driver, .handle = @ptrFromInt(3) }; + lines.done_epoch = 0; + lines.stats = .{}; + const result = lines.load(0, 1, src, .{ 3, 6, 9 }, 0, true, undefined); + try std.testing.expectEqualSlices(Fake.Call, &.{ .query, .wait }, Fake.calls[0..2]); + if (failure > 0) { + try std.testing.expectError(error.NotReady, result); + try std.testing.expectEqual(2 + failure, Fake.count); + try std.testing.expectEqual(failure, Fake.copied); + try std.testing.expect(!pending[0]); + try std.testing.expectEqual(@as(u64, 0), lines.stats.copies); + try std.testing.expectEqual(@as(u64, 0), lines.stats.bytes); + } else { + try result; + try std.testing.expectEqual(2 + copies + 18 + 1, Fake.count); + for (Fake.calls[2 + copies .. Fake.count - 1]) |call| try std.testing.expectEqual(Fake.Call.table, call); + try std.testing.expectEqual(Fake.Call.event, Fake.calls[Fake.count - 1]); + try std.testing.expect(pending[0]); + try std.testing.expectEqual(@as(u64, 1), lines.stats.copies); + try std.testing.expectEqual(@as(u64, bank.len), lines.stats.bytes); + try std.testing.expectEqual(@as(u64, 1), lines.stats.release_waits); + for (0..3) |which| { + const old = Fake.words[3 * which ..][0..3]; + const installed = Fake.words[9 + 3 * which ..][0..3]; + try std.testing.expectEqualDeep(Fake.Word{ .dst = lines.table(which), .value = 0 }, old[0]); + try std.testing.expectEqualDeep(Fake.Word{ .dst = lines.table(which) + 4, .value = 0 }, old[1]); + try std.testing.expectEqualDeep(Fake.Word{ .dst = lines.widths(which), .value = 0 }, old[2]); + try std.testing.expectEqualDeep(Fake.Word{ .dst = lines.table(which) + 8, .value = @intCast(4096 + lines.layout.off[which]) }, installed[0]); + try std.testing.expectEqualDeep(Fake.Word{ .dst = lines.table(which) + 12, .value = 0 }, installed[1]); + try std.testing.expectEqualDeep(Fake.Word{ .dst = lines.widths(which) + 4, .value = @intCast(3 * (which + 1)) }, installed[2]); + } + } + for (Fake.calls[2 .. 2 + Fake.copied]) |call| try std.testing.expectEqual(Fake.Call.copy, call); + } + } +} diff --git a/zig/src/cuda/exl3_pool.zig b/zig/src/cuda/exl3_pool.zig new file mode 100644 index 0000000000..4ac452ff2e --- /dev/null +++ b/zig/src/cuda/exl3_pool.zig @@ -0,0 +1,331 @@ +//! Device lines for one layer's EXL3 routed-expert offload on upstream's `exl3_experts.Experts`. + +const std = @import("std"); +const abi = @import("abi.zig"); +const memory = @import("memory.zig"); +const stream_ = @import("stream.zig"); +pub const dma = @import("exl3_dma.zig"); + +const Driver = @import("driver.zig").Driver; +const Error = @import("driver.zig").Error; +const DeviceBuffer = memory.DeviceBuffer; +const HostBuffer = memory.HostBuffer; +const Stream = stream_.Stream; +const Event = stream_.Event; + +fn align256(n: usize) usize { + return (n + 255) / 256 * 256; +} + +pub const Spec = struct { + /// Table length E (513: 512 routed experts and the shared expert as id 512). + experts: usize, + /// Largest trellis bytes of any expert, per projection (gate, up, down): a line fits every expert of the layer. + bytes: [3]usize, + lines: usize, + /// Pinned staging lines (a ring); 0 when every copy is `direct`. + staging: usize = 8, +}; + +/// Byte offsets of gate, up and down inside a line, each 256-aligned. +pub const Layout = struct { + off: [3]usize, + line_bytes: usize, + + pub fn init(bytes: [3]usize) Layout { + var l: Layout = .{ .off = undefined, .line_bytes = 0 }; + for (bytes, &l.off) |b, *o| { + o.* = l.line_bytes; + l.line_bytes += align256(b); + } + return l; + } +}; + +/// Where the six tables sit inside one allocation: three int64 [E] pointer tables, then three int32 [E] width tables. +pub const Tables = struct { + ptr: [3]usize, + k2: [3]usize, + bytes: usize, + + pub fn init(experts: usize) Tables { + var t: Tables = .{ .ptr = undefined, .k2 = undefined, .bytes = 0 }; + for (&t.ptr) |*o| { + o.* = t.bytes; + t.bytes += align256(experts * 8); + } + for (&t.k2) |*o| { + o.* = t.bytes; + t.bytes += align256(experts * 4); + } + return t; + } +}; + +pub const Stats = struct { + copies: u64 = 0, + bytes: u64 = 0, + staged_bytes: u64 = 0, // bytes the host memcpy'd into pinned staging + stage_ns: u64 = 0, // host time in those memcpys + staging_waits: u64 = 0, // a staging slot was still being copied from (the host blocked) + release_waits: u64 = 0, // a copy waited on the device for a pass that read its line + gate_waits: u64 = 0, // a pass waited on the device for a line copy +}; + +const release_ring = 64; + +pub const Lines = struct { + d: *const Driver, + spec: Spec, + layout: Layout, + tlay: Tables, + mem: DeviceBuffer, + tables: DeviceBuffer, + shared: ?DeviceBuffer = null, + pinned_bank: ?[]const u8 = null, + staging: ?HostBuffer, + stage_ev: []Event, + stage_busy: []bool, + stage_next: usize = 0, + copy: Stream, + line_ev: []Event, + line_pending: []bool, + /// Epoch of the last compute pass that read the line (0: never read). + line_epoch: []u64, + rel_ev: [release_ring]Event, + epoch: u64 = 0, // last epoch released + done_epoch: u64 = 0, // highest epoch known finished + stats: Stats = .{}, + gpa: std.mem.Allocator, + + pub fn init(gpa: std.mem.Allocator, d: *const Driver, spec: Spec) !Lines { + if (spec.lines == 0 or spec.experts == 0) return error.Invalid; + const layout = Layout.init(spec.bytes); + const tlay = Tables.init(spec.experts); + var mem = try DeviceBuffer.alloc(d, spec.lines * layout.line_bytes); + errdefer mem.free(); + var tables = try DeviceBuffer.alloc(d, tlay.bytes); + errdefer tables.free(); + try tables.fill8(0, null); // every expert absent; synchronous, before any stream uses the tables + var staging: ?HostBuffer = if (spec.staging > 0) try HostBuffer.alloc(d, spec.staging * layout.line_bytes) else null; + errdefer if (staging) |*h| h.free(); + var copy = try Stream.init(d, true); + errdefer copy.deinit(); + const stage_ev = try makeEvents(gpa, d, spec.staging); + errdefer freeEvents(gpa, stage_ev); + const stage_busy = try gpa.alloc(bool, spec.staging); + errdefer gpa.free(stage_busy); + @memset(stage_busy, false); + const line_ev = try makeEvents(gpa, d, spec.lines); + errdefer freeEvents(gpa, line_ev); + const line_pending = try gpa.alloc(bool, spec.lines); + errdefer gpa.free(line_pending); + @memset(line_pending, false); + const line_epoch = try gpa.alloc(u64, spec.lines); + errdefer gpa.free(line_epoch); + @memset(line_epoch, 0); + var rel: [release_ring]Event = undefined; + var made: usize = 0; + errdefer for (rel[0..made]) |*e| e.deinit(); + for (&rel) |*e| { + e.* = try Event.init(d, false); + made += 1; + } + return .{ + .d = d, + .spec = spec, + .layout = layout, + .tlay = tlay, + .mem = mem, + .tables = tables, + .staging = staging, + .stage_ev = stage_ev, + .stage_busy = stage_busy, + .copy = copy, + .line_ev = line_ev, + .line_pending = line_pending, + .line_epoch = line_epoch, + .rel_ev = rel, + .gpa = gpa, + }; + } + + fn makeEvents(gpa: std.mem.Allocator, d: *const Driver, n: usize) ![]Event { + const evs = try gpa.alloc(Event, n); + var made: usize = 0; + errdefer { + for (evs[0..made]) |*e| e.deinit(); + gpa.free(evs); + } + for (evs) |*e| { + e.* = try Event.init(d, false); + made += 1; + } + return evs; + } + + fn freeEvents(gpa: std.mem.Allocator, evs: []Event) void { + for (evs) |*e| e.deinit(); + gpa.free(evs); + } + + pub fn deinit(l: *Lines) void { + l.copy.synchronize() catch {}; + for (&l.rel_ev) |*e| e.deinit(); + freeEvents(l.gpa, l.line_ev); + freeEvents(l.gpa, l.stage_ev); + l.gpa.free(l.stage_busy); + l.gpa.free(l.line_pending); + l.gpa.free(l.line_epoch); + l.copy.deinit(); + if (l.staging) |*h| h.free(); + if (l.shared) |*b| b.free(); + l.tables.free(); + l.mem.free(); + l.* = undefined; + } + + /// Device bytes: the lines, the tables and the resident shared expert. + pub fn deviceBytes(l: Lines) usize { + return l.mem.len + l.tables.len + if (l.shared) |b| b.len else 0; + } + + pub fn table(l: Lines, which: usize) u64 { + return l.tables.ptr + l.tlay.ptr[which]; + } + + pub fn widths(l: Lines, which: usize) u64 { + return l.tables.ptr + l.tlay.k2[which]; + } + + /// Device address of projection `which` (0 gate, 1 up, 2 down) of line `line`. + pub fn at(l: Lines, line: usize, which: usize) u64 { + return l.mem.ptr + line * l.layout.line_bytes + l.layout.off[which]; + } + + /// Uploads the never-offloaded expert `id` (the shared expert) once into its own region (it may be wider than a + pub fn setResident(l: *Lines, id: usize, src: [3][]const u8, k2: [3]i32) !void { + if (id >= l.spec.experts or l.shared != null) return error.Invalid; + const lay = Layout.init(.{ src[0].len, src[1].len, src[2].len }); + var b = try DeviceBuffer.alloc(l.d, lay.line_bytes); + errdefer b.free(); + for (src, lay.off) |s, o| try b.upload(o, s); + l.shared = b; + for (0..3) |w| try l.writeEntry(id, b.ptr + lay.off[w], k2[w], w); + try l.copy.synchronize(); + } + + fn memset32(l: *Lines, addr: u64, value: u32, st: Stream) !void { + try l.d.check(l.d.api.cuMemsetD32Async(addr, value, 1, st.handle), "cuMemsetD32Async"); + } + + /// One expert's projection `which`: its pointer (two 32-bit words) and width, on the copy stream. + fn writeEntry(l: *Lines, id: usize, ptr: u64, k2: i32, which: usize) !void { + const st = l.copy; + const p = l.table(which) + id * 8; + try l.memset32(p, @truncate(ptr), st); + try l.memset32(p + 4, @truncate(ptr >> 32), st); + try l.memset32(l.widths(which) + id * 4, @bitCast(k2), st); + } + + /// The copy stream waits until the last compute pass that read `line` has finished (no wait if known done). + fn afterReaders(l: *Lines, line: usize) !void { + const e = l.line_epoch[line]; + if (e == 0 or e <= l.done_epoch) return; + const slot = e % release_ring; + const ev = l.rel_ev[slot]; + if (try ev.done()) { + // the slot holds epoch e or a newer one; either way e has finished + l.done_epoch = @max(l.done_epoch, e); + return; + } + try l.copy.wait(ev); // a newer epoch in the slot is a conservative wait + l.stats.release_waits += 1; + } + + /// Makes `id` resident in `line`: the copy waits for the line's readers, then H2D (staged or direct), then `id`'s + pub fn load(l: *Lines, line: usize, id: usize, src: [3][]const u8, k2: [3]i32, evicted: ?usize, direct: bool, io: std.Io) !void { + if (line >= l.spec.lines or id >= l.spec.experts) return error.Invalid; + for (src, l.spec.bytes) |s, cap| if (s.len > cap) return error.Invalid; + try l.afterReaders(line); + const dst = l.mem.ptr + line * l.layout.line_bytes; + var moved: usize = 0; + if (direct or l.staging == null) { + moved = try dma.directLines(l.copy, dst, l.layout.line_bytes, l.layout.off, src, l.pinned_bank); + } else { + const slot = l.stage_next; + l.stage_next = (l.stage_next + 1) % l.spec.staging; + if (l.stage_busy[slot]) { + if (!try l.stage_ev[slot].done()) l.stats.staging_waits += 1; + try l.stage_ev[slot].synchronize(); + } + const buf = l.staging.?.bytes[slot * l.layout.line_bytes ..][0..l.layout.line_bytes]; + const t0 = std.Io.Clock.awake.now(io).toNanoseconds(); + for (src, l.layout.off) |s, o| @memcpy(buf[o..][0..s.len], s); + l.stats.stage_ns += @intCast(std.Io.Clock.awake.now(io).toNanoseconds() - t0); + // one copy for the line: up to the end of down's bytes (alignment padding rides along) + const n = l.layout.off[2] + src[2].len; + try l.d.check(l.d.api.cuMemcpyHtoDAsync_v2(dst, buf.ptr, n, l.copy.handle), "cuMemcpyHtoDAsync"); + try l.stage_ev[slot].record(l.copy); + l.stage_busy[slot] = true; + l.stats.staged_bytes += src[0].len + src[1].len + src[2].len; + moved = n; + } + if (evicted) |old| for (0..3) |w| try l.writeEntry(old, 0, 0, w); + for (0..3) |w| try l.writeEntry(id, dst + l.layout.off[w], k2[w], w); + try l.line_ev[line].record(l.copy); + l.line_pending[line] = true; + l.stats.copies += 1; + l.stats.bytes += moved; + } + + /// True once the line's last copy finished (or it never had one); clears the pending flag. + pub fn landed(l: *Lines, line: usize) !bool { + if (!l.line_pending[line]) return true; + if (!try l.line_ev[line].done()) return false; + l.line_pending[line] = false; + return true; + } + + /// `compute` waits on the device for the line's copy, if one may still be running. The host does not block. + pub fn gate(l: *Lines, compute: Stream, line: usize) !void { + if (try l.landed(line)) return; + try compute.wait(l.line_ev[line]); + l.stats.gate_waits += 1; + } + + /// After a compute pass that read `lines`: a release event on `compute`; a later copy into any of those lines + pub fn release(l: *Lines, compute: Stream, lines: []const u32) !void { + l.epoch += 1; + const slot = l.epoch % release_ring; + try l.rel_ev[slot].record(compute); + for (lines) |i| l.line_epoch[i] = l.epoch; + } + + /// Blocks until every queued copy and table write finished. + pub fn drain(l: *Lines) !void { + try l.copy.synchronize(); + for (l.line_pending) |*p| p.* = false; + @memset(l.stage_busy, false); + } +}; + +test { + _ = dma; +} + +test "a line packs gate, up and down 256-aligned; the tables sit in one allocation" { + // Flash Next layer 0: 3-bit experts, K=2560 N=640 -> 160 x 40 tiles of 48 words + const tb = 160 * 40 * 48 * 2; + const l = Layout.init(.{ tb, tb, tb }); + try std.testing.expectEqual([3]usize{ 0, tb, 2 * tb }, l.off); + try std.testing.expectEqual(@as(usize, 3 * tb), l.line_bytes); + const odd = Layout.init(.{ 100, 300, 5 }); + try std.testing.expectEqual([3]usize{ 0, 256, 768 }, odd.off); + try std.testing.expectEqual(@as(usize, 1024), odd.line_bytes); + const t = Tables.init(513); + try std.testing.expectEqual([3]usize{ 0, 4352, 8704 }, t.ptr); + try std.testing.expectEqual(@as(usize, 13056), t.k2[0]); + try std.testing.expectEqual(@as(usize, 13056 + 3 * 2304), t.bytes); + for (t.ptr ++ t.k2) |o| try std.testing.expectEqual(@as(usize, 0), o % 256); +} diff --git a/zig/src/cuda/kernels.zig b/zig/src/cuda/kernels.zig index 33a3bbfe55..0594bdb4ce 100644 --- a/zig/src/cuda/kernels.zig +++ b/zig/src/cuda/kernels.zig @@ -18,6 +18,7 @@ pub const qwen4_gdn: []const u8 = if (available) &Blob("fatbin_qwen4_gdn").bytes pub const qwen4_projection: []const u8 = if (available) &Blob("fatbin_qwen4_projection").bytes else &.{}; pub const gdn: []const u8 = if (available) &Blob("fatbin_gdn").bytes else &.{}; pub const flashnext_gdn: []const u8 = if (available) &Blob("fatbin_flashnext_gdn").bytes else &.{}; +pub const flashnext_exl3: []const u8 = if (available) &Blob("fatbin_flashnext_exl3").bytes else &.{}; pub const flashnext_batch: []const u8 = if (available) &Blob("fatbin_flashnext_batch").bytes else &.{}; pub const flashnext_mtp: []const u8 = if (available) &Blob("fatbin_flashnext_mtp").bytes else &.{}; pub const probe: []const u8 = if (available) &Blob("fatbin_probe").bytes else &.{}; diff --git a/zig/src/cuda/root.zig b/zig/src/cuda/root.zig index 5489b019b8..98da1f7efa 100644 --- a/zig/src/cuda/root.zig +++ b/zig/src/cuda/root.zig @@ -69,6 +69,7 @@ pub const exl3_experts = @import("exl3_experts.zig"); pub const exl3_upload = @import("exl3_upload.zig"); pub const exl3_candidate = @import("exl3_candidate.zig"); pub const exl3_grouped = @import("exl3_grouped.zig"); +pub const exl3_pool = @import("exl3_pool.zig"); pub const arena = @import("arena.zig"); pub const pack = @import("pack.zig"); @@ -118,6 +119,7 @@ test { _ = exl3; _ = exl3_experts; _ = exl3_upload; + _ = exl3_pool; _ = @import("modelopt.zig"); _ = @import("ck.zig"); _ = @import("tp_watch.zig"); diff --git a/zig/src/families/flashnext/cuda_batch.zig b/zig/src/families/flashnext/cuda_batch.zig index 6066b1ea63..8ebdea1f63 100644 --- a/zig/src/families/flashnext/cuda_batch.zig +++ b/zig/src/families/flashnext/cuda_batch.zig @@ -8,6 +8,7 @@ const ple = @import("cuda_ple.zig"); const wt = @import("cuda_prompt_wt.zig"); const load = @import("cuda_prompt_load.zig"); const parts = @import("cuda_batch_state.zig"); +const exl3_pass = @import("exl3/model/cuda_exl3_pass.zig"); const batch_lanes = @import("cuda_batch_lanes.zig"); const tiles = @import("cuda_tiles.zig"); const tk = @import("cuda_tile_kernels.zig"); @@ -123,6 +124,10 @@ pub fn Pass(comptime Tri: type) type { /// qmm.cu's arithmetic over m rows into out (bf16) or part (fp32 K slices); returns the slices (1 when folded). pub fn lane(self: *Self, x: u64, xs: u64, w: Q4, out: u64, part: u64, m: usize, sk: usize) !usize { + if (w.exl3) |e| { + try exl3_pass.lane(self, x, e, out, part, m, sk); + return sk; // the first slice written, the other sk - 1 zero-filled (cuda_exl3_pass.lane) + } if (try zs.lane(self.k, self.stream, self.s, self.zero_slices, x, w, part, m, sk)) |one| return one; return tiles.matmul(self.k, self.stream, x, xs, w, out, part, self.s.native_part, m, sk); } @@ -228,27 +233,7 @@ pub fn Pass(comptime Tri: type) type { } pub fn pleBlock(self: *Self, lw: *const Layer, ids: []const u32) !void { - const s = self.s; - const n = self.rows; - // the rows' n-gram embeddings (host gathers from the mapped table) are already uploaded by `stage` - _ = ids; - _ = try self.lane(s.ple_emb, s.ple_xs, lw.key, s.ple_keys, s.dummy, n, 1); - const vsk = try self.lane(s.ple_emb, s.ple_xs, lw.val, s.dummy, s.ple_part, n, 4); - try self.sumSlices(s.ple_part, s.ple_vals, true, n * dims, vsk); - try self.tri.pleGate(s.ple_keys, s.ple_vals, s.h, lw.nk.ptr, lw.nq.ptr, s.ple_gated, s.ple_pss, eps, n, dims, streams); - if (self.lanes.len == 0) { - try self.tri.pleConv(s.ple_gated, s.ple_pss, lw.nc.ptr, self.seq.tail.ptr, lw.ple_cw.ptr, s.h, s.h, s.ple_nrow, eps, n, dims, streams, 4, 3); - try self.shiftWindow(self.seq.tail.ptr, s.ple_nrow, wide, n, ple_tail_n, wide); - return; - } - // shared rounds: each lane's rows convolved over its own stream's tail, the launch its lone pass makes - var row0: usize = 0; - for (self.lanes) |l| { - const g = row0 * wide * 2; - try self.tri.pleConv(s.ple_gated + g, s.ple_pss + row0 * streams * 4, lw.nc.ptr, l.seq.tail.ptr, lw.ple_cw.ptr, s.h + g, s.h + g, s.ple_nrow + g, eps, l.rows, dims, streams, 4, 3); - try self.shiftWindow(l.seq.tail.ptr, s.ple_nrow + g, wide, l.rows, ple_tail_n, wide); - row0 += l.rows; - } + return @import("cuda_batch_ple.zig").pleBlock(self, lw, ids); } pub fn gdnBlock(self: *Self, lw: *const Layer, layer: usize, li: usize, keep: LayerKeep) !void { @@ -427,6 +412,14 @@ pub fn Pass(comptime Tri: type) type { const s = self.s; const n = self.rows; const pairs = n * slots; + if (lw.exl3) |x| { + if (x.needsPool()) return error.Exl3PoolNotAttached; + if (keep.moe_x != 0) try self.copy(keep.moe_x, s.mixed, n * dims * 2); + try exl3_pass.moe(self, lw, x); + if (keep.moe_ids != 0) try self.copy(keep.moe_ids, s.picks, pairs * 4); + if (keep.moe_wts != 0) try self.copy(keep.moe_wts, s.wts, pairs * 4); + return; + } try self.copy(keep.moe_x, s.mixed, n * dims * 2); try self.tri.routerMany(s.mixed, lw.router.ptr, s.router_l, dims, n); try self.tri.topkMany(s.router_l, s.picks, s.wts, n); diff --git a/zig/src/families/flashnext/cuda_batch_gather.zig b/zig/src/families/flashnext/cuda_batch_gather.zig new file mode 100644 index 0000000000..b2a1ab8773 --- /dev/null +++ b/zig/src/families/flashnext/cuda_batch_gather.zig @@ -0,0 +1,40 @@ +//! The batched pass's n-gram gather: the rows' embeddings on a host thread while the GPU runs layer 0. +const std = @import("std"); +const ple = @import("cuda_ple.zig"); +const wt = @import("cuda_prompt_wt.zig"); +const Lane = @import("cuda_batch_state.zig").Lane; + +const dims = wt.dims; + +/// The rows' n-gram embeddings, gathered on a host thread while the GPU embeds and runs layer 0. +pub const Gather = struct { + ngram: *const ple.Ngram, + history: []i64, + ids: []const u32, + emb: []u16, + xs: []f32, + failed: ?anyerror = null, + ns: i96 = 0, + io: std.Io, + /// A round's rows, each gathered after its own stream's history (empty: every row after `history`). + lanes: []const Lane = &.{}, + + pub fn main(g: *Gather) void { + const t0 = std.Io.Clock.awake.now(g.io).toNanoseconds(); + defer g.ns = std.Io.Clock.awake.now(g.io).toNanoseconds() - t0; + if (g.lanes.len == 0) { + g.ngram.rows(g.history, g.ids, g.emb, g.xs) catch |err| { + g.failed = err; + }; + return; + } + var row0: usize = 0; + for (g.lanes) |l| { + g.ngram.rows(l.seq.history[0..], g.ids[row0..][0..l.rows], g.emb[row0 * dims ..][0 .. l.rows * dims], g.xs[row0 * (dims / 32) ..][0 .. l.rows * (dims / 32)]) catch |err| { + g.failed = err; + return; + }; + row0 += l.rows; + } + } +}; diff --git a/zig/src/families/flashnext/cuda_batch_ple.zig b/zig/src/families/flashnext/cuda_batch_ple.zig new file mode 100644 index 0000000000..4ff2ffc346 --- /dev/null +++ b/zig/src/families/flashnext/cuda_batch_ple.zig @@ -0,0 +1,32 @@ +//! The shared batched pass's per-layer embedding block: the n-gram keys and values, the gate and the convolution. +const wt = @import("cuda_prompt_wt.zig"); + +const dims = wt.dims; +const streams = wt.streams; +const wide = wt.wide; +const eps = wt.eps; +const ple_tail_n = wt.ple_tail_n; + +pub fn pleBlock(self: anytype, lw: anytype, ids: []const u32) !void { + const s = self.s; + const n = self.rows; + // the rows' n-gram embeddings (host gathers from the mapped table) are already uploaded by `stage` + _ = ids; + _ = try self.lane(s.ple_emb, s.ple_xs, lw.key, s.ple_keys, s.dummy, n, 1); + const vsk = try self.lane(s.ple_emb, s.ple_xs, lw.val, s.dummy, s.ple_part, n, 4); + try self.sumSlices(s.ple_part, s.ple_vals, true, n * dims, vsk); + try self.tri.pleGate(s.ple_keys, s.ple_vals, s.h, lw.nk.ptr, lw.nq.ptr, s.ple_gated, s.ple_pss, eps, n, dims, streams); + if (self.lanes.len == 0) { + try self.tri.pleConv(s.ple_gated, s.ple_pss, lw.nc.ptr, self.seq.tail.ptr, lw.ple_cw.ptr, s.h, s.h, s.ple_nrow, eps, n, dims, streams, 4, 3); + try self.shiftWindow(self.seq.tail.ptr, s.ple_nrow, wide, n, ple_tail_n, wide); + return; + } + // shared rounds: each lane's rows convolved over its own stream's tail, the launch its lone pass makes + var row0: usize = 0; + for (self.lanes) |l| { + const g = row0 * wide * 2; + try self.tri.pleConv(s.ple_gated + g, s.ple_pss + row0 * streams * 4, lw.nc.ptr, l.seq.tail.ptr, lw.ple_cw.ptr, s.h + g, s.h + g, s.ple_nrow + g, eps, l.rows, dims, streams, 4, 3); + try self.shiftWindow(l.seq.tail.ptr, s.ple_nrow + g, wide, l.rows, ple_tail_n, wide); + row0 += l.rows; + } +} diff --git a/zig/src/families/flashnext/cuda_batch_state.zig b/zig/src/families/flashnext/cuda_batch_state.zig index 98faec84be..14035aa5d9 100644 --- a/zig/src/families/flashnext/cuda_batch_state.zig +++ b/zig/src/families/flashnext/cuda_batch_state.zig @@ -192,8 +192,11 @@ pub fn Staging(comptime Q: type) type { } /// Device scratch for chunks of up to `cap` rows. +pub const ScratchFormat = enum { classic, exl3 }; + pub const Scratch = struct { cap: usize, + format: ScratchFormat = .classic, bufs: std.ArrayList(cuda.DeviceBuffer) = .empty, ids: u64 = 0, pos: u64 = 0, @@ -312,110 +315,127 @@ pub const Scratch = struct { return @alignCast(std.mem.bytesAsSlice(f32, s.host.?.bytes[hostAt(s.cap)[3]..][0 .. n * (dims / 32) * 4])); } - /// Bytes `alloc(cap)` takes on the device. - pub fn bytes(cap: usize) usize { + /// The buffers in the order `allocationSizes` lists their sizes. + const order = [_][]const u8{ + "ids", "pos", "h", "pss", "norm", "rinv", "dn", "hc_part", "hc_up", "native_part", "act", + "xs_act", "inj", "mixed", "xsm", "proj", "part", "gout", "gxs", "ys", "gq", "gk", + "gv", "ggb", "q", "iq", "po", "pm", "pl", "qids", "qnk", "qspr", "ao", + "gated", "gxs_attn", "router_l", "picks", "wts", "members", "items", "counts", "rank", "hist", "shared_gate", + "shared_up", "shared_act", "shared_out", "moe_act", "y", "branch", "moe_f32", "pocket_act", "ple_emb", "ple_xs", "ple_keys", + "ple_vals", "ple_part", "ple_gated", "ple_pss", "ple_nrow", "logits", "head_local", "inv", "dummy", + }; + + pub fn allocationSizes(cap: usize, format: ScratchFormat) [64]?usize { const pairs = cap * slots; const nblk = cdiv(pairs, plan_small); const items = cdiv(pairs, tile) + experts_n + 1; - var n: usize = 0; - n += cap * 4 + (cap + 1) * 16; - n += cap * streams * 4 + cap * wide * 2 * 2 + cap * pss_n * 4 + cap * (low + 4) * 2 + hc_sk * cap * (low + 4) * 4; - n += cap * wide * 2 + 8 * cap * dims * 4; - n += cap * low * 2 + cap * (low / 32) * 4 + cap * streams * 2 + cap * dims * 2 + cap * (dims / 32) * 4; - n += cap * proj_row * 2 + out_sk * cap * dims * 4 + cap * gdn.value_dim * (2 + 4 + 4) + cap * (gdn.value_dim / 32) * 4; - n += cap * (2 * gdn.nk * gdn.dk + 2 * gdn.nv) * 4; - n += cap * q_heads * head_dim * 2 + cap * index_heads * index_dim * 2; - n += att_rows * nch * q_heads * (head_dim + 2) * 4 + cap * out_k * 2 * 2 + cap * (out_k / 32) * 4; - n += fused_rows * (idw + 2) * 4; - n += cap * (experts_n + 1) * 4 + pairs * 4 * 2 + pairs * 4 + items * 3 * 4 + 8 + pairs * 4 + nblk * (experts_n + 1) * 4; - n += cap * moe_w * 6 + cap * dims * 4; - n += pairs * moe_w * 2 + pairs * dims * 4 + cap * dims * 2; - n += cap * dims * 2 + cap * (dims / 32) * 4 + cap * wide * 2 + cap * dims * 2 + 4 * cap * dims * 4; - n += cap * wide * 2 * 2 + cap * streams * 4 + vocab * 2 + cap * vocab + half * 4 + 4096; - n += cap * dims * 4 + cap * pocket_stride * 4; - return n; + return .{ + cap * 4, + (cap + 1) * 16, + cap * wide * 2, + cap * pss_n * 4, + cap * wide * 2, + cap * streams * 4, + cap * (low + 4) * 2, + hc_sk * cap * (low + 4) * 4, + cap * wide * 2, + 8 * cap * dims * 4, + cap * low * 2, + cap * (low / 32) * 4, + cap * streams * 2, + cap * dims * 2, + cap * (dims / 32) * 4, + cap * proj_row * 2, + out_sk * cap * dims * 4, + cap * gdn.value_dim * 2, + cap * (gdn.value_dim / 32) * 4, + cap * gdn.value_dim * 4, + cap * gdn.nk * gdn.dk * 4, + cap * gdn.nk * gdn.dk * 4, + cap * gdn.value_dim * 4, + cap * gdn.nv * 2 * 4, + cap * q_heads * head_dim * 2, + cap * index_heads * index_dim * 2, + att_rows * nch * q_heads * head_dim * 4, + att_rows * nch * q_heads * 4, + att_rows * nch * q_heads * 4, + fused_rows * idw * 4, + fused_rows * 4, + fused_rows * 4, + cap * out_k * 2, + cap * out_k * 2, + cap * (out_k / 32) * 4, + cap * (experts_n + 1) * 4, + pairs * 4, + pairs * 4, + if (format == .classic) pairs * 4 else null, + if (format == .classic) items * 3 * 4 else null, + if (format == .classic) 8 else null, + if (format == .classic) pairs * 4 else null, + if (format == .classic) nblk * (experts_n + 1) * 4 else null, + if (format == .classic) cap * moe_w * 2 else null, + if (format == .classic) cap * moe_w * 2 else null, + if (format == .classic) cap * moe_w * 2 else null, + if (format == .classic) cap * dims * 4 else null, + if (format == .classic) pairs * moe_w * 2 else null, + pairs * dims * 4, + cap * dims * 2, + cap * dims * 4, + cap * pocket_stride * 4, + cap * dims * 2, + cap * (dims / 32) * 4, + cap * wide * 2, + cap * dims * 2, + 4 * cap * dims * 4, + cap * wide * 2, + cap * streams * 4, + cap * wide * 2, + vocab * 2, + if (format == .classic) cap * vocab else null, + half * 4, + 4096, + }; + } + + pub fn bytes(cap: usize) usize { + return bytesForFormat(cap, .classic); + } + + pub fn bytesForFormat(cap: usize, format: ScratchFormat) usize { + var total: usize = 0; + for (allocationSizes(cap, format)) |optional| if (optional) |len| { + total += @max(len, 256); + }; + return total; } pub fn alloc(gpa: std.mem.Allocator, d: *const cuda.Driver, stream: cuda.Stream, cap: usize) !Scratch { - var s: Scratch = .{ .cap = cap }; + return allocForFormat(gpa, d, stream, cap, .classic); + } + + pub fn allocForFormat(gpa: std.mem.Allocator, d: *const cuda.Driver, stream: cuda.Stream, cap: usize, format: ScratchFormat) !Scratch { + var s: Scratch = .{ .cap = cap, .format = format }; errdefer { cuda.completion.drain(d, stream, "Flash Next scratch work did not drain before rollback"); s.free(gpa); } - const pairs = cap * slots; - const nblk = cdiv(pairs, plan_small); - const items = cdiv(pairs, tile) + experts_n + 1; - s.ids = try s.add(gpa, d, cap * 4); - // a position every 16 bytes: one a block of att_rows rows, or one a row of a decode round - s.pos = try s.add(gpa, d, (cap + 1) * 16); - s.h = try s.add(gpa, d, cap * wide * 2); - s.pss = try s.add(gpa, d, cap * pss_n * 4); - s.norm = try s.add(gpa, d, cap * wide * 2); - s.rinv = try s.add(gpa, d, cap * streams * 4); - s.dn = try s.add(gpa, d, cap * (low + 4) * 2); - s.hc_part = try s.add(gpa, d, hc_sk * cap * (low + 4) * 4); - s.hc_up = try s.add(gpa, d, cap * wide * 2); - s.native_part = try s.add(gpa, d, 8 * cap * dims * 4); - s.act = try s.add(gpa, d, cap * low * 2); - s.xs_act = try s.add(gpa, d, cap * (low / 32) * 4); - s.inj = try s.add(gpa, d, cap * streams * 2); - s.mixed = try s.add(gpa, d, cap * dims * 2); - s.xsm = try s.add(gpa, d, cap * (dims / 32) * 4); - s.proj = try s.add(gpa, d, cap * proj_row * 2); - s.part = try s.add(gpa, d, out_sk * cap * dims * 4); - try s.bufs.items[s.bufs.items.len - 1].fill8(0, stream.handle); - s.part_zero = true; - s.gout = try s.add(gpa, d, cap * gdn.value_dim * 2); - s.gxs = try s.add(gpa, d, cap * (gdn.value_dim / 32) * 4); - s.ys = try s.add(gpa, d, cap * gdn.value_dim * 4); - s.gq = try s.add(gpa, d, cap * gdn.nk * gdn.dk * 4); - s.gk = try s.add(gpa, d, cap * gdn.nk * gdn.dk * 4); - s.gv = try s.add(gpa, d, cap * gdn.value_dim * 4); - s.ggb = try s.add(gpa, d, cap * gdn.nv * 2 * 4); - s.q = try s.add(gpa, d, cap * q_heads * head_dim * 2); - s.iq = try s.add(gpa, d, cap * index_heads * index_dim * 2); - s.po = try s.add(gpa, d, att_rows * nch * q_heads * head_dim * 4); - s.pm = try s.add(gpa, d, att_rows * nch * q_heads * 4); - s.pl = try s.add(gpa, d, att_rows * nch * q_heads * 4); - s.qids = try s.add(gpa, d, fused_rows * idw * 4); - s.qnk = try s.add(gpa, d, fused_rows * 4); - s.qspr = try s.add(gpa, d, fused_rows * 4); - s.ao = try s.add(gpa, d, cap * out_k * 2); - s.gated = try s.add(gpa, d, cap * out_k * 2); - s.gxs_attn = try s.add(gpa, d, cap * (out_k / 32) * 4); - s.router_l = try s.add(gpa, d, cap * (experts_n + 1) * 4); - s.picks = try s.add(gpa, d, pairs * 4); - s.wts = try s.add(gpa, d, pairs * 4); - s.members = try s.add(gpa, d, pairs * 4); - s.items = try s.add(gpa, d, items * 3 * 4); - s.counts = try s.add(gpa, d, 8); - s.rank = try s.add(gpa, d, pairs * 4); - s.hist = try s.add(gpa, d, nblk * (experts_n + 1) * 4); - s.shared_gate = try s.add(gpa, d, cap * moe_w * 2); - s.shared_up = try s.add(gpa, d, cap * moe_w * 2); - s.shared_act = try s.add(gpa, d, cap * moe_w * 2); - s.shared_out = try s.add(gpa, d, cap * dims * 4); - s.moe_act = try s.add(gpa, d, pairs * moe_w * 2); - s.y = try s.add(gpa, d, pairs * dims * 4); - s.branch = try s.add(gpa, d, cap * dims * 2); - s.moe_f32 = try s.add(gpa, d, cap * dims * 4); - s.pocket_act = try s.add(gpa, d, cap * pocket_stride * 4); - s.ple_emb = try s.add(gpa, d, cap * dims * 2); - s.ple_xs = try s.add(gpa, d, cap * (dims / 32) * 4); - s.ple_keys = try s.add(gpa, d, cap * wide * 2); - s.ple_vals = try s.add(gpa, d, cap * dims * 2); - s.ple_part = try s.add(gpa, d, 4 * cap * dims * 4); - s.ple_gated = try s.add(gpa, d, cap * wide * 2); - s.ple_pss = try s.add(gpa, d, cap * streams * 4); - s.ple_nrow = try s.add(gpa, d, cap * wide * 2); - s.logits = try s.add(gpa, d, vocab * 2); - s.head_local = try s.add(gpa, d, cap * vocab); + const sizes = allocationSizes(cap, format); + // a row of a decode round or a block of att_rows rows takes `pos`'s 16 bytes var inv: [half]f32 = undefined; for (&inv, 0..) |*o, i| o.* = @floatCast(std.math.pow(f64, 10_000_000.0, -@as(f64, @floatFromInt(i)) / @as(f64, half))); - s.inv = try s.add(gpa, d, half * 4); - try s.bufs.items[s.bufs.items.len - 1].uploadAsync(0, std.mem.sliceAsBytes(&inv), stream.handle); - s.dummy = try s.add(gpa, d, 4096); - try s.bufs.items[s.bufs.items.len - 1].fill8(0, stream.handle); + inline for (order, 0..) |name, i| { + @field(s, name) = if (sizes[i]) |len| try s.add(gpa, d, len) else 0; + if (sizes[i] != null) { + const last = &s.bufs.items[s.bufs.items.len - 1]; + // part's second cap-stride slice stays zeros (native lanes write only the first); dummy is all zeros + if (comptime std.mem.eql(u8, name, "part")) { + try last.fill8(0, stream.handle); + s.part_zero = true; + } + if (comptime std.mem.eql(u8, name, "inv")) try last.uploadAsync(0, std.mem.sliceAsBytes(&inv), stream.handle); + if (comptime std.mem.eql(u8, name, "dummy")) try last.fill8(0, stream.handle); + } + } try stream.synchronize(); s.host = try cuda.HostBuffer.alloc(d, hostAt(cap)[4]); s.staging = .{ .q = .{ .stream = stream, .event = try cuda.Event.init(d, false) } }; @@ -483,42 +503,12 @@ pub const Model = struct { pocket: ?MoeAdd = null, }; -/// The rows' n-gram embeddings, gathered on a host thread while the GPU embeds and runs layer 0. -pub const Gather = struct { - ngram: *const ple.Ngram, - history: []i64, - ids: []const u32, - emb: []u16, - xs: []f32, - failed: ?anyerror = null, - ns: i96 = 0, - io: std.Io, - /// A round's rows, each gathered after its own stream's history (empty: every row after `history`). - lanes: []const Lane = &.{}, - - pub fn main(g: *Gather) void { - const t0 = std.Io.Clock.awake.now(g.io).toNanoseconds(); - defer g.ns = std.Io.Clock.awake.now(g.io).toNanoseconds() - t0; - if (g.lanes.len == 0) { - g.ngram.rows(g.history, g.ids, g.emb, g.xs) catch |err| { - g.failed = err; - }; - return; - } - var row0: usize = 0; - for (g.lanes) |l| { - g.ngram.rows(l.seq.history[0..], g.ids[row0..][0..l.rows], g.emb[row0 * dims ..][0 .. l.rows * dims], g.xs[row0 * (dims / 32) ..][0 .. l.rows * (dims / 32)]) catch |err| { - g.failed = err; - return; - }; - row0 += l.rows; - } - } -}; +pub const Gather = @import("cuda_batch_gather.zig").Gather; test { _ = @import("cuda_batch_state_test.zig"); _ = @import("cuda_nvfp4_tile.zig"); + _ = @import("cuda_scratch_layout_test.zig"); } pub const Pg = @import("cuda_glue.zig").Pg; diff --git a/zig/src/families/flashnext/cuda_config.zig b/zig/src/families/flashnext/cuda_config.zig index 41b4099220..659df5c70b 100644 --- a/zig/src/families/flashnext/cuda_config.zig +++ b/zig/src/families/flashnext/cuda_config.zig @@ -3,6 +3,7 @@ const std = @import("std"); const wt = @import("cuda_prompt_wt.zig"); const gdn = @import("cuda_gdn.zig"); +const exl3 = @import("exl3/pack/cuda_exl3_config.zig"); /// One config field and the value the kernels and buffers assume. const Want = struct { name: []const u8, value: i64 }; @@ -100,7 +101,7 @@ pub fn mismatch(text: std.json.ObjectMap, quant: ?std.json.ObjectMap) ?[]const u return geometry(text); } -fn geometry(text: std.json.ObjectMap) ?[]const u8 { +pub fn geometry(text: std.json.ObjectMap) ?[]const u8 { // a folder Living Weights taught widens the shared expert by its pocket (a whole number of groups, 128 at most) const learned = int(text, learned_key) orelse 0; if (learned < 0 or learned > 128 or @mod(learned, 32) != 0) return learned_key; @@ -111,7 +112,7 @@ fn geometry(text: std.json.ObjectMap) ?[]const u8 { return semantics(text); } -pub const Mode = enum { affine, modelopt }; +pub const Mode = enum { affine, modelopt, exl3 }; fn projection(q: std.json.ObjectMap, name: []const u8, algo: []const u8, group: i64) bool { const p = object(q, name) orelse return false; @@ -201,8 +202,8 @@ pub fn mode(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) !Mode { const root = parsed.value.object; const text = if (root.get("text_config")) |v| (if (v == .object) v.object else root) else root; const quant = if (root.get("quantization")) |v| (if (v == .object) v.object else null) else null; - const selected: Mode = if (quant != null) .affine else .modelopt; - var field = if (selected == .affine) mismatch(text, quant) else if (object(root, "quantization_config")) |q| modeloptMismatch(text, q) else "quantization"; + const selected: Mode = if (quant != null) .affine else if (exl3.named(root)) .exl3 else .modelopt; + var field = if (selected == .affine) mismatch(text, quant) else if (selected == .exl3) exl3.mismatch(root, text, geometry) else if (object(root, "quantization_config")) |q| modeloptMismatch(text, q) else "quantization"; if (field == null) if (given(root, "tie_word_embeddings")) |v| if (!holds(v, .{ .flag = false })) { field = "tie_word_embeddings"; }; diff --git a/zig/src/families/flashnext/cuda_embed.zig b/zig/src/families/flashnext/cuda_embed.zig index 453ba86271..d79a64f84c 100644 --- a/zig/src/families/flashnext/cuda_embed.zig +++ b/zig/src/families/flashnext/cuda_embed.zig @@ -83,6 +83,8 @@ pub const Table = struct { rows: usize, dims: usize, dense: bool = false, + /// A dense table kept in pinned host memory the kernels read in place (`w` borrows its device address). + host: ?cuda.HostBuffer = null, /// Copies the three tensors. `dims` is hidden size (eight values per stored word). pub fn upload(d: *const cuda.Driver, weight: Tensor, scales: Tensor, biases: Tensor) !Table { @@ -111,7 +113,27 @@ pub const Table = struct { return .{ .w = try cuda.DeviceBuffer.fromHost(d, weight.bytes), .s = zero, .b = zero, .rows = weight.dim(0), .dims = weight.dim(1), .dense = true }; } + /// A dense table in mapped host memory when `host`, else in VRAM as main has it. + pub fn uploadDenseIn(d: *const cuda.Driver, weight: Tensor, host: bool) !Table { + return if (host) uploadDenseHost(d, weight) else uploadDense(d, weight); + } + + /// Own a dense table in mapped host memory instead of VRAM: the gather reads its rows over PCIe. + pub fn uploadDenseHost(d: *const cuda.Driver, weight: Tensor) !Table { + if (weight.dtype != .bf16 or weight.rank != 2 or weight.dim(0) == 0 or weight.dim(1) == 0 or weight.dim(1) % 8 != 0) + return error.UnexpectedTensor; + const count = try std.math.mul(usize, weight.dim(0), weight.dim(1)); + if (weight.bytes.len != try std.math.mul(usize, count, 2)) return error.UnexpectedTensor; + var host = try cuda.HostBuffer.allocMapped(d, weight.bytes.len); + errdefer host.free(); + @memcpy(host.bytes, weight.bytes); + const zero: cuda.DeviceBuffer = .{ .d = d, .ptr = 0, .len = 0 }; + const w: cuda.DeviceBuffer = .{ .d = d, .ptr = try host.device(), .len = weight.bytes.len, .borrowed = true }; + return .{ .w = w, .s = zero, .b = zero, .rows = weight.dim(0), .dims = weight.dim(1), .dense = true, .host = host }; + } + pub fn deinit(self: *Table) void { + if (self.host) |*h| h.free(); self.w.free(); self.s.free(); self.b.free(); diff --git a/zig/src/families/flashnext/cuda_lanes.zig b/zig/src/families/flashnext/cuda_lanes.zig index 5fe5814d7d..fc28b9dcb3 100644 --- a/zig/src/families/flashnext/cuda_lanes.zig +++ b/zig/src/families/flashnext/cuda_lanes.zig @@ -27,18 +27,18 @@ pub fn Namespace(comptime Tri: type) type { return struct { pub const model_type = "qwen4_exp"; pub const tp_draft_snapshots = false; - pub const formats: []const []const u8 = &.{ "mlx-q4g32", "modelopt-mixed" }; + pub const formats: []const []const u8 = &.{ "mlx-q4g32", "modelopt-mixed", @import("exl3/pack/cuda_exl3_config.zig").format }; // the model's window: past 2,048 keys the block indexer picks each row's 512 best blocks of 4 (cuda_batch.zig) pub const default_context: i64 = @intCast(prompt.max_tokens); // a prompt runs on the batched pass in chunks of this many rows, one a prefill_step: no staggered segments pub const max_segments: u32 = 1; pub const prompt_rows: u32 = @intCast(prompt.batch.default_chunk); /// The prompt cache's copies of a stream's state (cuda_snap.zig), called with `Loaded.target`. - pub const snap_bytes = snapBytes; + pub const snap_bytes = copies.bytes; pub const snap_repeat_min: u32 = 16; - pub const snap_save = snapSave; - pub const snap_restore = snapRestore; - pub const snap_drop = snapDrop; + pub const snap_save = copies.save; + pub const snap_restore = copies.restore; + pub const snap_drop = copies.drop; pub const kept_disk = @import("cuda_disk.zig").Kept(@This()); // their disk image, as qwen3_5 declares its own /// TF_FLASHNEXT_KV_PAGES: the pool's lanes, the pages of a request's rows and a drafted window (kvl). pub const window_rows = kvl.window_rows; @@ -57,7 +57,10 @@ pub fn Namespace(comptime Tri: type) type { pub const min_gap: u32 = 256; pub const min_prompt: u32 = 512; }; - pub const Options = struct { vision: ?@import("cuda_images.zig").Limits = null, context: usize, drafts: bool, segments: usize = 1, slide: bool = false, tp: usize = 1, rank: usize = 0, master: ?[]const u8 = null, master_port: u16 = 29551 }; + + pub const Options = struct { vision: ?@import("cuda_images.zig").Limits = null, context: usize, drafts: bool, segments: usize = 1, slide: bool = false, lanes: u32 = 1, tp: usize = 1, rank: usize = 0, master: ?[]const u8 = null, master_port: u16 = 29551 }; + pub const preallocates_lanes = false; + /// Living Weights' learner (--slide): lessons put neurons in each layer's pocket beside its shared expert. pub const slide = slide_mod.Slide(Tri); /// The first file a lesson would rewrite through a link, or null: a teaching copy may hard-link the others. @@ -79,6 +82,10 @@ pub fn Namespace(comptime Tri: type) type { stream_bytes: usize, /// Every stream allocates its own sequence; none rides on the engine's. free_streams: u32 = 0, + /// Kept prompt states in pinned host memory (cuda_snap.on_host): the native host budgets them there. + host_snaps: bool = false, + /// The prefill step when the prompt arithmetic depends on the chunk cuts (EXL3): states keep on its starts + snap_grid: u32 = 0, ctx: *anyopaque, deinit: *const fn (*anyopaque) void, lone: ?LoneRun = null, @@ -134,6 +141,7 @@ pub fn Namespace(comptime Tri: type) type { }; const Ops = mtp.Ops(Tri); + const copies = @import("cuda_lanes_snap.zig").Copies(of); pub const Cuda = struct { gpa: std.mem.Allocator, @@ -142,6 +150,8 @@ pub fn Namespace(comptime Tri: type) type { vision: ?*@import("cuda_images.zig").Vision = null, ranks: ?*ranks.Session = null, head: ?*mtp.Head = null, + profile: ?@import("exl3/offload/cuda_exl3_frequency_state.zig").State = null, + profile_counts: ?@import("exl3/offload/cuda_exl3_frequency.zig").Counts = null, depth: usize = mtp.default_depth, confidence: f64 = mtp.default_confidence, /// A captured Triton set (TENSORFOLD_CUDA_KERNELS) for comparison, else null: our own glue. @@ -202,7 +212,7 @@ pub fn Namespace(comptime Tri: type) type { .score = @import("cuda_score.zig").Hook(Cuda, Tri).call, .first_masked = firstMaskedFn, }; - const ranked_table = @import("cuda_score.zig").withoutScore(scoring_table); + const ranked_table = comptime @import("cuda_score.zig").withoutScore(scoring_table); // static: backend() returns its address return .{ .ptr = self, .proposal_head = if (self.head) |h| &h.service else null, .vtable = if (self.ranks == null) &scoring_table else &ranked_table }; } @@ -320,6 +330,9 @@ pub fn Namespace(comptime Tri: type) type { }; const last = stop == ids.len; self.engine.ngram.warmAhead(ids, stop, @min(ids.len, stop + self.engine.chunk)); + const previous_telemetry_phase = self.engine.telemetry.forward_phase; + self.engine.telemetry.forward_phase = .prefill; + defer self.engine.telemetry.forward_phase = previous_telemetry_phase; if (l.draft) |*d| { try self.ops().prefillChunk(&l.seq, d, ids, at, stop); } else if (self.ranks) |r| { @@ -354,6 +367,7 @@ pub fn Namespace(comptime Tri: type) type { pub fn release(p: *anyopaque) void { const own: *Cuda = @ptrCast(@alignCast(p)); cuda.completion.drain(own.engine.driver, own.stream, "flashnext: release did not drain"); + if (own.profile) |*profile| profile.save(own); var it = own.lanes.valueIterator(); while (it.next()) |l| dropLane(own, l); own.lanes.deinit(own.gpa); @@ -367,6 +381,7 @@ pub fn Namespace(comptime Tri: type) type { if (own.set) |*x| x.deinit(); own.glue.deinit(); own.stream.deinit(); + if (own.profile) |*profile| profile.deinit(); own.gpa.destroy(own); } @@ -544,7 +559,7 @@ pub fn Namespace(comptime Tri: type) type { test { _ = @import("cuda_images_test.zig").Cases(Cuda, Cuda.open, dropLane); - _ = @import("cuda_lanes_restore_test.zig").Cases(Cuda, Cuda.open, snapRestore, dropLane); + _ = @import("cuda_lanes_restore_test.zig").Cases(Cuda, Cuda.open, copies.restore, dropLane); _ = @import("cuda_lanes_admission_test.zig").Cases(Cuda, Cuda.open, dropLane); _ = @import("cuda_score.zig").CapabilityCases(Cuda); } @@ -560,41 +575,5 @@ pub fn Namespace(comptime Tri: type) type { try std.testing.expectError(bad[2], open(std.testing.allocator, std.testing.io, &ctx, "no-such-checkpoint", null, .{ .context = 4096, .drafts = true })); } } - - // -- the prompt cache's copies ------------------------------------------------------------------------ - - fn snapBytes(ptr: *anyopaque, at: u32) u64 { - return snap.bytes(at, of(ptr).head != null); - } - - /// The owner stream's state, its prompt pass standing at `at` (a drafting stream's with the head's rows). - fn snapSave(ptr: *anyopaque, owner: ?*anyopaque, at: u32) anyerror!*anyopaque { - const self = of(ptr); - const s: *lanes.Stream = @ptrCast(@alignCast(owner orelse return error.NoStream)); - const l = try self.lane(s); - if (l.seq.fed != at) return error.NotAtMark; - if (self.ranks) |r| { - if (l.draft != null) return error.TensorRankDraftSnapshotUnavailable; - return try r.save(&l.seq); - } - return try snap.save(self.gpa, self.engine.driver, self.stream, &l.seq, if (l.draft) |*d| d else null); - } - - fn snapRestore(ptr: *anyopaque, owner: ?*anyopaque, saved: *anyopaque) anyerror!void { - const self = of(ptr); - const s: *lanes.Stream = @ptrCast(@alignCast(owner orelse return error.NoStream)); - const l = try self.lane(s); - if (self.ranks) |r| { - if (l.draft != null) return error.NoHeadState; - return r.restore(&l.seq, @ptrCast(@alignCast(saved))); - } - try snap.restore(self.stream, &l.seq, if (l.draft) |*d| d else null, @ptrCast(@alignCast(saved))); - } - - fn snapDrop(ptr: *anyopaque, saved: *anyopaque) void { - const self = of(ptr); - if (self.ranks) |r| return r.dropSaved(@ptrCast(@alignCast(saved))); - snap.drop(self.gpa, self.stream, @ptrCast(@alignCast(saved))); - } }; } diff --git a/zig/src/families/flashnext/cuda_lanes_bytes.zig b/zig/src/families/flashnext/cuda_lanes_bytes.zig index bb959edc89..556ef895a7 100644 --- a/zig/src/families/flashnext/cuda_lanes_bytes.zig +++ b/zig/src/families/flashnext/cuda_lanes_bytes.zig @@ -3,9 +3,11 @@ const std = @import("std"); const weights = @import("cuda_prompt_wt.zig"); const mtp = @import("cuda_mtp.zig"); const config = @import("cuda_config.zig"); +const exl3_memory = @import("exl3/pack/cuda_exl3_memory.zig"); /// Device bytes of the checkpoint: the n-gram tables stay mapped, the MTP head counts; 0 if unread. pub fn deviceWeightBytes(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) u64 { + if (exl3_memory.isExl3(gpa, io, dir)) return exl3_memory.weightBytes(gpa, io, dir); const main = weights.deviceBytes(gpa, io, dir) catch return 0; return main + mtp.headBytes(); } @@ -13,6 +15,7 @@ pub fn deviceWeightBytes(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) u6 pub fn placementWeightBytes(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, world: usize, rank: usize, drafts: bool) !u64 { if ((world != 1 and world != 2) or rank >= world) return error.BadRank; const mode = try config.mode(gpa, io, dir); + if (mode == .exl3) return exl3_memory.placement(gpa, io, dir, world, rank); if (mode == .modelopt) { const main = try @import("cuda_rank_memory.zig").estimate(gpa, io, dir, world, rank); // Dense draft projections and vocabulary rows exceed the affine head's estimate. diff --git a/zig/src/families/flashnext/cuda_lanes_load.zig b/zig/src/families/flashnext/cuda_lanes_load.zig index 508ce51c51..3513ebf653 100644 --- a/zig/src/families/flashnext/cuda_lanes_load.zig +++ b/zig/src/families/flashnext/cuda_lanes_load.zig @@ -7,6 +7,8 @@ const ranks = @import("cuda_rank_engine.zig"); const config = @import("cuda_config.zig"); const Glue = @import("cuda_glue.zig").Glue; const images = @import("cuda_images.zig"); +const memory = @import("exl3/pack/exl3_startup_memory.zig"); +const exl3_pass = @import("exl3/model/cuda_exl3_pass.zig"); pub fn open(comptime F: type, gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: []const u8, captured: ?[]const u8, o: F.Options) !F.Loaded { if (o.vision) |v| try @import("qwen_vision").media.checkLimits(v); const Cuda = F.Cuda; @@ -19,11 +21,12 @@ pub fn open(comptime F: type, gpa: std.mem.Allocator, io: std.Io, ctx: *const cu const chunk = try prompt.chunkRows(); _ = try @import("cuda_modelopt_threads.zig").configured(); // TF_FLASHNEXT_LOAD_THREADS, refused before any load const mode = try config.mode(gpa, io, dir); + if (mode == .exl3) try @import("exl3/model/cuda_exl3_arch.zig").check(ctx.d); // sm_120 only: refused before any weight loads if (o.tp != 1 and (mode != .modelopt or captured != null)) return error.UnsupportedQuantization; if (o.slide and (mode != .affine or o.tp != 1)) return error.UnsupportedLearningFormat; if (o.slide and captured != null) return error.TrainingOnCapturedGlue; const page_tokens = try @import("cuda_kv_lanes.zig").paging(); - if (page_tokens > 0 and (o.tp != 1 or o.slide or captured != null)) return error.FlashNextPagesUnsupported; + if (page_tokens > 0 and (o.tp != 1 or o.slide or captured != null or mode == .exl3)) return error.FlashNextPagesUnsupported; // the pool's pages are not wired to EXL3's K/V const driver: *cuda.Driver = @constCast(ctx.d); const own = try gpa.create(Cuda); errdefer gpa.destroy(own); @@ -61,11 +64,16 @@ pub fn open(comptime F: type, gpa: std.mem.Allocator, io: std.Io, ctx: *const cu own.engine.image_factory = own.vision; } errdefer if (own.vision) |v| v.deinit(); + var persistent_bytes = cuda.usage(false).device; + try @import("exl3/kv/cuda_exl3_kv.zig").setup(own.engine, mode == .exl3); // K10: the scratch format and the chunk of an EXL3 pack // the pass's scratch for a full chunk, once for the engine (every prompt, window, round and head pass) _ = try own.engine.reserve(@max(@min(o.context, own.engine.chunk), if (o.slide) @as(usize, @import("cuda_train.zig").pack_rows) else 0)); // the MTP head unless --no-drafts (TF_FLASHNEXT_DRAFTS and _CONFIDENCE: the chain's rule, 6 and 0.60) - if (o.drafts and (o.tp == 1 or o.rank == 0)) { + if (o.drafts and mode == .exl3) std.log.info("flashnext: an EXL3 pack decodes without the MTP head", .{}); + if (o.drafts and mode != .exl3 and (o.tp == 1 or o.rank == 0)) { + const head_before = cuda.usage(false).device; own.head = try mtp.Head.init(gpa, io, driver, &own.stream, dir); + persistent_bytes = try std.math.add(u64, persistent_bytes, cuda.usage(false).device -| head_before); own.depth = @max(1, @min(mtp.max_chain, F.envUsize("TF_FLASHNEXT_DRAFTS", mtp.default_depth))); if (std.c.getenv("TF_FLASHNEXT_CONFIDENCE")) |v| own.confidence = std.fmt.parseFloat(f64, std.mem.span(v)) catch mtp.default_confidence; } @@ -79,18 +87,38 @@ pub fn open(comptime F: type, gpa: std.mem.Allocator, io: std.Io, ctx: *const cu own.lanes = .empty; own.next = 0; if (o.slide) own.lay = try @import("cuda_slide.zig").layoutOf(gpa, io, dir); + const scratch_bytes = cuda.usage(false).device -| persistent_bytes; const own_seq_bytes = try prompt.seqBytesPaged(window, try @import("cuda_pages.zig").pageTest()); var stream_bytes = own_seq_bytes; if (own.head) |h| { - stream_bytes += mtp.Draft.bytes(window, h.ids.len, own.engine.chunk); + stream_bytes = try std.math.add(u64, stream_bytes, mtp.Draft.bytes(window, h.ids.len, own.engine.chunk)); } else if (o.drafts and o.tp == 2) stream_bytes += mtp.Draft.bytes(window, 1, own.engine.chunk); - if (o.vision) |v| stream_bytes += @as(usize, v.image_tokens) * @import("cuda_prompt_wt.zig").dims * 2 + 16 * o.context; + if (o.vision) |v| stream_bytes = try std.math.add(u64, stream_bytes, try std.math.add(u64, try std.math.mul(u64, v.image_tokens, @import("cuda_prompt_wt.zig").dims * 2), try std.math.mul(u64, 16, o.context))); if (page_tokens > 0) { // the lane's head runs at the engine's chunk: its cache on the stream's pages, its buffers in the run own.kv.head = if (own.head) |h| .{ .draft_n = h.ids.len, .rows = own.engine.chunk } else null; const head_bytes = if (own.head) |h| mtp.Draft.bytes(window, h.ids.len, own.engine.chunk) else 0; own.kv.layout = @import("cuda_kv_lanes.zig").layout(o.context, page_tokens, own.kv.head, stream_bytes - own_seq_bytes - head_bytes); } + try @import("exl3/kv/cuda_exl3_kv.zig").logStream(own.engine, window, stream_bytes); // K10: what a stream holds, by piece + const pool_budget = try exl3_pass.budgetForStreams(stream_bytes, o.lanes); + const placement = @import("exl3/offload/cuda_exl3_frequency_state.zig"); + if (try placement.prepare(gpa, io, dir, own.engine.layers[0..own.engine.n_layers], @intCast(o.tp), @intCast(o.rank))) |prepared| { + own.profile = prepared.state; + own.profile_counts = prepared.counts; + } + errdefer if (own.profile) |*profile| profile.deinit(); + const profile_counts = if (own.profile_counts) |*counts| &counts.decoder else null; + const attached = try exl3_pass.attach(gpa, io, driver, own.stream, own.engine.layers[0..own.engine.n_layers], own.engine.chunk, pool_budget, profile_counts); + if (own.profile_counts) |*counts| try placement.finish(own, counts); + const ledger: memory.Ledger = .{ + .persistent = persistent_bytes, + .scratch = scratch_bytes, + .pool_scratch = attached.scratch_bytes, + .expert_pools = attached.device_bytes, + .reservation = .{ .stream_bytes = stream_bytes, .streams = if (stream_bytes == 0) 0 else pool_budget.other / stream_bytes }, + }; + std.log.info("flashnext rank {d} memory ledger: persistent={d} scratch={d} pool_scratch={d} expert_pools={d} stream_bytes={d} streams={d} planned_total={d} free_after_pools={d}", .{ o.rank, ledger.persistent, ledger.scratch, ledger.pool_scratch, ledger.expert_pools, stream_bytes, ledger.reservation.streams, try ledger.total(), attached.free_after }); return .{ .layout = own.kv.layout, .vision = own.vision, @@ -98,7 +126,9 @@ pub fn open(comptime F: type, gpa: std.mem.Allocator, io: std.Io, ctx: *const cu .facts = own.facts(), .rows = if (own.head != null) mtp.max_rows else 2, .prompt_rows = @intCast(own.engine.chunk), + .snap_grid = if (mode == .exl3) @intCast(own.engine.chunk) else 0, .stream_bytes = stream_bytes, + .host_snaps = @import("exl3/kv/cuda_exl3_kv.zig").snapHost(o.context, mode == .exl3), .ctx = own, .deinit = F.release, .learns = if (o.slide) own else null, diff --git a/zig/src/families/flashnext/cuda_lanes_snap.zig b/zig/src/families/flashnext/cuda_lanes_snap.zig new file mode 100644 index 0000000000..fe8cc9b71c --- /dev/null +++ b/zig/src/families/flashnext/cuda_lanes_snap.zig @@ -0,0 +1,51 @@ +//! The prompt cache's copies of a Flash Next stream (cuda_snap.zig), called with `Loaded.target`. + +const std = @import("std"); +const lanes = @import("lanes"); +const snap = @import("cuda_snap.zig"); + +/// Size, save, restore and drop over the lane owner `of` finds behind a target pointer. +pub fn Copies(comptime of: anytype) type { + return struct { + pub fn bytes(ptr: *anyopaque, at: u32) u64 { + return snap.bytes(at, of(ptr).head != null); + } + + /// The owner stream's state, its prompt pass standing at `at` (a drafting stream's with the head's rows). + pub fn save(ptr: *anyopaque, owner: ?*anyopaque, at: u32) anyerror!*anyopaque { + const self = of(ptr); + const s: *lanes.Stream = @ptrCast(@alignCast(owner orelse return error.NoStream)); + const l = try self.lane(s); + if (l.seq.fed != at) return error.NotAtMark; + if (self.ranks) |r| { + if (l.draft != null) return error.TensorRankDraftSnapshotUnavailable; + return try r.save(&l.seq); + } + return try snap.save(self.gpa, self.engine.driver, self.stream, &l.seq, if (l.draft) |*d| d else null); + } + + pub fn restore(ptr: *anyopaque, owner: ?*anyopaque, saved: *anyopaque) anyerror!void { + const self = of(ptr); + const s: *lanes.Stream = @ptrCast(@alignCast(owner orelse return error.NoStream)); + const l = self.lane(s) catch |err| { + std.log.warn("flashnext: a kept state was not restored: no lane ({t})", .{err}); + return err; + }; + if (self.ranks) |r| { + if (l.draft != null) return error.NoHeadState; + return r.restore(&l.seq, @ptrCast(@alignCast(saved))); + } + self.engine.round_graph.clear(); + snap.restore(self.stream, &l.seq, if (l.draft) |*d| d else null, @ptrCast(@alignCast(saved))) catch |err| { + std.log.warn("flashnext: a kept state was not restored ({t})", .{err}); + return err; + }; + } + + pub fn drop(ptr: *anyopaque, saved: *anyopaque) void { + const self = of(ptr); + if (self.ranks) |r| return r.dropSaved(@ptrCast(@alignCast(saved))); + snap.drop(self.gpa, self.stream, @ptrCast(@alignCast(saved))); + } + }; +} diff --git a/zig/src/families/flashnext/cuda_launch.zig b/zig/src/families/flashnext/cuda_launch.zig index fb3fce8d91..4171f2e5fc 100644 --- a/zig/src/families/flashnext/cuda_launch.zig +++ b/zig/src/families/flashnext/cuda_launch.zig @@ -6,6 +6,7 @@ const aot = cuda.aot; pub const embed = @import("cuda_embed.zig"); pub const projection = @import("cuda_projection.zig"); +pub const exl3 = @import("exl3/model/cuda_exl3.zig"); pub const rank_engine = @import("cuda_rank_engine.zig"); pub const collective = @import("cuda_collective.zig"); pub const rank_wire = @import("cuda_rank_wire.zig"); @@ -22,8 +23,11 @@ pub const moe = @import("cuda_moe.zig"); pub const attn = @import("cuda_attn.zig"); pub const ple = @import("cuda_ple.zig"); pub const ple_check = @import("cuda_ple_check.zig"); +/// An EXL3 pack's n-gram table (cuda_ple.Ngram's `.exl3` format) +pub const exl3_ple = @import("exl3/model/cuda_exl3_ple.zig"); pub const mix = @import("cuda_mix.zig"); pub const prompt = @import("cuda_prompt.zig"); +pub const snapshot = @import("cuda_snap.zig"); /// The row-by-row reference path, for flashnext-prompt pub const row_path = @import("cuda_rows.zig"); /// The MTP head's drafts: its forward, the chain, the verify window's rows and their rollback @@ -32,6 +36,7 @@ pub const mtp = @import("cuda_mtp.zig"); pub const pocket = @import("cuda_pocket.zig"); /// Living Weights' trainer over the batched pass pub const train = @import("cuda_train.zig"); +pub const exl3_pool = @import("exl3/offload/cuda_exl3_pool.zig"); /// Decisions' label scores (the lanes backend's `score`) pub const score = @import("cuda_score.zig"); pub const native = @import("cuda_lanes.zig").Namespace(Tri); @@ -431,6 +436,7 @@ test { _ = @import("cuda_pages.zig"); _ = @import("cuda_kv_lanes.zig"); _ = @import("cuda_config.zig"); + _ = @import("exl3/model/cuda_exl3.zig"); _ = @import("cuda_projection_load.zig"); _ = @import("cuda_rank_engine.zig"); _ = @import("cuda_round_graph.zig"); @@ -448,6 +454,10 @@ test { _ = @import("living.zig"); } +test "exl3: dispatch host contracts" { + _ = @import("exl3/model/cuda_exl3_dispatch_test.zig"); +} + test "embed grid is one program per row and per group of 32" { try std.testing.expectEqual([3]u32{ 4, 80, 1 }, embedGrid(4, 2560)); } diff --git a/zig/src/families/flashnext/cuda_paging_batch_test.zig b/zig/src/families/flashnext/cuda_paging_batch_test.zig index 356cf096ac..70ae7ad70b 100644 --- a/zig/src/families/flashnext/cuda_paging_batch_test.zig +++ b/zig/src/families/flashnext/cuda_paging_batch_test.zig @@ -121,12 +121,14 @@ test "paged grouped prompts: every current attention route receives the sequence s.part = 0x1300000000; var lw: Layer = undefined; lw.proj.native = null; + lw.proj.exl3 = null; lw.proj.w.ptr = 0x1000; lw.proj.s.ptr = 0x2000; lw.proj.b.ptr = 0x3000; lw.proj.k = 7680; lw.proj.n = 13952; lw.out.native = null; + lw.out.exl3 = null; lw.out.w.ptr = 0x4000; lw.out.s.ptr = 0x5000; lw.out.b.ptr = 0x6000; diff --git a/zig/src/families/flashnext/cuda_ple.zig b/zig/src/families/flashnext/cuda_ple.zig index e4ba3ca438..f5e4b11b63 100644 --- a/zig/src/families/flashnext/cuda_ple.zig +++ b/zig/src/families/flashnext/cuda_ple.zig @@ -5,6 +5,7 @@ const cuda = @import("cuda"); const core = @import("core"); const embed = @import("cuda_embed.zig"); const hc = @import("cuda_hc.zig"); +const exl3_ple = @import("exl3/model/cuda_exl3_ple.zig"); const ple_pool = @import("cuda_ple_pool.zig"); const ple_reader = @import("cuda_ple_reader.zig"); @@ -343,11 +344,12 @@ pub const Ngram = struct { w: [max_shards][]const u8 = undefined, s: [max_shards][]const u8 = undefined, b: [max_shards][]const u8 = undefined, - format: enum { mlx, nvfp4 } = .mlx, + format: enum { mlx, nvfp4, exl3 } = .mlx, tables: [max_shards]core.nvfp4_rows.Table = undefined, readers: [max_shards]ple_reader.Reader = undefined, read_mode: ple_reader.Mode = .mmap, warm_ahead: bool = false, + x3: ?*exl3_ple.Table = null, pool: ?*ple_pool.Pool = null, // the gather's parked workers (null: every row on the caller's thread) pub fn init(gpa: std.mem.Allocator, io: std.Io, model_dir: []const u8) !Ngram { @@ -365,6 +367,7 @@ pub const Ngram = struct { .ple_index = cfg.ple_index, }), .scale = 1 }; errdefer g.deinit(); + if (try exl3_ple.Table.open(gpa, io, model_dir, g.hash, exl3_ple.optionsFromEnv())) |t| return .{ .gpa = gpa, .hash = g.hash, .scale = 1, .format = .exl3, .x3 = t }; const index_path = try std.fs.path.join(gpa, &.{ model_dir, "model.safetensors.index.json" }); defer gpa.free(index_path); const text = try std.Io.Dir.cwd().readFileAlloc(io, index_path, gpa, .limited(1 << 26)); @@ -424,6 +427,7 @@ pub const Ngram = struct { } pub fn deinit(g: *Ngram) void { + if (g.x3) |t| t.destroy(); if (g.pool) |p| p.destroy(); for (g.files.items) |*o| o.close(g.gpa); g.files.deinit(g.gpa); @@ -446,6 +450,7 @@ pub const Ngram = struct { try g.readers[loc.shard].read(g.gpa, g.tables[loc.shard], loc.local, row) else try g.tables[loc.shard].gather(&.{loc.local}, row), + .exl3 => try g.x3.?.row(loc.shard, loc.local, i % g.hash.heads, row), } } } @@ -481,7 +486,8 @@ pub const Ngram = struct { for (0..history.len - 1) |i| hist[i] = hist[i + 1]; hist[history.len - 1] = token; } - const n_threads = @min(ple_pool.max_parts, items / 4); + // EXL3 rows decode in about a microsecond: a decode step's few rows cost less than one thread spawn + const n_threads = if (g.format == .exl3 and items <= 1024) 1 else @min(ple_pool.max_parts, items / 4); if (g.format == .nvfp4 and items > 1) { try @import("cuda_ple_gather.zig").run(g, ids, out); } else if (n_threads <= 1 or g.pool == null) { @@ -498,7 +504,7 @@ pub const Ngram = struct { v.* = toBf16(f * g.scale); } } - sumGroups(row, xs[t * (w / 32) ..][0 .. w / 32]); + if (g.x3) |x3| x3.sums(row, xs[t * (w / 32) ..][0 .. w / 32]) else sumGroups(row, xs[t * (w / 32) ..][0 .. w / 32]); } @memcpy(history, hist[0..history.len]); } @@ -530,6 +536,7 @@ test "a fresh window hashes token 7 into sixteen rows" { test { _ = @import("cuda_ple_nvfp4_test.zig"); + _ = @import("exl3/model/cuda_exl3_ple.zig"); _ = @import("cuda_ple_pool.zig"); _ = @import("cuda_ple_pool_gather_test.zig"); } diff --git a/zig/src/families/flashnext/cuda_projection.zig b/zig/src/families/flashnext/cuda_projection.zig index c0b96adb3b..0d678b580b 100644 --- a/zig/src/families/flashnext/cuda_projection.zig +++ b/zig/src/families/flashnext/cuda_projection.zig @@ -10,8 +10,11 @@ pub const Projection = struct { n: usize, k: usize, native: ?cuda.qmmf.Weight = null, + /// An EXL3 checkpoint's projection: stored-layout EXL3 parts (and dense BF16 ones) a row (cuda_exl3_load.zig). + exl3: ?*@import("exl3/pack/cuda_exl3_load.zig").Proj = null, pub fn free(self: *Projection) void { + if (self.exl3) |e| e.destroy(); self.w.free(); self.s.free(); self.b.free(); diff --git a/zig/src/families/flashnext/cuda_projection_load.zig b/zig/src/families/flashnext/cuda_projection_load.zig index c0c37f1b79..461c698be0 100644 --- a/zig/src/families/flashnext/cuda_projection_load.zig +++ b/zig/src/families/flashnext/cuda_projection_load.zig @@ -3,6 +3,7 @@ const std = @import("std"); const cuda = @import("cuda"); const wt = @import("cuda_prompt_wt.zig"); const projection = @import("cuda_projection.zig"); +const exl3 = @import("exl3/pack/cuda_exl3_load.zig"); pub const Micro = struct { n: usize, @@ -69,6 +70,7 @@ fn uploadMicro(a: std.mem.Allocator, d: *const cuda.Driver, value: Micro) !wt.Q4 } pub fn load(a: std.mem.Allocator, d: *cuda.Driver, store: *wt.Store, prefix: []const u8, n: usize, k: usize) !wt.Q4 { + if (exl3.has(store, prefix)) return exl3.load(a, d, store, prefix, n, k); var buf: [512]u8 = undefined; const t = try store.tensor(try std.fmt.bufPrint(&buf, "{s}.weight", .{prefix})); switch (t.dtype) { @@ -96,6 +98,7 @@ pub fn load(a: std.mem.Allocator, d: *cuda.Driver, store: *wt.Store, prefix: []c pub fn loadStitched(a: std.mem.Allocator, d: *cuda.Driver, store: *wt.Store, prefix: []const u8, names: []const []const u8, n: usize, k: usize) !wt.Q4 { if (names.len == 0) return error.UnexpectedTensor; + if (exl3.hasAny(store, prefix, names)) return exl3.stitched(a, d, store, prefix, names, n, k); var buf: [512]u8 = undefined; const first = try store.tensor(try std.fmt.bufPrint(&buf, "{s}{s}.weight", .{ prefix, names[0] })); if (first.dtype == .f8_e4m3) { diff --git a/zig/src/families/flashnext/cuda_prompt.zig b/zig/src/families/flashnext/cuda_prompt.zig index ed2d82c9b1..93932f8803 100644 --- a/zig/src/families/flashnext/cuda_prompt.zig +++ b/zig/src/families/flashnext/cuda_prompt.zig @@ -115,6 +115,7 @@ pub fn pooledBytes(cap: usize) usize { return (cap + weights.index_ratio - 1) / weights.index_ratio * index_dim * 2; } const round_graph = @import("cuda_round_graph.zig"); +const cache_telemetry = @import("cuda_prompt_telemetry.zig"); const pages = @import("cuda_pages.zig"); /// One loaded checkpoint: the embedding table, every layer, the mixer and head, and the batched pass's scratch. pub const Engine = struct { @@ -135,6 +136,7 @@ pub const Engine = struct { ngram: ple.Ngram, /// The pass's scratch, grown to the widest chunk run. batch_s: ?batch.Scratch = null, + scratch_format: @import("cuda_batch_state.zig").ScratchFormat = .classic, /// Block scores for the sparse rows of an attention block, grown to the longest sequence's need. scores: ?cuda.DeviceBuffer = null, /// Rows a batched chunk takes at most (chunkRows). @@ -142,10 +144,14 @@ pub const Engine = struct { round_graph: round_graph.Cache = .{}, image_factory: ?*@import("cuda_images.zig").Vision = null, collective: ?*@import("cuda_collective.zig").Collective = null, + telemetry: cache_telemetry.State = .{}, /// Each layer's Living Weights pocket, from a taught folder's shards or the learner: every pass adds it. pockets: pocket.Pockets = .{}, /// During a round of taught and untaught streams: its rows' pocket scales (0 on an untaught row); 0: none. taught_rows: u64 = 0, + + const measured = cache_telemetry.Operations(Engine, Seq, batch.Hooks, .{ .forward = forwardUnmeasured, .round = roundUnmeasured }); + pub fn init(gpa: std.mem.Allocator, io: std.Io, driver: *cuda.Driver, stream: *cuda.Stream, model_dir: []const u8) !*Engine { return initRank(gpa, io, driver, stream, model_dir, 1, 0); } @@ -161,8 +167,10 @@ pub const Engine = struct { e.driver = driver; e.stream = stream; e.n_layers = 0; + e.scratch_format = .classic; e.collective = null; e.image_factory = null; + e.telemetry = cache_telemetry.State.fromEnv(rank, world); e.round_graph = .{}; errdefer for (e.layers[0..e.n_layers]) |*l| l.free(gpa); e.model_dir = try gpa.dupe(u8, model_dir); @@ -172,7 +180,8 @@ pub const Engine = struct { try e.store.detectNorms(); const weight = try e.store.tensor(embed.weight_name); if (weight.dtype == .bf16) { - e.table = try embed.Table.uploadDense(driver, weight); + // an EXL3 pack keeps its BF16 table in mapped host memory, every other checkpoint on the device + e.table = try embed.Table.uploadDenseIn(driver, weight, @import("exl3/pack/cuda_exl3_load.zig").has(&e.store, "language_model.lm_head")); } else { const scales = try e.store.tensor(embed.scales_name); const biases = try e.store.tensor(embed.biases_name); @@ -182,8 +191,9 @@ pub const Engine = struct { e.eos = try ple.eosOf(gpa, io, e.model_dir); e.mixer = try loadHc(gpa, driver, &e.store, "language_model.model.hyper_connection_mixer", false); errdefer e.mixer.free(); - const head_tensor = try e.store.tensor("language_model.lm_head.weight"); - if (head_tensor.dtype == .bf16) { + const exl3_head = try @import("exl3/pack/cuda_exl3_load.zig").head(gpa, driver, &e.store, world, rank); + const head_tensor = if (exl3_head == null) try e.store.tensor("language_model.lm_head.weight") else undefined; + if (exl3_head) |q| e.head_q = q else if (head_tensor.dtype == .bf16) { if (!head_tensor.is(.bf16, &.{ vocab, dims })) return error.UnexpectedTensor; e.head_q = try @import("cuda_projection.zig").dense(gpa, driver, if (try @import("cuda_head.zig").enabled()) try @import("cuda_head.zig").shard(head_tensor, world, rank) else head_tensor); } else { @@ -345,10 +355,7 @@ pub const Engine = struct { return e.forward(Tri, tri, seq, ids, .{}); } - pub fn complete(e: *Engine) !void { - if (e.collective) |c| return c.finish(e.stream.*); - try e.stream.synchronize(); - } + pub const complete = measured.complete; /// Block scores for a `cap`-token sequence's chunks of up to `rows` rows, reallocated only to grow. pub fn reserveScores(e: *Engine, cap: usize, rows: usize) !u64 { @@ -372,19 +379,22 @@ pub const Engine = struct { /// Scratch for chunks of `rows` rows, reallocated only to grow. pub fn reserve(e: *Engine, rows: usize) !*batch.Scratch { if (e.batch_s) |*bs| { + if (bs.format != e.scratch_format) return error.ScratchFormatChanged; if (bs.cap >= rows) return bs; try e.complete(); e.round_graph.clear(); bs.free(e.gpa); e.batch_s = null; } - e.batch_s = try batch.Scratch.alloc(e.gpa, e.driver, e.stream.*, rows); + e.batch_s = try batch.Scratch.allocForFormat(e.gpa, e.driver, e.stream.*, rows, e.scratch_format); return &e.batch_s.?; } /// The batched pass over `ids` from `seq.fed` in chunks of at most `chunk` rows; a window's hooks take one chunk. - pub fn forward(e: *Engine, comptime Tri: type, tri: Tri, seq: *Seq, ids: []const u32, hooks: batch.Hooks) !void { - if (ids.len == 0 or seq.fed + ids.len > seq.cap) return error.UnexpectedTensor; + pub const forward = measured.forward; + + fn forwardUnmeasured(e: *Engine, comptime Tri: type, tri: Tri, seq: *Seq, ids: []const u32, hooks: batch.Hooks) !void { + if (ids.len == 0 or seq.fed > seq.cap or ids.len > seq.cap - seq.fed) return error.UnexpectedTensor; if ((hooks.keep != null or hooks.logits == .all) and ids.len > e.chunk) return error.UnexpectedTensor; const s = try e.reserve(@min(ids.len, e.chunk)); const scores = try e.reserveScores(seq.cap, s.cap); @@ -417,7 +427,9 @@ pub const Engine = struct { pub const windows = @import("cuda_prompt_windows.zig").windows; /// One decode round: each of `seqs` takes its token of `ids` in one batched pass; scores as if alone. - pub fn round(e: *Engine, comptime Tri: type, tri: Tri, seqs: []const *Seq, ids: []const u32, out: u64) !void { + pub const round = measured.round; + + fn roundUnmeasured(e: *Engine, comptime Tri: type, tri: Tri, seqs: []const *Seq, ids: []const u32, out: u64) !void { if (seqs.len == 0 or seqs.len != ids.len) return error.UnexpectedTensor; var cap: usize = 0; for (seqs) |q| { @@ -500,23 +512,7 @@ pub const Seq = struct { } pub fn deinit(s: *Seq, gpa: std.mem.Allocator) void { - if (s.stream) |stream| cuda.completion.drain(stream.d, stream, "Flash Next sequence work did not drain before release"); - if (s.media) |*m| m.deinit(); - freeBufs(s.conv); - freeBufs(s.state); - freeBufs(s.kc); - freeBufs(s.vc); - freeBufs(s.ikc); - freeBufs(s.pooled); - gpa.free(s.conv); - gpa.free(s.state); - gpa.free(s.kc); - gpa.free(s.vc); - gpa.free(s.ikc); - gpa.free(s.pooled); - if (s.paged) |*p| p.deinit(gpa); - if (s.tail.ptr != 0) s.tail.free(); - gpa.free(s.logits); + @import("cuda_prompt_lifecycle.zig").deinit(s, gpa); } }; @@ -524,55 +520,21 @@ fn freeBufs(bs: []cuda.DeviceBuffer) void { for (bs) |*b| b.free(); } +fn freeOwned(gpa: std.mem.Allocator, bs: []cuda.DeviceBuffer) void { + freeBufs(bs); + gpa.free(bs); +} + test "linear layers skip every fourth index and attention starts at 3" { - try std.testing.expect(linearLayer(0)); - try std.testing.expect(!linearLayer(3)); - try std.testing.expect(!linearLayer(47)); - try std.testing.expectEqual(@as(usize, 0), linearIndex(0)); - try std.testing.expectEqual(@as(usize, 2), linearIndex(2)); - try std.testing.expectEqual(@as(usize, 35), linearIndex(46)); - try std.testing.expectEqual(@as(usize, 0), attnIndex(3)); - try std.testing.expectEqual(@as(usize, 11), attnIndex(47)); + try @import("cuda_prompt_test.zig").layerIndices(); } test "a sequence's device bytes grow by the twelve attention layers' keys, values, indexer keys and pooled blocks" { - const per_token = attn_n * (2 * kv_heads * head_dim * 2 + index_dim * 2); - try std.testing.expectEqual(per_token, seqBytes(6) - seqBytes(5)); - try std.testing.expectEqual(per_token + attn_n * index_dim * 2, seqBytes(5) - seqBytes(4)); - try std.testing.expect(seqBytes(max_tokens) > max_tokens * per_token); + try @import("cuda_prompt_test.zig").sequenceBytes(seqBytes); } test "a failed fill frees the buffer it was filling and every one before it" { - const Fake = struct { - live: *usize, - pub fn free(b: *@This()) void { - b.live.* -= 1; - } - }; - const Src = struct { - live: *usize, - fills: *usize, - fail_at: usize, - fn make(s: @This()) !Fake { - s.live.* += 1; - return .{ .live = s.live }; - } - fn fill(s: @This(), _: *Fake) !void { - s.fills.* += 1; - if (s.fills.* == s.fail_at) return error.FillFailed; - } - }; - for ([_]usize{ 1, 2, 4 }) |fail_at| { - var live: usize = 0; - var fills: usize = 0; - try std.testing.expectError(error.FillFailed, made(Fake, std.testing.allocator, 4, Src{ .live = &live, .fills = &fills, .fail_at = fail_at })); - try std.testing.expectEqual(@as(usize, 0), live); - } - var live: usize = 0; - var fills: usize = 0; - const all = try made(Fake, std.testing.allocator, 3, Src{ .live = &live, .fills = &fills, .fail_at = 0 }); - defer std.testing.allocator.free(all); - try std.testing.expectEqual(@as(usize, 3), live); + try @import("cuda_prompt_test.zig").failedFill(made); } /// TF_FLASHNEXT_CHUNK's rows a pass: unset, the default; otherwise 1 to batch.max_chunk, anything else refused. diff --git a/zig/src/families/flashnext/cuda_prompt_lifecycle.zig b/zig/src/families/flashnext/cuda_prompt_lifecycle.zig new file mode 100644 index 0000000000..1aa3e5e31e --- /dev/null +++ b/zig/src/families/flashnext/cuda_prompt_lifecycle.zig @@ -0,0 +1,14 @@ +const std = @import("std"); +const prompt = @import("cuda_prompt.zig"); + +pub fn deinit(sequence: *prompt.Seq, allocator: std.mem.Allocator) void { + if (sequence.stream) |stream| @import("cuda").completion.drain(stream.d, stream, "Flash Next sequence work did not drain before release"); + if (sequence.media) |*media| media.deinit(); + for ([_][]@import("cuda").DeviceBuffer{ sequence.conv, sequence.state, sequence.kc, sequence.vc, sequence.ikc, sequence.pooled }) |buffers| { + for (buffers) |*buffer| buffer.free(); + allocator.free(buffers); + } + if (sequence.tail.ptr != 0) sequence.tail.free(); + if (sequence.paged) |*table| table.deinit(allocator); + allocator.free(sequence.logits); +} diff --git a/zig/src/families/flashnext/cuda_prompt_load.zig b/zig/src/families/flashnext/cuda_prompt_load.zig index 45623e98cc..fc42f0a5df 100644 --- a/zig/src/families/flashnext/cuda_prompt_load.zig +++ b/zig/src/families/flashnext/cuda_prompt_load.zig @@ -7,6 +7,7 @@ const gdn = @import("cuda_gdn.zig"); const wt = @import("cuda_prompt_wt.zig"); const mixed = @import("cuda_mixed_experts.zig"); const projections = @import("cuda_projection_load.zig"); +const exl3_load = @import("exl3/pack/cuda_exl3_load.zig"); const layers_n = wt.layers_n; const dims = wt.dims; @@ -80,9 +81,12 @@ pub const Layer = struct { experts_up: cuda.DeviceBuffer, experts_down: cuda.DeviceBuffer, native_experts: ?mixed.Owned = null, + /// EXL3 routed experts with the shared one as id 512 (cuda_exl3_load.zig). + exl3: ?*exl3_load.Experts = null, pub fn free(self: *Layer, gpa: std.mem.Allocator) void { _ = gpa; + if (self.exl3) |x| x.destroy(); self.attn_hc.free(); self.mlp_hc.free(); self.proj.free(); @@ -343,7 +347,7 @@ fn loadLayerAtRank(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.S proj_q = try projections.loadStitched(gpa, driver, store, lp, &parts, gdn.proj_width, dims); out_q = try projections.load(gpa, driver, store, try std.fmt.bufPrint(&buf, "{s}.out_proj", .{lp}), dims, gdn.value_dim); const conv = try store.tensor(try std.fmt.bufPrint(&buf, "{s}.conv1d.weight", .{lp})); - if (conv.bytes.len != gdn.conv_dim * gdn.taps * 2) return error.UnexpectedTensor; + if (conv.dtype != .bf16 or conv.bytes.len != gdn.conv_dim * gdn.taps * 2) return error.UnexpectedTensor; cw_b = try cuda.DeviceBuffer.fromHost(driver, conv.bytes); const a_log = try asF32(gpa, try store.tensor(try std.fmt.bufPrint(&buf, "{s}.A_log", .{lp}))); defer gpa.free(a_log); @@ -384,6 +388,7 @@ fn loadLayerAtRank(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.S const pp = pre_buf[0..pre.len]; key_q = try projections.load(gpa, driver, store, try std.fmt.bufPrint(&buf, "{s}.key_proj", .{pp}), wide, dims); val_q = try projections.load(gpa, driver, store, try std.fmt.bufPrint(&buf, "{s}.value_proj", .{pp}), dims, dims); + try exl3_load.pleRows(gpa, driver, store, &key_q, &val_q); const nk = try store.normScale(gpa, try store.tensor(try std.fmt.bufPrint(&buf, "{s}.norm_key.weight", .{pp}))); defer gpa.free(nk); const nq = try store.normScale(gpa, try store.tensor(try std.fmt.bufPrint(&buf, "{s}.norm_query.weight", .{pp}))); @@ -395,7 +400,7 @@ fn loadLayerAtRank(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.S nq_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(nq)); nc_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(nc)); const conv = try store.tensor(try std.fmt.bufPrint(&buf, "{s}.conv1d.weight", .{pp})); - if (conv.bytes.len != wide * 4 * 2) return error.UnexpectedTensor; + if (conv.dtype != .bf16 or conv.bytes.len != wide * 4 * 2) return error.UnexpectedTensor; pcw_b = try cuda.DeviceBuffer.fromHost(driver, conv.bytes); } @@ -405,7 +410,11 @@ fn loadLayerAtRank(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.S errdefer experts_up.free(); var experts_down = deadBuf(driver); errdefer experts_down.free(); - if (store.has(try std.fmt.bufPrint(&buf, "{s}.experts.0.gate_proj.weight", .{mlp_root}))) { + var exl3_experts: ?*exl3_load.Experts = null; + errdefer if (exl3_experts) |x| x.destroy(); + if (exl3_load.hasExperts(store, mlp_root)) { + exl3_experts = try exl3_load.experts(gpa, driver, store, mlp_root, world, rank); + } else if (store.has(try std.fmt.bufPrint(&buf, "{s}.experts.0.gate_proj.weight", .{mlp_root}))) { native_experts = try mixed.load(gpa, driver, store, mlp_root, world, rank); } else { if (world != 1 or rank != 0) return error.UnsupportedQuantization; @@ -468,5 +477,6 @@ fn loadLayerAtRank(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.S .experts_up = experts_up, .experts_down = experts_down, .native_experts = native_experts, + .exl3 = exl3_experts, }; } diff --git a/zig/src/families/flashnext/cuda_prompt_telemetry.zig b/zig/src/families/flashnext/cuda_prompt_telemetry.zig new file mode 100644 index 0000000000..b346b9ee74 --- /dev/null +++ b/zig/src/families/flashnext/cuda_prompt_telemetry.zig @@ -0,0 +1,233 @@ +const std = @import("std"); +const cache = @import("exl3/model/cuda_exl3_telemetry.zig"); +const counters = cache.counters; + +pub const State = struct { + enabled: bool = false, + layers_enabled: bool = false, + layer_cursor: counters.LayerCursor(@import("cuda_prompt_wt.zig").layers_n) = .{}, + rank: counters.Rank = .{ .rank = 0, .world = 1 }, + forward_phase: counters.Phase = .unknown, + + pub fn fromEnv(rank: usize, world: usize) State { + const enabled = enabledFromEnv(); + const state: State = .{ .enabled = enabled, .layers_enabled = enabled and layersEnabledFromEnv(), .rank = .{ .rank = rank, .world = world } }; + if (state.enabled) std.log.warn("EXL3-CACHE-SAMPLING v=1 rank={d} world={d} optin_boundary_sync=1 bytes=logical_payload_not_pcie", .{ rank, world }); + return state; + } + + pub fn forwardPhase(state: State, fed: usize) counters.Phase { + return if (state.forward_phase != .unknown) state.forward_phase else if (fed == 0) .prefill else .unknown; + } + + pub fn enter(state: *State, phase: counters.Phase) Scope { + const previous = state.forward_phase; + state.forward_phase = phase; + return .{ .state = state, .previous = previous }; + } +}; + +pub const Scope = struct { + state: *State, + previous: counters.Phase, + + pub fn restore(scope: Scope) void { + scope.state.forward_phase = scope.previous; + } +}; + +pub fn Operations(comptime Engine: type, comptime Seq: type, comptime Hooks: type, comptime underlying: anytype) type { + return struct { + pub fn complete(engine: *Engine) !void { + if (engine.collective) |collective| return collective.finish(engine.stream.*); + try engine.stream.synchronize(); + } + + pub fn forward(engine: *Engine, comptime Tri: type, tri: Tri, sequence: *Seq, ids: []const u32, hooks: Hooks) !void { + if (!engine.telemetry.enabled) return underlying.forward(engine, Tri, tri, sequence, ids, hooks); + const phase = engine.telemetry.forwardPhase(sequence.fed); + const sequences = [_]*Seq{sequence}; + var before = try begin(engine, &sequences, phase, complete); + defer if (before) |*sample| sample.deinit(engine.gpa); + underlying.forward(engine, Tri, tri, sequence, ids, hooks) catch |failure| { + failed(engine, phase, failure); + return failure; + }; + if (before) |sample| try finish(engine, &sequences, sample, phase, ids.len, complete); + } + + pub fn round(engine: *Engine, comptime Tri: type, tri: Tri, sequences: []const *Seq, ids: []const u32, out: u64) !void { + if (!engine.telemetry.enabled) return underlying.round(engine, Tri, tri, sequences, ids, out); + var before = try begin(engine, sequences, .decode, complete); + defer if (before) |*sample| sample.deinit(engine.gpa); + underlying.round(engine, Tri, tri, sequences, ids, out) catch |failure| { + failed(engine, .decode, failure); + return failure; + }; + if (before) |sample| try finish(engine, sequences, sample, .decode, sequences.len, complete); + } + }; +} + +pub fn enabledFromEnv() bool { + const value = std.c.getenv("TF_FLASHNEXT_CACHE_TELEMETRY") orelse return false; + return std.mem.eql(u8, std.mem.span(value), "1"); +} + +pub fn layersEnabledFromEnv() bool { + const value = std.c.getenv("TF_FLASHNEXT_EXL3_LAYER_TELEMETRY") orelse return false; + return std.mem.eql(u8, std.mem.span(value), "1"); +} + +pub fn failed(engine: anytype, phase: counters.Phase, failure: anyerror) void { + std.log.warn("EXL3-CACHE-INVALID v=1 rank={d} world={d} phase={s} reason={s}", .{ engine.telemetry.rank.rank, engine.telemetry.rank.world, @tagName(phase), @errorName(failure) }); +} + +pub fn begin(engine: anytype, sequences: anytype, phase: counters.Phase, comptime complete: anytype) !?counters.Snapshot { + return beginWithSampler(engine, sequences, phase, complete, cache.sampleAfterConsumers); +} + +pub fn beginWithSampler(engine: anytype, sequences: anytype, phase: counters.Phase, comptime complete: anytype, comptime sample: anytype) !?counters.Snapshot { + if (!engine.telemetry.enabled) return null; + complete(engine) catch |failure| { + failed(engine, phase, failure); + return failure; + }; + return sample(engine.gpa, engine, sequences, engine.telemetry.rank, cache.Completion.all_rank_consumers_joined) catch |failure| { + failed(engine, phase, failure); + return null; + }; +} + +pub fn finish(engine: anytype, sequences: anytype, before: counters.Snapshot, phase: counters.Phase, rows: usize, comptime complete: anytype) !void { + return finishWithSampler(engine, sequences, before, phase, rows, complete, cache.sampleAfterConsumers, publish); +} + +pub fn finishWithSampler(engine: anytype, sequences: anytype, before: counters.Snapshot, phase: counters.Phase, rows: usize, comptime complete: anytype, comptime sample: anytype, comptime emit: anytype) !void { + complete(engine) catch |failure| { + failed(engine, phase, failure); + return failure; + }; + var after = sample(engine.gpa, engine, sequences, engine.telemetry.rank, cache.Completion.all_rank_consumers_joined) catch |failure| { + failed(engine, phase, failure); + return; + }; + defer after.deinit(engine.gpa); + var report = counters.Report.between(before, after, .{ .phase = phase, .model_rows = rows, .steps = 1 }) catch |failure| { + failed(engine, phase, failure); + return; + }; + const layers = if (engine.telemetry.layers_enabled) engine.telemetry.layer_cursor.claim(engine.gpa, before, after) catch |failure| { + failed(engine, phase, failure); + return; + } else &.{}; + defer engine.gpa.free(layers); + report.layers = layers; + emit(engine, report) catch |failure| { + failed(engine, phase, failure); + return; + }; +} + +fn publish(engine: anytype, report: counters.Report) !void { + const text = try report.formatAlloc(engine.gpa); + defer engine.gpa.free(text); + std.debug.print("{s}", .{text}); +} + +test "exl3: serving telemetry disabled adds no completion or readback and opt in samples only completed boundaries" { + const Fake = struct { + gpa: std.mem.Allocator, + telemetry: State = .{ .rank = .{ .rank = 0, .world = 1 } }, + completions: u64 = 0, + reads: u64 = 0, + records: u64 = 0, + + fn complete(engine: *@This()) !void { + engine.completions += 1; + } + + fn sample(_: std.mem.Allocator, engine: *@This(), _: @TypeOf(null), rank: counters.Rank, _: cache.Completion) !counters.Snapshot { + try std.testing.expect(engine.completions > engine.reads); + engine.reads += 1; + return .{ .rank = rank, .scope = 1, .experts = .{ .device_copies = engine.reads } }; + } + + fn emit(engine: *@This(), report: counters.Report) !void { + try std.testing.expectEqual(counters.Phase.decode, report.interval.phase); + try std.testing.expectEqual(@as(u64, 2), report.interval.model_rows); + try std.testing.expectEqual(@as(u64, 1), report.experts.device_copies); + try std.testing.expectEqual(@as(usize, 0), report.rank.rank); + engine.records += 1; + } + }; + var engine: Fake = .{ .gpa = std.testing.allocator }; + try std.testing.expect((try beginWithSampler(&engine, null, .decode, Fake.complete, Fake.sample)) == null); + try std.testing.expectEqual(@as(u64, 0), engine.completions); + try std.testing.expectEqual(@as(u64, 0), engine.reads); + engine.telemetry.enabled = true; + var first = (try beginWithSampler(&engine, null, .decode, Fake.complete, Fake.sample)).?; + defer first.deinit(engine.gpa); + try finishWithSampler(&engine, null, first, .decode, 2, Fake.complete, Fake.sample, Fake.emit); + try std.testing.expectEqual(@as(u64, 2), engine.completions); + try std.testing.expectEqual(@as(u64, 2), engine.reads); + try std.testing.expectEqual(@as(u64, 1), engine.records); +} + +test "exl3: forward and verify phase scopes restore nested caller state on error" { + const Call = struct { + fn fail(state: *State, phase: counters.Phase) !void { + const scope = state.enter(phase); + defer scope.restore(); + try std.testing.expectEqual(phase, state.forward_phase); + return error.ScopedOperationFailed; + } + }; + for ([_]counters.Phase{ .unknown, .prefill, .decode }) |previous| { + var state: State = .{ .forward_phase = previous }; + const parent = state.enter(.prefill); + try std.testing.expectError(error.ScopedOperationFailed, Call.fail(&state, .verify)); + try std.testing.expectEqual(counters.Phase.prefill, state.forward_phase); + try std.testing.expectError(error.ScopedOperationFailed, Call.fail(&state, .prefill)); + try std.testing.expectEqual(counters.Phase.prefill, state.forward_phase); + parent.restore(); + try std.testing.expectEqual(previous, state.forward_phase); + try std.testing.expect(!state.enabled); + } +} + +test "exl3: extracted disabled forward and round wrappers call underlying operations without joins" { + const prompt = @import("cuda_prompt.zig"); + const cuda = @import("cuda"); + const Fake = @import("exl3/model/cuda_exl3_fake_driver_test.zig").Fake; + const Calls = struct { + var forwards: usize = 0; + var rounds: usize = 0; + + fn forward(_: *prompt.Engine, comptime Tri: type, _: Tri, _: *prompt.Seq, _: []const u32, _: prompt.batch.Hooks) !void { + forwards += 1; + } + + fn round(_: *prompt.Engine, comptime Tri: type, _: Tri, _: []const *prompt.Seq, _: []const u32, _: u64) !void { + rounds += 1; + } + }; + Calls.forwards = 0; + Calls.rounds = 0; + var driver = Fake.driver(); + var stream: cuda.Stream = .{ .d = &driver, .handle = null }; + var engine: prompt.Engine = undefined; + engine.telemetry = .{}; + engine.gpa = std.testing.allocator; + engine.collective = null; + engine.stream = &stream; + var sequence: prompt.Seq = undefined; + sequence.fed = 9000; + const ops = Operations(prompt.Engine, prompt.Seq, prompt.batch.Hooks, .{ .forward = Calls.forward, .round = Calls.round }); + const syncs = Fake.sync_calls; + try ops.forward(&engine, void, {}, &sequence, &.{123}, .{}); + try ops.round(&engine, void, {}, &.{&sequence}, &.{123}, 0); + try std.testing.expectEqual(@as(usize, 1), Calls.forwards); + try std.testing.expectEqual(@as(usize, 1), Calls.rounds); + try std.testing.expectEqual(syncs, Fake.sync_calls); +} diff --git a/zig/src/families/flashnext/cuda_prompt_test.zig b/zig/src/families/flashnext/cuda_prompt_test.zig new file mode 100644 index 0000000000..ec581b9f4e --- /dev/null +++ b/zig/src/families/flashnext/cuda_prompt_test.zig @@ -0,0 +1,53 @@ +const std = @import("std"); +const weights = @import("cuda_prompt_wt.zig"); + +pub fn failedFill(comptime makeBuffers: anytype) !void { + const Fake = struct { + live: *usize, + pub fn free(buffer: *@This()) void { + buffer.live.* -= 1; + } + }; + const Src = struct { + live: *usize, + fills: *usize, + fail_at: usize, + pub fn make(source: @This()) !Fake { + source.live.* += 1; + return .{ .live = source.live }; + } + pub fn fill(source: @This(), _: *Fake) !void { + source.fills.* += 1; + if (source.fills.* == source.fail_at) return error.FillFailed; + } + }; + for ([_]usize{ 1, 2, 4 }) |fail_at| { + var live: usize = 0; + var fills: usize = 0; + try std.testing.expectError(error.FillFailed, makeBuffers(Fake, std.testing.allocator, 4, Src{ .live = &live, .fills = &fills, .fail_at = fail_at })); + try std.testing.expectEqual(@as(usize, 0), live); + } + var live: usize = 0; + var fills: usize = 0; + const all = try makeBuffers(Fake, std.testing.allocator, 3, Src{ .live = &live, .fills = &fills, .fail_at = 0 }); + defer std.testing.allocator.free(all); + try std.testing.expectEqual(@as(usize, 3), live); +} + +pub fn layerIndices() !void { + try std.testing.expect(weights.linearLayer(0)); + try std.testing.expect(!weights.linearLayer(3)); + try std.testing.expect(!weights.linearLayer(47)); + try std.testing.expectEqual(@as(usize, 0), weights.linearIndex(0)); + try std.testing.expectEqual(@as(usize, 2), weights.linearIndex(2)); + try std.testing.expectEqual(@as(usize, 35), weights.linearIndex(46)); + try std.testing.expectEqual(@as(usize, 0), weights.attnIndex(3)); + try std.testing.expectEqual(@as(usize, 11), weights.attnIndex(47)); +} + +pub fn sequenceBytes(comptime sequenceSize: anytype) !void { + const per_token = weights.attn_n * (2 * weights.kv_heads * weights.head_dim * 2 + weights.index_dim * 2); + try std.testing.expectEqual(per_token, sequenceSize(6) - sequenceSize(5)); + try std.testing.expectEqual(per_token + weights.attn_n * weights.index_dim * 2, sequenceSize(5) - sequenceSize(4)); + try std.testing.expect(sequenceSize(weights.window) > weights.window * per_token); +} diff --git a/zig/src/families/flashnext/cuda_prompt_wt.zig b/zig/src/families/flashnext/cuda_prompt_wt.zig index 732e0ce77e..c351aa42c2 100644 --- a/zig/src/families/flashnext/cuda_prompt_wt.zig +++ b/zig/src/families/flashnext/cuda_prompt_wt.zig @@ -6,6 +6,7 @@ const core = @import("core"); const hc = @import("cuda_hc.zig"); const qmm = @import("cuda_qmm.zig"); const gdn = @import("cuda_gdn.zig"); +const exl3_store = @import("exl3/pack/cuda_exl3_store.zig"); pub const layers_n: usize = 48; pub const linear_n: usize = 36; @@ -101,14 +102,19 @@ pub const Store = struct { released: [2]Held = undefined, nreleased: usize = 0, norm_offset: f32 = 0, + /// Dense F16 weights an EXL3 pack stores, widened to BF16 exactly (cuda_exl3_store.zig), freed at release(). + widened: std.ArrayList([]u16) = .empty, + /// The EXL3 kernels and scratch this checkpoint's projections and experts share (cuda_exl3_load.zig), if any. + exl3_ops: ?*@import("exl3/pack/cuda_exl3_load.zig").Ops = null, pub fn open(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) !Store { const index = try std.fs.path.join(gpa, &.{ dir, "model.safetensors.index.json" }); defer gpa.free(index); const text = try std.Io.Dir.cwd().readFileAlloc(io, index, gpa, .limited(1 << 26)); defer gpa.free(text); - const parsed = try std.json.parseFromSlice(std.json.Value, gpa, text, .{}); + var parsed = try std.json.parseFromSlice(std.json.Value, gpa, text, .{}); errdefer parsed.deinit(); + try exl3_store.addPatch(io, dir, &parsed); _ = try core.safetensors.index.weightMap(parsed.value); return .{ .gpa = gpa, .io = io, .dir = dir, .parsed = parsed }; } @@ -144,6 +150,7 @@ pub const Store = struct { pub fn deinit(self: *Store) void { self.dropStaging(); + if (self.exl3_ops) |o| o.release(); self.release(); self.held.deinit(self.gpa); self.closeReleased(); @@ -174,6 +181,8 @@ pub const Store = struct { /// Borrowers expire here; two indexed shards cover separate weight and calibration planes. pub fn release(self: *Store) void { + for (self.widened.items) |w| self.gpa.free(w); + self.widened.clearAndFree(self.gpa); for (self.held.items) |h| { if (self.nreleased == self.released.len) { self.released[0].file.close(); @@ -207,7 +216,7 @@ pub const Store = struct { const key = try self.resolveKey(requested, &mapped); const file = try self.shard(key); for (self.held.items) |*h| { - if (std.mem.eql(u8, h.rel, file)) return h.file.get(key) orelse return error.MissingTensor; + if (std.mem.eql(u8, h.rel, file)) return self.dense(key, h.file.get(key) orelse return error.MissingTensor); } try self.held.ensureUnusedCapacity(self.gpa, 1); for (self.released[0..self.nreleased], 0..) |cached, i| { @@ -215,7 +224,7 @@ pub const Store = struct { self.held.appendAssumeCapacity(cached); self.nreleased -= 1; if (i < self.nreleased) self.released[i] = self.released[self.nreleased]; - return self.held.items[self.held.items.len - 1].file.get(key) orelse return error.MissingTensor; + return self.dense(key, self.held.items[self.held.items.len - 1].file.get(key) orelse return error.MissingTensor); } } const rel = try self.gpa.dupe(u8, file); @@ -223,10 +232,10 @@ pub const Store = struct { const path = try std.fs.path.join(self.gpa, &.{ self.dir, file }); defer self.gpa.free(path); self.held.appendAssumeCapacity(.{ .rel = rel, .file = try Shard.load(self.gpa, self.io, path) }); - return self.held.items[self.held.items.len - 1].file.get(key) orelse { + return self.dense(key, self.held.items[self.held.items.len - 1].file.get(key) orelse { std.debug.print("absent {s}\n", .{key}); return error.MissingTensor; - }; + }); } pub fn closeFile(self: *Store) void { @@ -277,10 +286,16 @@ pub const Store = struct { var mapped: [512]u8 = undefined; const key = try self.resolveKey(requested, &mapped); try self.use(key); - return self.file.?.get(key) orelse { + return self.dense(key, self.file.?.get(key) orelse { std.debug.print("absent {s}\n", .{key}); return error.MissingTensor; - }; + }); + } + + /// An EXL3 pack's dense F16 weight as the BF16 the loaders read (exact or refused); any other tensor as stored. + fn dense(self: *Store, key: []const u8, t: Tensor) !Tensor { + if (!exl3_store.widens(key, t)) return t; + return exl3_store.widen(self.gpa, key, t, &self.widened); } /// Check the exact or compatibility name without mapping its shard. diff --git a/zig/src/families/flashnext/cuda_scratch_layout_test.zig b/zig/src/families/flashnext/cuda_scratch_layout_test.zig new file mode 100644 index 0000000000..589a972be2 --- /dev/null +++ b/zig/src/families/flashnext/cuda_scratch_layout_test.zig @@ -0,0 +1,41 @@ +const std = @import("std"); +const Scratch = @import("cuda_batch_state.zig").Scratch; + +test "exl3: scratch omits only classic MoE allocations and the split-head logits and keeps every common view" { + for ([_]usize{ 1, 8, 32, 64, 1024, 1536, 2048 }) |cap| { + const classic = Scratch.allocationSizes(cap, .classic); + const exl3 = Scratch.allocationSizes(cap, .exl3); + var removed: usize = 0; + var retained: usize = 0; + for (classic, exl3, 0..) |before, after, index| { + if ((index >= 38 and index <= 47) or index == 61) { + try std.testing.expectEqual(@as(?usize, null), after); + removed += @max(before.?, 256); + } else { + try std.testing.expectEqual(before, after); + retained += @max(after.?, 256); + } + } + try std.testing.expectEqual(retained, Scratch.bytesForFormat(cap, .exl3)); + try std.testing.expectEqual(removed, Scratch.bytes(cap) - retained); + try std.testing.expectEqual(@as(?usize, cap * 11 * 2560 * 4), exl3[48]); + try std.testing.expectEqual(@as(?usize, cap * 2560 * 4), exl3[50]); + try std.testing.expectEqual(@as(?usize, cap * 128 * 4), exl3[51]); + } +} + +test "exl3: scratch savings include the allocator minimum and grow with requested rows" { + var previous: usize = 0; + for ([_]usize{ 1, 8, 32, 64, 1024, 1536, 2048 }) |cap| { + const pairs = cap * 11; + const items = (pairs + 15) / 16 + 513; + const blocks = (pairs + 1023) / 1024; + const expected = 2 * @max(pairs * 4, 256) + @max(items * 12, 256) + 256 + + @max(blocks * 513 * 4, 256) + 3 * @max(cap * 640 * 2, 256) + + @max(cap * 2560 * 4, 256) + @max(pairs * 640 * 2, 256) + @max(cap * 248320, 256); + const saved = Scratch.bytes(cap) - Scratch.bytesForFormat(cap, .exl3); + try std.testing.expectEqual(expected, saved); + try std.testing.expect(saved > previous); + previous = saved; + } +} diff --git a/zig/src/families/flashnext/cuda_snap.zig b/zig/src/families/flashnext/cuda_snap.zig index 7db8fe4043..c8d03ccb1a 100644 --- a/zig/src/families/flashnext/cuda_snap.zig +++ b/zig/src/families/flashnext/cuda_snap.zig @@ -30,6 +30,8 @@ pub fn bytes(at: usize, head: bool) usize { pub const Snap = struct { rank_id: u64 = 0, buf: cuda.DeviceBuffer, + /// The pinned host memory `buf` borrows (states parked on the host: `on_host`), or null (device memory). + host: ?cuda.HostBuffer = null, at: usize, history: [2]i64, head: bool, @@ -102,20 +104,41 @@ fn move(stream: cuda.Stream, seq: *prompt.Seq, d: ?*Draft, base: u64, at: usize, std.debug.assert(off == bytes(at, d != null)); } +fn discard(gpa: std.mem.Allocator, stream: cuda.Stream, snap: *Snap) void { + cuda.completion.drain(stream.d, stream, "snapshot copies did not drain"); + if (snap.host) |*h| h.free() else snap.buf.free(); + gpa.destroy(snap); +} + +/// Kept states go to mapped pinned host memory, set once at load (TF_FLASHNEXT_SNAP_HOST); the host budgets the cache. +pub var on_host: bool = false; + /// A copy of `seq` (and of the head's cache, with `d`) as it stands after `seq.fed` tokens. pub fn save(gpa: std.mem.Allocator, driver: *const cuda.Driver, stream: cuda.Stream, seq: *prompt.Seq, d: ?*Draft) !*Snap { const at = seq.fed; - if (d) |h| if (h.len != at or h.drafted != 0) return error.NotAtMark; - const snap = try gpa.create(Snap); - errdefer gpa.destroy(snap); - snap.* = .{ .buf = try cuda.DeviceBuffer.alloc(driver, bytes(at, d != null)), .at = at, .history = seq.history, .head = d != null }; - errdefer snap.buf.free(); - errdefer cuda.completion.drain(driver, stream, "a failed snapshot save did not drain"); + if (at > seq.cap) return error.SequenceTooShort; + if (d) |h| if (h.len != at or h.drafted != 0 or h.open or at > h.cap) return error.NotAtMark; + const n = bytes(at, d != null); + const snap = try allocate(gpa, driver, n, at, seq.history, d != null); + errdefer discard(gpa, stream, snap); try move(stream, seq, d, snap.buf.ptr, at, .save); try stream.synchronize(); return snap; } +fn allocate(gpa: std.mem.Allocator, driver: *const cuda.Driver, n: usize, at: usize, history: [2]i64, head: bool) !*Snap { + const snap = try gpa.create(Snap); + errdefer gpa.destroy(snap); + snap.* = .{ .buf = undefined, .at = at, .history = history, .head = head }; + if (on_host) { + var host = try cuda.HostBuffer.allocMapped(driver, @max(n, 256)); + errdefer host.free(); + snap.buf = .{ .d = driver, .ptr = try host.device(), .len = n, .borrowed = true }; + snap.host = host; + } else snap.buf = try cuda.DeviceBuffer.alloc(driver, n); + return snap; +} + /// `seq` (and the head's cache, with `d`) becomes `snap`'s state; refused before any copy when it cannot hold it. pub fn restore(stream: cuda.Stream, seq: *prompt.Seq, d: ?*Draft, snap: *const Snap) !void { if (snap.at > seq.cap) return error.SequenceTooShort; @@ -128,13 +151,13 @@ pub fn restore(stream: cuda.Stream, seq: *prompt.Seq, d: ?*Draft, snap: *const S if (d) |h| { h.reset(); h.len = snap.at; + h.open = false; + h.fed0 = snap.at; } } pub fn drop(gpa: std.mem.Allocator, stream: cuda.Stream, snap: *Snap) void { - cuda.completion.drain(stream.d, stream, "a snapshot drop did not drain"); - snap.buf.free(); - gpa.destroy(snap); + discard(gpa, stream, snap); } test "a state's bytes: about 110 MiB of DeltaNet states and the n-gram tail, then 27 KiB a token" { diff --git a/zig/src/families/flashnext/exl3/kv/cuda_exl3_kv.zig b/zig/src/families/flashnext/exl3/kv/cuda_exl3_kv.zig new file mode 100644 index 0000000000..c87faba205 --- /dev/null +++ b/zig/src/families/flashnext/exl3/kv/cuda_exl3_kv.zig @@ -0,0 +1,83 @@ +//! KV memory of an EXL3 pack's streams (K10): bytes by piece, the scratch format, where kept states live. + +const std = @import("std"); +const wt = @import("../../cuda_prompt_wt.zig"); +const gdn = @import("../../cuda_gdn.zig"); + +const head_dim = wt.head_dim; +const kv_heads = wt.kv_heads; +const index_row: usize = wt.index_dim * 2; + +/// One token's K and V in one attention layer, bf16. +pub fn kvRowBytes() usize { + return 2 * kv_heads * head_dim * 2; +} + +/// Bytes a sequence of `cap` tokens holds, by piece. +pub const SeqBytes = struct { + gdn: u64, + kv: u64, + ikc: u64, + pooled: u64, + tail: u64, + + pub fn device(s: SeqBytes) u64 { + return s.gdn + s.kv + s.ikc + s.pooled + s.tail; + } +}; + +pub fn pooledBytes(cap: usize) u64 { + return (cap + wt.index_ratio - 1) / wt.index_ratio * index_row; +} + +/// cuda_prompt.zig seqBytes split by piece (a host test checks that the pieces sum to upstream's number). +pub fn seqBytes(cap: usize) SeqBytes { + const attn: u64 = wt.attn_n; + return .{ + .gdn = wt.linear_n * ((gdn.taps - 1) * gdn.conv_dim * 2 + wt.state_n * 4), + .kv = attn * cap * kvRowBytes(), + .ikc = attn * cap * index_row, + .pooled = attn * pooledBytes(cap), + .tail = wt.ple_tail_n * wt.wide * 2, + }; +} + +pub const gib: f64 = 1024.0 * 1024.0 * 1024.0; + +/// K10 lever 3: the pass's chunk for an EXL3 checkpoint unless TF_FLASHNEXT_CHUNK says. +pub const exl3_chunk: usize = 1024; + +/// The server's load hook (cuda_lanes_load.zig, right after the engine loads, before any scratch or sequence). +pub fn setup(e: anytype, exl3: bool) !void { + if (!exl3) return; + if (e.batch_s != null) return error.ScratchAlreadyReserved; + const decoder_exl3 = e.n_layers != 0 and for (e.layers[0..e.n_layers]) |layer| { + if (layer.exl3 == null) break false; + } else true; + e.scratch_format = if (decoder_exl3) .exl3 else .classic; + if (std.c.getenv("TF_FLASHNEXT_CHUNK") == null) e.chunk = exl3_chunk; +} + +/// One log line: a stream's device bytes at the window by piece. +pub fn logStream(e: anytype, window: usize, stream_bytes: usize) !void { + const b = seqBytes(window); + std.log.info("flashnext: a stream at a {d}-token window holds {d:.3} GiB on the device (K/V {d:.3} bf16, indexer keys {d:.3}, pooled {d:.3}, DeltaNet {d:.3}, draft and the rest {d:.3}); chunk {d} rows, scratch {d:.3} GiB", .{ + window, + @as(f64, @floatFromInt(stream_bytes)) / gib, + @as(f64, @floatFromInt(b.kv)) / gib, + @as(f64, @floatFromInt(b.ikc)) / gib, + @as(f64, @floatFromInt(b.pooled)) / gib, + @as(f64, @floatFromInt(b.gdn)) / gib, + @as(f64, @floatFromInt(stream_bytes -| b.device())) / gib, + e.chunk, + @as(f64, @floatFromInt(@import("../../cuda_batch_state.zig").Scratch.bytesForFormat(e.chunk, e.scratch_format))) / gib, + }); +} + +/// Whether kept prompt states go to pinned host memory (cuda_snap.on_host, set here): +pub fn snapHost(context: usize, exl3: bool) bool { + var on = exl3 and context > 32768; + if (std.c.getenv("TF_FLASHNEXT_SNAP_HOST")) |v| on = std.mem.eql(u8, std.mem.span(v), "1"); + @import("../../cuda_snap.zig").on_host = on; + return on; +} diff --git a/zig/src/families/flashnext/exl3/kv/cuda_exl3_kv_test.zig b/zig/src/families/flashnext/exl3/kv/cuda_exl3_kv_test.zig new file mode 100644 index 0000000000..95e577fd59 --- /dev/null +++ b/zig/src/families/flashnext/exl3/kv/cuda_exl3_kv_test.zig @@ -0,0 +1,20 @@ +//! Host tests of the KV accounting: the pieces of a sequence's bytes. + +const std = @import("std"); +const hub = @import("cuda_exl3_kv.zig"); + +const seqBytes = hub.seqBytes; + +test "exl3: kv the pieces of a sequence sum to upstream's seqBytes" { + const prompt = @import("../../cuda_prompt.zig"); + for ([_]usize{ 1, 5, 4096, 200_016, 262_144 }) |cap| { + const s = seqBytes(cap); + try std.testing.expectEqual(@as(u64, prompt.seqBytes(cap)), s.device()); + } +} + +test "exl3: kv a token's K and V are upstream's row twice" { + const prompt = @import("../../cuda_prompt.zig"); + try std.testing.expectEqual(@as(usize, 2048), hub.kvRowBytes()); + try std.testing.expectEqual(2 * prompt.kv_row_bytes, hub.kvRowBytes()); +} diff --git a/zig/src/families/flashnext/exl3/kv/exl3_kv.zig b/zig/src/families/flashnext/exl3/kv/exl3_kv.zig new file mode 100644 index 0000000000..d53973f8bb --- /dev/null +++ b/zig/src/families/flashnext/exl3/kv/exl3_kv.zig @@ -0,0 +1,98 @@ +//! The host oracle of the device K/V table: one layer's slots as the kernel addresses them. + +const std = @import("std"); +const shape = @import("../pack/exl3_shape.zig"); +const cache = @import("../model/exl3_cache.zig"); + +/// Full-attention layers only. GDN layers have no KV table. The size is `contextKv`. +const full_layers: usize = shape.layers / shape.full_attention_interval; + +pub const Table = struct { + heads: usize, + tokens: usize, + bytes: []u8, + + pub fn init(gpa: std.mem.Allocator, tokens: usize) error{OutOfMemory}!Table { + const n = cache.contextKv(tokens); + const bytes = try gpa.alloc(u8, @intCast(n)); + @memset(bytes, 0); + return .{ .heads = shape.kv_heads, .tokens = tokens, .bytes = bytes }; + } + + pub fn deinit(self: *Table, gpa: std.mem.Allocator) void { + gpa.free(self.bytes); + } + + /// One head, one token, one full-attention layer. + pub fn write(self: *Table, head: usize, layer: usize, token: usize, src: []const u8) error{ Head, Slot, Length }!void { + const off = try self.locate(head, layer, token, src.len); + @memcpy(self.bytes[off..][0..src.len], src); + } + + pub fn read(self: *const Table, head: usize, layer: usize, token: usize, dest: []u8) error{ Head, Slot, Length }!void { + const off = try self.locate(head, layer, token, dest.len); + @memcpy(dest, self.bytes[off..][0..dest.len]); + } + + fn locate(self: *const Table, head: usize, layer: usize, token: usize, len: usize) error{ Head, Slot, Length }!usize { + if (head >= self.heads) return error.Head; + if (self.tokens == 0 or full_layers == 0 or layer >= full_layers or token >= self.tokens) return error.Slot; + const denom = full_layers * self.tokens * self.heads; + if (self.bytes.len % denom != 0) return error.Slot; + const one = self.bytes.len / denom; + if (len != one) return error.Length; + return ((layer * self.tokens + token) * self.heads + head) * one; + } +}; + +test "the table is contextKv bytes, both heads" { + const gpa = std.testing.allocator; + const tokens: usize = 8192; + var t = try Table.init(gpa, tokens); + defer t.deinit(gpa); + try std.testing.expectEqual(shape.kv_heads, t.heads); + try std.testing.expectEqual(@as(u64, 201_326_592), @as(u64, t.bytes.len)); + try std.testing.expectEqual(cache.contextKv(tokens), @as(u64, t.bytes.len)); + var none = try Table.init(gpa, 0); + defer none.deinit(gpa); + try std.testing.expectEqual(@as(usize, 0), none.bytes.len); +} + +test "a write lands in its own head's slot and nowhere else" { + const gpa = std.testing.allocator; + const tokens: usize = 2; + const slot_bytes: usize = 1024; + const mark: u8 = 0xa5; + var src: [slot_bytes]u8 = undefined; + var got: [slot_bytes]u8 = undefined; + { + var t = try Table.init(gpa, tokens); + defer t.deinit(gpa); + const slot = src[0..slot_bytes]; + @memset(slot, mark); + try std.testing.expectError(error.Head, t.write(shape.kv_heads, 0, 0, slot)); + try std.testing.expectError(error.Length, t.write(0, 0, 0, slot[0 .. slot_bytes - 1])); + try std.testing.expectError(error.Slot, t.write(0, 0, tokens, slot)); + try std.testing.expectError(error.Slot, t.write(0, full_layers, 0, slot)); + try std.testing.expect(std.mem.allEqual(u8, t.bytes, 0)); + + try t.write(0, 0, 1, slot); + @memset(&got, 0); + try t.read(0, 0, 1, got[0..slot_bytes]); + try std.testing.expectEqualSlices(u8, slot, got[0..slot_bytes]); + try t.read(1, 0, 1, got[0..slot_bytes]); + try std.testing.expect(std.mem.allEqual(u8, got[0..slot_bytes], 0)); + + @memset(slot, mark +% 1); + try t.write(1, 1, 0, slot); + try t.read(1, 1, 0, got[0..slot_bytes]); + try std.testing.expectEqualSlices(u8, slot, got[0..slot_bytes]); + try t.read(0, 0, 1, got[0..slot_bytes]); + try std.testing.expect(got[0] == mark); + } + var none = try Table.init(gpa, 0); + defer none.deinit(gpa); + var scratch: [0]u8 = .{}; + try std.testing.expectError(error.Slot, none.write(0, 0, 0, &scratch)); + try std.testing.expectError(error.Slot, none.read(0, 0, 0, &scratch)); +} diff --git a/zig/src/families/flashnext/exl3/model/cuda_exl3.zig b/zig/src/families/flashnext/exl3/model/cuda_exl3.zig new file mode 100644 index 0000000000..5b681cdf05 --- /dev/null +++ b/zig/src/families/flashnext/exl3/model/cuda_exl3.zig @@ -0,0 +1,31 @@ +//! EXL3 checkpoints on the Flash Next CUDA family: what the upstream files' one-line hooks call. +pub const config = @import("../pack/cuda_exl3_config.zig"); +pub const arch = @import("cuda_exl3_arch.zig"); +pub const store = @import("../pack/cuda_exl3_store.zig"); +pub const memory = @import("../pack/cuda_exl3_memory.zig"); +pub const load = @import("../pack/cuda_exl3_load.zig"); +pub const pass = @import("cuda_exl3_pass.zig"); +pub const kv = @import("../kv/cuda_exl3_kv.zig"); +pub const fork = @import("../tp/cuda_exl3_fork.zig"); +pub const telemetry = @import("cuda_exl3_telemetry.zig"); +pub const frequency = @import("../offload/cuda_exl3_frequency.zig"); +/// The upstream modules an EXL3 check drives (tf-cuda-test flashnext-exl3): the store, loaders and the batched pass. +pub const prompt_wt = @import("../../cuda_prompt_wt.zig"); +pub const prompt_load = @import("../../cuda_prompt_load.zig"); +pub const projection_load = @import("../../cuda_projection_load.zig"); +pub const batch = @import("../../cuda_batch.zig"); + +test { + _ = config; + _ = arch; + _ = store; + _ = memory; + _ = load; + _ = pass; + _ = kv; + _ = @import("../kv/cuda_exl3_kv_test.zig"); + _ = fork; + _ = telemetry; + _ = frequency; + _ = @import("../offload/cuda_exl3_frequency_state.zig"); +} diff --git a/zig/src/families/flashnext/exl3/model/cuda_exl3_arch.zig b/zig/src/families/flashnext/exl3/model/cuda_exl3_arch.zig new file mode 100644 index 0000000000..fee7e05a47 --- /dev/null +++ b/zig/src/families/flashnext/exl3/model/cuda_exl3_arch.zig @@ -0,0 +1,32 @@ +//! EXL3 Flash Next runs on sm_120 only; other compute capabilities are refused until their gates have been run on them. +const std = @import("std"); +const cuda = @import("cuda"); + +/// The one compute capability (10 * major + minor) this family's EXL3 path is qualified on: Blackwell sm_120. +pub const qualified: u32 = 120; + +pub const message = "Flash Next EXL3 (expert offload, one rank) is qualified on sm_120 only; this device is sm_{d}: " ++ + "serve an NVFP4 or BF16 Flash Next checkpoint, or set TF_FLASHNEXT_EXL3_ARCH=unqualified to try it without a claim"; + +/// Whether a device of `capability` may serve an EXL3 pack; `override` is the TF_FLASHNEXT_EXL3_ARCH value for testers. +pub fn admits(capability: u32, override: ?[]const u8) bool { + if (capability == qualified) return true; + const v = override orelse return false; + return std.mem.eql(u8, v, "unqualified"); +} + +/// Refuses an EXL3 pack on an unqualified device before any weight loads. +pub fn check(driver: *const cuda.Driver) !void { + const capability = try driver.capability(0); + const v = std.c.getenv("TF_FLASHNEXT_EXL3_ARCH"); + if (admits(capability, if (v) |p| std.mem.span(p) else null)) return; + std.log.err(message, .{capability}); + return error.Exl3ArchUnqualified; +} + +test "only sm_120 admits an EXL3 pack unless a tester overrides the gate" { + try std.testing.expect(admits(120, null)); + for ([_]u32{ 80, 86, 89, 90, 100, 121 }) |c| try std.testing.expect(!admits(c, null)); + try std.testing.expect(!admits(121, "1")); + try std.testing.expect(admits(121, "unqualified")); +} diff --git a/zig/src/families/flashnext/exl3/model/cuda_exl3_dispatch_test.zig b/zig/src/families/flashnext/exl3/model/cuda_exl3_dispatch_test.zig new file mode 100644 index 0000000000..7c0253f1a3 --- /dev/null +++ b/zig/src/families/flashnext/exl3/model/cuda_exl3_dispatch_test.zig @@ -0,0 +1,74 @@ +const std = @import("std"); +const cuda = @import("cuda"); +const batch = @import("../../cuda_batch.zig"); +const load = @import("../../cuda_prompt_load.zig"); +const exl3_load = @import("../pack/cuda_exl3_load.zig"); + +const FakeDriver = struct { + var copies: usize = 0; + + fn copy(_: u64, _: u64, _: usize, _: cuda.abi.Stream) callconv(.c) cuda.abi.Result { + copies += 1; + return 1; + } + + fn errorName(_: cuda.abi.Result, out: *?[*:0]const u8) callconv(.c) cuda.abi.Result { + out.* = null; + return cuda.abi.success; + } + + fn errorString(_: cuda.abi.Result, out: *?[*:0]const u8) callconv(.c) cuda.abi.Result { + out.* = null; + return cuda.abi.success; + } + + fn make() cuda.Driver { + var driver: cuda.Driver = undefined; + driver.api.cuMemcpyDtoDAsync_v2 = copy; + driver.api.cuGetErrorName = errorName; + driver.api.cuGetErrorString = errorString; + return driver; + } +}; + +const MockTri = struct { + own: ?*const @import("../../cuda_glue.zig").Glue = null, + + pub fn routerMany(_: MockTri, _: u64, _: u64, _: u64, _: usize, _: usize) !void {} + pub fn topkMany(_: MockTri, _: u64, _: u64, _: u64, _: usize) !void {} +}; + +test "exl3: Keep dispatch rejects an unattached expert bank before legacy copy" { + FakeDriver.copies = 0; + var driver = FakeDriver.make(); + var host_bytes: [0]u8 align(16) = .{}; + var experts: exl3_load.Experts = .{ + .gpa = std.testing.allocator, + .ops = undefined, + .layer = undefined, + .bank = .{ .d = &driver, .bytes = &host_bytes }, + }; + var layer: load.Layer = undefined; + layer.router.ptr = 1; + layer.exl3 = &experts; + var scratch: batch.Scratch = undefined; + scratch.mixed = 2; + scratch.router_l = 3; + scratch.picks = 4; + scratch.wts = 5; + var pass: batch.Pass(MockTri) = .{ + .gpa = std.testing.allocator, + .io = undefined, + .stream = .{ .d = &driver, .handle = null }, + .tri = .{}, + .k = undefined, + .s = &scratch, + .m = undefined, + .seq = undefined, + .hooks = .{}, + .rows = 1, + }; + + try std.testing.expectError(error.Exl3PoolNotAttached, pass.moeBlockKeep(&layer, .{ .moe_x = 9 })); + try std.testing.expectEqual(@as(usize, 0), FakeDriver.copies); +} diff --git a/zig/src/families/flashnext/exl3/model/cuda_exl3_fake_driver_test.zig b/zig/src/families/flashnext/exl3/model/cuda_exl3_fake_driver_test.zig new file mode 100644 index 0000000000..f8d01f1dfc --- /dev/null +++ b/zig/src/families/flashnext/exl3/model/cuda_exl3_fake_driver_test.zig @@ -0,0 +1,23 @@ +const cuda = @import("cuda"); + +/// A driver whose only live entry points are the two synchronizations, which count themselves. +pub const Fake = struct { + pub var sync_calls: usize = 0; + + fn sync() callconv(.c) c_int { + sync_calls += 1; + return 0; + } + + fn streamSync(_: cuda.abi.Stream) callconv(.c) c_int { + return sync(); + } + + pub fn driver() cuda.Driver { + sync_calls = 0; + var result: cuda.Driver = undefined; + result.api.cuCtxSynchronize = sync; + result.api.cuStreamSynchronize = streamSync; + return result; + } +}; diff --git a/zig/src/families/flashnext/exl3/model/cuda_exl3_pass.zig b/zig/src/families/flashnext/exl3/model/cuda_exl3_pass.zig new file mode 100644 index 0000000000..71de32b822 --- /dev/null +++ b/zig/src/families/flashnext/exl3/model/cuda_exl3_pass.zig @@ -0,0 +1,374 @@ +//! The batched pass's EXL3 branches (cuda_batch.zig): + +const std = @import("std"); +const cuda = @import("cuda"); +const wt = @import("../../cuda_prompt_wt.zig"); +const load = @import("../pack/cuda_exl3_load.zig"); +const pool_ = @import("../offload/cuda_exl3_pool.zig"); +const memory = @import("../pack/exl3_startup_memory.zig"); +const line_deltas = @import("../offload/exl3_line_deltas.zig"); +const profile_line_budget = @import("../offload/exl3_profile_line_budget.zig"); + +const dims = wt.dims; +const slots = wt.slots; + +/// Rows a routed-experts call takes: its scratch grows with rows * 11 (fp16 inputs, fp32 split sums). +pub const moe_rows: usize = 256; + +fn copy2d(s: cuda.Stream, dst: u64, dpitch: usize, src: u64, spitch: usize, width: usize, height: usize) !void { + const m: cuda.abi.Memcpy2D = .{ .src = src, .src_pitch = spitch, .dst = dst, .dst_pitch = dpitch, .width = width, .height = height }; + try s.d.check(s.d.api.cuMemcpy2DAsync_v2(&m, s.handle), "cuMemcpy2DAsync"); +} + +/// One part's rows [m, part.n] into `y` (row pitch part.n), bf16 or fp32. +fn part(p: anytype, ops: *load.Ops, pt: load.Part, x: u64, k: usize, m: usize, y: u64, f32_out: bool) !void { + switch (pt.kind) { + .exl3 => |w| { + const de = try ops.linear(); + try de.decode(p.stream, w, x, .bf16, m, y, if (f32_out) .f32 else .bf16); + }, + .native => |w| { + if (f32_out) return p.k.native_lane.matmulF32(p.stream, x, k, m, w, y, p.s.native_part); + try p.k.native_lane.matmul(p.stream, x, k, m, w, y, p.s.native_part); + }, + } +} + +/// Pass.lane() for an EXL3 projection: out bf16 [m, n] (sk 1), or fp32 [m. +pub fn lane(p: anytype, x: u64, e: *const load.Proj, out: u64, slices: u64, m: usize, sk: usize) !void { + const s = p.stream; + if (e.x == .f16) return laneF16(p, x, e, out, slices, m, sk); + if (sk > 1) { + if (e.parts.len != 1) return error.UnexpectedTensor; + try part(p, e.ops, e.parts[0], x, e.k, m, slices, true); + const bytes = m * e.n * 4; + try s.d.check(s.d.api.cuMemsetD8Async(slices + bytes, 0, bytes * (sk - 1), s.handle), "cuMemsetD8Async"); + return; + } + if (try @import("../tp/cuda_exl3_fork.zig").lane(p, x, e, out, m, onePart)) return; // the parts on parallel streams + for (e.parts) |pt| try onePart(p, e, pt, x, m, out); +} + +/// One part of a bf16 sk-1 lane() on the pass's stream: one part over the whole row, or a single row, writes its +fn onePart(p: anytype, e: *const load.Proj, pt: load.Part, x: u64, m: usize, out: u64) !void { + const s = p.stream; + if (pt.n == e.n or m == 1) return part(p, e.ops, pt, x, e.k, m, out + pt.col * 2, false); + const tmp = try e.ops.scratch(s, &e.ops.tmp, m * pt.n * 2); + try part(p, e.ops, pt, x, e.k, m, tmp, false); + try copy2d(s, out + pt.col * 2, e.n * 2, tmp, pt.n * 2, pt.n * 2, m); +} + +/// fp16 rows through a dense BF16 lane weight: x = hi + lo exactly (fn_exl3_split_f16), hi . w and lo . w in fp32. +fn laneF16(p: anytype, x: u64, e: *const load.Proj, out: u64, slices: u64, m: usize, sk: usize) !void { + const s = p.stream; + if (e.parts.len != 1 or e.parts[0].kind != .native) return error.UnexpectedTensor; + const w = e.parts[0].kind.native; + const count = m * e.k; + const halves = try e.ops.scratch(s, &e.ops.halves, 2 * count * 2); + var a: cuda.Args = .{}; + a.add(x); + a.add(halves); + a.add(halves + count * 2); + a.add(@as(c_longlong, @intCast(count))); + try cuda.launch.launch(try e.ops.split(), .{ .grid = .{ .x = @intCast((count + 255) / 256) }, .block = .{ .x = 256 } }, s, &a); + const bytes = m * e.n * 4; + const sums = if (sk == 1) try e.ops.scratch(s, &e.ops.sums, 2 * bytes) else slices; + try p.k.native_lane.matmulF32(s, halves, e.k, m, w, sums, p.s.native_part); + try p.k.native_lane.matmulF32(s, halves + count * 2, e.k, m, w, sums + bytes, p.s.native_part); + if (sk == 1) return p.sumSlices(sums, out, true, m * e.n, 2); + if (sk > 2) try s.d.check(s.d.api.cuMemsetD8Async(slices + 2 * bytes, 0, bytes * (sk - 2), s.handle), "cuMemsetD8Async"); +} + +/// Pass.moeBlock() for EXL3 experts: the router's 513 logits (the BF16 rows the Store widened exactly), each row's ten +pub fn moe(p: anytype, lw: anytype, x: *const load.Experts) !void { + const s = p.s; + const n = p.rows; + try p.tri.routerMany(s.mixed, lw.router.ptr, s.router_l, dims, n); + try p.tri.topkMany(s.router_l, s.picks, s.wts, n); + if (x.needsPool()) return error.Exl3PoolNotAttached; + const ex = try x.ops.grouped(); + const block = @min(n, moe_rows); + const bytes = try cuda.exl3_experts.scratchBytes(x.layer, block, slots); + const base = try x.ops.scratch(p.stream, &x.ops.moe, bytes); + var r0: usize = 0; + while (r0 < n) : (r0 += block) { + const rows = @min(block, n - r0); + var sc = try cuda.exl3_experts.scratch(base, x.layer, block, slots); + // runSlots leaves each slot's fp32 output in sc.y [rows * slots, dims]: the pass's y rows are that layout + sc.y = s.y + r0 * slots * dims * 4; + const xin = s.mixed + r0 * dims * 2; + const picks = s.picks + r0 * slots * 4; + if (x.pool) |pool| { + try pooled(p.stream, x.ops, pool, ex, x.layer, xin, picks, rows, sc, 0, 0); + continue; + } + try ex.runSlots(p.stream, x.layer, xin, dims, .bf16, picks, rows, sc); + } +} + +/// Host time of the A/B cut (the picks D2H and its wait) and of Pool.plan, summed over every pooled block (read by +pub var pool_time: struct { cut_ns: u64 = 0, plan_ns: u64 = 0, blocks: u64 = 0 } = .{}; + +/// One block of `rows` rows through `pool`: per-slot outputs in sc.y, and with `wts` != 0 the last pass also combines +pub fn pooled(st: cuda.Stream, ops: *load.Ops, pool: *pool_.Pool, ex: *const cuda.exl3_experts.Experts, layer: cuda.exl3_experts.Layer, xin: u64, picks: u64, rows: usize, sc: cuda.exl3_experts.Scratch, wts: u64, out: u64) !void { + const host = (ops.picks orelse return error.Exl3ScratchNotReserved).slice(i32)[0 .. rows * slots]; + const io = pool.io; + const t0 = std.Io.Clock.awake.now(io).toNanoseconds(); + // end of segment A: the block's picks reach the host (the only host wait of the layer) + try st.d.check(st.d.api.cuMemcpyDtoHAsync_v2(host.ptr, picks, rows * slots * 4, st.handle), "cuMemcpyDtoHAsync"); + try st.synchronize(); + const t1 = std.Io.Clock.awake.now(io).toNanoseconds(); + const r = try pool.plan(host, rows); + pool_time.cut_ns += @intCast(t1 - t0); + pool_time.plan_ns += @intCast(std.Io.Clock.awake.now(io).toNanoseconds() - t1); + pool_time.blocks += 1; + // segment B, a pass at a time: its copies and table rewrites, then the experts on the pool's tables + for (0..r.n_passes) |j| { + try pool.stage(r, j, st); + const last = j + 1 == r.n_passes; + try ex.run(st, layer, xin, dims, .bf16, pool.picksFor(r, j, picks), if (last) wts else 0, rows, if (last) out else 0, sc); + try pool.release(r, j, st); + } +} + +/// Every offloaded experts layer of a loaded engine (the 48 decoder layers). +pub fn offloadedLayers(layers: []const @import("../../cuda_prompt_load.zig").Layer, out: []*load.Experts) usize { + var n: usize = 0; + for (layers) |*l| if (l.exl3) |x| if (x.needsPool()) { + out[n] = x; + n += 1; + }; + return n; +} + +pub const Budget = struct { + /// Device bytes left free after the pools; null: the native server's own reserve (TENSORFOLD_MEMORY_RESERVE_GIB). + reserve: ?u64 = null, + /// What the pools leave beyond the server's reserve: allocator rounding and the copy streams' own memory. + margin: u64 = 256 << 20, + /// Bytes the family still allocates after the pools: each admitted stream's caches. + other: u64 = 0, + /// Lines a layer, or null: from free memory (cuda_exl3_round.linesFor). + lines: ?usize = null, + /// Rows a planned round holds (the pass's block of routed rows). + max_rows: usize = moe_rows, +}; + +/// TF_FLASHNEXT_EXPERT_LINES and TF_FLASHNEXT_EXPERT_RESERVE_MB override pool sizing. +pub fn budgetFromEnv(stream_bytes: u64) !Budget { + return budgetFromEnvWithStreams(stream_bytes, null); +} + +pub fn budgetForStreams(stream_bytes: u64, requested_streams: u32) !Budget { + return budgetFromEnvWithStreams(stream_bytes, requested_streams); +} + +fn budgetFromEnvWithStreams(stream_bytes: u64, requested_streams: ?u32) !Budget { + const streams = requested_streams orelse 1; + var b: Budget = .{ .other = try streamReservation(stream_bytes, streams) }; + if (std.c.getenv("TF_FLASHNEXT_EXPERT_LINES")) |v| b.lines = try std.fmt.parseInt(usize, std.mem.span(v), 10); + if (std.c.getenv("TF_FLASHNEXT_EXPERT_RESERVE_MB")) |v| b.reserve = try memory.mebibytes(std.mem.span(v)); + if (requested_streams == null) if (std.c.getenv("TF_FLASHNEXT_EXL3_STREAMS")) |v| { + const env_streams = try std.fmt.parseInt(u64, std.mem.span(v), 10); + if (env_streams == 0) return error.InvalidStreamCount; + b.other = try std.math.mul(u64, stream_bytes, env_streams); + }; + return b; +} + +fn streamReservation(stream_bytes: u64, streams: u32) !u64 { + if (streams == 0) return error.InvalidStreamCount; + return std.math.mul(u64, stream_bytes, streams); +} + +test "EXL3 pool reserves the requested number of concurrent stream caches" { + try std.testing.expectEqual(@as(u64, 200), try streamReservation(100, 2)); + try std.testing.expectError(error.InvalidStreamCount, streamReservation(100, 0)); +} + +pub const Attached = struct { layers: usize = 0, lines: usize = 0, line_bytes: usize = 0, device_bytes: u64 = 0, free_before: u64 = 0, free_after: u64 = 0, scratch_bytes: u64 = 0, bank_bytes: u64 = 0, bank_ns: u64 = 0 }; + +pub fn freeBytes(d: *const cuda.Driver) !u64 { + var free: usize = 0; + var total: usize = 0; + try d.check(d.api.cuMemGetInfo_v2(&free, &total), "cuMemGetInfo"); + return free; +} + +/// The sizing hook (cuda_lanes_load.zig, once everything else of the family is allocated). +pub fn attach(gpa: std.mem.Allocator, io: std.Io, d: *const cuda.Driver, st: cuda.Stream, layers: []const @import("../../cuda_prompt_load.zig").Layer, rows: usize, b: Budget, profile_counts: ?*[pool_.shape.layers][pool_.shape.experts]u32) !Attached { + var all: [wt.layers_n]*load.Experts = undefined; + const n = offloadedLayers(layers, &all); + const delta_text = if (std.c.getenv("TF_FLASHNEXT_EXL3_LINE_DELTAS")) |value| std.mem.span(value) else null; + var out: Attached = .{}; + const scratch_before = try freeBytes(d); + // every distinct Ops (the engine's store) reserves its scratch now + var seen: [4]*load.Ops = undefined; + var n_seen: usize = 0; + for (layers) |*l| if (l.exl3) |x| try reserveOnce(x.ops, &seen, &n_seen, st, rows); + if (n == 0) { + if (delta_text != null) return error.InvalidLineDeltaIndex; + return out; + } + for (all[0..n]) |x| { + out.line_bytes = @max(out.line_bytes, x.lineBytes()); + out.bank_bytes += x.bank.?.bytes.len; + out.bank_ns += x.bank_ns; + } + var free: usize = 0; + var total: usize = 0; + try d.check(d.api.cuMemGetInfo_v2(&free, &total), "cuMemGetInfo"); + out.free_before = free; + out.scratch_bytes = scratch_before -| free; + const reserve_floor = b.reserve orelse serverReserve(total); + const reserve = b.reserve orelse try std.math.add(u64, reserve_floor, b.margin); + // beside its lines a layer holds the shared expert's region, its six tables and its masked picks (sized for the + var fixed: u64 = 0; + for (all[0..n]) |x| { + for (x.shared) |t| fixed = try std.math.add(u64, fixed, try std.math.mul(u64, (try std.math.add(u64, t.len, 255)) / 256, 256)); + const picks_bytes = try std.math.mul(u64, try std.math.mul(u64, pool_.passesFor(pool_.min_lines), b.max_rows), slots * 4); + fixed = try std.math.add(u64, fixed, try std.math.add(u64, cuda.exl3_pool.Tables.init(wt.experts_n + 1).bytes + (64 << 10), picks_bytes)); + } + const future = b.other; + const pool_room = try memory.available(free, reserve, future); + const lines = b.lines orelse try pool_.linesFor(free, reserve, try std.math.add(u64, future, fixed), out.line_bytes, n, pool_.min_lines); + if (lines < pool_.min_lines) return error.PoolTooSmall; + const startup_bytes = try memory.poolBytes(lines, out.line_bytes, n, fixed); + if (startup_bytes > pool_room) return error.InsufficientStartupMemory; + out.lines = lines; + var pool_lines: [wt.layers_n]usize = @splat(lines); + var line_bytes: [wt.layers_n]usize = undefined; + var profile_sized = false; + var profile_min_lines: usize = std.math.maxInt(usize); + var profile_max_lines: usize = 0; + for (all[0..n], 0..) |x, index| line_bytes[index] = x.lineBytes(); + if (delta_text == null and shouldProfileSize(profileSizingEnabled(), b.lines)) if (profile_counts) |prior| { + var count_rows: [wt.layers_n][]const u32 = undefined; + var allocation_costs: [wt.layers_n]usize = undefined; + var allocation_lines: [wt.layers_n]usize = undefined; + var decoder_index: usize = 0; + var has_profile_signal = false; + for (all[0..n], 0..) |_, pool_index| { + for (prior[decoder_index]) |count| { + has_profile_signal = has_profile_signal or count != 0; + } + count_rows[decoder_index] = &prior[decoder_index]; + allocation_costs[decoder_index] = line_bytes[pool_index]; + decoder_index += 1; + } + if (has_profile_signal and decoder_index == wt.layers_n) { + const maybe_plan: ?profile_line_budget.Budget = profile_line_budget.allocate(gpa, .{ + .counts = &count_rows, + .rank = 0, + .world = 1, + .line_bytes = &allocation_costs, + .uniform_lines = lines, + .min_lines = pool_.min_lines, + .max_rows = b.max_rows, + .slots = slots, + }, &allocation_lines) catch |err| blk: { + std.log.info("flashnext exl3: profile-sized pools unavailable ({s}); retaining uniform decoder pools", .{@errorName(err)}); + break :blk null; + }; + if (maybe_plan) |plan| { + decoder_index = 0; + for (all[0..n], 0..) |_, pool_index| { + pool_lines[pool_index] = allocation_lines[decoder_index]; + profile_min_lines = @min(profile_min_lines, allocation_lines[decoder_index]); + profile_max_lines = @max(profile_max_lines, allocation_lines[decoder_index]); + decoder_index += 1; + } + if (plan.uniform_fallback) { + @memset(pool_lines[0..n], lines); + std.log.info("flashnext exl3: profile candidate changed total variable CUDA bytes; retaining uniform decoder pools", .{}); + } else { + profile_sized = true; + } + std.log.info("flashnext exl3: profile-weighted decoder allocation used {d}/{d} variable CUDA bytes (expert lines and device picks); saved lifetime picks are a proxy, not latency", .{ plan.allocated_bytes, plan.baseline_bytes }); + std.log.info("flashnext exl3: profile decoder pool sizes range {d}..{d} lines; uniform-equivalent baseline was {d}", .{ profile_min_lines, profile_max_lines, lines }); + } + } + }; + if (delta_text) |text| { + const plan = try line_deltas.apply(text, lines, n, line_bytes[0..n], pool_.min_lines, fixed, startup_bytes, pool_lines[0..n]); + std.log.info("flashnext exl3: experimental rank-local decoder pool line deltas {s}; base={d} startup_bytes={d}", .{ text, lines, plan.pool_bytes }); + } + // the first decoder layer's pool owns the lifetime use table the other decoder layers share + var counts = profile_counts; + var decoder: u8 = 0; + for (all[0..n], pool_lines[0..n]) |x, count| { + try x.attach(io, d, decoder, count, b.max_rows, counts); + if (counts == null) counts = x.pool.?.counts; + decoder += 1; + out.device_bytes += x.pool.?.deviceBytes(); + out.layers += 1; + } + out.free_after = try freeBytes(d); + _ = memory.available(out.free_after, reserve_floor, future) catch |err| { + std.log.err("EXL3 pool reservation shortfall: free_before={d} free_after={d} pool_bytes={d} floor={d} allocator_margin={d} future={d}", .{ free, out.free_after, out.device_bytes, reserve_floor, b.margin, future }); + return err; + }; + const allocation = if (delta_text != null) "uniform-equivalent lines (experimental deltas applied)" else if (profile_sized) "uniform-equivalent lines (profile-sized decoder pools)" else "lines a layer"; + std.log.info("flashnext exl3: {d} offloaded expert layers, {d} {s} of {d:.2} MiB ({d:.2} GiB on the device, {d:.2} GiB free before); pinned host bank {d:.2} GiB copied in {d:.1} s", .{ out.layers, out.lines, allocation, @as(f64, @floatFromInt(out.line_bytes)) / (1 << 20), @as(f64, @floatFromInt(out.device_bytes)) / (1 << 30), @as(f64, @floatFromInt(out.free_before)) / (1 << 30), @as(f64, @floatFromInt(out.bank_bytes)) / (1 << 30), @as(f64, @floatFromInt(out.bank_ns)) / 1e9 }); + return out; +} + +test { + _ = line_deltas; + _ = profile_line_budget; +} + +fn profileSizingEnabled() bool { + const value = std.c.getenv("TF_FLASHNEXT_EXL3_PROFILE_SIZING") orelse return false; + return shouldProfileSize(std.mem.eql(u8, std.mem.span(value), "1"), null); +} + +fn shouldProfileSize(enabled: bool, explicit_lines: ?usize) bool { + return enabled and explicit_lines == null; +} + +test "exl3: profile sizing requires explicit opt-in and does not override fixed line budgets" { + try std.testing.expect(!shouldProfileSize(false, null)); + try std.testing.expect(!shouldProfileSize(true, 32)); + try std.testing.expect(shouldProfileSize(true, null)); +} + +/// native/cuda_memory.zig reserveBytes for a discrete card: TENSORFOLD_MEMORY_RESERVE_GIB, else max(1 GiB, total / 32). +fn serverReserve(total: u64) u64 { + if (std.c.getenv("TENSORFOLD_MEMORY_RESERVE_GIB")) |v| { + const g = std.fmt.parseFloat(f64, std.mem.trim(u8, std.mem.span(v), " ")) catch 0; + if (g >= 0.5) return @intFromFloat(g * (1 << 30)); + } + return @max(@as(u64, 1 << 30), total / 32); +} + +fn reserveOnce(o: *load.Ops, seen: []*load.Ops, n: *usize, st: cuda.Stream, rows: usize) !void { + for (seen[0..n.*]) |x| if (x == o) return; + if (n.* == seen.len) return error.Invalid; + seen[n.*] = o; + n.* += 1; + try o.reserve(st, rows, moe_rows, slots); +} + +/// Pool statistics summed over every attached layer. +pub fn poolStats(layers: []const @import("../../cuda_prompt_load.zig").Layer) pool_.Stats { + var t: pool_.Stats = .{}; + var all: [wt.layers_n]*load.Experts = undefined; + const n = offloadedPools(layers, &all); + for (all[0..n]) |x| { + const s = x.pool.?.stats; + inline for (@typeInfo(pool_.Stats).@"struct".field_names) |f| { + @field(t, f) += @field(s, f); + } + } + return t; +} + +pub fn offloadedPools(layers: []const @import("../../cuda_prompt_load.zig").Layer, out: []*load.Experts) usize { + var n: usize = 0; + for (layers) |*l| if (l.exl3) |x| if (x.pool != null) { + out[n] = x; + n += 1; + }; + return n; +} diff --git a/zig/src/families/flashnext/exl3/model/cuda_exl3_ple.zig b/zig/src/families/flashnext/exl3/model/cuda_exl3_ple.zig new file mode 100644 index 0000000000..b49662450e --- /dev/null +++ b/zig/src/families/flashnext/exl3/model/cuda_exl3_ple.zig @@ -0,0 +1,543 @@ +//! An EXL3 pack's n-gram table (ExLlamaV3 `ngram_embedding.safetensors`) behind `cuda_ple.Ngram` (`format = .exl3`). + +const std = @import("std"); +const core = @import("core"); +const ple = @import("../../cuda_ple.zig"); +const fmt = core.exl3_format; + +pub const file_name = "ngram_embedding.safetensors"; +/// Values a row holds: one head's slice of the 2,560-wide embedding. +pub const row_values: usize = 160; +pub const hot_slots: usize = 4096; +/// Bytes a hot slot holds (a row at 8 bits is 162). +pub const hot_cap: usize = 176; + +pub const Read = enum { map, pread }; +pub const Emb = enum { f16, bf16 }; +pub const Options = struct { read: Read = .map, emb: Emb = .f16 }; + +/// TF_FLASHNEXT_PLE_READ (map | pread); anything else keeps the default. The bf16 rows are a test option (`emb`). +pub fn optionsFromEnv() Options { + var o: Options = .{}; + if (std.c.getenv("TF_FLASHNEXT_PLE_READ")) |v| { + if (std.mem.eql(u8, std.mem.span(v), "pread")) o.read = .pread; + } + return o; +} + +// ---- the row codec (ported from exl3-flashnext-window-wip exl3_ngram.zig) ------------------------------------------ + +/// Int16 words in one row at `bits`: the fp16 scale, then 160 values. `bits` outside 2..8 is `NgramWidth`. +pub fn wordsPerRow(bits: usize) error{NgramWidth}!usize { + if (bits < 2 or bits > 8) return error.NgramWidth; + return 1 + row_values * bits / 16; +} + +/// Mul1 state of value `index`: bit m comes from position `index - m / bits` (plus 160 when negative), bit `m % bits`. +pub fn state(words: []const i16, bits: usize, index: usize) u16 { + var acc: u16 = 0; + for (0..16) |m| { + var pos: isize = @as(isize, @intCast(index)) - @as(isize, @intCast(m / bits)); + if (pos < 0) pos += @as(isize, @intCast(row_values)); + const sb: usize = @as(usize, @intCast(pos)) * bits + (m % bits); + const word: u16 = @bitCast(words[1 + (sb >> 4)]); + acc |= ((word >> @as(u4, @intCast(sb & 15))) & 1) << @as(u4, @intCast(m)); + } + return acc; +} + +/// One row into 160 fp16 values: one rounding of `f32(codebook) * f32(scale) + f32(bias[i])`. +pub fn decodeRow(words: []const i16, bits: usize, bias: []const f16, out: []f16) error{NgramWidth}!void { + const width = try wordsPerRow(bits); + if (words.len != width or bias.len != row_values or out.len != row_values) return error.NgramWidth; + const scale: f32 = @floatCast(@as(f16, @bitCast(words[0]))); + // `state(words, bits, i)` for every i at once: value j is stream bits [j * bits, (j + 1) * bits); state i is + var vals: [row_values]u32 = undefined; + const mask: u32 = (@as(u32, 1) << @intCast(bits)) - 1; + for (&vals, 0..) |*v, j| { + const sb = j * bits; + const wi = 1 + (sb >> 4); + const lo: u32 = @as(u16, @bitCast(words[wi])); + const hi: u32 = if (wi + 1 < words.len) @as(u16, @bitCast(words[wi + 1])) else 0; + v.* = ((lo | (hi << 16)) >> @intCast(sb & 15)) & mask; + } + for (out, bias, 0..) |*o, b, i| { + var acc: u32 = 0; + var k: usize = 0; + while (k * bits < 16) : (k += 1) acc |= vals[(i + row_values - k) % row_values] << @intCast(k * bits); + const cb: f32 = @floatCast(fmt.codebookValue(.mul1, @truncate(acc))); + o.* = @floatCast(cb * scale + @as(f32, @floatCast(b))); + } +} + +test "exl3 ple: the packed decode's states equal `state` for every value at every width" { + var prng = std.Random.DefaultPrng.init(11); + for (2..9) |bits| { + const n = try wordsPerRow(bits); + var words: [1 + row_values * 8 / 16]i16 = undefined; + for (words[0..n]) |*w| w.* = @bitCast(prng.random().int(u16)); + words[0] = @bitCast(@as(f16, 0.5)); + var bias: [row_values]f16 = undefined; + for (&bias) |*b| b.* = @floatCast(prng.random().float(f32) - 0.5); + var got: [row_values]f16 = undefined; + try decodeRow(words[0..n], bits, &bias, &got); + const scale: f32 = @floatCast(@as(f16, @bitCast(words[0]))); + for (got, bias, 0..) |g, b, i| { + const cb: f32 = @floatCast(fmt.codebookValue(.mul1, state(words[0..n], bits, i))); + const want: f16 = @floatCast(cb * scale + @as(f32, @floatCast(b))); + try std.testing.expectEqual(@as(u16, @bitCast(want)), @as(u16, @bitCast(g))); + } + } +} + +fn bf16Of(v: f32) u16 { + return ple.toBf16(v); +} + +/// Group sums of 32 fp16 values (the tree `cuda_ple.sumGroups` adds bf16 values in), for an fp16 row. +pub fn sumGroupsF16(row: []const u16, out: []f32) void { + for (0..row.len / 32) |g| { + var buf: [32]f32 = undefined; + for (0..32) |i| buf[i] = @floatCast(@as(f16, @bitCast(row[g * 32 + i]))); + var n: usize = 32; + while (n > 1) { + n /= 2; + for (0..n) |i| buf[i] = buf[2 * i] + buf[2 * i + 1]; + } + out[g] = buf[0]; + } +} + +// ---- the table ---------------------------------------------------------------------------------------------------- + +const c = std.c; +const is_windows = @import("builtin").os.tag == .windows; + +/// The open table. Heap-allocated (`open`/`destroy`): the pread reader's hot cache is 512 KiB. +pub const Table = struct { + gpa: std.mem.Allocator, + io: std.Io, + opts: Options, + heads: usize, + shards: usize, + bits: usize, + words: usize, + /// File byte offset of each shard's row 0, and its row count. + offset: [128]u64 = undefined, + rows: [128]u64 = undefined, + bias: []f16, + /// `.map`: the whole file's mapping (header parsed, nothing populated). + file: ?core.safetensors.File = null, + /// `.pread`: the descriptor; the file is never mapped past `open`. + fd: c.fd_t = if (is_windows) undefined else -1, + hot: ?*Hot = null, + + pub const Hot = struct { + lock: std.atomic.Value(u32) = .init(0), + lookups: u64 = 0, + misses: u64 = 0, + tags: [hot_slots]u64 = @splat(std.math.maxInt(u64)), + stored: [hot_slots][hot_cap]u8 = undefined, + + fn acquire(h: *Hot) void { + while (h.lock.cmpxchgWeak(0, 1, .acquire, .monotonic) != null) std.atomic.spinLoopHint(); + } + fn release(h: *Hot) void { + h.lock.store(0, .release); + } + }; + + fn prefix(buf: []u8, part: []const u8) ![]const u8 { + return std.fmt.bufPrint(buf, "model.language_model.layers.{d}.ple.ple_embedding.ngram_embedding.{s}", .{ ple.layer, part }); + } + + /// The pack's table when `dir` holds an EXL3 n-gram file whose constants equal `hash`'s; null. + pub fn open(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, hash: ple.Hash, opts: Options) !?*Table { + const path = try std.fs.path.join(gpa, &.{ dir, file_name }); + defer gpa.free(path); + var file = core.safetensors.File.open(gpa, io, path) catch |err| switch (err) { + error.FileNotFound => return null, + else => return err, + }; + var keep_file = false; + defer if (!keep_file) file.close(io); + var nb: [200]u8 = undefined; + const first = file.get(try prefix(&nb, "shard_0.trellis")) orelse return null; + if (first.dtype != .i16 or first.rank != 2) return error.UnexpectedTensor; + if (hash.shards > 128 or hash.dh != row_values) return error.UnexpectedTensor; + // the hash constants the pack ships against the ones the config derives + { + var got: [ple.max_heads]i64 = undefined; + try ple.readI64(file.get(try prefix(&nb, "layer_multipliers")) orelse return error.MissingTensor, got[0..hash.n]); + if (!std.mem.eql(i64, got[0..hash.n], hash.mult[0..hash.n])) return error.NgramConstants; + try ple.readI64(file.get(try prefix(&nb, "head_vocab_sizes")) orelse return error.MissingTensor, got[0..hash.heads]); + if (!std.mem.eql(i64, got[0..hash.heads], hash.sizes[0..hash.heads])) return error.NgramConstants; + try ple.readI64(file.get(try prefix(&nb, "head_offsets")) orelse return error.MissingTensor, got[0..hash.heads]); + if (!std.mem.eql(i64, got[0..hash.heads], hash.offsets[0..hash.heads])) return error.NgramConstants; + } + // the reference's EXL3 path applies no table scale: refuse a file that carries one + if (file.get(try prefix(&nb, "weight_scale")) != null) return error.UnexpectedTensor; + const bias_t = file.get(try prefix(&nb, "head_bias")) orelse return error.MissingTensor; + if (!bias_t.is(.f16, &.{ hash.heads, row_values })) return error.UnexpectedTensor; + const words = first.dim(1); + const bits = (words - 1) * 16 / row_values; + if ((wordsPerRow(bits) catch return error.UnexpectedTensor) != words) return error.UnexpectedTensor; + + const t = try gpa.create(Table); + errdefer gpa.destroy(t); + t.* = .{ .gpa = gpa, .io = io, .opts = opts, .heads = hash.heads, .shards = hash.shards, .bits = bits, .words = words, .bias = try gpa.alloc(f16, hash.heads * row_values) }; + errdefer gpa.free(t.bias); + @memcpy(std.mem.sliceAsBytes(t.bias), bias_t.bytes); + for (0..hash.shards) |i| { + var part: [32]u8 = undefined; + const e = file.names.get(try prefix(&nb, try std.fmt.bufPrint(&part, "shard_{d}.trellis", .{i}))) orelse return error.MissingTensor; + const want: usize = @intCast(hash.starts[i + 1] - hash.starts[i]); + if (e.dtype != .i16 or e.rank != 2 or e.shape[0] != want or e.shape[1] != words) return error.UnexpectedTensor; + t.offset[i] = file.data + e.begin; + t.rows[i] = want; + } + switch (opts.read) { + .map => { + if (@import("builtin").os.tag == .linux) { + if (file.map.section != null) _ = std.os.linux.madvise(file.map.memory.ptr, file.map.memory.len, std.os.linux.MADV.RANDOM); + } + t.file = file; + keep_file = true; + }, + .pread => { + if (comptime is_windows) return error.UnsupportedPlatform; + const z = try gpa.dupeSentinel(u8, path, 0); + defer gpa.free(z); + const fd = c.open(z.ptr, .{ .ACCMODE = .RDONLY }); + if (fd < 0) return error.OpenFailed; + t.fd = fd; + errdefer _ = c.close(fd); + if (@import("builtin").os.tag == .linux) _ = std.os.linux.fadvise(fd, 0, 0, std.os.linux.POSIX_FADV.RANDOM); + const hot = try gpa.create(Hot); + hot.* = .{}; + t.hot = hot; + }, + } + return t; + } + + pub fn destroy(t: *Table) void { + if (t.file) |*f| f.close(t.io); + if (comptime !is_windows) { + if (t.fd >= 0) _ = c.close(t.fd); + } + if (t.hot) |h| t.gpa.destroy(h); + t.gpa.free(t.bias); + t.gpa.destroy(t); + } + + pub fn rowBytes(t: *const Table) usize { + return t.words * 2; + } + + fn preadAll(t: *const Table, dest: []u8, at: u64) !void { + if (comptime is_windows) return error.UnsupportedPlatform; + var done: usize = 0; + while (done < dest.len) { + const n = c.pread(t.fd, dest.ptr + done, dest.len - done, @intCast(at + done)); + if (n <= 0) return error.ReadFailed; + done += @intCast(n); + } + } + + /// The stored bytes of row `local` of shard `shard` (`rowBytes` long), through the mapping or the hot cache. + pub fn raw(t: *const Table, shard: usize, local: usize, dest: []u8) !void { + if (shard >= t.shards or local >= t.rows[shard] or dest.len != t.rowBytes()) return error.UnexpectedTensor; + const at = t.offset[shard] + @as(u64, local) * t.rowBytes(); + if (t.file) |f| { + @memcpy(dest, f.map.memory[at..][0..dest.len]); + return; + } + const h = t.hot.?; + const slot: usize = @intCast(((at *% 0x9E3779B97F4A7C15) >> 20) % hot_slots); + h.acquire(); + h.lookups += 1; + if (h.tags[slot] == at) { + @memcpy(dest, h.stored[slot][0..dest.len]); + h.release(); + return; + } + h.release(); + try t.preadAll(dest, at); + h.acquire(); + h.misses += 1; + h.tags[slot] = at; + @memcpy(h.stored[slot][0..dest.len], dest); + h.release(); + } + + /// Row `local` of shard `shard`, decoded with head `head`'s bias into 160 fp16 (or bf16) bit patterns. + pub fn row(t: *const Table, shard: usize, local: usize, head: usize, out: []u16) !void { + if (head >= t.heads or out.len != row_values) return error.UnexpectedTensor; + var bytes: [hot_cap]u8 align(2) = undefined; + const rb = t.rowBytes(); + try t.raw(shard, local, bytes[0..rb]); + const words: []const i16 = @as([*]const i16, @ptrCast(&bytes))[0..t.words]; + var vals: [row_values]f16 = undefined; + try decodeRow(words, t.bits, t.bias[head * row_values ..][0..row_values], &vals); + switch (t.opts.emb) { + .f16 => for (out, vals) |*o, v| { + o.* = @bitCast(v); + }, + .bf16 => for (out, vals) |*o, v| { + o.* = bf16Of(@floatCast(v)); + }, + } + } + + /// Group sums of one decoded row in the table's element type. + pub fn sums(t: *const Table, row_bits: []const u16, out: []f32) void { + switch (t.opts.emb) { + .f16 => sumGroupsF16(row_bits, out), + .bf16 => ple.sumGroups(row_bits, out), + } + } + + pub fn hotStats(t: *const Table) struct { lookups: u64, misses: u64 } { + const h = t.hot orelse return .{ .lookups = 0, .misses = 0 }; + return .{ .lookups = h.lookups, .misses = h.misses }; + } +}; + +// ---- tests -------------------------------------------------------------------------------------------------------- + +/// The Triton `_ple_rows` formulation, a second decoder that shares only the three fp16 constants with `decodeRow`. +fn decodeRowRef(words: []const i16, bits: usize, bias: []const f16, out: []f16) void { + const k_inv: f32 = @floatCast(@as(f16, @bitCast(fmt.Mul1ScaleBits))); + const k_bias: f32 = @floatCast(@as(f16, @bitCast(fmt.Mul1BiasBits))); + const scale: f32 = @floatCast(@as(f16, @bitCast(words[0]))); + for (0..out.len) |i| { + var st: u64 = 0; + for (0..16) |m| { + var pos: isize = @as(isize, @intCast(i)) - @as(isize, @intCast(m / bits)); + if (pos < 0) pos += @as(isize, @intCast(out.len)); + const sb: usize = @as(usize, @intCast(pos)) * bits + (m % bits); + const word: u64 = @as(u16, @bitCast(words[1 + (sb >> 4)])); + st |= ((word >> @intCast(sb & 15)) & 1) << @intCast(m); + } + const prod: u64 = (st * 0x83DCD12D) & 0xFFFFFFFF; + const hs: u64 = (prod & 255) + ((prod >> 8) & 255) + ((prod >> 16) & 255) + ((prod >> 24) & 255); + const cb: f16 = @floatCast(@as(f32, @floatFromInt(1024 + hs)) * k_inv + k_bias); + out[i] = @floatCast(@as(f32, @floatCast(cb)) * scale + @as(f32, @floatCast(bias[i]))); + } +} + +test "exl3 ngram: wordsPerRow is one scale word plus 160 K-bit values" { + try std.testing.expectEqual(@as(usize, 21), try wordsPerRow(2)); + try std.testing.expectEqual(@as(usize, 31), try wordsPerRow(3)); + try std.testing.expectEqual(@as(usize, 51), try wordsPerRow(5)); + try std.testing.expectEqual(@as(usize, 81), try wordsPerRow(8)); + try std.testing.expectError(error.NgramWidth, wordsPerRow(1)); + try std.testing.expectError(error.NgramWidth, wordsPerRow(9)); +} + +test "exl3 ngram: patterned rows read tail-biting states" { + var storage: [81]i16 = undefined; + for (&storage, 0..) |*w, i| w.* = @bitCast(@as(u16, @intCast(i))); + try std.testing.expectEqual(@as(u16, 0x5001), state(storage[0..21], 2, 0)); + try std.testing.expectEqual(@as(u16, 0x9001), state(storage[0..31], 3, 0)); + try std.testing.expectEqual(@as(u16, 0x8001), state(storage[0..51], 5, 0)); + try std.testing.expectEqual(@as(u16, 0x0020), state(storage[0..51], 5, 1)); + try std.testing.expectEqual(@as(u16, 0x0400), state(storage[0..51], 5, 2)); + try std.testing.expectEqual(@as(u16, 0x0001), state(storage[0..81], 8, 0)); + // value 0 at 5 bits: bit 5 wraps to position 159, word 50 bit 11 + var words: [51]i16 = @splat(0); + words[50] = @bitCast(@as(u16, 1) << 11); + try std.testing.expectEqual(@as(u16, 1) << 5, state(&words, 5, 0)); +} + +test "exl3 ngram: scale and bias round in f32, then once to f16" { + var words: [51]i16 = @splat(0); + words[0] = @bitCast(@as(u16, 0x3C01)); + var bias: [160]f16 = @splat(0); + bias[0] = @bitCast(@as(u16, 0x1400)); + var out: [160]f16 = undefined; + try decodeRow(&words, 5, &bias, &out); + // codebook(0) is 0xC2E8; times 0x3C01 plus 0x1400 is 0xC2E9 (rounding the product first would give 0xC2EA) + try std.testing.expectEqual(@as(u16, 0xC2E8), @as(u16, @bitCast(fmt.codebookValue(.mul1, 0)))); + try std.testing.expectEqual(@as(u16, 0xC2E9), @as(u16, @bitCast(out[0]))); + try std.testing.expectError(error.NgramWidth, decodeRow(words[0..50], 5, &bias, &out)); + try std.testing.expectError(error.NgramWidth, decodeRow(&words, 2, &bias, &out)); +} + +test "exl3 ngram: decodeRow equals the Triton formulation on random rows at every width" { + var rng = std.Random.DefaultPrng.init(0x91e); + const r = rng.random(); + var words: [81]i16 = undefined; + var bias: [160]f16 = undefined; + var a: [160]f16 = undefined; + var b: [160]f16 = undefined; + for ([_]usize{ 2, 3, 4, 5, 6, 8 }) |bits| { + const n = try wordsPerRow(bits); + for (0..300) |_| { + for (words[0..n]) |*x| x.* = @bitCast(r.int(u16)); + words[0] = @bitCast(@as(f16, @floatCast((r.float(f32) - 0.5) * 4))); + for (&bias) |*x| x.* = @floatCast((r.float(f32) - 0.5) * 0.25); + try decodeRow(words[0..n], bits, &bias, &a); + decodeRowRef(words[0..n], bits, &bias, &b); + try std.testing.expectEqualSlices(u16, @ptrCast(&b), @ptrCast(&a)); + } + } +} + +/// A small table file plus its config: 2 shards, 4 heads of 160 (n-gram 3, 2 heads an n-gram), 5-bit rows. +const Synth = struct { + const vocab = 100; + const eos = 99; + const base = 11; + + fn spec() ple.Hash { + return ple.Hash.init(.{ .vocab = vocab, .ngram = 3, .per = 2, .base = base, .divisor = 4, .shards = 2, .seed = 1234, .eos = eos, .ple_dim = 4 * row_values, .ple_index = 0 }) catch unreachable; + } + + fn put(list: *std.ArrayList(u8), gpa: std.mem.Allocator, bytes: []const u8) !void { + try list.appendSlice(gpa, bytes); + } + + /// Writes config.json, an empty index and ngram_embedding.safetensors under `dir`; row r's words are a hash of r. + fn write(gpa: std.mem.Allocator, io: std.Io, dir: std.Io.Dir) !void { + const h = spec(); + var data: std.ArrayList(u8) = .empty; + defer data.deinit(gpa); + var header: std.ArrayList(u8) = .empty; + defer header.deinit(gpa); + try header.appendSlice(gpa, "{"); + const p = "model.language_model.layers.1.ple.ple_embedding.ngram_embedding."; + var first = true; + const Entry = struct { + fn add(hd: *std.ArrayList(u8), a: std.mem.Allocator, f: *bool, name: []const u8, dtype: []const u8, shape: []const usize, lo: usize, hi: usize) !void { + if (!f.*) try hd.appendSlice(a, ","); + f.* = false; + var buf: [400]u8 = undefined; + var sb: [64]u8 = undefined; + var sl: usize = 0; + for (shape, 0..) |d, i| { + const s = try std.fmt.bufPrint(sb[sl..], "{s}{d}", .{ if (i == 0) "" else ",", d }); + sl += s.len; + } + try hd.appendSlice(a, try std.fmt.bufPrint(&buf, "\"{s}\":{{\"dtype\":\"{s}\",\"shape\":[{s}],\"data_offsets\":[{d},{d}]}}", .{ name, dtype, sb[0..sl], lo, hi })); + } + }; + var nb: [200]u8 = undefined; + for (0..h.shards) |s| { + const rows: usize = @intCast(h.starts[s + 1] - h.starts[s]); + const lo = data.items.len; + for (0..rows) |r| { + const global: u64 = @intCast(h.starts[s] + @as(i64, @intCast(r))); + var prng = std.Random.DefaultPrng.init(global + 1); + for (0..51) |w| { + var v: u16 = prng.random().int(u16); + if (w == 0) v = @bitCast(@as(f16, @floatCast(0.25 + 0.01 * @as(f32, @floatFromInt(global))))); + try data.appendSlice(gpa, &std.mem.toBytes(v)); + } + } + try Entry.add(&header, gpa, &first, try std.fmt.bufPrint(&nb, "{s}shard_{d}.trellis", .{ p, s }), "I16", &.{ rows, 51 }, lo, data.items.len); + } + var lo = data.items.len; + for (0..h.heads * row_values) |i| try data.appendSlice(gpa, &std.mem.toBytes(@as(f16, @floatCast(@as(f32, @floatFromInt(i % 7)) * 0.01)))); + try Entry.add(&header, gpa, &first, try std.fmt.bufPrint(&nb, "{s}head_bias", .{p}), "F16", &.{ h.heads, row_values }, lo, data.items.len); + lo = data.items.len; + for (h.mult[0..h.n]) |v| try data.appendSlice(gpa, &std.mem.toBytes(v)); + try Entry.add(&header, gpa, &first, try std.fmt.bufPrint(&nb, "{s}layer_multipliers", .{p}), "I64", &.{h.n}, lo, data.items.len); + lo = data.items.len; + for (h.sizes[0..h.heads]) |v| try data.appendSlice(gpa, &std.mem.toBytes(v)); + try Entry.add(&header, gpa, &first, try std.fmt.bufPrint(&nb, "{s}head_vocab_sizes", .{p}), "I64", &.{h.heads}, lo, data.items.len); + lo = data.items.len; + for (h.offsets[0..h.heads]) |v| try data.appendSlice(gpa, &std.mem.toBytes(v)); + try Entry.add(&header, gpa, &first, try std.fmt.bufPrint(&nb, "{s}head_offsets", .{p}), "I64", &.{h.heads}, lo, data.items.len); + try header.appendSlice(gpa, "}"); + var out: std.ArrayList(u8) = .empty; + defer out.deinit(gpa); + try out.appendSlice(gpa, &std.mem.toBytes(@as(u64, header.items.len))); + try out.appendSlice(gpa, header.items); + try out.appendSlice(gpa, data.items); + try dir.writeFile(io, .{ .sub_path = file_name, .data = out.items }); + try dir.writeFile(io, .{ .sub_path = "model.safetensors.index.json", .data = "{\"weight_map\":{}}" }); + try dir.writeFile(io, .{ .sub_path = "config.json", .data = + \\{"text_config":{"vocab_size":100,"eos_token_id":99,"num_hidden_layers":4,"hidden_size":640,"ple_layer_ids":[2], + \\"ngram_size":3,"heads_per_ngram":2,"ngram_vocab_size_base":11,"make_ngram_vocab_size_divisible_by":4, + \\"split_ngram_parts":2,"ple_embed_dim":640}} + }); + } +}; + +test "exl3 ngram: Ngram.rows on a synthetic EXL3 table: hook, map and pread readers, hot cache" { + if (comptime is_windows) return error.SkipZigTest; // the pread reader is POSIX-only + const gpa = std.testing.allocator; + const io = std.testing.io; + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + try Synth.write(gpa, io, tmp.dir); + const dir = try tmp.dir.realPathFileAlloc(io, ".", gpa); + defer gpa.free(dir); + + var g = try ple.Ngram.init(gpa, io, dir); + defer g.deinit(); + try std.testing.expect(g.format == .exl3); + const h = g.hash; + try std.testing.expectEqual(@as(usize, 4), h.heads); + const tokens = [_]u32{ 5, 7, Synth.eos, 7, 7, 3 }; + const w = g.width(); + var out: [tokens.len * 4 * row_values]u16 = undefined; + var xs: [tokens.len * 4 * row_values / 32]f32 = undefined; + var history = [_]i64{ Synth.eos, Synth.eos }; + try g.rows(&history, &tokens, &out, &xs); + try std.testing.expectEqualSlices(i64, &.{ 7, 3 }, &history); + + // the same rows by hand: upstream's ids, the shard's bytes, decodeRow + var pread_t = (try Table.open(gpa, io, dir, h, .{ .read = .pread })).?; + defer pread_t.destroy(); + var hist = [_]i64{ Synth.eos, Synth.eos }; + for (tokens, 0..) |tok, t| { + var ids: [4]i64 = undefined; + try ple.lastIds(h, &hist, tok, &ids); + hist = .{ hist[1], tok }; + for (ids, 0..) |id, head| { + const loc = try h.locate(id); + var bytes: [102]u8 align(2) = undefined; + try pread_t.raw(loc.shard, loc.local, &bytes); + var mapped: [102]u8 = undefined; + try g.x3.?.raw(loc.shard, loc.local, &mapped); + try std.testing.expectEqualSlices(u8, &mapped, &bytes); + var vals: [row_values]f16 = undefined; + try decodeRow(@as([*]const i16, @ptrCast(&bytes))[0..51], 5, pread_t.bias[head * row_values ..][0..row_values], &vals); + const got = out[t * w + head * row_values ..][0..row_values]; + try std.testing.expectEqualSlices(u16, @ptrCast(&vals), got); + var gs: [5]f32 = undefined; + sumGroupsF16(got, &gs); + try std.testing.expectEqualSlices(f32, &gs, xs[(t * w + head * row_values) / 32 ..][0..5]); + } + } + // tokens 2 and 3 after an eos: token 7 twice in a fresh segment reads the same rows; the cache serves the repeat + const before = pread_t.hotStats(); + var again: [102]u8 = undefined; + const loc = try h.locate(h.offsets[0]); + try pread_t.raw(loc.shard, loc.local, &again); + try pread_t.raw(loc.shard, loc.local, &again); + const after = pread_t.hotStats(); + try std.testing.expectEqual(before.lookups + 2, after.lookups); + try std.testing.expect(after.misses <= before.misses + 1); + // bf16 elements: the fp16 values rounded once more + var bf_t = (try Table.open(gpa, io, dir, h, .{ .emb = .bf16 })).?; + defer bf_t.destroy(); + var r16: [row_values]u16 = undefined; + var rbf: [row_values]u16 = undefined; + try g.x3.?.row(loc.shard, loc.local, 0, &r16); + try bf_t.row(loc.shard, loc.local, 0, &rbf); + for (r16, rbf) |a, b| try std.testing.expectEqual(ple.toBf16(@floatCast(@as(f16, @bitCast(a)))), b); +} + +test "exl3 ngram: a directory without the n-gram file is not an EXL3 table" { + const gpa = std.testing.allocator; + const io = std.testing.io; + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + const dir = try tmp.dir.realPathFileAlloc(io, ".", gpa); + defer gpa.free(dir); + try std.testing.expect((try Table.open(gpa, io, dir, Synth.spec(), .{})) == null); +} diff --git a/zig/src/families/flashnext/exl3/model/cuda_exl3_telemetry.zig b/zig/src/families/flashnext/exl3/model/cuda_exl3_telemetry.zig new file mode 100644 index 0000000000..f6f245a852 --- /dev/null +++ b/zig/src/families/flashnext/exl3/model/cuda_exl3_telemetry.zig @@ -0,0 +1,85 @@ +const std = @import("std"); +const prompt = @import("../../cuda_prompt.zig"); +const pool = @import("../offload/cuda_exl3_pool.zig"); +const pass = @import("cuda_exl3_pass.zig"); +const weights = @import("../../cuda_prompt_wt.zig"); +pub const counters = @import("exl3_telemetry.zig"); +pub const Completion = enum { all_rank_consumers_joined }; + +pub fn sampleAfterConsumers(allocator: std.mem.Allocator, engine: *prompt.Engine, sequences: []const *prompt.Seq, rank: counters.Rank, completion: Completion) !counters.Snapshot { + _ = completion; + try rank.validate(); + var hash = std.hash.Wyhash.init(@intFromPtr(engine)); + var totals: counters.Experts = .{}; + var attached: [weights.layers_n]*@import("../pack/cuda_exl3_load.zig").Experts = undefined; + const attached_count = pass.offloadedPools(engine.layers[0..engine.n_layers], &attached); + var layers: []counters.LayerSample = if (engine.telemetry.layers_enabled) try allocator.alloc(counters.LayerSample, attached_count) else @constCast(&.{}); + errdefer allocator.free(layers); + var seen: [weights.layers_n]*pool.Pool = undefined; + var seen_count: usize = 0; + for (attached[0..attached_count]) |experts| { + const resident = experts.pool orelse continue; + var duplicate = false; + for (seen[0..seen_count]) |previous| if (previous == resident) { + duplicate = true; + break; + }; + if (duplicate) continue; + if (resident.telemetry_id == 0) return error.TelemetryOwnerChanged; + if (!resident.telemetry_valid) return error.TelemetryCounterRegression; + seen[seen_count] = resident; + seen_count += 1; + hash.update(std.mem.asBytes(&resident.telemetry_id)); + const device = resident.telemetry_device; + var minimum: u64 = std.math.maxInt(u64); + var maximum: u64 = 0; + for (resident.src.routed, resident.k2) |source, widths| { + if (widths[0] == 0) continue; + const bytes = try std.math.add(u64, try std.math.add(u64, source[0].len, source[1].len), source[2].len); + minimum = @min(minimum, bytes); + maximum = @max(maximum, bytes); + } + if (minimum == std.math.maxInt(u64)) { + if (device.copies != 0) return error.TelemetryInconsistentCounters; + minimum = 0; + } + const host = try counters.difference(counters.Device, resident.telemetry_device, .{ .plans = resident.stats.rounds, .hit_picks = resident.stats.hit_slots, .miss_picks = resident.stats.miss_slots, .copies = resident.stats.misses, .evictions = resident.stats.evictions, .inplace_experts = resident.telemetry_device.inplace_experts, .overflow_plans = resident.telemetry_device.overflow_plans }); + const contribution: counters.Experts = .{ + .host_plans = host.plans, + .host_hit_picks = host.hit_picks, + .host_miss_picks = host.miss_picks, + .host_mandatory_copies = host.copies, + .host_prefetch_copies = resident.stats.prefetches, + .pool_dma_copies = resident.lines.stats.copies, + .pool_dma_requested_bytes = resident.lines.stats.bytes, + .device_plans = device.plans, + .device_hit_picks = device.hit_picks, + .device_miss_picks = device.miss_picks, + .device_copies = device.copies, + .device_inplace_experts = device.inplace_experts, + .device_overflow_plans = device.overflow_plans, + .mapped_copy_payload_min_bytes = try std.math.mul(u64, device.copies, minimum), + .mapped_copy_payload_max_bytes = try std.math.mul(u64, device.copies, maximum), + }; + try counters.add(counters.Experts, &totals, contribution); + if (engine.telemetry.layers_enabled) { + var layer_sample: counters.LayerSample = .{ .owner = resident.telemetry_id, .kind = .decoder, .layer = 0, .lines = resident.opt.lines, .counts = counters.LayerCounts.fromExperts(contribution) }; + for (engine.layers[0..engine.n_layers], 0..) |layer, decoder_index| if (layer.exl3) |decoder| { + if (decoder.pool == resident) { + layer_sample.layer = decoder_index; + break; + } + }; + layers[seen_count - 1] = layer_sample; + } + } + for (sequences, 0..) |sequence, index| for (sequences[0..index]) |previous| if (previous == sequence) return error.TelemetryDuplicateOwner; + const sampled_layers = if (engine.telemetry.layers_enabled) try allocator.realloc(layers, seen_count) else layers; + return .{ .rank = rank, .scope = hash.final(), .experts = totals, .expert_pools = seen_count, .layers = sampled_layers }; +} + +test "exl3: cache telemetry pure host contracts" { + _ = counters; + _ = @import("cuda_exl3_telemetry_test.zig"); + _ = @import("../../cuda_prompt_telemetry.zig"); +} diff --git a/zig/src/families/flashnext/exl3/model/cuda_exl3_telemetry_test.zig b/zig/src/families/flashnext/exl3/model/cuda_exl3_telemetry_test.zig new file mode 100644 index 0000000000..69e0ded0c5 --- /dev/null +++ b/zig/src/families/flashnext/exl3/model/cuda_exl3_telemetry_test.zig @@ -0,0 +1,38 @@ +const std = @import("std"); +const cuda = @import("cuda"); +const pool = @import("../offload/cuda_exl3_pool.zig"); +const telemetry = @import("exl3_telemetry.zig"); + +test "exl3: aggregate telemetry finds attached pools deduplicates pool aliases and separates committed device plans" { + var resident: pool.Pool = undefined; + resident.lines.pinned_bank = null; + resident.telemetry_id = 991; + resident.telemetry_valid = true; + resident.telemetry_device = .{ .plans = 5, .hit_picks = 20, .miss_picks = 10, .copies = 3, .inplace_experts = 1 }; + resident.stats = .{ .rounds = 7, .hit_slots = 25, .miss_slots = 13, .misses = 5 }; + resident.lines.stats = .{ .copies = 2, .bytes = 96 }; + const payload: [16]u8 = @splat(0); + var sources = [_][3][]const u8{.{ &payload, &payload, &payload }}; + var widths = [_][3]i32{.{ 1, 1, 1 }}; + resident.src = .{ .routed = &sources }; + resident.k2 = &widths; + var experts: @import("../pack/cuda_exl3_load.zig").Experts = undefined; + experts.pool = &resident; + var engine: @import("../../cuda_prompt.zig").Engine = undefined; + engine.telemetry = .{}; + engine.n_layers = 1; + engine.layers[0].exl3 = &experts; + const collector = @import("cuda_exl3_telemetry.zig"); + var snapshot = try collector.sampleAfterConsumers(std.testing.allocator, &engine, &.{}, .{ .rank = 0, .world = 1 }, .all_rank_consumers_joined); + defer snapshot.deinit(std.testing.allocator); + try std.testing.expectEqual(@as(usize, 1), snapshot.expert_pools); + try std.testing.expectEqual(@as(u64, 2), snapshot.experts.host_plans); + try std.testing.expectEqual(@as(u64, 5), snapshot.experts.host_hit_picks); + try std.testing.expectEqual(@as(u64, 3), snapshot.experts.host_miss_picks); + try std.testing.expectEqual(@as(u64, 2), snapshot.experts.host_mandatory_copies); + try std.testing.expectEqual(@as(u64, 5), snapshot.experts.device_plans); + try std.testing.expectEqual(@as(u64, 1), snapshot.experts.device_inplace_experts); + try std.testing.expectEqual(@as(u64, 144), snapshot.experts.mapped_copy_payload_min_bytes); + try std.testing.expectEqual(@as(u64, 144), snapshot.experts.mapped_copy_payload_max_bytes); + try std.testing.expectEqual(@as(u64, 96), snapshot.experts.pool_dma_requested_bytes); +} diff --git a/zig/src/families/flashnext/exl3/model/exl3.zig b/zig/src/families/flashnext/exl3/model/exl3.zig new file mode 100644 index 0000000000..73e5d75ec4 --- /dev/null +++ b/zig/src/families/flashnext/exl3/model/exl3.zig @@ -0,0 +1,33 @@ +//! Host-side EXL3 index for Qwen3.8-Flash-Next: the float64 oracles and header rules kept as tests under option C + +pub const shape = @import("../pack/exl3_shape.zig"); +pub const scan = @import("../pack/exl3_scan.zig"); +pub const cache = @import("exl3_cache.zig"); +pub const cpu = @import("exl3_cpu.zig"); +pub const place = @import("../pack/exl3_place.zig"); +pub const route = @import("exl3_route.zig"); +pub const router = @import("exl3_router.zig"); +pub const kv = @import("../kv/exl3_kv.zig"); +pub const norm = @import("exl3_norm.zig"); +pub const hc = @import("exl3_hc.zig"); +pub const moe = @import("exl3_moe.zig"); +pub const gdn_cpu = @import("exl3_gdn_cpu.zig"); +pub const attn_cpu = @import("exl3_attn_cpu.zig"); +pub const fwd_cpu = @import("exl3_fwd_cpu.zig"); + +test { + _ = shape; + _ = scan; + _ = cache; + _ = cpu; + _ = place; + _ = route; + _ = router; + _ = kv; + _ = norm; + _ = hc; + _ = moe; + _ = gdn_cpu; + _ = attn_cpu; + _ = fwd_cpu; +} diff --git a/zig/src/families/flashnext/exl3/model/exl3_attn_cpu.zig b/zig/src/families/flashnext/exl3/model/exl3_attn_cpu.zig new file mode 100644 index 0000000000..f276d4c6f0 --- /dev/null +++ b/zig/src/families/flashnext/exl3/model/exl3_attn_cpu.zig @@ -0,0 +1,441 @@ +//! CPU oracle for one Qwen3.8-Flash-Next full-attention layer (layers 3, 7, ... 47). + +const std = @import("std"); +const shape = @import("../pack/exl3_shape.zig"); + +pub const q_heads: usize = shape.heads; +pub const kv_heads: usize = shape.kv_heads; +pub const head_dim: usize = shape.head_dim; +pub const rotary_dim: usize = 64; +pub const half: usize = rotary_dim / 2; +pub const index_heads: usize = 4; +pub const index_dim: usize = 128; +pub const ratio: usize = 4; +pub const top_blocks: usize = 512; +pub const theta: f64 = 1.0e7; +pub const eps: f64 = 1.0e-6; +/// Most keys one row ever reads: 512 selected blocks and a 3-key tail, or a dense row of up to 2051 keys. +pub const max_keys: usize = top_blocks * ratio + ratio; +pub const index_width: usize = (index_heads + 1) * index_dim; + +pub fn roundBf16(x: f64) f64 { + const f: f32 = @floatCast(x); + var u: u32 = @bitCast(f); + if (!std.math.isFinite(f)) return f; + // round to nearest even on the low 16 bits + const lsb = (u >> 16) & 1; + u +%= 0x7fff + lsb; + u &= 0xffff0000; + return @as(f32, @bitCast(u)); +} + +pub fn roundF16(x: f64) f64 { + const h: f16 = @floatCast(x); + return h; +} + +/// theta^(-2 i / rotary_dim) as the f32 table the kernels read. +pub fn invFreq() [half]f32 { + var t: [half]f32 = undefined; + for (&t, 0..) |*v, i| v.* = @floatCast(std.math.pow(f64, theta, -2.0 * @as(f64, @floatFromInt(i)) / @as(f64, rotary_dim))); + return t; +} + +pub const Norms = struct { + /// Scales with the centred offset already added: `w + offset`. + q: []const f32, + k: []const f32, + iq: []const f32, + ik: []const f32, + inv: [half]f32, +}; + +/// fp32 RMSNorm, bf16, rotate-half RoPE over the first 64 dims at `pos`, bf16. `x` and `out` are `width` long. +pub fn normRope(x: []const f64, scale: []const f32, pos: f64, inv: *const [half]f32, out: []f64) void { + const width = x.len; + var ss: f64 = 0; + for (x) |v| ss += v * v; + const rinv = 1.0 / @sqrt(ss / @as(f64, @floatFromInt(width)) + eps); + var xn: [head_dim]f64 = undefined; + for (x, 0..) |v, i| xn[i] = roundBf16(v * rinv * @as(f64, scale[i])); + const posf: f32 = @floatCast(pos); + for (0..width) |d| { + if (d < half) { + const ang: f32 = posf * inv[d]; + out[d] = roundBf16(xn[d] * @cos(@as(f64, ang)) - xn[d + half] * @sin(@as(f64, ang))); + } else if (d < 2 * half) { + const ang: f32 = posf * inv[d - half]; + out[d] = roundBf16(xn[d - half] * @sin(@as(f64, ang)) + xn[d] * @cos(@as(f64, ang))); + } else out[d] = xn[d]; + } +} + +/// How many of the 32 frequency indices each mrope axis owns with `mrope_interleaved` sections [11, 11, 10]. +pub fn mropeAxisCounts() [3]usize { + var c: [3]usize = .{ 0, 0, 0 }; + for (0..half) |i| { + if (i % 3 == 1 and i < 3 * 11) { + c[1] += 1; + } else if (i % 3 == 2 and i < 3 * 10) { + c[2] += 1; + } else c[0] += 1; + } + return c; +} + +/// One full-attention layer's state: bf16-rounded keys and values, raw index keys, pooled blocks. +pub const Cache = struct { + gpa: std.mem.Allocator, + tokens: usize, + keys: []f64, // [pos][kv head][256] + values: []f64, + raw: []f64, // [pos][128] + pooled: []f64, // [pos / 4][128] + + pub fn init(gpa: std.mem.Allocator, tokens: usize) !Cache { + const keys = try gpa.alloc(f64, tokens * kv_heads * head_dim); + errdefer gpa.free(keys); + const values = try gpa.alloc(f64, tokens * kv_heads * head_dim); + errdefer gpa.free(values); + const raw = try gpa.alloc(f64, tokens * index_dim); + errdefer gpa.free(raw); + const pooled = try gpa.alloc(f64, (tokens / ratio + 1) * index_dim); + return .{ .gpa = gpa, .tokens = tokens, .keys = keys, .values = values, .raw = raw, .pooled = pooled }; + } + + pub fn deinit(c: *Cache) void { + c.gpa.free(c.keys); + c.gpa.free(c.values); + c.gpa.free(c.raw); + c.gpa.free(c.pooled); + } +}; + +/// Projections of one token (what the EXL3 linears emit, fp16): q_proj [24 * 512], k_proj [512], v_proj [512]. +pub const Proj = struct { + qp: []const f16, + kp: []const f16, + vp: []const f16, + ip: []const f16, +}; + +/// The normed, rotated queries (bf16 values, [24 * 256]) and indexer queries ([4 * 128]) of one token. +pub fn prepare(c: *Cache, pos: usize, p: Proj, n: *const Norms, q_out: []f64, iq_out: []f64) void { + var buf: [head_dim]f64 = undefined; + for (0..q_heads) |h| { + for (0..head_dim) |d| buf[d] = p.qp[h * 2 * head_dim + d]; + normRope(buf[0..head_dim], n.q, @floatFromInt(pos), &n.inv, q_out[h * head_dim ..][0..head_dim]); + } + for (0..kv_heads) |j| { + for (0..head_dim) |d| buf[d] = p.kp[j * head_dim + d]; + normRope(buf[0..head_dim], n.k, @floatFromInt(pos), &n.inv, c.keys[(pos * kv_heads + j) * head_dim ..][0..head_dim]); + for (0..head_dim) |d| c.values[(pos * kv_heads + j) * head_dim + d] = roundBf16(p.vp[j * head_dim + d]); + } + for (0..index_heads) |h| { + for (0..index_dim) |d| buf[d] = p.ip[h * index_dim + d]; + normRope(buf[0..index_dim], n.iq, @floatFromInt(pos), &n.inv, iq_out[h * index_dim ..][0..index_dim]); + } + for (0..index_dim) |d| c.raw[pos * index_dim + d] = roundBf16(p.ip[index_heads * index_dim + d]); +} + +/// Pool block `b` (keys 4b .. 4b + 3 must be in `raw`). +pub fn poolBlock(c: *Cache, b: usize, n: *const Norms) void { + var mean: [index_dim]f64 = undefined; + for (0..index_dim) |d| { + var acc: f32 = @floatCast(c.raw[(ratio * b) * index_dim + d]); + for (1..ratio) |k| acc = acc + @as(f32, @floatCast(c.raw[(ratio * b + k) * index_dim + d])); + mean[d] = roundBf16(@as(f64, acc / @as(f32, ratio))); + } + normRope(&mean, n.ik, @floatFromInt(ratio * b), &n.inv, c.pooled[b * index_dim ..][0..index_dim]); +} + +/// Score of every complete block for one indexer query: sum_h relu(q_h . pooled_b) / sqrt(128). +pub fn blockScores(c: *const Cache, iq: []const f64, complete: usize, out: []f64) void { + for (0..complete) |b| { + var total: f64 = 0; + for (0..index_heads) |h| { + var dot: f64 = 0; + for (0..index_dim) |d| dot += iq[h * index_dim + d] * c.pooled[b * index_dim + d]; + total += @max(dot, 0); + } + out[b] = total / @sqrt(@as(f64, index_dim)); + } +} + +/// The `top_blocks` best blocks, ties to the lower id, in ascending block order. +pub fn selectBlocks(gpa: std.mem.Allocator, scores: []const f64, out: *[top_blocks]usize) !void { + const order = try gpa.alloc(usize, scores.len); + defer gpa.free(order); + for (order, 0..) |*o, i| o.* = i; + const Ctx = struct { + s: []const f64, + fn less(ctx: @This(), a: usize, b: usize) bool { + if (ctx.s[a] != ctx.s[b]) return ctx.s[a] > ctx.s[b]; + return a < b; + } + }; + std.mem.sort(usize, order, Ctx{ .s = scores }, Ctx.less); + @memcpy(out, order[0..top_blocks]); + std.mem.sort(usize, out, {}, std.sort.asc(usize)); +} + +/// Keys a row at `end = pos + 1` reads. +pub fn keyList(end: usize, blocks: ?*const [top_blocks]usize, out: *[max_keys]usize) usize { + const complete = end / ratio; + if (complete <= top_blocks) { + for (0..end) |i| out[i] = i; + return end; + } + const sel = blocks.?; + var n: usize = 0; + for (sel) |b| { + for (0..ratio) |k| { + out[n] = b * ratio + k; + n += 1; + } + } + var t = ratio * complete; + while (t < end) : (t += 1) { + out[n] = t; + n += 1; + } + return n; +} + +/// out[h * 256 + d] = fp16(bf16(bf16(o) / (1 + exp(-gate)))), o = softmax(q . K / 16) V over `keys`. +pub fn attend(c: *const Cache, q: []const f64, qp: []const f16, keys: []const usize, out: []f64) void { + var s: [max_keys]f64 = undefined; + const group = q_heads / kv_heads; + for (0..q_heads) |h| { + const kvh = h / group; + var m: f64 = -std.math.inf(f64); + for (keys, 0..) |key, j| { + var dot: f64 = 0; + for (0..head_dim) |d| dot += q[h * head_dim + d] * c.keys[(key * kv_heads + kvh) * head_dim + d]; + s[j] = dot * 0.0625; + m = @max(m, s[j]); + } + var l: f64 = 0; + for (keys, 0..) |_, j| { + s[j] = @exp(s[j] - m); + l += s[j]; + } + for (0..head_dim) |d| { + var acc: f64 = 0; + for (keys, 0..) |key, j| acc += s[j] * c.values[(key * kv_heads + kvh) * head_dim + d]; + const o = roundBf16(acc / l); + const g: f64 = qp[h * 2 * head_dim + head_dim + d]; + out[h * head_dim + d] = roundF16(roundBf16(o / (1.0 + @exp(-g)))); + } + } +} + +pub const RowInfo = struct { + sparse: bool = false, + nk: usize = 0, + /// The selected blocks of a sparse row (undefined when dense). + blocks: [top_blocks]usize = undefined, +}; + +/// The whole row: prepare, pool the blocks this token completes, select, attend. +pub fn row( + c: *Cache, + pos: usize, + p: Proj, + n: *const Norms, + scores: []f64, + q: []f64, + iq: []f64, + out: []f64, + info: *RowInfo, +) !void { + prepare(c, pos, p, n, q, iq); + const end = pos + 1; + if (end % ratio == 0) poolBlock(c, end / ratio - 1, n); + const complete = end / ratio; + info.sparse = complete > top_blocks; + if (info.sparse) { + blockScores(c, iq, complete, scores); + try selectBlocks(c.gpa, scores[0..complete], &info.blocks); + } + var keys: [max_keys]usize = undefined; + info.nk = keyList(end, if (info.sparse) &info.blocks else null, &keys); + attend(c, q, p.qp, keys[0..info.nk], out); +} + +// -- tests ----------------------------------------------------------------- + +test "mrope sections [11, 11, 10] interleave over 32 frequencies and need no per-axis table for text" { + const counts = mropeAxisCounts(); + try std.testing.expectEqual([3]usize{ 11, 11, 10 }, counts); + try std.testing.expectEqual(@as(usize, 32), counts[0] + counts[1] + counts[2]); + try std.testing.expectEqual(half, counts[0] + counts[1] + counts[2]); +} + +test "rotate-half RoPE matches an explicit rotation of the pairs (i, i + 32) and leaves dims 64.. alone" { + const inv = invFreq(); + try std.testing.expectApproxEqRel(@as(f32, 1.0), inv[0], 1e-7); + // theta^(-2*31/64) for the last frequency + try std.testing.expectApproxEqRel(@as(f32, @floatCast(std.math.pow(f64, 1e7, -62.0 / 64.0))), inv[31], 1e-6); + var x: [head_dim]f64 = undefined; + var rng = std.Random.DefaultPrng.init(7); + for (&x) |*v| v.* = rng.random().float(f64) * 2 - 1; + var scale: [head_dim]f32 = @splat(1); + scale[3] = 2.5; + var out: [head_dim]f64 = undefined; + const pos: usize = 1234; + normRope(&x, &scale, @floatFromInt(pos), &inv, &out); + + var ss: f64 = 0; + for (x) |v| ss += v * v; + const rinv = 1.0 / @sqrt(ss / head_dim + eps); + for (0..half) |i| { + const a = roundBf16(x[i] * rinv * scale[i]); + const b = roundBf16(x[i + half] * rinv * scale[i + half]); + const ang: f32 = @as(f32, @floatFromInt(pos)) * inv[i]; + try std.testing.expectEqual(roundBf16(a * @cos(@as(f64, ang)) - b * @sin(@as(f64, ang))), out[i]); + try std.testing.expectEqual(roundBf16(a * @sin(@as(f64, ang)) + b * @cos(@as(f64, ang))), out[i + half]); + } + for (rotary_dim..head_dim) |d| try std.testing.expectEqual(roundBf16(x[d] * rinv * scale[d]), out[d]); + // position 0 is the identity rotation + normRope(&x, &scale, 0, &inv, &out); + for (0..head_dim) |d| try std.testing.expectEqual(roundBf16(x[d] * rinv * scale[d]), out[d]); +} + +test "bf16 rounding is round-to-nearest-even" { + try std.testing.expectEqual(@as(f64, 1.0), roundBf16(1.0 + 1.0 / 512.0)); // below half an ulp (ulp 2^-7) + try std.testing.expectEqual(@as(f64, 1.0), roundBf16(1.0 + 1.0 / 256.0)); // tie goes to even (1.0) + try std.testing.expectEqual(@as(f64, 1.0 + 1.0 / 64.0), roundBf16(1.0 + 3.0 / 256.0)); // tie goes up to even mantissa + try std.testing.expectEqual(@as(f64, 1.0 + 1.0 / 128.0), roundBf16(1.0 + 1.0 / 128.0 + 1.0 / 1024.0)); +} + +test "selection takes the 512 best blocks, ties to the lower id, in block order" { + const gpa = std.testing.allocator; + const n: usize = 700; + const scores = try gpa.alloc(f64, n); + defer gpa.free(scores); + var rng = std.Random.DefaultPrng.init(11); + for (scores) |*v| v.* = @floor(rng.random().float(f64) * 40) / 4; // many exact ties + var sel: [top_blocks]usize = undefined; + try selectBlocks(gpa, scores, &sel); + try std.testing.expect(std.sort.isSorted(usize, &sel, {}, std.sort.asc(usize))); + // the cut: every unselected block scores no more than every selected one, and a tie at the cut keeps the lower id + var chosen: [700]bool = @splat(false); + var min_sel: f64 = std.math.inf(f64); + for (sel) |b| { + chosen[b] = true; + min_sel = @min(min_sel, scores[b]); + } + var count: usize = 0; + for (chosen) |c| count += @intFromBool(c); + try std.testing.expectEqual(top_blocks, count); + var max_out: f64 = -1; + var lowest_out_at_cut: usize = n; + for (0..n) |b| if (!chosen[b]) { + max_out = @max(max_out, scores[b]); + if (scores[b] == min_sel) lowest_out_at_cut = @min(lowest_out_at_cut, b); + }; + try std.testing.expect(max_out <= min_sel); + for (0..n) |b| if (chosen[b] and scores[b] == min_sel) try std.testing.expect(b < lowest_out_at_cut); +} + +test "key list: dense through 2051 keys, then 2048 selected plus the unfinished tail" { + var sel: [top_blocks]usize = undefined; + for (&sel, 0..) |*b, i| b.* = i + 1; // blocks 1..512 of 513 + var keys: [max_keys]usize = undefined; + try std.testing.expectEqual(@as(usize, 2051), keyList(2051, null, &keys)); + try std.testing.expectEqual(@as(usize, 2050), keys[2050]); + // end 2052: 513 complete blocks, no tail, block 0 dropped + try std.testing.expectEqual(@as(usize, 2048), keyList(2052, &sel, &keys)); + try std.testing.expectEqual(@as(usize, 4), keys[0]); + try std.testing.expectEqual(@as(usize, 2051), keys[2047]); + // end 2055: three tail keys 2052, 2053, 2054 + try std.testing.expectEqual(@as(usize, 2051), keyList(2055, &sel, &keys)); + try std.testing.expectEqual(@as(usize, 2052), keys[2048]); + try std.testing.expectEqual(@as(usize, 2054), keys[2050]); +} + +test "attend with one visible key returns that value gated" { + const gpa = std.testing.allocator; + var c = try Cache.init(gpa, 4); + defer c.deinit(); + for (0..kv_heads * head_dim) |i| { + c.keys[i] = 0.25; + c.values[i] = if (i < head_dim) 2.0 else -3.0; + } + var q: [q_heads * head_dim]f64 = @splat(0.5); + var qp: [q_heads * 2 * head_dim]f16 = @splat(0); + for (0..q_heads) |h| for (0..head_dim) |d| { + qp[h * 2 * head_dim + head_dim + d] = 0; // sigmoid(0) = 1/2 + }; + var out: [q_heads * head_dim]f64 = undefined; + attend(&c, &q, &qp, &.{0}, &out); + for (0..12) |h| try std.testing.expectEqual(@as(f64, 1.0), out[h * head_dim]); + for (12..q_heads) |h| try std.testing.expectEqual(@as(f64, -1.5), out[h * head_dim]); +} + +test "the device KV layout is exl3_kv.Table's for one layer" { + const kv = @import("../kv/exl3_kv.zig"); + const gpa = std.testing.allocator; + const tokens: usize = 8; + var t = try kv.Table.init(gpa, tokens); + defer t.deinit(gpa); + try std.testing.expectEqual(@as(usize, kv_heads), t.heads); + const slot_bytes = 2 * head_dim * 2; + for ([_]usize{ 0, 1, 2 }) |layer| for ([_]usize{ 0, 5, 7 }) |token| for ([_]usize{ 0, 1 }) |head| { + var src: [slot_bytes]u8 = undefined; + @memset(&src, @intCast(1 + layer * 40 + token * 4 + head)); + try t.write(head, layer, token, &src); + // the kernel's byte offset: layer base + (token * kv_heads + head) * slot + const off = layer * tokens * kv_heads * slot_bytes + (token * kv_heads + head) * slot_bytes; + try std.testing.expectEqualSlices(u8, &src, t.bytes[off..][0..slot_bytes]); + }; +} + +test "a full row past 2048 keys reads the selected blocks, and the dense row at 2051 keys does not" { + const gpa = std.testing.allocator; + const tokens: usize = 2100; + var c = try Cache.init(gpa, tokens); + defer c.deinit(); + var rng = std.Random.DefaultPrng.init(3); + const inv = invFreq(); + const q_scale: [head_dim]f32 = @splat(1); + const i_scale: [index_dim]f32 = @splat(1); + const norms = Norms{ .q = &q_scale, .k = &q_scale, .iq = &i_scale, .ik = &i_scale, .inv = inv }; + const qp = try gpa.alloc(f16, q_heads * 2 * head_dim); + defer gpa.free(qp); + const kp = try gpa.alloc(f16, kv_heads * head_dim); + defer gpa.free(kp); + const vp = try gpa.alloc(f16, kv_heads * head_dim); + defer gpa.free(vp); + const ip = try gpa.alloc(f16, index_width); + defer gpa.free(ip); + const scores = try gpa.alloc(f64, tokens / ratio + 1); + defer gpa.free(scores); + var q: [q_heads * head_dim]f64 = undefined; + var iq: [index_heads * index_dim]f64 = undefined; + var out: [q_heads * head_dim]f64 = undefined; + var info = RowInfo{}; + for (0..2056) |pos| { + for (qp) |*v| v.* = @floatCast(rng.random().floatNorm(f32)); + for (kp) |*v| v.* = @floatCast(rng.random().floatNorm(f32)); + for (vp) |*v| v.* = @floatCast(rng.random().floatNorm(f32)); + for (ip) |*v| v.* = @floatCast(rng.random().floatNorm(f32)); + const proj = Proj{ .qp = qp, .kp = kp, .vp = vp, .ip = ip }; + switch (pos) { + 2050, 2051, 2054 => { + try row(&c, pos, proj, &norms, scores, &q, &iq, &out, &info); + switch (pos) { + 2050 => try std.testing.expect(!info.sparse and info.nk == 2051), + 2051 => try std.testing.expect(info.sparse and info.nk == 2048), + else => try std.testing.expect(info.sparse and info.nk == 2051), + } + }, + else => { + prepare(&c, pos, proj, &norms, &q, &iq); + if ((pos + 1) % ratio == 0) poolBlock(&c, (pos + 1) / ratio - 1, &norms); + }, + } + } + for (out) |v| try std.testing.expect(std.math.isFinite(v)); +} diff --git a/zig/src/families/flashnext/exl3/model/exl3_cache.zig b/zig/src/families/flashnext/exl3/model/exl3_cache.zig new file mode 100644 index 0000000000..c193eefa6f --- /dev/null +++ b/zig/src/families/flashnext/exl3/model/exl3_cache.zig @@ -0,0 +1,168 @@ +//! The KV bytes and the resident-set budget. + +const std = @import("std"); +const shape = @import("../pack/exl3_shape.zig"); + +/// One attention layer's K and V for `tokens` tokens, bf16. +pub fn layerBytes(tokens: usize, kv_heads: usize, head_dim: usize) u64 { + const t: u64 = @intCast(tokens); + const h: u64 = @intCast(kv_heads); + const d: u64 = @intCast(head_dim); + return t * 2 * h * d * 2; +} + +/// A kept prefix stops one token before the prompt end. An empty prompt stays at 0. +pub fn resumeAt(prompt_len: u32) u32 { + if (prompt_len == 0) return 0; + return prompt_len - 1; +} + +pub const Budget = struct { + device_bytes: u64, + /// Dense + shared + router + HC + MTP workspace, already summed by the caller. + resident_bytes: u64, + kv_bytes: u64, + routed_expert_bytes: u64, + ngram_bytes: u64, +}; + +pub const Place = enum { host, refuse }; + +/// Resident weights plus KV. Routed experts and the n-gram table are not in this sum. +pub fn deviceSum(b: Budget) u64 { + return b.resident_bytes + b.kv_bytes; +} + +pub fn fits(b: Budget) error{ExceedsDevice}!void { + if (deviceSum(b) > b.device_bytes) return error.ExceedsDevice; +} + +/// One RTX 5060 Ti. The learned expert lines are whatever this does not already use. +pub const card_bytes: u64 = 16 << 30; +/// CUDA context and the batch-1 workspace. Not a weight, and not a second copy of the bank. +pub const device_reserve: u64 = 256 << 20; + +/// Bytes left for routed experts after pinned GEMMs, KV, and the reserve. +pub fn learnedBytes(device: u64, pinned: u64, kv: u64) u64 { + if (pinned >= device) return 0; + const after_pinned = device - pinned; + if (kv >= after_pinned) return 0; + const after_kv = after_pinned - kv; + if (after_kv <= device_reserve) return 0; + return after_kv - device_reserve; +} + +fn matrixBytes(k: usize, n: usize) u64 { + const kn: u64 = @as(u64, k) * @as(u64, n); + // 3.05 bpw trellis, plus one fp16 scale along K and one along N. + return kn * 3050 / 8000 + @as(u64, k + n) * 2; +} + +/// One routed expert as stored. +pub fn expertBytes() u64 { + const k = shape.hidden; + const n = shape.expert_intermediate; + return matrixBytes(k, n) * 2 + matrixBytes(n, k); +} + +pub fn expertLines(room: u64) u64 { + return expertLinesOf(room, expertBytes()); +} + +/// Budget stride for one owned expert until a header span is larger. Not a line count. +pub const full_expert_bytes: u64 = 1_901_213; + +/// Lines that fit in `room` at `line` bytes. A zero stride is no lines, not a panic. +pub fn expertLinesOf(room: u64, line: u64) u64 { + if (line == 0) return 0; + return room / line; +} + +/// KV for every full-attention layer. +pub fn contextKv(tokens: usize) u64 { + const layers: u64 = shape.layers / shape.full_attention_interval; + return layerBytes(tokens, shape.kv_heads, shape.head_dim) * layers; +} + +/// Routed experts and the n-gram table always live on the host. +pub fn placeWeights(b: Budget) Place { + _ = b; + return .host; +} + +/// Neither the routed-expert bank nor the n-gram table may be counted as resident. +pub fn refusesResident(which: enum { routed, ngram }) bool { + _ = which; + return true; +} + +test "resumeAt stops one token before the end" { + try std.testing.expectEqual(@as(u32, 0), resumeAt(0)); + try std.testing.expectEqual(@as(u32, 7), resumeAt(8)); + try std.testing.expectEqual(@as(u32, 0), resumeAt(1)); +} + +test "layer bytes match v0.6.6 row_bytes" { + try std.testing.expectEqual(@as(u64, 2048), layerBytes(1, 2, 256)); + try std.testing.expectEqual(@as(u64, 2048), layerBytes(1, shape.kv_heads, shape.head_dim)); + try std.testing.expectEqual(@as(u64, 6144), layerBytes(3, 2, 256)); + try std.testing.expectEqual(@as(u64, 192), layerBytes(1, 2, 24)); +} + +test "device sum is resident plus kv, and a 16 GiB card refuses the overflow" { + const device: u64 = 16 << 30; + const resident: u64 = 6 << 30; + const kv: u64 = 8 << 30; + const held = Budget{ + .device_bytes = device, + .resident_bytes = resident, + .kv_bytes = kv, + .routed_expert_bytes = 70 << 30, + .ngram_bytes = 29 << 30, + }; + try std.testing.expectEqual(resident + kv, deviceSum(held)); + try fits(held); + try std.testing.expectEqual(Place.host, placeWeights(held)); + try std.testing.expect(refusesResident(.routed)); + try std.testing.expect(refusesResident(.ngram)); + + try std.testing.expectError(error.ExceedsDevice, fits(.{ + .device_bytes = device, + .resident_bytes = resident, + .kv_bytes = 11 << 30, + .routed_expert_bytes = 0, + .ngram_bytes = 0, + })); + + const rest = device - resident; + try fits(.{ + .device_bytes = device, + .resident_bytes = resident, + .kv_bytes = rest, + .routed_expert_bytes = 1, + .ngram_bytes = 1, + }); + try std.testing.expectError(error.ExceedsDevice, fits(.{ + .device_bytes = device, + .resident_bytes = resident, + .kv_bytes = rest + 1, + .routed_expert_bytes = 0, + .ngram_bytes = 0, + })); + + try std.testing.expectEqual(@as(u64, 1_893_120), expertBytes()); + try std.testing.expectEqual(@as(u64, 1_901_213), full_expert_bytes); + try std.testing.expectEqual(@as(u64, 0), expertLinesOf(1, 0)); + try std.testing.expectEqual(@as(u64, 1), expertLinesOf(full_expert_bytes, full_expert_bytes)); + + try std.testing.expectEqual(@as(u64, 0), learnedBytes(card_bytes, card_bytes, 0)); + + // Counting the expert bank as resident is what blows the same card. + try std.testing.expectError(error.ExceedsDevice, fits(.{ + .device_bytes = device, + .resident_bytes = resident + (70 << 30), + .kv_bytes = kv, + .routed_expert_bytes = 0, + .ngram_bytes = 0, + })); +} diff --git a/zig/src/families/flashnext/exl3/model/exl3_cpu.zig b/zig/src/families/flashnext/exl3/model/exl3_cpu.zig new file mode 100644 index 0000000000..a40d6f5633 --- /dev/null +++ b/zig/src/families/flashnext/exl3/model/exl3_cpu.zig @@ -0,0 +1,380 @@ +//! CPU GEMV for one EXL3 matrix. + +const std = @import("std"); +const builtin = @import("builtin"); +const format = @import("../pack/exl3_format.zig"); + +/// AVX fp256, compiled only when this target actually has AVX2. No AVX-512. +const avx2_target = builtin.cpu.arch == .x86_64 and + builtin.zig_backend != .stage2_x86_64 and + builtin.cpu.has(.x86, .avx2); + +const Acc = @Vector(4, f64); + +/// fp32 SwiGLU from the v0.6.6 gate-up epilogue (`act_mode` 1). +pub fn swiglu(gate: []const f64, up: []const f64, out: []f64, limit: f64) error{Length}!void { + if (gate.len != up.len or gate.len != out.len) return error.Length; + for (gate, up, out) |g, u, *y| { + const gg = @min(g, limit); + const uu = @min(@max(u, -limit), limit); + y.* = (gg / (1.0 + @exp(-gg))) * uu; + } +} + +/// Weighted sum of `weights.len` expert rows. +pub fn combine(rows: []const f64, weights: []const f64, n: usize, out: []f64) error{Length}!void { + if (n == 0 or out.len != n or weights.len == 0 or rows.len != weights.len * n) return error.Length; + @memset(out, 0); + for (weights, 0..) |w, s| { + const row = rows[s * n ..][0..n]; + for (out, row) |*y, v| y.* += w * v; + } +} + +/// K = tiles_k * 16 and N = tiles_n * 16, both multiples of 128 (one Hadamard block). +pub fn gemv( + w: format.Width, + codebook: format.Codebook, + trellis: []const i16, + tiles_k: usize, + tiles_n: usize, + suh: []const f16, + svh: []const f16, + x: []const f64, + y: []f64, + scratch: []f64, +) error{Length}!void { + const k = tiles_k * 16; + const n = tiles_n * 16; + if (k % format.hadamard_dim != 0 or n % format.hadamard_dim != 0) return error.Length; + if (x.len != k or y.len != n or suh.len != k or svh.len != n) return error.Length; + const cells = k * n; + if (scratch.len < 3 * cells) return error.Length; + + const weight = scratch[0..cells]; + format.dequantize(w, codebook, trellis, tiles_k, tiles_n, suh, svh, weight, scratch[cells .. 3 * cells]); + + for (0..n) |j| { + var sum: f64 = 0; + for (0..k) |i| sum += x[i] * weight[i * n + j]; + y[j] = sum; + } +} + +/// Bit-exact with `gemv`. +pub fn gemvAvx2( + w: format.Width, + codebook: format.Codebook, + trellis: []const i16, + tiles_k: usize, + tiles_n: usize, + suh: []const f16, + svh: []const f16, + x: []const f64, + y: []f64, + scratch: []f64, +) error{Length}!void { + const k = tiles_k * 16; + const n = tiles_n * 16; + if (k % format.hadamard_dim != 0 or n % format.hadamard_dim != 0) return error.Length; + if (x.len != k or y.len != n or suh.len != k or svh.len != n) return error.Length; + const cells = k * n; + if (scratch.len < 3 * cells) return error.Length; + + const weight = scratch[0..cells]; + format.dequantize(w, codebook, trellis, tiles_k, tiles_n, suh, svh, weight, scratch[cells .. 3 * cells]); + + // N is a multiple of 128, so the four-column walk covers every column. + if (comptime avx2_target) { + accumulateAvx2(weight, x, y); + } else { + accumulatePortable(weight, x, y); + } +} + +/// Same four lanes as the AVX2 loop, one left-to-right sum per column. +fn accumulatePortable(weight: []const f64, x: []const f64, y: []f64) void { + @setFloatMode(.strict); + const k = x.len; + const n = y.len; + var j: usize = 0; + while (j < n) : (j += 4) { + var acc: [4]f64 = .{ 0, 0, 0, 0 }; + var i: usize = 0; + while (i < k) : (i += 1) { + const xi = x[i]; + const row = i * n + j; + inline for (0..4) |lane| acc[lane] += xi * weight[row + lane]; + } + y[j..][0..4].* = acc; + } +} + +fn avx2Load(p: [*]const f64) Acc { + return asm ( + \\vmovupd (%[p]), %[out] + : [out] "=x" (-> Acc), + : [p] "r" (p), + ); +} + +fn avx2Broadcast(p: *const f64) Acc { + return asm ( + \\vbroadcastsd (%[p]), %[out] + : [out] "=x" (-> Acc), + : [p] "r" (p), + ); +} + +fn avx2Mul(a: Acc, b: Acc) Acc { + return asm ( + \\vmulpd %[b], %[a], %[out] + : [out] "=x" (-> Acc), + : [a] "x" (a), + [b] "x" (b), + ); +} + +/// `out = sum + prod`, the same operand order as `sum += prod`. +fn avx2Add(sum: Acc, prod: Acc) Acc { + return asm ( + \\vaddpd %[prod], %[sum], %[out] + : [out] "=x" (-> Acc), + : [sum] "x" (sum), + [prod] "x" (prod), + ); +} + +fn avx2Store(p: [*]f64, v: Acc) void { + asm volatile ( + \\vmovupd %[v], (%[p]) + : + : [v] "x" (v), + [p] "r" (p), + : .{ .memory = true }); +} + +/// ymm lanes are columns j..j+3. vmovupd: +fn accumulateAvx2(weight: []const f64, x: []const f64, y: []f64) void { + comptime if (builtin.cpu.arch != .x86_64) @compileError("AVX2 GEMV is x86_64 only"); + const k = x.len; + const n = y.len; + var j: usize = 0; + while (j < n) : (j += 4) { + var acc: Acc = @splat(0); + var i: usize = 0; + while (i < k) : (i += 1) { + const prod = avx2Mul(avx2Broadcast(&x[i]), avx2Load(weight.ptr + i * n + j)); + acc = avx2Add(acc, prod); + } + avx2Store(y.ptr + j, acc); + } +} + +const sentinel: f64 = @bitCast(@as(u64, 0x7FF8_0000_0000_0001)); + +/// 128x128 mul1, width 2 halves. Trellis length is 8*8*Width.tileWords(). +const Square = struct { + allocator: std.mem.Allocator, + w: format.Width, + tiles_k: usize, + tiles_n: usize, + trellis: []i16, + suh: []f16, + svh: []f16, + x: []f64, + y: []f64, + scratch: []f64, + + fn init(allocator: std.mem.Allocator) !Square { + const w: format.Width = .{ .halves = 2 }; + const tiles_k: usize = 8; + const tiles_n: usize = 8; + const k = tiles_k * 16; + const n = tiles_n * 16; + const trellis = try allocator.alloc(i16, tiles_k * tiles_n * w.tileWords()); + errdefer allocator.free(trellis); + for (trellis, 0..) |*word, i| word.* = @bitCast(@as(u16, @truncate(i *% 4099 + 7))); + + const suh = try allocator.alloc(f16, k); + errdefer allocator.free(suh); + for (suh) |*v| v.* = 1.0; + const svh = try allocator.alloc(f16, n); + errdefer allocator.free(svh); + for (svh) |*v| v.* = 1.0; + + const x = try allocator.alloc(f64, k); + errdefer allocator.free(x); + for (x, 0..) |*v, i| v.* = @floatFromInt(i + 1); + + const y = try allocator.alloc(f64, n); + errdefer allocator.free(y); + for (y) |*v| v.* = sentinel; + + const scratch = try allocator.alloc(f64, 3 * k * n); + errdefer allocator.free(scratch); + return .{ + .allocator = allocator, + .w = w, + .tiles_k = tiles_k, + .tiles_n = tiles_n, + .trellis = trellis, + .suh = suh, + .svh = svh, + .x = x, + .y = y, + .scratch = scratch, + }; + } + + fn deinit(self: Square) void { + self.allocator.free(self.trellis); + self.allocator.free(self.suh); + self.allocator.free(self.svh); + self.allocator.free(self.x); + self.allocator.free(self.y); + self.allocator.free(self.scratch); + } +}; + +test "128x128 mul1 gemv matches dequantize then a naive matvec" { + const gpa = std.testing.allocator; + const s = try Square.init(gpa); + defer s.deinit(); + const k = s.tiles_k * 16; + const n = s.tiles_n * 16; + try std.testing.expectEqual(@as(usize, 128), k); + try std.testing.expectEqual(@as(usize, 128), n); + try std.testing.expectEqual(@as(usize, 8 * 8 * s.w.tileWords()), s.trellis.len); + + try gemv(s.w, .mul1, s.trellis, s.tiles_k, s.tiles_n, s.suh, s.svh, s.x, s.y, s.scratch); + + const weight = try gpa.alloc(f64, k * n); + defer gpa.free(weight); + const dq_scratch = try gpa.alloc(f64, 2 * k * n); + defer gpa.free(dq_scratch); + format.dequantize(s.w, .mul1, s.trellis, s.tiles_k, s.tiles_n, s.suh, s.svh, weight, dq_scratch); + + for (0..n) |j| { + var sum: f64 = 0; + for (0..k) |i| sum += s.x[i] * weight[i * n + j]; + try std.testing.expectEqual(@as(u64, @bitCast(sum)), @as(u64, @bitCast(s.y[j]))); + } +} + +test "scratch one f64 short returns Length and leaves y untouched" { + const gpa = std.testing.allocator; + const s = try Square.init(gpa); + defer s.deinit(); + for (s.y) |*v| v.* = sentinel; + const short = s.scratch[0 .. s.scratch.len - 1]; + try std.testing.expectError(error.Length, gemv(s.w, .mul1, s.trellis, s.tiles_k, s.tiles_n, s.suh, s.svh, s.x, s.y, short)); + for (s.y) |v| try std.testing.expectEqual(@as(u64, @bitCast(sentinel)), @as(u64, @bitCast(v))); +} + +test "swiglu is silu(gate) times up, and a finite limit clamps both sides" { + var out = [_]f64{ 9, 9 }; + const gate = [_]f64{ 0, 2 }; + const up = [_]f64{ 3, 4 }; + try swiglu(&gate, &up, &out, std.math.inf(f64)); + try std.testing.expectEqual(@as(f64, 0), out[0]); + const silu2 = 2.0 / (1.0 + @exp(-2.0)); + try std.testing.expectApproxEqAbs(silu2 * 4.0, out[1], 1e-12); + + const wide_gate = [_]f64{2}; + const wide_up = [_]f64{4}; + var clamped = [_]f64{9}; + try swiglu(&wide_gate, &wide_up, &clamped, 1); + const silu1 = 1.0 / (1.0 + @exp(-1.0)); + try std.testing.expectApproxEqAbs(silu1, clamped[0], 1e-12); + + var held = [_]f64{9}; + try std.testing.expectError(error.Length, swiglu(wide_gate[0..], wide_up[0..], held[0..0], 1)); + try std.testing.expectEqual(@as(f64, 9), held[0]); +} + +test "combine is the weighted sum of expert rows" { + const rows = [_]f64{ 1, 2, 3, 4 }; + const weights = [_]f64{ 0.25, 0.75 }; + var out = [_]f64{ 9, 9 }; + try combine(&rows, &weights, 2, &out); + try std.testing.expectApproxEqAbs(@as(f64, 2.5), out[0], 1e-12); + try std.testing.expectApproxEqAbs(@as(f64, 3.5), out[1], 1e-12); + try std.testing.expectError(error.Length, combine(rows[0..3], &weights, 2, &out)); +} + +test "x shorter than K returns Length" { + const gpa = std.testing.allocator; + const s = try Square.init(gpa); + defer s.deinit(); + for (s.y) |*v| v.* = sentinel; + const short_x = s.x[0 .. s.x.len - 1]; + try std.testing.expectError(error.Length, gemv(s.w, .mul1, s.trellis, s.tiles_k, s.tiles_n, s.suh, s.svh, short_x, s.y, s.scratch)); + for (s.y) |v| try std.testing.expectEqual(@as(u64, @bitCast(sentinel)), @as(u64, @bitCast(v))); +} + +fn expectGemvAvx2(allocator: std.mem.Allocator, tiles_k: usize, tiles_n: usize) !void { + const w: format.Width = .{ .halves = 2 }; + const k = tiles_k * 16; + const n = tiles_n * 16; + const trellis = try allocator.alloc(i16, tiles_k * tiles_n * w.tileWords()); + defer allocator.free(trellis); + for (trellis, 0..) |*word, i| word.* = @bitCast(@as(u16, @truncate(i *% 4099 + 7))); + + const suh = try allocator.alloc(f16, k); + defer allocator.free(suh); + for (suh) |*v| v.* = 1.0; + const svh = try allocator.alloc(f16, n); + defer allocator.free(svh); + for (svh) |*v| v.* = 1.0; + + const x = try allocator.alloc(f64, k); + defer allocator.free(x); + for (x, 0..) |*v, i| v.* = @floatFromInt(i + 1); + + const y = try allocator.alloc(f64, n); + defer allocator.free(y); + const y2 = try allocator.alloc(f64, n); + defer allocator.free(y2); + @memset(y, sentinel); + @memset(y2, sentinel); + + const scratch = try allocator.alloc(f64, 3 * k * n); + defer allocator.free(scratch); + try std.testing.expectEqual(@as(usize, 3 * k * n), scratch.len); + + try gemv(w, .mul1, trellis, tiles_k, tiles_n, suh, svh, x, y, scratch); + try gemvAvx2(w, .mul1, trellis, tiles_k, tiles_n, suh, svh, x, y2, scratch); + for (y, y2) |a, b| try std.testing.expectEqual(@as(u64, @bitCast(a)), @as(u64, @bitCast(b))); +} + +test "gemvAvx2 matches gemv on the 128x128 mul1 tile" { + const gpa = std.testing.allocator; + const s = try Square.init(gpa); + defer s.deinit(); + const y2 = try gpa.alloc(f64, s.y.len); + defer gpa.free(y2); + @memset(s.y, sentinel); + @memset(y2, sentinel); + try gemv(s.w, .mul1, s.trellis, s.tiles_k, s.tiles_n, s.suh, s.svh, s.x, s.y, s.scratch); + try gemvAvx2(s.w, .mul1, s.trellis, s.tiles_k, s.tiles_n, s.suh, s.svh, s.x, y2, s.scratch); + for (s.y, y2) |a, b| try std.testing.expectEqual(@as(u64, @bitCast(a)), @as(u64, @bitCast(b))); +} + +test "gemvAvx2 matches gemv on synthetic 2560x640 and 640x2560" { + // 2560 * 640 * 3 f64. One buffer, not a second scratch shape. + try std.testing.expectEqual(@as(usize, 39_321_600), 3 * 2560 * 640 * @sizeOf(f64)); + const gpa = std.testing.allocator; + try expectGemvAvx2(gpa, 160, 40); + try expectGemvAvx2(gpa, 40, 160); +} + +test "gemvAvx2 scratch one f64 short returns Length and leaves y untouched" { + const gpa = std.testing.allocator; + const s = try Square.init(gpa); + defer s.deinit(); + for (s.y) |*v| v.* = sentinel; + const short = s.scratch[0 .. s.scratch.len - 1]; + try std.testing.expectError(error.Length, gemvAvx2(s.w, .mul1, s.trellis, s.tiles_k, s.tiles_n, s.suh, s.svh, s.x, s.y, short)); + for (s.y) |v| try std.testing.expectEqual(@as(u64, @bitCast(sentinel)), @as(u64, @bitCast(v))); +} diff --git a/zig/src/families/flashnext/exl3/model/exl3_fwd_cpu.zig b/zig/src/families/flashnext/exl3/model/exl3_fwd_cpu.zig new file mode 100644 index 0000000000..a4ebfb70a6 --- /dev/null +++ b/zig/src/families/flashnext/exl3/model/exl3_fwd_cpu.zig @@ -0,0 +1,414 @@ +//! CPU f64 oracle for the whole Qwen3.8-Flash-Next EXL3 forward, assembled from the per-op oracles (`exl3_hc` + +const std = @import("std"); +const hc = @import("exl3_hc.zig"); +const gdn_cpu = @import("exl3_gdn_cpu.zig"); +const attn_cpu = @import("exl3_attn_cpu.zig"); +const moe = @import("exl3_moe.zig"); +const cpu = @import("exl3_cpu.zig"); +const qmat = @import("../pack/exl3_qmat.zig"); + +pub const hidden: usize = 2560; +pub const streams: usize = 4; +pub const wide: usize = streams * hidden; +pub const topk: usize = 10; +pub const experts: usize = 512; +pub const inter: usize = 640; + +pub const QMat = qmat.QMat; +pub const parallelFor = qmat.parallelFor; +pub const fwht128 = qmat.fwht128; + +// -- small ops --------------------------------------------------------------------------------------------- + +pub fn silu(x: f64) f64 { + return x / (1.0 + @exp(-x)); +} + +pub fn sigmoid(x: f64) f64 { + return 1.0 / (1.0 + @exp(-x)); +} + +/// f64 rows [t][hidden] -> mixed hidden of the hyper-connection (pre), per row. `inject` is [t][4] or null. +pub fn hcPreRows(gpa: std.mem.Allocator, w: hc.Weights, offset: f64, h: []const f64, rows: usize, mixed: []f64, inject: ?[]f64) !void { + const low = try gpa.alloc(f64, hc.real.lowrank); + defer gpa.free(low); + const mix = try gpa.alloc(f64, wide); + defer gpa.free(mix); + const normed = try gpa.alloc(f64, wide); + defer gpa.free(normed); + for (0..rows) |t| { + const inj: ?[]f64 = if (inject) |i| i[t * streams ..][0..streams] else null; + try hc.preF64(hc.real, w, offset, 1e-6, h[t * wide ..][0..wide], mixed[t * hidden ..][0..hidden], inj, low, mix, normed); + } +} + +pub fn hcPostRows(h: []f64, branch: []const f64, inject: []const f64, rows: usize) void { + for (0..rows) |t| { + for (0..streams) |s| { + const g = inject[t * streams + s]; + for (0..hidden) |d| h[t * wide + s * hidden + d] += branch[t * hidden + d] * g; + } + } +} + +// -- layer weights ----------------------------------------------------------------------------------------- + +pub const GdnW = struct { + qkv: QMat, + z: QMat, + out: QMat, + a_w: []const f64, // [48][2560] + b_w: []const f64, + params: gdn_cpu.Params, + + pub fn deinit(w: *GdnW, gpa: std.mem.Allocator) void { + w.qkv.deinit(gpa); + w.z.deinit(gpa); + w.out.deinit(gpa); + gpa.free(w.a_w); + gpa.free(w.b_w); + gpa.free(w.params.conv_w); + gpa.free(w.params.a_log); + gpa.free(w.params.dt_bias); + gpa.free(w.params.norm_w); + } +}; + +pub const AttnW = struct { + q: QMat, + k: QMat, + v: QMat, + o: QMat, + idx: QMat, + norms: attn_cpu.Norms, + + pub fn deinit(w: *AttnW, gpa: std.mem.Allocator) void { + w.q.deinit(gpa); + w.k.deinit(gpa); + w.v.deinit(gpa); + w.o.deinit(gpa); + w.idx.deinit(gpa); + gpa.free(w.norms.q); + gpa.free(w.norms.k); + gpa.free(w.norms.iq); + gpa.free(w.norms.ik); + } +}; + +/// Where the 512 routed experts of one layer come from: decode expert `e` (gate, up, down) into `out`. +pub const ExpertSource = struct { + ctx: *anyopaque, + load: *const fn (ctx: *anyopaque, e: usize, out: *[3]QMat) anyerror!void, +}; + +pub const MoeW = struct { + gate16: []const f16, // [512][2560] + sgate16: []const f16, // [2560] + shared: [3]QMat, + source: ExpertSource, + + pub fn deinit(w: *MoeW, gpa: std.mem.Allocator) void { + for (&w.shared) |*m| m.deinit(gpa); + } +}; + +pub const Mixer = union(enum) { gdn: *GdnW, attn: *AttnW }; + +pub const LayerW = struct { + attn_hc: hc.Weights, + mlp_hc: hc.Weights, + mixer: Mixer, + moe: *MoeW, +}; + +pub const LayerState = union(enum) { gdn: gdn_cpu.State, attn: attn_cpu.Cache }; + +/// What a layer pass records for comparison with the device. +pub const Trace = struct { + attn_branch: []f64, // [rows][2560] + mlp_branch: []f64, + picks: []u16, // [rows][10] + h_mid: []f64, // [rows][10240] streams after the attention block +}; + +// -- mixers ------------------------------------------------------------------------------------------------ + +fn roundF16Slice(dst: []f16, src: []const f64) void { + for (dst, src) |*d, s| d.* = @floatCast(s); +} + +const GdnProj = struct { + w: *const GdnW, + x: []const f64, + qkv: []f64, + z: []f64, + a: []f64, + b: []f64, + gpa: std.mem.Allocator, + err: ?anyerror = null, + + fn row(j: *GdnProj, t: usize) void { + const tmp = j.gpa.alloc(f64, hidden) catch |e| { + j.err = e; + return; + }; + defer j.gpa.free(tmp); + const xr = j.x[t * hidden ..][0..hidden]; + const c = gdn_cpu.qwen.channels(); + const zw = gdn_cpu.qwen.zWidth(); + j.w.qkv.apply(xr, j.qkv[t * c ..][0..c], tmp); + j.w.z.apply(xr, j.z[t * zw ..][0..zw], tmp); + gdn_cpu.rowsDot(j.w.a_w, xr, j.a[t * 48 ..][0..48]); + gdn_cpu.rowsDot(j.w.b_w, xr, j.b[t * 48 ..][0..48]); + } +}; + +const OutJob = struct { + m: *const QMat, + inp: []const f64, + in_w: usize, + out: []f64, + gpa: std.mem.Allocator, + err: ?anyerror = null, + fn row(j: *OutJob, t: usize) void { + const tmp = j.gpa.alloc(f64, j.m.k) catch |e| { + j.err = e; + return; + }; + defer j.gpa.free(tmp); + j.m.apply(j.inp[t * j.in_w ..][0..j.in_w], j.out[t * j.m.n ..][0..j.m.n], tmp); + } +}; + +pub fn gdnBranch(gpa: std.mem.Allocator, w: *const GdnW, st: gdn_cpu.State, x: []const f64, rows: usize, out: []f64) !void { + const d = gdn_cpu.qwen; + const c = d.channels(); + const zw = d.zWidth(); + const qkv = try gpa.alloc(f64, rows * c); + defer gpa.free(qkv); + const z = try gpa.alloc(f64, rows * zw); + defer gpa.free(z); + const a = try gpa.alloc(f64, rows * 48); + defer gpa.free(a); + const b = try gpa.alloc(f64, rows * 48); + defer gpa.free(b); + var pj: GdnProj = .{ .w = w, .x = x, .qkv = qkv, .z = z, .a = a, .b = b, .gpa = gpa }; + parallelFor(rows, 48, &pj, GdnProj.row); + if (pj.err) |e| return e; + const gated = try gpa.alloc(f64, rows * zw); + defer gpa.free(gated); + try gdn_cpu.tokens(gpa, d, w.params, st, rows, qkv, z, a, b, gated); + var oj: OutJob = .{ .m = &w.out, .inp = gated, .in_w = zw, .out = out, .gpa = gpa }; + parallelFor(rows, 48, &oj, OutJob.row); + if (oj.err) |e| return e; +} + +const AttnProj = struct { + w: *const AttnW, + x: []const f64, + qp: []f16, + kp: []f16, + vp: []f16, + ip: []f16, + gpa: std.mem.Allocator, + err: ?anyerror = null, + const q_width = attn_cpu.q_heads * 2 * attn_cpu.head_dim; + const kv_width = attn_cpu.kv_heads * attn_cpu.head_dim; + + fn row(j: *AttnProj, t: usize) void { + const tmp = j.gpa.alloc(f64, hidden) catch |e| { + j.err = e; + return; + }; + defer j.gpa.free(tmp); + const big = j.gpa.alloc(f64, q_width) catch |e| { + j.err = e; + return; + }; + defer j.gpa.free(big); + const xr = j.x[t * hidden ..][0..hidden]; + j.w.q.apply(xr, big, tmp); + roundF16Slice(j.qp[t * q_width ..][0..q_width], big); + j.w.k.apply(xr, big[0..kv_width], tmp); + roundF16Slice(j.kp[t * kv_width ..][0..kv_width], big[0..kv_width]); + j.w.v.apply(xr, big[0..kv_width], tmp); + roundF16Slice(j.vp[t * kv_width ..][0..kv_width], big[0..kv_width]); + j.w.idx.apply(xr, big[0..attn_cpu.index_width], tmp); + roundF16Slice(j.ip[t * attn_cpu.index_width ..][0..attn_cpu.index_width], big[0..attn_cpu.index_width]); + } +}; + +/// Attention rows pos0 .. pos0 + rows. The projections are rounded to fp16 where the EXL3 linears emit them. +pub fn attnBranch(gpa: std.mem.Allocator, w: *const AttnW, cache: *attn_cpu.Cache, pos0: usize, x: []const f64, rows: usize, out: []f64) !void { + const q_width = AttnProj.q_width; + const kv_width = AttnProj.kv_width; + const iw = attn_cpu.index_width; + const qp = try gpa.alloc(f16, rows * q_width); + defer gpa.free(qp); + const kp = try gpa.alloc(f16, rows * kv_width); + defer gpa.free(kp); + const vp = try gpa.alloc(f16, rows * kv_width); + defer gpa.free(vp); + const ip = try gpa.alloc(f16, rows * iw); + defer gpa.free(ip); + var pj: AttnProj = .{ .w = w, .x = x, .qp = qp, .kp = kp, .vp = vp, .ip = ip, .gpa = gpa }; + parallelFor(rows, 48, &pj, AttnProj.row); + if (pj.err) |e| return e; + + const attn_width = attn_cpu.q_heads * attn_cpu.head_dim; + const o_in = try gpa.alloc(f64, rows * attn_width); + defer gpa.free(o_in); + const scores = try gpa.alloc(f64, cache.tokens / attn_cpu.ratio + 2); + defer gpa.free(scores); + var oq: [attn_width]f64 = undefined; + var oiq: [attn_cpu.index_heads * attn_cpu.index_dim]f64 = undefined; + var info = attn_cpu.RowInfo{}; + for (0..rows) |t| { + try attn_cpu.row(cache, pos0 + t, .{ + .qp = qp[t * q_width ..][0..q_width], + .kp = kp[t * kv_width ..][0..kv_width], + .vp = vp[t * kv_width ..][0..kv_width], + .ip = ip[t * iw ..][0..iw], + }, &w.norms, scores, &oq, &oiq, o_in[t * attn_width ..][0..attn_width], &info); + } + var oj: OutJob = .{ .m = &w.o, .inp = o_in, .in_w = attn_width, .out = out, .gpa = gpa }; + parallelFor(rows, 48, &oj, OutJob.row); + if (oj.err) |e| return e; +} + +// -- MoE --------------------------------------------------------------------------------------------------- + +const Loaded = struct { + id: usize, + mats: [3]QMat, +}; + +const LoadJob = struct { + src: ExpertSource, + ids: []const usize, + slots: []Loaded, + err: ?anyerror = null, + fn one(j: *LoadJob, i: usize) void { + j.slots[i].id = j.ids[i]; + j.src.load(j.src.ctx, j.ids[i], &j.slots[i].mats) catch |e| { + j.err = e; + }; + } +}; + +const MoeRow = struct { + w: *const MoeW, + x: []const f64, + out: []f64, + picks: []const u16, // [rows][10] + weights: []const f64, // [rows][10] + shared_w: []const f64, + loaded: []const Loaded, + gpa: std.mem.Allocator, + err: ?anyerror = null, + + fn find(j: *const MoeRow, id: usize) *const Loaded { + for (j.loaded) |*l| if (l.id == id) return l; + unreachable; + } + + fn row(j: *MoeRow, t: usize) void { + const buf = j.gpa.alloc(f64, hidden + 4 * inter + 3 * hidden) catch |e| { + j.err = e; + return; + }; + defer j.gpa.free(buf); + const tmp = buf[0..hidden]; + const g = buf[hidden..][0..inter]; + const u = buf[hidden + inter ..][0..inter]; + const glu = buf[hidden + 2 * inter ..][0..inter]; + const y = buf[hidden + 3 * inter ..][0..hidden]; + const acc = j.out[t * hidden ..][0..hidden]; + @memset(acc, 0); + const xr = j.x[t * hidden ..][0..hidden]; + for (0..topk) |k| { + const l = j.find(j.picks[t * topk + k]); + expertApply(&l.mats, xr, tmp, g, u, glu, y); + const wk = j.weights[t * topk + k]; + for (acc, y) |*o, v| o.* += wk * v; + } + expertApply(&j.w.shared, xr, tmp, g, u, glu, y); + const sw = j.shared_w[t]; + for (acc, y) |*o, v| o.* += sw * v; + } +}; + +fn expertApply(m: *const [3]QMat, x: []const f64, tmp: []f64, g: []f64, u: []f64, glu: []f64, y: []f64) void { + m[0].apply(x, g, tmp); + m[1].apply(x, u, tmp); + for (glu, g, u) |*o, gv, uv| o.* = silu(gv) * uv; + m[2].apply(glu, y, tmp[0..inter]); +} + +/// Routed + shared experts for `rows` rows from the f64 hidden `x` (the f64 router, not the device's). +pub fn moeBranch(gpa: std.mem.Allocator, w: *const MoeW, x: []const f64, rows: usize, out: []f64, picks_out: []u16) !void { + const picks = picks_out; + const weights = try gpa.alloc(f64, rows * topk); + defer gpa.free(weights); + const shared_w = try gpa.alloc(f64, rows); + defer gpa.free(shared_w); + var logits: [experts + 1]f64 = undefined; + var ids: std.ArrayList(usize) = .empty; + defer ids.deinit(gpa); + for (0..rows) |t| { + try moe.logitsF16(w.gate16, experts, w.sgate16, x[t * hidden ..][0..hidden], &logits); + var sw: f64 = 0; + try moe.routeFromLogits(&logits, topk, picks[t * topk ..][0..topk], weights[t * topk ..][0..topk], &sw); + shared_w[t] = sw; + for (picks[t * topk ..][0..topk]) |p| { + var seen = false; + for (ids.items) |i| seen = seen or i == p; + if (!seen) try ids.append(gpa, p); + } + } + const loaded = try gpa.alloc(Loaded, ids.items.len); + defer gpa.free(loaded); + var lj: LoadJob = .{ .src = w.source, .ids = ids.items, .slots = loaded }; + parallelFor(ids.items.len, 48, &lj, LoadJob.one); + defer if (lj.err == null) for (loaded) |*l| for (&l.mats) |*m| m.deinit(gpa); + if (lj.err) |e| return e; + var mj: MoeRow = .{ .w = w, .x = x, .out = out, .picks = picks, .weights = weights, .shared_w = shared_w, .loaded = loaded, .gpa = gpa }; + parallelFor(rows, 48, &mj, MoeRow.row); + if (mj.err) |e| return e; +} + +// -- one layer --------------------------------------------------------------------------------------------- + +/// `h` is [rows][10240] f64 streams, updated in place. `pos0` is the absolute position of row 0. +pub fn layerRows(gpa: std.mem.Allocator, w: LayerW, st: *LayerState, pos0: usize, offset: f64, h: []f64, rows: usize, trace: ?Trace) !void { + const mixed = try gpa.alloc(f64, rows * hidden); + defer gpa.free(mixed); + const inject = try gpa.alloc(f64, rows * streams); + defer gpa.free(inject); + const branch = try gpa.alloc(f64, rows * hidden); + defer gpa.free(branch); + + try hcPreRows(gpa, w.attn_hc, offset, h, rows, mixed, inject); + switch (w.mixer) { + .gdn => |g| try gdnBranch(gpa, g, st.gdn, mixed, rows, branch), + .attn => |a| try attnBranch(gpa, a, &st.attn, pos0, mixed, rows, branch), + } + if (trace) |tr| @memcpy(tr.attn_branch[0 .. rows * hidden], branch); + hcPostRows(h, branch, inject, rows); + if (trace) |tr| @memcpy(tr.h_mid[0 .. rows * wide], h[0 .. rows * wide]); + + try hcPreRows(gpa, w.mlp_hc, offset, h, rows, mixed, inject); + const picks = try gpa.alloc(u16, rows * topk); + defer gpa.free(picks); + try moeBranch(gpa, w.moe, mixed, rows, branch, picks); + if (trace) |tr| { + @memcpy(tr.mlp_branch[0 .. rows * hidden], branch); + @memcpy(tr.picks[0 .. rows * topk], picks); + } + hcPostRows(h, branch, inject, rows); +} + +comptime { + _ = qmat; +} diff --git a/zig/src/families/flashnext/exl3/model/exl3_gdn_cpu.zig b/zig/src/families/flashnext/exl3/model/exl3_gdn_cpu.zig new file mode 100644 index 0000000000..cbf187544e --- /dev/null +++ b/zig/src/families/flashnext/exl3/model/exl3_gdn_cpu.zig @@ -0,0 +1,399 @@ +//! CPU oracle for one Gated DeltaNet layer of Qwen3.8-Flash-Next (qwen4_exp_text), everything between the two EXL3 + +const std = @import("std"); + +pub const Dims = struct { + nk: usize, + nv: usize, + dk: usize, + dv: usize, + taps: usize, + + /// conv channels: q (nk*dk), k (nk*dk), v (nv*dv) + pub fn channels(d: Dims) usize { + return 2 * d.nk * d.dk + d.nv * d.dv; + } + pub fn zWidth(d: Dims) usize { + return d.nv * d.dv; + } + /// conv state rows kept per layer + pub fn keep(d: Dims) usize { + return d.taps - 1; + } + pub fn recElems(d: Dims) usize { + return d.nv * d.dv * d.dk; + } + pub fn convElems(d: Dims) usize { + return d.keep() * d.channels(); + } +}; + +/// Flash Next: 16 key heads, 48 value heads, 128 x 128, kernel 4. +pub const qwen: Dims = .{ .nk = 16, .nv = 48, .dk = 128, .dv = 128, .taps = 4 }; + +/// Per-sequence, per-layer state. +pub const State = struct { + conv: []f64, + rec: []f64, + + pub fn init(gpa: std.mem.Allocator, d: Dims) !State { + const conv = try gpa.alloc(f64, d.convElems()); + errdefer gpa.free(conv); + const rec = try gpa.alloc(f64, d.recElems()); + @memset(conv, 0); + @memset(rec, 0); + return .{ .conv = conv, .rec = rec }; + } + pub fn deinit(s: *State, gpa: std.mem.Allocator) void { + gpa.free(s.conv); + gpa.free(s.rec); + } + pub fn reset(s: *State) void { + @memset(s.conv, 0); + @memset(s.rec, 0); + } +}; + +/// The layer's non-EXL3 parameters, as f64. +pub const Params = struct { + conv_w: []const f64, // [channels][taps] + a_log: []const f64, // [nv] + dt_bias: []const f64, // [nv] + norm_w: []const f64, // [dv] + eps: f64 = 1e-6, +}; + +pub fn sigmoid(x: f64) f64 { + return 1.0 / (1.0 + @exp(-x)); +} + +pub fn silu(x: f64) f64 { + return x / (1.0 + @exp(-x)); +} + +/// torch.nn.functional.softplus with threshold 20. +pub fn softplus(x: f64) f64 { + return if (x > 20.0) x else std.math.log1p(@exp(x)); +} + +/// y[j] = sum_i w[j * k + i] * x[i], for a dense [rows][k] weight (in_proj_a / in_proj_b). +pub fn rowsDot(w: []const f64, x: []const f64, y: []f64) void { + const k = x.len; + for (y, 0..) |*o, j| { + var s: f64 = 0; + for (x, w[j * k ..][0..k]) |xi, wi| s += xi * wi; + o.* = s; + } +} + +/// One token of the causal depthwise conv: out = silu(sum_tap w[c][tap] * x[t - (taps-1) + tap]), x before the +pub fn convStep(d: Dims, p: Params, conv: []f64, x: []const f64, out: []f64) void { + const c = d.channels(); + const kp = d.keep(); + for (0..c) |ch| { + var acc: f64 = 0; + for (0..d.taps) |tap| { + const v = if (tap < kp) conv[tap * c + ch] else x[ch]; + acc += p.conv_w[ch * d.taps + tap] * v; + } + out[ch] = silu(acc); + } + for (0..kp) |r| { + for (0..c) |ch| conv[r * c + ch] = if (r + 1 < kp) conv[(r + 1) * c + ch] else x[ch]; + } +} + +/// Decay exp(g) and beta for each value head from the raw a and b projections. +pub fn gates(d: Dims, p: Params, a: []const f64, b: []const f64, decay: []f64, beta: []f64) void { + for (0..d.nv) |h| { + const g = -@exp(p.a_log[h]) * softplus(a[h] + p.dt_bias[h]); + decay[h] = @exp(g); + beta[h] = sigmoid(b[h]); + } +} + +/// One token of the delta rule on all value heads. `act` is the conv output [q | k | v]; y is [nv][dv]. +pub fn deltaStep(d: Dims, rec: []f64, act: []const f64, decay: []const f64, beta: []const f64, y: []f64) void { + const group = d.nv / d.nk; + const kbase = d.nk * d.dk; + const vbase = 2 * kbase; + for (0..d.nv) |hv| { + const hk = hv / group; + const q = act[hk * d.dk ..][0..d.dk]; + const k = act[kbase + hk * d.dk ..][0..d.dk]; + var qq: f64 = 0; + var kk: f64 = 0; + for (q, k) |a, b| { + qq += a * a; + kk += b * b; + } + const iq = 1.0 / @sqrt(qq + 1e-6) / @sqrt(@as(f64, @floatFromInt(d.dk))); + const ik = 1.0 / @sqrt(kk + 1e-6); + for (0..d.dv) |v| { + const row = rec[(hv * d.dv + v) * d.dk ..][0..d.dk]; + var kv: f64 = 0; + for (row, k) |*s, ki| { + s.* *= decay[hv]; + kv += s.* * ki * ik; + } + const delta = (act[vbase + hv * d.dv + v] - kv) * beta[hv]; + var o: f64 = 0; + for (row, k, q) |*s, ki, qi| { + s.* += ki * ik * delta; + o += s.* * qi * iq; + } + y[hv * d.dv + v] = o; + } + } +} + +/// Gated RMS norm per value head: w * (y * rsqrt(mean(y^2) + eps)) * sigmoid(z). +pub fn gatedNorm(d: Dims, p: Params, y: []const f64, z: []const f64, out: []f64) void { + for (0..d.nv) |hv| { + const row = y[hv * d.dv ..][0..d.dv]; + var ss: f64 = 0; + for (row) |v| ss += v * v; + const inv = 1.0 / @sqrt(ss / @as(f64, @floatFromInt(d.dv)) + p.eps); + for (row, 0..) |v, i| out[hv * d.dv + i] = p.norm_w[i] * (v * inv) * sigmoid(z[hv * d.dv + i]); + } +} + +/// `n` tokens, in order, from the projection outputs: qkv [n][channels], z [n][nv*dv],. +pub fn tokens(gpa: std.mem.Allocator, d: Dims, p: Params, st: State, n: usize, qkv: []const f64, z: []const f64, a: []const f64, b: []const f64, out: []f64) !void { + const c = d.channels(); + const zw = d.zWidth(); + const act = try gpa.alloc(f64, c); + defer gpa.free(act); + const y = try gpa.alloc(f64, zw); + defer gpa.free(y); + const decay = try gpa.alloc(f64, d.nv); + defer gpa.free(decay); + const beta = try gpa.alloc(f64, d.nv); + defer gpa.free(beta); + for (0..n) |t| { + convStep(d, p, st.conv, qkv[t * c ..][0..c], act); + gates(d, p, a[t * d.nv ..][0..d.nv], b[t * d.nv ..][0..d.nv], decay, beta); + deltaStep(d, st.rec, act, decay, beta, y); + gatedNorm(d, p, y, z[t * zw ..][0..zw], out[t * zw ..][0..zw]); + } +} + +/// Byte sizes of the device state table: conv rows are fp16 (they are raw EXL3 projection outputs), the recurrent +pub const Table = struct { + dims: Dims, + layers: usize, + slots: usize, + + pub fn convBytesPer(t: Table) usize { + return t.dims.convElems() * 2; + } + pub fn recBytesPer(t: Table) usize { + return t.dims.recElems() * 4; + } + pub fn convBytes(t: Table) usize { + return t.convBytesPer() * t.layers * t.slots; + } + pub fn recBytes(t: Table) usize { + return t.recBytesPer() * t.layers * t.slots; + } + pub fn convOffset(t: Table, slot: usize, layer: usize) usize { + return (slot * t.layers + layer) * t.convBytesPer(); + } + pub fn recOffset(t: Table, slot: usize, layer: usize) usize { + return (slot * t.layers + layer) * t.recBytesPer(); + } +}; + +// ---- tests (small dims; the arithmetic does not depend on the Flash Next sizes) ---- + +const small: Dims = .{ .nk = 2, .nv = 4, .dk = 8, .dv = 8, .taps = 4 }; + +fn fill(rng: *std.Random.DefaultPrng, xs: []f64, scale: f64) void { + for (xs) |*x| x.* = (rng.random().float(f64) - 0.5) * 2.0 * scale; +} + +const TestParams = struct { + conv_w: []f64, + a_log: []f64, + dt_bias: []f64, + norm_w: []f64, + + fn init(gpa: std.mem.Allocator, d: Dims, rng: *std.Random.DefaultPrng) !TestParams { + const t: TestParams = .{ + .conv_w = try gpa.alloc(f64, d.channels() * d.taps), + .a_log = try gpa.alloc(f64, d.nv), + .dt_bias = try gpa.alloc(f64, d.nv), + .norm_w = try gpa.alloc(f64, d.dv), + }; + fill(rng, t.conv_w, 0.7); + fill(rng, t.a_log, 1.0); + fill(rng, t.dt_bias, 1.0); + for (t.norm_w) |*w| w.* = 0.5 + rng.random().float(f64); + return t; + } + fn deinit(t: TestParams, gpa: std.mem.Allocator) void { + gpa.free(t.conv_w); + gpa.free(t.a_log); + gpa.free(t.dt_bias); + gpa.free(t.norm_w); + } + fn view(t: TestParams) Params { + return .{ .conv_w = t.conv_w, .a_log = t.a_log, .dt_bias = t.dt_bias, .norm_w = t.norm_w }; + } +}; + +test "conv is a causal depthwise filter over the stored rows" { + const gpa = std.testing.allocator; + var rng = std.Random.DefaultPrng.init(1); + const tp = try TestParams.init(gpa, small, &rng); + defer tp.deinit(gpa); + const c = small.channels(); + const n = 7; + const xs = try gpa.alloc(f64, n * c); + defer gpa.free(xs); + fill(&rng, xs, 1.0); + var st = try State.init(gpa, small); + defer st.deinit(gpa); + const out = try gpa.alloc(f64, c); + defer gpa.free(out); + for (0..n) |t| { + convStep(small, tp.view(), st.conv, xs[t * c ..][0..c], out); + for (0..c) |ch| { + var acc: f64 = 0; + for (0..small.taps) |tap| { + // input index t - (taps - 1) + tap; before the sequence it is zero + const at: isize = @as(isize, @intCast(t)) - @as(isize, @intCast(small.taps - 1)) + @as(isize, @intCast(tap)); + if (at >= 0) acc += tp.conv_w[ch * small.taps + tap] * xs[@as(usize, @intCast(at)) * c + ch]; + } + try std.testing.expectApproxEqAbs(silu(acc), out[ch], 1e-12); + } + } + // the state ends as the last three raw rows + for (0..3) |r| for (0..c) |ch| try std.testing.expectEqual(xs[(n - 3 + r) * c + ch], st.conv[r * c + ch]); +} + +test "one scalar head follows S = g S + k (v - g S k) beta, o = S q" { + const d: Dims = .{ .nk = 1, .nv = 1, .dk = 1, .dv = 1, .taps = 4 }; + var rec = [_]f64{0.5}; + const act = [_]f64{ 2.0, 3.0, 0.75 }; // q, k, v + const decay = [_]f64{0.9}; + const beta = [_]f64{0.4}; + var y = [_]f64{0}; + deltaStep(d, &rec, &act, &decay, &beta, &y); + const ik = 1.0 / @sqrt(9.0 + 1e-6); + const iq = 1.0 / @sqrt(4.0 + 1e-6); + const k = 3.0 * ik; + const s0 = 0.5 * 0.9; + const delta = (0.75 - s0 * k) * 0.4; + const s1 = s0 + k * delta; + try std.testing.expectApproxEqAbs(s1, rec[0], 1e-14); + try std.testing.expectApproxEqAbs(s1 * 2.0 * iq, y[0], 1e-14); +} + +test "with beta 1 and no decay the state returns v for the key just written" { + const d: Dims = .{ .nk = 1, .nv = 2, .dk = 4, .dv = 3, .taps = 4 }; + var rng = std.Random.DefaultPrng.init(2); + var rec: [2 * 3 * 4]f64 = undefined; + fill(&rng, &rec, 1.0); + var act: [2 * 4 + 2 * 3]f64 = undefined; + fill(&rng, &act, 1.0); + const decay = [_]f64{ 1, 1 }; + const beta = [_]f64{ 1, 1 }; + var y: [6]f64 = undefined; + deltaStep(d, &rec, &act, &decay, &beta, &y); + var kk: f64 = 0; + for (act[4..8]) |k| kk += k * k; + const ik = 1.0 / @sqrt(kk + 1e-6); + for (0..2) |hv| for (0..3) |v| { + var read: f64 = 0; + for (0..4) |i| read += rec[(hv * 3 + v) * 4 + i] * act[4 + i] * ik; + // S k = v * (|k|^2 / (|k|^2 + 1e-6)), the l2norm eps leaves a tiny shortfall + try std.testing.expectApproxEqAbs(act[8 + hv * 3 + v], read, 1e-5); + }; +} + +test "value heads share their key head in groups (repeat_interleave)" { + const d: Dims = .{ .nk = 2, .nv = 4, .dk = 4, .dv = 2, .taps = 4 }; + var rng = std.Random.DefaultPrng.init(3); + var act: [2 * 2 * 4 + 4 * 2]f64 = undefined; + fill(&rng, &act, 1.0); + const decay = [_]f64{ 0.8, 0.8, 0.8, 0.8 }; + const beta = [_]f64{ 0.5, 0.5, 0.5, 0.5 }; + var rec: [4 * 2 * 4]f64 = @splat(0); + var y: [8]f64 = undefined; + deltaStep(d, &rec, &act, &decay, &beta, &y); + // zero state: kv = 0, so S = k (v beta)^T and o = (k . q) v beta, with q, k of key head hv / 2 + for (0..4) |hv| { + const hk = hv / 2; + const q = act[hk * 4 ..][0..4]; + const k = act[8 + hk * 4 ..][0..4]; + var qq: f64 = 0; + var kk: f64 = 0; + var qk: f64 = 0; + for (q, k) |a, b| { + qq += a * a; + kk += b * b; + qk += a * b; + } + const iq = 1.0 / @sqrt(qq + 1e-6) / @sqrt(4.0); + const ik = 1.0 / @sqrt(kk + 1e-6); + for (0..2) |v| try std.testing.expectApproxEqAbs(qk * iq * ik * act[16 + hv * 2 + v] * 0.5, y[hv * 2 + v], 1e-12); + } +} + +test "gated norm is w * rms-normalised y * sigmoid(z)" { + const d: Dims = .{ .nk = 1, .nv = 1, .dk = 1, .dv = 4, .taps = 4 }; + const p: Params = .{ .conv_w = &.{}, .a_log = &.{}, .dt_bias = &.{}, .norm_w = &.{ 1, 2, 3, 4 } }; + const y = [_]f64{ 1, -2, 3, -4 }; + const z = [_]f64{ 0, 1, -1, 2 }; + var out: [4]f64 = undefined; + gatedNorm(d, p, &y, &z, &out); + const inv = 1.0 / @sqrt((1.0 + 4 + 9 + 16) / 4.0 + 1e-6); + for (0..4) |i| try std.testing.expectApproxEqAbs(p.norm_w[i] * y[i] * inv * sigmoid(z[i]), out[i], 1e-14); +} + +test "processing tokens in two calls equals one call, state included" { + const gpa = std.testing.allocator; + var rng = std.Random.DefaultPrng.init(4); + const tp = try TestParams.init(gpa, small, &rng); + defer tp.deinit(gpa); + const n = 9; + const c = small.channels(); + const zw = small.zWidth(); + const qkv = try gpa.alloc(f64, n * c); + defer gpa.free(qkv); + const z = try gpa.alloc(f64, n * zw); + defer gpa.free(z); + const a = try gpa.alloc(f64, n * small.nv); + defer gpa.free(a); + const b = try gpa.alloc(f64, n * small.nv); + defer gpa.free(b); + fill(&rng, qkv, 1.0); + fill(&rng, z, 1.0); + fill(&rng, a, 1.0); + fill(&rng, b, 1.0); + var whole = try State.init(gpa, small); + defer whole.deinit(gpa); + var split = try State.init(gpa, small); + defer split.deinit(gpa); + const o1 = try gpa.alloc(f64, n * zw); + defer gpa.free(o1); + const o2 = try gpa.alloc(f64, n * zw); + defer gpa.free(o2); + try tokens(gpa, small, tp.view(), whole, n, qkv, z, a, b, o1); + const cut = 5; + try tokens(gpa, small, tp.view(), split, cut, qkv, z, a, b, o2); + try tokens(gpa, small, tp.view(), split, n - cut, qkv[cut * c ..], z[cut * zw ..], a[cut * small.nv ..], b[cut * small.nv ..], o2[cut * zw ..]); + try std.testing.expectEqualSlices(f64, o1, o2); + try std.testing.expectEqualSlices(f64, whole.rec, split.rec); + try std.testing.expectEqualSlices(f64, whole.conv, split.conv); +} + +test "state table addresses a (slot, layer) pair and Flash Next sizes" { + const t: Table = .{ .dims = qwen, .layers = 36, .slots = 2 }; + try std.testing.expectEqual(@as(usize, 10240), qwen.channels()); + try std.testing.expectEqual(@as(usize, 3 * 10240 * 2), t.convBytesPer()); + try std.testing.expectEqual(@as(usize, 48 * 128 * 128 * 4), t.recBytesPer()); + try std.testing.expectEqual(t.recBytesPer() * 37, t.recOffset(1, 1)); + try std.testing.expectEqual(t.convBytesPer() * 36 * 2, t.convBytes()); +} diff --git a/zig/src/families/flashnext/exl3/model/exl3_hc.zig b/zig/src/families/flashnext/exl3/model/exl3_hc.zig new file mode 100644 index 0000000000..1d4a25ef14 --- /dev/null +++ b/zig/src/families/flashnext/exl3/model/exl3_hc.zig @@ -0,0 +1,253 @@ +//! Host reference of the hyper-connection pre, post and expand steps over the unquantized hc weights. + +const std = @import("std"); +const shape = @import("../pack/exl3_shape.zig"); + +// --------------------------------------------------------------------------------------------- Hyper-connection pre /. + +pub const Dims = struct { streams: usize, dim: usize, lowrank: usize }; +pub const real: Dims = .{ .streams = shape.hc_count, .dim = shape.hidden, .lowrank = shape.hc_lowrank }; + +/// Row-major, out-major weights as f32 (exact from F16). `inject` is null for the final mixer. +pub const Weights = struct { + norm: []const f32, + down: []const f32, + up: []const f32, + inject: ?[]const f32, +}; + +pub const Error = error{Length}; + +fn silu32(x: f32) f32 { + return x / (1.0 + @exp(-x)); +} +fn sigmoid32(x: f32) f32 { + return 1.0 / (1.0 + @exp(-x)); +} + +/// f32 values `scratch` must hold: normed (wide), low (lowrank), mix (wide). +pub fn scratchLen(d: Dims) usize { + return 2 * d.streams * d.dim + d.lowrank; +} + +fn checkWeights(d: Dims, w: Weights) Error!void { + const wide = d.streams * d.dim; + if (w.norm.len != wide or w.down.len != d.lowrank * wide or w.up.len != wide * d.lowrank) return error.Length; + if (w.inject) |inj| if (inj.len != d.streams * wide) return error.Length; +} + +/// Block input `mixed` [dim] and, when the weights carry `block_inject_weight` and `inject` is non-null, the S gates. +pub fn pre(d: Dims, w: Weights, offset: f32, eps: f32, h: []const f32, mixed: []f32, inject: ?[]f32, scratch: []f32) Error!void { + const wide = d.streams * d.dim; + try checkWeights(d, w); + if (h.len != wide or mixed.len != d.dim or scratch.len < scratchLen(d)) return error.Length; + if (inject) |g| if (w.inject == null or g.len != d.streams) return error.Length; + if (!std.math.isFinite(eps) or eps < 0) return error.Length; + const normed = scratch[0..wide]; + const low = scratch[wide .. wide + d.lowrank]; + const mix = scratch[wide + d.lowrank .. 2 * wide + d.lowrank]; + + for (0..d.streams) |s| { + const row = h[s * d.dim ..][0..d.dim]; + var sum_sq: f32 = 0; + for (row) |v| sum_sq += v * v; + const rms = @sqrt(sum_sq / @as(f32, @floatFromInt(d.dim)) + eps); + for (0..d.dim) |i| normed[s * d.dim + i] = row[i] / rms * (w.norm[s * d.dim + i] + offset); + } + const inv_s: f32 = 1.0 / @as(f32, @floatFromInt(d.streams)); + for (0..d.lowrank) |r| { + var acc: f32 = 0; + const row = w.down[r * wide ..][0..wide]; + for (row, normed) |a, b| acc += a * b; + low[r] = silu32(acc * inv_s); + } + for (0..wide) |j| { + var acc: f32 = 0; + const row = w.up[j * d.lowrank ..][0..d.lowrank]; + for (row, low) |a, b| acc += a * b; + mix[j] = sigmoid32(acc); + } + for (0..d.dim) |i| { + var acc: f32 = 0; + for (0..d.streams) |s| acc += mix[s * d.dim + i] * normed[s * d.dim + i]; + mixed[i] = acc * inv_s; + } + if (inject) |g| { + const inj = w.inject.?; + for (0..d.streams) |s| { + var acc: f32 = 0; + const row = inj[s * wide ..][0..wide]; + for (row, normed) |a, b| acc += a * b; + g[s] = 2.0 * sigmoid32(acc * inv_s); + } + } +} + +/// `out[s, d] = h[s, d] + branch[d] * inject[s]`. `out` may be `h`. Lengths are checked before any store. +pub fn post(d: Dims, h: []const f32, branch: []const f32, inject: []const f32, out: []f32) Error!void { + if (h.len != d.streams * d.dim or out.len != h.len or branch.len != d.dim or inject.len != d.streams) return error.Length; + for (0..d.streams) |s| { + for (0..d.dim) |i| out[s * d.dim + i] = h[s * d.dim + i] + branch[i] * inject[s]; + } +} + +/// Model start: the embedding row copied into every stream. +pub fn expand(d: Dims, embed: []const f32, out: []f32) Error!void { + if (embed.len != d.dim or out.len != d.streams * d.dim) return error.Length; + for (0..d.streams) |s| @memcpy(out[s * d.dim ..][0..d.dim], embed); +} + +/// The same maths in f64, written as plain loops over (stream, dim) with no scratch, for the oracle. +pub fn preF64(d: Dims, w: Weights, offset: f64, eps: f64, h: []const f64, mixed: []f64, inject: ?[]f64, low: []f64, mix: []f64, normed: []f64) Error!void { + const wide = d.streams * d.dim; + try checkWeights(d, w); + if (h.len != wide or mixed.len != d.dim or low.len != d.lowrank or mix.len != wide or normed.len != wide) return error.Length; + for (0..d.streams) |s| { + var ms: f64 = 0; + for (0..d.dim) |i| ms += h[s * d.dim + i] * h[s * d.dim + i]; + ms /= @floatFromInt(d.dim); + const inv = 1.0 / @sqrt(ms + eps); + for (0..d.dim) |i| normed[s * d.dim + i] = h[s * d.dim + i] * inv * (@as(f64, w.norm[s * d.dim + i]) + offset); + } + const sf: f64 = @floatFromInt(d.streams); + for (0..d.lowrank) |r| { + var acc: f64 = 0; + for (0..wide) |c| acc += @as(f64, w.down[r * wide + c]) * normed[c]; + const z = acc / sf; + low[r] = z / (1.0 + @exp(-z)); + } + for (0..wide) |j| { + var acc: f64 = 0; + for (0..d.lowrank) |r| acc += @as(f64, w.up[j * d.lowrank + r]) * low[r]; + mix[j] = 1.0 / (1.0 + @exp(-acc)); + } + for (0..d.dim) |i| { + var acc: f64 = 0; + for (0..d.streams) |s| acc += mix[s * d.dim + i] * normed[s * d.dim + i]; + mixed[i] = acc / sf; + } + if (inject) |g| { + const inj = w.inject orelse return error.Length; + if (g.len != d.streams) return error.Length; + for (0..d.streams) |s| { + var acc: f64 = 0; + for (0..wide) |c| acc += @as(f64, inj[s * wide + c]) * normed[c]; + g[s] = 2.0 / (1.0 + @exp(-acc / sf)); + } + } +} + +pub fn postF64(d: Dims, h: []const f64, branch: []const f64, inject: []const f64, out: []f64) Error!void { + if (h.len != d.streams * d.dim or out.len != h.len or branch.len != d.dim or inject.len != d.streams) return error.Length; + for (0..d.streams) |s| for (0..d.dim) |i| { + out[s * d.dim + i] = h[s * d.dim + i] + branch[i] * inject[s]; + }; +} + +const tiny: Dims = .{ .streams = 2, .dim = 4, .lowrank = 3 }; + +fn tinyWeights(gpa: std.mem.Allocator, with_inject: bool) !struct { norm: []f32, down: []f32, up: []f32, inject: ?[]f32 } { + const wide = tiny.streams * tiny.dim; + var rng = std.Random.DefaultPrng.init(7); + const r = rng.random(); + const norm = try gpa.alloc(f32, wide); + const down = try gpa.alloc(f32, tiny.lowrank * wide); + const up = try gpa.alloc(f32, wide * tiny.lowrank); + for (norm) |*v| v.* = r.float(f32) - 0.5; + for (down) |*v| v.* = (r.float(f32) - 0.5) * 2; + for (up) |*v| v.* = (r.float(f32) - 0.5) * 2; + var inject: ?[]f32 = null; + if (with_inject) { + inject = try gpa.alloc(f32, tiny.streams * wide); + for (inject.?) |*v| v.* = (r.float(f32) - 0.5) * 2; + } + return .{ .norm = norm, .down = down, .up = up, .inject = inject }; +} + +test "pre and post match the f64 plain-loop oracle on a tiny shape" { + const gpa = std.testing.allocator; + const t = try tinyWeights(gpa, true); + defer gpa.free(t.norm); + defer gpa.free(t.down); + defer gpa.free(t.up); + defer gpa.free(t.inject.?); + const w: Weights = .{ .norm = t.norm, .down = t.down, .up = t.up, .inject = t.inject }; + const wide = tiny.streams * tiny.dim; + const h = [wide]f32{ 1, -2, 3, 0.5, 4, 1, -1, 2 }; + var mixed: [tiny.dim]f32 = undefined; + var gates: [tiny.streams]f32 = undefined; + var scratch: [2 * wide + tiny.lowrank]f32 = undefined; + try pre(tiny, w, 1.0, 1e-6, &h, &mixed, &gates, &scratch); + + var h64: [wide]f64 = undefined; + for (&h64, h) |*a, b| a.* = b; + var mixed64: [tiny.dim]f64 = undefined; + var gates64: [tiny.streams]f64 = undefined; + var low: [tiny.lowrank]f64 = undefined; + var mix: [wide]f64 = undefined; + var normed: [wide]f64 = undefined; + try preF64(tiny, w, 1.0, 1e-6, &h64, &mixed64, &gates64, &low, &mix, &normed); + for (mixed, mixed64) |a, b| try std.testing.expectApproxEqAbs(@as(f64, a), b, 1e-5); + for (gates, gates64) |a, b| try std.testing.expectApproxEqAbs(@as(f64, a), b, 1e-5); + + const branch = [tiny.dim]f32{ 0.5, -1, 2, 0 }; + var out: [wide]f32 = undefined; + try post(tiny, &h, &branch, &gates, &out); + for (0..tiny.streams) |s| for (0..tiny.dim) |i| { + try std.testing.expectApproxEqAbs(h[s * tiny.dim + i] + branch[i] * gates[s], out[s * tiny.dim + i], 1e-6); + }; + // aliasing out and h is allowed + var again = h; + try post(tiny, &again, &branch, &gates, &again); + try std.testing.expectEqualSlices(f32, &out, &again); +} + +test "zero mixing weights give mix 1/2, inject 1, and the offset-only norm" { + const gpa = std.testing.allocator; + const t = try tinyWeights(gpa, true); + defer gpa.free(t.norm); + defer gpa.free(t.down); + defer gpa.free(t.up); + defer gpa.free(t.inject.?); + @memset(t.down, 0); + @memset(t.up, 0); + @memset(t.inject.?, 0); + @memset(t.norm, 0); + const w: Weights = .{ .norm = t.norm, .down = t.down, .up = t.up, .inject = t.inject }; + const wide = tiny.streams * tiny.dim; + // stream 0 all 3s, stream 1 all -2s: each stream normalises to +-1 (offset 1, weight 0). + const h = [wide]f32{ 3, 3, 3, 3, -2, -2, -2, -2 }; + var mixed: [tiny.dim]f32 = undefined; + var gates: [tiny.streams]f32 = undefined; + var scratch: [2 * wide + tiny.lowrank]f32 = undefined; + try pre(tiny, w, 1.0, 0.0, &h, &mixed, &gates, &scratch); + // mean_s(0.5 * (+1, -1)) = 0 + for (mixed) |v| try std.testing.expectApproxEqAbs(@as(f32, 0), v, 1e-6); + for (gates) |g| try std.testing.expectApproxEqAbs(@as(f32, 1), g, 1e-6); + // gamma stored directly: offset 0 and weight 0 zeroes the normed streams + try pre(tiny, w, 0.0, 0.0, &h, &mixed, &gates, &scratch); + for (mixed) |v| try std.testing.expectEqual(@as(f32, 0), v); +} + +test "pre, post and expand refuse bad lengths before a store" { + const gpa = std.testing.allocator; + const t = try tinyWeights(gpa, false); + defer gpa.free(t.norm); + defer gpa.free(t.down); + defer gpa.free(t.up); + const w: Weights = .{ .norm = t.norm, .down = t.down, .up = t.up, .inject = null }; + const wide = tiny.streams * tiny.dim; + const h: [wide]f32 = @splat(1); + var mixed: [tiny.dim]f32 = @splat(9); + var gates: [tiny.streams]f32 = undefined; + var scratch: [2 * wide + tiny.lowrank]f32 = undefined; + try std.testing.expectError(error.Length, pre(tiny, w, 1, 1e-6, h[1..], &mixed, null, &scratch)); + try std.testing.expectError(error.Length, pre(tiny, w, 1, 1e-6, &h, &mixed, &gates, &scratch)); // no inject weights + try std.testing.expectError(error.Length, pre(tiny, w, 1, 1e-6, &h, &mixed, null, scratch[0..3])); + try std.testing.expectEqual(@as(f32, 9), mixed[0]); + var out: [wide]f32 = undefined; + try std.testing.expectError(error.Length, post(tiny, &h, mixed[0..3], &gates, &out)); + try expand(tiny, &.{ 1, 2, 3, 4 }, &out); + try std.testing.expectEqualSlices(f32, &.{ 1, 2, 3, 4, 1, 2, 3, 4 }, &out); + try std.testing.expectError(error.Length, expand(tiny, &.{ 1, 2, 3 }, &out)); +} diff --git a/zig/src/families/flashnext/exl3/model/exl3_moe.zig b/zig/src/families/flashnext/exl3/model/exl3_moe.zig new file mode 100644 index 0000000000..6220a4b9c1 --- /dev/null +++ b/zig/src/families/flashnext/exl3/model/exl3_moe.zig @@ -0,0 +1,147 @@ +//! CPU f64 oracle for one Flash Next MoE block: router, routed experts, shared expert, combine. + +const std = @import("std"); +const format = @import("../pack/exl3_format.zig"); +const cpu = @import("exl3_cpu.zig"); + +/// One decoded EXL3 matrix W[k][n] in f64. `y[j] = sum_i x[i] * w[i * n + j]`. +pub const Dense = struct { + k: usize, + n: usize, + w: []f64, + + pub fn init(gpa: std.mem.Allocator, halves: u16, trellis: []const i16, tiles_k: usize, tiles_n: usize, suh: []const f16, svh: []const f16) !Dense { + const k = tiles_k * 16; + const n = tiles_n * 16; + if (k % format.hadamard_dim != 0 or n % format.hadamard_dim != 0) return error.Length; + if (suh.len != k or svh.len != n or trellis.len != tiles_k * tiles_n * @as(usize, halves) * 8) return error.Length; + const w = try gpa.alloc(f64, k * n); + errdefer gpa.free(w); + const scratch = try gpa.alloc(f64, 2 * k * n); + defer gpa.free(scratch); + format.dequantize(.{ .halves = halves }, .mul1, trellis, tiles_k, tiles_n, suh, svh, w, scratch); + return .{ .k = k, .n = n, .w = w }; + } + + pub fn deinit(d: *Dense, gpa: std.mem.Allocator) void { + gpa.free(d.w); + d.* = undefined; + } + + pub fn apply(d: Dense, x: []const f64, y: []f64) error{Length}!void { + if (x.len != d.k or y.len != d.n) return error.Length; + for (0..d.n) |j| { + var sum: f64 = 0; + for (0..d.k) |i| sum += x[i] * d.w[i * d.n + j]; + y[j] = sum; + } + } +}; + +/// down(silu(gate x) * up x). `gate_buf`, `up_buf`, `glu_buf` are `gate.n` long; `out` is `down.n` long. +pub fn expert(gate: Dense, up: Dense, down: Dense, x: []const f64, gate_buf: []f64, up_buf: []f64, glu_buf: []f64, out: []f64) error{Length}!void { + if (gate.n != up.n or down.k != gate.n or down.n != gate.k) return error.Length; + if (gate_buf.len != gate.n or up_buf.len != gate.n or glu_buf.len != gate.n) return error.Length; + try gate.apply(x, gate_buf); + try up.apply(x, up_buf); + try cpu.swiglu(gate_buf, up_buf, glu_buf, std.math.inf(f64)); + try down.apply(glu_buf, out); +} + +/// `logits` is `experts + 1` values, the last being the shared-gate dot. +pub fn routeFromLogits(logits: []const f64, topk: usize, picks: []u16, weights: []f64, shared_weight: *f64) error{Length}!void { + if (logits.len < 2 or topk == 0 or topk > logits.len - 1 or picks.len != topk or weights.len != topk) return error.Length; + const experts = logits.len - 1; + var top = logits[0]; + for (logits[0..experts]) |v| top = @max(top, v); + var total: f64 = 0; + for (logits[0..experts]) |v| total += @exp(v - top); + var taken: [1024]bool = @splat(false); + if (experts > taken.len) return error.Length; + var kept: f64 = 0; + for (0..topk) |k| { + var best: usize = 0; + var best_v = -std.math.inf(f64); + for (logits[0..experts], 0..) |v, i| { + if (!taken[i] and v > best_v) { + best_v = v; + best = i; + } + } + taken[best] = true; + picks[k] = @intCast(best); + weights[k] = @exp(best_v - top) / total; + kept += weights[k]; + } + for (weights) |*w| w.* /= kept; + shared_weight.* = 1.0 / (1.0 + @exp(-logits[experts])); +} + +/// Row-major `[rows][x.len]` f16 gate dotted with `x`, then the one-row shared gate. +pub fn logitsF16(gate: []const f16, rows: usize, shared_gate: []const f16, x: []const f64, logits: []f64) error{Length}!void { + if (gate.len != rows * x.len or shared_gate.len != x.len or logits.len != rows + 1) return error.Length; + for (0..rows) |r| { + var sum: f64 = 0; + for (x, 0..) |v, i| sum += v * @as(f64, gate[r * x.len + i]); + logits[r] = sum; + } + var s: f64 = 0; + for (x, 0..) |v, i| s += v * @as(f64, shared_gate[i]); + logits[rows] = s; +} + +test "full softmax then renormalise is the softmax over the selected logits" { + const route = @import("exl3_route.zig"); + var rng = std.Random.DefaultPrng.init(11); + const r = rng.random(); + for (0..20) |_| { + var logits: [65]f64 = undefined; + var logits32: [65]f32 = undefined; + for (&logits, &logits32) |*a, *b| { + const v: f32 = (r.float(f32) - 0.5) * 8; + a.* = v; + b.* = v; + } + var picks: [6]u16 = undefined; + var w: [6]f64 = undefined; + var sw: f64 = 0; + try routeFromLogits(&logits, 6, &picks, &w, &sw); + var picks32: [7]u16 = undefined; + var w32: [7]f32 = undefined; + var scratch: [64]f32 = undefined; + try route.select(&logits32, 6, &picks32, &w32, &scratch); + try std.testing.expectEqualSlices(u16, &picks, picks32[0..6]); + var sum: f64 = 0; + for (w, w32[0..6]) |a, b| { + sum += a; + // select rounds through bf16: 8 bits of mantissa, relative 2^-8 + try std.testing.expectApproxEqRel(a, @as(f64, b), 1.0 / 128.0); + } + try std.testing.expectApproxEqAbs(@as(f64, 1), sum, 1e-12); + try std.testing.expectApproxEqRel(sw, @as(f64, w32[6]), 1.0 / 128.0); + try std.testing.expectEqual(@as(u16, 64), picks32[6]); + } +} + +test "ties take the lower id and bad lengths store nothing" { + const logits = [_]f64{ 1, 3, 3, 0, 0 }; + var picks = [_]u16{ 9, 9 }; + var w = [_]f64{ 0, 0 }; + var sw: f64 = -1; + try routeFromLogits(&logits, 2, &picks, &w, &sw); + try std.testing.expectEqualSlices(u16, &.{ 1, 2 }, &picks); + try std.testing.expectEqual(@as(f64, 0.5), w[0]); + try std.testing.expectEqual(@as(f64, 0.5), sw); + var one = [_]u16{7}; + try std.testing.expectError(error.Length, routeFromLogits(&logits, 2, &one, &w, &sw)); + try std.testing.expectEqual(@as(u16, 7), one[0]); +} + +test "logitsF16 is a row dot and the shared gate is the last logit" { + const gate = [_]f16{ 1, 0, 0, 2, 0, 0 }; + const shared = [_]f16{ 0, 0, 3 }; + var logits: [3]f64 = undefined; + try logitsF16(&gate, 2, &shared, &.{ 1, 5, 7 }, &logits); + try std.testing.expectEqualSlices(f64, &.{ 1, 2, 21 }, &logits); + try std.testing.expectError(error.Length, logitsF16(&gate, 2, shared[0..2], &.{ 1, 5, 7 }, &logits)); +} diff --git a/zig/src/families/flashnext/exl3/model/exl3_norm.zig b/zig/src/families/flashnext/exl3/model/exl3_norm.zig new file mode 100644 index 0000000000..9009fa0c16 --- /dev/null +++ b/zig/src/families/flashnext/exl3/model/exl3_norm.zig @@ -0,0 +1,90 @@ +//! Centred RMSNorm apply. + +const std = @import("std"); +const scan = @import("../pack/exl3_scan.zig"); + +/// `rms = sqrt(mean(x²) + eps)`, then `out[i] = x[i] / rms * (weight[i] + offset)`. +pub fn apply(x: []const f32, weight: []const f32, means: []const f32, eps: f32, out: []f32) error{ Length, NormAmbiguous }!void { + const offset = try scan.centredOffset(means); + if (!std.math.isFinite(eps) or eps < 0) return error.Length; + if (x.len == 0 or x.len != weight.len or x.len != out.len) return error.Length; + + var sum_sq: f32 = 0; + for (x) |v| sum_sq += v * v; + const rms = @sqrt(sum_sq / @as(f32, @floatFromInt(x.len)) + eps); + for (x, weight, out) |xi, wi, *yi| yi.* = xi / rms * (wi + offset); +} + +test "a unit vector changes by exactly the weight difference" { + // Same means `centredOffset` classifies as gamma-1 (1) and gamma (0). + const near_zero = [_]f32{ -0.02, 0.01, 0.0, 0.04 }; + const near_one = [_]f32{ 1.0, 0.98, 1.02, 1.01 }; + // x = [1], eps = 0: mean(x²)+eps is 1, so the RMS is 1 and out = weight + offset. + const x = [_]f32{1}; + const light = [_]f32{2}; + const heavy = [_]f32{5}; + var zero_light = [_]f32{0}; + var zero_heavy = [_]f32{0}; + var one_light = [_]f32{0}; + var one_heavy = [_]f32{0}; + try apply(&x, &light, &near_zero, 0, &zero_light); + try apply(&x, &heavy, &near_zero, 0, &zero_heavy); + try apply(&x, &light, &near_one, 0, &one_light); + try apply(&x, &heavy, &near_one, 0, &one_heavy); + try std.testing.expectEqual(@as(f32, 3), zero_light[0]); + try std.testing.expectEqual(@as(f32, 6), zero_heavy[0]); + try std.testing.expectEqual(@as(f32, 2), one_light[0]); + try std.testing.expectEqual(@as(f32, 5), one_heavy[0]); + try std.testing.expectEqual(heavy[0] - light[0], zero_heavy[0] - zero_light[0]); + try std.testing.expectEqual(heavy[0] - light[0], one_heavy[0] - one_light[0]); + + // Two ones still have mean 1, so a forgotten division by length cannot pass. + const pair_x = [_]f32{ 1, 1 }; + const pair_w = [_]f32{ 2, 7 }; + var pair_out = [_]f32{ 0, 0 }; + try apply(&pair_x, &pair_w, &near_zero, 0, &pair_out); + try std.testing.expectEqual(@as(f32, 3), pair_out[0]); + try std.testing.expectEqual(@as(f32, 8), pair_out[1]); +} + +test "an ambiguous norm leaves out untouched" { + var out = [_]f32{ 9, 8 }; + const many: [65]f32 = @splat(0); + try std.testing.expectError(error.NormAmbiguous, apply(&.{ 1, 1 }, &.{ 1, 1 }, &many, 0, &out)); + try std.testing.expectEqual(@as(f32, 9), out[0]); + try std.testing.expectEqual(@as(f32, 8), out[1]); + + // Length also disagrees. The locator runs first, so this is not Length. + try std.testing.expectError(error.NormAmbiguous, apply(&.{1}, &.{}, &.{ 0.4, 0.9 }, -1, &out)); + try std.testing.expectEqual(@as(f32, 9), out[0]); + try std.testing.expectEqual(@as(f32, 8), out[1]); +} + +test "a length mismatch leaves out untouched" { + const near_zero = [_]f32{ -0.02, 0.01, 0.0, 0.04 }; + const near_one = [_]f32{ 1.0, 0.98, 1.02, 1.01 }; + var out = [_]f32{ 9, 8 }; + try std.testing.expectError(error.Length, apply(&.{}, &.{}, &near_zero, 0, &out)); + try std.testing.expectEqual(@as(f32, 9), out[0]); + try std.testing.expectEqual(@as(f32, 8), out[1]); + + try std.testing.expectError(error.Length, apply(&.{ 1, 1 }, &.{1}, &near_zero, 0, &out)); + try std.testing.expectEqual(@as(f32, 9), out[0]); + try std.testing.expectEqual(@as(f32, 8), out[1]); + + try std.testing.expectError(error.Length, apply(&.{1}, &.{1}, &near_one, 0, &out)); + try std.testing.expectEqual(@as(f32, 9), out[0]); + try std.testing.expectEqual(@as(f32, 8), out[1]); +} + +test "a negative eps leaves out untouched" { + const near_zero = [_]f32{ -0.02, 0.01, 0.0, 0.04 }; + const near_one = [_]f32{ 1.0, 0.98, 1.02, 1.01 }; + var out = [_]f32{9}; + try std.testing.expectError(error.Length, apply(&.{1}, &.{2}, &near_zero, -1, &out)); + try std.testing.expectEqual(@as(f32, 9), out[0]); + try std.testing.expectError(error.Length, apply(&.{1}, &.{2}, &near_one, std.math.nan(f32), &out)); + try std.testing.expectEqual(@as(f32, 9), out[0]); + try std.testing.expectError(error.Length, apply(&.{1}, &.{2}, &near_zero, std.math.inf(f32), &out)); + try std.testing.expectEqual(@as(f32, 9), out[0]); +} diff --git a/zig/src/families/flashnext/exl3/model/exl3_round.zig b/zig/src/families/flashnext/exl3/model/exl3_round.zig new file mode 100644 index 0000000000..9d75bb336d --- /dev/null +++ b/zig/src/families/flashnext/exl3/model/exl3_round.zig @@ -0,0 +1,599 @@ +//! One round's miss plan for a per-layer EXL3 expert pool (the host half of `cuda_exl3_pool.Pool.plan`). + +const std = @import("std"); +const offload = @import("../offload/exl3_offload.zig"); + +const shape = offload.shape; +const Vram = offload.Vram; +const empty_id = offload.Cache.empty_id; +pub const shared_id = offload.shared_id; + +/// A copy of routed expert `id` into pool line `line`; `evicted` is the expert the line held (empty_id: none). +pub const Copy = struct { id: u16, line: u32, evicted: u16 }; + +/// Ranges into `Round.ids` (the experts the pass computes) and `Round.copies` (what it stages first). +pub const Pass = struct { id_lo: u32, id_hi: u32, copy_lo: u32, copy_hi: u32 }; + +pub const Options = struct { + swap_every: u16 = 4, + swap_cap: u16 = 96, + observe: bool = true, + /// Bytes a swap round may prefetch (the copy budget); 0 disables prefetch. + prefetch_bytes: u64 = 0, + /// Device bytes of one line (prefetch budget unit). + line_bytes: u64 = 1, + /// Double buffering: experts (of this layer) whose lines a single-pass round should not evict when it has room + protect: []const u16 = &.{}, +}; + +/// Caller-owned storage for one round's plan, sized once by `init` for the pool's line count and pick capacity. +pub const Round = struct { + ids: []u16, + /// The line each id of `ids` is read from in its pass (a later pass of the same round may reuse that line). + lines: []u32, + copies: []Copy, + passes: []Pass, + routed: []u16, // the round's routed picks as u16 (ordering scratch) + misses: []u16, + snap: []u16, // line -> id before a pass (empty_id for empty lines) + keep: []u16, // pass ids plus protected ids (stageUsedAt's keep list) + /// The round spared `protect`'s lines (double buffering held). + protected: bool = false, + n_ids: usize = 0, + n_copies: usize = 0, + n_passes: usize = 0, + n_prefetch: usize = 0, + /// Routed slots (pick in 0..511) whose expert was held / missing when the round started. + hit_slots: u32 = 0, + miss_slots: u32 = 0, + shared_slots: u32 = 0, + /// Distinct routed experts, of which `held` were already in a line (`pending` of those still in flight). + unique: u32 = 0, + held: u32 = 0, + pending: u32 = 0, + /// Copies that replaced a held expert (mandatory and prefetch). + evictions: u32 = 0, + + pub fn init(gpa: std.mem.Allocator, lines: usize, max_picks: usize) !Round { + const uniq = @min(max_picks, shape.experts); + var r: Round = undefined; + r.ids = try gpa.alloc(u16, uniq); + errdefer gpa.free(r.ids); + r.lines = try gpa.alloc(u32, uniq); + errdefer gpa.free(r.lines); + r.copies = try gpa.alloc(Copy, uniq + shape.experts); + errdefer gpa.free(r.copies); + r.passes = try gpa.alloc(Pass, uniq + 1); + errdefer gpa.free(r.passes); + r.routed = try gpa.alloc(u16, max_picks); + errdefer gpa.free(r.routed); + r.misses = try gpa.alloc(u16, uniq); + errdefer gpa.free(r.misses); + r.snap = try gpa.alloc(u16, lines); + errdefer gpa.free(r.snap); + r.keep = try gpa.alloc(u16, uniq + lines); + r.reset(); + return r; + } + + pub fn deinit(r: *Round, gpa: std.mem.Allocator) void { + gpa.free(r.ids); + gpa.free(r.lines); + gpa.free(r.copies); + gpa.free(r.passes); + gpa.free(r.routed); + gpa.free(r.misses); + gpa.free(r.snap); + gpa.free(r.keep); + r.* = undefined; + } + + fn reset(r: *Round) void { + r.n_ids = 0; + r.n_copies = 0; + r.n_passes = 0; + r.n_prefetch = 0; + r.hit_slots = 0; + r.miss_slots = 0; + r.shared_slots = 0; + r.unique = 0; + r.held = 0; + r.pending = 0; + r.evictions = 0; + r.protected = false; + } + + pub fn passIds(r: *const Round, j: usize) []const u16 { + const p = r.passes[j]; + return r.ids[p.id_lo..p.id_hi]; + } + + pub fn passLines(r: *const Round, j: usize) []const u32 { + const p = r.passes[j]; + return r.lines[p.id_lo..p.id_hi]; + } + + pub fn passCopies(r: *const Round, j: usize) []const Copy { + const p = r.passes[j]; + return r.copies[p.copy_lo..p.copy_hi]; + } + + /// Copies no pass waits for (swap-round prefetch); their lines stay in flight until reaped. + pub fn prefetch(r: *const Round) []const Copy { + return r.copies[r.n_copies - r.n_prefetch .. r.n_copies]; + } + + /// Mandatory copies of all passes. + pub fn misses_(r: *const Round) usize { + return r.n_copies - r.n_prefetch; + } + + /// More than one pass: each pass reads its own masked picks (picks of other passes' experts set to -1). + pub fn masked(r: *const Round) bool { + return r.n_passes > 1; + } + + /// True when the round needs no copy at all (the experts need no upload). + pub fn allHits(r: *const Round) bool { + return r.n_copies == 0; + } +}; + +fn contains32(xs: []const u32, x: u32) bool { + return std.mem.indexOfScalar(u32, xs, x) != null; +} + +fn contains(xs: []const u16, id: u16) bool { + return std.mem.indexOfScalar(u16, xs, id) != null; +} + +const Order = struct { + counts: []const u32, + ids: []const u16, + fn appearances(ids: []const u16, id: u16) u32 { + var n: u32 = 0; + for (ids) |x| n += @intFromBool(x == id); + return n; + } + /// Lifetime count descending, then this round's appearances, then lower id (the policy's `usedFirst`). + fn hotter(o: Order, a: u16, b: u16) bool { + if (o.counts[a] != o.counts[b]) return o.counts[a] > o.counts[b]; + const aa = appearances(o.ids, a); + const ab = appearances(o.ids, b); + if (aa != ab) return aa > ab; + return a < b; + } +}; + +/// Lines whose copy is in flight and whose expert is not in `use`: they can be neither read nor evicted this pass. +fn blocked(v: Vram, use: []const u16) usize { + var n: usize = 0; + for (v.line, v.slots) |state, id| { + if (state == .inflight and !contains(use, id)) n += 1; + } + return n; +} + +fn snapshot(v: Vram, out: []u16) void { + for (v.slots, v.line, out) |id, state, *o| o.* = if (state == .empty) empty_id else id; +} + +/// Plans one round of `layer` (see the file header). +pub fn plan( + v: *Vram, + layer: u8, + picks: []const i32, + counts: *[shape.layers][shape.experts]u32, + now: u32, + round: u32, + opt: Options, + r: *Round, +) error{ Busy, Length }!void { + r.reset(); + if (picks.len > r.routed.len or layer >= shape.layers or v.slots.len != r.snap.len) return error.Length; + // the routed picks and the distinct experts, in first-appearance order + var n_routed: usize = 0; + var n_unique: usize = 0; + for (picks) |p| { + if (p < 0) continue; + if (p > shared_id) return error.Length; + const id: u16 = @intCast(p); + if (id == shared_id) { + r.shared_slots += 1; + continue; + } + r.routed[n_routed] = id; + n_routed += 1; + if (!contains(r.ids[0..n_unique], id)) { + r.ids[n_unique] = id; + n_unique += 1; + } + } + const routed = r.routed[0..n_routed]; + const uniq = r.ids[0..n_unique]; + // capacity check before anything is written: pass 0 holds the held experts, so at least one more line. + var held_n: usize = 0; + for (uniq) |id| held_n += @intFromBool(offload.lineOf(v.*, layer, id) != null); + if (n_unique > held_n and blocked(v.*, &.{}) >= v.slots.len) return error.Busy; + + if (opt.observe) offload.observe(&counts[layer], routed) catch return error.Length; + const row: []const u32 = &counts[layer]; + + // held first (pass 0), misses hottest first; ids[] is rewritten in pass order + var n_held: usize = 0; + var n_miss: usize = 0; + for (uniq) |id| { + if (offload.lineOf(v.*, layer, id)) |line| { + r.held += 1; + if (v.line[line] == .inflight) r.pending += 1 else _ = offload.touchAt(v, layer, id, now); + } else { + r.misses[n_miss] = id; + n_miss += 1; + } + } + for (routed) |id| { + if (offload.lineOf(v.*, layer, id) != null) r.hit_slots += 1 else r.miss_slots += 1; + } + r.unique = @intCast(n_unique); + // pack held ids at the front of ids (stable), then the sorted misses + for (uniq) |id| { + if (offload.lineOf(v.*, layer, id) != null) { + r.ids[n_held] = id; + n_held += 1; + } + } + const misses = r.misses[0..n_miss]; + std.mem.sort(u16, misses, Order{ .counts = row, .ids = routed }, Order.hotter); + @memcpy(r.ids[n_held..][0..n_miss], misses); + + // passes + var at: usize = 0; // next id of ids[] to place + var mi: usize = 0; // next miss + while (at < n_unique) { + const first = r.n_passes == 0; + const lo = at; + const take_held: usize = if (first) n_held else 0; + at += take_held; + // lines usable this pass: all but in-flight lines of experts outside this pass + const usable = v.slots.len - blocked(v.*, r.ids[lo..at]); + const room = usable - take_held; + const n_new = @min(room, n_miss - mi); + std.debug.assert(n_new > 0 or take_held > 0); // the capacity check above + const hi = at + n_new; + // double buffering: also spare the protected experts' resident lines when the single pass still fits + var keep: []const u16 = r.ids[lo..hi]; + if (first and n_new == n_miss and opt.protect.len > 0) { + var spare: usize = 0; + var k: usize = hi - lo; + @memcpy(r.keep[0..k], r.ids[lo..hi]); + for (opt.protect) |id| { + if (contains(r.keep[0..k], id) or k == r.keep.len) continue; + const line = offload.lineOf(v.*, layer, id) orelse continue; + if (v.line[line] != .resident) continue; + r.keep[k] = id; + k += 1; + spare += 1; + } + if (room >= n_new + spare) { + keep = r.keep[0..k]; + r.protected = true; + } + } + snapshot(v.*, r.snap); + const copy_lo = r.n_copies; + for (r.ids[at..hi]) |id| { + const line = offload.stageUsedAt(v, layer, id, keep, counts) catch return error.Busy; + const old = r.snap[line]; + if (old != empty_id) r.evictions += 1; + r.copies[r.n_copies] = .{ .id = id, .line = @intCast(line), .evicted = old }; + r.n_copies += 1; + } + // the device orders the next pass's copies after this pass's compute: these lines are resident from here on + for (r.copies[copy_lo..r.n_copies]) |c| offload.reapAt(v, layer, c.id, now) catch unreachable; + for (r.ids[lo..hi], r.lines[lo..hi]) |id, *ln| ln.* = @intCast(offload.lineOf(v.*, layer, id).?); + r.passes[r.n_passes] = .{ .id_lo = @intCast(lo), .id_hi = @intCast(hi), .copy_lo = @intCast(copy_lo), .copy_hi = @intCast(r.n_copies) }; + r.n_passes += 1; + at = hi; + mi += n_new; + } + r.n_ids = n_unique; + if (r.n_passes == 0) { // shared and -1 picks only: one pass with no routed expert (the shared expert still runs) + r.passes[0] = .{ .id_lo = 0, .id_hi = 0, .copy_lo = 0, .copy_hi = 0 }; + r.n_passes = 1; + } + + // swap-round prefetch: hottest non-held experts into empty or strictly colder lines this round does not use + if (opt.swap_every == 0 or round == 0 or round % opt.swap_every != 0 or opt.prefetch_bytes == 0 or opt.line_bytes == 0) return; + const budget = @min(@as(u64, opt.swap_cap), opt.prefetch_bytes / opt.line_bytes); + var staged: u64 = 0; + snapshot(v.*, r.snap); + while (staged < budget) { + // hottest id not held on this layer + var best: ?u16 = null; + var id: u16 = 0; + while (id < shape.experts) : (id += 1) { + if (row[id] == 0 or offload.lineOf(v.*, layer, id) != null) continue; + if (best == null or row[id] > row[best.?]) best = id; + } + const cand = best orelse break; + // the line it would take: an empty one, else the coldest resident outside this round + var victim: ?usize = null; + for (v.line, 0..) |state, i| { + if (state == .empty) { + victim = i; + break; + } + } + if (victim == null) { + for (v.line, v.slots, 0..) |state, held_id, i| { + if (state != .resident or contains(uniq_view(r), held_id)) continue; + if (victim == null or colder(v.*, layer, counts, i, victim.?)) victim = i; + } + const vi = victim orelse break; + if (offload.countAt(counts, layer, v.slots[vi]) >= row[cand]) break; // nothing colder than the candidate + } + const line = offload.stageUsedAt(v, layer, cand, uniq_view(r), counts) catch break; + std.debug.assert(line == victim.?); + const old = r.snap[line]; + if (old != empty_id) r.evictions += 1; + r.snap[line] = cand; + r.copies[r.n_copies] = .{ .id = cand, .line = @intCast(line), .evicted = old }; + r.n_copies += 1; + r.n_prefetch += 1; + staged += 1; + } +} + +fn uniq_view(r: *const Round) []const u16 { + return r.ids[0..r.n_ids]; +} + +/// The policy's victim order (`vramCold`): lower lifetime count, then older stamp, then the higher line index. +fn colder(v: Vram, layer: u8, counts: *const [shape.layers][shape.experts]u32, i: usize, b: usize) bool { + const ci = offload.countAt(counts, layer, v.slots[i]); + const cb = offload.countAt(counts, layer, v.slots[b]); + return ci < cb or (ci == cb and v.stamp[i] < v.stamp[b]) or (ci == cb and v.stamp[i] == v.stamp[b] and i > b); +} + +/// Smallest pool that never needs a second pass for `rows` decode rows: every row's routed picks must fit at once. +pub fn minLines(top_k: usize) usize { + return top_k; +} + +/// Lines a layer gets from `free` device bytes after `reserve` and `other` (KV, scratch: the family's budget) +pub fn linesFor(free: u64, reserve: u64, other: u64, line_bytes: u64, layers: usize, min_lines: usize) error{PoolTooSmall}!usize { + if (line_bytes == 0 or layers == 0) return error.PoolTooSmall; + const spend = std.math.sub(u64, free, reserve +| other) catch return error.PoolTooSmall; + const per = spend / (line_bytes * layers); + if (per < min_lines) return error.PoolTooSmall; + return @intCast(per); +} + +test { + _ = offload; +} + +// ------------------------------------------------------------------------------------------------ tests + +const testing = std.testing; +const Line = offload.Line; + +const Fix = struct { + slots: [16]u16 = undefined, + line: [16]Line = undefined, + stamp: [16]u32 = undefined, + lay: [16]u8 = undefined, + v: Vram = undefined, + counts: [shape.layers][shape.experts]u32 = @splat(@splat(0)), + r: Round = undefined, + + fn init(f: *Fix, lines: usize, max_picks: usize) !void { + f.v = Vram.init(f.slots[0..lines], f.line[0..lines], f.stamp[0..lines], f.lay[0..lines]); + f.counts = @splat(@splat(0)); + f.r = try Round.init(testing.allocator, lines, max_picks); + } + fn deinit(f: *Fix) void { + f.r.deinit(testing.allocator); + } + fn go(f: *Fix, picks: []const i32, now: u32, round: u32, opt: Options) !void { + return plan(&f.v, 3, picks, &f.counts, now, round, opt, &f.r); + } + fn resident(f: *const Fix, id: u16) bool { + return offload.residentAt(f.v, 3, id); + } +}; + +test "misses fill empty lines in one pass, then the same picks are all hits" { + var f: Fix = .{}; + try f.init(10, 32); + defer f.deinit(); + const picks = [_]i32{ 5, 9, 5, 512, -1, 7 }; + try f.go(&picks, 1, 1, .{}); + try testing.expectEqual(@as(usize, 1), f.r.n_passes); + try testing.expectEqual(@as(usize, 3), f.r.misses_()); + try testing.expectEqual(@as(u32, 4), f.r.miss_slots); + try testing.expectEqual(@as(u32, 1), f.r.shared_slots); + try testing.expectEqual(@as(u32, 0), f.r.evictions); + for (f.r.passCopies(0)) |c| try testing.expectEqual(empty_id, c.evicted); + // 5 appears twice: hottest miss, copied first + try testing.expectEqual(@as(u16, 5), f.r.passCopies(0)[0].id); + try testing.expect(!f.r.masked()); + try testing.expect(f.resident(5) and f.resident(9) and f.resident(7)); + try testing.expectEqual(@as(u32, 2), f.counts[3][5]); + + try f.go(&picks, 2, 2, .{}); + try testing.expect(f.r.allHits()); + try testing.expectEqual(@as(u32, 4), f.r.hit_slots); + try testing.expectEqual(@as(u32, 3), f.r.held); + try testing.expectEqual(@as(usize, 1), f.r.n_passes); + try testing.expectEqual(@as(usize, 3), f.r.passIds(0).len); + // a hit touched its line + try testing.expectEqual(@as(u32, 2), f.stamp[offload.lineOf(f.v, 3, 9).?]); +} + +test "synthetic prefill picks do not heat expert history" { + var f: Fix = .{}; + try f.init(10, 16); + defer f.deinit(); + try f.go(&[_]i32{ 7, 8, 9 }, 1, 1, .{ .observe = false }); + try testing.expectEqual(@as(usize, 3), f.r.misses_()); + try testing.expectEqual(@as(u32, 0), f.counts[3][7]); + try testing.expectEqual(@as(u32, 0), f.counts[3][8]); + try testing.expectEqual(@as(u32, 0), f.counts[3][9]); +} + +test "a miss evicts the coldest line outside the round, never a line the round reads" { + var f: Fix = .{}; + try f.init(3, 16); + defer f.deinit(); + try f.go(&[_]i32{ 1, 2, 3 }, 1, 1, .{}); + // heat 1 and 3, leave 2 cold + try f.go(&[_]i32{ 1, 3 }, 2, 2, .{}); + try f.go(&[_]i32{ 1, 3 }, 3, 3, .{}); + try f.go(&[_]i32{ 1, 4 }, 4, 4, .{}); + try testing.expectEqual(@as(usize, 1), f.r.misses_()); + const c = f.r.passCopies(0)[0]; + try testing.expectEqual(@as(u16, 4), c.id); + try testing.expectEqual(@as(u16, 2), c.evicted); + try testing.expect(f.resident(1) and f.resident(3) and f.resident(4) and !f.resident(2)); + // 3 is hot but not in this round, 1 is in it; the cold one now is 4 (count 1) vs 3 (count 3) -> 4 goes + try f.go(&[_]i32{ 1, 3, 6 }, 5, 5, .{}); + try testing.expectEqual(@as(u16, 4), f.r.passCopies(0)[0].evicted); +} + +test "more distinct experts than lines split into passes that each fit, every expert computed once" { + var f: Fix = .{}; + try f.init(4, 64); + defer f.deinit(); + try f.go(&[_]i32{ 20, 21 }, 1, 1, .{}); + var picks: [30]i32 = undefined; + for (&picks, 0..) |*p, i| p.* = @intCast(10 + (i % 11)); // 10..20, 11 distinct; 20 held + picks[29] = 21; // held too + try f.go(&picks, 2, 2, .{}); + try testing.expect(f.r.masked()); + try testing.expectEqual(@as(u32, 12), f.r.unique); + try testing.expectEqual(@as(u32, 2), f.r.held); + try testing.expectEqual(@as(usize, 10), f.r.misses_()); + var seen: [shape.experts]u8 = @splat(0); + for (0..f.r.n_passes) |j| { + const ids = f.r.passIds(j); + try testing.expect(ids.len >= 1 and ids.len <= 4); + // the lines a pass copies into are distinct + const cs = f.r.passCopies(j); + for (cs, 0..) |a, i| for (cs[i + 1 ..]) |b| try testing.expect(a.line != b.line); + for (ids) |id| seen[id] += 1; + } + for (10..22) |id| try testing.expectEqual(@as(u8, 1), seen[id]); + // each pass reads distinct lines; a copy of a later pass may land in a line an earlier pass read + var reused = false; + for (0..f.r.n_passes) |j| { + const ls = f.r.passLines(j); + for (ls, 0..) |a, i| for (ls[i + 1 ..]) |b| try testing.expect(a != b); + for (f.r.passCopies(j)) |c| { + try testing.expect(contains32(ls, c.line)); + if (j > 0) for (0..j) |k| { + if (contains32(f.r.passLines(k), c.line)) reused = true; + }; + } + } + try testing.expect(reused); + // pass 0 holds the two held experts first + try testing.expect(contains(f.r.passIds(0), 20) and contains(f.r.passIds(0), 21)); + try testing.expectEqual(@as(usize, 3), f.r.n_passes); // 2 held + 2 misses, then 4, then 4 + // the last pass's experts are resident after the round + for (f.r.passIds(f.r.n_passes - 1)) |id| try testing.expect(f.resident(id)); +} + +test "an in-flight prefetch line is not a victim and counts as held when the round asks for it" { + var f: Fix = .{}; + try f.init(3, 16); + defer f.deinit(); + try f.go(&[_]i32{ 1, 2, 3 }, 1, 1, .{}); + // a prefetch of 9 into line of 2 (made by hand: stage it, leave it in flight) + const line = offload.lineOf(f.v, 3, 2).?; + f.slots[line] = 9; + f.line[line] = .inflight; + try f.go(&[_]i32{ 1, 4 }, 2, 2, .{}); + // 4 cannot take 9's in-flight line; it takes 3's (not in the round) + try testing.expectEqual(@as(u16, 3), f.r.passCopies(0)[0].evicted); + try testing.expectEqual(Line.inflight, f.line[line]); + try f.go(&[_]i32{ 9, 1 }, 3, 3, .{}); + try testing.expect(f.r.allHits()); + try testing.expectEqual(@as(u32, 1), f.r.pending); + // every line in flight for another expert: busy, and nothing written + for (f.line[0..3]) |*l| l.* = .inflight; + const before = f.counts[3][7]; + try testing.expectError(error.Busy, f.go(&[_]i32{7}, 4, 4, .{})); + try testing.expectEqual(before, f.counts[3][7]); +} + +test "double buffering spares the previous round's lines when the round fits without them" { + var f: Fix = .{}; + try f.init(6, 16); + defer f.deinit(); + // previous round 1,2,3 (hot), older residents 4,5 (cold), one empty line + try f.go(&[_]i32{ 4, 5 }, 1, 1, .{}); + for (0..3) |i| try f.go(&[_]i32{ 1, 2, 3 }, @intCast(2 + i), @intCast(2 + i), .{}); + // two misses: with 1, 2, 3 spared they go to the empty line and the colder of 4, 5 + const prev = [_]u16{ 1, 2, 3 }; + try f.go(&[_]i32{ 7, 8 }, 5, 5, .{ .protect = &prev }); + try testing.expect(f.r.protected); + for (f.r.passCopies(0)) |c| try testing.expect(!contains(&prev, c.evicted)); + try testing.expect(f.resident(1) and f.resident(2) and f.resident(3)); + // a round that only fits by evicting a protected line ignores the protection + try f.go(&[_]i32{ 9, 10, 11, 12 }, 6, 6, .{ .protect = &[_]u16{ 7, 8, 1, 2, 3 } }); + try testing.expect(!f.r.protected); + try testing.expectEqual(@as(usize, 1), f.r.n_passes); +} + +test "the shared expert and -1 are never staged; a pick above 512 is refused" { + var f: Fix = .{}; + try f.init(10, 16); + defer f.deinit(); + try f.go(&[_]i32{ 512, -1, 512 }, 1, 1, .{}); + try testing.expectEqual(@as(usize, 0), f.r.n_copies); + try testing.expectEqual(@as(u32, 2), f.r.shared_slots); + try testing.expectEqual(@as(usize, 1), f.r.n_passes); + try testing.expectEqual(@as(usize, 0), f.r.passIds(0).len); + for (f.slots[0..10]) |id| try testing.expect(id != shared_id); + try testing.expectError(error.Length, f.go(&[_]i32{513}, 2, 2, .{})); +} + +test "swap rounds prefetch the hottest non-held experts within the copy budget, only into colder lines" { + var f: Fix = .{}; + try f.init(4, 16); + defer f.deinit(); + const opt: Options = .{ .swap_every = 4, .prefetch_bytes = 200, .line_bytes = 100 }; + // lifetime counts: 30 hot (5), 31 warm (3), never in a line yet + f.counts[3][30] = 5; + f.counts[3][31] = 3; + f.counts[3][32] = 1; + try f.go(&[_]i32{ 1, 2 }, 1, 1, opt); + try testing.expectEqual(@as(usize, 0), f.r.n_prefetch); // round 1: not a swap round + try f.go(&[_]i32{ 1, 2 }, 4, 4, opt); + // two empty lines, budget 2: 30 and 31 in flight + try testing.expectEqual(@as(usize, 2), f.r.n_prefetch); + try testing.expectEqual(@as(u16, 30), f.r.prefetch()[0].id); + try testing.expectEqual(@as(u16, 31), f.r.prefetch()[1].id); + try testing.expectEqual(Line.inflight, f.line[f.r.prefetch()[0].line]); + try testing.expect(!f.resident(30)); + try offload.reapAt(&f.v, 3, 30, 5); + try offload.reapAt(&f.v, 3, 31, 5); + // round 8: 32 (count 1) is not hotter than any resident outside the round (30: 5, 31: 3) -> nothing evicted + try f.go(&[_]i32{ 1, 2 }, 8, 8, opt); + try testing.expectEqual(@as(usize, 0), f.r.n_prefetch); + // heat 32 above 31: it replaces 31, never 1 or 2 (in this round) + f.counts[3][32] = 9; + try f.go(&[_]i32{ 1, 2 }, 12, 12, opt); + try testing.expectEqual(@as(usize, 1), f.r.n_prefetch); + try testing.expectEqual(@as(u16, 31), f.r.prefetch()[0].evicted); +} + +test "line budget: free memory after reserve and the family's other bytes, at least min lines a layer" { + const line: u64 = 1_843_200; + // 12 GiB free, 1 GiB reserve, 5.5 GiB for KV and scratch: 5.5 GiB / (48 x 1,843,200 B) = 66 lines a layer + try testing.expectEqual(@as(usize, 66), try linesFor(12 << 30, 1 << 30, (11 << 30) / 2, line, 48, 10)); + try testing.expectEqual(@as(usize, 10), try linesFor(10 * 48 * line + 100, 100, 0, line, 48, 10)); + try testing.expectError(error.PoolTooSmall, linesFor(10 * 48 * line - 1, 0, 0, line, 48, 10)); + try testing.expectError(error.PoolTooSmall, linesFor(1 << 30, 2 << 30, 0, line, 48, 10)); +} diff --git a/zig/src/families/flashnext/exl3/model/exl3_route.zig b/zig/src/families/flashnext/exl3/model/exl3_route.zig new file mode 100644 index 0000000000..c6e7ab22df --- /dev/null +++ b/zig/src/families/flashnext/exl3/model/exl3_route.zig @@ -0,0 +1,94 @@ +//! One row of the Flash Next router. v0.6.6 `src/tensorfold/cuda/moe.py` `_topk_rows`: fp32 logits. + +const std = @import("std"); + +/// `logits` is `experts + 1` fp32 values. +pub fn select( + logits: []const f32, + topk: usize, + picks: []u16, + weights: []f32, + scratch: []f32, +) error{Length}!void { + if (logits.len < 2) return error.Length; + const experts = logits.len - 1; + if (experts > std.math.maxInt(u16)) return error.Length; + if (topk == 0 or topk > experts) return error.Length; + if (picks.len != topk + 1 or weights.len != topk + 1 or scratch.len != experts) return error.Length; + + @memcpy(scratch, logits[0..experts]); + var top = scratch[0]; + for (scratch[1..]) |v| top = @max(top, v); + + var total: f32 = 0; + for (0..topk) |k| { + var best_i: usize = 0; + var best = scratch[0]; + for (scratch[1..], 1..) |v, i| { + if (v > best) { + best = v; + best_i = i; + } + } + const ex = @exp(best - top); + picks[k] = @intCast(best_i); + weights[k] = ex; + total += ex; + scratch[best_i] = -std.math.inf(f32); + } + for (weights[0..topk]) |*w| w.* = bf16(w.* / total); + + const shared: u16 = @intCast(experts); + const gate = bf16(logits[experts]); + picks[topk] = shared; + weights[topk] = bf16(1.0 / (1.0 + @exp(-gate))); +} + +/// fp32 through bf16 and back. NaN stays a NaN. Same rounding as the contract helper. +fn bf16(v: f32) f32 { + const u: u32 = @bitCast(v); + const h: u16 = if (std.math.isNan(v)) + @truncate((u >> 16) | 0x40) + else + @truncate((u +% 0x7fff +% ((u >> 16) & 1)) >> 16); + return @bitCast(@as(u32, h) << 16); +} + +test "ties take the lower id, and the shared expert is the last slot" { + const logits = [_]f32{ 1, 3, 3, 0, 0 }; + var picks = [_]u16{ 9, 9, 9 }; + var weights = [_]f32{ 9, 9, 9 }; + var scratch: [4]f32 = undefined; + try select(&logits, 2, &picks, &weights, &scratch); + + try std.testing.expectEqual(@as(u16, 1), picks[0]); + try std.testing.expectEqual(@as(u16, 2), picks[1]); + try std.testing.expectEqual(@as(u16, 4), picks[2]); + try std.testing.expectEqual(@as(f32, 0.5), weights[0]); + try std.testing.expectEqual(@as(f32, 0.5), weights[1]); + try std.testing.expectEqual(@as(f32, 0.5), weights[2]); +} + +test "one winner keeps the whole routed weight" { + const logits = [_]f32{ 5, 5, 1, 2 }; + var picks = [_]u16{ 9, 9 }; + var weights = [_]f32{ 9, 9 }; + var scratch: [3]f32 = undefined; + try select(&logits, 1, &picks, &weights, &scratch); + try std.testing.expectEqual(@as(u16, 0), picks[0]); + try std.testing.expectEqual(@as(u16, 3), picks[1]); + try std.testing.expectEqual(@as(f32, 1), weights[0]); + const gate = bf16(2); + const want = bf16(1.0 / (1.0 + @exp(-gate))); + try std.testing.expectEqual(want, weights[1]); +} + +test "a short buffer is refused before picks change" { + const logits = [_]f32{ 1, 2, 0 }; + var picks = [_]u16{ 7, 7 }; + var weights = [_]f32{ 3, 3 }; + var scratch: [2]f32 = undefined; + try std.testing.expectError(error.Length, select(&logits, 2, picks[0..1], &weights, &scratch)); + try std.testing.expectEqual(@as(u16, 7), picks[0]); + try std.testing.expectEqual(@as(f32, 3), weights[0]); +} diff --git a/zig/src/families/flashnext/exl3/model/exl3_router.zig b/zig/src/families/flashnext/exl3/model/exl3_router.zig new file mode 100644 index 0000000000..48ee9f791f --- /dev/null +++ b/zig/src/families/flashnext/exl3/model/exl3_router.zig @@ -0,0 +1,252 @@ +//! Rank 0 of the Flash Next router. + +const std = @import("std"); +const route = @import("exl3_route.zig"); + +pub const gate_rows: usize = 512; +pub const gate_cols: usize = 2560; +pub const gate_bytes: usize = gate_rows * gate_cols * @sizeOf(u16); + +pub const Error = error{ + NotRank0, + EmptyGate, + GateBytes, + MissingSharedGate, + Length, +}; + +comptime { + std.debug.assert(gate_bytes == 2_621_440); + std.debug.assert((gate_rows + 1) * gate_cols * @sizeOf(u16) == 2_626_560); +} + +/// Storage type of the two replicated gate tensors. +pub const Dtype = enum { bf16, f16 }; + +fn dot(dtype: Dtype, weight: []const u8, x: []const f32) f32 { + var sum: f32 = 0; + for (x, 0..) |v, i| { + const h = std.mem.readInt(u16, weight[i * 2 ..][0..2], .little); + const w: f32 = switch (dtype) { + .bf16 => @bitCast(@as(u32, h) << 16), + .f16 => @floatCast(@as(f16, @bitCast(h))), + }; + sum += w * v; + } + return sum; +} + +/// `gate` is row-major bf16, exactly `gate_bytes`, times `hidden` (`gate_cols`). +pub fn launch( + rank: usize, + gate: []const u8, + hidden: []const f32, + shared_gate: []const u8, + shared_x: []const f32, + topk: usize, + logits: []f32, + picks: []u16, + weights: []f32, + scratch: []f32, +) Error!void { + return launchDtype(.bf16, rank, gate, hidden, shared_gate, shared_x, topk, logits, picks, weights, scratch); +} + +/// `launch` with the gate tensors' storage type named. +pub fn launchDtype( + dtype: Dtype, + rank: usize, + gate: []const u8, + hidden: []const f32, + shared_gate: []const u8, + shared_x: []const f32, + topk: usize, + logits: []f32, + picks: []u16, + weights: []f32, + scratch: []f32, +) Error!void { + if (rank != 0) return error.NotRank0; + if (gate.len == 0) return error.EmptyGate; + if (gate.len != gate_bytes) return error.GateBytes; + if (shared_gate.len == 0) return error.MissingSharedGate; + if (shared_gate.len % 2 != 0) return error.Length; + const shared_n = shared_gate.len / 2; + if (shared_x.len != shared_n) return error.Length; + if (hidden.len != gate_cols) return error.Length; + if (logits.len != gate_rows + 1) return error.Length; + if (topk == 0 or topk > gate_rows) return error.Length; + if (picks.len != topk + 1 or weights.len != topk + 1 or scratch.len != gate_rows) return error.Length; + + const row_bytes = gate_cols * @sizeOf(u16); + for (0..gate_rows) |r| { + const row = gate[r * row_bytes ..][0..row_bytes]; + logits[r] = dot(dtype, row, hidden); + } + logits[gate_rows] = dot(dtype, shared_gate, shared_x); + try route.select(logits, topk, picks, weights, scratch); +} + +const one_bf16: u16 = 0x3f80; +const winner: usize = 42; +const shared_span: usize = 3; + +fn paint(gate: []u8) void { + @memset(gate, 0); + const at = winner * gate_cols * @sizeOf(u16); + std.mem.writeInt(u16, gate[at..][0..2], one_bf16, .little); +} + +fn sharedGate() [shared_span * 2]u8 { + var bytes: [shared_span * 2]u8 = @splat(0); + for (0..shared_span) |i| std.mem.writeInt(u16, bytes[i * 2 ..][0..2], one_bf16, .little); + return bytes; +} + +fn hiddenVec() [gate_cols]f32 { + var hidden: [gate_cols]f32 = @splat(0); + hidden[0] = 1; + return hidden; +} + +fn expectUntouched(logits: []const f32, picks: []const u16, weights: []const f32) !void { + try std.testing.expectEqual(@as(f32, 8), logits[0]); + try std.testing.expectEqual(@as(f32, 8), logits[gate_rows]); + try std.testing.expectEqualSlices(u16, &.{ 9, 9 }, picks); + try std.testing.expectEqualSlices(f32, &.{ 3, 3 }, weights); +} + +test "one nonzero gate row and a shared span of 3 select id 512" { + try std.testing.expect(shared_span != gate_cols); + try std.testing.expect(shared_span != 640); + try std.testing.expect(shared_span != 1280); + try std.testing.expect(shared_span != 61_440); + try std.testing.expectEqual(@as(usize, 2_621_440), gate_bytes); + + const gpa = std.testing.allocator; + const gate = try gpa.alloc(u8, gate_bytes); + defer gpa.free(gate); + paint(gate); + + const hidden = hiddenVec(); + const shared_gate = sharedGate(); + const shared_x = [_]f32{ 0, 5, 4 }; + var logits: [gate_rows + 1]f32 = @splat(8); + var picks = [_]u16{ 9, 9 }; + var weights = [_]f32{ 3, 3 }; + var scratch: [gate_rows]f32 = undefined; + try launch(0, gate, &hidden, &shared_gate, &shared_x, 1, &logits, &picks, &weights, &scratch); + + try std.testing.expectEqual(@as(f32, 0), logits[0]); + try std.testing.expectEqual(@as(f32, 1), logits[winner]); + try std.testing.expectEqual(@as(f32, 0), logits[gate_rows - 1]); + try std.testing.expectEqual(@as(f32, 9), logits[gate_rows]); + try std.testing.expectEqual(@as(u16, winner), picks[0]); + try std.testing.expectEqual(@as(u16, 512), picks[1]); + try std.testing.expectEqual(@as(f32, 1), weights[0]); +} + +test "rank 1 leaves the pick buffer untouched" { + const gpa = std.testing.allocator; + const gate = try gpa.alloc(u8, gate_bytes); + defer gpa.free(gate); + paint(gate); + + const hidden = hiddenVec(); + const shared_gate = sharedGate(); + const shared_x = [_]f32{ 0, 5, 4 }; + var logits: [gate_rows + 1]f32 = @splat(8); + var picks = [_]u16{ 9, 9 }; + var weights = [_]f32{ 3, 3 }; + var scratch: [gate_rows]f32 = @splat(4); + try std.testing.expectError(error.NotRank0, launch(1, gate, &hidden, &shared_gate, &shared_x, 1, &logits, &picks, &weights, &scratch)); + try expectUntouched(&logits, &picks, &weights); + try std.testing.expectEqual(@as(f32, 4), scratch[0]); +} + +test "a missing shared gate does not call select" { + const gpa = std.testing.allocator; + const gate = try gpa.alloc(u8, gate_bytes); + defer gpa.free(gate); + paint(gate); + + const hidden = hiddenVec(); + var logits: [gate_rows + 1]f32 = @splat(8); + var picks = [_]u16{ 9, 9 }; + var weights = [_]f32{ 3, 3 }; + var scratch: [gate_rows]f32 = @splat(4); + try std.testing.expectError(error.MissingSharedGate, launch(0, gate, &hidden, &.{}, &.{}, 1, &logits, &picks, &weights, &scratch)); + try expectUntouched(&logits, &picks, &weights); + try std.testing.expectEqual(@as(f32, 4), scratch[0]); +} + +test "a 513-row gate buffer is refused" { + const gpa = std.testing.allocator; + const wide_bytes = (gate_rows + 1) * gate_cols * @sizeOf(u16); + try std.testing.expectEqual(@as(usize, 2_626_560), wide_bytes); + const gate = try gpa.alloc(u8, wide_bytes); + defer gpa.free(gate); + + const hidden = hiddenVec(); + const shared_gate = sharedGate(); + const shared_x = [_]f32{ 0, 5, 4 }; + var logits: [gate_rows + 1]f32 = @splat(8); + var picks = [_]u16{ 9, 9 }; + var weights = [_]f32{ 3, 3 }; + var scratch: [gate_rows]f32 = @splat(4); + try std.testing.expectError(error.GateBytes, launch(0, gate, &hidden, &shared_gate, &shared_x, 1, &logits, &picks, &weights, &scratch)); + try expectUntouched(&logits, &picks, &weights); +} + +test "an empty gate does not launch" { + const hidden = hiddenVec(); + const shared_gate = sharedGate(); + const shared_x = [_]f32{ 0, 5, 4 }; + var logits: [gate_rows + 1]f32 = @splat(8); + var picks = [_]u16{ 9, 9 }; + var weights = [_]f32{ 3, 3 }; + var scratch: [gate_rows]f32 = @splat(4); + try std.testing.expectError(error.EmptyGate, launch(0, &.{}, &hidden, &shared_gate, &shared_x, 1, &logits, &picks, &weights, &scratch)); + try expectUntouched(&logits, &picks, &weights); + try std.testing.expectEqual(@as(f32, 4), scratch[0]); +} + +test "a shared-span mismatch stores nothing" { + const gpa = std.testing.allocator; + const gate = try gpa.alloc(u8, gate_bytes); + defer gpa.free(gate); + @memset(gate, 0); + + const hidden = hiddenVec(); + const shared_gate = sharedGate(); + const short_x = [_]f32{1}; + var logits: [gate_rows + 1]f32 = @splat(8); + var picks = [_]u16{ 9, 9 }; + var weights = [_]f32{ 3, 3 }; + var scratch: [gate_rows]f32 = @splat(4); + try std.testing.expectError(error.Length, launch(0, gate, &hidden, &shared_gate, &short_x, 1, &logits, &picks, &weights, &scratch)); + try expectUntouched(&logits, &picks, &weights); +} + +test "the same bits read as F16 are not the bf16 value" { + const gpa = std.testing.allocator; + const gate = try gpa.alloc(u8, gate_bytes); + defer gpa.free(gate); + @memset(gate, 0); + const one_f16: u16 = 0x3c00; + std.mem.writeInt(u16, gate[7 * gate_cols * 2 ..][0..2], one_f16, .little); + var shared: [gate_cols * 2]u8 = @splat(0); + std.mem.writeInt(u16, shared[0..2], one_f16, .little); + var hidden: [gate_cols]f32 = @splat(0); + hidden[0] = 1; + var logits: [gate_rows + 1]f32 = undefined; + var picks = [_]u16{ 9, 9 }; + var weights = [_]f32{ 0, 0 }; + var scratch: [gate_rows]f32 = undefined; + try launchDtype(.f16, 0, gate, &hidden, &shared, &hidden, 1, &logits, &picks, &weights, &scratch); + try std.testing.expectEqual(@as(f32, 1), logits[7]); + try std.testing.expectEqual(@as(f32, 1), logits[gate_rows]); + try std.testing.expectEqual(@as(u16, 7), picks[0]); + try launchDtype(.bf16, 0, gate, &hidden, &shared, &hidden, 1, &logits, &picks, &weights, &scratch); + try std.testing.expect(logits[7] != 1); +} diff --git a/zig/src/families/flashnext/exl3/model/exl3_telemetry.zig b/zig/src/families/flashnext/exl3/model/exl3_telemetry.zig new file mode 100644 index 0000000000..8774d10242 --- /dev/null +++ b/zig/src/families/flashnext/exl3/model/exl3_telemetry.zig @@ -0,0 +1,226 @@ +const std = @import("std"); + +pub const Phase = enum { prefill, decode, verify, unknown }; +pub const Rank = struct { + rank: usize, + world: usize, + + pub fn validate(owner: Rank) !void { + if (owner.world == 0 or owner.rank >= owner.world) return error.TelemetryRank; + } +}; + +var next_identity: std.atomic.Value(u64) = .init(0); + +pub fn identity() !u64 { + var previous = next_identity.load(.monotonic); + while (true) { + const next = try std.math.add(u64, previous, 1); + previous = next_identity.cmpxchgWeak(previous, next, .monotonic, .monotonic) orelse return next; + } +} + +pub const Device = struct { + plans: u64 = 0, + hit_picks: u64 = 0, + miss_picks: u64 = 0, + copies: u64 = 0, + evictions: u64 = 0, + inplace_experts: u64 = 0, + overflow_plans: u64 = 0, + + pub fn pending(current: anytype, seen: @TypeOf(current)) Device { + return .{ .plans = current.rounds -% seen.rounds, .hit_picks = current.hit_slots -% seen.hit_slots, .miss_picks = current.miss_slots -% seen.miss_slots, .copies = current.copies -% seen.copies, .evictions = current.evictions -% seen.evictions, .inplace_experts = current.inplace -% seen.inplace, .overflow_plans = current.overflow -% seen.overflow }; + } + + pub fn pendingBounded(current: anytype, seen: @TypeOf(current), max_picks: u64) !Device { + const delta = pending(current, seen); + const bound = try std.math.mul(u64, delta.plans, max_picks); + if (max_picks == 0 or bound >= @as(u64, 1) << 32) return error.TelemetryExpertIntervalTooLong; + if (try std.math.add(u64, delta.hit_picks, delta.miss_picks) > bound or try std.math.add(u64, delta.copies, delta.inplace_experts) > bound or delta.overflow_plans > delta.plans) return error.TelemetryInconsistentCounters; + return delta; + } + + pub fn record(total: *Device, valid: *bool, current: anytype, seen: @TypeOf(current), max_picks: u64) void { + if (!valid.*) return; + const delta = pendingBounded(current, seen, max_picks) catch { + valid.* = false; + return; + }; + add(Device, total, delta) catch { + valid.* = false; + }; + } +}; + +pub fn add(comptime Counters: type, total: *Counters, contribution: Counters) !void { + var result = total.*; + inline for (@typeInfo(Counters).@"struct".field_names) |field| @field(result, field) = try std.math.add(u64, @field(result, field), @field(contribution, field)); + total.* = result; +} + +pub fn difference(comptime Counters: type, before: Counters, after: Counters) !Counters { + var result: Counters = undefined; + inline for (@typeInfo(Counters).@"struct".field_names) |field| { + if (@field(after, field) < @field(before, field)) return error.TelemetryCounterRegression; + @field(result, field) = @field(after, field) - @field(before, field); + } + return result; +} + +pub const Experts = struct { + host_plans: u64 = 0, + host_hit_picks: u64 = 0, + host_miss_picks: u64 = 0, + host_mandatory_copies: u64 = 0, + host_prefetch_copies: u64 = 0, + pool_dma_copies: u64 = 0, + pool_dma_requested_bytes: u64 = 0, + device_plans: u64 = 0, + device_hit_picks: u64 = 0, + device_miss_picks: u64 = 0, + device_copies: u64 = 0, + device_inplace_experts: u64 = 0, + device_overflow_plans: u64 = 0, + mapped_copy_payload_min_bytes: u64 = 0, + mapped_copy_payload_max_bytes: u64 = 0, +}; + +pub const Snapshot = struct { + rank: Rank, + scope: u64, + experts: Experts = .{}, + expert_pools: usize = 0, + layers: []LayerSample = &.{}, + + pub fn deinit(snapshot: *Snapshot, allocator: std.mem.Allocator) void { + allocator.free(snapshot.layers); + snapshot.* = undefined; + } +}; + +pub const LayerCounts = struct { + host_plans: u64 = 0, + host_hit_picks: u64 = 0, + host_miss_picks: u64 = 0, + host_mandatory_copies: u64 = 0, + host_prefetch_copies: u64 = 0, + pool_dma_copies: u64 = 0, + pool_dma_requested_bytes: u64 = 0, + device_plans: u64 = 0, + device_hit_picks: u64 = 0, + device_miss_picks: u64 = 0, + device_copies: u64 = 0, + device_inplace_experts: u64 = 0, + device_overflow_plans: u64 = 0, + mapped_copy_payload_min_bytes: u64 = 0, + mapped_copy_payload_max_bytes: u64 = 0, + + pub fn fromExperts(experts: Experts) LayerCounts { + var result: LayerCounts = undefined; + inline for (@typeInfo(LayerCounts).@"struct".field_names) |field| @field(result, field) = @field(experts, field); + return result; + } +}; + +pub const LayerSample = struct { + owner: u64, + kind: enum { decoder }, + layer: usize, + lines: usize, + counts: LayerCounts = .{}, + + fn samePool(first: LayerSample, last: LayerSample) bool { + return first.owner != 0 and first.owner == last.owner and first.kind == last.kind and first.layer == last.layer and first.lines == last.lines; + } +}; + +pub fn LayerCursor(comptime capacity: usize) type { + return struct { + lock: std.atomic.Mutex = .unlocked, + scope: u64 = 0, + rank: Rank = .{ .rank = 0, .world = 1 }, + samples: [capacity]LayerSample = undefined, + len: usize = 0, + + pub fn claim(cursor: *@This(), allocator: std.mem.Allocator, before: Snapshot, after: Snapshot) ![]LayerSample { + while (!cursor.lock.tryLock()) std.atomic.spinLoopHint(); + defer cursor.lock.unlock(); + if (before.layers.len != after.layers.len or after.layers.len > capacity) return error.TelemetryOwnerChanged; + if (!std.meta.eql(before.rank, after.rank) or before.scope != after.scope) return error.TelemetryOwnerChanged; + if (cursor.len != 0 and (cursor.scope != after.scope or !std.meta.eql(cursor.rank, after.rank) or cursor.len != after.layers.len)) return error.TelemetryOwnerChanged; + const result = try allocator.alloc(LayerSample, after.layers.len); + errdefer allocator.free(result); + for (before.layers, after.layers, result, 0..) |first, last, *delta, index| { + if (!first.samePool(last)) return error.TelemetryOwnerChanged; + for (after.layers[0..index]) |previous| if (previous.owner == last.owner) return error.TelemetryDuplicateOwner; + if (cursor.len != 0 and !cursor.samples[index].samePool(last)) return error.TelemetryOwnerChanged; + delta.* = last; + delta.counts = try difference(LayerCounts, first.counts, last.counts); + } + for (after.layers, 0..) |last, index| { + cursor.samples[index] = last; + } + cursor.scope = after.scope; + cursor.rank = after.rank; + cursor.len = after.layers.len; + return result; + } + }; +} + +pub const Interval = struct { + phase: Phase, + model_rows: u64, + steps: u64, +}; + +pub fn fraction(misses: u64, hits: u64) !?f64 { + const denominator = try std.math.add(u64, misses, hits); + return if (denominator == 0) null else @as(f64, @floatFromInt(misses)) / @as(f64, @floatFromInt(denominator)); +} + +pub const Report = struct { + rank: Rank, + interval: Interval, + experts: Experts, + expert_pools: usize = 0, + layers: []const LayerSample = &.{}, + + pub fn write(report: Report, writer: *std.Io.Writer) !void { + try report.header(writer, "expert"); + inline for (@typeInfo(Experts).@"struct".field_names) |field| try writer.print(" {s}={d}", .{ field, @field(report.experts, field) }); + try writer.writeByte('\n'); + for (report.layers) |layer| { + try report.header(writer, "expert-layer"); + try writer.print(" pool={d} kind={s} layer={d} lines={d}", .{ layer.owner, @tagName(layer.kind), layer.layer, layer.lines }); + inline for (@typeInfo(LayerCounts).@"struct".field_names) |field| try writer.print(" {s}={d}", .{ field, @field(layer.counts, field) }); + try writer.writeByte('\n'); + } + } + + fn header(report: Report, writer: *std.Io.Writer, cache: []const u8) !void { + try writer.print("EXL3-CACHE v=1 rank={d} world={d} phase={s} rows={d} steps={d} expert_pools={d} cache={s}", .{ report.rank.rank, report.rank.world, @tagName(report.interval.phase), report.interval.model_rows, report.interval.steps, report.expert_pools, cache }); + } + + pub fn formatAlloc(report: Report, allocator: std.mem.Allocator) ![]u8 { + var writer: std.Io.Writer.Allocating = .init(allocator); + errdefer writer.deinit(); + try report.write(&writer.writer); + return writer.toOwnedSlice(); + } + + pub fn between(before: Snapshot, after: Snapshot, interval: Interval) !Report { + try before.rank.validate(); + try after.rank.validate(); + if (!std.meta.eql(before.rank, after.rank) or before.scope != after.scope or before.expert_pools != after.expert_pools) return error.TelemetryOwnerChanged; + if (interval.model_rows == 0 or interval.steps == 0) return error.TelemetryDenominator; + const experts = try difference(Experts, before.experts, after.experts); + if (experts.device_overflow_plans != 0) return error.TelemetryExpertOverflow; + return .{ .rank = after.rank, .interval = interval, .experts = experts, .expert_pools = after.expert_pools }; + } +}; + +test "exl3: telemetry host contracts" { + _ = @import("exl3_telemetry_test.zig"); +} diff --git a/zig/src/families/flashnext/exl3/model/exl3_telemetry_test.zig b/zig/src/families/flashnext/exl3/model/exl3_telemetry_test.zig new file mode 100644 index 0000000000..459bace311 --- /dev/null +++ b/zig/src/families/flashnext/exl3/model/exl3_telemetry_test.zig @@ -0,0 +1,96 @@ +const std = @import("std"); +const telemetry = @import("exl3_telemetry.zig"); +const expectEqual = std.testing.expectEqual; +const expectError = std.testing.expectError; + +test "exl3: checked telemetry aggregation exceeds u32 without overflow or partial mutation" { + var totals: telemetry.Experts = .{ .device_hit_picks = std.math.maxInt(u32) }; + try telemetry.add(telemetry.Experts, &totals, .{ .device_hit_picks = 10 }); + try expectEqual(@as(u64, std.math.maxInt(u32)) + 10, totals.device_hit_picks); + totals.device_hit_picks = std.math.maxInt(u64); + const before = totals; + try expectError(error.Overflow, telemetry.add(telemetry.Experts, &totals, .{ .host_plans = 5, .device_hit_picks = 1 })); + try std.testing.expectEqualDeep(before, totals); + try expectError(error.TelemetryCounterRegression, telemetry.difference(telemetry.Experts, totals, .{})); + try expectEqual(@as(?f64, null), try telemetry.fraction(0, 0)); + try expectEqual(@as(?f64, 0.25), try telemetry.fraction(1, 3)); + try expectError(error.Overflow, telemetry.fraction(std.math.maxInt(u64), 1)); +} + +test "exl3: invalid device telemetry cannot fail normal ownership transfer or publish partial totals" { + const Misc = struct { now: u32 = 0, rounds: u32 = 0, hit_slots: u32 = 0, miss_slots: u32 = 0, copies: u32 = 0, evictions: u32 = 0, overflow: u32 = 0, inplace: u32 = 0 }; + var total: telemetry.Device = .{ .copies = 10 }; + var valid = true; + telemetry.Device.record(&total, &valid, Misc{ .rounds = 1, .hit_slots = 11 }, Misc{}, 10); + try std.testing.expect(!valid); + try expectEqual(@as(u64, 10), total.copies); + telemetry.Device.record(&total, &valid, Misc{ .rounds = 1, .copies = 1 }, Misc{}, 10); + try expectEqual(@as(u64, 10), total.copies); + try expectError(error.TelemetryExpertIntervalTooLong, telemetry.Device.pendingBounded(Misc{ .rounds = 1 << 30 }, Misc{}, 10)); +} + +test "exl3: phase report stable schema rank and model-row denominators" { + const before: telemetry.Snapshot = .{ .rank = .{ .rank = 1, .world = 2 }, .scope = 7 }; + var after = before; + after.experts.device_hit_picks = 40; + after.experts.device_miss_picks = 10; + const report = try telemetry.Report.between(before, after, .{ .phase = .decode, .model_rows = 16, .steps = 8 }); + const text = try report.formatAlloc(std.testing.allocator); + defer std.testing.allocator.free(text); + try std.testing.expect(std.mem.startsWith(u8, text, "EXL3-CACHE v=1 rank=1 world=2 phase=decode rows=16 steps=8 expert_pools=0 cache=expert ")); + try std.testing.expect(std.mem.indexOf(u8, text, "device_hit_picks=40 device_miss_picks=10") != null); + try expectError(error.TelemetryDenominator, telemetry.Report.between(before, after, .{ .phase = .prefill, .model_rows = 0, .steps = 1 })); + after.rank.rank = 0; + try expectError(error.TelemetryOwnerChanged, telemetry.Report.between(before, after, .{ .phase = .decode, .model_rows = 1, .steps = 1 })); + after = before; + after.scope += 1; + try expectError(error.TelemetryOwnerChanged, telemetry.Report.between(before, after, .{ .phase = .decode, .model_rows = 1, .steps = 1 })); +} + +test "exl3: per-layer deltas sum to aggregate counters and overlapping reports claim counters once" { + const rank: telemetry.Rank = .{ .rank = 0, .world = 2 }; + var before_layers = [_]telemetry.LayerSample{.{ .owner = 11, .kind = .decoder, .layer = 4, .lines = 20 }}; + var after_layers = before_layers; + after_layers[0].counts.host_miss_picks = 5; + after_layers[0].counts.host_mandatory_copies = 2; + after_layers[0].counts.device_miss_picks = 3; + after_layers[0].counts.device_copies = 1; + var before: telemetry.Snapshot = .{ .rank = rank, .scope = 42, .layers = &before_layers, .expert_pools = 1 }; + var after = before; + after.layers = &after_layers; + after.experts.host_miss_picks = 5; + after.experts.host_mandatory_copies = 2; + after.experts.device_miss_picks = 3; + after.experts.device_copies = 1; + var cursor: telemetry.LayerCursor(2) = .{}; + const first = try cursor.claim(std.testing.allocator, before, after); + defer std.testing.allocator.free(first); + const report = try telemetry.Report.between(before, after, .{ .phase = .decode, .model_rows = 1, .steps = 1 }); + try std.testing.expectEqual(report.experts.host_miss_picks, first[0].counts.host_miss_picks); + try std.testing.expectEqual(report.experts.host_mandatory_copies, first[0].counts.host_mandatory_copies); + try std.testing.expectEqual(report.experts.device_miss_picks, first[0].counts.device_miss_picks); + try std.testing.expectEqual(report.experts.device_copies, first[0].counts.device_copies); + + var concurrent_before_layers = after_layers; + var concurrent_after_layers = after_layers; + concurrent_before_layers[0].counts.host_miss_picks = 3; + concurrent_before_layers[0].counts.host_mandatory_copies = 1; + concurrent_before_layers[0].counts.device_miss_picks = 2; + concurrent_before_layers[0].counts.device_copies = 0; + concurrent_after_layers[0].counts.host_miss_picks = 8; + concurrent_after_layers[0].counts.host_mandatory_copies = 4; + concurrent_after_layers[0].counts.device_miss_picks = 6; + concurrent_after_layers[0].counts.device_copies = 2; + before.layers = &concurrent_before_layers; + after.layers = &concurrent_after_layers; + before.experts = .{ .host_miss_picks = 5, .host_mandatory_copies = 2, .device_miss_picks = 3, .device_copies = 1 }; + after.experts = .{ .host_miss_picks = 10, .host_mandatory_copies = 5, .device_miss_picks = 7, .device_copies = 3 }; + const second = try cursor.claim(std.testing.allocator, before, after); + defer std.testing.allocator.free(second); + const second_report = try telemetry.Report.between(before, after, .{ .phase = .decode, .model_rows = 1, .steps = 1 }); + try std.testing.expectEqual(second_report.experts.host_miss_picks, second[0].counts.host_miss_picks); + try std.testing.expectEqual(second_report.experts.host_mandatory_copies, second[0].counts.host_mandatory_copies); + try std.testing.expectEqual(second_report.experts.device_miss_picks, second[0].counts.device_miss_picks); + try std.testing.expectEqual(second_report.experts.device_copies, second[0].counts.device_copies); + try std.testing.expectEqual(@as(u64, 10), first[0].counts.host_miss_picks + second[0].counts.host_miss_picks); +} diff --git a/zig/src/families/flashnext/exl3/offload/cuda_exl3_frequency.zig b/zig/src/families/flashnext/exl3/offload/cuda_exl3_frequency.zig new file mode 100644 index 0000000000..4395f177df --- /dev/null +++ b/zig/src/families/flashnext/exl3/offload/cuda_exl3_frequency.zig @@ -0,0 +1,122 @@ +const std = @import("std"); +const Io = std.Io; +const Sha256 = std.crypto.hash.sha2.Sha256; + +pub const Identity = struct { + model: [32]u8, + world: u32, + rank: u32, +}; + +pub const decoder_layers = 48; +pub const experts = 512; +pub const format_version: u32 = 1; + +pub const Counts = struct { + decoder: [decoder_layers][experts]u32 = @splat(@splat(0)), +}; + +const magic = "TFEXL3FQ"; +const header_bytes = 64; +const decoder_bytes = decoder_layers * experts * @sizeOf(u32); +pub const max_file_bytes = header_bytes + decoder_bytes + Sha256.digest_length; + +pub fn read(gpa: std.mem.Allocator, io: Io, path: []const u8, identity: Identity) !Counts { + try validateIdentity(identity); + const file: Io.File = if (comptime @import("builtin").os.tag == .windows) return error.UnsupportedPlatform else .{ + .handle = try std.posix.openat(Io.Dir.cwd().handle, path, .{ + .ACCMODE = .RDONLY, + .NONBLOCK = true, + .CLOEXEC = true, + .NOCTTY = true, + }, 0), + .flags = .{ .nonblocking = true }, + }; + defer file.close(io); + const stat = try file.stat(io); + if (stat.kind != .file) return error.ProfileNotRegular; + if (stat.size > max_file_bytes) return error.ProfileTooLarge; + const buffer = try gpa.alloc(u8, max_file_bytes + 1); + defer gpa.free(buffer); + var reader = file.reader(io, &.{}); + const length = reader.interface.readSliceShort(buffer) catch return reader.err.?; + if (length > max_file_bytes) return error.ProfileTooLarge; + return decode(buffer[0..length], identity); +} + +pub fn write(gpa: std.mem.Allocator, io: Io, path: []const u8, identity: Identity, counts: *const Counts) !void { + try validateIdentity(identity); + if (comptime @import("builtin").os.tag == .windows) return error.UnsupportedPlatform; + const basename = std.fs.path.basename(path); + if (path.len == 0 or std.mem.endsWith(u8, path, "/") or + std.mem.eql(u8, basename, ".") or std.mem.eql(u8, basename, "..")) return error.InvalidPath; + const bytes = try gpa.alloc(u8, max_file_bytes); + defer gpa.free(bytes); + encode(bytes, identity, counts); + var atomic = try Io.Dir.cwd().createFileAtomic(io, path, .{ + .permissions = .fromMode(0o600), + .make_path = false, + .replace = true, + }); + defer atomic.deinit(io); + try atomic.file.writePositionalAll(io, bytes, 0); + try atomic.file.sync(io); + try atomic.replace(io); +} + +fn validateIdentity(identity: Identity) !void { + if (identity.world == 0 or identity.rank >= identity.world) return error.InvalidIdentity; +} + +fn encode(bytes: []u8, identity: Identity, counts: *const Counts) void { + @memcpy(bytes[0..8], magic); + putWord(bytes, 8, format_version); + putWord(bytes, 12, 0); + @memcpy(bytes[16..48], &identity.model); + putWord(bytes, 48, identity.world); + putWord(bytes, 52, identity.rank); + putWord(bytes, 56, decoder_layers); + putWord(bytes, 60, experts); + var offset: usize = header_bytes; + for (&counts.decoder) |*layer| for (layer) |count| { + putWord(bytes, offset, count); + offset += 4; + }; + var checksum: [Sha256.digest_length]u8 = undefined; + Sha256.hash(bytes[0..offset], &checksum, .{}); + @memcpy(bytes[offset..], &checksum); +} + +fn decode(bytes: []const u8, identity: Identity) !Counts { + if (bytes.len < header_bytes + Sha256.digest_length) return error.TruncatedProfile; + if (!std.mem.eql(u8, bytes[0..8], magic)) return error.MalformedProfile; + if (getWord(bytes, 8) != format_version) return error.UnsupportedVersion; + if (getWord(bytes, 12) != 0 or getWord(bytes, 56) != decoder_layers or getWord(bytes, 60) != experts) return error.MalformedProfile; + const payload_end = header_bytes + decoder_bytes; + if (bytes.len < max_file_bytes) return error.TruncatedProfile; + if (bytes.len != max_file_bytes) return error.MalformedProfile; + var checksum: [Sha256.digest_length]u8 = undefined; + Sha256.hash(bytes[0..payload_end], &checksum, .{}); + if (!std.mem.eql(u8, bytes[payload_end..], &checksum)) return error.ChecksumMismatch; + if (!std.mem.eql(u8, bytes[16..48], &identity.model) or + getWord(bytes, 48) != identity.world or getWord(bytes, 52) != identity.rank) return error.IdentityMismatch; + var counts: Counts = .{}; + var offset: usize = header_bytes; + for (&counts.decoder) |*layer| for (layer) |*count| { + count.* = getWord(bytes, offset); + offset += 4; + }; + return counts; +} + +fn putWord(bytes: []u8, offset: usize, value: u32) void { + std.mem.writeInt(u32, bytes[offset..][0..4], value, .little); +} + +fn getWord(bytes: []const u8, offset: usize) u32 { + return std.mem.readInt(u32, bytes[offset..][0..4], .little); +} + +test { + _ = @import("cuda_exl3_frequency_test.zig"); +} diff --git a/zig/src/families/flashnext/exl3/offload/cuda_exl3_frequency_state.zig b/zig/src/families/flashnext/exl3/offload/cuda_exl3_frequency_state.zig new file mode 100644 index 0000000000..18578046d8 --- /dev/null +++ b/zig/src/families/flashnext/exl3/offload/cuda_exl3_frequency_state.zig @@ -0,0 +1,444 @@ +//! Lifetime placement priors: routed/metadata fingerprint; startup warmth may be discarded by prefill. +const std = @import("std"); +const codec = @import("cuda_exl3_frequency.zig"); +const Sha256 = std.crypto.hash.sha2.Sha256; +const namespace = "flashnext-exl3-placement-v1"; + +pub const State = struct { + gpa: std.mem.Allocator, + io: std.Io, + path: []u8, + identity: codec.Identity, + lock: std.Io.File, + + pub fn deinit(state: *State) void { + state.lock.close(state.io); + state.gpa.free(state.path); + state.* = undefined; + } + + pub fn save(state: *State, own: anytype) void { + const counts = snapshot(own) catch |err| { + std.log.err("EXL3 placement profile rank {d}: unsafe snapshot skipped ({s})", .{ state.identity.rank, @errorName(err) }); + return; + }; + codec.write(state.gpa, state.io, state.path, state.identity, &counts) catch |err| { + std.log.warn("EXL3 placement profile rank {d}: save failed ({s})", .{ state.identity.rank, @errorName(err) }); + return; + }; + std.log.info("EXL3 placement profile rank {d}: saved decoder lifetime placement snapshot to {s}", .{ state.identity.rank, state.path }); + } +}; + +pub const Prepared = struct { + state: State, + counts: ?codec.Counts, +}; + +pub fn load(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, own: anytype, world: u32, rank: u32) !?State { + const prepared = try prepare(gpa, io, dir, own.engine.layers[0..own.engine.n_layers], world, rank) orelse return null; + var state = prepared.state; + errdefer state.deinit(); + if (prepared.counts) |*counts| try finish(own, counts); + return state; +} + +pub fn prepare(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, layers: anytype, world: u32, rank: u32) !?Prepared { + const root = std.c.getenv("TF_FLASHNEXT_EXL3_PROFILE_DIR") orelse return null; + if (std.mem.span(root).len == 0) return null; + if (comptime @import("builtin").os.tag == .windows) { + std.log.warn("EXL3 placement profiles require POSIX storage; persistence disabled", .{}); + return null; + } + if (!available(layers)) { + std.log.warn("EXL3 placement profile rank {d}: unavailable decoder pools; persistence disabled", .{rank}); + return null; + } + const before = std.Io.Clock.awake.now(io).toNanoseconds(); + const model = fingerprint(gpa, io, dir, layers) catch |err| { + std.log.warn("EXL3 placement profile rank {d}: fingerprint unavailable ({s}); persistence disabled", .{ rank, @errorName(err) }); + return null; + }; + std.log.info("EXL3 placement profile rank {d}: routed-bank/metadata fingerprint took {d:.3} s (excludes router/dense weights)", .{ rank, @as(f64, @floatFromInt(std.Io.Clock.awake.now(io).toNanoseconds() - before)) / 1e9 }); + var state = open(gpa, io, std.mem.span(root), .{ .model = model, .world = world, .rank = rank }) catch |err| { + std.log.warn("EXL3 placement profile rank {d}: lock/storage unavailable ({s}); persistence disabled", .{ rank, @errorName(err) }); + return null; + }; + errdefer state.deinit(); + var counts: ?codec.Counts = codec.read(gpa, io, state.path, state.identity) catch |err| blk: { + std.log.warn("EXL3 placement profile rank {d}: cold start ({s})", .{ rank, @errorName(err) }); + break :blk null; + }; + if (counts) |*value| { + decay(value); + std.log.info("EXL3 placement profile rank {d}: loaded and decayed rank-local placement counts before pool allocation", .{rank}); + } + return .{ .state = state, .counts = counts }; +} + +pub fn finish(own: anytype, counts: *const codec.Counts) !void { + const layers = own.engine.layers[0..own.engine.n_layers]; + if (!available(layers)) return error.ProfileUnavailable; + for (layers, 0..) |layer, index| { + const pool = layer.exl3.?.pool orelse return error.ProfileUnavailable; + pool.counts[pool.opt.layer] = counts.decoder[index]; + } + var warmed: usize = 0; + for (layers) |layer| warmed += try layer.exl3.?.pool.?.warmProfile(); + std.log.info("EXL3 placement profile: imported and warmed {d} expert lines after pool allocation", .{warmed}); +} + +fn open(gpa: std.mem.Allocator, io: std.Io, root: []const u8, identity: codec.Identity) !State { + const private = try std.fs.path.join(gpa, &.{ root, namespace }); + defer gpa.free(private); + var directory = try std.Io.Dir.cwd().createDirPathOpen(io, private, .{ + .permissions = .fromMode(0o700), + .open_options = .{ .iterate = true }, + }); + defer directory.close(io); + try directory.setPermissions(io, .fromMode(0o700)); + const name = try std.fmt.allocPrint(gpa, "{s}-w{d}-r{d}.bin", .{ std.fmt.bytesToHex(&identity.model, .lower), identity.world, identity.rank }); + defer gpa.free(name); + const path = try std.fs.path.join(gpa, &.{ private, name }); + errdefer gpa.free(path); + const lock_path = try std.fmt.allocPrint(gpa, "{s}.lock", .{path}); + defer gpa.free(lock_path); + const lock = try openLock(io, lock_path); + errdefer lock.close(io); + try lock.setPermissions(io, .fromMode(0o600)); + return .{ .gpa = gpa, .io = io, .path = path, .identity = identity, .lock = lock }; +} + +fn openLock(io: std.Io, path: []const u8) !std.Io.File { + const file: std.Io.File = if (comptime @import("builtin").os.tag == .windows) return error.UnsupportedPlatform else .{ + .handle = try std.posix.openat(std.Io.Dir.cwd().handle, path, .{ + .ACCMODE = .RDWR, + .CREAT = true, + .NONBLOCK = true, + .CLOEXEC = true, + .NOCTTY = true, + .NOFOLLOW = true, + }, 0o600), + .flags = .{ .nonblocking = true }, + }; + errdefer file.close(io); + if ((try file.stat(io)).kind != .file) return error.ProfileLockNotRegular; + if (!try file.tryLock(io, .exclusive)) return error.WouldBlock; + return file; +} + +fn restore(state: *State, own: anytype) !bool { + var counts = codec.read(state.gpa, state.io, state.path, state.identity) catch |err| { + std.log.warn("EXL3 placement profile rank {d}: cold start ({s})", .{ state.identity.rank, @errorName(err) }); + return false; + }; + decay(&counts); + const layers = own.engine.layers[0..own.engine.n_layers]; + for (layers, 0..) |layer, index| { + const pool = layer.exl3.?.pool.?; + pool.counts[pool.opt.layer] = counts.decoder[index]; + } + std.log.info("EXL3 placement profile rank {d}: imported decoder lifetime counts, decayed once by 2 (nonzero minimum 1)", .{state.identity.rank}); + var warmed: usize = 0; + for (layers) |layer| warmed += try layer.exl3.?.pool.?.warmProfile(); + std.log.info("EXL3 placement profile rank {d}: warmed {d} expert lines", .{ state.identity.rank, warmed }); + return true; +} + +fn available(layers: anytype) bool { + if (layers.len != 48) return false; + for (layers) |layer| { + const experts = layer.exl3 orelse return false; + if (experts.bank == null or experts.routed.len != codec.experts) return false; + if (experts.pool) |pool| if (pool.opt.layer >= 48) return false; + } + return true; +} + +fn decay(counts: *codec.Counts) void { + for (&counts.decoder) |*layer| decayRow(layer); +} + +fn decayRow(row: *[512]u32) void { + for (row) |*count| { + if (count.* != 0) count.* = @max(1, count.* / 2); + } +} + +fn snapshot(own: anytype) !codec.Counts { + const layers = own.engine.layers[0..own.engine.n_layers]; + if (!available(layers)) return error.ProfileUnavailable; + try own.stream.synchronize(); + var counts: codec.Counts = .{}; + for (layers, 0..) |layer, index| { + const pool = layer.exl3.?.pool.?; + try pool.lines.drain(); + counts.decoder[index] = pool.counts[pool.opt.layer]; + } + return counts; +} + +fn hashBytes(hash: *Sha256, bytes: []const u8) void { + var length: [8]u8 = undefined; + std.mem.writeInt(u64, &length, bytes.len, .little); + hash.update(&length); + hash.update(bytes); +} + +fn hashRouted(hash: *Sha256, routed: []const [3][]const u8) void { + for (routed) |expert| for (expert) |projection| hashBytes(hash, projection); +} + +fn fingerprint(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, layers: anytype) ![32]u8 { + var hash: Sha256 = .init(.{}); + hashBytes(&hash, namespace); + for ([_][]const u8{ "config.json", "model.safetensors.index.json" }) |name| { + const path = try std.fs.path.join(gpa, &.{ dir, name }); + defer gpa.free(path); + const bytes = try std.Io.Dir.cwd().readFileAlloc(io, path, gpa, .limited(1 << 26)); + defer gpa.free(bytes); + hashBytes(&hash, name); + hashBytes(&hash, bytes); + } + hashBytes(&hash, "decoder-routed-trellis"); + for (layers) |layer| hashRouted(&hash, expertRouted(layer.exl3.?)); + return hash.finalResult(); +} + +fn expertRouted(experts: anytype) []const [3][]const u8 { + if (experts.pool) |pool| return pool.src.routed; + return experts.routed; +} + +const MockPool = struct { + counts: *[48][512]u32, + opt: struct { layer: u8, lines: usize = 10 }, + src: struct { routed: []const [3][]const u8 = &.{.{ "gate", "up", "down" }} } = .{}, + lines: struct { + drained: bool = false, + pub fn drain(copies: *@This()) !void { + copies.drained = true; + } + } = .{}, + warm_calls: usize = 0, + fail_warm: bool = false, + + pub fn warmProfile(pool: *MockPool) !usize { + pool.warm_calls += 1; + if (pool.fail_warm) return error.MockCudaWarmFailure; + return 1; + } +}; + +const mock_routed: [512][3][]const u8 = @splat(.{ "gate", "up", "down" }); +const MockExperts = struct { pool: ?*MockPool, bank: ?[]const u8 = "bank", routed: []const [3][]const u8 = &mock_routed }; +const MockLayer = struct { exl3: ?*MockExperts }; +const Mock = struct { + decoder: [48][512]u32 = @splat(@splat(0)), + pools: [48]MockPool = undefined, + experts: [48]MockExperts = undefined, + engine: struct { layers: [48]MockLayer, n_layers: usize = 48 } = undefined, + + fn init(mock: *Mock) void { + mock.* = .{}; + mock.engine.n_layers = 48; + for (&mock.pools, &mock.experts, 0..) |*pool, *experts, index| { + pool.* = .{ .counts = &mock.decoder, .opt = .{ .layer = @intCast(index) }, .src = .{ .routed = &mock_routed } }; + experts.* = .{ .pool = pool }; + mock.engine.layers[index] = .{ .exl3 = experts }; + } + } + + const Stream = struct { + syncs: usize = 0, + fail: bool = false, + pub fn synchronize(stream: *@This()) !void { + if (stream.fail) return error.MockCudaSyncFailure; + stream.syncs += 1; + } + }; + fn owner(mock: *Mock) struct { engine: *@TypeOf(mock.engine), stream: Stream } { + return .{ .engine = &mock.engine, .stream = .{} }; + } +}; + +const test_identity: codec.Identity = .{ .model = @splat(42), .world = 2, .rank = 0 }; + +test "exl3: placement state lifetime lock conflicts, rank separation and private storage" { + if (comptime @import("builtin").os.tag == .windows) return error.SkipZigTest; + const allocator = std.testing.allocator; + const io = std.testing.io; + var temporary = std.testing.tmpDir(.{}); + defer temporary.cleanup(); + const root = try temporary.dir.realPathFileAlloc(io, ".", allocator); + defer allocator.free(root); + var state = try open(allocator, io, root, test_identity); + var held = true; + defer if (held) state.deinit(); + try std.testing.expectError(error.WouldBlock, open(allocator, io, root, test_identity)); + const empty: codec.Counts = .{}; + try codec.write(allocator, io, state.path, state.identity, &empty); + try std.testing.expectError(error.WouldBlock, open(allocator, io, root, test_identity)); + var rank_identity = test_identity; + rank_identity.rank = 1; + var other_rank = try open(allocator, io, root, rank_identity); + defer other_rank.deinit(); + try std.testing.expect(!std.mem.eql(u8, state.path, other_rank.path)); + const lock_stat = try state.lock.stat(io); + try std.testing.expectEqual(@as(u32, 0o600), lock_stat.permissions.toMode() & 0o777); + const directory = try std.Io.Dir.cwd().openDir(io, std.fs.path.dirname(state.path).?, .{}); + defer directory.close(io); + const directory_stat = try directory.stat(io); + try std.testing.expectEqual(@as(u32, 0o700), directory_stat.permissions.toMode() & 0o777); + state.deinit(); + held = false; + var reopened = try open(allocator, io, root, test_identity); + defer reopened.deinit(); +} + +test "exl3: placement state cold IO never imports or warms; CUDA warm errors propagate" { + if (comptime @import("builtin").os.tag == .windows) return error.SkipZigTest; + const allocator = std.testing.allocator; + const io = std.testing.io; + var temporary = std.testing.tmpDir(.{}); + defer temporary.cleanup(); + const root = try temporary.dir.realPathFileAlloc(io, ".", allocator); + defer allocator.free(root); + var state = try open(allocator, io, root, test_identity); + defer state.deinit(); + var mock: Mock = undefined; + mock.init(); + var own = mock.owner(); + try std.testing.expect(!try restore(&state, &own)); + try std.Io.Dir.cwd().writeFile(io, .{ .sub_path = state.path, .data = "corrupt" }); + try std.testing.expect(!try restore(&state, &own)); + var counts: codec.Counts = .{}; + counts.decoder[0][1] = 123; + try codec.write(allocator, io, state.path, state.identity, &counts); + mock.pools[0].fail_warm = true; + try std.testing.expectError(error.MockCudaWarmFailure, restore(&state, &own)); + try std.testing.expectEqual(@as(u32, 61), mock.decoder[0][1]); + mock.engine.n_layers = 47; + try std.testing.expect(!available(mock.engine.layers[0..mock.engine.n_layers])); +} + +test "exl3: placement state decay and replacement snapshot do not add imported prior twice" { + if (comptime @import("builtin").os.tag == .windows) return error.SkipZigTest; + const allocator = std.testing.allocator; + const io = std.testing.io; + var temporary = std.testing.tmpDir(.{}); + defer temporary.cleanup(); + const root = try temporary.dir.realPathFileAlloc(io, ".", allocator); + defer allocator.free(root); + var state = try open(allocator, io, root, test_identity); + defer state.deinit(); + var counts: codec.Counts = .{}; + counts.decoder[0][0..5].* = .{ 0, 1, 2, 3, std.math.maxInt(u32) }; + try codec.write(allocator, io, state.path, state.identity, &counts); + var mock: Mock = undefined; + mock.init(); + var own = mock.owner(); + try std.testing.expect(try restore(&state, &own)); + try std.testing.expectEqualSlices(u32, &.{ 0, 1, 1, 1, 2147483647 }, mock.decoder[0][0..5]); + for (mock.pools) |pool| try std.testing.expectEqual(@as(usize, 1), pool.warm_calls); + mock.decoder[0][1] += 8; + const saved = try snapshot(&own); + try codec.write(allocator, io, state.path, state.identity, &saved); + const disk = try codec.read(allocator, io, state.path, state.identity); + try std.testing.expectEqual(@as(u32, 9), disk.decoder[0][1]); + mock.init(); + own = mock.owner(); + try std.testing.expect(try restore(&state, &own)); + try std.testing.expectEqual(@as(u32, 4), mock.decoder[0][1]); + const missing = try std.fs.path.join(allocator, &.{ root, "missing", "profile" }); + defer allocator.free(missing); + try std.testing.expectError(error.FileNotFound, codec.write(allocator, io, missing, state.identity, &saved)); +} + +test "exl3: placement state snapshot synchronizes once and drains decoder counts" { + if (comptime @import("builtin").os.tag == .windows) return error.SkipZigTest; + var mock: Mock = undefined; + mock.init(); + var own = mock.owner(); + for (&mock.pools, 0..) |*pool, index| pool.counts[pool.opt.layer][7] = @intCast(100 + index); + const counts = try snapshot(&own); + try std.testing.expectEqual(@as(usize, 1), own.stream.syncs); + try std.testing.expectEqual(@as(u32, 100), counts.decoder[0][7]); + try std.testing.expectEqual(@as(u32, 147), counts.decoder[47][7]); + for (mock.pools) |pool| try std.testing.expect(pool.lines.drained); + own.stream.fail = true; + try std.testing.expectError(error.MockCudaSyncFailure, snapshot(&own)); +} + +test "exl3: prepared placement counts import and warm only after pools attach" { + var mock: Mock = undefined; + mock.init(); + var own = mock.owner(); + var counts: codec.Counts = .{}; + for (&counts.decoder, 0..) |*row, layer| row[7] = @intCast(layer + 1); + + try finish(&own, &counts); + + for (mock.pools[0..48], 0..) |pool, layer| { + try std.testing.expectEqual(@as(u32, @intCast(layer + 1)), mock.decoder[layer][7]); + try std.testing.expectEqual(@as(usize, 1), pool.warm_calls); + } +} + +test "exl3: placement fingerprint reads routed and metadata bytes and ignores pool size" { + if (comptime @import("builtin").os.tag == .windows) return error.SkipZigTest; + const allocator = std.testing.allocator; + const io = std.testing.io; + var temporary = std.testing.tmpDir(.{}); + defer temporary.cleanup(); + const root = try temporary.dir.realPathFileAlloc(io, ".", allocator); + defer allocator.free(root); + try temporary.dir.writeFile(io, .{ .sub_path = "config.json", .data = "config" }); + try temporary.dir.writeFile(io, .{ .sub_path = "model.safetensors.index.json", .data = "index" }); + var mock: Mock = undefined; + mock.init(); + const original = try fingerprint(allocator, io, root, &mock.engine.layers); + for (&mock.experts) |*experts| experts.pool = null; + try std.testing.expectEqual(original, try fingerprint(allocator, io, root, &mock.engine.layers)); + for (&mock.experts, &mock.pools) |*experts, *pool| experts.pool = pool; + mock.pools[0].opt.lines = 23; + try std.testing.expectEqual(original, try fingerprint(allocator, io, root, &mock.engine.layers)); + mock.pools[47].src.routed = &.{.{ "changed gate", "up", "down" }}; + try std.testing.expect(!std.mem.eql(u8, &original, &(try fingerprint(allocator, io, root, &mock.engine.layers)))); + mock.pools[47].src.routed = &.{.{ "gate", "up", "down" }}; + try temporary.dir.writeFile(io, .{ .sub_path = "config.json", .data = "changed config" }); + try std.testing.expect(!std.mem.eql(u8, &original, &(try fingerprint(allocator, io, root, &mock.engine.layers)))); + try temporary.dir.writeFile(io, .{ .sub_path = "config.json", .data = "config" }); + try temporary.dir.writeFile(io, .{ .sub_path = "model.safetensors.index.json", .data = "changed index" }); + try std.testing.expect(!std.mem.eql(u8, &original, &(try fingerprint(allocator, io, root, &mock.engine.layers)))); + try temporary.dir.writeFile(io, .{ .sub_path = "model.safetensors.index.json", .data = "index" }); +} + +test "exl3: placement lock rejects nonregular storage without waiting for FIFO readers" { + if (comptime @import("builtin").os.tag == .windows) return error.SkipZigTest; + const allocator = std.testing.allocator; + const io = std.testing.io; + var temporary = std.testing.tmpDir(.{}); + defer temporary.cleanup(); + const path = try temporary.dir.realPathFileAlloc(io, ".", allocator); + defer allocator.free(path); + try std.testing.expectError(error.ProfileLockNotRegular, openLock(io, "/dev/null")); + const fifo = try std.fs.path.join(allocator, &.{ path, "profile.lock" }); + defer allocator.free(fifo); + const terminated = try allocator.dupeSentinel(u8, fifo, 0); + defer allocator.free(terminated); + switch (@import("builtin").os.tag) { + .linux => { + const result = std.os.linux.mknodat(temporary.dir.handle, "profile.lock", std.posix.S.IFIFO | 0o600, 0); + if (std.posix.errno(result) != .SUCCESS) return error.FifoCreateFailed; + }, + .macos => { + const Native = struct { + extern "c" fn mkfifo(pathname: [*:0]const u8, mode: std.posix.mode_t) c_int; + }; + if (Native.mkfifo(terminated, 0o600) != 0) return error.FifoCreateFailed; + }, + else => return error.SkipZigTest, + } + try std.testing.expectError(error.ProfileLockNotRegular, openLock(io, fifo)); +} diff --git a/zig/src/families/flashnext/exl3/offload/cuda_exl3_frequency_test.zig b/zig/src/families/flashnext/exl3/offload/cuda_exl3_frequency_test.zig new file mode 100644 index 0000000000..f39617e4a8 --- /dev/null +++ b/zig/src/families/flashnext/exl3/offload/cuda_exl3_frequency_test.zig @@ -0,0 +1,299 @@ +const std = @import("std"); +const profile = @import("cuda_exl3_frequency.zig"); +const Io = std.Io; +const allocator = std.testing.allocator; +const io = std.testing.io; +const identity: profile.Identity = .{ .model = @splat(0xa7), .world = 2, .rank = 1 }; +const header_bytes = 64; +const checksum_bytes = 32; + +fn profilePath(directory: Io.Dir) ![]u8 { + const root = try directory.realPathFileAlloc(io, ".", allocator); + defer allocator.free(root); + return std.fs.path.join(allocator, &.{ root, "frequency.bin" }); +} + +fn sample(counts: *profile.Counts) void { + counts.* = .{}; + for (&counts.decoder, 0..) |*layer, layer_index| for (layer, 0..) |*count, expert_index| { + count.* = @intCast((layer_index + 1) * 100_000 + expert_index); + }; + counts.decoder[0][0] = 0; + counts.decoder[47][511] = std.math.maxInt(u32); +} + +fn expectCounts(expected: *const profile.Counts, actual: *const profile.Counts) !void { + for (&expected.decoder, &actual.decoder) |*expected_layer, *actual_layer| { + try std.testing.expectEqualSlices(u32, expected_layer, actual_layer); + } +} + +fn expectOnlyProfile(directory: Io.Dir) !void { + var iterator = directory.iterate(); + var entries: usize = 0; + while (try iterator.next(io)) |entry| { + entries += 1; + try std.testing.expectEqualStrings("frequency.bin", entry.name); + } + try std.testing.expectEqual(@as(usize, 1), entries); +} + +fn bytesAt(path: []const u8) ![]u8 { + return Io.Dir.cwd().readFileAlloc(io, path, allocator, .limited(profile.max_file_bytes + 2)); +} + +fn expectPreserved(path: []const u8, before: []const u8) !void { + const after = try bytesAt(path); + defer allocator.free(after); + try std.testing.expectEqualSlices(u8, before, after); +} + +test "exl3: frequency full decoder round trip with private permissions" { + if (comptime @import("builtin").os.tag == .windows) return error.SkipZigTest; + var temporary = std.testing.tmpDir(.{ .iterate = true }); + defer temporary.cleanup(); + const path = try profilePath(temporary.dir); + defer allocator.free(path); + var counts: profile.Counts = .{}; + sample(&counts); + try profile.write(allocator, io, path, identity, &counts); + const restored = try profile.read(allocator, io, path, identity); + try expectCounts(&counts, &restored); + const stat = try temporary.dir.statFile(io, "frequency.bin", .{}); + try std.testing.expectEqual(@as(u64, profile.max_file_bytes), stat.size); + try std.testing.expectEqual(@as(std.posix.mode_t, 0o600), stat.permissions.toMode() & 0o777); + const bytes = try bytesAt(path); + defer allocator.free(bytes); + try std.testing.expectEqualSlices(u8, "TFEXL3FQ", bytes[0..8]); + try std.testing.expectEqualSlices(u8, &identity.model, bytes[16..48]); + try std.testing.expectEqualSlices(u8, &.{ 1, 0, 0, 0 }, bytes[8..12]); + try std.testing.expectEqualSlices(u8, &.{ 0, 0, 0, 0 }, bytes[64..68]); + try std.testing.expectEqualSlices(u8, &.{ 0xa1, 0x86, 0x01, 0x00 }, bytes[68..72]); + try std.testing.expectEqualSlices(u8, &.{ 0xff, 0xff, 0xff, 0xff }, bytes[64 + 48 * 512 * 4 - 4 ..][0..4]); + try expectOnlyProfile(temporary.dir); +} + +test "exl3: frequency missing profile and every exact identity mismatch are recoverable" { + if (comptime @import("builtin").os.tag == .windows) return error.SkipZigTest; + var temporary = std.testing.tmpDir(.{}); + defer temporary.cleanup(); + const path = try profilePath(temporary.dir); + defer allocator.free(path); + try std.testing.expectError(error.FileNotFound, profile.read(allocator, io, path, identity)); + const counts: profile.Counts = .{}; + try profile.write(allocator, io, path, identity, &counts); + for (0..32) |model_index| { + var different = identity; + different.model[model_index] ^= 1; + try std.testing.expectError(error.IdentityMismatch, profile.read(allocator, io, path, different)); + } + var different = identity; + different.world = 3; + try std.testing.expectError(error.IdentityMismatch, profile.read(allocator, io, path, different)); + different = identity; + different.rank = 0; + try std.testing.expectError(error.IdentityMismatch, profile.read(allocator, io, path, different)); + const restored = try profile.read(allocator, io, path, identity); + try expectCounts(&counts, &restored); +} + +test "exl3: frequency rejects truncated oversized malformed and corrupt profiles" { + if (comptime @import("builtin").os.tag == .windows) return error.SkipZigTest; + var temporary = std.testing.tmpDir(.{}); + defer temporary.cleanup(); + const path = try profilePath(temporary.dir); + defer allocator.free(path); + var counts: profile.Counts = .{}; + sample(&counts); + try profile.write(allocator, io, path, identity, &counts); + const original = try bytesAt(path); + defer allocator.free(original); + for ([_]usize{ 0, 7, 63, 64, 95, 96, profile.max_file_bytes - 1 }) |length| { + try temporary.dir.writeFile(io, .{ .sub_path = "frequency.bin", .data = original[0..length] }); + try std.testing.expectError(error.TruncatedProfile, profile.read(allocator, io, path, identity)); + } + const changed = try allocator.dupe(u8, original); + defer allocator.free(changed); + const Mutation = struct { offset: usize, value: u8, expected: anyerror }; + for ([_]Mutation{ + .{ .offset = 0, .value = 0, .expected = error.MalformedProfile }, + .{ .offset = 8, .value = 2, .expected = error.UnsupportedVersion }, + .{ .offset = 12, .value = 2, .expected = error.MalformedProfile }, + .{ .offset = 56, .value = 47, .expected = error.MalformedProfile }, + .{ .offset = 60, .value = 1, .expected = error.MalformedProfile }, + .{ .offset = 16, .value = 0, .expected = error.ChecksumMismatch }, + .{ .offset = 48, .value = 3, .expected = error.ChecksumMismatch }, + .{ .offset = 64, .value = 1, .expected = error.ChecksumMismatch }, + .{ .offset = 64 + 48 * 512 * 4, .value = 1, .expected = error.ChecksumMismatch }, + .{ .offset = profile.max_file_bytes - 1, .value = original[profile.max_file_bytes - 1] ^ 1, .expected = error.ChecksumMismatch }, + }) |mutation| { + @memcpy(changed, original); + changed[mutation.offset] = mutation.value; + try temporary.dir.writeFile(io, .{ .sub_path = "frequency.bin", .data = changed }); + try std.testing.expectError(mutation.expected, profile.read(allocator, io, path, identity)); + } + const oversized = try allocator.alloc(u8, profile.max_file_bytes + 1); + defer allocator.free(oversized); + @memcpy(oversized[0..original.len], original); + oversized[original.len] = 0; + try temporary.dir.writeFile(io, .{ .sub_path = "frequency.bin", .data = oversized }); + try std.testing.expectError(error.ProfileTooLarge, profile.read(allocator, io, path, identity)); +} + +fn failWrite(userdata: ?*anyopaque, file: Io.File, header: []const u8, data: []const []const u8, splat: usize, offset: u64) Io.File.WritePositionalError!usize { + if (offset != 0) return error.NoSpaceLeft; + const prefix = data[0][0..13]; + return io.vtable.fileWritePositional(userdata, file, header, &.{prefix}, splat, offset); +} + +fn failSync(_: ?*anyopaque, file: Io.File) Io.File.SyncError!void { + const stat = file.stat(io) catch return error.InputOutput; + if (stat.permissions.toMode() & 0o777 != 0o600 or stat.size != profile.max_file_bytes) return error.AccessDenied; + return error.InputOutput; +} + +fn failRename(_: ?*anyopaque, _: Io.Dir, _: []const u8, _: Io.Dir, _: []const u8) Io.Dir.RenameError!void { + return error.AccessDenied; +} + +test "exl3: frequency partial write sync and rename failures preserve previous bytes and remove temporaries" { + if (comptime @import("builtin").os.tag == .windows) return error.SkipZigTest; + var temporary = std.testing.tmpDir(.{ .iterate = true }); + defer temporary.cleanup(); + const path = try profilePath(temporary.dir); + defer allocator.free(path); + var counts: profile.Counts = .{}; + try profile.write(allocator, io, path, identity, &counts); + const before = try bytesAt(path); + defer allocator.free(before); + sample(&counts); + for (0..3) |failure_index| { + var failing_vtable = io.vtable.*; + const expected: anyerror = switch (failure_index) { + 0 => failure: { + failing_vtable.fileWritePositional = failWrite; + break :failure error.NoSpaceLeft; + }, + 1 => failure: { + failing_vtable.fileSync = failSync; + break :failure error.InputOutput; + }, + else => failure: { + failing_vtable.dirRename = failRename; + break :failure error.AccessDenied; + }, + }; + var failing_io = io; + failing_io.vtable = &failing_vtable; + try std.testing.expectError(expected, profile.write(allocator, failing_io, path, identity, &counts)); + try expectPreserved(path, before); + try expectOnlyProfile(temporary.dir); + const restored = try profile.read(allocator, io, path, identity); + try std.testing.expectEqual(@as(u32, 0), restored.decoder[47][511]); + } + try profile.write(allocator, io, path, identity, &counts); + const restored = try profile.read(allocator, io, path, identity); + try expectCounts(&counts, &restored); +} + +test "exl3: frequency allocation and invalid identity failures preserve previous profile" { + if (comptime @import("builtin").os.tag == .windows) return error.SkipZigTest; + var temporary = std.testing.tmpDir(.{ .iterate = true }); + defer temporary.cleanup(); + const path = try profilePath(temporary.dir); + defer allocator.free(path); + const counts: profile.Counts = .{}; + try profile.write(allocator, io, path, identity, &counts); + const before = try bytesAt(path); + defer allocator.free(before); + var failing = std.testing.FailingAllocator.init(allocator, .{ .fail_index = 0 }); + try std.testing.expectError(error.OutOfMemory, profile.write(failing.allocator(), io, path, identity, &counts)); + try std.testing.expectError(error.OutOfMemory, profile.read(failing.allocator(), io, path, identity)); + for ([_]profile.Identity{ + .{ .model = identity.model, .world = 0, .rank = 0 }, + .{ .model = identity.model, .world = 2, .rank = 2 }, + }) |invalid| { + try std.testing.expectError(error.InvalidIdentity, profile.write(allocator, io, path, invalid, &counts)); + try std.testing.expectError(error.InvalidIdentity, profile.read(allocator, io, path, invalid)); + } + try std.testing.expectError(error.InvalidPath, profile.write(allocator, io, "", identity, &counts)); + try std.testing.expectError(error.InvalidPath, profile.write(allocator, io, ".", identity, &counts)); + try std.testing.expectError(error.InvalidPath, profile.write(allocator, io, "..", identity, &counts)); + try expectPreserved(path, before); + try expectOnlyProfile(temporary.dir); +} + +test "exl3: frequency read allocation stays fixed and oversized input needs no allocation" { + if (comptime @import("builtin").os.tag == .windows) return error.SkipZigTest; + var temporary = std.testing.tmpDir(.{}); + defer temporary.cleanup(); + const path = try profilePath(temporary.dir); + defer allocator.free(path); + const single: profile.Identity = .{ .model = @splat(0), .world = 1, .rank = 0 }; + var counts: profile.Counts = .{}; + sample(&counts); + try profile.write(allocator, io, path, single, &counts); + const backing = try allocator.alloc(u8, profile.max_file_bytes + 1); + defer allocator.free(backing); + var bounded = std.heap.FixedBufferAllocator.init(backing); + const restored = try profile.read(bounded.allocator(), io, path, single); + try expectCounts(&counts, &restored); + const file = try temporary.dir.openFile(io, "frequency.bin", .{ .mode = .read_write }); + defer file.close(io); + try file.setLength(io, profile.max_file_bytes + 1); + var failing = std.testing.FailingAllocator.init(allocator, .{ .fail_index = 0 }); + try std.testing.expectError(error.ProfileTooLarge, profile.read(failing.allocator(), io, path, single)); + try std.testing.expect(!failing.has_induced_failure); +} + +test "exl3: frequency failed first save leaves no destination or temporary file" { + if (comptime @import("builtin").os.tag == .windows) return error.SkipZigTest; + var temporary = std.testing.tmpDir(.{ .iterate = true }); + defer temporary.cleanup(); + const path = try profilePath(temporary.dir); + defer allocator.free(path); + var counts: profile.Counts = .{}; + sample(&counts); + var failing_vtable = io.vtable.*; + failing_vtable.fileSync = failSync; + var failing_io = io; + failing_io.vtable = &failing_vtable; + try std.testing.expectError(error.InputOutput, profile.write(allocator, failing_io, path, identity, &counts)); + try std.testing.expectError(error.FileNotFound, profile.read(allocator, io, path, identity)); + var iterator = temporary.dir.iterate(); + try std.testing.expectEqual(@as(?Io.Dir.Entry, null), try iterator.next(io)); + try profile.write(allocator, io, path, identity, &counts); + const restored = try profile.read(allocator, io, path, identity); + try expectCounts(&counts, &restored); + try expectOnlyProfile(temporary.dir); +} + +test "exl3: frequency rejects a FIFO without a writer and other nonregular input" { + if (comptime @import("builtin").os.tag == .windows) return error.SkipZigTest; + var temporary = std.testing.tmpDir(.{}); + defer temporary.cleanup(); + const path = try profilePath(temporary.dir); + defer allocator.free(path); + const terminated_path = try allocator.dupeSentinel(u8, path, 0); + defer allocator.free(terminated_path); + switch (@import("builtin").os.tag) { + .linux => { + const result = std.os.linux.mknodat(temporary.dir.handle, "frequency.bin", std.posix.S.IFIFO | 0o600, 0); + if (std.posix.errno(result) != .SUCCESS) return error.FifoCreateFailed; + }, + .macos => { + const Native = struct { + extern "c" fn mkfifo(pathname: [*:0]const u8, mode: std.posix.mode_t) c_int; + }; + if (Native.mkfifo(terminated_path, 0o600) != 0) return error.FifoCreateFailed; + }, + else => return error.SkipZigTest, + } + var failing = std.testing.FailingAllocator.init(allocator, .{ .fail_index = 0 }); + try std.testing.expectError(error.ProfileNotRegular, profile.read(failing.allocator(), io, path, identity)); + try std.testing.expectError(error.ProfileNotRegular, profile.read(allocator, io, "/dev/null", identity)); + const root = try temporary.dir.realPathFileAlloc(io, ".", allocator); + defer allocator.free(root); + try std.testing.expectError(error.ProfileNotRegular, profile.read(allocator, io, root, identity)); + try std.testing.expect(!failing.has_induced_failure); +} diff --git a/zig/src/families/flashnext/exl3/offload/cuda_exl3_pool.zig b/zig/src/families/flashnext/exl3/offload/cuda_exl3_pool.zig new file mode 100644 index 0000000000..da076ee7e6 --- /dev/null +++ b/zig/src/families/flashnext/exl3/offload/cuda_exl3_pool.zig @@ -0,0 +1,345 @@ +//! One layer's EXL3 routed-expert offload pool on upstream's `Experts` (plan step 6 of + +const std = @import("std"); +const cuda = @import("cuda"); +const round_ = @import("../model/exl3_round.zig"); +const exl3e = cuda.exl3_experts; +const Lines = cuda.exl3_pool.Lines; +const telemetry = @import("../model/exl3_telemetry.zig"); + +pub const offload = @import("exl3_offload.zig"); +pub const round = round_; +pub const Round = round_.Round; +pub const linesFor = round_.linesFor; +pub const shape = offload.shape; +pub const shared_id = offload.shared_id; +/// The fewest lines a layer may have: one row's ten routed picks always fit one pass. +pub const min_lines: usize = 10; + +test { + _ = @import("exl3_frequency_seed.zig"); +} + +/// The layer's experts in host memory: routed [512] gate/up/down trellis bytes (stored layout), the shared expert's. +pub const Source = struct { + routed: []const [3][]const u8, + shared: ?[3][]const u8 = null, + pinned_bank: ?[]const u8 = null, +}; + +/// Device addresses of the resident, stacked fp16 scales ([E, D] or [E, I], E = 513). +pub const Scales = struct { suh_g: u64, suh_u: u64, svh_g: u64, svh_u: u64, suh_d: u64, svh_d: u64 }; + +pub const Options = struct { + layer: u8, + lines: usize, + dims: usize = 2560, + width: usize = 640, + /// Rows and slots a round may hold (sizes the masked-pick buffers). + max_rows: usize = 16, + slots: usize = 11, + staging: usize = 8, + /// Copy from the source bytes directly (they are pinned already); otherwise through pinned staging lines. + direct: bool = false, + codebook: cuda.exl3.Codebook = .mul1, + act_mode: i32 = 1, + swap_every: u16 = 4, + swap_cap: u16 = 96, + /// Bytes a swap round may prefetch (the copy budget); 0: no prefetch. + prefetch_bytes: u64 = 0, + /// Double buffering: a round that fits without them does not evict the previous round's lines. + double_buffer: bool = true, + /// Masked-pick buffers for this many passes of a split round (0: the round's own bound, uniq + 1). + max_passes: usize = 0, +}; + +/// Passes a round of a `lines`-line pool can split into when nothing prefetches (`prefetch_bytes` 0): pass 0 and one +pub fn passesFor(lines: usize) usize { + return 2 + (shape.experts + lines - 1) / lines; +} + +pub const Stats = struct { + rounds: u64 = 0, + passes: u64 = 0, + hit_slots: u64 = 0, + miss_slots: u64 = 0, + unique: u64 = 0, + misses: u64 = 0, // mandatory copies + prefetches: u64 = 0, + evictions: u64 = 0, + busy: u64 = 0, // plans that had to drain in-flight copies first + protected: u64 = 0, // rounds that spared the previous round's lines +}; + +pub const Pool = struct { + gpa: std.mem.Allocator, + io: std.Io, + opt: Options, + src: Source, + scales: Scales, + lines: Lines, + k2: [][3]i32, // per routed expert, from its trellis sizes + slots_tab: []u16, + line_tab: []offload.Line, + stamp_tab: []u32, + layer_tab: []u8, + v: offload.Vram, + counts: *[shape.layers][shape.experts]u32, + own_counts: bool, + round: Round, + now: u32 = 0, + round_no: u32 = 0, + picks_host: cuda.HostBuffer, + picks_dev: cuda.DeviceBuffer, + picks_ev: cuda.Event, + picks_pending: bool = false, + picks_n: usize = 0, // the planned round's picks (rows * slots) + prev: []u16, // the previous round's routed experts (double buffering) + prev_n: usize = 0, + stats: Stats = .{}, + telemetry_id: u64 = 0, + telemetry_valid: bool = true, + telemetry_device: telemetry.Device = .{}, + + /// `counts`: the family's lifetime use table shared by every layer's pool (null: the pool owns one). + pub fn init(gpa: std.mem.Allocator, io: std.Io, d: *const cuda.Driver, src: Source, scales: Scales, opt: Options, counts: ?*[shape.layers][shape.experts]u32) !Pool { + if (opt.lines < min_lines) return error.PoolTooSmall; + if (src.routed.len != shape.experts) return error.Invalid; + const tiles = (opt.dims / 16) * (opt.width / 16); + var bytes: [3]usize = .{ 0, 0, 0 }; + const k2 = try gpa.alloc([3]i32, shape.experts); + errdefer gpa.free(k2); + var held: usize = 0; + for (src.routed, k2) |e, *w| { + // an expert with no source bytes: never picked here + if (e[0].len == 0 and e[1].len == 0 and e[2].len == 0) { + w.* = .{ 0, 0, 0 }; + continue; + } + held += 1; + for (0..3) |p| { + // trellis bytes = tiles * 16 * k2 (k2 half-bits a value, 8 * k2 int16 words a 16x16 tile) + if (e[p].len == 0 or e[p].len % (tiles * 16) != 0) return error.Invalid; + w[p] = @intCast(e[p].len / (tiles * 16)); + bytes[p] = @max(bytes[p], e[p].len); + } + } + if (held == 0) return error.Invalid; + const table_e: usize = if (src.shared != null) shape.experts + 1 else shape.experts; + var lines = try Lines.init(gpa, d, .{ .experts = table_e, .bytes = bytes, .lines = opt.lines, .staging = if (opt.direct) 0 else opt.staging }); + errdefer lines.deinit(); + lines.pinned_bank = if (opt.direct) src.pinned_bank else null; + if (src.shared) |sh| { + var w: [3]i32 = undefined; + for (0..3) |p| { + if (sh[p].len == 0 or sh[p].len % (tiles * 16) != 0) return error.Invalid; + w[p] = @intCast(sh[p].len / (tiles * 16)); + } + // the shared expert gets its own resident region (it may be wider than the routed experts) + try lines.setResident(shared_id, sh, w); + } + const n = opt.lines; + const slots_tab = try gpa.alloc(u16, n); + errdefer gpa.free(slots_tab); + const line_tab = try gpa.alloc(offload.Line, n); + errdefer gpa.free(line_tab); + const stamp_tab = try gpa.alloc(u32, n); + errdefer gpa.free(stamp_tab); + const layer_tab = try gpa.alloc(u8, n); + errdefer gpa.free(layer_tab); + const own = counts == null; + const c = counts orelse blk: { + const t = try gpa.create([shape.layers][shape.experts]u32); + t.* = @splat(@splat(0)); + break :blk t; + }; + errdefer if (own) gpa.destroy(c); + var rnd = try Round.init(gpa, n, opt.max_rows * opt.slots); + errdefer rnd.deinit(gpa); + const passes_n = if (opt.max_passes == 0) rnd.passes.len else @min(opt.max_passes, rnd.passes.len); + if (opt.max_passes != 0 and opt.prefetch_bytes != 0) return error.Invalid; // in-flight prefetches lengthen rounds + const pick_bytes = passes_n * opt.max_rows * opt.slots * 4; + var picks_host = try cuda.HostBuffer.alloc(d, pick_bytes); + errdefer picks_host.free(); + var picks_dev = try cuda.DeviceBuffer.alloc(d, pick_bytes); + errdefer picks_dev.free(); + var picks_ev = try cuda.Event.init(d, false); + errdefer picks_ev.deinit(); + const prev = try gpa.alloc(u16, shape.experts); + errdefer gpa.free(prev); + return .{ + .telemetry_id = try telemetry.identity(), + .prev = prev, + .gpa = gpa, + .io = io, + .opt = opt, + .src = src, + .scales = scales, + .lines = lines, + .k2 = k2, + .slots_tab = slots_tab, + .line_tab = line_tab, + .stamp_tab = stamp_tab, + .layer_tab = layer_tab, + .v = offload.Vram.init(slots_tab, line_tab, stamp_tab, layer_tab), + .counts = c, + .own_counts = own, + .round = rnd, + .picks_host = picks_host, + .picks_dev = picks_dev, + .picks_ev = picks_ev, + }; + } + + pub fn deinit(p: *Pool) void { + p.lines.deinit(); + p.picks_ev.deinit(); + p.picks_dev.free(); + p.picks_host.free(); + p.round.deinit(p.gpa); + if (p.own_counts) p.gpa.destroy(p.counts); + p.gpa.free(p.prev); + p.gpa.free(p.layer_tab); + p.gpa.free(p.stamp_tab); + p.gpa.free(p.line_tab); + p.gpa.free(p.slots_tab); + p.gpa.free(p.k2); + p.* = undefined; + } + + /// The `Experts` layer: the pool's device tables (fixed addresses) and the resident scales. + pub fn tables(p: *const Pool) exl3e.Layer { + return .{ + .gate = p.lines.table(0), + .up = p.lines.table(1), + .down = p.lines.table(2), + .gate_k2 = p.lines.widths(0), + .up_k2 = p.lines.widths(1), + .down_k2 = p.lines.widths(2), + .suh_g = p.scales.suh_g, + .suh_u = p.scales.suh_u, + .svh_g = p.scales.svh_g, + .svh_u = p.scales.svh_u, + .suh_d = p.scales.suh_d, + .svh_d = p.scales.svh_d, + .experts = p.lines.spec.experts, + .dims = p.opt.dims, + .width = p.opt.width, + .codebook = p.opt.codebook, + .act_mode = p.opt.act_mode, + }; + } + + /// Promotes prefetch lines whose copy landed (the policy's reap); never blocks. + fn reap(p: *Pool) !void { + for (p.line_tab, p.slots_tab, 0..) |state, id, i| { + if (state != .inflight) continue; + if (try p.lines.landed(i)) offload.reapAt(&p.v, p.opt.layer, id, p.now) catch unreachable; + } + } + + /// Host step of a round: `picks` [rows * slots] (int32, -1 skipped, 512 the shared expert), as the device holds + pub fn plan(p: *Pool, picks: []const i32, rows: usize) !*const Round { + return p.planInternal(picks, rows, true); + } + + fn planInternal(p: *Pool, picks: []const i32, rows: usize, observe: bool) !*const Round { + if (rows == 0 or rows > p.opt.max_rows or picks.len != rows * p.opt.slots) return error.Invalid; + p.now += 1; + p.round_no += 1; + try p.reap(); + const ropt: round_.Options = .{ + .swap_every = p.opt.swap_every, + .swap_cap = p.opt.swap_cap, + .prefetch_bytes = p.opt.prefetch_bytes, + .line_bytes = p.lines.layout.line_bytes, + .protect = if (p.opt.double_buffer) p.prev[0..p.prev_n] else &.{}, + .observe = observe, + }; + round_.plan(&p.v, p.opt.layer, picks, p.counts, p.now, p.round_no, ropt, &p.round) catch |e| switch (e) { + error.Busy => { + // every line is held by a copy in flight: let them land, then plan again (nothing was written) + p.stats.busy += 1; + try p.lines.drain(); + try p.reap(); + try round_.plan(&p.v, p.opt.layer, picks, p.counts, p.now, p.round_no, ropt, &p.round); + }, + else => return e, + }; + const r = &p.round; + if (r.n_passes * p.opt.max_rows * p.opt.slots * 4 > p.picks_dev.len) return error.PoolPasses; + p.picks_n = picks.len; + p.prev_n = r.n_ids; + @memcpy(p.prev[0..r.n_ids], r.ids[0..r.n_ids]); + if (r.protected) p.stats.protected += 1; + if (r.masked()) { + if (p.picks_pending) try p.picks_ev.synchronize(); // the previous round's uploads read these bytes + p.picks_pending = false; + const stride = p.opt.max_rows * p.opt.slots; + const host = p.picks_host.slice(i32); + var member: [shape.experts + 1]bool = undefined; + for (0..r.n_passes) |j| { + @memset(&member, false); + for (r.passIds(j)) |id| member[id] = true; + member[shared_id] = j == 0; + for (picks, host[j * stride ..][0..picks.len]) |pk, *o| { + o.* = if (pk >= 0 and pk <= shared_id and member[@intCast(pk)]) pk else -1; + } + } + } + p.stats.rounds += 1; + p.stats.passes += r.n_passes; + p.stats.hit_slots += r.hit_slots; + p.stats.miss_slots += r.miss_slots; + p.stats.unique += r.unique; + p.stats.misses += r.misses_(); + p.stats.evictions += r.evictions; + return r; + } + + fn loadCopy(p: *Pool, c: round_.Copy) !void { + if (p.k2[c.id][0] == 0) return error.Invalid; // an expert without source bytes reached the planner + const ev: ?usize = if (c.evicted == offload.Cache.empty_id) null else c.evicted; + try p.lines.load(c.line, c.id, p.src.routed[c.id], p.k2[c.id], ev, p.opt.direct, p.io); + } + + /// Device step of pass `j`: its copies and table rewrites on the copy stream, its masked picks (if any) on `s` + pub fn stage(p: *Pool, r: *const Round, j: usize, s: cuda.Stream) !void { + for (r.passCopies(j)) |c| try p.loadCopy(c); + if (r.masked()) { + const stride = p.opt.max_rows * p.opt.slots * 4; + const n = p.picks_n * 4; + try p.picks_dev.d.check(p.picks_dev.d.api.cuMemcpyHtoDAsync_v2(p.picks_dev.ptr + j * stride, p.picks_host.bytes[j * stride ..].ptr, n, s.handle), "cuMemcpyHtoDAsync"); + try p.picks_ev.record(s); + p.picks_pending = true; + } + for (r.passLines(j)) |line| try p.lines.gate(s, line); + } + + /// The picks pass `j` reads: the round's own device picks, or the pool's masked copy for a split round. + pub fn picksFor(p: *const Pool, r: *const Round, j: usize, own: u64) u64 { + if (!r.masked()) return own; + return p.picks_dev.ptr + j * p.opt.max_rows * p.opt.slots * 4; + } + + /// After pass `j` was enqueued on `s`: its lines are protected until it finishes. + pub fn release(p: *Pool, r: *const Round, j: usize, s: cuda.Stream) !void { + try p.lines.release(s, r.passLines(j)); + if (j + 1 == r.n_passes) { + for (r.prefetch()) |c| try p.loadCopy(c); + p.stats.prefetches += r.n_prefetch; + } + } + + pub fn warmProfile(pool: *Pool) !usize { + return @import("exl3_frequency_seed.zig").warm(pool); + } + + pub fn deviceBytes(p: *const Pool) usize { + return p.lines.deviceBytes() + p.picks_dev.len; + } + + pub fn lineBytes(p: *const Pool) usize { + return p.lines.layout.line_bytes; + } +}; diff --git a/zig/src/families/flashnext/exl3/offload/exl3_frequency_seed.zig b/zig/src/families/flashnext/exl3/offload/exl3_frequency_seed.zig new file mode 100644 index 0000000000..274e8aa982 --- /dev/null +++ b/zig/src/families/flashnext/exl3/offload/exl3_frequency_seed.zig @@ -0,0 +1,139 @@ +const std = @import("std"); + +pub fn select(counts: []const u32, widths: []const [3]i32, output: []u16) !usize { + if (counts.len != 512 or widths.len != counts.len or output.len > counts.len) return error.InvalidProfileShape; + var order: [512]u16 = undefined; + var used: usize = 0; + for (counts, widths, 0..) |count, width, expert| { + if (count == 0 or width[0] <= 0 or width[1] <= 0 or width[2] <= 0) continue; + order[used] = @intCast(expert); + used += 1; + } + const Ranking = struct { + fn hotter(histogram: []const u32, left: u16, right: u16) bool { + if (histogram[left] != histogram[right]) return histogram[left] > histogram[right]; + return left < right; + } + }; + std.mem.sort(u16, order[0..used], counts, Ranking.hotter); + const selected = @min(output.len, used); + @memcpy(output[0..selected], order[0..selected]); + return selected; +} + +pub fn warm(pool: anytype) !usize { + const policy = @import("exl3_offload.zig"); + for (pool.line_tab, pool.slots_tab) |state, expert| { + if (state != .empty or expert != policy.Cache.empty_id) return error.ProfilePoolNotEmpty; + } + if (pool.round_no != 0) return error.ProfilePoolNotEmpty; + var selected: [512]u16 = undefined; + const count = try select(&pool.counts[pool.opt.layer], pool.k2, selected[0..@min(pool.opt.lines, selected.len)]); + for (selected[0..count], 0..) |expert, held| { + const line = try policy.stageUsedAt(&pool.v, pool.opt.layer, expert, selected[0..held], pool.counts); + try pool.lines.load(line, expert, pool.src.routed[expert], pool.k2[expert], null, pool.opt.direct, pool.io); + } + try pool.lines.drain(); + for (selected[0..count]) |expert| try policy.reapAt(&pool.v, pool.opt.layer, expert, pool.now); + return count; +} + +test "exl3: frequency seed ranks local positive experts without heating observations" { + var histogram: [512]u32 = @splat(0); + var widths: [512][3]i32 = @splat(.{ 0, 0, 0 }); + histogram[2] = 19; + histogram[3] = 19; + histogram[4] = 20; + histogram[511] = std.math.maxInt(u32); + widths[2] = .{ 1, 1, 1 }; + widths[3] = .{ 1, 1, 1 }; + widths[4] = .{ 1, 0, 1 }; + var output: [2]u16 = undefined; + const before = histogram; + try std.testing.expectEqual(@as(usize, 2), try select(&histogram, &widths, &output)); + try std.testing.expectEqualSlices(u16, &.{ 2, 3 }, &output); + try std.testing.expectEqualSlices(u32, &before, &histogram); +} + +test "exl3: frequency seed leaves spare lines for unknown and unobserved experts" { + var histogram: [512]u32 = @splat(0); + const widths: [512][3]i32 = @splat(.{ 1, 1, 1 }); + var output: [512]u16 = undefined; + try std.testing.expectEqual(@as(usize, 0), try select(&histogram, &widths, &output)); + histogram[511] = 1; + try std.testing.expectEqual(@as(usize, 1), try select(&histogram, &widths, &output)); + try std.testing.expectEqual(@as(u16, 511), output[0]); + try std.testing.expectEqual(@as(usize, 0), try select(&histogram, &widths, output[0..0])); + try std.testing.expectError(error.InvalidProfileShape, select(histogram[0..511], &widths, &output)); +} + +test "exl3: frequency warm publishes residency only after copies drain and never heats counts" { + const policy = @import("exl3_offload.zig"); + const Copies = struct { + ids: [2]u16 = undefined, + used: usize = 0, + drained: bool = false, + states: []policy.Line, + fail: bool = false, + + pub fn load(copies: *@This(), line: usize, expert: usize, _: [3][]const u8, _: [3]i32, evicted: ?usize, _: bool, _: void) !void { + try std.testing.expect(evicted == null); + try std.testing.expectEqual(policy.Line.inflight, copies.states[line]); + if (copies.fail) return error.CopyFailed; + copies.ids[copies.used] = @intCast(expert); + copies.used += 1; + } + + pub fn drain(copies: *@This()) !void { + for (copies.states) |state| try std.testing.expectEqual(policy.Line.inflight, state); + copies.drained = true; + } + }; + const Fake = struct { + line_tab: []policy.Line, + slots_tab: []u16, + round_no: u32 = 0, + counts: *[48][512]u32, + opt: struct { layer: u8 = 7, lines: usize = 2, direct: bool = true } = .{}, + k2: []const [3]i32, + v: policy.Vram, + lines: Copies, + src: struct { routed: []const [3][]const u8 }, + io: void = {}, + now: u32 = 0, + }; + const allocator = std.testing.allocator; + const counts = try allocator.create([48][512]u32); + defer allocator.destroy(counts); + counts.* = std.mem.zeroes([48][512]u32); + counts[7][2] = 100; + counts[7][4] = 200; + const before = counts[7]; + const widths: [512][3]i32 = @splat(.{ 1, 1, 1 }); + const sources: [512][3][]const u8 = @splat(.{ &.{}, &.{}, &.{} }); + var ids: [2]u16 = undefined; + var states: [2]policy.Line = undefined; + var stamps: [2]u32 = undefined; + var layers: [2]u8 = undefined; + var fake: Fake = .{ + .line_tab = &states, + .slots_tab = &ids, + .counts = counts, + .k2 = &widths, + .v = policy.Vram.init(&ids, &states, &stamps, &layers), + .lines = .{ .states = &states }, + .src = .{ .routed = &sources }, + }; + try std.testing.expectEqual(@as(usize, 2), try warm(&fake)); + try std.testing.expect(fake.lines.drained); + try std.testing.expectEqualSlices(u16, &.{ 4, 2 }, &ids); + try std.testing.expectEqualSlices(u32, &before, &counts[7]); + for (states) |state| try std.testing.expectEqual(policy.Line.resident, state); + try std.testing.expectError(error.ProfilePoolNotEmpty, warm(&fake)); + fake.v = policy.Vram.init(&ids, &states, &stamps, &layers); + fake.lines = .{ .states = &states, .fail = true }; + try std.testing.expectError(error.CopyFailed, warm(&fake)); + try std.testing.expect(!fake.lines.drained); + try std.testing.expectEqual(policy.Line.inflight, states[0]); + try std.testing.expectEqualSlices(u32, &before, &counts[7]); +} diff --git a/zig/src/families/flashnext/exl3/offload/exl3_line_deltas.zig b/zig/src/families/flashnext/exl3/offload/exl3_line_deltas.zig new file mode 100644 index 0000000000..50e9eca705 --- /dev/null +++ b/zig/src/families/flashnext/exl3/offload/exl3_line_deltas.zig @@ -0,0 +1,120 @@ +const std = @import("std"); + +pub const Plan = struct { + pool_bytes: u64, +}; + +pub fn apply(text: []const u8, uniform_lines: usize, decoder_pools: usize, line_bytes: []const usize, min_lines: usize, fixed_bytes: u64, startup_reservation: u64, counts: []usize) !Plan { + if (counts.len != line_bytes.len or decoder_pools > counts.len or min_lines == 0) return error.InvalidLineDeltaIndex; + if (uniform_lines < min_lines) return error.PoolTooSmall; + @memset(counts, 0); + var added: u64 = 0; + var removed: u64 = 0; + var entries = std.mem.splitScalar(u8, text, ','); + while (entries.next()) |entry| { + const colon = std.mem.indexOfScalar(u8, entry, ':') orelse return error.InvalidLineDeltas; + const index = try decimal(entry[0..colon]); + if (index >= decoder_pools) return error.InvalidLineDeltaIndex; + if (counts[index] != 0) return error.DuplicateLineDelta; + const delta = entry[colon + 1 ..]; + if (delta.len < 2 or (delta[0] != '+' and delta[0] != '-')) return error.InvalidLineDeltas; + const magnitude = try decimal(delta[1..]); + const result = if (delta[0] == '+') result: { + added = try std.math.add(u64, added, magnitude); + break :result try std.math.add(usize, uniform_lines, magnitude); + } else result: { + removed = try std.math.add(u64, removed, magnitude); + if (magnitude > uniform_lines) return error.PoolTooSmall; + break :result uniform_lines - magnitude; + }; + if (result < min_lines) return error.PoolTooSmall; + counts[index] = result; + } + if (added != removed) return error.NonZeroLineDeltas; + var before: u64 = 0; + var after: u64 = 0; + var expert_bytes: u64 = 0; + for (counts, line_bytes) |*count, bytes| { + if (count.* == 0) count.* = uniform_lines; + before = try std.math.add(u64, before, try std.math.mul(u64, uniform_lines, bytes)); + after = try std.math.add(u64, after, try std.math.mul(u64, count.*, bytes)); + expert_bytes = try std.math.add(u64, expert_bytes, try std.math.mul(u64, count.*, bytes)); + } + if (before != after) return error.LineDeltaByteBudgetChanged; + const pool_bytes = try std.math.add(u64, expert_bytes, fixed_bytes); + if (pool_bytes > startup_reservation) return error.InsufficientStartupMemory; + return .{ .pool_bytes = pool_bytes }; +} + +fn decimal(text: []const u8) !usize { + if (text.len == 0) return error.InvalidLineDeltas; + for (text) |digit| if (digit < '0' or digit > '9') return error.InvalidLineDeltas; + return std.fmt.parseInt(usize, text, 10); +} + +test "exl3: line deltas parse signed decoder ordinals and leave MTP unchanged" { + var counts: [22]usize = undefined; + const bytes: [22]usize = @splat(256); + const plan = try apply("20:+1,15:-1", 32, 21, &bytes, 10, 100, std.math.maxInt(u64), &counts); + for (counts, 0..) |count, index| try std.testing.expectEqual(@as(usize, if (index == 20) 33 else if (index == 15) 31 else 32), count); + try std.testing.expectEqual(22 * 32 * 256 + 100, plan.pool_bytes); + try std.testing.expectEqual(@as(usize, 32), counts[21]); +} + +test "exl3: line deltas reject malformed empty duplicate and out of range maps" { + var counts: [3]usize = undefined; + const bytes = [_]usize{ 256, 256, 512 }; + for ([_][]const u8{ "", "0", "0:1", "0:+", ":+1", "0:++1", "0:+1,", ",0:+1", "0:+1,,1:-1", " 0:+1,1:-1", "0:+1, 1:-1", "0:+1 ,1:-1", "-1:+1", "+0:+1", "0x0:+1", "0:+1_0", "0:+1:0" }) |text| { + try std.testing.expectError(error.InvalidLineDeltas, apply(text, 32, 2, &bytes, 10, 0, std.math.maxInt(u64), &counts)); + } + for ([_][]const u8{ "0:+1,0:-1", "0:+0,0:-0", "00:+0,0:+0" }) |text| { + try std.testing.expectError(error.DuplicateLineDelta, apply(text, 32, 2, &bytes, 10, 0, std.math.maxInt(u64), &counts)); + } + for ([_][]const u8{ "2:+0", "3:+0", "99:+0" }) |text| { + try std.testing.expectError(error.InvalidLineDeltaIndex, apply(text, 32, 2, &bytes, 10, 0, std.math.maxInt(u64), &counts)); + } + try std.testing.expectError(error.InvalidLineDeltaIndex, apply("0:+0", 32, 0, &bytes, 10, 0, std.math.maxInt(u64), &counts)); +} + +test "exl3: line deltas conserve lines and actual bytes independently" { + var counts: [4]usize = undefined; + const bytes = [_]usize{ 256, 512, 256, 1024 }; + for ([_][]const u8{ "0:+1", "0:-1", "0:+2,2:-1" }) |text| { + try std.testing.expectError(error.NonZeroLineDeltas, apply(text, 32, 3, &bytes, 10, 0, std.math.maxInt(u64), &counts)); + } + for ([_][]const u8{ "0:+1,1:-1", "0:-1,1:+1" }) |text| { + try std.testing.expectError(error.LineDeltaByteBudgetChanged, apply(text, 32, 3, &bytes, 10, 0, std.math.maxInt(u64), &counts)); + } + const expected_counts = [_]usize{ 33, 32, 31, 32 }; + const expected_bytes = 32 * (256 + 512 + 256 + 1024) + 100; + const plan = try apply("0:+1,2:-1", 32, 3, &bytes, 10, 100, expected_bytes, &counts); + try std.testing.expectEqualSlices(usize, &expected_counts, &counts); + try std.testing.expectEqual(expected_bytes, plan.pool_bytes); + try std.testing.expectError(error.InsufficientStartupMemory, apply("0:+1,2:-1", 32, 3, &bytes, 10, 100, expected_bytes - 1, &counts)); +} + +test "exl3: line deltas enforce floor before allocation and reject arithmetic overflow" { + var counts: [2]usize = undefined; + const bytes = [_]usize{ 256, 256 }; + _ = try apply("0:+22,1:-22", 32, 2, &bytes, 10, 0, std.math.maxInt(u64), &counts); + try std.testing.expectEqualSlices(usize, &.{ 54, 10 }, &counts); + for ([_][]const u8{ "0:+23,1:-23", "0:+32,1:-32", "0:+33,1:-33" }) |text| { + try std.testing.expectError(error.PoolTooSmall, apply(text, 32, 2, &bytes, 10, 0, std.math.maxInt(u64), &counts)); + } + try std.testing.expectError(error.PoolTooSmall, apply("0:+0", 9, 2, &bytes, 10, 0, std.math.maxInt(u64), &counts)); + try std.testing.expectError(error.Overflow, apply("0:+18446744073709551615,1:-18446744073709551615", 32, 2, &bytes, 10, 0, std.math.maxInt(u64), &counts)); + try std.testing.expectError(error.Overflow, apply("0:+1,1:-1", 32, 2, &.{ std.math.maxInt(usize), std.math.maxInt(usize) }, 10, 0, std.math.maxInt(u64), &counts)); +} + +test "exl3: line deltas allow byte neutral multiway transfers within conservative startup sizing" { + var counts: [4]usize = undefined; + const bytes = [_]usize{ 256, 512, 768, 1024 }; + const fixed: u64 = 100; + const reservation: u64 = 32 * (256 + 512 + 768 + 1024) + fixed; + const plan = try apply("0:+1,1:-2,2:+1", 32, 3, &bytes, 10, fixed, reservation, &counts); + try std.testing.expectEqualSlices(usize, &.{ 33, 30, 33, 32 }, &counts); + try std.testing.expect(plan.pool_bytes <= reservation); + try std.testing.expectEqual(reservation, plan.pool_bytes); + _ = try apply("0:+0", 32, 3, &bytes, 10, fixed, reservation, &counts); + try std.testing.expectEqualSlices(usize, &.{ 32, 32, 32, 32 }, &counts); +} diff --git a/zig/src/families/flashnext/exl3/offload/exl3_offload.zig b/zig/src/families/flashnext/exl3/offload/exl3_offload.zig new file mode 100644 index 0000000000..29376ad764 --- /dev/null +++ b/zig/src/families/flashnext/exl3/offload/exl3_offload.zig @@ -0,0 +1,52 @@ +//! Flash Next EXL3 offload hub: re-exports the policy, VRAM, count, fill and scheduling modules under one name. + +const policy = @import("exl3_offload_policy.zig"); +const vram = @import("exl3_offload_vram.zig"); +const counts = @import("exl3_offload_counts.zig"); +const fill = @import("exl3_offload_fill.zig"); +const used = @import("exl3_offload_used.zig"); + +pub const shape = policy.shape; +pub const shared_id = policy.shared_id; +pub const Decision = policy.Decision; +pub const Policy = policy.Policy; +pub const Cache = policy.Cache; +pub const init = policy.init; +pub const find = policy.find; +pub const planRound = policy.planRound; +pub const swap = policy.swap; +pub const installProfile = policy.installProfile; +pub const Line = vram.Line; +pub const Vram = vram.Vram; +pub const onDevice = vram.onDevice; +pub const residentAt = vram.residentAt; +pub const touch = vram.touch; +pub const stage = vram.stage; +pub const stageKeeping = vram.stageKeeping; +pub const schedule = vram.schedule; +pub const reapAt = vram.reapAt; +pub const reap = vram.reap; +pub const vramIndexAt = vram.vramIndexAt; +pub const touchAt = vram.touchAt; +pub const lineOf = vram.lineOf; +pub const Copy = vram.Copy; +pub const lineIds = vram.lineIds; +pub const stagedCopies = vram.stagedCopies; +pub const countAt = counts.countAt; +pub const observe = counts.observe; +pub const Owned = counts.Owned; +pub const allOwned = counts.allOwned; +pub const seedHottest = fill.seedHottest; +pub const fillOwned = fill.fillOwned; +pub const fillOwnedLayers = fill.fillOwnedLayers; +pub const stageUsed = used.stageUsed; +pub const stageUsedAt = used.stageUsedAt; +pub const scheduleUsed = used.scheduleUsed; +pub const scheduleUsedAt = used.scheduleUsedAt; + +comptime { + _ = @import("exl3_offload_policy_test.zig"); + _ = @import("exl3_offload_vram_test.zig"); + _ = @import("exl3_offload_used_test.zig"); + _ = @import("exl3_offload_fill_test.zig"); +} diff --git a/zig/src/families/flashnext/exl3/offload/exl3_offload_counts.zig b/zig/src/families/flashnext/exl3/offload/exl3_offload_counts.zig new file mode 100644 index 0000000000..4b212ab68c --- /dev/null +++ b/zig/src/families/flashnext/exl3/offload/exl3_offload_counts.zig @@ -0,0 +1,53 @@ +//! Flash Next EXL3 offload: lifetime routing counts and the expert-id span a rank owns. + +const std = @import("std"); +const pol = @import("exl3_offload_policy.zig"); + +const shape = pol.shape; +const appearances = pol.appearances; + +pub fn countOf(counts: []const u32, id: u16) u32 { + if (id >= counts.len) return 0; + return counts[id]; +} + +pub fn countAt(counts: *const [shape.layers][shape.experts]u32, layer: u8, id: u16) u32 { + if (layer >= shape.layers or id >= shape.experts) return 0; + return counts[layer][id]; +} + +pub const UsedOrder = struct { + counts: []const u32, + ids: []const u16, +}; + +/// Lifetime count descending, then this round's appearances, then lower id. +pub fn usedFirst(ctx: UsedOrder, a: u16, b: u16) bool { + const ca = countOf(ctx.counts, a); + const cb = countOf(ctx.counts, b); + if (ca != cb) return ca > cb; + const aa = appearances(ctx.ids, a); + const ab = appearances(ctx.ids, b); + if (aa != ab) return aa > ab; + return a < b; +} + +/// Add one to each routed id in `ids`. +pub fn observe(counts: []u32, ids: []const u16) error{Length}!void { + if (counts.len != shape.experts) return error.Length; + for (ids) |id| { + if (id >= shape.experts) continue; + counts[id] +|= 1; + } +} + +/// Expert ids this rank stores. Rank 0 is `0..256`, rank 1 is `256..512`. +pub const Owned = struct { lo: u16, hi: u16 }; + +pub fn allOwned() Owned { + return .{ .lo = 0, .hi = @intCast(shape.experts) }; +} + +pub fn inOwned(owned: Owned, id: u16) bool { + return id >= owned.lo and id < owned.hi; +} diff --git a/zig/src/families/flashnext/exl3/offload/exl3_offload_fill.zig b/zig/src/families/flashnext/exl3/offload/exl3_offload_fill.zig new file mode 100644 index 0000000000..c81bd6df4b --- /dev/null +++ b/zig/src/families/flashnext/exl3/offload/exl3_offload_fill.zig @@ -0,0 +1,145 @@ +//! Flash Next EXL3 offload: first fill of the VRAM lines, hottest-first and striped across layers. + +const std = @import("std"); +const pol = @import("exl3_offload_policy.zig"); +const vram_ = @import("exl3_offload_vram.zig"); +const counts_ = @import("exl3_offload_counts.zig"); + +const shape = pol.shape; +const shared_id = pol.shared_id; +const Cache = pol.Cache; +const Vram = vram_.Vram; +const stage = vram_.stage; +const vramEmpty = vram_.vramEmpty; +const vramIndex = vram_.vramIndex; +const vramIndexAt = vram_.vramIndexAt; +const Owned = counts_.Owned; +const allOwned = counts_.allOwned; + +/// Fill empty VRAM lines with the hottest counted ids. +pub fn seedHottest(v: *Vram, counts: []const u32, limit: u16) error{Length}!u16 { + return seedSpan(v, counts, limit, allOwned()); +} + +fn seedSpan(v: *Vram, counts: []const u32, limit: u16, owned: Owned) error{Length}!u16 { + if (counts.len != shape.experts) return error.Length; + var seeded: u16 = 0; + while (seeded < limit) { + if (vramEmpty(v.*) == null) break; + var best: ?u16 = null; + var best_n: u32 = 0; + var id: u16 = owned.lo; + while (id < owned.hi) : (id += 1) { + const count = counts[id]; + if (count == 0) continue; + if (vramIndex(v.*, id) != null) continue; + if (best == null or count > best_n) { + best = id; + best_n = count; + } + } + if (best == null) break; + _ = stage(v, best.?) catch break; + seeded += 1; + } + return seeded; +} + +/// Use every empty line. Hottest owned ids go in first. Lines that are still empty take the remaining owned ids. +pub fn fillOwned(v: *Vram, counts: []const u32, owned: Owned) error{Length}!u16 { + if (owned.hi > shape.experts or owned.lo > owned.hi) return error.Length; + var n: u16 = try seedSpan(v, counts, std.math.maxInt(u16), owned); + var id: u16 = owned.lo; + while (id < owned.hi) : (id += 1) { + if (vramEmpty(v.*) == null) break; + if (vramIndex(v.*, id) != null) continue; + _ = stage(v, id) catch break; + n += 1; + } + return n; +} + +fn quotaAt(lines: u16, layer: usize) u16 { + const n_layers: u16 = @intCast(shape.layers); + const base = lines / n_layers; + if (layer < @as(usize, lines % n_layers)) return base + 1; + return base; +} + +fn placeOwned(v: *Vram, layer: u8, id: u16) bool { + // Shared 512 is not a routed line. Callers only pass an owned id. + std.debug.assert(id != shared_id and id != Cache.empty_id); + const slot = vramEmpty(v.*) orelse return false; + v.slots[slot] = id; + v.layer[slot] = layer; + v.line[slot] = .inflight; + v.stamp[slot] = 0; + return true; +} + +/// Hottest owned `(layer, id)` that still has stripe room, one line at a time. +fn placeHottest( + v: *Vram, + lines: u16, + counts: []const [shape.experts]u32, + owned: Owned, + used: *[shape.layers]u16, +) u16 { + var placed: u16 = 0; + while (placed < lines and vramEmpty(v.*) != null) { + var found = false; + var best_layer: u8 = 0; + var best_id: u16 = 0; + var best_c: u32 = 0; + for (counts, 0..) |row, layer_i| { + if (used[layer_i] >= quotaAt(lines, layer_i)) continue; + var id = owned.lo; + while (id < owned.hi) : (id += 1) { + const c = row[id]; + if (c == 0) continue; + if (vramIndexAt(v.*, @intCast(layer_i), id) != null) continue; + const layer: u8 = @intCast(layer_i); + const better = !found or c > best_c or + (c == best_c and id < best_id) or + (c == best_c and id == best_id and layer < best_layer); + if (!better) continue; + found = true; + best_layer = layer; + best_id = id; + best_c = c; + } + } + if (!found) break; + if (!placeOwned(v, best_layer, best_id)) break; + used[best_layer] += 1; + placed += 1; + } + return placed; +} + +/// Id, then layer. An uncapped layer-major walk would park every line on layer 0. +fn fillStripe(v: *Vram, lines: u16, owned: Owned, used: *[shape.layers]u16, placed: u16) u16 { + var n = placed; + var id = owned.lo; + while (id < owned.hi and n < lines) : (id += 1) { + var layer: usize = 0; + while (layer < shape.layers and n < lines) : (layer += 1) { + if (used[layer] >= quotaAt(lines, layer)) continue; + if (vramIndexAt(v.*, @intCast(layer), id) != null) continue; + if (!placeOwned(v, @intCast(layer), id)) return n; + used[layer] += 1; + n += 1; + } + } + return n; +} + +/// Stage up to `lines` owned experts. +pub fn fillOwnedLayers(v: *Vram, lines: u16, counts: []const [shape.experts]u32, owned: Owned) error{Length}!u16 { + if (counts.len != shape.layers) return error.Length; + if (owned.hi > shape.experts or owned.lo > owned.hi) return error.Length; + std.debug.assert(v.layer.len == v.slots.len and v.line.len == v.slots.len); + var used: [shape.layers]u16 = @splat(0); + const hot = placeHottest(v, lines, counts, owned, &used); + return fillStripe(v, lines, owned, &used, hot); +} diff --git a/zig/src/families/flashnext/exl3/offload/exl3_offload_fill_test.zig b/zig/src/families/flashnext/exl3/offload/exl3_offload_fill_test.zig new file mode 100644 index 0000000000..d886f29332 --- /dev/null +++ b/zig/src/families/flashnext/exl3/offload/exl3_offload_fill_test.zig @@ -0,0 +1,207 @@ +//! Flash Next EXL3 offload tests: the first fill of the VRAM lines. + +const std = @import("std"); +const hub = @import("exl3_offload.zig"); +const vram_ = @import("exl3_offload_vram.zig"); + +const shape = hub.shape; +const shared_id = hub.shared_id; +const Cache = hub.Cache; +const Decision = hub.Decision; +const Policy = hub.Policy; +const Line = hub.Line; +const Vram = hub.Vram; +const Owned = hub.Owned; +const Copy = hub.Copy; +const init = hub.init; +const find = hub.find; +const planRound = hub.planRound; +const swap = hub.swap; +const installProfile = hub.installProfile; +const onDevice = hub.onDevice; +const residentAt = hub.residentAt; +const touch = hub.touch; +const stage = hub.stage; +const stageKeeping = hub.stageKeeping; +const schedule = hub.schedule; +const reap = hub.reap; +const reapAt = hub.reapAt; +const vramIndexAt = hub.vramIndexAt; +const vramIndex = vram_.vramIndex; +const observe = hub.observe; +const allOwned = hub.allOwned; +const seedHottest = hub.seedHottest; +const fillOwned = hub.fillOwned; +const fillOwnedLayers = hub.fillOwnedLayers; +const scheduleUsed = hub.scheduleUsed; +const scheduleUsedAt = hub.scheduleUsedAt; +const lineOf = hub.lineOf; +const lineIds = hub.lineIds; +const stagedCopies = hub.stagedCopies; + +fn benchPolicy(slot_n: u16, expert_bytes: u64, copy_bytes: u64) Policy { + return .{ + .slots = slot_n, + .expert_bytes = expert_bytes, + .copy_bytes = copy_bytes, + }; +} + +test "observed counts seed empty lines and a later round keeps the hot id" { + var counts: [shape.experts]u32 = @splat(0); + try observe(&counts, &.{ 1, shared_id, 1, 3 }); + try std.testing.expectEqual(@as(u32, 2), counts[1]); + try std.testing.expectEqual(@as(u32, 1), counts[3]); + try observe(&counts, &.{1}); + try std.testing.expectEqual(@as(u32, 3), counts[1]); + var short: [1]u32 = .{9}; + try std.testing.expectError(error.Length, observe(&short, &.{1})); + try std.testing.expectEqual(@as(u32, 9), short[0]); + + var seeded_counts: [shape.experts]u32 = @splat(0); + seeded_counts[3] = 10; + seeded_counts[1] = 10; + seeded_counts[8] = 2; + var slots = [_]u16{ Cache.empty_id, Cache.empty_id }; + var line = [_]Line{ .empty, .empty }; + var stamp = [_]u32{ 0, 0 }; + var layer = [_]u8{ 0, 0 }; + var vram = Vram.init(&slots, &line, &stamp, &layer); + const n = try seedHottest(&vram, &seeded_counts, 2); + try std.testing.expectEqual(@as(u16, 2), n); + try std.testing.expectEqual(@as(u16, 1), slots[0]); + try std.testing.expectEqual(@as(u16, 3), slots[1]); + try std.testing.expectEqual(Line.inflight, line[0]); + try std.testing.expectEqual(Line.inflight, line[1]); + try std.testing.expect(!onDevice(vram, 1)); + try std.testing.expectEqual(@as(?usize, null), vramIndex(vram, 8)); + try reap(&vram, 1, 1); + try std.testing.expect(onDevice(vram, 1)); + + var wide: [3]u16 = .{ Cache.empty_id, Cache.empty_id, Cache.empty_id }; + var wide_line: [3]Line = .{ .empty, .empty, .empty }; + var wide_stamp: [3]u32 = .{ 0, 0, 0 }; + var wide_layer: [3]u8 = .{ 0, 0, 0 }; + var wide_vram = Vram.init(&wide, &wide_line, &wide_stamp, &wide_layer); + try std.testing.expectError(error.Length, seedHottest(&wide_vram, counts[0..3], 1)); + try std.testing.expectEqual(Line.empty, wide_line[0]); +} + +test "a rank fills its own lines and does not copy the other rank's expert" { + var slots = [_]u16{ Cache.empty_id, Cache.empty_id }; + var line = [_]Line{ .empty, .empty }; + var stamp = [_]u32{ 0, 0 }; + var layer = [_]u8{ 0, 0 }; + var vram = Vram.init(&slots, &line, &stamp, &layer); + var counts: [shape.experts]u32 = @splat(0); + counts[400] = 100; + counts[7] = 3; + const rank0 = Owned{ .lo = 0, .hi = 256 }; + const n = try fillOwned(&vram, &counts, rank0); + try std.testing.expectEqual(@as(u16, 2), n); + try std.testing.expectEqual(@as(u16, 7), slots[0]); + try std.testing.expectEqual(@as(u16, 0), slots[1]); + try std.testing.expectEqual(@as(?usize, null), vramIndex(vram, 400)); + try std.testing.expectEqual(Line.inflight, line[0]); + try reap(&vram, 7, 1); + try reap(&vram, 0, 1); + + var decisions: [2]Decision = undefined; + var scratch: [2]u16 = undefined; + const asked = [_]u16{ 400, 9 }; + try scheduleUsed(&vram, &asked, &counts, benchPolicy(2, 100, 10_000), 4, &decisions, &scratch, rank0); + try std.testing.expectEqual(Decision.cpu, decisions[0]); + try std.testing.expectEqual(Decision.copy, decisions[1]); + try std.testing.expect(onDevice(vram, 7)); + try std.testing.expectEqual(@as(?usize, null), vramIndex(vram, 400)); + try std.testing.expect(vramIndex(vram, 9) != null); +} + +test "fillOwnedLayers stripes an explicit line count and skips the other rank" { + const wide: u16 = @intCast(shape.layers + 2); + var slots: [shape.layers + 2]u16 = @splat(Cache.empty_id); + var line: [shape.layers + 2]Line = @splat(.empty); + var stamp: [shape.layers + 2]u32 = @splat(0); + var layer: [shape.layers + 2]u8 = @splat(0); + var vram = Vram.init(&slots, &line, &stamp, &layer); + var counts: [shape.layers][shape.experts]u32 = @splat(@splat(0)); + const rank0 = Owned{ .lo = 0, .hi = 256 }; + var short: [2][shape.experts]u32 = @splat(@splat(0)); + try std.testing.expectError(error.Length, fillOwnedLayers(&vram, wide, &short, rank0)); + try std.testing.expectError(error.Length, fillOwnedLayers(&vram, wide, &counts, .{ .lo = 4, .hi = 2 })); + try std.testing.expectEqual(Line.empty, line[0]); + + const n = try fillOwnedLayers(&vram, wide, &counts, rank0); + try std.testing.expectEqual(wide, n); + // Id, then layer: slot 1 is layer 1's lowest id, not layer 0's second id. + try std.testing.expectEqual(@as(u16, 0), slots[0]); + try std.testing.expectEqual(@as(u8, 0), layer[0]); + try std.testing.expectEqual(@as(u16, 0), slots[1]); + try std.testing.expectEqual(@as(u8, 1), layer[1]); + try std.testing.expectEqual(@as(u16, 0), slots[shape.layers - 1]); + try std.testing.expectEqual(@as(u8, @intCast(shape.layers - 1)), layer[shape.layers - 1]); + try std.testing.expectEqual(@as(u16, 1), slots[shape.layers]); + try std.testing.expectEqual(@as(u8, 0), layer[shape.layers]); + try std.testing.expectEqual(@as(u16, 1), slots[shape.layers + 1]); + try std.testing.expectEqual(@as(u8, 1), layer[shape.layers + 1]); + var on_layer0: usize = 0; + for (slots, layer) |id, lay| { + try std.testing.expect(id < 256); + try std.testing.expect(id != shared_id); + if (lay == 0) on_layer0 += 1; + } + try std.testing.expectEqual(@as(usize, 2), on_layer0); + + // Fewer lines than layers: the remainder is layers 0..3, one lowest id each. + var few_slots: [4]u16 = @splat(Cache.empty_id); + var few_line: [4]Line = @splat(.empty); + var few_stamp: [4]u32 = @splat(0); + var few_layer: [4]u8 = @splat(0); + var few = Vram.init(&few_slots, &few_line, &few_stamp, &few_layer); + try std.testing.expectEqual(@as(u16, 4), try fillOwnedLayers(&few, 4, &counts, rank0)); + try std.testing.expectEqual([_]u16{ 0, 0, 0, 0 }, few_slots); + try std.testing.expectEqual([_]u8{ 0, 1, 2, 3 }, few_layer); + + // Rank 1. A hotter id outside Owned is not staged. Layer 0 keeps its hottest owned id. + var rank_slots: [4]u16 = @splat(Cache.empty_id); + var rank_line: [4]Line = @splat(.empty); + var rank_stamp: [4]u32 = @splat(0); + var rank_layer: [4]u8 = @splat(0); + var rank = Vram.init(&rank_slots, &rank_line, &rank_stamp, &rank_layer); + counts[0][1] = 999; + counts[0][300] = 5; + const rank1 = Owned{ .lo = 256, .hi = 512 }; + try std.testing.expectEqual(@as(u16, 4), try fillOwnedLayers(&rank, 4, &counts, rank1)); + try std.testing.expect(vramIndexAt(rank, 0, 300) != null); + try std.testing.expect(vramIndexAt(rank, 1, 256) != null); + try std.testing.expect(vramIndexAt(rank, 2, 256) != null); + try std.testing.expect(vramIndexAt(rank, 3, 256) != null); + try std.testing.expectEqual(@as(?usize, null), vramIndex(rank, 1)); + for (rank_slots) |id| { + try std.testing.expect(id >= 256 and id < 512); + try std.testing.expect(id != shared_id); + } +} + +test "fillOwnedLayers keeps a hot id inside the stripe and does not give its layer a second line" { + const lines: u16 = @intCast(shape.layers); + var slots: [shape.layers]u16 = @splat(Cache.empty_id); + var line: [shape.layers]Line = @splat(.empty); + var stamp: [shape.layers]u32 = @splat(0); + var layer: [shape.layers]u8 = @splat(0); + var vram = Vram.init(&slots, &line, &stamp, &layer); + var counts: [shape.layers][shape.experts]u32 = @splat(@splat(0)); + counts[0][9] = 10; + counts[0][8] = 9; + const n = try fillOwnedLayers(&vram, lines, &counts, .{ .lo = 0, .hi = 256 }); + try std.testing.expectEqual(lines, n); + try std.testing.expect(vramIndexAt(vram, 0, 9) != null); + try std.testing.expectEqual(@as(?usize, null), vramIndexAt(vram, 0, 8)); + try std.testing.expect(vramIndexAt(vram, 1, 0) != null); + try std.testing.expect(vramIndexAt(vram, 40, 0) != null); + var on_layer0: usize = 0; + for (layer) |lay| { + if (lay == 0) on_layer0 += 1; + } + try std.testing.expectEqual(@as(usize, 1), on_layer0); +} diff --git a/zig/src/families/flashnext/exl3/offload/exl3_offload_policy.zig b/zig/src/families/flashnext/exl3/offload/exl3_offload_policy.zig new file mode 100644 index 0000000000..5c8dd9dc79 --- /dev/null +++ b/zig/src/families/flashnext/exl3/offload/exl3_offload_policy.zig @@ -0,0 +1,191 @@ +//! Flash Next EXL3 offload: the host slot cache and its per-round copy/cpu/hit planner. + +const std = @import("std"); +/// Flash Next's routed geometry (was exl3_shape.zig): 48 decoder layers, 512 routed experts. +pub const shape = struct { + pub const layers: usize = 48; + pub const experts: usize = 512; + pub const shared_expert: usize = 512; +}; + +pub const shared_id: u16 = @intCast(shape.shared_expert); + +pub const Decision = enum { hit, copy, cpu }; + +pub const Policy = struct { + slots: u16, + expert_bytes: u64, + /// Bytes the copy engine may move this round. copy_count = expert_bytes==0 ? 0 : copy_bytes / expert_bytes. + copy_bytes: u64, + swap_every: u16 = 4, + swap_cap: u16 = 96, +}; + +pub const Cache = struct { + slots: []u16, // expert id, or empty_id + freq: []u32, // frequency of the resident id, parallel to slots + pub const empty_id: u16 = 0xffff; +}; + +/// Clear every slot to `empty_id` and every frequency to 0. Lengths must match. +pub fn init(slots: []u16, freq: []u32) Cache { + std.debug.assert(slots.len == freq.len); + @memset(slots, Cache.empty_id); + @memset(freq, 0); + return .{ .slots = slots, .freq = freq }; +} + +/// Slot index of a resident expert, or null. `empty_id` is a vacant line, not an expert. +pub fn find(c: Cache, id: u16) ?usize { + if (id == Cache.empty_id) return null; + return indexOf(c.slots, id); +} + +/// One decision per id. Does not write the slot table. Unique misses sort by count descending. +pub fn planRound( + c: Cache, + ids: []const u16, + policy: Policy, + decisions: []Decision, + scratch: []u16, +) error{ Length, Scratch }!void { + if (decisions.len != ids.len) return error.Length; + std.debug.assert(c.slots.len == c.freq.len); + + var n: usize = 0; + for (ids) |id| { + if (id == shared_id) continue; + if (find(c, id) != null) continue; + try pushUnique(scratch, &n, id); + } + std.mem.sort(u16, scratch[0..n], ids, hotterFirst); + + const copies = copyCount(policy, n); + for (ids, decisions) |id, *decision| { + if (id == shared_id or find(c, id) != null) { + decision.* = .hit; + continue; + } + const rank = indexOf(scratch[0..n], id).?; + decision.* = if (rank < copies) .copy else .cpu; + } +} + +/// Install on swap rounds only (`round > 0` and `round % swap_every == 0`). +pub fn swap( + c: *Cache, + requested: []const u16, + round: u32, + policy: Policy, + scratch: []u16, +) error{Scratch}!u16 { + if (policy.swap_every == 0 or round == 0 or round % @as(u32, policy.swap_every) != 0) return 0; + std.debug.assert(c.slots.len == c.freq.len); + + var n: usize = 0; + for (requested) |id| { + if (id == shared_id) continue; + if (find(c.*, id) != null) continue; + try pushUnique(scratch, &n, id); + } + std.mem.sort(u16, scratch[0..n], requested, hotterFirst); + heatResidents(c, requested); + + const cap = @min(n, @as(usize, policy.swap_cap)); + var installed: usize = 0; + for (scratch[0..cap]) |id| { + const slot = indexOf(c.slots, Cache.empty_id) orelse victim(c.*, scratch[0..installed]) orelse break; + c.slots[slot] = id; + c.freq[slot] = appearances(requested, id); + installed += 1; + } + return @intCast(installed); +} + +/// Fill slots from a routed histogram of length 512. +pub fn installProfile(c: *Cache, histogram: []const u32) void { + std.debug.assert(histogram.len == shape.experts); + std.debug.assert(c.slots.len == c.freq.len); + @memset(c.slots, Cache.empty_id); + @memset(c.freq, 0); + + const n = @min(c.slots.len, histogram.len); + var filled: usize = 0; + while (filled < n) : (filled += 1) { + var best: u16 = 0; + var best_n: u32 = 0; + var any = false; + for (histogram, 0..) |count, i| { + const id: u16 = @intCast(i); + if (indexOf(c.slots[0..filled], id) != null) continue; + // Scan is id-ascending, so the first id that reaches the max count wins ties. + if (!any or count > best_n) { + best = id; + best_n = count; + any = true; + } + } + if (!any) break; + c.slots[filled] = best; + c.freq[filled] = best_n; + } +} + +pub fn copyCount(policy: Policy, unique_misses: usize) usize { + if (policy.expert_bytes == 0 or unique_misses == 0) return 0; + const raw = policy.copy_bytes / policy.expert_bytes; + if (raw >= unique_misses) return unique_misses; + return @intCast(raw); +} + +pub fn appearances(ids: []const u16, id: u16) u32 { + var n: u32 = 0; + for (ids) |x| { + if (x == id) n +|= 1; + } + return n; +} + +/// Count descending, id ascending. `ctx` is the multiset the counts are taken from. +pub fn hotterFirst(ctx: []const u16, a: u16, b: u16) bool { + const ca = appearances(ctx, a); + const cb = appearances(ctx, b); + if (ca != cb) return ca > cb; + return a < b; +} + +pub fn indexOf(ids: []const u16, id: u16) ?usize { + for (ids, 0..) |x, i| { + if (x == id) return i; + } + return null; +} + +pub fn pushUnique(dst: []u16, n: *usize, id: u16) error{Scratch}!void { + if (indexOf(dst[0..n.*], id) != null) return; + if (n.* >= dst.len) return error.Scratch; + dst[n.*] = id; + n.* += 1; +} + +fn heatResidents(c: *Cache, requested: []const u16) void { + for (c.slots, 0..) |id, i| { + if (id == Cache.empty_id) continue; + const n = appearances(requested, id); + if (n == 0) continue; + c.freq[i] = c.freq[i] +| n; + } +} + +/// Lowest frequency, skipping lines already placed in this call. +fn victim(c: Cache, skip: []const u16) ?usize { + var best: ?usize = null; + for (c.slots, 0..) |id, i| { + if (id == Cache.empty_id) continue; + if (indexOf(skip, id) != null) continue; + if (best) |b| { + if (c.freq[i] < c.freq[b] or (c.freq[i] == c.freq[b] and i > b)) best = i; + } else best = i; + } + return best; +} diff --git a/zig/src/families/flashnext/exl3/offload/exl3_offload_policy_test.zig b/zig/src/families/flashnext/exl3/offload/exl3_offload_policy_test.zig new file mode 100644 index 0000000000..a12152217e --- /dev/null +++ b/zig/src/families/flashnext/exl3/offload/exl3_offload_policy_test.zig @@ -0,0 +1,359 @@ +//! Flash Next EXL3 offload tests: the slot cache planner, swap and profile install. + +const std = @import("std"); +const hub = @import("exl3_offload.zig"); +const vram_ = @import("exl3_offload_vram.zig"); + +const shape = hub.shape; +const shared_id = hub.shared_id; +const Cache = hub.Cache; +const Decision = hub.Decision; +const Policy = hub.Policy; +const Line = hub.Line; +const Vram = hub.Vram; +const Owned = hub.Owned; +const Copy = hub.Copy; +const init = hub.init; +const find = hub.find; +const planRound = hub.planRound; +const swap = hub.swap; +const installProfile = hub.installProfile; +const onDevice = hub.onDevice; +const residentAt = hub.residentAt; +const touch = hub.touch; +const stage = hub.stage; +const stageKeeping = hub.stageKeeping; +const schedule = hub.schedule; +const reap = hub.reap; +const reapAt = hub.reapAt; +const vramIndexAt = hub.vramIndexAt; +const vramIndex = vram_.vramIndex; +const observe = hub.observe; +const allOwned = hub.allOwned; +const seedHottest = hub.seedHottest; +const fillOwned = hub.fillOwned; +const fillOwnedLayers = hub.fillOwnedLayers; +const scheduleUsed = hub.scheduleUsed; +const scheduleUsedAt = hub.scheduleUsedAt; +const lineOf = hub.lineOf; +const lineIds = hub.lineIds; +const stagedCopies = hub.stagedCopies; + +fn benchPolicy(slot_n: u16, expert_bytes: u64, copy_bytes: u64) Policy { + return .{ + .slots = slot_n, + .expert_bytes = expert_bytes, + .copy_bytes = copy_bytes, + }; +} + +test "resident ids are hits and planRound does not change slots" { + var slots = [_]u16{ 4, 1, 9 }; + var freq = [_]u32{ 3, 8, 1 }; + const cache = init(&slots, &freq); + slots = .{ 4, 1, 9 }; + freq = .{ 3, 8, 1 }; + const before_slots = slots; + const before_freq = freq; + + const ids = [_]u16{ 4, shared_id, 1, 9, 4 }; + var decisions: [5]Decision = undefined; + try planRound(cache, &ids, benchPolicy(3, 100, 10_000), &decisions, &[_]u16{}); + + for (decisions) |d| try std.testing.expectEqual(Decision.hit, d); + try std.testing.expectEqual(before_slots, slots); + try std.testing.expectEqual(before_freq, freq); + try std.testing.expectEqual(@as(?usize, 1), find(cache, 1)); + try std.testing.expectEqual(@as(?usize, null), find(cache, shared_id)); +} + +test "misses past the copy budget are cpu and slots stay put" { + var slots = [_]u16{42}; + var freq = [_]u32{6}; + const cache = init(&slots, &freq); + slots[0] = 42; + freq[0] = 6; + const before_slots = slots; + const before_freq = freq; + + // Counts: 5 x4, 1 x3, 8 x1, 3 x1. Budget 250/100 = 2, so 5 and 1 are copied. + const ids = [_]u16{ 8, 5, 1, 5, 3, 5, 1, 42, 5, 1 }; + var decisions: [ids.len]Decision = undefined; + var scratch: [8]u16 = undefined; + try planRound(cache, &ids, benchPolicy(1, 100, 250), &decisions, &scratch); + + const want = [_]Decision{ .cpu, .copy, .copy, .copy, .cpu, .copy, .copy, .hit, .copy, .copy }; + try std.testing.expectEqualSlices(Decision, &want, &decisions); + try std.testing.expectEqual(before_slots, slots); + try std.testing.expectEqual(before_freq, freq); +} + +test "duplicate miss ids share one decision" { + var slots: [1]u16 = undefined; + var freq: [1]u32 = undefined; + const cache = init(&slots, &freq); + const ids = [_]u16{ 6, 2, 6 }; + var decisions: [3]Decision = undefined; + var scratch: [4]u16 = undefined; + // One copy. 6 appears twice, 2 once, so both 6s are copy and 2 is cpu. + try planRound(cache, &ids, benchPolicy(1, 100, 100), &decisions, &scratch); + try std.testing.expectEqual(Decision.copy, decisions[0]); + try std.testing.expectEqual(Decision.cpu, decisions[1]); + try std.testing.expectEqual(Decision.copy, decisions[2]); +} + +test "equal miss counts break ties toward the lower id" { + var slots: [0]u16 = .{}; + var freq: [0]u32 = .{}; + const cache = init(&slots, &freq); + // 10 x5, 6 x4, 2 x4. One copy would be 10; two copies are 10 and 2, not 6. + const ids = [_]u16{ 6, 10, 2, 6, 10, 2, 10, 6, 2, 10, 6, 10, 2 }; + var decisions: [ids.len]Decision = undefined; + var scratch: [4]u16 = undefined; + try planRound(cache, &ids, benchPolicy(0, 50, 100), &decisions, &scratch); + for (ids, decisions) |id, d| { + const want: Decision = if (id == 10 or id == 2) .copy else .cpu; + try std.testing.expectEqual(want, d); + } +} + +test "shared id is a hit and is never stored" { + var slots: [4]u16 = @splat(Cache.empty_id); + var freq: [4]u32 = @splat(0); + var cache = init(&slots, &freq); + slots[0] = 7; + freq[0] = 2; + + const ids = [_]u16{ shared_id, 4, shared_id }; + var decisions: [3]Decision = undefined; + var scratch: [4]u16 = undefined; + try planRound(cache, &ids, benchPolicy(4, 100, 0), &decisions, &scratch); + try std.testing.expectEqual(Decision.hit, decisions[0]); + try std.testing.expectEqual(Decision.cpu, decisions[1]); + try std.testing.expectEqual(Decision.hit, decisions[2]); + try std.testing.expectEqual(@as(u16, 7), slots[0]); + + const requested = [_]u16{ shared_id, 4, shared_id, 9 }; + const installed = try swap(&cache, &requested, 4, benchPolicy(4, 1, 0), &scratch); + try std.testing.expectEqual(@as(u16, 2), installed); + try std.testing.expectEqual(@as(?usize, null), find(cache, shared_id)); + for (slots) |id| try std.testing.expect(id != shared_id); + + var hist: [shape.experts]u32 = @splat(0); + hist[4] = 3; + installProfile(&cache, &hist); + try std.testing.expectEqual(@as(?usize, null), find(cache, shared_id)); + for (slots) |id| try std.testing.expect(id != shared_id); +} + +test "round 2 installs nothing and round 4 installs min of cap and missing" { + var slots: [8]u16 = @splat(Cache.empty_id); + var freq: [8]u32 = @splat(0); + var cache = init(&slots, &freq); + const requested = [_]u16{ 3, 1, 3, 8 }; + var scratch: [8]u16 = undefined; + var policy = benchPolicy(8, 1, 0); + policy.swap_every = 4; + + try std.testing.expectEqual(@as(u16, 0), try swap(&cache, &requested, 0, policy, &scratch)); + try std.testing.expectEqual(@as(u16, 0), try swap(&cache, &requested, 2, policy, &scratch)); + try std.testing.expectEqual(@as(u16, Cache.empty_id), slots[0]); + try std.testing.expectEqual(@as(u32, 0), freq[0]); + + // Three unique misses, under the cap. + const n = try swap(&cache, &requested, 4, policy, &scratch); + try std.testing.expectEqual(@as(u16, 3), n); + try std.testing.expectEqual(@as(u16, 3), slots[0]); + try std.testing.expectEqual(@as(u32, 2), freq[0]); + try std.testing.expectEqual(@as(u16, 1), slots[1]); + try std.testing.expectEqual(@as(u16, 8), slots[2]); + try std.testing.expectEqual(@as(u16, Cache.empty_id), slots[3]); +} + +test "a cold line is replaced by the missing id" { + var slots = [_]u16{41}; + var freq = [_]u32{0}; + var cache = init(&slots, &freq); + slots[0] = 41; + freq[0] = 0; + var scratch: [2]u16 = undefined; + const requested = [_]u16{9}; + const n = try swap(&cache, &requested, 4, benchPolicy(1, 1, 0), &scratch); + try std.testing.expectEqual(@as(u16, 1), n); + try std.testing.expectEqual(@as(u16, 9), slots[0]); + try std.testing.expectEqual(@as(u32, 1), freq[0]); + try std.testing.expectEqual(@as(?usize, null), find(cache, 41)); +} + +test "eviction ties take the later slot" { + var slots = [_]u16{ 30, 10, 20 }; + var freq = [_]u32{ 4, 4, 4 }; + var cache = init(&slots, &freq); + slots = .{ 30, 10, 20 }; + freq = .{ 4, 4, 4 }; + var scratch: [2]u16 = undefined; + const n = try swap(&cache, &[_]u16{7}, 4, benchPolicy(3, 1, 0), &scratch); + try std.testing.expectEqual(@as(u16, 1), n); + // Higher id would have dropped 30. Higher slot index drops 20. + try std.testing.expectEqual([_]u16{ 30, 10, 7 }, slots); + try std.testing.expectEqual(@as(u32, 1), freq[2]); +} + +test "a resident named this round is heated before victims are chosen" { + var slots = [_]u16{ 2, 1 }; + var freq = [_]u32{ 0, 0 }; + var cache = init(&slots, &freq); + slots = .{ 2, 1 }; + freq = .{ 0, 0 }; + var scratch: [2]u16 = undefined; + // Without the heat, both freqs stay 0 and the later slot (id 1) would be the victim. + const n = try swap(&cache, &[_]u16{ 1, 99 }, 4, benchPolicy(2, 1, 0), &scratch); + try std.testing.expectEqual(@as(u16, 1), n); + try std.testing.expectEqual([_]u16{ 99, 1 }, slots); + try std.testing.expectEqual(@as(u32, 1), freq[0]); + try std.testing.expectEqual(@as(u32, 1), freq[1]); +} + +test "resident frequency adds this round's request count" { + var slots = [_]u16{ 5, Cache.empty_id }; + var freq = [_]u32{ 4, 0 }; + var cache = init(&slots, &freq); + slots = .{ 5, Cache.empty_id }; + freq = .{ 4, 0 }; + var scratch: [2]u16 = undefined; + const requested = [_]u16{ 5, 5, 5, 4 }; + const n = try swap(&cache, &requested, 4, benchPolicy(2, 1, 0), &scratch); + try std.testing.expectEqual(@as(u16, 1), n); + try std.testing.expectEqual(@as(u16, 5), slots[0]); + try std.testing.expectEqual(@as(u32, 7), freq[0]); + try std.testing.expectEqual(@as(u16, 4), slots[1]); + try std.testing.expectEqual(@as(u32, 1), freq[1]); +} + +test "empty lines are filled before a cold resident is evicted" { + var slots = [_]u16{ 30, Cache.empty_id, 40 }; + var freq = [_]u32{ 5, 0, 1 }; + var cache = init(&slots, &freq); + slots = .{ 30, Cache.empty_id, 40 }; + freq = .{ 5, 0, 1 }; + var scratch: [4]u16 = undefined; + const n = try swap(&cache, &[_]u16{ 7, 8 }, 4, benchPolicy(3, 1, 0), &scratch); + try std.testing.expectEqual(@as(u16, 2), n); + try std.testing.expectEqual([_]u16{ 30, 7, 8 }, slots); + try std.testing.expectEqual(@as(u32, 5), freq[0]); +} + +test "a candidate placed this call is not evicted by a later one" { + var slots = [_]u16{ 1, 2, 3, 4 }; + var freq = [_]u32{ 1, 50, 2, 40 }; + var cache = init(&slots, &freq); + slots = .{ 1, 2, 3, 4 }; + freq = .{ 1, 50, 2, 40 }; + var scratch: [4]u16 = undefined; + const n = try swap(&cache, &[_]u16{ 10, 11, 12 }, 4, benchPolicy(4, 1, 0), &scratch); + try std.testing.expectEqual(@as(u16, 3), n); + try std.testing.expectEqual([_]u16{ 10, 2, 11, 12 }, slots); + try std.testing.expectEqual(@as(u32, 50), freq[1]); +} + +test "installProfile keeps the hottest routed ids" { + var slots: [3]u16 = @splat(shared_id); + var freq = [_]u32{ 9, 9, 9 }; + var cache = init(&slots, &freq); + var hist: [shape.experts]u32 = @splat(0); + hist[10] = 1; + hist[4] = 7; + hist[8] = 7; + hist[2] = 3; + installProfile(&cache, &hist); + try std.testing.expectEqual([_]u16{ 4, 8, 2 }, slots); + try std.testing.expectEqual([_]u32{ 7, 7, 3 }, freq); + try std.testing.expectEqual(@as(?usize, null), find(cache, shared_id)); + + var wide_slots: [513]u16 = undefined; + var wide_freq: [513]u32 = undefined; + var wide = init(&wide_slots, &wide_freq); + var flat: [shape.experts]u32 = @splat(0); + flat[0] = 1; + installProfile(&wide, &flat); + for (wide_slots[0..shape.experts], 0..) |id, i| { + try std.testing.expectEqual(@as(u16, @intCast(i)), id); + } + try std.testing.expectEqual(Cache.empty_id, wide_slots[512]); + try std.testing.expectEqual(@as(u32, 1), wide_freq[0]); + try std.testing.expectEqual(@as(u32, 0), wide_freq[1]); + try std.testing.expectEqual(@as(?usize, null), find(wide, shared_id)); +} + +test "swap installs at most the cap even when slots and misses are larger" { + var slots: [200]u16 = @splat(Cache.empty_id); + var freq: [200]u32 = @splat(0); + var cache = init(&slots, &freq); + var requested: [200]u16 = undefined; + for (&requested, 0..) |*id, i| id.* = @intCast(i); + var scratch: [200]u16 = undefined; + const n = try swap(&cache, &requested, 4, benchPolicy(200, 1, 0), &scratch); + try std.testing.expectEqual(@as(u16, 96), n); + for (slots[0..96], 0..) |id, i| try std.testing.expectEqual(@as(u16, @intCast(i)), id); + try std.testing.expectEqual(Cache.empty_id, slots[96]); + try std.testing.expectEqual(@as(u32, 1), freq[0]); + try std.testing.expectEqual(@as(u32, 0), freq[96]); +} + +test "an empty slot table misses every routed id and installs nothing" { + var slots: [0]u16 = .{}; + var freq: [0]u32 = .{}; + var cache = init(&slots, &freq); + const ids = [_]u16{ shared_id, 4, 4, 8 }; + var decisions: [4]Decision = undefined; + var scratch: [2]u16 = undefined; + const policy = benchPolicy(0, 10, 10); + try planRound(cache, &ids, policy, &decisions, &scratch); + try std.testing.expectEqual(Decision.hit, decisions[0]); + try std.testing.expectEqual(Decision.copy, decisions[1]); + try std.testing.expectEqual(Decision.copy, decisions[2]); + try std.testing.expectEqual(Decision.cpu, decisions[3]); + + try std.testing.expectEqual(@as(u16, 0), try swap(&cache, &ids, 4, policy, &scratch)); + try std.testing.expectEqual(@as(usize, 0), slots.len); +} + +test "zero expert bytes copies nothing" { + var slots: [0]u16 = .{}; + var freq: [0]u32 = .{}; + const cache = init(&slots, &freq); + const ids = [_]u16{ 1, 2 }; + var decisions: [2]Decision = undefined; + var scratch: [2]u16 = undefined; + try planRound(cache, &ids, benchPolicy(0, 0, 1000), &decisions, &scratch); + try std.testing.expectEqual(Decision.cpu, decisions[0]); + try std.testing.expectEqual(Decision.cpu, decisions[1]); +} + +test "a short scratch or a short decision buffer is refused without writes" { + var slots = [_]u16{1}; + var freq = [_]u32{5}; + var cache = init(&slots, &freq); + slots[0] = 1; + freq[0] = 5; + + var decisions: [1]Decision = undefined; + try std.testing.expectError(error.Length, planRound(cache, &[_]u16{ 1, 2 }, benchPolicy(1, 1, 1), &decisions, &[_]u16{})); + try std.testing.expectEqual(@as(u16, 1), slots[0]); + try std.testing.expectEqual(@as(u32, 5), freq[0]); + + var one: [2]Decision = undefined; + try std.testing.expectError(error.Scratch, planRound(cache, &[_]u16{ 2, 3 }, benchPolicy(1, 1, 1), &one, &[_]u16{})); + try std.testing.expectEqual(@as(u16, 1), slots[0]); + + var tiny: [1]u16 = undefined; + try std.testing.expectError(error.Scratch, swap(&cache, &[_]u16{ 2, 3 }, 4, benchPolicy(1, 1, 0), &tiny)); + try std.testing.expectEqual(@as(u16, 1), slots[0]); + try std.testing.expectEqual(@as(u32, 5), freq[0]); + + var off = benchPolicy(1, 1, 0); + off.swap_every = 0; + var scratch: [2]u16 = undefined; + try std.testing.expectEqual(@as(u16, 0), try swap(&cache, &[_]u16{2}, 4, off, &scratch)); + try std.testing.expectEqual(@as(u32, 5), freq[0]); +} diff --git a/zig/src/families/flashnext/exl3/offload/exl3_offload_used.zig b/zig/src/families/flashnext/exl3/offload/exl3_offload_used.zig new file mode 100644 index 0000000000..602935cd22 --- /dev/null +++ b/zig/src/families/flashnext/exl3/offload/exl3_offload_used.zig @@ -0,0 +1,146 @@ +//! Flash Next EXL3 offload: VRAM scheduling ordered by lifetime counts, per layer and per rank. + +const std = @import("std"); +const pol = @import("exl3_offload_policy.zig"); +const vram_ = @import("exl3_offload_vram.zig"); +const counts_ = @import("exl3_offload_counts.zig"); + +const shape = pol.shape; +const shared_id = pol.shared_id; +const Cache = pol.Cache; +const Decision = pol.Decision; +const Policy = pol.Policy; +const indexOf = pol.indexOf; +const pushUnique = pol.pushUnique; +const copyCount = pol.copyCount; +const Vram = vram_.Vram; +const vramEmpty = vram_.vramEmpty; +const vramIndexAt = vram_.vramIndexAt; +const residentAt = vram_.residentAt; +const touchAt = vram_.touchAt; +const Owned = counts_.Owned; +const inOwned = counts_.inOwned; +const countAt = counts_.countAt; +const UsedOrder = counts_.UsedOrder; +const usedFirst = counts_.usedFirst; + +/// Lowest `counts[layer][id]`, then oldest stamp, then the higher slot index. +fn vramCold( + v: Vram, + counts: *const [shape.layers][shape.experts]u32, + keep_layer: u8, + keep_ids: []const u16, +) ?usize { + std.debug.assert(v.layer.len == v.slots.len); + var best: ?usize = null; + for (v.line, 0..) |state, i| { + if (state != .resident) continue; + if (v.layer[i] == keep_layer and indexOf(keep_ids, v.slots[i]) != null) continue; + if (best) |b| { + const ci = countAt(counts, v.layer[i], v.slots[i]); + const cb = countAt(counts, v.layer[b], v.slots[b]); + const colder = ci < cb or + (ci == cb and v.stamp[i] < v.stamp[b]) or + (ci == cb and v.stamp[i] == v.stamp[b] and i > b); + if (colder) best = i; + } else best = i; + } + return best; +} + +/// `stageKeeping`, but the victim is the coldest lifetime count. +pub fn stageUsed(v: *Vram, id: u16, keep: []const u16, counts: []const u32) error{ Busy, Shared, Length }!usize { + if (counts.len != shape.experts) return error.Length; + var table: [shape.layers][shape.experts]u32 = std.mem.zeroes([shape.layers][shape.experts]u32); + @memcpy(table[0][0..], counts); + return stageUsedAt(v, 0, id, keep, &table); +} + +pub fn stageUsedAt( + v: *Vram, + layer: u8, + id: u16, + keep_ids: []const u16, + counts: *const [shape.layers][shape.experts]u32, +) error{ Busy, Shared, Length }!usize { + if (layer >= shape.layers) return error.Length; + if (id == shared_id or id == Cache.empty_id) return error.Shared; + if (vramIndexAt(v.*, layer, id)) |i| { + if (v.line[i] == .inflight) return i; + v.line[i] = .inflight; + return i; + } + const slot = vramEmpty(v.*) orelse vramCold(v.*, counts, layer, keep_ids) orelse return error.Busy; + v.slots[slot] = id; + v.layer[slot] = layer; + v.line[slot] = .inflight; + v.stamp[slot] = 0; + return slot; +} + +/// `schedule`, ordered by lifetime counts from `observe`. +pub fn scheduleUsed( + v: *Vram, + ids: []const u16, + counts: []const u32, + policy: Policy, + now: u32, + decisions: []Decision, + scratch: []u16, + owned: Owned, +) error{ Length, Scratch }!void { + if (decisions.len != ids.len) return error.Length; + if (counts.len != shape.experts) return error.Length; + var table: [shape.layers][shape.experts]u32 = std.mem.zeroes([shape.layers][shape.experts]u32); + @memcpy(table[0][0..], counts); + return scheduleUsedAt(v, 0, ids, &table, policy, now, decisions, scratch, owned); +} + +/// Generate path. Misses sort by `counts[layer]`. A victim scores `counts[victim.layer][victim.id]`. Keep is `(layer). +pub fn scheduleUsedAt( + v: *Vram, + layer: u8, + ids: []const u16, + counts: *const [shape.layers][shape.experts]u32, + policy: Policy, + now: u32, + decisions: []Decision, + scratch: []u16, + owned: Owned, +) error{ Length, Scratch }!void { + if (decisions.len != ids.len) return error.Length; + if (layer >= shape.layers) return error.Length; + + var n: usize = 0; + for (ids, decisions) |id, *decision| { + if (id == shared_id) { + decision.* = .hit; + continue; + } + if (residentAt(v.*, layer, id)) { + _ = touchAt(v, layer, id, now); + decision.* = .hit; + continue; + } + decision.* = .cpu; + if (!inOwned(owned, id)) continue; + if (vramIndexAt(v.*, layer, id) != null) continue; + try pushUnique(scratch, &n, id); + } + const row: []const u32 = &counts[layer]; + std.mem.sort(u16, scratch[0..n], UsedOrder{ .counts = row, .ids = ids }, usedFirst); + + const copies = copyCount(policy, n); + var staged: usize = 0; + for (scratch[0..n]) |id| { + if (staged >= copies) break; + _ = stageUsedAt(v, layer, id, ids, counts) catch |err| switch (err) { + error.Length => return error.Length, + error.Busy, error.Shared => break, + }; + staged += 1; + for (ids, decisions) |asked, *decision| { + if (asked == id) decision.* = .copy; + } + } +} diff --git a/zig/src/families/flashnext/exl3/offload/exl3_offload_used_test.zig b/zig/src/families/flashnext/exl3/offload/exl3_offload_used_test.zig new file mode 100644 index 0000000000..0256922ca4 --- /dev/null +++ b/zig/src/families/flashnext/exl3/offload/exl3_offload_used_test.zig @@ -0,0 +1,276 @@ +//! Flash Next EXL3 offload tests: count-ordered scheduling and line addressing. + +const std = @import("std"); +const hub = @import("exl3_offload.zig"); +const vram_ = @import("exl3_offload_vram.zig"); + +const shape = hub.shape; +const shared_id = hub.shared_id; +const Cache = hub.Cache; +const Decision = hub.Decision; +const Policy = hub.Policy; +const Line = hub.Line; +const Vram = hub.Vram; +const Owned = hub.Owned; +const Copy = hub.Copy; +const init = hub.init; +const find = hub.find; +const planRound = hub.planRound; +const swap = hub.swap; +const installProfile = hub.installProfile; +const onDevice = hub.onDevice; +const residentAt = hub.residentAt; +const touch = hub.touch; +const stage = hub.stage; +const stageKeeping = hub.stageKeeping; +const schedule = hub.schedule; +const reap = hub.reap; +const reapAt = hub.reapAt; +const vramIndexAt = hub.vramIndexAt; +const vramIndex = vram_.vramIndex; +const observe = hub.observe; +const allOwned = hub.allOwned; +const seedHottest = hub.seedHottest; +const fillOwned = hub.fillOwned; +const fillOwnedLayers = hub.fillOwnedLayers; +const scheduleUsed = hub.scheduleUsed; +const scheduleUsedAt = hub.scheduleUsedAt; +const lineOf = hub.lineOf; +const lineIds = hub.lineIds; +const stagedCopies = hub.stagedCopies; + +fn benchPolicy(slot_n: u16, expert_bytes: u64, copy_bytes: u64) Policy { + return .{ + .slots = slot_n, + .expert_bytes = expert_bytes, + .copy_bytes = copy_bytes, + }; +} + +test "a high count keeps an old stamp and a one-off is the victim" { + var slots = [_]u16{ Cache.empty_id, Cache.empty_id }; + var line = [_]Line{ .empty, .empty }; + var stamp = [_]u32{ 0, 0 }; + var layer = [_]u8{ 0, 0 }; + var vram = Vram.init(&slots, &line, &stamp, &layer); + slots = .{ 1, 4 }; + line = .{ .resident, .resident }; + stamp = .{ 1, 100 }; + var counts: [shape.experts]u32 = @splat(0); + counts[1] = 50; + counts[4] = 1; + var decisions: [1]Decision = undefined; + var scratch: [2]u16 = undefined; + + const miss = [_]u16{9}; + try scheduleUsed(&vram, &miss, &counts, benchPolicy(2, 100, 100), 6, &decisions, &scratch, allOwned()); + try std.testing.expectEqual(Decision.copy, decisions[0]); + try std.testing.expectEqual(@as(u16, 1), slots[0]); + try std.testing.expectEqual(Line.resident, line[0]); + try std.testing.expectEqual(@as(u16, 9), slots[1]); + try std.testing.expectEqual(Line.inflight, line[1]); + try std.testing.expect(onDevice(vram, 1)); + try std.testing.expect(!onDevice(vram, 9)); +} + +test "zero lifetime counts fall back to stamp LRU" { + var slots = [_]u16{ Cache.empty_id, Cache.empty_id }; + var line = [_]Line{ .empty, .empty }; + var stamp = [_]u32{ 0, 0 }; + var layer = [_]u8{ 0, 0 }; + var vram = Vram.init(&slots, &line, &stamp, &layer); + slots = .{ 4, 7 }; + line = .{ .resident, .resident }; + stamp = .{ 1, 5 }; + var counts: [shape.experts]u32 = @splat(0); + var decisions: [1]Decision = undefined; + var scratch: [2]u16 = undefined; + + const miss = [_]u16{9}; + try scheduleUsed(&vram, &miss, &counts, benchPolicy(2, 100, 100), 6, &decisions, &scratch, allOwned()); + try std.testing.expectEqual(@as(u16, 9), slots[0]); + try std.testing.expectEqual(Line.inflight, line[0]); + try std.testing.expect(onDevice(vram, 7)); +} + +test "lifetime count beats this round's appearances" { + var slots = [_]u16{Cache.empty_id}; + var line = [_]Line{.empty}; + var stamp = [_]u32{0}; + var layer = [_]u8{0}; + var vram = Vram.init(&slots, &line, &stamp, &layer); + var counts: [shape.experts]u32 = @splat(0); + counts[8] = 5; + const ids = [_]u16{ 3, 3, 3, 3, 8 }; + var decisions: [ids.len]Decision = undefined; + var scratch: [2]u16 = undefined; + try scheduleUsed(&vram, &ids, &counts, benchPolicy(1, 100, 100), 1, &decisions, &scratch, allOwned()); + try std.testing.expectEqual(Decision.cpu, decisions[0]); + try std.testing.expectEqual(Decision.copy, decisions[4]); + try std.testing.expectEqual(@as(u16, 8), slots[0]); + try std.testing.expectEqual(Line.inflight, line[0]); +} + +test "a resident named this step is kept even when its count is low" { + var slots = [_]u16{ Cache.empty_id, Cache.empty_id }; + var line = [_]Line{ .empty, .empty }; + var stamp = [_]u32{ 0, 0 }; + var layer = [_]u8{ 0, 0 }; + var vram = Vram.init(&slots, &line, &stamp, &layer); + slots = .{ 1, 4 }; + line = .{ .resident, .resident }; + stamp = .{ 1, 100 }; + var counts: [shape.experts]u32 = @splat(0); + counts[4] = 1; + const ids = [_]u16{ 1, 4, 9 }; + var decisions: [3]Decision = undefined; + var scratch: [2]u16 = undefined; + try scheduleUsed(&vram, &ids, &counts, benchPolicy(2, 100, 100), 6, &decisions, &scratch, allOwned()); + try std.testing.expectEqual(Decision.hit, decisions[0]); + try std.testing.expectEqual(Decision.hit, decisions[1]); + try std.testing.expectEqual(Decision.cpu, decisions[2]); + try std.testing.expect(onDevice(vram, 1)); + try std.testing.expect(onDevice(vram, 4)); + try std.testing.expectEqual(@as(?usize, null), vramIndex(vram, 9)); +} + +test "a keep on layer 0 does not protect that id on layer 40" { + var slots = [_]u16{ 7, 7, 3 }; + var line = [_]Line{ .resident, .resident, .resident }; + var stamp = [_]u32{ 4, 4, 4 }; + var layer = [_]u8{ 0, 40, 0 }; + var vram = Vram.init(&slots, &line, &stamp, &layer); + slots = .{ 7, 7, 3 }; + line = .{ .resident, .resident, .resident }; + stamp = .{ 4, 4, 4 }; + layer = .{ 0, 40, 0 }; + + var counts: [shape.layers][shape.experts]u32 = @splat(@splat(0)); + counts[0][7] = 10; + counts[40][7] = 1; + counts[0][3] = 8; + const ids = [_]u16{ 7, 9 }; + var decisions: [2]Decision = undefined; + var scratch: [2]u16 = undefined; + try scheduleUsedAt(&vram, 0, &ids, &counts, benchPolicy(3, 100, 100), 8, &decisions, &scratch, allOwned()); + + try std.testing.expectEqual(Decision.hit, decisions[0]); + try std.testing.expectEqual(Decision.copy, decisions[1]); + try std.testing.expect(residentAt(vram, 0, 7)); + try std.testing.expect(residentAt(vram, 0, 3)); + try std.testing.expect(!residentAt(vram, 40, 7)); + try std.testing.expect(vramIndexAt(vram, 0, 9) != null); + try std.testing.expectEqual(@as(u32, 8), stamp[0]); + try std.testing.expectEqual(@as(u32, 4), stamp[2]); + try std.testing.expectEqual(@as(u16, 7), slots[0]); + try std.testing.expectEqual(@as(u16, 3), slots[2]); +} + +test "scheduleUsed is the layer-0 wrapper and is not the multi-layer path" { + var slots = [_]u16{ 1, 4 }; + var line = [_]Line{ .resident, .resident }; + var stamp = [_]u32{ 1, 50 }; + var layer = [_]u8{ 0, 40 }; + var vram = Vram.init(&slots, &line, &stamp, &layer); + slots = .{ 1, 4 }; + line = .{ .resident, .resident }; + stamp = .{ 1, 50 }; + layer = .{ 0, 40 }; + + // The flat vector says id 4 is hot. + var flat: [shape.experts]u32 = @splat(0); + flat[1] = 1; + flat[4] = 100; + var decisions: [1]Decision = undefined; + var scratch: [2]u16 = undefined; + const miss = [_]u16{9}; + try scheduleUsed(&vram, &miss, &flat, benchPolicy(2, 100, 100), 6, &decisions, &scratch, allOwned()); + try std.testing.expectEqual(Decision.copy, decisions[0]); + try std.testing.expectEqual(@as(u16, 1), slots[0]); + try std.testing.expectEqual(@as(u8, 0), layer[0]); + try std.testing.expectEqual(@as(u16, 9), slots[1]); + try std.testing.expectEqual(@as(u8, 0), layer[1]); + try std.testing.expect(!residentAt(vram, 40, 4)); + + slots = .{ 1, 4 }; + line = .{ .resident, .resident }; + stamp = .{ 1, 50 }; + layer = .{ 0, 40 }; + var counts: [shape.layers][shape.experts]u32 = @splat(@splat(0)); + counts[0][1] = 1; + counts[40][4] = 100; + try scheduleUsedAt(&vram, 0, &miss, &counts, benchPolicy(2, 100, 100), 7, &decisions, &scratch, allOwned()); + try std.testing.expectEqual(@as(u16, 9), slots[0]); + try std.testing.expectEqual(@as(u8, 0), layer[0]); + try std.testing.expectEqual(@as(u16, 4), slots[1]); + try std.testing.expectEqual(@as(u8, 40), layer[1]); + try std.testing.expect(residentAt(vram, 40, 4)); + try std.testing.expect(!residentAt(vram, 0, 1)); +} + +test "scheduleUsedAt does not stage another rank or the shared expert" { + var slots = [_]u16{ Cache.empty_id, Cache.empty_id }; + var line = [_]Line{ .empty, .empty }; + var stamp = [_]u32{ 0, 0 }; + var layer = [_]u8{ 0, 0 }; + var vram = Vram.init(&slots, &line, &stamp, &layer); + var counts: [shape.layers][shape.experts]u32 = @splat(@splat(0)); + counts[3][400] = 50; + counts[3][9] = 1; + const ids = [_]u16{ 400, shared_id, 9 }; + var decisions: [3]Decision = undefined; + var scratch: [4]u16 = undefined; + const rank0 = Owned{ .lo = 0, .hi = 256 }; + const policy = benchPolicy(2, 100, 10_000); + try std.testing.expectError(error.Length, scheduleUsedAt(&vram, @intCast(shape.layers), &ids, &counts, policy, 4, &decisions, &scratch, rank0)); + try std.testing.expectEqual(Line.empty, line[0]); + + try scheduleUsedAt(&vram, 3, &ids, &counts, policy, 4, &decisions, &scratch, rank0); + try std.testing.expectEqual(Decision.cpu, decisions[0]); + try std.testing.expectEqual(Decision.hit, decisions[1]); + try std.testing.expectEqual(Decision.copy, decisions[2]); + try std.testing.expectEqual(@as(?usize, null), vramIndex(vram, 400)); + try std.testing.expectEqual(@as(?usize, null), vramIndex(vram, shared_id)); + try std.testing.expect(!residentAt(vram, 3, 9)); + try std.testing.expect(vramIndexAt(vram, 3, 9) != null); + try std.testing.expectEqual(@as(u16, 9), slots[0]); + try std.testing.expectEqual(@as(u8, 3), layer[0]); + try std.testing.expectEqual(Line.empty, line[1]); + try std.testing.expect(slots[0] != shared_id and slots[1] != shared_id); +} + +test "lineIds maps hits to their line and everything else to -1, stagedCopies lists each staged line once" { + var slots: [3]u16 = undefined; + var line: [3]Line = undefined; + var stamp: [3]u32 = undefined; + var layer: [3]u8 = undefined; + var v = Vram.init(&slots, &line, &stamp, &layer); + var counts: [shape.layers][shape.experts]u32 = @splat(@splat(0)); + const pol = Policy{ .slots = 3, .expert_bytes = 10, .copy_bytes = 20 }; + var ids = [_]u16{ 7, 9, 9, 11, 512 }; + var dec: [5]Decision = undefined; + var scratch: [8]u16 = undefined; + try observe(&counts[2], &ids); + try scheduleUsedAt(&v, 2, &ids, &counts, pol, 1, &dec, &scratch, allOwned()); + // 3 unique misses (9 twice is hottest), budget for two copies: 9 and 7 copy, 11 cpu, 512 hit. + try std.testing.expectEqual(Decision.copy, dec[0]); + try std.testing.expectEqual(Decision.copy, dec[1]); + try std.testing.expectEqual(Decision.copy, dec[2]); + try std.testing.expectEqual(Decision.cpu, dec[3]); + try std.testing.expectEqual(Decision.hit, dec[4]); + var copies: [4]Copy = undefined; + const n = try stagedCopies(v, 2, &ids, &dec, &copies); + try std.testing.expectEqual(@as(usize, 2), n); + var ln: [5]i32 = undefined; + try lineIds(v, 2, &ids, &dec, &ln); + try std.testing.expectEqualSlices(i32, &.{ -1, -1, -1, -1, -1 }, &ln); + try std.testing.expectError(error.Length, stagedCopies(v, 2, &ids, &dec, copies[0..1])); + for (copies[0..n]) |c| try reapAt(&v, 2, c.id, 2); + try scheduleUsedAt(&v, 2, &ids, &counts, pol, 3, &dec, &scratch, allOwned()); + try lineIds(v, 2, &ids, &dec, &ln); + try std.testing.expectEqual(@as(i32, @intCast(lineOf(v, 2, 7).?)), ln[0]); + try std.testing.expectEqual(@as(i32, @intCast(lineOf(v, 2, 9).?)), ln[1]); + try std.testing.expectEqual(ln[1], ln[2]); + try std.testing.expect(ln[3] == -1 and ln[4] == -1); + try std.testing.expectEqual(@as(?usize, null), lineOf(v, 3, 7)); +} diff --git a/zig/src/families/flashnext/exl3/offload/exl3_offload_vram.zig b/zig/src/families/flashnext/exl3/offload/exl3_offload_vram.zig new file mode 100644 index 0000000000..e20bd2c4a0 --- /dev/null +++ b/zig/src/families/flashnext/exl3/offload/exl3_offload_vram.zig @@ -0,0 +1,216 @@ +//! Flash Next EXL3 offload: the VRAM line table, its LRU scheduling, and the line addressing the device pool reads. + +const std = @import("std"); +const pol = @import("exl3_offload_policy.zig"); + +const Cache = pol.Cache; +const Decision = pol.Decision; +const Policy = pol.Policy; +const shared_id = pol.shared_id; +const indexOf = pol.indexOf; +const pushUnique = pol.pushUnique; +const hotterFirst = pol.hotterFirst; +const copyCount = pol.copyCount; + +pub const Line = enum { empty, inflight, resident }; + +/// VRAM working set. Host RAM still holds every routed expert. `stamp` is the last reap or touch time. A higher stamp. +pub const Vram = struct { + slots: []u16, + line: []Line, + stamp: []u32, + /// Layer of `slots[i]`. Parallel to the other three slices. + layer: []u8, + + pub fn init(slots: []u16, line: []Line, stamp: []u32, layer: []u8) Vram { + std.debug.assert(slots.len == line.len and slots.len == stamp.len and slots.len == layer.len); + @memset(slots, Cache.empty_id); + @memset(line, .empty); + @memset(stamp, 0); + @memset(layer, 0); + return .{ .slots = slots, .line = line, .stamp = stamp, .layer = layer }; + } +}; + +/// True once `id` has been reaped into VRAM. Id-only: the first matching line wins. +pub fn onDevice(v: Vram, id: u16) bool { + const i = vramIndex(v, id) orelse return false; + return v.line[i] == .resident; +} + +/// True when this layer's copy of `id` is resident. Another layer's copy does not count. +pub fn residentAt(v: Vram, layer: u8, id: u16) bool { + const i = vramIndexAt(v, layer, id) orelse return false; + return v.line[i] == .resident; +} + +/// A resident hit. Moves that line to `now` so LRU keeps it. +pub fn touch(v: *Vram, id: u16, now: u32) bool { + const i = vramIndex(v.*, id) orelse return false; + if (v.line[i] != .resident) return false; + v.stamp[i] = now; + return true; +} + +/// Start a host-to-VRAM copy of `id`. +pub fn stage(v: *Vram, id: u16) error{ Busy, Shared }!usize { + return stageKeeping(v, id, &.{}); +} + +/// Same as `stage`, except a resident whose id is in `keep` is not a victim. +pub fn stageKeeping(v: *Vram, id: u16, keep: []const u16) error{ Busy, Shared }!usize { + if (id == shared_id or id == Cache.empty_id) return error.Shared; + if (vramIndex(v.*, id)) |i| { + if (v.line[i] == .inflight) return i; + v.line[i] = .inflight; + return i; + } + const slot = vramEmpty(v.*) orelse vramLruKeep(v.*, keep) orelse return error.Busy; + v.slots[slot] = id; + v.layer[slot] = 0; + v.line[slot] = .inflight; + v.stamp[slot] = 0; + return slot; +} + +/// One step of the VRAM cache. +pub fn schedule( + v: *Vram, + ids: []const u16, + policy: Policy, + now: u32, + decisions: []Decision, + scratch: []u16, +) error{ Length, Scratch }!void { + if (decisions.len != ids.len) return error.Length; + + var n: usize = 0; + for (ids, decisions) |id, *decision| { + if (id == shared_id) { + decision.* = .hit; + continue; + } + if (onDevice(v.*, id)) { + _ = touch(v, id, now); + decision.* = .hit; + continue; + } + decision.* = .cpu; + if (vramIndex(v.*, id) != null) continue; + try pushUnique(scratch, &n, id); + } + std.mem.sort(u16, scratch[0..n], ids, hotterFirst); + + const copies = copyCount(policy, n); + var staged: usize = 0; + for (scratch[0..n]) |id| { + if (staged >= copies) break; + _ = stageKeeping(v, id, ids) catch break; + staged += 1; + for (ids, decisions) |asked, *decision| { + if (asked == id) decision.* = .copy; + } + } +} + +/// The copy of `(layer, id)` has landed. +pub fn reapAt(v: *Vram, layer: u8, id: u16, now: u32) error{NotInflight}!void { + const i = vramIndexAt(v.*, layer, id) orelse return error.NotInflight; + if (v.line[i] != .inflight) return error.NotInflight; + v.line[i] = .resident; + v.stamp[i] = now; +} + +/// Layer-0 wrapper. Does not promote another layer's inflight copy. Old tests are layer 0. Generate calls `reapAt`. +pub fn reap(v: *Vram, id: u16, now: u32) error{NotInflight}!void { + return reapAt(v, 0, id, now); +} + +/// First occupied slot with this id, on any layer. +pub fn vramIndex(v: Vram, id: u16) ?usize { + if (id == Cache.empty_id) return null; + for (v.slots, 0..) |slot, i| { + if (slot == id and v.line[i] != .empty) return i; + } + return null; +} + +/// Slot whose layer and id both match. The id-only search aliases two layers. +pub fn vramIndexAt(v: Vram, layer: u8, id: u16) ?usize { + if (id == Cache.empty_id) return null; + std.debug.assert(v.layer.len == v.slots.len and v.line.len == v.slots.len); + for (v.slots, v.layer, 0..) |slot, lay, i| { + if (slot == id and lay == layer and v.line[i] != .empty) return i; + } + return null; +} + +pub fn touchAt(v: *Vram, layer: u8, id: u16, now: u32) bool { + const i = vramIndexAt(v.*, layer, id) orelse return false; + if (v.line[i] != .resident) return false; + v.stamp[i] = now; + return true; +} + +pub fn vramEmpty(v: Vram) ?usize { + for (v.line, 0..) |state, i| if (state == .empty) return i; + return null; +} + +/// Oldest resident stamp. Equal stamps take the higher slot index. +fn vramLru(v: Vram) ?usize { + return vramLruKeep(v, &.{}); +} + +/// Oldest resident whose id is not in `keep`. Equal stamps take the higher slot index. +fn vramLruKeep(v: Vram, keep: []const u16) ?usize { + var best: ?usize = null; + for (v.line, 0..) |state, i| { + if (state != .resident) continue; + if (indexOf(keep, v.slots[i]) != null) continue; + if (best) |b| { + if (v.stamp[i] < v.stamp[b] or (v.stamp[i] == v.stamp[b] and i > b)) best = i; + } else best = i; + } + return best; +} + +/// Line index of `(layer, id)` in any state, or null. The device pool addresses its buffers by this index. +pub fn lineOf(v: Vram, layer: u8, id: u16) ?usize { + return vramIndexAt(v, layer, id); +} + +/// A host-to-device copy the policy asked for: expert `id` into pool line `line`. +pub const Copy = struct { id: u16, line: u32 }; + +/// Per pick, the pool line the hit kernel reads, or -1 when the kernel must skip the slot (a `.copy` or `.cpu` +pub fn lineIds(v: Vram, layer: u8, ids: []const u16, decisions: []const Decision, out: []i32) error{Length}!void { + if (decisions.len != ids.len or out.len != ids.len) return error.Length; + for (ids, decisions, out) |id, d, *o| { + o.* = -1; + if (d != .hit or id == shared_id) continue; + const i = vramIndexAt(v, layer, id) orelse continue; + if (v.line[i] != .resident) continue; + o.* = @intCast(i); + } +} + +/// Unique `(id, line)` pairs the round staged: a `.copy` decision whose line is in flight. +pub fn stagedCopies(v: Vram, layer: u8, ids: []const u16, decisions: []const Decision, out: []Copy) error{Length}!usize { + if (decisions.len != ids.len) return error.Length; + var n: usize = 0; + for (ids, decisions) |id, d| { + if (d != .copy) continue; + const i = vramIndexAt(v, layer, id) orelse continue; + if (v.line[i] != .inflight) continue; + var seen = false; + for (out[0..n]) |c| { + if (c.id == id) seen = true; + } + if (seen) continue; + if (n >= out.len) return error.Length; + out[n] = .{ .id = id, .line = @intCast(i) }; + n += 1; + } + return n; +} diff --git a/zig/src/families/flashnext/exl3/offload/exl3_offload_vram_test.zig b/zig/src/families/flashnext/exl3/offload/exl3_offload_vram_test.zig new file mode 100644 index 0000000000..8a1d4f8f87 --- /dev/null +++ b/zig/src/families/flashnext/exl3/offload/exl3_offload_vram_test.zig @@ -0,0 +1,246 @@ +//! Flash Next EXL3 offload tests: the VRAM line states, LRU and layer tagging. + +const std = @import("std"); +const hub = @import("exl3_offload.zig"); +const vram_ = @import("exl3_offload_vram.zig"); + +const shape = hub.shape; +const shared_id = hub.shared_id; +const Cache = hub.Cache; +const Decision = hub.Decision; +const Policy = hub.Policy; +const Line = hub.Line; +const Vram = hub.Vram; +const Owned = hub.Owned; +const Copy = hub.Copy; +const init = hub.init; +const find = hub.find; +const planRound = hub.planRound; +const swap = hub.swap; +const installProfile = hub.installProfile; +const onDevice = hub.onDevice; +const residentAt = hub.residentAt; +const touch = hub.touch; +const stage = hub.stage; +const stageKeeping = hub.stageKeeping; +const schedule = hub.schedule; +const reap = hub.reap; +const reapAt = hub.reapAt; +const vramIndexAt = hub.vramIndexAt; +const vramIndex = vram_.vramIndex; +const observe = hub.observe; +const allOwned = hub.allOwned; +const seedHottest = hub.seedHottest; +const fillOwned = hub.fillOwned; +const fillOwnedLayers = hub.fillOwnedLayers; +const scheduleUsed = hub.scheduleUsed; +const scheduleUsedAt = hub.scheduleUsedAt; +const lineOf = hub.lineOf; +const lineIds = hub.lineIds; +const stagedCopies = hub.stagedCopies; + +fn benchPolicy(slot_n: u16, expert_bytes: u64, copy_bytes: u64) Policy { + return .{ + .slots = slot_n, + .expert_bytes = expert_bytes, + .copy_bytes = copy_bytes, + }; +} + +test "a staged copy is not on device until reap, and LRU drops the older resident" { + var slots = [_]u16{ Cache.empty_id, Cache.empty_id }; + var line = [_]Line{ .empty, .empty }; + var stamp = [_]u32{ 0, 0 }; + var layer = [_]u8{ 0, 0 }; + var vram = Vram.init(&slots, &line, &stamp, &layer); + + try std.testing.expectEqual(@as(usize, 0), try stage(&vram, 4)); + try std.testing.expectEqual(@as(usize, 1), try stage(&vram, 7)); + try std.testing.expect(!onDevice(vram, 4)); + try std.testing.expectError(error.NotInflight, reap(&vram, 9, 1)); + try reap(&vram, 4, 1); + try reap(&vram, 7, 2); + try std.testing.expect(onDevice(vram, 4)); + try std.testing.expect(touch(&vram, 7, 5)); + try std.testing.expect(onDevice(vram, 4)); + + // 4 was reaped at 1 and not touched again. 7 was touched at 5. 4 is the victim. + try std.testing.expectEqual(@as(usize, 0), try stage(&vram, 8)); + try std.testing.expect(!onDevice(vram, 8)); + try std.testing.expectEqual(Line.inflight, line[0]); + try std.testing.expectEqual(@as(u16, 7), slots[1]); + try reap(&vram, 8, 6); + try std.testing.expect(onDevice(vram, 8)); + try std.testing.expectEqual(@as(?usize, null), vramIndex(vram, 4)); +} + +test "an in-flight line is not reused, and a full in-flight set is busy" { + var slots = [_]u16{ Cache.empty_id, Cache.empty_id }; + var line = [_]Line{ .empty, .empty }; + var stamp = [_]u32{ 0, 0 }; + var layer = [_]u8{ 0, 0 }; + var vram = Vram.init(&slots, &line, &stamp, &layer); + _ = try stage(&vram, 1); + _ = try stage(&vram, 2); + try std.testing.expectError(error.Busy, stage(&vram, 3)); + try std.testing.expectEqual(@as(u16, 1), slots[0]); + try std.testing.expectEqual(@as(u16, 2), slots[1]); + try std.testing.expectError(error.Shared, stage(&vram, shared_id)); +} + +test "schedule touches residents, stages the hottest miss, and will not evict a hit" { + var slots = [_]u16{ Cache.empty_id, Cache.empty_id }; + var line = [_]Line{ .empty, .empty }; + var stamp = [_]u32{ 0, 0 }; + var layer = [_]u8{ 0, 0 }; + var vram = Vram.init(&slots, &line, &stamp, &layer); + var decisions: [4]Decision = undefined; + var scratch: [4]u16 = undefined; + const policy = benchPolicy(2, 100, 100); + + const first = [_]u16{ 7, shared_id, 4, 7 }; + try schedule(&vram, &first, policy, 3, &decisions, &scratch); + try std.testing.expectEqual(Decision.hit, decisions[1]); + try std.testing.expectEqual(Decision.copy, decisions[0]); + try std.testing.expectEqual(Decision.copy, decisions[3]); + try std.testing.expectEqual(Decision.cpu, decisions[2]); + try std.testing.expect(!onDevice(vram, 7)); + try std.testing.expectEqual(Line.inflight, line[0]); + try std.testing.expectEqual(@as(u16, 7), slots[0]); + try std.testing.expectEqual(Line.empty, line[1]); + + try reap(&vram, 7, 4); + const held = [_]u16{7}; + try schedule(&vram, &held, policy, 9, decisions[0..1], &scratch); + try std.testing.expectEqual(Decision.hit, decisions[0]); + try std.testing.expectEqual(@as(u32, 9), stamp[0]); + + // Both lines are residents this step asked for, so the miss cannot take either. + _ = try stage(&vram, 3); + try reap(&vram, 3, 10); + const protected = [_]u16{ 7, 3, 8 }; + try schedule(&vram, &protected, benchPolicy(2, 100, 10_000), 11, decisions[0..3], &scratch); + try std.testing.expectEqual(Decision.hit, decisions[0]); + try std.testing.expectEqual(Decision.hit, decisions[1]); + try std.testing.expectEqual(Decision.cpu, decisions[2]); + try std.testing.expect(onDevice(vram, 7)); + try std.testing.expect(onDevice(vram, 3)); + try std.testing.expectEqual(@as(?usize, null), vramIndex(vram, 8)); +} + +test "a zero copy budget stages nothing, and an older resident is the one that moves" { + var slots = [_]u16{ 4, 7 }; + var line = [_]Line{ .resident, .resident }; + var stamp = [_]u32{ 1, 5 }; + var layer = [_]u8{ 0, 0 }; + var vram = Vram.init(&slots, &line, &stamp, &layer); + slots = .{ 4, 7 }; + line = .{ .resident, .resident }; + stamp = .{ 1, 5 }; + var decisions: [1]Decision = undefined; + var scratch: [2]u16 = undefined; + + const miss = [_]u16{9}; + try schedule(&vram, &miss, benchPolicy(2, 0, 10_000), 6, &decisions, &scratch); + try std.testing.expectEqual(Decision.cpu, decisions[0]); + try std.testing.expectEqual(@as(u16, 4), slots[0]); + try std.testing.expectEqual(Line.resident, line[0]); + + try schedule(&vram, &miss, benchPolicy(2, 100, 100), 6, &decisions, &scratch); + try std.testing.expectEqual(Decision.copy, decisions[0]); + try std.testing.expectEqual(@as(u16, 9), slots[0]); + try std.testing.expectEqual(Line.inflight, line[0]); + try std.testing.expect(onDevice(vram, 7)); + try std.testing.expect(!onDevice(vram, 9)); +} + +test "residentAt distinguishes the same expert id on two layers" { + var slots = [_]u16{ 7, 7 }; + var line = [_]Line{ .resident, .resident }; + var stamp = [_]u32{ 1, 2 }; + var layer = [_]u8{ 0, 40 }; + const vram = Vram.init(&slots, &line, &stamp, &layer); + slots = .{ 7, 7 }; + line = .{ .resident, .resident }; + stamp = .{ 1, 2 }; + layer = .{ 0, 40 }; + + try std.testing.expect(residentAt(vram, 0, 7)); + try std.testing.expect(residentAt(vram, 40, 7)); + try std.testing.expect(!residentAt(vram, 1, 7)); + try std.testing.expect(!residentAt(vram, 0, 8)); + // Id-only lookup still aliases. Generate does not call it on this table. + try std.testing.expect(onDevice(vram, 7)); + + line[1] = .inflight; + try std.testing.expect(residentAt(vram, 0, 7)); + try std.testing.expect(!residentAt(vram, 40, 7)); +} + +test "reapAt promotes one inflight layer and leaves the other inflight" { + var slots = [_]u16{ 7, 7 }; + var line = [_]Line{ .inflight, .inflight }; + var stamp = [_]u32{ 0, 0 }; + var layer = [_]u8{ 0, 40 }; + var vram = Vram.init(&slots, &line, &stamp, &layer); + slots = .{ 7, 7 }; + line = .{ .inflight, .inflight }; + stamp = .{ 0, 0 }; + layer = .{ 0, 40 }; + + try reapAt(&vram, 0, 7, 3); + try std.testing.expect(residentAt(vram, 0, 7)); + try std.testing.expectEqual(Line.inflight, line[1]); + try std.testing.expect(!residentAt(vram, 40, 7)); + try std.testing.expectEqual(@as(u32, 3), stamp[0]); + + // The other layer's copy is still inflight, then promotable on its own reap. + try reapAt(&vram, 40, 7, 4); + try std.testing.expect(residentAt(vram, 40, 7)); + try std.testing.expect(residentAt(vram, 0, 7)); + try std.testing.expectEqual(@as(u32, 4), stamp[1]); + try std.testing.expectEqual(@as(u32, 3), stamp[0]); + + // Layer 0 is already resident. + line[1] = .inflight; + stamp[1] = 0; + try std.testing.expectError(error.NotInflight, reap(&vram, 7, 9)); + try std.testing.expectEqual(Line.resident, line[0]); + try std.testing.expectEqual(Line.inflight, line[1]); + try std.testing.expectEqual(@as(u32, 3), stamp[0]); + try reapAt(&vram, 40, 7, 9); + try std.testing.expect(residentAt(vram, 40, 7)); + try std.testing.expectEqual(@as(u32, 9), stamp[1]); + + // An inflight copy that exists only on layer 40 is not promoted by `reap`. + slots = .{ 7, Cache.empty_id }; + line = .{ .inflight, .empty }; + stamp = .{ 0, 0 }; + layer = .{ 40, 0 }; + try std.testing.expectError(error.NotInflight, reap(&vram, 7, 1)); + try std.testing.expectEqual(Line.inflight, line[0]); + try std.testing.expect(!residentAt(vram, 40, 7)); + try reapAt(&vram, 40, 7, 2); + try std.testing.expect(residentAt(vram, 40, 7)); + try std.testing.expectEqual(@as(u8, 40), layer[0]); +} + +test "stageKeeping tags a newly assigned line as layer 0" { + var slots = [_]u16{4}; + var line = [_]Line{.resident}; + var stamp = [_]u32{1}; + var layer = [_]u8{40}; + var vram = Vram.init(&slots, &line, &stamp, &layer); + slots = .{4}; + line = .{.resident}; + stamp = .{1}; + layer = .{40}; + + try std.testing.expectEqual(@as(usize, 0), try stageKeeping(&vram, 9, &.{})); + try std.testing.expectEqual(@as(u16, 9), slots[0]); + try std.testing.expectEqual(@as(u8, 0), layer[0]); + try std.testing.expectEqual(Line.inflight, line[0]); + try std.testing.expect(vramIndexAt(vram, 0, 9) != null); + try std.testing.expectEqual(@as(?usize, null), vramIndexAt(vram, 40, 9)); + try std.testing.expect(slots[0] != shared_id); +} diff --git a/zig/src/families/flashnext/exl3/offload/exl3_profile_line_budget.zig b/zig/src/families/flashnext/exl3/offload/exl3_profile_line_budget.zig new file mode 100644 index 0000000000..16ba9055e9 --- /dev/null +++ b/zig/src/families/flashnext/exl3/offload/exl3_profile_line_budget.zig @@ -0,0 +1,468 @@ +const std = @import("std"); + +pub const heuristic = "Greedy marginal histogram coverage per line byte, not real miss-cost optimization. Equal ratios prefer fewer lines, then the lower decoder ordinal. Exact line completion uses bounded search, potentially exponential with many distinct costs. Candidates must also preserve exact device picks bytes; otherwise sizing stays uniform."; +pub const experts: usize = 512; + +pub const Request = struct { + counts: []const []const u32, + rank: usize, + world: usize, + line_bytes: []const usize, + uniform_lines: usize, + min_lines: usize, + max_rows: usize, + slots: usize, +}; + +pub const Budget = struct { + baseline_bytes: u64, + allocated_bytes: u64, + uniform_fallback: bool = false, +}; + +pub const Error = std.mem.Allocator.Error || error{ InvalidDimensions, InvalidRank, InvalidLineCost, InvalidBounds, Overflow, UnfillableBudget }; + +fn poolBytes(line_bytes: usize, lines: usize, max_rows: usize, slots: usize, min_lines: usize) Error!u64 { + if (lines == 0) return error.InvalidBounds; + if (min_lines == 0) return error.InvalidBounds; + if (max_rows == 0 or slots == 0) return error.InvalidDimensions; + const picks = try std.math.mul(usize, max_rows, slots); + const max_passes = 2 + experts / min_lines + @intFromBool(experts % min_lines != 0); + const passes = @min(max_passes, @min(picks, experts) + 1); + const picks_bytes = try std.math.mul(usize, try std.math.mul(usize, passes, picks), @sizeOf(i32)); + return std.math.add(u64, try std.math.mul(u64, lines, line_bytes), picks_bytes); +} + +const Group = struct { + bytes: u64, + available: usize, + + fn descending(_: void, left: Group, right: Group) bool { + return left.bytes > right.bytes; + } +}; + +pub fn allocate(gpa: std.mem.Allocator, request: Request, output: []usize) Error!Budget { + if ((request.world != 1 and request.world != 2) or request.rank >= request.world) return error.InvalidRank; + const local_experts = experts / request.world; + if (request.min_lines == 0 or request.min_lines > request.uniform_lines or request.uniform_lines > local_experts) return error.InvalidBounds; + if (request.counts.len == 0 or request.counts.len != request.line_bytes.len or output.len != request.counts.len) return error.InvalidDimensions; + var baseline_bytes: u64 = 0; + var line_budget_bytes: u64 = 0; + var floor_bytes: u64 = 0; + for (request.counts, request.line_bytes) |histogram, bytes| { + if (histogram.len != experts) return error.InvalidDimensions; + if (bytes == 0) return error.InvalidLineCost; + baseline_bytes = try std.math.add(u64, baseline_bytes, try poolBytes(bytes, request.uniform_lines, request.max_rows, request.slots, request.min_lines)); + line_budget_bytes = try std.math.add(u64, line_budget_bytes, try std.math.mul(u64, request.uniform_lines, bytes)); + floor_bytes = try std.math.add(u64, floor_bytes, try std.math.mul(u64, request.min_lines, bytes)); + } + const ranked = try gpa.alloc([experts]u32, request.counts.len); + defer gpa.free(ranked); + const lines = try gpa.alloc(usize, request.counts.len); + defer gpa.free(lines); + const storage = try gpa.alloc(Group, request.counts.len); + defer gpa.free(storage); + const layer_groups = try gpa.alloc(usize, request.counts.len); + defer gpa.free(layer_groups); + const attempted = try gpa.alloc(bool, request.counts.len); + defer gpa.free(attempted); + const suffix_length = try std.math.add(usize, request.counts.len, 1); + const suffix_bytes = try gpa.alloc(u64, suffix_length); + defer gpa.free(suffix_bytes); + const suffix_gcd = try gpa.alloc(u64, suffix_length); + defer gpa.free(suffix_gcd); + const first = request.rank * local_experts; + for (request.counts, ranked) |histogram, *ordered| { + @memcpy(ordered[0..local_experts], histogram[first..][0..local_experts]); + std.mem.sort(u32, ordered[0..local_experts], {}, std.sort.desc(u32)); + } + @memset(lines, request.min_lines); + var group_count: usize = 0; + for (request.line_bytes) |bytes| { + var found: ?usize = null; + for (storage[0..group_count], 0..) |group, index| { + if (group.bytes == bytes) { + found = index; + break; + } + } + const group_index = found orelse new: { + storage[group_count] = .{ .bytes = bytes, .available = 0 }; + group_count += 1; + break :new group_count - 1; + }; + storage[group_index].available = try std.math.add(usize, storage[group_index].available, local_experts - request.min_lines); + } + const groups = storage[0..group_count]; + std.mem.sort(Group, groups, {}, Group.descending); + for (request.line_bytes, layer_groups) |bytes, *group_index| { + for (groups, 0..) |group, index| { + if (group.bytes == bytes) { + group_index.* = index; + break; + } + } + } + var remaining = line_budget_bytes - floor_bytes; + while (remaining != 0) { + @memset(attempted, false); + var selected: ?usize = null; + for (0..lines.len) |_| { + var best: ?usize = null; + for (lines, request.line_bytes, 0..) |count, bytes, layer| { + if (attempted[layer] or count == local_experts or bytes > remaining) continue; + if (best == null or better(layer, best.?, ranked, lines, request.line_bytes)) best = layer; + } + const candidate = best orelse break; + attempted[candidate] = true; + const group_index = layer_groups[candidate]; + groups[group_index].available -= 1; + const possible = canFill(groups, remaining - request.line_bytes[candidate], suffix_bytes, suffix_gcd); + groups[group_index].available += 1; + if (possible) { + selected = candidate; + break; + } + } + const layer = selected orelse return error.UnfillableBudget; + lines[layer] += 1; + groups[layer_groups[layer]].available -= 1; + remaining -= request.line_bytes[layer]; + } + var allocated_bytes: u64 = 0; + for (lines, request.line_bytes) |count, bytes| allocated_bytes = try std.math.add(u64, allocated_bytes, try poolBytes(bytes, count, request.max_rows, request.slots, request.min_lines)); + const uniform_fallback = allocated_bytes != baseline_bytes; + if (uniform_fallback) { + @memset(lines, request.uniform_lines); + allocated_bytes = baseline_bytes; + } + @memcpy(output, lines); + return .{ .baseline_bytes = baseline_bytes, .allocated_bytes = allocated_bytes, .uniform_fallback = uniform_fallback }; +} + +fn better(left: usize, right: usize, ranked: []const [experts]u32, lines: []const usize, bytes: []const usize) bool { + const left_score = @as(u128, ranked[left][lines[left]]) * bytes[right]; + const right_score = @as(u128, ranked[right][lines[right]]) * bytes[left]; + if (left_score != right_score) return left_score > right_score; + if (lines[left] != lines[right]) return lines[left] < lines[right]; + return left < right; +} + +fn gcd(left: u64, right: u64) u64 { + var numerator = left; + var denominator = right; + while (denominator != 0) { + const remainder = numerator % denominator; + numerator = denominator; + denominator = remainder; + } + return numerator; +} + +fn canFill(groups: []const Group, target: u64, suffix_bytes: []u64, suffix_gcd: []u64) bool { + suffix_bytes[groups.len] = 0; + suffix_gcd[groups.len] = 0; + var index = groups.len; + while (index != 0) { + index -= 1; + const group = groups[index]; + const capacity = @as(u128, group.bytes) * group.available + suffix_bytes[index + 1]; + suffix_bytes[index] = @intCast(@min(capacity, target)); + suffix_gcd[index] = if (group.available == 0) suffix_gcd[index + 1] else gcd(group.bytes, suffix_gcd[index + 1]); + } + return fill(groups, 0, target, suffix_bytes, suffix_gcd); +} + +fn fill(groups: []const Group, index: usize, target: u64, suffix_bytes: []const u64, suffix_gcd: []const u64) bool { + if (target == 0) return true; + if (index == groups.len or target > suffix_bytes[index] or suffix_gcd[index] == 0 or target % suffix_gcd[index] != 0) return false; + const group = groups[index]; + if (group.available == 0) return fill(groups, index + 1, target, suffix_bytes, suffix_gcd); + const max_take = @min(group.available, target / group.bytes); + const required = target - @min(target, suffix_bytes[index + 1]); + const min_take = required / group.bytes + @intFromBool(required % group.bytes != 0); + if (min_take > max_take) return false; + var take = max_take; + while (true) { + if (fill(groups, index + 1, target - take * group.bytes, suffix_bytes, suffix_gcd)) return true; + if (take == min_take) break; + take -= 1; + } + return false; +} + +fn testRequest(counts: []const []const u32, bytes: []const usize, uniform: usize, minimum: usize) Request { + return .{ .counts = counts, .rank = 0, .world = 2, .line_bytes = bytes, .uniform_lines = uniform, .min_lines = minimum, .max_rows = 1, .slots = 1 }; +} + +test "profile budget: sorted marginal counts favor hot layers with exact bytes" { + var hot: [experts]u32 = @splat(0); + const cold: [experts]u32 = @splat(1); + @memset(hot[0..6], 100); + var output: [2]usize = undefined; + const plan = try allocate(std.testing.allocator, testRequest(&.{ &hot, &cold }, &.{ 4, 6 }, 3, 1), &output); + try std.testing.expectEqualSlices(usize, &.{ 6, 1 }, &output); + try std.testing.expectEqual(@as(u64, 30 + 2 * 8), plan.baseline_bytes); + try std.testing.expectEqual(plan.baseline_bytes, plan.allocated_bytes); + hot[0] = 0; + hot[255] = 100; + _ = try allocate(std.testing.allocator, testRequest(&.{ &hot, &cold }, &.{ 4, 6 }, 3, 1), &output); + try std.testing.expectEqualSlices(usize, &.{ 6, 1 }, &output); +} + +test "profile budget: rank filtering handles asymmetric counts and TP1" { + var first: [experts]u32 = @splat(0); + var second: [experts]u32 = @splat(0); + @memset(first[0..4], 10); + @memset(second[257..260], 20); + second[511] = 20; + var output: [2]usize = undefined; + var request = testRequest(&.{ &first, &second }, &.{ 8, 8 }, 2, 1); + _ = try allocate(std.testing.allocator, request, &output); + try std.testing.expectEqualSlices(usize, &.{ 3, 1 }, &output); + request.rank = 1; + _ = try allocate(std.testing.allocator, request, &output); + try std.testing.expectEqualSlices(usize, &.{ 1, 3 }, &output); + request.rank = 0; + request.world = 1; + _ = try allocate(std.testing.allocator, request, &output); + try std.testing.expectEqualSlices(usize, &.{ 1, 3 }, &output); +} + +test "profile budget: zero and flat histograms retain uniform sizing deterministically" { + for ([_]u32{ 0, 7, std.math.maxInt(u32) }) |value| { + const histogram: [experts]u32 = @splat(value); + var output: [3]usize = undefined; + for (0..3) |_| { + _ = try allocate(std.testing.allocator, testRequest(&.{ &histogram, &histogram, &histogram }, &.{ 16, 16, 16 }, 17, 10), &output); + try std.testing.expectEqualSlices(usize, &.{ 17, 17, 17 }, &output); + } + } +} + +test "profile budget: tied ratios prefer fewer lines then lower decoder ordinal" { + var first: [experts]u32 = @splat(0); + var second: [experts]u32 = @splat(0); + const third: [experts]u32 = @splat(0); + @memset(first[0..3], 10); + @memset(second[0..3], 10); + var output: [3]usize = undefined; + _ = try allocate(std.testing.allocator, testRequest(&.{ &first, &second, &third }, &.{ 8, 8, 8 }, 2, 1), &output); + try std.testing.expectEqualSlices(usize, &.{ 3, 2, 1 }, &output); +} + +test "profile budget: ratios use wide products and avoid greedy dead ends" { + const hot: [experts]u32 = @splat(std.math.maxInt(u32) - 1); + const cold: [experts]u32 = @splat(std.math.maxInt(u32)); + const unit: usize = @as(usize, 1) << 32; + var output: [2]usize = undefined; + const plan = try allocate(std.testing.allocator, testRequest(&.{ &hot, &cold }, &.{ 2 * unit, 3 * unit }, 3, 1), &output); + try std.testing.expectEqualSlices(usize, &.{ 6, 1 }, &output); + try std.testing.expectEqual(@as(u64, 15) * unit + 2 * 8, plan.allocated_bytes); + const low: [experts]u32 = @splat(1); + _ = try allocate(std.testing.allocator, testRequest(&.{ &hot, &low }, &.{ 4, 6 }, 2, 1), &output); + try std.testing.expectEqualSlices(usize, &.{ 2, 2 }, &output); +} + +test "profile budget: floors caps and single layer budgets" { + const hot: [experts]u32 = @splat(100); + const cold: [experts]u32 = @splat(0); + var output: [2]usize = undefined; + _ = try allocate(std.testing.allocator, testRequest(&.{ &hot, &cold }, &.{ 8, 8 }, 200, 10), &output); + try std.testing.expectEqualSlices(usize, &.{ 256, 144 }, &output); + var request = testRequest(&.{ &hot, &cold }, &.{ 8, 8 }, 400, 10); + request.world = 1; + _ = try allocate(std.testing.allocator, request, &output); + try std.testing.expectEqualSlices(usize, &.{ 512, 288 }, &output); + _ = try allocate(std.testing.allocator, testRequest(&.{ &hot, &cold }, &.{ 7, 11 }, 10, 10), &output); + try std.testing.expectEqualSlices(usize, &.{ 10, 10 }, &output); + _ = try allocate(std.testing.allocator, testRequest(&.{&hot}, &.{13}, 47, 10), output[0..1]); + try std.testing.expectEqual(@as(usize, 47), output[0]); +} + +test "profile budget: malformed inputs and overflow leave output unchanged" { + const histogram: [experts]u32 = @splat(0); + const shared: [513]u32 = @splat(0); + var output = [_]usize{ 999, 999 }; + const valid = testRequest(&.{&histogram}, &.{8}, 20, 10); + for ([_][]const []const u32{ &.{}, &.{histogram[0..511]}, &.{&shared} }) |counts| { + var request = valid; + request.counts = counts; + try std.testing.expectError(error.InvalidDimensions, allocate(std.testing.allocator, request, output[0..1])); + } + var request = valid; + request.line_bytes = &.{}; + try std.testing.expectError(error.InvalidDimensions, allocate(std.testing.allocator, request, output[0..1])); + try std.testing.expectError(error.InvalidDimensions, allocate(std.testing.allocator, valid, &output)); + request = valid; + request.line_bytes = &.{0}; + try std.testing.expectError(error.InvalidLineCost, allocate(std.testing.allocator, request, output[0..1])); + for ([_][2]usize{ .{ 0, 0 }, .{ 3, 0 }, .{ 2, 2 }, .{ 1, 1 }, .{ 2, std.math.maxInt(usize) } }) |identity| { + request = valid; + request.world = identity[0]; + request.rank = identity[1]; + try std.testing.expectError(error.InvalidRank, allocate(std.testing.allocator, request, output[0..1])); + } + for ([_][2]usize{ .{ 20, 0 }, .{ 9, 10 }, .{ 257, 10 }, .{ 0, 0 } }) |bounds| { + request = valid; + request.uniform_lines = bounds[0]; + request.min_lines = bounds[1]; + try std.testing.expectError(error.InvalidBounds, allocate(std.testing.allocator, request, output[0..1])); + } + try std.testing.expectError(error.Overflow, allocate(std.testing.allocator, testRequest(&.{&histogram}, &.{std.math.maxInt(usize)}, 2, 1), output[0..1])); + const bytes = std.math.maxInt(usize) / 2 + 1; + try std.testing.expectError(error.Overflow, allocate(std.testing.allocator, testRequest(&.{ &histogram, &histogram }, &.{ bytes, bytes }, 1, 1), &output)); + try std.testing.expectEqualSlices(usize, &.{ 999, 999 }, &output); +} + +fn allocationFailureProbe(gpa: std.mem.Allocator) !void { + const histogram: [experts]u32 = @splat(7); + var output = [_]usize{ 999, 999 }; + const result = allocate(gpa, testRequest(&.{ &histogram, &histogram }, &.{ 4, 6 }, 3, 1), &output) catch |err| { + try std.testing.expectEqualSlices(usize, &.{ 999, 999 }, &output); + return err; + }; + try std.testing.expectEqual(result.baseline_bytes, result.allocated_bytes); +} + +test "profile budget: allocation failures release scratch and preserve output" { + try std.testing.checkAllAllocationFailures(std.testing.allocator, allocationFailureProbe, .{}); +} + +test "profile budget: feasibility matches exhaustive bounded enumeration" { + const groups = [_]Group{ .{ .bytes = 11, .available = 1 }, .{ .bytes = 7, .available = 2 }, .{ .bytes = 5, .available = 3 }, .{ .bytes = 3, .available = 4 } }; + var suffix_bytes: [5]u64 = undefined; + var suffix_gcd: [5]u64 = undefined; + for (0..65) |target| { + var expected = false; + for (0..2) |first| { + for (0..3) |second| { + for (0..4) |third| { + for (0..5) |fourth| { + if (11 * first + 7 * second + 5 * third + 3 * fourth == target) expected = true; + } + } + } + } + try std.testing.expectEqual(expected, canFill(&groups, target, &suffix_bytes, &suffix_gcd)); + } + try std.testing.expect(!canFill(&.{ .{ .bytes = 6, .available = 10 }, .{ .bytes = 4, .available = 10 } }, 2, &suffix_bytes, &suffix_gcd)); +} + +test "profile budget: heterogeneous allocations preserve the feasible uniform witness" { + const first: [experts]u32 = @splat(17); + const second: [experts]u32 = @splat(11); + const third: [experts]u32 = @splat(5); + var output: [3]usize = undefined; + for (1..6) |first_cost| { + for (1..6) |second_cost| { + for (1..6) |third_cost| { + const bytes = [_]usize{ first_cost, second_cost, third_cost }; + const plan = try allocate(std.testing.allocator, testRequest(&.{ &first, &second, &third }, &bytes, 2, 1), &output); + try std.testing.expectEqual(@as(u64, 2 * (first_cost + second_cost + third_cost) + 3 * 8), plan.allocated_bytes); + try std.testing.expectEqual(plan.baseline_bytes, plan.allocated_bytes); + var observed: u64 = 0; + for (output, bytes) |count, cost| { + try std.testing.expect(count >= 1 and count <= 256); + observed += try poolBytes(cost, count, 1, 1, 1); + } + try std.testing.expectEqual(plan.allocated_bytes, observed); + } + } + } +} + +test "profile budget: device picks capacity stays uniform across pool sizes" { + const line_bytes = (1 << 20) + 8; + for ([_]usize{ 10, 11, 31, 32, 33, 64, 256, 512 }) |lines| { + const actual = try poolBytes(line_bytes, lines, 64, 11, 10); + try std.testing.expectEqual(@as(u64, lines * line_bytes + 54 * 64 * 11 * 4), actual); + try std.testing.expectEqual(@as(usize, 54), @divExact(@as(usize, @intCast(actual - lines * line_bytes)), 64 * 11 * 4)); + } + try std.testing.expectEqual(@as(u64, 10 * line_bytes + 12 * 11 * 4), try poolBytes(line_bytes, 10, 1, 11, 10)); + try std.testing.expectEqual(@as(u64, 10 * line_bytes + 2 * 4), try poolBytes(line_bytes, 10, 1, 1, 10)); + try std.testing.expectEqual(@as(u64, 256 * line_bytes + 12 * 11 * 4), try poolBytes(line_bytes, 256, 1, 11, 10)); +} + +test "profile budget: redistribution preserves expert planner and picks bytes exactly" { + var hot: [experts]u32 = @splat(0); + @memset(hot[0..55], 100); + const cold: [experts]u32 = @splat(1); + const line_bytes = (1 << 20) + 8; + var request = testRequest(&.{ &hot, &cold }, &.{ line_bytes, line_bytes }, 54, 53); + request.max_rows = 64; + request.slots = 11; + var output: [2]usize = undefined; + const plan = try allocate(std.testing.allocator, request, &output); + try std.testing.expectEqualSlices(usize, &.{ 55, 53 }, &output); + try std.testing.expect(!plan.uniform_fallback); + try std.testing.expectEqual(@as(u64, 108 * line_bytes + 2 * 12 * 64 * 11 * 4), plan.baseline_bytes); + var observed: u64 = 0; + for (output) |lines| { + try std.testing.expect(lines >= request.min_lines); + observed += try poolBytes(line_bytes, lines, request.max_rows, request.slots, request.min_lines); + } + try std.testing.expectEqual(plan.baseline_bytes, observed); + try std.testing.expectEqual(observed, plan.allocated_bytes); +} + +test "profile budget: fixed picks capacity preserves exact bytes across pass thresholds" { + var hot: [experts]u32 = @splat(0); + @memset(hot[0..33], 100); + const cold: [experts]u32 = @splat(1); + const line_bytes = (1 << 20) + 8; + var request = testRequest(&.{ &hot, &cold }, &.{ line_bytes, line_bytes }, 32, 31); + var output: [2]usize = undefined; + _ = try allocate(std.testing.allocator, request, &output); + try std.testing.expectEqualSlices(usize, &.{ 33, 31 }, &output); + request.max_rows = 64; + request.slots = 11; + const plan = try allocate(std.testing.allocator, request, &output); + try std.testing.expectEqualSlices(usize, &.{ 33, 31 }, &output); + try std.testing.expect(!plan.uniform_fallback); + try std.testing.expectEqual(@as(u64, 64 * line_bytes + 2 * 19 * 64 * 11 * 4), plan.baseline_bytes); + try std.testing.expectEqual(plan.baseline_bytes, plan.allocated_bytes); +} + +test "profile budget: fewer picks bytes also retain the exact uniform budget" { + const hot: [experts]u32 = @splat(100); + const cold: [experts]u32 = @splat(1); + var request = testRequest(&.{ &hot, &cold }, &.{ 4, 6 }, 3, 1); + request.max_rows = 16; + request.slots = 11; + var output: [2]usize = undefined; + const plan = try allocate(std.testing.allocator, request, &output); + try std.testing.expectEqualSlices(usize, &.{ 6, 1 }, &output); + try std.testing.expect(!plan.uniform_fallback); + try std.testing.expectEqual(@as(u64, 30 + 2 * 177 * 16 * 11 * 4), plan.baseline_bytes); + try std.testing.expectEqual(plan.baseline_bytes, plan.allocated_bytes); +} + +test "profile budget: TP2 line count above the local expert bank is rejected for caller fallback" { + const histogram: [experts]u32 = @splat(1); + var request = testRequest(&.{&histogram}, &.{8}, 257, 10); + var output: [1]usize = .{257}; + try std.testing.expectError(error.InvalidBounds, allocate(std.testing.allocator, request, &output)); + try std.testing.expectEqual(@as(usize, 257), output[0]); + request.uniform_lines = 256; + _ = try allocate(std.testing.allocator, request, &output); +} + +test "profile budget: picks dimensions and byte overflow preserve output" { + const histogram: [experts]u32 = @splat(0); + var request = testRequest(&.{&histogram}, &.{8}, 20, 10); + var output = [_]usize{999}; + request.max_rows = 0; + try std.testing.expectError(error.InvalidDimensions, allocate(std.testing.allocator, request, &output)); + request.max_rows = 16; + request.slots = 0; + try std.testing.expectError(error.InvalidDimensions, allocate(std.testing.allocator, request, &output)); + request.slots = 11; + request.max_rows = std.math.maxInt(usize); + try std.testing.expectError(error.Overflow, allocate(std.testing.allocator, request, &output)); + request.slots = 1; + request.max_rows = std.math.maxInt(usize) / 4; + try std.testing.expectError(error.Overflow, allocate(std.testing.allocator, request, &output)); + try std.testing.expectEqual(@as(usize, 999), output[0]); +} diff --git a/zig/src/families/flashnext/exl3/offload/exl3_recent_demand.zig b/zig/src/families/flashnext/exl3/offload/exl3_recent_demand.zig new file mode 100644 index 0000000000..a836408b20 --- /dev/null +++ b/zig/src/families/flashnext/exl3/offload/exl3_recent_demand.zig @@ -0,0 +1,212 @@ +const std = @import("std"); + +pub const expert_count: usize = 512; +pub const Count = u32; +const Row = [expert_count / @bitSizeOf(u64)]u64; + +pub const RecentDemand = struct { + history: []Row, + counts: [expert_count]Count = @splat(0), + next_row: usize = 0, + retained_rows: usize = 0, + + pub fn init(allocator: std.mem.Allocator, window_rows: usize) !RecentDemand { + if (window_rows == 0 or window_rows > std.math.maxInt(Count)) return error.InvalidWindowSize; + const history = try allocator.alloc(Row, window_rows); + @memset(history, @splat(0)); + return .{ .history = history }; + } + + pub fn deinit(self: *RecentDemand, allocator: std.mem.Allocator) void { + allocator.free(self.history); + self.* = undefined; + } + + pub fn reset(self: *RecentDemand) void { + @memset(self.history, @splat(0)); + self.counts = @splat(0); + self.next_row = 0; + self.retained_rows = 0; + } + + pub fn windowRows(self: *const RecentDemand) usize { + return self.history.len; + } + + pub fn rowCount(self: *const RecentDemand) usize { + return self.retained_rows; + } + + pub fn scores(self: *const RecentDemand) *const [expert_count]Count { + return &self.counts; + } + + pub fn score(self: *const RecentDemand, expert_id: i32) Count { + if (expert_id < 0 or expert_id >= expert_count) return 0; + return self.counts[@intCast(expert_id)]; + } + + pub fn observeRow(self: *RecentDemand, picks: []const i32) void { + var incoming: Row = @splat(0); + for (picks) |expert_id| { + if (expert_id < 0 or expert_id >= expert_count) continue; + const expert: usize = @intCast(expert_id); + incoming[expert / @bitSizeOf(u64)] |= @as(u64, 1) << @as(u6, @intCast(expert % @bitSizeOf(u64))); + } + + const outgoing = self.history[self.next_row]; + for (outgoing, incoming, 0..) |expired_word, incoming_word, word_index| { + var removed = expired_word & ~incoming_word; + while (removed != 0) { + const expert = word_index * @bitSizeOf(u64) + @ctz(removed); + self.counts[expert] -= 1; + removed &= removed - 1; + } + var added = incoming_word & ~expired_word; + while (added != 0) { + const expert = word_index * @bitSizeOf(u64) + @ctz(added); + self.counts[expert] += 1; + added &= added - 1; + } + } + self.history[self.next_row] = incoming; + self.next_row += 1; + if (self.next_row == self.history.len) self.next_row = 0; + if (self.retained_rows < self.history.len) self.retained_rows += 1; + } +}; + +test "demand expires after exactly the configured number of rows" { + var demand = try RecentDemand.init(std.testing.allocator, 3); + defer demand.deinit(std.testing.allocator); + + demand.observeRow(&.{0}); + demand.observeRow(&.{ 1, 511 }); + demand.observeRow(&.{0}); + try std.testing.expectEqual(@as(Count, 2), demand.score(0)); + try std.testing.expectEqual(@as(Count, 1), demand.score(511)); + try std.testing.expectEqual(@as(usize, 3), demand.rowCount()); + + demand.observeRow(&.{}); + try std.testing.expectEqual(@as(Count, 1), demand.score(0)); + try std.testing.expectEqual(@as(Count, 1), demand.score(1)); + demand.observeRow(&.{ -1, 512 }); + try std.testing.expectEqual(@as(Count, 0), demand.score(1)); + try std.testing.expectEqual(@as(Count, 0), demand.score(511)); + demand.observeRow(&.{}); + try std.testing.expectEqual(@as([expert_count]Count, @splat(0)), demand.scores().*); + try std.testing.expectEqual(@as(usize, 3), demand.rowCount()); +} + +test "duplicates contribute once and shared masked and invalid picks are ignored" { + var demand = try RecentDemand.init(std.testing.allocator, 2); + defer demand.deinit(std.testing.allocator); + + demand.observeRow(&.{ -1, 0, 0, 63, 64, 64, 511, 511, 512, 512, -2, 513, std.math.maxInt(i32), std.math.minInt(i32) }); + for (demand.scores(), 0..) |count, expert| { + const expected: Count = if (expert == 0 or expert == 63 or expert == 64 or expert == 511) 1 else 0; + try std.testing.expectEqual(expected, count); + } + try std.testing.expectEqual(@as(Count, 0), demand.score(-1)); + try std.testing.expectEqual(@as(Count, 0), demand.score(512)); + try std.testing.expectEqual(@as(Count, 0), demand.score(std.math.maxInt(i32))); +} + +test "pick permutations produce identical state through wraparound" { + var ordered = try RecentDemand.init(std.testing.allocator, 3); + defer ordered.deinit(std.testing.allocator); + var permuted = try RecentDemand.init(std.testing.allocator, 3); + defer permuted.deinit(std.testing.allocator); + + for (0..10) |_| { + ordered.observeRow(&.{ 0, 63, 64, 511 }); + permuted.observeRow(&.{ 511, 64, 0, 63, 0, 512, -1 }); + ordered.observeRow(&.{ 7, 8 }); + permuted.observeRow(&.{ 8, 7, 8 }); + ordered.observeRow(&.{}); + permuted.observeRow(&.{ -1, 512 }); + try std.testing.expectEqual(ordered.counts, permuted.counts); + try std.testing.expectEqualSlices(Row, ordered.history, permuted.history); + try std.testing.expectEqual(ordered.next_row, permuted.next_row); + } +} + +test "expiration follows row order" { + var forward = try RecentDemand.init(std.testing.allocator, 2); + defer forward.deinit(std.testing.allocator); + var reverse = try RecentDemand.init(std.testing.allocator, 2); + defer reverse.deinit(std.testing.allocator); + + forward.observeRow(&.{1}); + forward.observeRow(&.{2}); + reverse.observeRow(&.{2}); + reverse.observeRow(&.{1}); + try std.testing.expectEqual(forward.counts, reverse.counts); + forward.observeRow(&.{}); + reverse.observeRow(&.{}); + try std.testing.expectEqual(@as(Count, 0), forward.score(1)); + try std.testing.expectEqual(@as(Count, 1), forward.score(2)); + try std.testing.expectEqual(@as(Count, 1), reverse.score(1)); + try std.testing.expectEqual(@as(Count, 0), reverse.score(2)); +} + +test "reset and replay reproduce every state transition" { + const rows = [_][]const i32{ &.{ 1, 2, 2 }, &.{511}, &.{}, &.{ -1, 512 }, &.{ 2, 511 }, &.{7}, &.{ 1, 7 } }; + var replayed = try RecentDemand.init(std.testing.allocator, 3); + defer replayed.deinit(std.testing.allocator); + var fresh = try RecentDemand.init(std.testing.allocator, 3); + defer fresh.deinit(std.testing.allocator); + + for (rows) |row| replayed.observeRow(row); + replayed.reset(); + try std.testing.expectEqual(fresh.counts, replayed.counts); + try std.testing.expectEqualSlices(Row, fresh.history, replayed.history); + try std.testing.expectEqual(@as(usize, 0), replayed.rowCount()); + try std.testing.expectEqual(@as(usize, 0), replayed.next_row); + try std.testing.expectEqual(@as(usize, 3), replayed.windowRows()); + + for (rows) |row| { + replayed.observeRow(row); + fresh.observeRow(row); + try std.testing.expectEqual(fresh.counts, replayed.counts); + try std.testing.expectEqualSlices(Row, fresh.history, replayed.history); + try std.testing.expectEqual(fresh.next_row, replayed.next_row); + try std.testing.expectEqual(fresh.rowCount(), replayed.rowCount()); + } + replayed.reset(); + replayed.reset(); + try std.testing.expectEqual(@as([expert_count]Count, @splat(0)), replayed.scores().*); +} + +test "all expert scores stay exact at capacity over repeated replacements" { + var demand = try RecentDemand.init(std.testing.allocator, 1); + defer demand.deinit(std.testing.allocator); + var picks: [expert_count]i32 = undefined; + for (&picks, 0..) |*expert, index| expert.* = @intCast(index); + for (0..1024) |_| { + demand.observeRow(&picks); + try std.testing.expectEqual(@as([expert_count]Count, @splat(1)), demand.scores().*); + } + demand.observeRow(&.{}); + try std.testing.expectEqual(@as([expert_count]Count, @splat(0)), demand.scores().*); + + var larger = try RecentDemand.init(std.testing.allocator, 4); + defer larger.deinit(std.testing.allocator); + for (0..16) |index| { + larger.observeRow(&picks); + const expected: Count = @intCast(@min(index + 1, larger.windowRows())); + for (larger.scores()) |count| try std.testing.expectEqual(expected, count); + } + for (0..4) |index| { + larger.observeRow(&.{}); + const expected: Count = @intCast(3 - index); + for (larger.scores()) |count| try std.testing.expectEqual(expected, count); + } +} + +test "invalid capacities are rejected before allocation" { + try std.testing.expectError(error.InvalidWindowSize, RecentDemand.init(std.testing.allocator, 0)); + if (@bitSizeOf(usize) > @bitSizeOf(Count)) { + try std.testing.expectError(error.InvalidWindowSize, RecentDemand.init(std.testing.allocator, @as(usize, std.math.maxInt(Count)) + 1)); + } +} diff --git a/zig/src/families/flashnext/exl3/pack/cuda_exl3_config.zig b/zig/src/families/flashnext/exl3/pack/cuda_exl3_config.zig new file mode 100644 index 0000000000..b4db09d186 --- /dev/null +++ b/zig/src/families/flashnext/exl3/pack/cuda_exl3_config.zig @@ -0,0 +1,113 @@ +//! EXL3 admission for the Flash Next CUDA family: config.json's quantization_config names quant_method "exl3" and the + +const std = @import("std"); + +/// The one codebook the family reads (the kernels' `mul1`, cuda.exl3.Codebook.mul1). +pub const codebook = "mul1"; +/// The weight format the family lists for an EXL3 checkpoint. +pub const format = "exl3-mul1"; + +fn object(o: std.json.ObjectMap, name: []const u8) ?std.json.ObjectMap { + const v = o.get(name) orelse return null; + return if (v == .object) v.object else null; +} + +fn word(o: std.json.ObjectMap, name: []const u8) ?[]const u8 { + const v = o.get(name) orelse return null; + return if (v == .string) v.string else null; +} + +/// config.json's quantization_config: top level, else the text config's. +pub fn block(root: std.json.ObjectMap) ?std.json.ObjectMap { + if (object(root, "quantization_config")) |q| return q; + if (object(root, "text_config")) |t| if (object(t, "quantization_config")) |q| return q; + return null; +} + +/// Whether config.json declares an EXL3 checkpoint (quant_method "exl3"), whatever its codebook. +pub fn named(root: std.json.ObjectMap) bool { + const q = block(root) orelse return false; + const m = word(q, "quant_method") orelse return false; + return std.mem.eql(u8, m, "exl3"); +} + +/// The first field an EXL3 checkpoint differs in from what the family serves, else null. +pub fn mismatch(root: std.json.ObjectMap, text: std.json.ObjectMap, comptime geometry: fn (std.json.ObjectMap) ?[]const u8) ?[]const u8 { + const q = block(root) orelse return "quantization_config"; + const m = word(q, "quant_method") orelse return "quantization_config.quant_method"; + if (!std.mem.eql(u8, m, "exl3")) return "quantization_config.quant_method"; + const book = word(q, "codebook") orelse return "quantization_config.codebook"; + if (!std.mem.eql(u8, book, codebook)) return "quantization_config.codebook"; + return geometry(text); +} + +const fixture = @embedFile("fixtures/exl3_config.json"); + +test "exl3: the Qwen3.8-Flash-Next EXL3 pack's config.json is admitted as Mode.exl3" { + const config = @import("../../cuda_config.zig"); + const a = std.testing.allocator; + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + try tmp.dir.writeFile(std.testing.io, .{ .sub_path = "config.json", .data = fixture }); + const dir = try tmp.dir.realPathFileAlloc(std.testing.io, ".", a); + defer a.free(dir); + try std.testing.expectEqual(config.Mode.exl3, try config.mode(a, std.testing.io, dir)); +} + +test "exl3: other codebooks and other methods are refused by name; bits are never read" { + const config = @import("../../cuda_config.zig"); + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const Case = struct { field: []const u8, value: std.json.Value, refused: ?[]const u8 }; + const cases = [_]Case{ + .{ .field = "codebook", .value = .{ .string = "mcg" }, .refused = "quantization_config.codebook" }, + .{ .field = "codebook", .value = .{ .string = "inst3" }, .refused = "quantization_config.codebook" }, + .{ .field = "codebook", .value = .null, .refused = "quantization_config.codebook" }, + .{ .field = "quant_method", .value = .{ .string = "gptq" }, .refused = "quantization_config.quant_method" }, + // the average width and the head's are not a contract: the trellis says each matrix's own width + .{ .field = "bits", .value = .{ .float = 4.25 }, .refused = null }, + .{ .field = "bits", .value = .{ .integer = 2 }, .refused = null }, + .{ .field = "head_bits", .value = .{ .integer = 8 }, .refused = null }, + }; + for (cases) |c| { + const parsed = try std.json.parseFromSlice(std.json.Value, a, fixture, .{}); + const root = parsed.value.object; + const text = root.get("text_config").?.object; + try std.testing.expect(named(root)); + try std.testing.expectEqual(@as(?[]const u8, null), mismatch(root, text, config.geometry)); + try root.getPtr("quantization_config").?.object.put(a, c.field, c.value); + const got = mismatch(parsed.value.object, text, config.geometry); + if (c.refused) |want| try std.testing.expectEqualStrings(want, got orelse "accepted") else try std.testing.expectEqual(@as(?[]const u8, null), got); + } +} + +test "exl3: the shapes and semantics checks still apply to an EXL3 checkpoint" { + const config = @import("../../cuda_config.zig"); + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const Case = struct { field: []const u8, value: std.json.Value, refused: []const u8 }; + const cases = [_]Case{ + .{ .field = "hidden_size", .value = .{ .integer = 4096 }, .refused = "hidden_size" }, + .{ .field = "num_experts_per_tok", .value = .{ .integer = 8 }, .refused = "num_experts_per_tok" }, + .{ .field = "rms_norm_eps", .value = .{ .float = 1e-5 }, .refused = "rms_norm_eps" }, + .{ .field = "output_gate_type", .value = .{ .string = "silu" }, .refused = "output_gate_type" }, + }; + for (cases) |c| { + const parsed = try std.json.parseFromSlice(std.json.Value, a, fixture, .{}); + const root = parsed.value.object; + const text = &root.getPtr("text_config").?.object; + try text.put(a, c.field, c.value); + try std.testing.expectEqualStrings(c.refused, mismatch(root, text.*, config.geometry) orelse "accepted"); + } +} + +test "exl3: a non-EXL3 config is not named EXL3" { + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const parsed = try std.json.parseFromSlice(std.json.Value, arena.allocator(), "{\"quantization_config\":{\"quant_method\":\"modelopt\"}}", .{}); + try std.testing.expect(!named(parsed.value.object)); + const bare = try std.json.parseFromSlice(std.json.Value, arena.allocator(), "{\"text_config\":{}}", .{}); + try std.testing.expect(!named(bare.value.object)); +} diff --git a/zig/src/families/flashnext/exl3/pack/cuda_exl3_load.zig b/zig/src/families/flashnext/exl3/pack/cuda_exl3_load.zig new file mode 100644 index 0000000000..cec7e8da81 --- /dev/null +++ b/zig/src/families/flashnext/exl3/pack/cuda_exl3_load.zig @@ -0,0 +1,527 @@ +//! EXL3 weights for the Flash Next CUDA family's own loaders (cuda_projection_load.zig, cuda_prompt_load.zig) + +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const wt = @import("../../cuda_prompt_wt.zig"); +const projection = @import("../../cuda_projection.zig"); +const pool_ = @import("../offload/cuda_exl3_pool.zig"); + +const Tensor = core.safetensors.Tensor; +const Store = wt.Store; + +/// format.py's mul1 marker value (core.exl3_load.markers). +const mul1_magic: u32 = blk: { + for (core.exl3_load.markers) |m| if (m.codebook == .mul1) break :blk m.value; + unreachable; +}; + +/// The EXL3 kernels and scratch one engine's calls share, on the engine's one compute stream. +pub const Ops = struct { + gpa: std.mem.Allocator, + d: *const cuda.Driver, + refs: usize = 1, + fork: ?*@import("../tp/cuda_exl3_fork.zig").Fork = null, + dense: ?cuda.exl3.Dense = null, + experts: ?cuda.exl3_experts.Experts = null, + /// A part's rows before its columns land in the stitched row (bf16 [rows, part n]). + tmp: ?cuda.DeviceBuffer = null, + /// The routed experts' scratch for `moe_rows` rows of 11 slots. + moe: ?cuda.DeviceBuffer = null, + moe_rows: usize = 0, + /// fp16 rows split into bf16 hi and lo [2, rows, k], and their fp32 products [2, rows, n] (an fp16-input Proj). + halves: ?cuda.DeviceBuffer = null, + sums: ?cuda.DeviceBuffer = null, + glue: ?cuda.Module = null, + split_fn: cuda.Function = undefined, + /// Sizing maxima over what was loaded (`reserve` allocates every scratch for them once, before any round): the + max_k: usize = 0, + max_z: usize = 0, + max_counters: usize = 0, + max_part_n: usize = 0, + max_f16_k: usize = 0, + max_f16_n: usize = 0, + moe_layer: ?cuda.exl3_experts.Layer = null, + /// Pinned [moe_rows, slots] int32: one block's router picks on the host (the offload pool's A/B cut). + picks: ?cuda.HostBuffer = null, + + pub fn create(gpa: std.mem.Allocator, d: *const cuda.Driver) !*Ops { + const o = try gpa.create(Ops); + o.* = .{ .gpa = gpa, .d = d }; + return o; + } + + pub fn retain(o: *Ops) *Ops { + o.refs += 1; + return o; + } + + pub fn release(o: *Ops) void { + o.refs -= 1; + if (o.refs != 0) return; + @import("../tp/cuda_exl3_fork.zig").release(o); + if (o.dense) |*x| x.unload(); + if (o.experts) |*x| x.unload(); + if (o.tmp) |*b| b.free(); + if (o.moe) |*b| b.free(); + if (o.halves) |*b| b.free(); + if (o.sums) |*b| b.free(); + if (o.picks) |*b| b.free(); + if (o.glue) |*m| m.unload(); + o.gpa.destroy(o); + } + + /// The dense EXL3 linear, loaded on first use (no GPU work at checkpoint load). + pub fn linear(o: *Ops) !*cuda.exl3.Dense { + if (o.dense == null) o.dense = try cuda.exl3.Dense.load(o.d); + return &o.dense.?; + } + + /// fn_exl3_split_f16 (flashnext_exl3.cu), loaded on first use. + pub fn split(o: *Ops) !cuda.Function { + if (o.glue == null) { + if (!cuda.kernels.available) return error.BuiltWithoutKernels; + var m = try cuda.Module.load(o.d, cuda.kernels.flashnext_exl3); + errdefer m.unload(); + o.split_fn = try m.function("fn_exl3_split_f16"); + o.glue = m; + } + return o.split_fn; + } + + pub fn grouped(o: *Ops) !*const cuda.exl3_experts.Experts { + if (o.experts == null) o.experts = try cuda.exl3_experts.Experts.load(o.d); + return &o.experts.?; + } + + pub fn noteWeight(o: *Ops, w: cuda.exl3.Weight, row_n: usize) void { + o.max_k = @max(o.max_k, w.k); + o.max_z = @max(o.max_z, cuda.exl3.Linear.zBytes(w, 128)); + o.max_counters = @max(o.max_counters, cuda.exl3.Linear.counterBytes(w)); + if (w.n != row_n) o.max_part_n = @max(o.max_part_n, w.n); + } + + /// Every scratch the EXL3 calls of a pass of up to `rows` rows read, allocated now (after `s` drains), so no call + pub fn reserve(o: *Ops, s: cuda.Stream, rows: usize, moe_rows: usize, slots: usize) !void { + try s.synchronize(); + if (o.max_k > 0) { + const de = try o.linear(); + try grow(o.d, &de.xh, 128 * o.max_k * 2); + if (o.max_z > 0) try grow(o.d, &de.z, o.max_z); + try grow(o.d, &de.counters, o.max_counters); + try @import("../tp/cuda_exl3_fork.zig").reserve(o); // parallel parts' branches (decode projections) + } + if (o.max_part_n > 0) try grow(o.d, &o.tmp, rows * o.max_part_n * 2); + if (o.max_f16_k > 0) { + _ = try o.split(); + try grow(o.d, &o.halves, 2 * rows * o.max_f16_k * 2); + try grow(o.d, &o.sums, 2 * rows * o.max_f16_n * 4); + } + if (o.moe_layer) |l| { + _ = try o.grouped(); + try grow(o.d, &o.moe, try cuda.exl3_experts.scratchBytes(l, moe_rows, slots)); + if (o.picks == null) o.picks = try cuda.HostBuffer.alloc(o.d, moe_rows * slots * 4); + } + } + + fn grow(d: *const cuda.Driver, slot: *?cuda.DeviceBuffer, bytes: usize) !void { + if (bytes == 0) return; + if (slot.*) |b| if (b.len >= bytes) return; + if (slot.*) |*b| b.free(); + slot.* = null; + slot.* = try cuda.DeviceBuffer.alloc(d, bytes); + } + + /// Every scratch buffer's address and size: unchanged across a pass means nothing grew (nor synchronised) in it. + pub fn fingerprint(o: *const Ops) [8][2]u64 { + var f: [8][2]u64 = @splat(.{ 0, 0 }); + const de: ?cuda.exl3.Dense = o.dense; + const bufs = [8]?cuda.DeviceBuffer{ o.tmp, o.moe, o.halves, o.sums, if (de) |x| x.xh else null, if (de) |x| x.z else null, if (de) |x| x.counters else null, null }; + for (bufs, &f) |b, *e| if (b) |x| { + e.* = .{ x.ptr, x.len }; + }; + if (o.picks) |h| f[7] = .{ @intFromPtr(h.bytes.ptr), h.bytes.len }; + return f; + } + + /// Device bytes `reserve` holds (for the pool's budget). + pub fn reservedBytes(o: *const Ops) usize { + var n: usize = 0; + for ([_]?cuda.DeviceBuffer{ o.tmp, o.moe, o.halves, o.sums }) |b| if (b) |x| { + n += x.len; + }; + if (o.dense) |de| for ([_]?cuda.DeviceBuffer{ de.xh, de.z, de.counters }) |b| if (b) |x| { + n += x.len; + }; + return n; + } + + /// At least `bytes` of part scratch; a growth waits for `s` first (an older call may still read the old buffer). + pub fn scratch(o: *Ops, s: cuda.Stream, slot: *?cuda.DeviceBuffer, bytes: usize) !u64 { + if (slot.*) |b| if (b.len >= bytes) return b.ptr; + try s.synchronize(); + if (slot.*) |*b| b.free(); + slot.* = null; + slot.* = try cuda.DeviceBuffer.alloc(o.d, bytes); + return slot.*.?.ptr; + } +}; + +/// The store's shared Ops, made on first use; every holder retains it. +pub fn opsOf(store: *Store, d: *const cuda.Driver) !*Ops { + if (store.exl3_ops == null) store.exl3_ops = try Ops.create(store.gpa, d); + return store.exl3_ops.?.retain(); +} + +/// One EXL3 matrix read from the store: words as stored, its width and shape. +pub const Mat = struct { + trellis: []const u8, + suh: []const u8, + svh: []const u8, + k: usize, + n: usize, + halves: u32, +}; + +/// Whether `prefix` names an EXL3 matrix (its trellis is indexed). +pub fn has(store: *const Store, prefix: []const u8) bool { + var buf: [256]u8 = undefined; + return store.has(std.fmt.bufPrint(&buf, "{s}.trellis", .{prefix}) catch return false); +} + +/// `prefix`'s trellis, scales and mul1 marker, checked: int16 [K/16, N/16, 16 * bits], fp16 suh [K] and svh [N]. +pub fn mat(store: *Store, prefix: []const u8) !Mat { + var buf: [256]u8 = undefined; + const t = try store.hold(try std.fmt.bufPrint(&buf, "{s}.trellis", .{prefix})); + if (t.dtype != .i16 or t.rank != 3 or t.dim(2) % 8 != 0) return error.UnexpectedTensor; + const halves = t.dim(2) / 8; + // the mul1 half steps are odd halves; 2..16 halves (1 to 8 bits) are what the kernels instantiate + if (halves < 2 or halves > 16) return error.UnsupportedExl3Width; + const k = t.dim(0) * 16; + const n = t.dim(1) * 16; + if (k % 128 != 0 or n % 128 != 0) return error.UnexpectedTensor; + const marker = store.hold(try std.fmt.bufPrint(&buf, "{s}.mul1", .{prefix})) catch return error.NotExl3Mul1; + if (marker.dtype != .i32 or marker.bytes.len != 4 or std.mem.readInt(u32, marker.bytes[0..4], .little) != mul1_magic) return error.NotExl3Mul1; + const suh = try store.hold(try std.fmt.bufPrint(&buf, "{s}.suh", .{prefix})); + const svh = try store.hold(try std.fmt.bufPrint(&buf, "{s}.svh", .{prefix})); + if (!suh.is(.f16, &.{k}) or !svh.is(.f16, &.{n})) return error.UnexpectedTensor; + if (store.has(try std.fmt.bufPrint(&buf, "{s}.bias", .{prefix}))) return error.UnexpectedTensor; // none in the text model + return .{ .trellis = t.bytes, .suh = suh.bytes, .svh = svh.bytes, .k = k, .n = n, .halves = @intCast(halves) }; +} + +/// One column range of a projection's output row. +pub const Part = struct { + col: usize, + n: usize, + kind: union(enum) { exl3: cuda.exl3.Weight, native: cuda.qmmf.Weight }, +}; + +/// A projection's parts and the device buffers they own. +pub const Proj = struct { + gpa: std.mem.Allocator, + ops: *Ops, + k: usize, + n: usize, + parts: []Part, + /// The rows the pass hands in: bf16, or fp16 bit patterns (an EXL3 pack's n-gram rows, cuda_exl3_ple.zig). + x: enum { bf16, f16 } = .bf16, + bufs: std.ArrayList(cuda.DeviceBuffer) = .empty, + natives: std.ArrayList(projection.Projection) = .empty, + + pub fn destroy(p: *Proj) void { + for (p.bufs.items) |*b| b.free(); + p.bufs.deinit(p.gpa); + for (p.natives.items) |*q| q.free(); + p.natives.deinit(p.gpa); + p.gpa.free(p.parts); + p.ops.release(); + p.gpa.destroy(p); + } + + fn up(p: *Proj, d: *const cuda.Driver, bytes: []const u8) !u64 { + try p.bufs.ensureUnusedCapacity(p.gpa, 1); + const b = try cuda.DeviceBuffer.fromHost(d, bytes); + p.bufs.appendAssumeCapacity(b); + return b.ptr; + } + + /// An EXL3 matrix as the stored-layout weight (trellis [K/16][N/16][W]: strides(k, n, halves, false)). + fn exl3Part(p: *Proj, d: *const cuda.Driver, m: Mat) !cuda.exl3.Weight { + const st = cuda.exl3.strides(m.k, m.n, m.halves, false); + return .{ .trellis = try p.up(d, m.trellis), .suh = try p.up(d, m.suh), .svh = try p.up(d, m.svh), .k = @intCast(m.k), .n = @intCast(m.n), .halves = m.halves, .codebook = .mul1, .stride_k = st.k, .stride_nb = st.nb }; + } +}; + +/// A Projection whose rows come from `parts` (EXL3 matrices or dense BF16 rows), `n` columns of `k` inputs. +fn wrap(gpa: std.mem.Allocator, d: *cuda.Driver, store: *Store, prefix: []const u8, names: []const []const u8, n: usize, k: usize) !projection.Projection { + const p = try gpa.create(Proj); + p.* = .{ .gpa = gpa, .ops = try opsOf(store, d), .k = k, .n = n, .parts = &.{} }; + errdefer p.destroy(); + var parts: std.ArrayList(Part) = .empty; + errdefer parts.deinit(gpa); + var buf: [256]u8 = undefined; + var col: usize = 0; + var i: usize = 0; + while (i < names.len) { + const name = try std.fmt.bufPrint(&buf, "{s}{s}", .{ prefix, names[i] }); + if (has(store, name)) { + const m = try mat(store, name); + if (m.k != k) return error.UnexpectedTensor; + const w = try p.exl3Part(d, m); + p.ops.noteWeight(w, n); + try parts.append(gpa, .{ .col = col, .n = m.n, .kind = .{ .exl3 = w } }); + col += m.n; + i += 1; + continue; + } + // a run of dense BF16 parts (after the Store's exact widening) packs as one lane weight + var rows: usize = 0; + var j = i; + var bytes: std.ArrayList(u8) = .empty; + defer bytes.deinit(gpa); + while (j < names.len) : (j += 1) { + var nb: [256]u8 = undefined; + const dense_name = try std.fmt.bufPrint(&nb, "{s}{s}", .{ prefix, names[j] }); + if (has(store, dense_name)) break; + const t = try store.tensor(try std.fmt.bufPrint(&nb, "{s}{s}.weight", .{ prefix, names[j] })); + if (t.dtype != .bf16 or t.rank != 2 or t.dim(1) != k or t.bytes.len != t.dim(0) * k * 2) return error.UnexpectedTensor; + try bytes.appendSlice(gpa, t.bytes); + rows += t.dim(0); + } + try p.natives.ensureUnusedCapacity(gpa, 1); + const q = try projection.dense(gpa, d, .{ .dtype = .bf16, .rank = 2, .shape = .{ rows, k, 1, 1, 1 }, .bytes = bytes.items }); + p.natives.appendAssumeCapacity(q); + try parts.append(gpa, .{ .col = col, .n = rows, .kind = .{ .native = q.native.? } }); + col += rows; + i = j; + } + if (col != n) return error.UnexpectedTensor; + p.parts = try parts.toOwnedSlice(gpa); + var out = projection.Projection.empty(d); + out.n = n; + out.k = k; + out.exl3 = p; + return out; +} + +/// cuda_projection_load.load's EXL3 branch: `prefix` one EXL3 matrix of n outputs and k inputs. +pub fn load(gpa: std.mem.Allocator, d: *cuda.Driver, store: *Store, prefix: []const u8, n: usize, k: usize) !projection.Projection { + return wrap(gpa, d, store, prefix, &.{""}, n, k); +} + +/// Engine.initRank's head for an EXL3 pack: one rank holds the whole vocabulary. +pub fn head(gpa: std.mem.Allocator, d: *cuda.Driver, store: *Store, world: usize, rank: usize) !?projection.Projection { + if (!has(store, "language_model.lm_head")) return null; + if (world != 1 or rank != 0) return error.UnsupportedQuantization; + if (try @import("../../cuda_head.zig").enabled()) return error.BadHeadSharding; + return try load(gpa, d, store, "language_model.lm_head", wt.vocab, wt.dims); +} + +/// Whether any of `prefix`'s parts is an EXL3 matrix (cuda_projection_load.loadStitched's branch). +pub fn hasAny(store: *const Store, prefix: []const u8, names: []const []const u8) bool { + var buf: [256]u8 = undefined; + for (names) |name| if (has(store, std.fmt.bufPrint(&buf, "{s}{s}", .{ prefix, name }) catch return false)) return true; + return false; +} + +/// cuda_projection_load.loadStitched's EXL3 branch: parts side by side in one output row, EXL3 or dense BF16. +pub fn stitched(gpa: std.mem.Allocator, d: *cuda.Driver, store: *Store, prefix: []const u8, names: []const []const u8, n: usize, k: usize) !projection.Projection { + return wrap(gpa, d, store, prefix, names, n, k); +} + +/// cuda_prompt_load's PLE branch: an EXL3 pack's n-gram rows reach key_proj and value_proj as fp16 bit patterns +pub fn pleRows(gpa: std.mem.Allocator, d: *cuda.Driver, store: *Store, key: *projection.Projection, value: *projection.Projection) !void { + if (store.exl3_ops == null) return; // not an EXL3 pack: the n-gram rows are bf16 + for ([_]*projection.Projection{ key, value }) |q| { + const native = q.native orelse return error.UnexpectedTensor; + const p = try gpa.create(Proj); + p.* = .{ .gpa = gpa, .ops = try opsOf(store, d), .k = q.k, .n = q.n, .parts = &.{}, .x = .f16 }; + errdefer p.destroy(); + p.ops.max_f16_k = @max(p.ops.max_f16_k, q.k); + p.ops.max_f16_n = @max(p.ops.max_f16_n, q.n); + p.parts = try gpa.alloc(Part, 1); + p.parts[0] = .{ .col = 0, .n = q.n, .kind = .{ .native = native } }; + try p.natives.append(gpa, q.*); + var out = projection.Projection.empty(d); + out.n = q.n; + out.k = q.k; + out.exl3 = p; + q.* = out; + } +} + +/// One layer's routed experts and the shared expert (id `wt.experts_n`) for exl3_experts.Experts. +pub const Experts = struct { + gpa: std.mem.Allocator, + ops: *Ops, + layer: cuda.exl3_experts.Layer, + bufs: std.ArrayList(cuda.DeviceBuffer) = .empty, + /// Each expert's and projection's trellis words on the device, for a pool to rewrite (gate, up, down; 513 each). + trellis: [3][]u64 = .{ &.{}, &.{}, &.{} }, + halves: [3][]i32 = .{ &.{}, &.{}, &.{} }, + bank: ?cuda.HostBuffer = null, + /// The pool's Source: every routed expert's gate/up/down bytes inside `bank`, and the shared expert's. + routed: [][3][]const u8 = &.{}, + shared: [3][]const u8 = .{ &.{}, &.{}, &.{} }, + pool: ?*pool_.Pool = null, + /// Host time the bank's copies took at load (ns). + bank_ns: u64 = 0, + + pub fn destroy(x: *Experts) void { + if (x.pool) |p| { + p.deinit(); + x.gpa.destroy(p); + } + for (x.bufs.items) |*b| b.free(); + x.bufs.deinit(x.gpa); + for (x.trellis) |t| x.gpa.free(t); + for (x.halves) |h| x.gpa.free(h); + x.gpa.free(x.routed); + if (x.bank) |*b| b.free(); + x.ops.release(); + x.gpa.destroy(x); + } + + fn up(x: *Experts, d: *const cuda.Driver, bytes: []const u8) !u64 { + try x.bufs.ensureUnusedCapacity(x.gpa, 1); + const b = try cuda.DeviceBuffer.fromHost(d, bytes); + x.bufs.appendAssumeCapacity(b); + return b.ptr; + } + + /// Whether the routed experts wait for a pool (offloaded, none attached yet). + pub fn needsPool(x: *const Experts) bool { + return x.bank != null and x.pool == null; + } + + /// Bytes of one pool line of this layer (gate, up and down of its widest expert, 256-aligned). + pub fn lineBytes(x: *const Experts) usize { + var bytes: [3]usize = .{ 0, 0, 0 }; + for (x.routed) |e| for (0..3) |j| { + bytes[j] = @max(bytes[j], e[j].len); + }; + return cuda.exl3_pool.Layout.init(bytes).line_bytes; + } + + /// Makes the pool: `lines` device lines, `counts` the family's lifetime use table (null: the pool's own), `layer` + pub fn attach(x: *Experts, io: std.Io, d: *const cuda.Driver, which: u8, lines: usize, max_rows: usize, counts: ?*[pool_.shape.layers][pool_.shape.experts]u32) !void { + if (x.bank == null or x.pool != null) return error.Invalid; + const p = try x.gpa.create(pool_.Pool); + errdefer x.gpa.destroy(p); + const l = x.layer; + p.* = try pool_.Pool.init(x.gpa, io, d, .{ .routed = x.routed, .shared = if (x.shared[0].len == 0) null else x.shared, .pinned_bank = x.bank.?.bytes }, .{ .suh_g = l.suh_g, .suh_u = l.suh_u, .svh_g = l.svh_g, .svh_u = l.svh_u, .suh_d = l.suh_d, .svh_d = l.svh_d }, .{ + .layer = which, + .lines = lines, + .dims = wt.dims, + .width = wt.moe_w, + .max_rows = max_rows, + .slots = wt.slots, + .direct = true, // the bank is pinned: one H2D a projection on the copy engine, no staging memcpy + .codebook = .mul1, + .act_mode = l.act_mode, + .max_passes = pool_.passesFor(pool_.min_lines), + }, counts); + x.pool = p; + x.layer = p.tables(); + } +}; + +/// Whether `mlp` (".mlp" of a layer) holds EXL3 experts. +pub fn hasExperts(store: *const Store, mlp: []const u8) bool { + var buf: [256]u8 = undefined; + return store.has(std.fmt.bufPrint(&buf, "{s}.shared_expert.gate_proj.trellis", .{mlp}) catch return false); +} + +/// TF_FLASHNEXT_EXL3_ROUTED=device keeps every routed expert on the device (cuda_exl3_memory.zig reads the same). +pub fn offloaded() bool { + return @import("cuda_exl3_memory.zig").options().routed_on_host; +} + +fn align256(n: usize) usize { + return (n + 255) / 256 * 256; +} + +/// Every routed expert of `mlp`, the shared expert as id 512 (the same 2560 x 640 x 2560 shape): resident, or with +pub fn experts(gpa: std.mem.Allocator, d: *cuda.Driver, store: *Store, mlp: []const u8, world: usize, rank: usize) !*Experts { + if (world != 1 or rank != 0) return error.UnsupportedQuantization; // one rank holds all 513 + const count = wt.experts_n + 1; + const offload = offloaded(); + const x = try gpa.create(Experts); + x.* = .{ .gpa = gpa, .ops = try opsOf(store, d), .layer = undefined }; + errdefer x.destroy(); + for (&x.trellis, &x.halves) |*t, *h| { + t.* = try gpa.alloc(u64, count); + @memset(t.*, 0); + h.* = try gpa.alloc(i32, count); + @memset(h.*, 0); + } + const projs = [3][]const u8{ "gate_proj", "up_proj", "down_proj" }; + const ks = [3]usize{ wt.dims, wt.dims, wt.moe_w }; + const ns = [3]usize{ wt.moe_w, wt.moe_w, wt.dims }; + var scales: [6]std.ArrayList(u8) = @splat(.empty); + defer for (&scales) |*l| l.deinit(gpa); + var buf: [256]u8 = undefined; + const nameOf = struct { + fn f(b: []u8, m: []const u8, e: usize, proj: []const u8) ![]const u8 { + if (e == wt.experts_n) return std.fmt.bufPrint(b, "{s}.shared_expert.{s}", .{ m, proj }); + return std.fmt.bufPrint(b, "{s}.experts.{d}.{s}", .{ m, e, proj }); + } + }.f; + var at: usize = 0; + if (offload) { + // the bank's size first (headers only), then one pinned region for the layer + var total: usize = 0; + for (0..count) |e| { + for (projs) |proj| total += align256((try mat(store, try nameOf(&buf, mlp, e, proj))).trellis.len); + store.release(); + } + x.bank = try cuda.HostBuffer.alloc(d, total); + x.routed = try gpa.alloc([3][]const u8, wt.experts_n); + } + const t0 = std.Io.Clock.awake.now(store.io).toNanoseconds(); + for (0..count) |e| { + for (projs, 0..) |proj, j| { + const m = try mat(store, try nameOf(&buf, mlp, e, proj)); + if (m.k != ks[j] or m.n != ns[j]) return error.UnexpectedTensor; + if (x.bank) |bank| { + const dst = bank.bytes[at..][0..m.trellis.len]; + @memcpy(dst, m.trellis); + at += align256(m.trellis.len); + if (e == wt.experts_n) x.shared[j] = dst else x.routed[e][j] = dst; + } else x.trellis[j][e] = try x.up(d, m.trellis); + x.halves[j][e] = @intCast(m.halves); + try scales[2 * j].appendSlice(gpa, m.suh); + try scales[2 * j + 1].appendSlice(gpa, m.svh); + } + // each expert's mapped pages are read once: release them so the 513 experts never pin a shard's worth + store.release(); + } + x.bank_ns = @intCast(std.Io.Clock.awake.now(store.io).toNanoseconds() - t0); + x.layer = .{ + // offloaded: the pool's tables replace these at attach (0 until then; a pass refuses an unattached layer) + .gate = if (offload) 0 else try x.up(d, std.mem.sliceAsBytes(x.trellis[0])), + .up = if (offload) 0 else try x.up(d, std.mem.sliceAsBytes(x.trellis[1])), + .down = if (offload) 0 else try x.up(d, std.mem.sliceAsBytes(x.trellis[2])), + .gate_k2 = if (offload) 0 else try x.up(d, std.mem.sliceAsBytes(x.halves[0])), + .up_k2 = if (offload) 0 else try x.up(d, std.mem.sliceAsBytes(x.halves[1])), + .down_k2 = if (offload) 0 else try x.up(d, std.mem.sliceAsBytes(x.halves[2])), + .suh_g = try x.up(d, scales[0].items), + .svh_g = try x.up(d, scales[1].items), + .suh_u = try x.up(d, scales[2].items), + .svh_u = try x.up(d, scales[3].items), + .suh_d = try x.up(d, scales[4].items), + .svh_d = try x.up(d, scales[5].items), + .experts = count, + .dims = wt.dims, + .width = wt.moe_w, + .codebook = .mul1, + // fp32 SwiGLU between the grouped GEMMs (GLM's mode 0 rounds to bf16 there) + .act_mode = 1, + }; + x.ops.moe_layer = x.layer; + return x; +} + +test "exl3: the mul1 marker value is format.py's" { + try std.testing.expectEqual(@as(u32, 0x83DCD12D), mul1_magic); +} diff --git a/zig/src/families/flashnext/exl3/pack/cuda_exl3_memory.zig b/zig/src/families/flashnext/exl3/pack/cuda_exl3_memory.zig new file mode 100644 index 0000000000..ae55c588b9 --- /dev/null +++ b/zig/src/families/flashnext/exl3/pack/cuda_exl3_memory.zig @@ -0,0 +1,134 @@ +//! Device bytes of an EXL3 Flash Next checkpoint as cuda_exl3_load.zig places it: + +const std = @import("std"); +const core = @import("core"); +const cuda = @import("cuda"); +const wt = @import("../../cuda_prompt_wt.zig"); +const config = @import("../../cuda_config.zig"); +const Entry = core.safetensors.Entry; + +pub const Options = struct { + /// Routed experts' trellis words live in pinned host memory (the offload pool's), not counted here. + routed_on_host: bool = true, +}; + +/// TF_FLASHNEXT_EXL3_ROUTED=device counts every routed expert as resident; the default is the offload pool's layout. +pub fn options() Options { + const v = std.c.getenv("TF_FLASHNEXT_EXL3_ROUTED") orelse return .{}; + return .{ .routed_on_host = !std.mem.eql(u8, std.mem.span(v), "device") }; +} + +/// Whether `dir`'s config.json admits as Mode.exl3 (any other mode, or a refusal: false). +pub fn isExl3(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) bool { + return (config.mode(gpa, io, dir) catch return false) == .exl3; +} + +/// One tensor's device bytes (0: kept on the host). +pub fn tensorBytes(name: []const u8, e: Entry, o: Options) !u64 { + const bytes: u64 = e.end - e.begin; + if (wt.hostOnly(name)) return 0; + // the BF16 embedding table is pinned host memory the gather reads in place (cuda_embed.Table.uploadDenseHost) + if (e.dtype == .bf16 and std.mem.indexOf(u8, name, "embed_tokens") != null) return 0; + // the codebook marker is checked at load and never uploaded + if (std.mem.endsWith(u8, name, ".mul1") or std.mem.endsWith(u8, name, ".mcg")) return 0; + if (std.mem.indexOf(u8, name, ".mlp.experts.") != null) { + if (std.mem.endsWith(u8, name, ".trellis")) return if (o.routed_on_host) 0 else bytes; + return bytes; // stacked suh / svh, resident + } + if (std.mem.endsWith(u8, name, ".trellis") or std.mem.endsWith(u8, name, ".suh") or std.mem.endsWith(u8, name, ".svh")) return bytes; + if (!std.mem.endsWith(u8, name, ".weight") and e.rank == 1) return try std.math.mul(u64, e.dim(0), 4); // A_log, dt_bias as fp32 + if ((e.dtype == .f16 or e.dtype == .bf16) and e.rank == 2 and std.mem.endsWith(u8, name, ".weight")) { + // the router's rows and the embedding table go up as stored (BF16). + if (std.mem.indexOf(u8, name, ".mlp.gate.") != null or std.mem.indexOf(u8, name, ".shared_expert_gate.") != null or std.mem.indexOf(u8, name, "embed_tokens") != null) return bytes; + return try std.math.mul(u64, try std.math.mul(u64, cuda.nvfp4.padded(e.dim(0)), e.dim(1)), 2); + } + // norms the loaders scale go up as fp32 (cuda_rank_memory.zig's list), the GDN gated norm and convs as stored + if (e.rank == 1 and centered(name)) return try std.math.mul(u64, e.dim(0), 4); + return bytes; +} + +fn centered(name: []const u8) bool { + for ([_][]const u8{ ".hc_norm.weight", ".q_norm.weight", ".k_norm.weight", ".q_layernorm.weight", ".k_layernorm.weight", ".norm_key.weight", ".norm_query.weight", ".norm_conv.weight" }) |suffix| { + if (std.mem.endsWith(u8, name, suffix)) return true; + } + return false; +} + +/// Per layer: three 513-entry pointer tables (int64) and three width tables (int32) for the experts' grouped GEMMs. +const tables_per_layer: u64 = 3 * (wt.experts_n + 1) * 8 + 3 * (wt.experts_n + 1) * 4; + +/// Every shard the index names (and the mixer patch, which Store.open adds), each read for its header alone. +pub fn deviceBytes(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, o: Options) !u64 { + var store = try wt.Store.open(gpa, io, dir); + defer store.deinit(); + const wm = store.parsed.value.object.get("weight_map") orelse return error.MissingTensor; + var files: std.StringHashMapUnmanaged(void) = .empty; + defer files.deinit(gpa); + var it = wm.object.iterator(); + while (it.next()) |kv| try files.put(gpa, kv.value_ptr.string, {}); + var total: u64 = tables_per_layer * wt.layers_n; + var fi = files.keyIterator(); + while (fi.next()) |rel| { + const path = try std.fs.path.join(gpa, &.{ dir, rel.* }); + defer gpa.free(path); + var file = try core.safetensors.File.open(gpa, io, path); + defer file.close(io); + var ni = file.names.iterator(); + while (ni.next()) |e| total = try std.math.add(u64, total, try tensorBytes(e.key_ptr.*, e.value_ptr.*, o)); + } + return total; +} + +/// The family's deviceWeightBytes for an EXL3 checkpoint; 0 if unread. +pub fn weightBytes(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) u64 { + return deviceBytes(gpa, io, dir, options()) catch 0; +} + +/// A rank's placement bytes: one rank holds the whole model (no MTP head: an EXL3 pack decodes plain). +pub fn placement(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, world: usize, rank: usize) !u64 { + if (world != 1 or rank != 0) return error.BadRank; + return deviceBytes(gpa, io, dir, options()); +} + +fn entry(dtype: core.safetensors.DType, shape: []const usize, bytes: u64) Entry { + var e: Entry = .{ .dtype = dtype, .rank = @intCast(shape.len), .shape = @splat(1), .begin = 0, .end = bytes }; + for (shape, 0..) |d, i| e.shape[i] = d; + return e; +} + +test "exl3: routed trellis stays on the host under offload, its scales and every other EXL3 matrix are resident" { + const tr = entry(.i16, &.{ 160, 40, 24 }, 160 * 40 * 24 * 2); + const routed = "model.language_model.layers.0.mlp.experts.7.gate_proj.trellis"; + try std.testing.expectEqual(@as(u64, 0), try tensorBytes(routed, tr, .{})); + try std.testing.expectEqual(tr.end, try tensorBytes(routed, tr, .{ .routed_on_host = false })); + const suh = entry(.f16, &.{2560}, 5120); + try std.testing.expectEqual(@as(u64, 5120), try tensorBytes("model.language_model.layers.0.mlp.experts.7.gate_proj.suh", suh, .{})); + try std.testing.expectEqual(tr.end, try tensorBytes("model.language_model.layers.0.mlp.shared_expert.gate_proj.trellis", tr, .{})); + try std.testing.expectEqual(tr.end, try tensorBytes("lm_head.trellis", tr, .{})); + try std.testing.expectEqual(@as(u64, 0), try tensorBytes("lm_head.mul1", entry(.i32, &.{}, 4), .{})); +} + +test "exl3: dense F16 weights count as their BF16 device packing, norms as fp32" { + const router = entry(.f16, &.{ 512, 2560 }, 512 * 2560 * 2); + try std.testing.expectEqual(router.end, try tensorBytes("model.language_model.layers.3.mlp.gate.weight", router, .{})); + const down = entry(.f16, &.{ 320, 10240 }, 320 * 10240 * 2); + try std.testing.expectEqual(@as(u64, cuda.nvfp4.padded(320) * 10240 * 2), try tensorBytes("model.language_model.layers.3.attn_hyper_connection.input_mix_weight_down.weight", down, .{})); + const norm = entry(.f16, &.{10240}, 20480); + try std.testing.expectEqual(@as(u64, 40960), try tensorBytes("model.language_model.layers.3.attn_hyper_connection.hc_norm.weight", norm, .{})); + const gnorm = entry(.bf16, &.{128}, 256); + try std.testing.expectEqual(@as(u64, 256), try tensorBytes("model.language_model.layers.0.linear_attn.norm.weight", gnorm, .{})); + try std.testing.expectEqual(@as(u64, 48 * 4), try tensorBytes("model.language_model.layers.0.linear_attn.A_log", entry(.bf16, &.{48}, 96), .{})); + try std.testing.expectEqual(@as(u64, 0), try tensorBytes("model.visual.blocks.0.attn.qkv.weight", entry(.bf16, &.{ 3456, 1152 }, 3456 * 1152 * 2), .{})); +} + +test "exl3: placement names an EXL3 checkpoint and refuses any world but one" { + const a = std.testing.allocator; + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + try tmp.dir.writeFile(std.testing.io, .{ .sub_path = "config.json", .data = @embedFile("fixtures/exl3_config.json") }); + const dir = try tmp.dir.realPathFileAlloc(std.testing.io, ".", a); + defer a.free(dir); + try std.testing.expect(isExl3(a, std.testing.io, dir)); + try std.testing.expectError(error.BadRank, placement(a, std.testing.io, dir, 2, 0)); + try std.testing.expectError(error.BadRank, placement(a, std.testing.io, dir, 1, 1)); +} diff --git a/zig/src/families/flashnext/exl3/pack/cuda_exl3_store.zig b/zig/src/families/flashnext/exl3/pack/cuda_exl3_store.zig new file mode 100644 index 0000000000..147d2dbffb --- /dev/null +++ b/zig/src/families/flashnext/exl3/pack/cuda_exl3_store.zig @@ -0,0 +1,121 @@ +//! The EXL3 pack's tensors as the Flash Next loaders read them (cuda_prompt_wt.zig `Store`). + +const std = @import("std"); +const core = @import("core"); + +/// The mixer file the index does not name (the MTP head's `mtp.hyper_connection_mixer.*`). +pub const patch_file = "mtp_hyper_connection_mixer_patch.safetensors"; + +/// Adds every tensor of `dir`'s mixer patch to the index's weight map (the file absent: nothing). +pub fn addPatch(io: std.Io, dir: []const u8, parsed: *std.json.Parsed(std.json.Value)) !void { + const a = parsed.arena.allocator(); + const wm = (parsed.value.object.getPtr("weight_map") orelse return); + if (wm.* != .object) return; + const path = try std.fs.path.join(a, &.{ dir, patch_file }); + var file = core.safetensors.File.open(a, io, path) catch |err| switch (err) { + error.FileNotFound => return, + else => return err, + }; + defer file.close(io); + var it = file.names.iterator(); + while (it.next()) |e| { + if (wm.object.contains(e.key_ptr.*)) continue; + try wm.object.put(a, try a.dupe(u8, e.key_ptr.*), .{ .string = patch_file }); + } +} + +/// Whether the loaders read `name` widened: a dense F16 `.weight` (an EXL3 part never ends in `.weight`). +pub fn widens(name: []const u8, t: core.safetensors.Tensor) bool { + return t.dtype == .f16 and std.mem.endsWith(u8, name, ".weight"); +} + +/// RNE to BF16 of an F32 (cuda_prompt_wt.toBf16's rounding, without its NaN quieting: a NaN is inexact here). +fn bf16Of(v: f32) u16 { + const bits: u32 = @bitCast(v); + return @intCast((bits + 0x7FFF + ((bits >> 16) & 1)) >> 16); +} + +/// The F16 words as BF16, exactly, or error.F16NotBf16Exact at the first value BF16 cannot hold (NaN included). +pub fn exactBf16(out: []u16, f16_bytes: []const u8) !void { + if (f16_bytes.len != out.len * 2) return error.UnexpectedTensor; + for (out, 0..) |*o, i| { + const h: f16 = @bitCast(std.mem.readInt(u16, f16_bytes[2 * i ..][0..2], .little)); + const v: f32 = h; + const b = bf16Of(v); + const back: f32 = @bitCast(@as(u32, b) << 16); + if (back != v or std.math.isNan(v)) return error.F16NotBf16Exact; + o.* = b; + } +} + +/// A widened copy of `t` as a BF16 tensor of the same shape, in `owned` (the caller frees it). +pub fn widen(gpa: std.mem.Allocator, name: []const u8, t: core.safetensors.Tensor, owned: *std.ArrayList([]u16)) !core.safetensors.Tensor { + const count = t.bytes.len / 2; + const words = try gpa.alloc(u16, count); + errdefer gpa.free(words); + exactBf16(words, t.bytes) catch |err| { + if (err == error.F16NotBf16Exact) std.log.err("flashnext: {s} is F16 with a value BF16 cannot hold exactly", .{name}); + return err; + }; + try owned.append(gpa, words); + var out = t; + out.dtype = .bf16; + out.bytes = std.mem.sliceAsBytes(words); + return out; +} + +test "exl3: an F16 that came from BF16 widens bit for bit and one extra mantissa bit is refused" { + const values = [_]f32{ 0, -0.0, 1, -2.5, 0.0078125, 65280, -3.0517578125e-05, 5.9604644775390625e-08, std.math.inf(f32) }; + var bytes: [values.len * 2]u8 = undefined; + for (values, 0..) |v, i| { + const b: f32 = @bitCast(@as(u32, bf16Of(v)) << 16); + try std.testing.expectEqual(v, b); // every case is a BF16 value + std.mem.writeInt(u16, bytes[2 * i ..][0..2], @bitCast(@as(f16, @floatCast(v))), .little); + } + var out: [values.len]u16 = undefined; + try exactBf16(&out, &bytes); + for (values, out) |v, o| try std.testing.expectEqual(v, @as(f32, @bitCast(@as(u32, o) << 16))); + // 1 + 2^-10 needs 11 significant bits: F16 holds it, BF16 does not + std.mem.writeInt(u16, bytes[0..2], @bitCast(@as(f16, 1.0009765625)), .little); + try std.testing.expectError(error.F16NotBf16Exact, exactBf16(&out, &bytes)); + std.mem.writeInt(u16, bytes[0..2], @bitCast(std.math.nan(f16)), .little); + try std.testing.expectError(error.F16NotBf16Exact, exactBf16(&out, &bytes)); +} + +test "exl3: only dense F16 weights widen, EXL3 parts keep their dtype" { + const t: core.safetensors.Tensor = .{ .dtype = .f16, .rank = 1, .shape = .{ 4, 1, 1, 1, 1 }, .bytes = &.{} }; + try std.testing.expect(widens("model.language_model.layers.0.mlp.gate.weight", t)); + try std.testing.expect(widens("mtp.hyper_connection_mixer.hc_norm.weight", t)); + try std.testing.expect(!widens("model.language_model.layers.0.self_attn.q_proj.suh", t)); + try std.testing.expect(!widens("lm_head.svh", t)); + var b = t; + b.dtype = .bf16; + try std.testing.expect(!widens("model.language_model.layers.0.linear_attn.norm.weight", b)); +} + +test "exl3: the mixer patch's names join the weight map, an indexed name keeps its shard" { + const a = std.testing.allocator; + const io = std.testing.io; + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + const header = "{\"mtp.hyper_connection_mixer.hc_norm.weight\":{\"dtype\":\"F16\",\"shape\":[1],\"data_offsets\":[0,2]},\"x\":{\"dtype\":\"U8\",\"shape\":[1],\"data_offsets\":[2,3]}}"; + var data: [8 + header.len + 3]u8 = undefined; + std.mem.writeInt(u64, data[0..8], header.len, .little); + @memcpy(data[8..][0..header.len], header); + @memset(data[8 + header.len ..], 0); + try tmp.dir.writeFile(io, .{ .sub_path = patch_file, .data = &data }); + const dir = try tmp.dir.realPathFileAlloc(io, ".", a); + defer a.free(dir); + var parsed = try std.json.parseFromSlice(std.json.Value, a, "{\"weight_map\":{\"x\":\"model-00001-of-00007.safetensors\"}}", .{}); + defer parsed.deinit(); + try addPatch(io, dir, &parsed); + const wm = parsed.value.object.get("weight_map").?.object; + try std.testing.expectEqualStrings(patch_file, wm.get("mtp.hyper_connection_mixer.hc_norm.weight").?.string); + try std.testing.expectEqualStrings("model-00001-of-00007.safetensors", wm.get("x").?.string); + // a folder without the patch file opens as before + var other = std.testing.tmpDir(.{}); + defer other.cleanup(); + const empty = try other.dir.realPathFileAlloc(io, ".", a); + defer a.free(empty); + try addPatch(io, empty, &parsed); +} diff --git a/zig/src/families/flashnext/exl3/pack/exl3_format.zig b/zig/src/families/flashnext/exl3/pack/exl3_format.zig new file mode 120000 index 0000000000..66d9fdff5e --- /dev/null +++ b/zig/src/families/flashnext/exl3/pack/exl3_format.zig @@ -0,0 +1 @@ +../../../../core/exl3_format.zig \ No newline at end of file diff --git a/zig/src/families/flashnext/exl3/pack/exl3_place.zig b/zig/src/families/flashnext/exl3/pack/exl3_place.zig new file mode 100644 index 0000000000..554ccf6881 --- /dev/null +++ b/zig/src/families/flashnext/exl3/pack/exl3_place.zig @@ -0,0 +1,309 @@ +//! Header-only home for one tensor name. + +const std = @import("std"); +const scan = @import("exl3_scan.zig"); +const shape = @import("exl3_shape.zig"); +const cache = @import("../model/exl3_cache.zig"); + +pub const Home = enum { + ngram_host, // PLE rows, page cache or SSD, never a VRAM alloc + routed_host, // routed expert bank, host RAM, cache policy decides GPU copies + shared_resident, // shared expert, replicated, stays on device + resident, // dense, router, HC, MTP head, unquantized fp16 linears + centred_norm, // add centredOffset in fp32 at load + mtp_patch, // the mixer file kept outside the index + vision_skip, // vision encoder, not this milestone + plain, +}; + +pub const Tensor = struct { + name: []const u8, + rank: usize, + /// Filename of the shard, used only for the MTP mixer patch. + file: []const u8 = "", +}; + +const centred_names = [_][]const u8{ + "hc_norm", + "q_norm", + "k_norm", + "q_layernorm", + "k_layernorm", + "norm_key", + "norm_query", + "norm_conv", + "pre_fc_norm_embedding", + "pre_fc_norm_hidden", +}; + +const resident_names = [_][]const u8{ + "lm_head", + "q_proj", + "k_proj", + "v_proj", + "o_proj", + "out_proj", + "router", + "hyper_connection", + "in_proj", + "conv1d", + "A_log", + "dt_bias", + "indexer", + "mtp", + // Router logits and the PLE pair. `k_proj` is not a substring of `key_proj`. + "mlp.gate", + "key_proj", + "value_proj", +}; + +comptime { + std.debug.assert(shape.experts == 512); + std.debug.assert(shape.shared_expert == shape.experts); +} + +pub fn home(t: Tensor) Home { + if (scan.isMtpMixerPatch(t.file)) return .mtp_patch; + + const kind = scan.classifyName(t.name, t.rank); + if (kind == .ngram_row) return .ngram_host; + if (isVision(t.name)) return .vision_skip; + if (contains(t.name, "shared_expert")) return .shared_resident; + if (isRouted(t.name)) return .routed_host; + + // Gamma is a rank-1 `.plain` weight. + if (isCentred(t.name)) { + if (kind == .exl3_part) return .plain; + return .centred_norm; + } + + // `mtp` pins the MTP head. + if (contains(t.name, "mtp.fc_embedding")) return .plain; + + if ((kind == .exl3_part or kind == .plain) and isResident(t.name)) return .resident; + return .plain; +} + +/// True when `which` is the cache's host place and both host banks refuse a resident copy. +pub fn hostBank(which: cache.Place) bool { + const budget = cache.Budget{ + .device_bytes = 0, + .resident_bytes = 0, + .kv_bytes = 0, + .routed_expert_bytes = 0, + .ngram_bytes = 0, + }; + const host = cache.refusesResident(.routed) and cache.refusesResident(.ngram) and cache.placeWeights(budget) == .host; + return host and which == .host; +} + +fn contains(hay: []const u8, needle: []const u8) bool { + return std.mem.indexOf(u8, hay, needle) != null; +} + +fn isVision(name: []const u8) bool { + return contains(name, "visual") or std.mem.startsWith(u8, name, "vision") or contains(name, ".vision."); +} + +/// `.experts.` / `experts.` gate, up, and down for ids 0..511. Any `shared` substring stays out. +fn isRouted(name: []const u8) bool { + if (contains(name, "shared")) return false; + return contains(name, ".experts.") or contains(name, "experts."); +} + +/// Token is a full dot-component, or the name's suffix. +pub fn isCentred(name: []const u8) bool { + for (centred_names) |token| { + if (std.mem.endsWith(u8, name, token)) { + if (name.len == token.len or name[name.len - token.len - 1] == '.') return true; + } + var it = std.mem.splitScalar(u8, name, '.'); + while (it.next()) |part| if (std.mem.eql(u8, part, token)) return true; + } + return false; +} + +fn isResident(name: []const u8) bool { + for (resident_names) |token| { + if (!contains(name, token)) continue; + // `mlp.switch_mlp.gate_proj` contains `mlp.gate` and is not the router logits. + if (std.mem.eql(u8, token, "mlp.gate") and contains(name, "gate_proj")) continue; + return true; + } + return false; +} + +test "a tensor name picks one home and the host banks stay off device" { + const ngram = "model.language_model.layers.2.ple.ngram_embedding.shard_0.trellis"; + try std.testing.expectEqual(Home.ngram_host, home(.{ .name = ngram, .rank = 2 })); + + const wide = home(.{ .name = ngram, .rank = 3 }); + try std.testing.expect(wide != .ngram_host); + try std.testing.expectEqual(Home.plain, wide); + + try std.testing.expectEqual(Home.routed_host, home(.{ + .name = "model.layers.0.mlp.experts.0.up_proj.trellis", + .rank = 3, + })); + try std.testing.expectEqual(Home.shared_resident, home(.{ + .name = "model.layers.0.mlp.shared_expert.up_proj.trellis", + .rank = 3, + })); + try std.testing.expectEqual(Home.centred_norm, home(.{ + .name = "model.layers.0.self_attn.q_norm.weight", + .rank = 1, + })); + try std.testing.expectEqual(Home.centred_norm, home(.{ + .name = "model.layers.0.attn_hyper_connection.hc_norm.weight", + .rank = 1, + })); + try std.testing.expectEqual(Home.centred_norm, home(.{ + .name = "language_model.mtp.pre_fc_norm_embedding.weight", + .rank = 1, + })); + try std.testing.expectEqual(Home.plain, home(.{ + .name = "model.layers.0.self_attn.q_norm.trellis", + .rank = 3, + })); + try std.testing.expectEqual(Home.mtp_patch, home(.{ + .name = "weight", + .rank = 1, + .file = "mtp_hyper_connection_mixer_patch.safetensors", + })); + try std.testing.expectEqual(Home.vision_skip, home(.{ + .name = "model.visual.blocks.0.weight", + .rank = 1, + })); + try std.testing.expectEqual(Home.plain, home(.{ + .name = "model.embed_tokens.weight", + .rank = 2, + })); + try std.testing.expectEqual(Home.resident, home(.{ + .name = "model.layers.0.linear_attn.out_proj.weight", + .rank = 2, + })); + try std.testing.expectEqual(Home.plain, home(.{ .name = "readme", .rank = 0 })); + try std.testing.expect(hostBank(.host)); +} + +/// Where the bytes live. +pub const Stay = enum { vram_pinned, vram_learned, host, disk, skip }; + +/// Integer after `.experts.`, or null when the name is not a routed expert. +pub fn expertId(name: []const u8) ?u16 { + const key = ".experts."; + const at = std.mem.indexOf(u8, name, key) orelse return null; + var i = at + key.len; + if (i >= name.len or name[i] < '0' or name[i] > '9') return null; + var n: u32 = 0; + while (i < name.len and name[i] >= '0' and name[i] <= '9') : (i += 1) { + n = n * 10 + (name[i] - '0'); + if (n > std.math.maxInt(u16)) return null; + } + if (i >= name.len or name[i] != '.') return null; + return @intCast(n); +} + +pub fn stay(which: Home, id: ?u16) error{IllegalCut}!Stay { + switch (which) { + .ngram_host => return .disk, + .vision_skip => return .skip, + .routed_host => { + _ = id orelse return error.IllegalCut; + return .vram_learned; + }, + .plain => return .host, + .shared_resident, .resident, .centred_norm, .mtp_patch => return .vram_pinned, + } +} + +pub fn stayOf(t: Tensor) error{IllegalCut}!Stay { + return stay(home(t), expertId(t.name)); +} + +test "the decode path stays pinned, the n-gram on disk, and the routed experts are learned" { + const ngram = "model.language_model.layers.2.ple.ngram_embedding.shard_0.trellis"; + try std.testing.expectEqual(Stay.disk, try stayOf(.{ .name = ngram, .rank = 2 })); + + const pinned = [_][]const u8{ + "model.layers.0.mlp.shared_expert.up_proj.trellis", + "lm_head.weight", + "model.layers.0.mlp.router.weight", + "model.layers.0.mlp_hyper_connection.weight", + "model.layers.0.self_attn.q_proj.weight", + "model.layers.0.self_attn.k_proj.weight", + "model.layers.0.self_attn.v_proj.weight", + "model.layers.0.self_attn.o_proj.weight", + "model.layers.0.linear_attn.in_proj_qkv.weight", + "model.layers.0.linear_attn.out_proj.weight", + "model.layers.0.linear_attn.conv1d.weight", + "model.layers.0.linear_attn.A_log", + "model.layers.0.linear_attn.dt_bias", + "model.language_model.layers.0.mlp.gate.weight", + "model.language_model.layers.1.ple.key_proj.weight", + "model.language_model.layers.1.ple.value_proj.weight", + }; + try std.testing.expectEqual(Stay.host, try stayOf(.{ + .name = "model.embed_tokens.weight", + .rank = 2, + })); + try std.testing.expectEqual(Stay.host, try stayOf(.{ .name = "readme", .rank = 0 })); + for (pinned) |name| { + try std.testing.expectEqual(Stay.vram_pinned, try stayOf(.{ .name = name, .rank = 2 })); + } + try std.testing.expectEqual(Stay.vram_pinned, try stayOf(.{ + .name = "model.layers.0.self_attn.q_norm.weight", + .rank = 1, + })); + try std.testing.expectEqual(Stay.vram_pinned, try stayOf(.{ + .name = "weight", + .rank = 1, + .file = "mtp_hyper_connection_mixer_patch.safetensors", + })); + try std.testing.expectEqual(Stay.skip, try stayOf(.{ + .name = "model.visual.blocks.0.weight", + .rank = 1, + })); + + try std.testing.expectEqual(@as(?u16, 0), expertId("model.layers.0.mlp.experts.0.up_proj.trellis")); + try std.testing.expectEqual(@as(?u16, 255), expertId("model.layers.0.mlp.experts.255.down_proj.trellis")); + try std.testing.expectEqual(@as(?u16, 511), expertId("model.layers.0.mlp.experts.511.gate_proj.trellis")); + try std.testing.expectEqual(@as(?u16, null), expertId("model.layers.0.mlp.shared_expert.up_proj.trellis")); + + const low = "model.layers.0.mlp.experts.0.up_proj.trellis"; + const high = "model.layers.0.mlp.experts.400.up_proj.trellis"; + try std.testing.expectEqual(Stay.vram_learned, try stayOf(.{ .name = low, .rank = 3 })); + try std.testing.expectEqual(Stay.vram_learned, try stayOf(.{ .name = high, .rank = 3 })); + + try std.testing.expectError(error.IllegalCut, stay(.routed_host, null)); + try std.testing.expect(std.mem.indexOf(u8, "key_proj", "k_proj") == null); + try std.testing.expect(std.mem.indexOf(u8, "value_proj", "v_proj") == null); + try std.testing.expect(std.mem.indexOf(u8, "mlp.shared_expert_gate", "mlp.gate") == null); +} + +test "mlp.gate, key_proj, and value_proj stay pinned, and so does mlp.router.weight" { + const gate = "model.language_model.layers.0.mlp.gate.weight"; + const key = "model.language_model.layers.1.ple.key_proj.weight"; + const value = "model.language_model.layers.1.ple.value_proj.weight"; + const router = "model.layers.0.mlp.router.weight"; + const shared_gate = "model.layers.0.mlp.shared_expert_gate.weight"; + for ([_][]const u8{ gate, key, value, router }) |name| { + try std.testing.expectEqual(Home.resident, home(.{ .name = name, .rank = 2 })); + try std.testing.expectEqual(Stay.vram_pinned, try stayOf(.{ .name = name, .rank = 2 })); + } + // `shared_expert` matches before the resident tokens. This is not the 2560×640 expert. + try std.testing.expectEqual(Home.shared_resident, home(.{ .name = shared_gate, .rank = 2 })); + try std.testing.expectEqual(Stay.vram_pinned, try stayOf(.{ .name = shared_gate, .rank = 1 })); + + const switch_gate = "model.layers.0.mlp.switch_mlp.gate_proj.weight"; + try std.testing.expectEqual(Home.plain, home(.{ .name = switch_gate, .rank = 2 })); + try std.testing.expectEqual(Stay.host, try stayOf(.{ .name = switch_gate, .rank = 2 })); + try std.testing.expectEqual(Home.routed_host, home(.{ + .name = "model.layers.0.mlp.experts.0.gate_proj.trellis", + .rank = 3, + })); + try std.testing.expectEqual(Home.plain, home(.{ .name = "mtp.fc_embedding.weight", .rank = 2 })); + try std.testing.expectEqual(Stay.host, try stayOf(.{ .name = "mtp.fc_embedding.weight", .rank = 2 })); + try std.testing.expectEqual(Home.resident, home(.{ .name = "language_model.mtp.fc_hidden.weight", .rank = 2 })); + try std.testing.expectEqual(Stay.vram_pinned, try stayOf(.{ .name = "language_model.mtp.fc_hidden.weight", .rank = 2 })); +} diff --git a/zig/src/families/flashnext/exl3/pack/exl3_qmat.zig b/zig/src/families/flashnext/exl3/pack/exl3_qmat.zig new file mode 100644 index 0000000000..34e4fc3b5b --- /dev/null +++ b/zig/src/families/flashnext/exl3/pack/exl3_qmat.zig @@ -0,0 +1,229 @@ +//! Flash Next EXL3 CPU oracle: the unpacked quantized matrix, its Hadamard transform and the row-parallel helper. + +const std = @import("std"); +const format = @import("exl3_format.zig"); + +// -- parallel for ------------------------------------------------------------------------------------------ + +/// Runs `f(ctx, i)` for i in 0..n over up to `max_threads` threads (the caller included). +pub fn parallelFor(n: usize, max_threads: usize, ctx: anytype, comptime f: anytype) void { + const Ctx = @TypeOf(ctx); + const W = struct { + fn run(next: *std.atomic.Value(usize), count: usize, c: Ctx) void { + while (true) { + const i = next.fetchAdd(1, .monotonic); + if (i >= count) return; + f(c, i); + } + } + }; + var next = std.atomic.Value(usize).init(0); + const cpus = std.Thread.getCpuCount() catch 4; + const want = @min(@min(cpus, max_threads), n); + var handles: [128]std.Thread = undefined; + var spawned: usize = 0; + while (spawned + 1 < want and spawned < handles.len) : (spawned += 1) { + handles[spawned] = std.Thread.spawn(.{}, W.run, .{ &next, n, ctx }) catch break; + } + W.run(&next, n, ctx); + for (handles[0..spawned]) |h| h.join(); +} + +// -- QMat -------------------------------------------------------------------------------------------------- + +/// In-place orthonormal Walsh-Hadamard transform of each 128 block of `x` (Sylvester order, H[i][j] =). +pub fn fwht128(x: []f64) void { + std.debug.assert(x.len % 128 == 0); + const inv = 1.0 / @sqrt(128.0); + var base: usize = 0; + while (base < x.len) : (base += 128) { + const b = x[base..][0..128]; + var h: usize = 1; + while (h < 128) : (h *= 2) { + var i: usize = 0; + while (i < 128) : (i += 2 * h) { + for (i..i + h) |j| { + const lo = b[j]; + const hi = b[j + h]; + b[j] = lo + hi; + b[j + h] = lo - hi; + } + } + } + for (b) |*v| v.* *= inv; + } +} + +pub const QMat = struct { + k: usize, + n: usize, + wq: []f32, + suh: []f32, + svh: []f32, + + /// `threads` parallelises the trellis unpack over tile rows. + pub fn init(gpa: std.mem.Allocator, halves: u16, trellis: []const i16, tiles_k: usize, tiles_n: usize, suh: []const f16, svh: []const f16, threads: usize) !QMat { + const k = tiles_k * 16; + const n = tiles_n * 16; + if (k % 128 != 0 or n % 128 != 0 or suh.len != k or svh.len != n) return error.Length; + if (trellis.len != tiles_k * tiles_n * @as(usize, halves) * 8) return error.Length; + const wq = try gpa.alloc(f32, k * n); + errdefer gpa.free(wq); + const su = try gpa.alloc(f32, k); + errdefer gpa.free(su); + const sv = try gpa.alloc(f32, n); + errdefer gpa.free(sv); + for (su, suh) |*d, s| d.* = @floatCast(s); + for (sv, svh) |*d, s| d.* = @floatCast(s); + var job: UnpackJob = .{ .halves = halves, .trellis = trellis, .tiles_n = tiles_n, .wq = wq, .n = n }; + parallelFor(tiles_k, threads, &job, UnpackJob.row); + return .{ .k = k, .n = n, .wq = wq, .suh = su, .svh = sv }; + } + + pub fn deinit(m: *QMat, gpa: std.mem.Allocator) void { + gpa.free(m.wq); + gpa.free(m.suh); + gpa.free(m.svh); + m.* = undefined; + } + + /// y[0..n] = x[0..k] W. `tmp` is k long. + pub fn apply(m: QMat, x: []const f64, y: []f64, tmp: []f64) void { + std.debug.assert(x.len == m.k and y.len == m.n and tmp.len == m.k); + for (tmp, x, m.suh) |*d, v, s| d.* = v * s; + fwht128(tmp); + @memset(y, 0); + for (0..m.k) |i| { + const u = tmp[i]; + if (u == 0) continue; + const row = m.wq[i * m.n ..][0..m.n]; + for (y, row) |*o, w| o.* += u * @as(f64, w); + } + fwht128(y); + for (y, m.svh) |*o, s| o.* *= s; + } + + /// Same as `apply` with the output columns [lo, hi) only (the Hadamard blocks must be whole: lo. + pub fn applyCols(m: QMat, tmp_h: []const f64, y: []f64, lo: usize, hi: usize) void { + std.debug.assert(lo % 128 == 0 and hi % 128 == 0 and y.len == hi - lo); + @memset(y, 0); + for (0..m.k) |i| { + const u = tmp_h[i]; + if (u == 0) continue; + const row = m.wq[i * m.n + lo ..][0 .. hi - lo]; + for (y, row) |*o, w| o.* += u * @as(f64, w); + } + fwht128(y); + for (y, m.svh[lo..hi]) |*o, s| o.* *= s; + } + + /// The transformed input (x * suh) H_K that `applyCols` consumes. + pub fn prepare(m: QMat, x: []const f64, out: []f64) void { + for (out, x, m.suh) |*d, v, s| d.* = v * s; + fwht128(out); + } +}; + +const UnpackJob = struct { + halves: u16, + trellis: []const i16, + tiles_n: usize, + wq: []f32, + n: usize, + fn row(j: *UnpackJob, tk: usize) void { + var placement: [256]format.Placement = undefined; + format.tilePlacement(&placement); + const w: format.Width = .{ .halves = j.halves }; + const words = w.tileWords(); + for (0..j.tiles_n) |tn| { + const tile = j.trellis[(tk * j.tiles_n + tn) * words ..][0..words]; + var states: [256]u32 = undefined; + format.tileStates(w, tile, &states); + for (0..256) |p| { + const v: f16 = format.codebookValue(.mul1, @intCast(states[p])); + j.wq[(tk * 16 + placement[p].row) * j.n + tn * 16 + placement[p].col] = @floatCast(v); + } + } + } +}; + +test "fwht128 is orthonormal and equals the dense Sylvester matrix" { + var rng = std.Random.DefaultPrng.init(3); + var x: [256]f64 = undefined; + for (&x) |*v| v.* = rng.random().float(f64) * 2 - 1; + var y = x; + fwht128(&y); + for (0..2) |blk| { + for (0..128) |i| { + var s: f64 = 0; + for (0..128) |j| { + const sign: f64 = if (@popCount(i & j) & 1 == 1) -1 else 1; + s += sign * x[blk * 128 + j]; + } + try std.testing.expectApproxEqAbs(s / @sqrt(128.0), y[blk * 128 + i], 1e-12); + } + } + fwht128(&y); + for (x, y) |a, b| try std.testing.expectApproxEqAbs(a, b, 1e-12); +} + +test "QMat.apply equals the dense decoder's dequantize then matvec" { + const gpa = std.testing.allocator; + var rng = std.Random.DefaultPrng.init(17); + const r = rng.random(); + for ([_]u16{ 4, 6, 10 }) |halves| { + const tiles_k: usize = 8; + const tiles_n: usize = 16; + const k = tiles_k * 16; + const n = tiles_n * 16; + const words = @as(usize, halves) * 8; + const trellis = try gpa.alloc(i16, tiles_k * tiles_n * words); + defer gpa.free(trellis); + for (trellis) |*v| v.* = @bitCast(r.int(u16)); + const suh = try gpa.alloc(f16, k); + defer gpa.free(suh); + const svh = try gpa.alloc(f16, n); + defer gpa.free(svh); + for (suh) |*v| v.* = if (r.boolean()) 1.0 else -1.0; + for (svh) |*v| v.* = @floatCast(0.5 + r.float(f32)); + const dense = try gpa.alloc(f64, k * n); + defer gpa.free(dense); + const scratch = try gpa.alloc(f64, 2 * k * n); + defer gpa.free(scratch); + format.dequantize(.{ .halves = halves }, .mul1, trellis, tiles_k, tiles_n, suh, svh, dense, scratch); + var q = try QMat.init(gpa, halves, trellis, tiles_k, tiles_n, suh, svh, 4); + defer q.deinit(gpa); + var x: [k]f64 = undefined; + for (&x) |*v| v.* = r.floatNorm(f64); + var want: [n]f64 = undefined; + for (0..n) |j| { + var s: f64 = 0; + for (0..k) |i| s += x[i] * dense[i * n + j]; + want[j] = s; + } + var got: [n]f64 = undefined; + var tmp: [k]f64 = undefined; + q.apply(&x, &got, &tmp); + for (want, got) |a, b| try std.testing.expectApproxEqAbs(a, b, 1e-9 * (1 + @abs(a))); + // column slices equal the whole + var prep: [k]f64 = undefined; + q.prepare(&x, &prep); + var part: [128]f64 = undefined; + q.applyCols(&prep, &part, 128, 256); + for (part, got[128..256]) |a, b| try std.testing.expectApproxEqAbs(a, b, 1e-12); + } +} + +test "parallelFor visits every index once" { + const Ctx = struct { + hits: []std.atomic.Value(u8), + fn f(c: *@This(), i: usize) void { + _ = c.hits[i].fetchAdd(1, .monotonic); + } + }; + var hits: [1000]std.atomic.Value(u8) = undefined; + for (&hits) |*h| h.* = std.atomic.Value(u8).init(0); + var c: Ctx = .{ .hits = &hits }; + parallelFor(hits.len, 8, &c, Ctx.f); + for (&hits) |*h| try std.testing.expectEqual(@as(u8, 1), h.load(.monotonic)); +} diff --git a/zig/src/families/flashnext/exl3/pack/exl3_scan.zig b/zig/src/families/flashnext/exl3/pack/exl3_scan.zig new file mode 100644 index 0000000000..06b53571d8 --- /dev/null +++ b/zig/src/families/flashnext/exl3/pack/exl3_scan.zig @@ -0,0 +1,206 @@ +//! EXL3 safetensors header scan. + +const std = @import("std"); + +/// Same tags, in the same order, as `exl3_format.Codebook`. +pub const format = struct { + pub const Codebook = enum { inst3, mcg, mul1 }; +}; + +pub const Scales = enum { suh, su, svh, sv }; + +pub const Group = struct { + halves: u16, + codebook: format.Codebook, + k: usize, + n: usize, + in_scales: Scales, + out_scales: Scales, + bias: bool, +}; + +pub const Parts = struct { + trellis_words: usize, + k_tiles: usize, + n_tiles: usize, + suh_len: ?usize = null, + su_len: ?usize = null, + svh_len: ?usize = null, + sv_len: ?usize = null, + bias: bool = false, + mcg: bool = false, + mul1: bool = false, +}; + +pub const Kind = enum { exl3_part, ngram_row, plain }; + +const suffixes = [_][]const u8{ "trellis", "suh", "su", "svh", "sv", "bias", "mcg", "mul1" }; + +/// One EXL3 linear group from its header parts. Half-bit widths need the mul1 marker. +pub fn parseGroup(parts: Parts) error{NotExl3}!Group { + if (parts.trellis_words == 0 or parts.trellis_words % 8 != 0) return error.NotExl3; + const halves_usize = parts.trellis_words / 8; + if (halves_usize < 2 or halves_usize > 16) return error.NotExl3; + const halves: u16 = @intCast(halves_usize); + if (parts.mcg and parts.mul1) return error.NotExl3; + const codebook: format.Codebook = if (parts.mul1) .mul1 else if (parts.mcg) .mcg else .inst3; + if (halves % 2 == 1 and codebook != .mul1) return error.NotExl3; + + const k = parts.k_tiles * 16; + const n = parts.n_tiles * 16; + if (k == 0 or n == 0 or k % 128 != 0 or n % 128 != 0) return error.NotExl3; + + const in_scales: Scales = if (parts.suh_len != null and parts.su_len == null) + .suh + else if (parts.su_len != null and parts.suh_len == null) + .su + else + return error.NotExl3; + if (in_scales == .suh and parts.suh_len.? != k) return error.NotExl3; + if (in_scales == .su and parts.su_len.? != k / 16) return error.NotExl3; + + const out_scales: Scales = if (parts.svh_len != null and parts.sv_len == null) + .svh + else if (parts.sv_len != null and parts.svh_len == null) + .sv + else + return error.NotExl3; + if (out_scales == .svh and parts.svh_len.? != n) return error.NotExl3; + if (out_scales == .sv and parts.sv_len.? != n / 16) return error.NotExl3; + + return .{ + .halves = halves, + .codebook = codebook, + .k = k, + .n = n, + .in_scales = in_scales, + .out_scales = out_scales, + .bias = parts.bias, + }; +} + +/// An n-gram shard is a rank-2 `ngram_embedding` `.trellis`, not an EXL3 tile tensor. +pub fn classifyName(name: []const u8, rank: usize) Kind { + if (rank == 2 and std.mem.indexOf(u8, name, "ngram_embedding") != null and std.mem.endsWith(u8, name, ".trellis")) { + return .ngram_row; + } + const dot = std.mem.lastIndexOfScalar(u8, name, '.') orelse return .plain; + const suffix = name[dot + 1 ..]; + for (suffixes) |part| if (std.mem.eql(u8, suffix, part)) return .exl3_part; + return .plain; +} + +/// The MTP mixer's extra file, the one turboderp keeps outside the index. +pub fn isMtpMixerPatch(filename: []const u8) bool { + return std.mem.eql(u8, filename, "mtp_hyper_connection_mixer_patch.safetensors"); +} + +fn codebookName(name: []const u8) error{QuantRefuse}!void { + if (std.mem.eql(u8, name, "mul1") or std.mem.eql(u8, name, "mcg") or std.mem.eql(u8, name, "3inst") or std.mem.eql(u8, name, "inst3")) return; + return error.QuantRefuse; +} + +fn integerBits(bits: f32) error{QuantRefuse}!void { + if (!std.math.isFinite(bits) or bits < 1 or bits > 8) return error.QuantRefuse; + if (bits != @round(bits)) return error.QuantRefuse; +} + +/// Admit a pack from config fields alone. Average bits may be fractional. Head and MTP bits may not. +pub fn admit(codebook: []const u8, average_bits: f32, head_bits: ?f32, mtp_bits: ?f32) error{QuantRefuse}!void { + try codebookName(codebook); + if (!std.math.isFinite(average_bits) or average_bits < 1 or average_bits > 8) return error.QuantRefuse; + if (head_bits) |bits| try integerBits(bits); + if (mtp_bits) |bits| try integerBits(bits); +} + +/// 1 when centred norms are stored as gamma-1, 0 when stored as gamma. +pub fn centredOffset(means: []const f32) error{NormAmbiguous}!f32 { + if (means.len == 0) return 1; + if (means.len > 64) return error.NormAmbiguous; + var copy: [64]f32 = undefined; + @memcpy(copy[0..means.len], means); + var i: usize = 1; + while (i < means.len) : (i += 1) { + const v = copy[i]; + var j = i; + while (j > 0 and copy[j - 1] > v) : (j -= 1) copy[j] = copy[j - 1]; + copy[j] = v; + } + const median = copy[means.len / 2]; + var above: usize = 0; + for (means) |m| if (m > 0.5) { + above += 1; + }; + const frac = @as(f32, @floatFromInt(above)) / @as(f32, @floatFromInt(means.len)); + const around_zero = frac <= 0.1 and median >= -0.5 and median <= 0.25; + const around_one = frac >= 0.9 and median >= 0.75 and median <= 1.5; + if (around_zero == around_one) return error.NormAmbiguous; + return if (around_zero) 1 else 0; +} + +test "a 3-bit mul1 gate group parses and a half-bit group needs mul1" { + const gate = try parseGroup(.{ + .trellis_words = 48, + .k_tiles = 160, + .n_tiles = 40, + .suh_len = 2560, + .svh_len = 640, + .mul1 = true, + }); + try std.testing.expectEqual(@as(u16, 6), gate.halves); + try std.testing.expectEqual(format.Codebook.mul1, gate.codebook); + try std.testing.expectEqual(@as(usize, 2560), gate.k); + try std.testing.expectEqual(@as(usize, 640), gate.n); + try std.testing.expectEqual(Scales.suh, gate.in_scales); + try std.testing.expectEqual(Scales.svh, gate.out_scales); + + const half = Parts{ + .trellis_words = 40, + .k_tiles = 160, + .n_tiles = 40, + .suh_len = 2560, + .svh_len = 640, + .mul1 = true, + }; + try std.testing.expectEqual(@as(u16, 5), (try parseGroup(half)).halves); + var no_marker = half; + no_marker.mul1 = false; + try std.testing.expectError(error.NotExl3, parseGroup(no_marker)); + var both = half; + both.mcg = true; + try std.testing.expectError(error.NotExl3, parseGroup(both)); +} + +test "K must be a multiple of 128" { + try std.testing.expectError(error.NotExl3, parseGroup(.{ + .trellis_words = 48, + .k_tiles = 20, + .n_tiles = 40, + .suh_len = 320, + .svh_len = 640, + .mul1 = true, + })); +} + +test "names sort into n-gram rows, EXL3 parts and plain tensors" { + try std.testing.expectEqual(Kind.exl3_part, classifyName("model.language_model.layers.0.mlp.experts.0.gate_proj.trellis", 3)); + try std.testing.expectEqual(Kind.ngram_row, classifyName("ple.ngram_embedding.shard_3.trellis", 2)); + try std.testing.expectEqual(Kind.exl3_part, classifyName("ple.ngram_embedding.shard_3.trellis", 3)); + try std.testing.expectEqual(Kind.plain, classifyName("model.embed_tokens.weight", 2)); + try std.testing.expect(isMtpMixerPatch("mtp_hyper_connection_mixer_patch.safetensors")); + try std.testing.expect(!isMtpMixerPatch("model.safetensors")); +} + +test "the 3.05 config is admitted and a fractional head width is not" { + try admit("mul1", 3.05, 5, 3); + try admit("3inst", 4, null, null); + try std.testing.expectError(error.QuantRefuse, admit("mul1", 3.05, 3.05, null)); + try std.testing.expectError(error.QuantRefuse, admit("exl2", 3, null, null)); +} + +test "centred norms are gamma-1 when the means sit near zero" { + try std.testing.expectEqual(@as(f32, 1), try centredOffset(&.{ -0.02, 0.01, 0.0, 0.04 })); + try std.testing.expectEqual(@as(f32, 0), try centredOffset(&.{ 1.0, 0.98, 1.02, 1.01 })); + try std.testing.expectError(error.NormAmbiguous, centredOffset(&.{ 0.4, 0.9 })); + try std.testing.expectEqual(@as(f32, 1), try centredOffset(&.{})); +} diff --git a/zig/src/families/flashnext/exl3/pack/exl3_shape.zig b/zig/src/families/flashnext/exl3/pack/exl3_shape.zig new file mode 100644 index 0000000000..ec97e5d1df --- /dev/null +++ b/zig/src/families/flashnext/exl3/pack/exl3_shape.zig @@ -0,0 +1,66 @@ +//! Published text shape of turboderp/Qwen3.8-Flash-Next-exl3, revision 3.05bpw_h5_ng5. + +const std = @import("std"); + +pub const model = "turboderp/Qwen3.8-Flash-Next-exl3"; +pub const revision = "3.05bpw_h5_ng5"; + +pub const hidden: usize = 2560; +pub const layers: usize = 48; +pub const vocab: usize = 248320; +pub const heads: usize = 24; +pub const kv_heads: usize = 2; +pub const head_dim: usize = 256; +pub const experts: usize = 512; +pub const topk: usize = 10; +pub const expert_intermediate: usize = 640; +pub const shared_intermediate: usize = 640; +/// Routed ids are 0..511. The shared expert is the next id in the same table. +pub const shared_expert: usize = experts; +pub const hc_count: usize = 4; +pub const hc_lowrank: usize = 320; + +pub const ngram: usize = 3; +pub const ngram_heads_per: usize = 8; +pub const ngram_dh: usize = 160; +/// Branch name `ng5`: the n-gram row codec is 5 bits. Confirmed by words-per-row once a shard header is read. +pub const ngram_bits: usize = 5; +pub const ngram_shards: usize = 128; +pub const ngram_vocab_base: usize = 20_000_000; + +pub const quant_method = "exl3"; +pub const quant_version = "1.4.4"; +pub const average_bits: f32 = 3.05; +pub const head_bits: usize = 5; +pub const mtp_bits: usize = 3; +pub const codebook = "mul1"; + +pub const tile: usize = 16; +pub const hadamard: usize = 128; + +pub const linear_layers_per_cycle: usize = 3; +pub const full_attention_interval: usize = 4; + +/// Heads in one n-gram lookup: (ngram - 1) * heads_per_ngram. +pub fn ngramHeadCount() usize { + return (ngram - 1) * ngram_heads_per; +} + +/// q_proj outputs. 24 * 256 = 6144. +pub fn qOut() usize { + return heads * head_dim; +} + +/// One attention layer's kv width, keys or values. 2 * 256 = 512. +pub fn kvOut() usize { + return kv_heads * head_dim; +} + +test "the published 3.05 shape is the one the budget is written against" { + try std.testing.expectEqual(@as(usize, 16), ngramHeadCount()); + try std.testing.expectEqual(@as(usize, 6144), qOut()); + try std.testing.expectEqual(@as(usize, 512), kvOut()); + try std.testing.expect(hidden % hadamard == 0); + try std.testing.expect(expert_intermediate % hadamard == 0); + try std.testing.expect(qOut() % hadamard == 0); +} diff --git a/zig/src/families/flashnext/exl3/pack/exl3_startup_memory.zig b/zig/src/families/flashnext/exl3/pack/exl3_startup_memory.zig new file mode 100644 index 0000000000..d924759dec --- /dev/null +++ b/zig/src/families/flashnext/exl3/pack/exl3_startup_memory.zig @@ -0,0 +1,72 @@ +const std = @import("std"); + +pub const Reservation = struct { + stream_bytes: u64, + streams: u64, + + pub fn total(self: Reservation) !u64 { + return std.math.mul(u64, self.stream_bytes, self.streams); + } +}; + +pub const Ledger = struct { + persistent: u64, + scratch: u64, + pool_scratch: u64, + expert_pools: u64, + reservation: Reservation, + + pub fn total(self: Ledger) !u64 { + var bytes = try self.reservation.total(); + inline for (.{ self.persistent, self.scratch, self.pool_scratch, self.expert_pools }) |part| bytes = try std.math.add(u64, bytes, part); + return bytes; + } +}; + +pub fn mebibytes(text: []const u8) !u64 { + return std.math.mul(u64, try std.fmt.parseInt(u64, text, 10), 1 << 20); +} + +pub fn available(free: u64, reserve: u64, future: u64) !u64 { + const held = try std.math.add(u64, reserve, future); + if (held > free) return error.InsufficientStartupMemory; + return free - held; +} + +pub fn poolBytes(lines: u64, line_bytes: u64, layers: u64, fixed: u64) !u64 { + return std.math.add(u64, try std.math.mul(u64, try std.math.mul(u64, lines, line_bytes), layers), fixed); +} + +test "BF16 c1 c8 preserve the requested stream bytes" { + const bf16_stream: u64 = 262144 * 256 * 2; + for ([_]u64{ 1, 8 }) |streams| { + const reservation: Reservation = .{ .stream_bytes = bf16_stream, .streams = streams }; + try std.testing.expectEqual(bf16_stream * streams, try reservation.total()); + const free = try available(16 << 30, 1 << 30, try reservation.total()); + try std.testing.expectEqual(@as(u64, 16 << 30), free + (1 << 30) + try reservation.total()); + } +} + +test "startup accounting fails closed on overflow and underfunded explicit pools" { + try std.testing.expectError(error.Overflow, (Reservation{ .stream_bytes = std.math.maxInt(u64), .streams = 8 }).total()); + try std.testing.expectError(error.Overflow, mebibytes("18446744073709551615")); + try std.testing.expectError(error.Overflow, poolBytes(std.math.maxInt(u64), 2, 49, 0)); + try std.testing.expectError(error.InsufficientStartupMemory, available(100, 20, try poolBytes(10, 1, 8, 1))); + try std.testing.expectError(error.Overflow, available(100, std.math.maxInt(u64), 1)); +} + +test "ledger includes all allocations after the pool checkpoint" { + const ledger: Ledger = .{ .persistent = 100, .scratch = 20, .pool_scratch = 40, .expert_pools = 50, .reservation = .{ .stream_bytes = 10, .streams = 8 } }; + try std.testing.expectEqual(@as(u64, 290), try ledger.total()); +} + +test "pool setup may consume its allocator margin but not lane or server reservations" { + const server: u64 = 100; + const margin: u64 = 25; + const future: u64 = 80; + const planner: u64 = 10; + const planned_free = server + margin + future + planner; + try std.testing.expectEqual(margin, try available(planned_free, server, future + planner)); + try std.testing.expectEqual(@as(u64, 0), try available(planned_free - margin, server, future + planner)); + try std.testing.expectError(error.InsufficientStartupMemory, available(planned_free - margin - 1, server, future + planner)); +} diff --git a/zig/src/families/flashnext/exl3/pack/fixtures/exl3_config.json b/zig/src/families/flashnext/exl3/pack/fixtures/exl3_config.json new file mode 100644 index 0000000000..ea4b025ebf --- /dev/null +++ b/zig/src/families/flashnext/exl3/pack/fixtures/exl3_config.json @@ -0,0 +1,168 @@ +{ + "architectures": [ + "Qwen4ExpForConditionalGeneration" + ], + "image_token_id": 248056, + "language_model_only": false, + "model_type": "qwen4_exp", + "text_config": { + "attention_bias": false, + "attention_dropout": 0.0, + "bos_token_id": 248044, + "dtype": "bfloat16", + "eos_token_id": 248044, + "full_attention_interval": 4, + "hc_count": 4, + "hc_lowrank": 320, + "head_dim": 256, + "heads_per_ngram": 8, + "hidden_act": "silu", + "hidden_size": 2560, + "indexer_budget": 2048, + "indexer_compress_ratio": 4, + "indexer_head_dim": 128, + "indexer_kv_heads": 1, + "indexer_n_heads": 4, + "initializer_range": 0.02, + "layer_types": [ + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention" + ], + "linear_conv_kernel_dim": 4, + "linear_key_head_dim": 128, + "linear_num_key_heads": 16, + "linear_num_value_heads": 48, + "linear_value_head_dim": 128, + "make_ngram_vocab_size_divisible_by": 128, + "mamba_ssm_dtype": "float32", + "max_position_embeddings": 262144, + "model_type": "qwen4_exp_text", + "moe_intermediate_size": 640, + "mtp": { + "hybrid": true, + "layer_types": [ + "full_attention" + ], + "mtp_use_hidden_state_from_layer": null, + "num_hidden_layers": 1, + "rope_theta": 10000000 + }, + "mtp_num_hidden_layers": 1, + "mtp_use_dedicated_embeddings": false, + "ngram_size": 3, + "ngram_vocab_size_base": 20000000, + "num_attention_heads": 24, + "num_experts": 512, + "num_experts_per_tok": 10, + "num_hidden_layers": 48, + "num_key_value_heads": 2, + "output_gate_type": "sigmoid", + "output_router_logits": false, + "pad_token_id": null, + "partial_rotary_factor": 0.25, + "ple_conv_kernel_size": 4, + "ple_embed_dim": 2560, + "ple_layer_ids": [ + 2 + ], + "rms_norm_eps": 1e-06, + "rope_parameters": { + "mrope_interleaved": true, + "mrope_section": [ + 11, + 11, + 10 + ], + "partial_rotary_factor": 0.25, + "rope_theta": 10000000, + "rope_type": "default" + }, + "router_aux_loss_coef": 0.001, + "shared_expert_intermediate_size": 640, + "split_ngram_parts": 128, + "tie_word_embeddings": false, + "use_cache": true, + "vocab_size": 248320 + }, + "tie_word_embeddings": false, + "transformers_version": "5.8.0.dev0", + "video_token_id": 248057, + "vision_config": { + "deepstack_visual_indexes": [], + "depth": 27, + "hidden_act": "gelu_pytorch_tanh", + "hidden_size": 1152, + "in_channels": 3, + "initializer_range": 0.02, + "intermediate_size": 4304, + "model_type": "qwen4_exp", + "num_heads": 16, + "num_position_embeddings": 2304, + "out_hidden_size": 2560, + "patch_size": 16, + "spatial_merge_size": 2, + "temporal_patch_size": 2 + }, + "vision_end_token_id": 248054, + "vision_start_token_id": 248053, + "quantization_config": { + "quant_method": "exl3", + "version": "1.4.4", + "bits": 3.05, + "head_bits": 5, + "calibration": { + "rows": 250, + "cols": 2048 + }, + "out_scales": "always", + "codebook": "mul1", + "vision_bits": 5, + "mtp_bits": 3 + } +} \ No newline at end of file diff --git a/zig/src/families/flashnext/exl3/tp/cuda_exl3_fork.zig b/zig/src/families/flashnext/exl3/tp/cuda_exl3_fork.zig new file mode 100644 index 0000000000..4f296e551d --- /dev/null +++ b/zig/src/families/flashnext/exl3/tp/cuda_exl3_fork.zig @@ -0,0 +1,126 @@ +//! A multi-part EXL3 projection's parts on parallel streams (a DeltaNet layer's q|k|v, z, b, a; attention's q, k, v). + +const std = @import("std"); +const cuda = @import("cuda"); +const load = @import("../pack/cuda_exl3_load.zig"); + +pub const branches = 3; +/// Rows a forked projection takes at most (decode rounds and verify windows; prompts keep one stream). +pub const max_rows: usize = 16; + +pub const Fork = struct { + side: [branches]cuda.Stream, + dense: [branches]cuda.exl3.Dense, + tmp: [branches]cuda.DeviceBuffer, + ev_fork: cuda.Event, + ev_join: [branches]cuda.Event, + tmp_bytes: usize, +}; + +pub fn enabled() bool { + return force orelse true; +} + +/// A test's override of the fork. +pub var force: ?bool = null; + +/// Ops.reserve's hook: the branches' streams, kernels and scratch for `o`'s widest projection, before any pass. +pub fn reserve(o: *load.Ops) !void { + if (o.max_k == 0 or o.fork != null) return; + const f = try o.gpa.create(Fork); + errdefer o.gpa.destroy(f); + f.tmp_bytes = max_rows * @max(o.max_part_n, 128) * 2; + var made: usize = 0; + errdefer for (0..made) |i| { + f.dense[i].unload(); + f.side[i].deinit(); + f.tmp[i].free(); + f.ev_join[i].deinit(); + }; + f.ev_fork = try cuda.Event.init(o.d, false); + errdefer f.ev_fork.deinit(); + for (0..branches) |i| { + var side = try cuda.Stream.init(o.d, true); + errdefer side.deinit(); + var dense = try cuda.exl3.Dense.load(o.d); + errdefer dense.unload(); + var tmp = try cuda.DeviceBuffer.alloc(o.d, f.tmp_bytes); + errdefer tmp.free(); + var join = try cuda.Event.init(o.d, false); + errdefer join.deinit(); + dense.xh = try cuda.DeviceBuffer.alloc(o.d, 128 * o.max_k * 2); + if (o.max_z > 0) dense.z = try cuda.DeviceBuffer.alloc(o.d, o.max_z); + dense.counters = try cuda.DeviceBuffer.alloc(o.d, o.max_counters); + f.side[i] = side; + f.dense[i] = dense; + f.tmp[i] = tmp; + f.ev_join[i] = join; + made += 1; + } + o.fork = f; +} + +pub fn release(o: *load.Ops) void { + const f = o.fork orelse return; + o.fork = null; + for (0..branches) |i| { + f.side[i].synchronize() catch {}; + f.dense[i].unload(); + f.side[i].deinit(); + f.tmp[i].free(); + f.ev_join[i].deinit(); + } + f.ev_fork.deinit(); + o.gpa.destroy(f); +} + +fn copy2d(s: cuda.Stream, dst: u64, dpitch: usize, src: u64, spitch: usize, width: usize, height: usize) !void { + const m: cuda.abi.Memcpy2D = .{ .src = src, .src_pitch = spitch, .dst = dst, .dst_pitch = dpitch, .width = width, .height = height }; + try s.d.check(s.d.api.cuMemcpy2DAsync_v2(&m, s.handle), "cuMemcpy2DAsync"); +} + +/// Pass.lane()'s parts loop (bf16 out, sk 1) with the EXL3 parts after the first on side streams; false when the +pub fn lane(p: anytype, x: u64, e: *const load.Proj, out: u64, m: usize, comptime runMain: anytype) !bool { + if (!enabled() or m > max_rows) return false; + var n_exl3: usize = 0; + for (e.parts) |pt| n_exl3 += @intFromBool(pt.kind == .exl3); + if (n_exl3 < 2) return false; + const f = e.ops.fork orelse return false; + for (e.parts) |pt| if (pt.kind == .exl3 and m * pt.n * 2 > f.tmp_bytes and pt.n != e.n and m != 1) return false; + const s = p.stream; + try f.ev_fork.record(s); + var b: usize = 0; + var first = true; + for (e.parts) |pt| { + const w = switch (pt.kind) { + .exl3 => |w| w, + .native => { + try runMain(p, e, pt, x, m, out); + continue; + }, + }; + if (first or b == branches) { + first = false; + try runMain(p, e, pt, x, m, out); + continue; + } + const st = f.side[b]; + try st.wait(f.ev_fork); + const de = &f.dense[b]; + if (pt.n == e.n or m == 1) { + try de.decode(st, w, x, .bf16, m, out + pt.col * 2, .bf16); + } else { + const tmp = f.tmp[b].ptr; + try de.decode(st, w, x, .bf16, m, tmp, .bf16); + try copy2d(st, out + pt.col * 2, e.n * 2, tmp, pt.n * 2, pt.n * 2, m); + } + try f.ev_join[b].record(st); + b += 1; + } + for (f.ev_join[0..b]) |ev| try s.wait(ev); + return true; +} + +test { + _ = @import("cuda_exl3_fork_test.zig"); +} diff --git a/zig/src/families/flashnext/exl3/tp/cuda_exl3_fork_test.zig b/zig/src/families/flashnext/exl3/tp/cuda_exl3_fork_test.zig new file mode 100644 index 0000000000..e690f9e32e --- /dev/null +++ b/zig/src/families/flashnext/exl3/tp/cuda_exl3_fork_test.zig @@ -0,0 +1,105 @@ +const std = @import("std"); +const cuda = @import("cuda"); +const load = @import("../pack/cuda_exl3_load.zig"); +const fork = @import("cuda_exl3_fork.zig"); + +const Fake = struct { + var streams: usize = 0; + var modules: usize = 0; + var events: usize = 0; + var buffers: usize = 0; + fn sync(_: cuda.abi.Stream) callconv(.c) c_int { + return 0; + } + fn stream(_: cuda.abi.Stream) callconv(.c) c_int { + streams += 1; + return 0; + } + fn module(_: cuda.abi.Module) callconv(.c) c_int { + modules += 1; + return 0; + } + fn event(_: cuda.abi.Event) callconv(.c) c_int { + events += 1; + return 0; + } + fn alloc(out: *u64, _: usize) callconv(.c) c_int { + out.* = 0x1000; + return 0; + } + fn buffer(_: u64) callconv(.c) c_int { + buffers += 1; + return 0; + } + fn createStream(out: *cuda.abi.Stream, _: c_uint) callconv(.c) c_int { + out.* = null; + return 0; + } + fn createEvent(out: *cuda.abi.Event, _: c_uint) callconv(.c) c_int { + out.* = null; + return 0; + } + fn driver() cuda.Driver { + streams = 0; + modules = 0; + events = 0; + buffers = 0; + var d: cuda.Driver = undefined; + d.api.cuStreamCreate = createStream; + d.api.cuEventCreate = createEvent; + d.api.cuStreamSynchronize = sync; + d.api.cuStreamDestroy_v2 = stream; + d.api.cuModuleUnload = module; + d.api.cuEventDestroy_v2 = event; + d.api.cuMemAlloc_v2 = alloc; + d.api.cuMemFree_v2 = buffer; + return d; + } +}; + +fn attach(o: *load.Ops) !*fork.Fork { + const f = try o.gpa.create(fork.Fork); + f.tmp_bytes = 64; + f.ev_fork = .{ .d = o.d, .handle = null }; + for (0..fork.branches) |i| { + f.side[i] = .{ .d = o.d, .handle = null }; + f.dense[i] = .{ .d = o.d, .lin = .{ .mod = .{ .d = o.d, .handle = null }, .fns = undefined, .rot_in = undefined } }; + f.tmp[i] = try cuda.DeviceBuffer.alloc(o.d, 64); + f.ev_join[i] = .{ .d = o.d, .handle = null }; + } + o.fork = f; + return f; +} + +test "exl3: projection forks retire with their Ops owner and preserve other owners" { + var d = Fake.driver(); + const left = try load.Ops.create(std.testing.allocator, &d); + const right = try load.Ops.create(std.testing.allocator, &d); + const l = try attach(left); + const r = try attach(right); + _ = left.retain(); + left.release(); + try std.testing.expect(left.fork == l and right.fork == r); + try std.testing.expectEqual(@as(usize, 0), Fake.modules); + left.release(); + try std.testing.expect(right.fork == r); + try std.testing.expectEqual(fork.branches, Fake.modules); + right.release(); + try std.testing.expectEqual(2 * fork.branches, Fake.streams); + try std.testing.expectEqual(2 * fork.branches, Fake.modules); + try std.testing.expectEqual(2 * (fork.branches + 1), Fake.events); + try std.testing.expectEqual(2 * fork.branches, Fake.buffers); +} + +test "exl3: fork reservation failure releases the current branch before publication" { + if (cuda.kernels.available) return error.SkipZigTest; + var d = Fake.driver(); + const o = try load.Ops.create(std.testing.allocator, &d); + defer o.release(); + o.max_k = 128; + try std.testing.expectError(error.BuiltWithoutKernels, fork.reserve(o)); + try std.testing.expect(o.fork == null); + try std.testing.expectEqual(@as(usize, 1), Fake.streams); + try std.testing.expectEqual(@as(usize, 1), Fake.events); + try std.testing.expectEqual(@as(usize, 0), Fake.modules); +} diff --git a/zig/src/native/cuda.zig b/zig/src/native/cuda.zig index ab6a6e9443..e6b5d9c796 100644 --- a/zig/src/native/cuda.zig +++ b/zig/src/native/cuda.zig @@ -358,7 +358,8 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O low = try std.fmt.allocPrint(a, " on rank {d}, the smallest room of the ranks", .{pr.rank}); }; const free = loaded.free_streams; - const grow = budget.growRoom(room, !@hasField(F.Options, "lanes")); + const preallocates_lanes = if (@hasDecl(F, "preallocates_lanes")) F.preallocates_lanes else @hasField(F.Options, "lanes"); + const grow = budget.growRoom(room, !preallocates_lanes); const wrapped = kv_pages.wrapped(F, loaded, free, served.stream_bytes, room); const layout = wrapped.layout; var paged: ?kv_pages.Plan = null; @@ -412,7 +413,7 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O // rules and budget from the family's snap_* declarations (cuda_snaps.zig) const image: ?*const pc.rows.Image = if (F == nemotron.native) &cuda_disk.image else if (comptime cuda_disk.keptDisk(F)) |K| &K.image else null; const copies: pc.Snapshots = .{ .ptr = loaded.target, .vtable = &snaps.Snaps(F, image).vtable }; - store.* = pc.Store.init(gpa, copies, snaps.passRules(F, @import("cuda_startup.zig").rows(F, loaded)), if (@hasDecl(F, "snap_slots")) cache else try snaps.cacheBudget(F, loaded.target, cache)); + store.* = pc.Store.init(gpa, copies, snaps.passRules(F, @import("cuda_startup.zig").rows(F, loaded), snaps.loadedGrid(loaded)), if (@hasDecl(F, "snap_slots")) cache else try snaps.cacheBudget(F, loaded.target, cache)); errdefer { const disk = store.disk; store.deinit(); diff --git a/zig/src/native/cuda_snaps.zig b/zig/src/native/cuda_snaps.zig index 35a5785f1a..eda5ba048f 100644 --- a/zig/src/native/cuda_snaps.zig +++ b/zig/src/native/cuda_snaps.zig @@ -15,15 +15,24 @@ pub fn cacheRules(comptime F: type) api.prompt_cache.Rules { return rules; } -/// Native CUDA families own their grid, independently of the server's template chunk starts. -pub fn passRules(comptime F: type, prompt_rows: u32) api.prompt_cache.Rules { +/// Native CUDA families own their grid; a loaded `snap_grid` (the prefill step) keeps states on chunk starts only. +pub fn passRules(comptime F: type, prompt_rows: u32, snap_grid: u32) api.prompt_cache.Rules { var rules = cacheRules(F); rules.chunk = prompt_rows; if (@hasDecl(F, "snap_repeat_min")) rules.repeat_min = if (std.c.getenv("TF_FLASHNEXT_REPEAT_CACHE")) |v| (if (std.mem.eql(u8, std.mem.span(v), "0")) 0 else F.snap_repeat_min) else F.snap_repeat_min; rules.chunk_starts = false; + if (snap_grid > 0) { + rules.planned = true; + rules.grid = snap_grid; + } return rules; } +/// The grid a loaded family declares (0: any row). +pub fn loadedGrid(loaded: anytype) u32 { + return if (@hasField(@TypeOf(loaded), "snap_grid")) loaded.snap_grid else 0; +} + /// The cache's bytes less what the family sets aside for its states (`snap_ready`, as Nemotron's spare buffers). pub fn cacheBudget(comptime F: type, target: anytype, cache: u64) !u64 { return if (@hasDecl(F, "snap_ready")) cache - try F.snap_ready(target, cache) else cache; @@ -91,3 +100,34 @@ test "a family with snap_grid keeps on its chunk starts, any other at any row as test { _ = @import("cuda_snaps_test.zig"); } + +test "a family that names its prefill step keeps marks on chunk starts, so a resumed pass cuts as a fresh one" { + const Exl3 = struct {}; + const Stub = struct { + fn bytes(_: *anyopaque, _: u32) u64 { + return 0; + } + fn save(_: *anyopaque, _: ?*anyopaque, _: u32) anyerror!api.prompt_cache.Saved { + return error.Unused; + } + fn restore(_: *anyopaque, _: ?*anyopaque, _: api.prompt_cache.Saved) anyerror!void {} + fn drop(_: *anyopaque, _: api.prompt_cache.Saved) void {} + }; + var unused: u8 = 0; + var arena = std.heap.ArenaAllocator.init(std.testing.allocator); + defer arena.deinit(); + const prompt = try arena.allocator().alloc(u32, 12780); + for (prompt, 0..) |*t, i| t.* = @intCast(i); + const vtable: api.prompt_cache.Snapshots.VTable = .{ .bytes = Stub.bytes, .save = Stub.save, .restore = Stub.restore, .drop = Stub.drop }; + const copies: api.prompt_cache.Snapshots = .{ .ptr = &unused, .vtable = &vtable }; + // any row: the history mark cuts the cold pass's last chunk to 7 rows, a different kernel path than the fresh pass + var any = api.prompt_cache.Store.init(std.testing.allocator, copies, passRules(Exl3, 4096, 0), 1 << 30); + defer any.deinit(); + const loose = try any.marks(arena.allocator(), prompt, 0, 12773, &.{}, &.{}, &.{}); + try std.testing.expect(std.mem.indexOfScalar(u32, loose, 12773) != null); + var grid = api.prompt_cache.Store.init(std.testing.allocator, copies, passRules(Exl3, 4096, 4096), 1 << 30); + defer grid.deinit(); + const exact = try grid.marks(arena.allocator(), prompt, 0, 12773, &.{}, &.{}, &.{}); + try std.testing.expect(exact.len > 0); + for (exact) |m| try std.testing.expectEqual(@as(u32, 0), m % 4096); +} diff --git a/zig/src/native/cuda_startup_test.zig b/zig/src/native/cuda_startup_test.zig index 1a7d265a81..d9aea9735b 100644 --- a/zig/src/native/cuda_startup_test.zig +++ b/zig/src/native/cuda_startup_test.zig @@ -99,7 +99,17 @@ test "CUDA startup: every other registered family retains its declared pass size test "CUDA startup: Flash Next cache cuts use the loaded cold pass size" { var loaded: flash.native.Loaded = undefined; loaded.prompt_rows = 4096; - const rules = @import("cuda_snaps.zig").passRules(flash.native, startup.rows(flash.native, loaded)); + loaded.snap_grid = 0; + const rules = @import("cuda_snaps.zig").passRules(flash.native, startup.rows(flash.native, loaded), @import("cuda_snaps.zig").loadedGrid(loaded)); try std.testing.expectEqual(@as(u32, 4096), rules.chunk); - try std.testing.expect(!rules.chunk_starts and rules.lookahead == 1); + try std.testing.expect(!rules.chunk_starts and rules.lookahead == 1 and !rules.planned); +} + +test "CUDA startup: an EXL3 Flash Next keeps and resumes states on its prefill step only" { + var loaded: flash.native.Loaded = undefined; + loaded.prompt_rows = 4096; + loaded.snap_grid = 4096; + const snaps = @import("cuda_snaps.zig"); + const rules = snaps.passRules(flash.native, startup.rows(flash.native, loaded), snaps.loadedGrid(loaded)); + try std.testing.expect(rules.planned and rules.grid == 4096 and !rules.chunk_starts); } diff --git a/zig/tests/cuda/contract_limits.zig b/zig/tests/cuda/contract_limits.zig index 2a16afa5b6..48bb601d76 100644 --- a/zig/tests/cuda/contract_limits.zig +++ b/zig/tests/cuda/contract_limits.zig @@ -26,6 +26,8 @@ pub fn image(name: []const u8) error{UnknownContractImage}![]const u8 { inline for (.{ "qwen4_projection", "qwen4_hc", "qwen4_gdn", "qwen4_ple" }) |n| if (std.mem.eql(u8, name, n)) return @field(kk, n); inline for (.{ "tree_round", "act_rows", "gdn_prep", "qk_rows", "norms", "moe", "experts", "qwen36_mtp", "nemotron_attention", "nemotron_route", "nemotron_latent", "qmm_group", "qmm_prefill", "lane_gemv", "prompt16", "qmmf", "ck_experts", "nvfp4_experts", "nvfp4_experts_tile", "rank_sum", "prefill_attention", "vit", "embed_rows", "route_rows", "flashnext_gdn", "flashnext_batch", "flashnext_hc", "flashnext_attn", "flashnext_mtp", "qwen_vision", "exl3_linear", "exl3_experts", "exl3_candidate", "experts3", "experts3_prefill", "drama3", "drama3_fp8", "drama3_codec", "glm_mm", "glm_dflash", "glm_glue", "glm_ops", "glm_index", "nemotron_mamba", "sample", "qwen35", "qwen35_bf16_rows", "qwen35_keep", "qwen35_kv8", "w4a8_prefill", "qwen35_draft", "ck_lane", "gdn_train", "train_q4", "train_attn", "flashnext_train", "train", "attn_train", "qwen27_train", "tree_attention", "block_attention", "window_attention", "window_pages", "moe_glue", "dense_rows", "qmm", "flashnext_fact", "flashnext_hcf", "flashnext_fexp", "experts_rows" }) |n| if (std.mem.eql(u8, name, n)) return @field(kk, n); + inline for (.{"flashnext_exl3"}) |n| + if (std.mem.eql(u8, name, n)) return @field(kk, n); if (std.mem.eql(u8, name, "contract_controls")) return if (options.controls) &controls.bytes else &.{}; std.debug.print("FAIL kernel-contract: unknown image name: {s}\n", .{name}); return error.UnknownContractImage; diff --git a/zig/tests/cuda/flashnext_exl3/kv/exl3_kv.zig b/zig/tests/cuda/flashnext_exl3/kv/exl3_kv.zig new file mode 120000 index 0000000000..f5414b3e7b --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/kv/exl3_kv.zig @@ -0,0 +1 @@ +../../../../src/families/flashnext/exl3/kv/exl3_kv.zig \ No newline at end of file diff --git a/zig/tests/cuda/flashnext_exl3/model/exl3_attn_cpu.zig b/zig/tests/cuda/flashnext_exl3/model/exl3_attn_cpu.zig new file mode 120000 index 0000000000..cb95d2d7a0 --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/model/exl3_attn_cpu.zig @@ -0,0 +1 @@ +../../../../src/families/flashnext/exl3/model/exl3_attn_cpu.zig \ No newline at end of file diff --git a/zig/tests/cuda/flashnext_exl3/model/exl3_cache.zig b/zig/tests/cuda/flashnext_exl3/model/exl3_cache.zig new file mode 120000 index 0000000000..a3cf4c6f46 --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/model/exl3_cache.zig @@ -0,0 +1 @@ +../../../../src/families/flashnext/exl3/model/exl3_cache.zig \ No newline at end of file diff --git a/zig/tests/cuda/flashnext_exl3/model/exl3_cpu.zig b/zig/tests/cuda/flashnext_exl3/model/exl3_cpu.zig new file mode 120000 index 0000000000..012c3b18c0 --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/model/exl3_cpu.zig @@ -0,0 +1 @@ +../../../../src/families/flashnext/exl3/model/exl3_cpu.zig \ No newline at end of file diff --git a/zig/tests/cuda/flashnext_exl3/model/exl3_fwd_cpu.zig b/zig/tests/cuda/flashnext_exl3/model/exl3_fwd_cpu.zig new file mode 120000 index 0000000000..8e732b9fe6 --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/model/exl3_fwd_cpu.zig @@ -0,0 +1 @@ +../../../../src/families/flashnext/exl3/model/exl3_fwd_cpu.zig \ No newline at end of file diff --git a/zig/tests/cuda/flashnext_exl3/model/exl3_gdn_cpu.zig b/zig/tests/cuda/flashnext_exl3/model/exl3_gdn_cpu.zig new file mode 120000 index 0000000000..3c9a8a52a3 --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/model/exl3_gdn_cpu.zig @@ -0,0 +1 @@ +../../../../src/families/flashnext/exl3/model/exl3_gdn_cpu.zig \ No newline at end of file diff --git a/zig/tests/cuda/flashnext_exl3/model/exl3_hc.zig b/zig/tests/cuda/flashnext_exl3/model/exl3_hc.zig new file mode 120000 index 0000000000..47b748a982 --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/model/exl3_hc.zig @@ -0,0 +1 @@ +../../../../src/families/flashnext/exl3/model/exl3_hc.zig \ No newline at end of file diff --git a/zig/tests/cuda/flashnext_exl3/model/exl3_moe.zig b/zig/tests/cuda/flashnext_exl3/model/exl3_moe.zig new file mode 120000 index 0000000000..cbbbe19f51 --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/model/exl3_moe.zig @@ -0,0 +1 @@ +../../../../src/families/flashnext/exl3/model/exl3_moe.zig \ No newline at end of file diff --git a/zig/tests/cuda/flashnext_exl3/model/exl3_ple2_tests.zig b/zig/tests/cuda/flashnext_exl3/model/exl3_ple2_tests.zig new file mode 100644 index 0000000000..a3e338829d --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/model/exl3_ple2_tests.zig @@ -0,0 +1,598 @@ +//! exl3-ple2: an EXL3 pack's n-gram path on v1.1.0's Flash Next family (cuda_ple.Ngram with `format = .exl3`). 1. + +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const flash = @import("flashnext"); +const check = @import("../../check.zig"); +const st = core.safetensors; +const ple = flash.ple; +const x3 = flash.exl3_ple; +const Gpu = check.Gpu; + +const py_ids_crc: u32 = 0xd04d5584; +const py_raw_crc: u32 = 0x2d57eafe; +const py_decoded_crc: u32 = 0xdcf520dc; + +const streams = 4; +const hidden = 2560; +const wide = streams * hidden; +const taps = 4; +const dilation = 3; +const tail_rows = (taps - 1) * dilation; +const eps: f32 = 1e-6; +const prefix = "model.language_model.layers.1.ple."; + +// ---- RSS ---------------------------------------------------------------------------------------------------------- + +const Rss = struct { total: u64 = 0, file: u64 = 0, anon: u64 = 0 }; // kB + +fn rss() Rss { + var buf: [1 << 14]u8 = undefined; + const fd = std.c.open("/proc/self/status", .{ .ACCMODE = .RDONLY }); + if (fd < 0) return .{}; + defer _ = std.c.close(fd); + var len: usize = 0; + while (len < buf.len) { + const n = std.c.read(fd, buf[len..].ptr, buf.len - len); + if (n <= 0) break; + len += @intCast(n); + } + var out: Rss = .{}; + var it = std.mem.splitScalar(u8, buf[0..len], '\n'); + while (it.next()) |line| { + const fields = [_]struct { key: []const u8, dst: *u64 }{ .{ .key = "VmRSS:", .dst = &out.total }, .{ .key = "RssFile:", .dst = &out.file }, .{ .key = "RssAnon:", .dst = &out.anon } }; + for (fields) |f| { + if (!std.mem.startsWith(u8, line, f.key)) continue; + var parts = std.mem.tokenizeAny(u8, line[f.key.len..], " \tkB"); + if (parts.next()) |n| f.dst.* = std.fmt.parseInt(u64, n, 10) catch 0; + } + } + return out; +} + +fn showRss(what: []const u8) Rss { + const r = rss(); + std.debug.print("RSS {s}: total {d} MiB, file-backed {d} MiB, anon {d} MiB\n", .{ what, r.total >> 10, r.file >> 10, r.anon >> 10 }); + return r; +} + +// ---- the branch's ids, ported from exl3-flashnext-window-wip zig/src/families/flashnext/exl3_ple_ids.zig ---------- + +/// Row ids [tokens.len][16] after `history` (incremental form; a token equal to eos starts a segment after itself). +fn oursIds(h: ple.Hash, history: [2]i64, tokens: []const u32, out: [][16]u64) void { + var seq: [3]i64 = undefined; + @memcpy(seq[0..2], &history); + var seg: usize = 0; + for (history) |t| seg = if (t == h.eos) 0 else @min(seg + 1, 2); + for (tokens, out[0..tokens.len]) |tok, *row| { + seq[2] = tok; + for (row, 0..) |*id, head| { + var mixed: u64 = 0; + for (0..2 + head / h.per) |q| { + const t = if (seg >= q) seq[2 - q] else h.eos; + const term = @as(u64, @bitCast(t)) *% @as(u64, @bitCast(h.mult[q])); + mixed = if (q == 0) term else mixed ^ term; + } + id.* = @as(u64, @intCast(h.offsets[head])) + mixed % @as(u64, @intCast(h.sizes[head])); + } + seg = if (tok == h.eos) 0 else @min(seg + 1, 2); + seq[0] = seq[1]; + seq[1] = seq[2]; + } +} + +// ---- weights -------------------------------------------------------------------------------------------------------- + +const Index = struct { + parsed: std.json.Parsed(std.json.Value), + files: std.StringHashMapUnmanaged(st.File) = .empty, + dir: []const u8, + + fn open(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) !Index { + const path = try std.fs.path.join(gpa, &.{ dir, "model.safetensors.index.json" }); + defer gpa.free(path); + const text = try std.Io.Dir.cwd().readFileAlloc(io, path, gpa, .limited(1 << 26)); + defer gpa.free(text); + return .{ .parsed = try std.json.parseFromSlice(std.json.Value, gpa, text, .{}), .dir = dir }; + } + + fn deinit(x: *Index, gpa: std.mem.Allocator, io: std.Io) void { + var it = x.files.valueIterator(); + while (it.next()) |f| f.close(io); + x.files.deinit(gpa); + x.parsed.deinit(); + } + + fn get(x: *Index, gpa: std.mem.Allocator, io: std.Io, name: []const u8) !st.Tensor { + const rel = (x.parsed.value.object.get("weight_map").?.object.get(name) orelse return error.MissingTensor).string; + if (x.files.get(rel)) |f| return f.get(name) orelse error.MissingTensor; + const path = try std.fs.path.join(gpa, &.{ x.dir, rel }); + defer gpa.free(path); + const f = try st.File.open(gpa, io, path); + try x.files.put(gpa, rel, f); + return f.get(name) orelse error.MissingTensor; + } +}; + +fn widen(gpa: std.mem.Allocator, t: st.Tensor, offset: f32) ![]f32 { + const out = try gpa.alloc(f32, t.numel()); + for (out, 0..) |*o, i| { + const w = std.mem.readInt(u16, t.bytes[2 * i ..][0..2], .little); + const v: f32 = switch (t.dtype) { + .f16 => @floatCast(@as(f16, @bitCast(w))), + .bf16 => @bitCast(@as(u32, w) << 16), + else => return error.UnexpectedTensor, + }; + o.* = v + offset; + } + return out; +} + +fn bf(x: f32) f32 { + return @bitCast(@as(u32, ple.toBf16(x)) << 16); +} + +fn bfBits(x: f32) u16 { + return ple.toBf16(x); +} + +fn bfVal(b: u16) f32 { + return @bitCast(@as(u32, b) << 16); +} + +// ---- references of the injection (from exl3-flashnext-window-wip exl3_ple_cpu.zig) --------------------------------- + +const Weights = struct { + key: []const f32, // [wide][hidden] + value: []const f32, // [hidden][hidden] + norm_key: []const f32, // [wide], 1 + w + norm_query: []const f32, + norm_conv: []const f32, + conv: []const f32, // [wide][taps], as stored (F16) +}; + +fn dot32(w: []const f32, x: []const f32) f32 { + var acc: f32 = 0; + for (w, x) |a, b| acc += a * b; + return acc; +} + +/// The keys and values the reference's F16 dense projections store (fp32 sums, bf16 outputs). +fn project(w: Weights, emb: []const f32, keys: []f32, vals: []f32) void { + for (0..wide) |o| keys[o] = bf(dot32(w.key[o * hidden ..][0..hidden], emb)); + for (0..hidden) |o| vals[o] = bf(dot32(w.value[o * hidden ..][0..hidden], emb)); +} + +/// The reference's roundings (glue.py `_ple_gate` / `_ple_conv`) on the host, one token at a time. +const Branch = struct { + w: Weights, + conv_bf: []f32, + gated: []f32, + nrow: []f32, + tail: []f32, + + fn init(gpa: std.mem.Allocator, w: Weights) !Branch { + const conv_bf = try gpa.alloc(f32, w.conv.len); + for (conv_bf, w.conv) |*d, s| d.* = bf(s); + const tail = try gpa.alloc(f32, tail_rows * wide); + @memset(tail, 0); + return .{ .w = w, .conv_bf = conv_bf, .gated = try gpa.alloc(f32, wide), .nrow = try gpa.alloc(f32, wide), .tail = tail }; + } + + fn step(b: *Branch, keys: []const f32, vals: []const f32, h: []const f32, out: []f32) void { + const w = b.w; + const d: f32 = hidden; + for (0..streams) |s| { + const k = keys[s * hidden ..][0..hidden]; + const q = h[s * hidden ..][0..hidden]; + var kss: f32 = 0; + var qss: f32 = 0; + for (k, q) |a, c| { + kss += a * a; + qss += c * c; + } + const krinv = 1.0 / @sqrt(kss / d + eps); + const qrinv = 1.0 / @sqrt(qss / d + eps); + var dot: f32 = 0; + for (0..hidden) |i| dot += bf(bf(k[i] * krinv * w.norm_key[s * hidden + i]) * bf(q[i] * qrinv * w.norm_query[s * hidden + i])); + var gate = bf(dot / @sqrt(d)); + const mag = @sqrt(@max(@abs(gate), 1e-6)); + gate = bf(if (gate > 0) mag else if (gate < 0) -mag else 0); + const sg = bf(1.0 / (1.0 + @exp(-gate))); + var gss: f32 = 0; + for (0..hidden) |i| { + const g = bf(sg * vals[i]); + b.gated[s * hidden + i] = g; + gss += g * g; + } + const rinv = 1.0 / @sqrt(gss / d + eps); + for (0..hidden) |i| b.nrow[s * hidden + i] = bf(b.gated[s * hidden + i] * rinv * w.norm_conv[s * hidden + i]); + } + for (0..wide) |e| { + var acc: f32 = 0; + for (0..taps) |t| { + const at = t * dilation; + acc += b.conv_bf[e * taps + t] * (if (at < tail_rows) b.tail[at * wide + e] else b.nrow[e]); + } + const cv = bf(acc); + out[e] = bf(b.gated[e] + bf(cv / (1.0 + @exp(-cv)))); + } + std.mem.copyForwards(f32, b.tail[0 .. (tail_rows - 1) * wide], b.tail[wide..]); + @memcpy(b.tail[(tail_rows - 1) * wide ..], b.nrow); + } +}; + +/// model.py `PLELayer` in float64 with no intermediate rounding (emb as decoded fp16, weights as stored). +const Oracle = struct { + w: Weights, + tail: []f64, + keys: []f64, + vals: []f64, + gated: []f64, + normed: []f64, + + fn init(gpa: std.mem.Allocator, w: Weights) !Oracle { + const tail = try gpa.alloc(f64, tail_rows * wide); + @memset(tail, 0); + return .{ .w = w, .tail = tail, .keys = try gpa.alloc(f64, wide), .vals = try gpa.alloc(f64, hidden), .gated = try gpa.alloc(f64, wide), .normed = try gpa.alloc(f64, wide) }; + } + + fn centred(x: []const f64, scale: []const f32, out: []f64) void { + var ss: f64 = 0; + for (x) |v| ss += v * v; + const r = 1.0 / @sqrt(ss / @as(f64, @floatFromInt(x.len)) + @as(f64, eps)); + for (out, x, scale) |*o, v, s| o.* = v * r * s; + } + + fn step(o: *Oracle, emb: []const f32, h: []const f32, out: []f64) void { + const w = o.w; + for (0..wide) |r| { + var acc: f64 = 0; + for (w.key[r * hidden ..][0..hidden], emb) |a, b| acc += @as(f64, a) * b; + o.keys[r] = acc; + } + for (0..hidden) |r| { + var acc: f64 = 0; + for (w.value[r * hidden ..][0..hidden], emb) |a, b| acc += @as(f64, a) * b; + o.vals[r] = acc; + } + var kn: [hidden]f64 = undefined; + var qn: [hidden]f64 = undefined; + var hs: [hidden]f64 = undefined; + for (0..streams) |s| { + centred(o.keys[s * hidden ..][0..hidden], w.norm_key[s * hidden ..][0..hidden], &kn); + for (&hs, h[s * hidden ..][0..hidden]) |*d, v| d.* = v; + centred(&hs, w.norm_query[s * hidden ..][0..hidden], &qn); + var dot: f64 = 0; + for (kn, qn) |a, b| dot += a * b; + var gate = dot / @sqrt(@as(f64, hidden)); + const sgn: f64 = if (gate > 0) 1 else if (gate < 0) -1 else 0; + gate = sgn * @sqrt(@max(@abs(gate), 1e-6)); + const sg = 1.0 / (1.0 + @exp(-gate)); + for (0..hidden) |i| o.gated[s * hidden + i] = sg * o.vals[i]; + centred(o.gated[s * hidden ..][0..hidden], w.norm_conv[s * hidden ..][0..hidden], o.normed[s * hidden ..][0..hidden]); + } + for (0..wide) |e| { + var acc: f64 = 0; + for (0..taps) |t| { + const at = t * dilation; + acc += @as(f64, w.conv[e * taps + t]) * (if (at < tail_rows) o.tail[at * wide + e] else o.normed[e]); + } + out[e] = o.gated[e] + acc / (1.0 + @exp(-acc)); + } + std.mem.copyForwards(f64, o.tail[0 .. (tail_rows - 1) * wide], o.tail[wide..]); + @memcpy(o.tail[(tail_rows - 1) * wide ..], o.normed); + } +}; + +const Qbench = struct { rows: []const struct { input_ids: []const u32 } }; + +fn crc(bytes: []const u8) u32 { + return std.hash.Crc32.hash(bytes); +} + +fn bytesOf(comptime T: type, v: []const T) []const u8 { + return std.mem.sliceAsBytes(v); +} + +pub fn run(gpu: Gpu, rest: []const [:0]const u8) !void { + if (rest.len < 1) return error.MissingArgument; + const dir = rest[0]; + const qrow: usize = if (rest.len > 1) try std.fmt.parseInt(usize, rest[1], 10) else 3; + const count: usize = if (rest.len > 2) try std.fmt.parseInt(usize, rest[2], 10) else 24; + var arena_state = std.heap.ArenaAllocator.init(gpu.gpa); + defer arena_state.deinit(); + const a = arena_state.allocator(); + const io = gpu.io; + const rss_start = showRss("start"); + + const qpath = try std.fs.path.join(a, &.{ dir, "qbench_prompts.json" }); + const qtext = try std.Io.Dir.cwd().readFileAlloc(io, qpath, a, .limited(1 << 26)); + const parsed = try std.json.parseFromSlice(Qbench, a, qtext, .{ .ignore_unknown_fields = true }); + const qrows = parsed.value.rows; + try check.expect(qrow < qrows.len and count <= qrows[qrow].input_ids.len, "qbench row {d} with {d} tokens", .{ qrow, count }); + const tokens = qrows[qrow].input_ids[0..count]; + + // ---- the table through upstream's Ngram (mapping) ----------------------------------------------------------- + var g = try ple.Ngram.init(gpu.gpa, io, dir); + defer g.deinit(); + try check.expect(g.format == .exl3, "Ngram.init took the EXL3 table", .{}); + const h = g.hash; + const heads = h.heads; + const dh = h.dh; + std.debug.print("hash: n {d}, heads {d} x {d}, eos {d}, {d} shards, {d} rows; table {d}-bit, {d} bytes a row, read {t}, emb {t}\n", .{ h.n, heads, dh, h.eos, h.shards, h.starts[h.shards], g.x3.?.bits, g.x3.?.rowBytes(), g.x3.?.opts.read, g.x3.?.opts.emb }); + check.pass("Ngram.init on the pack: config-derived multipliers, head sizes and offsets equal the pack's shipped tensors; 128 shard row counts equal Hash.starts", .{}); + const rss_open = showRss("after Ngram.init (30.4 GiB table mapped, header only)"); + + // ---- (1) ids --------------------------------------------------------------------------------------------------- + const ids = try a.alloc(i64, count * heads); + { + var hist = [_]i64{ h.eos, h.eos }; + for (tokens, 0..) |tok, t| { + try ple.lastIds(h, &hist, tok, ids[t * heads ..][0..heads]); + hist = .{ hist[1], tok }; + } + const got = crc(bytesOf(i64, ids)); + std.debug.print("ids[0][0..4] = {any}; ids crc32 {x:0>8} (numpy NGram.ids {x:0>8})\n", .{ ids[0..4], got, py_ids_crc }); + if (qrow == 3 and count == 24) try check.expect(got == py_ids_crc, "upstream lastIds crc32 {x:0>8} != numpy {x:0>8}", .{ got, py_ids_crc }); + // every token of every qbench prompt: upstream's lastIds against the branch's incremental ids. + var compared: usize = 0; + var eos_seen: usize = 0; + for (qrows) |r| { + const ours = try gpu.gpa.alloc([16]u64, r.input_ids.len); + defer gpu.gpa.free(ours); + oursIds(h, .{ h.eos, h.eos }, r.input_ids, ours); + var hist2 = [_]i64{ h.eos, h.eos }; + for (r.input_ids, ours) |tok, want| { + if (tok == h.eos) eos_seen += 1; + var up: [16]i64 = undefined; + try ple.lastIds(h, &hist2, tok, up[0..heads]); + hist2 = .{ hist2[1], tok }; + for (up[0..heads], want[0..heads]) |u, o| { + if (u < 0 or @as(u64, @intCast(u)) != o) { + std.debug.print("FAIL ids differ at token {d}: upstream {any} ours {any}\n", .{ tok, up[0..heads], want[0..heads] }); + return error.TestFailed; + } + } + compared += 1; + } + } + check.pass("row ids: upstream Hash/lastIds byte-equal to the numpy NGram.ids crc{s} and to the branch's exl3_ple_ids on all {d} tokens of {d} qbench prompts ({d} eos inside)", .{ if (qrow == 3 and count == 24) "" else " (crc not checked: not row 3 x 24)", compared, qrows.len, eos_seen }); + } + + // ---- (2) raw rows, two readers; decoded rows through Ngram.rows, two readers ------------------------------------ + var pread_t = (try x3.Table.open(gpu.gpa, io, dir, h, .{ .read = .pread })).?; + defer pread_t.destroy(); + const rb = g.x3.?.rowBytes(); + const raw_map = try a.alloc(u8, ids.len * rb); + const raw_ssd = try a.alloc(u8, ids.len * rb); + for (ids, 0..) |id, i| { + const loc = try h.locate(id); + try g.x3.?.raw(loc.shard, loc.local, raw_map[i * rb ..][0..rb]); + try pread_t.raw(loc.shard, loc.local, raw_ssd[i * rb ..][0..rb]); + } + try check.sameBytes("rows: mapping vs pread", raw_ssd, raw_map); + const raw_crc = crc(raw_map); + std.debug.print("{d} rows x {d} bytes: raw crc32 {x:0>8} (numpy {x:0>8})\n", .{ ids.len, rb, raw_crc, py_raw_crc }); + if (qrow == 3 and count == 24) try check.expect(raw_crc == py_raw_crc, "raw rows crc32 {x:0>8} != numpy {x:0>8}", .{ raw_crc, py_raw_crc }); + const rss_map = showRss("after the mapped row reads"); + + const w = g.width(); + const emb_map = try a.alloc(u16, count * w); + const xs_map = try a.alloc(f32, count * w / 32); + { + var hist = [_]i64{ h.eos, h.eos }; + try g.rows(&hist, tokens, emb_map, xs_map); + } + const dec_crc = crc(bytesOf(u16, emb_map)); + var peak_emb: f32 = 0; + for (emb_map) |v| peak_emb = @max(peak_emb, @abs(@as(f32, @floatCast(@as(f16, @bitCast(v)))))); + std.debug.print("Ngram.rows: {d} fp16 values, max |value| {d:.4}, crc32 {x:0>8} (numpy _ple_rows {x:0>8})\n", .{ emb_map.len, peak_emb, dec_crc, py_decoded_crc }); + if (qrow == 3 and count == 24) try check.expect(dec_crc == py_decoded_crc, "decoded crc32 {x:0>8} != numpy {x:0>8}", .{ dec_crc, py_decoded_crc }); + // the same call with the pread reader swapped in behind Ngram + const emb_ssd = try a.alloc(u16, count * w); + const xs_ssd = try a.alloc(f32, count * w / 32); + { + const mapped_t = g.x3.?; + g.x3 = pread_t; + defer g.x3 = mapped_t; + var hist = [_]i64{ h.eos, h.eos }; + try g.rows(&hist, tokens, emb_ssd, xs_ssd); + } + try check.sameBytes("Ngram.rows: mapping vs pread", bytesOf(u16, emb_ssd), bytesOf(u16, emb_map)); + try check.sameBytes("Ngram.rows group sums: mapping vs pread", bytesOf(f32, xs_ssd), bytesOf(f32, xs_map)); + check.pass("rows: {d} ids, stored bytes and Ngram.rows fp16 values byte-identical through the mapping and through pread + hot cache{s}", .{ ids.len, if (qrow == 3 and count == 24) "; raw and decoded crc32 equal the numpy references" else "" }); + + // ---- (3) hot cache over real text; RSS ------------------------------------------------------------------------- + { + var fresh = (try x3.Table.open(gpu.gpa, io, dir, h, .{ .read = .pread })).?; + defer fresh.destroy(); + const r0 = rss(); + var buf: [256]u8 = undefined; + var ntok: usize = 0; + for (qrows) |r| { + var hist = [_]i64{ h.eos, h.eos }; + for (r.input_ids) |tok| { + var row_ids: [16]i64 = undefined; + try ple.lastIds(h, &hist, tok, row_ids[0..heads]); + hist = .{ hist[1], tok }; + for (row_ids[0..heads]) |id| { + const loc = try h.locate(id); + try fresh.raw(loc.shard, loc.local, buf[0..rb]); + } + ntok += 1; + } + } + const s = fresh.hotStats(); + const r1 = rss(); + std.debug.print("pread + 4096-row hot cache over all {d} qbench tokens: {d} lookups, {d} hits ({d:.2}%); RSS file-backed {d} -> {d} MiB, anon {d} -> {d} MiB\n", .{ ntok, s.lookups, s.lookups - s.misses, 100.0 * @as(f64, @floatFromInt(s.lookups - s.misses)) / @as(f64, @floatFromInt(s.lookups)), r0.file >> 10, r1.file >> 10, r0.anon >> 10, r1.anon >> 10 }); + try check.expect(r1.file <= r0.file + 16 * 1024, "pread reads grew file-backed RSS by {d} MiB", .{(r1.file - r0.file) >> 10}); + } + const rss_rows = showRss("after rows (both readers) and the hot-cache pass"); + const table_mib = g.x3.?.file.?.map.memory.len >> 20; + std.debug.print("table {d} MiB; file-backed RSS: start {d}, after open {d}, after mapped reads {d}, now {d} MiB\n", .{ table_mib, rss_start.file >> 10, rss_open.file >> 10, rss_map.file >> 10, rss_rows.file >> 10 }); + try check.expect(rss_map.file <= rss_open.file + 128 * 1024 and rss_rows.file <= 512 * 1024, "the table became resident", .{}); + check.pass("30.4 GiB table stays out of memory: file-backed RSS {d} MiB after the mapped reads (+{d} MiB), anon {d} MiB", .{ rss_rows.file >> 10, (rss_map.file - @min(rss_map.file, rss_open.file)) >> 10, rss_rows.anon >> 10 }); + + // ---- (4) layer 1's injection through upstream's fn_ple_gate / fn_ple_conv --------------------------------------- + try injection(gpu, a, dir, tokens, emb_map, h.eos); +} + +fn injection(gpu: Gpu, a: std.mem.Allocator, dir: []const u8, tokens: []const u32, emb_bits: []const u16, eos: i64) !void { + _ = eos; + const io = gpu.io; + const n = tokens.len; + var index = try Index.open(gpu.gpa, io, dir); + defer index.deinit(gpu.gpa, io); + const key_t = try index.get(gpu.gpa, io, prefix ++ "key_proj.weight"); + const val_t = try index.get(gpu.gpa, io, prefix ++ "value_proj.weight"); + const conv_t = try index.get(gpu.gpa, io, prefix ++ "conv1d.weight"); + try check.expect(key_t.is(.f16, &.{ wide, hidden }) and val_t.is(.f16, &.{ hidden, hidden }) and conv_t.is(.f16, &.{ wide, 1, taps }), "PLE key/value/conv dtypes or shapes (want F16)", .{}); + var norms: [3][]f32 = undefined; + for ([_][]const u8{ "norm_key", "norm_query", "norm_conv" }, 0..) |nm, i| { + var nb: [96]u8 = undefined; + const t = try index.get(gpu.gpa, io, try std.fmt.bufPrint(&nb, prefix ++ "{s}.weight", .{nm})); + try check.expect(t.is(.bf16, &.{wide}), "{s} dtype/shape (want bf16 [10240])", .{nm}); + const raw = try widen(a, t, 0); + var mean: f64 = 0; + for (raw) |v| mean += v; + mean /= wide; + std.debug.print("{s} stored mean {d:.4} (gamma - 1 storage: scale 1 + w)\n", .{ nm, mean }); + try check.expect(@abs(mean) < 0.5, "{s} does not look gamma - 1", .{nm}); + norms[i] = try widen(a, t, 1.0); + } + const weights: Weights = .{ .key = try widen(a, key_t, 0), .value = try widen(a, val_t, 0), .norm_key = norms[0], .norm_query = norms[1], .norm_conv = norms[2], .conv = try widen(a, conv_t, 0) }; + // what upstream's loader (cuda_prompt_load.zig with_ple) would upload as CW: the raw bytes, read as bf16 + var raw_as_bf: usize = 0; + for (0..weights.conv.len) |i| { + const bits = std.mem.readInt(u16, conv_t.bytes[2 * i ..][0..2], .little); + if (bfVal(bits) != bf(weights.conv[i])) raw_as_bf += 1; + } + std.debug.print("conv1d.weight is F16: {d} of {d} values change if its bytes are read as bf16 (upstream's loader does), so it is converted (RNE) as exl3.py does\n", .{ raw_as_bf, weights.conv.len }); + + // inputs: the decoded rows (fp16), the layer input stand-in (token embedding tiled over 4 streams, bf16) + const embed_t = try index.get(gpu.gpa, io, "model.language_model.embed_tokens.weight"); + try check.expect(embed_t.dtype == .bf16 and embed_t.dim(1) == hidden, "embed_tokens", .{}); + const emb = try a.alloc(f32, n * hidden); + for (emb, emb_bits) |*e, b| e.* = @floatCast(@as(f16, @bitCast(b))); + const hq = try a.alloc(f32, n * wide); + const hq_bits = try a.alloc(u16, n * wide); + for (tokens, 0..) |tok, t| for (0..streams) |s| for (0..hidden) |d| { + const bits = std.mem.readInt(u16, embed_t.bytes[(@as(usize, tok) * hidden + d) * 2 ..][0..2], .little); + hq_bits[t * wide + s * hidden + d] = bits; + hq[t * wide + s * hidden + d] = bfVal(bits); + }; + // the projections the reference's F16 dense stores (bf16) + const keys = try a.alloc(f32, n * wide); + const vals = try a.alloc(f32, n * hidden); + for (0..n) |t| project(weights, emb[t * hidden ..][0..hidden], keys[t * wide ..][0..wide], vals[t * hidden ..][0..hidden]); + const keys_bits = try a.alloc(u16, keys.len); + for (keys_bits, keys) |*b, v| b.* = bfBits(v); + const vals_bits = try a.alloc(u16, vals.len); + for (vals_bits, vals) |*b, v| b.* = bfBits(v); + const cw_bits = try a.alloc(u16, weights.conv.len); + for (cw_bits, weights.conv) |*b, v| b.* = bfBits(v); + + // host references + const want = try a.alloc(f64, n * wide); + const emu = try a.alloc(f32, n * wide); + { + var oracle = try Oracle.init(a, weights); + var branch = try Branch.init(a, weights); + for (0..n) |t| { + oracle.step(emb[t * hidden ..][0..hidden], hq[t * wide ..][0..wide], want[t * wide ..][0..wide]); + branch.step(keys[t * wide ..][0..wide], vals[t * hidden ..][0..hidden], hq[t * wide ..][0..wide], emu[t * wide ..][0..wide]); + } + } + + // the device: upstream's glue kernels + var glue = try flash.Glue.load(gpu.d); + defer glue.deinit(); + var s = try cuda.Stream.init(gpu.d, false); + defer s.deinit(); + const D = cuda.DeviceBuffer; + var d_keys = try D.fromHost(gpu.d, bytesOf(u16, keys_bits)); + defer d_keys.free(); + var d_vals = try D.fromHost(gpu.d, bytesOf(u16, vals_bits)); + defer d_vals.free(); + var d_h = try D.fromHost(gpu.d, bytesOf(u16, hq_bits)); + defer d_h.free(); + var d_nk = try D.fromHost(gpu.d, bytesOf(f32, weights.norm_key)); + defer d_nk.free(); + var d_nq = try D.fromHost(gpu.d, bytesOf(f32, weights.norm_query)); + defer d_nq.free(); + var d_nc = try D.fromHost(gpu.d, bytesOf(f32, weights.norm_conv)); + defer d_nc.free(); + var d_cw = try D.fromHost(gpu.d, bytesOf(u16, cw_bits)); + defer d_cw.free(); + const zeros = try a.alloc(u8, n * wide * 2); + @memset(zeros, 0); + var d_zero = try D.fromHost(gpu.d, zeros); // conv's H: zero, so HOUT is the injection itself + defer d_zero.free(); + var d_tail = try D.fromHost(gpu.d, zeros[0 .. tail_rows * wide * 2]); + defer d_tail.free(); + var d_gated = try D.alloc(gpu.d, n * wide * 2); + defer d_gated.free(); + var d_pss = try D.alloc(gpu.d, n * streams * 4); + defer d_pss.free(); + var d_out = try D.alloc(gpu.d, n * wide * 2); + defer d_out.free(); + var d_nrow = try D.alloc(gpu.d, n * wide * 2); + defer d_nrow.free(); + try gpu.ctx.synchronize(); + + // whole prompt: one launch each, the conv reading its own earlier rows + try glue.pleGate(s, d_keys.ptr, d_vals.ptr, d_h.ptr, d_nk.ptr, d_nq.ptr, d_gated.ptr, d_pss.ptr, eps, n, hidden, streams); + try glue.pleConv(s, d_gated.ptr, d_pss.ptr, d_nc.ptr, d_tail.ptr, d_cw.ptr, d_zero.ptr, d_out.ptr, d_nrow.ptr, eps, n, hidden, streams, taps, dilation); + try s.synchronize(); + const prompt_bits = try a.alloc(u16, n * wide); + try d_out.download(0, std.mem.sliceAsBytes(prompt_bits)); + + // token by token: the tail carried on the host from each launch's NROW, as shiftWindow keeps it + const tail = try a.alloc(u16, tail_rows * wide); + @memset(tail, 0); + const decode_bits = try a.alloc(u16, n * wide); + const nrow = try a.alloc(u16, wide); + for (0..n) |t| { + try d_tail.upload(0, bytesOf(u16, tail)); + try gpu.ctx.synchronize(); + try glue.pleGate(s, d_keys.ptr + t * wide * 2, d_vals.ptr + t * hidden * 2, d_h.ptr + t * wide * 2, d_nk.ptr, d_nq.ptr, d_gated.ptr, d_pss.ptr, eps, 1, hidden, streams); + try glue.pleConv(s, d_gated.ptr, d_pss.ptr, d_nc.ptr, d_tail.ptr, d_cw.ptr, d_zero.ptr, d_out.ptr, d_nrow.ptr, eps, 1, hidden, streams, taps, dilation); + try s.synchronize(); + try d_out.download(0, std.mem.sliceAsBytes(decode_bits[t * wide ..][0..wide])); + try d_nrow.download(0, std.mem.sliceAsBytes(nrow)); + std.mem.copyForwards(u16, tail[0 .. (tail_rows - 1) * wide], tail[wide..]); + @memcpy(tail[(tail_rows - 1) * wide ..], nrow); + } + + var worst: f64 = 0; + var peak: f64 = 0; + var worst_rel_rms: f64 = 0; + var emu_worst: f64 = 0; + var emu_off: usize = 0; + var emu_vs_oracle: f64 = 0; + for (0..n) |t| { + var se: f64 = 0; + var sw: f64 = 0; + for (prompt_bits[t * wide ..][0..wide], want[t * wide ..][0..wide], emu[t * wide ..][0..wide]) |gb, o, e| { + const gv: f64 = bfVal(gb); + const err = gv - o; + worst = @max(worst, @abs(err)); + peak = @max(peak, @abs(o)); + se += err * err; + sw += o * o; + emu_worst = @max(emu_worst, @abs(gv - e)); + if (gv != e) emu_off += 1; + emu_vs_oracle = @max(emu_vs_oracle, @abs(@as(f64, e) - o)); + } + worst_rel_rms = @max(worst_rel_rms, @sqrt(se / sw)); + } + const same_modes = std.mem.eql(u16, prompt_bits, decode_bits); + std.debug.print("injection over {d} tokens: GPU max |gpu - f64| {e:.3} at peak {e:.3} ({d:.3}% of peak), worst per-token rms rel {e:.3}\n", .{ n, worst, peak, 100 * worst / peak, worst_rel_rms }); + std.debug.print(" GPU vs the reference's roundings emulated on the host: max |diff| {e:.3}, {d} of {d} values differ; emulation vs f64 max {e:.3}\n", .{ emu_worst, emu_off, n * wide, emu_vs_oracle }); + std.debug.print(" whole-prompt launch vs token-by-token launches: {s}\n", .{if (same_modes) "bit-identical" else "DIFFER"}); + try check.expect(peak > 0 and worst <= 0.03 * peak and worst_rel_rms <= 0.02, "injection beyond bf16 rounding of the f64 oracle", .{}); + try check.expect(same_modes, "prompt and decode launches differ", .{}); + check.pass("layer-1 PLE injection through upstream's fn_ple_gate/fn_ple_conv on real EXL3 weights: max |err| {e:.2} <= 3% of peak {e:.2}, rms rel {e:.2} <= 2e-2; prompt == decode bits", .{ worst, peak, worst_rel_rms }); +} diff --git a/zig/tests/cuda/flashnext_exl3/model/exl3_route.zig b/zig/tests/cuda/flashnext_exl3/model/exl3_route.zig new file mode 120000 index 0000000000..587fabf2a4 --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/model/exl3_route.zig @@ -0,0 +1 @@ +../../../../src/families/flashnext/exl3/model/exl3_route.zig \ No newline at end of file diff --git a/zig/tests/cuda/flashnext_exl3/model/exl3_router.zig b/zig/tests/cuda/flashnext_exl3/model/exl3_router.zig new file mode 120000 index 0000000000..cd9eb177b7 --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/model/exl3_router.zig @@ -0,0 +1 @@ +../../../../src/families/flashnext/exl3/model/exl3_router.zig \ No newline at end of file diff --git a/zig/tests/cuda/flashnext_exl3/model/flashnext_exl3_engine_tests.zig b/zig/tests/cuda/flashnext_exl3/model/flashnext_exl3_engine_tests.zig new file mode 100644 index 0000000000..15236b0bad --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/model/flashnext_exl3_engine_tests.zig @@ -0,0 +1,536 @@ +//! flashnext-exl3-engine: the Qwen3.8-Flash-Next EXL3 pack through the Flash Next family's own Engine +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const flash = @import("flashnext"); +const check = @import("../../check.zig"); + +const Gpu = check.Gpu; +const st = core.safetensors; +const fmt = core.exl3_format; +const prompt = flash.prompt; +const xpass = flash.exl3.pass; +const vocab = prompt.vocabulary; +const dims: usize = 2560; +const cap_default: usize = 32768; +/// Rows of logits a teacher-forced pass returns at most. +const out_rows: usize = 64; + +fn promote(bits: u16) f32 { + return @bitCast(@as(u32, bits) << 16); +} + +fn ms(t0: i96, t1: i96) f64 { + return @as(f64, @floatFromInt(t1 - t0)) / 1e6; +} + +const Eng = struct { + gpu: Gpu, + stream: cuda.Stream, + glue: flash.Glue, + e: *prompt.Engine, + seq: prompt.Seq, + out: cuda.DeviceBuffer, + cap: usize, + attached: xpass.Attached, + load_s: f64, + + fn init(gpu: Gpu, dir: []const u8, lines: ?usize, cap: usize) !*Eng { + const g = try gpu.gpa.create(Eng); + errdefer gpu.gpa.destroy(g); + g.gpu = gpu; + g.cap = cap; + g.stream = try cuda.Stream.init(gpu.d, true); + g.glue = try flash.Glue.load(gpu.d); + const t0 = check.now(gpu.io); + g.e = try prompt.Engine.initRank(gpu.gpa, gpu.io, @constCast(gpu.d), &g.stream, dir, 1, 0); + const t1 = check.now(gpu.io); + // as cuda_lanes_load.open: the pass's scratch for a full chunk and the block scores, then the stream's caches + _ = try g.e.reserve(@min(cap, g.e.chunk)); + _ = try g.e.reserveScores(cap, @min(cap, g.e.chunk)); + g.seq = try g.e.newSeq(cap); + g.out = try cuda.DeviceBuffer.alloc(gpu.d, out_rows * vocab * 2); + var b = try xpass.budgetFromEnv(0); + if (lines) |n| b.lines = n; + g.attached = try xpass.attach(gpu.gpa, gpu.io, gpu.d, g.stream, g.e.layers[0..g.e.n_layers], g.e.chunk, b, null); + const t2 = check.now(gpu.io); + g.load_s = ms(t0, t2) / 1e3; + var free: usize = 0; + var total: usize = 0; + try gpu.d.check(gpu.d.api.cuMemGetInfo_v2(&free, &total), "cuMemGetInfo"); + std.debug.print("engine: loaded in {d:.1} s (Engine.initRank {d:.1} s incl. the pinned bank's copies {d:.1} s; pools {d:.1} s); {d} lines a layer x {d} layers ({d:.2} GiB), {d:.2} GiB free after, chunk {d}, context {d}\n", .{ g.load_s, ms(t0, t1) / 1e3, @as(f64, @floatFromInt(g.attached.bank_ns)) / 1e9, ms(t1, t2) / 1e3, g.attached.lines, g.attached.layers, @as(f64, @floatFromInt(g.attached.device_bytes)) / (1 << 30), @as(f64, @floatFromInt(free)) / (1 << 30), g.e.chunk, cap }); + return g; + } + + fn deinit(g: *Eng) void { + g.out.free(); + g.seq.deinit(g.gpu.gpa); + g.e.deinit(); + g.glue.deinit(); + g.stream.deinit(); + g.gpu.gpa.destroy(g); + } + + fn tri(g: *Eng) flash.Tri { + return .{ .own = &g.glue, .s = g.stream }; + } + + /// The sequence as Engine.newSeq makes it, without a new allocation (the pools took the free memory). + fn reset(g: *Eng) !void { + const s = &g.seq; + for ([_][]cuda.DeviceBuffer{ s.conv, s.state, s.kc, s.vc, s.ikc, s.pooled }) |bs| for (bs) |b| try b.fill8(0, g.stream.handle); + try s.tail.fill8(0, g.stream.handle); + try g.stream.synchronize(); + s.fed = 0; + s.history = .{ g.e.eos, g.e.eos }; + } + + fn stats(g: *Eng) flash.exl3_pool.Stats { + return xpass.poolStats(g.e.layers[0..g.e.n_layers]); + } +}; + +const Golden = struct { + parsed: std.json.Parsed(std.json.Value), + + fn open(gpu: Gpu, path: []const u8) !Golden { + const text = try std.Io.Dir.cwd().readFileAlloc(gpu.io, path, gpu.gpa, .limited(1 << 26)); + defer gpu.gpa.free(text); + return .{ .parsed = try std.json.parseFromSlice(std.json.Value, gpu.gpa, text, .{}) }; + } + + fn prompts(g: *const Golden) []std.json.Value { + return g.parsed.value.object.get("prompts").?.array.items; + } +}; + +fn ints(a: std.mem.Allocator, v: std.json.Value) ![]u32 { + const items = v.array.items; + const out = try a.alloc(u32, items.len); + for (items, out) |x, *o| o.* = @intCast(x.integer); + return out; +} + +fn num(v: std.json.Value) f64 { + return switch (v) { + .float => |f| f, + .integer => |i| @floatFromInt(i), + else => std.math.nan(f64), + }; +} + +fn top5(row: []const u16) [5]u32 { + var ids: [5]u32 = @splat(0); + var vals: [5]f32 = @splat(-std.math.inf(f32)); + for (row, 0..) |b, i| { + const v = promote(b); + if (!(v > vals[4])) continue; + var at: usize = 4; + while (at > 0 and v > vals[at - 1]) : (at -= 1) { + vals[at] = vals[at - 1]; + ids[at] = ids[at - 1]; + } + vals[at] = v; + ids[at] = @intCast(i); + } + return ids; +} + +// -- the EXL3 head in float64 ------------------------------------------------------------------------------------ + +const Raw = struct { + files: [7]st.File = undefined, + n: usize = 0, + + fn open(gpu: Gpu, dir: []const u8) !Raw { + var r: Raw = .{}; + var buf: [512]u8 = undefined; + for (0..7) |i| { + r.files[i] = try st.File.open(gpu.gpa, gpu.io, try std.fmt.bufPrint(&buf, "{s}/model-{d:0>5}-of-00007.safetensors", .{ dir, i + 1 })); + r.n += 1; + } + return r; + } + + fn close(r: *Raw, io: std.Io) void { + for (r.files[0..r.n]) |*f| f.close(io); + } + + fn get(r: *const Raw, name: []const u8) !st.Tensor { + for (r.files[0..r.n]) |*f| if (f.get(name)) |t| return t; + return error.MissingTensor; + } +}; + +/// lm_head as stored: trellis int16 [K/16, N/16, 8 * halves] (bytes), suh fp16 [K], svh fp16 [N]. +const Head = struct { + trellis: []const u8, + suh: []const u8, + svh: []const u8, + k: usize, + n: usize, + halves: u16, + + fn words(h: Head) usize { + return @as(usize, h.halves) * 8; + } + + fn of(trellis: st.Tensor, suh: st.Tensor, svh: st.Tensor, mul1: ?st.Tensor) !Head { + if (trellis.dtype != .i16 or trellis.rank != 3 or trellis.dim(2) % 8 != 0) return error.UnexpectedTensor; + const h: Head = .{ .trellis = trellis.bytes, .suh = suh.bytes, .svh = svh.bytes, .k = trellis.dim(0) * 16, .n = trellis.dim(1) * 16, .halves = @intCast(trellis.dim(2) / 8) }; + if (suh.dtype != .f16 or suh.numel() != h.k or svh.dtype != .f16 or svh.numel() != h.n) return error.UnexpectedTensor; + if (h.k % fmt.hadamard_dim != 0 or h.n % fmt.hadamard_dim != 0 or h.halves < 2 or h.halves > 16) return error.UnexpectedTensor; + if (mul1) |m| { + if (m.bytes.len != 4 or std.mem.readInt(u32, m.bytes[0..4], .little) != fmt.Mul1Mul) return error.UnexpectedTensor; + } + return h; + } +}; + +const HeadWork = struct { + head: Head, + x: []const f64, // [rows, K] + rows: usize, + out: []f64, // [rows, N] + next: std.atomic.Value(usize) = .init(0), + failed: ?anyerror = null, +}; + +fn headWorker(w: *HeadWork, gpa: std.mem.Allocator) void { + headWork(w, gpa) catch |err| { + w.failed = err; + }; +} + +/// logits[r][v] = sum_k x[r][k] * W[k][v], W decoded in float64 a 128-column block at a time (exl3_format.dequantize). +fn headWork(w: *HeadWork, gpa: std.mem.Allocator) !void { + const h = w.head; + const bn: usize = fmt.hadamard_dim; + const tiles_k = h.k / 16; + const tiles_n = h.n / 16; + const per = bn / 16; + const words = h.words(); + const sub = try gpa.alloc(i16, tiles_k * per * words); + defer gpa.free(sub); + const vals = try gpa.alloc(f64, h.k * bn); + defer gpa.free(vals); + const scratch = try gpa.alloc(f64, 2 * h.k * bn); + defer gpa.free(scratch); + const suh = try gpa.alloc(f16, h.k); + defer gpa.free(suh); + for (suh, 0..) |*o, i| o.* = @bitCast(std.mem.readInt(u16, h.suh[2 * i ..][0..2], .little)); + var svh: [fmt.hadamard_dim]f16 = undefined; + while (true) { + const nb = w.next.fetchAdd(1, .monotonic); + if (nb >= h.n / bn) return; + for (0..tiles_k) |tk| for (0..per) |j| { + const src = ((tk * tiles_n) + nb * per + j) * words * 2; + const dst = sub[(tk * per + j) * words ..][0..words]; + for (dst, 0..) |*d, i| d.* = std.mem.readInt(i16, h.trellis[src + 2 * i ..][0..2], .little); + }; + for (&svh, 0..) |*o, i| o.* = @bitCast(std.mem.readInt(u16, h.svh[2 * (nb * bn + i) ..][0..2], .little)); + fmt.dequantize(.{ .halves = h.halves }, .mul1, sub, tiles_k, per, suh, &svh, vals, scratch); + for (0..w.rows) |r| { + const xr = w.x[r * h.k ..][0..h.k]; + var acc: [fmt.hadamard_dim]f64 = @splat(0); + for (0..h.k) |k| { + const xv = xr[k]; + const row = vals[k * bn ..][0..bn]; + for (&acc, row) |*a, v| a.* += xv * v; + } + @memcpy(w.out[r * h.n + nb * bn ..][0..bn], &acc); + } + } +} + +fn headF64(gpa: std.mem.Allocator, raw: *const Raw, x: []const f64, rows: usize) ![]f64 { + const head = try Head.of(try raw.get("lm_head.trellis"), try raw.get("lm_head.suh"), try raw.get("lm_head.svh"), raw.get("lm_head.mul1") catch null); + const out = try gpa.alloc(f64, rows * head.n); + errdefer gpa.free(out); + var w: HeadWork = .{ .head = head, .x = x, .rows = rows, .out = out }; + var pool: [64]std.Thread = undefined; + const n = @min(pool.len, @max(1, std.Thread.getCpuCount() catch 1)); + var spawned: usize = 0; + for (0..n) |_| { + pool[spawned] = try std.Thread.spawn(.{}, headWorker, .{ &w, gpa }); + spawned += 1; + } + for (pool[0..spawned]) |t| t.join(); + if (w.failed) |err| return err; + return out; +} + +// -- modes --------------------------------------------------------------------------------------------------------- + +fn tfMode(gpu: Gpu, dir: []const u8, golden_path: []const u8, lines: ?usize) !void { + const a = gpu.gpa; + var golden = try Golden.open(gpu, golden_path); + defer golden.parsed.deinit(); + const p = golden.prompts()[0]; + const ids = try ints(a, p.object.get("ids").?); + defer a.free(ids); + const tf = p.object.get("teacher_forced").?.array.items; + const g = try Eng.init(gpu, dir, lines, cap_default); + defer g.deinit(); + try g.reset(); + const n = ids.len; + const ops = g.e.layers[0].exl3.?.ops; + const before = ops.fingerprint(); + const t0 = check.now(gpu.io); + try g.e.forward(flash.Tri, g.tri(), &g.seq, ids, .{ .logits = .all, .logits_out = g.out.ptr }); + try g.stream.synchronize(); + const t1 = check.now(gpu.io); + const logits = try a.alloc(u16, n * vocab); + defer a.free(logits); + try g.out.download(0, std.mem.sliceAsBytes(logits)); + // the head rows the pass fed lm_head (the final mixer's read-out, bf16 [n, dims]) + const hrows = try a.alloc(u16, n * dims); + defer a.free(hrows); + try gpu.d.check(gpu.d.api.cuMemcpyDtoH_v2(hrows.ptr, g.e.batch_s.?.mixed, n * dims * 2), "cuMemcpyDtoH"); + std.debug.print("teacher-forced pass: {d} rows in {d:.1} ms (first pass: includes every miss's copy)\n", .{ n, ms(t0, t1) }); + // one decode round too, then every EXL3 scratch buffer must be where Ops.reserve put it + { + var tok = prompt.argmax(g.seq.logits); + const q = [_]*prompt.Seq{&g.seq}; + try g.e.round(flash.Tri, g.tri(), &q, (&tok)[0..1], g.out.ptr); + tok = prompt.argmax(g.seq.logits); + } + const after = ops.fingerprint(); + try check.expect(std.mem.eql(u8, std.mem.sliceAsBytes(&before), std.mem.sliceAsBytes(&after)), "an EXL3 scratch buffer grew inside a pass", .{}); + check.pass("EXL3 scratch reserved up front: the {d} scratch buffers kept their addresses and sizes through a {d}-row prompt pass and a decode round", .{ blk: { + var k: usize = 0; + for (after) |e| k += @intFromBool(e[1] != 0); + break :blk k; + }, n }); + + std.debug.print(" pos input ours top1 (logit) golden top1 (logit) match top5 overlap max|dlogit| on golden top5\n", .{}); + var top1_ok: usize = 0; + var top5_same: usize = 0; + var worst: f64 = 0; + var sum_abs: f64 = 0; + var count: usize = 0; + for (tf, 0..) |row, pos| { + const r = logits[pos * vocab ..][0..vocab]; + const ours = top5(r); + const g5 = try ints(a, row.object.get("top5_ids").?); + defer a.free(g5); + const gl = row.object.get("top5_logits").?.array.items; + var overlap: usize = 0; + for (ours) |o| for (g5) |x| { + overlap += @intFromBool(o == x); + }; + var dmax: f64 = 0; + for (g5, gl) |id, lv| { + const d = @abs(@as(f64, promote(r[id])) - num(lv)); + dmax = @max(dmax, d); + sum_abs += d; + count += 1; + } + worst = @max(worst, dmax); + const ok = ours[0] == g5[0]; + top1_ok += @intFromBool(ok); + top5_same += @intFromBool(std.mem.eql(u32, &ours, g5)); + std.debug.print("{d:4} {d:7} {d:9} ({d:7.3}) {d:9} ({d:7.3}) {s:5} {d}/5 {s} {d:.4}\n", .{ pos, ids[pos], ours[0], promote(r[ours[0]]), g5[0], num(gl[0]), if (ok) "yes" else "NO", overlap, if (std.mem.eql(u32, &ours, g5)) "(same order)" else " ", dmax }); + } + std.debug.print("TF top1 {d}/{d}, top5 identical in order {d}/{d}, |our logit - golden bf16 logit| on the golden's top-5 ids: max {d:.4}, mean {d:.4}\n", .{ top1_ok, tf.len, top5_same, tf.len, worst, sum_abs / @as(f64, @floatFromInt(count)) }); + + // the EXL3 lm_head itself, on the same head rows, against an exact float64 product over all 248,320 columns + var raw = try Raw.open(gpu, dir); + defer raw.close(gpu.io); + const xf = try a.alloc(f64, n * dims); + defer a.free(xf); + for (xf, hrows) |*o, b| o.* = promote(b); + const t2 = check.now(gpu.io); + const exact = try headF64(a, &raw, xf, n); + defer a.free(exact); + const t3 = check.now(gpu.io); + var head_abs: f64 = 0; + var head_rel: f64 = 0; + var head_top1: usize = 0; + for (0..n) |r| { + var peak: f64 = 0; + var best: usize = 0; + var dmax: f64 = 0; + for (0..vocab) |v| { + const want = exact[r * vocab + v]; + peak = @max(peak, @abs(want)); + if (want > exact[r * vocab + best]) best = v; + dmax = @max(dmax, @abs(@as(f64, promote(logits[r * vocab + v])) - want)); + } + head_abs = @max(head_abs, dmax); + head_rel = @max(head_rel, dmax / peak); + const ours = top5(logits[r * vocab ..][0..vocab])[0]; + // a bf16 tie at the top is no disagreement + head_top1 += @intFromBool(ours == best or promote(logits[r * vocab + ours]) == promote(logits[r * vocab + best])); + } + std.debug.print("EXL3 lm_head (Dense, bf16 out) vs exact float64 on the pass's {d} head rows, {d} columns ({d:.1} s on the host): max |d| {d:.4} (rel to the row peak {e:.2}), top-1 equal on {d}/{d}\n", .{ n, vocab, ms(t2, t3) / 1e3, head_abs, head_rel, head_top1, n }); + // bf16 output rounding is 2^-9 of a value (0.06 at a logit of 30); the fp16 rotated input adds ~2.4e-4 of the peak + try check.expect(head_rel < 4e-3 and head_top1 == n, "lm_head within bf16 rounding of float64", .{}); + check.pass("lm_head: EXL3 hook through exl3.Dense, {d} rows, max |d| {d:.4} (rel {e:.2}) vs float64, top-1 equal on every row", .{ n, head_abs, head_rel }); + try check.expect(top1_ok == tf.len, "teacher-forced top-1 equals the golden on every position ({d}/{d})", .{ top1_ok, tf.len }); + check.pass("teacher-forced top-1 {d}/{d} vs the golden; max |dlogit| {d:.4} on the golden's top-5", .{ top1_ok, tf.len, worst }); + printPool(g, "after tf"); +} + +fn printPool(g: *Eng, what: []const u8) void { + const s = g.stats(); + const routed = s.hit_slots + s.miss_slots; + std.debug.print("pool {s}: {d} rounds, {d} passes, hit rate {d:.1}% ({d}/{d} routed slots), {d} copies, {d} evictions; host cut {d:.1} ms over {d} blocks, plan {d:.1} ms\n", .{ what, s.rounds, s.passes, 100 * @as(f64, @floatFromInt(s.hit_slots)) / @as(f64, @floatFromInt(@max(routed, 1))), s.hit_slots, routed, s.misses, s.evictions, @as(f64, @floatFromInt(xpass.pool_time.cut_ns)) / 1e6, xpass.pool_time.blocks, @as(f64, @floatFromInt(xpass.pool_time.plan_ns)) / 1e6 }); +} + +/// Greedy `steps` tokens after `ids`: the prefill's last row, then a decode round a token. +fn greedyRun(g: *Eng, ids: []const u32, out: []u32, step_ms: ?*f64) !void { + try g.reset(); + try g.e.forward(flash.Tri, g.tri(), &g.seq, ids, .{ .logits = .last }); + out[0] = prompt.argmax(g.seq.logits); + const t0 = check.now(g.gpu.io); + for (1..out.len) |i| { + const q = [_]*prompt.Seq{&g.seq}; + try g.e.round(flash.Tri, g.tri(), &q, out[i - 1 ..][0..1], g.out.ptr); + out[i] = prompt.argmax(g.seq.logits); + } + if (step_ms) |p| p.* = ms(t0, check.now(g.gpu.io)) / @as(f64, @floatFromInt(@max(out.len - 1, 1))); +} + +/// Each golden prompt followed by the golden's own 24 greedy ids, teacher-forced in one pass: our top-1 at every +fn tfcMode(gpu: Gpu, dir: []const u8, golden_path: []const u8, lines: ?usize) !void { + const a = gpu.gpa; + var golden = try Golden.open(gpu, golden_path); + defer golden.parsed.deinit(); + const g = try Eng.init(gpu, dir, lines, cap_default); + defer g.deinit(); + var all_ok = true; + for (golden.prompts()) |p| { + const ids = try ints(a, p.object.get("ids").?); + defer a.free(ids); + const cont = try ints(a, p.object.get("greedy_ids").?); + defer a.free(cont); + const steps = p.object.get("greedy_steps").?.array.items; + const all = try a.alloc(u32, ids.len + cont.len - 1); + defer a.free(all); + @memcpy(all[0..ids.len], ids); + @memcpy(all[ids.len..], cont[0 .. cont.len - 1]); + if (all.len > out_rows) return error.UnexpectedTensor; + try g.reset(); + try g.e.forward(flash.Tri, g.tri(), &g.seq, all, .{ .logits = .all, .logits_out = g.out.ptr }); + try g.stream.synchronize(); + const logits = try a.alloc(u16, all.len * vocab); + defer a.free(logits); + try g.out.download(0, std.mem.sliceAsBytes(logits[0 .. all.len * vocab])); + std.debug.print("TFC {s}\n step golden id ours top1 our margin golden margin max|dlogit| golden top5\n", .{p.object.get("name").?.string}); + var ok_n: usize = 0; + for (cont, 0..) |want, k| { + const r = logits[(ids.len - 1 + k) * vocab ..][0..vocab]; + const t = top5(r); + const margin = promote(r[t[0]]) - promote(r[t[1]]); + const g5 = try ints(a, steps[k].object.get("top5_ids").?); + defer a.free(g5); + var dmax: f64 = 0; + for (g5, steps[k].object.get("top5_logits").?.array.items) |id, lv| dmax = @max(dmax, @abs(@as(f64, promote(r[id])) - num(lv))); + const ok = t[0] == want; + ok_n += @intFromBool(ok); + std.debug.print(" {d:4} {d:10} {d:9} {s} {d:10.4} {d:14.4} {d:10.4}\n", .{ k, want, t[0], if (ok) " " else "*", margin, num(steps[k].object.get("margin12").?), dmax }); + } + std.debug.print("TFC {s}: top-1 equals the golden's greedy id on {d}/{d} continuation positions\n", .{ p.object.get("name").?.string, ok_n, cont.len }); + if (ok_n != cont.len) all_ok = false; + } + printPool(g, "after tfc"); + try check.expect(all_ok, "teacher-forced continuation top-1 equals the golden everywhere", .{}); + check.pass("teacher-forced continuation: top-1 equals the golden's greedy id at every step of all {d} prompts", .{golden.prompts().len}); +} + +fn genMode(gpu: Gpu, dir: []const u8, golden_path: []const u8, lines: ?usize) !void { + const a = gpu.gpa; + var golden = try Golden.open(gpu, golden_path); + defer golden.parsed.deinit(); + const g = try Eng.init(gpu, dir, lines, cap_default); + defer g.deinit(); + var all_ok = true; + for (golden.prompts()) |p| { + const ids = try ints(a, p.object.get("ids").?); + defer a.free(ids); + const want = try ints(a, p.object.get("greedy_ids").?); + defer a.free(want); + const got = try a.alloc(u32, want.len); + defer a.free(got); + var step: f64 = 0; + try greedyRun(g, ids, got, &step); + var first_diff: ?usize = null; + for (got, want, 0..) |x, y, i| if (x != y and first_diff == null) { + first_diff = i; + }; + std.debug.print("GEN {s}: ", .{p.object.get("name").?.string}); + for (got) |x| std.debug.print("{d},", .{x}); + std.debug.print(" {s} ({d:.1} ms/step decode)\n", .{ if (first_diff != null) "DIFFERS" else "equal to the golden", step }); + if (first_diff) |d| { + std.debug.print(" first difference at step {d}: ours {d}, golden {d}; golden margin there {d:.4}\n", .{ d, got[d], want[d], num(p.object.get("greedy_steps").?.array.items[d].object.get("margin12").?) }); + all_ok = false; + } + } + printPool(g, "after gen"); + try check.expect(all_ok, "greedy 24 equals the golden on every prompt", .{}); + check.pass("greedy 24 tokens equal the golden on all {d} prompts", .{golden.prompts().len}); +} + +fn benchMode(gpu: Gpu, dir: []const u8, lines: ?usize) !void { + const a = gpu.gpa; + const g = try Eng.init(gpu, dir, lines, cap_default); + defer g.deinit(); + const max_n: usize = 2048; + const ids = try a.alloc(u32, max_n); + defer a.free(ids); + // text-like ids: a fixed pseudo-random walk over the common part of the vocabulary + var prng = std.Random.DefaultPrng.init(0x5eed); + for (ids) |*x| x.* = 200 + prng.random().uintLessThan(u32, 60000); + // warm-up: one pass to load kernels, settle the pool + try g.reset(); + try g.e.forward(flash.Tri, g.tri(), &g.seq, ids[0..64], .{ .logits = .last }); + for ([_]usize{ 16, 256, 2048 }) |n| { + var best: f64 = std.math.inf(f64); + for (0..2) |_| { + try g.reset(); + const before = g.stats(); + const t0 = check.now(gpu.io); + try g.e.forward(flash.Tri, g.tri(), &g.seq, ids[0..n], .{ .logits = .last }); + const t = ms(t0, check.now(gpu.io)); + const after = g.stats(); + const hits = after.hit_slots - before.hit_slots; + const all = hits + after.miss_slots - before.miss_slots; + std.debug.print(" prefill {d:5} tokens: {d:8.1} ms = {d:7.1} rows/s (hit {d:.1}%, {d} copies)\n", .{ n, t, @as(f64, @floatFromInt(n)) / (t / 1e3), 100 * @as(f64, @floatFromInt(hits)) / @as(f64, @floatFromInt(@max(all, 1))), after.misses - before.misses }); + best = @min(best, t); + } + std.debug.print("BENCH prefill {d} tokens: best {d:.1} ms, {d:.1} rows/s\n", .{ n, best, @as(f64, @floatFromInt(n)) / (best / 1e3) }); + // decode after this prompt: 32 tokens, greedy + var tok = prompt.argmax(g.seq.logits); + const before = g.stats(); + const t0 = check.now(gpu.io); + const steps: usize = 32; + for (0..steps) |_| { + const q = [_]*prompt.Seq{&g.seq}; + try g.e.round(flash.Tri, g.tri(), &q, (&tok)[0..1], g.out.ptr); + tok = prompt.argmax(g.seq.logits); + } + const t = ms(t0, check.now(gpu.io)) / steps; + const after = g.stats(); + const hits = after.hit_slots - before.hit_slots; + const all = hits + after.miss_slots - before.miss_slots; + std.debug.print("BENCH decode after {d} tokens: {d:.2} ms/token over {d} steps (hit {d:.1}%, {d:.1} copies a token)\n", .{ n, t, steps, 100 * @as(f64, @floatFromInt(hits)) / @as(f64, @floatFromInt(@max(all, 1))), @as(f64, @floatFromInt(after.misses - before.misses)) / steps }); + } + printPool(g, "after bench"); +} + +pub fn run(gpu: Gpu, rest: []const [:0]const u8) !void { + if (rest.len < 2) return error.MissingArgument; + const dir = rest[0]; + const mode = rest[1]; + const lineArg = struct { + fn f(r: []const [:0]const u8, i: usize) !?usize { + return if (r.len > i) try std.fmt.parseInt(usize, r[i], 10) else null; + } + }.f; + if (std.mem.eql(u8, mode, "tf")) return tfMode(gpu, dir, if (rest.len > 2) rest[2] else return error.MissingArgument, try lineArg(rest, 3)); + if (std.mem.eql(u8, mode, "tfc")) return tfcMode(gpu, dir, if (rest.len > 2) rest[2] else return error.MissingArgument, try lineArg(rest, 3)); + if (std.mem.eql(u8, mode, "gen")) return genMode(gpu, dir, if (rest.len > 2) rest[2] else return error.MissingArgument, try lineArg(rest, 3)); + if (std.mem.eql(u8, mode, "bench")) return benchMode(gpu, dir, try lineArg(rest, 2)); + return error.UnknownMode; +} diff --git a/zig/tests/cuda/flashnext_exl3/model/flashnext_exl3_layer_tests.zig b/zig/tests/cuda/flashnext_exl3/model/flashnext_exl3_layer_tests.zig new file mode 100644 index 0000000000..cbc2cb518d --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/model/flashnext_exl3_layer_tests.zig @@ -0,0 +1,386 @@ +//! flashnext-exl3: one decoder layer through the family pass against the float64 layer (gate 5). +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const flash = @import("flashnext"); +const check = @import("../../check.zig"); +const fwd = @import("exl3_fwd_cpu.zig"); +const hc = @import("exl3_hc.zig"); +const moe = @import("exl3_moe.zig"); +const router_cpu = @import("exl3_router.zig"); +const route_cpu = @import("exl3_route.zig"); +const scan = @import("../pack/exl3_scan.zig"); +const place = @import("../pack/exl3_place.zig"); +const gdn_cpu = @import("exl3_gdn_cpu.zig"); +const attn_cpu = @import("exl3_attn_cpu.zig"); + +const X = flash.exl3; +const wt = X.prompt_wt; +const pl = X.prompt_load; +const batch = X.batch; +const Gpu = check.Gpu; +const st = core.safetensors; + +const dims = wt.dims; +const wide = wt.wide; +const streams = wt.streams; +const slots = wt.slots; +const topk = slots - 1; +const experts_n = wt.experts_n; +const shared = @import("flashnext_exl3_rig.zig"); +const typed = shared.typed; +const valueOf = shared.valueOf; +const f64Of = shared.f64Of; +const f32Of = shared.f32Of; +const f16Of = shared.f16Of; +const qmat = shared.qmat; +const bf16Bits = shared.bf16Bits; +const bfValue = shared.bfValue; +const Err = shared.Err; +const errOf = shared.errOf; +const download = shared.download; +const uploadTo = shared.uploadTo; +const RawPack = shared.RawPack; +const Pass = shared.Pass; +const Rig = shared.Rig; +const bf16Rows = shared.bf16Rows; +const default_ids = shared.default_ids; +const projections = @import("flashnext_exl3_proj_tests.zig").projections; + +// -- the float64 layer -------------------------------------------------------------------------------------------- + +const HcHost = struct { + norm: []f32, + down: []f32, + up: []f32, + inject: ?[]f32, + + fn load(a: std.mem.Allocator, raw: *const RawPack, comptime base: []const u8, args: anytype, with_inject: bool) !HcHost { + return .{ + .norm = try f32Of(a, try raw.get(base ++ ".hc_norm.weight", args)), + .down = try f32Of(a, try raw.get(base ++ ".input_mix_weight_down.weight", args)), + .up = try f32Of(a, try raw.get(base ++ ".input_mix_weight_up.weight", args)), + .inject = if (with_inject) try f32Of(a, try raw.get(base ++ ".block_inject_weight.weight", args)) else null, + }; + } + + fn weights(h: HcHost) hc.Weights { + return .{ .norm = h.norm, .down = h.down, .up = h.up, .inject = h.inject }; + } +}; + +const ExpertCtx = struct { + gpa: std.mem.Allocator, + raw: *const RawPack, + layer: usize, + + fn load(ctx: *anyopaque, e: usize, out: *[3]fwd.QMat) anyerror!void { + const self: *ExpertCtx = @ptrCast(@alignCast(ctx)); + const projs = [3][]const u8{ "gate_proj", "up_proj", "down_proj" }; + var made: usize = 0; + errdefer for (0..made) |i| out[i].deinit(self.gpa); + while (made < 3) : (made += 1) { + var nb: [160]u8 = undefined; + out[made] = try qmat(self.gpa, self.raw, try std.fmt.bufPrint(&nb, "model.language_model.layers.{d}.mlp.experts.{d}.{s}", .{ self.layer, e, projs[made] }), 1); + } + } +}; + +/// The oracle's layer (exl3_fwd_cpu.LayerW) from the pack's stored tensors, decoded by the Zig f64 decoder. +const Oracle = struct { + attn_hc: HcHost, + mlp_hc: HcHost, + gdn: fwd.GdnW = undefined, + attn: fwd.AttnW = undefined, + moe: fwd.MoeW, + ctx: ExpertCtx, + state: fwd.LayerState, + /// The mixer's state for the staged comparison (fed the device's own inputs), apart from the chained one. + staged: fwd.LayerState, + linear: bool, + + fn create(a: std.mem.Allocator, raw: *const RawPack, l: usize, offset: f32, cap: usize) !*Oracle { + const o = try a.create(Oracle); + o.linear = wt.linearLayer(l); + o.attn_hc = try HcHost.load(a, raw, "model.language_model.layers.{d}.attn_hyper_connection", .{l}, true); + o.mlp_hc = try HcHost.load(a, raw, "model.language_model.layers.{d}.mlp_hyper_connection", .{l}, true); + o.ctx = .{ .gpa = a, .raw = raw, .layer = l }; + var nb: [4][160]u8 = undefined; + o.moe = .{ + .gate16 = try f16Of(a, try raw.get("model.language_model.layers.{d}.mlp.gate.weight", .{l})), + .sgate16 = try f16Of(a, try raw.get("model.language_model.layers.{d}.mlp.shared_expert_gate.weight", .{l})), + .shared = .{ + try qmat(a, raw, try std.fmt.bufPrint(&nb[0], "model.language_model.layers.{d}.mlp.shared_expert.gate_proj", .{l}), 48), + try qmat(a, raw, try std.fmt.bufPrint(&nb[1], "model.language_model.layers.{d}.mlp.shared_expert.up_proj", .{l}), 48), + try qmat(a, raw, try std.fmt.bufPrint(&nb[2], "model.language_model.layers.{d}.mlp.shared_expert.down_proj", .{l}), 48), + }, + .source = .{ .ctx = &o.ctx, .load = ExpertCtx.load }, + }; + if (o.linear) { + const stem = "model.language_model.layers.{d}.linear_attn"; + o.gdn = .{ + .qkv = try qmat(a, raw, try std.fmt.bufPrint(&nb[0], stem ++ ".in_proj_qkv", .{l}), 48), + .z = try qmat(a, raw, try std.fmt.bufPrint(&nb[1], stem ++ ".in_proj_z", .{l}), 48), + .out = try qmat(a, raw, try std.fmt.bufPrint(&nb[2], stem ++ ".out_proj", .{l}), 48), + .a_w = try f64Of(a, try raw.get(stem ++ ".in_proj_a.weight", .{l})), + .b_w = try f64Of(a, try raw.get(stem ++ ".in_proj_b.weight", .{l})), + .params = .{ + .conv_w = try f64Of(a, try raw.get(stem ++ ".conv1d.weight", .{l})), + .a_log = try f64Of(a, try raw.get(stem ++ ".A_log", .{l})), + .dt_bias = try f64Of(a, try raw.get(stem ++ ".dt_bias", .{l})), + .norm_w = try f64Of(a, try raw.get(stem ++ ".norm.weight", .{l})), + }, + }; + o.state = .{ .gdn = try gdn_cpu.State.init(a, gdn_cpu.qwen) }; + o.staged = .{ .gdn = try gdn_cpu.State.init(a, gdn_cpu.qwen) }; + } else { + const stem = "model.language_model.layers.{d}.self_attn"; + var sc: [4][]f32 = undefined; + for ([_][]const u8{ "q_norm", "k_norm", "indexer.q_layernorm", "indexer.k_layernorm" }, 0..) |nm, i| { + var buf: [160]u8 = undefined; + const name = try std.fmt.bufPrint(&buf, "model.language_model.layers.{d}.self_attn.{s}.weight", .{ l, nm }); + sc[i] = try f32Of(a, try raw.get("{s}", .{name})); + for (sc[i]) |*v| v.* += offset; + } + o.attn = .{ + .q = try qmat(a, raw, try std.fmt.bufPrint(&nb[0], stem ++ ".q_proj", .{l}), 48), + .k = try qmat(a, raw, try std.fmt.bufPrint(&nb[1], stem ++ ".k_proj", .{l}), 48), + .v = try qmat(a, raw, try std.fmt.bufPrint(&nb[2], stem ++ ".v_proj", .{l}), 48), + .o = try qmat(a, raw, try std.fmt.bufPrint(&nb[3], stem ++ ".o_proj", .{l}), 48), + .idx = try qmat(a, raw, try std.fmt.bufPrint(&nb[0], stem ++ ".indexer.index_qk_proj", .{l}), 48), + .norms = .{ .q = sc[0], .k = sc[1], .iq = sc[2], .ik = sc[3], .inv = attn_cpu.invFreq() }, + }; + o.state = .{ .attn = try attn_cpu.Cache.init(a, cap) }; + o.staged = .{ .attn = try attn_cpu.Cache.init(a, cap) }; + } + return o; + } + + fn weights(o: *Oracle) fwd.LayerW { + return .{ .attn_hc = o.attn_hc.weights(), .mlp_hc = o.mlp_hc.weights(), .mixer = if (o.linear) .{ .gdn = &o.gdn } else .{ .attn = &o.attn }, .moe = &o.moe }; + } +}; + +// -- one decoder layer through the family's pass (gate 5) ---------------------------------------------------------- + +const Stage = struct { mixer: f64 = 0, streams: f64 = 0, moe: f64 = 0, out: f64 = 0, picks_same: usize = 0, rows: usize = 0 }; + +/// What the device fed each block: the read-outs' mixed rows and inject gates (bf16), attention's then the MoE's. +pub const Taps = struct { attn_mixed: []u16, attn_inj: []u16, mlp_mixed: []u16, mlp_inj: []u16 }; + +/// The pass's per-layer sequence (cuda_batch.zig Pass.chunk's loop body), stepped with read-backs between blocks. +pub fn deviceLayer(gpu: Gpu, p: *Pass, lw: *const pl.Layer, l: usize, ids: []const u32, pos0: usize, h_in: []const u16, tr: fwd.Trace, h_out: []u16, taps: Taps) !void { + const s = p.s; + const n = ids.len; + p.rows = n; + p.pos0 = pos0; + p.pending = null; + try p.stage(ids); + try uploadTo(gpu, p.stream, s.h, std.mem.sliceAsBytes(h_in)); + try p.writeback(0, n); + try p.readout(lw.attn_hc, 0, n); + try tap(gpu, p, s.mixed, taps.attn_mixed[0 .. n * dims]); + try tap(gpu, p, s.inj, taps.attn_inj[0 .. n * streams]); + if (lw.linear) try p.gdnBlock(lw, l, wt.linearIndex(l), .{}) else try p.attnBlock(lw, wt.attnIndex(l)); + const branch = try download(f32, gpu, p.stream, s.part, n * dims); + defer gpu.gpa.free(branch); + for (tr.attn_branch[0 .. n * dims], branch) |*d, v| d.* = v; + try p.tri.hcWriteback(s.h, s.h, s.pss, s.part, "*fp32", s.inj, s.dummy, "*bf16", s.dummy, "*bf16", n * dims, n, dims, streams, 4, 1, 1, wt.out_sk); + const mid = try download(u16, gpu, p.stream, s.h, n * wide); + defer gpu.gpa.free(mid); + for (tr.h_mid[0 .. n * wide], mid) |*d, v| d.* = bfValue(v); + try p.readout(lw.mlp_hc, 0, n); + try tap(gpu, p, s.mixed, taps.mlp_mixed[0 .. n * dims]); + try tap(gpu, p, s.inj, taps.mlp_inj[0 .. n * streams]); + try p.moeBlock(lw); + const picks = try download(i32, gpu, p.stream, s.picks, n * slots); + defer gpu.gpa.free(picks); + const wts = try download(f32, gpu, p.stream, s.wts, n * slots); + defer gpu.gpa.free(wts); + const y = try download(f32, gpu, p.stream, s.y, n * slots * dims); + defer gpu.gpa.free(y); + for (0..n) |r| { + for (0..topk) |k| tr.picks[r * topk + k] = @intCast(picks[r * slots + k]); + for (0..dims) |d| { + var acc: f64 = 0; + for (0..slots) |k| acc += @as(f64, wts[r * slots + k]) * y[(r * slots + k) * dims + d]; + tr.mlp_branch[r * dims + d] = acc; + } + } + p.pending = l; + try p.writeback(0, n); + const out = try download(u16, gpu, p.stream, s.h, n * wide); + defer gpu.gpa.free(out); + @memcpy(h_out, out); +} + +fn tap(gpu: Gpu, p: *Pass, ptr: u64, out: []u16) !void { + const got = try download(u16, gpu, p.stream, ptr, out.len); + defer gpu.gpa.free(got); + @memcpy(out, got); +} + +fn toF64(a: std.mem.Allocator, bits: []const u16) ![]f64 { + const out = try a.alloc(f64, bits.len); + for (out, bits) |*o, b| o.* = bfValue(b); + return out; +} + +/// Each block against the float64 oracle fed the DEVICE's own input to that block: its error alone. +const Staged = struct { attn_mixed: f64 = 0, attn_inj: f64 = 0, mixer: f64 = 0, mlp_mixed: f64 = 0, mlp_inj: f64 = 0, moe: f64 = 0, writeback: f64 = 0 }; + +pub fn layerCheck(gpu: Gpu, r: *Rig, raw: *const RawPack, lw: *const pl.Layer, l: usize, tokens: []const u32) !void { + var arena: std.heap.ArenaAllocator = .init(gpu.gpa); + defer arena.deinit(); + const a = arena.allocator(); + const prefill = 8; + var t0 = check.now(gpu.io); + const o = try Oracle.create(a, raw, l, r.store.norm_offset, r.cap); + std.debug.print(" float64 layer {d} decoded in {d:.0} ms (offset {d})\n", .{ l, @as(f64, @floatFromInt(check.now(gpu.io) - t0)) / 1e6, r.store.norm_offset }); + try r.reset(); + var p = r.pass((lw)[0..1]); + const R = prefill; + const mk = struct { + fn trace(al: std.mem.Allocator) !fwd.Trace { + return .{ .attn_branch = try al.alloc(f64, R * dims), .mlp_branch = try al.alloc(f64, R * dims), .picks = try al.alloc(u16, R * topk), .h_mid = try al.alloc(f64, R * wide) }; + } + }; + const dt = try mk.trace(a); + const ot = try mk.trace(a); + const table = try raw.get("model.language_model.embed_tokens.weight", .{}); + const calls = [_][2]usize{ .{ 0, prefill }, .{ 8, 1 }, .{ 9, 1 }, .{ 10, 1 }, .{ 11, 1 } }; + var worst: Stage = .{}; + var alone: Staged = .{}; + var device_ms: f64 = 0; + const taps: Taps = .{ .attn_mixed = try a.alloc(u16, R * dims), .attn_inj = try a.alloc(u16, R * streams), .mlp_mixed = try a.alloc(u16, R * dims), .mlp_inj = try a.alloc(u16, R * streams) }; + for (calls) |c| { + const toks = tokens[c[0] .. c[0] + c[1]]; + const n = c[1]; + // the layer's input: each token's embedding row in all four streams (bf16, as the pass's expand writes it) + const hb = try a.alloc(u16, n * wide); + const ho = try a.alloc(f64, n * wide); + for (toks, 0..) |tk, t| for (0..dims) |i| { + const bits = std.mem.readInt(u16, table.bytes[(tk * dims + i) * 2 ..][0..2], .little); + for (0..streams) |sidx| { + hb[t * wide + sidx * dims + i] = bits; + ho[t * wide + sidx * dims + i] = bfValue(bits); + } + }; + const hg = try a.alloc(u16, n * wide); + t0 = check.now(gpu.io); + try deviceLayer(gpu, &p, lw, l, toks, c[0], hb, dt, hg, taps); + device_ms += @as(f64, @floatFromInt(check.now(gpu.io) - t0)) / 1e6; + // staged: every block from the device's own input (before the chained oracle updates `ho` in place) + { + const w = o.weights(); + const off = r.store.norm_offset; + const mixed = try a.alloc(f64, n * dims); + const inj = try a.alloc(f64, n * streams); + try fwd.hcPreRows(a, w.attn_hc, off, ho, n, mixed, inj); + const e_am = errOf(@as([]const u16, taps.attn_mixed[0 .. n * dims]), mixed); + const e_ai = errOf(@as([]const u16, taps.attn_inj[0 .. n * streams]), inj); + const dev_mixed = try toF64(a, taps.attn_mixed[0 .. n * dims]); + const branch = try a.alloc(f64, n * dims); + switch (w.mixer) { + .gdn => |g| try fwd.gdnBranch(a, g, o.staged.gdn, dev_mixed, n, branch), + .attn => |at| try fwd.attnBranch(a, at, &o.staged.attn, c[0], dev_mixed, n, branch), + } + const e_mx = errOf(dt.attn_branch[0 .. n * dims], branch); + const dev_mid = dt.h_mid[0 .. n * wide]; + try fwd.hcPreRows(a, w.mlp_hc, off, dev_mid, n, mixed, inj); + const e_mm = errOf(@as([]const u16, taps.mlp_mixed[0 .. n * dims]), mixed); + const e_mi = errOf(@as([]const u16, taps.mlp_inj[0 .. n * streams]), inj); + const dev_mlp = try toF64(a, taps.mlp_mixed[0 .. n * dims]); + const moe_out = try a.alloc(f64, n * dims); + const picks = try a.alloc(u16, n * topk); + try fwd.moeBranch(a, w.moe, dev_mlp, n, moe_out, picks); + var same_picks = true; + for (0..n) |t| { + var x = dt.picks[t * topk ..][0..topk].*; + var y = picks[t * topk ..][0..topk].*; + std.mem.sort(u16, &x, {}, std.sort.asc(u16)); + std.mem.sort(u16, &y, {}, std.sort.asc(u16)); + same_picks = same_picks and std.mem.eql(u16, &x, &y); + } + const e_moe = errOf(dt.mlp_branch[0 .. n * dims], moe_out); + // the write-back alone: the device's streams after the mixer plus its MoE branch times its inject gates + const wb = try a.alloc(f64, n * wide); + for (0..n) |t| for (0..streams) |sidx| for (0..dims) |d| { + const at = t * wide + sidx * dims + d; + wb[at] = dev_mid[at] + dt.mlp_branch[t * dims + d] * bfValue(taps.mlp_inj[t * streams + sidx]); + }; + const e_wb = errOf(@as([]const u16, hg), wb); + std.debug.print(" staged (device input to each block): attn readout mixed rel {e:.3}, inject rel {e:.3}; mixer rel {e:.3}; mlp readout mixed rel {e:.3}, inject rel {e:.3}; MoE rel {e:.3} (picks {s}); write-back rel {e:.3}\n", .{ e_am.rel, e_ai.rel, e_mx.rel, e_mm.rel, e_mi.rel, e_moe.rel, if (same_picks) "equal" else "DIFFER", e_wb.rel }); + alone.attn_mixed = @max(alone.attn_mixed, e_am.rel); + alone.attn_inj = @max(alone.attn_inj, e_ai.rel); + alone.mixer = @max(alone.mixer, e_mx.rel); + alone.mlp_mixed = @max(alone.mlp_mixed, e_mm.rel); + alone.mlp_inj = @max(alone.mlp_inj, e_mi.rel); + if (same_picks) alone.moe = @max(alone.moe, e_moe.rel); + alone.writeback = @max(alone.writeback, e_wb.rel); + } + try fwd.layerRows(a, o.weights(), &o.state, c[0], r.store.norm_offset, ho, n, ot); + const e_mix = errOf(dt.attn_branch[0 .. n * dims], ot.attn_branch[0 .. n * dims]); + const e_mid = errOf(dt.h_mid[0 .. n * wide], ot.h_mid[0 .. n * wide]); + var same: usize = 0; + for (0..n) |t| { + var x = dt.picks[t * topk ..][0..topk].*; + var y = ot.picks[t * topk ..][0..topk].*; + std.mem.sort(u16, &x, {}, std.sort.asc(u16)); + std.mem.sort(u16, &y, {}, std.sort.asc(u16)); + same += @intFromBool(std.mem.eql(u16, &x, &y)); + } + const e_moe = errOf(dt.mlp_branch[0 .. n * dims], ot.mlp_branch[0 .. n * dims]); + const e_out = errOf(@as([]const u16, hg), ho); + std.debug.print(" rows {d}..{d}: mixer branch rel {e:.3} (peak {d:.3}); streams after mixer rel {e:.3}; MoE picks same set {d}/{d}, MoE branch rel {e:.3} (peak {d:.3}); layer output streams rel {e:.3} (peak {d:.2}, max_abs {e:.2})\n", .{ c[0], c[0] + n - 1, e_mix.rel, e_mix.peak, e_mid.rel, same, n, e_moe.rel, e_moe.peak, e_out.rel, e_out.peak, e_out.max_abs }); + worst.mixer = @max(worst.mixer, e_mix.rel); + worst.streams = @max(worst.streams, e_mid.rel); + worst.picks_same += same; + worst.rows += n; + if (same == n) { + worst.moe = @max(worst.moe, e_moe.rel); + worst.out = @max(worst.out, e_out.rel); + } + } + std.debug.print(" device time for the 5 calls (read-backs between blocks included): {d:.0} ms\n", .{device_ms}); + std.debug.print(" worst staged: attn readout mixed {e:.3} inject {e:.3}; mixer {e:.3}; mlp readout mixed {e:.3} inject {e:.3}; MoE {e:.3}; write-back {e:.3}\n", .{ alone.attn_mixed, alone.attn_inj, alone.mixer, alone.mlp_mixed, alone.mlp_inj, alone.moe, alone.writeback }); + std.debug.print(" worst chained: mixer {e:.3}, streams {e:.3}, MoE {e:.3}, output {e:.3}\n", .{ worst.mixer, worst.streams, worst.moe, worst.out }); + // Each block alone, fed the device's own input: its bf16 output and the roundings inside it. + const staged = [_]f64{ alone.attn_mixed, alone.attn_inj, alone.mixer, alone.mlp_mixed, alone.mlp_inj, alone.moe, alone.writeback }; + for (staged) |v| try check.expect(v <= 0.01, "layer {d}: a block alone differs from float64 by {d:.4} of its peak (> 0.01)", .{ l, v }); + // Chained from the embedding: the old engine's gates (mixer 1% GDN / 3% attention, streams 1%, MoE 2%) hold; its + const mixer_tol: f64 = if (wt.linearLayer(l)) 0.01 else 0.03; + try check.expect(worst.mixer <= mixer_tol, "layer {d}: mixer branch differs from the oracle ({d:.4} > {d})", .{ l, worst.mixer, mixer_tol }); + try check.expect(worst.streams <= 0.01, "layer {d}: streams after the mixer differ ({d:.4})", .{ l, worst.streams }); + try check.expect(worst.moe <= 0.02, "layer {d}: MoE branch differs from the oracle ({d:.4})", .{ l, worst.moe }); + try check.expect(worst.out <= 0.02, "layer {d}: layer output streams differ ({d:.4})", .{ l, worst.out }); + if (worst.out > 0.01) std.debug.print(" NOTE layer {d}: chained output {e:.3} is above the old engine's 1% gate (f32 streams there)\n", .{ l, worst.out }); + check.pass("layer {d} ({s}) through the family's pass (writeback, readout, {s}, write-back, readout, EXL3 moeBlock, write-back) vs float64: mixer rel {e:.3}, streams rel {e:.3}, MoE rel {e:.3}, output rel {e:.3}; picks equal on {d}/{d} rows", .{ l, if (wt.linearLayer(l)) "GDN" else "attention", if (wt.linearLayer(l)) "gdnBlock" else "attnBlock", worst.mixer, worst.streams, worst.moe, worst.out, worst.picks_same, worst.rows }); +} + +pub fn layerMode(gpu: Gpu, dir: []const u8, which: []const usize) !void { + var rig: Rig = undefined; + try rig.init(gpu, dir, 8, 64); + defer rig.deinit(); + var raw = try RawPack.open(gpu, dir); + defer raw.close(gpu); + var prng = std.Random.DefaultPrng.init(0xe713_c03e); + var failed: usize = 0; + for (which) |l| { + const t0 = check.now(gpu.io); + var lw = try pl.loadLayerRank(gpu.gpa, @constCast(gpu.d), &rig.stream, &rig.store, rig.pack_fn, l, 1, 0, null); + defer lw.free(gpu.gpa); + std.debug.print("== layer {d} ({s}): the family's loader (cuda_prompt_load.loadLayerRank) in {d:.0} ms, every expert resident\n", .{ l, if (wt.linearLayer(l)) "GDN" else "attention", @as(f64, @floatFromInt(check.now(gpu.io) - t0)) / 1e6 }); + try check.expect(lw.exl3 != null, "layer {d}: experts did not load as EXL3", .{l}); + var arena: std.heap.ArenaAllocator = .init(gpu.gpa); + defer arena.deinit(); + // every layer runs, a failed check reported and counted + projections(gpu, &rig, &raw, &lw, l, arena.allocator(), prng.random()) catch |err| if (err == error.TestFailed) { + failed += 1; + } else return err; + layerCheck(gpu, &rig, &raw, &lw, l, &default_ids) catch |err| if (err == error.TestFailed) { + failed += 1; + } else return err; + } + try check.expect(failed == 0, "flashnext-exl3 layer: {d} checks failed", .{failed}); +} diff --git a/zig/tests/cuda/flashnext_exl3/model/flashnext_exl3_names_tests.zig b/zig/tests/cuda/flashnext_exl3/model/flashnext_exl3_names_tests.zig new file mode 100644 index 0000000000..49a8d78eb7 --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/model/flashnext_exl3_names_tests.zig @@ -0,0 +1,374 @@ +//! flashnext-exl3: tensor names, norm statistics and the router against the float64 oracle. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const flash = @import("flashnext"); +const check = @import("../../check.zig"); +const fwd = @import("exl3_fwd_cpu.zig"); +const hc = @import("exl3_hc.zig"); +const moe = @import("exl3_moe.zig"); +const router_cpu = @import("exl3_router.zig"); +const route_cpu = @import("exl3_route.zig"); +const scan = @import("../pack/exl3_scan.zig"); +const place = @import("../pack/exl3_place.zig"); +const gdn_cpu = @import("exl3_gdn_cpu.zig"); +const attn_cpu = @import("exl3_attn_cpu.zig"); + +const X = flash.exl3; +const wt = X.prompt_wt; +const pl = X.prompt_load; +const batch = X.batch; +const Gpu = check.Gpu; +const st = core.safetensors; + +const dims = wt.dims; +const wide = wt.wide; +const streams = wt.streams; +const slots = wt.slots; +const topk = slots - 1; +const experts_n = wt.experts_n; +const shared = @import("flashnext_exl3_rig.zig"); +const typed = shared.typed; +const valueOf = shared.valueOf; +const f64Of = shared.f64Of; +const f32Of = shared.f32Of; +const f16Of = shared.f16Of; +const qmat = shared.qmat; +const bf16Bits = shared.bf16Bits; +const bfValue = shared.bfValue; +const Err = shared.Err; +const errOf = shared.errOf; +const download = shared.download; +const uploadTo = shared.uploadTo; +const RawPack = shared.RawPack; +const Pass = shared.Pass; +const Rig = shared.Rig; +const bf16Rows = shared.bf16Rows; +const default_ids = shared.default_ids; + +// -- names ---------------------------------------------------------------------------------------------------------- + +const Names = struct { + gpa: std.mem.Allocator, + store: *wt.Store, + claimed: std.StringHashMapUnmanaged(void) = .empty, + missing: usize = 0, + matrices: usize = 0, + dense: usize = 0, + widened: usize = 0, + by_halves: [17]usize = @splat(0), + + fn claim(n: *Names, requested: []const u8) !void { + var buf: [512]u8 = undefined; + const key = n.store.resolveKey(requested, &buf) catch { + n.missing += 1; + if (n.missing <= 20) std.debug.print(" missing: {s}\n", .{requested}); + return; + }; + if (n.claimed.contains(key)) return; + try n.claimed.put(n.gpa, try n.gpa.dupe(u8, key), {}); + } + + /// An EXL3 matrix: its four parts claimed and the matrix admitted by the family's own reader (marker, widths). + fn matrix(n: *Names, comptime fmt: []const u8, args: anytype) !void { + var buf: [256]u8 = undefined; + const prefix = try std.fmt.bufPrint(&buf, fmt, args); + var pb: [280]u8 = undefined; + for ([_][]const u8{ "trellis", "suh", "svh", "mul1" }) |part| try n.claim(try std.fmt.bufPrint(&pb, "{s}.{s}", .{ prefix, part })); + if (!X.load.has(n.store, prefix)) { + n.missing += 1; + std.debug.print(" not EXL3: {s}\n", .{prefix}); + return; + } + const m = try X.load.mat(n.store, prefix); + n.by_halves[m.halves] += 1; + n.matrices += 1; + } + + /// A dense tensor as the loaders read it: through Store.tensor, so an F16 `.weight` is widened (exact or refused). + fn tensor(n: *Names, comptime fmt: []const u8, args: anytype) !void { + var buf: [256]u8 = undefined; + const name = try std.fmt.bufPrint(&buf, fmt, args); + try n.claim(name); + if (!n.store.has(name)) return; + const before = n.store.widened.items.len; + _ = try n.store.tensor(name); + if (n.store.widened.items.len > before) n.widened += 1; + n.dense += 1; + } + + fn hcw(n: *Names, base: []const u8, inject: bool) !void { + try n.tensor("{s}.hc_norm.weight", .{base}); + try n.tensor("{s}.input_mix_weight_down.weight", .{base}); + try n.tensor("{s}.input_mix_weight_up.weight", .{base}); + if (inject) try n.tensor("{s}.block_inject_weight.weight", .{base}); + } + + /// What cuda_prompt_load.loadLayerAtRank (and the EXL3 hooks it reaches) read for one decoder layer at `base`. + fn layer(n: *Names, base: []const u8, linear: bool, with_ple: bool) !void { + var hb: [128]u8 = undefined; + try n.hcw(try std.fmt.bufPrint(&hb, "{s}.attn_hyper_connection", .{base}), true); + try n.hcw(try std.fmt.bufPrint(&hb, "{s}.mlp_hyper_connection", .{base}), true); + try n.tensor("{s}.mlp.gate.weight", .{base}); + try n.tensor("{s}.mlp.shared_expert_gate.weight", .{base}); + if (linear) { + for ([_][]const u8{ "in_proj_qkv", "in_proj_z", "out_proj" }) |p| try n.matrix("{s}.linear_attn.{s}", .{ base, p }); + for ([_][]const u8{ "in_proj_b.weight", "in_proj_a.weight", "conv1d.weight", "A_log", "dt_bias", "norm.weight" }) |p| try n.tensor("{s}.linear_attn.{s}", .{ base, p }); + } else { + for ([_][]const u8{ "q_proj", "k_proj", "v_proj", "indexer.index_qk_proj", "o_proj" }) |p| try n.matrix("{s}.self_attn.{s}", .{ base, p }); + for ([_][]const u8{ "q_norm", "k_norm", "indexer.q_layernorm", "indexer.k_layernorm" }) |p| try n.tensor("{s}.self_attn.{s}.weight", .{ base, p }); + } + if (with_ple) { + for ([_][]const u8{ "key_proj", "value_proj", "norm_key", "norm_query", "norm_conv", "conv1d" }) |p| try n.tensor("{s}.ple.{s}.weight", .{ base, p }); + } + for ([_][]const u8{ "gate_proj", "up_proj", "down_proj" }) |p| try n.matrix("{s}.mlp.shared_expert.{s}", .{ base, p }); + for (0..experts_n) |e| { + for ([_][]const u8{ "gate_proj", "up_proj", "down_proj" }) |p| try n.matrix("{s}.mlp.experts.{d}.{s}", .{ base, e, p }); + if (e % 64 == 63) n.store.release(); + } + n.store.release(); + } +}; + +pub fn names(gpu: Gpu, dir: []const u8) !void { + const gpa = gpu.gpa; + var store = try wt.Store.open(gpa, gpu.io, dir); + defer store.deinit(); + var n: Names = .{ .gpa = gpa, .store = &store }; + defer { + var it = n.claimed.keyIterator(); + while (it.next()) |k| gpa.free(k.*); + n.claimed.deinit(gpa); + } + const t0 = check.now(gpu.io); + // the engine (cuda_prompt.zig Engine.initRank): embedding, the final mixer, the head (EXL3 here: plan step 9) + try n.tensor("language_model.model.embed_tokens.weight", .{}); + try n.hcw("language_model.model.hyper_connection_mixer", false); + try n.matrix("language_model.lm_head", .{}); + for (0..wt.layers_n) |l| { + var bb: [64]u8 = undefined; + try n.layer(try std.fmt.bufPrint(&bb, "language_model.model.layers.{d}", .{l}), wt.linearLayer(l), l == wt.ple_layer); + } + // the MTP head (cuda_mtp.zig Head.init): norms, two EXL3 fc projections, its mixer (the patch file), one layer + try n.tensor("language_model.mtp.pre_fc_norm_embedding.weight", .{}); + try n.tensor("language_model.mtp.pre_fc_norm_hidden.weight", .{}); + try n.matrix("language_model.mtp.fc_embedding", .{}); + try n.matrix("language_model.mtp.fc_hidden", .{}); + try n.hcw("language_model.mtp.hyper_connection_mixer", false); + try n.layer("language_model.mtp.layers.0", false, false); + // everything the index (and the patch) names: claimed, or the vision tower; and the header rules' homes agree + const wm = store.parsed.value.object.get("weight_map").?.object; + var vision: usize = 0; + var unclaimed: usize = 0; + var homes: [8]usize = @splat(0); + var disagree: usize = 0; + var it = wm.iterator(); + while (it.next()) |kv| { + const name = kv.key_ptr.*; + const home = place.home(.{ .name = name, .rank = if (std.mem.endsWith(u8, name, ".trellis")) 3 else 1, .file = kv.value_ptr.string }); + homes[@backingInt(home)] += 1; + if (n.claimed.contains(name)) { + // a routed expert's part must be the routed bank's, a shared expert's part the resident one's + if (std.mem.indexOf(u8, name, ".mlp.experts.") != null and home != .routed_host) disagree += 1; + if (std.mem.indexOf(u8, name, ".shared_expert.") != null and home != .shared_resident) disagree += 1; + continue; + } + if (std.mem.startsWith(u8, name, "model.visual.")) { + vision += 1; + continue; + } + unclaimed += 1; + if (unclaimed <= 20) std.debug.print(" unclaimed: {s}\n", .{name}); + } + const ms = @as(f64, @floatFromInt(check.now(gpu.io) - t0)) / 1e6; + std.debug.print("names: {d} indexed (+ the mixer patch), {d} claimed, {d} vision, {d} unclaimed, {d} missing; {d} EXL3 matrices admitted (mul1 marker, widths from the trellis), {d} dense tensors read, {d} F16 weights widened to BF16 exactly; {d:.0} ms\n", .{ wm.count(), n.claimed.count(), vision, unclaimed, n.missing, n.matrices, n.dense, n.widened, ms }); + std.debug.print(" EXL3 widths (halves: count):", .{}); + for (n.by_halves, 0..) |c, h| if (c > 0) std.debug.print(" {d}:{d}", .{ h, c }); + std.debug.print("\n header homes (exl3_place):", .{}); + for (homes, 0..) |c, i| if (c > 0) std.debug.print(" {s}={d}", .{ @tagName(@as(place.Home, @fromBackingInt(@intCast(i)))), c }); + std.debug.print("\n", .{}); + try check.expect(n.missing == 0, "names: {d} tensors the family reads do not resolve", .{n.missing}); + try check.expect(unclaimed == 0, "names: {d} indexed tensors no family loader reads (vision excepted)", .{unclaimed}); + try check.expect(disagree == 0, "names: {d} expert tensors whose exl3_place home disagrees", .{disagree}); + // the family's device sizing (cuda_exl3_memory.zig) on the same pack + const gib = 1024.0 * 1024.0 * 1024.0; + const on_host = try X.memory.deviceBytes(gpa, gpu.io, dir, .{}); + const resident = try X.memory.deviceBytes(gpa, gpu.io, dir, .{ .routed_on_host = false }); + std.debug.print(" device bytes: {d:.2} GiB with routed experts on the host, {d:.2} GiB all resident\n", .{ @as(f64, @floatFromInt(on_host)) / gib, @as(f64, @floatFromInt(resident)) / gib }); + check.pass("names: every tensor the family reads resolves ({d}), nothing unclaimed but {d} vision tensors; {d} EXL3 matrices admitted; {d} F16 weights widened exactly", .{ n.claimed.count(), vision, n.matrices, n.widened }); +} + +// -- norms ---------------------------------------------------------------------------------------------------------- + +fn meanOf(t: st.Tensor) f64 { + var s: f64 = 0; + const count = t.bytes.len / 2; + for (0..count) |i| s += valueOf(t, i); + return s / @as(f64, @floatFromInt(count)); +} + +pub fn norms(gpu: Gpu, dir: []const u8) !void { + const gpa = gpu.gpa; + var store = try wt.Store.open(gpa, gpu.io, dir); + defer store.deinit(); + try store.detectNorms(); + var raw = try RawPack.open(gpu, dir); + defer raw.close(gpu); + var means: [wt.layers_n]f64 = undefined; + for (&means, 0..) |*m, l| m.* = meanOf(try raw.get("model.language_model.layers.{d}.attn_hyper_connection.hc_norm.weight", .{l})); + var sorted = means; + std.mem.sort(f64, &sorted, {}, std.sort.asc(f64)); + std.debug.print("norms: 48 attn hc_norm means min {d:.4} median {d:.4} max {d:.4} -> offset {d} (cuda_norms.offset via Store.detectNorms)\n", .{ sorted[0], (sorted[23] + sorted[24]) / 2, sorted[47], store.norm_offset }); + try check.expect(store.norm_offset == 1, "norms: the pack stores gamma - 1; the convention came out {d}", .{store.norm_offset}); + // the independent cross-check: the indexer norms' own 24 means locate the same storage + var idx: [24]f32 = undefined; + var c: usize = 0; + for (0..wt.layers_n) |l| { + if (wt.linearLayer(l)) continue; + for ([_][]const u8{ "q_layernorm", "k_layernorm" }) |nm| { + idx[c] = @floatCast(meanOf(try raw.get("model.language_model.layers.{d}.self_attn.indexer.{s}.weight", .{ l, nm }))); + c += 1; + } + } + const indexer = try scan.centredOffset(&idx); + std.debug.print("norms: indexer q/k_layernorm means over the 12 attention layers -> exl3_scan.centredOffset {d}\n", .{indexer}); + try check.expect(indexer == store.norm_offset, "norms: the indexer norms locate offset {d}, the hc_norm means {d}", .{ indexer, store.norm_offset }); + // every centred norm the family scales, MTP's included, loads through Store.normScale with that one offset + const mtp_names = [_][]const u8{ "language_model.mtp.pre_fc_norm_embedding.weight", "language_model.mtp.pre_fc_norm_hidden.weight", "language_model.mtp.hyper_connection_mixer.hc_norm.weight", "language_model.mtp.layers.0.attn_hyper_connection.hc_norm.weight", "language_model.mtp.layers.0.mlp_hyper_connection.hc_norm.weight", "language_model.mtp.layers.0.self_attn.q_norm.weight", "language_model.mtp.layers.0.self_attn.k_norm.weight", "language_model.mtp.layers.0.self_attn.indexer.q_layernorm.weight", "language_model.mtp.layers.0.self_attn.indexer.k_layernorm.weight", "language_model.model.layers.1.ple.norm_key.weight", "language_model.model.layers.1.ple.norm_query.weight", "language_model.model.layers.1.ple.norm_conv.weight", "language_model.model.layers.3.self_attn.q_norm.weight", "language_model.model.layers.3.self_attn.k_norm.weight", "language_model.model.hyper_connection_mixer.hc_norm.weight" }; + for (mtp_names) |name| { + const scales = try store.normScale(gpa, try store.tensor(name)); + defer gpa.free(scales); + var sum: f64 = 0; + var lo: f32 = std.math.inf(f32); + for (scales) |v| { + sum += v; + lo = @min(lo, v); + } + std.debug.print(" {s:<72} stored mean {d:>7.4} -> scale mean {d:.4} (min {d:.4})\n", .{ name, sum / @as(f64, @floatFromInt(scales.len)) - store.norm_offset, sum / @as(f64, @floatFromInt(scales.len)), lo }); + } + store.release(); + check.pass("norms: one convention (offset {d}) from the 48 hc_norm means, confirmed by the indexer norms; every centred norm incl. MTP's scales finite", .{store.norm_offset}); +} + +// -- router --------------------------------------------------------------------------------------------------------- + +pub fn router(gpu: Gpu, dir: []const u8, layers: []const usize, rows: usize) !void { + const gpa = gpu.gpa; + var store = try wt.Store.open(gpa, gpu.io, dir); + defer store.deinit(); + var raw = try RawPack.open(gpu, dir); + defer raw.close(gpu); + var glue = try flash.Glue.load(gpu.d); + defer glue.deinit(); + var s = try cuda.Stream.init(gpu.d, true); + defer s.deinit(); + var prng = std.Random.DefaultPrng.init(0x5eed_2026); + const rng = prng.random(); + var worst_rows_diff: usize = 0; + for (layers) |l| { + // the family's router buffer, as cuda_prompt_load builds it: mlp.gate's 512 rows then the shared gate. + var nb: [128]u8 = undefined; + const gate = try store.tensor(try std.fmt.bufPrint(&nb, "language_model.model.layers.{d}.mlp.gate.weight", .{l})); + try check.expect(gate.is(.bf16, &.{ experts_n, dims }), "router: layer {d} gate is not the loaders' bf16 [512, 2560]", .{l}); + const w = try gpa.alloc(u8, (experts_n + 1) * dims * 2); + defer gpa.free(w); + @memcpy(w[0 .. experts_n * dims * 2], gate.bytes); + const sg = try store.tensor(try std.fmt.bufPrint(&nb, "language_model.model.layers.{d}.mlp.shared_expert_gate.weight", .{l})); + @memcpy(w[experts_n * dims * 2 ..], sg.bytes); + store.release(); + var wd = try cuda.DeviceBuffer.fromHost(gpu.d, w); + defer wd.free(); + // the oracle's weights: the pack's F16 bits, never the widened copy + const g16 = try f16Of(gpa, try raw.get("model.language_model.layers.{d}.mlp.gate.weight", .{l})); + defer gpa.free(g16); + const s16 = try f16Of(gpa, try raw.get("model.language_model.layers.{d}.mlp.shared_expert_gate.weight", .{l})); + defer gpa.free(s16); + // rows: N(0, sigma) per row, sigma in {0.25, 1, 4}, rounded to bf16 as the pass's `mixed` is + const xb = try gpa.alloc(u16, rows * dims); + defer gpa.free(xb); + for (0..rows) |r| { + const sigma: f32 = switch (r % 3) { + 0 => 0.25, + 1 => 1, + else => 4, + }; + for (0..dims) |i| xb[r * dims + i] = bf16Bits(rng.floatNorm(f32) * sigma); + } + var xd = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(xb)); + defer xd.free(); + var ld = try cuda.DeviceBuffer.alloc(gpu.d, rows * (experts_n + 1) * 4); + defer ld.free(); + var pd = try cuda.DeviceBuffer.alloc(gpu.d, rows * slots * 4); + defer pd.free(); + var wdd = try cuda.DeviceBuffer.alloc(gpu.d, rows * slots * 4); + defer wdd.free(); + try glue.routerMany(s, xd.ptr, wd.ptr, ld.ptr, dims, rows); + try glue.topkMany(s, ld.ptr, pd.ptr, wdd.ptr, rows); + const gl = try download(f32, gpu, s, ld.ptr, rows * (experts_n + 1)); + defer gpa.free(gl); + const gp = try download(i32, gpu, s, pd.ptr, rows * slots); + defer gpa.free(gp); + const gw = try download(f32, gpu, s, wdd.ptr, rows * slots); + defer gpa.free(gw); + var diff_rows: usize = 0; + var logit_err: f64 = 0; + var min_gap: f64 = std.math.inf(f64); + var worst_w: f64 = 0; + var worst_sw: f64 = 0; + var f32_agree: usize = 0; + const x = try gpa.alloc(f64, dims); + defer gpa.free(x); + const xf = try gpa.alloc(f32, dims); + defer gpa.free(xf); + var logits: [experts_n + 1]f64 = undefined; + for (0..rows) |r| { + for (x, xf, 0..) |*v, *vf, i| { + v.* = bfValue(xb[r * dims + i]); + vf.* = @floatCast(v.*); + } + try moe.logitsF16(g16, experts_n, s16, x, &logits); + var picks: [topk]u16 = undefined; + var wts: [topk]f64 = undefined; + var sw: f64 = 0; + try moe.routeFromLogits(&logits, topk, &picks, &wts, &sw); + for (0..experts_n + 1) |i| logit_err = @max(logit_err, @abs(@as(f64, gl[r * (experts_n + 1) + i]) - logits[i])); + var sorted = logits[0..experts_n].*; + std.mem.sort(f64, &sorted, {}, std.sort.desc(f64)); + const gap = sorted[topk - 1] - sorted[topk]; + min_gap = @min(min_gap, gap); + var a: [topk]u16 = undefined; + for (&a, 0..) |*v, k| v.* = @intCast(gp[r * slots + k]); + var b = picks; + std.mem.sort(u16, &a, {}, std.sort.asc(u16)); + std.mem.sort(u16, &b, {}, std.sort.asc(u16)); + if (!std.mem.eql(u16, &a, &b)) { + diff_rows += 1; + if (diff_rows <= 10) std.debug.print(" layer {d} row {d}: picks differ; f64 logit gap between the 10th and 11th {e:.3}\n", .{ l, r, gap }); + } else { + for (0..topk) |k| { + // the same expert's weight: the device's bf16 against the float64 softmax + const id = gp[r * slots + k]; + for (picks, wts) |pk, wv| if (pk == id) { + worst_w = @max(worst_w, @abs(@as(f64, gw[r * slots + k]) - wv)); + }; + } + try check.expect(gp[r * slots + topk] == experts_n, "router: the shared slot holds {d}", .{gp[r * slots + topk]}); + worst_sw = @max(worst_sw, @abs(@as(f64, gw[r * slots + topk]) - sw)); + } + // the f32 host reference (exl3_router.launchDtype(.f16)) on the same row: its picks as a set + var l32: [experts_n + 1]f32 = undefined; + var p32: [topk + 1]u16 = undefined; + var w32: [topk + 1]f32 = undefined; + var scratch: [experts_n]f32 = undefined; + try router_cpu.launchDtype(.f16, 0, std.mem.sliceAsBytes(g16), xf, std.mem.sliceAsBytes(s16), xf, topk, &l32, &p32, &w32, &scratch); + var c32 = p32[0..topk].*; + std.mem.sort(u16, &c32, {}, std.sort.asc(u16)); + f32_agree += @intFromBool(std.mem.eql(u16, &c32, &b)); + } + std.debug.print("router layer {d}: {d} rows, picks equal to float64 on {d}/{d} (f32 host reference equal on {d}/{d}); max |device logit - f64| {e:.3}; smallest f64 10th/11th gap {e:.3}; weights max |bf16 - f64| {e:.3}, shared {e:.3}\n", .{ l, rows, rows - diff_rows, rows, f32_agree, rows, logit_err, min_gap, worst_w, worst_sw }); + worst_rows_diff = @max(worst_rows_diff, diff_rows); + try check.expect(diff_rows == 0 or min_gap <= 4 * logit_err, "router layer {d}: {d} rows pick differently with every gap above the logit error", .{ l, diff_rows }); + } + try check.expect(worst_rows_diff == 0, "router: picks differ from float64 on {d} rows (see the gaps above)", .{worst_rows_diff}); + check.pass("router: the exactly widened F16 router's top-10 equals the float64 oracle on every row of layers {any}", .{layers}); +} diff --git a/zig/tests/cuda/flashnext_exl3/model/flashnext_exl3_proj_tests.zig b/zig/tests/cuda/flashnext_exl3/model/flashnext_exl3_proj_tests.zig new file mode 100644 index 0000000000..0ed5e8c8d8 --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/model/flashnext_exl3_proj_tests.zig @@ -0,0 +1,269 @@ +//! flashnext-exl3: the projections (gate 4) and the PLE key/value projections. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const flash = @import("flashnext"); +const check = @import("../../check.zig"); +const fwd = @import("exl3_fwd_cpu.zig"); +const hc = @import("exl3_hc.zig"); +const moe = @import("exl3_moe.zig"); +const router_cpu = @import("exl3_router.zig"); +const route_cpu = @import("exl3_route.zig"); +const scan = @import("../pack/exl3_scan.zig"); +const place = @import("../pack/exl3_place.zig"); +const gdn_cpu = @import("exl3_gdn_cpu.zig"); +const attn_cpu = @import("exl3_attn_cpu.zig"); + +const X = flash.exl3; +const wt = X.prompt_wt; +const pl = X.prompt_load; +const batch = X.batch; +const Gpu = check.Gpu; +const st = core.safetensors; + +const dims = wt.dims; +const wide = wt.wide; +const streams = wt.streams; +const slots = wt.slots; +const topk = slots - 1; +const experts_n = wt.experts_n; +const shared = @import("flashnext_exl3_rig.zig"); +const typed = shared.typed; +const valueOf = shared.valueOf; +const f64Of = shared.f64Of; +const f32Of = shared.f32Of; +const f16Of = shared.f16Of; +const qmat = shared.qmat; +const bf16Bits = shared.bf16Bits; +const bfValue = shared.bfValue; +const Err = shared.Err; +const errOf = shared.errOf; +const download = shared.download; +const uploadTo = shared.uploadTo; +const RawPack = shared.RawPack; +const Pass = shared.Pass; +const Rig = shared.Rig; +const bf16Rows = shared.bf16Rows; +const default_ids = shared.default_ids; + +// -- projections (gate 4) -------------------------------------------------------------------------------------------- + +/// One Projection through Pass.lane(): `sk` 1 is bf16 [m, n]; `sk` > 1 fp32 [m, n] in the first slice, others zeroed. +fn laneOut(gpu: Gpu, p: *Pass, w: wt.Q4, x: u64, m: usize, sk: usize) !struct { bf: ?[]u16 = null, f: ?[]f32 = null, rest_zero: bool = true } { + const s = p.s; + if (sk == 1) { + _ = try p.lane(x, s.dummy, w, s.proj, s.dummy, m, 1); + return .{ .bf = try download(u16, gpu, p.stream, s.proj, m * w.n) }; + } + // the slices in a buffer of their own: a 12,288-wide part's exceeds the pass's out_sk x dims `part` + var slices = try cuda.DeviceBuffer.alloc(gpu.d, sk * m * w.n * 4); + defer slices.free(); + try slices.fill8(0xff, p.stream.handle); + _ = try p.lane(x, s.dummy, w, s.dummy, slices.ptr, m, sk); + const all = try download(f32, gpu, p.stream, slices.ptr, sk * m * w.n); + defer gpu.gpa.free(all); + var zero = true; + for (all[m * w.n ..]) |v| zero = zero and v == 0; + return .{ .f = try gpu.gpa.dupe(f32, all[0 .. m * w.n]), .rest_zero = zero }; +} + +/// The f64 product of `rows` bf16 rows through the EXL3 matrix `q`. +fn exactRows(a: std.mem.Allocator, q: *const fwd.QMat, xb: []const u16, rows: usize) ![]f64 { + const y = try a.alloc(f64, rows * q.n); + const x = try a.alloc(f64, q.k); + defer a.free(x); + const tmp = try a.alloc(f64, q.k); + defer a.free(tmp); + for (0..rows) |r| { + for (x, 0..) |*v, i| v.* = bfValue(xb[r * q.k + i]); + q.apply(x, y[r * q.n ..][0..q.n], tmp); + } + return y; +} + +/// The f64 product of `rows` bf16 rows through a dense row-major [n, k] weight. +fn denseRows(a: std.mem.Allocator, w: []const f64, n: usize, k: usize, xb: []const u16, rows: usize) ![]f64 { + const y = try a.alloc(f64, rows * n); + for (0..rows) |r| for (0..n) |j| { + var acc: f64 = 0; + for (0..k) |i| acc += bfValue(xb[r * k + i]) * w[j * k + i]; + y[r * n + j] = acc; + }; + return y; +} + +/// The fp16 rotated-input contract: x * suh after the Hadamard is rounded to fp16 (2^-11 of each term), so a column's +const proj_f32_tol = 4e-4; +const proj_bf16_tol = 4e-3; // bf16 output rounding: half an ulp is 2^-9 of a value, at most 1.95e-3 of the peak + +pub fn projections(gpu: Gpu, r: *Rig, raw: *const RawPack, lw: *const pl.Layer, l: usize, a: std.mem.Allocator, rng: std.Random) !void { + var p = r.pass((lw)[0..1]); + const linear = wt.linearLayer(l); + const kind = if (linear) "linear_attn" else "self_attn"; + const e = lw.proj.exl3 orelse return error.NotExl3; + try check.expect(lw.out.exl3 != null and lw.proj.native == null and lw.out.native == null, "layer {d}: projections did not load as EXL3", .{l}); + var worst32: f64 = 0; + var worst16: f64 = 0; + for ([_]usize{ 1, 5, 8 }) |m| { + const xb = try bf16Rows(a, rng, m * dims, 1); + var xd = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(xb)); + defer xd.free(); + // the stitched input projection (bf16, as the pass reads it) and each of its parts alone (fp32) + const whole = try laneOut(gpu, &p, lw.proj, xd.ptr, m, 1); + for (e.parts) |pt| { + const want = switch (pt.kind) { + .exl3 => blk: { + var names_buf: [128]u8 = undefined; + const name = partName(&names_buf, l, kind, pt.col, linear); + var q = try qmat(a, raw, name, 48); + defer q.deinit(a); + const y = try exactRows(a, &q, xb, m); + // the part alone in fp32: Projection.exl3 over one part + var one = e.*; + var single = [1]X.load.Part{pt}; + single[0].col = 0; + one.parts = &single; + one.n = pt.n; + var view = wt.Q4.empty(gpu.d); + view.n = pt.n; + view.k = dims; + view.exl3 = &one; + const got = try laneOut(gpu, &p, view, xd.ptr, m, 2); + const ef = errOf(got.f.?, y); + try check.expect(got.rest_zero, "layer {d}: {s} left the other K slice unzeroed", .{ l, name }); + std.debug.print(" L{d} {s:<50} rows={d} fp32 out: max|err| {e:.3} peak {d:.3} rel {e:.3}\n", .{ l, name, m, ef.max_abs, ef.peak, ef.rel }); + worst32 = @max(worst32, ef.rel); + break :blk y; + }, + .native => blk: { + // in_proj_b then in_proj_a: the F16 rows the Store widened, the oracle reads the stored F16 + const b = try f64Of(a, try raw.get("model.language_model.layers.{d}.linear_attn.in_proj_b.weight", .{l})); + const av = try f64Of(a, try raw.get("model.language_model.layers.{d}.linear_attn.in_proj_a.weight", .{l})); + const w = try std.mem.concat(a, f64, &.{ b, av }); + break :blk try denseRows(a, w, pt.n, dims, xb, m); + }, + }; + // the part's columns of the stitched bf16 row + const got = try a.alloc(u16, m * pt.n); + for (0..m) |row| @memcpy(got[row * pt.n ..][0..pt.n], whole.bf.?[row * e.n + pt.col ..][0..pt.n]); + const eb = errOf(got, want); + std.debug.print(" L{d} stitched in-proj columns {d}..{d} ({s}) rows={d} bf16 out: rel {e:.3}\n", .{ l, pt.col, pt.col + pt.n, @tagName(pt.kind), m, eb.rel }); + worst16 = @max(worst16, eb.rel); + } + // the output projection as the pass runs it: fp32 into the first of out_sk K slices + const ok = lw.out.k; + const ob = try bf16Rows(a, rng, m * ok, 1); + var od = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(ob)); + defer od.free(); + var nb: [128]u8 = undefined; + const oname = try std.fmt.bufPrint(&nb, "model.language_model.layers.{d}.{s}.{s}", .{ l, kind, if (linear) "out_proj" else "o_proj" }); + var q = try qmat(a, raw, oname, 48); + defer q.deinit(a); + const want = try exactRows(a, &q, ob, m); + const got = try laneOut(gpu, &p, lw.out, od.ptr, m, wt.out_sk); + const eo = errOf(got.f.?, want); + try check.expect(got.rest_zero, "layer {d}: the output projection left its other K slices unzeroed", .{l}); + std.debug.print(" L{d} {s:<50} rows={d} fp32 out (slice 0 of {d}): max|err| {e:.3} peak {d:.3} rel {e:.3}\n", .{ l, oname, m, wt.out_sk, eo.max_abs, eo.peak, eo.rel }); + worst32 = @max(worst32, eo.rel); + } + try check.expect(worst32 <= proj_f32_tol, "layer {d}: an EXL3 projection's fp32 output is {e:.3} of peak from float64 (> {e})", .{ l, worst32, proj_f32_tol }); + try check.expect(worst16 <= proj_bf16_tol, "layer {d}: the stitched bf16 input projection is {e:.3} of peak from float64 (> {e})", .{ l, worst16, proj_bf16_tol }); + check.pass("layer {d} projections through Pass.lane(): fp32 out worst {e:.3} of peak, stitched bf16 in-proj worst {e:.3}, rows 1/5/8", .{ l, worst32, worst16 }); +} + +/// The EXL3 matrix behind a stitched part, from its first column. +fn partName(buf: []u8, l: usize, kind: []const u8, col: usize, linear: bool) []const u8 { + const name: []const u8 = if (linear) switch (col) { + 0 => "in_proj_qkv", + else => "in_proj_z", + } else switch (col) { + 0 => "q_proj", + 12288 => "k_proj", + 12800 => "v_proj", + else => "indexer.index_qk_proj", + }; + return std.fmt.bufPrint(buf, "model.language_model.layers.{d}.{s}.{s}", .{ l, kind, name }) catch unreachable; +} + +/// The PLE key/value projections of an EXL3 pack take the n-gram rows as fp16 (cuda_exl3_ple.zig's default): through +pub fn pleMode(gpu: Gpu, dir: []const u8) !void { + var rig: Rig = undefined; + try rig.init(gpu, dir, 8, 64); + defer rig.deinit(); + var raw = try RawPack.open(gpu, dir); + defer raw.close(gpu); + var arena: std.heap.ArenaAllocator = .init(gpu.gpa); + defer arena.deinit(); + const a = arena.allocator(); + const d: *cuda.Driver = @constCast(gpu.d); + // an EXL3 pack's store holds the shared Ops once its first EXL3 matrix loads; this check loads none + (try X.load.opsOf(&rig.store, gpu.d)).release(); + var key = try X.projection_load.load(a, d, &rig.store, "language_model.model.layers.1.ple.key_proj", wide, dims); + defer key.free(); + var value = try X.projection_load.load(a, d, &rig.store, "language_model.model.layers.1.ple.value_proj", dims, dims); + defer value.free(); + try X.load.pleRows(a, d, &rig.store, &key, &value); + try check.expect(key.exl3 != null and key.exl3.?.x == .f16 and value.exl3.?.x == .f16, "ple: key/value did not take fp16 rows", .{}); + var p = rig.pass(&.{}); + const kw = try f64Of(a, try raw.get("model.language_model.layers.1.ple.key_proj.weight", .{})); + const vw = try f64Of(a, try raw.get("model.language_model.layers.1.ple.value_proj.weight", .{})); + var prng = std.Random.DefaultPrng.init(0x91e_f16); + const rng = prng.random(); + var worst_key: f64 = 0; + var worst_value: f64 = 0; + var worst_bf16: f64 = 0; + for ([_]usize{ 1, 5, 8 }) |m| { + const xh = try a.alloc(f16, m * dims); + const xv = try a.alloc(f64, m * dims); + const xr = try a.alloc(u16, m * dims); // the same rows rounded to bf16 + for (xh, xv, xr) |*h, *v, *b| { + h.* = @floatCast(rng.floatNorm(f32) * 0.5); + v.* = @floatCast(h.*); + b.* = bf16Bits(@floatCast(h.*)); + } + var xd = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(xh)); + defer xd.free(); + const want_k = try a.alloc(f64, m * wide); + const want_v = try a.alloc(f64, m * dims); + for (0..m) |r| { + for (0..wide) |j| { + var acc: f64 = 0; + for (0..dims) |i| acc += xv[r * dims + i] * kw[j * dims + i]; + want_k[r * wide + j] = acc; + } + for (0..dims) |j| { + var acc: f64 = 0; + for (0..dims) |i| acc += xv[r * dims + i] * vw[j * dims + i]; + want_v[r * dims + j] = acc; + } + } + const got_k = try laneOut(gpu, &p, key, xd.ptr, m, 1); + const ek = errOf(got_k.bf.?, want_k); + // the value projection as pleBlock reads it: the sum of its 4 fp32 K slices + var slices = try cuda.DeviceBuffer.alloc(gpu.d, 4 * m * dims * 4); + defer slices.free(); + try slices.fill8(0xff, p.stream.handle); + _ = try p.lane(xd.ptr, p.s.dummy, value, p.s.dummy, slices.ptr, m, 4); + const all = try download(f32, gpu, p.stream, slices.ptr, 4 * m * dims); + defer gpu.gpa.free(all); + const sum = try a.alloc(f64, m * dims); + for (sum, 0..) |*o, i| o.* = @as(f64, all[i]) + all[m * dims + i] + all[2 * m * dims + i] + all[3 * m * dims + i]; + const ev = errOf(sum, want_v); + // the bf16-row alternative, on the host in float64: rows rounded to bf16, weights exact + const alt = try a.alloc(f64, m * dims); + for (0..m) |r| for (0..dims) |j| { + var acc: f64 = 0; + for (0..dims) |i| acc += bfValue(xr[r * dims + i]) * vw[j * dims + i]; + alt[r * dims + j] = acc; + }; + const eb = errOf(alt, want_v); + std.debug.print(" ple rows={d}: key_proj bf16 out rel {e:.3}; value_proj fp32 slices rel {e:.3}; (bf16-rounded rows would give {e:.3})\n", .{ m, ek.rel, ev.rel, eb.rel }); + worst_key = @max(worst_key, ek.rel); + worst_value = @max(worst_value, ev.rel); + worst_bf16 = @max(worst_bf16, eb.rel); + } + // key: bf16 output rounding (2^-9 of a value); value: fp32 sums of exact products + try check.expect(worst_key <= 4e-3, "ple: key_proj from fp16 rows is {e:.3} of peak from float64", .{worst_key}); + try check.expect(worst_value <= 1e-5, "ple: value_proj from fp16 rows is {e:.3} of peak from float64", .{worst_value}); + check.pass("ple: fp16 n-gram rows through key/value (exact bf16 split): key {e:.3}, value {e:.3} of peak (bf16-rounded rows: {e:.3})", .{ worst_key, worst_value, worst_bf16 }); +} diff --git a/zig/tests/cuda/flashnext_exl3/model/flashnext_exl3_rig.zig b/zig/tests/cuda/flashnext_exl3/model/flashnext_exl3_rig.zig new file mode 100644 index 0000000000..55f85d788f --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/model/flashnext_exl3_rig.zig @@ -0,0 +1,251 @@ +//! flashnext-exl3: the raw-pack reader, f64 conversions, error metrics and the one-layer pass rig the gates share. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const flash = @import("flashnext"); +const check = @import("../../check.zig"); +const fwd = @import("exl3_fwd_cpu.zig"); +const hc = @import("exl3_hc.zig"); +const moe = @import("exl3_moe.zig"); +const router_cpu = @import("exl3_router.zig"); +const route_cpu = @import("exl3_route.zig"); +const scan = @import("../pack/exl3_scan.zig"); +const place = @import("../pack/exl3_place.zig"); +const gdn_cpu = @import("exl3_gdn_cpu.zig"); +const attn_cpu = @import("exl3_attn_cpu.zig"); + +const X = flash.exl3; +const wt = X.prompt_wt; +const pl = X.prompt_load; +const batch = X.batch; +const Gpu = check.Gpu; +const st = core.safetensors; + +const dims = wt.dims; +const wide = wt.wide; +const streams = wt.streams; +const slots = wt.slots; +const topk = slots - 1; +const experts_n = wt.experts_n; + +/// "The capital of France is Paris. It is known for the Eiffel Tower." (exl3_fwd_tests.zig's default ids) +pub const default_ids = [_]u32{ 760, 6511, 314, 9338, 369, 11751, 13, 1049, 369, 3750, 364, 279, 242476, 300, 21262, 13 }; + +// -- the raw pack (the oracle's view: tensors exactly as stored, no widening) ------------------------------------- + +pub const RawPack = struct { + files: std.ArrayList(st.File) = .empty, + + pub fn open(gpu: Gpu, dir: []const u8) !RawPack { + var p: RawPack = .{}; + var buf: [512]u8 = undefined; + for (0..7) |i| { + const path = try std.fmt.bufPrint(&buf, "{s}/model-{d:0>5}-of-00007.safetensors", .{ dir, i + 1 }); + try p.files.append(gpu.gpa, try st.File.open(gpu.gpa, gpu.io, path)); + } + return p; + } + + pub fn close(p: *RawPack, gpu: Gpu) void { + for (p.files.items) |*f| f.close(gpu.io); + p.files.deinit(gpu.gpa); + } + + pub fn get(p: *const RawPack, comptime fmt: []const u8, args: anytype) !st.Tensor { + var buf: [256]u8 = undefined; + const name = try std.fmt.bufPrint(&buf, fmt, args); + for (p.files.items) |*f| if (f.get(name)) |t| return t; + std.debug.print("raw pack: no {s}\n", .{name}); + return error.MissingTensor; + } +}; + +pub fn typed(comptime T: type, a: std.mem.Allocator, bytes: []const u8) ![]T { + const out = try a.alloc(T, bytes.len / @sizeOf(T)); + @memcpy(std.mem.sliceAsBytes(out), bytes[0 .. out.len * @sizeOf(T)]); + return out; +} + +pub fn valueOf(t: st.Tensor, i: usize) f64 { + const h = std.mem.readInt(u16, t.bytes[2 * i ..][0..2], .little); + return switch (t.dtype) { + .f16 => @floatCast(@as(f16, @bitCast(h))), + .bf16 => @as(f32, @bitCast(@as(u32, h) << 16)), + else => unreachable, + }; +} + +pub fn f64Of(a: std.mem.Allocator, t: st.Tensor) ![]f64 { + if (t.dtype != .f16 and t.dtype != .bf16) return error.UnexpectedTensor; + const out = try a.alloc(f64, t.bytes.len / 2); + for (out, 0..) |*o, i| o.* = valueOf(t, i); + return out; +} + +pub fn f32Of(a: std.mem.Allocator, t: st.Tensor) ![]f32 { + if (t.dtype != .f16 and t.dtype != .bf16) return error.UnexpectedTensor; + const out = try a.alloc(f32, t.bytes.len / 2); + for (out, 0..) |*o, i| o.* = @floatCast(valueOf(t, i)); + return out; +} + +pub fn f16Of(a: std.mem.Allocator, t: st.Tensor) ![]f16 { + if (t.dtype != .f16) return error.UnexpectedTensor; + return typed(f16, a, t.bytes); +} + +/// The EXL3 matrix `prefix` as the float64 oracle's QMat (W = diag(suh) H Wq H diag(svh), exl3_format's decoder). +pub fn qmat(a: std.mem.Allocator, raw: *const RawPack, prefix: []const u8, threads: usize) !fwd.QMat { + const t = try raw.get("{s}.trellis", .{prefix}); + const trellis = try typed(i16, a, t.bytes); + defer a.free(trellis); + const suh = try typed(f16, a, (try raw.get("{s}.suh", .{prefix})).bytes); + defer a.free(suh); + const svh = try typed(f16, a, (try raw.get("{s}.svh", .{prefix})).bytes); + defer a.free(svh); + return fwd.QMat.init(a, @intCast(t.dim(2) / 8), trellis, t.dim(0), t.dim(1), suh, svh, threads); +} + +pub fn bf16Bits(v: f32) u16 { + const b: u32 = @bitCast(v); + return @intCast((b + 0x7fff + ((b >> 16) & 1)) >> 16); +} + +pub fn bfValue(bits: u16) f64 { + return @as(f32, @bitCast(@as(u32, bits) << 16)); +} + +pub const Err = struct { max_abs: f64 = 0, peak: f64 = 0, rel: f64 = 0 }; + +pub fn errOf(got: anytype, want: []const f64) Err { + var e: Err = .{}; + for (want, 0..) |w, i| { + const g: f64 = switch (@TypeOf(got)) { + []const u16, []u16 => bfValue(got[i]), + else => @floatCast(got[i]), + }; + e.max_abs = @max(e.max_abs, @abs(g - w)); + e.peak = @max(e.peak, @abs(w)); + } + e.rel = if (e.peak > 0) e.max_abs / e.peak else 0; + return e; +} + +pub fn download(comptime T: type, gpu: Gpu, s: cuda.Stream, ptr: u64, n: usize) ![]T { + try s.synchronize(); + const out = try gpu.gpa.alloc(T, n); + const b: cuda.DeviceBuffer = .{ .d = gpu.d, .ptr = ptr, .len = n * @sizeOf(T) }; + try b.download(0, std.mem.sliceAsBytes(out)); + return out; +} + +/// A pageable copy into the pass's scratch: the legacy-stream copy is ordered against nothing, so drain both sides. +pub fn uploadTo(gpu: Gpu, s: cuda.Stream, ptr: u64, bytes: []const u8) !void { + try s.synchronize(); + const b: cuda.DeviceBuffer = .{ .d = gpu.d, .ptr = ptr, .len = bytes.len }; + try b.upload(0, bytes); + try gpu.ctx.synchronize(); +} + +// -- the pass rig ----------------------------------------------------------------------------------------------------- + +pub const Pass = batch.Pass(flash.Tri); + +pub const Rig = struct { + gpu: Gpu, + stream: cuda.Stream, + glue: flash.Glue, + k: batch.Kernels, + s: batch.Scratch, + store: wt.Store, + pack_mod: cuda.Module, + pack_fn: cuda.Function, + bufs: std.ArrayList(cuda.DeviceBuffer) = .empty, + conv: [wt.linear_n]cuda.DeviceBuffer = undefined, + state: [wt.linear_n]cuda.DeviceBuffer = undefined, + kc: [wt.attn_n]cuda.DeviceBuffer = undefined, + vc: [wt.attn_n]cuda.DeviceBuffer = undefined, + ikc: [wt.attn_n]cuda.DeviceBuffer = undefined, + pooled: [wt.attn_n]cuda.DeviceBuffer = undefined, + tail: cuda.DeviceBuffer = undefined, + history: [2]i64 = .{ 0, 0 }, + scores: cuda.DeviceBuffer = undefined, + cap: usize, + + pub fn init(r: *Rig, gpu: Gpu, dir: []const u8, rows: usize, cap: usize) !void { + r.* = .{ .gpu = gpu, .stream = try cuda.Stream.init(gpu.d, true), .glue = undefined, .k = undefined, .s = undefined, .store = undefined, .pack_mod = undefined, .pack_fn = undefined, .cap = cap }; + r.glue = try flash.Glue.load(gpu.d); + r.k = try batch.Kernels.load(gpu.d); + r.s = try batch.Scratch.alloc(gpu.gpa, gpu.d, r.stream, rows); + r.store = try wt.Store.open(gpu.gpa, gpu.io, dir); + try r.store.detectNorms(); + r.pack_mod = try cuda.Module.load(gpu.d, cuda.kernels.experts_pack); + r.pack_fn = try r.pack_mod.function(wt.pack_symbol); + const gdn = flash.gdn; + const z: cuda.DeviceBuffer = .{ .d = gpu.d, .ptr = 0, .len = 0 }; + r.conv = @splat(z); + r.state = @splat(z); + r.kc = @splat(z); + r.vc = @splat(z); + r.ikc = @splat(z); + r.pooled = @splat(z); + // every layer's state, as Engine.newSeq makes them (a check may step any layer) + for (&r.conv, &r.state) |*c, *x| { + c.* = try r.zeroed((gdn.taps - 1) * gdn.conv_dim * 2); + x.* = try r.zeroed(wt.state_n * 4); + } + for (&r.kc, &r.vc, &r.ikc, &r.pooled) |*kc, *vc, *ikc, *pooled| { + kc.* = try r.zeroed(cap * wt.kv_heads * wt.head_dim * 2); + vc.* = try r.zeroed(cap * wt.kv_heads * wt.head_dim * 2); + ikc.* = try r.zeroed(cap * wt.index_dim * 2); + pooled.* = try r.zeroed(flash.prompt.pooledBytes(cap)); + } + r.tail = try r.zeroed(wt.ple_tail_n * wide * 2); + r.scores = try r.zeroed(batch.scoreBytes(cap, rows)); + } + + fn zeroed(r: *Rig, bytes: usize) !cuda.DeviceBuffer { + const b = try cuda.DeviceBuffer.alloc(r.gpu.d, bytes); + try r.bufs.append(r.gpu.gpa, b); + try b.fill8(0, r.stream.handle); + try r.stream.synchronize(); + return b; + } + + pub fn reset(r: *Rig) !void { + for (r.bufs.items) |b| try b.fill8(0, r.stream.handle); + try r.stream.synchronize(); + } + + pub fn deinit(r: *Rig) void { + r.stream.synchronize() catch {}; + for (r.bufs.items) |*b| b.free(); + r.bufs.deinit(r.gpu.gpa); + r.pack_mod.unload(); + r.store.deinit(); + r.s.free(r.gpu.gpa); + r.k.deinit(); + r.glue.deinit(); + r.stream.deinit(); + } + + pub fn pass(r: *Rig, layers: []const pl.Layer) Pass { + return .{ + .gpa = r.gpu.gpa, + .io = r.gpu.io, + .stream = r.stream, + .tri = .{ .own = &r.glue, .s = r.stream }, + .k = &r.k, + .s = &r.s, + .m = .{ .table_w = 0, .table_s = 0, .table_b = 0, .layers = layers, .mixer = undefined, .head = undefined, .ngram = undefined, .collective = null }, + .seq = .{ .conv = &r.conv, .state = &r.state, .kc = &r.kc, .vc = &r.vc, .ikc = &r.ikc, .pooled = &r.pooled, .tail = r.tail, .history = &r.history, .cap = r.cap, .scores = r.scores.ptr }, + .hooks = .{}, + }; + } +}; + +pub fn bf16Rows(a: std.mem.Allocator, rng: std.Random, n: usize, sigma: f32) ![]u16 { + const v = try a.alloc(u16, n); + for (v) |*x| x.* = bf16Bits(rng.floatNorm(f32) * sigma); + return v; +} diff --git a/zig/tests/cuda/flashnext_exl3/model/flashnext_exl3_tests.zig b/zig/tests/cuda/flashnext_exl3/model/flashnext_exl3_tests.zig new file mode 100644 index 0000000000..1713d06354 --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/model/flashnext_exl3_tests.zig @@ -0,0 +1,50 @@ +//! flashnext-exl3: the Flash Next CUDA family on a real EXL3 pack against our float64 oracle; this hub parses the mode. +const std = @import("std"); +const check = @import("../../check.zig"); +const shared = @import("flashnext_exl3_rig.zig"); +const names_tests = @import("flashnext_exl3_names_tests.zig"); +const proj_tests = @import("flashnext_exl3_proj_tests.zig"); +const layer_tests = @import("flashnext_exl3_layer_tests.zig"); + +const Gpu = check.Gpu; + +pub const default_ids = shared.default_ids; +pub const RawPack = shared.RawPack; +pub const bfValue = shared.bfValue; +pub const Err = shared.Err; +pub const errOf = shared.errOf; +pub const download = shared.download; +pub const uploadTo = shared.uploadTo; +pub const Pass = shared.Pass; +pub const Rig = shared.Rig; +pub const Taps = layer_tests.Taps; +pub const deviceLayer = layer_tests.deviceLayer; +pub const layerCheck = layer_tests.layerCheck; + +fn parseList(gpa: std.mem.Allocator, text: []const u8) ![]usize { + var list: std.ArrayList(usize) = .empty; + var it = std.mem.splitScalar(u8, text, ','); + while (it.next()) |tok| if (tok.len > 0) try list.append(gpa, try std.fmt.parseInt(usize, tok, 10)); + return list.toOwnedSlice(gpa); +} + +pub fn run(gpu: Gpu, rest: []const [:0]const u8) !void { + if (rest.len < 2) return error.MissingArgument; + const dir = rest[0]; + const mode = rest[1]; + if (std.mem.eql(u8, mode, "names")) return names_tests.names(gpu, dir); + if (std.mem.eql(u8, mode, "norms")) return names_tests.norms(gpu, dir); + if (std.mem.eql(u8, mode, "ple")) return proj_tests.pleMode(gpu, dir); + if (std.mem.eql(u8, mode, "router")) { + const which = try parseList(gpu.gpa, if (rest.len > 2) rest[2] else "0,3,17"); + defer gpu.gpa.free(which); + const rows = if (rest.len > 3) try std.fmt.parseInt(usize, rest[3], 10) else 4096; + return names_tests.router(gpu, dir, which, rows); + } + if (std.mem.eql(u8, mode, "layer")) { + const which = try parseList(gpu.gpa, if (rest.len > 2) rest[2] else "0,3"); + defer gpu.gpa.free(which); + return layer_tests.layerMode(gpu, dir, which); + } + return error.UnknownMode; +} diff --git a/zig/tests/cuda/flashnext_exl3/offload/exl3_pool_bench_tests.zig b/zig/tests/cuda/flashnext_exl3/offload/exl3_pool_bench_tests.zig new file mode 100644 index 0000000000..a23f3fcb55 --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/offload/exl3_pool_bench_tests.zig @@ -0,0 +1,255 @@ +//! exl3-pool: gate 4, miss cost, copy/compute overlap and the decode pipeline. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const flash = @import("flashnext"); +const check = @import("../../check.zig"); + +const st = core.safetensors; +const exl3e = cuda.exl3_experts; +const fpool = flash.exl3_pool; +const policy = fpool.offload; +const Gpu = check.Gpu; + +const data = @import("exl3_pool_data.zig"); +const dims = data.dims; +const width = data.width; +const routed_n = data.routed_n; +const table_n = data.table_n; +const top_k = data.top_k; +const slots = data.slots; +const max_rows = data.max_rows; +const shard_count = data.shard_count; +const projs = data.projs; +const Pack = data.Pack; +const LayerData = data.LayerData; +const Zipf = data.Zipf; +const bf16Bits = data.bf16Bits; +const Io = data.Io; +const call = data.call; +const poolRound = data.poolRound; +const genRound = data.genRound; + +// ---------------------------------------------------------------------------------------------------- gate 4 + +fn ms(t0: i96, t1: i96) f64 { + return @as(f64, @floatFromInt(t1 - t0)) / 1e6; +} + +fn median(xs: []f64) f64 { + std.mem.sort(f64, xs, {}, std.sort.asc(f64)); + return xs[xs.len / 2]; +} + +/// A row whose first `k` routed picks are not held by the pool and the rest are. +fn rowWithMisses(pool: *const fpool.Pool, rng: std.Random, k: usize, out: []i32) void { + var n: usize = 0; + // held experts first (up to 10 - k) + for (pool.slots_tab, pool.line_tab) |id, state| { + if (n == top_k - k) break; + if (state == .resident) { + out[k + n] = id; + n += 1; + } + } + var m: usize = 0; + while (m < k) { + const id: u16 = rng.intRangeLessThan(u16, 0, routed_n); + if (policy.lineOf(pool.v, pool.opt.layer, id) != null) continue; + if (std.mem.indexOfScalar(i32, out[0..m], id) != null) continue; + out[m] = id; + m += 1; + } + out[top_k] = 512; +} + +fn missCost(gpu: Gpu, ex: *const exl3e.Experts, s: cuda.Stream, ld: *const LayerData, src: fpool.Source, direct: bool, label: []const u8) !void { + const gpa = gpu.gpa; + var pool = try fpool.Pool.init(gpa, gpu.io, gpu.d, src, ld.scales(), .{ .layer = @intCast(ld.layer), .lines = 64, .direct = direct, .double_buffer = false }, null); + defer pool.deinit(); + var vp = try Io.init(gpu, pool.tables()); + defer vp.free(); + var prng = std.Random.DefaultPrng.init(0xbe7c4); + const rng = prng.random(); + var x: [dims]u16 = undefined; + for (&x) |*b| b.* = bf16Bits((rng.float(f32) - 0.5) * 2.0); + const wts: [slots]f32 = @splat(0.1); + var picks: [slots]i32 = undefined; + // warm: fill the pool with 64 experts + for (0..7) |_| { + rowWithMisses(&pool, rng, top_k, &picks); + try vp.set(&x, &picks, &wts); + _ = try poolRound(&pool, ex, s, vp, &picks, 1, false); + try s.synchronize(); + } + var c0 = try cuda.Event.init(gpu.d, true); + defer c0.deinit(); + var c1 = try cuda.Event.init(gpu.d, true); + defer c1.deinit(); + var k0 = try cuda.Event.init(gpu.d, true); + defer k0.deinit(); + var k1 = try cuda.Event.init(gpu.d, true); + defer k1.deinit(); + std.debug.print(" [{s}] one decode row (10 routed + shared), 64 lines, rounds of k misses (median of 9; idle card):\n", .{label}); + var base: f64 = 0; + for ([_]usize{ 0, 1, 2, 4, 10 }) |k| { + var wall: [9]f64 = undefined; + var copy: [9]f64 = undefined; + var comp: [9]f64 = undefined; + var stage_host: [9]f64 = undefined; + for (0..9) |it| { + rowWithMisses(&pool, rng, k, &picks); + try vp.set(&x, &picks, &wts); + try s.synchronize(); + const t0 = check.now(gpu.io); + const r = try pool.plan(&picks, 1); + const t1 = check.now(gpu.io); + try c0.record(pool.lines.copy); + try pool.stage(r, 0, s); + try c1.record(pool.lines.copy); + const t2 = check.now(gpu.io); + try k0.record(s); + try call(ex, s, pool.tables(), vp, pool.picksFor(r, 0, vp.picks.ptr), 1, false); + try k1.record(s); + try pool.release(r, 0, s); + try s.synchronize(); + try pool.lines.copy.synchronize(); + const t3 = check.now(gpu.io); + wall[it] = ms(t0, t3); + copy[it] = try c0.elapsedMs(c1); + comp[it] = try k0.elapsedMs(k1); + stage_host[it] = ms(t1, t2); + } + const w = median(&wall); + if (k == 0) base = w; + std.debug.print(" k={d:>2}: round wall {d:.3} ms (+{d:.3} vs all-hit), copy stream {d:.3} ms ({d:.3} ms/expert), experts compute {d:.3} ms (measured after the stream's wait on the copies), host stage {d:.3} ms\n", .{ + k, w, w - base, median(©), if (k > 0) median(©) / @as(f64, @floatFromInt(k)) else 0, median(&comp), median(&stage_host), + }); + } + const ls = pool.lines.stats; + if (!direct) std.debug.print(" host staging memcpy (mmap -> pinned): {d:.3} ms/expert average over {d} copies ({d:.2} GB/s)\n", .{ @as(f64, @floatFromInt(ls.stage_ns)) / 1e6 / @as(f64, @floatFromInt(ls.copies)), ls.copies, @as(f64, @floatFromInt(ls.staged_bytes)) / @as(f64, @floatFromInt(ls.stage_ns)) }); +} + +/// A 16-row chunk through a 24-line pool (several passes): copies of pass j+1 overlap compute of pass j or not. +fn overlapChunk(gpu: Gpu, ex: *const exl3e.Experts, s: cuda.Stream, ld: *const LayerData, src: fpool.Source, serial: bool) !struct { wall: f64, copy: f64, comp: f64, passes: usize, copies: usize } { + const gpa = gpu.gpa; + var pool = try fpool.Pool.init(gpa, gpu.io, gpu.d, src, ld.scales(), .{ .layer = @intCast(ld.layer), .lines = 24, .direct = true }, null); + defer pool.deinit(); + var vp = try Io.init(gpu, pool.tables()); + defer vp.free(); + const z = Zipf.init(0x0e1a9 + ld.layer, 0.8); + var prng = std.Random.DefaultPrng.init(0x0e1a9); + const rng = prng.random(); + var x: [max_rows * dims]u16 = undefined; + var picks: [max_rows * slots]i32 = undefined; + var wts: [max_rows * slots]f32 = undefined; + genRound(&z, rng, max_rows, &x, &picks, &wts); + try vp.set(&x, &picks, &wts); + var cs: [64]cuda.Event = undefined; + for (&cs) |*e| e.* = try cuda.Event.init(gpu.d, true); + defer for (&cs) |*e| e.deinit(); + var ks: [64]cuda.Event = undefined; + for (&ks) |*e| e.* = try cuda.Event.init(gpu.d, true); + defer for (&ks) |*e| e.deinit(); + try s.synchronize(); + const t0 = check.now(gpu.io); + const r = try pool.plan(&picks, max_rows); + if (r.n_passes > 32) return error.TooManyPasses; + const l = pool.tables(); + for (0..r.n_passes) |j| { + try cs[2 * j].record(pool.lines.copy); + try pool.stage(r, j, s); + try cs[2 * j + 1].record(pool.lines.copy); + if (serial) try pool.lines.copy.synchronize(); + try ks[2 * j].record(s); + try call(ex, s, l, vp, pool.picksFor(r, j, vp.picks.ptr), max_rows, false); + try ks[2 * j + 1].record(s); + try pool.release(r, j, s); + if (serial) try s.synchronize(); + } + try s.synchronize(); + try pool.lines.copy.synchronize(); + const wall = ms(t0, check.now(gpu.io)); + var copy: f64 = 0; + var comp: f64 = 0; + for (0..r.n_passes) |j| { + copy += try cs[2 * j].elapsedMs(cs[2 * j + 1]); + comp += try ks[2 * j].elapsedMs(ks[2 * j + 1]); + } + return .{ .wall = wall, .copy = copy, .comp = comp, .passes = r.n_passes, .copies = r.n_copies }; +} + +/// 60 decode rounds (one row each, skewed picks known ahead), sequential (sync every round) vs pipelined (the host +fn decodePipe(gpu: Gpu, ex: *const exl3e.Experts, s: cuda.Stream, ld: *const LayerData, src: fpool.Source, lines: usize, pipelined: bool, double: bool) !struct { wall: f64, hit: f64, release_waits: u64, protected: u64 } { + const gpa = gpu.gpa; + var pool = try fpool.Pool.init(gpa, gpu.io, gpu.d, src, ld.scales(), .{ .layer = @intCast(ld.layer), .lines = lines, .direct = true, .double_buffer = double }, null); + defer pool.deinit(); + const rounds = 60; + var ios: [rounds]Io = undefined; + var made: usize = 0; + defer for (ios[0..made]) |*v| v.free(); + const z = Zipf.init(0xdec0 + ld.layer, 1.0); + var prng = std.Random.DefaultPrng.init(0xdec0); + const rng = prng.random(); + var picks: [rounds][slots]i32 = undefined; + var x: [dims]u16 = undefined; + var wts: [slots]f32 = undefined; + for (0..rounds) |i| { + ios[i] = try Io.init(gpu, pool.tables()); + made += 1; + genRound(&z, rng, 1, &x, &picks[i], &wts); + try ios[i].set(&x, &picks[i], &wts); + } + try s.synchronize(); + const t0 = check.now(gpu.io); + for (0..rounds) |i| { + _ = try poolRound(&pool, ex, s, ios[i], &picks[i], 1, false); + if (!pipelined) try s.synchronize(); + } + try s.synchronize(); + const wall = ms(t0, check.now(gpu.io)); + const ps = pool.stats; + return .{ .wall = wall, .hit = @as(f64, @floatFromInt(ps.hit_slots)) / @as(f64, @floatFromInt(ps.hit_slots + ps.miss_slots)), .release_waits = pool.lines.stats.release_waits, .protected = ps.protected }; +} + +pub fn benchRun(gpu: Gpu, ex: *const exl3e.Experts, s: cuda.Stream, ld: *const LayerData) !void { + const gpa = gpu.gpa; + // (a) miss cost, staged from the mapped checkpoint + try missCost(gpu, ex, s, ld, ld.source(), false, "staged: mmap -> pinned staging -> H2D"); + // the same experts in one pinned bank (what keeping the routed trellis in pinned host RAM gives): direct H2D + var total: usize = 0; + for (ld.routed) |e| for (e) |b| { + total += b.len; + }; + var bank = try cuda.HostBuffer.alloc(gpu.d, total); + defer bank.free(); + const pinned = try gpa.alloc([3][]const u8, routed_n); + defer gpa.free(pinned); + var at: usize = 0; + for (ld.routed, pinned) |e, *o| for (e, 0..) |b, p| { + @memcpy(bank.bytes[at..][0..b.len], b); + o[p] = bank.bytes[at..][0..b.len]; + at += b.len; + }; + std.debug.print(" pinned bank for layer {d}: {d:.0} MiB of routed trellis copied into page-locked host memory\n", .{ ld.layer, @as(f64, @floatFromInt(total)) / 1048576.0 }); + const psrc: fpool.Source = .{ .routed = pinned, .shared = ld.shared }; + try missCost(gpu, ex, s, ld, psrc, true, "pinned bank: direct H2D"); + + // (b) a 16-row chunk split into passes: serial vs overlapped copies + const ser = try overlapChunk(gpu, ex, s, ld, psrc, true); + const ovl = try overlapChunk(gpu, ex, s, ld, psrc, false); + const hidden_ms = (ser.wall - ovl.wall); + std.debug.print(" [16-row chunk, 24 lines, pinned bank] {d} passes, {d} copies: serial {d:.3} ms (copy {d:.3} + compute {d:.3}), overlapped {d:.3} ms; overlap hides {d:.3} ms = {d:.0}% of the smaller of copy/compute\n", .{ + ser.passes, ser.copies, ser.wall, ser.copy, ser.comp, ovl.wall, hidden_ms, 100.0 * hidden_ms / @min(ser.copy, ser.comp), + }); + // (c) decode rounds: pipelined staging, double buffering on/off + for ([_]usize{ 24, 48 }) |lines| { + const seq = try decodePipe(gpu, ex, s, ld, psrc, lines, false, true); + const pipe1 = try decodePipe(gpu, ex, s, ld, psrc, lines, true, false); + const pipe2 = try decodePipe(gpu, ex, s, ld, psrc, lines, true, true); + std.debug.print(" [60 decode rounds, {d} lines, pinned bank, hit {d:.3}] sequential {d:.2} ms; pipelined {d:.2} ms single-buffered ({d} copies waited for the previous round's compute), {d:.2} ms double-buffered ({d} waited, {d} rounds spared the previous round's lines)\n", .{ + lines, seq.hit, seq.wall, pipe1.wall, pipe1.release_waits, pipe2.wall, pipe2.release_waits, pipe2.protected, + }); + } + check.pass("layer {d}: miss cost, chunk overlap and decode pipelining measured (see lines above)", .{ld.layer}); +} diff --git a/zig/tests/cuda/flashnext_exl3/offload/exl3_pool_data.zig b/zig/tests/cuda/flashnext_exl3/offload/exl3_pool_data.zig new file mode 100644 index 0000000000..783bbc8a76 --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/offload/exl3_pool_data.zig @@ -0,0 +1,284 @@ +//! exl3-pool: the pack reader, one layer's experts and the synthetic routing the pool gates share. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const flash = @import("flashnext"); +const check = @import("../../check.zig"); + +const st = core.safetensors; +const exl3e = cuda.exl3_experts; +const fpool = flash.exl3_pool; +const policy = fpool.offload; +const Gpu = check.Gpu; + +pub const dims = 2560; +pub const width = 640; +pub const routed_n = 512; +pub const table_n = 513; +pub const top_k = 10; +pub const slots = 11; +pub const max_rows = 16; +pub const shard_count = 7; +pub const projs = [3][]const u8{ "gate_proj", "up_proj", "down_proj" }; + +pub const Pack = struct { + files: [shard_count]st.File = undefined, + open_n: usize = 0, + + pub fn open(p: *Pack, gpu: Gpu, dir: []const u8) !void { + var buf: [512]u8 = undefined; + for (0..shard_count) |i| { + const path = try std.fmt.bufPrint(&buf, "{s}/model-{d:0>5}-of-{d:0>5}.safetensors", .{ dir, i + 1, shard_count }); + p.files[i] = try st.File.open(gpu.gpa, gpu.io, path); + p.open_n += 1; + } + } + + pub fn close(p: *Pack, io: std.Io) void { + for (p.files[0..p.open_n]) |*f| f.close(io); + } + + pub fn get(p: *const Pack, comptime fmt: []const u8, args: anytype) !st.Tensor { + var buf: [200]u8 = undefined; + const name = try std.fmt.bufPrint(&buf, fmt, args); + for (p.files[0..p.open_n]) |*f| if (f.get(name)) |t| return t; + std.debug.print("missing tensor {s}\n", .{name}); + return error.MissingTensor; + } +}; + +// ---------------------------------------------------------------------------------------------------- one layer + +/// One layer's experts: host trellis slices (mmap), the stacked scales and the fully resident reference on the device. +pub const LayerData = struct { + layer: usize, + routed: [routed_n][3][]const u8, + shared: [3][]const u8, + scale_bufs: [6]cuda.DeviceBuffer, + res_mem: cuda.DeviceBuffer, + res_tab: cuda.DeviceBuffer, + res: exl3e.Layer, + + pub fn load(gpu: Gpu, pack: *const Pack, layer: usize) !*LayerData { + const gpa = gpu.gpa; + const ld = try gpa.create(LayerData); + errdefer gpa.destroy(ld); + ld.layer = layer; + // stacked scales [E, ·]: gate suh [D] svh [I], up suh [D] svh [I], down suh [I] svh [D] + const sizes = [6]usize{ dims, dims, width, width, width, dims }; // suh_g suh_u svh_g svh_u suh_d svh_d + var host: [6][]u8 = undefined; + for (&host, sizes) |*h, n| h.* = try gpa.alloc(u8, table_n * n * 2); + defer for (host) |h| gpa.free(h); + var total: usize = 0; + for (0..table_n) |e| { + for (projs, 0..) |pn, p| { + const tr = if (e < routed_n) + try pack.get("model.language_model.layers.{d}.mlp.experts.{d}.{s}.trellis", .{ layer, e, pn }) + else + try pack.get("model.language_model.layers.{d}.mlp.shared_expert.{s}.trellis", .{ layer, pn }); + const su = if (e < routed_n) + try pack.get("model.language_model.layers.{d}.mlp.experts.{d}.{s}.suh", .{ layer, e, pn }) + else + try pack.get("model.language_model.layers.{d}.mlp.shared_expert.{s}.suh", .{ layer, pn }); + const sv = if (e < routed_n) + try pack.get("model.language_model.layers.{d}.mlp.experts.{d}.{s}.svh", .{ layer, e, pn }) + else + try pack.get("model.language_model.layers.{d}.mlp.shared_expert.{s}.svh", .{ layer, pn }); + const k: usize = if (p == 2) width else dims; + const n: usize = if (p == 2) dims else width; + if (tr.rank != 3 or tr.dim(0) * 16 != k or tr.dim(1) * 16 != n) return error.BadShape; + if (su.bytes.len != k * 2 or sv.bytes.len != n * 2) return error.BadShape; + if (e < routed_n) ld.routed[e][p] = tr.bytes else ld.shared[p] = tr.bytes; + total += (tr.bytes.len + 255) / 256 * 256; + const su_i: usize = switch (p) { + 0 => 0, + 1 => 1, + else => 4, + }; + const sv_i: usize = switch (p) { + 0 => 2, + 1 => 3, + else => 5, + }; + @memcpy(host[su_i][e * k * 2 ..][0 .. k * 2], su.bytes); + @memcpy(host[sv_i][e * n * 2 ..][0 .. n * 2], sv.bytes); + } + } + var made: usize = 0; + errdefer for (ld.scale_bufs[0..made]) |*b| b.free(); + for (&ld.scale_bufs, host) |*b, h| { + b.* = try cuda.DeviceBuffer.fromHost(gpu.d, h); + made += 1; + } + // the fully resident reference: every expert's trellis back to back, its own tables + ld.res_mem = try cuda.DeviceBuffer.alloc(gpu.d, total); + errdefer ld.res_mem.free(); + var ptrs: [3][table_n]u64 = undefined; + var k2s: [3][table_n]i32 = undefined; + var at: usize = 0; + const tiles = (dims / 16) * (width / 16); + for (0..table_n) |e| for (0..3) |p| { + const bytes = if (e < routed_n) ld.routed[e][p] else ld.shared[p]; + try ld.res_mem.upload(at, bytes); + ptrs[p][e] = ld.res_mem.ptr + at; + k2s[p][e] = @intCast(bytes.len / (tiles * 16)); + at += (bytes.len + 255) / 256 * 256; + }; + const tl = cuda.exl3_pool.Tables.init(table_n); + const tab = try gpa.alloc(u8, tl.bytes); + defer gpa.free(tab); + @memset(tab, 0); + for (0..3) |p| { + @memcpy(tab[tl.ptr[p]..][0 .. table_n * 8], std.mem.sliceAsBytes(&ptrs[p])); + @memcpy(tab[tl.k2[p]..][0 .. table_n * 4], std.mem.sliceAsBytes(&k2s[p])); + } + ld.res_tab = try cuda.DeviceBuffer.fromHost(gpu.d, tab); + ld.res = .{ + .gate = ld.res_tab.ptr + tl.ptr[0], + .up = ld.res_tab.ptr + tl.ptr[1], + .down = ld.res_tab.ptr + tl.ptr[2], + .gate_k2 = ld.res_tab.ptr + tl.k2[0], + .up_k2 = ld.res_tab.ptr + tl.k2[1], + .down_k2 = ld.res_tab.ptr + tl.k2[2], + .suh_g = ld.scale_bufs[0].ptr, + .suh_u = ld.scale_bufs[1].ptr, + .svh_g = ld.scale_bufs[2].ptr, + .svh_u = ld.scale_bufs[3].ptr, + .suh_d = ld.scale_bufs[4].ptr, + .svh_d = ld.scale_bufs[5].ptr, + .experts = table_n, + .dims = dims, + .width = width, + .codebook = .mul1, + .act_mode = 1, + }; + std.debug.print(" layer {d}: 513 experts mapped, fully resident reference {d:.0} MiB; routed k2 gate/up/down {d}/{d}/{d}, shared {d}/{d}/{d}\n", .{ + layer, @as(f64, @floatFromInt(total)) / 1048576.0, k2s[0][0], k2s[1][0], k2s[2][0], k2s[0][512], k2s[1][512], k2s[2][512], + }); + return ld; + } + + pub fn free(ld: *LayerData, gpa: std.mem.Allocator) void { + ld.res_tab.free(); + ld.res_mem.free(); + for (&ld.scale_bufs) |*b| b.free(); + gpa.destroy(ld); + } + + pub fn scales(ld: *const LayerData) fpool.Scales { + return .{ .suh_g = ld.scale_bufs[0].ptr, .suh_u = ld.scale_bufs[1].ptr, .svh_g = ld.scale_bufs[2].ptr, .svh_u = ld.scale_bufs[3].ptr, .suh_d = ld.scale_bufs[4].ptr, .svh_d = ld.scale_bufs[5].ptr }; + } + + pub fn source(ld: *const LayerData) fpool.Source { + return .{ .routed = &ld.routed, .shared = ld.shared }; + } +}; + +// ---------------------------------------------------------------------------------------------------- inputs + +/// Skewed expert choice: Zipf(s) over a fixed random permutation of the 512 ids. +pub const Zipf = struct { + perm: [routed_n]u16, + cdf: [routed_n]f64, + + pub fn init(seed: u64, s: f64) Zipf { + var z: Zipf = undefined; + var prng = std.Random.DefaultPrng.init(seed); + for (&z.perm, 0..) |*p, i| p.* = @intCast(i); + prng.random().shuffle(u16, &z.perm); + var sum: f64 = 0; + for (&z.cdf, 0..) |*c, i| { + sum += 1.0 / std.math.pow(f64, @floatFromInt(i + 1), s); + c.* = sum; + } + for (&z.cdf) |*c| c.* /= sum; + return z; + } + + /// One row: 10 distinct routed ids then the shared expert. + pub fn row(z: *const Zipf, r: std.Random, out: []i32) void { + var n: usize = 0; + while (n < top_k) { + const u = r.float(f64); + var i: usize = 0; + while (i + 1 < routed_n and z.cdf[i] < u) i += 1; + const id: i32 = z.perm[i]; + if (std.mem.indexOfScalar(i32, out[0..n], id) != null) continue; + out[n] = id; + n += 1; + } + out[top_k] = 512; + } +}; + +pub fn bf16Bits(v: f32) u16 { + const b: u32 = @bitCast(v); + return @intCast((b + 0x7fff + ((b >> 16) & 1)) >> 16); +} + +/// The device buffers of one engine side (pool or resident): x, picks, weights, scratch, out. +pub const Io = struct { + x: cuda.DeviceBuffer, + picks: cuda.DeviceBuffer, + wts: cuda.DeviceBuffer, + out: cuda.DeviceBuffer, + sbuf: cuda.DeviceBuffer, + sc: exl3e.Scratch, + + pub fn init(gpu: Gpu, l: exl3e.Layer) !Io { + var v: Io = undefined; + v.x = try cuda.DeviceBuffer.alloc(gpu.d, max_rows * dims * 2); + v.picks = try cuda.DeviceBuffer.alloc(gpu.d, max_rows * slots * 4); + v.wts = try cuda.DeviceBuffer.alloc(gpu.d, max_rows * slots * 4); + v.out = try cuda.DeviceBuffer.alloc(gpu.d, max_rows * dims * 4); + v.sbuf = try cuda.DeviceBuffer.alloc(gpu.d, try exl3e.scratchBytes(l, max_rows, slots)); + v.sc = try exl3e.scratch(v.sbuf.ptr, l, max_rows, slots); + return v; + } + + pub fn free(v: *Io) void { + v.x.free(); + v.picks.free(); + v.wts.free(); + v.out.free(); + v.sbuf.free(); + } + + pub fn set(v: *Io, x: []const u16, picks: []const i32, wts: []const f32) !void { + try v.x.upload(0, std.mem.sliceAsBytes(x)); + try v.picks.upload(0, std.mem.sliceAsBytes(picks)); + try v.wts.upload(0, std.mem.sliceAsBytes(wts)); + } + + pub fn y(v: Io, gpa: std.mem.Allocator, rows: usize) ![]u8 { + const out = try gpa.alloc(u8, rows * slots * dims * 4); + try v.sbuf.d.check(v.sbuf.d.api.cuMemcpyDtoH_v2(out.ptr, v.sc.y, out.len), "cuMemcpyDtoH"); + return out; + } +}; + +pub fn call(ex: *const exl3e.Experts, s: cuda.Stream, l: exl3e.Layer, v: Io, picks: u64, rows: usize, combine: bool) !void { + if (combine) { + try ex.run(s, l, v.x.ptr, dims, .bf16, picks, v.wts.ptr, rows, v.out.ptr, v.sc); + } else { + try ex.runSlots(s, l, v.x.ptr, dims, .bf16, picks, rows, v.sc); + } +} + +/// One round through the pool: plan, then each pass staged, computed and released. +pub fn poolRound(pool: *fpool.Pool, ex: *const exl3e.Experts, s: cuda.Stream, v: Io, picks: []const i32, rows: usize, combine: bool) !*const fpool.Round { + const r = try pool.plan(picks, rows); + const l = pool.tables(); + for (0..r.n_passes) |j| { + try pool.stage(r, j, s); + try call(ex, s, l, v, pool.picksFor(r, j, v.picks.ptr), rows, combine); + try pool.release(r, j, s); + } + return r; +} + +pub fn genRound(z: *const Zipf, rng: std.Random, rows: usize, x: []u16, picks: []i32, wts: []f32) void { + for (x[0 .. rows * dims]) |*b| b.* = bf16Bits((rng.float(f32) - 0.5) * 2.0); + for (0..rows) |i| z.row(rng, picks[i * slots ..][0..slots]); + for (wts[0 .. rows * slots]) |*w| w.* = rng.float(f32) * 0.2; +} diff --git a/zig/tests/cuda/flashnext_exl3/offload/exl3_pool_equal_tests.zig b/zig/tests/cuda/flashnext_exl3/offload/exl3_pool_equal_tests.zig new file mode 100644 index 0000000000..2dcd93790d --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/offload/exl3_pool_equal_tests.zig @@ -0,0 +1,151 @@ +//! exl3-pool: gate 2, the pool against the resident experts. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const flash = @import("flashnext"); +const check = @import("../../check.zig"); + +const st = core.safetensors; +const exl3e = cuda.exl3_experts; +const fpool = flash.exl3_pool; +const policy = fpool.offload; +const Gpu = check.Gpu; + +const data = @import("exl3_pool_data.zig"); +const dims = data.dims; +const width = data.width; +const routed_n = data.routed_n; +const table_n = data.table_n; +const top_k = data.top_k; +const slots = data.slots; +const max_rows = data.max_rows; +const shard_count = data.shard_count; +const projs = data.projs; +const Pack = data.Pack; +const LayerData = data.LayerData; +const Zipf = data.Zipf; +const bf16Bits = data.bf16Bits; +const Io = data.Io; +const call = data.call; +const poolRound = data.poolRound; +const genRound = data.genRound; + +// ---------------------------------------------------------------------------------------------------- gate 2 + +fn verifyPool(gpu: Gpu, pool: *fpool.Pool, ld: *const LayerData) !usize { + const gpa = gpu.gpa; + try pool.lines.drain(); + const lb = pool.lines.layout.line_bytes; + const got = try gpa.alloc(u8, lb); + defer gpa.free(got); + // tables + const tl = pool.lines.tlay; + const tab = try gpa.alloc(u8, tl.bytes); + defer gpa.free(tab); + try pool.lines.tables.download(0, tab); + var checked: usize = 0; + var resident: [routed_n]?usize = @splat(null); + for (pool.slots_tab, pool.line_tab, 0..) |id, state, i| { + if (state == .empty) continue; + resident[id] = i; + try pool.lines.mem.download(i * lb, got); + for (0..3) |p| { + const want = ld.routed[id][p]; + try check.sameBytes("pool line vs mapped checkpoint", got[pool.lines.layout.off[p]..][0..want.len], want); + } + checked += 1; + } + for (0..table_n) |e| for (0..3) |p| { + const ptr = std.mem.readInt(u64, tab[tl.ptr[p] + e * 8 ..][0..8], .little); + const k2 = std.mem.readInt(i32, tab[tl.k2[p] + e * 4 ..][0..4], .little); + if (e == 512) { + try check.expect(ptr != 0 and k2 > 0, "shared expert entry missing", .{}); + } else if (resident[e]) |line| { + try check.expect(ptr == pool.lines.at(line, p) and k2 == pool.k2[e][p], "expert {d} proj {d}: table points at {x}, line {d} is at {x}", .{ e, p, ptr, line, pool.lines.at(line, p) }); + } else { + try check.expect(ptr == 0 and k2 == 0, "absent expert {d} has a stale table entry", .{e}); + } + }; + return checked; +} + +const Script = struct { rows: usize, combine: bool }; + +pub fn equalRun(gpu: Gpu, ex: *const exl3e.Experts, s: cuda.Stream, ld: *const LayerData, lines: usize, prefetch_lines: usize, seed: u64) !void { + const gpa = gpu.gpa; + var pool = try fpool.Pool.init(gpa, gpu.io, gpu.d, ld.source(), ld.scales(), .{ .layer = @intCast(ld.layer), .lines = lines, .prefetch_bytes = 0 }, null); + defer pool.deinit(); + pool.opt.prefetch_bytes = prefetch_lines * pool.lineBytes(); + var vp = try Io.init(gpu, pool.tables()); + defer vp.free(); + var vr = try Io.init(gpu, ld.res); + defer vr.free(); + const z = Zipf.init(0x5eed + ld.layer, 1.1); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + // decode rows, two 16-row prompt chunks, a 4-row chunk; runSlots and run (with combine) alternate + var script: [40]Script = undefined; + for (&script, 0..) |*sc, i| sc.* = .{ .rows = 1, .combine = i % 2 == 1 }; + script[6] = .{ .rows = max_rows, .combine = false }; + script[17] = .{ .rows = max_rows, .combine = true }; + script[25] = .{ .rows = 4, .combine = false }; + var x: [max_rows * dims]u16 = undefined; + var picks: [max_rows * slots]i32 = undefined; + var wts: [max_rows * slots]f32 = undefined; + var ledger_staged: u64 = 0; + var ledger_h2d: u64 = 0; + var ledger_copies: u64 = 0; + var split_rounds: usize = 0; + var max_passes: usize = 0; + var slot_bytes: usize = 0; + for (script, 0..) |sc, ri| { + genRound(&z, rng, sc.rows, &x, &picks, &wts); + const n = sc.rows * slots; + try vp.set(x[0 .. sc.rows * dims], picks[0..n], wts[0..n]); + try vr.set(x[0 .. sc.rows * dims], picks[0..n], wts[0..n]); + try vp.sbuf.fill8(0xff, s.handle); // a slot no pass wrote shows up + const r = try poolRound(&pool, ex, s, vp, picks[0..n], sc.rows, sc.combine); + for (r.copies[0..r.n_copies]) |c| { + const b = ld.routed[c.id]; + ledger_staged += b[0].len + b[1].len + b[2].len; + ledger_h2d += pool.lines.layout.off[2] + b[2].len; + ledger_copies += 1; + } + if (r.masked()) split_rounds += 1; + max_passes = @max(max_passes, r.n_passes); + try call(ex, s, ld.res, vr, vr.picks.ptr, sc.rows, sc.combine); + try s.synchronize(); + const yp = try vp.y(gpa, sc.rows); + defer gpa.free(yp); + const yr = try vr.y(gpa, sc.rows); + defer gpa.free(yr); + check.sameBytes("pool y vs resident y", yp, yr) catch |e| { + std.debug.print(" round {d} ({d} rows, {d} passes) differs\n", .{ ri, sc.rows, r.n_passes }); + return e; + }; + slot_bytes += yp.len; + if (sc.combine) { + const op = try gpa.alloc(u8, sc.rows * dims * 4); + defer gpa.free(op); + const or_ = try gpa.alloc(u8, sc.rows * dims * 4); + defer gpa.free(or_); + try vp.out.download(0, op); + try vr.out.download(0, or_); + try check.sameBytes("pool combined out vs resident", op, or_); + } + } + const lines_checked = try verifyPool(gpu, &pool, ld); + const ps = pool.stats; + const ls = pool.lines.stats; + const hit = @as(f64, @floatFromInt(ps.hit_slots)) / @as(f64, @floatFromInt(ps.hit_slots + ps.miss_slots)); + std.debug.print(" [{d} lines, prefetch {d}/swap round] {d} rounds ({d} split into passes, up to {d}): routed slots hit {d} miss {d} ({d:.3} hit); {d} copies ({d} misses + {d} prefetch), {d} evictions, {d} protected rounds; H2D {d:.1} MiB, staged {d:.1} MiB (host memcpy {d:.2} ms/expert), staging waits {d}, release waits {d}, gate waits {d}; y {d} bytes compared\n", .{ + lines, prefetch_lines, ps.rounds, split_rounds, max_passes, ps.hit_slots, ps.miss_slots, hit, ls.copies, ps.misses, ps.prefetches, ps.evictions, ps.protected, + @as(f64, @floatFromInt(ls.bytes)) / 1048576.0, @as(f64, @floatFromInt(ls.staged_bytes)) / 1048576.0, @as(f64, @floatFromInt(ls.stage_ns)) / 1e6 / @as(f64, @floatFromInt(@max(ls.copies, 1))), ls.staging_waits, ls.release_waits, ls.gate_waits, slot_bytes, + }); + try check.expect(ls.copies == ledger_copies and ls.copies == ps.misses + ps.prefetches, "copy count: device {d}, plans {d}, misses+prefetch {d}", .{ ls.copies, ledger_copies, ps.misses + ps.prefetches }); + try check.expect(ls.staged_bytes == ledger_staged, "staged bytes {d} vs plan ledger {d}", .{ ls.staged_bytes, ledger_staged }); + try check.expect(ls.bytes == ledger_h2d, "H2D bytes {d} vs plan ledger {d}", .{ ls.bytes, ledger_h2d }); + try check.expect(ps.evictions > 0 and ps.miss_slots > 0 and ps.hit_slots > 0 and split_rounds > 0, "a path was not exercised (evictions {d}, misses {d}, hits {d}, split {d})", .{ ps.evictions, ps.miss_slots, ps.hit_slots, split_rounds }); + if (prefetch_lines > 0) try check.expect(ps.prefetches > 0, "no prefetch happened", .{}); + check.pass("layer {d}, {d} lines: every round's runSlots/run output is byte-identical to the fully resident layer; {d} lines and 513 table entries match the checkpoint; staged bytes match the plans", .{ ld.layer, lines, lines_checked }); +} diff --git a/zig/tests/cuda/flashnext_exl3/offload/exl3_pool_hits_tests.zig b/zig/tests/cuda/flashnext_exl3/offload/exl3_pool_hits_tests.zig new file mode 100644 index 0000000000..fc36a6d404 --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/offload/exl3_pool_hits_tests.zig @@ -0,0 +1,166 @@ +//! exl3-pool: hit rates of the cache policy replayed over recorded routing. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const flash = @import("flashnext"); +const check = @import("../../check.zig"); + +const st = core.safetensors; +const exl3e = cuda.exl3_experts; +const fpool = flash.exl3_pool; +const policy = fpool.offload; +const Gpu = check.Gpu; + +const data = @import("exl3_pool_data.zig"); +const dims = data.dims; +const width = data.width; +const routed_n = data.routed_n; +const table_n = data.table_n; +const top_k = data.top_k; +const slots = data.slots; +const max_rows = data.max_rows; +const shard_count = data.shard_count; +const projs = data.projs; +const Pack = data.Pack; +const LayerData = data.LayerData; +const Zipf = data.Zipf; +const bf16Bits = data.bf16Bits; +const Io = data.Io; +const call = data.call; +const poolRound = data.poolRound; +const genRound = data.genRound; + +// ---------------------------------------------------------------------------------------------------- hit rates + +const line_counts = [_]usize{ 10, 24, 48, 96, 160 }; + +const Sim = struct { + hit: u64 = 0, + miss: u64 = 0, + dec_hit: u64 = 0, + dec_miss: u64 = 0, + copies: u64 = 0, +}; + +/// Host-only replay of a pick sequence through one layer's policy. +fn simulate(gpa: std.mem.Allocator, layer: u8, picks: []const u16, tokens: usize, prompt: usize, lines: usize, prefetch_lines: usize, counts: *[policy.shape.layers][policy.shape.experts]u32) !Sim { + const slots_tab = try gpa.alloc(u16, lines); + defer gpa.free(slots_tab); + const line_tab = try gpa.alloc(policy.Line, lines); + defer gpa.free(line_tab); + const stamp = try gpa.alloc(u32, lines); + defer gpa.free(stamp); + const lay = try gpa.alloc(u8, lines); + defer gpa.free(lay); + var v = policy.Vram.init(slots_tab, line_tab, stamp, lay); + var r = try fpool.Round.init(gpa, lines, max_rows * slots); + defer r.deinit(gpa); + var buf: [max_rows * slots]i32 = undefined; + var sim: Sim = .{}; + var prev: [routed_n]u16 = undefined; + var prev_n: usize = 0; + var t: usize = 0; + var round: u32 = 0; + while (t < tokens) { + const rows: usize = if (t < prompt) @min(max_rows, prompt - t) else 1; + for (0..rows) |i| { + for (0..top_k) |k| buf[i * slots + k] = picks[(t + i) * top_k + k]; + buf[i * slots + top_k] = 512; + } + round += 1; + try fpool.round.plan(&v, layer, buf[0 .. rows * slots], counts, round, round, .{ .prefetch_bytes = prefetch_lines, .line_bytes = 1, .protect = prev[0..prev_n] }, &r); + @memcpy(prev[0..r.n_ids], r.ids[0..r.n_ids]); + prev_n = r.n_ids; + // prefetch lines land before the next round (the device has a whole layer stack of time) + for (r.prefetch()) |c| policy.reapAt(&v, layer, c.id, round) catch {}; + sim.hit += r.hit_slots; + sim.miss += r.miss_slots; + sim.copies += r.n_copies; + if (t >= prompt) { + sim.dec_hit += r.hit_slots; + sim.dec_miss += r.miss_slots; + } + t += rows; + } + return sim; +} + +fn rate(h: u64, m: u64) f64 { + return if (h + m == 0) 0 else @as(f64, @floatFromInt(h)) / @as(f64, @floatFromInt(h + m)); +} + +pub fn hitRun(gpu: Gpu, files: []const [:0]const u8) !void { + const gpa = gpu.gpa; + const counts = try gpa.create([policy.shape.layers][policy.shape.experts]u32); + defer gpa.destroy(counts); + // synthetic: 200 tokens of Zipf(1.1) picks, decode rows only + { + const z = Zipf.init(0x200, 1.1); + var prng = std.Random.DefaultPrng.init(0x200); + var row: [slots]i32 = undefined; + var picks: [200 * top_k]u16 = undefined; + for (0..200) |t| { + z.row(prng.random(), &row); + for (0..top_k) |k| picks[t * top_k + k] = @intCast(row[k]); + } + std.debug.print(" synthetic Zipf(1.1) over 512 experts, 200 decode tokens, one layer:\n", .{}); + for (line_counts) |lines| { + counts.* = @splat(@splat(0)); + const a = try simulate(gpa, 0, &picks, 200, 0, lines, 0, counts); + counts.* = @splat(@splat(0)); + const b = try simulate(gpa, 0, &picks, 200, 0, lines, 4, counts); + std.debug.print(" {d:>3} lines ({d:>5.1} MiB/layer at 1.76 MiB a line): hit {d:.3} (with 4-line prefetch every 4 rounds {d:.3})\n", .{ lines, @as(f64, @floatFromInt(lines)) * 1.7578, rate(a.hit, a.miss), rate(b.hit, b.miss) }); + } + } + for (files) |path| { + const blob = try std.Io.Dir.cwd().readFileAlloc(gpu.io, path, gpa, .limited(1 << 30)); + defer gpa.free(blob); + const hdr = std.mem.bytesAsSlice(u32, blob[0..12]); + const tokens: usize = hdr[0]; + const prompt: usize = hdr[1]; + const layers: usize = hdr[2]; + const raw = try gpa.alloc(u16, tokens * layers * top_k); + defer gpa.free(raw); + @memcpy(std.mem.sliceAsBytes(raw), blob[12..][0 .. raw.len * 2]); + std.debug.print(" real router picks {s}: {d} tokens ({d} prompt as 16-row chunks, {d} decode rows) x {d} layers\n", .{ path, tokens, prompt, tokens - prompt, layers }); + const per = try gpa.alloc(u16, tokens * top_k); + defer gpa.free(per); + // distinct experts a layer over the decode part (an upper bound on what any pool must hold) + var distinct_sum: usize = 0; + for (0..layers) |l| { + var seen: [routed_n]bool = @splat(false); + for (prompt..tokens) |t| for (0..top_k) |k| { + seen[raw[(t * layers + l) * top_k + k]] = true; + }; + for (seen) |b| distinct_sum += @intFromBool(b); + } + std.debug.print(" distinct experts a layer over the {d} decode tokens: {d:.1} on average\n", .{ tokens - prompt, @as(f64, @floatFromInt(distinct_sum)) / @as(f64, @floatFromInt(layers)) }); + for (line_counts) |lines| { + for ([_]usize{ 0, 4 }) |pf| { + counts.* = @splat(@splat(0)); + var tot: Sim = .{}; + var lo: f64 = 1; + var hi: f64 = 0; + var copies: u64 = 0; + for (0..layers) |l| { + for (0..tokens) |t| @memcpy(per[t * top_k ..][0..top_k], raw[(t * layers + l) * top_k ..][0..top_k]); + const s = try simulate(gpa, @intCast(l), per, tokens, prompt, lines, pf, counts); + tot.hit += s.hit; + tot.miss += s.miss; + tot.dec_hit += s.dec_hit; + tot.dec_miss += s.dec_miss; + copies += s.copies; + const dr = rate(s.dec_hit, s.dec_miss); + lo = @min(lo, dr); + hi = @max(hi, dr); + } + const dec_tokens = tokens - prompt; + const miss_per_tok = @as(f64, @floatFromInt(tot.dec_miss)) / @as(f64, @floatFromInt(dec_tokens)); + std.debug.print(" {d:>3} lines{s}: decode hit {d:.3} (layers {d:.3}..{d:.3}), {d:.1} routed misses a decode token over 48 layers; all rows hit {d:.3}; pool {d:.2} GiB for 48 layers\n", .{ + lines, if (pf > 0) " +prefetch" else " ", rate(tot.dec_hit, tot.dec_miss), lo, hi, miss_per_tok, rate(tot.hit, tot.miss), @as(f64, @floatFromInt(lines * 48)) * 1.7578 / 1024.0, + }); + } + } + } + check.pass("hit rates printed (host simulation of the pool's own planner)", .{}); +} diff --git a/zig/tests/cuda/flashnext_exl3/offload/exl3_pool_tests.zig b/zig/tests/cuda/flashnext_exl3/offload/exl3_pool_tests.zig new file mode 100644 index 0000000000..53a2a3e594 --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/offload/exl3_pool_tests.zig @@ -0,0 +1,57 @@ +//! exl3-pool: the Flash Next EXL3 expert offload pool (cuda_exl3_pool.zig) on a real pack; this hub parses the mode. +const std = @import("std"); +const cuda = @import("cuda"); +const check = @import("../../check.zig"); +const data = @import("exl3_pool_data.zig"); +const equal = @import("exl3_pool_equal_tests.zig"); +const bench = @import("exl3_pool_bench_tests.zig"); +const hits = @import("exl3_pool_hits_tests.zig"); + +const exl3e = cuda.exl3_experts; +const Gpu = check.Gpu; +const Pack = data.Pack; +const LayerData = data.LayerData; + +fn parseLayers(spec: []const u8, out: []usize) !usize { + var n: usize = 0; + var it = std.mem.splitScalar(u8, spec, ','); + while (it.next()) |tok| { + if (n == out.len) return error.TooManyLayers; + out[n] = try std.fmt.parseInt(usize, tok, 10); + n += 1; + } + return n; +} + +pub fn run(gpu: Gpu, rest: []const [:0]const u8) !void { + if (rest.len < 1) return error.MissingArgument; + const dir = rest[0]; + var mode: []const u8 = "all"; + var layer_spec: []const u8 = "0,17"; + if (rest.len > 1) { + if (std.ascii.isDigit(rest[1][0])) layer_spec = rest[1] else mode = rest[1]; + } + if (std.mem.eql(u8, mode, "hit")) return hits.hitRun(gpu, rest[2..]); + if (rest.len > 2) layer_spec = rest[2]; + var layers: [8]usize = undefined; + const n_layers = try parseLayers(layer_spec, &layers); + var pack: Pack = .{}; + defer pack.close(gpu.io); + try pack.open(gpu, dir); + var ex = try exl3e.Experts.load(gpu.d); + defer ex.unload(); + var s = try cuda.Stream.init(gpu.d, true); + defer s.deinit(); + for (layers[0..n_layers]) |layer| { + std.debug.print("== layer {d}\n", .{layer}); + const ld = try LayerData.load(gpu, &pack, layer); + defer ld.free(gpu.gpa); + const all = std.mem.eql(u8, mode, "all"); + if (all or std.mem.eql(u8, mode, "equal")) { + try equal.equalRun(gpu, &ex, s, ld, 10, 0, 0x10 + layer); + try equal.equalRun(gpu, &ex, s, ld, 24, 4, 0x24 + layer); + try equal.equalRun(gpu, &ex, s, ld, 64, 0, 0x64 + layer); + } + if (std.mem.eql(u8, mode, "bench")) try bench.benchRun(gpu, &ex, s, ld); + } +} diff --git a/zig/tests/cuda/flashnext_exl3/pack/exl3_format.zig b/zig/tests/cuda/flashnext_exl3/pack/exl3_format.zig new file mode 120000 index 0000000000..0618faed80 --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/pack/exl3_format.zig @@ -0,0 +1 @@ +../../../../src/core/exl3_format.zig \ No newline at end of file diff --git a/zig/tests/cuda/flashnext_exl3/pack/exl3_pack_tests.zig b/zig/tests/cuda/flashnext_exl3/pack/exl3_pack_tests.zig new file mode 100644 index 0000000000..a7c195b25d --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/pack/exl3_pack_tests.zig @@ -0,0 +1,522 @@ +//! exl3-pack: v1.1.0's EXL3 operators (exl3.Dense, exl3_experts.Experts) on a real Flash Next EXL3 pack's shapes. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const check = @import("../../check.zig"); +const exl3fmt = core.exl3_format; +const st = core.safetensors; +const exl3 = cuda.exl3; +const exl3e = cuda.exl3_experts; +const Gpu = check.Gpu; + +const shard_count = 7; + +const Pack = struct { + files: [shard_count]st.File = undefined, + open_n: usize = 0, + + fn open(p: *Pack, gpu: Gpu, dir: []const u8) !void { + var buf: [512]u8 = undefined; + for (0..shard_count) |i| { + const path = try std.fmt.bufPrint(&buf, "{s}/model-{d:0>5}-of-{d:0>5}.safetensors", .{ dir, i + 1, shard_count }); + p.files[i] = try st.File.open(gpu.gpa, gpu.io, path); + p.open_n += 1; + } + } + + fn close(p: *Pack, io: std.Io) void { + for (p.files[0..p.open_n]) |*f| f.close(io); + } + + fn get(p: *const Pack, name: []const u8) ?st.Tensor { + for (p.files[0..p.open_n]) |*f| if (f.get(name)) |t| return t; + return null; + } +}; + +/// One EXL3 matrix copied to aligned host memory: trellis int16 [K/16, N/16, 8 * halves], fp16 suh [K], svh [N]. +const Mat = struct { + trellis: []i16, + suh: []f16, + svh: []f16, + tiles_k: usize, + tiles_n: usize, + halves: u16, + + fn k(m: Mat) usize { + return m.tiles_k * 16; + } + fn n(m: Mat) usize { + return m.tiles_n * 16; + } +}; + +fn copyAs(comptime T: type, a: std.mem.Allocator, bytes: []const u8) ![]T { + const out = try a.alloc(T, bytes.len / @sizeOf(T)); + @memcpy(std.mem.sliceAsBytes(out), bytes[0 .. out.len * @sizeOf(T)]); + return out; +} + +fn mat(p: *const Pack, a: std.mem.Allocator, base: []const u8) !Mat { + var buf: [256]u8 = undefined; + const t = p.get(try std.fmt.bufPrint(&buf, "{s}.trellis", .{base})) orelse return error.MissingTensor; + const su = p.get(try std.fmt.bufPrint(&buf, "{s}.suh", .{base})) orelse return error.MissingTensor; + const sv = p.get(try std.fmt.bufPrint(&buf, "{s}.svh", .{base})) orelse return error.MissingTensor; + if (p.get(try std.fmt.bufPrint(&buf, "{s}.mul1", .{base})) == null) return error.NotMul1; + if (t.rank != 3 or t.dtype != .i16) return error.UnexpectedTensor; + return .{ + .trellis = try copyAs(i16, a, t.bytes), + .suh = try copyAs(f16, a, su.bytes), + .svh = try copyAs(f16, a, sv.bytes), + .tiles_k = t.dim(0), + .tiles_n = t.dim(1), + .halves = @intCast(t.dim(2) / 8), + }; +} + +fn bf16Bits(v: f32) u16 { + const b: u32 = @bitCast(v); + return @intCast((b + 0x7fff + ((b >> 16) & 1)) >> 16); +} + +fn bfValue(bits: u16) f64 { + return @floatCast(@as(f32, @bitCast(@as(u32, bits) << 16))); +} + +/// y [rows, N] = x [rows, K] · diag(suh) H Wq H diag(svh), all in float64 (x taken as given, no fp16 step). +fn exact(a: std.mem.Allocator, m: Mat, x: []const f64, rows: usize, y: []f64) !void { + const k = m.k(); + const n = m.n(); + const wq = try a.alloc(f16, k * n); + defer a.free(wq); + exl3fmt.unpack(.{ .halves = m.halves }, .mul1, m.trellis, m.tiles_k, m.tiles_n, wq); + const xs = try a.alloc(f64, rows * k); + defer a.free(xs); + for (0..rows) |r| for (0..k) |i| { + xs[r * k + i] = x[r * k + i] * @as(f64, @floatCast(m.suh[i])); + }; + exl3fmt.rotate(xs, rows, k, 1); + @memset(y, 0); + for (0..k) |ki| { + const row_w = wq[ki * n ..][0..n]; + for (0..rows) |r| { + const xv = xs[r * k + ki]; + const yr = y[r * n ..][0..n]; + for (yr, row_w) |*o, w| o.* += xv * @as(f64, @floatCast(w)); + } + } + exl3fmt.rotate(y, rows, n, 1); + for (0..rows) |r| for (0..n) |c| { + y[r * n + c] *= @as(f64, @floatCast(m.svh[c])); + }; +} + +const Err = struct { max_abs: f64 = 0, peak: f64 = 0, worst_rel: f64 = 0 }; + +fn compareBf16(got: []const u16, want: []const f64) Err { + var e: Err = .{}; + for (got, want) |g, w| { + const d = @abs(bfValue(g) - w); + e.max_abs = @max(e.max_abs, d); + e.peak = @max(e.peak, @abs(w)); + } + e.worst_rel = e.max_abs / @max(e.peak, 1e-30); + return e; +} + +fn compareF32(got: []const f32, want: []const f64) Err { + var e: Err = .{}; + for (got, want) |g, w| { + const d = @abs(@as(f64, g) - w); + e.max_abs = @max(e.max_abs, d); + e.peak = @max(e.peak, @abs(w)); + } + e.worst_rel = e.max_abs / @max(e.peak, 1e-30); + return e; +} + +const Bench = struct { + gpu: Gpu, + s: cuda.Stream, + ev0: cuda.Event, + ev1: cuda.Event, + reps: usize, + total_ms: f64 = 0, + + fn time(b: *Bench, ctx: anytype, comptime f: fn (@TypeOf(ctx)) anyerror!void) !f64 { + for (0..3) |_| try f(ctx); + try b.s.synchronize(); + var best: f64 = std.math.inf(f64); + for (0..5) |_| { + try b.ev0.record(b.s); + for (0..b.reps) |_| try f(ctx); + try b.ev1.record(b.s); + try b.ev1.synchronize(); + const ms = try b.ev0.elapsedMs(b.ev1); + best = @min(best, @as(f64, ms) / @as(f64, @floatFromInt(b.reps))); + } + return best; + } +}; + +const Call = struct { + de: *exl3.Dense, + s: cuda.Stream, + w: exl3.Weight, + x: u64, + y: u64, + rows: usize, + fn go(c: Call) anyerror!void { + try c.de.decode(c.s, c.w, c.x, .bf16, c.rows, c.y, .bf16); + } +}; + +fn randX(a: std.mem.Allocator, rng: std.Random, n: usize) !struct { bits: []u16, vals: []f64 } { + const bits = try a.alloc(u16, n); + const vals = try a.alloc(f64, n); + for (bits, vals) |*b, *v| { + b.* = bf16Bits((rng.float(f32) - 0.5) * 2.0); + v.* = bfValue(b.*); + } + return .{ .bits = bits, .vals = vals }; +} + +fn linear(b: *Bench, de: *exl3.Dense, pack: *const Pack, a: std.mem.Allocator, rng: std.Random, name: []const u8, impl: @TypeOf(de.implementation)) !void { + const m = mat(pack, a, name) catch |e| { + std.debug.print("{s}: skipped ({t})\n", .{ name, e }); + return; + }; + const k = m.k(); + const n = m.n(); + const strides = exl3.strides(k, n, m.halves, false); + var tr = try cuda.DeviceBuffer.fromHost(b.gpu.d, std.mem.sliceAsBytes(m.trellis)); + defer tr.free(); + var su = try cuda.DeviceBuffer.fromHost(b.gpu.d, std.mem.sliceAsBytes(m.suh)); + defer su.free(); + var sv = try cuda.DeviceBuffer.fromHost(b.gpu.d, std.mem.sliceAsBytes(m.svh)); + defer sv.free(); + const w: exl3.Weight = .{ .trellis = tr.ptr, .suh = su.ptr, .svh = sv.ptr, .k = @intCast(k), .n = @intCast(n), .halves = m.halves, .codebook = .mul1, .stride_k = strides.k, .stride_nb = strides.nb }; + if (impl == .tiles) { + _ = exl3.Dense.candidateWeight(w) catch |e| { + std.debug.print("{s:<58} tiles: refused ({t})\n", .{ name, e }); + return; + }; + } + de.implementation = impl; + defer de.implementation = .reference; + for ([_]usize{ 1, 5 }) |rows| { + const x = try randX(a, rng, rows * k); + var xd = try cuda.DeviceBuffer.fromHost(b.gpu.d, std.mem.sliceAsBytes(x.bits)); + defer xd.free(); + var yd = try cuda.DeviceBuffer.alloc(b.gpu.d, rows * n * 2); + defer yd.free(); + try yd.fill8(0xff, b.s.handle); + const call: Call = .{ .de = de, .s = b.s, .w = w, .x = xd.ptr, .y = yd.ptr, .rows = rows }; + try call.go(); + try b.s.synchronize(); + const got = try a.alloc(u16, rows * n); + try yd.download(0, std.mem.sliceAsBytes(got)); + // the reference row count: lm_head's 636M products a row, so only its first row + const ref_rows: usize = if (n > 100_000) 1 else rows; + const want = try a.alloc(f64, ref_rows * n); + try exact(a, m, x.vals[0 .. ref_rows * k], ref_rows, want); + const e = compareBf16(got[0 .. ref_rows * n], want); + // a second call's bytes equal the first's (deterministic) + try call.go(); + try b.s.synchronize(); + const again = try a.alloc(u16, rows * n); + try yd.download(0, std.mem.sliceAsBytes(again)); + const same = std.mem.eql(u16, got, again); + const ms = try b.time(call, Call.go); + if (rows == 1) b.total_ms += ms; + std.debug.print("{s:<58} {s:<9} K={d:<5} N={d:<6} halves={d:<2} rows={d} ms={d:>8.4} max|err|={e:.3} peak={d:.3} rel={e:.2} repeat={s}\n", .{ name, @tagName(impl), k, n, m.halves, rows, ms, e.max_abs, e.peak, e.worst_rel, if (same) "same" else "DIFF" }); + if (rows == 1) { + // the same row with fp32 outputs: the accumulation's own error, without bf16's output rounding + var yf = try cuda.DeviceBuffer.alloc(b.gpu.d, n * 4); + defer yf.free(); + try de.decode(b.s, w, xd.ptr, .bf16, 1, yf.ptr, .f32); + try b.s.synchronize(); + const gf = try a.alloc(f32, n); + try yf.download(0, std.mem.sliceAsBytes(gf)); + const ef = compareF32(gf, want[0..n]); + std.debug.print("{s:<58} {s:<9} fp32 out rows=1: max|err|={e:.3} peak={d:.3} rel={e:.2}\n", .{ name, @tagName(impl), ef.max_abs, ef.peak, ef.worst_rel }); + } + } +} + +const ExpCall = struct { + ex: *const exl3e.Experts, + s: cuda.Stream, + l: exl3e.Layer, + x: u64, + pick: u64, + wts: u64, + out: u64, + sc: exl3e.Scratch, + fn go(c: ExpCall) anyerror!void { + try c.ex.run(c.s, c.l, c.x, c.l.dims, .bf16, c.pick, c.wts, 1, c.out, c.sc); + } +}; + +/// Ten experts of `layer` as E = 10, one row picking all ten (weights 0.1): fp32 out vs float64. +fn experts(b: *Bench, ex: *const exl3e.Experts, pack: *const Pack, a: std.mem.Allocator, rng: std.Random, layer: usize, act_mode: i32, mapped: usize) !void { + const count = 10; + var buf: [160]u8 = undefined; + var g: [count]Mat = undefined; + var u: [count]Mat = undefined; + var dn: [count]Mat = undefined; + for (0..count) |e| { + g[e] = try mat(pack, a, try std.fmt.bufPrint(&buf, "model.language_model.layers.{d}.mlp.experts.{d}.gate_proj", .{ layer, e })); + u[e] = try mat(pack, a, try std.fmt.bufPrint(&buf, "model.language_model.layers.{d}.mlp.experts.{d}.up_proj", .{ layer, e })); + dn[e] = try mat(pack, a, try std.fmt.bufPrint(&buf, "model.language_model.layers.{d}.mlp.experts.{d}.down_proj", .{ layer, e })); + } + const dims = g[0].k(); + const width = g[0].n(); + var bufs: std.ArrayList(cuda.DeviceBuffer) = .empty; + defer { + for (bufs.items) |*x| x.free(); + bufs.deinit(b.gpu.gpa); + } + const up = struct { + fn f(gpu: Gpu, list: *std.ArrayList(cuda.DeviceBuffer), bytes: []const u8) !u64 { + const d = try cuda.DeviceBuffer.fromHost(gpu.d, bytes); + try list.append(gpu.gpa, d); + return d.ptr; + } + }.f; + var hosts: std.ArrayList(cuda.HostBuffer) = .empty; + defer { + for (hosts.items) |*h| h.free(); + hosts.deinit(b.gpu.gpa); + } + const hostUp = struct { + fn f(gpu: Gpu, list: *std.ArrayList(cuda.HostBuffer), bytes: []const u8) !u64 { + var h = try cuda.HostBuffer.allocMapped(gpu.d, bytes.len); + errdefer h.free(); + @memcpy(h.bytes[0..bytes.len], bytes); + try list.append(gpu.gpa, h); + return try h.device(); + } + }.f; + var gp: [count]u64 = undefined; + var upp: [count]u64 = undefined; + var dp: [count]u64 = undefined; + var gk2: [count]i32 = undefined; + var uk2: [count]i32 = undefined; + var dk2: [count]i32 = undefined; + var suh_g: std.ArrayList(f16) = .empty; + var suh_u: std.ArrayList(f16) = .empty; + var svh_g: std.ArrayList(f16) = .empty; + var svh_u: std.ArrayList(f16) = .empty; + var suh_d: std.ArrayList(f16) = .empty; + var svh_d: std.ArrayList(f16) = .empty; + for (0..count) |e| { + if (e < mapped) { + gp[e] = try hostUp(b.gpu, &hosts, std.mem.sliceAsBytes(g[e].trellis)); + upp[e] = try hostUp(b.gpu, &hosts, std.mem.sliceAsBytes(u[e].trellis)); + dp[e] = try hostUp(b.gpu, &hosts, std.mem.sliceAsBytes(dn[e].trellis)); + } else { + gp[e] = try up(b.gpu, &bufs, std.mem.sliceAsBytes(g[e].trellis)); + upp[e] = try up(b.gpu, &bufs, std.mem.sliceAsBytes(u[e].trellis)); + dp[e] = try up(b.gpu, &bufs, std.mem.sliceAsBytes(dn[e].trellis)); + } + gk2[e] = g[e].halves; + uk2[e] = u[e].halves; + dk2[e] = dn[e].halves; + try suh_g.appendSlice(a, g[e].suh); + try suh_u.appendSlice(a, u[e].suh); + try svh_g.appendSlice(a, g[e].svh); + try svh_u.appendSlice(a, u[e].svh); + try suh_d.appendSlice(a, dn[e].suh); + try svh_d.appendSlice(a, dn[e].svh); + } + const l: exl3e.Layer = .{ + .gate = try up(b.gpu, &bufs, std.mem.sliceAsBytes(&gp)), + .up = try up(b.gpu, &bufs, std.mem.sliceAsBytes(&upp)), + .down = try up(b.gpu, &bufs, std.mem.sliceAsBytes(&dp)), + .gate_k2 = try up(b.gpu, &bufs, std.mem.sliceAsBytes(&gk2)), + .up_k2 = try up(b.gpu, &bufs, std.mem.sliceAsBytes(&uk2)), + .down_k2 = try up(b.gpu, &bufs, std.mem.sliceAsBytes(&dk2)), + .suh_g = try up(b.gpu, &bufs, std.mem.sliceAsBytes(suh_g.items)), + .suh_u = try up(b.gpu, &bufs, std.mem.sliceAsBytes(suh_u.items)), + .svh_g = try up(b.gpu, &bufs, std.mem.sliceAsBytes(svh_g.items)), + .svh_u = try up(b.gpu, &bufs, std.mem.sliceAsBytes(svh_u.items)), + .suh_d = try up(b.gpu, &bufs, std.mem.sliceAsBytes(suh_d.items)), + .svh_d = try up(b.gpu, &bufs, std.mem.sliceAsBytes(svh_d.items)), + .experts = count, + .dims = dims, + .width = width, + .codebook = .mul1, + .act_mode = act_mode, + }; + const x = try randX(a, rng, dims); + var pick: [count]i32 = undefined; + for (&pick, 0..) |*p, i| p.* = @intCast((i * 7) % count); + var wts: [count]f32 = @splat(0.1); + const sbytes = try exl3e.scratchBytes(l, 1, count); + var sbuf = try cuda.DeviceBuffer.alloc(b.gpu.d, sbytes); + defer sbuf.free(); + var outd = try cuda.DeviceBuffer.alloc(b.gpu.d, dims * 4); + defer outd.free(); + try outd.fill8(0xff, b.s.handle); + const call: ExpCall = .{ + .ex = ex, + .s = b.s, + .l = l, + .x = try up(b.gpu, &bufs, std.mem.sliceAsBytes(x.bits)), + .pick = try up(b.gpu, &bufs, std.mem.sliceAsBytes(&pick)), + .wts = try up(b.gpu, &bufs, std.mem.sliceAsBytes(&wts)), + .out = outd.ptr, + .sc = try exl3e.scratch(sbuf.ptr, l, 1, count), + }; + try call.go(); + try b.s.synchronize(); + const got = try a.alloc(f32, dims); + try outd.download(0, std.mem.sliceAsBytes(got)); + // float64: out = sum_e 0.1 * down_e(silu(gate_e(x)) * up_e(x)) + const want = try a.alloc(f64, dims); + @memset(want, 0); + const gy = try a.alloc(f64, width); + const uy = try a.alloc(f64, width); + const dy = try a.alloc(f64, dims); + for (pick) |pe| { + const e: usize = @intCast(pe); + try exact(a, g[e], x.vals, 1, gy); + try exact(a, u[e], x.vals, 1, uy); + for (gy, uy) |*gv, uv| gv.* = gv.* / (1.0 + @exp(-gv.*)) * uv; + try exact(a, dn[e], gy, 1, dy); + for (want, dy) |*o, v| o.* += 0.1 * v; + } + const err = compareF32(got, want); + const ms = try b.time(call, ExpCall.go); + // with a 64 MiB memset between calls (evicts the 32 MiB L2, so mapped host bytes come over PCIe each call). + var flush = try cuda.DeviceBuffer.alloc(b.gpu.d, 64 << 20); + defer flush.free(); + const Flushed = struct { + c: ExpCall, + f: u64, + with_call: bool, + fn go(fl: @This()) anyerror!void { + try fl.c.s.d.check(fl.c.s.d.api.cuMemsetD8Async(fl.f, 0, 64 << 20, fl.c.s.handle), "cuMemsetD8Async"); + if (fl.with_call) try fl.c.go(); + } + }; + const with = try b.time(Flushed{ .c = call, .f = flush.ptr, .with_call = true }, Flushed.go); + const alone = try b.time(Flushed{ .c = call, .f = flush.ptr, .with_call = false }, Flushed.go); + std.debug.print(" cold L2 (64 MiB memset between calls, its time subtracted): {d:.4} ms\n", .{with - alone}); + std.debug.print("routed experts L{d} (1 row x 10 picks, act_mode {d}, {d} of 10 from mapped host memory) K={d} I={d} halves={d}/{d}/{d} ms={d:>8.4} max|err|={e:.3} peak={d:.3} rel={e:.2}\n", .{ layer, act_mode, mapped, dims, width, g[0].halves, u[0].halves, dn[0].halves, ms, err.max_abs, err.peak, err.worst_rel }); + if (mapped > 0) { + // the copy-engine alternative: the same experts' trellis bytes, pinned host to device. + var total: usize = 0; + for (0..mapped) |e| total += (g[e].trellis.len + u[e].trellis.len + dn[e].trellis.len) * 2; + var src = try cuda.HostBuffer.alloc(b.gpu.d, total); + defer src.free(); + var dst = try cuda.DeviceBuffer.alloc(b.gpu.d, total); + defer dst.free(); + const Copy = struct { + d: *const cuda.Driver, + s: cuda.Stream, + src: []u8, + dst: u64, + fn go(c: @This()) anyerror!void { + try c.d.check(c.d.api.cuMemcpyHtoDAsync_v2(c.dst, c.src.ptr, c.src.len, c.s.handle), "cuMemcpyHtoDAsync"); + } + }; + const c: Copy = .{ .d = b.gpu.d, .s = b.s, .src = src.bytes[0..total], .dst = dst.ptr }; + const cms = try b.time(c, Copy.go); + std.debug.print(" pinned H2D copy of those {d} experts' trellis ({d} bytes): {d:.4} ms ({d:.1} GB/s)\n", .{ mapped, total, cms, @as(f64, @floatFromInt(total)) / cms / 1e6 }); + } +} + +/// Rows=1 timing only, warm and with a 64 MiB memset between calls (the memset's own time subtracted). +fn coldLinear(b: *Bench, de: *exl3.Dense, pack: *const Pack, a: std.mem.Allocator, rng: std.Random, name: []const u8) !void { + const m = try mat(pack, a, name); + const k = m.k(); + const n = m.n(); + const strides = exl3.strides(k, n, m.halves, false); + var tr = try cuda.DeviceBuffer.fromHost(b.gpu.d, std.mem.sliceAsBytes(m.trellis)); + defer tr.free(); + var su = try cuda.DeviceBuffer.fromHost(b.gpu.d, std.mem.sliceAsBytes(m.suh)); + defer su.free(); + var sv = try cuda.DeviceBuffer.fromHost(b.gpu.d, std.mem.sliceAsBytes(m.svh)); + defer sv.free(); + const w: exl3.Weight = .{ .trellis = tr.ptr, .suh = su.ptr, .svh = sv.ptr, .k = @intCast(k), .n = @intCast(n), .halves = m.halves, .codebook = .mul1, .stride_k = strides.k, .stride_nb = strides.nb }; + const x = try randX(a, rng, k); + var xd = try cuda.DeviceBuffer.fromHost(b.gpu.d, std.mem.sliceAsBytes(x.bits)); + defer xd.free(); + var yd = try cuda.DeviceBuffer.alloc(b.gpu.d, n * 2); + defer yd.free(); + var flush = try cuda.DeviceBuffer.alloc(b.gpu.d, 64 << 20); + defer flush.free(); + const call: Call = .{ .de = de, .s = b.s, .w = w, .x = xd.ptr, .y = yd.ptr, .rows = 1 }; + const warm = try b.time(call, Call.go); + const F = struct { + c: Call, + f: u64, + with_call: bool, + fn go(fl: @This()) anyerror!void { + try fl.c.s.d.check(fl.c.s.d.api.cuMemsetD8Async(fl.f, 0, 64 << 20, fl.c.s.handle), "cuMemsetD8Async"); + if (fl.with_call) try fl.c.go(); + } + }; + const with = try b.time(F{ .c = call, .f = flush.ptr, .with_call = true }, F.go); + const alone = try b.time(F{ .c = call, .f = flush.ptr, .with_call = false }, F.go); + const bytes = m.trellis.len * 2; + std.debug.print("{s:<58} rows=1 bytes={d:>10} warm {d:.4} ms cold {d:.4} ms ({d:.0} GB/s cold)\n", .{ name, bytes, warm, with - alone, @as(f64, @floatFromInt(bytes)) / (with - alone) / 1e6 }); +} + +pub fn run(gpu: Gpu, rest: []const [:0]const u8) !void { + if (rest.len < 1) return error.MissingArgument; + const dir = rest[0]; + const reps: usize = if (rest.len > 1) try std.fmt.parseInt(usize, rest[1], 10) else 20; + var arena = std.heap.ArenaAllocator.init(gpu.gpa); + defer arena.deinit(); + const a = arena.allocator(); + var pack: Pack = .{}; + defer pack.close(gpu.io); + try pack.open(gpu, dir); + var s = try cuda.Stream.init(gpu.d, true); + defer s.deinit(); + var b: Bench = .{ .gpu = gpu, .s = s, .ev0 = try cuda.Event.init(gpu.d, true), .ev1 = try cuda.Event.init(gpu.d, true), .reps = reps }; + defer b.ev0.deinit(); + defer b.ev1.deinit(); + var de = try exl3.Dense.load(gpu.d); + defer de.unload(); + var ex = try exl3e.Experts.load(gpu.d); + defer ex.unload(); + var prng = std.Random.DefaultPrng.init(0x5eed); + const rng = prng.random(); + + const names = [_][]const u8{ + "model.language_model.layers.0.linear_attn.in_proj_qkv", + "model.language_model.layers.0.linear_attn.in_proj_z", + "model.language_model.layers.0.linear_attn.out_proj", + "model.language_model.layers.3.self_attn.q_proj", + "model.language_model.layers.3.self_attn.k_proj", + "model.language_model.layers.3.self_attn.v_proj", + "model.language_model.layers.3.self_attn.o_proj", + "model.language_model.layers.3.mlp.shared_expert.gate_proj", + "model.language_model.layers.3.mlp.shared_expert.down_proj", + "model.language_model.layers.3.self_attn.indexer.index_qk_proj", + "model.language_model.layers.0.mlp.experts.0.gate_proj", + "model.language_model.layers.0.mlp.experts.0.down_proj", + "lm_head", + }; + if (rest.len > 2 and std.mem.eql(u8, rest[2], "cold")) { + for (names) |nm| try coldLinear(&b, &de, &pack, a, rng, nm); + try experts(&b, &ex, &pack, a, rng, 17, 1, 0); + return; + } + if (rest.len > 2 and std.mem.eql(u8, rest[2], "experts")) { + for ([_]usize{ 0, 1, 4, 10 }) |mapped| try experts(&b, &ex, &pack, a, rng, 0, 1, mapped); + return; + } + for (names) |nm| try linear(&b, &de, &pack, a, rng, nm, .reference); + std.debug.print("sum of rows=1 ms over the projections above (reference Dense): {d:.3}\n", .{b.total_ms}); + b.total_ms = 0; + for (names) |nm| try linear(&b, &de, &pack, a, rng, nm, .tiles); + try experts(&b, &ex, &pack, a, rng, 0, 1, 0); + try experts(&b, &ex, &pack, a, rng, 0, 0, 0); + try experts(&b, &ex, &pack, a, rng, 17, 1, 0); + check.pass("exl3-pack: v1.1.0 EXL3 operators ran on the pack's shapes (see the error and timing lines)", .{}); +} diff --git a/zig/tests/cuda/flashnext_exl3/pack/exl3_place.zig b/zig/tests/cuda/flashnext_exl3/pack/exl3_place.zig new file mode 120000 index 0000000000..0a791de308 --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/pack/exl3_place.zig @@ -0,0 +1 @@ +../../../../src/families/flashnext/exl3/pack/exl3_place.zig \ No newline at end of file diff --git a/zig/tests/cuda/flashnext_exl3/pack/exl3_qmat.zig b/zig/tests/cuda/flashnext_exl3/pack/exl3_qmat.zig new file mode 120000 index 0000000000..e690658604 --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/pack/exl3_qmat.zig @@ -0,0 +1 @@ +../../../../src/families/flashnext/exl3/pack/exl3_qmat.zig \ No newline at end of file diff --git a/zig/tests/cuda/flashnext_exl3/pack/exl3_scan.zig b/zig/tests/cuda/flashnext_exl3/pack/exl3_scan.zig new file mode 120000 index 0000000000..640d68021d --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/pack/exl3_scan.zig @@ -0,0 +1 @@ +../../../../src/families/flashnext/exl3/pack/exl3_scan.zig \ No newline at end of file diff --git a/zig/tests/cuda/flashnext_exl3/pack/exl3_shape.zig b/zig/tests/cuda/flashnext_exl3/pack/exl3_shape.zig new file mode 120000 index 0000000000..e4dccac406 --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/pack/exl3_shape.zig @@ -0,0 +1 @@ +../../../../src/families/flashnext/exl3/pack/exl3_shape.zig \ No newline at end of file diff --git a/zig/tests/cuda/flashnext_exl3/prefill/exl3_prefill_copy.zig b/zig/tests/cuda/flashnext_exl3/prefill/exl3_prefill_copy.zig new file mode 100644 index 0000000000..0aeb73061d --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/prefill/exl3_prefill_copy.zig @@ -0,0 +1,157 @@ +//! exl3-prefill copy: the expert line copies of the offload pool, timed and checked on one layer. + +const std = @import("std"); +const cuda = @import("cuda"); +const check = @import("../../check.zig"); +const exl3e = cuda.exl3_experts; +const Gpu = check.Gpu; +const packs = @import("exl3_prefill_pack.zig"); +const experts = @import("exl3_prefill_experts.zig"); + +const dims = 2560; +const width = 640; +const routed_n = 512; +const table_n = 513; +const slots = 11; + +const LayerData = packs.LayerData; +const bf16Bits = packs.bf16Bits; +const genPicks = packs.genPicks; +const SlotsCall = experts.SlotsCall; + +// ---------------------------------------------------------------------------------------------------- copies + +pub fn copyRun(gpu: Gpu, ex: *const exl3e.Experts, s: cuda.Stream, ld: *const LayerData) !void { + const a = gpu.gpa; + // the layer's routed trellis in one pinned bank (what S9 keeps), lines laid out as Lines does + var total: usize = 0; + for (ld.routed) |e| for (e) |b| { + total += b.len; + }; + var bank = try cuda.HostBuffer.alloc(gpu.d, total); + defer bank.free(); + const pinned = try a.alloc([3][]const u8, routed_n); + defer a.free(pinned); + var at: usize = 0; + for (ld.routed, pinned) |e, *o| for (e, 0..) |b, p| { + @memcpy(bank.bytes[at..][0..b.len], b); + o[p] = bank.bytes[at..][0..b.len]; + at += b.len; + }; + const n_lines: usize = 64; + var bytes: [3]usize = .{ 0, 0, 0 }; + var k2: [routed_n][3]i32 = undefined; + const tiles = (dims / 16) * (width / 16); + for (pinned, &k2) |e, *w| for (0..3) |p| { + bytes[p] = @max(bytes[p], e[p].len); + w[p] = @intCast(e[p].len / (tiles * 16)); + }; + var lines = try cuda.exl3_pool.Lines.init(a, gpu.d, .{ .experts = table_n, .bytes = bytes, .lines = n_lines, .staging = 0 }); + defer lines.deinit(); + var c0 = try cuda.Event.init(gpu.d, true); + defer c0.deinit(); + var c1 = try cuda.Event.init(gpu.d, true); + defer c1.deinit(); + var raw = try cuda.DeviceBuffer.alloc(gpu.d, n_lines * lines.layout.line_bytes); + defer raw.free(); + const moved = blk: { + var m: usize = 0; + for (pinned[0..n_lines]) |e| m += e[0].len + e[1].len + e[2].len; + break :blk m; + }; + std.debug.print(" {d} expert lines ({d:.2} MiB, {d:.3} MiB a line) from the pinned bank, best of 5:\n", .{ n_lines, @as(f64, @floatFromInt(moved)) / 1048576.0, @as(f64, @floatFromInt(moved)) / 1048576.0 / n_lines }); + var best_load: f64 = std.math.inf(f64); + var best_raw: f64 = std.math.inf(f64); + var best_raw1: f64 = std.math.inf(f64); + for (0..5) |it| { + try lines.drain(); + try c0.record(lines.copy); + for (0..n_lines) |i| { + const id = (it * n_lines + i) % routed_n; + const prev = (id + routed_n - n_lines) % routed_n; + try lines.load(i, id, pinned[id], k2[id], if (it == 0) null else prev, true, gpu.io); + } + try c1.record(lines.copy); + try c1.synchronize(); + best_load = @min(best_load, try c0.elapsedMs(c1)); + // the same bytes as three H2D a line, no table writes + try c0.record(lines.copy); + for (0..n_lines) |i| for (0..3) |p| { + const e = pinned[i]; + try gpu.d.check(gpu.d.api.cuMemcpyHtoDAsync_v2(raw.ptr + i * lines.layout.line_bytes + lines.layout.off[p], e[p].ptr, e[p].len, lines.copy.handle), "cuMemcpyHtoDAsync"); + }; + try c1.record(lines.copy); + try c1.synchronize(); + best_raw = @min(best_raw, try c0.elapsedMs(c1)); + // one H2D of the bank's first n_lines experts (contiguous in the bank) + try c0.record(lines.copy); + try gpu.d.check(gpu.d.api.cuMemcpyHtoDAsync_v2(raw.ptr, bank.bytes.ptr, moved, lines.copy.handle), "cuMemcpyHtoDAsync"); + try c1.record(lines.copy); + try c1.synchronize(); + best_raw1 = @min(best_raw1, try c0.elapsedMs(c1)); + } + const gbs = struct { + fn f(b: usize, ms: f64) f64 { + return @as(f64, @floatFromInt(b)) / ms / 1e6; + } + }.f; + std.debug.print(" Lines.load (3 H2D + 6-9 table memsets + event a line): {d:.3} ms = {d:.3} ms a line, {d:.2} GB/s\n", .{ best_load, best_load / n_lines, gbs(moved, best_load) }); + std.debug.print(" 3 bare H2D a line: {d:.3} ms = {d:.3} ms a line, {d:.2} GB/s\n", .{ best_raw, best_raw / n_lines, gbs(moved, best_raw) }); + std.debug.print(" one H2D of the same bytes: {d:.3} ms, {d:.2} GB/s\n", .{ best_raw1, gbs(moved, best_raw1) }); + + // copies beside compute: a 1024-row experts call on the resident layer while the copy stream moves 64 lines + const rows: usize = 1024; + var prng = std.Random.DefaultPrng.init(0xc0b1); + const rng = prng.random(); + const xh = try a.alloc(u16, rows * dims); + defer a.free(xh); + for (xh) |*b| b.* = bf16Bits((rng.float(f32) - 0.5) * 2.0); + var x = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(xh)); + defer x.free(); + const picks = try a.alloc(i32, rows * slots); + defer a.free(picks); + genPicks(rng, 0, rows, picks); + var pk = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(picks)); + defer pk.free(); + var sbuf = try cuda.DeviceBuffer.alloc(gpu.d, try exl3e.scratchBytes(ld.res, rows, slots)); + defer sbuf.free(); + const sc = try exl3e.scratch(sbuf.ptr, ld.res, rows, slots); + const call: SlotsCall = .{ .ex = ex, .s = s, .l = ld.res, .x = x.ptr, .picks = pk.ptr, .rows = rows, .sc = sc }; + var k0 = try cuda.Event.init(gpu.d, true); + defer k0.deinit(); + var k1 = try cuda.Event.init(gpu.d, true); + defer k1.deinit(); + var alone_c: f64 = std.math.inf(f64); + var alone_k: f64 = std.math.inf(f64); + var both_c: f64 = std.math.inf(f64); + var both_k: f64 = std.math.inf(f64); + var both_w: f64 = std.math.inf(f64); + for (0..5) |_| { + try k0.record(s); + try call.go(); + try k1.record(s); + try k1.synchronize(); + alone_k = @min(alone_k, try k0.elapsedMs(k1)); + try c0.record(lines.copy); + try gpu.d.check(gpu.d.api.cuMemcpyHtoDAsync_v2(raw.ptr, bank.bytes.ptr, moved, lines.copy.handle), "cuMemcpyHtoDAsync"); + try c1.record(lines.copy); + try c1.synchronize(); + alone_c = @min(alone_c, try c0.elapsedMs(c1)); + // both at once + try s.synchronize(); + const t0 = check.now(gpu.io); + try c0.record(lines.copy); + try gpu.d.check(gpu.d.api.cuMemcpyHtoDAsync_v2(raw.ptr, bank.bytes.ptr, moved, lines.copy.handle), "cuMemcpyHtoDAsync"); + try c1.record(lines.copy); + try k0.record(s); + try call.go(); + try k1.record(s); + try k1.synchronize(); + try c1.synchronize(); + both_w = @min(both_w, @as(f64, @floatFromInt(check.now(gpu.io) - t0)) / 1e6); + both_c = @min(both_c, try c0.elapsedMs(c1)); + both_k = @min(both_k, try k0.elapsedMs(k1)); + } + std.debug.print(" overlap: {d}-row experts call alone {d:.3} ms, 64-line H2D alone {d:.3} ms; together: experts {d:.3} ms, H2D {d:.3} ms, wall {d:.3} ms (serial would be {d:.3})\n", .{ rows, alone_k, alone_c, both_k, both_c, both_w, alone_k + alone_c }); + check.pass("layer {d}: copy cost and copy/compute overlap measured", .{ld.layer}); +} diff --git a/zig/tests/cuda/flashnext_exl3/prefill/exl3_prefill_dense.zig b/zig/tests/cuda/flashnext_exl3/prefill/exl3_prefill_dense.zig new file mode 100644 index 0000000000..833f380083 --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/prefill/exl3_prefill_dense.zig @@ -0,0 +1,171 @@ +//! exl3-prefill dense: the prompt tiles of a dense EXL3 projection against a float64 reference. + +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const check = @import("../../check.zig"); +const exl3fmt = core.exl3_format; +const exl3 = cuda.exl3; +const Gpu = check.Gpu; +const packs = @import("exl3_prefill_pack.zig"); + +const Pack = packs.Pack; +const Timer = packs.Timer; +const bf16Bits = packs.bf16Bits; +const bfValue = packs.bfValue; + +// ---------------------------------------------------------------------------------------------------- dense + +const DenseCall = struct { + de: *exl3.Dense, + s: cuda.Stream, + w: exl3.Weight, + x: u64, + y: u64, + rows: usize, + pub fn go(c: DenseCall) anyerror!void { + try c.de.prompt(c.s, c.w, c.x, .bf16, c.rows, c.y, .bf16); + } +}; + +/// y [rows, N] = x [rows, K] . diag(suh) H Wq H diag(svh) in float64. +fn exact(a: std.mem.Allocator, trellis: []const i16, suh: []const f16, svh: []const f16, tk: usize, tn: usize, halves: u16, x: []const f64, rows: usize, y: []f64) !void { + const k = tk * 16; + const n = tn * 16; + const wq = try a.alloc(f16, k * n); + defer a.free(wq); + exl3fmt.unpack(.{ .halves = halves }, .mul1, trellis, tk, tn, wq); + const xs = try a.alloc(f64, rows * k); + defer a.free(xs); + for (0..rows) |r| for (0..k) |i| { + xs[r * k + i] = x[r * k + i] * @as(f64, @floatCast(suh[i])); + }; + exl3fmt.rotate(xs, rows, k, 1); + @memset(y, 0); + for (0..k) |ki| { + const row_w = wq[ki * n ..][0..n]; + for (0..rows) |r| { + const xv = xs[r * k + ki]; + const yr = y[r * n ..][0..n]; + for (yr, row_w) |*o, w| o.* += xv * @as(f64, @floatCast(w)); + } + } + exl3fmt.rotate(y, rows, n, 1); + for (0..rows) |r| for (0..n) |c| { + y[r * n + c] *= @as(f64, @floatCast(svh[c])); + }; +} + +fn copyAs(comptime T: type, a: std.mem.Allocator, bytes: []const u8) ![]T { + const out = try a.alloc(T, bytes.len / @sizeOf(T)); + @memcpy(std.mem.sliceAsBytes(out), bytes[0 .. out.len * @sizeOf(T)]); + return out; +} + +pub fn denseRun(gpu: Gpu, pack: *const Pack, s: cuda.Stream) !void { + const a = gpu.gpa; + var de = try exl3.Dense.load(gpu.d); + defer de.unload(); + var t = try Timer.init(gpu.d); + defer t.deinit(); + const names = [_][]const u8{ + "model.language_model.layers.0.linear_attn.in_proj_qkv", + "model.language_model.layers.0.linear_attn.in_proj_z", + "model.language_model.layers.0.linear_attn.out_proj", + "model.language_model.layers.3.self_attn.q_proj", + "model.language_model.layers.3.self_attn.k_proj", + "model.language_model.layers.3.self_attn.o_proj", + "model.language_model.layers.3.mlp.shared_expert.gate_proj", + "model.language_model.layers.3.mlp.shared_expert.down_proj", + }; + const row_counts = [_]usize{ 16, 128, 256, 1024, 2048 }; + const max_rows: usize = 2048; + var prng = std.Random.DefaultPrng.init(0xde45e); + const rng = prng.random(); + var sum_ref: [row_counts.len]f64 = @splat(0); + var sum_tiles: [row_counts.len]f64 = @splat(0); + for (names) |name| { + const tr = try pack.get("{s}.trellis", .{name}); + const su = try pack.get("{s}.suh", .{name}); + const sv = try pack.get("{s}.svh", .{name}); + const tk = tr.dim(0); + const tn = tr.dim(1); + const halves: u16 = @intCast(tr.dim(2) / 8); + const k = tk * 16; + const n = tn * 16; + const trellis = try copyAs(i16, a, tr.bytes); + defer a.free(trellis); + const suh = try copyAs(f16, a, su.bytes); + defer a.free(suh); + const svh = try copyAs(f16, a, sv.bytes); + defer a.free(svh); + var trd = try cuda.DeviceBuffer.fromHost(gpu.d, tr.bytes); + defer trd.free(); + var sud = try cuda.DeviceBuffer.fromHost(gpu.d, su.bytes); + defer sud.free(); + var svd = try cuda.DeviceBuffer.fromHost(gpu.d, sv.bytes); + defer svd.free(); + const strides = exl3.strides(k, n, halves, false); + const w: exl3.Weight = .{ .trellis = trd.ptr, .suh = sud.ptr, .svh = svd.ptr, .k = @intCast(k), .n = @intCast(n), .halves = halves, .codebook = .mul1, .stride_k = strides.k, .stride_nb = strides.nb }; + const tiles_ok = if (exl3.Dense.candidateWeight(w)) |_| true else |_| false; + const xh = try a.alloc(u16, max_rows * k); + defer a.free(xh); + for (xh) |*b| b.* = bf16Bits((rng.float(f32) - 0.5) * 2.0); + var x = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(xh)); + defer x.free(); + var yr = try cuda.DeviceBuffer.alloc(gpu.d, max_rows * n * 2); + defer yr.free(); + var yt = try cuda.DeviceBuffer.alloc(gpu.d, max_rows * n * 2); + defer yt.free(); + // f64 of rows 0 and 1 + const xf = try a.alloc(f64, 2 * k); + defer a.free(xf); + for (xf, xh[0 .. 2 * k]) |*o, b| o.* = bfValue(b); + const want = try a.alloc(f64, 2 * n); + defer a.free(want); + try exact(a, trellis, suh, svh, tk, tn, halves, xf, 2, want); + std.debug.print(" {s} K={d} N={d} halves={d}{s}\n", .{ name, k, n, halves, if (tiles_ok) "" else " (tiles refuse it)" }); + for (row_counts, 0..) |rows, ri| { + de.implementation = .reference; + const rc: DenseCall = .{ .de = &de, .s = s, .w = w, .x = x.ptr, .y = yr.ptr, .rows = rows }; + const ms_ref = try t.best(s, 3, rc, DenseCall.go); + sum_ref[ri] += ms_ref; + var ms_tiles: f64 = std.math.nan(f64); + var diff: f64 = std.math.nan(f64); + var err_ref: f64 = 0; + var err_tiles: f64 = std.math.nan(f64); + var peak: f64 = 0; + const got_r = try a.alloc(u16, rows * n); + defer a.free(got_r); + try yr.download(0, std.mem.sliceAsBytes(got_r)); + for (got_r[0 .. 2 * n], want) |g, wv| { + err_ref = @max(err_ref, @abs(bfValue(g) - wv)); + peak = @max(peak, @abs(wv)); + } + if (tiles_ok) { + de.implementation = .tiles; + const tc: DenseCall = .{ .de = &de, .s = s, .w = w, .x = x.ptr, .y = yt.ptr, .rows = rows }; + ms_tiles = try t.best(s, 3, tc, DenseCall.go); + sum_tiles[ri] += ms_tiles; + const got_t = try a.alloc(u16, rows * n); + defer a.free(got_t); + try yt.download(0, std.mem.sliceAsBytes(got_t)); + diff = 0; + var rpeak: f64 = 0; + for (got_r, got_t) |g, h| { + diff = @max(diff, @abs(bfValue(g) - bfValue(h))); + rpeak = @max(rpeak, @abs(bfValue(g))); + } + diff /= rpeak; + err_tiles = 0; + for (got_t[0 .. 2 * n], want) |g, wv| err_tiles = @max(err_tiles, @abs(bfValue(g) - wv)); + de.implementation = .reference; + } else sum_tiles[ri] += ms_ref; + std.debug.print(" rows {d:5}: reference {d:8.4} ms ({d:6.2} us/row) tiles {d:8.4} ms ({d:6.2} us/row) max|ref-tiles|/peak {e:.2} rows 0-1 vs f64: ref {e:.2} tiles {e:.2} (rel to peak {d:.3})\n", .{ rows, ms_ref, 1e3 * ms_ref / @as(f64, @floatFromInt(rows)), ms_tiles, 1e3 * ms_tiles / @as(f64, @floatFromInt(rows)), diff, err_ref / peak, err_tiles / peak, peak }); + } + } + std.debug.print(" sum over the {d} shapes (ms): ", .{names.len}); + for (row_counts, sum_ref, sum_tiles) |r, x, y| std.debug.print("rows {d}: reference {d:.3} best {d:.3}; ", .{ r, x, y }); + std.debug.print("\n", .{}); + check.pass("dense prompt kernels timed", .{}); +} diff --git a/zig/tests/cuda/flashnext_exl3/prefill/exl3_prefill_experts.zig b/zig/tests/cuda/flashnext_exl3/prefill/exl3_prefill_experts.zig new file mode 100644 index 0000000000..2dcca86822 --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/prefill/exl3_prefill_experts.zig @@ -0,0 +1,118 @@ +//! exl3-prefill experts: the grouped, slot and expert-major masked passes timed on one layer. + +const std = @import("std"); +const cuda = @import("cuda"); +const check = @import("../../check.zig"); +const exl3e = cuda.exl3_experts; +const Gpu = check.Gpu; +const packs = @import("exl3_prefill_pack.zig"); + +const dims = 2560; +const width = 640; +const routed_n = 512; +const table_n = 513; +const slots = 11; +const max_call_rows: usize = 1408; + +const LayerData = packs.LayerData; +const Timer = packs.Timer; +const bf16Bits = packs.bf16Bits; +const genPicks = packs.genPicks; + +pub const SlotsCall = struct { + ex: *const exl3e.Experts, + s: cuda.Stream, + l: exl3e.Layer, + x: u64, + picks: u64, + rows: usize, + sc: exl3e.Scratch, + pub fn go(c: SlotsCall) anyerror!void { + try c.ex.runSlots(c.s, c.l, c.x, dims, .bf16, c.picks, c.rows, c.sc); + } +}; + +/// Expert-major passes over one block: pass j runs the block's picks of experts [j*per, (j+1)*per) (others -1). +const PassesCall = struct { + ex: *const exl3e.Experts, + s: cuda.Stream, + l: exl3e.Layer, + x: u64, + masked: u64, // int32 [passes, rows * slots] + rows: usize, + passes: usize, + sc: exl3e.Scratch, + pub fn go(c: PassesCall) anyerror!void { + for (0..c.passes) |j| try c.ex.runSlots(c.s, c.l, c.x, dims, .bf16, c.masked + j * c.rows * slots * 4, c.rows, c.sc); + } +}; + +pub fn expertsRun(gpu: Gpu, ex: *const exl3e.Experts, s: cuda.Stream, ld: *const LayerData) !void { + const a = gpu.gpa; + var t = try Timer.init(gpu.d); + defer t.deinit(); + const max_rows = max_call_rows; + var prng = std.Random.DefaultPrng.init(0x5eed); + const rng = prng.random(); + const xh = try a.alloc(u16, max_rows * dims); + defer a.free(xh); + for (xh) |*b| b.* = bf16Bits((rng.float(f32) - 0.5) * 2.0); + var x = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(xh)); + defer x.free(); + const picks = try a.alloc(i32, max_rows * slots); + defer a.free(picks); + var pk = try cuda.DeviceBuffer.alloc(gpu.d, max_rows * slots * 4); + defer pk.free(); + var sbuf = try cuda.DeviceBuffer.alloc(gpu.d, try exl3e.scratchBytes(ld.res, max_rows, slots)); + defer sbuf.free(); + const sc = try exl3e.scratch(sbuf.ptr, ld.res, max_rows, slots); + std.debug.print(" runSlots, one resident layer (10 routed + shared a row), ms a call (best of 5 x 5):\n", .{}); + for ([_]f64{ 0.8, 0 }) |zs| { + genPicks(rng, zs, max_rows, picks); + try pk.upload(0, std.mem.sliceAsBytes(picks)); + for ([_]usize{ 16, 64, 256, 512, 1024, max_rows }) |rows| { + var seen: [table_n]bool = @splat(false); + var uniq: usize = 0; + for (picks[0 .. rows * slots]) |p| if (!seen[@intCast(p)]) { + seen[@intCast(p)] = true; + uniq += 1; + }; + const ms = try t.best(s, 5, SlotsCall{ .ex = ex, .s = s, .l = ld.res, .x = x.ptr, .picks = pk.ptr, .rows = rows, .sc = sc }, SlotsCall.go); + std.debug.print(" {s:<8} rows {d:5}: {d:8.3} ms {d:7.2} us/row {d:6.0} rows/s/layer ({d} experts incl. shared, {d:.1} member rows each)\n", .{ if (zs == 0) "uniform" else "zipf0.8", rows, ms, 1e3 * ms / @as(f64, @floatFromInt(rows)), @as(f64, @floatFromInt(rows)) / ms * 1e3, uniq, @as(f64, @floatFromInt(rows * slots)) / @as(f64, @floatFromInt(uniq)) }); + } + } + // expert-major passes over one 1024-row block (uniform picks): bytes equal to one call, and the time + const rows: usize = 1024; + genPicks(rng, 0, rows, picks); + try pk.upload(0, std.mem.sliceAsBytes(picks[0 .. rows * slots])); + const one_bytes = try a.alloc(u8, rows * slots * dims * 4); + defer a.free(one_bytes); + try (SlotsCall{ .ex = ex, .s = s, .l = ld.res, .x = x.ptr, .picks = pk.ptr, .rows = rows, .sc = sc }).go(); + try s.synchronize(); + try gpu.d.check(gpu.d.api.cuMemcpyDtoH_v2(one_bytes.ptr, sc.y, one_bytes.len), "cuMemcpyDtoH"); + const base_ms = try t.best(s, 5, SlotsCall{ .ex = ex, .s = s, .l = ld.res, .x = x.ptr, .picks = pk.ptr, .rows = rows, .sc = sc }, SlotsCall.go); + std.debug.print(" expert-major passes over one {d}-row block (uniform picks; pass j: experts [j*512/P, (j+1)*512/P) plus the shared one in pass 0):\n", .{rows}); + for ([_]usize{ 1, 2, 4, 8, 16, 32 }) |passes| { + const per = routed_n / passes; + const host = try a.alloc(i32, passes * rows * slots); + defer a.free(host); + for (0..passes) |j| for (picks[0 .. rows * slots], host[j * rows * slots ..][0 .. rows * slots]) |p, *o| { + const in = if (p == 512) j == 0 else @as(usize, @intCast(p)) / per == j; + o.* = if (in) p else -1; + }; + var masked = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(host)); + defer masked.free(); + try gpu.d.check(gpu.d.api.cuMemsetD8_v2(sc.y, 0xff, rows * slots * dims * 4), "cuMemsetD8"); + const call: PassesCall = .{ .ex = ex, .s = s, .l = ld.res, .x = x.ptr, .masked = masked.ptr, .rows = rows, .passes = passes, .sc = sc }; + try call.go(); + try s.synchronize(); + const got = try a.alloc(u8, one_bytes.len); + defer a.free(got); + try gpu.d.check(gpu.d.api.cuMemcpyDtoH_v2(got.ptr, sc.y, got.len), "cuMemcpyDtoH"); + const same = std.mem.eql(u8, got, one_bytes); + const ms = try t.best(s, 3, call, PassesCall.go); + std.debug.print(" {d:2} passes of {d:3} experts: {d:8.3} ms (+{d:.3} ms vs one call, {d:.3} ms a pass) y {s}\n", .{ passes, per, ms, ms - base_ms, ms / @as(f64, @floatFromInt(passes)), if (same) "bit-identical to one call" else "DIFFERS" }); + try check.expect(same, "masked passes reproduce one call's bytes", .{}); + } + check.pass("layer {d}: runSlots timed at 16..{d} rows; expert-major masked passes are bit-identical to one call", .{ ld.layer, max_rows }); +} diff --git a/zig/tests/cuda/flashnext_exl3/prefill/exl3_prefill_pack.zig b/zig/tests/cuda/flashnext_exl3/prefill/exl3_prefill_pack.zig new file mode 100644 index 0000000000..366f189eca --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/prefill/exl3_prefill_pack.zig @@ -0,0 +1,218 @@ +//! exl3-prefill fixtures: the checkpoint pack, a timer, one layer's experts on the device and random picks. + +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const check = @import("../../check.zig"); +const st = core.safetensors; +const exl3e = cuda.exl3_experts; +const Gpu = check.Gpu; + +const dims = 2560; +const width = 640; +const routed_n = 512; +const table_n = 513; +const top_k = 10; +const slots = 11; +const shard_count = 7; +const projs = [3][]const u8{ "gate_proj", "up_proj", "down_proj" }; + +pub const Pack = struct { + files: [shard_count]st.File = undefined, + open_n: usize = 0, + + pub fn open(p: *Pack, gpu: Gpu, dir: []const u8) !void { + var buf: [512]u8 = undefined; + for (0..shard_count) |i| { + const path = try std.fmt.bufPrint(&buf, "{s}/model-{d:0>5}-of-{d:0>5}.safetensors", .{ dir, i + 1, shard_count }); + p.files[i] = try st.File.open(gpu.gpa, gpu.io, path); + p.open_n += 1; + } + } + + pub fn close(p: *Pack, io: std.Io) void { + for (p.files[0..p.open_n]) |*f| f.close(io); + } + + pub fn get(p: *const Pack, comptime fmt: []const u8, args: anytype) !st.Tensor { + var buf: [200]u8 = undefined; + const name = try std.fmt.bufPrint(&buf, fmt, args); + for (p.files[0..p.open_n]) |*f| if (f.get(name)) |t| return t; + std.debug.print("missing tensor {s}\n", .{name}); + return error.MissingTensor; + } +}; + +pub fn bf16Bits(v: f32) u16 { + const b: u32 = @bitCast(v); + return @intCast((b + 0x7fff + ((b >> 16) & 1)) >> 16); +} + +pub fn bfValue(bits: u16) f64 { + return @floatCast(@as(f32, @bitCast(@as(u32, bits) << 16))); +} + +pub const Timer = struct { + a: cuda.Event, + b: cuda.Event, + + pub fn init(d: *const cuda.Driver) !Timer { + return .{ .a = try cuda.Event.init(d, true), .b = try cuda.Event.init(d, true) }; + } + pub fn deinit(t: *Timer) void { + t.a.deinit(); + t.b.deinit(); + } + /// Best of 5 batches of `reps` calls, ms a call (after 2 warm-up calls). + pub fn best(t: *Timer, s: cuda.Stream, reps: usize, ctx: anytype, comptime f: fn (@TypeOf(ctx)) anyerror!void) !f64 { + for (0..2) |_| try f(ctx); + try s.synchronize(); + var b: f64 = std.math.inf(f64); + for (0..5) |_| { + try t.a.record(s); + for (0..reps) |_| try f(ctx); + try t.b.record(s); + try t.b.synchronize(); + b = @min(b, @as(f64, try t.a.elapsedMs(t.b)) / @as(f64, @floatFromInt(reps))); + } + return b; + } +}; + +/// One layer's 513 experts: the mmapped trellis, the stacked scales and every expert resident on the device. +pub const LayerData = struct { + layer: usize, + routed: [routed_n][3][]const u8, + shared: [3][]const u8, + scale_bufs: [6]cuda.DeviceBuffer, + res_mem: cuda.DeviceBuffer, + res_tab: cuda.DeviceBuffer, + res: exl3e.Layer, + line_bytes: usize, + + pub fn load(gpu: Gpu, pack: *const Pack, layer: usize) !*LayerData { + const gpa = gpu.gpa; + const ld = try gpa.create(LayerData); + errdefer gpa.destroy(ld); + ld.layer = layer; + const sizes = [6]usize{ dims, dims, width, width, width, dims }; // suh_g suh_u svh_g svh_u suh_d svh_d + var host: [6][]u8 = undefined; + for (&host, sizes) |*h, n| h.* = try gpa.alloc(u8, table_n * n * 2); + defer for (host) |h| gpa.free(h); + var total: usize = 0; + ld.line_bytes = 0; + for (0..table_n) |e| { + var line: usize = 0; + for (projs, 0..) |pn, p| { + const base = "model.language_model.layers.{d}.mlp."; + const tr = if (e < routed_n) try pack.get(base ++ "experts.{d}.{s}.trellis", .{ layer, e, pn }) else try pack.get(base ++ "shared_expert.{s}.trellis", .{ layer, pn }); + const su = if (e < routed_n) try pack.get(base ++ "experts.{d}.{s}.suh", .{ layer, e, pn }) else try pack.get(base ++ "shared_expert.{s}.suh", .{ layer, pn }); + const sv = if (e < routed_n) try pack.get(base ++ "experts.{d}.{s}.svh", .{ layer, e, pn }) else try pack.get(base ++ "shared_expert.{s}.svh", .{ layer, pn }); + const k: usize = if (p == 2) width else dims; + const n: usize = if (p == 2) dims else width; + if (tr.rank != 3 or tr.dim(0) * 16 != k or tr.dim(1) * 16 != n) return error.BadShape; + if (e < routed_n) ld.routed[e][p] = tr.bytes else ld.shared[p] = tr.bytes; + total += (tr.bytes.len + 255) / 256 * 256; + line += (tr.bytes.len + 255) / 256 * 256; + const su_i: usize = switch (p) { + 0 => 0, + 1 => 1, + else => 4, + }; + const sv_i: usize = switch (p) { + 0 => 2, + 1 => 3, + else => 5, + }; + @memcpy(host[su_i][e * k * 2 ..][0 .. k * 2], su.bytes); + @memcpy(host[sv_i][e * n * 2 ..][0 .. n * 2], sv.bytes); + } + if (e < routed_n) ld.line_bytes = @max(ld.line_bytes, line); + } + var made: usize = 0; + errdefer for (ld.scale_bufs[0..made]) |*b| b.free(); + for (&ld.scale_bufs, host) |*b, h| { + b.* = try cuda.DeviceBuffer.fromHost(gpu.d, h); + made += 1; + } + ld.res_mem = try cuda.DeviceBuffer.alloc(gpu.d, total); + errdefer ld.res_mem.free(); + var ptrs: [3][table_n]u64 = undefined; + var k2s: [3][table_n]i32 = undefined; + var at: usize = 0; + const tiles = (dims / 16) * (width / 16); + for (0..table_n) |e| for (0..3) |p| { + const bytes = if (e < routed_n) ld.routed[e][p] else ld.shared[p]; + try ld.res_mem.upload(at, bytes); + ptrs[p][e] = ld.res_mem.ptr + at; + k2s[p][e] = @intCast(bytes.len / (tiles * 16)); + at += (bytes.len + 255) / 256 * 256; + }; + const tl = cuda.exl3_pool.Tables.init(table_n); + const tab = try gpa.alloc(u8, tl.bytes); + defer gpa.free(tab); + @memset(tab, 0); + for (0..3) |p| { + @memcpy(tab[tl.ptr[p]..][0 .. table_n * 8], std.mem.sliceAsBytes(&ptrs[p])); + @memcpy(tab[tl.k2[p]..][0 .. table_n * 4], std.mem.sliceAsBytes(&k2s[p])); + } + ld.res_tab = try cuda.DeviceBuffer.fromHost(gpu.d, tab); + ld.res = .{ + .gate = ld.res_tab.ptr + tl.ptr[0], + .up = ld.res_tab.ptr + tl.ptr[1], + .down = ld.res_tab.ptr + tl.ptr[2], + .gate_k2 = ld.res_tab.ptr + tl.k2[0], + .up_k2 = ld.res_tab.ptr + tl.k2[1], + .down_k2 = ld.res_tab.ptr + tl.k2[2], + .suh_g = ld.scale_bufs[0].ptr, + .suh_u = ld.scale_bufs[1].ptr, + .svh_g = ld.scale_bufs[2].ptr, + .svh_u = ld.scale_bufs[3].ptr, + .suh_d = ld.scale_bufs[4].ptr, + .svh_d = ld.scale_bufs[5].ptr, + .experts = table_n, + .dims = dims, + .width = width, + .codebook = .mul1, + .act_mode = 1, + }; + std.debug.print(" layer {d}: 513 experts resident ({d:.0} MiB), routed k2 gate/up/down {d}/{d}/{d}, a routed line {d:.3} MiB\n", .{ layer, @as(f64, @floatFromInt(total)) / 1048576.0, k2s[0][0], k2s[1][0], k2s[2][0], @as(f64, @floatFromInt(ld.line_bytes)) / 1048576.0 }); + return ld; + } + + pub fn free(ld: *LayerData, gpa: std.mem.Allocator) void { + ld.res_tab.free(); + ld.res_mem.free(); + for (&ld.scale_bufs) |*b| b.free(); + gpa.destroy(ld); + } +}; + +/// Picks for `rows` rows: ten distinct routed ids then the shared expert; +pub fn genPicks(rng: std.Random, s: f64, rows: usize, out: []i32) void { + var perm: [routed_n]u16 = undefined; + for (&perm, 0..) |*p, i| p.* = @intCast(i); + var prng = std.Random.DefaultPrng.init(0x9e3779b9); + prng.random().shuffle(u16, &perm); + var cdf: [routed_n]f64 = undefined; + var sum: f64 = 0; + for (&cdf, 0..) |*c, i| { + sum += if (s == 0) 1.0 else 1.0 / std.math.pow(f64, @floatFromInt(i + 1), s); + c.* = sum; + } + for (&cdf) |*c| c.* /= sum; + for (0..rows) |r| { + const row = out[r * slots ..][0..slots]; + var n: usize = 0; + while (n < top_k) { + const u = rng.float(f64); + var i: usize = 0; + while (i + 1 < routed_n and cdf[i] < u) i += 1; + const id: i32 = perm[i]; + if (std.mem.indexOfScalar(i32, row[0..n], id) != null) continue; + row[n] = id; + n += 1; + } + row[top_k] = 512; + } +} diff --git a/zig/tests/cuda/flashnext_exl3/prefill/exl3_prefill_tests.zig b/zig/tests/cuda/flashnext_exl3/prefill/exl3_prefill_tests.zig new file mode 100644 index 0000000000..5695cf4cf4 --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/prefill/exl3_prefill_tests.zig @@ -0,0 +1,37 @@ +//! exl3-prefill: the pieces of a Flash Next EXL3 prefill, timed on the pack's real tensors (one idle card). + +const std = @import("std"); +const cuda = @import("cuda"); +const check = @import("../../check.zig"); +const exl3e = cuda.exl3_experts; +const Gpu = check.Gpu; +const packs = @import("exl3_prefill_pack.zig"); +const experts = @import("exl3_prefill_experts.zig"); +const copies = @import("exl3_prefill_copy.zig"); +const dense = @import("exl3_prefill_dense.zig"); + +const LayerData = packs.LayerData; +const Pack = packs.Pack; +const expertsRun = experts.expertsRun; +const copyRun = copies.copyRun; +const denseRun = dense.denseRun; + +pub fn run(gpu: Gpu, rest: []const [:0]const u8) !void { + if (rest.len < 2) return error.MissingArgument; + const dir = rest[0]; + const mode = rest[1]; + const layer: usize = if (rest.len > 2) try std.fmt.parseInt(usize, rest[2], 10) else 17; + var pack: Pack = .{}; + defer pack.close(gpu.io); + try pack.open(gpu, dir); + var s = try cuda.Stream.init(gpu.d, true); + defer s.deinit(); + if (std.mem.eql(u8, mode, "dense")) return denseRun(gpu, &pack, s); + var ex = try exl3e.Experts.load(gpu.d); + defer ex.unload(); + const ld = try LayerData.load(gpu, &pack, layer); + defer ld.free(gpu.gpa); + if (std.mem.eql(u8, mode, "experts")) return expertsRun(gpu, &ex, s, ld); + if (std.mem.eql(u8, mode, "copy")) return copyRun(gpu, &ex, s, ld); + return error.UnknownMode; +} diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index 2d37fc734a..14adcb4c19 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -95,6 +95,12 @@ const lane_half_tests = @import("lane_half_tests.zig"); const flashnext_dense_hc_tests = @import("flashnext_dense_hc_tests.zig"); const flashnext_hc_tests = @import("flashnext_hc_tests.zig"); const flashnext_attn_tests = @import("flashnext_attn_tests.zig"); +const exl3_pack_tests = @import("flashnext_exl3/pack/exl3_pack_tests.zig"); +const exl3_ple2_tests = @import("flashnext_exl3/model/exl3_ple2_tests.zig"); +const exl3_pool_tests = @import("flashnext_exl3/offload/exl3_pool_tests.zig"); +const exl3_prefill_tests = @import("flashnext_exl3/prefill/exl3_prefill_tests.zig"); +const flashnext_exl3_tests = @import("flashnext_exl3/model/flashnext_exl3_tests.zig"); +const flashnext_exl3_engine_tests = @import("flashnext_exl3/model/flashnext_exl3_engine_tests.zig"); const ggml_tests = @import("ggml_tests.zig"); const paged_tests = @import("paged_tests.zig"); const glm_loader_gate = @import("glm_loader_gate.zig"); @@ -135,6 +141,8 @@ const usage = \\ lane-half lane half tiles (MXFP8, BF16): bytes equal to the 64-column tile, FP64, cold-weight time \\ flashnext-dense-hc dense connection weights and FP32 normalization against FP64 \\ flashnext-hc Flash Next's own hyper-connection and n-gram glue against float64 (synthetic rows) + \\ flashnext-exl3 DIR names|norms|router [layers] [rows]|layer [layers] Flash Next on an EXL3 pack vs float64 + \\ flashnext-exl3-engine DIR tf|gen GOLDEN [LINES] | bench [LINES] the family's Engine with the offload pool vs the reference golden \\ flashnext-attn Flash Next's own attention glue, pooling and block scores against float64 (synthetic) \\ flashnext-pages the attention glue's caches through a page table: paged == contiguous, float64, controls \\ info driver, device and library versions @@ -223,6 +231,10 @@ const usage = \\ dense-rows bf16 rows times weight rows (routers, heads) against float64 references (synthetic) \\ exl3 [dir] EXL3 decode-row linear: float64 checks, row invariance, the Python engine's fixture (dir none skips it) \\ exl3-tp2 [dir] EXL3 tensor parallel: rank shards vs the Python engine's, all-gathered (dir none skips them) + \\ exl3-pack DIR [REPS] EXL3 Dense and grouped experts on a real EXL3 pack's shapes: float64 error and timing + \\ exl3-ple2 DIR [ROW] [N] EXL3 n-gram table through cuda_ple.Ngram (ids, rows, RSS) and the PLE kernels vs float64 + \\ exl3-prefill DIR experts|copy [LAYER]|dense EXL3 prefill pieces on the pack: grouped experts by rows, expert-major passes, line copies, prompt Dense + \\ exl3-pool DIR [equal|bench|hit FILES] [LAYERS] Flash Next EXL3 expert offload pool on Experts' pointer tables vs fully resident \\ exl3-tiles [DIR] EXL3's tile candidate: codebooks and tiles bit-equal, FP64, rows through qlinear and Dense \\ train-pack MODEL EXAMPLES [COUNT] training passes holding several sequences against one a pass: differences, time \\ kernel-contract the registered kernels through the contract's gates, each planted control red (needs -Dnvcc) @@ -333,6 +345,8 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "flashnext-dense-hc")) return flashnext_dense_hc_tests.run(gpu); if (std.mem.eql(u8, cmd, "flashnext-hc")) return flashnext_hc_tests.run(gpu); if (std.mem.eql(u8, cmd, "flashnext-attn")) return flashnext_attn_tests.run(gpu); + if (std.mem.eql(u8, cmd, "flashnext-exl3")) return flashnext_exl3_tests.run(gpu, rest); + if (std.mem.eql(u8, cmd, "flashnext-exl3-engine")) return flashnext_exl3_engine_tests.run(gpu, rest); if (std.mem.eql(u8, cmd, "flashnext-pages")) return @import("contract_fn_pages.zig").run(gpu); if (std.mem.eql(u8, cmd, "info")) return info(gpu); if (std.mem.eql(u8, cmd, "smoke")) return runtime_tests.smoke(gpu); @@ -386,6 +400,10 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "exlr8-dense")) return exlr8_dense_tests.run(gpu, if (rest.len > 0) rest[0] else null); if (std.mem.eql(u8, cmd, "exlr8-bench")) return exlr8_tests.bench(gpu, try arg(rest, 0), if (rest.len > 1) rest[1] else "home", if (rest.len > 2) try std.fmt.parseInt(u16, rest[2], 10) else 40, if (rest.len > 3) try std.fmt.parseInt(usize, rest[3], 10) else 4); if (std.mem.eql(u8, cmd, "exl3")) return exl3_tests.run(gpu, if (rest.len > 0) rest[0] else "none"); + if (std.mem.eql(u8, cmd, "exl3-pack")) return exl3_pack_tests.run(gpu, rest); + if (std.mem.eql(u8, cmd, "exl3-ple2")) return exl3_ple2_tests.run(gpu, rest); + if (std.mem.eql(u8, cmd, "exl3-pool")) return exl3_pool_tests.run(gpu, rest); + if (std.mem.eql(u8, cmd, "exl3-prefill")) return exl3_prefill_tests.run(gpu, rest); if (std.mem.eql(u8, cmd, "exl3-tiles")) return exl3_candidate_tests.run(gpu, if (rest.len > 0) rest[0] else null); if (std.mem.eql(u8, cmd, "exl3-tp2")) return exl3_tp2_tests.run(gpu, try arg(rest, 0), try arg(rest, 1), if (rest.len > 2) rest[2] else "none"); if (std.mem.eql(u8, cmd, "sample")) return sample_tests.draws(gpu); From bd23ac24e4d418219e81841a48a87263eb1875c7 Mon Sep 17 00:00:00 2001 From: BlivionIaG <15323295+BlivionIaG@users.noreply.github.com> Date: Wed, 28 Oct 2026 16:12:38 +0100 Subject: [PATCH 2/3] EXL3 Flash Next: expert-parallel TP and MTP drafts Ranks split the experts by contiguous runs of leaves (the routed ids by eights, the shared expert in the last leaf), the routers are replicated, and MTP drafts run on every rank. The MoE branch has one canonical reduction order at any tensor-parallel degree: each leaf is an fp32 fma chain over the row's slots in pick order starting at +0, and the 64 leaves are added in a fixed balanced binary tree. One rank evaluates the whole tree. A rank of d computes the maximal complete subtrees of its run (one for a power-of-two d), one fixed-size all-gather of those fp32 [rows, 2560] partials follows, and every rank finishes the same tree, so TP=1 and TP=d write the same bits. Two ranks exchange 10 KB a row a layer. A host test emulates d = 1 to 8 against the whole tree and against the definition without shortcuts, and mutations of the order, a slot, a leaf and the tree shape must change bits. --- docs/kernel-contract.md | 4 +- docs/multi-rank-receipts.md | 3 +- tools/exl3_serve.sh | 13 - tools/exl3_serve_tp2.sh | 16 + zig/build/cuda_kernels.zig | 1 + zig/kernels/cuda/flashnext_batch.cu | 73 +++ .../cuda/flashnext_exl3/flashnext_exl3_tp.cu | 9 + .../core/contract/flashnext_exl3_pending.zig | 8 +- zig/src/cuda/kernels.zig | 1 + zig/src/families/flashnext/cuda_batch.zig | 3 + .../families/flashnext/cuda_batch_head.zig | 2 +- .../families/flashnext/cuda_batch_kernels.zig | 4 + zig/src/families/flashnext/cuda_batch_moe.zig | 35 +- .../families/flashnext/cuda_batch_state.zig | 26 +- .../families/flashnext/cuda_collective.zig | 20 + zig/src/families/flashnext/cuda_lanes.zig | 22 +- .../families/flashnext/cuda_lanes_bytes.zig | 2 +- .../families/flashnext/cuda_lanes_load.zig | 15 +- .../families/flashnext/cuda_lanes_snap.zig | 8 +- zig/src/families/flashnext/cuda_launch.zig | 8 +- zig/src/families/flashnext/cuda_mtp.zig | 8 +- zig/src/families/flashnext/cuda_mtp_load.zig | 1 + zig/src/families/flashnext/cuda_mtp_ops.zig | 4 +- .../families/flashnext/cuda_mtp_ops_test.zig | 75 +++ zig/src/families/flashnext/cuda_prompt.zig | 5 +- .../flashnext/cuda_prompt_telemetry.zig | 69 ++- .../families/flashnext/cuda_rank_engine.zig | 33 +- .../flashnext/cuda_scratch_layout_test.zig | 45 +- .../flashnext/exl3/kv/cuda_exl3_kv.zig | 31 +- .../flashnext/exl3/model/cuda_exl3.zig | 3 + .../flashnext/exl3/model/cuda_exl3_arch.zig | 2 +- .../flashnext/exl3/model/cuda_exl3_pass.zig | 71 ++- .../exl3/model/cuda_exl3_telemetry.zig | 12 +- .../exl3/model/cuda_exl3_telemetry_test.zig | 4 +- .../families/flashnext/exl3/model/exl3.zig | 3 + .../families/flashnext/exl3/model/exl3_hc.zig | 113 +++- .../flashnext/exl3/model/exl3_telemetry.zig | 2 +- .../exl3/offload/cuda_exl3_frequency.zig | 30 +- .../offload/cuda_exl3_frequency_state.zig | 115 +++- .../exl3/offload/cuda_exl3_frequency_test.zig | 65 ++- .../flashnext/exl3/offload/cuda_exl3_pool.zig | 4 +- .../flashnext/exl3/pack/cuda_exl3_config.zig | 13 + .../flashnext/exl3/pack/cuda_exl3_load.zig | 32 +- .../flashnext/exl3/pack/cuda_exl3_memory.zig | 54 +- .../flashnext/exl3/pack/exl3_place.zig | 59 ++- .../flashnext/exl3/pack/exl3_shape.zig | 6 +- .../flashnext/exl3/tp/cuda_exl3_mtp.zig | 266 ++++++++++ .../flashnext/exl3/tp/cuda_exl3_tp.zig | 500 ++++++++++++++++++ .../families/flashnext/exl3/tp/exl3_canon.zig | 243 +++++++++ .../families/flashnext/exl3/tp/exl3_tp.zig | 148 ++++++ zig/src/native/cuda.zig | 10 +- .../native/cuda_flashnext_refusal_test.zig | 15 + zig/src/native/cuda_refusal.zig | 23 +- .../cuda/contract_flashnext_exl3_suite.zig | 26 + zig/tests/cuda/contract_flashnext_exl3_tp.zig | 138 +++++ zig/tests/cuda/contract_limits.zig | 2 +- zig/tests/cuda/contract_tests.zig | 2 + .../model/flashnext_exl3_engine_tests.zig | 33 +- .../model/flashnext_exl3_names_tests.zig | 3 +- .../model/flashnext_exl3_rig.zig | 4 +- .../offload/exl3_frequency_tests.zig | 82 +++ .../flashnext_exl3/pack/exl3_pack_tests.zig | 55 ++ .../flashnext_exl3/tp/exl3_draft_tests.zig | 159 ++++++ zig/tests/cuda/flashnext_exl3/tp/exl3_tp.zig | 1 + .../tp/flashnext_exl3_tp_engine_tests.zig | 22 + .../tp/flashnext_exl3_tp_tests.zig | 452 ++++++++++++++++ .../tp/flashnext_exl3_tpe_decode_tests.zig | 212 ++++++++ .../tp/flashnext_exl3_tpe_engine.zig | 214 ++++++++ zig/tests/cuda/flashnext_hc_tests.zig | 46 ++ zig/tests/cuda/flashnext_model_rank_tests.zig | 6 +- zig/tests/cuda/flashnext_window_tests.zig | 6 +- zig/tests/cuda/main.zig | 10 + 72 files changed, 3570 insertions(+), 240 deletions(-) delete mode 100755 tools/exl3_serve.sh create mode 100755 tools/exl3_serve_tp2.sh create mode 100644 zig/kernels/cuda/flashnext_exl3/flashnext_exl3_tp.cu create mode 100644 zig/src/families/flashnext/cuda_mtp_ops_test.zig create mode 100644 zig/src/families/flashnext/exl3/tp/cuda_exl3_mtp.zig create mode 100644 zig/src/families/flashnext/exl3/tp/cuda_exl3_tp.zig create mode 100644 zig/src/families/flashnext/exl3/tp/exl3_canon.zig create mode 100644 zig/src/families/flashnext/exl3/tp/exl3_tp.zig create mode 100644 zig/tests/cuda/contract_flashnext_exl3_suite.zig create mode 100644 zig/tests/cuda/contract_flashnext_exl3_tp.zig create mode 100644 zig/tests/cuda/flashnext_exl3/offload/exl3_frequency_tests.zig create mode 100644 zig/tests/cuda/flashnext_exl3/tp/exl3_draft_tests.zig create mode 120000 zig/tests/cuda/flashnext_exl3/tp/exl3_tp.zig create mode 100644 zig/tests/cuda/flashnext_exl3/tp/flashnext_exl3_tp_engine_tests.zig create mode 100644 zig/tests/cuda/flashnext_exl3/tp/flashnext_exl3_tp_tests.zig create mode 100644 zig/tests/cuda/flashnext_exl3/tp/flashnext_exl3_tpe_decode_tests.zig create mode 100644 zig/tests/cuda/flashnext_exl3/tp/flashnext_exl3_tpe_engine.zig diff --git a/docs/kernel-contract.md b/docs/kernel-contract.md index 8f3a676fc4..7323e4b44e 100644 --- a/docs/kernel-contract.md +++ b/docs/kernel-contract.md @@ -215,8 +215,8 @@ shared memory, the device's limits and the planted controls are checked the same A restored 13.0 receipt can arrive after the later compiler sets. Such records mark `later_pending` so 13.0 and 13.1 use their base usage, while 13.2 and later keep their independently measured or pending selection. A missing later receipt remains pending; a 13.0 measurement never substitutes for it. -Flash Next's EXL3 launch (`flashnext_exl3_pending.zig`) carries only an nvcc 13.4 receipt: it stays pending under -13.0 to 13.3 and takes its 13.4 usage from `cuda_134`, which `pending_132.lookup` serves for that release. +Flash Next's 4 EXL3 launches (`flashnext_exl3_pending.zig`) carry only an nvcc 13.4 receipt: they stay pending under +13.0 to 13.3 and take their 13.4 usage from `cuda_134`, which `pending_132.lookup` serves for that release. Gemma 4 CUDA tower QKV norms and nine-term pooling register FP32 sums and BF16 outputs. Their contract keeps the exact FP32 mirror and adds independent FP64 ranges at the BF16 nodes. diff --git a/docs/multi-rank-receipts.md b/docs/multi-rank-receipts.md index a761377670..f48836b451 100644 --- a/docs/multi-rank-receipts.md +++ b/docs/multi-rank-receipts.md @@ -13,7 +13,8 @@ Each path places its ranks its own way: | Native ROCm build | `--tp`, `--rank`, `--master` | [Native ROCm placement](../zig/src/native/hip.zig) | | The `fabric/tp` transport's harness | `TF_TP_WORLD`, `TF_TP_RANK`, `TF_TP_MASTER` | [The transport's README](../zig/src/fabric/tp/README.md) | -The native CUDA build serves neither `--tp` nor `--rank`: a family's own variables place its ranks. `tf-tp-test` is +The native CUDA build serves neither `--tp` nor `--rank` for a checkpoint other than a Flash Next EXL3 pack, which they +place across two GPUs (`--tp 2 --rank N --master HOST`): a family's own variables place its ranks. `tf-tp-test` is the transport's harness rather than a server. `tensorfold benchmark` accepts no group placement flags and records 1 for a managed run. Its child inherits the parent's rank environment, so start managed runs from a single-rank configuration. Measure a group as an attached server: start the group, then point the client at the rank that serves. diff --git a/tools/exl3_serve.sh b/tools/exl3_serve.sh deleted file mode 100755 index 82c68f254a..0000000000 --- a/tools/exl3_serve.sh +++ /dev/null @@ -1,13 +0,0 @@ -#!/bin/bash -# Serve the EXL3 Qwen3.8-Flash-Next pack on one 16 GB CUDA GPU (sm_120): the routed experts sit in pinned host memory and -# each layer keeps LINES of them on the device (default 95). Plain decode (no MTP head); 0.5 GiB device reserve. -# Usage: tools/exl3_serve.sh MODEL_DIR PORT [extra serve args]; CUDA_VISIBLE_DEVICES picks the card. -# Environment: LINES, RESERVE_MB (planner reserve), MEMORY_RESERVE_GIB, CONTEXT, PARALLEL, TENSORFOLD_BIN. -set -euo pipefail -model=$1 port=$2 -shift 2 -bin=${TENSORFOLD_BIN:-./zig-out/native/bin/tensorfold-native} -export TENSORFOLD_CHECKPOINT_SHA256=${TENSORFOLD_CHECKPOINT_SHA256:-$(cat "$model/model.safetensors.index.json" "$model/config.json" | sha256sum | cut -c1-64)} -export TF_FLASHNEXT_EXPERT_LINES=${LINES:-95} TF_FLASHNEXT_EXPERT_RESERVE_MB=${RESERVE_MB:-256} TENSORFOLD_MEMORY_RESERVE_GIB=${MEMORY_RESERVE_GIB:-0.5} -exec "$bin" serve "$model" --backend cuda --context "${CONTEXT:-32768}" --port "$port" --no-update-check \ - --parallel "${PARALLEL:-1}" --prompt-cache-gib "${CACHE_GIB:-0}" --no-drafts "$@" diff --git a/tools/exl3_serve_tp2.sh b/tools/exl3_serve_tp2.sh new file mode 100755 index 0000000000..e051f0d5ac --- /dev/null +++ b/tools/exl3_serve_tp2.sh @@ -0,0 +1,16 @@ +#!/bin/bash +# Launch one TP=2 rank of the EXL3 Qwen3.8-Flash-Next server on two RTX 5060 Ti (16 GB), drafts on at every stream count. +# Usage: tools/exl3_serve_tp2.sh MODEL_DIR RANK PORT MASTER_PORT [extra serve args]; start rank 1 first, then rank 0. +# The MTP head is loaded by default (rank 0 holds it); `--no-drafts` in the extra arguments serves without it. +# Environment: LINES (expert lines a layer, default 38), RESERVE_MB (planner reserve), MEMORY_RESERVE_GIB, CONTEXT, +# PARALLEL (streams, default 2), TENSORFOLD_BIN. +set -euo pipefail +model=$1 rank=$2 port=$3 master=$4 +shift 4 +bin=${TENSORFOLD_BIN:-./zig-out/native/bin/tensorfold-native} +export TENSORFOLD_CHECKPOINT_SHA256=${TENSORFOLD_CHECKPOINT_SHA256:-$(cat "$model/model.safetensors.index.json" "$model/config.json" | sha256sum | cut -c1-64)} +export TF_FLASHNEXT_EXPERT_LINES=${LINES:-38} +[ -n "${RESERVE_MB:-}" ] && export TF_FLASHNEXT_EXPERT_RESERVE_MB=$RESERVE_MB +[ -n "${MEMORY_RESERVE_GIB:-}" ] && export TENSORFOLD_MEMORY_RESERVE_GIB=$MEMORY_RESERVE_GIB +CUDA_VISIBLE_DEVICES=$rank exec "$bin" serve "$model" --backend cuda --context "${CONTEXT:-32768}" --tp 2 --rank "$rank" \ + --master 127.0.0.1 --master-port "$master" --port "$port" --no-update-check --parallel "${PARALLEL:-2}" --prompt-cache-gib 0 "$@" diff --git a/zig/build/cuda_kernels.zig b/zig/build/cuda_kernels.zig index 3dc4c2a816..4faa679777 100644 --- a/zig/build/cuda_kernels.zig +++ b/zig/build/cuda_kernels.zig @@ -121,6 +121,7 @@ pub const kernels = [_]Kernel{ .{ .name = "drama3_fp8", .flags = &.{"-O3"} }, .{ .name = "drama3_codec", .flags = &.{"-O3"} }, .{ .name = "flashnext_exl3", .src = "flashnext_exl3/flashnext_exl3", .flags = glue }, // ours: Flash Next's EXL3 glue (an fp16 row split exactly into two bf16) + .{ .name = "flashnext_exl3_tp", .src = "flashnext_exl3/flashnext_exl3_tp", .flags = glue }, // ours: Flash Next's EXL3 expert-parallel pick mask }; /// torch.utils.cpp_extension's own nvcc flags (torch 2.13): C++20 and which half/bf16 operators the headers define. diff --git a/zig/kernels/cuda/flashnext_batch.cu b/zig/kernels/cuda/flashnext_batch.cu index f7064ed911..b314191803 100644 --- a/zig/kernels/cuda/flashnext_batch.cu +++ b/zig/kernels/cuda/flashnext_batch.cu @@ -326,6 +326,79 @@ extern "C" __global__ void fn_moe_sum(const float* __restrict__ y, const float* out[static_cast(r) * d + col] = acc; } +// EXL3 canonical expert reduction (receipts canon-reduction.md): the value of the leaf subtree [start, start + 2^log). +extern "C" __global__ void fn_exl3_leaf_tree(const float* __restrict__ y, const float* __restrict__ wts, const int* __restrict__ pick, + float* __restrict__ out, int d, int slots, int start, int log) { + const int r = blockIdx.x, col = blockIdx.y * blockDim.x + threadIdx.x; + if (col >= d) return; + const int end = start + (1 << log); + int leaf[11]; + float val[11]; + int n = 0; + for (int k = 0; k < slots; ++k) { + const int p = pick[r * slots + k]; + const int l = p >= 512 ? 63 : (p >> 3); + if (l < start || l >= end) continue; + int j = 0; + while (j < n && leaf[j] != l) ++j; + if (j == n) { + leaf[n] = l; + val[n] = 0.0f; + ++n; + } + val[j] = __fmaf_rn(y[(static_cast(r) * slots + k) * d + col], wts[r * slots + k], val[j]); + } + for (int i = 1; i < n; ++i) + for (int j = i; j > 0 && leaf[j - 1] > leaf[j]; --j) { + const int tl = leaf[j - 1]; + leaf[j - 1] = leaf[j]; + leaf[j] = tl; + const float tv = val[j - 1]; + val[j - 1] = val[j]; + val[j] = tv; + } + for (int i = 0; i < n; ++i) leaf[i] -= start; + for (int lev = 0; lev < log; ++lev) { + int m = 0; + for (int q = 0; q < n;) { + const int parent = leaf[q] >> 1; + if (q + 1 < n && parent == (leaf[q + 1] >> 1)) { + val[m] = __fadd_rn(val[q], val[q + 1]); + q += 2; + } else { + val[m] = val[q]; + q += 1; + } + leaf[m] = parent; + ++m; + } + n = m; + } + out[static_cast(r) * d + col] = n == 0 ? 0.0f : val[0]; +} + +// The finish: `pieces` partial subtrees (meta: start leaf, log, source piece), in leaf order, merged left to right. +extern "C" __global__ void fn_exl3_leaf_finish(const float* __restrict__ g, float* __restrict__ out, const int* __restrict__ meta, + int pieces, int d, long long stride) { + const int r = blockIdx.x, col = blockIdx.y * blockDim.x + threadIdx.x; + if (col >= d) return; + int st[32], lg[32]; + float sv[32]; + int top = 0; + for (int i = 0; i < pieces; ++i) { + st[top] = meta[3 * i]; + lg[top] = meta[3 * i + 1]; + sv[top] = g[static_cast(meta[3 * i + 2]) * stride + static_cast(r) * d + col]; + ++top; + while (top >= 2 && lg[top - 1] == lg[top - 2] && st[top - 2] % (2 << lg[top - 1]) == 0) { + sv[top - 2] = __fadd_rn(sv[top - 2], sv[top - 1]); + lg[top - 2] += 1; + --top; + } + } + out[static_cast(r) * d + col] = sv[0]; +} + // dst [rows, width] (rows dst_pitch apart) = the columns src [rows, width] (src_pitch apart), 16-bit values. extern "C" __global__ void fn_copy_cols(const uint16_t* __restrict__ src, long long src_pitch, uint16_t* __restrict__ dst, long long dst_pitch, int width) { diff --git a/zig/kernels/cuda/flashnext_exl3/flashnext_exl3_tp.cu b/zig/kernels/cuda/flashnext_exl3/flashnext_exl3_tp.cu new file mode 100644 index 0000000000..02701e0aeb --- /dev/null +++ b/zig/kernels/cuda/flashnext_exl3/flashnext_exl3_tp.cu @@ -0,0 +1,9 @@ +// The Flash Next family's EXL3 expert-parallel glue (cuda_exl3_tp.zig): a rank's masked picks. + +// The other rank's slots (outside [LO, HI)) become -1, which the experts' kernels skip; slots keep their positions. +extern "C" __global__ void fn_exl3_tp_mask(const int* __restrict__ pick, int* __restrict__ mpick, int n, int lo, int hi) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= n) return; + const int p = pick[i]; + mpick[i] = (p >= lo && p < hi) ? p : -1; +} diff --git a/zig/src/core/contract/flashnext_exl3_pending.zig b/zig/src/core/contract/flashnext_exl3_pending.zig index e8d0f82fec..4178c20231 100644 --- a/zig/src/core/contract/flashnext_exl3_pending.zig +++ b/zig/src/core/contract/flashnext_exl3_pending.zig @@ -1,4 +1,4 @@ -//! Flash Next EXL3 launches (the fp16 row split); nvcc 13.4 measured, older pending. +//! Flash Next EXL3 launches (the fp16 row split, the expert-parallel pick mask); nvcc 13.4 measured, older pending. const std = @import("std"); const limits = @import("limits.zig"); const registry = @import("registry.zig"); @@ -6,11 +6,17 @@ const registry = @import("registry.zig"); /// Launches with no nvcc 13.0 to 13.3 receipt yet: the geometry is recorded and usage stays pending until measured. pub const records = [_]registry.CudaRecord{ .{ .launch = .{ .name = "flashnext_exl3 fn_exl3_split_f16", .symbol = "fn_exl3_split_f16", .image = "flashnext_exl3", .threads = 256 }, .usage = registry.none }, + .{ .launch = .{ .name = "flashnext_batch fn_exl3_leaf_tree", .symbol = "fn_exl3_leaf_tree", .image = "flashnext_batch", .threads = 256 }, .usage = registry.none }, + .{ .launch = .{ .name = "flashnext_batch fn_exl3_leaf_finish", .symbol = "fn_exl3_leaf_finish", .image = "flashnext_batch", .threads = 256 }, .usage = registry.none }, + .{ .launch = .{ .name = "flashnext_exl3_tp fn_exl3_tp_mask", .symbol = "fn_exl3_tp_mask", .image = "flashnext_exl3_tp", .threads = 256 }, .usage = registry.none }, }; /// nvcc 13.4's measured registers, static shared bytes and launch bounds for sm_80, sm_86, sm_89, sm_120, sm_121. pub const cuda_134 = [_]registry.Override{ .{ .name = "flashnext_exl3 fn_exl3_split_f16", .usage = .{ .{ .regs = 10, .static = 0, .bound = 0 }, .{ .regs = 10, .static = 0, .bound = 0 }, .{ .regs = 10, .static = 0, .bound = 0 }, .{ .regs = 12, .static = 0, .bound = 0 }, .{ .regs = 12, .static = 0, .bound = 0 } } }, + .{ .name = "flashnext_batch fn_exl3_leaf_tree", .usage = .{ .{ .regs = 32, .static = 0, .bound = 0 }, .{ .regs = 34, .static = 0, .bound = 0 }, .{ .regs = 34, .static = 0, .bound = 0 }, .{ .regs = 24, .static = 0, .bound = 0 }, .{ .regs = 24, .static = 0, .bound = 0 } } }, + .{ .name = "flashnext_batch fn_exl3_leaf_finish", .usage = .{ .{ .regs = 23, .static = 0, .bound = 0 }, .{ .regs = 22, .static = 0, .bound = 0 }, .{ .regs = 22, .static = 0, .bound = 0 }, .{ .regs = 19, .static = 0, .bound = 0 }, .{ .regs = 19, .static = 0, .bound = 0 } } }, + .{ .name = "flashnext_exl3_tp fn_exl3_tp_mask", .usage = .{ .{ .regs = 10, .static = 0, .bound = 0 }, .{ .regs = 10, .static = 0, .bound = 0 }, .{ .regs = 10, .static = 0, .bound = 0 }, .{ .regs = 10, .static = 0, .bound = 0 }, .{ .regs = 10, .static = 0, .bound = 0 } } }, }; test "Flash Next EXL3 launches: pending below nvcc 13.4, measured 13.4 resources fit every SM launch limit" { diff --git a/zig/src/cuda/kernels.zig b/zig/src/cuda/kernels.zig index 0594bdb4ce..9e1402f1c0 100644 --- a/zig/src/cuda/kernels.zig +++ b/zig/src/cuda/kernels.zig @@ -19,6 +19,7 @@ pub const qwen4_projection: []const u8 = if (available) &Blob("fatbin_qwen4_proj pub const gdn: []const u8 = if (available) &Blob("fatbin_gdn").bytes else &.{}; pub const flashnext_gdn: []const u8 = if (available) &Blob("fatbin_flashnext_gdn").bytes else &.{}; pub const flashnext_exl3: []const u8 = if (available) &Blob("fatbin_flashnext_exl3").bytes else &.{}; +pub const flashnext_exl3_tp: []const u8 = if (available) &Blob("fatbin_flashnext_exl3_tp").bytes else &.{}; pub const flashnext_batch: []const u8 = if (available) &Blob("fatbin_flashnext_batch").bytes else &.{}; pub const flashnext_mtp: []const u8 = if (available) &Blob("fatbin_flashnext_mtp").bytes else &.{}; pub const probe: []const u8 = if (available) &Blob("fatbin_probe").bytes else &.{}; diff --git a/zig/src/families/flashnext/cuda_batch.zig b/zig/src/families/flashnext/cuda_batch.zig index 8ebdea1f63..c387d05d6b 100644 --- a/zig/src/families/flashnext/cuda_batch.zig +++ b/zig/src/families/flashnext/cuda_batch.zig @@ -117,6 +117,8 @@ pub fn Pass(comptime Tri: type) type { prepared: bool = false, /// The pending write-back's branch sits in fp32 in moe_f32 (something added to it, or it is kept). pending_f32: bool = false, + /// moe_f32 holds the EXL3 canonical grouped branch (moeAdds keeps it instead of re-summing the slots). + canon_f32: bool = false, /// This layer's out projection leaves part's other slices alone; its writeback adds the zero slice once. zero_slices: bool = false, @@ -501,6 +503,7 @@ pub fn Pass(comptime Tri: type) type { self.pos0 = pos0; self.pending = null; self.pending_f32 = false; + self.canon_f32 = false; // the gather writes the scratch's pinned rows, so the last pass's uploads must have read them first if (!self.prepared) try s.ready(); const emb = s.hostEmb(n); diff --git a/zig/src/families/flashnext/cuda_batch_head.zig b/zig/src/families/flashnext/cuda_batch_head.zig index c33fe9bc1a..0b3518c501 100644 --- a/zig/src/families/flashnext/cuda_batch_head.zig +++ b/zig/src/families/flashnext/cuda_batch_head.zig @@ -81,7 +81,7 @@ pub fn projectHead(self: anytype, out: u64, n: usize) !void { _ = try self.lane(s.mixed, s.xsm, self.m.head, out, s.dummy, n, 1); return; } - if (self.m.head.n != vocab / 2 or n > s.cap) return error.UnexpectedTensor; + if (self.m.head.n != vocab / 2 or n > s.headRows()) return error.UnexpectedTensor; const c = self.m.collective orelse return error.TensorParallelUnavailable; _ = try self.lane(s.mixed, s.xsm, self.m.head, s.head_local, s.dummy, n, 1); for (0..n) |r| try c.gather(self.stream, s.head_local + r * vocab, out + r * vocab * 2, vocab / 2); diff --git a/zig/src/families/flashnext/cuda_batch_kernels.zig b/zig/src/families/flashnext/cuda_batch_kernels.zig index 1443cf1398..3bc3765785 100644 --- a/zig/src/families/flashnext/cuda_batch_kernels.zig +++ b/zig/src/families/flashnext/cuda_batch_kernels.zig @@ -30,6 +30,8 @@ pub const Kernels = struct { gnorm: cuda.Function, sum: cuda.Function, moe_sum: cuda.Function, + leaf_tree: cuda.Function, + leaf_finish: cuda.Function, copy_cols: cuda.Function, shift: cuda.Function, /// Each sparse row's top blocks (attention._select / _select_tiles), native. @@ -66,6 +68,8 @@ pub const Kernels = struct { errdefer if (k.expert_tile) |*tile| tile.unload(); k.sum = try k.batch.function("fn_sum_slices"); k.moe_sum = try k.batch.function("fn_moe_sum"); + k.leaf_tree = try k.batch.function("fn_exl3_leaf_tree"); + k.leaf_finish = try k.batch.function("fn_exl3_leaf_finish"); k.copy_cols = try k.batch.function("fn_copy_cols"); k.shift = try k.batch.function("fn_shift_window"); k.select = try k.batch.function("fn_qsa_select"); diff --git a/zig/src/families/flashnext/cuda_batch_moe.zig b/zig/src/families/flashnext/cuda_batch_moe.zig index a595380b0a..505bf69bc6 100644 --- a/zig/src/families/flashnext/cuda_batch_moe.zig +++ b/zig/src/families/flashnext/cuda_batch_moe.zig @@ -8,6 +8,8 @@ const slots = wt.slots; /// After a layer's experts, when something adds to or keeps its output: the slots' fp32 sum, adds, kept branch. pub fn moeAdds(self: anytype, layer: usize, keep: LayerKeep) !void { + const canon = self.canon_f32; + self.canon_f32 = false; const add = self.addAt(layer); if (add == null and keep.moe_branch == 0) return; const s = self.s; @@ -18,8 +20,39 @@ pub fn moeAdds(self: anytype, layer: usize, keep: LayerKeep) !void { a.add(s.moe_f32); a.add(batch.cint(dims)); a.add(batch.cint(slots)); - try cuda.launch.launch(self.k.moe_sum, .{ .grid = .{ .x = @intCast(n), .y = dims / 256 }, .block = .{ .x = 256 } }, self.stream, &a); + // the EXL3 canonical branch is already in moe_f32 + if (!canon) try cuda.launch.launch(self.k.moe_sum, .{ .grid = .{ .x = @intCast(n), .y = dims / 256 }, .block = .{ .x = 256 } }, self.stream, &a); if (add) |m| try m.add(m.ctx, layer, s.mixed, s.picks, s.wts, s.moe_f32, n); if (keep.moe_branch != 0) try self.sumSlices(s.moe_f32, keep.moe_branch, true, n * dims, 1); self.pending_f32 = true; } + +/// The EXL3 canonical value of the rows over the leaf subtree [start, start + 2^log) into `out` [rows, dims] fp32. +pub fn leafTree(self: anytype, out: u64, start: usize, log: usize) !void { + const s = self.s; + var a: cuda.Args = .{}; + a.add(s.y); + a.add(s.wts); + a.add(s.picks); + a.add(out); + for ([_]usize{ dims, slots, start, log }) |v| a.add(batch.cint(v)); + try cuda.launch.launch(self.k.leaf_tree, .{ .grid = .{ .x = @intCast(self.rows), .y = dims / 256 }, .block = .{ .x = 256 } }, self.stream, &a); +} + +/// The pieces in `g` (piece i at i * stride floats) merged by `meta` into `out`. +pub fn leafFinish(self: anytype, g: u64, out: u64, meta: u64, pieces: usize, stride: usize) !void { + var a: cuda.Args = .{}; + a.add(g); + a.add(out); + a.add(meta); + a.add(batch.cint(pieces)); + a.add(batch.cint(dims)); + a.add(@as(i64, @intCast(stride))); + try cuda.launch.launch(self.k.leaf_finish, .{ .grid = .{ .x = @intCast(self.rows), .y = dims / 256 }, .block = .{ .x = 256 } }, self.stream, &a); +} + +/// moe_f32 holds the canonical branch: the write-back reads it and moeAdds keeps it. +pub fn kept(self: anytype) void { + self.canon_f32 = true; + self.pending_f32 = true; +} diff --git a/zig/src/families/flashnext/cuda_batch_state.zig b/zig/src/families/flashnext/cuda_batch_state.zig index 14035aa5d9..b49c64539a 100644 --- a/zig/src/families/flashnext/cuda_batch_state.zig +++ b/zig/src/families/flashnext/cuda_batch_state.zig @@ -197,6 +197,8 @@ pub const ScratchFormat = enum { classic, exl3 }; pub const Scratch = struct { cap: usize, format: ScratchFormat = .classic, + /// The allocation's own head bound; another engine cannot change its launch admission. + head_local_rows: usize = std.math.maxInt(usize), bufs: std.ArrayList(cuda.DeviceBuffer) = .empty, ids: u64 = 0, pos: u64 = 0, @@ -325,7 +327,15 @@ pub const Scratch = struct { "ple_vals", "ple_part", "ple_gated", "ple_pss", "ple_nrow", "logits", "head_local", "inv", "dummy", }; + pub fn headRows(s: *const Scratch) usize { + return @min(s.cap, s.head_local_rows); + } + pub fn allocationSizes(cap: usize, format: ScratchFormat) [64]?usize { + return allocationSizesForHead(cap, format, cap); + } + + pub fn allocationSizesForHead(cap: usize, format: ScratchFormat, head_rows: usize) [64]?usize { const pairs = cap * slots; const nblk = cdiv(pairs, plan_small); const items = cdiv(pairs, tile) + experts_n + 1; @@ -391,7 +401,7 @@ pub const Scratch = struct { cap * streams * 4, cap * wide * 2, vocab * 2, - if (format == .classic) cap * vocab else null, + @min(cap, head_rows) * vocab, half * 4, 4096, }; @@ -402,8 +412,12 @@ pub const Scratch = struct { } pub fn bytesForFormat(cap: usize, format: ScratchFormat) usize { + return bytesForHead(cap, format, cap); + } + + pub fn bytesForHead(cap: usize, format: ScratchFormat, head_rows: usize) usize { var total: usize = 0; - for (allocationSizes(cap, format)) |optional| if (optional) |len| { + for (allocationSizesForHead(cap, format, head_rows)) |optional| if (optional) |len| { total += @max(len, 256); }; return total; @@ -414,12 +428,16 @@ pub const Scratch = struct { } pub fn allocForFormat(gpa: std.mem.Allocator, d: *const cuda.Driver, stream: cuda.Stream, cap: usize, format: ScratchFormat) !Scratch { - var s: Scratch = .{ .cap = cap, .format = format }; + return allocForHead(gpa, d, stream, cap, format, cap); + } + + pub fn allocForHead(gpa: std.mem.Allocator, d: *const cuda.Driver, stream: cuda.Stream, cap: usize, format: ScratchFormat, head_rows: usize) !Scratch { + var s: Scratch = .{ .cap = cap, .format = format, .head_local_rows = head_rows }; errdefer { cuda.completion.drain(d, stream, "Flash Next scratch work did not drain before rollback"); s.free(gpa); } - const sizes = allocationSizes(cap, format); + const sizes = allocationSizesForHead(cap, format, head_rows); // a row of a decode round or a block of att_rows rows takes `pos`'s 16 bytes var inv: [half]f32 = undefined; for (&inv, 0..) |*o, i| o.* = @floatCast(std.math.pow(f64, 10_000_000.0, -@as(f64, @floatFromInt(i)) / @as(f64, half))); diff --git a/zig/src/families/flashnext/cuda_collective.zig b/zig/src/families/flashnext/cuda_collective.zig index fd6d930836..77bd17b386 100644 --- a/zig/src/families/flashnext/cuda_collective.zig +++ b/zig/src/families/flashnext/cuda_collective.zig @@ -17,6 +17,8 @@ pub const Collective = struct { link: wire.Link, timeout_ms: u64, failed: bool = false, + /// How long a wait polls before it sleeps between polls (0: sleeps from the first poll); an EXL3 engine sets it. + spin_ns: u64 = 0, /// `link` has completed the source/settings handshake; it stays owned by the caller. pub fn init(link: wire.Link, rank: usize, timeout_ms: u64) !Collective { @@ -63,6 +65,7 @@ pub const Collective = struct { fn wait(c: *Collective, stream: ?cuda.Stream, until: u64) !void { errdefer c.abort(); + var spin_until: u64 = 0; while (true) { try c.link.checkPeer(); if (wire.now() >= until) return error.RankTimeout; @@ -74,6 +77,13 @@ pub const Collective = struct { if (try s.done()) return; } else return; } + if (c.spin_ns != 0) { + if (spin_until == 0) spin_until = wire.now() + c.spin_ns; + if (wire.now() < spin_until) { + std.atomic.spinLoopHint(); + continue; + } + } wire.pause(); } } @@ -94,6 +104,16 @@ pub const Collective = struct { try c.wait(null, wire.deadline(c.timeout_ms)); } + /// Every rank's `count` fp32 values side by side in `out` (rank r at r * count); `local` may sit there already. + pub fn gatherF32(c: *Collective, stream: cuda.Stream, local: u64, out: u64, count: usize) !void { + if (c.failed or c.comm == null) return error.RankDisconnected; + if (count == 0) return; + errdefer c.abort(); + try c.link.checkPeer(); + try c.submitted(c.lib.api.ncclAllGather(local, out, count, .f32, c.comm, stream.handle)); + try c.wait(null, wire.deadline(c.timeout_ms)); + } + pub fn broadcast(c: *Collective, stream: cuda.Stream, values: u64, bytes: usize) !void { if (c.failed or c.comm == null) return error.RankDisconnected; errdefer c.abort(); diff --git a/zig/src/families/flashnext/cuda_lanes.zig b/zig/src/families/flashnext/cuda_lanes.zig index fc28b9dcb3..99a57f0c5f 100644 --- a/zig/src/families/flashnext/cuda_lanes.zig +++ b/zig/src/families/flashnext/cuda_lanes.zig @@ -9,6 +9,8 @@ const mtp_ops = @import("cuda_mtp_ops.zig"); const lanes_draft = @import("cuda_lanes_draft.zig"); const snap = @import("cuda_snap.zig"); const ranks = @import("cuda_rank_engine.zig"); +const config = @import("cuda_config.zig"); +const exl3_memory = @import("exl3/pack/cuda_exl3_memory.zig"); const Glue = @import("cuda_glue.zig").Glue; const living = @import("living"); const slide_mod = @import("cuda_slide.zig"); @@ -27,6 +29,10 @@ pub fn Namespace(comptime Tri: type) type { return struct { pub const model_type = "qwen4_exp"; pub const tp_draft_snapshots = false; + /// Whether `dir` holds an EXL3 pack: --tp is served for those, and their ranks keep drafting states. + pub fn exl3_pack(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) bool { + return @import("exl3/pack/cuda_exl3_config.zig").packAt(gpa, io, dir); + } pub const formats: []const []const u8 = &.{ "mlx-q4g32", "modelopt-mixed", @import("exl3/pack/cuda_exl3_config.zig").format }; // the model's window: past 2,048 keys the block indexer picks each row's 512 best blocks of 4 (cuda_batch.zig) pub const default_context: i64 = @intCast(prompt.max_tokens); @@ -120,7 +126,8 @@ pub fn Namespace(comptime Tri: type) type { } pub fn rankRoom(ptr: *anyopaque, local: u64) !u64 { - if (of(ptr).ranks) |r| return r.room(local); + const own = of(ptr); + if (own.ranks) |r| return if (own.exl3) @import("exl3/kv/cuda_exl3_kv.zig").rankRoom(r, local, own.stream_bytes) else r.room(local); return local; } @@ -163,6 +170,9 @@ pub fn Namespace(comptime Tri: type) type { logits: cuda.DeviceBuffer, drawn: [ring]u32 = undefined, next: u64 = 0, + /// An EXL3 pack: its ranks agree their room in streams of stream_bytes and keep drafting states. + exl3: bool = false, + stream_bytes: usize = 0, /// --slide: the pocket's layout, from config.json at load. lay: living.Layout = undefined, kv: kvl.State = .{}, @@ -303,14 +313,18 @@ pub fn Namespace(comptime Tri: type) type { const saved = s.reuse.saved orelse return l; const kept: *snap.Snap = @ptrCast(@alignCast(saved)); if (kept.at != s.reuse.at or kept.at == 0 or kept.at >= ids.len) { + std.log.warn("flashnext: kept state at {d} not restored for reuse at {d} of {d} ids", .{ kept.at, s.reuse.at, ids.len }); s.reuse_failed = true; - } else if (self.ranks != null and l.draft != null) { + } else if (self.ranks != null and l.draft != null and !self.exl3) { s.reuse_failed = true; - } else if (if (self.ranks) |r| r.restore(&l.seq, kept) else snap.restore(self.stream, &l.seq, if (l.draft) |*d| d else null, kept)) |_| { + } else if (if (self.ranks) |r| r.restore(&l.seq, if (l.draft) |*d| d else null, kept) else snap.restore(self.stream, &l.seq, if (l.draft) |*d| d else null, kept)) |_| { l.filled = kept.at; s.cached = @intCast(kept.at); } else |err| switch (err) { - error.SequenceTooShort, error.NoHeadState => s.reuse_failed = true, + error.SequenceTooShort, error.NoHeadState => { + std.log.warn("flashnext: kept state at {d} not restored ({t})", .{ kept.at, err }); + s.reuse_failed = true; + }, else => return err, } return l; diff --git a/zig/src/families/flashnext/cuda_lanes_bytes.zig b/zig/src/families/flashnext/cuda_lanes_bytes.zig index 556ef895a7..04d8f9513c 100644 --- a/zig/src/families/flashnext/cuda_lanes_bytes.zig +++ b/zig/src/families/flashnext/cuda_lanes_bytes.zig @@ -15,7 +15,7 @@ pub fn deviceWeightBytes(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) u6 pub fn placementWeightBytes(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, world: usize, rank: usize, drafts: bool) !u64 { if ((world != 1 and world != 2) or rank >= world) return error.BadRank; const mode = try config.mode(gpa, io, dir); - if (mode == .exl3) return exl3_memory.placement(gpa, io, dir, world, rank); + if (mode == .exl3) return exl3_memory.placement(gpa, io, dir, world, rank, drafts); if (mode == .modelopt) { const main = try @import("cuda_rank_memory.zig").estimate(gpa, io, dir, world, rank); // Dense draft projections and vocabulary rows exceed the affine head's estimate. diff --git a/zig/src/families/flashnext/cuda_lanes_load.zig b/zig/src/families/flashnext/cuda_lanes_load.zig index 3513ebf653..12a605c570 100644 --- a/zig/src/families/flashnext/cuda_lanes_load.zig +++ b/zig/src/families/flashnext/cuda_lanes_load.zig @@ -22,7 +22,7 @@ pub fn open(comptime F: type, gpa: std.mem.Allocator, io: std.Io, ctx: *const cu _ = try @import("cuda_modelopt_threads.zig").configured(); // TF_FLASHNEXT_LOAD_THREADS, refused before any load const mode = try config.mode(gpa, io, dir); if (mode == .exl3) try @import("exl3/model/cuda_exl3_arch.zig").check(ctx.d); // sm_120 only: refused before any weight loads - if (o.tp != 1 and (mode != .modelopt or captured != null)) return error.UnsupportedQuantization; + if (o.tp != 1 and ((mode != .modelopt and mode != .exl3) or captured != null)) return error.UnsupportedQuantization; if (o.slide and (mode != .affine or o.tp != 1)) return error.UnsupportedLearningFormat; if (o.slide and captured != null) return error.TrainingOnCapturedGlue; const page_tokens = try @import("cuda_kv_lanes.zig").paging(); @@ -47,6 +47,8 @@ pub fn open(comptime F: type, gpa: std.mem.Allocator, io: std.Io, ctx: *const cu id.max_images = v.max_images; } own.ranks = try ranks.Session.init(gpa, io, o.master orelse "127.0.0.1", o.master_port, o.rank, id); + // an EXL3 pack's decode step ends its collectives within tens of microseconds: poll before sleeping + if (mode == .exl3) own.ranks.?.collective.spin_ns = 3 * std.time.ns_per_ms; } errdefer if (own.ranks) |r| { r.poison(); @@ -69,8 +71,7 @@ pub fn open(comptime F: type, gpa: std.mem.Allocator, io: std.Io, ctx: *const cu // the pass's scratch for a full chunk, once for the engine (every prompt, window, round and head pass) _ = try own.engine.reserve(@max(@min(o.context, own.engine.chunk), if (o.slide) @as(usize, @import("cuda_train.zig").pack_rows) else 0)); // the MTP head unless --no-drafts (TF_FLASHNEXT_DRAFTS and _CONFIDENCE: the chain's rule, 6 and 0.60) - if (o.drafts and mode == .exl3) std.log.info("flashnext: an EXL3 pack decodes without the MTP head", .{}); - if (o.drafts and mode != .exl3 and (o.tp == 1 or o.rank == 0)) { + if (o.drafts and (o.tp == 1 or o.rank == 0)) { const head_before = cuda.usage(false).device; own.head = try mtp.Head.init(gpa, io, driver, &own.stream, dir); persistent_bytes = try std.math.add(u64, persistent_bytes, cuda.usage(false).device -| head_before); @@ -92,7 +93,7 @@ pub fn open(comptime F: type, gpa: std.mem.Allocator, io: std.Io, ctx: *const cu var stream_bytes = own_seq_bytes; if (own.head) |h| { stream_bytes = try std.math.add(u64, stream_bytes, mtp.Draft.bytes(window, h.ids.len, own.engine.chunk)); - } else if (o.drafts and o.tp == 2) stream_bytes += mtp.Draft.bytes(window, 1, own.engine.chunk); + } else if (o.drafts and o.tp == 2) stream_bytes = try std.math.add(u64, stream_bytes, mtp.Draft.bytes(if (mode == .exl3) 1 else window, 1, own.engine.chunk)); // an EXL3 rank 1 keeps no head cache if (o.vision) |v| stream_bytes = try std.math.add(u64, stream_bytes, try std.math.add(u64, try std.math.mul(u64, v.image_tokens, @import("cuda_prompt_wt.zig").dims * 2), try std.math.mul(u64, 16, o.context))); if (page_tokens > 0) { // the lane's head runs at the engine's chunk: its cache on the stream's pages, its buffers in the run @@ -101,15 +102,17 @@ pub fn open(comptime F: type, gpa: std.mem.Allocator, io: std.Io, ctx: *const cu own.kv.layout = @import("cuda_kv_lanes.zig").layout(o.context, page_tokens, own.kv.head, stream_bytes - own_seq_bytes - head_bytes); } try @import("exl3/kv/cuda_exl3_kv.zig").logStream(own.engine, window, stream_bytes); // K10: what a stream holds, by piece + own.exl3 = mode == .exl3; + own.stream_bytes = stream_bytes; const pool_budget = try exl3_pass.budgetForStreams(stream_bytes, o.lanes); const placement = @import("exl3/offload/cuda_exl3_frequency_state.zig"); - if (try placement.prepare(gpa, io, dir, own.engine.layers[0..own.engine.n_layers], @intCast(o.tp), @intCast(o.rank))) |prepared| { + if (try placement.prepare(gpa, io, dir, own.engine.layers[0..own.engine.n_layers], own.head, @intCast(o.tp), @intCast(o.rank))) |prepared| { own.profile = prepared.state; own.profile_counts = prepared.counts; } errdefer if (own.profile) |*profile| profile.deinit(); const profile_counts = if (own.profile_counts) |*counts| &counts.decoder else null; - const attached = try exl3_pass.attach(gpa, io, driver, own.stream, own.engine.layers[0..own.engine.n_layers], own.engine.chunk, pool_budget, profile_counts); + const attached = try exl3_pass.attach(gpa, io, driver, own.stream, own.engine.layers[0..own.engine.n_layers], if (own.head) |h| &h.layer else null, own.engine.chunk, pool_budget, profile_counts); if (own.profile_counts) |*counts| try placement.finish(own, counts); const ledger: memory.Ledger = .{ .persistent = persistent_bytes, diff --git a/zig/src/families/flashnext/cuda_lanes_snap.zig b/zig/src/families/flashnext/cuda_lanes_snap.zig index fe8cc9b71c..cf626e6c36 100644 --- a/zig/src/families/flashnext/cuda_lanes_snap.zig +++ b/zig/src/families/flashnext/cuda_lanes_snap.zig @@ -18,8 +18,8 @@ pub fn Copies(comptime of: anytype) type { const l = try self.lane(s); if (l.seq.fed != at) return error.NotAtMark; if (self.ranks) |r| { - if (l.draft != null) return error.TensorRankDraftSnapshotUnavailable; - return try r.save(&l.seq); + if (l.draft != null and !self.exl3) return error.TensorRankDraftSnapshotUnavailable; + return try r.save(&l.seq, if (l.draft) |*d| d else null); // an EXL3 pack's state carries the head's rows on rank 0 } return try snap.save(self.gpa, self.engine.driver, self.stream, &l.seq, if (l.draft) |*d| d else null); } @@ -32,8 +32,8 @@ pub fn Copies(comptime of: anytype) type { return err; }; if (self.ranks) |r| { - if (l.draft != null) return error.NoHeadState; - return r.restore(&l.seq, @ptrCast(@alignCast(saved))); + if (l.draft != null and !self.exl3) return error.NoHeadState; + return r.restore(&l.seq, if (l.draft) |*d| d else null, @ptrCast(@alignCast(saved))); } self.engine.round_graph.clear(); snap.restore(self.stream, &l.seq, if (l.draft) |*d| d else null, @ptrCast(@alignCast(saved))) catch |err| { diff --git a/zig/src/families/flashnext/cuda_launch.zig b/zig/src/families/flashnext/cuda_launch.zig index 4171f2e5fc..7e69b9eb33 100644 --- a/zig/src/families/flashnext/cuda_launch.zig +++ b/zig/src/families/flashnext/cuda_launch.zig @@ -23,8 +23,9 @@ pub const moe = @import("cuda_moe.zig"); pub const attn = @import("cuda_attn.zig"); pub const ple = @import("cuda_ple.zig"); pub const ple_check = @import("cuda_ple_check.zig"); -/// An EXL3 pack's n-gram table (cuda_ple.Ngram's `.exl3` format) +/// An EXL3 pack's n-gram table and draft head rows (cuda_ple.Ngram's `.exl3` format, cuda_mtp_load.draft's EXL3 branch) pub const exl3_ple = @import("exl3/model/cuda_exl3_ple.zig"); +pub const exl3_mtp = @import("exl3/tp/cuda_exl3_mtp.zig"); pub const mix = @import("cuda_mix.zig"); pub const prompt = @import("cuda_prompt.zig"); pub const snapshot = @import("cuda_snap.zig"); @@ -447,6 +448,7 @@ test { _ = row_path; _ = @import("cuda_mtp_check.zig"); _ = @import("cuda_draft_affine.zig"); + _ = @import("exl3/tp/cuda_exl3_mtp.zig"); _ = @import("cuda_head.zig"); _ = fn_back; _ = pocket; @@ -454,7 +456,9 @@ test { _ = @import("living.zig"); } -test "exl3: dispatch host contracts" { +test "exl3: greedy verification host contracts" { + _ = @import("cuda_mtp_ops.zig"); + _ = @import("cuda_mtp_draft.zig"); _ = @import("exl3/model/cuda_exl3_dispatch_test.zig"); } diff --git a/zig/src/families/flashnext/cuda_mtp.zig b/zig/src/families/flashnext/cuda_mtp.zig index 2d822aa88b..f7cd7433cb 100644 --- a/zig/src/families/flashnext/cuda_mtp.zig +++ b/zig/src/families/flashnext/cuda_mtp.zig @@ -186,8 +186,12 @@ pub const Head = struct { defer gpa.free(ids); const head_load = @import("cuda_mtp_load.zig"); const role = try head_load.mode(); - const dtype = (try store.tensor("language_model.lm_head.weight")).dtype; - h.artifact = try proposal.legacy.artifact(gpa, ids, .{ .vocab = vocab, .hidden = dims, .format = if (dtype == .bf16) "bf16" else "affine4", .group = if (dtype == .bf16) 0 else 32, .scale_recipe = if (dtype == .bf16) @tagName(role) else "stored affine4 group32" }); + // an EXL3 pack has no dense lm_head tensor: its draft rows follow the requested head role + const bf16 = if (store.tensor("language_model.lm_head.weight")) |tensor| tensor.dtype == .bf16 else |err| switch (err) { + error.MissingTensor => role == .bf16, + else => return err, + }; + h.artifact = try proposal.legacy.artifact(gpa, ids, .{ .vocab = vocab, .hidden = dims, .format = if (bf16) "bf16" else "affine4", .group = if (bf16) 0 else 32, .scale_recipe = if (bf16) @tagName(role) else "stored affine4 group32" }); errdefer h.artifact.deinit(); h.ids = h.artifact.ids(); h.service = try lanes.proposal_head.Service.legacy(gpa, &h.artifact); diff --git a/zig/src/families/flashnext/cuda_mtp_load.zig b/zig/src/families/flashnext/cuda_mtp_load.zig index 37062a694c..da43e06c8e 100644 --- a/zig/src/families/flashnext/cuda_mtp_load.zig +++ b/zig/src/families/flashnext/cuda_mtp_load.zig @@ -15,6 +15,7 @@ pub fn mode() !Mode { pub fn draft(a: std.mem.Allocator, d: *cuda.Driver, store: *wt.Store, artifact: *const core.proposal_vocab.Artifact, role: Mode) !wt.Q4 { const ids = artifact.ids(); const affine = role == .affine4; + if (try @import("exl3/tp/cuda_exl3_mtp.zig").draft(a, d, store, ids, affine)) |q| return q; const tensor = try store.tensor("language_model.lm_head.weight"); if (tensor.dtype == .bf16) { if (!tensor.is(.bf16, &.{ wt.vocab, wt.dims })) return error.UnexpectedTensor; diff --git a/zig/src/families/flashnext/cuda_mtp_ops.zig b/zig/src/families/flashnext/cuda_mtp_ops.zig index 3366bf7c8f..9afbe7ee65 100644 --- a/zig/src/families/flashnext/cuda_mtp_ops.zig +++ b/zig/src/families/flashnext/cuda_mtp_ops.zig @@ -102,7 +102,7 @@ pub fn Ops(comptime Tri: type) type { try o.forwardMain(seq, ids[at..stop], .{ .logits = if (last) .last else .none, .streams = d.streams.ptr }); const head_to = @min(stop, ids.len - 1); if (head_to > at) try o.forwardHead(d, ids[at + 1 .. head_to + 1], d.streams.ptr, false); - if (last) d.prompt_last = d.streams.ptr + (stop - 1 - at) * wide * 2; + d.prompt_last = d.streams.ptr + (stop - 1 - at) * wide * 2; } /// The whole prompt a chunk of the engine's rows at a time, from an empty head cache. @@ -413,6 +413,8 @@ pub fn Ops(comptime Tri: type) type { /// One forward over the window `ids` from seq.fed: each row's scores in d.logits, its streams in d.streams. pub fn verify(o: Self, seq: *prompt.Seq, d: *Draft, ids: []const u32) !void { + const telemetry_scope = o.e.telemetry.enter(.verify); + defer telemetry_scope.restore(); if (o.main) |m| return m.verify(m.ptr, seq, d, ids); o.e.round_graph.clear(); const n = ids.len; diff --git a/zig/src/families/flashnext/cuda_mtp_ops_test.zig b/zig/src/families/flashnext/cuda_mtp_ops_test.zig new file mode 100644 index 0000000000..bf7fa1578c --- /dev/null +++ b/zig/src/families/flashnext/cuda_mtp_ops_test.zig @@ -0,0 +1,75 @@ +const std = @import("std"); +const cuda = @import("cuda"); +const batch = @import("cuda_batch.zig"); +const prompt = @import("cuda_prompt.zig"); +const mtp = @import("cuda_mtp.zig"); +const ops_ = @import("cuda_mtp_ops.zig"); + +const Ops = ops_.Ops; +const Draft = mtp.Draft; +const max_rows = mtp.max_rows; + +test "exl3: common verification scopes TP1 and TP2 validation and delegation without joins" { + const Tri = @import("cuda_launch.zig").Tri; + const Fake = @import("exl3/model/cuda_exl3_fake_driver_test.zig").Fake; + const Phase = @import("exl3/model/exl3_telemetry.zig").Phase; + const Delegate = struct { + engine: *prompt.Engine, + calls: usize = 0, + fail: bool = false, + + fn verify(ptr: *anyopaque, _: *prompt.Seq, _: *Draft, _: []const u32) !void { + const delegate: *@This() = @ptrCast(@alignCast(ptr)); + delegate.calls += 1; + try std.testing.expectEqual(Phase.verify, delegate.engine.telemetry.forward_phase); + if (delegate.fail) return error.VerificationFailed; + } + + fn forward(_: *anyopaque, _: *prompt.Seq, _: []const u32, _: batch.Hooks) !void { + return error.UnexpectedDelegation; + } + + fn rollback(_: *anyopaque, _: *prompt.Seq, _: *Draft, _: usize) !void { + return error.UnexpectedDelegation; + } + }; + var driver = Fake.driver(); + var stream: cuda.Stream = .{ .d = &driver, .handle = null }; + var engine: prompt.Engine = undefined; + engine.stream = &stream; + engine.round_graph = .{}; + var sequence: prompt.Seq = undefined; + sequence.cap = max_rows; + var window: Draft = undefined; + var delegate: Delegate = .{ .engine = &engine }; + var ops: Ops(Tri) = .{ .e = &engine, .tri = .{ .s = stream }, .h = undefined }; + const too_wide: [max_rows + 1]u32 = @splat(1); + const syncs = Fake.sync_calls; + for ([_]usize{ 1, 2 }) |world| { + for ([_]bool{ false, true }) |enabled| { + for ([_]Phase{ .unknown, .prefill, .decode, .verify }) |previous| { + engine.telemetry = .{ .enabled = enabled, .rank = .{ .rank = 0, .world = world }, .forward_phase = previous }; + sequence.fed = 0; + ops.main = null; + try std.testing.expectError(error.WindowTooWide, ops.verify(&sequence, &window, &.{})); + try std.testing.expectEqual(previous, engine.telemetry.forward_phase); + try std.testing.expectError(error.WindowTooWide, ops.verify(&sequence, &window, &too_wide)); + try std.testing.expectEqual(previous, engine.telemetry.forward_phase); + sequence.fed = sequence.cap; + try std.testing.expectError(error.PromptTooLong, ops.verify(&sequence, &window, &.{1})); + try std.testing.expectEqual(previous, engine.telemetry.forward_phase); + ops.main = .{ .ptr = &delegate, .forward = Delegate.forward, .verify = Delegate.verify, .rollback = Delegate.rollback }; + const calls = delegate.calls; + delegate.fail = false; + try ops.verify(&sequence, &window, &.{1}); + try std.testing.expectEqual(previous, engine.telemetry.forward_phase); + delegate.fail = true; + try std.testing.expectError(error.VerificationFailed, ops.verify(&sequence, &window, &.{1})); + try std.testing.expectEqual(previous, engine.telemetry.forward_phase); + try std.testing.expectEqual(calls + 2, delegate.calls); + try std.testing.expectEqual(enabled, engine.telemetry.enabled); + try std.testing.expectEqual(syncs, Fake.sync_calls); + } + } + } +} diff --git a/zig/src/families/flashnext/cuda_prompt.zig b/zig/src/families/flashnext/cuda_prompt.zig index 93932f8803..f606115540 100644 --- a/zig/src/families/flashnext/cuda_prompt.zig +++ b/zig/src/families/flashnext/cuda_prompt.zig @@ -137,6 +137,8 @@ pub const Engine = struct { /// The pass's scratch, grown to the widest chunk run. batch_s: ?batch.Scratch = null, scratch_format: @import("cuda_batch_state.zig").ScratchFormat = .classic, + /// Local head staging bound for this engine, fixed before its first scratch allocation. + head_local_rows: usize = std.math.maxInt(usize), /// Block scores for the sparse rows of an attention block, grown to the longest sequence's need. scores: ?cuda.DeviceBuffer = null, /// Rows a batched chunk takes at most (chunkRows). @@ -168,6 +170,7 @@ pub const Engine = struct { e.stream = stream; e.n_layers = 0; e.scratch_format = .classic; + e.head_local_rows = std.math.maxInt(usize); e.collective = null; e.image_factory = null; e.telemetry = cache_telemetry.State.fromEnv(rank, world); @@ -386,7 +389,7 @@ pub const Engine = struct { bs.free(e.gpa); e.batch_s = null; } - e.batch_s = try batch.Scratch.allocForFormat(e.gpa, e.driver, e.stream.*, rows, e.scratch_format); + e.batch_s = try batch.Scratch.allocForHead(e.gpa, e.driver, e.stream.*, rows, e.scratch_format, e.head_local_rows); return &e.batch_s.?; } diff --git a/zig/src/families/flashnext/cuda_prompt_telemetry.zig b/zig/src/families/flashnext/cuda_prompt_telemetry.zig index b346b9ee74..17d1c55a9a 100644 --- a/zig/src/families/flashnext/cuda_prompt_telemetry.zig +++ b/zig/src/families/flashnext/cuda_prompt_telemetry.zig @@ -5,7 +5,7 @@ const counters = cache.counters; pub const State = struct { enabled: bool = false, layers_enabled: bool = false, - layer_cursor: counters.LayerCursor(@import("cuda_prompt_wt.zig").layers_n) = .{}, + layer_cursor: counters.LayerCursor(@import("cuda_prompt_wt.zig").layers_n + 1) = .{}, rank: counters.Rank = .{ .rank = 0, .world = 1 }, forward_phase: counters.Phase = .unknown, @@ -36,6 +36,26 @@ pub const Scope = struct { } }; +pub fn forwardArg(state: State, fed: usize) u64 { + if (!state.enabled) return 0; + return switch (state.forwardPhase(fed)) { + .unknown => 0, + .prefill => 1, + .decode => 2, + .verify => 3, + }; +} + +pub fn phaseFromArg(value: u64) !counters.Phase { + return switch (value) { + 0 => .unknown, + 1 => .prefill, + 2 => .decode, + 3 => .verify, + else => error.BadRankFrame, + }; +} + pub fn Operations(comptime Engine: type, comptime Seq: type, comptime Hooks: type, comptime underlying: anytype) type { return struct { pub fn complete(engine: *Engine) !void { @@ -93,7 +113,7 @@ pub fn beginWithSampler(engine: anytype, sequences: anytype, phase: counters.Pha failed(engine, phase, failure); return failure; }; - return sample(engine.gpa, engine, sequences, engine.telemetry.rank, cache.Completion.all_rank_consumers_joined) catch |failure| { + return sample(engine.gpa, engine, null, sequences, engine.telemetry.rank, cache.Completion.all_rank_consumers_joined) catch |failure| { failed(engine, phase, failure); return null; }; @@ -108,7 +128,7 @@ pub fn finishWithSampler(engine: anytype, sequences: anytype, before: counters.S failed(engine, phase, failure); return failure; }; - var after = sample(engine.gpa, engine, sequences, engine.telemetry.rank, cache.Completion.all_rank_consumers_joined) catch |failure| { + var after = sample(engine.gpa, engine, null, sequences, engine.telemetry.rank, cache.Completion.all_rank_consumers_joined) catch |failure| { failed(engine, phase, failure); return; }; @@ -138,7 +158,7 @@ fn publish(engine: anytype, report: counters.Report) !void { test "exl3: serving telemetry disabled adds no completion or readback and opt in samples only completed boundaries" { const Fake = struct { gpa: std.mem.Allocator, - telemetry: State = .{ .rank = .{ .rank = 0, .world = 1 } }, + telemetry: State = .{ .rank = .{ .rank = 1, .world = 2 } }, completions: u64 = 0, reads: u64 = 0, records: u64 = 0, @@ -147,7 +167,7 @@ test "exl3: serving telemetry disabled adds no completion or readback and opt in engine.completions += 1; } - fn sample(_: std.mem.Allocator, engine: *@This(), _: @TypeOf(null), rank: counters.Rank, _: cache.Completion) !counters.Snapshot { + fn sample(_: std.mem.Allocator, engine: *@This(), _: @TypeOf(null), _: @TypeOf(null), rank: counters.Rank, _: cache.Completion) !counters.Snapshot { try std.testing.expect(engine.completions > engine.reads); engine.reads += 1; return .{ .rank = rank, .scope = 1, .experts = .{ .device_copies = engine.reads } }; @@ -157,7 +177,7 @@ test "exl3: serving telemetry disabled adds no completion or readback and opt in try std.testing.expectEqual(counters.Phase.decode, report.interval.phase); try std.testing.expectEqual(@as(u64, 2), report.interval.model_rows); try std.testing.expectEqual(@as(u64, 1), report.experts.device_copies); - try std.testing.expectEqual(@as(usize, 0), report.rank.rank); + try std.testing.expectEqual(@as(usize, 1), report.rank.rank); engine.records += 1; } }; @@ -174,6 +194,41 @@ test "exl3: serving telemetry disabled adds no completion or readback and opt in try std.testing.expectEqual(@as(u64, 1), engine.records); } +test "exl3: TP forward metadata has explicit bounded phase mapping independent of token count" { + const rank_engine = @import("cuda_rank_engine.zig"); + const wire = @import("cuda_rank_wire.zig"); + const logits = @import("cuda_prompt.zig").batch.Logits; + const phases = [_]counters.Phase{ .unknown, .prefill, .decode, .verify }; + var sender: State = .{ .enabled = true }; + var worker: State = .{ .enabled = true, .forward_phase = .verify }; + for (phases, 0..) |phase, index| { + sender.forward_phase = phase; + for ([_]u32{ 4, 8192 }) |bytes| { + const header: wire.Header = .{ .kind = .forward, .object = 17, .arg = forwardArg(sender, 16384), .bytes = bytes, .flags = @backingInt(logits.last) }; + try std.testing.expectEqual(@as(u64, index), header.arg); + const frame = try header.encode(); + const received = try wire.Header.decode(&frame); + try rank_engine.shape(received); + { + const scope = worker.enter(try phaseFromArg(received.arg)); + defer scope.restore(); + try std.testing.expectEqual(sender.forwardPhase(16384), worker.forwardPhase(16384)); + } + try std.testing.expectEqual(counters.Phase.verify, worker.forward_phase); + } + } + sender.enabled = false; + sender.forward_phase = .prefill; + try std.testing.expectEqual(@as(u64, 0), forwardArg(sender, 0)); + try rank_engine.shape(.{ .kind = .forward, .object = 17, .bytes = 4 }); + for ([_]u64{ 4, 255, std.math.maxInt(u64) }) |invalid_arg| { + try std.testing.expectError(error.BadRankFrame, phaseFromArg(invalid_arg)); + try std.testing.expectError(error.BadRankFrame, rank_engine.shape(.{ .kind = .forward, .object = 17, .arg = invalid_arg, .bytes = 4 })); + } + try std.testing.expectError(error.BadRankFrame, rank_engine.shape(.{ .kind = .verify, .object = 17, .arg = 3, .bytes = 4 })); + try std.testing.expectError(error.BadRankFrame, rank_engine.shape(.{ .kind = .round, .arg = 1, .bytes = 12 })); +} + test "exl3: forward and verify phase scopes restore nested caller state on error" { const Call = struct { fn fail(state: *State, phase: counters.Phase) !void { @@ -188,7 +243,7 @@ test "exl3: forward and verify phase scopes restore nested caller state on error const parent = state.enter(.prefill); try std.testing.expectError(error.ScopedOperationFailed, Call.fail(&state, .verify)); try std.testing.expectEqual(counters.Phase.prefill, state.forward_phase); - try std.testing.expectError(error.ScopedOperationFailed, Call.fail(&state, .prefill)); + try std.testing.expectError(error.ScopedOperationFailed, Call.fail(&state, try phaseFromArg(1))); try std.testing.expectEqual(counters.Phase.prefill, state.forward_phase); parent.restore(); try std.testing.expectEqual(previous, state.forward_phase); diff --git a/zig/src/families/flashnext/cuda_rank_engine.zig b/zig/src/families/flashnext/cuda_rank_engine.zig index 35a1650541..31e3056d1e 100644 --- a/zig/src/families/flashnext/cuda_rank_engine.zig +++ b/zig/src/families/flashnext/cuda_rank_engine.zig @@ -7,6 +7,7 @@ const wire = @import("cuda_rank_wire.zig"); const Collective = @import("cuda_collective.zig").Collective; const wt = @import("cuda_prompt_wt.zig"); const mtp = @import("cuda_mtp.zig"); +const telemetry = @import("cuda_prompt_telemetry.zig"); pub const Identity = extern struct { config: [32]u8, @@ -187,7 +188,7 @@ pub const Session = struct { pub fn forward(s: *Session, comptime Tri: type, tri: Tri, q: *prompt.Seq, ids: []const u32, hooks: prompt.batch.Hooks) !void { if (ids.len == 0 or ids.len > wire.max_payload / 4 or q.fed + ids.len > q.cap or hooks.keep != null or hooks.logits == .all) return error.UnsupportedRankOperation; try tokenIds(ids); - const h = try s.command(.{ .kind = .forward, .object = q.rank_id, .bytes = @intCast(ids.len * 4), .flags = @backingInt(hooks.logits) }, std.mem.sliceAsBytes(ids)); + const h = try s.command(.{ .kind = .forward, .object = q.rank_id, .arg = telemetry.forwardArg(s.engine.telemetry, q.fed), .bytes = @intCast(ids.len * 4), .flags = @backingInt(hooks.logits) }, std.mem.sliceAsBytes(ids)); errdefer s.poison(); try s.engine.forward(Tri, tri, q, ids, hooks); try s.ack(h); @@ -235,26 +236,32 @@ pub const Session = struct { const item = try s.windows.getOrPut(s.gpa, q.rank_id); if (!item.found_existing) { errdefer _ = s.windows.remove(q.rank_id); - item.value_ptr.* = try mtp.Draft.init(s.gpa, s.engine.driver, s.engine.stream.*, q.cap, 1, s.engine.chunk); + // an EXL3 rank 1 never runs the head: its window keeps the rollback buffers and a one-entry head cache + const exl3 = s.engine.scratch_format == .exl3; + item.value_ptr.* = try mtp.Draft.init(s.gpa, s.engine.driver, s.engine.stream.*, if (exl3) 1 else q.cap, 1, s.engine.chunk); } return item.value_ptr; } - pub fn save(s: *Session, q: *prompt.Seq) !*snaps.Snap { - const kept = try snaps.save(s.gpa, s.engine.driver, s.engine.stream.*, q, null); + /// `d`: rank 0's head cache (a drafting stream); rank 1 keeps the sequence alone (it never runs the head). + pub fn save(s: *Session, q: *prompt.Seq, d: ?*mtp.Draft) !*snaps.Snap { + const kept = try snaps.save(s.gpa, s.engine.driver, s.engine.stream.*, q, d); errdefer snaps.drop(s.gpa, s.engine.stream.*, kept); kept.rank_id = s.next_object; s.next_object = std.math.add(u64, s.next_object, 1) catch return error.BadRankFrame; const h = try s.command(.{ .kind = .save, .object = q.rank_id, .arg = kept.rank_id }, &.{}); + errdefer s.poison(); try s.ack(h); return kept; } - pub fn restore(s: *Session, q: *prompt.Seq, kept: *const snaps.Snap) !void { + pub fn restore(s: *Session, q: *prompt.Seq, d: ?*mtp.Draft, kept: *const snaps.Snap) !void { if (kept.at > q.cap) return error.SequenceTooShort; + // refused before rank 1 hears of it (snaps.restore's own check, made first) + if (d) |hd| if (!kept.head or kept.at > hd.cap) return error.NoHeadState; const h = try s.command(.{ .kind = .restore, .object = q.rank_id, .arg = kept.rank_id }, &.{}); errdefer s.poison(); - try snaps.restore(s.engine.stream.*, q, null, kept); + try snaps.restore(s.engine.stream.*, q, d, kept); try s.ack(h); } @@ -315,6 +322,8 @@ pub const Session = struct { defer s.gpa.free(ids); for (ids, 0..) |*v, i| v.* = std.mem.readInt(u32, payload[i * 4 ..][0..4], .little); try tokenIds(ids); + const telemetry_scope = s.engine.telemetry.enter(try telemetry.phaseFromArg(h.arg)); + defer telemetry_scope.restore(); try s.engine.forward(Tri, tri, q, ids, .{ .logits = @fromBackingInt(@intCast(h.flags)) }); }, .round => { @@ -360,6 +369,7 @@ pub const Session = struct { .save => { if (s.saved.contains(h.arg)) return error.BadRankFrame; const q = s.seqs.get(h.object) orelse return error.UnknownRankSequence; + if (s.windows.getPtr(h.object)) |d| if (d.open) return error.NotAtMark; const kept = try snaps.save(s.gpa, s.engine.driver, s.engine.stream.*, q, null); errdefer snaps.drop(s.gpa, s.engine.stream.*, kept); try s.saved.put(s.gpa, h.arg, kept); @@ -367,7 +377,13 @@ pub const Session = struct { .restore => { const q = s.seqs.get(h.object) orelse return error.UnknownRankSequence; const kept = s.saved.get(h.arg) orelse return error.UnknownRankSnapshot; + s.engine.round_graph.clear(); try snaps.restore(s.engine.stream.*, q, null, kept); + // a drafting sequence's verify window starts over at the restored mark, as rank 0's head cache does + if (s.windows.getPtr(h.object)) |d| { + d.reset(); + d.len = kept.at; + } }, .drop_snap => { const entry = s.saved.fetchRemove(h.object) orelse return error.UnknownRankSnapshot; @@ -406,7 +422,10 @@ pub fn shape(h: wire.Header) !void { switch (h.kind) { .create => if (h.object == 0 or h.arg == 0 or h.arg > prompt.max_tokens or h.bytes != 0 or h.flags != 0) return error.BadRankFrame, .media => if (h.object == 0 or h.arg == 0 or h.arg > prompt.max_tokens or h.flags != 0 or h.bytes == 0 or h.bytes % 12 != 0 or h.bytes / 12 > 64) return error.BadRankFrame, - .forward => if (h.object == 0 or h.arg != 0 or h.bytes == 0 or h.bytes % 4 != 0 or (h.flags != @backingInt(prompt.batch.Logits.last) and h.flags != @backingInt(prompt.batch.Logits.none))) return error.BadRankFrame, + .forward => { + _ = try telemetry.phaseFromArg(h.arg); + if (h.object == 0 or h.bytes == 0 or h.bytes % 4 != 0 or (h.flags != @backingInt(prompt.batch.Logits.last) and h.flags != @backingInt(prompt.batch.Logits.none))) return error.BadRankFrame; + }, .verify => if (h.object == 0 or h.arg != 0 or h.flags != 0 or h.bytes == 0 or h.bytes % 4 != 0 or h.bytes / 4 > mtp.max_rows) return error.BadRankFrame, .keep => if (h.object == 0 or h.arg == 0 or h.arg > mtp.max_rows or h.bytes != 0 or h.flags != 0) return error.BadRankFrame, .round => if (h.object != 0 or h.arg != 0 or h.flags != 0 or h.bytes == 0 or h.bytes % 12 != 0 or h.bytes / 12 > 32) return error.BadRankFrame, diff --git a/zig/src/families/flashnext/cuda_scratch_layout_test.zig b/zig/src/families/flashnext/cuda_scratch_layout_test.zig index 589a972be2..52abdd731b 100644 --- a/zig/src/families/flashnext/cuda_scratch_layout_test.zig +++ b/zig/src/families/flashnext/cuda_scratch_layout_test.zig @@ -1,14 +1,14 @@ const std = @import("std"); const Scratch = @import("cuda_batch_state.zig").Scratch; -test "exl3: scratch omits only classic MoE allocations and the split-head logits and keeps every common view" { +test "exl3: scratch omits only classic MoE allocations and keeps every common view" { for ([_]usize{ 1, 8, 32, 64, 1024, 1536, 2048 }) |cap| { const classic = Scratch.allocationSizes(cap, .classic); const exl3 = Scratch.allocationSizes(cap, .exl3); var removed: usize = 0; var retained: usize = 0; for (classic, exl3, 0..) |before, after, index| { - if ((index >= 38 and index <= 47) or index == 61) { + if (index >= 38 and index <= 47) { try std.testing.expectEqual(@as(?usize, null), after); removed += @max(before.?, 256); } else { @@ -32,10 +32,49 @@ test "exl3: scratch savings include the allocator minimum and grow with requeste const blocks = (pairs + 1023) / 1024; const expected = 2 * @max(pairs * 4, 256) + @max(items * 12, 256) + 256 + @max(blocks * 513 * 4, 256) + 3 * @max(cap * 640 * 2, 256) + - @max(cap * 2560 * 4, 256) + @max(pairs * 640 * 2, 256) + @max(cap * 248320, 256); + @max(cap * 2560 * 4, 256) + @max(pairs * 640 * 2, 256); const saved = Scratch.bytes(cap) - Scratch.bytesForFormat(cap, .exl3); try std.testing.expectEqual(expected, saved); try std.testing.expect(saved > previous); previous = saved; } } + +const HeadGuard = struct { + const Collective = struct { + pub fn gather(_: *Collective, _: @import("cuda").Stream, _: u64, _: u64, _: usize) !void { + return error.UnexpectedLaunch; + } + }; + s: *Scratch, + m: struct { head: struct { n: usize }, collective: ?*Collective = null } = .{ .head = .{ .n = @import("cuda_prompt_wt.zig").vocab / 2 } }, + stream: @import("cuda").Stream = undefined, + + pub fn lane(_: *HeadGuard, _: u64, _: u64, _: anytype, _: u64, _: u64, _: usize, _: usize) !void { + return error.UnexpectedLaunch; + } +}; + +test "exl3: each scratch owns the same head-row bound its launch guard admits" { + const vocab = @import("cuda_prompt_wt.zig").vocab; + var small: Scratch = .{ .cap = 1024, .format = .exl3, .head_local_rows = 32 }; + var split: Scratch = .{ .cap = 1024, .format = .exl3, .head_local_rows = 64 }; + var classic: Scratch = .{ .cap = 1536 }; + const classic_bytes = Scratch.bytes(1536); + const small_sizes = Scratch.allocationSizesForHead(small.cap, small.format, small.head_local_rows); + const split_sizes = Scratch.allocationSizesForHead(split.cap, split.format, split.head_local_rows); + try std.testing.expectEqual(@as(?usize, 32 * vocab), small_sizes[61]); + try std.testing.expectEqual(@as(?usize, 64 * vocab), split_sizes[61]); + var a: HeadGuard = .{ .s = &small }; + var b: HeadGuard = .{ .s = &split }; + var c: HeadGuard = .{ .s = &classic }; + const head = @import("cuda_batch_head.zig"); + try std.testing.expectError(error.UnexpectedTensor, head.projectHead(&a, 0, 40)); + try std.testing.expectError(error.TensorParallelUnavailable, head.projectHead(&b, 0, 40)); + try std.testing.expectError(error.TensorParallelUnavailable, head.projectHead(&c, 0, 1024)); + split.head_local_rows = 16; + try std.testing.expectError(error.UnexpectedTensor, head.projectHead(&b, 0, 40)); + try std.testing.expectError(error.TensorParallelUnavailable, head.projectHead(&a, 0, 32)); + try std.testing.expectEqual(@as(usize, 1536), classic.headRows()); + try std.testing.expectEqual(classic_bytes, Scratch.bytes(1536)); +} diff --git a/zig/src/families/flashnext/exl3/kv/cuda_exl3_kv.zig b/zig/src/families/flashnext/exl3/kv/cuda_exl3_kv.zig index c87faba205..d695c9aa7d 100644 --- a/zig/src/families/flashnext/exl3/kv/cuda_exl3_kv.zig +++ b/zig/src/families/flashnext/exl3/kv/cuda_exl3_kv.zig @@ -44,6 +44,13 @@ pub fn seqBytes(cap: usize) SeqBytes { pub const gib: f64 = 1024.0 * 1024.0 * 1024.0; +/// Widest logits pass the server runs (a decode round of 32 streams, an MTP window): `head_local`'s rows. +pub const head_local_serve_rows: usize = 64; + +/// A full-vocabulary rank still stages every batched MTP lane's BF16 scores in the shared head buffer. +pub fn headLocalRows(full_vocab: bool) usize { + return if (full_vocab) 2 * @import("../../cuda_mtp_ops.zig").max_lanes else head_local_serve_rows; +} /// K10 lever 3: the pass's chunk for an EXL3 checkpoint unless TF_FLASHNEXT_CHUNK says. pub const exl3_chunk: usize = 1024; @@ -54,8 +61,19 @@ pub fn setup(e: anytype, exl3: bool) !void { const decoder_exl3 = e.n_layers != 0 and for (e.layers[0..e.n_layers]) |layer| { if (layer.exl3 == null) break false; } else true; - e.scratch_format = if (decoder_exl3) .exl3 else .classic; + const draft_exl3 = @import("../pack/cuda_exl3_load.zig").hasExperts(&e.store, "language_model.mtp.layers.0.mlp"); + e.scratch_format = if (decoder_exl3 and draft_exl3) .exl3 else .classic; if (std.c.getenv("TF_FLASHNEXT_CHUNK") == null) e.chunk = exl3_chunk; + e.head_local_rows = headLocalRows(e.head_q.n == wt.vocab); +} + +test "exl3: full-vocabulary scratch admits every shared MTP score row" { + const state = @import("../../cuda_batch_state.zig"); + const lanes = @import("../../cuda_mtp_ops.zig").max_lanes; + const empty = state.Scratch.bytesForHead(exl3_chunk, .exl3, 0); + const full = state.Scratch.bytesForHead(exl3_chunk, .exl3, headLocalRows(true)); + try std.testing.expect(full - empty + 256 >= lanes * wt.vocab * 2); + try std.testing.expect(headLocalRows(false) >= headLocalRows(true)); } /// One log line: a stream's device bytes at the window by piece. @@ -70,7 +88,7 @@ pub fn logStream(e: anytype, window: usize, stream_bytes: usize) !void { @as(f64, @floatFromInt(b.gdn)) / gib, @as(f64, @floatFromInt(stream_bytes -| b.device())) / gib, e.chunk, - @as(f64, @floatFromInt(@import("../../cuda_batch_state.zig").Scratch.bytesForFormat(e.chunk, e.scratch_format))) / gib, + @as(f64, @floatFromInt(@import("../../cuda_batch_state.zig").Scratch.bytesForHead(e.chunk, e.scratch_format, e.head_local_rows))) / gib, }); } @@ -81,3 +99,12 @@ pub fn snapHost(context: usize, exl3: bool) bool { @import("../../cuda_snap.zig").on_host = on; return on; } + +/// cuda_lanes.zig rankRoom: the ranks agree how many STREAMS fit (each its own stream_bytes: rank 1 holds no MTP head). +pub fn rankRoom(r: anytype, local: u64, stream_bytes: usize) !u64 { + if (stream_bytes == 0) return r.room(local); + const sb: u64 = stream_bytes; + const q = (local / sb) << 16 | ((local % sb) << 16) / sb; + const m = try r.room(q); + return (m >> 16) * sb + ((m & 0xffff) * sb >> 16); +} diff --git a/zig/src/families/flashnext/exl3/model/cuda_exl3.zig b/zig/src/families/flashnext/exl3/model/cuda_exl3.zig index 5b681cdf05..4ea7418da1 100644 --- a/zig/src/families/flashnext/exl3/model/cuda_exl3.zig +++ b/zig/src/families/flashnext/exl3/model/cuda_exl3.zig @@ -5,6 +5,7 @@ pub const store = @import("../pack/cuda_exl3_store.zig"); pub const memory = @import("../pack/cuda_exl3_memory.zig"); pub const load = @import("../pack/cuda_exl3_load.zig"); pub const pass = @import("cuda_exl3_pass.zig"); +pub const tp = @import("../tp/cuda_exl3_tp.zig"); pub const kv = @import("../kv/cuda_exl3_kv.zig"); pub const fork = @import("../tp/cuda_exl3_fork.zig"); pub const telemetry = @import("cuda_exl3_telemetry.zig"); @@ -22,8 +23,10 @@ test { _ = memory; _ = load; _ = pass; + _ = tp; _ = kv; _ = @import("../kv/cuda_exl3_kv_test.zig"); + _ = @import("../../cuda_mtp_ops_test.zig"); _ = fork; _ = telemetry; _ = frequency; diff --git a/zig/src/families/flashnext/exl3/model/cuda_exl3_arch.zig b/zig/src/families/flashnext/exl3/model/cuda_exl3_arch.zig index fee7e05a47..9e87c88dc9 100644 --- a/zig/src/families/flashnext/exl3/model/cuda_exl3_arch.zig +++ b/zig/src/families/flashnext/exl3/model/cuda_exl3_arch.zig @@ -5,7 +5,7 @@ const cuda = @import("cuda"); /// The one compute capability (10 * major + minor) this family's EXL3 path is qualified on: Blackwell sm_120. pub const qualified: u32 = 120; -pub const message = "Flash Next EXL3 (expert offload, one rank) is qualified on sm_120 only; this device is sm_{d}: " ++ +pub const message = "Flash Next EXL3 (expert offload, TP=2) is qualified on sm_120 only; this device is sm_{d}: " ++ "serve an NVFP4 or BF16 Flash Next checkpoint, or set TF_FLASHNEXT_EXL3_ARCH=unqualified to try it without a claim"; /// Whether a device of `capability` may serve an EXL3 pack; `override` is the TF_FLASHNEXT_EXL3_ARCH value for testers. diff --git a/zig/src/families/flashnext/exl3/model/cuda_exl3_pass.zig b/zig/src/families/flashnext/exl3/model/cuda_exl3_pass.zig index 71de32b822..dc0c966371 100644 --- a/zig/src/families/flashnext/exl3/model/cuda_exl3_pass.zig +++ b/zig/src/families/flashnext/exl3/model/cuda_exl3_pass.zig @@ -81,6 +81,7 @@ fn laneF16(p: anytype, x: u64, e: *const load.Proj, out: u64, slices: u64, m: us /// Pass.moeBlock() for EXL3 experts: the router's 513 logits (the BF16 rows the Store widened exactly), each row's ten pub fn moe(p: anytype, lw: anytype, x: *const load.Experts) !void { + if (x.tp != null) return @import("../tp/cuda_exl3_tp.zig").moe(p, lw, x); // expert parallel (plan step 10) const s = p.s; const n = p.rows; try p.tri.routerMany(s.mixed, lw.router.ptr, s.router_l, dims, n); @@ -104,6 +105,7 @@ pub fn moe(p: anytype, lw: anytype, x: *const load.Experts) !void { } try ex.runSlots(p.stream, x.layer, xin, dims, .bf16, picks, rows, sc); } + try @import("../tp/cuda_exl3_tp.zig").canonical(p); } /// Host time of the A/B cut (the picks D2H and its wait) and of Pool.plan, summed over every pooled block (read by @@ -131,13 +133,17 @@ pub fn pooled(st: cuda.Stream, ops: *load.Ops, pool: *pool_.Pool, ex: *const cud } } -/// Every offloaded experts layer of a loaded engine (the 48 decoder layers). -pub fn offloadedLayers(layers: []const @import("../../cuda_prompt_load.zig").Layer, out: []*load.Experts) usize { +/// Every offloaded experts layer of a loaded engine (the 48 decoder layers, then the MTP layer if any). +pub fn offloadedLayers(layers: []const @import("../../cuda_prompt_load.zig").Layer, mtp: ?*const @import("../../cuda_prompt_load.zig").Layer, out: []*load.Experts) usize { var n: usize = 0; for (layers) |*l| if (l.exl3) |x| if (x.needsPool()) { out[n] = x; n += 1; }; + if (mtp) |l| if (l.exl3) |x| if (x.needsPool()) { + out[n] = x; + n += 1; + }; return n; } @@ -196,16 +202,22 @@ pub fn freeBytes(d: *const cuda.Driver) !u64 { } /// The sizing hook (cuda_lanes_load.zig, once everything else of the family is allocated). -pub fn attach(gpa: std.mem.Allocator, io: std.Io, d: *const cuda.Driver, st: cuda.Stream, layers: []const @import("../../cuda_prompt_load.zig").Layer, rows: usize, b: Budget, profile_counts: ?*[pool_.shape.layers][pool_.shape.experts]u32) !Attached { - var all: [wt.layers_n]*load.Experts = undefined; - const n = offloadedLayers(layers, &all); +pub fn attach(gpa: std.mem.Allocator, io: std.Io, d: *const cuda.Driver, st: cuda.Stream, layers: []const @import("../../cuda_prompt_load.zig").Layer, mtp: ?*const @import("../../cuda_prompt_load.zig").Layer, rows: usize, b: Budget, profile_counts: ?*[pool_.shape.layers][pool_.shape.experts]u32) !Attached { + var all: [wt.layers_n + 1]*load.Experts = undefined; + const n = offloadedLayers(layers, mtp, &all); const delta_text = if (std.c.getenv("TF_FLASHNEXT_EXL3_LINE_DELTAS")) |value| std.mem.span(value) else null; var out: Attached = .{}; const scratch_before = try freeBytes(d); - // every distinct Ops (the engine's store) reserves its scratch now + // every distinct Ops (the engine's store and the MTP head's) reserves its scratch now var seen: [4]*load.Ops = undefined; var n_seen: usize = 0; for (layers) |*l| if (l.exl3) |x| try reserveOnce(x.ops, &seen, &n_seen, st, rows); + if (mtp) |l| if (l.exl3) |x| try reserveOnce(x.ops, &seen, &n_seen, st, rows); + // expert parallel (cuda_exl3_tp.zig): its masks, sums and block scratch, before the pools take what is free + for (layers) |*l| if (l.exl3) |x| if (x.tp != null) { + try @import("../tp/cuda_exl3_tp.zig").reserve(x.ops, st, x.layer, rows, x.tp.?.world); + break; + }; if (n == 0) { if (delta_text != null) return error.InvalidLineDeltaIndex; return out; @@ -236,19 +248,26 @@ pub fn attach(gpa: std.mem.Allocator, io: std.Io, d: *const cuda.Driver, st: cud const startup_bytes = try memory.poolBytes(lines, out.line_bytes, n, fixed); if (startup_bytes > pool_room) return error.InsufficientStartupMemory; out.lines = lines; - var pool_lines: [wt.layers_n]usize = @splat(lines); - var line_bytes: [wt.layers_n]usize = undefined; + var pool_lines: [wt.layers_n + 1]usize = @splat(lines); + var line_bytes: [wt.layers_n + 1]usize = undefined; var profile_sized = false; var profile_min_lines: usize = std.math.maxInt(usize); var profile_max_lines: usize = 0; - for (all[0..n], 0..) |x, index| line_bytes[index] = x.lineBytes(); + var decoder_pool_count: usize = 0; + for (all[0..n], 0..) |x, index| { + line_bytes[index] = x.lineBytes(); + const is_mtp = if (mtp) |l| (if (l.exl3) |m| m == x else false) else false; + if (!is_mtp) decoder_pool_count += 1; + } if (delta_text == null and shouldProfileSize(profileSizingEnabled(), b.lines)) if (profile_counts) |prior| { var count_rows: [wt.layers_n][]const u32 = undefined; var allocation_costs: [wt.layers_n]usize = undefined; var allocation_lines: [wt.layers_n]usize = undefined; var decoder_index: usize = 0; var has_profile_signal = false; - for (all[0..n], 0..) |_, pool_index| { + for (all[0..n], 0..) |x, pool_index| { + const is_mtp = if (mtp) |l| (if (l.exl3) |m| m == x else false) else false; + if (is_mtp) continue; for (prior[decoder_index]) |count| { has_profile_signal = has_profile_signal or count != 0; } @@ -257,10 +276,11 @@ pub fn attach(gpa: std.mem.Allocator, io: std.Io, d: *const cuda.Driver, st: cud decoder_index += 1; } if (has_profile_signal and decoder_index == wt.layers_n) { + const rank = all[0].tp orelse try @import("../tp/cuda_exl3_tp.zig").Rank.init(1, 0); const maybe_plan: ?profile_line_budget.Budget = profile_line_budget.allocate(gpa, .{ .counts = &count_rows, - .rank = 0, - .world = 1, + .rank = rank.rank, + .world = rank.world, .line_bytes = &allocation_costs, .uniform_lines = lines, .min_lines = pool_.min_lines, @@ -272,7 +292,9 @@ pub fn attach(gpa: std.mem.Allocator, io: std.Io, d: *const cuda.Driver, st: cud }; if (maybe_plan) |plan| { decoder_index = 0; - for (all[0..n], 0..) |_, pool_index| { + for (all[0..n], 0..) |x, pool_index| { + const is_mtp = if (mtp) |l| (if (l.exl3) |m| m == x else false) else false; + if (is_mtp) continue; pool_lines[pool_index] = allocation_lines[decoder_index]; profile_min_lines = @min(profile_min_lines, allocation_lines[decoder_index]); profile_max_lines = @max(profile_max_lines, allocation_lines[decoder_index]); @@ -284,22 +306,23 @@ pub fn attach(gpa: std.mem.Allocator, io: std.Io, d: *const cuda.Driver, st: cud } else { profile_sized = true; } - std.log.info("flashnext exl3: profile-weighted decoder allocation used {d}/{d} variable CUDA bytes (expert lines and device picks); saved lifetime picks are a proxy, not latency", .{ plan.allocated_bytes, plan.baseline_bytes }); + std.log.info("flashnext exl3: profile-weighted decoder allocation used {d}/{d} variable CUDA bytes (expert lines and device picks); saved lifetime picks are a proxy, not latency; MTP remains uniform", .{ plan.allocated_bytes, plan.baseline_bytes }); std.log.info("flashnext exl3: profile decoder pool sizes range {d}..{d} lines; uniform-equivalent baseline was {d}", .{ profile_min_lines, profile_max_lines, lines }); } } }; if (delta_text) |text| { - const plan = try line_deltas.apply(text, lines, n, line_bytes[0..n], pool_.min_lines, fixed, startup_bytes, pool_lines[0..n]); + const plan = try line_deltas.apply(text, lines, decoder_pool_count, line_bytes[0..n], pool_.min_lines, fixed, startup_bytes, pool_lines[0..n]); std.log.info("flashnext exl3: experimental rank-local decoder pool line deltas {s}; base={d} startup_bytes={d}", .{ text, lines, plan.pool_bytes }); } // the first decoder layer's pool owns the lifetime use table the other decoder layers share var counts = profile_counts; var decoder: u8 = 0; for (all[0..n], pool_lines[0..n]) |x, count| { - try x.attach(io, d, decoder, count, b.max_rows, counts); - if (counts == null) counts = x.pool.?.counts; - decoder += 1; + const is_mtp = if (mtp) |l| (if (l.exl3) |m| m == x else false) else false; + try x.attach(io, d, if (is_mtp) 0 else decoder, count, b.max_rows, if (is_mtp) null else counts); + if (!is_mtp and counts == null) counts = x.pool.?.counts; + if (!is_mtp) decoder += 1; out.device_bytes += x.pool.?.deviceBytes(); out.layers += 1; } @@ -351,10 +374,10 @@ fn reserveOnce(o: *load.Ops, seen: []*load.Ops, n: *usize, st: cuda.Stream, rows } /// Pool statistics summed over every attached layer. -pub fn poolStats(layers: []const @import("../../cuda_prompt_load.zig").Layer) pool_.Stats { +pub fn poolStats(layers: []const @import("../../cuda_prompt_load.zig").Layer, mtp: ?*const @import("../../cuda_prompt_load.zig").Layer) pool_.Stats { var t: pool_.Stats = .{}; - var all: [wt.layers_n]*load.Experts = undefined; - const n = offloadedPools(layers, &all); + var all: [wt.layers_n + 1]*load.Experts = undefined; + const n = offloadedPools(layers, mtp, &all); for (all[0..n]) |x| { const s = x.pool.?.stats; inline for (@typeInfo(pool_.Stats).@"struct".field_names) |f| { @@ -364,11 +387,15 @@ pub fn poolStats(layers: []const @import("../../cuda_prompt_load.zig").Layer) po return t; } -pub fn offloadedPools(layers: []const @import("../../cuda_prompt_load.zig").Layer, out: []*load.Experts) usize { +pub fn offloadedPools(layers: []const @import("../../cuda_prompt_load.zig").Layer, mtp: ?*const @import("../../cuda_prompt_load.zig").Layer, out: []*load.Experts) usize { var n: usize = 0; for (layers) |*l| if (l.exl3) |x| if (x.pool != null) { out[n] = x; n += 1; }; + if (mtp) |l| if (l.exl3) |x| if (x.pool != null) { + out[n] = x; + n += 1; + }; return n; } diff --git a/zig/src/families/flashnext/exl3/model/cuda_exl3_telemetry.zig b/zig/src/families/flashnext/exl3/model/cuda_exl3_telemetry.zig index f6f245a852..e55b1cc1c6 100644 --- a/zig/src/families/flashnext/exl3/model/cuda_exl3_telemetry.zig +++ b/zig/src/families/flashnext/exl3/model/cuda_exl3_telemetry.zig @@ -1,21 +1,22 @@ const std = @import("std"); const prompt = @import("../../cuda_prompt.zig"); +const load = @import("../../cuda_prompt_load.zig"); const pool = @import("../offload/cuda_exl3_pool.zig"); const pass = @import("cuda_exl3_pass.zig"); const weights = @import("../../cuda_prompt_wt.zig"); pub const counters = @import("exl3_telemetry.zig"); pub const Completion = enum { all_rank_consumers_joined }; -pub fn sampleAfterConsumers(allocator: std.mem.Allocator, engine: *prompt.Engine, sequences: []const *prompt.Seq, rank: counters.Rank, completion: Completion) !counters.Snapshot { +pub fn sampleAfterConsumers(allocator: std.mem.Allocator, engine: *prompt.Engine, mtp: ?*const load.Layer, sequences: []const *prompt.Seq, rank: counters.Rank, completion: Completion) !counters.Snapshot { _ = completion; try rank.validate(); var hash = std.hash.Wyhash.init(@intFromPtr(engine)); var totals: counters.Experts = .{}; - var attached: [weights.layers_n]*@import("../pack/cuda_exl3_load.zig").Experts = undefined; - const attached_count = pass.offloadedPools(engine.layers[0..engine.n_layers], &attached); + var attached: [weights.layers_n + 1]*@import("../pack/cuda_exl3_load.zig").Experts = undefined; + const attached_count = pass.offloadedPools(engine.layers[0..engine.n_layers], mtp, &attached); var layers: []counters.LayerSample = if (engine.telemetry.layers_enabled) try allocator.alloc(counters.LayerSample, attached_count) else @constCast(&.{}); errdefer allocator.free(layers); - var seen: [weights.layers_n]*pool.Pool = undefined; + var seen: [weights.layers_n + 1]*pool.Pool = undefined; var seen_count: usize = 0; for (attached[0..attached_count]) |experts| { const resident = experts.pool orelse continue; @@ -63,9 +64,10 @@ pub fn sampleAfterConsumers(allocator: std.mem.Allocator, engine: *prompt.Engine }; try counters.add(counters.Experts, &totals, contribution); if (engine.telemetry.layers_enabled) { - var layer_sample: counters.LayerSample = .{ .owner = resident.telemetry_id, .kind = .decoder, .layer = 0, .lines = resident.opt.lines, .counts = counters.LayerCounts.fromExperts(contribution) }; + var layer_sample: counters.LayerSample = .{ .owner = resident.telemetry_id, .kind = .mtp, .layer = 0, .lines = resident.opt.lines, .counts = counters.LayerCounts.fromExperts(contribution) }; for (engine.layers[0..engine.n_layers], 0..) |layer, decoder_index| if (layer.exl3) |decoder| { if (decoder.pool == resident) { + layer_sample.kind = .decoder; layer_sample.layer = decoder_index; break; } diff --git a/zig/src/families/flashnext/exl3/model/cuda_exl3_telemetry_test.zig b/zig/src/families/flashnext/exl3/model/cuda_exl3_telemetry_test.zig index 69e0ded0c5..4f8e4166b9 100644 --- a/zig/src/families/flashnext/exl3/model/cuda_exl3_telemetry_test.zig +++ b/zig/src/families/flashnext/exl3/model/cuda_exl3_telemetry_test.zig @@ -3,7 +3,7 @@ const cuda = @import("cuda"); const pool = @import("../offload/cuda_exl3_pool.zig"); const telemetry = @import("exl3_telemetry.zig"); -test "exl3: aggregate telemetry finds attached pools deduplicates pool aliases and separates committed device plans" { +test "exl3: aggregate telemetry finds attached pools deduplicates MTP aliases and separates committed device plans" { var resident: pool.Pool = undefined; resident.lines.pinned_bank = null; resident.telemetry_id = 991; @@ -23,7 +23,7 @@ test "exl3: aggregate telemetry finds attached pools deduplicates pool aliases a engine.n_layers = 1; engine.layers[0].exl3 = &experts; const collector = @import("cuda_exl3_telemetry.zig"); - var snapshot = try collector.sampleAfterConsumers(std.testing.allocator, &engine, &.{}, .{ .rank = 0, .world = 1 }, .all_rank_consumers_joined); + var snapshot = try collector.sampleAfterConsumers(std.testing.allocator, &engine, &engine.layers[0], &.{}, .{ .rank = 0, .world = 2 }, .all_rank_consumers_joined); defer snapshot.deinit(std.testing.allocator); try std.testing.expectEqual(@as(usize, 1), snapshot.expert_pools); try std.testing.expectEqual(@as(u64, 2), snapshot.experts.host_plans); diff --git a/zig/src/families/flashnext/exl3/model/exl3.zig b/zig/src/families/flashnext/exl3/model/exl3.zig index 73e5d75ec4..32926bd349 100644 --- a/zig/src/families/flashnext/exl3/model/exl3.zig +++ b/zig/src/families/flashnext/exl3/model/exl3.zig @@ -2,6 +2,8 @@ pub const shape = @import("../pack/exl3_shape.zig"); pub const scan = @import("../pack/exl3_scan.zig"); +pub const tp = @import("../tp/exl3_tp.zig"); +pub const canon = @import("../tp/exl3_canon.zig"); pub const cache = @import("exl3_cache.zig"); pub const cpu = @import("exl3_cpu.zig"); pub const place = @import("../pack/exl3_place.zig"); @@ -18,6 +20,7 @@ pub const fwd_cpu = @import("exl3_fwd_cpu.zig"); test { _ = shape; _ = scan; + _ = tp; _ = cache; _ = cpu; _ = place; diff --git a/zig/src/families/flashnext/exl3/model/exl3_hc.zig b/zig/src/families/flashnext/exl3/model/exl3_hc.zig index 1d4a25ef14..517c5c125d 100644 --- a/zig/src/families/flashnext/exl3/model/exl3_hc.zig +++ b/zig/src/families/flashnext/exl3/model/exl3_hc.zig @@ -1,7 +1,32 @@ -//! Host reference of the hyper-connection pre, post and expand steps over the unquantized hc weights. +//! Host side of the unquantized hc_lowrank split. const std = @import("std"); const shape = @import("../pack/exl3_shape.zig"); +const tp = @import("../tp/exl3_tp.zig"); + +/// One rank's dense product. +pub fn partial(w: []const f32, x: []const f32, out: []f32) error{Length}!void { + const part = shape.hc_lowrank / 2; + const n = shape.hidden; + if (w.len != part * n or x.len != part or out.len != n) return error.Length; + for (0..n) |j| { + var sum: f32 = 0; + for (0..part) |i| sum += x[i] * w[i * n + j]; + out[j] = sum; + } +} + +/// Hidden-sized sum of the two partials. +pub fn add(partial0: []const f32, partial1: []const f32, out: []f32) error{Length}!void { + const n = shape.hidden; + if (partial0.len != n or partial1.len != n or out.len != n) return error.Length; + for (partial0, partial1, out) |a, b, *y| y.* = a + b; +} + +/// `exl3_bounce_bytes` for hc_lowrank only. +pub fn bounceBytes(tokens: usize) u64 { + return @as(u64, @intCast(tokens)) * 5120 * 2; +} // --------------------------------------------------------------------------------------------- Hyper-connection pre /. @@ -251,3 +276,89 @@ test "pre, post and expand refuse bad lengths before a store" { try std.testing.expectEqualSlices(f32, &.{ 1, 2, 3, 4, 1, 2, 3, 4 }, &out); try std.testing.expectError(error.Length, expand(tiny, &.{ 1, 2, 3 }, &out)); } + +test "one-hot x picks the row of the hc_lowrank part" { + const cut0 = try tp.cut(.hc_lowrank, 0, 2); + try std.testing.expectEqual(@as(usize, 160), cut0.part); + const r0 = try tp.range(.hc_lowrank, 0, 2); + const r1 = try tp.range(.hc_lowrank, 1, 2); + try std.testing.expectEqual(@as(usize, 0), r0.lo); + try std.testing.expectEqual(@as(usize, 160), r0.hi); + try std.testing.expectEqual(@as(usize, 160), r1.lo); + try std.testing.expectEqual(@as(usize, 320), r1.hi); + try std.testing.expectEqual(cut0.part, r0.hi - r0.lo); + try std.testing.expectEqual(cut0.part, r1.hi - r1.lo); + + const part = cut0.part; + const n = shape.hidden; + try std.testing.expectEqual(@as(usize, 160 * 2560), part * n); + + const gpa = std.testing.allocator; + const w = try gpa.alloc(f32, part * n); + defer gpa.free(w); + @memset(w, 0); + + var x: [shape.hc_lowrank / 2]f32 = @splat(0); + try std.testing.expectEqual(part, x.len); + + var rank0: [shape.hidden]f32 = @splat(42); + const kept0 = rank0; + try std.testing.expectError(error.Length, partial(w[0..part], &x, &rank0)); + try std.testing.expectError(error.Length, partial(w, x[0 .. part - 1], &rank0)); + try std.testing.expectError(error.Length, partial(w, &x, rank0[0 .. n - 1])); + try std.testing.expectEqualSlices(f32, &kept0, &rank0); + + // Rank 0 owns global [0, 160). A one-hot at local 0 is that first row, not the poisoned neighbor. + const local0: usize = 0; + try std.testing.expectEqual(r0.lo + local0, local0); + for (0..n) |j| w[local0 * n + j] = @floatFromInt(j + 3); + for (0..n) |j| w[1 * n + j] = -1000; + x[local0] = 1; + try partial(w, &x, &rank0); + try std.testing.expectEqualSlices(f32, w[local0 * n .. local0 * n + n], &rank0); + + // Rank 1 owns global [160, 320). A one-hot at the last local row is global 319. + var rank1: [shape.hidden]f32 = @splat(-7); + @memset(w, 0); + const local1 = part - 1; + try std.testing.expectEqual(@as(usize, 319), r1.lo + local1); + try std.testing.expect(r1.lo + local1 < r1.hi); + for (0..n) |j| w[0 * n + j] = 80; + for (0..n) |j| w[local1 * n + j] = -@as(f32, @floatFromInt(j)) - 1; + @memset(&x, 0); + x[local1] = 1; + try partial(w, &x, &rank1); + try std.testing.expectEqualSlices(f32, w[local1 * n .. local1 * n + n], &rank1); + + // Two nonzeros so the product is a dot, not a copy of the hot row. + @memset(w, 0); + @memset(&x, 0); + const row_a: usize = 2; + const row_b: usize = 5; + for (0..n) |j| w[row_a * n + j] = 1; + for (0..n) |j| w[row_b * n + j] = 4; + x[row_a] = 2; + x[row_b] = 0.5; + var dotted: [shape.hidden]f32 = @splat(0); + try partial(w, &x, &dotted); + const fours: [shape.hidden]f32 = @splat(4); + try std.testing.expectEqualSlices(f32, &fours, &dotted); + + var out: [shape.hidden]f32 = @splat(7); + const kept = out; + try std.testing.expectError(error.Length, add(rank0[0 .. n - 1], &rank1, &out)); + try std.testing.expectError(error.Length, add(&rank0, rank1[0..0], &out)); + try std.testing.expectError(error.Length, add(&rank0, &rank1, out[0 .. n - 1])); + try std.testing.expectEqualSlices(f32, &kept, &out); + + // (j + 3) + (-j - 1) = 2. Rank 0, then rank 1. + try add(&rank0, &rank1, &out); + const twos: [shape.hidden]f32 = @splat(2); + try std.testing.expectEqualSlices(f32, &twos, &out); +} + +test "bounceBytes is the hc_lowrank exchange for one layer" { + try std.testing.expectEqual(@as(u64, 10240), bounceBytes(1)); + try std.testing.expectEqual(@as(u64, 0), bounceBytes(0)); + try std.testing.expectEqual(@as(u64, 10240) * 3, bounceBytes(3)); +} diff --git a/zig/src/families/flashnext/exl3/model/exl3_telemetry.zig b/zig/src/families/flashnext/exl3/model/exl3_telemetry.zig index 8774d10242..bb03c39158 100644 --- a/zig/src/families/flashnext/exl3/model/exl3_telemetry.zig +++ b/zig/src/families/flashnext/exl3/model/exl3_telemetry.zig @@ -125,7 +125,7 @@ pub const LayerCounts = struct { pub const LayerSample = struct { owner: u64, - kind: enum { decoder }, + kind: enum { decoder, mtp }, layer: usize, lines: usize, counts: LayerCounts = .{}, diff --git a/zig/src/families/flashnext/exl3/offload/cuda_exl3_frequency.zig b/zig/src/families/flashnext/exl3/offload/cuda_exl3_frequency.zig index 4395f177df..912894ba62 100644 --- a/zig/src/families/flashnext/exl3/offload/cuda_exl3_frequency.zig +++ b/zig/src/families/flashnext/exl3/offload/cuda_exl3_frequency.zig @@ -14,12 +14,15 @@ pub const format_version: u32 = 1; pub const Counts = struct { decoder: [decoder_layers][experts]u32 = @splat(@splat(0)), + mtp: [experts]u32 = @splat(0), + mtp_present: bool = false, }; const magic = "TFEXL3FQ"; const header_bytes = 64; const decoder_bytes = decoder_layers * experts * @sizeOf(u32); -pub const max_file_bytes = header_bytes + decoder_bytes + Sha256.digest_length; +const mtp_bytes = experts * @sizeOf(u32); +pub const max_file_bytes = header_bytes + decoder_bytes + mtp_bytes + Sha256.digest_length; pub fn read(gpa: std.mem.Allocator, io: Io, path: []const u8, identity: Identity) !Counts { try validateIdentity(identity); @@ -50,7 +53,8 @@ pub fn write(gpa: std.mem.Allocator, io: Io, path: []const u8, identity: Identit const basename = std.fs.path.basename(path); if (path.len == 0 or std.mem.endsWith(u8, path, "/") or std.mem.eql(u8, basename, ".") or std.mem.eql(u8, basename, "..")) return error.InvalidPath; - const bytes = try gpa.alloc(u8, max_file_bytes); + const length = header_bytes + decoder_bytes + (if (counts.mtp_present) @as(usize, mtp_bytes) else 0) + Sha256.digest_length; + const bytes = try gpa.alloc(u8, length); defer gpa.free(bytes); encode(bytes, identity, counts); var atomic = try Io.Dir.cwd().createFileAtomic(io, path, .{ @@ -71,7 +75,7 @@ fn validateIdentity(identity: Identity) !void { fn encode(bytes: []u8, identity: Identity, counts: *const Counts) void { @memcpy(bytes[0..8], magic); putWord(bytes, 8, format_version); - putWord(bytes, 12, 0); + putWord(bytes, 12, @intFromBool(counts.mtp_present)); @memcpy(bytes[16..48], &identity.model); putWord(bytes, 48, identity.world); putWord(bytes, 52, identity.rank); @@ -82,6 +86,10 @@ fn encode(bytes: []u8, identity: Identity, counts: *const Counts) void { putWord(bytes, offset, count); offset += 4; }; + if (counts.mtp_present) for (&counts.mtp) |count| { + putWord(bytes, offset, count); + offset += 4; + }; var checksum: [Sha256.digest_length]u8 = undefined; Sha256.hash(bytes[0..offset], &checksum, .{}); @memcpy(bytes[offset..], &checksum); @@ -91,21 +99,27 @@ fn decode(bytes: []const u8, identity: Identity) !Counts { if (bytes.len < header_bytes + Sha256.digest_length) return error.TruncatedProfile; if (!std.mem.eql(u8, bytes[0..8], magic)) return error.MalformedProfile; if (getWord(bytes, 8) != format_version) return error.UnsupportedVersion; - if (getWord(bytes, 12) != 0 or getWord(bytes, 56) != decoder_layers or getWord(bytes, 60) != experts) return error.MalformedProfile; - const payload_end = header_bytes + decoder_bytes; - if (bytes.len < max_file_bytes) return error.TruncatedProfile; - if (bytes.len != max_file_bytes) return error.MalformedProfile; + const flags = getWord(bytes, 12); + if (flags > 1 or getWord(bytes, 56) != decoder_layers or getWord(bytes, 60) != experts) return error.MalformedProfile; + const payload_end = header_bytes + decoder_bytes + (if (flags == 1) @as(usize, mtp_bytes) else 0); + const expected_length = payload_end + Sha256.digest_length; + if (bytes.len < expected_length) return error.TruncatedProfile; + if (bytes.len != expected_length) return error.MalformedProfile; var checksum: [Sha256.digest_length]u8 = undefined; Sha256.hash(bytes[0..payload_end], &checksum, .{}); if (!std.mem.eql(u8, bytes[payload_end..], &checksum)) return error.ChecksumMismatch; if (!std.mem.eql(u8, bytes[16..48], &identity.model) or getWord(bytes, 48) != identity.world or getWord(bytes, 52) != identity.rank) return error.IdentityMismatch; - var counts: Counts = .{}; + var counts: Counts = .{ .mtp_present = flags == 1 }; var offset: usize = header_bytes; for (&counts.decoder) |*layer| for (layer) |*count| { count.* = getWord(bytes, offset); offset += 4; }; + if (counts.mtp_present) for (&counts.mtp) |*count| { + count.* = getWord(bytes, offset); + offset += 4; + }; return counts; } diff --git a/zig/src/families/flashnext/exl3/offload/cuda_exl3_frequency_state.zig b/zig/src/families/flashnext/exl3/offload/cuda_exl3_frequency_state.zig index 18578046d8..6331276501 100644 --- a/zig/src/families/flashnext/exl3/offload/cuda_exl3_frequency_state.zig +++ b/zig/src/families/flashnext/exl3/offload/cuda_exl3_frequency_state.zig @@ -26,7 +26,7 @@ pub const State = struct { std.log.warn("EXL3 placement profile rank {d}: save failed ({s})", .{ state.identity.rank, @errorName(err) }); return; }; - std.log.info("EXL3 placement profile rank {d}: saved decoder lifetime placement snapshot to {s}", .{ state.identity.rank, state.path }); + std.log.info("EXL3 placement profile rank {d}: saved decoder/MTP lifetime placement snapshot to {s}", .{ state.identity.rank, state.path }); } }; @@ -36,14 +36,14 @@ pub const Prepared = struct { }; pub fn load(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, own: anytype, world: u32, rank: u32) !?State { - const prepared = try prepare(gpa, io, dir, own.engine.layers[0..own.engine.n_layers], world, rank) orelse return null; + const prepared = try prepare(gpa, io, dir, own.engine.layers[0..own.engine.n_layers], own.head, world, rank) orelse return null; var state = prepared.state; errdefer state.deinit(); if (prepared.counts) |*counts| try finish(own, counts); return state; } -pub fn prepare(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, layers: anytype, world: u32, rank: u32) !?Prepared { +pub fn prepare(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, layers: anytype, head: anytype, world: u32, rank: u32) !?Prepared { const root = std.c.getenv("TF_FLASHNEXT_EXL3_PROFILE_DIR") orelse return null; if (std.mem.span(root).len == 0) return null; if (comptime @import("builtin").os.tag == .windows) { @@ -55,7 +55,7 @@ pub fn prepare(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, layers: anyt return null; } const before = std.Io.Clock.awake.now(io).toNanoseconds(); - const model = fingerprint(gpa, io, dir, layers) catch |err| { + const model = fingerprint(gpa, io, dir, layers, head) catch |err| { std.log.warn("EXL3 placement profile rank {d}: fingerprint unavailable ({s}); persistence disabled", .{ rank, @errorName(err) }); return null; }; @@ -70,8 +70,13 @@ pub fn prepare(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, layers: anyt break :blk null; }; if (counts) |*value| { - decay(value); - std.log.info("EXL3 placement profile rank {d}: loaded and decayed rank-local placement counts before pool allocation", .{rank}); + if (value.mtp_present != (mtpExperts(head) != null)) { + std.log.warn("EXL3 placement profile rank {d}: cold start (MTP schema mismatch)", .{rank}); + counts = null; + } else { + decay(value); + std.log.info("EXL3 placement profile rank {d}: loaded and decayed rank-local placement counts before pool allocation", .{rank}); + } } return .{ .state = state, .counts = counts }; } @@ -83,8 +88,10 @@ pub fn finish(own: anytype, counts: *const codec.Counts) !void { const pool = layer.exl3.?.pool orelse return error.ProfileUnavailable; pool.counts[pool.opt.layer] = counts.decoder[index]; } + if (mtpPool(own.head)) |pool| pool.counts[pool.opt.layer] = counts.mtp; var warmed: usize = 0; for (layers) |layer| warmed += try layer.exl3.?.pool.?.warmProfile(); + if (mtpPool(own.head)) |pool| warmed += try pool.warmProfile(); std.log.info("EXL3 placement profile: imported and warmed {d} expert lines after pool allocation", .{warmed}); } @@ -132,15 +139,21 @@ fn restore(state: *State, own: anytype) !bool { std.log.warn("EXL3 placement profile rank {d}: cold start ({s})", .{ state.identity.rank, @errorName(err) }); return false; }; + if (counts.mtp_present != (mtpPool(own.head) != null)) { + std.log.warn("EXL3 placement profile rank {d}: cold start (MTP schema mismatch)", .{state.identity.rank}); + return false; + } decay(&counts); const layers = own.engine.layers[0..own.engine.n_layers]; for (layers, 0..) |layer, index| { const pool = layer.exl3.?.pool.?; pool.counts[pool.opt.layer] = counts.decoder[index]; } - std.log.info("EXL3 placement profile rank {d}: imported decoder lifetime counts, decayed once by 2 (nonzero minimum 1)", .{state.identity.rank}); + if (mtpPool(own.head)) |pool| pool.counts[pool.opt.layer] = counts.mtp; + std.log.info("EXL3 placement profile rank {d}: imported decoder/MTP lifetime counts, decayed once by 2 (nonzero minimum 1)", .{state.identity.rank}); var warmed: usize = 0; for (layers) |layer| warmed += try layer.exl3.?.pool.?.warmProfile(); + if (mtpPool(own.head)) |pool| warmed += try pool.warmProfile(); std.log.info("EXL3 placement profile rank {d}: warmed {d} expert lines", .{ state.identity.rank, warmed }); return true; } @@ -155,8 +168,20 @@ fn available(layers: anytype) bool { return true; } +fn mtpPool(head: anytype) @TypeOf(head.?.layer.exl3.?.pool) { + const present = head orelse return null; + const experts = present.layer.exl3 orelse return null; + return experts.pool; +} + +fn mtpExperts(head: anytype) @TypeOf(head.?.layer.exl3) { + const present = head orelse return null; + return present.layer.exl3; +} + fn decay(counts: *codec.Counts) void { for (&counts.decoder) |*layer| decayRow(layer); + decayRow(&counts.mtp); } fn decayRow(row: *[512]u32) void { @@ -175,6 +200,11 @@ fn snapshot(own: anytype) !codec.Counts { try pool.lines.drain(); counts.decoder[index] = pool.counts[pool.opt.layer]; } + if (mtpPool(own.head)) |pool| { + try pool.lines.drain(); + counts.mtp = pool.counts[pool.opt.layer]; + counts.mtp_present = true; + } return counts; } @@ -189,7 +219,7 @@ fn hashRouted(hash: *Sha256, routed: []const [3][]const u8) void { for (routed) |expert| for (expert) |projection| hashBytes(hash, projection); } -fn fingerprint(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, layers: anytype) ![32]u8 { +fn fingerprint(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, layers: anytype, head: anytype) ![32]u8 { var hash: Sha256 = .init(.{}); hashBytes(&hash, namespace); for ([_][]const u8{ "config.json", "model.safetensors.index.json" }) |name| { @@ -202,6 +232,10 @@ fn fingerprint(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, layers: anyt } hashBytes(&hash, "decoder-routed-trellis"); for (layers) |layer| hashRouted(&hash, expertRouted(layer.exl3.?)); + if (mtpExperts(head)) |experts| { + hashBytes(&hash, "mtp-routed-trellis"); + hashRouted(&hash, expertRouted(experts)); + } return hash.finalResult(); } @@ -235,18 +269,21 @@ const MockExperts = struct { pool: ?*MockPool, bank: ?[]const u8 = "bank", route const MockLayer = struct { exl3: ?*MockExperts }; const Mock = struct { decoder: [48][512]u32 = @splat(@splat(0)), - pools: [48]MockPool = undefined, - experts: [48]MockExperts = undefined, + mtp: [48][512]u32 = @splat(@splat(0)), + pools: [49]MockPool = undefined, + experts: [49]MockExperts = undefined, + head: struct { layer: MockLayer } = undefined, engine: struct { layers: [48]MockLayer, n_layers: usize = 48 } = undefined, fn init(mock: *Mock) void { mock.* = .{}; mock.engine.n_layers = 48; for (&mock.pools, &mock.experts, 0..) |*pool, *experts, index| { - pool.* = .{ .counts = &mock.decoder, .opt = .{ .layer = @intCast(index) }, .src = .{ .routed = &mock_routed } }; + pool.* = .{ .counts = if (index == 48) &mock.mtp else &mock.decoder, .opt = .{ .layer = if (index == 48) 0 else @intCast(index) }, .src = .{ .routed = &mock_routed } }; experts.* = .{ .pool = pool }; - mock.engine.layers[index] = .{ .exl3 = experts }; + if (index < 48) mock.engine.layers[index] = .{ .exl3 = experts }; } + mock.head = .{ .layer = .{ .exl3 = &mock.experts[48] } }; } const Stream = struct { @@ -257,8 +294,8 @@ const Mock = struct { stream.syncs += 1; } }; - fn owner(mock: *Mock) struct { engine: *@TypeOf(mock.engine), stream: Stream } { - return .{ .engine = &mock.engine, .stream = .{} }; + fn owner(mock: *Mock, drafting: bool) struct { engine: *@TypeOf(mock.engine), head: ?*@TypeOf(mock.head), stream: Stream } { + return .{ .engine = &mock.engine, .head = if (drafting) &mock.head else null, .stream = .{} }; } }; @@ -296,7 +333,7 @@ test "exl3: placement state lifetime lock conflicts, rank separation and private defer reopened.deinit(); } -test "exl3: placement state cold IO never imports or warms; CUDA warm errors propagate" { +test "exl3: placement state cold IO and MTP mismatch never import or warm; CUDA warm errors propagate" { if (comptime @import("builtin").os.tag == .windows) return error.SkipZigTest; const allocator = std.testing.allocator; const io = std.testing.io; @@ -308,13 +345,18 @@ test "exl3: placement state cold IO never imports or warms; CUDA warm errors pro defer state.deinit(); var mock: Mock = undefined; mock.init(); - var own = mock.owner(); + var own = mock.owner(true); try std.testing.expect(!try restore(&state, &own)); try std.Io.Dir.cwd().writeFile(io, .{ .sub_path = state.path, .data = "corrupt" }); try std.testing.expect(!try restore(&state, &own)); var counts: codec.Counts = .{}; counts.decoder[0][1] = 123; try codec.write(allocator, io, state.path, state.identity, &counts); + try std.testing.expect(!try restore(&state, &own)); + try std.testing.expectEqual(@as(u32, 0), mock.decoder[0][1]); + try std.testing.expectEqual(@as(usize, 0), mock.pools[0].warm_calls); + counts.mtp_present = true; + try codec.write(allocator, io, state.path, state.identity, &counts); mock.pools[0].fail_warm = true; try std.testing.expectError(error.MockCudaWarmFailure, restore(&state, &own)); try std.testing.expectEqual(@as(u32, 61), mock.decoder[0][1]); @@ -332,39 +374,44 @@ test "exl3: placement state decay and replacement snapshot do not add imported p defer allocator.free(root); var state = try open(allocator, io, root, test_identity); defer state.deinit(); - var counts: codec.Counts = .{}; + var counts: codec.Counts = .{ .mtp_present = true }; counts.decoder[0][0..5].* = .{ 0, 1, 2, 3, std.math.maxInt(u32) }; + counts.mtp[1] = 9; try codec.write(allocator, io, state.path, state.identity, &counts); var mock: Mock = undefined; mock.init(); - var own = mock.owner(); + var own = mock.owner(true); try std.testing.expect(try restore(&state, &own)); try std.testing.expectEqualSlices(u32, &.{ 0, 1, 1, 1, 2147483647 }, mock.decoder[0][0..5]); + try std.testing.expectEqual(@as(u32, 4), mock.mtp[0][1]); for (mock.pools) |pool| try std.testing.expectEqual(@as(usize, 1), pool.warm_calls); mock.decoder[0][1] += 8; + mock.mtp[0][1] += 2; const saved = try snapshot(&own); try codec.write(allocator, io, state.path, state.identity, &saved); const disk = try codec.read(allocator, io, state.path, state.identity); try std.testing.expectEqual(@as(u32, 9), disk.decoder[0][1]); mock.init(); - own = mock.owner(); + own = mock.owner(true); try std.testing.expect(try restore(&state, &own)); try std.testing.expectEqual(@as(u32, 4), mock.decoder[0][1]); + try std.testing.expectEqual(@as(u32, 3), mock.mtp[0][1]); const missing = try std.fs.path.join(allocator, &.{ root, "missing", "profile" }); defer allocator.free(missing); try std.testing.expectError(error.FileNotFound, codec.write(allocator, io, missing, state.identity, &saved)); } -test "exl3: placement state snapshot synchronizes once and drains decoder counts" { +test "exl3: placement state snapshot synchronizes once and drains decoder and MTP counts" { if (comptime @import("builtin").os.tag == .windows) return error.SkipZigTest; var mock: Mock = undefined; mock.init(); - var own = mock.owner(); + var own = mock.owner(true); for (&mock.pools, 0..) |*pool, index| pool.counts[pool.opt.layer][7] = @intCast(100 + index); const counts = try snapshot(&own); try std.testing.expectEqual(@as(usize, 1), own.stream.syncs); try std.testing.expectEqual(@as(u32, 100), counts.decoder[0][7]); try std.testing.expectEqual(@as(u32, 147), counts.decoder[47][7]); + try std.testing.expectEqual(@as(u32, 148), counts.mtp[7]); for (mock.pools) |pool| try std.testing.expect(pool.lines.drained); own.stream.fail = true; try std.testing.expectError(error.MockCudaSyncFailure, snapshot(&own)); @@ -373,9 +420,11 @@ test "exl3: placement state snapshot synchronizes once and drains decoder counts test "exl3: prepared placement counts import and warm only after pools attach" { var mock: Mock = undefined; mock.init(); - var own = mock.owner(); + var own = mock.owner(true); var counts: codec.Counts = .{}; for (&counts.decoder, 0..) |*row, layer| row[7] = @intCast(layer + 1); + counts.mtp[7] = 73; + counts.mtp_present = true; try finish(&own, &counts); @@ -383,9 +432,11 @@ test "exl3: prepared placement counts import and warm only after pools attach" { try std.testing.expectEqual(@as(u32, @intCast(layer + 1)), mock.decoder[layer][7]); try std.testing.expectEqual(@as(usize, 1), pool.warm_calls); } + try std.testing.expectEqual(@as(u32, 73), mock.mtp[0][7]); + try std.testing.expectEqual(@as(usize, 1), mock.pools[48].warm_calls); } -test "exl3: placement fingerprint reads routed and metadata bytes and ignores pool size" { +test "exl3: placement fingerprint reads routed and metadata bytes, ignores pool size and distinguishes MTP" { if (comptime @import("builtin").os.tag == .windows) return error.SkipZigTest; const allocator = std.testing.allocator; const io = std.testing.io; @@ -397,21 +448,27 @@ test "exl3: placement fingerprint reads routed and metadata bytes and ignores po try temporary.dir.writeFile(io, .{ .sub_path = "model.safetensors.index.json", .data = "index" }); var mock: Mock = undefined; mock.init(); - const original = try fingerprint(allocator, io, root, &mock.engine.layers); + const own = mock.owner(false); + const original = try fingerprint(allocator, io, root, &mock.engine.layers, own.head); for (&mock.experts) |*experts| experts.pool = null; - try std.testing.expectEqual(original, try fingerprint(allocator, io, root, &mock.engine.layers)); + try std.testing.expectEqual(original, try fingerprint(allocator, io, root, &mock.engine.layers, own.head)); for (&mock.experts, &mock.pools) |*experts, *pool| experts.pool = pool; mock.pools[0].opt.lines = 23; - try std.testing.expectEqual(original, try fingerprint(allocator, io, root, &mock.engine.layers)); + try std.testing.expectEqual(original, try fingerprint(allocator, io, root, &mock.engine.layers, own.head)); mock.pools[47].src.routed = &.{.{ "changed gate", "up", "down" }}; - try std.testing.expect(!std.mem.eql(u8, &original, &(try fingerprint(allocator, io, root, &mock.engine.layers)))); + try std.testing.expect(!std.mem.eql(u8, &original, &(try fingerprint(allocator, io, root, &mock.engine.layers, own.head)))); mock.pools[47].src.routed = &.{.{ "gate", "up", "down" }}; try temporary.dir.writeFile(io, .{ .sub_path = "config.json", .data = "changed config" }); - try std.testing.expect(!std.mem.eql(u8, &original, &(try fingerprint(allocator, io, root, &mock.engine.layers)))); + try std.testing.expect(!std.mem.eql(u8, &original, &(try fingerprint(allocator, io, root, &mock.engine.layers, own.head)))); try temporary.dir.writeFile(io, .{ .sub_path = "config.json", .data = "config" }); try temporary.dir.writeFile(io, .{ .sub_path = "model.safetensors.index.json", .data = "changed index" }); - try std.testing.expect(!std.mem.eql(u8, &original, &(try fingerprint(allocator, io, root, &mock.engine.layers)))); + try std.testing.expect(!std.mem.eql(u8, &original, &(try fingerprint(allocator, io, root, &mock.engine.layers, own.head)))); try temporary.dir.writeFile(io, .{ .sub_path = "model.safetensors.index.json", .data = "index" }); + const drafting = mock.owner(true); + const with_mtp = try fingerprint(allocator, io, root, &mock.engine.layers, drafting.head); + try std.testing.expect(!std.mem.eql(u8, &original, &with_mtp)); + mock.pools[48].src.routed = &.{.{ "changed mtp", "up", "down" }}; + try std.testing.expect(!std.mem.eql(u8, &with_mtp, &(try fingerprint(allocator, io, root, &mock.engine.layers, drafting.head)))); } test "exl3: placement lock rejects nonregular storage without waiting for FIFO readers" { diff --git a/zig/src/families/flashnext/exl3/offload/cuda_exl3_frequency_test.zig b/zig/src/families/flashnext/exl3/offload/cuda_exl3_frequency_test.zig index f39617e4a8..7d2b2a0f7a 100644 --- a/zig/src/families/flashnext/exl3/offload/cuda_exl3_frequency_test.zig +++ b/zig/src/families/flashnext/exl3/offload/cuda_exl3_frequency_test.zig @@ -6,6 +6,7 @@ const io = std.testing.io; const identity: profile.Identity = .{ .model = @splat(0xa7), .world = 2, .rank = 1 }; const header_bytes = 64; const checksum_bytes = 32; +const decoder_file_bytes = header_bytes + 48 * 512 * 4 + checksum_bytes; fn profilePath(directory: Io.Dir) ![]u8 { const root = try directory.realPathFileAlloc(io, ".", allocator); @@ -13,19 +14,27 @@ fn profilePath(directory: Io.Dir) ![]u8 { return std.fs.path.join(allocator, &.{ root, "frequency.bin" }); } -fn sample(counts: *profile.Counts) void { - counts.* = .{}; +fn sample(counts: *profile.Counts, mtp_present: bool) void { + counts.* = .{ .mtp_present = mtp_present }; for (&counts.decoder, 0..) |*layer, layer_index| for (layer, 0..) |*count, expert_index| { count.* = @intCast((layer_index + 1) * 100_000 + expert_index); }; + for (&counts.mtp, 0..) |*count, expert_index| count.* = @intCast(7_000_000 + expert_index); counts.decoder[0][0] = 0; counts.decoder[47][511] = std.math.maxInt(u32); + counts.mtp[511] = std.math.maxInt(u32); } fn expectCounts(expected: *const profile.Counts, actual: *const profile.Counts) !void { + try std.testing.expectEqual(expected.mtp_present, actual.mtp_present); for (&expected.decoder, &actual.decoder) |*expected_layer, *actual_layer| { try std.testing.expectEqualSlices(u32, expected_layer, actual_layer); } + if (expected.mtp_present) { + try std.testing.expectEqualSlices(u32, &expected.mtp, &actual.mtp); + } else { + for (actual.mtp) |count| try std.testing.expectEqual(@as(u32, 0), count); + } } fn expectOnlyProfile(directory: Io.Dir) !void { @@ -48,29 +57,31 @@ fn expectPreserved(path: []const u8, before: []const u8) !void { try std.testing.expectEqualSlices(u8, before, after); } -test "exl3: frequency full decoder round trip with private permissions" { +test "exl3: frequency full decoder and optional MTP round trip with private permissions" { if (comptime @import("builtin").os.tag == .windows) return error.SkipZigTest; var temporary = std.testing.tmpDir(.{ .iterate = true }); defer temporary.cleanup(); const path = try profilePath(temporary.dir); defer allocator.free(path); var counts: profile.Counts = .{}; - sample(&counts); - try profile.write(allocator, io, path, identity, &counts); - const restored = try profile.read(allocator, io, path, identity); - try expectCounts(&counts, &restored); - const stat = try temporary.dir.statFile(io, "frequency.bin", .{}); - try std.testing.expectEqual(@as(u64, profile.max_file_bytes), stat.size); - try std.testing.expectEqual(@as(std.posix.mode_t, 0o600), stat.permissions.toMode() & 0o777); - const bytes = try bytesAt(path); - defer allocator.free(bytes); - try std.testing.expectEqualSlices(u8, "TFEXL3FQ", bytes[0..8]); - try std.testing.expectEqualSlices(u8, &identity.model, bytes[16..48]); - try std.testing.expectEqualSlices(u8, &.{ 1, 0, 0, 0 }, bytes[8..12]); - try std.testing.expectEqualSlices(u8, &.{ 0, 0, 0, 0 }, bytes[64..68]); - try std.testing.expectEqualSlices(u8, &.{ 0xa1, 0x86, 0x01, 0x00 }, bytes[68..72]); - try std.testing.expectEqualSlices(u8, &.{ 0xff, 0xff, 0xff, 0xff }, bytes[64 + 48 * 512 * 4 - 4 ..][0..4]); - try expectOnlyProfile(temporary.dir); + for ([_]bool{ false, true, false }) |mtp_present| { + sample(&counts, mtp_present); + try profile.write(allocator, io, path, identity, &counts); + const restored = try profile.read(allocator, io, path, identity); + try expectCounts(&counts, &restored); + const stat = try temporary.dir.statFile(io, "frequency.bin", .{}); + try std.testing.expectEqual(@as(u64, if (mtp_present) profile.max_file_bytes else decoder_file_bytes), stat.size); + try std.testing.expectEqual(@as(std.posix.mode_t, 0o600), stat.permissions.toMode() & 0o777); + const bytes = try bytesAt(path); + defer allocator.free(bytes); + try std.testing.expectEqualSlices(u8, "TFEXL3FQ", bytes[0..8]); + try std.testing.expectEqualSlices(u8, &identity.model, bytes[16..48]); + try std.testing.expectEqualSlices(u8, &.{ 1, 0, 0, 0 }, bytes[8..12]); + try std.testing.expectEqualSlices(u8, &.{ 0, 0, 0, 0 }, bytes[64..68]); + try std.testing.expectEqualSlices(u8, &.{ 0xa1, 0x86, 0x01, 0x00 }, bytes[68..72]); + try std.testing.expectEqualSlices(u8, &.{ 0xff, 0xff, 0xff, 0xff }, bytes[64 + 48 * 512 * 4 - 4 ..][0..4]); + try expectOnlyProfile(temporary.dir); + } } test "exl3: frequency missing profile and every exact identity mismatch are recoverable" { @@ -104,11 +115,11 @@ test "exl3: frequency rejects truncated oversized malformed and corrupt profiles const path = try profilePath(temporary.dir); defer allocator.free(path); var counts: profile.Counts = .{}; - sample(&counts); + sample(&counts, true); try profile.write(allocator, io, path, identity, &counts); const original = try bytesAt(path); defer allocator.free(original); - for ([_]usize{ 0, 7, 63, 64, 95, 96, profile.max_file_bytes - 1 }) |length| { + for ([_]usize{ 0, 7, 63, 64, 95, 96, decoder_file_bytes, profile.max_file_bytes - 1 }) |length| { try temporary.dir.writeFile(io, .{ .sub_path = "frequency.bin", .data = original[0..length] }); try std.testing.expectError(error.TruncatedProfile, profile.read(allocator, io, path, identity)); } @@ -138,6 +149,12 @@ test "exl3: frequency rejects truncated oversized malformed and corrupt profiles oversized[original.len] = 0; try temporary.dir.writeFile(io, .{ .sub_path = "frequency.bin", .data = oversized }); try std.testing.expectError(error.ProfileTooLarge, profile.read(allocator, io, path, identity)); + counts.mtp_present = false; + try profile.write(allocator, io, path, identity, &counts); + const file = try temporary.dir.openFile(io, "frequency.bin", .{ .mode = .read_write }); + defer file.close(io); + try file.writePositionalAll(io, &.{0}, decoder_file_bytes); + try std.testing.expectError(error.MalformedProfile, profile.read(allocator, io, path, identity)); } fn failWrite(userdata: ?*anyopaque, file: Io.File, header: []const u8, data: []const []const u8, splat: usize, offset: u64) Io.File.WritePositionalError!usize { @@ -166,7 +183,7 @@ test "exl3: frequency partial write sync and rename failures preserve previous b try profile.write(allocator, io, path, identity, &counts); const before = try bytesAt(path); defer allocator.free(before); - sample(&counts); + sample(&counts, true); for (0..3) |failure_index| { var failing_vtable = io.vtable.*; const expected: anyerror = switch (failure_index) { @@ -231,7 +248,7 @@ test "exl3: frequency read allocation stays fixed and oversized input needs no a defer allocator.free(path); const single: profile.Identity = .{ .model = @splat(0), .world = 1, .rank = 0 }; var counts: profile.Counts = .{}; - sample(&counts); + sample(&counts, true); try profile.write(allocator, io, path, single, &counts); const backing = try allocator.alloc(u8, profile.max_file_bytes + 1); defer allocator.free(backing); @@ -253,7 +270,7 @@ test "exl3: frequency failed first save leaves no destination or temporary file" const path = try profilePath(temporary.dir); defer allocator.free(path); var counts: profile.Counts = .{}; - sample(&counts); + sample(&counts, true); var failing_vtable = io.vtable.*; failing_vtable.fileSync = failSync; var failing_io = io; diff --git a/zig/src/families/flashnext/exl3/offload/cuda_exl3_pool.zig b/zig/src/families/flashnext/exl3/offload/cuda_exl3_pool.zig index da076ee7e6..647cfce0f9 100644 --- a/zig/src/families/flashnext/exl3/offload/cuda_exl3_pool.zig +++ b/zig/src/families/flashnext/exl3/offload/cuda_exl3_pool.zig @@ -111,7 +111,7 @@ pub const Pool = struct { errdefer gpa.free(k2); var held: usize = 0; for (src.routed, k2) |e, *w| { - // an expert with no source bytes: never picked here + // an expert another rank holds (expert parallel, cuda_exl3_tp.zig): no bytes, never picked here if (e[0].len == 0 and e[1].len == 0 and e[2].len == 0) { w.* = .{ 0, 0, 0 }; continue; @@ -298,7 +298,7 @@ pub const Pool = struct { } fn loadCopy(p: *Pool, c: round_.Copy) !void { - if (p.k2[c.id][0] == 0) return error.Invalid; // an expert without source bytes reached the planner + if (p.k2[c.id][0] == 0) return error.Invalid; // another rank's expert reached the planner unmasked const ev: ?usize = if (c.evicted == offload.Cache.empty_id) null else c.evicted; try p.lines.load(c.line, c.id, p.src.routed[c.id], p.k2[c.id], ev, p.opt.direct, p.io); } diff --git a/zig/src/families/flashnext/exl3/pack/cuda_exl3_config.zig b/zig/src/families/flashnext/exl3/pack/cuda_exl3_config.zig index b4db09d186..8e8a027114 100644 --- a/zig/src/families/flashnext/exl3/pack/cuda_exl3_config.zig +++ b/zig/src/families/flashnext/exl3/pack/cuda_exl3_config.zig @@ -31,6 +31,19 @@ pub fn named(root: std.json.ObjectMap) bool { return std.mem.eql(u8, m, "exl3"); } +/// Whether the checkpoint at `dir` declares an EXL3 pack (cuda_config.mode's rule); false if config.json is unreadable. +pub fn packAt(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) bool { + const path = std.fs.path.join(gpa, &.{ dir, "config.json" }) catch return false; + defer gpa.free(path); + const bytes = std.Io.Dir.cwd().readFileAlloc(io, path, gpa, .limited(1 << 24)) catch return false; + defer gpa.free(bytes); + const parsed = std.json.parseFromSlice(std.json.Value, gpa, bytes, .{}) catch return false; + defer parsed.deinit(); + if (parsed.value != .object) return false; + if (parsed.value.object.get("quantization")) |v| if (v == .object) return false; // an MLX affine pack + return named(parsed.value.object); +} + /// The first field an EXL3 checkpoint differs in from what the family serves, else null. pub fn mismatch(root: std.json.ObjectMap, text: std.json.ObjectMap, comptime geometry: fn (std.json.ObjectMap) ?[]const u8) ?[]const u8 { const q = block(root) orelse return "quantization_config"; diff --git a/zig/src/families/flashnext/exl3/pack/cuda_exl3_load.zig b/zig/src/families/flashnext/exl3/pack/cuda_exl3_load.zig index cec7e8da81..87da7c761f 100644 --- a/zig/src/families/flashnext/exl3/pack/cuda_exl3_load.zig +++ b/zig/src/families/flashnext/exl3/pack/cuda_exl3_load.zig @@ -34,6 +34,8 @@ pub const Ops = struct { sums: ?cuda.DeviceBuffer = null, glue: ?cuda.Module = null, split_fn: cuda.Function = undefined, + /// Expert parallel's mask kernel and block scratch (cuda_exl3_tp.zig), made on first use. + tp: ?*@import("../tp/cuda_exl3_tp.zig").Scratch = null, /// Sizing maxima over what was loaded (`reserve` allocates every scratch for them once, before any round): the max_k: usize = 0, max_z: usize = 0, @@ -68,6 +70,7 @@ pub const Ops = struct { if (o.sums) |*b| b.free(); if (o.picks) |*b| b.free(); if (o.glue) |*m| m.unload(); + if (o.tp) |t| t.destroy(); o.gpa.destroy(o); } @@ -133,11 +136,12 @@ pub const Ops = struct { } /// Every scratch buffer's address and size: unchanged across a pass means nothing grew (nor synchronised) in it. - pub fn fingerprint(o: *const Ops) [8][2]u64 { - var f: [8][2]u64 = @splat(.{ 0, 0 }); + pub fn fingerprint(o: *const Ops) [10][2]u64 { + var f: [10][2]u64 = @splat(.{ 0, 0 }); + if (o.tp) |t| f[8..10].* = @import("../tp/cuda_exl3_tp.zig").fingerprint(t); const de: ?cuda.exl3.Dense = o.dense; const bufs = [8]?cuda.DeviceBuffer{ o.tmp, o.moe, o.halves, o.sums, if (de) |x| x.xh else null, if (de) |x| x.z else null, if (de) |x| x.counters else null, null }; - for (bufs, &f) |b, *e| if (b) |x| { + for (bufs, f[0..8]) |b, *e| if (b) |x| { e.* = .{ x.ptr, x.len }; }; if (o.picks) |h| f[7] = .{ @intFromPtr(h.bytes.ptr), h.bytes.len }; @@ -153,6 +157,7 @@ pub const Ops = struct { if (o.dense) |de| for ([_]?cuda.DeviceBuffer{ de.xh, de.z, de.counters }) |b| if (b) |x| { n += x.len; }; + if (o.tp) |t| n += @import("../tp/cuda_exl3_tp.zig").reservedBytes(t); return n; } @@ -309,14 +314,21 @@ pub fn load(gpa: std.mem.Allocator, d: *cuda.Driver, store: *Store, prefix: []co return wrap(gpa, d, store, prefix, &.{""}, n, k); } -/// Engine.initRank's head for an EXL3 pack: one rank holds the whole vocabulary. +/// Engine.initRank's head for an EXL3 pack: pub fn head(gpa: std.mem.Allocator, d: *cuda.Driver, store: *Store, world: usize, rank: usize) !?projection.Projection { if (!has(store, "language_model.lm_head")) return null; - if (world != 1 or rank != 0) return error.UnsupportedQuantization; - if (try @import("../../cuda_head.zig").enabled()) return error.BadHeadSharding; + if (shardHead(world)) return try @import("../tp/cuda_exl3_tp.zig").headProjection(gpa, d, store, "language_model.lm_head", world, rank); + if (world == 1 and try @import("../../cuda_head.zig").enabled()) return error.BadHeadSharding; return try load(gpa, d, store, "language_model.lm_head", wt.vocab, wt.dims); } +/// Whether an EXL3 head is split at `world`: at 2 unless TF_FLASHNEXT_SHARD_HEAD=0; never at 1. +pub fn shardHead(world: usize) bool { + if (world != 2) return false; + const v = std.c.getenv("TF_FLASHNEXT_SHARD_HEAD") orelse return true; + return !std.mem.eql(u8, std.mem.span(v), "0"); +} + /// Whether any of `prefix`'s parts is an EXL3 matrix (cuda_projection_load.loadStitched's branch). pub fn hasAny(store: *const Store, prefix: []const u8, names: []const []const u8) bool { var buf: [256]u8 = undefined; @@ -359,6 +371,8 @@ pub const Experts = struct { /// Each expert's and projection's trellis words on the device, for a pool to rewrite (gate, up, down; 513 each). trellis: [3][]u64 = .{ &.{}, &.{}, &.{} }, halves: [3][]i32 = .{ &.{}, &.{}, &.{} }, + /// The ids this rank holds under expert parallel (cuda_exl3_tp.zig); null: one rank holds all 513. + tp: ?@import("../tp/cuda_exl3_tp.zig").Rank = null, bank: ?cuda.HostBuffer = null, /// The pool's Source: every routed expert's gate/up/down bytes inside `bank`, and the shared expert's. routed: [][3][]const u8 = &.{}, @@ -443,9 +457,11 @@ fn align256(n: usize) usize { /// Every routed expert of `mlp`, the shared expert as id 512 (the same 2560 x 640 x 2560 shape): resident, or with pub fn experts(gpa: std.mem.Allocator, d: *cuda.Driver, store: *Store, mlp: []const u8, world: usize, rank: usize) !*Experts { - if (world != 1 or rank != 0) return error.UnsupportedQuantization; // one rank holds all 513 + if (world != 1 or rank != 0) return @import("../tp/cuda_exl3_tp.zig").experts(gpa, d, store, mlp, world, rank); // plan step 10 const count = wt.experts_n + 1; - const offload = offloaded(); + // TF_FLASHNEXT_EXL3_MTP_RESIDENT=1: the MTP layer's experts stay on the device (0.88 GiB) instead of a pool + const mtp_resident = std.mem.indexOf(u8, mlp, "mtp.") != null and if (std.c.getenv("TF_FLASHNEXT_EXL3_MTP_RESIDENT")) |v| std.mem.eql(u8, std.mem.span(v), "1") else false; + const offload = offloaded() and !mtp_resident; const x = try gpa.create(Experts); x.* = .{ .gpa = gpa, .ops = try opsOf(store, d), .layer = undefined }; errdefer x.destroy(); diff --git a/zig/src/families/flashnext/exl3/pack/cuda_exl3_memory.zig b/zig/src/families/flashnext/exl3/pack/cuda_exl3_memory.zig index ae55c588b9..22d84a5f3b 100644 --- a/zig/src/families/flashnext/exl3/pack/cuda_exl3_memory.zig +++ b/zig/src/families/flashnext/exl3/pack/cuda_exl3_memory.zig @@ -79,15 +79,51 @@ pub fn deviceBytes(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, o: Optio return total; } -/// The family's deviceWeightBytes for an EXL3 checkpoint; 0 if unread. +/// The MTP head's device bytes: its `mtp.*` tensors placed as the main layers' are (its routed experts under the same +pub fn headBytes(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, o: Options) !u64 { + var store = try wt.Store.open(gpa, io, dir); + defer store.deinit(); + const wm = store.parsed.value.object.get("weight_map") orelse return error.MissingTensor; + var files: std.StringHashMapUnmanaged(void) = .empty; + defer files.deinit(gpa); + var it = wm.object.iterator(); + while (it.next()) |kv| if (std.mem.startsWith(u8, kv.key_ptr.*, "mtp.")) try files.put(gpa, kv.value_ptr.string, {}); + var total: u64 = tables_per_layer + draft_rows * wt.dims * 2; + var fi = files.keyIterator(); + while (fi.next()) |rel| { + const path = try std.fs.path.join(gpa, &.{ dir, rel.* }); + defer gpa.free(path); + var file = try core.safetensors.File.open(gpa, io, path); + defer file.close(io); + var ni = file.names.iterator(); + while (ni.next()) |e| { + const name = e.key_ptr.*; + if (!std.mem.startsWith(u8, name, "mtp.")) continue; + // the head's tensors as a decoder layer's: drop the `mtp.` prefix the host-only rule keys on + var buf: [256]u8 = undefined; + const as_layer = try std.fmt.bufPrint(&buf, "model.language_model.{s}", .{name["mtp.".len..]}); + total = try std.math.add(u64, total, try tensorBytes(as_layer, e.value_ptr.*, o)); + } + } + return total; +} + +/// The draft head's vocabulary rows (cuda_mtp.zig headBytes' count). +const draft_rows: u64 = 79_591; + +/// The family's deviceWeightBytes for an EXL3 checkpoint: the main model and the MTP head; 0 if unread. pub fn weightBytes(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) u64 { - return deviceBytes(gpa, io, dir, options()) catch 0; + const o = options(); + const main = deviceBytes(gpa, io, dir, o) catch return 0; + return main + (headBytes(gpa, io, dir, o) catch return 0); } -/// A rank's placement bytes: one rank holds the whole model (no MTP head: an EXL3 pack decodes plain). -pub fn placement(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, world: usize, rank: usize) !u64 { - if (world != 1 or rank != 0) return error.BadRank; - return deviceBytes(gpa, io, dir, options()); +/// A rank's placement bytes; at world 2 the expert-parallel rank's (cuda_exl3_tp.zig, plan step 10). +pub fn placement(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, world: usize, rank: usize, drafts: bool) !u64 { + if ((world != 1 and world != 2) or rank >= world) return error.BadRank; + if (world != 1) return @import("../tp/cuda_exl3_tp.zig").placement(gpa, io, dir, world, rank, drafts); + const o = options(); + return try deviceBytes(gpa, io, dir, o) + if (drafts) try headBytes(gpa, io, dir, o) else @as(u64, 0); } fn entry(dtype: core.safetensors.DType, shape: []const usize, bytes: u64) Entry { @@ -106,6 +142,7 @@ test "exl3: routed trellis stays on the host under offload, its scales and every try std.testing.expectEqual(tr.end, try tensorBytes("model.language_model.layers.0.mlp.shared_expert.gate_proj.trellis", tr, .{})); try std.testing.expectEqual(tr.end, try tensorBytes("lm_head.trellis", tr, .{})); try std.testing.expectEqual(@as(u64, 0), try tensorBytes("lm_head.mul1", entry(.i32, &.{}, 4), .{})); + try std.testing.expectEqual(@as(u64, 0), try tensorBytes("mtp.layers.0.mlp.experts.3.up_proj.trellis", tr, .{ .routed_on_host = false })); } test "exl3: dense F16 weights count as their BF16 device packing, norms as fp32" { @@ -121,7 +158,7 @@ test "exl3: dense F16 weights count as their BF16 device packing, norms as fp32" try std.testing.expectEqual(@as(u64, 0), try tensorBytes("model.visual.blocks.0.attn.qkv.weight", entry(.bf16, &.{ 3456, 1152 }, 3456 * 1152 * 2), .{})); } -test "exl3: placement names an EXL3 checkpoint and refuses any world but one" { +test "exl3: placement names an EXL3 checkpoint and refuses worlds past 2 (world 2: cuda_exl3_tp.zig's test)" { const a = std.testing.allocator; var tmp = std.testing.tmpDir(.{}); defer tmp.cleanup(); @@ -129,6 +166,5 @@ test "exl3: placement names an EXL3 checkpoint and refuses any world but one" { const dir = try tmp.dir.realPathFileAlloc(std.testing.io, ".", a); defer a.free(dir); try std.testing.expect(isExl3(a, std.testing.io, dir)); - try std.testing.expectError(error.BadRank, placement(a, std.testing.io, dir, 2, 0)); - try std.testing.expectError(error.BadRank, placement(a, std.testing.io, dir, 1, 1)); + try std.testing.expectError(error.BadRank, placement(a, std.testing.io, dir, 3, 0, false)); } diff --git a/zig/src/families/flashnext/exl3/pack/exl3_place.zig b/zig/src/families/flashnext/exl3/pack/exl3_place.zig index 554ccf6881..f2b4d37e96 100644 --- a/zig/src/families/flashnext/exl3/pack/exl3_place.zig +++ b/zig/src/families/flashnext/exl3/pack/exl3_place.zig @@ -4,6 +4,7 @@ const std = @import("std"); const scan = @import("exl3_scan.zig"); const shape = @import("exl3_shape.zig"); const cache = @import("../model/exl3_cache.zig"); +const tp = @import("../tp/exl3_tp.zig"); pub const Home = enum { ngram_host, // PLE rows, page cache or SSD, never a VRAM alloc @@ -186,8 +187,8 @@ test "a tensor name picks one home and the host banks stay off device" { try std.testing.expect(hostBank(.host)); } -/// Where the bytes live. -pub const Stay = enum { vram_pinned, vram_learned, host, disk, skip }; +/// Where the bytes for this rank live. +pub const Stay = enum { vram_pinned, vram_learned, host, disk, other_rank, skip }; /// Integer after `.experts.`, or null when the name is not a routed expert. pub fn expertId(name: []const u8) ?u16 { @@ -204,26 +205,31 @@ pub fn expertId(name: []const u8) ?u16 { return @intCast(n); } -pub fn stay(which: Home, id: ?u16) error{IllegalCut}!Stay { +pub fn stay(which: Home, id: ?u16, rank: usize, ranks: usize) error{ IllegalCut, RankOutOfRange }!Stay { + if (ranks != 1 and ranks != 2) return error.IllegalCut; + if (rank >= ranks) return error.RankOutOfRange; switch (which) { .ngram_host => return .disk, .vision_skip => return .skip, .routed_host => { - _ = id orelse return error.IllegalCut; - return .vram_learned; + const expert = id orelse return error.IllegalCut; + const span = try tp.range(.routed_set, rank, ranks); + if (@as(usize, expert) >= span.lo and @as(usize, expert) < span.hi) return .vram_learned; + return .other_rank; }, .plain => return .host, .shared_resident, .resident, .centred_norm, .mtp_patch => return .vram_pinned, } } -pub fn stayOf(t: Tensor) error{IllegalCut}!Stay { - return stay(home(t), expertId(t.name)); +pub fn stayOf(t: Tensor, rank: usize, ranks: usize) error{ IllegalCut, RankOutOfRange }!Stay { + return stay(home(t), expertId(t.name), rank, ranks); } -test "the decode path stays pinned, the n-gram on disk, and the routed experts are learned" { +test "tp 2 pins the decode path, leaves the n-gram on disk, and learns only this rank's experts" { const ngram = "model.language_model.layers.2.ple.ngram_embedding.shard_0.trellis"; - try std.testing.expectEqual(Stay.disk, try stayOf(.{ .name = ngram, .rank = 2 })); + try std.testing.expectEqual(Stay.disk, try stayOf(.{ .name = ngram, .rank = 2 }, 0, 2)); + try std.testing.expectEqual(Stay.disk, try stayOf(.{ .name = ngram, .rank = 2 }, 1, 2)); const pinned = [_][]const u8{ "model.layers.0.mlp.shared_expert.up_proj.trellis", @@ -246,24 +252,25 @@ test "the decode path stays pinned, the n-gram on disk, and the routed experts a try std.testing.expectEqual(Stay.host, try stayOf(.{ .name = "model.embed_tokens.weight", .rank = 2, - })); - try std.testing.expectEqual(Stay.host, try stayOf(.{ .name = "readme", .rank = 0 })); + }, 0, 2)); + try std.testing.expectEqual(Stay.host, try stayOf(.{ .name = "readme", .rank = 0 }, 1, 2)); for (pinned) |name| { - try std.testing.expectEqual(Stay.vram_pinned, try stayOf(.{ .name = name, .rank = 2 })); + try std.testing.expectEqual(Stay.vram_pinned, try stayOf(.{ .name = name, .rank = 2 }, 0, 2)); + try std.testing.expectEqual(Stay.vram_pinned, try stayOf(.{ .name = name, .rank = 2 }, 1, 2)); } try std.testing.expectEqual(Stay.vram_pinned, try stayOf(.{ .name = "model.layers.0.self_attn.q_norm.weight", .rank = 1, - })); + }, 0, 2)); try std.testing.expectEqual(Stay.vram_pinned, try stayOf(.{ .name = "weight", .rank = 1, .file = "mtp_hyper_connection_mixer_patch.safetensors", - })); + }, 1, 2)); try std.testing.expectEqual(Stay.skip, try stayOf(.{ .name = "model.visual.blocks.0.weight", .rank = 1, - })); + }, 0, 2)); try std.testing.expectEqual(@as(?u16, 0), expertId("model.layers.0.mlp.experts.0.up_proj.trellis")); try std.testing.expectEqual(@as(?u16, 255), expertId("model.layers.0.mlp.experts.255.down_proj.trellis")); @@ -272,10 +279,15 @@ test "the decode path stays pinned, the n-gram on disk, and the routed experts a const low = "model.layers.0.mlp.experts.0.up_proj.trellis"; const high = "model.layers.0.mlp.experts.400.up_proj.trellis"; - try std.testing.expectEqual(Stay.vram_learned, try stayOf(.{ .name = low, .rank = 3 })); - try std.testing.expectEqual(Stay.vram_learned, try stayOf(.{ .name = high, .rank = 3 })); + try std.testing.expectEqual(Stay.vram_learned, try stayOf(.{ .name = low, .rank = 3 }, 0, 2)); + try std.testing.expectEqual(Stay.other_rank, try stayOf(.{ .name = low, .rank = 3 }, 1, 2)); + try std.testing.expectEqual(Stay.other_rank, try stayOf(.{ .name = high, .rank = 3 }, 0, 2)); + try std.testing.expectEqual(Stay.vram_learned, try stayOf(.{ .name = high, .rank = 3 }, 1, 2)); + try std.testing.expectEqual(Stay.vram_learned, try stayOf(.{ .name = high, .rank = 3 }, 0, 1)); - try std.testing.expectError(error.IllegalCut, stay(.routed_host, null)); + try std.testing.expectError(error.IllegalCut, stayOf(.{ .name = low, .rank = 3 }, 0, 4)); + try std.testing.expectError(error.RankOutOfRange, stayOf(.{ .name = low, .rank = 3 }, 2, 2)); + try std.testing.expectError(error.IllegalCut, stay(.routed_host, null, 0, 2)); try std.testing.expect(std.mem.indexOf(u8, "key_proj", "k_proj") == null); try std.testing.expect(std.mem.indexOf(u8, "value_proj", "v_proj") == null); try std.testing.expect(std.mem.indexOf(u8, "mlp.shared_expert_gate", "mlp.gate") == null); @@ -289,21 +301,22 @@ test "mlp.gate, key_proj, and value_proj stay pinned, and so does mlp.router.wei const shared_gate = "model.layers.0.mlp.shared_expert_gate.weight"; for ([_][]const u8{ gate, key, value, router }) |name| { try std.testing.expectEqual(Home.resident, home(.{ .name = name, .rank = 2 })); - try std.testing.expectEqual(Stay.vram_pinned, try stayOf(.{ .name = name, .rank = 2 })); + try std.testing.expectEqual(Stay.vram_pinned, try stayOf(.{ .name = name, .rank = 2 }, 0, 2)); + try std.testing.expectEqual(Stay.vram_pinned, try stayOf(.{ .name = name, .rank = 2 }, 1, 2)); } // `shared_expert` matches before the resident tokens. This is not the 2560×640 expert. try std.testing.expectEqual(Home.shared_resident, home(.{ .name = shared_gate, .rank = 2 })); - try std.testing.expectEqual(Stay.vram_pinned, try stayOf(.{ .name = shared_gate, .rank = 1 })); + try std.testing.expectEqual(Stay.vram_pinned, try stayOf(.{ .name = shared_gate, .rank = 1 }, 0, 2)); const switch_gate = "model.layers.0.mlp.switch_mlp.gate_proj.weight"; try std.testing.expectEqual(Home.plain, home(.{ .name = switch_gate, .rank = 2 })); - try std.testing.expectEqual(Stay.host, try stayOf(.{ .name = switch_gate, .rank = 2 })); + try std.testing.expectEqual(Stay.host, try stayOf(.{ .name = switch_gate, .rank = 2 }, 0, 2)); try std.testing.expectEqual(Home.routed_host, home(.{ .name = "model.layers.0.mlp.experts.0.gate_proj.trellis", .rank = 3, })); try std.testing.expectEqual(Home.plain, home(.{ .name = "mtp.fc_embedding.weight", .rank = 2 })); - try std.testing.expectEqual(Stay.host, try stayOf(.{ .name = "mtp.fc_embedding.weight", .rank = 2 })); + try std.testing.expectEqual(Stay.host, try stayOf(.{ .name = "mtp.fc_embedding.weight", .rank = 2 }, 0, 2)); try std.testing.expectEqual(Home.resident, home(.{ .name = "language_model.mtp.fc_hidden.weight", .rank = 2 })); - try std.testing.expectEqual(Stay.vram_pinned, try stayOf(.{ .name = "language_model.mtp.fc_hidden.weight", .rank = 2 })); + try std.testing.expectEqual(Stay.vram_pinned, try stayOf(.{ .name = "language_model.mtp.fc_hidden.weight", .rank = 2 }, 0, 2)); } diff --git a/zig/src/families/flashnext/exl3/pack/exl3_shape.zig b/zig/src/families/flashnext/exl3/pack/exl3_shape.zig index ec97e5d1df..a488ca4e8b 100644 --- a/zig/src/families/flashnext/exl3/pack/exl3_shape.zig +++ b/zig/src/families/flashnext/exl3/pack/exl3_shape.zig @@ -37,6 +37,7 @@ pub const codebook = "mul1"; pub const tile: usize = 16; pub const hadamard: usize = 128; +pub const tp_ranks: usize = 2; pub const linear_layers_per_cycle: usize = 3; pub const full_attention_interval: usize = 4; @@ -56,11 +57,14 @@ pub fn kvOut() usize { return kv_heads * head_dim; } -test "the published 3.05 shape is the one the budget is written against" { +test "the published 3.05 shape is the one the cuts and the budget are written against" { try std.testing.expectEqual(@as(usize, 16), ngramHeadCount()); try std.testing.expectEqual(@as(usize, 6144), qOut()); try std.testing.expectEqual(@as(usize, 512), kvOut()); try std.testing.expect(hidden % hadamard == 0); try std.testing.expect(expert_intermediate % hadamard == 0); try std.testing.expect(qOut() % hadamard == 0); + // Half of 640 is 320, which is not a Hadamard block. Column TP inside one expert is illegal. + try std.testing.expect(expert_intermediate / tp_ranks % hadamard != 0); + try std.testing.expect(hidden / tp_ranks % hadamard == 0); } diff --git a/zig/src/families/flashnext/exl3/tp/cuda_exl3_mtp.zig b/zig/src/families/flashnext/exl3/tp/cuda_exl3_mtp.zig new file mode 100644 index 0000000000..000d8adba8 --- /dev/null +++ b/zig/src/families/flashnext/exl3/tp/cuda_exl3_mtp.zig @@ -0,0 +1,266 @@ +//! An EXL3 pack's MTP head on v1.1.0's Flash Next family (cuda_mtp.zig): the draft vocabulary rows here. + +const std = @import("std"); +const core = @import("core"); +const cuda = @import("cuda"); +const wt = @import("../../cuda_prompt_wt.zig"); +const projection = @import("../../cuda_projection.zig"); +const fmt = core.exl3_format; + +const block_n: usize = fmt.hadamard_dim; + +/// lm_head as stored: trellis int16 [K/16, N/16, 8 * halves] (bytes, possibly unaligned), suh fp16 [K], svh fp16 [N]. +pub const Head = struct { + trellis: []const u8, + suh: []const u8, + svh: []const u8, + k: usize, + n: usize, + halves: u16, + + pub fn words(h: Head) usize { + return @as(usize, h.halves) * 8; + } + + /// From the stored tensors; `mul1` (when present) must hold the mul1 multiplier. + pub fn of(trellis: core.safetensors.Tensor, suh: core.safetensors.Tensor, svh: core.safetensors.Tensor, mul1: ?core.safetensors.Tensor) !Head { + if (trellis.dtype != .i16 or trellis.rank != 3 or trellis.dim(2) % 8 != 0) return error.UnexpectedTensor; + const h: Head = .{ .trellis = trellis.bytes, .suh = suh.bytes, .svh = svh.bytes, .k = trellis.dim(0) * 16, .n = trellis.dim(1) * 16, .halves = @intCast(trellis.dim(2) / 8) }; + if (suh.dtype != .f16 or suh.numel() != h.k or svh.dtype != .f16 or svh.numel() != h.n) return error.UnexpectedTensor; + if (h.k % block_n != 0 or h.n % block_n != 0 or h.halves < 2 or h.halves > 16) return error.UnexpectedTensor; + if (mul1) |m| { + if (m.bytes.len != 4 or std.mem.readInt(u32, m.bytes[0..4], .little) != fmt.Mul1Mul) return error.UnexpectedTensor; + } + return h; + } +}; + +/// f64 -> bf16 bits, one rounding to nearest even (no f32 step, so no double rounding). +pub fn bf16Of(v: f64) u16 { + const near = wt.toBf16(@floatCast(v)); + if (std.math.isNan(v) or std.math.isInf(v)) return near; + var best = near; + var best_d = @abs(@as(f64, wt.promote(near)) - v); + for ([_]u16{ near -% 1, near +% 1 }) |cand| { + const cv: f64 = wt.promote(cand); + if (std.math.isNan(cv) or std.math.isInf(cv)) continue; + if (std.math.signbit(cv) != std.math.signbit(@as(f64, wt.promote(near))) and cv != 0) continue; + const d = @abs(cv - v); + if (d < best_d or (d == best_d and cand & 1 == 0)) { + best = cand; + best_d = d; + } + } + return best; +} + +const Work = struct { + head: Head, + ids: []const u32, + /// Indices into `ids`, sorted by id. + order: []const u32, + /// Start of each block's run in `order`; `blocks` entries plus one. + starts: []const u32, + out: []u16, + next: std.atomic.Value(usize) = .init(0), + failed: ?anyerror = null, +}; + +fn copyF16(bytes: []const u8, out: []f16) void { + for (out, 0..) |*o, i| o.* = @bitCast(std.mem.readInt(u16, bytes[2 * i ..][0..2], .little)); +} + +fn worker(w: *Work, gpa: std.mem.Allocator) void { + work(w, gpa) catch |err| { + w.failed = err; + }; +} + +fn work(w: *Work, gpa: std.mem.Allocator) !void { + const h = w.head; + const tiles_k = h.k / 16; + const tiles_n = h.n / 16; + const per = block_n / 16; + const words = h.words(); + const sub = try gpa.alloc(i16, tiles_k * per * words); + defer gpa.free(sub); + const vals = try gpa.alloc(f64, h.k * block_n); + defer gpa.free(vals); + const scratch = try gpa.alloc(f64, 2 * h.k * block_n); + defer gpa.free(scratch); + const suh = try gpa.alloc(f16, h.k); + defer gpa.free(suh); + copyF16(h.suh, suh); + var svh: [block_n]f16 = undefined; + while (true) { + const b = w.next.fetchAdd(1, .monotonic); + if (b + 1 >= w.starts.len) return; + const lo = w.starts[b]; + const hi = w.starts[b + 1]; + if (lo == hi) continue; + const nb = w.ids[w.order[lo]] / block_n; + // the block's tiles, k-major as a [tiles_k, per] trellis of their own + for (0..tiles_k) |tk| for (0..per) |j| { + const src = ((tk * tiles_n) + nb * per + j) * words * 2; + const dst = sub[(tk * per + j) * words ..][0..words]; + for (dst, 0..) |*d, i| d.* = std.mem.readInt(i16, h.trellis[src + 2 * i ..][0..2], .little); + }; + copyF16(h.svh[nb * block_n * 2 ..][0 .. block_n * 2], &svh); + fmt.dequantize(.{ .halves = h.halves }, .mul1, sub, tiles_k, per, suh, &svh, vals, scratch); + for (w.order[lo..hi]) |i| { + const col = w.ids[i] - nb * block_n; + const row = w.out[@as(usize, i) * h.k ..][0..h.k]; + for (row, 0..) |*o, k| o.* = bf16Of(vals[k * block_n + col]); + } + } +} + +/// Row i of `out` ([ids.len, K] bf16) is lm_head's column `ids[i]` (the weight row of token ids[i]). +pub fn draftRows(gpa: std.mem.Allocator, head: Head, ids: []const u32, out: []u16, threads: usize) !void { + if (out.len != ids.len * head.k) return error.UnexpectedTensor; + for (ids) |id| if (id >= head.n) return error.InvalidToken; + const order = try gpa.alloc(u32, ids.len); + defer gpa.free(order); + for (order, 0..) |*o, i| o.* = @intCast(i); + std.mem.sort(u32, order, ids, struct { + fn lt(x: []const u32, a: u32, b: u32) bool { + return x[a] < x[b]; + } + }.lt); + // runs of one block each + var starts: std.ArrayList(u32) = .empty; + defer starts.deinit(gpa); + var prev: ?usize = null; + for (order, 0..) |i, at| { + const nb = ids[i] / block_n; + if (prev == null or prev.? != nb) try starts.append(gpa, @intCast(at)); + prev = nb; + } + try starts.append(gpa, @intCast(order.len)); + var w: Work = .{ .head = head, .ids = ids, .order = order, .starts = starts.items, .out = out }; + const n_threads = @max(1, @min(threads, starts.items.len - 1)); + if (n_threads == 1) { + try work(&w, gpa); + return; + } + var pool: [64]std.Thread = undefined; + var spawned: usize = 0; + defer for (pool[0..spawned]) |t| t.join(); + for (0..@min(n_threads, pool.len)) |_| { + pool[spawned] = try std.Thread.spawn(.{}, worker, .{ &w, gpa }); + spawned += 1; + } + for (pool[0..spawned]) |t| t.join(); + spawned = 0; + if (w.failed) |err| return err; +} + +/// Threads for a load: the machine's cores, at most 32 (each holds about 8 MB of float64 scratch). +pub fn loadThreads() usize { + return @min(32, std.Thread.getCpuCount() catch 1); +} + +/// The draft head of an EXL3 pack (null for any other): bf16 rows uploaded as a dense projection, or affine4 rows. +pub fn draft(a: std.mem.Allocator, d: *cuda.Driver, store: *wt.Store, ids: []const u32, affine: bool) !?wt.Q4 { + if (!store.has("lm_head.trellis")) return null; + const head = try Head.of(try store.hold("lm_head.trellis"), try store.hold("lm_head.suh"), try store.hold("lm_head.svh"), if (store.has("lm_head.mul1")) try store.hold("lm_head.mul1") else null); + if (head.k != wt.dims or head.n != wt.vocab) return error.UnexpectedTensor; + const rows = try a.alloc(u16, ids.len * wt.dims); + defer a.free(rows); + try draftRows(a, head, ids, rows, loadThreads()); + const bytes = std.mem.sliceAsBytes(rows); + if (!affine) return try projection.dense(a, d, .{ .dtype = .bf16, .rank = 2, .shape = .{ ids.len, wt.dims, 1, 1, 1 }, .bytes = bytes }); + const words = try a.alloc(u8, ids.len * wt.dims / 2); + defer a.free(words); + const scales = try a.alloc(u8, ids.len * wt.dims / 16); + defer a.free(scales); + const biases = try a.alloc(u8, scales.len); + defer a.free(biases); + for (0..ids.len) |i| try @import("../../cuda_draft_affine.zig").row(bytes[i * wt.dims * 2 ..][0 .. wt.dims * 2], words[i * wt.dims / 2 ..][0 .. wt.dims / 2], scales[i * wt.dims / 16 ..][0 .. wt.dims / 16], biases[i * wt.dims / 16 ..][0 .. wt.dims / 16]); + return try wt.uploadQ4(a, d, words, scales, biases, ids.len, wt.dims / 8); +} + +// ---- tests -------------------------------------------------------------------------------------------------------- + +test "exl3 draft: bf16Of rounds float64 once, to nearest even" { + try std.testing.expectEqual(@as(u16, 0x3F80), bf16Of(1.0)); + try std.testing.expectEqual(@as(u16, 0x3F80), bf16Of(1.0 + 1.0 / 256.0)); // tie: even (down) + try std.testing.expectEqual(@as(u16, 0x3F82), bf16Of(1.0 + 3.0 / 256.0)); // tie: even (up) + // just above a tie after an f32 step would round to the tie: float64 sees it above + const above = 1.0 + 1.0 / 256.0 + 1.0e-12; + try std.testing.expectEqual(@as(u16, 0x3F81), bf16Of(above)); + try std.testing.expectEqual(@as(u16, 0xBF80), bf16Of(-1.0)); + try std.testing.expectEqual(@as(u16, 0), bf16Of(0)); +} + +fn synthHead(gpa: std.mem.Allocator, r: std.Random, k: usize, n: usize, halves: u16) !struct { t: []i16, su: []f16, sv: []f16, head: Head } { + const t = try gpa.alloc(i16, (k / 16) * (n / 16) * halves * 8); + for (t) |*x| x.* = @bitCast(r.int(u16)); + const su = try gpa.alloc(f16, k); + for (su) |*x| x.* = @floatCast((r.float(f32) * 0.02 + 0.005) * (if (r.boolean()) @as(f32, 1) else -1)); + const sv = try gpa.alloc(f16, n); + for (sv) |*x| x.* = @floatCast((r.float(f32) * 0.5 + 0.5) * (if (r.boolean()) @as(f32, 1) else -1)); + return .{ .t = t, .su = su, .sv = sv, .head = .{ .trellis = std.mem.sliceAsBytes(t), .suh = std.mem.sliceAsBytes(su), .svh = std.mem.sliceAsBytes(sv), .k = k, .n = n, .halves = halves } }; +} + +test "exl3 draft: rows equal the whole matrix's float64 decode rounded to bf16, any id order, any thread count" { + const gpa = std.testing.allocator; + var prng = std.Random.DefaultPrng.init(0xd4af); + const r = prng.random(); + const k = 256; + const n = 384; + const s = try synthHead(gpa, r, k, n, 10); + defer gpa.free(s.t); + defer gpa.free(s.su); + defer gpa.free(s.sv); + const full = try gpa.alloc(f64, k * n); + defer gpa.free(full); + const scratch = try gpa.alloc(f64, 2 * k * n); + defer gpa.free(scratch); + fmt.dequantize(.{ .halves = 10 }, .mul1, s.t, k / 16, n / 16, s.su, s.sv, full, scratch); + const ids = [_]u32{ 300, 5, 127, 128, 383, 5, 0, 200, 255, 256 }; + for ([_]usize{ 1, 3 }) |threads| { + var rows: [ids.len * k]u16 = undefined; + try draftRows(gpa, s.head, &ids, &rows, threads); + for (ids, 0..) |id, i| for (0..k) |kk| { + try std.testing.expectEqual(bf16Of(full[kk * n + id]), rows[i * k + kk]); + }; + } + var bad: [k]u16 = undefined; + try std.testing.expectError(error.InvalidToken, draftRows(gpa, s.head, &.{384}, &bad, 1)); +} + +test "exl3 draft: a draft logit is the full head's within the bf16 weight bound" { + const gpa = std.testing.allocator; + var prng = std.Random.DefaultPrng.init(0x10917); + const r = prng.random(); + const k = 256; + const n = 256; + const s = try synthHead(gpa, r, k, n, 6); + defer gpa.free(s.t); + defer gpa.free(s.su); + defer gpa.free(s.sv); + const full = try gpa.alloc(f64, k * n); + defer gpa.free(full); + const scratch = try gpa.alloc(f64, 2 * k * n); + defer gpa.free(scratch); + fmt.dequantize(.{ .halves = 6 }, .mul1, s.t, k / 16, n / 16, s.su, s.sv, full, scratch); + var x: [k]f64 = undefined; + for (&x) |*v| v.* = (r.float(f64) - 0.5) * 4; + var ids: [n]u32 = undefined; + for (&ids, 0..) |*id, i| id.* = @intCast(i); + const rows = try gpa.alloc(u16, n * k); + defer gpa.free(rows); + try draftRows(gpa, s.head, &ids, rows, 2); + for (0..n) |v| { + var exact: f64 = 0; + var mag: f64 = 0; + var got: f64 = 0; + for (0..k) |kk| { + exact += x[kk] * full[kk * n + v]; + mag += @abs(x[kk] * full[kk * n + v]); + got += x[kk] * @as(f64, wt.promote(rows[v * k + kk])); + } + try std.testing.expect(@abs(got - exact) <= mag / 512.0 + 1e-12); + } +} diff --git a/zig/src/families/flashnext/exl3/tp/cuda_exl3_tp.zig b/zig/src/families/flashnext/exl3/tp/cuda_exl3_tp.zig new file mode 100644 index 0000000000..9cbd489f57 --- /dev/null +++ b/zig/src/families/flashnext/exl3/tp/cuda_exl3_tp.zig @@ -0,0 +1,500 @@ +//! Two-rank expert parallel for an EXL3 Flash Next checkpoint (plan step 10 of docs/plans/exl3-on-upstream-1.1.0.md) + +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const wt = @import("../../cuda_prompt_wt.zig"); +const load = @import("../pack/cuda_exl3_load.zig"); +const memory = @import("../pack/cuda_exl3_memory.zig"); +const projection = @import("../../cuda_projection.zig"); +const cuts = @import("exl3_tp.zig"); +const canon = @import("exl3_canon.zig"); + +const dims = wt.dims; +const slots = wt.slots; +const width = wt.moe_w; +/// Table entries: 512 routed experts and the shared one as id 512. +const table_n = wt.experts_n + 1; +/// The rank that holds the shared expert (id 512). +pub const shared_rank: usize = 1; +/// Rows an `Experts.run` call takes (cuda_exl3_pass.zig's block). +pub const moe_rows: usize = @import("../model/cuda_exl3_pass.zig").moe_rows; + +/// The ids [lo, hi) rank `rank` of `world` holds. +pub const Rank = struct { + world: usize, + rank: usize, + lo: usize, + hi: usize, + + pub fn init(world: usize, rank: usize) !Rank { + if ((world != 1 and world != 2) or rank >= world) return error.BadRank; + if (world == 1) return .{ .world = 1, .rank = 0, .lo = 0, .hi = table_n }; + const r = cuts.range(.routed_set, rank, world) catch return error.BadRank; + return .{ .world = world, .rank = rank, .lo = r.lo, .hi = if (rank == shared_rank) r.hi + 1 else r.hi }; + } + + pub fn owns(r: Rank, id: usize) bool { + return id >= r.lo and id < r.hi; + } +}; + +// ---- loading ----------------------------------------------------------------------------------------------------- + +fn up(x: *load.Experts, d: *const cuda.Driver, bytes: []const u8) !u64 { + try x.bufs.ensureUnusedCapacity(x.gpa, 1); + const b = try cuda.DeviceBuffer.fromHost(d, bytes); + x.bufs.appendAssumeCapacity(b); + return b.ptr; +} + +fn align256(n: usize) usize { + return (n + 255) / 256 * 256; +} + +fn expertName(b: []u8, mlp: []const u8, e: usize, proj: []const u8) ![]const u8 { + if (e == wt.experts_n) return std.fmt.bufPrint(b, "{s}.shared_expert.{s}", .{ mlp, proj }); + return std.fmt.bufPrint(b, "{s}.experts.{d}.{s}", .{ mlp, e, proj }); +} + +/// cuda_exl3_load.experts at world 2: the rank's experts, the 513-entry tables (other ids zero), the scales' rows +pub fn experts(gpa: std.mem.Allocator, d: *cuda.Driver, store: *wt.Store, mlp: []const u8, world: usize, rank: usize) !*load.Experts { + const r = try Rank.init(world, rank); + const offload = load.offloaded(); + const x = try gpa.create(load.Experts); + x.* = .{ .gpa = gpa, .ops = try load.opsOf(store, d), .layer = undefined, .tp = r }; + errdefer x.destroy(); + for (&x.trellis, &x.halves) |*t, *h| { + t.* = try gpa.alloc(u64, table_n); + @memset(t.*, 0); + h.* = try gpa.alloc(i32, table_n); + @memset(h.*, 0); + } + const projs = [3][]const u8{ "gate_proj", "up_proj", "down_proj" }; + const ks = [3]usize{ dims, dims, width }; + const ns = [3]usize{ width, width, dims }; + var scales: [6]std.ArrayList(u8) = @splat(.empty); + defer for (&scales) |*l| l.deinit(gpa); + var buf: [256]u8 = undefined; + var at: usize = 0; + if (offload) { + // the bank's size from the headers of the rank's own experts, then one pinned region for the layer + var total: usize = 0; + for (r.lo..r.hi) |e| { + for (projs) |proj| total += align256((try load.mat(store, try expertName(&buf, mlp, e, proj))).trellis.len); + store.release(); + } + x.bank = try cuda.HostBuffer.alloc(d, total); + x.routed = try gpa.alloc([3][]const u8, wt.experts_n); + for (x.routed) |*e| e.* = .{ &.{}, &.{}, &.{} }; + } + const t0 = std.Io.Clock.awake.now(store.io).toNanoseconds(); + for (r.lo..r.hi) |e| { + for (projs, 0..) |proj, j| { + const m = try load.mat(store, try expertName(&buf, mlp, e, proj)); + if (m.k != ks[j] or m.n != ns[j]) return error.UnexpectedTensor; + if (x.bank) |bank| { + const dst = bank.bytes[at..][0..m.trellis.len]; + @memcpy(dst, m.trellis); + at += align256(m.trellis.len); + if (e == wt.experts_n) x.shared[j] = dst else x.routed[e][j] = dst; + } else x.trellis[j][e] = try up(x, d, m.trellis); + x.halves[j][e] = @intCast(m.halves); + try scales[2 * j].appendSlice(gpa, m.suh); + try scales[2 * j + 1].appendSlice(gpa, m.svh); + } + store.release(); + } + x.bank_ns = @intCast(std.Io.Clock.awake.now(store.io).toNanoseconds() - t0); + // a scale row's fp16 count: suh_g D, svh_g I, suh_u D, svh_u I, suh_d I, svh_d D + const rows = [6]usize{ dims, width, dims, width, width, dims }; + var base: [6]u64 = undefined; + for (&base, scales, rows) |*b, l, n| b.* = try up(x, d, l.items) - r.lo * n * 2; + x.layer = .{ + // offloaded: the pool's tables replace these at attach (0 until then; a pass refuses an unattached layer) + .gate = if (offload) 0 else try up(x, d, std.mem.sliceAsBytes(x.trellis[0])), + .up = if (offload) 0 else try up(x, d, std.mem.sliceAsBytes(x.trellis[1])), + .down = if (offload) 0 else try up(x, d, std.mem.sliceAsBytes(x.trellis[2])), + .gate_k2 = if (offload) 0 else try up(x, d, std.mem.sliceAsBytes(x.halves[0])), + .up_k2 = if (offload) 0 else try up(x, d, std.mem.sliceAsBytes(x.halves[1])), + .down_k2 = if (offload) 0 else try up(x, d, std.mem.sliceAsBytes(x.halves[2])), + .suh_g = base[0], + .svh_g = base[1], + .suh_u = base[2], + .svh_u = base[3], + .suh_d = base[4], + .svh_d = base[5], + .experts = table_n, + .dims = dims, + .width = width, + .codebook = .mul1, + .act_mode = 1, + }; + return x; +} + +// ---- the pass ---------------------------------------------------------------------------------------------------- + +/// The engine's expert-parallel scratch and kernels (held by the shared cuda_exl3_load.Ops, released with it). +pub const Scratch = struct { + gpa: std.mem.Allocator, + d: *const cuda.Driver, + mod: cuda.Module, + mask_fn: cuda.Function, + /// The rank's masked picks [rows, 11] int32, then the finish's piece table and the all-gather buffer (gatherAt). + picks: ?cuda.DeviceBuffer = null, + rows: usize = 0, + world: usize = 1, + /// Experts.run's scratch for `moe_cap` rows, zeroed when allocated (a skipped slot reads its stale, finite y). + moe: ?cuda.DeviceBuffer = null, + moe_cap: usize = 0, + + fn create(gpa: std.mem.Allocator, d: *const cuda.Driver) !*Scratch { + if (!cuda.kernels.available) return error.BuiltWithoutKernels; + var m = try cuda.Module.load(d, cuda.kernels.flashnext_exl3_tp); + errdefer m.unload(); + const t = try gpa.create(Scratch); + t.* = .{ .gpa = gpa, .d = d, .mod = m, .mask_fn = try m.function("fn_exl3_tp_mask") }; + return t; + } + + pub fn destroy(t: *Scratch) void { + inline for (.{ "picks", "moe" }) |f| if (@field(t, f)) |*b| b.free(); + t.mod.unload(); + t.gpa.destroy(t); + } + + /// Room for `rows` rows; a growth waits for `s` (an earlier call may still read the old buffers). + fn ensure(t: *Scratch, s: cuda.Stream, l: cuda.exl3_experts.Layer, rows: usize, world: usize) !void { + if (rows > t.rows) { + try s.synchronize(); + if (t.picks) |*b| b.free(); + t.picks = null; + t.rows = 0; + t.world = world; + t.picks = try cuda.DeviceBuffer.alloc(t.d, gatherAt(rows) + world * canon.maxPartials(world) * rows * dims * 4); + t.rows = rows; + var meta: [3 * canon.max_pieces]i32 = undefined; + var ps: [canon.max_pieces]canon.Piece = undefined; + const np = canon.pieces(&ps, world); + if (np > 32) return error.TensorParallelUnavailable; + for (ps[0..np], 0..) |q, i| meta[3 * i ..][0..3].* = .{ q.start, q.log, @as(i32, q.rank) * @as(i32, @intCast(canon.maxPartials(world))) + q.idx }; + try t.picks.?.upload(gatherAt(rows) - meta_bytes, std.mem.sliceAsBytes(meta[0 .. 3 * np])); + } + const cap = @min(rows, moe_rows); + if (cap > t.moe_cap) { + try s.synchronize(); + if (t.moe) |*b| b.free(); + t.moe = null; + t.moe_cap = 0; + const bytes = try cuda.exl3_experts.scratchBytes(l, cap, slots); + t.moe = try cuda.DeviceBuffer.alloc(t.d, bytes); + try s.d.check(s.d.api.cuMemsetD8Async(t.moe.?.ptr, 0, bytes, s.handle), "cuMemsetD8Async"); + t.moe_cap = cap; + } + } + + /// pick -> -1 outside [lo, hi), for `n` slots. + pub fn mask(t: *const Scratch, s: cuda.Stream, picks: u64, mpicks: u64, n: usize, r: Rank) !void { + var a: cuda.Args = .{}; + a.add(picks); + a.add(mpicks); + a.add(@as(c_int, @intCast(n))); + a.add(@as(c_int, @intCast(r.lo))); + a.add(@as(c_int, @intCast(r.hi))); + try cuda.launch.launch(t.mask_fn, .{ .grid = .{ .x = @intCast((n + 255) / 256) }, .block = .{ .x = 256 } }, s, &a); + } +}; + +/// Bytes of the finish's piece table, kept in front of the gather buffer. +const meta_bytes: usize = 3 * 32 * 4; + +/// Where the gather buffer starts in the picks allocation of `rows` rows: after the picks and the piece table. +fn gatherAt(rows: usize) usize { + return align256(rows * slots * 4) + meta_bytes; +} + +/// The Ops' expert-parallel scratch, made on first use. +pub fn scratchOf(ops: *load.Ops) !*Scratch { + if (ops.tp == null) ops.tp = try Scratch.create(ops.gpa, ops.d); + return ops.tp.?; +} + +/// Every expert-parallel scratch a pass of up to `rows` rows reads, allocated now (cuda_exl3_pass.attach). +pub fn reserve(ops: *load.Ops, s: cuda.Stream, l: cuda.exl3_experts.Layer, rows: usize, world: usize) !void { + const t = try scratchOf(ops); + try t.ensure(s, l, rows, world); + if (ops.picks == null) ops.picks = try cuda.HostBuffer.alloc(ops.d, moe_rows * slots * 4); + try s.synchronize(); +} + +/// The expert-parallel buffers' addresses and sizes (cuda_exl3_load.Ops.fingerprint's tail). +pub fn fingerprint(t: *const Scratch) [2][2]u64 { + var f: [2][2]u64 = @splat(.{ 0, 0 }); + for ([2]?cuda.DeviceBuffer{ t.picks, t.moe }, &f) |b, *e| if (b) |x| { + e.* = .{ x.ptr, x.len }; + }; + return f; +} + +/// Device bytes the expert-parallel scratch holds. +pub fn reservedBytes(t: *const Scratch) usize { + var n: usize = 0; + for ([2]?cuda.DeviceBuffer{ t.picks, t.moe }) |b| if (b) |x| { + n += x.len; + }; + return n; +} + +/// cuda_exl3_pass.moe for a layer whose experts carry a rank: its slots into y, one all-reduce, the one-rank sum. +pub fn moe(p: anytype, lw: anytype, x: *const load.Experts) !void { + const r = x.tp orelse return error.UnexpectedTensor; + const s = p.s; + const n = p.rows; + const st = p.stream; + const c = p.m.collective orelse return error.TensorParallelUnavailable; + try p.tri.routerMany(s.mixed, lw.router.ptr, s.router_l, dims, n); + try p.tri.topkMany(s.router_l, s.picks, s.wts, n); + const ex = try x.ops.grouped(); + const t = try scratchOf(x.ops); + try t.ensure(st, x.layer, n, r.world); + const mp = t.picks.?.ptr; + try t.mask(st, s.picks, mp, n * slots, r); + if (x.needsPool()) return error.Exl3PoolNotAttached; + var r0: usize = 0; + while (r0 < n) : (r0 += t.moe_cap) { + const rows = @min(t.moe_cap, n - r0); + var sc = try cuda.exl3_experts.scratch(t.moe.?.ptr, x.layer, t.moe_cap, slots); + // runSlots leaves each owned slot's fp32 output in sc.y [rows * slots, dims]: the pass's y rows are that layout + sc.y = s.y + r0 * slots * dims * 4; + const xin = s.mixed + r0 * dims * 2; + // offloaded: the rank's own pool plans on the rank's masked picks (the other rank's slots are -1) + if (x.pool) |pool| { + try @import("../model/cuda_exl3_pass.zig").pooled(st, x.ops, pool, ex, x.layer, xin, mp + r0 * slots * 4, rows, sc, 0, 0); + continue; + } + try ex.runSlots(st, x.layer, xin, dims, .bf16, mp + r0 * slots * 4, rows, sc); + } + try exchange(p, t, c, r); +} + +/// The canonical branch on one rank: the whole leaf tree (receipts canon-reduction.md) into moe_f32. +pub fn canonical(p: anytype) !void { + const m = @import("../../cuda_batch_moe.zig"); + try m.leafTree(p, p.s.moe_f32, 0, canon.leaf_log); + m.kept(p); +} + +/// The canonical branch on `world` ranks: own subtrees, one fixed-size all-gather, every rank finishes the tree. +fn exchange(p: anytype, t: *Scratch, c: anytype, r: Rank) !void { + const m = @import("../../cuda_batch_moe.zig"); + const n = p.rows; + const most = canon.maxPartials(r.world); + const piece = n * dims; // floats a partial + const base = t.picks.?.ptr + gatherAt(t.rows); + const run = canon.run(r.world, r.rank); + var blk: [canon.max_blocks]canon.Block = undefined; + const nb = canon.blocks(&blk, run.lo, run.hi); + for (blk[0..nb], 0..) |b, j| try m.leafTree(p, base + (r.rank * most + j) * piece * 4, b.start, b.log); + try c.gatherF32(p.stream, base + r.rank * most * piece * 4, base, most * piece); + var ps: [canon.max_pieces]canon.Piece = undefined; + const np = canon.pieces(&ps, r.world); + try m.leafFinish(p, base, p.s.moe_f32, t.picks.?.ptr + gatherAt(t.rows) - meta_bytes, np, piece); + m.kept(p); +} + +// ---- memory ------------------------------------------------------------------------------------------------------ + +/// The expert id a tensor belongs to (routed N, or 512 for the shared expert), else null (replicated). +pub fn expertOf(name: []const u8) ?usize { + if (std.mem.indexOf(u8, name, ".mlp.shared_expert.") != null) return wt.experts_n; + const key = ".mlp.experts."; + const at = std.mem.indexOf(u8, name, key) orelse return null; + const rest = name[at + key.len ..]; + const end = std.mem.indexOfScalar(u8, rest, '.') orelse return null; + return std.fmt.parseInt(usize, rest[0..end], 10) catch null; +} + +/// One tensor's device bytes on rank `r`: memory.tensorBytes for what the rank holds, 0 for the other rank's experts. +pub fn tensorBytes(name: []const u8, e: core.safetensors.Entry, o: memory.Options, r: Rank) !u64 { + if (expertOf(name)) |id| if (!r.owns(id)) return 0; + // the vocabulary-split head (load.shardHead): half the trellis and svh a rank, suh whole + if (r.world == 2 and load.shardHead(2) and (std.mem.eql(u8, name, "lm_head.trellis") or std.mem.eql(u8, name, "lm_head.svh"))) + return (try memory.tensorBytes(name, e, o)) / 2; + return memory.tensorBytes(name, e, o); +} + +/// Per layer: three 513-entry pointer tables (int64) and three width tables (int32), on every rank. +const tables_per_layer: u64 = 3 * table_n * 8 + 3 * table_n * 4; + +/// The main model's device bytes on rank `r` (the MTP head is memory.headBytes, rank 0's). +pub fn deviceBytes(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, o: memory.Options, r: Rank) !u64 { + var store = try wt.Store.open(gpa, io, dir); + defer store.deinit(); + const wm = store.parsed.value.object.get("weight_map") orelse return error.MissingTensor; + var files: std.StringHashMapUnmanaged(void) = .empty; + defer files.deinit(gpa); + var it = wm.object.iterator(); + while (it.next()) |kv| try files.put(gpa, kv.value_ptr.string, {}); + var total: u64 = tables_per_layer * wt.layers_n; + var fi = files.keyIterator(); + while (fi.next()) |rel| { + const path = try std.fs.path.join(gpa, &.{ dir, rel.* }); + defer gpa.free(path); + var file = try core.safetensors.File.open(gpa, io, path); + defer file.close(io); + var ni = file.names.iterator(); + while (ni.next()) |e| total = try std.math.add(u64, total, try tensorBytes(e.key_ptr.*, e.value_ptr.*, o, r)); + } + return total; +} + +/// cuda_exl3_memory.placement at world 2: the rank's experts, everything else replicated, the MTP head on rank 0 +pub fn placement(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, world: usize, rank: usize, drafts: bool) !u64 { + const r = try Rank.init(world, rank); + const o = memory.options(); + const head = if (drafts and rank == 0) try memory.headBytes(gpa, io, dir, o) else 0; + return try deviceBytes(gpa, io, dir, o, r) + head; +} + +// ---- the head's vocabulary split. + +/// Rank `rank`'s output columns [lo, hi) of an `n`-column EXL3 matrix: whole 128-column blocks, equal shares. +pub fn headCut(n: usize, world: usize, rank: usize) ![2]usize { + if ((world != 1 and world != 2) or rank >= world or n % (128 * world) != 0) return error.BadHeadSharding; + const part = n / world; + return .{ rank * part, (rank + 1) * part }; +} + +/// The trellis tiles of output columns [lo, hi), in the stored layout [K/16][N/16][tile words] made dense. +pub fn shardTrellis(gpa: std.mem.Allocator, m: load.Mat, lo: usize, hi: usize) ![]u8 { + if (lo % 128 != 0 or hi % 128 != 0 or hi <= lo or hi > m.n) return error.BadHeadSharding; + const tile = @as(usize, m.halves) * 16; // bytes a 16x16 tile: 256 values at halves/2 bits + const row = (m.n / 16) * tile; + if (m.trellis.len != (m.k / 16) * row) return error.UnexpectedTensor; + const part = ((hi - lo) / 16) * tile; + const out = try gpa.alloc(u8, (m.k / 16) * part); + for (0..m.k / 16) |kt| @memcpy(out[kt * part ..][0..part], m.trellis[kt * row + (lo / 16) * tile ..][0..part]); + return out; +} + +/// The EXL3 head `prefix` (lm_head) as rank `rank`'s Projection of n / world columns (a dense shard on the device). +pub fn headProjection(gpa: std.mem.Allocator, d: *cuda.Driver, store: *wt.Store, prefix: []const u8, world: usize, rank: usize) !projection.Projection { + const m = try load.mat(store, prefix); + const cut = try headCut(m.n, world, rank); + const n = cut[1] - cut[0]; + const trellis = try shardTrellis(gpa, m, cut[0], cut[1]); + defer gpa.free(trellis); + const p = try gpa.create(load.Proj); + p.* = .{ .gpa = gpa, .ops = try load.opsOf(store, d), .k = m.k, .n = n, .parts = &.{} }; + errdefer p.destroy(); + var bufs: [3]u64 = undefined; + for (&bufs, [3][]const u8{ trellis, m.suh, m.svh[cut[0] * 2 .. cut[1] * 2] }) |*b, bytes| { + try p.bufs.ensureUnusedCapacity(gpa, 1); + const db = try cuda.DeviceBuffer.fromHost(d, bytes); + p.bufs.appendAssumeCapacity(db); + b.* = db.ptr; + } + const st = cuda.exl3.strides(m.k, n, m.halves, false); + p.parts = try gpa.alloc(load.Part, 1); + p.parts[0] = .{ .col = 0, .n = n, .kind = .{ .exl3 = .{ .trellis = bufs[0], .suh = bufs[1], .svh = bufs[2], .k = @intCast(m.k), .n = @intCast(n), .halves = m.halves, .codebook = .mul1, .stride_k = st.k, .stride_nb = st.nb } } }; + // the shard's Dense scratch (split partials, counters) is reserved with the rest at attach + p.ops.noteWeight(p.parts[0].kind.exl3, n); + store.release(); + var out = projection.Projection.empty(d); + out.n = n; + out.k = m.k; + out.exl3 = p; + return out; +} + +// ---- host tests -------------------------------------------------------------------------------------------------- + +test "exl3: expert-parallel ranks own one contiguous id range each, the shared expert on rank 1" { + const a = try Rank.init(2, 0); + const b = try Rank.init(2, 1); + try std.testing.expectEqual(@as(usize, 0), a.lo); + try std.testing.expectEqual(@as(usize, 256), a.hi); + try std.testing.expectEqual(a.hi, b.lo); + try std.testing.expectEqual(@as(usize, 513), b.hi); + for (0..table_n) |id| try std.testing.expect(a.owns(id) != b.owns(id)); // every id exactly once + try std.testing.expect(b.owns(wt.experts_n) and !a.owns(wt.experts_n)); + const one = try Rank.init(1, 0); + try std.testing.expect(one.owns(0) and one.owns(511) and one.owns(512)); + try std.testing.expectError(error.BadRank, Rank.init(2, 2)); + try std.testing.expectError(error.BadRank, Rank.init(4, 0)); +} + +test "exl3: expert tensors count on their own rank only; the router, the shared gate and HC on both" { + try std.testing.expectEqual(@as(?usize, 7), expertOf("model.language_model.layers.0.mlp.experts.7.gate_proj.trellis")); + try std.testing.expectEqual(@as(?usize, 511), expertOf("model.language_model.layers.47.mlp.experts.511.down_proj.svh")); + try std.testing.expectEqual(@as(?usize, 512), expertOf("model.language_model.layers.3.mlp.shared_expert.up_proj.suh")); + try std.testing.expectEqual(@as(?usize, null), expertOf("model.language_model.layers.3.mlp.shared_expert_gate.weight")); + try std.testing.expectEqual(@as(?usize, null), expertOf("model.language_model.layers.3.mlp.gate.weight")); + var e: core.safetensors.Entry = .{ .dtype = .f16, .rank = 1, .shape = .{ 2560, 1, 1, 1, 1 }, .begin = 0, .end = 5120 }; + const r0 = try Rank.init(2, 0); + const r1 = try Rank.init(2, 1); + const routed = "model.language_model.layers.0.mlp.experts.300.gate_proj.suh"; + try std.testing.expectEqual(@as(u64, 0), try tensorBytes(routed, e, .{}, r0)); + try std.testing.expectEqual(@as(u64, 5120), try tensorBytes(routed, e, .{}, r1)); + const shared = "model.language_model.layers.0.mlp.shared_expert.gate_proj.suh"; + try std.testing.expectEqual(@as(u64, 0), try tensorBytes(shared, e, .{}, r0)); + try std.testing.expectEqual(@as(u64, 5120), try tensorBytes(shared, e, .{}, r1)); + e.dtype = .f16; + e.rank = 2; + e.shape = .{ 512, 2560, 1, 1, 1 }; + e.end = 512 * 2560 * 2; + const router = "model.language_model.layers.0.mlp.gate.weight"; + try std.testing.expectEqual(try tensorBytes(router, e, .{}, r0), try tensorBytes(router, e, .{}, r1)); + try std.testing.expect(try tensorBytes(router, e, .{}, r0) > 0); +} + +test "exl3: the split head counts half its trellis and svh a rank" { + const r0 = try Rank.init(2, 0); + const tr: core.safetensors.Entry = .{ .dtype = .i16, .rank = 3, .shape = .{ 160, 15520, 80, 1, 1 }, .begin = 0, .end = 160 * 15520 * 80 * 2 }; + if (load.shardHead(2)) try std.testing.expectEqual(tr.end / 2, try tensorBytes("lm_head.trellis", tr, .{}, r0)); + try std.testing.expectEqual(tr.end, try tensorBytes("lm_head.trellis", tr, .{}, try Rank.init(1, 0))); +} + +test "exl3: tensor parallel placement opens at world 2 and still refuses other worlds" { + const a = std.testing.allocator; + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + try tmp.dir.writeFile(std.testing.io, .{ .sub_path = "config.json", .data = @embedFile("../pack/fixtures/exl3_config.json") }); + const dir = try tmp.dir.realPathFileAlloc(std.testing.io, ".", a); + defer a.free(dir); + // past the gate: no index in the fixture dir, so it fails reading the checkpoint, not at the quantization + for ([_]usize{ 0, 1 }) |rank| { + if (memory.placement(a, std.testing.io, dir, 2, rank, rank == 0)) |_| return error.TestUnexpectedResult else |err| try std.testing.expect(err != error.UnsupportedQuantization and err != error.BadRank); + } + try std.testing.expectError(error.BadRank, memory.placement(a, std.testing.io, dir, 3, 0, false)); + try std.testing.expectError(error.BadRank, memory.placement(a, std.testing.io, dir, 2, 2, false)); +} + +test "exl3: the head splits at 128-column blocks and a shard's tiles are the full matrix's columns" { + try std.testing.expectEqual([2]usize{ 0, 124_160 }, try headCut(248_320, 2, 0)); + try std.testing.expectEqual([2]usize{ 124_160, 248_320 }, try headCut(248_320, 2, 1)); + try std.testing.expectError(error.BadHeadSharding, headCut(248_320 + 128, 2, 0)); + try std.testing.expectError(error.BadHeadSharding, headCut(256, 2, 2)); + // K 32 (2 tile rows), N 512 (32 tiles a row), 3 halves: 48 bytes a tile + const k = 32; + const n = 512; + const halves = 3; + const tile = halves * 16; + var bytes: [(k / 16) * (n / 16) * tile]u8 = undefined; + for (&bytes, 0..) |*v, i| v.* = @truncate(i * 7 + i / 255); + const m: load.Mat = .{ .trellis = &bytes, .suh = &.{}, .svh = &.{}, .k = k, .n = n, .halves = halves }; + const a = std.testing.allocator; + for (0..2) |rank| { + const cut = try headCut(n, 2, rank); + const got = try shardTrellis(a, m, cut[0], cut[1]); + defer a.free(got); + try std.testing.expectEqual(bytes.len / 2, got.len); + for (0..k / 16) |kt| for (0..(cut[1] - cut[0]) / 16) |j| { + const src = (kt * (n / 16) + cut[0] / 16 + j) * tile; + const dst = (kt * ((cut[1] - cut[0]) / 16) + j) * tile; + try std.testing.expectEqualSlices(u8, bytes[src..][0..tile], got[dst..][0..tile]); + }; + } + try std.testing.expectError(error.BadHeadSharding, shardTrellis(a, m, 0, 100)); +} diff --git a/zig/src/families/flashnext/exl3/tp/exl3_canon.zig b/zig/src/families/flashnext/exl3/tp/exl3_canon.zig new file mode 100644 index 0000000000..a35c7599c6 --- /dev/null +++ b/zig/src/families/flashnext/exl3/tp/exl3_canon.zig @@ -0,0 +1,243 @@ +//! The canonical expert reduction at any TP degree: 64 leaves of 8 ids, an fma chain each, a fixed binary tree. +const std = @import("std"); + +pub const leaves: usize = 64; +pub const leaf_log: u8 = 6; +pub const leaf_ids: usize = 8; +pub const slots: usize = 11; +/// Most subtrees one contiguous leaf run of the tree can need (two a level). +pub const max_blocks: usize = 12; +/// Most pieces for up to 8 ranks (a run of 8 leaves is at most 4 blocks; 64 covers any split of the 64 leaves). +pub const max_pieces: usize = 64; + +/// The leaf of an expert id: routed ids by eights, the shared expert (512) in the last leaf. +pub fn leafOf(id: usize) usize { + return if (id >= leaves * leaf_ids) leaves - 1 else id / leaf_ids; +} + +/// A complete subtree: leaves [start, start + 2^log). +pub const Block = struct { start: u8, log: u8 }; + +/// A subtree that rank `rank` computes as its `idx`th partial. +pub const Piece = struct { start: u8, log: u8, rank: u8, idx: u8 }; + +pub const Run = struct { lo: usize, hi: usize }; + +/// The contiguous leaves rank `rank` of `world` owns. +pub fn run(world: usize, rank: usize) Run { + return .{ .lo = rank * leaves / world, .hi = (rank + 1) * leaves / world }; +} + +/// The maximal complete subtrees of the run [lo, hi), left to right. +pub fn blocks(out: *[max_blocks]Block, lo: usize, hi: usize) usize { + var n: usize = 0; + var at = lo; + while (at < hi) { + var log: u8 = leaf_log; + while ((at & ((@as(usize, 1) << @intCast(log)) - 1)) != 0 or at + (@as(usize, 1) << @intCast(log)) > hi) log -= 1; + out[n] = .{ .start = @intCast(at), .log = log }; + n += 1; + at += @as(usize, 1) << @intCast(log); + } + return n; +} + +/// The most partials any rank of `world` computes (1 for a power of two): the all-gather's fixed width. +pub fn maxPartials(world: usize) usize { + var most: usize = 0; + var b: [max_blocks]Block = undefined; + for (0..world) |r| { + const w = run(world, r); + most = @max(most, blocks(&b, w.lo, w.hi)); + } + return most; +} + +/// Every rank's pieces in leaf order (the order the finish merges them); returns the count. +pub fn pieces(out: *[max_pieces]Piece, world: usize) usize { + var n: usize = 0; + var b: [max_blocks]Block = undefined; + for (0..world) |r| { + const w = run(world, r); + const m = blocks(&b, w.lo, w.hi); + for (b[0..m], 0..) |blk, i| { + out[n] = .{ .start = blk.start, .log = blk.log, .rank = @intCast(r), .idx = @intCast(i) }; + n += 1; + } + } + return n; +} + +/// One column of a row: the value of the subtree [start, start + 2^log) (the reference of fn_exl3_leaf_tree). +pub fn subtree(y: *const [slots]f32, w: *const [slots]f32, picks: *const [slots]u16, start: usize, log: u8) f32 { + var leaf: [slots]u8 = undefined; + var val: [slots]f32 = undefined; + var n: usize = 0; + const end = start + (@as(usize, 1) << @intCast(log)); + for (0..slots) |k| { + const l = leafOf(picks[k]); + if (l < start or l >= end) continue; + var j: usize = 0; + while (j < n and leaf[j] != l) j += 1; + if (j == n) { + leaf[n] = @intCast(l); + val[n] = 0.0; + n += 1; + } + val[j] = @mulAdd(f32, y[k], w[k], val[j]); + } + // ascending leaf order, then level by level: adjacent siblings add, a lone node is its own parent (x + 0 = x) + var i: usize = 1; + while (i < n) : (i += 1) { + var j = i; + while (j > 0 and leaf[j - 1] > leaf[j]) : (j -= 1) { + std.mem.swap(u8, &leaf[j - 1], &leaf[j]); + std.mem.swap(f32, &val[j - 1], &val[j]); + } + } + var pos: [slots]u8 = undefined; + for (0..n) |q| pos[q] = leaf[q] - @as(u8, @intCast(start)); + for (0..log) |_| { + var m: usize = 0; + var q: usize = 0; + while (q < n) { + const parent = pos[q] >> 1; + if (q + 1 < n and parent == pos[q + 1] >> 1) { + val[m] = val[q] + val[q + 1]; + q += 2; + } else { + val[m] = val[q]; + q += 1; + } + pos[m] = parent; + m += 1; + } + n = m; + } + return if (n == 0) 0.0 else val[0]; +} + +/// The finish: pieces in leaf order merged left to right, two equal adjacent blocks with an aligned left start adding. +pub fn finish(ps: []const Piece, values: []const f32) f32 { + var s_start: [max_pieces]u8 = undefined; + var s_log: [max_pieces]u8 = undefined; + var s_val: [max_pieces]f32 = undefined; + var top: usize = 0; + for (ps, values) |p, v| { + s_start[top] = p.start; + s_log[top] = p.log; + s_val[top] = v; + top += 1; + while (top >= 2 and s_log[top - 1] == s_log[top - 2] and s_start[top - 2] % (@as(usize, 2) << @intCast(s_log[top - 1])) == 0) { + s_val[top - 2] = s_val[top - 2] + s_val[top - 1]; + s_log[top - 2] += 1; + top -= 1; + } + } + return s_val[0]; +} + +// ---- tests ------------------------------------------------------------------------------------------------------- + +const Row = struct { y: [slots]f32, w: [slots]f32, p: [slots]u16 }; + +fn randomRow(rng: std.Random) Row { + var r: Row = undefined; + for (0..slots) |k| { + r.y[k] = rng.floatNorm(f32) * 0.05; + r.w[k] = rng.floatNorm(f32) * 0.2; + if (rng.uintLessThan(u8, 23) == 0) r.y[k] = -0.0; + if (rng.uintLessThan(u8, 29) == 0) r.w[k] = -0.0; + // clustered picks make several slots share a leaf; the last slot is the shared expert + r.p[k] = if (k == slots - 1) 512 else if (rng.boolean()) rng.uintLessThan(u16, 512) else 256 + rng.uintLessThan(u16, 24); + } + return r; +} + +/// The definition with no shortcut: all 64 leaf sums (empty = +0), then the full tree level by level. +fn naive(r: Row) f32 { + var level: [leaves]f32 = @splat(0.0); + for (0..slots) |k| { + const l = leafOf(r.p[k]); + level[l] = @mulAdd(f32, r.y[k], r.w[k], level[l]); + } + var n: usize = leaves; + while (n > 1) : (n /= 2) for (0..n / 2) |i| { + level[i] = level[2 * i] + level[2 * i + 1]; + }; + return level[0]; +} + +/// One virtual rank a piece: every piece of `world` computed by its own subtree, then merged. +fn emulate(world: usize, r: Row) f32 { + var ps: [max_pieces]Piece = undefined; + const n = pieces(&ps, world); + var vals: [max_pieces]f32 = undefined; + for (ps[0..n], vals[0..n]) |p, *v| v.* = subtree(&r.y, &r.w, &r.p, p.start, p.log); + return finish(ps[0..n], vals[0..n]); +} + +fn bits(f: f32) u32 { + return @bitCast(f); +} + +test "exl3: canon pieces cover the leaves in order for every degree to 8" { + for (1..9) |world| { + var ps: [max_pieces]Piece = undefined; + const n = pieces(&ps, world); + var at: usize = 0; + for (ps[0..n]) |p| { + try std.testing.expectEqual(at, p.start); + try std.testing.expect(p.start % (@as(usize, 1) << @intCast(p.log)) == 0); + at += @as(usize, 1) << @intCast(p.log); + } + try std.testing.expectEqual(leaves, at); + if (std.math.isPowerOfTwo(world)) { + try std.testing.expectEqual(world, n); + try std.testing.expectEqual(@as(usize, 1), maxPartials(world)); + } else try std.testing.expect(maxPartials(world) > 1); + } +} + +test "exl3: canon every degree 1 to 8 is bit-equal to the whole tree and to the unshortcut definition" { + var prng = std.Random.DefaultPrng.init(0x7ee); + const rng = prng.random(); + for (0..4000) |_| { + const row = randomRow(rng); + const whole = subtree(&row.y, &row.w, &row.p, 0, leaf_log); + try std.testing.expectEqual(bits(naive(row)), bits(whole)); + for (1..9) |world| try std.testing.expectEqual(bits(whole), bits(emulate(world, row))); + } +} + +test "exl3: canon mutations of the order, a slot, a leaf or a level change bits" { + var prng = std.Random.DefaultPrng.init(0x7ef); + const rng = prng.random(); + var rev_slots: usize = 0; + var dropped: usize = 0; + var flat: usize = 0; + var wrong_leaf: usize = 0; + for (0..4000) |_| { + const row = randomRow(rng); + const want = bits(subtree(&row.y, &row.w, &row.p, 0, leaf_log)); + var rev = row; + for (0..slots) |k| { + rev.y[k] = row.y[slots - 1 - k]; + rev.w[k] = row.w[slots - 1 - k]; + rev.p[k] = row.p[slots - 1 - k]; + } + rev_slots += @intFromBool(bits(subtree(&rev.y, &rev.w, &rev.p, 0, leaf_log)) != want); + var drop = row; + drop.w[rng.uintLessThan(usize, slots)] = 0.0; + dropped += @intFromBool(bits(subtree(&drop.y, &drop.w, &drop.p, 0, leaf_log)) != want); + // one flat chain over the 11 slots is not the tree + var chain: f32 = 0.0; + for (0..slots) |k| chain = @mulAdd(f32, row.y[k], row.w[k], chain); + flat += @intFromBool(bits(chain) != want); + // a slot filed under the neighbouring leaf + var moved = row; + moved.p[0] = (row.p[0] + 8) % 512; + wrong_leaf += @intFromBool(bits(subtree(&moved.y, &moved.w, &moved.p, 0, leaf_log)) != want); + } + try std.testing.expect(rev_slots > 0 and dropped > 0 and flat > 0 and wrong_leaf > 0); +} diff --git a/zig/src/families/flashnext/exl3/tp/exl3_tp.zig b/zig/src/families/flashnext/exl3/tp/exl3_tp.zig new file mode 100644 index 0000000000..1559f7f009 --- /dev/null +++ b/zig/src/families/flashnext/exl3/tp/exl3_tp.zig @@ -0,0 +1,148 @@ +//! Two-rank cuts for the Qwen3.8-Flash-Next EXL3 3.05 pack. + +const std = @import("std"); +const shape = @import("../pack/exl3_shape.zig"); + +pub const Axis = enum { replicate, split_k, split_n, expert_parallel }; + +pub const Role = enum { gate, up, down, q, k, v, o, head, shared, routed_set, hc_lowrank }; + +pub const Cut = struct { + role: Role, + axis: Axis, + k: usize, + n: usize, + part: usize, +}; + +pub const Span = struct { lo: usize, hi: usize }; + +/// True when `axis_len` can be cut into `ranks` EXL3 pieces. +pub fn legal(axis_len: usize, ranks: usize) bool { + if (ranks == 0) return false; + return axis_len % (shape.hadamard * ranks) == 0; +} + +/// The 3.05 cut for `role` on `rank`. Ranks are 1 (whole tensor) or 2. +pub fn cut(role: Role, rank: usize, ranks: usize) error{ IllegalCut, RankOutOfRange }!Cut { + if (ranks != 1 and ranks != 2) return error.IllegalCut; + if (rank >= ranks) return error.RankOutOfRange; + if (ranks == 1) return unsplit(role); + return switch (role) { + .gate, .up => .{ .role = role, .axis = .split_k, .k = shape.hidden, .n = shape.expert_intermediate, .part = shape.hidden / 2 }, + .down => .{ .role = role, .axis = .split_n, .k = shape.expert_intermediate, .n = shape.hidden, .part = shape.hidden / 2 }, + .shared => .{ .role = role, .axis = .replicate, .k = shape.hidden, .n = shape.expert_intermediate, .part = 0 }, + .routed_set => .{ .role = role, .axis = .expert_parallel, .k = 0, .n = 0, .part = 0 }, + .q => .{ .role = role, .axis = .split_n, .k = shape.hidden, .n = shape.qOut(), .part = shape.qOut() / 2 }, + .k, .v => .{ .role = role, .axis = .split_n, .k = shape.hidden, .n = shape.kvOut(), .part = shape.kvOut() / 2 }, + // o keeps the heads this rank produced in q, so the split is along K. + .o => .{ .role = role, .axis = .split_k, .k = shape.qOut(), .n = shape.hidden, .part = shape.qOut() / 2 }, + .head => .{ .role = role, .axis = .split_n, .k = shape.hidden, .n = shape.vocab, .part = shape.vocab / 2 }, + // Unquantized. The 320-wide low-rank axis is split in half. 128 alignment does not apply. + .hc_lowrank => .{ .role = role, .axis = .split_n, .k = shape.hc_lowrank, .n = shape.hidden, .part = shape.hc_lowrank / 2 }, + }; +} + +fn unsplit(role: Role) Cut { + return switch (role) { + .gate, .up, .shared => .{ .role = role, .axis = .replicate, .k = shape.hidden, .n = shape.expert_intermediate, .part = shape.hidden }, + .down => .{ .role = role, .axis = .replicate, .k = shape.expert_intermediate, .n = shape.hidden, .part = shape.hidden }, + .routed_set => .{ .role = role, .axis = .expert_parallel, .k = 0, .n = 0, .part = 0 }, + .q => .{ .role = role, .axis = .replicate, .k = shape.hidden, .n = shape.qOut(), .part = shape.qOut() }, + .k, .v => .{ .role = role, .axis = .replicate, .k = shape.hidden, .n = shape.kvOut(), .part = shape.kvOut() }, + .o => .{ .role = role, .axis = .replicate, .k = shape.qOut(), .n = shape.hidden, .part = shape.qOut() }, + .head => .{ .role = role, .axis = .replicate, .k = shape.hidden, .n = shape.vocab, .part = shape.vocab }, + .hc_lowrank => .{ .role = role, .axis = .replicate, .k = shape.hc_lowrank, .n = shape.hidden, .part = shape.hc_lowrank }, + }; +} + +/// `[lo, hi)` on the split axis. Expert-parallel spans are expert ids. +pub fn range(role: Role, rank: usize, ranks: usize) error{ IllegalCut, RankOutOfRange }!Span { + const c = try cut(role, rank, ranks); + if (role == .routed_set) { + if (shape.experts % ranks != 0) return error.IllegalCut; + const part = shape.experts / ranks; + return .{ .lo = rank * part, .hi = (rank + 1) * part }; + } + const along: usize = switch (role) { + .gate, .up, .o => c.k, + .hc_lowrank => c.k, + else => c.n, + }; + if (c.axis == .replicate or ranks == 1) return .{ .lo = 0, .hi = along }; + const part = along / ranks; + return .{ .lo = rank * part, .hi = (rank + 1) * part }; +} + +/// True when an EXL3 N-split at `ranks` would break a Hadamard block. +pub fn refusesColumnSplit(k: usize, n: usize, ranks: usize) bool { + _ = k; + return !legal(n, ranks); +} + +/// `axis_len` is this tensor's header axis, not a class total. `hc_lowrank` is not a Hadamard cut. +pub fn share(role: Role, span: u64, axis_len: usize, rank: usize, ranks: usize) error{ IllegalCut, RankOutOfRange }!u64 { + const c = try cut(role, rank, ranks); + if (span == 0 or c.axis == .replicate or c.axis == .expert_parallel) return span; + if (role == .hc_lowrank) { + if (axis_len == 0 or axis_len % ranks != 0) return span; + return span / @as(u64, ranks); + } + if (axis_len == 0 or !legal(axis_len, ranks)) return span; + return span / @as(u64, ranks); +} + +test "gate is K-split and down is N-split" { + const gate = try cut(.gate, 0, 2); + try std.testing.expectEqual(Axis.split_k, gate.axis); + try std.testing.expectEqual(@as(usize, 1280), gate.part); + try std.testing.expect(refusesColumnSplit(2560, 640, 2)); + try std.testing.expect(!legal(640, 2)); + try std.testing.expect(legal(2560, 2)); + try std.testing.expect(!legal(320, 2)); + + const down = try cut(.down, 1, 2); + try std.testing.expectEqual(Axis.split_n, down.axis); + try std.testing.expect(down.axis != .split_k); + const span = try range(.down, 1, 2); + try std.testing.expectEqual(@as(usize, 1280), span.lo); + try std.testing.expectEqual(@as(usize, 2560), span.hi); +} + +test "expert ids abut and q, k, hc take their parts" { + const a = try range(.routed_set, 0, 2); + const b = try range(.routed_set, 1, 2); + try std.testing.expectEqual(@as(usize, 0), a.lo); + try std.testing.expectEqual(@as(usize, 256), a.hi); + try std.testing.expectEqual(a.hi, b.lo); + try std.testing.expectEqual(@as(usize, 512), b.hi); + + const q = try cut(.q, 0, 2); + try std.testing.expectEqual(@as(usize, 3072), q.part); + try std.testing.expect(q.part % 128 == 0); + try std.testing.expectEqual(@as(usize, 256), (try cut(.k, 0, 2)).part); + try std.testing.expectEqual(@as(usize, 160), (try cut(.hc_lowrank, 0, 2)).part); + try std.testing.expect((try cut(.hc_lowrank, 0, 2)).part % 128 != 0); + try std.testing.expectError(error.RankOutOfRange, cut(.gate, 2, 2)); +} + +test "one rank keeps the whole tensor and both ranks of q stay 128-aligned" { + try std.testing.expectEqual(Axis.replicate, (try cut(.gate, 0, 1)).axis); + try std.testing.expectEqual(Axis.expert_parallel, (try cut(.routed_set, 0, 1)).axis); + const all = try range(.routed_set, 0, 1); + try std.testing.expectEqual(@as(usize, 0), all.lo); + try std.testing.expectEqual(@as(usize, 512), all.hi); + try std.testing.expectEqual(@as(usize, 124160), (try cut(.head, 0, 2)).part); + try std.testing.expect((try cut(.head, 0, 2)).part % 128 == 0); +} + +test "share follows the header axis and hc_lowrank still halves 320" { + try std.testing.expectEqual(@as(u64, 500), try share(.q, 1000, shape.qOut(), 0, 2)); + try std.testing.expectEqual(@as(u64, 1000), try share(.q, 1000, 100, 0, 2)); + try std.testing.expect(!legal(shape.hc_lowrank, 2)); + try std.testing.expectEqual(@as(u64, 160), try share(.hc_lowrank, 320, shape.hc_lowrank, 1, 2)); + try std.testing.expectEqual(@as(u64, 640), try share(.shared, 640, shape.expert_intermediate, 0, 2)); + try std.testing.expectEqual(@as(u64, 1000), try share(.q, 1000, shape.qOut(), 0, 1)); + try std.testing.expectError(error.IllegalCut, share(.q, 1000, shape.qOut(), 0, 0)); + try std.testing.expectError(error.IllegalCut, share(.gate, 1000, shape.hidden, 0, 4)); +} diff --git a/zig/src/native/cuda.zig b/zig/src/native/cuda.zig index e6b5d9c796..4facc9f937 100644 --- a/zig/src/native/cuda.zig +++ b/zig/src/native/cuda.zig @@ -32,7 +32,8 @@ const registry = .{ nemotron.native, nemotron.omni, glm.native, qwen35.native, q /// --record opens here (engines.zig refuses it for a backend without this). pub const records = true; pub const backends: []const []const u8 = &.{"cuda"}; -pub const serves: []const []const u8 = &.{ "--mtp-drafts", "--mtp-confidence", "--policy", "--checkpoint-slots", "--spill-gib" }; +// --tp and the flags it places ranks with: a Flash Next EXL3 pack only (cuda_refusal.packRanks) +pub const serves: []const []const u8 = &.{ "--tp", "--mtp-drafts", "--mtp-confidence", "--policy", "--checkpoint-slots", "--spill-gib" }; pub const families: []const api.Family = blk: { var out: [registry.len + 4]api.Family = undefined; for (registry, 0..) |F, i| out[i] = .{ .model_type = F.model_type, .formats = F.formats }; @@ -205,7 +206,8 @@ const no_stream = "the CUDA memory budget fits no stream: one at a {d}-token win const too_many = "--parallel {d} needs {d:.2} GiB for its streams at a {d}-token window, " ++ "and the CUDA memory budget leaves {d:.2} GiB{s}: serve --parallel {d}, or lower --context"; fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.Open, problem: *[]const u8) !?api.Opened { - if (refusals.family(F, o)) |why| { + const exl3_pack = if (comptime @hasDecl(F, "exl3_pack")) o.tp > 1 and F.exl3_pack(gpa, io, o.dir) else false; + if (refusals.admitted(F, o, exl3_pack) orelse refusals.packRanks(F, o, exl3_pack)) |why| { problem.* = why; return null; } @@ -531,6 +533,10 @@ test "Flash Next refuses tensor-rank draft caching before files or CUDA" { try @import("cuda_flashnext_refusal_test.zig").cache(openWith); } +test "Flash Next serves --tp for EXL3 packs only and keeps the tensor-rank draft refusal for any other pack" { + try @import("cuda_flashnext_refusal_test.zig").packs(); +} + test "Flash Next takes --vision at family admission" { try @import("cuda_flashnext_refusal_test.zig").vision(); } diff --git a/zig/src/native/cuda_flashnext_refusal_test.zig b/zig/src/native/cuda_flashnext_refusal_test.zig index 3ed47d9b34..bab853ac46 100644 --- a/zig/src/native/cuda_flashnext_refusal_test.zig +++ b/zig/src/native/cuda_flashnext_refusal_test.zig @@ -40,3 +40,18 @@ pub fn vision() !void { try std.testing.expect(refusals.family(VisionFamily, o) == null); try std.testing.expect(refusals.family(F, .{ .dir = "d", .model_type = F.model_type }) == null); } + +/// --tp opens for an EXL3 pack only: the drafts-with-cache refusal lifts for it, any other Flash Next pack is refused. +pub fn packs() !void { + const o: api.Open = .{ .dir = "/missing-flashnext-test-checkpoint", .model_type = F.model_type, .tp = 2, .rank = 1, .drafts = true, .prompt_cache_gib = 4 }; + try std.testing.expect(refusals.admitted(F, o, true) == null); + try std.testing.expectStringStartsWith(refusals.admitted(F, o, false) orelse "", "Flash Next tensor-rank MTP snapshots are unavailable"); + try std.testing.expect(refusals.packRanks(F, o, true) == null); + try std.testing.expectStringStartsWith(refusals.packRanks(F, o, false) orelse "", "--tp: on CUDA only an EXL3 pack"); + var one = o; + one.tp = 1; + try std.testing.expect(refusals.packRanks(F, one, false) == null); + try std.testing.expect(refusals.admitted(F, one, false) == null); + // a directory that holds no checkpoint is no EXL3 pack + try std.testing.expect(!F.exl3_pack(std.testing.allocator, std.testing.io, o.dir)); +} diff --git a/zig/src/native/cuda_refusal.zig b/zig/src/native/cuda_refusal.zig index 70c2342065..e727d714e8 100644 --- a/zig/src/native/cuda_refusal.zig +++ b/zig/src/native/cuda_refusal.zig @@ -8,6 +8,7 @@ pub const deepseek_release_message = "DeepSeek-V4.1 is not served by this releas pub fn backend(o: api.Open) ?[]const u8 { if (std.mem.eql(u8, o.model_type, "deepseek_v4") or std.mem.eql(u8, o.model_type, "deepseek_v41")) return deepseek_release_message; + if (o.tp != 1 and o.tp != 2) return "--tp: the CUDA engine runs one or two ranks (one process a GPU): serve --tp 1 or --tp 2"; if (o.mtp_drafts) |n| if (n > nemotron.native.max_drafts) return std.fmt.comptimePrint("--mtp-drafts: the CUDA engine drafts at most {d}", .{nemotron.native.max_drafts}); if (o.mtp_confidence) |c| if (!(c >= 0 and c <= 1)) return "--mtp-confidence must be from 0 to 1"; if (o.mtp_drafts != null or o.mtp_confidence != null) { @@ -33,12 +34,17 @@ const lfm2_vision = "--vision: LFM2 reads the checkpoint's images through its ow /// A family takes --drafter only if it says so, and --learn-conversations only as Nemotron; null: `o` may open `F`. pub fn family(comptime F: type, o: api.Open) ?[]const u8 { + return admitted(F, o, false); +} + +/// `family` for a checkpoint that is or is not an EXL3 pack: its ranks keep drafting states, so --tp with drafts opens. +pub fn admitted(comptime F: type, o: api.Open, exl3_pack: bool) ?[]const u8 { if (o.vision != null and !(@hasDecl(F, "Options") and @hasField(F.Options, "vision"))) return "--vision: this CUDA family reads no images; serve it without --vision"; const ranked = @hasDecl(F, "Options") and @hasField(F.Options, "rank"); if (o.tp != 1 and !ranked) return "this CUDA family has no tensor-parallel ranks: serve without --tp"; if (@hasDecl(F, "tp_draft_snapshots")) { - if (!F.tp_draft_snapshots and o.tp > 1 and o.drafts and (o.prompt_cache_gib orelse 1) > 0) + if (!F.tp_draft_snapshots and !exl3_pack and o.tp > 1 and o.drafts and (o.prompt_cache_gib orelse 1) > 0) return "Flash Next tensor-rank MTP snapshots are unavailable: use --prompt-cache-gib 0 with --tp and drafting"; } // ahead of the drafter's and --learn-conversations' refusals (--spill-gib sets that flag), on a line of its own @@ -55,6 +61,12 @@ pub fn family(comptime F: type, o: api.Open) ?[]const u8 { return null; } +/// --tp is served on CUDA for a family's EXL3 packs (`F.exl3_pack`); any other pack keeps the flag refused. +pub fn packRanks(comptime F: type, o: api.Open, exl3_pack: bool) ?[]const u8 { + if (o.tp == 1 or !@hasDecl(F, "exl3_pack") or exl3_pack) return null; + return "--tp: on CUDA only an EXL3 pack is placed across ranks by --tp --rank; this checkpoint is not one: serve without --tp"; +} + test "a family must consume rank options before tensor parallelism is admitted" { const Single = struct { pub const Options = struct { context: usize }; @@ -66,6 +78,15 @@ test "a family must consume rank options before tensor parallelism is admitted" try std.testing.expect(family(Ranked, .{ .dir = "d", .model_type = "m", .tp = 2 }) == null); } +test "--tp is served on CUDA for two ranks of a ranked family (Flash Next) only; 4 and 8 stay refused" { + const flashnext = @import("flashnext").native; + for ([_]u32{ 4, 8 }) |tp| try std.testing.expectStringStartsWith(backend(.{ .dir = "d", .model_type = flashnext.model_type, .tp = tp }) orelse "", "--tp: the CUDA engine runs one or two ranks"); + for ([_]u32{ 1, 2 }) |tp| try std.testing.expect(backend(.{ .dir = "d", .model_type = flashnext.model_type, .tp = tp }) == null); + try std.testing.expect(family(flashnext, .{ .dir = "d", .model_type = flashnext.model_type, .tp = 2, .rank = 1, .drafts = false }) == null); + // a family without rank options (Nemotron) still refuses --tp 2 at its admission + try std.testing.expectStringStartsWith(family(nemotron.native, .{ .dir = "d", .model_type = nemotron.native.model_type, .tp = 2 }) orelse "", "this CUDA family has no tensor-parallel ranks"); +} + test "--learn is refused whatever the family, and --slide or a learned-disk size alone is not" { try std.testing.expectStringStartsWith(backend(.{ .dir = "d", .model_type = "nemotron_h", .learn = "/x" }) orelse "", "--learn learns shared prompt prefixes on Macs only"); try std.testing.expect(backend(.{ .dir = "d", .model_type = "nemotron_h", .learn = "/x", .learn_conversations = true }) == null); diff --git a/zig/tests/cuda/contract_flashnext_exl3_suite.zig b/zig/tests/cuda/contract_flashnext_exl3_suite.zig new file mode 100644 index 0000000000..3cb954e766 --- /dev/null +++ b/zig/tests/cuda/contract_flashnext_exl3_suite.zig @@ -0,0 +1,26 @@ +//! Flash Next EXL3 contract kernels the runner builds together and releases together. +const std = @import("std"); +const cuda = @import("cuda"); +const check = @import("check.zig"); +const g = @import("core").contract.gates; +const tp_ = @import("contract_flashnext_exl3_tp.zig"); + +pub const Suite = struct { + a: std.mem.Allocator, + masks: tp_.Masks, + + pub fn create(a: std.mem.Allocator, gpu: check.Gpu, s: cuda.Stream, kernels: *std.ArrayList(g.Kernel)) !*Suite { + const self = try a.create(Suite); + errdefer a.destroy(self); + self.a = a; + self.masks = try tp_.Masks.init(gpu, s); + errdefer self.masks.deinit(); + try self.masks.add(a, kernels); + return self; + } + + pub fn destroy(self: *Suite) void { + self.masks.deinit(); + self.a.destroy(self); + } +}; diff --git a/zig/tests/cuda/contract_flashnext_exl3_tp.zig b/zig/tests/cuda/contract_flashnext_exl3_tp.zig new file mode 100644 index 0000000000..0210c91bfa --- /dev/null +++ b/zig/tests/cuda/contract_flashnext_exl3_tp.zig @@ -0,0 +1,138 @@ +//! Flash Next EXL3 expert parallel: fn_exl3_tp_mask keeps the picks a rank holds and turns the others into -1. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const check = @import("check.zig"); +const rig_ = @import("contract_rig.zig"); +const c = core.contract.spec; +const g = core.contract.gates; + +/// Slots a row (ten routed experts and the shared one) and the ids a table holds (512 routed, the shared one as 512). +const slots = 11; +const table_n = 513; + +fn noSums(_: c.Shape) usize { + return 0; +} + +/// The ids [lo, hi) of rank `rank` of two (cuda_exl3_tp.Rank). +fn range(rank: usize) [2]i32 { + return if (rank == 0) .{ 0, 256 } else .{ 256, table_n }; +} + +fn spec(rank: usize) c.Spec { + return .{ + .name = if (rank == 0) "cuda Flash Next EXL3 fn_exl3_tp_mask (rank 0)" else "cuda Flash Next EXL3 fn_exl3_tp_mask (rank 1)", + .accum = .none, + .depth = noSums, + .out = .raw, + .nan = .{ .stated = "integer picks only: a pick in the rank's [lo, hi) keeps its value and every other one, negatives and the int32 extremes included, becomes -1" }, + // the masks run over rows * 11 picks in blocks of 256: 24 rows is the first count that crosses one + .tiles = .{ .{ .size = 24, .step = 1, .min = 1 }, .{}, .{} }, + .served = &.{ .{ .rows = 1, .n = slots, .k = 0 }, .{ .rows = 8, .n = slots, .k = 0 }, .{ .rows = 64, .n = slots, .k = 0 }, .{ .rows = 4096, .n = slots, .k = 0 } }, + }; +} + +pub const Mask = struct { + rig: rig_.Rig, + f: cuda.Function, + rank: usize, + shape: c.Shape = .{ .rows = 0, .n = 0, .k = 0 }, + picks: []i32 = &.{}, + din: u64 = 0, + dout: u64 = 0, + msg: [160]u8 = undefined, + + pub fn deinit(self: *Mask) void { + self.rig.gpu.gpa.free(self.picks); + self.rig.deinit(); + } + + const vt = g.hooks(Mask); + + pub fn kernel(self: *Mask) g.Kernel { + return .{ .spec = spec(self.rank), .ctx = self, .vt = &vt }; + } + + pub fn prepare(self: *Mask, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows == 0 or shape.rows > 8192) return false; + self.rig.reset(); + const a = self.rig.gpu.gpa; + a.free(self.picks); + self.picks = &.{}; + self.shape = shape; + self.picks = try a.alloc(i32, shape.rows * slots); + var prng = std.Random.DefaultPrng.init(seed ^ 0x6d61736b); + const rng = prng.random(); + for (self.picks) |*p| p.* = rng.intRangeLessThan(i32, 0, table_n); + // the edges of both ranks' ranges, and values no router makes (the kernel masks them like any other id) + const edge = [_]i32{ 0, 255, 256, 511, 512, 513, -1, -2, std.math.minInt(i32), std.math.maxInt(i32) }; + if (fill != .random) for (self.picks, 0..) |*p, i| { + if (i % 3 == 0) p.* = edge[(i / 3) % edge.len]; + }; + self.din = try self.rig.input(std.mem.sliceAsBytes(self.picks)); + self.dout = try self.rig.output(self.picks.len * 4); + return true; + } + + pub fn run(self: *Mask) !void { + const r = range(self.rank); + var args: cuda.Args = .{}; + args.add(self.din); + args.add(self.dout); + args.add(@as(c_int, @intCast(self.picks.len))); + args.add(@as(c_int, r[0])); + args.add(@as(c_int, r[1])); + try cuda.launch.launch(self.f, .{ .grid = .{ .x = @intCast((self.picks.len + 255) / 256) }, .block = .{ .x = 256 } }, self.rig.s, &args); + try self.rig.s.synchronize(); + } + + pub fn count(self: *Mask) usize { + return self.picks.len; + } + + fn expected(self: *Mask, p: i32) i32 { + const r = range(self.rank); + return if (p >= r[0] and p < r[1]) p else -1; + } + + pub fn reference(self: *Mask, want: []f64, mag: []f64) !void { + for (self.picks, want, mag) |p, *w, *m| { + w.* = @floatFromInt(self.expected(p)); + m.* = @abs(w.*); + } + } + + pub fn check(self: *Mask, _: g.Gate) !?[]const u8 { + if (!try self.rig.intact()) return "a payload guard changed"; + if (!std.mem.eql(u8, std.mem.sliceAsBytes(self.picks), try self.rig.inputBytes())) return "the picks changed"; + const out = try self.rig.outputBytes(); + for (self.picks, 0..) |p, i| { + const got = std.mem.readInt(i32, out[i * 4 ..][0..4], .little); + if (got != self.expected(p)) return try std.fmt.bufPrint(&self.msg, "pick {d} ({d}) became {d}, not {d}", .{ i, p, got, self.expected(p) }); + } + return null; + } +}; + +/// Both ranks' masks over one loaded image, for the runner's list. +pub const Masks = struct { + mod: cuda.Module, + each: [2]Mask, + + pub fn init(gpu: check.Gpu, s: cuda.Stream) !Masks { + var mod = try cuda.Module.load(gpu.d, cuda.kernels.flashnext_exl3_tp); + errdefer mod.unload(); + const f = try mod.function("fn_exl3_tp_mask"); + return .{ .mod = mod, .each = .{ .{ .rig = .{ .gpu = gpu, .s = s }, .f = f, .rank = 0 }, .{ .rig = .{ .gpu = gpu, .s = s }, .f = f, .rank = 1 } } }; + } + + pub fn deinit(self: *Masks) void { + for (&self.each) |*m| m.deinit(); + self.mod.unload(); + } + + pub fn add(self: *Masks, a: std.mem.Allocator, list: *std.ArrayList(g.Kernel)) !void { + for (&self.each) |*m| try list.append(a, m.kernel()); + } +}; diff --git a/zig/tests/cuda/contract_limits.zig b/zig/tests/cuda/contract_limits.zig index 48bb601d76..b25ab8b403 100644 --- a/zig/tests/cuda/contract_limits.zig +++ b/zig/tests/cuda/contract_limits.zig @@ -26,7 +26,7 @@ pub fn image(name: []const u8) error{UnknownContractImage}![]const u8 { inline for (.{ "qwen4_projection", "qwen4_hc", "qwen4_gdn", "qwen4_ple" }) |n| if (std.mem.eql(u8, name, n)) return @field(kk, n); inline for (.{ "tree_round", "act_rows", "gdn_prep", "qk_rows", "norms", "moe", "experts", "qwen36_mtp", "nemotron_attention", "nemotron_route", "nemotron_latent", "qmm_group", "qmm_prefill", "lane_gemv", "prompt16", "qmmf", "ck_experts", "nvfp4_experts", "nvfp4_experts_tile", "rank_sum", "prefill_attention", "vit", "embed_rows", "route_rows", "flashnext_gdn", "flashnext_batch", "flashnext_hc", "flashnext_attn", "flashnext_mtp", "qwen_vision", "exl3_linear", "exl3_experts", "exl3_candidate", "experts3", "experts3_prefill", "drama3", "drama3_fp8", "drama3_codec", "glm_mm", "glm_dflash", "glm_glue", "glm_ops", "glm_index", "nemotron_mamba", "sample", "qwen35", "qwen35_bf16_rows", "qwen35_keep", "qwen35_kv8", "w4a8_prefill", "qwen35_draft", "ck_lane", "gdn_train", "train_q4", "train_attn", "flashnext_train", "train", "attn_train", "qwen27_train", "tree_attention", "block_attention", "window_attention", "window_pages", "moe_glue", "dense_rows", "qmm", "flashnext_fact", "flashnext_hcf", "flashnext_fexp", "experts_rows" }) |n| if (std.mem.eql(u8, name, n)) return @field(kk, n); - inline for (.{"flashnext_exl3"}) |n| + inline for (.{ "flashnext_exl3", "flashnext_exl3_tp" }) |n| if (std.mem.eql(u8, name, n)) return @field(kk, n); if (std.mem.eql(u8, name, "contract_controls")) return if (options.controls) &controls.bytes else &.{}; std.debug.print("FAIL kernel-contract: unknown image name: {s}\n", .{name}); diff --git a/zig/tests/cuda/contract_tests.zig b/zig/tests/cuda/contract_tests.zig index fd53c75789..57f5252a05 100644 --- a/zig/tests/cuda/contract_tests.zig +++ b/zig/tests/cuda/contract_tests.zig @@ -255,6 +255,8 @@ pub fn run(gpu: check.Gpu) !void { } defer for (&writebacks) |*w| w.deinit(); for (&writebacks) |*w| try kernels.append(a, w.kernel()); + const exl3 = try @import("contract_flashnext_exl3_suite.zig").Suite.create(a, gpu, s, &kernels); + defer exl3.destroy(); var hcnorms: [2]hcnorm_k.Normalize = undefined; for (std.enums.values(hcnorm_k.Kind), &hcnorms) |kind, *x| x.* = .{ .rig = .{ .gpu = gpu, .s = s }, .kind = kind, .cols_fn = try hc.function("fn_hc_normalize_cols"), .old_fn = try hc.function("fn_hc_normalize") }; defer for (&hcnorms) |*x| x.deinit(); diff --git a/zig/tests/cuda/flashnext_exl3/model/flashnext_exl3_engine_tests.zig b/zig/tests/cuda/flashnext_exl3/model/flashnext_exl3_engine_tests.zig index 15236b0bad..33f43bc8e0 100644 --- a/zig/tests/cuda/flashnext_exl3/model/flashnext_exl3_engine_tests.zig +++ b/zig/tests/cuda/flashnext_exl3/model/flashnext_exl3_engine_tests.zig @@ -52,7 +52,7 @@ const Eng = struct { g.out = try cuda.DeviceBuffer.alloc(gpu.d, out_rows * vocab * 2); var b = try xpass.budgetFromEnv(0); if (lines) |n| b.lines = n; - g.attached = try xpass.attach(gpu.gpa, gpu.io, gpu.d, g.stream, g.e.layers[0..g.e.n_layers], g.e.chunk, b, null); + g.attached = try xpass.attach(gpu.gpa, gpu.io, gpu.d, g.stream, g.e.layers[0..g.e.n_layers], null, g.e.chunk, b, null); const t2 = check.now(gpu.io); g.load_s = ms(t0, t2) / 1e3; var free: usize = 0; @@ -86,7 +86,7 @@ const Eng = struct { } fn stats(g: *Eng) flash.exl3_pool.Stats { - return xpass.poolStats(g.e.layers[0..g.e.n_layers]); + return xpass.poolStats(g.e.layers[0..g.e.n_layers], null); } }; @@ -162,33 +162,8 @@ const Raw = struct { } }; -/// lm_head as stored: trellis int16 [K/16, N/16, 8 * halves] (bytes), suh fp16 [K], svh fp16 [N]. -const Head = struct { - trellis: []const u8, - suh: []const u8, - svh: []const u8, - k: usize, - n: usize, - halves: u16, - - fn words(h: Head) usize { - return @as(usize, h.halves) * 8; - } - - fn of(trellis: st.Tensor, suh: st.Tensor, svh: st.Tensor, mul1: ?st.Tensor) !Head { - if (trellis.dtype != .i16 or trellis.rank != 3 or trellis.dim(2) % 8 != 0) return error.UnexpectedTensor; - const h: Head = .{ .trellis = trellis.bytes, .suh = suh.bytes, .svh = svh.bytes, .k = trellis.dim(0) * 16, .n = trellis.dim(1) * 16, .halves = @intCast(trellis.dim(2) / 8) }; - if (suh.dtype != .f16 or suh.numel() != h.k or svh.dtype != .f16 or svh.numel() != h.n) return error.UnexpectedTensor; - if (h.k % fmt.hadamard_dim != 0 or h.n % fmt.hadamard_dim != 0 or h.halves < 2 or h.halves > 16) return error.UnexpectedTensor; - if (mul1) |m| { - if (m.bytes.len != 4 or std.mem.readInt(u32, m.bytes[0..4], .little) != fmt.Mul1Mul) return error.UnexpectedTensor; - } - return h; - } -}; - const HeadWork = struct { - head: Head, + head: flash.exl3_mtp.Head, x: []const f64, // [rows, K] rows: usize, out: []f64, // [rows, N] @@ -244,7 +219,7 @@ fn headWork(w: *HeadWork, gpa: std.mem.Allocator) !void { } fn headF64(gpa: std.mem.Allocator, raw: *const Raw, x: []const f64, rows: usize) ![]f64 { - const head = try Head.of(try raw.get("lm_head.trellis"), try raw.get("lm_head.suh"), try raw.get("lm_head.svh"), raw.get("lm_head.mul1") catch null); + const head = try flash.exl3_mtp.Head.of(try raw.get("lm_head.trellis"), try raw.get("lm_head.suh"), try raw.get("lm_head.svh"), raw.get("lm_head.mul1") catch null); const out = try gpa.alloc(f64, rows * head.n); errdefer gpa.free(out); var w: HeadWork = .{ .head = head, .x = x, .rows = rows, .out = out }; diff --git a/zig/tests/cuda/flashnext_exl3/model/flashnext_exl3_names_tests.zig b/zig/tests/cuda/flashnext_exl3/model/flashnext_exl3_names_tests.zig index 49a8d78eb7..afbf84ab7f 100644 --- a/zig/tests/cuda/flashnext_exl3/model/flashnext_exl3_names_tests.zig +++ b/zig/tests/cuda/flashnext_exl3/model/flashnext_exl3_names_tests.zig @@ -194,7 +194,8 @@ pub fn names(gpu: Gpu, dir: []const u8) !void { const gib = 1024.0 * 1024.0 * 1024.0; const on_host = try X.memory.deviceBytes(gpa, gpu.io, dir, .{}); const resident = try X.memory.deviceBytes(gpa, gpu.io, dir, .{ .routed_on_host = false }); - std.debug.print(" device bytes: {d:.2} GiB with routed experts on the host, {d:.2} GiB all resident\n", .{ @as(f64, @floatFromInt(on_host)) / gib, @as(f64, @floatFromInt(resident)) / gib }); + const head = try X.memory.headBytes(gpa, gpu.io, dir, .{}); + std.debug.print(" device bytes: {d:.2} GiB with routed experts on the host, {d:.2} GiB all resident; MTP head {d:.2} GiB (routed on the host)\n", .{ @as(f64, @floatFromInt(on_host)) / gib, @as(f64, @floatFromInt(resident)) / gib, @as(f64, @floatFromInt(head)) / gib }); check.pass("names: every tensor the family reads resolves ({d}), nothing unclaimed but {d} vision tensors; {d} EXL3 matrices admitted; {d} F16 weights widened exactly", .{ n.claimed.count(), vision, n.matrices, n.widened }); } diff --git a/zig/tests/cuda/flashnext_exl3/model/flashnext_exl3_rig.zig b/zig/tests/cuda/flashnext_exl3/model/flashnext_exl3_rig.zig index 55f85d788f..311c7f50a0 100644 --- a/zig/tests/cuda/flashnext_exl3/model/flashnext_exl3_rig.zig +++ b/zig/tests/cuda/flashnext_exl3/model/flashnext_exl3_rig.zig @@ -171,6 +171,8 @@ pub const Rig = struct { history: [2]i64 = .{ 0, 0 }, scores: cuda.DeviceBuffer = undefined, cap: usize, + /// The ranks' NCCL owner when a check runs two ranks (flashnext_exl3_tp_tests.zig); null: one rank. + collective: ?*flash.collective.Collective = null, pub fn init(r: *Rig, gpu: Gpu, dir: []const u8, rows: usize, cap: usize) !void { r.* = .{ .gpu = gpu, .stream = try cuda.Stream.init(gpu.d, true), .glue = undefined, .k = undefined, .s = undefined, .store = undefined, .pack_mod = undefined, .pack_fn = undefined, .cap = cap }; @@ -237,7 +239,7 @@ pub const Rig = struct { .tri = .{ .own = &r.glue, .s = r.stream }, .k = &r.k, .s = &r.s, - .m = .{ .table_w = 0, .table_s = 0, .table_b = 0, .layers = layers, .mixer = undefined, .head = undefined, .ngram = undefined, .collective = null }, + .m = .{ .table_w = 0, .table_s = 0, .table_b = 0, .layers = layers, .mixer = undefined, .head = undefined, .ngram = undefined, .collective = r.collective }, .seq = .{ .conv = &r.conv, .state = &r.state, .kc = &r.kc, .vc = &r.vc, .ikc = &r.ikc, .pooled = &r.pooled, .tail = r.tail, .history = &r.history, .cap = r.cap, .scores = r.scores.ptr }, .hooks = .{}, }; diff --git a/zig/tests/cuda/flashnext_exl3/offload/exl3_frequency_tests.zig b/zig/tests/cuda/flashnext_exl3/offload/exl3_frequency_tests.zig new file mode 100644 index 0000000000..0dfacd8475 --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/offload/exl3_frequency_tests.zig @@ -0,0 +1,82 @@ +const std = @import("std"); +const cuda = @import("cuda"); +const flash = @import("flashnext"); +const check = @import("../../check.zig"); +const profile = flash.exl3.frequency; +const Pool = flash.exl3_pool.Pool; + +fn verify(gpu: check.Gpu, pool: *Pool, expected: usize) !void { + const tables = try gpu.gpa.alloc(u8, pool.lines.tables.len); + defer gpu.gpa.free(tables); + try pool.lines.tables.download(0, tables); + var resident: usize = 0; + var bytes: [16]u8 = undefined; + for (0..512) |expert| { + var found: ?usize = null; + for (pool.slots_tab, pool.line_tab, 0..) |held, state, line| { + if (held == expert and state == .resident) found = line; + } + if (found != null) resident += 1; + for (0..3) |projection| { + const pointer = std.mem.readInt(u64, tables[pool.lines.tlay.ptr[projection] + expert * 8 ..][0..8], .little); + const width = std.mem.readInt(i32, tables[pool.lines.tlay.k2[projection] + expert * 4 ..][0..4], .little); + if (found) |line| { + const offset = line * pool.lines.layout.line_bytes + pool.lines.layout.off[projection]; + try check.expect(pointer == pool.lines.mem.ptr + offset and width == 1, "frequency warm table expert {d} projection {d}", .{ expert, projection }); + try pool.lines.mem.download(offset, &bytes); + try check.sameBytes("frequency warm expert bytes", &bytes, pool.src.routed[expert][projection]); + } else try check.expect(pointer == 0 and width == 0, "unloaded frequency expert has no device entry", .{}); + } + } + try check.expect(resident == expected, "frequency resident count {d} expected {d}", .{ resident, expected }); +} + +pub fn run(gpu: check.Gpu, arguments: []const [:0]const u8) !void { + if (arguments.len != 2) return error.MissingArgument; + const rank = try std.fmt.parseInt(u32, arguments[0], 10); + if (rank >= 2) return error.BadRank; + const path = arguments[1]; + const identity: profile.Identity = .{ .model = @splat(0x57), .world = 2, .rank = rank }; + const start = @as(usize, rank) * 256; + var bank = try cuda.HostBuffer.alloc(gpu.d, 512 * 3 * 16); + defer bank.free(); + var source: [512][3][]const u8 = @splat(.{ &.{}, &.{}, &.{} }); + for (start..start + 256) |expert| for (0..3) |projection| { + const offset = (expert * 3 + projection) * 16; + const bytes = bank.bytes[offset..][0..16]; + @memset(bytes, @truncate(expert + projection)); + source[expert][projection] = bytes; + }; + var saved: profile.Counts = .{}; + for (0..20) |index| saved.decoder[7][start + index] = @intCast(100 + index); + saved.decoder[7][(start + 256) % 512] = std.math.maxInt(u32); + try profile.write(gpu.gpa, gpu.io, path, identity, &saved); + const scales: flash.exl3_pool.Scales = .{ .suh_g = 0, .suh_u = 0, .svh_g = 0, .svh_u = 0, .suh_d = 0, .svh_d = 0 }; + for ([_]usize{ 10, 23 }, 0..) |capacity, restart| { + const restored = try profile.read(gpu.gpa, gpu.io, path, identity); + var pool = try Pool.init(gpu.gpa, gpu.io, gpu.d, .{ .routed = &source }, scales, .{ + .layer = 7, + .lines = capacity, + .dims = 16, + .width = 16, + .max_rows = 1, + .direct = true, + }, null); + defer pool.deinit(); + @memcpy(&pool.counts[7], &restored.decoder[7]); + const before = pool.counts[7]; + const heated = try pool.warmProfile(); + try check.expect(heated == @min(capacity, 20), "frequency restart fills known local experts only", .{}); + try check.expect(std.mem.eql(u32, &before, &pool.counts[7]), "warming preserves observation counts", .{}); + try verify(gpu, &pool, heated); + @memcpy(&saved.decoder[7], &pool.counts[7]); + try profile.write(gpu.gpa, gpu.io, path, identity, &saved); + std.debug.print("FREQUENCY rank={d} restart={d} capacity={d} warmed={d}\n", .{ rank, restart, capacity, heated }); + } + var other_rank = identity; + other_rank.rank = 1 - rank; + if (profile.read(gpu.gpa, gpu.io, path, other_rank)) |_| return error.AcceptedForeignProfile else |failure| { + try check.expect(failure == error.IdentityMismatch, "frequency rejects other TP rank", .{}); + } + check.pass("frequency disk restart, rank isolation and resized warm pools", .{}); +} diff --git a/zig/tests/cuda/flashnext_exl3/pack/exl3_pack_tests.zig b/zig/tests/cuda/flashnext_exl3/pack/exl3_pack_tests.zig index a7c195b25d..4db05eef1a 100644 --- a/zig/tests/cuda/flashnext_exl3/pack/exl3_pack_tests.zig +++ b/zig/tests/cuda/flashnext_exl3/pack/exl3_pack_tests.zig @@ -520,3 +520,58 @@ pub fn run(gpu: Gpu, rest: []const [:0]const u8) !void { try experts(&b, &ex, &pack, a, rng, 17, 1, 0); check.pass("exl3-pack: v1.1.0 EXL3 operators ran on the pack's shapes (see the error and timing lines)", .{}); } + +/// nccl-lat : two processes, one GPU each; fp32 all-reduce (sum) latency at a few sizes, cudaEvent +pub fn ncclLatency(gpu: Gpu, rest: []const [:0]const u8) !void { + if (rest.len < 2) return error.MissingArgument; + const rank = try std.fmt.parseInt(c_int, rest[0], 10); + const id_path = rest[1]; + var lib = try cuda.nccl.Library.open(); + defer lib.close(); + var id: cuda.nccl.UniqueId = undefined; + if (rank == 0) { + try lib.check(lib.api.ncclGetUniqueId(&id), "ncclGetUniqueId"); + const hex = std.fmt.bytesToHex(id.internal, .lower); + try std.Io.Dir.cwd().writeFile(gpu.io, .{ .sub_path = id_path, .data = &hex }); + } else { + var hex: [256]u8 = undefined; + var tries: usize = 0; + while (true) : (tries += 1) { + const got = std.Io.Dir.cwd().readFile(gpu.io, id_path, &hex) catch { + if (tries > 200) return error.NoUniqueId; + try std.Io.sleep(gpu.io, .fromMilliseconds(50), .awake); + continue; + }; + if (got.len == 256) break; + if (tries > 200) return error.NoUniqueId; + try std.Io.sleep(gpu.io, .fromMilliseconds(50), .awake); + } + _ = try std.fmt.hexToBytes(&id.internal, &hex); + } + var comm: cuda.nccl.Comm = null; + try lib.check(lib.api.ncclCommInitRank(&comm, 2, id, rank), "ncclCommInitRank"); + defer _ = lib.api.ncclCommDestroy(comm); + var s = try cuda.Stream.init(gpu.d, true); + defer s.deinit(); + var ev0 = try cuda.Event.init(gpu.d, true); + defer ev0.deinit(); + var ev1 = try cuda.Event.init(gpu.d, true); + defer ev1.deinit(); + const sizes = [_]usize{ 2560, 11 * 2560, 4 * 11 * 2560, 16 * 11 * 2560, 512 * 2560, 4096 * 2560 }; + var buf = try cuda.DeviceBuffer.alloc(gpu.d, sizes[sizes.len - 1] * 4); + defer buf.free(); + try buf.fill8(0, s.handle); + for (sizes) |count| { + for (0..10) |_| try lib.check(lib.api.ncclAllReduce(buf.ptr, buf.ptr, count, .f32, .sum, comm, s.handle), "ncclAllReduce"); + try s.synchronize(); + var best: f64 = std.math.inf(f64); + for (0..5) |_| { + try ev0.record(s); + for (0..100) |_| try lib.check(lib.api.ncclAllReduce(buf.ptr, buf.ptr, count, .f32, .sum, comm, s.handle), "ncclAllReduce"); + try ev1.record(s); + try ev1.synchronize(); + best = @min(best, @as(f64, try ev0.elapsedMs(ev1)) / 100.0); + } + if (rank == 0) std.debug.print("RESULT nccl all-reduce f32 {d:>9} values ({d:>9} bytes): {d:.4} ms\n", .{ count, count * 4, best }); + } +} diff --git a/zig/tests/cuda/flashnext_exl3/tp/exl3_draft_tests.zig b/zig/tests/cuda/flashnext_exl3/tp/exl3_draft_tests.zig new file mode 100644 index 0000000000..64e7bd58fa --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/tp/exl3_draft_tests.zig @@ -0,0 +1,159 @@ +//! exl3-draft: the MTP draft head rows of an EXL3 pack (cuda_exl3_mtp.draftRows) against the full EXL3 lm_head. + +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const flash = @import("flashnext"); +const check = @import("../../check.zig"); +const st = core.safetensors; +const exl3 = cuda.exl3; +const fmt = core.exl3_format; +const mtp3 = flash.exl3_mtp; +const Gpu = check.Gpu; + +fn bfVal(b: u16) f64 { + return @floatCast(@as(f32, @bitCast(@as(u32, b) << 16))); +} + +fn bf16Bits(v: f32) u16 { + return flash.ple.toBf16(v); +} + +pub fn run(gpu: Gpu, rest: []const [:0]const u8) !void { + if (rest.len < 1) return error.MissingArgument; + const dir = rest[0]; + const threads: usize = if (rest.len > 1) try std.fmt.parseInt(usize, rest[1], 10) else mtp3.loadThreads(); + var arena_state = std.heap.ArenaAllocator.init(gpu.gpa); + defer arena_state.deinit(); + const a = arena_state.allocator(); + const io = gpu.io; + const path = try std.fs.path.join(a, &.{ dir, "model-00007-of-00007.safetensors" }); + var file = try st.File.open(gpu.gpa, io, path); + defer file.close(io); + const tr = file.get("lm_head.trellis") orelse return error.MissingTensor; + const su = file.get("lm_head.suh") orelse return error.MissingTensor; + const sv = file.get("lm_head.svh") orelse return error.MissingTensor; + const head = try mtp3.Head.of(tr, su, sv, file.get("lm_head.mul1")); + const k = head.k; + const n = head.n; + std.debug.print("lm_head: K {d}, N {d}, halves {d} ({d}-bit), {d} MB trellis\n", .{ k, n, head.halves, head.halves / 2, tr.bytes.len >> 20 }); + + const ids = try flash.mtp.draftIds(a); + const rows = try a.alloc(u16, ids.len * k); + const t0 = std.Io.Clock.awake.now(io).toNanoseconds(); + try mtp3.draftRows(gpu.gpa, head, ids, rows, threads); + const ms = @as(f64, @floatFromInt(std.Io.Clock.awake.now(io).toNanoseconds() - t0)) / 1e6; + std.debug.print("draft rows: {d} ids in {d} column blocks, {d} threads, {d:.0} ms (host, ReleaseFast)\n", .{ ids.len, blk: { + var c: usize = 0; + var prev: usize = std.math.maxInt(usize); + for (ids) |id| { + if (id / 128 != prev) c += 1; + prev = id / 128; + } + break :blk c; + }, threads, ms }); + + // x: a real embedding row (bf16) and a random bf16 row + const embed_path = try std.fs.path.join(a, &.{ dir, "model-00001-of-00007.safetensors" }); + var ef = try st.File.open(gpu.gpa, io, embed_path); + defer ef.close(io); + const emb = ef.get("model.language_model.embed_tokens.weight") orelse return error.MissingTensor; + const xs_bits = try a.alloc(u16, 2 * k); + for (0..k) |i| xs_bits[i] = std.mem.readInt(u16, emb.bytes[(9338 * k + i) * 2 ..][0..2], .little); + var prng = std.Random.DefaultPrng.init(0xd1a); + for (xs_bits[k..]) |*b| b.* = bf16Bits((prng.random().float(f32) - 0.5) * 2); + + // full head on the GPU, fp32 out + var s = try cuda.Stream.init(gpu.d, false); + defer s.deinit(); + var de = try exl3.Dense.load(gpu.d); + defer de.unload(); + const strides = exl3.strides(k, n, head.halves, false); + const trellis_al = try a.alloc(i16, tr.bytes.len / 2); + @memcpy(std.mem.sliceAsBytes(trellis_al), tr.bytes); + var d_tr = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(trellis_al)); + defer d_tr.free(); + var d_su = try cuda.DeviceBuffer.fromHost(gpu.d, su.bytes); + defer d_su.free(); + var d_sv = try cuda.DeviceBuffer.fromHost(gpu.d, sv.bytes); + defer d_sv.free(); + const wgt: exl3.Weight = .{ .trellis = d_tr.ptr, .suh = d_su.ptr, .svh = d_sv.ptr, .k = @intCast(k), .n = @intCast(n), .halves = head.halves, .codebook = .mul1, .stride_k = strides.k, .stride_nb = strides.nb }; + var d_x = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(xs_bits)); + defer d_x.free(); + var d_y = try cuda.DeviceBuffer.alloc(gpu.d, 2 * n * 4); + defer d_y.free(); + try gpu.ctx.synchronize(); + try de.decode(s, wgt, d_x.ptr, .bf16, 2, d_y.ptr, .f32); + try s.synchronize(); + const full = try a.alloc(f32, 2 * n); + try d_y.download(0, std.mem.sliceAsBytes(full)); + + // exact float64 for the first x (the exl3-pack test's `exact` product, one row) + const exact = try a.alloc(f64, n); + { + const tiles_k = k / 16; + const tiles_n = n / 16; + const tw: []const i16 = trellis_al; + const wq = try a.alloc(f16, k * n); + fmt.unpack(.{ .halves = head.halves }, .mul1, tw, tiles_k, tiles_n, wq); + const xr = try a.alloc(f64, k); + for (xr, 0..) |*v, i| v.* = bfVal(xs_bits[i]) * @as(f64, @floatCast(@as(f16, @bitCast(std.mem.readInt(u16, su.bytes[2 * i ..][0..2], .little))))); + fmt.rotate(xr, 1, k, 1); + @memset(exact, 0); + for (0..k) |ki| { + const xv = xr[ki]; + for (exact, wq[ki * n ..][0..n]) |*o, wv| o.* += xv * @as(f64, @floatCast(wv)); + } + fmt.rotate(exact, 1, n, 1); + for (exact, 0..) |*v, c| v.* *= @as(f64, @floatCast(@as(f16, @bitCast(std.mem.readInt(u16, sv.bytes[2 * c ..][0..2], .little))))); + } + + for (0..2) |xi| { + const x = xs_bits[xi * k ..][0..k]; + var worst_gpu: f64 = 0; // in units of the bf16 bound + var worst_exact: f64 = 0; + var max_abs_gpu: f64 = 0; + var peak: f64 = 0; + var arg_draft: usize = 0; + var arg_full: usize = 0; + var best_draft = -std.math.inf(f64); + var best_full = -std.math.inf(f64); + var gpu_exact_rel: f64 = 0; + var peak_full: f64 = 0; + for (full[xi * n ..][0..n]) |v| peak_full = @max(peak_full, @abs(@as(f64, v))); + for (ids, 0..) |id, i| { + const row = rows[i * k ..][0..k]; + var got: f64 = 0; + var mag: f64 = 0; + for (row, x) |wb, xb| { + const wv = bfVal(wb); + got += wv * bfVal(xb); + mag += @abs(wv * bfVal(xb)); + } + const bound = mag / 512.0 * (1.0 + 1.0 / 128.0) + 1e-9; // |W - bf16(W)| <= 2^-9 |bf16(W)| (1 + 2^-7) + const g: f64 = full[xi * n + id]; + // plus the kernel's fp16-input error, measured at 1.75e-4 - 2.4e-4 of the row's peak (exl3-pack) + const tol_gpu = bound + 5e-4 * peak_full; + worst_gpu = @max(worst_gpu, @abs(got - g) / tol_gpu); + max_abs_gpu = @max(max_abs_gpu, @abs(got - g)); + peak = @max(peak, @abs(g)); + if (got > best_draft) { + best_draft = got; + arg_draft = id; + } + if (g > best_full) { + best_full = g; + arg_full = id; + } + if (xi == 0) { + worst_exact = @max(worst_exact, @abs(got - exact[id]) / bound); + gpu_exact_rel = @max(gpu_exact_rel, @abs(g - exact[id])); + } + } + std.debug.print("x {s}: draft vs full head (GPU Dense fp32): max |diff| {e:.3} at draft peak {d:.3} (row peak {d:.3}), worst diff / tolerance {d:.3}; argmax over draft ids: draft {d}, full {d}\n", .{ if (xi == 0) "embed row 9338" else "random", max_abs_gpu, peak, peak_full, worst_gpu, arg_draft, arg_full }); + if (xi == 0) std.debug.print(" draft vs exact float64: worst diff / bf16 bound {d:.3}; GPU fp32 vs exact max |diff| {e:.3}\n", .{ worst_exact, gpu_exact_rel }); + try check.expect(worst_gpu <= 1.0, "draft logits off the full EXL3 head beyond bf16 weight rounding", .{}); + if (xi == 0) try check.expect(worst_exact <= 1.0, "draft logits off the exact product beyond the bf16 bound", .{}); + } + check.pass("EXL3 draft head: {d} rows dequantised from lm_head; every draft logit equals the full head's (GPU Dense fp32 and exact f64) within bf16 weight rounding", .{ids.len}); +} diff --git a/zig/tests/cuda/flashnext_exl3/tp/exl3_tp.zig b/zig/tests/cuda/flashnext_exl3/tp/exl3_tp.zig new file mode 120000 index 0000000000..19ef409826 --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/tp/exl3_tp.zig @@ -0,0 +1 @@ +../../../../src/families/flashnext/exl3/tp/exl3_tp.zig \ No newline at end of file diff --git a/zig/tests/cuda/flashnext_exl3/tp/flashnext_exl3_tp_engine_tests.zig b/zig/tests/cuda/flashnext_exl3/tp/flashnext_exl3_tp_engine_tests.zig new file mode 100644 index 0000000000..cb71c5f15c --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/tp/flashnext_exl3_tp_engine_tests.zig @@ -0,0 +1,22 @@ +//! flashnext-exl3-tpe: the Flash Next EXL3 pack through the family Engine at tp=1 or tp=2; this hub parses the mode. +const std = @import("std"); +const check = @import("../../check.zig"); +const decode = @import("flashnext_exl3_tpe_decode_tests.zig"); + +const Gpu = check.Gpu; + +pub fn run(gpu: Gpu, rest: []const [:0]const u8) !void { + if (rest.len < 5) return error.MissingArgument; + const rank = try std.fmt.parseInt(usize, rest[0], 10); + const world = try std.fmt.parseInt(usize, rest[1], 10); + const port = try std.fmt.parseInt(u16, rest[2], 10); + const dir = rest[3]; + const mode = rest[4]; + if ((world != 1 and world != 2) or rank >= world) return error.BadRank; + if (std.mem.eql(u8, mode, "bench")) return decode.benchMode(gpu, rank, world, port, dir); + if (rest.len < 7) return error.MissingArgument; + if (std.mem.eql(u8, mode, "tf")) return decode.tfMode(gpu, rank, world, port, dir, rest[5], rest[6], false); + if (std.mem.eql(u8, mode, "tfc")) return decode.tfMode(gpu, rank, world, port, dir, rest[5], rest[6], true); + if (std.mem.eql(u8, mode, "gen")) return decode.genMode(gpu, rank, world, port, dir, rest[5], rest[6]); + return error.UnknownMode; +} diff --git a/zig/tests/cuda/flashnext_exl3/tp/flashnext_exl3_tp_tests.zig b/zig/tests/cuda/flashnext_exl3/tp/flashnext_exl3_tp_tests.zig new file mode 100644 index 0000000000..bfc09aff68 --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/tp/flashnext_exl3_tp_tests.zig @@ -0,0 +1,452 @@ +//! flashnext-exl3-tp: EXL3 expert parallel (cuda_exl3_tp.zig) on two ranks, one process a GPU, through the Flash Next +const std = @import("std"); +const cuda = @import("cuda"); +const flash = @import("flashnext"); +const check = @import("../../check.zig"); +const T = @import("../model/flashnext_exl3_tests.zig"); + +const X = flash.exl3; +const tp = X.tp; +const wt = X.prompt_wt; +const pl = X.prompt_load; +const wire = flash.rank_wire; +const Collective = flash.collective.Collective; +const Gpu = check.Gpu; + +const dims = wt.dims; +const wide = wt.wide; +const streams = wt.streams; +const slots = wt.slots; + +// -- the two ranks' control link ----------------------------------------------------------------------------------- + +const Ranks = struct { + rank: usize, + link: wire.Link, + collective: Collective, + seq: u64 = 10, + + fn open(gpu: Gpu, rank: usize, port: u16) !Ranks { + const link = try wire.Link.open(gpu.io, "127.0.0.1", port, rank, 120_000); + errdefer link.close(); + // rank 0 does float64 work between collectives: a generous deadline (the waits are bounded by the stream) + return .{ .rank = rank, .link = link, .collective = try Collective.init(link, rank, 3_600_000) }; + } + + fn close(r: *Ranks) void { + r.collective.deinit(); + r.link.close(); + } + + /// Both ranks here; rank 1's `value` reaches rank 0 (returned there), rank 0's reaches rank 1. + fn swap(r: *Ranks, value: u64) !u64 { + r.seq += 1; + const until = wire.deadline(3_600_000); + if (r.rank == 0) { + const got = try r.link.receiveHeader(until); + if (got.kind != .ready or got.sequence != r.seq) return error.RankOutOfStep; + try r.link.sendHeader(.{ .kind = .ready, .sequence = r.seq, .arg = value }, until); + return got.arg; + } + try r.link.sendHeader(.{ .kind = .ready, .sequence = r.seq, .arg = value }, until); + const got = try r.link.receiveHeader(until); + if (got.kind != .ready or got.sequence != r.seq) return error.RankOutOfStep; + return got.arg; + } +}; + +// -- one decoder layer, stepped ------------------------------------------------------------------------------------- + +/// What one call of a layer leaves: the output streams (bf16), and the MoE's picks, weights and per-slot y. +const Out = struct { h: []u16, picks: []i32, wts: []f32, y: []f32 }; + +fn grab(comptime E: type, a: std.mem.Allocator, gpu: Gpu, p: *T.Pass, ptr: u64, n: usize) ![]E { + const got = try T.download(E, gpu, p.stream, ptr, n); + defer gpu.gpa.free(got); + return a.dupe(E, got); +} + +/// cuda_batch.zig Pass.chunk's loop body for one layer (T.deviceLayer's sequence), the MoE's outputs read back. +fn step(gpu: Gpu, a: std.mem.Allocator, p: *T.Pass, lw: *const pl.Layer, l: usize, ids: []const u32, pos0: usize, h_in: []const u16) !Out { + const s = p.s; + const n = ids.len; + p.rows = n; + p.pos0 = pos0; + p.pending = null; + try p.stage(ids); + try T.uploadTo(gpu, p.stream, s.h, std.mem.sliceAsBytes(h_in)); + try p.writeback(0, n); + try p.readout(lw.attn_hc, 0, n); + if (lw.linear) try p.gdnBlock(lw, l, wt.linearIndex(l), .{}) else try p.attnBlock(lw, wt.attnIndex(l)); + try p.tri.hcWriteback(s.h, s.h, s.pss, s.part, "*fp32", s.inj, s.dummy, "*bf16", s.dummy, "*bf16", n * dims, n, dims, streams, 4, 1, 1, wt.out_sk); + try p.readout(lw.mlp_hc, 0, n); + try p.moeBlock(lw); + var o: Out = undefined; + o.picks = try grab(i32, a, gpu, p, s.picks, n * slots); + o.wts = try grab(f32, a, gpu, p, s.wts, n * slots); + o.y = try grab(f32, a, gpu, p, s.y, n * slots * dims); + p.pending = l; + try p.writeback(0, n); + o.h = try grab(u16, a, gpu, p, s.h, n * wide); + return o; +} + +/// prefill 8 rows then 4 single decode rows (T.layerCheck's calls) +const calls = [_][2]usize{ .{ 0, 8 }, .{ 8, 1 }, .{ 9, 1 }, .{ 10, 1 }, .{ 11, 1 } }; + +fn embedRows(a: std.mem.Allocator, raw: *const T.RawPack, toks: []const u32) ![]u16 { + const table = try raw.get("model.language_model.embed_tokens.weight", .{}); + const hb = try a.alloc(u16, toks.len * wide); + for (toks, 0..) |tk, t| for (0..dims) |i| { + const bits = std.mem.readInt(u16, table.bytes[(tk * dims + i) * 2 ..][0..2], .little); + for (0..streams) |sidx| hb[t * wide + sidx * dims + i] = bits; + }; + return hb; +} + +/// The 5 calls on `lw` from a reset rig. +fn sequence(gpu: Gpu, a: std.mem.Allocator, rig: *T.Rig, raw: *const T.RawPack, lw: *const pl.Layer, l: usize) ![calls.len]Out { + try rig.reset(); + var p = rig.pass(lw[0..1]); + var outs: [calls.len]Out = undefined; + for (calls, &outs) |c, *o| { + const toks = T.default_ids[c[0] .. c[0] + c[1]]; + o.* = try step(gpu, a, &p, lw, l, toks, c[0], try embedRows(a, raw, toks)); + } + return outs; +} + +fn hashOf(o: Out) u64 { + var h = std.hash.Wyhash.init(0x7e_2); + h.update(std.mem.sliceAsBytes(o.h)); + h.update(std.mem.sliceAsBytes(o.picks)); + h.update(std.mem.sliceAsBytes(o.wts)); + h.update(std.mem.sliceAsBytes(o.y)); + return h.final(); +} + +fn layerMode(gpu: Gpu, rk: *Ranks, dir: []const u8, which: []const usize) !void { + const rank = rk.rank; + var rig: T.Rig = undefined; + try rig.init(gpu, dir, 8, 64); + defer rig.deinit(); + var raw = try T.RawPack.open(gpu, dir); + defer raw.close(gpu); + var failed: usize = 0; + for (which) |l| { + var arena: std.heap.ArenaAllocator = .init(gpu.gpa); + defer arena.deinit(); + const a = arena.allocator(); + var t0 = check.now(gpu.io); + var lw2 = try pl.loadLayerRank(gpu.gpa, @constCast(gpu.d), &rig.stream, &rig.store, rig.pack_fn, l, 2, rank, null); + defer lw2.free(gpu.gpa); + const x2 = lw2.exl3 orelse return error.UnexpectedTensor; + try check.expect(x2.tp != null and x2.tp.?.rank == rank, "layer {d}: rank {d} experts carry no rank", .{ l, rank }); + var held: usize = 0; + for (x2.trellis[0]) |ptr| held += @intFromBool(ptr != 0); + std.debug.print("== rank {d} layer {d} ({s}): tp=2 layer loaded in {d:.0} ms, {d} of 513 experts held (ids {d}..{d})\n", .{ rank, l, if (wt.linearLayer(l)) "GDN" else "attention", @as(f64, @floatFromInt(check.now(gpu.io) - t0)) / 1e6, held, x2.tp.?.lo, x2.tp.?.hi - 1 }); + try check.expect(held == x2.tp.?.hi - x2.tp.?.lo, "layer {d}: rank {d} holds {d} experts", .{ l, rank, held }); + + // tp=1 on rank 0 alone (rank 1 waits at the swap) + var ref: [calls.len]Out = undefined; + if (rank == 0) { + t0 = check.now(gpu.io); + var lw1 = try pl.loadLayerRank(gpu.gpa, @constCast(gpu.d), &rig.stream, &rig.store, rig.pack_fn, l, 1, 0, null); + defer lw1.free(gpu.gpa); + rig.collective = null; + ref = try sequence(gpu, a, &rig, &raw, &lw1, l); + std.debug.print(" tp=1 reference (all 513 experts on this card) in {d:.0} ms\n", .{@as(f64, @floatFromInt(check.now(gpu.io) - t0)) / 1e6}); + } + _ = try rk.swap(0); + rig.collective = &rk.collective; + t0 = check.now(gpu.io); + const got = try sequence(gpu, a, &rig, &raw, &lw2, l); + const tp_ms = @as(f64, @floatFromInt(check.now(gpu.io) - t0)) / 1e6; + // both ranks must hold the same bits (replicated streams): rank 1's hashes to rank 0 + var same_ranks = true; + for (got) |o| { + const theirs = try rk.swap(hashOf(o)); + if (rank == 0) same_ranks = same_ranks and theirs == hashOf(o); + } + if (rank == 1) { + check.pass("rank 1 layer {d}: tp=2 mirrored the 5 calls in {d:.0} ms", .{ l, tp_ms }); + } else { + try check.expect(same_ranks, "layer {d}: the ranks' outputs or sums differ in bits", .{l}); + var picks_equal = true; + var wts_equal = true; + var y_equal: usize = 0; + var y_total: usize = 0; + var h_equal: usize = 0; + var h_total: usize = 0; + for (ref, got) |want, have| { + picks_equal = picks_equal and std.mem.eql(i32, want.picks, have.picks); + wts_equal = wts_equal and std.mem.eql(u8, std.mem.sliceAsBytes(want.wts), std.mem.sliceAsBytes(have.wts)); + // every slot's fp32 row, bit for bit (the write-back sums them in pick order on both) + for (want.y, have.y) |u, v| { + y_equal += @intFromBool(@as(u32, @bitCast(u)) == @as(u32, @bitCast(v))); + y_total += 1; + } + for (want.h, have.h) |u, v| { + h_equal += @intFromBool(u == v); + h_total += 1; + } + } + std.debug.print(" tp=2 ({d:.0} ms for the 5 calls): picks {s}, weights {s}; per-slot fp32 MoE rows equal tp=1's on {d}/{d} values; layer output bf16 equal {d}/{d}\n", .{ tp_ms, if (picks_equal) "equal" else "DIFFER", if (wts_equal) "equal" else "DIFFER", y_equal, y_total, h_equal, h_total }); + var ok = true; + check.expect(picks_equal and wts_equal, "layer {d}: tp=2 router picks or weights differ from tp=1", .{l}) catch { + ok = false; + }; + check.expect(y_equal == y_total, "layer {d}: tp=2's per-slot MoE rows are not tp=1's bits ({d}/{d})", .{ l, y_equal, y_total }) catch { + ok = false; + }; + check.expect(h_equal == h_total, "layer {d}: tp=2's layer output is not tp=1's bits ({d}/{d})", .{ l, h_equal, h_total }) catch { + ok = false; + }; + if (ok) check.pass("layer {d} tp=2 vs tp=1: picks, weights, every per-slot MoE row and the layer output bit-identical ({d}/{d} values, {d}/{d} outputs); ranks bit-identical", .{ l, y_equal, y_total, h_equal, h_total }) else failed += 1; + } + // tp=2 against float64 (T.layerCheck on rank 0; rank 1 mirrors its device calls) + _ = try rk.swap(1); + if (rank == 0) { + T.layerCheck(gpu, &rig, &raw, &lw2, l, &T.default_ids) catch |err| if (err == error.TestFailed) { + failed += 1; + } else return err; + } else { + _ = try sequence(gpu, a, &rig, &raw, &lw2, l); + } + _ = try rk.swap(2); + } + const all = try rk.swap(failed); + try check.expect(failed == 0 and all == 0, "flashnext-exl3-tp layer: {d} checks failed on rank {d}", .{ failed, rank }); +} + +// -- the per-slot exchange's cost and bits --------------------------------------------------------------------------- + +fn fill(a: std.mem.Allocator, rows: usize, seed: u64) ![]f32 { + const out = try a.alloc(f32, rows * slots * dims); + var prng = std.Random.DefaultPrng.init(seed); + for (out) |*v| v.* = prng.random().floatNorm(f32) * 0.05; + return out; +} + +/// A nonblocking communicator takes its next call only once the last one's enqueue finished (state 0, not 7). +fn ready(c: *Collective) !void { + while (true) { + var state: cuda.nccl.Result = 0; + try c.lib.check(c.api.ncclCommGetAsyncError(c.comm, &state), "ncclCommGetAsyncError"); + if (state == 0) break; + if (state != 7) { + std.debug.print("ready: NCCL async state {d}\n", .{state}); + return c.lib.check(state, "NCCL async"); + } + } +} + +fn timeIt(s: cuda.Stream, e0: *cuda.Event, e1: *cuda.Event, reps: usize, ctx: anytype, comptime f: fn (@TypeOf(ctx)) anyerror!void) !f64 { + for (0..3) |_| try f(ctx); + try s.synchronize(); + var best: f64 = std.math.inf(f64); + for (0..5) |_| { + try e0.record(s); + for (0..reps) |_| try f(ctx); + try e1.record(s); + try e1.synchronize(); + best = @min(best, @as(f64, try e0.elapsedMs(e1.*)) / @as(f64, @floatFromInt(reps))); + } + return best; +} + +/// Each rank owns the slots of its parity and zeros elsewhere (moe's y rows); the all-reduce keeps each value. +fn sumMode(gpu: Gpu, rk: *Ranks) !void { + const rank = rk.rank; + var arena: std.heap.ArenaAllocator = .init(gpu.gpa); + defer arena.deinit(); + const a = arena.allocator(); + var s = try cuda.Stream.init(gpu.d, true); + defer s.deinit(); + var e0 = try cuda.Event.init(gpu.d, true); + defer e0.deinit(); + var e1 = try cuda.Event.init(gpu.d, true); + defer e1.deinit(); + var ops = try X.load.Ops.create(gpu.gpa, gpu.d); + defer ops.release(); + const t = try tp.scratchOf(ops); + const c = &rk.collective; + const sizes = [_]usize{ 1, 4, 128, 512, 2048 }; + var all_equal = true; + for (sizes) |rows| { + const n = rows * slots * dims; + const whole = try fill(a, rows, 0x5eed); + // negative zeros and denormals among the owned values: the sum with the other rank's +0 keeps their value + whole[3] = -0.0; + whole[n - 1] = std.math.floatMin(f32) / 4; + const mine = try a.alloc(f32, n); + for (0..rows * slots) |rs| for (0..dims) |d| { + mine[rs * dims + d] = if (rs % 2 == rank) whole[rs * dims + d] else 0; + }; + var buf = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(mine)); + defer buf.free(); + try c.sum(s, buf.ptr, n); + try c.finish(s); + const got = try a.alloc(f32, n); + try buf.download(0, std.mem.sliceAsBytes(got)); + var equal: usize = 0; + for (got, whole) |g, w| equal += @intFromBool(g == w); + all_equal = all_equal and equal == n; + // times: this exchange (the pass's y rows), the old one (the summed branch [rows, 2560]), the mask + const reps: usize = if (rows <= 128) 100 else if (rows <= 512) 20 else 5; + const Ctx = struct { c: *Collective, s: cuda.Stream, y: u64, n: usize, t: *const tp.Scratch, picks: u64, rank: usize, rows: usize }; + var picks = try cuda.DeviceBuffer.alloc(gpu.d, rows * slots * 4); + defer picks.free(); + try picks.fill8(0, s.handle); + try s.synchronize(); + const ctx: Ctx = .{ .c = c, .s = s, .y = buf.ptr, .n = n, .t = t, .picks = picks.ptr, .rank = rank, .rows = rows }; + const t_y = try timeIt(s, &e0, &e1, reps, ctx, struct { + fn f(x: Ctx) anyerror!void { + try x.c.sum(x.s, x.y, x.n); + } + }.f); + const t_old = try timeIt(s, &e0, &e1, reps, ctx, struct { + fn f(x: Ctx) anyerror!void { + try x.c.sum(x.s, x.y, x.n / slots); + } + }.f); + const t_mask = try timeIt(s, &e0, &e1, reps, ctx, struct { + fn f(x: Ctx) anyerror!void { + try x.t.mask(x.s, x.picks, x.picks, x.rows * slots, try tp.Rank.init(2, x.rank)); + } + }.f); + if (rank == 0) std.debug.print("RESULT tp2 per-slot exchange rows {d:>4}: all-reduce of y [rows, 11, 2560] ({d:>9} B) {d:.4} ms | the summed branch [rows, 2560] ({d:>9} B, the previous exchange) {d:.4} ms | mask {d:.4} ms | equal to the union of the ranks' slots {d}/{d}\n", .{ rows, n * 4, t_y, n * 4 / slots, t_old, t_mask, equal, n }); + _ = try rk.swap(rows); + try ready(c); + } + try check.expect(all_equal, "rank {d}: the all-reduce of zero-padded slot rows differs from the union of the ranks' slots", .{rank}); + check.pass("rank {d}: the fp32 all-reduce of zero-padded per-slot rows leaves every slot's value exactly at 1..2048 rows", .{rank}); +} + +// -- the head's vocabulary split (one card) ------------------------------------------------------------------------- + +fn headMode(gpu: Gpu, dir: []const u8) !void { + var rig: T.Rig = undefined; + try rig.init(gpu, dir, 8, 64); + defer rig.deinit(); + const d: *cuda.Driver = @constCast(gpu.d); + var arena: std.heap.ArenaAllocator = .init(gpu.gpa); + defer arena.deinit(); + const a = arena.allocator(); + var full = try X.projection_load.load(gpu.gpa, d, &rig.store, "lm_head", wt.vocab, dims); + defer full.free(); + var halves: [2]flash.projection.Projection = undefined; + for (&halves, 0..) |*h, r| h.* = try tp.headProjection(gpu.gpa, d, &rig.store, "lm_head", 2, r); + defer for (&halves) |*h| h.free(); + var p = rig.pass(&.{}); + var prng = std.Random.DefaultPrng.init(0x4ead); + var all_equal = true; + var e0 = try cuda.Event.init(gpu.d, true); + defer e0.deinit(); + var e1 = try cuda.Event.init(gpu.d, true); + defer e1.deinit(); + for ([_]usize{ 1, 5 }) |m| { + const xb = try a.alloc(u16, m * dims); + for (xb) |*v| { + const f: f32 = prng.random().floatNorm(f32); + v.* = @truncate(@as(u32, @bitCast(f)) >> 16); + } + var x = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(xb)); + defer x.free(); + var yf = try cuda.DeviceBuffer.alloc(gpu.d, m * wt.vocab * 2); + defer yf.free(); + var yh = try cuda.DeviceBuffer.alloc(gpu.d, m * wt.vocab * 2); + defer yh.free(); + _ = try p.lane(x.ptr, p.s.dummy, full, yf.ptr, p.s.dummy, m, 1); + try e0.record(p.stream); + _ = try p.lane(x.ptr, p.s.dummy, full, yf.ptr, p.s.dummy, m, 1); + try e1.record(p.stream); + try e1.synchronize(); + const t_full = try e0.elapsedMs(e1); + var t_half: f32 = 0; + for (halves, 0..) |h, r| { + // the gathered layout: rank r's half at columns [r * N/2, (r + 1) * N/2) of every row + const tmp = try a.alloc(u16, m * h.n); + var yr = try cuda.DeviceBuffer.alloc(gpu.d, m * h.n * 2); + defer yr.free(); + _ = try p.lane(x.ptr, p.s.dummy, h, yr.ptr, p.s.dummy, m, 1); + try e0.record(p.stream); + _ = try p.lane(x.ptr, p.s.dummy, h, yr.ptr, p.s.dummy, m, 1); + try e1.record(p.stream); + try e1.synchronize(); + t_half = @max(t_half, try e0.elapsedMs(e1)); + try yr.download(0, std.mem.sliceAsBytes(tmp)); + for (0..m) |row| { + const dst = (row * wt.vocab + r * h.n) * 2; + try yh.upload(dst, std.mem.sliceAsBytes(tmp[row * h.n ..][0..h.n])); + } + } + const got = try a.alloc(u16, m * wt.vocab); + try yh.download(0, std.mem.sliceAsBytes(got)); + const want = try a.alloc(u16, m * wt.vocab); + try yf.download(0, std.mem.sliceAsBytes(want)); + var equal: usize = 0; + var worst: f64 = 0; + for (got, want) |g, w| { + equal += @intFromBool(g == w); + worst = @max(worst, @abs(T.bfValue(g) - T.bfValue(w))); + } + all_equal = all_equal and equal == got.len; + std.debug.print("RESULT head split rows {d}: halves (124,160 columns a rank) joined in vocabulary order equal the whole head on {d}/{d} bf16 logits (max diff {e:.3}); whole {d:.3} ms, one half {d:.3} ms\n", .{ m, equal, got.len, worst, t_full, t_half }); + } + try check.expect(all_equal, "head: the halves' logits differ from the whole head's", .{}); + check.pass("head: the EXL3 lm_head split at 128-column blocks (2 x 124,160) gives the whole head's bf16 logits bit for bit", .{}); +} + +// -- memory --------------------------------------------------------------------------------------------------------- + +fn gib(b: u64) f64 { + return @as(f64, @floatFromInt(b)) / (1 << 30); +} + +fn memoryMode(gpu: Gpu, dir: []const u8) !void { + const a = gpu.gpa; + const io = gpu.io; + const o = X.memory.options(); + const head = try X.memory.headBytes(a, io, dir, o); + const one = try tp.deviceBytes(a, io, dir, o, try tp.Rank.init(1, 0)); + const ranks = [2]u64{ try tp.deviceBytes(a, io, dir, o, try tp.Rank.init(2, 0)), try tp.deviceBytes(a, io, dir, o, try tp.Rank.init(2, 1)) }; + const all_dev = try tp.deviceBytes(a, io, dir, .{ .routed_on_host = false }, try tp.Rank.init(2, 0)); + std.debug.print("RESULT memory (routed trellis {s}): tp=1 main {d:.3} GiB; tp=2 rank 0 {d:.3} GiB + MTP head {d:.3} GiB, rank 1 {d:.3} GiB; (rank 0 with its 256 routed experts resident: {d:.3} GiB)\n", .{ if (o.routed_on_host) "on the host (offload)" else "resident", gib(one), gib(ranks[0]), gib(head), gib(ranks[1]), gib(all_dev) }); + // the lm_head's split would save half its bytes on each rank + var store = try wt.Store.open(a, io, dir); + defer store.deinit(); + const lm = try X.load.mat(&store, "lm_head"); + const head_bytes: u64 = lm.trellis.len + lm.suh.len + lm.svh.len; + std.debug.print("RESULT memory lm_head {d:.3} GiB a rank replicated; split: {d:.3} GiB a rank\n", .{ gib(head_bytes), gib(head_bytes / 2) }); + const prompt = flash.prompt; + const window: usize = 200_000; + const seq = prompt.seqBytes(window); + const kv = wt.attn_n * window * 2 * prompt.kv_row_bytes; + const kv8 = kv / 2 + wt.attn_n * window * 2 * wt.kv_heads * 4; // int8 rows plus an fp32 scale a row and head + const draft = flash.mtp.Draft.bytes(window + flash.mtp.max_rows, 1, flash.prompt.batch.default_chunk); + const scratch = flash.prompt.batch.Scratch.bytes(flash.prompt.batch.default_chunk); + std.debug.print("RESULT memory a sequence at {d} tokens: {d:.3} GiB (bf16 K/V {d:.3} GiB of it; int8 K/V would be {d:.3}); MTP draft cache {d:.3} GiB; pass scratch at chunk {d}: {d:.3} GiB\n", .{ window, gib(seq), gib(kv), gib(kv8), gib(draft), flash.prompt.batch.default_chunk, gib(scratch) }); + check.pass("memory: per-rank placement read from the pack's headers", .{}); +} + +pub fn run(gpu: Gpu, rest: []const [:0]const u8) !void { + if (rest.len < 4) return error.MissingArgument; + const rank = try std.fmt.parseInt(usize, rest[0], 10); + const port = try std.fmt.parseInt(u16, rest[1], 10); + const dir = rest[2]; + const mode = rest[3]; + if (std.mem.eql(u8, mode, "head")) return headMode(gpu, dir); + if (std.mem.eql(u8, mode, "memory")) return memoryMode(gpu, dir); + var rk = try Ranks.open(gpu, rank, port); + defer rk.close(); + if (std.mem.eql(u8, mode, "sum")) return sumMode(gpu, &rk); + if (std.mem.eql(u8, mode, "layer")) { + var list: std.ArrayList(usize) = .empty; + defer list.deinit(gpu.gpa); + var it = std.mem.splitScalar(u8, if (rest.len > 4) rest[4] else "0,3,17", ','); + while (it.next()) |tok| if (tok.len > 0) try list.append(gpu.gpa, try std.fmt.parseInt(usize, tok, 10)); + return layerMode(gpu, &rk, dir, list.items); + } + return error.UnknownMode; +} diff --git a/zig/tests/cuda/flashnext_exl3/tp/flashnext_exl3_tpe_decode_tests.zig b/zig/tests/cuda/flashnext_exl3/tp/flashnext_exl3_tpe_decode_tests.zig new file mode 100644 index 0000000000..dea2050339 --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/tp/flashnext_exl3_tpe_decode_tests.zig @@ -0,0 +1,212 @@ +//! flashnext-exl3-tpe: teacher-forced, free-running and timed decode through the engine. +const std = @import("std"); +const cuda = @import("cuda"); +const flash = @import("flashnext"); +const check = @import("../../check.zig"); + +const Gpu = check.Gpu; +const prompt = flash.prompt; +const X = flash.exl3; +const xpass = X.pass; +const tp = X.tp; +const fpool = flash.exl3_pool; +const wire = flash.rank_wire; +const Collective = flash.collective.Collective; +const vocab = prompt.vocabulary; + +const engine = @import("flashnext_exl3_tpe_engine.zig"); +const cap_default = engine.cap_default; +const out_rows = engine.out_rows; +const promote = engine.promote; +const ms = engine.ms; +const gib = engine.gib; +const Ranks = engine.Ranks; +const Eng = engine.Eng; +const printPool = engine.printPool; +const Golden = engine.Golden; +const ints = engine.ints; +const num = engine.num; +const top1 = engine.top1; +const margin = engine.margin; + +pub fn tfMode(gpu: Gpu, rank: usize, world: usize, port: u16, dir: []const u8, golden_path: []const u8, dump: []const u8, cont: bool) !void { + const a = gpu.gpa; + var golden = try Golden.open(gpu, golden_path); + defer golden.parsed.deinit(); + const p = golden.prompts()[0]; + const prompt_ids = try ints(a, p.object.get("ids").?); + defer a.free(prompt_ids); + const greedy = try ints(a, p.object.get("greedy_ids").?); + defer a.free(greedy); + const ids = try a.alloc(u32, prompt_ids.len + if (cont) greedy.len - 1 else 0); + defer a.free(ids); + @memcpy(ids[0..prompt_ids.len], prompt_ids); + if (cont) @memcpy(ids[prompt_ids.len..], greedy[0 .. greedy.len - 1]); + const tf = p.object.get("teacher_forced").?.array.items; + const g = try Eng.init(gpu, rank, world, port, dir); + defer g.deinit(); + try g.reset(); + const n = ids.len; + const ops = g.e.layers[0].exl3.?.ops; + const before = ops.fingerprint(); + const t0 = check.now(gpu.io); + try g.forward(ids, .{ .logits = .all, .logits_out = g.out.ptr }); + const t1 = check.now(gpu.io); + const logits = try a.alloc(u16, n * vocab); + defer a.free(logits); + try g.out.download(0, std.mem.sliceAsBytes(logits)); + std.debug.print("rank {d}/{d}: teacher-forced pass of {d} rows in {d:.1} ms (cold pool)\n", .{ rank, world, n, ms(t0, t1) }); + _ = try g.round(prompt.argmax(g.seq.logits)); + const after = ops.fingerprint(); + try check.expect(std.mem.eql(u8, std.mem.sliceAsBytes(&before), std.mem.sliceAsBytes(&after)), "rank {d}: an EXL3 scratch buffer grew inside a pass", .{rank}); + check.pass("rank {d}/{d}: EXL3 scratch (incl. expert parallel's) kept its addresses through a {d}-row prompt pass and a decode round", .{ rank, world, n }); + // both ranks hold the same gathered logits + const h = std.hash.Wyhash.hash(0, std.mem.sliceAsBytes(logits)); + const other = try g.swap(h); + if (world == 2) try check.expect(other == h, "the two ranks' gathered logits differ", .{}); + if (rank != 0) return; + // vs the golden + var ok: usize = 0; + var worst: f64 = 0; + for (tf, 0..) |row, pos| { + const r = logits[pos * vocab ..][0..vocab]; + const g5 = try ints(a, row.object.get("top5_ids").?); + defer a.free(g5); + ok += @intFromBool(top1(r) == g5[0]); + for (g5, row.object.get("top5_logits").?.array.items) |id, lv| worst = @max(worst, @abs(@as(f64, promote(r[id])) - num(lv))); + } + std.debug.print("RESULT tp={d} TF vs golden: top-1 {d}/{d}, max |our logit - golden| on the golden top-5 {d:.4}\n", .{ world, ok, tf.len, worst }); + if (world == 1) { + try std.Io.Dir.cwd().writeFile(gpu.io, .{ .sub_path = dump, .data = std.mem.sliceAsBytes(logits) }); + check.pass("tp=1 logits written to {s}", .{dump}); + return; + } + const ref_bytes = try std.Io.Dir.cwd().readFileAlloc(gpu.io, dump, a, .limited(1 << 30)); + defer a.free(ref_bytes); + if (ref_bytes.len != logits.len * 2) return error.UnexpectedTensor; + const ref = try a.alloc(u16, logits.len); + defer a.free(ref); + @memcpy(std.mem.sliceAsBytes(ref), ref_bytes); + var top_eq: usize = 0; + var bits: usize = 0; + var dmax: f64 = 0; + var dsum: f64 = 0; + std.debug.print(" pos tp=1 top1 (margin) tp=2 top1 (margin) max|d| over the row bit-equal of {d}\n", .{vocab}); + for (0..n) |pos| { + const x = ref[pos * vocab ..][0..vocab]; + const y = logits[pos * vocab ..][0..vocab]; + var row_d: f64 = 0; + var row_bits: usize = 0; + for (x, y) |u, v| { + row_bits += @intFromBool(u == v); + const d = @abs(@as(f64, promote(u)) - promote(v)); + row_d = @max(row_d, d); + dsum += d; + } + bits += row_bits; + dmax = @max(dmax, row_d); + const same = top1(x) == top1(y); + top_eq += @intFromBool(same); + std.debug.print(" {d:3} {d:9} ({d:6.3}) {d:9} ({d:6.3}) {d:.4} {d}{s}\n", .{ pos, top1(x), margin(x), top1(y), margin(y), row_d, row_bits, if (same) "" else " TOP-1 DIFFERS" }); + } + std.debug.print("RESULT tp=2 vs tp=1 logits: top-1 equal {d}/{d}; bf16 bit-equal {d}/{d} ({d:.3}%); max |d| {d:.4}, mean |d| {e:.3}\n", .{ top_eq, n, bits, n * vocab, 100 * @as(f64, @floatFromInt(bits)) / @as(f64, @floatFromInt(n * vocab)), dmax, dsum / @as(f64, @floatFromInt(n * vocab)) }); + try check.expect(top_eq == n, "tp=2 top-1 differs from tp=1", .{}); + check.pass("tp=2 teacher-forced top-1 equals tp=1 at all {d} positions", .{n}); + printPool(g, "after tf"); +} + +pub fn genMode(gpu: Gpu, rank: usize, world: usize, port: u16, dir: []const u8, golden_path: []const u8, dump: []const u8) !void { + const a = gpu.gpa; + var golden = try Golden.open(gpu, golden_path); + defer golden.parsed.deinit(); + const g = try Eng.init(gpu, rank, world, port, dir); + defer g.deinit(); + var text: std.ArrayList(u8) = .empty; + defer text.deinit(a); + var all_golden = true; + for (golden.prompts()) |p| { + const ids = try ints(a, p.object.get("ids").?); + defer a.free(ids); + const want = try ints(a, p.object.get("greedy_ids").?); + defer a.free(want); + const got = try a.alloc(u32, want.len); + defer a.free(got); + try g.reset(); + try g.forward(ids, .{ .logits = .last }); + got[0] = prompt.argmax(g.seq.logits); + const t0 = check.now(gpu.io); + for (1..got.len) |i| got[i] = try g.round(got[i - 1]); + const step = ms(t0, check.now(gpu.io)) / @as(f64, @floatFromInt(got.len - 1)); + const h = std.hash.Wyhash.hash(0, std.mem.sliceAsBytes(got)); + const other = try g.swap(h); + if (world == 2) try check.expect(other == h, "the two ranks' greedy ids differ", .{}); + var first: ?usize = null; + for (got, want, 0..) |x, y, i| if (x != y and first == null) { + first = i; + }; + if (first != null) all_golden = false; + for (got) |x| try text.print(a, "{d},", .{x}); + try text.append(a, '\n'); + if (rank == 0) { + std.debug.print("GEN tp={d} {s}: ", .{ world, p.object.get("name").?.string }); + for (got) |x| std.debug.print("{d},", .{x}); + std.debug.print(" {s} ({d:.1} ms/step)\n", .{ if (first) |_| "differs from the golden" else "equal to the golden", step }); + if (first) |d| std.debug.print(" first difference from the golden at step {d}: ours {d}, golden {d}\n", .{ d, got[d], want[d] }); + } + } + if (rank != 0) return; + printPool(g, "after gen"); + if (world == 1) { + try std.Io.Dir.cwd().writeFile(gpu.io, .{ .sub_path = dump, .data = text.items }); + check.pass("tp=1 greedy ids written to {s}", .{dump}); + return; + } + const ref = try std.Io.Dir.cwd().readFileAlloc(gpu.io, dump, a, .limited(1 << 20)); + defer a.free(ref); + const same = std.mem.eql(u8, ref, text.items); + std.debug.print("RESULT tp=2 greedy 24 on the {d} golden prompts: {s} tp=1's ids; all equal the golden: {}\n", .{ golden.prompts().len, if (same) "identical to" else "DIFFERENT from", all_golden }); + try check.expect(same, "tp=2 greedy ids differ from tp=1's", .{}); + check.pass("tp=2 greedy 24 tokens identical to tp=1 on all prompts", .{}); +} + +pub fn benchMode(gpu: Gpu, rank: usize, world: usize, port: u16, dir: []const u8) !void { + const a = gpu.gpa; + const g = try Eng.init(gpu, rank, world, port, dir); + defer g.deinit(); + const max_n: usize = 2048; + const ids = try a.alloc(u32, max_n); + defer a.free(ids); + var prng = std.Random.DefaultPrng.init(0x5eed); + for (ids) |*x| x.* = 200 + prng.random().uintLessThan(u32, 60000); + try g.reset(); + try g.forward(ids[0..64], .{ .logits = .last }); + for ([_]usize{ 16, 256, 2048 }) |n| { + var best: f64 = std.math.inf(f64); + for (0..2) |_| { + try g.reset(); + _ = try g.swap(0); + const before = g.stats(); + const t0 = check.now(gpu.io); + try g.forward(ids[0..n], .{ .logits = .last }); + const t = ms(t0, check.now(gpu.io)); + const after = g.stats(); + const hits = after.hit_slots - before.hit_slots; + const all = hits + after.miss_slots - before.miss_slots; + std.debug.print(" rank {d}/{d} prefill {d:5} tokens: {d:8.1} ms = {d:7.1} rows/s (hit {d:.1}%, {d} copies)\n", .{ rank, world, n, t, @as(f64, @floatFromInt(n)) / (t / 1e3), 100 * @as(f64, @floatFromInt(hits)) / @as(f64, @floatFromInt(@max(all, 1))), after.misses - before.misses }); + best = @min(best, t); + } + std.debug.print("BENCH rank {d}/{d} prefill {d} tokens: best {d:.1} ms, {d:.1} rows/s\n", .{ rank, world, n, best, @as(f64, @floatFromInt(n)) / (best / 1e3) }); + var tok = prompt.argmax(g.seq.logits); + _ = try g.swap(1); + const before = g.stats(); + const t0 = check.now(gpu.io); + const steps: usize = 32; + for (0..steps) |_| tok = try g.round(tok); + const t = ms(t0, check.now(gpu.io)) / steps; + const after = g.stats(); + const hits = after.hit_slots - before.hit_slots; + const all = hits + after.miss_slots - before.miss_slots; + std.debug.print("BENCH rank {d}/{d} decode after {d} tokens: {d:.2} ms/token over {d} steps (this rank's hit {d:.1}%, {d:.1} copies a token)\n", .{ rank, world, n, t, steps, 100 * @as(f64, @floatFromInt(hits)) / @as(f64, @floatFromInt(@max(all, 1))), @as(f64, @floatFromInt(after.misses - before.misses)) / steps }); + } + printPool(g, "after bench"); +} diff --git a/zig/tests/cuda/flashnext_exl3/tp/flashnext_exl3_tpe_engine.zig b/zig/tests/cuda/flashnext_exl3/tp/flashnext_exl3_tpe_engine.zig new file mode 100644 index 0000000000..51f25aa643 --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/tp/flashnext_exl3_tpe_engine.zig @@ -0,0 +1,214 @@ +//! flashnext-exl3-tpe: the ranks' link, the engine wrapper and the golden reader the modes share. +const std = @import("std"); +const cuda = @import("cuda"); +const flash = @import("flashnext"); +const check = @import("../../check.zig"); + +const Gpu = check.Gpu; +const prompt = flash.prompt; +const X = flash.exl3; +const xpass = X.pass; +const tp = X.tp; +const fpool = flash.exl3_pool; +const wire = flash.rank_wire; +const Collective = flash.collective.Collective; +const vocab = prompt.vocabulary; +pub const cap_default: usize = 32768; +pub const out_rows: usize = 64; + +pub fn promote(bits: u16) f32 { + return @bitCast(@as(u32, bits) << 16); +} + +pub fn ms(t0: i96, t1: i96) f64 { + return @as(f64, @floatFromInt(t1 - t0)) / 1e6; +} + +pub fn gib(b: u64) f64 { + return @as(f64, @floatFromInt(b)) / (1 << 30); +} + +// -- the ranks' link ------------------------------------------------------------------------------------------------ + +pub const Ranks = struct { + rank: usize, + link: wire.Link, + collective: Collective, + seq: u64 = 10, + + pub fn open(gpu: Gpu, rank: usize, port: u16) !*Ranks { + const r = try gpu.gpa.create(Ranks); + errdefer gpu.gpa.destroy(r); + const link = try wire.Link.open(gpu.io, "127.0.0.1", port, rank, 600_000); + errdefer link.close(); + r.* = .{ .rank = rank, .link = link, .collective = try Collective.init(link, rank, 3_600_000) }; + return r; + } + + pub fn close(r: *Ranks, gpa: std.mem.Allocator) void { + r.collective.deinit(); + r.link.close(); + gpa.destroy(r); + } + + /// Both ranks here; returns the other rank's `value`. + pub fn swap(r: *Ranks, value: u64) !u64 { + r.seq += 1; + const until = wire.deadline(3_600_000); + if (r.rank == 0) { + const got = try r.link.receiveHeader(until); + if (got.kind != .ready or got.sequence != r.seq) return error.RankOutOfStep; + try r.link.sendHeader(.{ .kind = .ready, .sequence = r.seq, .arg = value }, until); + return got.arg; + } + try r.link.sendHeader(.{ .kind = .ready, .sequence = r.seq, .arg = value }, until); + const got = try r.link.receiveHeader(until); + if (got.kind != .ready or got.sequence != r.seq) return error.RankOutOfStep; + return got.arg; + } +}; + +// -- the engine ----------------------------------------------------------------------------------------------------- + +pub const Eng = struct { + gpu: Gpu, + rank: usize, + world: usize, + ranks: ?*Ranks, + stream: cuda.Stream, + glue: flash.Glue, + e: *prompt.Engine, + seq: prompt.Seq, + seq_live: bool, + out: cuda.DeviceBuffer, + attached: xpass.Attached, + + pub fn init(gpu: Gpu, rank: usize, world: usize, port: u16, dir: []const u8) !*Eng { + const g = try gpu.gpa.create(Eng); + errdefer gpu.gpa.destroy(g); + g.gpu = gpu; + g.rank = rank; + g.world = world; + g.ranks = if (world == 2) try Ranks.open(gpu, rank, port) else null; + g.stream = try cuda.Stream.init(gpu.d, true); + g.glue = try flash.Glue.load(gpu.d); + const t0 = check.now(gpu.io); + g.e = try prompt.Engine.initRank(gpu.gpa, gpu.io, @constCast(gpu.d), &g.stream, dir, world, rank); + if (g.ranks) |r| g.e.collective = &r.collective; + const t1 = check.now(gpu.io); + _ = try g.e.reserve(@min(cap_default, g.e.chunk)); + _ = try g.e.reserveScores(cap_default, @min(cap_default, g.e.chunk)); + g.seq = try g.e.newSeq(cap_default); + g.seq_live = true; + g.out = try cuda.DeviceBuffer.alloc(gpu.d, out_rows * vocab * 2); + g.attached = try xpass.attach(gpu.gpa, gpu.io, gpu.d, g.stream, g.e.layers[0..g.e.n_layers], null, g.e.chunk, try xpass.budgetFromEnv(0), null); + const t2 = check.now(gpu.io); + var free: usize = 0; + var total: usize = 0; + try gpu.d.check(gpu.d.api.cuMemGetInfo_v2(&free, &total), "cuMemGetInfo"); + const head_n = g.e.head_q.n; + std.debug.print("RESULT rank {d}/{d} load: {d:.1} s (initRank {d:.1} s, pinned bank {d:.2} GiB copied in {d:.1} s); {d} lines a layer x {d} layers of {d:.2} MiB = {d:.2} GiB of pool; {d:.2} GiB free before the pools, {d:.2} GiB free after (card {d:.2} GiB); head columns {d}; chunk {d}\n", .{ rank, world, ms(t0, t2) / 1e3, ms(t0, t1) / 1e3, gib(g.attached.bank_bytes), @as(f64, @floatFromInt(g.attached.bank_ns)) / 1e9, g.attached.lines, g.attached.layers, @as(f64, @floatFromInt(g.attached.line_bytes)) / (1 << 20), gib(g.attached.device_bytes), gib(g.attached.free_before), gib(free), gib(total), head_n, g.e.chunk }); + return g; + } + + pub fn deinit(g: *Eng) void { + g.out.free(); + if (g.seq_live) g.seq.deinit(g.gpu.gpa); + g.e.deinit(); + g.glue.deinit(); + g.stream.deinit(); + if (g.ranks) |r| r.close(g.gpu.gpa); + g.gpu.gpa.destroy(g); + } + + pub fn tri(g: *Eng) flash.Tri { + return .{ .own = &g.glue, .s = g.stream }; + } + + pub fn reset(g: *Eng) !void { + const s = &g.seq; + for ([_][]cuda.DeviceBuffer{ s.conv, s.state, s.kc, s.vc, s.ikc, s.pooled }) |bs| for (bs) |b| try b.fill8(0, g.stream.handle); + try s.tail.fill8(0, g.stream.handle); + try g.stream.synchronize(); + s.fed = 0; + s.history = .{ g.e.eos, g.e.eos }; + } + + pub fn stats(g: *Eng) fpool.Stats { + return xpass.poolStats(g.e.layers[0..g.e.n_layers], null); + } + + /// Both ranks reach this point; the other rank's value (or `v` at tp=1). + pub fn swap(g: *Eng, v: u64) !u64 { + return if (g.ranks) |r| r.swap(v) else v; + } + + pub fn forward(g: *Eng, ids: []const u32, hooks: prompt.batch.Hooks) !void { + try g.e.forward(flash.Tri, g.tri(), &g.seq, ids, hooks); + if (g.ranks) |r| try r.collective.finish(g.stream) else try g.stream.synchronize(); + } + + pub fn round(g: *Eng, tok: u32) !u32 { + const q = [_]*prompt.Seq{&g.seq}; + var t = tok; + try g.e.round(flash.Tri, g.tri(), &q, (&t)[0..1], g.out.ptr); + if (g.ranks) |r| try r.collective.finish(g.stream); + return prompt.argmax(g.seq.logits); + } +}; + +pub fn printPool(g: *Eng, what: []const u8) void { + const s = g.stats(); + const routed = s.hit_slots + s.miss_slots; + std.debug.print("RESULT rank {d}/{d} pool {s}: {d} rounds, {d} passes, hit rate {d:.1}% ({d}/{d} routed slots of this rank), {d} copies, {d} evictions; host cut {d:.1} ms over {d} blocks, plan {d:.1} ms\n", .{ g.rank, g.world, what, s.rounds, s.passes, 100 * @as(f64, @floatFromInt(s.hit_slots)) / @as(f64, @floatFromInt(@max(routed, 1))), s.hit_slots, routed, s.misses, s.evictions, @as(f64, @floatFromInt(xpass.pool_time.cut_ns)) / 1e6, xpass.pool_time.blocks, @as(f64, @floatFromInt(xpass.pool_time.plan_ns)) / 1e6 }); +} + +// -- golden --------------------------------------------------------------------------------------------------------- + +pub const Golden = struct { + parsed: std.json.Parsed(std.json.Value), + + pub fn open(gpu: Gpu, path: []const u8) !Golden { + const text = try std.Io.Dir.cwd().readFileAlloc(gpu.io, path, gpu.gpa, .limited(1 << 26)); + defer gpu.gpa.free(text); + return .{ .parsed = try std.json.parseFromSlice(std.json.Value, gpu.gpa, text, .{}) }; + } + + pub fn prompts(g: *const Golden) []std.json.Value { + return g.parsed.value.object.get("prompts").?.array.items; + } +}; + +pub fn ints(a: std.mem.Allocator, v: std.json.Value) ![]u32 { + const items = v.array.items; + const out = try a.alloc(u32, items.len); + for (items, out) |x, *o| o.* = @intCast(x.integer); + return out; +} + +pub fn num(v: std.json.Value) f64 { + return switch (v) { + .float => |f| f, + .integer => |i| @floatFromInt(i), + else => std.math.nan(f64), + }; +} + +pub fn top1(row: []const u16) u32 { + var best: usize = 0; + for (row, 0..) |b, i| if (promote(b) > promote(row[best])) { + best = i; + }; + return @intCast(best); +} + +// -- modes ---------------------------------------------------------------------------------------------------------- + +pub fn margin(row: []const u16) f32 { + const b = top1(row); + var second: f32 = -std.math.inf(f32); + for (row, 0..) |v, i| if (i != b) { + second = @max(second, promote(v)); + }; + return promote(row[b]) - second; +} diff --git a/zig/tests/cuda/flashnext_hc_tests.zig b/zig/tests/cuda/flashnext_hc_tests.zig index 974ce95d5e..5f590fe675 100644 --- a/zig/tests/cuda/flashnext_hc_tests.zig +++ b/zig/tests/cuda/flashnext_hc_tests.zig @@ -369,6 +369,51 @@ fn writebackRounding(r: *Rig, g: *const flash.Glue) !void { } } +fn reducedWriteback(r: *Rig, g: *const flash.Glue) !void { + const rows: usize = 17; + const count = rows * dims; + const hidden = r.bfs(rows * wide, 1); + const inject = r.bfs(rows * streams, 1); + const branch = r.f32s(count, -2, 2); + const edges = [_]f32{ 0.0, -0.0, 0x1.01p0, -0x1.01p0, 0x1.03p0, -0x1.03p0, 0x1p-126, -0x1p-126 }; + for (branch, 0..) |*value, index| if (index % 7 == 0) { + value.* = edges[index % edges.len]; + }; + hidden[1] = 0x8000; + inject[0] = 0x3f80; + branch[1] = -0.0; + const y = r.f32s(count * 11, -13, 13); + const weights = try r.a.alloc(f32, rows * 11); + @memset(weights, 0); + const ones = try r.a.alloc(f32, rows); + @memset(ones, 1); + for (0..rows) |row| { + @memcpy(y[row * 11 * dims ..][0..dims], branch[row * dims ..][0..dims]); + weights[row * 11] = 1; + } + const dh = try r.dev(std.mem.sliceAsBytes(hidden)); + const di = try r.dev(std.mem.sliceAsBytes(inject)); + const dy = try r.dev(std.mem.sliceAsBytes(y)); + const dw = try r.dev(std.mem.sliceAsBytes(weights)); + const db = try r.dev(std.mem.sliceAsBytes(branch)); + const du = try r.dev(std.mem.sliceAsBytes(ones)); + const reference = try r.out(hidden.len * 2); + const result = try r.out(hidden.len * 2); + const reference_sums = try r.out(rows * (dims / 256) * streams * 4); + const result_sums = try r.out(rows * (dims / 256) * streams * 4); + try g.hcWriteback(r.s, dh, reference, reference_sums, 0, di, dy, dw, count, rows, dims, streams, 2, 10, 11, 1); + for ([_]usize{ 0, 3, 16 }, [_]usize{ 3, 13, 1 }) |at, nrows| { + try g.hcWriteback(r.s, dh + at * wide * 2, result + at * wide * 2, result_sums + at * (dims / 256) * streams * 4, 0, di + at * streams * 2, db + at * dims * 4, du + at * 4, nrows * dims, nrows, dims, streams, 2, 0, 1, 1); + } + const want = try r.back(u16, reference, hidden.len); + const got = try r.back(u16, result, hidden.len); + try check.expect(std.mem.eql(u16, want, got), "reduced writeback: hidden BF16 bits equal the placed 11-slot path", .{}); + const want_sums = try r.back(u32, reference_sums, rows * (dims / 256) * streams); + const got_sums = try r.back(u32, result_sums, want_sums.len); + try check.expect(std.mem.eql(u32, want_sums, got_sums), "reduced writeback: FP32 square-sum bits equal the placed 11-slot path", .{}); + check.pass("reduced writeback: exact hidden BF16 and FP32 square sums, finite stale slots, signed zero, BF16 ties and short tail", .{}); +} + pub fn run(gpu: check.Gpu) !void { var arena: std.heap.ArenaAllocator = .init(gpu.gpa); defer arena.deinit(); @@ -377,6 +422,7 @@ pub fn run(gpu: check.Gpu) !void { defer r.deinit(); var g = try flash.Glue.load(gpu.d); defer g.deinit(); + try reducedWriteback(&r, &g); try writebackRounding(&r, &g); for ([_]usize{ 0, 2, 4 }) |mode| try writeback(&r, &g, mode); try readout(&r, &g); diff --git a/zig/tests/cuda/flashnext_model_rank_tests.zig b/zig/tests/cuda/flashnext_model_rank_tests.zig index ebb937ae71..6f0d80d30e 100644 --- a/zig/tests/cuda/flashnext_model_rank_tests.zig +++ b/zig/tests/cuda/flashnext_model_rank_tests.zig @@ -52,16 +52,16 @@ pub fn run(gpu: check.Gpu, dir: []const u8, rank: usize, port: u16, ids_path: [] const value: f32 = @bitCast(@as(u32, std.mem.readInt(u16, first.logits[i * 2 ..][0..2], .little)) << 16); try check.expect(std.math.isFinite(value), "finite initial logits at {d}", .{i}); } - const kept = try session.save(&first); + const kept = try session.save(&first, null); defer session.dropSaved(kept); const token = prompt.argmax(first.logits); try session.forward(flash.Tri, tri, &first, &.{token}, .{}); const expected = try gpu.gpa.dupe(u8, first.logits); defer gpu.gpa.free(expected); - try session.restore(&first, kept); + try session.restore(&first, null, kept); try session.forward(flash.Tri, tri, &first, &.{token}, .{}); try check.sameBytes("two-rank snapshot restore and replay", first.logits, expected); - try session.restore(&first, kept); + try session.restore(&first, null, kept); var second = try session.create(context); defer session.drop(&second); try session.forward(flash.Tri, tri, &second, ids.items, .{}); diff --git a/zig/tests/cuda/flashnext_window_tests.zig b/zig/tests/cuda/flashnext_window_tests.zig index 999d773cd0..c915a19d2f 100644 --- a/zig/tests/cuda/flashnext_window_tests.zig +++ b/zig/tests/cuda/flashnext_window_tests.zig @@ -45,7 +45,7 @@ pub fn run(gpu: check.Gpu, dir: []const u8, rank: usize, port: u16, ids_path: [] var serial = try session.create(context); defer session.drop(&serial); try session.forward(flash.Tri, tri, &serial, ids.items, .{}); - const base = try session.save(&serial); + const base = try session.save(&serial, null); defer session.dropSaved(base); const forced = [_]u32{ 37, 44, 102, 1, 320, 17, 19, 71, 99, 4, 80, 122, 7, 25, 61, 94 }; const row_bytes = prompt.vocabulary * 2; @@ -62,14 +62,14 @@ pub fn run(gpu: check.Gpu, dir: []const u8, rank: usize, port: u16, ids_path: [] const rows = try gpu.gpa.alloc(u8, expected.len); defer gpu.gpa.free(rows); for ([_]usize{ 1, 3, 7, 15, 16 }) |keep| { - try session.restore(&window, base); + try session.restore(&window, null, base); try session.verify(flash.Tri, tri, &window, &draft, &forced); try stream.synchronize(); try e.driver.check(e.driver.api.cuMemcpyDtoH_v2(rows.ptr, draft.logits.ptr, rows.len), "cuMemcpyDtoH"); try check.sameBytes("each two-rank window row equals its serial teacher-forced row", rows, expected); try session.keep(flash.Tri, tri, &window, &draft, keep); try check.expect(window.fed == ids.items.len + keep, "kept prefix position is exact", .{}); - try session.restore(&serial, base); + try session.restore(&serial, null, base); for (forced[0..keep]) |token| try session.forward(flash.Tri, tri, &serial, &.{token}, .{}); try check.expect(std.mem.eql(i64, &window.history, &serial.history), "kept host n-gram history matches serial", .{}); try session.forward(flash.Tri, tri, &serial, &.{13}, .{}); diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index 14adcb4c19..1f90b32f7e 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -97,10 +97,13 @@ const flashnext_hc_tests = @import("flashnext_hc_tests.zig"); const flashnext_attn_tests = @import("flashnext_attn_tests.zig"); const exl3_pack_tests = @import("flashnext_exl3/pack/exl3_pack_tests.zig"); const exl3_ple2_tests = @import("flashnext_exl3/model/exl3_ple2_tests.zig"); +const exl3_draft_tests = @import("flashnext_exl3/tp/exl3_draft_tests.zig"); const exl3_pool_tests = @import("flashnext_exl3/offload/exl3_pool_tests.zig"); const exl3_prefill_tests = @import("flashnext_exl3/prefill/exl3_prefill_tests.zig"); const flashnext_exl3_tests = @import("flashnext_exl3/model/flashnext_exl3_tests.zig"); +const flashnext_exl3_tp_tests = @import("flashnext_exl3/tp/flashnext_exl3_tp_tests.zig"); const flashnext_exl3_engine_tests = @import("flashnext_exl3/model/flashnext_exl3_engine_tests.zig"); +const flashnext_exl3_tp_engine_tests = @import("flashnext_exl3/tp/flashnext_exl3_tp_engine_tests.zig"); const ggml_tests = @import("ggml_tests.zig"); const paged_tests = @import("paged_tests.zig"); const glm_loader_gate = @import("glm_loader_gate.zig"); @@ -142,7 +145,9 @@ const usage = \\ flashnext-dense-hc dense connection weights and FP32 normalization against FP64 \\ flashnext-hc Flash Next's own hyper-connection and n-gram glue against float64 (synthetic rows) \\ flashnext-exl3 DIR names|norms|router [layers] [rows]|layer [layers] Flash Next on an EXL3 pack vs float64 + \\ flashnext-exl3-tp RANK PORT DIR layer [layers]|sum|head|memory EXL3 expert parallel on two ranks vs tp=1 and float64 \\ flashnext-exl3-engine DIR tf|gen GOLDEN [LINES] | bench [LINES] the family's Engine with the offload pool vs the reference golden + \\ flashnext-exl3-tpe RANK WORLD PORT DIR tf|gen GOLDEN DUMP | bench the Engine at tp=1/2 with per-rank pools \\ flashnext-attn Flash Next's own attention glue, pooling and block scores against float64 (synthetic) \\ flashnext-pages the attention glue's caches through a page table: paged == contiguous, float64, controls \\ info driver, device and library versions @@ -233,6 +238,7 @@ const usage = \\ exl3-tp2 [dir] EXL3 tensor parallel: rank shards vs the Python engine's, all-gathered (dir none skips them) \\ exl3-pack DIR [REPS] EXL3 Dense and grouped experts on a real EXL3 pack's shapes: float64 error and timing \\ exl3-ple2 DIR [ROW] [N] EXL3 n-gram table through cuda_ple.Ngram (ids, rows, RSS) and the PLE kernels vs float64 + \\ exl3-draft DIR [THREADS] EXL3 draft head rows dequantised from lm_head vs the full head (Dense) and float64 \\ exl3-prefill DIR experts|copy [LAYER]|dense EXL3 prefill pieces on the pack: grouped experts by rows, expert-major passes, line copies, prompt Dense \\ exl3-pool DIR [equal|bench|hit FILES] [LAYERS] Flash Next EXL3 expert offload pool on Experts' pointer tables vs fully resident \\ exl3-tiles [DIR] EXL3's tile candidate: codebooks and tiles bit-equal, FP64, rows through qlinear and Dense @@ -346,7 +352,9 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "flashnext-hc")) return flashnext_hc_tests.run(gpu); if (std.mem.eql(u8, cmd, "flashnext-attn")) return flashnext_attn_tests.run(gpu); if (std.mem.eql(u8, cmd, "flashnext-exl3")) return flashnext_exl3_tests.run(gpu, rest); + if (std.mem.eql(u8, cmd, "flashnext-exl3-tp")) return flashnext_exl3_tp_tests.run(gpu, rest); if (std.mem.eql(u8, cmd, "flashnext-exl3-engine")) return flashnext_exl3_engine_tests.run(gpu, rest); + if (std.mem.eql(u8, cmd, "flashnext-exl3-tpe")) return flashnext_exl3_tp_engine_tests.run(gpu, rest); if (std.mem.eql(u8, cmd, "flashnext-pages")) return @import("contract_fn_pages.zig").run(gpu); if (std.mem.eql(u8, cmd, "info")) return info(gpu); if (std.mem.eql(u8, cmd, "smoke")) return runtime_tests.smoke(gpu); @@ -402,8 +410,10 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "exl3")) return exl3_tests.run(gpu, if (rest.len > 0) rest[0] else "none"); if (std.mem.eql(u8, cmd, "exl3-pack")) return exl3_pack_tests.run(gpu, rest); if (std.mem.eql(u8, cmd, "exl3-ple2")) return exl3_ple2_tests.run(gpu, rest); + if (std.mem.eql(u8, cmd, "exl3-draft")) return exl3_draft_tests.run(gpu, rest); if (std.mem.eql(u8, cmd, "exl3-pool")) return exl3_pool_tests.run(gpu, rest); if (std.mem.eql(u8, cmd, "exl3-prefill")) return exl3_prefill_tests.run(gpu, rest); + if (std.mem.eql(u8, cmd, "nccl-lat")) return exl3_pack_tests.ncclLatency(gpu, rest); if (std.mem.eql(u8, cmd, "exl3-tiles")) return exl3_candidate_tests.run(gpu, if (rest.len > 0) rest[0] else null); if (std.mem.eql(u8, cmd, "exl3-tp2")) return exl3_tp2_tests.run(gpu, try arg(rest, 0), try arg(rest, 1), if (rest.len > 2) rest[2] else "none"); if (std.mem.eql(u8, cmd, "sample")) return sample_tests.draws(gpu); From 505a649e7d1ca794c5d74691ca046cc742761bd7 Mon Sep 17 00:00:00 2001 From: BlivionIaG <15323295+BlivionIaG@users.noreply.github.com> Date: Wed, 28 Oct 2026 16:23:30 +0100 Subject: [PATCH 3/3] EXL3 Flash Next: deep prefill Expert-major prefill for a Flash Next EXL3 checkpoint whose routed experts live in host memory: deep chunks of up to 4096 rows, the experts streamed through a ring of pool lines while the rows compute, and the dense projections of a chunk on tensor tiles. Every row is the decode kernels' bytes whatever the chunk or segment length: * Dense projections of 64 rows or more run tf_exl3_promptd_*: the rows are rotated by Linear.matmul's own tf_exl3_rot_in, each 64-row block sums the same chains matmul's plan cuts (sk splits of wk chains of k tiles, an fp32 mma chain from zero each, added in warp order and then split order) and finishes with its transform, scale and one rounding. * The 64-member expert tiles (tf_exl3pf_gu, tf_exl3pf_down) take 32 members a pass and sum K as the grouped kernels do (four chains of a split, then the splits, then 0 + each), with the same rotation and epilogue code. Transposed mma operands give the same fp32 bits as the decode kernels' roles. * The branch of a row is slice 2's canonical reduction at any tensor-parallel degree: 64 leaves of eight routed ids (the shared expert in the last), a leaf an fmaf chain in slot order from +0, a fixed balanced tree over the leaves. A rank sums the maximal complete subtrees of its contiguous leaf run (tf_exl3pf_tp_combine), the block partials travel in one explicit fp32 all-gather per layer (fixed size, [rows, 2560] a rank at a power of two) and every rank finishes the tree the same way (tf_exl3pf_tp_finish). One rank sums the whole tree and exchanges nothing; TP=1 and TP=2 write the same bits, and the host tests replay the leaf sums and the gather for 1 to 8 virtual ranks against the single-rank tree bit for bit. Tests: exl3-prefill dense and prompt compare the tiles with Linear.matmul and runSlots byte for byte on the pack's layers; exl3-prefill equal runs one prompt through the batched pass and through deep chunks of several sizes and compares the last row's logits and the greedy tokens; exl3-tpprefill requires the same bytes for deep against batched and for two ranks against one. tf_exl3pf_gu and tf_exl3pf_down are registered as replacements of the grouped decode form and pass gates 1 to 3 of the kernel contract on random layers, with a control; gate 4 stays a named finding. tf-cuda-test exl3-prefill-contract runs those gates alone. --- docs/kernel-contract.md | 2 +- zig/build/cuda.zig | 2 +- zig/build/cuda_kernels.zig | 1 + zig/build/exl3_prefill_body.zig | 16 + zig/kernels/cuda/exl3_candidate.cu | 114 ++++ .../flashnext_exl3/flashnext_exl3_prefill.cu | 476 +++++++++++++++ .../core/contract/flashnext_exl3_pending.zig | 20 +- zig/src/cuda/exl3.zig | 23 +- zig/src/cuda/exl3_candidate.zig | 38 +- zig/src/cuda/kernels.zig | 1 + zig/src/families/flashnext/cuda_batch.zig | 7 + .../families/flashnext/cuda_collective.zig | 20 +- zig/src/families/flashnext/cuda_lanes.zig | 4 +- .../families/flashnext/cuda_lanes_load.zig | 12 +- zig/src/families/flashnext/cuda_mtp_ops.zig | 1 + zig/src/families/flashnext/cuda_prompt.zig | 2 + .../flashnext/exl3/model/cuda_exl3.zig | 7 + .../flashnext/exl3/model/cuda_exl3_pass.zig | 14 +- .../exl3/model/cuda_exl3_telemetry.zig | 14 +- .../exl3/model/cuda_exl3_telemetry_test.zig | 13 + .../flashnext/exl3/model/exl3_telemetry.zig | 21 + .../exl3/model/exl3_telemetry_test.zig | 7 +- .../flashnext/exl3/offload/cuda_exl3_pool.zig | 4 + .../exl3/pack/exl3_startup_memory.zig | 7 +- .../exl3/prefill/cuda_exl3_prefill.zig | 333 +++++++++++ .../prefill/cuda_exl3_prefill_attribution.zig | 49 ++ .../cuda_exl3_prefill_attribution_test.zig | 310 ++++++++++ .../exl3/prefill/cuda_exl3_prefill_chunk.zig | 390 ++++++++++++ .../exl3/prefill/cuda_exl3_prefill_head.zig | 149 +++++ .../prefill/cuda_exl3_prefill_history.zig | 63 ++ .../exl3/prefill/cuda_exl3_prefill_k.zig | 243 ++++++++ .../exl3/prefill/cuda_exl3_prefill_leaves.zig | 187 ++++++ .../exl3/prefill/cuda_exl3_prefill_moe.zig | 302 ++++++++++ .../exl3/prefill/cuda_exl3_prefill_ring.zig | 562 ++++++++++++++++++ .../cuda_exl3_prefill_ring_fixture.zig | 57 ++ .../prefill/cuda_exl3_prefill_ringmoe.zig | 200 +++++++ .../exl3/prefill/cuda_exl3_prefill_state.zig | 293 +++++++++ .../exl3/prefill/cuda_exl3_prefill_test.zig | 302 ++++++++++ .../exl3/prefill/cuda_exl3_prefill_tiles.zig | 35 ++ .../cuda_exl3_prefill_writeback_test.zig | 131 ++++ .../exl3/prefill/exl3_prefill_window.zig | 96 +++ .../cuda/contract_flashnext_exl3_prefill.zig | 420 +++++++++++++ .../cuda/contract_flashnext_exl3_suite.zig | 4 + zig/tests/cuda/contract_limits.zig | 2 +- .../prefill/exl3_prefill_dense.zig | 26 +- .../prefill/exl3_prefill_engine.zig | 360 +++++++++++ .../prefill/exl3_prefill_experts.zig | 131 +++- .../prefill/exl3_prefill_tests.zig | 9 + .../tp/exl3_tpprefill_tests.zig | 339 +++++++++++ zig/tests/cuda/main.zig | 6 + zig/tests/exl3_prefill_body_compile.zig | 5 + 51 files changed, 5792 insertions(+), 38 deletions(-) create mode 100644 zig/build/exl3_prefill_body.zig create mode 100644 zig/kernels/cuda/flashnext_exl3/flashnext_exl3_prefill.cu create mode 100644 zig/src/families/flashnext/exl3/prefill/cuda_exl3_prefill.zig create mode 100644 zig/src/families/flashnext/exl3/prefill/cuda_exl3_prefill_attribution.zig create mode 100644 zig/src/families/flashnext/exl3/prefill/cuda_exl3_prefill_attribution_test.zig create mode 100644 zig/src/families/flashnext/exl3/prefill/cuda_exl3_prefill_chunk.zig create mode 100644 zig/src/families/flashnext/exl3/prefill/cuda_exl3_prefill_head.zig create mode 100644 zig/src/families/flashnext/exl3/prefill/cuda_exl3_prefill_history.zig create mode 100644 zig/src/families/flashnext/exl3/prefill/cuda_exl3_prefill_k.zig create mode 100644 zig/src/families/flashnext/exl3/prefill/cuda_exl3_prefill_leaves.zig create mode 100644 zig/src/families/flashnext/exl3/prefill/cuda_exl3_prefill_moe.zig create mode 100644 zig/src/families/flashnext/exl3/prefill/cuda_exl3_prefill_ring.zig create mode 100644 zig/src/families/flashnext/exl3/prefill/cuda_exl3_prefill_ring_fixture.zig create mode 100644 zig/src/families/flashnext/exl3/prefill/cuda_exl3_prefill_ringmoe.zig create mode 100644 zig/src/families/flashnext/exl3/prefill/cuda_exl3_prefill_state.zig create mode 100644 zig/src/families/flashnext/exl3/prefill/cuda_exl3_prefill_test.zig create mode 100644 zig/src/families/flashnext/exl3/prefill/cuda_exl3_prefill_tiles.zig create mode 100644 zig/src/families/flashnext/exl3/prefill/cuda_exl3_prefill_writeback_test.zig create mode 100644 zig/src/families/flashnext/exl3/prefill/exl3_prefill_window.zig create mode 100644 zig/tests/cuda/contract_flashnext_exl3_prefill.zig create mode 100644 zig/tests/cuda/flashnext_exl3/prefill/exl3_prefill_engine.zig create mode 100644 zig/tests/cuda/flashnext_exl3/tp/exl3_tpprefill_tests.zig create mode 100644 zig/tests/exl3_prefill_body_compile.zig diff --git a/docs/kernel-contract.md b/docs/kernel-contract.md index 7323e4b44e..c89ab1afd4 100644 --- a/docs/kernel-contract.md +++ b/docs/kernel-contract.md @@ -215,7 +215,7 @@ shared memory, the device's limits and the planted controls are checked the same A restored 13.0 receipt can arrive after the later compiler sets. Such records mark `later_pending` so 13.0 and 13.1 use their base usage, while 13.2 and later keep their independently measured or pending selection. A missing later receipt remains pending; a 13.0 measurement never substitutes for it. -Flash Next's 4 EXL3 launches (`flashnext_exl3_pending.zig`) carry only an nvcc 13.4 receipt: they stay pending under +Flash Next's 13 EXL3 launches (`flashnext_exl3_pending.zig`) carry only an nvcc 13.4 receipt: they stay pending under 13.0 to 13.3 and take their 13.4 usage from `cuda_134`, which `pending_132.lookup` serves for that release. Gemma 4 CUDA tower QKV norms and nine-term pooling register FP32 sums and BF16 outputs. diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index e53e74b715..897dfee7a9 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -350,7 +350,7 @@ pub fn hostTests(b: *std.Build, draft_ids: *std.Build.Module, build_options: *st b.step("test-api-boundaries", "Runtime thread initialization failure regression").dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = native_modules.api, .filters = &.{"Failure boundaries:"} })).step); b.step("test-review-api", "Ownership rollback: ownership and final batching regressions").dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = native_modules.api, .filters = &.{"Ownership rollback:"} })).step); const check_native = nativeServer(b, cuda_target, .debug, check_cuda, check_mods.core, check_mods.lanes, check_mods.nemotron, check_mods.heat, check_mods.qwen35, check_mods.qwen_vision, check_mods.tokenizer, check_mods.living, build_options, false, null); - b.step("check-native-cuda-host", "Compile the real native CUDA server without kernels or GPU execution").dependOn(&check_native.step); + @import("exl3_prefill_body.zig").check(b, cuda_target, check_native, flashnextModule(b, cuda_target, .debug, check_cuda, check_mods.core, check_mods.lanes, check_mods.qwen_vision, native_modules.api)); const offer_tests = b.addRunArtifact(b.addTest(.{ .root_module = b.createModule(.{ .root_source_file = b.path("zig/src/native/cuda_vision_offer.zig"), .target = host, .link_libc = true, .imports = &.{ .{ .name = "engine_api", .module = native_modules.api }, .{ .name = "nemotron", .module = mods.nemotron }, .{ .name = "qwen35", .module = mods.qwen35 } } }), .filters = if (video_filter) |f| &.{f} else &.{} })); offer_tests.has_side_effects = true; diff --git a/zig/build/cuda_kernels.zig b/zig/build/cuda_kernels.zig index 4faa679777..62e688d7f9 100644 --- a/zig/build/cuda_kernels.zig +++ b/zig/build/cuda_kernels.zig @@ -122,6 +122,7 @@ pub const kernels = [_]Kernel{ .{ .name = "drama3_codec", .flags = &.{"-O3"} }, .{ .name = "flashnext_exl3", .src = "flashnext_exl3/flashnext_exl3", .flags = glue }, // ours: Flash Next's EXL3 glue (an fp16 row split exactly into two bf16) .{ .name = "flashnext_exl3_tp", .src = "flashnext_exl3/flashnext_exl3_tp", .flags = glue }, // ours: Flash Next's EXL3 expert-parallel pick mask + .{ .name = "flashnext_exl3_prefill", .src = "flashnext_exl3/flashnext_exl3_prefill", .flags = &.{ "-O3", "--expt-relaxed-constexpr" } }, // ours: Flash Next EXL3 prefill experts, 64-member tiles any width }; /// torch.utils.cpp_extension's own nvcc flags (torch 2.13): C++20 and which half/bf16 operators the headers define. diff --git a/zig/build/exl3_prefill_body.zig b/zig/build/exl3_prefill_body.zig new file mode 100644 index 0000000000..a783451972 --- /dev/null +++ b/zig/build/exl3_prefill_body.zig @@ -0,0 +1,16 @@ +//! The native CUDA host check also compiles the EXL3 prefill body object. +const std = @import("std"); + +pub fn check(b: *std.Build, target: std.Build.ResolvedTarget, check_native: *std.Build.Step.Compile, flashnext: *std.Build.Module) void { + const native_check = b.step("check-native-cuda-host", "Compile the real native CUDA server without kernels or GPU execution"); + native_check.dependOn(&check_native.step); + const prefill_body = b.addObject(.{ .name = "exl3-prefill-body", .root_module = b.createModule(.{ + .root_source_file = b.path("zig/tests/exl3_prefill_body_compile.zig"), + .target = target, + .optimize = .debug, + .link_libc = true, + .imports = &.{.{ .name = "flashnext", .module = flashnext }}, + }) }); + _ = prefill_body.getEmittedBin(); + native_check.dependOn(&prefill_body.step); +} diff --git a/zig/kernels/cuda/exl3_candidate.cu b/zig/kernels/cuda/exl3_candidate.cu index a288cb5b45..6a59303741 100644 --- a/zig/kernels/cuda/exl3_candidate.cu +++ b/zig/kernels/cuda/exl3_candidate.cu @@ -1,5 +1,6 @@ // EXL3 projections on CUDA: trellis tiles as stored ([K/16, N/16, 16 bits] int16), core/exl3_format.zig's codebooks. #include "exl3_candidate.cuh" +#include "exl3_decode.cuh" struct TfExl3Rows { unsigned rows, x_stride, y_stride, y_f32; @@ -176,6 +177,119 @@ TF_EXL3_PROMPT(0) TF_EXL3_PROMPT(1) TF_EXL3_PROMPT(2) +// Where decode (exl3_linear.cu) cuts the K sum: chains of `per` k tiles, `wk` chains a split, added in that order. +struct TfExl3Chains { + unsigned per, wk; +}; + +// The prompt tile with decode's bytes at any row count: tf_exl3_rot_in's rows, chains summed as a decode warp does. +template +__device__ __forceinline__ void tf_exl3_promptd_job(const __half* Xh, unsigned rows, const uint8_t* T, const __half* svh, + const TfExl3Matrix m, const TfExl3Chains ch, uint8_t* Y, unsigned y_f32, + float* buf) { + constexpr unsigned TB = 32 * K; + const unsigned tid = threadIdx.x, sg = tid / 32, lane = tid % 32, g8 = lane / 4, t4 = lane % 4, ob = blockIdx.x; + const unsigned row0 = blockIdx.y * 64, kb = m.k / 128; + __half* xs = reinterpret_cast<__half*>(buf); + const uint64_t row_bytes = (uint64_t)m.tiles_n * TB; + const uint8_t* col = T + (uint64_t)(8 * ob + sg) * TB; + float acc[4][2][4], part[4][2][4], tot[4][2][4]; +#pragma unroll + for (int i = 0; i < 4; ++i) +#pragma unroll + for (int n = 0; n < 2; ++n) +#pragma unroll + for (int q = 0; q < 4; ++q) acc[i][n][q] = part[i][n][q] = tot[i][n][q] = 0.0f; + unsigned left = ch.per, w = 0, s = 0; + uint4 next = tf_exl3_words(reinterpret_cast(col), lane); + for (unsigned ib = 0; ib < kb; ++ib) { + for (unsigned i = tid; i < 64 * 16; i += 256) { + const unsigned r = i / 16, c = (i % 16) * 8; + const uint4 v = row0 + r < rows ? *reinterpret_cast(Xh + (uint64_t)(row0 + r) * m.k + ib * 128 + c) : make_uint4(0, 0, 0, 0); + *reinterpret_cast(xs + r * TF_EXL3_XS + c) = v; + } + __syncthreads(); + for (unsigned r = 0; r < 8; ++r) { + const uint4 now = next; + next = tf_exl3_words(reinterpret_cast(col + (uint64_t)(r < 7 ? 8 * ib + r + 1 : (ib + 1 < kb ? 8 * ib + 8 : 8 * ib + 7)) * row_bytes), lane); + __half wv[8]; + tf_exl3_unpack(now, lane, wv); + const uint32_t b[2][2] = {{tf_exl3_pair(wv[0], wv[1]), tf_exl3_pair(wv[2], wv[3])}, {tf_exl3_pair(wv[4], wv[5]), tf_exl3_pair(wv[6], wv[7])}}; +#pragma unroll + for (int i = 0; i < 4; ++i) { + const __half* x = xs + (16 * i + g8) * TF_EXL3_XS + 16 * r + 2 * t4; + const uint32_t a[4] = {*reinterpret_cast(x), *reinterpret_cast(x + 8 * TF_EXL3_XS), + *reinterpret_cast(x + 8), *reinterpret_cast(x + 8 * TF_EXL3_XS + 8)}; + tf_exl3_mma(acc[i][0], a, b[0][0], b[0][1]); + tf_exl3_mma(acc[i][1], a, b[1][0], b[1][1]); + } + if (--left == 0) { // a chain ends: into its split's warp sum, and the last chain of a split into the total + left = ch.per; +#pragma unroll + for (int i = 0; i < 4; ++i) +#pragma unroll + for (int n = 0; n < 2; ++n) +#pragma unroll + for (int q = 0; q < 4; ++q) { + part[i][n][q] = w == 0 ? acc[i][n][q] : __fadd_rn(part[i][n][q], acc[i][n][q]); + acc[i][n][q] = 0.0f; + } + if (++w == ch.wk) { + w = 0; +#pragma unroll + for (int i = 0; i < 4; ++i) +#pragma unroll + for (int n = 0; n < 2; ++n) +#pragma unroll + for (int q = 0; q < 4; ++q) tot[i][n][q] = s == 0 ? part[i][n][q] : __fadd_rn(tot[i][n][q], part[i][n][q]); + ++s; + } + } + } + __syncthreads(); + } +#pragma unroll + for (int i = 0; i < 4; ++i) + for (int n = 0; n < 2; ++n) + for (int h = 0; h < 2; ++h) { + const unsigned c = 16 * sg + 8 * n + 2 * t4; + buf[(16 * i + g8 + 8 * h) * 128 + c] = tot[i][n][2 * h]; + buf[(16 * i + g8 + 8 * h) * 128 + c + 1] = tot[i][n][2 * h + 1]; + } + __syncthreads(); + for (unsigned i = sg; i < 64 && row0 + i < rows; i += 8) { + float v[4]; + for (int q = 0; q < 4; ++q) v[q] = buf[i * 128 + 4 * lane + q]; + tf_exl3::fwht128(v, lane); + for (int q = 0; q < 4; ++q) { + const unsigned c = ob * 128 + 4 * lane + q; + tf_exl3_store(Y, (uint64_t)(row0 + i) * m.n + c, __fmul_rn(__fmul_rn(v[q], tf_exl3::HAD_SCALE), __half2float(svh[c])), y_f32); + } + } +} + +#define TF_EXL3_PROMPTD(C) \ + extern "C" __global__ void __launch_bounds__(256) tf_exl3_promptd_##C(const __half* Xh, unsigned rows, \ + const uint8_t* T, const __half* svh, \ + const TfExl3Matrix m, const TfExl3Chains ch, \ + uint8_t* Y, unsigned y_f32) { \ + __shared__ __align__(16) float buf[64 * 128]; \ + switch (m.bits) { \ + case 1: tf_exl3_promptd_job<1, C>(Xh, rows, T, svh, m, ch, Y, y_f32, buf); break; \ + case 2: tf_exl3_promptd_job<2, C>(Xh, rows, T, svh, m, ch, Y, y_f32, buf); break; \ + case 3: tf_exl3_promptd_job<3, C>(Xh, rows, T, svh, m, ch, Y, y_f32, buf); break; \ + case 4: tf_exl3_promptd_job<4, C>(Xh, rows, T, svh, m, ch, Y, y_f32, buf); break; \ + case 5: tf_exl3_promptd_job<5, C>(Xh, rows, T, svh, m, ch, Y, y_f32, buf); break; \ + case 6: tf_exl3_promptd_job<6, C>(Xh, rows, T, svh, m, ch, Y, y_f32, buf); break; \ + case 7: tf_exl3_promptd_job<7, C>(Xh, rows, T, svh, m, ch, Y, y_f32, buf); break; \ + case 8: tf_exl3_promptd_job<8, C>(Xh, rows, T, svh, m, ch, Y, y_f32, buf); break; \ + default: break; \ + } \ + } +TF_EXL3_PROMPTD(0) +TF_EXL3_PROMPTD(1) +TF_EXL3_PROMPTD(2) + // Decode rows for the tensor units: tf_exl3_rotate_h's values, each 128-block in fragment order (t4, r, 4 halves). extern "C" __global__ void __launch_bounds__(32) tf_exl3_rotate_f(const __nv_bfloat16* X, unsigned x_stride, const __half* suh, diff --git a/zig/kernels/cuda/flashnext_exl3/flashnext_exl3_prefill.cu b/zig/kernels/cuda/flashnext_exl3/flashnext_exl3_prefill.cu new file mode 100644 index 0000000000..408710c51b --- /dev/null +++ b/zig/kernels/cuda/flashnext_exl3/flashnext_exl3_prefill.cu @@ -0,0 +1,476 @@ +// Flash Next EXL3 prefill experts: upstream's transposed 64-member prompt tiles (exl3_prompt.cuh, after #976 and + +#include +#include +#include + +#include "../exl3_decode.cuh" +#include "../exl3_prompt.cuh" + +namespace tf_exl3pf { +using namespace tf_exl3; +using tf_exl3_prompt::cp_async16; +using tf_exl3_prompt::cp_async_commit; +using tf_exl3_prompt::cp_async_wait; +using tf_exl3_prompt::ldmatrix_x4; +using tf_exl3_prompt::tspill; +using tf_exl3_prompt::ttile; +using tf_exl3_prompt::TPM; +using tf_exl3_prompt::TNW; +using tf_exl3_prompt::TMT; +using tf_exl3_prompt::TNB; +using tf_exl3_prompt::TT; +using tf_exl3_prompt::TPH; +using tf_exl3_prompt::TBK; +using tf_exl3_prompt::TKT; +using tf_exl3_prompt::TSTAGES; +using tf_exl3_prompt::TXS; +using tf_exl3_prompt::GU_ES; +using tf_exl3_prompt::DN_ES; + +constexpr int KMAX2 = 10; // the widest tile a stage holds (5 bits): Flash Next's routed 3-bit and shared 5-bit experts +// Members a pass: decode's sums (exl3_experts.cu grouped_body) keep three accumulator sets live; 32 fit the registers. +constexpr int PM = 32; +constexpr int PNB = PM / 8; +constexpr int CHAIN_WARPS = 4; +template +constexpr int stage_words() { return 16 * TKT * tile_words(); } +template +constexpr int stage_halves() { return XR * TXS + 2 * stage_words(); } +constexpr int smem_bytes() { + const int gu = TSTAGES * stage_halves<2 * PM>() * 2; + const int ep = 2 * PM * GU_ES * 4, dn = PM * DN_ES * 4; + return gu > ep ? (gu > dn ? gu : dn) : (ep > dn ? ep : dn); +} + +// K splits of exl3_experts.zig's plan: gate/up 4 (else 2, else 1), down 1. +__device__ __forceinline__ int k_splits(int K, bool gateup) { + if (!gateup) return 1; + return K % 256 == 0 ? 4 : (K % 128 == 0 ? 2 : 1); +} + +template +__device__ __forceinline__ void decode_tile_k(const uint32_t* tile, int lane, uint32_t (&b0)[2], uint32_t (&b1)[2]) { + uint32_t w[lane_words()]; + int word, offset; + lane_start(lane, word, offset); +#pragma unroll + for (int i = 0; i < lane_words(); ++i) w[i] = tile[(word + i) % tile_words()]; + decode_lane(w, lane, b0, b1); +} + +// res[i][nb] = warp tile i x members nb * 8 .. + 7 over K, summed as grouped_body does; member i reads X[xfirst + i]. +template +__device__ __forceinline__ void pmain(const half* __restrict__ X0, const half* __restrict__ X1, const uint32_t* __restrict__ T0, + const uint32_t* __restrict__ T1, int xfirst, int cnt, int K, int NTILES, int cb, int chain_chunks, half* stage, + float (&res)[TMT][PNB][4]) { + constexpr int SH = stage_halves(); + constexpr int TW = tile_words(); + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int chunks = K / TBK; + const int nblocks = (cnt + 15) >> 4; + float acc[TMT][PNB][4], part[TMT][PNB][4]; +#pragma unroll + for (int i = 0; i < TMT; ++i) +#pragma unroll + for (int j = 0; j < PNB; ++j) +#pragma unroll + for (int c = 0; c < 4; ++c) acc[i][j][c] = part[i][j][c] = res[i][j][c] = 0.f; + auto load = [&](int chunk) { + half* dst = stage + (size_t)(chunk % TSTAGES) * SH; + const int k0 = chunk * TBK; + for (int idx = threadIdx.x; idx < XR * (TBK / 8); idx += TT) { + const int m = idx / (PM * (TBK / 8)), rem = idx % (PM * (TBK / 8)), i = rem / (TBK / 8), j = rem % (TBK / 8); + if (i >= nblocks * 16) continue; + const bool ok = i < cnt; + const half* src = (m ? X1 : X0) + (size_t)(xfirst + (ok ? i : 0)) * K + k0 + j * 8; + cp_async16(dst + ((size_t)m * PM + i) * TXS + j * 8, src, ok); + } + uint32_t* wdst = reinterpret_cast(dst + (size_t)XR * TXS); + for (int idx = threadIdx.x; idx < 16 * TKT * K2; idx += TT) { // a tile is K2 pieces of 16 bytes + const int q = idx / (TKT * K2), kk = (idx / K2) % TKT, piece = idx % K2; + int wm, nt; + ttile(q, cb, wm, nt); + const uint32_t* src = (wm ? T1 : T0) + (((size_t)chunk * TKT + kk) * NTILES + nt) * TW + piece * 4; + cp_async16(wdst + ((size_t)q * TKT + kk) * TW + piece * 4, src, true); + } + }; + int mat, nt_unused; + ttile(warp * TMT, cb, mat, nt_unused); +#pragma unroll + for (int c = 0; c < TSTAGES - 1; ++c) { + if (c < chunks) load(c); + cp_async_commit(); + } + int left = chain_chunks, w = 0; + for (int c = 0; c < chunks; ++c) { + cp_async_wait(); + __syncthreads(); + if (c + TSTAGES - 1 < chunks) load(c + TSTAGES - 1); + cp_async_commit(); + const half* base = stage + (size_t)(c % TSTAGES) * SH; + const half* xs = base + (size_t)(XR == PM ? 0 : mat) * PM * TXS; + const uint32_t* ws = reinterpret_cast(base + (size_t)XR * TXS); +#pragma unroll + for (int kk = 0; kk < TKT; ++kk) { + uint32_t a[TMT][4]; +#pragma unroll + for (int i = 0; i < TMT; ++i) { + uint32_t b0[2], b1[2]; + decode_tile_k(ws + ((size_t)(warp * TMT + i) * TKT + kk) * TW, lane, b0, b1); + a[i][0] = b0[0]; + a[i][1] = b1[0]; + a[i][2] = b0[1]; + a[i][3] = b1[1]; + } +#pragma unroll + for (int jb = 0; jb < PM / 16; ++jb) { + if (jb < nblocks) { + uint32_t r[4]; + const int row = jb * 16 + (lane & 7) + ((lane >> 4) & 1) * 8; + ldmatrix_x4(r, xs + (size_t)row * TXS + kk * 16 + ((lane >> 3) & 1) * 8); + const uint32_t bl[2] = {r[0], r[1]}, bh[2] = {r[2], r[3]}; +#pragma unroll + for (int i = 0; i < TMT; ++i) { + mma16816(acc[i][2 * jb], a[i], bl); + mma16816(acc[i][2 * jb + 1], a[i], bh); + } + } + } + } + if (--left == 0) { // a chain ends: into its split's warp sum; a split's last into the total (0 + each, as the epilogue) + left = chain_chunks; +#pragma unroll + for (int i = 0; i < TMT; ++i) +#pragma unroll + for (int j = 0; j < PNB; ++j) +#pragma unroll + for (int c4 = 0; c4 < 4; ++c4) { + part[i][j][c4] = w == 0 ? acc[i][j][c4] : __fadd_rn(part[i][j][c4], acc[i][j][c4]); + acc[i][j][c4] = 0.f; + } + if (++w == CHAIN_WARPS) { + w = 0; +#pragma unroll + for (int i = 0; i < TMT; ++i) +#pragma unroll + for (int j = 0; j < PNB; ++j) +#pragma unroll + for (int c4 = 0; c4 < 4; ++c4) res[i][j][c4] = __fadd_rn(res[i][j][c4], part[i][j][c4]); + } + } + } + cp_async_wait<0>(); + __syncthreads(); // the stages become the epilogue's sums +} + +// res -> E[member][col] fp32, row stride ES; col0: the warp's first column. +__device__ __forceinline__ void spill(float* E, int ES, int col0, const float (&res)[TMT][PNB][4]) { + const int lane = threadIdx.x & 31, g = lane >> 2, t = lane & 3; +#pragma unroll + for (int i = 0; i < TMT; ++i) +#pragma unroll + for (int q = 0; q < PNB; ++q) { + const int tok = q * 8 + 2 * t, col = col0 + i * 16 + g; + E[(size_t)tok * ES + col] = res[i][q][0]; + E[(size_t)(tok + 1) * ES + col] = res[i][q][1]; + E[(size_t)tok * ES + col + 8] = res[i][q][2]; + E[(size_t)(tok + 1) * ES + col + 8] = res[i][q][3]; + } +} + +__device__ __forceinline__ float upTo(float x, float l) { return x > l ? l : x; } +__device__ __forceinline__ float within(float x, float l) { return x > l ? l : (x < -l ? -l : x); } +__device__ __forceinline__ float bf16r(float x) { return __bfloat162float(__float2bfloat16_rn(x)); } + +template +__device__ __forceinline__ void gateup(const half* __restrict__ Xg, const half* __restrict__ Xu, const uint32_t* Tg, const uint32_t* Tu, + int e, int first, int cnt, const half* __restrict__ svh_g, const half* __restrict__ svh_u, const half* __restrict__ suh_d, + half* __restrict__ xd, int K, int N, float limit, int act_mode, unsigned char* smem_raw) { + const int cb = blockIdx.x; + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int chain_chunks = (K / 16 / k_splits(K, true) / CHAIN_WARPS) / TKT; + int mat, nt; + ttile<2>(warp * TMT, cb, mat, nt); + const int col0 = (nt - cb * 8) * 16; + const int n = cb * 128 + 4 * lane; + float* Eg = reinterpret_cast(smem_raw); + for (int m0 = 0; m0 < cnt; m0 += PM) { + const int pc = cnt - m0 < PM ? cnt - m0 : PM; + if (m0) __syncthreads(); // the last pass's sums are read before the stages are refilled + float res[TMT][PNB][4]; + pmain(Xg, Xu, Tg, Tu, first + m0, pc, K, N >> 4, cb, chain_chunks, reinterpret_cast(smem_raw), res); + spill(Eg + (size_t)mat * PM * GU_ES, GU_ES, col0, res); + __syncthreads(); + for (int i = warp; i < pc; i += TNW) { + float gv[4], uv[4]; +#pragma unroll + for (int j = 0; j < 4; ++j) { + gv[j] = Eg[(size_t)i * GU_ES + 4 * lane + j]; + uv[j] = Eg[(size_t)(PM + i) * GU_ES + 4 * lane + j]; + } + fwht128(gv, lane); + fwht128(uv, lane); + float v[4]; +#pragma unroll + for (int j = 0; j < 4; ++j) { + float act; + if (act_mode == 0) { + float gg = bf16r(gv[j] * HAD_SCALE * __half2float(svh_g[(size_t)e * N + n + j])); + float uu = bf16r(uv[j] * HAD_SCALE * __half2float(svh_u[(size_t)e * N + n + j])); + if (limit > 0) { gg = upTo(gg, limit); uu = within(uu, limit); } + act = bf16r(bf16r(gg / (1.f + expf(-gg))) * uu); + } else { // runSlots' tf_exl3_gateup_epilogue, act_mode 1, operation for operation + float gg = gv[j] * HAD_SCALE * __half2float(svh_g[(size_t)e * N + n + j]); + float uu = uv[j] * HAD_SCALE * __half2float(svh_u[(size_t)e * N + n + j]); + if (limit > 0) { gg = upTo(gg, limit); uu = within(uu, limit); } + act = gg / (1.f + expf(-gg)) * uu; + } + v[j] = act * __half2float(suh_d[(size_t)e * N + n + j]); + } + fwht128(v, lane); + half* o = xd + (size_t)(first + m0 + i) * N + n; +#pragma unroll + for (int j = 0; j < 4; ++j) o[j] = __float2half_rn(v[j] * HAD_SCALE); + } + } +} + +template +__device__ __forceinline__ void down(const half* __restrict__ Xd, const uint32_t* Td, int e, int first, int cnt, const int* rows_sh, + const half* __restrict__ svh_d, float* __restrict__ y, int K, int D, unsigned char* smem_raw) { + const int cb = blockIdx.x; + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int chain_chunks = (K / 16 / k_splits(K, false) / CHAIN_WARPS) / TKT; + float* Ed = reinterpret_cast(smem_raw); + for (int m0 = 0; m0 < cnt; m0 += PM) { + const int pc = cnt - m0 < PM ? cnt - m0 : PM; + if (m0) __syncthreads(); + float res[TMT][PNB][4]; + pmain(Xd, Xd, Td, Td, first + m0, pc, K, D >> 4, cb, chain_chunks, reinterpret_cast(smem_raw), res); + spill(Ed, DN_ES, warp * TMT * 16, res); + __syncthreads(); + for (int i = warp; i < pc; i += TNW) { + const int p = rows_sh[m0 + i]; +#pragma unroll + for (int blk = 0; blk < 2; ++blk) { + const int n = cb * 256 + blk * 128 + 4 * lane; + float v[4]; +#pragma unroll + for (int j = 0; j < 4; ++j) v[j] = Ed[(size_t)i * DN_ES + blk * 128 + 4 * lane + j]; + fwht128(v, lane); + float* o = y + (size_t)p * D + n; +#pragma unroll + for (int j = 0; j < 4; ++j) o[j] = v[j] * HAD_SCALE * __half2float(svh_d[(size_t)e * D + n + j]); + } + } + } +} +} // namespace tf_exl3pf + +using namespace tf_exl3pf; + +// Program (compact member c, 128-block of D, matrix): xg/xu[c] = fp16((x[row] * suh[e]) @ H) for member c's pair. +extern "C" __global__ void tf_exl3pf_rotate(const __nv_bfloat16* __restrict__ x, int x_stride, const int* __restrict__ pick, + const int* __restrict__ members, int slots, const half* __restrict__ suh_g, const half* __restrict__ suh_u, + half* __restrict__ xg, half* __restrict__ xu, int D) { + const int c = blockIdx.x, blk = blockIdx.y, mat = blockIdx.z; + const int p = members[c]; + const int e = pick[p]; + const int row = p / slots; + const int lane = threadIdx.x; + const half* suh = (mat ? suh_u : suh_g) + (size_t)e * D + blk * 128 + 4 * lane; + const __nv_bfloat16* xr = x + (size_t)row * x_stride + blk * 128 + 4 * lane; + float v[4]; +#pragma unroll + for (int j = 0; j < 4; ++j) v[j] = __bfloat162float(xr[j]) * __half2float(suh[j]); + fwht128(v, lane); + half* o = (mat ? xu : xg) + (size_t)c * D + blk * 128 + 4 * lane; +#pragma unroll + for (int j = 0; j < 4; ++j) o[j] = __float2half_rn(v[j] * HAD_SCALE); +} + +#define TF_EXL3PF_K2(X) \ + X(4) \ + X(6) \ + X(8) \ + X(10) + +// Block (128 width columns of gate and up, item): SwiGLU of the item's members into compact xd rows. +extern "C" __global__ void __launch_bounds__(256, 2) tf_exl3pf_gu(const half* __restrict__ Xg, const half* __restrict__ Xu, + const uint64_t* __restrict__ Tg, const uint64_t* __restrict__ Tu, const int* __restrict__ k2g, const int* __restrict__ items, + const half* __restrict__ svh_g, const half* __restrict__ svh_u, const half* __restrict__ suh_d, half* __restrict__ xd, int K, + int N, float limit, int act_mode) { + __shared__ __align__(16) unsigned char smem_raw[smem_bytes()]; + const int item = blockIdx.y; + const int e = items[3 * item], first = items[3 * item + 1], cnt = items[3 * item + 2]; + const uint32_t* G = reinterpret_cast(Tg[e]); + const uint32_t* U = reinterpret_cast(Tu[e]); + switch (k2g[e]) { +#define TF_CASE(K2_) \ + case K2_: gateup(Xg, Xu, G, U, e, first, cnt, svh_g, svh_u, suh_d, xd, K, N, limit, act_mode, smem_raw); break; + TF_EXL3PF_K2(TF_CASE) +#undef TF_CASE + default: __trap(); + } +} + +// Block (256 model-width columns, item): y[pair] = the down epilogue of the item's members. +extern "C" __global__ void __launch_bounds__(256, 2) tf_exl3pf_down(const half* __restrict__ Xd, const uint64_t* __restrict__ Td, + const int* __restrict__ k2d, const int* __restrict__ items, const int* __restrict__ members, const half* __restrict__ svh_d, + float* __restrict__ y, int K, int D) { + __shared__ __align__(16) unsigned char smem_raw[smem_bytes()]; + __shared__ int rows_sh[TPM]; + const int item = blockIdx.y; + const int e = items[3 * item], first = items[3 * item + 1], cnt = items[3 * item + 2]; + if (threadIdx.x < TPM) rows_sh[threadIdx.x] = (int)threadIdx.x < cnt ? members[first + threadIdx.x] : -1; + __syncthreads(); + const uint32_t* T = reinterpret_cast(Td[e]); + switch (k2d[e]) { +#define TF_CASE(K2_) \ + case K2_: down(Xd, T, e, first, cnt, rows_sh, svh_d, y, K, D, smem_raw); break; + TF_EXL3PF_K2(TF_CASE) +#undef TF_CASE + default: __trap(); + } +} + +// The canonical expert reduction at any degree: 64 leaves, a balanced tree; a rank sums its maximal complete subtrees. +__device__ __forceinline__ int pf_blocks(int lo, int hi, int* st, int* sz) { + int n = 0; + for (int at = lo; at < hi;) { + int size = 1; + while (at % (size * 2) == 0 && at + size * 2 <= hi) size *= 2; + st[n] = at; + sz[n] = size; + ++n; + at += size; + } + return n; +} + +// Block partials of one rank into out[sub-block][rank * M + j][row][D]: a sub-block's slice is one all-gather's region. +extern "C" __global__ void __launch_bounds__(256) tf_exl3pf_tp_combine(const float* __restrict__ y, const float* __restrict__ wts, + const int* __restrict__ pick, float* __restrict__ out, int slots, int D, int n, int sub, int world, int rank, int M) { + __shared__ float w_sh[16]; + __shared__ int nl_sh; + __shared__ int leaf_sh[16]; + __shared__ unsigned mask_sh[16]; + const int r = blockIdx.x; + const int llo = 64 * rank / world, lhi = 64 * (rank + 1) / world; + if ((int)threadIdx.x < slots) w_sh[threadIdx.x] = wts[r * slots + threadIdx.x]; + if (threadIdx.x == 0) { + int nl = 0; + for (int k = 0; k < slots; ++k) { + const int p = pick[r * slots + k]; + if (p < 0) continue; + const int leaf = p >= 512 ? 63 : p / 8; + if (leaf < llo || leaf >= lhi) continue; + int at = 0; + while (at < nl && leaf_sh[at] < leaf) ++at; + if (at < nl && leaf_sh[at] == leaf) { mask_sh[at] |= 1u << k; continue; } + for (int i = nl; i > at; --i) { leaf_sh[i] = leaf_sh[i - 1]; mask_sh[i] = mask_sh[i - 1]; } + leaf_sh[at] = leaf; + mask_sh[at] = 1u << k; + ++nl; + } + nl_sh = nl; + } + __syncthreads(); + const int nl = nl_sh; + int bst[16], bsz[16]; + const int nb = pf_blocks(llo, lhi, bst, bsz); + const int b = r / sub, rr = r - b * sub; + const int mb = min(sub, n - b * sub); + const float* yr = y + (size_t)r * slots * D; + for (int d = threadIdx.x; d < D; d += blockDim.x) { + float val[16]; + int st[16]; + for (int i = 0; i < nl; ++i) { + float acc = 0.0f; + for (unsigned m = mask_sh[i]; m; m &= m - 1) { + const int k = __ffs(m) - 1; + acc = fmaf(yr[(size_t)k * D + d], w_sh[k], acc); + } + val[i] = acc; + st[i] = leaf_sh[i]; + } + for (int j = 0; j < nb; ++j) { + float v[16]; + int s[16]; + int cnt = 0; + for (int i = 0; i < nl; ++i) if (st[i] >= bst[j] && st[i] < bst[j] + bsz[j]) { v[cnt] = val[i]; s[cnt] = st[i]; ++cnt; } + for (int size = 1; size < bsz[j]; size *= 2) { + int m = 0; + for (int i = 0; i < cnt;) { + const int start = s[i] & ~(2 * size - 1); + if (s[i] % (2 * size) == 0 && i + 1 < cnt && s[i + 1] == s[i] + size) { + v[m] = __fadd_rn(v[i], v[i + 1]); + i += 2; + } else { + v[m] = v[i]; + ++i; + } + s[m] = start; + ++m; + } + cnt = m; + } + out[(((size_t)b * sub * world * M + (size_t)(rank * M + j) * mb + rr) * D) + d] = cnt ? v[0] : 0.0f; + } + } +} + +// The tree finished from the gathered partials: equal adjacent aligned blocks merge by fadd_rn until one is left. +extern "C" __global__ void __launch_bounds__(256) tf_exl3pf_tp_finish(const float* __restrict__ gath, float* __restrict__ sum, + int D, int m, int world, int M) { + __shared__ int bs_sh[96], bz_sh[96], bq_sh[96]; + __shared__ int nb_sh; + if (threadIdx.x == 0) { + int n = 0; + for (int q = 0; q < world; ++q) { + int st[16], sz[16]; + const int c = pf_blocks(64 * q / world, 64 * (q + 1) / world, st, sz); + for (int j = 0; j < c; ++j) { bs_sh[n] = st[j]; bz_sh[n] = sz[j]; bq_sh[n] = q * M + j; ++n; } + } + nb_sh = n; + } + __syncthreads(); + const int rr = blockIdx.x; + const int nb = nb_sh; + for (int d = threadIdx.x; d < D; d += blockDim.x) { + float v[96]; + int s[96], z[96]; + for (int i = 0; i < nb; ++i) { v[i] = gath[((size_t)bq_sh[i] * m + rr) * D + d]; s[i] = bs_sh[i]; z[i] = bz_sh[i]; } + int cnt = nb; + while (cnt > 1) { + int k = 0; + for (int i = 0; i < cnt;) { + if (i + 1 < cnt && z[i] == z[i + 1] && s[i] % (z[i] * 2) == 0) { + v[k] = __fadd_rn(v[i], v[i + 1]); + s[k] = s[i]; + z[k] = z[i] * 2; + i += 2; + } else { + v[k] = v[i]; + s[k] = s[i]; + z[k] = z[i]; + ++i; + } + ++k; + } + cnt = k; + } + sum[(size_t)rr * D + d] = v[0]; + } +} + +// The ranks' summed branch back where the write-back reads it: y[r, 0, :] = sum[r], weights {1, 0, ..., 0}. The other +extern "C" __global__ void __launch_bounds__(256) tf_exl3pf_tp_place(const float* __restrict__ sum, float* __restrict__ y, + float* __restrict__ wts, int slots, int D) { + const int r = blockIdx.x; + float* out = y + (size_t)r * slots * D; + const float* in = sum + (size_t)r * D; + for (int d = threadIdx.x; d < D; d += blockDim.x) out[d] = in[d]; + if ((int)threadIdx.x < slots) wts[r * slots + threadIdx.x] = threadIdx.x == 0 ? 1.0f : 0.0f; +} diff --git a/zig/src/core/contract/flashnext_exl3_pending.zig b/zig/src/core/contract/flashnext_exl3_pending.zig index 4178c20231..1d57be4c76 100644 --- a/zig/src/core/contract/flashnext_exl3_pending.zig +++ b/zig/src/core/contract/flashnext_exl3_pending.zig @@ -1,4 +1,4 @@ -//! Flash Next EXL3 launches (the fp16 row split, the expert-parallel pick mask); nvcc 13.4 measured, older pending. +//! Flash Next EXL3 launches (pick mask, prefill tiles in decode's bytes); nvcc 13.4 measured, older pending. const std = @import("std"); const limits = @import("limits.zig"); const registry = @import("registry.zig"); @@ -9,6 +9,15 @@ pub const records = [_]registry.CudaRecord{ .{ .launch = .{ .name = "flashnext_batch fn_exl3_leaf_tree", .symbol = "fn_exl3_leaf_tree", .image = "flashnext_batch", .threads = 256 }, .usage = registry.none }, .{ .launch = .{ .name = "flashnext_batch fn_exl3_leaf_finish", .symbol = "fn_exl3_leaf_finish", .image = "flashnext_batch", .threads = 256 }, .usage = registry.none }, .{ .launch = .{ .name = "flashnext_exl3_tp fn_exl3_tp_mask", .symbol = "fn_exl3_tp_mask", .image = "flashnext_exl3_tp", .threads = 256 }, .usage = registry.none }, + .{ .launch = .{ .name = "flashnext_exl3_prefill tf_exl3pf_rotate", .symbol = "tf_exl3pf_rotate", .image = "flashnext_exl3_prefill", .threads = 32 }, .usage = registry.none }, + .{ .launch = .{ .name = "flashnext_exl3_prefill tf_exl3pf_gu", .symbol = "tf_exl3pf_gu", .image = "flashnext_exl3_prefill", .threads = 256 }, .usage = registry.none }, + .{ .launch = .{ .name = "flashnext_exl3_prefill tf_exl3pf_down", .symbol = "tf_exl3pf_down", .image = "flashnext_exl3_prefill", .threads = 256 }, .usage = registry.none }, + .{ .launch = .{ .name = "flashnext_exl3_prefill tf_exl3pf_tp_combine", .symbol = "tf_exl3pf_tp_combine", .image = "flashnext_exl3_prefill", .threads = 256 }, .usage = registry.none }, + .{ .launch = .{ .name = "flashnext_exl3_prefill tf_exl3pf_tp_finish", .symbol = "tf_exl3pf_tp_finish", .image = "flashnext_exl3_prefill", .threads = 256 }, .usage = registry.none }, + .{ .launch = .{ .name = "flashnext_exl3_prefill tf_exl3pf_tp_place", .symbol = "tf_exl3pf_tp_place", .image = "flashnext_exl3_prefill", .threads = 256 }, .usage = registry.none }, + .{ .launch = .{ .name = "exl3_candidate tf_exl3_promptd_0", .symbol = "tf_exl3_promptd_0", .image = "exl3_candidate", .threads = 256 }, .usage = registry.none }, + .{ .launch = .{ .name = "exl3_candidate tf_exl3_promptd_1", .symbol = "tf_exl3_promptd_1", .image = "exl3_candidate", .threads = 256 }, .usage = registry.none }, + .{ .launch = .{ .name = "exl3_candidate tf_exl3_promptd_2", .symbol = "tf_exl3_promptd_2", .image = "exl3_candidate", .threads = 256 }, .usage = registry.none }, }; /// nvcc 13.4's measured registers, static shared bytes and launch bounds for sm_80, sm_86, sm_89, sm_120, sm_121. @@ -17,6 +26,15 @@ pub const cuda_134 = [_]registry.Override{ .{ .name = "flashnext_batch fn_exl3_leaf_tree", .usage = .{ .{ .regs = 32, .static = 0, .bound = 0 }, .{ .regs = 34, .static = 0, .bound = 0 }, .{ .regs = 34, .static = 0, .bound = 0 }, .{ .regs = 24, .static = 0, .bound = 0 }, .{ .regs = 24, .static = 0, .bound = 0 } } }, .{ .name = "flashnext_batch fn_exl3_leaf_finish", .usage = .{ .{ .regs = 23, .static = 0, .bound = 0 }, .{ .regs = 22, .static = 0, .bound = 0 }, .{ .regs = 22, .static = 0, .bound = 0 }, .{ .regs = 19, .static = 0, .bound = 0 }, .{ .regs = 19, .static = 0, .bound = 0 } } }, .{ .name = "flashnext_exl3_tp fn_exl3_tp_mask", .usage = .{ .{ .regs = 10, .static = 0, .bound = 0 }, .{ .regs = 10, .static = 0, .bound = 0 }, .{ .regs = 10, .static = 0, .bound = 0 }, .{ .regs = 10, .static = 0, .bound = 0 }, .{ .regs = 10, .static = 0, .bound = 0 } } }, + .{ .name = "flashnext_exl3_prefill tf_exl3pf_rotate", .usage = .{ .{ .regs = 22, .static = 0, .bound = 0 }, .{ .regs = 22, .static = 0, .bound = 0 }, .{ .regs = 22, .static = 0, .bound = 0 }, .{ .regs = 19, .static = 0, .bound = 0 }, .{ .regs = 19, .static = 0, .bound = 0 } } }, + .{ .name = "flashnext_exl3_prefill tf_exl3pf_gu", .usage = .{ .{ .regs = 128, .static = 33792, .bound = 256 }, .{ .regs = 128, .static = 33792, .bound = 256 }, .{ .regs = 128, .static = 33792, .bound = 256 }, .{ .regs = 128, .static = 33792, .bound = 256 }, .{ .regs = 128, .static = 33792, .bound = 256 } } }, + .{ .name = "flashnext_exl3_prefill tf_exl3pf_down", .usage = .{ .{ .regs = 128, .static = 34048, .bound = 256 }, .{ .regs = 128, .static = 34048, .bound = 256 }, .{ .regs = 128, .static = 34048, .bound = 256 }, .{ .regs = 128, .static = 34048, .bound = 256 }, .{ .regs = 128, .static = 34048, .bound = 256 } } }, + .{ .name = "flashnext_exl3_prefill tf_exl3pf_tp_combine", .usage = .{ .{ .regs = 32, .static = 196, .bound = 256 }, .{ .regs = 40, .static = 196, .bound = 256 }, .{ .regs = 40, .static = 196, .bound = 256 }, .{ .regs = 29, .static = 196, .bound = 256 }, .{ .regs = 29, .static = 196, .bound = 256 } } }, + .{ .name = "flashnext_exl3_prefill tf_exl3pf_tp_finish", .usage = .{ .{ .regs = 32, .static = 1156, .bound = 256 }, .{ .regs = 40, .static = 1156, .bound = 256 }, .{ .regs = 40, .static = 1156, .bound = 256 }, .{ .regs = 38, .static = 1156, .bound = 256 }, .{ .regs = 38, .static = 1156, .bound = 256 } } }, + .{ .name = "flashnext_exl3_prefill tf_exl3pf_tp_place", .usage = .{ .{ .regs = 18, .static = 0, .bound = 256 }, .{ .regs = 18, .static = 0, .bound = 256 }, .{ .regs = 18, .static = 0, .bound = 256 }, .{ .regs = 24, .static = 0, .bound = 256 }, .{ .regs = 24, .static = 0, .bound = 256 } } }, + .{ .name = "exl3_candidate tf_exl3_promptd_0", .usage = .{ .{ .regs = 202, .static = 32768, .bound = 256 }, .{ .regs = 204, .static = 32768, .bound = 256 }, .{ .regs = 204, .static = 32768, .bound = 256 }, .{ .regs = 180, .static = 32768, .bound = 256 }, .{ .regs = 180, .static = 32768, .bound = 256 } } }, + .{ .name = "exl3_candidate tf_exl3_promptd_1", .usage = .{ .{ .regs = 202, .static = 32768, .bound = 256 }, .{ .regs = 204, .static = 32768, .bound = 256 }, .{ .regs = 204, .static = 32768, .bound = 256 }, .{ .regs = 180, .static = 32768, .bound = 256 }, .{ .regs = 180, .static = 32768, .bound = 256 } } }, + .{ .name = "exl3_candidate tf_exl3_promptd_2", .usage = .{ .{ .regs = 202, .static = 32768, .bound = 256 }, .{ .regs = 204, .static = 32768, .bound = 256 }, .{ .regs = 204, .static = 32768, .bound = 256 }, .{ .regs = 179, .static = 32768, .bound = 256 }, .{ .regs = 179, .static = 32768, .bound = 256 } } }, }; test "Flash Next EXL3 launches: pending below nvcc 13.4, measured 13.4 resources fit every SM launch limit" { diff --git a/zig/src/cuda/exl3.zig b/zig/src/cuda/exl3.zig index 54176e9f64..5acdd75ce0 100644 --- a/zig/src/cuda/exl3.zig +++ b/zig/src/cuda/exl3.zig @@ -243,19 +243,34 @@ pub const Dense = struct { return m; } - fn tiles(de: *Dense, s: Stream, w: Weight, x: u64, xd: DType, rows: usize, y: u64, yd: DType, prompt_path: bool) !void { + /// Prompt rows on the tensor units with `matmul`'s bytes: the decode kernel's rotation, chains and epilogue. + pub fn promptExact(de: *Dense, s: Stream, w: Weight, x: u64, rows: usize, y: u64, yd: DType) !void { const m = try candidateWeight(w); - if (xd != .bf16 or yd == .f16) return error.CandidateFormat; - if (de.candidate_k == null or de.candidate_k.?.k_max < w.k) { + if (yd == .f16) return error.CandidateFormat; + const p = plan(w.k, w.n, 192, 8); + if (variantIndex(w.halves, w.codebook, p.wk) == null) return error.UnsupportedExl3Width; + try de.candidates(w.k); + return de.candidate_k.?.promptExact(s, de.lin.rot_in, x, m, y, rows, yd == .f32, p.sk, p.wk); + } + + /// Loads the candidate module's kernels (fp16 rotated rows for inputs up to `k_in`) if not loaded wide enough. + pub fn candidates(de: *Dense, k_in: usize) !void { + if (de.candidate_k == null or de.candidate_k.?.k_max < k_in) { if (!kernels.available) return error.BuiltWithoutKernels; var mod = try Module.load(de.d, kernels.exl3_candidate); errdefer mod.unload(); - const replacement = try candidate.Kernels.resolve(de.d, mod, w.k); + const replacement = try candidate.Kernels.resolve(de.d, mod, k_in); if (de.candidate_k) |*k| k.deinit(); if (de.candidate_mod) |*old| old.unload(); de.candidate_mod = mod; de.candidate_k = replacement; } + } + + fn tiles(de: *Dense, s: Stream, w: Weight, x: u64, xd: DType, rows: usize, y: u64, yd: DType, prompt_path: bool) !void { + const m = try candidateWeight(w); + if (xd != .bf16 or yd == .f16) return error.CandidateFormat; + try de.candidates(w.k); const k = &de.candidate_k.?; if (prompt_path) return k.prompt(s, x, w.k, m, y, rows, yd == .f32); return k.rows(s, x, w.k, m, y, rows, yd == .f32); diff --git a/zig/src/cuda/exl3_candidate.zig b/zig/src/cuda/exl3_candidate.zig index c894f33096..31550ff7f4 100644 --- a/zig/src/cuda/exl3_candidate.zig +++ b/zig/src/cuda/exl3_candidate.zig @@ -40,6 +40,8 @@ pub const Matrix = struct { /// exl3.cu's TfExl3Rows and TfExl3Matrix, passed by value. pub const RowsArgs = extern struct { rows: u32, x_stride: u32, y_stride: u32, y_f32: u32 }; pub const MatrixArgs = extern struct { k: u32, n: u32, bits: u32, tiles_n: u32, split: u32 }; +/// exl3_candidate.cu's TfExl3Chains: where exl3_linear.cu cuts the K sum (k tiles a chain, chains a split). +pub const ChainsArgs = extern struct { per: u32, wk: u32 }; /// Decode rows a launch; longer decode runs go window by window. pub const window = 16; @@ -63,6 +65,7 @@ pub fn split(k: usize, n: usize) usize { pub const Kernels = struct { decode_fns: [3]Function, // 1-16 rows on the tensor units, by codebook prompt_fns: [3]Function, // 64 rows a block on the tensor units, by codebook + promptd_fns: [3]Function, // the same with decode's bytes (exl3_linear.cu's sums), by codebook rotate_f_fn: Function, // decode rows, each 128-block in fragment order rotate_h_fn: Function, // prompt rows, plain order codebook_fn: Function, @@ -73,11 +76,12 @@ pub const Kernels = struct { k_max: usize, pub fn resolve(d: *const Driver, mod: module.Module, k_max: usize) !Kernels { - var k: Kernels = .{ .decode_fns = undefined, .prompt_fns = undefined, .rotate_f_fn = try mod.function("tf_exl3_rotate_f"), .rotate_h_fn = try mod.function("tf_exl3_rotate_h"), .codebook_fn = try mod.function("tf_exl3_codebook"), .unpack_fn = try mod.function("tf_exl3_unpack"), .rotated = undefined, .parts = undefined, .tickets = undefined, .k_max = k_max }; - for (&k.decode_fns, &k.prompt_fns, 0..) |*f, *p, c| { + var k: Kernels = .{ .decode_fns = undefined, .prompt_fns = undefined, .promptd_fns = undefined, .rotate_f_fn = try mod.function("tf_exl3_rotate_f"), .rotate_h_fn = try mod.function("tf_exl3_rotate_h"), .codebook_fn = try mod.function("tf_exl3_codebook"), .unpack_fn = try mod.function("tf_exl3_unpack"), .rotated = undefined, .parts = undefined, .tickets = undefined, .k_max = k_max }; + for (&k.decode_fns, &k.prompt_fns, &k.promptd_fns, 0..) |*f, *p, *e, c| { var name: [32]u8 = undefined; f.* = try mod.function(try std.fmt.bufPrintSentinel(&name, "tf_exl3_tc_{d}", .{c}, 0)); p.* = try mod.function(try std.fmt.bufPrintSentinel(&name, "tf_exl3_prompt_{d}", .{c}, 0)); + e.* = try mod.function(try std.fmt.bufPrintSentinel(&name, "tf_exl3_promptd_{d}", .{c}, 0)); } k.rotated = try memory.DeviceBuffer.alloc(d, prompt_chunk * k_max * 2); errdefer k.rotated.free(); @@ -137,6 +141,36 @@ pub const Kernels = struct { } } + /// Prompt rows with `Linear.matmul`'s bytes at any row count; plan_sk and plan_wk are exl3.plan(k, n, 192, 8)'s. + pub fn promptExact(k: *const Kernels, s: Stream, rot_in: Function, x: u64, m: Matrix, out: u64, count: usize, f32_out: bool, plan_sk: u32, plan_wk: u32) !void { + const elem: usize = if (f32_out) 4 else 2; + try k.fits(m, m.k); + if (plan_sk == 0 or plan_wk == 0 or (m.k / 16) % (plan_sk * plan_wk) != 0) return error.Exl3Shape; + const chains: ChainsArgs = .{ .per = @intCast(m.k / 16 / plan_sk / plan_wk), .wk = plan_wk }; + var at: usize = 0; + while (at < count) : (at += prompt_chunk) { + const r = @min(prompt_chunk, count - at); + var ra: launch_.Args = .{}; + ra.add(x + at * m.k * 2); + ra.add(m.suh); + ra.add(k.rotated.ptr); + ra.add(@as(i32, 1)); // bf16 rows + ra.add(@as(i32, @intCast(r))); + ra.add(@as(i32, @intCast(m.k))); + try launch_.launch(rot_in, .{ .grid = .{ .x = @intCast((m.k / 128 + 3) / 4), .y = @intCast(r), .z = 1 }, .block = .{ .x = 128, .y = 1, .z = 1 } }, s, &ra); + var a: launch_.Args = .{}; + a.add(k.rotated.ptr); + a.add(@as(u32, @intCast(r))); + a.add(m.trellis); + a.add(m.svh); + a.add(m.args()); + a.add(chains); + a.add(out + at * m.n * elem); + a.add(@as(u32, @intFromBool(f32_out))); + try launch_.launch(k.promptd_fns[@backingInt(m.codebook)], .{ .grid = .{ .x = @intCast(m.n / 128), .y = @intCast((r + 63) / 64), .z = 1 }, .block = .{ .x = 256 } }, s, &a); + } + } + fn fits(k: *const Kernels, m: Matrix, ldx: usize) !void { try m.check(); if (m.k > k.k_max) return error.Exl3InputsOverScratch; diff --git a/zig/src/cuda/kernels.zig b/zig/src/cuda/kernels.zig index 9e1402f1c0..888ec9d67d 100644 --- a/zig/src/cuda/kernels.zig +++ b/zig/src/cuda/kernels.zig @@ -20,6 +20,7 @@ pub const gdn: []const u8 = if (available) &Blob("fatbin_gdn").bytes else &.{}; pub const flashnext_gdn: []const u8 = if (available) &Blob("fatbin_flashnext_gdn").bytes else &.{}; pub const flashnext_exl3: []const u8 = if (available) &Blob("fatbin_flashnext_exl3").bytes else &.{}; pub const flashnext_exl3_tp: []const u8 = if (available) &Blob("fatbin_flashnext_exl3_tp").bytes else &.{}; +pub const flashnext_exl3_prefill: []const u8 = if (available) &Blob("fatbin_flashnext_exl3_prefill").bytes else &.{}; pub const flashnext_batch: []const u8 = if (available) &Blob("fatbin_flashnext_batch").bytes else &.{}; pub const flashnext_mtp: []const u8 = if (available) &Blob("fatbin_flashnext_mtp").bytes else &.{}; pub const probe: []const u8 = if (available) &Blob("fatbin_probe").bytes else &.{}; diff --git a/zig/src/families/flashnext/cuda_batch.zig b/zig/src/families/flashnext/cuda_batch.zig index c387d05d6b..adc5b674be 100644 --- a/zig/src/families/flashnext/cuda_batch.zig +++ b/zig/src/families/flashnext/cuda_batch.zig @@ -117,6 +117,7 @@ pub fn Pass(comptime Tri: type) type { prepared: bool = false, /// The pending write-back's branch sits in fp32 in moe_f32 (something added to it, or it is kept). pending_f32: bool = false, + pending_f32_weights: ?u64 = null, /// moe_f32 holds the EXL3 canonical grouped branch (moeAdds keeps it instead of re-summing the slots). canon_f32: bool = false, /// This layer's out projection leaves part's other slices alone; its writeback adds the zero slice once. @@ -182,6 +183,11 @@ pub fn Pass(comptime Tri: type) type { if (self.pending_f32) { // the fp32 branch rounded once, as MODE 2 rounds its sum (a slice of one) self.pending_f32 = false; + if (self.pending_f32_weights) |weights| { + self.pending_f32_weights = null; + try self.tri.hcWriteback(h, h, pss, dummy, "*bf16", inj, s.moe_f32 + r0 * dims * 4, "*fp32", weights + r0 * 4, "*fp32", rs, n, dims, streams, 2, 0, 1, 1); + return; + } try self.tri.hcWriteback(h, h, pss, s.moe_f32 + r0 * dims * 4, "*fp32", inj, dummy, "*bf16", dummy, "*bf16", rs, n, dims, streams, 4, 1, 1, 1); return; } @@ -503,6 +509,7 @@ pub fn Pass(comptime Tri: type) type { self.pos0 = pos0; self.pending = null; self.pending_f32 = false; + self.pending_f32_weights = null; self.canon_f32 = false; // the gather writes the scratch's pinned rows, so the last pass's uploads must have read them first if (!self.prepared) try s.ready(); diff --git a/zig/src/families/flashnext/cuda_collective.zig b/zig/src/families/flashnext/cuda_collective.zig index 77bd17b386..f7aec59250 100644 --- a/zig/src/families/flashnext/cuda_collective.zig +++ b/zig/src/families/flashnext/cuda_collective.zig @@ -16,6 +16,7 @@ pub const Collective = struct { comm: nccl.Comm = null, link: wire.Link, timeout_ms: u64, + rank: usize = 0, failed: bool = false, /// How long a wait polls before it sleeps between polls (0: sleeps from the first poll); an EXL3 engine sets it. spin_ns: u64 = 0, @@ -23,7 +24,7 @@ pub const Collective = struct { /// `link` has completed the source/settings handshake; it stays owned by the caller. pub fn init(link: wire.Link, rank: usize, timeout_ms: u64) !Collective { if (rank > 1 or timeout_ms == 0) return error.BadRank; - var c: Collective = .{ .lib = try nccl.Library.open(), .api = undefined, .link = link, .timeout_ms = timeout_ms }; + var c: Collective = .{ .lib = try nccl.Library.open(), .api = undefined, .link = link, .timeout_ms = timeout_ms, .rank = rank }; errdefer c.lib.close(); try c.lib.requireSupportedVersion(); const info = @typeInfo(Advanced).@"struct"; @@ -122,6 +123,23 @@ pub const Collective = struct { try c.wait(null, wire.deadline(c.timeout_ms)); } + /// FP32 values sent to the peer and the peer's received, enqueued together: no arithmetic, the bytes as they are. + pub fn exchange(c: *Collective, stream: cuda.Stream, send: u64, send_count: usize, recv: u64, recv_count: usize) !void { + if (c.failed or c.comm == null) return error.RankDisconnected; + if (send_count == 0 and recv_count == 0) return; + errdefer c.abort(); + try c.link.checkPeer(); + const peer: c_int = @intCast(1 - c.rank); + try c.submitted(c.lib.api.ncclGroupStart()); + var posted: nccl.Result = 0; + if (send_count != 0) posted = c.lib.api.ncclSend(send, send_count, .f32, peer, c.comm, stream.handle); + if (posted == 0 and recv_count != 0) posted = c.lib.api.ncclRecv(recv, recv_count, .f32, peer, c.comm, stream.handle); + const ended = c.lib.api.ncclGroupEnd(); + try c.submitted(posted); + try c.submitted(ended); + try c.wait(null, wire.deadline(c.timeout_ms)); + } + /// FP32 sum is enqueued on the compute stream; later kernels on that stream observe its result. pub fn sum(c: *Collective, stream: cuda.Stream, values: u64, count: usize) !void { if (c.failed or c.comm == null) return error.RankDisconnected; diff --git a/zig/src/families/flashnext/cuda_lanes.zig b/zig/src/families/flashnext/cuda_lanes.zig index 99a57f0c5f..8292d2b2dd 100644 --- a/zig/src/families/flashnext/cuda_lanes.zig +++ b/zig/src/families/flashnext/cuda_lanes.zig @@ -301,7 +301,7 @@ pub fn Namespace(comptime Tri: type) type { dropLane(self, l); self.lanes.removeByPtr(gop.key_ptr); } - if (drafting) l.draft = if (from == .pool) try self.kv.pool.?.draft(&l.seq, cap) else try mtp.Draft.init(self.gpa, self.engine.driver, self.stream, cap, self.head.?.ids.len, self.engine.chunk); + if (drafting) l.draft = if (from == .pool) try self.kv.pool.?.draft(&l.seq, cap) else try mtp.Draft.init(self.gpa, self.engine.driver, self.stream, cap, self.head.?.ids.len, @import("exl3/prefill/cuda_exl3_prefill.zig").draftRows(self.engine)); // EXL3 deep prefill if (s.images.len != 0) { const v = self.vision orelse return error.VisionNotEnabled; l.seq.media = try v.state(ids, s.images); @@ -337,7 +337,7 @@ pub fn Namespace(comptime Tri: type) type { s.keepInterruptedPrefill(@intCast(at)); return error.Cancelled; } - var stop = @min(ids.len, at + self.engine.chunk); + var stop = @min(ids.len, at + @import("exl3/prefill/cuda_exl3_prefill.zig").stepRows(self.engine, if (l.draft) |*d| d.rows_cap else null)); // EXL3 deep prefill for (s.reuse.marks) |m| if (m > at and m < stop) { stop = m; break; diff --git a/zig/src/families/flashnext/cuda_lanes_load.zig b/zig/src/families/flashnext/cuda_lanes_load.zig index 12a605c570..9193aa0dcd 100644 --- a/zig/src/families/flashnext/cuda_lanes_load.zig +++ b/zig/src/families/flashnext/cuda_lanes_load.zig @@ -89,11 +89,14 @@ pub fn open(comptime F: type, gpa: std.mem.Allocator, io: std.Io, ctx: *const cu own.next = 0; if (o.slide) own.lay = try @import("cuda_slide.zig").layoutOf(gpa, io, dir); const scratch_bytes = cuda.usage(false).device -| persistent_bytes; + const deep_before = cuda.usage(false).device; + @import("exl3/prefill/cuda_exl3_prefill.zig").reserveFor(own.engine, o.context, o.tp); // EXL3 offload: deep prefill rows before the pools + const deep_bytes = cuda.usage(false).device -| deep_before; const own_seq_bytes = try prompt.seqBytesPaged(window, try @import("cuda_pages.zig").pageTest()); var stream_bytes = own_seq_bytes; if (own.head) |h| { - stream_bytes = try std.math.add(u64, stream_bytes, mtp.Draft.bytes(window, h.ids.len, own.engine.chunk)); - } else if (o.drafts and o.tp == 2) stream_bytes = try std.math.add(u64, stream_bytes, mtp.Draft.bytes(if (mode == .exl3) 1 else window, 1, own.engine.chunk)); // an EXL3 rank 1 keeps no head cache + stream_bytes = try std.math.add(u64, stream_bytes, mtp.Draft.bytes(window, h.ids.len, @import("exl3/prefill/cuda_exl3_prefill.zig").draftRows(own.engine))); + } else if (o.drafts and o.tp == 2) stream_bytes = try std.math.add(u64, stream_bytes, mtp.Draft.bytes(if (mode == .exl3) 1 else window, 1, @import("exl3/prefill/cuda_exl3_prefill.zig").draftRows(own.engine))); // an EXL3 rank 1 keeps no head cache if (o.vision) |v| stream_bytes = try std.math.add(u64, stream_bytes, try std.math.add(u64, try std.math.mul(u64, v.image_tokens, @import("cuda_prompt_wt.zig").dims * 2), try std.math.mul(u64, 16, o.context))); if (page_tokens > 0) { // the lane's head runs at the engine's chunk: its cache on the stream's pages, its buffers in the run @@ -117,11 +120,12 @@ pub fn open(comptime F: type, gpa: std.mem.Allocator, io: std.Io, ctx: *const cu const ledger: memory.Ledger = .{ .persistent = persistent_bytes, .scratch = scratch_bytes, + .deep_prefill = try std.math.add(u64, deep_bytes, attached.prefill_window_bytes), .pool_scratch = attached.scratch_bytes, .expert_pools = attached.device_bytes, .reservation = .{ .stream_bytes = stream_bytes, .streams = if (stream_bytes == 0) 0 else pool_budget.other / stream_bytes }, }; - std.log.info("flashnext rank {d} memory ledger: persistent={d} scratch={d} pool_scratch={d} expert_pools={d} stream_bytes={d} streams={d} planned_total={d} free_after_pools={d}", .{ o.rank, ledger.persistent, ledger.scratch, ledger.pool_scratch, ledger.expert_pools, stream_bytes, ledger.reservation.streams, try ledger.total(), attached.free_after }); + std.log.info("flashnext rank {d} memory ledger: persistent={d} scratch={d} deep_prefill={d} pool_scratch={d} expert_pools={d} stream_bytes={d} streams={d} planned_total={d} free_after_pools={d}", .{ o.rank, ledger.persistent, ledger.scratch, ledger.deep_prefill, ledger.pool_scratch, ledger.expert_pools, stream_bytes, ledger.reservation.streams, try ledger.total(), attached.free_after }); return .{ .layout = own.kv.layout, .vision = own.vision, @@ -129,7 +133,7 @@ pub fn open(comptime F: type, gpa: std.mem.Allocator, io: std.Io, ctx: *const cu .facts = own.facts(), .rows = if (own.head != null) mtp.max_rows else 2, .prompt_rows = @intCast(own.engine.chunk), - .snap_grid = if (mode == .exl3) @intCast(own.engine.chunk) else 0, + .snap_grid = if (mode == .exl3) @intCast(@import("exl3/prefill/cuda_exl3_prefill.zig").stepRows(own.engine, null)) else 0, .stream_bytes = stream_bytes, .host_snaps = @import("exl3/kv/cuda_exl3_kv.zig").snapHost(o.context, mode == .exl3), .ctx = own, diff --git a/zig/src/families/flashnext/cuda_mtp_ops.zig b/zig/src/families/flashnext/cuda_mtp_ops.zig index 9afbe7ee65..7df07efe1b 100644 --- a/zig/src/families/flashnext/cuda_mtp_ops.zig +++ b/zig/src/families/flashnext/cuda_mtp_ops.zig @@ -123,6 +123,7 @@ pub fn Ops(comptime Tri: type) type { const h = o.h; if (next.len == 0) return; if (d.len + next.len > d.cap) return error.PromptTooLong; + if (try @import("exl3/prefill/cuda_exl3_prefill.zig").head(Tri, o, d, next, rows, scores)) return; // EXL3: the head expert-major over a deep chunk const want = @max(@as(usize, 32), 2 * @min(next.len, piece_rows)); const s = try e.reserve(want); const piece = @max(@as(usize, 1), @min(piece_rows, s.cap / 2)); diff --git a/zig/src/families/flashnext/cuda_prompt.zig b/zig/src/families/flashnext/cuda_prompt.zig index f606115540..7558192074 100644 --- a/zig/src/families/flashnext/cuda_prompt.zig +++ b/zig/src/families/flashnext/cuda_prompt.zig @@ -241,6 +241,7 @@ pub const Engine = struct { } pub fn deinit(e: *Engine) void { + @import("exl3/prefill/cuda_exl3_prefill.zig").releaseFor(e); e.round_graph.clear(); e.pockets.deinit(); if (e.batch_s) |*bs| bs.free(e.gpa); @@ -398,6 +399,7 @@ pub const Engine = struct { fn forwardUnmeasured(e: *Engine, comptime Tri: type, tri: Tri, seq: *Seq, ids: []const u32, hooks: batch.Hooks) !void { if (ids.len == 0 or seq.fed > seq.cap or ids.len > seq.cap - seq.fed) return error.UnexpectedTensor; + if (try @import("exl3/prefill/cuda_exl3_prefill.zig").forward(e, Tri, tri, seq, ids, hooks)) return; if ((hooks.keep != null or hooks.logits == .all) and ids.len > e.chunk) return error.UnexpectedTensor; const s = try e.reserve(@min(ids.len, e.chunk)); const scores = try e.reserveScores(seq.cap, s.cap); diff --git a/zig/src/families/flashnext/exl3/model/cuda_exl3.zig b/zig/src/families/flashnext/exl3/model/cuda_exl3.zig index 4ea7418da1..26e4251248 100644 --- a/zig/src/families/flashnext/exl3/model/cuda_exl3.zig +++ b/zig/src/families/flashnext/exl3/model/cuda_exl3.zig @@ -7,6 +7,8 @@ pub const load = @import("../pack/cuda_exl3_load.zig"); pub const pass = @import("cuda_exl3_pass.zig"); pub const tp = @import("../tp/cuda_exl3_tp.zig"); pub const kv = @import("../kv/cuda_exl3_kv.zig"); +pub const prefill_k = @import("../prefill/cuda_exl3_prefill_k.zig"); +pub const prefill = @import("../prefill/cuda_exl3_prefill.zig"); pub const fork = @import("../tp/cuda_exl3_fork.zig"); pub const telemetry = @import("cuda_exl3_telemetry.zig"); pub const frequency = @import("../offload/cuda_exl3_frequency.zig"); @@ -27,6 +29,11 @@ test { _ = kv; _ = @import("../kv/cuda_exl3_kv_test.zig"); _ = @import("../../cuda_mtp_ops_test.zig"); + _ = prefill_k; + _ = prefill; + _ = @import("../prefill/cuda_exl3_prefill_test.zig"); + _ = @import("../prefill/cuda_exl3_prefill_head.zig"); + _ = @import("../prefill/cuda_exl3_prefill_ring.zig"); _ = fork; _ = telemetry; _ = frequency; diff --git a/zig/src/families/flashnext/exl3/model/cuda_exl3_pass.zig b/zig/src/families/flashnext/exl3/model/cuda_exl3_pass.zig index dc0c966371..0e9fd23574 100644 --- a/zig/src/families/flashnext/exl3/model/cuda_exl3_pass.zig +++ b/zig/src/families/flashnext/exl3/model/cuda_exl3_pass.zig @@ -25,6 +25,7 @@ fn part(p: anytype, ops: *load.Ops, pt: load.Part, x: u64, k: usize, m: usize, y switch (pt.kind) { .exl3 => |w| { const de = try ops.linear(); + if (try @import("../prefill/cuda_exl3_prefill.zig").tilesPart(de, p.stream, w, x, m, y, f32_out)) return; // prefill rows on #806's tiles try de.decode(p.stream, w, x, .bf16, m, y, if (f32_out) .f32 else .bf16); }, .native => |w| { @@ -192,7 +193,7 @@ test "EXL3 pool reserves the requested number of concurrent stream caches" { try std.testing.expectError(error.InvalidStreamCount, streamReservation(100, 0)); } -pub const Attached = struct { layers: usize = 0, lines: usize = 0, line_bytes: usize = 0, device_bytes: u64 = 0, free_before: u64 = 0, free_after: u64 = 0, scratch_bytes: u64 = 0, bank_bytes: u64 = 0, bank_ns: u64 = 0 }; +pub const Attached = struct { layers: usize = 0, lines: usize = 0, line_bytes: usize = 0, device_bytes: u64 = 0, free_before: u64 = 0, free_after: u64 = 0, scratch_bytes: u64 = 0, bank_bytes: u64 = 0, bank_ns: u64 = 0, prefill_window_bytes: u64 = 0 }; pub fn freeBytes(d: *const cuda.Driver) !u64 { var free: usize = 0; @@ -242,6 +243,17 @@ pub fn attach(gpa: std.mem.Allocator, io: std.Io, d: *const cuda.Driver, st: cud fixed = try std.math.add(u64, fixed, try std.math.add(u64, cuda.exl3_pool.Tables.init(wt.experts_n + 1).bytes + (64 << 10), picks_bytes)); } const future = b.other; + const prefill = @import("../prefill/cuda_exl3_prefill.zig"); + const minimum_pools = try memory.poolBytes(b.lines orelse pool_.min_lines, out.line_bytes, n, fixed); + const protected = try std.math.add(u64, try std.math.add(u64, reserve, future), minimum_pools); + out.prefill_window_bytes = prefill.reserveTwoChunk(d, free, protected) catch |err| declined: { + std.log.warn("flashnext exl3 prefill: two-chunk reservation refused ({t}); keeping single-chunk prefill", .{err}); + break :declined 0; + }; + if (out.prefill_window_bytes != 0) { + free = @intCast(try freeBytes(d)); + out.free_before = free; + } const pool_room = try memory.available(free, reserve, future); const lines = b.lines orelse try pool_.linesFor(free, reserve, try std.math.add(u64, future, fixed), out.line_bytes, n, pool_.min_lines); if (lines < pool_.min_lines) return error.PoolTooSmall; diff --git a/zig/src/families/flashnext/exl3/model/cuda_exl3_telemetry.zig b/zig/src/families/flashnext/exl3/model/cuda_exl3_telemetry.zig index e55b1cc1c6..af7527dfe2 100644 --- a/zig/src/families/flashnext/exl3/model/cuda_exl3_telemetry.zig +++ b/zig/src/families/flashnext/exl3/model/cuda_exl3_telemetry.zig @@ -3,6 +3,7 @@ const prompt = @import("../../cuda_prompt.zig"); const load = @import("../../cuda_prompt_load.zig"); const pool = @import("../offload/cuda_exl3_pool.zig"); const pass = @import("cuda_exl3_pass.zig"); +const prefill = @import("../prefill/cuda_exl3_prefill.zig"); const weights = @import("../../cuda_prompt_wt.zig"); pub const counters = @import("exl3_telemetry.zig"); pub const Completion = enum { all_rank_consumers_joined }; @@ -76,8 +77,19 @@ pub fn sampleAfterConsumers(allocator: std.mem.Allocator, engine: *prompt.Engine } } for (sequences, 0..) |sequence, index| for (sequences[0..index]) |previous| if (previous == sequence) return error.TelemetryDuplicateOwner; + const ring = prefill.ringTelemetry(engine); + if (ring) |observed| { + totals.ring_dma_copies = observed.copies; + totals.ring_dma_requested_bytes = observed.dma_requested_bytes; + totals.ring_unused_copies = observed.unused_copies; + totals.ring_hit_picks = observed.hit_picks; + totals.ring_miss_picks = observed.miss_picks; + totals.ring_hit_expert_requests = observed.hit_expert_requests; + totals.ring_miss_expert_requests = observed.miss_expert_requests; + if (observed.attribution) |by_source| @import("../prefill/cuda_exl3_prefill_attribution.zig").publish(by_source, &totals); + } const sampled_layers = if (engine.telemetry.layers_enabled) try allocator.realloc(layers, seen_count) else layers; - return .{ .rank = rank, .scope = hash.final(), .experts = totals, .expert_pools = seen_count, .layers = sampled_layers }; + return .{ .rank = rank, .scope = hash.final(), .ring_owner = if (ring) |observed| observed.owner else 0, .experts = totals, .expert_pools = seen_count, .layers = sampled_layers }; } test "exl3: cache telemetry pure host contracts" { diff --git a/zig/src/families/flashnext/exl3/model/cuda_exl3_telemetry_test.zig b/zig/src/families/flashnext/exl3/model/cuda_exl3_telemetry_test.zig index 4f8e4166b9..3e48bb222e 100644 --- a/zig/src/families/flashnext/exl3/model/cuda_exl3_telemetry_test.zig +++ b/zig/src/families/flashnext/exl3/model/cuda_exl3_telemetry_test.zig @@ -3,6 +3,19 @@ const cuda = @import("cuda"); const pool = @import("../offload/cuda_exl3_pool.zig"); const telemetry = @import("exl3_telemetry.zig"); +test "exl3: ring demand counts routed pick appearances separately from per-chunk expert requests and speculative copies" { + var stats: @import("../prefill/cuda_exl3_prefill_ring.zig").Stats = .{}; + stats.observeDemand(true, 17); + stats.observeDemand(false, 31); + stats.observeDemand(false, 0); + try std.testing.expectEqual(@as(u64, 17), stats.hit_picks); + try std.testing.expectEqual(@as(u64, 31), stats.miss_picks); + try std.testing.expectEqual(@as(u64, 1), stats.hit_expert_requests); + try std.testing.expectEqual(@as(u64, 1), stats.miss_expert_requests); + try std.testing.expectEqual(@as(u64, 0), stats.copies); + try std.testing.expectEqual(@as(u64, 0), stats.dma_requested_bytes); +} + test "exl3: aggregate telemetry finds attached pools deduplicates MTP aliases and separates committed device plans" { var resident: pool.Pool = undefined; resident.lines.pinned_bank = null; diff --git a/zig/src/families/flashnext/exl3/model/exl3_telemetry.zig b/zig/src/families/flashnext/exl3/model/exl3_telemetry.zig index bb03c39158..9ce723700f 100644 --- a/zig/src/families/flashnext/exl3/model/exl3_telemetry.zig +++ b/zig/src/families/flashnext/exl3/model/exl3_telemetry.zig @@ -84,11 +84,31 @@ pub const Experts = struct { device_overflow_plans: u64 = 0, mapped_copy_payload_min_bytes: u64 = 0, mapped_copy_payload_max_bytes: u64 = 0, + ring_dma_copies: u64 = 0, + ring_dma_requested_bytes: u64 = 0, + ring_unused_copies: u64 = 0, + ring_hit_picks: u64 = 0, + ring_miss_picks: u64 = 0, + ring_hit_expert_requests: u64 = 0, + ring_miss_expert_requests: u64 = 0, + ring_demand_copies: u64 = 0, + ring_demand_requested_bytes: u64 = 0, + ring_demand_unused_copies: u64 = 0, + ring_demand_unused_requested_bytes: u64 = 0, + ring_initial_prefetch_copies: u64 = 0, + ring_initial_prefetch_requested_bytes: u64 = 0, + ring_initial_prefetch_unused_copies: u64 = 0, + ring_initial_prefetch_unused_requested_bytes: u64 = 0, + ring_cross_layer_prefetch_copies: u64 = 0, + ring_cross_layer_prefetch_requested_bytes: u64 = 0, + ring_cross_layer_prefetch_unused_copies: u64 = 0, + ring_cross_layer_prefetch_unused_requested_bytes: u64 = 0, }; pub const Snapshot = struct { rank: Rank, scope: u64, + ring_owner: u64 = 0, experts: Experts = .{}, expert_pools: usize = 0, layers: []LayerSample = &.{}, @@ -214,6 +234,7 @@ pub const Report = struct { try before.rank.validate(); try after.rank.validate(); if (!std.meta.eql(before.rank, after.rank) or before.scope != after.scope or before.expert_pools != after.expert_pools) return error.TelemetryOwnerChanged; + if (before.ring_owner != 0 and before.ring_owner != after.ring_owner) return error.TelemetryOwnerChanged; if (interval.model_rows == 0 or interval.steps == 0) return error.TelemetryDenominator; const experts = try difference(Experts, before.experts, after.experts); if (experts.device_overflow_plans != 0) return error.TelemetryExpertOverflow; diff --git a/zig/src/families/flashnext/exl3/model/exl3_telemetry_test.zig b/zig/src/families/flashnext/exl3/model/exl3_telemetry_test.zig index 459bace311..84b375bf4d 100644 --- a/zig/src/families/flashnext/exl3/model/exl3_telemetry_test.zig +++ b/zig/src/families/flashnext/exl3/model/exl3_telemetry_test.zig @@ -30,7 +30,7 @@ test "exl3: invalid device telemetry cannot fail normal ownership transfer or pu } test "exl3: phase report stable schema rank and model-row denominators" { - const before: telemetry.Snapshot = .{ .rank = .{ .rank = 1, .world = 2 }, .scope = 7 }; + var before: telemetry.Snapshot = .{ .rank = .{ .rank = 1, .world = 2 }, .scope = 7 }; var after = before; after.experts.device_hit_picks = 40; after.experts.device_miss_picks = 10; @@ -45,6 +45,11 @@ test "exl3: phase report stable schema rank and model-row denominators" { after = before; after.scope += 1; try expectError(error.TelemetryOwnerChanged, telemetry.Report.between(before, after, .{ .phase = .decode, .model_rows = 1, .steps = 1 })); + after = before; + after.ring_owner = 13; + _ = try telemetry.Report.between(before, after, .{ .phase = .prefill, .model_rows = 1, .steps = 1 }); + before.ring_owner = 12; + try expectError(error.TelemetryOwnerChanged, telemetry.Report.between(before, after, .{ .phase = .prefill, .model_rows = 1, .steps = 1 })); } test "exl3: per-layer deltas sum to aggregate counters and overlapping reports claim counters once" { diff --git a/zig/src/families/flashnext/exl3/offload/cuda_exl3_pool.zig b/zig/src/families/flashnext/exl3/offload/cuda_exl3_pool.zig index 647cfce0f9..9ae062fe28 100644 --- a/zig/src/families/flashnext/exl3/offload/cuda_exl3_pool.zig +++ b/zig/src/families/flashnext/exl3/offload/cuda_exl3_pool.zig @@ -243,6 +243,10 @@ pub const Pool = struct { return p.planInternal(picks, rows, true); } + pub fn planUnobserved(p: *Pool, picks: []const i32, rows: usize) !*const Round { + return p.planInternal(picks, rows, false); + } + fn planInternal(p: *Pool, picks: []const i32, rows: usize, observe: bool) !*const Round { if (rows == 0 or rows > p.opt.max_rows or picks.len != rows * p.opt.slots) return error.Invalid; p.now += 1; diff --git a/zig/src/families/flashnext/exl3/pack/exl3_startup_memory.zig b/zig/src/families/flashnext/exl3/pack/exl3_startup_memory.zig index d924759dec..7bf9ed4655 100644 --- a/zig/src/families/flashnext/exl3/pack/exl3_startup_memory.zig +++ b/zig/src/families/flashnext/exl3/pack/exl3_startup_memory.zig @@ -12,13 +12,14 @@ pub const Reservation = struct { pub const Ledger = struct { persistent: u64, scratch: u64, + deep_prefill: u64, pool_scratch: u64, expert_pools: u64, reservation: Reservation, pub fn total(self: Ledger) !u64 { var bytes = try self.reservation.total(); - inline for (.{ self.persistent, self.scratch, self.pool_scratch, self.expert_pools }) |part| bytes = try std.math.add(u64, bytes, part); + inline for (.{ self.persistent, self.scratch, self.deep_prefill, self.pool_scratch, self.expert_pools }) |part| bytes = try std.math.add(u64, bytes, part); return bytes; } }; @@ -56,8 +57,8 @@ test "startup accounting fails closed on overflow and underfunded explicit pools } test "ledger includes all allocations after the pool checkpoint" { - const ledger: Ledger = .{ .persistent = 100, .scratch = 20, .pool_scratch = 40, .expert_pools = 50, .reservation = .{ .stream_bytes = 10, .streams = 8 } }; - try std.testing.expectEqual(@as(u64, 290), try ledger.total()); + const ledger: Ledger = .{ .persistent = 100, .scratch = 20, .deep_prefill = 30, .pool_scratch = 40, .expert_pools = 50, .reservation = .{ .stream_bytes = 10, .streams = 8 } }; + try std.testing.expectEqual(@as(u64, 320), try ledger.total()); } test "pool setup may consume its allocator margin but not lane or server reservations" { diff --git a/zig/src/families/flashnext/exl3/prefill/cuda_exl3_prefill.zig b/zig/src/families/flashnext/exl3/prefill/cuda_exl3_prefill.zig new file mode 100644 index 0000000000..de7438efb9 --- /dev/null +++ b/zig/src/families/flashnext/exl3/prefill/cuda_exl3_prefill.zig @@ -0,0 +1,333 @@ +//! Expert-major prefill for an EXL3 Flash Next checkpoint whose routed experts live in host memory (the offload pool). + +const std = @import("std"); +const cuda = @import("cuda"); +const batch = @import("../../cuda_batch.zig"); +const wt = @import("../../cuda_prompt_wt.zig"); +const pk = @import("cuda_exl3_prefill_k.zig"); +const prompt = @import("../../cuda_prompt.zig"); +const window_ = @import("exl3_prefill_window.zig"); +pub const state = @import("cuda_exl3_prefill_state.zig"); +const chunks = @import("cuda_exl3_prefill_chunk.zig"); +const ringmoe = @import("cuda_exl3_prefill_ringmoe.zig"); +const heads = @import("cuda_exl3_prefill_head.zig"); +const tiles = @import("cuda_exl3_prefill_tiles.zig"); + +const dims = wt.dims; +const wide = wt.wide; +const slots = wt.slots; +const vocab = wt.vocab; +const att_rows = batch.att_rows; + +const Deep = state.Deep; +const HostLayout = state.HostLayout; +const Layout = state.Layout; +const Persistent = state.Persistent; +const opts = state.opts; +const chunk = chunks.chunk; +const selectWriteback = chunks.selectWriteback; +const twoChunk = chunks.twoChunk; +const ringOf = ringmoe.ringOf; +const warmTiles = tiles.warmTiles; + +/// Frees the deep chunk's buffers (the engine's deinit, or a test). +pub fn release() void { + if (state.deep) |dp| { + if (dp.second) |second| { + dp.aux.synchronize() catch @panic("two-chunk aux retirement synchronization failed"); + dp.dev.d.synchronize() catch @panic("two-chunk compute retirement synchronization failed"); + second.host.free(); + if (dp.red != null) { + for (&second.red_ev) |*event| event.deinit(); + second.red_go.deinit(); + } + if (second.sum_dev) |*buffer| buffer.free(); + second.dev.free(); + dp.gpa.destroy(second); + } + if (dp.ring) |r| r.deinit(); + if (dp.red) |*rs| { + rs.synchronize() catch {}; + for (&dp.red_ev) |*ev| ev.deinit(); + dp.red_go.deinit(); + rs.deinit(); + } + if (dp.unit_weights) |*buffer| buffer.free(); + dp.kern.unload(); + dp.aux.deinit(); + dp.host.free(); + dp.dev.free(); + dp.gpa.destroy(dp); + state.deep = null; + } + if (agree_buf) |*b| b.free(); + agree_buf = null; +} + +/// The server's sizing hook (cuda_lanes_load.zig, before the pools take the free memory): the deep chunk's buffers for +pub fn reserveFor(e: *prompt.Engine, context: usize, tp: usize) void { + const o = opts(); + if (!o.enabled or (tp != 1 and tp != 2) or e.n_layers != wt.layers_n) return; + for (e.layers[0..e.n_layers]) |*l| { + const x = l.exl3 orelse return; + if (x.bank == null) return; + // expert parallel: every layer's experts carry this world's rank + if ((if (x.tp) |r| r.world else 1) != tp) return; + } + const rows = @min(o.serve_rows, (context + att_rows - 1) / att_rows * att_rows); + if (rows < o.min_rows) return; + reserveAt(e.gpa, e.driver, e, rows, tp) catch |err| std.log.warn("flashnext exl3 prefill: deep chunks of {d} rows refused ({t}); prompts take the batched pass", .{ rows, err }); +} + +/// Rows a serving lane prefills a step (cuda_lanes.zig `step`): a deep chunk when this engine holds deep buffers. +pub fn stepRows(e: *const prompt.Engine, draft_rows: ?usize) usize { + var n = e.chunk; + if (state.deep) |dp| if (dp.owner == @as(*const anyopaque, e) and opts().enabled) { + n = @max(n, stepCapacity(dp.rows, opts().two_chunk and dp.second != null)); + }; + if (draft_rows) |r| n = @min(n, @max(r, 1)); + return n; +} + +pub fn stepCapacity(rows: usize, paired: bool) usize { + return if (paired and rows <= window_.max_rows) rows * 2 else rows; +} + +/// Rows of streams a lane's MTP draft keeps (cuda_lanes.zig, cuda_lanes_load.zig): a whole step's. +pub fn draftRows(e: *const prompt.Engine) usize { + var n = stepRows(e, null); + if (state.deep) |dp| if (dp.owner == @as(*const anyopaque, e)) { + const o = opts(); + if (o.two_chunk and o.enabled and o.ring and !o.reduced_branch and dp.world == 2) { + n = @max(n, stepCapacity(dp.rows, true)); + } + }; + return n; +} + +/// The engine's deinit: frees the deep chunk's buffers if `owner` made them. +pub fn releaseFor(owner: *const anyopaque) void { + if (state.deep) |dp| if (dp.owner == owner) release(); +} + +/// Allocates the deep chunk's buffers for `rows` rows now (before the pool takes the free memory). +pub fn reserve(gpa: std.mem.Allocator, d: *const cuda.Driver, owner: *const anyopaque, rows_in: usize) !void { + return reserveAt(gpa, d, owner, rows_in, 1); +} + +/// `reserve` for a rank of `world` (expert parallel at 2: the ranks' sum, and per-slot outputs zeroed once so a slot +pub fn reserveAt(gpa: std.mem.Allocator, d: *const cuda.Driver, owner: *const anyopaque, rows_in: usize, world: usize) !void { + const rows = rows_in / att_rows * att_rows; + if (rows == 0 or (world != 1 and world != 2)) return error.Invalid; + if (state.deep) |dp| { + if (dp.owner == owner and dp.rows >= rows and dp.world == world) return; + release(); + } + const dp = try gpa.create(Deep); + errdefer gpa.destroy(dp); + dp.* = .{ .gpa = gpa, .owner = owner, .rows = rows, .world = world, .lay = Layout.of(rows, world), .hl = HostLayout.of(rows), .dev = undefined, .host = undefined, .kern = undefined, .aux = undefined }; + dp.dev = try cuda.DeviceBuffer.alloc(d, dp.lay.total); + errdefer dp.dev.free(); + dp.host = try cuda.HostBuffer.alloc(d, dp.hl.total); + errdefer dp.host.free(); + dp.kern = try pk.Kernels.load(d); + errdefer dp.kern.unload(); + dp.aux = try cuda.Stream.init(d, true); + errdefer dp.aux.deinit(); + if (world > 1) { + try d.check(d.api.cuMemsetD8Async(dp.dptr(dp.lay.y), 0, rows * slots * dims * 4, dp.aux.handle), "cuMemsetD8Async"); + try dp.aux.synchronize(); + var rs = try cuda.Stream.init(d, true); + errdefer rs.deinit(); + var made: usize = 0; + errdefer for (dp.red_ev[0..made]) |*ev| ev.deinit(); + for (&dp.red_ev) |*ev| { + ev.* = try cuda.Event.init(d, false); + made += 1; + } + dp.red_go = try cuda.Event.init(d, false); + dp.red = rs; + } + state.deep = dp; + std.log.info("flashnext exl3 prefill: deep chunks of {d} rows, {d:.2} GiB on the device, {d:.1} MiB pinned (world {d})", .{ rows, @as(f64, @floatFromInt(dp.lay.total)) / (1 << 30), @as(f64, @floatFromInt(dp.hl.total)) / (1 << 20), world }); +} + +pub fn reserveTwoChunk(d: *const cuda.Driver, free: u64, protected: u64) !u64 { + const o = opts(); + if (!o.two_chunk or !o.enabled or !o.ring or o.reduced_branch) return 0; + const dp = state.deep orelse return 0; + if (dp.dev.d != d or dp.world != 2 or dp.rows > window_.max_rows) return 0; + if (dp.second != null) return 0; + const persistent = Persistent.of(dp.rows); + const sum_bytes = if (dp.world > 1) state.sumBytes(dp.rows, dp.world) else 0; + const device_bytes = try std.math.add(usize, persistent.total, sum_bytes); + if (!window_.admits(free, protected, device_bytes)) { + std.log.info("flashnext exl3 prefill: two-chunk window declined, extra_device={d} full_layout={d} extra_pinned=0 protected={d} free={d}", .{ device_bytes, dp.lay.total, protected, free }); + return 0; + } + const second = try dp.gpa.create(Deep); + second.* = .{ .gpa = dp.gpa, .owner = dp.owner, .rows = dp.rows, .world = dp.world, .lay = dp.lay, .hl = dp.hl, .dev = undefined, .host = undefined, .kern = dp.kern, .aux = dp.aux, .red = dp.red, .red_ev = undefined, .red_go = undefined, .shared = dp }; + var device_allocated = false; + var host_allocated = false; + var made_events: usize = 0; + var reduce_event_allocated = false; + errdefer { + if (reduce_event_allocated) second.red_go.deinit(); + for (second.red_ev[0..made_events]) |*event| event.deinit(); + if (host_allocated) second.host.free(); + if (second.sum_dev) |*buffer| buffer.free(); + if (device_allocated) second.dev.free(); + dp.gpa.destroy(second); + } + second.dev = try cuda.DeviceBuffer.alloc(d, persistent.total); + device_allocated = true; + if (sum_bytes != 0) { + second.sum_dev = try cuda.DeviceBuffer.alloc(d, sum_bytes); + } + second.host = try cuda.HostBuffer.alloc(d, second.hl.total); + host_allocated = true; + if (dp.red != null) { + for (&second.red_ev) |*event| { + event.* = try cuda.Event.init(d, false); + made_events += 1; + } + second.red_go = try cuda.Event.init(d, false); + reduce_event_allocated = true; + } else { + second.red_ev = dp.red_ev; + second.red_go = dp.red_go; + } + try d.check(d.api.cuMemsetD8Async(second.dptr(second.lay.y), 0, dp.rows * slots * dims * 4, dp.aux.handle), "cuMemsetD8Async"); + try dp.aux.synchronize(); + dp.second = second; + std.log.info("flashnext exl3 prefill: experimental two-chunk window, chunk_rows={d} hidden_bytes={d} full_layout={d} extra_device={d} extra_pinned={d}", .{ dp.rows, 2 * dp.rows * wide * 2, dp.lay.total, device_bytes, second.hl.total }); + return device_bytes; +} + +pub fn twoChunkAllocated() bool { + return if (state.deep) |buffers| buffers.second != null else false; +} + +fn applies(e: *prompt.Engine, seq: *prompt.Seq, ids: []const u32, hooks: batch.Hooks, o: Options) bool { + if (!o.enabled or ids.len < o.min_rows) return false; + if (hooks.keep != null or hooks.logits == .all) return false; + if (seq.media != null or e.n_layers != wt.layers_n) return false; + // expert parallel (cuda_exl3_tp.zig) at two ranks over the engine's collective. + const world = worldOf(e); + if (e.head_q.n != vocab and !(world == 2 and e.head_q.n == vocab / 2)) return false; + for (e.layers[0..e.n_layers]) |*l| { + const x = l.exl3 orelse return false; + if ((if (x.tp) |r| r.world else 1) != world or x.needsPool()) return false; + if (x.layer.act_mode != 1 and x.layer.act_mode != 0) return false; + } + return e.chunk >= att_rows; +} + +fn worldOf(e: *const prompt.Engine) usize { + return if (e.collective != null) 2 else 1; +} + +/// Expert parallel: both ranks take the deep path or neither, with the same deep chunk (the smaller rank's): each +fn agree(e: *prompt.Engine, mine: usize) !usize { + const c = e.collective orelse return mine; + const st = e.stream.*; + const rank = e.layers[0].exl3.?.tp.?.rank; + if (agree_buf == null) agree_buf = try cuda.DeviceBuffer.alloc(e.driver, 256); + const b = agree_buf.?; + var v: [2]f32 = .{ 0, 0 }; + v[rank] = @floatFromInt(mine); + try st.synchronize(); + try b.upload(0, std.mem.asBytes(&v)); + try c.sum(st, b.ptr, 2); + try c.finish(st); + try b.download(0, std.mem.asBytes(&v)); + if (v[0] < 1 or v[1] < 1) return 0; + return @intFromFloat(@min(v[0], v[1])); +} +var agree_buf: ?cuda.DeviceBuffer = null; +var tp_logged = false; + +/// The hook (Engine.forward): runs `ids` from `seq.fed` and returns true, or returns false having done nothing. +pub fn forward(e: *prompt.Engine, comptime Tri: type, tri: Tri, seq: *prompt.Seq, ids: []const u32, hooks: batch.Hooks) !bool { + const o = opts(); + if (!applies(e, seq, ids, hooks, o)) return false; + if (ids.len == 0 or seq.fed + ids.len > seq.cap) return error.UnexpectedTensor; + const s = try e.reserve(@min(ids.len, e.chunk)); + const sub = s.cap / att_rows * att_rows; + if (sub == 0) return false; + // the deep chunk: as large as asked (whole attention blocks), no larger than the prompt needs; its sub-blocks are + const rows = @max(att_rows, @min(o.rows, (ids.len + att_rows - 1) / att_rows * att_rows) / att_rows * att_rows); + // the buffers reserved at load (reserveFor) are kept whatever their size; a first prompt without them allocates + const world = worldOf(e); + const have = if (state.deep) |d| d.owner == @as(*const anyopaque, e) and d.world == world else false; + var room: usize = if (have) state.deep.?.rows else 0; + if (!have) { + if (reserveAt(e.gpa, e.driver, e, rows, world)) |_| { + room = state.deep.?.rows; + } else |err| std.log.warn("flashnext exl3 prefill: no room for deep chunks of {d} rows ({t}); the batched pass takes the prompt", .{ rows, err }); + } + // both ranks the same way (a collective a rank skips would hang the other) + room = try agree(e, room); + if (room == 0) return false; + if (world == 2 and !tp_logged) std.log.info("flashnext exl3 prefill: expert-parallel deep chunks of up to {d} rows (rank {d}: its own experts, one all-reduce a layer)", .{ room, e.layers[0].exl3.?.tp.?.rank }); + tp_logged = world == 2 or tp_logged; + const dp = state.deep.?; + try selectWriteback(dp, o.reduced_branch and world == 2); + const scores = try e.reserveScores(seq.cap, s.cap); + var p: batch.Pass(Tri) = .{ + .gpa = e.gpa, + .io = e.io, + .stream = e.stream.*, + .tri = tri, + .k = &e.batch_k, + .s = s, + .m = .{ .table_w = e.table.w.ptr, .table_s = e.table.s.ptr, .table_b = e.table.b.ptr, .table_dense = e.table.dense, .layers = e.layers[0..e.n_layers], .mixer = e.mixer, .head = e.head_q, .ngram = &e.ngram, .collective = e.collective }, + .seq = seq.view(scores), + .hooks = hooks, + }; + // equal deep chunks (a short last chunk would pay a whole chunk's expert copies for a few rows) + const cap_rows = @min(rows, room); + const n_chunks = (ids.len + cap_rows - 1) / cap_rows; + const deep_rows = @min(cap_rows, ((ids.len + n_chunks - 1) / n_chunks + att_rows - 1) / att_rows * att_rows); + const rg = if (o.ring) try ringOf(dp, e) else null; + var window_room: usize = 0; + if (o.two_chunk) { + const supported = dp.second != null and rg != null and world == 2 and !dp.reduced_branch and deep_rows <= window_.max_rows and ids.len > deep_rows; + window_room = try agree(e, if (supported) deep_rows else 0); + } + if (o.tiles and !o.slots_kernel) { + try warmTiles(e); + tiles.tiles_on = true; + } + defer tiles.tiles_on = false; + if (rg) |r| try r.begin(p.stream); + defer if (rg) |r| r.end(p.stream, o.reseed) catch |err| std.log.err("flashnext exl3 prefill: ring end: {t}", .{err}); + var at: usize = 0; + while (at < ids.len) { + if (window_room != 0) { + const plan_ = try window_.Plan.init(ids.len - at, window_room, att_rows); + try twoChunk(Tri, dp, &p, s, ids[at..][0..plan_.rows], seq.fed + at, sub, hooks, at, at + plan_.rows == ids.len, seq.logits, o, rg.?, plan_); + at += plan_.rows; + continue; + } + const n = @min(deep_rows, ids.len - at); + const last = at + n == ids.len; + const lg: batch.Logits = if (last) hooks.logits else .none; + const st: u64 = if (hooks.streams != 0) hooks.streams + at * wide * 2 else 0; + try chunk(Tri, dp, &p, s, ids[at..][0..n], seq.fed + at, sub, lg, st, if (last and hooks.logits != .none) seq.logits else null, o, rg); + at += n; + } + if (hooks.logits == .none) try p.stream.synchronize(); + seq.fed += ids.len; + return true; +} + +pub const Options = state.Options; +pub const optionsFromEnv = state.optionsFromEnv; +pub const Stats = state.Stats; +pub const RingTelemetry = state.RingTelemetry; +pub const ringTelemetry = state.ringTelemetry; +pub const deviceBytes = state.deviceBytes; +pub const deviceBytesAt = state.deviceBytesAt; +pub const head = heads.head; +pub const tilesPart = tiles.tilesPart; diff --git a/zig/src/families/flashnext/exl3/prefill/cuda_exl3_prefill_attribution.zig b/zig/src/families/flashnext/exl3/prefill/cuda_exl3_prefill_attribution.zig new file mode 100644 index 0000000000..927d23cdee --- /dev/null +++ b/zig/src/families/flashnext/exl3/prefill/cuda_exl3_prefill_attribution.zig @@ -0,0 +1,49 @@ +const std = @import("std"); +const telemetry = @import("../model/exl3_telemetry.zig"); + +pub const Source = enum { demand, initial_prefetch, cross_layer_prefetch }; +pub const Counts = struct { + copies: u64 = 0, + requested_bytes: u64 = 0, + unused_copies: u64 = 0, + unused_requested_bytes: u64 = 0, +}; +pub const Totals = [3]Counts; + +const Residency = struct { source: Source, bytes: u64 }; + +pub const State = struct { + slots: []Residency, + totals: Totals = @splat(.{}), + + pub fn init(allocator: std.mem.Allocator, slots: usize, enabled: bool) !?State { + if (!enabled) return null; + return .{ .slots = try allocator.alloc(Residency, slots) }; + } + + pub fn deinit(state: *State, allocator: std.mem.Allocator) void { + allocator.free(state.slots); + } + + pub fn copied(state: *State, slot: usize, source: Source, bytes: u64) void { + state.slots[slot] = .{ .source = source, .bytes = bytes }; + const counts = &state.totals[@backingInt(source)]; + counts.copies += 1; + counts.requested_bytes += bytes; + } + + pub fn unused(state: *State, slot: usize) void { + const residency = state.slots[slot]; + const counts = &state.totals[@backingInt(residency.source)]; + counts.unused_copies += 1; + counts.unused_requested_bytes += residency.bytes; + } +}; + +pub fn publish(totals: Totals, experts: *telemetry.Experts) void { + inline for (std.enums.values(Source), 0..) |source, index| { + inline for (@typeInfo(Counts).@"struct".field_names) |field| { + @field(experts.*, "ring_" ++ @tagName(source) ++ "_" ++ field) = @field(totals[index], field); + } + } +} diff --git a/zig/src/families/flashnext/exl3/prefill/cuda_exl3_prefill_attribution_test.zig b/zig/src/families/flashnext/exl3/prefill/cuda_exl3_prefill_attribution_test.zig new file mode 100644 index 0000000000..b27cb39220 --- /dev/null +++ b/zig/src/families/flashnext/exl3/prefill/cuda_exl3_prefill_attribution_test.zig @@ -0,0 +1,310 @@ +const std = @import("std"); +const cuda = @import("cuda"); +const attribution = @import("cuda_exl3_prefill_attribution.zig"); +const ring_module = @import("cuda_exl3_prefill_ring.zig"); +const pool_module = @import("../offload/cuda_exl3_pool.zig"); +const telemetry = @import("../model/exl3_telemetry.zig"); +const State = std.meta.Child(@FieldType(ring_module.Ring, "state")); +const allocator = std.testing.allocator; + +const Calls = struct { + copies: usize = 0, + payload_bytes: usize = 0, + records: usize = 0, + queries: usize = 0, + waits: usize = 0, + syncs: usize = 0, +}; + +const Fixture = struct { + driver: cuda.Driver, + ring: *ring_module.Ring, + pool: *pool_module.Pool, + payload: [8]u8, + sources: [512][3][]const u8, + counts: [48][512]u32, + addresses: [2]u64, + offsets: [2][3]u32, + rooms: [2][3]u32, + states: [2]State, + owners: [2]u32, + epochs: [2]u64, + + var calls: Calls = .{}; + var copy_result: cuda.abi.Result = cuda.abi.success; + + fn copy(destination: u64, _: ?*const anyopaque, bytes: usize, _: cuda.abi.Stream) callconv(.c) cuda.abi.Result { + calls.copies += 1; + if (destination < 0x9000) calls.payload_bytes += bytes; + return copy_result; + } + + fn record(_: cuda.abi.Event, _: cuda.abi.Stream) callconv(.c) cuda.abi.Result { + calls.records += 1; + return cuda.abi.success; + } + + fn query(_: cuda.abi.Event) callconv(.c) cuda.abi.Result { + calls.queries += 1; + return cuda.abi.error_not_ready; + } + + fn wait(_: cuda.abi.Stream, _: cuda.abi.Event, _: c_uint) callconv(.c) cuda.abi.Result { + calls.waits += 1; + return cuda.abi.success; + } + + fn sync(_: cuda.abi.Event) callconv(.c) cuda.abi.Result { + calls.syncs += 1; + return cuda.abi.success; + } + + fn init(enabled: bool) !*Fixture { + const fixture = try allocator.create(Fixture); + errdefer allocator.destroy(fixture); + fixture.* = undefined; + fixture.ring = try allocator.create(ring_module.Ring); + errdefer allocator.destroy(fixture.ring); + fixture.pool = try allocator.create(pool_module.Pool); + errdefer allocator.destroy(fixture.pool); + fixture.pool.* = undefined; + fixture.pool.opt = .{ .layer = 0, .lines = 2 }; + fixture.pool.src = .{ .routed = &fixture.sources }; + fixture.pool.counts = &fixture.counts; + fixture.pool.lines.tables.ptr = 0x9000; + fixture.payload = @splat(7); + for (&fixture.sources, 0..) |*source, index| { + const length = index % 4 + 1; + source.* = .{ fixture.payload[0..length], fixture.payload[0 .. length + 1], fixture.payload[0 .. length + 2] }; + } + for (&fixture.counts) |*row| @memset(row, 0); + fixture.addresses = .{ 0x1000, 0x2000 }; + fixture.offsets = @splat(.{ 0, 8, 16 }); + fixture.rooms = @splat(.{ 8, 8, 8 }); + fixture.states = @splat(.free); + fixture.owners = @splat(0); + fixture.epochs = @splat(0); + fixture.driver = undefined; + fixture.driver.api.cuMemcpyHtoDAsync_v2 = copy; + fixture.driver.api.cuEventRecord = record; + fixture.driver.api.cuEventQuery = query; + fixture.driver.api.cuStreamWaitEvent = wait; + fixture.driver.api.cuEventSynchronize = sync; + const ring = fixture.ring; + ring.* = undefined; + ring.gpa = allocator; + ring.d = &fixture.driver; + ring.pools = @splat(fixture.pool); + ring.addr = &fixture.addresses; + ring.off = &fixture.offsets; + ring.room_of = &fixture.rooms; + ring.state = &fixture.states; + ring.owner = &fixture.owners; + ring.epoch_of = &fixture.epochs; + ring.next = 0; + ring.pending_n = 0; + ring.epoch = 0; + ring.copy_ordered_epoch = 0; + ring.batches = 0; + ring.open_layer = null; + ring.open_n = 0; + ring.own_lo = 0; + ring.own_hi = 512; + ring.copy = .{ .d = &fixture.driver, .handle = @ptrFromInt(1) }; + for (&ring.cev, 0..) |*event, index| event.* = .{ .d = &fixture.driver, .handle = @ptrFromInt(index + 1) }; + for (&ring.rev, 0..) |*event, index| event.* = .{ .d = &fixture.driver, .handle = @ptrFromInt(index + 1) }; + ring.tlay = cuda.exl3_pool.Tables.init(513); + const shadow = try allocator.alignedAlloc(u8, .@"16", 48 * ring.tlay.bytes); + errdefer allocator.free(shadow); + @memset(shadow, 0); + ring.shadow = .{ .d = &fixture.driver, .bytes = shadow }; + const snapshots = try allocator.alignedAlloc(u8, .@"16", 64 * ring.tlay.k2[0]); + errdefer allocator.free(snapshots); + @memset(snapshots, 0); + ring.snaps = .{ .d = &fixture.driver, .bytes = snapshots }; + for (&ring.slot) |*row| @memset(row, -1); + for (&ring.counts) |*row| @memset(row, 0); + ring.stats = .{}; + ring.attribution = try attribution.State.init(allocator, 2, enabled); + calls = .{}; + copy_result = cuda.abi.success; + return fixture; + } + + fn deinit(fixture: *Fixture) void { + if (fixture.ring.attribution) |*state| state.deinit(allocator); + allocator.free(fixture.ring.snaps.bytes); + allocator.free(fixture.ring.shadow.bytes); + allocator.destroy(fixture.pool); + allocator.destroy(fixture.ring); + allocator.destroy(fixture); + } + + fn checksum(fixture: *Fixture, unused_bytes: u64) !void { + var sum: attribution.Counts = .{}; + for (fixture.ring.attribution.?.totals) |counts| try telemetry.add(attribution.Counts, &sum, counts); + try std.testing.expectEqual(fixture.ring.stats.copies, sum.copies); + try std.testing.expectEqual(fixture.ring.stats.dma_requested_bytes, sum.requested_bytes); + try std.testing.expectEqual(fixture.ring.stats.wasted, sum.unused_copies); + try std.testing.expectEqual(unused_bytes, sum.unused_requested_bytes); + } + + fn snapshot(fixture: *Fixture) telemetry.Snapshot { + const stats = fixture.ring.stats; + var experts: telemetry.Experts = .{ .ring_dma_copies = stats.copies, .ring_dma_requested_bytes = stats.dma_requested_bytes, .ring_unused_copies = stats.wasted }; + attribution.publish(fixture.ring.attribution.?.totals, &experts); + return .{ .rank = .{ .rank = 1, .world = 2 }, .scope = 1, .ring_owner = 1, .experts = experts }; + } +}; + +test "exl3: prefill attribution counts actual copies once across held retries full rings drop and slot reuse" { + const fixture = try Fixture.init(true); + defer fixture.deinit(); + const ring = fixture.ring; + try std.testing.expect(try ring.load(0, 0, .initial_prefetch)); + try std.testing.expect(try ring.load(1, 1, .cross_layer_prefetch)); + try fixture.checksum(0); + try std.testing.expectEqual(@as(usize, 15), Fixture.calls.payload_bytes); + const before_retry = ring.attribution.?.totals; + const calls_before_retry = Fixture.calls; + try std.testing.expect(!try ring.load(2, 2, .demand)); + try std.testing.expect(try ring.load(0, 0, .demand)); + try std.testing.expectEqualDeep(before_retry, ring.attribution.?.totals); + try std.testing.expectEqualDeep(calls_before_retry, Fixture.calls); + try std.testing.expectEqual(@as(i32, 0), ring.slot[0][0]); + try std.testing.expectEqual(@as(i32, 1), ring.slot[1][1]); + const used: [512]u32 = @splat(0); + ring.drop(0, &used); + ring.drop(0, &used); + try fixture.checksum(6); + try std.testing.expect(try ring.load(0, 0, .cross_layer_prefetch)); + try fixture.checksum(6); + try std.testing.expect(try ring.load(2, 2, .demand)); + try std.testing.expectEqual(@as(i32, -1), ring.slot[0][0]); + try std.testing.expectEqual(@as(i32, 1), ring.slot[1][1]); + try fixture.checksum(6); + try ring.newChunk(); + try fixture.checksum(27); + try ring.newChunk(); + try fixture.checksum(27); + const totals = ring.attribution.?.totals; + try std.testing.expectEqualDeep(attribution.Counts{ .copies = 1, .requested_bytes = 12, .unused_copies = 1, .unused_requested_bytes = 12 }, totals[0]); + try std.testing.expectEqualDeep(attribution.Counts{ .copies = 1, .requested_bytes = 6, .unused_copies = 1, .unused_requested_bytes = 6 }, totals[1]); + try std.testing.expectEqualDeep(attribution.Counts{ .copies = 1, .requested_bytes = 9, .unused_copies = 1, .unused_requested_bytes = 9 }, totals[2]); + try std.testing.expectEqualDeep(telemetry.Experts{ + .ring_dma_copies = 3, + .ring_dma_requested_bytes = 27, + .ring_unused_copies = 3, + .ring_demand_copies = 1, + .ring_demand_requested_bytes = 12, + .ring_demand_unused_copies = 1, + .ring_demand_unused_requested_bytes = 12, + .ring_initial_prefetch_copies = 1, + .ring_initial_prefetch_requested_bytes = 6, + .ring_initial_prefetch_unused_copies = 1, + .ring_initial_prefetch_unused_requested_bytes = 6, + .ring_cross_layer_prefetch_copies = 1, + .ring_cross_layer_prefetch_requested_bytes = 9, + .ring_cross_layer_prefetch_unused_copies = 1, + .ring_cross_layer_prefetch_unused_requested_bytes = 9, + }, fixture.snapshot().experts); + try std.testing.expectEqual(@as(usize, 0), Fixture.calls.syncs); +} + +test "exl3: prefill attribution retains outstanding provenance across snapshots releases and later chunk reset" { + const fixture = try Fixture.init(true); + defer fixture.deinit(); + const ring = fixture.ring; + try std.testing.expect(try ring.load(0, 0, .initial_prefetch)); + try std.testing.expect(try ring.load(0, 1, .cross_layer_prefetch)); + const before = fixture.snapshot(); + try ring.gate(ring.copy, 0, &.{ 0, 1 }); + try ring.release(ring.copy, 0, &.{0}); + const used: [512]u32 = @splat(0); + ring.drop(0, &used); + try fixture.checksum(9); + try std.testing.expect(try ring.load(0, 0, .demand)); + try ring.newChunk(); + try fixture.checksum(9); + const after = fixture.snapshot(); + const report = try telemetry.Report.between(before, after, .{ .phase = .prefill, .model_rows = 2048, .steps = 1 }); + try std.testing.expectEqual(@as(u64, 0), report.experts.ring_dma_copies); + try std.testing.expectEqual(@as(u64, 0), report.experts.ring_dma_requested_bytes); + try std.testing.expectEqual(@as(u64, 1), report.experts.ring_cross_layer_prefetch_unused_copies); + try std.testing.expectEqual(@as(u64, 9), report.experts.ring_cross_layer_prefetch_unused_requested_bytes); + try std.testing.expectEqual(@as(u64, 0), report.experts.ring_initial_prefetch_unused_copies); + const text = try report.formatAlloc(allocator); + defer allocator.free(text); + try std.testing.expect(std.mem.indexOf(u8, text, "ring_cross_layer_prefetch_unused_requested_bytes=9") != null); + try std.testing.expect(try ring.load(2, 2, .demand)); + try fixture.checksum(9); + try std.testing.expectEqual(@as(usize, 1), Fixture.calls.queries); + try std.testing.expectEqual(@as(usize, 2), Fixture.calls.waits); + try std.testing.expectEqual(@as(usize, 0), Fixture.calls.syncs); +} + +test "exl3: prefill attribution preserves prefetch ownership and first-copy source during cross-layer topups" { + const fixture = try Fixture.init(true); + defer fixture.deinit(); + const ring = fixture.ring; + ring.own_lo = 1; + ring.own_hi = 3; + ring.counts[4][0] = 100; + ring.counts[4][2] = 20; + ring.counts[4][1] = 10; + try std.testing.expectEqual(@as(usize, 1), try ring.prefetch(4, 1, 0, .initial_prefetch)); + try std.testing.expectEqual(@as(usize, 1), try ring.prefetch(4, 2, 0, .cross_layer_prefetch)); + try std.testing.expectEqual(@as(usize, 0), try ring.prefetch(4, 2, 0, .cross_layer_prefetch)); + try std.testing.expect(!ring.isLoaded(4, 0)); + try fixture.checksum(0); + try std.testing.expectEqual(@as(u64, 1), ring.attribution.?.totals[1].copies); + try std.testing.expectEqual(@as(u64, 1), ring.attribution.?.totals[2].copies); + try ring.newChunk(); + try fixture.checksum(21); +} + +test "exl3: disabled prefill attribution allocates nothing and bookkeeping adds no driver operations" { + var failing = std.testing.FailingAllocator.init(allocator, .{ .fail_index = 0 }); + try std.testing.expect((try attribution.State.init(failing.allocator(), 1000, false)) == null); + try std.testing.expectError(error.OutOfMemory, attribution.State.init(failing.allocator(), 1000, true)); + var previous_calls: Calls = undefined; + var previous_stats: ring_module.Stats = undefined; + var previous_image: u64 = 0; + for ([_]bool{ false, true }) |enabled| { + const fixture = try Fixture.init(enabled); + defer fixture.deinit(); + const ring = fixture.ring; + try std.testing.expect(try ring.load(0, 0, .initial_prefetch)); + try ring.gate(ring.copy, 0, &.{0}); + try ring.release(ring.copy, 0, &.{0}); + try std.testing.expect(try ring.load(1, 1, .cross_layer_prefetch)); + try ring.newChunk(); + const image = std.hash.Wyhash.hash(0, ring.shadow.bytes); + if (enabled) { + try std.testing.expectEqualDeep(previous_calls, Fixture.calls); + try std.testing.expectEqualDeep(previous_stats, ring.stats); + try std.testing.expectEqual(previous_image, image); + try fixture.checksum(9); + } else { + try std.testing.expect(ring.attribution == null); + previous_calls = Fixture.calls; + previous_stats = ring.stats; + previous_image = image; + } + } +} + +test "exl3: prefill attribution ignores failed transfers without suppressing retries" { + const fixture = try Fixture.init(true); + defer fixture.deinit(); + const ring = fixture.ring; + Fixture.copy_result = cuda.abi.error_not_ready; + try std.testing.expectError(error.NotReady, ring.load(0, 0, .initial_prefetch)); + try std.testing.expect(!ring.isLoaded(0, 0)); + try fixture.checksum(0); + Fixture.copy_result = cuda.abi.success; + try std.testing.expect(try ring.load(0, 0, .demand)); + try fixture.checksum(0); + try std.testing.expectEqual(@as(u64, 1), ring.attribution.?.totals[0].copies); + try std.testing.expectEqual(@as(u64, 0), ring.attribution.?.totals[1].copies); +} diff --git a/zig/src/families/flashnext/exl3/prefill/cuda_exl3_prefill_chunk.zig b/zig/src/families/flashnext/exl3/prefill/cuda_exl3_prefill_chunk.zig new file mode 100644 index 0000000000..632bf87cdd --- /dev/null +++ b/zig/src/families/flashnext/exl3/prefill/cuda_exl3_prefill_chunk.zig @@ -0,0 +1,390 @@ +//! Deep-prefill chunks: the sub-block views, writeback, the window stages and one chunk's layers and head. + +const std = @import("std"); +const cuda = @import("cuda"); +const batch = @import("../../cuda_batch.zig"); +const wt = @import("../../cuda_prompt_wt.zig"); +const ring_ = @import("cuda_exl3_prefill_ring.zig"); +const history_ = @import("cuda_exl3_prefill_history.zig"); +const window_ = @import("exl3_prefill_window.zig"); +const state = @import("cuda_exl3_prefill_state.zig"); +const experts = @import("cuda_exl3_prefill_moe.zig"); +const ringmoe = @import("cuda_exl3_prefill_ringmoe.zig"); +const tiles = @import("cuda_exl3_prefill_tiles.zig"); + +const dims = wt.dims; +const wide = wt.wide; +const streams = wt.streams; +const slots = wt.slots; +const eps = wt.eps; +const vocab = wt.vocab; +const out_sk = wt.out_sk; +const ple_layer = wt.ple_layer; +const ple_tail_n = wt.ple_tail_n; +const routed_n_ = wt.experts_n; +const att_rows = batch.att_rows; +const pss_n = batch.pss_n; + +const Deep = state.Deep; +const Options = state.Options; +const max_sub = state.max_sub; +const now = state.now; +const profile = state.profile; +const stats = &state.stats; +const upload = state.upload; +const moe = experts.moe; +const rankSum = experts.rankSum; +const keepFree = ringmoe.keepFree; +const ringMoe = ringmoe.ringMoe; + +/// A sub-block's view of the pass's scratch: the per-row buffers a layer keeps across its sub-blocks point into the +pub fn view(base: *const batch.Scratch, dp: *const Deep, r0: usize) batch.Scratch { + var v = base.*; + const l = dp.lay; + v.h = dp.dptr(l.h) + r0 * wide * 2; + v.pss = dp.dptr(l.pss) + r0 * pss_n * 4; + v.inj = dp.dptr(l.inj) + r0 * streams * 2; + v.mixed = dp.dptr(l.mixed) + r0 * dims * 2; + v.xsm = dp.dptr(l.xsm) + r0 * (dims / 32) * 4; + v.y = dp.dptr(l.y) + r0 * slots * dims * 4; + v.wts = dp.dptr(l.wts) + r0 * slots * 4; + v.picks = dp.dptr(l.picks) + r0 * slots * 4; + if (dp.world > 1) v.moe_f32 = dp.dptr(l.sum) + r0 * dims * 4; + v.ple_emb = dp.dptr(l.ple_emb) + r0 * dims * 2; + v.ple_xs = dp.dptr(l.ple_xs) + r0 * (dims / 32) * 4; + v.ids = dp.dptr(l.ids) + r0 * 4; + v.pos = dp.dptr(l.pos) + (r0 / att_rows) * 16; + return v; +} + +pub fn selectWriteback(dp: *Deep, reduced_branch: bool) !void { + if (reduced_branch and dp.unit_weights == null) { + const weights: []f32 = @alignCast(std.mem.bytesAsSlice(f32, dp.hbytes(dp.hl.ids, dp.rows * 4))); + @memset(weights, 1); + dp.unit_weights = try cuda.DeviceBuffer.fromHost(dp.dev.d, std.mem.sliceAsBytes(weights)); + } + dp.reduced_branch = reduced_branch; +} + +pub fn pendingWriteback(p: anytype, dp: *const Deep, layer: ?usize, r0: usize) void { + p.pending = layer; + p.pending_f32 = layer != null and dp.reduced_branch; + p.pending_f32_weights = if (p.pending_f32) dp.unit_weights.?.ptr + r0 * 4 else null; +} + +pub fn compactLast(p: anytype, reduced_branch: bool, row: usize) !void { + if (row == 0) return; + const s = p.s; + try p.copy(s.h, s.h + row * wide * 2, wide * 2); + try p.copy(s.pss, s.pss + row * pss_n * 4, pss_n * 4); + try p.copy(s.inj, s.inj + row * streams * 2, streams * 2); + if (reduced_branch) { + try p.copy(s.moe_f32, s.moe_f32 + row * dims * 4, dims * 4); + } else { + try p.copy(s.y, s.y + row * slots * dims * 4, slots * dims * 4); + try p.copy(s.wts, s.wts + row * slots * 4, slots * 4); + } +} + +const Gather = struct { + ngram: *const @import("../../cuda_ple.zig").Ngram, + history: []i64, + ids: []const u32, + emb: []u16, + xs: []f32, + failed: ?anyerror = null, + ns: i96 = 0, + io: std.Io, + + fn main(g: *Gather) void { + const t0 = now(g.io); + defer g.ns = now(g.io) - t0; + g.ngram.rows(g.history, g.ids, g.emb, g.xs) catch |err| { + g.failed = err; + }; + } +}; + +fn windowStage(dp: *Deep, p: anytype, ids: []const u32, pos0: usize) !void { + const n = ids.len; + try p.stream.synchronize(); + const id32 = std.mem.bytesAsSlice(i32, dp.hbytes(dp.hl.ids, n * 4)); + for (id32, ids) |*destination, id| destination.* = @intCast(id); + const blocks = (n + att_rows - 1) / att_rows; + const positions = std.mem.bytesAsSlice(i32, dp.hbytes(dp.hl.pos, blocks * 16)); + @memset(positions, 0); + for (0..blocks) |block| positions[4 * block] = @intCast(pos0 + block * att_rows); + try upload(p.stream, dp.dptr(dp.lay.ids), dp.hbytes(dp.hl.ids, n * 4)); + try upload(p.stream, dp.dptr(dp.lay.pos), dp.hbytes(dp.hl.pos, blocks * 16)); + if (p.m.table_dense) { + try p.tri.embedDense(dp.dptr(dp.lay.ids), p.m.table_w, dp.dptr(dp.lay.h), dims, streams, n); + } else { + try p.tri.embed(dp.dptr(dp.lay.ids), p.m.table_w, p.m.table_s, p.m.table_b, dp.dptr(dp.lay.h), dims, streams, n); + } +} + +pub fn windowInputs(dp: *Deep, p: anytype, ids: []const u32, pos0: usize, layer: usize) !void { + const blocks = (ids.len + att_rows - 1) / att_rows; + const positions_offset = dp.hl.pos + (layer + 1) * dp.hl.pos_stride; + const positions: []i32 = @alignCast(std.mem.bytesAsSlice(i32, dp.hbytes(positions_offset, blocks * 16))); + @memset(positions, 0); + for (0..blocks) |block| positions[4 * block] = @intCast(pos0 + block * att_rows); + try upload(p.stream, dp.dptr(dp.lay.pos), dp.hbytes(positions_offset, blocks * 16)); + if (layer == ple_layer) { + const embedding: []u16 = @alignCast(std.mem.bytesAsSlice(u16, dp.hbytes(dp.hl.emb, ids.len * dims * 2))); + const sums: []f32 = @alignCast(std.mem.bytesAsSlice(f32, dp.hbytes(dp.hl.xs, ids.len * (dims / 32) * 4))); + try p.m.ngram.rows(p.seq.history[0..], ids, embedding, sums); + try upload(p.stream, dp.dptr(dp.lay.ple_emb), std.mem.sliceAsBytes(embedding)); + try upload(p.stream, dp.dptr(dp.lay.ple_xs), std.mem.sliceAsBytes(sums)); + } +} + +fn windowDense(comptime Tri: type, dp: *Deep, p: *batch.Pass(Tri), base: *batch.Scratch, layer: usize, n: usize, pos0: usize, sub: usize, waits: []const cuda.Event) !void { + const lw = &p.m.layers[layer]; + const n_sub = (n + sub - 1) / sub; + var views: [max_sub]batch.Scratch = undefined; + for (0..n_sub) |block| views[block] = view(base, dp, block * sub); + defer p.s = base; + p.s = base; + p.rows = n; + p.pos0 = pos0; + for (0..n_sub) |block| { + const at = block * sub; + const count = @min(sub, n - at); + p.s = &views[block]; + p.rows = count; + p.pos0 = pos0 + at; + pendingWriteback(p, dp, if (layer == 0) null else layer - 1, at); + if (block < waits.len) try p.stream.wait(waits[block]); + try layerBlock(Tri, p, lw, layer, count); + } + dp.red_n = 0; +} + +fn windowFinish(comptime Tri: type, dp: *Deep, p: *batch.Pass(Tri), base: *batch.Scratch, n: usize, sub: usize, logits: batch.Logits, streams_out: u64, logits_host: ?[]u8) !void { + const n_sub = (n + sub - 1) / sub; + var views: [max_sub]batch.Scratch = undefined; + for (0..n_sub) |block| views[block] = view(base, dp, block * sub); + defer p.s = base; + stats.chunks += 1; + stats.rows += n; + const last_block = n_sub - 1; + const last_rows = n - last_block * sub; + if (streams_out != 0) { + for (0..n_sub) |block| { + const count = @min(sub, n - block * sub); + p.s = &views[block]; + p.rows = count; + pendingWriteback(p, dp, wt.layers_n - 1, block * sub); + try p.writeback(0, count); + try p.copy(streams_out + block * sub * wide * 2, views[block].h, count * wide * 2); + } + } + if (logits == .none) return; + p.s = &views[last_block]; + p.rows = last_rows; + const scratch = p.s; + const row = last_rows - 1; + if (streams_out != 0) { + if (row > 0) { + try p.copy(scratch.h, scratch.h + row * wide * 2, wide * 2); + try p.copy(scratch.pss, scratch.pss + row * pss_n * 4, pss_n * 4); + } + } else { + try compactLast(p, dp.reduced_branch, row); + pendingWriteback(p, dp, wt.layers_n - 1, last_block * sub); + try p.writeback(0, 1); + } + try p.readout(p.m.mixer, 0, 1); + if (p.m.head.n == vocab) { + _ = try p.lane(scratch.mixed, scratch.xsm, p.m.head, scratch.logits, scratch.dummy, 1, 1); + } else { + const collective = p.m.collective orelse return error.TensorParallelUnavailable; + _ = try p.lane(scratch.mixed, scratch.xsm, p.m.head, scratch.head_local, scratch.dummy, 1, 1); + try collective.gather(p.stream, scratch.head_local, scratch.logits, vocab / 2); + } + if (logits_host) |destination| { + if (p.m.collective) |collective| try collective.finish(p.stream); + try p.stream.synchronize(); + try p.stream.d.check(p.stream.d.api.cuMemcpyDtoH_v2(destination.ptr, scratch.logits, vocab * 2), "cuMemcpyDtoH"); + } +} + +pub fn twoChunk(comptime Tri: type, dp: *Deep, p: *batch.Pass(Tri), base: *batch.Scratch, ids: []const u32, pos0: usize, sub: usize, hooks: batch.Hooks, output_at: usize, last: bool, logits_host: []u8, o: Options, ring: *ring_.Ring, plan_: window_.Plan) !void { + const buffers = [_]*Deep{ dp, dp.second.? }; + for (plan_.segments[0..plan_.count]) |segment| { + if ((segment.rows + sub - 1) / sub > max_sub) return error.UnexpectedTensor; + } + try p.stream.synchronize(); + try ring.newChunk(); + history_.prepare(&dp.history, &dp.history_shape, .{ .sub = sub, .lo = ring.own_lo, .hi = ring.own_hi }); + for (0..@min(o.ring_ahead + 1, wt.layers_n)) |layer| _ = try ring.prefetch(layer, @min(routed_n_, dp.used[layer]), keepFree(ring, o), .initial_prefetch); + for (plan_.segments[0..plan_.count], 0..) |segment, index| { + try windowStage(buffers[index], p, ids[segment.offset..][0..segment.rows], pos0 + segment.offset); + } + var schedule = plan_.iterator(p.m.layers.len); + while (schedule.next()) |step| { + const active = buffers[step.chunk]; + const segment = plan_.segments[step.chunk]; + try windowInputs(active, p, ids[segment.offset..][0..segment.rows], pos0 + segment.offset, step.layer); + try windowDense(Tri, active, p, base, step.layer, segment.rows, pos0 + segment.offset, sub, active.red_ev[0..active.red_n]); + if (active != dp) { + active.used = dp.used; + active.history = dp.history; + } + const layer = &p.m.layers[step.layer]; + try ringMoe(Tri, active, ring, p, layer, step.layer, segment.rows, o); + if (active != dp) { + dp.used = active.used; + dp.history = active.history; + } + try rankSum(active, p, layer.exl3.?, segment.rows, sub); + if (step.layer + 1 == p.m.layers.len) { + for (active.red_ev[0..active.red_n]) |event| try p.stream.wait(event); + active.red_n = 0; + } + } + for (plan_.segments[0..plan_.count], 0..) |segment, index| { + const final = last and index + 1 == plan_.count; + const output = if (hooks.streams == 0) 0 else hooks.streams + (output_at + segment.offset) * wide * 2; + try windowFinish(Tri, buffers[index], p, base, segment.rows, sub, if (final) hooks.logits else .none, output, if (final and hooks.logits != .none) logits_host else null); + } +} + +pub fn chunk(comptime Tri: type, dp: *Deep, p: *batch.Pass(Tri), base: *batch.Scratch, ids: []const u32, pos0: usize, sub: usize, logits: batch.Logits, streams_out: u64, logits_host: ?[]u8, o: Options, rg: ?*ring_.Ring) !void { + const n = ids.len; + const st = p.stream; + const hl = dp.hl; + // the previous chunk's uploads read the pinned rows: they landed before its head (a synchronize), or below + try st.synchronize(); + if (rg) |r| { + try r.newChunk(); + history_.prepare(&dp.history, &dp.history_shape, .{ .sub = sub, .lo = r.own_lo, .hi = r.own_hi }); + // the first layers' experts start streaming now, under the embedding and layer 0's blocks + for (0..@min(o.ring_ahead + 1, wt.layers_n)) |l| _ = try r.prefetch(l, @min(routed_n_, dp.used[l]), keepFree(r, o), .initial_prefetch); + } + const id32 = std.mem.bytesAsSlice(i32, dp.hbytes(hl.ids, n * 4)); + for (id32, ids) |*x, v| x.* = @intCast(v); + const blocks = (n + att_rows - 1) / att_rows; + const pos = std.mem.bytesAsSlice(i32, dp.hbytes(hl.pos, blocks * 16)); + @memset(pos, 0); + for (0..blocks) |b| pos[4 * b] = @intCast(pos0 + b * att_rows); + try upload(st, dp.dptr(dp.lay.ids), dp.hbytes(hl.ids, n * 4)); + try upload(st, dp.dptr(dp.lay.pos), dp.hbytes(hl.pos, blocks * 16)); + const emb: []u16 = @alignCast(std.mem.bytesAsSlice(u16, dp.hbytes(hl.emb, n * dims * 2))); + const xs: []f32 = @alignCast(std.mem.bytesAsSlice(f32, dp.hbytes(hl.xs, n * (dims / 32) * 4))); + var gather: Gather = .{ .ngram = p.m.ngram, .history = p.seq.history[0..], .ids = ids, .emb = emb, .xs = xs, .io = p.io }; + var thread: ?std.Thread = try std.Thread.spawn(.{}, Gather.main, .{&gather}); + defer if (thread) |th| th.join(); + const t0 = now(p.io); + defer if (profile()) { + st.synchronize() catch {}; + std.debug.print("exl3 prefill: {d} rows at {d} in {d:.1} ms (n-gram gather {d:.1} ms on the host)\n", .{ n, pos0, @as(f64, @floatFromInt(now(p.io) - t0)) / 1e6, @as(f64, @floatFromInt(gather.ns)) / 1e6 }); + }; + p.rows = n; + p.pos0 = pos0; + pendingWriteback(p, dp, null, 0); + if (p.m.table_dense) { + try p.tri.embedDense(dp.dptr(dp.lay.ids), p.m.table_w, dp.dptr(dp.lay.h), dims, streams, n); + } else { + try p.tri.embed(dp.dptr(dp.lay.ids), p.m.table_w, p.m.table_s, p.m.table_b, dp.dptr(dp.lay.h), dims, streams, n); + } + const n_sub = (n + sub - 1) / sub; + var views: [max_sub]batch.Scratch = undefined; + if (n_sub > views.len) return error.UnexpectedTensor; + for (0..n_sub) |b| views[b] = view(base, dp, b * sub); + defer p.s = base; + for (p.m.layers, 0..) |*lw, layer| { + if (layer == ple_layer) { + if (thread) |th| th.join(); + thread = null; + if (gather.failed) |err| return err; + try upload(st, dp.dptr(dp.lay.ple_emb), std.mem.sliceAsBytes(emb)); + try upload(st, dp.dptr(dp.lay.ple_xs), std.mem.sliceAsBytes(xs)); + } + for (0..n_sub) |b| { + const m = @min(sub, n - b * sub); + p.s = &views[b]; + p.rows = m; + p.pos0 = pos0 + b * sub; + pendingWriteback(p, dp, if (layer > 0) layer - 1 else null, b * sub); + if (b < dp.red_n) try st.wait(dp.red_ev[b]); + try layerBlock(Tri, p, lw, layer, m); + } + dp.red_n = 0; + if (rg) |r| try ringMoe(Tri, dp, r, p, lw, layer, n, o) else try moe(Tri, dp, p, lw, layer, n, o); + try rankSum(dp, p, lw.exl3.?, n, sub); + } + for (dp.red_ev[0..dp.red_n]) |ev| try st.wait(ev); + dp.red_n = 0; + stats.chunks += 1; + stats.rows += n; + // the last layer's experts written back where the caller keeps the streams, then the head on the last row + const last_b = n_sub - 1; + const m_last = n - last_b * sub; + if (streams_out != 0) { + for (0..n_sub) |b| { + const m = @min(sub, n - b * sub); + p.s = &views[b]; + p.rows = m; + pendingWriteback(p, dp, wt.layers_n - 1, b * sub); + try p.writeback(0, m); + try p.copy(streams_out + b * sub * wide * 2, views[b].h, m * wide * 2); + } + } + if (logits == .none) return; + p.s = &views[last_b]; + p.rows = m_last; + const s = p.s; + const r = m_last - 1; + if (streams_out != 0) { + if (r > 0) { + try p.copy(s.h, s.h + r * wide * 2, wide * 2); + try p.copy(s.pss, s.pss + r * pss_n * 4, pss_n * 4); + } + } else { + try compactLast(p, dp.reduced_branch, r); + pendingWriteback(p, dp, wt.layers_n - 1, last_b * sub); + try p.writeback(0, 1); + } + try p.readout(p.m.mixer, 0, 1); + if (p.m.head.n == vocab) { + _ = try p.lane(s.mixed, s.xsm, p.m.head, s.logits, s.dummy, 1, 1); + } else { + // the head split by vocabulary (expert parallel): each rank's half, gathered in vocabulary order (projectHead) + const c = p.m.collective orelse return error.TensorParallelUnavailable; + _ = try p.lane(s.mixed, s.xsm, p.m.head, s.head_local, s.dummy, 1, 1); + try c.gather(st, s.head_local, s.logits, vocab / 2); + } + if (logits_host) |lh| { + if (p.m.collective) |c| try c.finish(st); + try st.synchronize(); + try st.d.check(st.d.api.cuMemcpyDtoH_v2(lh.ptr, s.logits, vocab * 2), "cuMemcpyDtoH"); + } +} + +/// One sub-block of one layer, as Pass.chunk runs it, up to the router's picks (the experts come after every +fn layerBlock(comptime Tri: type, p: *batch.Pass(Tri), lw: *const @import("../../cuda_prompt_load.zig").Layer, layer: usize, m: usize) !void { + const s = p.s; + if (layer == ple_layer) { + if (p.pending != null) try p.writeback(0, m); + // Pass.pleBlock for a prompt's rows (no lanes) + _ = try p.lane(s.ple_emb, s.ple_xs, lw.key, s.ple_keys, s.dummy, m, 1); + _ = try p.lane(s.ple_emb, s.ple_xs, lw.val, s.dummy, s.ple_part, m, 4); + try p.sumSlices(s.ple_part, s.ple_vals, true, m * dims, 4); + try p.tri.pleGate(s.ple_keys, s.ple_vals, s.h, lw.nk.ptr, lw.nq.ptr, s.ple_gated, s.ple_pss, eps, m, dims, streams); + try p.tri.pleConv(s.ple_gated, s.ple_pss, lw.nc.ptr, p.seq.tail.ptr, lw.ple_cw.ptr, s.h, s.h, s.ple_nrow, eps, m, dims, streams, 4, 3); + try p.shiftWindow(p.seq.tail.ptr, s.ple_nrow, wide, m, ple_tail_n, wide); + } + try p.writeback(0, m); + try p.readout(lw.attn_hc, 0, m); + if (lw.linear) { + try p.gdnBlock(lw, layer, wt.linearIndex(layer), .{}); + } else { + try p.attnBlock(lw, wt.attnIndex(layer)); + } + try p.tri.hcWriteback(s.h, s.h, s.pss, s.part, "*fp32", s.inj, s.dummy, "*bf16", s.dummy, "*bf16", m * dims, m, dims, streams, 4, 1, 1, out_sk); + try p.readout(lw.mlp_hc, 0, m); + try p.tri.routerMany(s.mixed, lw.router.ptr, s.router_l, dims, m); + try p.tri.topkMany(s.router_l, s.picks, s.wts, m); +} diff --git a/zig/src/families/flashnext/exl3/prefill/cuda_exl3_prefill_head.zig b/zig/src/families/flashnext/exl3/prefill/cuda_exl3_prefill_head.zig new file mode 100644 index 0000000000..128bac0788 --- /dev/null +++ b/zig/src/families/flashnext/exl3/prefill/cuda_exl3_prefill_head.zig @@ -0,0 +1,149 @@ +//! The MTP head run expert-major over deep chunks (the hook cuda_mtp_ops.zig calls). + +const std = @import("std"); +const batch = @import("../../cuda_batch.zig"); +const wt = @import("../../cuda_prompt_wt.zig"); +const state = @import("cuda_exl3_prefill_state.zig"); +const chunks = @import("cuda_exl3_prefill_chunk.zig"); +const experts = @import("cuda_exl3_prefill_moe.zig"); +const tiles = @import("cuda_exl3_prefill_tiles.zig"); + +const dims = wt.dims; +const wide = wt.wide; +const streams = wt.streams; +const eps = wt.eps; +const out_sk = wt.out_sk; +const att_rows = batch.att_rows; + +const max_sub = state.max_sub; +const opts = state.opts; +const stats = &state.stats; +const view = chunks.view; +const moe = experts.moe; +const warmTiles = tiles.warmTiles; + +/// The hook in cuda_mtp_ops.zig `forwardHead` (a lane's prompt chunk through the MTP head, after the main pass kept +pub fn head(comptime Tri: type, o: anytype, d: *@import("../../cuda_mtp.zig").Draft, next: []const u32, rows: u64, scores: bool) !bool { + const op = opts(); + const n = next.len; + // A paired window hands the head up to two deep chunks of rows: it runs them as two slices (headSplit). + if (!scores and state.deep != null and n > state.deep.?.rows) { + const first = headSplit(n, state.deep.?.rows, op.min_rows) orelse return false; + if (!try headSlice(Tri, o, d, next[0..first], rows, false)) return false; + if (!try headSlice(Tri, o, d, next[first..], rows + first * wide * 2, false)) return error.UnexpectedTensor; + return true; + } + return headSlice(Tri, o, d, next, rows, scores); +} + +/// Where a head call of `n` rows over a paired window splits into two deep-sized slices, or null when it cannot. +fn headSplit(n: usize, deep_rows: usize, min_rows: usize) ?usize { + if (n <= deep_rows or n > 2 * deep_rows) return null; + const first = ((n + 1) / 2 + att_rows - 1) / att_rows * att_rows; + if (first >= n or first > deep_rows or n - first < min_rows) return null; + return first; +} + +test "exl3: a paired window's head splits into two deep slices that cannot decline after the first runs" { + try std.testing.expectEqual(@as(?usize, 4096), headSplit(8192, 4096, 512)); + try std.testing.expectEqual(@as(?usize, null), headSplit(4096, 4096, 512)); + try std.testing.expectEqual(@as(?usize, null), headSplit(8193, 4096, 512)); + var n: usize = 4097; + while (n <= 8192) : (n += 61) { + const first = headSplit(n, 4096, 512) orelse continue; + try std.testing.expect(first <= 4096 and first % att_rows == 0 and n - first <= first and n - first >= 512); + } +} + +fn headSlice(comptime Tri: type, o: anytype, d: *@import("../../cuda_mtp.zig").Draft, next: []const u32, rows: u64, scores: bool) !bool { + const op = opts(); + const n = next.len; + if (!op.enabled or !headEnabled() or scores or n < op.min_rows or d.media != null) return false; + const e = o.e; + const dp = state.deep orelse return false; + if (dp.owner != @as(*const anyopaque, e) or dp.rows < n) return false; + const h = o.h; + const x = h.layer.exl3 orelse return false; + if (x.tp != null or x.pool == null or (x.layer.act_mode != 1 and x.layer.act_mode != 0)) return false; + if (d.len + n > d.cap) return error.PromptTooLong; + const mtp = @import("../../cuda_mtp.zig"); + const tri = o.tri; + const want = @max(@as(usize, 32), 2 * @min(n, mtp.piece_rows)); + const s = try e.reserve(want); + const piece = @max(@as(usize, 1), @min(mtp.piece_rows, s.cap / 2)) / att_rows * att_rows; + if (piece == 0) return false; + const n_pieces = (n + piece - 1) / piece; + if (n_pieces > max_sub) return false; + const score_rows = try e.reserveScores(d.cap, piece); + var p: batch.Pass(Tri) = .{ + .gpa = e.gpa, + .io = e.io, + .stream = e.stream.*, + .tri = tri, + .k = &e.batch_k, + .s = s, + .m = .{ .table_w = e.table.w.ptr, .table_s = e.table.s.ptr, .table_b = e.table.b.ptr, .table_dense = e.table.dense, .layers = (&h.layer)[0..1], .mixer = h.mixer, .head = h.draft, .ngram = &e.ngram }, + .seq = .{ .media = null, .conv = &.{}, .state = &.{}, .kc = (&d.kc)[0..1], .vc = (&d.vc)[0..1], .ikc = (&d.ikc)[0..1], .pooled = (&d.pooled)[0..1], .tail = d.kc, .history = &d.unused_history, .cap = d.cap, .scores = score_rows }, + .hooks = .{ .logits = .none }, + }; + var views: [max_sub]batch.Scratch = undefined; + for (0..n_pieces) |b| views[b] = view(s, dp, b * piece); + defer p.s = s; + if (op.tiles and !op.slots_kernel) { + try warmTiles(e); + tiles.tiles_on = true; + } + defer tiles.tiles_on = false; + const lw = &h.layer; + for (0..n_pieces) |b| { + const at = b * piece; + const m = @min(piece, n - at); + p.s = &views[b]; + const v = p.s; + p.rows = m; + p.pos0 = d.len + at; + p.pending = null; + try p.stage(next[at..][0..m]); + // forwardHead's piece up to the router (the embedding branch, the hidden branch). + if (e.table.dense) { + try tri.embedDense(v.ids, e.table.w.ptr, v.ple_keys, dims, streams, m); + } else { + try tri.embed(v.ids, e.table.w.ptr, e.table.s.ptr, e.table.b.ptr, v.ple_keys, dims, streams, m); + } + try tri.rmsnorm(v.ple_keys, h.norm_e.ptr, v.mixed, v.xsm, eps, wide, m, dims, dims); + _ = try p.lane(v.mixed, v.xsm, h.fc_e, v.dummy, v.part, m, mtp.fc_sk); + try p.sumSlices(v.part, v.branch, true, m * dims, mtp.fc_sk); + try tri.rmsnorm(rows + at * wide * 2, h.norm_h.ptr, v.norm, v.gq, eps, wide, m, wide, wide); + _ = try p.lane(v.norm, v.gq, h.fc_h, v.dummy, v.part, streams * m, mtp.fc_sk); + try p.sumSlices(v.part, v.ple_gated, true, streams * m * dims, mtp.fc_sk); + try tri.addStreams(v.branch, v.ple_gated, v.h, m, dims, streams); + try p.writeback(0, m); + try p.readout(lw.attn_hc, 0, m); + try p.attnBlock(lw, 0); + try tri.hcWriteback(v.h, v.h, v.pss, v.part, "*fp32", v.inj, v.dummy, "*bf16", v.dummy, "*bf16", m * dims, m, dims, streams, 4, 1, 1, out_sk); + try p.readout(lw.mlp_hc, 0, m); + try tri.routerMany(v.mixed, lw.router.ptr, v.router_l, dims, m); + try tri.topkMany(v.router_l, v.picks, v.wts, m); + } + // the head's experts once over the chunk (its own pool, passes of half its lines) + try moe(Tri, dp, &p, lw, 0, n, op); + for (0..n_pieces) |b| { + const m = @min(piece, n - b * piece); + p.s = &views[b]; + p.rows = m; + p.pending = 0; + try p.writeback(0, m); + } + const last = &views[n_pieces - 1]; + try p.copy(d.out.ptr, last.h + (n - 1 - (n_pieces - 1) * piece) * wide * 2, wide * 2); + d.len += n; + stats.head_rows += n; + if (!head_logged) std.log.info("flashnext exl3 prefill: the MTP head runs expert-major over deep chunks ({d} rows first)", .{n}); + head_logged = true; + return true; +} +var head_logged = false; + +fn headEnabled() bool { + return true; +} diff --git a/zig/src/families/flashnext/exl3/prefill/cuda_exl3_prefill_history.zig b/zig/src/families/flashnext/exl3/prefill/cuda_exl3_prefill_history.zig new file mode 100644 index 0000000000..b762bdc296 --- /dev/null +++ b/zig/src/families/flashnext/exl3/prefill/cuda_exl3_prefill_history.zig @@ -0,0 +1,63 @@ +const std = @import("std"); + +pub const Slot = struct { + rows: usize = 0, + used: usize = 0, +}; + +pub const Shape = struct { sub: usize, lo: usize, hi: usize }; + +pub fn prepare(history: []Slot, previous: *?Shape, shape: Shape) void { + if (previous.* == null or !std.meta.eql(previous.*.?, shape)) @memset(history, .{}); + previous.* = shape; +} + +pub fn want(history: Slot, rows: usize, current_used: usize, cap: usize) usize { + const used = if (rows != 0 and history.rows == rows) history.used else current_used; + return @min(cap, used +| 16); +} + +test "exl3: prefetch history unseen and different rows retain the current-layer fallback" { + try std.testing.expectEqual(@as(usize, 416), want(.{}, 4096, 400, 512)); + try std.testing.expectEqual(@as(usize, 416), want(.{ .rows = 2048, .used = 50 }, 4096, 400, 512)); + try std.testing.expectEqual(@as(usize, 416), want(.{}, 0, 400, 512)); +} + +test "exl3: prefetch history uses the target layer's own unique demand including valid zero" { + var history: [4]Slot = @splat(.{}); + history[2] = .{ .rows = 4096, .used = 80 }; + history[0] = .{ .rows = 4096, .used = 300 }; + try std.testing.expectEqual(@as(usize, 316), want(history[1], 4096, 300, 512)); + try std.testing.expectEqual(@as(usize, 96), want(history[2], 4096, 300, 512)); + history[1] = .{ .rows = 4096, .used = 200 }; + try std.testing.expectEqual(@as(usize, 96), want(history[2], 4096, 200, 512)); + history[2] = .{ .rows = 4096, .used = 0 }; + try std.testing.expectEqual(@as(usize, 16), want(history[2], 4096, 200, 512)); + try std.testing.expectEqual(@as(usize, 216), want(history[2], 2048, 200, 512)); +} + +test "exl3: prefetch history caps both paths without overflow" { + for ([_]usize{ 0, 1, 16, 256, 512, std.math.maxInt(usize) }) |cap| { + try std.testing.expectEqual(cap, want(.{}, 4096, std.math.maxInt(usize), cap)); + try std.testing.expectEqual(cap, want(.{ .rows = 4096, .used = std.math.maxInt(usize) }, 4096, 0, cap)); + } + for (0..513) |used| { + try std.testing.expectEqual(@min(@as(usize, 512), used + 16), want(.{}, 4096, used, 512)); + } +} + +test "exl3: prefetch history resets for sub-block or rank shape changes" { + const shape: Shape = .{ .sub = 256, .lo = 0, .hi = 256 }; + var previous: ?Shape = null; + var history: [4]Slot = @splat(.{ .rows = 4096, .used = 80 }); + prepare(&history, &previous, shape); + for (history) |entry| try std.testing.expectEqual(@as(usize, 416), want(entry, 4096, 400, 512)); + history[2] = .{ .rows = 4096, .used = 80 }; + prepare(&history, &previous, shape); + try std.testing.expectEqual(@as(usize, 96), want(history[2], 4096, 400, 512)); + prepare(&history, &previous, .{ .sub = 512, .lo = 0, .hi = 256 }); + try std.testing.expectEqual(@as(usize, 416), want(history[2], 4096, 400, 512)); + history[2] = .{ .rows = 4096, .used = 80 }; + prepare(&history, &previous, .{ .sub = 512, .lo = 256, .hi = 512 }); + try std.testing.expectEqual(@as(usize, 416), want(history[2], 4096, 400, 512)); +} diff --git a/zig/src/families/flashnext/exl3/prefill/cuda_exl3_prefill_k.zig b/zig/src/families/flashnext/exl3/prefill/cuda_exl3_prefill_k.zig new file mode 100644 index 0000000000..687ad12c91 --- /dev/null +++ b/zig/src/families/flashnext/exl3/prefill/cuda_exl3_prefill_k.zig @@ -0,0 +1,243 @@ +//! The Flash Next EXL3 prefill experts (flashnext_exl3_prefill.cu): transposed prompt tiles in decode's bytes. +const std = @import("std"); +const cuda = @import("cuda"); + +const exl3e = cuda.exl3_experts; + +pub const tile: usize = 64; + +pub const Kernels = struct { + mod: cuda.Module, + rotate: cuda.Function, + gu: cuda.Function, + down: cuda.Function, + /// Expert parallel: a rank's combined branch over its own slots, and the ranks' sum placed for the write-back. + tp_combine: cuda.Function, + tp_finish: cuda.Function, + tp_place: cuda.Function, + + pub fn load(d: *const cuda.Driver) !Kernels { + if (!cuda.kernels.available) return error.BuiltWithoutKernels; + var mod = try cuda.Module.load(d, cuda.kernels.flashnext_exl3_prefill); + errdefer mod.unload(); + return .{ .mod = mod, .rotate = try mod.function("tf_exl3pf_rotate"), .gu = try mod.function("tf_exl3pf_gu"), .down = try mod.function("tf_exl3pf_down"), .tp_combine = try mod.function("tf_exl3pf_tp_combine"), .tp_finish = try mod.function("tf_exl3pf_tp_finish"), .tp_place = try mod.function("tf_exl3pf_tp_place") }; + } + + pub fn unload(k: *Kernels) void { + k.mod.unload(); + } +}; + +/// Members and items of one pass, host side (int32, the layouts the kernels read). +pub const Plan = struct { + members: []i32, + items: []i32, // 3 a item + n_members: usize = 0, + n_items: usize = 0, + + /// Room for `pairs` members and their items over `experts` table entries. + pub fn init(gpa: std.mem.Allocator, pairs: usize, experts: usize) !Plan { + const members = try gpa.alloc(i32, pairs); + errdefer gpa.free(members); + return .{ .members = members, .items = try gpa.alloc(i32, 3 * maxItems(pairs, experts)) }; + } + + pub fn deinit(p: *Plan, gpa: std.mem.Allocator) void { + gpa.free(p.members); + gpa.free(p.items); + p.* = undefined; + } + + pub fn itemBytes(p: *const Plan) []const u8 { + return std.mem.sliceAsBytes(p.items[0 .. 3 * p.n_items]); + } + + pub fn memberBytes(p: *const Plan) []const u8 { + return std.mem.sliceAsBytes(p.members[0..p.n_members]); + } +}; + +pub fn maxItems(pairs: usize, experts: usize) usize { + return @min(pairs, experts) + pairs / tile; +} + +/// Counting sort of `picks` (int32 [pairs], -1 or ids outside the table skipped) by expert, keeping only experts with +pub fn build(picks: []const i32, take: ?[]const bool, counts: []u32, p: *Plan) void { + @memset(counts, 0); + const e_n = counts.len; + for (picks) |pk| { + if (pk < 0 or @as(usize, @intCast(pk)) >= e_n) continue; + const id: usize = @intCast(pk); + if (take) |t| if (!t[id]) continue; + counts[id] += 1; + } + // exclusive prefix into counts (the next free member slot of each expert), items as runs of `tile` + var at: u32 = 0; + var items: usize = 0; + for (counts, 0..) |*c, id| { + const n = c.*; + c.* = at; + var off: u32 = 0; + while (off < n) : (off += tile) { + p.items[3 * items] = @intCast(id); + p.items[3 * items + 1] = @intCast(at + off); + p.items[3 * items + 2] = @intCast(@min(@as(u32, tile), n - off)); + items += 1; + } + at += n; + } + for (picks, 0..) |pk, i| { + if (pk < 0 or @as(usize, @intCast(pk)) >= e_n) continue; + const id: usize = @intCast(pk); + if (take) |t| if (!t[id]) continue; + p.members[counts[id]] = @intCast(i); + counts[id] += 1; + } + p.n_members = at; + p.n_items = items; +} + +/// Device scratch a pass of `members` members needs: xg, xu [members, dims] and xd [members, width] fp16, 256-aligned. +pub fn scratchBytes(members: usize, dims: usize, width: usize) usize { + return 2 * align256(members * dims * 2) + align256(members * width * 2); +} + +fn align256(n: usize) usize { + return (n + 255) / 256 * 256; +} + +pub const Device = struct { + /// Device copies of the plan (members, items) and the compact scratch base (scratchBytes of the plan's members). + members: u64, + items: u64, + scratch: u64, +}; + +/// True when the tiles cut K as the grouped kernels do (plan's splits of four warps, chains of whole 32-input steps). +pub fn sumsAsDecode(d: usize, w: usize) bool { + const g_sk = exl3e.plan(d, w, true) catch return false; + const d_sk = exl3e.plan(w, d, false) catch return false; + return d_sk == 1 and (d / 16 / g_sk / 4) % 2 == 0 and (w / 16 / 4) % 2 == 0; +} + +/// One pass: rotate every member's input, gate/up + SwiGLU, down into y [pairs, dims] fp32 (y[p] for member p). +pub fn run(k: *const Kernels, s: cuda.Stream, l: exl3e.Layer, x: u64, stride: usize, picks: u64, slots: usize, n_members: usize, n_items: usize, dv: Device, y: u64) !void { + if (n_members == 0) return; + const d = l.dims; + const w = l.width; + if (d % 256 != 0 or w % 128 != 0 or (l.act_mode != 1 and l.act_mode != 0)) return error.Exl3Shape; + if (!sumsAsDecode(d, w)) return error.Exl3Shape; + const xg = dv.scratch; + const xu = xg + align256(n_members * d * 2); + const xd = xu + align256(n_members * d * 2); + var r: cuda.Args = .{}; + r.add(x); + r.add(@as(i32, @intCast(stride))); + r.add(picks); + r.add(dv.members); + r.add(@as(i32, @intCast(slots))); + r.add(l.suh_g); + r.add(l.suh_u); + r.add(xg); + r.add(xu); + r.add(@as(i32, @intCast(d))); + try cuda.launch.launch(k.rotate, .{ .grid = .{ .x = @intCast(n_members), .y = @intCast(d / 128), .z = 2 }, .block = .{ .x = 32 } }, s, &r); + var g: cuda.Args = .{}; + g.add(xg); + g.add(xu); + g.add(l.gate); + g.add(l.up); + g.add(l.gate_k2); + g.add(dv.items); + g.add(l.svh_g); + g.add(l.svh_u); + g.add(l.suh_d); + g.add(xd); + g.add(@as(i32, @intCast(d))); + g.add(@as(i32, @intCast(w))); + g.add(l.limit); + g.add(l.act_mode); + try cuda.launch.launch(k.gu, .{ .grid = .{ .x = @intCast(w / 128), .y = @intCast(n_items) }, .block = .{ .x = 256 } }, s, &g); + var dn: cuda.Args = .{}; + dn.add(xd); + dn.add(l.down); + dn.add(l.down_k2); + dn.add(dv.items); + dn.add(dv.members); + dn.add(l.svh_d); + dn.add(y); + dn.add(@as(i32, @intCast(w))); + dn.add(@as(i32, @intCast(d))); + try cuda.launch.launch(k.down, .{ .grid = .{ .x = @intCast(d / 256), .y = @intCast(n_items) }, .block = .{ .x = 256 } }, s, &dn); +} + +/// The canonical reduction's block partials of one rank: `sub` rows per gather region, `m` partials per rank. +pub fn tpCombine(k: *const Kernels, s: cuda.Stream, y: u64, wts: u64, picks: u64, out: u64, rows: usize, slots: usize, d: usize, sub: usize, world: usize, rank: usize, m: usize) !void { + if (rows == 0) return; + var a: cuda.Args = .{}; + a.add(y); + a.add(wts); + a.add(picks); + a.add(out); + a.add(@as(i32, @intCast(slots))); + a.add(@as(i32, @intCast(d))); + a.add(@as(i32, @intCast(rows))); + a.add(@as(i32, @intCast(sub))); + a.add(@as(i32, @intCast(world))); + a.add(@as(i32, @intCast(rank))); + a.add(@as(i32, @intCast(m))); + try cuda.launch.launch(k.tp_combine, .{ .grid = .{ .x = @intCast(rows) }, .block = .{ .x = 256 } }, s, &a); +} + +/// The tree finished from the gathered partials of `rows` rows (one gather region): sum[r] for each. +pub fn tpFinish(k: *const Kernels, s: cuda.Stream, gath: u64, sum: u64, rows: usize, d: usize, world: usize, m: usize) !void { + if (rows == 0) return; + var a: cuda.Args = .{}; + a.add(gath); + a.add(sum); + a.add(@as(i32, @intCast(d))); + a.add(@as(i32, @intCast(rows))); + a.add(@as(i32, @intCast(world))); + a.add(@as(i32, @intCast(m))); + try cuda.launch.launch(k.tp_finish, .{ .grid = .{ .x = @intCast(rows) }, .block = .{ .x = 256 } }, s, &a); +} + +/// Expert parallel: y[r, 0] = sum[r] and weights {1, 0, ...} (the write-back then writes the sum's bf16). +pub fn tpPlace(k: *const Kernels, s: cuda.Stream, sum: u64, y: u64, wts: u64, rows: usize, slots: usize, d: usize) !void { + if (rows == 0) return; + var a: cuda.Args = .{}; + a.add(sum); + a.add(y); + a.add(wts); + a.add(@as(i32, @intCast(slots))); + a.add(@as(i32, @intCast(d))); + try cuda.launch.launch(k.tp_place, .{ .grid = .{ .x = @intCast(rows) }, .block = .{ .x = 256 } }, s, &a); +} + +test "exl3: prefill plan groups pairs by expert into items of at most 64" { + const gpa = std.testing.allocator; + // 3 rows of 3 slots: experts 2, 5, -1 / 5, 2, 7 / 2, 2, 513 (out of a 9-entry table: skipped) + const picks = [_]i32{ 2, 5, -1, 5, 2, 7, 2, 2, 513 }; + var counts: [9]u32 = undefined; + var p = try Plan.init(gpa, picks.len, 9); + defer p.deinit(gpa); + build(&picks, null, &counts, &p); + try std.testing.expectEqual(@as(usize, 7), p.n_members); + try std.testing.expectEqualSlices(i32, &.{ 0, 4, 6, 7, 1, 3, 5 }, p.members[0..7]); + try std.testing.expectEqual(@as(usize, 3), p.n_items); + try std.testing.expectEqualSlices(i32, &.{ 2, 0, 4, 5, 4, 2, 7, 6, 1 }, p.items[0..9]); + // a pass of expert 5 alone + var take: [9]bool = @splat(false); + take[5] = true; + build(&picks, &take, &counts, &p); + try std.testing.expectEqualSlices(i32, &.{ 1, 3 }, p.members[0..p.n_members]); + try std.testing.expectEqualSlices(i32, &.{ 5, 0, 2 }, p.items[0 .. 3 * p.n_items]); + // 150 pairs of one expert: items of 64, 64, 22 + var many: [150]i32 = @splat(1); + var p2 = try Plan.init(gpa, many.len, 9); + defer p2.deinit(gpa); + many[0] = -1; + build(&many, null, &counts, &p2); + try std.testing.expectEqual(@as(usize, 149), p2.n_members); + try std.testing.expectEqualSlices(i32, &.{ 1, 0, 64, 1, 64, 64, 1, 128, 21 }, p2.items[0 .. 3 * p2.n_items]); +} diff --git a/zig/src/families/flashnext/exl3/prefill/cuda_exl3_prefill_leaves.zig b/zig/src/families/flashnext/exl3/prefill/cuda_exl3_prefill_leaves.zig new file mode 100644 index 0000000000..1cdd38ac4b --- /dev/null +++ b/zig/src/families/flashnext/exl3/prefill/cuda_exl3_prefill_leaves.zig @@ -0,0 +1,187 @@ +//! The canonical expert reduction at any degree: 64 leaves, fmaf chains, a balanced tree, gathered block partials. + +const std = @import("std"); + +pub const leaves: usize = 64; +pub const ids_per_leaf: usize = 8; +pub const shared_id: usize = 512; + +/// The leaf an expert id belongs to: eight routed ids each, the shared expert with the last leaf. +pub fn leafOf(id: usize) usize { + return if (id >= shared_id) leaves - 1 else id / ids_per_leaf; +} + +/// A complete subtree of the leaf tree: `size` leaves (a power of two) from the multiple of `size` at `start`. +pub const Block = struct { start: usize, size: usize }; + +/// The contiguous leaf run [lo, hi) of `rank` of `world` ranks. +pub fn run(world: usize, rank: usize) [2]usize { + return .{ leaves * rank / world, leaves * (rank + 1) / world }; +} + +/// The maximal complete subtrees covering the leaf run [lo, hi), in leaf order. +pub fn blocks(lo: usize, hi: usize, out: []Block) usize { + var n: usize = 0; + var at = lo; + while (at < hi) { + var size: usize = 1; + while (at % (size * 2) == 0 and at + size * 2 <= hi) size *= 2; + out[n] = .{ .start = at, .size = size }; + n += 1; + at += size; + } + return n; +} + +/// The most blocks a run of leaves needs (two per level). +pub const max_blocks: usize = 12; + +/// The partials every rank sends at `world` ranks: the largest block count of any rank (1 for a power of two). +pub fn partials(world: usize) usize { + var most: usize = 1; + for (0..world) |r| { + const lr = run(world, r); + var bl: [max_blocks]Block = undefined; + most = @max(most, blocks(lr[0], lr[1], &bl)); + } + return most; +} + +/// Sum of the leaf values `v[0..n]` (n a power of two) in the tree's order: siblings added left plus right. +pub fn treeSum(v: []f32) f32 { + var n = v.len; + while (n > 1) : (n /= 2) for (0..n / 2) |i| { + v[i] = v[2 * i] + v[2 * i + 1]; + }; + return v[0]; +} + +/// The branch from all ranks' block partials in leaf order: equal aligned neighbours merge until 64 leaves are left. +pub fn finish(starts: []const Block, vals: []f32) f32 { + var blk: [leaves]Block = undefined; + var n = starts.len; + @memcpy(blk[0..n], starts); + while (n > 1) { + var i: usize = 0; + var m: usize = 0; + while (i < n) { + if (i + 1 < n and blk[i].size == blk[i + 1].size and blk[i].start % (blk[i].size * 2) == 0) { + vals[m] = vals[i] + vals[i + 1]; + blk[m] = .{ .start = blk[i].start, .size = blk[i].size * 2 }; + i += 2; + } else { + vals[m] = vals[i]; + blk[m] = blk[i]; + i += 1; + } + m += 1; + } + n = m; + } + return vals[0]; +} + +/// Leaf sums of one output element: leaf[j] = fmaf chain of y[k] * w[k] over slots in leaf j, in slot order, from +0. +pub fn leafSums(picks: []const i32, y: []const f32, w: []const f32, out: *[leaves]f32) void { + @memset(out, 0); + for (picks, 0..) |p, k| { + if (p < 0) continue; + const j = leafOf(@intCast(p)); + out[j] = @mulAdd(f32, y[k], w[k], out[j]); + } +} + +/// The whole tree over the 64 leaf sums on one rank. +pub fn single(leaf: *const [leaves]f32) f32 { + var v = leaf.*; + return treeSum(&v); +} + +/// The same branch from `world` virtual ranks: each sums its blocks, the partials are gathered, the tree is finished. +pub fn gathered(world: usize, leaf: *const [leaves]f32) f32 { + var starts: [leaves]Block = undefined; + var vals: [leaves]f32 = undefined; + var n: usize = 0; + for (0..world) |r| { + const lr = run(world, r); + var bl: [max_blocks]Block = undefined; + const nb = blocks(lr[0], lr[1], &bl); + for (bl[0..nb]) |b| { + var part: [leaves]f32 = undefined; + @memcpy(part[0..b.size], leaf[b.start..][0..b.size]); + starts[n] = b; + vals[n] = treeSum(part[0..b.size]); + n += 1; + } + } + return finish(starts[0..n], vals[0..n]); +} + +test "exl3: leaves ids map eight to a leaf and the shared expert joins the last" { + try std.testing.expectEqual(@as(usize, 0), leafOf(0)); + try std.testing.expectEqual(@as(usize, 0), leafOf(7)); + try std.testing.expectEqual(@as(usize, 1), leafOf(8)); + try std.testing.expectEqual(@as(usize, 63), leafOf(511)); + try std.testing.expectEqual(@as(usize, 63), leafOf(512)); +} + +test "exl3: leaves the runs of any degree cover the 64 leaves once, in maximal complete subtrees" { + var world: usize = 1; + while (world <= 8) : (world += 1) { + var next: usize = 0; + for (0..world) |r| { + const lr = run(world, r); + try std.testing.expectEqual(next, lr[0]); + next = lr[1]; + var bl: [max_blocks]Block = undefined; + const nb = blocks(lr[0], lr[1], &bl); + var at = lr[0]; + for (bl[0..nb]) |b| { + try std.testing.expectEqual(at, b.start); + try std.testing.expect(std.math.isPowerOfTwo(b.size) and b.start % b.size == 0); + at += b.size; + } + try std.testing.expectEqual(lr[1], at); + } + try std.testing.expectEqual(leaves, next); + } +} + +test "exl3: leaves the gathered tree equals the single-rank tree bit for bit at every degree" { + var prng = std.Random.DefaultPrng.init(0x1ea5); + const rng = prng.random(); + for (0..2000) |_| { + var leaf: [leaves]f32 = undefined; + for (&leaf) |*v| v.* = if (rng.uintLessThan(u8, 5) == 0) 0 else (rng.float(f32) - 0.5) * @exp(rng.float(f32) * 20 - 10); + const want = single(&leaf); + var world: usize = 1; + while (world <= 8) : (world += 1) { + const got = gathered(world, &leaf); + try std.testing.expectEqual(@as(u32, @bitCast(want)), @as(u32, @bitCast(got))); + } + } +} + +test "exl3: leaves a row's leaf sums chain the picks in slot order from +0" { + const picks = [_]i32{ 3, 7, 300, 5, -1, 512, 9, 20, 21, 22, 23 }; + var y: [11]f32 = undefined; + var w: [11]f32 = undefined; + for (&y, &w, 0..) |*a, *b, i| { + a.* = @as(f32, @floatFromInt(i)) * 0.37 + 0.1; + b.* = 1.0 / (@as(f32, @floatFromInt(i)) + 1.5); + } + var out: [leaves]f32 = undefined; + leafSums(&picks, &y, &w, &out); + var l0: f32 = 0; + for ([_]usize{ 0, 1, 3 }) |k| l0 = @mulAdd(f32, y[k], w[k], l0); + try std.testing.expectEqual(@as(u32, @bitCast(l0)), @as(u32, @bitCast(out[0]))); + var l63: f32 = 0; + l63 = @mulAdd(f32, y[5], w[5], l63); + try std.testing.expectEqual(@as(u32, @bitCast(l63)), @as(u32, @bitCast(out[63]))); + try std.testing.expectEqual(@as(f32, 0), out[4]); +} + +test "exl3: leaves a power-of-two degree sends one partial per rank" { + for ([_]usize{ 1, 2, 4, 8 }) |world| try std.testing.expectEqual(@as(usize, 1), partials(world)); + try std.testing.expect(partials(3) > 1 and partials(3) <= max_blocks); +} diff --git a/zig/src/families/flashnext/exl3/prefill/cuda_exl3_prefill_moe.zig b/zig/src/families/flashnext/exl3/prefill/cuda_exl3_prefill_moe.zig new file mode 100644 index 0000000000..7e6ee05083 --- /dev/null +++ b/zig/src/families/flashnext/exl3/prefill/cuda_exl3_prefill_moe.zig @@ -0,0 +1,302 @@ +//! Deep-prefill experts, pass by pass: the pass plan, members and items, prefetch and the rank sum. + +const std = @import("std"); +const cuda = @import("cuda"); +const batch = @import("../../cuda_batch.zig"); +const wt = @import("../../cuda_prompt_wt.zig"); +const load = @import("../pack/cuda_exl3_load.zig"); +const pool_ = @import("../offload/cuda_exl3_pool.zig"); +const pk = @import("cuda_exl3_prefill_k.zig"); +const tp_ = @import("../tp/cuda_exl3_tp.zig"); +const state = @import("cuda_exl3_prefill_state.zig"); +const leaves = @import("cuda_exl3_prefill_leaves.zig"); + +const dims = wt.dims; +const slots = wt.slots; +const table_n = wt.experts_n + 1; +const shared_id = wt.experts_n; + +const Deep = state.Deep; +const Options = state.Options; +const Pass = state.Pass; +const now = state.now; +const profile = state.profile; +const stats = &state.stats; +const upload = state.upload; + +/// Lines a pass may fill: half the layer's lines (the other half holds the previous pass, still computing). +fn passLines(pool: *const pool_.Pool, o: Options) usize { + const lines = pool.opt.lines; + return @max(1, if (o.pass_lines_half) lines / 2 else lines); +} + +/// The layer's experts over all `n` rows of the deep chunk, expert-major. +pub fn moe(comptime Tri: type, dp: *Deep, p: *batch.Pass(Tri), lw: *const @import("../../cuda_prompt_load.zig").Layer, layer: usize, n: usize, o: Options) !void { + const x = lw.exl3.?; + const st = p.stream; + const pairs = n * slots; + const hl = dp.hl; + const t0 = now(p.io); + // the picks of every row on the host: the layer's only wait + const host_picks: []i32 = @alignCast(std.mem.bytesAsSlice(i32, dp.hbytes(hl.picks, pairs * 4))); + try st.d.check(st.d.api.cuMemcpyDtoHAsync_v2(host_picks.ptr, dp.dptr(dp.lay.picks), pairs * 4, st.handle), "cuMemcpyDtoHAsync"); + try st.synchronize(); + maskOwned(host_picks, x.tp); + const t1 = now(p.io); + try plan(dp, x, host_picks, o); + // members and items of every pass in one upload (the pinned bytes stay put until the next layer's wait) + const members: []i32 = @alignCast(std.mem.bytesAsSlice(i32, dp.hbytes(hl.members, pairs * 4))); + const items: []i32 = @alignCast(std.mem.bytesAsSlice(i32, dp.hbytes(hl.items, 3 * pk.maxItems(pairs, table_n) * 4))); + const counts = build(dp, host_picks, members, items); + try upload(st, dp.dptr(dp.lay.members), std.mem.sliceAsBytes(members[0..counts.members])); + try upload(st, dp.dptr(dp.lay.items), std.mem.sliceAsBytes(items[0 .. 3 * counts.items])); + const t2 = now(p.io); + stats.host_wait_ns += @intCast(t1 - t0); + stats.plan_ns += @intCast(t2 - t1); + stats.members += counts.members; + const pool = x.pool; + var fake_rows: usize = 0; + for (dp.passes[0..dp.n_passes]) |ps| { + var r: ?*const pool_.Round = null; + if (pool) |pl| { + const ids_ = dp.order[ps.id_lo..ps.id_hi]; + fake_rows = fakePicks(dp, ids_); + if (fake_rows > 0) { + r = try pl.planUnobserved(dp.fake[0 .. fake_rows * slots], fake_rows); + if (r.?.n_passes != 1) return error.Exl3PrefillPass; + try pl.stage(r.?, 0, st); + stats.copies += r.?.misses_(); + } + } + const l = if (pool) |pl| pl.tables() else x.layer; + if (o.slots_kernel) { + try slotsPass(dp, x, st, l, host_picks, ps, n); + if (r) |rr| try pool.?.release(rr, 0, st); + stats.passes += 1; + continue; + } + const dv: pk.Device = .{ .members = dp.dptr(dp.lay.members) + ps.m_lo * 4, .items = dp.dptr(dp.lay.items) + ps.i_lo * 12, .scratch = dp.dptr(dp.lay.kscratch) }; + try pk.run(&dp.kern, st, l, dp.dptr(dp.lay.mixed), dims, dp.dptr(dp.lay.picks), slots, ps.m_n, ps.i_n, dv, dp.dptr(dp.lay.y)); + if (r) |rr| try pool.?.release(rr, 0, st); + stats.passes += 1; + } + // the next layer's first pass, staged on the copy engine while this layer's experts and the next blocks run + if (o.prefetch and layer + 1 < p.m.layers.len) { + if (p.m.layers[layer + 1].exl3) |nx| if (nx.pool) |npl| try prefetch(dp, npl, x.tp, o); + } + p.pending = layer; +} + +/// Expert parallel: another rank's picks become -1 on the host (this rank plans, copies and computes its own pairs). +pub fn maskOwned(picks: []i32, r: ?tp_.Rank) void { + const rk = r orelse return; + for (picks) |*pk_| if (pk_.* >= 0 and !rk.owns(@intCast(pk_.*))) { + pk_.* = -1; + }; +} + +/// The canonical reduction at any degree: block partials, an fp32 all-gather, the tree finished on every rank, placed. +pub fn rankSum(dp: *Deep, p: anytype, x: *const load.Experts, n: usize, sub: usize) !void { + const st = p.stream; + const sum = dp.dptr(dp.lay.sum); + const y = dp.dptr(dp.lay.y); + const wts = dp.dptr(dp.lay.wts); + const world: usize = if (x.tp) |r| r.world else 1; + const rank: usize = if (x.tp) |r| r.rank else 0; + const m = leaves.partials(world); + const gath = sum + dp.rows * dims * 4; + const t0 = if (profile()) now(p.io) else 0; + const use_overlap = world > 1 and overlap() and dp.red != null; + const per = if (use_overlap) sub else n; + try pk.tpCombine(&dp.kern, st, y, wts, dp.dptr(dp.lay.picks), if (world == 1) sum else gath, n, slots, dims, per, world, rank, m); + if (world == 1) { + if (!dp.reduced_branch) try pk.tpPlace(&dp.kern, st, sum, y, wts, n, slots, dims); + return; + } + const c = p.m.collective orelse return error.TensorParallelUnavailable; + if (!use_overlap) { + try c.gatherF32(st, gath + rank * m * n * dims * 4, gath, m * n * dims); + try pk.tpFinish(&dp.kern, st, gath, sum, n, dims, world, m); + if (!dp.reduced_branch) try pk.tpPlace(&dp.kern, st, sum, y, wts, n, slots, dims); + } else { + const rs = dp.red.?; + try dp.red_go.record(st); + try rs.wait(dp.red_go); + const n_sub = (n + sub - 1) / sub; + for (0..n_sub) |b| { + const r0 = b * sub; + const rows = @min(sub, n - r0); + const region = gath + r0 * world * m * dims * 4; + try c.gatherF32(rs, region + rank * m * rows * dims * 4, region, m * rows * dims); + try pk.tpFinish(&dp.kern, rs, region, sum + r0 * dims * 4, rows, dims, world, m); + if (!dp.reduced_branch) try pk.tpPlace(&dp.kern, rs, sum + r0 * dims * 4, y + r0 * slots * dims * 4, wts + r0 * slots * 4, rows, slots, dims); + try dp.red_ev[b].record(rs); + } + dp.red_n = n_sub; + } + if (profile()) { + if (dp.red) |rs| try c.finish(rs); + try c.finish(st); + stats.reduce_ns += @intCast(now(p.io) - t0); + } + stats.reduces += 1; +} + +fn overlap() bool { + const v = std.c.getenv("TF_FLASHNEXT_EXL3_PREFILL_OVERLAP") orelse return true; + return !std.mem.eql(u8, std.mem.span(v), "0"); +} + +/// Picks for a pool round that holds `ids` (routed, < 512): rows of 11 slots, -1 after the last. +fn fakePicks(dp: *Deep, ids: []const u16) usize { + var routed: usize = 0; + for (ids) |id| routed += @intFromBool(id != shared_id); + if (routed == 0) return 0; + const rows = (routed + slots - 1) / slots; + std.debug.assert(rows * slots <= dp.fake.len); + @memset(dp.fake[0 .. rows * slots], -1); + var k: usize = 0; + for (ids) |id| { + if (id == shared_id) continue; + dp.fake[k] = id; + k += 1; + } + return rows; +} + +/// Stages up to a pass's worth of the next layer's experts that its pool does not hold (lowest ids first): a deep +pub fn prefetch(dp: *Deep, pool: *pool_.Pool, own: ?tp_.Rank, o: Options) !void { + const want = passLines(pool, o); + var ids: [table_n]u16 = undefined; + var k: usize = 0; + for (0..wt.experts_n) |id| { + if (k == want) break; + if (own) |rk| if (!rk.owns(id)) continue; + if (pool_.offload.lineOf(pool.v, pool.opt.layer, @intCast(id)) != null) continue; + ids[k] = @intCast(id); + k += 1; + } + const rows = fakePicks(dp, ids[0..k]); + if (rows == 0) return; + const r = try pool.plan(dp.fake[0 .. rows * slots], rows); + if (r.n_passes != 1) return error.Exl3PrefillPass; + // the copies wait for nothing on the compute stream: the aux stream takes the gates and the release + try pool.stage(r, 0, dp.aux); + try pool.release(r, 0, dp.aux); + stats.prefetched += r.misses_(); + stats.copies += r.misses_(); +} + +/// The layer's passes: pass 0 the shared expert and every expert the pool holds (no copy), then the others in id +fn plan(dp: *Deep, x: *const load.Experts, picks: []const i32, o: Options) !void { + @memset(&dp.cnt, 0); + for (picks) |pk_| { + if (pk_ < 0 or pk_ >= table_n) continue; + dp.cnt[@intCast(pk_)] += 1; + } + const cap = dp.lay.member_cap; + var k: usize = 0; + dp.n_passes = 0; + var held: [table_n]bool = @splat(false); + held[shared_id] = true; + const per: usize = if (x.pool) |pl| passLines(pl, o) else table_n; + if (x.pool) |pl| { + for (0..wt.experts_n) |id| held[id] = pool_.offload.lineOf(pl.v, pl.opt.layer, @intCast(id)) != null; + } else { + @memset(&held, true); + } + // pass 0 (and more if the held experts' members exceed the cap): held experts, the shared one first + var lo: usize = 0; + var members: usize = 0; + const order_ids = struct { + fn close(d: *Deep, lo_: *usize, k_: usize, members_: *usize) void { + if (k_ == lo_.*) return; + d.passes[d.n_passes] = .{ .id_lo = @intCast(lo_.*), .id_hi = @intCast(k_), .m_lo = 0, .m_n = 0, .i_lo = 0, .i_n = 0 }; + d.n_passes += 1; + lo_.* = k_; + members_.* = 0; + } + }; + if (dp.cnt[shared_id] > 0) { + dp.order[k] = shared_id; + k += 1; + members += dp.cnt[shared_id]; + } + for (0..wt.experts_n) |id| { + if (!held[id] or dp.cnt[id] == 0) continue; + if (members + dp.cnt[id] > cap) order_ids.close(dp, &lo, k, &members); + dp.order[k] = @intCast(id); + k += 1; + members += dp.cnt[id]; + } + order_ids.close(dp, &lo, k, &members); + // the misses, `per` new experts a pass + var in_pass: usize = 0; + for (0..wt.experts_n) |id| { + if (held[id] or dp.cnt[id] == 0) continue; + if (in_pass == per or members + dp.cnt[id] > cap) { + order_ids.close(dp, &lo, k, &members); + in_pass = 0; + } + dp.order[k] = @intCast(id); + k += 1; + members += dp.cnt[id]; + in_pass += 1; + } + order_ids.close(dp, &lo, k, &members); +} + +/// Members (pairs grouped by expert in pass order, pairs ascending) and items (first relative to the pass's first +pub fn build(dp: *Deep, picks: []const i32, members: []i32, items: []i32) struct { members: usize, items: usize } { + var m_at: u32 = 0; + var i_at: u32 = 0; + @memset(&dp.start, 0); + for (dp.passes[0..dp.n_passes]) |*ps| { + ps.m_lo = m_at; + ps.i_lo = i_at; + for (dp.order[ps.id_lo..ps.id_hi]) |id| { + const c = dp.cnt[id]; + dp.start[id] = m_at; + var off: u32 = 0; + while (off < c) : (off += pk.tile) { + items[3 * i_at] = id; + items[3 * i_at + 1] = @intCast(m_at - ps.m_lo + off); + items[3 * i_at + 2] = @intCast(@min(@as(u32, pk.tile), c - off)); + i_at += 1; + } + m_at += c; + } + ps.m_n = m_at - ps.m_lo; + ps.i_n = i_at - ps.i_lo; + } + for (picks, 0..) |pk_, i| { + if (pk_ < 0 or pk_ >= table_n) continue; + const id: usize = @intCast(pk_); + members[dp.start[id]] = @intCast(i); + dp.start[id] += 1; + } + return .{ .members = m_at, .items = i_at }; +} + +/// One pass through Experts.runSlots: the deep chunk's picks with every expert outside the pass masked to -1, in the +pub fn slotsPass(dp: *Deep, x: *load.Experts, st: cuda.Stream, l: cuda.exl3_experts.Layer, picks: []const i32, ps: Pass, n: usize) !void { + var in_pass: [table_n]bool = @splat(false); + for (dp.order[ps.id_lo..ps.id_hi]) |id| in_pass[id] = true; + const pairs = n * slots; + // the members' pinned region holds the masked picks (members are not used in this mode); the previous pass's + try st.synchronize(); + const masked: []i32 = @alignCast(std.mem.bytesAsSlice(i32, dp.hbytes(dp.hl.members, pairs * 4))); + for (picks[0..pairs], masked) |pk_, *o| o.* = if (pk_ >= 0 and pk_ < table_n and in_pass[@intCast(pk_)]) pk_ else -1; + try upload(st, dp.dptr(dp.lay.members), std.mem.sliceAsBytes(masked)); + const ex = try x.ops.grouped(); + const block: usize = @import("../model/cuda_exl3_pass.zig").moe_rows; + const bytes = try cuda.exl3_experts.scratchBytes(l, block, slots); + const base = try x.ops.scratch(st, &x.ops.moe, bytes); + var r0: usize = 0; + while (r0 < n) : (r0 += block) { + const rows = @min(block, n - r0); + var sc = try cuda.exl3_experts.scratch(base, l, block, slots); + sc.y = dp.dptr(dp.lay.y) + r0 * slots * dims * 4; + try ex.runSlots(st, l, dp.dptr(dp.lay.mixed) + r0 * dims * 2, dims, .bf16, dp.dptr(dp.lay.members) + r0 * slots * 4, rows, sc); + } +} diff --git a/zig/src/families/flashnext/exl3/prefill/cuda_exl3_prefill_ring.zig b/zig/src/families/flashnext/exl3/prefill/cuda_exl3_prefill_ring.zig new file mode 100644 index 0000000000..9e1948c978 --- /dev/null +++ b/zig/src/families/flashnext/exl3/prefill/cuda_exl3_prefill_ring.zig @@ -0,0 +1,562 @@ +//! The prefill line ring: during an expert-major prefill every decoder layer's pool lines. + +const std = @import("std"); +const cuda = @import("cuda"); +const pool_ = @import("../offload/cuda_exl3_pool.zig"); +const wt = @import("../../cuda_prompt_wt.zig"); +const attribution_ = @import("cuda_exl3_prefill_attribution.zig"); + +const layers_n = wt.layers_n; +const routed_n = wt.experts_n; +const table_n = routed_n + 1; +const batch_ring = 1024; +const release_ring = 4096; +/// Copies a batch holds before its table upload and event. +const batch_copies = 16; +/// Pinned snapshots of a layer's pointer tables, one a batch in flight (a batch uploads its snapshot, never the live +const snap_ring = 64; + +const SlotState = enum(u8) { free, pending }; + +pub const Stats = struct { + copies: u64 = 0, + dma_requested_bytes: u64 = 0, + hit_picks: u64 = 0, + miss_picks: u64 = 0, + hit_expert_requests: u64 = 0, + miss_expert_requests: u64 = 0, + wasted: u64 = 0, + ring_full: u64 = 0, + release_waits: u64 = 0, + compute_waits: u64 = 0, + reseeded: u64 = 0, + + pub fn observeDemand(stats: *Stats, present: bool, picks: u32) void { + if (picks == 0) return; + if (present) { + stats.hit_picks += picks; + stats.hit_expert_requests += 1; + } else { + stats.miss_picks += picks; + stats.miss_expert_requests += 1; + } + } +}; + +pub const Ring = struct { + gpa: std.mem.Allocator, + d: *const cuda.Driver, + pools: [layers_n]*pool_.Pool, + addr: []u64, + /// Each slot's projection offsets and room (its pool's line layout: under expert parallel a rank's layers hold + off: [][3]u32, + room_of: [][3]u32, + state: []SlotState, + owner: []u32, // (layer << 16 | id) + 1, 0: none + epoch_of: []u64, + next: usize = 0, + pending_n: usize = 0, + copy: cuda.Stream, + cev: [batch_ring]cuda.Event, + batches: u64 = 0, + rev: [release_ring]cuda.Event, + epoch: u64 = 0, + copy_ordered_epoch: u64 = 0, + tlay: cuda.exl3_pool.Tables, + /// Pinned images of each layer's tables (pointers and widths) as the device holds them during the prefill. + shadow: cuda.HostBuffer, + snaps: cuda.HostBuffer, + /// Slot of each layer's expert in the current chunk, or -1; the batch that copied it. + slot: [layers_n][routed_n]i32 = undefined, + batch_of: [layers_n][routed_n]u64 = undefined, + /// Prefill member counts a layer has seen (the guess order for the next chunks). + counts: [layers_n][routed_n]u32 = undefined, + /// The open batch: its layer and copies. + open_layer: ?usize = null, + open_n: usize = 0, + /// The routed ids this rank holds (expert parallel, cuda_exl3_tp.zig): a prefetch never guesses another rank's. + own_lo: usize = 0, + own_hi: usize = routed_n, + stats: Stats = .{}, + telemetry_id: u64 = 0, + attribution: ?attribution_.State = null, + + /// A ring over `pools` (one per decoder layer, every line the same layout); error.Invalid otherwise. + pub fn init(gpa: std.mem.Allocator, d: *const cuda.Driver, pools: []const *pool_.Pool, telemetry: bool) !*Ring { + if (pools.len != layers_n) return error.Invalid; + // the tables' entries: 513 (routed and shared), or 512 on an expert-parallel rank without the shared expert + const te = pools[0].lines.spec.experts; + if (te != table_n and te != routed_n) return error.Invalid; + var n: usize = 0; + for (pools) |p| { + if (p.lines.spec.experts != te) return error.Invalid; + n += p.opt.lines; + } + const r = try gpa.create(Ring); + errdefer gpa.destroy(r); + r.* = .{ .gpa = gpa, .d = d, .pools = undefined, .off = undefined, .room_of = undefined, .addr = undefined, .state = undefined, .owner = undefined, .epoch_of = undefined, .copy = undefined, .cev = undefined, .rev = undefined, .tlay = cuda.exl3_pool.Tables.init(te), .shadow = undefined, .snaps = undefined }; + r.telemetry_id = try @import("../model/exl3_telemetry.zig").identity(); + @memcpy(&r.pools, pools); + r.addr = try gpa.alloc(u64, n); + errdefer gpa.free(r.addr); + r.off = try gpa.alloc([3]u32, n); + errdefer gpa.free(r.off); + r.room_of = try gpa.alloc([3]u32, n); + errdefer gpa.free(r.room_of); + r.state = try gpa.alloc(SlotState, n); + errdefer gpa.free(r.state); + r.owner = try gpa.alloc(u32, n); + errdefer gpa.free(r.owner); + r.epoch_of = try gpa.alloc(u64, n); + errdefer gpa.free(r.epoch_of); + r.attribution = try attribution_.State.init(gpa, n, telemetry); + errdefer if (r.attribution) |*state| state.deinit(gpa); + var k: usize = 0; + for (pools) |p| for (0..p.opt.lines) |i| { + const lay = p.lines.layout; + r.addr[k] = p.lines.at(i, 0); + for (0..3) |w| { + const stop = if (w < 2) lay.off[w + 1] else lay.line_bytes; + r.off[k][w] = @intCast(lay.off[w]); + r.room_of[k][w] = @intCast(stop - lay.off[w]); + } + k += 1; + }; + r.copy = try cuda.Stream.init(d, true); + errdefer r.copy.deinit(); + var made_c: usize = 0; + errdefer for (r.cev[0..made_c]) |*e| e.deinit(); + for (&r.cev) |*e| { + e.* = try cuda.Event.init(d, false); + made_c += 1; + } + var made_r: usize = 0; + errdefer for (r.rev[0..made_r]) |*e| e.deinit(); + for (&r.rev) |*e| { + e.* = try cuda.Event.init(d, false); + made_r += 1; + } + r.shadow = try cuda.HostBuffer.alloc(d, layers_n * r.tlay.bytes); + errdefer r.shadow.free(); + r.snaps = try cuda.HostBuffer.alloc(d, snap_ring * r.tlay.k2[0]); + for (&r.counts) |*c| @memset(c, 0); + return r; + } + + pub fn deinit(r: *Ring) void { + r.copy.synchronize() catch {}; + r.snaps.free(); + r.shadow.free(); + for (&r.rev) |*e| e.deinit(); + for (&r.cev) |*e| e.deinit(); + r.copy.deinit(); + if (r.attribution) |*state| state.deinit(r.gpa); + r.gpa.free(r.epoch_of); + r.gpa.free(r.room_of); + r.gpa.free(r.off); + r.gpa.free(r.owner); + r.gpa.free(r.state); + r.gpa.free(r.addr); + r.gpa.destroy(r); + } + + pub fn slots(r: *const Ring) usize { + return r.addr.len; + } + + fn shadowOf(r: *Ring, layer: usize) []u8 { + return r.shadow.bytes[layer * r.tlay.bytes ..][0..r.tlay.bytes]; + } + + fn ptrEntry(r: *Ring, layer: usize, which: usize, id: usize) *align(1) u64 { + return @ptrCast(r.shadowOf(layer)[r.tlay.ptr[which] + id * 8 ..][0..8]); + } + + fn k2Entry(r: *Ring, layer: usize, which: usize, id: usize) *align(1) i32 { + return @ptrCast(r.shadowOf(layer)[r.tlay.k2[which] + id * 4 ..][0..4]); + } + + /// Before a prefill: nothing of the pools' may run (`s` drained), every slot free, each layer's table image read + pub fn begin(r: *Ring, s: cuda.Stream) !void { + try s.synchronize(); + for (r.pools) |p| try p.lines.drain(); + @memset(r.state, .free); + @memset(r.owner, 0); + @memset(r.epoch_of, 0); + r.next = 0; + r.pending_n = 0; + r.copy_ordered_epoch = r.epoch; + r.open_layer = null; + r.open_n = 0; + for (&r.slot) |*row| @memset(row, -1); + for (r.pools, 0..) |p, l| { + const img = r.shadowOf(l); + try p.lines.tables.download(0, img); + for (0..routed_n) |id| for (0..3) |w| { + r.ptrEntry(l, w, id).* = 0; + r.k2Entry(l, w, id).* = p.k2[id][w]; + }; + try upload(r.copy, p.lines.tables.ptr, img); + } + // the images above are read before the host rewrites them + try r.copy.synchronize(); + } + + /// A new deep chunk: every layer's slots of the last one are stale (nothing reads them any more). + pub fn newChunk(r: *Ring) !void { + try r.closeBatch(); + for (r.state, r.owner, 0..) |*st, *ow, index| { + if (st.* == .pending) { + r.stats.wasted += 1; + if (r.attribution) |*state| state.unused(index); + } + st.* = .free; + ow.* = 0; + } + r.pending_n = 0; + for (&r.slot, 0..) |*row, l| { + @memset(row, -1); + for (0..routed_n) |id| for (0..3) |w| { + r.ptrEntry(l, w, id).* = 0; + }; + } + } + + /// After a prefill: every pool empty (policy state, routed table entries), the device idle; then, with `reseed` + pub fn end(r: *Ring, s: cuda.Stream, reseed: bool) !void { + try r.closeBatch(); + try r.copy.synchronize(); + try s.synchronize(); + for (r.pools, 0..) |p, l| { + const img = r.shadowOf(l); + for (0..routed_n) |id| for (0..3) |w| { + r.ptrEntry(l, w, id).* = 0; + r.k2Entry(l, w, id).* = 0; + }; + try p.lines.tables.upload(0, img); + p.v = pool_.offload.Vram.init(p.slots_tab, p.line_tab, p.stamp_tab, p.layer_tab); + p.prev_n = 0; + @memset(p.lines.line_pending, false); + @memset(p.lines.line_epoch, 0); + } + if (!reseed) return; + for (r.pools, 0..) |p, l| { + var order: [routed_n]u16 = undefined; + var n: usize = 0; + for (0..routed_n) |id| { + if (r.counts[l][id] == 0) continue; + order[n] = @intCast(id); + n += 1; + } + const Ctx = struct { + c: *const [routed_n]u32, + fn hotter(c: @This(), a: u16, b: u16) bool { + if (c.c[a] != c.c[b]) return c.c[a] > c.c[b]; + return a < b; + } + }; + std.mem.sort(u16, order[0..n], Ctx{ .c = &r.counts[l] }, Ctx.hotter); + const take = @min(n, p.opt.lines, p.opt.max_rows * p.opt.slots); + if (take == 0) continue; + const rows = (take + p.opt.slots - 1) / p.opt.slots; + var fake: [routed_n + 16]i32 = undefined; + const picks = fake[0 .. rows * p.opt.slots]; + @memset(picks, -1); + for (order[0..take], 0..) |id, i| picks[i] = id; + const rd = try p.plan(picks, rows); + // the pool's copy stream does the copies; the ring's stream takes the gates and the release + for (0..rd.n_passes) |j| { + try p.stage(rd, j, r.copy); + try p.release(rd, j, r.copy); + } + r.stats.reseeded += rd.misses_(); + } + } + + fn closeBatch(r: *Ring) !void { + const l = r.open_layer orelse return; + const p = r.pools[l]; + // the pointer tables only (widths were written at begin), from this batch's own snapshot + const n = r.tlay.k2[0]; + const b = r.batches; + if (b >= snap_ring) try r.cev[(b - snap_ring) % batch_ring].synchronize(); + const snap = r.snaps.bytes[(b % snap_ring) * n ..][0..n]; + @memcpy(snap, r.shadowOf(l)[0..n]); + try upload(r.copy, p.lines.tables.ptr, snap); + try r.cev[r.batches % batch_ring].record(r.copy); + r.batches += 1; + r.open_layer = null; + r.open_n = 0; + } + + /// The copy stream waits until the pass that last read slot `k` finished. + fn afterReaders(r: *Ring, k: usize) !void { + const e = r.epoch_of[k]; + if (e == 0 or e <= r.copy_ordered_epoch) return; + const ev = r.rev[e % release_ring]; + if (r.epoch - e >= release_ring) { + // its event was recorded over by a newer pass: wait for that one on the host (no device cycle) + try ev.synchronize(); + r.copy_ordered_epoch = e; + return; + } + if (try ev.done()) { + r.copy_ordered_epoch = e; + return; + } + try r.copy.wait(ev); + r.copy_ordered_epoch = e; + r.stats.release_waits += 1; + } + + pub fn isLoaded(r: *const Ring, layer: usize, id: usize) bool { + return r.slot[layer][id] >= 0; + } + + /// Free slots ahead of the cursor before a pending one (the ring's room). + pub fn room(r: *const Ring) usize { + return r.slots() - r.pending_n; + } + + /// Copies expert `id` of `layer` into the next free slot; false when the ring holds no free slot. + pub fn load(r: *Ring, layer: usize, id: usize, source: attribution_.Source) !bool { + if (r.slot[layer][id] >= 0) return true; + if (r.pending_n == r.slots()) { + r.stats.ring_full += 1; + return false; + } + const p = r.pools[layer]; + const src = p.src.routed[id]; + // the oldest slots are consumed first; skip any still pending, or too small for this expert + var tried: usize = 0; + while (r.state[r.next] == .pending or !fits(r.room_of[r.next], src)) { + r.next = (r.next + 1) % r.slots(); + tried += 1; + if (tried == r.slots()) { + r.stats.ring_full += 1; + return false; + } + } + const k = r.next; + r.next = (r.next + 1) % r.slots(); + // the open batch's snapshot first (it must not see the evicted expert gone). + if (r.open_layer != null and (r.open_layer.? != layer or r.open_n == batch_copies)) try r.closeBatch(); + try r.afterReaders(k); + if (r.owner[k] != 0) { + const o = r.owner[k] - 1; + const ol = o >> 16; + const oid = o & 0xffff; + r.slot[ol][oid] = -1; + for (0..3) |w| r.ptrEntry(ol, w, oid).* = 0; + } + const dst = r.addr[k]; + const offsets: [3]usize = .{ r.off[k][0], r.off[k][1], r.off[k][2] }; + const capacity = try std.math.add(usize, offsets[2], r.room_of[k][2]); + _ = try cuda.exl3_pool.dma.directLines(r.copy, dst, capacity, offsets, src, p.src.pinned_bank); + for (0..3) |w| r.ptrEntry(layer, w, id).* = dst + r.off[k][w]; + r.slot[layer][id] = @intCast(k); + r.batch_of[layer][id] = r.batches; + r.owner[k] = @intCast(((layer << 16) | id) + 1); + r.state[k] = .pending; + r.pending_n += 1; + r.open_layer = layer; + r.open_n += 1; + r.stats.copies += 1; + r.stats.dma_requested_bytes += src[0].len + src[1].len + src[2].len; + if (r.attribution) |*state| state.copied(k, source, src[0].len + src[1].len + src[2].len); + return true; + } + + /// The compute stream waits for every copy (and table write) of `ids` of `layer`. + pub fn gate(r: *Ring, s: cuda.Stream, layer: usize, ids: []const u16) !void { + var last: ?u64 = null; + for (ids) |id| { + if (id >= routed_n) continue; + const b = r.batch_of[layer][id]; + last = if (last) |x| @max(x, b) else b; + } + const b = last orelse return; + if (b == r.batches) try r.closeBatch(); // its batch is still open + const ev = r.cev[b % batch_ring]; + if (r.batches - b > batch_ring) { + try r.copy.synchronize(); + return; + } + try s.wait(ev); + r.stats.compute_waits += 1; + } + + /// After a pass that read `ids` of `layer` was enqueued on `s`: its slots become free once it finishes. + pub fn release(r: *Ring, s: cuda.Stream, layer: usize, ids: []const u16) !void { + r.epoch += 1; + try r.rev[r.epoch % release_ring].record(s); + for (ids) |id| { + if (id >= routed_n) continue; + const k: usize = @intCast(r.slot[layer][id]); + r.epoch_of[k] = r.epoch; + if (r.state[k] == .pending) { + r.state[k] = .free; + r.pending_n -= 1; + } + } + } + + /// Experts loaded for `layer` that its picks did not use: their slots are free again (never read). + pub fn drop(r: *Ring, layer: usize, used: []const u32) void { + for (0..routed_n) |id| { + if (used[id] != 0) continue; + const k = r.slot[layer][id]; + if (k < 0) continue; + const ku: usize = @intCast(k); + if (r.state[ku] == .pending) { + r.state[ku] = .free; + r.pending_n -= 1; + r.stats.wasted += 1; + if (r.attribution) |*state| state.unused(ku); + } + } + } + + /// Tops `layer` up to `want` loaded experts with the ones not loaded, hottest (prefill counts). + pub fn prefetch(r: *Ring, layer: usize, want: usize, keep_free: usize, source: attribution_.Source) !usize { + var order: [routed_n]u16 = undefined; + var n: usize = 0; + var have: usize = 0; + for (0..routed_n) |id| { + if (r.slot[layer][id] >= 0) { + have += 1; + continue; + } + if (id < r.own_lo or id >= r.own_hi) continue; + order[n] = @intCast(id); + n += 1; + } + if (have >= want) return 0; + const Ctx = struct { + a: *const [routed_n]u32, + b: *const [routed_n]u32, + fn hotter(c: @This(), x: u16, y: u16) bool { + if (c.a[x] != c.a[y]) return c.a[x] > c.a[y]; + if (c.b[x] != c.b[y]) return c.b[x] > c.b[y]; + return x < y; + } + }; + const p = r.pools[layer]; + std.mem.sort(u16, order[0..n], Ctx{ .a = &r.counts[layer], .b = &p.counts[p.opt.layer] }, Ctx.hotter); + var queued: usize = 0; + for (order[0..n]) |id| { + if (have + queued == want or r.room() <= keep_free) break; + if (!try r.load(layer, id, source)) break; + queued += 1; + } + try r.closeBatch(); + return queued; + } + + pub fn flush(r: *Ring) !void { + try r.closeBatch(); + } +}; + +fn fits(room: [3]u32, src: [3][]const u8) bool { + for (room, src) |n, s| if (s.len > n) return false; + return true; +} + +fn upload(s: cuda.Stream, dst: u64, bytes: []const u8) !void { + try s.d.check(s.d.api.cuMemcpyHtoDAsync_v2(dst, bytes.ptr, bytes.len, s.handle), "cuMemcpyHtoDAsync"); +} + +test "exl3: a ring slot takes an expert only if each projection fits its line's room" { + const big: [300]u8 = @splat(0); + const room: [3]u32 = .{ 256, 512, 256 }; + try std.testing.expect(fits(room, .{ big[0..256], big[0..300], big[0..10] })); + try std.testing.expect(!fits(room, .{ big[0..300], big[0..10], big[0..10] })); + try std.testing.expect(fits(room, .{ &.{}, &.{}, &.{} })); +} + +const ReleaseFixture = @import("cuda_exl3_prefill_ring_fixture.zig").Fixture; + +test "exl3: prefill attribution contracts" { + _ = @import("cuda_exl3_prefill_attribution_test.zig"); +} + +test "exl3: prefill ring queues one release fence for a pass's pending slot copies" { + var driver: cuda.Driver = undefined; + const ring = try std.testing.allocator.create(Ring); + defer std.testing.allocator.destroy(ring); + var epochs: [32]u64 = @splat(7); + ReleaseFixture.init(&driver, ring, &epochs, 7); + for (epochs, 0..) |_, index| { + try ring.afterReaders(index); + try upload(ring.copy, 42, &.{1}); + } + try std.testing.expectEqualSlices(ReleaseFixture.Call, &.{ .query, .wait }, ReleaseFixture.calls[0..2]); + try std.testing.expectEqual(@as(usize, 34), ReleaseFixture.count); + for (ReleaseFixture.calls[2..ReleaseFixture.count]) |call| try std.testing.expectEqual(ReleaseFixture.Call.copy, call); + try std.testing.expectEqual(@as(u64, 7), ring.copy_ordered_epoch); + try std.testing.expectEqual(@as(u64, 1), ring.stats.release_waits); +} + +test "exl3: prefill ring queued fences cover older epochs but never a newer reader" { + var driver: cuda.Driver = undefined; + const ring = try std.testing.allocator.create(Ring); + defer std.testing.allocator.destroy(ring); + var epochs = [_]u64{ 4, 2, 4, 5, 0, 1, 6 }; + ReleaseFixture.init(&driver, ring, &epochs, 6); + for (0..6) |index| try ring.afterReaders(index); + try std.testing.expectEqualSlices(ReleaseFixture.Call, &.{ .query, .wait, .query, .wait }, ReleaseFixture.calls[0..ReleaseFixture.count]); + try std.testing.expectEqual(ring.rev[5].handle, ReleaseFixture.last_event); + ReleaseFixture.query_result = cuda.abi.success; + try ring.afterReaders(6); + try ring.afterReaders(3); + try std.testing.expectEqual(@as(usize, 5), ReleaseFixture.count); + try std.testing.expectEqual(@as(u64, 6), ring.copy_ordered_epoch); + try std.testing.expectEqual(@as(u64, 2), ring.stats.release_waits); +} + +test "exl3: prefill ring retains the host fallback at the release event wrap boundary" { + var driver: cuda.Driver = undefined; + const ring = try std.testing.allocator.create(Ring); + defer std.testing.allocator.destroy(ring); + var epochs = [_]u64{ 7, 7 }; + for ([_]u64{ release_ring - 1, release_ring, release_ring + 1 }) |distance| { + ReleaseFixture.init(&driver, ring, &epochs, 7 + distance); + try ring.afterReaders(0); + try ring.afterReaders(1); + if (distance < release_ring) { + try std.testing.expectEqualSlices(ReleaseFixture.Call, &.{ .query, .wait }, ReleaseFixture.calls[0..ReleaseFixture.count]); + try std.testing.expectEqual(@as(u64, 1), ring.stats.release_waits); + } else { + try std.testing.expectEqualSlices(ReleaseFixture.Call, &.{.sync}, ReleaseFixture.calls[0..ReleaseFixture.count]); + try std.testing.expectEqual(@as(u64, 0), ring.stats.release_waits); + } + try std.testing.expectEqual(ring.rev[7].handle, ReleaseFixture.last_event); + try std.testing.expectEqual(@as(u64, 7), ring.copy_ordered_epoch); + } +} + +test "exl3: prefill ring failed release dependencies cannot authorize subsequent copies" { + var driver: cuda.Driver = undefined; + const ring = try std.testing.allocator.create(Ring); + defer std.testing.allocator.destroy(ring); + var epochs = [_]u64{7}; + ReleaseFixture.init(&driver, ring, &epochs, 7); + ReleaseFixture.wait_result = cuda.abi.error_not_ready; + try std.testing.expectError(error.NotReady, ring.afterReaders(0)); + try std.testing.expectEqual(@as(u64, 0), ring.copy_ordered_epoch); + try std.testing.expectEqual(@as(u64, 0), ring.stats.release_waits); + ReleaseFixture.wait_result = cuda.abi.success; + try ring.afterReaders(0); + try std.testing.expectEqualSlices(ReleaseFixture.Call, &.{ .query, .wait, .query, .wait }, ReleaseFixture.calls[0..ReleaseFixture.count]); + try std.testing.expectEqual(@as(u64, 7), ring.copy_ordered_epoch); + + ReleaseFixture.init(&driver, ring, &epochs, 7 + release_ring); + ReleaseFixture.sync_result = cuda.abi.error_not_ready; + try std.testing.expectError(error.NotReady, ring.afterReaders(0)); + try std.testing.expectEqual(@as(u64, 0), ring.copy_ordered_epoch); + ReleaseFixture.sync_result = cuda.abi.success; + try ring.afterReaders(0); + try std.testing.expectEqualSlices(ReleaseFixture.Call, &.{ .sync, .sync }, ReleaseFixture.calls[0..ReleaseFixture.count]); + try std.testing.expectEqual(@as(u64, 7), ring.copy_ordered_epoch); +} diff --git a/zig/src/families/flashnext/exl3/prefill/cuda_exl3_prefill_ring_fixture.zig b/zig/src/families/flashnext/exl3/prefill/cuda_exl3_prefill_ring_fixture.zig new file mode 100644 index 0000000000..62d11ad190 --- /dev/null +++ b/zig/src/families/flashnext/exl3/prefill/cuda_exl3_prefill_ring_fixture.zig @@ -0,0 +1,57 @@ +const cuda = @import("cuda"); + +pub const Fixture = struct { + pub const Call = enum { query, wait, sync, copy }; + pub var calls: [128]Call = undefined; + pub var count: usize = 0; + pub var query_result: cuda.abi.Result = cuda.abi.error_not_ready; + pub var wait_result: cuda.abi.Result = cuda.abi.success; + pub var sync_result: cuda.abi.Result = cuda.abi.success; + pub var last_event: cuda.abi.Event = null; + + fn note(call: Call, event: cuda.abi.Event) void { + calls[count] = call; + count += 1; + if (event != null) last_event = event; + } + + fn query(event: cuda.abi.Event) callconv(.c) cuda.abi.Result { + note(.query, event); + return query_result; + } + + fn wait(_: cuda.abi.Stream, event: cuda.abi.Event, _: c_uint) callconv(.c) cuda.abi.Result { + note(.wait, event); + return wait_result; + } + + fn sync(event: cuda.abi.Event) callconv(.c) cuda.abi.Result { + note(.sync, event); + return sync_result; + } + + fn copy(_: cuda.abi.DevicePtr, _: ?*const anyopaque, _: usize, _: cuda.abi.Stream) callconv(.c) cuda.abi.Result { + note(.copy, null); + return cuda.abi.success; + } + + pub fn init(driver: *cuda.Driver, ring: anytype, epochs: []u64, epoch: u64) void { + count = 0; + query_result = cuda.abi.error_not_ready; + wait_result = cuda.abi.success; + sync_result = cuda.abi.success; + last_event = null; + driver.* = undefined; + driver.api.cuEventQuery = query; + driver.api.cuStreamWaitEvent = wait; + driver.api.cuEventSynchronize = sync; + driver.api.cuMemcpyHtoDAsync_v2 = copy; + ring.* = undefined; + ring.copy = .{ .d = driver, .handle = @ptrFromInt(1) }; + ring.epoch_of = epochs; + ring.epoch = epoch; + ring.copy_ordered_epoch = 0; + ring.stats = .{}; + for (&ring.rev, 0..) |*event, index| event.* = .{ .d = driver, .handle = @ptrFromInt(index + 1) }; + } +}; diff --git a/zig/src/families/flashnext/exl3/prefill/cuda_exl3_prefill_ringmoe.zig b/zig/src/families/flashnext/exl3/prefill/cuda_exl3_prefill_ringmoe.zig new file mode 100644 index 0000000000..abf2c35893 --- /dev/null +++ b/zig/src/families/flashnext/exl3/prefill/cuda_exl3_prefill_ringmoe.zig @@ -0,0 +1,200 @@ +//! Deep-prefill experts through the line ring: its pass order, copies and cross-layer prefetch. + +const std = @import("std"); +const batch = @import("../../cuda_batch.zig"); +const wt = @import("../../cuda_prompt_wt.zig"); +const pool_ = @import("../offload/cuda_exl3_pool.zig"); +const pk = @import("cuda_exl3_prefill_k.zig"); +const prompt = @import("../../cuda_prompt.zig"); +const ring_ = @import("cuda_exl3_prefill_ring.zig"); +const history_ = @import("cuda_exl3_prefill_history.zig"); +const state = @import("cuda_exl3_prefill_state.zig"); +const experts = @import("cuda_exl3_prefill_moe.zig"); + +const dims = wt.dims; +const slots = wt.slots; +const table_n = wt.experts_n + 1; +const routed_n_ = wt.experts_n; +const shared_id = wt.experts_n; + +const Deep = state.Deep; +const Options = state.Options; +const now = state.now; +const stats = &state.stats; +const upload = state.upload; +const build = experts.build; +const maskOwned = experts.maskOwned; +const slotsPass = experts.slotsPass; + +/// The ring over the engine's decoder pools, made once; null (and the per-layer passes) if they cannot form one. +pub fn ringOf(dp: *Deep, e: *prompt.Engine) !?*ring_.Ring { + if (dp.ring) |r| return r; + if (dp.ring_failed) return null; + var pools: [wt.layers_n]*pool_.Pool = undefined; + for (e.layers[0..e.n_layers], 0..) |*l, i| { + const x = l.exl3 orelse { + dp.ring_failed = true; + return null; + }; + pools[i] = x.pool orelse { + dp.ring_failed = true; + return null; + }; + } + dp.ring = ring_.Ring.init(dp.gpa, e.driver, &pools, e.telemetry.enabled) catch |err| { + std.log.warn("flashnext exl3 prefill: no line ring ({t}); each layer stages into its own lines", .{err}); + dp.ring_failed = true; + return null; + }; + if (e.layers[0].exl3.?.tp) |rk| { + dp.ring.?.own_lo = rk.lo; + dp.ring.?.own_hi = @min(rk.hi, routed_n_); + } + std.log.info("flashnext exl3 prefill: a ring of {d} expert slots over the {d} layers' pool lines", .{ dp.ring.?.slots(), wt.layers_n }); + return dp.ring; +} + +/// Slots a prefetch leaves free: a layer's worth of misses and two passes. +pub fn keepFree(r: *const ring_.Ring, o: Options) usize { + return @min(r.slots() / 2, routed_n_ + 2 * o.ring_pass); +} + +/// The layer's experts over all `n` rows of the deep chunk, through the ring. +pub fn ringMoe(comptime Tri: type, dp: *Deep, r: *ring_.Ring, p: *batch.Pass(Tri), lw: *const @import("../../cuda_prompt_load.zig").Layer, layer: usize, n: usize, o: Options) !void { + const x = lw.exl3.?; + const pool = x.pool.?; + const st = p.stream; + const pairs = n * slots; + const hl = dp.hl; + const t0 = now(p.io); + const host_picks: []i32 = @alignCast(std.mem.bytesAsSlice(i32, dp.hbytes(hl.picks, pairs * 4))); + try st.d.check(st.d.api.cuMemcpyDtoHAsync_v2(host_picks.ptr, dp.dptr(dp.lay.picks), pairs * 4, st.handle), "cuMemcpyDtoHAsync"); + try st.synchronize(); + maskOwned(host_picks, x.tp); + const t1 = now(p.io); + @memset(&dp.cnt, 0); + for (host_picks) |pk_| { + if (pk_ < 0 or pk_ >= table_n) continue; + dp.cnt[@intCast(pk_)] += 1; + } + r.drop(layer, dp.cnt[0..routed_n_]); + // pass order: shared, then the held experts by load batch, then the misses hottest first + var held: [routed_n_]u16 = undefined; + var n_held: usize = 0; + var miss: [routed_n_]u16 = undefined; + var n_miss: usize = 0; + var used: usize = 0; + for (0..routed_n_) |id| { + if (dp.cnt[id] == 0) continue; + used += 1; + r.counts[layer][id] +|= dp.cnt[id]; + const present = r.isLoaded(layer, id); + r.stats.observeDemand(present, dp.cnt[id]); + if (present) { + held[n_held] = @intCast(id); + n_held += 1; + } else { + miss[n_miss] = @intCast(id); + n_miss += 1; + } + } + dp.used[layer] = used; + const ByBatch = struct { + r: *const ring_.Ring, + l: usize, + fn lt(c: @This(), a: u16, b: u16) bool { + return c.r.batch_of[c.l][a] < c.r.batch_of[c.l][b]; + } + }; + std.mem.sort(u16, held[0..n_held], ByBatch{ .r = r, .l = layer }, ByBatch.lt); + const ByCount = struct { + c: *const [table_n]u32, + fn gt(c: @This(), a: u16, b: u16) bool { + if (c.c[a] != c.c[b]) return c.c[a] > c.c[b]; + return a < b; + } + }; + std.mem.sort(u16, miss[0..n_miss], ByCount{ .c = &dp.cnt }, ByCount.gt); + // passes: at most ring_pass experts and member_cap members each; the shared expert opens pass 0 + var k: usize = 0; + var lo: usize = 0; + var members: usize = 0; + var in_pass: usize = 0; + dp.n_passes = 0; + const cap = dp.lay.member_cap; + const close = struct { + fn f(d: *Deep, lo_: *usize, k_: usize, members_: *usize, in_pass_: *usize) void { + if (k_ == lo_.*) return; + d.passes[d.n_passes] = .{ .id_lo = @intCast(lo_.*), .id_hi = @intCast(k_), .m_lo = 0, .m_n = 0, .i_lo = 0, .i_n = 0 }; + d.n_passes += 1; + lo_.* = k_; + members_.* = 0; + in_pass_.* = 0; + } + }.f; + if (dp.cnt[shared_id] > 0) { + dp.order[k] = shared_id; + k += 1; + members += dp.cnt[shared_id]; + } + for ([2][]const u16{ held[0..n_held], miss[0..n_miss] }) |list| for (list) |id| { + if (in_pass == o.ring_pass or members + dp.cnt[id] > cap) close(dp, &lo, k, &members, &in_pass); + dp.order[k] = id; + k += 1; + members += dp.cnt[id]; + in_pass += 1; + }; + close(dp, &lo, k, &members, &in_pass); + const staging0 = now(p.io); + const copies0 = r.stats.copies; + for (0..@min(2, dp.n_passes)) |j| try loadPass(dp, r, layer, j, false); + const staging_ns = now(p.io) - staging0; + const members_h: []i32 = @alignCast(std.mem.bytesAsSlice(i32, dp.hbytes(hl.members, pairs * 4))); + const items_h: []i32 = @alignCast(std.mem.bytesAsSlice(i32, dp.hbytes(hl.items, 3 * pk.maxItems(pairs, table_n) * 4))); + const counts = build(dp, host_picks, members_h, items_h); + try upload(st, dp.dptr(dp.lay.members), std.mem.sliceAsBytes(members_h[0..counts.members])); + try upload(st, dp.dptr(dp.lay.items), std.mem.sliceAsBytes(items_h[0 .. 3 * counts.items])); + const t2 = now(p.io); + stats.host_wait_ns += @intCast(t1 - t0); + stats.plan_ns += @intCast(t2 - t1 - staging_ns); + stats.members += counts.members; + const l = pool.tables(); + // the first two passes' copies, then each pass: the next pass's copies, the gate, the kernel, the release + for (dp.passes[0..dp.n_passes], 0..) |ps, j| { + if (j + 2 < dp.n_passes) try loadPass(dp, r, layer, j + 2, true); + try loadPass(dp, r, layer, j, false); + const ids_ = dp.order[ps.id_lo..ps.id_hi]; + try r.gate(st, layer, ids_); + if (o.slots_kernel) { + try slotsPass(dp, x, st, l, host_picks, ps, n); + } else { + const dv: pk.Device = .{ .members = dp.dptr(dp.lay.members) + ps.m_lo * 4, .items = dp.dptr(dp.lay.items) + ps.i_lo * 12, .scratch = dp.dptr(dp.lay.kscratch) }; + try pk.run(&dp.kern, st, l, dp.dptr(dp.lay.mixed), dims, dp.dptr(dp.lay.picks), slots, ps.m_n, ps.i_n, dv, dp.dptr(dp.lay.y)); + } + try r.release(st, layer, ids_); + stats.passes += 1; + } + stats.copies += r.stats.copies - copies0; + dp.history[layer] = .{ .rows = n, .used = used }; + const c1 = r.stats.copies; + for (1..o.ring_ahead + 1) |ahead| { + const nl = layer + ahead; + if (nl >= wt.layers_n) break; + _ = try r.prefetch(nl, history_.want(dp.history[nl], n, used, routed_n_), keepFree(r, o), .cross_layer_prefetch); + } + stats.prefetched += r.stats.copies - c1; + stats.copies += r.stats.copies - c1; + p.pending = layer; +} + +/// Queues the copies of pass `j`'s routed experts the ring does not hold; `best_effort`: stop quietly when full. +fn loadPass(dp: *Deep, r: *ring_.Ring, layer: usize, j: usize, best_effort: bool) !void { + const ps = dp.passes[j]; + for (dp.order[ps.id_lo..ps.id_hi]) |id| { + if (id == shared_id) continue; + if (try r.load(layer, id, .demand)) continue; + if (best_effort) return; + return error.Exl3PrefillRingFull; + } + try r.flush(); +} diff --git a/zig/src/families/flashnext/exl3/prefill/cuda_exl3_prefill_state.zig b/zig/src/families/flashnext/exl3/prefill/cuda_exl3_prefill_state.zig new file mode 100644 index 0000000000..8d2a537246 --- /dev/null +++ b/zig/src/families/flashnext/exl3/prefill/cuda_exl3_prefill_state.zig @@ -0,0 +1,293 @@ +//! Deep-prefill state: options, stats, the chunk's device and pinned layouts, and the buffers one engine owns. + +const std = @import("std"); +const cuda = @import("cuda"); +const batch = @import("../../cuda_batch.zig"); +const wt = @import("../../cuda_prompt_wt.zig"); +const pk = @import("cuda_exl3_prefill_k.zig"); +const ring_ = @import("cuda_exl3_prefill_ring.zig"); +const history_ = @import("cuda_exl3_prefill_history.zig"); +const window_ = @import("exl3_prefill_window.zig"); +const leaves = @import("cuda_exl3_prefill_leaves.zig"); + +const dims = wt.dims; +const wide = wt.wide; +const streams = wt.streams; +const slots = wt.slots; +const table_n = wt.experts_n + 1; +const att_rows = batch.att_rows; +const routed_guess: usize = 400; +const pss_n = batch.pss_n; + +pub const Options = struct { + enabled: bool = true, + /// Rows a deep chunk holds (whole attention blocks). + rows: usize = 8192, + /// The deep chunk the server reserves at load (reserveFor): its buffers come out of the pools' share of the card + serve_rows: usize = 4096, + /// The shortest prompt that takes this path. + min_rows: usize = 512, + /// New experts a pass stages at most, as a fraction of the layer's lines (0.5: double buffering). + pass_lines_half: bool = true, + /// Stage the next layer's first pass while this layer's blocks compute. + prefetch: bool = true, + /// Run each pass through Experts.runSlots (256-row blocks over masked picks) instead of the 64-member prefill + slots_kernel: bool = false, + /// Every decoder layer's pool lines as one ring the copy engine streams the coming layers' experts into + ring: bool = true, + /// Experts a ring pass computes at most. + ring_pass: usize = 32, + /// Layers ahead the ring prefetches (the guess: the hottest experts, as many as the current layer used). + ring_ahead: usize = 1, + /// Dense EXL3 projections of a deep chunk's sub-blocks on #806's prompt tiles (Dense.prompt, .tiles) instead of + tiles: bool = true, + /// After a ring prefill, refill every pool with its layer's hottest experts of the prefills (measured: decode hit + reseed: bool = false, + reduced_branch: bool = false, + two_chunk: bool = false, +}; + +pub fn optionsFromEnv() Options { + var o: Options = .{}; + if (std.c.getenv("TF_FLASHNEXT_EXL3_PREFILL")) |v| o.enabled = !std.mem.eql(u8, std.mem.span(v), "0"); + if (std.c.getenv("TF_FLASHNEXT_EXL3_PREFILL_ROWS")) |v| { + o.rows = std.fmt.parseInt(usize, std.mem.span(v), 10) catch o.rows; + o.serve_rows = o.rows; + } + if (std.c.getenv("TF_FLASHNEXT_EXL3_PREFILL_REDUCED_BRANCH")) |v| o.reduced_branch = std.mem.eql(u8, std.mem.span(v), "1"); + o.two_chunk = window_.enabled(if (std.c.getenv("TF_FLASHNEXT_EXL3_PREFILL_TWO_CHUNK")) |value| std.mem.span(value) else null); + if (std.c.getenv("TF_FLASHNEXT_EXL3_PREFILL_RING")) |v| o.ring = !std.mem.eql(u8, std.mem.span(v), "0"); + if (std.c.getenv("TF_FLASHNEXT_EXL3_PREFILL_RING_AHEAD")) |v| o.ring_ahead = std.fmt.parseInt(usize, std.mem.span(v), 10) catch o.ring_ahead; + return o; +} + +/// The options every call reads (a test may change them between calls). +pub var options: ?Options = null; + +pub fn opts() Options { + if (options == null) options = optionsFromEnv(); + return options.?; +} + +/// What the last deep forward did (read by tf-cuda-test exl3-prefill engine). +pub const Stats = struct { + chunks: u64 = 0, + rows: u64 = 0, + passes: u64 = 0, + copies: u64 = 0, + prefetched: u64 = 0, + members: u64 = 0, + host_wait_ns: u64 = 0, // the picks' D2H and its wait, every layer + plan_ns: u64 = 0, // host plan of the passes (members, items, pool rounds) + reduces: u64 = 0, // expert parallel: the ranks' all-reduces (one a layer a deep chunk) + reduce_ns: u64 = 0, // their combine + all-reduce + place, measured only with TF_FLASHNEXT_PROFILE=1 (a wait each) + head_rows: u64 = 0, // MTP head rows run expert-major (`head`) +}; +pub var stats: Stats = .{}; + +/// Device bytes of a deep chunk's own rows: streams, sums, gates, mixed rows, slot outputs, picks, n-gram rows, and +pub fn deviceBytes(rows: usize) usize { + return Layout.of(rows, 1).total; +} + +/// The same at `world` ranks (expert parallel adds the ranks' sum, rows x 2560 fp32). +pub fn deviceBytesAt(rows: usize, world: usize) usize { + return Layout.of(rows, world).total; +} + +/// The reduction's branch [rows, D] fp32, then (expert parallel) the gathered block partials of every rank. +pub fn sumBytes(rows: usize, world: usize) usize { + const one = rows * dims * 4; + return if (world > 1) one * (1 + world * leaves.partials(world)) else one; +} + +pub const Layout = struct { + h: usize, + pss: usize, + inj: usize, + mixed: usize, + xsm: usize, + y: usize, + wts: usize, + picks: usize, + ple_emb: usize, + ple_xs: usize, + ids: usize, + pos: usize, + members: usize, + items: usize, + kscratch: usize, + /// The canonical reduction's branch and the gathered partials (`sumBytes`). + sum: usize, + total: usize, + member_cap: usize, + + pub fn of(rows: usize, world: usize) Layout { + var l: Layout = undefined; + var at: usize = 0; + const take = struct { + fn f(a: *usize, n: usize) usize { + const r = a.*; + a.* = r + (n + 255) / 256 * 256; + return r; + } + }.f; + const pairs = rows * slots; + l.member_cap = 2 * rows; + l.h = take(&at, rows * wide * 2); + l.pss = take(&at, rows * pss_n * 4); + l.inj = take(&at, rows * streams * 2); + l.mixed = take(&at, rows * dims * 2); + l.xsm = take(&at, rows * (dims / 32) * 4); + l.y = take(&at, pairs * dims * 4); + l.wts = take(&at, pairs * 4); + l.picks = take(&at, pairs * 4); + l.ple_emb = take(&at, rows * dims * 2); + l.ple_xs = take(&at, rows * (dims / 32) * 4); + l.ids = take(&at, rows * 4); + l.pos = take(&at, (rows / att_rows + 1) * 16); + l.members = take(&at, pairs * 4); + l.items = take(&at, 3 * pk.maxItems(pairs, table_n) * 4); + l.kscratch = take(&at, pk.scratchBytes(l.member_cap, dims, wt.moe_w)); + l.sum = take(&at, sumBytes(rows, world)); + l.total = at; + return l; + } +}; + +/// Host side: pinned uploads (ids, positions, n-gram rows, members and items) and the picks' download. +pub const HostLayout = struct { + ids: usize, + pos: usize, + pos_stride: usize, + emb: usize, + xs: usize, + picks: usize, + members: usize, + items: usize, + total: usize, + + pub fn of(rows: usize) HostLayout { + var l: HostLayout = undefined; + var at: usize = 0; + const take = struct { + fn f(a: *usize, n: usize) usize { + const r = a.*; + a.* = r + (n + 63) / 64 * 64; + return r; + } + }.f; + const pairs = rows * slots; + l.ids = take(&at, rows * 4); + l.pos_stride = (rows / att_rows + 1) * 16; + l.pos = take(&at, l.pos_stride * (wt.layers_n + 1)); + l.emb = take(&at, rows * dims * 2); + l.xs = take(&at, rows * (dims / 32) * 4); + l.picks = take(&at, pairs * 4); + l.members = take(&at, pairs * 4); + l.items = take(&at, 3 * pk.maxItems(pairs, table_n) * 4); + l.total = at; + return l; + } +}; + +/// One pass of a layer: its experts (the pool round's routed ids), its members and items. +pub const Pass = struct { id_lo: u32, id_hi: u32, m_lo: u32, m_n: u32, i_lo: u32, i_n: u32 }; + +/// Sub-blocks a deep chunk holds at most (chunk's view array). +pub const max_sub = 64; + +pub const Persistent = struct { + h: usize, + pss: usize, + inj: usize, + y: usize, + wts: usize, + total: usize, + + pub fn of(rows: usize) Persistent { + const h = rows * wide * 2; + const pss = std.mem.alignForward(usize, rows * pss_n * 4, 256); + const inj = std.mem.alignForward(usize, rows * streams * 2, 256); + const y = rows * slots * dims * 4; + const wts = std.mem.alignForward(usize, rows * slots * 4, 256); + return .{ .h = 0, .pss = h, .inj = h + pss, .y = h + pss + inj, .wts = h + pss + inj + y, .total = h + pss + inj + y + wts }; + } +}; + +pub const Deep = struct { + gpa: std.mem.Allocator, + owner: *const anyopaque, + rows: usize, + world: usize = 1, + lay: Layout, + hl: HostLayout, + dev: cuda.DeviceBuffer, + sum_dev: ?cuda.DeviceBuffer = null, + host: cuda.HostBuffer, + kern: pk.Kernels, + aux: cuda.Stream, + /// Expert parallel: the ranks' all-reduces run on `red`, one a sub-block. + red: ?cuda.Stream = null, + red_ev: [max_sub]cuda.Event = undefined, + red_go: cuda.Event = undefined, + red_n: usize = 0, + reduced_branch: bool = false, + unit_weights: ?cuda.DeviceBuffer = null, + /// The layer's pass plan: expert ids in pass order, the passes, member counts and starts. + order: [table_n]u16 = undefined, + passes: [table_n + 1]Pass = undefined, + n_passes: usize = 0, + cnt: [table_n]u32 = undefined, + start: [table_n]u32 = undefined, + /// Experts the previous layer's prefetch staged for each layer (held at its MoE). + fake: [48 * slots]i32 = undefined, + /// The prefill line ring over the pools (made at the first prefill, once the pools exist). + ring: ?*ring_.Ring = null, + ring_failed: bool = false, + /// Experts each layer used in the current chunk (the next layers' prefetch size). + used: [wt.layers_n]usize = @splat(routed_guess), + history: [wt.layers_n]history_.Slot = @splat(.{}), + history_shape: ?history_.Shape = null, + second: ?*Deep = null, + shared: ?*const Deep = null, + + pub fn dptr(dp: *const Deep, off: usize) u64 { + if (dp.shared) |source| { + const persistent = Persistent.of(dp.rows); + inline for (.{ "h", "pss", "inj", "y", "wts" }) |field| { + if (off == @field(dp.lay, field)) return dp.dev.ptr + @field(persistent, field); + } + if (off == dp.lay.sum) if (dp.sum_dev) |buffer| return buffer.ptr; + return source.dptr(off); + } + return dp.dev.ptr + off; + } + pub fn hbytes(dp: *const Deep, off: usize, n: usize) []u8 { + return dp.host.bytes[off..][0..n]; + } +}; + +pub var deep: ?*Deep = null; + +pub const RingTelemetry = struct { owner: u64, copies: u64, dma_requested_bytes: u64, unused_copies: u64, hit_picks: u64, miss_picks: u64, hit_expert_requests: u64, miss_expert_requests: u64, attribution: ?@import("cuda_exl3_prefill_attribution.zig").Totals }; + +pub fn ringTelemetry(owner: *const anyopaque) ?RingTelemetry { + const buffers = deep orelse return null; + if (buffers.owner != owner) return null; + const ring = buffers.ring orelse return null; + return .{ .owner = ring.telemetry_id, .copies = ring.stats.copies, .dma_requested_bytes = ring.stats.dma_requested_bytes, .unused_copies = ring.stats.wasted, .hit_picks = ring.stats.hit_picks, .miss_picks = ring.stats.miss_picks, .hit_expert_requests = ring.stats.hit_expert_requests, .miss_expert_requests = ring.stats.miss_expert_requests, .attribution = if (ring.attribution) |state| state.totals else null }; +} + +pub fn now(io: std.Io) i96 { + return std.Io.Clock.awake.now(io).toNanoseconds(); +} + +pub fn upload(st: cuda.Stream, dst: u64, bytes: []const u8) !void { + try st.d.check(st.d.api.cuMemcpyHtoDAsync_v2(dst, bytes.ptr, bytes.len, st.handle), "cuMemcpyHtoDAsync"); +} + +pub fn profile() bool { + const v = std.c.getenv("TF_FLASHNEXT_PROFILE") orelse return false; + return std.mem.eql(u8, std.mem.span(v), "1"); +} diff --git a/zig/src/families/flashnext/exl3/prefill/cuda_exl3_prefill_test.zig b/zig/src/families/flashnext/exl3/prefill/cuda_exl3_prefill_test.zig new file mode 100644 index 0000000000..847e5e178e --- /dev/null +++ b/zig/src/families/flashnext/exl3/prefill/cuda_exl3_prefill_test.zig @@ -0,0 +1,302 @@ +//! Host tests of the deep prefill: layouts, arenas, window staging, passes, rank masks and buffer ownership. + +const std = @import("std"); +const cuda = @import("cuda"); +const batch = @import("../../cuda_batch.zig"); +const wt = @import("../../cuda_prompt_wt.zig"); +const history_ = @import("cuda_exl3_prefill_history.zig"); +const tp_ = @import("../tp/cuda_exl3_tp.zig"); +const window_ = @import("exl3_prefill_window.zig"); +const state = @import("cuda_exl3_prefill_state.zig"); +const hub = @import("cuda_exl3_prefill.zig"); +const chunks = @import("cuda_exl3_prefill_chunk.zig"); +const experts = @import("cuda_exl3_prefill_moe.zig"); + +const dims = wt.dims; +const wide = wt.wide; +const ple_layer = wt.ple_layer; +const table_n = wt.experts_n + 1; +const routed_n_ = wt.experts_n; +const shared_id = wt.experts_n; + +const Deep = state.Deep; +const HostLayout = state.HostLayout; +const Layout = state.Layout; +const Options = state.Options; +const Pass = state.Pass; +const Persistent = state.Persistent; +const deviceBytes = state.deviceBytes; +const deviceBytesAt = state.deviceBytesAt; +const release = hub.release; +const releaseFor = hub.releaseFor; +const reserveAt = hub.reserveAt; +const reserveTwoChunk = hub.reserveTwoChunk; +const stepCapacity = hub.stepCapacity; +const compactLast = chunks.compactLast; +const pendingWriteback = chunks.pendingWriteback; +const selectWriteback = chunks.selectWriteback; +const view = chunks.view; +const windowInputs = chunks.windowInputs; +const build = experts.build; +const maskOwned = experts.maskOwned; + +test "exl3: reduced prefill writeback preserves slot rounding sub-blocks and tail selection" { + try @import("cuda_exl3_prefill_writeback_test.zig").run(Deep, Layout.of, view, pendingWriteback, compactLast, selectWriteback); + try std.testing.expect(!(Options{}).reduced_branch); +} + +test "exl3: two-chunk scheduling and reservation contracts" { + _ = window_; + try std.testing.expect(!(Options{}).two_chunk); + try std.testing.expectEqual(@as(usize, 4096), stepCapacity(4096, false)); + try std.testing.expectEqual(@as(usize, 8192), stepCapacity(4096, true)); + try std.testing.expectEqual(@as(usize, 8192), stepCapacity(8192, true)); +} + +test "exl3: two-chunk arenas retain hidden and deferred HC operands without aliasing" { + const rows = 4096; + const layout = Layout.of(rows, 2); + const persistent = Persistent.of(rows); + try std.testing.expectEqual(@as(usize, 811285504), layout.total); + const host_layout = HostLayout.of(rows); + try std.testing.expectEqual(@as(usize, 160 << 20), 2 * rows * wide * 2); + var driver: cuda.Driver = undefined; + var first: Deep = undefined; + first.world = 2; + first.rows = rows; + first.shared = null; + first.lay = layout; + first.dev = .{ .d = &driver, .ptr = 0x100000, .len = layout.total }; + first.sum_dev = null; + var second = first; + second.dev.ptr += layout.total; + second.dev.len = persistent.total; + second.rows = rows; + second.shared = &first; + var base: batch.Scratch = undefined; + base.moe_f32 = 0; + const fields = .{ "h", "pss", "inj", "mixed", "xsm", "y", "wts", "picks", "ple_emb", "ple_xs", "ids", "pos", "members", "items", "kscratch", "sum" }; + inline for (fields, 0..) |field, index| { + const start = @field(layout, field); + const end = if (index + 1 < fields.len) @field(layout, fields[index + 1]) else layout.total; + try std.testing.expectEqual(@as(usize, 0), start % 256); + try std.testing.expect(start < end and end <= layout.total); + try std.testing.expect(first.dev.ptr + end <= second.dev.ptr); + } + for ([_]usize{ 0, 1024, 3072 }) |offset| { + const a = view(&base, &first, offset); + const b = view(&base, &second, offset); + inline for (.{ "h", "pss", "inj", "y", "wts" }) |field| { + const second_start = second.dev.ptr + @field(persistent, field); + try std.testing.expect(@field(b, field) >= second_start and @field(b, field) < second.dev.ptr + persistent.total); + try std.testing.expect(@field(a, field) < second.dev.ptr); + } + inline for (.{ "mixed", "xsm", "picks", "ple_emb", "ple_xs", "ids", "pos", "moe_f32" }) |field| { + try std.testing.expectEqual(@field(a, field), @field(b, field)); + } + } + inline for (.{ "members", "items", "kscratch", "sum" }) |field| try std.testing.expectEqual(first.dptr(@field(layout, field)), second.dptr(@field(layout, field))); + try std.testing.expectEqual(@as(usize, 546127872), persistent.total); + try std.testing.expect(persistent.total < layout.total); + try std.testing.expect(host_layout.total > rows * dims * 2); + try std.testing.expect(!window_.admits(layout.total - 1, 0, layout.total)); + try std.testing.expect(window_.admits(layout.total, 0, layout.total)); +} + +test "exl3: two-chunk memory refusal performs no CUDA allocation" { + const saved_deep = state.deep; + const saved_options = state.options; + defer { + state.deep = saved_deep; + state.options = saved_options; + } + state.options = .{ .two_chunk = true }; + var driver: cuda.Driver = undefined; + var buffers: Deep = undefined; + buffers.dev = .{ .d = &driver, .ptr = 0x100000, .len = 1 }; + buffers.world = 2; + buffers.rows = 4096; + buffers.lay = Layout.of(4096, 2); + buffers.second = null; + state.deep = &buffers; + try std.testing.expectEqual(@as(u64, 0), try reserveTwoChunk(&driver, Persistent.of(4096).total - 1, 0)); + try std.testing.expect(buffers.second == null); +} + +test "exl3: two-chunk PLE input staging advances mutable history only in chronological PLE segments" { + const Fixture = struct { + calls: usize = 0, + before: [2][2]i64 = undefined, + + fn sync(_: cuda.abi.Stream) callconv(.c) cuda.abi.Result { + return cuda.abi.success; + } + + fn copy(_: cuda.abi.DevicePtr, _: ?*const anyopaque, _: usize, _: cuda.abi.Stream) callconv(.c) cuda.abi.Result { + return cuda.abi.success; + } + + pub fn rows(self: *@This(), history: []i64, ids: []const u32, embedding: []u16, sums: []f32) !void { + self.before[self.calls] = history[0..2].*; + self.calls += 1; + for (ids) |id| { + history[0] = history[1]; + history[1] = id; + } + @memset(embedding, 0); + @memset(sums, 0); + } + }; + var driver: cuda.Driver = undefined; + driver.api.cuStreamSynchronize = Fixture.sync; + driver.api.cuMemcpyHtoDAsync_v2 = Fixture.copy; + const host_layout = HostLayout.of(256); + const host = try std.testing.allocator.alignedAlloc(u8, .@"16", host_layout.total); + defer std.testing.allocator.free(host); + var buffers: Deep = undefined; + buffers.rows = 256; + buffers.shared = null; + buffers.lay = Layout.of(256, 2); + buffers.hl = host_layout; + buffers.host = .{ .d = &driver, .bytes = host }; + buffers.dev = .{ .d = &driver, .ptr = 0x100000, .len = buffers.lay.total }; + var fixture: Fixture = .{}; + var history: [2]i64 = .{ -2, -1 }; + var pass = .{ .stream = cuda.Stream{ .d = &driver, .handle = null }, .m = .{ .ngram = &fixture }, .seq = .{ .history = &history } }; + var ids: [512]u32 = undefined; + for (&ids, 1..) |*id, value| id.* = @intCast(value); + for (0..2) |index| try windowInputs(&buffers, &pass, ids[index * 256 ..][0..256], index * 256, 0); + try std.testing.expectEqual(@as(usize, 0), fixture.calls); + try std.testing.expectEqualSlices(i64, &.{ -2, -1 }, &history); + for (0..2) |index| try windowInputs(&buffers, &pass, ids[index * 256 ..][0..256], index * 256, ple_layer); + try std.testing.expectEqual(@as(usize, 2), fixture.calls); + try std.testing.expectEqualSlices(i64, &.{ -2, -1 }, &fixture.before[0]); + try std.testing.expectEqualSlices(i64, &.{ 255, 256 }, &fixture.before[1]); + try std.testing.expectEqualSlices(i64, &.{ 511, 512 }, &history); +} + +test "exl3: prefill passes keep each expert's pairs together, items relative to their pass" { + var dp: Deep = undefined; + // pairs 0..8: experts 5, 512, 2 / 2, 512, 7 / 5, 512, 5 (-1 and out-of-table ids skipped) + const picks = [_]i32{ 5, 512, 2, 2, 512, 7, 5, 512, 5, -1, 600 }; + @memset(&dp.cnt, 0); + for (picks) |pk_| { + if (pk_ < 0 or pk_ >= table_n) continue; + dp.cnt[@intCast(pk_)] += 1; + } + // pass 0: the shared expert and 5; pass 1: 2 and 7 + dp.order[0] = shared_id; + dp.order[1] = 5; + dp.order[2] = 2; + dp.order[3] = 7; + dp.passes[0] = .{ .id_lo = 0, .id_hi = 2, .m_lo = 0, .m_n = 0, .i_lo = 0, .i_n = 0 }; + dp.passes[1] = .{ .id_lo = 2, .id_hi = 4, .m_lo = 0, .m_n = 0, .i_lo = 0, .i_n = 0 }; + dp.n_passes = 2; + var members: [16]i32 = undefined; + var items: [3 * 16]i32 = undefined; + const c = build(&dp, &picks, &members, &items); + try std.testing.expectEqual(@as(usize, 9), c.members); + try std.testing.expectEqual(@as(usize, 4), c.items); + try std.testing.expectEqualSlices(i32, &.{ 1, 4, 7, 0, 6, 8, 2, 3, 5 }, members[0..9]); + try std.testing.expectEqual(Pass{ .id_lo = 0, .id_hi = 2, .m_lo = 0, .m_n = 6, .i_lo = 0, .i_n = 2 }, dp.passes[0]); + try std.testing.expectEqual(Pass{ .id_lo = 2, .id_hi = 4, .m_lo = 6, .m_n = 3, .i_lo = 2, .i_n = 2 }, dp.passes[1]); + // (expert, first within its pass, count) + try std.testing.expectEqualSlices(i32, &.{ 512, 0, 3, 5, 3, 3, 2, 0, 2, 7, 2, 1 }, items[0..12]); +} + +test "exl3: expert-parallel prefill masks the other rank's picks and keeps the shared expert on rank 1" { + const r0 = try tp_.Rank.init(2, 0); + const r1 = try tp_.Rank.init(2, 1); + const picks = [_]i32{ 5, 300, 512, -1, 255, 256, 511, 0 }; + var a = picks; + maskOwned(&a, r0); + try std.testing.expectEqualSlices(i32, &.{ 5, -1, -1, -1, 255, -1, -1, 0 }, &a); + var b = picks; + maskOwned(&b, r1); + try std.testing.expectEqualSlices(i32, &.{ -1, 300, 512, -1, -1, 256, 511, -1 }, &b); + var c = picks; + maskOwned(&c, null); + try std.testing.expectEqualSlices(i32, &picks, &c); + // every pick lands on exactly one rank + for (picks, a, b) |p, x, y| if (p >= 0) try std.testing.expect((x == p) != (y == p)); + // two ranks add the gathered partials: one [rows, D] fp32 per rank + try std.testing.expectEqual(deviceBytesAt(4096, 1) + 2 * 4096 * 2560 * 4, deviceBytesAt(4096, 2)); +} + +test "exl3: prefetch history follows deep engine ownership reuse replacement and release" { + const Fixture = struct { + fn unload(_: cuda.abi.Module) callconv(.c) cuda.abi.Result { + return cuda.abi.success; + } + fn stream(_: cuda.abi.Stream) callconv(.c) cuda.abi.Result { + return cuda.abi.success; + } + fn freeHost(_: ?*anyopaque) callconv(.c) cuda.abi.Result { + return cuda.abi.success; + } + fn refuse(_: *cuda.abi.DevicePtr, _: usize) callconv(.c) cuda.abi.Result { + return cuda.abi.error_not_ready; + } + fn install(driver: *const cuda.Driver, owner: *const anyopaque) !*Deep { + const buffers = try std.testing.allocator.create(Deep); + buffers.* = .{ + .gpa = std.testing.allocator, + .owner = owner, + .rows = 4096, + .lay = undefined, + .hl = undefined, + .dev = .{ .d = driver, .ptr = 0, .len = 0 }, + .host = .{ .d = driver, .bytes = @constCast(&.{}) }, + .kern = undefined, + .aux = .{ .d = driver, .handle = null }, + }; + buffers.kern.mod = .{ .d = driver, .handle = null }; + state.deep = buffers; + return buffers; + } + }; + try std.testing.expect(state.deep == null); + defer release(); + var driver: cuda.Driver = undefined; + driver.api.cuModuleUnload = Fixture.unload; + driver.api.cuStreamSynchronize = Fixture.stream; + driver.api.cuStreamDestroy_v2 = Fixture.stream; + driver.api.cuMemFreeHost = Fixture.freeHost; + driver.api.cuMemAlloc_v2 = Fixture.refuse; + var owners: [2]u8 = .{ 0, 0 }; + const first = try Fixture.install(&driver, &owners[0]); + first.history[3] = .{ .rows = 4096, .used = 80 }; + first.history_shape = .{ .sub = 256, .lo = 0, .hi = 256 }; + releaseFor(&owners[1]); + try std.testing.expect(state.deep.? == first); + try reserveAt(std.testing.allocator, &driver, &owners[0], 2048, 1); + try std.testing.expectEqual(@as(usize, 96), history_.want(state.deep.?.history[3], 4096, 400, routed_n_)); + try std.testing.expectError(error.NotReady, reserveAt(std.testing.allocator, &driver, &owners[1], 4096, 1)); + try std.testing.expect(state.deep == null); + const second = try Fixture.install(&driver, &owners[1]); + try std.testing.expectEqual(@as(usize, 416), history_.want(second.history[3], 4096, 400, routed_n_)); + try std.testing.expect(second.history_shape == null); + second.history[3] = .{ .rows = 4096, .used = 80 }; + try std.testing.expectError(error.NotReady, reserveAt(std.testing.allocator, &driver, &owners[1], 4096, 2)); + try std.testing.expect(state.deep == null); + const third = try Fixture.install(&driver, &owners[1]); + try std.testing.expectEqual(@as(usize, 416), history_.want(third.history[3], 4096, 400, routed_n_)); + third.history[3] = .{ .rows = 4096, .used = 80 }; + releaseFor(&owners[1]); + try std.testing.expect(state.deep == null); + const fourth = try Fixture.install(&driver, &owners[1]); + try std.testing.expectEqual(@as(usize, 416), history_.want(fourth.history[3], 4096, 400, routed_n_)); +} + +test "exl3: the prefill tiles cut K as decode does at Flash Next's shapes, and the two expert groups split at 256" { + const pk = @import("cuda_exl3_prefill_k.zig"); + try std.testing.expect(pk.sumsAsDecode(dims, wt.moe_w)); + try std.testing.expect(!pk.sumsAsDecode(dims + 64, wt.moe_w + 64)); // no plan of four 8-tile warps: refused + // decode's plan here: gate/up in 4 splits of 4 warps, down in 4 warps (chains of 10 k tiles) + try std.testing.expectEqual(@as(u32, 4), try cuda.exl3_experts.plan(dims, wt.moe_w, true)); + try std.testing.expectEqual(@as(u32, 1), try cuda.exl3_experts.plan(wt.moe_w, dims, false)); +} + +test "exl3: leaves the any-degree reduction tests run with the prefill tests" { + _ = @import("cuda_exl3_prefill_leaves.zig"); +} diff --git a/zig/src/families/flashnext/exl3/prefill/cuda_exl3_prefill_tiles.zig b/zig/src/families/flashnext/exl3/prefill/cuda_exl3_prefill_tiles.zig new file mode 100644 index 0000000000..b32cf381d4 --- /dev/null +++ b/zig/src/families/flashnext/exl3/prefill/cuda_exl3_prefill_tiles.zig @@ -0,0 +1,35 @@ +//! Dense EXL3 projections of a deep prefill on the prompt tiles (the hook cuda_exl3_pass.zig calls), in decode's bytes. + +const cuda = @import("cuda"); +const wt = @import("../../cuda_prompt_wt.zig"); +const prompt = @import("../../cuda_prompt.zig"); + +/// Set while a deep prefill runs with Options.tiles: cuda_exl3_pass.part sends prompt rows to Dense.prompt on tiles. +pub var tiles_on: bool = false; +/// Rows below which a projection keeps the reference kernels (tiles lose at a few rows). +const tiles_min_rows: usize = 64; + +/// The hook in cuda_exl3_pass.part: true when it ran the part on the tiles (bf16 or fp32 out). +pub fn tilesPart(de: *cuda.exl3.Dense, s: cuda.Stream, w: cuda.exl3.Weight, x: u64, m: usize, y: u64, f32_out: bool) !bool { + if (!tiles_on or m < tiles_min_rows) return false; + _ = cuda.exl3.Dense.candidateWeight(w) catch return false; + if (de.candidate_k == null or de.candidate_k.?.k_max < w.k) return false; + try de.promptExact(s, w, x, m, y, if (f32_out) .f32 else .bf16); + return true; +} + +/// Loads the tiles once with the widest input any EXL3 projection has (no allocation inside a pass). +pub fn warmTiles(e: *prompt.Engine) !void { + const x = e.layers[0].exl3 orelse return; + const de = try x.ops.linear(); + const k_max = @max(x.ops.max_k, wt.dims * 3); + if (de.candidate_k != null and de.candidate_k.?.k_max >= k_max) return; + if (!cuda.kernels.available) return error.BuiltWithoutKernels; + var mod = try cuda.Module.load(de.d, cuda.kernels.exl3_candidate); + errdefer mod.unload(); + const k = try cuda.exl3_candidate.Kernels.resolve(de.d, mod, k_max); + if (de.candidate_k) |*old| old.deinit(); + if (de.candidate_mod) |*old| old.unload(); + de.candidate_mod = mod; + de.candidate_k = k; +} diff --git a/zig/src/families/flashnext/exl3/prefill/cuda_exl3_prefill_writeback_test.zig b/zig/src/families/flashnext/exl3/prefill/cuda_exl3_prefill_writeback_test.zig new file mode 100644 index 0000000000..052763db79 --- /dev/null +++ b/zig/src/families/flashnext/exl3/prefill/cuda_exl3_prefill_writeback_test.zig @@ -0,0 +1,131 @@ +const std = @import("std"); +const cuda = @import("cuda"); +const batch = @import("../../cuda_batch.zig"); +const wt = @import("../../cuda_prompt_wt.zig"); +const dims = wt.dims; +const slots = wt.slots; +const streams = wt.streams; + +const Call = struct { branch: u64, y: u64, weights: u64, mode: usize, topk: usize, slots: usize, rows: usize }; + +const Tri = struct { + call: *Call, + + pub fn hcWriteback(self: Tri, _: u64, _: u64, _: u64, branch: u64, _: []const u8, _: u64, y: u64, _: []const u8, weights: u64, _: []const u8, _: usize, rows: usize, _: usize, _: usize, mode: usize, topk: usize, count: usize, _: usize) !void { + self.call.* = .{ .branch = branch, .y = y, .weights = weights, .mode = mode, .topk = topk, .slots = count, .rows = rows }; + } +}; + +const Copies = struct { + s: *batch.Scratch, + calls: usize = 0, + last: struct { dst: u64, src: u64, bytes: usize } = undefined, + + pub fn copy(self: *Copies, dst: u64, src: u64, bytes: usize) !void { + self.calls += 1; + self.last = .{ .dst = dst, .src = src, .bytes = bytes }; + } +}; + +pub fn run(comptime Deep: type, comptime layout: anytype, comptime view: anytype, comptime pending: anytype, comptime compact: anytype, comptime select: anytype) !void { + var driver: cuda.Driver = undefined; + var dp: Deep = undefined; + dp.shared = null; + dp.world = 2; + dp.rows = 7; + dp.lay = layout(dp.rows, dp.world); + dp.dev = .{ .d = &driver, .ptr = 0x100000, .len = dp.lay.total }; + dp.reduced_branch = false; + dp.unit_weights = .{ .d = &driver, .ptr = 0x900000, .len = dp.rows * 4 }; + try select(&dp, true); + try std.testing.expect(dp.reduced_branch); + var base: batch.Scratch = undefined; + base.dummy = 0xdead; + var call: Call = undefined; + var pass: batch.Pass(Tri) = undefined; + pass.tri = .{ .call = &call }; + const sum = dp.dev.ptr + dp.lay.sum; + const sub: usize = 3; + for (0..3) |block| { + const at = block * sub; + const rows = @min(sub, dp.rows - at); + var scratch = view(&base, &dp, at); + pass.s = &scratch; + pass.rows = rows; + try std.testing.expectEqual(sum + at * dims * 4, scratch.moe_f32); + pending(&pass, &dp, null, at); + try std.testing.expect(!pass.pending_f32); + pending(&pass, &dp, 4, at); + try pass.writeback(0, rows); + try std.testing.expectEqual(Call{ .branch = base.dummy, .y = sum + at * dims * 4, .weights = dp.unit_weights.?.ptr + at * 4, .mode = 2, .topk = 0, .slots = 1, .rows = rows }, call); + try std.testing.expect(pass.pending == null and !pass.pending_f32 and pass.pending_f32_weights == null); + if (rows > 1) { + pending(&pass, &dp, 4, at); + try pass.writeback(1, rows - 1); + try std.testing.expectEqual(sum + (at + 1) * dims * 4, call.y); + try std.testing.expectEqual(dp.unit_weights.?.ptr + (at + 1) * 4, call.weights); + } + var copies: Copies = .{ .s = &scratch }; + try compact(&copies, true, rows - 1); + try std.testing.expectEqual(@as(usize, if (rows > 1) 4 else 0), copies.calls); + if (rows > 1) { + try std.testing.expectEqual(scratch.moe_f32, copies.last.dst); + try std.testing.expectEqual(scratch.moe_f32 + (rows - 1) * dims * 4, copies.last.src); + try std.testing.expectEqual(dims * 4, copies.last.bytes); + } + } + var scratch = view(&base, &dp, 0); + pass.s = &scratch; + pass.rows = sub; + try select(&dp, false); + pending(&pass, &dp, 4, 0); + try pass.writeback(0, sub); + try std.testing.expectEqual(@as(usize, slots), call.slots); + try std.testing.expectEqual(@as(usize, slots - 1), call.topk); + try std.testing.expectEqual(scratch.y, call.y); + pass.pending = 4; + pass.pending_f32 = true; + pass.pending_f32_weights = null; + try pass.writeback(0, sub); + try std.testing.expectEqual(@as(usize, 4), call.mode); + try std.testing.expectEqual(scratch.moe_f32, call.branch); + var copies: Copies = .{ .s = &scratch }; + try compact(&copies, false, 2); + try std.testing.expectEqual(@as(usize, 5), copies.calls); + try std.testing.expectEqual(scratch.wts, copies.last.dst); + try std.testing.expectEqual(scratch.wts + 2 * slots * 4, copies.last.src); + dp.world = 1; + base.moe_f32 = 0xbeef; + const single_rank = view(&base, &dp, sub); + try std.testing.expectEqual(base.moe_f32, single_rank.moe_f32); + try rounding(); +} + +fn bf16(value: f32) u16 { + const bits: u32 = @bitCast(value); + return @truncate((bits +% (0x7fff + ((bits >> 16) & 1))) >> 16); +} + +fn promote(value: u16) f32 { + return @bitCast(@as(u32, value) << 16); +} + +fn rounding() !void { + const values = [_]f32{ 0.0, -0.0, 0x1p-126, -0x1p-126, 0x1p-133, -0x1p-133, 0x1.01p0, -0x1.01p0, 0x1.03p0, -0x1.03p0, 0x1p40, -0x1p40 }; + const hidden = [_]u16{ 0, 0x8000, 0x3b34, 0xbe16, 0x3f80, 0xbf80 }; + const inject = [_]u16{ 0, 0x8000, 0x3c9a, 0x3f80, 0xbf80 }; + for (values) |sum| { + const direct = @mulAdd(f32, sum, 1, 0); + var old = direct; + for (1..slots) |slot| old = @mulAdd(f32, if (slot % 2 == 0) @as(f32, -7) else 13, 0, old); + for (hidden) |bits| for (inject) |gate| { + const want = bf16(@mulAdd(f32, promote(bf16(old)), promote(gate), promote(bits))); + const got = bf16(@mulAdd(f32, promote(bf16(direct)), promote(gate), promote(bits))); + try std.testing.expectEqual(want, got); + }; + } + const wrong = bf16(@mulAdd(f32, promote(bf16(-@as(f32, 0.0))), 1, -@as(f32, 0.0))); + const correct = bf16(@mulAdd(f32, @mulAdd(f32, -0.0, 1, 0), 1, -@as(f32, 0.0))); + try std.testing.expectEqual(@as(u16, 0x8000), wrong); + try std.testing.expectEqual(@as(u16, 0), correct); +} diff --git a/zig/src/families/flashnext/exl3/prefill/exl3_prefill_window.zig b/zig/src/families/flashnext/exl3/prefill/exl3_prefill_window.zig new file mode 100644 index 0000000000..353a4b84a0 --- /dev/null +++ b/zig/src/families/flashnext/exl3/prefill/exl3_prefill_window.zig @@ -0,0 +1,96 @@ +const std = @import("std"); + +pub const max_rows = 4096; + +pub fn enabled(value: ?[]const u8) bool { + return if (value) |text| std.mem.eql(u8, text, "1") else false; +} + +pub fn admits(free: u64, protected: u64, extra: u64) bool { + const required = std.math.add(u64, protected, extra) catch return false; + return extra != 0 and free >= required; +} + +pub const Segment = struct { offset: usize, rows: usize }; +pub const Step = struct { layer: usize, chunk: usize }; + +pub const Plan = struct { + segments: [2]Segment, + count: usize, + rows: usize, + + pub fn init(remaining: usize, chunk_rows: usize, alignment: usize) !Plan { + if (remaining == 0 or alignment == 0 or chunk_rows == 0 or chunk_rows > max_rows or chunk_rows % alignment != 0) return error.InvalidPrefillWindow; + const first = @min(remaining, chunk_rows); + const second = @min(remaining - first, chunk_rows); + return .{ .segments = .{ .{ .offset = 0, .rows = first }, .{ .offset = first, .rows = second } }, .count = if (second == 0) 1 else 2, .rows = first + second }; + } + + pub fn iterator(plan: Plan, layers: usize) Iterator { + return .{ .count = plan.count, .layers = layers }; + } +}; + +pub const Iterator = struct { + count: usize, + layers: usize, + layer: usize = 0, + chunk: usize = 0, + + pub fn next(iterator: *Iterator) ?Step { + if (iterator.layer >= iterator.layers) return null; + const step: Step = .{ .layer = iterator.layer, .chunk = iterator.chunk }; + iterator.chunk += 1; + if (iterator.chunk == iterator.count) { + iterator.chunk = 0; + iterator.layer += 1; + } + return step; + } +}; + +test "exl3: two chunk windows require explicit opt-in and bounded aligned rows" { + for ([_]?[]const u8{ null, "", "0", "true", "2" }) |value| try std.testing.expect(!enabled(value)); + try std.testing.expect(enabled("1")); + for ([_]usize{ 0, 1, 255, 4097, 8192 }) |rows| try std.testing.expectError(error.InvalidPrefillWindow, Plan.init(16384, rows, 256)); + try std.testing.expectError(error.InvalidPrefillWindow, Plan.init(0, 4096, 256)); + try std.testing.expectError(error.InvalidPrefillWindow, Plan.init(4096, 4096, 0)); +} + +test "exl3: two chunk scheduling keeps each layer chronological and retains chunk boundaries" { + for ([_]usize{ 1, 255, 256, 4095, 4096, 4097, 8192, 8193, 16384, 16641 }) |length| { + var position: [4]usize = @splat(0); + var cursor: usize = 0; + var windows: usize = 0; + while (cursor < length) { + const plan = try Plan.init(length - cursor, 4096, 256); + try std.testing.expect(plan.rows <= 8192); + var steps = plan.iterator(position.len); + var previous: ?Step = null; + while (steps.next()) |step| { + const segment = plan.segments[step.chunk]; + try std.testing.expectEqual(position[step.layer], cursor + segment.offset); + if (previous) |before| { + if (before.layer == step.layer) try std.testing.expectEqual(before.chunk + 1, step.chunk) else { + try std.testing.expectEqual(before.layer + 1, step.layer); + try std.testing.expectEqual(@as(usize, 0), step.chunk); + } + } + position[step.layer] += segment.rows; + previous = step; + } + cursor += plan.rows; + windows += 1; + } + for (position) |end| try std.testing.expectEqual(length, end); + try std.testing.expectEqual((length + 8191) / 8192, windows); + } +} + +test "exl3: two chunk admission protects all existing reservations and rejects overflow" { + try std.testing.expect(admits(100, 60, 40)); + try std.testing.expect(!admits(99, 60, 40)); + try std.testing.expect(!admits(100, 101, 1)); + try std.testing.expect(!admits(100, 0, 0)); + try std.testing.expect(!admits(std.math.maxInt(u64), std.math.maxInt(u64), 1)); +} diff --git a/zig/tests/cuda/contract_flashnext_exl3_prefill.zig b/zig/tests/cuda/contract_flashnext_exl3_prefill.zig new file mode 100644 index 0000000000..1b320befd3 --- /dev/null +++ b/zig/tests/cuda/contract_flashnext_exl3_prefill.zig @@ -0,0 +1,420 @@ +//! Flash Next EXL3 prefill experts under the contract: tf_exl3pf_gu and tf_exl3pf_down replace the decode form. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const flash = @import("flashnext"); +const check = @import("check.zig"); +const rig_ = @import("contract_rig.zig"); + +const c = core.contract.spec; +const g = core.contract.gates; +const pk_ = flash.exl3.prefill_k; +const exl3e = cuda.exl3_experts; + +/// The gate-up pass (rotated inputs, SwiGLU outputs), the down pass (y), and the gate-up pass against a wrong baseline. +pub const Kind = enum { gu, down, control }; + +const dims = 2560; +const width = 640; +const table_n = 16; // 15 routed experts and the shared one as id 15 +const slots = 11; // ten routed picks and the shared expert a row +const max_rows = 1408; +const tiles = (dims / 16) * (width / 16); +const k2s = [3]u32{ 6, 6, 8 }; // half-bits a weight of gate, up and down: 3, 3 and 4 bits +const rule = "as the decode form it replaces gives them: a NaN or infinite activation reaches every value its row's picks make"; + +fn depth(s: c.Shape) usize { + return s.k / 16 + 16; +} + +// Zero trellis words decode to one dyadic weight; the stress dots are exact FP32 integers. +const stress_svh: f16 = -0x1.094p-14; +const stress_suh: f16 = 0.1104736328125; +const had: f64 = 1.0 / @sqrt(128.0); + +/// Rows of 11 picks over 640-wide experts: the gate-up pass reads 2,560 inputs, the down pass 640. +pub fn spec(kind: Kind) c.Spec { + const up = kind != .down; + const common = .{ .accum = c.Accum.f32, .depth = depth, .nan = c.Nan{ .stated = rule }, .tiles = [3]c.Tile{ .{ .size = 6 }, .{}, .{} } }; + const served: []const c.Shape = if (up) &.{ .{ .rows = 130, .n = width, .k = dims }, .{ .rows = 1, .n = width, .k = dims }, .{ .rows = max_rows, .n = width, .k = dims } } else &.{ .{ .rows = 130, .n = dims, .k = width }, .{ .rows = 1, .n = dims, .k = width }, .{ .rows = max_rows, .n = dims, .k = width } }; + return .{ + .name = switch (kind) { + .gu => "cuda Flash Next EXL3 prefill rotate and gate-up tf_exl3pf_gu (decode sums)", + .down => "cuda Flash Next EXL3 prefill down tf_exl3pf_down (decode sums)", + .control => "control: tf_exl3pf_gu against a decode form with the other SwiGLU roundings", + }, + .accum = common.accum, + .depth = common.depth, + .out = if (up) .raw else .f32, + .replaces = if (up) "grouped gate/up and SwiGLU of exl3_experts" else "grouped down of exl3_experts", + .nan = common.nan, + // Down has one rounded Hadamard normalization constant and one FP32 multiply. + .ops = if (up) 0 else 2, + .tiles = common.tiles, + .served = served, + }; +} + +fn align256(n: usize) usize { + return (n + 255) / 256 * 256; +} + +/// One layer of random trellis words (any bits decode), random signs and scales: 15 routed experts and the shared one. +const Weights = struct { + mem: cuda.DeviceBuffer, + tab: cuda.DeviceBuffer, + scales: [6]cuda.DeviceBuffer, + layer: exl3e.Layer, + + fn init(gpu: check.Gpu, seed: u64, stress: bool) !Weights { + const a = gpu.gpa; + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + var made: usize = 0; + var sc: [6]cuda.DeviceBuffer = undefined; + errdefer for (sc[0..made]) |*b| b.free(); + for (&sc, [6]usize{ dims, dims, width, width, width, dims }, 0..) |*b, n, si| { + const v = try a.alloc(u16, table_n * n); + defer a.free(v); + for (v) |*h| { + const m: f16 = @floatCast(0.5 + rng.float(f32)); + h.* = @bitCast(if (stress) switch (si) { + 2, 3 => stress_svh, + 4 => stress_suh, + else => @as(f16, 1), + } else if (rng.boolean()) m else -m); + } + b.* = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(v)); + made += 1; + } + var total: usize = 0; + for (k2s) |k2| total += align256(tiles * 16 * k2) * table_n; + const host = try a.alloc(u8, total); + defer a.free(host); + if (stress) @memset(host, 0) else rng.bytes(host); + var mem = try cuda.DeviceBuffer.fromHost(gpu.d, host); + errdefer mem.free(); + var ptrs: [3][table_n]u64 = undefined; + var k2v: [3][table_n]i32 = undefined; + var at: usize = 0; + for (0..table_n) |e| for (0..3) |p| { + ptrs[p][e] = mem.ptr + at; + k2v[p][e] = @intCast(k2s[p]); + at += align256(tiles * 16 * k2s[p]); + }; + const tl = cuda.exl3_pool.Tables.init(table_n); + const tb = try a.alloc(u8, tl.bytes); + defer a.free(tb); + @memset(tb, 0); + for (0..3) |p| { + @memcpy(tb[tl.ptr[p]..][0 .. table_n * 8], std.mem.sliceAsBytes(&ptrs[p])); + @memcpy(tb[tl.k2[p]..][0 .. table_n * 4], std.mem.sliceAsBytes(&k2v[p])); + } + const tab = try cuda.DeviceBuffer.fromHost(gpu.d, tb); + return .{ .mem = mem, .tab = tab, .scales = sc, .layer = .{ + .gate = tab.ptr + tl.ptr[0], + .up = tab.ptr + tl.ptr[1], + .down = tab.ptr + tl.ptr[2], + .gate_k2 = tab.ptr + tl.k2[0], + .up_k2 = tab.ptr + tl.k2[1], + .down_k2 = tab.ptr + tl.k2[2], + .suh_g = sc[0].ptr, + .suh_u = sc[1].ptr, + .svh_g = sc[2].ptr, + .svh_u = sc[3].ptr, + .suh_d = sc[4].ptr, + .svh_d = sc[5].ptr, + .experts = table_n, + .dims = dims, + .width = width, + .codebook = .mul1, + .act_mode = 1, + } }; + } + + fn deinit(w: *Weights) void { + w.tab.free(); + w.mem.free(); + for (&w.scales) |*b| b.free(); + } +}; + +pub const Prefill = struct { + rig: rig_.Rig, + kind: Kind, + w: *const Weights, + sw: *const Weights, + stress: bool = false, + k: *const pk_.Kernels, + ex: *const exl3e.Experts, + shape: c.Shape = undefined, + x: []u16 = &.{}, + picks: []i32 = &.{}, + plan: pk_.Plan = undefined, + planned: bool = false, + want: []u8 = &.{}, + xd: u64 = 0, + pd: u64 = 0, + dv: pk_.Device = undefined, + scr: u64 = 0, + y: u64 = 0, + base: u64 = 0, + out_len: usize = 0, + msg: [160]u8 = undefined, + + pub fn deinit(self: *Prefill) void { + self.free(); + self.rig.deinit(); + } + + fn free(self: *Prefill) void { + const a = self.rig.gpu.gpa; + a.free(self.x); + a.free(self.picks); + a.free(self.want); + if (self.planned) self.plan.deinit(a); + self.x = &.{}; + self.picks = &.{}; + self.want = &.{}; + self.planned = false; + } + + pub fn kernel(self: *Prefill) g.Kernel { + return .{ .spec = spec(self.kind), .ctx = self, .vt = &vt, .control = if (self.kind == .control) .shapes else null }; + } + + const vt = g.hooks(Prefill); + + fn up(self: *const Prefill) bool { + return self.kind != .down; + } + + /// Random bf16 rows, ten distinct routed picks a row and the shared expert; non-finite: a NaN and an infinity. + pub fn prepare(self: *Prefill, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + const n: usize = if (self.up()) width else dims; + const k: usize = if (self.up()) dims else width; + if (shape.rows == 0 or shape.rows > max_rows or shape.n != n or shape.k != k) return false; + self.free(); + self.rig.reset(); + const a = self.rig.gpu.gpa; + self.shape = shape; + self.stress = fill == .stress; + const rows = shape.rows; + const pairs = rows * slots; + var prng = std.Random.DefaultPrng.init(seed ^ 0x70667830); + const rng = prng.random(); + self.x = try a.alloc(u16, rows * dims); + for (self.x, 0..) |*v, i| { + const col = i % dims; + v.* = if (self.stress) c.toBf16(if (col % 128 != 0 or col >= dims - 128) 0 else if (col < dims / 2) 11584 else -11584) else c.toBf16((rng.float(f32) - 0.5) * 2); + } + if (fill == .nonfinite) { + self.x[(rows - 1) * dims] = c.toBf16(std.math.nan(f32)); + if (rows > 1) self.x[(rows - 2) * dims + 1] = c.toBf16(std.math.inf(f32)); + } + self.picks = try a.alloc(i32, pairs); + for (0..rows) |r| { + var ids: [table_n - 1]i32 = undefined; + for (&ids, 0..) |*v, i| v.* = @intCast(i); + rng.shuffle(i32, &ids); + @memcpy(self.picks[r * slots ..][0 .. slots - 1], ids[0 .. slots - 1]); + self.picks[r * slots + slots - 1] = table_n - 1; + } + self.plan = try pk_.Plan.init(a, pairs, table_n); + self.planned = true; + var counts: [table_n]u32 = undefined; + pk_.build(self.picks, null, &counts, &self.plan); + const nm = self.plan.n_members; + self.xd = try self.rig.input(std.mem.sliceAsBytes(self.x)); + self.pd = try self.rig.input(std.mem.sliceAsBytes(self.picks)); + const sb = pk_.scratchBytes(nm, dims, width); + self.dv = .{ .members = try self.rig.input(self.plan.memberBytes()), .items = try self.rig.input(self.plan.itemBytes()), .scratch = 0 }; + if (self.up()) { + self.scr = try self.rig.output(sb); + self.y = try self.rig.scratch(pairs * dims * 4); + self.out_len = sb; + } else { + self.scr = try self.rig.scratch(sb); + self.y = try self.rig.output(pairs * dims * 4); + self.out_len = pairs * dims * 4; + } + self.dv.scratch = self.scr; + self.base = try self.rig.scratch(try exl3e.scratchBytes(self.layer(), rows, slots)); + return true; + } + + fn layer(self: *const Prefill) exl3e.Layer { + return if (self.stress) self.sw.layer else self.w.layer; + } + + pub fn run(self: *Prefill) !void { + if (self.stress and !self.up()) { + // Isolate the down sum from the nonlinear GU epilogue with exact half inputs. + const a = self.rig.gpu.gpa; + const x = try a.alloc(f16, self.plan.n_members * width); + defer a.free(x); + @memset(x, 1024); + const xd = try self.rig.input(std.mem.sliceAsBytes(x)); + const l = self.layer(); + var dn: cuda.Args = .{}; + dn.add(xd); + dn.add(l.down); + dn.add(l.down_k2); + dn.add(self.dv.items); + dn.add(self.dv.members); + dn.add(l.svh_d); + dn.add(self.y); + dn.add(@as(i32, width)); + dn.add(@as(i32, dims)); + try cuda.launch.launch(self.k.down, .{ .grid = .{ .x = dims / 256, .y = @intCast(self.plan.n_items) }, .block = .{ .x = 256 } }, self.rig.s, &dn); + } else try pk_.run(self.k, self.rig.s, self.layer(), self.xd, dims, self.pd, slots, self.plan.n_members, self.plan.n_items, self.dv, self.y); + try self.rig.s.synchronize(); + } + + pub fn count(self: *const Prefill) usize { + return self.out_len / c.size(spec(self.kind).out); + } + + pub fn reference(self: *Prefill, out: []f64, mag: []f64) !void { + try self.stressOutput(.none, .sequential, out); + for (mag, out) |*m, v| m.* = @abs(v); + } + + pub fn simulate(self: *Prefill, acc: c.Accum, order: g.Order, out: []f64) !void { + try self.stressOutput(acc, order, out); + } + + // This reference decodes the zero-state weight in FP64, not through either GPU kernel. + fn stressOutput(self: *Prefill, acc: c.Accum, order: g.Order, out: []f64) !void { + if (!self.stress) return error.ExpectedStressInputs; + const w: f64 = core.exl3_format.codebookValue(.mul1, 0); + var zeros: [128]i16 = @splat(0); + var decoded: [256]f16 = undefined; + for (k2s) |k2| { + const bits: core.exl3_format.Width = .{ .halves = @intCast(k2) }; + core.exl3_format.unpack(bits, .mul1, zeros[0..bits.tileWords()], 1, 1, &decoded); + for (decoded) |v| if (v != w) return error.StressWeightIsNotConstant; + } + const rotated: f16 = @floatCast(11584 * had); + if (rotated != 1024 or @abs(w * 1024) * dims >= 0x1p24) return error.StressIsNotExactFP32; + var terms: [dims]f64 = undefined; + const n: usize = if (self.up()) dims else width; + for (terms[0..n], 0..) |*v, i| v.* = w * (if (!self.up()) @as(f64, 1024) else @as(f64, if (i >= dims - 128) 0 else if (i < dims / 2) 1024 else -1024)); + const dot = core.contract.rows.sumIn(acc, order, terms[0..n]); + if (!self.up()) { + // Constant dot rows transform to a single nonzero value per 128-column block. + for (out, 0..) |*v, i| v.* = if (i % 128 == 0) dot / had else 0; + return; + } + const gg = dot / had * @as(f64, stress_svh); + const act = gg / (1 + @exp(-gg)) * gg; + const raw = act * @as(f64, stress_suh) * had; + // A wide margin to both half midpoints absorbs the FP32 nonlinear epilogue's roundings. + if (acc == .none and !(raw > 1023.9 and raw < 1024.1)) return error.StressHalfMargin; + const xd: f16 = @floatCast(raw); + const a = self.rig.gpu.gpa; + const bytes = try a.alloc(u8, self.out_len); + defer a.free(bytes); + @memset(bytes, 0xff); + const nm = self.plan.n_members; + const xu_at = align256(nm * dims * 2); + for (0..2) |p| for (0..nm * dims) |i| { + const col = i % dims; + const xh: f16 = if (col >= dims - 128) 0 else if (col < dims / 2) 1024 else -1024; + std.mem.writeInt(u16, bytes[p * xu_at + i * 2 ..][0..2], @bitCast(xh), .little); + }; + for (0..nm * width) |i| std.mem.writeInt(u16, bytes[2 * xu_at + i * 2 ..][0..2], @bitCast(xd), .little); + for (out, bytes) |*v, b| v.* = @floatFromInt(b); + } + + /// The bytes of `n` rows of `row` bytes at device `ptr`, member m taking the decode form's row of pair members[m]. + fn gather(self: *Prefill, dst: []u8, ptr: u64, row: usize) !void { + const src = try self.rig.bytes(ptr, self.shape.rows * slots * row); + for (self.plan.members[0..self.plan.n_members], 0..) |pair, m| @memcpy(dst[m * row ..][0..row], src[@as(usize, @intCast(pair)) * row ..][0..row]); + } + + /// The decode form (runSlots) on the same picks, in this pass's bytes: xg, xu, xd by member, or y. + pub fn replaced(self: *Prefill) ![]const u8 { + const a = self.rig.gpu.gpa; + var l = self.layer(); + if (self.kind == .control) l.act_mode = 0; + const sc = try exl3e.scratch(self.base, l, self.shape.rows, slots); + try self.ex.runSlots(self.rig.s, l, self.xd, dims, .bf16, self.pd, self.shape.rows, sc); + try self.rig.s.synchronize(); + a.free(self.want); + self.want = try a.alloc(u8, self.out_len); + @memset(self.want, 0xff); + if (!self.up()) { + @memcpy(self.want, try self.rig.bytes(sc.y, self.out_len)); + return self.want; + } + const nm = self.plan.n_members; + const xu_at = align256(nm * dims * 2); + try self.gather(self.want[0 .. nm * dims * 2], sc.xg, dims * 2); + try self.gather(self.want[xu_at..][0 .. nm * dims * 2], sc.xu, dims * 2); + try self.gather(self.want[2 * xu_at ..][0 .. nm * width * 2], sc.xd, width * 2); + return self.want; + } + + /// Gate 1: the decode form's bytes, and some output non-finite (the planted values reach it). + pub fn check(self: *Prefill, _: g.Gate) !?[]const u8 { + const got = try self.rig.outputBytes(); + const want = try self.replaced(); + if (std.mem.indexOfDiff(u8, got, want)) |at| return try std.fmt.bufPrint(&self.msg, "byte {d} of {d} differs from {s}'s", .{ at, got.len, spec(self.kind).replaces.? }); + if (self.up()) { + for (std.mem.bytesAsSlice(u16, got[0 .. self.plan.n_members * dims * 2])) |h| if (h & 0x7c00 == 0x7c00) return null; + } else for (std.mem.bytesAsSlice(f32, got)) |v| if (!std.math.isFinite(v)) return null; + return try std.fmt.bufPrint(&self.msg, "no output is non-finite: the planted values prove nothing", .{}); + } +}; + +/// The three tests over one layer of random weights and the decode form's experts. +pub const Group = struct { + a: std.mem.Allocator, + w: Weights, + sw: Weights, + k: pk_.Kernels, + ex: exl3e.Experts, + entries: [3]Prefill, + + pub fn create(a: std.mem.Allocator, gpu: check.Gpu, s: cuda.Stream, kernels: *std.ArrayList(g.Kernel)) !*Group { + const self = try a.create(Group); + errdefer a.destroy(self); + self.a = a; + self.w = try Weights.init(gpu, 0x65786c33, false); + errdefer self.w.deinit(); + self.sw = try Weights.init(gpu, 0, true); + errdefer self.sw.deinit(); + self.k = try pk_.Kernels.load(gpu.d); + errdefer self.k.unload(); + self.ex = try exl3e.Experts.load(gpu.d); + errdefer self.ex.unload(); + for (&self.entries, std.enums.values(Kind)) |*x, kind| x.* = .{ .rig = .{ .gpu = gpu, .s = s }, .kind = kind, .w = &self.w, .sw = &self.sw, .k = &self.k, .ex = &self.ex }; + errdefer for (&self.entries) |*x| x.deinit(); + for (&self.entries) |*x| try kernels.append(a, x.kernel()); + return self; + } + + pub fn destroy(self: *Group) void { + for (&self.entries) |*x| x.deinit(); + self.ex.unload(); + self.k.unload(); + self.w.deinit(); + self.sw.deinit(); + self.a.destroy(self); + } +}; + +/// tf-cuda-test exl3-prefill-contract: just these kernels through the gates and the control (the full runner is hours). +pub fn run(gpu: check.Gpu) !void { + var s = try cuda.Stream.init(gpu.d, false); + defer s.deinit(); + var list: std.ArrayList(g.Kernel) = .empty; + defer list.deinit(gpu.gpa); + const group = try Group.create(gpu.gpa, gpu, s, &list); + defer group.destroy(); + var t: g.Tally = .{}; + try g.runAll(gpu.gpa, list.items, 0x6b63, &t); + std.debug.print("SUMMARY {d} kernels: {d} pass, {d} fail, {d} known; {d} controls: {d} red, {d} missed\n", .{ t.kernels, t.passed, t.failed, t.known, t.controls, t.caught, t.missed }); + try check.expect(t.problems() == 0, "exl3-prefill-contract: {d} gate failures, {d} controls missed", .{ t.failed, t.missed }); +} diff --git a/zig/tests/cuda/contract_flashnext_exl3_suite.zig b/zig/tests/cuda/contract_flashnext_exl3_suite.zig index 3cb954e766..376dfcd33e 100644 --- a/zig/tests/cuda/contract_flashnext_exl3_suite.zig +++ b/zig/tests/cuda/contract_flashnext_exl3_suite.zig @@ -4,10 +4,12 @@ const cuda = @import("cuda"); const check = @import("check.zig"); const g = @import("core").contract.gates; const tp_ = @import("contract_flashnext_exl3_tp.zig"); +const prefill_ = @import("contract_flashnext_exl3_prefill.zig"); pub const Suite = struct { a: std.mem.Allocator, masks: tp_.Masks, + prefill: *prefill_.Group, pub fn create(a: std.mem.Allocator, gpu: check.Gpu, s: cuda.Stream, kernels: *std.ArrayList(g.Kernel)) !*Suite { const self = try a.create(Suite); @@ -16,10 +18,12 @@ pub const Suite = struct { self.masks = try tp_.Masks.init(gpu, s); errdefer self.masks.deinit(); try self.masks.add(a, kernels); + self.prefill = try prefill_.Group.create(a, gpu, s, kernels); return self; } pub fn destroy(self: *Suite) void { + self.prefill.destroy(); self.masks.deinit(); self.a.destroy(self); } diff --git a/zig/tests/cuda/contract_limits.zig b/zig/tests/cuda/contract_limits.zig index b25ab8b403..91f2100e17 100644 --- a/zig/tests/cuda/contract_limits.zig +++ b/zig/tests/cuda/contract_limits.zig @@ -26,7 +26,7 @@ pub fn image(name: []const u8) error{UnknownContractImage}![]const u8 { inline for (.{ "qwen4_projection", "qwen4_hc", "qwen4_gdn", "qwen4_ple" }) |n| if (std.mem.eql(u8, name, n)) return @field(kk, n); inline for (.{ "tree_round", "act_rows", "gdn_prep", "qk_rows", "norms", "moe", "experts", "qwen36_mtp", "nemotron_attention", "nemotron_route", "nemotron_latent", "qmm_group", "qmm_prefill", "lane_gemv", "prompt16", "qmmf", "ck_experts", "nvfp4_experts", "nvfp4_experts_tile", "rank_sum", "prefill_attention", "vit", "embed_rows", "route_rows", "flashnext_gdn", "flashnext_batch", "flashnext_hc", "flashnext_attn", "flashnext_mtp", "qwen_vision", "exl3_linear", "exl3_experts", "exl3_candidate", "experts3", "experts3_prefill", "drama3", "drama3_fp8", "drama3_codec", "glm_mm", "glm_dflash", "glm_glue", "glm_ops", "glm_index", "nemotron_mamba", "sample", "qwen35", "qwen35_bf16_rows", "qwen35_keep", "qwen35_kv8", "w4a8_prefill", "qwen35_draft", "ck_lane", "gdn_train", "train_q4", "train_attn", "flashnext_train", "train", "attn_train", "qwen27_train", "tree_attention", "block_attention", "window_attention", "window_pages", "moe_glue", "dense_rows", "qmm", "flashnext_fact", "flashnext_hcf", "flashnext_fexp", "experts_rows" }) |n| if (std.mem.eql(u8, name, n)) return @field(kk, n); - inline for (.{ "flashnext_exl3", "flashnext_exl3_tp" }) |n| + inline for (.{ "flashnext_exl3", "flashnext_exl3_tp", "flashnext_exl3_prefill" }) |n| if (std.mem.eql(u8, name, n)) return @field(kk, n); if (std.mem.eql(u8, name, "contract_controls")) return if (options.controls) &controls.bytes else &.{}; std.debug.print("FAIL kernel-contract: unknown image name: {s}\n", .{name}); diff --git a/zig/tests/cuda/flashnext_exl3/prefill/exl3_prefill_dense.zig b/zig/tests/cuda/flashnext_exl3/prefill/exl3_prefill_dense.zig index 833f380083..6b1bf2afbd 100644 --- a/zig/tests/cuda/flashnext_exl3/prefill/exl3_prefill_dense.zig +++ b/zig/tests/cuda/flashnext_exl3/prefill/exl3_prefill_dense.zig @@ -1,4 +1,4 @@ -//! exl3-prefill dense: the prompt tiles of a dense EXL3 projection against a float64 reference. +//! exl3-prefill dense: the prompt tiles of a dense EXL3 projection, bit-equal to the decode kernel at every row count. const std = @import("std"); const cuda = @import("cuda"); @@ -23,7 +23,9 @@ const DenseCall = struct { x: u64, y: u64, rows: usize, + exact: bool = false, pub fn go(c: DenseCall) anyerror!void { + if (c.exact) return c.de.promptExact(c.s, c.w, c.x, c.rows, c.y, .bf16); try c.de.prompt(c.s, c.w, c.x, .bf16, c.rows, c.y, .bf16); } }; @@ -78,7 +80,7 @@ pub fn denseRun(gpu: Gpu, pack: *const Pack, s: cuda.Stream) !void { "model.language_model.layers.3.mlp.shared_expert.gate_proj", "model.language_model.layers.3.mlp.shared_expert.down_proj", }; - const row_counts = [_]usize{ 16, 128, 256, 1024, 2048 }; + const row_counts = [_]usize{ 16, 63, 64, 65, 128, 256, 1024, 2048 }; const max_rows: usize = 2048; var prng = std.Random.DefaultPrng.init(0xde45e); const rng = prng.random(); @@ -143,29 +145,25 @@ pub fn denseRun(gpu: Gpu, pack: *const Pack, s: cuda.Stream) !void { peak = @max(peak, @abs(wv)); } if (tiles_ok) { - de.implementation = .tiles; - const tc: DenseCall = .{ .de = &de, .s = s, .w = w, .x = x.ptr, .y = yt.ptr, .rows = rows }; + try de.candidates(w.k); + const tc: DenseCall = .{ .de = &de, .s = s, .w = w, .x = x.ptr, .y = yt.ptr, .rows = rows, .exact = true }; ms_tiles = try t.best(s, 3, tc, DenseCall.go); sum_tiles[ri] += ms_tiles; const got_t = try a.alloc(u16, rows * n); defer a.free(got_t); try yt.download(0, std.mem.sliceAsBytes(got_t)); - diff = 0; - var rpeak: f64 = 0; - for (got_r, got_t) |g, h| { - diff = @max(diff, @abs(bfValue(g) - bfValue(h))); - rpeak = @max(rpeak, @abs(bfValue(g))); - } - diff /= rpeak; + var unequal: usize = 0; + for (got_r, got_t) |g, h| unequal += @intFromBool(g != h); + diff = @floatFromInt(unequal); err_tiles = 0; for (got_t[0 .. 2 * n], want) |g, wv| err_tiles = @max(err_tiles, @abs(bfValue(g) - wv)); - de.implementation = .reference; + try check.expect(unequal == 0, "{s}: the prompt tiles' bytes are the decode kernel's at {d} rows ({d} of {d} outputs differ)", .{ name, rows, unequal, got_r.len }); } else sum_tiles[ri] += ms_ref; - std.debug.print(" rows {d:5}: reference {d:8.4} ms ({d:6.2} us/row) tiles {d:8.4} ms ({d:6.2} us/row) max|ref-tiles|/peak {e:.2} rows 0-1 vs f64: ref {e:.2} tiles {e:.2} (rel to peak {d:.3})\n", .{ rows, ms_ref, 1e3 * ms_ref / @as(f64, @floatFromInt(rows)), ms_tiles, 1e3 * ms_tiles / @as(f64, @floatFromInt(rows)), diff, err_ref / peak, err_tiles / peak, peak }); + std.debug.print(" rows {d:5}: reference {d:8.4} ms ({d:6.2} us/row) tiles {d:8.4} ms ({d:6.2} us/row) outputs differing from the decode kernel's {d:.0} rows 0-1 vs f64: ref {e:.2} tiles {e:.2} (rel to peak {d:.3})\n", .{ rows, ms_ref, 1e3 * ms_ref / @as(f64, @floatFromInt(rows)), ms_tiles, 1e3 * ms_tiles / @as(f64, @floatFromInt(rows)), diff, err_ref / peak, err_tiles / peak, peak }); } } std.debug.print(" sum over the {d} shapes (ms): ", .{names.len}); for (row_counts, sum_ref, sum_tiles) |r, x, y| std.debug.print("rows {d}: reference {d:.3} best {d:.3}; ", .{ r, x, y }); std.debug.print("\n", .{}); - check.pass("dense prompt kernels timed", .{}); + check.pass("dense prompt tiles are the decode kernel's bytes at {d} row counts, timed", .{row_counts.len}); } diff --git a/zig/tests/cuda/flashnext_exl3/prefill/exl3_prefill_engine.zig b/zig/tests/cuda/flashnext_exl3/prefill/exl3_prefill_engine.zig new file mode 100644 index 0000000000..ac49300e55 --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/prefill/exl3_prefill_engine.zig @@ -0,0 +1,360 @@ +//! exl3-prefill engine and agree: whole prefills through the deep chunks, timed and compared with the batched pass. + +const std = @import("std"); +const cuda = @import("cuda"); +const check = @import("../../check.zig"); +const flash = @import("flashnext"); +const Gpu = check.Gpu; + +// ---------------------------------------------------------------------------------------------------- engine + +const prompt = flash.prompt; +const xpass = flash.exl3.pass; +const xpre = flash.exl3.prefill; +const vocab = prompt.vocabulary; + +fn promote(bits: u16) f32 { + return @bitCast(@as(u32, bits) << 16); +} + +const Eng = struct { + gpu: Gpu, + stream: cuda.Stream, + glue: flash.Glue, + e: *prompt.Engine, + seq: prompt.Seq, + attached: xpass.Attached, + + pub fn init(gpu: Gpu, dir: []const u8, lines: ?usize, cap: usize, deep_rows: usize) !*Eng { + const g = try gpu.gpa.create(Eng); + errdefer gpu.gpa.destroy(g); + g.gpu = gpu; + g.stream = try cuda.Stream.init(gpu.d, true); + g.glue = try flash.Glue.load(gpu.d); + const t0 = check.now(gpu.io); + g.e = try prompt.Engine.initRank(gpu.gpa, gpu.io, @constCast(gpu.d), &g.stream, dir, 1, 0); + var free0: usize = 0; + var total: usize = 0; + try gpu.d.check(gpu.d.api.cuMemGetInfo_v2(&free0, &total), "cuMemGetInfo"); + _ = try g.e.reserve(@min(cap, g.e.chunk)); + _ = try g.e.reserveScores(cap, @min(cap, g.e.chunk)); + var free1: usize = 0; + try gpu.d.check(gpu.d.api.cuMemGetInfo_v2(&free1, &total), "cuMemGetInfo"); + g.seq = try g.e.newSeq(cap); + var free2: usize = 0; + try gpu.d.check(gpu.d.api.cuMemGetInfo_v2(&free2, &total), "cuMemGetInfo"); + if (deep_rows > 0) try xpre.reserve(gpu.gpa, gpu.d, g.e, deep_rows); + var free3: usize = 0; + try gpu.d.check(gpu.d.api.cuMemGetInfo_v2(&free3, &total), "cuMemGetInfo"); + var b = try xpass.budgetFromEnv(0); + if (lines) |n| b.lines = n; + g.attached = try xpass.attach(gpu.gpa, gpu.io, gpu.d, g.stream, g.e.layers[0..g.e.n_layers], null, g.e.chunk, b, null); + var free4: usize = 0; + try gpu.d.check(gpu.d.api.cuMemGetInfo_v2(&free4, &total), "cuMemGetInfo"); + const gib = struct { + fn f(x: usize) f64 { + return @as(f64, @floatFromInt(x)) / (1 << 30); + } + }.f; + std.debug.print("engine: loaded in {d:.1} s; card {d:.2} GiB; after weights {d:.2} free; pass scratch (chunk {d}) {d:.2} GiB; caches for {d} tokens {d:.2} GiB; deep chunks of {d} rows {d:.2} GiB; pool {d} lines a layer {d:.2} GiB; {d:.2} GiB free after\n", .{ @as(f64, @floatFromInt(check.now(gpu.io) - t0)) / 1e9, gib(total), gib(free0), g.e.chunk, gib(free0 - free1), cap, gib(free1 - free2), deep_rows, gib(free2 - free3), g.attached.lines, gib(g.attached.device_bytes), gib(free4) }); + return g; + } + + pub fn deinit(g: *Eng) void { + xpre.release(); + g.seq.deinit(g.gpu.gpa); + g.e.deinit(); + g.glue.deinit(); + g.stream.deinit(); + g.gpu.gpa.destroy(g); + } + + fn tri(g: *Eng) flash.Tri { + return .{ .own = &g.glue, .s = g.stream }; + } + + fn reset(g: *Eng) !void { + const s = &g.seq; + for ([_][]cuda.DeviceBuffer{ s.conv, s.state, s.kc, s.vc, s.ikc, s.pooled }) |bs| for (bs) |b| try b.fill8(0, g.stream.handle); + try s.tail.fill8(0, g.stream.handle); + try g.stream.synchronize(); + s.fed = 0; + s.history = .{ g.e.eos, g.e.eos }; + } + + fn poolStats(g: *Eng) flash.exl3_pool.Stats { + return xpass.poolStats(g.e.layers[0..g.e.n_layers], null); + } +}; + +fn top5(row: []const u16) [5]u32 { + var ids: [5]u32 = @splat(0); + var vals: [5]f32 = @splat(-std.math.inf(f32)); + for (row, 0..) |b, i| { + const v = promote(b); + if (!(v > vals[4])) continue; + var at: usize = 4; + while (at > 0 and v > vals[at - 1]) : (at -= 1) { + vals[at] = vals[at - 1]; + ids[at] = ids[at - 1]; + } + vals[at] = v; + ids[at] = @intCast(i); + } + return ids; +} + +/// One prefill of `ids` from a fresh sequence: ms, and the last row's logits copied into `out`. +fn prefillOnce(g: *Eng, ids: []const u32, deep_on: bool, out: []u16) !f64 { + var o = xpre.optionsFromEnv(); + o.enabled = deep_on; + if (rows_override) |r| o.rows = r; + xpre.state.options = o; + try g.reset(); + const t0 = check.now(g.gpu.io); + try g.e.forward(flash.Tri, g.tri(), &g.seq, ids, .{ .logits = .last }); + const ms_ = @as(f64, @floatFromInt(check.now(g.gpu.io) - t0)) / 1e6; + @memcpy(std.mem.sliceAsBytes(out), g.seq.logits); + return ms_; +} + +/// Greedy decode rounds after each engine-mode prefill (TF_TEST_DECODE, default 32; 0: none). +var decode_steps: usize = 32; + +fn decodeAfter(g: *Eng, n: usize) !void { + var out = try cuda.DeviceBuffer.alloc(g.gpu.d, vocab * 2); + defer out.free(); + var tok = prompt.argmax(g.seq.logits); + const before = g.poolStats(); + const t0 = check.now(g.gpu.io); + var first: f64 = 0; + for (0..decode_steps) |i| { + const q = [_]*prompt.Seq{&g.seq}; + try g.e.round(flash.Tri, g.tri(), &q, (&tok)[0..1], out.ptr); + tok = prompt.argmax(g.seq.logits); + if (i == 7) first = @as(f64, @floatFromInt(check.now(g.gpu.io) - t0)) / 1e6; + } + const t = @as(f64, @floatFromInt(check.now(g.gpu.io) - t0)) / 1e6; + const after = g.poolStats(); + const hits = after.hit_slots - before.hit_slots; + const all = hits + after.miss_slots - before.miss_slots; + std.debug.print(" decode after {d} tokens: {d:.2} ms/token over {d} steps (first 8: {d:.2} ms/token; hit {d:.1}%)\n", .{ n, t / @as(f64, @floatFromInt(decode_steps)), decode_steps, first / 8, 100 * @as(f64, @floatFromInt(hits)) / @as(f64, @floatFromInt(@max(all, 1))) }); +} + +/// The deep chunk's rows for the next prefills (engine mode's ROWS list), else the environment's. +var rows_override: ?usize = null; + +pub fn engineRun(gpu: Gpu, dir: []const u8, rest: []const [:0]const u8) !void { + const a = gpu.gpa; + if (rest.len < 2) return error.MissingArgument; + const text = try std.Io.Dir.cwd().readFileAlloc(gpu.io, rest[0], a, .limited(1 << 28)); + defer a.free(text); + const all_ids = try a.alloc(u32, text.len / 4); + defer a.free(all_ids); + for (all_ids, 0..) |*x, i| x.* = std.mem.readInt(u32, text[4 * i ..][0..4], .little); + var ns: [16]usize = undefined; + var n_ns: usize = 0; + var it = std.mem.splitScalar(u8, rest[1], ','); + while (it.next()) |tok| { + ns[n_ns] = try std.fmt.parseInt(usize, tok, 10); + n_ns += 1; + } + const lines: ?usize = if (rest.len > 2 and !std.mem.eql(u8, rest[2], "-")) try std.fmt.parseInt(usize, rest[2], 10) else null; + const cmp = rest.len > 3 and std.mem.eql(u8, rest[3], "cmp"); + // ROWS: deep chunk sizes to sweep (rest[3] when not "cmp", or rest[4]) + var rs: [16]usize = undefined; + var n_rs: usize = 0; + const rows_arg: ?[]const u8 = if (rest.len > 3 and !cmp) rest[3] else if (rest.len > 4) rest[4] else null; + if (rows_arg) |ra| { + var rit = std.mem.splitScalar(u8, ra, ','); + while (rit.next()) |tok| { + rs[n_rs] = try std.fmt.parseInt(usize, tok, 10); + n_rs += 1; + } + } + var max_n: usize = 0; + for (ns[0..n_ns]) |n| max_n = @max(max_n, n); + if (max_n > all_ids.len) return error.NotEnoughIds; + if (std.c.getenv("TF_TEST_DECODE")) |v| decode_steps = std.fmt.parseInt(usize, std.mem.span(v), 10) catch decode_steps; + const o = xpre.optionsFromEnv(); + var max_r = o.rows; + for (rs[0..n_rs]) |r| max_r = @max(max_r, r); + if (n_rs > 0) max_r = 0; + for (rs[0..n_rs]) |r| max_r = @max(max_r, r); + const deep_rows = @min(max_r, (max_n + 255) / 256 * 256); + const g = try Eng.init(gpu, dir, lines, max_n + 64, if (o.enabled) deep_rows else 0); + defer g.deinit(); + const old = try a.alloc(u16, vocab); + defer a.free(old); + const new = try a.alloc(u16, vocab); + defer a.free(new); + // warm-up: kernels, the pool's first lines + _ = try prefillOnce(g, all_ids[0..@min(max_n, 1024)], o.enabled, new); + var all_ok = true; + if (n_rs == 0) { + rs[0] = 0; + n_rs = 1; + } + for (rs[0..n_rs]) |rr| for (ns[0..n_ns]) |n| { + rows_override = if (rr == 0) null else rr; + if (rr != 0) std.debug.print("deep chunk rows {d}:\n", .{rr}); + const ids = all_ids[0..n]; + const p0 = g.poolStats(); + xpre.state.stats = .{}; + const t_new = try prefillOnce(g, ids, o.enabled, new); + const p1 = g.poolStats(); + const xs_ = xpre.state.stats; + std.debug.print("PREFILL {d:6} tokens ({s}): {d:9.1} ms = {d:7.1} rows/s; pool copies {d} ({d:.1} GiB, {d:.2} a layer per 1k rows), deep chunks {d}, passes {d}, prefetched {d}, members {d}; host: picks wait {d:.1} ms, plan {d:.1} ms\n", .{ n, if (o.enabled) "expert-major deep" else "batched pass", t_new, @as(f64, @floatFromInt(n)) / (t_new / 1e3), p1.misses - p0.misses, @as(f64, @floatFromInt(p1.misses - p0.misses)) * 1.758 / 1024, @as(f64, @floatFromInt(p1.misses - p0.misses)) / 48 / (@as(f64, @floatFromInt(n)) / 1000), xs_.chunks, xs_.passes, xs_.prefetched, xs_.members, @as(f64, @floatFromInt(xs_.host_wait_ns)) / 1e6, @as(f64, @floatFromInt(xs_.plan_ns)) / 1e6 }); + // decode after the prefill: the pools were left empty by the ring (or hold the batched pass's lines) + if (decode_steps > 0) try decodeAfter(g, n); + if (!cmp) continue; + const q0 = g.poolStats(); + const t_old = try prefillOnce(g, ids, false, old); + const q1 = g.poolStats(); + std.debug.print(" same prompt, batched pass: {d:9.1} ms = {d:7.1} rows/s; pool copies {d}\n", .{ t_old, @as(f64, @floatFromInt(n)) / (t_old / 1e3), q1.misses - q0.misses }); + if (decode_steps > 0) try decodeAfter(g, n); + var dmax: f32 = 0; + var peak: f32 = 0; + var same: usize = 0; + for (old, new) |x, y| { + dmax = @max(dmax, @abs(promote(x) - promote(y))); + peak = @max(peak, @abs(promote(x))); + same += @intFromBool(x == y); + } + const to = top5(old); + const tn = top5(new); + var overlap: usize = 0; + for (to) |x| for (tn) |y| { + overlap += @intFromBool(x == y); + }; + const margin_old = promote(old[to[0]]) - promote(old[to[1]]); + std.debug.print("CMP {d:6} tokens: last-row logits max |d| {d:.4} (peak {d:.2}, {d:.1}% of logits bit-equal); top-1 {d} vs {d} ({s}, batched-pass margin {d:.4}); top-5 overlap {d}/5\n", .{ n, dmax, peak, 100 * @as(f64, @floatFromInt(same)) / vocab, to[0], tn[0], if (to[0] == tn[0]) "equal" else "DIFFERENT", margin_old, overlap }); + if (to[0] != tn[0] and margin_old > 0.25) all_ok = false; + }; + const ps = g.poolStats(); + std.debug.print("pool overall: {d} rounds, {d} copies, hit {d:.1}%\n", .{ ps.rounds, ps.misses, 100 * @as(f64, @floatFromInt(ps.hit_slots)) / @as(f64, @floatFromInt(@max(1, ps.hit_slots + ps.miss_slots))) }); + try check.expect(all_ok, "top-1 equal wherever the batched pass's margin exceeds bf16 noise", .{}); + check.pass("engine prefill measured", .{}); +} + +fn readIds(gpu: Gpu, path: []const u8) ![]u32 { + const a = gpu.gpa; + const text = try std.Io.Dir.cwd().readFileAlloc(gpu.io, path, a, .limited(1 << 28)); + defer a.free(text); + const ids = try a.alloc(u32, text.len / 4); + for (ids, 0..) |*x, i| x.* = std.mem.readInt(u32, text[4 * i ..][0..4], .little); + return ids; +} + +pub fn agreeRun(gpu: Gpu, dir: []const u8, rest: []const [:0]const u8) !void { + const a = gpu.gpa; + if (rest.len < 4) return error.MissingArgument; + const all_ids = try readIds(gpu, rest[0]); + defer a.free(all_ids); + const first = try std.fmt.parseInt(usize, rest[1], 10); + const step = try std.fmt.parseInt(usize, rest[2], 10); + const count = try std.fmt.parseInt(usize, rest[3], 10); + const lines: ?usize = if (rest.len > 4) try std.fmt.parseInt(usize, rest[4], 10) else null; + const max_n = first + step * (count - 1); + if (max_n + 1 > all_ids.len) return error.NotEnoughIds; + const o = xpre.optionsFromEnv(); + const g = try Eng.init(gpu, dir, lines, max_n + 64, @min(o.rows, (max_n + 255) / 256 * 256)); + defer g.deinit(); + const old = try a.alloc(u16, vocab); + defer a.free(old); + const new = try a.alloc(u16, vocab); + defer a.free(new); + var agree: usize = 0; + var acc_old: usize = 0; + var acc_new: usize = 0; + var top5_same: usize = 0; + var sum_d: f64 = 0; + var max_d: f64 = 0; + var n_d: usize = 0; + var disagree_margin_max: f32 = 0; + var t_old: f64 = 0; + var t_new: f64 = 0; + for (0..count) |i| { + const n = first + step * i; + const ids = all_ids[0..n]; + t_old += try prefillOnce(g, ids, false, old); + t_new += try prefillOnce(g, ids, o.enabled, new); + const to = top5(old); + const tn = top5(new); + const want = all_ids[n]; + agree += @intFromBool(to[0] == tn[0]); + acc_old += @intFromBool(to[0] == want); + acc_new += @intFromBool(tn[0] == want); + top5_same += @intFromBool(std.mem.eql(u32, &to, &tn)); + for (to) |id| { + const d = @abs(@as(f64, promote(old[id])) - promote(new[id])); + sum_d += d; + max_d = @max(max_d, d); + n_d += 1; + } + if (to[0] != tn[0]) { + const m = promote(old[to[0]]) - promote(old[to[1]]); + disagree_margin_max = @max(disagree_margin_max, m); + std.debug.print(" length {d}: top-1 {d} (batched) vs {d} (deep), batched margin {d:.4}, deep margin {d:.4}\n", .{ n, to[0], tn[0], m, promote(new[tn[0]]) - promote(new[tn[1]]) }); + } + } + std.debug.print("AGREE {d} prompt lengths {d}..{d}: top-1 equal {d}/{d}, top-5 identical in order {d}/{d}; next-token accuracy batched {d}/{d}, deep {d}/{d}; |d| on the batched top-5: mean {d:.4}, max {d:.4}; largest batched margin where top-1 differs {d:.4}; time batched {d:.1} s, deep {d:.1} s\n", .{ count, first, max_n, agree, count, top5_same, count, acc_old, count, acc_new, count, sum_d / @as(f64, @floatFromInt(n_d)), max_d, disagree_margin_max, t_old / 1e3, t_new / 1e3 }); + check.pass("agreement measured", .{}); +} + +/// Greedy tokens after the prefill that just ran: the next `out.len` ids, each fed back through a decode round. +fn greedy(g: *Eng, out: []u32) !void { + var buf = try cuda.DeviceBuffer.alloc(g.gpu.d, vocab * 2); + defer buf.free(); + var tok = prompt.argmax(g.seq.logits); + for (out) |*o| { + o.* = tok; + const q = [_]*prompt.Seq{&g.seq}; + try g.e.round(flash.Tri, g.tri(), &q, (&tok)[0..1], buf.ptr); + tok = prompt.argmax(g.seq.logits); + } +} + +/// One prompt through the batched pass and through the deep chunks at each chunk size: same logits bytes and tokens. +pub fn equalRun(gpu: Gpu, dir: []const u8, rest: []const [:0]const u8) !void { + const a = gpu.gpa; + if (rest.len < 3) return error.MissingArgument; + const all_ids = try readIds(gpu, rest[0]); + defer a.free(all_ids); + const n = try std.fmt.parseInt(usize, rest[1], 10); + if (n == 0 or n > all_ids.len) return error.NotEnoughIds; + var sizes: [8]usize = undefined; + var n_sizes: usize = 0; + var it = std.mem.splitScalar(u8, rest[2], ','); + while (it.next()) |tok| : (n_sizes += 1) sizes[n_sizes] = try std.fmt.parseInt(usize, tok, 10); + const lines: ?usize = if (rest.len > 3 and !std.mem.eql(u8, rest[3], "-")) try std.fmt.parseInt(usize, rest[3], 10) else null; + var max_rows: usize = 0; + for (sizes[0..n_sizes]) |r| max_rows = @max(max_rows, r); + const g = try Eng.init(gpu, dir, lines, n + 64, @min(max_rows, (n + 255) / 256 * 256)); + defer g.deinit(); + const ids = all_ids[0..n]; + const want = try a.alloc(u16, vocab); + defer a.free(want); + const got = try a.alloc(u16, vocab); + defer a.free(got); + var want_tokens: [24]u32 = undefined; + var got_tokens: [24]u32 = undefined; + // TF_EQUAL_FIRST_SIZE=1: the first chunk size is the reference (the batched pass is not the canonical reduction's) + const first = std.c.getenv("TF_EQUAL_FIRST_SIZE") != null; + rows_override = if (first) sizes[0] else null; + _ = try prefillOnce(g, ids, first, want); + try greedy(g, &want_tokens); + var all_equal = true; + for (sizes[if (first) 1 else 0..n_sizes]) |rows| { + rows_override = rows; + _ = try prefillOnce(g, ids, true, got); + try greedy(g, &got_tokens); + const logits_equal = std.mem.eql(u8, std.mem.sliceAsBytes(want), std.mem.sliceAsBytes(got)); + const tokens_equal = std.mem.eql(u32, &want_tokens, &got_tokens); + std.debug.print("EQUAL {d} tokens, deep chunks of {d} rows vs the batched pass: last-row logits {s}, {d} greedy tokens {s}\n", .{ n, rows, if (logits_equal) "bit-equal" else "DIFFERENT", want_tokens.len, if (tokens_equal) "identical" else "DIFFERENT" }); + all_equal = all_equal and logits_equal and tokens_equal; + } + rows_override = null; + try check.expect(all_equal, "deep prefill on (at {d} chunk sizes) and off give the same logits bytes and tokens", .{n_sizes}); + check.pass("deep prefill is the batched pass's bytes at {d} chunk sizes", .{n_sizes}); +} diff --git a/zig/tests/cuda/flashnext_exl3/prefill/exl3_prefill_experts.zig b/zig/tests/cuda/flashnext_exl3/prefill/exl3_prefill_experts.zig index 2dcca86822..b16ff3eabc 100644 --- a/zig/tests/cuda/flashnext_exl3/prefill/exl3_prefill_experts.zig +++ b/zig/tests/cuda/flashnext_exl3/prefill/exl3_prefill_experts.zig @@ -1,8 +1,10 @@ -//! exl3-prefill experts: the grouped, slot and expert-major masked passes timed on one layer. +//! exl3-prefill experts and prompt: the grouped, slot and expert-major prefill kernels timed on one layer. const std = @import("std"); const cuda = @import("cuda"); const check = @import("../../check.zig"); +const flash = @import("flashnext"); +const pk_ = flash.exl3.prefill_k; const exl3e = cuda.exl3_experts; const Gpu = check.Gpu; const packs = @import("exl3_prefill_pack.zig"); @@ -116,3 +118,130 @@ pub fn expertsRun(gpu: Gpu, ex: *const exl3e.Experts, s: cuda.Stream, ld: *const } check.pass("layer {d}: runSlots timed at 16..{d} rows; expert-major masked passes are bit-identical to one call", .{ ld.layer, max_rows }); } + +// ---------------------------------------------------------------------------------------------------- prompt kernel + +const PromptCall = struct { + k: *const pk_.Kernels, + s: cuda.Stream, + l: exl3e.Layer, + x: u64, + picks: u64, + n_members: usize, + n_items: usize, + dv: pk_.Device, + y: u64, + pub fn go(c: PromptCall) anyerror!void { + try pk_.run(c.k, c.s, c.l, c.x, dims, c.picks, slots, c.n_members, c.n_items, c.dv, c.y); + } +}; + +/// P expert-major passes, each its own uploaded plan (members, items) over the same picks. +const PromptPasses = struct { + calls: []const PromptCall, + pub fn go(c: PromptPasses) anyerror!void { + for (c.calls) |call| try call.go(); + } +}; + +pub fn promptRun(gpu: Gpu, ex: *const exl3e.Experts, s: cuda.Stream, ld: *const LayerData) !void { + const a = gpu.gpa; + var k = try pk_.Kernels.load(gpu.d); + defer k.unload(); + var t = try Timer.init(gpu.d); + defer t.deinit(); + const max_rows: usize = 8192; + const max_pairs = max_rows * slots; + var prng = std.Random.DefaultPrng.init(0x9a11); + const rng = prng.random(); + const xh = try a.alloc(u16, max_rows * dims); + defer a.free(xh); + for (xh) |*b| b.* = bf16Bits((rng.float(f32) - 0.5) * 2.0); + var x = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(xh)); + defer x.free(); + const picks = try a.alloc(i32, max_pairs); + defer a.free(picks); + genPicks(rng, 0.8, max_rows, picks); + var pkd = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(picks)); + defer pkd.free(); + var plan = try pk_.Plan.init(a, max_pairs, table_n); + defer plan.deinit(a); + var counts: [table_n]u32 = undefined; + var mem_d = try cuda.DeviceBuffer.alloc(gpu.d, max_pairs * 4); + defer mem_d.free(); + var items_d = try cuda.DeviceBuffer.alloc(gpu.d, 3 * pk_.maxItems(max_pairs, table_n) * 4); + defer items_d.free(); + var scr = try cuda.DeviceBuffer.alloc(gpu.d, pk_.scratchBytes(max_pairs, dims, width)); + defer scr.free(); + var yk = try cuda.DeviceBuffer.alloc(gpu.d, max_pairs * dims * 4); + defer yk.free(); + // runSlots' reference (at most max_call_rows rows a call) + var sbuf = try cuda.DeviceBuffer.alloc(gpu.d, try exl3e.scratchBytes(ld.res, max_call_rows, slots)); + defer sbuf.free(); + const sc = try exl3e.scratch(sbuf.ptr, ld.res, max_call_rows, slots); + std.debug.print(" prefill experts kernel (64-member tiles, compact pairs) vs runSlots, one resident layer, zipf0.8 picks:\n", .{}); + for ([_]usize{ 16, 64, 256, 1024, 1408, 2048, 4096, 8192 }) |rows| { + const pairs = rows * slots; + pk_.build(picks[0..pairs], null, &counts, &plan); + try mem_d.upload(0, plan.memberBytes()); + try items_d.upload(0, plan.itemBytes()); + const dv: pk_.Device = .{ .members = mem_d.ptr, .items = items_d.ptr, .scratch = scr.ptr }; + const call: PromptCall = .{ .k = &k, .s = s, .l = ld.res, .x = x.ptr, .picks = pkd.ptr, .n_members = plan.n_members, .n_items = plan.n_items, .dv = dv, .y = yk.ptr }; + try call.go(); + try s.synchronize(); + const ms_k = try t.best(s, 3, call, PromptCall.go); + var ms_r: f64 = std.math.nan(f64); + var rel: f64 = std.math.nan(f64); + var same_words: usize = 0; + if (rows <= max_call_rows) { + const rc: SlotsCall = .{ .ex = ex, .s = s, .l = ld.res, .x = x.ptr, .picks = pkd.ptr, .rows = rows, .sc = sc }; + try rc.go(); + try s.synchronize(); + const want = try a.alloc(f32, pairs * dims); + defer a.free(want); + const got = try a.alloc(f32, pairs * dims); + defer a.free(got); + try gpu.d.check(gpu.d.api.cuMemcpyDtoH_v2(want.ptr, sc.y, want.len * 4), "cuMemcpyDtoH"); + try gpu.d.check(gpu.d.api.cuMemcpyDtoH_v2(got.ptr, yk.ptr, got.len * 4), "cuMemcpyDtoH"); + var dmax: f64 = 0; + var peak: f64 = 0; + for (got, want) |g, w| { + dmax = @max(dmax, @abs(@as(f64, g) - w)); + peak = @max(peak, @abs(@as(f64, w))); + same_words += @intFromBool(@as(u32, @bitCast(g)) == @as(u32, @bitCast(w))); + } + rel = dmax / peak; + ms_r = try t.best(s, 3, rc, SlotsCall.go); + try check.expect(same_words == pairs * dims, "prefill kernel's y is runSlots' bytes at {d} rows ({d} of {d} words differ)", .{ rows, pairs * dims - same_words, pairs * dims }); + } + std.debug.print(" rows {d:5}: prefill kernel {d:8.3} ms ({d:6.2} us/row, {d} items) runSlots {d:8.3} ms ({d:6.2} us/row) speedup {d:.2}x max|dy|/peak {e:.2}, {d:.1}% of words bit-equal\n", .{ rows, ms_k, 1e3 * ms_k / @as(f64, @floatFromInt(rows)), plan.n_items, ms_r, 1e3 * ms_r / @as(f64, @floatFromInt(rows)), ms_r / ms_k, rel, 100 * @as(f64, @floatFromInt(same_words)) / @as(f64, @floatFromInt(pairs * dims)) }); + } + // expert-major passes over 8192 rows: the plan of each pass alone, back to back + const rows = max_rows; + const pairs = rows * slots; + std.debug.print(" expert-major passes over {d} rows (pass j: experts [j*512/P, (j+1)*512/P), the shared one in pass 0):\n", .{rows}); + for ([_]usize{ 1, 8, 16, 32 }) |passes| { + const per = routed_n / passes; + var calls: [32]PromptCall = undefined; + var mems: [32]cuda.DeviceBuffer = undefined; + var its: [32]cuda.DeviceBuffer = undefined; + var made: usize = 0; + defer for (mems[0..made], its[0..made]) |*m, *i| { + m.free(); + i.free(); + }; + for (0..passes) |j| { + var take: [table_n]bool = @splat(false); + for (j * per..(j + 1) * per) |id| take[id] = true; + take[512] = j == 0; + pk_.build(picks[0..pairs], &take, &counts, &plan); + mems[j] = try cuda.DeviceBuffer.fromHost(gpu.d, plan.memberBytes()); + its[j] = try cuda.DeviceBuffer.fromHost(gpu.d, plan.itemBytes()); + made += 1; + calls[j] = .{ .k = &k, .s = s, .l = ld.res, .x = x.ptr, .picks = pkd.ptr, .n_members = plan.n_members, .n_items = plan.n_items, .dv = .{ .members = mems[j].ptr, .items = its[j].ptr, .scratch = scr.ptr }, .y = yk.ptr }; + } + const ms = try t.best(s, 2, PromptPasses{ .calls = calls[0..passes] }, PromptPasses.go); + std.debug.print(" {d:2} passes: {d:8.3} ms ({d:6.2} us/row)\n", .{ passes, ms, 1e3 * ms / @as(f64, @floatFromInt(rows)) }); + } + check.pass("layer {d}: prefill experts kernel is runSlots' bytes, timed to {d} rows", .{ ld.layer, max_rows }); +} diff --git a/zig/tests/cuda/flashnext_exl3/prefill/exl3_prefill_tests.zig b/zig/tests/cuda/flashnext_exl3/prefill/exl3_prefill_tests.zig index 5695cf4cf4..cecf237e1d 100644 --- a/zig/tests/cuda/flashnext_exl3/prefill/exl3_prefill_tests.zig +++ b/zig/tests/cuda/flashnext_exl3/prefill/exl3_prefill_tests.zig @@ -9,17 +9,25 @@ const packs = @import("exl3_prefill_pack.zig"); const experts = @import("exl3_prefill_experts.zig"); const copies = @import("exl3_prefill_copy.zig"); const dense = @import("exl3_prefill_dense.zig"); +const engine = @import("exl3_prefill_engine.zig"); const LayerData = packs.LayerData; const Pack = packs.Pack; const expertsRun = experts.expertsRun; +const promptRun = experts.promptRun; const copyRun = copies.copyRun; const denseRun = dense.denseRun; +const agreeRun = engine.agreeRun; +const equalRun = engine.equalRun; +const engineRun = engine.engineRun; pub fn run(gpu: Gpu, rest: []const [:0]const u8) !void { if (rest.len < 2) return error.MissingArgument; const dir = rest[0]; const mode = rest[1]; + if (std.mem.eql(u8, mode, "engine")) return engineRun(gpu, dir, rest[2..]); + if (std.mem.eql(u8, mode, "agree")) return agreeRun(gpu, dir, rest[2..]); + if (std.mem.eql(u8, mode, "equal")) return equalRun(gpu, dir, rest[2..]); const layer: usize = if (rest.len > 2) try std.fmt.parseInt(usize, rest[2], 10) else 17; var pack: Pack = .{}; defer pack.close(gpu.io); @@ -33,5 +41,6 @@ pub fn run(gpu: Gpu, rest: []const [:0]const u8) !void { defer ld.free(gpu.gpa); if (std.mem.eql(u8, mode, "experts")) return expertsRun(gpu, &ex, s, ld); if (std.mem.eql(u8, mode, "copy")) return copyRun(gpu, &ex, s, ld); + if (std.mem.eql(u8, mode, "prompt")) return promptRun(gpu, &ex, s, ld); return error.UnknownMode; } diff --git a/zig/tests/cuda/flashnext_exl3/tp/exl3_tpprefill_tests.zig b/zig/tests/cuda/flashnext_exl3/tp/exl3_tpprefill_tests.zig new file mode 100644 index 0000000000..f142b76931 --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/tp/exl3_tpprefill_tests.zig @@ -0,0 +1,339 @@ +//! exl3-tpprefill: the expert-major deep-chunk prefill (cuda_exl3_prefill.zig) at tp=1 and at tp=2 expert parallel +const std = @import("std"); +const cuda = @import("cuda"); +const flash = @import("flashnext"); +const check = @import("../../check.zig"); + +const Gpu = check.Gpu; +const prompt = flash.prompt; +const X = flash.exl3; +const xpass = X.pass; +const xpre = X.prefill; +const wire = flash.rank_wire; +const Collective = flash.collective.Collective; +const vocab = prompt.vocabulary; + +fn promote(bits: u16) f32 { + return @bitCast(@as(u32, bits) << 16); +} + +fn ms(t0: i96, t1: i96) f64 { + return @as(f64, @floatFromInt(t1 - t0)) / 1e6; +} + +fn gib(b: u64) f64 { + return @as(f64, @floatFromInt(b)) / (1 << 30); +} + +const Ranks = struct { + rank: usize, + link: wire.Link, + collective: Collective, + seq: u64 = 10, + + fn open(gpu: Gpu, rank: usize, port: u16) !*Ranks { + const r = try gpu.gpa.create(Ranks); + errdefer gpu.gpa.destroy(r); + const link = try wire.Link.open(gpu.io, "127.0.0.1", port, rank, 600_000); + errdefer link.close(); + r.* = .{ .rank = rank, .link = link, .collective = try Collective.init(link, rank, 3_600_000) }; + return r; + } + + fn close(r: *Ranks, gpa: std.mem.Allocator) void { + r.collective.deinit(); + r.link.close(); + gpa.destroy(r); + } + + fn swap(r: *Ranks, value: u64) !u64 { + r.seq += 1; + const until = wire.deadline(3_600_000); + if (r.rank == 0) { + const got = try r.link.receiveHeader(until); + if (got.kind != .ready or got.sequence != r.seq) return error.RankOutOfStep; + try r.link.sendHeader(.{ .kind = .ready, .sequence = r.seq, .arg = value }, until); + return got.arg; + } + try r.link.sendHeader(.{ .kind = .ready, .sequence = r.seq, .arg = value }, until); + const got = try r.link.receiveHeader(until); + if (got.kind != .ready or got.sequence != r.seq) return error.RankOutOfStep; + return got.arg; + } +}; + +const Eng = struct { + gpu: Gpu, + rank: usize, + world: usize, + ranks: ?*Ranks, + stream: cuda.Stream, + glue: flash.Glue, + e: *prompt.Engine, + seq: prompt.Seq, + out: cuda.DeviceBuffer, + attached: xpass.Attached, + + fn init(gpu: Gpu, rank: usize, world: usize, port: u16, dir: []const u8, cap: usize, deep_rows: usize) !*Eng { + const g = try gpu.gpa.create(Eng); + errdefer gpu.gpa.destroy(g); + g.gpu = gpu; + g.rank = rank; + g.world = world; + g.ranks = if (world == 2) try Ranks.open(gpu, rank, port) else null; + g.stream = try cuda.Stream.init(gpu.d, true); + g.glue = try flash.Glue.load(gpu.d); + const t0 = check.now(gpu.io); + g.e = try prompt.Engine.initRank(gpu.gpa, gpu.io, @constCast(gpu.d), &g.stream, dir, world, rank); + if (g.ranks) |r| g.e.collective = &r.collective; + var free: [6]usize = undefined; + var total: usize = 0; + try gpu.d.check(gpu.d.api.cuMemGetInfo_v2(&free[0], &total), "cuMemGetInfo"); + _ = try g.e.reserve(@min(cap, g.e.chunk)); + _ = try g.e.reserveScores(cap, @min(cap, g.e.chunk)); + try gpu.d.check(gpu.d.api.cuMemGetInfo_v2(&free[1], &total), "cuMemGetInfo"); + g.seq = try g.e.newSeq(cap); + try gpu.d.check(gpu.d.api.cuMemGetInfo_v2(&free[2], &total), "cuMemGetInfo"); + g.out = try cuda.DeviceBuffer.alloc(gpu.d, vocab * 2); + if (deep_rows > 0) try xpre.reserveAt(gpu.gpa, gpu.d, g.e, deep_rows, world); + try gpu.d.check(gpu.d.api.cuMemGetInfo_v2(&free[3], &total), "cuMemGetInfo"); + g.attached = try xpass.attach(gpu.gpa, gpu.io, gpu.d, g.stream, g.e.layers[0..g.e.n_layers], null, g.e.chunk, try xpass.budgetFromEnv(0), null); + try gpu.d.check(gpu.d.api.cuMemGetInfo_v2(&free[4], &total), "cuMemGetInfo"); + std.debug.print("MEM rank {d}/{d}: loaded in {d:.1} s; card {d:.2} GiB; free after weights {d:.2}; pass scratch (chunk {d}) + scores {d:.2} GiB; caches for {d} tokens {d:.2} GiB; deep chunks of {d} rows {d:.2} GiB (layout {d:.2}); pool {d} lines a layer x {d} layers = {d:.2} GiB (line {d:.2} MiB, ring {d} slots); {d:.2} GiB free after; head columns {d}\n", .{ rank, world, ms(t0, check.now(gpu.io)) / 1e3, gib(total), gib(free[0]), g.e.chunk, gib(free[0] - free[1]), cap, gib(free[1] - free[2]), deep_rows, gib(free[2] - free[3]), gib(xpre.deviceBytesAt(deep_rows, world)), g.attached.lines, g.attached.layers, gib(g.attached.device_bytes), @as(f64, @floatFromInt(g.attached.line_bytes)) / (1 << 20), g.attached.lines * 48, gib(free[4]), g.e.head_q.n }); + return g; + } + + fn deinit(g: *Eng) void { + xpre.release(); + g.out.free(); + g.seq.deinit(g.gpu.gpa); + g.e.deinit(); + g.glue.deinit(); + g.stream.deinit(); + if (g.ranks) |r| r.close(g.gpu.gpa); + g.gpu.gpa.destroy(g); + } + + fn tri(g: *Eng) flash.Tri { + return .{ .own = &g.glue, .s = g.stream }; + } + + fn reset(g: *Eng) !void { + const s = &g.seq; + for ([_][]cuda.DeviceBuffer{ s.conv, s.state, s.kc, s.vc, s.ikc, s.pooled }) |bs| for (bs) |b| try b.fill8(0, g.stream.handle); + try s.tail.fill8(0, g.stream.handle); + try g.stream.synchronize(); + s.fed = 0; + s.history = .{ g.e.eos, g.e.eos }; + } + + fn swap(g: *Eng, v: u64) !u64 { + return if (g.ranks) |r| r.swap(v) else v; + } + + fn finish(g: *Eng) !void { + if (g.ranks) |r| try r.collective.finish(g.stream) else try g.stream.synchronize(); + } + + fn forward(g: *Eng, ids: []const u32) !void { + try g.e.forward(flash.Tri, g.tri(), &g.seq, ids, .{ .logits = .last }); + try g.finish(); + } + + fn round(g: *Eng, tok: u32) !u32 { + const q = [_]*prompt.Seq{&g.seq}; + var t = tok; + try g.e.round(flash.Tri, g.tri(), &q, (&t)[0..1], g.out.ptr); + try g.finish(); + return prompt.argmax(g.seq.logits); + } +}; + +fn top1(row: []const u16) u32 { + var best: usize = 0; + for (row, 0..) |b, i| if (promote(b) > promote(row[best])) { + best = i; + }; + return @intCast(best); +} + +fn margin(row: []const u16) f32 { + const b = top1(row); + var second: f32 = -std.math.inf(f32); + for (row, 0..) |v, i| if (i != b) { + second = @max(second, promote(v)); + }; + return promote(row[b]) - second; +} + +fn top5(row: []const u16) [5]u32 { + var ids: [5]u32 = @splat(0); + var vals: [5]f32 = @splat(-std.math.inf(f32)); + for (row, 0..) |b, i| { + const v = promote(b); + if (!(v > vals[4])) continue; + var at: usize = 4; + while (at > 0 and v > vals[at - 1]) : (at -= 1) { + vals[at] = vals[at - 1]; + ids[at] = ids[at - 1]; + } + vals[at] = v; + ids[at] = @intCast(i); + } + return ids; +} + +const Run = struct { ms: f64, logits: []u16, greedy: []u32, decode_ms: f64, stats: xpre.Stats }; + +/// A fresh sequence, `ids` prefilled (deep or the batched pass), then `steps` greedy rounds. +fn once(g: *Eng, ids: []const u32, deep: bool, rows: usize, steps: usize) !Run { + const a = g.gpu.gpa; + var o = xpre.optionsFromEnv(); + o.enabled = deep; + o.rows = rows; + xpre.state.options = o; + xpre.state.stats = .{}; + try g.reset(); + _ = try g.swap(ids.len); + const t0 = check.now(g.gpu.io); + try g.forward(ids); + const t = ms(t0, check.now(g.gpu.io)); + const st = xpre.state.stats; + const logits = try a.alloc(u16, vocab); + @memcpy(std.mem.sliceAsBytes(logits), g.seq.logits); + const greedy = try a.alloc(u32, steps + 1); + greedy[0] = top1(logits); + const t1 = check.now(g.gpu.io); + for (1..steps + 1) |i| greedy[i] = try g.round(greedy[i - 1]); + const dt = if (steps > 0) ms(t1, check.now(g.gpu.io)) / @as(f64, @floatFromInt(steps)) else 0; + return .{ .ms = t, .logits = logits, .greedy = greedy, .decode_ms = dt, .stats = st }; +} + +fn freeRun(a: std.mem.Allocator, r: Run) void { + a.free(r.logits); + a.free(r.greedy); +} + +/// `exact`: the same bytes and greedy ids (deep chunks against the batched pass), else the same real-margin top-1. +fn compare(what: []const u8, n: usize, ref: []const u16, got: []const u16, ref_g: []const u32, got_g: []const u32, exact: bool) bool { + var dmax: f32 = 0; + var peak: f32 = 0; + var same: usize = 0; + for (ref, got) |x, y| { + dmax = @max(dmax, @abs(promote(x) - promote(y))); + peak = @max(peak, @abs(promote(x))); + same += @intFromBool(x == y); + } + var first: ?usize = null; + for (ref_g, got_g, 0..) |x, y, i| if (x != y and first == null) { + first = i; + }; + const t_eq = top1(ref) == top1(got); + std.debug.print("CMP {s} {d} tokens: last-row logits max |d| {d:.4} (peak {d:.2}, {d:.1}% bit-equal); top-1 {d} vs {d} ({s}, reference margin {d:.4}); greedy {d}: {s}", .{ what, n, dmax, peak, 100 * @as(f64, @floatFromInt(same)) / vocab, top1(ref), top1(got), if (t_eq) "equal" else "DIFFERENT", margin(ref), ref_g.len, if (first == null) "identical" else "differ" }); + if (first) |f| std.debug.print(" from step {d} ({d} vs {d})", .{ f, ref_g[f], got_g[f] }); + std.debug.print("\n", .{}); + if (!t_eq) { + // the two candidates in both rows, and the top-5 overlap + const a_ = top1(ref); + const b_ = top1(got); + const r5 = top5(ref); + const g5 = top5(got); + var overlap: usize = 0; + for (r5) |x| for (g5) |y| { + overlap += @intFromBool(x == y); + }; + std.debug.print(" reference: logit[{d}] {d:.4}, logit[{d}] {d:.4}; ours: logit[{d}] {d:.4}, logit[{d}] {d:.4}; top-5 overlap {d}/5 (reference {any}, ours {any})\n", .{ a_, promote(ref[a_]), b_, promote(ref[b_]), a_, promote(got[a_]), b_, promote(got[b_]), overlap, r5, g5 }); + } + if (exact) return same == vocab and first == null; + return t_eq or margin(ref) < 0.25; +} + +pub fn run(gpu: Gpu, rest: []const [:0]const u8) !void { + if (rest.len < 10) return error.MissingArgument; + const a = gpu.gpa; + const rank = try std.fmt.parseInt(usize, rest[0], 10); + const world = try std.fmt.parseInt(usize, rest[1], 10); + const port = try std.fmt.parseInt(u16, rest[2], 10); + const dir = rest[3]; + const text = try std.Io.Dir.cwd().readFileAlloc(gpu.io, rest[4], a, .limited(1 << 28)); + defer a.free(text); + const all_ids = try a.alloc(u32, text.len / 4); + defer a.free(all_ids); + for (all_ids, 0..) |*x, i| x.* = std.mem.readInt(u32, text[4 * i ..][0..4], .little); + var ns: [16]usize = undefined; + var n_ns: usize = 0; + var it = std.mem.splitScalar(u8, rest[5], ','); + while (it.next()) |tok| { + ns[n_ns] = try std.fmt.parseInt(usize, tok, 10); + n_ns += 1; + } + const rows = try std.fmt.parseInt(usize, rest[6], 10); + const cap = try std.fmt.parseInt(usize, rest[7], 10); + const both = std.mem.eql(u8, rest[8], "both"); + const dump = rest[9]; + var steps: usize = 16; + if (std.c.getenv("TF_TEST_GREEDY")) |v| steps = std.fmt.parseInt(usize, std.mem.span(v), 10) catch steps; + var both_max: usize = std.math.maxInt(usize); + if (std.c.getenv("TF_TEST_BOTH_MAX")) |v| both_max = std.fmt.parseInt(usize, std.mem.span(v), 10) catch both_max; + for (ns[0..n_ns]) |n| if (n + steps + 1 > cap or n > all_ids.len) return error.NotEnoughRoom; + if ((world != 1 and world != 2) or rank >= world) return error.BadRank; + const g = try Eng.init(gpu, rank, world, port, dir, cap, rows); + defer g.deinit(); + // warm-up: kernels, tiles, the ring + { + const w = try once(g, all_ids[0..@min(1024, all_ids.len)], true, rows, 1); + freeRun(a, w); + } + var all_ok = true; + for (ns[0..n_ns]) |n| { + const ids = all_ids[0..n]; + const d = try once(g, ids, true, rows, steps); + defer freeRun(a, d); + const s = d.stats; + std.debug.print("PREFILL rank {d}/{d} {d:6} tokens deep: {d:9.1} ms = {d:7.1} rows/s; chunks {d}, passes {d}, copies {d} ({d} prefetched), all-reduces {d}; host picks wait {d:.1} ms, plan {d:.1} ms; decode after {d:.2} ms/token\n", .{ rank, world, n, d.ms, @as(f64, @floatFromInt(n)) / (d.ms / 1e3), s.chunks, s.passes, s.copies, s.prefetched, s.reduces, @as(f64, @floatFromInt(s.host_wait_ns)) / 1e6, @as(f64, @floatFromInt(s.plan_ns)) / 1e6, d.decode_ms }); + if (s.reduce_ns > 0) std.debug.print(" rank {d}: combine + all-reduce + place {d:.1} ms over {d} layers (profiled, each waited)\n", .{ rank, @as(f64, @floatFromInt(s.reduce_ns)) / 1e6, s.reduces }); + if (both and n <= both_max) { + const b = try once(g, ids, false, rows, steps); + defer freeRun(a, b); + std.debug.print("PREFILL rank {d}/{d} {d:6} tokens batched: {d:9.1} ms = {d:7.1} rows/s; decode after {d:.2} ms/token\n", .{ rank, world, n, b.ms, @as(f64, @floatFromInt(n)) / (b.ms / 1e3), b.decode_ms }); + if (rank == 0) all_ok = compare("deep vs batched (same world)", n, b.logits, d.logits, b.greedy, d.greedy, true) and all_ok; + } + // both ranks hold the same logits and greedy ids + const h = std.hash.Wyhash.hash(std.hash.Wyhash.hash(0, std.mem.sliceAsBytes(d.logits)), std.mem.sliceAsBytes(d.greedy)); + const other = try g.swap(h); + if (world == 2 and other != h) { + std.debug.print("rank {d}: the ranks' logits or greedy ids DIFFER at {d} tokens\n", .{ rank, n }); + all_ok = false; + } + if (rank != 0) continue; + var pb: [512]u8 = undefined; + const path = try std.fmt.bufPrint(&pb, "{s}-{d}.bin", .{ dump, n }); + if (world == 1) { + const bytes = try a.alloc(u8, vocab * 2 + d.greedy.len * 4); + defer a.free(bytes); + @memcpy(bytes[0 .. vocab * 2], std.mem.sliceAsBytes(d.logits)); + @memcpy(bytes[vocab * 2 ..], std.mem.sliceAsBytes(d.greedy)); + try std.Io.Dir.cwd().writeFile(gpu.io, .{ .sub_path = path, .data = bytes }); + continue; + } + const ref = std.Io.Dir.cwd().readFileAlloc(gpu.io, path, a, .limited(1 << 24)) catch |err| { + std.debug.print("no tp=1 dump at {s} ({t})\n", .{ path, err }); + continue; + }; + defer a.free(ref); + if (ref.len != vocab * 2 + d.greedy.len * 4) { + std.debug.print("tp=1 dump {s} has another greedy length; skipped\n", .{path}); + continue; + } + const rl = try a.alloc(u16, vocab); + defer a.free(rl); + @memcpy(std.mem.sliceAsBytes(rl), ref[0 .. vocab * 2]); + const rg = try a.alloc(u32, d.greedy.len); + defer a.free(rg); + @memcpy(std.mem.sliceAsBytes(rg), ref[vocab * 2 ..]); + all_ok = compare("tp=2 deep vs tp=1 deep", n, rl, d.logits, rg, d.greedy, false) and all_ok; + } + const both_ok = try g.swap(@intFromBool(all_ok)); + try check.expect(all_ok and both_ok == 1, "rank {d}: a top-1 differs past bf16 noise or the ranks disagree", .{rank}); + check.pass("rank {d}/{d}: deep prefill measured and compared", .{ rank, world }); +} diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index 1f90b32f7e..26f60a1d30 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -100,6 +100,7 @@ const exl3_ple2_tests = @import("flashnext_exl3/model/exl3_ple2_tests.zig"); const exl3_draft_tests = @import("flashnext_exl3/tp/exl3_draft_tests.zig"); const exl3_pool_tests = @import("flashnext_exl3/offload/exl3_pool_tests.zig"); const exl3_prefill_tests = @import("flashnext_exl3/prefill/exl3_prefill_tests.zig"); +const exl3_tpprefill_tests = @import("flashnext_exl3/tp/exl3_tpprefill_tests.zig"); const flashnext_exl3_tests = @import("flashnext_exl3/model/flashnext_exl3_tests.zig"); const flashnext_exl3_tp_tests = @import("flashnext_exl3/tp/flashnext_exl3_tp_tests.zig"); const flashnext_exl3_engine_tests = @import("flashnext_exl3/model/flashnext_exl3_engine_tests.zig"); @@ -240,9 +241,12 @@ const usage = \\ exl3-ple2 DIR [ROW] [N] EXL3 n-gram table through cuda_ple.Ngram (ids, rows, RSS) and the PLE kernels vs float64 \\ exl3-draft DIR [THREADS] EXL3 draft head rows dequantised from lm_head vs the full head (Dense) and float64 \\ exl3-prefill DIR experts|copy [LAYER]|dense EXL3 prefill pieces on the pack: grouped experts by rows, expert-major passes, line copies, prompt Dense + \\ exl3-prefill DIR equal IDS N ROWS[,ROWS] [LINES] deep prefill at each chunk size vs the batched pass: last-row logits bytes and greedy tokens + \\ exl3-tpprefill RANK WORLD PORT DIR IDS NS ROWS CAP deep|both DUMP EXL3 deep prefill at tp=1/2 vs the batched pass and tp=1 \\ exl3-pool DIR [equal|bench|hit FILES] [LAYERS] Flash Next EXL3 expert offload pool on Experts' pointer tables vs fully resident \\ exl3-tiles [DIR] EXL3's tile candidate: codebooks and tiles bit-equal, FP64, rows through qlinear and Dense \\ train-pack MODEL EXAMPLES [COUNT] training passes holding several sequences against one a pass: differences, time + \\ exl3-prefill-contract the prefill expert tiles alone through the contract (gates 1-4 and the control) \\ kernel-contract the registered kernels through the contract's gates, each planted control red (needs -Dnvcc) \\ qmm-prefill-many-contract the GB10 bit-exact multi-projection prompt launch against separate launches \\ qmm-prefill-tile22-contract the GB10 large-prompt 2x2 tile against the default tile @@ -351,6 +355,7 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "flashnext-dense-hc")) return flashnext_dense_hc_tests.run(gpu); if (std.mem.eql(u8, cmd, "flashnext-hc")) return flashnext_hc_tests.run(gpu); if (std.mem.eql(u8, cmd, "flashnext-attn")) return flashnext_attn_tests.run(gpu); + if (std.mem.eql(u8, cmd, "exl3-prefill-contract")) return @import("contract_flashnext_exl3_prefill.zig").run(gpu); if (std.mem.eql(u8, cmd, "flashnext-exl3")) return flashnext_exl3_tests.run(gpu, rest); if (std.mem.eql(u8, cmd, "flashnext-exl3-tp")) return flashnext_exl3_tp_tests.run(gpu, rest); if (std.mem.eql(u8, cmd, "flashnext-exl3-engine")) return flashnext_exl3_engine_tests.run(gpu, rest); @@ -413,6 +418,7 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "exl3-draft")) return exl3_draft_tests.run(gpu, rest); if (std.mem.eql(u8, cmd, "exl3-pool")) return exl3_pool_tests.run(gpu, rest); if (std.mem.eql(u8, cmd, "exl3-prefill")) return exl3_prefill_tests.run(gpu, rest); + if (std.mem.eql(u8, cmd, "exl3-tpprefill")) return exl3_tpprefill_tests.run(gpu, rest); if (std.mem.eql(u8, cmd, "nccl-lat")) return exl3_pack_tests.ncclLatency(gpu, rest); if (std.mem.eql(u8, cmd, "exl3-tiles")) return exl3_candidate_tests.run(gpu, if (rest.len > 0) rest[0] else null); if (std.mem.eql(u8, cmd, "exl3-tp2")) return exl3_tp2_tests.run(gpu, try arg(rest, 0), try arg(rest, 1), if (rest.len > 2) rest[2] else "none"); diff --git a/zig/tests/exl3_prefill_body_compile.zig b/zig/tests/exl3_prefill_body_compile.zig new file mode 100644 index 0000000000..7302f7f285 --- /dev/null +++ b/zig/tests/exl3_prefill_body_compile.zig @@ -0,0 +1,5 @@ +const flash = @import("flashnext"); + +export fn exl3PrefillForwardBody(engine: *flash.prompt.Engine, sequence: *flash.prompt.Seq, launch: *const flash.Tri, ids: [*]const u32, count: usize) void { + engine.forward(flash.Tri, launch.*, sequence, ids[0..count], .{ .logits = .last }) catch {}; +}