From 688922781f2639c1ec711257a6bd6ac24cd0698c 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/7] 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 | 39 +- .../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 | 94 +++ .../flashnext/exl3/kv/cuda_exl3_kv_test.zig | 38 ++ .../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 | 530 ++++++++++++++++ .../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 | 127 ++++ .../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, 18001 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 b30ed4c065..56da1d836f 100644 --- a/docs/kernel-contract.md +++ b/docs/kernel-contract.md @@ -256,6 +256,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 cfb83c933b..bd1fcb24f8 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -450,7 +450,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 04bded45c9..24c55dff5f 100644 --- a/zig/build/cuda_kernels.zig +++ b/zig/build/cuda_kernels.zig @@ -122,6 +122,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 288899e034..f9b1bf286b 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 ac87ea207c..eff8c7f4ce 100644 --- a/zig/src/core/contract/registry.zig +++ b/zig/src/core/contract/registry.zig @@ -319,6 +319,7 @@ pub const cuda_sets = [_]CudaSet{ .{ .release = .{ .major = 13, .minor = 3 }, .usage = &([_]Override{ .{ .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 ++ @import("minimax_cuda_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); @@ -574,7 +575,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("exl3_grouped_records.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("exl3_grouped_records.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 743c4b1a8c..8ad6420343 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 }, 121, 45 }, .{ .{ .major = 13, .minor = 2 }, .{ .major = 13, .minor = 2 }, 121, 43 }, .{ .{ .major = 13, .minor = 3 }, .{ .major = 13, .minor = 3 }, 122, 43 }, - .{ .{ .major = 14, .minor = 0 }, .{ .major = 13, .minor = 3 }, 122, 43 }, + .{ .{ .major = 13, .minor = 4 }, .{ .major = 13, .minor = 4 }, 122, 43 }, + .{ .{ .major = 14, .minor = 0 }, .{ .major = 13, .minor = 4 }, 122, 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 db13a74f50..276b6bc5c7 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 69ac0f595d..2aad15ecbc 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 }; } @@ -321,6 +331,9 @@ pub fn Namespace(comptime Tri: type) type { }; const last = stop == ids.len; self.engine.ngram.warmAhead(ids, stop, @min(ids.len, stop + rows)); + 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 5830e8705f..8913e3f136 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,13 @@ pub fn open(comptime F: type, gpa: std.mem.Allocator, io: std.Io, ctx: *const cu const chunk = try @import("cuda_chunk.zig").soloRows(try prompt.chunkRows()); // the largest pass, solo or not _ = 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 + try @import("exl3/kv/cuda_exl3_kv.zig").refuseSolo(mode == .exl3, std.c.getenv("TF_FLASHNEXT_CHUNK_SOLO") != null); 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); @@ -63,11 +67,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; } @@ -81,18 +90,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, @@ -100,7 +129,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 27d78bba79..a75f9a06ce 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 fc70793092..98681c0ef5 100644 --- a/zig/src/families/flashnext/cuda_prompt.zig +++ b/zig/src/families/flashnext/cuda_prompt.zig @@ -116,6 +116,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 { @@ -136,6 +137,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: TF_FLASHNEXT_CHUNK_SOLO if set, else TF_FLASHNEXT_CHUNK (cuda_chunk.zig). @@ -145,10 +147,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); } @@ -165,8 +171,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); @@ -176,7 +184,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); @@ -186,8 +195,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 { @@ -350,10 +360,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 { @@ -377,19 +384,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); @@ -422,7 +432,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| { @@ -505,23 +517,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); } }; @@ -529,55 +525,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); } pub const parseChunk = chunks.parseChunk; 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..d450b5c8bf --- /dev/null +++ b/zig/src/families/flashnext/exl3/kv/cuda_exl3_kv.zig @@ -0,0 +1,94 @@ +//! 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; + +/// One prompt cut whatever else decodes (chunks never depend on batching): busy and solo passes take the same rows. +pub fn cut(e: anytype, chunk_set: bool) void { + if (!chunk_set) e.chunk = exl3_chunk; + e.busy_chunk = e.chunk; +} + +/// TF_FLASHNEXT_CHUNK_SOLO would cut a lone stream's prompt apart from a batched one's. +pub fn refuseSolo(exl3: bool, solo_set: bool) !void { + if (exl3 and solo_set) return error.Exl3ChunkSoloUnsupported; +} + +/// 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; + cut(e, std.c.getenv("TF_FLASHNEXT_CHUNK") != null); +} + +/// 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..551a2f44e8 --- /dev/null +++ b/zig/src/families/flashnext/exl3/kv/cuda_exl3_kv_test.zig @@ -0,0 +1,38 @@ +//! 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()); +} + +test "exl3: kv a prompt cut is the same busy or solo, with or without TF_FLASHNEXT_CHUNK" { + const Rows = struct { chunk: usize, busy_chunk: usize }; + var unset = Rows{ .chunk = 1536, .busy_chunk = 1536 }; + hub.cut(&unset, false); + try std.testing.expectEqual(hub.exl3_chunk, unset.chunk); + try std.testing.expectEqual(unset.chunk, unset.busy_chunk); + var set = Rows{ .chunk = 4096, .busy_chunk = 4096 }; + hub.cut(&set, true); + try std.testing.expectEqual(@as(usize, 4096), set.chunk); + try std.testing.expectEqual(set.chunk, set.busy_chunk); +} + +test "exl3: kv TF_FLASHNEXT_CHUNK_SOLO is refused for an EXL3 pack only" { + try std.testing.expectError(error.Exl3ChunkSoloUnsupported, hub.refuseSolo(true, true)); + try hub.refuseSolo(true, false); + try hub.refuseSolo(false, true); +} 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..1685a86349 --- /dev/null +++ b/zig/src/families/flashnext/exl3/pack/cuda_exl3_load.zig @@ -0,0 +1,530 @@ +//! 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); + // exl3.Dense zeroes its split-K counters once, as they grow (each block's last program resets its own) + const counters_grew = if (de.counters) |b| b.len < o.max_counters else o.max_counters > 0; + try grow(o.d, &de.counters, o.max_counters); + if (counters_grew) try de.counters.?.fill8(0, s.handle); + 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..d838a570fe --- /dev/null +++ b/zig/src/families/flashnext/exl3/tp/cuda_exl3_fork.zig @@ -0,0 +1,127 @@ +//! 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); + try dense.counters.?.fill8(0, side.handle); // zeroed once on the branch's stream: each block's last program resets its own + 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 70ad49f13b..dd36d54788 100644 --- a/zig/src/native/cuda.zig +++ b/zig/src/native/cuda.zig @@ -385,7 +385,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; @@ -441,7 +442,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 ca9e662d17..3dac74e654 100644 --- a/zig/tests/cuda/contract_limits.zig +++ b/zig/tests/cuda/contract_limits.zig @@ -27,6 +27,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", "nemotron_ops", "minimax_m3", "minimax_m3_attn", "moss_image" }) |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 a3ad2d2dad..366e15af57 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -97,6 +97,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"); @@ -137,6 +143,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 @@ -226,6 +234,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) @@ -342,6 +354,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); @@ -395,6 +409,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 6a6f12061b490290d5cacc600190baba57bf8e92 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/7] 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 | 5 +- .../cuda_flashnext_exl3_refusal_test.zig | 19 + zig/src/native/cuda_refusal.zig | 29 +- .../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 | 3 +- .../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, 3575 insertions(+), 241 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/src/native/cuda_flashnext_exl3_refusal_test.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 56da1d836f..eb72dbf708 100644 --- a/docs/kernel-contract.md +++ b/docs/kernel-contract.md @@ -256,8 +256,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 24c55dff5f..6b3cd75ab7 100644 --- a/zig/build/cuda_kernels.zig +++ b/zig/build/cuda_kernels.zig @@ -123,6 +123,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 8b66e200aa..ea25859d55 100644 --- a/zig/kernels/cuda/flashnext_batch.cu +++ b/zig/kernels/cuda/flashnext_batch.cu @@ -324,6 +324,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 276b6bc5c7..b1e2d15683 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 2aad15ecbc..85b5f19cc1 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 8913e3f136..d8d5add901 100644 --- a/zig/src/families/flashnext/cuda_lanes_load.zig +++ b/zig/src/families/flashnext/cuda_lanes_load.zig @@ -23,7 +23,7 @@ pub fn open(comptime F: type, gpa: std.mem.Allocator, io: std.Io, ctx: *const cu 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 try @import("exl3/kv/cuda_exl3_kv.zig").refuseSolo(mode == .exl3, std.c.getenv("TF_FLASHNEXT_CHUNK_SOLO") != null); - 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(); @@ -48,6 +48,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(); @@ -72,8 +74,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); @@ -95,7 +96,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 @@ -104,15 +105,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 a75f9a06ce..9ba07c4106 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 8fe177ec9a..35202a1724 100644 --- a/zig/src/families/flashnext/cuda_mtp.zig +++ b/zig/src/families/flashnext/cuda_mtp.zig @@ -202,8 +202,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 993e5328ce..32b3bf7a70 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. @@ -412,6 +412,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 98681c0ef5..68b8b92285 100644 --- a/zig/src/families/flashnext/cuda_prompt.zig +++ b/zig/src/families/flashnext/cuda_prompt.zig @@ -138,6 +138,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: TF_FLASHNEXT_CHUNK_SOLO if set, else TF_FLASHNEXT_CHUNK (cuda_chunk.zig). @@ -172,6 +174,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); @@ -391,7 +394,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 d450b5c8bf..27c2193e82 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; @@ -65,8 +72,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; cut(e, std.c.getenv("TF_FLASHNEXT_CHUNK") != null); + 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. @@ -81,7 +99,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, }); } @@ -92,3 +110,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 1685a86349..34998a299c 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); } @@ -136,11 +139,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 }; @@ -156,6 +160,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; } @@ -312,14 +317,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; @@ -362,6 +374,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 = &.{}, @@ -446,9 +460,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 dd36d54788..8c9bf0a026 100644 --- a/zig/src/native/cuda.zig +++ b/zig/src/native/cuda.zig @@ -34,7 +34,7 @@ const registry = .{ nemotron.native, nemotron.omni, glm.native, minimax.native, /// --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" }; +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 }; @@ -211,7 +211,7 @@ 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| { + if (refusals.open(F, gpa, io, o)) |why| { problem.* = why; return null; } @@ -575,6 +575,7 @@ test "Flash Next takes --vision at family admission" { } test { _ = @import("cuda_health_test.zig"); + _ = @import("cuda_flashnext_exl3_refusal_test.zig"); _ = @import("cuda_vision_offer.zig"); _ = kv_pages; _ = drafted; diff --git a/zig/src/native/cuda_flashnext_exl3_refusal_test.zig b/zig/src/native/cuda_flashnext_exl3_refusal_test.zig new file mode 100644 index 0000000000..4313bfb755 --- /dev/null +++ b/zig/src/native/cuda_flashnext_exl3_refusal_test.zig @@ -0,0 +1,19 @@ +//! Flash Next EXL3 rank admission and draft-cache refusals. +const std = @import("std"); +const api = @import("engine_api"); +const F = @import("flashnext").native; +const refusals = @import("cuda_refusal.zig"); + +test "Flash Next serves --tp for EXL3 packs only and keeps the tensor-rank draft refusal for any other pack" { + 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 d86d1bbfb2..e602284b0b 100644 --- a/zig/src/native/cuda_refusal.zig +++ b/zig/src/native/cuda_refusal.zig @@ -9,6 +9,7 @@ const glm_model_type = "glm5_next"; // GLM-5.3-Flash's (glm.cuda_native.model_ty 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 (std.mem.eql(u8, o.model_type, glm_model_type)) { @@ -39,12 +40,23 @@ 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); +} + +/// Apply family and rank refusals after identifying a multi-rank EXL3 checkpoint. +pub fn open(comptime F: type, gpa: std.mem.Allocator, io: std.Io, o: api.Open) ?[]const u8 { + const exl3_pack = if (comptime @hasDecl(F, "exl3_pack")) o.tp > 1 and F.exl3_pack(gpa, io, o.dir) else false; + return admitted(F, o, exl3_pack) orelse packRanks(F, o, exl3_pack); +} + +/// `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 @@ -61,6 +73,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 }; @@ -72,6 +90,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 3dac74e654..793de75e90 100644 --- a/zig/tests/cuda/contract_limits.zig +++ b/zig/tests/cuda/contract_limits.zig @@ -27,7 +27,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", "nemotron_ops", "minimax_m3", "minimax_m3_attn", "moss_image" }) |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 b9158471f4..d7a32a1524 100644 --- a/zig/tests/cuda/contract_tests.zig +++ b/zig/tests/cuda/contract_tests.zig @@ -62,7 +62,6 @@ const q35_contiguous = @import("contract_q35_contiguous.zig"); const contract = core.contract; const g = contract.gates; const kk = cuda.kernels; - const images = struct { pub const controls align(16) = @embedFile("fatbin_contract_controls").*; pub const h3 align(16) = @embedFile("fatbin_h3").*; @@ -253,6 +252,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 366e15af57..6fe49038a5 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -99,10 +99,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"); @@ -144,7 +147,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 @@ -236,6 +241,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 @@ -355,7 +361,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); @@ -411,8 +419,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 2924a922f6325376c360600eb1864fbe276d6d63 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/7] 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 eb72dbf708..4c1dae23ad 100644 --- a/docs/kernel-contract.md +++ b/docs/kernel-contract.md @@ -256,7 +256,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 bd1fcb24f8..eba95dda8b 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -365,7 +365,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, check_mods.qwen35.import_table.get("engine_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 6b3cd75ab7..e5cbdcb67c 100644 --- a/zig/build/cuda_kernels.zig +++ b/zig/build/cuda_kernels.zig @@ -124,6 +124,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..adabcf56ff --- /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_commit; +using tf_exl3_prompt::cp_wait; +using tf_exl3_prompt::ldmatrix4; +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_commit(); + } + int left = chain_chunks, w = 0; + for (int c = 0; c < chunks; ++c) { + cp_wait(); + __syncthreads(); + if (c + TSTAGES - 1 < chunks) load(c + TSTAGES - 1); + cp_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; + ldmatrix4(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_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 da5c535dce..021a82eead 100644 --- a/zig/src/cuda/exl3.zig +++ b/zig/src/cuda/exl3.zig @@ -265,19 +265,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 b1e2d15683..132624ecc7 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 85b5f19cc1..8196e0299a 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); @@ -338,7 +338,7 @@ pub fn Namespace(comptime Tri: type) type { return error.Cancelled; } const rows = @import("cuda_chunk.zig").passRows(&self.lanes, l, self.engine.chunk, self.engine.busy_chunk); - var stop = @min(ids.len, at + rows); + var stop = @min(ids.len, at + @import("exl3/prefill/cuda_exl3_prefill.zig").stepRows(self.engine, rows, 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 d8d5add901..0b38a490bc 100644 --- a/zig/src/families/flashnext/cuda_lanes_load.zig +++ b/zig/src/families/flashnext/cuda_lanes_load.zig @@ -92,11 +92,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 @@ -120,11 +123,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, @@ -132,7 +136,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, own.engine.chunk, 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 32b3bf7a70..eeb1a42107 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 68b8b92285..0d22973c8c 100644 --- a/zig/src/families/flashnext/cuda_prompt.zig +++ b/zig/src/families/flashnext/cuda_prompt.zig @@ -246,6 +246,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); @@ -403,6 +404,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..87f443cd48 --- /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 }); +} + +/// A serving step's batched `base` rows, or the deep chunk when this engine holds deep buffers. +pub fn stepRows(e: *const prompt.Engine, base: usize, draft_rows: ?usize) usize { + var n = base; + 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, e.chunk, 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 793de75e90..bb48bb5e08 100644 --- a/zig/tests/cuda/contract_limits.zig +++ b/zig/tests/cuda/contract_limits.zig @@ -27,7 +27,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", "nemotron_ops", "minimax_m3", "minimax_m3_attn", "moss_image" }) |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 6fe49038a5..bdb6bb3a5b 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -102,6 +102,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"); @@ -243,9 +244,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 @@ -360,6 +364,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); @@ -422,6 +427,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 {}; +} From 70b4527f538e486c9a9feac216f2e6fd9b5384e4 Mon Sep 17 00:00:00 2001 From: BlivionIaG <15323295+BlivionIaG@users.noreply.github.com> Date: Tue, 27 Oct 2026 16:32:38 +0100 Subject: [PATCH 4/7] EXL3 Flash Next: KV pages from the core pool TF_FLASHNEXT_KV_PAGES now serves the EXL3 checkpoint, one rank or two: streams beside the first take pages of n positions from the core pool, as the NVFP4 path does. At TP=2 each rank owns a pool and the ranks agree lanes, pages and run pages at attach, so a create frame flagged for the pool opens the same page ids on both. Split attention heads stay refused with pages. --- zig/src/families/flashnext/cuda_kv_lanes.zig | 28 ++++++++ .../flashnext/cuda_kv_lifetime_test.zig | 64 +++++++++++++++++++ zig/src/families/flashnext/cuda_lanes.zig | 6 +- .../families/flashnext/cuda_lanes_explain.zig | 2 +- .../families/flashnext/cuda_lanes_load.zig | 7 +- zig/src/families/flashnext/cuda_mtp_draft.zig | 5 ++ zig/src/families/flashnext/cuda_mtp_ops.zig | 4 +- zig/src/families/flashnext/cuda_prompt.zig | 6 +- .../families/flashnext/cuda_rank_engine.zig | 55 ++++++++++++---- .../exl3/prefill/cuda_exl3_prefill_head.zig | 2 +- 10 files changed, 156 insertions(+), 23 deletions(-) diff --git a/zig/src/families/flashnext/cuda_kv_lanes.zig b/zig/src/families/flashnext/cuda_kv_lanes.zig index 296a396be0..72b53ec61a 100644 --- a/zig/src/families/flashnext/cuda_kv_lanes.zig +++ b/zig/src/families/flashnext/cuda_kv_lanes.zig @@ -53,6 +53,29 @@ pub fn layout(capacity: usize, tokens: usize, head: ?Head, extra: usize) kv.Layo return .{ .fixed_bytes = mostPages(capacity, tokens) * 4 + extra, .token_bytes = pages.pageBytesOf(tokens, layers) / tokens, .page_tokens = @intCast(tokens), .snap_bytes = snap.bytes(0, head != null), .run_bytes = run }; } +/// `l` with every stream's run `k` pages: the ranks of a pair agree the larger, so their pools open the same ids. +pub fn withRunPages(l: kv.Layout, k: usize) kv.Layout { + var out = l; + out.run_bytes = k * l.pageBytes(); + return out; +} + +/// The lanes' pool on `own`; a pair's ranks agree lanes, pages and run pages, so their page ids are equal. +pub fn attachTo(own: anytype, streams: u32, count: usize) !void { + var lay = own.kv.layout.?; + var lane_count = streams; + var page_count = count; + if (own.ranks) |r| { + lane_count = @intCast(try r.room(streams)); + page_count = @intCast(try r.room(count)); + lay = withRunPages(lay, @intCast(~(try r.room(~@as(u64, lay.runPages()))))); + } + own.kv.layout = lay; + own.kv.pool = try Pool.init(own.gpa, own.engine, lay, own.context, own.kv.head, lane_count, page_count); + own.kv.paged_only = pagedOnly() and lane_count > 0; + if (own.ranks) |r| r.pool = own.kv.pool; +} + /// Where a stream's caches came from: no pool, the one stream on its own caches beside a pool, or a pool lane. pub const From = enum { none, own, pool }; @@ -233,6 +256,11 @@ pub const Pool = struct { return mtp.Draft.initIn(p.gpa, p.e.driver, cap, h.draft_n, h.rows, &x.seq.paged.?, p.runBase(x) + seqRunBytes(), run_len - seqRunBytes()); } + /// Whether `seq` is a lane's cache. + pub fn owns(p: *Pool, seq: *const prompt.Seq) bool { + return p.laneOf(seq) != null; + } + fn laneOf(p: *Pool, seq: *const prompt.Seq) ?*Lane { const t = (seq.paged orelse return null).table.ptr; for (p.lanes) |*x| if (x.seq.paged.?.table.ptr == t) return x; diff --git a/zig/src/families/flashnext/cuda_kv_lifetime_test.zig b/zig/src/families/flashnext/cuda_kv_lifetime_test.zig index 8984d54301..062cfbe4de 100644 --- a/zig/src/families/flashnext/cuda_kv_lifetime_test.zig +++ b/zig/src/families/flashnext/cuda_kv_lifetime_test.zig @@ -83,3 +83,67 @@ test "KV lifetime: Flash Next close and deinit retain pages until context proof" try t.expectEqual(@as(usize, 0), D.freed_pending); } } + +test "KV lanes: two ranks whose pages differ in bytes open and close the same page ids in the same order" { + var d = D.driver(); + var e = engine(&d); + defer cuda.completion.drain(&d, e.stream.*, "the test driver must drain before it leaves"); + // rank 0 keeps the MTP head's layer and a bigger run, rank 1 neither: both take the larger run in pages + const head: pages.Head = .{ .draft_n = 4096, .rows = 64 }; + const a = pages.layout(256, 16, head, 0); + const b = pages.layout(256, 16, null, 0); + try t.expect(a.pageBytes() != b.pageBytes()); + const k = @max(a.runPages(), b.runPages()); + const left = pages.withRunPages(a, k); + const right = pages.withRunPages(b, k); + try t.expectEqual(left.runPages(), right.runPages()); + const count = 4 * k + 64; + const p0 = try pages.Pool.init(t.allocator, &e, left, 256, head, 3, count); + defer p0.deinit(); + const p1 = try pages.Pool.init(t.allocator, &e, right, 256, null, 3, count); + defer p1.deinit(); + const caps = [_]usize{ 40, 17, 100 }; + var s0: [3]@import("cuda_prompt.zig").Seq = undefined; + var s1: [3]@import("cuda_prompt.zig").Seq = undefined; + for (caps, 0..) |cap, i| { + s0[i] = try p0.open(cap); + s1[i] = try p1.open(cap); + } + for (0..3) |i| { + try t.expectEqualSlices(u32, p0.lanes[i].lease.ids.items, p1.lanes[i].lease.ids.items); + try t.expectEqual(p0.lanes[i].run, p1.lanes[i].run); + } + try t.expect(p0.close(&s0[1]) and p1.close(&s1[1])); + s0[1] = try p0.open(60); + s1[1] = try p1.open(60); + try t.expectEqualSlices(u32, p0.lanes[1].lease.ids.items, p1.lanes[1].lease.ids.items); + try t.expectEqual(p0.lanes[1].run, p1.lanes[1].run); + for (0..3) |i| { + try t.expect(p0.close(&s0[i]) and p1.close(&s1[i])); + } + try t.expect(p0.pool.idle() and p1.pool.idle()); +} + +test "KV lanes: every pass over a lane's MTP head reads its cache through the lane's pages, across a page boundary" { + const wt = @import("cuda_prompt_wt.zig"); + const pg = @import("cuda_pages.zig"); + var d = D.driver(); + var e = engine(&d); + defer cuda.completion.drain(&d, e.stream.*, "the test driver must drain before it leaves"); + const head: pages.Head = .{ .draft_n = 4096, .rows = 64 }; + const l = pages.layout(256, 16, head, 0); + const p = try pages.Pool.init(t.allocator, &e, l, 256, head, 1, 4 * l.runPages() + 64); + defer p.deinit(); + var seq = try p.open(100); + var dr = try p.draft(&seq, 100); + defer dr.deinit(); + const paged = seq.paged.?; + // the head's rows 15 and 16 are in two pages: not a contiguous cache's next row + try t.expect(paged.rowAt(wt.attn_n, .key, 16) != paged.rowAt(wt.attn_n, .key, 15) + pg.rowBytes(.key)); + // the head's pass (a prompt chunk, a window's absorb, a shared round, the EXL3 deep chunk) takes the lane's table + const hs = dr.headSeq(null, 0); + try t.expect(hs.pg.pt != 0); + try t.expectEqual(paged.pg(), hs.pg); + try t.expectEqual(dr.kc.ptr, hs.kc[0].ptr); + try t.expect(p.close(&seq)); +} diff --git a/zig/src/families/flashnext/cuda_lanes.zig b/zig/src/families/flashnext/cuda_lanes.zig index 8196e0299a..1939d049d1 100644 --- a/zig/src/families/flashnext/cuda_lanes.zig +++ b/zig/src/families/flashnext/cuda_lanes.zig @@ -49,9 +49,7 @@ pub fn Namespace(comptime Tri: type) type { /// 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; pub fn attach(p: *anyopaque, streams: u32, count: usize) !void { - const own = of(p); - own.kv.pool = try kvl.Pool.init(own.gpa, own.engine, own.kv.layout.?, own.context, own.kv.head, streams, count); - own.kv.paged_only = kvl.pagedOnly() and streams > 0; + try kvl.attachTo(of(p), streams, count); } pub fn fits(p: *anyopaque, rows: u32, max_tokens: u32) bool { return of(p).kv.fits(rows, max_tokens); @@ -290,7 +288,7 @@ pub fn Namespace(comptime Tri: type) type { // a window's rows past the reply's last token (its drafts) still need cache slots const cap = @min(ids.len + s.max_new + (if (drafting) mtp.max_rows else 0), prompt.max_tokens); const from = self.kv.choose(); - gop.value_ptr.* = .{ .seq = (if (self.ranks) |r| r.create(cap) else if (from == .pool) self.kv.pool.?.open(cap) else self.engine.newSeq(cap)) catch |err| { + gop.value_ptr.* = .{ .seq = (if (self.ranks) |r| r.createOn(cap, from) else if (from == .pool) self.kv.pool.?.open(cap) else self.engine.newSeq(cap)) catch |err| { self.lanes.removeByPtr(gop.key_ptr); return err; }, .from = from }; diff --git a/zig/src/families/flashnext/cuda_lanes_explain.zig b/zig/src/families/flashnext/cuda_lanes_explain.zig index 4459d56c04..e5c5d7448a 100644 --- a/zig/src/families/flashnext/cuda_lanes_explain.zig +++ b/zig/src/families/flashnext/cuda_lanes_explain.zig @@ -6,7 +6,7 @@ pub fn explain(err: anyerror) ?[]const u8 { error.PromptTooLong => "the prompt and its reply exceed this server's context window: shorten it or lower max_tokens", error.OutOfDeviceMemory => "the GPU had no memory left for this request's caches: retry once another request ends", error.BadFlashNextKvPages, error.BadFlashNextPageTest => "TF_FLASHNEXT_KV_PAGES and _PAGE_TEST take positions a page, a multiple of 4", - error.FlashNextPagesUnsupported => "TF_FLASHNEXT_KV_PAGES needs one rank, our own glue and no --slide: unset it", + error.FlashNextPagesUnsupported => "TF_FLASHNEXT_KV_PAGES needs our own glue and no --slide: unset it", error.UnsupportedCheckpoint => "the checkpoint differs from the supported Flash Next CUDA layout; the log names its config field", error.TensorRankDraftSnapshotUnavailable => "tensor-rank MTP snapshots are unavailable: disable the prompt cache for this configuration", error.RankDisconnected, error.RankTimeout, error.RemoteRank, error.BadRankFrame => "a tensor-parallel rank failed; restart both ranks from the same checkpoint before retrying", diff --git a/zig/src/families/flashnext/cuda_lanes_load.zig b/zig/src/families/flashnext/cuda_lanes_load.zig index 0b38a490bc..6d0a05c19a 100644 --- a/zig/src/families/flashnext/cuda_lanes_load.zig +++ b/zig/src/families/flashnext/cuda_lanes_load.zig @@ -27,7 +27,7 @@ pub fn open(comptime F: type, gpa: std.mem.Allocator, io: std.Io, ctx: *const cu 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 or mode == .exl3)) return error.FlashNextPagesUnsupported; // the pool's pages are not wired to EXL3's K/V + if (page_tokens > 0 and (o.slide or captured != null)) return error.FlashNextPagesUnsupported; const driver: *cuda.Driver = @constCast(ctx.d); const own = try gpa.create(Cuda); errdefer gpa.destroy(own); @@ -103,8 +103,9 @@ pub fn open(comptime F: type, gpa: std.mem.Allocator, io: std.Io, ctx: *const cu 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; + const draft_rows = @import("exl3/prefill/cuda_exl3_prefill.zig").draftRows(own.engine); + own.kv.head = if (own.head) |h| .{ .draft_n = h.ids.len, .rows = draft_rows } else null; + const head_bytes = if (own.head) |h| mtp.Draft.bytes(window, h.ids.len, draft_rows) 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 diff --git a/zig/src/families/flashnext/cuda_mtp_draft.zig b/zig/src/families/flashnext/cuda_mtp_draft.zig index c2f23d0598..e95ffb221b 100644 --- a/zig/src/families/flashnext/cuda_mtp_draft.zig +++ b/zig/src/families/flashnext/cuda_mtp_draft.zig @@ -179,6 +179,11 @@ pub const Draft = struct { return d; } + /// The head's one layer as a pass reads its cache: contiguous, or the lane's pages (`pg`) after the model's. + pub fn headSeq(d: *Draft, media: ?@import("cuda_media.zig").View, scores: u64) batch.SeqState { + return .{ .media = media, .conv = &.{}, .state = &.{}, .kc = (&d.kc)[0..1], .vc = (&d.vc)[0..1], .ikc = (&d.ikc)[0..1], .pooled = (&d.pooled)[0..1], .pg = d.pg, .tail = d.kc, .history = &d.unused_history, .cap = d.cap, .scores = scores }; + } + /// Each DeltaNet layer's window end state into its spare buffer. fn aim(d: *Draft) void { var li: usize = 0; diff --git a/zig/src/families/flashnext/cuda_mtp_ops.zig b/zig/src/families/flashnext/cuda_mtp_ops.zig index eeb1a42107..0809faa1d1 100644 --- a/zig/src/families/flashnext/cuda_mtp_ops.zig +++ b/zig/src/families/flashnext/cuda_mtp_ops.zig @@ -137,7 +137,7 @@ pub fn Ops(comptime Tri: type) type { .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 = d.media, .conv = &.{}, .state = &.{}, .kc = (&d.kc)[0..1], .vc = (&d.vc)[0..1], .ikc = (&d.ikc)[0..1], .pooled = (&d.pooled)[0..1], .pg = d.pg, .tail = d.kc, .history = &d.unused_history, .cap = d.cap, .scores = score_rows }, + .seq = d.headSeq(d.media, score_rows), .hooks = .{ .logits = .none }, }; var at: usize = 0; @@ -279,7 +279,7 @@ pub fn Ops(comptime Tri: type) type { var at: usize = 0; for (ds, nexts, 0..) |d, nx, i| { total += nx.len; - lane_buf[i] = .{ .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 }, .pos = d.len, .rows = nx.len }; + lane_buf[i] = .{ .seq = d.headSeq(null, score_rows), .pos = d.len, .rows = nx.len }; @memcpy(ids_buf[at..][0..nx.len], nx); at += nx.len; } diff --git a/zig/src/families/flashnext/cuda_prompt.zig b/zig/src/families/flashnext/cuda_prompt.zig index 0d22973c8c..38a6ef7043 100644 --- a/zig/src/families/flashnext/cuda_prompt.zig +++ b/zig/src/families/flashnext/cuda_prompt.zig @@ -476,7 +476,11 @@ pub const Engine = struct { const has_media = for (seqs) |q| { if (q.media != null) break true; } else false; - if (e.collective != null and round_graph.enabled() and !has_media) { + const has_paged = for (seqs) |q| { + if (q.paged != null) break true; + } else false; + // a captured round fixes the lanes' tables: paged streams run the plain pass + if (e.collective != null and round_graph.enabled() and !has_media and !has_paged) { const key = try round_graph.Key.init(seqs, s.ids, scores, out); try e.round_graph.run(key, &p, ids); } else try p.chunk(ids, 0, null); diff --git a/zig/src/families/flashnext/cuda_rank_engine.zig b/zig/src/families/flashnext/cuda_rank_engine.zig index 31e3056d1e..d1f8565cbd 100644 --- a/zig/src/families/flashnext/cuda_rank_engine.zig +++ b/zig/src/families/flashnext/cuda_rank_engine.zig @@ -8,6 +8,7 @@ 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"); +const kvl = @import("cuda_kv_lanes.zig"); pub const Identity = extern struct { config: [32]u8, @@ -21,6 +22,9 @@ pub const Identity = extern struct { head_sharded: u64 = 0, image_tokens: u64 = 0, max_images: u64 = 0, + /// Positions a page of the core KV pool holds on this rank (0: no pool), and whether every stream is on pages. + kv_pages: u64 = 0, + kv_pages_all: u64 = 0, }; fn fileHash(a: std.mem.Allocator, io: std.Io, path: []const u8) ![32]u8 { @@ -42,7 +46,7 @@ pub fn identity(a: std.mem.Allocator, io: std.Io, dir: []const u8, context: usiz defer a.free(index); const program = try std.process.executablePathAlloc(io, a); defer a.free(program); - return .{ .config = try fileHash(a, io, config), .index = try fileHash(a, io, index), .program = try fileHash(a, io, program), .checkpoint = checkpoint, .context = context, .chunk = chunk, .decode_graph = @intFromBool(@import("cuda_round_graph.zig").enabled()), .head_sharded = @intFromBool(try @import("cuda_head.zig").enabled()) }; + return .{ .config = try fileHash(a, io, config), .index = try fileHash(a, io, index), .program = try fileHash(a, io, program), .checkpoint = checkpoint, .context = context, .chunk = chunk, .decode_graph = @intFromBool(@import("cuda_round_graph.zig").enabled()), .head_sharded = @intFromBool(try @import("cuda_head.zig").enabled()), .kv_pages = try kvl.paging(), .kv_pages_all = @intFromBool(kvl.pagedOnly()) }; } pub const Session = struct { @@ -56,6 +60,8 @@ pub const Session = struct { seqs: std.AutoHashMapUnmanaged(u64, *prompt.Seq) = .empty, saved: std.AutoHashMapUnmanaged(u64, *snaps.Snap) = .empty, windows: std.AutoHashMapUnmanaged(u64, mtp.Draft) = .empty, + /// This rank's core KV pool, when it pages (cuda_lanes.zig attach): a create with the pool flag takes a lane. + pool: ?*kvl.Pool = null, drafting: bool = false, stopped: bool = false, available: std.atomic.Value(bool) = .init(true), @@ -110,7 +116,7 @@ pub const Session = struct { s.collective.deinit(); var qi = s.seqs.valueIterator(); while (qi.next()) |q| { - q.*.deinit(s.gpa); + s.release(q.*); s.gpa.destroy(q.*); } s.seqs.deinit(s.gpa); @@ -159,12 +165,31 @@ pub const Session = struct { if (h.flags != 0) return error.RemoteRank; } + /// A stream's caches given back: a pool lane's to the pool, else freed (a window's swapped-in states first). + fn release(s: *Session, q: *prompt.Seq) void { + if (s.pool) |p| { + if (p.owns(q)) { + cuda.completion.drain(s.engine.driver, s.engine.stream.*, "flashnext: a closed rank stream did not drain"); + for (q.state) |*b| if (!b.borrowed) b.free(); // a window swapped its own states in + _ = p.close(q); + return; + } + } + q.deinit(s.gpa); + } + pub fn create(s: *Session, cap: usize) !prompt.Seq { - var q = try s.engine.newSeq(cap); - errdefer q.deinit(s.gpa); + return s.createOn(cap, .none); + } + + /// A stream's caches on both ranks: a lane of each rank's pool (`from` .pool), else each rank's own. + pub fn createOn(s: *Session, cap: usize, from: kvl.From) !prompt.Seq { + const pooled = from == .pool; + var q = if (pooled) try s.pool.?.open(cap) else try s.engine.newSeq(cap); + errdefer s.release(&q); const id = s.next_object; s.next_object = std.math.add(u64, id, 1) catch return error.BadRankFrame; - const h = try s.command(.{ .kind = .create, .object = id, .arg = cap }, &.{}); + const h = try s.command(.{ .kind = .create, .object = id, .arg = cap, .flags = @intFromBool(pooled) }, &.{}); try s.ack(h); q.rank_id = id; return q; @@ -175,12 +200,12 @@ pub const Session = struct { if (!s.stopped) { const h = s.command(.{ .kind = .drop_seq, .object = q.rank_id }, &.{}) catch { s.poison(); - q.deinit(s.gpa); + s.release(q); return; }; s.ack(h) catch s.poison(); } - q.deinit(s.gpa); + s.release(q); } pub const setMedia = @import("cuda_rank_media.zig").send; @@ -307,8 +332,8 @@ pub const Session = struct { if (s.seqs.contains(h.object)) return error.BadRankFrame; const q = try s.gpa.create(prompt.Seq); errdefer s.gpa.destroy(q); - q.* = try s.engine.newSeq(@intCast(h.arg)); - errdefer q.deinit(s.gpa); + q.* = if (h.flags == 1) try (s.pool orelse return error.BadRankFrame).open(@intCast(h.arg)) else try s.engine.newSeq(@intCast(h.arg)); + errdefer s.release(q); q.rank_id = h.object; try s.seqs.put(s.gpa, h.object, q); }, @@ -363,7 +388,7 @@ pub const Session = struct { d.deinit(); } const entry = s.seqs.fetchRemove(h.object) orelse return error.UnknownRankSequence; - entry.value.deinit(s.gpa); + s.release(entry.value); s.gpa.destroy(entry.value); }, .save => { @@ -420,7 +445,7 @@ fn sameIdentity(a: Identity, b: Identity) bool { /// Refuse malformed operations before looking up an owner or touching the GPU. 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, + .create => if (h.object == 0 or h.arg == 0 or h.arg > prompt.max_tokens or h.bytes != 0 or h.flags > 1) 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 => { _ = try telemetry.phaseFromArg(h.arg); @@ -459,10 +484,18 @@ test "rank identity binds checkpoint, config, index, program and execution setti other = base; other.head_sharded = 1; try std.testing.expect(!sameIdentity(base, other)); + other = base; + other.kv_pages = 256; + try std.testing.expect(!sameIdentity(base, other)); + other = base; + other.kv_pages_all = 1; + try std.testing.expect(!sameIdentity(base, other)); } test "rank operations refuse invalid geometry and unsupported hooks before execution" { try shape(.{ .kind = .create, .object = 1, .arg = 4096 }); + try shape(.{ .kind = .create, .object = 1, .arg = 4096, .flags = 1 }); // a lane of the core KV pool + try std.testing.expectError(error.BadRankFrame, shape(.{ .kind = .create, .object = 1, .arg = 4096, .flags = 2 })); try shape(.{ .kind = .forward, .object = 1, .bytes = 4 }); try shape(.{ .kind = .round, .bytes = 24 }); try shape(.{ .kind = .verify, .object = 1, .bytes = mtp.max_rows * 4 }); 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 index 128bac0788..f754048c1c 100644 --- a/zig/src/families/flashnext/exl3/prefill/cuda_exl3_prefill_head.zig +++ b/zig/src/families/flashnext/exl3/prefill/cuda_exl3_prefill_head.zig @@ -83,7 +83,7 @@ fn headSlice(comptime Tri: type, o: anytype, d: *@import("../../cuda_mtp.zig").D .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 }, + .seq = d.headSeq(null, score_rows), .hooks = .{ .logits = .none }, }; var views: [max_sub]batch.Scratch = undefined; From 5ed763bcfa287e7ec9bf1800d862aa9c56c775f1 Mon Sep 17 00:00:00 2001 From: BlivionIaG <15323295+BlivionIaG@users.noreply.github.com> Date: Tue, 27 Oct 2026 19:48:13 +0100 Subject: [PATCH 5/7] KV pool: host tier A Tier beside the Pool for families whose pages are cached on the device from a host store. The pool stays the only owner of page ids and of what is promised; the tier adds a second currency (device frames a stream may pin at once), the page-to-frame mapping, an LRU order, pins for the frames a call reads, a dirty flag so a frame is dropped only once its rows reached the host store, and a generation per frame so a table written before a frame was recycled is told from a current one. A stream can also be lent frames for its life, mapped to no page, for a family that manages them itself. The pool tells the tier when a page loses its last holder. It moves no bytes and uses no device API: the family copies what a miss names, and a snapshot reads canonical rows from the host store. Tests: reservation conservation in both currencies, eviction order, pin rules, lent frames, generation safety, canonical rows, a randomised serving loop that audits the mapping and checks every row read against the last write, two ranks in lockstep (and a step or a lent count taken by one rank caught by the digest), and mutants (no generation, no pin rule, dirty frames unpinned) that the checks fail on. A failed batch unmaps the misses it mapped (their generation moves on), a page is marked dirty only under a pin of the marking lease, a stream cannot end while it holds the last pin of a dirty page, and a family can lend a run of consecutive frames (lendRun). --- zig/src/core/kv_pool.zig | 15 + zig/src/core/kv_tier.zig | 432 ++++++++++++++++++++++++++ zig/src/core/kv_tier_test.zig | 564 ++++++++++++++++++++++++++++++++++ 3 files changed, 1011 insertions(+) create mode 100644 zig/src/core/kv_tier.zig create mode 100644 zig/src/core/kv_tier_test.zig diff --git a/zig/src/core/kv_pool.zig b/zig/src/core/kv_pool.zig index 54b0322fa0..7d47fde91f 100644 --- a/zig/src/core/kv_pool.zig +++ b/zig/src/core/kv_pool.zig @@ -142,6 +142,12 @@ pub const Notify = struct { /// A family's pool as the load ledger reads it: tokens a page, and the contiguous run it needs (0: none). pub const Source = struct { pool: *Pool, page_tokens: u32, run_pages: u32 = 0 }; +/// A tier's notice that a page lost its last holder (its frame, if any, is free again). +pub const Freed = struct { + ctx: *anyopaque, + call: *const fn (ctx: *anyopaque, id: u32) void, +}; + /// What an idle pool must show: nothing promised, every page free or the store's alone. pub const Census = struct { total: usize, free: usize, promised: usize, kept: usize, leased: usize }; @@ -160,6 +166,7 @@ pub const Pool = struct { reclaim: ?Reclaim = null, notify: ?Notify = null, depth: u8 = 0, // public changes in progress: the outermost one notifies, after nested calls and rollbacks + freed: ?Freed = null, pub fn init(gpa: Allocator, count: usize, page_bytes: u64) !Pool { const refs = try gpa.alloc(u32, count); @@ -220,6 +227,7 @@ pub const Pool = struct { if (p.leased[id] == 0 or p.refs[id] < p.leased[id]) std.debug.panic("kv pool: page {d} released with no holder", .{id}); p.ownership(id, p.refs[id] - 1, p.leased[id] - 1); if (p.refs[id] > 0) return; + if (p.freed) |f| f.call(f.ctx, id); p.free[p.nfree] = id; p.nfree += 1; } @@ -256,6 +264,7 @@ pub const Pool = struct { if (p.refs[id] <= p.leased[id]) return error.NotHeld; // no holder past the leases' (callers drop theirs first) p.ownership(id, p.refs[id] - 1, p.leased[id]); if (p.refs[id] > 0) return; + if (p.freed) |f| f.call(f.ctx, id); p.free[p.nfree] = id; p.nfree += 1; } @@ -424,6 +433,7 @@ pub const Pool = struct { for (start..start + k) |id| { if (p.refs[id] != 1 or p.leased[id] != 1) std.debug.panic("kv pool: run page {d} released while not a run's", .{id}); p.ownership(id, 0, 0); + if (p.freed) |f| f.call(f.ctx, @intCast(id)); p.free[p.nfree] = @intCast(id); p.nfree += 1; } @@ -449,8 +459,13 @@ pub const Pool = struct { } }; +/// The host tier: logical pages in host RAM and device frames that cache them (kv_tier.zig). +pub const tier = @import("kv_tier.zig"); +pub const Tier = tier.Tier; + test { _ = @import("kv_pool_test.zig"); + _ = @import("kv_tier.zig"); _ = @import("kv_pool_counts_test.zig"); _ = @import("kv_pool_walk_test.zig"); _ = @import("kv_pool_notify_test.zig"); diff --git a/zig/src/core/kv_tier.zig b/zig/src/core/kv_tier.zig new file mode 100644 index 0000000000..3d40f1d98c --- /dev/null +++ b/zig/src/core/kv_tier.zig @@ -0,0 +1,432 @@ +//! A host tier for the page pool: logical pages whose canonical rows live in host RAM, cached by device frames. +const std = @import("std"); +const kv = @import("kv_pool.zig"); +const Allocator = std.mem.Allocator; + +/// No frame (for a page) or no page (for a frame). +pub const none: u32 = std.math.maxInt(u32); + +/// What the ranks of one model agree on before they open a tier: the smallest of each. +pub const Shape = struct { + host_pages: usize, + frames: usize, + + pub fn agree(a: Shape, b: Shape) Shape { + return .{ .host_pages = @min(a.host_pages, b.host_pages), .frames = @min(a.frames, b.frames) }; + } +}; + +/// A deliberate bug a test turns on to show that the check beside it fails without the rule: never set outside tests. +pub const Fault = enum { none, keep_generation, ignore_pins, ignore_dirty, unpinned_dirty, keep_failed_miss }; + +/// A page a call pinned: its frame, the frame's generation, and whether the family must copy the page in. +pub const Fill = struct { page: u32, frame: u32, gen: u32, miss: bool }; + +/// A page's frame as a device table saw it: current while the frame's generation is the same. +pub const Ticket = struct { frame: u32, gen: u32 }; + +/// A stream's reservation in both currencies: pages (the pool's lease) and the frames it may pin at once. +pub const Lease = struct { + pages: kv.Lease = .{}, + frames: usize = 0, + pins: std.ArrayList(u32) = .empty, // one entry per pin, a page id + lent: std.ArrayList(u32) = .empty, // frames lent to the stream: mapped to no page, the family's to manage +}; + +pub const Census = struct { frames: usize, free: usize, resident: usize, pinned: usize, dirty: usize, promised: usize, lent: usize }; + +pub const Tier = struct { + gpa: Allocator, + pool: kv.Pool, // logical pages; the host store holds `pool.total()` of them + host: []u8, // canonical rows, `pool.total() * page_bytes`, borrowed (pinned host memory in a server) + frame_of: []u32, // per page: its frame or `none` + dirty: []bool, // per page: the frame holds rows the host store lacks + page_of: []u32, // per frame: its page or `none` + gen: []u32, // per frame: changes whenever the frame stops holding the page it held + stamp: []u64, // per frame: the LRU clock at its last use + pinned: []u32, // per frame: calls reading it + free_frames: []u32, + lent_mark: []bool, // per frame: lent to a lease + nfree: usize, + promised: usize = 0, // frames promised, summed over leases + nlent: usize = 0, // frames lent, summed over leases + clock: u64 = 0, + fault: Fault = .none, + + /// In place (the pool's release hook points here): `host.len / page_bytes` logical pages, `frames` frames. + pub fn init(t: *Tier, gpa: Allocator, host: []u8, frames: usize, page_bytes: u64) !void { + const npages: usize = @intCast(host.len / page_bytes); + var pool = try kv.Pool.init(gpa, npages, page_bytes); + errdefer pool.deinit(); + const frame_of = try gpa.alloc(u32, npages); + errdefer gpa.free(frame_of); + const dirty = try gpa.alloc(bool, npages); + errdefer gpa.free(dirty); + const page_of = try gpa.alloc(u32, frames); + errdefer gpa.free(page_of); + const gen = try gpa.alloc(u32, frames); + errdefer gpa.free(gen); + const stamp = try gpa.alloc(u64, frames); + errdefer gpa.free(stamp); + const pinned = try gpa.alloc(u32, frames); + errdefer gpa.free(pinned); + const free_frames = try gpa.alloc(u32, frames); + errdefer gpa.free(free_frames); + const lent_mark = try gpa.alloc(bool, frames); + @memset(frame_of, none); + @memset(dirty, false); + @memset(page_of, none); + @memset(gen, 0); + @memset(stamp, 0); + @memset(pinned, 0); + @memset(lent_mark, false); + for (free_frames, 0..) |*f, i| f.* = @intCast(frames - 1 - i); // frame 0 first + t.* = .{ .gpa = gpa, .pool = pool, .host = host, .frame_of = frame_of, .dirty = dirty, .page_of = page_of, .gen = gen, .stamp = stamp, .pinned = pinned, .free_frames = free_frames, .lent_mark = lent_mark, .nfree = frames }; + t.pool.freed = .{ .ctx = t, .call = freed }; + } + + pub fn deinit(t: *Tier) void { + t.pool.deinit(); + t.gpa.free(t.frame_of); + t.gpa.free(t.dirty); + t.gpa.free(t.page_of); + t.gpa.free(t.gen); + t.gpa.free(t.stamp); + t.gpa.free(t.pinned); + t.gpa.free(t.free_frames); + t.gpa.free(t.lent_mark); + t.* = undefined; + } + + pub fn count(t: *const Tier) usize { + return t.page_of.len; + } + + pub fn shape(t: *const Tier) Shape { + return .{ .host_pages = t.pool.total(), .frames = t.count() }; + } + + /// Promise `pages` pages and `n` frames to `l`, or neither: LaneMemory or FrameMemory, whichever does not fit. + pub fn reserve(t: *Tier, l: *Lease, pages: usize, n: usize) error{ LaneMemory, FrameMemory }!void { + if (t.promised + n > t.count()) return error.FrameMemory; + try t.pool.promise(&l.pages, pages); + t.promised += n; + l.frames += n; + } + + /// Whether `pages` and `n` frames could be promised now (the store's idle pages counted). + pub fn fits(t: *const Tier, pages: usize, n: usize) bool { + return t.promised + n <= t.count() and t.pool.fits(pages); + } + + /// `l` covers `n` pages, as `Pool.cover`. + pub fn cover(t: *Tier, l: *Lease, n: usize, table: ?kv.Table) !void { + return t.pool.cover(&l.pages, n, table); + } + + /// The frame of `id`, if it has one. + pub fn resident(t: *const Tier, id: u32) ?Ticket { + const f = t.frame_of[id]; + return if (f == none) null else .{ .frame = f, .gen = t.gen[f] }; + } + + /// Whether a table entry written from `tk` still names `id`'s rows. + pub fn valid(t: *const Tier, id: u32, tk: Ticket) bool { + return t.page_of[tk.frame] == id and t.gen[tk.frame] == tk.gen; + } + + /// Pin `id` for a call: a frame if it has none (a miss: copy `host_page(id)` in); FrameMemory past the promise. + pub fn pin(t: *Tier, l: *Lease, id: u32) error{ FrameMemory, OutOfMemory }!Fill { + if (t.pool.refs[id] == 0) std.debug.panic("kv tier: page {d} pinned while free", .{id}); + if (l.pins.items.len + l.lent.items.len >= l.frames) return error.FrameMemory; + try l.pins.ensureUnusedCapacity(t.gpa, 1); + var miss = false; + if (t.frame_of[id] == none) { + const f = t.takeFrame() orelse return error.FrameMemory; + t.page_of[f] = id; + t.frame_of[id] = f; + miss = true; + } + const f = t.frame_of[id]; + t.pinned[f] += 1; + t.touch(id); + l.pins.appendAssumeCapacity(id); + return .{ .page = id, .frame = f, .gen = t.gen[f], .miss = miss }; + } + + /// Pin every page of `ids` into `out`, or none of them. + pub fn fill(t: *Tier, l: *Lease, ids: []const u32, out: []Fill) !void { + std.debug.assert(out.len >= ids.len); + var done: usize = 0; + errdefer { + for (ids[0..done]) |id| t.unpin(l, id) catch {}; + for (out[0..done]) |o| t.undo(o); // a miss was never copied in: its frame must not stay mapped + } + for (ids, 0..) |id, i| { + out[i] = try t.pin(l, id); + done += 1; + } + } + + /// The frame a failed batch took for a miss goes back unmapped (generation moved on), unless something pins it. + fn undo(t: *Tier, o: Fill) void { + if (t.fault == .keep_failed_miss or !o.miss or t.frame_of[o.page] != o.frame or t.pinned[o.frame] != 0) return; + t.unmap(o.frame); + t.free_frames[t.nfree] = o.frame; + t.nfree += 1; + } + + /// One pin of `id` fewer. Dirty: the last pin of a page whose rows the host store lacks cannot go (sync it first). + pub fn unpin(t: *Tier, l: *Lease, id: u32) error{ NotPinned, Dirty }!void { + var i = l.pins.items.len; + while (i > 0) : (i -= 1) if (l.pins.items[i - 1] == id) break; + if (i == 0) return error.NotPinned; + const f = t.frame_of[id]; + if (t.pinned[f] == 1 and t.dirty[id] and t.fault != .ignore_dirty) return error.Dirty; + _ = l.pins.orderedRemove(i - 1); + t.pinned[f] -= 1; + } + + /// Lend `n` frames to `l` for its life: mapped to no page, never a victim, counted in its promise like pins. + pub fn lend(t: *Tier, l: *Lease, n: usize) error{ FrameMemory, OutOfMemory }!void { + if (l.pins.items.len + l.lent.items.len + n > l.frames) return error.FrameMemory; + try l.lent.ensureUnusedCapacity(t.gpa, n); + var taken: usize = 0; + errdefer t.giveBackTo(l, l.lent.items.len - taken); + while (taken < n) : (taken += 1) { + const f = t.takeFrame() orelse return error.FrameMemory; + t.gen[f] +%= 1; + t.lent_mark[f] = true; + l.lent.appendAssumeCapacity(f); + t.nlent += 1; + } + } + + /// Lend `l` `n` consecutive free frames (a family's state run) for its life; the first frame, or FrameMemory. + pub fn lendRun(t: *Tier, l: *Lease, n: usize) error{ FrameMemory, OutOfMemory }!u32 { + if (n == 0 or l.pins.items.len + l.lent.items.len + n > l.frames) return error.FrameMemory; + var start: usize = 0; + var run: usize = 0; + for (t.page_of, t.lent_mark, 0..) |p, m, f| { + if (p == none and !m) run += 1 else run = 0; + if (run == n) { + start = f + 1 - n; + break; + } + } else return error.FrameMemory; + try l.lent.ensureUnusedCapacity(t.gpa, n); + var kept: usize = 0; + for (t.free_frames[0..t.nfree]) |f| if (f < start or f >= start + n) { + t.free_frames[kept] = f; + kept += 1; + }; + t.nfree = kept; + for (start..start + n) |f| { + t.gen[f] +%= 1; + t.lent_mark[f] = true; + l.lent.appendAssumeCapacity(@intCast(f)); + } + t.nlent += n; + return @intCast(start); + } + + /// The frames lent to `l` go back to the free frames, their generation moved on. + pub fn giveBack(t: *Tier, l: *Lease) void { + t.giveBackTo(l, 0); + } + + fn giveBackTo(t: *Tier, l: *Lease, keep: usize) void { + while (l.lent.items.len > keep) { + const f = l.lent.pop().?; + t.gen[f] +%= 1; + t.lent_mark[f] = false; + t.free_frames[t.nfree] = f; + t.nfree += 1; + t.nlent -= 1; + } + } + + /// The LRU clock moves on for `id` (a hit the device saw, or a call that read it). + pub fn touch(t: *Tier, id: u32) void { + const f = t.frame_of[id]; + if (f == none) return; + t.clock += 1; + t.stamp[f] = t.clock; + } + + /// Rows were written into `id`'s frame, which `l` pins: the host store lacks them until `sync`. + pub fn markDirty(t: *Tier, l: *const Lease, id: u32) error{NotPinned}!void { + if (t.fault != .unpinned_dirty and std.mem.indexOfScalar(u32, l.pins.items, id) == null) return error.NotPinned; + if (t.frame_of[id] == none) std.debug.panic("kv tier: page {d} written with no frame", .{id}); + t.dirty[id] = true; + } + + /// The family copied `id`'s frame to `host_page(id)`: the frame may be dropped. + pub fn sync(t: *Tier, id: u32) void { + t.dirty[id] = false; + } + + /// Where a page's canonical rows go (a flush writes here). + pub fn host_page(t: *Tier, id: u32) []u8 { + const b: usize = @intCast(t.pool.page_bytes); + return t.host[id * b ..][0..b]; + } + + /// The canonical rows of `id` for a snapshot or a fill; Unsynced while its frame holds rows not flushed yet. + pub fn canonical(t: *const Tier, id: u32) error{Unsynced}![]const u8 { + if (t.dirty[id]) return error.Unsynced; + const b: usize = @intCast(t.pool.page_bytes); + return t.host[id * b ..][0..b]; + } + + /// Drop up to `n` unpinned frames, least recently used first; the count dropped. + pub fn evict(t: *Tier, n: usize) usize { + var done: usize = 0; + while (done < n) : (done += 1) { + const f = t.victim() orelse break; + t.unmap(f); + t.free_frames[t.nfree] = f; + t.nfree += 1; + } + return done; + } + + /// Whether `l` can end: none of its pinned pages is left dirty and unpinned (a page only it holds goes with it). + pub fn canFinish(t: *const Tier, l: *const Lease) bool { + if (t.fault == .ignore_dirty) return true; + for (l.pins.items) |id| { + if (!t.dirty[id]) continue; + const f = t.frame_of[id]; + var mine: usize = 0; // this lease's own pins of the page + for (l.pins.items) |q| mine += @intFromBool(q == id); + if (f != none and t.pinned[f] > mine) continue; // another stream still pins it + const own = std.mem.indexOfScalar(u32, l.pages.ids.items, id) != null; + if (!own or t.pool.refs[id] > 1) return false; + } + for (l.pages.ids.items) |id| if (t.dirty[id] and t.pool.refs[id] > 1) return false; + return true; + } + + /// The stream ended: its pins, frames, promise and pages go back (a page nobody else holds loses its frame). + pub fn finish(t: *Tier, l: *Lease) void { + if (t.fault == .none and !t.canFinish(l)) std.debug.panic("kv tier: a stream ended with unflushed rows pinned or shared", .{}); + for (l.pins.items) |id| { + const f = t.frame_of[id]; + if (f != none) t.pinned[f] -= 1 else if (t.fault == .none) std.debug.panic("kv tier: pinned page {d} lost its frame", .{id}); + } + l.pins.deinit(t.gpa); + t.giveBack(l); + l.lent.deinit(t.gpa); + t.promised -= l.frames; + t.pool.finish(&l.pages); + l.* = .{}; + } + + pub fn census(t: *const Tier) Census { + var c: Census = .{ .frames = t.count(), .free = t.nfree, .resident = 0, .pinned = 0, .dirty = 0, .promised = t.promised, .lent = t.nlent }; + for (t.page_of, t.pinned) |p, n| { + c.resident += @intFromBool(p != none); + c.pinned += @intFromBool(n > 0); + } + for (t.dirty) |d| c.dirty += @intFromBool(d); + return c; + } + + /// Nothing promised, pinned or dirty, and every frame free or holding a page still held (the store's). + pub fn idle(t: *const Tier) bool { + const c = t.census(); + return t.pool.idle() and c.promised == 0 and c.pinned == 0 and c.dirty == 0 and c.lent == 0 and c.free + c.resident == c.frames; + } + + /// The mapping's own consistency; an error names the broken rule. Tests call it after every step. + pub fn audit(t: *const Tier) error{Broken}!void { + var mapped: usize = 0; + for (t.page_of, 0..) |p, f| { + if (p == none) { + if (t.pinned[f] != 0) return error.Broken; // a pin on an empty frame + continue; + } + mapped += 1; + if (t.frame_of[p] != f) return error.Broken; // page and frame disagree + if (t.pool.refs[p] == 0) return error.Broken; // a free page keeps a frame + } + if (mapped + t.nfree + t.nlent != t.count()) return error.Broken; // frames neither free, mapped nor lent, or two + for (t.free_frames[0..t.nfree]) |f| if (t.page_of[f] != none or t.lent_mark[f]) return error.Broken; + var marked: usize = 0; + for (t.lent_mark, t.page_of) |m, p| { + marked += @intFromBool(m); + if (m and p != none) return error.Broken; // a lent frame holds a page + } + if (marked != t.nlent) return error.Broken; + for (t.frame_of, 0..) |f, p| { + if (f != none and t.page_of[f] != p) return error.Broken; + if (t.dirty[p] and f == none) return error.Broken; // unflushed rows with no frame + } + if (t.promised > t.count()) return error.Broken; + } + + /// A hash of what two ranks must hold alike: mapping, generations, LRU order, pins, holders and promises. + pub fn digest(t: *const Tier) u64 { + var h = std.hash.Wyhash.init(0); + h.update(std.mem.sliceAsBytes(t.frame_of)); + h.update(std.mem.sliceAsBytes(t.page_of)); + h.update(std.mem.sliceAsBytes(t.gen)); + h.update(std.mem.sliceAsBytes(t.stamp)); + h.update(std.mem.sliceAsBytes(t.pinned)); + h.update(std.mem.sliceAsBytes(t.free_frames[0..t.nfree])); + h.update(std.mem.sliceAsBytes(t.pool.refs)); + h.update(std.mem.sliceAsBytes(t.pool.free[0..t.pool.nfree])); + const counts = [_]u64{ t.promised, t.pool.promised, t.clock, t.nlent }; + h.update(std.mem.sliceAsBytes(&counts)); + return h.final(); + } + + /// A free frame, or the least recently used unpinned one dropped. + fn takeFrame(t: *Tier) ?u32 { + if (t.nfree > 0) { + t.nfree -= 1; + return t.free_frames[t.nfree]; + } + const f = t.victim() orelse return null; + t.unmap(f); + return f; + } + + /// The mapped frame with the oldest stamp, not pinned and with no unflushed rows (the lowest index among equals). + fn victim(t: *const Tier) ?u32 { + var best: ?u32 = null; + for (t.page_of, 0..) |p, f| { + if (p == none) continue; + if (t.pinned[f] > 0 and t.fault != .ignore_pins) continue; + if (t.dirty[p] and t.fault != .ignore_dirty) continue; + if (best == null or t.stamp[f] < t.stamp[best.?]) best = @intCast(f); + } + return best; + } + + /// The frame stops holding its page: the generation moves on so a table naming it is stale. + fn unmap(t: *Tier, f: u32) void { + const p = t.page_of[f]; + if (t.dirty[p] and t.fault != .ignore_dirty) std.debug.panic("kv tier: frame {d} dropped with page {d} unflushed", .{ f, p }); + t.frame_of[p] = none; + t.dirty[p] = false; + t.page_of[f] = none; + if (t.fault != .keep_generation) t.gen[f] +%= 1; + } + + /// The pool freed `id` (its last holder let go): its frame is free again. + fn freed(ctx: *anyopaque, id: u32) void { + const t: *Tier = @ptrCast(@alignCast(ctx)); + const f = t.frame_of[id]; + if (f == none) return; + if (t.pinned[f] != 0 and t.fault == .none) std.debug.panic("kv tier: page {d} freed while pinned", .{id}); + t.dirty[id] = false; // the page is gone: rows nobody can read need no flush + t.unmap(f); + t.free_frames[t.nfree] = f; + t.nfree += 1; + } +}; + +test { + _ = @import("kv_tier_test.zig"); +} diff --git a/zig/src/core/kv_tier_test.zig b/zig/src/core/kv_tier_test.zig new file mode 100644 index 0000000000..95437b2bea --- /dev/null +++ b/zig/src/core/kv_tier_test.zig @@ -0,0 +1,564 @@ +//! The host tier: two currencies, eviction order, pins, generations, canonical rows, lockstep ranks, mutants. +const std = @import("std"); +const kv = @import("kv_pool.zig"); +const kt = @import("kv_tier.zig"); +const Tier = kt.Tier; +const Lease = kt.Lease; +const expect = std.testing.expect; +const expectEqual = std.testing.expectEqual; +const gpa = std.testing.allocator; + +const page_bytes = 8; + +/// A tier of `pages` logical pages and `frames` frames over a host store and a fake device. +const Rig = struct { + host: []u8, + dev: []u8, + tier: Tier, + + fn open(r: *Rig, pages: usize, frames: usize) !void { + r.host = try gpa.alloc(u8, pages * page_bytes); + errdefer gpa.free(r.host); + @memset(r.host, 0); + r.dev = try gpa.alloc(u8, frames * page_bytes); + errdefer gpa.free(r.dev); + @memset(r.dev, 0); + try r.tier.init(gpa, r.host, frames, page_bytes); + } + + fn close(r: *Rig) void { + r.tier.deinit(); + gpa.free(r.host); + gpa.free(r.dev); + } + + fn frame(r: *Rig, f: u32) []u8 { + return r.dev[f * page_bytes ..][0..page_bytes]; + } + + /// A stream with `pages` pages covered and `frames` frames promised. + fn lease(r: *Rig, pages: usize, frames: usize) !Lease { + return r.leaseLent(pages, frames, 0); + } + + /// A stream that also has `lent` of its frames lent for its life. + fn leaseLent(r: *Rig, pages: usize, frames: usize, lent: usize) !Lease { + var l: Lease = .{}; + errdefer r.tier.finish(&l); + try r.tier.reserve(&l, pages, frames); + try r.tier.cover(&l, pages, null); + try r.tier.lend(&l, lent); + return l; + } + + /// Pin `id`, copying its canonical rows in on a miss. + fn pin(r: *Rig, l: *Lease, id: u32) !kt.Fill { + const f = try r.tier.pin(l, id); + if (f.miss) @memcpy(r.frame(f.frame), r.tier.host_page(id)); + return f; + } + + /// Write `b` into every row byte of the pinned page and flush it to the host store. + fn write(r: *Rig, l: *const Lease, id: u32, b: u8) void { + const tk = r.tier.resident(id).?; + @memset(r.frame(tk.frame), b); + r.tier.markDirty(l, id) catch unreachable; + } + + fn flush(r: *Rig, id: u32) void { + const tk = r.tier.resident(id).?; + @memcpy(r.tier.host_page(id), r.frame(tk.frame)); + r.tier.sync(id); + } +}; + +test "two currencies: a promise that fits one and not the other takes nothing, and finishing returns both" { + var r: Rig = undefined; + try r.open(8, 3); + defer r.close(); + const t = &r.tier; + var a = try r.lease(4, 2); + try std.testing.expectError(error.FrameMemory, t.reserve(&a, 1, 2)); // frames: 2 + 2 > 3 + try expectEqual(@as(usize, 0), a.pages.promised); + try expectEqual(@as(usize, 2), t.promised); + var b: Lease = .{}; + try std.testing.expectError(error.LaneMemory, t.reserve(&b, 5, 1)); // pages: 4 + 5 > 8 + try expectEqual(@as(usize, 2), t.promised); // the frame promise was not kept + try expectEqual(@as(usize, 0), t.pool.promised); + try expect(!t.fits(5, 1) and t.fits(4, 1) and !t.fits(4, 2)); + var c = try r.lease(4, 1); + try expect(!t.fits(1, 0)); // every page is promised or taken + t.finish(&a); + t.finish(&c); + t.finish(&b); + try expect(t.idle()); + try t.audit(); +} + +test "a lease cannot pin more frames than it promised, and a failed batch pins nothing" { + var r: Rig = undefined; + try r.open(8, 4); + defer r.close(); + const t = &r.tier; + var a = try r.lease(4, 2); + var out: [3]kt.Fill = undefined; + const ids = a.pages.ids.items; + try std.testing.expectError(error.FrameMemory, t.fill(&a, ids[0..3], &out)); + try expectEqual(@as(usize, 0), t.census().pinned); + try expectEqual(@as(usize, 0), a.pins.items.len); + try expect(t.resident(ids[0]) == null and t.resident(ids[1]) == null); // the misses of the failed batch were unmapped + try t.fill(&a, ids[0..2], out[0..2]); + try expect(out[0].miss and out[1].miss); + try expectEqual(@as(usize, 2), t.census().pinned); + try std.testing.expectError(error.FrameMemory, t.pin(&a, ids[2])); + t.finish(&a); + try expect(t.idle()); +} + +test "eviction takes the least recently used unpinned frame, and a touch moves a page to the back" { + var r: Rig = undefined; + try r.open(8, 3); + defer r.close(); + const t = &r.tier; + var a = try r.lease(6, 1); + const ids = a.pages.ids.items; + for (ids[0..3]) |id| { // one frame promised: pin, read, unpin, three times over + _ = try r.pin(&a, id); + try t.unpin(&a, id); + } + try expectEqual(@as(usize, 3), t.census().resident); + t.touch(ids[0]); // order now 1, 2, 0 + const f = try r.pin(&a, ids[3]); // evicts page 1 + try expect(f.miss); + try expect(t.resident(ids[1]) == null and t.resident(ids[0]) != null and t.resident(ids[2]) != null); + try t.unpin(&a, ids[3]); + try expectEqual(@as(usize, 1), t.evict(1)); // page 2 is next + try expect(t.resident(ids[2]) == null and t.resident(ids[0]) != null); + try expectEqual(@as(usize, 2), t.evict(5)); // 0 and 3: all that is left + try expectEqual(@as(usize, 3), t.census().free); + try t.audit(); + t.finish(&a); +} + +test "pinned frames are never evicted, and a dirty page stays until its rows reach the host store" { + var r: Rig = undefined; + try r.open(8, 2); + defer r.close(); + const t = &r.tier; + var a = try r.lease(2, 1); + var b = try r.lease(2, 1); + const pa = a.pages.ids.items[0]; + const pb = b.pages.ids.items; + _ = try r.pin(&a, pa); + r.write(&a, pa, 7); + _ = try r.pin(&b, pb[0]); + try t.unpin(&b, pb[0]); + _ = try r.pin(&b, pb[1]); // the one frame that can go is pb[0]'s + try expect(t.resident(pa) != null and t.resident(pb[0]) == null); + try std.testing.expectError(error.Dirty, t.unpin(&a, pa)); + try expectEqual(@as(usize, 0), t.evict(2)); // pa pinned and dirty, pb[1] pinned + try t.unpin(&b, pb[1]); + try expectEqual(@as(usize, 1), t.evict(2)); // only pb[1] + try expect(t.resident(pa) != null); + try std.testing.expectError(error.Unsynced, t.canonical(pa)); + r.flush(pa); + try t.unpin(&a, pa); + try expectEqual(@as(u8, 7), (try t.canonical(pa))[3]); + try expectEqual(@as(usize, 1), t.evict(2)); + t.finish(&a); + t.finish(&b); + try expect(t.idle()); +} + +/// Whether a table entry written before a frame was dropped and filled again with the same page is taken for current. +fn staleAccepted(fault: kt.Fault) !bool { + var r: Rig = undefined; + try r.open(4, 1); + defer r.close(); + r.tier.fault = fault; + var a = try r.lease(2, 1); + defer r.tier.finish(&a); + const p = a.pages.ids.items; + const f0 = try r.pin(&a, p[0]); + const old: kt.Ticket = .{ .frame = f0.frame, .gen = f0.gen }; + try r.tier.unpin(&a, p[0]); + _ = try r.pin(&a, p[1]); // the only frame is recycled + try r.tier.unpin(&a, p[1]); + _ = try r.pin(&a, p[0]); // and returns to the first page: same frame, new contents + try r.tier.unpin(&a, p[0]); + return r.tier.valid(p[0], old); +} + +test "generations: a recycled frame is not read through a table written before" { + try expect(!try staleAccepted(.none)); + try expect(try staleAccepted(.keep_generation)); // mutation: without the generation the stale table passes +} + +/// Whether a frame a call pins survives another stream's miss. +fn pinnedSurvives(fault: kt.Fault) !bool { + var r: Rig = undefined; + try r.open(4, 2); + defer r.close(); + r.tier.fault = fault; + var a = try r.lease(1, 1); + defer r.tier.finish(&a); + var b = try r.lease(2, 1); + defer r.tier.finish(&b); + const pa = a.pages.ids.items[0]; + _ = try r.pin(&a, pa); // the oldest stamp, and pinned + _ = try r.pin(&b, b.pages.ids.items[0]); + try r.tier.unpin(&b, b.pages.ids.items[0]); + _ = try r.pin(&b, b.pages.ids.items[1]); + return r.tier.resident(pa) != null; +} + +test "pin rule: the frame a call holds survives a miss; mutation: ignoring pins evicts it" { + try expect(try pinnedSurvives(.none)); + try expect(!try pinnedSurvives(.ignore_pins)); +} + +test "freeing a page (its last holder lets go) frees its frame; a page the store still holds keeps it" { + var r: Rig = undefined; + try r.open(4, 2); + defer r.close(); + const t = &r.tier; + var a = try r.lease(2, 2); + const p = try gpa.dupe(u32, a.pages.ids.items); // ids outlive the lease + defer gpa.free(p); + _ = try r.pin(&a, p[0]); + r.write(&a, p[0], 5); + r.flush(p[0]); + _ = try r.pin(&a, p[1]); + t.pool.retain(p[0]); // the prompt store keeps the first page + t.finish(&a); + try expect(t.resident(p[0]) != null and t.resident(p[1]) == null); + try expectEqual(@as(usize, 1), t.census().free); + try expectEqual(@as(u8, 5), (try t.canonical(p[0]))[0]); + try t.pool.drop(p[0]); + try expect(t.resident(p[0]) == null and t.idle()); + try t.audit(); +} + +/// A fake serving loop: streams open, pin, write, flush, unpin, finish; every row read equals the page's last write. +const Sim = struct { + r: Rig, + rng: std.Random.DefaultPrng, + slots: [3]?Lease = .{ null, null, null }, + model: [16]u8 = @splat(0), // the byte last written to a page (0: none) + next: u8 = 1, + misses: usize = 0, + reads: usize = 0, + + fn open(s: *Sim, seed: u64, frames: usize) !void { + s.* = .{ .r = undefined, .rng = .init(seed) }; + try s.r.open(16, frames); + } + + fn close(s: *Sim) void { + for (&s.slots) |*l| if (l.*) |*x| s.r.tier.finish(x); + s.r.close(); + } + + fn check(s: *Sim, id: u32, tk: kt.Ticket) !void { + if (s.model[id] == 0) return; + for (s.r.frame(tk.frame)) |x| try expectEqual(s.model[id], x); + s.reads += 1; + } + + fn step(s: *Sim) !void { + const t = &s.r.tier; + const rnd = s.rng.random(); + const op = rnd.uintLessThan(u32, 100); + const k = rnd.uintLessThan(usize, s.slots.len); + const l: *Lease = if (s.slots[k]) |*x| x else { + if (op < 60) s.slots[k] = s.r.leaseLent(4, 3, 1) catch null; + return; + }; + const ids = l.pages.ids.items; + const id = ids[rnd.uintLessThan(usize, ids.len)]; + switch (op) { + 0...34 => { + const f = s.r.pin(l, id) catch |e| { + if (e == error.FrameMemory and l.pins.items.len + l.lent.items.len >= l.frames) return; + return e; + }; + s.misses += @intFromBool(f.miss); + try s.check(id, .{ .frame = f.frame, .gen = f.gen }); + }, + 35...59 => if (l.pins.items.len > 0) { + const p = l.pins.items[rnd.uintLessThan(usize, l.pins.items.len)]; + s.model[p] = s.next; + s.next = if (s.next == 255) 1 else s.next + 1; + s.r.write(l, p, s.model[p]); + }, + 60...79 => if (l.pins.items.len > 0) { + const p = l.pins.items[rnd.uintLessThan(usize, l.pins.items.len)]; + if (t.dirty[p]) s.r.flush(p); + try t.unpin(l, p); + }, + 80...89 => _ = t.evict(1), + else => { + for (l.pins.items) |p| if (t.dirty[p]) s.r.flush(p); + for (ids) |p| s.model[p] = 0; + t.finish(l); + s.slots[k] = null; + }, + } + try t.audit(); + for (t.page_of, 0..) |p, f| if (p != kt.none and !t.dirty[p]) try s.check(p, .{ .frame = @intCast(f), .gen = t.gen[f] }); + } +}; + +test "a random serving loop with fewer frames than pages: the audit holds, every read equals the last write, and it ends idle" { + inline for (.{ 1, 2, 3 }) |seed| { + var s: Sim = undefined; + try s.open(seed, 6); + defer s.close(); + for (0..4000) |_| try s.step(); + try expect(s.misses > 100 and s.reads > 500); + for (&s.slots) |*l| if (l.*) |*x| { + for (x.pins.items) |p| if (s.r.tier.dirty[p]) s.r.flush(p); + s.r.tier.finish(x); + l.* = null; + }; + _ = s.r.tier.evict(100); + try expect(s.r.tier.idle()); + } +} + +test "two ranks in lockstep: the same openings and pins give the same frames, generations and order, step after step" { + var a: Sim = undefined; + var b: Sim = undefined; + try a.open(9, 6); + defer a.close(); + try b.open(9, 6); + defer b.close(); + for (0..3000) |_| { + try a.step(); + try b.step(); + try expectEqual(a.r.tier.digest(), b.r.tier.digest()); + } +} + +test "two ranks: a step only one rank takes is caught by the digest at once" { + var a: Sim = undefined; + var b: Sim = undefined; + try a.open(9, 6); + defer a.close(); + try b.open(9, 6); + defer b.close(); + var caught: ?usize = null; + for (0..3000) |i| { + try a.step(); + if (i == 1500) _ = b.r.tier.evict(1) else try b.step(); // rank 1 evicts instead of the step rank 0 took + if (a.r.tier.digest() != b.r.tier.digest()) { + caught = i; + break; + } + } + try expect(caught != null); +} + +test "ranks agree on the smaller shape before they open" { + const a: kt.Shape = .{ .host_pages = 100, .frames = 40 }; + const b: kt.Shape = .{ .host_pages = 90, .frames = 48 }; + try expectEqual(kt.Shape{ .host_pages = 90, .frames = 40 }, a.agree(b)); + try expectEqual(a.agree(b), b.agree(a)); +} + +test "lent frames: counted in the frame promise, never victims or mapped, back at finish with a new generation" { + var r: Rig = undefined; + try r.open(8, 4); + defer r.close(); + const t = &r.tier; + var a = try r.leaseLent(4, 3, 2); // two lent, one left to pin + try expectEqual(@as(usize, 2), t.census().lent); + try expectEqual(@as(usize, 2), a.lent.items.len); + try std.testing.expectError(error.FrameMemory, t.lend(&a, 2)); + const gen0 = t.gen[a.lent.items[0]]; + var b = try r.lease(4, 1); + const ids = a.pages.ids.items; + _ = try r.pin(&a, ids[0]); + try std.testing.expectError(error.FrameMemory, t.pin(&a, ids[1])); // the lent frames used the promise + try t.unpin(&a, ids[0]); + for (b.pages.ids.items) |id| { // misses cycle through the two frames left; the lent ones are never taken + _ = try r.pin(&b, id); + try t.unpin(&b, id); + _ = try r.pin(&a, ids[1]); + try t.unpin(&a, ids[1]); + for (a.lent.items) |f| try expect(t.page_of[f] == kt.none); + } + try expect(t.evict(10) <= 2); // only the mapped frames can go, the lent ones are not among them + try t.audit(); + const lent0 = a.lent.items[0]; + t.finish(&a); + try expect(t.gen[lent0] != gen0); + try expectEqual(@as(usize, 0), t.census().lent); + t.finish(&b); + try expect(t.idle()); + try t.audit(); +} + +test "ranks lending a different count are told apart by the digest" { + var a: Rig = undefined; + var b: Rig = undefined; + try a.open(8, 4); + defer a.close(); + try b.open(8, 4); + defer b.close(); + var la = try a.leaseLent(4, 3, 2); + defer a.tier.finish(&la); + var lb = try b.leaseLent(4, 3, 2); + defer b.tier.finish(&lb); + try expectEqual(a.tier.digest(), b.tier.digest()); + try b.tier.lend(&lb, 0); + var lc = try b.leaseLent(0, 1, 1); + defer b.tier.finish(&lc); + try expect(a.tier.digest() != b.tier.digest()); +} + +test "a lent run is consecutive frames, leaves the free frames, and comes back at finish" { + var r: Rig = undefined; + try r.open(8, 6); + defer r.close(); + const t = &r.tier; + var a = try r.lease(2, 4); + _ = try r.pin(&a, a.pages.ids.items[0]); // a mapped frame in the way + const first = try t.lendRun(&a, 3); + try expectEqual(@as(usize, 3), t.census().lent); + for (a.lent.items, 0..) |f, i| try expectEqual(first + i, f); + try t.audit(); + var b = try r.lease(2, 2); + try std.testing.expectError(error.FrameMemory, t.lendRun(&b, 3)); // past b's promise + _ = try r.pin(&b, b.pages.ids.items[0]); + _ = try r.pin(&b, b.pages.ids.items[1]); + try std.testing.expectEqual(@as(usize, 0), t.nfree); // 1 mapped (a) + 3 lent + 2 mapped (b) + t.finish(&a); + t.finish(&b); + try expect(t.idle()); + try t.audit(); + var c = try r.lease(1, 6); + defer t.finish(&c); + try expectEqual(@as(u32, 0), try t.lendRun(&c, 6)); +} + +/// A failed batch, then the same pages pinned again through the rig: the frames must equal the host pages. +fn failedFillReads(fault: kt.Fault) !bool { + var r: Rig = undefined; + try r.open(8, 4); + defer r.close(); + r.tier.fault = fault; + for (r.host, 0..) |*b, i| b.* = @intCast(1 + i / page_bytes); // page k holds k + 1 + @memset(r.dev, 0xee); // the frames hold something else + var a = try r.lease(4, 2); + defer r.tier.finish(&a); + const ids = a.pages.ids.items; + var out: [3]kt.Fill = undefined; + try std.testing.expectError(error.FrameMemory, r.tier.fill(&a, ids[0..3], &out)); + try r.tier.audit(); + for (ids[0..2]) |id| { + const f = try r.pin(&a, id); + if (!std.mem.eql(u8, r.frame(f.frame), r.tier.host_page(id))) return false; + } + return true; +} + +test "a failed batch leaves no miss mapped: the pages pinned again read the host rows; mutation: keeping the mapping reads stale frames" { + try expect(try failedFillReads(.none)); + try expect(!try failedFillReads(.keep_failed_miss)); +} + +test "a failed batch moves the generation of the frames it unmapped" { + var r: Rig = undefined; + try r.open(4, 2); + defer r.close(); + var a = try r.lease(3, 2); + defer r.tier.finish(&a); + const ids = a.pages.ids.items; + var out: [3]kt.Fill = undefined; + const before = [2]u32{ r.tier.gen[0], r.tier.gen[1] }; + try std.testing.expectError(error.FrameMemory, r.tier.fill(&a, ids[0..3], &out)); + try expect(r.tier.gen[0] != before[0] and r.tier.gen[1] != before[1]); + try expectEqual(@as(usize, 2), r.tier.census().free); +} + +test "dirty needs a pin by the lease: an unpinned frame cannot be marked; mutation: a dirty frame nobody pins starves a lease inside its promise" { + var r: Rig = undefined; + try r.open(4, 2); + defer r.close(); + var a = try r.lease(1, 1); + defer r.tier.finish(&a); + var b = try r.lease(1, 1); + defer r.tier.finish(&b); + const pa = a.pages.ids.items[0]; + _ = try r.pin(&a, pa); + try r.tier.unpin(&a, pa); + try std.testing.expectError(error.NotPinned, r.tier.markDirty(&a, pa)); + try std.testing.expectError(error.NotPinned, r.tier.markDirty(&b, pa)); // pinned by nobody of b + try r.tier.audit(); + // the mutant lets the mark through on a frame no lease pins, and the only frame is then stuck + var m: Rig = undefined; + try m.open(4, 2); + defer m.close(); + m.tier.fault = .unpinned_dirty; + var c = try m.lease(2, 1); + defer m.tier.finish(&c); + var d = try m.lease(1, 1); + defer m.tier.finish(&d); + for (c.pages.ids.items) |pc| { // both frames hold c's pages, unpinned + _ = try m.pin(&c, pc); + try m.tier.unpin(&c, pc); + } + for (c.pages.ids.items) |pc| try m.tier.markDirty(&c, pc); + try std.testing.expectError(error.FrameMemory, m.pin(&d, d.pages.ids.items[0])); +} + +/// A lease that pins a dirty page it does not own, after the owner let go of its pin: whether it may end. +fn endsWithForeignDirty(fault: kt.Fault) !bool { + var r: Rig = undefined; + try r.open(4, 2); + defer r.close(); + r.tier.fault = fault; + var a = try r.lease(1, 1); + var b = try r.lease(1, 1); + const pa = a.pages.ids.items[0]; + _ = try r.pin(&a, pa); + _ = try r.pin(&b, pa); + r.write(&a, pa, 9); + try r.tier.unpin(&a, pa); // b still pins it + const ok = r.tier.canFinish(&b); + if (ok) { + r.tier.finish(&b); + r.flush(pa); // too late for a real stream; the test only wants the tier consistent to close + } else { + r.flush(pa); + r.tier.finish(&b); + } + r.tier.finish(&a); + return ok; +} + +test "a stream cannot end while it holds the last pin of a dirty page, its own table or not; mutation: ignoring dirty lets it" { + try expect(!try endsWithForeignDirty(.none)); + try expect(try endsWithForeignDirty(.ignore_dirty)); +} + +test "a stream ends cleanly once the dirty pages it pins are flushed, and one that shares a dirty page of its own cannot" { + var r: Rig = undefined; + try r.open(4, 2); + defer r.close(); + const t = &r.tier; + var a = try r.lease(2, 2); + const p = a.pages.ids.items; + _ = try r.pin(&a, p[0]); + r.write(&a, p[0], 3); + try expect(t.canFinish(&a)); // its own page and its only holder: the page goes with it + t.pool.retain(p[0]); + try expect(!t.canFinish(&a)); // now another holder would read rows the host lacks + r.flush(p[0]); + try expect(t.canFinish(&a)); + const keep = p[0]; + t.finish(&a); + try t.pool.drop(keep); + try expect(t.idle()); +} From 05f8933fd615f658f933a0b9a6a8423dc7fff883 Mon Sep 17 00:00:00 2001 From: BlivionIaG <15323295+BlivionIaG@users.noreply.github.com> Date: Wed, 28 Oct 2026 10:18:22 +0100 Subject: [PATCH 6/7] EXL3 Flash Next: host-resident K/V on the tier A Flash Next EXL3 stream's K/V pages can live in host RAM on the KV pool's host tier (kv_tier.zig), cached by device frames. Off by default; TF_FLASHNEXT_KV_HOST= turns it on, TF_FLASHNEXT_KV_FRAMES and TF_FLASHNEXT_KV_SLOTS size the frames and the per-layer fetch slots. A page is two planes: K/V rows (in frames) and the indexer and pooled rows (by logical page). With the tier off the layout is the one-plane page and every kernel is byte-identical (SASS compared); the tier launches use separate _planes kernel instances. The first pages and a window of recent pages stay pinned in frames. Older K/V rows are canonical in the pinned mapped host store; the indexer's top-512 blocks are fetched into slots in frames lent to the stream (kvt_plan: device-side hit/miss/LRU, kvt_fetch: the GPU reads host memory, kvt_remap: positions to table rows), so the sparse attention kernels read through the page table unchanged. A page leaves the window only when every attention layer has written it. No host sync per token or layer. Tests: exl3-kvtier runs the three kernels against a host LRU model at 128, 16 and 9 slot frames; the pages contract cases cover the planes kernels; host tests cover the two-plane layout and the lane lifetimes. --- docs/kernel-contract.md | 2 +- zig/build/cuda_kernels.zig | 1 + zig/kernels/cuda/flashnext_attn.cu | 17 +- .../flashnext_exl3/flashnext_exl3_kvtier.cu | 189 +++++++++++ zig/kernels/cuda/flashnext_pages.cuh | 22 ++ zig/kernels/cuda/flashnext_planes.cuh | 47 +++ zig/kernels/cuda/flashnext_qsa_tile.cuh | 6 +- .../core/contract/flashnext_exl3_pending.zig | 8 +- .../core/contract/flashnext_pages_pending.zig | 7 +- zig/src/core/kv_tier.zig | 31 ++ zig/src/core/kv_tier_ranks_test.zig | 49 +++ zig/src/core/kv_tier_test.zig | 20 +- zig/src/cuda/kernels.zig | 1 + zig/src/families/flashnext/cuda_batch.zig | 8 +- .../families/flashnext/cuda_batch_state.zig | 4 + zig/src/families/flashnext/cuda_glue.zig | 51 ++- zig/src/families/flashnext/cuda_images.zig | 5 +- zig/src/families/flashnext/cuda_kv_lanes.zig | 185 +++++++++-- .../flashnext/cuda_kv_lifetime_test.zig | 8 +- .../families/flashnext/cuda_lanes_load.zig | 1 + zig/src/families/flashnext/cuda_media.zig | 15 +- zig/src/families/flashnext/cuda_mtp_draft.zig | 3 +- zig/src/families/flashnext/cuda_pages.zig | 95 +++++- .../flashnext/cuda_paging_batch_test.zig | 1 + .../flashnext/cuda_paging_launch_test.zig | 6 +- zig/src/families/flashnext/cuda_prompt.zig | 4 +- zig/src/families/flashnext/cuda_snap.zig | 4 +- .../flashnext/exl3/kv/cuda_exl3_kvhost.zig | 91 +++++ .../flashnext/exl3/kv/cuda_exl3_kvtier.zig | 314 ++++++++++++++++++ .../flashnext/exl3/model/cuda_exl3.zig | 3 + zig/tests/cuda/contract_fn_pages.zig | 22 +- zig/tests/cuda/contract_fn_pages_media.zig | 27 +- zig/tests/cuda/contract_fn_pages_sums.zig | 88 +++-- zig/tests/cuda/contract_limits.zig | 2 +- .../flashnext_exl3/kv/exl3_kvtier_tests.zig | 221 ++++++++++++ zig/tests/cuda/main.zig | 2 + 36 files changed, 1432 insertions(+), 128 deletions(-) create mode 100644 zig/kernels/cuda/flashnext_exl3/flashnext_exl3_kvtier.cu create mode 100644 zig/kernels/cuda/flashnext_planes.cuh create mode 100644 zig/src/core/kv_tier_ranks_test.zig create mode 100644 zig/src/families/flashnext/exl3/kv/cuda_exl3_kvhost.zig create mode 100644 zig/src/families/flashnext/exl3/kv/cuda_exl3_kvtier.zig create mode 100644 zig/tests/cuda/flashnext_exl3/kv/exl3_kvtier_tests.zig diff --git a/docs/kernel-contract.md b/docs/kernel-contract.md index 4c1dae23ad..64e6e8d4cc 100644 --- a/docs/kernel-contract.md +++ b/docs/kernel-contract.md @@ -256,7 +256,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 13 EXL3 launches (`flashnext_exl3_pending.zig`) carry only an nvcc 13.4 receipt: they stay pending under +Flash Next's 17 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_kernels.zig b/zig/build/cuda_kernels.zig index e5cbdcb67c..05454f9f7b 100644 --- a/zig/build/cuda_kernels.zig +++ b/zig/build/cuda_kernels.zig @@ -124,6 +124,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_kvtier", .src = "flashnext_exl3/flashnext_exl3_kvtier", .flags = glue }, // ours: Flash Next K/V on the host tier, a device block cache in lent frames .{ .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 }; diff --git a/zig/kernels/cuda/flashnext_attn.cu b/zig/kernels/cuda/flashnext_attn.cu index ff4f7aabb7..d688326427 100644 --- a/zig/kernels/cuda/flashnext_attn.cu +++ b/zig/kernels/cuda/flashnext_attn.cu @@ -61,7 +61,7 @@ __device__ __forceinline__ int media_pos(int pos, int d, int half, const int* ro } // glue._attn_prep (text positions, bf16 cache), program (row, head): heads normed, rounded, roped, rounded, stored. -template +template __device__ __forceinline__ void fn_attn_prep_body(const __nv_bfloat16* __restrict__ P, const int* __restrict__ POS0, const float* __restrict__ QW, const float* __restrict__ KW, const float* __restrict__ IW, const float* __restrict__ INV, @@ -69,13 +69,13 @@ __device__ __forceinline__ void fn_attn_prep_body(const __nv_bfloat16* __restric __nv_bfloat16* __restrict__ VC, __nv_bfloat16* __restrict__ IQ, __nv_bfloat16* __restrict__ IKC, float eps, int PW, int NQ, int NKV, int NI, int IHD, int HALF, const int* rope3, int n, - int delta, const Pg& pg) { + int delta, const G& pg) { __shared__ float scratch[32]; const int r = blockIdx.x, head = blockIdx.y, d = threadIdx.x; const int pos = POS0[0] + r; const long long row = static_cast(r) * PW; if (head >= NQ + NKV + NI) { - const long long at = PAGED ? page_row(pg, pos, pg.ptok, IHD) : static_cast(pos) * IHD; + const long long at = PAGED ? ix_row(pg, pos, pg.ptok, IHD) : static_cast(pos) * IHD; if (d < IHD) IKC[at + d] = P[row + NQ * 2 * HD + 2 * NKV * HD + NI * IHD + d]; return; } @@ -160,19 +160,19 @@ extern "C" __global__ void __launch_bounds__(256) fn_attn_gate(const __nv_bfloat } // attention._pool, a block a completed block of RATIO indexer keys: their mean, rounded, normed and roped at its start. -template +template __device__ __forceinline__ void fn_qsa_pool_body(const __nv_bfloat16* __restrict__ IKC, __nv_bfloat16* __restrict__ POOLED, const int* __restrict__ POS0, const float* __restrict__ W, const float* __restrict__ INV, float eps, int R, int DI, int HALF, int RATIO, const int* rope3, int n, int delta, - const Pg& pg) { + const G& pg) { __shared__ float scratch[32]; __shared__ float xs[256]; const int p0 = POS0[0]; const int b = p0 / RATIO + blockIdx.x, d = threadIdx.x; if (RATIO * b + RATIO > p0 + R) return; // the block's keys: rows RATIO * b on, or (PAGED) the same rows of its one page (a page holds whole blocks) - const __nv_bfloat16* keys = IKC + (PAGED ? page_row(pg, RATIO * b, pg.ptok, DI) : 0); + const __nv_bfloat16* keys = IKC + (PAGED ? ix_row(pg, RATIO * b, pg.ptok, DI) : 0); const long long first = PAGED ? 0 : static_cast(RATIO) * b; float acc = bf(keys[first * DI + d]); for (int k = 1; k < RATIO; ++k) acc = __fadd_rn(acc, bf(keys[(first + k) * DI + d])); @@ -182,7 +182,7 @@ __device__ __forceinline__ void fn_qsa_pool_body(const __nv_bfloat16* __restrict const int p = partner(d, HALF); const float xn = rbf(__fmul_rn(__fmul_rn(x, rinv), W[d])); const float xpn = rbf(__fmul_rn(__fmul_rn(xs[p], rinv), W[p])); - const long long at = PAGED ? page_row(pg, b, pg.ptok / RATIO, DI) : static_cast(b) * DI; + const long long at = PAGED ? ix_row(pg, b, pg.ptok / RATIO, DI) : static_cast(b) * DI; POOLED[at + d] = __float2bfloat16_rn(rope(xn, xpn, d, HALF, media_pos(RATIO * b, d, HALF, rope3, n, delta), INV)); } @@ -584,3 +584,6 @@ extern "C" __global__ void __launch_bounds__(256) Pg pg) { qsa_tile::scores_body<8, 4, true>(IQ, POOLED, POS0, SC, NB, RATIO, TOP, rows, pg); } + +// The EXL3 launches through a host tier's two planes (PgP): the same bodies, the indexer keys by logical page. +#include "flashnext_planes.cuh" diff --git a/zig/kernels/cuda/flashnext_exl3/flashnext_exl3_kvtier.cu b/zig/kernels/cuda/flashnext_exl3/flashnext_exl3_kvtier.cu new file mode 100644 index 0000000000..a06b0222cc --- /dev/null +++ b/zig/kernels/cuda/flashnext_exl3/flashnext_exl3_kvtier.cu @@ -0,0 +1,189 @@ +// Host-tier K/V for Flash Next (EXL3): a layer's 4-row slots in lent frames; kvt_plan, kvt_fetch, kvt_remap. +#include +#include + +// One block of 1024 threads. +__device__ __forceinline__ void kvt_plan_body(const int* __restrict__ ids, const int* __restrict__ nkr, + const int* __restrict__ spr, int m, int idw, int top, + int ratio, int win0, unsigned* clock, int* blk2slot, + int* blk2slot_all, int key_off, + int* __restrict__ slot2blk, unsigned* __restrict__ stamp, + unsigned* __restrict__ mark, int nslots, int* __restrict__ ctl, + int* __restrict__ need, int* __restrict__ miss, + int* __restrict__ fetch) { + __shared__ int s_need, s_miss, s_avail; + __shared__ int hist[256]; + __shared__ unsigned cut_s; + __shared__ int left_s; + __shared__ unsigned epoch; + (void)nkr; + if (threadIdx.x == 0) { + s_need = 0; + s_miss = 0; + s_avail = 0; + if (*clock == UINT32_MAX) { + ctl[2] = 1; + epoch = 0; + } else { + epoch = ++*clock; + } + ctl[6] = static_cast(epoch); + ctl[1] = 0; + } + __syncthreads(); + if (epoch == 0 || ctl[2] != 0) return; + bool any = false; + for (int r = 0; r < m; ++r) any |= spr[r] != 0; + if (!any) { + if (threadIdx.x == 0) ctl[1] = 0; + return; + } + // the call's unique blocks before the window (a row lists its 512 blocks as 4 consecutive keys each) + for (int e = threadIdx.x; e < m * top; e += blockDim.x) { + const int r = e / top, j = e - r * top; + if (spr[r] == 0) continue; + const int b = ids[static_cast(r) * idw + j * ratio] / ratio; + if (b * ratio >= win0) continue; + if (atomicExch(&mark[b], epoch) != epoch) need[atomicAdd(&s_need, 1)] = b; + } + __syncthreads(); + const int n_need = s_need; + for (int i = threadIdx.x; i < n_need; i += blockDim.x) { + const int b = need[i]; + const int s = blk2slot[b]; + bool hit = s >= 0 && s < nslots; + if (hit) + stamp[s] = epoch; + else { + blk2slot[b] = -1; + miss[atomicAdd(&s_miss, 1)] = b; + } + } + __syncthreads(); + int n_miss = s_miss; + // slots this call does not use + for (int s = threadIdx.x; s < nslots; s += blockDim.x) + if (stamp[s] != epoch) atomicAdd(&s_avail, 1); + __syncthreads(); + if (threadIdx.x == 0) { + if (n_miss > s_avail) ctl[2] = 1; + ctl[1] = n_miss > s_avail ? 0 : n_miss; + unsigned* counters = reinterpret_cast(ctl); + counters[3] += ctl[1]; + counters[4] += n_need; + counters[5] += 1; + } + if (n_miss > s_avail) return; + n_miss = min(n_miss, s_avail); + if (n_miss == 0) return; + // the n_miss smallest stamps among unused slots: the largest keys of ~stamp (used slots key 0, never needed) + auto key = [&](int s) -> unsigned { const unsigned t = stamp[s]; return t == epoch ? 0u : ~t; }; + unsigned cut = 0; + int left = n_miss; + for (int p = 0; p < 4; ++p) { + const int shift = 24 - 8 * p; + for (int i = threadIdx.x; i < 256; i += blockDim.x) hist[i] = 0; + __syncthreads(); + for (int s = threadIdx.x; s < nslots; s += blockDim.x) { + const unsigned k = key(s); + if (p == 0 || (k >> (shift + 8)) == (cut >> (shift + 8))) atomicAdd(&hist[(k >> shift) & 0xFF], 1); + } + __syncthreads(); + if (threadIdx.x == 0) { + int above = 0; + int d = 255; + for (; d > 0; --d) { + if (above + hist[d] >= left) break; + above += hist[d]; + } + left_s = left - above; + cut_s = cut | (static_cast(d) << shift); + } + __syncthreads(); + cut = cut_s; + left = left_s; + __syncthreads(); + } + // every slot above the cut and `left` of those equal to it (any of them: equal stamps are equally old) take the + __shared__ int s_eq, s_place; + if (threadIdx.x == 0) { + s_eq = 0; + s_place = 0; + } + __syncthreads(); + for (int s = threadIdx.x; s < nslots; s += blockDim.x) { + const unsigned k = key(s); + bool chosen = k > cut; + if (!chosen && k == cut) chosen = atomicAdd(&s_eq, 1) < left; + if (!chosen) continue; + const int place = atomicAdd(&s_place, 1); + if (place >= n_miss) continue; + const int b = miss[place]; + const int old = slot2blk[s]; + if (old >= 0) blk2slot_all[old] = -1; + blk2slot[b] = s; + slot2blk[s] = b; + stamp[s] = epoch; + fetch[2 * place] = b; + fetch[2 * place + 1] = s; + } +} + +extern "C" __global__ void __launch_bounds__(1024) kvt_plan( + const int* ids, const int* counts, const int* sparse, int rows, int id_width, int top_blocks, + int ratio, int window_origin, unsigned* clock, int* inverse, int* all_inverse, int key_offset, + int* owners, unsigned* stamps, unsigned* marks, int slots, int* controls, int* needed, + int* missing, int* fetched) { + kvt_plan_body(ids, counts, sparse, rows, id_width, top_blocks, ratio, window_origin, + clock, inverse, all_inverse, key_offset, owners, stamps, marks, slots, controls, needed, + missing, fetched); +} + +// Each miss's 4 rows of K and V: host store (hk, hv) to slot s: row (s % 64) * 4 of entry vpage + s / 64. +extern "C" __global__ void __launch_bounds__(256) kvt_fetch(const uint4* __restrict__ hk, const uint4* __restrict__ hv, + uint4* __restrict__ dk, uint4* __restrict__ dv, + const int* __restrict__ ctl, const int* __restrict__ fetch, + const int* __restrict__ lids, const int* __restrict__ pt, + int vpage, long long ps, int row16, int ratio) { + const int n = ctl[1]; + const long long per = static_cast(ratio) * row16; + for (int i = blockIdx.x; i < n; i += gridDim.x) { + const long long b = fetch[2 * i], s = fetch[2 * i + 1]; + const long long page = b * ratio / 256, in_page = b * ratio % 256; + const long long src = static_cast(lids[page]) * ps + in_page * row16; + const long long dst = static_cast(pt[vpage + (s >> 6)]) * ps + (s & 63) * per; + for (int e = threadIdx.x; e < per; e += blockDim.x) { + dk[dst + e] = hk[src + e]; + dv[dst + e] = hv[src + e]; + } + } +} + +// A row a block: its key list from positions to table rows: the window's stay, older ones take their slot's. +extern "C" __global__ void __launch_bounds__(256) kvt_remap(int* __restrict__ ids, int* __restrict__ nkr, + const int* __restrict__ spr, int idw, int win0, int vrow, + const int* __restrict__ blk2slot, int ratio, int* ctl) { + const int r = blockIdx.x; + if (spr[r] == 0) return; + const int n = nkr[r]; + int* row = ids + static_cast(r) * idw; + __shared__ int missing; + if (threadIdx.x == 0) missing = atomicAdd(&ctl[2], 0) != 0; + __syncthreads(); + for (int k = threadIdx.x; k < n; k += blockDim.x) { + const int p = row[k]; + if (p < win0 && blk2slot[p / ratio] < 0) atomicExch(&missing, 1); + } + __syncthreads(); + if (missing) { + if (threadIdx.x == 0) { + atomicExch(&ctl[2], 1); + nkr[r] = 0; + } + return; + } + for (int k = threadIdx.x; k < n; k += blockDim.x) { + const int p = row[k]; + row[k] = p >= win0 ? p : vrow + blk2slot[p / ratio] * ratio + p % ratio; + } +} diff --git a/zig/kernels/cuda/flashnext_pages.cuh b/zig/kernels/cuda/flashnext_pages.cuh index b25aed86fd..c97bcb31bc 100644 --- a/zig/kernels/cuda/flashnext_pages.cuh +++ b/zig/kernels/cuda/flashnext_pages.cuh @@ -9,3 +9,25 @@ struct Pg { __device__ __forceinline__ long long page_row(const Pg& pg, int p, int per, int row) { return static_cast(pg.pt[p / per]) * pg.ps + static_cast(p % per) * row; } + +// Pages of a host tier: K/V rows by frame (pt, ps as in Pg), indexer keys and pooled rows by logical page (ipt, ips). +struct PgP { + const int* pt; + long long ps; + int ptok; + long long ips; // elements (bf16) a page holds of the index plane + const int* ipt; // the index plane's table, an int a page +}; + +__device__ __forceinline__ long long page_row(const PgP& pg, int p, int per, int row) { + return static_cast(pg.pt[p / per]) * pg.ps + static_cast(p % per) * row; +} + +// Row p of an indexer-key or pooled cache: its page's start in the index plane (a Pg has one plane), then its own row. +__device__ __forceinline__ long long ix_row(const Pg& pg, int p, int per, int row) { + return page_row(pg, p, per, row); +} + +__device__ __forceinline__ long long ix_row(const PgP& pg, int p, int per, int row) { + return static_cast(pg.ipt[p / per]) * pg.ips + static_cast(p % per) * row; +} diff --git a/zig/kernels/cuda/flashnext_planes.cuh b/zig/kernels/cuda/flashnext_planes.cuh new file mode 100644 index 0000000000..3ab4bf5bd1 --- /dev/null +++ b/zig/kernels/cuda/flashnext_planes.cuh @@ -0,0 +1,47 @@ +// Flash Next attention launches for a stream whose K/V rows sit in frames and whose indexer keys sit by logical page. +#pragma once + +// The same two through a host tier's planes (PgP): the K/V rows by frame, the indexer keys by logical page. +extern "C" __global__ void __launch_bounds__(256) + fn_attn_prep_planes(const __nv_bfloat16* __restrict__ P, const int* __restrict__ POS0, const float* __restrict__ QW, + const float* __restrict__ KW, const float* __restrict__ IW, const float* __restrict__ INV, + __nv_bfloat16* __restrict__ Q, __nv_bfloat16* __restrict__ KC, __nv_bfloat16* __restrict__ VC, + __nv_bfloat16* __restrict__ IQ, __nv_bfloat16* __restrict__ IKC, float eps, int PW, int NQ, + int NKV, int NI, int IHD, int HALF, PgP pg) { + fn_attn_prep_body(P, POS0, QW, KW, IW, INV, Q, KC, VC, IQ, IKC, eps, PW, NQ, NKV, NI, IHD, HALF, nullptr, 0, + 0, pg); +} + +extern "C" __global__ void __launch_bounds__(256) + fn_attn_prep_media_planes(const __nv_bfloat16* __restrict__ P, const int* __restrict__ POS0, + const float* __restrict__ QW, const float* __restrict__ KW, const float* __restrict__ IW, + const float* __restrict__ INV, __nv_bfloat16* __restrict__ Q, + __nv_bfloat16* __restrict__ KC, __nv_bfloat16* __restrict__ VC, + __nv_bfloat16* __restrict__ IQ, __nv_bfloat16* __restrict__ IKC, float eps, int PW, + int NQ, int NKV, int NI, int IHD, int HALF, const int* rope3, int n, int delta, PgP pg) { + fn_attn_prep_body(P, POS0, QW, KW, IW, INV, Q, KC, VC, IQ, IKC, eps, PW, NQ, NKV, NI, IHD, HALF, rope3, n, + delta, pg); +} + +extern "C" __global__ void __launch_bounds__(128) + fn_qsa_pool_planes(const __nv_bfloat16* __restrict__ IKC, __nv_bfloat16* __restrict__ POOLED, + const int* __restrict__ POS0, const float* __restrict__ W, const float* __restrict__ INV, + float eps, int R, int DI, int HALF, int RATIO, PgP pg) { + fn_qsa_pool_body(IKC, POOLED, POS0, W, INV, eps, R, DI, HALF, RATIO, nullptr, 0, 0, pg); +} + +extern "C" __global__ void __launch_bounds__(128) + fn_qsa_pool_media_planes(const __nv_bfloat16* __restrict__ IKC, __nv_bfloat16* __restrict__ POOLED, + const int* __restrict__ POS0, const float* __restrict__ W, const float* __restrict__ INV, + float eps, int R, int DI, int HALF, int RATIO, const int* rope3, int n, int delta, + PgP pg) { + fn_qsa_pool_body(IKC, POOLED, POS0, W, INV, eps, R, DI, HALF, RATIO, rope3, n, delta, pg); +} + +// The same sums through a host tier's planes: the pooled rows by logical page. +extern "C" __global__ void __launch_bounds__(256) + fn_qsa_scores_planes(const __nv_bfloat16* __restrict__ IQ, const __nv_bfloat16* __restrict__ POOLED, + const int* __restrict__ POS0, float* __restrict__ SC, int NB, int RATIO, int TOP, int rows, + PgP pg) { + qsa_tile::scores_body<8, 4, true>(IQ, POOLED, POS0, SC, NB, RATIO, TOP, rows, pg); +} diff --git a/zig/kernels/cuda/flashnext_qsa_tile.cuh b/zig/kernels/cuda/flashnext_qsa_tile.cuh index 5ead983593..28a57e8cbd 100644 --- a/zig/kernels/cuda/flashnext_qsa_tile.cuh +++ b/zig/kernels/cuda/flashnext_qsa_tile.cuh @@ -16,11 +16,11 @@ struct Stack { } }; // PAGED: POOLED is the arena's slot in page 0, a block's row in its page (fn_attn's Pg): one lookup a thread. -template +template __device__ __forceinline__ void scores_body(const __nv_bfloat16* __restrict__ IQ, const __nv_bfloat16* __restrict__ POOLED, const int* __restrict__ POS0, float* __restrict__ SC, int NB, int RATIO, int TOP, int rows, - const Pg& pg) { + const G& pg) { __shared__ float4 qs[R][HI][DI / 4]; const int r0 = blockIdx.x * R, b = blockIdx.y * 256 + threadIdx.x; const int pos0 = POS0[0]; @@ -36,7 +36,7 @@ __device__ __forceinline__ void scores_body(const __nv_bfloat16* __restrict__ IQ __syncthreads(); if (b >= cmax) return; Stack st[R][HI]; - const __nv_bfloat16* kb = POOLED + (PAGED ? page_row(pg, b, pg.ptok / RATIO, DI) : static_cast(b) * DI); + const __nv_bfloat16* kb = POOLED + (PAGED ? ix_row(pg, b, pg.ptok / RATIO, DI) : static_cast(b) * DI); #pragma unroll for (int half = 0; half < 2; ++half) { float k[16][4]; // the 16 lanes l (bit 0 == half) a half's leaves read diff --git a/zig/src/core/contract/flashnext_exl3_pending.zig b/zig/src/core/contract/flashnext_exl3_pending.zig index 1d57be4c76..1f7461dcb7 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 (pick mask, prefill tiles in decode's bytes); nvcc 13.4 measured, older pending. +//! Flash Next EXL3 launches (pick mask, prefill tiles, host-tier K/V); nvcc 13.4 measured, older pending. const std = @import("std"); const limits = @import("limits.zig"); const registry = @import("registry.zig"); @@ -18,6 +18,9 @@ pub const records = [_]registry.CudaRecord{ .{ .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 }, + .{ .launch = .{ .name = "flashnext_exl3_kvtier kvt_plan", .symbol = "kvt_plan", .image = "flashnext_exl3_kvtier", .threads = 1024 }, .usage = registry.none }, + .{ .launch = .{ .name = "flashnext_exl3_kvtier kvt_fetch", .symbol = "kvt_fetch", .image = "flashnext_exl3_kvtier", .threads = 256 }, .usage = registry.none }, + .{ .launch = .{ .name = "flashnext_exl3_kvtier kvt_remap", .symbol = "kvt_remap", .image = "flashnext_exl3_kvtier", .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. @@ -35,6 +38,9 @@ pub const cuda_134 = [_]registry.Override{ .{ .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 } } }, + .{ .name = "flashnext_exl3_kvtier kvt_plan", .usage = .{ .{ .regs = 27, .static = 1056, .bound = 1024 }, .{ .regs = 27, .static = 1056, .bound = 1024 }, .{ .regs = 27, .static = 1056, .bound = 1024 }, .{ .regs = 34, .static = 1056, .bound = 1024 }, .{ .regs = 34, .static = 1056, .bound = 1024 } } }, + .{ .name = "flashnext_exl3_kvtier kvt_fetch", .usage = .{ .{ .regs = 28, .static = 0, .bound = 256 }, .{ .regs = 28, .static = 0, .bound = 256 }, .{ .regs = 28, .static = 0, .bound = 256 }, .{ .regs = 28, .static = 0, .bound = 256 }, .{ .regs = 28, .static = 0, .bound = 256 } } }, + .{ .name = "flashnext_exl3_kvtier kvt_remap", .usage = .{ .{ .regs = 19, .static = 4, .bound = 256 }, .{ .regs = 19, .static = 4, .bound = 256 }, .{ .regs = 19, .static = 4, .bound = 256 }, .{ .regs = 24, .static = 4, .bound = 256 }, .{ .regs = 24, .static = 4, .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/core/contract/flashnext_pages_pending.zig b/zig/src/core/contract/flashnext_pages_pending.zig index 5b8826d7a1..a9c4de2459 100644 --- a/zig/src/core/contract/flashnext_pages_pending.zig +++ b/zig/src/core/contract/flashnext_pages_pending.zig @@ -1,4 +1,4 @@ -//! Paged Flash Next cache launches: measured nvcc 13.0 resources for all five SM targets. +//! Paged Flash Next cache launches: measured nvcc resources for all five SM targets (the planes launches: nvcc 13.4). const registry = @import("registry.zig"); pub const records = [_]registry.CudaRecord{ .{ .launch = .{ .name = "flashnext_attn fn_attn_prep_paged", .symbol = "fn_attn_prep_paged", .image = "flashnext_attn", .threads = 256 }, .usage = .{ .{ .regs = 32, .static = 128, .bound = 256 }, .{ .regs = 36, .static = 128, .bound = 256 }, .{ .regs = 36, .static = 128, .bound = 256 }, .{ .regs = 30, .static = 128, .bound = 256 }, .{ .regs = 30, .static = 128, .bound = 256 } } }, @@ -9,6 +9,11 @@ pub const records = [_]registry.CudaRecord{ .{ .launch = .{ .name = "flashnext_attn fn_attn_chunks_ahead_paged", .symbol = "fn_attn_chunks_ahead_paged", .image = "flashnext_attn", .threads = 128, .dynamic = 77312, .opt_in = true }, .usage = .{ .{ .regs = 76, .static = 1280, .bound = 128 }, .{ .regs = 76, .static = 1280, .bound = 128 }, .{ .regs = 76, .static = 1280, .bound = 128 }, .{ .regs = 73, .static = 1280, .bound = 128 }, .{ .regs = 73, .static = 1280, .bound = 128 } } }, .{ .launch = .{ .name = "flashnext_attn fn_attn_chunks_one_paged", .symbol = "fn_attn_chunks_one_paged", .image = "flashnext_attn", .threads = 128, .dynamic = 26624, .opt_in = true }, .usage = .{ .{ .regs = 72, .static = 896, .bound = 128 }, .{ .regs = 72, .static = 896, .bound = 128 }, .{ .regs = 72, .static = 896, .bound = 128 }, .{ .regs = 96, .static = 896, .bound = 128 }, .{ .regs = 96, .static = 896, .bound = 128 } } }, .{ .launch = .{ .name = "flashnext_attn fn_attn_fused_paged", .symbol = "fn_attn_fused_paged", .image = "flashnext_attn", .threads = 128, .dynamic = 26624 }, .usage = .{ .{ .regs = 118, .static = 896, .bound = 128 }, .{ .regs = 118, .static = 896, .bound = 128 }, .{ .regs = 118, .static = 896, .bound = 128 }, .{ .regs = 123, .static = 896, .bound = 128 }, .{ .regs = 123, .static = 896, .bound = 128 } } }, + .{ .launch = .{ .name = "flashnext_attn fn_attn_prep_planes", .symbol = "fn_attn_prep_planes", .image = "flashnext_attn", .threads = 256 }, .usage = .{ .{ .regs = 32, .static = 128, .bound = 256 }, .{ .regs = 36, .static = 128, .bound = 256 }, .{ .regs = 36, .static = 128, .bound = 256 }, .{ .regs = 32, .static = 128, .bound = 256 }, .{ .regs = 32, .static = 128, .bound = 256 } } }, + .{ .launch = .{ .name = "flashnext_attn fn_attn_prep_media_planes", .symbol = "fn_attn_prep_media_planes", .image = "flashnext_attn", .threads = 256 }, .usage = .{ .{ .regs = 32, .static = 128, .bound = 256 }, .{ .regs = 36, .static = 128, .bound = 256 }, .{ .regs = 36, .static = 128, .bound = 256 }, .{ .regs = 32, .static = 128, .bound = 256 }, .{ .regs = 32, .static = 128, .bound = 256 } } }, + .{ .launch = .{ .name = "flashnext_attn fn_qsa_pool_planes", .symbol = "fn_qsa_pool_planes", .image = "flashnext_attn", .threads = 128 }, .usage = .{ .{ .regs = 32, .static = 1152, .bound = 128 }, .{ .regs = 40, .static = 1152, .bound = 128 }, .{ .regs = 40, .static = 1152, .bound = 128 }, .{ .regs = 40, .static = 1152, .bound = 128 }, .{ .regs = 40, .static = 1152, .bound = 128 } } }, + .{ .launch = .{ .name = "flashnext_attn fn_qsa_pool_media_planes", .symbol = "fn_qsa_pool_media_planes", .image = "flashnext_attn", .threads = 128 }, .usage = .{ .{ .regs = 32, .static = 1152, .bound = 128 }, .{ .regs = 40, .static = 1152, .bound = 128 }, .{ .regs = 40, .static = 1152, .bound = 128 }, .{ .regs = 40, .static = 1152, .bound = 128 }, .{ .regs = 40, .static = 1152, .bound = 128 } } }, + .{ .launch = .{ .name = "flashnext_attn fn_qsa_scores_planes", .symbol = "fn_qsa_scores_planes", .image = "flashnext_attn", .threads = 256 }, .usage = .{ .{ .regs = 220, .static = 16384, .bound = 256 }, .{ .regs = 219, .static = 16384, .bound = 256 }, .{ .regs = 219, .static = 16384, .bound = 256 }, .{ .regs = 219, .static = 16384, .bound = 256 }, .{ .regs = 219, .static = 16384, .bound = 256 } } }, }; test "Flash Next paged launches: measured resources satisfy every SM launch limit" { const t = @import("std").testing; diff --git a/zig/src/core/kv_tier.zig b/zig/src/core/kv_tier.zig index 3d40f1d98c..ed447dbbcc 100644 --- a/zig/src/core/kv_tier.zig +++ b/zig/src/core/kv_tier.zig @@ -16,6 +16,25 @@ pub const Shape = struct { } }; +/// A rank's shape, slots and digest (16-bit float pieces), told to the other ranks at open; a differing one is refused. +pub const Lockstep = struct { + pub const len = 8; + + pub fn vector(t: *const Tier, slots: usize) [len]f32 { + const d = t.digest(); + var v: [len]f32 = @splat(0); + v[0] = @floatFromInt(t.pool.total()); + v[1] = @floatFromInt(t.count()); + v[2] = @floatFromInt(slots); + for (0..4) |i| v[3 + i] = @floatFromInt((d >> @intCast(16 * i)) & 0xffff); + return v; + } + + pub fn agree(a: [len]f32, b: [len]f32) error{RanksDisagree}!void { + if (!std.mem.eql(f32, &a, &b)) return error.RanksDisagree; + } +}; + /// A deliberate bug a test turns on to show that the check beside it fails without the rule: never set outside tests. pub const Fault = enum { none, keep_generation, ignore_pins, ignore_dirty, unpinned_dirty, keep_failed_miss }; @@ -279,6 +298,17 @@ pub const Tier = struct { return t.host[id * b ..][0..b]; } + /// `id`'s host rows were rewritten (a restored snapshot): its frame is dropped; Pinned while a call reads it. + pub fn invalidate(t: *Tier, id: u32) error{Pinned}!void { + const f = t.frame_of[id]; + if (f == none) return; + if (t.pinned[f] > 0) return error.Pinned; + t.dirty[id] = false; // the frame's rows are being replaced + t.unmap(f); + t.free_frames[t.nfree] = f; + t.nfree += 1; + } + /// Drop up to `n` unpinned frames, least recently used first; the count dropped. pub fn evict(t: *Tier, n: usize) usize { var done: usize = 0; @@ -429,4 +459,5 @@ pub const Tier = struct { test { _ = @import("kv_tier_test.zig"); + _ = @import("kv_tier_ranks_test.zig"); } diff --git a/zig/src/core/kv_tier_ranks_test.zig b/zig/src/core/kv_tier_ranks_test.zig new file mode 100644 index 0000000000..8b5250f811 --- /dev/null +++ b/zig/src/core/kv_tier_ranks_test.zig @@ -0,0 +1,49 @@ +//! The host tier: a rewritten page loses its frame, and ranks open only in lockstep. +const std = @import("std"); +const kt = @import("kv_tier.zig"); +const Rig = @import("kv_tier_test.zig").Rig; +const Lease = kt.Lease; +const expect = std.testing.expect; + +test "a page whose host rows were rewritten loses its frame, so the next pin copies them in; a pinned one refuses" { + var r: Rig = undefined; + try r.open(4, 2); + defer r.close(); + const t = &r.tier; + var a = try r.lease(2, 2); + defer t.finish(&a); + const p = a.pages.ids.items; + _ = try r.pin(&a, p[0]); + r.write(&a, p[0], 4); + r.flush(p[0]); + try std.testing.expectError(error.Pinned, t.invalidate(p[0])); + try t.unpin(&a, p[0]); + @memset(r.tier.host_page(p[0]), 9); // the restore + try t.invalidate(p[0]); + try expect(t.resident(p[0]) == null); + try t.audit(); + const f = try r.pin(&a, p[0]); + try expect(f.miss); + try t.invalidate(p[1]); // no frame: nothing to do +} + +test "ranks open only on the same shape, slots and digest; a rank that differs in any is refused" { + var a: Rig = undefined; + var b: Rig = undefined; + try a.open(8, 4); + defer a.close(); + try b.open(8, 4); + defer b.close(); + const va = kt.Lockstep.vector(&a.tier, 16); + try kt.Lockstep.agree(va, kt.Lockstep.vector(&b.tier, 16)); + for ([_][2]usize{ .{ 8, 3 }, .{ 10, 4 } }) |shape| { // fewer frames, more pages + var c: Rig = undefined; + try c.open(shape[0], shape[1]); + defer c.close(); + try std.testing.expectError(error.RanksDisagree, kt.Lockstep.agree(va, kt.Lockstep.vector(&c.tier, 16))); + } + try std.testing.expectError(error.RanksDisagree, kt.Lockstep.agree(va, kt.Lockstep.vector(&b.tier, 17))); + var l = try b.lease(1, 1); // a different history: the digest moves + defer b.tier.finish(&l); + try std.testing.expectError(error.RanksDisagree, kt.Lockstep.agree(va, kt.Lockstep.vector(&b.tier, 16))); +} diff --git a/zig/src/core/kv_tier_test.zig b/zig/src/core/kv_tier_test.zig index 95437b2bea..e365d4f52e 100644 --- a/zig/src/core/kv_tier_test.zig +++ b/zig/src/core/kv_tier_test.zig @@ -8,15 +8,15 @@ const expect = std.testing.expect; const expectEqual = std.testing.expectEqual; const gpa = std.testing.allocator; -const page_bytes = 8; +pub const page_bytes = 8; /// A tier of `pages` logical pages and `frames` frames over a host store and a fake device. -const Rig = struct { +pub const Rig = struct { host: []u8, dev: []u8, tier: Tier, - fn open(r: *Rig, pages: usize, frames: usize) !void { + pub fn open(r: *Rig, pages: usize, frames: usize) !void { r.host = try gpa.alloc(u8, pages * page_bytes); errdefer gpa.free(r.host); @memset(r.host, 0); @@ -26,23 +26,23 @@ const Rig = struct { try r.tier.init(gpa, r.host, frames, page_bytes); } - fn close(r: *Rig) void { + pub fn close(r: *Rig) void { r.tier.deinit(); gpa.free(r.host); gpa.free(r.dev); } - fn frame(r: *Rig, f: u32) []u8 { + pub fn frame(r: *Rig, f: u32) []u8 { return r.dev[f * page_bytes ..][0..page_bytes]; } /// A stream with `pages` pages covered and `frames` frames promised. - fn lease(r: *Rig, pages: usize, frames: usize) !Lease { + pub fn lease(r: *Rig, pages: usize, frames: usize) !Lease { return r.leaseLent(pages, frames, 0); } /// A stream that also has `lent` of its frames lent for its life. - fn leaseLent(r: *Rig, pages: usize, frames: usize, lent: usize) !Lease { + pub fn leaseLent(r: *Rig, pages: usize, frames: usize, lent: usize) !Lease { var l: Lease = .{}; errdefer r.tier.finish(&l); try r.tier.reserve(&l, pages, frames); @@ -52,20 +52,20 @@ const Rig = struct { } /// Pin `id`, copying its canonical rows in on a miss. - fn pin(r: *Rig, l: *Lease, id: u32) !kt.Fill { + pub fn pin(r: *Rig, l: *Lease, id: u32) !kt.Fill { const f = try r.tier.pin(l, id); if (f.miss) @memcpy(r.frame(f.frame), r.tier.host_page(id)); return f; } /// Write `b` into every row byte of the pinned page and flush it to the host store. - fn write(r: *Rig, l: *const Lease, id: u32, b: u8) void { + pub fn write(r: *Rig, l: *const Lease, id: u32, b: u8) void { const tk = r.tier.resident(id).?; @memset(r.frame(tk.frame), b); r.tier.markDirty(l, id) catch unreachable; } - fn flush(r: *Rig, id: u32) void { + pub fn flush(r: *Rig, id: u32) void { const tk = r.tier.resident(id).?; @memcpy(r.tier.host_page(id), r.frame(tk.frame)); r.tier.sync(id); diff --git a/zig/src/cuda/kernels.zig b/zig/src/cuda/kernels.zig index 132624ecc7..32aa9d4077 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_kvtier: []const u8 = if (available) &Blob("fatbin_flashnext_exl3_kvtier").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 &.{}; diff --git a/zig/src/families/flashnext/cuda_batch.zig b/zig/src/families/flashnext/cuda_batch.zig index adc5b674be..4652b924bb 100644 --- a/zig/src/families/flashnext/cuda_batch.zig +++ b/zig/src/families/flashnext/cuda_batch.zig @@ -369,12 +369,14 @@ pub fn Pass(comptime Tri: type) type { const ikc = self.seq.ikc[ai].ptr; const pooled = self.seq.pooled[ai].ptr; const pg = self.seq.pg; + const sl = self.seq.sel_layer + ai; + if (self.seq.sel) |t| try t.begin(self.stream, sl, self.pos0, n); if (self.seq.media) |m| try m.attn(self.stream, pa, at.pos, lw.q_scale.ptr, lw.k_scale.ptr, lw.i_scale.ptr, s.inv, at.q, kc, vc, at.iq, ikc, eps, n, pg) else try self.tri.attnPrep(pa, at.pos, lw.q_scale.ptr, lw.k_scale.ptr, lw.i_scale.ptr, s.inv, at.q, kc, vc, at.iq, ikc, s.dummy, s.dummy, eps, n, 0, pg); // attention.qsa_pool for the blocks these rows complete, then each block of rows' scores and attention if (self.seq.media) |m| try m.pooled(self.stream, ikc, pooled, at.pos, lw.ik_scale.ptr, s.inv, eps, n, pg) else try self.tri.qsaPool(ikc, pooled, at.pos, lw.ik_scale.ptr, s.inv, eps, n, pg); const nb = scoreStride(self.seq.cap); // prompt rows: selections for up to fused_rows rows, then one fused launch (chunks merged in the program) - const fused = n > att_rows and self.tri.fusedAttentionOwned(); + const fused = n > att_rows and self.tri.fusedAttentionOwned() and self.seq.sel == null; const group: usize = if (fused) parts.fused_rows else att_rows; var g0: usize = 0; while (g0 < n) : (g0 += group) { @@ -400,7 +402,9 @@ pub fn Pass(comptime Tri: type) type { a.add(cint(top)); a.add(cint(ratio)); try cuda.launch.launch(self.k.select, .{ .grid = .{ .x = @intCast(m) }, .block = .{ .x = 1024 } }, self.stream, &a); - if (!fused) { + if (self.seq.sel) |t| { + try @import("exl3/kv/cuda_exl3_kvtier.zig").attend(t, self, sl, .{ .q = at.q + r0 * q_heads * head_dim * 2, .ao = at.ao + r0 * out_k * 2, .kc = kc, .vc = vc, .pos = pos, .sel = sel, .m = m, .ends = ends, .first_row = self.pos0 + r0 }); + } else if (!fused) { const chunks = cdiv(@min(ends, wt.dense_keys), 512); try self.tri.attnChunks(at.q + r0 * q_heads * head_dim * 2, kc, vc, pos, s.po, s.pm, s.pl, s.qids, s.qnk, s.qspr, m, kv_heads, chunks, true, pg); try self.tri.attnMerge(s.po, s.pm, s.pl, pos, at.ao + r0 * out_k * 2, s.qnk, s.qspr, m, kv_heads, true); diff --git a/zig/src/families/flashnext/cuda_batch_state.zig b/zig/src/families/flashnext/cuda_batch_state.zig index b49c64539a..00a69045b5 100644 --- a/zig/src/families/flashnext/cuda_batch_state.zig +++ b/zig/src/families/flashnext/cuda_batch_state.zig @@ -491,6 +491,10 @@ pub const SeqState = struct { scores: u64, /// The attention caches' pages (pt 0: kc, vc, ikc and pooled are contiguous rows). pg: @import("cuda_glue.zig").Pg = .{}, + /// A host tier's selection of blocks in slots (null: all in frames). + sel: ?*@import("exl3/kv/cuda_exl3_kvtier.zig").LaneSel = null, + /// The selection's layer of this cache's layer 0 (the MTP head's cache is one layer after the model's). + sel_layer: usize = 0, }; /// The block scores' row stride for a cache of `cap` tokens: its blocks rounded up to 16 (the captured _scores'). diff --git a/zig/src/families/flashnext/cuda_glue.zig b/zig/src/families/flashnext/cuda_glue.zig index 14a17dc451..87b855549d 100644 --- a/zig/src/families/flashnext/cuda_glue.zig +++ b/zig/src/families/flashnext/cuda_glue.zig @@ -27,7 +27,25 @@ fn cdiv(a: usize, b: usize) usize { } /// flashnext_attn.cu's Pg, each cache launch's `pg`: its table, elements and positions a page (pt 0: contiguous). -pub const Pg = extern struct { pt: u64 = 0, ps: i64 = 0, ptok: c_int = 0 }; +pub const PgMain = extern struct { pt: u64 = 0, ps: i64 = 0, ptok: c_int = 0 }; + +/// A launch's pages: `PgMain` and (ips, ipt) for indexer keys in their own plane (PgP); one-plane kernels: `PgMain`. +pub const Pg = extern struct { + pt: u64 = 0, + ps: i64 = 0, + ptok: c_int = 0, + ips: i64 = 0, + ipt: u64 = 0, + + pub fn main(pg: Pg) PgMain { + return .{ .pt = pg.pt, .ps = pg.ps, .ptok = pg.ptok }; + } + + /// Whether the indexer keys have their own plane: the five kernels that touch them take the whole struct. + pub fn planes(pg: Pg) bool { + return pg.ips != 0; + } +}; /// The native kernels, by their extern "C" names. const Fns = struct { @@ -63,6 +81,11 @@ const Fns = struct { merge: cuda.Function, /// The cache entries through a page table (TF_FLASHNEXT_KV_PAGES). prep_paged: cuda.Function, + prep_planes: cuda.Function, + prep_media_planes: cuda.Function, + pool_planes: cuda.Function, + pool_media_planes: cuda.Function, + scores_planes: cuda.Function, prep_media_paged: cuda.Function, pool_paged: cuda.Function, pool_media_paged: cuda.Function, @@ -149,6 +172,11 @@ pub const Glue = struct { .fused = try g.attn.function("fn_attn_fused"), .merge = try g.attn.function("fn_attn_merge"), .prep_paged = try g.attn.function("fn_attn_prep_paged"), + .prep_planes = try g.attn.function("fn_attn_prep_planes"), + .prep_media_planes = try g.attn.function("fn_attn_prep_media_planes"), + .pool_planes = try g.attn.function("fn_qsa_pool_planes"), + .pool_media_planes = try g.attn.function("fn_qsa_pool_media_planes"), + .scores_planes = try g.attn.function("fn_qsa_scores_planes"), .prep_media_paged = try g.attn.function("fn_attn_prep_media_paged"), .pool_paged = try g.attn.function("fn_qsa_pool_paged"), .pool_media_paged = try g.attn.function("fn_qsa_pool_media_paged"), @@ -226,8 +254,8 @@ pub const Glue = struct { ptrs(&a, &.{ proj, pos, q_scale, k_scale, i_scale, inv, q, kc, vc, iq, ikc }); a.add(eps); ints(&a, &.{ wt.proj_n, wt.q_heads, wt.kv_heads, wt.index_heads, wt.index_dim, wt.half }); - if (pg.pt != 0) a.add(pg); - const f = if (pg.pt != 0) g.f.prep_paged else g.f.prep; + if (pg.pt != 0) if (pg.planes()) a.add(pg) else a.add(pg.main()); + const f = if (pg.pt != 0) (if (pg.planes()) g.f.prep_planes else g.f.prep_paged) else g.f.prep; try go(f, s, .{ rows, wt.q_heads + wt.kv_heads + wt.index_heads + 1, 1 }, tile_head_dim, 0, &a); } @@ -240,7 +268,7 @@ pub const Glue = struct { ints(&a, &.{ wt.q_heads, kv, group, 512, if (qsa) wt.nch else wt.dense_nch }); a.add(@as(f32, 0.0625)); ints(&a, &.{ wt.index_idw, @intFromBool(qsa) }); - if (pg.pt != 0) a.add(pg); + if (pg.pt != 0) a.add(pg.main()); const ahead = rows <= ahead_rows; const f = if (ahead) g.f.chunks_ahead else g.f.chunks_one; const paged = if (ahead) g.f.chunks_ahead_paged else g.f.chunks_one_paged; @@ -256,7 +284,7 @@ pub const Glue = struct { ints(&a, &.{ wt.q_heads, kv, group, 512 }); a.add(@as(f32, 0.0625)); ints(&a, &.{ wt.index_idw, 1 }); - if (pg.pt != 0) a.add(pg); + if (pg.pt != 0) a.add(pg.main()); try go(if (pg.pt != 0) g.f.fused_paged else g.f.fused, s, .{ rows, kv, 1 }, 128, oneSmem(), &a); } @@ -272,8 +300,8 @@ pub const Glue = struct { ptrs(&a, &.{ ikc, pooled, pos, k_scale, inv }); a.add(eps); ints(&a, &.{ rows, wt.index_dim, wt.half, wt.index_ratio }); - if (pg.pt != 0) a.add(pg); - const f = if (pg.pt != 0) g.f.pool_paged else g.f.pool; + if (pg.pt != 0) if (pg.planes()) a.add(pg) else a.add(pg.main()); + const f = if (pg.pt != 0) (if (pg.planes()) g.f.pool_planes else g.f.pool_paged) else g.f.pool; try go(f, s, .{ rows / wt.index_ratio + 2, 1, 1 }, wt.index_dim, 0, &a); } @@ -284,8 +312,8 @@ pub const Glue = struct { if (wt.index_heads == 4) { // 8 rows and 256 blocks a CTA, every head, fn_qsa_scores' sums in its order (its bits) ints(&a, &.{ nb, wt.index_ratio, wt.index_top, rows }); - if (pg.pt != 0) a.add(pg); - const f = if (pg.pt != 0) g.f.scores_paged else g.f.scores_tile; + if (pg.pt != 0) if (pg.planes()) a.add(pg) else a.add(pg.main()); + const f = if (pg.pt != 0) (if (pg.planes()) g.f.scores_planes else g.f.scores_paged) else g.f.scores_tile; return go(f, s, .{ cdiv(rows, 8), cdiv(blocks, 256), 1 }, 256, 0, &a); } if (pg.pt != 0) return error.RowShape; // pages hold the tiles' four indexer heads only @@ -452,9 +480,12 @@ test "the attention tiles hold the model's heads" { } test "Pg lays out as flashnext_attn.cu's: the table, elements a page at 8, positions a page at 16, 24 bytes" { - try std.testing.expectEqual(@as(usize, 24), @sizeOf(Pg)); + try std.testing.expectEqual(@as(usize, 24), @sizeOf(PgMain)); + try std.testing.expectEqual(@as(usize, 40), @sizeOf(Pg)); try std.testing.expectEqual(@as(usize, 8), @offsetOf(Pg, "ps")); try std.testing.expectEqual(@as(usize, 16), @offsetOf(Pg, "ptok")); + try std.testing.expectEqual(@as(usize, 24), @offsetOf(Pg, "ips")); + try std.testing.expectEqual(@as(usize, 32), @offsetOf(Pg, "ipt")); } test { diff --git a/zig/src/families/flashnext/cuda_images.zig b/zig/src/families/flashnext/cuda_images.zig index 0e2997552d..a889b26471 100644 --- a/zig/src/families/flashnext/cuda_images.zig +++ b/zig/src/families/flashnext/cuda_images.zig @@ -22,6 +22,7 @@ pub const Vision = struct { /// prep and pool through a stream's page table (TF_FLASHNEXT_KV_PAGES). prep_paged: cuda.Function, pool_paged: cuda.Function, + planes: ?media.Planes = null, pub fn load(a: std.mem.Allocator, io: std.Io, d: *const cuda.Driver, s: cuda.Stream, dir: []const u8, limits: Limits, rank: usize, prep: cuda.Function, pool: cuda.Function, prep_paged: cuda.Function, pool_paged: cuda.Function) !*Vision { try shared.media.checkLimits(limits); const settings = (try shared.Settings.read(a, io, dir)) orelse return error.NoVisionTower; @@ -70,11 +71,11 @@ pub const Vision = struct { return .{ .pixels = p.pixels, .gh = p.gh, .gw = p.gw, .tokens = p.tokens, .hash = p.hash }; } pub fn state(v: *Vision, ids: []const u32, images: []const @import("lanes").stream.Image) !media.State { - return .{ .image = try shared.media.Media.init(v.a, v.encoder, ids, images, v.limits), .prep = v.prep, .pool = v.pool, .prep_paged = v.prep_paged, .pool_paged = v.pool_paged }; + return .{ .image = try shared.media.Media.init(v.a, v.encoder, ids, images, v.limits), .prep = v.prep, .pool = v.pool, .prep_paged = v.prep_paged, .pool_paged = v.pool_paged, .planes = v.planes }; } pub fn receive(v: *Vision, n: usize, runs: []const shared.media.Run) !media.State { const p = try shared.media.Plan.init(v.a, n, runs, v.limits); - return .{ .image = try shared.media.Media.allocate(p, v.encoder), .prep = v.prep, .pool = v.pool, .prep_paged = v.prep_paged, .pool_paged = v.pool_paged }; + return .{ .image = try shared.media.Media.allocate(p, v.encoder), .prep = v.prep, .pool = v.pool, .prep_paged = v.prep_paged, .pool_paged = v.pool_paged, .planes = v.planes }; } pub fn explain(err: anyerror) ?[]const u8 { return switch (err) { diff --git a/zig/src/families/flashnext/cuda_kv_lanes.zig b/zig/src/families/flashnext/cuda_kv_lanes.zig index 72b53ec61a..7277f9d176 100644 --- a/zig/src/families/flashnext/cuda_kv_lanes.zig +++ b/zig/src/families/flashnext/cuda_kv_lanes.zig @@ -7,6 +7,8 @@ const pages = @import("cuda_pages.zig"); const mtp = @import("cuda_mtp.zig"); const snap = @import("cuda_snap.zig"); const wt = @import("cuda_prompt_wt.zig"); +const host_tier = @import("exl3/kv/cuda_exl3_kvhost.zig"); +const kvt = @import("exl3/kv/cuda_exl3_kvtier.zig"); /// Rows past a request's prompt and reply its pages hold: one drafted window's. pub const window_rows: u32 = mtp.max_rows; @@ -50,7 +52,11 @@ pub const Head = struct { draft_n: usize, rows: usize }; pub fn layout(capacity: usize, tokens: usize, head: ?Head, extra: usize) kv.Layout { const layers = wt.attn_n + @intFromBool(head != null); const run = seqRunBytes() + if (head) |h| mtp.Draft.runBytes(h.draft_n, h.rows) else 0; - return .{ .fixed_bytes = mostPages(capacity, tokens) * 4 + extra, .token_bytes = pages.pageBytesOf(tokens, layers) / tokens, .page_tokens = @intCast(tokens), .snap_bytes = snap.bytes(0, head != null), .run_bytes = run }; + // a run is whole pages (a host tier: of the K/V plane, which alone takes fewer bytes) + const page = pages.pageBytesOf(tokens, layers); + const run_page = if (host_tier.tiered()) pages.kvPageBytes(tokens, layers) else page; + const run_pages = (run + run_page - 1) / run_page; + return .{ .fixed_bytes = mostPages(capacity, tokens) * 4 + extra, .token_bytes = page / tokens, .page_tokens = @intCast(tokens), .snap_bytes = snap.bytes(0, head != null), .run_bytes = run_pages * page }; } /// `l` with every stream's run `k` pages: the ranks of a pair agree the larger, so their pools open the same ids. @@ -115,7 +121,15 @@ pub const State = struct { }; /// One pool lane: its cache's views (its run's states, its stream's pages), that stream's lease and its run. -const Lane = struct { seq: prompt.Seq, lease: kv.Lease = .{}, used: bool = false, run: u32 = 0 }; +const Lane = struct { + seq: prompt.Seq, + tl: kv.tier.Lease = .{}, // the stream's reservation: its pages (and, on a host tier, frames) + used: bool = false, + run: u32 = 0, // the first page (or frame, on a host tier) of its run + fids: std.ArrayList(u32) = .empty, // on a host tier: the frame of each page's K/V rows (none: out of frames) + sel: ?*kvt.LaneSel = null, // on a host tier: its selection of blocks in slots + itable: cuda.DeviceBuffer = undefined, // on a host tier: the index plane's table, logical page ids +}; pub const Pool = struct { gpa: std.mem.Allocator, @@ -123,8 +137,14 @@ pub const Pool = struct { layout: kv.Layout, capacity: usize, head: ?Head, - pool: kv.Pool, + plain: kv.Pool, // the pool when there is no host tier + pool: *kv.Pool, // `plain`, or the host tier's logical pages + host: ?*host_tier.Host = null, + kern: ?kvt.Kernels = null, + slot_frames: usize = 0, + planes: bool = false, // a host tier: the arena holds the K/V plane (frames), then the index plane (logical pages) arena: cuda.DeviceBuffer, + ibase: u64, // where the index plane starts in the arena: after `count` pages of the K/V plane lanes: []Lane, /// `streams` lanes and `count` pages of layout `l`, made and touched now; a lane holds only its table idle. @@ -135,10 +155,37 @@ pub const Pool = struct { _ = try std.math.add(u64, l.run_bytes, page_bytes - 1); const p = try gpa.create(Pool); errdefer gpa.destroy(p); - var arena = try cuda.DeviceBuffer.alloc(e.driver, try std.math.mul(usize, count, page_bytes)); + const layers = wt.attn_n + @intFromBool(head != null); + const kv_page = pages.kvPageBytes(l.page_tokens, layers); + const ix_page = pages.indexPageBytes(l.page_tokens, layers); + // a host tier: `logical` pages in host RAM, `frames` of them on the device at once; else every page is a frame + const mib = try host_tier.hostMiB(); + const slot_frames = if (mib > 0) try kvt.slotFrames() else 0; + const per_stream = kvt.framesFor(l.runPages(), slot_frames); + const mine_pages = if (mib > 0) mib * (1 << 20) / kv_page else count; + const mine_frames = if (mib > 0) (try host_tier.framesCap()) orelse streams * per_stream else count; + // two ranks hold the same logical pages and frames: the smaller of each + const agreed = if (mib > 0) try host_tier.agreeShape(e, mine_pages, mine_frames) else kv.tier.Shape{ .host_pages = mine_pages, .frames = mine_frames }; + const logical = agreed.host_pages; + const frames = agreed.frames; + if (logical == 0 or frames == 0 or (mib > 0 and frames < per_stream)) return error.BadFlashNextKvHost; + var arena = try cuda.DeviceBuffer.alloc(e.driver, try std.math.add(usize, try std.math.mul(usize, frames, kv_page), try std.math.mul(usize, logical, ix_page))); errdefer arena.free(); - p.* = .{ .gpa = gpa, .e = e, .layout = l, .capacity = capacity, .head = head, .pool = try kv.Pool.init(gpa, count, page_bytes), .arena = arena, .lanes = &.{} }; - errdefer p.pool.deinit(); + p.* = .{ .gpa = gpa, .e = e, .layout = l, .capacity = capacity, .head = head, .plain = undefined, .pool = undefined, .arena = arena, .ibase = arena.ptr + frames * kv_page, .lanes = &.{} }; + p.planes = mib > 0; + if (mib > 0) { + p.host = try host_tier.Host.create(gpa, e.driver, logical, frames, kv_page); + p.kern = try kvt.Kernels.load(e.driver, slot_frames * 64); + p.slot_frames = slot_frames; + std.log.info("flashnext: host tier: {d} logical pages ({d} MiB pinned), {d} frames on the device", .{ logical, logical * kv_page >> 20, frames }); + p.pool = &p.host.?.tier.pool; + } else { + p.plain = try kv.Pool.init(gpa, count, page_bytes); + p.pool = &p.plain; + } + errdefer if (p.host) |h| h.destroy(gpa) else p.plain.deinit(); + errdefer if (p.kern) |*kn| kn.unload(); + if (p.host) |h| try h.lockstep(e, slot_frames); p.lanes = try gpa.alloc(Lane, streams); errdefer gpa.free(p.lanes); var made: usize = 0; @@ -146,10 +193,17 @@ pub const Pool = struct { for (p.lanes) |*x| { x.* = .{ .seq = try p.laneSeq() }; made += 1; + if (p.host != null) { + x.itable = try cuda.DeviceBuffer.alloc(e.driver, mostPages(capacity, l.page_tokens) * 4); + x.seq.paged.?.itable = x.itable.ptr; + x.sel = try kvt.LaneSel.init(gpa, e.driver, &p.kern.?, capacity + window_rows, mostPages(capacity, l.page_tokens), slot_frames, layers, p.own(x)); + x.seq.sel = x.sel; + } } errdefer cuda.completion.drain(e.driver, e.stream.*, "flashnext: failed pool initialization did not drain"); try arena.fill8(0, e.stream.handle); for (p.lanes) |x| try x.seq.paged.?.table.fill8(0, e.stream.handle); + if (p.host != null) for (p.lanes) |x| try x.itable.fill8(0, e.stream.handle); try e.stream.synchronize(); return p; } @@ -158,9 +212,9 @@ pub const Pool = struct { fn laneSeq(p: *Pool) !prompt.Seq { const gpa = p.gpa; const d = p.e.driver; - var table = try cuda.DeviceBuffer.alloc(d, mostPages(p.capacity, p.layout.page_tokens) * 4); + var table = try cuda.DeviceBuffer.alloc(d, (mostPages(p.capacity, p.layout.page_tokens) + p.slot_frames) * 4); errdefer table.free(); - const paged: pages.Paged = .{ .base = p.arena.ptr, .table = table, .ids = &.{}, .tokens = p.layout.page_tokens, .own_ids = false, .layers = wt.attn_n + @intFromBool(p.head != null) }; + const paged: pages.Paged = .{ .base = p.arena.ptr, .ibase = p.ibase, .planes = p.planes, .hbase = if (p.host) |h| (h.store.device() catch 0) else 0, .table = table, .ids = &.{}, .tokens = p.layout.page_tokens, .own_ids = false, .layers = wt.attn_n + @intFromBool(p.head != null) }; const conv = try gpa.alloc(cuda.DeviceBuffer, wt.linear_n); errdefer gpa.free(conv); const state = try gpa.alloc(cuda.DeviceBuffer, wt.linear_n); @@ -182,15 +236,16 @@ pub const Pool = struct { pub fn deinit(p: *Pool) void { cuda.completion.drain(p.e.driver, p.e.stream.*, "flashnext: the pool did not drain at close"); for (p.lanes) |*x| { - if (x.used) { - p.pool.releaseRun(x.run, p.layout.runPages()); - p.pool.finish(&x.lease); - } + if (x.used) p.giveBack(x); + if (x.sel) |sel| sel.deinit(); + if (p.host != null) x.itable.free(); + x.fids.deinit(p.gpa); x.seq.deinit(p.gpa); } p.gpa.free(p.lanes); p.arena.free(); - p.pool.deinit(); + if (p.kern) |*kn| kn.unload(); + if (p.host) |h| h.destroy(p.gpa) else p.plain.deinit(); p.gpa.destroy(p); } @@ -198,7 +253,12 @@ pub const Pool = struct { pub fn fits(p: *const Pool, cap: usize) bool { if (cap > p.capacity + window_rows) return false; const k = p.layout.runPages(); - for (p.lanes) |x| if (!x.used) return p.pool.fits(p.layout.pages(cap) + k) and p.pool.findRun(k) != null; + const n = p.layout.pages(cap); + if (p.host) |h| { + for (p.lanes) |x| if (!x.used) return h.tier.fits(n, n + k); + return false; + } + for (p.lanes) |x| if (!x.used) return p.pool.fits(n + k) and p.pool.findRun(k) != null; return false; } @@ -211,15 +271,29 @@ pub const Pool = struct { const k = p.layout.runPages(); const n = p.layout.pages(cap); if (n > mostPages(p.capacity, p.layout.page_tokens)) return error.PromptTooLong; - try p.pool.promise(&x.lease, n + k); - errdefer p.pool.finish(&x.lease); + var run: u32 = undefined; + var up: Upload = .{ .p = p, .x = x, .table = x.seq.paged.?.table }; + if (p.host) |h| { + // pages in host RAM and a frame each, the run in consecutive frames of its own + h.tier.reserve(&x.tl, n, kvt.framesFor(k, p.slot_frames)) catch return error.LaneMemory; + } else try p.pool.promise(&x.tl.pages, n + k); + errdefer if (p.host) |h| h.tier.finish(&x.tl) else p.pool.finish(&x.tl.pages); // the run first: pages fill from the bottom, runs from the top; a fragmented pool waits (LaneMemory) - const run = p.pool.claimRun(&x.lease, k) orelse return error.LaneMemory; - errdefer p.pool.releaseRun(run, k); - errdefer cuda.completion.drain(p.e.driver, p.e.stream.*, "flashnext: a failed lane open did not drain"); - var up: Upload = .{ .p = p, .table = x.seq.paged.?.table }; - try p.pool.cover(&x.lease, n, .{ .ctx = &up, .write = Upload.write }); - x.seq.paged.?.ids = x.lease.ids.items; + if (p.host) |h| { + run = h.tier.lendRun(&x.tl, k) catch return error.LaneMemory; + h.tier.lend(&x.tl, p.slot_frames) catch return error.LaneMemory; + } else run = p.pool.claimRun(&x.tl.pages, k) orelse return error.LaneMemory; + errdefer if (p.host == null) p.pool.releaseRun(run, k); + errdefer p.drainFailed(); + if (p.host) |h| { + x.fids.clearRetainingCapacity(); + try h.tier.cover(&x.tl, n, .{ .ctx = &up, .write = Upload.writeFrames }); + x.sel.?.npages = n; + x.sel.?.reset_pending = true; + try x.sel.?.nameSlots(p.e.stream.*); + } else try p.pool.cover(&x.tl.pages, n, .{ .ctx = &up, .write = Upload.write }); + x.seq.paged.?.ids = x.tl.pages.ids.items; + x.seq.paged.?.fids = x.fids.items; x.run = run; try p.mapStates(&x.seq); try p.e.clearSeq(&x.seq, cap); @@ -227,13 +301,24 @@ pub const Pool = struct { return x.seq; } + /// Bytes a run's pages hold: whole pages of the K/V plane. + fn runLen(p: *const Pool) usize { + return p.layout.runPages() * p.runPage(); + } + + /// Bytes of one page a run holds: a host tier's frames hold the K/V plane alone. + fn runPage(p: *const Pool) usize { + const layers = wt.attn_n + @intFromBool(p.head != null); + return if (p.host != null) pages.kvPageBytes(p.layout.page_tokens, layers) else pages.pageBytesOf(p.layout.page_tokens, layers); + } + fn runBase(p: *const Pool, x: *const Lane) u64 { - return p.arena.ptr + @as(u64, x.run) * p.layout.pageBytes(); + return p.arena.ptr + @as(u64, x.run) * p.runPage(); } /// The lane's states as views of its run, laid out afresh (a kept window swaps state buffers with the head's). fn mapStates(p: *Pool, seq: *prompt.Seq) !void { - if (p.layout.runPages() * p.layout.pageBytes() < seqRunBytes()) return error.FlashNextRunTooShort; + if (p.runLen() < seqRunBytes()) return error.FlashNextRunTooShort; const x = p.laneOf(seq).?; var at = p.runBase(x); for (seq.conv) |*b| { @@ -251,9 +336,11 @@ pub const Pool = struct { pub fn draft(p: *Pool, seq: *const prompt.Seq, cap: usize) !mtp.Draft { const h = p.head orelse return error.FlashNextRunTooShort; const x = p.laneOf(seq) orelse return error.UnknownStream; - const run_len = p.layout.runPages() * p.layout.pageBytes(); + const run_len = p.runLen(); if (run_len < seqRunBytes()) return error.FlashNextRunTooShort; - return mtp.Draft.initIn(p.gpa, p.e.driver, cap, h.draft_n, h.rows, &x.seq.paged.?, p.runBase(x) + seqRunBytes(), run_len - seqRunBytes()); + var d = try mtp.Draft.initIn(p.gpa, p.e.driver, cap, h.draft_n, h.rows, &x.seq.paged.?, p.runBase(x) + seqRunBytes(), run_len - seqRunBytes()); + d.sel = x.sel; + return d; } /// Whether `seq` is a lane's cache. @@ -272,21 +359,63 @@ pub const Pool = struct { const x = p.laneOf(seq) orelse return false; if (!x.used) return true; cuda.completion.drain(p.e.driver, p.e.stream.*, "flashnext: a closed stream did not drain"); - p.pool.releaseRun(x.run, p.layout.runPages()); - p.pool.finish(&x.lease); + p.giveBack(x); if (seq.media) |*m| m.deinit(); // the stream's image state, held by its copy seq.media = null; x.seq.paged.?.ids = &.{}; + x.seq.paged.?.fids = &.{}; x.seq.cap = 0; x.used = false; return true; } + /// What a lane's selection moves: its tables, lease, frames and cache views. + fn own(p: *Pool, x: *Lane) kvt.Own { + return .{ .table = x.seq.paged.?.table, .itable = x.itable, .host = p.host.?, .arena = p.arena, .tl = &x.tl, .fids = &x.fids, .kc = x.seq.kc, .vc = x.seq.vc }; + } + + fn drainFailed(p: *Pool) void { + cuda.completion.drain(p.e.driver, p.e.stream.*, "flashnext: a failed lane open did not drain"); + } + + /// A lane's pages, frames and run go back (a host tier: the pages the prompt store keeps are made canonical first). + fn giveBack(p: *Pool, x: *Lane) void { + if (p.host) |h| { + const t = &h.tier; + var flushed = false; + for (x.tl.pages.ids.items) |id| if (t.pool.refs[id] > 1 and t.dirty[id]) { + h.flush(p.arena, id, p.e.stream.*) catch std.debug.panic("flashnext: a kept page could not reach the host store", .{}); + flushed = true; + }; + if (flushed) p.e.stream.synchronize() catch std.debug.panic("flashnext: the host store flush did not finish", .{}); + for (x.tl.pages.ids.items) |id| if (t.pool.refs[id] > 1) h.done(id); + if (x.sel) |sel| if (sel.stats()) |st| { + const c = t.census(); + std.log.info("flashnext: host tier stream: {d} pages ({d} flushed), {d} lent frames, sparse blocks needed {d} fetched {d}{s}; frames {d} resident {d} free", .{ x.tl.pages.ids.items.len, sel.flushed, x.tl.lent.items.len, st.needed, st.fetched, if (st.overflow) " OVERFLOW" else "", c.resident, c.free }); + } else |_| {}; + t.finish(&x.tl); + x.fids.clearRetainingCapacity(); + } else { + p.pool.releaseRun(x.run, p.layout.runPages()); + p.pool.finish(&x.tl.pages); + } + } + /// A lane's new table entries uploaded to its device table. const Upload = struct { p: *Pool, + x: *Lane, table: cuda.DeviceBuffer, + /// A host tier: the index table gets the page ids; the pages take frames as the stream's window reaches them. + fn writeFrames(ctx: *anyopaque, ids: []const u32, from: usize) anyerror!void { + const u: *Upload = @ptrCast(@alignCast(ctx)); + for (ids[from..]) |_| try u.x.fids.append(u.p.gpa, std.math.maxInt(u32)); + errdefer cuda.completion.drain(u.p.e.driver, u.p.e.stream.*, "flashnext: a failed table upload did not drain"); + try u.x.itable.uploadAsync(from * 4, std.mem.sliceAsBytes(ids[from..]), u.p.e.stream.handle); + try u.p.e.stream.synchronize(); // the ids slice may move as the lease grows + } + fn write(ctx: *anyopaque, ids: []const u32, from: usize) anyerror!void { const u: *Upload = @ptrCast(@alignCast(ctx)); errdefer cuda.completion.drain(u.p.e.driver, u.p.e.stream.*, "flashnext: a failed table upload did not drain"); diff --git a/zig/src/families/flashnext/cuda_kv_lifetime_test.zig b/zig/src/families/flashnext/cuda_kv_lifetime_test.zig index 062cfbe4de..b45c626682 100644 --- a/zig/src/families/flashnext/cuda_kv_lifetime_test.zig +++ b/zig/src/families/flashnext/cuda_kv_lifetime_test.zig @@ -40,7 +40,7 @@ test "KV lifetime: Flash Next failed table wait and later fill retain pages unti cuda.completion.drain(&d, e.stream.*, "the test queue must drain before its owners leave"); p.deinit(); } - Watch.pool = &p.pool; + Watch.pool = p.pool; D.observe = Watch.held; D.expected = l.pages(1) + l.runPages(); D.proofs = 0; @@ -66,7 +66,7 @@ test "KV lifetime: Flash Next close and deinit retain pages until context proof" const l = pages.layout(64, 16, null, 0); const p = try pages.Pool.init(t.allocator, &e, l, 64, null, 1, l.runPages() + 16); var c = try p.open(1); - Watch.pool = &p.pool; + Watch.pool = p.pool; D.observe = Watch.held; D.expected = l.pages(1) + l.runPages(); D.proofs = 0; @@ -110,13 +110,13 @@ test "KV lanes: two ranks whose pages differ in bytes open and close the same pa s1[i] = try p1.open(cap); } for (0..3) |i| { - try t.expectEqualSlices(u32, p0.lanes[i].lease.ids.items, p1.lanes[i].lease.ids.items); + try t.expectEqualSlices(u32, p0.lanes[i].tl.pages.ids.items, p1.lanes[i].tl.pages.ids.items); try t.expectEqual(p0.lanes[i].run, p1.lanes[i].run); } try t.expect(p0.close(&s0[1]) and p1.close(&s1[1])); s0[1] = try p0.open(60); s1[1] = try p1.open(60); - try t.expectEqualSlices(u32, p0.lanes[1].lease.ids.items, p1.lanes[1].lease.ids.items); + try t.expectEqualSlices(u32, p0.lanes[1].tl.pages.ids.items, p1.lanes[1].tl.pages.ids.items); try t.expectEqual(p0.lanes[1].run, p1.lanes[1].run); for (0..3) |i| { try t.expect(p0.close(&s0[i]) and p1.close(&s1[i])); diff --git a/zig/src/families/flashnext/cuda_lanes_load.zig b/zig/src/families/flashnext/cuda_lanes_load.zig index 6d0a05c19a..b6e64da73a 100644 --- a/zig/src/families/flashnext/cuda_lanes_load.zig +++ b/zig/src/families/flashnext/cuda_lanes_load.zig @@ -66,6 +66,7 @@ pub fn open(comptime F: type, gpa: std.mem.Allocator, io: std.Io, ctx: *const cu if (o.vision) |limits| { if (own.set != null) return error.VisionCapturedSetUnsupported; own.vision = try images.Vision.load(gpa, io, driver, own.stream, dir, limits, o.rank, own.glue.f.prep_media, own.glue.f.pool_media, own.glue.f.prep_media_paged, own.glue.f.pool_media_paged); + own.vision.?.planes = .{ .prep = own.glue.f.prep_media_planes, .pool = own.glue.f.pool_media_planes }; own.engine.image_factory = own.vision; } errdefer if (own.vision) |v| v.deinit(); diff --git a/zig/src/families/flashnext/cuda_media.zig b/zig/src/families/flashnext/cuda_media.zig index d8eac19a5a..66cb872dbf 100644 --- a/zig/src/families/flashnext/cuda_media.zig +++ b/zig/src/families/flashnext/cuda_media.zig @@ -3,6 +3,9 @@ const cuda = @import("cuda"); const shared = @import("qwen_vision"); const wt = @import("cuda_prompt_wt.zig"); const Pg = @import("cuda_glue.zig").Pg; +/// prep and pool through a host tier's planes (the indexer keys by logical page). +pub const Planes = struct { prep: cuda.Function, pool: cuda.Function }; + pub const State = struct { image: shared.media.Media, prep: cuda.Function, @@ -10,8 +13,9 @@ pub const State = struct { /// prep and pool through a stream's page table (TF_FLASHNEXT_KV_PAGES). prep_paged: cuda.Function, pool_paged: cuda.Function, + planes: ?Planes = null, pub fn view(s: *const State) View { - return .{ .image = s.image.view(), .prep = s.prep, .pool = s.pool, .prep_paged = s.prep_paged, .pool_paged = s.pool_paged }; + return .{ .image = s.image.view(), .prep = s.prep, .pool = s.pool, .prep_paged = s.prep_paged, .pool_paged = s.pool_paged, .planes = s.planes }; } pub fn deinit(s: *State) void { s.image.deinit(); @@ -23,6 +27,7 @@ pub const View = struct { pool: cuda.Function, prep_paged: cuda.Function, pool_paged: cuda.Function, + planes: ?Planes = null, pub fn attn(v: View, stream: cuda.Stream, proj: u64, pos: u64, qw: u64, kw: u64, iw: u64, inv: u64, q: u64, kc: u64, vc: u64, iq: u64, ikc: u64, eps: f32, rows: usize, pg: Pg) !void { var a: cuda.Args = .{}; for ([_]u64{ proj, pos, qw, kw, iw, inv, q, kc, vc, iq, ikc }) |x| a.add(x); @@ -31,8 +36,8 @@ pub const View = struct { a.add(v.image.rope); a.add(@as(i32, @intCast(v.image.n))); a.add(v.image.delta); - if (pg.pt != 0) a.add(pg); - const f = if (pg.pt != 0) v.prep_paged else v.prep; + if (pg.pt != 0) if (pg.planes()) a.add(pg) else a.add(pg.main()); + const f = if (pg.pt != 0) (if (pg.planes()) (v.planes orelse return error.NoPlanesKernel).prep else v.prep_paged) else v.prep; try cuda.launch.launch(f, .{ .grid = .{ .x = @intCast(rows), .y = wt.q_heads + wt.kv_heads + wt.index_heads + 1, .z = 1 }, .block = .{ .x = 256, .y = 1, .z = 1 } }, stream, &a); } pub fn pooled(v: View, stream: cuda.Stream, ikc: u64, out: u64, pos: u64, w: u64, inv: u64, eps: f32, rows: usize, pg: Pg) !void { @@ -43,8 +48,8 @@ pub const View = struct { a.add(v.image.rope); a.add(@as(i32, @intCast(v.image.n))); a.add(v.image.delta); - if (pg.pt != 0) a.add(pg); - const f = if (pg.pt != 0) v.pool_paged else v.pool; + if (pg.pt != 0) if (pg.planes()) a.add(pg) else a.add(pg.main()); + const f = if (pg.pt != 0) (if (pg.planes()) (v.planes orelse return error.NoPlanesKernel).pool else v.pool_paged) else v.pool; try cuda.launch.launch(f, .{ .grid = .{ .x = @intCast(rows / wt.index_ratio + 2), .y = 1, .z = 1 }, .block = .{ .x = wt.index_dim, .y = 1, .z = 1 } }, stream, &a); } }; diff --git a/zig/src/families/flashnext/cuda_mtp_draft.zig b/zig/src/families/flashnext/cuda_mtp_draft.zig index e95ffb221b..770502fe64 100644 --- a/zig/src/families/flashnext/cuda_mtp_draft.zig +++ b/zig/src/families/flashnext/cuda_mtp_draft.zig @@ -45,6 +45,7 @@ pub const Chain = struct { /// A sequence's head cache and the buffers a window needs: its rows' scores and streams, and what it rolls back to. pub const Draft = struct { + sel: ?*@import("exl3/kv/cuda_exl3_kvtier.zig").LaneSel = null, // a host tier's selection of the stream's blocks stream: ?cuda.Stream = null, media: ?@import("cuda_media.zig").View = null, gpa: std.mem.Allocator, @@ -181,7 +182,7 @@ pub const Draft = struct { /// The head's one layer as a pass reads its cache: contiguous, or the lane's pages (`pg`) after the model's. pub fn headSeq(d: *Draft, media: ?@import("cuda_media.zig").View, scores: u64) batch.SeqState { - return .{ .media = media, .conv = &.{}, .state = &.{}, .kc = (&d.kc)[0..1], .vc = (&d.vc)[0..1], .ikc = (&d.ikc)[0..1], .pooled = (&d.pooled)[0..1], .pg = d.pg, .tail = d.kc, .history = &d.unused_history, .cap = d.cap, .scores = scores }; + return .{ .media = media, .conv = &.{}, .state = &.{}, .kc = (&d.kc)[0..1], .vc = (&d.vc)[0..1], .ikc = (&d.ikc)[0..1], .pooled = (&d.pooled)[0..1], .pg = d.pg, .tail = d.kc, .history = &d.unused_history, .cap = d.cap, .scores = scores, .sel = d.sel, .sel_layer = wt.attn_n }; } /// Each DeltaNet layer's window end state into its spare buffer. diff --git a/zig/src/families/flashnext/cuda_pages.zig b/zig/src/families/flashnext/cuda_pages.zig index 1e03ae2bf5..3f9f5d39c2 100644 --- a/zig/src/families/flashnext/cuda_pages.zig +++ b/zig/src/families/flashnext/cuda_pages.zig @@ -1,5 +1,4 @@ -//! Flash Next's attention caches as pages: per attention layer, keys, values, indexer keys, then pooled rows (bf16). - +//! Flash Next attention caches as pages: per layer keys, values, indexer keys, pooled rows (bf16); a tier: two planes. const std = @import("std"); const cuda = @import("cuda"); const wt = @import("cuda_prompt_wt.zig"); @@ -24,9 +23,19 @@ pub fn rows(part: Part, tokens: usize) usize { return if (part == .pooled) tokens / ratio else tokens; } -/// Bytes one attention layer's share of a page takes. +/// Bytes one attention layer's share of a page takes in the K/V plane. +pub fn kvLayerBytes(tokens: usize) usize { + return 2 * tokens * kv_row; +} + +/// Bytes one attention layer's share of a page takes in the index plane. +pub fn indexLayerBytes(tokens: usize) usize { + return tokens * index_row + tokens / ratio * index_row; +} + +/// Bytes one attention layer's share of a page takes, both planes. pub fn layerBytes(tokens: usize) usize { - return tokens * (2 * kv_row + index_row) + tokens / ratio * index_row; + return kvLayerBytes(tokens) + indexLayerBytes(tokens); } /// Bytes a page takes: every attention layer's share. @@ -34,12 +43,27 @@ pub fn pageBytes(tokens: usize) usize { return pageBytesOf(tokens, wt.attn_n); } -/// Bytes a page of `layers` attention layers takes (the model's, and the MTP head's cache when it drafts). +/// Bytes a page of `layers` attention layers takes (the model's, and the MTP head's cache when it drafts), both planes. pub fn pageBytesOf(tokens: usize, layers: usize) usize { return layers * layerBytes(tokens); } -/// Where `part` of attention layer `layer` starts inside a page. +/// Bytes a page of `layers` layers takes in the K/V plane. +pub fn kvPageBytes(tokens: usize, layers: usize) usize { + return layers * kvLayerBytes(tokens); +} + +/// Bytes a page of `layers` layers takes in the index plane. +pub fn indexPageBytes(tokens: usize, layers: usize) usize { + return layers * indexLayerBytes(tokens); +} + +/// Whether `part` is in the index plane. +pub fn inIndexPlane(part: Part) bool { + return part == .index or part == .pooled; +} + +/// Where `part` of layer `layer` starts in a one-plane page: per layer keys, values, indexer keys, pooled rows. pub fn offset(tokens: usize, layer: usize, part: Part) usize { const at = layer * layerBytes(tokens); return at + switch (part) { @@ -50,6 +74,16 @@ pub fn offset(tokens: usize, layer: usize, part: Part) usize { }; } +/// Where `part` of attention layer `layer` starts inside its plane's page when a host tier keeps two planes. +pub fn planeOffset(tokens: usize, layer: usize, part: Part) usize { + return switch (part) { + .key => layer * kvLayerBytes(tokens), + .value => layer * kvLayerBytes(tokens) + tokens * kv_row, + .index => layer * indexLayerBytes(tokens), + .pooled => layer * indexLayerBytes(tokens) + tokens * index_row, + }; +} + /// A page size the kernels take: whole blocks of the indexer's ratio. pub fn valid(tokens: usize) bool { return tokens > 0 and tokens <= wt.window and tokens <= std.math.maxInt(c_int) and tokens % ratio == 0; @@ -80,11 +114,21 @@ pub fn checkedLayout(cap: usize, tokens: usize) !Layout { /// One stream's pages: page i of its positions at `base` + ids[i] pages. pub const Paged = struct { - /// Where page 0 of the arena starts. + /// Where page 0 starts (a host tier: of the K/V plane). base: u64, + /// A host tier: where page 0 of the index plane starts. + ibase: u64 = 0, + /// A host tier keeps two planes: the K/V rows in frames, the indexer keys by page id. + planes: bool = false, /// The stream's table on the device, an int a page (`ids` uploaded). table: cuda.DeviceBuffer, ids: []u32, + /// The frame of each page's K/V rows when a host tier maps them (empty: the page id is the frame). + fids: []const u32 = &.{}, + /// A host tier's K/V store as the device addresses it (0: none): snapshots use its rows, not the frames'. + hbase: u64 = 0, + /// The index plane's table on the device, an int a page naming the logical page (0: the K/V table's). + itable: u64 = 0, /// Positions a page holds. tokens: usize, /// The paged state owns `ids` (the page test's); a pool's lane maps its lease's. @@ -96,18 +140,29 @@ pub const Paged = struct { stream: ?cuda.Stream = null, pub fn pg(p: *const Paged) Pg { - return .{ .pt = p.table.ptr, .ps = @intCast(pageBytesOf(p.tokens, p.layers) / 2), .ptok = @intCast(p.tokens) }; + if (!p.planes) return .{ .pt = p.table.ptr, .ps = @intCast(pageBytesOf(p.tokens, p.layers) / 2), .ptok = @intCast(p.tokens) }; + return .{ .pt = p.table.ptr, .ps = @intCast(kvPageBytes(p.tokens, p.layers) / 2), .ptok = @intCast(p.tokens), .ips = @intCast(indexPageBytes(p.tokens, p.layers) / 2), .ipt = p.itable }; } /// `part` of attention layer `layer` in page 0: the cache pointer a kernel takes beside pg(). pub fn at(p: *const Paged, layer: usize, part: Part) u64 { - return p.base + offset(p.tokens, layer, part); + if (!p.planes) return p.base + offset(p.tokens, layer, part); + return (if (inIndexPlane(part)) p.ibase else p.base) + planeOffset(p.tokens, layer, part); } /// Byte address of row `row` of `part` of `layer` (a position, or a block for pooled). pub fn rowAt(p: *const Paged, layer: usize, part: Part, row: usize) u64 { const per = rows(part, p.tokens); - return p.base + @as(u64, p.ids[row / per]) * pageBytesOf(p.tokens, p.layers) + offset(p.tokens, layer, part) + (row % per) * rowBytes(part); + if (!p.planes) return p.base + @as(u64, p.ids[row / per]) * pageBytesOf(p.tokens, p.layers) + offset(p.tokens, layer, part) + (row % per) * rowBytes(part); + const kv_id = if (p.fids.len > 0) p.fids[row / per] else p.ids[row / per]; + const plane = if (inIndexPlane(part)) p.ibase + @as(u64, p.ids[row / per]) * indexPageBytes(p.tokens, p.layers) else p.base + @as(u64, kv_id) * kvPageBytes(p.tokens, p.layers); + return plane + planeOffset(p.tokens, layer, part) + (row % per) * rowBytes(part); + } + + /// Byte address of row `row` of a K/V `part` in the host store (a host tier's canonical rows). + pub fn hostRowAt(p: *const Paged, layer: usize, part: Part, row: usize) u64 { + const per = rows(part, p.tokens); + return p.hbase + @as(u64, p.ids[row / per]) * kvPageBytes(p.tokens, p.layers) + planeOffset(p.tokens, layer, part) + (row % per) * rowBytes(part); } /// Borrowed views of each attention layer's `part` in page 0 (what kc, vc, ikc and pooled hold when paged). @@ -156,6 +211,26 @@ test "a page's parts tile it in order and pooled rows follow their positions' pa try std.testing.expectEqual(p.base + offset(t, 3, .value) + kv_row, p.rowAt(3, .value, 17)); try std.testing.expectEqual(p.base + offset(t, 3, .pooled) + index_row, p.rowAt(3, .pooled, 5)); try std.testing.expectEqual(p.base + 2 * pageBytes(t) + offset(t, 0, .index), p.rowAt(0, .index, 0)); + try std.testing.expectEqual(@as(i64, 0), p.pg().ips); +} + +test "with a host tier a page's parts tile its two planes; the K/V rows follow their frames, the rest their pages" { + const t = 16; + try std.testing.expectEqual(@as(usize, 0), planeOffset(t, 0, .key)); + try std.testing.expectEqual(t * kv_row, planeOffset(t, 0, .value)); + try std.testing.expectEqual(kvLayerBytes(t), planeOffset(t, 1, .key)); + try std.testing.expectEqual(@as(usize, 0), planeOffset(t, 0, .index)); + try std.testing.expectEqual(planeOffset(t, 0, .pooled) + t / ratio * index_row, planeOffset(t, 1, .index)); + try std.testing.expectEqual(kvPageBytes(t, wt.attn_n), planeOffset(t, wt.attn_n - 1, .value) + t * kv_row); + try std.testing.expectEqual(indexPageBytes(t, wt.attn_n), planeOffset(t, wt.attn_n - 1, .pooled) + t / ratio * index_row); + try std.testing.expectEqual(pageBytes(t), kvPageBytes(t, wt.attn_n) + indexPageBytes(t, wt.attn_n)); + var ids = [_]u32{ 2, 0, 1 }; + var frames = [_]u32{ 5, 4, 3 }; + const p: Paged = .{ .base = 1 << 20, .ibase = 1 << 30, .planes = true, .table = .{ .d = undefined, .ptr = 0, .len = 0, .borrowed = true }, .ids = &ids, .fids = &frames, .tokens = t }; + try std.testing.expectEqual(p.base + 4 * kvPageBytes(t, wt.attn_n) + planeOffset(t, 3, .value) + kv_row, p.rowAt(3, .value, 17)); + try std.testing.expectEqual(p.ibase + 0 * indexPageBytes(t, wt.attn_n) + planeOffset(t, 3, .pooled) + index_row, p.rowAt(3, .pooled, 5)); + try std.testing.expectEqual(p.ibase + 2 * indexPageBytes(t, wt.attn_n) + planeOffset(t, 0, .index), p.rowAt(0, .index, 0)); + try std.testing.expect(p.pg().ips != 0); } test { diff --git a/zig/src/families/flashnext/cuda_paging_batch_test.zig b/zig/src/families/flashnext/cuda_paging_batch_test.zig index 70ae7ad70b..b7284b7533 100644 --- a/zig/src/families/flashnext/cuda_paging_batch_test.zig +++ b/zig/src/families/flashnext/cuda_paging_batch_test.zig @@ -149,6 +149,7 @@ test "paged grouped prompts: every current attention route receives the sequence p.pos0 = 0; p.seq.cap = 4096; p.seq.media = null; + p.seq.sel = null; p.seq.pg = .{}; p.seq.kc = &buffers; p.seq.vc = &buffers; diff --git a/zig/src/families/flashnext/cuda_paging_launch_test.zig b/zig/src/families/flashnext/cuda_paging_launch_test.zig index f012392f6c..969421c734 100644 --- a/zig/src/families/flashnext/cuda_paging_launch_test.zig +++ b/zig/src/families/flashnext/cuda_paging_launch_test.zig @@ -6,7 +6,7 @@ const glue = @import("cuda_glue.zig"); const Capture = struct { var count: usize = 0; var symbols: [6]usize = @splat(0); - var pages: [6]glue.Pg = @splat(.{}); + var pages: [6]glue.PgMain = @splat(.{}); fn launch(f: cuda.abi.Function, _: u32, _: u32, _: u32, _: u32, _: u32, _: u32, _: u32, _: cuda.abi.Stream, args: ?[*]?*anyopaque, _: ?[*]?*anyopaque) callconv(.c) c_int { const id = @intFromPtr(f); symbols[count] = id; @@ -19,7 +19,7 @@ const Capture = struct { 12 => 15, else => unreachable, }; - pages[count] = @as(*const glue.Pg, @ptrCast(@alignCast(args.?[at].?))).*; + pages[count] = @as(*const glue.PgMain, @ptrCast(@alignCast(args.?[at].?))).*; } count += 1; return cuda.abi.success; @@ -60,7 +60,7 @@ test "paged launches: prep, pool, tiled QSA, ahead, one-buffer and fused select try std.testing.expectEqual(@as(usize, 6), Capture.count); for (Capture.symbols, Capture.pages, 0..) |id, actual, i| { try std.testing.expectEqual(2 * i + 1 + @intFromBool(pg.pt != 0), id); - try std.testing.expectEqual(pg, actual); + try std.testing.expectEqual(pg.main(), actual); } } } diff --git a/zig/src/families/flashnext/cuda_prompt.zig b/zig/src/families/flashnext/cuda_prompt.zig index 38a6ef7043..8067cc667b 100644 --- a/zig/src/families/flashnext/cuda_prompt.zig +++ b/zig/src/families/flashnext/cuda_prompt.zig @@ -507,6 +507,8 @@ pub const Seq = struct { pooled: []cuda.DeviceBuffer, /// kc, vc, ikc and pooled as borrowed views of page 0 of these pages (null: each its own contiguous buffer). paged: ?pages.Paged = null, + /// On a host tier: which blocks of the stream's K/V are in slots (null: every page is in a frame). + sel: ?*@import("exl3/kv/cuda_exl3_kvtier.zig").LaneSel = null, tail: cuda.DeviceBuffer, history: [2]i64, fed: usize, @@ -517,7 +519,7 @@ pub const Seq = struct { /// The caches a batched pass reads and leaves; `scores` the engine's block scores (Engine.reserveScores). pub fn view(s: *Seq, scores: u64) batch.SeqState { - return .{ .media = if (s.media) |*m| m.view() else null, .conv = s.conv, .state = s.state, .kc = s.kc, .vc = s.vc, .ikc = s.ikc, .pooled = s.pooled, .tail = s.tail, .history = &s.history, .cap = s.cap, .scores = scores, .pg = s.pg() }; + return .{ .media = if (s.media) |*m| m.view() else null, .conv = s.conv, .state = s.state, .kc = s.kc, .vc = s.vc, .ikc = s.ikc, .pooled = s.pooled, .tail = s.tail, .history = &s.history, .cap = s.cap, .scores = scores, .pg = s.pg(), .sel = s.sel }; } /// The attention kernels' page argument (zero: contiguous caches). diff --git a/zig/src/families/flashnext/cuda_snap.zig b/zig/src/families/flashnext/cuda_snap.zig index c8d03ccb1a..820376b392 100644 --- a/zig/src/families/flashnext/cuda_snap.zig +++ b/zig/src/families/flashnext/cuda_snap.zig @@ -88,7 +88,7 @@ fn move(stream: cuda.Stream, seq: *prompt.Seq, d: ?*Draft, base: u64, at: usize, var r: usize = 0; while (r < n) { const run = @min(per - r % per, n - r); - const at_row = pg.rowAt(at_layer, part, r); + const at_row = if (pg.hbase != 0 and (part == .key or part == .value)) pg.hostRowAt(at_layer, part, r) else pg.rowAt(at_layer, part, r); const dst, const src = if (way == .save) .{ base + off, at_row } else .{ at_row, base + off }; try dr.check(dr.api.cuMemcpyDtoDAsync_v2(dst, src, run * row, stream.handle), "cuMemcpyDtoDAsync"); off += run * row; @@ -118,6 +118,7 @@ pub fn save(gpa: std.mem.Allocator, driver: *const cuda.Driver, stream: cuda.Str const at = seq.fed; 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; + if (seq.sel) |t| try t.flushAll(stream, at); // a host tier: the window's rows reach the host store first const n = bytes(at, d != null); const snap = try allocate(gpa, driver, n, at, seq.history, d != null); errdefer discard(gpa, stream, snap); @@ -146,6 +147,7 @@ pub fn restore(stream: cuda.Stream, seq: *prompt.Seq, d: ?*Draft, snap: *const S // a state with the head's rows serves a stream without drafts too: those rows stay in the buffer try move(stream, seq, d, snap.buf.ptr, snap.at, .restore); try stream.synchronize(); + if (seq.sel) |t| t.restored(); seq.history = snap.history; seq.fed = snap.at; if (d) |h| { diff --git a/zig/src/families/flashnext/exl3/kv/cuda_exl3_kvhost.zig b/zig/src/families/flashnext/exl3/kv/cuda_exl3_kvhost.zig new file mode 100644 index 0000000000..c4c485d0d9 --- /dev/null +++ b/zig/src/families/flashnext/exl3/kv/cuda_exl3_kvhost.zig @@ -0,0 +1,91 @@ +//! The host tier of the Flash Next pool (TF_FLASHNEXT_KV_HOST=MiB): canonical K/V in pinned host RAM, cached in frames. +const std = @import("std"); +const cuda = @import("cuda"); +const kv = @import("core").kv_pool; +const pages = @import("../../cuda_pages.zig"); + +/// TF_FLASHNEXT_KV_HOST=n: n MiB of pinned host RAM hold the K/V plane of every logical page (unset or 0: no tier). +pub fn hostMiB() !usize { + const v = std.c.getenv("TF_FLASHNEXT_KV_HOST") orelse return 0; + return std.fmt.parseInt(usize, std.mem.span(v), 10) catch return error.BadFlashNextKvHost; +} + +/// Whether the host tier is on (the page layout then keeps two planes). +pub fn tiered() bool { + return (hostMiB() catch 0) > 0; +} + +/// TF_FLASHNEXT_KV_FRAMES=n: at most n device frames of the K/V plane (unset: the pages the plan gave). +pub fn framesCap() !?usize { + const v = std.c.getenv("TF_FLASHNEXT_KV_FRAMES") orelse return null; + return std.fmt.parseInt(usize, std.mem.span(v), 10) catch return error.BadFlashNextKvHost; +} + +/// The tier of one pool: logical pages, frames, the pinned host store they are canonical in. +pub const Host = struct { + tier: kv.tier.Tier, + store: cuda.HostBuffer, + kv_page: usize, // bytes of one page's K/V plane + + pub fn create(gpa: std.mem.Allocator, d: *const cuda.Driver, logical: usize, frames: usize, kv_page: usize) !*Host { + const h = try gpa.create(Host); + errdefer gpa.destroy(h); + var store = try cuda.HostBuffer.allocMapped(d, try std.math.mul(usize, logical, kv_page)); + errdefer store.free(); + h.* = .{ .tier = undefined, .store = store, .kv_page = kv_page }; + try h.tier.init(gpa, store.bytes, frames, kv_page); + return h; + } + + /// The ranks exchange the tier's shape, slot count and digest at open; one that differs refuses. + pub fn lockstep(h: *const Host, e: anytype, slots: usize) !void { + const c = e.collective orelse return; + const both = try exchange(e, c, kv.tier.Lockstep.vector(&h.tier, slots)); + kv.tier.Lockstep.agree(both[0], both[1]) catch return error.FlashNextTierRanksDisagree; + } + + pub fn destroy(h: *Host, gpa: std.mem.Allocator) void { + h.tier.deinit(); + h.store.free(); + gpa.destroy(h); + } + + /// Page `id`'s K/V plane from its frame (in `arena`) to the host store; sync the stream, then call `done`. + pub fn flush(h: *Host, arena: cuda.DeviceBuffer, id: u32, stream: cuda.Stream) !void { + const tk = h.tier.resident(id) orelse return; + try arena.downloadAsync(@as(usize, tk.frame) * h.kv_page, h.tier.host_page(id), stream.handle); + } + + pub fn done(h: *Host, id: u32) void { + h.tier.sync(id); + } +}; + +/// Each rank's `mine`, side by side, from one sum across the ranks. +fn exchange(e: anytype, c: anytype, mine: [kv.tier.Lockstep.len]f32) ![2][kv.tier.Lockstep.len]f32 { + const n = kv.tier.Lockstep.len; + const st = e.stream.*; + var b = try cuda.DeviceBuffer.alloc(e.driver, 256); + defer b.free(); + var v: [2][n]f32 = .{ @splat(0), @splat(0) }; + v[e.layers[0].exl3.?.tp.?.rank] = mine; + try st.synchronize(); + try b.upload(0, std.mem.asBytes(&v)); + try c.sum(st, b.ptr, 2 * n); + try c.finish(st); + try b.download(0, std.mem.asBytes(&v)); + return v; +} + +/// The pages and frames both ranks can hold: the smaller of each (rank 1 has no MTP head, so its pages are smaller). +pub fn agreeShape(e: anytype, logical: usize, frames: usize) !kv.tier.Shape { + const mine: kv.tier.Shape = .{ .host_pages = logical, .frames = frames }; + const c = e.collective orelse return mine; + var m: [kv.tier.Lockstep.len]f32 = @splat(0); + m[0] = @floatFromInt(logical); + m[1] = @floatFromInt(frames); + const both = try exchange(e, c, m); + const a: kv.tier.Shape = .{ .host_pages = @intFromFloat(both[0][0]), .frames = @intFromFloat(both[0][1]) }; + const b: kv.tier.Shape = .{ .host_pages = @intFromFloat(both[1][0]), .frames = @intFromFloat(both[1][1]) }; + return a.agree(b); +} diff --git a/zig/src/families/flashnext/exl3/kv/cuda_exl3_kvtier.zig b/zig/src/families/flashnext/exl3/kv/cuda_exl3_kvtier.zig new file mode 100644 index 0000000000..49e9a88e50 --- /dev/null +++ b/zig/src/families/flashnext/exl3/kv/cuda_exl3_kvtier.zig @@ -0,0 +1,314 @@ +//! A stream's K/V on the host tier: first pages and a window in frames, the rest canonical in host RAM, cached in slots +const std = @import("std"); +const cuda = @import("cuda"); +const kv = @import("core").kv_pool; +const wt = @import("../../cuda_prompt_wt.zig"); +const pages = @import("../../cuda_pages.zig"); +const host_tier = @import("cuda_exl3_kvhost.zig"); + +const ratio = wt.index_ratio; +const top = wt.index_top; +const idw = wt.index_idw; +const attn_n = wt.attn_n; +const ctl_n: usize = 8; +const none = std.math.maxInt(u32); + +/// Positions a page holds on a host tier: 64 four-row slots a lent frame and layer. +pub const page_rows: usize = 256; +/// Pages kept in frames from the first row: a dense row's keys (the first `dense_keys` positions) are all in them. +pub const head_pages: usize = (wt.dense_keys + page_rows - 1) / page_rows; +/// The most rows one attention call writes (a deep prefill's chunk and a drafted window). +pub const call_rows: usize = 4608; + +/// TF_FLASHNEXT_KV_SLOTS=n: n lent frames a stream, 64 four-row slots a frame for each attention layer (default 128). +pub fn slotFrames() !usize { + const v = std.c.getenv("TF_FLASHNEXT_KV_SLOTS") orelse return 128; + const n = std.fmt.parseInt(usize, std.mem.span(v), 10) catch return error.BadFlashNextKvHost; + if (n == 0) return error.BadFlashNextKvHost; + return n; +} + +/// Frames a stream may hold at once: the first pages, the window and a call's rows, the state run, the slot frames. +pub fn framesFor(run_pages: usize, slot_frames: usize) usize { + return head_pages + (call_rows + 2 * page_rows) / page_rows + run_pages + slot_frames; +} + +pub const Kernels = struct { + module: cuda.Module, + plan: cuda.Function, + fetch: cuda.Function, + remap: cuda.Function, + /// One call's lists (need, miss, fetch), shared by the layers and streams. + work: cuda.DeviceBuffer, + + pub fn load(d: *const cuda.Driver, nslots: usize) !Kernels { + if (!cuda.kernels.available) return error.BuiltWithoutKernels; + var k: Kernels = undefined; + k.module = try cuda.Module.load(d, cuda.kernels.flashnext_exl3_kvtier); + errdefer k.module.unload(); + k.plan = try k.module.function("kvt_plan"); + k.fetch = try k.module.function("kvt_fetch"); + k.remap = try k.module.function("kvt_remap"); + k.work = try cuda.DeviceBuffer.alloc(d, 4 * 4 * @max(nslots, 2048)); + return k; + } + + pub fn unload(k: *Kernels) void { + k.work.free(); + k.module.unload(); + } +}; + +/// What the pool lends a stream's selection: its tables, tier lease, frame list and cache views. +pub const Own = struct { + table: cuda.DeviceBuffer, + itable: cuda.DeviceBuffer, + host: *host_tier.Host, + arena: cuda.DeviceBuffer, + tl: *kv.tier.Lease, + fids: *std.ArrayList(u32), + kc: []cuda.DeviceBuffer, + vc: []cuda.DeviceBuffer, +}; + +pub const LaneSel = struct { + gpa: std.mem.Allocator, + d: *const cuda.Driver, + k: *const Kernels, + dev: cuda.DeviceBuffer, // blk2slot, mark [layers][nblocks]; ctl [layers][8]; slot2blk, stamp [layers][nslots] + stage: cuda.HostBuffer, // table entries on their way to the device + itable: u64, // the index plane's table: logical page ids + table: u64, + tbuf: cuda.DeviceBuffer, // the K/V table: a frame a page, then the lent slot frames + host: *host_tier.Host, + arena: cuda.DeviceBuffer, + tl: *kv.tier.Lease, + fids: *std.ArrayList(u32), + kc: []cuda.DeviceBuffer, // each layer's K and V in page 0 of the frames + vc: []cuda.DeviceBuffer, + nblocks: usize, + nslots: usize, + npages: usize, // logical pages of the stream: the lent entries follow them in the table + layers: usize, + flushed: usize = 0, // pages [0, flushed) are canonical in the host store and out of the window + win0: usize = 0, // the first row not canonical in the host store + done: [attn_n + 1]usize = @splat(0), // rows each layer (the MTP head last) has been given to write + head_on: bool = false, // the MTP head has written (or will write) rows of this stream + hw: usize = 0, // pages [0, hw) have been flushed once: their host rows are valid + reset_pending: bool = true, + + pub fn init(gpa: std.mem.Allocator, d: *const cuda.Driver, k: *const Kernels, cap: usize, npages: usize, slot_frames: usize, layers: usize, own: Own) !*LaneSel { + const s = try gpa.create(LaneSel); + errdefer gpa.destroy(s); + const nblocks = (cap + ratio - 1) / ratio; + const nslots = slot_frames * 64; + var dev = try cuda.DeviceBuffer.alloc(d, 4 * layers * (2 * nblocks + ctl_n + 2 * nslots)); + errdefer dev.free(); + var stage = try cuda.HostBuffer.alloc(d, 4 * (npages + slot_frames)); + errdefer stage.free(); + s.* = .{ .gpa = gpa, .d = d, .k = k, .dev = dev, .stage = stage, .itable = own.itable.ptr, .table = own.table.ptr, .tbuf = own.table, .host = own.host, .arena = own.arena, .tl = own.tl, .fids = own.fids, .kc = own.kc, .vc = own.vc, .nblocks = nblocks, .nslots = nslots, .npages = npages, .layers = layers }; + return s; + } + + pub fn deinit(s: *LaneSel) void { + s.stage.free(); + s.dev.free(); + s.gpa.destroy(s); + } + + /// Slots a layer's plan can take: its rows' blocks at most `top` each. + pub fn rowsPerPlan(s: *const LaneSel) usize { + return @max(1, s.nslots / top); + } + + const Tab = struct { blk2slot: u64, mark: u64, ctl: u64, slot2blk: u64, stamp: u64 }; + + fn tab(s: *const LaneSel, ai: usize) Tab { + const nb = s.nblocks * 4; + const marks = s.dev.ptr + s.layers * nb; + const ctls = marks + s.layers * nb; + const s2b = ctls + s.layers * ctl_n * 4; + return .{ .blk2slot = s.dev.ptr + ai * nb, .mark = marks + ai * nb, .ctl = ctls + ai * ctl_n * 4, .slot2blk = s2b + ai * s.nslots * 4, .stamp = s2b + (s.layers + ai) * s.nslots * 4 }; + } + + fn memset32(s: *const LaneSel, ptr: u64, value: u32, words: usize, stream: cuda.Stream) !void { + try s.d.check(s.d.api.cuMemsetD32Async(ptr, value, words, stream.handle), "cuMemsetD32Async"); + } + + /// No block is in a slot, every clock and counter at 0. + pub fn reset(s: *LaneSel, stream: cuda.Stream) !void { + const t = s.tab(0); + try s.memset32(t.blk2slot, 0xFFFF_FFFF, s.layers * s.nblocks, stream); + try s.memset32(t.mark, 0, s.layers * s.nblocks, stream); + try s.memset32(t.ctl, 0, s.layers * ctl_n, stream); + try s.memset32(t.slot2blk, 0xFFFF_FFFF, s.layers * s.nslots, stream); + try s.memset32(s.tab(0).stamp, 0, s.layers * s.nslots, stream); + s.reset_pending = false; + } + + fn setEntry(s: *LaneSel, stream: cuda.Stream, index: usize, frame: u32) !void { + const words = s.stage.slice(u32); + words[index] = frame; + try s.tbuf.uploadAsync(index * 4, std.mem.sliceAsBytes(words[index..][0..1]), stream.handle); + } + + /// The lent frames named at the end of the K/V table: the slots' rows are rows past the stream's. + pub fn nameSlots(s: *LaneSel, stream: cuda.Stream) !void { + for (s.tl.lent.items[s.tl.lent.items.len - (s.nslots / 64) ..], 0..) |f, j| try s.setEntry(stream, s.npages + j, f); + } + + /// Page `pidx` in a frame: pinned, its rows copied in from the host store when they were out. + fn ensure(s: *LaneSel, stream: cuda.Stream, pidx: usize) !void { + if (s.fids.items[pidx] != none) return; + const id = s.tl.pages.ids.items[pidx]; + const fill = try s.host.tier.pin(s.tl, id); + if (fill.miss and pidx < s.hw) { + const at = @as(usize, fill.frame) * s.host.kv_page; + try s.arena.uploadAsync(at, s.host.tier.host_page(id), stream.handle); + } + try s.host.tier.markDirty(s.tl, id); // rows of the window are written; a refilled page is flushed again, equal + s.fids.items[pidx] = fill.frame; + try s.setEntry(stream, pidx, fill.frame); + } + + /// Layer `ai` writes [pos0, pos0 + n): a page leaves the window once every layer wrote it; a rewind brings it back. + pub fn begin(s: *LaneSel, stream: cuda.Stream, ai: usize, pos0: usize, n: usize) !void { + if (ai >= s.layers) return; + if (s.reset_pending) { + // a new, resumed or restored stream: rows before pos0 are canonical in the host store + try s.reset(stream); + s.flushed = pos0 / page_rows; + s.win0 = s.flushed * page_rows; + s.hw = (pos0 + page_rows - 1) / page_rows; + s.done = @splat(pos0); + s.head_on = false; + } + const first = pos0 / page_rows; + if (first < s.flushed) { + // a resume, a rewind or the head priming: the window restarts at this page (brought in below) + s.flushed = first; + s.win0 = first * page_rows; + try s.reset(stream); // blocks at and past the rewind point hold rows that are about to change + } + s.done[ai] = @min(s.done[ai], pos0); + if (ai == attn_n) s.head_on = true; // from its first call the MTP head holds pages back until it has written them + var low = s.done[0]; + for (s.done[1..if (s.head_on) attn_n + 1 else attn_n]) |d| low = @min(low, d); + const target = low / page_rows; + var p = s.flushed; + while (p < target) : (p += 1) { + if (s.fids.items[p] == none) continue; + const id = s.tl.pages.ids.items[p]; + try s.host.flush(s.arena, id, stream); + s.host.done(id); + s.hw = @max(s.hw, p + 1); + if (p >= head_pages) { + try s.host.tier.unpin(s.tl, id); + s.fids.items[p] = none; + } + } + if (target > s.flushed) s.flushed = target; + s.win0 = s.flushed * page_rows; + s.done[ai] = pos0 + n; + const last = (pos0 + n + page_rows - 1) / page_rows; + if (last > s.npages) return error.PromptTooLong; + for (0..last) |q| if (q < head_pages or q >= s.flushed) try s.ensure(stream, q); + } + + /// The snapshot restored the stream's K/V rows into the host store: no page keeps a frame, the window starts over. + pub fn restored(s: *LaneSel) void { + const t = &s.host.tier; + for (s.tl.pages.ids.items, 0..) |id, p| { + if (p < s.fids.items.len and s.fids.items[p] != none) { + t.sync(id); + t.unpin(s.tl, id) catch {}; + s.fids.items[p] = none; + } + t.invalidate(id) catch {}; + } + s.reset_pending = true; + } + + /// Rows [fed / page_rows * page_rows, fed) and everything before them are canonical in the host store. + pub fn flushAll(s: *LaneSel, stream: cuda.Stream, fed: usize) !void { + const upto = (fed + page_rows - 1) / page_rows; + for (0..@min(upto, s.npages)) |p| { + if (s.fids.items[p] == none) continue; + const id = s.tl.pages.ids.items[p]; + try s.host.flush(s.arena, id, stream); + s.host.done(id); + } + } + + /// One plan: `m` rows' selections in layer `ai`: older blocks to slots (fetched if absent), positions to rows. + pub fn page(s: *LaneSel, stream: cuda.Stream, ai: usize, qids: u64, qnk: u64, qspr: u64, m: usize) !void { + const t = s.tab(ai); + const k = s.k; + const need = k.work.ptr; + const miss = need + 4 * @max(s.nslots, 2048); + const fetch = miss + 4 * @max(s.nslots, 2048); + var a: cuda.Args = .{}; + for ([_]u64{ qids, qnk, qspr }) |v| a.add(v); + for ([_]usize{ m, idw, top, ratio, s.win0 }) |v| a.add(@as(c_int, @intCast(v))); + a.add(s.tab(0).ctl); + a.add(t.blk2slot); + a.add(t.blk2slot); + a.add(@as(c_int, 0)); + for ([_]u64{ t.slot2blk, t.stamp, t.mark }) |v| a.add(v); + a.add(@as(c_int, @intCast(s.nslots))); + for ([_]u64{ t.ctl, need, miss, fetch }) |v| a.add(v); + try cuda.launch.launch(k.plan, .{ .grid = .{ .x = 1 }, .block = .{ .x = 1024 } }, stream, &a); + const hk = s.host.store.device() catch return error.NoMappedHostStore; + const layer_k = pages.planeOffset(page_rows, ai, .key); + const layer_v = pages.planeOffset(page_rows, ai, .value); + var f: cuda.Args = .{}; + for ([_]u64{ hk + layer_k, hk + layer_v, s.arena.ptr + layer_k, s.arena.ptr + layer_v, t.ctl, fetch, s.itable, s.table }) |v| f.add(v); + f.add(@as(c_int, @intCast(s.npages))); + f.add(@as(i64, @intCast(s.host.kv_page / 16))); + f.add(@as(c_int, @intCast(pages.rowBytes(.key) / 16))); + f.add(@as(c_int, ratio)); + try cuda.launch.launch(k.fetch, .{ .grid = .{ .x = 128 }, .block = .{ .x = 256 } }, stream, &f); + var r: cuda.Args = .{}; + for ([_]u64{ qids, qnk, qspr }) |v| r.add(v); + for ([_]usize{ idw, s.win0, s.npages * page_rows }) |v| r.add(@as(c_int, @intCast(v))); + r.add(t.blk2slot); + r.add(@as(c_int, ratio)); + r.add(t.ctl); + try cuda.launch.launch(k.remap, .{ .grid = .{ .x = @intCast(m) }, .block = .{ .x = 256 } }, stream, &r); + } + + /// The blocks the plans needed, and those fetched, since open (read back: a sync). + pub fn stats(s: *const LaneSel) !struct { needed: u64, fetched: u64, overflow: bool } { + var needed: u64 = 0; + var fetched: u64 = 0; + var overflow = false; + for (0..s.layers) |ai| { + var ctl: [ctl_n]i32 = undefined; + try s.d.check(s.d.api.cuMemcpyDtoH_v2(&ctl, s.tab(ai).ctl, ctl_n * 4), "cuMemcpyDtoH"); + needed += @as(u32, @bitCast(ctl[4])); + fetched += @as(u32, @bitCast(ctl[3])); + overflow = overflow or ctl[2] != 0; + } + return .{ .needed = needed, .fetched = fetched, .overflow = overflow }; + } +}; + +/// One block of attention rows: where its queries and outputs are, its rows' selection in the batch's lists. +pub const Block = struct { q: u64, ao: u64, kc: u64, vc: u64, pos: u64, sel: usize, m: usize, ends: usize, first_row: usize }; + +/// Attention for `a.m` rows: plan, fetch, remap, chunks, merge; sparse rows a few at a time, within the layer's slots. +pub fn attend(s: *LaneSel, b: anytype, ai: usize, a: Block) !void { + const st = b.s; + const per = if (a.first_row / ratio > top) s.rowsPerPlan() else a.m; + var first: usize = 0; + while (first < a.m) : (first += per) { + const cnt = @min(per, a.m - first); + const ids = st.qids + (a.sel + first) * idw * 4; + const nk = st.qnk + (a.sel + first) * 4; + const spr = st.qspr + (a.sel + first) * 4; + if (a.ends / ratio > top) try s.page(b.stream, ai, ids, nk, spr, cnt); + const chunks = (@min(a.ends, wt.dense_keys) + 511) / 512; + try b.tri.attnChunks(a.q + first * wt.q_heads * wt.head_dim * 2, a.kc, a.vc, a.pos, st.po, st.pm, st.pl, ids, nk, spr, cnt, wt.kv_heads, chunks, true, b.seq.pg); + try b.tri.attnMerge(st.po, st.pm, st.pl, a.pos, a.ao + first * wt.out_k * 2, nk, spr, cnt, wt.kv_heads, true); + } +} diff --git a/zig/src/families/flashnext/exl3/model/cuda_exl3.zig b/zig/src/families/flashnext/exl3/model/cuda_exl3.zig index 26e4251248..fb49c3b268 100644 --- a/zig/src/families/flashnext/exl3/model/cuda_exl3.zig +++ b/zig/src/families/flashnext/exl3/model/cuda_exl3.zig @@ -7,6 +7,9 @@ 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 kvhost = @import("../kv/cuda_exl3_kvhost.zig"); +pub const kvtier = @import("../kv/cuda_exl3_kvtier.zig"); +pub const pages = @import("../../cuda_pages.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"); diff --git a/zig/tests/cuda/contract_fn_pages.zig b/zig/tests/cuda/contract_fn_pages.zig index c0e10e2f06..bf2beaeeb1 100644 --- a/zig/tests/cuda/contract_fn_pages.zig +++ b/zig/tests/cuda/contract_fn_pages.zig @@ -28,8 +28,15 @@ const row_kv = kv_heads * head_dim; const slot_elems = 2 * row_kv + index_dim + index_dim / ratio; /// Two streams' cached positions, the rows each adds, positions a page, the slot the layer takes and slots a page. -pub const Case = struct { lens: [2]usize, rows: usize, ptok: usize, slot: usize, slots: usize, image: ?media.Spec = null }; -pub const cases = [_]Case{ +pub const Case = struct { lens: [2]usize, rows: usize, ptok: usize, slot: usize, slots: usize, image: ?media.Spec = null, planes: bool = false }; +/// Every case twice: through one page table, and through the host tier's two planes (ips, ipt: the PgP launches). +pub const cases = one_plane ++ two_planes(); +fn two_planes() [one_plane.len]Case { + var out = one_plane; + for (&out) |*c| c.planes = true; + return out; +} +const one_plane = [_]Case{ .{ .lens = .{ 0, 3 }, .rows = 1, .ptok = 16, .slot = 1, .slots = 12 }, // a stream's first key .{ .lens = .{ 127, 128 }, .rows = 16, .ptok = 32, .slot = 1, .slots = 12 }, // a verify window over page edges .{ .lens = .{ 700, 33 }, .rows = 40, .ptok = 48, .slot = 0, .slots = 12 }, // a page no 512-key chunk divides @@ -319,7 +326,14 @@ pub fn launch(gpu: check.Gpu, s: cuda.Stream, g: *const flash.Glue, select: cuda .pl = try rig.scratch(att_rows * nch * q_heads * 4), .ao = try rig.output(n * out_k * 2), }; - if (given orelse placed) |t| pages[st] = .{ .pt = try rig.input(std.mem.sliceAsBytes(t[st])), .ps = @intCast(f.pageElems()), .ptok = @intCast(f.c.ptok) }; + if (given orelse placed) |t| { + const table = try rig.input(std.mem.sliceAsBytes(t[st])); + pages[st] = .{ .pt = table, .ps = @intCast(f.pageElems()), .ptok = @intCast(f.c.ptok) }; + if (f.c.planes) { + pages[st].ips = @intCast(f.pageElems()); + pages[st].ipt = table; + } + } } var images: [2]?media.View = .{ null, null }; if (f.image) |im| for (0..2) |st| { @@ -465,7 +479,7 @@ test "paged media contract: prompt, verify and sparse rows cover image spans, pa for (spans) |span| crossing = crossing or span.at / c.ptok != (span.at + span.tokens() - 1) / c.ptok; } }; - try std.testing.expectEqual(@as(usize, 6), count); + try std.testing.expectEqual(@as(usize, 12), count); // each of the six media cases through one plane and through two try std.testing.expect(crossing and suffix and sparse and fused and decode and partial); for (sizes) |present| try std.testing.expect(present); } diff --git a/zig/tests/cuda/contract_fn_pages_media.zig b/zig/tests/cuda/contract_fn_pages_media.zig index 306fa7c6be..38ef156476 100644 --- a/zig/tests/cuda/contract_fn_pages_media.zig +++ b/zig/tests/cuda/contract_fn_pages_media.zig @@ -33,6 +33,7 @@ pub fn at(d: *const cuda.Driver, f: anytype, rope: u64, p: Plan) View { .pool = f.pool_media, .prep_paged = f.prep_media_paged, .pool_paged = f.pool_media_paged, + .planes = .{ .prep = f.prep_media_planes, .pool = f.pool_media_planes }, .image = .{ .features = 0, .src = 0, .rope = rope, .n = p.src.len, .delta = p.delta, .width = 5120, .splice_fn = .{ .d = d, .handle = null } }, }; } @@ -61,29 +62,32 @@ test "paged media contract: real plans retain image axes, text suffixes and a ne } const Capture = struct { + const MainPages = @TypeOf((@as(flash.Pg, .{})).main()); var calls: usize = 0; var symbols: [2]usize = undefined; - var pages: [2]flash.Pg = undefined; + var pages: [2]MainPages = undefined; + var plane_pages: [2]flash.Pg = undefined; var configs: [2][7]u32 = undefined; var rope: [2]u64 = undefined; var n: [2]i32 = undefined; var delta: [2]i32 = undefined; fn launch(f: cuda.abi.Function, gx: u32, gy: u32, gz: u32, bx: u32, by: u32, bz: u32, sm: u32, _: cuda.abi.Stream, args: ?[*]?*anyopaque, _: ?[*]?*anyopaque) callconv(.c) c_int { - const prep = @intFromPtr(f) == 3; + const symbol = @intFromPtr(f); + const prep = symbol == 3 or symbol == 5; const i: usize = if (prep) 18 else 10; symbols[calls] = @intFromPtr(f); rope[calls] = @as(*const u64, @ptrCast(@alignCast(args.?[i].?))).*; n[calls] = @as(*const i32, @ptrCast(@alignCast(args.?[i + 1].?))).*; delta[calls] = @as(*const i32, @ptrCast(@alignCast(args.?[i + 2].?))).*; - pages[calls] = @as(*const flash.Pg, @ptrCast(@alignCast(args.?[i + 3].?))).*; + if (symbol >= 5) plane_pages[calls] = @as(*const flash.Pg, @ptrCast(@alignCast(args.?[i + 3].?))).* else pages[calls] = @as(*const MainPages, @ptrCast(@alignCast(args.?[i + 3].?))).*; configs[calls] = .{ gx, gy, gz, bx, by, bz, sm }; calls += 1; return 0; } }; -test "paged media contract: production wrappers select both media symbols with full Pages and image arguments" { +test "paged media contract: production wrappers select each page ABI and its media symbols" { var d: cuda.Driver = undefined; d.api.cuLaunchKernel = Capture.launch; const s: cuda.Stream = .{ .d = &d, .handle = null }; @@ -92,23 +96,28 @@ test "paged media contract: production wrappers select both media symbols with f .pool_media = cuda.Function{ .d = &d, .handle = @ptrFromInt(2) }, .prep_media_paged = cuda.Function{ .d = &d, .handle = @ptrFromInt(3) }, .pool_media_paged = cuda.Function{ .d = &d, .handle = @ptrFromInt(4) }, + .prep_media_planes = cuda.Function{ .d = &d, .handle = @ptrFromInt(5) }, + .pool_media_planes = cuda.Function{ .d = &d, .handle = @ptrFromInt(6) }, }; var plans = try specs[1].plans(std.testing.allocator); defer for (&plans) |*p| p.deinit(); - const pg: flash.Pg = .{ .pt = 0x7000, .ps = 13 * 64 * 1184, .ptok = 64 }; - for (plans) |p| { + const main: flash.Pg = .{ .pt = 0x7000, .ps = 13 * 64 * 1184, .ptok = 64 }; + var planes = main; + planes.ips = 13 * 64 * 160; + planes.ipt = 0x9000; + for ([_]flash.Pg{ main, planes }) |pg| for (plans) |p| { Capture.calls = 0; const v = at(&d, f, 0x8000, p); try v.attn(s, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 1e-6, 16, pg); try v.pooled(s, 1, 2, 3, 4, 5, 1e-6, 16, pg); - try std.testing.expectEqualSlices(usize, &.{ 3, 4 }, &Capture.symbols); + try std.testing.expectEqualSlices(usize, if (pg.planes()) &.{ 5, 6 } else &.{ 3, 4 }, &Capture.symbols); for (0..2) |i| { - try std.testing.expectEqual(pg, Capture.pages[i]); + if (pg.planes()) try std.testing.expectEqual(pg, Capture.plane_pages[i]) else try std.testing.expectEqual(pg.main(), Capture.pages[i]); try std.testing.expectEqual(@as(u64, 0x8000), Capture.rope[i]); try std.testing.expectEqual(@as(i32, @intCast(p.src.len)), Capture.n[i]); try std.testing.expectEqual(p.delta, Capture.delta[i]); } try std.testing.expectEqual([7]u32{ 16, 31, 1, 256, 1, 1, 0 }, Capture.configs[0]); try std.testing.expectEqual([7]u32{ 6, 1, 1, 128, 1, 1, 0 }, Capture.configs[1]); - } + }; } diff --git a/zig/tests/cuda/contract_fn_pages_sums.zig b/zig/tests/cuda/contract_fn_pages_sums.zig index b5fbed9fc8..ad2e7196e2 100644 --- a/zig/tests/cuda/contract_fn_pages_sums.zig +++ b/zig/tests/cuda/contract_fn_pages_sums.zig @@ -22,7 +22,7 @@ const keys = 2048; const blocks = 600; const eps: f32 = 1e-6; -pub const Mode = enum { prep, prep_media, pool, pool_media, scores, ahead, one, fused }; +pub const Mode = enum { prep, prep_media, pool, pool_media, scores, ahead, one, fused, prep_planes, prep_media_planes, pool_planes, pool_media_planes, scores_planes }; const modes = std.enums.values(Mode); const vt = g.hooks(Runner); const rms_why = "BF16 RMS output rounding can hide a BF16 sum error; FP16 overflow is proved; BF16 margins are printed"; @@ -31,10 +31,29 @@ fn depth(s: c.Shape) usize { return s.k + 32; } -pub fn spec(mode: Mode) c.Spec { +/// The one-plane mode a planes mode checks the same sums as. +fn baseOf(mode: Mode) Mode { + return switch (mode) { + .prep_planes => .prep, + .prep_media_planes => .prep_media, + .pool_planes => .pool, + .pool_media_planes => .pool_media, + .scores_planes => .scores, + else => mode, + }; +} + +pub fn spec(full: Mode) c.Spec { + const mode = baseOf(full); const norm = mode == .prep or mode == .prep_media or mode == .pool or mode == .pool_media; return .{ - .name = switch (mode) { + .name = if (mode != full) switch (full) { + .prep_planes => "cuda planes fn_attn_prep FP32 norms", + .prep_media_planes => "cuda planes fn_attn_prep_media FP32 norms", + .pool_planes => "cuda planes fn_qsa_pool FP32 pool and norms", + .pool_media_planes => "cuda planes fn_qsa_pool_media FP32 pool and norms", + else => "cuda planes fn_qsa_scores FP32 dots", + } else switch (mode) { .prep => "cuda paged fn_attn_prep FP32 norms", .prep_media => "cuda paged fn_attn_prep_media FP32 norms", .pool => "cuda paged fn_qsa_pool FP32 pool and norms", @@ -43,6 +62,7 @@ pub fn spec(mode: Mode) c.Spec { .ahead => "cuda paged fn_attn_chunks_ahead FP32 attention", .one => "cuda paged fn_attn_chunks_one FP32 attention", .fused => "cuda paged fn_attn_fused FP32 attention", + .prep_planes, .prep_media_planes, .pool_planes, .pool_media_planes, .scores_planes => unreachable, }, .accum = .f32, .depth = depth, @@ -103,6 +123,15 @@ const Runner = struct { terms: []f64 = &.{}, dev: [16]u64 = @splat(0), pg: flash.Pg = .{}, + itab: []u32 = &.{}, // a planes mode: the index plane's table, the page ids in order + + fn base(self: *const Runner) Mode { + return baseOf(self.mode); + } + + fn planes(self: *const Runner) bool { + return self.base() != self.mode; + } fn deinit(self: *Runner) void { cuda.completion.drain(self.rig.gpu.d, self.rig.s, "paged sum checks did not drain before buffer release"); @@ -121,9 +150,9 @@ const Runner = struct { const a = self.arena.?.allocator(); self.rows = shape.rows; self.dev = @splat(0); - const pos: i32 = if (self.mode == .scores) 4 * blocks - 1 else if (self.isAttention()) 4095 else 0; + const pos: i32 = if (self.base() == .scores) 4 * blocks - 1 else if (self.isAttention()) 4095 else 0; self.dev[0] = try self.rig.input(std.mem.asBytes(&pos)); - const cap = if (self.mode == .scores) 4 * blocks else if (self.isAttention()) 4096 + self.rows else ptok; + const cap = if (self.base() == .scores) 4 * blocks else if (self.isAttention()) 4096 + self.rows else ptok; const np = (cap + ptok - 1) / ptok; const table = try a.alloc(u32, np); for (table, 0..) |*id, i| id.* = @intCast(np - 1 - i); @@ -131,6 +160,13 @@ const Runner = struct { self.kv = try a.alloc(u16, np * page_elems); @memset(self.kv, 0); self.pg = .{ .pt = try self.rig.input(std.mem.sliceAsBytes(table)), .ps = @intCast(page_elems), .ptok = ptok }; + if (self.planes()) { + // the indexer keys by logical page, the K/V rows by (reversed) frame: two tables over the one buffer + self.itab = try a.alloc(u32, np); + for (self.itab, 0..) |*id, i| id.* = @intCast(i); + self.pg.ips = @intCast(page_elems); + self.pg.ipt = try self.rig.input(std.mem.sliceAsBytes(self.itab)); + } if (self.isAttention()) { self.q = try a.alloc(u16, self.rows * qh * hd); @memset(self.q, 0); @@ -154,7 +190,7 @@ const Runner = struct { self.dev[6] = try self.rig.scratch(self.rows * 5 * qh * 4); self.dev[7] = try self.rig.scratch(self.rows * 5 * qh * 4); self.dev[8] = try self.rig.output(self.count() * 2); - } else if (self.mode == .scores) { + } else if (self.base() == .scores) { self.q = try a.alloc(u16, self.rows * ih * di); for (self.q, 0..) |*v, i| v.* = c.toBf16(if (i % di < di / 2) 127 else 126); self.terms = try a.alloc(f64, di); @@ -171,7 +207,7 @@ const Runner = struct { } const Part = enum { key, value, index, pooled }; - fn pageAt(_: *const Runner, table: []const u32, row: usize, part: Part) usize { + fn pageAt(self: *const Runner, table: []const u32, row: usize, part: Part) usize { const stride = ptok * (2 * kh * hd + di) + ptok / 4 * di; const per: usize = if (part == .pooled) ptok / 4 else ptok; const row_len: usize = if (part == .key or part == .value) kh * hd else di; @@ -181,11 +217,12 @@ const Runner = struct { .index => 2 * ptok * kh * hd, .pooled => ptok * (2 * kh * hd + di), }; - return table[row / per] * stride + off + row % per * row_len; + const id = if (self.planes() and (part == .index or part == .pooled)) self.itab[row / per] else table[row / per]; + return id * stride + off + row % per * row_len; } fn isAttention(self: *const Runner) bool { - return self.mode == .ahead or self.mode == .one or self.mode == .fused; + return self.base() == .ahead or self.base() == .one or self.base() == .fused; } fn prepareNorm(self: *Runner, a: std.mem.Allocator, table: []const u32) !void { @@ -202,7 +239,7 @@ const Runner = struct { for (0..di) |d| proj[row * pw + qh * 2 * hd + 2 * kh * hd + h * di + d] = c.toBf16(@floatCast(normValue(d))); } } - if (self.mode == .pool or self.mode == .pool_media) { + if (self.base() == .pool or self.base() == .pool_media) { for (0..4) |j| { for (0..di) |d| self.kv[self.pageAt(table, j, .index) + d] = c.toBf16(@floatCast(normValue(d))); } @@ -215,7 +252,7 @@ const Runner = struct { self.dev[3] = try self.rig.input(std.mem.asBytes(&inv)); const rope: [12]i32 = @splat(0); self.dev[4] = try self.rig.input(std.mem.asBytes(&rope)); - if (self.mode == .prep or self.mode == .prep_media) { + if (self.base() == .prep or self.base() == .prep_media) { self.dev[8] = try self.rig.output(self.rows * qh * hd * 2); self.dev[10] = try self.rig.output(self.rows * ih * di * 2); } else { @@ -233,10 +270,10 @@ const Runner = struct { const vc = kc + ptok * kh * hd * 2; const ikc = kc + 2 * ptok * kh * hd * 2; const pooled = ikc + ptok * di * 2; - if (self.mode == .scores) { + if (self.base() == .scores) { try self.glue.qsaScores(self.rig.s, p[1], pooled, p[0], p[8], blocks, self.rows, blocks, self.pg); } else if (self.isAttention()) { - if (self.mode == .fused) { + if (self.base() == .fused) { try self.glue.attnFused(self.rig.s, p[1], kc, vc, p[0], p[8], p[2], p[3], p[4], self.rows, kh, self.pg); } else { var args: cuda.Args = .{}; @@ -245,16 +282,16 @@ const Runner = struct { args.add(@as(f32, 0.0625)); args.add(@as(c_int, idw)); args.add(@as(c_int, 1)); - args.add(self.pg); - const ahead = self.mode == .ahead; + args.add(self.pg.main()); + const ahead = self.base() == .ahead; const shared: usize = (16 + (if (ahead) @as(usize, 128) else 32)) * (hd + 8) * 2 + 16 * 40 * 2; try self.launch(if (ahead) self.glue.f.chunks_ahead_paged else self.glue.f.chunks_one_paged, .{ self.rows, kh, 4 }, 128, shared, &args); try self.glue.attnMerge(self.rig.s, p[5], p[6], p[7], p[0], p[8], p[3], p[4], self.rows, kh, true); } } else { var args: cuda.Args = .{}; - const prep = self.mode == .prep or self.mode == .prep_media; - const media = self.mode == .prep_media or self.mode == .pool_media; + const prep = self.base() == .prep or self.base() == .prep_media; + const media = self.base() == .prep_media or self.base() == .pool_media; if (prep) { for ([_]u64{ p[1], p[0], p[2], p[2], p[2], p[3], p[8], kc, vc, p[10], ikc }) |v| args.add(v); args.add(eps); @@ -269,8 +306,11 @@ const Runner = struct { args.add(@as(c_int, 4)); args.add(@as(c_int, -1)); } - args.add(self.pg); - const f = if (prep) (if (media) self.glue.f.prep_media_paged else self.glue.f.prep_paged) else (if (media) self.glue.f.pool_media_paged else self.glue.f.pool_paged); + if (self.planes()) args.add(self.pg) else args.add(self.pg.main()); + const f = if (self.planes()) + (if (prep) (if (media) self.glue.f.prep_media_planes else self.glue.f.prep_planes) else (if (media) self.glue.f.pool_media_planes else self.glue.f.pool_planes)) + else + (if (prep) (if (media) self.glue.f.prep_media_paged else self.glue.f.prep_paged) else (if (media) self.glue.f.pool_media_paged else self.glue.f.pool_paged)); try self.launch(f, if (prep) .{ self.rows, qh + kh + ih + 1, 1 } else .{ 1, 1, 1 }, if (prep) hd else di, 0, &args); if (!prep) { const np = self.kv.len / @as(usize, @intCast(self.pg.ps)); @@ -282,7 +322,7 @@ const Runner = struct { } pub fn count(self: *const Runner) usize { - return if (self.isAttention()) self.rows * qh * hd else if (self.mode == .scores) self.rows * blocks else if (self.mode == .pool or self.mode == .pool_media) di else self.rows * (qh * hd + ih * di); + return if (self.isAttention()) self.rows * qh * hd else if (self.base() == .scores) self.rows * blocks else if (self.base() == .pool or self.base() == .pool_media) di else self.rows * (qh * hd + ih * di); } pub fn load(self: *Runner, i: usize, head: *attn.Head) !void { @@ -301,7 +341,7 @@ const Runner = struct { pub fn reference(self: *Runner, want: []f64, mag: []f64) !void { if (self.isAttention()) return self.r.?.reference(self, self.rows * qh, 0.0625, want, mag); - if (self.mode == .scores) { + if (self.base() == .scores) { var dot: f64 = 0; var abs: f64 = 0; for (self.terms) |v| { @@ -312,7 +352,7 @@ const Runner = struct { @memset(mag, abs * ih / @sqrt(@as(f64, di))); return; } - const prep = self.mode == .prep or self.mode == .prep_media; + const prep = self.base() == .prep or self.base() == .prep_media; for (want, mag, 0..) |*w, *m, i| { const width: usize = if (!prep or i >= self.rows * qh * hd) di else hd; w.* = normReference(width, i % width); @@ -322,7 +362,7 @@ const Runner = struct { pub fn simulate(self: *Runner, acc: c.Accum, order: g.Order, out: []f64) !void { if (self.isAttention()) return self.r.?.simulate(self, self.rows * qh, 0.0625, acc, order, out); - if (self.mode == .scores) { + if (self.base() == .scores) { const dot = sums.sumIn(acc, order, self.terms); const terms: [ih]f64 = @splat(@max(dot, 0)); @memset(out, sums.sumIn(acc, order, &terms) / @sqrt(@as(f64, di))); @@ -331,7 +371,7 @@ const Runner = struct { const a = self.arena.?.allocator(); const wide = try normalized(a, hd, acc, order); const narrow = try normalized(a, di, acc, order); - const prep = self.mode == .prep or self.mode == .prep_media; + const prep = self.base() == .prep or self.base() == .prep_media; for (out, 0..) |*o, i| o.* = if (!prep or i >= self.rows * qh * hd) narrow[i % di] else wide[i % hd]; } }; diff --git a/zig/tests/cuda/contract_limits.zig b/zig/tests/cuda/contract_limits.zig index bb48bb5e08..4b556a4463 100644 --- a/zig/tests/cuda/contract_limits.zig +++ b/zig/tests/cuda/contract_limits.zig @@ -27,7 +27,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", "nemotron_ops", "minimax_m3", "minimax_m3_attn", "moss_image" }) |n| if (std.mem.eql(u8, name, n)) return @field(kk, n); - inline for (.{ "flashnext_exl3", "flashnext_exl3_tp", "flashnext_exl3_prefill" }) |n| + inline for (.{ "flashnext_exl3", "flashnext_exl3_tp", "flashnext_exl3_kvtier", "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/kv/exl3_kvtier_tests.zig b/zig/tests/cuda/flashnext_exl3/kv/exl3_kvtier_tests.zig new file mode 100644 index 0000000000..3adc571c5f --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/kv/exl3_kvtier_tests.zig @@ -0,0 +1,221 @@ +//! exl3-kvtier: the host tier slots (kvt_plan, kvt_fetch, kvt_remap) against the host LRU model and store bytes. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const check = @import("../../check.zig"); +const flash = @import("flashnext"); +const Gpu = check.Gpu; +const kvt = flash.exl3.kvtier; +const pages = flash.exl3.pages; +const host_tier = flash.exl3.kvhost; +const kv = core.kv_pool; + +const ratio = 4; +const idw = flash.exl3.prompt_wt.index_idw; + +/// The host's model of kvt_plan: exact LRU at call granularity, ties to the lower slot. +const Model = struct { + stamp: []u32, + slot_of: std.AutoHashMap(u32, u32), + blk: []i64, + epoch: u32 = 0, + + fn use(mo: *Model, a: std.mem.Allocator, need: []const u32) !usize { + mo.epoch += 1; + var misses: std.ArrayList(u32) = .empty; + defer misses.deinit(a); + for (need) |b| { + if (mo.slot_of.get(b)) |sl| mo.stamp[sl] = mo.epoch else try misses.append(a, b); + } + const order = try a.alloc(u32, mo.stamp.len); + defer a.free(order); + for (order, 0..) |*o, i| o.* = @intCast(i); + const Ctx = struct { + st: []const u32, + ep: u32, + fn key(c: @This(), i: u32) u64 { + const t = c.st[i]; + return if (t == c.ep) std.math.maxInt(u64) else (@as(u64, t) << 32) | i; + } + fn less(c: @This(), x: u32, y: u32) bool { + return c.key(x) < c.key(y); + } + }; + std.mem.sort(u32, order, Ctx{ .st = mo.stamp, .ep = mo.epoch }, Ctx.less); + for (misses.items, 0..) |b, i| { + const sl = order[i]; + if (mo.stamp[sl] == mo.epoch) return error.ModelOverflow; + if (mo.blk[sl] >= 0) _ = mo.slot_of.remove(@intCast(mo.blk[sl])); + mo.blk[sl] = b; + try mo.slot_of.put(b, sl); + mo.stamp[sl] = mo.epoch; + } + return misses.items.len; + } +}; + +/// `slot_frames` lent frames of a one-layer tier over `npages` pages: drifting 512-block selections. +pub fn kernelRun(gpu: Gpu, slot_frames: usize) !void { + const a = gpu.gpa; + const npages: usize = 200; + const w0: usize = 40_000; + const nb_before = w0 / ratio; + const nslots = slot_frames * 64; + const kv_page = pages.kvPageBytes(kvt.page_rows, 1); + const row = pages.rowBytes(.key); + var stream = try cuda.Stream.init(gpu.d, true); + defer stream.deinit(); + var host = try host_tier.Host.create(a, gpu.d, npages, 1 + slot_frames, kv_page); + defer host.destroy(a); + var lease: kv.tier.Lease = .{}; + defer host.tier.finish(&lease); + try host.tier.reserve(&lease, npages, 1 + slot_frames); + try host.tier.cover(&lease, npages, null); + try host.tier.lend(&lease, slot_frames); + var arena = try cuda.DeviceBuffer.alloc(gpu.d, (1 + slot_frames) * kv_page); + defer arena.free(); + var table = try cuda.DeviceBuffer.alloc(gpu.d, (npages + slot_frames) * 4); + defer table.free(); + var itable = try cuda.DeviceBuffer.alloc(gpu.d, npages * 4); + defer itable.free(); + try itable.upload(0, std.mem.sliceAsBytes(lease.pages.ids.items)); + var fids: std.ArrayList(u32) = .empty; + defer fids.deinit(a); + try fids.appendNTimes(a, std.math.maxInt(u32), npages); + var views: [1]cuda.DeviceBuffer = .{.{ .d = gpu.d, .ptr = arena.ptr + pages.planeOffset(kvt.page_rows, 0, .key), .len = 0, .borrowed = true }}; + var vviews: [1]cuda.DeviceBuffer = .{.{ .d = gpu.d, .ptr = arena.ptr + pages.planeOffset(kvt.page_rows, 0, .value), .len = 0, .borrowed = true }}; + var kern = try kvt.Kernels.load(gpu.d, nslots); + defer kern.unload(); + const sel = try kvt.LaneSel.init(a, gpu.d, &kern, npages * kvt.page_rows, npages, slot_frames, 1, .{ .table = table, .itable = itable, .host = host, .arena = arena, .tl = &lease, .fids = &fids, .kc = &views, .vc = &vviews }); + defer sel.deinit(); + try sel.reset(stream); + try sel.nameSlots(stream); + try stream.synchronize(); + sel.win0 = w0; + // each host row's words unique to (position, K or V) + for (0..npages) |pg| { + const id = lease.pages.ids.items[pg]; + const words = std.mem.bytesAsSlice(u32, host.tier.host_page(id)); + for (0..kvt.page_rows) |r| for (0..2) |v| for (0..row / 4) |i| { + words[(v * kvt.page_rows + r) * (row / 4) + i] = @intCast(((pg * 256 + r) *% 2654435761 +% i *% 40503 +% v *% 977) & 0xFFFF_FFFF); + }; + } + var ids_d = try cuda.DeviceBuffer.alloc(gpu.d, 256 * idw * 4); + defer ids_d.free(); + var nk_d = try cuda.DeviceBuffer.alloc(gpu.d, 256 * 4); + defer nk_d.free(); + var sp_d = try cuda.DeviceBuffer.alloc(gpu.d, 256 * 4); + defer sp_d.free(); + const ids = try a.alloc(i32, 256 * idw); + defer a.free(ids); + const back = try a.alloc(i32, 256 * idw); + defer a.free(back); + const nk = try a.alloc(i32, 256); + defer a.free(nk); + const sp = try a.alloc(i32, 256); + defer a.free(sp); + @memset(sp, 1); + const tbl = try a.alloc(i32, npages + slot_frames); + defer a.free(tbl); + try gpu.d.check(gpu.d.api.cuMemcpyDtoH_v2(tbl.ptr, table.ptr, tbl.len * 4), "cuMemcpyDtoH"); + var model: Model = .{ .stamp = try a.alloc(u32, nslots), .slot_of = .init(a), .blk = try a.alloc(i64, nslots) }; + defer { + a.free(model.stamp); + a.free(model.blk); + model.slot_of.deinit(); + } + @memset(model.stamp, 0); + @memset(model.blk, -1); + var prng = std.Random.DefaultPrng.init(42); + const rng = prng.random(); + const per_plan = sel.rowsPerPlan(); + var base: [512]u32 = undefined; + var used = std.AutoHashMap(u32, void).init(a); + defer used.deinit(); + var k: usize = 0; + while (k < 512) { + const b = rng.uintLessThan(u32, @intCast(nb_before)); + if ((try used.getOrPut(b)).found_existing) continue; + base[k] = b; + k += 1; + } + var gpu_miss: usize = 0; + var model_miss: usize = 0; + var bad: usize = 0; + var keys_checked: usize = 0; + const calls: usize = 60; + for (0..calls) |c| { + const m: usize = @min(per_plan, @as(usize, if (c % 10 == 9) 64 else 1)); + var need = std.AutoHashMap(u32, void).init(a); + defer need.deinit(); + for (0..m) |r| { + for (0..77) |_| { // drift: ~15% of the selection replaced a row + const j = rng.uintLessThan(usize, 512); + while (true) { + const b = rng.uintLessThan(u32, @intCast(nb_before)); + if ((try used.getOrPut(b)).found_existing) continue; + _ = used.remove(base[j]); + base[j] = b; + break; + } + } + var sorted = base; + std.mem.sort(u32, &sorted, {}, std.sort.asc(u32)); + const out = ids[r * idw ..][0..idw]; + for (sorted, 0..) |b, j| for (0..4) |t| { + out[j * 4 + t] = @intCast(b * 4 + t); + }; + for (0..3) |t| out[2048 + t] = @intCast(w0 + 100 + c * 4 + t); // keys inside the window + nk[r] = 2051; + for (sorted) |b| try need.put(b, {}); + } + var need_list: std.ArrayList(u32) = .empty; + defer need_list.deinit(a); + var it = need.keyIterator(); + while (it.next()) |b| try need_list.append(a, b.*); + model_miss += try model.use(a, need_list.items); + try gpu.d.check(gpu.d.api.cuMemcpyHtoD_v2(ids_d.ptr, ids.ptr, m * idw * 4), "cuMemcpyHtoD"); + try gpu.d.check(gpu.d.api.cuMemcpyHtoD_v2(nk_d.ptr, nk.ptr, m * 4), "cuMemcpyHtoD"); + try gpu.d.check(gpu.d.api.cuMemcpyHtoD_v2(sp_d.ptr, sp.ptr, m * 4), "cuMemcpyHtoD"); + const before = try sel.stats(); + try sel.page(stream, 0, ids_d.ptr, nk_d.ptr, sp_d.ptr, m); + try stream.synchronize(); + const after = try sel.stats(); + gpu_miss += after.fetched - before.fetched; + try gpu.d.check(gpu.d.api.cuMemcpyDtoH_v2(back.ptr, ids_d.ptr, m * idw * 4), "cuMemcpyDtoH"); + for (0..m) |r| for (0..2051) |ki| { + const p: usize = @intCast(ids[r * idw + ki]); + const q: usize = @intCast(back[r * idw + ki]); + keys_checked += 1; + if (p >= w0) { + bad += @intFromBool(q != p); + continue; + } + // row q of the table: entry q / 256 is a lent frame, the row is where the slot's bytes sit in the arena + const fr: usize = @intCast(tbl[q / kvt.page_rows]); + if (q < npages * kvt.page_rows or q / kvt.page_rows >= npages + slot_frames) { + bad += 1; + continue; + } + var got: [2][1024]u8 = undefined; + for ([_]pages.Part{ .key, .value }, 0..) |part, v| { + const at = fr * kv_page + pages.planeOffset(kvt.page_rows, 0, part) + (q % kvt.page_rows) * row; + try gpu.d.check(gpu.d.api.cuMemcpyDtoH_v2(&got[v], arena.ptr + at, row), "cuMemcpyDtoH"); + const hp = host.tier.host_page(lease.pages.ids.items[p / kvt.page_rows]); + const want = hp[pages.planeOffset(kvt.page_rows, 0, part) + (p % kvt.page_rows) * row ..][0..row]; + bad += @intFromBool(!std.mem.eql(u8, &got[v], want)); + } + }; + } + const st = try sel.stats(); + std.debug.print("kvtier kernel ({d} slot frames, {d} slots): {d} calls ({d} keys): bad {d}; misses GPU {d}, host model {d}; needed {d}; overflow {}\n", .{ slot_frames, nslots, calls, keys_checked, bad, gpu_miss, model_miss, st.needed, st.overflow }); + try check.expect(bad == 0 and !st.overflow, "every remapped key reads its own host row", .{}); + const diff = @as(f64, @floatFromInt(@max(gpu_miss, model_miss) - @min(gpu_miss, model_miss))) / @as(f64, @floatFromInt(@max(model_miss, 1))); + try check.expect(diff < 0.05, "misses within 5% of the LRU model ({d:.3})", .{diff}); + check.pass("exl3-kvtier kernel, {d} slot frames", .{slot_frames}); +} + +pub fn run(gpu: Gpu, args: []const []const u8) !void { + _ = args; + for ([_]usize{ 128, 16, 9 }) |frames| try kernelRun(gpu, frames); +} diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index bdb6bb3a5b..b713bbd283 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -103,6 +103,7 @@ 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 exl3_kvtier_tests = @import("flashnext_exl3/kv/exl3_kvtier_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"); @@ -366,6 +367,7 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { 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, "exl3-kvtier")) return exl3_kvtier_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); From cb63fec925d1c75cc4386312e8db3f2845612ee7 Mon Sep 17 00:00:00 2001 From: BlivionIaG <15323295+BlivionIaG@users.noreply.github.com> Date: Thu, 29 Oct 2026 20:44:59 +0100 Subject: [PATCH 7/7] EXL3 Flash Next: CUDA graphs and the device planner --- docs/flash-next-exl3-graphs.md | 32 + docs/kernel-contract.md | 2 +- tools/exl3_nsys_step_breakdown.py | 41 ++ tools/exl3_serve_tp2.sh | 6 +- tools/exl3_snapshot_gate.py | 7 +- tools/tests/test_exl3_snapshot_gate.py | 18 +- zig/build/cuda.zig | 2 +- zig/build/cuda_kernels.zig | 1 + zig/build/exl3_host.zig | 19 + .../flashnext_exl3/flashnext_exl3_dplan.cu | 430 +++++++++++++ .../core/contract/flashnext_exl3_pending.zig | 14 +- zig/src/cuda/abi.zig | 1 + zig/src/cuda/kernels.zig | 1 + zig/src/cuda/stream.zig | 4 + .../families/flashnext/cuda_batch_state.zig | 28 +- .../families/flashnext/cuda_batch_uploads.zig | 21 + .../families/flashnext/cuda_collective.zig | 199 ++++++ .../families/flashnext/cuda_images_test.zig | 4 + zig/src/families/flashnext/cuda_lanes.zig | 16 +- .../flashnext/cuda_lanes_admission_test.zig | 2 + .../families/flashnext/cuda_lanes_load.zig | 13 +- .../flashnext/cuda_lanes_restore_test.zig | 2 + zig/src/families/flashnext/cuda_launch.zig | 1 + zig/src/families/flashnext/cuda_mtp_check.zig | 2 +- zig/src/families/flashnext/cuda_mtp_draft.zig | 25 + zig/src/families/flashnext/cuda_mtp_ops.zig | 51 +- .../families/flashnext/cuda_mtp_ops_test.zig | 144 ++++- .../flashnext/cuda_mtp_verify_cache.zig | 65 ++ .../flashnext/cuda_mtp_verify_tail.zig | 28 + zig/src/families/flashnext/cuda_prompt.zig | 23 +- .../flashnext/cuda_prompt_lifecycle.zig | 2 + .../families/flashnext/cuda_rank_engine.zig | 19 +- .../flashnext/cuda_rank_media_test.zig | 4 + .../cuda_sequence_completion_test.zig | 6 + .../flashnext/exl3/dplan/cuda_exl3_dplan.zig | 516 +++++++++++++++ .../exl3/dplan/cuda_exl3_dplan_test.zig | 193 ++++++ .../exl3/dplan/exl3_dplan_layout.zig | 128 ++++ .../exl3/dplan/exl3_dplan_policy_eval.zig | 374 +++++++++++ .../dplan/exl3_dplan_policy_eval_test.zig | 289 +++++++++ .../exl3/dplan/exl3_dplan_replay.zig | 392 ++++++++++++ .../flashnext/exl3/dplan/exl3_dplan_trace.zig | 189 ++++++ .../dplan/exl3_dplan_trace_kernel_test.cpp | 261 ++++++++ .../dplan/exl3_dplan_trace_replay_test.zig | 294 +++++++++ .../exl3/dplan/exl3_dplan_trace_test.zig | 262 ++++++++ .../dplan/exl3_dplan_trace_writer_test.zig | 31 + .../flashnext/exl3/graph/cuda_exl3_graph.zig | 98 +++ .../exl3/graph/cuda_exl3_graph_budget.zig | 128 ++++ .../exl3/graph/cuda_exl3_graph_cache.zig | 196 ++++++ .../exl3/graph/cuda_exl3_graph_cache_test.zig | 158 +++++ .../exl3/graph/cuda_exl3_graph_fallback.zig | 140 +++++ .../graph/cuda_exl3_graph_fallback_test.zig | 253 ++++++++ .../exl3/graph/cuda_exl3_graph_hooks.zig | 127 ++++ .../exl3/graph/cuda_exl3_graph_key.zig | 74 +++ .../exl3/graph/cuda_exl3_graph_key_test.zig | 129 ++++ .../exl3/graph/cuda_exl3_graph_log.zig | 28 + .../exl3/graph/cuda_exl3_graph_memory.zig | 398 ++++++++++++ .../exl3/graph/cuda_exl3_graph_pages.zig | 10 + .../exl3/graph/cuda_exl3_graph_pending.zig | 75 +++ .../exl3/graph/cuda_exl3_graph_prefill.zig | 286 +++++++++ .../graph/cuda_exl3_graph_prefill_test.zig | 137 ++++ .../exl3/graph/cuda_exl3_graph_rebind.zig | 148 +++++ .../graph/cuda_exl3_graph_rebind_test.zig | 145 +++++ .../exl3/graph/cuda_exl3_graph_run.zig | 140 +++++ .../graph/cuda_exl3_graph_test_driver.zig | 117 ++++ .../exl3/graph/cuda_exl3_graph_together.zig | 190 ++++++ .../graph/cuda_exl3_graph_together_test.zig | 287 +++++++++ .../flashnext/exl3/graph/exl3_graph.zig | 414 ++++++++++++ .../exl3/graph/exl3_graph_budget.zig | 206 ++++++ .../flashnext/exl3/kv/cuda_exl3_kvtier.zig | 2 + .../exl3/kv/cuda_exl3_state_binding.zig | 239 +++++++ .../exl3/kv/cuda_exl3_state_request_test.zig | 72 +++ .../flashnext/exl3/kv/exl3_state_binding.zig | 177 ++++++ .../exl3/kv/exl3_state_binding_test.zig | 298 +++++++++ .../flashnext/exl3/model/cuda_exl3.zig | 8 + .../flashnext/exl3/model/cuda_exl3_arch.zig | 2 +- .../flashnext/exl3/model/cuda_exl3_chain.zig | 594 ++++++++++++++++++ .../exl3/model/cuda_exl3_fake_driver_test.zig | 15 +- .../flashnext/exl3/model/cuda_exl3_pass.zig | 90 ++- .../exl3/model/cuda_exl3_telemetry.zig | 3 +- .../exl3/model/cuda_exl3_telemetry_test.zig | 113 ++++ .../flashnext/exl3/model/exl3_cache.zig | 2 +- .../flashnext/exl3/model/exl3_round.zig | 5 +- .../offload/cuda_exl3_frequency_state.zig | 54 +- .../flashnext/exl3/offload/cuda_exl3_pool.zig | 4 + .../exl3/offload/cuda_exl3_recycle.zig | 144 +++++ .../exl3/offload/cuda_exl3_recycle_test.zig | 300 +++++++++ .../exl3/offload/exl3_line_deltas.zig | 62 +- .../exl3/pack/exl3_startup_memory.zig | 150 ++++- .../exl3/prefill/cuda_exl3_prefill.zig | 100 +-- .../exl3/prefill/cuda_exl3_prefill_chunk.zig | 97 ++- .../exl3/prefill/cuda_exl3_prefill_head.zig | 21 +- .../exl3/prefill/cuda_exl3_prefill_moe.zig | 21 +- .../prefill/cuda_exl3_prefill_ringmoe.zig | 15 +- .../exl3/prefill/cuda_exl3_prefill_state.zig | 28 +- .../exl3/prefill/cuda_exl3_prefill_test.zig | 63 +- .../exl3/prefill/cuda_exl3_prefill_tiles.zig | 4 +- .../exl3/prefill/exl3_prefill_window.zig | 8 + .../flashnext/exl3/tp/cuda_exl3_tp.zig | 8 +- zig/src/native/draft_sources_cuda_test.zig | 3 + .../cuda/contract_flashnext_exl3_dplan.zig | 419 ++++++++++++ .../contract_flashnext_exl3_dplan_copy.zig | 213 +++++++ .../contract_flashnext_exl3_dplan_draft.zig | 293 +++++++++ .../contract_flashnext_exl3_dplan_util.zig | 40 ++ .../cuda/contract_flashnext_exl3_suite.zig | 12 +- zig/tests/cuda/contract_limits.zig | 2 +- .../graph/flashnext_exl3_graph_engine.zig | 159 +++++ .../graph/flashnext_exl3_graph_eq_tests.zig | 237 +++++++ .../graph/flashnext_exl3_graph_plan_tests.zig | 207 ++++++ .../flashnext_exl3_graph_prefill_tests.zig | 188 ++++++ .../flashnext_exl3_graph_snapshot_tests.zig | 93 +++ .../graph/flashnext_exl3_graph_tests.zig | 39 ++ .../kv/exl3_state_rebind_tests.zig | 117 ++++ .../model/flashnext_exl3_engine_tests.zig | 4 +- .../offload/exl3_frequency_tests.zig | 19 +- .../offload/exl3_pool_equal_tests.zig | 83 ++- .../offload/exl3_pool_tests.zig | 1 + .../prefill/exl3_prefill_engine.zig | 12 +- .../tp/exl3_tpprefill_tests.zig | 12 +- .../tp/flashnext_exl3_tp_engine_tests.zig | 14 + .../tp/flashnext_exl3_tpe_engine.zig | 4 +- .../tp/flashnext_exl3_tpe_graph_tests.zig | 139 ++++ .../tp/flashnext_exl3_tpe_hits_tests.zig | 154 +++++ .../tp/flashnext_exl3_tpe_trace_tests.zig | 139 ++++ zig/tests/cuda/main.zig | 12 +- 124 files changed, 13068 insertions(+), 292 deletions(-) create mode 100644 docs/flash-next-exl3-graphs.md create mode 100644 tools/exl3_nsys_step_breakdown.py create mode 100644 zig/build/exl3_host.zig create mode 100644 zig/kernels/cuda/flashnext_exl3/flashnext_exl3_dplan.cu create mode 100644 zig/src/families/flashnext/cuda_batch_uploads.zig create mode 100644 zig/src/families/flashnext/cuda_mtp_verify_cache.zig create mode 100644 zig/src/families/flashnext/cuda_mtp_verify_tail.zig create mode 100644 zig/src/families/flashnext/exl3/dplan/cuda_exl3_dplan.zig create mode 100644 zig/src/families/flashnext/exl3/dplan/cuda_exl3_dplan_test.zig create mode 100644 zig/src/families/flashnext/exl3/dplan/exl3_dplan_layout.zig create mode 100644 zig/src/families/flashnext/exl3/dplan/exl3_dplan_policy_eval.zig create mode 100644 zig/src/families/flashnext/exl3/dplan/exl3_dplan_policy_eval_test.zig create mode 100644 zig/src/families/flashnext/exl3/dplan/exl3_dplan_replay.zig create mode 100644 zig/src/families/flashnext/exl3/dplan/exl3_dplan_trace.zig create mode 100644 zig/src/families/flashnext/exl3/dplan/exl3_dplan_trace_kernel_test.cpp create mode 100644 zig/src/families/flashnext/exl3/dplan/exl3_dplan_trace_replay_test.zig create mode 100644 zig/src/families/flashnext/exl3/dplan/exl3_dplan_trace_test.zig create mode 100644 zig/src/families/flashnext/exl3/dplan/exl3_dplan_trace_writer_test.zig create mode 100644 zig/src/families/flashnext/exl3/graph/cuda_exl3_graph.zig create mode 100644 zig/src/families/flashnext/exl3/graph/cuda_exl3_graph_budget.zig create mode 100644 zig/src/families/flashnext/exl3/graph/cuda_exl3_graph_cache.zig create mode 100644 zig/src/families/flashnext/exl3/graph/cuda_exl3_graph_cache_test.zig create mode 100644 zig/src/families/flashnext/exl3/graph/cuda_exl3_graph_fallback.zig create mode 100644 zig/src/families/flashnext/exl3/graph/cuda_exl3_graph_fallback_test.zig create mode 100644 zig/src/families/flashnext/exl3/graph/cuda_exl3_graph_hooks.zig create mode 100644 zig/src/families/flashnext/exl3/graph/cuda_exl3_graph_key.zig create mode 100644 zig/src/families/flashnext/exl3/graph/cuda_exl3_graph_key_test.zig create mode 100644 zig/src/families/flashnext/exl3/graph/cuda_exl3_graph_log.zig create mode 100644 zig/src/families/flashnext/exl3/graph/cuda_exl3_graph_memory.zig create mode 100644 zig/src/families/flashnext/exl3/graph/cuda_exl3_graph_pages.zig create mode 100644 zig/src/families/flashnext/exl3/graph/cuda_exl3_graph_pending.zig create mode 100644 zig/src/families/flashnext/exl3/graph/cuda_exl3_graph_prefill.zig create mode 100644 zig/src/families/flashnext/exl3/graph/cuda_exl3_graph_prefill_test.zig create mode 100644 zig/src/families/flashnext/exl3/graph/cuda_exl3_graph_rebind.zig create mode 100644 zig/src/families/flashnext/exl3/graph/cuda_exl3_graph_rebind_test.zig create mode 100644 zig/src/families/flashnext/exl3/graph/cuda_exl3_graph_run.zig create mode 100644 zig/src/families/flashnext/exl3/graph/cuda_exl3_graph_test_driver.zig create mode 100644 zig/src/families/flashnext/exl3/graph/cuda_exl3_graph_together.zig create mode 100644 zig/src/families/flashnext/exl3/graph/cuda_exl3_graph_together_test.zig create mode 100644 zig/src/families/flashnext/exl3/graph/exl3_graph.zig create mode 100644 zig/src/families/flashnext/exl3/graph/exl3_graph_budget.zig create mode 100644 zig/src/families/flashnext/exl3/kv/cuda_exl3_state_binding.zig create mode 100644 zig/src/families/flashnext/exl3/kv/cuda_exl3_state_request_test.zig create mode 100644 zig/src/families/flashnext/exl3/kv/exl3_state_binding.zig create mode 100644 zig/src/families/flashnext/exl3/kv/exl3_state_binding_test.zig create mode 100644 zig/src/families/flashnext/exl3/model/cuda_exl3_chain.zig create mode 100644 zig/src/families/flashnext/exl3/offload/cuda_exl3_recycle.zig create mode 100644 zig/src/families/flashnext/exl3/offload/cuda_exl3_recycle_test.zig create mode 100644 zig/tests/cuda/contract_flashnext_exl3_dplan.zig create mode 100644 zig/tests/cuda/contract_flashnext_exl3_dplan_copy.zig create mode 100644 zig/tests/cuda/contract_flashnext_exl3_dplan_draft.zig create mode 100644 zig/tests/cuda/contract_flashnext_exl3_dplan_util.zig create mode 100644 zig/tests/cuda/flashnext_exl3/graph/flashnext_exl3_graph_engine.zig create mode 100644 zig/tests/cuda/flashnext_exl3/graph/flashnext_exl3_graph_eq_tests.zig create mode 100644 zig/tests/cuda/flashnext_exl3/graph/flashnext_exl3_graph_plan_tests.zig create mode 100644 zig/tests/cuda/flashnext_exl3/graph/flashnext_exl3_graph_prefill_tests.zig create mode 100644 zig/tests/cuda/flashnext_exl3/graph/flashnext_exl3_graph_snapshot_tests.zig create mode 100644 zig/tests/cuda/flashnext_exl3/graph/flashnext_exl3_graph_tests.zig create mode 100644 zig/tests/cuda/flashnext_exl3/kv/exl3_state_rebind_tests.zig create mode 100644 zig/tests/cuda/flashnext_exl3/tp/flashnext_exl3_tpe_graph_tests.zig create mode 100644 zig/tests/cuda/flashnext_exl3/tp/flashnext_exl3_tpe_hits_tests.zig create mode 100644 zig/tests/cuda/flashnext_exl3/tp/flashnext_exl3_tpe_trace_tests.zig diff --git a/docs/flash-next-exl3-graphs.md b/docs/flash-next-exl3-graphs.md new file mode 100644 index 0000000000..319f1ee849 --- /dev/null +++ b/docs/flash-next-exl3-graphs.md @@ -0,0 +1,32 @@ +# Flash Next EXL3: CUDA graphs and the device planner + +With the expert offload pool, a decode round of Flash Next EXL3 plans its experts on the device +(`flashnext_exl3_dplan.cu`) instead of waiting for the host, and can replay the whole round as a CUDA graph. Acceptance +stays in the core lane engine; the device code proposes and plans. + +## Switches + +| Variable | Default | Meaning | +| --- | --- | --- | +| `TF_FLASHNEXT_EXL3_GRAPH` | `0` | `1` plans decode rounds on the device and replays graphs for decode rounds, verify windows, the MTP draft chain and prefill layers. Output is identical to `0`. | +| `TF_FLASHNEXT_EXL3_DPLAN` | `0` | `1` plans on the device without graphs (each planned round launches kernel by kernel). Implied by `GRAPH=1`. This is the eager arm of an eager-vs-graph comparison. | +| `TF_FLASHNEXT_EXL3_GRAPH_ASYNC` | `1` | `0` instantiates a captured graph on the step's thread; `1` on a worker thread while that step runs eagerly. | + +Everything else is fixed in the code: a shape is captured at its second sighting, a device-planned round holds at most +16 rows (the MTP window), and the graph memory model (reservation, headroom, per-graph charges) follows the stream +count. `tools/exl3_serve_tp2.sh` sets `GRAPH=1` and `GRAPH_ASYNC=0`. + +## A graph that cannot be built + +A capture or an instantiation that fails runs its round eagerly and the process keeps its other graphs. The failed +shape backs off (64, 128, 256 rounds) and is tried at most three more times before it stays eager. The prefill layer +shape and the draft chain do the same. + +On two ranks the capture outcome is agreed over the rank wire (`graphAgree`) before either rank launches: a graph is +kept only when both ranks hold one, otherwise both run the round eagerly, so the communicator is never left +half-launched. A communicator error (a lost peer, a timeout) still ends the step, as before. + +## State the engine owns + +The device planner, the plan active during a round and the MTP head's pool belong to the engine +(`Engine.exl3_dplan`), so two engines in one process share none of them. diff --git a/docs/kernel-contract.md b/docs/kernel-contract.md index 64e6e8d4cc..f995fdf3a6 100644 --- a/docs/kernel-contract.md +++ b/docs/kernel-contract.md @@ -256,7 +256,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 17 EXL3 launches (`flashnext_exl3_pending.zig`) carry only an nvcc 13.4 receipt: they stay pending under +Flash Next's 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/tools/exl3_nsys_step_breakdown.py b/tools/exl3_nsys_step_breakdown.py new file mode 100644 index 0000000000..bb8090151b --- /dev/null +++ b/tools/exl3_nsys_step_breakdown.py @@ -0,0 +1,41 @@ +"""Per-decode-step kernel time by category from an nsys SQLite export of a rank (counts fn_exl3_dplan launches / 48 as steps). usage: exl3_nsys_step_breakdown.py EXPORT_PREFIX.""" +import sqlite3, collections, re, subprocess, os, sys +f = sys.argv[1] # the export's path without its .sqlite / .nsys-rep suffix +if not os.path.exists(f + ".sqlite"): + subprocess.run(["nsys", "stats", "--report", "cuda_gpu_kern_sum", "-o", "/dev/null", f + ".nsys-rep"], capture_output=True) +c = sqlite3.connect(f + ".sqlite") +names = dict(c.execute("select id,value from StringIds")) +rows = c.execute("select start,end,demangledName from CUPTI_ACTIVITY_KIND_KERNEL order by start").fetchall() +dp = [r for r in rows if "fn_exl3_dplan" in names[r[2]]] +steps = len(dp) // 48 +t0, t1 = dp[0][0], dp[-1][1] +print("dplan launches", len(dp), "decode steps", steps, "window s", round((t1 - t0) / 1e9, 2), "ms/step wall", round((t1 - t0) / 1e6 / steps, 2)) +def cat(n): + if "AllReduce" in n or "AllGather" in n: return "nccl all-reduce/gather" + if "fn_exl3_dcopy" in n: return "expert miss copy (side stream)" + if "tf_exl3_linear" in n or "tf_exl3_grouped" in n or "tf_exl3_group" in n or "tf_exl3_gateup" in n or "tf_exl3_down" in n or "tf_exl3_rot" in n: return "expert/dense exl3 matmuls" + if "qmmf" in n: return "qmmf matmuls" + if "fn_gdn" in n: return "DeltaNet (gdn)" + if "fn_hc_" in n: return "hyper-connection" + if "attn" in n or "fn_qsa" in n: return "attention" + if "tf_route" in n or "router" in n: return "router" + if "dplan" in n: return "device planner" + return "other" +tot = collections.Counter(); cnt = collections.Counter(); other = collections.Counter() +for s, e, n in rows: + if s < t0 or e > t1: continue + nm = names[n]; k = cat(nm); tot[k] += e - s; cnt[k] += 1 + if k == "other": other[nm[:70]] += e - s +for k, v in tot.most_common(): + print("%-34s %7.2f ms/step %5d launches/step" % (k, v / 1e6 / steps, cnt[k] // steps)) +print("sum of kernel time (all streams) ms/step", round(sum(tot.values()) / 1e6 / steps, 2)) +print("top 'other':") +for k, v in other.most_common(8): print(" %6.2f ms/step %s" % (v / 1e6 / steps, k)) +# main stream busy fraction: union over all non-dcopy kernels +iv = sorted((s, e) for s, e, n in rows if t0 <= s and e <= t1 and "dcopy" not in names[n]) +busy = 0; cs, ce = iv[0] +for s, e in iv[1:]: + if s > ce: busy += ce - cs; cs, ce = s, e + else: ce = max(ce, e) +busy += ce - cs +print("busy fraction (non-copy kernels)", round(busy / (t1 - t0), 3)) diff --git a/tools/exl3_serve_tp2.sh b/tools/exl3_serve_tp2.sh index e051f0d5ac..b9d390bcb4 100755 --- a/tools/exl3_serve_tp2.sh +++ b/tools/exl3_serve_tp2.sh @@ -3,12 +3,16 @@ # 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. +# PARALLEL (streams, default 2), TENSORFOLD_BIN, TF_FLASHNEXT_EXL3_GRAPH (1: graphs, 0: the host cut). 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)} +# Graphs (decode rounds, verify windows, the draft chain, prefill layers) and the device planner they replay: one switch. +# TF_FLASHNEXT_EXL3_GRAPH_ASYNC=1 instantiates new graphs on a worker thread (default here: 0, on the step's thread). +export TF_FLASHNEXT_EXL3_GRAPH=${TF_FLASHNEXT_EXL3_GRAPH:-1} +export TF_FLASHNEXT_EXL3_GRAPH_ASYNC=${TF_FLASHNEXT_EXL3_GRAPH_ASYNC:-0} 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 diff --git a/tools/exl3_snapshot_gate.py b/tools/exl3_snapshot_gate.py index eb1ef83481..554fd4d82b 100644 --- a/tools/exl3_snapshot_gate.py +++ b/tools/exl3_snapshot_gate.py @@ -4,6 +4,7 @@ import argparse import concurrent.futures import json +import re import threading import time import urllib.request @@ -50,9 +51,9 @@ def ask(ids): 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") + digest = runtime.get("token_sha256") + if not isinstance(digest, str) or not re.fullmatch(r"[0-9a-f]{64}", digest): + raise AssertionError("missing full output token SHA256") return { "sha": digest, "cached_tokens": usage["prompt_tokens_details"]["cached_tokens"], diff --git a/tools/tests/test_exl3_snapshot_gate.py b/tools/tests/test_exl3_snapshot_gate.py index 7363c7c345..4969d8c4f4 100644 --- a/tools/tests/test_exl3_snapshot_gate.py +++ b/tools/tests/test_exl3_snapshot_gate.py @@ -32,17 +32,20 @@ def read(self): class SnapshotGateTests(unittest.TestCase): - def invoke(self, args, cached): + def invoke(self, args, cached, runtime_override=None): 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() + runtime = {"token_sha256": digest} + if runtime_override is not None: + runtime = runtime_override return Response({ "usage": {"prompt_tokens": 80, "completion_tokens": 2, "prompt_tokens_details": {"cached_tokens": cached}}, - "tensorfold": {"token_sha256": digest}, + "tensorfold": runtime, }) output, errors = io.StringIO(), io.StringIO() @@ -81,6 +84,17 @@ def test_missing_warm_cache_hit_fails(self): self.assertEqual(code, 1) self.assertIn("did not restore a cached prefix", errors) + def test_missing_or_inexact_full_digest_cannot_record_a_reference(self): + for runtime in ({"token_sha": "a" * 12}, {"token_sha256": "a" * 12}, + {"token_sha256": "z" * 64}, {"token_sha256": None}): + with self.subTest(runtime=runtime), tempfile.TemporaryDirectory() as directory: + reference = Path(directory) / "reference.json" + code, _, errors, _ = self.invoke( + ["--record-reference", "--output", str(reference)], 0, runtime) + self.assertEqual(code, 1) + self.assertIn("missing full output token SHA256", errors) + self.assertFalse(reference.exists()) + if __name__ == "__main__": unittest.main() diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index eba95dda8b..5268f03a85 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -451,7 +451,7 @@ pub fn hostTests(b: *std.Build, draft_ids: *std.Build.Module, build_options: *st @import("cortheon_fuzz.zig").add(b, host, native_modules.api, mods.tokenizer, template); } // 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); + @import("exl3_host.zig").steps(b, host, flashnextModule(b, host, .debug, cuda, mods.core, mods.lanes, mods.qwen_vision, mods.api)); 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 05454f9f7b..1c9a145bc0 100644 --- a/zig/build/cuda_kernels.zig +++ b/zig/build/cuda_kernels.zig @@ -124,6 +124,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_dplan", .src = "flashnext_exl3/flashnext_exl3_dplan", .flags = &.{"-O3"} }, // ours: Flash Next's EXL3 pool planned on the device (decode graphs) .{ .name = "flashnext_exl3_kvtier", .src = "flashnext_exl3/flashnext_exl3_kvtier", .flags = glue }, // ours: Flash Next K/V on the host tier, a device block cache in lent frames .{ .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 }; diff --git a/zig/build/exl3_host.zig b/zig/build/exl3_host.zig new file mode 100644 index 0000000000..70ab051fac --- /dev/null +++ b/zig/build/exl3_host.zig @@ -0,0 +1,19 @@ +//! The EXL3 host tests of the Flash Next module and the CPU-only planner trace writer parity test. +const std = @import("std"); + +pub fn steps(b: *std.Build, host: std.Build.ResolvedTarget, flashnext: *std.Build.Module) void { + b.step("test-flashnext-exl3-host", "Flash Next EXL3 admission, store, sizing and loader host tests").dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = flashnext, .filters = &.{"exl3: "} })).step); + const trace_writer_build = b.addSystemCommand(&.{ "c++", "-std=c++17", "-O1" }); + trace_writer_build.addFileArg(b.path("zig/src/families/flashnext/exl3/dplan/exl3_dplan_trace_kernel_test.cpp")); + trace_writer_build.addArg("-o"); + const trace_writer_binary = trace_writer_build.addOutputFileArg("dplan-trace-writer"); + trace_writer_build.addFileInput(b.path("zig/kernels/cuda/flashnext_exl3/flashnext_exl3_dplan.cu")); + const trace_writer_run = std.Build.Step.Run.create(b, "dplan-trace-writer"); + trace_writer_run.addFileArg(trace_writer_binary); + const trace_fixture = trace_writer_run.addOutputFileArg("demand-v4.exdpt"); + const checkpoint_fixture = trace_writer_run.addOutputFileArg("demand-v4-checkpoints.exdpt"); + const trace_writer_tests = b.createModule(.{ .root_source_file = b.path("zig/src/families/flashnext/exl3/dplan/exl3_dplan_trace_writer_test.zig"), .target = host, .optimize = .debug }); + trace_writer_tests.addAnonymousImport("dplan_trace_fixture", .{ .root_source_file = trace_fixture }); + trace_writer_tests.addAnonymousImport("dplan_checkpoint_fixture", .{ .root_source_file = checkpoint_fixture }); + b.step("test-flashnext-exl3-trace-writer-host", "CPU-only CUDA planner writer parity and exact Zig replay; requires c++").dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = trace_writer_tests })).step); +} diff --git a/zig/kernels/cuda/flashnext_exl3/flashnext_exl3_dplan.cu b/zig/kernels/cuda/flashnext_exl3/flashnext_exl3_dplan.cu new file mode 100644 index 0000000000..e7bb1df06e --- /dev/null +++ b/zig/kernels/cuda/flashnext_exl3/flashnext_exl3_dplan.cu @@ -0,0 +1,430 @@ +// The Flash Next family's EXL3 device-side pool plan (cuda_exl3_dplan.zig): one layer's offload round planned on the + +#include + +#define DP_EXPERTS 512 +#define DP_SHARED 512 +#define DP_MAXU 256 +#define DP_RECENT_WINDOW 256 +#define DP_RECENT_WORDS_PER_ROW (DP_EXPERTS / 32) +#define DP_RECENT_WORDS (DP_RECENT_WINDOW * DP_RECENT_WORDS_PER_ROW) + +struct DpMisc { + unsigned now, rounds, hit_slots, miss_slots, copies, evictions, overflow, inplace, recent_cursor, recent_filled; +}; + +struct DpTraceHeader { + unsigned magic, version, pool_index, layer, lines, max_rows, seed_words, record_words; + unsigned capacity, recorded, dropped_lo, dropped_hi, discontinuity, invalid_launch, checkpoints, dropped_checkpoints; +}; + +static_assert(sizeof(DpTraceHeader) == 64, "trace header ABI"); +static_assert(sizeof(DpMisc) == 40, "planner misc ABI"); + +static constexpr unsigned DP_TRACE_WORDS = 1024 * 1024 / 4; + +__device__ __forceinline__ bool dp_trace_valid(DpTraceHeader* header, int lines) { + const unsigned long long state_words = 2ull * static_cast(lines) + 3 * DP_EXPERTS + DP_RECENT_WORDS + 10; + const unsigned long long record_words = 15ull + 11ull * header->max_rows; + const unsigned long long base = 16 + 2 * state_words; + const unsigned long long checkpoint_words = 4 + 2 * state_words; + if (lines <= 0 || header->magic != 0x44505431 || header->version != 4 || + header->lines != static_cast(lines) || header->max_rows == 0 || + base >= DP_TRACE_WORDS || record_words > DP_TRACE_WORDS - base || + header->seed_words != state_words || header->record_words != record_words || + header->capacity != (DP_TRACE_WORDS - base) / record_words || header->recorded > header->capacity || + header->discontinuity > 1 || header->invalid_launch > 1 || + header->checkpoints > (DP_TRACE_WORDS - base) / checkpoint_words || + base + static_cast(header->recorded) * record_words > + DP_TRACE_WORDS - static_cast(header->checkpoints) * checkpoint_words) { + header->invalid_launch = 1; + return false; + } + return true; +} + +extern "C" __global__ void fn_exl3_dcheckpoint(const unsigned* st, int lines, unsigned* trace, int transfer) { + if (trace == nullptr || threadIdx.x != 0 || blockIdx.x != 0) return; + auto* header = reinterpret_cast(trace); + if (!dp_trace_valid(header, lines)) return; + if (transfer != 0 && transfer != 1) { + header->invalid_launch = 1; + return; + } + unsigned* seed = trace + 16; + unsigned* terminal = seed + header->seed_words; + const bool initialized = header->recorded != 0 || header->checkpoints != 0; + if (transfer == 0) { + if (initialized) + for (unsigned index = 0; index < header->seed_words; ++index) + if (terminal[index] != st[index]) header->discontinuity = 1; + return; + } + const unsigned stride = 4 + 2 * header->seed_words; + const unsigned end = 16 + 2 * header->seed_words + header->recorded * header->record_words; + const unsigned start = DP_TRACE_WORDS - header->checkpoints * stride; + if (start - end < stride || header->dropped_lo != 0 || header->dropped_hi != 0 || header->dropped_checkpoints != 0) { + if (header->dropped_checkpoints != 0xffffffffu) ++header->dropped_checkpoints; + return; + } + unsigned* checkpoint = trace + start - stride; + checkpoint[0] = header->recorded; + checkpoint[1] = 1; + checkpoint[2] = header->seed_words; + checkpoint[3] = header->checkpoints; + for (unsigned index = 0; index < header->seed_words; ++index) { + if (!initialized) seed[index] = terminal[index] = st[index]; + checkpoint[4 + index] = terminal[index]; + checkpoint[4 + header->seed_words + index] = st[index]; + terminal[index] = st[index]; + } + ++header->checkpoints; +} + +__device__ __forceinline__ unsigned* dp_trace_begin(const int* picks, int n, unsigned* trace, + const unsigned* st, int recent_enabled, int lines) { + if (trace == nullptr) return nullptr; + DpTraceHeader* header = reinterpret_cast(trace); + if (!dp_trace_valid(header, lines)) return nullptr; + if (n <= 0 || n % 11 != 0 || static_cast(n / 11) > header->max_rows) { + header->invalid_launch = 1; + return nullptr; + } + const unsigned end = 16 + 2 * header->seed_words + header->recorded * header->record_words; + const unsigned start = DP_TRACE_WORDS - header->checkpoints * (4 + 2 * header->seed_words); + if (start - end < header->record_words || header->dropped_checkpoints != 0) { + unsigned long long dropped = static_cast(header->dropped_hi) << 32 | header->dropped_lo; + if (dropped != ~0ull) ++dropped; + header->dropped_lo = static_cast(dropped); + header->dropped_hi = static_cast(dropped >> 32); + return nullptr; + } + unsigned* seed = trace + 16; + unsigned* terminal = seed + header->seed_words; + for (unsigned index = 0; index < header->seed_words; ++index) { + if (header->recorded == 0 && header->checkpoints == 0) seed[index] = st[index]; + else if (terminal[index] != st[index]) header->discontinuity = 1; + } + unsigned* round = terminal + header->seed_words + header->recorded * header->record_words; + round[0] = static_cast(n / 11); + round[1] = 0; + round[2] = 0; + round[3] = static_cast(recent_enabled); + unsigned seen[DP_EXPERTS / 32] = {}; + unsigned* raw = round + 15; + for (int pick = 0; pick < n; ++pick) { + const int expert = picks[pick]; + raw[pick] = static_cast(expert); + if (expert < 0 || expert >= DP_SHARED) continue; + round[1] += 1; + const unsigned bit = 1u << (expert & 31); + if ((seen[expert >> 5] & bit) == 0) round[2] += 1; + seen[expert >> 5] |= bit; + } + return round; +} + +__device__ __forceinline__ void dp_trace_end(unsigned* trace, unsigned* round, const unsigned* st, + const DpMisc* misc, const int* picks, int n, + const int* ovf, int n_ovf) { + if (round == nullptr) return; + DpTraceHeader* header = reinterpret_cast(trace); + const unsigned* counters = reinterpret_cast(misc); + for (unsigned index = 0; index < 10; ++index) round[4 + index] = counters[index]; + unsigned mapped_experts[DP_EXPERTS / 32] = {}; + for (int mapped = 0; mapped < n_ovf; ++mapped) { + const int expert = ovf[mapped]; + mapped_experts[expert >> 5] |= 1u << (expert & 31); + } + unsigned mapped_picks = 0; + for (int pick = 0; pick < n; ++pick) { + const int expert = picks[pick]; + if (expert >= 0 && expert < DP_SHARED && (mapped_experts[expert >> 5] & (1u << (expert & 31)))) + ++mapped_picks; + } + round[14] = mapped_picks; + unsigned* terminal = trace + 16 + header->seed_words; + for (unsigned index = 0; index < header->seed_words; ++index) terminal[index] = st[index]; + header->recorded += 1; +} + +__device__ __forceinline__ bool dp_colder(unsigned ca, unsigned sa, int ia, unsigned cb, unsigned sb, int ib) { + if (ca != cb) return ca < cb; + if (sa != sb) return sa < sb; + return ia > ib; +} + +// One block of 256 threads. picks [n] int32 (-1 skipped, 512 the shared expert, masked to the rank's ids at tp=2). +extern "C" __global__ void __launch_bounds__(256) fn_exl3_dplan( + const int* __restrict__ picks, int n, unsigned* __restrict__ st, int lines, int recent_enabled, long long* __restrict__ tab0, + long long* __restrict__ tab1, long long* __restrict__ tab2, int* __restrict__ k2t0, int* __restrict__ k2t1, + int* __restrict__ k2t2, const long long* __restrict__ src, const long long* __restrict__ lens, + const int* __restrict__ k2, unsigned long long line_base, unsigned long long line_bytes, unsigned long long off0, + unsigned long long off1, unsigned long long off2, long long* __restrict__ work, int* __restrict__ work_n, + int* __restrict__ hit_picks, int* __restrict__ miss_picks, int* __restrict__ ovf, int* __restrict__ ovf_n, + unsigned* __restrict__ trace) { + int* line_id = reinterpret_cast(st); + unsigned* stamp = st + lines; + int* line_of = reinterpret_cast(st + 2 * lines); + unsigned* counts = st + 2 * lines + DP_EXPERTS; + unsigned* recent_counts = counts + DP_EXPERTS; + unsigned* recent_rows = recent_counts + DP_EXPERTS; + DpMisc* misc = reinterpret_cast(recent_rows + DP_RECENT_WORDS); + const unsigned* score = recent_enabled ? recent_counts : counts; + + __shared__ unsigned short uniq[DP_MAXU]; + __shared__ unsigned short app[DP_MAXU]; + __shared__ unsigned short miss[DP_MAXU]; + __shared__ unsigned keep[DP_EXPERTS / 32]; + __shared__ int n_uniq, n_miss, n_copy, n_ovf; + __shared__ unsigned now; + __shared__ unsigned rkey_c[256], rkey_s[256]; + __shared__ int rkey_i[256]; + const int t = threadIdx.x; + unsigned* trace_round = nullptr; + + for (int word = t; word < DP_EXPERTS / 32; word += blockDim.x) keep[word] = 0; + __syncthreads(); + if (t == 0) { + trace_round = dp_trace_begin(picks, n, trace, st, recent_enabled, lines); + now = ++misc->now; + misc->rounds += 1; + if (recent_enabled) { + const int rows = n / 11; + for (int row = 0; row < rows; ++row) { + unsigned* expired = recent_rows + misc->recent_cursor * DP_RECENT_WORDS_PER_ROW; + for (int word = 0; word < DP_RECENT_WORDS_PER_ROW; ++word) { + unsigned removed = expired[word]; + while (removed != 0) { + const int bit = __ffs(removed) - 1; + const int expert = word * 32 + bit; + recent_counts[expert] -= 1; + removed &= removed - 1; + } + expired[word] = 0; + } + for (int slot = 0; slot < 11; ++slot) { + const int expert = picks[row * 11 + slot]; + if (expert < 0 || expert >= DP_SHARED) continue; + const unsigned bit = 1u << (expert & 31); + unsigned& word = expired[expert >> 5]; + if ((word & bit) == 0) { + word |= bit; + recent_counts[expert] += 1; + } + } + misc->recent_cursor = (misc->recent_cursor + 1) % DP_RECENT_WINDOW; + if (misc->recent_filled < DP_RECENT_WINDOW) misc->recent_filled += 1; + } + } + int nu = 0; + bool over = false; + for (int i = 0; i < n; ++i) { + const int p = picks[i]; + if (p < 0 || p >= DP_SHARED) continue; + // observe: the lifetime count (saturating) and this round's appearances + if (counts[p] != 0xffffffffu) counts[p] += 1; + if (line_of[p] >= 0) misc->hit_slots += 1; else misc->miss_slots += 1; + int j = 0; + while (j < nu && uniq[j] != p) ++j; + if (j == nu) { + if (nu == DP_MAXU) { over = true; continue; } + uniq[nu] = static_cast(p); + app[nu] = 0; + keep[p >> 5] |= 1u << (p & 31); + ++nu; + } + app[j] += 1; + } + // held experts are touched; misses keep first-appearance order, then sort hottest first (insertion sort) + int nm = 0; + for (int j = 0; j < nu; ++j) { + const int id = uniq[j]; + const int ln = line_of[id]; + if (ln >= 0) { + stamp[ln] = now; + } else { + miss[nm] = static_cast(j); // index into uniq/app + ++nm; + } + } + for (int a = 1; a < nm; ++a) { + const unsigned short x = miss[a]; + int b = a - 1; + while (b >= 0) { + const unsigned short y = miss[b]; + const unsigned cx = score[uniq[x]], cy = score[uniq[y]]; + const bool hotter = cx != cy ? cx > cy : (app[x] != app[y] ? app[x] > app[y] : uniq[x] < uniq[y]); + if (!hotter) break; + miss[b + 1] = y; + --b; + } + miss[b + 1] = x; + } + if (over) misc->overflow += 1; // more distinct experts than DP_MAXU: the caller's row bound prevents it + n_uniq = nu; + n_miss = over ? 0 : nm; + n_copy = 0; + n_ovf = 0; + } + __syncthreads(); + // the split (optional): picks held when the round started (and the shared expert) vs the misses this round copies. + if (hit_picks != nullptr) { + for (int i = t; i < n; i += blockDim.x) { + const int p = picks[i]; + const bool held = p >= DP_SHARED || (p >= 0 && line_of[p] >= 0); + hit_picks[i] = p < 0 ? -1 : (held ? p : -1); + miss_picks[i] = p < 0 ? -1 : (held ? -1 : p); + } + __syncthreads(); + } + const unsigned long long offs[3] = {off0, off1, off2}; + long long* tabs[3] = {tab0, tab1, tab2}; + int* k2s[3] = {k2t0, k2t1, k2t2}; + for (int m = 0; m < n_miss; ++m) { + // victim: the first empty line, else the coldest resident not in this round (count, stamp, higher line) + unsigned bc = 0xffffffffu, bs = 0xffffffffu; + int bi = -1; + bool bempty = false; + for (int i = t; i < lines; i += blockDim.x) { + const int id = line_id[i]; + if (id < 0) { + if (!bempty || i < bi) { bempty = true; bi = i; } + continue; + } + if (bempty) continue; + if (keep[id >> 5] & (1u << (id & 31))) continue; + const unsigned c = score[id], s = stamp[i]; + if (bi < 0 || dp_colder(c, s, i, bc, bs, bi)) { bc = c; bs = s; bi = i; } + } + // empty lines win (encoded as count 0 with a flag in the stamp key's top bit), lowest index among them + rkey_c[t] = bempty ? 0u : bc; + rkey_s[t] = bempty ? 0u : bs; + rkey_i[t] = bempty ? -2 - bi : bi; // -2 - i < -1 marks an empty line i + __syncthreads(); + for (int w = blockDim.x / 2; w > 0; w >>= 1) { + if (t < w) { + const int ia = rkey_i[t], ib = rkey_i[t + w]; + bool take_b; + if (ib == -1) take_b = false; + else if (ia == -1) take_b = true; + else if (ia <= -2 || ib <= -2) { + // an empty line beats a resident one; between empties the lower line + if (ia <= -2 && ib <= -2) take_b = (-2 - ib) < (-2 - ia); + else take_b = ib <= -2; + } else take_b = dp_colder(rkey_c[t + w], rkey_s[t + w], ib, rkey_c[t], rkey_s[t], ia); + if (take_b) { + rkey_c[t] = rkey_c[t + w]; + rkey_s[t] = rkey_s[t + w]; + rkey_i[t] = ib; + } + } + __syncthreads(); + } + if (t == 0) { + const int r = rkey_i[0]; + const int v = r <= -2 ? -2 - r : r; + const int id = uniq[miss[m]]; + if (v >= 0) { + const int old = line_id[v]; + if (old >= 0) { + line_of[old] = -1; + for (int w = 0; w < 3; ++w) { + tabs[w][old] = 0; + k2s[w][old] = 0; + } + misc->evictions += 1; + } + line_id[v] = id; + line_of[id] = v; + stamp[v] = now; + const unsigned long long base = line_base + static_cast(v) * line_bytes; + long long* wk = work + static_cast(n_copy) * 9; + for (int w = 0; w < 3; ++w) { + tabs[w][id] = static_cast(base + offs[w]); + k2s[w][id] = k2[id * 3 + w]; + wk[w * 3 + 0] = src[id * 3 + w]; + wk[w * 3 + 1] = static_cast(base + offs[w]); + wk[w * 3 + 2] = lens[id * 3 + w]; + } + n_copy += 1; + misc->copies += 1; + } else { + // no line is free of this round's experts: the expert is read in place from the mapped host bank. + for (int w = 0; w < 3; ++w) { + tabs[w][id] = src[id * 3 + w]; + k2s[w][id] = k2[id * 3 + w]; + } + ovf[n_ovf] = id; + n_ovf += 1; + misc->inplace += 1; + } + } + __syncthreads(); + } + if (t == 0) { + *work_n = n_copy; + *ovf_n = n_ovf; + dp_trace_end(trace, trace_round, st, misc, picks, n, ovf, n_ovf); + } +} + +// After the round's experts: the in-place (host-bank) experts' table entries cleared again (they hold no line). +extern "C" __global__ void fn_exl3_drestore(const int* __restrict__ ovf, const int* __restrict__ ovf_n, + long long* __restrict__ tab0, long long* __restrict__ tab1, + long long* __restrict__ tab2, int* __restrict__ k2t0, int* __restrict__ k2t1, + int* __restrict__ k2t2) { + const int n = *ovf_n; + for (int i = threadIdx.x; i < n; i += blockDim.x) { + const int id = ovf[i]; + tab0[id] = 0; + tab1[id] = 0; + tab2[id] = 0; + k2t0[id] = 0; + k2t1[id] = 0; + k2t2[id] = 0; + } +} + +// The planned copies: each segment {src (host bank, mapped), dst (a line), len (bytes, a multiple of 16)} read as +extern "C" __global__ void __launch_bounds__(256) fn_exl3_dcopy(const long long* __restrict__ work, + const int* __restrict__ work_n) { + const int segs = *work_n * 3; + const long long tid = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const long long stride = static_cast(gridDim.x) * blockDim.x; + for (int g = 0; g < segs; ++g) { + const uint4* __restrict__ s = reinterpret_cast(work[g * 3 + 0]); + uint4* __restrict__ d = reinterpret_cast(work[g * 3 + 1]); + const long long words = work[g * 3 + 2] / 16; + long long i = tid; + for (; i + 3 * stride < words; i += 4 * stride) { + const uint4 a = s[i], b = s[i + stride], c = s[i + 2 * stride], e = s[i + 3 * stride]; + d[i] = a; + d[i + stride] = b; + d[i + 2 * stride] = c; + d[i + 3 * stride] = e; + } + for (; i < words; i += stride) d[i] = s[i]; + } +} + +// The draft chain on the device (cuda_exl3_chain.zig): fn_row_best's record [best, log-sum-exp, column bits, -] of a +extern "C" __global__ void fn_exl3_draft_token(const float* __restrict__ best, const unsigned* __restrict__ ids, + int* __restrict__ token) { + if (threadIdx.x == 0 && blockIdx.x == 0) token[0] = static_cast(ids[__float_as_int(best[2])]); +} + +// The chain's step decision on the device: the next token as fn_exl3_draft_token, then upstream's stop rule +extern "C" __global__ void fn_exl3_draft_next(const float* __restrict__ best, const unsigned* __restrict__ ids, + int* __restrict__ token, int* __restrict__ alive, + int* __restrict__ alive_out, cudaGraphConditionalHandle handle, + double confidence) { + if (threadIdx.x != 0 || blockIdx.x != 0) return; + token[0] = static_cast(ids[__float_as_int(best[2])]); + const double prob = exp(static_cast(best[0]) - static_cast(best[1])); + const int low = confidence > 0.0 && prob < confidence; + const int a = alive[0] && !low; + alive[0] = a; + alive_out[0] = a; + cudaGraphSetConditional(handle, static_cast(a)); +} diff --git a/zig/src/core/contract/flashnext_exl3_pending.zig b/zig/src/core/contract/flashnext_exl3_pending.zig index 1f7461dcb7..6f84095cf6 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 (pick mask, prefill tiles, host-tier K/V); nvcc 13.4 measured, older pending. +//! Flash Next EXL3 launches (pick mask, pool plan, prefill tiles, host-tier K/V); nvcc 13.4 measured, older pending. const std = @import("std"); const limits = @import("limits.zig"); const registry = @import("registry.zig"); @@ -9,6 +9,12 @@ 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_dplan fn_exl3_dcheckpoint", .symbol = "fn_exl3_dcheckpoint", .image = "flashnext_exl3_dplan", .threads = 1 }, .usage = registry.none }, + .{ .launch = .{ .name = "flashnext_exl3_dplan fn_exl3_dplan", .symbol = "fn_exl3_dplan", .image = "flashnext_exl3_dplan", .threads = 256 }, .usage = registry.none }, + .{ .launch = .{ .name = "flashnext_exl3_dplan fn_exl3_drestore", .symbol = "fn_exl3_drestore", .image = "flashnext_exl3_dplan", .threads = 256 }, .usage = registry.none }, + .{ .launch = .{ .name = "flashnext_exl3_dplan fn_exl3_dcopy", .symbol = "fn_exl3_dcopy", .image = "flashnext_exl3_dplan", .threads = 256 }, .usage = registry.none }, + .{ .launch = .{ .name = "flashnext_exl3_dplan fn_exl3_draft_token", .symbol = "fn_exl3_draft_token", .image = "flashnext_exl3_dplan", .threads = 32 }, .usage = registry.none }, + .{ .launch = .{ .name = "flashnext_exl3_dplan fn_exl3_draft_next", .symbol = "fn_exl3_draft_next", .image = "flashnext_exl3_dplan", .threads = 32 }, .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 }, @@ -29,6 +35,12 @@ 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_dplan fn_exl3_dcheckpoint", .usage = .{ .{ .regs = 32, .static = 0, .bound = 0 }, .{ .regs = 34, .static = 0, .bound = 0 }, .{ .regs = 34, .static = 0, .bound = 0 }, .{ .regs = 34, .static = 0, .bound = 0 }, .{ .regs = 34, .static = 0, .bound = 0 } } }, + .{ .name = "flashnext_exl3_dplan fn_exl3_dplan", .usage = .{ .{ .regs = 32, .static = 4688, .bound = 256 }, .{ .regs = 40, .static = 4688, .bound = 256 }, .{ .regs = 40, .static = 4688, .bound = 256 }, .{ .regs = 48, .static = 4688, .bound = 256 }, .{ .regs = 48, .static = 4688, .bound = 256 } } }, + .{ .name = "flashnext_exl3_dplan fn_exl3_drestore", .usage = .{ .{ .regs = 20, .static = 0, .bound = 0 }, .{ .regs = 20, .static = 0, .bound = 0 }, .{ .regs = 20, .static = 0, .bound = 0 }, .{ .regs = 34, .static = 0, .bound = 0 }, .{ .regs = 34, .static = 0, .bound = 0 } } }, + .{ .name = "flashnext_exl3_dplan fn_exl3_dcopy", .usage = .{ .{ .regs = 40, .static = 0, .bound = 256 }, .{ .regs = 38, .static = 0, .bound = 256 }, .{ .regs = 38, .static = 0, .bound = 256 }, .{ .regs = 40, .static = 0, .bound = 256 }, .{ .regs = 40, .static = 0, .bound = 256 } } }, + .{ .name = "flashnext_exl3_dplan fn_exl3_draft_token", .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_dplan fn_exl3_draft_next", .usage = .{ .{ .regs = 24, .static = 0, .bound = 0 }, .{ .regs = 24, .static = 0, .bound = 0 }, .{ .regs = 24, .static = 0, .bound = 0 }, .{ .regs = 24, .static = 0, .bound = 0 }, .{ .regs = 24, .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 } } }, diff --git a/zig/src/cuda/abi.zig b/zig/src/cuda/abi.zig index 10747bd88e..20846b7910 100644 --- a/zig/src/cuda/abi.zig +++ b/zig/src/cuda/abi.zig @@ -20,6 +20,7 @@ pub const error_not_ready: Result = 600; pub const stream_non_blocking: c_uint = 1; pub const event_disable_timing: c_uint = 2; +pub const event_wait_external: c_uint = 1; pub const host_alloc_portable: c_uint = 1; pub const host_alloc_devicemap: c_uint = 2; pub const host_register_devicemap: c_uint = 2; diff --git a/zig/src/cuda/kernels.zig b/zig/src/cuda/kernels.zig index 32aa9d4077..581c25b220 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_dplan: []const u8 = if (available) &Blob("fatbin_flashnext_exl3_dplan").bytes else &.{}; pub const flashnext_exl3_kvtier: []const u8 = if (available) &Blob("fatbin_flashnext_exl3_kvtier").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 &.{}; diff --git a/zig/src/cuda/stream.zig b/zig/src/cuda/stream.zig index 1733965193..bb7fbda772 100644 --- a/zig/src/cuda/stream.zig +++ b/zig/src/cuda/stream.zig @@ -54,6 +54,10 @@ pub const Stream = struct { pub fn wait(self: Stream, event: Event) Error!void { try self.d.check(self.d.api.cuStreamWaitEvent(self.handle, event.handle, 0), "cuStreamWaitEvent"); } + + pub fn waitExternal(self: Stream, event: Event) Error!void { + try self.d.check(self.d.api.cuStreamWaitEvent(self.handle, event.handle, abi.event_wait_external), "cuStreamWaitEvent"); + } }; pub const Event = struct { diff --git a/zig/src/families/flashnext/cuda_batch_state.zig b/zig/src/families/flashnext/cuda_batch_state.zig index 00a69045b5..aaeb4dbc0d 100644 --- a/zig/src/families/flashnext/cuda_batch_state.zig +++ b/zig/src/families/flashnext/cuda_batch_state.zig @@ -130,12 +130,18 @@ comptime { /// A caller's view into the pass. Without any, the pass is the server's prompt pass. pub const Hooks = struct { + pub const OutputTail = struct { + ptr: *anyopaque, + enqueue: *const fn (*anyopaque, cuda.Stream) anyerror!void, + }; + logits: Logits = .last, /// Device [rows, vocab] bf16 for `.all` (else the pass's own buffer for the last row). logits_out: u64 = 0, keep: ?*Keep = null, /// Device [rows, wide] bf16: every row's streams after the last MoE write-back (the MTP head's input); 0: none. streams: u64 = 0, + output_tail: ?OutputTail = null, /// A training pass's adds to each layer's MoE output, in place of the model's own (Model.pocket). moe_add: ?MoeAdd = null, }; @@ -143,24 +149,7 @@ pub const Hooks = struct { /// The kernels the pass launches besides the captured Triton ones, loaded once. pub const Kernels = @import("cuda_batch_kernels.zig").Kernels; -/// The compute stream as a pass's pinned staging uses it: copies from it, an event behind them, and a drain. -pub const Uploads = struct { - stream: cuda.Stream, - event: cuda.Event, - - pub fn copy(u: Uploads, dst: u64, bytes: []const u8) !void { - try cuda.copy_source.enqueue(u.stream.d, dst, bytes, u.stream); - } - pub fn record(u: Uploads) !void { - try u.event.record(u.stream); - } - pub fn wait(u: Uploads) !void { - try u.event.synchronize(); - } - pub fn drain(u: Uploads) !void { - try u.stream.synchronize(); - } -}; +pub const Uploads = @import("cuda_batch_uploads.zig").Uploads; /// Pinned bytes' uploads and the fence that says when the bytes may change (`Q`: `Uploads`, or a test's stand-in). pub fn Staging(comptime Q: type) type { @@ -512,7 +501,10 @@ pub const Lane = struct { seq: SeqState, pos: usize, rows: usize = 1, keep: ?*co /// The engine's resident weights the pass reads. pub const Model = struct { + prefill: ?*@import("exl3/prefill/cuda_exl3_prefill_state.zig").Runtime = null, collective: ?*@import("cuda_collective.zig").Collective = null, + /// The engine's device planner state (EXL3 offload); null: the host cut. + dplan: ?*@import("exl3/dplan/cuda_exl3_dplan.zig").State = null, table_w: u64, table_s: u64, table_b: u64, diff --git a/zig/src/families/flashnext/cuda_batch_uploads.zig b/zig/src/families/flashnext/cuda_batch_uploads.zig new file mode 100644 index 0000000000..5a59c2e193 --- /dev/null +++ b/zig/src/families/flashnext/cuda_batch_uploads.zig @@ -0,0 +1,21 @@ +const cuda = @import("cuda"); + +/// The compute stream as a pass's pinned staging uses it: copies from it, an event behind them, and a drain. +pub const Uploads = struct { + stream: cuda.Stream, + event: cuda.Event, + + pub fn copy(u: Uploads, dst: u64, bytes: []const u8) !void { + if (@import("exl3/graph/cuda_exl3_graph.zig").capturing) return @import("exl3/graph/cuda_exl3_graph.zig").pinnedCopy(u.stream, dst, bytes); // a copy node reads the pinned bytes at each replay + try cuda.copy_source.enqueue(u.stream.d, dst, bytes, u.stream); + } + pub fn record(u: Uploads) !void { + try u.event.record(u.stream); + } + pub fn wait(u: Uploads) !void { + try u.event.synchronize(); + } + pub fn drain(u: Uploads) !void { + try u.stream.synchronize(); + } +}; diff --git a/zig/src/families/flashnext/cuda_collective.zig b/zig/src/families/flashnext/cuda_collective.zig index f7aec59250..b9c8ed32ff 100644 --- a/zig/src/families/flashnext/cuda_collective.zig +++ b/zig/src/families/flashnext/cuda_collective.zig @@ -10,6 +10,60 @@ const Advanced = struct { ncclCommGetAsyncError: *const fn (nccl.Comm, *nccl.Result) callconv(.c) nccl.Result, }; +pub const GraphState = enum(u32) { eager, capture, pending, ready, failed, timed_out }; +pub const GraphChoice = enum { eager, capture, replay }; +pub const GraphVote = struct { state: GraphState, cohort: u64 = 0 }; +pub const GraphDecision = struct { choice: GraphChoice, cohort: u64 }; + +pub fn selectGraphVictim(local: u64, peer: u64) u64 { + if (local == 0) return peer; + if (peer == 0) return local; + return @min(local, peer); +} + +test "exl3: selective retirement agrees stable union victims for asymmetric cohort sets" { + var left = [_]u64{ 3, 7, 11 }; + var right = [_]u64{ 5, 7, 11 }; + for ([_]u64{ 3, 5 }) |expected| { + var local: u64 = 0; + var peer: u64 = 0; + for (left) |cohort| local = selectGraphVictim(local, cohort); + for (right) |cohort| peer = selectGraphVictim(peer, cohort); + const victim = selectGraphVictim(local, peer); + try std.testing.expectEqual(expected, victim); + try std.testing.expectEqual(victim, selectGraphVictim(peer, local)); + for (&left) |*cohort| if (cohort.* == victim) { + cohort.* = 0; + }; + for (&right) |*cohort| if (cohort.* == victim) { + cohort.* = 0; + }; + } + try std.testing.expectEqual(@as(u64, 7), left[1]); + try std.testing.expectEqual(@as(u64, 11), right[2]); + try std.testing.expectEqual(@as(u64, 0), selectGraphVictim(0, 0)); +} + +pub fn chooseGraph(local: GraphVote, peer: GraphVote, step: u64) !GraphDecision { + if (local.state == .timed_out or peer.state == .timed_out) return error.GraphInstantiationTimeout; + if (local.state == .failed or peer.state == .failed) return error.GraphInstantiationFailed; + if (local.state == .eager or peer.state == .eager or local.state == .pending or peer.state == .pending) + return .{ .choice = .eager, .cohort = 0 }; + if (local.state == .ready and peer.state == .ready and local.cohort != 0 and local.cohort == peer.cohort) + return .{ .choice = .replay, .cohort = local.cohort }; + return .{ .choice = .capture, .cohort = step }; +} + +pub fn graphExchange(link: wire.Link, step: u64, signature: u64, vote: GraphVote, until: u64) !GraphDecision { + const magic: u32 = 0x47520000; + try link.sendHeader(.{ .kind = .ready, .sequence = step, .object = vote.cohort, .arg = signature, .flags = magic | @backingInt(vote.state) }, until); + const peer = try link.receiveHeader(until); + if (peer.kind != .ready or peer.sequence != step or peer.arg != signature or peer.bytes != 0 or peer.flags & 0xffff0000 != magic) + return error.GraphRendezvousMismatch; + const state = std.enums.fromInt(GraphState, peer.flags & 0xffff) orelse return error.GraphRendezvousMismatch; + return chooseGraph(vote, .{ .state = state, .cohort = peer.object }, step); +} + pub const Collective = struct { lib: nccl.Library, api: Advanced, @@ -20,6 +74,7 @@ pub const Collective = struct { 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, + graph_step: 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 { @@ -95,6 +150,34 @@ pub const Collective = struct { try c.wait(stream, wire.deadline(c.timeout_ms)); } + pub fn failGraph(c: *Collective) void { + c.abort(); + _ = std.c.shutdown(c.link.fd, std.posix.SHUT.RDWR); + } + + pub fn graphTimeoutMs(c: *const Collective) u64 { + return @min(c.timeout_ms, 30_000); + } + + pub fn graphAgree(c: *Collective, signature: u64, vote: GraphVote) !GraphDecision { + if (c.failed or c.comm == null) return error.RankDisconnected; + errdefer c.failGraph(); + c.graph_step = try std.math.add(u64, c.graph_step, 1); + return graphExchange(c.link, c.graph_step, signature, vote, wire.deadline(c.graphTimeoutMs())); + } + + pub fn graphVictim(c: *Collective, signature: u64, candidate: u64) !u64 { + if (c.failed or c.comm == null) return error.RankDisconnected; + errdefer c.failGraph(); + c.graph_step = try std.math.add(u64, c.graph_step, 1); + const until = wire.deadline(c.graphTimeoutMs()); + try c.link.sendHeader(.{ .kind = .ready, .sequence = c.graph_step, .object = candidate, .arg = signature, .flags = 0x47560000 }, until); + const peer = try c.link.receiveHeader(until); + if (peer.kind != .ready or peer.sequence != c.graph_step or peer.arg != signature or peer.bytes != 0 or peer.flags != 0x47560000) + return error.GraphRendezvousMismatch; + return selectGraphVictim(candidate, peer.object); + } + /// Disjoint BF16 vocabulary columns are concatenated without an arithmetic reduction. pub fn gather(c: *Collective, stream: cuda.Stream, local: u64, out: u64, count: usize) !void { if (c.failed or c.comm == null) return error.RankDisconnected; @@ -150,3 +233,119 @@ pub const Collective = struct { try c.wait(null, wire.deadline(c.timeout_ms)); } }; + +test "exl3: graph readiness requires matching completed capture cohorts" { + const states = std.enums.values(GraphState); + for (states) |local| for (states) |peer| { + const left: GraphVote = .{ .state = local, .cohort = 7 }; + const right: GraphVote = .{ .state = peer, .cohort = 7 }; + if (local == .timed_out or peer == .timed_out) { + try std.testing.expectError(error.GraphInstantiationTimeout, chooseGraph(left, right, 10)); + } else if (local == .failed or peer == .failed) { + try std.testing.expectError(error.GraphInstantiationFailed, chooseGraph(left, right, 10)); + } else { + const choice: GraphChoice = if (local == .eager or peer == .eager or local == .pending or peer == .pending) + .eager + else if (local == .ready and peer == .ready) .replay else .capture; + const decision = try chooseGraph(left, right, 10); + try std.testing.expectEqual(choice, decision.choice); + try std.testing.expectEqualDeep(decision, try chooseGraph(right, left, 10)); + } + }; + try std.testing.expectEqual(GraphChoice.capture, (try chooseGraph(.{ .state = .ready, .cohort = 7 }, .{ .state = .ready, .cohort = 8 }, 10)).choice); +} + +fn graphTestPair() ![2]std.posix.socket_t { + var pair: [2]std.posix.socket_t = undefined; + if (std.c.socketpair(std.posix.AF.UNIX, std.posix.SOCK.STREAM, 0, &pair) != 0) return error.SocketPair; + errdefer for (pair) |fd| { + _ = std.c.close(fd); + }; + const nonblock: c_int = @bitCast(@as(u32, @bitCast(std.c.O{ .NONBLOCK = true }))); + for (pair) |fd| { + const flags = std.c.fcntl(fd, std.c.F.GETFL); + if (flags < 0 or std.c.fcntl(fd, std.c.F.SETFL, flags | nonblock) < 0) return error.Nonblocking; + } + return pair; +} + +test "exl3: graph rendezvous keeps a delayed rank eager until both workers finish" { + if (@import("builtin").os.tag == .windows) return error.SkipZigTest; // The rank wire is POSIX-only. + const pair = try graphTestPair(); + defer for (pair) |fd| { + _ = std.c.close(fd); + }; + const Peer = struct { + fn run(link: wire.Link) void { + const votes = [_]GraphVote{ .{ .state = .pending, .cohort = 3 }, .{ .state = .ready, .cohort = 3 }, .{ .state = .ready, .cohort = 3 } }; + for (votes, 1..) |vote, step| { + const decision = graphExchange(link, step, 42, vote, wire.deadline(1000)) catch @panic("peer rendezvous failed"); + std.debug.assert(decision.choice == (if (step < 3) GraphChoice.eager else GraphChoice.replay)); + } + } + }; + const thread = try std.Thread.spawn(.{}, Peer.run, .{wire.Link{ .fd = pair[1] }}); + defer thread.join(); + const votes = [_]GraphVote{ .{ .state = .ready, .cohort = 3 }, .{ .state = .pending, .cohort = 3 }, .{ .state = .ready, .cohort = 3 } }; + for (votes, 1..) |vote, step| { + const decision = try graphExchange(.{ .fd = pair[0] }, step, 42, vote, wire.deadline(1000)); + try std.testing.expectEqual(if (step < 3) GraphChoice.eager else GraphChoice.replay, decision.choice); + } +} + +test "exl3: graph rendezvous rejects stale frames failures timeouts and disconnected peers" { + if (@import("builtin").os.tag == .windows) return error.SkipZigTest; // The rank wire is POSIX-only. + const pair = try graphTestPair(); + defer for (pair) |fd| { + _ = std.c.close(fd); + }; + const local: wire.Link = .{ .fd = pair[0] }; + const peer: wire.Link = .{ .fd = pair[1] }; + const vote: GraphVote = .{ .state = .ready, .cohort = 2 }; + for ([_]wire.Header{ + .{ .kind = .ack, .sequence = 1, .arg = 42, .flags = 0x47520003 }, + .{ .kind = .ready, .sequence = 2, .arg = 42, .flags = 0x47520003 }, + .{ .kind = .ready, .sequence = 1, .arg = 43, .flags = 0x47520003 }, + .{ .kind = .ready, .sequence = 1, .arg = 42, .flags = 0x475200ff }, + }) |header| { + try peer.sendHeader(header, wire.deadline(1000)); + try std.testing.expectError(error.GraphRendezvousMismatch, graphExchange(local, 1, 42, vote, wire.deadline(1000))); + _ = try peer.receiveHeader(wire.deadline(1000)); + } + try peer.sendHeader(.{ .kind = .ready, .sequence = 1, .arg = 42, .flags = 0x47520004 }, wire.deadline(1000)); + try std.testing.expectError(error.GraphInstantiationFailed, graphExchange(local, 1, 42, vote, wire.deadline(1000))); + _ = try peer.receiveHeader(wire.deadline(1000)); + try std.testing.expectError(error.RankTimeout, graphExchange(local, 1, 42, vote, wire.deadline(5))); + _ = std.c.shutdown(peer.fd, std.posix.SHUT.RDWR); + try std.testing.expectError(error.RankDisconnected, graphExchange(local, 1, 42, vote, wire.deadline(1000))); +} + +test "exl3: graph rendezvous failure aborts the communicator and disconnects its control peer" { + if (@import("builtin").os.tag == .windows) return error.SkipZigTest; // The rank wire is POSIX-only. + const pair = try graphTestPair(); + defer for (pair) |fd| { + _ = std.c.close(fd); + }; + const Abort = struct { + var calls: usize = 0; + fn run(_: nccl.Comm) callconv(.c) nccl.Result { + calls += 1; + return 0; + } + }; + Abort.calls = 0; + var c: Collective = undefined; + c.comm = @ptrFromInt(1); + c.api.ncclCommAbort = Abort.run; + c.link = .{ .fd = pair[0] }; + c.failed = false; + c.timeout_ms = 5; + c.graph_step = 0; + try std.testing.expectError(error.RankTimeout, c.graphAgree(42, .{ .state = .pending })); + try std.testing.expect(c.failed and c.comm == null); + try std.testing.expectEqual(@as(usize, 1), Abort.calls); + try std.testing.expectError(error.RankDisconnected, c.graphAgree(42, .{ .state = .ready })); + const peer: wire.Link = .{ .fd = pair[1] }; + _ = try peer.receiveHeader(wire.deadline(1000)); + try std.testing.expectError(error.RankDisconnected, peer.checkPeer()); +} diff --git a/zig/src/families/flashnext/cuda_images_test.zig b/zig/src/families/flashnext/cuda_images_test.zig index a4a5b4a823..b81fc6a6d6 100644 --- a/zig/src/families/flashnext/cuda_images_test.zig +++ b/zig/src/families/flashnext/cuda_images_test.zig @@ -36,6 +36,8 @@ fn run(comptime Cuda: type, comptime open: anytype, comptime drop: anytype, draf var stream: cuda.Stream = .{ .d = &driver, .handle = null }; const f: cuda.Function = .{ .d = &driver, .handle = null }; var engine: prompt.Engine = undefined; + engine.exl3_prefill = .{}; + engine.exl3_chain = .{}; engine.gpa = a; engine.driver = &driver; engine.stream = &stream; @@ -134,6 +136,8 @@ fn release(comptime Cuda: type, comptime open: anytype) !void { driver.api.cuCtxSynchronize = Wait.context; var stream: cuda.Stream = .{ .d = &driver, .handle = null }; var engine: prompt.Engine = undefined; + engine.exl3_prefill = .{}; + engine.exl3_chain = .{}; engine.gpa = a; engine.driver = &driver; engine.stream = &stream; diff --git a/zig/src/families/flashnext/cuda_lanes.zig b/zig/src/families/flashnext/cuda_lanes.zig index 1939d049d1..defa83b400 100644 --- a/zig/src/families/flashnext/cuda_lanes.zig +++ b/zig/src/families/flashnext/cuda_lanes.zig @@ -125,8 +125,9 @@ pub fn Namespace(comptime Tri: type) type { pub fn rankRoom(ptr: *anyopaque, local: u64) !u64 { 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; + const room = @import("exl3/pack/exl3_startup_memory.zig").laneRoom(local, own.graph_reservation); + if (own.ranks) |r| return if (own.exl3) @import("exl3/kv/cuda_exl3_kv.zig").rankRoom(r, room, own.stream_bytes) else r.room(room); + return room; } /// A request or load this engine refuses, in words; null: none of its own. @@ -171,6 +172,7 @@ pub fn Namespace(comptime Tri: type) type { /// An EXL3 pack: its ranks agree their room in streams of stream_bytes and keep drafting states. exl3: bool = false, stream_bytes: usize = 0, + graph_reservation: u64 = 0, /// --slide: the pocket's layout, from config.json at load. lay: living.Layout = undefined, kv: kvl.State = .{}, @@ -182,6 +184,7 @@ pub fn Namespace(comptime Tri: type) type { pub fn ops(self: *Cuda) Ops { return .{ .e = self.engine, .tri = self.tri(), .h = self.head.?, .main = if (self.ranks != null) .{ .ptr = self, + .rank = self.ranks.?.rank, .forward = mainForward, .verify = mainVerify, .rollback = mainRollback, @@ -193,9 +196,9 @@ pub fn Namespace(comptime Tri: type) type { return self.ranks.?.forward(Tri, self.tri(), q, ids, hooks); } - fn mainVerify(ptr: *anyopaque, q: *prompt.Seq, d: *mtp.Draft, ids: []const u32) anyerror!void { + fn mainVerify(ptr: *anyopaque, q: *prompt.Seq, d: *mtp.Draft, ids: []const u32, tail: ?prompt.batch.Hooks.OutputTail) anyerror!void { const self = of(ptr); - return self.ranks.?.verify(Tri, self.tri(), q, d, ids); + return self.ranks.?.verifyWithTail(Tri, self.tri(), q, d, ids, tail); } fn mainRollback(ptr: *anyopaque, q: *prompt.Seq, d: *mtp.Draft, keep: usize) anyerror!void { @@ -286,8 +289,8 @@ pub fn Namespace(comptime Tri: type) type { s.proposer = null; } // a window's rows past the reply's last token (its drafts) still need cache slots - const cap = @min(ids.len + s.max_new + (if (drafting) mtp.max_rows else 0), prompt.max_tokens); const from = self.kv.choose(); + const cap = @import("exl3/offload/cuda_exl3_recycle.zig").capFor(from == .pool, ids.len + s.max_new + (if (drafting) mtp.max_rows else 0), self.context + mtp.max_rows); gop.value_ptr.* = .{ .seq = (if (self.ranks) |r| r.createOn(cap, from) else if (from == .pool) self.kv.pool.?.open(cap) else self.engine.newSeq(cap)) catch |err| { self.lanes.removeByPtr(gop.key_ptr); return err; @@ -384,7 +387,9 @@ pub fn Namespace(comptime Tri: type) type { while (it.next()) |l| dropLane(own, l); own.lanes.deinit(own.gpa); own.engine.round_graph.clear(); + @import("exl3/graph/cuda_exl3_graph.zig").releaseFor(own.engine); own.logits.free(); + @import("exl3/model/cuda_exl3_chain.zig").releaseRuntime(&own.engine.exl3_chain); if (own.head) |h| h.deinit(); if (own.vision) |v| v.deinit(); if (own.ranks) |r| r.deinit(); @@ -399,6 +404,7 @@ pub fn Namespace(comptime Tri: type) type { fn dropLane(self: *Cuda, l: *Lane) void { cuda.completion.drain(self.engine.driver, self.stream, "flashnext: a failed or closed lane did not drain"); + if (l.seq.paged != null) @import("exl3/graph/cuda_exl3_graph.zig").releaseDecodeSeq(&l.seq); if (l.draft) |*d| d.deinit(); if (self.ranks) |r| r.drop(&l.seq) else if (!self.kv.close(l.from, &l.seq)) l.seq.deinit(self.gpa); } diff --git a/zig/src/families/flashnext/cuda_lanes_admission_test.zig b/zig/src/families/flashnext/cuda_lanes_admission_test.zig index 898c0bd9c6..692b6ab886 100644 --- a/zig/src/families/flashnext/cuda_lanes_admission_test.zig +++ b/zig/src/families/flashnext/cuda_lanes_admission_test.zig @@ -11,6 +11,8 @@ fn run(comptime Cuda: type, comptime open: anytype, comptime drop: anytype, coun var driver = fixtures.fakeDriver(); var stream: cuda.Stream = .{ .d = &driver, .handle = null }; var engine: prompt.Engine = undefined; + engine.exl3_prefill = .{}; + engine.exl3_chain = .{}; engine.gpa = a; engine.driver = &driver; engine.stream = &stream; diff --git a/zig/src/families/flashnext/cuda_lanes_load.zig b/zig/src/families/flashnext/cuda_lanes_load.zig index b6e64da73a..7fc4b3e82c 100644 --- a/zig/src/families/flashnext/cuda_lanes_load.zig +++ b/zig/src/families/flashnext/cuda_lanes_load.zig @@ -112,7 +112,9 @@ pub fn open(comptime F: type, gpa: std.mem.Allocator, io: std.Io, ctx: *const cu 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); + var pool_budget = try exl3_pass.budgetForStreams(stream_bytes, o.lanes); + if (mode == .exl3 and @import("exl3/graph/cuda_exl3_graph.zig").enabled()) pool_budget.graph_reservation = try @import("exl3/graph/cuda_exl3_graph_memory.zig").startupReservationForStreams(o.lanes); + own.graph_reservation = pool_budget.graph_reservation; 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], own.head, @intCast(o.tp), @intCast(o.rank))) |prepared| { own.profile = prepared.state; @@ -120,7 +122,10 @@ pub fn open(comptime F: type, gpa: std.mem.Allocator, io: std.Io, ctx: *const cu } 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], if (own.head) |h| &h.layer else null, own.engine.chunk, pool_budget, profile_counts); + const attached = try exl3_pass.attach(gpa, io, driver, own.stream, &own.engine.exl3_prefill, own.engine.layers[0..own.engine.n_layers], if (own.head) |h| &h.layer else null, own.engine.chunk, pool_budget, profile_counts); + own.graph_reservation = attached.graph_reservation; + own.engine.exl3_dplan.mtp_pool = attached.mtp_pool; + if (own.graph_reservation != 0) std.log.info("flashnext rank {d}: pools and lanes keep {d} bytes free for the retained graphs (measured exec memory of a full run: under 64 MiB a stream; captures draw on the device reserve)", .{ o.rank, own.graph_reservation }); if (own.profile_counts) |*counts| try placement.finish(own, counts); const ledger: memory.Ledger = .{ .persistent = persistent_bytes, @@ -128,9 +133,9 @@ pub fn open(comptime F: type, gpa: std.mem.Allocator, io: std.Io, ctx: *const cu .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 }, + .reservation = .{ .stream_bytes = stream_bytes, .streams = if (stream_bytes == 0) 0 else pool_budget.other / stream_bytes, .graph_bytes = attached.graph_reservation, .after_pool_bytes = attached.after_pool_bytes }, }; - 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 }); + 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} graph_reservation={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, pool_budget.graph_reservation, try ledger.total(), attached.free_after }); return .{ .layout = own.kv.layout, .vision = own.vision, diff --git a/zig/src/families/flashnext/cuda_lanes_restore_test.zig b/zig/src/families/flashnext/cuda_lanes_restore_test.zig index aa441f6166..c340ba2e18 100644 --- a/zig/src/families/flashnext/cuda_lanes_restore_test.zig +++ b/zig/src/families/flashnext/cuda_lanes_restore_test.zig @@ -56,6 +56,8 @@ pub fn run(comptime Cuda: type, comptime open: anytype, comptime restore: anytyp var driver = Device.driver(); var stream: cuda.Stream = .{ .d = &driver, .handle = null }; var engine: prompt.Engine = undefined; + engine.exl3_prefill = .{}; + engine.exl3_chain = .{}; engine.gpa = a; engine.driver = &driver; engine.stream = &stream; diff --git a/zig/src/families/flashnext/cuda_launch.zig b/zig/src/families/flashnext/cuda_launch.zig index 9ba07c4106..0abc056c22 100644 --- a/zig/src/families/flashnext/cuda_launch.zig +++ b/zig/src/families/flashnext/cuda_launch.zig @@ -459,6 +459,7 @@ test { test "exl3: greedy verification host contracts" { _ = @import("cuda_mtp_ops.zig"); _ = @import("cuda_mtp_draft.zig"); + _ = @import("cuda_mtp_verify_cache.zig"); _ = @import("exl3/model/cuda_exl3_dispatch_test.zig"); } diff --git a/zig/src/families/flashnext/cuda_mtp_check.zig b/zig/src/families/flashnext/cuda_mtp_check.zig index 53e3b2d057..67a0d121ad 100644 --- a/zig/src/families/flashnext/cuda_mtp_check.zig +++ b/zig/src/families/flashnext/cuda_mtp_check.zig @@ -47,7 +47,7 @@ pub fn decode(comptime Tri: type, o: Ops(Tri), seq: *prompt.Seq, d: *Draft, out: @memcpy(ids[1..][0..ch.n], ch.slice()); const rows = 1 + ch.n; const tv = std.Io.Clock.awake.now(o.e.io).toNanoseconds(); - try o.verify(seq, d, ids[0..rows]); + try o.verifySampled(seq, d, ids[0..rows], sm); try o.draws(d, sm, &.{}, sampled[0..rows]); const td = std.Io.Clock.awake.now(o.e.io).toNanoseconds(); st.verify_ms += @as(f64, @floatFromInt(td - tv)) / 1e6; diff --git a/zig/src/families/flashnext/cuda_mtp_draft.zig b/zig/src/families/flashnext/cuda_mtp_draft.zig index 770502fe64..927a682577 100644 --- a/zig/src/families/flashnext/cuda_mtp_draft.zig +++ b/zig/src/families/flashnext/cuda_mtp_draft.zig @@ -43,6 +43,18 @@ pub const Chain = struct { } }; +test "exl3: reset and settlement discard completed verification records" { + var window: Draft = undefined; + window.verify_cache = .{}; + window.verify_cache.publish(window.verify_cache.generation, 4, 10); + window.reset(); + try std.testing.expect(!window.verify_cache.take(4, 10, true)); + window.verify_cache.publish(window.verify_cache.generation, 4, 10); + window.open = false; + window.settle(undefined); + try std.testing.expect(!window.verify_cache.take(4, 10, true)); +} + /// A sequence's head cache and the buffers a window needs: its rows' scores and streams, and what it rolls back to. pub const Draft = struct { sel: ?*@import("exl3/kv/cuda_exl3_kvtier.zig").LaneSel = null, // a host tier's selection of the stream's blocks @@ -86,6 +98,9 @@ pub const Draft = struct { /// The scores a sampled draw reads on the host (a window's rows, or the draft head's row). host: []u16, best_host: [max_rows * 4]f32 = undefined, + verify_host: ?cuda.HostBuffer = null, + verify_pending: ?cuda.Stream = null, + verify_cache: @import("cuda_mtp_verify_cache.zig").Cache = .{}, const hist_at: usize = 0; const proj_at: usize = hist_at + linear_n * conv_bytes; @@ -101,6 +116,7 @@ pub const Draft = struct { /// A head cache of `cap` entries; `rows` the widest prompt chunk whose streams it keeps. pub fn init(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: cuda.Stream, cap: usize, draft_n: usize, rows: usize) !Draft { + if (try @import("exl3/offload/cuda_exl3_recycle.zig").takeDraft(stream, cap, draft_n, rows)) |kept| return kept; // graphs: stable addresses var d: Draft = .{ .stream = stream, .gpa = gpa, .cap = cap, .rows_cap = @max(rows, max_rows), .kc = undefined, .vc = undefined, .ikc = undefined, .pooled = undefined, .streams = undefined, .out = undefined, .logits = undefined, .dlogits = undefined, .best = undefined, .saved = undefined, .host = undefined }; d.kc = try cuda.DeviceBuffer.alloc(driver, cap * kv_heads * head_dim * 2); errdefer d.kc.free(); @@ -197,6 +213,7 @@ pub const Draft = struct { /// An open window kept whole: its end states become the sequence's (buffers swapped, nothing copied). pub fn settle(d: *Draft, seq: *prompt.Seq) void { + d.verify_cache.invalidate(); if (!d.open) return; for (seq.state, &d.spare) |*a, *b| std.mem.swap(cuda.DeviceBuffer, a, b); d.aim(); @@ -204,7 +221,14 @@ pub const Draft = struct { } pub fn deinit(d: *Draft) void { + d.verify_cache.invalidate(); + if (d.verify_pending) |stream| { + stream.synchronize() catch {}; + d.verify_pending = null; + } if (d.stream) |stream| cuda.completion.drain(stream.d, stream, "Flash Next draft work did not drain before release"); + if (@import("exl3/offload/cuda_exl3_recycle.zig").giveDraft(d)) return; // graphs: kept for the next stream + if (d.verify_host) |*host| host.free(); d.gpa.free(d.host); for (&d.spare) |*b| b.free(); d.saved.free(); @@ -221,6 +245,7 @@ pub const Draft = struct { /// Drop the head's cache (a new prompt). pub fn reset(d: *Draft) void { + d.verify_cache.invalidate(); d.len = 0; d.drafted = 0; d.rows = 0; diff --git a/zig/src/families/flashnext/cuda_mtp_ops.zig b/zig/src/families/flashnext/cuda_mtp_ops.zig index 0809faa1d1..edebbbe512 100644 --- a/zig/src/families/flashnext/cuda_mtp_ops.zig +++ b/zig/src/families/flashnext/cuda_mtp_ops.zig @@ -69,8 +69,9 @@ fn download(stream: cuda.Stream, dst: []u8, src: u64) !void { /// An owner can mirror main-model mutations to a peer without exposing device pointers on the control wire. pub const Main = struct { ptr: *anyopaque, + rank: usize, forward: *const fn (*anyopaque, *prompt.Seq, []const u32, batch.Hooks) anyerror!void, - verify: *const fn (*anyopaque, *prompt.Seq, *Draft, []const u32) anyerror!void, + verify: *const fn (*anyopaque, *prompt.Seq, *Draft, []const u32, ?batch.Hooks.OutputTail) anyerror!void, rollback: *const fn (*anyopaque, *prompt.Seq, *Draft, usize) anyerror!void, }; @@ -95,6 +96,7 @@ pub fn Ops(comptime Tri: type) type { /// Prompt rows [at, stop): the pass with its streams kept, then the head over each row with a known next token. pub fn prefillChunk(o: Self, seq: *prompt.Seq, d: *Draft, ids: []const u32, at: usize, stop: usize) !void { + d.verify_cache.invalidate(); if (stop > d.cap or stop <= at or stop - at > d.rows_cap) return error.PromptTooLong; if (seq.fed != at or d.len != at or d.drafted != 0) return error.PositionMismatch; d.media = if (seq.media) |*m| m.view() else null; @@ -118,12 +120,15 @@ pub fn Ops(comptime Tri: type) type { /// The head over `next.len` rows of streams `rows`: cache at d.len, last row in d.out, scores in d.dlogits. pub fn forwardHead(o: Self, d: *Draft, next: []const u32, rows: u64, scores: bool) !void { + d.verify_cache.invalidate(); const e = o.e; const tri = o.tri; 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 planned = try @import("exl3/model/cuda_exl3_chain.zig").enter(o.e, @min(next.len, piece_rows)); // graphs: the head's experts on the device planner + defer @import("exl3/model/cuda_exl3_chain.zig").leave(o.e, planned); 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)); @@ -136,7 +141,7 @@ pub fn Ops(comptime Tri: type) type { .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 }, + .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, .dplan = &e.exl3_dplan, .prefill = &e.exl3_prefill }, .seq = d.headSeq(d.media, score_rows), .hooks = .{ .logits = .none }, }; @@ -188,6 +193,7 @@ pub fn Ops(comptime Tri: type) type { /// Python draft.absorb: the chain's entries trimmed, then the kept rows (main streams `rows`, next tokens). pub fn absorb(o: Self, d: *Draft, next: []const u32, rows: u64, scores: bool) !void { + d.verify_cache.invalidate(); if (d.drafted > 0) { d.len -= d.drafted; d.drafted = 0; @@ -197,6 +203,7 @@ pub fn Ops(comptime Tri: type) type { /// The draft at `position` and its temperature-1 probability (Python Engine.sample_draft). pub fn pick(o: Self, d: *Draft, position: u64, sm: ?Sampling) !Pick { + d.verify_cache.invalidate(); const h = o.h; const view_ = h.view(); const n = view_.logical; @@ -216,6 +223,8 @@ pub fn Ops(comptime Tri: type) type { /// Python decode.draft: absorb the kept rows, then chain up to `count` drafts, the first always kept. pub fn draft(o: Self, d: *Draft, next: []const u32, rows: u64, position: u64, count: usize, sm: ?Sampling, confidence: f64, out: *Chain) !void { + d.verify_cache.invalidate(); + if (try @import("exl3/model/cuda_exl3_chain.zig").draft(Tri, o, d, next, rows, position, count, sm, confidence, out)) return; // graphs: the greedy chain as one graph out.n = 0; try o.absorb(d, next, rows, count > 0); var j: usize = 0; @@ -291,7 +300,7 @@ pub fn Ops(comptime Tri: type) type { .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 }, + .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, .dplan = &e.exl3_dplan, .prefill = &e.exl3_prefill }, .seq = lane_buf[0].seq, .hooks = .{ .logits = .none }, .lanes = lane_buf[0..n_l], @@ -413,9 +422,31 @@ 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 { + return o.verifyWithTail(seq, d, ids, null); + } + + pub fn verifySampled(ops: Self, seq: *prompt.Seq, window: *Draft, ids: []const u32, sm: ?Sampling) !void { + window.verify_cache.invalidate(); + const greedy = if (sm) |sampling| sampling.temperature <= 0 else true; + const leader = if (ops.main) |main| main.rank == 0 else false; + if (!@import("cuda_mtp_verify_tail.zig").enabled or !greedy or !leader or ops.e.collective == null or !@import("exl3/graph/cuda_exl3_graph.zig").enabled()) return ops.verify(seq, window, ids); + if (window.verify_host == null) window.verify_host = try cuda.HostBuffer.alloc(ops.e.driver, max_rows * 16); + var tail: VerifyTail = .{ .head = ops.h, .draft = window }; + errdefer window.verify_cache.invalidate(); + try ops.verifyWithTail(seq, window, ids, .{ .ptr = &tail, .enqueue = VerifyTail.enqueue }); + if (tail.enqueued) { + window.verify_pending = null; + window.verify_cache.publish(tail.generation, window.rows, window.fed0); + } + } + + pub const VerifyTail = @import("cuda_mtp_verify_tail.zig").VerifyTail; + + pub fn verifyWithTail(o: Self, seq: *prompt.Seq, d: *Draft, ids: []const u32, tail: ?batch.Hooks.OutputTail) !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); + d.verify_cache.invalidate(); + if (o.main) |m| return m.verify(m.ptr, seq, d, ids, tail); o.e.round_graph.clear(); const n = ids.len; if (n == 0 or n > max_rows) return error.WindowTooWide; @@ -427,7 +458,7 @@ pub fn Ops(comptime Tri: type) type { @memcpy(d.ids[0..n], ids); const keeps = n > 1; if (keeps) try copyAsync(o.stream(), d.saved.ptr + Draft.tail_at, seq.tail.ptr, ple_tail_n * wide * 2); - try o.e.forward(Tri, o.tri, seq, ids, .{ .logits = .all, .logits_out = d.logits.ptr, .keep = if (keeps) &d.keep else null, .streams = d.streams.ptr }); + try o.e.forward(Tri, o.tri, seq, ids, .{ .logits = .all, .logits_out = d.logits.ptr, .keep = if (keeps) &d.keep else null, .streams = d.streams.ptr, .output_tail = tail }); if (keeps) try copyAsync(o.stream(), d.saved.ptr + Draft.nrow_at, o.e.batch_s.?.ple_nrow, n * wide * 2); d.open = keeps; } @@ -468,9 +499,14 @@ pub fn Ops(comptime Tri: type) type { pub fn draws(o: Self, d: *Draft, sm: ?Sampling, masks: []const u32, out: []u32) !void { const n = d.rows; const greedy = if (sm) |s| s.temperature <= 0 else true; + const cached = d.verify_cache.take(n, d.fed0, greedy); if (greedy and masks.len == 0) { - try o.h.best(o.stream(), d.logits.ptr, vocab, vocab, n, d.best.ptr); - try download(o.stream(), std.mem.sliceAsBytes(d.best_host[0 .. 4 * n]), d.best.ptr); + if (cached) { + @memcpy(d.best_host[0 .. 4 * n], d.verify_host.?.slice(f32)[0 .. 4 * n]); + } else { + try o.h.best(o.stream(), d.logits.ptr, vocab, vocab, n, d.best.ptr); + try download(o.stream(), std.mem.sliceAsBytes(d.best_host[0 .. 4 * n]), d.best.ptr); + } for (out[0..n], 0..) |*t, r| t.* = @intCast(@as(i32, @bitCast(d.best_host[4 * r + 2]))); return; } @@ -486,6 +522,7 @@ pub fn Ops(comptime Tri: type) type { /// Keep the window's first `keep` rows (Python forward.commit) from its saved rows; a repeat is a no-op. pub fn rollback(o: Self, seq: *prompt.Seq, d: *Draft, keep: usize) !void { + d.verify_cache.invalidate(); if (o.main) |m| return m.rollback(m.ptr, seq, d, keep); o.e.round_graph.clear(); if (keep == 0 or keep > d.rows) return error.CannotRewind; diff --git a/zig/src/families/flashnext/cuda_mtp_ops_test.zig b/zig/src/families/flashnext/cuda_mtp_ops_test.zig index bf7fa1578c..e11e49c598 100644 --- a/zig/src/families/flashnext/cuda_mtp_ops_test.zig +++ b/zig/src/families/flashnext/cuda_mtp_ops_test.zig @@ -8,6 +8,92 @@ const ops_ = @import("cuda_mtp_ops.zig"); const Ops = ops_.Ops; const Draft = mtp.Draft; const max_rows = mtp.max_rows; +const Head = mtp.Head; +const vocab = @import("cuda_prompt_wt.zig").vocab; + +test "exl3: greedy verification consumes completed records without touching the head or CUDA" { + var records: [max_rows * 16]u8 align(16) = undefined; + const values = std.mem.bytesAsSlice(f32, &records); + const tokens = [_]u32{ 0, 6124, 9564, @intCast(vocab - 1) }; + for (tokens, 0..) |token, row| values[4 * row + 2] = @bitCast(@as(i32, @intCast(token))); + var window: Draft = undefined; + window.rows = tokens.len; + window.fed0 = 37; + window.verify_cache = .{}; + window.verify_host = .{ .d = undefined, .bytes = &records }; + const ops: Ops(struct {}) = .{ .e = undefined, .tri = .{}, .h = undefined }; + var drawn: [tokens.len]u32 = undefined; + window.verify_cache.publish(window.verify_cache.generation, window.rows, window.fed0); + try ops.draws(&window, null, &.{}, &drawn); + try std.testing.expectEqualSlices(u32, &tokens, &drawn); + try std.testing.expect(window.verify_cache.ready == null); + window.verify_cache.publish(window.verify_cache.generation, window.rows, window.fed0); + try ops.draws(&window, .{ .seed = 0, .temperature = 0 }, &.{}, &drawn); + try std.testing.expectEqualSlices(u32, &tokens, &drawn); +} + +test "exl3: verification tail enqueues the existing top-1 kernel then only its records copy" { + const Fake = struct { + var launches: usize = 0; + var copies: usize = 0; + var rows: usize = 0; + var copied: usize = 0; + var fail_copy: bool = false; + + fn launch(_: cuda.abi.Function, grid: c_uint, _: c_uint, _: c_uint, block: c_uint, _: c_uint, _: c_uint, _: c_uint, _: cuda.abi.Stream, params: ?[*]?*anyopaque, _: ?[*]?*anyopaque) callconv(.c) cuda.abi.Result { + std.debug.assert(block == 1024); + std.debug.assert(launches == copies); + const scores: *const u64 = @ptrCast(@alignCast(params.?[0].?)); + const width: *const c_int = @ptrCast(@alignCast(params.?[1].?)); + const stride: *const c_longlong = @ptrCast(@alignCast(params.?[2].?)); + const output: *const u64 = @ptrCast(@alignCast(params.?[3].?)); + std.debug.assert(scores.* == 0x2000 and width.* == vocab and stride.* == vocab and output.* == 0x4000); + rows = grid; + launches += 1; + return 0; + } + + fn copy(_: ?*anyopaque, src: u64, bytes: usize, _: cuda.abi.Stream) callconv(.c) cuda.abi.Result { + std.debug.assert(launches == copies + 1 and src == 0x4000); + copies += 1; + copied = bytes; + if (fail_copy) return cuda.abi.error_not_ready; + return 0; + } + }; + Fake.launches = 0; + Fake.copies = 0; + Fake.fail_copy = false; + var driver: cuda.Driver = undefined; + driver.api.cuLaunchKernel = Fake.launch; + driver.api.cuMemcpyDtoHAsync_v2 = Fake.copy; + const stream: cuda.Stream = .{ .d = &driver, .handle = null }; + var head: Head = undefined; + head.best_fn = .{ .d = &driver, .handle = null }; + var records: [max_rows * 16]u8 align(16) = undefined; + var window: Draft = undefined; + window.logits = .{ .d = &driver, .ptr = 0x2000, .len = max_rows * vocab * 2 }; + window.best = .{ .d = &driver, .ptr = 0x4000, .len = max_rows * 16 }; + window.verify_host = .{ .d = &driver, .bytes = &records }; + window.verify_cache = .{}; + for ([_]usize{ 1, 2, 3, 4, 8, 16 }) |rows| { + window.rows = rows; + window.verify_cache.invalidate(); + var tail: Ops(struct {}).VerifyTail = .{ .head = &head, .draft = &window }; + try Ops(struct {}).VerifyTail.enqueue(&tail, stream); + try std.testing.expectEqual(rows, Fake.rows); + try std.testing.expectEqual(rows * 16, Fake.copied); + try std.testing.expect(tail.enqueued); + try std.testing.expectEqual(window.verify_cache.generation, tail.generation); + try std.testing.expect(window.verify_cache.ready == null); + } + Fake.fail_copy = true; + var failed: Ops(struct {}).VerifyTail = .{ .head = &head, .draft = &window }; + try std.testing.expectError(error.NotReady, Ops(struct {}).VerifyTail.enqueue(&failed, stream)); + try std.testing.expect(failed.enqueued); + try std.testing.expect(window.verify_pending != null); + try std.testing.expect(window.verify_cache.ready == null); +} test "exl3: common verification scopes TP1 and TP2 validation and delegation without joins" { const Tri = @import("cuda_launch.zig").Tri; @@ -18,7 +104,7 @@ test "exl3: common verification scopes TP1 and TP2 validation and delegation wit calls: usize = 0, fail: bool = false, - fn verify(ptr: *anyopaque, _: *prompt.Seq, _: *Draft, _: []const u32) !void { + fn verify(ptr: *anyopaque, _: *prompt.Seq, _: *Draft, _: []const u32, _: ?batch.Hooks.OutputTail) !void { const delegate: *@This() = @ptrCast(@alignCast(ptr)); delegate.calls += 1; try std.testing.expectEqual(Phase.verify, delegate.engine.telemetry.forward_phase); @@ -36,11 +122,14 @@ test "exl3: common verification scopes TP1 and TP2 validation and delegation wit var driver = Fake.driver(); var stream: cuda.Stream = .{ .d = &driver, .handle = null }; var engine: prompt.Engine = undefined; + engine.exl3_prefill = .{}; + engine.exl3_chain = .{}; engine.stream = &stream; engine.round_graph = .{}; var sequence: prompt.Seq = undefined; sequence.cap = max_rows; var window: Draft = undefined; + window.verify_cache = .{}; 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); @@ -51,20 +140,20 @@ test "exl3: common verification scopes TP1 and TP2 validation and delegation wit 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.expectError(error.WindowTooWide, ops.verifyWithTail(&sequence, &window, &.{}, null)); try std.testing.expectEqual(previous, engine.telemetry.forward_phase); - try std.testing.expectError(error.WindowTooWide, ops.verify(&sequence, &window, &too_wide)); + try std.testing.expectError(error.WindowTooWide, ops.verifyWithTail(&sequence, &window, &too_wide, null)); 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.expectError(error.PromptTooLong, ops.verifyWithTail(&sequence, &window, &.{1}, null)); try std.testing.expectEqual(previous, engine.telemetry.forward_phase); - ops.main = .{ .ptr = &delegate, .forward = Delegate.forward, .verify = Delegate.verify, .rollback = Delegate.rollback }; + ops.main = .{ .ptr = &delegate, .rank = 0, .forward = Delegate.forward, .verify = Delegate.verify, .rollback = Delegate.rollback }; const calls = delegate.calls; delegate.fail = false; - try ops.verify(&sequence, &window, &.{1}); + try ops.verifyWithTail(&sequence, &window, &.{1}, null); 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.expectError(error.VerificationFailed, ops.verifyWithTail(&sequence, &window, &.{1}, null)); 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); @@ -73,3 +162,44 @@ test "exl3: common verification scopes TP1 and TP2 validation and delegation wit } } } + +test "exl3: paged draft heads enter conditional graph dispatch" { + const graph = @import("exl3/graph/cuda_exl3_graph.zig"); + const chain = @import("exl3/model/cuda_exl3_chain.zig"); + const Tri = @import("cuda_launch.zig").Tri; + const Mock = struct { + e: *prompt.Engine, + h: *Head = undefined, + tri: Tri = undefined, + + pub fn absorb(_: @This(), _: *Draft, _: []const u32, _: u64, _: bool) !void { + return error.UnexpectedGraphAbsorb; + } + pub fn pick(_: @This(), _: *Draft, _: u64, _: ?mtp.Sampling) !mtp.Pick { + return error.UnexpectedGraphPick; + } + pub fn forwardHead(_: @This(), _: *Draft, _: []const u32, _: u64, _: bool) !void { + return error.UnexpectedGraphForward; + } + }; + const Pool = @import("exl3/offload/cuda_exl3_pool.zig").Pool; + var pool: Pool = undefined; + pool.opt.layer = 0; + var pools = [_]*Pool{&pool}; + var plan: @import("exl3/dplan/cuda_exl3_dplan.zig").DevPlan = undefined; + plan.pools = &pools; + var engine: prompt.Engine = undefined; + engine.n_layers = 0; + engine.gpa = std.testing.allocator; + engine.driver = undefined; + engine.exl3_dplan = .{ .plan = &plan, .mtp_pool = &pool }; + var d: Draft = undefined; + d.media = null; + d.sel = null; + d.pg = .{ .pt = 0x1000, .ps = 4096, .ptok = 256 }; + var out: mtp.Chain = .{}; + const previous = graph.force; + graph.force = .graph; + defer graph.force = previous; + try std.testing.expectError(error.UnexpectedGraphAbsorb, chain.draft(Tri, Mock{ .e = &engine }, &d, &.{1}, 0, 0, 2, null, 0, &out)); +} diff --git a/zig/src/families/flashnext/cuda_mtp_verify_cache.zig b/zig/src/families/flashnext/cuda_mtp_verify_cache.zig new file mode 100644 index 0000000000..c612943b7c --- /dev/null +++ b/zig/src/families/flashnext/cuda_mtp_verify_cache.zig @@ -0,0 +1,65 @@ +const std = @import("std"); + +pub const Cache = struct { + generation: u64 = 0, + ready: ?u64 = null, + rows: usize = 0, + position: usize = 0, + + pub fn invalidate(cache: *Cache) void { + cache.generation +%= 1; + cache.ready = null; + } + + pub fn publish(cache: *Cache, generation: u64, rows: usize, position: usize) void { + if (generation != cache.generation) return; + cache.rows = rows; + cache.position = position; + cache.ready = generation; + } + + pub fn take(cache: *Cache, rows: usize, position: usize, greedy: bool) bool { + defer cache.invalidate(); + return greedy and cache.ready == cache.generation and cache.rows == rows and cache.position == position; + } +}; + +test "exl3: verification records require completion in the current generation and are consumed once" { + var cache: Cache = .{}; + try std.testing.expect(!cache.take(4, 10, true)); + const generation = cache.generation; + cache.publish(generation, 4, 10); + try std.testing.expect(cache.take(4, 10, true)); + try std.testing.expect(!cache.take(4, 10, true)); + cache.publish(generation, 4, 10); + try std.testing.expect(!cache.take(4, 10, true)); +} + +test "exl3: sampled draws, changed windows, rollback and reuse invalidate verification records" { + var cache: Cache = .{}; + cache.publish(cache.generation, 4, 10); + try std.testing.expect(!cache.take(4, 10, false)); + try std.testing.expect(!cache.take(4, 10, true)); + cache.publish(cache.generation, 4, 10); + try std.testing.expect(!cache.take(3, 10, true)); + cache.publish(cache.generation, 4, 10); + try std.testing.expect(!cache.take(4, 11, true)); + const prior = cache.generation; + cache.invalidate(); + cache.publish(prior, 4, 10); + try std.testing.expect(!cache.take(4, 10, true)); + cache.generation = std.math.maxInt(u64); + cache.publish(cache.generation, 4, 10); + cache.invalidate(); + try std.testing.expect(!cache.take(4, 10, true)); +} + +test "exl3: verification generations belong to each draft independently" { + var first: Cache = .{}; + var second: Cache = .{}; + first.publish(first.generation, 4, 10); + second.publish(second.generation, 8, 20); + first.invalidate(); + try std.testing.expect(!first.take(4, 10, true)); + try std.testing.expect(second.take(8, 20, true)); +} diff --git a/zig/src/families/flashnext/cuda_mtp_verify_tail.zig b/zig/src/families/flashnext/cuda_mtp_verify_tail.zig new file mode 100644 index 0000000000..719c41fda5 --- /dev/null +++ b/zig/src/families/flashnext/cuda_mtp_verify_tail.zig @@ -0,0 +1,28 @@ +//! The verification window's output tail: the top-1 kernel and its records copy enqueued behind the main pass. +const cuda = @import("cuda"); +const mtp = @import("cuda_mtp.zig"); +const wt = @import("cuda_prompt_wt.zig"); + +const vocab = wt.vocab; +const Head = mtp.Head; +const Draft = mtp.Draft; + +/// Not a user switch: the tail is enqueued only when a test or tool turns this on (greedy drafts, graphs, two ranks). +pub var enabled: bool = false; + +pub const VerifyTail = struct { + head: *Head, + draft: *Draft, + generation: u64 = 0, + enqueued: bool = false, + + pub fn enqueue(ptr: *anyopaque, st: cuda.Stream) !void { + const tail: *VerifyTail = @ptrCast(@alignCast(ptr)); + const window = tail.draft; + tail.generation = window.verify_cache.generation; + try tail.head.best(st, window.logits.ptr, vocab, vocab, window.rows, window.best.ptr); + tail.enqueued = true; + window.verify_pending = st; + try window.best.downloadAsync(0, window.verify_host.?.bytes[0 .. window.rows * 16], st.handle); + } +}; diff --git a/zig/src/families/flashnext/cuda_prompt.zig b/zig/src/families/flashnext/cuda_prompt.zig index 8067cc667b..859f7c9755 100644 --- a/zig/src/families/flashnext/cuda_prompt.zig +++ b/zig/src/families/flashnext/cuda_prompt.zig @@ -149,6 +149,10 @@ pub const Engine = struct { round_graph: round_graph.Cache = .{}, image_factory: ?*@import("cuda_images.zig").Vision = null, collective: ?*@import("cuda_collective.zig").Collective = null, + /// This engine's EXL3 device planner (cuda_exl3_dplan.zig): nothing of it is shared with another engine. + exl3_chain: @import("exl3/model/cuda_exl3_chain.zig").Runtime = .{}, + exl3_prefill: @import("exl3/prefill/cuda_exl3_prefill_state.zig").Runtime = .{}, + exl3_dplan: @import("exl3/dplan/cuda_exl3_dplan.zig").State = .{}, 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 = .{}, @@ -176,6 +180,9 @@ pub const Engine = struct { e.scratch_format = .classic; e.head_local_rows = std.math.maxInt(usize); e.collective = null; + e.exl3_dplan = .{}; + e.exl3_chain = .{}; + e.exl3_prefill = .{ .options = @import("exl3/prefill/cuda_exl3_prefill_state.zig").optionsFromEnv() }; e.image_factory = null; e.telemetry = cache_telemetry.State.fromEnv(rank, world); e.round_graph = .{}; @@ -246,7 +253,11 @@ pub const Engine = struct { } pub fn deinit(e: *Engine) void { - @import("exl3/prefill/cuda_exl3_prefill.zig").releaseFor(e); + @import("exl3/model/cuda_exl3_chain.zig").releaseRuntime(&e.exl3_chain); + @import("exl3/prefill/cuda_exl3_prefill.zig").releaseFor(&e.exl3_prefill, e); + @import("exl3/graph/cuda_exl3_graph.zig").releaseFor(e); + @import("exl3/dplan/cuda_exl3_dplan.zig").release(&e.exl3_dplan); + @import("exl3/offload/cuda_exl3_recycle.zig").drain(e); e.round_graph.clear(); e.pockets.deinit(); if (e.batch_s) |*bs| bs.free(e.gpa); @@ -293,7 +304,7 @@ pub const Engine = struct { /// The resident weights a pass reads, with the taught layers' pockets. pub fn model(e: *Engine) batch.Model { const m = e.pockets.mask(); - return .{ .table_w = e.table.w.ptr, .table_s = e.table.s.ptr, .table_b = e.table.b.ptr, .table_dense = e.table.dense, .collective = e.collective, .layers = e.layers[0..e.n_layers], .mixer = e.mixer, .head = e.head_q, .ngram = &e.ngram, .pocket = if (m != 0) .{ .ctx = e, .add = servePocket, .layers = m } else null }; + return .{ .table_w = e.table.w.ptr, .table_s = e.table.s.ptr, .table_b = e.table.b.ptr, .table_dense = e.table.dense, .collective = e.collective, .dplan = &e.exl3_dplan, .prefill = &e.exl3_prefill, .layers = e.layers[0..e.n_layers], .mixer = e.mixer, .head = e.head_q, .ngram = &e.ngram, .pocket = if (m != 0) .{ .ctx = e, .add = servePocket, .layers = m } else null }; } /// The weights a pass of streams that read the model as shipped reads (untaught): the pockets left out. @@ -315,6 +326,9 @@ pub const Engine = struct { var paged = given; errdefer if (paged) |*p| p.deinit(e.gpa); e.round_graph.clear(); + if (paged == null and cap != 0) { + if (try @import("exl3/offload/cuda_exl3_recycle.zig").takeSeq(e, cap)) |kept| return kept; + } if (cap == 0 or cap > max_tokens) return error.UnexpectedTensor; _ = try seqBytesPaged(cap, if (paged) |p| p.tokens else 0); var seq: Seq = .{ @@ -331,6 +345,7 @@ pub const Engine = struct { .fed = 0, .logits = &.{}, .cap = cap, + .recycle_owner = if (paged == null) e else null, }; paged = null; // seq owns the pages from here // One rollback owns every published array, and drains before releasing any destination. @@ -419,6 +434,7 @@ pub const Engine = struct { .seq = seq.view(scores), .hooks = hooks, }; + if (try @import("exl3/graph/cuda_exl3_graph.zig").forward(e, &p, seq, ids)) return; var at: usize = 0; while (at < ids.len) { const n = @min(s.cap, ids.len - at); @@ -480,7 +496,7 @@ pub const Engine = struct { if (q.paged != null) break true; } else false; // a captured round fixes the lanes' tables: paged streams run the plain pass - if (e.collective != null and round_graph.enabled() and !has_media and !has_paged) { + if (try @import("exl3/graph/cuda_exl3_graph.zig").round(e, &p, seqs, ids, scores, out)) {} else if (e.collective != null and round_graph.enabled() and !has_media and !has_paged) { const key = try round_graph.Key.init(seqs, s.ids, scores, out); try e.round_graph.run(key, &p, ids); } else try p.chunk(ids, 0, null); @@ -495,6 +511,7 @@ pub const Engine = struct { /// One stream's caches. `fed` tokens have already gone through the model. pub const Seq = struct { + recycle_owner: ?*Engine = null, /// The owning compute stream; synthetic non-device fixtures may leave it absent. stream: ?cuda.Stream = null, media: ?@import("cuda_media.zig").State = null, diff --git a/zig/src/families/flashnext/cuda_prompt_lifecycle.zig b/zig/src/families/flashnext/cuda_prompt_lifecycle.zig index 1aa3e5e31e..6268306d4b 100644 --- a/zig/src/families/flashnext/cuda_prompt_lifecycle.zig +++ b/zig/src/families/flashnext/cuda_prompt_lifecycle.zig @@ -2,7 +2,9 @@ const std = @import("std"); const prompt = @import("cuda_prompt.zig"); pub fn deinit(sequence: *prompt.Seq, allocator: std.mem.Allocator) void { + if (@import("exl3/offload/cuda_exl3_recycle.zig").giveSeq(sequence)) return; if (sequence.stream) |stream| @import("cuda").completion.drain(stream.d, stream, "Flash Next sequence work did not drain before release"); + @import("exl3/graph/cuda_exl3_graph.zig").releaseSeq(sequence); 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(); diff --git a/zig/src/families/flashnext/cuda_rank_engine.zig b/zig/src/families/flashnext/cuda_rank_engine.zig index d1f8565cbd..7f4ca2fb0c 100644 --- a/zig/src/families/flashnext/cuda_rank_engine.zig +++ b/zig/src/families/flashnext/cuda_rank_engine.zig @@ -238,13 +238,18 @@ pub const Session = struct { /// Both ranks keep their own recurrent window; only the leader consumes draft scores and hidden rows. pub fn verify(s: *Session, comptime Tri: type, tri: Tri, q: *prompt.Seq, d: *mtp.Draft, ids: []const u32) !void { - if (!s.drafting or ids.len == 0 or ids.len > mtp.max_rows or q.fed + ids.len > q.cap) return error.UnsupportedRankOperation; + return s.verifyWithTail(Tri, tri, q, d, ids, null); + } + + pub fn verifyWithTail(session: *Session, comptime Tri: type, tri: Tri, seq: *prompt.Seq, draft_state: *mtp.Draft, ids: []const u32, tail: ?prompt.batch.Hooks.OutputTail) !void { + if (tail != null and session.rank != 0) return error.UnsupportedRankOperation; + if (!session.drafting or ids.len == 0 or ids.len > mtp.max_rows or seq.fed + ids.len > seq.cap) return error.UnsupportedRankOperation; try tokenIds(ids); - const h = try s.command(.{ .kind = .verify, .object = q.rank_id, .bytes = @intCast(ids.len * 4) }, std.mem.sliceAsBytes(ids)); - errdefer s.poison(); - const ops: mtp.Ops(Tri) = .{ .e = s.engine, .tri = tri, .h = undefined }; - try ops.verify(q, d, ids); - try s.ack(h); + const header = try session.command(.{ .kind = .verify, .object = seq.rank_id, .bytes = @intCast(ids.len * 4) }, std.mem.sliceAsBytes(ids)); + errdefer session.poison(); + const ops: mtp.Ops(Tri) = .{ .e = session.engine, .tri = tri, .h = undefined }; + try ops.verifyWithTail(seq, draft_state, ids, tail); + try session.ack(header); } pub fn keep(s: *Session, comptime Tri: type, tri: Tri, q: *prompt.Seq, d: *mtp.Draft, count: usize) !void { @@ -383,6 +388,8 @@ pub const Session = struct { }, .drop_seq => { s.engine.round_graph.clear(); + const sequence = s.seqs.get(h.object) orelse return error.UnknownRankSequence; + if (sequence.paged != null) @import("exl3/graph/cuda_exl3_graph.zig").releaseDecodeSeq(sequence); if (s.windows.fetchRemove(h.object)) |entry| { var d = entry.value; d.deinit(); diff --git a/zig/src/families/flashnext/cuda_rank_media_test.zig b/zig/src/families/flashnext/cuda_rank_media_test.zig index 77c2d54cff..32fca0e7ad 100644 --- a/zig/src/families/flashnext/cuda_rank_media_test.zig +++ b/zig/src/families/flashnext/cuda_rank_media_test.zig @@ -111,6 +111,8 @@ test "a leader whose peer leaves after the image broadcast is enqueued drains be var d = Fake.driver(); var stream: cuda.Stream = .{ .d = &d, .handle = null }; var engine: prompt.Engine = undefined; + engine.exl3_prefill = .{}; + engine.exl3_chain = .{}; engine.driver = &d; engine.stream = &stream; var pair: [2]c_int = undefined; @@ -132,6 +134,8 @@ test "a follower whose image broadcast fails after it is enqueued keeps the feat var d = Fake.driver(); var stream: cuda.Stream = .{ .d = &d, .handle = null }; var engine: prompt.Engine = undefined; + engine.exl3_prefill = .{}; + engine.exl3_chain = .{}; engine.driver = &d; engine.stream = &stream; var v: images.Vision = undefined; diff --git a/zig/src/families/flashnext/cuda_sequence_completion_test.zig b/zig/src/families/flashnext/cuda_sequence_completion_test.zig index ddf4e69bfa..360264522d 100644 --- a/zig/src/families/flashnext/cuda_sequence_completion_test.zig +++ b/zig/src/families/flashnext/cuda_sequence_completion_test.zig @@ -93,6 +93,8 @@ test "Flash Next completion: allocation refusals drain fills before every constr var d = Fake.init(deny, 0); var stream: cuda.Stream = .{ .d = &d, .handle = @ptrFromInt(325) }; var e: prompt.Engine = undefined; + e.exl3_prefill = .{}; + e.exl3_chain = .{}; e.gpa = t.allocator; e.driver = &d; e.stream = &stream; @@ -114,6 +116,8 @@ test "Flash Next completion: failed asynchronous fill and an interrupted prefill var d = Fake.init(0, fail_fill); var stream: cuda.Stream = .{ .d = &d, .handle = @ptrFromInt(326) }; var e: prompt.Engine = undefined; + e.exl3_prefill = .{}; + e.exl3_chain = .{}; e.gpa = t.allocator; e.driver = &d; e.stream = &stream; @@ -127,6 +131,8 @@ test "Flash Next completion: failed asynchronous fill and an interrupted prefill var d = Fake.init(0, 0); var stream: cuda.Stream = .{ .d = &d, .handle = @ptrFromInt(327) }; var e: prompt.Engine = undefined; + e.exl3_prefill = .{}; + e.exl3_chain = .{}; e.gpa = t.allocator; e.driver = &d; e.stream = &stream; diff --git a/zig/src/families/flashnext/exl3/dplan/cuda_exl3_dplan.zig b/zig/src/families/flashnext/exl3/dplan/cuda_exl3_dplan.zig new file mode 100644 index 0000000000..044b5b4f90 --- /dev/null +++ b/zig/src/families/flashnext/exl3/dplan/cuda_exl3_dplan.zig @@ -0,0 +1,516 @@ +//! The offload pool planned on the device (flashnext_exl3_dplan.cu): a decode round's MoE needs no host wait between + +const std = @import("std"); +const cuda = @import("cuda"); +const pool_ = @import("../offload/cuda_exl3_pool.zig"); +const offload = pool_.offload; +const telemetry = @import("../model/exl3_telemetry.zig"); +const recent_demand = @import("../offload/exl3_recent_demand.zig"); +const layout = @import("exl3_dplan_layout.zig"); +pub const demand_trace = @import("exl3_dplan_trace.zig"); +pub const replay = @import("exl3_dplan_replay.zig"); +pub const policy_eval = @import("exl3_dplan_policy_eval.zig"); + +const experts_n = layout.experts_n; +const misc_words = layout.misc_words; +const recentCountsOffset = layout.recentCountsOffset; +pub const miscOffset = layout.miscOffset; +pub const PoolState = layout.PoolState; +const empty_id = offload.Cache.empty_id; + +comptime { + if (recent_demand.expert_count != experts_n) @compileError("recent expert history shape mismatch"); +} + +test { + _ = recent_demand; + _ = layout; + _ = demand_trace; + _ = replay; + _ = @import("exl3_dplan_trace_test.zig"); + _ = @import("exl3_dplan_trace_replay_test.zig"); + _ = @import("cuda_exl3_dplan_test.zig"); + _ = @import("exl3_dplan_policy_eval_test.zig"); + _ = policy_eval; +} + +/// u32 words of one layer's device state. +pub const stateWords = layout.stateWords; + +/// The planned copies a round can hold: a row's ten routed picks each. +fn workCap(max_rows: usize) usize { + return max_rows * 10; +} + +pub const Owner = enum { host, device }; + +/// Device statistics of a layer (the misc words after now). +pub const Misc = extern struct { now: u32, rounds: u32, hit_slots: u32, miss_slots: u32, copies: u32, evictions: u32, overflow: u32, inplace: u32, recent_cursor: u32, recent_filled: u32 }; + +pub const DevPlan = struct { + gpa: std.mem.Allocator, + d: *const cuda.Driver, + mod: cuda.Module, + plan_fn: cuda.Function, + copy_fn: cuda.Function, + restore_fn: cuda.Function, + pools: []*pool_.Pool, + pool_states: []PoolState, + /// Rows a planned round may hold: rows * 10 distinct experts must fit the lines. + max_rows: usize, + state: cuda.DeviceBuffer, + mirror: cuda.HostBuffer, + /// Per pool: src [512 * 3] int64, lens [512 * 3] int64, k2 [512 * 3] int32 (static). + statics: cuda.DeviceBuffer, + work: cuda.DeviceBuffer, + owner: Owner = .host, + /// The stream the last device round ran on (toHost waits for it). + stream: ?cuda.Stream = null, + copy_blocks: u32 = 144, + /// Overlap: the held experts compute while the misses' copies run on `side`, then the misses' experts. + overlap: bool = false, + recent_eviction: bool = false, + side: cuda.Stream, + ev_plan: cuda.Event, + ev_copy: cuda.Event, + /// The split's masked picks [max_rows * 11] each: held (and the shared expert), then the misses. + split: cuda.DeviceBuffer, + /// Device statistics already added to the pools' (a pull adds the difference). + seen: []Misc, + pushes: u64 = 0, + pulls: u64 = 0, + history_initialized: bool = false, + /// Experts read in place from the host bank (more misses than free lines in a round), summed at each pull. + inplace: u64 = 0, + trace_mem: ?cuda.DeviceBuffer = null, + trace_checkpoint_fn: ?cuda.Function = null, + + const static_bytes = layout.static_bytes; + + pub fn init(gpa: std.mem.Allocator, d: *const cuda.Driver, pools: []const *pool_.Pool) !*DevPlan { + if (!cuda.kernels.available) return error.BuiltWithoutKernels; + if (pools.len == 0) return error.Invalid; + const pool_states = try gpa.alloc(PoolState, pools.len); + errdefer gpa.free(pool_states); + var total_words: usize = 0; + for (pools, pool_states) |p, *entry| { + if (p.lines.spec.experts < experts_n) return error.Invalid; + entry.* = try PoolState.init(p.opt.lines, total_words); + total_words = entry.endWords(); + } + var mod = try cuda.Module.load(d, cuda.kernels.flashnext_exl3_dplan); + errdefer mod.unload(); + const dp = try gpa.create(DevPlan); + errdefer gpa.destroy(dp); + const own = try gpa.dupe(*pool_.Pool, pools); + errdefer gpa.free(own); + var state = try cuda.DeviceBuffer.alloc(d, total_words * @sizeOf(u32)); + errdefer state.free(); + var mirror = try cuda.HostBuffer.alloc(d, state.len); + errdefer mirror.free(); + var statics = try cuda.DeviceBuffer.alloc(d, pools.len * static_bytes); + errdefer statics.free(); + // rows a planned round may hold (MTP verify windows reach 16): more distinct misses than lines free of the + const max_rows = maxRows(); + var trace_mem: ?cuda.DeviceBuffer = null; + errdefer if (trace_mem) |*buffer| buffer.free(); + if (demand_trace.requested) { + for (pools, pool_states, 0..) |p, entry, index| _ = try demand_trace.header(index, p.opt.layer, entry.lines, max_rows); + trace_mem = try cuda.DeviceBuffer.alloc(d, try std.math.mul(usize, pools.len, demand_trace.pool_bytes)); + try trace_mem.?.fill32(0, null); + for (pools, pool_states, 0..) |p, entry, index| { + const header = try demand_trace.header(index, p.opt.layer, entry.lines, max_rows); + try trace_mem.?.upload(index * demand_trace.pool_bytes, std.mem.asBytes(&header)); + } + } + var work = try cuda.DeviceBuffer.alloc(d, workCap(max_rows) * 9 * 8 + 256 + (experts_n + 64) * 4); + errdefer work.free(); + const seen = try gpa.alloc(Misc, pools.len); + errdefer gpa.free(seen); + var side = try cuda.Stream.init(d, true); + errdefer side.deinit(); + var ev_plan = try cuda.Event.init(d, false); + errdefer ev_plan.deinit(); + var ev_copy = try cuda.Event.init(d, false); + errdefer ev_copy.deinit(); + var split = try cuda.DeviceBuffer.alloc(d, 2 * max_rows * 11 * 4); + errdefer split.free(); + // every expert's three host-bank addresses (mapped: unified addressing), byte counts and widths + const host = try gpa.alloc(u8, static_bytes); + defer gpa.free(host); + for (pools, 0..) |p, i| { + const src = std.mem.bytesAsSlice(u64, host[0 .. experts_n * 24]); + const lens = std.mem.bytesAsSlice(u64, host[experts_n * 24 .. experts_n * 48]); + const k2 = std.mem.bytesAsSlice(i32, host[experts_n * 48 ..]); + for (0..experts_n) |e| for (0..3) |w| { + const b = p.src.routed[e][w]; + src[e * 3 + w] = if (b.len == 0) 0 else @intFromPtr(b.ptr); + lens[e * 3 + w] = b.len; + k2[e * 3 + w] = p.k2[e][w]; + if (b.len != 0 and (b.len % 16 != 0 or @intFromPtr(b.ptr) % 16 != 0)) return error.Invalid; + }; + try statics.upload(i * static_bytes, host); + } + dp.* = .{ + .gpa = gpa, + .d = d, + .mod = mod, + .plan_fn = try mod.function("fn_exl3_dplan"), + .copy_fn = try mod.function("fn_exl3_dcopy"), + .restore_fn = try mod.function("fn_exl3_drestore"), + .pools = own, + .pool_states = pool_states, + .max_rows = max_rows, + .state = state, + .mirror = mirror, + .statics = statics, + .work = work, + .seen = seen, + .side = side, + .ev_plan = ev_plan, + .ev_copy = ev_copy, + .split = split, + .overlap = true, // the held experts compute beside the misses' copies: +3-4% decode on two RTX 5060 Ti (see the report) + .trace_mem = trace_mem, + .trace_checkpoint_fn = if (trace_mem != null) try mod.function("fn_exl3_dcheckpoint") else null, + }; + if (dp.overlap) dp.copy_blocks = overlapCopyBlocks(d) catch dp.copy_blocks; + return dp; + } + + pub fn deinit(dp: *DevPlan) void { + dp.side.synchronize() catch {}; + if (dp.trace_mem) |*buffer| buffer.free(); + dp.split.free(); + dp.ev_copy.deinit(); + dp.ev_plan.deinit(); + dp.side.deinit(); + dp.work.free(); + dp.statics.free(); + dp.mirror.free(); + dp.state.free(); + dp.mod.unload(); + dp.gpa.free(dp.seen); + dp.gpa.free(dp.pool_states); + dp.gpa.free(dp.pools); + dp.gpa.destroy(dp); + } + + fn mirrorOf(dp: *DevPlan, i: usize) []u32 { + const all: []u32 = @alignCast(std.mem.bytesAsSlice(u32, dp.mirror.bytes[0..dp.state.len])); + return dp.pool_states[i].slice(all); + } + + /// The host's residency up to the device, on `s` (a no-op if the device owns it). Waits for every pool's copies. + pub fn toDevice(dp: *DevPlan, s: cuda.Stream) !void { + dp.stream = s; + if (dp.owner == .device) return; + for (dp.pools, 0..) |p, i| { + const lines = dp.pool_states[i].lines; + // copies in flight land, prefetched lines count as resident (the device knows no in-flight state) + try p.lines.drain(); + for (p.line_tab, p.slots_tab) |*state, id| if (state.* == .inflight) { + offload.reapAt(&p.v, p.opt.layer, id, p.now) catch {}; + }; + const m = dp.mirrorOf(i); + const line_id: []i32 = @ptrCast(m[0..lines]); + const stamp = m[lines .. 2 * lines]; + const line_of: []i32 = @ptrCast(m[2 * lines ..][0..experts_n]); + const counts = m[2 * lines + experts_n ..][0..experts_n]; + if (!dp.history_initialized) @memset(m[recentCountsOffset(lines)..miscOffset(lines)], 0); + @memset(line_of, -1); + for (0..lines) |ln| { + const id = p.slots_tab[ln]; + const held = p.line_tab[ln] == .resident and id != empty_id and id < experts_n; + line_id[ln] = if (held) id else -1; + stamp[ln] = p.stamp_tab[ln]; + if (held) line_of[id] = @intCast(ln); + } + @memcpy(counts, &p.counts[p.opt.layer]); + const misc: *Misc = @ptrCast(m[miscOffset(lines)..][0..misc_words]); + const recent_cursor = if (dp.history_initialized) misc.recent_cursor else 0; + const recent_filled = if (dp.history_initialized) misc.recent_filled else 0; + misc.* = .{ .now = p.now, .rounds = 0, .hit_slots = 0, .miss_slots = 0, .copies = 0, .evictions = 0, .overflow = 0, .inplace = 0, .recent_cursor = recent_cursor, .recent_filled = recent_filled }; + dp.seen[i] = misc.*; + } + try s.d.check(s.d.api.cuMemcpyHtoDAsync_v2(dp.state.ptr, dp.mirror.bytes.ptr, dp.state.len, s.handle), "cuMemcpyHtoDAsync"); + try dp.traceOwnership(s, true); + dp.history_initialized = true; + dp.owner = .device; + dp.pushes += 1; + } + + /// The device's residency back to the host pools (a no-op if the host owns it): waits for the last round's stream. + pub fn toHost(dp: *DevPlan) !void { + if (dp.owner == .host) return; + if (dp.stream) |s| try dp.traceOwnership(s, false); + if (dp.stream) |s| try s.synchronize(); + try dp.side.synchronize(); + try dp.d.check(dp.d.api.cuMemcpyDtoH_v2(dp.mirror.bytes.ptr, dp.state.ptr, dp.state.len), "cuMemcpyDtoH"); + var overflow: u64 = 0; + for (dp.pools, 0..) |p, i| { + const lines = dp.pool_states[i].lines; + const m = dp.mirrorOf(i); + const line_id: []const i32 = @ptrCast(m[0..lines]); + const stamp = m[lines .. 2 * lines]; + const counts = m[2 * lines + experts_n ..][0..experts_n]; + const misc: *const Misc = @ptrCast(m[miscOffset(lines)..][0..misc_words]); + for (0..lines) |ln| { + const id = line_id[ln]; + p.slots_tab[ln] = if (id < 0) empty_id else @intCast(id); + p.line_tab[ln] = if (id < 0) .empty else .resident; + p.stamp_tab[ln] = stamp[ln]; + p.layer_tab[ln] = p.opt.layer; + } + @memcpy(&p.counts[p.opt.layer], counts); + const was = dp.seen[i]; + telemetry.Device.record(&p.telemetry_device, &p.telemetry_valid, misc.*, was, dp.max_rows * 10); + p.now = misc.now; + p.stats.rounds += misc.rounds -% was.rounds; + p.stats.passes += misc.rounds -% was.rounds; + p.stats.hit_slots += misc.hit_slots -% was.hit_slots; + p.stats.miss_slots += misc.miss_slots -% was.miss_slots; + p.stats.misses += misc.copies -% was.copies; + p.stats.evictions += misc.evictions -% was.evictions; + overflow += misc.overflow; + dp.inplace += misc.inplace -% was.inplace; + dp.seen[i] = misc.*; + // the device's rounds finished (synchronized above): no line waits for a reader, nothing is in flight + @memset(p.lines.line_epoch, 0); + @memset(p.lines.line_pending, false); + p.prev_n = 0; + p.picks_pending = false; + } + dp.owner = .host; + dp.pulls += 1; + if (overflow != 0) { + std.log.err("flashnext exl3 dplan: {d} rounds had more distinct experts than the planner holds (256): results of those rounds are wrong", .{overflow}); + return error.PoolPasses; + } + } + + /// Device statistics so far, summed over the layers (synchronizes; does not change the owner). + pub fn peek(dp: *DevPlan) !Misc { + var t: Misc = std.mem.zeroes(Misc); + if (dp.owner == .host) return t; + if (dp.stream) |s| try s.synchronize(); + try dp.d.check(dp.d.api.cuMemcpyDtoH_v2(dp.mirror.bytes.ptr, dp.state.ptr, dp.state.len), "cuMemcpyDtoH"); + for (dp.pools, 0..) |_, i| { + const m = dp.mirrorOf(i); + const misc: *const Misc = @ptrCast(m[miscOffset(dp.pool_states[i].lines)..][0..misc_words]); + inline for (.{ "rounds", "hit_slots", "miss_slots", "copies", "evictions", "overflow", "inplace" }) |f| @field(t, f) += @field(misc, f) - @field(dp.seen[i], f); + } + return t; + } + + pub fn samplePoolAfterConsumers(dp: *const DevPlan, pool: *const pool_.Pool) !telemetry.Device { + if (!pool.telemetry_valid) return error.TelemetryCounterRegression; + var total = pool.telemetry_device; + if (dp.owner != .device) return total; + for (dp.pools, 0..) |candidate, index| { + if (candidate != pool) continue; + var observed: Misc = undefined; + const entry = dp.pool_states[index]; + const offset = entry.byteOffset() + miscOffset(entry.lines) * @sizeOf(u32); + try dp.state.download(offset, std.mem.asBytes(&observed)); + try telemetry.add(telemetry.Device, &total, try telemetry.Device.pendingBounded(observed, dp.seen[index], dp.max_rows * 10)); + return total; + } + return total; + } + + /// Exact additional device allocation; tracing is fixed at init and costs zero bytes when disabled. + pub fn traceDeviceBytes(dp: *const DevPlan) usize { + return if (dp.trace_mem) |buffer| buffer.len else 0; + } + + pub fn traceOwnership(dp: *const DevPlan, stream: cuda.Stream, transfer: bool) !void { + const buffer = dp.trace_mem orelse return; + for (dp.pool_states, 0..) |entry, index| { + var args: cuda.launch.Args = .{}; + args.add(dp.state.ptr + entry.byteOffset()); + args.add(@as(c_int, @intCast(entry.lines))); + args.add(buffer.ptr + index * demand_trace.pool_bytes); + args.add(@as(c_int, @intFromBool(transfer))); + try cuda.launch.launch(dp.trace_checkpoint_fn.?, .{ .grid = .{ .x = 1 }, .block = .{ .x = 1 } }, stream, &args); + } + } + + /// Caller must join all planner/expert consumers and stop submissions before this blocking diagnostic read. + pub fn readTraceAfterConsumers(dp: *const DevPlan, index: usize, out: []u32) !demand_trace.View { + const buffer = dp.trace_mem orelse return error.TraceDisabled; + if (index >= dp.pools.len or out.len != demand_trace.pool_words) return error.Invalid; + try buffer.download(index * demand_trace.pool_bytes, std.mem.sliceAsBytes(out)); + return demand_trace.View.init(out); + } + + /// Allocating diagnostic read with the same joined/quiescent contract. + pub fn downloadTraceAfterConsumers(dp: *const DevPlan, allocator: std.mem.Allocator, index: usize) !demand_trace.Snapshot { + if (dp.trace_mem == null) return error.TraceDisabled; + if (index >= dp.pools.len) return error.Invalid; + const words = try allocator.alloc(u32, demand_trace.pool_words); + errdefer allocator.free(words); + _ = try dp.readTraceAfterConsumers(index, words); + return .{ .allocator = allocator, .words = words }; + } + + /// The in-place experts of the last plan ([512] int32) and their count after it. + fn ovfList(dp: *const DevPlan) u64 { + return dp.work.ptr + workCap(dp.max_rows) * 9 * 8 + 256; + } + + /// After the round's experts (on `s`): the in-place experts' table entries cleared. + pub fn restore(dp: *DevPlan, s: cuda.Stream, i: usize) !void { + if (i >= dp.pools.len) return error.Invalid; + const p = dp.pools[i]; + var a: cuda.Args = .{}; + a.add(dp.ovfList()); + a.add(dp.ovfList() + experts_n * 4); + for (0..3) |w| a.add(p.lines.table(w)); + for (0..3) |w| a.add(p.lines.widths(w)); + try cuda.launch.launch(dp.restore_fn, .{ .grid = .{ .x = 1 }, .block = .{ .x = 256 } }, s, &a); + } + + /// `s` waits for the last plan's copies (runLayer joins them itself; a caller of `plan` alone needs this). + pub fn join(dp: *DevPlan, s: cuda.Stream) !void { + if (dp.overlap) try s.wait(dp.ev_copy); + } + + /// The index of `p` among the planned pools, or null (the MTP layer's pool, another engine's). + pub fn indexOf(dp: *const DevPlan, p: *const pool_.Pool) ?usize { + const i = p.opt.layer; + if (i < dp.pools.len and dp.pools[i] == p) return i; + for (dp.pools, 0..) |q, j| if (q == p) return j; + return null; + } + + /// Plans pool `i`'s round of `rows` rows on the device and copies its misses into their lines, on `s`: after this + pub fn plan(dp: *DevPlan, s: cuda.Stream, i: usize, picks: u64, rows: usize) !void { + if (rows == 0 or rows > dp.max_rows) return error.PoolPasses; + if (i >= dp.pools.len) return error.Invalid; + const p = dp.pools[i]; + const entry = dp.pool_states[i]; + const st = dp.state.ptr + entry.byteOffset(); + const sb = dp.statics.ptr + i * static_bytes; + const work = dp.work.ptr; + const work_n = dp.work.ptr + workCap(dp.max_rows) * 9 * 8; + var a: cuda.Args = .{}; + a.add(picks); + a.add(@as(c_int, @intCast(rows * 11))); + a.add(st); + a.add(@as(c_int, @intCast(entry.lines))); + a.add(@as(c_int, @intFromBool(dp.recent_eviction))); + for (0..3) |w| a.add(p.lines.table(w)); + for (0..3) |w| a.add(p.lines.widths(w)); + a.add(sb); + a.add(sb + experts_n * 24); + a.add(sb + experts_n * 48); + a.add(@as(u64, p.lines.mem.ptr)); + a.add(@as(u64, p.lines.layout.line_bytes)); + for (p.lines.layout.off) |o| a.add(@as(u64, o)); + a.add(work); + a.add(work_n); + a.add(if (dp.overlap) dp.split.ptr else @as(u64, 0)); + a.add(if (dp.overlap) dp.split.ptr + dp.max_rows * 11 * 4 else @as(u64, 0)); + a.add(dp.ovfList()); + a.add(dp.ovfList() + experts_n * 4); + a.add(if (dp.trace_mem) |buffer| buffer.ptr + i * demand_trace.pool_bytes else @as(u64, 0)); + try cuda.launch.launch(dp.plan_fn, .{ .grid = .{ .x = 1 }, .block = .{ .x = 256 } }, s, &a); + var c: cuda.Args = .{}; + c.add(work); + c.add(work_n); + if (!dp.overlap) return cuda.launch.launch(dp.copy_fn, .{ .grid = .{ .x = dp.copy_blocks }, .block = .{ .x = 256 } }, s, &c); + // the copies on the side stream (a fork in a captured graph); `s` joins them before the misses' experts + try dp.ev_plan.record(s); + try dp.side.wait(dp.ev_plan); + try cuda.launch.launch(dp.copy_fn, .{ .grid = .{ .x = dp.copy_blocks }, .block = .{ .x = 256 } }, dp.side, &c); + try dp.ev_copy.record(dp.side); + } +}; + +// ---- the engine's planner ------------------------------------------------------------------------------------------- + +/// What one engine owns of the device planner: two engines in a process never share it. +pub const State = struct { + /// The planner over the engine's pools, made on its first device round. + plan: ?*DevPlan = null, + /// Set while a decode round of this engine runs with device plans (cuda_exl3_graph_hooks.zig). + active: ?*DevPlan = null, + /// The draft head's (MTP layer's) pool, registered by cuda_exl3_lanes_load after the pools attach. + mtp_pool: ?*pool_.Pool = null, +}; + +/// TF_FLASHNEXT_EXL3_DPLAN: 1 plans decode rounds on the device (implied by TF_FLASHNEXT_EXL3_GRAPH=1). +pub fn enabled() bool { + if (@import("../graph/cuda_exl3_graph.zig").force) |f| return f != .off; + if (std.c.getenv("TF_FLASHNEXT_EXL3_DPLAN")) |v| return std.mem.eql(u8, std.mem.span(v), "1"); + return @import("../graph/cuda_exl3_graph.zig").enabled(); +} + +/// The engine's device planner over its decoder layers' pools, made once; null when a layer has no pool. +pub fn forLayers(gpa: std.mem.Allocator, d: *const cuda.Driver, st: *State, layers: []const @import("../../cuda_prompt_load.zig").Layer) !?*DevPlan { + if (st.plan) |g| return g; + var pools: [64]*pool_.Pool = undefined; + var n: usize = 0; + for (layers) |*l| { + const x = l.exl3 orelse return null; + const p = x.pool orelse return null; + if (n == pools.len) return null; + pools[n] = p; + n += 1; + } + if (n == 0) return null; + // the MTP layer's pool last (its own lifetime counts; the draft chain plans it on the device) + if (st.mtp_pool) |mp| if (n < pools.len) { + pools[n] = mp; + n += 1; + }; + st.plan = try DevPlan.init(gpa, d, pools[0..n]); + return st.plan; +} + +/// The engine's planner if made. +pub fn current(st: *const State) ?*DevPlan { + return st.plan; +} + +/// Before any host plan: the engine's device residency back to its pools (no-op without a device planner). +pub fn toHost(st: *const State) !void { + if (st.plan) |g| try g.toHost(); +} + +/// The engine's deinit: its planner goes. +pub fn release(st: *State) void { + if (st.plan) |g| g.deinit(); + st.* = .{}; +} + +/// One offloaded layer's experts on the device plan: plan + copies, then `Experts.run` on the pool's tables, as the +pub fn runLayer(dp: *DevPlan, st: cuda.Stream, i: usize, 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 { + try dp.plan(st, i, picks, rows); + if (!dp.overlap) { + try ex.run(st, layer, xin, 2560, .bf16, picks, wts, rows, out, sc); + } else { + // held experts while the copies run, then the misses' experts and (expert parallel) the combine of every slot + try ex.run(st, layer, xin, 2560, .bf16, dp.split.ptr, 0, rows, 0, sc); + try st.wait(dp.ev_copy); + try ex.run(st, layer, xin, 2560, .bf16, dp.split.ptr + dp.max_rows * 11 * 4, wts, rows, out, sc); + } + try dp.restore(st, i); +} + +/// Rows a device-planned round holds at most: the MTP verify window. +pub const planner_rows: usize = 16; + +fn maxRows() usize { + return planner_rows; +} + +/// With the split the copies run beside the held experts: a quarter of the SMs, at least 8 blocks, keeps the link busy. +fn overlapCopyBlocks(d: *const cuda.Driver) !u32 { + var device: c_int = 0; + var sms: c_int = 0; + try d.check(d.api.cuDeviceGet(&device, 0), "cuDeviceGet"); + try d.check(d.api.cuDeviceGetAttribute(&sms, .multiprocessor_count, device), "cuDeviceGetAttribute"); + return @max(8, @as(u32, @intCast(sms)) / 4); +} diff --git a/zig/src/families/flashnext/exl3/dplan/cuda_exl3_dplan_test.zig b/zig/src/families/flashnext/exl3/dplan/cuda_exl3_dplan_test.zig new file mode 100644 index 0000000000..978dfd9351 --- /dev/null +++ b/zig/src/families/flashnext/exl3/dplan/cuda_exl3_dplan_test.zig @@ -0,0 +1,193 @@ +//! Host tests of the EXL3 device planner: the demand trace, its ownership checkpoints and the planner's state words. + +const std = @import("std"); +const cuda = @import("cuda"); +const pool_ = @import("../offload/cuda_exl3_pool.zig"); +const layout = @import("exl3_dplan_layout.zig"); +const dplan = @import("cuda_exl3_dplan.zig"); +const DevPlan = dplan.DevPlan; +const PoolState = dplan.PoolState; +const demand_trace = dplan.demand_trace; +const stateWords = dplan.stateWords; +const experts_n = layout.experts_n; +const recent_window = layout.recent_window; +const recent_words_per_row = experts_n / 32; +const misc_words = layout.misc_words; +const recentOffset = layout.recentOffset; +const miscOffset = layout.miscOffset; + +test "exl3: demand trace diagnostic APIs reject disabled capture" { + var dp: DevPlan = undefined; + dp.trace_mem = null; + try std.testing.expectEqual(@as(usize, 0), dp.traceDeviceBytes()); + try std.testing.expectError(error.TraceDisabled, dp.readTraceAfterConsumers(0, &.{})); + try std.testing.expectError(error.TraceDisabled, dp.downloadTraceAfterConsumers(std.testing.allocator, 0)); + try dp.traceOwnership(undefined, true); + try dp.traceOwnership(undefined, false); +} + +test "exl3: ownership trace checkpoints use ragged pool state and serialized compute stream" { + const Probe = struct { + var calls: usize = 0; + var arguments: [4][4]u64 = undefined; + var streams: [4]cuda.abi.Stream = undefined; + + fn launch(_: cuda.abi.Function, _: c_uint, _: c_uint, _: c_uint, _: c_uint, _: c_uint, _: c_uint, _: c_uint, stream: cuda.abi.Stream, parameters: ?[*]?*anyopaque, _: ?[*]?*anyopaque) callconv(.c) c_int { + inline for (.{ u64, c_int, u64, c_int }, 0..) |Type, index| { + const value: *const Type = @ptrCast(@alignCast(parameters.?[index].?)); + arguments[calls][index] = @intCast(value.*); + } + streams[calls] = stream; + calls += 1; + return 0; + } + }; + var driver: cuda.Driver = undefined; + driver.api.cuLaunchKernel = Probe.launch; + var entries = [_]PoolState{ try PoolState.init(10, 0), try PoolState.init(31, stateWords(10)) }; + var device: DevPlan = undefined; + device.pool_states = &entries; + device.state.ptr = 0x10000; + device.trace_mem = .{ .d = &driver, .ptr = 0x100000, .len = 2 * demand_trace.pool_bytes, .borrowed = true }; + device.trace_checkpoint_fn = .{ .d = &driver, .handle = @ptrFromInt(1) }; + const stream: cuda.Stream = .{ .d = &driver, .handle = @ptrFromInt(2) }; + Probe.calls = 0; + try device.traceOwnership(stream, true); + try device.traceOwnership(stream, false); + for (0..4) |call| { + const index = call % 2; + try std.testing.expectEqual([4]u64{ device.state.ptr + entries[index].byteOffset(), entries[index].lines, 0x100000 + index * demand_trace.pool_bytes, @intFromBool(call < 2) }, Probe.arguments[call]); + try std.testing.expectEqual(stream.handle, Probe.streams[call]); + } + device.trace_mem = null; + device.trace_checkpoint_fn = null; + try device.traceOwnership(stream, true); + try std.testing.expectEqual(@as(usize, 4), Probe.calls); +} + +test "exl3: recent-demand planner state is separate from lifetime profile counts" { + try std.testing.expectEqual(@as(usize, 3 * experts_n), recentOffset(10) - 20); + try std.testing.expectEqual(@as(usize, recent_window * recent_words_per_row), miscOffset(10) - recentOffset(10)); + try std.testing.expectEqual(@as(usize, misc_words), stateWords(10) - miscOffset(10)); +} + +test "exl3: demand trace planner arguments stay fixed and joined reads isolate pools" { + const Probe = struct { + var trace_address: u64 = 0; + var launches: usize = 0; + + fn launch(function: cuda.abi.Function, _: c_uint, _: c_uint, _: c_uint, _: c_uint, _: c_uint, _: c_uint, _: c_uint, _: cuda.abi.Stream, parameters: ?[*]?*anyopaque, _: ?[*]?*anyopaque) callconv(.c) c_int { + launches += 1; + if (@intFromPtr(function.?) == 1) { + const address: *const u64 = @ptrCast(@alignCast(parameters.?[25].?)); + trace_address = address.*; + } + return 0; + } + + fn download(destination: ?*anyopaque, source: u64, length: usize) callconv(.c) c_int { + const target: [*]u8 = @ptrCast(destination.?); + const origin: [*]const u8 = @ptrFromInt(source); + @memcpy(target[0..length], origin[0..length]); + return 0; + } + }; + var driver: cuda.Driver = undefined; + driver.api.cuLaunchKernel = Probe.launch; + driver.api.cuMemcpyDtoH_v2 = Probe.download; + const words = try std.testing.allocator.alloc(u32, 2 * demand_trace.pool_words); + defer std.testing.allocator.free(words); + for (0..2) |index| @memcpy(words[index * demand_trace.pool_words ..][0..demand_trace.header_words], &try demand_trace.header(index, index, 10, 16)); + var resident: pool_.Pool = undefined; + resident.lines.tables.ptr = 4096; + resident.lines.mem.ptr = 8192; + resident.lines.tlay = cuda.exl3_pool.Tables.init(513); + resident.lines.layout = cuda.exl3_pool.Layout.init(.{ 256, 256, 256 }); + var pools = [_]*pool_.Pool{ &resident, &resident }; + var states = [_]PoolState{ try PoolState.init(10, 0), try PoolState.init(10, stateWords(10)) }; + var device: DevPlan = undefined; + device.pools = &pools; + device.pool_states = &states; + device.max_rows = 16; + device.state.ptr = 16384; + device.statics.ptr = 32768; + device.work.ptr = 65536; + device.plan_fn = .{ .d = &driver, .handle = @ptrFromInt(1) }; + device.copy_fn = .{ .d = &driver, .handle = @ptrFromInt(2) }; + device.copy_blocks = 1; + device.overlap = false; + device.trace_mem = .{ .d = &driver, .ptr = @intFromPtr(words.ptr), .len = words.len * 4, .borrowed = true }; + const out = try std.testing.allocator.alloc(u32, demand_trace.pool_words); + defer std.testing.allocator.free(out); + const stream: cuda.Stream = .{ .d = &driver, .handle = null }; + Probe.launches = 0; + for (0..2) |index| { + for (0..2) |_| { + try device.plan(stream, index, 1024, 1); + try std.testing.expectEqual(@intFromPtr(words.ptr) + index * demand_trace.pool_bytes, Probe.trace_address); + } + const view = try device.readTraceAfterConsumers(index, out); + try std.testing.expectEqual(@as(usize, 0), view.recorded()); + try std.testing.expectEqual(@as(u32, @intCast(index)), view.words[2]); + } + try std.testing.expectEqual(@as(usize, 8), Probe.launches); + try std.testing.expectEqual(words.len * 4, device.traceDeviceBytes()); + try std.testing.expectError(error.Invalid, device.readTraceAfterConsumers(2, out)); + try std.testing.expectError(error.Invalid, device.readTraceAfterConsumers(0, out[1..])); + device.trace_mem = null; + try device.plan(stream, 0, 1024, 1); + try std.testing.expectEqual(@as(u64, 0), Probe.trace_address); + try std.testing.expectEqual(@as(usize, 0), device.traceDeviceBytes()); + try std.testing.expectError(error.TraceDisabled, device.readTraceAfterConsumers(0, out)); +} + +test "exl3: two engines in one process keep separate planner state" { + const Layer = @import("../../cuda_prompt_load.zig").Layer; + const driver: *const cuda.Driver = undefined; // never reached: no layer here has a pool + var layers: [1]Layer = undefined; + layers[0].exl3 = null; + const plan_b: *DevPlan = @ptrFromInt(0x2000); + const pool_b: *pool_.Pool = @ptrFromInt(0x3000); + var a: dplan.State = .{}; + var b: dplan.State = .{ .plan = plan_b, .active = plan_b, .mtp_pool = pool_b }; + // engine a has no pools: it plans nothing, and engine b keeps the planner it made + try std.testing.expectEqual(@as(?*DevPlan, null), try dplan.forLayers(std.testing.allocator, driver, &a, &layers)); + try std.testing.expectEqual(@as(?*DevPlan, plan_b), try dplan.forLayers(std.testing.allocator, driver, &b, &layers)); + try std.testing.expectEqual(@as(?*DevPlan, null), dplan.current(&a)); + try std.testing.expectEqual(@as(?*DevPlan, plan_b), dplan.current(&b)); + // a round of a enters and leaves its plan; b's active plan and head pool are untouched + a.active = @ptrFromInt(0x1000); + a.mtp_pool = @ptrFromInt(0x4000); + try std.testing.expect(b.active.? == plan_b and b.mtp_pool.? == pool_b); + a.active = null; + try dplan.toHost(&a); // nothing planned: no device work, and b is not asked either + dplan.release(&a); + try std.testing.expectEqual(@as(?*pool_.Pool, null), a.mtp_pool); + try std.testing.expect(b.active.? == plan_b and b.mtp_pool.? == pool_b and b.plan.? == plan_b); +} + +test "exl3: every pass built from an engine carries the engine's planner state" { + // a pass without it runs the host pool plan without pulling the device planner's residency back first + const sources = [_][]const u8{ + @embedFile("../../cuda_prompt.zig"), + @embedFile("../../cuda_mtp_ops.zig"), + @embedFile("../prefill/cuda_exl3_prefill.zig"), + @embedFile("../prefill/cuda_exl3_prefill_head.zig"), + @embedFile("../model/cuda_exl3_chain.zig"), + }; + for (sources) |text| { + var literals: usize = 0; + var with_state: usize = 0; + var from: usize = 0; + while (std.mem.indexOfPos(u8, text, from, ".table_dense = ")) |at| { + const end = std.mem.indexOfPos(u8, text, at, "\n") orelse text.len; + literals += 1; + if (std.mem.indexOf(u8, text[at..end], ".dplan = &e.exl3_dplan") != null) with_state += 1; + from = end; + } + // the chain sets it on the pass after building it + if (literals == 1 and std.mem.indexOf(u8, text, "p.m.dplan = &e.exl3_dplan") != null) continue; + try std.testing.expect(literals > 0); + try std.testing.expectEqual(literals, with_state); + } +} diff --git a/zig/src/families/flashnext/exl3/dplan/exl3_dplan_layout.zig b/zig/src/families/flashnext/exl3/dplan/exl3_dplan_layout.zig new file mode 100644 index 0000000000..e0d2da8d62 --- /dev/null +++ b/zig/src/families/flashnext/exl3/dplan/exl3_dplan_layout.zig @@ -0,0 +1,128 @@ +const std = @import("std"); + +pub const experts_n: usize = 512; +pub const recent_window: usize = 256; +pub const recent_words: usize = recent_window * (experts_n / 32); +pub const misc_words: usize = 10; +pub const static_bytes: usize = experts_n * 3 * (8 + 8 + 4); + +pub fn recentCountsOffset(lines: usize) usize { + return 2 * lines + 2 * experts_n; +} + +pub fn recentOffset(lines: usize) usize { + return 2 * lines + 3 * experts_n; +} + +pub fn miscOffset(lines: usize) usize { + return recentOffset(lines) + recent_words; +} + +pub fn stateWords(lines: usize) usize { + return miscOffset(lines) + misc_words; +} + +pub const PoolState = struct { + lines: usize, + offset_words: usize, + words: usize, + + pub fn init(lines: usize, offset_words: usize) !PoolState { + if (lines == 0 or lines > (std.math.maxInt(c_int) - stateWords(0)) / 2) return error.Invalid; + const words = stateWords(lines); + const end = try std.math.add(usize, offset_words, words); + _ = try std.math.mul(usize, end, @sizeOf(u32)); + return .{ .lines = lines, .offset_words = offset_words, .words = words }; + } + + pub fn endWords(self: PoolState) usize { + return self.offset_words + self.words; + } + + pub fn byteOffset(self: PoolState) usize { + return self.offset_words * @sizeOf(u32); + } + + pub fn slice(self: PoolState, all: []u32) []u32 { + return all[self.offset_words..self.endWords()]; + } +}; + +pub fn deviceBytes(layers: u64, total_lines: u64, rows: u64) !u64 { + const fixed = try std.math.mul(u64, layers, stateWords(0) * @sizeOf(u32) + static_bytes); + const state_lines = try std.math.mul(u64, total_lines, 2 * @sizeOf(u32)); + const work = try std.math.add(u64, try std.math.mul(u64, rows, 10 * 9 * 8 + 2 * 11 * 4), 256 + (experts_n + 64) * 4); + return std.math.add(u64, try std.math.add(u64, fixed, state_lines), work); +} + +test "exl3: ragged planner offsets are contiguous and exactly bounded" { + const counts = [_]usize{ 10, 73, 11, 32 }; + var end: usize = 0; + for (counts) |lines| { + const entry = try PoolState.init(lines, end); + try std.testing.expectEqual(end, entry.offset_words); + try std.testing.expectEqual(stateWords(lines), entry.words); + try std.testing.expectEqual(end * 4, entry.byteOffset()); + try std.testing.expectEqual(entry.words, miscOffset(lines) + misc_words); + end = entry.endWords(); + } + try std.testing.expectEqual(counts.len * stateWords(0) + 2 * (10 + 73 + 11 + 32), end); + try std.testing.expectError(error.Invalid, PoolState.init(0, 0)); + const max_lines = (std.math.maxInt(c_int) - stateWords(0)) / 2; + _ = try PoolState.init(max_lines, 0); + try std.testing.expectError(error.Invalid, PoolState.init(max_lines + 1, 0)); + try std.testing.expectError(error.Invalid, PoolState.init(std.math.maxInt(usize), 0)); + try std.testing.expectError(error.Overflow, PoolState.init(10, std.math.maxInt(usize))); + try std.testing.expectError(error.Overflow, PoolState.init(10, std.math.maxInt(usize) / 4)); + var uniform_end: usize = 0; + for (0..4) |index| { + const entry = try PoolState.init(32, uniform_end); + try std.testing.expectEqual(index * stateWords(32), entry.offset_words); + uniform_end = entry.endWords(); + } +} + +test "exl3: ragged planner state round trip isolates line maps history and counters" { + const counts = [_]usize{ 10, 57, 13 }; + var entries: [counts.len]PoolState = undefined; + var end: usize = 1; + for (counts, &entries) |lines, *entry| { + entry.* = try PoolState.init(lines, end); + end = entry.endWords(); + } + const host = try std.testing.allocator.alloc(u32, end + 1); + defer std.testing.allocator.free(host); + @memset(host, 0xfeedface); + for (entries, 0..) |entry, index| { + const state = entry.slice(host); + @memset(state, 0); + for (0..entry.lines) |line| { + const expert = (line + index * 100) % experts_n; + state[line] = @intCast(expert); + state[entry.lines + line] = @intCast(line + 1); + state[2 * entry.lines + expert] = @intCast(line); + } + state[recentCountsOffset(entry.lines)] = @intCast(index + 1); + state[recentOffset(entry.lines)] = @intCast(index + 10); + state[miscOffset(entry.lines)] = @intCast(index + 20); + state[state.len - 1] = @intCast(index + 30); + } + const device = try std.testing.allocator.dupe(u32, host); + defer std.testing.allocator.free(device); + @memset(host[1..end], 0); + for (entries) |entry| @memcpy(entry.slice(host), entry.slice(device)); + try std.testing.expectEqual(@as(u32, 0xfeedface), host[0]); + try std.testing.expectEqual(@as(u32, 0xfeedface), host[end]); + for (entries, 0..) |entry, index| { + const state = entry.slice(host); + for (0..entry.lines) |line| { + const expert = state[line]; + try std.testing.expectEqual(@as(u32, @intCast(line)), state[2 * entry.lines + expert]); + try std.testing.expectEqual(@as(u32, @intCast(line + 1)), state[entry.lines + line]); + } + try std.testing.expectEqual(@as(u32, @intCast(index + 1)), state[recentCountsOffset(entry.lines)]); + try std.testing.expectEqual(@as(u32, @intCast(index + 10)), state[recentOffset(entry.lines)]); + try std.testing.expectEqual(@as(u32, @intCast(index + 20)), state[miscOffset(entry.lines)]); + try std.testing.expectEqual(@as(u32, @intCast(index + 30)), state[state.len - 1]); + } +} diff --git a/zig/src/families/flashnext/exl3/dplan/exl3_dplan_policy_eval.zig b/zig/src/families/flashnext/exl3/dplan/exl3_dplan_policy_eval.zig new file mode 100644 index 0000000000..90cbb03af2 --- /dev/null +++ b/zig/src/families/flashnext/exl3/dplan/exl3_dplan_policy_eval.zig @@ -0,0 +1,374 @@ +//! CPU counterfactuals for one ordered pool-demand timeline. + +const std = @import("std"); +const trace = @import("exl3_dplan_trace.zig"); + +pub const expert_count: usize = 512; +pub const max_unique: usize = 256; + +pub const Policy = enum { lifetime, lru, decayed_frequency, decayed_admission }; + +pub const Seed = struct { + line_ids: []const i32, + stamps: []const u32, + counts: [expert_count]u32, + now: u32, +}; + +pub const Options = struct { + capacity: usize, + seed: ?Seed = null, + decay_interval_rounds: u64 = 64, + decay_shift: u5 = 1, +}; + +pub const Metrics = struct { + rounds: u64 = 0, + hit_picks: u64 = 0, + miss_picks: u64 = 0, + hit_experts: u64 = 0, + miss_experts: u64 = 0, + admissions: u64 = 0, + copies: u64 = 0, + evictions: u64 = 0, + mapped_experts: u64 = 0, + mapped_picks: u64 = 0, + + fn plus(self: Metrics, other: Metrics) !Metrics { + var result: Metrics = undefined; + inline for (@typeInfo(Metrics).@"struct".field_names) |field| { + @field(result, field) = try std.math.add(u64, @field(self, field), @field(other, field)); + } + return result; + } +}; + +pub const Comparison = struct { + lifetime: Metrics, + lru: Metrics, + decayed_frequency: Metrics, + decayed_admission: Metrics, + belady_unique_misses_lower_bound: u64, +}; + +const Demand = struct { + counts: [expert_count]u32 = @splat(0), + ordered: [expert_count]usize = undefined, + unique: usize = 0, + valid_picks: u64 = 0, + picks: []const u32, +}; + +pub const State = struct { + allocator: std.mem.Allocator, + policy: Policy, + options: Options, + line_ids: []i32, + stamps: []u32, + scores: [expert_count]u32, + last_access: [expert_count]u64, + access_clock: u64, + now: u32, + total: Metrics = .{}, + + pub fn init(allocator: std.mem.Allocator, policy: Policy, options: Options) !State { + try validateOptions(options); + const line_ids = try allocator.alloc(i32, options.capacity); + errdefer allocator.free(line_ids); + const stamps = try allocator.alloc(u32, options.capacity); + @memset(line_ids, -1); + @memset(stamps, 0); + var self: State = .{ + .allocator = allocator, + .policy = policy, + .options = options, + .line_ids = line_ids, + .stamps = stamps, + .scores = @splat(0), + .last_access = @splat(0), + .access_clock = 0, + .now = 0, + }; + if (options.seed) |seed| { + @memcpy(line_ids, seed.line_ids); + @memcpy(stamps, seed.stamps); + self.scores = seed.counts; + self.now = seed.now; + self.access_clock = seed.now; + for (line_ids, stamps) |expert, stamp| { + self.access_clock = @max(self.access_clock, stamp); + if (expert >= 0) self.last_access[@intCast(expert)] = stamp; + } + } + self.options.seed = null; + return self; + } + + pub fn deinit(self: *State) void { + self.allocator.free(self.line_ids); + self.allocator.free(self.stamps); + self.* = undefined; + } + + pub fn observe(self: *State, record: trace.Round) !Metrics { + const demand = try readDemand(record); + const round_number = try std.math.add(u64, self.total.rounds, 1); + const next_clock = try std.math.add(u64, self.access_clock, demand.valid_picks); + var delta: Metrics = .{ .rounds = 1 }; + for (demand.ordered[0..demand.unique]) |expert| { + if (self.lineOf(expert) != null) { + delta.hit_picks += demand.counts[expert]; + delta.hit_experts += 1; + } else { + delta.miss_picks += demand.counts[expert]; + delta.miss_experts += 1; + } + } + _ = try self.total.plus(delta); + if ((self.policy == .decayed_frequency or self.policy == .decayed_admission) and round_number % self.options.decay_interval_rounds == 0) { + for (&self.scores) |*score| score.* >>= self.options.decay_shift; + } + self.now +%= 1; + for (demand.ordered[0..demand.unique]) |expert| { + self.scores[expert] +|= demand.counts[expert]; + if (self.lineOf(expert)) |line| self.stamps[line] = self.now; + } + for (demand.picks) |bits| { + const expert = routed(bits) orelse continue; + self.access_clock += 1; + self.last_access[expert] = self.access_clock; + } + std.debug.assert(self.access_clock == next_clock); + + var misses: [max_unique]usize = undefined; + var miss_count: usize = 0; + for (demand.ordered[0..demand.unique]) |expert| { + if (self.lineOf(expert) == null) { + misses[miss_count] = expert; + miss_count += 1; + } + } + var index: usize = 1; + while (index < miss_count) : (index += 1) { + const expert = misses[index]; + var position = index; + while (position > 0 and self.preferAdmission(expert, misses[position - 1], &demand)) : (position -= 1) { + misses[position] = misses[position - 1]; + } + misses[position] = expert; + } + for (misses[0..miss_count]) |expert| { + if (self.victim(&demand)) |line| { + if (self.policy == .decayed_admission and self.line_ids[line] >= 0 and self.scores[expert] < self.scores[@intCast(self.line_ids[line])]) { + delta.mapped_experts += 1; + delta.mapped_picks += demand.counts[expert]; + continue; + } + if (self.line_ids[line] >= 0) delta.evictions += 1; + self.line_ids[line] = @intCast(expert); + self.stamps[line] = self.now; + delta.admissions += 1; + delta.copies += 1; + } else { + delta.mapped_experts += 1; + delta.mapped_picks += demand.counts[expert]; + } + } + self.total = try self.total.plus(delta); + return delta; + } + + fn lineOf(self: *const State, expert: usize) ?usize { + for (self.line_ids, 0..) |resident, line| { + if (resident == @as(i32, @intCast(expert))) return line; + } + return null; + } + + fn preferAdmission(self: *const State, left: usize, right: usize, demand: *const Demand) bool { + if (self.policy == .lru) { + if (self.last_access[left] != self.last_access[right]) return self.last_access[left] > self.last_access[right]; + } else { + if (self.scores[left] != self.scores[right]) return self.scores[left] > self.scores[right]; + if (demand.counts[left] != demand.counts[right]) return demand.counts[left] > demand.counts[right]; + } + return left < right; + } + + fn victim(self: *const State, demand: *const Demand) ?usize { + var coldest: ?usize = null; + for (self.line_ids, 0..) |expert, line| { + if (expert < 0) return line; + if (demand.counts[@intCast(expert)] != 0) continue; + if (coldest == null or self.colder(line, coldest.?)) coldest = line; + } + return coldest; + } + + fn colder(self: *const State, left: usize, right: usize) bool { + const left_expert: usize = @intCast(self.line_ids[left]); + const right_expert: usize = @intCast(self.line_ids[right]); + if (self.policy != .lru and self.scores[left_expert] != self.scores[right_expert]) return self.scores[left_expert] < self.scores[right_expert]; + if (self.policy == .lifetime) { + if (self.stamps[left] != self.stamps[right]) return self.stamps[left] < self.stamps[right]; + } else { + if (self.last_access[left_expert] != self.last_access[right_expert]) return self.last_access[left_expert] < self.last_access[right_expert]; + } + return left > right; + } +}; + +pub fn evaluate(allocator: std.mem.Allocator, rounds: []const trace.Round, options: Options) !Comparison { + var lifetime = try State.init(allocator, .lifetime, options); + defer lifetime.deinit(); + var lru = try State.init(allocator, .lru, options); + defer lru.deinit(); + var decayed = try State.init(allocator, .decayed_frequency, options); + defer decayed.deinit(); + var decayed_admission = try State.init(allocator, .decayed_admission, options); + defer decayed_admission.deinit(); + for (rounds) |record| { + _ = try lifetime.observe(record); + _ = try lru.observe(record); + _ = try decayed.observe(record); + _ = try decayed_admission.observe(record); + } + return .{ + .lifetime = lifetime.total, + .lru = lru.total, + .decayed_frequency = decayed.total, + .decayed_admission = decayed_admission.total, + .belady_unique_misses_lower_bound = try beladyLowerBound(allocator, rounds, options), + }; +} + +pub fn evaluateTrace(allocator: std.mem.Allocator, view: trace.View) !Comparison { + const verification = try @import("exl3_dplan_replay.zig").verifyTrace(allocator, view); + const rounds = try allocator.alloc(trace.Round, view.recorded()); + defer allocator.free(rounds); + for (rounds, 0..) |*round, index| round.* = try view.round(index); + var options: Options = .{ .capacity = view.words[4], .seed = view.seedState(Seed) }; + var total: Comparison = .{ .lifetime = .{}, .lru = .{}, .decayed_frequency = .{}, .decayed_admission = .{}, .belady_unique_misses_lower_bound = 0 }; + var begin: usize = 0; + for (0..view.checkpoints() + 1) |index| { + const checkpoint = if (index < view.checkpoints()) try view.checkpoint(index) else null; + const end = if (checkpoint) |entry| entry.roundIndex() else view.recorded(); + if (end > begin) { + const segment = try evaluate(allocator, rounds[begin..end], options); + inline for (.{ "lifetime", "lru", "decayed_frequency", "decayed_admission" }) |field| { + @field(total, field) = try @field(total, field).plus(@field(segment, field)); + } + total.belady_unique_misses_lower_bound = try std.math.add(u64, total.belady_unique_misses_lower_bound, segment.belady_unique_misses_lower_bound); + } + if (checkpoint) |entry| options.seed = trace.seedFromState(Seed, view.words[4], entry.after()); + begin = end; + } + const baseline = total.lifetime; + const exact = verification.delta; + if (baseline.rounds != @as(u64, @intCast(verification.rounds)) or + baseline.hit_picks != exact.hit_picks or + baseline.miss_picks != exact.miss_picks or + baseline.copies != exact.copies or + baseline.evictions != exact.evictions or + baseline.mapped_experts != exact.inplace or + baseline.mapped_picks != exact.mapped_picks) + return error.PolicyBaselineMismatch; + return total; +} + +/// Belady with optional admission is optimal for uniform-cost sequential requests: on a full-cache miss. +pub fn beladyLowerBound(allocator: std.mem.Allocator, rounds: []const trace.Round, options: Options) !u64 { + try validateOptions(options); + var count: usize = 0; + for (rounds) |record| count = try std.math.add(usize, count, (try readDemand(record)).unique); + const events = try allocator.alloc(Event, count); + defer allocator.free(events); + var offset: usize = 0; + for (rounds) |record| { + const demand = try readDemand(record); + for (demand.ordered[0..demand.unique]) |expert| { + events[offset] = .{ .expert = expert }; + offset += 1; + } + } + return belady(events, options.capacity, if (options.seed) |seed| seed.line_ids else &.{}); +} + +pub const Event = struct { expert: usize, next: usize = std.math.maxInt(usize) }; + +pub fn belady(events: []Event, capacity: usize, initial: []const i32) u64 { + var next_use: [expert_count]usize = @splat(std.math.maxInt(usize)); + var position = events.len; + while (position > 0) { + position -= 1; + const event = &events[position]; + event.next = next_use[event.expert]; + next_use[event.expert] = position; + } + var resident: [expert_count]i32 = @splat(-1); + @memcpy(resident[0..initial.len], initial); + var misses: u64 = 0; + for (events) |event| { + var hit = false; + var empty: ?usize = null; + var farthest: ?usize = null; + for (resident[0..capacity], 0..) |expert, line| { + if (expert == @as(i32, @intCast(event.expert))) hit = true; + if (expert < 0) { + if (empty == null) empty = line; + } else if (farthest == null or next_use[@intCast(expert)] > next_use[@intCast(resident[farthest.?])]) { + farthest = line; + } + } + next_use[event.expert] = event.next; + if (hit) continue; + misses += 1; + if (empty) |line| { + resident[line] = @intCast(event.expert); + } else if (farthest) |line| { + if (event.next < next_use[@intCast(resident[line])]) resident[line] = @intCast(event.expert); + } + } + return misses; +} + +fn validateOptions(options: Options) !void { + if (options.capacity > expert_count or options.decay_interval_rounds == 0 or options.decay_shift == 0) return error.InvalidOptions; + if (options.seed) |seed| { + if (seed.line_ids.len != options.capacity or seed.stamps.len != options.capacity) return error.InvalidSeed; + var seen: [expert_count]bool = @splat(false); + for (seed.line_ids) |expert| { + if (expert < -1 or expert >= expert_count) return error.InvalidSeed; + if (expert >= 0) { + const index: usize = @intCast(expert); + if (seen[index]) return error.InvalidSeed; + seen[index] = true; + } + } + } +} + +fn readDemand(record: trace.Round) !Demand { + const shape = trace.header(0, 0, record.lines, record.max_rows) catch return error.InvalidRound; + if (record.words.len != shape[7]) return error.InvalidRound; + if (record.recentEnabled()) return error.UnsupportedRecentPolicy; + if (record.rows() == 0 or record.rows() > record.max_rows or record.words[3] > 1) return error.InvalidRound; + var demand: Demand = .{ .picks = record.picks() }; + for (demand.picks) |bits| { + const expert = routed(bits) orelse continue; + if (demand.counts[expert] == 0) { + demand.ordered[demand.unique] = expert; + demand.unique += 1; + } + demand.counts[expert] += 1; + demand.valid_picks += 1; + } + if (demand.valid_picks != record.validPicks() or demand.unique != record.uniqueCount()) return error.InvalidRound; + if (demand.unique > max_unique) return error.PlannerOverflow; + return demand; +} + +pub fn routed(bits: u32) ?usize { + const expert: i32 = @bitCast(bits); + return if (expert >= 0 and expert < expert_count) @intCast(expert) else null; +} diff --git a/zig/src/families/flashnext/exl3/dplan/exl3_dplan_policy_eval_test.zig b/zig/src/families/flashnext/exl3/dplan/exl3_dplan_policy_eval_test.zig new file mode 100644 index 0000000000..f6bc78a536 --- /dev/null +++ b/zig/src/families/flashnext/exl3/dplan/exl3_dplan_policy_eval_test.zig @@ -0,0 +1,289 @@ +//! Host tests of the CPU policy counterfactuals: round hits and admissions, decay, LRU, Belady and bad input. + +const std = @import("std"); +const trace = @import("exl3_dplan_trace.zig"); +const policy = @import("exl3_dplan_policy_eval.zig"); +const expert_count = policy.expert_count; +const State = policy.State; +const Metrics = policy.Metrics; +const Options = policy.Options; +const Event = policy.Event; +const evaluate = policy.evaluate; +const belady = policy.belady; +const routed = policy.routed; + +const Synthetic = struct { + words: []u32, + max_rows: usize, + + fn init(picks: []const i32) !Synthetic { + const rows = @max(1, (picks.len + 10) / 11); + const shape = try trace.header(0, 0, 2, rows); + const words = try std.testing.allocator.alloc(u32, shape[7]); + @memset(words, 0); + words[0] = @intCast(rows); + const self: Synthetic = .{ .words = words, .max_rows = rows }; + const raw_offset = trace.record_header_words; + @memset(words[raw_offset..][0 .. rows * 11], std.math.maxInt(u32)); + var counts: [expert_count]u32 = @splat(0); + for (picks, 0..) |expert, index| { + words[raw_offset + index] = @bitCast(expert); + if (routed(@bitCast(expert))) |valid| { + if (counts[valid] == 0) words[2] += 1; + counts[valid] += 1; + words[1] += 1; + } + } + return self; + } + + fn round(self: Synthetic) trace.Round { + return .{ .words = self.words, .lines = 2, .max_rows = self.max_rows }; + } + + fn deinit(self: Synthetic) void { + std.testing.allocator.free(self.words); + } +}; + +test "exl3: policy evaluation pins round hits and admissions and separates mapped misses" { + var fixture = try Synthetic.init(&.{ 0, 2, 2, 2, 3, 512, -1, 999, -9 }); + defer fixture.deinit(); + var state = try State.init(std.testing.allocator, .lifetime, .{ + .capacity = 2, + .seed = .{ .line_ids = &.{ 0, 1 }, .stamps = &.{ 1, 2 }, .counts = @splat(0), .now = 2 }, + }); + defer state.deinit(); + const result = try state.observe(fixture.round()); + try std.testing.expectEqual(Metrics{ .rounds = 1, .hit_picks = 1, .miss_picks = 4, .hit_experts = 1, .miss_experts = 2, .admissions = 1, .copies = 1, .evictions = 1, .mapped_experts = 1, .mapped_picks = 1 }, result); + try std.testing.expectEqualSlices(i32, &.{ 0, 2 }, state.line_ids); +} + +test "exl3: lifetime admission saturation and victim ties match current planner" { + var fixture = try Synthetic.init(&.{ 2, 2, 2, 2, 2, 3 }); + defer fixture.deinit(); + var counts: [expert_count]u32 = @splat(0); + counts[0] = 5; + counts[1] = 3; + counts[2] = std.math.maxInt(u32) - 2; + var state = try State.init(std.testing.allocator, .lifetime, .{ + .capacity = 2, + .seed = .{ .line_ids = &.{ 0, 1 }, .stamps = &.{ 1, 2 }, .counts = counts, .now = std.math.maxInt(u32) }, + }); + defer state.deinit(); + const result = try state.observe(fixture.round()); + try std.testing.expectEqual(@as(u64, 2), result.copies); + try std.testing.expectEqualSlices(i32, &.{ 3, 2 }, state.line_ids); + try std.testing.expectEqual(std.math.maxInt(u32), state.scores[2]); + try std.testing.expectEqual(@as(u32, 0), state.now); + + var tie = try State.init(std.testing.allocator, .lifetime, .{ + .capacity = 2, + .seed = .{ .line_ids = &.{ 0, 1 }, .stamps = &.{ 7, 7 }, .counts = @splat(0), .now = 7 }, + }); + defer tie.deinit(); + var single = try Synthetic.init(&.{9}); + defer single.deinit(); + _ = try tie.observe(single.round()); + try std.testing.expectEqualSlices(i32, &.{ 0, 9 }, tie.line_ids); +} + +test "exl3: LRU preserves raw last occurrence order when the round exceeds capacity" { + var first = try Synthetic.init(&.{ 1, 2, 1 }); + defer first.deinit(); + var second = try Synthetic.init(&.{ 1, 1, 2 }); + defer second.deinit(); + var left = try State.init(std.testing.allocator, .lru, .{ .capacity = 1 }); + defer left.deinit(); + var right = try State.init(std.testing.allocator, .lru, .{ .capacity = 1 }); + defer right.deinit(); + _ = try left.observe(first.round()); + _ = try right.observe(second.round()); + try std.testing.expectEqualSlices(i32, &.{1}, left.line_ids); + try std.testing.expectEqualSlices(i32, &.{2}, right.line_ids); + _ = try left.observe(second.round()); + try std.testing.expectEqualSlices(i32, &.{1}, left.line_ids); +} + +test "exl3: frequency decay halves scores at configured round boundaries before new demand" { + var fixture = try Synthetic.init(&.{ 0, 1, 1, 1 }); + defer fixture.deinit(); + var empty = try Synthetic.init(&.{}); + defer empty.deinit(); + var counts: [expert_count]u32 = @splat(0); + counts[0] = 64; + const options: Options = .{ + .capacity = 1, + .seed = .{ .line_ids = &.{-1}, .stamps = &.{0}, .counts = counts, .now = 0 }, + .decay_interval_rounds = 1, + .decay_shift = 6, + }; + var lifetime = try State.init(std.testing.allocator, .lifetime, options); + defer lifetime.deinit(); + var decayed = try State.init(std.testing.allocator, .decayed_frequency, options); + defer decayed.deinit(); + _ = try lifetime.observe(fixture.round()); + _ = try decayed.observe(fixture.round()); + try std.testing.expectEqualSlices(i32, &.{0}, lifetime.line_ids); + try std.testing.expectEqualSlices(i32, &.{1}, decayed.line_ids); + + var interval = try State.init(std.testing.allocator, .decayed_frequency, .{ .capacity = 0, .decay_interval_rounds = 2 }); + defer interval.deinit(); + _ = try interval.observe(fixture.round()); + try std.testing.expectEqual(@as(u32, 3), interval.scores[1]); + _ = try interval.observe(empty.round()); + try std.testing.expectEqual(@as(u32, 1), interval.scores[1]); +} + +test "exl3: LRU evicts by within-round recency and decayed eviction forgets old frequency" { + var touching = try Synthetic.init(&.{ 0, 1, 0 }); + defer touching.deinit(); + var incoming = try Synthetic.init(&.{2}); + defer incoming.deinit(); + var empty = try Synthetic.init(&.{}); + defer empty.deinit(); + var counts: [expert_count]u32 = @splat(0); + counts[0] = 128; + counts[1] = 4; + const options: Options = .{ + .capacity = 2, + .seed = .{ .line_ids = &.{ 0, 1 }, .stamps = &.{ 1, 2 }, .counts = counts, .now = 2 }, + .decay_interval_rounds = 1, + .decay_shift = 6, + }; + var lru = try State.init(std.testing.allocator, .lru, options); + defer lru.deinit(); + _ = try lru.observe(touching.round()); + _ = try lru.observe(incoming.round()); + try std.testing.expectEqualSlices(i32, &.{ 0, 2 }, lru.line_ids); + + var lifetime = try State.init(std.testing.allocator, .lifetime, options); + defer lifetime.deinit(); + var decayed = try State.init(std.testing.allocator, .decayed_frequency, options); + defer decayed.deinit(); + for (0..2) |_| { + _ = try lifetime.observe(empty.round()); + _ = try decayed.observe(empty.round()); + } + _ = try lifetime.observe(incoming.round()); + _ = try decayed.observe(incoming.round()); + try std.testing.expectEqualSlices(i32, &.{ 0, 2 }, lifetime.line_ids); + try std.testing.expectEqualSlices(i32, &.{ 2, 1 }, decayed.line_ids); +} + +test "exl3: frequency admission bypasses a colder newcomer and later admits it when hot" { + var warm = try Synthetic.init(&.{ 1, 1 }); + defer warm.deinit(); + var hot = try Synthetic.init(&.{ 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1 }); + defer hot.deinit(); + var counts: [expert_count]u32 = @splat(0); + counts[0] = 10; + var state = try State.init(std.testing.allocator, .decayed_admission, .{ + .capacity = 1, + .seed = .{ .line_ids = &.{0}, .stamps = &.{1}, .counts = counts, .now = 1 }, + }); + defer state.deinit(); + const cold = try state.observe(warm.round()); + try std.testing.expectEqual(@as(u64, 0), cold.copies); + try std.testing.expectEqual(@as(u64, 1), cold.mapped_experts); + try std.testing.expectEqualSlices(i32, &.{0}, state.line_ids); + const now_hot = try state.observe(hot.round()); + try std.testing.expectEqual(@as(u64, 1), now_hot.copies); + try std.testing.expectEqual(@as(u64, 1), now_hot.evictions); + try std.testing.expectEqualSlices(i32, &.{1}, state.line_ids); +} + +test "exl3: comparison evolves independently from the initial seed and supports zero capacity" { + var fixture = try Synthetic.init(&.{ 5, 5 }); + defer fixture.deinit(); + const rounds = [_]trace.Round{ fixture.round(), fixture.round() }; + const result = try evaluate(std.testing.allocator, &rounds, .{ .capacity = 1 }); + try std.testing.expectEqual(@as(u64, 2), result.lifetime.miss_picks); + try std.testing.expectEqual(@as(u64, 2), result.lifetime.hit_picks); + try std.testing.expectEqual(@as(u64, 1), result.lifetime.copies); + try std.testing.expectEqual(result.lifetime, result.lru); + try std.testing.expectEqual(result.lifetime, result.decayed_frequency); + try std.testing.expectEqual(result.lifetime, result.decayed_admission); + try std.testing.expectEqual(@as(u64, 1), result.belady_unique_misses_lower_bound); + const zero = try evaluate(std.testing.allocator, &rounds, .{ .capacity = 0 }); + try std.testing.expectEqual(@as(u64, 4), zero.lifetime.mapped_picks); + try std.testing.expectEqual(@as(u64, 0), zero.lifetime.copies); + try std.testing.expectEqual(@as(u64, 4), zero.decayed_admission.mapped_picks); + try std.testing.expectEqual(@as(u64, 0), zero.decayed_admission.copies); + try std.testing.expectEqual(@as(u64, 2), zero.belady_unique_misses_lower_bound); + const none = try evaluate(std.testing.allocator, &.{}, .{ .capacity = 2 }); + try std.testing.expectEqual(Metrics{}, none.lifetime); + try std.testing.expectEqual(@as(u64, 0), none.belady_unique_misses_lower_bound); +} + +test "exl3: Belady bypass retains a warm expert and bounds pinned round misses only" { + var first = try Synthetic.init(&.{ 1, 1, 2 }); + defer first.deinit(); + var second = try Synthetic.init(&.{ 0, 0 }); + defer second.deinit(); + const options: Options = .{ + .capacity = 1, + .seed = .{ .line_ids = &.{0}, .stamps = &.{0}, .counts = @splat(0), .now = 0 }, + }; + const result = try evaluate(std.testing.allocator, &.{ first.round(), second.round() }, options); + try std.testing.expectEqual(@as(u64, 2), result.belady_unique_misses_lower_bound); + try std.testing.expectEqual(@as(u64, 3), result.lifetime.miss_experts); + try std.testing.expect(result.belady_unique_misses_lower_bound <= result.lru.miss_experts); + try std.testing.expect(result.belady_unique_misses_lower_bound <= result.decayed_frequency.miss_experts); +} + +fn exhaustiveMisses(events: []const Event, position: usize, mask: u8, capacity: usize) u64 { + if (position == events.len) return 0; + const bit = @as(u8, 1) << @as(u3, @intCast(events[position].expert)); + if (mask & bit != 0) return exhaustiveMisses(events, position + 1, mask, capacity); + var best = exhaustiveMisses(events, position + 1, mask, capacity); + if (@popCount(mask) < capacity) { + best = @min(best, exhaustiveMisses(events, position + 1, mask | bit, capacity)); + } else { + for (0..3) |expert| { + const victim_bit = @as(u8, 1) << @as(u3, @intCast(expert)); + if (mask & victim_bit != 0) best = @min(best, exhaustiveMisses(events, position + 1, (mask & ~victim_bit) | bit, capacity)); + } + } + return 1 + best; +} + +test "exl3: Belady equals exhaustive optional-admission optimum for short cold and warm streams" { + for (0..243) |sequence| { + var encoded = sequence; + var events: [5]Event = undefined; + for (&events) |*event| { + event.* = .{ .expert = encoded % 3 }; + encoded /= 3; + } + for (0..3) |capacity| { + try std.testing.expectEqual(exhaustiveMisses(&events, 0, 0, capacity), belady(&events, capacity, &.{})); + if (capacity > 0) try std.testing.expectEqual(exhaustiveMisses(&events, 0, 1, capacity), belady(&events, capacity, &.{0})); + } + } +} + +test "exl3: invalid input is rejected before state changes including planner unique overflow" { + var state = try State.init(std.testing.allocator, .lifetime, .{ .capacity = 1 }); + defer state.deinit(); + var fixture = try Synthetic.init(&.{1}); + defer fixture.deinit(); + var truncated = fixture.round(); + truncated.words = truncated.words[0..3]; + try std.testing.expectError(error.InvalidRound, state.observe(truncated)); + fixture.words[1] = 2; + try std.testing.expectError(error.InvalidRound, state.observe(fixture.round())); + var picks: [257]i32 = undefined; + for (&picks, 0..) |*expert, index| expert.* = @intCast(index); + var overflow = try Synthetic.init(&picks); + defer overflow.deinit(); + try std.testing.expectError(error.PlannerOverflow, state.observe(overflow.round())); + try std.testing.expectEqual(Metrics{}, state.total); + try std.testing.expectEqualSlices(i32, &.{-1}, state.line_ids); + try std.testing.expectError(error.InvalidOptions, State.init(std.testing.allocator, .lru, .{ .capacity = 513 })); + try std.testing.expectError(error.InvalidOptions, State.init(std.testing.allocator, .lru, .{ .capacity = 1, .decay_interval_rounds = 0 })); + try std.testing.expectError(error.InvalidSeed, State.init(std.testing.allocator, .lru, .{ + .capacity = 2, + .seed = .{ .line_ids = &.{ 1, 1 }, .stamps = &.{ 0, 0 }, .counts = @splat(0), .now = 0 }, + })); +} diff --git a/zig/src/families/flashnext/exl3/dplan/exl3_dplan_replay.zig b/zig/src/families/flashnext/exl3/dplan/exl3_dplan_replay.zig new file mode 100644 index 0000000000..e3855df0b0 --- /dev/null +++ b/zig/src/families/flashnext/exl3/dplan/exl3_dplan_replay.zig @@ -0,0 +1,392 @@ +const std = @import("std"); +const trace = @import("exl3_dplan_trace.zig"); +const layout = @import("exl3_dplan_layout.zig"); + +pub const expert_count: usize = 512; +pub const max_unique: usize = 256; + +pub const Entry = struct { expert: u32, appearances: u32 }; +pub const Round = struct { + rows: u32, + valid_picks: u32, + entries: []const Entry, + planner_overflow: bool = false, +}; + +pub const Seed = struct { + line_ids: []const i32, + stamps: []const u32, + counts: [expert_count]u32, + now: u32, +}; + +pub const Delta = struct { + hit_picks: u64 = 0, + miss_picks: u64 = 0, + copies: u64 = 0, + evictions: u64 = 0, + inplace: u64 = 0, + mapped_picks: u64 = 0, + overflow: u64 = 0, +}; + +pub const Verification = struct { rounds: usize, matched_transitions: usize, delta: Delta }; + +pub const State = struct { + allocator: std.mem.Allocator, + line_ids: []i32, + stamps: []u32, + counts: [expert_count]u32, + now: u32, + + pub fn init(allocator: std.mem.Allocator, seed: Seed) !State { + if (seed.line_ids.len == 0 or seed.line_ids.len != seed.stamps.len) return error.InvalidSeed; + var resident: [expert_count]bool = @splat(false); + for (seed.line_ids) |id| { + if (id < -1 or id >= @as(i32, expert_count)) return error.InvalidSeed; + if (id >= 0) { + const index: usize = @intCast(id); + if (resident[index]) return error.InvalidSeed; + resident[index] = true; + } + } + const line_ids = try allocator.dupe(i32, seed.line_ids); + errdefer allocator.free(line_ids); + const stamps = try allocator.dupe(u32, seed.stamps); + return .{ .allocator = allocator, .line_ids = line_ids, .stamps = stamps, .counts = seed.counts, .now = seed.now }; + } + + pub fn deinit(self: *State) void { + self.allocator.free(self.line_ids); + self.allocator.free(self.stamps); + self.* = undefined; + } + + pub fn replay(self: *State, round: Round) !Delta { + if (round.planner_overflow or round.rows == 0 or round.rows > std.math.maxInt(u32) / 11 or round.valid_picks > round.rows * 11 or round.entries.len > max_unique) return error.InvalidRound; + var seen: [expert_count]bool = @splat(false); + var appearances: u64 = 0; + for (round.entries) |entry| { + if (entry.expert >= @as(u32, expert_count) or entry.appearances == 0 or seen[@intCast(entry.expert)]) return error.InvalidRound; + seen[@intCast(entry.expert)] = true; + appearances += entry.appearances; + } + if (appearances != round.valid_picks) return error.InvalidRound; + + self.now +%= 1; + var delta: Delta = .{}; + for (round.entries) |entry| { + const expert: usize = @intCast(entry.expert); + self.counts[expert] = self.counts[expert] +| entry.appearances; + if (self.lineOf(entry.expert) != null) { + delta.hit_picks += entry.appearances; + } else { + delta.miss_picks += entry.appearances; + } + } + + for (round.entries) |entry| if (self.lineOf(entry.expert)) |line| { + self.stamps[line] = self.now; + }; + + var misses: [max_unique]Entry = undefined; + var miss_count: usize = 0; + for (round.entries) |entry| if (self.lineOf(entry.expert) == null) { + misses[miss_count] = entry; + miss_count += 1; + }; + var index: usize = 1; + while (index < miss_count) : (index += 1) { + const item = misses[index]; + var insert = index; + while (insert > 0 and hotter(self, item, misses[insert - 1])) : (insert -= 1) { + misses[insert] = misses[insert - 1]; + } + misses[insert] = item; + } + + for (misses[0..miss_count]) |entry| { + if (self.victim(&seen)) |line| { + if (self.line_ids[line] >= 0) delta.evictions += 1; + self.line_ids[line] = @intCast(entry.expert); + self.stamps[line] = self.now; + delta.copies += 1; + } else { + delta.inplace += 1; + delta.mapped_picks += entry.appearances; + } + } + return delta; + } + + fn lineOf(self: *const State, expert: u32) ?usize { + const target: i32 = @intCast(expert); + for (self.line_ids, 0..) |id, line| if (id == target) return line; + return null; + } + + fn victim(self: *const State, keep: *const [expert_count]bool) ?usize { + var empty: ?usize = null; + var coldest: ?usize = null; + for (self.line_ids, 0..) |id, line| { + if (id < 0) { + if (empty == null) empty = line; + continue; + } + if (keep[@intCast(id)]) continue; + if (coldest == null or colder(self, line, coldest.?)) coldest = line; + } + return empty orelse coldest; + } +}; + +pub const ExactState = struct { + allocator: std.mem.Allocator, + lines: usize, + words: []u32, + + pub fn init(allocator: std.mem.Allocator, lines: usize, seed: []const u32) !ExactState { + const entry = layout.PoolState.init(lines, 0) catch return error.InvalidSeed; + if (seed.len != entry.words) return error.InvalidSeed; + var reverse: [expert_count]u32 = @splat(std.math.maxInt(u32)); + for (seed[0..lines], 0..) |raw, line| { + const id: i32 = @bitCast(raw); + if (id < -1 or id >= expert_count) return error.InvalidSeed; + if (id >= 0) { + const expert: usize = @intCast(id); + if (reverse[expert] != std.math.maxInt(u32)) return error.InvalidSeed; + reverse[expert] = @intCast(line); + } + } + if (!std.mem.eql(u32, &reverse, seed[2 * lines ..][0..expert_count])) return error.InvalidSeed; + const misc = seed[layout.miscOffset(lines)..]; + if (misc[8] >= layout.recent_window or misc[9] > layout.recent_window) return error.InvalidSeed; + var recent_counts: [expert_count]u32 = @splat(0); + for (seed[layout.recentOffset(lines)..layout.miscOffset(lines)], 0..) |bits, index| { + for (0..32) |bit| { + if (bits & (@as(u32, 1) << @intCast(bit)) != 0) recent_counts[(index % 16) * 32 + bit] += 1; + } + } + if (!std.mem.eql(u32, &recent_counts, seed[layout.recentCountsOffset(lines)..][0..expert_count])) return error.InvalidSeed; + for (recent_counts) |count| if (count > misc[9]) return error.InvalidSeed; + return .{ .allocator = allocator, .lines = lines, .words = try allocator.dupe(u32, seed) }; + } + + pub fn deinit(self: *ExactState) void { + self.allocator.free(self.words); + self.* = undefined; + } + + pub fn replay(self: *ExactState, record: trace.Round) !Delta { + const lines = self.lines; + const ids = self.words[0..lines]; + const stamps = self.words[lines .. 2 * lines]; + const reverse = self.words[2 * lines ..][0..expert_count]; + const counts = self.words[2 * lines + expert_count ..][0..expert_count]; + const recent = self.words[layout.recentCountsOffset(lines)..][0..expert_count]; + const history = self.words[layout.recentOffset(lines)..layout.miscOffset(lines)]; + const misc = self.words[layout.miscOffset(lines)..]; + misc[0] +%= 1; + misc[1] +%= 1; + if (record.recentEnabled()) { + for (0..record.rows()) |row| { + const expired = history[misc[8] * 16 ..][0..16]; + for (expired, 0..) |bits, word| { + for (0..32) |bit| { + if (bits & (@as(u32, 1) << @intCast(bit)) != 0) recent[word * 32 + bit] -%= 1; + } + } + @memset(expired, 0); + for (record.picks()[row * 11 ..][0..11]) |raw| { + const expert: i32 = @bitCast(raw); + if (expert < 0 or expert >= expert_count) continue; + const index: usize = @intCast(expert); + const bit = @as(u32, 1) << @intCast(index & 31); + if (expired[index >> 5] & bit == 0) { + expired[index >> 5] |= bit; + recent[index] += 1; + } + } + misc[8] = (misc[8] + 1) % @as(u32, layout.recent_window); + misc[9] = @min(misc[9] + 1, @as(u32, layout.recent_window)); + } + } + var delta: Delta = .{}; + var ordered: [max_unique]usize = undefined; + var appearances: [expert_count]u16 = @splat(0); + var keep: [expert_count]bool = @splat(false); + var unique: usize = 0; + var overflow = false; + for (record.picks()) |raw| { + const expert: i32 = @bitCast(raw); + if (expert < 0 or expert >= expert_count) continue; + const index: usize = @intCast(expert); + counts[index] +|= 1; + if (reverse[index] != std.math.maxInt(u32)) delta.hit_picks += 1 else delta.miss_picks += 1; + if (!keep[index]) { + if (unique == max_unique) { + overflow = true; + continue; + } + keep[index] = true; + ordered[unique] = index; + unique += 1; + } + appearances[index] +%= 1; + } + var misses: [max_unique]usize = undefined; + var missing: usize = 0; + for (ordered[0..unique]) |expert| { + if (reverse[expert] != std.math.maxInt(u32)) { + stamps[reverse[expert]] = misc[0]; + } else { + misses[missing] = expert; + missing += 1; + } + } + const scores = if (record.recentEnabled()) recent else counts; + var cursor: usize = 1; + while (cursor < missing) : (cursor += 1) { + const item = misses[cursor]; + var insert = cursor; + while (insert > 0) : (insert -= 1) { + const previous = misses[insert - 1]; + const hotter_item = if (scores[item] != scores[previous]) scores[item] > scores[previous] else if (appearances[item] != appearances[previous]) appearances[item] > appearances[previous] else item < previous; + if (!hotter_item) break; + misses[insert] = previous; + } + misses[insert] = item; + } + if (overflow) { + delta.overflow = 1; + } else for (misses[0..missing]) |expert| { + var victim: ?usize = null; + for (ids, 0..) |raw, line| { + if (raw == std.math.maxInt(u32)) { + victim = line; + break; + } + if (keep[raw]) continue; + if (victim) |previous| { + const old = ids[previous]; + if (scores[raw] > scores[old] or (scores[raw] == scores[old] and (stamps[line] > stamps[previous] or (stamps[line] == stamps[previous] and line < previous)))) continue; + } + victim = line; + } + if (victim) |line| { + if (ids[line] != std.math.maxInt(u32)) { + reverse[ids[line]] = std.math.maxInt(u32); + delta.evictions += 1; + } + ids[line] = @intCast(expert); + reverse[expert] = @intCast(line); + stamps[line] = misc[0]; + delta.copies += 1; + } else { + delta.inplace += 1; + for (record.picks()) |raw| { + if (raw == expert) delta.mapped_picks += 1; + } + } + } + inline for (.{ .{ 2, "hit_picks" }, .{ 3, "miss_picks" }, .{ 4, "copies" }, .{ 5, "evictions" }, .{ 6, "overflow" }, .{ 7, "inplace" } }) |field| { + misc[field[0]] +%= @truncate(@field(delta, field[1])); + } + return delta; + } +}; + +pub fn verifyTrace(allocator: std.mem.Allocator, view: trace.View) !Verification { + return verifyTraceWithFinal(allocator, view, null); +} + +pub fn verifyTraceWithFinal(allocator: std.mem.Allocator, view: trace.View, final_words: ?[]const u32) !Verification { + if (view.dropped() != 0 or view.droppedCheckpoints() != 0) return error.TraceDroppedRecords; + if (view.words[12] != 0) return error.TraceDiscontinuous; + if (view.words[13] != 0) return error.TraceInvalidLaunch; + if (view.recorded() == 0) return error.EmptyTrace; + var state = try ExactState.init(allocator, view.words[4], view.initialSeed()); + defer state.deinit(); + var total: Delta = .{}; + var checkpoint_index: usize = 0; + for (0..view.recorded() + 1) |index| { + while (checkpoint_index < view.checkpoints()) { + const checkpoint = try view.checkpoint(checkpoint_index); + if (checkpoint.roundIndex() != index) break; + if (!std.mem.eql(u32, state.words, checkpoint.before())) return error.PlannerCheckpointMismatch; + var incoming = try ExactState.init(allocator, state.lines, checkpoint.after()); + std.mem.swap(ExactState, &state, &incoming); + incoming.deinit(); + checkpoint_index += 1; + } + if (index == view.recorded()) break; + const record = try view.round(index); + const delta = try state.replay(record); + if (!std.mem.eql(u32, state.words[layout.miscOffset(state.lines)..], record.misc()) or delta.mapped_picks != record.mappedPicks()) return error.PlannerCounterMismatch; + try add(&total, delta); + } + if (!std.mem.eql(u32, state.words, view.terminalSeed())) return error.PlannerFinalStateMismatch; + if (final_words) |expected| { + if (!std.mem.eql(u32, state.words, expected)) return error.PlannerFinalStateMismatch; + } + if (total.overflow != 0) return error.PlannerOverflow; + return .{ .rounds = view.recorded(), .matched_transitions = view.recorded(), .delta = total }; +} + +fn add(total: *Delta, delta: Delta) !void { + inline for (.{ "hit_picks", "miss_picks", "copies", "evictions", "inplace", "mapped_picks", "overflow" }) |field| { + @field(total, field) = try std.math.add(u64, @field(total, field), @field(delta, field)); + } +} + +fn hotter(state: *const State, a: Entry, b: Entry) bool { + const a_count = state.counts[@intCast(a.expert)]; + const b_count = state.counts[@intCast(b.expert)]; + if (a_count != b_count) return a_count > b_count; + if (a.appearances != b.appearances) return a.appearances > b.appearances; + if (a.expert != b.expert) return a.expert < b.expert; + return false; +} + +fn colder(state: *const State, a: usize, b: usize) bool { + const a_id: usize = @intCast(state.line_ids[a]); + const b_id: usize = @intCast(state.line_ids[b]); + if (state.counts[a_id] != state.counts[b_id]) return state.counts[a_id] < state.counts[b_id]; + if (state.stamps[a] != state.stamps[b]) return state.stamps[a] < state.stamps[b]; + return a > b; +} + +test "planner replay reproduces current lifetime-count admission and eviction" { + var counts: [expert_count]u32 = @splat(0); + counts[0] = 5; + counts[1] = 3; + var state = try State.init(std.testing.allocator, .{ .line_ids = &.{ 0, 1 }, .stamps = &.{ 1, 2 }, .counts = counts, .now = 2 }); + defer state.deinit(); + + const delta = try state.replay(.{ .rows = 1, .valid_picks = 6, .entries = &.{ .{ .expert = 2, .appearances = 5 }, .{ .expert = 3, .appearances = 1 } } }); + try std.testing.expectEqual(Delta{ .miss_picks = 6, .copies = 2, .evictions = 2 }, delta); + try std.testing.expectEqualSlices(i32, &.{ 3, 2 }, state.line_ids); + try std.testing.expectEqualSlices(u32, &.{ 3, 3 }, state.stamps); + try std.testing.expectEqual(@as(u32, 5), state.counts[2]); +} + +test "planner replay counts repeated slot misses and preserves held lines under overflow pressure" { + var counts: [expert_count]u32 = @splat(0); + counts[0] = 9; + var state = try State.init(std.testing.allocator, .{ .line_ids = &.{0}, .stamps = &.{5}, .counts = counts, .now = 5 }); + defer state.deinit(); + + const delta = try state.replay(.{ .rows = 1, .valid_picks = 4, .entries = &.{ .{ .expert = 0, .appearances = 3 }, .{ .expert = 1, .appearances = 1 } } }); + try std.testing.expectEqual(Delta{ .hit_picks = 3, .miss_picks = 1, .inplace = 1, .mapped_picks = 1 }, delta); + try std.testing.expectEqualSlices(i32, &.{0}, state.line_ids); + try std.testing.expectEqual(@as(u32, 6), state.stamps[0]); +} + +test "planner replay consumes empty lines first and rejects incomplete evidence" { + var state = try State.init(std.testing.allocator, .{ .line_ids = &.{ -1, -1 }, .stamps = &.{ 90, 91 }, .counts = @splat(0), .now = 0 }); + defer state.deinit(); + const delta = try state.replay(.{ .rows = 1, .valid_picks = 2, .entries = &.{ .{ .expert = 10, .appearances = 1 }, .{ .expert = 11, .appearances = 1 } } }); + try std.testing.expectEqual(Delta{ .miss_picks = 2, .copies = 2 }, delta); + try std.testing.expectEqualSlices(i32, &.{ 10, 11 }, state.line_ids); + try std.testing.expectError(error.InvalidRound, state.replay(.{ .rows = 1, .valid_picks = 1, .entries = &.{.{ .expert = 12, .appearances = 1 }}, .planner_overflow = true })); +} diff --git a/zig/src/families/flashnext/exl3/dplan/exl3_dplan_trace.zig b/zig/src/families/flashnext/exl3/dplan/exl3_dplan_trace.zig new file mode 100644 index 0000000000..c73f48e250 --- /dev/null +++ b/zig/src/families/flashnext/exl3/dplan/exl3_dplan_trace.zig @@ -0,0 +1,189 @@ +//! DPLAN v4: default off, little-endian u32 words, fixed 1 MiB per device pool. + +const std = @import("std"); +const layout = @import("exl3_dplan_layout.zig"); + +pub const magic: u32 = 0x44505431; +pub const version: u32 = 4; +pub const pool_bytes: usize = 1024 * 1024; +pub const pool_words: usize = pool_bytes / 4; +pub const header_words: usize = 16; +pub const record_header_words: usize = 4 + layout.misc_words + 1; +pub const checkpoint_header_words: usize = 4; + +pub fn checkpointWords(state_words: usize) usize { + return checkpoint_header_words + 2 * state_words; +} + +pub fn seedFromState(comptime Seed: type, lines: usize, seed: []const u32) Seed { + return .{ .line_ids = @ptrCast(seed[0..lines]), .stamps = seed[lines .. 2 * lines], .counts = seed[2 * lines + layout.experts_n ..][0..layout.experts_n].*, .now = seed[layout.miscOffset(lines)] }; +} + +pub const Checkpoint = struct { + words: []const u32, + + pub fn roundIndex(self: Checkpoint) usize { + return self.words[0]; + } + pub fn before(self: Checkpoint) []const u32 { + return self.words[checkpoint_header_words..][0..self.words[2]]; + } + pub fn after(self: Checkpoint) []const u32 { + return self.words[checkpoint_header_words + self.words[2] ..][0..self.words[2]]; + } +}; + +/// Set by a test before the planner is made: the planner then keeps a bounded demand trace per pool. Not a user switch. +pub var requested: bool = false; + +pub fn enabledValue(value: ?[]const u8) bool { + return if (value) |value_text| std.mem.eql(u8, value_text, "1") else false; +} + +pub fn header(pool_index: usize, layer: usize, lines: usize, max_rows: usize) ![header_words]u32 { + if (max_rows == 0 or max_rows > pool_words / 11) return error.TraceShapeTooLarge; + const entry = layout.PoolState.init(lines, 0) catch return error.TraceShapeTooLarge; + if (entry.words > (pool_words - header_words) / 2) return error.TraceShapeTooLarge; + const stride = record_header_words + max_rows * 11; + const available = pool_words - header_words - 2 * entry.words; + if (stride > available or pool_index > std.math.maxInt(u32) or layer > std.math.maxInt(u32)) return error.TraceShapeTooLarge; + return .{ magic, version, @intCast(pool_index), @intCast(layer), @intCast(lines), @intCast(max_rows), @intCast(entry.words), @intCast(stride), @intCast(available / stride), 0, 0, 0, 0, 0, 0, 0 }; +} + +pub const Round = struct { + words: []const u32, + lines: usize, + max_rows: usize, + + pub fn rows(self: Round) usize { + return self.words[0]; + } + pub fn validPicks(self: Round) u32 { + return self.words[1]; + } + pub fn uniqueCount(self: Round) usize { + return self.words[2]; + } + pub fn recentEnabled(self: Round) bool { + return self.words[3] == 1; + } + pub fn misc(self: Round) []const u32 { + return self.words[4..][0..layout.misc_words]; + } + pub fn mappedPicks(self: Round) u32 { + return self.words[4 + layout.misc_words]; + } + pub fn picks(self: Round) []const u32 { + return self.words[record_header_words..][0 .. self.rows() * 11]; + } + + fn validate(self: Round) !void { + if (self.rows() == 0 or self.rows() > self.max_rows or self.uniqueCount() > layout.experts_n or self.words[3] > 1) return error.InvalidTrace; + var seen: [layout.experts_n]bool = @splat(false); + var unique: usize = 0; + var valid: u32 = 0; + for (self.picks()) |bits| { + const expert: i32 = @bitCast(bits); + if (expert < 0 or expert >= layout.experts_n) continue; + const index: usize = @intCast(expert); + if (!seen[index]) unique += 1; + seen[index] = true; + valid += 1; + } + if (unique != self.uniqueCount() or valid != self.validPicks() or self.mappedPicks() > valid) return error.InvalidTrace; + if (self.misc()[8] >= layout.recent_window or self.misc()[9] > layout.recent_window) return error.InvalidTrace; + } +}; + +pub const View = struct { + words: []const u32, + + pub fn init(words: []const u32) !View { + if (words.len < header_words or words.len > pool_words or words[0] != magic or words[1] != version) return error.InvalidTrace; + const expected = header(words[2], words[3], words[4], words[5]) catch return error.InvalidTrace; + if (!std.mem.eql(u32, words[0..9], expected[0..9]) or words[9] > words[8]) return error.InvalidTrace; + if (words[12] > 1 or words[13] > 1) return error.InvalidTrace; + const self: View = .{ .words = words }; + const stride = checkpointWords(words[6]); + if (self.checkpoints() > (pool_words - self.recordsOffset()) / stride) return error.InvalidTrace; + const records_end = self.recordsOffset() + self.recorded() * words[7]; + const checkpoints_start = pool_words - self.checkpoints() * stride; + if (records_end > checkpoints_start) return error.InvalidTrace; + if (words.len < self.usedWords()) return error.InvalidTrace; + if (self.dropped() != 0 and self.droppedCheckpoints() == 0 and checkpoints_start - records_end >= words[7]) return error.InvalidTrace; + if (self.droppedCheckpoints() != 0 and self.dropped() == 0 and checkpoints_start - records_end >= stride) return error.InvalidTrace; + var previous_round: usize = 0; + for (0..self.checkpoints()) |index| { + const entry = try self.checkpoint(index); + if (entry.words[1] != 1 or entry.words[2] != words[6] or entry.words[3] != index or entry.roundIndex() > self.recorded() or entry.roundIndex() < previous_round) return error.InvalidTrace; + previous_round = entry.roundIndex(); + } + for (0..self.recorded()) |index| try (try self.round(index)).validate(); + return self; + } + + pub fn recorded(self: View) usize { + return self.words[9]; + } + pub fn dropped(self: View) u64 { + return @as(u64, self.words[10]) | (@as(u64, self.words[11]) << 32); + } + pub fn checkpoints(self: View) usize { + return self.words[14]; + } + pub fn droppedCheckpoints(self: View) u32 { + return self.words[15]; + } + pub fn checkpoint(self: View, index: usize) !Checkpoint { + if (index >= self.checkpoints()) return error.TraceCheckpointOutOfRange; + const stride = checkpointWords(self.words[6]); + return .{ .words = self.words[pool_words - (index + 1) * stride ..][0..stride] }; + } + pub fn initialSeed(self: View) []const u32 { + return self.words[header_words..][0..self.words[6]]; + } + pub fn terminalSeed(self: View) []const u32 { + return self.words[header_words + self.words[6] ..][0..self.words[6]]; + } + pub fn seedState(self: View, comptime Seed: type) Seed { + return seedFromState(Seed, self.words[4], self.initialSeed()); + } + pub fn recordsOffset(self: View) usize { + return header_words + 2 * @as(usize, self.words[6]); + } + pub fn round(self: View, index: usize) !Round { + if (index >= self.recorded()) return error.TraceRoundOutOfRange; + const stride = self.words[7]; + return .{ .words = self.words[self.recordsOffset() + index * stride ..][0..stride], .lines = self.words[4], .max_rows = self.words[5] }; + } + fn usedWords(self: View) usize { + if (self.checkpoints() != 0) return pool_words; + return self.recordsOffset() + self.recorded() * self.words[7]; + } + pub fn encode(self: View, allocator: std.mem.Allocator) ![]u8 { + const bytes = try allocator.alloc(u8, self.usedWords() * 4); + for (self.words[0..self.usedWords()], 0..) |word, index| std.mem.writeInt(u32, bytes[index * 4 ..][0..4], word, .little); + return bytes; + } +}; + +pub const Snapshot = struct { + allocator: std.mem.Allocator, + words: []u32, + + pub fn deinit(self: *Snapshot) void { + self.allocator.free(self.words); + self.* = undefined; + } + pub fn view(self: Snapshot) !View { + return View.init(self.words); + } + pub fn decode(allocator: std.mem.Allocator, bytes: []const u8) !Snapshot { + if (bytes.len % 4 != 0 or bytes.len > pool_bytes or bytes.len < header_words * 4) return error.InvalidTrace; + const words = try allocator.alloc(u32, bytes.len / 4); + errdefer allocator.free(words); + for (words, 0..) |*word, index| word.* = std.mem.readInt(u32, bytes[index * 4 ..][0..4], .little); + _ = try View.init(words); + return .{ .allocator = allocator, .words = words }; + } +}; diff --git a/zig/src/families/flashnext/exl3/dplan/exl3_dplan_trace_kernel_test.cpp b/zig/src/families/flashnext/exl3/dplan/exl3_dplan_trace_kernel_test.cpp new file mode 100644 index 0000000000..25689fd978 --- /dev/null +++ b/zig/src/families/flashnext/exl3/dplan/exl3_dplan_trace_kernel_test.cpp @@ -0,0 +1,261 @@ +#include +#include +#include +#include +#include +#include + +#define __device__ +#define __global__ +#define __shared__ +#define __forceinline__ inline +#define __launch_bounds__(value) +struct TestDim { int x; }; +static TestDim threadIdx{0}, blockIdx{0}, blockDim{1}, gridDim{1}; +struct uint4 { unsigned x, y, z, w; }; +using cudaGraphConditionalHandle = unsigned long long; +static void __syncthreads() {} +static int __ffs(unsigned value) { return __builtin_ffs(value); } +static int __float_as_int(float value) { + int bits; + std::memcpy(&bits, &value, sizeof(bits)); + return bits; +} +static void cudaGraphSetConditional(cudaGraphConditionalHandle, unsigned) {} + +#include "../../../../../kernels/cuda/flashnext_exl3/flashnext_exl3_dplan.cu" + +static constexpr unsigned pool_words = 1024 * 1024 / 4; +static constexpr unsigned sentinel = 0xfeedface; + +static std::vector make_trace(unsigned lines, unsigned rows) { + std::vector trace(pool_words + 1, 0); + const unsigned seed_words = 2 * lines + 3 * DP_EXPERTS + DP_RECENT_WORDS + 10; + const unsigned stride = 15 + rows * 11; + auto* header = reinterpret_cast(trace.data()); + *header = {0x44505431, 4, 0, 0, lines, rows, seed_words, stride, + (pool_words - 16 - 2 * seed_words) / stride, 0, 0, 0, 0, 0, 0, 0}; + trace[pool_words] = sentinel; + return trace; +} + +static std::vector make_state(unsigned lines) { + std::vector state(2 * lines + 3 * DP_EXPERTS + DP_RECENT_WORDS + 10, 0); + for (unsigned line = 0; line < lines; ++line) state[line] = 0xffffffffu; + for (unsigned expert = 0; expert < DP_EXPERTS; ++expert) state[2 * lines + expert] = 0xffffffffu; + return state; +} + +struct PlanOutput { + long long tabs[3][513]{}; + int widths[3][513]{}; + long long work[512 * 9]{}; + int work_n = 0, overflow[512]{}, overflow_n = 0; +}; + +static void run_plan(std::vector& state, unsigned lines, const std::vector& picks, + unsigned* trace, int recent, PlanOutput& output) { + long long src[512 * 3], lens[512 * 3]; + int widths[512 * 3]; + for (int index = 0; index < 512 * 3; ++index) { + src[index] = 0x100000 + index * 256; + lens[index] = 256; + widths[index] = 3; + } + fn_exl3_dplan(picks.data(), static_cast(picks.size()), state.data(), lines, recent, + output.tabs[0], output.tabs[1], output.tabs[2], + output.widths[0], output.widths[1], output.widths[2], + src, lens, widths, 1024, 1024, 0, 256, 512, + output.work, &output.work_n, nullptr, nullptr, output.overflow, &output.overflow_n, trace); +} + +static void assert_same(const PlanOutput& first, const PlanOutput& second) { + assert(std::memcmp(first.tabs, second.tabs, sizeof(first.tabs)) == 0); + assert(std::memcmp(first.widths, second.widths, sizeof(first.widths)) == 0); + assert(std::memcmp(first.work, second.work, sizeof(first.work)) == 0); + assert(std::memcmp(first.overflow, second.overflow, sizeof(first.overflow)) == 0); + assert(first.work_n == second.work_n && first.overflow_n == second.overflow_n); +} + +static void test_bounds_and_flags() { + auto trace = make_trace(2, 47); + auto* header = reinterpret_cast(trace.data()); + auto state = make_state(2); + std::vector picks(47 * 11, 512); + for (int index = 0; index < 512; ++index) picks[index] = 511 - index; + picks[512] = -1; + picks[513] = 511; + PlanOutput output; + const auto seed = state; + run_plan(state, 2, picks, trace.data(), 1, output); + const unsigned offset = 16 + 2 * header->seed_words; + const unsigned* round = trace.data() + offset; + assert(round[0] == 47 && round[1] == 513 && round[2] == 512 && round[3] == 1); + assert(round[10] == 1 && round[8] == 0 && round[14] == 0); + assert(std::memcmp(trace.data() + 16, seed.data(), seed.size() * 4) == 0); + assert(std::memcmp(round + 15, picks.data(), picks.size() * 4) == 0); + assert(std::memcmp(trace.data() + 16 + seed.size(), state.data(), state.size() * 4) == 0); + const auto first = std::vector(round, round + header->record_words); + const unsigned capacity = header->capacity; + for (unsigned index = 1; index < capacity + 3; ++index) + run_plan(state, 2, picks, trace.data(), 1, output); + assert(header->recorded == capacity && header->dropped_lo == 3 && header->dropped_hi == 0); + assert(std::memcmp(trace.data() + offset, first.data(), first.size() * 4) == 0); + assert(trace[pool_words] == sentinel); + const auto terminal = std::vector(trace.data() + 16 + seed.size(), trace.data() + offset); + header->dropped_lo = 0xffffffffu; + header->dropped_hi = 4; + run_plan(state, 2, picks, trace.data(), 0, output); + assert(header->dropped_lo == 0 && header->dropped_hi == 5); + header->dropped_lo = header->dropped_hi = 0xffffffffu; + run_plan(state, 2, picks, trace.data(), 0, output); + assert(header->dropped_lo == 0xffffffffu && header->dropped_hi == 0xffffffffu); + assert(std::memcmp(trace.data() + 16 + seed.size(), terminal.data(), terminal.size() * 4) == 0); + assert(dp_trace_begin(nullptr, 0, nullptr, nullptr, 0, 2) == nullptr); + auto invalid = make_trace(2, 1); + auto* invalid_header = reinterpret_cast(invalid.data()); + assert(dp_trace_begin(picks.data(), 12, invalid.data(), state.data(), 0, 2) == nullptr); + assert(invalid_header->invalid_launch == 1 && invalid_header->recorded == 0); + auto broken = make_trace(2, 47); + run_plan(state, 2, picks, broken.data(), 0, output); + state.back() ^= 1; + run_plan(state, 2, picks, broken.data(), 0, output); + assert(reinterpret_cast(broken.data())->discontinuity == 1); + + for (unsigned field : {0u, 1u, 4u, 5u, 6u, 7u, 8u, 9u, 12u, 13u, 14u}) { + auto corrupt = make_trace(2, 1); + corrupt[field] = 0xffffffffu; + assert(dp_trace_begin(nullptr, 11, corrupt.data(), nullptr, 0, 2) == nullptr); + assert(corrupt[13] == 1 && corrupt[pool_words] == sentinel); + corrupt[field] = 0xffffffffu; + fn_exl3_dcheckpoint(nullptr, 2, corrupt.data(), 1); + assert(corrupt[pool_words] == sentinel); + } + auto full = make_trace(2, 16); + auto* full_header = reinterpret_cast(full.data()); + const unsigned checkpoint_words = 4 + 2 * full_header->seed_words; + const unsigned checkpoint_capacity = (pool_words - 16 - 2 * full_header->seed_words) / checkpoint_words; + for (unsigned index = 0; index < checkpoint_capacity + 1; ++index) + fn_exl3_dcheckpoint(state.data(), 2, full.data(), 1); + assert(full_header->checkpoints == checkpoint_capacity && full_header->dropped_checkpoints == 1); + full_header->dropped_checkpoints = 0xffffffffu; + fn_exl3_dcheckpoint(state.data(), 2, full.data(), 1); + assert(full_header->dropped_checkpoints == 0xffffffffu && full[pool_words] == sentinel); + run_plan(state, 2, std::vector(11, -1), full.data(), 0, output); + assert(full_header->recorded == 0 && full_header->dropped_lo == 1); + + auto sealed = make_trace(2, 16); + fn_exl3_dcheckpoint(state.data(), 2, sealed.data(), 1); + state.back() ^= 1; + fn_exl3_dcheckpoint(state.data(), 2, sealed.data(), 0); + assert(reinterpret_cast(sealed.data())->discontinuity == 1); +} + +static std::vector test_checkpoint_capture() { + const unsigned lines = 10; + auto trace = make_trace(lines, 16); + auto* header = reinterpret_cast(trace.data()); + auto state = make_state(lines); + auto untraced = state; + PlanOutput captured, disabled; + fn_exl3_dcheckpoint(state.data(), lines, trace.data(), 1); + for (unsigned round = 0; round < 1024; ++round) { + if (round == 512 || round == 768) { + fn_exl3_dcheckpoint(state.data(), lines, trace.data(), 0); + for (unsigned host_round = 0; host_round < 3; ++host_round) { + std::vector host_picks(11, static_cast((round + host_round * 13) % 512)); + run_plan(state, lines, host_picks, nullptr, 0, captured); + run_plan(untraced, lines, host_picks, nullptr, 0, disabled); + } + const unsigned misc = 2 * lines + 3 * DP_EXPERTS + DP_RECENT_WORDS; + for (unsigned counter = 1; counter < 8; ++counter) + state[misc + counter] = untraced[misc + counter] = 0; + fn_exl3_dcheckpoint(state.data(), lines, trace.data(), 1); + } + std::vector picks(11, static_cast((round * 7) % 512)); + picks.back() = 512; + run_plan(state, lines, picks, trace.data(), round % 3 != 0, captured); + run_plan(untraced, lines, picks, nullptr, round % 3 != 0, disabled); + assert(state == untraced); + assert_same(captured, disabled); + } + fn_exl3_dcheckpoint(state.data(), lines, trace.data(), 0); + auto* misc = reinterpret_cast(state.data() + 2 * lines + 3 * DP_EXPERTS + DP_RECENT_WORDS); + misc->now += 100; + fn_exl3_dcheckpoint(state.data(), lines, trace.data(), 1); + assert(header->recorded == 1024 && header->checkpoints == 4); + assert(header->dropped_lo == 0 && header->dropped_hi == 0 && header->dropped_checkpoints == 0); + assert(header->discontinuity == 0 && header->invalid_launch == 0); + assert(std::memcmp(trace.data() + 16 + header->seed_words, state.data(), state.size() * 4) == 0); + assert(trace[pool_words] == sentinel); + return trace; +} + +static std::vector test_long_capture() { + const unsigned lines = 10; + auto trace = make_trace(lines, 16); + auto* header = reinterpret_cast(trace.data()); + assert(header->capacity >= 1200); + assert(16 + 2 * header->seed_words + header->capacity * header->record_words <= pool_words); + auto state = make_state(lines); + const unsigned misc_at = 2 * lines + 3 * DP_EXPERTS + DP_RECENT_WORDS; + for (unsigned index = 0; index < 8; ++index) state[misc_at + index] = 0xfffffffeu; + state[2 * lines + DP_EXPERTS + 7] = 0xffffffffu; + auto untraced = state; + const auto initial = state; + PlanOutput captured, disabled; + for (int index = 0; index < 1200; ++index) { + std::vector picks(16 * 11); + for (unsigned slot = 0; slot < picks.size(); ++slot) { + switch (slot % 11) { + case 0: picks[slot] = -1; break; + case 1: picks[slot] = 512; break; + case 2: picks[slot] = -99; break; + case 3: picks[slot] = 800; break; + default: picks[slot] = (index * 7 + slot * 3) % 512; break; + } + } + const int recent = index % 3 != 0; + run_plan(state, lines, picks, trace.data(), recent, captured); + run_plan(untraced, lines, picks, nullptr, recent, disabled); + assert(state == untraced); + assert_same(captured, disabled); + assert(header->recorded == static_cast(index + 1)); + const unsigned* round = trace.data() + 16 + 2 * header->seed_words + index * header->record_words; + assert(std::memcmp(round + 4, state.data() + misc_at, sizeof(DpMisc)) == 0); + unsigned mapped = 0; + for (int pick : picks) + for (int overflow = 0; overflow < captured.overflow_n; ++overflow) + if (pick == captured.overflow[overflow]) ++mapped; + assert(round[14] == mapped); + assert(std::memcmp(round + 15, picks.data(), picks.size() * 4) == 0); + } + assert(header->discontinuity == 0 && header->invalid_launch == 0); + assert(std::memcmp(trace.data() + 16, initial.data(), initial.size() * 4) == 0); + assert(std::memcmp(trace.data() + 16 + header->seed_words, state.data(), state.size() * 4) == 0); + assert(trace[pool_words] == sentinel); + return trace; +} + +static void write_trace(const std::vector& trace, const char* path) { + const auto* header = reinterpret_cast(trace.data()); + const unsigned used = header->checkpoints ? pool_words : 16 + 2 * header->seed_words + header->recorded * header->record_words; + FILE* output = std::fopen(path, "wb"); + assert(output != nullptr); + for (unsigned index = 0; index < used; ++index) { + const unsigned word = trace[index]; + const unsigned char bytes[4] = { + static_cast(word), static_cast(word >> 8), + static_cast(word >> 16), static_cast(word >> 24)}; + assert(std::fwrite(bytes, 1, 4, output) == 4); + } + assert(std::fclose(output) == 0); +} + +int main(int argc, char** argv) { + test_bounds_and_flags(); + const auto trace = test_long_capture(); + const auto checkpoint_trace = test_checkpoint_capture(); + if (argc >= 2) write_trace(trace, argv[1]); + if (argc >= 3) write_trace(checkpoint_trace, argv[2]); +} diff --git a/zig/src/families/flashnext/exl3/dplan/exl3_dplan_trace_replay_test.zig b/zig/src/families/flashnext/exl3/dplan/exl3_dplan_trace_replay_test.zig new file mode 100644 index 0000000000..4c5a41137d --- /dev/null +++ b/zig/src/families/flashnext/exl3/dplan/exl3_dplan_trace_replay_test.zig @@ -0,0 +1,294 @@ +const std = @import("std"); +const trace = @import("exl3_dplan_trace.zig"); +const layout = @import("exl3_dplan_layout.zig"); +const replay = @import("exl3_dplan_replay.zig"); +const policies = @import("exl3_dplan_policy_eval.zig"); +const fixture = @import("exl3_dplan_trace_test.zig"); + +fn capture(words: []u32, state: *replay.ExactState, picks: []const i32, recent: bool) !replay.Delta { + const before = try std.testing.allocator.dupe(u32, state.words); + defer std.testing.allocator.free(before); + const record_words = try std.testing.allocator.alloc(u32, words[7]); + defer std.testing.allocator.free(record_words); + @memset(record_words, 0); + record_words[0] = @intCast(picks.len / 11); + record_words[3] = @intFromBool(recent); + for (picks, 0..) |pick, index| record_words[trace.record_header_words + index] = @bitCast(pick); + const delta = try state.replay(.{ .words = record_words, .lines = state.lines, .max_rows = words[5] }); + fixture.append(words, before, state.words, picks, recent, @intCast(delta.mapped_picks)); + return delta; +} + +test "exl3: ownership checkpoints verify every segment including repeated and final transfers" { + const allocator = std.testing.allocator; + const words = try fixture.make(allocator, 2, 1); + defer allocator.free(words); + const seed = try fixture.emptySeed(allocator, 2); + defer allocator.free(seed); + var state = try replay.ExactState.init(allocator, 2, seed); + defer state.deinit(); + fixture.checkpoint(words, state.words); + _ = try capture(words, &state, &.{ 2, 2, 3, -1, -1, -1, -1, -1, -1, -1, 512 }, false); + @memcpy(state.words, seed); + state.words[0] = 77; + state.words[4 + 77] = 0; + state.words[4 + layout.experts_n + 77] = 90; + state.words[layout.miscOffset(2)] = 500; + fixture.checkpoint(words, state.words); + state.words[layout.miscOffset(2)] += 10; + fixture.checkpoint(words, state.words); + _ = try capture(words, &state, &.{ 77, 77, 88, -1, -1, -1, -1, -1, -1, -1, 512 }, false); + state.words[layout.miscOffset(2)] += 10; + fixture.checkpoint(words, state.words); + const view = try trace.View.init(words); + const result = try replay.verifyTraceWithFinal(allocator, view, state.words); + try std.testing.expectEqual(@as(usize, 2), result.matched_transitions); + try std.testing.expectEqual(@as(u64, 2), result.delta.hit_picks); + try std.testing.expectEqual(@as(u64, 4), result.delta.miss_picks); + const comparison = try policies.evaluateTrace(allocator, view); + try std.testing.expectEqual(result.delta.hit_picks, comparison.lifetime.hit_picks); + try std.testing.expectEqual(result.delta.miss_picks, comparison.lifetime.miss_picks); + try std.testing.expectEqual(result.delta.copies, comparison.lifetime.copies); + try std.testing.expectEqual(result.delta.evictions, comparison.lifetime.evictions); + try std.testing.expectEqual(result.delta.inplace, comparison.lifetime.mapped_experts); + try std.testing.expectEqual(result.delta.mapped_picks, comparison.lifetime.mapped_picks); + try std.testing.expectEqual(result.rounds, comparison.lifetime.rounds); + try std.testing.expectEqual(@as(usize, 4), view.checkpoints()); + for ([_]usize{ 0, 1, 1, 2 }, 0..) |round, index| try std.testing.expectEqual(round, (try view.checkpoint(index)).roundIndex()); + const bytes = try view.encode(allocator); + defer allocator.free(bytes); + try std.testing.expectEqual(trace.pool_bytes, bytes.len); + var snapshot = try trace.Snapshot.decode(allocator, bytes); + defer snapshot.deinit(); + _ = try replay.verifyTraceWithFinal(allocator, try snapshot.view(), state.words); + state.words[0] ^= 1; + try std.testing.expectError(error.PlannerFinalStateMismatch, replay.verifyTraceWithFinal(allocator, view, state.words)); +} + +test "exl3: unexplained handoff and corrupted outgoing checkpoint cannot be reseeded away" { + const allocator = std.testing.allocator; + const words = try fixture.make(allocator, 2, 1); + defer allocator.free(words); + const seed = try fixture.emptySeed(allocator, 2); + defer allocator.free(seed); + var state = try replay.ExactState.init(allocator, 2, seed); + defer state.deinit(); + const picks: [11]i32 = @splat(-1); + _ = try capture(words, &state, &picks, false); + state.words[layout.miscOffset(2)] += 100; + _ = try capture(words, &state, &picks, false); + try std.testing.expectError(error.TraceDiscontinuous, replay.verifyTrace(allocator, try trace.View.init(words))); + words[12] = 0; + try std.testing.expectError(error.PlannerCounterMismatch, replay.verifyTrace(allocator, try trace.View.init(words))); + words[12] = 1; + fixture.checkpoint(words, state.words); + try std.testing.expectError(error.TraceDiscontinuous, replay.verifyTrace(allocator, try trace.View.init(words))); + + @memset(words, 0); + @memcpy(words[0..trace.header_words], &try trace.header(0, 0, 2, 1)); + @memcpy(state.words, seed); + _ = try capture(words, &state, &picks, false); + state.words[layout.miscOffset(2)] += 100; + fixture.checkpoint(words, state.words); + _ = try capture(words, &state, &picks, false); + const view = try trace.View.init(words); + _ = try replay.verifyTrace(allocator, view); + const offset = trace.pool_words - trace.checkpointWords(words[6]); + for (0..words[6]) |index| { + words[offset + 4 + index] ^= 1; + try std.testing.expectError(error.PlannerCheckpointMismatch, replay.verifyTrace(allocator, try trace.View.init(words))); + words[offset + 4 + index] ^= 1; + } + words[offset + 4 + words[6]] = 9999; + try std.testing.expectError(error.InvalidSeed, replay.verifyTrace(allocator, try trace.View.init(words))); +} + +test "exl3: dropped ownership checkpoints invalidate exact replay even without decode rounds" { + const allocator = std.testing.allocator; + const words = try fixture.make(allocator, 2, 16); + defer allocator.free(words); + const seed = try fixture.emptySeed(allocator, 2); + defer allocator.free(seed); + const capacity = (trace.pool_words - trace.header_words - 2 * words[6]) / trace.checkpointWords(words[6]); + for (0..capacity + 1) |_| fixture.checkpoint(words, seed); + try std.testing.expectError(error.TraceDroppedRecords, replay.verifyTrace(allocator, try trace.View.init(words))); +} + +test "exl3: v4 exact warm replay matches known decisions counters and all terminal words" { + const allocator = std.testing.allocator; + const words = try fixture.make(allocator, 2, 1); + defer allocator.free(words); + const seed = try fixture.emptySeed(allocator, 2); + defer allocator.free(seed); + seed[0] = 0; + seed[1] = 1; + seed[2] = 10; + seed[3] = 20; + seed[4] = 0; + seed[5] = 1; + seed[4 + 512] = 5; + seed[5 + 512] = 3; + seed[layout.miscOffset(2)] = 20; + var state = try replay.ExactState.init(allocator, 2, seed); + defer state.deinit(); + const first = try capture(words, &state, &.{ 2, 2, 2, 2, 2, 3, -1, -1, -1, -1, 512 }, false); + try std.testing.expectEqual(replay.Delta{ .miss_picks = 6, .copies = 2, .evictions = 2 }, first); + try std.testing.expectEqualSlices(u32, &.{ 3, 2, 21, 21 }, state.words[0..4]); + const second = try capture(words, &state, &.{ 2, 2, 4, -1, -1, -1, -1, -1, -1, -1, -1 }, false); + try std.testing.expectEqual(replay.Delta{ .hit_picks = 2, .miss_picks = 1, .copies = 1, .evictions = 1 }, second); + const view = try trace.View.init(words); + const result = try replay.verifyTraceWithFinal(allocator, view, state.words); + try std.testing.expectEqual(@as(usize, 2), result.matched_transitions); + try std.testing.expectEqual(replay.Delta{ .hit_picks = 2, .miss_picks = 7, .copies = 3, .evictions = 3 }, result.delta); + const comparison = try policies.evaluate(allocator, &.{ try view.round(0), try view.round(1) }, .{ .capacity = 2, .seed = view.seedState(policies.Seed) }); + try std.testing.expectEqual(result.delta.hit_picks, comparison.lifetime.hit_picks); + try std.testing.expectEqual(result.delta.miss_picks, comparison.lifetime.miss_picks); + try std.testing.expectEqual(result.delta.copies, comparison.lifetime.copies); + try std.testing.expectEqual(result.delta.evictions, comparison.lifetime.evictions); + state.words[state.words.len - 1] = 1; + try std.testing.expectError(error.PlannerFinalStateMismatch, replay.verifyTraceWithFinal(allocator, view, state.words)); +} + +test "exl3: v4 1200 full rounds replay recent wrap saturated counts and modular counters" { + const allocator = std.testing.allocator; + const words = try fixture.make(allocator, 10, 16); + defer allocator.free(words); + const seed = try fixture.emptySeed(allocator, 10); + defer allocator.free(seed); + seed[20 + 512 + 7] = std.math.maxInt(u32); + const misc = layout.miscOffset(10); + @memset(seed[misc .. misc + 8], std.math.maxInt(u32) - 1); + var state = try replay.ExactState.init(allocator, 10, seed); + defer state.deinit(); + var total: replay.Delta = .{}; + for (0..1200) |round| { + var picks: [16 * 11]i32 = undefined; + for (&picks, 0..) |*pick, index| { + pick.* = switch (index % 11) { + 0 => -1, + 1 => 512, + 2 => -99, + 3 => 800, + else => @intCast((round * 7 + index * 3) % 512), + }; + } + const delta = try capture(words, &state, &picks, round % 3 != 0); + inline for (comptime std.meta.fieldNames(replay.Delta)) |field| @field(total, field) += @field(delta, field); + } + const view = try trace.View.init(words); + try std.testing.expectEqual(@as(usize, 1200), view.recorded()); + try std.testing.expectEqual(@as(u64, 0), view.dropped()); + const result = try replay.verifyTraceWithFinal(allocator, view, state.words); + try std.testing.expectEqual(total, result.delta); + try std.testing.expectEqual(std.math.maxInt(u32), state.words[20 + 512 + 7]); + try std.testing.expectEqual(@as(u32, 1198), state.words[misc]); + try std.testing.expectEqual(@as(u32, 256), state.words[misc + 9]); + const bytes = try view.encode(allocator); + defer allocator.free(bytes); + try std.testing.expect(bytes.len <= trace.pool_bytes); + var snapshot = try trace.Snapshot.decode(allocator, bytes); + defer snapshot.deinit(); + _ = try replay.verifyTrace(allocator, try snapshot.view()); +} + +test "exl3: v4 starts from warm recent history and retains empty masked rounds" { + const allocator = std.testing.allocator; + const words = try fixture.make(allocator, 2, 1); + defer allocator.free(words); + const seed = try fixture.emptySeed(allocator, 2); + defer allocator.free(seed); + seed[layout.recentCountsOffset(2) + 7] = 2; + seed[layout.recentOffset(2)] = 1 << 7; + seed[layout.recentOffset(2) + 16] = 1 << 7; + seed[layout.miscOffset(2) + 8] = 2; + seed[layout.miscOffset(2) + 9] = 2; + var state = try replay.ExactState.init(allocator, 2, seed); + defer state.deinit(); + _ = try capture(words, &state, &.{ 7, 7, -1, 512, 999, -1, -1, -1, -1, -1, -1 }, true); + try std.testing.expectEqual(@as(u32, 3), state.words[layout.recentCountsOffset(2) + 7]); + try std.testing.expectEqual(@as(u32, 3), state.words[layout.miscOffset(2) + 8]); + const masked = try capture(words, &state, &@as([11]i32, @splat(-1)), true); + try std.testing.expectEqual(replay.Delta{}, masked); + const result = try replay.verifyTraceWithFinal(allocator, try trace.View.init(words), state.words); + try std.testing.expectEqual(@as(usize, 2), result.rounds); +} + +test "exl3: v4 mapped experts count repeated reads and mask shared and invalid slots" { + const allocator = std.testing.allocator; + const words = try fixture.make(allocator, 1, 1); + defer allocator.free(words); + const seed = try fixture.emptySeed(allocator, 1); + defer allocator.free(seed); + seed[0] = 0; + seed[2] = 0; + var state = try replay.ExactState.init(allocator, 1, seed); + defer state.deinit(); + const delta = try capture(words, &state, &.{ 0, 0, 1, 1, 1, 512, -1, 700, -9, 512, -1 }, false); + try std.testing.expectEqual(replay.Delta{ .hit_picks = 2, .miss_picks = 3, .inplace = 1, .mapped_picks = 3 }, delta); + const view = try trace.View.init(words); + _ = try replay.verifyTrace(allocator, view); + const comparison = try policies.evaluateTrace(allocator, view); + try std.testing.expectEqual(delta.hit_picks, comparison.lifetime.hit_picks); + try std.testing.expectEqual(delta.miss_picks, comparison.lifetime.miss_picks); + try std.testing.expectEqual(delta.inplace, comparison.lifetime.mapped_experts); + try std.testing.expectEqual(delta.mapped_picks, comparison.lifetime.mapped_picks); + const offset = (try trace.View.init(words)).recordsOffset(); + words[offset + 14] = 2; + try std.testing.expectError(error.PlannerCounterMismatch, replay.verifyTrace(allocator, try trace.View.init(words))); +} + +test "exl3: v4 replay rejects altered production outcomes every terminal field and bad seeds" { + const allocator = std.testing.allocator; + const words = try fixture.make(allocator, 2, 1); + defer allocator.free(words); + const seed = try fixture.emptySeed(allocator, 2); + defer allocator.free(seed); + var state = try replay.ExactState.init(allocator, 2, seed); + defer state.deinit(); + _ = try capture(words, &state, &.{ 2, 2, 3, -1, -1, -1, -1, -1, -1, -1, 512 }, false); + const view = try trace.View.init(words); + for (0..layout.misc_words) |index| { + const offset = view.recordsOffset() + 4 + index; + words[offset] ^= 1; + try std.testing.expectError(error.PlannerCounterMismatch, replay.verifyTrace(allocator, try trace.View.init(words))); + words[offset] ^= 1; + } + const terminal_offset = trace.header_words + words[6]; + for (0..words[6]) |index| { + words[terminal_offset + index] ^= 1; + try std.testing.expectError(error.PlannerFinalStateMismatch, replay.verifyTrace(allocator, try trace.View.init(words))); + words[terminal_offset + index] ^= 1; + } + for ([_]usize{ 0, 4, layout.recentCountsOffset(2), layout.miscOffset(2) + 8 }) |index| { + const saved = words[trace.header_words + index]; + words[trace.header_words + index] = 9999; + try std.testing.expectError(error.InvalidSeed, replay.verifyTrace(allocator, try trace.View.init(words))); + words[trace.header_words + index] = saved; + } +} + +test "exl3: v4 overflow dropped empty invalid and discontinuous timelines cannot pass" { + const allocator = std.testing.allocator; + const words = try fixture.make(allocator, 2, 47); + defer allocator.free(words); + try std.testing.expectError(error.EmptyTrace, replay.verifyTrace(allocator, try trace.View.init(words))); + const seed = try fixture.emptySeed(allocator, 2); + defer allocator.free(seed); + var state = try replay.ExactState.init(allocator, 2, seed); + defer state.deinit(); + var picks: [47 * 11]i32 = @splat(512); + for (0..512) |index| picks[index] = @intCast(index); + const delta = try capture(words, &state, &picks, false); + try std.testing.expectEqual(@as(u64, 512), delta.miss_picks); + try std.testing.expectEqual(@as(u64, 0), delta.copies); + try std.testing.expectEqual(@as(u64, 1), delta.overflow); + try std.testing.expectError(error.PlannerOverflow, replay.verifyTrace(allocator, try trace.View.init(words))); + words[12] = 1; + try std.testing.expectError(error.TraceDiscontinuous, replay.verifyTrace(allocator, try trace.View.init(words))); + words[12] = 0; + words[13] = 1; + try std.testing.expectError(error.TraceInvalidLaunch, replay.verifyTrace(allocator, try trace.View.init(words))); + words[13] = 0; + for (1..words[8] + 1) |_| _ = try capture(words, &state, &picks, false); + try std.testing.expectError(error.TraceDroppedRecords, replay.verifyTrace(allocator, try trace.View.init(words))); +} diff --git a/zig/src/families/flashnext/exl3/dplan/exl3_dplan_trace_test.zig b/zig/src/families/flashnext/exl3/dplan/exl3_dplan_trace_test.zig new file mode 100644 index 0000000000..bc9ee14235 --- /dev/null +++ b/zig/src/families/flashnext/exl3/dplan/exl3_dplan_trace_test.zig @@ -0,0 +1,262 @@ +const std = @import("std"); +const trace = @import("exl3_dplan_trace.zig"); +const layout = @import("exl3_dplan_layout.zig"); + +pub fn emptySeed(allocator: std.mem.Allocator, lines: usize) ![]u32 { + const seed = try allocator.alloc(u32, layout.stateWords(lines)); + @memset(seed, 0); + @memset(seed[0..lines], std.math.maxInt(u32)); + @memset(seed[2 * lines ..][0..layout.experts_n], std.math.maxInt(u32)); + return seed; +} + +pub fn make(allocator: std.mem.Allocator, lines: usize, rows: usize) ![]u32 { + const words = try allocator.alloc(u32, trace.pool_words); + errdefer allocator.free(words); + @memset(words, 0); + @memcpy(words[0..trace.header_words], &try trace.header(3, 48, lines, rows)); + return words; +} + +pub fn append(words: []u32, before: []const u32, after: []const u32, picks: []const i32, recent: bool, mapped: u32) void { + const records_end = trace.header_words + 2 * words[6] + words[9] * words[7]; + const checkpoints_start = trace.pool_words - words[14] * trace.checkpointWords(words[6]); + if (checkpoints_start - records_end < words[7] or words[15] != 0) { + const dropped = (@as(u64, words[11]) << 32 | words[10]) +| 1; + words[10] = @truncate(dropped); + words[11] = @truncate(dropped >> 32); + return; + } + const seed = words[trace.header_words..][0..words[6]]; + const terminal = words[trace.header_words + words[6] ..][0..words[6]]; + if (words[9] == 0 and words[14] == 0) { + @memcpy(seed, before); + } else if (!std.mem.eql(u32, terminal, before)) { + words[12] = 1; + } + const offset = trace.header_words + 2 * words[6] + words[9] * words[7]; + const record = words[offset..][0..words[7]]; + @memset(record, 0); + record[0] = @intCast(picks.len / 11); + record[3] = @intFromBool(recent); + @memcpy(record[4..][0..layout.misc_words], after[layout.miscOffset(words[4])..]); + record[14] = mapped; + var seen: [layout.experts_n]bool = @splat(false); + for (picks, 0..) |expert, index| { + record[trace.record_header_words + index] = @bitCast(expert); + if (expert < 0 or expert >= layout.experts_n) continue; + const valid: usize = @intCast(expert); + record[1] += 1; + if (!seen[valid]) record[2] += 1; + seen[valid] = true; + } + @memcpy(terminal, after); + words[9] += 1; +} + +pub fn checkpoint(words: []u32, incoming: []const u32) void { + const stride = trace.checkpointWords(words[6]); + const end = trace.header_words + 2 * words[6] + words[9] * words[7]; + const start = trace.pool_words - words[14] * stride; + if (start - end < stride or words[10] != 0 or words[11] != 0 or words[15] != 0) { + words[15] +|= 1; + return; + } + const initial = words[trace.header_words..][0..words[6]]; + const terminal = words[trace.header_words + words[6] ..][0..words[6]]; + if (words[9] == 0 and words[14] == 0) { + @memcpy(initial, incoming); + @memcpy(terminal, incoming); + } + const entry = words[start - stride .. start]; + entry[0..4].* = .{ words[9], 1, words[6], words[14] }; + @memcpy(entry[4..][0..words[6]], terminal); + @memcpy(entry[4 + words[6] ..], incoming); + @memcpy(terminal, incoming); + words[14] += 1; +} + +test "exl3: checkpoints reject corrupt ordinals shapes order overlap and truncation" { + const allocator = std.testing.allocator; + const words = try make(allocator, 2, 1); + defer allocator.free(words); + const seed = try emptySeed(allocator, 2); + defer allocator.free(seed); + const picks: [11]i32 = @splat(-1); + append(words, seed, seed, &picks, false, 0); + checkpoint(words, seed); + append(words, seed, seed, &picks, false, 0); + checkpoint(words, seed); + const view = try trace.View.init(words); + try std.testing.expectError(error.TraceCheckpointOutOfRange, view.checkpoint(2)); + const stride = trace.checkpointWords(words[6]); + const offset = trace.pool_words - stride; + for (0..4) |field| { + const saved = words[offset + field]; + words[offset + field] = std.math.maxInt(u32); + try std.testing.expectError(error.InvalidTrace, trace.View.init(words)); + words[offset + field] = saved; + } + words[offset - stride] = 0; + try std.testing.expectError(error.InvalidTrace, trace.View.init(words)); + words[offset - stride] = 2; + const saved_count = words[14]; + words[14] = std.math.maxInt(u32); + try std.testing.expectError(error.InvalidTrace, trace.View.init(words)); + words[14] = @intCast((trace.pool_words - view.recordsOffset()) / stride); + words[9] = words[8]; + try std.testing.expectError(error.InvalidTrace, trace.View.init(words)); + words[14] = saved_count; + words[9] = 2; + const bytes = try view.encode(allocator); + defer allocator.free(bytes); + try std.testing.expectError(error.InvalidTrace, trace.Snapshot.decode(allocator, bytes[0 .. bytes.len - 4])); + try std.testing.expectError(error.InvalidTrace, trace.View.init(words[0 .. words.len - 1])); + words[15] = 1; + try std.testing.expectError(error.InvalidTrace, trace.View.init(words)); +} + +test "exl3: checkpoint exhaustion saturates without overwriting retained records or states" { + const allocator = std.testing.allocator; + const words = try make(allocator, 2, 16); + defer allocator.free(words); + const seed = try emptySeed(allocator, 2); + defer allocator.free(seed); + const capacity = (trace.pool_words - trace.header_words - 2 * words[6]) / trace.checkpointWords(words[6]); + for (0..capacity) |_| checkpoint(words, seed); + const retained = try allocator.dupe(u32, words); + defer allocator.free(retained); + seed[layout.miscOffset(2)] += 20; + checkpoint(words, seed); + try std.testing.expectEqual(@as(u32, 1), (try trace.View.init(words)).droppedCheckpoints()); + try std.testing.expectEqualSlices(u32, retained[trace.header_words..], words[trace.header_words..]); + words[15] = std.math.maxInt(u32); + checkpoint(words, seed); + try std.testing.expectEqual(std.math.maxInt(u32), (try trace.View.init(words)).droppedCheckpoints()); + append(words, seed, seed, &@as([11]i32, @splat(-1)), false, 0); + try std.testing.expectEqual(@as(u64, 1), (try trace.View.init(words)).dropped()); +} + +test "exl3: v4 trace preserves raw row order seed checkpoint and post counters" { + const allocator = std.testing.allocator; + const words = try make(allocator, 2, 2); + defer allocator.free(words); + const seed = try emptySeed(allocator, 2); + defer allocator.free(seed); + seed[layout.miscOffset(2)] = 0xfffffffe; + const picks = [_]i32{ 7, 2, 7, -1, 512, 4, 2, 2, 0, -9, 800, 4, 7, 2, 511, -1, 512, 4, 0, 0, 7, 511 }; + append(words, seed, seed, &picks, true, 3); + const view = try trace.View.init(words); + const bytes = try view.encode(allocator); + defer allocator.free(bytes); + var decoded = try trace.Snapshot.decode(allocator, bytes); + defer decoded.deinit(); + const restored = try decoded.view(); + const record = try restored.round(0); + try std.testing.expectEqual(@as(usize, 2), record.rows()); + try std.testing.expectEqual(@as(u32, 16), record.validPicks()); + try std.testing.expectEqual(@as(usize, 5), record.uniqueCount()); + try std.testing.expect(record.recentEnabled()); + try std.testing.expectEqual(@as(u32, 3), record.mappedPicks()); + try std.testing.expectEqualSlices(u32, seed, restored.initialSeed()); + try std.testing.expectEqualSlices(u32, seed, restored.terminalSeed()); + try std.testing.expectEqual(@as(u32, 0xfffffffe), record.misc()[0]); + for (picks, record.picks()) |expert, bits| try std.testing.expectEqual(expert, @as(i32, @bitCast(bits))); + try std.testing.expectEqual(trace.magic, std.mem.readInt(u32, bytes[0..4], .little)); + try std.testing.expectEqual((trace.header_words + 2 * seed.len + words[7]) * 4, bytes.len); + try std.testing.expectError(error.TraceRoundOutOfRange, view.round(1)); +} + +test "exl3: v4 capacity exceeds 1000 with full 16-row records before ownership checkpoints" { + for ([_]usize{ 1, 10, 73, 256, 512 }) |lines| { + const shape = try trace.header(0, 0, lines, 16); + try std.testing.expect(shape[8] >= 1000); + const used = trace.header_words + 2 * shape[6] + shape[8] * shape[7]; + try std.testing.expect(used <= trace.pool_words); + try std.testing.expect(used + shape[7] > trace.pool_words); + try std.testing.expectEqual(@as(u32, 191), shape[7]); + } + try std.testing.expectError(error.TraceShapeTooLarge, trace.header(0, 0, trace.pool_words, 1)); + try std.testing.expectError(error.TraceShapeTooLarge, trace.header(0, 0, 2, std.math.maxInt(usize))); + try std.testing.expectError(error.TraceShapeTooLarge, trace.header(0, 0, 2, 0)); + try std.testing.expect(!trace.enabledValue(null)); + try std.testing.expect(!trace.enabledValue("true")); + try std.testing.expect(!trace.enabledValue("0")); + try std.testing.expect(trace.enabledValue("1")); +} + +test "exl3: v4 keeps initial seed detects discontinuity and saturates bounded drops" { + const allocator = std.testing.allocator; + const words = try make(allocator, 2, 16); + defer allocator.free(words); + const seed = try emptySeed(allocator, 2); + defer allocator.free(seed); + const picks: [11]i32 = @splat(-1); + append(words, seed, seed, &picks, false, 0); + seed[layout.miscOffset(2)] = 700; + append(words, seed, seed, &picks, false, 0); + try std.testing.expectEqual(@as(u32, 1), words[12]); + try std.testing.expectEqual(@as(u32, 0), (try trace.View.init(words)).initialSeed()[layout.miscOffset(2)]); + for (2..words[8] + 3) |_| append(words, seed, seed, &picks, false, 0); + var view = try trace.View.init(words); + try std.testing.expectEqual(@as(usize, words[8]), view.recorded()); + try std.testing.expectEqual(@as(u64, 3), view.dropped()); + words[10] = std.math.maxInt(u32); + words[11] = 4; + append(words, seed, seed, &picks, false, 0); + try std.testing.expectEqual(@as(u64, 5) << 32, (try trace.View.init(words)).dropped()); + words[10] = std.math.maxInt(u32); + words[11] = std.math.maxInt(u32); + append(words, seed, seed, &picks, false, 0); + view = try trace.View.init(words); + try std.testing.expectEqual(std.math.maxInt(u64), view.dropped()); +} + +test "exl3: v4 rejects metadata corruption truncation old versions and false pick counts" { + const allocator = std.testing.allocator; + const words = try make(allocator, 2, 1); + defer allocator.free(words); + const seed = try emptySeed(allocator, 2); + defer allocator.free(seed); + append(words, seed, seed, &@as([11]i32, @splat(7)), false, 0); + const view = try trace.View.init(words); + const bytes = try view.encode(allocator); + defer allocator.free(bytes); + try std.testing.expectError(error.InvalidTrace, trace.Snapshot.decode(allocator, bytes[0 .. bytes.len - 4])); + try std.testing.expectError(error.InvalidTrace, trace.Snapshot.decode(allocator, bytes[0 .. bytes.len - 1])); + const offset = view.recordsOffset(); + for ([_]usize{ 1, 2, 3, 14 }) |field| { + const saved = words[offset + field]; + words[offset + field] = std.math.maxInt(u32); + try std.testing.expectError(error.InvalidTrace, trace.View.init(words)); + words[offset + field] = saved; + } + for ([_]usize{ 1, 6, 7, 8, 9, 12, 13, 14, 15 }) |field| { + const saved = words[field]; + words[field] = std.math.maxInt(u32); + try std.testing.expectError(error.InvalidTrace, trace.View.init(words)); + words[field] = saved; + } + words[1] = 2; + try std.testing.expectError(error.InvalidTrace, trace.View.init(words)); + words[1] = 3; + try std.testing.expectError(error.InvalidTrace, trace.View.init(words)); + words[1] = trace.version; + words[10] = 1; + try std.testing.expectError(error.InvalidTrace, trace.View.init(words)); +} + +test "exl3: v4 raw picks cover all 512 experts beyond the planner unique bound" { + const allocator = std.testing.allocator; + const words = try make(allocator, 2, 47); + defer allocator.free(words); + const seed = try emptySeed(allocator, 2); + defer allocator.free(seed); + var picks: [47 * 11]i32 = @splat(512); + for (0..512) |index| picks[index] = @intCast(511 - index); + append(words, seed, seed, &picks, false, 0); + const record = try (try trace.View.init(words)).round(0); + try std.testing.expectEqual(@as(usize, 512), record.uniqueCount()); + try std.testing.expectEqual(@as(u32, 512), record.validPicks()); + for (picks, record.picks()) |expert, raw| try std.testing.expectEqual(expert, @as(i32, @bitCast(raw))); +} diff --git a/zig/src/families/flashnext/exl3/dplan/exl3_dplan_trace_writer_test.zig b/zig/src/families/flashnext/exl3/dplan/exl3_dplan_trace_writer_test.zig new file mode 100644 index 0000000000..3ab0088f56 --- /dev/null +++ b/zig/src/families/flashnext/exl3/dplan/exl3_dplan_trace_writer_test.zig @@ -0,0 +1,31 @@ +const std = @import("std"); +const trace = @import("exl3_dplan_trace.zig"); +const replay = @import("exl3_dplan_replay.zig"); + +test "exl3: CUDA writer bytes replay all 1200 rounds and terminal words exactly on CPU" { + var snapshot = try trace.Snapshot.decode(std.testing.allocator, @embedFile("dplan_trace_fixture")); + defer snapshot.deinit(); + const view = try snapshot.view(); + const result = try replay.verifyTrace(std.testing.allocator, view); + try std.testing.expectEqual(@as(usize, 1200), result.rounds); + try std.testing.expectEqual(@as(usize, 1200), result.matched_transitions); + try std.testing.expectEqual(@as(u64, 0), result.delta.overflow); + try std.testing.expect(result.delta.mapped_picks > 0); + try std.testing.expect(result.delta.copies > 0); + try std.testing.expect(result.delta.evictions > 0); +} + +test "exl3: CUDA ownership checkpoints replay every segment and final transferred state" { + var snapshot = try trace.Snapshot.decode(std.testing.allocator, @embedFile("dplan_checkpoint_fixture")); + defer snapshot.deinit(); + const view = try snapshot.view(); + const result = try replay.verifyTraceWithFinal(std.testing.allocator, view, view.terminalSeed()); + try std.testing.expectEqual(@as(usize, 1024), result.rounds); + try std.testing.expectEqual(@as(usize, 1024), result.matched_transitions); + try std.testing.expectEqual(@as(usize, 4), view.checkpoints()); + try std.testing.expectEqual(@as(usize, 0), (try view.checkpoint(0)).roundIndex()); + try std.testing.expectEqual(@as(usize, 512), (try view.checkpoint(1)).roundIndex()); + try std.testing.expectEqual(@as(usize, 768), (try view.checkpoint(2)).roundIndex()); + try std.testing.expectEqual(@as(usize, 1024), (try view.checkpoint(3)).roundIndex()); + try std.testing.expect(result.delta.copies > 0); +} diff --git a/zig/src/families/flashnext/exl3/graph/cuda_exl3_graph.zig b/zig/src/families/flashnext/exl3/graph/cuda_exl3_graph.zig new file mode 100644 index 0000000000..dc9002c784 --- /dev/null +++ b/zig/src/families/flashnext/exl3/graph/cuda_exl3_graph.zig @@ -0,0 +1,98 @@ +//! CUDA graphs for an EXL3 Flash Next decode step with the expert offload pool (TF_FLASHNEXT_EXL3_GRAPH=1). + +const std = @import("std"); +const cuda = @import("cuda"); +const state_driver = @import("../kv/cuda_exl3_state_binding.zig"); +const gkey = @import("cuda_exl3_graph_key.zig"); +const gpending = @import("cuda_exl3_graph_pending.zig"); +const gcache = @import("cuda_exl3_graph_cache.zig"); +const gbudget = @import("cuda_exl3_graph_budget.zig"); +const ghooks = @import("cuda_exl3_graph_hooks.zig"); +const gprefill = @import("cuda_exl3_graph_prefill.zig"); + +pub const memory = @import("cuda_exl3_graph_memory.zig"); + +/// TF_FLASHNEXT_EXL3_GRAPH=1: decode rounds and verify windows replay captured graphs. +pub fn enabled() bool { + if (force) |f| return f == .graph; + const v = std.c.getenv("TF_FLASHNEXT_EXL3_GRAPH") orelse return false; + return std.mem.eql(u8, std.mem.span(v), "1"); +} + +/// A test's override of the environment: the host cut, the device plan launched eagerly, or graphs. +pub const Mode = enum { off, dplan, graph }; +pub var force: ?Mode = null; + +pub const Key = gkey.Key; +pub const Pending = gpending.Pending; +pub const fallback = @import("cuda_exl3_graph_fallback.zig"); +pub const cache = gcache; +pub const rebind = @import("cuda_exl3_graph_rebind.zig"); +pub const asyncInstantiate = gpending.asyncInstantiate; +pub const startInstantiate = gpending.startInstantiate; +pub const takePending = gpending.takePending; +pub const uploadResident = gbudget.uploadResident; +pub const reserveMemory = gbudget.reserveMemory; +pub const clear = gcache.clear; +pub const releaseSeq = gcache.releaseSeq; +pub const releaseDecodeSeq = gcache.releaseDecodeSeq; +pub const releaseFor = gcache.releaseFor; +pub const round = ghooks.round; +pub const forward = ghooks.forward; +pub const Prefill = gprefill.Prefill; +pub const prefillBegin = gprefill.prefillBegin; +pub const prefillEnd = gprefill.prefillEnd; +pub const prefillAbort = gprefill.prefillAbort; +pub const prefillLayer = gprefill.runLayer; +pub const clearPrefill = gprefill.clearPrefill; +pub const clearPrefillRuntime = gprefill.clearRuntime; + +/// True while this module captures a step: the pass's pinned uploads become copy nodes (Uploads.copy). +pub var capturing: bool = false; + +/// A copy node from pinned host bytes (the pass's staging): the graph reads them at every replay, as `prepare` left +pub fn pinnedCopy(s: cuda.Stream, dst: u64, bytes: []const u8) !void { + if (bytes.len == 0) return; + try s.d.check(s.d.api.cuMemcpyHtoDAsync_v2(dst, bytes.ptr, bytes.len, s.handle), "cuMemcpyHtoDAsync"); +} + +pub const Stats = struct { + captures: u64 = 0, + replays: u64 = 0, + rebind_attempts: u64 = 0, + rebinds: u64 = 0, + rebind_failures: u64 = 0, + rebind_declined: u64 = 0, + /// Steps on the device plan launched one by one (TF_FLASHNEXT_EXL3_DPLAN=1 without graphs). + eager_dplan: u64 = 0, + /// Steps that left for the host cut (too many rows, media, no pools). + declined: u64 = 0, + capture_ns: u64 = 0, + /// Captures that updated a cached graph of the same shape in place instead of instantiating one. + updates: u64 = 0, + /// Host time of `prepare` (ids, positions, n-gram rows) before a step's launch, and how many. + prepare_ns: u64 = 0, + prepares: u64 = 0, + /// Captures or instantiations that failed; their rounds ran eagerly (cuda_exl3_graph_fallback.zig). + capture_failures: u64 = 0, +}; + +pub var stats: Stats = .{}; + +/// Captures `pass.chunk(ids, pos0, null)` (prepared) under `key`, or replays the graph captured under it. +pub const Outcome = enum { replay, eager, instantiated, updated }; +pub var coverage: [2][17][4]u32 = @splat(@splat(@splat(0))); + +pub var prefill_stats: struct { captures: u64 = 0, updates: u64 = 0, replays: u64 = 0, sighting_hits: u64 = 0, sighting_misses: u64 = 0, eager: u64 = 0, admission_misses: u64 = 0, capture_ns: u64 = 0 } = .{}; + +test "exl3: state binding helper contracts" { + _ = @import("../kv/exl3_state_binding_test.zig"); + _ = state_driver; + _ = fallback; + _ = @import("cuda_exl3_graph_fallback_test.zig"); + _ = @import("cuda_exl3_graph_key_test.zig"); + _ = @import("cuda_exl3_graph_rebind_test.zig"); + _ = @import("cuda_exl3_graph_cache_test.zig"); + _ = @import("cuda_exl3_graph_prefill_test.zig"); + _ = @import("cuda_exl3_graph_together_test.zig"); +} diff --git a/zig/src/families/flashnext/exl3/graph/cuda_exl3_graph_budget.zig b/zig/src/families/flashnext/exl3/graph/cuda_exl3_graph_budget.zig new file mode 100644 index 0000000000..cf0ac4eb96 --- /dev/null +++ b/zig/src/families/flashnext/exl3/graph/cuda_exl3_graph_budget.zig @@ -0,0 +1,128 @@ +//! The graph memory budget of the EXL3 decode graphs: admission, eviction and cohort retirement. + +const std = @import("std"); +const cuda = @import("cuda"); +const collective = @import("../../cuda_collective.zig"); +const memory = @import("cuda_exl3_graph_memory.zig"); +const hub = @import("cuda_exl3_graph.zig"); +const gcache = @import("cuda_exl3_graph_cache.zig"); +const gprefill = @import("cuda_exl3_graph_prefill.zig"); + +pub fn freeMemory(stream: cuda.Stream) !u64 { + var free: usize = 0; + var total: usize = 0; + try stream.d.check(stream.d.api.cuMemGetInfo_v2(&free, &total), "cuMemGetInfo"); + return free; +} + +pub fn uploadResident(exec: *cuda.graph.Exec, stream: cuda.Stream) !void { + try exec.upload(stream); + try stream.synchronize(); +} + +pub fn reserveMemory(stream: cuda.Stream, class: memory.Class, cohort: u64) !memory.Ticket { + _ = try memory.startupReservation(); + while (true) { + return memory.admitWithQuery(class, cohort, stream, freeMemory) catch |err| switch (err) { + error.GraphBudgetBusy, error.GraphMemoryHeadroom, error.GraphEvictionRequired => { + if (cohort != 0) return err; + try stream.d.check(stream.d.api.cuCtxSynchronize(), "cuCtxSynchronize(graph retirement)"); + if (!evictOne()) return err; + continue; + }, + else => return err, + }; + } +} + +pub fn retirementBarrier(c: *collective.Collective, signature: u64, cohort: u64) !void { + const decision = try c.graphAgree(signature, .{ .state = .ready, .cohort = cohort }); + if (decision.choice != .replay) return error.GraphCohortMismatch; +} + +pub fn oldestCohort() u64 { + var cohort: u64 = 0; + var used: u64 = std.math.maxInt(u64); + for (gcache.entries) |maybe| if (maybe) |entry| { + if (entry.cohort != 0 and (entry.used < used or (entry.used == used and (cohort == 0 or entry.cohort < cohort)))) { + cohort = entry.cohort; + used = entry.used; + } + }; + return cohort; +} + +pub fn retireCohort(cohort: u64) void { + for (&gcache.entries) |*maybe| if (maybe.*) |*entry| { + if (entry.cohort == cohort) { + gcache.drop(entry); + maybe.* = null; + } + }; +} + +pub fn reserveTogether(c: *collective.Collective, stream: cuda.Stream, signature: u64, cohort: u64) !memory.Ticket { + _ = try memory.startupReservation(); + var retired: usize = 0; + while (true) { + var ticket: ?memory.Ticket = memory.admitWithQuery(.decode, cohort, stream, freeMemory) catch |err| switch (err) { + error.GraphBudgetBusy, error.GraphBudgetExceeded, error.GraphMemoryHeadroom, error.GraphEvictionRequired => null, + else => return err, + }; + errdefer if (ticket) |held| memory.release(held) catch unreachable; + var free_slot = false; + for (gcache.entries) |entry| free_slot = free_slot or entry == null; + const agreement = try c.graphAgree(signature ^ 0x61646d6974, .{ .state = if (ticket != null and free_slot) .ready else .eager, .cohort = cohort }); + if (agreement.choice == .replay) return ticket.?; + if (ticket) |held| { + try memory.release(held); + ticket = null; + } + const victim = try c.graphVictim(signature ^ 0x76696374696d, oldestCohort()); + if (victim == 0) { + try stream.d.check(stream.d.api.cuCtxSynchronize(), "cuCtxSynchronize(TP local graph retirement)"); + try retirementBarrier(c, signature ^ 0x6c6f63616c7175, cohort); + const reclaimed = evictOne(); + const any_reclaimed = try c.graphVictim(signature ^ 0x6c6f63616c646f, @intFromBool(reclaimed)); + if (any_reclaimed == 0) return error.GraphBudgetExceeded; + continue; + } + if (retired >= gcache.entries.len * 2) return error.GraphBudgetExceeded; + try stream.d.check(stream.d.api.cuCtxSynchronize(), "cuCtxSynchronize(TP graph retirement)"); + try retirementBarrier(c, signature ^ 0x71756965736365, victim); + retireCohort(victim); + try retirementBarrier(c, signature ^ 0x72657469726564, victim); + retired += 1; + } +} + +pub fn evictOne() bool { + var oldest: u64 = std.math.maxInt(u64); + var decode: ?usize = null; + var prefill: ?usize = null; + for (gcache.entries, 0..) |maybe, index| if (maybe) |entry| { + if (entry.cohort == 0 and entry.used < oldest) { + oldest = entry.used; + decode = index; + } + }; + for (gprefill.pentries, 0..) |maybe, index| if (maybe) |entry| { + if (entry.used < oldest) { + oldest = entry.used; + decode = null; + prefill = index; + } + }; + if (decode) |index| { + gcache.drop(&gcache.entries[index].?); + gcache.entries[index] = null; + return true; + } + if (prefill) |index| { + const entry = &gprefill.pentries[index].?; + gprefill.dropPrefill(entry); + gprefill.pentries[index] = null; + return true; + } + return @import("../model/cuda_exl3_chain.zig").evictOne(); +} diff --git a/zig/src/families/flashnext/exl3/graph/cuda_exl3_graph_cache.zig b/zig/src/families/flashnext/exl3/graph/cuda_exl3_graph_cache.zig new file mode 100644 index 0000000000..8ecbdd009e --- /dev/null +++ b/zig/src/families/flashnext/exl3/graph/cuda_exl3_graph_cache.zig @@ -0,0 +1,196 @@ +//! The cache of captured EXL3 decode graphs: entries, LRU slots, sightings and the retirement of a sequence's graphs. + +const std = @import("std"); +const cuda = @import("cuda"); +const state_binding = @import("../kv/exl3_state_binding.zig"); +const memory = @import("cuda_exl3_graph_memory.zig"); +const hub = @import("cuda_exl3_graph.zig"); +const gbudget = @import("cuda_exl3_graph_budget.zig"); +const glog = @import("cuda_exl3_graph_log.zig"); +const gprefill = @import("cuda_exl3_graph_prefill.zig"); +const gkey = @import("cuda_exl3_graph_key.zig"); +const gpending = @import("cuda_exl3_graph_pending.zig"); +const Key = gkey.Key; +const Pending = gpending.Pending; + +/// `template`: the graph an update took this exec's parameters from, kept while the exec uses them. +pub const Entry = struct { key: Key, exec: ?cuda.graph.Exec = null, used: u64, template: ?cuda.graph.Graph = null, pending: ?*Pending = null, cohort: u64 = 0, admission: ?memory.Ticket = null, stream: ?cuda.Stream = null, binding: ?state_binding.Binding = null, poisoned: bool = false }; + +/// A pending entry whose instantiation finished becomes a usable graph (or is dropped if it failed); true if usable. +pub fn settle(x: *Entry) bool { + if (x.poisoned) return false; + const pd = x.pending orelse return x.exec != null; + if (!pd.done.load(.acquire)) return false; + pd.thread.join(); + x.pending = null; + defer std.heap.c_allocator.destroy(pd); + if (pd.exec) |e| { + var exec = e; + if (x.stream) |stream| gbudget.uploadResident(&exec, stream) catch { + exec.deinit(); + pd.graph.deinit(); + return false; + }; + x.exec = exec; + if (pd.keep_template) x.template = pd.graph else pd.graph.deinit(); + if (x.admission) |ticket| memory.installed(ticket) catch unreachable; + return true; + } + pd.graph.deinit(); + return false; +} + +pub fn drop(x: *Entry) void { + if (x.binding) |*binding| binding.invalidate(); + if (x.stream) |stream| stream.synchronize() catch @panic("graph retirement synchronization failed"); + defer if (x.admission) |ticket| { + memory.release(ticket) catch unreachable; + x.admission = null; + }; + if (x.pending) |pd| { + pd.thread.join(); + if (pd.exec) |*e| e.deinit(); + pd.graph.deinit(); + std.heap.c_allocator.destroy(pd); + return; + } + if (x.exec) |*e| e.deinit(); + if (x.template) |*g| g.deinit(); +} + +/// Steps after which a graph nobody replayed counts as stale (another stream's, or a finished one's). +pub const stale_after: u64 = 512; + +/// Graphs kept (LRU): a round and up to `max_rows` window sizes per stream, at two seams. +pub const cache_n = 128; +pub var entries: [cache_n]?Entry = @splat(null); +pub var clock: u64 = 0; + +pub fn clear() void { + gprefill.clearPrefill(); + for (&entries) |*e| if (e.*) |*x| { + drop(x); + e.* = null; + }; + seen = @splat(null); +} + +pub fn references(key: Key, tail: u64) bool { + if (tail == 0) return false; + for (key.seq) |seq| if (seq[2] == tail) return true; + return false; +} + +pub fn releaseSeqImpl(seq: anytype, keep_prefill: bool) void { + const tail = seq.tail.ptr; + if (tail == 0) return; + for (&entries) |*entry| if (entry.*) |*x| if (references(x.key, tail)) { + drop(x); + entry.* = null; + }; + for (&seen) |*entry| if (entry.*) |x| if (references(x.key, tail)) { + entry.* = null; + }; + if (!keep_prefill) { + for (&gprefill.pentries) |*entry| if (entry.*) |*x| if (x.key.seq[2] == tail) { + gprefill.dropPrefill(x); + entry.* = null; + }; + for (&gprefill.pseen) |*entry| if (entry.*) |*x| if (x.key.seq[2] == tail) { + entry.* = null; + }; + } +} + +pub fn releaseSeq(seq: anytype) void { + releaseSeqImpl(seq, false); +} + +pub fn releaseDecodeSeq(seq: anytype) void { + releaseSeqImpl(seq, true); +} + +/// The engine's deinit: the graphs go (the engine releases its own planner state). +pub fn releaseFor(owner: *@import("../../cuda_prompt.zig").Engine) void { + const id = @intFromPtr(owner); + for (&entries) |*entry| if (entry.*) |*held| if (held.key.owner == id) { + drop(held); + entry.* = null; + }; + for (&seen) |*entry| if (entry.*) |held| if (held.key.owner == id) { + entry.* = null; + }; + gprefill.clearRuntime(@intFromPtr(&owner.exl3_prefill)); +} + +/// Drops `x` and empties its slot. +pub fn forget(x: *Entry) void { + for (&entries) |*en| if (en.*) |*y| if (y == x) { + drop(y); + en.* = null; + return; + }; +} + +pub fn find(key: Key) ?*Entry { + for (&entries) |*e| if (e.*) |*x| if (std.meta.eql(x.key, key)) return x; + return null; +} + +pub fn slot() *?Entry { + var best: usize = 0; + var best_used: u64 = std.math.maxInt(u64); + for (&entries, 0..) |*e, i| { + if (e.* == null) return e; + if (e.*.?.used < best_used) { + best_used = e.*.?.used; + best = i; + } + } + drop(&entries[best].?); + entries[best] = null; + return &entries[best]; +} + +/// Keys seen without a graph, and how often (a key is captured at its `after`-th sighting; one-off shapes, e.g. a +pub const Seen = struct { key: Key, n: u32, used: u64 }; +pub var seen: [256]?Seen = @splat(null); + +pub fn sighting(key: Key) u32 { + var oldest: usize = 0; + var oldest_used: u64 = std.math.maxInt(u64); + for (&seen, 0..) |*e, i| { + if (e.*) |*x| { + if (std.meta.eql(x.key, key)) { + x.n += 1; + x.used = clock; + return x.n; + } + if (x.used < oldest_used) { + oldest_used = x.used; + oldest = i; + } + } else { + oldest = i; + oldest_used = 0; + } + } + seen[oldest] = .{ .key = key, .n = 1, .used = clock }; + return 1; +} + +/// The sighting of a key that captures it: a one-off shape never pays for a capture. +pub const capture_after: u32 = 2; +/// A test's override of `capture_after` (1 captures at once); not a user switch. +pub var capture_after_override: ?u32 = null; + +pub fn after() u32 { + return capture_after_override orelse capture_after; +} + +/// A window (kind 2) is keyed by rows, state set and seams, so each is seen a few times and a capture costs more. +pub const window_capture_after: u32 = 12; + +pub fn afterKey(key: Key) u32 { + return capture_after_override orelse if (key.kind == 2) window_capture_after else capture_after; +} diff --git a/zig/src/families/flashnext/exl3/graph/cuda_exl3_graph_cache_test.zig b/zig/src/families/flashnext/exl3/graph/cuda_exl3_graph_cache_test.zig new file mode 100644 index 0000000000..05c0000a2c --- /dev/null +++ b/zig/src/families/flashnext/exl3/graph/cuda_exl3_graph_cache_test.zig @@ -0,0 +1,158 @@ +//! Host tests of the EXL3 graph cache: sequence retirement, pending deadlines and the TP template. + +const std = @import("std"); +const cuda = @import("cuda"); +const collective = @import("../../cuda_collective.zig"); +const wire = @import("../../cuda_rank_wire.zig"); +const gpending = @import("cuda_exl3_graph_pending.zig"); +const gcache = @import("cuda_exl3_graph_cache.zig"); +const gtogether = @import("cuda_exl3_graph_together.zig"); +const gprefill = @import("cuda_exl3_graph_prefill.zig"); +const gkey = @import("cuda_exl3_graph_key.zig"); +const gdriver = @import("cuda_exl3_graph_test_driver.zig"); +const Key = gkey.Key; +const Entry = gcache.Entry; +const Pending = gpending.Pending; +const PKey = gprefill.PKey; +const GraphTestDriver = gdriver.GraphTestDriver; + +test "exl3: sequence graph retirement removes mixed lanes and pending work without evicting a live owner" { + gcache.clear(); + defer gcache.clear(); + GraphTestDriver.exec_drops = 0; + GraphTestDriver.graph_drops = 0; + var driver = GraphTestDriver.driver(); + const exec: cuda.graph.Exec = .{ .d = &driver, .handle = @ptrFromInt(1) }; + var key: Key = .{ .kind = 1, .rows = 1, .scratch = 1, .scores = 2, .output = 3 }; + key.seq[0][2] = 100; + const owner_key = key; + gcache.entries[0] = .{ .key = key, .exec = exec, .used = 1 }; + key.seq[0][2] = 200; + const live_key = key; + gcache.entries[1] = .{ .key = key, .exec = exec, .used = 1 }; + key.rows = 2; + key.seq[1][2] = 100; + gcache.entries[2] = .{ .key = key, .exec = exec, .used = 1 }; + const pending = try std.heap.c_allocator.create(Pending); + pending.* = .{ .graph = .{ .d = &driver, .handle = @ptrFromInt(1) }, .ctx = null, .keep_template = true }; + pending.thread = try std.Thread.spawn(.{}, GraphTestDriver.finishPending, .{ pending, &driver }); + gcache.entries[3] = .{ .key = owner_key, .pending = pending, .used = 1 }; + const live_pending = try std.heap.c_allocator.create(Pending); + live_pending.* = .{ .graph = .{ .d = &driver, .handle = @ptrFromInt(1) }, .ctx = null, .keep_template = true }; + live_pending.thread = try std.Thread.spawn(.{}, GraphTestDriver.finishPending, .{ live_pending, &driver }); + var live_pending_key = live_key; + live_pending_key.rows = 3; + gcache.entries[4] = .{ .key = live_pending_key, .pending = live_pending, .used = 1 }; + const pkey: PKey = .{ .layer = 0, .rows = 32, .sub = 16, .pos0 = 0, .deep = 1, .base = 2, .base_cap = 32, .seq = .{ 1, 2, 100, 32, 3 }, .wait_count = 0, .wait_key = 0, .scores = 4, .tiles = false }; + gprefill.pentries[0] = .{ .key = pkey, .exec = exec, .used = 1 }; + var live_pkey = pkey; + live_pkey.seq[2] = 200; + gprefill.pentries[1] = .{ .key = live_pkey, .exec = exec, .used = 1 }; + gcache.seen[0] = .{ .key = owner_key, .n = 10, .used = 1 }; + gcache.seen[1] = .{ .key = live_key, .n = 10, .used = 1 }; + gprefill.pseen[0] = .{ .key = pkey, .n = 10 }; + gprefill.pseen[1] = .{ .key = live_pkey, .n = 10 }; + const Seq = struct { tail: struct { ptr: u64 } }; + gcache.releaseSeq(Seq{ .tail = .{ .ptr = 100 } }); + try std.testing.expect(gcache.find(owner_key) == null); + try std.testing.expect(gcache.find(live_key) != null); + try std.testing.expect(gcache.entries[2] == null and gcache.entries[3] == null); + try std.testing.expect(gcache.entries[4].?.pending == live_pending); + try std.testing.expect(gprefill.pentries[0] == null and gprefill.pentries[1] != null); + try std.testing.expect(gcache.seen[0] == null and gcache.seen[1] != null); + try std.testing.expect(gprefill.pseen[0] == null and gprefill.pseen[1] != null); + try std.testing.expectEqual(@as(usize, 4), GraphTestDriver.exec_drops); + try std.testing.expectEqual(@as(usize, 1), GraphTestDriver.graph_drops); + try std.testing.expectEqual(@as(u32, 1), gcache.sighting(owner_key)); + gcache.releaseSeq(Seq{ .tail = .{ .ptr = 100 } }); + try std.testing.expectEqual(@as(usize, 4), GraphTestDriver.exec_drops); +} + +test "exl3: snapshot graph retirement keeps prefill templates but drops decode graphs" { + gcache.clear(); + defer gcache.clear(); + GraphTestDriver.exec_drops = 0; + GraphTestDriver.graph_drops = 0; + var driver = GraphTestDriver.driver(); + const exec: cuda.graph.Exec = .{ .d = &driver, .handle = @ptrFromInt(1) }; + var key: Key = .{ .kind = 1, .rows = 1, .scratch = 1, .scores = 2, .output = 3 }; + key.seq[0][2] = 100; + gcache.entries[0] = .{ .key = key, .exec = exec, .used = 1 }; + const pkey: PKey = .{ .layer = 0, .rows = 32, .sub = 16, .pos0 = 0, .deep = 1, .base = 2, .base_cap = 32, .seq = .{ 1, 2, 100, 32, 3 }, .wait_count = 0, .wait_key = 0, .scores = 4, .tiles = false }; + gprefill.pentries[0] = .{ .key = pkey, .exec = exec, .used = 1 }; + gprefill.pseen[0] = .{ .key = pkey, .n = 2 }; + const Seq = struct { tail: struct { ptr: u64 } }; + gcache.releaseDecodeSeq(Seq{ .tail = .{ .ptr = 100 } }); + try std.testing.expect(gcache.entries[0] == null); + try std.testing.expect(gprefill.pentries[0] != null); + try std.testing.expect(gprefill.pseen[0] != null); +} + +test "exl3: pending graph deadlines and failed instantiation never expose an executable" { + var pending: Pending = .{ .graph = undefined, .ctx = null, .keep_template = true, .started_ns = wire.now() }; + var entry: Entry = .{ .key = .{ .kind = 1, .rows = 1, .scratch = 1, .scores = 2, .output = 3 }, .pending = &pending, .used = 1, .cohort = 7 }; + try std.testing.expectEqual(collective.GraphState.pending, gtogether.graphVote(&entry, 1000).state); + pending.started_ns = wire.now() -| 2 * std.time.ns_per_ms; + try std.testing.expectEqual(collective.GraphState.timed_out, gtogether.graphVote(&entry, 1).state); + entry.pending = null; + try std.testing.expectEqual(collective.GraphState.failed, gtogether.graphVote(&entry, 1000).state); + gcache.drop(&entry); +} + +test "exl3: asynchronous TP instantiation preserves its NCCL template until owner retirement" { + if (!gpending.asyncInstantiate()) return error.SkipZigTest; + const previous_api = gpending.ctx_api; + defer gpending.ctx_api = previous_api; + gpending.ctx_api = .{ .get = GraphTestDriver.getContext, .set = GraphTestDriver.setContext }; + var driver = GraphTestDriver.driver(); + const graph: cuda.graph.Graph = .{ .d = &driver, .handle = @ptrFromInt(1) }; + GraphTestDriver.exec_drops = 0; + GraphTestDriver.graph_drops = 0; + const pending = (try gpending.startInstantiate(graph, true)).?; + var entry: Entry = .{ .key = .{ .kind = 1, .rows = 1, .scratch = 1, .scores = 2, .output = 3 }, .pending = pending, .used = 1, .cohort = 9 }; + defer gcache.drop(&entry); + const until = wire.deadline(1000); + while (!pending.done.load(.acquire)) { + if (wire.now() >= until) return error.RankTimeout; + wire.pause(); + } + const vote = gtogether.graphVote(&entry, 1000); + try std.testing.expectEqual(collective.GraphState.ready, vote.state); + try std.testing.expectEqual(@as(u64, 9), vote.cohort); + try std.testing.expect(entry.pending == null and entry.exec != null and entry.template != null); + try std.testing.expectEqual(@as(usize, 0), GraphTestDriver.graph_drops); +} + +test "exl3: a decode round captures on its second sighting and a window on its twelfth" { + const round: Key = .{ .kind = 1, .rows = 1, .scratch = 1, .scores = 2, .output = 3 }; + const window: Key = .{ .kind = 2, .rows = 3, .scratch = 1, .scores = 2, .output = 3 }; + try std.testing.expectEqual(@as(u32, 2), gcache.afterKey(round)); + try std.testing.expectEqual(@as(u32, 12), gcache.afterKey(window)); + gcache.capture_after_override = 1; + defer gcache.capture_after_override = null; + try std.testing.expectEqual(@as(u32, 1), gcache.afterKey(window)); +} + +test "exl3: engine retirement drops only that owner's graphs and prevents cross-owner updates" { + gcache.clear(); + defer gcache.clear(); + const Engine = @import("../../cuda_prompt.zig").Engine; + var first: Engine = undefined; + var second: Engine = undefined; + var driver = GraphTestDriver.driver(); + GraphTestDriver.exec_drops = 0; + const a: Key = .{ .owner = @intFromPtr(&first), .kind = 1, .rows = 1, .scratch = 1, .scores = 2, .output = 3 }; + var b = a; + b.owner = @intFromPtr(&second); + try std.testing.expect(!gkey.sameShape(a, b)); + gcache.entries[0] = .{ .key = a, .exec = .{ .d = &driver, .handle = @ptrFromInt(1) }, .used = 1 }; + gcache.entries[1] = .{ .key = b, .exec = .{ .d = &driver, .handle = @ptrFromInt(2) }, .used = 2 }; + gcache.seen[0] = .{ .key = a, .n = 2, .used = 1 }; + gcache.seen[1] = .{ .key = b, .n = 3, .used = 2 }; + gcache.releaseFor(&first); + try std.testing.expectEqual(@as(usize, 1), GraphTestDriver.exec_drops); + try std.testing.expect(gcache.entries[0] == null and gcache.seen[0] == null); + try std.testing.expect(gcache.entries[1] != null and gcache.seen[1] != null); + gcache.releaseFor(&second); + try std.testing.expectEqual(@as(usize, 2), GraphTestDriver.exec_drops); +} diff --git a/zig/src/families/flashnext/exl3/graph/cuda_exl3_graph_fallback.zig b/zig/src/families/flashnext/exl3/graph/cuda_exl3_graph_fallback.zig new file mode 100644 index 0000000000..005a7853ba --- /dev/null +++ b/zig/src/families/flashnext/exl3/graph/cuda_exl3_graph_fallback.zig @@ -0,0 +1,140 @@ +//! A failed capture costs its round an eager run: the key backs off and retries a bounded number of times. + +const std = @import("std"); +const cuda = @import("cuda"); +const hub = @import("cuda_exl3_graph.zig"); + +/// A graph is optional when its memory admission is refused; other driver and collective errors still propagate. +pub fn admissionDenied(err: anyerror) bool { + return switch (err) { + error.GraphBudgetBusy, error.GraphBudgetExceeded, error.GraphMemoryHeadroom, error.GraphEvictionRequired => true, + else => false, + }; +} + +/// Failures a key tolerates; the one after the last stays eager until the cache is cleared. +pub const max_failures: u8 = 3; +/// Rounds the first failure waits before the key is captured again; each further failure doubles the wait. +pub const first_backoff: u64 = 32; +const slots = 64; + +pub const Failure = struct { hash: u64, failures: u8, retry_at: u64 }; +var table: [slots]?Failure = @splat(null); +var next_slot: usize = 0; + +pub fn hashOf(salt: u64, key: anytype) u64 { + var h = std.hash.Wyhash.init(salt); + std.hash.autoHashStrat(&h, key, .Deep); + return h.final(); +} + +/// True while `hash` waits out a backoff (`now`: the graph cache's clock) or has used its retries. +pub fn blocked(hash: u64, now: u64) bool { + for (table) |slot| if (slot) |f| if (f.hash == hash) return f.failures > max_failures or now < f.retry_at; + return false; +} + +/// A capture or instantiation of `hash` failed: its next attempt waits `first_backoff << failures` rounds. +pub fn record(hash: u64, now: u64) void { + hub.stats.capture_failures += 1; + for (&table) |*slot| if (slot.*) |*f| if (f.hash == hash) { + f.failures +|= 1; + f.retry_at = now +| (first_backoff << @intCast(@min(f.failures, 8))); + return; + }; + for (&table) |*slot| if (slot.* == null) { + slot.* = .{ .hash = hash, .failures = 1, .retry_at = now +| (first_backoff << 1) }; + return; + }; + table[next_slot % slots] = .{ .hash = hash, .failures = 1, .retry_at = now +| (first_backoff << 1) }; + next_slot += 1; +} + +pub fn failures(hash: u64) u8 { + for (table) |slot| if (slot) |f| if (f.hash == hash) return f.failures; + return 0; +} + +pub fn clear() void { + table = @splat(null); + next_slot = 0; +} + +/// A test's way to make the next captures fail after they ran (the graph is then discarded as a failed capture). +pub var inject_capture_failures: u32 = 0; +/// A test's way to make the next instantiations fail. +pub var inject_instantiate_failures: u32 = 0; + +pub fn take(counter: *u32) bool { + if (counter.* == 0) return false; + counter.* -= 1; + return true; +} + +/// Captures `pass.chunk(ids, pos0, null)` once; null when it could not (closed and dropped, nothing ran). +pub fn captureOnce(pass: anytype, ids: []const u32, pos0: usize) ?cuda.graph.Graph { + cuda.graph.beginCapture(pass.stream, .thread_local) catch return null; + pass.prepared = true; + hub.capturing = true; + const ran = pass.chunk(ids, pos0, null); + hub.capturing = false; + ran catch { + if (cuda.graph.endCapture(pass.stream)) |g| { + var discarded = g; + discarded.deinit(); + } else |_| {} + return null; + }; + var graph = cuda.graph.endCapture(pass.stream) catch return null; + if (take(&inject_capture_failures)) { + graph.deinit(); + return null; + } + return graph; +} + +/// One rank's capture attempt (a graph or null), agreed through `c.graphAgree`: kept only if both ranks hold one. +pub fn agreedCapture(c: anytype, signature: u64, cohort: u64, attempt: anytype) !@TypeOf(attempt.run()) { + var mine = attempt.run(); + const decision = c.graphAgree(signature ^ capture_tag, .{ .state = if (mine != null) .ready else .eager, .cohort = cohort }) catch |err| { + if (mine) |*g| g.deinit(); + return err; + }; + if (decision.choice == .replay) return mine; + if (mine) |*g| g.deinit(); + return null; +} + +const capture_tag: u64 = 0x63617074757265; + +test "exl3: a failed key backs off, retries a bounded number of times and then stays eager" { + clear(); + defer clear(); + const h: u64 = 0x1234; + try std.testing.expect(!blocked(h, 0)); + var now: u64 = 100; + var tries: u32 = 0; + // a key that always fails: every attempt waits longer, and the attempts stop + while (tries < 20) : (tries += 1) { + while (blocked(h, now) and now < 100_000) now += 1; + if (blocked(h, now)) break; + record(h, now); + } + try std.testing.expectEqual(@as(u8, max_failures + 1), failures(h)); + try std.testing.expect(blocked(h, now + 1_000_000)); + // other keys are untouched + try std.testing.expect(!blocked(h +% 1, now)); + clear(); + try std.testing.expect(!blocked(h, now)); +} + +test "exl3: retry waits double with each failure" { + clear(); + defer clear(); + record(7, 0); + try std.testing.expect(blocked(7, first_backoff * 2 - 1)); + try std.testing.expect(!blocked(7, first_backoff * 2)); + record(7, first_backoff * 2); + try std.testing.expect(blocked(7, first_backoff * 2 + first_backoff * 4 - 1)); + try std.testing.expect(!blocked(7, first_backoff * 2 + first_backoff * 4)); +} diff --git a/zig/src/families/flashnext/exl3/graph/cuda_exl3_graph_fallback_test.zig b/zig/src/families/flashnext/exl3/graph/cuda_exl3_graph_fallback_test.zig new file mode 100644 index 0000000000..2b8b2d2251 --- /dev/null +++ b/zig/src/families/flashnext/exl3/graph/cuda_exl3_graph_fallback_test.zig @@ -0,0 +1,253 @@ +//! The capture-failure fallback: a failed capture costs one eager round; two ranks decide it together. + +const std = @import("std"); +const windows = @import("builtin").os.tag == .windows; +const collective = @import("../../cuda_collective.zig"); +const wire = @import("../../cuda_rank_wire.zig"); +const gfall = @import("cuda_exl3_graph_fallback.zig"); + +const FakeGraph = struct { + var live: std.atomic.Value(i32) = .init(0); + id: u32, + + fn make(id: u32) FakeGraph { + _ = live.fetchAdd(1, .monotonic); + return .{ .id = id }; + } + + pub fn deinit(_: *FakeGraph) void { + _ = live.fetchSub(1, .monotonic); + } +}; + +/// One rank's capture attempts: `fail[i]` is whether its i-th attempt fails. +const Attempt = struct { + fail: []const bool, + calls: *usize, + + pub fn run(a: Attempt) ?FakeGraph { + defer a.calls.* += 1; + return if (a.fail[a.calls.*]) null else FakeGraph.make(@intCast(a.calls.*)); + } +}; + +/// A rank's side of the wire: `graphAgree` is the exchange `Collective.graphAgree` makes. +const Rank = struct { + link: wire.Link, + step: u64 = 0, + + pub fn graphAgree(r: *Rank, signature: u64, vote: collective.GraphVote) !collective.GraphDecision { + r.step += 1; + return collective.graphExchange(r.link, r.step, signature, vote, wire.deadline(2000)); + } +}; + +fn socketPair() ![2]std.posix.socket_t { + var pair: [2]std.posix.socket_t = undefined; + if (std.c.socketpair(std.posix.AF.UNIX, std.posix.SOCK.STREAM, 0, &pair) != 0) return error.SocketPair; + errdefer for (pair) |fd| { + _ = std.c.close(fd); + }; + const nonblock: c_int = @bitCast(@as(u32, @bitCast(std.c.O{ .NONBLOCK = true }))); + for (pair) |fd| { + const flags = std.c.fcntl(fd, std.c.F.GETFL); + if (flags < 0 or std.c.fcntl(fd, std.c.F.SETFL, flags | nonblock) < 0) return error.Nonblocking; + } + return pair; +} + +const rounds = 4; + +/// What a rank saw of each round: whether it kept a graph. +const Outcome = struct { kept: [rounds]bool = @splat(false), err: ?anyerror = null }; + +fn rankMain(link: wire.Link, fail: []const bool, out: *Outcome) void { + var rank: Rank = .{ .link = link }; + var calls: usize = 0; + for (0..rounds) |round| { + const got = gfall.agreedCapture(&rank, 0x5eed + round, 7, Attempt{ .fail = fail, .calls = &calls }) catch |err| { + out.err = err; + return; + }; + out.kept[round] = got != null; + if (got) |g| { + var graph = g; + graph.deinit(); // the round's replay is done with it + } + } +} + +fn runPair(fail0: []const bool, fail1: []const bool) !struct { Outcome, Outcome } { + const pair = try socketPair(); + defer for (pair) |fd| { + _ = std.c.close(fd); + }; + var peer: Outcome = .{}; + var mine: Outcome = .{}; + const thread = try std.Thread.spawn(.{}, rankMain, .{ wire.Link{ .fd = pair[1] }, fail1, &peer }); + rankMain(.{ .fd = pair[0] }, fail0, &mine); + thread.join(); + return .{ mine, peer }; +} + +test "exl3: both ranks keep a capture only when both captured" { + if (windows) return error.SkipZigTest; // The rank wire explicitly refuses Windows. + FakeGraph.live.store(0, .monotonic); + const none = [_]bool{ false, false, false, false }; + const second = [_]bool{ false, true, false, false }; + const cases = [_][2][]const bool{ .{ &none, &none }, .{ &second, &none }, .{ &none, &second }, .{ &second, &second } }; + for (cases) |c| { + const a, const b = try runPair(c[0], c[1]); + try std.testing.expectEqual(@as(?anyerror, null), a.err); + try std.testing.expectEqual(@as(?anyerror, null), b.err); + // lockstep: both ranks decided every round the same way + try std.testing.expectEqualSlices(bool, &a.kept, &b.kept); + // round 2 failed on at least one rank in three of the cases: both ran it eagerly; later rounds are not poisoned + const failed = c[0][1] or c[1][1]; + try std.testing.expectEqualSlices(bool, &.{ true, !failed, true, true }, &a.kept); + // no capture is left behind on either rank + try std.testing.expectEqual(@as(i32, 0), FakeGraph.live.load(.monotonic)); + } +} + +test "exl3: a rank that fails every capture leaves its peer eager every round, the communicator untouched" { + if (windows) return error.SkipZigTest; // The rank wire explicitly refuses Windows. + FakeGraph.live.store(0, .monotonic); + const all = [_]bool{ true, true, true, true }; + const none = [_]bool{ false, false, false, false }; + const a, const b = try runPair(&all, &none); + try std.testing.expectEqual(@as(?anyerror, null), a.err); + try std.testing.expectEqual(@as(?anyerror, null), b.err); + try std.testing.expectEqualSlices(bool, &.{ false, false, false, false }, &a.kept); + try std.testing.expectEqualSlices(bool, &a.kept, &b.kept); + try std.testing.expectEqual(@as(i32, 0), FakeGraph.live.load(.monotonic)); +} + +test "exl3: a lost peer during the capture agreement drops the local capture and reports the loss" { + if (windows) return error.SkipZigTest; // The rank wire explicitly refuses Windows. + FakeGraph.live.store(0, .monotonic); + const pair = try socketPair(); + defer for (pair) |fd| { + _ = std.c.close(fd); + }; + _ = std.c.shutdown(pair[1], std.posix.SHUT.RDWR); + var rank: Rank = .{ .link = .{ .fd = pair[0] } }; + var calls: usize = 0; + const none = [_]bool{false}; + try std.testing.expectError(error.RankDisconnected, gfall.agreedCapture(&rank, 1, 7, Attempt{ .fail = &none, .calls = &calls })); + try std.testing.expectEqual(@as(i32, 0), FakeGraph.live.load(.monotonic)); +} + +const cuda = @import("cuda"); +const memory = @import("cuda_exl3_graph_memory.zig"); +const gcache = @import("cuda_exl3_graph_cache.zig"); +const gkey = @import("cuda_exl3_graph_key.zig"); +const grun = @import("cuda_exl3_graph_run.zig"); +const TestDriver = @import("cuda_exl3_graph_test_driver.zig").GraphTestDriver; +const admission_key: gkey.Key = .{ .kind = 1, .rows = 1, .scratch = 1, .scores = 2, .output = 3 }; + +const AdmissionPass = struct { + io: std.Io = std.testing.io, + stream: cuda.Stream, + m: struct { collective: ?*collective.Collective }, + prepared: bool = false, + chunks: usize = 0, + + pub fn chunk(p: *@This(), _: []const u32, _: usize, _: ?[]u8) !void { + p.chunks += 1; + } +}; + +fn noHeadroom() !u64 { + return 0; +} + +fn queryFailure() !u64 { + return error.CudaFailed; +} + +fn admitTest() !void { + gcache.clear(); + gfall.clear(); + try memory.configure(.{ .reservation_bytes = 100, .headroom_bytes = 20, .charge = .{ .exec = 20, .template = 10, .transient = 10 } }); + gcache.seen[0] = .{ .key = admission_key, .n = 1000, .used = 0 }; +} + +test "exl3: decode graph headroom refusal runs eagerly and leaves unexpected query errors visible" { + try admitTest(); + defer { + gcache.clear(); + gfall.clear(); + } + memory.setFreeMemoryHook(noHeadroom); + defer memory.setFreeMemoryHook(null); + var driver = TestDriver.driver(); + TestDriver.launches = 0; + var pass: AdmissionPass = .{ .stream = cuda.Stream.borrowed(&driver, null), .m = .{ .collective = null } }; + try grun.run(admission_key, &pass, &.{1}, 0); + try std.testing.expectEqual(@as(usize, 1), pass.chunks); + try std.testing.expect(pass.prepared); + try std.testing.expectEqual(@as(usize, 0), TestDriver.launches); + try std.testing.expectEqual(@as(u64, 0), memory.charged()); + try grun.run(admission_key, &pass, &.{1}, 0); + try std.testing.expectEqual(@as(usize, 2), pass.chunks); + try admitTest(); + memory.setFreeMemoryHook(queryFailure); + try std.testing.expectError(error.CudaFailed, grun.run(admission_key, &pass, &.{1}, 0)); + try std.testing.expectEqual(@as(usize, 2), pass.chunks); + try std.testing.expectEqual(@as(u64, 0), memory.charged()); +} + +test "exl3: exhausted TP graph admission runs eagerly without poisoning the next collective step" { + if (windows) return error.SkipZigTest; // The rank wire explicitly refuses Windows. + try admitTest(); + defer { + gcache.clear(); + gfall.clear(); + } + memory.setFreeMemoryHook(noHeadroom); + defer memory.setFreeMemoryHook(null); + const pair = try socketPair(); + defer for (pair) |fd| { + _ = std.c.close(fd); + }; + const Peer = struct { + fn run(link: wire.Link) !void { + var c = TestDriver.control(link); + c.timeout_ms = 5000; + const signature = gkey.graphSignature(admission_key, 0); + const capture = try c.graphAgree(signature, .{ .state = .capture }); + try std.testing.expectEqual(collective.GraphChoice.capture, capture.choice); + const refused = try c.graphAgree(signature ^ 0x61646d6974, .{ .state = .ready, .cohort = capture.cohort }); + try std.testing.expectEqual(collective.GraphChoice.eager, refused.choice); + try std.testing.expectEqual(@as(u64, 0), try c.graphVictim(signature ^ 0x76696374696d, 0)); + const quiet = try c.graphAgree(signature ^ 0x6c6f63616c7175, .{ .state = .ready, .cohort = capture.cohort }); + try std.testing.expectEqual(collective.GraphChoice.replay, quiet.choice); + try std.testing.expectEqual(@as(u64, 0), try c.graphVictim(signature ^ 0x6c6f63616c646f, 0)); + const next = try c.graphAgree(signature, .{ .state = .eager }); + try std.testing.expectEqual(collective.GraphChoice.eager, next.choice); + } + fn main(link: wire.Link, err: *?anyerror) void { + run(link) catch |e| { + err.* = e; + }; + } + }; + var peer_error: ?anyerror = null; + const thread = try std.Thread.spawn(.{}, Peer.main, .{ wire.Link{ .fd = pair[1] }, &peer_error }); + var joined = false; + defer if (!joined) thread.join(); + var control = TestDriver.control(.{ .fd = pair[0] }); + control.timeout_ms = 5000; + var driver = TestDriver.driver(); + var pass: AdmissionPass = .{ .stream = cuda.Stream.borrowed(&driver, null), .m = .{ .collective = &control } }; + try grun.run(admission_key, &pass, &.{1}, 0); + try std.testing.expect(!control.failed); + try grun.run(admission_key, &pass, &.{1}, 0); + thread.join(); + joined = true; + try std.testing.expectEqual(@as(?anyerror, null), peer_error); + try std.testing.expectEqual(@as(usize, 2), pass.chunks); + try std.testing.expect(!control.failed); + try std.testing.expectEqual(@as(u64, 0), memory.charged()); +} diff --git a/zig/src/families/flashnext/exl3/graph/cuda_exl3_graph_hooks.zig b/zig/src/families/flashnext/exl3/graph/cuda_exl3_graph_hooks.zig new file mode 100644 index 0000000000..cb1811b554 --- /dev/null +++ b/zig/src/families/flashnext/exl3/graph/cuda_exl3_graph_hooks.zig @@ -0,0 +1,127 @@ +//! The engine's hooks for EXL3 graphs: a decode round and a verify window prepare their host inputs and run. + +const std = @import("std"); +const wt = @import("../../cuda_prompt_wt.zig"); +const state_binding = @import("../kv/exl3_state_binding.zig"); +const dplan = @import("../dplan/cuda_exl3_dplan.zig"); +const hub = @import("cuda_exl3_graph.zig"); +const gkey = @import("cuda_exl3_graph_key.zig"); +const grebind = @import("cuda_exl3_graph_rebind.zig"); +const grun = @import("cuda_exl3_graph_run.zig"); +const Key = gkey.Key; +const page_graph = @import("cuda_exl3_graph_pages.zig"); + +const dims = wt.dims; +const vocab = wt.vocab; +const att_rows = @import("../../cuda_batch_state.zig").att_rows; + +/// A round's host inputs (what chunk's stage and gather would write), outside capture and replay. +pub fn prepareRound(pass: anytype, ids: []const u32) !void { + const s = pass.s; + try s.ready(); + for (s.hostIds(ids.len), ids) |*out, token| out.* = @intCast(token); + const positions = s.hostPos(ids.len * 4); + @memset(positions, 0); + for (pass.lanes, 0..) |lane, i| { + positions[i * 4] = @intCast(lane.pos); + try pass.m.ngram.rows(lane.seq.history[0..], ids[i..][0..1], s.hostEmb(ids.len)[i * dims ..][0..dims], s.hostXs(ids.len)[i * (dims / 32) ..][0 .. dims / 32]); + } +} + +/// A window's host inputs: ids, one position a block of attention rows, the n-gram rows after the sequence's history. +pub fn prepareWindow(pass: anytype, ids: []const u32, pos0: usize) !void { + const s = pass.s; + const n = ids.len; + try s.ready(); + for (s.hostIds(n), ids) |*out, token| out.* = @intCast(token); + const blocks = (n + att_rows - 1) / att_rows; + const positions = s.hostPos(blocks * 4); + @memset(positions, 0); + for (0..blocks) |b| positions[4 * b] = @intCast(pos0 + b * att_rows); + try pass.m.ngram.rows(pass.seq.history[0..], ids, s.hostEmb(n), s.hostXs(n)); +} + +/// The device planner for `e`, or null when its decode steps keep the host cut. +pub fn planner(e: anytype, rows: usize) !?*dplan.DevPlan { + if (!dplan.enabled()) return null; + const dp = (try dplan.forLayers(e.gpa, e.driver, &e.exl3_dplan, e.layers[0..e.n_layers])) orelse return null; + if (rows > dp.max_rows) { + hub.stats.declined += 1; + return null; + } + return dp; +} + +/// Engine.round's hook: true when this module ran the round (captured or not); false leaves it to the plain pass. +pub fn round(e: anytype, p: anytype, seqs: anytype, ids: []const u32, scores: u64, out: u64) !bool { + for (seqs) |q| if (q.media != null) return false; + if (seqs.len > gkey.max_lanes) return false; + const dp = (try planner(e, ids.len)) orelse return false; + var staged: usize = 0; + defer for (seqs[0..staged]) |q| if (q.sel) |selection| page_graph.finish(selection); + for (seqs) |q| { + if (q.sel) |selection| try page_graph.stage(selection, p.stream, 0, wt.attn_n, q.fed, 1); + staged += 1; + } + try dp.toDevice(p.stream); + e.exl3_dplan.active = dp; + defer e.exl3_dplan.active = null; + if (!hub.enabled()) { + hub.stats.eager_dplan += 1; + try p.chunk(ids, 0, null); + return true; + } + var key: Key = .{ .owner = @intFromPtr(e), .kind = 1, .rows = ids.len, .scratch = p.s.ids, .scratch_cap = p.s.cap, .scores = scores, .output = out, .plan = dp.state.ptr }; + for (seqs, 0..) |q, i| { + key.seq[i] = Key.seqOf(q); + key.pages[i] = Key.pageOf(q.pg()); + if (q.sel) |selection| key.tiers[i] = Key.tierOf(selection, q.fed, 1); + Key.seams(q.fed + 1, &key.chunks[i], &key.score_tiles[i]); + } + const tp0 = std.Io.Clock.awake.now(p.io).toNanoseconds(); + try prepareRound(p, ids); + hub.stats.prepare_ns += @intCast(std.Io.Clock.awake.now(p.io).toNanoseconds() - tp0); + hub.stats.prepares += 1; + try grun.run(key, p, ids, 0); + return true; +} + +/// Engine.forward's hook for a window: every row's logits into hooks.logits_out (an MTP verify window, or a one-row +pub fn forward(e: anytype, p: anytype, seq: anytype, ids: []const u32) !bool { + const hooks = p.hooks; + if (hooks.logits != .all or hooks.logits_out == 0 or seq.media != null or ids.len == 0 or ids.len > att_rows) return false; + if (seq.fed + ids.len > seq.cap or ids.len > p.s.cap) return false; + const dp = (try planner(e, ids.len)) orelse return false; + if (seq.sel) |selection| try page_graph.stage(selection, p.stream, 0, wt.attn_n, seq.fed, ids.len); + defer if (seq.sel) |selection| page_graph.finish(selection); + try dp.toDevice(p.stream); + e.exl3_dplan.active = dp; + defer e.exl3_dplan.active = null; + const n = ids.len; + if (!hub.enabled()) { + hub.stats.eager_dplan += 1; + try p.chunk(ids, seq.fed, null); + } else { + var key: Key = .{ .owner = @intFromPtr(e), .kind = 2, .rows = n, .scratch = p.s.ids, .scratch_cap = p.s.cap, .scores = p.seq.scores, .output = hooks.logits_out, .streams = hooks.streams, .plan = dp.state.ptr }; + key.seq[0] = Key.seqOf(seq); + key.pages[0] = Key.pageOf(seq.pg()); + if (seq.sel) |selection| key.tiers[0] = Key.tierOf(selection, seq.fed, n); + Key.seams(seq.fed + n, &key.chunks[0], &key.score_tiles[0]); + if (hooks.keep) |kp| { + key.keep = std.hash.Wyhash.hash(1, std.mem.asBytes(kp)); + if (grebind.stateRebindEnabled()) key.keep_nonstate = state_binding.keepHash(kp.*); + } + const tp0 = std.Io.Clock.awake.now(p.io).toNanoseconds(); + try prepareWindow(p, ids, seq.fed); + hub.stats.prepare_ns += @intCast(std.Io.Clock.awake.now(p.io).toNanoseconds() - tp0); + hub.stats.prepares += 1; + try grun.run(key, p, ids, seq.fed); + } + if (hooks.output_tail) |tail| try tail.enqueue(tail.ptr, p.stream); + // the pass's head(.all) copies the last row's scores to the host after the step (outside the graph) + if (p.m.collective) |c| try c.finish(p.stream) else try p.stream.synchronize(); + const d = p.stream.d; + try d.check(d.api.cuMemcpyDtoH_v2(seq.logits.ptr, hooks.logits_out + (n - 1) * vocab * 2, vocab * 2), "cuMemcpyDtoH"); + seq.fed += n; + return true; +} diff --git a/zig/src/families/flashnext/exl3/graph/cuda_exl3_graph_key.zig b/zig/src/families/flashnext/exl3/graph/cuda_exl3_graph_key.zig new file mode 100644 index 0000000000..a5c1c094d6 --- /dev/null +++ b/zig/src/families/flashnext/exl3/graph/cuda_exl3_graph_key.zig @@ -0,0 +1,74 @@ +//! The identity of a captured EXL3 decode graph: its key, the shape that may be updated in place, the rendezvous hash. + +const std = @import("std"); +const wt = @import("../../cuda_prompt_wt.zig"); +const grebind = @import("cuda_exl3_graph_rebind.zig"); + +pub const max_lanes = 32; + +pub const Key = struct { + owner: u64 = 0, + /// 1: a decode round; 2: a window. + kind: u8, + rows: usize, + /// Per lane (a window: lane 0): the sequence's buffers, its attention chunks and score tiles at the step's end. + seq: [max_lanes][5]u64 = @splat(@splat(0)), + pages: [max_lanes][5]u64 = @splat(@splat(0)), + tiers: [max_lanes][5]u64 = @splat(@splat(0)), + chunks: [max_lanes]u32 = @splat(0), + score_tiles: [max_lanes]u32 = @splat(0), + /// The pass's scratch (its first buffer and its row capacity, which sets every other buffer's offset). + scratch: u64, + scratch_cap: usize = 0, + scores: u64, + output: u64, + streams: u64 = 0, + keep: u64 = 0, + keep_nonstate: u64 = 0, + /// The device planner's state (a new planner is new memory). + plan: u64 = 0, + + pub fn seams(end: usize, chunks: *u32, tiles: *u32) void { + chunks.* = @intCast((@min(end, wt.dense_keys) + 511) / 512); + tiles.* = if (end / wt.index_ratio > wt.index_top) @intCast(((end + 3) / 4 + 63) / 64) else 0; + } + + /// The DeltaNet state set is in the key: a drafting stream swaps it with its window's spare set (Draft.settle). + pub fn seqOf(q: anytype) [5]u64 { + return .{ if (q.kc.len > 0) q.kc[0].ptr else 0, if (q.conv.len > 0) q.conv[0].ptr else 0, q.tail.ptr, q.cap, if (q.state.len > 0) q.state[0].ptr else 0 }; + } + + pub fn pageOf(pg: anytype) [5]u64 { + return .{ pg.pt, @bitCast(pg.ps), @intCast(pg.ptok), @bitCast(pg.ips), pg.ipt }; + } + + pub fn tierOf(selection: anytype, position: usize, rows: usize) [5]u64 { + const sparse_start: u64 = @intFromBool(position / wt.index_ratio > wt.index_top); + const sparse_end: u64 = @intFromBool((position + rows) / wt.index_ratio > wt.index_top); + return .{ @intFromPtr(selection), selection.win0, selection.npages, selection.nslots, sparse_start | (sparse_end << 1) }; + } +}; + +fn pageShape(p: [5]u64) [5]u64 { + return .{ @intFromBool(p[0] != 0), p[1], p[2], p[3], @intFromBool(p[4] != 0) }; +} + +/// The key without the addresses a cudaGraphExecUpdate may change: equal shapes capture the same topology. +pub fn sameShape(a: Key, b: Key) bool { + for (a.pages, b.pages) |x, y| if (!std.meta.eql(pageShape(x), pageShape(y))) return false; + for (a.tiers, b.tiers) |x, y| if (!std.mem.eql(u64, x[1..], y[1..])) return false; + return a.owner == b.owner and a.kind == b.kind and a.rows == b.rows and std.mem.eql(u32, &a.chunks, &b.chunks) and + std.mem.eql(u32, &a.score_tiles, &b.score_tiles) and a.scratch_cap == b.scratch_cap and + (a.keep != 0) == (b.keep != 0) and (a.streams != 0) == (b.streams != 0); +} + +pub fn graphSignature(key: Key, pos0: usize) u64 { + var hash = std.hash.Wyhash.init(0); + std.hash.autoHash(&hash, .{ key.kind, key.rows, key.chunks, key.score_tiles, key.scratch_cap, key.keep != 0, key.streams != 0, pos0 }); + std.hash.autoHash(&hash, grebind.stateRebindEnabled()); + for (key.seq) |seq| std.hash.autoHash(&hash, seq[3]); + // Rank 0's page contains the MTP head too: strides and addresses may differ across ranks. + for (key.pages) |page| std.hash.autoHash(&hash, .{ page[0] != 0, page[2], page[4] != 0 }); + for (key.tiers) |tier| std.hash.autoHash(&hash, .{ tier[0] != 0, tier[3], tier[4] }); + return hash.final(); +} diff --git a/zig/src/families/flashnext/exl3/graph/cuda_exl3_graph_key_test.zig b/zig/src/families/flashnext/exl3/graph/cuda_exl3_graph_key_test.zig new file mode 100644 index 0000000000..8d9e6600dc --- /dev/null +++ b/zig/src/families/flashnext/exl3/graph/cuda_exl3_graph_key_test.zig @@ -0,0 +1,129 @@ +//! Host tests of the EXL3 graph key: its state-binding identity and the rendezvous signature. + +const std = @import("std"); +const state_binding = @import("../kv/exl3_state_binding.zig"); +const gkey = @import("cuda_exl3_graph_key.zig"); +const Key = gkey.Key; + +test "exl3: paged graph identities bind tables while rank agreement compares page layout" { + const Pg = @import("../../cuda_glue.zig").Pg; + var first: Key = .{ .kind = 1, .rows = 1, .scratch = 1, .scores = 2, .output = 3 }; + first.pages[0] = Key.pageOf(Pg{ .pt = 100, .ps = 4096, .ptok = 256, .ips = 1024, .ipt = 200 }); + var other = first; + other.pages[0][0] += 1000; + other.pages[0][4] += 1000; + try std.testing.expect(!std.meta.eql(first, other)); + try std.testing.expect(gkey.sameShape(first, other)); + try std.testing.expectEqual(gkey.graphSignature(first, 0), gkey.graphSignature(other, 0)); + for ([_]usize{ 1, 2, 3 }) |field| { + other = first; + other.pages[0][field] += 1; + try std.testing.expect(!gkey.sameShape(first, other)); + if (field == 2) { + try std.testing.expect(gkey.graphSignature(first, 0) != gkey.graphSignature(other, 0)); + } else { + try std.testing.expectEqual(gkey.graphSignature(first, 0), gkey.graphSignature(other, 0)); + } + } + other = first; + other.pages[0] = @splat(0); + try std.testing.expect(!gkey.sameShape(first, other)); + try std.testing.expect(gkey.graphSignature(first, 0) != gkey.graphSignature(other, 0)); +} + +test "exl3: state binding normalized identity preserves all nonstate key guards" { + const original: Key = .{ .kind = 2, .rows = 2, .scratch = 11, .scratch_cap = 16, .scores = 12, .output = 13, .keep = 14, .keep_nonstate = 15, .streams = 16, .plan = 17 }; + var flipped = original; + flipped.seq[0][4] = 100; + flipped.keep = 101; + try std.testing.expect(!std.meta.eql(original, flipped)); + try std.testing.expectEqualDeep(state_binding.normalized(original), state_binding.normalized(flipped)); + inline for (@typeInfo(Key).@"struct".field_names) |field| { + if (comptime !std.mem.eql(u8, field, "seq") and + !std.mem.eql(u8, field, "pages") and !std.mem.eql(u8, field, "tiers") and !std.mem.eql(u8, field, "chunks") and !std.mem.eql(u8, field, "score_tiles") and !std.mem.eql(u8, field, "keep")) + { + var changed = original; + @field(changed, field) += 1; + try std.testing.expect(!std.meta.eql(state_binding.normalized(original), state_binding.normalized(changed))); + } + } + for (0..gkey.max_lanes) |lane| { + for (0..5) |component| { + var changed = original; + changed.pages[lane][component] += 1; + try std.testing.expect(!std.meta.eql(state_binding.normalized(original), state_binding.normalized(changed))); + changed = original; + changed.tiers[lane][component] += 1; + try std.testing.expect(!std.meta.eql(state_binding.normalized(original), state_binding.normalized(changed))); + } + for (0..5) |component| { + if (lane == 0 and component == 4) continue; + var changed = original; + changed.seq[lane][component] += 1; + try std.testing.expect(!std.meta.eql(state_binding.normalized(original), state_binding.normalized(changed))); + } + var changed = original; + changed.chunks[lane] += 1; + try std.testing.expect(!std.meta.eql(state_binding.normalized(original), state_binding.normalized(changed))); + changed = original; + changed.score_tiles[lane] += 1; + try std.testing.expect(!std.meta.eql(state_binding.normalized(original), state_binding.normalized(changed))); + } +} + +test "exl3: rendezvous signatures compare topology and capacity without device addresses" { + var key: Key = .{ .kind = 1, .rows = 1, .scratch = 1, .scores = 2, .output = 3 }; + key.seq[0] = .{ 1, 2, 100, 4096, 3 }; + const signature = gkey.graphSignature(key, 4096); + key.seq[0][0] = 1000; + key.seq[0][2] = 1001; + key.output = 1003; + try std.testing.expectEqual(signature, gkey.graphSignature(key, 4096)); + try std.testing.expect(signature != gkey.graphSignature(key, 4097)); + key.seq[0][3] += 1; + try std.testing.expect(signature != gkey.graphSignature(key, 4096)); +} + +test "exl3: tier frame staging runs before replay and unwinds failed preparation" { + const page_graph = @import("cuda_exl3_graph_pages.zig"); + const Selection = struct { + graph_staged: bool = false, + calls: usize = 0, + fail_at: ?usize = null, + pub fn begin(s: *@This(), _: void, layer: usize, position: usize, rows: usize) !void { + try std.testing.expect(!s.graph_staged); + try std.testing.expectEqual(s.calls + 2, layer); + try std.testing.expectEqual(@as(usize, 2048), position); + try std.testing.expectEqual(@as(usize, 4), rows); + if (s.fail_at == s.calls) return error.OutOfDeviceMemory; + s.calls += 1; + } + }; + var selection: Selection = .{}; + try page_graph.stage(&selection, {}, 2, 3, 2048, 4); + try std.testing.expect(selection.graph_staged and selection.calls == 3); + try std.testing.expectError(error.NestedPageGraphStage, page_graph.stage(&selection, {}, 2, 3, 2048, 4)); + page_graph.finish(&selection); + try std.testing.expect(!selection.graph_staged); + selection = .{ .fail_at = 1 }; + try std.testing.expectError(error.OutOfDeviceMemory, page_graph.stage(&selection, {}, 2, 3, 2048, 4)); + try std.testing.expect(!selection.graph_staged); +} + +test "exl3: tier graph keys distinguish sparse split topology and mutable frame windows" { + const Selection = struct { win0: usize = 0, npages: usize = 129, nslots: usize = 576 }; + var selection: Selection = .{}; + var first: Key = .{ .kind = 2, .rows = 4, .scratch = 1, .scores = 2, .output = 3 }; + first.tiers[0] = Key.tierOf(&selection, 2048, 4); + var other = first; + other.tiers[0] = Key.tierOf(&selection, 2052, 4); + try std.testing.expect(!gkey.sameShape(first, other)); + selection.win0 = 256; + other.tiers[0] = Key.tierOf(&selection, 2048, 4); + try std.testing.expect(!std.meta.eql(first, other)); + try std.testing.expect(!gkey.sameShape(first, other)); + try std.testing.expectEqual(gkey.graphSignature(first, 2048), gkey.graphSignature(other, 2048)); + selection.nslots += 64; + other.tiers[0] = Key.tierOf(&selection, 2048, 4); + try std.testing.expect(gkey.graphSignature(first, 2048) != gkey.graphSignature(other, 2048)); +} diff --git a/zig/src/families/flashnext/exl3/graph/cuda_exl3_graph_log.zig b/zig/src/families/flashnext/exl3/graph/cuda_exl3_graph_log.zig new file mode 100644 index 0000000000..39404aeecf --- /dev/null +++ b/zig/src/families/flashnext/exl3/graph/cuda_exl3_graph_log.zig @@ -0,0 +1,28 @@ +//! Coverage counters of the EXL3 decode graphs, logged every `log_every` steps. + +const std = @import("std"); +const hub = @import("cuda_exl3_graph.zig"); +const gkey = @import("cuda_exl3_graph_key.zig"); +const Key = gkey.Key; +const Outcome = hub.Outcome; + +pub var steps_seen: u64 = 0; +const log_every: u64 = 256; + +pub fn note(key: Key, o: Outcome) void { + hub.coverage[key.kind - 1][@min(key.rows, 16)][@backingInt(o)] += 1; + steps_seen += 1; + if (steps_seen % log_every != 0) return; + std.log.info("EXL3 state rebind attempts/succeeded/failed/declined {d}/{d}/{d}/{d}", .{ hub.stats.rebind_attempts, hub.stats.rebinds, hub.stats.rebind_failures, hub.stats.rebind_declined }); + var buf: [1024]u8 = undefined; + var w: std.Io.Writer = .fixed(&buf); + for (0..2) |k| for (1..17) |r| { + const c = hub.coverage[k][r]; + if (c[0] + c[1] + c[2] + c[3] == 0) continue; + w.print(" {s}{d}:{d}/{d}/{d}/{d}", .{ if (k == 0) "round" else "win", r, c[0], c[1], c[2], c[3] }) catch break; + }; + const rs = @import("../offload/cuda_exl3_recycle.zig").stats; + const cs = @import("../model/cuda_exl3_chain.zig").stats; + std.log.info("flashnext exl3 graph coverage after {d} steps (replay/eager/instantiated/updated):{s}; recycled seqs {d}/{d} drafts {d}/{d} (reused/kept); draft chains {d}: {d} replays, {d} captures ({d} updates), {d} eager; decode prepare {d} ms/{d}, capture {d} ms; prefill graphs {d}/{d}/{d}/{d}, capture {d} ms, sightings {d}/{d}, admission misses {d}", .{ steps_seen, w.buffered(), rs.seq_reused, rs.seq_kept, rs.draft_reused, rs.draft_kept, cs.chains, cs.replays, cs.captures, cs.updates, cs.eager, hub.stats.prepare_ns / std.time.ns_per_ms, hub.stats.prepares, hub.stats.capture_ns / std.time.ns_per_ms, hub.prefill_stats.replays, hub.prefill_stats.eager, hub.prefill_stats.captures, hub.prefill_stats.updates, hub.prefill_stats.capture_ns / std.time.ns_per_ms, hub.prefill_stats.sighting_hits, hub.prefill_stats.sighting_misses, hub.prefill_stats.admission_misses }); + std.log.info("EXL3 sequence retained-buffer evictions {d}", .{rs.seq_evicted}); +} diff --git a/zig/src/families/flashnext/exl3/graph/cuda_exl3_graph_memory.zig b/zig/src/families/flashnext/exl3/graph/cuda_exl3_graph_memory.zig new file mode 100644 index 0000000000..4a235ddf4f --- /dev/null +++ b/zig/src/families/flashnext/exl3/graph/cuda_exl3_graph_memory.zig @@ -0,0 +1,398 @@ +const std = @import("std"); +const policy = @import("exl3_graph_budget.zig"); + +pub const Charge = policy.Charge; +pub const Class = policy.Class; +pub const Policy = struct { reservation_bytes: u64, headroom_bytes: u64, auxiliary_bytes: u64 = 0, charge: Charge, prefill_charge: ?Charge = null, chain_charge: ?Charge = null, two_chunk_reserved: bool = false }; +const mib: u64 = 1024 * 1024; +/// Ledger slack above the retained graphs: a ledger bound, not device memory (graphs measured 48-62 MiB at one stream). +const ledger_slack_bytes: u64 = 192 * mib; + +/// What startup keeps free of pools and lanes for the graphs: the ledger limit without its slack. +fn protectedBytes(reservation: u64) u64 { + return reservation -| ledger_slack_bytes; +} + +pub fn defaultPolicy() Policy { + return workloadPolicy(8, 48) catch unreachable; +} + +pub fn workloadPolicy(streams: u64, layers: u64) !Policy { + if (streams == 0 or layers == 0) return error.InvalidGraphPolicy; + const decode = try std.math.mul(u64, try std.math.mul(u64, streams, 2), 32 * mib); + const prefill = try std.math.mul(u64, layers, mib); + const chain = try std.math.mul(u64, try std.math.mul(u64, streams, 2), mib); + const auxiliary = try std.math.add(u64, prefill, chain); + const retained = try std.math.add(u64, decode, auxiliary); + return .{ .reservation_bytes = try std.math.add(u64, retained, ledger_slack_bytes), .headroom_bytes = 128 * mib, .auxiliary_bytes = try std.math.add(u64, auxiliary, 32 * mib), .charge = .{ .exec = 16 * mib, .template = 16 * mib, .transient = 32 * mib }, .prefill_charge = .{ .exec = mib, .template = 0, .transient = 32 * mib }, .chain_charge = .{ .exec = mib, .template = 0, .transient = 32 * mib } }; +} + +pub fn policyFromEnv() !Policy { + const stream_text = std.mem.span(std.c.getenv("TF_FLASHNEXT_EXL3_STREAMS") orelse "1"); + return policyFromEnvForStreams(try std.fmt.parseInt(u64, stream_text, 10)); +} + +fn policyFromEnvForStreams(streams: u64) !Policy { + const settings = try workloadPolicy(streams, 48); + const window = @import("../prefill/exl3_prefill_window.zig"); + const two_chunk = window.enabled(if (std.c.getenv("TF_FLASHNEXT_EXL3_PREFILL_TWO_CHUNK")) |value| std.mem.span(value) else null); + return withWindowReservation(settings, two_chunk, null); +} + +fn withWindowReservation(base: Policy, enabled: bool, limit: ?u64) !Policy { + var settings = base; + const charge = settings.prefill_charge orelse settings.charge; + const extra = if (enabled) try @import("../prefill/exl3_prefill_window.zig").graphBytes(charge.exec, charge.template, charge.transient) else 0; + const required = try std.math.add(u64, base.reservation_bytes, extra); + settings.reservation_bytes = limit orelse required; + settings.two_chunk_reserved = enabled and settings.reservation_bytes >= required; + if (settings.two_chunk_reserved) settings.auxiliary_bytes = try std.math.add(u64, base.auxiliary_bytes, extra); + return settings; +} + +test "exl3: two-chunk graph reservation includes bindings peaks and declines explicit underfunding" { + const base = try workloadPolicy(2, 48); + try std.testing.expectEqualDeep(base, try withWindowReservation(base, false, null)); + const expanded = try withWindowReservation(base, true, null); + try std.testing.expect(expanded.two_chunk_reserved); + try std.testing.expectEqual(@as(u64, 288 * mib), expanded.reservation_bytes - base.reservation_bytes); + try std.testing.expectEqual(@as(u64, 288 * mib), expanded.auxiliary_bytes - base.auxiliary_bytes); + const short = try withWindowReservation(base, true, expanded.reservation_bytes - 1); + try std.testing.expect(!short.two_chunk_reserved); + try std.testing.expectEqual(base.auxiliary_bytes, short.auxiliary_bytes); + var overflow = base; + overflow.reservation_bytes = std.math.maxInt(u64); + try std.testing.expectError(error.Overflow, withWindowReservation(overflow, true, null)); +} + +const Ledger = policy.Budget(256); +pub const Ticket = Ledger.Ticket; +var guard: std.atomic.Value(bool) = .init(false); +var ledger: ?Ledger = null; +var configured: ?Policy = null; +var free_memory_hook: ?*const fn () anyerror!u64 = null; + +fn lock() void { + while (guard.swap(true, .acquire)) std.atomic.spinLoopHint(); +} + +fn unlock() void { + guard.store(false, .release); +} + +pub fn configure(settings: Policy) !void { + lock(); + defer unlock(); + try configureLocked(settings); +} + +fn configureLocked(settings: Policy) !void { + if (ledger) |*current| if (current.charged() != 0) return error.GraphBudgetBusy; + var next = try Ledger.init(settings.reservation_bytes, settings.headroom_bytes); + if (settings.auxiliary_bytes > next.limit or try settings.charge.peak() > next.limit - settings.auxiliary_bytes) return error.InvalidGraphPolicy; + _ = try next.plan(settings.charge); + if (settings.prefill_charge) |charge| _ = try next.plan(charge); + if (settings.chain_charge) |charge| _ = try next.plan(charge); + if (ledger) |*current| next.generation = current.generation; + ledger = next; + configured = settings; +} + +pub fn startupReservation() !u64 { + return startupReservationWithStreams(null); +} + +pub fn twoChunkReservationBytes() u64 { + lock(); + defer unlock(); + return if (configured) |settings| (if (settings.two_chunk_reserved) protectedBytes(settings.reservation_bytes) else 0) else 0; +} + +pub fn declineTwoChunkReservation() !u64 { + lock(); + defer unlock(); + var settings = configured orelse return 0; + if (!settings.two_chunk_reserved) return protectedBytes(settings.reservation_bytes); + const charge = settings.prefill_charge orelse settings.charge; + const extra = try @import("../prefill/exl3_prefill_window.zig").graphBytes(charge.exec, charge.template, charge.transient); + settings.reservation_bytes = try std.math.sub(u64, settings.reservation_bytes, extra); + settings.auxiliary_bytes = try std.math.sub(u64, settings.auxiliary_bytes, extra); + settings.two_chunk_reserved = false; + try configureLocked(settings); + return protectedBytes(settings.reservation_bytes); +} + +test "exl3: declined two-chunk allocation returns the original graph reservation before capture" { + const base = try workloadPolicy(2, 48); + try configure(try withWindowReservation(base, true, null)); + try std.testing.expect(twoChunkReservationBytes() > protectedBytes(base.reservation_bytes)); + try std.testing.expectEqual(protectedBytes(base.reservation_bytes), try declineTwoChunkReservation()); + try std.testing.expectEqual(@as(u64, 0), twoChunkReservationBytes()); + try std.testing.expectEqual(protectedBytes(base.reservation_bytes), try declineTwoChunkReservation()); +} + +pub fn startupReservationForStreams(streams: u32) !u64 { + if (streams == 0) return error.InvalidGraphPolicy; + return startupReservationWithStreams(streams); +} + +fn startupReservationWithStreams(streams: ?u32) !u64 { + lock(); + defer unlock(); + if (configured == null) { + const proposed = if (streams) |count| try policyFromEnvForStreams(count) else try policyFromEnv(); + try configureLocked(proposed); + const settings = configured.?; + std.log.info("EXL3 graph budget auxiliary logical allowance: {d} MiB (heuristic, only the unmaterialized remainder is protected from TP admission)", .{settings.auxiliary_bytes / mib}); + std.log.info("EXL3 graph budget heuristic allowances (not measured exec bytes): reservation {d} MiB, headroom {d} MiB, exec/template/transient {d}/{d}/{d} MiB", .{ settings.reservation_bytes / mib, settings.headroom_bytes / mib, settings.charge.exec / mib, settings.charge.template / mib, settings.charge.transient / mib }); + std.log.info("EXL3 workload heuristic follows requested server concurrency (or TF_FLASHNEXT_EXL3_STREAMS for standalone tests): two decode variants per stream, 48 prefill layer execs and two MTP chain variants per stream; local exec allowances are heuristic, not measured GPU sizes; local templates are released after instantiation", .{}); + } + return protectedBytes(configured.?.reservation_bytes); +} + +pub fn setFreeMemoryHook(hook: ?*const fn () anyerror!u64) void { + lock(); + defer unlock(); + free_memory_hook = hook; +} + +pub fn admit(class: Class, cohort: u64, free_bytes: u64) !Ticket { + lock(); + defer unlock(); + const available = if (free_memory_hook) |hook| try hook() else free_bytes; + return admitLocked(class, cohort, available); +} + +pub fn admitWithQuery(class: Class, cohort: u64, context: anytype, comptime query: anytype) !Ticket { + lock(); + defer unlock(); + const available = if (free_memory_hook) |hook| try hook() else try query(context); + return admitLocked(class, cohort, available); +} + +fn admitLocked(class: Class, cohort: u64, available: u64) !Ticket { + const settings = configured orelse return error.GraphBudgetUnconfigured; + const current = &ledger.?; + const charge = switch (class) { + .decode => settings.charge, + .prefill => settings.prefill_charge orelse settings.charge, + .chain => settings.chain_charge orelse settings.charge, + }; + const peak = try charge.peak(); + if (class == .decode and cohort != 0) { + const unmaterialized = settings.auxiliary_bytes -| current.localCharged(); + const allowance = current.limit - unmaterialized; + if (peak > allowance or current.charged() > allowance - peak) return error.GraphBudgetBusy; + } + const needed = try std.math.add(u64, try std.math.add(u64, current.pending(), peak), settings.headroom_bytes); + if (available < needed) return error.GraphMemoryHeadroom; + return current.reserve(class, cohort, charge, 0); +} + +pub fn installed(ticket: Ticket) !void { + lock(); + defer unlock(); + const current = if (ledger) |*value| value else return error.GraphBudgetUnconfigured; + try current.installed(ticket); +} + +pub fn release(ticket: Ticket) !void { + lock(); + defer unlock(); + const current = if (ledger) |*value| value else return error.GraphBudgetUnconfigured; + const entry = try current.getForRelease(ticket); + if (entry) { + try current.touch(ticket, 0, true); + try current.beginRetire(ticket, try current.cohortOf(ticket)); + } + try current.destroyed(ticket); +} + +pub fn charged() u64 { + lock(); + defer unlock(); + return if (ledger) |*current| current.charged() else 0; +} + +test "exl3: runtime admissions charge pending overlap and reject reconfiguration" { + try configure(.{ .reservation_bytes = 100, .headroom_bytes = 20, .charge = .{ .exec = 20, .template = 10, .transient = 10 } }); + const first = try admit(.decode, 1, 100); + defer release(first) catch unreachable; + try std.testing.expectError(error.GraphMemoryHeadroom, admit(.chain, 2, 99)); + try std.testing.expectError(error.GraphBudgetBusy, configure(configured.?)); + try installed(first); + try std.testing.expectEqual(@as(u64, 30), charged()); + try std.testing.expectError(error.GraphMemoryHeadroom, admit(.prefill, 2, 59)); + const second = try admit(.prefill, 2, 60); + try std.testing.expectEqual(@as(u64, 70), charged()); + try release(second); + try std.testing.expectError(error.StaleGraphTicket, release(second)); +} + +test "exl3: startup protects the retained graph allowance without the ledger slack" { + const one = try workloadPolicy(1, 48); + try std.testing.expectEqual(@as(u64, 114 * mib), protectedBytes(one.reservation_bytes)); + try std.testing.expectEqual(@as(u64, 0), protectedBytes(100)); +} + +test "exl3: runtime policy defaults are bounded heuristics" { + const settings = defaultPolicy(); + const bounded = try Ledger.init(settings.reservation_bytes, settings.headroom_bytes); + _ = try bounded.plan(settings.charge); + try std.testing.expectEqual(@as(u64, 768 * mib), settings.reservation_bytes); + try std.testing.expectEqual(@as(u64, 128 * mib), settings.headroom_bytes); +} + +test "exl3: graph reservation scales with the requested concurrent stream count" { + const one = try workloadPolicy(1, 48); + const two = try workloadPolicy(2, 48); + try std.testing.expectEqual(@as(u64, 66 * mib), two.reservation_bytes - one.reservation_bytes); +} + +test "exl3: startup reservation configures the environment policy without a separate configure call" { + try std.testing.expectEqual(@as(u64, 0), charged()); + const previous_settings = configured; + const previous_ledger = ledger; + defer { + configured = previous_settings; + ledger = previous_ledger; + } + configured = null; + ledger = null; + const expected = try policyFromEnv(); + try std.testing.expectEqual(protectedBytes(expected.reservation_bytes), try startupReservation()); + try std.testing.expectEqual(protectedBytes(expected.reservation_bytes), try startupReservation()); +} + +test "exl3: installed physical allocations are not subtracted twice by the free query" { + try configure(.{ .reservation_bytes = 100, .headroom_bytes = 20, .charge = .{ .exec = 20, .template = 10, .transient = 10 } }); + const Query = struct { + fn free(available: u64) !u64 { + return available; + } + }; + const first = try admitWithQuery(.decode, 1, @as(u64, 100), Query.free); + defer release(first) catch unreachable; + try std.testing.expectError(error.GraphMemoryHeadroom, admitWithQuery(.chain, 0, @as(u64, 60), Query.free)); + try installed(first); + const second = try admitWithQuery(.chain, 0, @as(u64, 60), Query.free); + try release(second); +} + +test "exl3: runtime policy reconfiguration cannot reuse a released ticket generation" { + const settings: Policy = .{ .reservation_bytes = 100, .headroom_bytes = 20, .charge = .{ .exec = 20, .template = 10, .transient = 10 } }; + try configure(settings); + const old = try admit(.decode, 1, 100); + try release(old); + try configure(settings); + const current = try admit(.decode, 1, 100); + defer release(current) catch unreachable; + try std.testing.expectError(error.StaleGraphTicket, release(old)); +} + +test "exl3: simultaneous runtime admissions cannot spend the same free headroom" { + try configure(.{ .reservation_bytes = 100, .headroom_bytes = 20, .charge = .{ .exec = 20, .template = 10, .transient = 10 } }); + const Worker = struct { + fn run(result: *?Ticket) void { + result.* = admit(.chain, 0, 100) catch null; + } + }; + var results: [8]?Ticket = @splat(null); + var threads: [8]std.Thread = undefined; + for (&threads, &results) |*thread, *result| thread.* = try std.Thread.spawn(.{}, Worker.run, .{result}); + for (threads) |thread| thread.join(); + var admitted: usize = 0; + for (results) |result| if (result) |ticket| { + admitted += 1; + try release(ticket); + }; + try std.testing.expectEqual(@as(usize, 2), admitted); + try std.testing.expectEqual(@as(u64, 0), charged()); +} + +test "exl3: default TP cohorts leave a complete local admission peak" { + const settings = defaultPolicy(); + try configure(settings); + var tickets: [256]?Ticket = @splat(null); + defer for (tickets) |maybe| if (maybe) |ticket| release(ticket) catch unreachable; + var count: usize = 0; + while (count < tickets.len) : (count += 1) { + const ticket = admit(.decode, count + 1, 1024 * mib) catch |err| { + try std.testing.expect(err == error.GraphBudgetBusy); + break; + }; + tickets[count] = ticket; + try installed(ticket); + } + try std.testing.expectEqual(@as(usize, 16), count); + for ([_]Class{ .prefill, .chain }) |class| { + const local = try admit(class, 0, 1024 * mib); + try release(local); + } +} + +test "exl3: bounded prefill admission retains reusable keys across 48 layer scans" { + try configure(defaultPolicy()); + var cached: [48]?Ticket = @splat(null); + defer for (cached) |maybe| if (maybe) |ticket| release(ticket) catch unreachable; + var replays: usize = 0; + var misses: usize = 0; + var captures: usize = 0; + for (0..3) |_| for (&cached) |*entry| { + if (entry.* != null) { + replays += 1; + continue; + } + const ticket = admit(.prefill, 0, 1024 * mib) catch |err| switch (err) { + error.GraphBudgetBusy, error.GraphEvictionRequired => { + misses += 1; + continue; + }, + else => return err, + }; + entry.* = ticket; + try installed(ticket); + captures += 1; + }; + try std.testing.expectEqual(@as(usize, 48), captures); + try std.testing.expectEqual(@as(usize, 96), replays); + try std.testing.expectEqual(@as(usize, 0), misses); +} + +test "exl3: workload heuristic scales with concurrency and shape with checked arithmetic" { + try std.testing.expectEqual(@as(u64, 768 * mib), (try workloadPolicy(8, 48)).reservation_bytes); + try std.testing.expectEqual(@as(u64, 504 * mib), (try workloadPolicy(4, 48)).reservation_bytes); + try std.testing.expectEqual(@as(u64, 792 * mib), (try workloadPolicy(8, 72)).reservation_bytes); + try std.testing.expectError(error.InvalidGraphPolicy, workloadPolicy(0, 48)); + try std.testing.expectError(error.InvalidGraphPolicy, workloadPolicy(8, 0)); + try std.testing.expectError(error.Overflow, workloadPolicy(std.math.maxInt(u64), 48)); + try std.testing.expectError(error.Overflow, workloadPolicy(8, std.math.maxInt(u64))); + var invalid = defaultPolicy(); + invalid.prefill_charge = .{ .exec = std.math.maxInt(u64), .template = 0, .transient = 1 }; + try std.testing.expectError(error.Overflow, configure(invalid)); +} + +test "exl3: saturated TP cache retains all local workload variants without double protection" { + try configure(defaultPolicy()); + var tickets: [80]?Ticket = @splat(null); + defer for (tickets) |maybe| if (maybe) |ticket| release(ticket) catch unreachable; + for (0..16) |index| { + const ticket = try admit(.decode, index + 1, 1024 * mib); + tickets[index] = ticket; + try installed(ticket); + } + for (16..80) |index| { + const ticket = try admit(if (index < 64) .prefill else .chain, 0, 1024 * mib); + tickets[index] = ticket; + try installed(ticket); + } + try std.testing.expectEqual(@as(u64, 576 * mib), charged()); + try std.testing.expectError(error.GraphBudgetBusy, admit(.decode, 99, 1024 * mib)); + try release(tickets[0].?); + tickets[0] = null; + tickets[0] = try admit(.decode, 99, 1024 * mib); + try installed(tickets[0].?); + try std.testing.expectEqual(@as(u64, 576 * mib), charged()); + const local = try admit(.chain, 0, 1024 * mib); + try release(local); +} diff --git a/zig/src/families/flashnext/exl3/graph/cuda_exl3_graph_pages.zig b/zig/src/families/flashnext/exl3/graph/cuda_exl3_graph_pages.zig new file mode 100644 index 0000000000..01b491661d --- /dev/null +++ b/zig/src/families/flashnext/exl3/graph/cuda_exl3_graph_pages.zig @@ -0,0 +1,10 @@ +//! Host-tier frame ownership advances before capture or replay; kernels consume the resulting device tables. +pub fn stage(selection: anytype, stream: anytype, first: usize, count: usize, position: usize, rows: usize) !void { + if (selection.graph_staged) return error.NestedPageGraphStage; + for (first..first + count) |layer| try selection.begin(stream, layer, position, rows); + selection.graph_staged = true; +} + +pub fn finish(selection: anytype) void { + selection.graph_staged = false; +} diff --git a/zig/src/families/flashnext/exl3/graph/cuda_exl3_graph_pending.zig b/zig/src/families/flashnext/exl3/graph/cuda_exl3_graph_pending.zig new file mode 100644 index 0000000000..aac70ba2b2 --- /dev/null +++ b/zig/src/families/flashnext/exl3/graph/cuda_exl3_graph_pending.zig @@ -0,0 +1,75 @@ +//! A captured EXL3 graph instantiated on a worker thread while its step runs eagerly. + +const std = @import("std"); +const platform = @import("sys"); +const cuda = @import("cuda"); +const wire = @import("../../cuda_rank_wire.zig"); + +/// A captured graph being instantiated on a worker thread (20-35 ms that the step need not wait for). +pub const Pending = struct { + thread: std.Thread = undefined, + done: std.atomic.Value(bool) = .init(false), + graph: cuda.graph.Graph, + exec: ?cuda.graph.Exec = null, + ctx: ?*anyopaque, + keep_template: bool, + started_ns: u64 = 0, + + pub fn main(pd: *Pending) void { + const api = ctxApi() orelse { + pd.done.store(true, .release); + return; + }; + if (api.set(pd.ctx) != 0) { + pd.done.store(true, .release); + return; + } + pd.exec = pd.graph.instantiate() catch null; + pd.done.store(true, .release); + } +}; + +const CtxApi = struct { get: *const fn (*?*anyopaque) callconv(.c) c_int, set: *const fn (?*anyopaque) callconv(.c) c_int }; +pub var ctx_api: ?CtxApi = null; +var ctx_lib: ?platform.DynLib = null; +pub fn ctxApi() ?CtxApi { + if (ctx_api) |a| return a; + var lib = platform.DynLib.open(platform.cuda_driver) catch return null; + const a: CtxApi = .{ + .get = lib.lookup(*const fn (*?*anyopaque) callconv(.c) c_int, "cuCtxGetCurrent") orelse return null, + .set = lib.lookup(*const fn (?*anyopaque) callconv(.c) c_int, "cuCtxSetCurrent") orelse return null, + }; + ctx_lib = lib; + ctx_api = a; + return a; +} + +/// TF_FLASHNEXT_EXL3_GRAPH_ASYNC=0 instantiates on the step's thread. +pub fn asyncInstantiate() bool { + const v = std.c.getenv("TF_FLASHNEXT_EXL3_GRAPH_ASYNC") orelse return true; + return !std.mem.eql(u8, std.mem.span(v), "0"); +} + +/// Starts instantiating `g` on a worker thread (it owns `g` from here); null when that is not possible. +pub fn startInstantiate(g: cuda.graph.Graph, keep_template: bool) !?*Pending { + if (!asyncInstantiate()) return null; + const api = ctxApi() orelse return null; + var ctx: ?*anyopaque = null; + if (api.get(&ctx) != 0) return null; + const pd = try std.heap.c_allocator.create(Pending); + pd.* = .{ .graph = g, .ctx = ctx, .keep_template = keep_template, .started_ns = wire.now() }; + pd.thread = std.Thread.spawn(.{}, Pending.main, .{pd}) catch |err| { + std.heap.c_allocator.destroy(pd); + return err; + }; + return pd; +} + +/// The finished instantiation's exec (the graph template freed), or null while it runs; error when it failed. +pub fn takePending(pd: *Pending) !?cuda.graph.Exec { + if (!pd.done.load(.acquire)) return null; + pd.thread.join(); + defer std.heap.c_allocator.destroy(pd); + pd.graph.deinit(); + return pd.exec orelse error.CudaFailed; +} diff --git a/zig/src/families/flashnext/exl3/graph/cuda_exl3_graph_prefill.zig b/zig/src/families/flashnext/exl3/graph/cuda_exl3_graph_prefill.zig new file mode 100644 index 0000000000..1ba30d6828 --- /dev/null +++ b/zig/src/families/flashnext/exl3/graph/cuda_exl3_graph_prefill.zig @@ -0,0 +1,286 @@ +//! Prefill graphs of the EXL3 engine: one layer's sub-blocks of a deep chunk captured and replayed. + +const std = @import("std"); +const cuda = @import("cuda"); +const wt = @import("../../cuda_prompt_wt.zig"); +const memory = @import("cuda_exl3_graph_memory.zig"); +const hub = @import("cuda_exl3_graph.zig"); +const gcache = @import("cuda_exl3_graph_cache.zig"); +const gfall = @import("cuda_exl3_graph_fallback.zig"); +const gbudget = @import("cuda_exl3_graph_budget.zig"); + +// ---- prefill: one layer's sub-blocks of a deep chunk. + +pub const Prefill = enum { eager, replayed, capturing }; + +pub const PBuffer = struct { + driver: *const cuda.Driver, + address: u64, + bytes: usize, + + pub fn of(buffer: cuda.DeviceBuffer) PBuffer { + return .{ .driver = buffer.d, .address = buffer.ptr, .bytes = buffer.len }; + } +}; + +pub const PGdn = struct { + conv: PBuffer, + state: PBuffer, + + pub fn of(sequence: anytype, layer: usize) PGdn { + const index = wt.linearIndex(layer); + return .{ .conv = PBuffer.of(sequence.conv[index]), .state = PBuffer.of(sequence.state[index]) }; + } +}; + +pub const PKey = struct { + owner: u64 = 0, + layer: usize, + rows: usize, + sub: usize, + /// The chunk's first position (attention layers only; 0 for DeltaNet layers). + pos0: usize, + deep: u64, + base: u64, + base_cap: usize, + seq: [5]u64, + pages: [5]u64 = @splat(0), + tier: [5]u64 = @splat(0), + wait_count: usize, + wait_key: u64, + scores: u64, + tiles: bool, + positions: u64 = 0, + stream: u64 = 0, + gdn: ?PGdn = null, +}; + +pub const PEntry = struct { key: PKey, exec: cuda.graph.Exec, used: u64, admission: ?memory.Ticket = null, stream: ?cuda.Stream = null }; + +pub fn dropPrefill(entry: *PEntry) void { + if (entry.stream) |stream| stream.synchronize() catch @panic("prefill retirement synchronization failed"); + entry.exec.deinit(); + if (entry.admission) |ticket| memory.release(ticket) catch unreachable; +} +pub const pcache_n = 112; +pub var pentries: [pcache_n]?PEntry = @splat(null); +pub var pseen: [pcache_n * 2]?struct { key: PKey, n: u32 } = @splat(null); +pub var pseen_at: usize = 0; +pub var pcur: ?PKey = null; +pub var pcur_admission: ?memory.Ticket = null; + +pub fn prefillAdmissionMiss() Prefill { + hub.prefill_stats.eager += 1; + hub.prefill_stats.admission_misses += 1; + const misses = hub.prefill_stats.admission_misses; + if (misses == 1 or misses % 256 == 0) std.log.info("flashnext exl3 prefill graph budget admission declines: {d}; retained cache preserved, layer runs eagerly", .{misses}); + return .eager; +} + +/// Prefill layers are captured whenever graphs are on (TF_FLASHNEXT_EXL3_GRAPH=1). +pub fn prefillEnabled() bool { + return hub.enabled(); +} + +pub fn pshape(k: PKey) PKey { + var s = k; + s.pos0 = 0; + return s; +} + +pub fn psighting(k: PKey) u32 { + const s = pshape(k); + for (&pseen) |*e| if (e.*) |*x| if (std.meta.eql(x.key, s)) { + x.n += 1; + hub.prefill_stats.sighting_hits += 1; + return x.n; + }; + hub.prefill_stats.sighting_misses += 1; + pseen[pseen_at % pseen.len] = .{ .key = s, .n = 1 }; + pseen_at += 1; + return 1; +} + +/// The sighting of a layer shape that captures it (a chunk takes every layer's shape twice). +pub fn prefillAfter() u32 { + return gcache.after(); +} + +pub fn clearPrefill() void { + for (&pentries) |*e| if (e.*) |*x| { + dropPrefill(x); + e.* = null; + }; + pseen = @splat(null); +} + +pub fn clearRuntime(owner: u64) void { + for (&pentries) |*entry| if (entry.*) |*held| if (held.key.owner == owner) { + dropPrefill(held); + entry.* = null; + }; + for (&pseen) |*entry| if (entry.*) |held| if (held.key.owner == owner) { + entry.* = null; + }; +} + +/// Before a layer's sub-blocks: +pub fn prefillBegin(p: anytype, linear: bool, layer: usize, rows: usize, sub: usize, pos0: usize, deep: u64, base: anytype, tiles: bool, waits: []const cuda.Event, positions: u64) !Prefill { + if (!prefillEnabled()) return .eager; + if (!linear) if (p.seq.sel) |selection| if (!selection.graph_staged) return .eager; + const q = p.seq; + var wait_key: u64 = 0; + for (waits) |ev| { + const handle = if (ev.handle) |h| @intFromPtr(h) else 0; + wait_key = std.hash.Wyhash.hash(wait_key, std.mem.asBytes(&handle)); + } + const key: PKey = .{ + .owner = if (p.m.prefill) |runtime| @intFromPtr(runtime) else 0, + .layer = layer, + .rows = rows, + .sub = sub, + .pos0 = if (linear) 0 else pos0, + .deep = deep, + .base = base.ids, + .base_cap = base.cap, + .seq = .{ if (q.conv.len > 0) q.conv[0].ptr else 0, if (q.kc.len > 0) q.kc[0].ptr else 0, q.tail.ptr, q.cap, if (q.state.len > 0) q.state[0].ptr else 0 }, + .pages = @import("cuda_exl3_graph_key.zig").Key.pageOf(q.pg), + .tier = if (q.sel) |selection| @import("cuda_exl3_graph_key.zig").Key.tierOf(selection, pos0, rows) else @splat(0), + .wait_count = waits.len, + .wait_key = wait_key, + .scores = q.scores, + .tiles = tiles, + .positions = positions, + .stream = @intFromPtr(p.stream.handle), + .gdn = if (linear) PGdn.of(q, layer) else null, + }; + gcache.clock += 1; + for (&pentries) |*e| if (e.*) |*x| if (std.meta.eql(x.key, key)) { + x.used = gcache.clock; + hub.prefill_stats.replays += 1; + try x.exec.launchOn(p.stream); + return .replayed; + }; + if (psighting(key) < prefillAfter() or gfall.blocked(gfall.hashOf(2, pshape(key)), gcache.clock)) { + hub.prefill_stats.eager += 1; + return .eager; + } + var free_slot = false; + for (pentries) |entry| free_slot = free_slot or entry == null; + if (!free_slot) return prefillAdmissionMiss(); + _ = try memory.startupReservation(); + const admission = memory.admitWithQuery(.prefill, 0, p.stream, gbudget.freeMemory) catch |err| switch (err) { + error.GraphBudgetBusy, error.GraphBudgetExceeded, error.GraphMemoryHeadroom, error.GraphEvictionRequired => return prefillAdmissionMiss(), + else => return err, + }; + errdefer memory.release(admission) catch unreachable; + try p.stream.synchronize(); + try cuda.graph.beginCapture(p.stream, .thread_local); + hub.capturing = true; + pcur = key; + pcur_admission = admission; + return .capturing; +} + +/// After a layer's sub-blocks ran under `.capturing`: true once the layer ran from its graph, false: run it eagerly. +pub fn prefillEnd(p: anytype, how: Prefill) !bool { + if (how != .capturing) return true; + const t0 = std.Io.Clock.awake.now(p.io).toNanoseconds(); + hub.capturing = false; + const key = pcur.?; + pcur = null; + const admission = pcur_admission.?; + pcur_admission = null; + var admitted = true; + defer if (admitted) memory.release(admission) catch unreachable; + var graph = cuda.graph.endCapture(p.stream) catch return unusable(key); + defer graph.deinit(); + if (gfall.take(&gfall.inject_capture_failures)) return unusable(key); + const at: *PEntry = blk: { + for (&pentries) |*e| if (e.*) |*x| if (std.meta.eql(pshape(x.key), pshape(key))) { + if ((try x.exec.update(graph)) == .success) { + try gbudget.uploadResident(&x.exec, p.stream); + x.key = key; + x.used = gcache.clock; + hub.prefill_stats.updates += 1; + break :blk x; + } + }; + var best: usize = 0; + var best_used: u64 = std.math.maxInt(u64); + var same_shape: ?usize = null; + for (pentries, 0..) |entry, index| if (entry) |cached| { + if (std.meta.eql(pshape(cached.key), pshape(key))) { + same_shape = index; + break; + } + }; + for (&pentries, 0..) |*e, i| { + if (e.* == null) { + best = i; + break; + } + if (e.*.?.used < best_used) { + best_used = e.*.?.used; + best = i; + } + } + if (same_shape) |index| best = index; + if (pentries[best]) |*old| { + dropPrefill(old); + pentries[best] = null; + } + var exec = (if (gfall.take(&gfall.inject_instantiate_failures)) error.CudaFailed else graph.instantiate()) catch return unusable(key); + errdefer exec.deinit(); + try gbudget.uploadResident(&exec, p.stream); + try memory.installed(admission); + pentries[best] = .{ .key = key, .exec = exec, .used = gcache.clock, .admission = admission, .stream = p.stream }; + admitted = false; + break :blk &pentries[best].?; + }; + hub.prefill_stats.captures += 1; + hub.prefill_stats.capture_ns += @intCast(std.Io.Clock.awake.now(p.io).toNanoseconds() - t0); + try at.exec.launchOn(p.stream); + return true; +} + +fn unusable(key: PKey) bool { + gfall.record(gfall.hashOf(2, pshape(key)), gcache.clock); + hub.prefill_stats.eager += 1; + return false; +} + +/// Runs one layer's sub-blocks (`blocks.run(mode)`); a layer that cannot be captured runs eagerly, its shape backs off. +pub fn runLayer(p: anytype, how: Prefill, blocks: anytype) !void { + var mode = how; + while (mode != .replayed) { + blocks.run(mode) catch |err| { + if (mode != .capturing) return err; + prefillAbort(p, mode); + mode = .eager; + continue; + }; + if (try prefillEnd(p, mode)) return; + mode = .eager; + } +} + +/// A capture that failed while its sub-blocks were enqueued: closed and dropped, the shape backs off. +pub fn prefillAbort(p: anytype, how: Prefill) void { + if (how != .capturing) return; + hub.capturing = false; + if (pcur) |key| { + gfall.record(gfall.hashOf(2, pshape(key)), gcache.clock); + hub.prefill_stats.eager += 1; + } + pcur = null; + defer if (pcur_admission) |ticket| { + memory.release(ticket) catch unreachable; + pcur_admission = null; + }; + if (cuda.graph.endCapture(p.stream)) |g| { + var d = g; + d.deinit(); + } else |_| {} + std.log.warn("flashnext exl3 graph: a prefill layer could not be captured; it runs eagerly and is tried again later", .{}); +} diff --git a/zig/src/families/flashnext/exl3/graph/cuda_exl3_graph_prefill_test.zig b/zig/src/families/flashnext/exl3/graph/cuda_exl3_graph_prefill_test.zig new file mode 100644 index 0000000000..ee08266c65 --- /dev/null +++ b/zig/src/families/flashnext/exl3/graph/cuda_exl3_graph_prefill_test.zig @@ -0,0 +1,137 @@ +//! Host tests of the EXL3 prefill graph keys, retention and sighting history. + +const std = @import("std"); +const cuda = @import("cuda"); +const wt = @import("../../cuda_prompt_wt.zig"); +const memory = @import("cuda_exl3_graph_memory.zig"); +const gprefill = @import("cuda_exl3_graph_prefill.zig"); +const gdriver = @import("cuda_exl3_graph_test_driver.zig"); +const PKey = gprefill.PKey; +const PGdn = gprefill.PGdn; +const GraphTestDriver = gdriver.GraphTestDriver; + +test "exl3: GDN prefill keys bind the actual layer buffers including driver and extent" { + var driver: cuda.Driver = undefined; + var other_driver: cuda.Driver = undefined; + var conv = [_]cuda.DeviceBuffer{ + .{ .d = &driver, .ptr = 100, .len = 16 }, + .{ .d = &driver, .ptr = 200, .len = 16 }, + .{ .d = &driver, .ptr = 300, .len = 16 }, + }; + var state = conv; + for (&state) |*buffer| buffer.ptr += 1000; + const sequence = .{ .conv = &conv, .state = &state }; + var key = std.mem.zeroes(PKey); + key.layer = 2; + key.seq = .{ conv[0].ptr, 0, 500, 4096, state[0].ptr }; + key.gdn = PGdn.of(sequence, key.layer); + const original = key.gdn.?; + try std.testing.expectEqual(@as(u64, 300), original.conv.address); + try std.testing.expectEqual(@as(u64, 1300), original.state.address); + for ([_]bool{ false, true }) |recurrent| { + const buffer = if (recurrent) &state[2] else &conv[2]; + const saved = buffer.*; + defer buffer.* = saved; + for (0..3) |field| { + buffer.* = saved; + switch (field) { + 0 => buffer.ptr += 1, + 1 => buffer.len += 1, + 2 => buffer.d = &other_driver, + else => unreachable, + } + var changed = key; + changed.gdn = PGdn.of(sequence, key.layer); + try std.testing.expectEqual(key.seq, changed.seq); + try std.testing.expect(!std.meta.eql(gprefill.pshape(key), gprefill.pshape(changed))); + } + } + try std.testing.expectEqualDeep(original, PGdn.of(sequence, key.layer)); +} + +test "exl3: prefill retained graphs release their graph budget ticket on retirement" { + gprefill.clearPrefill(); + try memory.configure(.{ .reservation_bytes = 100, .headroom_bytes = 0, .charge = .{ .exec = 10, .template = 10, .transient = 10 }, .prefill_charge = .{ .exec = 5, .template = 0, .transient = 10 } }); + const exec_ticket = try memory.admit(.prefill, 0, 100); + try memory.installed(exec_ticket); + try std.testing.expectEqual(@as(u64, 5), memory.charged()); + var driver = GraphTestDriver.driver(); + GraphTestDriver.exec_drops = 0; + gprefill.pentries[0] = .{ .key = std.mem.zeroes(PKey), .exec = .{ .d = &driver, .handle = @ptrFromInt(1) }, .used = 1, .admission = exec_ticket, .stream = .{ .d = &driver, .handle = @ptrFromInt(1) } }; + gprefill.clearPrefill(); + try std.testing.expectEqual(@as(usize, 1), GraphTestDriver.exec_drops); + try std.testing.expectEqual(@as(u64, 0), memory.charged()); + try std.testing.expectError(error.StaleGraphTicket, memory.release(exec_ticket)); +} + +test "prefill graph update candidates tolerate a dynamic position only" { + const base: PKey = .{ .layer = 7, .rows = 4096, .sub = 256, .pos0 = 4096, .deep = 1, .base = 2, .base_cap = 4096, .seq = .{ 3, 4, 5, 4096, 6 }, .wait_count = 0, .wait_key = 0, .scores = 13, .tiles = true }; + var moved = base; + moved.pos0 = 8192; + try std.testing.expect(!std.meta.eql(base, moved)); + try std.testing.expect(std.meta.eql(gprefill.pshape(base), gprefill.pshape(moved))); + + moved.seq[0] += 1; + try std.testing.expect(!std.meta.eql(gprefill.pshape(base), gprefill.pshape(moved))); + moved = base; + moved.rows += 1; + try std.testing.expect(!std.meta.eql(gprefill.pshape(base), gprefill.pshape(moved))); + moved = base; + moved.wait_key += 1; + try std.testing.expect(!std.meta.eql(gprefill.pshape(base), gprefill.pshape(moved))); +} + +test "exl3: prefill warmup history retains two lanes and two chunk banks" { + gprefill.pseen = @splat(null); + gprefill.pseen_at = 0; + defer { + gprefill.pseen = @splat(null); + gprefill.pseen_at = 0; + } + var keys: [2 * 2 * wt.layers_n]PKey = undefined; + var count: usize = 0; + for (0..2) |sequence| for (0..2) |bank| for (0..wt.layers_n) |layer| { + keys[count] = .{ + .layer = layer, + .rows = 4096, + .sub = 256, + .pos0 = bank * 4096, + .deep = 100 + bank, + .base = 200, + .base_cap = 8192, + .seq = .{ 1, 2, 300 + sequence, 8192, 4 }, + .wait_count = 0, + .wait_key = 0, + .scores = 10, + .tiles = true, + .positions = 11, + .stream = 12, + }; + count += 1; + }; + try std.testing.expectEqual(@as(usize, 192), count); + for (keys) |key| try std.testing.expectEqual(@as(u32, 1), gprefill.psighting(key)); + for (keys) |key| try std.testing.expectEqual(@as(u32, 2), gprefill.psighting(key)); +} + +test "exl3: retiring one prefill runtime preserves another runtime's graph and sightings" { + gprefill.clearPrefill(); + defer gprefill.clearPrefill(); + var driver = GraphTestDriver.driver(); + GraphTestDriver.exec_drops = 0; + var first = std.mem.zeroes(PKey); + first.owner = 1; + var second = first; + second.owner = 2; + gprefill.pentries[0] = .{ .key = first, .exec = .{ .d = &driver, .handle = @ptrFromInt(1) }, .used = 1 }; + gprefill.pentries[1] = .{ .key = second, .exec = .{ .d = &driver, .handle = @ptrFromInt(2) }, .used = 2 }; + gprefill.pseen[0] = .{ .key = first, .n = 3 }; + gprefill.pseen[1] = .{ .key = second, .n = 4 }; + gprefill.clearRuntime(1); + try std.testing.expectEqual(@as(usize, 1), GraphTestDriver.exec_drops); + try std.testing.expect(gprefill.pentries[0] == null and gprefill.pseen[0] == null); + try std.testing.expect(gprefill.pentries[1] != null); + try std.testing.expectEqual(@as(u32, 4), gprefill.pseen[1].?.n); + gprefill.clearRuntime(1); + try std.testing.expectEqual(@as(usize, 1), GraphTestDriver.exec_drops); +} diff --git a/zig/src/families/flashnext/exl3/graph/cuda_exl3_graph_rebind.zig b/zig/src/families/flashnext/exl3/graph/cuda_exl3_graph_rebind.zig new file mode 100644 index 0000000000..84ff923be2 --- /dev/null +++ b/zig/src/families/flashnext/exl3/graph/cuda_exl3_graph_rebind.zig @@ -0,0 +1,148 @@ +//! Rebinding a cached window graph to another sequence's DeltaNet state (off unless a test turns `state_rebind` on). + +const std = @import("std"); +const cuda = @import("cuda"); +const wt = @import("../../cuda_prompt_wt.zig"); +const collective = @import("../../cuda_collective.zig"); +const state_binding = @import("../kv/exl3_state_binding.zig"); +const state_driver = @import("../kv/cuda_exl3_state_binding.zig"); +const hub = @import("cuda_exl3_graph.zig"); +const gcache = @import("cuda_exl3_graph_cache.zig"); +const glog = @import("cuda_exl3_graph_log.zig"); +const gkey = @import("cuda_exl3_graph_key.zig"); +const Key = gkey.Key; +const Entry = gcache.Entry; + +/// Not a user switch: the rebind is exercised by its tests only. +pub var state_rebind: bool = false; + +pub fn stateRebindEnabled() bool { + return state_rebind; +} + +pub fn stateRequest(key: Key, pass: anytype) ?state_binding.Request { + if (!stateRebindEnabled()) return null; + return windowStateRequest(key, pass); +} + +pub fn windowStateRequest(key: Key, pass: anytype) ?state_binding.Request { + const Pass = @typeInfo(@TypeOf(pass)).pointer.child; + if (comptime !@hasField(Pass, "lanes")) return null; + if (key.kind != 2 or !state_binding.supportsWindow(key.rows) or pass.lanes.len != 0) return null; + const gdn = @import("../../cuda_gdn.zig"); + var request: state_binding.Request = .{ + .width = key.rows, + .chain = @intFromPtr(pass.k.chain.handle), + .chain_to = @intFromPtr(pass.k.chain_to.handle), + .geometry = .{ .grid = .{ gdn.nv, gdn.dv / 16, 1 }, .block = .{ 128, 1, 1 } }, + .nonstate = .{ .gq = pass.s.gq, .gk = pass.s.gk, .gv = pass.s.gv, .ggb = pass.s.ggb, .rows = @intCast(key.rows), .ys = pass.s.ys }, + }; + for (pass.m.layers, 0..) |layer, index| { + if (!layer.linear) continue; + const linear = wt.linearIndex(index); + if (request.count == state_binding.max_layers or linear >= pass.seq.state.len) return null; + request.layers[request.count] = .{ + .layer = index, + .state_in = pass.seq.state[linear].ptr, + .state_out = if (pass.hooks.keep) |keep| keep.layers[index].state_out else 0, + }; + if (request.layers[request.count].state_out == 0) return null; + request.count += 1; + } + request.validate() catch return null; + return request; +} + +pub fn capturedBinding(key: Key, pass: anytype, graph: cuda.graph.Graph) ?state_binding.Binding { + const request = stateRequest(key, pass) orelse return null; + if (state_driver.nodeType(graph.d) == null) return null; + return state_binding.Binding.init(@intFromPtr(graph.handle), request) catch null; +} + +pub fn rebindCandidate(key: Key, request: state_binding.Request, tp: bool) ?*Entry { + var candidate: ?*Entry = null; + for (&gcache.entries) |*maybe| if (maybe.*) |*entry| { + if (entry.poisoned or entry.pending != null or entry.exec == null or entry.template == null or (tp and entry.cohort == 0)) continue; + const binding = entry.binding orelse continue; + if (binding.poisoned or !std.meta.eql(state_binding.normalized(entry.key), state_binding.normalized(key)) or !binding.request.compatible(request)) continue; + if (candidate == null or entry.cohort < candidate.?.cohort or (entry.cohort == candidate.?.cohort and entry.used > candidate.?.used)) candidate = entry; + }; + return candidate; +} + +pub fn prepareBinding(entry: *Entry) !?state_driver.Backend { + const backend = state_driver.Backend.init(entry.template.?, entry.exec.?) orelse return null; + entry.binding.?.load(backend) catch |err| switch (err) { + error.UnsupportedBinding => { + entry.binding.?.invalidate(); + entry.binding = null; + return null; + }, + else => return err, + }; + return backend; +} + +pub fn tryRebind(key: Key, pass: anytype) !bool { + const request = stateRequest(key, pass) orelse return false; + const entry = rebindCandidate(key, request, false) orelse { + hub.stats.rebind_declined += 1; + return false; + }; + const prepared = prepareBinding(entry) catch |err| { + hub.stats.rebind_failures += 1; + return err; + }; + const backend = prepared orelse { + hub.stats.rebind_declined += 1; + return false; + }; + hub.stats.rebind_attempts += 1; + entry.binding.?.apply(backend, request) catch |err| { + entry.poisoned = true; + hub.stats.rebind_failures += 1; + return err; + }; + entry.key = key; + entry.used = gcache.clock; + hub.stats.rebinds += 1; + hub.stats.replays += 1; + glog.note(key, .replay); + try entry.exec.?.launchOn(pass.stream); + return true; +} + +pub fn tryRebindTogether(c: *collective.Collective, signature: u64, key: Key, pass: anytype) !bool { + if (!stateRebindEnabled() or key.kind != 2 or !state_binding.supportsWindow(key.rows)) return false; + const request = stateRequest(key, pass); + const candidate = if (request) |target| rebindCandidate(key, target, true) else null; + var backend: ?state_driver.Backend = null; + var failed = false; + if (candidate) |entry| backend = prepareBinding(entry) catch blk: { + hub.stats.rebind_failures += 1; + failed = true; + break :blk null; + }; + const selected = try c.graphAgree(signature ^ 0x737461746573656c, .{ + .state = if (failed) .failed else if (backend != null) .ready else .capture, + .cohort = if (candidate) |entry| entry.cohort else 0, + }); + if (selected.choice != .replay) { + hub.stats.rebind_declined += 1; + return false; + } + const entry = candidate.?; + hub.stats.rebind_attempts += 1; + entry.binding.?.applyTogether(backend.?, request.?, c, signature ^ 0x7374617465646f6e, selected.cohort) catch |err| { + entry.poisoned = true; + hub.stats.rebind_failures += 1; + return err; + }; + entry.key = key; + entry.used = gcache.clock; + hub.stats.rebinds += 1; + hub.stats.replays += 1; + glog.note(key, .replay); + try entry.exec.?.launchOn(pass.stream); + return true; +} diff --git a/zig/src/families/flashnext/exl3/graph/cuda_exl3_graph_rebind_test.zig b/zig/src/families/flashnext/exl3/graph/cuda_exl3_graph_rebind_test.zig new file mode 100644 index 0000000000..0df34327ec --- /dev/null +++ b/zig/src/families/flashnext/exl3/graph/cuda_exl3_graph_rebind_test.zig @@ -0,0 +1,145 @@ +//! Host tests of rebinding a cached EXL3 window graph to another sequence's state. + +const std = @import("std"); +const windows = @import("builtin").os.tag == .windows; +const wire = @import("../../cuda_rank_wire.zig"); +const state_binding = @import("../kv/exl3_state_binding.zig"); +const gcache = @import("cuda_exl3_graph_cache.zig"); +const grebind = @import("cuda_exl3_graph_rebind.zig"); +const gkey = @import("cuda_exl3_graph_key.zig"); +const gpending = @import("cuda_exl3_graph_pending.zig"); +const gdriver = @import("cuda_exl3_graph_test_driver.zig"); +const Key = gkey.Key; +const Pending = gpending.Pending; +const GraphTestDriver = gdriver.GraphTestDriver; + +test "exl3: state binding candidate excludes pending poisoned retired and incompatible generations" { + gcache.entries = @splat(null); + defer gcache.entries = @splat(null); + var request: state_binding.Request = .{ + .width = 2, + .chain = 1, + .chain_to = 2, + .geometry = .{ .grid = .{ 16, 8, 1 }, .block = .{ 128, 1, 1 } }, + .nonstate = .{ .gq = 1, .gk = 2, .gv = 3, .ggb = 4, .rows = 2, .ys = 5 }, + .count = 1, + }; + request.layers[0] = .{ .layer = 0, .state_in = 100, .state_out = 200 }; + for (2..state_binding.max_width + 1) |width| { + request.width = width; + request.nonstate.rows = @intCast(width); + const key: Key = .{ .kind = 2, .rows = width, .scratch = 1, .scores = 2, .output = 3, .keep = 4, .keep_nonstate = 5 }; + gcache.entries[0] = .{ + .key = key, + .used = 1, + .cohort = 7, + .exec = .{ .d = undefined, .handle = @ptrFromInt(1) }, + .template = .{ .d = undefined, .handle = @ptrFromInt(2) }, + .binding = try state_binding.Binding.init(2, request), + }; + var target = request; + std.mem.swap(u64, &target.layers[0].state_in, &target.layers[0].state_out); + var flipped = key; + flipped.seq[0][4] = 200; + flipped.keep = 9; + try std.testing.expect(gcache.find(flipped) == null); + try std.testing.expect(grebind.rebindCandidate(flipped, target, true) == &gcache.entries[0].?); + gcache.entries[0].?.pending = @ptrFromInt(@alignOf(Pending)); + try std.testing.expect(grebind.rebindCandidate(flipped, target, true) == null); + gcache.entries[0].?.pending = null; + gcache.entries[0].?.poisoned = true; + try std.testing.expect(grebind.rebindCandidate(flipped, target, true) == null); + gcache.entries[0].?.poisoned = false; + target.layers[0].state_in = 300; + try std.testing.expect(grebind.rebindCandidate(flipped, target, true) == null); + target = request; + flipped.output += 1; + try std.testing.expect(grebind.rebindCandidate(flipped, target, true) == null); + gcache.entries[0].?.binding.?.invalidate(); + try std.testing.expect(grebind.rebindCandidate(key, request, true) == null); + } +} + +test "exl3: state binding transaction uses real graphAgree before publication" { + if (windows) return error.SkipZigTest; // The rank wire explicitly refuses Windows. + const Fake = struct { + calls: usize = 0, + fail: bool = false, + + pub fn set(self: *@This(), _: usize, _: usize, _: state_binding.Geometry, _: state_binding.Parameters) !void { + self.calls += 1; + if (self.fail and self.calls == 2) return error.SetterFailed; + } + }; + var request: state_binding.Request = .{ + .width = 2, + .chain = 10, + .chain_to = 11, + .geometry = .{ .grid = .{ 16, 8, 1 }, .block = .{ 128, 1, 1 } }, + .nonstate = .{ .gq = 1, .gk = 2, .gv = 3, .ggb = 4, .rows = 2, .ys = 5 }, + .count = 2, + }; + request.layers[0] = .{ .layer = 0, .state_in = 100, .state_out = 200 }; + request.layers[1] = .{ .layer = 2, .state_in = 300, .state_out = 400 }; + var target = request; + for (target.layers[0..target.count]) |*mapping| std.mem.swap(u64, &mapping.state_in, &mapping.state_out); + for (2..state_binding.max_width + 1) |width| for (0..6) |scenario| { + request.width = width; + request.nonstate.rows = @intCast(width); + target.width = width; + target.nonstate.rows = @intCast(width); + var pair: [2]std.posix.socket_t = undefined; + if (std.c.socketpair(std.posix.AF.UNIX, std.posix.SOCK.STREAM, 0, &pair) != 0) return error.SocketPair; + defer for (pair) |descriptor| { + _ = std.c.close(descriptor); + }; + const nonblock: c_int = @bitCast(@as(u32, @bitCast(std.c.O{ .NONBLOCK = true }))); + for (pair) |descriptor| { + const flags = std.c.fcntl(descriptor, std.c.F.GETFL); + if (flags < 0 or std.c.fcntl(descriptor, std.c.F.SETFL, flags | nonblock) < 0) return error.Nonblocking; + } + const Peer = struct { + fn run(link: wire.Link, current: usize) void { + var control = GraphTestDriver.control(link); + const selected = control.graphAgree(51, .{ + .state = if (current == 2) .pending else if (current == 3) .capture else .ready, + .cohort = if (current == 1) 8 else 7, + }) catch @panic("binding selection failed"); + if (selected.choice != .replay) return; + const result = control.graphAgree(52, .{ .state = if (current == 5) .failed else .ready, .cohort = 7 }); + if (current == 4 or current == 5) { + if (result) |_| @panic("binding failure accepted") else |err| std.debug.assert(err == error.GraphInstantiationFailed); + } else { + std.debug.assert((result catch @panic("binding commit failed")).choice == .replay); + } + } + }; + const peer = try std.Thread.spawn(.{}, Peer.run, .{ wire.Link{ .fd = pair[1] }, scenario }); + defer peer.join(); + var control = GraphTestDriver.control(.{ .fd = pair[0] }); + var state = try state_binding.Binding.init(42, request); + for (request.layers[0..request.count], 0..) |mapping, index| { + state.nodes[index] = index; + state.parameters[index] = .{ .nonstate = request.nonstate, .state_in = mapping.state_in, .state_out = mapping.state_out }; + } + state.loaded = true; + var fake: Fake = .{ .fail = scenario == 4 }; + const selected = try control.graphAgree(51, .{ .state = .ready, .cohort = 7 }); + var published = false; + if (selected.choice == .replay) { + if (scenario == 4 or scenario == 5) { + try std.testing.expectError(error.GraphInstantiationFailed, state.applyTogether(&fake, target, &control, 52, selected.cohort)); + try std.testing.expect(state.poisoned); + } else { + try state.applyTogether(&fake, target, &control, 52, selected.cohort); + published = true; + } + } + try std.testing.expectEqual(scenario == 0, published); + try std.testing.expectEqual(@as(usize, if (scenario == 0 or scenario >= 4) 2 else 0), fake.calls); + }; +} + +test "exl3: state binding current pass supports exact window layouts through the MTP row bound" { + try @import("../kv/cuda_exl3_state_request_test.zig").currentPass(Key, grebind.windowStateRequest); +} diff --git a/zig/src/families/flashnext/exl3/graph/cuda_exl3_graph_run.zig b/zig/src/families/flashnext/exl3/graph/cuda_exl3_graph_run.zig new file mode 100644 index 0000000000..06d6f456fa --- /dev/null +++ b/zig/src/families/flashnext/exl3/graph/cuda_exl3_graph_run.zig @@ -0,0 +1,140 @@ +//! One rank's EXL3 decode step: replay the graph cached under its key, or capture it after enough sightings. + +const std = @import("std"); +const cuda = @import("cuda"); +const wire = @import("../../cuda_rank_wire.zig"); +const memory = @import("cuda_exl3_graph_memory.zig"); +const hub = @import("cuda_exl3_graph.zig"); +const gkey = @import("cuda_exl3_graph_key.zig"); +const gpending = @import("cuda_exl3_graph_pending.zig"); +const gcache = @import("cuda_exl3_graph_cache.zig"); +const gbudget = @import("cuda_exl3_graph_budget.zig"); +const grebind = @import("cuda_exl3_graph_rebind.zig"); +const glog = @import("cuda_exl3_graph_log.zig"); +const gfall = @import("cuda_exl3_graph_fallback.zig"); +const gtogether = @import("cuda_exl3_graph_together.zig"); +const Key = gkey.Key; +const Entry = gcache.Entry; +const Pending = gpending.Pending; + +/// The key's capture or instantiation failed: it backs off, and this round runs eagerly (nothing else is affected). +fn fellBack(key: Key, pass: anytype, ids: []const u32, pos0: usize) !void { + gfall.record(gfall.hashOf(1, key), gcache.clock); + hub.stats.eager_dplan += 1; + glog.note(key, .eager); + pass.prepared = true; + return pass.chunk(ids, pos0, null); +} + +pub fn run(key: Key, pass: anytype, ids: []const u32, pos0: usize) !void { + gcache.clock += 1; + if (pass.m.collective) |c| return gtogether.runTogether(c, key, pass, ids, pos0); + if (gcache.find(key)) |e| { + e.used = gcache.clock; + if (!gcache.settle(e)) { + if (e.pending == null) { + // the instantiation failed: forget the key, back it off, run this round eagerly + gcache.forget(e); + return fellBack(key, pass, ids, pos0); + } + hub.stats.eager_dplan += 1; + glog.note(key, .eager); + pass.prepared = true; + return pass.chunk(ids, pos0, null); + } + hub.stats.replays += 1; + glog.note(key, .replay); + return e.exec.?.launchOn(pass.stream); + } + if (try grebind.tryRebind(key, pass)) return; + if (gcache.sighting(key) < gcache.afterKey(key) or gfall.blocked(gfall.hashOf(1, key), gcache.clock)) { + hub.stats.eager_dplan += 1; + glog.note(key, .eager); + pass.prepared = true; + return pass.chunk(ids, pos0, null); + } + const t0 = std.Io.Clock.awake.now(pass.io).toNanoseconds(); + try pass.stream.synchronize(); + const admission = gbudget.reserveMemory(pass.stream, .decode, 0) catch |err| { + if (gfall.admissionDenied(err)) return fellBack(key, pass, ids, pos0); + return err; + }; + var admitted = true; + defer if (admitted) memory.release(admission) catch unreachable; + var graph = gfall.captureOnce(pass, ids, pos0) orelse return fellBack(key, pass, ids, pos0); + var kept = false; + defer if (!kept) graph.deinit(); + const t1 = std.Io.Clock.awake.now(pass.io).toNanoseconds(); + var how: []const u8 = "instantiated"; + const tp = pass.m.collective != null; + const at: *Entry = blk: { + // a same-shape graph is updated in place only when it is stale (its key not used for a while, e.g. a finished + var free_slot = false; + for (&gcache.entries) |*e| free_slot = free_slot or e.* == null; + if (grebind.capturedBinding(key, pass, graph) == null) { + for (&gcache.entries) |*e| if (e.*) |*x| if (!x.poisoned and x.exec != null and x.pending == null and gkey.sameShape(x.key, key) and (!free_slot or gcache.clock - x.used > gcache.stale_after)) { + // every earlier launch finished (each step ends on the host), so the old parameters are free to go + const r = try x.exec.?.update(graph); + if (r == .success) { + if (x.binding) |*binding| binding.invalidate(); + x.binding = null; + try gbudget.uploadResident(&x.exec.?, pass.stream); + x.key = key; + x.used = gcache.clock; + if (x.template) |*old| old.deinit(); + x.template = null; + if (tp) { + x.template = graph; + kept = true; + } + hub.stats.updates += 1; + glog.note(key, .updated); + how = "updated"; + break :blk x; + } + }; + } + if (gpending.asyncInstantiate()) if (gpending.ctxApi()) |api| { + var ctx: ?*anyopaque = null; + if (api.get(&ctx) == 0) { + const sl = gcache.slot(); + const pd = try std.heap.c_allocator.create(Pending); + const binding = grebind.capturedBinding(key, pass, graph); + pd.* = .{ .graph = graph, .ctx = ctx, .keep_template = tp or binding != null, .started_ns = wire.now() }; + kept = true; + pd.thread = std.Thread.spawn(.{}, Pending.main, .{pd}) catch |err| { + kept = false; + std.heap.c_allocator.destroy(pd); + return err; + }; + sl.* = .{ .key = key, .used = gcache.clock, .pending = pd, .admission = admission, .stream = pass.stream, .binding = binding }; + admitted = false; + glog.note(key, .instantiated); + hub.stats.captures += 1; + // this step runs eagerly while the graph instantiates + hub.capturing = false; + pass.prepared = true; + return pass.chunk(ids, pos0, null); + } + }; + const sl = gcache.slot(); + var exec = (if (gfall.take(&gfall.inject_instantiate_failures)) error.CudaFailed else graph.instantiate()) catch return fellBack(key, pass, ids, pos0); + errdefer exec.deinit(); + try gbudget.uploadResident(&exec, pass.stream); + try memory.installed(admission); + sl.* = .{ .key = key, .exec = exec, .used = gcache.clock, .admission = admission, .stream = pass.stream }; + sl.*.?.binding = grebind.capturedBinding(key, pass, graph); + if (sl.*.?.binding != null) { + sl.*.?.template = graph; + kept = true; + } + admitted = false; + glog.note(key, .instantiated); + break :blk &sl.*.?; + }; + hub.stats.captures += 1; + const t2 = std.Io.Clock.awake.now(pass.io).toNanoseconds(); + hub.stats.capture_ns += @intCast(t2 - t0); + std.log.info("flashnext exl3 graph: captured a {s} of {d} rows in {d:.1} ms + {s} in {d:.1} ms ({d} captures, {d} updates, {d} replays, {d} eager, {d} declined steps)", .{ if (key.kind == 1) "decode round" else "window", key.rows, @as(f64, @floatFromInt(t1 - t0)) / 1e6, how, @as(f64, @floatFromInt(t2 - t1)) / 1e6, hub.stats.captures, hub.stats.updates, hub.stats.replays, hub.stats.eager_dplan, hub.stats.declined }); + try at.exec.?.launchOn(pass.stream); +} diff --git a/zig/src/families/flashnext/exl3/graph/cuda_exl3_graph_test_driver.zig b/zig/src/families/flashnext/exl3/graph/cuda_exl3_graph_test_driver.zig new file mode 100644 index 0000000000..2323fab5fc --- /dev/null +++ b/zig/src/families/flashnext/exl3/graph/cuda_exl3_graph_test_driver.zig @@ -0,0 +1,117 @@ +//! A stand-in CUDA driver and rank link for the EXL3 graph cache's host tests. + +const std = @import("std"); +const cuda = @import("cuda"); +const collective = @import("../../cuda_collective.zig"); +const wire = @import("../../cuda_rank_wire.zig"); +const gpending = @import("cuda_exl3_graph_pending.zig"); +const Pending = gpending.Pending; + +pub const GraphTestDriver = struct { + pub var exec_drops: usize = 0; + pub var graph_drops: usize = 0; + pub var launches: usize = 0; + pub var exec_updates: usize = 0; + + pub fn dropExec(_: cuda.abi.GraphExec) callconv(.c) c_int { + exec_drops += 1; + return 0; + } + + pub fn dropGraph(_: cuda.abi.Graph) callconv(.c) c_int { + graph_drops += 1; + return 0; + } + + pub fn instantiate(out: *cuda.abi.GraphExec, _: cuda.abi.Graph, _: u64) callconv(.c) c_int { + out.* = @ptrFromInt(1); + return 0; + } + + pub fn getContext(out: *?*anyopaque) callconv(.c) c_int { + out.* = @ptrFromInt(1); + return 0; + } + + pub fn noContext(_: *?*anyopaque) callconv(.c) c_int { + return 1; + } + + pub fn setContext(_: ?*anyopaque) callconv(.c) c_int { + return 0; + } + + pub fn driver() cuda.Driver { + var d: cuda.Driver = undefined; + d.api.cuGraphExecDestroy = dropExec; + d.api.cuGraphDestroy = dropGraph; + d.api.cuGraphLaunch = launch; + d.api.cuGraphInstantiateWithFlags = instantiate; + d.api.cuMemGetInfo_v2 = memoryInfo; + d.api.cuCtxSynchronize = synchronize; + d.api.cuStreamSynchronize = synchronizeStream; + d.api.cuStreamBeginCapture_v2 = beginCapture; + d.api.cuStreamEndCapture = endCapture; + d.api.cuGraphUpload = upload; + d.api.cuGraphExecUpdate_v2 = update; + return d; + } + + pub fn memoryInfo(free: *usize, total: *usize) callconv(.c) c_int { + free.* = 100; + total.* = 100; + return 0; + } + + pub fn synchronize() callconv(.c) c_int { + return 0; + } + + pub fn synchronizeStream(_: cuda.abi.Stream) callconv(.c) c_int { + return 0; + } + + pub fn beginCapture(_: cuda.abi.Stream, _: cuda.abi.CaptureMode) callconv(.c) c_int { + return 0; + } + + pub fn endCapture(_: cuda.abi.Stream, out: *cuda.abi.Graph) callconv(.c) c_int { + out.* = @ptrFromInt(2); + return 0; + } + + pub fn upload(_: cuda.abi.GraphExec, _: cuda.abi.Stream) callconv(.c) c_int { + return 0; + } + + pub fn update(_: cuda.abi.GraphExec, _: cuda.abi.Graph, info: *cuda.abi.ExecUpdateResultInfo) callconv(.c) c_int { + exec_updates += 1; + info.result = .success; + return 0; + } + + pub fn launch(_: cuda.abi.GraphExec, _: cuda.abi.Stream) callconv(.c) c_int { + launches += 1; + return 0; + } + + pub fn abort(_: cuda.nccl.Comm) callconv(.c) cuda.nccl.Result { + return 0; + } + + pub fn control(link: wire.Link) collective.Collective { + var c: collective.Collective = undefined; + c.comm = @ptrFromInt(1); + c.api.ncclCommAbort = abort; + c.link = link; + c.failed = false; + c.timeout_ms = 1000; + c.graph_step = 0; + return c; + } + + pub fn finishPending(pd: *Pending, driver_: *cuda.Driver) void { + pd.exec = .{ .d = driver_, .handle = @ptrFromInt(1) }; + pd.done.store(true, .release); + } +}; diff --git a/zig/src/families/flashnext/exl3/graph/cuda_exl3_graph_together.zig b/zig/src/families/flashnext/exl3/graph/cuda_exl3_graph_together.zig new file mode 100644 index 0000000000..ccdecd3e26 --- /dev/null +++ b/zig/src/families/flashnext/exl3/graph/cuda_exl3_graph_together.zig @@ -0,0 +1,190 @@ +//! A tensor-parallel EXL3 decode step: the ranks agree on replay, capture, update and retirement together. + +const std = @import("std"); +const cuda = @import("cuda"); +const collective = @import("../../cuda_collective.zig"); +const wire = @import("../../cuda_rank_wire.zig"); +const memory = @import("cuda_exl3_graph_memory.zig"); +const hub = @import("cuda_exl3_graph.zig"); +const gkey = @import("cuda_exl3_graph_key.zig"); +const gpending = @import("cuda_exl3_graph_pending.zig"); +const gcache = @import("cuda_exl3_graph_cache.zig"); +const gbudget = @import("cuda_exl3_graph_budget.zig"); +const grebind = @import("cuda_exl3_graph_rebind.zig"); +const glog = @import("cuda_exl3_graph_log.zig"); +const gfall = @import("cuda_exl3_graph_fallback.zig"); +const Key = gkey.Key; +const Entry = gcache.Entry; + +pub fn graphVote(entry: ?*Entry, timeout_ms: u64) collective.GraphVote { + const x = entry orelse return .{ .state = .capture }; + if (gcache.settle(x)) return .{ .state = .ready, .cohort = x.cohort }; + if (x.pending) |pd| { + const elapsed = wire.now() -| pd.started_ns; + return .{ .state = if (elapsed >= timeout_ms *| std.time.ns_per_ms) .timed_out else .pending, .cohort = x.cohort }; + } + return .{ .state = .failed, .cohort = x.cohort }; +} + +pub fn runTogether(c: *collective.Collective, key: Key, pass: anytype, ids: []const u32, pos0: usize) !void { + errdefer c.failGraph(); + const signature = gkey.graphSignature(key, pos0); + var entry = gcache.find(key); + if (entry) |x| x.used = gcache.clock; + const fail_hash = gfall.hashOf(1, key); + var vote = if (entry != null) graphVote(entry, c.graphTimeoutMs()) else collective.GraphVote{ + .state = if (grebind.stateRequest(key, pass)) |request| + if (grebind.rebindCandidate(key, request, true) != null) .capture else if (gcache.sighting(key) < gcache.afterKey(key) or gfall.blocked(fail_hash, gcache.clock)) .eager else .capture + else if (gcache.sighting(key) < gcache.afterKey(key) or gfall.blocked(fail_hash, gcache.clock)) .eager else .capture, + }; + // this rank's instantiation failed: it votes eager (the peer's graph stays), the key backs off, both run eagerly + const instantiation_failed = vote.state == .failed; + if (instantiation_failed) vote.state = .eager; + const decision = try c.graphAgree(signature, vote); + if (instantiation_failed) { + gcache.forget(entry.?); + entry = null; + gfall.record(fail_hash, gcache.clock); + } + switch (decision.choice) { + .eager => { + hub.stats.eager_dplan += 1; + glog.note(key, .eager); + pass.prepared = true; + return pass.chunk(ids, pos0, null); + }, + .replay => { + hub.stats.replays += 1; + glog.note(key, .replay); + return entry.?.exec.?.launchOn(pass.stream); + }, + .capture => {}, + } + if (try grebind.tryRebindTogether(c, signature, key, pass)) return; + const started = std.Io.Clock.awake.now(pass.io).toNanoseconds(); + try pass.stream.synchronize(); + const admission = gbudget.reserveTogether(c, pass.stream, signature, decision.cohort) catch |err| { + if (!gfall.admissionDenied(err)) return err; + gfall.record(fail_hash, gcache.clock); + hub.stats.eager_dplan += 1; + glog.note(key, .eager); + pass.prepared = true; + return pass.chunk(ids, pos0, null); + }; + entry = gcache.find(key); + var admitted = true; + defer if (admitted) memory.release(admission) catch unreachable; + const Attempt = struct { + pass: @TypeOf(pass), + ids: []const u32, + pos0: usize, + pub fn run(a: @This()) ?cuda.graph.Graph { + return gfall.captureOnce(a.pass, a.ids, a.pos0); + } + }; + // both ranks hold a capture or neither does: a failure on either rank sends both to the eager round + var graph = (try gfall.agreedCapture(c, signature, decision.cohort, Attempt{ .pass = pass, .ids = ids, .pos0 = pos0 })) orelse { + gfall.record(fail_hash, gcache.clock); + hub.stats.eager_dplan += 1; + glog.note(key, .eager); + pass.prepared = true; + return pass.chunk(ids, pos0, null); + }; + var kept = false; + defer if (!kept) graph.deinit(); + var updated = false; + var install_failed = false; + var candidate: ?*Entry = null; + if (grebind.capturedBinding(key, pass, graph) == null) for (&gcache.entries) |*maybe| if (maybe.*) |*current| { + if (!current.poisoned and current.cohort != 0 and current.exec != null and current.pending == null and gkey.sameShape(current.key, key) and + (std.meta.eql(current.key, key) or gcache.clock - current.used > gcache.stale_after)) + { + if (candidate == null or current.used < candidate.?.used) candidate = current; + } + }; + const update_agreement = try c.graphAgree(signature ^ 0x757064617465, if (candidate) |current| + .{ .state = .ready, .cohort = current.cohort } + else + .{ .state = .capture }); + if (update_agreement.choice == .replay) { + const current = candidate.?; + const result = try current.exec.?.update(graph); + if (result == .success) { + if (current.binding) |*binding| binding.invalidate(); + current.binding = null; + if (current.template) |*old| old.deinit(); + current.template = graph; + kept = true; + updated = true; + } + const result_agreement = try c.graphAgree(signature ^ 0x75706461746564, .{ .state = if (updated) .ready else .capture, .cohort = decision.cohort }); + if (result_agreement.choice == .replay) { + if (current.admission) |old| try memory.release(old); + current.admission = admission; + admitted = false; + try gbudget.uploadResident(¤t.exec.?, pass.stream); + try memory.installed(admission); + current.key = key; + current.used = gcache.clock; + current.cohort = decision.cohort; + current.stream = pass.stream; + entry = current; + } else { + if (updated) { + current.template = null; + kept = false; + } + gcache.drop(current); + for (&gcache.entries) |*slot_entry| if (slot_entry.*) |*cached| if (cached == current) { + slot_entry.* = null; + }; + entry = gcache.find(key); + updated = false; + try gbudget.retirementBarrier(c, signature ^ 0x72656275696c64, decision.cohort); + } + } + if (!updated) { + const target = if (entry) |old| blk: { + gcache.drop(old); + for (&gcache.entries) |*slot_entry| if (slot_entry.*) |*x| if (x == old) { + slot_entry.* = null; + break :blk slot_entry; + }; + unreachable; + } else blk: { + for (&gcache.entries) |*slot_entry| if (slot_entry.* == null) break :blk slot_entry; + return error.GraphBudgetBusy; + }; + if (try gpending.startInstantiate(graph, true)) |pending| { + target.* = .{ .key = key, .used = gcache.clock, .pending = pending, .cohort = decision.cohort, .admission = admission, .stream = pass.stream, .binding = grebind.capturedBinding(key, pass, graph) }; + } else if (if (gfall.take(&gfall.inject_instantiate_failures)) error.CudaFailed else graph.instantiate()) |instantiated| { + var exec = instantiated; + errdefer exec.deinit(); + try gbudget.uploadResident(&exec, pass.stream); + try memory.installed(admission); + target.* = .{ .key = key, .exec = exec, .used = gcache.clock, .template = graph, .cohort = decision.cohort, .admission = admission, .stream = pass.stream, .binding = grebind.capturedBinding(key, pass, graph) }; + } else |_| install_failed = true; + if (!install_failed) { + admitted = false; + kept = true; + entry = &target.*.?; + } else entry = null; + } + if (!install_failed) { + hub.stats.captures += 1; + if (updated) hub.stats.updates += 1; + glog.note(key, if (updated) .updated else .instantiated); + hub.stats.capture_ns += @intCast(std.Io.Clock.awake.now(pass.io).toNanoseconds() - started); + } else gfall.record(fail_hash, gcache.clock); + // a failed local instantiation votes eager: the peer's ready graph stays, and both run this round eagerly + const installed = try c.graphAgree(signature ^ 0x696e7374616c6c, if (install_failed) .{ .state = .eager } else graphVote(entry, c.graphTimeoutMs())); + switch (installed.choice) { + .replay => try entry.?.exec.?.launchOn(pass.stream), + .eager => { + hub.stats.eager_dplan += 1; + pass.prepared = true; + try pass.chunk(ids, pos0, null); + }, + .capture => return error.GraphRendezvousMismatch, + } +} diff --git a/zig/src/families/flashnext/exl3/graph/cuda_exl3_graph_together_test.zig b/zig/src/families/flashnext/exl3/graph/cuda_exl3_graph_together_test.zig new file mode 100644 index 0000000000..d2f9a77667 --- /dev/null +++ b/zig/src/families/flashnext/exl3/graph/cuda_exl3_graph_together_test.zig @@ -0,0 +1,287 @@ +//! Host tests of the tensor-parallel EXL3 graph step: dispatch, eviction and update agreements. + +const std = @import("std"); +const windows = @import("builtin").os.tag == .windows; +const cuda = @import("cuda"); +const collective = @import("../../cuda_collective.zig"); +const wire = @import("../../cuda_rank_wire.zig"); +const state_binding = @import("../kv/exl3_state_binding.zig"); +const memory = @import("cuda_exl3_graph_memory.zig"); +const gkey = @import("cuda_exl3_graph_key.zig"); +const gpending = @import("cuda_exl3_graph_pending.zig"); +const gcache = @import("cuda_exl3_graph_cache.zig"); +const gbudget = @import("cuda_exl3_graph_budget.zig"); +const grun = @import("cuda_exl3_graph_run.zig"); +const gtogether = @import("cuda_exl3_graph_together.zig"); +const gprefill = @import("cuda_exl3_graph_prefill.zig"); +const gdriver = @import("cuda_exl3_graph_test_driver.zig"); +const Key = gkey.Key; +const PKey = gprefill.PKey; +const GraphTestDriver = gdriver.GraphTestDriver; + +test "exl3: TP eviction proposes the least recently used cohort and ignores local graphs" { + gcache.entries = @splat(null); + defer gcache.entries = @splat(null); + const key: Key = .{ .kind = 1, .rows = 1, .scratch = 1, .scores = 2, .output = 3 }; + gcache.entries[0] = .{ .key = key, .used = 30, .cohort = 3 }; + gcache.entries[1] = .{ .key = key, .used = 10, .cohort = 7 }; + gcache.entries[2] = .{ .key = key, .used = 0, .cohort = 0 }; + try std.testing.expectEqual(@as(u64, 7), gbudget.oldestCohort()); + gcache.entries[0].?.used = 10; + try std.testing.expectEqual(@as(u64, 3), gbudget.oldestCohort()); +} + +test "exl3: TP decode dispatch stays eager while its peer instantiates and replays only after agreement" { + if (windows) return error.SkipZigTest; // The rank wire explicitly refuses Windows. + gcache.clear(); + defer gcache.clear(); + var pair: [2]std.posix.socket_t = undefined; + if (std.c.socketpair(std.posix.AF.UNIX, std.posix.SOCK.STREAM, 0, &pair) != 0) return error.SocketPair; + defer for (pair) |fd| { + _ = std.c.close(fd); + }; + const nonblock: c_int = @bitCast(@as(u32, @bitCast(std.c.O{ .NONBLOCK = true }))); + for (pair) |fd| { + const flags = std.c.fcntl(fd, std.c.F.GETFL); + if (flags < 0 or std.c.fcntl(fd, std.c.F.SETFL, flags | nonblock) < 0) return error.Nonblocking; + } + const key: Key = .{ .kind = 1, .rows = 1, .scratch = 1, .scores = 2, .output = 3 }; + const Peer = struct { + fn main(link: wire.Link, signature: u64) void { + var c = GraphTestDriver.control(link); + const first = c.graphAgree(signature, .{ .state = .pending, .cohort = 7 }) catch @panic("first peer agreement failed"); + std.debug.assert(first.choice == .eager); + const second = c.graphAgree(signature, .{ .state = .ready, .cohort = 7 }) catch @panic("second peer agreement failed"); + std.debug.assert(second.choice == .replay); + } + }; + const thread = try std.Thread.spawn(.{}, Peer.main, .{ wire.Link{ .fd = pair[1] }, gkey.graphSignature(key, 0) }); + defer thread.join(); + var c = GraphTestDriver.control(.{ .fd = pair[0] }); + var driver = GraphTestDriver.driver(); + GraphTestDriver.launches = 0; + gcache.entries[0] = .{ .key = key, .used = 1, .exec = .{ .d = &driver, .handle = @ptrFromInt(1) }, .cohort = 7 }; + const Pass = struct { + io: std.Io = std.testing.io, + stream: cuda.Stream, + m: struct { collective: ?*collective.Collective }, + prepared: bool = false, + chunks: usize = 0, + + pub fn chunk(p: *@This(), _: []const u32, _: usize, _: ?[]u8) !void { + p.chunks += 1; + } + }; + var pass: Pass = .{ .stream = cuda.Stream.borrowed(&driver, null), .m = .{ .collective = &c } }; + try grun.run(key, &pass, &.{1}, 0); + try std.testing.expectEqual(@as(usize, 1), pass.chunks); + try std.testing.expectEqual(@as(usize, 0), GraphTestDriver.launches); + try grun.run(key, &pass, &.{1}, 0); + try std.testing.expectEqual(@as(usize, 1), pass.chunks); + try std.testing.expectEqual(@as(usize, 1), GraphTestDriver.launches); +} + +test "exl3: TP pressure on either rank retires cohorts before retrying admission" { + if (windows) return error.SkipZigTest; // The rank wire explicitly refuses Windows. + gcache.clear(); + defer gcache.clear(); + for ([_]bool{ true, false }) |local_pressure| { + try memory.configure(.{ .reservation_bytes = 100, .headroom_bytes = 20, .charge = .{ .exec = 20, .template = 10, .transient = 10 }, .prefill_charge = .{ .exec = 1, .template = 0, .transient = 0 } }); + var driver = GraphTestDriver.driver(); + GraphTestDriver.exec_drops = 0; + const local = try memory.admit(.prefill, 0, 100); + try memory.installed(local); + gprefill.pentries[0] = .{ .key = std.mem.zeroes(PKey), .exec = .{ .d = &driver, .handle = @ptrFromInt(1) }, .used = 0, .admission = local }; + if (local_pressure) for (0..2) |index| { + const ticket = try memory.admit(.decode, index + 7, 100); + try memory.installed(ticket); + gcache.entries[index] = .{ .key = .{ .kind = 1, .rows = index + 1, .scratch = 1, .scores = 2, .output = 3 }, .exec = .{ .d = &driver, .handle = @ptrFromInt(1) }, .used = index + 1, .cohort = index + 7, .admission = ticket }; + }; + var pair: [2]std.posix.socket_t = undefined; + if (std.c.socketpair(std.posix.AF.UNIX, std.posix.SOCK.STREAM, 0, &pair) != 0) return error.SocketPair; + defer for (pair) |fd| { + _ = std.c.close(fd); + }; + const nonblock: c_int = @bitCast(@as(u32, @bitCast(std.c.O{ .NONBLOCK = true }))); + for (pair) |fd| { + const flags = std.c.fcntl(fd, std.c.F.GETFL); + if (flags < 0 or std.c.fcntl(fd, std.c.F.SETFL, flags | nonblock) < 0) return error.Nonblocking; + } + const Peer = struct { + fn run(link: wire.Link, pressure: bool) void { + var control = GraphTestDriver.control(link); + const first = control.graphAgree(42 ^ 0x61646d6974, .{ .state = if (pressure) .ready else .eager, .cohort = 11 }) catch @panic("pressure agreement failed"); + std.debug.assert(first.choice == .eager); + const victim = control.graphVictim(42 ^ 0x76696374696d, if (pressure) 8 else 7) catch @panic("victim agreement failed"); + std.debug.assert(victim == 7); + gbudget.retirementBarrier(&control, 42 ^ 0x71756965736365, victim) catch @panic("quiescence agreement failed"); + gbudget.retirementBarrier(&control, 42 ^ 0x72657469726564, victim) catch @panic("retirement agreement failed"); + const admitted = control.graphAgree(42 ^ 0x61646d6974, .{ .state = .ready, .cohort = 11 }) catch @panic("retry agreement failed"); + std.debug.assert(admitted.choice == .replay); + } + }; + const peer = try std.Thread.spawn(.{}, Peer.run, .{ wire.Link{ .fd = pair[1] }, local_pressure }); + defer peer.join(); + var control = GraphTestDriver.control(.{ .fd = pair[0] }); + const ticket = try gbudget.reserveTogether(&control, cuda.Stream.borrowed(&driver, null), 42, 11); + try std.testing.expectEqual(@as(u64, if (local_pressure) 71 else 41), memory.charged()); + try std.testing.expectEqual(@as(usize, if (local_pressure) 1 else 0), GraphTestDriver.exec_drops); + try std.testing.expect(gprefill.pentries[0] != null); + if (local_pressure) try std.testing.expectEqual(@as(u64, 8), gcache.entries[1].?.cohort); + try memory.release(ticket); + gcache.clear(); + } +} + +test "exl3: exhausted TP victims reclaim locals incrementally with asymmetric peer caches" { + if (windows) return error.SkipZigTest; // The rank wire explicitly refuses Windows. + gcache.clear(); + defer gcache.clear(); + for ([_]bool{ false, true }) |local_cached| for ([_]bool{ false, true }) |peer_cached| { + try memory.configure(.{ .reservation_bytes = 100, .headroom_bytes = 20, .charge = .{ .exec = 20, .template = 10, .transient = 10 } }); + var driver = GraphTestDriver.driver(); + GraphTestDriver.exec_drops = 0; + if (local_cached) { + const prefill = try memory.admit(.prefill, 0, 100); + try memory.installed(prefill); + gprefill.pentries[0] = .{ .key = std.mem.zeroes(PKey), .exec = .{ .d = &driver, .handle = @ptrFromInt(1) }, .used = 1, .admission = prefill }; + const local = try memory.admit(.decode, 0, 100); + try memory.installed(local); + gcache.entries[0] = .{ .key = .{ .kind = 1, .rows = 1, .scratch = 1, .scores = 2, .output = 3 }, .exec = .{ .d = &driver, .handle = @ptrFromInt(1) }, .used = 2, .admission = local }; + } + var pair: [2]std.posix.socket_t = undefined; + if (std.c.socketpair(std.posix.AF.UNIX, std.posix.SOCK.STREAM, 0, &pair) != 0) return error.SocketPair; + defer for (pair) |fd| { + _ = std.c.close(fd); + }; + const nonblock: c_int = @bitCast(@as(u32, @bitCast(std.c.O{ .NONBLOCK = true }))); + for (pair) |fd| { + const flags = std.c.fcntl(fd, std.c.F.GETFL); + if (flags < 0 or std.c.fcntl(fd, std.c.F.SETFL, flags | nonblock) < 0) return error.Nonblocking; + } + const Peer = struct { + fn run(link: wire.Link, local_cached_peer: bool, cached: bool) void { + var control = GraphTestDriver.control(link); + const first = control.graphAgree(42 ^ 0x61646d6974, .{ .state = .eager, .cohort = 11 }) catch @panic("pressure agreement failed"); + std.debug.assert(first.choice == .eager); + const victim = control.graphVictim(42 ^ 0x76696374696d, 0) catch @panic("victim agreement failed"); + std.debug.assert(victim == 0); + gbudget.retirementBarrier(&control, 42 ^ 0x6c6f63616c7175, 11) catch @panic("local quiescence failed"); + const reclaimed = control.graphVictim(42 ^ 0x6c6f63616c646f, @intFromBool(cached)) catch @panic("local destruction acknowledgment failed"); + std.debug.assert(reclaimed == @as(u64, @intFromBool(local_cached_peer or cached))); + if (reclaimed == 0) return; + const admitted = control.graphAgree(42 ^ 0x61646d6974, .{ .state = .ready, .cohort = 11 }) catch @panic("retry agreement failed"); + std.debug.assert(admitted.choice == .replay); + } + }; + const peer = try std.Thread.spawn(.{}, Peer.run, .{ wire.Link{ .fd = pair[1] }, local_cached, peer_cached }); + defer peer.join(); + var control = GraphTestDriver.control(.{ .fd = pair[0] }); + const result = gbudget.reserveTogether(&control, cuda.Stream.borrowed(&driver, null), 42, 11); + if (local_cached or peer_cached) { + const ticket = try result; + try std.testing.expectEqual(@as(u64, if (local_cached) 70 else 40), memory.charged()); + try memory.release(ticket); + } else { + try std.testing.expectError(error.GraphBudgetExceeded, result); + try std.testing.expectEqual(@as(u64, 0), memory.charged()); + } + try std.testing.expectEqual(@as(usize, @intFromBool(local_cached)), GraphTestDriver.exec_drops); + try std.testing.expect(gprefill.pentries[0] == null); + try std.testing.expectEqual(local_cached, gcache.entries[0] != null); + gcache.clear(); + }; +} + +test "exl3: TP updates agree the old cohort and transfer template and admission together" { + if (windows) return error.SkipZigTest; // The rank wire explicitly refuses Windows. + try testTPUpdate(true); +} + +test "exl3: asymmetric TP update rejection retires both execs before re-instantiating" { + if (windows) return error.SkipZigTest; // The rank wire explicitly refuses Windows. + try testTPUpdate(false); +} + +fn testTPUpdate(peer_accept: bool) !void { + if (windows) return error.SkipZigTest; + gcache.clear(); + defer gcache.clear(); + const previous_api = gpending.ctx_api; + defer gpending.ctx_api = previous_api; + gpending.ctx_api = .{ .get = GraphTestDriver.noContext, .set = GraphTestDriver.setContext }; + try memory.configure(.{ .reservation_bytes = 120, .headroom_bytes = 20, .charge = .{ .exec = 20, .template = 10, .transient = 10 } }); + var driver = GraphTestDriver.driver(); + GraphTestDriver.exec_updates = 0; + GraphTestDriver.exec_drops = 0; + GraphTestDriver.graph_drops = 0; + const old_ticket = try memory.admit(.decode, 7, 100); + try memory.installed(old_ticket); + const key: Key = .{ .kind = 1, .rows = 1, .scratch = 2, .scores = 2, .output = 3 }; + var old_key = key; + old_key.scratch = 1; + gcache.entries[0] = .{ .key = old_key, .exec = .{ .d = &driver, .handle = @ptrFromInt(1) }, .template = .{ .d = &driver, .handle = @ptrFromInt(1) }, .used = 1, .cohort = 7, .admission = old_ticket }; + var old_binding_request: state_binding.Request = .{ + .width = 2, + .chain = 10, + .chain_to = 11, + .geometry = .{ .grid = .{ 16, 8, 1 }, .block = .{ 128, 1, 1 } }, + .nonstate = .{ .gq = 1, .gk = 2, .gv = 3, .ggb = 4, .rows = 2, .ys = 5 }, + .count = 1, + }; + old_binding_request.layers[0] = .{ .layer = 0, .state_in = 100, .state_out = 200 }; + gcache.entries[0].?.binding = try state_binding.Binding.init(1, old_binding_request); + gcache.seen[0] = .{ .key = key, .n = std.math.maxInt(u32) - 1, .used = 1 }; + const previous_clock = gcache.clock; + defer gcache.clock = previous_clock; + gcache.clock = 1024; + var pair: [2]std.posix.socket_t = undefined; + if (std.c.socketpair(std.posix.AF.UNIX, std.posix.SOCK.STREAM, 0, &pair) != 0) return error.SocketPair; + defer for (pair) |fd| { + _ = std.c.close(fd); + }; + const nonblock: c_int = @bitCast(@as(u32, @bitCast(std.c.O{ .NONBLOCK = true }))); + for (pair) |fd| { + const flags = std.c.fcntl(fd, std.c.F.GETFL); + if (flags < 0 or std.c.fcntl(fd, std.c.F.SETFL, flags | nonblock) < 0) return error.Nonblocking; + } + const Peer = struct { + fn run(link: wire.Link, signature: u64, accept: bool) void { + var control = GraphTestDriver.control(link); + control.graph_step = 10; + const capture = control.graphAgree(signature, .{ .state = .capture }) catch @panic("capture agreement failed"); + std.debug.assert(capture.choice == .capture and capture.cohort == 11); + gbudget.retirementBarrier(&control, signature ^ 0x61646d6974, 11) catch @panic("admission failed"); + const captured = control.graphAgree(signature ^ 0x63617074757265, .{ .state = .ready, .cohort = 11 }) catch @panic("capture outcome agreement failed"); + std.debug.assert(captured.choice == .replay); + gbudget.retirementBarrier(&control, signature ^ 0x757064617465, 7) catch @panic("old cohort agreement failed"); + const result = control.graphAgree(signature ^ 0x75706461746564, .{ .state = if (accept) .ready else .capture, .cohort = 11 }) catch @panic("update result agreement failed"); + std.debug.assert(result.choice == if (accept) collective.GraphChoice.replay else collective.GraphChoice.capture); + if (!accept) gbudget.retirementBarrier(&control, signature ^ 0x72656275696c64, 11) catch @panic("rebuild retirement agreement failed"); + gbudget.retirementBarrier(&control, signature ^ 0x696e7374616c6c, 11) catch @panic("installation agreement failed"); + } + }; + const peer = try std.Thread.spawn(.{}, Peer.run, .{ wire.Link{ .fd = pair[1] }, gkey.graphSignature(key, 0), peer_accept }); + defer peer.join(); + var control = GraphTestDriver.control(.{ .fd = pair[0] }); + control.graph_step = 10; + const Pass = struct { + io: std.Io = std.testing.io, + stream: cuda.Stream, + prepared: bool = false, + chunks: usize = 0, + pub fn chunk(pass: *@This(), _: []const u32, _: usize, _: ?[]u8) !void { + pass.chunks += 1; + } + }; + var pass: Pass = .{ .stream = cuda.Stream.borrowed(&driver, null) }; + try gtogether.runTogether(&control, key, &pass, &.{1}, 0); + try std.testing.expectEqual(@as(usize, 1), GraphTestDriver.exec_updates); + try std.testing.expectEqual(@as(usize, 1), GraphTestDriver.graph_drops); + try std.testing.expectEqual(@as(usize, if (peer_accept) 0 else 1), GraphTestDriver.exec_drops); + try std.testing.expectEqual(@as(u64, 30), memory.charged()); + try std.testing.expectEqual(@as(u64, 11), gcache.find(key).?.cohort); + try std.testing.expect(gcache.find(key).?.binding == null); + try std.testing.expectError(error.StaleGraphTicket, memory.release(old_ticket)); + try std.testing.expectEqual(@as(usize, 1), pass.chunks); +} diff --git a/zig/src/families/flashnext/exl3/graph/exl3_graph.zig b/zig/src/families/flashnext/exl3/graph/exl3_graph.zig new file mode 100644 index 0000000000..bd536fac63 --- /dev/null +++ b/zig/src/families/flashnext/exl3/graph/exl3_graph.zig @@ -0,0 +1,414 @@ +//! Host-side description of what a CUDA graph may bake in versus what replay overwrites (ported from + +const std = @import("std"); +const offload = @import("../offload/exl3_offload.zig"); +const round_ = @import("../model/exl3_round.zig"); + +/// The card and the device bytes kept free (were exl3_cache.zig's): a 16 GiB RTX 5060 Ti, 256 MiB reserve. +pub const card_bytes: u64 = 16 << 30; +pub const device_reserve: u64 = 256 << 20; + +pub const Binding = struct { + name: []const u8, + address: u64, + /// True: replay input. False: fixed for the life of the executable graph. + replay: bool, +}; + +pub const Capture = struct { + window: u32, + bindings: []const Binding, +}; + +/// One executable graph. Rank, layer, and window are the whole key. +pub const Key = struct { + rank: u8, + layer: u8, + window: u32, +}; + +const replay_names = [_][]const u8{ + "expert_ids", + "expert_tables", + "picks", + "ple_rows", + "kv_index", +}; + +const fixed_names = [_][]const u8{ + "dense_weights", + "shared_weights", + "pool_lines", + "expert_scales", +}; + +fn listed(name: []const u8, table: []const []const u8) bool { + for (table) |entry| if (std.mem.eql(u8, name, entry)) return true; + return false; +} + +/// Expert ids (pool line indices, the old engine), the expert pointer tables, picks, PLE rows. +pub fn isReplayName(name: []const u8) bool { + return listed(name, &replay_names); +} + +fn isFixedName(name: []const u8) bool { + return listed(name, &fixed_names); +} + +/// Known names must carry the flag this module requires. Unknown names keep the caller's flag. +pub fn valid(c: Capture) bool { + for (c.bindings) |b| { + if (isReplayName(b.name)) { + if (!b.replay) return false; + } else if (isFixedName(b.name)) { + if (b.replay) return false; + } + } + return true; +} + +/// Same window, same binding order, same replay flags, and the same address on every fixed binding. +pub fn sameTopology(a: Capture, b: Capture) bool { + if (a.window != b.window or a.bindings.len != b.bindings.len) return false; + for (a.bindings, b.bindings) |ba, bb| { + if (!std.mem.eql(u8, ba.name, bb.name) or ba.replay != bb.replay) return false; + if (!ba.replay and ba.address != bb.address) return false; + } + return true; +} + +/// `sameTopology`, and every binding address matches. Eager and the graph stand-in share ids and weights. +pub fn argumentsEqual(a: Capture, b: Capture) bool { + if (!sameTopology(a, b)) return false; + for (a.bindings, b.bindings) |ba, bb| { + if (ba.address != bb.address) return false; + } + return true; +} + +fn sameKey(a: Key, b: Key) bool { + return a.rank == b.rank and a.layer == b.layer and a.window == b.window; +} + +/// Cold start stays eager through the first reap wave. +pub fn captureReady(startup: bool, decisions: []const offload.Decision, lines: []const offload.Line) bool { + if (startup) return false; + for (decisions) |d| if (d == .copy) return false; + for (lines) |line| if (line == .inflight) return false; + return true; +} + +/// Upstream `Experts` + offload pool: a captured experts graph replays as is when the round copies nothing. +pub fn roundReady(startup: bool, r: *const round_.Round) bool { + if (startup) return false; + return r.allHits() and !r.masked(); +} + +pub const Promote = enum { + /// Rewrite the expert_ids line-index buffer and replay. + rewrite_ids, + /// This key has no capture yet. + recapture, +}; + +/// One capture per key. A stored key rewrites expert_ids and does not recapture. +pub fn promote(keys: []const Key, key: Key) Promote { + for (keys) |have| if (sameKey(have, key)) return .rewrite_ids; + return .recapture; +} + +/// One routed line. `count` is counts[layer][id]. +pub const LineUse = struct { + layer: u8, + id: u16, + count: u32, +}; + +fn colder(a: LineUse, b: LineUse) bool { + if (a.count != b.count) return a.count < b.count; + if (a.layer != b.layer) return a.layer < b.layer; + return a.id < b.id; +} + +/// `rank_pinned + kv_rank + n * stride + workspace` fits on the card. +fn cardHolds(rank_pinned: u64, kv_rank: u64, n: usize, stride: u64, workspace: u64) bool { + const rows: u64 = @intCast(n); + const line_bytes = std.math.mul(u64, rows, stride) catch return false; + const fixed = std.math.add(u64, rank_pinned, kv_rank) catch return false; + const with_lines = std.math.add(u64, fixed, line_bytes) catch return false; + const total = std.math.add(u64, with_lines, workspace) catch return false; + return total <= card_bytes; +} + +/// Lines kept once the batch workspace is placed. +pub fn fitWorkspace( + rank_pinned: u64, + kv_rank: u64, + lines: []LineUse, + stride: u64, + workspace: u64, +) error{GraphWorkspace}!usize { + if (workspace <= device_reserve) return lines.len; + var active: usize = lines.len; + while (active > 0 and !cardHolds(rank_pinned, kv_rank, active, stride, workspace)) { + var cold: usize = 0; + for (lines[1..active], 1..) |line, i| { + if (colder(line, lines[cold])) cold = i; + } + const last = active - 1; + const held = lines[cold]; + lines[cold] = lines[last]; + lines[last] = held; + active = last; + } + if (!cardHolds(rank_pinned, kv_rank, active, stride, workspace)) return error.GraphWorkspace; + return active; +} + +fn capture(window: u32, bindings: []const Binding) Capture { + return .{ .window = window, .bindings = bindings }; +} + +test "replay address change keeps the topology and drops argument equality" { + const fixed: u64 = 0x2000; + const a = [_]Binding{ + .{ .name = "expert_ids", .address = 0x1000, .replay = true }, + .{ .name = "dense_weights", .address = fixed, .replay = false }, + }; + const b = [_]Binding{ + .{ .name = "expert_ids", .address = 0x3000, .replay = true }, + .{ .name = "dense_weights", .address = fixed, .replay = false }, + }; + const left = capture(4, &a); + const right = capture(4, &b); + try std.testing.expect(sameTopology(left, right)); + try std.testing.expect(!argumentsEqual(left, right)); +} + +test "a new fixed address is a different topology" { + const a = [_]Binding{ + .{ .name = "expert_ids", .address = 0x1000, .replay = true }, + .{ .name = "dense_weights", .address = 0x2000, .replay = false }, + }; + const b = [_]Binding{ + .{ .name = "expert_ids", .address = 0x1000, .replay = true }, + .{ .name = "dense_weights", .address = 0x4000, .replay = false }, + }; + try std.testing.expect(!sameTopology(capture(4, &a), capture(4, &b))); + try std.testing.expect(!argumentsEqual(capture(4, &a), capture(4, &b))); +} + +test "a known replay name marked fixed is not a valid capture" { + const bad = [_]Binding{ + .{ .name = "expert_ids", .address = 0x1000, .replay = false }, + }; + try std.testing.expect(!valid(capture(1, &bad))); + try std.testing.expect(isReplayName("expert_ids")); +} + +test "the three replay names and the two fixed names" { + const good = [_]Binding{ + .{ .name = "expert_ids", .address = 1, .replay = true }, + .{ .name = "ple_rows", .address = 3, .replay = true }, + .{ .name = "kv_index", .address = 5, .replay = true }, + .{ .name = "dense_weights", .address = 6, .replay = false }, + .{ .name = "shared_weights", .address = 7, .replay = false }, + }; + try std.testing.expect(valid(capture(2, &good))); + for (replay_names) |name| try std.testing.expect(isReplayName(name)); + try std.testing.expect(!isReplayName("dense_weights")); + try std.testing.expect(!isReplayName("shared_weights")); + + const shared_replay = [_]Binding{ + .{ .name = "shared_weights", .address = 7, .replay = true }, + }; + try std.testing.expect(!valid(capture(2, &shared_replay))); + + const unknown = [_]Binding{ + .{ .name = "token_ids", .address = 9, .replay = true }, + }; + try std.testing.expect(valid(capture(2, &unknown))); + try std.testing.expect(!isReplayName("token_ids")); +} + +test "the pointer tables and picks are replay inputs; pool lines and scales are fixed" { + const t = [_]Binding{ + .{ .name = "expert_tables", .address = 0x100, .replay = true }, + .{ .name = "picks", .address = 0x200, .replay = true }, + .{ .name = "pool_lines", .address = 0x300, .replay = false }, + .{ .name = "expert_scales", .address = 0x400, .replay = false }, + }; + try std.testing.expect(valid(capture(1, &t))); + var bad = t; + bad[2].replay = true; + try std.testing.expect(!valid(capture(1, &bad))); + bad = t; + bad[0].replay = false; + try std.testing.expect(!valid(capture(1, &bad))); + // a pool that reallocated its lines is a new topology (recapture); new table contents are not + var moved = t; + moved[2].address += 256; + try std.testing.expect(!sameTopology(capture(1, &t), capture(1, &moved))); +} + +test "roundReady: an all-hit single-pass round replays straight away, a copy or a second pass stages first" { + const gpa = std.testing.allocator; + var slots: [4]u16 = undefined; + var line: [4]offload.Line = undefined; + var stamp: [4]u32 = undefined; + var lay: [4]u8 = undefined; + var v = offload.Vram.init(&slots, &line, &stamp, &lay); + var counts: [offload.shape.layers][offload.shape.experts]u32 = @splat(@splat(0)); + var r = try round_.Round.init(gpa, 4, 16); + defer r.deinit(gpa); + try round_.plan(&v, 0, &[_]i32{ 1, 2, 512 }, &counts, 1, 1, .{}, &r); + try std.testing.expect(!roundReady(false, &r)); + try round_.plan(&v, 0, &[_]i32{ 2, 1, 512 }, &counts, 2, 2, .{}, &r); + try std.testing.expect(roundReady(false, &r)); + try std.testing.expect(!roundReady(true, &r)); + try round_.plan(&v, 0, &[_]i32{ 1, 2, 3, 4, 5, 6 }, &counts, 3, 3, .{}, &r); + try std.testing.expect(r.masked() and !roundReady(false, &r)); +} + +test "slot_ptrs and hit_mask are not replay names" { + try std.testing.expect(!isReplayName("slot_ptrs")); + try std.testing.expect(!isReplayName("hit_mask")); + try std.testing.expect(isReplayName("ple_rows")); + try std.testing.expect(isReplayName("kv_index")); +} + +test "matching addresses are equal arguments" { + const rows = [_]Binding{ + .{ .name = "expert_ids", .address = 0x11, .replay = true }, + .{ .name = "slot_ptrs", .address = 0x22, .replay = true }, + .{ .name = "dense_weights", .address = 0x33, .replay = false }, + .{ .name = "shared_weights", .address = 0x44, .replay = false }, + }; + const again = [_]Binding{ + .{ .name = "expert_ids", .address = 0x11, .replay = true }, + .{ .name = "slot_ptrs", .address = 0x22, .replay = true }, + .{ .name = "dense_weights", .address = 0x33, .replay = false }, + .{ .name = "shared_weights", .address = 0x44, .replay = false }, + }; + try std.testing.expect(argumentsEqual(capture(1, &rows), capture(1, &again))); + + const other_window = capture(2, &again); + try std.testing.expect(!sameTopology(capture(1, &rows), other_window)); +} + +test "startup blocks capture when every decision is hit and every line is resident" { + const decisions = [_]offload.Decision{ .hit, .hit }; + const lines = [_]offload.Line{ .resident, .resident }; + try std.testing.expect(!captureReady(true, &decisions, &lines)); + try std.testing.expect(captureReady(false, &decisions, &lines)); +} + +test "a copy decision or an inflight line is not capture ready" { + const copy = [_]offload.Decision{ .hit, .copy, .cpu }; + const resident = [_]offload.Line{.resident}; + try std.testing.expect(!captureReady(false, ©, &resident)); + + const hits = [_]offload.Decision{.hit}; + const inflight = [_]offload.Line{ .resident, .inflight, .empty }; + try std.testing.expect(!captureReady(false, &hits, &inflight)); + + const cpu = [_]offload.Decision{.cpu}; + const quiet = [_]offload.Line{ .resident, .empty }; + try std.testing.expect(captureReady(false, &cpu, &quiet)); +} + +test "promote on a ready key is rewrite_ids, not recapture" { + const decisions = [_]offload.Decision{.hit}; + const resident = [_]offload.Line{.resident}; + const key: Key = .{ .rank = 2, .layer = 7, .window = 4 }; + try std.testing.expect(captureReady(false, &decisions, &resident)); + + const have = [_]Key{key}; + try std.testing.expectEqual(Promote.rewrite_ids, promote(&have, key)); + try std.testing.expect(promote(&have, key) != .recapture); + + const other_rank: Key = .{ .rank = 3, .layer = 7, .window = 4 }; + const other_layer: Key = .{ .rank = 2, .layer = 8, .window = 4 }; + const other_window: Key = .{ .rank = 2, .layer = 7, .window = 5 }; + try std.testing.expectEqual(Promote.recapture, promote(&have, other_rank)); + try std.testing.expectEqual(Promote.recapture, promote(&have, other_layer)); + try std.testing.expectEqual(Promote.recapture, promote(&have, other_window)); +} + +test "workspace at device_reserve keeps every line" { + var lines = [_]LineUse{ + .{ .layer = 1, .id = 2, .count = 1 }, + .{ .layer = 0, .id = 9, .count = 50 }, + }; + const stride: u64 = 8 << 30; + try std.testing.expect(2 * stride + device_reserve > card_bytes); + const kept = try fitWorkspace(0, 0, &lines, stride, device_reserve); + try std.testing.expectEqual(@as(usize, 2), kept); + try std.testing.expectEqual(@as(u8, 1), lines[0].layer); + try std.testing.expectEqual(@as(u16, 2), lines[0].id); + try std.testing.expectEqual(@as(u8, 0), lines[1].layer); +} + +test "workspace above device_reserve drops the coldest line first" { + var lines = [_]LineUse{ + .{ .layer = 2, .id = 1, .count = 10 }, + .{ .layer = 4, .id = 3, .count = 1 }, + .{ .layer = 1, .id = 0, .count = 7 }, + }; + const stride: u64 = 6 << 30; + const workspace = device_reserve + 1; + try std.testing.expect(3 * stride + workspace > card_bytes); + try std.testing.expect(2 * stride + workspace <= card_bytes); + const kept = try fitWorkspace(0, 0, &lines, stride, workspace); + try std.testing.expectEqual(@as(usize, 2), kept); + try std.testing.expectEqual(@as(u32, 10), lines[0].count); + try std.testing.expectEqual(@as(u8, 2), lines[0].layer); + try std.testing.expectEqual(@as(u32, 7), lines[1].count); + try std.testing.expectEqual(@as(u8, 1), lines[1].layer); +} + +test "equal counts drop the smaller layer first" { + var lines = [_]LineUse{ + .{ .layer = 5, .id = 1, .count = 4 }, + .{ .layer = 2, .id = 9, .count = 4 }, + }; + const stride: u64 = 8 << 30; + const workspace = device_reserve + 1; + try std.testing.expect(2 * stride + workspace > card_bytes); + try std.testing.expect(stride + workspace <= card_bytes); + const kept = try fitWorkspace(0, 0, &lines, stride, workspace); + try std.testing.expectEqual(@as(usize, 1), kept); + try std.testing.expectEqual(@as(u8, 5), lines[0].layer); + try std.testing.expectEqual(@as(u16, 1), lines[0].id); + try std.testing.expectEqual(@as(u32, 4), lines[0].count); +} + +test "pins and kv stay in the sum" { + const pin: u64 = 8 << 30; + const kv: u64 = 6 << 30; + const stride: u64 = 1 << 30; + const workspace = device_reserve + 1; + // Original pin and kv. Two lines miss the card; one line fits. Dropping either input would keep both. + try std.testing.expect(pin + kv + 2 * stride + workspace > card_bytes); + try std.testing.expect(pin + kv + stride + workspace <= card_bytes); + try std.testing.expect(2 * stride + workspace <= card_bytes); + var lines = [_]LineUse{ + .{ .layer = 1, .id = 1, .count = 3 }, + .{ .layer = 0, .id = 2, .count = 1 }, + }; + const kept = try fitWorkspace(pin, kv, &lines, stride, workspace); + try std.testing.expectEqual(@as(usize, 1), kept); + try std.testing.expectEqual(@as(u8, 1), lines[0].layer); + try std.testing.expectEqual(@as(u16, 1), lines[0].id); + try std.testing.expectEqual(@as(u32, 3), lines[0].count); +} + +test "zero lines past the card return GraphWorkspace" { + const pin: u64 = 10 << 30; + const kv: u64 = 6 << 30; + const workspace = device_reserve + 1; + try std.testing.expect(pin + kv + workspace > card_bytes); + var lines = [_]LineUse{}; + try std.testing.expectError(error.GraphWorkspace, fitWorkspace(pin, kv, &lines, 1, workspace)); +} diff --git a/zig/src/families/flashnext/exl3/graph/exl3_graph_budget.zig b/zig/src/families/flashnext/exl3/graph/exl3_graph_budget.zig new file mode 100644 index 0000000000..2a24a7355b --- /dev/null +++ b/zig/src/families/flashnext/exl3/graph/exl3_graph_budget.zig @@ -0,0 +1,206 @@ +const std = @import("std"); + +pub const Class = enum { decode, prefill, chain }; +pub const Charge = struct { + exec: u64, + template: u64, + transient: u64, + + pub fn peak(self: Charge) !u64 { + return std.math.add(u64, try std.math.add(u64, self.exec, self.template), self.transient); + } +}; + +pub fn Budget(comptime capacity: usize) type { + return struct { + const Self = @This(); + pub const Ticket = struct { slot: usize, generation: u64 }; + const Entry = struct { + generation: u64, + cohort: u64, + class: Class, + charge: Charge, + phase: enum { pending, retained, retiring } = .pending, + quiescent: bool = false, + used: u64, + }; + pub const Plan = struct { + victims: [capacity]Ticket = undefined, + len: usize = 0, + }; + + limit: u64, + headroom: u64, + entries: [capacity]?Entry = @splat(null), + generation: u64 = 0, + + pub fn init(reservation: u64, headroom: u64) !Self { + if (reservation <= headroom) return error.InvalidGraphReservation; + return .{ .limit = reservation - headroom, .headroom = headroom }; + } + + pub fn charged(self: *const Self) u64 { + var total: u64 = 0; + for (self.entries) |maybe| if (maybe) |entry| { + total += entry.charge.peak() catch unreachable; + }; + return total; + } + + pub fn pending(self: *const Self) u64 { + var total: u64 = 0; + for (self.entries) |maybe| if (maybe) |entry| { + if (entry.phase == .pending) total += entry.charge.peak() catch unreachable; + }; + return total; + } + + pub fn localCharged(self: *const Self) u64 { + var total: u64 = 0; + for (self.entries) |maybe| if (maybe) |entry| { + if (entry.cohort == 0) total += entry.charge.peak() catch unreachable; + }; + return total; + } + + pub fn plan(self: *const Self, charge: Charge) !Plan { + const needed = try charge.peak(); + if (charge.exec == 0 or needed > self.limit) return error.GraphBudgetExceeded; + var result: Plan = .{}; + var remaining = self.charged(); + var free_slot = false; + var selected: [capacity]bool = @splat(false); + for (self.entries) |entry| free_slot = free_slot or entry == null; + while (!free_slot or needed > self.limit - remaining) { + var oldest: ?usize = null; + for (self.entries, 0..) |maybe, index| if (maybe) |entry| { + if (selected[index] or entry.phase != .retained or !entry.quiescent) continue; + if (oldest == null or entry.used < self.entries[oldest.?].?.used) oldest = index; + }; + const index = oldest orelse return error.GraphBudgetBusy; + const entry = self.entries[index].?; + result.victims[result.len] = .{ .slot = index, .generation = entry.generation }; + result.len += 1; + selected[index] = true; + remaining -= try entry.charge.peak(); + free_slot = true; + } + return result; + } + + pub fn reserve(self: *Self, class: Class, cohort: u64, charge: Charge, used: u64) !Ticket { + const proposed = try self.plan(charge); + if (proposed.len != 0) return error.GraphEvictionRequired; + const generation = try std.math.add(u64, self.generation, 1); + for (&self.entries, 0..) |*entry, index| if (entry.* == null) { + entry.* = .{ .generation = generation, .cohort = cohort, .class = class, .charge = charge, .used = used }; + self.generation = generation; + return .{ .slot = index, .generation = generation }; + }; + unreachable; + } + + fn get(self: *Self, ticket: Ticket) !*Entry { + if (ticket.slot >= capacity) return error.StaleGraphTicket; + const entry = &(self.entries[ticket.slot] orelse return error.StaleGraphTicket); + if (entry.generation != ticket.generation) return error.StaleGraphTicket; + return &self.entries[ticket.slot].?; + } + + pub fn installed(self: *Self, ticket: Ticket) !void { + const entry = try self.get(ticket); + if (entry.phase != .pending) return error.InvalidGraphPhase; + entry.phase = .retained; + entry.charge.transient = 0; + } + + pub fn getForRelease(self: *Self, ticket: Ticket) !bool { + return (try self.get(ticket)).phase == .retained; + } + + pub fn cohortOf(self: *Self, ticket: Ticket) !u64 { + return (try self.get(ticket)).cohort; + } + + pub fn touch(self: *Self, ticket: Ticket, used: u64, quiescent: bool) !void { + const entry = try self.get(ticket); + if (entry.phase != .retained) return error.InvalidGraphPhase; + entry.used = used; + entry.quiescent = quiescent; + } + + pub fn beginRetire(self: *Self, ticket: Ticket, agreed_cohort: u64) !void { + const entry = try self.get(ticket); + if (entry.cohort != agreed_cohort) return error.GraphCohortMismatch; + if (entry.phase != .retained or !entry.quiescent) return error.GraphBudgetBusy; + entry.phase = .retiring; + } + + pub fn destroyed(self: *Self, ticket: Ticket) !void { + const entry = try self.get(ticket); + if (entry.phase != .retiring and entry.phase != .pending) return error.InvalidGraphPhase; + self.entries[ticket.slot] = null; + } + }; +} + +test "exl3: graph admission reserves pending peak and explicit headroom" { + var budget = try Budget(4).init(100, 20); + const ticket = try budget.reserve(.decode, 7, .{ .exec = 30, .template = 10, .transient = 20 }, 1); + try std.testing.expectEqual(@as(u64, 60), budget.charged()); + try std.testing.expectError(error.GraphBudgetBusy, budget.plan(.{ .exec = 30, .template = 0, .transient = 0 })); + try budget.installed(ticket); + try std.testing.expectEqual(@as(u64, 40), budget.charged()); + try std.testing.expectError(error.GraphBudgetExceeded, budget.plan(.{ .exec = 81, .template = 0, .transient = 0 })); +} + +test "exl3: pre eviction requires quiescence cohort agreement and completed destruction" { + var budget = try Budget(2).init(120, 20); + const first = try budget.reserve(.prefill, 11, .{ .exec = 40, .template = 10, .transient = 0 }, 1); + try budget.installed(first); + const second = try budget.reserve(.chain, 12, .{ .exec = 40, .template = 10, .transient = 0 }, 2); + try budget.installed(second); + try std.testing.expectError(error.GraphBudgetBusy, budget.plan(.{ .exec = 50, .template = 0, .transient = 0 })); + try budget.touch(first, 1, true); + const plan = try budget.plan(.{ .exec = 50, .template = 0, .transient = 0 }); + try std.testing.expectEqual(@as(usize, 1), plan.len); + try std.testing.expectEqual(first, plan.victims[0]); + try std.testing.expectError(error.GraphCohortMismatch, budget.beginRetire(first, 12)); + try budget.beginRetire(first, 11); + try std.testing.expectEqual(@as(u64, 100), budget.charged()); + try std.testing.expectError(error.GraphBudgetBusy, budget.reserve(.decode, 13, .{ .exec = 50, .template = 0, .transient = 0 }, 3)); + try budget.destroyed(first); + const replacement = try budget.reserve(.decode, 13, .{ .exec = 50, .template = 0, .transient = 0 }, 3); + try std.testing.expectError(error.StaleGraphTicket, budget.destroyed(first)); + try budget.destroyed(replacement); + try std.testing.expectEqual(@as(u64, 50), budget.charged()); +} + +test "exl3: graph estimates reject overflow and invalid reservations" { + try std.testing.expectError(error.InvalidGraphReservation, Budget(1).init(10, 10)); + var budget = try Budget(1).init(100, 10); + try std.testing.expectError(error.Overflow, budget.plan(.{ .exec = std.math.maxInt(u64), .template = 1, .transient = 0 })); + try std.testing.expectError(error.GraphBudgetExceeded, budget.plan(.{ .exec = 0, .template = 0, .transient = 0 })); +} + +test "exl3: cross class LRU planning is read only and launch revokes retirement eligibility" { + var budget = try Budget(3).init(130, 10); + const charge: Charge = .{ .exec = 30, .template = 10, .transient = 0 }; + const decode = try budget.reserve(.decode, 1, charge, 3); + const prefill = try budget.reserve(.prefill, 2, charge, 1); + const chain = try budget.reserve(.chain, 3, charge, 2); + for ([_]Budget(3).Ticket{ decode, prefill, chain }, [_]u64{ 3, 1, 2 }) |ticket, used| { + try budget.installed(ticket); + try budget.touch(ticket, used, true); + } + const request: Charge = .{ .exec = 70, .template = 0, .transient = 0 }; + const plan = try budget.plan(request); + try std.testing.expectEqual(@as(usize, 2), plan.len); + try std.testing.expectEqual(prefill, plan.victims[0]); + try std.testing.expectEqual(chain, plan.victims[1]); + try std.testing.expectEqual(@as(u64, 120), budget.charged()); + try std.testing.expectError(error.GraphEvictionRequired, budget.reserve(.decode, 4, request, 4)); + try budget.touch(prefill, 4, false); + try std.testing.expectError(error.GraphBudgetBusy, budget.beginRetire(prefill, 2)); + try std.testing.expectError(error.InvalidGraphPhase, budget.destroyed(chain)); +} diff --git a/zig/src/families/flashnext/exl3/kv/cuda_exl3_kvtier.zig b/zig/src/families/flashnext/exl3/kv/cuda_exl3_kvtier.zig index 49e9a88e50..b5e8b351f9 100644 --- a/zig/src/families/flashnext/exl3/kv/cuda_exl3_kvtier.zig +++ b/zig/src/families/flashnext/exl3/kv/cuda_exl3_kvtier.zig @@ -96,6 +96,7 @@ pub const LaneSel = struct { head_on: bool = false, // the MTP head has written (or will write) rows of this stream hw: usize = 0, // pages [0, hw) have been flushed once: their host rows are valid reset_pending: bool = true, + graph_staged: bool = false, pub fn init(gpa: std.mem.Allocator, d: *const cuda.Driver, k: *const Kernels, cap: usize, npages: usize, slot_frames: usize, layers: usize, own: Own) !*LaneSel { const s = try gpa.create(LaneSel); @@ -173,6 +174,7 @@ pub const LaneSel = struct { /// Layer `ai` writes [pos0, pos0 + n): a page leaves the window once every layer wrote it; a rewind brings it back. pub fn begin(s: *LaneSel, stream: cuda.Stream, ai: usize, pos0: usize, n: usize) !void { + if (s.graph_staged) return; if (ai >= s.layers) return; if (s.reset_pending) { // a new, resumed or restored stream: rows before pos0 are canonical in the host store diff --git a/zig/src/families/flashnext/exl3/kv/cuda_exl3_state_binding.zig b/zig/src/families/flashnext/exl3/kv/cuda_exl3_state_binding.zig new file mode 100644 index 0000000000..0b49ba199b --- /dev/null +++ b/zig/src/families/flashnext/exl3/kv/cuda_exl3_state_binding.zig @@ -0,0 +1,239 @@ +const std = @import("std"); +const cuda = @import("cuda"); +const binding = @import("exl3_state_binding.zig"); + +pub const NodeType = *const fn (cuda.graph.Node, *c_uint) callconv(.c) cuda.abi.Result; + +pub fn nodeType(driver: *const cuda.Driver) ?NodeType { + var library = driver.lib; + return library.lookup(NodeType, "cuGraphNodeGetType"); +} + +pub const Backend = struct { + graph: cuda.graph.Graph, + exec: cuda.graph.Exec, + node_type: NodeType, + + pub fn init(graph: cuda.graph.Graph, exec: cuda.graph.Exec) ?Backend { + return initWithNodeType(graph, exec, nodeType(graph.d)); + } + + pub fn initWithNodeType(graph: cuda.graph.Graph, exec: cuda.graph.Exec, node_type: ?NodeType) ?Backend { + return .{ .graph = graph, .exec = exec, .node_type = node_type orelse return null }; + } + + pub fn nodes(self: Backend, template: usize) ![]usize { + if (template != @intFromPtr(self.graph.handle)) return error.UnsupportedBinding; + const count = try self.graph.nodeCount(); + const storage = try std.heap.c_allocator.alloc(cuda.graph.Node, count); + defer std.heap.c_allocator.free(storage); + const actual = try self.graph.nodes(storage); + if (actual.len != count) return error.UnsupportedBinding; + const identifiers = try std.heap.c_allocator.alloc(usize, count); + for (actual, identifiers) |node, *identifier| identifier.* = @intFromPtr(node); + return identifiers; + } + + pub fn freeNodes(_: Backend, nodes_list: []usize) void { + std.heap.c_allocator.free(nodes_list); + } + + pub fn kernel(self: Backend, node: usize) !?binding.Header { + var kind: c_uint = undefined; + try self.graph.d.check(self.node_type(@ptrFromInt(node), &kind), "cuGraphNodeGetType(state binding)"); + if (kind == 4) return error.UnsupportedBinding; + if (kind != 0) return null; + var params: cuda.abi.KernelNodeParams = undefined; + try self.graph.d.check(self.graph.d.api.cuGraphKernelNodeGetParams_v2(@ptrFromInt(node), ¶ms), "cuGraphKernelNodeGetParams(state binding)"); + return .{ + .function = @intFromPtr(params.func), + .geometry = .{ .grid = .{ params.grid_x, params.grid_y, params.grid_z }, .block = .{ params.block_x, params.block_y, params.block_z }, .shared = params.shared_bytes }, + .params = params.params, + .extra = params.extra != null, + .alternate_kernel = params.func == null, + }; + } + + pub fn set(self: Backend, node: usize, function: usize, geometry: binding.Geometry, parameters: binding.Parameters) !void { + var args: cuda.Args = .{}; + args.add(parameters.nonstate.gq); + args.add(parameters.nonstate.gk); + args.add(parameters.nonstate.gv); + args.add(parameters.nonstate.ggb); + args.add(parameters.state_in); + if (parameters.state_out != 0) args.add(parameters.state_out); + args.add(parameters.nonstate.rows); + args.add(parameters.nonstate.ys); + try self.exec.setKernel(@ptrFromInt(node), .{ .d = self.exec.d, .handle = @ptrFromInt(function) }, .{ + .grid = .{ .x = geometry.grid[0], .y = geometry.grid[1], .z = geometry.grid[2] }, + .block = .{ .x = geometry.block[0], .y = geometry.block[1], .z = geometry.block[2] }, + .shared = geometry.shared, + }, &args); + } +}; + +test "exl3: state binding adapter declines absent type query and propagates query errors" { + std.testing.refAllDecls(Backend); + const Fake = struct { + var getter_calls: usize = 0; + var kind: c_uint = 1; + var result: cuda.abi.Result = 0; + + fn nodeType(_: cuda.graph.Node, output: *c_uint) callconv(.c) cuda.abi.Result { + output.* = kind; + return result; + } + + fn get(_: cuda.graph.Node, _: *cuda.abi.KernelNodeParams) callconv(.c) cuda.abi.Result { + getter_calls += 1; + return cuda.abi.error_not_ready; + } + }; + var driver: cuda.Driver = undefined; + driver.api.cuGraphKernelNodeGetParams_v2 = Fake.get; + const graph: cuda.graph.Graph = .{ .d = &driver, .handle = @ptrFromInt(42) }; + const exec: cuda.graph.Exec = .{ .d = &driver, .handle = @ptrFromInt(43) }; + try std.testing.expect(Backend.initWithNodeType(graph, exec, null) == null); + const backend = Backend.initWithNodeType(graph, exec, Fake.nodeType).?; + Fake.getter_calls = 0; + Fake.kind = 1; + Fake.result = 0; + try std.testing.expect(try backend.kernel(1) == null); + try std.testing.expectEqual(@as(usize, 0), Fake.getter_calls); + Fake.result = cuda.abi.error_not_ready; + try std.testing.expectError(error.NotReady, backend.kernel(1)); + try std.testing.expectEqual(@as(usize, 0), Fake.getter_calls); + Fake.result = 0; + Fake.kind = 0; + try std.testing.expectError(error.NotReady, backend.kernel(1)); + try std.testing.expectEqual(@as(usize, 1), Fake.getter_calls); + Fake.kind = 4; + try std.testing.expectError(error.UnsupportedBinding, backend.kernel(1)); +} + +test "exl3: state binding adapter setters retain exact schema7 and schema8 types" { + const Fake = struct { + var output: bool = false; + var captured: binding.Parameters = undefined; + var geometry: binding.Geometry = undefined; + var function: usize = 0; + var node: usize = 0; + + fn value(comptime Value: type, params: [*]?*anyopaque, index: usize) Value { + var copied: Value = undefined; + @memcpy(std.mem.asBytes(&copied), @as([*]const u8, @ptrCast(params[index].?))[0..@sizeOf(Value)]); + return copied; + } + + fn set(_: cuda.abi.GraphExec, target: cuda.graph.Node, params: *const cuda.abi.KernelNodeParams) callconv(.c) cuda.abi.Result { + const args = params.params.?; + captured = .{ + .nonstate = .{ + .gq = value(u64, args, 0), + .gk = value(u64, args, 1), + .gv = value(u64, args, 2), + .ggb = value(u64, args, 3), + .rows = value(c_int, args, if (output) 6 else 5), + .ys = value(u64, args, if (output) 7 else 6), + }, + .state_in = value(u64, args, 4), + .state_out = if (output) value(u64, args, 5) else 0, + }; + geometry = .{ .grid = .{ params.grid_x, params.grid_y, params.grid_z }, .block = .{ params.block_x, params.block_y, params.block_z }, .shared = params.shared_bytes }; + function = @intFromPtr(params.func); + node = @intFromPtr(target); + return 0; + } + + fn nodeType(_: cuda.graph.Node, _: *c_uint) callconv(.c) cuda.abi.Result { + return 0; + } + }; + var driver: cuda.Driver = undefined; + driver.api.cuGraphExecKernelNodeSetParams_v2 = Fake.set; + const graph: cuda.graph.Graph = .{ .d = &driver, .handle = @ptrFromInt(42) }; + const exec: cuda.graph.Exec = .{ .d = &driver, .handle = @ptrFromInt(43) }; + const backend = Backend.initWithNodeType(graph, exec, Fake.nodeType).?; + const geometry: binding.Geometry = .{ .grid = .{ 16, 8, 1 }, .block = .{ 128, 1, 1 } }; + for (1..binding.max_width + 1) |width| for ([_]bool{ false, true }) |output| for ([_]bool{ false, true }) |flipped| { + Fake.output = output; + const parameters: binding.Parameters = .{ + .nonstate = .{ .gq = 0x100000001, .gk = 0x200000002, .gv = 0x300000003, .ggb = 0x400000004, .rows = @intCast(width), .ys = 0x500000005 }, + .state_in = if (flipped) 2000 else 1000, + .state_out = if (!output) 0 else if (flipped) 1000 else 2000, + }; + try backend.set(9, if (output) 11 else 10, geometry, parameters); + try std.testing.expectEqualDeep(parameters, Fake.captured); + try std.testing.expectEqualDeep(geometry, Fake.geometry); + try std.testing.expectEqual(@as(usize, if (output) 11 else 10), Fake.function); + try std.testing.expectEqual(@as(usize, 9), Fake.node); + }; +} + +test "exl3: state binding whitelisted CUfunction ignores ctx and kern fields" { + const Fake = struct { + var args: cuda.Args = .{}; + + fn nodeType(_: cuda.graph.Node, output: *c_uint) callconv(.c) cuda.abi.Result { + output.* = 0; + return 0; + } + + fn get(_: cuda.graph.Node, params: *cuda.abi.KernelNodeParams) callconv(.c) cuda.abi.Result { + params.* = .{ + .func = @ptrFromInt(11), + .grid_x = 16, + .grid_y = 8, + .grid_z = 1, + .block_x = 128, + .block_y = 1, + .block_z = 1, + .shared_bytes = 0, + .params = args.pointers(), + .extra = null, + .kern = @ptrFromInt(20), + .ctx = @ptrFromInt(30), + }; + return 0; + } + }; + var driver: cuda.Driver = undefined; + driver.api.cuGraphKernelNodeGetParams_v2 = Fake.get; + const graph: cuda.graph.Graph = .{ .d = &driver, .handle = @ptrFromInt(42) }; + const exec: cuda.graph.Exec = .{ .d = &driver, .handle = @ptrFromInt(43) }; + const backend = Backend.initWithNodeType(graph, exec, Fake.nodeType).?; + Fake.args = .{}; + for ([_]u64{ 1, 2, 3, 4, 100, 200 }) |value| Fake.args.add(value); + Fake.args.add(@as(c_int, 2)); + Fake.args.add(@as(u64, 5)); + const header = (try backend.kernel(1)).?; + try std.testing.expectEqual(@as(usize, 11), header.function); + try std.testing.expect(!header.alternate_kernel); + const Source = struct { + backend: Backend, + identifiers: [1]usize = .{1}, + + pub fn nodes(self: *@This(), _: usize) ![]usize { + return &self.identifiers; + } + + pub fn freeNodes(_: *@This(), _: []usize) void {} + + pub fn kernel(self: *@This(), node: usize) !?binding.Header { + return self.backend.kernel(node); + } + }; + var request: binding.Request = .{ + .width = 2, + .chain = 10, + .chain_to = 11, + .geometry = .{ .grid = .{ 16, 8, 1 }, .block = .{ 128, 1, 1 } }, + .nonstate = .{ .gq = 1, .gk = 2, .gv = 3, .ggb = 4, .rows = 2, .ys = 5 }, + .count = 1, + }; + request.layers[0] = .{ .layer = 0, .state_in = 100, .state_out = 200 }; + var state = try binding.Binding.init(42, request); + var source: Source = .{ .backend = backend }; + try state.load(&source); + try std.testing.expect(state.loaded); +} diff --git a/zig/src/families/flashnext/exl3/kv/cuda_exl3_state_request_test.zig b/zig/src/families/flashnext/exl3/kv/cuda_exl3_state_request_test.zig new file mode 100644 index 0000000000..8ec2ccaca5 --- /dev/null +++ b/zig/src/families/flashnext/exl3/kv/cuda_exl3_state_request_test.zig @@ -0,0 +1,72 @@ +const std = @import("std"); +const cuda = @import("cuda"); +const wt = @import("../../cuda_prompt_wt.zig"); +const batch = @import("../../cuda_batch.zig"); +const mtp = @import("../../cuda_mtp.zig"); +const binding = @import("exl3_state_binding.zig"); + +pub fn currentPass(comptime Key: type, comptime windowRequest: anytype) !void { + try std.testing.expectEqual(mtp.max_rows, binding.max_width); + var pass: batch.Pass(struct {}) = undefined; + var kernels: batch.Kernels = undefined; + kernels.chain = .{ .d = undefined, .handle = @ptrFromInt(10) }; + kernels.chain_to = .{ .d = undefined, .handle = @ptrFromInt(11) }; + var scratch: batch.Scratch = undefined; + scratch.gq = 1; + scratch.gk = 2; + scratch.gv = 3; + scratch.ggb = 4; + scratch.ys = 5; + var states: [wt.linear_n]cuda.DeviceBuffer = undefined; + var layers: [wt.layers_n]@import("../../cuda_prompt_load.zig").Layer = undefined; + var keep: batch.Keep = .{}; + for (&layers, 0..) |*layer, index| layer.linear = wt.linearLayer(index); + pass.lanes = &.{}; + pass.k = &kernels; + pass.s = &scratch; + pass.seq.state = &states; + pass.m.layers = &layers; + pass.hooks.keep = &keep; + var key: Key = .{ .kind = 2, .rows = 2, .scratch = 1, .scores = 2, .output = 3 }; + for (2..mtp.max_rows + 1) |width| { + key.rows = width; + for ([_]bool{ false, true }) |flipped| { + for (&states, 0..) |*state, index| state.ptr = (if (flipped) @as(u64, 2000) else 1000) + index; + for (&layers, 0..) |layer, index| { + if (layer.linear) keep.layers[index].state_out = (if (flipped) @as(u64, 1000) else 2000) + wt.linearIndex(index); + } + const maybe = windowRequest(key, &pass); + try std.testing.expect(maybe != null); + const request = maybe.?; + try std.testing.expectEqual(width, request.width); + try std.testing.expectEqual(@as(c_int, @intCast(width)), request.nonstate.rows); + try std.testing.expectEqual(wt.linear_n, request.count); + for (request.layers[0..request.count], 0..) |mapping, index| { + try std.testing.expectEqual(index, wt.linearIndex(mapping.layer)); + try std.testing.expectEqual(states[index].ptr, mapping.state_in); + try std.testing.expectEqual(keep.layers[mapping.layer].state_out, mapping.state_out); + } + } + } + for ([_]usize{ 0, 1, mtp.max_rows + 1, batch.att_rows + 1 }) |width| { + key.rows = width; + try std.testing.expect(windowRequest(key, &pass) == null); + } + key.rows = 3; + key.kind = 1; + try std.testing.expect(windowRequest(key, &pass) == null); + key.kind = 2; + const lanes = [_]batch.Lane{undefined}; + pass.lanes = &lanes; + try std.testing.expect(windowRequest(key, &pass) == null); + pass.lanes = &.{}; + pass.seq.state = states[0 .. wt.linear_n - 1]; + try std.testing.expect(windowRequest(key, &pass) == null); + pass.seq.state = &states; + keep.layers[0].state_out = states[0].ptr; + try std.testing.expect(windowRequest(key, &pass) == null); + keep.layers[0].state_out = 0; + try std.testing.expect(windowRequest(key, &pass) == null); + pass.hooks.keep = null; + try std.testing.expect(windowRequest(key, &pass) == null); +} diff --git a/zig/src/families/flashnext/exl3/kv/exl3_state_binding.zig b/zig/src/families/flashnext/exl3/kv/exl3_state_binding.zig new file mode 100644 index 0000000000..32b6151d69 --- /dev/null +++ b/zig/src/families/flashnext/exl3/kv/exl3_state_binding.zig @@ -0,0 +1,177 @@ +const std = @import("std"); + +pub const max_layers = 48; +pub const max_width = @import("../../cuda_mtp.zig").max_rows; + +pub fn supportsWindow(width: usize) bool { + return width >= 2 and width <= max_width; +} + +pub const Mapping = struct { layer: usize, state_in: u64, state_out: u64 }; +pub const Geometry = struct { + grid: [3]u32, + block: [3]u32, + shared: u32 = 0, +}; +pub const Nonstate = struct { gq: u64, gk: u64, gv: u64, ggb: u64, rows: c_int, ys: u64 }; +pub const Request = struct { + width: usize, + chain: usize, + chain_to: usize, + geometry: Geometry, + nonstate: Nonstate, + layers: [max_layers]Mapping = undefined, + count: usize = 0, + + pub fn validate(self: Request) !void { + if (self.width < 1 or self.width > max_width or self.nonstate.rows != self.width or + self.chain == 0 or self.chain_to == 0 or self.chain == self.chain_to or + self.count == 0 or self.count > max_layers) return error.UnsupportedBinding; + for (self.layers[0..self.count], 0..) |mapping, index| { + if (mapping.state_in == 0 or mapping.state_in == mapping.state_out) return error.UnsupportedBinding; + for (self.layers[0..index]) |previous| { + if (previous.layer == mapping.layer or previous.state_in == mapping.state_in or + previous.state_out == mapping.state_in or + (mapping.state_out != 0 and (previous.state_in == mapping.state_out or previous.state_out == mapping.state_out))) + return error.UnsupportedBinding; + } + } + } + + pub fn compatible(self: Request, target: Request) bool { + target.validate() catch return false; + if (self.width != target.width or self.chain != target.chain or self.chain_to != target.chain_to or + !std.meta.eql(self.geometry, target.geometry) or !std.meta.eql(self.nonstate, target.nonstate) or + self.count != target.count) return false; + for (self.layers[0..self.count], target.layers[0..target.count]) |before, after| { + if (before.layer != after.layer) return false; + const same = before.state_in == after.state_in and before.state_out == after.state_out; + const flipped = before.state_out != 0 and before.state_in == after.state_out and before.state_out == after.state_in; + if (!same and !flipped) return false; + } + return true; + } +}; + +pub const Header = struct { + function: usize, + geometry: Geometry, + params: ?[*]?*anyopaque, + extra: bool = false, + alternate_kernel: bool = false, +}; + +pub const Parameters = struct { + nonstate: Nonstate, + state_in: u64, + state_out: u64, +}; + +fn argument(comptime Value: type, params: [*]?*anyopaque, index: usize) !Value { + const pointer = params[index] orelse return error.UnsupportedBinding; + var value: Value = undefined; + @memcpy(std.mem.asBytes(&value), @as([*]const u8, @ptrCast(pointer))[0..@sizeOf(Value)]); + return value; +} + +pub const Binding = struct { + template: usize, + request: Request, + nodes: [max_layers]usize = undefined, + parameters: [max_layers]Parameters = undefined, + loaded: bool = false, + poisoned: bool = false, + + pub fn init(template: usize, request: Request) !Binding { + try request.validate(); + if (template == 0) return error.UnsupportedBinding; + return .{ .template = template, .request = request }; + } + + pub fn invalidate(self: *Binding) void { + self.loaded = false; + self.poisoned = true; + self.template = 0; + } + + pub fn load(self: *Binding, backend: anytype) !void { + if (self.poisoned or self.template == 0) return error.StateBindingPoisoned; + if (self.loaded) return; + var matched: [max_layers]bool = @splat(false); + const nodes = try backend.nodes(self.template); + defer backend.freeNodes(nodes); + for (nodes) |node| { + const header = (try backend.kernel(node)) orelse continue; + if (header.function != self.request.chain and header.function != self.request.chain_to) continue; + if (header.extra or header.alternate_kernel or !std.meta.eql(header.geometry, self.request.geometry)) return error.UnsupportedBinding; + const params = header.params orelse return error.UnsupportedBinding; + const output = header.function == self.request.chain_to; + const values: Parameters = .{ + .nonstate = .{ + .gq = try argument(u64, params, 0), + .gk = try argument(u64, params, 1), + .gv = try argument(u64, params, 2), + .ggb = try argument(u64, params, 3), + .rows = try argument(c_int, params, if (output) 6 else 5), + .ys = try argument(u64, params, if (output) 7 else 6), + }, + .state_in = try argument(u64, params, 4), + .state_out = if (output) try argument(u64, params, 5) else 0, + }; + if (!std.meta.eql(values.nonstate, self.request.nonstate)) return error.UnsupportedBinding; + var found: ?usize = null; + for (self.request.layers[0..self.request.count], 0..) |mapping, index| { + if (values.state_in == mapping.state_in and values.state_out == mapping.state_out) { + if (matched[index] or found != null) return error.UnsupportedBinding; + found = index; + } + } + const index = found orelse return error.UnsupportedBinding; + matched[index] = true; + self.nodes[index] = node; + self.parameters[index] = values; + } + for (matched[0..self.request.count]) |present| if (!present) return error.UnsupportedBinding; + self.loaded = true; + } + + pub fn apply(self: *Binding, backend: anytype, target: Request) !void { + if (self.poisoned) return error.StateBindingPoisoned; + if (!self.loaded or !self.request.compatible(target)) return error.UnsupportedBinding; + errdefer self.poisoned = true; + for (target.layers[0..target.count], 0..) |mapping, index| { + var parameters = self.parameters[index]; + parameters.state_in = mapping.state_in; + parameters.state_out = mapping.state_out; + try backend.set(self.nodes[index], if (mapping.state_out == 0) target.chain else target.chain_to, target.geometry, parameters); + } + } + + pub fn applyTogether(self: *Binding, backend: anytype, target: Request, peer: anytype, signature: u64, cohort: u64) !void { + var succeeded = true; + self.apply(backend, target) catch { + succeeded = false; + }; + errdefer self.poisoned = true; + const agreement = try peer.graphAgree(signature, .{ .state = if (succeeded) .ready else .failed, .cohort = cohort }); + if (!succeeded or agreement.choice != .replay or agreement.cohort != cohort) return error.StateRebindFailed; + } +}; + +pub fn normalized(key: anytype) @TypeOf(key) { + var candidate = key; + candidate.seq[0][4] = 0; + candidate.keep = key.keep_nonstate; + return candidate; +} + +pub fn keepHash(keep: anytype) u64 { + var hash = std.hash.Wyhash.init(0x73746174656b6565); + for (keep.layers) |layer| { + inline for (@typeInfo(@TypeOf(layer)).@"struct".field_names) |field| { + if (comptime !std.mem.eql(u8, field, "state_out")) std.hash.autoHash(&hash, @field(layer, field)); + } + } + std.hash.autoHash(&hash, keep.mixer); + return hash.final(); +} diff --git a/zig/src/families/flashnext/exl3/kv/exl3_state_binding_test.zig b/zig/src/families/flashnext/exl3/kv/exl3_state_binding_test.zig new file mode 100644 index 0000000000..bcaed32840 --- /dev/null +++ b/zig/src/families/flashnext/exl3/kv/exl3_state_binding_test.zig @@ -0,0 +1,298 @@ +const std = @import("std"); +const binding = @import("exl3_state_binding.zig"); + +fn requestFor(width: usize, output: bool) binding.Request { + var request: binding.Request = .{ + .width = width, + .chain = 10, + .chain_to = 11, + .geometry = .{ .grid = .{ 16, 8, 1 }, .block = .{ 128, 1, 1 } }, + .nonstate = .{ .gq = 100, .gk = 200, .gv = 300, .ggb = 400, .rows = @intCast(width), .ys = 500 }, + .count = 2, + }; + request.layers[0] = .{ .layer = 0, .state_in = 1000, .state_out = if (output) 2000 else 0 }; + request.layers[1] = .{ .layer = 2, .state_in = 3000, .state_out = if (output) 4000 else 0 }; + return request; +} + +fn flip(request: binding.Request) binding.Request { + var target = request; + for (target.layers[0..target.count]) |*mapping| std.mem.swap(u64, &mapping.state_in, &mapping.state_out); + return target; +} + +const FakeNode = struct { + function: usize, + geometry: binding.Geometry, + values: binding.Parameters, + pointers: [8]?*anyopaque = @splat(null), + kernel: bool = true, + extra: bool = false, + + fn header(self: *FakeNode) binding.Header { + self.pointers[0] = &self.values.nonstate.gq; + self.pointers[1] = &self.values.nonstate.gk; + self.pointers[2] = &self.values.nonstate.gv; + self.pointers[3] = &self.values.nonstate.ggb; + self.pointers[4] = &self.values.state_in; + const output = self.values.state_out != 0; + if (output) self.pointers[5] = &self.values.state_out; + self.pointers[if (output) 6 else 5] = &self.values.nonstate.rows; + self.pointers[if (output) 7 else 6] = &self.values.nonstate.ys; + return .{ .function = self.function, .geometry = self.geometry, .params = &self.pointers, .extra = self.extra }; + } +}; + +const Fake = struct { + template: usize = 42, + records: [4]FakeNode = undefined, + ids: [4]usize = .{ 0, 1, 2, 3 }, + count: usize = 2, + calls: usize = 0, + fail_at: ?usize = null, + query_failure: bool = false, + writes: [2]binding.Parameters = undefined, + + fn init(request: binding.Request) Fake { + var fake: Fake = .{}; + for (request.layers[0..request.count], 0..) |mapping, index| { + fake.records[index] = .{ + .function = if (mapping.state_out == 0) request.chain else request.chain_to, + .geometry = request.geometry, + .values = .{ .nonstate = request.nonstate, .state_in = mapping.state_in, .state_out = mapping.state_out }, + }; + } + return fake; + } + + pub fn nodes(self: *Fake, template: usize) ![]usize { + if (template != self.template) return error.WrongTemplate; + return self.ids[0..self.count]; + } + + pub fn freeNodes(_: *Fake, _: []usize) void {} + + pub fn kernel(self: *Fake, node: usize) !?binding.Header { + if (self.query_failure) return error.PriorAsyncFailure; + if (!self.records[node].kernel) return null; + return self.records[node].header(); + } + + pub fn set(self: *Fake, node: usize, function: usize, geometry: binding.Geometry, parameters: binding.Parameters) !void { + const call = self.calls; + self.calls += 1; + if (self.fail_at == call) return error.SetterFailed; + try std.testing.expectEqual(self.records[node].function, function); + try std.testing.expectEqualDeep(self.records[node].geometry, geometry); + self.writes[node] = parameters; + } +}; + +test "exl3: state binding bounded widths and orientations preserve typed nonstate values" { + for (1..binding.max_width + 1) |width| { + const request = requestFor(width, true); + var fake = Fake.init(request); + var state = try binding.Binding.init(fake.template, request); + try state.load(&fake); + const target = flip(request); + try state.apply(&fake, target); + for (target.layers[0..target.count], fake.writes) |mapping, written| { + try std.testing.expectEqualDeep(request.nonstate, written.nonstate); + try std.testing.expectEqual(mapping.state_in, written.state_in); + try std.testing.expectEqual(mapping.state_out, written.state_out); + } + fake.records[0].values.nonstate.rows = 99; + try state.apply(&fake, request); + for (request.layers[0..request.count], fake.writes) |mapping, written| { + try std.testing.expectEqualDeep(request.nonstate, written.nonstate); + try std.testing.expectEqual(mapping.state_in, written.state_in); + try std.testing.expectEqual(mapping.state_out, written.state_out); + } + try std.testing.expectEqual(@as(usize, 4), fake.calls); + } +} + +test "exl3: state binding schema7 reads c_int rows and u64 ys without a state output" { + for (1..binding.max_width + 1) |width| { + const request = requestFor(width, false); + var fake = Fake.init(request); + var state = try binding.Binding.init(fake.template, request); + try state.load(&fake); + try state.apply(&fake, request); + try std.testing.expectEqualDeep(request.nonstate, fake.writes[0].nonstate); + var changed = request; + changed.layers[0].state_in += 1; + try std.testing.expectError(error.UnsupportedBinding, state.apply(&fake, changed)); + try std.testing.expectEqual(@as(usize, 2), fake.calls); + } +} + +test "exl3: state binding wrong functions are never inspected and missing duplicate mappings decline" { + const request = requestFor(2, true); + var fake = Fake.init(request); + fake.records[0].function = 999; + fake.records[0].extra = true; + var state = try binding.Binding.init(fake.template, request); + try std.testing.expectError(error.UnsupportedBinding, state.load(&fake)); + fake = Fake.init(request); + fake.count = 1; + try std.testing.expectError(error.UnsupportedBinding, state.load(&fake)); + fake = Fake.init(request); + fake.records[1] = fake.records[0]; + try std.testing.expectError(error.UnsupportedBinding, state.load(&fake)); + var duplicate = request; + duplicate.layers[1].layer = duplicate.layers[0].layer; + try std.testing.expectError(error.UnsupportedBinding, binding.Binding.init(42, duplicate)); + duplicate = request; + duplicate.layers[1].state_in = duplicate.layers[0].state_out; + try std.testing.expectError(error.UnsupportedBinding, binding.Binding.init(42, duplicate)); + try std.testing.expectEqual(@as(usize, 0), fake.calls); +} + +test "exl3: state binding full validation precedes every setter" { + const request = requestFor(1, true); + var fake = Fake.init(request); + var state = try binding.Binding.init(fake.template, request); + try state.load(&fake); + var target = flip(request); + target.layers[1].state_out += 1; + try std.testing.expectError(error.UnsupportedBinding, state.apply(&fake, target)); + target = flip(request); + target.nonstate.gk += 1; + try std.testing.expectError(error.UnsupportedBinding, state.apply(&fake, target)); + target = flip(request); + target.geometry.block[0] += 1; + try std.testing.expectError(error.UnsupportedBinding, state.apply(&fake, target)); + target = flip(request); + target.chain_to += 1; + try std.testing.expectError(error.UnsupportedBinding, state.apply(&fake, target)); + try std.testing.expectEqual(@as(usize, 0), fake.calls); + fake = Fake.init(request); + fake.records[1].values.nonstate.ys += 1; + state = try binding.Binding.init(fake.template, request); + try std.testing.expectError(error.UnsupportedBinding, state.load(&fake)); + fake = Fake.init(request); + fake.query_failure = true; + try std.testing.expectError(error.PriorAsyncFailure, state.load(&fake)); +} + +test "exl3: state binding partial update poisons all future uses" { + const request = requestFor(2, true); + var fake = Fake.init(request); + var state = try binding.Binding.init(fake.template, request); + try state.load(&fake); + fake.fail_at = 1; + try std.testing.expectError(error.SetterFailed, state.apply(&fake, flip(request))); + try std.testing.expect(state.poisoned); + try std.testing.expectError(error.StateBindingPoisoned, state.apply(&fake, request)); + try std.testing.expectEqual(@as(usize, 2), fake.calls); +} + +const Peer = struct { + const Vote = struct { state: enum { ready, failed }, cohort: u64 }; + const Decision = struct { choice: enum { replay, capture, eager }, cohort: u64 }; + choice: Decision = .{ .choice = .replay, .cohort = 7 }, + observed: ?Vote = null, + disconnected: bool = false, + + pub fn graphAgree(self: *Peer, _: u64, vote: Vote) !Decision { + self.observed = vote; + if (self.disconnected) return error.RankDisconnected; + return self.choice; + } +}; + +test "exl3: state binding peer failure disagreement and disconnect poison successful local updates" { + const request = requestFor(2, true); + for (0..4) |scenario| { + var fake = Fake.init(request); + var state = try binding.Binding.init(fake.template, request); + try state.load(&fake); + var peer: Peer = .{}; + switch (scenario) { + 0 => peer.choice.choice = .capture, + 1 => peer.choice.cohort = 8, + 2 => peer.disconnected = true, + 3 => fake.fail_at = 1, + else => unreachable, + } + if (scenario == 2) { + try std.testing.expectError(error.RankDisconnected, state.applyTogether(&fake, flip(request), &peer, 3, 7)); + } else try std.testing.expectError(error.StateRebindFailed, state.applyTogether(&fake, flip(request), &peer, 3, 7)); + try std.testing.expect(state.poisoned); + try std.testing.expectEqual(scenario == 3, peer.observed.?.state == .failed); + } + var fake = Fake.init(request); + var state = try binding.Binding.init(fake.template, request); + try state.load(&fake); + var peer: Peer = .{}; + try state.applyTogether(&fake, flip(request), &peer, 3, 7); + try std.testing.expect(!state.poisoned); +} + +test "exl3: state binding extended windows reject incompatible shapes and poison partial rank updates" { + for (2..binding.max_width + 1) |width| { + const request = requestFor(width, true); + var fake = Fake.init(request); + var state = try binding.Binding.init(fake.template, request); + try state.load(&fake); + var changed = flip(request); + changed.width = if (width == binding.max_width) 2 else width + 1; + changed.nonstate.rows = @intCast(changed.width); + try std.testing.expectError(error.UnsupportedBinding, state.apply(&fake, changed)); + changed = flip(request); + changed.nonstate.rows += 1; + try std.testing.expectError(error.UnsupportedBinding, state.apply(&fake, changed)); + changed = flip(request); + changed.layers[1].state_in += 5000; + try std.testing.expectError(error.UnsupportedBinding, state.apply(&fake, changed)); + changed = flip(request); + changed.layers[1].layer += 1; + try std.testing.expectError(error.UnsupportedBinding, state.apply(&fake, changed)); + changed = flip(request); + changed.nonstate.ys += 1; + try std.testing.expectError(error.UnsupportedBinding, state.apply(&fake, changed)); + try std.testing.expectEqual(@as(usize, 0), fake.calls); + fake.fail_at = 1; + var peer: Peer = .{}; + try std.testing.expectError(error.StateRebindFailed, state.applyTogether(&fake, flip(request), &peer, 3, 7)); + try std.testing.expect(state.poisoned); + try std.testing.expect(peer.observed.?.state == .failed); + try std.testing.expectEqual(@as(usize, 2), fake.calls); + try std.testing.expectError(error.StateBindingPoisoned, state.apply(&fake, request)); + } + for ([_]usize{ 0, binding.max_width + 1, std.math.maxInt(usize) }) |width| { + var request = requestFor(2, true); + request.width = width; + try std.testing.expectError(error.UnsupportedBinding, binding.Binding.init(42, request)); + } +} + +test "exl3: state binding template replacement requires new metadata" { + const request = requestFor(1, true); + var fake = Fake.init(request); + var state = try binding.Binding.init(fake.template, request); + try state.load(&fake); + state.invalidate(); + fake.template = 43; + try std.testing.expectError(error.StateBindingPoisoned, state.load(&fake)); + try std.testing.expectError(error.StateBindingPoisoned, state.apply(&fake, flip(request))); + state = try binding.Binding.init(fake.template, request); + try std.testing.expectError(error.UnsupportedBinding, state.apply(&fake, flip(request))); + try state.load(&fake); + try state.apply(&fake, flip(request)); +} + +test "exl3: state binding keep hash excludes only layer state_out" { + const Layer = struct { proj: u64 = 3, state_out: u64 = 4, hist: u64 = 5 }; + const Keep = struct { layers: [2]Layer = @splat(.{}), mixer: u64 = 8 }; + const original: Keep = .{}; + var changed = original; + changed.layers[0].state_out += 1; + try std.testing.expectEqual(binding.keepHash(original), binding.keepHash(changed)); + changed.layers[1].hist += 1; + try std.testing.expect(binding.keepHash(original) != binding.keepHash(changed)); + changed = original; + changed.mixer += 1; + try std.testing.expect(binding.keepHash(original) != binding.keepHash(changed)); +} diff --git a/zig/src/families/flashnext/exl3/model/cuda_exl3.zig b/zig/src/families/flashnext/exl3/model/cuda_exl3.zig index fb49c3b268..667b5475e8 100644 --- a/zig/src/families/flashnext/exl3/model/cuda_exl3.zig +++ b/zig/src/families/flashnext/exl3/model/cuda_exl3.zig @@ -12,7 +12,11 @@ pub const kvtier = @import("../kv/cuda_exl3_kvtier.zig"); pub const pages = @import("../../cuda_pages.zig"); pub const prefill_k = @import("../prefill/cuda_exl3_prefill_k.zig"); pub const prefill = @import("../prefill/cuda_exl3_prefill.zig"); +pub const dplan = @import("../dplan/cuda_exl3_dplan.zig"); +pub const graph = @import("../graph/cuda_exl3_graph.zig"); pub const fork = @import("../tp/cuda_exl3_fork.zig"); +pub const chain = @import("cuda_exl3_chain.zig"); +pub const recycle = @import("../offload/cuda_exl3_recycle.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. @@ -37,7 +41,11 @@ test { _ = @import("../prefill/cuda_exl3_prefill_test.zig"); _ = @import("../prefill/cuda_exl3_prefill_head.zig"); _ = @import("../prefill/cuda_exl3_prefill_ring.zig"); + _ = dplan; + _ = graph; _ = fork; + _ = chain; + _ = recycle; _ = 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 index 9e87c88dc9..09b26d622b 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, TP=2) is qualified on sm_120 only; this device is sm_{d}: " ++ +pub const message = "Flash Next EXL3 (expert offload, TP=2, graphs) 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_chain.zig b/zig/src/families/flashnext/exl3/model/cuda_exl3_chain.zig new file mode 100644 index 0000000000..8411837124 --- /dev/null +++ b/zig/src/families/flashnext/exl3/model/cuda_exl3_chain.zig @@ -0,0 +1,594 @@ +//! The MTP draft chain as one CUDA graph (TF_FLASHNEXT_EXL3_GRAPH=1, greedy drafts). + +const std = @import("std"); +const platform = @import("sys"); +const cuda = @import("cuda"); +const wt = @import("../../cuda_prompt_wt.zig"); +const batch = @import("../../cuda_batch.zig"); +const prompt = @import("../../cuda_prompt.zig"); +const mtp = @import("../../cuda_mtp.zig"); +const dplan = @import("../dplan/cuda_exl3_dplan.zig"); +const graph = @import("../graph/cuda_exl3_graph.zig"); + +const dims = wt.dims; +const streams = wt.streams; +const wide = wt.wide; +const eps = wt.eps; +const out_sk = wt.out_sk; +const max_chain = mtp.max_chain; +const fc_sk = mtp.fc_sk; +const piece_rows = mtp.piece_rows; + +/// forwardHead's hook: the head's pool on the device planner for a step of `rows` rows; returns what to restore. +pub fn enter(e: *prompt.Engine, rows: usize) !?*dplan.DevPlan { + const was = e.exl3_dplan.active; + if (!graph.enabled() or rows == 0) return was; + const dp = (try dplan.forLayers(e.gpa, e.driver, &e.exl3_dplan, e.layers[0..e.n_layers])) orelse return was; + const mp = e.exl3_dplan.mtp_pool orelse return was; + if (rows > dp.max_rows or dp.indexOf(mp) == null) return was; + try dp.toDevice(e.stream.*); + e.exl3_dplan.active = dp; + return was; +} + +pub fn leave(e: *prompt.Engine, was: ?*dplan.DevPlan) void { + e.exl3_dplan.active = was; +} + +const Key = struct { + picks: usize, + heads: usize, + /// Attention seams of the first and the last step (chunk counts and score tiles move with the head's length). + seams: [4]u32, + head: usize, + kc: u64, + pages: @import("../../cuda_glue.zig").Pg, + tier: [5]u64 = @splat(0), + out: u64, + dlogits: u64, + cap: usize, + scratch: u64, + scratch_cap: usize, + scores: u64, + plan: u64, + confidence: u64, +}; + +const Entry = struct { key: Key, exec: cuda.graph.Exec = undefined, used: u64, pending: ?*graph.Pending = null, admission: ?graph.memory.Ticket = null, stream: ?cuda.Stream = null }; + +const State = struct { + head: usize, + ids: cuda.DeviceBuffer, + token: cuda.DeviceBuffer, + best: cuda.DeviceBuffer, + pos_host: cuda.HostBuffer, + best_host: cuda.HostBuffer, + module: cuda.Module, + token_fn: cuda.Function, + next_fn: cuda.Function, + /// [0] the chain's running decision, [1 + j] the decision after pick j (read back with the records). + alive: cuda.DeviceBuffer, + alive_host: cuda.HostBuffer, + /// Where a step's conditional body is captured (cuStreamBeginCaptureToGraph). + side: cuda.Stream, + cond: Cond, +}; + +/// The conditional-node driver entry points (CUDA 12.4+), looked up from libcuda. +const NodeParams = extern struct { + type: c_int, + reserved0: [3]c_int = .{ 0, 0, 0 }, + u: extern union { cond: CondParams, pad: [29]i64 }, + reserved2: i64 = 0, +}; +const CondParams = extern struct { handle: u64, type: c_int, size: c_uint, graphs: ?[*]?*anyopaque, ctx: ?*anyopaque }; +pub const Cond = struct { + lib: platform.DynLib, + handle_create: *const fn (*u64, ?*anyopaque, ?*anyopaque, c_uint, c_uint) callconv(.c) c_int, + add_node: *const fn (*?*anyopaque, ?*anyopaque, ?[*]const ?*anyopaque, ?*const anyopaque, usize, *NodeParams) callconv(.c) c_int, + capture_info: *const fn (?*anyopaque, *c_int, *u64, *?*anyopaque, *?[*]const ?*anyopaque, *?*const anyopaque, *usize) callconv(.c) c_int, + update_deps: *const fn (?*anyopaque, ?[*]const ?*anyopaque, ?*const anyopaque, usize, c_uint) callconv(.c) c_int, + begin_to_graph: *const fn (?*anyopaque, ?*anyopaque, ?[*]const ?*anyopaque, ?*const anyopaque, usize, c_int) callconv(.c) c_int, + ctx_current: *const fn (*?*anyopaque) callconv(.c) c_int, + + pub fn open() !Cond { + var lib = try platform.DynLib.open(platform.cuda_driver); + errdefer lib.close(); + var c: Cond = undefined; + c.lib = lib; + c.handle_create = lib.lookup(@TypeOf(c.handle_create), "cuGraphConditionalHandleCreate") orelse return error.MissingSymbol; + c.add_node = lib.lookup(@TypeOf(c.add_node), "cuGraphAddNode_v2") orelse return error.MissingSymbol; + c.capture_info = lib.lookup(@TypeOf(c.capture_info), "cuStreamGetCaptureInfo_v3") orelse return error.MissingSymbol; + c.update_deps = lib.lookup(@TypeOf(c.update_deps), "cuStreamUpdateCaptureDependencies_v2") orelse return error.MissingSymbol; + c.begin_to_graph = lib.lookup(@TypeOf(c.begin_to_graph), "cuStreamBeginCaptureToGraph") orelse return error.MissingSymbol; + c.ctx_current = lib.lookup(@TypeOf(c.ctx_current), "cuCtxGetCurrent") orelse return error.MissingSymbol; + return c; + } + + pub fn ok(r: c_int, what: []const u8) !void { + if (r == 0) return; + std.log.err("flashnext exl3 chain: {s}: CUDA error {d}", .{ what, r }); + return error.CudaFailed; + } + + /// The capturing graph of `s` (main capture) and a new conditional handle on it. + pub fn handle(c: *const Cond, s: cuda.Stream) !u64 { + var status: c_int = 0; + var id: u64 = 0; + var g: ?*anyopaque = null; + var deps: ?[*]const ?*anyopaque = null; + var edges: ?*const anyopaque = null; + var n: usize = 0; + try ok(c.capture_info(@ptrCast(s.handle), &status, &id, &g, &deps, &edges, &n), "cuStreamGetCaptureInfo"); + var ctx: ?*anyopaque = null; + try ok(c.ctx_current(&ctx), "cuCtxGetCurrent"); + var h: u64 = 0; + try ok(c.handle_create(&h, g, ctx, 0, 1), "cuGraphConditionalHandleCreate"); + return h; + } + + /// An IF node on `h` after `s`'s captured work; `s` continues after it; returns its body graph. + pub fn ifNode(c: *const Cond, s: cuda.Stream, h: u64) !?*anyopaque { + var status: c_int = 0; + var id: u64 = 0; + var g: ?*anyopaque = null; + var deps: ?[*]const ?*anyopaque = null; + var edges: ?*const anyopaque = null; + var n: usize = 0; + try ok(c.capture_info(@ptrCast(s.handle), &status, &id, &g, &deps, &edges, &n), "cuStreamGetCaptureInfo"); + var ctx: ?*anyopaque = null; + try ok(c.ctx_current(&ctx), "cuCtxGetCurrent"); + var params: NodeParams = .{ .type = 13, .u = .{ .pad = @splat(0) } }; + params.u.cond = .{ .handle = h, .type = 0, .size = 1, .graphs = null, .ctx = ctx }; + var node: ?*anyopaque = null; + try ok(c.add_node(&node, g, deps, null, n, ¶ms), "cuGraphAddNode(conditional)"); + const one = [1]?*anyopaque{node}; + try ok(c.update_deps(@ptrCast(s.handle), &one, null, 1, 1), "cuStreamUpdateCaptureDependencies"); + return params.u.cond.graphs.?[0]; + } +}; + +pub const Runtime = struct { state: ?State = null }; +var entries: [16]?Entry = @splat(null); +var clock: u64 = 0; + +pub var stats: struct { chains: u64 = 0, replays: u64 = 0, captures: u64 = 0, updates: u64 = 0, eager: u64 = 0 } = .{}; + +fn stateFor(e: *prompt.Engine, h: *const mtp.Head) !*State { + const runtime = &e.exl3_chain; + if (runtime.state) |*s| if (s.head == @intFromPtr(h)) return s; + if (runtime.state) |*s| release(runtime, s); + var st: State = undefined; + st.head = @intFromPtr(h); + st.ids = try cuda.DeviceBuffer.fromHost(e.driver, std.mem.sliceAsBytes(h.ids)); + errdefer st.ids.free(); + st.token = try cuda.DeviceBuffer.alloc(e.driver, 256); + errdefer st.token.free(); + st.best = try cuda.DeviceBuffer.alloc(e.driver, max_chain * 16); + errdefer st.best.free(); + st.pos_host = try cuda.HostBuffer.alloc(e.driver, max_chain * 16); + errdefer st.pos_host.free(); + st.best_host = try cuda.HostBuffer.alloc(e.driver, max_chain * 16); + errdefer st.best_host.free(); + st.module = try cuda.Module.load(e.driver, cuda.kernels.flashnext_exl3_dplan); + errdefer st.module.unload(); + st.token_fn = try st.module.function("fn_exl3_draft_token"); + st.next_fn = try st.module.function("fn_exl3_draft_next"); + st.alive = try cuda.DeviceBuffer.alloc(e.driver, (1 + max_chain) * 4 + 256); + errdefer st.alive.free(); + st.alive_host = try cuda.HostBuffer.alloc(e.driver, (1 + max_chain) * 4); + errdefer st.alive_host.free(); + st.side = try cuda.Stream.init(e.driver, true); + errdefer st.side.deinit(); + st.cond = try Cond.open(); + runtime.state = st; + return &runtime.state.?; +} + +fn release(runtime: *Runtime, s: *State) void { + for (&entries) |*en| if (en.*) |*x| if (x.key.head == s.head) { + dropEntry(x); + en.* = null; + }; + s.cond.lib.close(); + s.side.deinit(); + s.alive_host.free(); + s.alive.free(); + s.module.unload(); + s.best_host.free(); + s.pos_host.free(); + s.best.free(); + s.token.free(); + s.ids.free(); + runtime.state = null; +} + +pub fn releaseRuntime(runtime: *Runtime) void { + if (runtime.state) |*s| release(runtime, s); +} + +fn seams(end: usize, chunks: *u32, tiles: *u32) void { + chunks.* = @intCast((@min(end, wt.dense_keys) + 511) / 512); + tiles.* = if (end / wt.index_ratio > wt.index_top) @intCast(((end + 3) / 4 + 63) / 64) else 0; +} + +/// One head step for one row whose token is on the device: forwardHead's launches for `next.len == 1` (streams from +fn step(comptime Tri: type, o: anytype, tri: Tri, p: *batch.Pass(Tri), d: *mtp.Draft, st: *State, j: usize, len0: usize) !void { + const e = o.e; + const h = o.h; + const s = p.s; + const stream = p.stream; + p.rows = 1; + p.pos0 = len0 + j; + p.pending = null; + try stream.d.check(stream.d.api.cuMemcpyDtoDAsync_v2(s.ids, st.token.ptr, 4, stream.handle), "cuMemcpyDtoDAsync"); + try graph.pinnedCopy(stream, s.pos, st.pos_host.bytes[j * 16 ..][0..16]); + if (e.table.dense) { + try tri.embedDense(s.ids, e.table.w.ptr, s.ple_keys, dims, streams, 1); + } else { + try tri.embed(s.ids, e.table.w.ptr, e.table.s.ptr, e.table.b.ptr, s.ple_keys, dims, streams, 1); + } + try tri.rmsnorm(s.ple_keys, h.norm_e.ptr, s.mixed, s.xsm, eps, wide, 1, dims, dims); + _ = try p.lane(s.mixed, s.xsm, h.fc_e, s.dummy, s.part, 1, fc_sk); + try p.sumSlices(s.part, s.branch, true, dims, fc_sk); + try tri.rmsnorm(d.out.ptr, h.norm_h.ptr, s.norm, s.gq, eps, wide, 1, wide, wide); + _ = try p.lane(s.norm, s.gq, h.fc_h, s.dummy, s.part, streams, fc_sk); + try p.sumSlices(s.part, s.ple_gated, true, streams * dims, fc_sk); + try tri.addStreams(s.branch, s.ple_gated, s.h, 1, dims, streams); + try p.writeback(0, 1); + try p.readout(h.layer.attn_hc, 0, 1); + try p.attnBlock(&h.layer, 0); + try tri.hcWriteback(s.h, s.h, s.pss, s.part, "*fp32", s.inj, s.dummy, "*bf16", s.dummy, "*bf16", dims, 1, dims, streams, 4, 1, 1, out_sk); + try p.readout(h.layer.mlp_hc, 0, 1); + try p.moeBlock(&h.layer); + p.pending = 0; + try p.writeback(0, 1); + try p.copy(d.out.ptr, s.h, wide * 2); + try p.readout(h.mixer, 0, 1); + _ = try p.lane(s.mixed, s.xsm, h.draft, d.dlogits.ptr, s.dummy, 1, 1); +} + +/// The chain's launches on `p`'s stream: the decision reset, every pick's record into st.best; after a pick that +fn chainBody(comptime Tri: type, o: anytype, p: *batch.Pass(Tri), ps: *batch.Pass(Tri), d: *mtp.Draft, st: *State, picks: usize, heads: usize, len0: usize, confidence: f64) !void { + const h = o.h; + const stream = p.stream; + try stream.d.check(stream.d.api.cuMemsetD32Async(st.alive.ptr, 1, 1 + max_chain, stream.handle), "cuMemsetD32Async"); + for (0..picks) |j| { + try h.best(stream, d.dlogits.ptr, h.ids.len, h.ids.len, 1, st.best.ptr + j * 16); + if (j >= heads) continue; + const handle = try st.cond.handle(stream); + var a: cuda.Args = .{}; + a.add(st.best.ptr + j * 16); + a.add(st.ids.ptr); + a.add(st.token.ptr); + a.add(st.alive.ptr); + a.add(st.alive.ptr + (1 + j) * 4); + a.add(handle); + a.add(confidence); + try cuda.launch.launch(st.next_fn, .{ .grid = .{ .x = 1 }, .block = .{ .x = 32 } }, stream, &a); + const body = try st.cond.ifNode(stream, handle); + try Cond.ok(st.cond.begin_to_graph(@ptrCast(st.side.handle), body, null, null, 0, 1), "cuStreamBeginCaptureToGraph"); + step(Tri, o, ps.tri, ps, d, st, j, len0) catch |err| { + var discarded: cuda.abi.Graph = null; + _ = stream.d.api.cuStreamEndCapture(st.side.handle, &discarded); + return err; + }; + var g: cuda.abi.Graph = null; + try stream.d.check(stream.d.api.cuStreamEndCapture(st.side.handle, &g), "cuStreamEndCapture(body)"); + } +} + +/// Ops.draft's hook (cuda_mtp_ops.zig): the chain for greedy drafts as one graph; false leaves it to upstream. +pub fn draft(comptime Tri: type, o: anytype, d: *mtp.Draft, next: []const u32, rows: u64, position: u64, count: usize, sm: ?mtp.Sampling, confidence: f64, out: *mtp.Chain) !bool { + _ = position; + if (!graph.enabled() or count < 2 or d.media != null) return false; + if (sm) |s| if (s.temperature > 0) return false; + const e = o.e; + const dp = (try dplan.forLayers(e.gpa, e.driver, &e.exl3_dplan, e.layers[0..e.n_layers])) orelse return false; + const mp = e.exl3_dplan.mtp_pool orelse return false; + if (dp.indexOf(mp) == null) return false; + const picks: usize = @min(count, max_chain); + // upstream runs a head step after pick j while j + 1 < count + const heads: usize = if (count > max_chain) max_chain else count - 1; + if (heads > picks) return false; + out.n = 0; + try o.absorb(d, next, rows, true); + if (d.len + heads > d.cap) { + // upstream's loop, which reports PromptTooLong where it would + stats.eager += 1; + return eagerChain(o, d, count, picks, confidence, out); + } + const st = try stateFor(e, o.h); + const len0 = d.len; + const page_graph = @import("../graph/cuda_exl3_graph_pages.zig"); + if (d.sel) |selection| try page_graph.stage(selection, e.stream.*, wt.attn_n, 1, len0, heads); + defer if (d.sel) |selection| page_graph.finish(selection); + const pos: []i32 = @alignCast(std.mem.bytesAsSlice(i32, st.pos_host.bytes[0 .. max_chain * 16])); + @memset(pos, 0); + for (0..heads) |j| pos[4 * j] = @intCast(len0 + j); + const s = try e.reserve(@max(@as(usize, 32), 2)); + const piece = @max(@as(usize, 1), @min(piece_rows, s.cap / 2)); + const score_rows = try e.reserveScores(d.cap, piece); + var p: batch.Pass(Tri) = .{ + .gpa = e.gpa, + .io = e.io, + .stream = e.stream.*, + .tri = o.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 = (&o.h.layer)[0..1], .mixer = o.h.mixer, .head = o.h.draft, .ngram = &e.ngram }, + .seq = d.headSeq(d.media, score_rows), + .hooks = .{ .logits = .none }, + }; + p.m.dplan = &e.exl3_dplan; + var key: Key = .{ .picks = picks, .heads = heads, .seams = undefined, .head = @intFromPtr(o.h), .kc = d.kc.ptr, .pages = d.pg, .out = d.out.ptr, .dlogits = d.dlogits.ptr, .cap = d.cap, .scratch = s.ids, .scratch_cap = s.cap, .scores = score_rows, .plan = dp.state.ptr, .confidence = @bitCast(confidence) }; + if (d.sel) |selection| key.tier = @import("../graph/cuda_exl3_graph_key.zig").Key.tierOf(selection, len0, heads); + seams(len0 + 1, &key.seams[0], &key.seams[1]); + seams(len0 + @max(heads, 1), &key.seams[2], &key.seams[3]); + try dp.toDevice(p.stream); + const was = e.exl3_dplan.active; + e.exl3_dplan.active = dp; + defer e.exl3_dplan.active = was; + stats.chains += 1; + var tri_side = o.tri; + tri_side.s = st.side; + var ps = p; + ps.stream = st.side; + ps.tri = tri_side; + const launched = run(Tri, key, o, &p, &ps, d, st, picks, heads, len0, confidence) catch |err| { + std.log.err("flashnext exl3 chain: the draft chain failed ({t})", .{err}); + return err; + }; + if (!launched) { + stats.eager += 1; + return eagerChain(o, d, count, picks, confidence, out); + } + // the chain's records, once + const stream = p.stream; + try stream.d.check(stream.d.api.cuMemcpyDtoHAsync_v2(st.best_host.bytes.ptr, st.best.ptr, picks * 16, stream.handle), "cuMemcpyDtoHAsync"); + try stream.d.check(stream.d.api.cuMemcpyDtoHAsync_v2(st.alive_host.bytes.ptr, st.alive.ptr, (1 + max_chain) * 4, stream.handle), "cuMemcpyDtoHAsync"); + try stream.synchronize(); + const alive: []const i32 = @alignCast(std.mem.bytesAsSlice(i32, st.alive_host.bytes[0 .. (1 + max_chain) * 4])); + d.len += heads; + d.drafted += heads; + const rec: []const f32 = @alignCast(std.mem.bytesAsSlice(f32, st.best_host.bytes[0 .. picks * 16])); + for (0..picks) |j| { + const top = rec[4 * j]; + const lse = rec[4 * j + 1]; + const col: usize = @intCast(@as(i32, @bitCast(rec[4 * j + 2]))); + const got: mtp.Pick = .{ .token = o.h.ids[col], .prob = @exp(@as(f64, top) - @as(f64, lse)) }; + const low = confidence > 0 and got.prob < confidence; + if (low and j > 0) break; + out.tokens[out.n] = got.token; + out.probs[out.n] = got.prob; + out.n += 1; + if (low) break; + // the device stopped the chain here (its decision is the same rule; a step past it did not run) + if (j < heads and alive[1 + j] == 0) break; + } + return true; +} + +/// upstream's loop (Ops.draft after its absorb), for chains that leave the graph path. +fn eagerChain(o: anytype, d: *mtp.Draft, count: usize, picks: usize, confidence: f64, out: *mtp.Chain) !bool { + var j: usize = 0; + while (j < picks) : (j += 1) { + const got = try o.pick(d, 0, null); + const low = confidence > 0 and got.prob < confidence; + if (low and j > 0) break; + out.tokens[out.n] = got.token; + out.probs[out.n] = got.prob; + out.n += 1; + if (low) break; + if (j + 1 < count) { + try o.forwardHead(d, &.{got.token}, d.out.ptr, true); + d.drafted += 1; + } + } + return true; +} + +fn sameShape(a: Key, b: Key) bool { + var x = a; + x.kc = b.kc; + x.out = b.out; + x.dlogits = b.dlogits; + return std.meta.eql(x, b); +} + +/// Replays the chain's graph (true), or captures it; a new graph instantiates on a worker thread and this chain runs +fn run(comptime Tri: type, key: Key, o: anytype, p: *batch.Pass(Tri), ps: *batch.Pass(Tri), d: *mtp.Draft, st: *State, picks: usize, heads: usize, len0: usize, confidence: f64) !bool { + clock += 1; + for (&entries) |*en| if (en.*) |*x| if (std.meta.eql(x.key, key)) { + x.used = clock; + if (x.pending) |pd| { + const got = graph.takePending(pd) catch { + if (x.admission) |ticket| graph.memory.release(ticket) catch unreachable; + en.* = null; + return failed(key); + }; + var exec = got orelse return false; + x.pending = null; + graph.uploadResident(&exec, p.stream) catch { + exec.deinit(); + if (x.admission) |ticket| graph.memory.release(ticket) catch unreachable; + en.* = null; + return failed(key); + }; + x.exec = exec; + if (x.admission) |ticket| try graph.memory.installed(ticket); + } + stats.replays += 1; + try x.exec.launchOn(p.stream); + return true; + }; + if (graph.fallback.blocked(graph.fallback.hashOf(3, key), clock)) return false; + try p.stream.synchronize(); + const admission = graph.reserveMemory(p.stream, .chain, 0) catch |err| { + if (graph.fallback.admissionDenied(err)) return failed(key); + return err; + }; + var admitted = true; + defer if (admitted) graph.memory.release(admission) catch unreachable; + cuda.graph.beginCapture(p.stream, .thread_local) catch return failed(key); + graph.capturing = true; + p.prepared = true; + chainBody(Tri, o, p, ps, d, st, picks, heads, len0, confidence) catch { + graph.capturing = false; + if (cuda.graph.endCapture(p.stream)) |g| { + var x = g; + x.deinit(); + } else |_| {} + return failed(key); + }; + graph.capturing = false; + var g = cuda.graph.endCapture(p.stream) catch return failed(key); + if (graph.fallback.take(&graph.fallback.inject_capture_failures)) { + g.deinit(); + return failed(key); + } + var owned = true; + defer if (owned) g.deinit(); + stats.captures += 1; + // a stale graph of the same shape (another stream's buffers) is updated in place; else a new one + var oldest: usize = 0; + var oldest_used: u64 = std.math.maxInt(u64); + for (&entries, 0..) |*en, i| { + if (en.* == null) { + oldest = i; + oldest_used = 0; + break; + } + if (en.*.?.used < oldest_used) { + oldest_used = en.*.?.used; + oldest = i; + } + } + for (&entries) |*en| if (en.*) |*x| if (x.pending == null and sameShape(x.key, key) and clock - x.used > 64) { + if ((try x.exec.update(g)) == .success) { + try graph.uploadResident(&x.exec, p.stream); + x.key = key; + x.used = clock; + stats.updates += 1; + try x.exec.launchOn(p.stream); + return true; + } + }; + if (entries[oldest]) |*x| dropEntry(x); + entries[oldest] = null; + if (try graph.startInstantiate(g, false)) |pd| { + owned = false; + entries[oldest] = .{ .key = key, .used = clock, .pending = pd, .admission = admission, .stream = p.stream }; + admitted = false; + return false; + } + installSync(&entries[oldest], g, .{ .key = key, .used = clock, .admission = admission, .stream = p.stream }) catch return failed(key); + admitted = false; + try entries[oldest].?.exec.launchOn(p.stream); + return true; +} + +/// The chain's capture or instantiation failed: its key backs off and this chain runs as upstream's loop does. +fn failed(key: Key) bool { + graph.fallback.record(graph.fallback.hashOf(3, key), clock); + return false; +} + +fn installSync(target: *?Entry, template: cuda.graph.Graph, entry: Entry) !void { + var exec = try (if (graph.fallback.take(&graph.fallback.inject_instantiate_failures)) error.CudaFailed else template.instantiate()); + errdefer exec.deinit(); + try graph.uploadResident(&exec, entry.stream.?); + try graph.memory.installed(entry.admission.?); + target.* = entry; + target.*.?.exec = exec; +} + +fn dropEntry(x: *Entry) void { + if (x.stream) |stream| stream.synchronize() catch @panic("chain retirement synchronization failed"); + defer if (x.admission) |ticket| { + graph.memory.release(ticket) catch unreachable; + x.admission = null; + }; + if (x.pending) |pd| { + while (true) { + const got = graph.takePending(pd) catch break; + if (got) |e| { + var ex = e; + ex.deinit(); + break; + } + std.Thread.yield() catch {}; + } + return; + } + x.exec.deinit(); +} + +pub fn evictOne() bool { + var oldest: ?usize = null; + for (entries, 0..) |maybe, index| if (maybe) |entry| { + if (oldest == null or entry.used < entries[oldest.?].?.used) oldest = index; + }; + const index = oldest orelse return false; + dropEntry(&entries[index].?); + entries[index] = null; + return true; +} + +test "exl3: synchronous chain launch failure retains ownership until retirement" { + const Mock = struct { + var drops: usize = 0; + + fn instantiate(out: *cuda.abi.GraphExec, _: cuda.abi.Graph, _: u64) callconv(.c) c_int { + out.* = @ptrFromInt(1); + return 0; + } + + fn upload(_: cuda.abi.GraphExec, _: cuda.abi.Stream) callconv(.c) c_int { + return 0; + } + + fn synchronize(_: cuda.abi.Stream) callconv(.c) c_int { + return 0; + } + + fn launch(_: cuda.abi.GraphExec, _: cuda.abi.Stream) callconv(.c) c_int { + return cuda.abi.error_not_ready; + } + + fn destroy(_: cuda.abi.GraphExec) callconv(.c) c_int { + drops += 1; + return 0; + } + + fn installAndLaunch(target: *?Entry, template: cuda.graph.Graph, entry: Entry) !void { + try installSync(target, template, entry); + try target.*.?.exec.launchOn(entry.stream.?); + } + }; + Mock.drops = 0; + var driver: cuda.Driver = undefined; + driver.api.cuGraphInstantiateWithFlags = Mock.instantiate; + driver.api.cuGraphUpload = Mock.upload; + driver.api.cuStreamSynchronize = Mock.synchronize; + driver.api.cuGraphLaunch = Mock.launch; + driver.api.cuGraphExecDestroy = Mock.destroy; + const stream: cuda.Stream = .{ .d = &driver, .handle = @ptrFromInt(1) }; + try graph.memory.configure(.{ .reservation_bytes = 100, .headroom_bytes = 20, .charge = .{ .exec = 20, .template = 10, .transient = 10 } }); + const ticket = try graph.memory.admit(.chain, 0, 100); + var ticket_owned = true; + defer if (ticket_owned) graph.memory.release(ticket) catch unreachable; + var entry: ?Entry = null; + defer if (entry) |*cached| dropEntry(cached); + const result = Mock.installAndLaunch(&entry, .{ .d = &driver, .handle = @ptrFromInt(1) }, .{ .key = undefined, .used = 1, .admission = ticket, .stream = stream }); + if (entry != null) ticket_owned = false; + try std.testing.expectError(error.NotReady, result); + try std.testing.expect(entry != null); + try std.testing.expectEqual(@as(usize, 0), Mock.drops); + try std.testing.expectEqual(@as(u64, 30), graph.memory.charged()); + dropEntry(&entry.?); + entry = null; + try std.testing.expectEqual(@as(usize, 1), Mock.drops); + try std.testing.expectEqual(@as(u64, 0), graph.memory.charged()); + try std.testing.expectError(error.StaleGraphTicket, graph.memory.release(ticket)); +} 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 index f8d01f1dfc..5ce98943d9 100644 --- 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 @@ -1,6 +1,6 @@ const cuda = @import("cuda"); -/// A driver whose only live entry points are the two synchronizations, which count themselves. +/// A driver whose live entry points are the two counted synchronizations and device-to-host reads of host memory. pub const Fake = struct { pub var sync_calls: usize = 0; @@ -13,11 +13,24 @@ pub const Fake = struct { return sync(); } + fn download(destination: ?*anyopaque, source: u64, length: usize) callconv(.c) c_int { + const target: [*]u8 = @ptrCast(destination.?); + const origin: [*]const u8 = @ptrFromInt(source); + @memcpy(target[0..length], origin[0..length]); + return 0; + } + + fn downloadAsync(destination: ?*anyopaque, source: u64, length: usize, _: cuda.abi.Stream) callconv(.c) c_int { + return download(destination, source, length); + } + pub fn driver() cuda.Driver { sync_calls = 0; var result: cuda.Driver = undefined; result.api.cuCtxSynchronize = sync; result.api.cuStreamSynchronize = streamSync; + result.api.cuMemcpyDtoH_v2 = download; + result.api.cuMemcpyDtoHAsync_v2 = downloadAsync; 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 index 0e9fd23574..aca031ad4f 100644 --- a/zig/src/families/flashnext/exl3/model/cuda_exl3_pass.zig +++ b/zig/src/families/flashnext/exl3/model/cuda_exl3_pass.zig @@ -5,6 +5,7 @@ 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 dplan = @import("../dplan/cuda_exl3_dplan.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"); @@ -25,7 +26,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 + if (try @import("../prefill/cuda_exl3_prefill.zig").tilesPart(if (p.m.prefill) |rt| rt.tiles_on else false, 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| { @@ -101,7 +102,12 @@ pub fn moe(p: anytype, lw: anytype, x: *const load.Experts) !void { 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); + const plans = p.m.dplan; + if (plans) |st| if (st.active) |dp| if (dp.indexOf(pool)) |i| { // a decode round planned on the device (no host cut) + try dplan.runLayer(dp, p.stream, i, ex, x.layer, xin, picks, rows, sc, 0, 0); + continue; + }; + try pooled(p.stream, x.ops, pool, plans, ex, x.layer, xin, picks, rows, sc, 0, 0); continue; } try ex.runSlots(p.stream, x.layer, xin, dims, .bf16, picks, rows, sc); @@ -109,12 +115,17 @@ pub fn moe(p: anytype, lw: anytype, x: *const load.Experts) !void { 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 -pub var pool_time: struct { cut_ns: u64 = 0, plan_ns: u64 = 0, blocks: u64 = 0 } = .{}; +/// Whether this engine has an attached expert offload pool. +pub fn offloaded(e: anytype) bool { + for (e.layers[0..e.n_layers]) |layer| if (layer.exl3) |x| if (x.pool != null) return true; + return false; +} /// 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 { +pub fn pooled(st: cuda.Stream, ops: *load.Ops, pool: *pool_.Pool, plans: ?*const dplan.State, 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]; + // the device planned the last decode rounds of this pool: its residency back to the host pools first (the MTP + if (plans) |state| if (dplan.current(state)) |dp| if (dp.indexOf(pool) != null) try dp.toHost(); 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) @@ -122,9 +133,9 @@ pub fn pooled(st: cuda.Stream, ops: *load.Ops, pool: *pool_.Pool, ex: *const cud 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; + pool.stats.cut_ns += @intCast(t1 - t0); + pool.stats.plan_ns += @intCast(std.Io.Clock.awake.now(io).toNanoseconds() - t1); + pool.stats.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); @@ -149,6 +160,8 @@ pub fn offloadedLayers(layers: []const @import("../../cuda_prompt_load.zig").Lay } pub const Budget = struct { + graph_reservation: u64 = 0, + planner_rows: ?u64 = null, /// 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. @@ -173,6 +186,9 @@ pub fn budgetForStreams(stream_bytes: u64, requested_streams: u32) !Budget { 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 (dplan.enabled()) { + b.planner_rows = dplan.planner_rows; + } 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| { @@ -193,7 +209,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, prefill_window_bytes: 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, graph_reservation: u64 = 0, after_pool_bytes: u64 = 0, bank_bytes: u64 = 0, bank_ns: u64 = 0, prefill_window_bytes: u64 = 0, mtp_pool: ?*pool_.Pool = null }; pub fn freeBytes(d: *const cuda.Driver) !u64 { var free: usize = 0; @@ -203,12 +219,13 @@ 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, mtp: ?*const @import("../../cuda_prompt_load.zig").Layer, rows: usize, b: Budget, profile_counts: ?*[pool_.shape.layers][pool_.shape.experts]u32) !Attached { +pub fn attach(gpa: std.mem.Allocator, io: std.Io, d: *const cuda.Driver, st: cuda.Stream, runtime: ?*@import("../prefill/cuda_exl3_prefill_state.zig").Runtime, 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); + out.graph_reservation = b.graph_reservation; // 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; @@ -242,23 +259,37 @@ pub fn attach(gpa: std.mem.Allocator, io: std.Io, d: *const cuda.Driver, st: cud 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; + var future = try std.math.add(u64, b.other, b.graph_reservation); + const planner_fixed = if (b.planner_rows) |planner_rows| try memory.plannerBytes(n, 0, planner_rows) else 0; + const line_cost = try std.math.add(u64, out.line_bytes, if (b.planner_rows != null) 8 else 0); + const all_fixed = try std.math.add(u64, fixed, planner_fixed); 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 minimum_pools = try memory.poolBytes(b.lines orelse pool_.min_lines, line_cost, n, all_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: { + out.prefill_window_bytes = prefill.reserveTwoChunk(runtime, d, free, protected, b.graph_reservation) catch |err| declined: { std.log.warn("flashnext exl3 prefill: two-chunk reservation refused ({t}); keeping single-chunk prefill", .{err}); break :declined 0; }; + const graph_memory = @import("../graph/cuda_exl3_graph_memory.zig"); + const window_reservation = graph_memory.twoChunkReservationBytes(); + if (!prefill.twoChunkAllocated(runtime) and window_reservation != 0 and b.graph_reservation == window_reservation) { + out.graph_reservation = try graph_memory.declineTwoChunkReservation(); + future = try std.math.add(u64, b.other, out.graph_reservation); + } 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); + const world: usize = if (all[0].tp) |t| t.world else 1; + const fit = pool_.linesFor(free, try std.math.add(u64, serverReserve(total), b.margin), try std.math.add(u64, future, all_fixed), line_cost, n, pool_.min_lines) catch null; + const uniform = std.c.getenv("TF_FLASHNEXT_EXL3_UNIFORM_LINES") != null; + const lines = if (b.lines) |explicit| (if (delta_text != null or uniform) explicit else memory.rankLines(explicit, world, fit)) else try pool_.linesFor(free, reserve, try std.math.add(u64, future, all_fixed), line_cost, n, pool_.min_lines); + if (b.lines) |explicit| if (lines != explicit) std.log.info("flashnext exl3: this rank holds {d} lines a layer, {d} more than the {d} asked for: it has the room (residency never changes output bytes)", .{ lines, lines - explicit, explicit }); if (lines < pool_.min_lines) return error.PoolTooSmall; - const startup_bytes = try memory.poolBytes(lines, out.line_bytes, n, fixed); + const startup_bytes = try memory.poolBytes(lines, line_cost, n, all_fixed); if (startup_bytes > pool_room) return error.InsufficientStartupMemory; + out.after_pool_bytes = if (b.planner_rows) |planner_rows| try memory.plannerBytes(n, lines, planner_rows) else 0; out.lines = lines; var pool_lines: [wt.layers_n + 1]usize = @splat(lines); var line_bytes: [wt.layers_n + 1]usize = undefined; @@ -284,7 +315,7 @@ pub fn attach(gpa: std.mem.Allocator, io: std.Io, d: *const cuda.Driver, st: cud has_profile_signal = has_profile_signal or count != 0; } count_rows[decoder_index] = &prior[decoder_index]; - allocation_costs[decoder_index] = line_bytes[pool_index]; + allocation_costs[decoder_index] = try std.math.add(usize, line_bytes[pool_index], if (b.planner_rows != null) 8 else 0); decoder_index += 1; } if (has_profile_signal and decoder_index == wt.layers_n) { @@ -316,16 +347,18 @@ pub fn attach(gpa: std.mem.Allocator, io: std.Io, d: *const cuda.Driver, st: cud @memset(pool_lines[0..n], lines); std.log.info("flashnext exl3: profile candidate changed total variable CUDA bytes; retaining uniform decoder pools", .{}); } else { + out.after_pool_bytes = if (b.planner_rows) |planner_rows| try memory.plannerBytesForPools(pool_lines[0..n], planner_rows) else 0; 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; MTP remains uniform", .{ plan.allocated_bytes, plan.baseline_bytes }); + std.log.info("flashnext exl3: profile-weighted decoder allocation used {d}/{d} variable CUDA bytes (expert lines, per-line planner state 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, 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 }); + const plan = try line_deltas.apply(text, lines, decoder_pool_count, line_bytes[0..n], pool_.min_lines, b.planner_rows, fixed, startup_bytes, pool_lines[0..n]); + out.after_pool_bytes = plan.planner_bytes; + std.log.info("flashnext exl3: experimental rank-local decoder pool line deltas {s}; base={d} planner_bytes={d} startup_bytes={d}", .{ text, lines, plan.planner_bytes, plan.pool_bytes }); } // the first decoder layer's pool owns the lifetime use table the other decoder layers share var counts = profile_counts; @@ -333,14 +366,15 @@ pub fn attach(gpa: std.mem.Allocator, io: std.Io, d: *const cuda.Driver, st: cud for (all[0..n], pool_lines[0..n]) |x, count| { 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) out.mtp_pool = x.pool.?; // the draft head's experts plan on the device too (cuda_exl3_chain.zig) if (!is_mtp and counts == null) counts = x.pool.?.counts; if (!is_mtp) 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 }); + _ = memory.available(out.free_after, reserve_floor, try std.math.add(u64, future, out.after_pool_bytes)) 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} planner={d}", .{ free, out.free_after, out.device_bytes, reserve_floor, b.margin, future, out.after_pool_bytes }); 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"; @@ -411,3 +445,17 @@ pub fn offloadedPools(layers: []const @import("../../cuda_prompt_load.zig").Laye }; return n; } + +test "exl3: offload facts belong to the engine whose layers are inspected" { + const Layer = struct { exl3: ?*load.Experts = null }; + const Engine = struct { n_layers: usize = 1, layers: [1]Layer = .{.{}} }; + var a: Engine = .{}; + var b: Engine = .{}; + var experts: load.Experts = undefined; + experts.pool = @ptrFromInt(256); + a.layers[0].exl3 = &experts; + try std.testing.expect(offloaded(&a)); + try std.testing.expect(!offloaded(&b)); + experts.pool = null; + try std.testing.expect(!offloaded(&a)); +} 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 af7527dfe2..9f24b66afa 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 dplan = @import("../dplan/cuda_exl3_dplan.zig"); const prefill = @import("../prefill/cuda_exl3_prefill.zig"); const weights = @import("../../cuda_prompt_wt.zig"); pub const counters = @import("exl3_telemetry.zig"); @@ -32,7 +33,7 @@ pub fn sampleAfterConsumers(allocator: std.mem.Allocator, engine: *prompt.Engine seen[seen_count] = resident; seen_count += 1; hash.update(std.mem.asBytes(&resident.telemetry_id)); - const device = resident.telemetry_device; + const device = if (dplan.current(&engine.exl3_dplan)) |planner| try planner.samplePoolAfterConsumers(resident) else resident.telemetry_device; var minimum: u64 = std.math.maxInt(u64); var maximum: u64 = 0; for (resident.src.routed, resident.k2) |source, widths| { 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 3e48bb222e..aa5d8b1347 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 @@ -1,7 +1,117 @@ const std = @import("std"); const cuda = @import("cuda"); +const planner = @import("../dplan/cuda_exl3_dplan.zig"); const pool = @import("../offload/cuda_exl3_pool.zig"); const telemetry = @import("exl3_telemetry.zig"); +const Fake = @import("cuda_exl3_fake_driver_test.zig").Fake; + +test "exl3: replay device sampling reads executed counters without pulling ownership or synchronizing" { + var driver = Fake.driver(); + var resident: pool.Pool = undefined; + resident.lines.pinned_bank = null; + resident.telemetry_device = .{ .plans = 5, .hit_picks = 20, .copies = 3 }; + resident.telemetry_valid = true; + var pools = [_]*pool.Pool{&resident}; + var state: [planner.stateWords(1)]u32 = @splat(0); + const executed: planner.Misc = .{ .now = 104, .rounds = 4, .hit_slots = 30, .miss_slots = 10, .copies = 8, .evictions = 2, .overflow = 0, .inplace = 1, .recent_cursor = 9, .recent_filled = 8 }; + const offset = planner.miscOffset(1) * 4; + @memcpy(std.mem.sliceAsBytes(&state)[offset..][0..@sizeOf(planner.Misc)], std.mem.asBytes(&executed)); + var seen = [_]planner.Misc{std.mem.zeroes(planner.Misc)}; + var device: planner.DevPlan = undefined; + device.owner = .device; + device.pools = &pools; + var pool_states = [_]planner.PoolState{try planner.PoolState.init(1, 0)}; + device.pool_states = &pool_states; + device.max_rows = 1; + device.state = .{ .d = &driver, .ptr = @intFromPtr(&state), .len = @sizeOf(@TypeOf(state)), .borrowed = true }; + device.seen = &seen; + const syncs = Fake.sync_calls; + const first = try device.samplePoolAfterConsumers(&resident); + const again = try device.samplePoolAfterConsumers(&resident); + try std.testing.expectEqualDeep(first, again); + try std.testing.expectEqual(@as(u64, 9), first.plans); + try std.testing.expectEqual(@as(u64, 50), first.hit_picks); + try std.testing.expectEqual(@as(u64, 11), first.copies); + try std.testing.expectEqual(@as(u64, 1), first.inplace_experts); + try std.testing.expectEqual(planner.Owner.device, device.owner); + try std.testing.expectEqual(@as(u32, 0), seen[0].copies); + try std.testing.expectEqual(syncs, Fake.sync_calls); + resident.telemetry_device = first; + seen[0] = executed; + device.owner = .host; + try std.testing.expectEqualDeep(first, try device.samplePoolAfterConsumers(&resident)); + device.owner = .device; + try std.testing.expectEqualDeep(first, try device.samplePoolAfterConsumers(&resident)); +} + +test "exl3: ragged planner telemetry and MTP lookup use each pool's state bounds" { + var driver = Fake.driver(); + const counts = [_]usize{ 10, 37, 12 }; + var residents: [counts.len]pool.Pool = undefined; + var pools: [counts.len]*pool.Pool = undefined; + var pool_states: [counts.len]planner.PoolState = undefined; + var seen: [counts.len]planner.Misc = @splat(std.mem.zeroes(planner.Misc)); + var words: usize = 0; + for (counts, 0..) |lines, index| { + residents[index].opt = .{ .layer = if (index == 2) 0 else @intCast(index), .lines = lines }; + residents[index].telemetry_valid = true; + residents[index].telemetry_device = .{}; + pools[index] = &residents[index]; + pool_states[index] = try planner.PoolState.init(lines, words); + words = pool_states[index].endWords(); + } + const state = try std.testing.allocator.alloc(u32, words); + defer std.testing.allocator.free(state); + const mirror = try std.testing.allocator.alignedAlloc(u8, .@"16", words * 4); + defer std.testing.allocator.free(mirror); + @memset(state, 0); + for (pool_states, 0..) |entry, index| { + const misc: *planner.Misc = @ptrCast(entry.slice(state)[planner.miscOffset(entry.lines)..][0..10]); + misc.* = std.mem.zeroes(planner.Misc); + misc.rounds = @intCast(index + 1); + misc.hit_slots = @intCast(10 * (index + 1)); + misc.copies = @intCast(index + 2); + } + var device: planner.DevPlan = undefined; + device.d = &driver; + device.owner = .device; + device.pools = &pools; + device.pool_states = &pool_states; + device.max_rows = 16; + device.state = .{ .d = &driver, .ptr = @intFromPtr(state.ptr), .len = state.len * 4, .borrowed = true }; + device.mirror = .{ .d = &driver, .bytes = mirror }; + device.stream = null; + device.seen = &seen; + const syncs = Fake.sync_calls; + for (&residents, 0..) |*resident, index| { + try std.testing.expectEqual(@as(?usize, index), device.indexOf(resident)); + const sample = try device.samplePoolAfterConsumers(resident); + try std.testing.expectEqual(@as(u64, index + 1), sample.plans); + try std.testing.expectEqual(@as(u64, 10 * (index + 1)), sample.hit_picks); + try std.testing.expectEqual(@as(u64, index + 2), sample.copies); + } + var outside: pool.Pool = undefined; + outside.opt = .{ .layer = 0, .lines = 18 }; + outside.telemetry_valid = true; + outside.telemetry_device = .{ .plans = 7 }; + try std.testing.expectEqual(@as(?usize, null), device.indexOf(&outside)); + try std.testing.expectEqualDeep(outside.telemetry_device, try device.samplePoolAfterConsumers(&outside)); + const totals = try device.peek(); + try std.testing.expectEqual(@as(u32, 6), totals.rounds); + try std.testing.expectEqual(@as(u32, 60), totals.hit_slots); + try std.testing.expectEqual(@as(u32, 9), totals.copies); + try std.testing.expectEqualSlices(u8, std.mem.sliceAsBytes(state), mirror); + try std.testing.expectEqual(syncs, Fake.sync_calls); +} + +test "exl3: device counter widening handles one u32 rollover before summation" { + var previous = std.mem.zeroes(planner.Misc); + previous.hit_slots = std.math.maxInt(u32) - 2; + var current = previous; + current.hit_slots = 4; + const delta = telemetry.Device.pending(current, previous); + try std.testing.expectEqual(@as(u64, 7), delta.hit_picks); +} 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 = .{}; @@ -32,7 +142,10 @@ test "exl3: aggregate telemetry finds attached pools deduplicates MTP aliases an var experts: @import("../pack/cuda_exl3_load.zig").Experts = undefined; experts.pool = &resident; var engine: @import("../../cuda_prompt.zig").Engine = undefined; + engine.exl3_prefill = .{}; + engine.exl3_chain = .{}; engine.telemetry = .{}; + engine.exl3_dplan = .{}; engine.n_layers = 1; engine.layers[0].exl3 = &experts; const collector = @import("cuda_exl3_telemetry.zig"); diff --git a/zig/src/families/flashnext/exl3/model/exl3_cache.zig b/zig/src/families/flashnext/exl3/model/exl3_cache.zig index c193eefa6f..f27075ceea 100644 --- a/zig/src/families/flashnext/exl3/model/exl3_cache.zig +++ b/zig/src/families/flashnext/exl3/model/exl3_cache.zig @@ -19,7 +19,7 @@ pub fn resumeAt(prompt_len: u32) u32 { pub const Budget = struct { device_bytes: u64, - /// Dense + shared + router + HC + MTP workspace, already summed by the caller. + /// Dense + shared + router + HC + MTP + graph workspace, already summed by the caller. resident_bytes: u64, kv_bytes: u64, routed_expert_bytes: u64, diff --git a/zig/src/families/flashnext/exl3/model/exl3_round.zig b/zig/src/families/flashnext/exl3/model/exl3_round.zig index 9d75bb336d..1538666c4c 100644 --- a/zig/src/families/flashnext/exl3/model/exl3_round.zig +++ b/zig/src/families/flashnext/exl3/model/exl3_round.zig @@ -133,7 +133,7 @@ pub const Round = struct { return r.n_passes > 1; } - /// True when the round needs no copy at all (the experts need no upload). + /// True when the round needs no copy at all (a graph of the experts may replay straight away). pub fn allHits(r: *const Round) bool { return r.n_copies == 0; } @@ -360,7 +360,7 @@ 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) +/// Lines a layer gets from `free` device bytes after `reserve` and `other` (KV, scratch, graphs: 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; @@ -371,6 +371,7 @@ pub fn linesFor(free: u64, reserve: u64, other: u64, line_bytes: u64, layers: us test { _ = offload; + _ = @import("../graph/exl3_graph.zig"); } // ------------------------------------------------------------------------------------------------ tests 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 6331276501..95c53aeb63 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 @@ -18,7 +18,7 @@ pub const State = struct { } pub fn save(state: *State, own: anytype) void { - const counts = snapshot(own) catch |err| { + const counts = snapshot(own, @import("../dplan/cuda_exl3_dplan.zig").current(&own.engine.exl3_dplan)) catch |err| { std.log.err("EXL3 placement profile rank {d}: unsafe snapshot skipped ({s})", .{ state.identity.rank, @errorName(err) }); return; }; @@ -190,10 +190,20 @@ fn decayRow(row: *[512]u32) void { } } -fn snapshot(own: anytype) !codec.Counts { +fn snapshot(own: anytype, planner: anytype) !codec.Counts { const layers = own.engine.layers[0..own.engine.n_layers]; if (!available(layers)) return error.ProfileUnavailable; try own.stream.synchronize(); + if (planner) |current| { + if (current.indexOf(layers[0].exl3.?.pool.?) != null) { + for (current.pools) |pool| { + var ours = pool == mtpPool(own.head); + for (layers) |layer| ours = ours or pool == layer.exl3.?.pool.?; + if (!ours) return error.ProfilePlannerOwnerMismatch; + } + try current.toHost(); + } + } var counts: codec.Counts = .{}; for (layers, 0..) |layer, index| { const pool = layer.exl3.?.pool.?; @@ -299,6 +309,21 @@ const Mock = struct { } }; +const MockPlan = struct { + pools: []const *MockPool, + pulls: usize = 0, + fail: bool = false, + pub fn indexOf(plan: *MockPlan, pool: *const MockPool) ?usize { + for (plan.pools, 0..) |held, index| if (held == pool) return index; + return null; + } + pub fn toHost(plan: *MockPlan) !void { + if (plan.fail) return error.MockCudaPullFailure; + plan.pulls += 1; + for (plan.pools, 0..) |pool, index| pool.counts[pool.opt.layer][7] = @intCast(100 + index); + } +}; + const test_identity: codec.Identity = .{ .model = @splat(42), .world = 2, .rank = 0 }; test "exl3: placement state lifetime lock conflicts, rank separation and private storage" { @@ -387,7 +412,7 @@ test "exl3: placement state decay and replacement snapshot do not add imported p 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); + const saved = try snapshot(&own, @as(?*MockPlan, null)); 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]); @@ -401,20 +426,35 @@ test "exl3: placement state decay and replacement snapshot do not add imported p try std.testing.expectError(error.FileNotFound, codec.write(allocator, io, missing, state.identity, &saved)); } -test "exl3: placement state snapshot synchronizes once and drains decoder and MTP counts" { +test "exl3: placement state pulls only owning planner before copying decoder and MTP counts" { if (comptime @import("builtin").os.tag == .windows) return error.SkipZigTest; var mock: Mock = undefined; mock.init(); 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); + var owned: [49]*MockPool = undefined; + for (&mock.pools, &owned) |*pool, *target| target.* = pool; + var plan: MockPlan = .{ .pools = &owned }; + const counts = try snapshot(&own, @as(?*MockPlan, &plan)); try std.testing.expectEqual(@as(usize, 1), own.stream.syncs); + try std.testing.expectEqual(@as(usize, 1), plan.pulls); 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); + plan.fail = true; + try std.testing.expectError(error.MockCudaPullFailure, snapshot(&own, @as(?*MockPlan, &plan))); own.stream.fail = true; - try std.testing.expectError(error.MockCudaSyncFailure, snapshot(&own)); + try std.testing.expectError(error.MockCudaSyncFailure, snapshot(&own, @as(?*MockPlan, &plan))); + own.stream.fail = false; + var other: Mock = undefined; + other.init(); + const foreign = [_]*MockPool{&other.pools[0]}; + var unrelated: MockPlan = .{ .pools = &foreign, .fail = true }; + _ = try snapshot(&own, @as(?*MockPlan, &unrelated)); + try std.testing.expectEqual(@as(usize, 0), unrelated.pulls); + const mixed = [_]*MockPool{ &mock.pools[0], &other.pools[0] }; + unrelated.pools = &mixed; + try std.testing.expectError(error.ProfilePlannerOwnerMismatch, snapshot(&own, @as(?*MockPlan, &unrelated))); } test "exl3: prepared placement counts import and warm only after pools attach" { 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 9ae062fe28..8da842c215 100644 --- a/zig/src/families/flashnext/exl3/offload/cuda_exl3_pool.zig +++ b/zig/src/families/flashnext/exl3/offload/cuda_exl3_pool.zig @@ -9,6 +9,7 @@ const telemetry = @import("../model/exl3_telemetry.zig"); pub const offload = @import("exl3_offload.zig"); pub const round = round_; +pub const graph = @import("../graph/exl3_graph.zig"); pub const Round = round_.Round; pub const linesFor = round_.linesFor; pub const shape = offload.shape; @@ -59,6 +60,9 @@ pub fn passesFor(lines: usize) usize { } pub const Stats = struct { + cut_ns: u64 = 0, + plan_ns: u64 = 0, + blocks: u64 = 0, rounds: u64 = 0, passes: u64 = 0, hit_slots: u64 = 0, diff --git a/zig/src/families/flashnext/exl3/offload/cuda_exl3_recycle.zig b/zig/src/families/flashnext/exl3/offload/cuda_exl3_recycle.zig new file mode 100644 index 0000000000..2184bad0e0 --- /dev/null +++ b/zig/src/families/flashnext/exl3/offload/cuda_exl3_recycle.zig @@ -0,0 +1,144 @@ +//! Stable buffers for captured decode graphs: a closed stream's caches and draft window are kept for the next one. + +const std = @import("std"); +const cuda = @import("cuda"); +const prompt = @import("../../cuda_prompt.zig"); +const Draft = @import("../../cuda_mtp_draft.zig").Draft; + +pub const bucket: usize = 4096; +const pool_n = 4; + +pub fn enabled() bool { + return @import("../graph/cuda_exl3_graph.zig").enabled(); +} + +/// A stream's capacity for `want` positions: a pool lane's own size, else rounded up to the bucket (at most `limit`). +pub fn capFor(pooled: bool, want: usize, limit: usize) usize { + const cap = @min(want, prompt.max_tokens); + if (pooled or !enabled()) return cap; + return @min(@min(limit, prompt.max_tokens), (cap + bucket - 1) / bucket * bucket); +} + +var seqs: [pool_n]?prompt.Seq = @splat(null); +var drafts: [pool_n]?Draft = @splat(null); +var draining = false; +pub var stats: struct { seq_reused: u64 = 0, seq_kept: u64 = 0, seq_evicted: u64 = 0, draft_reused: u64 = 0, draft_kept: u64 = 0 } = .{}; + +fn zero(bufs: []cuda.DeviceBuffer, s: cuda.Stream) !void { + for (bufs) |b| try b.fill8(0, s.handle); +} + +/// Engine.newSeq's hook: a kept sequence of capacity `cap`, zeroed as newSeq zeroes a new one; null: allocate. +pub fn takeSeq(e: *prompt.Engine, cap: usize) !?prompt.Seq { + const reuse = enabled(); + for (&seqs) |*slot| if (slot.*) |*sequence| { + if (sequence.recycle_owner == e and (!reuse or !matches(sequence, e, cap))) { + discardSeq(slot, e.gpa); + stats.seq_evicted += 1; + } + }; + if (!reuse) return null; + for (&seqs) |*slot| if (slot.*) |*q| if (matches(q, e, cap)) { + var s = q.*; + slot.* = null; + errdefer { + s.tail.d.synchronize() catch @panic("recycled sequence reset did not drain"); + s.recycle_owner = null; // a half-reset sequence is freed, not kept again + s.deinit(e.gpa); + } + const st = e.stream.*; + for ([_][]cuda.DeviceBuffer{ s.conv, s.state, s.kc, s.vc, s.ikc, s.pooled }) |bs| try zero(bs, st); + try s.tail.fill8(0, st.handle); + @memset(s.logits, 0); + s.history = .{ e.eos, e.eos }; + s.fed = 0; + s.rank_id = 0; + s.media = null; + s.untaught = false; + stats.seq_reused += 1; + return s; + }; + return null; +} + +fn discardSeq(slot: *?prompt.Seq, allocator: std.mem.Allocator) void { + var sequence = slot.*.?; + slot.* = null; + sequence.recycle_owner = null; + sequence.deinit(allocator); +} + +fn canRetainSeq(s: *const prompt.Seq) bool { + return s.media == null and s.paged == null and s.recycle_owner != null; +} + +fn matches(sequence: *const prompt.Seq, engine: *prompt.Engine, cap: usize) bool { + if (sequence.recycle_owner != engine or sequence.tail.d != engine.driver or sequence.cap != cap) return false; + return sequence.paged == null; +} + +/// Seq.deinit's hook: true when the sequence's BF16 allocations were kept for reuse. +pub fn giveSeq(s: *prompt.Seq) bool { + if (draining or !enabled() or !canRetainSeq(s)) return false; + for (&seqs) |*slot| if (slot.* == null) { + s.tail.d.synchronize() catch @panic("recycled sequence consumers did not drain"); + slot.* = s.*; + stats.seq_kept += 1; + return true; + }; + return false; +} + +/// Draft.init's hook: a kept draft window of the same shape, its caches zeroed and its state reset; null: allocate. +pub fn takeDraft(stream: cuda.Stream, cap: usize, draft_n: usize, rows: usize) !?Draft { + if (!enabled()) return null; + for (&drafts) |*slot| if (slot.*) |*d| if (d.cap == cap and d.dlogits.len == draft_n * 2 and d.rows_cap == @max(rows, @import("../../cuda_mtp.zig").max_rows)) { + var x = d.*; + x.stream = stream; + x.verify_cache.invalidate(); + slot.* = null; + for ([_]cuda.DeviceBuffer{ x.kc, x.vc, x.ikc, x.pooled }) |b| try b.fill8(0, stream.handle); + x.media = null; + x.len = 0; + x.drafted = 0; + x.rows = 0; + x.open = false; + x.prompt_last = 0; + x.fed0 = 0; + x.history = .{ 0, 0 }; + x.unused_history = .{ 0, 0 }; + stats.draft_reused += 1; + return x; + }; + return null; +} + +/// Draft.deinit's hook: true when the window was kept. +pub fn giveDraft(d: *Draft) bool { + if (draining or !enabled() or d.media != null or d.pg.pt != 0) return false; // a pool lane's head is the pool's + // A peer's contiguous window can swap in a paged sequence's borrowed run states. + for (d.spare) |state| if (state.borrowed) return false; + for (&drafts) |*slot| if (slot.* == null) { + slot.* = d.*; + stats.draft_kept += 1; + return true; + }; + return false; +} + +/// Frees everything kept (the engine's deinit). +pub fn drain(engine: *prompt.Engine) void { + draining = true; + defer draining = false; + for (&seqs) |*slot| if (slot.*) |*sequence| if (sequence.recycle_owner == engine) { + discardSeq(slot, engine.gpa); + }; + for (&drafts) |*slot| if (slot.*) |*d| { + d.deinit(); + slot.* = null; + }; +} + +test { + _ = @import("cuda_exl3_recycle_test.zig"); +} diff --git a/zig/src/families/flashnext/exl3/offload/cuda_exl3_recycle_test.zig b/zig/src/families/flashnext/exl3/offload/cuda_exl3_recycle_test.zig new file mode 100644 index 0000000000..bc2cada99b --- /dev/null +++ b/zig/src/families/flashnext/exl3/offload/cuda_exl3_recycle_test.zig @@ -0,0 +1,300 @@ +const std = @import("std"); +const cuda = @import("cuda"); +const prompt = @import("../../cuda_prompt.zig"); +const recycle = @import("cuda_exl3_recycle.zig"); +const graphs = @import("../graph/cuda_exl3_graph.zig"); +const allocator = std.testing.allocator; + +test "exl3: settled peer windows cannot retain borrowed pool states" { + const Draft = @import("../../cuda_mtp_draft.zig").Draft; + const previous = graphs.force; + graphs.force = .graph; + defer graphs.force = previous; + var driver = Recording.driver(); + const stream: cuda.Stream = .{ .d = &driver, .handle = @ptrFromInt(9) }; + var window = try Draft.init(allocator, &driver, stream, 1, 1, 1); + defer { + graphs.force = .off; + window.deinit(); + } + var run = try buffer(&driver, window.spare.len * 16); + defer run.free(); + var states: @TypeOf(window.spare) = undefined; + for (&states, 0..) |*state, i| state.* = .{ .d = &driver, .ptr = run.ptr + i * 16, .len = 16, .borrowed = true }; + var sequence_view: prompt.Seq = undefined; + sequence_view.state = &states; + window.open = true; + window.settle(&sequence_view); + defer for (&states) |*state| state.free(); + try std.testing.expect(window.pg.pt == 0); + const retained = recycle.giveDraft(&window); + // Recover an incorrectly retained window so this regression fails without a dangling recycler entry. + if (retained) window = (try recycle.takeDraft(stream, 1, 1, 1)).?; + try std.testing.expect(!retained); +} + +/// A driver of host memory: allocations, fills and the two synchronizations, counting what the recycler allocates. +const Recording = struct { + var allocations: usize = 0; + var fills: usize = 0; + var fail_fill_at: ?usize = null; + var live: [256]?[]align(16) u8 = @splat(null); + + fn alloc(output: *u64, len: usize) callconv(.c) c_int { + allocations += 1; + for (&live) |*slot| if (slot.* == null) { + const bytes = allocator.alignedAlloc(u8, .@"16", @max(len, 1)) catch return 2; + slot.* = bytes; + output.* = @intFromPtr(bytes.ptr); + return 0; + }; + return 2; + } + + fn free(address: u64) callconv(.c) c_int { + for (&live) |*slot| if (slot.*) |bytes| if (@intFromPtr(bytes.ptr) == address) { + allocator.free(bytes); + slot.* = null; + return 0; + }; + return 1; + } + + fn fill(address: u64, value: u8, len: usize, _: cuda.abi.Stream) callconv(.c) c_int { + fills += 1; + if (fail_fill_at == fills) return cuda.abi.error_not_ready; + const target: [*]u8 = @ptrFromInt(address); + @memset(target[0..len], value); + return 0; + } + + fn fillNow(address: u64, value: u8, len: usize) callconv(.c) c_int { + const target: [*]u8 = @ptrFromInt(address); + @memset(target[0..len], value); + return 0; + } + + fn sync() callconv(.c) c_int { + return 0; + } + + fn streamSync(_: cuda.abi.Stream) callconv(.c) c_int { + return 0; + } + + fn driver() cuda.Driver { + allocations = 0; + fills = 0; + fail_fill_at = null; + var result: cuda.Driver = undefined; + result.api.cuMemAlloc_v2 = alloc; + result.api.cuMemFree_v2 = free; + result.api.cuMemsetD8Async = fill; + result.api.cuMemsetD8_v2 = fillNow; + result.api.cuCtxSynchronize = sync; + result.api.cuStreamSynchronize = streamSync; + return result; + } +}; + +fn engine(driver: *cuda.Driver, stream: *cuda.Stream) prompt.Engine { + var result: prompt.Engine = undefined; + result.exl3_prefill = .{}; + result.exl3_chain = .{}; + result.gpa = allocator; + result.driver = driver; + result.stream = stream; + result.chunk = 1; + result.eos = 42; + result.round_graph = .{}; + return result; +} + +fn buffer(driver: *const cuda.Driver, len: usize) !cuda.DeviceBuffer { + const result = try cuda.DeviceBuffer.alloc(driver, len); + try result.fill8(0x7f, null); + return result; +} + +fn buffers(driver: *const cuda.Driver) ![]cuda.DeviceBuffer { + const result = try allocator.alloc(cuda.DeviceBuffer, 1); + result[0] = try buffer(driver, 16); + return result; +} + +fn sequence(owner: *prompt.Engine, cap: usize) !prompt.Seq { + const driver = owner.driver; + return .{ + .recycle_owner = owner, + .conv = try buffers(driver), + .state = try buffers(driver), + .kc = try buffers(driver), + .vc = try buffers(driver), + .ikc = try buffers(driver), + .pooled = try buffers(driver), + .tail = try buffer(driver, 16), + .history = .{ 1, 2 }, + .fed = cap, + .rank_id = 99, + .untaught = true, + .logits = try allocator.alloc(u8, 16), + .cap = cap, + }; +} + +fn zero(bytes: []const u8) !void { + for (bytes) |value| try std.testing.expectEqual(@as(u8, 0), value); +} + +test "exl3: recycled sequences keep their allocations and come back zeroed" { + const previous = graphs.force; + graphs.force = .graph; + defer graphs.force = previous; + var driver = Recording.driver(); + var stream: cuda.Stream = .{ .d = &driver, .handle = @ptrFromInt(9) }; + var owner = engine(&driver, &stream); + defer recycle.drain(&owner); + var current = try sequence(&owner, 16); + const addresses = [_]u64{ current.tail.ptr, current.kc[0].ptr, current.state[0].ptr }; + const count = Recording.allocations; + for (0..8) |_| { + current.fed = 16; + current.history = .{ 1, 2 }; + current.rank_id = 99; + current.untaught = true; + @memset(current.logits, 0xff); + current.deinit(allocator); + current = (try recycle.takeSeq(&owner, 16)).?; + try std.testing.expectEqual(addresses, [_]u64{ current.tail.ptr, current.kc[0].ptr, current.state[0].ptr }); + try std.testing.expectEqual(count, Recording.allocations); + try std.testing.expectEqual(@as(usize, 0), current.fed); + try std.testing.expectEqual([2]i64{ 42, 42 }, current.history); + try std.testing.expectEqual(@as(u64, 0), current.rank_id); + try std.testing.expect(!current.untaught); + try zero(current.logits); + for ([_][]cuda.DeviceBuffer{ current.conv, current.state, current.kc, current.vc, current.ikc, current.pooled }) |group| for (group) |item| try zero(@as([*]const u8, @ptrFromInt(item.ptr))[0..item.len]); + try zero(@as([*]const u8, @ptrFromInt(current.tail.ptr))[0..current.tail.len]); + } + current.deinit(allocator); +} + +test "exl3: request allocation retention is bounded to four sequences" { + const previous = graphs.force; + graphs.force = .graph; + defer graphs.force = previous; + var driver = Recording.driver(); + var stream: cuda.Stream = .{ .d = &driver, .handle = @ptrFromInt(9) }; + var owner = engine(&driver, &stream); + defer recycle.drain(&owner); + const before = cuda.memory.usage(false); + var requests: [5]prompt.Seq = undefined; + for (&requests) |*request| request.* = try sequence(&owner, 16); + for (&requests) |*request| request.deinit(allocator); + var retained: [4]prompt.Seq = undefined; + for (&retained) |*request| request.* = (try recycle.takeSeq(&owner, 16)).?; + try std.testing.expect(try recycle.takeSeq(&owner, 16) == null); + for (&retained) |*request| { + request.recycle_owner = null; + request.deinit(allocator); + } + const after = cuda.memory.usage(false); + try std.testing.expectEqual(before.device, after.device); + try std.testing.expectEqual(before.host, after.host); +} + +test "exl3: recycler isolates engine shape and driver and drains only its owner's sequences" { + const previous = graphs.force; + graphs.force = .graph; + defer graphs.force = previous; + var driver = Recording.driver(); + var stream: cuda.Stream = .{ .d = &driver, .handle = @ptrFromInt(9) }; + var owner = engine(&driver, &stream); + var other = engine(&driver, &stream); + defer recycle.drain(&owner); + defer recycle.drain(&other); + var first = try sequence(&owner, 16); + first.deinit(allocator); + try std.testing.expect(try recycle.takeSeq(&other, 16) == null); + const evicted = recycle.stats.seq_evicted; + try std.testing.expect(try recycle.takeSeq(&owner, 17) == null); + try std.testing.expectEqual(evicted + 1, recycle.stats.seq_evicted); + first = try sequence(&owner, 16); + first.deinit(allocator); + var wrong_driver: cuda.Driver = undefined; + owner.driver = &wrong_driver; + try std.testing.expect(try recycle.takeSeq(&owner, 16) == null); + try std.testing.expectEqual(evicted + 2, recycle.stats.seq_evicted); + owner.driver = &driver; + first = try sequence(&owner, 16); + const address = first.tail.ptr; + first.deinit(allocator); + var second = try sequence(&other, 16); + second.deinit(allocator); + recycle.drain(&other); + first = (try recycle.takeSeq(&owner, 16)).?; + try std.testing.expectEqual(address, first.tail.ptr); + first.recycle_owner = null; + try std.testing.expect(!recycle.giveSeq(&first)); + first.deinit(allocator); +} + +test "exl3: capacity churn evicts retained allocations under the original limit" { + const previous = graphs.force; + graphs.force = .graph; + defer graphs.force = previous; + var driver = Recording.driver(); + var stream: cuda.Stream = .{ .d = &driver, .handle = @ptrFromInt(9) }; + var owner = engine(&driver, &stream); + defer recycle.drain(&owner); + var active: [8]?prompt.Seq = @splat(null); + defer for (&active) |*entry| if (entry.*) |*request| request.deinit(allocator); + for (&active) |*request| request.* = try sequence(&owner, 64); + const admitted = cuda.memory.usage(false); + const previous_limit = try cuda.memory.setAllocationLimit(admitted.device + admitted.host); + defer _ = cuda.memory.setAllocationLimit(previous_limit) catch unreachable; + for (active[0..4]) |*entry| { + entry.*.?.deinit(allocator); + entry.* = null; + } + const evicted = recycle.stats.seq_evicted; + for (active[0..4]) |*entry| { + try std.testing.expect(try recycle.takeSeq(&owner, 16) == null); + entry.* = try sequence(&owner, 16); + } + try std.testing.expectEqual(evicted + 4, recycle.stats.seq_evicted); + const mixed = cuda.memory.usage(false); + try std.testing.expect(mixed.device + mixed.host <= admitted.device + admitted.host); + for (active[0..4]) |*entry| { + entry.*.?.deinit(allocator); + entry.* = null; + } + const count = Recording.allocations; + for (active[0..4]) |*entry| entry.* = (try recycle.takeSeq(&owner, 16)).?; + try std.testing.expectEqual(count, Recording.allocations); + for (active[0..4]) |*entry| { + entry.*.?.deinit(allocator); + entry.* = null; + } + graphs.force = .off; + try std.testing.expect(try recycle.takeSeq(&owner, 16) == null); + try std.testing.expectEqual(evicted + 8, recycle.stats.seq_evicted); +} + +test "exl3: a partial recycled reset drains and frees its allocations without retention" { + const previous = graphs.force; + graphs.force = .graph; + defer graphs.force = previous; + var driver = Recording.driver(); + var stream: cuda.Stream = .{ .d = &driver, .handle = @ptrFromInt(9) }; + var owner = engine(&driver, &stream); + defer recycle.drain(&owner); + const before = cuda.memory.usage(false); + var current = try sequence(&owner, 16); + current.deinit(allocator); + Recording.fail_fill_at = Recording.fills + 3; + try std.testing.expectError(error.NotReady, recycle.takeSeq(&owner, 16)); + try std.testing.expect(try recycle.takeSeq(&owner, 16) == null); + const after = cuda.memory.usage(false); + try std.testing.expectEqual(before.device, after.device); + try std.testing.expectEqual(before.host, after.host); +} diff --git a/zig/src/families/flashnext/exl3/offload/exl3_line_deltas.zig b/zig/src/families/flashnext/exl3/offload/exl3_line_deltas.zig index 50e9eca705..c03e5c1091 100644 --- a/zig/src/families/flashnext/exl3/offload/exl3_line_deltas.zig +++ b/zig/src/families/flashnext/exl3/offload/exl3_line_deltas.zig @@ -1,10 +1,13 @@ const std = @import("std"); +const memory = @import("../pack/exl3_startup_memory.zig"); +const layout = @import("../dplan/exl3_dplan_layout.zig"); pub const Plan = struct { pool_bytes: u64, + planner_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 { +pub fn apply(text: []const u8, uniform_lines: usize, decoder_pools: usize, line_bytes: []const usize, min_lines: usize, planner_rows: ?u64, 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); @@ -36,14 +39,17 @@ pub fn apply(text: []const u8, uniform_lines: usize, decoder_pools: usize, line_ 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)); + _ = try layout.PoolState.init(count.*, 0); + const cost = try std.math.add(u64, bytes, 8); + before = try std.math.add(u64, before, try std.math.mul(u64, uniform_lines, cost)); + after = try std.math.add(u64, after, try std.math.mul(u64, count.*, cost)); 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); + const planner_bytes = if (planner_rows) |rows| try memory.plannerBytesForPools(counts, rows) else 0; + const pool_bytes = try std.math.add(u64, try std.math.add(u64, expert_bytes, fixed_bytes), planner_bytes); if (pool_bytes > startup_reservation) return error.InsufficientStartupMemory; - return .{ .pool_bytes = pool_bytes }; + return .{ .pool_bytes = pool_bytes, .planner_bytes = planner_bytes }; } fn decimal(text: []const u8) !usize { @@ -55,9 +61,10 @@ fn decimal(text: []const u8) !usize { 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); + const plan = try apply("20:+1,15:-1", 32, 21, &bytes, 10, 16, 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(try memory.plannerBytes(22, 32, 16), plan.planner_bytes); + try std.testing.expectEqual(22 * 32 * 256 + 100 + plan.planner_bytes, plan.pool_bytes); try std.testing.expectEqual(@as(usize, 32), counts[21]); } @@ -65,56 +72,61 @@ 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)); + try std.testing.expectError(error.InvalidLineDeltas, apply(text, 32, 2, &bytes, 10, 16, 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)); + try std.testing.expectError(error.DuplicateLineDelta, apply(text, 32, 2, &bytes, 10, 16, 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(text, 32, 2, &bytes, 10, 16, 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)); + try std.testing.expectError(error.InvalidLineDeltaIndex, apply("0:+0", 32, 0, &bytes, 10, 16, 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)); + try std.testing.expectError(error.NonZeroLineDeltas, apply(text, 32, 3, &bytes, 10, 16, 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)); + try std.testing.expectError(error.LineDeltaByteBudgetChanged, apply(text, 32, 3, &bytes, 10, 16, 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); + const planner = try memory.plannerBytes(4, 32, 16); + const expected_bytes = 32 * (256 + 512 + 256 + 1024) + planner + 100; + const plan = try apply("0:+1,2:-1", 32, 3, &bytes, 10, 16, 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)); + try std.testing.expectEqual(planner, plan.planner_bytes); + try std.testing.expectError(error.InsufficientStartupMemory, apply("0:+1,2:-1", 32, 3, &bytes, 10, 16, 100, expected_bytes - 1, &counts)); + const without_planner = try apply("0:+1,2:-1", 32, 3, &bytes, 10, null, 100, expected_bytes - planner, &counts); + try std.testing.expectEqual(@as(u64, 0), without_planner.planner_bytes); + try std.testing.expectEqual(expected_bytes - planner, without_planner.pool_bytes); } 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 apply("0:+22,1:-22", 32, 2, &bytes, 10, 16, 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(text, 32, 2, &bytes, 10, 16, 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)); + try std.testing.expectError(error.PoolTooSmall, apply("0:+0", 9, 2, &bytes, 10, 16, 0, std.math.maxInt(u64), &counts)); + try std.testing.expectError(error.Overflow, apply("0:+18446744073709551615,1:-18446744073709551615", 32, 2, &bytes, 10, 16, 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, 16, 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); + const reservation = try memory.poolBytes(32, 1024 + 8, 4, fixed + try memory.plannerBytes(4, 0, 16)); + const plan = try apply("0:+1,1:-2,2:+1", 32, 3, &bytes, 10, 16, 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.expectEqual(@as(u64, 32 * (256 + 512 + 768 + 1024) + fixed) + try memory.plannerBytes(4, 32, 16), plan.pool_bytes); + _ = try apply("0:+0", 32, 3, &bytes, 10, 16, fixed, reservation, &counts); try std.testing.expectEqualSlices(usize, &.{ 32, 32, 32, 32 }, &counts); } 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 7bf9ed4655..cc502b8f48 100644 --- a/zig/src/families/flashnext/exl3/pack/exl3_startup_memory.zig +++ b/zig/src/families/flashnext/exl3/pack/exl3_startup_memory.zig @@ -1,11 +1,15 @@ const std = @import("std"); +const planner_layout = @import("../dplan/exl3_dplan_layout.zig"); +const demand_trace = @import("../dplan/exl3_dplan_trace.zig"); pub const Reservation = struct { stream_bytes: u64, streams: u64, + graph_bytes: u64, + after_pool_bytes: u64 = 0, pub fn total(self: Reservation) !u64 { - return std.math.mul(u64, self.stream_bytes, self.streams); + return std.math.add(u64, try std.math.add(u64, try std.math.mul(u64, self.stream_bytes, self.streams), self.graph_bytes), self.after_pool_bytes); } }; @@ -28,40 +32,90 @@ pub fn mebibytes(text: []const u8) !u64 { return std.math.mul(u64, try std.fmt.parseInt(u64, text, 10), 1 << 20); } +pub fn graphReservation(enabled: bool, text: ?[]const u8) !u64 { + if (!enabled) return 0; + const bytes = try mebibytes(text orelse return error.GraphReservationRequired); + if (bytes == 0) return error.GraphReservationRequired; + return bytes; +} + 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; } +/// Lines a rank holds: the asked count is the floor; at world 2 a rank with room for more takes them. +pub fn rankLines(explicit: usize, world: usize, fit: ?usize) usize { + if (world != 2) return explicit; + return @max(explicit, fit orelse explicit); +} + +test "exl3: a rank with room beyond the asked lines takes them, the binding rank and one rank keep the ask" { + try std.testing.expectEqual(@as(usize, 80), rankLines(80, 2, 70)); + try std.testing.expectEqual(@as(usize, 90), rankLines(80, 2, 90)); + try std.testing.expectEqual(@as(usize, 80), rankLines(80, 2, null)); + try std.testing.expectEqual(@as(usize, 80), rankLines(80, 1, 90)); +} + +pub fn laneRoom(local: u64, graph_reservation: u64) u64 { + return local -| graph_reservation; +} + 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" { +pub fn plannerBytes(layers: u64, lines: u64, rows: u64) !u64 { + const state = try planner_layout.deviceBytes(layers, try std.math.mul(u64, layers, lines), rows); + return std.math.add(u64, state, try plannerTraceBytes(layers, traceEnabled())); +} + +pub fn plannerBytesForPools(lines: []const usize, rows: u64) !u64 { + var total_lines: u64 = 0; + for (lines) |count| total_lines = try std.math.add(u64, total_lines, count); + const state = try planner_layout.deviceBytes(lines.len, total_lines, rows); + return std.math.add(u64, state, try plannerTraceBytes(lines.len, traceEnabled())); +} + +pub fn plannerTraceBytes(pools: u64, enabled: bool) !u64 { + return if (enabled) std.math.mul(u64, pools, demand_trace.pool_bytes) else 0; +} + +fn traceEnabled() bool { + return demand_trace.requested; +} + +test "BF16 c1 c8 and graph on off 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()); + for ([_]bool{ false, true }) |graphs| { + const graph = try graphReservation(graphs, "512"); + const reservation: Reservation = .{ .stream_bytes = bf16_stream, .streams = streams, .graph_bytes = graph, .after_pool_bytes = 4096 }; + try std.testing.expectEqual(bf16_stream * streams + graph + 4096, 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, (Reservation{ .stream_bytes = std.math.maxInt(u64), .streams = 8, .graph_bytes = 0 }).total()); try std.testing.expectError(error.Overflow, mebibytes("18446744073709551615")); + try std.testing.expectError(error.GraphReservationRequired, graphReservation(true, null)); + try std.testing.expectError(error.GraphReservationRequired, graphReservation(true, "0")); + try std.testing.expectEqual(@as(u64, 0), try graphReservation(false, null)); 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, .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()); + const ledger: Ledger = .{ .persistent = 100, .scratch = 20, .deep_prefill = 30, .pool_scratch = 40, .expert_pools = 50, .reservation = .{ .stream_bytes = 10, .streams = 8, .graph_bytes = 60, .after_pool_bytes = 70 } }; + try std.testing.expectEqual(@as(u64, 450), try ledger.total()); } -test "pool setup may consume its allocator margin but not lane or server reservations" { +test "pool setup may consume its allocator margin but not lane graph or server reservations" { const server: u64 = 100; const margin: u64 = 25; const future: u64 = 80; @@ -71,3 +125,79 @@ test "pool setup may consume its allocator margin but not lane or server reserva 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)); } + +test "planner allocation ledger agrees with device state statics work and split buffers" { + const layers: u64 = 49; + const lines: u64 = 10; + const rows: u64 = 16; + const expected = layers * planner_layout.stateWords(lines) * 4 + layers * 512 * 3 * (8 + 8 + 4) + rows * 10 * 9 * 8 + 256 + (512 + 64) * 4 + 2 * rows * 11 * 4; + try std.testing.expectEqual(expected, try plannerBytes(layers, lines, rows)); + try std.testing.expectEqual(try plannerBytes(layers, 0, rows) + layers * lines * 8, try plannerBytes(layers, lines, rows)); + try std.testing.expectError(error.Overflow, plannerBytes(layers, std.math.maxInt(u64), rows)); +} + +test "opt-in demand traces are included in the startup planner reservation" { + try std.testing.expectEqual(@as(u64, 0), try plannerTraceBytes(48, false)); + try std.testing.expectEqual(@as(u64, 48 << 20), try plannerTraceBytes(48, true)); + const base = try planner_layout.deviceBytes(48, 0, 16); + try std.testing.expectEqual(base + (48 << 20), try std.math.add(u64, base, try plannerTraceBytes(48, true))); +} + +test "exl3: ragged planner redistribution conserves each rank's exact device budget" { + const uniform = [_]usize{ 32, 32, 32, 32 }; + const ragged = [_]usize{ 10, 73, 11, 34 }; + const other_rank = [_]usize{ 12, 28, 52, 36 }; + for ([_][]const usize{ &uniform, &ragged, &other_rank }) |counts| { + var state_words: usize = 0; + for (counts) |lines| { + const entry = try planner_layout.PoolState.init(lines, state_words); + state_words = entry.endWords(); + } + const rows = 16; + const work = rows * 10 * 9 * 8 + 256 + (512 + 64) * 4; + const split = 2 * rows * 11 * 4; + const allocated = state_words * 4 + counts.len * planner_layout.static_bytes + work + split; + try std.testing.expectEqual(allocated, try plannerBytesForPools(counts, rows)); + try std.testing.expectEqual(try plannerBytes(uniform.len, uniform[0], rows), try plannerBytesForPools(counts, rows)); + const line_bytes = 3 << 20; + const fixed = try plannerBytesForPools(counts, rows); + var total = fixed; + for (counts) |lines| total += lines * line_bytes; + try std.testing.expectEqual(try poolBytes(32, line_bytes, 4, fixed), total); + } + try std.testing.expectError(error.Overflow, plannerBytesForPools(&.{ std.math.maxInt(usize), std.math.maxInt(usize) }, 16)); + try std.testing.expectError(error.Overflow, plannerBytesForPools(&ragged, std.math.maxInt(u64))); +} + +test "lane admission protects the full graph reservation without changing BF16 stream sizes" { + const stream_bytes: u64 = 262144 * 256 * 2; + const graph_bytes: u64 = 512 << 20; + for ([_]u64{ 1, 8 }) |streams| { + const requested = try std.math.mul(u64, streams, stream_bytes); + try std.testing.expectEqual(requested, laneRoom(requested, 0)); + try std.testing.expectEqual(requested, laneRoom(try std.math.add(u64, requested, graph_bytes), graph_bytes)); + try std.testing.expect(laneRoom(requested, graph_bytes) < requested); + } + try std.testing.expectEqual(@as(u64, 0), laneRoom(graph_bytes - 1, graph_bytes)); + try std.testing.expectEqual(@as(u64, 0), laneRoom(graph_bytes, graph_bytes)); + try std.testing.expectEqual(@as(u64, 0), laneRoom(0, graph_bytes)); +} + +test "lane room subtraction cannot overflow at integer boundaries" { + const maximum = std.math.maxInt(u64); + try std.testing.expectEqual(maximum, laneRoom(maximum, 0)); + try std.testing.expectEqual(@as(u64, 0), laneRoom(maximum, maximum)); + try std.testing.expectEqual(@as(u64, 0), laneRoom(1, maximum)); + try std.testing.expectEqual(maximum - 1, laneRoom(maximum, 1)); + try std.testing.expectError(error.Overflow, (Reservation{ .stream_bytes = maximum, .streams = 1, .graph_bytes = 1 }).total()); +} + +test "per rank protection precedes agreement of lane counts" { + const graph_bytes: u64 = 512 << 20; + const rank_zero_stream: u64 = 128 << 20; + const rank_one_stream: u64 = 64 << 20; + const rank_zero_room = laneRoom(8 * rank_zero_stream + graph_bytes, graph_bytes); + const rank_one_room = laneRoom(8 * rank_one_stream + graph_bytes, graph_bytes); + try std.testing.expectEqual(@as(u64, 8), @min(rank_zero_room / rank_zero_stream, rank_one_room / rank_one_stream)); + try std.testing.expectEqual(@as(u64, 7), @min(rank_zero_room / rank_zero_stream, laneRoom(8 * rank_one_stream + graph_bytes - 1, graph_bytes) / rank_one_stream)); +} diff --git a/zig/src/families/flashnext/exl3/prefill/cuda_exl3_prefill.zig b/zig/src/families/flashnext/exl3/prefill/cuda_exl3_prefill.zig index 87f443cd48..0b2b83a0c9 100644 --- a/zig/src/families/flashnext/exl3/prefill/cuda_exl3_prefill.zig +++ b/zig/src/families/flashnext/exl3/prefill/cuda_exl3_prefill.zig @@ -23,7 +23,7 @@ const Deep = state.Deep; const HostLayout = state.HostLayout; const Layout = state.Layout; const Persistent = state.Persistent; -const opts = state.opts; +const Runtime = state.Runtime; const chunk = chunks.chunk; const selectWriteback = chunks.selectWriteback; const twoChunk = chunks.twoChunk; @@ -31,8 +31,9 @@ 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| { +pub fn release(runtime: *Runtime) void { + if (runtime.deep) |dp| { + @import("../graph/cuda_exl3_graph.zig").clearPrefillRuntime(@intFromPtr(runtime)); 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"); @@ -58,15 +59,16 @@ pub fn release() void { dp.host.free(); dp.dev.free(); dp.gpa.destroy(dp); - state.deep = null; + runtime.deep = null; } - if (agree_buf) |*b| b.free(); - agree_buf = null; + if (runtime.agree_buf) |*b| b.free(); + runtime.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(); + const runtime = &e.exl3_prefill; + const o = runtime.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; @@ -76,14 +78,15 @@ pub fn reserveFor(e: *prompt.Engine, context: usize, tp: usize) void { } 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 }); + reserveAt(e.gpa, e.driver, runtime, 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 }); } /// A serving step's batched `base` rows, or the deep chunk when this engine holds deep buffers. pub fn stepRows(e: *const prompt.Engine, base: usize, draft_rows: ?usize) usize { + const runtime = &e.exl3_prefill; var n = base; - 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 (runtime.deep) |dp| if (dp.owner == @as(*const anyopaque, e) and runtime.opts().enabled) { + n = @max(n, stepCapacity(dp.rows, runtime.opts().two_chunk and dp.second != null)); }; if (draft_rows) |r| n = @min(n, @max(r, 1)); return n; @@ -95,9 +98,10 @@ pub fn stepCapacity(rows: usize, paired: bool) usize { /// 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 { + const runtime = &e.exl3_prefill; var n = stepRows(e, e.chunk, null); - if (state.deep) |dp| if (dp.owner == @as(*const anyopaque, e)) { - const o = opts(); + if (runtime.deep) |dp| if (dp.owner == @as(*const anyopaque, e)) { + const o = runtime.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)); } @@ -106,26 +110,26 @@ pub fn draftRows(e: *const prompt.Engine) usize { } /// 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(); +pub fn releaseFor(runtime: *Runtime, owner: *const anyopaque) void { + if (runtime.deep) |dp| if (dp.owner == owner) release(runtime); } /// 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); +pub fn reserve(gpa: std.mem.Allocator, d: *const cuda.Driver, runtime: *Runtime, owner: *const anyopaque, rows_in: usize) !void { + return reserveAt(gpa, d, runtime, 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 { +pub fn reserveAt(gpa: std.mem.Allocator, d: *const cuda.Driver, runtime: *Runtime, 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 (runtime.deep) |dp| { if (dp.owner == owner and dp.rows >= rows and dp.world == world) return; - release(); + release(runtime); } 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.* = .{ .runtime = runtime, .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); @@ -148,16 +152,22 @@ pub fn reserveAt(gpa: std.mem.Allocator, d: *const cuda.Driver, owner: *const an dp.red_go = try cuda.Event.init(d, false); dp.red = rs; } - state.deep = dp; + runtime.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(); +pub fn reserveTwoChunk(optional: ?*Runtime, d: *const cuda.Driver, free: u64, protected: u64, graph_reservation: u64) !u64 { + const runtime = optional orelse return 0; + const o = runtime.opts(); if (!o.two_chunk or !o.enabled or !o.ring or o.reduced_branch) return 0; - const dp = state.deep orelse return 0; + const dp = runtime.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 graph_required = @import("../graph/cuda_exl3_graph_memory.zig").twoChunkReservationBytes(); + if (@import("../graph/cuda_exl3_graph.zig").enabled() and (graph_required == 0 or graph_reservation < graph_required)) { + std.log.info("flashnext exl3 prefill: two-chunk window declined, incremental graph reservation unavailable (required={d} supplied={d})", .{ graph_required, graph_reservation }); + 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); @@ -200,12 +210,14 @@ pub fn reserveTwoChunk(d: *const cuda.Driver, free: u64, protected: u64) !u64 { 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; + dp.window_graph_reservation = graph_reservation; 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; +pub fn twoChunkAllocated(optional: ?*const Runtime) bool { + const runtime = optional orelse return false; + return if (runtime.deep) |buffers| buffers.second != null else false; } fn applies(e: *prompt.Engine, seq: *prompt.Seq, ids: []const u32, hooks: batch.Hooks, o: Options) bool { @@ -229,11 +241,12 @@ fn worldOf(e: *const prompt.Engine) usize { /// 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 runtime = &e.exl3_prefill; 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.?; + if (runtime.agree_buf == null) runtime.agree_buf = try cuda.DeviceBuffer.alloc(e.driver, 256); + const b = runtime.agree_buf.?; var v: [2]f32 = .{ 0, 0 }; v[rank] = @floatFromInt(mine); try st.synchronize(); @@ -244,13 +257,14 @@ fn agree(e: *prompt.Engine, mine: usize) !usize { 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(); + const runtime = &e.exl3_prefill; + const o = runtime.opts(); if (!applies(e, seq, ids, hooks, o)) return false; + try @import("../dplan/cuda_exl3_dplan.zig").toHost(&e.exl3_dplan); // the pools' residency on the host before planning a prompt 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; @@ -259,11 +273,11 @@ pub fn forward(e: *prompt.Engine, comptime Tri: type, tri: Tri, seq: *prompt.Seq 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; + const have = if (runtime.deep) |d| d.owner == @as(*const anyopaque, e) and d.world == world else false; + var room: usize = if (have) runtime.deep.?.rows else 0; if (!have) { - if (reserveAt(e.gpa, e.driver, e, rows, world)) |_| { - room = state.deep.?.rows; + if (reserveAt(e.gpa, e.driver, runtime, e, rows, world)) |_| { + room = runtime.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) @@ -271,7 +285,7 @@ pub fn forward(e: *prompt.Engine, comptime Tri: type, tri: Tri, seq: *prompt.Seq 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.?; + const dp = runtime.deep.?; try selectWriteback(dp, o.reduced_branch and world == 2); const scores = try e.reserveScores(seq.cap, s.cap); var p: batch.Pass(Tri) = .{ @@ -281,7 +295,7 @@ pub fn forward(e: *prompt.Engine, comptime Tri: type, tri: Tri, seq: *prompt.Seq .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 }, + .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, .dplan = &e.exl3_dplan, .collective = e.collective }, .seq = seq.view(scores), .hooks = hooks, }; @@ -292,14 +306,18 @@ pub fn forward(e: *prompt.Engine, comptime Tri: type, tri: Tri, seq: *prompt.Seq 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; + var 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; + if (@import("../graph/cuda_exl3_graph.zig").enabled()) { + const required = @import("../graph/cuda_exl3_graph_memory.zig").twoChunkReservationBytes(); + if (required == 0 or dp.window_graph_reservation < required) supported = false; + } 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; + runtime.tiles_on = true; } - defer tiles.tiles_on = false; + defer runtime.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; @@ -326,7 +344,9 @@ 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 fn ringTelemetry(e: *const prompt.Engine) ?RingTelemetry { + return state.ringTelemetry(&e.exl3_prefill, e); +} pub const deviceBytes = state.deviceBytes; pub const deviceBytesAt = state.deviceBytesAt; pub const head = heads.head; 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 index 632bf87cdd..2cc99f12a5 100644 --- a/zig/src/families/flashnext/exl3/prefill/cuda_exl3_prefill_chunk.zig +++ b/zig/src/families/flashnext/exl3/prefill/cuda_exl3_prefill_chunk.zig @@ -30,7 +30,6 @@ 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; @@ -63,6 +62,7 @@ pub fn selectWriteback(dp: *Deep, reduced_branch: bool) !void { @memset(weights, 1); dp.unit_weights = try cuda.DeviceBuffer.fromHost(dp.dev.d, std.mem.sliceAsBytes(weights)); } + if (dp.reduced_branch != reduced_branch) @import("../graph/cuda_exl3_graph.zig").clearPrefill(); dp.reduced_branch = reduced_branch; } @@ -140,7 +140,10 @@ pub fn windowInputs(dp: *Deep, p: anytype, ids: []const u32, pos0: usize, layer: } 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 graph = @import("../graph/cuda_exl3_graph.zig"); const lw = &p.m.layers[layer]; + const staged = try stagePages(p, lw.linear, layer, pos0, n); + defer if (staged) @import("../graph/cuda_exl3_graph_pages.zig").finish(p.seq.sel.?); 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); @@ -148,16 +151,35 @@ fn windowDense(comptime Tri: type, dp: *Deep, p: *batch.Pass(Tri), base: *batch. 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); - } + const capture = try graph.prefillBegin(p, lw.linear, layer, n, sub, pos0, dp.dev.ptr, base, dp.runtime.tiles_on, waits, dp.dptr(dp.lay.pos)); + const Blocks = struct { + dp: *Deep, + p: *batch.Pass(Tri), + views: []batch.Scratch, + waits: []const cuda.Event, + layer: usize, + n: usize, + sub: usize, + pos0: usize, + pub fn run(b: @This(), mode: @TypeOf(capture)) !void { + const lw_ = &b.p.m.layers[b.layer]; + for (0..b.views.len) |block| { + const at = block * b.sub; + const count = @min(b.sub, b.n - at); + b.p.s = &b.views[block]; + b.p.rows = count; + b.p.pos0 = b.pos0 + at; + pendingWriteback(b.p, b.dp, if (b.layer == 0) null else b.layer - 1, at); + if (block < b.waits.len) switch (mode) { + .eager => try b.p.stream.wait(b.waits[block]), + .capturing => try b.p.stream.waitExternal(b.waits[block]), + .replayed => unreachable, + }; + try layerBlock(Tri, b.p, lw_, b.layer, count); + } + } + }; + try graph.prefillLayer(p, capture, Blocks{ .dp = dp, .p = p, .views = views[0..n_sub], .waits = waits, .layer = layer, .n = n, .sub = sub, .pos0 = pos0 }); dp.red_n = 0; } @@ -166,8 +188,8 @@ fn windowFinish(comptime Tri: type, dp: *Deep, p: *batch.Pass(Tri), base: *batch 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; + dp.runtime.stats.chunks += 1; + dp.runtime.stats.rows += n; const last_block = n_sub - 1; const last_rows = n - last_block * sub; if (streams_out != 0) { @@ -302,23 +324,43 @@ pub fn chunk(comptime Tri: type, dp: *Deep, p: *batch.Pass(Tri), base: *batch.Sc 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); - } + const staged = try stagePages(p, lw.linear, layer, pos0, n); + defer if (staged) @import("../graph/cuda_exl3_graph_pages.zig").finish(p.seq.sel.?); + const pg = try @import("../graph/cuda_exl3_graph.zig").prefillBegin(p, lw.linear, layer, n, sub, pos0, dp.dev.ptr, base, dp.runtime.tiles_on, dp.red_ev[0..dp.red_n], dp.dptr(dp.lay.pos)); + const Blocks = struct { + dp: *Deep, + p: *batch.Pass(Tri), + views: []batch.Scratch, + lw: *const @TypeOf(lw.*), + layer: usize, + n: usize, + sub: usize, + pos0: usize, + pub fn run(b: @This(), mode: @TypeOf(pg)) !void { + for (0..b.views.len) |blk| { + const m = @min(b.sub, b.n - blk * b.sub); + b.p.s = &b.views[blk]; + b.p.rows = m; + b.p.pos0 = b.pos0 + blk * b.sub; + pendingWriteback(b.p, b.dp, if (b.layer > 0) b.layer - 1 else null, blk * b.sub); + if (blk < b.dp.red_n) switch (mode) { + .eager => try b.p.stream.wait(b.dp.red_ev[blk]), + .capturing => try b.p.stream.waitExternal(b.dp.red_ev[blk]), + .replayed => unreachable, + }; + try layerBlock(Tri, b.p, b.lw, b.layer, m); + } + } + }; + try @import("../graph/cuda_exl3_graph.zig").prefillLayer(p, pg, Blocks{ .dp = dp, .p = p, .views = views[0..n_sub], .lw = lw, .layer = layer, .n = n, .sub = sub, .pos0 = pos0 }); 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; + dp.runtime.stats.chunks += 1; + dp.runtime.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; @@ -388,3 +430,10 @@ fn layerBlock(comptime Tri: type, p: *batch.Pass(Tri), lw: *const @import("../.. 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); } + +fn stagePages(p: anytype, linear: bool, layer: usize, position: usize, rows: usize) !bool { + if (linear or !@import("../graph/cuda_exl3_graph.zig").enabled()) return false; + const selection = p.seq.sel orelse return false; + try @import("../graph/cuda_exl3_graph_pages.zig").stage(selection, p.stream, p.seq.sel_layer + wt.attnIndex(layer), 1, position, rows); + return true; +} 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 index f754048c1c..0acd642f32 100644 --- a/zig/src/families/flashnext/exl3/prefill/cuda_exl3_prefill_head.zig +++ b/zig/src/families/flashnext/exl3/prefill/cuda_exl3_prefill_head.zig @@ -16,19 +16,17 @@ 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 op = o.e.exl3_prefill.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 (!scores and o.e.exl3_prefill.deep != null and n > o.e.exl3_prefill.deep.?.rows) { + const first = headSplit(n, o.e.exl3_prefill.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; @@ -56,16 +54,17 @@ test "exl3: a paired window's head splits into two deep slices that cannot decli } 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 op = o.e.exl3_prefill.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; + const dp = o.e.exl3_prefill.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; + try @import("../dplan/cuda_exl3_dplan.zig").toHost(&e.exl3_dplan); // the head's pool may be device-planned (cuda_exl3_chain.zig) const mtp = @import("../../cuda_mtp.zig"); const tri = o.tri; const want = @max(@as(usize, 32), 2 * @min(n, mtp.piece_rows)); @@ -82,7 +81,7 @@ fn headSlice(comptime Tri: type, o: anytype, d: *@import("../../cuda_mtp.zig").D .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 }, + .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, .dplan = &e.exl3_dplan }, .seq = d.headSeq(null, score_rows), .hooks = .{ .logits = .none }, }; @@ -91,9 +90,9 @@ fn headSlice(comptime Tri: type, o: anytype, d: *@import("../../cuda_mtp.zig").D defer p.s = s; if (op.tiles and !op.slots_kernel) { try warmTiles(e); - tiles.tiles_on = true; + dp.runtime.tiles_on = true; } - defer tiles.tiles_on = false; + defer dp.runtime.tiles_on = false; const lw = &h.layer; for (0..n_pieces) |b| { const at = b * piece; @@ -137,7 +136,7 @@ fn headSlice(comptime Tri: type, o: anytype, d: *@import("../../cuda_mtp.zig").D 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; + dp.runtime.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; 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 index 7e6ee05083..5abee280d4 100644 --- a/zig/src/families/flashnext/exl3/prefill/cuda_exl3_prefill_moe.zig +++ b/zig/src/families/flashnext/exl3/prefill/cuda_exl3_prefill_moe.zig @@ -21,7 +21,6 @@ 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). @@ -51,9 +50,9 @@ pub fn moe(comptime Tri: type, dp: *Deep, p: *batch.Pass(Tri), lw: *const @impor 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; + dp.runtime.stats.host_wait_ns += @intCast(t1 - t0); + dp.runtime.stats.plan_ns += @intCast(t2 - t1); + dp.runtime.stats.members += counts.members; const pool = x.pool; var fake_rows: usize = 0; for (dp.passes[0..dp.n_passes]) |ps| { @@ -65,20 +64,20 @@ pub fn moe(comptime Tri: type, dp: *Deep, p: *batch.Pass(Tri), lw: *const @impor 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_(); + dp.runtime.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; + dp.runtime.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; + dp.runtime.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) { @@ -137,9 +136,9 @@ pub fn rankSum(dp: *Deep, p: anytype, x: *const load.Experts, n: usize, sub: usi if (profile()) { if (dp.red) |rs| try c.finish(rs); try c.finish(st); - stats.reduce_ns += @intCast(now(p.io) - t0); + dp.runtime.stats.reduce_ns += @intCast(now(p.io) - t0); } - stats.reduces += 1; + dp.runtime.stats.reduces += 1; } fn overlap() bool { @@ -183,8 +182,8 @@ pub fn prefetch(dp: *Deep, pool: *pool_.Pool, own: ?tp_.Rank, o: Options) !void // 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_(); + dp.runtime.stats.prefetched += r.misses_(); + dp.runtime.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 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 index abf2c35893..498be57763 100644 --- a/zig/src/families/flashnext/exl3/prefill/cuda_exl3_prefill_ringmoe.zig +++ b/zig/src/families/flashnext/exl3/prefill/cuda_exl3_prefill_ringmoe.zig @@ -20,7 +20,6 @@ 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; @@ -155,9 +154,9 @@ pub fn ringMoe(comptime Tri: type, dp: *Deep, r: *ring_.Ring, p: *batch.Pass(Tri 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; + dp.runtime.stats.host_wait_ns += @intCast(t1 - t0); + dp.runtime.stats.plan_ns += @intCast(t2 - t1 - staging_ns); + dp.runtime.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| { @@ -172,9 +171,9 @@ pub fn ringMoe(comptime Tri: type, dp: *Deep, r: *ring_.Ring, p: *batch.Pass(Tri 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; + dp.runtime.stats.passes += 1; } - stats.copies += r.stats.copies - copies0; + dp.runtime.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| { @@ -182,8 +181,8 @@ pub fn ringMoe(comptime Tri: type, dp: *Deep, r: *ring_.Ring, p: *batch.Pass(Tri 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; + dp.runtime.stats.prefetched += r.stats.copies - c1; + dp.runtime.stats.copies += r.stats.copies - c1; p.pending = layer; } 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 index 8d2a537246..ddc68cdf66 100644 --- a/zig/src/families/flashnext/exl3/prefill/cuda_exl3_prefill_state.zig +++ b/zig/src/families/flashnext/exl3/prefill/cuda_exl3_prefill_state.zig @@ -61,14 +61,6 @@ pub fn optionsFromEnv() Options { 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, @@ -83,7 +75,17 @@ pub const Stats = struct { 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 = .{}; +pub const Runtime = struct { + options: ?Options = null, + stats: Stats = .{}, + deep: ?*Deep = null, + agree_buf: ?cuda.DeviceBuffer = null, + tiles_on: bool = false, + + pub fn opts(runtime: *const Runtime) Options { + return runtime.options orelse optionsFromEnv(); + } +}; /// 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 { @@ -216,6 +218,7 @@ pub const Persistent = struct { }; pub const Deep = struct { + runtime: *Runtime = undefined, gpa: std.mem.Allocator, owner: *const anyopaque, rows: usize, @@ -251,6 +254,7 @@ pub const Deep = struct { history_shape: ?history_.Shape = null, second: ?*Deep = null, shared: ?*const Deep = null, + window_graph_reservation: u64 = 0, pub fn dptr(dp: *const Deep, off: usize) u64 { if (dp.shared) |source| { @@ -268,12 +272,10 @@ pub const Deep = struct { } }; -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; +pub fn ringTelemetry(runtime: *const Runtime, owner: *const anyopaque) ?RingTelemetry { + const buffers = runtime.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 }; 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 index 847e5e178e..270531fcb3 100644 --- a/zig/src/families/flashnext/exl3/prefill/cuda_exl3_prefill_test.zig +++ b/zig/src/families/flashnext/exl3/prefill/cuda_exl3_prefill_test.zig @@ -104,13 +104,14 @@ test "exl3: two-chunk arenas retain hidden and deferred HC operands without alia } test "exl3: two-chunk memory refusal performs no CUDA allocation" { - const saved_deep = state.deep; - const saved_options = state.options; + var runtime: state.Runtime = .{}; + const graph = @import("../graph/cuda_exl3_graph.zig"); + const saved_force = graph.force; defer { - state.deep = saved_deep; - state.options = saved_options; + graph.force = saved_force; } - state.options = .{ .two_chunk = true }; + graph.force = .off; + runtime.options = .{ .two_chunk = true }; var driver: cuda.Driver = undefined; var buffers: Deep = undefined; buffers.dev = .{ .d = &driver, .ptr = 0x100000, .len = 1 }; @@ -118,8 +119,8 @@ test "exl3: two-chunk memory refusal performs no CUDA allocation" { 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)); + runtime.deep = &buffers; + try std.testing.expectEqual(@as(u64, 0), try reserveTwoChunk(&runtime, &driver, Persistent.of(4096).total - 1, 0, 0)); try std.testing.expect(buffers.second == null); } @@ -237,9 +238,10 @@ test "exl3: prefetch history follows deep engine ownership reuse replacement and 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 { + fn install(runtime: *state.Runtime, driver: *const cuda.Driver, owner: *const anyopaque) !*Deep { const buffers = try std.testing.allocator.create(Deep); buffers.* = .{ + .runtime = runtime, .gpa = std.testing.allocator, .owner = owner, .rows = 4096, @@ -251,12 +253,13 @@ test "exl3: prefetch history follows deep engine ownership reuse replacement and .aux = .{ .d = driver, .handle = null }, }; buffers.kern.mod = .{ .d = driver, .handle = null }; - state.deep = buffers; + runtime.deep = buffers; return buffers; } }; - try std.testing.expect(state.deep == null); - defer release(); + var runtime: state.Runtime = .{}; + try std.testing.expect(runtime.deep == null); + defer release(&runtime); var driver: cuda.Driver = undefined; driver.api.cuModuleUnload = Fixture.unload; driver.api.cuStreamSynchronize = Fixture.stream; @@ -264,28 +267,38 @@ test "exl3: prefetch history follows deep engine ownership reuse replacement and 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]); + const first = try Fixture.install(&runtime, &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]); + releaseFor(&runtime, &owners[1]); + try std.testing.expect(runtime.deep.? == first); + try reserveAt(std.testing.allocator, &driver, &runtime, &owners[0], 2048, 1); + try std.testing.expectEqual(@as(usize, 96), history_.want(runtime.deep.?.history[3], 4096, 400, routed_n_)); + try std.testing.expectError(error.NotReady, reserveAt(std.testing.allocator, &driver, &runtime, &owners[1], 4096, 1)); + try std.testing.expect(runtime.deep == null); + const second = try Fixture.install(&runtime, &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.expectError(error.NotReady, reserveAt(std.testing.allocator, &driver, &runtime, &owners[1], 4096, 2)); + try std.testing.expect(runtime.deep == null); + const third = try Fixture.install(&runtime, &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]); + releaseFor(&runtime, &owners[1]); + try std.testing.expect(runtime.deep == null); + const fourth = try Fixture.install(&runtime, &driver, &owners[1]); try std.testing.expectEqual(@as(usize, 416), history_.want(fourth.history[3], 4096, 400, routed_n_)); + var other_runtime: state.Runtime = .{}; + defer release(&other_runtime); + const other = try Fixture.install(&other_runtime, &driver, &owners[0]); + releaseFor(&runtime, &owners[1]); + try std.testing.expect(runtime.deep == null); + try std.testing.expect(other_runtime.deep.? == other); + releaseFor(&other_runtime, &owners[1]); + try std.testing.expect(other_runtime.deep.? == other); + releaseFor(&other_runtime, &owners[0]); + try std.testing.expect(other_runtime.deep == null); } test "exl3: the prefill tiles cut K as decode does at Flash Next's shapes, and the two expert groups split at 256" { 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 index b32cf381d4..717340543a 100644 --- a/zig/src/families/flashnext/exl3/prefill/cuda_exl3_prefill_tiles.zig +++ b/zig/src/families/flashnext/exl3/prefill/cuda_exl3_prefill_tiles.zig @@ -4,13 +4,11 @@ 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 { +pub fn tilesPart(tiles_on: bool, 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; diff --git a/zig/src/families/flashnext/exl3/prefill/exl3_prefill_window.zig b/zig/src/families/flashnext/exl3/prefill/exl3_prefill_window.zig index 353a4b84a0..23ea4e3a8c 100644 --- a/zig/src/families/flashnext/exl3/prefill/exl3_prefill_window.zig +++ b/zig/src/families/flashnext/exl3/prefill/exl3_prefill_window.zig @@ -1,6 +1,12 @@ const std = @import("std"); pub const max_rows = 4096; +pub const graph_entries = 112; + +pub fn graphBytes(exec: u64, template: u64, transient: u64) !u64 { + const retained = try std.math.mul(u64, graph_entries * 2, try std.math.add(u64, exec, template)); + return std.math.add(u64, retained, try std.math.mul(u64, 2, transient)); +} pub fn enabled(value: ?[]const u8) bool { return if (value) |text| std.mem.eql(u8, text, "1") else false; @@ -93,4 +99,6 @@ test "exl3: two chunk admission protects all existing reservations and rejects o 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)); + try std.testing.expectEqual(@as(u64, 288 << 20), try graphBytes(1 << 20, 0, 32 << 20)); + try std.testing.expectError(error.Overflow, graphBytes(std.math.maxInt(u64), 0, 0)); } diff --git a/zig/src/families/flashnext/exl3/tp/cuda_exl3_tp.zig b/zig/src/families/flashnext/exl3/tp/cuda_exl3_tp.zig index 9cbd489f57..a8c4ad9be2 100644 --- a/zig/src/families/flashnext/exl3/tp/cuda_exl3_tp.zig +++ b/zig/src/families/flashnext/exl3/tp/cuda_exl3_tp.zig @@ -270,7 +270,13 @@ pub fn moe(p: anytype, lw: anytype, x: *const load.Experts) !void { 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); + const dplan = @import("../dplan/cuda_exl3_dplan.zig"); + const plans = p.m.dplan; + if (plans) |state| if (state.active) |dp| if (dp.indexOf(pool)) |i| { // a decode round planned on the device (no host cut) + try dplan.runLayer(dp, st, i, ex, x.layer, xin, mp + r0 * slots * 4, rows, sc, 0, 0); + continue; + }; + try @import("../model/cuda_exl3_pass.zig").pooled(st, x.ops, pool, plans, 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); diff --git a/zig/src/native/draft_sources_cuda_test.zig b/zig/src/native/draft_sources_cuda_test.zig index 2f652631a6..bdeec01c2f 100644 --- a/zig/src/native/draft_sources_cuda_test.zig +++ b/zig/src/native/draft_sources_cuda_test.zig @@ -95,6 +95,9 @@ test "draft sources: CUDA Qwen dense and MoE share copy-only target facts" { test "draft sources: CUDA Flash Next follows the loaded checkpoint head" { var b: @import("flashnext").native.Cuda = undefined; + var engine: @typeInfo(@TypeOf(b.engine)).pointer.child = undefined; + engine.n_layers = 0; + b.engine = &engine; b.head = null; b.depth = 8; try check(b.facts(), null, true, false); diff --git a/zig/tests/cuda/contract_flashnext_exl3_dplan.zig b/zig/tests/cuda/contract_flashnext_exl3_dplan.zig new file mode 100644 index 0000000000..1931e87462 --- /dev/null +++ b/zig/tests/cuda/contract_flashnext_exl3_dplan.zig @@ -0,0 +1,419 @@ +//! Flash Next EXL3 device pool plan under gates 1-4: the planner against a host model, and the kernel list. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const rig_ = @import("contract_rig.zig"); +const u = @import("contract_flashnext_exl3_dplan_util.zig"); +const c = core.contract.spec; +const g = core.contract.gates; +const noSums = u.noSums; +const put = u.put; +const fetch = u.fetch; +const same = u.same; +const experts = u.experts; +const check = @import("check.zig"); +const tp_ = @import("contract_flashnext_exl3_tp.zig"); +const small = @import("contract_flashnext_exl3_dplan_copy.zig"); +const draft = @import("contract_flashnext_exl3_dplan_draft.zig"); +const stateWords = u.stateWords; +const recent_rows_words = u.recent_rows_words; + +/// The planner's state, tables and outputs after one round: the kernel's formula on the same words. +const Model = struct { + st: []u32, + tab: [3][experts]i64, + k2t: [3][experts]i32, + work: []i64, + hit: []i32, + miss: []i32, + ovf: []i32, + work_n: i32 = 0, + ovf_n: i32 = 0, + + fn colder(ca: u32, sa: u32, ia: usize, cb: u32, sb: u32, ib: usize) bool { + if (ca != cb) return ca < cb; + if (sa != sb) return sa < sb; + return ia > ib; + } + + fn round(m: *Model, lines: usize, recent: bool, picks: []const i32, src: []const i64, lens: []const i64, k2: []const i32, line_base: u64, line_bytes: u64, offs: [3]u64) void { + const line_id: []i32 = @ptrCast(m.st[0..lines]); + const stamp = m.st[lines .. 2 * lines]; + const line_of: []i32 = @ptrCast(m.st[2 * lines ..][0..experts]); + const counts = m.st[2 * lines + experts ..][0..experts]; + const recent_counts = m.st[2 * lines + 2 * experts ..][0..experts]; + const recent_rows = m.st[2 * lines + 3 * experts ..][0..recent_rows_words]; + const misc = m.st[2 * lines + 3 * experts + recent_rows_words ..][0..10]; + const score = if (recent) recent_counts else counts; + const n = picks.len; + misc[0] += 1; + const now = misc[0]; + misc[1] += 1; + if (recent) for (0..n / 11) |row| { + const expired = recent_rows[misc[8] * (experts / 32) ..][0 .. experts / 32]; + for (expired, 0..) |*word, w| { + var removed = word.*; + while (removed != 0) : (removed &= removed - 1) recent_counts[w * 32 + @ctz(removed)] -%= 1; + word.* = 0; + } + for (0..11) |slot| { + const e = picks[row * 11 + slot]; + if (e < 0 or e >= experts) continue; + const bit = @as(u32, 1) << @intCast(e & 31); + if (expired[@intCast(e >> 5)] & bit == 0) { + expired[@intCast(e >> 5)] |= bit; + recent_counts[@intCast(e)] += 1; + } + } + misc[8] = (misc[8] + 1) % 256; + if (misc[9] < 256) misc[9] += 1; + }; + var uniq: [256]u16 = undefined; + var app: [256]u16 = undefined; + var miss: [256]u16 = undefined; + var keep: [experts / 32]u32 = @splat(0); + var nu: usize = 0; + var over = false; + for (picks) |p| { + if (p < 0 or p >= experts) continue; + const e: usize = @intCast(p); + if (counts[e] != 0xffffffff) counts[e] += 1; + if (line_of[e] >= 0) misc[2] += 1 else misc[3] += 1; + var j: usize = 0; + while (j < nu and uniq[j] != e) j += 1; + if (j == nu) { + if (nu == 256) { + over = true; + continue; + } + uniq[nu] = @intCast(e); + app[nu] = 0; + keep[e >> 5] |= @as(u32, 1) << @intCast(e & 31); + nu += 1; + } + app[j] += 1; + } + var nm: usize = 0; + for (0..nu) |j| { + const ln = line_of[uniq[j]]; + if (ln >= 0) stamp[@intCast(ln)] = now else { + miss[nm] = @intCast(j); + nm += 1; + } + } + for (1..nm) |a| { + const x = miss[a]; + var b: isize = @as(isize, @intCast(a)) - 1; + while (b >= 0) : (b -= 1) { + const y = miss[@intCast(b)]; + const cx = score[uniq[x]]; + const cy = score[uniq[y]]; + const hotter = if (cx != cy) cx > cy else if (app[x] != app[y]) app[x] > app[y] else uniq[x] < uniq[y]; + if (!hotter) break; + miss[@intCast(b + 1)] = y; + } + miss[@intCast(b + 1)] = x; + } + if (over) misc[6] += 1; + const n_miss = if (over) 0 else nm; + for (picks, 0..) |p, i| { + const held = p >= experts or (p >= 0 and line_of[@intCast(p)] >= 0); + m.hit[i] = if (p < 0) -1 else if (held) p else -1; + m.miss[i] = if (p < 0) -1 else if (held) -1 else p; + } + for (miss[0..n_miss]) |mi| { + var victim: ?usize = null; + for (line_id, 0..) |id, i| if (id < 0) { + victim = i; + break; + }; + if (victim == null) for (line_id, 0..) |id, i| { + const e: usize = @intCast(id); + if (keep[e >> 5] & (@as(u32, 1) << @intCast(e & 31)) != 0) continue; + if (victim == null or colder(score[e], stamp[i], i, score[@intCast(line_id[victim.?])], stamp[victim.?], victim.?)) victim = i; + }; + const id: usize = uniq[mi]; + if (victim) |v| { + const old = line_id[v]; + if (old >= 0) { + line_of[@intCast(old)] = -1; + for (0..3) |w| { + m.tab[w][@intCast(old)] = 0; + m.k2t[w][@intCast(old)] = 0; + } + misc[5] += 1; + } + line_id[v] = @intCast(id); + line_of[id] = @intCast(v); + stamp[v] = now; + const base = line_base + @as(u64, v) * line_bytes; + for (0..3) |w| { + const dst: i64 = @bitCast(base + offs[w]); + m.tab[w][id] = dst; + m.k2t[w][id] = k2[id * 3 + w]; + m.work[@as(usize, @intCast(m.work_n)) * 9 + w * 3 ..][0..3].* = .{ src[id * 3 + w], dst, lens[id * 3 + w] }; + } + m.work_n += 1; + misc[4] += 1; + } else { + for (0..3) |w| { + m.tab[w][id] = src[id * 3 + w]; + m.k2t[w][id] = k2[id * 3 + w]; + } + m.ovf[@intCast(m.ovf_n)] = @intCast(id); + m.ovf_n += 1; + misc[7] += 1; + } + } + } +}; + +pub const Plan = struct { + rig: rig_.Rig, + f: cuda.Function, + recent: bool = false, + lines: usize = 0, + picks: []i32 = &.{}, + st0: []u32 = &.{}, + tab0: [3][experts]i64 = undefined, + k2t0: [3][experts]i32 = undefined, + src: [experts * 3]i64 = undefined, + lens: [experts * 3]i64 = undefined, + k2: [experts * 3]i32 = undefined, + // device buffers: picks src lens k2 | st tab0-2 k2t0-2 work work_n hit miss ovf ovf_n + d: [17]u64 = @splat(0), + msg: [160]u8 = undefined, + const line_base: u64 = 0x7000_0000_0000; + const line_bytes: u64 = 1 << 20; + const offs: [3]u64 = .{ 0, 4096, 8192 }; + + pub fn deinit(self: *Plan) void { + const a = self.rig.gpu.gpa; + a.free(self.picks); + a.free(self.st0); + self.rig.deinit(); + } + + const vt = g.hooks(Plan); + + pub fn kernel(self: *Plan) g.Kernel { + return .{ + .spec = .{ + .name = "cuda Flash Next EXL3 fn_exl3_dplan", + .accum = .none, + .depth = noSums, + .out = .raw, + .nan = .{ .stated = "integer picks only: a pick outside [0, 512) is skipped (512 and above counts as held, as the shared expert), and no pick changes a table entry it does not plan" }, + // the victim search strides the lines by the 256 threads of the block + .tiles = .{ .{}, .{ .size = 256, .step = 1, .min = 1 }, .{} }, + .served = &.{ .{ .rows = 1, .n = 38, .k = 0 }, .{ .rows = 4, .n = 38, .k = 0 }, .{ .rows = 16, .n = 160, .k = 0 }, .{ .rows = 16, .n = 6, .k = 0 } }, + }, + .ctx = self, + .vt = &vt, + }; + } + + pub fn prepare(self: *Plan, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows == 0 or shape.rows > 24 or shape.n == 0 or shape.n > 2048) return false; + self.rig.reset(); + const a = self.rig.gpu.gpa; + a.free(self.picks); + a.free(self.st0); + const lines = shape.n; + self.lines = lines; + self.recent = seed & 1 == 1; + var prng = std.Random.DefaultPrng.init(seed ^ 0x706c616e); + const rng = prng.random(); + self.picks = try a.alloc(i32, shape.rows * 11); + const edge = [_]i32{ -1, 0, 511, 512, 513, std.math.minInt(i32), std.math.maxInt(i32) }; + for (self.picks, 0..) |*p, i| { + p.* = if (rng.uintLessThan(u8, 10) < 6) rng.intRangeLessThan(i32, 0, 48) else rng.intRangeLessThan(i32, 0, experts); + if (i % 11 == 10) p.* = 512; + if (fill == .nonfinite and i % 3 == 0) p.* = edge[(i / 3) % edge.len]; + } + self.st0 = try a.alloc(u32, stateWords(lines)); + @memset(self.st0, 0); + const line_id: []i32 = @ptrCast(self.st0[0..lines]); + const line_of: []i32 = @ptrCast(self.st0[2 * lines ..][0..experts]); + @memset(line_of, -1); + // held lines: a stress fill leaves none empty (every miss must evict, or be read in place) + var order: [experts]u16 = undefined; + for (&order, 0..) |*o, i| o.* = @intCast(i); + rng.shuffle(u16, &order); + for (line_id, 0..) |*id, i| { + if (i >= experts or (fill != .stress and rng.uintLessThan(u8, 10) < 3)) { + id.* = -1; + continue; + } + id.* = order[i]; + line_of[order[i]] = @intCast(i); + } + for (self.st0[lines .. 2 * lines]) |*s| s.* = rng.uintLessThan(u32, 8); + for (self.st0[2 * lines + experts ..][0..experts]) |*cnt| cnt.* = rng.uintLessThan(u32, 4); + for (self.st0[2 * lines + 2 * experts ..][0..experts]) |*cnt| cnt.* = 0; + self.st0[2 * lines + 3 * experts + recent_rows_words] = rng.uintLessThan(u32, 100); // now + for (0..3) |w| for (0..experts) |i| { + self.tab0[w][i] = 0x1111 + @as(i64, @intCast(rng.uintLessThan(u16, 1000))); + }; + for (0..3) |w| for (0..experts) |i| { + self.k2t0[w][i] = 7; + }; + for (&self.src, &self.lens, 0..) |*s, *l, i| { + s.* = 0x6000_0000_0000 + @as(i64, @intCast(i)) * 4096; + l.* = 16 * (1 + @as(i64, @intCast(i % 7))); + } + for (&self.k2, 0..) |*x, i| x.* = @intCast(2 + i % 5); + const r = &self.rig; + self.d[0] = try r.input(std.mem.sliceAsBytes(self.picks)); + self.d[1] = try r.input(std.mem.sliceAsBytes(&self.src)); + self.d[2] = try r.input(std.mem.sliceAsBytes(&self.lens)); + self.d[3] = try r.input(std.mem.sliceAsBytes(&self.k2)); + self.d[4] = try r.output(self.st0.len * 4); + for (0..3) |w| { + self.d[5 + w] = try r.output(experts * 8); + self.d[8 + w] = try r.output(experts * 4); + } + self.d[11] = try r.output(self.picks.len * 9 * 8); + self.d[12] = try r.output(4); + self.d[13] = try r.output(self.picks.len * 4); + self.d[14] = try r.output(self.picks.len * 4); + self.d[15] = try r.output(self.picks.len * 4); + self.d[16] = try r.output(4); + return true; + } + + /// The in-out buffers go back to their prepared contents, so every run plans the same round. + fn reset(self: *Plan) !void { + const r = &self.rig; + try put(r, self.d[4], std.mem.sliceAsBytes(self.st0)); + for (0..3) |w| { + try put(r, self.d[5 + w], std.mem.asBytes(&self.tab0[w])); + try put(r, self.d[8 + w], std.mem.asBytes(&self.k2t0[w])); + } + const zeros = try r.gpu.gpa.alloc(u8, self.picks.len * 9 * 8); + defer r.gpu.gpa.free(zeros); + @memset(zeros, 0); + try put(r, self.d[11], zeros); + try put(r, self.d[12], zeros[0..4]); + try put(r, self.d[13], zeros[0 .. self.picks.len * 4]); + try put(r, self.d[14], zeros[0 .. self.picks.len * 4]); + try put(r, self.d[15], zeros[0 .. self.picks.len * 4]); + try put(r, self.d[16], zeros[0..4]); + } + + pub fn run(self: *Plan) !void { + try self.reset(); + const d = self.d; + var a: cuda.Args = .{}; + a.add(d[0]); + a.add(@as(c_int, @intCast(self.picks.len))); + a.add(d[4]); + a.add(@as(c_int, @intCast(self.lines))); + a.add(@as(c_int, @intFromBool(self.recent))); + for (5..11) |i| a.add(d[i]); + a.add(d[1]); + a.add(d[2]); + a.add(d[3]); + a.add(line_base); + a.add(line_bytes); + for (offs) |o| a.add(o); + a.add(d[11]); + a.add(d[12]); + a.add(d[13]); + a.add(d[14]); + a.add(d[15]); + a.add(d[16]); + a.add(@as(u64, 0)); // no trace + try cuda.launch.launch(self.f, .{ .grid = .{ .x = 1 }, .block = .{ .x = 256 } }, self.rig.s, &a); + try self.rig.s.synchronize(); + } + + pub fn count(self: *Plan) usize { + return self.picks.len; + } + + pub fn reference(self: *Plan, want: []f64, mag: []f64) !void { + for (self.picks, want, mag) |p, *w, *m| { + w.* = @floatFromInt(p); + m.* = @abs(w.*); + } + } + + pub fn check(self: *Plan, _: g.Gate) !?[]const u8 { + if (!try self.rig.intact()) return "a payload guard changed"; + const a = self.rig.gpu.gpa; + var m: Model = .{ .st = try a.dupe(u32, self.st0), .tab = self.tab0, .k2t = self.k2t0, .work = try a.alloc(i64, self.picks.len * 9), .hit = try a.alloc(i32, self.picks.len), .miss = try a.alloc(i32, self.picks.len), .ovf = try a.alloc(i32, self.picks.len) }; + defer a.free(m.st); + defer a.free(m.work); + defer a.free(m.hit); + defer a.free(m.miss); + defer a.free(m.ovf); + @memset(m.work, 0); + @memset(m.hit, 0); + @memset(m.miss, 0); + @memset(m.ovf, 0); + m.round(self.lines, self.recent, self.picks, &self.src, &self.lens, &self.k2, line_base, line_bytes, offs); + const r = &self.rig; + const msg = &self.msg; + if (try same(r, u32, self.d[4], m.st, "state", msg)) |why| return why; + for (0..3) |w| { + if (try same(r, i64, self.d[5 + w], &m.tab[w], "table", msg)) |why| return why; + if (try same(r, i32, self.d[8 + w], &m.k2t[w], "k2 table", msg)) |why| return why; + } + if (try same(r, i64, self.d[11], m.work, "copy work", msg)) |why| return why; + if (try same(r, i32, self.d[12], &.{m.work_n}, "copies", msg)) |why| return why; + if (try same(r, i32, self.d[13], m.hit, "held picks", msg)) |why| return why; + if (try same(r, i32, self.d[14], m.miss, "missed picks", msg)) |why| return why; + if (try same(r, i32, self.d[15], m.ovf, "in-place experts", msg)) |why| return why; + return try same(r, i32, self.d[16], &.{m.ovf_n}, "in-place count", msg); + } +}; + +/// Every kernel of the image, plus both ranks' pick masks, for the runner's list. +pub const Set = struct { + masks: tp_.Masks, + mod: cuda.Module, + plan: Plan, + restore: small.Restore, + copy: small.Copy, + drafts: [2]draft.Draft, + checkpoint: draft.Checkpoint, + + pub fn init(gpu: check.Gpu, s: cuda.Stream) !Set { + var masks = try tp_.Masks.init(gpu, s); + errdefer masks.deinit(); + var mod = try cuda.Module.load(gpu.d, cuda.kernels.flashnext_exl3_dplan); + errdefer mod.unload(); + const rig: rig_.Rig = .{ .gpu = gpu, .s = s }; + const token_fn = try mod.function("fn_exl3_draft_token"); + const next_fn = try mod.function("fn_exl3_draft_next"); + return .{ + .masks = masks, + .mod = mod, + .plan = .{ .rig = rig, .f = try mod.function("fn_exl3_dplan") }, + .restore = .{ .rig = rig, .f = try mod.function("fn_exl3_drestore") }, + .copy = .{ .rig = rig, .f = try mod.function("fn_exl3_dcopy") }, + .drafts = .{ .{ .rig = rig, .token_fn = token_fn, .next_fn = next_fn, .decide = false }, .{ .rig = rig, .token_fn = token_fn, .next_fn = next_fn, .decide = true } }, + .checkpoint = .{ .rig = rig, .f = try mod.function("fn_exl3_dcheckpoint") }, + }; + } + + pub fn deinit(self: *Set) void { + self.plan.deinit(); + self.restore.deinit(); + self.copy.deinit(); + for (&self.drafts) |*d| d.deinit(); + self.checkpoint.deinit(); + self.mod.unload(); + self.masks.deinit(); + } + + pub fn add(self: *Set, a: std.mem.Allocator, list: *std.ArrayList(g.Kernel)) !void { + try self.masks.add(a, list); + try list.append(a, self.plan.kernel()); + try list.append(a, self.restore.kernel()); + try list.append(a, self.copy.kernel()); + for (&self.drafts) |*d| try list.append(a, d.kernel()); + try list.append(a, self.checkpoint.kernel()); + } +}; diff --git a/zig/tests/cuda/contract_flashnext_exl3_dplan_copy.zig b/zig/tests/cuda/contract_flashnext_exl3_dplan_copy.zig new file mode 100644 index 0000000000..f9f97ae758 --- /dev/null +++ b/zig/tests/cuda/contract_flashnext_exl3_dplan_copy.zig @@ -0,0 +1,213 @@ +//! Flash Next EXL3 fn_exl3_drestore and fn_exl3_dcopy under gates 1-4. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const rig_ = @import("contract_rig.zig"); +const u = @import("contract_flashnext_exl3_dplan_util.zig"); +const c = core.contract.spec; +const g = core.contract.gates; +const noSums = u.noSums; +const put = u.put; +const fetch = u.fetch; +const same = u.same; +const experts = u.experts; + +pub const Restore = struct { + rig: rig_.Rig, + f: cuda.Function, + ids: []i32 = &.{}, + tab0: [3][experts]i64 = undefined, + k2t0: [3][experts]i32 = undefined, + d: [8]u64 = @splat(0), + msg: [160]u8 = undefined, + + pub fn deinit(self: *Restore) void { + self.rig.gpu.gpa.free(self.ids); + self.rig.deinit(); + } + + const vt = g.hooks(Restore); + + pub fn kernel(self: *Restore) g.Kernel { + return .{ + .spec = .{ + .name = "cuda Flash Next EXL3 fn_exl3_drestore", + .accum = .none, + .depth = noSums, + .out = .raw, + .nan = .{ .stated = "integer ids only: each listed expert's six table entries become 0 and every other entry keeps its value" }, + // the block strides the list by its 256 threads + .tiles = .{ .{ .size = 256, .step = 1, .min = 1 }, .{}, .{} }, + .served = &.{ .{ .rows = 1, .n = 0, .k = 0 }, .{ .rows = 11, .n = 0, .k = 0 }, .{ .rows = 160, .n = 0, .k = 0 }, .{ .rows = 512, .n = 0, .k = 0 } }, + }, + .ctx = self, + .vt = &vt, + }; + } + + pub fn prepare(self: *Restore, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows == 0 or shape.rows > experts) return false; + self.rig.reset(); + const a = self.rig.gpu.gpa; + a.free(self.ids); + var prng = std.Random.DefaultPrng.init(seed ^ 0x72657374); + const rng = prng.random(); + var order: [experts]i32 = undefined; + for (&order, 0..) |*o, i| o.* = @intCast(i); + if (fill != .random) std.mem.reverse(i32, &order) else rng.shuffle(i32, &order); + self.ids = try a.dupe(i32, order[0..shape.rows]); + for (0..3) |w| for (0..experts) |i| { + self.tab0[w][i] = 0x2222 + @as(i64, @intCast(rng.uintLessThan(u16, 1000))); + }; + for (0..3) |w| for (0..experts) |i| { + self.k2t0[w][i] = 9; + }; + const n: i32 = @intCast(shape.rows); + self.d[0] = try self.rig.input(std.mem.sliceAsBytes(self.ids)); + self.d[1] = try self.rig.input(std.mem.asBytes(&n)); + for (0..3) |w| { + self.d[2 + w] = try self.rig.output(experts * 8); + self.d[5 + w] = try self.rig.output(experts * 4); + } + return true; + } + + pub fn run(self: *Restore) !void { + for (0..3) |w| { + try put(&self.rig, self.d[2 + w], std.mem.asBytes(&self.tab0[w])); + try put(&self.rig, self.d[5 + w], std.mem.asBytes(&self.k2t0[w])); + } + var a: cuda.Args = .{}; + for (self.d) |p| a.add(p); + try cuda.launch.launch(self.f, .{ .grid = .{ .x = 1 }, .block = .{ .x = 256 } }, self.rig.s, &a); + try self.rig.s.synchronize(); + } + + pub fn count(self: *Restore) usize { + return self.ids.len; + } + + pub fn reference(self: *Restore, want: []f64, mag: []f64) !void { + for (self.ids, want, mag) |id, *w, *m| { + w.* = @floatFromInt(id); + m.* = @abs(w.*); + } + } + + pub fn check(self: *Restore, _: g.Gate) !?[]const u8 { + if (!try self.rig.intact()) return "a payload guard changed"; + var tab = self.tab0; + var k2t = self.k2t0; + for (self.ids) |id| for (0..3) |w| { + tab[w][@intCast(id)] = 0; + k2t[w][@intCast(id)] = 0; + }; + for (0..3) |w| { + if (try same(&self.rig, i64, self.d[2 + w], &tab[w], "table", &self.msg)) |why| return why; + if (try same(&self.rig, i32, self.d[5 + w], &k2t[w], "k2 table", &self.msg)) |why| return why; + } + return null; + } +}; + +pub const Copy = struct { + rig: rig_.Rig, + f: cuda.Function, + src: []u8 = &.{}, + blocks: u32 = 1, + dst: u64 = 0, + work_ptrs: [2]u64 = .{ 0, 0 }, + msg: [160]u8 = undefined, + + pub fn deinit(self: *Copy) void { + self.rig.gpu.gpa.free(self.src); + self.rig.deinit(); + } + + const vt = g.hooks(Copy); + + pub fn kernel(self: *Copy) g.Kernel { + return .{ + .spec = .{ + .name = "cuda Flash Next EXL3 fn_exl3_dcopy", + .accum = .none, + .depth = noSums, + .out = .raw, + .nan = .{ .stated = "bytes only: a segment's 16-byte words arrive unchanged, NaN and infinity patterns included" }, + // the copy unrolls by four once a segment passes three strides of the grid (256 threads a block) + .tiles = .{ .{}, .{ .size = 768, .step = 1, .min = 1 }, .{} }, + .served = &.{ .{ .rows = 1, .n = 8, .k = 0 }, .{ .rows = 3, .n = 800, .k = 0 }, .{ .rows = 9, .n = 4096, .k = 0 }, .{ .rows = 30, .n = 9000, .k = 0 } }, + }, + .ctx = self, + .vt = &vt, + }; + } + + /// The kernel copies `work_n` lines of three segments (gate, up, down): the shape's rows are lines. + const per_line = 3; + + /// Words (16 bytes) of segment `i`: the shape's longest, then shorter ones. + fn words(shape: c.Shape, i: usize) usize { + return @max(1, shape.n - (i * 37) % @max(shape.n / 2, 1)); + } + + pub fn prepare(self: *Copy, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows == 0 or shape.n == 0 or shape.rows > 64 or shape.n > 20000) return false; + self.rig.reset(); + const a = self.rig.gpu.gpa; + a.free(self.src); + var total: usize = 0; + const segs = shape.rows * per_line; + for (0..segs) |i| total += words(shape, i) * 16; + self.src = try a.alloc(u8, total); + var prng = std.Random.DefaultPrng.init(seed ^ 0x636f7079); + prng.random().bytes(self.src); + if (fill != .random) for (0..self.src.len / 4) |i| if (i % 5 == 0) { + const pattern = [_]u32{ 0x7fc00000, 0xffc00000, 0x7f800000, 0xff800000, 0x7fc01234 }; + std.mem.writeInt(u32, self.src[i * 4 ..][0..4], pattern[(i / 5) % pattern.len], .little); + }; + self.blocks = if (seed & 2 == 0) 1 else 3; + const dsrc = try self.rig.input(self.src); + self.dst = try self.rig.output(total); + const work = try a.alloc(i64, segs * 3); + defer a.free(work); + var at: usize = 0; + for (0..segs) |i| { + const len = words(shape, i) * 16; + work[i * 3 + 0] = @intCast(dsrc + at); + work[i * 3 + 1] = @intCast(self.dst + at); + work[i * 3 + 2] = @intCast(len); + at += len; + } + const n: i32 = @intCast(shape.rows); + const dwork = try self.rig.input(std.mem.sliceAsBytes(work)); + const dn = try self.rig.input(std.mem.asBytes(&n)); + self.work_ptrs = .{ dwork, dn }; + return true; + } + + pub fn run(self: *Copy) !void { + var a: cuda.Args = .{}; + a.add(self.work_ptrs[0]); + a.add(self.work_ptrs[1]); + try cuda.launch.launch(self.f, .{ .grid = .{ .x = self.blocks }, .block = .{ .x = 256 } }, self.rig.s, &a); + try self.rig.s.synchronize(); + } + + pub fn count(self: *Copy) usize { + return self.src.len; + } + + pub fn reference(self: *Copy, want: []f64, mag: []f64) !void { + for (self.src, want, mag) |b, *w, *m| { + w.* = @floatFromInt(b); + m.* = w.*; + } + } + + pub fn check(self: *Copy, _: g.Gate) !?[]const u8 { + if (!try self.rig.intact()) return "a payload guard changed"; + if (!std.mem.eql(u8, self.src, (try self.rig.inputBytes())[0..self.src.len])) return "the source changed"; + return try same(&self.rig, u8, self.dst, self.src, "copied byte", &self.msg); + } +}; diff --git a/zig/tests/cuda/contract_flashnext_exl3_dplan_draft.zig b/zig/tests/cuda/contract_flashnext_exl3_dplan_draft.zig new file mode 100644 index 0000000000..e223b66f41 --- /dev/null +++ b/zig/tests/cuda/contract_flashnext_exl3_dplan_draft.zig @@ -0,0 +1,293 @@ +//! Flash Next EXL3 draft chain decisions and the trace checkpoint under gates 1-4. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const rig_ = @import("contract_rig.zig"); +const u = @import("contract_flashnext_exl3_dplan_util.zig"); +const c = core.contract.spec; +const g = core.contract.gates; +const noSums = u.noSums; +const put = u.put; +const fetch = u.fetch; +const same = u.same; +const experts = u.experts; +const flash = @import("flashnext"); +const stateWords = u.stateWords; +const recent_rows_words = u.recent_rows_words; +const trace_words = u.trace_words; + +pub const Draft = struct { + rig: rig_.Rig, + token_fn: cuda.Function, + next_fn: cuda.Function, + /// The decision kernel, which sets a graph conditional, runs inside a graph with an IF node. + decide: bool, + ids: []u32 = &.{}, + best: [4]f32 = undefined, + index: u32 = 0, + alive_in: i32 = 1, + confidence: f64 = 0, + d: [3]u64 = @splat(0), + msg: [160]u8 = undefined, + + pub fn deinit(self: *Draft) void { + self.rig.gpu.gpa.free(self.ids); + self.rig.deinit(); + } + + const vt = g.hooks(Draft); + + pub fn kernel(self: *Draft) g.Kernel { + return .{ .spec = .{ + .name = if (self.decide) "cuda Flash Next EXL3 fn_exl3_draft_next" else "cuda Flash Next EXL3 fn_exl3_draft_token", + .accum = .none, + .depth = noSums, + .out = .raw, + .nan = .{ .stated = "the best record's column index picks the token whatever its score fields hold; the chain ends only when a probability exp(best - lse) is below a positive confidence, so a NaN score never ends it and a probability of 0 does" }, + .served = &.{ .{ .rows = 1, .n = 1, .k = 0 }, .{ .rows = 1, .n = 4096, .k = 0 }, .{ .rows = 1, .n = 248320, .k = 0 } }, + }, .ctx = self, .vt = &vt }; + } + + fn expectedAlive(self: *Draft) i32 { + const prob = @exp(@as(f64, self.best[0]) - @as(f64, self.best[1])); + const low = self.confidence > 0.0 and prob < self.confidence; + return @intFromBool(self.alive_in != 0 and !low); + } + + pub fn prepare(self: *Draft, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.n == 0 or shape.n > 1 << 20) return false; + self.rig.reset(); + const a = self.rig.gpu.gpa; + a.free(self.ids); + var prng = std.Random.DefaultPrng.init(seed ^ 0x64726166); + const rng = prng.random(); + self.ids = try a.alloc(u32, shape.n); + rng.bytes(std.mem.sliceAsBytes(self.ids)); + self.index = if (seed % 3 == 0) 0 else if (seed % 3 == 1) @intCast(shape.n - 1) else rng.uintLessThan(u32, @intCast(shape.n)); + const top = -rng.float(f32) * 4; + self.best = .{ top, top + rng.float(f32) * 3, @bitCast(self.index), 0 }; + if (fill != .random) { + const specials = [_]f32{ std.math.nan(f32), std.math.inf(f32), -std.math.inf(f32) }; + self.best[0] = specials[seed % 3]; + self.best[1] = specials[(seed / 3) % 3]; + } + self.alive_in = if (seed & 4 == 0) 1 else 0; + const prob = @exp(@as(f64, self.best[0]) - @as(f64, self.best[1])); + self.confidence = switch (seed % 4) { + 0 => 0, + 1 => if (std.math.isFinite(prob) and prob > 0) prob * 0.9 else 0.5, + 2 => if (std.math.isFinite(prob) and prob > 0) prob * 1.1 else 0.5, + else => 0.5, + }; + self.d[0] = try self.rig.input(std.mem.asBytes(&self.best)); + self.d[1] = try self.rig.input(std.mem.sliceAsBytes(self.ids)); + self.d[2] = try self.rig.output(16); + return true; + } + + pub fn run(self: *Draft) !void { + const init = [4]i32{ -1, self.alive_in, -1, 0 }; + try put(&self.rig, self.d[2], std.mem.asBytes(&init)); + const s = self.rig.s; + if (!self.decide) { + var a: cuda.Args = .{}; + a.add(self.d[0]); + a.add(self.d[1]); + a.add(self.d[2]); + try cuda.launch.launch(self.token_fn, .{ .grid = .{ .x = 1 }, .block = .{ .x = 32 } }, s, &a); + return s.synchronize(); + } + const chain = flash.exl3.chain; + var cond = try chain.Cond.open(); + defer cond.lib.close(); + var side = try cuda.Stream.init(self.rig.gpu.d, true); + defer side.deinit(); + try cuda.graph.beginCapture(s, .thread_local); + const handle = try cond.handle(s); + var a: cuda.Args = .{}; + a.add(self.d[0]); + a.add(self.d[1]); + a.add(self.d[2]); + a.add(self.d[2] + 4); + a.add(self.d[2] + 8); + a.add(handle); + a.add(self.confidence); + try cuda.launch.launch(self.next_fn, .{ .grid = .{ .x = 1 }, .block = .{ .x = 32 } }, s, &a); + const body = try cond.ifNode(s, handle); + const d = self.rig.gpu.d; + try chain.Cond.ok(cond.begin_to_graph(@ptrCast(side.handle), body, null, null, 0, 1), "cuStreamBeginCaptureToGraph"); + try d.check(d.api.cuMemsetD32Async(self.d[2] + 12, 1, 1, side.handle), "cuMemsetD32Async"); + var inner: cuda.abi.Graph = null; + try d.check(d.api.cuStreamEndCapture(side.handle, &inner), "cuStreamEndCapture"); + var graph = try cuda.graph.endCapture(s); + defer graph.deinit(); + var exec = try graph.instantiate(); + defer exec.deinit(); + try exec.launchOn(s); + try s.synchronize(); + } + + pub fn count(_: *Draft) usize { + return 4; + } + + pub fn reference(self: *Draft, want: []f64, mag: []f64) !void { + for (want, mag, 0..) |*w, *m, i| { + w.* = if (i == 0) @floatFromInt(self.ids[self.index]) else @floatFromInt(self.expectedAlive()); + m.* = @abs(w.*); + } + } + + pub fn check(self: *Draft, _: g.Gate) !?[]const u8 { + if (!try self.rig.intact()) return "a payload guard changed"; + const out = try fetch(&self.rig, i32, self.d[2], 4); + defer self.rig.gpu.gpa.free(out); + const token: i32 = @bitCast(self.ids[self.index]); + if (out[0] != token) return try std.fmt.bufPrint(&self.msg, "token {d}, not ids[{d}] = {d}", .{ out[0], self.index, token }); + if (!self.decide) return null; + const alive = self.expectedAlive(); + if (out[1] != alive or out[2] != alive) return try std.fmt.bufPrint(&self.msg, "alive {d} and {d}, not {d} (best {d} {d}, confidence {d})", .{ out[1], out[2], alive, self.best[0], self.best[1], self.confidence }); + if (out[3] != alive) return try std.fmt.bufPrint(&self.msg, "the conditional body ran {d}, not {d}", .{ out[3], alive }); + return null; + } +}; + +// ---- fn_exl3_dcheckpoint ---------------------------------------------------------------------------------- + +/// The kernel's checkpoint on a host copy of the trace buffer (dp_trace_valid, then the transfer 0 / 1 rules). +fn modelCheckpoint(tr: []u32, st: []const u32, lines: u32, transfer: i32) void { + const seed_words: u64 = 2 * @as(u64, lines) + 3 * experts + recent_rows_words + 10; + const record_words: u64 = 15 + 11 * @as(u64, tr[5]); + const base = 16 + 2 * seed_words; + const stride: u64 = 4 + 2 * seed_words; + const capacity = if (record_words == 0) 0 else (trace_words - base) / record_words; + const valid = lines > 0 and tr[0] == 0x44505431 and tr[1] == 4 and tr[4] == lines and tr[5] != 0 and base < trace_words and + record_words <= trace_words - base and tr[6] == seed_words and tr[7] == record_words and tr[8] == capacity and tr[9] <= capacity and + tr[12] <= 1 and tr[13] <= 1 and tr[14] <= (trace_words - base) / stride and + base + @as(u64, tr[9]) * record_words <= trace_words - @as(u64, tr[14]) * stride; + if (!valid or (transfer != 0 and transfer != 1)) { + tr[13] = 1; + return; + } + const sw: usize = @intCast(seed_words); + const initialized = tr[9] != 0 or tr[14] != 0; + if (transfer == 0) { + if (initialized) for (0..sw) |i| if (tr[16 + sw + i] != st[i]) { + tr[12] = 1; + }; + return; + } + const end: u32 = @intCast(16 + 2 * seed_words + tr[9] * record_words); + const start: u32 = @intCast(trace_words - tr[14] * stride); + if (start -% end < stride or tr[10] != 0 or tr[11] != 0 or tr[15] != 0) { + if (tr[15] != 0xffffffff) tr[15] += 1; + return; + } + const at: usize = start - @as(usize, @intCast(stride)); + tr[at] = tr[9]; + tr[at + 1] = 1; + tr[at + 2] = @intCast(sw); + tr[at + 3] = tr[14]; + for (0..sw) |i| { + if (!initialized) { + tr[16 + i] = st[i]; + tr[16 + sw + i] = st[i]; + } + tr[at + 4 + i] = tr[16 + sw + i]; + tr[at + 4 + sw + i] = st[i]; + tr[16 + sw + i] = st[i]; + } + tr[14] += 1; +} + +pub const Checkpoint = struct { + rig: rig_.Rig, + f: cuda.Function, + lines: u32 = 0, + st: [2][]u32 = .{ &.{}, &.{} }, + header: [16]u32 = undefined, + d: [3]u64 = @splat(0), + msg: [160]u8 = undefined, + + pub fn deinit(self: *Checkpoint) void { + for (self.st) |s| self.rig.gpu.gpa.free(s); + self.rig.deinit(); + } + + const vt = g.hooks(Checkpoint); + + pub fn kernel(self: *Checkpoint) g.Kernel { + return .{ .spec = .{ + .name = "cuda Flash Next EXL3 fn_exl3_dcheckpoint", + .accum = .none, + .depth = noSums, + .out = .raw, + .nan = .{ .stated = "integer state words only: a checkpoint copies the planner's words, flags a discontinuity when they differ from the last record's end, and counts an invalid launch" }, + .served = &.{ .{ .rows = 1, .n = 6, .k = 0 }, .{ .rows = 1, .n = 38, .k = 0 }, .{ .rows = 1, .n = 160, .k = 0 } }, + }, .ctx = self, .vt = &vt }; + } + + pub fn prepare(self: *Checkpoint, shape: c.Shape, _: g.Fill, seed: u64) !bool { + if (shape.n == 0 or shape.n > 512) return false; + self.rig.reset(); + const a = self.rig.gpu.gpa; + self.lines = @intCast(shape.n); + var prng = std.Random.DefaultPrng.init(seed ^ 0x636b7074); + const words = stateWords(shape.n); + for (&self.st) |*s| { + a.free(s.*); + s.* = try a.alloc(u32, words); + prng.random().bytes(std.mem.sliceAsBytes(s.*)); + } + const max_rows: u32 = 16; + const record_words = 15 + 11 * max_rows; + const base = 16 + 2 * @as(u32, @intCast(words)); + self.header = .{ 0x44505431, 4, 0, 3, self.lines, max_rows, @intCast(words), record_words, (trace_words - base) / record_words, 0, 0, 0, 0, 0, 0, 0 }; + self.d[0] = try self.rig.input(std.mem.sliceAsBytes(self.st[0])); + self.d[1] = try self.rig.input(std.mem.sliceAsBytes(self.st[1])); + self.d[2] = try self.rig.output(trace_words * 4); + return true; + } + + /// A first checkpoint, a repeat, a different state, its checkpoint and a bad transfer. + const steps = [_][2]i32{ .{ 0, 1 }, .{ 0, 0 }, .{ 1, 0 }, .{ 1, 1 }, .{ 1, 2 } }; + + pub fn run(self: *Checkpoint) !void { + const zeros = try self.rig.gpu.gpa.alloc(u8, trace_words * 4); + defer self.rig.gpu.gpa.free(zeros); + @memset(zeros, 0); + @memcpy(zeros[0..64], std.mem.sliceAsBytes(&self.header)); + try put(&self.rig, self.d[2], zeros); + for (steps) |step| { + var a: cuda.Args = .{}; + a.add(self.d[@intCast(step[0])]); + a.add(@as(c_int, @intCast(self.lines))); + a.add(self.d[2]); + a.add(@as(c_int, step[1])); + try cuda.launch.launch(self.f, .{ .grid = .{ .x = 1 }, .block = .{ .x = 1 } }, self.rig.s, &a); + } + try self.rig.s.synchronize(); + } + + pub fn count(self: *Checkpoint) usize { + return self.st[0].len; + } + + pub fn reference(self: *Checkpoint, want: []f64, mag: []f64) !void { + for (self.st[0], want, mag) |x, *w, *m| { + w.* = @floatFromInt(x); + m.* = w.*; + } + } + + pub fn check(self: *Checkpoint, _: g.Gate) !?[]const u8 { + if (!try self.rig.intact()) return "a payload guard changed"; + const a = self.rig.gpu.gpa; + const tr = try a.alloc(u32, trace_words); + defer a.free(tr); + @memset(tr, 0); + @memcpy(tr[0..16], &self.header); + for (steps) |step| modelCheckpoint(tr, self.st[@intCast(step[0])], self.lines, step[1]); + return try same(&self.rig, u32, self.d[2], tr, "trace word", &self.msg); + } +}; diff --git a/zig/tests/cuda/contract_flashnext_exl3_dplan_util.zig b/zig/tests/cuda/contract_flashnext_exl3_dplan_util.zig new file mode 100644 index 0000000000..a984449165 --- /dev/null +++ b/zig/tests/cuda/contract_flashnext_exl3_dplan_util.zig @@ -0,0 +1,40 @@ +//! Shared pieces of the Flash Next EXL3 device-plan contract runners. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const rig_ = @import("contract_rig.zig"); +const c = core.contract.spec; + +pub const experts = 512; +pub const recent_rows_words = 256 * (experts / 32); +pub const trace_words = 1024 * 1024 / 4; + +pub fn stateWords(lines: usize) usize { + return 2 * lines + 3 * experts + recent_rows_words + 10; +} + +pub fn noSums(_: c.Shape) usize { + return 0; +} + +pub fn put(rig: *rig_.Rig, ptr: u64, bytes: []const u8) !void { + const b: cuda.DeviceBuffer = .{ .d = rig.gpu.d, .ptr = ptr, .len = bytes.len }; + try b.upload(0, bytes); +} + +pub fn fetch(rig: *rig_.Rig, comptime T: type, ptr: u64, n: usize) ![]T { + const out = try rig.gpu.gpa.alloc(T, n); + errdefer rig.gpu.gpa.free(out); + try rig.s.synchronize(); + const b: cuda.DeviceBuffer = .{ .d = rig.gpu.d, .ptr = ptr, .len = n * @sizeOf(T) }; + try b.download(0, std.mem.sliceAsBytes(out)); + return out; +} + +/// `got` (a device buffer's `n` elements) equals `want`, or what first differs. +pub fn same(rig: *rig_.Rig, comptime T: type, ptr: u64, want: []const T, what: []const u8, msg: []u8) !?[]const u8 { + const got = try fetch(rig, T, ptr, want.len); + defer rig.gpu.gpa.free(got); + for (got, want, 0..) |x, y, i| if (x != y) return try std.fmt.bufPrint(msg, "{s}[{d}] is {d}, not {d}", .{ what, i, x, y }); + return null; +} diff --git a/zig/tests/cuda/contract_flashnext_exl3_suite.zig b/zig/tests/cuda/contract_flashnext_exl3_suite.zig index 376dfcd33e..5c06ad6f87 100644 --- a/zig/tests/cuda/contract_flashnext_exl3_suite.zig +++ b/zig/tests/cuda/contract_flashnext_exl3_suite.zig @@ -3,28 +3,28 @@ 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"); +const dplan_ = @import("contract_flashnext_exl3_dplan.zig"); const prefill_ = @import("contract_flashnext_exl3_prefill.zig"); pub const Suite = struct { a: std.mem.Allocator, - masks: tp_.Masks, + set: dplan_.Set, 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); 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); + self.set = try dplan_.Set.init(gpu, s); + errdefer self.set.deinit(); + try self.set.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.set.deinit(); self.a.destroy(self); } }; diff --git a/zig/tests/cuda/contract_limits.zig b/zig/tests/cuda/contract_limits.zig index 4b556a4463..46326a81b7 100644 --- a/zig/tests/cuda/contract_limits.zig +++ b/zig/tests/cuda/contract_limits.zig @@ -27,7 +27,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", "nemotron_ops", "minimax_m3", "minimax_m3_attn", "moss_image" }) |n| if (std.mem.eql(u8, name, n)) return @field(kk, n); - inline for (.{ "flashnext_exl3", "flashnext_exl3_tp", "flashnext_exl3_kvtier", "flashnext_exl3_prefill" }) |n| + inline for (.{ "flashnext_exl3", "flashnext_exl3_tp", "flashnext_exl3_dplan", "flashnext_exl3_kvtier", "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/graph/flashnext_exl3_graph_engine.zig b/zig/tests/cuda/flashnext_exl3/graph/flashnext_exl3_graph_engine.zig new file mode 100644 index 0000000000..24311b1361 --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/graph/flashnext_exl3_graph_engine.zig @@ -0,0 +1,159 @@ +//! The EXL3 graph tests' engine harness: one rank's engine with its sequence, TP link and output buffer. + +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 dplan = X.dplan; +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 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 engine at tp=1/2 ------------------------------------------------------------------------------------------- + +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); + } + + 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; + } +}; + +pub 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, + + 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; + errdefer if (g.ranks) |ranks| ranks.close(gpu.gpa); + g.stream = try cuda.Stream.init(gpu.d, true); + errdefer g.stream.deinit(); + g.glue = try flash.Glue.load(gpu.d); + errdefer g.glue.deinit(); + const t0 = check.now(gpu.io); + g.e = try prompt.Engine.initRank(gpu.gpa, gpu.io, @constCast(gpu.d), &g.stream, dir, world, rank); + errdefer g.e.deinit(); + if (g.ranks) |r| g.e.collective = &r.collective; + _ = 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); + errdefer g.seq.deinit(gpu.gpa); + g.out = try cuda.DeviceBuffer.alloc(gpu.d, out_rows * vocab * 2); + errdefer g.out.free(); + X.prefill.reserveFor(g.e, cap_default, world); + const budget = try xpass.budgetFromEnv(0); + g.attached = try xpass.attach(gpu.gpa, gpu.io, gpu.d, g.stream, &g.e.exl3_prefill, g.e.layers[0..g.e.n_layers], null, g.e.chunk, budget, null); + std.debug.print("RESULT rank {d}/{d} load: {d:.1} s; {d} lines a layer x {d} layers ({d:.2} GiB of pool)\n", .{ rank, world, ms(t0, check.now(gpu.io)) / 1e3, g.attached.lines, g.attached.layers, gib(g.attached.device_bytes) }); + return g; + } + + pub fn deinit(g: *Eng) void { + 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); + } + + 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 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 stats(g: *Eng) !fpool.Stats { + try dplan.toHost(&g.e.exl3_dplan); + return xpass.poolStats(g.e.layers[0..g.e.n_layers], null); + } +}; + +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; +} diff --git a/zig/tests/cuda/flashnext_exl3/graph/flashnext_exl3_graph_eq_tests.zig b/zig/tests/cuda/flashnext_exl3/graph/flashnext_exl3_graph_eq_tests.zig new file mode 100644 index 0000000000..ad062ba226 --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/graph/flashnext_exl3_graph_eq_tests.zig @@ -0,0 +1,237 @@ +//! flashnext-exl3-graph eq and fork: device-plan and graph steps against the host cut, bit for bit. + +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 dplan = X.dplan; +const graph = X.graph; +const vocab = prompt.vocabulary; +const harness = @import("flashnext_exl3_graph_engine.zig"); +const Eng = harness.Eng; +const ms = harness.ms; +const ints = harness.ints; +const cap_default = harness.cap_default; + +const modes = [_]graph.Mode{ .off, .dplan, .graph }; + +const Run = struct { + ids: []u32, + /// Per decode step, the hash of the step's whole bf16 logits row; then per window row the same. + round_h: []u64, + window_h: []u64, + ms_token: f64 = 0, + ms_window_row: f64 = 0, + hit: f64 = 0, + copies: f64 = 0, +}; + +const win_sizes = [_]usize{ 3, 1, 4, 2, 8, 16 }; + +fn runMode(g: *Eng, ids: []const u32, steps: usize, r: *Run) !void { + const a = g.gpu.gpa; + // decode rounds + try dplan.toHost(&g.e.exl3_dplan); + try g.reset(); + try g.forward(ids, .{ .logits = .last }); + var tok = prompt.argmax(g.seq.logits); + r.ids[0] = tok; + _ = try g.swap(1); + const before = try g.stats(); + const t0 = check.now(g.gpu.io); + for (0..steps) |i| { + tok = try g.round(tok); + r.ids[i + 1] = tok; + r.round_h[i] = std.hash.Wyhash.hash(0, g.seq.logits); + } + const t1 = check.now(g.gpu.io); + const after = try g.stats(); + r.ms_token = ms(t0, t1) / @as(f64, @floatFromInt(steps)); + const hits = after.hit_slots - before.hit_slots; + const all = hits + after.miss_slots - before.miss_slots; + r.hit = 100 * @as(f64, @floatFromInt(hits)) / @as(f64, @floatFromInt(@max(all, 1))); + r.copies = @as(f64, @floatFromInt(after.misses - before.misses)) / @as(f64, @floatFromInt(steps)); + // the same continuation teacher-forced in windows of 3, 1, 4, 2 rows (every row's logits) + try dplan.toHost(&g.e.exl3_dplan); + try g.reset(); + try g.forward(ids, .{ .logits = .last }); + var at: usize = 0; + var k: usize = 0; + const host = try a.alloc(u16, 16 * vocab); + defer a.free(host); + const tw = check.now(g.gpu.io); + while (at < steps) : (k += 1) { + const n = @min(win_sizes[k % win_sizes.len], steps - at); + try g.forward(r.ids[at..][0..n], .{ .logits = .all, .logits_out = g.out.ptr }); + try g.out.download(0, std.mem.sliceAsBytes(host[0 .. n * vocab])); + for (0..n) |j| r.window_h[at + j] = std.hash.Wyhash.hash(0, std.mem.sliceAsBytes(host[j * vocab ..][0..vocab])); + at += n; + } + r.ms_window_row = ms(tw, check.now(g.gpu.io)) / @as(f64, @floatFromInt(steps)); +} + +/// Two sequences in one decode round (two rows, each its own caches): +fn pairRun(g: *Eng, b: *prompt.Seq, ids_a: []const u32, ids_b: []const u32, steps: usize, out: []u64) !f64 { + try dplan.toHost(&g.e.exl3_dplan); + try g.reset(); + try g.forward(ids_a, .{ .logits = .last }); + for ([_][]cuda.DeviceBuffer{ b.conv, b.state, b.kc, b.vc, b.ikc, b.pooled }) |bs| for (bs) |x| try x.fill8(0, g.stream.handle); + try b.tail.fill8(0, g.stream.handle); + try g.stream.synchronize(); + b.fed = 0; + b.history = .{ g.e.eos, g.e.eos }; + try g.e.forward(flash.Tri, g.tri(), b, ids_b, .{ .logits = .last }); + if (g.ranks) |r| try r.collective.finish(g.stream) else try g.stream.synchronize(); + var t = [2]u32{ prompt.argmax(g.seq.logits), prompt.argmax(b.logits) }; + const q = [_]*prompt.Seq{ &g.seq, b }; + const t0 = check.now(g.gpu.io); + for (0..steps) |i| { + try g.e.round(flash.Tri, g.tri(), &q, &t, g.out.ptr); + if (g.ranks) |r| try r.collective.finish(g.stream); + out[i] = std.hash.Wyhash.hash(0, g.seq.logits) ^ std.hash.Wyhash.hash(7, b.logits); + t = .{ prompt.argmax(g.seq.logits), prompt.argmax(b.logits) }; + } + return ms(t0, check.now(g.gpu.io)) / @as(f64, @floatFromInt(steps)); +} + +pub fn eqMode(gpu: Gpu, rank: usize, world: usize, port: u16, dir: []const u8, golden_path: []const u8, steps: usize, which: []const u8) !void { + const a = gpu.gpa; + const text = try std.Io.Dir.cwd().readFileAlloc(gpu.io, golden_path, a, .limited(1 << 26)); + defer a.free(text); + var parsed = try std.json.parseFromSlice(std.json.Value, a, text, .{}); + defer parsed.deinit(); + const g = try Eng.init(gpu, rank, world, port, dir); + defer g.deinit(); + var all_ok = true; + for (parsed.value.object.get("prompts").?.array.items) |p| { + const ids = try ints(a, p.object.get("ids").?); + defer a.free(ids); + var runs: [modes.len]Run = undefined; + for (&runs) |*r| r.* = .{ .ids = try a.alloc(u32, steps + 1), .round_h = try a.alloc(u64, steps), .window_h = try a.alloc(u64, steps) }; + defer for (&runs) |*r| { + a.free(r.ids); + a.free(r.round_h); + a.free(r.window_h); + }; + // warm the pools once (first-touch copies are not any mode's cost) + graph.force = .off; + try runMode(g, ids, steps, &runs[0]); + for (modes, 0..) |m, mi| { + if (std.mem.indexOf(u8, which, @tagName(m)) == null and m != .off) continue; + graph.force = m; + const g0 = graph.stats; + try runMode(g, ids, steps, &runs[mi]); + const g1 = graph.stats; + const h = std.hash.Wyhash.hash(0, std.mem.sliceAsBytes(runs[mi].round_h)) ^ std.hash.Wyhash.hash(1, std.mem.sliceAsBytes(runs[mi].window_h)); + const other = try g.swap(h); + if (world == 2) try check.expect(other == h, "rank {d}: the two ranks' logits differ ({s})", .{ rank, @tagName(m) }); + const rs = runs[mi]; + var same_r: usize = 0; + var same_w: usize = 0; + for (rs.round_h, runs[0].round_h) |x, y| same_r += @intFromBool(x == y); + for (rs.window_h, runs[0].window_h) |x, y| same_w += @intFromBool(x == y); + const ok = same_r == steps and same_w == steps and std.mem.eql(u32, rs.ids, runs[0].ids); + if (mi != 0) all_ok = all_ok and ok; + if (rank == 0) std.debug.print("RESULT tp={d} {s} {s}: decode {d:.2} ms/token (hit {d:.1}%, {d:.1} copies a token), windows {d:.2} ms/row; logits bit-equal to off: rounds {d}/{d}, window rows {d}/{d}; greedy ids {s}; graph captures {d} (updates {d}) replays {d}, eager-dplan steps {d}, declined {d}, capture {d:.1} ms\n", .{ world, p.object.get("name").?.string, @tagName(m), rs.ms_token, rs.hit, rs.copies, rs.ms_window_row, same_r, steps, same_w, steps, if (std.mem.eql(u32, rs.ids, runs[0].ids)) "equal" else "DIFFER", g1.captures - g0.captures, g1.updates - g0.updates, g1.replays - g0.replays, g1.eager_dplan - g0.eager_dplan, g1.declined - g0.declined, @as(f64, @floatFromInt(g1.capture_ns - g0.capture_ns)) / 1e6 }); + } + if (std.mem.indexOf(u8, which, "inject") != null) { + // captures fail on rank 1 and instantiations on rank 0: both ranks run those rounds eagerly, bits unchanged + graph.clear(); + graph.fallback.clear(); + graph.force = .graph; + graph.fallback.inject_capture_failures = if (rank == 1 or world == 1) 3 else 0; + graph.fallback.inject_instantiate_failures = if (rank == 0) 2 else 0; + var injected: Run = .{ .ids = try a.alloc(u32, steps + 1), .round_h = try a.alloc(u64, steps), .window_h = try a.alloc(u64, steps) }; + defer a.free(injected.ids); + defer a.free(injected.round_h); + defer a.free(injected.window_h); + const before = graph.stats; + try runMode(g, ids, steps, &injected); + graph.fallback.inject_capture_failures = 0; + graph.fallback.inject_instantiate_failures = 0; + const h = std.hash.Wyhash.hash(0, std.mem.sliceAsBytes(injected.round_h)) ^ std.hash.Wyhash.hash(1, std.mem.sliceAsBytes(injected.window_h)); + if (world == 2) try check.expect(try g.swap(h) == h, "rank {d}: the two ranks' logits differ (injected capture failures)", .{rank}); + var same_r: usize = 0; + var same_w: usize = 0; + for (injected.round_h, runs[0].round_h) |x, y| same_r += @intFromBool(x == y); + for (injected.window_h, runs[0].window_h) |x, y| same_w += @intFromBool(x == y); + const failures = graph.stats.capture_failures - before.capture_failures; + const ok = same_r == steps and same_w == steps and std.mem.eql(u32, injected.ids, runs[0].ids) and failures > 0; + all_ok = all_ok and ok; + if (rank == 0) std.debug.print("RESULT tp={d} injected failures: {d} captures/instantiations failed on this rank, {d} captures and {d} replays after; logits bit-equal to off: rounds {d}/{d}, window rows {d}/{d}\n", .{ world, failures, graph.stats.captures - before.captures, graph.stats.replays - before.replays, same_r, steps, same_w, steps }); + } + graph.force = null; + } + // two streams in one round + { + const ps = parsed.value.object.get("prompts").?.array.items; + const ia = try ints(a, ps[0].object.get("ids").?); + defer a.free(ia); + const ib = try ints(a, ps[ps.len - 1].object.get("ids").?); + defer a.free(ib); + var b = try g.e.newSeq(cap_default); + defer b.deinit(a); + const pair_steps = @min(steps, 32); + var ref: [64]u64 = undefined; + var got: [64]u64 = undefined; + graph.force = .off; + _ = try pairRun(g, &b, ia, ib, pair_steps, &ref); + for (modes) |m| { + if (std.mem.indexOf(u8, which, @tagName(m)) == null and m != .off) continue; + graph.force = m; + const t = try pairRun(g, &b, ia, ib, pair_steps, &got); + var same: usize = 0; + for (got[0..pair_steps], ref[0..pair_steps]) |x, y| same += @intFromBool(x == y); + if (m != .off) all_ok = all_ok and same == pair_steps; + const h = std.hash.Wyhash.hash(0, std.mem.sliceAsBytes(got[0..pair_steps])); + if (world == 2) try check.expect(try g.swap(h) == h, "rank {d}: the two ranks' two-stream logits differ", .{rank}); + if (rank == 0) std.debug.print("RESULT tp={d} two streams {s}: {d:.2} ms a round of 2 rows; both rows' logits bit-equal to off in {d}/{d} rounds\n", .{ world, @tagName(m), t, same, pair_steps }); + } + graph.force = null; + } + if (dplan.current(&g.e.exl3_dplan)) |dp| if (rank == 0) std.debug.print("RESULT dplan owner switches: {d} pushes, {d} pulls; {d} experts read in place (rounds with more misses than free lines)\n", .{ dp.pushes, dp.pulls, dp.inplace }); + try check.expect(all_ok, "rank {d}: a device-plan or graph step's logits differ from the host cut's", .{rank}); + check.pass("rank {d}/{d}: decode rounds and verify windows on the device plan and in graphs are bit-identical to the host cut ({d} steps a prompt)", .{ rank, world, steps }); +} + +/// A scheduling change's bits: every golden prompt's decode rounds and windows (runMode) under `which` (off, dplan). +pub fn forkMode(gpu: Gpu, rank: usize, world: usize, port: u16, dir: []const u8, golden_path: []const u8, steps: usize, which: []const u8) !void { + const a = gpu.gpa; + const text = try std.Io.Dir.cwd().readFileAlloc(gpu.io, golden_path, a, .limited(1 << 26)); + defer a.free(text); + var parsed = try std.json.parseFromSlice(std.json.Value, a, text, .{}); + defer parsed.deinit(); + const g = try Eng.init(gpu, rank, world, port, dir); + defer g.deinit(); + graph.force = if (std.mem.eql(u8, which, "graph")) .graph else if (std.mem.eql(u8, which, "dplan")) .dplan else .off; + defer graph.force = null; + defer X.fork.force = null; + var ok = true; + for (parsed.value.object.get("prompts").?.array.items) |p| { + const ids = try ints(a, p.object.get("ids").?); + defer a.free(ids); + var runs: [3]Run = undefined; + for (&runs) |*r| r.* = .{ .ids = try a.alloc(u32, steps + 1), .round_h = try a.alloc(u64, steps), .window_h = try a.alloc(u64, steps) }; + defer for (&runs) |*r| { + a.free(r.ids); + a.free(r.round_h); + a.free(r.window_h); + }; + for ([_]bool{ false, false, true }, 0..) |on, i| { + X.fork.force = on; + try runMode(g, ids, steps, &runs[i]); + } + const same = std.mem.eql(u64, runs[1].round_h, runs[2].round_h) and std.mem.eql(u64, runs[1].window_h, runs[2].window_h) and std.mem.eql(u32, runs[1].ids, runs[2].ids); + ok = ok and same; + if (world == 2) { + const h = std.hash.Wyhash.hash(0, std.mem.sliceAsBytes(runs[2].round_h)); + try check.expect(try g.swap(h) == h, "rank {d}: the ranks differ", .{rank}); + } + if (rank == 0) std.debug.print("RESULT tp={d} {s} {s}: decode {d:.2} -> {d:.2} ms/token, windows {d:.2} -> {d:.2} ms/row (serial -> parallel parts); logits {s}\n", .{ world, which, p.object.get("name").?.string, runs[1].ms_token, runs[2].ms_token, runs[1].ms_window_row, runs[2].ms_window_row, if (same) "bit-equal" else "DIFFER" }); + } + try check.expect(ok, "rank {d}: parallel parts changed the logits", .{rank}); + check.pass("rank {d}/{d}: parallel projection parts give the same bits ({s})", .{ rank, world, which }); +} diff --git a/zig/tests/cuda/flashnext_exl3/graph/flashnext_exl3_graph_plan_tests.zig b/zig/tests/cuda/flashnext_exl3/graph/flashnext_exl3_graph_plan_tests.zig new file mode 100644 index 0000000000..d75ec95aa1 --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/graph/flashnext_exl3_graph_plan_tests.zig @@ -0,0 +1,207 @@ +//! flashnext-exl3-graph plan: the device planner against the host planner on synthetic pools, and the copy kernel. + +const std = @import("std"); +const cuda = @import("cuda"); +const flash = @import("flashnext"); +const check = @import("../../check.zig"); +const Gpu = check.Gpu; +const X = flash.exl3; +const dplan = X.dplan; +const fpool = flash.exl3_pool; +const harness = @import("flashnext_exl3_graph_engine.zig"); +const ms = harness.ms; + +// -- synthetic pools: device planner == host planner ----------------------------------------------------------------- + +const syn_dims = 256; +const syn_width = 64; + +fn synPool(gpu: Gpu, routed: [][3][]const u8, lines: usize, max_rows: usize) !fpool.Pool { + return fpool.Pool.init(gpu.gpa, gpu.io, gpu.d, .{ .routed = routed }, .{ .suh_g = 0, .suh_u = 0, .svh_g = 0, .svh_u = 0, .suh_d = 0, .svh_d = 0 }, .{ + .layer = 0, + .lines = lines, + .dims = syn_dims, + .width = syn_width, + .max_rows = max_rows, + .slots = 11, + .direct = true, + .double_buffer = false, + }, null); +} + +fn planCase(gpu: Gpu, lines: usize, max_rows: usize, rounds: usize, seed: u64) !void { + const a = gpu.gpa; + const tiles = (syn_dims / 16) * (syn_width / 16); + // a pinned bank of 512 experts, widths 2..5 half-bits, random bytes + var sizes: [512][3]usize = undefined; + var prng = std.Random.DefaultPrng.init(seed); + const rnd = prng.random(); + var total: usize = 0; + for (&sizes) |*e| for (e) |*b| { + b.* = tiles * 16 * (2 + rnd.uintLessThan(usize, 4)); + total += (b.* + 255) / 256 * 256; + }; + var bank = try cuda.HostBuffer.alloc(gpu.d, total); + defer bank.free(); + rnd.bytes(bank.bytes); + const routed = try a.alloc([3][]const u8, 512); + defer a.free(routed); + var at: usize = 0; + for (routed, sizes) |*r, sz| for (0..3) |w| { + r[w] = bank.bytes[at..][0..sz[w]]; + at += (sz[w] + 255) / 256 * 256; + }; + var host = try synPool(gpu, routed, lines, max_rows); + defer host.deinit(); + var dev = try synPool(gpu, routed, lines, max_rows); + defer dev.deinit(); + var s = try cuda.Stream.init(gpu.d, true); + defer s.deinit(); + var pools = [_]*fpool.Pool{&dev}; + const dp = try dplan.DevPlan.init(a, gpu.d, &pools); + defer dp.deinit(); + var picks_dev = try cuda.DeviceBuffer.alloc(gpu.d, max_rows * 11 * 4); + defer picks_dev.free(); + const picks = try a.alloc(i32, max_rows * 11); + defer a.free(picks); + // skewed routing: a hot set of 24 experts takes most picks, so rounds hit, miss and evict + var host_ns: u64 = 0; + for (0..rounds) |round| { + const rows = 1 + rnd.uintLessThan(usize, max_rows); + for (0..rows) |r| { + var k: usize = 0; + while (k < 10) { + const id: i32 = if (rnd.uintLessThan(u32, 4) != 0) @intCast(rnd.uintLessThan(u32, 24) * 7) else @intCast(rnd.uintLessThan(u32, 512)); + var dup = false; + for (picks[r * 11 ..][0..k]) |x| dup = dup or x == id; + if (dup) continue; + picks[r * 11 + k] = id; + k += 1; + } + picks[r * 11 + 10] = if (round % 3 == 0) -1 else 512; + } + const n = rows * 11; + // host + const t0 = check.now(gpu.io); + const rp = try host.plan(picks[0..n], rows); + host_ns += @intCast(check.now(gpu.io) - t0); + try check.expect(rp.n_passes == 1, "a synthetic round split into {d} passes", .{rp.n_passes}); + try host.stage(rp, 0, s); + try host.release(rp, 0, s); + // device + try dp.toDevice(s); + try picks_dev.upload(0, std.mem.sliceAsBytes(picks[0..n])); + try dp.plan(s, 0, picks_dev.ptr, rows); + try dp.join(s); + try dp.restore(s, 0); + if (round % 37 == 36) try dp.toHost(); // owner switches mid-stream keep both views consistent + } + try s.synchronize(); + try host.lines.drain(); + try dp.toHost(); + // residency, stamps, counts + var same_lines: usize = 0; + for (host.slots_tab, dev.slots_tab, host.stamp_tab, dev.stamp_tab) |hs, ds, ht, dt| same_lines += @intFromBool(hs == ds and ht == dt); + const same_counts = std.mem.eql(u32, &host.counts[0], &dev.counts[0]); + std.debug.print("RESULT plan lines={d} max_rows={d} rounds={d}: lines+stamps equal {d}/{d}, counts equal {}, host hits {d}/{d} copies {d} evictions {d}; device hits {d}/{d} copies {d} evictions {d}; host plan {d:.1} us a round\n", .{ lines, max_rows, rounds, same_lines, lines, same_counts, host.stats.hit_slots, host.stats.hit_slots + host.stats.miss_slots, host.stats.misses, host.stats.evictions, dev.stats.hit_slots, dev.stats.hit_slots + dev.stats.miss_slots, dev.stats.misses, dev.stats.evictions, @as(f64, @floatFromInt(host_ns)) / 1e3 / @as(f64, @floatFromInt(rounds)) }); + try check.expect(same_lines == lines and same_counts, "device planner diverged from the host planner", .{}); + try check.expect(host.stats.hit_slots == dev.stats.hit_slots and host.stats.misses == dev.stats.misses and host.stats.evictions == dev.stats.evictions, "device planner statistics differ", .{}); + // tables: the same line (relative to each pool's lines) and width for every expert + const tb = host.lines.tables.len; + const th = try a.alloc(u8, tb); + defer a.free(th); + const td = try a.alloc(u8, tb); + defer a.free(td); + try host.lines.tables.download(0, th); + try dev.lines.tables.download(0, td); + var bad_tab: usize = 0; + for (0..512) |id| for (0..3) |w| { + const ph = std.mem.readInt(u64, th[host.lines.tlay.ptr[w] + id * 8 ..][0..8], .little); + const pd = std.mem.readInt(u64, td[dev.lines.tlay.ptr[w] + id * 8 ..][0..8], .little); + const kh = std.mem.readInt(i32, th[host.lines.tlay.k2[w] + id * 4 ..][0..4], .little); + const kd = std.mem.readInt(i32, td[dev.lines.tlay.k2[w] + id * 4 ..][0..4], .little); + const rh = if (ph == 0) 0 else ph - host.lines.mem.ptr + 1; + const rd = if (pd == 0) 0 else pd - dev.lines.mem.ptr + 1; + bad_tab += @intFromBool(rh != rd or kh != kd); + }; + // every resident line holds its expert's bytes (the copy kernel read the mapped bank) + const mem = try a.alloc(u8, dev.lines.mem.len); + defer a.free(mem); + try dev.lines.mem.download(0, mem); + var bad_lines: usize = 0; + var resident: usize = 0; + for (dev.slots_tab, 0..) |id, ln| { + if (id == fpool.offload.Cache.empty_id) continue; + resident += 1; + for (0..3) |w| { + const want = routed[id][w]; + const got = mem[ln * dev.lines.layout.line_bytes + dev.lines.layout.off[w] ..][0..want.len]; + if (!std.mem.eql(u8, want, got)) bad_lines += 1; + } + } + std.debug.print("RESULT plan tables: {d} of 1536 entries differ; {d} resident lines, {d} projections with wrong bytes\n", .{ bad_tab, resident, bad_lines }); + try check.expect(bad_tab == 0 and bad_lines == 0, "device tables or line bytes wrong", .{}); + check.pass("device planner == host planner over {d} rounds of up to {d} rows on {d} lines (tables, lines, stamps, counts, line bytes)", .{ rounds, max_rows, lines }); +} + +fn copyBench(gpu: Gpu) !void { + // Flash Next line: gate/up/down 3-bit, K 2560 x N 640: 160 x 40 tiles of 48 bytes = 614,400 bytes each + const seg: usize = 614_400; + const n_lines: usize = 30; + var bank = try cuda.HostBuffer.alloc(gpu.d, n_lines * 3 * seg); + defer bank.free(); + @memset(bank.bytes, 0x5a); + var dst = try cuda.DeviceBuffer.alloc(gpu.d, n_lines * 3 * seg); + defer dst.free(); + var mod = try cuda.Module.load(gpu.d, cuda.kernels.flashnext_exl3_dplan); + defer mod.unload(); + const f = try mod.function("fn_exl3_dcopy"); + 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(); + for ([_]usize{ 1, 3, 10 }) |k| { + // work list: k lines of three segments + var work: [10 * 9]i64 = undefined; + for (0..k) |i| for (0..3) |w| { + const off = (i * 3 + w) * seg; + work[i * 9 + w * 3 + 0] = @intCast(@intFromPtr(bank.bytes.ptr) + off); + work[i * 9 + w * 3 + 1] = @intCast(dst.ptr + off); + work[i * 9 + w * 3 + 2] = @intCast(seg); + }; + var wbuf = try cuda.DeviceBuffer.alloc(gpu.d, 10 * 9 * 8 + 16); + defer wbuf.free(); + try wbuf.upload(0, std.mem.sliceAsBytes(work[0 .. k * 9])); + const cnt = [1]i32{@intCast(k)}; + try wbuf.upload(10 * 9 * 8, std.mem.sliceAsBytes(&cnt)); + for ([_]u32{ 36, 72, 144, 288, 576 }) |blocks| { + const reps = 20; + var args: cuda.Args = .{}; + args.add(wbuf.ptr); + args.add(wbuf.ptr + 10 * 9 * 8); + try cuda.launch.launch(f, .{ .grid = .{ .x = blocks }, .block = .{ .x = 256 } }, s, &args); + try e0.record(s); + for (0..reps) |_| try cuda.launch.launch(f, .{ .grid = .{ .x = blocks }, .block = .{ .x = 256 } }, s, &args); + try e1.record(s); + try s.synchronize(); + const t = (try e0.elapsedMs(e1)) / reps; + std.debug.print("BENCH dcopy {d} lines ({d:.2} MiB) with {d} blocks: {d:.3} ms = {d:.1} GB/s\n", .{ k, @as(f64, @floatFromInt(k * 3 * seg)) / (1 << 20), blocks, t, @as(f64, @floatFromInt(k * 3 * seg)) / (t * 1e6) }); + } + // the copy engine, one H2D a segment + try e0.record(s); + for (0..20) |_| for (0..k * 3) |j| try s.d.check(s.d.api.cuMemcpyHtoDAsync_v2(dst.ptr + j * seg, bank.bytes[j * seg ..].ptr, seg, s.handle), "cuMemcpyHtoDAsync"); + try e1.record(s); + try s.synchronize(); + const t = (try e0.elapsedMs(e1)) / 20; + std.debug.print("BENCH dma {d} lines ({d:.2} MiB): {d:.3} ms = {d:.1} GB/s\n", .{ k, @as(f64, @floatFromInt(k * 3 * seg)) / (1 << 20), t, @as(f64, @floatFromInt(k * 3 * seg)) / (t * 1e6) }); + } +} + +pub fn planMode(gpu: Gpu) !void { + try copyBench(gpu); + try planCase(gpu, 10, 1, 400, 1); + try planCase(gpu, 47, 4, 600, 2); + try planCase(gpu, 63, 6, 600, 3); +} diff --git a/zig/tests/cuda/flashnext_exl3/graph/flashnext_exl3_graph_prefill_tests.zig b/zig/tests/cuda/flashnext_exl3/graph/flashnext_exl3_graph_prefill_tests.zig new file mode 100644 index 0000000000..2b27aa244c --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/graph/flashnext_exl3_graph_prefill_tests.zig @@ -0,0 +1,188 @@ +//! flashnext-exl3-graph prefill, prefill-prof, prof and ngram: layer graphs in the deep prefill and profiling windows. + +const std = @import("std"); +const platform = @import("sys"); +const flash = @import("flashnext"); +const check = @import("../../check.zig"); +const Gpu = check.Gpu; +const prompt = flash.prompt; +const X = flash.exl3; +const dplan = X.dplan; +const graph = X.graph; +const harness = @import("flashnext_exl3_graph_engine.zig"); +const Eng = harness.Eng; +const ms = harness.ms; + +/// A long prompt through the deep prefill (TF_FLASHNEXT_EXL3_PREFILL_ROWS sets its chunks), then `steps` greedy +pub fn prefillMode(gpu: Gpu, rank: usize, world: usize, port: u16, dir: []const u8, n: usize, steps: usize) !void { + const a = gpu.gpa; + const g = try Eng.init(gpu, rank, world, port, dir); + defer g.deinit(); + const ids = try a.alloc(u32, n); + defer a.free(ids); + var prng = std.Random.DefaultPrng.init(0x9e3779b9); + for (ids) |*x| x.* = 200 + prng.random().uintLessThan(u32, 60000); + const hs = try a.alloc(u64, steps + 1); + defer a.free(hs); + var ref: ?u64 = null; + var ok = true; + for ([_]graph.Mode{ .off, .off, .graph, .graph, .graph }) |m| { + graph.force = m; + try dplan.toHost(&g.e.exl3_dplan); + try g.reset(); + _ = try g.swap(3); + const p0 = graph.prefill_stats; + const f0 = g.e.exl3_prefill.stats; + const t0 = check.now(gpu.io); + try g.forward(ids, .{ .logits = .last }); + const t1 = check.now(gpu.io); + const p1 = graph.prefill_stats; + const f1 = g.e.exl3_prefill.stats; + try check.expect(f1.chunks > f0.chunks and f1.rows >= f0.rows + n, "rank {d}: deep prefill ran instead of falling back", .{rank}); + if (m == .graph) try check.expect(p1.captures + p1.updates + p1.replays > p0.captures + p0.updates + p0.replays, "rank {d}: prefill graph captured or replayed a layer", .{rank}); + hs[0] = std.hash.Wyhash.hash(0, g.seq.logits); + var tok = prompt.argmax(g.seq.logits); + for (0..steps) |i| { + tok = try g.round(tok); + hs[i + 1] = std.hash.Wyhash.hash(0, g.seq.logits); + } + const h = std.hash.Wyhash.hash(0, std.mem.sliceAsBytes(hs)); + if (world == 2) try check.expect(try g.swap(h) == h, "rank {d}: the ranks' prefill logits differ", .{rank}); + if (ref == null) ref = h; + const same = ref.? == h; + ok = ok and same; + if (rank == 0) std.debug.print("RESULT tp={d} prefill {s}: {d} rows in {d:.1} ms = {d:.1} rows/s; layer graphs captured {d} (updates {d}) replayed {d} eager {d} (capture {d:.1} ms); last-row + {d} decode logits {s} the first eager run\n", .{ world, @tagName(m), n, ms(t0, t1), @as(f64, @floatFromInt(n)) / (ms(t0, t1) / 1e3), p1.captures - p0.captures, p1.updates - p0.updates, p1.replays - p0.replays, p1.eager - p0.eager, @as(f64, @floatFromInt(p1.capture_ns - p0.capture_ns)) / 1e6, steps, if (same) "bit-equal to" else "DIFFER from" }); + } + graph.force = null; + try check.expect(ok, "rank {d}: prefill under graphs differs from eager", .{rank}); + check.pass("rank {d}/{d}: deep prefill of {d} rows with layer graphs and the decode after it bit-identical to eager", .{ rank, world, n }); +} + +/// One warmed TP2 deep-prefill forward inside a CUDA-profiler API range for Nsight Systems. +pub fn prefillProfMode(gpu: Gpu, rank: usize, world: usize, port: u16, dir: []const u8, n: usize, which: []const u8) !void { + var stage: []const u8 = "engine initialization"; + errdefer std.debug.print("prefill-prof rank={d} failed during {s}\n", .{ rank, stage }); + const a = gpu.gpa; + const g = try Eng.init(gpu, rank, world, port, dir); + defer g.deinit(); + defer graph.force = null; + const ids = try a.alloc(u32, n); + defer a.free(ids); + var prng = std.Random.DefaultPrng.init(0x9e3779b9); + for (ids) |*x| x.* = 200 + prng.random().uintLessThan(u32, 60000); + const use_graphs = std.mem.eql(u8, which, "graph"); + if (!use_graphs and !std.mem.eql(u8, which, "eager")) return error.BadMode; + graph.force = if (use_graphs) .graph else .off; + try dplan.toHost(&g.e.exl3_dplan); + stage = "graph warm-up"; + // For graph mode, capture each layer shape before the measured forward. + for (0..if (use_graphs) @as(usize, 2) else 1) |_| { + try g.reset(); + _ = try g.swap(3); + try g.forward(ids, .{ .logits = .last }); + try g.stream.synchronize(); + } + const before = graph.prefill_stats; + const f_before = g.e.exl3_prefill.stats; + stage = "profile setup"; + try g.reset(); + _ = try g.swap(3); + try g.stream.synchronize(); + var prof = try Profiler.open(); + _ = prof.start(); + stage = "measured forward"; + const t0 = check.now(gpu.io); + try g.forward(ids, .{ .logits = .last }); + const t1 = check.now(gpu.io); + _ = prof.stop(); + try g.stream.synchronize(); + const after = graph.prefill_stats; + const f_after = g.e.exl3_prefill.stats; + if (use_graphs and after.replays == before.replays) return error.NoPrefillGraphReplay; + const hash = std.hash.Wyhash.hash(0, g.seq.logits); + stage = "rank hash check"; + if (world == 2) try check.expect(try g.swap(hash) == hash, "rank {d}: profiled prefill logits differ across ranks", .{rank}); + std.debug.print("RESULT prefill-prof rank={d}/{d} mode={s} rows={d}: {d:.1} ms wall, graphs replayed {d}, eager {d}, captured {d}, updated {d}; expert copies {d} (prefetch {d}), passes {d}, pick wait {d:.1} ms, plan {d:.1} ms; logits {x}; Nsight range excludes warm-up\n", .{ rank, world, which, n, ms(t0, t1), after.replays - before.replays, after.eager - before.eager, after.captures - before.captures, after.updates - before.updates, f_after.copies - f_before.copies, f_after.prefetched - f_before.prefetched, f_after.passes - f_before.passes, @as(f64, @floatFromInt(f_after.host_wait_ns - f_before.host_wait_ns)) / 1e6, @as(f64, @floatFromInt(f_after.plan_ns - f_before.plan_ns)) / 1e6, hash }); +} + +/// cuProfilerStart/Stop from libcuda (nsys --capture-range=cudaProfilerApi), looked up at run time. +const Profiler = struct { + lib: platform.DynLib, + start: *const fn () callconv(.c) c_int, + stop: *const fn () callconv(.c) c_int, + fn open() !Profiler { + var lib = try platform.DynLib.open(platform.cuda_driver); + return .{ .start = lib.lookup(*const fn () callconv(.c) c_int, "cuProfilerStart") orelse return error.MissingSymbol, .stop = lib.lookup(*const fn () callconv(.c) c_int, "cuProfilerStop") orelse return error.MissingSymbol, .lib = lib }; + } +}; + +/// A profiling window for nsys: +pub fn profMode(gpu: Gpu, rank: usize, world: usize, port: u16, dir: []const u8, steps: usize, rows: usize, which: []const u8) !void { + const a = gpu.gpa; + const g = try Eng.init(gpu, rank, world, port, dir); + defer g.deinit(); + const ids = [_]u32{ 760, 6511, 314, 9338, 369, 11751, 13, 1049, 369, 3750, 364, 279, 242476, 300, 21262, 13 }; + graph.force = if (std.mem.eql(u8, which, "graph")) .graph else if (std.mem.eql(u8, which, "dplan")) .dplan else .off; + var prof = try Profiler.open(); + const toks = try a.alloc(u32, 4 * steps * rows + 64); + defer a.free(toks); + var phase: usize = 0; + while (phase < 2) : (phase += 1) { + try dplan.toHost(&g.e.exl3_dplan); + try g.reset(); + try g.forward(&ids, .{ .logits = .last }); + var tok = prompt.argmax(g.seq.logits); + _ = try g.swap(5); + if (phase == 1) _ = prof.start(); + const t0 = check.now(gpu.io); + for (0..steps) |i| { + tok = try g.round(tok); + toks[i] = tok; + } + const t1 = check.now(gpu.io); + for (0..steps) |i| { + for (0..rows) |j| toks[steps + i * rows + j] = toks[(i * rows + j) % steps]; + try g.forward(toks[steps + i * rows ..][0..rows], .{ .logits = .all, .logits_out = g.out.ptr }); + } + const t2 = check.now(gpu.io); + if (phase == 1) _ = prof.stop(); + if (rank == 0) std.debug.print("RESULT prof tp={d} {s} phase {d}: rounds {d:.2} ms each, windows of {d} rows {d:.2} ms each; graph captures {d} replays {d}; host prepare {d:.3} ms a step\n", .{ world, which, phase, ms(t0, t1) / @as(f64, @floatFromInt(steps)), rows, ms(t1, t2) / @as(f64, @floatFromInt(steps)), graph.stats.captures, graph.stats.replays, @as(f64, @floatFromInt(graph.stats.prepare_ns)) / 1e6 / @as(f64, @floatFromInt(@max(graph.stats.prepares, 1))) }); + } + graph.force = null; +} + +/// Host cost of the n-gram rows a step gathers (Ngram.rows), 1 and 3 tokens, fresh random histories. +pub fn ngramMode(gpu: Gpu, dir: []const u8) !void { + const a = gpu.gpa; + var ng = try flash.ple.Ngram.init(a, gpu.io, dir); + defer ng.deinit(); + const w = ng.width(); + const emb = try a.alloc(u16, 4 * w); + defer a.free(emb); + const xs = try a.alloc(f32, 4 * w / 32); + defer a.free(xs); + var prng = std.Random.DefaultPrng.init(7); + for ([_]usize{ 1, 3 }) |n| for (0..2) |rep| { + var tot: f64 = 0; + var worst: f64 = 0; + const reps = 64; + for (0..reps) |_| { + var hist = [2]i64{ prng.random().uintLessThan(u32, 240000), prng.random().uintLessThan(u32, 240000) }; + var ids: [3]u32 = undefined; + for (&ids) |*x| x.* = prng.random().uintLessThan(u32, 240000); + const h0 = hist; + const t0 = check.now(gpu.io); + try ng.rows(&hist, ids[0..n], emb[0 .. n * w], xs[0 .. n * w / 32]); + var t = ms(t0, check.now(gpu.io)); + if (rep == 1) { // the same rows again: their pages are resident now + hist = h0; + const t1 = check.now(gpu.io); + try ng.rows(&hist, ids[0..n], emb[0 .. n * w], xs[0 .. n * w / 32]); + t = ms(t1, check.now(gpu.io)); + } + tot += t; + worst = @max(worst, t); + } + std.debug.print("RESULT ngram rows of {d} tokens (pass {d}, 1 = warm repeat): {d:.3} ms mean, {d:.3} ms worst\n", .{ n, rep, tot / reps, worst }); + }; +} diff --git a/zig/tests/cuda/flashnext_exl3/graph/flashnext_exl3_graph_snapshot_tests.zig b/zig/tests/cuda/flashnext_exl3/graph/flashnext_exl3_graph_snapshot_tests.zig new file mode 100644 index 0000000000..917ffd5dbe --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/graph/flashnext_exl3_graph_snapshot_tests.zig @@ -0,0 +1,93 @@ +//! flashnext-exl3-graph snapshot: two lanes saved, restored and continued under eager and graph steps. + +const std = @import("std"); +const flash = @import("flashnext"); +const check = @import("../../check.zig"); +const Gpu = check.Gpu; +const prompt = flash.prompt; +const X = flash.exl3; +const dplan = X.dplan; +const graph = X.graph; +const vocab = prompt.vocabulary; +const harness = @import("flashnext_exl3_graph_engine.zig"); +const Eng = harness.Eng; +const cap_default = harness.cap_default; + +/// A snapshot's bytes on the host: its pinned buffer, or a copy of its device buffer. +fn bytesOf(gpu: Gpu, saved: *const flash.snapshot.Snap) ![]u8 { + const out = try gpu.gpa.alloc(u8, saved.buf.len); + errdefer gpu.gpa.free(out); + if (saved.host) |host| @memcpy(out, host.bytes[0..out.len]) else try gpu.d.check(gpu.d.api.cuMemcpyDtoH_v2(out.ptr, saved.buf.ptr, out.len), "cuMemcpyDtoH"); + return out; +} + +pub fn snapshotMode(gpu: Gpu, rank: usize, world: usize, port: u16, dir: []const u8, steps: usize) !void { + const snaps = flash.snapshot; + const allocator = gpu.gpa; + const g = try Eng.init(gpu, rank, world, port, dir); + defer g.deinit(); + defer graph.force = null; + var second = try g.e.newSeq(cap_default); + defer second.deinit(allocator); + const ids = try allocator.alloc(u32, 16384); + defer allocator.free(ids); + var random = std.Random.DefaultPrng.init(0x736e6170); + for (ids) |*token| token.* = 200 + random.random().uintLessThan(u32, 60000); + graph.force = .graph; + const first_end: usize = 3075; + const second_end: usize = 6917; + try g.forward(ids[0..first_end], .{ .logits = .last }); + try g.e.forward(flash.Tri, g.tri(), &second, ids[4096..second_end], .{ .logits = .last }); + const first_saved = try snaps.save(allocator, gpu.d, g.stream, &g.seq, null); + defer snaps.drop(allocator, g.stream, first_saved); + const second_saved = try snaps.save(allocator, gpu.d, g.stream, &second, null); + defer snaps.drop(allocator, g.stream, second_saved); + const expected = try allocator.alloc(u8, (steps + 1) * 2 * vocab * 2); + defer allocator.free(expected); + var reference: [2]?*snaps.Snap = .{ null, null }; + defer for (reference) |item| if (item) |saved| snaps.drop(allocator, g.stream, saved); + for ([_]graph.Mode{ .off, .graph, .graph, .graph }, 0..) |mode, run_index| { + graph.force = mode; + try dplan.toHost(&g.e.exl3_dplan); + g.e.round_graph.clear(); + try snaps.restore(g.stream, &g.seq, null, first_saved); + try snaps.restore(g.stream, &second, null, second_saved); + const before = graph.stats; + const prefill_before = graph.prefill_stats; + try g.forward(ids[first_end..][0..8192], .{ .logits = .last }); + try g.e.forward(flash.Tri, g.tri(), &second, ids[second_end..][0..8192], .{ .logits = .last }); + var tokens = [_]u32{ prompt.argmax(g.seq.logits), prompt.argmax(second.logits) }; + const sequences = [_]*prompt.Seq{ &g.seq, &second }; + for (0..steps + 1) |step| { + if (step > 0) { + try g.e.round(flash.Tri, g.tri(), &sequences, &tokens, g.out.ptr); + tokens = .{ prompt.argmax(g.seq.logits), prompt.argmax(second.logits) }; + } + for (sequences, 0..) |sequence, lane| { + const wanted = expected[(step * 2 + lane) * vocab * 2 ..][0 .. vocab * 2]; + if (run_index == 0) @memcpy(wanted, sequence.logits) else try check.expect(std.mem.eql(u8, wanted, sequence.logits), "rank {d} snapshot {s} lane {d} step {d} BF16 logits", .{ rank, @tagName(mode), lane, step }); + } + } + for (sequences, 0..) |sequence, lane| { + const saved = try snaps.save(allocator, gpu.d, g.stream, sequence, null); + if (run_index == 0) { + reference[lane] = saved; + } else { + defer snaps.drop(allocator, g.stream, saved); + const previous = reference[lane].?; + try check.expect(previous.at == saved.at and std.meta.eql(previous.history, saved.history), "rank {d} snapshot lane {d} exact position and history", .{ rank, lane }); + const old_bytes = try bytesOf(gpu, previous); + defer allocator.free(old_bytes); + const new_bytes = try bytesOf(gpu, saved); + defer allocator.free(new_bytes); + try check.expect(std.mem.eql(u8, old_bytes, new_bytes), "rank {d} snapshot lane {d} exact cache and recurrent state", .{ rank, lane }); + } + } + if (mode == .graph) { + std.debug.print("SNAPSHOT rank {d}/{d} graph cache run={d}: prefill captures +{d}, updates +{d}, replays +{d}, eager +{d}; decode replays +{d}\n", .{ rank, world, run_index, graph.prefill_stats.captures - prefill_before.captures, graph.prefill_stats.updates - prefill_before.updates, graph.prefill_stats.replays - prefill_before.replays, graph.prefill_stats.eager - prefill_before.eager, graph.stats.replays - before.replays }); + try check.expect(graph.stats.replays > before.replays, "rank {d} concurrent decode graph replay after restore", .{rank}); + } + std.debug.print("SNAPSHOT rank {d}/{d} {s}: two lanes, {d} decode steps, exact state and logits\n", .{ rank, world, @tagName(mode), steps }); + } + check.pass("rank {d}/{d}: concurrent snapshots preserve the BF16 continuation under eager and graph steps", .{ rank, world }); +} diff --git a/zig/tests/cuda/flashnext_exl3/graph/flashnext_exl3_graph_tests.zig b/zig/tests/cuda/flashnext_exl3/graph/flashnext_exl3_graph_tests.zig new file mode 100644 index 0000000000..c1b145ae06 --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/graph/flashnext_exl3_graph_tests.zig @@ -0,0 +1,39 @@ +//! flashnext-exl3-graph: the EXL3 offload pool planned on the device (cuda_exl3_dplan.zig) and the decode step as one +const std = @import("std"); +const check = @import("../../check.zig"); + +const Gpu = check.Gpu; +const harness = @import("flashnext_exl3_graph_engine.zig"); +const Eng = harness.Eng; +const plan_tests = @import("flashnext_exl3_graph_plan_tests.zig"); +const eq_tests = @import("flashnext_exl3_graph_eq_tests.zig"); +const prefill_tests = @import("flashnext_exl3_graph_prefill_tests.zig"); +const snapshot_tests = @import("flashnext_exl3_graph_snapshot_tests.zig"); + +pub fn run(gpu: Gpu, rest: []const [:0]const u8) !void { + if (rest.len >= 2 and std.mem.eql(u8, rest[0], "ngram")) return prefill_tests.ngramMode(gpu, rest[1]); + if (rest.len < 1) return error.MissingArgument; + if (std.mem.eql(u8, rest[0], "plan")) return plan_tests.planMode(gpu); + if (std.mem.eql(u8, rest[0], "frequency")) return @import("../offload/exl3_frequency_tests.zig").run(gpu, rest[1..]); + if (rest.len < 6) 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; + const steps = if (rest.len > 6 and !std.mem.eql(u8, mode, "prefill-prof")) try std.fmt.parseInt(usize, rest[6], 10) else 48; + const which: []const u8 = if (rest.len > 7) rest[7] else "dplan,graph"; + if (std.mem.eql(u8, mode, "rebind")) { + const engine = try Eng.init(gpu, rank, world, port, dir); + defer engine.deinit(); + return @import("../kv/exl3_state_rebind_tests.zig").run(engine, .{ .reset = Eng.reset, .forward = Eng.forward, .swap = Eng.swap }); + } + if (std.mem.eql(u8, mode, "snapshot")) return snapshot_tests.snapshotMode(gpu, rank, world, port, dir, steps); + if (std.mem.eql(u8, mode, "fork")) return eq_tests.forkMode(gpu, rank, world, port, dir, rest[5], steps, which); + if (std.mem.eql(u8, mode, "prof")) return prefill_tests.profMode(gpu, rank, world, port, dir, try std.fmt.parseInt(usize, rest[5], 10), steps, which); + if (std.mem.eql(u8, mode, "prefill")) return prefill_tests.prefillMode(gpu, rank, world, port, dir, try std.fmt.parseInt(usize, rest[5], 10), steps); + if (std.mem.eql(u8, mode, "prefill-prof")) return prefill_tests.prefillProfMode(gpu, rank, world, port, dir, try std.fmt.parseInt(usize, rest[5], 10), if (rest.len > 6) rest[6] else "graph"); + if (std.mem.eql(u8, mode, "eq")) return eq_tests.eqMode(gpu, rank, world, port, dir, rest[5], steps, which); + return error.UnknownMode; +} diff --git a/zig/tests/cuda/flashnext_exl3/kv/exl3_state_rebind_tests.zig b/zig/tests/cuda/flashnext_exl3/kv/exl3_state_rebind_tests.zig new file mode 100644 index 0000000000..0e83974f8c --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/kv/exl3_state_rebind_tests.zig @@ -0,0 +1,117 @@ +const std = @import("std"); +const cuda = @import("cuda"); +const flash = @import("flashnext"); +const check = @import("../../check.zig"); + +const graph = flash.exl3.graph; +const iterations = 4; +const prefix = [_]u32{ 760, 6511, 314, 9338, 369, 11751, 13, 1049, 369, 3750, 364, 279, 242476, 300, 21262, 13 }; + +fn require(name: [:0]const u8, value: []const u8) !void { + const setting = std.c.getenv(name); + try check.expect(setting != null and std.mem.eql(u8, std.mem.span(setting.?), value), "rebind fixture requires {s}={s}", .{ name, value }); +} + +fn stateBytes(sequence: *const flash.prompt.Seq) usize { + var total = sequence.tail.len; + for ([_][]cuda.DeviceBuffer{ sequence.state, sequence.conv }) |buffers| for (buffers) |buffer| { + total += buffer.len; + }; + return total; +} + +fn readState(sequence: *const flash.prompt.Seq, target: []u8) !void { + var offset: usize = 0; + for ([_][]cuda.DeviceBuffer{ sequence.state, sequence.conv }) |buffers| for (buffers) |buffer| { + try buffer.download(0, target[offset..][0..buffer.len]); + offset += buffer.len; + }; + try sequence.tail.download(0, target[offset..][0..sequence.tail.len]); +} + +pub fn run(engine: anytype, comptime operations: anytype) !void { + try require("TF_FLASHNEXT_EXL3_GRAPH_ASYNC", "0"); + graph.rebind.state_rebind = true; + graph.cache.capture_after_override = 1; + defer { + graph.rebind.state_rebind = false; + graph.cache.capture_after_override = null; + } + try check.expect(engine.seq.paged == null and engine.world == 2, "rebind fixture requires two ranks and the dense control KV layout", .{}); + var draft: flash.mtp.Draft = undefined; + draft.keep = .{}; + draft.verify_cache = .{}; + try check.expect(engine.seq.state.len == draft.spare.len, "rebind fixture requires the complete recurrent layer layout", .{}); + var made: usize = 0; + defer { + cuda.completion.drain(engine.gpu.d, engine.stream, "state rebind fixture consumers did not drain"); + graph.clear(); + for (draft.spare[0..made]) |*buffer| buffer.free(); + } + for (engine.seq.state) |state| { + draft.spare[made] = try cuda.DeviceBuffer.alloc(engine.gpu.d, state.len); + made += 1; + } + var linear: usize = 0; + for (engine.e.layers[0..engine.e.n_layers], 0..) |layer, index| { + if (!layer.linear) continue; + draft.keep.layers[index].state_out = draft.spare[linear].ptr; + linear += 1; + } + try check.expect(linear == made, "rebind fixture maps every physical recurrent layer", .{}); + const allocator = engine.gpu.gpa; + const state_bytes = stateBytes(&engine.seq); + const reference_state = try allocator.alloc(u8, iterations * state_bytes); + defer allocator.free(reference_state); + const actual_state = try allocator.alloc(u8, state_bytes); + defer allocator.free(actual_state); + const reference_logits = try allocator.alloc(u8, iterations * 16 * flash.prompt.vocabulary * 2); + defer allocator.free(reference_logits); + const actual_logits = try allocator.alloc(u8, 16 * flash.prompt.vocabulary * 2); + defer allocator.free(actual_logits); + const previous_mode = graph.force; + defer graph.force = previous_mode; + for ([_]usize{ 3, 4, 16 }) |width| { + const row_bytes = width * flash.prompt.vocabulary * 2; + var before: graph.Stats = undefined; + for ([_]bool{ false, true }) |captured| { + try flash.exl3.dplan.toHost(&engine.e.exl3_dplan); + graph.clear(); + graph.force = .off; + try operations.reset(engine); + for (draft.spare) |buffer| try buffer.fill8(0, engine.stream.handle); + try operations.forward(engine, &prefix, .{ .logits = .last }); + graph.force = if (captured) .graph else .dplan; + before = graph.stats; + for (0..iterations) |step| { + try operations.forward(engine, prefix[0..width], .{ .logits = .all, .logits_out = engine.out.ptr, .keep = &draft.keep }); + try engine.out.download(0, actual_logits[0..row_bytes]); + const input_before = engine.seq.state[0].ptr; + const output_before = draft.spare[0].ptr; + draft.open = true; + draft.settle(&engine.seq); + try check.expect(engine.seq.state[0].ptr == output_before and draft.spare[0].ptr == input_before, "rank {d}: width {d} step {d} actually swaps the live state pair", .{ engine.rank, width, step }); + try readState(&engine.seq, actual_state); + const wanted_state = reference_state[step * state_bytes ..][0..state_bytes]; + const wanted_logits = reference_logits[step * row_bytes ..][0..row_bytes]; + if (!captured) { + @memcpy(wanted_state, actual_state); + @memcpy(wanted_logits, actual_logits[0..row_bytes]); + } + const equal = std.mem.eql(u8, wanted_state, actual_state) and std.mem.eql(u8, wanted_logits, actual_logits[0..row_bytes]); + const other_equal = try operations.swap(engine, @intFromBool(equal)); + try check.expect(equal and other_equal == 1, "rank {d}: width {d} step {d} exact eager/graph logits, recurrent state, convolution and PLE tail", .{ engine.rank, width, step }); + const digest = std.hash.Wyhash.hash(0, actual_logits[0..row_bytes]); + const other_digest = try operations.swap(engine, digest); + try check.expect(other_digest == digest, "rank {d}: width {d} step {d} cross-rank BF16 logits", .{ engine.rank, width, step }); + } + } + const after = graph.stats; + const succeeded = after.rebinds - before.rebinds; + const failures = after.rebind_failures - before.rebind_failures; + const proved = succeeded == iterations - 1 and failures == 0 and after.rebind_attempts - before.rebind_attempts == succeeded and after.captures - before.captures == 1 and after.replays - before.replays == succeeded; + const other_proved = try operations.swap(engine, @intFromBool(proved)); + try check.expect(proved and other_proved == 1, "rank {d}: width {d} needs one capture and three actual successful widened rebind replays, zero failures; got {d}/{d}", .{ engine.rank, width, succeeded, failures }); + check.pass("rank {d}/2 width {d}: real state flips; exact eager/graph logits and state; actual widened rebinds {d}, failures {d}", .{ engine.rank, width, succeeded, failures }); + } +} 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 33f43bc8e0..9050b7f9b2 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], null, g.e.chunk, b, null); + g.attached = try xpass.attach(gpu.gpa, gpu.io, gpu.d, g.stream, &g.e.exl3_prefill, 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; @@ -350,7 +350,7 @@ fn tfMode(gpu: Gpu, dir: []const u8, golden_path: []const u8, lines: ?usize) !vo 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 }); + 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(s.cut_ns)) / 1e6, s.blocks, @as(f64, @floatFromInt(s.plan_ns)) / 1e6 }); } /// Greedy `steps` tokens after `ids`: the prefill's last row, then a decode round a token. diff --git a/zig/tests/cuda/flashnext_exl3/offload/exl3_frequency_tests.zig b/zig/tests/cuda/flashnext_exl3/offload/exl3_frequency_tests.zig index 0dfacd8475..ba47512af0 100644 --- a/zig/tests/cuda/flashnext_exl3/offload/exl3_frequency_tests.zig +++ b/zig/tests/cuda/flashnext_exl3/offload/exl3_frequency_tests.zig @@ -51,6 +51,8 @@ pub fn run(gpu: check.Gpu, arguments: []const [:0]const u8) !void { 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); + var stream = try cuda.Stream.init(gpu.d, true); + defer stream.deinit(); 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); @@ -69,6 +71,21 @@ pub fn run(gpu: check.Gpu, arguments: []const [:0]const u8) !void { 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); + var picks: [11]i32 = @splat(-1); + for (0..10) |index| picks[index] = @intCast(start + 19 - index); + var device_picks = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.asBytes(&picks)); + defer device_picks.free(); + var planner = try flash.exl3.dplan.DevPlan.init(gpu.gpa, gpu.d, &.{&pool}); + defer planner.deinit(); + try planner.toDevice(stream); + try planner.plan(stream, 0, device_picks.ptr, 1); + try planner.join(stream); + try planner.restore(stream, 0); + try planner.toHost(); + for (0..20) |index| { + const expected = before[start + index] + @as(u32, @intFromBool(index >= 10)); + try check.expect(pool.counts[7][start + index] == expected, "device learned count survives profile restart", .{}); + } @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 }); @@ -78,5 +95,5 @@ pub fn run(gpu: check.Gpu, arguments: []const [:0]const u8) !void { 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", .{}); + check.pass("frequency disk restart, live device counts, rank isolation and resized warm pools", .{}); } 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 index 2dcd93790d..369ed7d746 100644 --- a/zig/tests/cuda/flashnext_exl3/offload/exl3_pool_equal_tests.zig +++ b/zig/tests/cuda/flashnext_exl3/offload/exl3_pool_equal_tests.zig @@ -1,4 +1,4 @@ -//! exl3-pool: gate 2, the pool against the resident experts. +//! exl3-pool: gate 2 and 3, the pool against the resident experts and the graph replay. const std = @import("std"); const cuda = @import("cuda"); const core = @import("core"); @@ -149,3 +149,84 @@ pub fn equalRun(gpu: Gpu, ex: *const exl3e.Experts, s: cuda.Stream, ld: *const L 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 }); } + +// ---------------------------------------------------------------------------------------------------- gate 3 + +pub fn graphRun(gpu: Gpu, ex: *const exl3e.Experts, s: cuda.Stream, ld: *const LayerData) !void { + const gpa = gpu.gpa; + const lines = 10; + var pool = try fpool.Pool.init(gpa, gpu.io, gpu.d, ld.source(), ld.scales(), .{ .layer = @intCast(ld.layer), .lines = lines }, null); + defer pool.deinit(); + 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(0xc0ffee + ld.layer, 0.6); + var prng = std.Random.DefaultPrng.init(0x96a9 + ld.layer); + const rng = prng.random(); + var x: [dims]u16 = undefined; + var picks: [slots]i32 = undefined; + var wts: [slots]f32 = undefined; + var exec: ?cuda.graph.Exec = null; + defer if (exec) |*e| e.deinit(); + var g: ?cuda.graph.Graph = null; + defer if (g) |*gg| gg.deinit(); + var swaps: usize = 0; + var last: ?[]u8 = null; + defer if (last) |b| gpa.free(b); + var distinct: usize = 0; + const states = 8; + for (0..states) |si| { + genRound(&z, rng, 1, &x, &picks, &wts); + if (si == 5) for (picks[0..5], 0..) |*p, i| { + p.* = @intCast(pool.slots_tab[i]); // half the row from what the lines hold: a partial hit + }; + try vp.set(&x, &picks, &wts); + try vr.set(&x, &picks, &wts); + const r = try pool.plan(&picks, 1); + try check.expect(!r.masked(), "one row must fit 10 lines", .{}); + try pool.stage(r, 0, s); + swaps += r.n_copies; + const l = pool.tables(); + if (exec == null) { + // capture once over the tables, with the outputs poisoned so a capture that ran shows up + try s.synchronize(); + try vp.sbuf.fill8(0xff, s.handle); + try s.synchronize(); + try cuda.graph.beginCapture(s, .thread_local); + try call(ex, s, l, vp, vp.picks.ptr, 1, false); + g = try cuda.graph.endCapture(s); + exec = try g.?.instantiate(); + try exec.?.upload(s); + const still = try vp.y(gpa, 1); + defer gpa.free(still); + try check.expect(std.mem.allEqual(u8, still, 0xff), "capture executed the kernels", .{}); + std.debug.print(" [graph] runSlots captured once: {d} nodes over the pool's fixed table addresses\n", .{try g.?.nodeCount()}); + } + try vp.sbuf.fill8(0xff, s.handle); + try exec.?.launchOn(s); + try pool.release(r, 0, s); + try s.synchronize(); + const yg = try vp.y(gpa, 1); + defer gpa.free(yg); + // eager over the same tables, then the fully resident layer + try vp.sbuf.fill8(0xff, s.handle); + try call(ex, s, l, vp, vp.picks.ptr, 1, false); + try call(ex, s, ld.res, vr, vr.picks.ptr, 1, false); + try s.synchronize(); + const ye = try vp.y(gpa, 1); + defer gpa.free(ye); + const yr = try vr.y(gpa, 1); + defer gpa.free(yr); + try check.sameBytes("graph replay vs eager", yg, ye); + try check.sameBytes("graph replay vs fully resident", yg, yr); + if (last) |b| { + if (!std.mem.eql(u8, b, yg)) distinct += 1; + gpa.free(b); + } + last = try gpa.dupe(u8, yg); + std.debug.print(" [graph] state {d}: {d} lines swapped by H2D before the replay ({d} held); replay == eager == resident, {d} slot rows byte for byte\n", .{ si, r.n_copies, r.held, slots }); + } + try check.expect(swaps > lines and distinct == states - 1, "the states did not exercise swaps ({d}) or distinct outputs ({d})", .{ swaps, distinct }); + check.pass("layer {d}: one captured runSlots graph replayed {d} states after {d} line swaps (no recapture), byte-identical to eager and to the resident layer", .{ ld.layer, states, swaps }); +} diff --git a/zig/tests/cuda/flashnext_exl3/offload/exl3_pool_tests.zig b/zig/tests/cuda/flashnext_exl3/offload/exl3_pool_tests.zig index 53a2a3e594..1adc593586 100644 --- a/zig/tests/cuda/flashnext_exl3/offload/exl3_pool_tests.zig +++ b/zig/tests/cuda/flashnext_exl3/offload/exl3_pool_tests.zig @@ -52,6 +52,7 @@ pub fn run(gpu: Gpu, rest: []const [:0]const u8) !void { try equal.equalRun(gpu, &ex, s, ld, 24, 4, 0x24 + layer); try equal.equalRun(gpu, &ex, s, ld, 64, 0, 0x64 + layer); } + if (all or std.mem.eql(u8, mode, "graph")) try equal.graphRun(gpu, &ex, s, ld); if (std.mem.eql(u8, mode, "bench")) try bench.benchRun(gpu, &ex, s, ld); } } diff --git a/zig/tests/cuda/flashnext_exl3/prefill/exl3_prefill_engine.zig b/zig/tests/cuda/flashnext_exl3/prefill/exl3_prefill_engine.zig index ac49300e55..56272a09aa 100644 --- a/zig/tests/cuda/flashnext_exl3/prefill/exl3_prefill_engine.zig +++ b/zig/tests/cuda/flashnext_exl3/prefill/exl3_prefill_engine.zig @@ -43,12 +43,12 @@ const Eng = struct { 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); + if (deep_rows > 0) try xpre.reserve(gpu.gpa, gpu.d, &g.e.exl3_prefill, 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); + g.attached = try xpass.attach(gpu.gpa, gpu.io, gpu.d, g.stream, &g.e.exl3_prefill, 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 { @@ -61,7 +61,7 @@ const Eng = struct { } pub fn deinit(g: *Eng) void { - xpre.release(); + xpre.release(&g.e.exl3_prefill); g.seq.deinit(g.gpu.gpa); g.e.deinit(); g.glue.deinit(); @@ -109,7 +109,7 @@ 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; + g.e.exl3_prefill.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 }); @@ -200,10 +200,10 @@ pub fn engineRun(gpu: Gpu, dir: []const u8, rest: []const [:0]const u8) !void { 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 = .{}; + g.e.exl3_prefill.stats = .{}; const t_new = try prefillOnce(g, ids, o.enabled, new); const p1 = g.poolStats(); - const xs_ = xpre.state.stats; + const xs_ = g.e.exl3_prefill.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); diff --git a/zig/tests/cuda/flashnext_exl3/tp/exl3_tpprefill_tests.zig b/zig/tests/cuda/flashnext_exl3/tp/exl3_tpprefill_tests.zig index f142b76931..fedbe9570f 100644 --- a/zig/tests/cuda/flashnext_exl3/tp/exl3_tpprefill_tests.zig +++ b/zig/tests/cuda/flashnext_exl3/tp/exl3_tpprefill_tests.zig @@ -95,16 +95,16 @@ const Eng = struct { 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); + if (deep_rows > 0) try xpre.reserveAt(gpu.gpa, gpu.d, &g.e.exl3_prefill, 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); + g.attached = try xpass.attach(gpu.gpa, gpu.io, gpu.d, g.stream, &g.e.exl3_prefill, 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(); + xpre.release(&g.e.exl3_prefill); g.out.free(); g.seq.deinit(g.gpu.gpa); g.e.deinit(); @@ -191,14 +191,14 @@ fn once(g: *Eng, ids: []const u32, deep: bool, rows: usize, steps: usize) !Run { var o = xpre.optionsFromEnv(); o.enabled = deep; o.rows = rows; - xpre.state.options = o; - xpre.state.stats = .{}; + g.e.exl3_prefill.options = o; + g.e.exl3_prefill.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 st = g.e.exl3_prefill.stats; const logits = try a.alloc(u16, vocab); @memcpy(std.mem.sliceAsBytes(logits), g.seq.logits); const greedy = try a.alloc(u32, steps + 1); 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 index cb71c5f15c..1d23f42895 100644 --- 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 @@ -2,6 +2,9 @@ const std = @import("std"); const check = @import("../../check.zig"); const decode = @import("flashnext_exl3_tpe_decode_tests.zig"); +const trace = @import("flashnext_exl3_tpe_trace_tests.zig"); +const hits = @import("flashnext_exl3_tpe_hits_tests.zig"); +const graph = @import("flashnext_exl3_tpe_graph_tests.zig"); const Gpu = check.Gpu; @@ -13,7 +16,18 @@ pub fn run(gpu: Gpu, rest: []const [:0]const u8) !void { 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, "hit")) return hits.hitMode(gpu, rest[5..]); + if (std.mem.eql(u8, mode, "policy")) return hits.policyMode(gpu, rest[5..]); + if (std.mem.eql(u8, mode, "graph")) return graph.graphMode(gpu, rank, port); if (std.mem.eql(u8, mode, "bench")) return decode.benchMode(gpu, rank, world, port, dir); + if (std.mem.eql(u8, mode, "trace")) { + if (rest.len < 7) return error.MissingArgument; + const rounds = if (rest.len > 7) try std.fmt.parseInt(usize, rest[7], 10) else 1024; + if (rest.len > 10) return error.UnexpectedArgument; + const schedule = if (rest.len > 8) rest[8] else "0"; + const repeats = if (rest.len > 9) try std.fmt.parseInt(usize, rest[9], 10) else 1; + return trace.traceMode(gpu, rank, world, port, dir, rest[5], rest[6], rounds, schedule, repeats); + } 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); 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 index 51f25aa643..37d1371b5e 100644 --- a/zig/tests/cuda/flashnext_exl3/tp/flashnext_exl3_tpe_engine.zig +++ b/zig/tests/cuda/flashnext_exl3/tp/flashnext_exl3_tpe_engine.zig @@ -101,7 +101,7 @@ pub const Eng = struct { 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); + g.attached = try xpass.attach(gpu.gpa, gpu.io, gpu.d, g.stream, &g.e.exl3_prefill, 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; @@ -160,7 +160,7 @@ pub const Eng = struct { 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 }); + 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(s.cut_ns)) / 1e6, s.blocks, @as(f64, @floatFromInt(s.plan_ns)) / 1e6 }); } // -- golden --------------------------------------------------------------------------------------------------------- diff --git a/zig/tests/cuda/flashnext_exl3/tp/flashnext_exl3_tpe_graph_tests.zig b/zig/tests/cuda/flashnext_exl3/tp/flashnext_exl3_tpe_graph_tests.zig new file mode 100644 index 0000000000..827c8818b9 --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/tp/flashnext_exl3_tpe_graph_tests.zig @@ -0,0 +1,139 @@ +//! flashnext-exl3-tpe: NCCL all-reduce inside a captured graph. +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; + +// -- NCCL in a captured graph --------------------------------------------------------------------------------------- + +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) return; + if (state != 7) return c.lib.check(state, "NCCL async"); + } +} + +fn allReduce(c: *Collective, s: cuda.Stream, src: u64, dst: u64, n: usize) !void { + const r = c.lib.api.ncclAllReduce(src, dst, n, .f32, .sum, c.comm, s.handle); + if (r != 0 and r != 7) return c.lib.check(r, "ncclAllReduce"); + try ready(c); +} + +pub fn graphMode(gpu: Gpu, rank: usize, port: u16) !void { + const a = gpu.gpa; + const rk = try Ranks.open(gpu, rank, port); + defer rk.close(a); + const c = &rk.collective; + 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 all_ok = true; + // one decode row's MoE sum (2,560 fp32), one row's per-slot y, and a 128-row block + for ([_]usize{ 2560, 28160, 327680 }) |n| for ([_]usize{ 1, 48 }) |per_graph| { + const mine = try a.alloc(f32, n * per_graph); + defer a.free(mine); + const theirs = try a.alloc(f32, n * per_graph); + defer a.free(theirs); + var p0 = std.Random.DefaultPrng.init(0x9a + rank); + var p1 = std.Random.DefaultPrng.init(0x9a + (1 - rank)); + for (mine) |*v| v.* = p0.random().floatNorm(f32); + for (theirs) |*v| v.* = p1.random().floatNorm(f32); + var src = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(mine)); + defer src.free(); + var dst = try cuda.DeviceBuffer.alloc(gpu.d, n * per_graph * 4); + defer dst.free(); + try dst.fill8(0, s.handle); + try s.synchronize(); + _ = try rk.swap(n); + // eager: per_graph all-reduces in a row + const reps: usize = if (per_graph == 1) 50 else 5; + for (0..3) |_| for (0..per_graph) |k| try allReduce(c, s, src.ptr + k * n * 4, dst.ptr + k * n * 4, n); + try s.synchronize(); + try e0.record(s); + for (0..reps) |_| for (0..per_graph) |k| try allReduce(c, s, src.ptr + k * n * 4, dst.ptr + k * n * 4, n); + try e1.record(s); + try e1.synchronize(); + const t_eager = (try e0.elapsedMs(e1)) / @as(f32, @floatFromInt(reps)); + try dst.fill8(0, s.handle); + try s.synchronize(); + _ = try rk.swap(n + 1); + // captured: the same calls in one graph, instantiated, replayed + const cap_ok = blk: { + cuda.graph.beginCapture(s, .thread_local) catch |err| { + std.debug.print("rank {d}: beginCapture failed: {s}\n", .{ rank, @errorName(err) }); + break :blk false; + }; + var failed = false; + for (0..per_graph) |k| allReduce(c, s, src.ptr + k * n * 4, dst.ptr + k * n * 4, n) catch |err| { + std.debug.print("rank {d}: ncclAllReduce under capture failed: {s}\n", .{ rank, @errorName(err) }); + failed = true; + break; + }; + var g = cuda.graph.endCapture(s) catch |err| { + std.debug.print("rank {d}: endCapture failed: {s}\n", .{ rank, @errorName(err) }); + break :blk false; + }; + defer g.deinit(); + if (failed) break :blk false; + const nodes = g.nodeCount() catch 0; + var exec = g.instantiate() catch |err| { + std.debug.print("rank {d}: instantiate failed: {s}\n", .{ rank, @errorName(err) }); + break :blk false; + }; + defer exec.deinit(); + exec.upload(s) catch {}; + for (0..3) |_| try exec.launchOn(s); + try s.synchronize(); + const got = try a.alloc(f32, n * per_graph); + defer a.free(got); + try dst.download(0, std.mem.sliceAsBytes(got)); + var eq: usize = 0; + for (got, mine, theirs) |x, m, t| { + const want = if (rank == 0) m + t else t + m; + eq += @intFromBool(@as(u32, @bitCast(x)) == @as(u32, @bitCast(want))); + } + try e0.record(s); + for (0..reps) |_| try exec.launchOn(s); + try e1.record(s); + try e1.synchronize(); + const t_graph = (try e0.elapsedMs(e1)) / @as(f32, @floatFromInt(reps)); + std.debug.print("RESULT rank {d} graph: {d} x ncclAllReduce of {d} fp32 ({d} graph nodes): eager {d:.4} ms, graph replay {d:.4} ms ({d:.4} vs {d:.4} ms an all-reduce); replayed sums == fp32 a+b on {d}/{d}\n", .{ rank, per_graph, n, nodes, t_eager, t_graph, t_eager / @as(f32, @floatFromInt(per_graph)), t_graph / @as(f32, @floatFromInt(per_graph)), eq, got.len }); + break :blk eq == got.len; + }; + all_ok = all_ok and cap_ok; + const both = try rk.swap(@intFromBool(cap_ok)); + all_ok = all_ok and both == 1; + }; + try check.expect(all_ok, "rank {d}: an all-reduce captured in a graph failed or gave other bits", .{rank}); + check.pass("rank {d}: ncclAllReduce on the nonblocking communicator captured in a CUDA graph replays on both ranks with the eager bits (SHM transport)", .{rank}); +} diff --git a/zig/tests/cuda/flashnext_exl3/tp/flashnext_exl3_tpe_hits_tests.zig b/zig/tests/cuda/flashnext_exl3/tp/flashnext_exl3_tpe_hits_tests.zig new file mode 100644 index 0000000000..8568126d3a --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/tp/flashnext_exl3_tpe_hits_tests.zig @@ -0,0 +1,154 @@ +//! flashnext-exl3-tpe: per-rank hit rates and the cache policy replay (host only). +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; + +// -- hit rates per rank (host only) --------------------------------------------------------------------------------- + +const top_k = 10; +const slots = 11; +const max_rows = 16; + +fn simulate(gpa: std.mem.Allocator, layer: u8, picks: []const u16, tokens: usize, prompt_n: usize, lines: usize, r: tp.Rank, counts: *[fpool.shape.layers][fpool.shape.experts]u32, dec: *[2]u64, all_: *[2]u64) !void { + const policy = fpool.offload; + 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 rnd = try fpool.Round.init(gpa, lines, max_rows * slots); + defer rnd.deinit(gpa); + var buf: [max_rows * slots]i32 = undefined; + var prev: [512]u16 = undefined; + var prev_n: usize = 0; + var t: usize = 0; + var round: u32 = 0; + while (t < tokens) { + const rows: usize = if (t < prompt_n) @min(max_rows, prompt_n - t) else 1; + for (0..rows) |i| { + for (0..top_k) |k| { + const id: usize = picks[(t + i) * top_k + k]; + buf[i * slots + k] = if (r.owns(id)) @intCast(id) else -1; + } + buf[i * slots + top_k] = if (r.owns(512)) 512 else -1; + } + round += 1; + try fpool.round.plan(&v, layer, buf[0 .. rows * slots], counts, round, round, .{ .protect = prev[0..prev_n] }, &rnd); + @memcpy(prev[0..rnd.n_ids], rnd.ids[0..rnd.n_ids]); + prev_n = rnd.n_ids; + all_[0] += rnd.hit_slots; + all_[1] += rnd.miss_slots; + if (t >= prompt_n) { + dec[0] += rnd.hit_slots; + dec[1] += rnd.miss_slots; + } + t += rows; + } +} + +fn rate(x: [2]u64) f64 { + return if (x[0] + x[1] == 0) 0 else @as(f64, @floatFromInt(x[0])) / @as(f64, @floatFromInt(x[0] + x[1])); +} + +pub fn hitMode(gpu: Gpu, files: []const [:0]const u8) !void { + const gpa = gpu.gpa; + const counts = try gpa.create([fpool.shape.layers][fpool.shape.experts]u32); + defer gpa.destroy(counts); + 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_n: 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]); + const per = try gpa.alloc(u16, tokens * top_k); + defer gpa.free(per); + std.debug.print("real router picks {s}: {d} tokens ({d} prompt in 16-row rounds, {d} decode rows) x {d} layers\n", .{ path, tokens, prompt_n, tokens - prompt_n, layers }); + // routed slots a decode token each rank owns + var own: [2]u64 = .{ 0, 0 }; + for (prompt_n..tokens) |t| for (0..layers) |l| for (0..top_k) |k| { + own[@intFromBool(raw[(t * layers + l) * top_k + k] >= 256)] += 1; + }; + std.debug.print(" decode routed slots on rank 0 / rank 1: {d} / {d} ({d:.1}% on rank 0)\n", .{ own[0], own[1], 100 * @as(f64, @floatFromInt(own[0])) / @as(f64, @floatFromInt(own[0] + own[1])) }); + for ([_]usize{ 10, 24, 48, 56, 64, 96, 128, 160 }) |lines| { + var line_txt: [3][64]u8 = undefined; + var outs: [3][]const u8 = undefined; + const worlds = [3][2]usize{ .{ 1, 0 }, .{ 2, 0 }, .{ 2, 1 } }; + var miss_tok: [3]f64 = undefined; + for (worlds, 0..) |w, i| { + const r = try tp.Rank.init(w[0], w[1]); + counts.* = @splat(@splat(0)); + var dec: [2]u64 = .{ 0, 0 }; + var all_: [2]u64 = .{ 0, 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]); + try simulate(gpa, @intCast(l), per, tokens, prompt_n, lines, r, counts, &dec, &all_); + } + miss_tok[i] = @as(f64, @floatFromInt(dec[1])) / @as(f64, @floatFromInt(tokens - prompt_n)); + outs[i] = try std.fmt.bufPrint(&line_txt[i], "hit {d:.3} ({d:.1} misses/token)", .{ rate(dec), miss_tok[i] }); + } + std.debug.print(" {d:>3} lines a layer: tp=1 {s} | tp=2 rank 0 {s} | rank 1 {s} | slowest rank's misses/token {d:.1} vs tp=1 {d:.1}\n", .{ lines, outs[0], outs[1], outs[2], @max(miss_tok[1], miss_tok[2]), miss_tok[0] }); + } + } + check.pass("per-rank hit rates printed (host replay through the pool's own planner)", .{}); +} + +pub fn policyMode(gpu: Gpu, files: []const [:0]const u8) !void { + if (files.len == 0) return error.MissingTraceFiles; + const allocator = gpu.gpa; + for (files) |path| { + const bytes = try std.Io.Dir.cwd().readFileAlloc(gpu.io, path, allocator, .limited(X.dplan.demand_trace.pool_bytes + 1)); + defer allocator.free(bytes); + var snapshot = try X.dplan.demand_trace.Snapshot.decode(allocator, bytes); + defer snapshot.deinit(); + const view = try snapshot.view(); + if (view.dropped() != 0) return error.TraceDroppedRecords; + if (view.recorded() == 0) return error.EmptyTrace; + const first = try view.round(0); + const result = try X.dplan.policy_eval.evaluateTrace(allocator, view); + const lifetime = result.lifetime; + const lru = result.lru; + const decayed = result.decayed_frequency; + const admission = result.decayed_admission; + std.debug.print("POLICY {s} pool={d} layer={d} lines={d} rounds={d} drops={d}\n", .{ path, view.words[2], view.words[3], first.lines, view.recorded(), view.dropped() }); + if (view.checkpoints() != 0) std.debug.print(" {d} ownership checkpoints: each device segment starts from its recorded state; host intervals excluded\n", .{view.checkpoints()}); + std.debug.print(" lifetime picks miss={d} copies={d} evictions={d} mapped={d}\n", .{ lifetime.miss_picks, lifetime.copies, lifetime.evictions, lifetime.mapped_experts }); + std.debug.print(" lru picks miss={d} copies={d} evictions={d} mapped={d}\n", .{ lru.miss_picks, lru.copies, lru.evictions, lru.mapped_experts }); + std.debug.print(" decayed picks miss={d} copies={d} evictions={d} mapped={d}\n", .{ decayed.miss_picks, decayed.copies, decayed.evictions, decayed.mapped_experts }); + std.debug.print(" guarded picks miss={d} copies={d} evictions={d} mapped={d} mapped_picks={d}\n", .{ admission.miss_picks, admission.copies, admission.evictions, admission.mapped_experts, admission.mapped_picks }); + std.debug.print(" Belady unique-miss lower bound={d} (not copy/time bound)\n", .{result.belady_unique_misses_lower_bound}); + } +} diff --git a/zig/tests/cuda/flashnext_exl3/tp/flashnext_exl3_tpe_trace_tests.zig b/zig/tests/cuda/flashnext_exl3/tp/flashnext_exl3_tpe_trace_tests.zig new file mode 100644 index 0000000000..91aae400b1 --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/tp/flashnext_exl3_tpe_trace_tests.zig @@ -0,0 +1,139 @@ +//! flashnext-exl3-tpe: the routing trace over a prompt schedule. +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; + +fn traceSchedule(allocator: std.mem.Allocator, text: []const u8, prompts: usize) ![]usize { + var result: std.ArrayList(usize) = .empty; + errdefer result.deinit(allocator); + var fields = std.mem.splitScalar(u8, text, ','); + while (fields.next()) |field| { + const index = std.fmt.parseInt(usize, field, 10) catch return error.InvalidTraceSchedule; + if (index >= prompts) return error.InvalidTraceSchedule; + try result.append(allocator, index); + } + return result.toOwnedSlice(allocator); +} + +test "trace schedule preserves prompt order and repeated selections" { + const schedule = try traceSchedule(std.testing.allocator, "2,0,2,1", 3); + defer std.testing.allocator.free(schedule); + try std.testing.expectEqualSlices(usize, &.{ 2, 0, 2, 1 }, schedule); + for ([_][]const u8{ "", "0,", "3", "-1" }) |invalid| + try std.testing.expectError(error.InvalidTraceSchedule, traceSchedule(std.testing.allocator, invalid, 3)); +} + +pub fn traceMode(gpu: Gpu, rank: usize, world: usize, port: u16, dir: []const u8, golden_path: []const u8, prefix: []const u8, rounds: usize, schedule_text: []const u8, repeats: usize) !void { + if (rounds == 0) return error.InvalidTraceRounds; + if (repeats == 0) return error.InvalidTraceSchedule; + if (world != 2) return error.TraceRequiresTp2; + if (!X.graph.enabled()) return error.TraceRequiresGraphs; + var golden = try Golden.open(gpu, golden_path); + defer golden.parsed.deinit(); + const a = gpu.gpa; + const schedule = try traceSchedule(a, schedule_text, golden.prompts().len); + defer a.free(schedule); + const segments = try std.math.mul(usize, schedule.len, repeats); + const total_rounds = try std.math.mul(usize, segments, rounds); + var record_bound: usize = 0; + for (schedule) |index| { + const count = golden.prompts()[index].object.get("ids").?.array.items.len; + if (count == 0 or count >= cap_default or rounds >= cap_default - count) return error.InvalidTraceRounds; + record_bound = try std.math.add(usize, record_bound, try std.math.add(usize, count, rounds)); + } + record_bound = try std.math.mul(usize, record_bound, repeats); + X.dplan.demand_trace.requested = true; // the planner keeps its demand trace + defer X.dplan.demand_trace.requested = false; + const g = try Eng.init(gpu, rank, world, port, dir); + defer g.deinit(); + const initial_planner = (try X.dplan.forLayers(a, gpu.d, &g.e.exl3_dplan, g.e.layers[0..g.e.n_layers])) orelse return error.TracePlannerMissing; + if (initial_planner.trace_mem == null) return error.TraceDisabled; + for (initial_planner.pools, initial_planner.pool_states, 0..) |pool, entry, index| { + const header = try X.dplan.demand_trace.header(index, pool.opt.layer, entry.lines, initial_planner.max_rows); + if (record_bound > header[8]) return error.TraceScheduleExceedsCapacity; + } + std.debug.print("TRACE sequential shared-pool workload (not concurrent): schedule={s}, repeats={d}, segments={d}, decode rounds={d}, conservative record bound={d}\n", .{ schedule_text, repeats, segments, total_rounds, record_bound }); + const generated = try a.alloc(u32, try std.math.add(usize, rounds, 1)); + defer a.free(generated); + var output_hash: u64 = 0; + for (0..segments) |segment| { + const selection = schedule[segment % schedule.len]; + if (segment != 0) { + try gpu.d.synchronize(); + g.seq.deinit(a); + g.seq_live = false; + g.seq = try g.e.newSeq(cap_default); + g.seq_live = true; + } + const ids = try ints(a, golden.prompts()[selection].object.get("ids").?); + defer a.free(ids); + try g.reset(); + try g.forward(ids, .{ .logits = .last }); + var token = prompt.argmax(g.seq.logits); + const before = X.graph.coverage[0][1][@backingInt(X.graph.Outcome.replay)]; + generated[0] = token; + for (1..generated.len) |index| { + token = try g.round(token); + generated[index] = token; + } + if (X.graph.coverage[0][1][@backingInt(X.graph.Outcome.replay)] == before) return error.TraceRequiresGraphReplay; + const segment_hash = std.hash.Wyhash.hash(0, std.mem.sliceAsBytes(generated)); + try check.expect(try g.swap(segment_hash) == segment_hash, "trace segment {d}: tp=2 output hashes differ", .{segment}); + output_hash = std.hash.Wyhash.hash(output_hash, std.mem.sliceAsBytes(generated)); + std.debug.print("TRACE segment={d} prompt={d} rounds={d} output={x}\n", .{ segment, selection, rounds, segment_hash }); + } + try check.expect(try g.swap(output_hash) == output_hash, "trace run rank {d}: tp=2 output hashes differ", .{rank}); + if (X.dplan.current(&g.e.exl3_dplan)) |planner| { + if (planner != initial_planner) return error.TracePlannerChanged; + try X.dplan.toHost(&g.e.exl3_dplan); + var decoder_pools: usize = 0; + for (planner.pools, 0..) |pool, index| { + var snapshot = try planner.downloadTraceAfterConsumers(a, index); + defer snapshot.deinit(); + const view = try snapshot.view(); + if (view.dropped() != 0) return error.TraceDroppedRecords; + if (pool.opt.layer < 48) { + decoder_pools += 1; + if (view.recorded() < total_rounds) return error.TraceMissingDecoderPool; + } else if (view.recorded() == 0) continue; + const entry = planner.pool_states[index]; + const terminal_words = try a.alloc(u32, entry.words); + defer a.free(terminal_words); + try planner.state.download(entry.byteOffset(), std.mem.sliceAsBytes(terminal_words)); + const replay = try X.dplan.replay.verifyTraceWithFinal(a, view, terminal_words); + const bytes = try view.encode(a); + defer a.free(bytes); + const path = try std.fmt.allocPrintSentinel(a, "{s}.rank-{d}.pool-{d}.exdpt", .{ prefix, rank, index }, 0); + defer a.free(path); + try std.Io.Dir.cwd().writeFile(gpu.io, .{ .sub_path = path, .data = bytes }); + std.debug.print("TRACE v{d} rank {d} pool {d} layer {d}: rounds={d}/{d}, checkpoints={d}, verified transitions={d}, picks hit/miss={d}/{d}, copies={d}, evictions={d}, in-place experts/picks={d}/{d}, bytes={d}, output={x}\n", .{ X.dplan.demand_trace.version, rank, index, pool.opt.layer, view.recorded(), view.words[8], view.checkpoints(), replay.matched_transitions, replay.delta.hit_picks, replay.delta.miss_picks, replay.delta.copies, replay.delta.evictions, replay.delta.inplace, replay.delta.mapped_picks, bytes.len, output_hash }); + } + if (decoder_pools != 48) return error.TraceDecoderPoolCount; + } else return error.TracePlannerMissing; +} diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index b713bbd283..6e83323ffa 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -108,6 +108,7 @@ const flashnext_exl3_tests = @import("flashnext_exl3/model/flashnext_exl3_tests. 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 flashnext_exl3_graph_tests = @import("flashnext_exl3/graph/flashnext_exl3_graph_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"); @@ -151,7 +152,13 @@ const usage = \\ 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-exl3-tpe RANK 2 PORT DIR trace GOLDEN PREFIX [ROUNDS [SCHEDULE [REPEATS]]] + \\ graph-on TP2 v3; rounds/segment=1024, CSV zero-based prompt schedule=0, repeats=1 + \\ sequential shared-pool workload, independent sequence/KV; not concurrent requests + \\ e.g. trace GOLDEN PREFIX 128 0,1,2 2; rejects schedules exceeding 1 MiB/pool + \\ flashnext-exl3-tpe 0 1 0 - policy FILES... compare lifetime/LRU/frequency/bypass on expert demand + \\ flashnext-exl3-tpe RANK WORLD PORT DIR tf|gen GOLDEN DUMP | bench | hit FILES | graph the Engine at tp=1/2 with per-rank pools + \\ flashnext-exl3-graph plan | RANK WORLD PORT DIR eq GOLDEN [STEPS] [MODES] the pool planned on the device and decode graphs vs the host cut \\ 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 @@ -247,7 +254,7 @@ const usage = \\ 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-pool DIR [equal|graph|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) @@ -371,6 +378,7 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { 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-exl3-graph")) return flashnext_exl3_graph_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);