From 5d7684c3dd06d1af60fe208d42ef23d823613ae5 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 | 8 + docs/multi-rank-receipts.md | 2 + docs/recipes/exl3.md | 4 +- docs/recipes/flash-next-cuda.md | 21 + tools/exl3_len_abba.py | 97 +++ tools/exl3_longctx_client.py | 120 ++++ tools/exl3_pack_sizes.py | 26 + tools/exl3_serve.sh | 15 + tools/exl3_serve_gate.py | 53 ++ tools/exl3_single_request_bench.py | 49 ++ tools/exl3_snapshot_gate.py | 91 +++ tools/exl3_text_abba.py | 46 ++ tools/tests/test_exl3_snapshot_gate.py | 86 +++ zig/build/cuda.zig | 49 +- zig/build/cuda_kernels.zig | 2 + zig/build/exl3_host.zig | 6 + zig/build/flashnext_checks.zig | 47 +- .../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 | 4 +- .../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 | 36 +- .../families/flashnext/cuda_batch_gather.zig | 40 ++ zig/src/families/flashnext/cuda_batch_ple.zig | 32 + .../families/flashnext/cuda_batch_state.zig | 241 ++++--- zig/src/families/flashnext/cuda_config.zig | 9 +- zig/src/families/flashnext/cuda_embed.zig | 22 + zig/src/families/flashnext/cuda_lanes.zig | 43 +- .../flashnext/cuda_lanes_admission_test.zig | 14 +- .../families/flashnext/cuda_lanes_bytes.zig | 3 + .../families/flashnext/cuda_lanes_load.zig | 43 +- .../families/flashnext/cuda_lanes_snap.zig | 22 +- zig/src/families/flashnext/cuda_launch.zig | 10 + .../flashnext/cuda_paging_batch_test.zig | 2 + zig/src/families/flashnext/cuda_ple.zig | 17 +- .../families/flashnext/cuda_projection.zig | 3 + .../flashnext/cuda_projection_load.zig | 3 + zig/src/families/flashnext/cuda_prompt.zig | 111 ++-- .../flashnext/cuda_prompt_lifecycle.zig | 14 + .../families/flashnext/cuda_prompt_load.zig | 16 +- .../flashnext/cuda_prompt_telemetry.zig | 197 ++++++ .../families/flashnext/cuda_prompt_test.zig | 53 ++ zig/src/families/flashnext/cuda_prompt_wt.zig | 29 +- .../flashnext/cuda_scratch_layout_test.zig | 53 ++ .../families/flashnext/cuda_scratch_sizes.zig | 12 +- 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 +++ .../flashnext/exl3/model/cuda_exl3_pass.zig | 374 +++++++++++ .../flashnext/exl3/model/cuda_exl3_ple.zig | 520 +++++++++++++++ .../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 | 424 +++++++++++++ .../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 | 346 ++++++++++ .../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 | 532 ++++++++++++++++ .../flashnext/exl3/pack/cuda_exl3_memory.zig | 135 ++++ .../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 | 128 ++++ .../flashnext/exl3/tp/cuda_exl3_fork_test.zig | 105 +++ zig/src/native/cuda.zig | 2 +- zig/src/native/cuda_snaps.zig | 44 +- zig/src/native/cuda_startup_test.zig | 14 +- zig/src/native/cuda_streams.zig | 3 +- zig/tests/cuda/contract_family_sets.zig | 6 +- zig/tests/cuda/contract_flashnext_exl3.zig | 261 ++++++++ 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 | 537 ++++++++++++++++ .../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 | 32 +- zig/tests/cuda/runtime_tests.zig | 12 + 150 files changed, 18267 insertions(+), 362 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/build/exl3_host.zig 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_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_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 100644 zig/tests/cuda/contract_flashnext_exl3.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 51baac0601..15792dd44b 100644 --- a/docs/kernel-contract.md +++ b/docs/kernel-contract.md @@ -343,6 +343,14 @@ 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. +That launch, `fn_exl3_split_f16`, also runs gates 1 to 4 on the CUDA runner (`contract_flashnext_exl3.zig`): an fp16 +row splits into bf16 hi = bf16(x) and lo = bf16(x - hi), every output's bits must equal float64's (NaN stays NaN; an +infinity gives hi = Inf and lo = NaN), and the device's hi + lo must equal x for every finite pattern. Its shapes +hold all 65,536 fp16 patterns and partial last blocks (n = 1,000 and 65,537). Three controls go red: one flushing NaN +and Inf to zero (gate 1), one launching whole blocks only (gate 3, at every tail) and one writing hi over its input +(gate 2). Gate 4 is N/A: it has no sums. 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 fe12f71a6b..78e2b19bdc 100644 --- a/docs/multi-rank-receipts.md +++ b/docs/multi-rank-receipts.md @@ -19,6 +19,8 @@ Placement is family-specific; use the selected family's current native recipe. ` New community-v2 receipts carry an optional `group` object with 1-8 loaded ranks, per-node hardware and a `single-node` or `multi-node` category. The server gathers it at startup and HTTP reads its frozen snapshot. `settings.rank_count` must agree. [Loaded group receipts](loaded-group-receipts.md) defines the fields, validation and receiver partitions. Ranks sharing a node are not multiple nodes; identical hardware names on two nodes do not merge them. +Flash Next registers engine cleanup before constructing its loaded-rank inventory, so a failed startup inventory releases the loaded engine. + V1 remains byte-for-byte unchanged, accepting 1, 2 or null as before. Legacy v2 receipts without `group` keep that bound. Missing, unsupported or inconsistent inventories remain unranked. No client or environment rank hint can fill them. ROCm distributed inventory remains unavailable. Community-v2 publication stays disabled until the receiver accepts and separates its new fields. ## The settings the ranks agreed on diff --git a/docs/recipes/exl3.md b/docs/recipes/exl3.md index c1184f7827..cd7a081015 100644 --- a/docs/recipes/exl3.md +++ b/docs/recipes/exl3.md @@ -7,8 +7,8 @@ GLM-5.3-Flash. ## Where it stands - The native engine has the format and a reference decoder in core (`zig/src/core/exl3_format.zig`, @akol1, #496), - and Metal has a standalone EXL3 projection (`zig/src/core/exl3_projection.zig`, @akol1, #681). No model serves an - EXL3 checkpoint in this release. + and Metal has a standalone EXL3 projection (`zig/src/core/exl3_projection.zig`, @akol1, #681). Flash Next serves an + EXL3 pack on CUDA (sm_120, one rank; see [its recipe](flash-next-cuda.md#exl3-checkpoints)); no other model does. - CUDA projections and routed experts on this format are in review (#740, #741, #806). A model needs them and a loader for its EXL3 tensors before it can serve a checkpoint; the 27B on CUDA is the first planned. - The Python engine served EXL3 on CUDA from 0.3.6 to 0.6.6 through the module @vcruz305 wrote (#42): every codebook, diff --git a/docs/recipes/flash-next-cuda.md b/docs/recipes/flash-next-cuda.md index 0554230046..646c339309 100644 --- a/docs/recipes/flash-next-cuda.md +++ b/docs/recipes/flash-next-cuda.md @@ -110,6 +110,27 @@ tensorfold-native serve /path/to/checkpoint --temperature 0 `_chunks,_merge`) runs those entry points on our kernels and the rest on the set's, to put a difference down to one kernel. +## EXL3 checkpoints + +An EXL3 pack (see [EXL3 weights](exl3.md)) serves on one sm_120 rank with its routed experts in pinned host memory and +a pool of expert lines a layer on the device. `tools/exl3_serve.sh` starts it. These switches tune it: + +| Name | Default | Effect | +| --- | --- | --- | +| `TF_FLASHNEXT_EXPERT_LINES` | from free memory | Expert lines a layer keeps on the device. | +| `TF_FLASHNEXT_EXPERT_RESERVE_MB` | the server's reserve (`TENSORFOLD_MEMORY_RESERVE_GIB`) | MiB left free after the pools. | +| `TF_FLASHNEXT_EXL3_STREAMS` | 1 | Streams whose caches the pool leaves room for, when `--parallel` does not say. | +| `TF_FLASHNEXT_EXL3_LINE_DELTAS` | unset | `POOL:+N,POOL:-N` lines added to or removed from single pools of the uniform count. | +| `TF_FLASHNEXT_EXL3_PROFILE_DIR` | unset | Directory where expert-use profiles persist between runs. | +| `TF_FLASHNEXT_EXL3_PROFILE_SIZING` | `0` | `1` sizes the pools from the stored profile; ignored with `TF_FLASHNEXT_EXPERT_LINES`. | +| `TF_FLASHNEXT_EXL3_ROUTED` | host | `device` keeps every routed expert on the device. | +| `TF_FLASHNEXT_EXL3_ARCH` | unset | `unqualified` serves on a device that is not sm_120, without a claim. | +| `TF_FLASHNEXT_PLE_READ` | `map` | `pread` reads the n-gram embedding rows by positional read. | +| `TF_FLASHNEXT_SNAP_HOST` | `1` above 32,768 tokens of context, else `0` | Whether kept prompt-cache states live in pinned host memory. | +| `TF_FLASHNEXT_CACHE_TELEMETRY` | `0` | Diagnostics only: `1` logs `EXL3-CACHE` expert-cache counters per phase, with a sync at each boundary. | +| `TF_FLASHNEXT_EXL3_LAYER_TELEMETRY` | `0` | Diagnostics only: `1`, with the switch above, adds per-layer counters. | +| `TENSORFOLD_BIN`, `TF_SRC` | built binary, the checkout | Diagnostics only: the binary `tools/exl3_*.sh` run and the checkout `tools/exl3_*.py` read their corpus from. | + ## What the core runs and what the family runs The family plugs into the core's pieces and adds no scheduler: diff --git a/tools/exl3_len_abba.py b/tools/exl3_len_abba.py new file mode 100644 index 0000000000..b97ca1d928 --- /dev/null +++ b/tools/exl3_len_abba.py @@ -0,0 +1,97 @@ +"""Concurrent two-request HTTP bench (repeated-ID prompts) for 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..6f268682f1 --- /dev/null +++ b/tools/exl3_longctx_client.py @@ -0,0 +1,120 @@ +#!/usr/bin/env python3 +"""Long-context checks: PORT needle CTX[,CTX..] CONC TEXT | PORT cont N[,N..] GEN IDS (needle recall, continuation).""" +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..87c476807e --- /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 from the headers: 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")) +# The checkout whose sources are the corpus. +root = os.environ.get("TF_SRC") or os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +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..cfa5ec343c --- /dev/null +++ b/tools/exl3_text_abba.py @@ -0,0 +1,46 @@ +"""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")) + # The checkout whose sources are the corpus. + root = os.environ.get("TF_SRC") or os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + 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 0e6e7ef813..328eb6db11 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -116,7 +116,7 @@ pub fn targets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.bu // tf-tp-watch-check: the ranks' lines with no GPU or NCCL; a failing or dying rank stops every rank b.installArtifact(b.addExecutable(.{ .name = "tf-tp-watch-check", .root_module = b.createModule(.{ .root_source_file = b.path("zig/tests/cuda/tp_watch_check.zig"), .target = target, .optimize = optimize, .link_libc = true, .strip = strip, .imports = &.{.{ .name = "cuda", .module = cuda }} }) })); @import("living_dense.zig").target(b, target, optimize, cuda, mods.qwen35); - flashnextChecks(b, target, optimize, cuda, mods.core, mods.lanes, mods.qwen_vision, mods.api, opts); + @import("flashnext_checks.zig").cudaTargets(b, target, optimize, cuda, mods.core, flashnextModule(b, target, optimize, cuda, mods.core, mods.lanes, mods.qwen_vision, mods.api), opts); b.step("compile-cuda-tests", "Compile CUDA runner without executing a GPU").dependOn(&runner_exe.step); const startup = b.addObject(.{ .name = "modelopt-startup", .root_module = b.createModule(.{ .root_source_file = b.path("zig/modelopt_startup_compile.zig"), .target = target, .optimize = optimize, .link_libc = true, .imports = &.{.{ .name = "native", .module = native.engines }} }) }); _ = startup.getEmittedBin(); @@ -473,6 +473,7 @@ 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, mods.core); } + @import("exl3_host.zig").steps(b, flashnextModule(b, host, .debug, cuda, mods.core, mods.lanes, mods.qwen_vision, mods.api)); // the tokenizer's own tests (memo, steps, regex, unicode) run here: no other test binary reaches them 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:"; @@ -552,49 +553,3 @@ pub fn checkDistFatbins(b: *std.Build, dir: []const u8, privacy_tool: *std.Build privacy.dependOn(&check.step); return privacy; } - -/// Explicit model diagnostics; these tools do not qualify a checkpoint merely by compiling. -fn flashnextChecks(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, cuda: *std.Build.Module, core: *std.Build.Module, lanes: *std.Build.Module, qwen_vision: *std.Build.Module, api: *std.Build.Module, opts: @import("options.zig").Cuda) void { - const flashnext = flashnextModule(b, target, optimize, cuda, core, lanes, qwen_vision, api); - // not installed: one embedding row through the captured _embed cubin - const embed_check = b.createModule(.{ .root_source_file = b.path("zig/src/families/flashnext/cuda_embed_main.zig"), .target = target, .optimize = optimize, .link_libc = true }); - embed_check.addImport("cuda", cuda); - embed_check.addImport("core", core); - embed_check.addImport("flashnext", flashnext); - const embed_run = b.addRunArtifact(b.addExecutable(.{ .name = "flashnext-embed", .root_module = embed_check })); - const embed_model = opts.flashnext_model; - const embed_kernels = opts.flashnext_kernels; - const embed_token = opts.flashnext_token; - if (embed_model != null and embed_kernels != null) embed_run.addArgs(&.{ embed_model.?, embed_kernels.?, embed_token }); - b.step("flashnext-embed", "Upload the Flash Next embedding table and run _embed for one token").dependOn(&embed_run.step); - const prompt_check = b.createModule(.{ .root_source_file = b.path("zig/src/families/flashnext/cuda_prompt_main.zig"), .target = target, .optimize = optimize, .link_libc = true }); - prompt_check.addImport("cuda", cuda); - prompt_check.addImport("core", core); - prompt_check.addImport("flashnext", flashnext); - const prompt_run = b.addRunArtifact(b.addExecutable(.{ .name = "flashnext-prompt", .root_module = prompt_check })); - const prompt_ids = opts.flashnext_ids; - const prompt_count = opts.flashnext_count; - if (embed_model != null and embed_kernels != null and prompt_ids != null) prompt_run.addArgs(&.{ embed_model.?, embed_kernels.?, prompt_ids.?, prompt_count }); - b.step("flashnext-prompt", "Run one Flash Next prompt on the captured CUDA kernels").dependOn(&prompt_run.step); - // not installed: our own glue against a captured set on one loaded engine (scores of every prompt row, replies) - const compare_check = b.createModule(.{ .root_source_file = b.path("zig/src/families/flashnext/cuda_compare_main.zig"), .target = target, .optimize = optimize, .link_libc = true }); - compare_check.addImport("cuda", cuda); - compare_check.addImport("flashnext", flashnext); - const compare_run = b.addRunArtifact(b.addExecutable(.{ .name = "flashnext-compare", .root_module = compare_check })); - if (embed_model != null and embed_kernels != null and prompt_ids != null) compare_run.addArgs(&.{ embed_model.?, embed_kernels.?, prompt_ids.? }); - b.step("flashnext-compare", "Our Flash Next glue against a captured set: every prompt row's scores, perplexity, greedy replies").dependOn(&compare_run.step); - // not installed: the native block selection against Triton's lists (capture_sparse.sh's fixtures) - const select_check = b.createModule(.{ .root_source_file = b.path("zig/src/families/flashnext/cuda_select_main.zig"), .target = target, .optimize = optimize, .link_libc = true }); - select_check.addImport("cuda", cuda); - const select_run = b.addRunArtifact(b.addExecutable(.{ .name = "flashnext-select", .root_module = select_check })); - if (opts.flashnext_fixtures) |dir| select_run.addArgs(&.{dir}); - b.step("flashnext-select", "fn_qsa_select against Triton's _select / _select_tiles lists on captured score rows").dependOn(&select_run.step); - const learn_check = b.createModule(.{ .root_source_file = b.path("zig/src/families/flashnext/cuda_learn_main.zig"), .target = target, .optimize = optimize, .link_libc = true }); - learn_check.addImport("cuda", cuda); - learn_check.addImport("flashnext", flashnext); - const learn_exe = b.addExecutable(.{ .name = "flashnext-learn", .root_module = learn_check }); - const learn_run = b.addRunArtifact(learn_exe); - if (embed_model != null and prompt_ids != null) learn_run.addArgs(&.{ embed_model.?, prompt_ids.? }); - b.step("flashnext-learn", "Trainer forward identity and all-path finite-difference gates").dependOn(&learn_run.step); - b.step("compile-flashnext-learn", "Compile the real trainer runner without executing it").dependOn(&learn_exe.step); -} diff --git a/zig/build/cuda_kernels.zig b/zig/build/cuda_kernels.zig index c0d496d61c..946cd3cd79 100644 --- a/zig/build/cuda_kernels.zig +++ b/zig/build/cuda_kernels.zig @@ -148,6 +148,8 @@ pub const kernels = [_]Kernel{ .{ .name = "drama3", .flags = &.{"-O3"} }, .{ .name = "drama3_fp8", .flags = &.{"-O3"} }, .{ .name = "drama3_codec", .flags = &.{"-O3"} }, + // ours: Flash Next's EXL3 glue (an fp16 row split exactly into two bf16) + .{ .name = "flashnext_exl3", .src = "flashnext_exl3/flashnext_exl3", .flags = glue }, }; /// 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_host.zig b/zig/build/exl3_host.zig new file mode 100644 index 0000000000..3e46cb84d5 --- /dev/null +++ b/zig/build/exl3_host.zig @@ -0,0 +1,6 @@ +//! The EXL3 host tests of the Flash Next module. +const std = @import("std"); + +pub fn steps(b: *std.Build, 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); +} diff --git a/zig/build/flashnext_checks.zig b/zig/build/flashnext_checks.zig index 2b0e7a8386..28a8d6d1d0 100644 --- a/zig/build/flashnext_checks.zig +++ b/zig/build/flashnext_checks.zig @@ -1,4 +1,4 @@ -//! Flash Next's device checks on synthetic weights: programs that run on the GPU, each under its own step. +//! Flash Next's device checks and explicit model diagnostics, each under its own build step. const std = @import("std"); @@ -14,3 +14,48 @@ pub fn targets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.bu const about = "Flash Next's products at every oQ width against float64 (worst column), every row's bits in every window"; b.step("tf-flashnext-widths", about).dependOn(&b.addInstallArtifact(exe, .{}).step); } + +/// Explicit model diagnostics; these tools do not qualify a checkpoint merely by compiling. +pub fn cudaTargets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, cuda: *std.Build.Module, core: *std.Build.Module, flashnext: *std.Build.Module, opts: @import("options.zig").Cuda) void { + // not installed: one embedding row through the captured _embed cubin + const embed_check = b.createModule(.{ .root_source_file = b.path("zig/src/families/flashnext/cuda_embed_main.zig"), .target = target, .optimize = optimize, .link_libc = true }); + embed_check.addImport("cuda", cuda); + embed_check.addImport("core", core); + embed_check.addImport("flashnext", flashnext); + const embed_run = b.addRunArtifact(b.addExecutable(.{ .name = "flashnext-embed", .root_module = embed_check })); + const embed_model = opts.flashnext_model; + const embed_kernels = opts.flashnext_kernels; + const embed_token = opts.flashnext_token; + if (embed_model != null and embed_kernels != null) embed_run.addArgs(&.{ embed_model.?, embed_kernels.?, embed_token }); + b.step("flashnext-embed", "Upload the Flash Next embedding table and run _embed for one token").dependOn(&embed_run.step); + const prompt_check = b.createModule(.{ .root_source_file = b.path("zig/src/families/flashnext/cuda_prompt_main.zig"), .target = target, .optimize = optimize, .link_libc = true }); + prompt_check.addImport("cuda", cuda); + prompt_check.addImport("core", core); + prompt_check.addImport("flashnext", flashnext); + const prompt_run = b.addRunArtifact(b.addExecutable(.{ .name = "flashnext-prompt", .root_module = prompt_check })); + const prompt_ids = opts.flashnext_ids; + const prompt_count = opts.flashnext_count; + if (embed_model != null and embed_kernels != null and prompt_ids != null) prompt_run.addArgs(&.{ embed_model.?, embed_kernels.?, prompt_ids.?, prompt_count }); + b.step("flashnext-prompt", "Run one Flash Next prompt on the captured CUDA kernels").dependOn(&prompt_run.step); + // not installed: our own glue against a captured set on one loaded engine (scores of every prompt row, replies) + const compare_check = b.createModule(.{ .root_source_file = b.path("zig/src/families/flashnext/cuda_compare_main.zig"), .target = target, .optimize = optimize, .link_libc = true }); + compare_check.addImport("cuda", cuda); + compare_check.addImport("flashnext", flashnext); + const compare_run = b.addRunArtifact(b.addExecutable(.{ .name = "flashnext-compare", .root_module = compare_check })); + if (embed_model != null and embed_kernels != null and prompt_ids != null) compare_run.addArgs(&.{ embed_model.?, embed_kernels.?, prompt_ids.? }); + b.step("flashnext-compare", "Our Flash Next glue against a captured set: every prompt row's scores, perplexity, greedy replies").dependOn(&compare_run.step); + // not installed: the native block selection against Triton's lists (capture_sparse.sh's fixtures) + const select_check = b.createModule(.{ .root_source_file = b.path("zig/src/families/flashnext/cuda_select_main.zig"), .target = target, .optimize = optimize, .link_libc = true }); + select_check.addImport("cuda", cuda); + const select_run = b.addRunArtifact(b.addExecutable(.{ .name = "flashnext-select", .root_module = select_check })); + if (opts.flashnext_fixtures) |dir| select_run.addArgs(&.{dir}); + b.step("flashnext-select", "fn_qsa_select against Triton's _select / _select_tiles lists on captured score rows").dependOn(&select_run.step); + const learn_check = b.createModule(.{ .root_source_file = b.path("zig/src/families/flashnext/cuda_learn_main.zig"), .target = target, .optimize = optimize, .link_libc = true }); + learn_check.addImport("cuda", cuda); + learn_check.addImport("flashnext", flashnext); + const learn_exe = b.addExecutable(.{ .name = "flashnext-learn", .root_module = learn_check }); + const learn_run = b.addRunArtifact(learn_exe); + if (embed_model != null and prompt_ids != null) learn_run.addArgs(&.{ embed_model.?, prompt_ids.? }); + b.step("flashnext-learn", "Trainer forward identity and all-path finite-difference gates").dependOn(&learn_run.step); + b.step("compile-flashnext-learn", "Compile the real trainer runner without executing it").dependOn(&learn_exe.step); +} 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 dfa823c433..1d647fa163 100644 --- a/zig/src/core/contract/registry.zig +++ b/zig/src/core/contract/registry.zig @@ -320,6 +320,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 = 52, .static = 1168, .bound = 1024 }, .{ .regs = 52, .static = 1168, .bound = 1024 }, .{ .regs = 63, .static = 1168, .bound = 1024 }, .{ .regs = 63, .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); @@ -350,7 +351,6 @@ pub fn usageFor(r: CudaRecord, release: ?limits.Nvcc) [limits.sms.len]U { } return u; } -/// Restored Qwen launches retain a pending entry wherever a compiler receipt is absent. pub const cuda_pending = @import("qwen_whole_records.zig").records ++ @import("cuda_lane_pending.zig").records ++ @import("weight_only_records.zig").cuda ++ @import("hc_activation_records.zig").cuda ++ @import("qwen4_records.zig").cuda; /// Gate 5's planted CUDA controls: a launch past its __launch_bounds__, one needing more registers than a block has. pub const cuda_controls = [_]CudaRecord{ @@ -576,7 +576,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 ++ @import("flashnext_attn_lanes_pending.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 ++ @import("flashnext_attn_lanes_pending.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 ea33d031ae..e86b427d30 100644 --- a/zig/src/cuda/kernels.zig +++ b/zig/src/cuda/kernels.zig @@ -20,6 +20,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 be49ffbe2e..6f489cb0e2 100644 --- a/zig/src/cuda/root.zig +++ b/zig/src/cuda/root.zig @@ -70,6 +70,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"); @@ -120,6 +121,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 ee28d09152..71913a3cac 100644 --- a/zig/src/families/flashnext/cuda_batch.zig +++ b/zig/src/families/flashnext/cuda_batch.zig @@ -10,6 +10,8 @@ const load = @import("cuda_prompt_load.zig"); const parts = @import("cuda_batch_state.zig"); const late = @import("cuda_batch_late.zig"); const hc = @import("cuda_batch_hc.zig"); +const shared = @import("cuda_batch_side.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"); @@ -129,6 +131,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); } @@ -225,27 +231,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 { @@ -436,6 +422,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 f98fbf5057..c2b561b7cb 100644 --- a/zig/src/families/flashnext/cuda_batch_state.zig +++ b/zig/src/families/flashnext/cuda_batch_state.zig @@ -219,14 +219,16 @@ pub fn Staging(comptime Q: type) type { const sizes = @import("cuda_scratch_sizes.zig"); pub const head_rows_one = sizes.head_rows_one; pub const Sizes = sizes.Sizes; +pub const ScratchFormat = sizes.ScratchFormat; const hcPartRows = sizes.hcPartRows; -const partRows = sizes.partRows; +const partRowsFor = sizes.partRowsFor; /// Device scratch for chunks of up to `cap` rows. pub const Scratch = struct { cap: usize, /// The shared expert's stream (cuda_batch_side.zig); null: it runs after the routed experts on the pass's stream. side: ?shared.Side = null, + format: ScratchFormat = .classic, bufs: std.ArrayList(cuda.DeviceBuffer) = .empty, ids: u64 = 0, pos: u64 = 0, @@ -357,119 +359,126 @@ 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, z: Sizes) usize { - const head_rows = z.head_rows; + /// 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", "picks_ep", "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", "lanes_tab", + }; + + pub fn allocationSizes(cap: usize, z: Sizes) [66]?usize { + const format = z.format; 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 + hcPartRows(cap) * (low + 4) * 4; - n += cap * wide * 2 + (if (z.lane_reduce) 8 * cap * dims * 4 else 0); - n += cap * low * 2 + cap * (low / 32) * 4 + cap * streams * 2 + cap * dims * 2 + cap * (dims / 32) * 4; - n += cap * proj_row * 2 + partRows(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 + lanes_bytes; - n += cap * (experts_n + 1) * 4 + pairs * 4 * 2 + pairs * 4 + items * 3 * 4 + 8 + pairs * 4 + nblk * (experts_n + 1) * 4 + pairs * 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 + head_rows * 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, + hcPartRows(cap) * (low + 4) * 4, + cap * wide * 2, + if (z.lane_reduce) 8 * cap * dims * 4 else null, + cap * low * 2, + cap * (low / 32) * 4, + cap * streams * 2, + cap * dims * 2, + cap * (dims / 32) * 4, + cap * proj_row * 2, + partRowsFor(format, 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, + 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) z.head_rows * vocab else null, + half * 4, + 4096, + lanes_bytes, + }; + } + + pub fn bytes(cap: usize, z: Sizes) usize { + var total: usize = 0; + for (allocationSizes(cap, z)) |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, z: Sizes) !Scratch { - const head_rows = z.head_rows; - var s: Scratch = .{ .cap = cap, .head_rows = head_rows }; + var s: Scratch = .{ .cap = cap, .format = z.format, .head_rows = z.head_rows, .hc_part_bytes = hcPartRows(cap) * (low + 4) * 4 }; 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_bytes = hcPartRows(cap) * (low + 4) * 4; - s.hc_part = try s.add(gpa, d, s.hc_part_bytes); - s.hc_up = try s.add(gpa, d, cap * wide * 2); - if (z.lane_reduce) 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, partRows(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.lanes_tab = try s.add(gpa, d, lanes_bytes); - 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.picks_ep = 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, head_rows * vocab); + const lens = allocationSizes(cap, z); + // 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 (lens[i]) |len| try s.add(gpa, d, len) else 0; + if (lens[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)[5]); s.staging = .{ .q = .{ .stream = stream, .event = try cuda.Event.init(d, false) } }; - if (shared.on()) s.side = try shared.Side.init(d); // last: a refused allocation above never makes a stream + if (z.format == .classic and shared.on()) s.side = try shared.Side.init(d); // EXL3 runs its shared expert itself return s; } @@ -535,42 +544,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 1054083b32..af523cda86 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; @@ -214,8 +215,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 86014e4143..efec601d9b 100644 --- a/zig/src/families/flashnext/cuda_embed.zig +++ b/zig/src/families/flashnext/cuda_embed.zig @@ -80,6 +80,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 { @@ -108,7 +110,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 feb9a74057..231a513559 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 = snaps.bytes; + pub const snap_bytes = copies.bytes; pub const snap_repeat_min: u32 = 16; - pub const snap_save = snaps.save; - pub const snap_restore = snaps.restore; - pub const snap_drop = snaps.drop; + 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,9 @@ 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 +81,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, @@ -147,7 +153,7 @@ pub fn Namespace(comptime Tri: type) type { const seq_parts = @import("cuda_seq_parts.zig"); const Ops = mtp.Ops(Tri); - const snaps = @import("cuda_lanes_snap.zig").Snaps(Cuda); + const copies = @import("cuda_lanes_snap.zig").Copies(of); pub const Cuda = struct { gpa: std.mem.Allocator, @@ -156,6 +162,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. @@ -343,6 +351,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| { @@ -376,6 +387,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); @@ -390,6 +402,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); } @@ -572,21 +585,9 @@ 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, snaps.restore, dropLane); - _ = @import("cuda_lanes_admission_test.zig").Cases(Cuda, Cuda.open, 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, open); _ = @import("cuda_score.zig").CapabilityCases(Cuda); } - - test "a bad TF_FLASHNEXT_CHUNK, _CHUNK_SOLO or TF_FLASHNEXT_LOAD_THREADS is refused at load before the loader takes a stream, a module or a buffer" { - const env = @import("sys").env; - // every driver entry point undefined: the refusal must come before the first call - var driver: cuda.Driver = undefined; - const ctx: cuda.Context = .{ .d = &driver, .device = 0, .handle = null }; - for ([_]struct { [*:0]const u8, [*:0]const u8, anyerror }{ .{ "TF_FLASHNEXT_CHUNK", "0", error.BadChunk }, .{ "TF_FLASHNEXT_CHUNK", "16385", error.BadChunk }, .{ "TF_FLASHNEXT_CHUNK", "4k", error.BadChunk }, .{ "TF_FLASHNEXT_CHUNK_SOLO", "1024", error.BadChunk }, .{ "TF_FLASHNEXT_LOAD_THREADS", "0", error.BadLoadThreads }, .{ "TF_FLASHNEXT_LOAD_THREADS", "33", error.BadLoadThreads }, .{ "TF_FLASHNEXT_LOAD_THREADS", "8x", error.BadLoadThreads } }) |bad| { - _ = env.setenv(bad[0], bad[1], 1); - defer _ = env.unsetenv(bad[0]); - try std.testing.expectError(bad[2], open(std.testing.allocator, std.testing.io, &ctx, "no-such-checkpoint", null, .{ .context = 4096, .drafts = true })); - } - } }; } diff --git a/zig/src/families/flashnext/cuda_lanes_admission_test.zig b/zig/src/families/flashnext/cuda_lanes_admission_test.zig index 7666ad741c..a1f47d700e 100644 --- a/zig/src/families/flashnext/cuda_lanes_admission_test.zig +++ b/zig/src/families/flashnext/cuda_lanes_admission_test.zig @@ -58,7 +58,7 @@ fn run(comptime Cuda: type, comptime open: anytype, comptime drop: anytype, coun } } -pub fn Cases(comptime Cuda: type, comptime open: anytype, comptime drop: anytype) type { +pub fn Cases(comptime Cuda: type, comptime open: anytype, comptime drop: anytype, comptime load: anytype) type { return struct { test "four concurrent streams all retain MTP draft state" { try run(Cuda, open, drop, 4); @@ -67,5 +67,17 @@ pub fn Cases(comptime Cuda: type, comptime open: anytype, comptime drop: anytype test "eight concurrent streams all retain MTP draft state" { try run(Cuda, open, drop, 8); } + + test "a bad TF_FLASHNEXT_CHUNK, _CHUNK_SOLO or TF_FLASHNEXT_LOAD_THREADS is refused at load before the loader takes a stream, a module or a buffer" { + const env = @import("sys").env; + // every driver entry point undefined: the refusal must come before the first call + var driver: cuda.Driver = undefined; + const ctx: cuda.Context = .{ .d = &driver, .device = 0, .handle = null }; + for ([_]struct { [*:0]const u8, [*:0]const u8, anyerror }{ .{ "TF_FLASHNEXT_CHUNK", "0", error.BadChunk }, .{ "TF_FLASHNEXT_CHUNK", "16385", error.BadChunk }, .{ "TF_FLASHNEXT_CHUNK", "4k", error.BadChunk }, .{ "TF_FLASHNEXT_CHUNK_SOLO", "1024", error.BadChunk }, .{ "TF_FLASHNEXT_LOAD_THREADS", "0", error.BadLoadThreads }, .{ "TF_FLASHNEXT_LOAD_THREADS", "33", error.BadLoadThreads }, .{ "TF_FLASHNEXT_LOAD_THREADS", "8x", error.BadLoadThreads } }) |bad| { + _ = env.setenv(bad[0], bad[1], 1); + defer _ = env.unsetenv(bad[0]); + try std.testing.expectError(bad[2], load(std.testing.allocator, std.testing.io, &ctx, "no-such-checkpoint", null, .{ .context = 4096, .drafts = true })); + } + } }; } diff --git a/zig/src/families/flashnext/cuda_lanes_bytes.zig b/zig/src/families/flashnext/cuda_lanes_bytes.zig index ba1a163a96..77e70fa78d 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(headFp8(gpa, io, dir) catch return 0); } @@ -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 8a107c3674..6507a716d6 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; @@ -22,10 +24,14 @@ 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 (o.tp != 1 and mode != .modelopt) return error.FlashNextRanksNeedModelopt; if (o.tp != 1 and captured != null) return error.FlashNextRanksOnCapturedGlue; + // sm_120 only: refused before any weight loads + if (mode == .exl3) try @import("exl3/model/cuda_exl3_arch.zig").check(ctx.d); + try @import("exl3/kv/cuda_exl3_kv.zig").refuseSolo(mode == .exl3, std.c.getenv("TF_FLASHNEXT_CHUNK_SOLO") != null); 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; + // the pool's pages are not wired to EXL3's K/V + if (page_tokens > 0 and (o.tp != 1 or o.slide or captured != null or mode == .exl3)) return error.FlashNextPagesUnsupported; // pending DeltaNet rows run on one GPU only: TF_FLASHNEXT_KEEP_LATE=1 with tensor parallel is refused const Draft = @import("cuda_mtp_draft.zig").Draft; const late_fits = o.tp == 1; @@ -59,12 +65,12 @@ pub fn open(comptime F: type, gpa: std.mem.Allocator, io: std.Io, ctx: *const cu r.deinit(); }; own.engine = try prompt.Engine.initRank(gpa, io, driver, &own.stream, dir, o.tp, o.rank); + errdefer own.engine.deinit(); if (own.ranks) |r| { r.engine = own.engine; own.engine.collective = &r.collective; r.benchmark_group = try @import("core").benchmark_group.cuda(gpa, io, ctx, &r.collective, own.stream); } - errdefer own.engine.deinit(); const e = own.engine; if (e.chunk != e.busy_chunk) std.log.info("flashnext: prompt passes of {d} rows, {d} while another stream decodes", .{ e.chunk, e.busy_chunk }); if (o.vision) |limits| { @@ -73,11 +79,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); // 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; } @@ -92,18 +103,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); // 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, @@ -111,7 +142,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 index 9ce4e6d0ca..fffcdcbad0 100644 --- a/zig/src/families/flashnext/cuda_lanes_snap.zig +++ b/zig/src/families/flashnext/cuda_lanes_snap.zig @@ -1,13 +1,12 @@ -//! The prompt cache's copies of a Flash Next lane's state (cuda_snap.zig), called with the lane namespace's target. +//! 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"); -pub fn Snaps(comptime Cuda: type) type { +/// Size, save, restore and drop over the lane owner `of` finds behind a target pointer. +pub fn Copies(comptime of: anytype) type { return struct { - fn of(ptr: *anyopaque) *Cuda { - return @ptrCast(@alignCast(ptr)); - } - pub fn bytes(ptr: *anyopaque, at: u32) u64 { return snap.bytes(at, of(ptr).head != null); } @@ -30,12 +29,19 @@ pub fn Snaps(comptime Cuda: type) type { 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 = try self.lane(s); + 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))); } - try snap.restore(self.stream, &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| { + std.log.warn("flashnext: a kept state was not restored ({t})", .{err}); + return err; + }; } pub fn drop(ptr: *anyopaque, saved: *anyopaque) void { diff --git a/zig/src/families/flashnext/cuda_launch.zig b/zig/src/families/flashnext/cuda_launch.zig index 4301a1eacf..b53778e9ad 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"); @@ -24,8 +25,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 @@ -34,6 +38,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); @@ -452,6 +457,7 @@ test { _ = @import("cuda_kv_lanes.zig"); _ = @import("cuda_config.zig"); _ = @import("cuda_lanes_explain.zig"); + _ = @import("exl3/model/cuda_exl3.zig"); _ = @import("cuda_projection_load.zig"); _ = @import("cuda_rank_engine.zig"); _ = @import("cuda_round_graph.zig"); @@ -469,6 +475,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 094d61e575..660efd8075 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"); const ple_fp8 = @import("cuda_ple_fp8.zig"); @@ -332,7 +333,7 @@ pub fn eosOf(gpa: std.mem.Allocator, io: std.Io, model_dir: []const u8) !i64 { } /// A table's row format: MLX affine words, NVFP4 codes with E4M3 group scales, or FP8 rows with one table scale. -pub const Format = enum { mlx, nvfp4, fp8 }; +pub const Format = enum { mlx, nvfp4, fp8, exl3 }; /// The n-gram table held open for an engine's life: the hash, the scale, and each shard's slices of its mapped file. pub const Ngram = struct { @@ -349,6 +350,7 @@ pub const Ngram = struct { 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 { @@ -366,6 +368,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)); @@ -382,6 +385,7 @@ pub const Ngram = struct { .mlx => try tableScale(gpa, io, model_dir, wm), .nvfp4 => try readScale(gpa, io, model_dir, wm, "model.language_model.layers.1.ple.ple_embedding.ngram_embedding.weight_scale_2", true), .fp8 => try readScale(gpa, io, model_dir, wm, "model.language_model.layers.1.ple.ple_embedding.ngram_embedding.weight_scale", true), + .exl3 => unreachable, // an EXL3 table returns above }; const dh = g.hash.dh; var key: [180]u8 = undefined; @@ -406,6 +410,7 @@ pub const Ngram = struct { .mlx => .u32, .nvfp4 => .u8, .fp8 => .f8_e4m3, + .exl3 => unreachable, // an EXL3 table returns above }; const w = try file.tensor(try partNameFor(&key, shard, nested, "weight", modelopt), dtype); const expected_rows: usize = @intCast(g.hash.starts[shard + 1] - g.hash.starts[shard]); @@ -427,6 +432,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); @@ -451,6 +457,7 @@ pub const Ngram = struct { try g.tables8[loc.shard].gather(&.{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), } } } @@ -486,8 +493,9 @@ 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); - if (g.format != .mlx and items > 1) { + // 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 or g.format == .fp8) and items > 1) { try @import("cuda_ple_gather.zig").run(g, ids, out); } else if (n_threads <= 1 or g.pool == null) { try g.fill(ids, out, 0, items); @@ -503,7 +511,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]); } @@ -535,6 +543,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 1cacb9b18b..82a16897e3 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,22 +384,28 @@ 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; } - // head_local: the draft head's lanes on one GPU, a pass's rows split; native_part only where the lane reduces - const parts = @import("cuda_batch_state.zig"); - const z: parts.Sizes = .{ .head_rows = if (e.collective != null) @max(rows, parts.head_rows_one) else parts.head_rows_one, .lane_reduce = !e.batch_k.native_lane.clusters }; - e.batch_s = try batch.Scratch.alloc(e.gpa, e.driver, e.stream.*, rows, z); + e.batch_s = try batch.Scratch.alloc(e.gpa, e.driver, e.stream.*, rows, e.scratchSizes(rows)); return &e.batch_s.?; } + /// The scratch's sizes for `rows`: the format, head_local's rows and whether the lane reduces through native_part. + pub fn scratchSizes(e: *const Engine, rows: usize) @import("cuda_batch_state.zig").Sizes { + const one = @import("cuda_batch_state.zig").head_rows_one; + return .{ .format = e.scratch_format, .head_rows = if (e.collective != null) @max(rows, one) else one, .lane_reduce = !e.batch_k.native_lane.clusters }; + } + /// 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); @@ -425,7 +438,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| { @@ -508,23 +523,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); } }; @@ -532,55 +531,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..6a8dc5315d --- /dev/null +++ b/zig/src/families/flashnext/cuda_prompt_telemetry.zig @@ -0,0 +1,197 @@ +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); + } +} 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 423194ba60..69af087775 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; @@ -112,14 +113,19 @@ pub const Store = struct { /// Shards a load already opened, kept open until it ends (loaded or deinit): reopening one parses its header again. released: std.ArrayList(Held) = .empty, 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 }; } @@ -155,6 +161,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(); @@ -187,6 +194,8 @@ pub const Store = struct { /// Borrowers expire here; their shards stay open, pages dropped, until the load ends (or close if not kept). 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| { // the borrowers' pages leave this process (their bytes are copied), except the n-gram table's (Shard.close) if (!runtimePages(h.file.f.names.keys())) h.file.f.release(self.gpa); @@ -226,21 +235,21 @@ 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.items, 0..) |cached, i| { if (std.mem.eql(u8, cached.rel, file)) { self.held.appendAssumeCapacity(cached); _ = self.released.swapRemove(i); - 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); } } self.held.appendAssumeCapacity(try self.openShard(file)); - 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 { @@ -306,10 +315,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..56ef59dda2 --- /dev/null +++ b/zig/src/families/flashnext/cuda_scratch_layout_test.zig @@ -0,0 +1,53 @@ +const std = @import("std"); +const Scratch = @import("cuda_batch_state.zig").Scratch; + +const sizes = @import("cuda_scratch_sizes.zig"); +const part_index = 16; + +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, .{}); + const exl3 = Scratch.allocationSizes(cap, .{ .format = .exl3 }); + var removed: usize = 0; + var common: usize = 0; + for (classic, exl3, 0..) |before, after, index| { + if ((index >= 39 and index <= 48) or index == 62) { + try std.testing.expectEqual(@as(?usize, null), after); + removed += @max(before.?, 256); + } else if (index == part_index) { + // an EXL3 lane zero-fills every slice it is asked for: out_sk whole slices, not the classic zero slice + try std.testing.expectEqual(@as(?usize, sizes.partRowsFor(.exl3, cap) * 2560 * 4), after); + } else { + try std.testing.expectEqual(before, after); + common += @max(after.?, 256); + } + } + const exl3_part = @max(exl3[part_index].?, 256); + const classic_part = @max(classic[part_index].?, 256); + try std.testing.expectEqual(common + exl3_part, Scratch.bytes(cap, .{ .format = .exl3 })); + try std.testing.expectEqual(removed + classic_part, Scratch.bytes(cap, .{}) - common); + try std.testing.expectEqual(@as(?usize, cap * 11 * 2560 * 4), exl3[49]); + try std.testing.expectEqual(@as(?usize, cap * 2560 * 4), exl3[51]); + try std.testing.expectEqual(@as(?usize, cap * 128 * 4), exl3[52]); + try std.testing.expectEqual(@as(?usize, @import("cuda_batch_state.zig").lanes_bytes), exl3[65]); + } +} + +test "exl3: the classic MoE buffers it omits 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(sizes.head_rows_one * 248320, 256); + var omitted: usize = 0; + for (Scratch.allocationSizes(cap, .{}), 0..) |len, index| { + if ((index >= 39 and index <= 48) or index == 62) omitted += @max(len.?, 256); + } + try std.testing.expectEqual(expected, omitted); + try std.testing.expect(omitted > previous); + previous = omitted; + } +} diff --git a/zig/src/families/flashnext/cuda_scratch_sizes.zig b/zig/src/families/flashnext/cuda_scratch_sizes.zig index 2646140991..21f6cd55e4 100644 --- a/zig/src/families/flashnext/cuda_scratch_sizes.zig +++ b/zig/src/families/flashnext/cuda_scratch_sizes.zig @@ -7,8 +7,16 @@ const out_sk = wt.out_sk; /// head_local's rows on one GPU: the draft head's lanes (<= 16) times its ids; a split head needs a pass's rows. pub const head_rows_one: usize = 32; -/// head_local's rows, and whether the lane reduces K slices through native_part (with clusters it never does). -pub const Sizes = struct { head_rows: usize = head_rows_one, lane_reduce: bool = true }; +/// part's rows for `format`: an EXL3 lane zero-fills every slice it is asked for, so it keeps out_sk whole slices. +pub fn partRowsFor(format: ScratchFormat, cap: usize) usize { + return if (format == .exl3) out_sk * cap else partRows(cap); +} + +/// An EXL3 checkpoint's scratch leaves out the classic MoE buffers its own experts never touch. +pub const ScratchFormat = enum { classic, exl3 }; + +/// head_local's rows, whether the lane reduces K slices through native_part (never with clusters), and the format. +pub const Sizes = struct { head_rows: usize = head_rows_one, lane_reduce: bool = true, format: ScratchFormat = .classic }; /// hc_part's rows of (low + 4) floats: hc_sk slices below 64 rows (legacy kernel), <= 8 below 256 (lane), one beyond. pub fn hcPartRows(cap: usize) usize { 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..8f2850384b --- /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: 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; + +/// 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.bytes(e.chunk, e.scratchSizes(e.chunk)))) / 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_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..086f4683fd --- /dev/null +++ b/zig/src/families/flashnext/exl3/model/cuda_exl3_ple.zig @@ -0,0 +1,520 @@ +//! 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 Options = struct { read: Read = .map }; + +/// TF_FLASHNEXT_PLE_READ (map | pread); anything else keeps the default. +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 ---------------------------------------------------------------------------------------------- + +/// 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))); + } + } +} + +/// 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 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); + for (out, vals) |*o, v| o.* = @bitCast(v); + } + + /// Group sums of one decoded row. + pub fn sums(_: *const Table, row_bits: []const u16, out: []f32) void { + sumGroupsF16(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); +} + +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..1a3ee8a12b --- /dev/null +++ b/zig/src/families/flashnext/exl3/model/exl3_attn_cpu.zig @@ -0,0 +1,424 @@ +//! 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) + // tie goes up to even mantissa + try std.testing.expectEqual(@as(f64, 1.0 + 1.0 / 64.0), roundBf16(1.0 + 3.0 / 256.0)); + 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 "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..0b20cf2996 --- /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; +} + +/// A 16 GiB card. 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..d81fc0d24a --- /dev/null +++ b/zig/src/families/flashnext/exl3/offload/cuda_exl3_pool.zig @@ -0,0 +1,346 @@ +//! One layer's EXL3 routed-expert offload pool on upstream's `Experts`. + +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); + // in-flight prefetches lengthen rounds + if (opt.max_passes != 0 and opt.prefetch_bytes != 0) return error.Invalid; + 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..5774b8de78 --- /dev/null +++ b/zig/src/families/flashnext/exl3/pack/cuda_exl3_load.zig @@ -0,0 +1,532 @@ +//! 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; + // none in the text model + if (store.has(try std.fmt.bufPrint(&buf, "{s}.bias", .{prefix}))) return error.UnexpectedTensor; + 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); + } + // one release for the layer, not per expert: it drops a shard's mapped and cached pages and scans its names + 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); + } + } + // the layer's mapped pages are read once, so they are released together + 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..7e4ce455ec --- /dev/null +++ b/zig/src/families/flashnext/exl3/pack/cuda_exl3_memory.zig @@ -0,0 +1,135 @@ +//! 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; + // A_log, dt_bias as fp32 + if (!std.mem.endsWith(u8, name, ".weight") and e.rank == 1) return try std.math.mul(u64, e.dim(0), 4); + 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..14fb767219 --- /dev/null +++ b/zig/src/families/flashnext/exl3/tp/cuda_exl3_fork.zig @@ -0,0 +1,128 @@ +//! 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); + // zeroed once on the branch's stream: each block's last program resets its own + try dense.counters.?.fill8(0, side.handle); + 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 81405d3198..e6b8dd5186 100644 --- a/zig/src/native/cuda.zig +++ b/zig/src/native/cuda.zig @@ -415,7 +415,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/src/native/cuda_streams.zig b/zig/src/native/cuda_streams.zig index 7eb6de3111..7ebe48e2e0 100644 --- a/zig/src/native/cuda_streams.zig +++ b/zig/src/native/cuda_streams.zig @@ -34,7 +34,8 @@ pub const Count = struct { streams: u32, paged: ?kv_pages.Plan, layout: ?kv_pool /// The family's streams in `b.room`; null with `problem` set when none fit or --parallel asks for more than fit. pub fn count(comptime F: type, a: Allocator, loaded: anytype, o: api.Open, b: Budget, problem: *[]const u8) !?Count { const free = loaded.free_streams; - const grow = budget.growRoom(b.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(b.room, !preallocates_lanes); const wrapped = kv_pages.wrapped(F, loaded, free, b.stream_bytes, b.room); if (wrapped.layout) |l| { const paged = try kv_pages.plan(F, a, l, wrapped.room, o, @intCast(b.window), problem, b.host_cache, kv_pages.enginePaged(loaded), if (b.share) 1_000_000 else 0) orelse return null; diff --git a/zig/tests/cuda/contract_family_sets.zig b/zig/tests/cuda/contract_family_sets.zig index 84da81e3d4..b626289e33 100644 --- a/zig/tests/cuda/contract_family_sets.zig +++ b/zig/tests/cuda/contract_family_sets.zig @@ -1,4 +1,4 @@ -//! The contract's family sets: pocket, MiniMax-M3, MOSS-Image, Flash Next NVFP4, grouped affine-4, Flash Next, EXL3. +//! The contract's family sets: pocket, MiniMax-M3, MOSS-Image, Flash Next (NVFP4, EXL3), grouped affine-4, EXL3. const std = @import("std"); const cuda = @import("cuda"); const check = @import("check.zig"); @@ -15,6 +15,7 @@ pub const Sets = struct { grouped: *Grouped, flashnext: *@import("contract_flashnext_more.zig").Set, exl3: *@import("contract_exl3.zig").Set, + flashnext_exl3: *@import("contract_flashnext_exl3.zig").Set, /// Each set's kernels appended to `list`, in this order. pub fn init(gpu: check.Gpu, s: cuda.Stream, a: std.mem.Allocator, list: *std.ArrayList(g.Kernel), affine4: *cuda.qlinear.Affine4Kernels, cap: u32) !Sets { @@ -40,6 +41,7 @@ pub const Sets = struct { gpu.gpa.destroy(exl3); } try exl3.add(gpu, s, list); + const flashnext_exl3 = try @import("contract_flashnext_exl3.zig").Set.init(gpu, s, a, list); return .{ .gpa = gpu.gpa, .pocket = pocket, @@ -49,10 +51,12 @@ pub const Sets = struct { .grouped = grouped, .flashnext = flashnext, .exl3 = exl3, + .flashnext_exl3 = flashnext_exl3, }; } pub fn deinit(sets: Sets) void { + sets.flashnext_exl3.deinit(); sets.exl3.deinit(); sets.gpa.destroy(sets.exl3); sets.flashnext.deinit(); diff --git a/zig/tests/cuda/contract_flashnext_exl3.zig b/zig/tests/cuda/contract_flashnext_exl3.zig new file mode 100644 index 0000000000..98654b1ea3 --- /dev/null +++ b/zig/tests/cuda/contract_flashnext_exl3.zig @@ -0,0 +1,261 @@ +//! fn_exl3_split_f16 under the contract: an fp16 row as bf16 hi = bf16(x) and lo = bf16(x - hi). +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; +const bfr = core.bf16_round; + +const block = 256; +const patterns = 1 << 16; + +/// A planted fault: the control poses as the kernel and must go red at the gate it names. +pub const Control = enum { + /// A kernel that flushes NaN and the infinities to zero before it splits (gate 1). + flush, + /// A grid rounded down, so the partial last block is never launched (gate 3: whole blocks pass, every tail fails). + tail, + /// hi written over x, the kernel's own input (gate 2). + over, +}; + +fn none(_: c.Shape) usize { + return 0; +} + +/// Shape: rows x n fp16 values (k is 1), outputs hi then lo, no sums; the last two shapes end in a partial block. +pub fn spec(control: ?Control) c.Spec { + return .{ + .name = if (control) |x| switch (x) { + .flush => "control: flashnext_exl3 fn_exl3_split_f16 flushing NaN and Inf to zero", + .tail => "control: flashnext_exl3 fn_exl3_split_f16 launching whole blocks only", + .over => "control: flashnext_exl3 fn_exl3_split_f16 writing hi over its input", + } else "cuda flashnext_exl3 fn_exl3_split_f16 (fp16 row to bf16 hi + lo)", + .accum = .none, + .depth = none, + .ops = 1, // x - hi in fp32 (exact here: at most 3 bits remain) + .out = .bf16, + .nan = .propagate, // as float64 does on x - hi: NaN gives NaN twice, and +-Inf gives hi = +-Inf and lo = NaN + .tiles = .{ .{}, .{ .size = block, .step = 1, .min = 1 }, .{} }, + .served = &.{ + .{ .rows = 1, .n = 2560, .k = 1 }, + .{ .rows = 128, .n = 2560, .k = 1 }, + .{ .rows = 1, .n = 1000, .k = 1 }, + .{ .rows = 1, .n = patterns + 1, .k = 1 }, + }, + }; +} + +/// Planted at every shape's start and, reversed, its end: zeros, subnormal and normal limits, 1, Inf and NaNs. +pub const edges = [_]u16{ + 0x0000, 0x8000, 0x0001, 0x8001, 0x03ff, 0x83ff, 0x0400, 0x8400, 0x7bff, 0xfbff, 0x3c00, 0x3c01, + 0x7c00, 0xfc00, 0x7e00, 0xfe00, 0x7c01, 0xfdff, +}; + +/// Non-finite fill: NaN and both infinities in the first elements, the rest finite. +const planted = [_]u16{ 0x7e00, 0x7c00, 0xfc00, 0xfe01, 0x7c01 }; + +fn isFinite(bits: u16) bool { + return bits & 0x7c00 != 0x7c00; +} + +/// The odd stride makes i -> i * stride + start a bijection on 16 bits: n = 65536 holds every pattern once. +fn pattern(i: usize, start: u16) u16 { + return @truncate(i *% 40503 +% start); +} + +/// The float64 reference: hi = bf16(x), lo = bf16(x - hi), each rounded once in float64; NaN and Inf follow. +pub fn split(x: u16) [2]f64 { + const v: f64 = @as(f16, @bitCast(x)); + const hi = bfr.bf16(v); + return .{ hi, bfr.bf16(v - hi) }; +} + +pub const Split = struct { + rig: rig_.Rig, + f: cuda.Function, + control: ?Control = null, + arena: ?std.heap.ArenaAllocator = null, + total: usize = 0, + x: []u16 = &.{}, // the fp16 bits as the reference reads them + hi: u64 = 0, + lo: u64 = 0, + xd: u64 = 0, + msg: [160]u8 = undefined, + + pub fn deinit(self: *Split) void { + if (self.arena) |*a| a.deinit(); + self.rig.deinit(); + } + + const vt = g.hooks(Split); + + pub fn kernel(self: *Split) g.Kernel { + const gate: ?g.Gate = if (self.control) |x| switch (x) { + .flush => .nonfinite, + .tail => .shapes, + .over => .repeat, + } else null; + return .{ .spec = spec(self.control), .ctx = self, .vt = &vt, .control = gate }; + } + + /// Random fill spreads all 65,536 patterns (n = 65,536 holds each once) with the edges at both ends. + pub fn prepare(self: *Split, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows == 0 or shape.n == 0 or shape.k != 1) return false; + self.rig.reset(); + if (self.arena) |*a| a.deinit(); + self.arena = std.heap.ArenaAllocator.init(self.rig.gpu.gpa); + const a = self.arena.?.allocator(); + var prng = std.Random.DefaultPrng.init(seed); + const start = prng.random().int(u16); + self.total = shape.rows * shape.n; + self.x = try a.alloc(u16, self.total); + for (self.x, 0..) |*v, i| { + v.* = pattern(i, start); + if (fill == .nonfinite and !isFinite(v.*)) v.* &= ~@as(u16, 0x0400); + } + if (fill == .nonfinite) { + for (planted, 0..) |bits, i| if (i < self.total) { + self.x[i] = bits; + }; + } else for (edges, 0..) |bits, i| if (i < self.total) { + self.x[i] = bits; + self.x[self.total - 1 - i] = bits; + }; + const sent = try a.dupe(u16, self.x); + if (self.control == .flush) for (sent) |*v| { + if (!isFinite(v.*)) v.* = 0; + }; + self.xd = try self.rig.input(std.mem.sliceAsBytes(sent)); + self.hi = try self.rig.output(self.total * 2); + self.lo = try self.rig.output(self.total * 2); + return true; + } + + pub fn run(self: *Split) !void { + const blocks = if (self.control == .tail) self.total / block else (self.total + block - 1) / block; + if (blocks != 0) { + var args: cuda.Args = .{}; + args.add(self.xd); + args.add(if (self.control == .over) self.xd else self.hi); + args.add(self.lo); + args.add(@as(c_longlong, @intCast(self.total))); + try cuda.launch.launch(self.f, .{ .grid = .{ .x = @intCast(blocks) }, .block = .{ .x = block } }, self.rig.s, &args); + } + try self.rig.s.synchronize(); + } + + pub fn count(self: *Split) usize { + return 2 * self.total; + } + + /// hi then lo for every element; no sums, so no magnitude. + pub fn reference(self: *Split, want: []f64, mag: []f64) !void { + for (self.x, 0..) |x, i| { + const r = split(x); + want[i] = r[0]; + want[self.total + i] = r[1]; + } + @memset(mag, 0); + } + + /// Bit-exact at gates 1 and 3 (a half-ulp bound would let hi be one ulp off), any NaN counting, and hi + lo == x. + pub fn check(self: *Split, at: g.Gate) !?[]const u8 { + const out = try self.rig.outputBytes(); + var nonfinite: usize = 0; + for (self.x, 0..) |x, i| { + const got = [2]u16{ + std.mem.readInt(u16, out[2 * i ..][0..2], .little), + std.mem.readInt(u16, out[2 * (self.total + i) ..][0..2], .little), + }; + const want = split(x); + if (!isFinite(x)) nonfinite += 1; + for (got, want, [_][]const u8{ "hi", "lo" }) |bits, w, name| { + const ok = if (std.math.isNan(w)) std.math.isNan(bfr.value(bits)) else bits == bfr.bits(@floatCast(w)); + if (!ok) return std.fmt.bufPrint(&self.msg, "x {x:0>4} at {d}: {s} is {x:0>4} where float64 gives {e}", .{ x, i, name, bits, w }) catch "mismatch"; + } + const sum: f64 = @as(f16, @bitCast(x)); + if (isFinite(x) and @as(f64, bfr.value(got[0])) + @as(f64, bfr.value(got[1])) != sum) + return std.fmt.bufPrint(&self.msg, "x {x:0>4} at {d}: hi + lo is not x", .{ x, i }) catch "hi + lo differs from x"; + } + if (at == .nonfinite and (nonfinite == 0 or nonfinite == self.total)) return "the planted values make no proof"; + return null; + } +}; + +/// The module and the kernel's registrations: itself and its three controls. +pub const Set = struct { + a: std.mem.Allocator, + mod: cuda.Module, + cases: [4]Split, + + pub fn init(gpu: check.Gpu, s: cuda.Stream, a: std.mem.Allocator, kernels: *std.ArrayList(g.Kernel)) !*Set { + const self = try a.create(Set); + errdefer a.destroy(self); + self.a = a; + self.mod = try cuda.Module.load(gpu.d, cuda.kernels.flashnext_exl3); + errdefer self.mod.unload(); + const f = try self.mod.function("fn_exl3_split_f16"); + const controls = [_]?Control{ null, .flush, .tail, .over }; + for (&self.cases, controls) |*x, ctl| x.* = .{ .rig = .{ .gpu = gpu, .s = s }, .f = f, .control = ctl }; + for (&self.cases) |*x| try kernels.append(a, x.kernel()); + return self; + } + + pub fn deinit(self: *Set) void { + for (&self.cases) |*x| x.deinit(); + self.mod.unload(); + self.a.destroy(self); + } +}; + +test "the served shapes hold every fp16 pattern and a partial last block, and the bijection spreads small shapes" { + const t = std.testing; + const s = spec(null); + try t.expect(s.served[3].n >= patterns and s.served[3].n % block != 0 and s.served[2].n % block != 0); + var seen = try t.allocator.alloc(bool, patterns); + defer t.allocator.free(seen); + @memset(seen, false); + for (0..patterns) |i| seen[pattern(i, 0x1234)] = true; + for (seen) |v| try t.expect(v); + var spread: usize = 0; + for (0..255) |i| spread += @intFromBool(pattern(i, 7) & 0x7c00 != pattern(i + 1, 7) & 0x7c00); + try t.expect(spread > 200); +} + +test "the float64 reference splits every fp16 pattern: hi + lo is x when x is finite, Inf gives a NaN lo" { + const t = std.testing; + for (0..patterns) |i| { + const x: u16 = @intCast(i); + const r = split(x); + const v: f64 = @as(f16, @bitCast(x)); + // the kernel's fp32 path and the float64 reference round the same + const f: f32 = @floatCast(v); + const hi = bfr.bits(f); + const lo = bfr.bits(f - bfr.value(hi)); + if (std.math.isNan(v)) { + try t.expect(std.math.isNan(r[0]) and std.math.isNan(r[1]) and std.math.isNan(bfr.value(hi))); + } else if (std.math.isInf(v)) { + try t.expect(r[0] == v and std.math.isNan(r[1]) and std.math.isNan(bfr.value(lo))); + } else { + try t.expectEqual(v, r[0] + r[1]); + try t.expectEqual(bfr.bits(@floatCast(r[0])), hi); + try t.expectEqual(bfr.bits(@floatCast(r[1])), lo); + } + } +} + +test "the edges are the boundaries: zeros, subnormals, normals, the greatest finite, infinities, NaNs" { + const t = std.testing; + var finite: usize = 0; + for (edges) |x| { + if (isFinite(x)) finite += 1; + const r = split(x); + if (isFinite(x)) try t.expectEqual(@as(f64, @as(f16, @bitCast(x))), r[0] + r[1]); + } + try t.expectEqual(@as(usize, 12), finite); + try t.expect(planted.len <= 5 and !isFinite(planted[0]) and !isFinite(planted[4])); +} diff --git a/zig/tests/cuda/contract_limits.zig b/zig/tests/cuda/contract_limits.zig index 45091417d9..7b6c80bcb4 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 &.{}; 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..56fc02a44e --- /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; +} + +// ---- incremental row ids, a second form of the hash to check lastIds against ---------------------------------- + +/// 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 ---------------------------------------------------------------------------------- + +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}\n", .{ h.n, heads, dh, h.eos, h.shards, h.starts[h.shards], g.x3.?.bits, g.x3.?.rowBytes(), g.x3.?.opts.read }); + 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 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 incremental 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..6e7d343a9a --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/model/flashnext_exl3_engine_tests.zig @@ -0,0 +1,537 @@ +//! 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); + try flash.exl3.kv.setup(g.e, true); // the server's load hook: EXL3 scratch format, chunk and head rows + 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..b588fa53d7 --- /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: mixer 1% (GDN) or 3% (attention), streams 1%, MoE 2% and output 2% of the peak hold. + 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 1%\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..afb8f48546 --- /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 + 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..6f40415f77 --- /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 cdc4b0e2c3..ecb4a69e87 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -103,6 +103,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"); @@ -144,6 +150,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 @@ -237,6 +245,10 @@ const usage = \\ exl3-plans [shape] the 27B's EXL3 shapes: every (K splits, warps) and layout timed at 1 and 16 rows \\ exl3-contract EXL3's replacements of linear_body through gates 1-4, each control red (kernel-contract's subset) \\ 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) @@ -361,10 +373,12 @@ 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, "flashnext-attn-lanes")) return @import("contract_fn_attn_lanes.zig").run(gpu); if (std.mem.eql(u8, cmd, "flashnext-readout-cut")) return @import("flashnext_readout_cut.zig").run(gpu); - if (std.mem.eql(u8, cmd, "info")) return info(gpu); + if (std.mem.eql(u8, cmd, "info")) return runtime_tests.info(gpu); if (std.mem.eql(u8, cmd, "smoke")) return runtime_tests.smoke(gpu); if (std.mem.eql(u8, cmd, "shared-context")) return runtime_tests.sharedContext(gpu); if (std.mem.eql(u8, cmd, "graph")) return runtime_tests.graphs(gpu); @@ -418,6 +432,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-shapes")) return exl3_plan_bench.profileShapes(gpu, if (rest.len > 0) rest[0] else null); if (std.mem.eql(u8, cmd, "exl3-midm")) return exl3_plan_bench.midm(gpu); @@ -531,18 +549,6 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { return error.UnknownCommand; } -fn info(gpu: check.Gpu) !void { - var name_buf: [256]u8 = undefined; - const name = try gpu.ctx.name(&name_buf); - const mem = try gpu.ctx.memInfo(); - std.debug.print("RESULT driver CUDA {d}, device {s}, sm_{d}, {d} SMs, {d} MiB total, {d} MiB free, kernels embedded {}\n", .{ - try gpu.d.version(), name, - try gpu.ctx.capability(), try gpu.ctx.attribute(.multiprocessor_count), - mem.total >> 20, mem.free >> 20, - cuda.kernels.available, - }); -} - fn hasFlag(flags: []const [:0]const u8, name: []const u8) bool { for (flags) |f| if (std.mem.eql(u8, f, name)) return true; return false; diff --git a/zig/tests/cuda/runtime_tests.zig b/zig/tests/cuda/runtime_tests.zig index 53a843c812..9b092eb5fc 100644 --- a/zig/tests/cuda/runtime_tests.zig +++ b/zig/tests/cuda/runtime_tests.zig @@ -8,6 +8,18 @@ const expect = check.expect; const ProbeView = extern struct { src: u64, n: c_int, scale: f32 }; +pub fn info(gpu: Gpu) !void { + var name_buf: [256]u8 = undefined; + const name = try gpu.ctx.name(&name_buf); + const mem = try gpu.ctx.memInfo(); + std.debug.print("RESULT driver CUDA {d}, device {s}, sm_{d}, {d} SMs, {d} MiB total, {d} MiB free, kernels embedded {}\n", .{ + try gpu.d.version(), name, + try gpu.ctx.capability(), try gpu.ctx.attribute(.multiprocessor_count), + mem.total >> 20, mem.free >> 20, + cuda.kernels.available, + }); +} + pub fn smoke(gpu: Gpu) !void { const d = gpu.d; const n: usize = 1 << 20; From 36b8db618e09ddb3ed10b68c594688f11746325f 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 | 18 +- docs/multi-rank-receipts.md | 1 + docs/recipes/exl3.md | 3 +- docs/recipes/flash-next-cuda.md | 6 +- tools/exl3_serve.sh | 15 - tools/exl3_serve_tp2.sh | 16 + zig/build/cuda_kernels.zig | 2 + zig/kernels/cuda/contract_controls.cu | 3 + zig/kernels/cuda/contract_leaf_controls.cuh | 122 ++++ zig/kernels/cuda/flashnext_batch.cu | 2 + .../cuda/flashnext_exl3/flashnext_exl3_tp.cu | 9 + zig/kernels/cuda/flashnext_exl3_leaf.cuh | 72 +++ .../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_kernels.zig | 4 + zig/src/families/flashnext/cuda_batch_moe.zig | 35 +- .../families/flashnext/cuda_batch_state.zig | 4 +- .../families/flashnext/cuda_collective.zig | 20 + zig/src/families/flashnext/cuda_lanes.zig | 28 +- .../families/flashnext/cuda_lanes_bytes.zig | 2 +- .../families/flashnext/cuda_lanes_load.zig | 16 +- .../flashnext/cuda_lanes_restore_test.zig | 2 + .../families/flashnext/cuda_lanes_snap.zig | 31 +- zig/src/families/flashnext/cuda_launch.zig | 10 +- 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 | 6 +- .../families/flashnext/cuda_mtp_ops_test.zig | 75 +++ zig/src/families/flashnext/cuda_prompt.zig | 9 +- .../flashnext/cuda_prompt_telemetry.zig | 69 ++- .../families/flashnext/cuda_rank_engine.zig | 33 +- .../flashnext/cuda_scratch_layout_test.zig | 46 +- .../families/flashnext/cuda_scratch_sizes.zig | 10 +- .../flashnext/exl3/kv/cuda_exl3_kv.zig | 48 +- .../flashnext/exl3/model/cuda_exl3.zig | 3 + .../flashnext/exl3/model/cuda_exl3_arch.zig | 2 +- .../exl3/model/cuda_exl3_fake_driver_test.zig | 23 + .../flashnext/exl3/model/cuda_exl3_pass.zig | 72 ++- .../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 | 33 +- .../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 | 247 ++++++++ .../exl3/tp/exl3_canon_fp64_test.zig | 245 ++++++++ .../families/flashnext/exl3/tp/exl3_tp.zig | 148 +++++ zig/src/native/cuda.zig | 2 +- .../cuda_flashnext_exl3_refusal_test.zig | 19 + zig/src/native/cuda_refusal.zig | 33 +- .../cuda/contract_flashnext_exl3_canon.zig | 551 ++++++++++++++++++ .../cuda/contract_flashnext_exl3_suite.zig | 45 ++ zig/tests/cuda/contract_flashnext_exl3_tp.zig | 138 +++++ zig/tests/cuda/contract_limits.zig | 2 +- zig/tests/cuda/contract_tests.zig | 15 +- .../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 | 347 +++++++++++ .../tp/flashnext_exl3_tpe_decode_tests.zig | 212 +++++++ .../tp/flashnext_exl3_tpe_engine.zig | 215 +++++++ 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 + 81 files changed, 4489 insertions(+), 274 deletions(-) delete mode 100755 tools/exl3_serve.sh create mode 100755 tools/exl3_serve_tp2.sh create mode 100644 zig/kernels/cuda/contract_leaf_controls.cuh create mode 100644 zig/kernels/cuda/flashnext_exl3/flashnext_exl3_tp.cu create mode 100644 zig/kernels/cuda/flashnext_exl3_leaf.cuh create mode 100644 zig/src/families/flashnext/cuda_mtp_ops_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/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_canon_fp64_test.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_canon.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 15792dd44b..ce6c2338e7 100644 --- a/docs/kernel-contract.md +++ b/docs/kernel-contract.md @@ -299,6 +299,18 @@ CUDA also checks the affine-3 experts' split decode pair, bit-equal to the pair one dropping a last round of K groups short of the warps, one chaining the groups in reverse). CUDA also checks Flash Next's NVFP4 experts on prompt tiles with the decode form's sums (SwiGLU gate-up and fp32 down, each bit-equal to the decode instance it replaces; the scaled-code prompt tile is the planted control). +CUDA also checks Flash Next EXL3's canonical expert reduction (`flashnext_batch.cu`): `fn_exl3_leaf_tree` over leaves +0-31, 32-63 and all 64 (rows 1, 2, 5 and 16 at width 2560) and `fn_exl3_leaf_finish` over the pieces of 2, 3, 5 and 8 +ranks, each output byte-equal to the host reference `exl3_canon.zig` (NaN payloads aside), its class equal to float64's +and its value within float64's bound. Six controls, each one planted fault: the tree as one flat fma chain, NaN products +dropped, a write into the input, bf16 tree adds, and the finish as one flat chain and with bf16 merges. +Its order is TP-independent: 64 leaves of 8 expert ids (the shared expert in the last), an fp32 fma chain a leaf in +slot order from +0, then a fixed balanced fp32 tree, so every TP degree 1 to 8 gives the same bits. Against exact integer +sums of the products (`exl3_canon_fp64_test.zig`) the host tests check each output at |err| <= n u / (1 - n u) times the +sum of |y w| with n = 11 and u = 2^-24, not the contract's depth 17: a term meets at most c roundings in its leaf's chain +of c and one per tree level whose sibling holds terms, at most 11 - c, so n <= 11, as for the old one chain over the 11 +slots (the bound is the chain's, not a lower one). On router-like and wide-range rows the tree's worst output is no +worse than the old chain's or a rank-split sum's and its mean within a quarter of the best of them. CUDA also checks Nemotron's pocket (`nemotron_ops.cu`, the learned neurons the living-weights learner adds beside the shared expert): relu^2 up into act rows and down added into the shared expert's slot, a block row a row and tiled, the down to fp32 and bf16, against float64, the neurons past the count NaN so a read of one shows (two controls: the up's @@ -343,9 +355,9 @@ 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. -That launch, `fn_exl3_split_f16`, also runs gates 1 to 4 on the CUDA runner (`contract_flashnext_exl3.zig`): an fp16 +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. +The first, `fn_exl3_split_f16`, also runs gates 1 to 4 on the CUDA runner (`contract_flashnext_exl3.zig`): an fp16 row splits into bf16 hi = bf16(x) and lo = bf16(x - hi), every output's bits must equal float64's (NaN stays NaN; an infinity gives hi = Inf and lo = NaN), and the device's hi + lo must equal x for every finite pattern. Its shapes hold all 65,536 fp16 patterns and partial last blocks (n = 1,000 and 65,537). Three controls go red: one flushing NaN diff --git a/docs/multi-rank-receipts.md b/docs/multi-rank-receipts.md index 78e2b19bdc..ad6204c39f 100644 --- a/docs/multi-rank-receipts.md +++ b/docs/multi-rank-receipts.md @@ -7,6 +7,7 @@ Each path places its ranks its own way: | Path | Places ranks by | Documented by | | --- | --- | --- | | Flash Next on two Macs | `--speed-up FILE`, one file per rank | [Speed-up mode](speed-up-mode.md) | +| Flash Next EXL3 on CUDA | `--tp 2 --rank N --master HOST` | [Its CUDA recipe](recipes/flash-next-cuda.md) places the pack across two GPUs | | Native GLM-5.3-Flash on CUDA | `TF_GLM_RANK` of `TF_GLM_WORLD`, with `NCCL_COMM_ID` on rank 0's address | [Its recipe](recipes/glm-5.3-flash-cuda.md) starts two; [native placement and limits](../zig/src/families/glm/cuda_native.zig): the engine refuses a world past 16 | | Nemotron 3.5 Super on CUDA | `TF_NEMOTRON_RANK` of `TF_NEMOTRON_WORLD`, with `NCCL_COMM_ID` | [Its recipe](recipes/nemotron-3.5-super-vl.md) starts two; a world must divide the expert shares | | Legacy Python 0.6.6 engine | `--tp`, `--rank`, `--master` | [The Python GLM recipe](recipes/glm-5.3-flash.md) | diff --git a/docs/recipes/exl3.md b/docs/recipes/exl3.md index cd7a081015..7619bfcc3f 100644 --- a/docs/recipes/exl3.md +++ b/docs/recipes/exl3.md @@ -8,7 +8,8 @@ GLM-5.3-Flash. - The native engine has the format and a reference decoder in core (`zig/src/core/exl3_format.zig`, @akol1, #496), and Metal has a standalone EXL3 projection (`zig/src/core/exl3_projection.zig`, @akol1, #681). Flash Next serves an - EXL3 pack on CUDA (sm_120, one rank; see [its recipe](flash-next-cuda.md#exl3-checkpoints)); no other model does. + EXL3 pack on CUDA (sm_120, one rank or two; see [its recipe](flash-next-cuda.md#exl3-checkpoints)); no other model + does. - CUDA projections and routed experts on this format are in review (#740, #741, #806). A model needs them and a loader for its EXL3 tensors before it can serve a checkpoint; the 27B on CUDA is the first planned. - The Python engine served EXL3 on CUDA from 0.3.6 to 0.6.6 through the module @vcruz305 wrote (#42): every codebook, diff --git a/docs/recipes/flash-next-cuda.md b/docs/recipes/flash-next-cuda.md index 646c339309..b0351aae65 100644 --- a/docs/recipes/flash-next-cuda.md +++ b/docs/recipes/flash-next-cuda.md @@ -112,8 +112,9 @@ tensorfold-native serve /path/to/checkpoint --temperature 0 ## EXL3 checkpoints -An EXL3 pack (see [EXL3 weights](exl3.md)) serves on one sm_120 rank with its routed experts in pinned host memory and -a pool of expert lines a layer on the device. `tools/exl3_serve.sh` starts it. These switches tune it: +An EXL3 pack (see [EXL3 weights](exl3.md)) serves on one sm_120 rank, or on two with expert parallel (`--tp 2`), with +its routed experts in pinned host memory and a pool of expert lines a layer on the device. `tools/exl3_serve_tp2.sh` +starts a rank. These switches tune it: | Name | Default | Effect | | --- | --- | --- | @@ -125,6 +126,7 @@ a pool of expert lines a layer on the device. `tools/exl3_serve.sh` starts it. T | `TF_FLASHNEXT_EXL3_PROFILE_SIZING` | `0` | `1` sizes the pools from the stored profile; ignored with `TF_FLASHNEXT_EXPERT_LINES`. | | `TF_FLASHNEXT_EXL3_ROUTED` | host | `device` keeps every routed expert on the device. | | `TF_FLASHNEXT_EXL3_ARCH` | unset | `unqualified` serves on a device that is not sm_120, without a claim. | +| `TF_FLASHNEXT_EXL3_MTP_RESIDENT` | `0` | `1` keeps the MTP layer's experts on the device instead of a pool (one rank). | | `TF_FLASHNEXT_PLE_READ` | `map` | `pread` reads the n-gram embedding rows by positional read. | | `TF_FLASHNEXT_SNAP_HOST` | `1` above 32,768 tokens of context, else `0` | Whether kept prompt-cache states live in pinned host memory. | | `TF_FLASHNEXT_CACHE_TELEMETRY` | `0` | Diagnostics only: `1` logs `EXL3-CACHE` expert-cache counters per phase, with a sync at each boundary. | diff --git a/tools/exl3_serve.sh b/tools/exl3_serve.sh deleted file mode 100755 index 90a211fa0c..0000000000 --- a/tools/exl3_serve.sh +++ /dev/null @@ -1,15 +0,0 @@ -#!/bin/bash -# Serve the EXL3 Flash Next pack on one CUDA GPU with routed experts in pinned host memory: MODEL_DIR PORT [args]. -# Environment: LINES, RESERVE_MB, MEMORY_RESERVE_GIB, CONTEXT, PARALLEL, CACHE_GIB, TENSORFOLD_BIN. -set -euo pipefail -model=$1 port=$2 -shift 2 -bin=${TENSORFOLD_BIN:-./zig-out/native/bin/tensorfold-native} -if [ -z "${TENSORFOLD_CHECKPOINT_SHA256:-}" ]; then - TENSORFOLD_CHECKPOINT_SHA256=$(cat "$model/model.safetensors.index.json" "$model/config.json" | sha256sum | cut -c1-64) -fi -export TENSORFOLD_CHECKPOINT_SHA256 -export TF_FLASHNEXT_EXPERT_LINES=${LINES:-95} TF_FLASHNEXT_EXPERT_RESERVE_MB=${RESERVE_MB:-256} -export 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..a82b155596 --- /dev/null +++ b/tools/exl3_serve_tp2.sh @@ -0,0 +1,16 @@ +#!/bin/bash +# One TP=2 rank of the EXL3 Flash Next server: MODEL_DIR RANK PORT MASTER_PORT [serve args]; start rank 1 first. +# Environment: LINES (default 38), RESERVE_MB, MEMORY_RESERVE_GIB, CONTEXT, PARALLEL (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} +if [ -z "${TENSORFOLD_CHECKPOINT_SHA256:-}" ]; then + TENSORFOLD_CHECKPOINT_SHA256=$(cat "$model/model.safetensors.index.json" "$model/config.json" | sha256sum | cut -c1-64) +fi +export TENSORFOLD_CHECKPOINT_SHA256 +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 946cd3cd79..0bd612df8b 100644 --- a/zig/build/cuda_kernels.zig +++ b/zig/build/cuda_kernels.zig @@ -150,6 +150,8 @@ pub const kernels = [_]Kernel{ .{ .name = "drama3_codec", .flags = &.{"-O3"} }, // ours: Flash Next's EXL3 glue (an fp16 row split exactly into two bf16) .{ .name = "flashnext_exl3", .src = "flashnext_exl3/flashnext_exl3", .flags = glue }, + // ours: Flash Next's EXL3 expert-parallel pick mask + .{ .name = "flashnext_exl3_tp", .src = "flashnext_exl3/flashnext_exl3_tp", .flags = glue }, }; /// 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/contract_controls.cu b/zig/kernels/cuda/contract_controls.cu index d4b30e07d8..8fe55c7f3d 100644 --- a/zig/kernels/cuda/contract_controls.cu +++ b/zig/kernels/cuda/contract_controls.cu @@ -533,3 +533,6 @@ TF_CTL_EXL3_GROUPED(tf_ctl_exl3_grouped_old_2, 2, tf_exl3::PLAIN) TF_CTL_EXL3_GROUPED(tf_ctl_exl3_grouped_skip, 2, tf_exl3::SKIP) TF_CTL_EXL3_GROUPED(tf_ctl_exl3_grouped_stale, 2, tf_exl3::STALE) #undef TF_CTL_EXL3_GROUPED + +// Gates 1-4: the canonical expert reduction's leaf tree and finish, each with one planted fault. +#include "contract_leaf_controls.cuh" diff --git a/zig/kernels/cuda/contract_leaf_controls.cuh b/zig/kernels/cuda/contract_leaf_controls.cuh new file mode 100644 index 0000000000..2e20d3584a --- /dev/null +++ b/zig/kernels/cuda/contract_leaf_controls.cuh @@ -0,0 +1,122 @@ +// Planted faults for fn_exl3_leaf_tree and fn_exl3_leaf_finish: each control is the kernel with one change. +#pragma once +#include + +namespace tf_leaf_ctl { + +enum Fault { FLAT = 0, NAN_DROPPED = 1, WRITES_INPUT = 2, BF16_ADDS = 3 }; + +__device__ __forceinline__ float add(float a, float b, bool bf16) { + const float s = __fadd_rn(a, b); + return bf16 ? __bfloat162float(__float2bfloat16_rn(s)) : s; +} + +// fn_exl3_leaf_tree, then the fault: FLAT one chain (gate 3), NAN_DROPPED (1), WRITES_INPUT (2), BF16_ADDS (4). +template +__device__ __forceinline__ void leaf_tree(float* y, const float* wts, const int* pick, float* 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; + float chain = 0.0f; + 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; + const float yv = y[(static_cast(r) * slots + k) * d + col]; + const float wv = wts[r * slots + k]; + if (F == NAN_DROPPED) { + const float product = __fmul_rn(yv, wv); + if (product != product) continue; + } + chain = __fmaf_rn(yv, wv, chain); + 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(yv, wv, val[j]); + } + if (F == FLAT) { + out[static_cast(r) * d + col] = chain; + return; + } + 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] = add(val[q], val[q + 1], F == BF16_ADDS); + 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]; + if (F == WRITES_INPUT && r == 0 && col == 0) y[0] = __fadd_rn(y[0], 1.0f); +} + +// fn_exl3_leaf_finish, then the fault: FLAT adds every piece in one chain (gate 3), BF16_ADDS rounds merges (gate 4). +template +__device__ __forceinline__ void leaf_finish(const float* g, float* out, const int* 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; + float chain = 0.0f; + 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]; + chain = __fadd_rn(chain, sv[top]); + ++top; + while (top >= 2 && lg[top - 1] == lg[top - 2] && st[top - 2] % (2 << lg[top - 1]) == 0) { + sv[top - 2] = add(sv[top - 2], sv[top - 1], F == BF16_ADDS); + lg[top - 2] += 1; + --top; + } + } + out[static_cast(r) * d + col] = F == FLAT ? chain : sv[0]; +} + +} // namespace tf_leaf_ctl + +#define TF_CTL_LEAF_TREE(NAME, F) \ + extern "C" __global__ void NAME(float* y, const float* wts, const int* pick, float* out, int d, int slots, \ + int start, int log) { \ + tf_leaf_ctl::leaf_tree(y, wts, pick, out, d, slots, start, log); \ + } +TF_CTL_LEAF_TREE(tf_ctl_leaf_tree_flat, tf_leaf_ctl::FLAT) +TF_CTL_LEAF_TREE(tf_ctl_leaf_tree_nan_dropped, tf_leaf_ctl::NAN_DROPPED) +TF_CTL_LEAF_TREE(tf_ctl_leaf_tree_writes_input, tf_leaf_ctl::WRITES_INPUT) +TF_CTL_LEAF_TREE(tf_ctl_leaf_tree_bf16_adds, tf_leaf_ctl::BF16_ADDS) +#undef TF_CTL_LEAF_TREE + +#define TF_CTL_LEAF_FINISH(NAME, F) \ + extern "C" __global__ void NAME(const float* g, float* out, const int* meta, int pieces, int d, long long stride) { \ + tf_leaf_ctl::leaf_finish(g, out, meta, pieces, d, stride); \ + } +TF_CTL_LEAF_FINISH(tf_ctl_leaf_finish_flat, tf_leaf_ctl::FLAT) +TF_CTL_LEAF_FINISH(tf_ctl_leaf_finish_bf16_adds, tf_leaf_ctl::BF16_ADDS) +#undef TF_CTL_LEAF_FINISH diff --git a/zig/kernels/cuda/flashnext_batch.cu b/zig/kernels/cuda/flashnext_batch.cu index 6fec5e28cb..4b9f960977 100644 --- a/zig/kernels/cuda/flashnext_batch.cu +++ b/zig/kernels/cuda/flashnext_batch.cu @@ -519,6 +519,8 @@ extern "C" __global__ void fn_moe_sum(const float* __restrict__ y, const float* out[static_cast(r) * d + col] = acc; } +#include "flashnext_exl3_leaf.cuh" + // 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/kernels/cuda/flashnext_exl3_leaf.cuh b/zig/kernels/cuda/flashnext_exl3_leaf.cuh new file mode 100644 index 0000000000..f9d3846016 --- /dev/null +++ b/zig/kernels/cuda/flashnext_exl3_leaf.cuh @@ -0,0 +1,72 @@ +// 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]; +} 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 e86b427d30..e5e3f79f36 100644 --- a/zig/src/cuda/kernels.zig +++ b/zig/src/cuda/kernels.zig @@ -21,6 +21,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 71913a3cac..66f467a760 100644 --- a/zig/src/families/flashnext/cuda_batch.zig +++ b/zig/src/families/flashnext/cuda_batch.zig @@ -122,6 +122,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, /// This round's lanes take their attention in one launch a kernel (batch_lanes, TF_FLASHNEXT_ATTN_LANES). @@ -492,6 +494,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_kernels.zig b/zig/src/families/flashnext/cuda_batch_kernels.zig index a17d43c8a1..3ed1c1e1fd 100644 --- a/zig/src/families/flashnext/cuda_batch_kernels.zig +++ b/zig/src/families/flashnext/cuda_batch_kernels.zig @@ -40,6 +40,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, ep_picks: cuda.Function, shift: cuda.Function, @@ -83,6 +85,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.ep_picks = try k.batch.function("fn_ep_picks"); k.shift = try k.batch.function("fn_shift_window"); 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 c2b561b7cb..d525b7ed1f 100644 --- a/zig/src/families/flashnext/cuda_batch_state.zig +++ b/zig/src/families/flashnext/cuda_batch_state.zig @@ -391,7 +391,7 @@ pub const Scratch = struct { cap * dims * 2, cap * (dims / 32) * 4, cap * proj_row * 2, - partRowsFor(format, cap) * dims * 4, + partRowsFor(z, cap) * dims * 4, cap * gdn.value_dim * 2, cap * (gdn.value_dim / 32) * 4, cap * gdn.value_dim * 4, @@ -437,7 +437,7 @@ pub const Scratch = struct { cap * streams * 4, cap * wide * 2, vocab * 2, - if (format == .classic) z.head_rows * vocab else null, + z.head_rows * vocab, half * 4, 4096, lanes_bytes, diff --git a/zig/src/families/flashnext/cuda_collective.zig b/zig/src/families/flashnext/cuda_collective.zig index 3f150502da..70b52c916a 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(); } } @@ -116,6 +126,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 231a513559..25c7778b9f 100644 --- a/zig/src/families/flashnext/cuda_lanes.zig +++ b/zig/src/families/flashnext/cuda_lanes.zig @@ -7,8 +7,9 @@ const weights = @import("cuda_prompt_wt.zig"); const mtp = @import("cuda_mtp.zig"); 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 +28,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); @@ -124,7 +129,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; } @@ -176,6 +182,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 = .{}, @@ -319,20 +328,7 @@ pub fn Namespace(comptime Tri: type) type { } if (l.draft) |*d| d.media = if (l.seq.media) |*m| m.view() else null; - s.cached = 0; - 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) { - s.reuse_failed = true; - } else if (self.ranks != null and l.draft != null) { - 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)) |_| { - l.filled = kept.at; - s.cached = @intCast(kept.at); - } else |err| switch (err) { - error.SequenceTooShort, error.NoHeadState => s.reuse_failed = true, - else => return err, - } + try @import("cuda_lanes_snap.zig").reuse(self, l, s, ids); return l; } /// One prompt chunk, its rows picked as it starts, cut at the next cache mark; true once it was the last. diff --git a/zig/src/families/flashnext/cuda_lanes_bytes.zig b/zig/src/families/flashnext/cuda_lanes_bytes.zig index 77e70fa78d..9f754741d6 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 6507a716d6..939ba5cc72 100644 --- a/zig/src/families/flashnext/cuda_lanes_load.zig +++ b/zig/src/families/flashnext/cuda_lanes_load.zig @@ -22,7 +22,7 @@ pub fn open(comptime F: type, gpa: std.mem.Allocator, io: std.Io, ctx: *const cu 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 (o.tp != 1 and mode != .modelopt) return error.FlashNextRanksNeedModelopt; + if (o.tp != 1 and mode != .modelopt and mode != .exl3) return error.FlashNextRanksNeedModelopt; if (o.tp != 1 and captured != null) return error.FlashNextRanksOnCapturedGlue; // sm_120 only: refused before any weight loads if (mode == .exl3) try @import("exl3/model/cuda_exl3_arch.zig").check(ctx.d); @@ -59,6 +59,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, ranks.linkWait()); + // 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(); @@ -84,8 +86,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); @@ -106,9 +107,10 @@ pub fn open(comptime F: type, gpa: std.mem.Allocator, io: std.Io, ctx: *const cu 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; + // an EXL3 rank 1 keeps no head cache 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)); if (o.vision) |v| stream_bytes = try std.math.add(u64, stream_bytes, try std.math.add(u64, try std.math.mul(u64, v.image_tokens, @import("cuda_prompt_wt.zig").dims * 2), try std.math.mul(u64, 16, o.context))); if (page_tokens > 0) { // the lane's head runs at the engine's chunk: its cache on the stream's pages, its buffers in the run @@ -117,15 +119,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); // 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_restore_test.zig b/zig/src/families/flashnext/cuda_lanes_restore_test.zig index 92176e0dc7..abdf51c2f1 100644 --- a/zig/src/families/flashnext/cuda_lanes_restore_test.zig +++ b/zig/src/families/flashnext/cuda_lanes_restore_test.zig @@ -87,6 +87,8 @@ pub fn run(comptime Cuda: type, comptime open: anytype, comptime restore: anytyp self.stream = stream; self.ranks = &rank; self.head = if (drafting) &head else null; + self.exl3 = false; + self.kv = .{}; self.lanes = .empty; self.parts = .{}; defer self.lanes.deinit(a); diff --git a/zig/src/families/flashnext/cuda_lanes_snap.zig b/zig/src/families/flashnext/cuda_lanes_snap.zig index fffcdcbad0..53800febe3 100644 --- a/zig/src/families/flashnext/cuda_lanes_snap.zig +++ b/zig/src/families/flashnext/cuda_lanes_snap.zig @@ -18,8 +18,9 @@ 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; + // an EXL3 pack's state carries the head's rows on rank 0 + return try r.save(&l.seq, if (l.draft) |*d| d else null); } // pending DeltaNet rows into the states first: a snapshot holds them as they stand if (l.draft) |*d| try self.ops().flushLate(&l.seq, d); @@ -34,8 +35,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| { @@ -51,3 +52,25 @@ pub fn Copies(comptime of: anytype) type { } }; } + +/// Restore the stream's kept prompt state when its cache mark can be reused. +pub fn reuse(self: anytype, l: anytype, s: *lanes.Stream, ids: []const u32) !void { + s.cached = 0; + const saved = s.reuse.saved orelse return; + 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 and !self.exl3) { + s.reuse_failed = true; + } 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 => { + std.log.warn("flashnext: kept state at {d} not restored ({t})", .{ kept.at, err }); + s.reuse_failed = true; + }, + else => return err, + } +} diff --git a/zig/src/families/flashnext/cuda_launch.zig b/zig/src/families/flashnext/cuda_launch.zig index b53778e9ad..4fd2fef572 100644 --- a/zig/src/families/flashnext/cuda_launch.zig +++ b/zig/src/families/flashnext/cuda_launch.zig @@ -25,8 +25,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) +/// 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"); +/// The canonical expert reduction's host reference: the oracle of the contract test for its two kernels +pub const exl3_canon = @import("exl3/tp/exl3_canon.zig"); pub const mix = @import("cuda_mix.zig"); pub const prompt = @import("cuda_prompt.zig"); pub const snapshot = @import("cuda_snap.zig"); @@ -468,6 +471,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; @@ -475,7 +479,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 f1191cb48e..ae2257e94b 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 ec380c5060..97512aecef 100644 --- a/zig/src/families/flashnext/cuda_mtp_ops.zig +++ b/zig/src/families/flashnext/cuda_mtp_ops.zig @@ -103,7 +103,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, .prompt = true }); 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. @@ -251,10 +251,8 @@ pub fn Ops(comptime Tri: type) type { cap = @max(cap, d.cap); } if (n_l > max_lanes) return error.UnexpectedTensor; - // the head's plan: the admitted capacity bounds the pass's rows, so the head never grows the shared scratch const pl = plan.plan(e.admittedRows(), total); const s = try e.reserve(pl.pass_rows); - // the engine's block scores, priced to the piece (the single-stream head's rule) const score_rows = try e.reserveScores(cap, pl.score_rows); // the lanes in groups whose rows fit the piece var lane_rows: [max_lanes]usize = undefined; @@ -416,6 +414,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 82a16897e3..91120debf8 100644 --- a/zig/src/families/flashnext/cuda_prompt.zig +++ b/zig/src/families/flashnext/cuda_prompt.zig @@ -138,6 +138,10 @@ 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, + /// The decoder's projections are EXL3's whatever the scratch format (cuda_exl3_kv.setup). + exl3_lanes: bool = false, + /// 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 +176,8 @@ pub const Engine = struct { e.stream = stream; e.n_layers = 0; e.scratch_format = .classic; + e.exl3_lanes = false; + e.head_local_rows = std.math.maxInt(usize); e.collective = null; e.image_factory = null; e.telemetry = cache_telemetry.State.fromEnv(rank, world); @@ -398,7 +404,8 @@ pub const Engine = struct { /// The scratch's sizes for `rows`: the format, head_local's rows and whether the lane reduces through native_part. pub fn scratchSizes(e: *const Engine, rows: usize) @import("cuda_batch_state.zig").Sizes { const one = @import("cuda_batch_state.zig").head_rows_one; - return .{ .format = e.scratch_format, .head_rows = if (e.collective != null) @max(rows, one) else one, .lane_reduce = !e.batch_k.native_lane.clusters }; + const classic = if (e.collective != null) @max(rows, one) else one; + return .{ .exl3_lanes = e.exl3_lanes, .format = e.scratch_format, .head_rows = @min(rows, @min(classic, e.head_local_rows)), .lane_reduce = !e.batch_k.native_lane.clusters }; } /// The batched pass over `ids` from `seq.fed` in chunks of at most `chunk` rows; a window's hooks take one chunk. diff --git a/zig/src/families/flashnext/cuda_prompt_telemetry.zig b/zig/src/families/flashnext/cuda_prompt_telemetry.zig index 6a8dc5315d..87a05f827a 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 964bea3bec..1adaad17c0 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, @@ -200,7 +201,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); @@ -248,26 +249,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); } @@ -328,6 +335,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 => { @@ -373,6 +382,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); @@ -380,7 +390,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; @@ -419,7 +435,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 56ef59dda2..731d36b116 100644 --- a/zig/src/families/flashnext/cuda_scratch_layout_test.zig +++ b/zig/src/families/flashnext/cuda_scratch_layout_test.zig @@ -4,19 +4,19 @@ const Scratch = @import("cuda_batch_state.zig").Scratch; const sizes = @import("cuda_scratch_sizes.zig"); const part_index = 16; -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, .{}); const exl3 = Scratch.allocationSizes(cap, .{ .format = .exl3 }); var removed: usize = 0; var common: usize = 0; for (classic, exl3, 0..) |before, after, index| { - if ((index >= 39 and index <= 48) or index == 62) { + if (index >= 39 and index <= 48) { try std.testing.expectEqual(@as(?usize, null), after); removed += @max(before.?, 256); } else if (index == part_index) { // an EXL3 lane zero-fills every slice it is asked for: out_sk whole slices, not the classic zero slice - try std.testing.expectEqual(@as(?usize, sizes.partRowsFor(.exl3, cap) * 2560 * 4), after); + try std.testing.expectEqual(@as(?usize, sizes.partRowsFor(.{ .format = .exl3 }, cap) * 2560 * 4), after); } else { try std.testing.expectEqual(before, after); common += @max(after.?, 256); @@ -41,13 +41,49 @@ test "exl3: the classic MoE buffers it omits include the allocator minimum and g 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(sizes.head_rows_one * 248320, 256); + @max(cap * 2560 * 4, 256) + @max(pairs * 640 * 2, 256); var omitted: usize = 0; for (Scratch.allocationSizes(cap, .{}), 0..) |len, index| { - if ((index >= 39 and index <= 48) or index == 62) omitted += @max(len.?, 256); + if (index >= 39 and index <= 48) omitted += @max(len.?, 256); } try std.testing.expectEqual(expected, omitted); try std.testing.expect(omitted > previous); previous = omitted; } } + +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_rows = 32 }; + var split: Scratch = .{ .cap = 1024, .format = .exl3, .head_rows = 64 }; + var classic: Scratch = .{ .cap = 1536, .head_rows = 1536 }; + const small_sizes = Scratch.allocationSizes(small.cap, .{ .format = .exl3, .head_rows = small.head_rows }); + const split_sizes = Scratch.allocationSizes(split.cap, .{ .format = .exl3, .head_rows = split.head_rows }); + try std.testing.expectEqual(@as(?usize, 32 * vocab), small_sizes[62]); + try std.testing.expectEqual(@as(?usize, 64 * vocab), split_sizes[62]); + 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_rows = 16; + try std.testing.expectError(error.UnexpectedTensor, head.projectHead(&b, 0, 40)); + try std.testing.expectError(error.TensorParallelUnavailable, head.projectHead(&a, 0, 32)); +} diff --git a/zig/src/families/flashnext/cuda_scratch_sizes.zig b/zig/src/families/flashnext/cuda_scratch_sizes.zig index 21f6cd55e4..b6b33114fc 100644 --- a/zig/src/families/flashnext/cuda_scratch_sizes.zig +++ b/zig/src/families/flashnext/cuda_scratch_sizes.zig @@ -7,16 +7,16 @@ const out_sk = wt.out_sk; /// head_local's rows on one GPU: the draft head's lanes (<= 16) times its ids; a split head needs a pass's rows. pub const head_rows_one: usize = 32; -/// part's rows for `format`: an EXL3 lane zero-fills every slice it is asked for, so it keeps out_sk whole slices. -pub fn partRowsFor(format: ScratchFormat, cap: usize) usize { - return if (format == .exl3) out_sk * cap else partRows(cap); +/// part's rows: an EXL3 decoder's lane zero-fills every slice it is asked for, so it keeps out_sk whole slices. +pub fn partRowsFor(z: Sizes, cap: usize) usize { + return if (z.format == .exl3 or z.exl3_lanes) out_sk * cap else partRows(cap); } /// An EXL3 checkpoint's scratch leaves out the classic MoE buffers its own experts never touch. pub const ScratchFormat = enum { classic, exl3 }; -/// head_local's rows, whether the lane reduces K slices through native_part (never with clusters), and the format. -pub const Sizes = struct { head_rows: usize = head_rows_one, lane_reduce: bool = true, format: ScratchFormat = .classic }; +/// head_local's rows, whether the lane reduces through native_part, the format, whether the decoder's lanes are EXL3's. +pub const Sizes = struct { head_rows: usize = head_rows_one, lane_reduce: bool = true, format: ScratchFormat = .classic, exl3_lanes: bool = false }; /// hc_part's rows of (low + 4) floats: hc_sk slices below 64 rows (legacy kernel), <= 8 below 256 (lane), one beyond. pub fn hcPartRows(cap: usize) usize { 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 8f2850384b..9e6f9d7f1b 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; +} /// The pass's chunk for an EXL3 checkpoint unless TF_FLASHNEXT_CHUNK says. pub const exl3_chunk: usize = 1024; @@ -58,6 +65,23 @@ pub fn refuseSolo(exl3: bool, solo_set: bool) !void { if (exl3 and solo_set) return error.Exl3ChunkSoloUnsupported; } +/// The scratch format (classic MoE buffers stay unless the draft head is EXL3) and whether the decoder is EXL3. +pub fn scratchKind(decoder_exl3: bool, draft_exl3: bool) struct { format: @import("../../cuda_batch_state.zig").ScratchFormat, lanes: bool } { + return .{ .format = if (decoder_exl3 and draft_exl3) .exl3 else .classic, .lanes = decoder_exl3 }; +} + +test "exl3: an EXL3 decoder keeps whole part slices whatever its draft head is" { + const state = @import("../../cuda_batch_state.zig"); + for ([_]bool{ true, false }) |draft| { + const kind = scratchKind(true, draft); + for ([_]usize{ 1, 64, 1024, exl3_chunk }) |cap| { + const z: state.Sizes = .{ .format = kind.format, .exl3_lanes = kind.lanes }; + try std.testing.expectEqual(@as(?usize, wt.out_sk * cap * wt.dims * 4), state.Scratch.allocationSizes(cap, z)[16]); + } + } + try std.testing.expect(!scratchKind(false, true).lanes); +} + /// 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; @@ -65,8 +89,21 @@ 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"); + const kind = scratchKind(decoder_exl3, draft_exl3); + e.scratch_format = kind.format; + e.exl3_lanes = kind.lanes; 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.bytes(exl3_chunk, .{ .format = .exl3, .head_rows = 0 }); + const full = state.Scratch.bytes(exl3_chunk, .{ .format = .exl3, .head_rows = 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. @@ -92,3 +129,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_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 index 71de32b822..7a6be0aa9f 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,8 @@ 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 { + // expert parallel + if (x.tp != null) return @import("../tp/cuda_exl3_tp.zig").moe(p, lw, x); const s = p.s; const n = p.rows; try p.tri.routerMany(s.mixed, lw.router.ptr, s.router_l, dims, n); @@ -104,6 +106,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 +134,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 +203,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 +249,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 +277,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 +293,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 +307,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 +375,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 +388,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 d81fc0d24a..47540a2436 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; @@ -299,7 +299,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 5774b8de78..5197ce16b1 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; } @@ -313,14 +318,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; @@ -363,6 +375,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 = &.{}, @@ -447,9 +461,12 @@ 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 + // two-rank expert parallel + if (world != 1 or rank != 0) return @import("../tp/cuda_exl3_tp.zig").experts(gpa, d, store, mlp, world, rank); const count = wt.experts_n + 1; - const offload = offloaded(); + // TF_FLASHNEXT_EXL3_MTP_RESIDENT=1: the MTP layer's experts stay on the device 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 7e4ce455ec..06023601d1 100644 --- a/zig/src/families/flashnext/exl3/pack/cuda_exl3_memory.zig +++ b/zig/src/families/flashnext/exl3/pack/cuda_exl3_memory.zig @@ -80,15 +80,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). +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 { @@ -107,6 +143,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" { @@ -122,7 +159,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(); @@ -130,6 +167,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..5ea9a2e22a --- /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. + +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(); // once for the layer, not once an expert: it drops a whole shard's pages and scans its names + 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(); // the layer's mapped pages are read once, so they are released together + 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..e8a1e0f769 --- /dev/null +++ b/zig/src/families/flashnext/exl3/tp/exl3_canon.zig @@ -0,0 +1,247 @@ +//! 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 ------------------------------------------------------------------------------------------------------- + +pub 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. +pub 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); +} + +test { + _ = @import("exl3_canon_fp64_test.zig"); +} diff --git a/zig/src/families/flashnext/exl3/tp/exl3_canon_fp64_test.zig b/zig/src/families/flashnext/exl3/tp/exl3_canon_fp64_test.zig new file mode 100644 index 0000000000..3a7f74f9b3 --- /dev/null +++ b/zig/src/families/flashnext/exl3/tp/exl3_canon_fp64_test.zig @@ -0,0 +1,245 @@ +//! The canonical expert reduction against exact integer sums: the n u / (1 - n u) bound, other orders, TP independence. +const std = @import("std"); +const canon = @import("exl3_canon.zig"); + +const slots = canon.slots; +/// Every f32 product and sum below is a multiple of 2^-170 and under 2^30 in these tests: i256 holds them exactly. +const lsb: i32 = -170; + +const Parts = struct { m: u64, e: i32, neg: bool }; + +fn parts(x: f32) Parts { + const b: u32 = @bitCast(x); + const ex: i32 = @intCast((b >> 23) & 0xff); + const frac: u64 = b & 0x7fffff; + std.debug.assert(ex != 0xff); + const m: u64 = if (ex == 0) frac else frac | 0x800000; + // drop the mantissa's trailing zeros: the lowest set bit is what has to fit above 2^lsb + const z: u6 = if (m == 0) 0 else @intCast(@ctz(m)); + return .{ .m = m >> z, .e = (if (ex == 0) 1 else ex) - 150 + z, .neg = b >> 31 != 0 }; +} + +fn place(m: u64, e: i32, neg: bool) i256 { + std.debug.assert(e >= lsb and e - lsb < 200); + const v: i256 = @as(i256, m) << @intCast(e - lsb); + return if (neg) -v else v; +} + +fn exact(x: f32) i256 { + const p = parts(x); + return place(p.m, p.e, p.neg); +} + +/// y * w with no rounding (24 bit by 24 bit mantissas). +fn product(y: f32, w: f32) i256 { + const a = parts(y); + const b = parts(w); + return place(a.m * b.m, a.e + b.e, a.neg != b.neg); +} + +/// What one row's order produced against the exact sum of its products. +const Miss = struct { + /// |result - sum|, in 2^-170. + abs: u256, + /// the sum of |y w|, in 2^-170. + mag: u256, +}; + +fn miss(r: canon.Row, got: f32) Miss { + var sum: i256 = 0; + var mag: u256 = 0; + for (0..slots) |k| { + const t = product(r.y[k], r.w[k]); + sum += t; + mag += @abs(t); + } + return .{ .abs = @abs(exact(got) - sum), .mag = mag }; +} + +/// |err| <= n u / (1 - n u) mag with u = 2^-24, in integers: |err| (2^24 - n) <= n mag. +fn withinGamma(m: Miss, comptime n: u256) bool { + return (m.abs << 24) - m.abs * n <= m.mag * n; +} + +/// The leading 64 bits of v as a float (the rest is under 2^-63 of it; the error figures need no more). +fn approx(v: u256) f64 { + var sh: u8 = 0; + while (v >> sh >> 64 != 0) sh += 1; + return std.math.ldexp(@as(f64, @floatFromInt(@as(u64, @truncate(v >> sh)))), @as(i32, sh)); +} + +/// The error in units of u times the magnitude of the terms (the number of roundings that mattered). +fn units(m: Miss) f64 { + if (m.mag == 0) return 0; + return approx(m.abs) * 0x1p24 / approx(m.mag); +} + +// ---- the orders -------------------------------------------------------------------------------------------------- + +fn tree(r: canon.Row) f32 { + return canon.subtree(&r.y, &r.w, &r.p, 0, canon.leaf_log); +} + +/// The order before the tree: one fma chain over the slots, in slot order, from +0. +fn chain(r: canon.Row) f32 { + var acc: f32 = 0.0; + for (0..slots) |k| acc = @mulAdd(f32, r.y[k], r.w[k], acc); + return acc; +} + +/// A rank-split sum: each rank chains the slots whose leaf it owns; partials add in rank order (moves with the degree). +fn rankSplit(r: canon.Row, world: usize) f32 { + var acc: f32 = 0.0; + for (0..world) |rank| { + const own = canon.run(world, rank); + var part: f32 = 0.0; + for (0..slots) |k| { + const l = canon.leafOf(r.p[k]); + if (l >= own.lo and l < own.hi) part = @mulAdd(f32, r.y[k], r.w[k], part); + } + acc = if (rank == 0) part else acc + part; + } + return acc; +} + +// ---- data -------------------------------------------------------------------------------------------------------- + +const Kind = enum { + /// Ten distinct routed experts, softmax weights (sum 1), a sigmoid-gated shared expert, outputs of column scales. + router, + /// As `router` with output magnitudes over 2^-12 to 2^12: sums that cancel and absorb. + wide, + /// Picks on both sides of the leaf 31 / 32 boundary and in the shared leaf (the rank split of two). + boundary, + /// All 11 slots in one leaf (a chain of 11, the tree adds only zeros). + one_leaf, +}; + +fn row(rng: std.Random, kind: Kind) canon.Row { + var r: canon.Row = undefined; + var logits: [slots - 1]f32 = undefined; + var top: f32 = -std.math.inf(f32); + for (&logits) |*l| { + l.* = rng.floatNorm(f32); + top = @max(top, l.*); + } + var norm: f32 = 0.0; + for (&logits) |*l| { + l.* = @exp(l.* - top); + norm += l.*; + } + const leaf = rng.uintLessThan(u16, canon.leaves - 1); + for (0..slots) |k| { + const scale: f32 = switch (kind) { + .wide => std.math.exp2(@as(f32, @floatFromInt(rng.intRangeAtMost(i32, -12, 12)))), + else => std.math.exp2(-rng.float(f32) * 4.0) * 0.1, + }; + r.y[k] = rng.floatNorm(f32) * scale; + r.w[k] = if (k == slots - 1) 1.0 / (1.0 + @exp(-rng.floatNorm(f32))) else logits[k] / norm; + r.p[k] = switch (kind) { + .router, .wide => if (k == slots - 1) 512 else distinct(rng, r.p[0..k]), + .boundary => if (k == slots - 1) 512 else 248 + rng.uintLessThan(u16, 16), + .one_leaf => @as(u16, @intCast(leaf * canon.leaf_ids)) + rng.uintLessThan(u16, canon.leaf_ids), + }; + } + return r; +} + +fn distinct(rng: std.Random, taken: []const u16) u16 { + while (true) { + const id = rng.uintLessThan(u16, 512); + if (std.mem.indexOfScalar(u16, taken, id) == null) return id; + } +} + +const Stat = struct { + worst: f64 = 0.0, + sum: f64 = 0.0, + n: usize = 0, + + fn add(s: *Stat, m: Miss) !void { + // n = 11 roundings at most on any term, for the tree as for the chain and the rank splits + try std.testing.expect(withinGamma(m, slots)); + const u = units(m); + s.worst = @max(s.worst, u); + s.sum += u; + s.n += 1; + } + + fn mean(s: Stat) f64 { + return s.sum / @as(f64, @floatFromInt(s.n)); + } +}; + +fn bits(f: f32) u32 { + return @bitCast(f); +} + +const kinds = [_]Kind{ .router, .wide, .boundary, .one_leaf }; + +test "exl3: canon fp64 oracle sees what fp32 loses and nothing where fp32 is exact" { + var r: canon.Row = .{ .y = @splat(0.0), .w = @splat(1.0), .p = @splat(0) }; + // 2^24 + 1 is not an fp32 number: the chain rounds it to 2^24 (a miss of exactly 1) + r.y[0] = 0x1p24; + r.y[1] = 1.0; + const m = miss(r, chain(r)); + try std.testing.expectEqual(@as(u256, 1) << -lsb, m.abs); + try std.testing.expectEqual((@as(u256, 0x1p24 + 1)) << -lsb, m.mag); + try std.testing.expect(withinGamma(m, 1)); + try std.testing.expect(!withinGamma(.{ .abs = 2 * m.abs, .mag = m.mag }, 1)); + // integers whose sums stay under 2^24 are exact in every order + var prng = std.Random.DefaultPrng.init(0xf640); + const rng = prng.random(); + for (0..2000) |_| { + var q = row(rng, .router); + for (0..slots) |k| { + q.y[k] = @floatFromInt(rng.intRangeAtMost(i32, -4096, 4096)); + q.w[k] = @floatFromInt(rng.intRangeAtMost(i32, -64, 64)); + } + for ([_]f32{ tree(q), chain(q), rankSplit(q, 2), rankSplit(q, 8), canon.emulate(5, q) }) |got| { + try std.testing.expectEqual(@as(u256, 0), miss(q, got).abs); + } + } +} + +test "exl3: canon fp64 tree meets n u / (1 - n u) at n = 11 and is no worse than the slot chain or rank splits" { + var prng = std.Random.DefaultPrng.init(0xf64c); + const rng = prng.random(); + for (kinds) |kind| { + var t: Stat = .{}; + var c: Stat = .{}; + var s2: Stat = .{}; + var s8: Stat = .{}; + for (0..20000) |_| { + const r = row(rng, kind); + try t.add(miss(r, tree(r))); + try c.add(miss(r, chain(r))); + try s2.add(miss(r, rankSplit(r, 2))); + try s8.add(miss(r, rankSplit(r, 8))); + // a leaf's chain is the slot chain: with every slot in one leaf the tree adds only zeros + if (kind == .one_leaf) try std.testing.expectEqual(bits(chain(r)), bits(tree(r))); + } + // the worst output of the tree over these rows is at or under each other order's worst + try std.testing.expect(t.worst <= c.worst and t.worst <= s2.worst and t.worst <= s8.worst); + // and its mean is within a quarter of the best other order's + try std.testing.expect(t.mean() <= 1.25 * @min(c.mean(), @min(s2.mean(), s8.mean()))); + } +} + +test "exl3: canon fp64 tree bits are the same at TP 1 to 8 and as the halves of leaves 0-31 and 32-63" { + var prng = std.Random.DefaultPrng.init(0xf64d); + const rng = prng.random(); + var split_differs: usize = 0; + for (kinds) |kind| for (0..20000) |_| { + const r = row(rng, kind); + const whole = tree(r); + for (1..9) |world| try std.testing.expectEqual(bits(whole), bits(canon.emulate(world, r))); + const low = canon.subtree(&r.y, &r.w, &r.p, 0, canon.leaf_log - 1); + const high = canon.subtree(&r.y, &r.w, &r.p, canon.leaves / 2, canon.leaf_log - 1); + try std.testing.expectEqual(bits(whole), bits(low + high)); + // the rank-split sums are not that: their bits move with the degree, so the check above has teeth + const moved = bits(rankSplit(r, 2)) != bits(rankSplit(r, 1)) or bits(rankSplit(r, 8)) != bits(rankSplit(r, 3)); + split_differs += @intFromBool(moved); + }; + try std.testing.expect(split_differs > 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 e6b8dd5186..b1253768d5 100644 --- a/zig/src/native/cuda.zig +++ b/zig/src/native/cuda.zig @@ -206,7 +206,7 @@ pub fn open(a: Allocator, gpa: Allocator, io: std.Io, o: api.Open, problem: *[]c return null; } 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; } 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 72b58f9c8d..2d5ae47d2d 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)) { @@ -49,12 +50,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 no_ranks; 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 @@ -71,6 +83,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 }; @@ -99,6 +117,15 @@ test "the CUDA server serves --tp, and Flash Next's two ranks pass family admiss try std.testing.expect(family(flashnext, .{ .dir = "d", .model_type = flashnext.model_type, .tp = 2, .prompt_cache_gib = 0 }) == 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); @@ -259,3 +286,7 @@ test "--spill-gib and --checkpoint-slots open the families with a disk image; th try std.testing.expectStringStartsWith(backend(.{ .dir = "d", .model_type = "m", .keep = 2 }) orelse "", "--checkpoint-slots sets how many"); try std.testing.expect(backend(.{ .dir = "d", .model_type = "m" }) == null); } + +test { + _ = @import("cuda_flashnext_exl3_refusal_test.zig"); +} diff --git a/zig/tests/cuda/contract_flashnext_exl3_canon.zig b/zig/tests/cuda/contract_flashnext_exl3_canon.zig new file mode 100644 index 0000000000..5b3e0ad6e4 --- /dev/null +++ b/zig/tests/cuda/contract_flashnext_exl3_canon.zig @@ -0,0 +1,551 @@ +//! Flash Next EXL3's canonical expert reduction (leaf tree and finish) against the host exl3_canon.zig, bit for bit. +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 canon = flash.exl3_canon; + +const slots = canon.slots; + +/// The subtree a launch computes: the leaves a rank of two owns, or all 64 (one rank, the whole tree). +pub const Split = enum { low, high, full }; + +/// A planted fault of a control kernel (contract_leaf_controls.cuh); none is the shipped kernel. +pub const Fault = enum { none, flat, nan_dropped, writes_input, bf16_adds }; + +fn subtreeOf(s: Split) struct { start: usize, log: u8 } { + return switch (s) { + .low => .{ .start = 0, .log = 5 }, + .high => .{ .start = 32, .log = 5 }, + .full => .{ .start = 0, .log = canon.leaf_log }, + }; +} + +fn controlGate(f: Fault) ?g.Gate { + return switch (f) { + .none => null, + .flat => .shapes, + .nan_dropped => .nonfinite, + .writes_input => .repeat, + .bf16_adds => .sums, + }; +} + +/// A leaf's fma chain (at most 11 slots) and the 6 levels of the tree above it. +fn treeDepth(_: c.Shape) usize { + return slots + canon.leaf_log; +} + +/// Each piece meets at most 6 adds on its way up the tree. +fn finishDepth(_: c.Shape) usize { + return canon.leaf_log; +} + +fn treeName(s: Split, f: Fault) []const u8 { + return switch (f) { + .none => switch (s) { + .low => "cuda Flash Next EXL3 fn_exl3_leaf_tree (leaves 0-31)", + .high => "cuda Flash Next EXL3 fn_exl3_leaf_tree (leaves 32-63)", + .full => "cuda Flash Next EXL3 fn_exl3_leaf_tree (all 64 leaves)", + }, + .flat => "cuda Flash Next EXL3 fn_exl3_leaf_tree control: one chain, no tree", + .nan_dropped => "cuda Flash Next EXL3 fn_exl3_leaf_tree control: NaN products dropped", + .writes_input => "cuda Flash Next EXL3 fn_exl3_leaf_tree control: writes its input", + .bf16_adds => "cuda Flash Next EXL3 fn_exl3_leaf_tree control: bf16 tree adds", + }; +} + +fn finishName(f: Fault) []const u8 { + return switch (f) { + .none => "cuda Flash Next EXL3 fn_exl3_leaf_finish", + .flat => "cuda Flash Next EXL3 fn_exl3_leaf_finish control: one chain, no tree", + else => "cuda Flash Next EXL3 fn_exl3_leaf_finish control: bf16 merges", + }; +} + +const stated_rule = "the bytes of the host canonical tree (NaN payloads aside) and float64's class on every output"; + +/// Rows 1, 2, 5 and 16 at the hidden width 2560 (the engine's launches, dims / 256 column blocks); k is unused. +const tree_served = [_]c.Shape{ + .{ .rows = 5, .n = 2560, .k = 0 }, + .{ .rows = 1, .n = 2560, .k = 0 }, + .{ .rows = 2, .n = 2560, .k = 0 }, + .{ .rows = 16, .n = 2560, .k = 0 }, +}; + +fn treeSpec(s: Split, f: Fault) c.Spec { + return .{ + .name = treeName(s, f), + .accum = .f32, + .depth = treeDepth, + .out = .f32, + .nan = .{ .stated = stated_rule }, + // a block takes 256 columns: one block, two, and the whole width + .tiles = .{ .{}, .{ .size = 256, .step = 256, .min = 256 }, .{} }, + .served = &tree_served, + }; +} + +/// k is the TP degree whose pieces the finish merges (2 and 8 give one piece a rank, 3 and 5 give uneven runs). +const finish_served = [_]c.Shape{ + .{ .rows = 5, .n = 2560, .k = 8 }, + .{ .rows = 1, .n = 2560, .k = 2 }, + .{ .rows = 2, .n = 2560, .k = 3 }, + .{ .rows = 16, .n = 2560, .k = 5 }, +}; + +fn finishSpec(f: Fault) c.Spec { + return .{ + .name = finishName(f), + .accum = .f32, + .depth = finishDepth, + .out = .f32, + .nan = .{ .stated = stated_rule }, + .tiles = .{ .{}, .{ .size = 256, .step = 256, .min = 256 }, .{} }, + .served = &finish_served, + }; +} + +fn refused(shape: c.Shape) bool { + return shape.rows == 0 or shape.rows > 64 or shape.n == 0 or shape.n % 256 != 0 or shape.n > 2560; +} + +/// Bit-equal to the host's result, or both NaN (a payload is not part of the contract). +fn sameBits(got: f32, want: f32) bool { + return @as(u32, @bitCast(got)) == @as(u32, @bitCast(want)) or (std.math.isNan(got) and std.math.isNan(want)); +} + +fn bitsOf(out: []const u8, i: usize) f32 { + return @bitCast(std.mem.readInt(u32, out[i * 4 ..][0..4], .little)); +} + +/// Fp32 values that make fp16 sums overflow and bf16 ones round away their low bits: 2^19 to 2^20, either sign. +fn big(rng: std.Random) f32 { + const v: f32 = @floatFromInt((@as(u32, 1) << 19) + rng.uintLessThan(u32, 1 << 19)); + return if (rng.boolean()) v else -v; +} + +pub const Tree = struct { + rig: rig_.Rig, + f: cuda.Function, + split: Split, + fault: Fault, + arena: ?std.heap.ArenaAllocator = null, + shape: c.Shape = .{ .rows = 0, .n = 0, .k = 0 }, + y: []f32 = &.{}, + w: []f32 = &.{}, + p: []i32 = &.{}, + out: u64 = 0, + args: [3]u64 = .{ 0, 0, 0 }, + msg: [220]u8 = undefined, + + pub fn deinit(self: *Tree) void { + if (self.arena) |*a| a.deinit(); + self.rig.deinit(); + } + + const vt = g.hooks(Tree); + + pub fn kernel(self: *Tree) g.Kernel { + return .{ .spec = treeSpec(self.split, self.fault), .ctx = self, .vt = &vt, .control = controlGate(self.fault) }; + } + + fn inRange(self: *const Tree, id: i32) bool { + const s = subtreeOf(self.split); + const l = canon.leafOf(@intCast(id)); + return l >= s.start and l < s.start + (@as(usize, 1) << @intCast(s.log)); + } + + pub fn prepare(self: *Tree, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (refused(shape)) return false; + self.rig.reset(); + if (self.arena) |*a| a.deinit(); + self.arena = std.heap.ArenaAllocator.init(self.rig.gpu.gpa); + const a = self.arena.?.allocator(); + self.shape = shape; + self.y = try a.alloc(f32, shape.rows * slots * shape.n); + self.w = try a.alloc(f32, shape.rows * slots); + self.p = try a.alloc(i32, shape.rows * slots); + var prng = std.Random.DefaultPrng.init(seed ^ 0x6c656166); + const rng = prng.random(); + if (fill == .stress) self.stress(rng) else self.random(rng, fill); + const dy = try self.rig.input(std.mem.sliceAsBytes(self.y)); + const dw = try self.rig.input(std.mem.sliceAsBytes(self.w)); + const dp = try self.rig.input(std.mem.sliceAsBytes(self.p)); + self.out = try self.rig.output(self.count() * 4); + self.args = .{ dy, dw, dp }; + return true; + } + + /// Router-like picks (clustered leaves, the shared expert 512 last), small weights, a few negative zeros. + fn random(self: *Tree, rng: std.Random, fill: g.Fill) void { + const n = self.shape.n; + for (0..self.shape.rows) |r| { + const base = [2]u16{ rng.uintLessThan(u16, 64) * 8, rng.uintLessThan(u16, 64) * 8 }; + for (0..slots) |k| { + const id: u16 = if (k == slots - 1) 512 else if (rng.boolean()) rng.uintLessThan(u16, 512) else base[rng.uintLessThan(usize, 2)] + rng.uintLessThan(u16, 8); + self.p[r * slots + k] = id; + } + } + for (self.y) |*v| v.* = if (rng.uintLessThan(u8, 23) == 0) -0.0 else rng.floatNorm(f32) * 0.05; + for (self.w) |*v| v.* = if (rng.uintLessThan(u8, 29) == 0) -0.0 else rng.floatNorm(f32) * 0.2; + if (fill != .nonfinite) return; + // a row each of NaN, +Inf, -Inf, both infinities and Inf x 0: every 7th column, a slot a column + for (0..self.shape.rows) |r| for (0..n) |col| { + if (col % 7 != 0) continue; + const k = (col / 7) % slots; + const at = (r * slots + k) * n + col; + switch (r % 5) { + 0 => self.y[at] = std.math.nan(f32), + 1 => self.y[at] = std.math.inf(f32), + 2 => self.y[at] = -std.math.inf(f32), + 3 => { + self.y[at] = std.math.inf(f32); + self.y[(r * slots + (k + 1) % slots) * n + col] = -std.math.inf(f32); + }, + else => { + self.y[at] = std.math.inf(f32); + self.w[r * slots + k] = 0; + }, + } + }; + } + + /// Integers: picks inside the subtree, routed slots mirrored (same y, negated weight), the shared slot 3 x 1. + fn stress(self: *Tree, rng: std.Random) void { + const s = subtreeOf(self.split); + const lo: u16 = @intCast(s.start * canon.leaf_ids); + const span: u16 = @intCast(@min((@as(usize, 1) << @intCast(s.log)) * canon.leaf_ids, 512 - s.start * canon.leaf_ids)); + const n = self.shape.n; + for (0..self.shape.rows) |r| { + for (0..slots - 1) |k| self.p[r * slots + k] = lo + rng.uintLessThan(u16, span); + self.p[r * slots + slots - 1] = if (s.start + (@as(usize, 1) << @intCast(s.log)) == canon.leaves) 512 else lo + rng.uintLessThan(u16, span); + for (0..(slots - 1) / 2) |k| { + const wv: f32 = @floatFromInt(512 + rng.uintLessThan(u32, 512)); + self.w[r * slots + k] = wv; + self.w[r * slots + slots - 2 - k] = -wv; + for (0..n) |col| { + const yv: f32 = @floatFromInt(512 + rng.uintLessThan(u32, 512)); + self.y[(r * slots + k) * n + col] = yv; + self.y[(r * slots + slots - 2 - k) * n + col] = yv; + } + } + self.w[r * slots + slots - 1] = 1; + for (0..n) |col| self.y[(r * slots + slots - 1) * n + col] = 3; + } + } + + pub fn run(self: *Tree) !void { + const s = subtreeOf(self.split); + var a: cuda.Args = .{}; + for (self.args) |ptr| a.add(ptr); + a.add(self.out); + a.add(@as(c_int, @intCast(self.shape.n))); + a.add(@as(c_int, slots)); + a.add(@as(c_int, @intCast(s.start))); + a.add(@as(c_int, s.log)); + try cuda.launch.launch(self.f, .{ .grid = .{ .x = @intCast(self.shape.rows), .y = @intCast(self.shape.n / 256) }, .block = .{ .x = 256 } }, self.rig.s, &a); + try self.rig.s.synchronize(); + } + + pub fn count(self: *Tree) usize { + return self.shape.rows * self.shape.n; + } + + /// The host canonical subtree of one output (exl3_canon.subtree, the fp32 reference of this kernel). + fn expected(self: *const Tree, i: usize) f32 { + const r = i / self.shape.n; + const col = i % self.shape.n; + var y: [slots]f32 = undefined; + var w: [slots]f32 = undefined; + var p: [slots]u16 = undefined; + for (0..slots) |k| { + y[k] = self.y[(r * slots + k) * self.shape.n + col]; + w[k] = self.w[r * slots + k]; + p[k] = @intCast(self.p[r * slots + k]); + } + const s = subtreeOf(self.split); + return canon.subtree(&y, &w, &p, s.start, s.log); + } + + /// Float64's sum of the products of the slots whose leaf the subtree holds, and the sum of their magnitudes. + fn terms(self: *const Tree, i: usize, buf: *[slots]f64) usize { + const r = i / self.shape.n; + const col = i % self.shape.n; + var m: usize = 0; + for (0..slots) |k| { + if (!self.inRange(self.p[r * slots + k])) continue; + buf[m] = @as(f64, self.y[(r * slots + k) * self.shape.n + col]) * @as(f64, self.w[r * slots + k]); + m += 1; + } + return m; + } + + pub fn reference(self: *Tree, want: []f64, mag: []f64) !void { + var t: [slots]f64 = undefined; + for (want, mag, 0..) |*w, *m, i| { + const cnt = self.terms(i, &t); + w.* = 0; + m.* = 0; + for (t[0..cnt]) |x| { + w.* += x; + m.* += @abs(x); + } + } + } + + pub fn simulate(self: *Tree, acc: c.Accum, order: g.Order, out: []f64) !void { + var t: [slots]f64 = undefined; + for (out, 0..) |*o, i| o.* = core.contract.rows.sumIn(acc, order, t[0..self.terms(i, &t)]); + } + + pub fn check(self: *Tree, gate: g.Gate) !?[]const u8 { + const a = self.arena.?.allocator(); + const want = try a.alloc(f64, self.count()); + const mag = try a.alloc(f64, self.count()); + try self.reference(want, mag); + const bytes = [_][]const u8{ std.mem.sliceAsBytes(self.y), std.mem.sliceAsBytes(self.w), std.mem.sliceAsBytes(self.p) }; + return verify(&self.rig, &self.msg, treeSpec(self.split, self.fault), self.shape, gate, &bytes, want, mag, self, Tree.expected); + } +}; + +/// Guards, inputs, then each output against the host tree, float64's class and bound (gate 1: finite and not). +fn verify(rig: *rig_.Rig, msg: []u8, s: c.Spec, shape: c.Shape, gate: g.Gate, inputs: []const []const u8, want: []const f64, mag: []const f64, ctx: anytype, comptime expected: anytype) !?[]const u8 { + if (!try rig.intact()) return "a payload guard changed"; + const seen = try rig.inputBytes(); + var at: usize = 0; + for (inputs) |bytes| { + if (!std.mem.eql(u8, seen[at..][0..bytes.len], bytes)) return "an immutable input changed"; + at += bytes.len; + } + const out = try rig.outputBytes(); + var finite: usize = 0; + for (want, mag, 0..) |w, m, i| { + const got = bitsOf(out, i); + const host = expected(ctx, i); + if (!sameBits(got, host)) return try std.fmt.bufPrint(msg, "output {d}: {x:0>8} where the host canonical reduction gives {x:0>8}", .{ i, @as(u32, @bitCast(got)), @as(u32, @bitCast(host)) }); + const have: f64 = got; + if (c.classOf(have) != c.classOf(w)) return try std.fmt.bufPrint(msg, "output {d} is {t} where float64 gives {t}", .{ i, c.classOf(have), c.classOf(w) }); + if (std.math.isFinite(w)) { + finite += 1; + if (@abs(have - w) > c.bound(s, shape, w, m)) return try std.fmt.bufPrint(msg, "output {d}: {d} is outside float64's bound around {d}", .{ i, have, w }); + } + } + if (gate == .nonfinite and (finite == 0 or finite == want.len)) return "the planted values leave all outputs finite or none"; + return null; +} + +pub const Finish = struct { + rig: rig_.Rig, + f: cuda.Function, + fault: Fault, + arena: ?std.heap.ArenaAllocator = null, + shape: c.Shape = .{ .rows = 0, .n = 0, .k = 0 }, + ps: [canon.max_pieces]canon.Piece = undefined, + np: usize = 0, + /// Piece i's rows are stored at slot np - 1 - i: the meta's source index is exercised, not the identity. + vals: []f32 = &.{}, + meta: []i32 = &.{}, + out: u64 = 0, + args: [2]u64 = .{ 0, 0 }, + msg: [220]u8 = undefined, + + pub fn deinit(self: *Finish) void { + if (self.arena) |*a| a.deinit(); + self.rig.deinit(); + } + + const vt = g.hooks(Finish); + + pub fn kernel(self: *Finish) g.Kernel { + return .{ .spec = finishSpec(self.fault), .ctx = self, .vt = &vt, .control = controlGate(self.fault) }; + } + + fn at(self: *const Finish, piece: usize, i: usize) usize { + return (self.np - 1 - piece) * self.shape.rows * self.shape.n + i; + } + + pub fn prepare(self: *Finish, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (refused(shape) or shape.k < 1 or shape.k > 8) return false; + self.rig.reset(); + if (self.arena) |*a| a.deinit(); + self.arena = std.heap.ArenaAllocator.init(self.rig.gpu.gpa); + const a = self.arena.?.allocator(); + self.shape = shape; + self.np = canon.pieces(&self.ps, shape.k); + const per = shape.rows * shape.n; + self.vals = try a.alloc(f32, self.np * per); + self.meta = try a.alloc(i32, 3 * self.np); + for (self.ps[0..self.np], 0..) |p, i| self.meta[3 * i ..][0..3].* = .{ p.start, p.log, @intCast(self.np - 1 - i) }; + var prng = std.Random.DefaultPrng.init(seed ^ 0x66696e69); + const rng = prng.random(); + switch (fill) { + .stress => self.stress(rng), + else => self.random(rng, fill), + } + const dg = try self.rig.input(std.mem.sliceAsBytes(self.vals)); + const dmeta = try self.rig.input(std.mem.sliceAsBytes(self.meta)); + self.args = .{ dg, dmeta }; + self.out = try self.rig.output(per * 4); + return true; + } + + /// Partials of different magnitudes (an fp32 sum of them depends on the order of the merges), a few negative zeros. + fn random(self: *Finish, rng: std.Random, fill: g.Fill) void { + for (self.vals) |*v| { + const scale: f32 = @floatFromInt(@as(u32, 1) << rng.uintLessThan(u5, 10)); + v.* = if (rng.uintLessThan(u8, 23) == 0) -0.0 else rng.floatNorm(f32) * scale; + } + if (fill != .nonfinite) return; + const n = self.shape.n; + for (0..self.shape.rows) |r| for (0..n) |col| { + if (col % 7 != 0) continue; + const piece = (col / 7) % self.np; + const i = r * n + col; + switch (r % 5) { + 0 => self.vals[self.at(piece, i)] = std.math.nan(f32), + 1 => self.vals[self.at(piece, i)] = std.math.inf(f32), + 2 => self.vals[self.at(piece, i)] = -std.math.inf(f32), + 3 => { + self.vals[self.at(piece, i)] = std.math.inf(f32); + self.vals[self.at((piece + 1) % self.np, i)] = -std.math.inf(f32); + }, + else => self.vals[self.at(self.np - 1, i)] = std.math.nan(f32), + } + }; + } + + /// Integer partials of 2^19 to 2^20 in shuffled negated pairs (odd one out 3): exact in fp32, ruinous in fp16/bf16. + fn stress(self: *Finish, rng: std.Random) void { + for (0..self.shape.rows * self.shape.n) |i| { + var v: [canon.max_pieces]f32 = undefined; + for (0..self.np / 2) |j| { + v[2 * j] = big(rng); + v[2 * j + 1] = -v[2 * j]; + } + if (self.np % 2 == 1) v[self.np - 1] = 3; + // a mirrored order would cancel its own roundings (halves and lanes add pairs symmetrically): break it + rng.shuffle(f32, v[0..self.np]); + for (v[0..self.np], 0..) |x, piece| self.vals[self.at(piece, i)] = x; + } + } + + pub fn run(self: *Finish) !void { + var a: cuda.Args = .{}; + a.add(self.args[0]); + a.add(self.out); + a.add(self.args[1]); + a.add(@as(c_int, @intCast(self.np))); + a.add(@as(c_int, @intCast(self.shape.n))); + a.add(@as(c_longlong, @intCast(self.shape.rows * self.shape.n))); + try cuda.launch.launch(self.f, .{ .grid = .{ .x = @intCast(self.shape.rows), .y = @intCast(self.shape.n / 256) }, .block = .{ .x = 256 } }, self.rig.s, &a); + try self.rig.s.synchronize(); + } + + pub fn count(self: *Finish) usize { + return self.shape.rows * self.shape.n; + } + + /// The pieces of one output, in the order the finish merges them. + fn pieceValues(self: *const Finish, i: usize, buf: *[canon.max_pieces]f32) []const f32 { + for (0..self.np) |piece| buf[piece] = self.vals[self.at(piece, i)]; + return buf[0..self.np]; + } + + fn expected(self: *const Finish, i: usize) f32 { + var buf: [canon.max_pieces]f32 = undefined; + return canon.finish(self.ps[0..self.np], self.pieceValues(i, &buf)); + } + + pub fn reference(self: *Finish, want: []f64, mag: []f64) !void { + var buf: [canon.max_pieces]f32 = undefined; + for (want, mag, 0..) |*w, *m, i| { + w.* = 0; + m.* = 0; + for (self.pieceValues(i, &buf)) |v| { + w.* += v; + m.* += @abs(@as(f64, v)); + } + } + } + + pub fn simulate(self: *Finish, acc: c.Accum, order: g.Order, out: []f64) !void { + var buf: [canon.max_pieces]f32 = undefined; + var t: [canon.max_pieces]f64 = undefined; + for (out, 0..) |*o, i| { + const v = self.pieceValues(i, &buf); + for (v, t[0..v.len]) |x, *y| y.* = x; + o.* = core.contract.rows.sumIn(acc, order, t[0..v.len]); + } + } + + pub fn check(self: *Finish, gate: g.Gate) !?[]const u8 { + const a = self.arena.?.allocator(); + const want = try a.alloc(f64, self.count()); + const mag = try a.alloc(f64, self.count()); + try self.reference(want, mag); + const bytes = [_][]const u8{ std.mem.sliceAsBytes(self.vals), std.mem.sliceAsBytes(self.meta) }; + return verify(&self.rig, &self.msg, finishSpec(self.fault), self.shape, gate, &bytes, want, mag, self, Finish.expected); + } +}; + +/// Both kernels, each fault of each as a control, over the flashnext_batch image and the controls image. +pub const Canons = struct { + mod: cuda.Module, + trees: [7]Tree, + finishes: [3]Finish, + + const tree_plan = [7]struct { Split, Fault }{ + .{ .low, .none }, .{ .high, .none }, .{ .full, .none }, .{ .full, .flat }, + .{ .full, .nan_dropped }, .{ .full, .writes_input }, .{ .full, .bf16_adds }, + }; + const finish_plan = [3]Fault{ .none, .flat, .bf16_adds }; + + fn treeSymbol(f: Fault) [:0]const u8 { + return switch (f) { + .none => "fn_exl3_leaf_tree", + .flat => "tf_ctl_leaf_tree_flat", + .nan_dropped => "tf_ctl_leaf_tree_nan_dropped", + .writes_input => "tf_ctl_leaf_tree_writes_input", + .bf16_adds => "tf_ctl_leaf_tree_bf16_adds", + }; + } + + fn finishSymbol(f: Fault) [:0]const u8 { + return switch (f) { + .none => "fn_exl3_leaf_finish", + .flat => "tf_ctl_leaf_finish_flat", + else => "tf_ctl_leaf_finish_bf16_adds", + }; + } + + pub fn init(gpu: check.Gpu, s: cuda.Stream, controls: cuda.Module) !Canons { + var mod = try cuda.Module.load(gpu.d, cuda.kernels.flashnext_batch); + errdefer mod.unload(); + var self: Canons = .{ .mod = mod, .trees = undefined, .finishes = undefined }; + for (tree_plan, &self.trees) |plan, *t| { + const f = try (if (plan[1] == .none) mod else controls).function(treeSymbol(plan[1])); + t.* = .{ .rig = .{ .gpu = gpu, .s = s }, .f = f, .split = plan[0], .fault = plan[1] }; + } + for (finish_plan, &self.finishes) |fault, *x| { + const f = try (if (fault == .none) mod else controls).function(finishSymbol(fault)); + x.* = .{ .rig = .{ .gpu = gpu, .s = s }, .f = f, .fault = fault }; + } + return self; + } + + pub fn deinit(self: *Canons) void { + for (&self.trees) |*t| t.deinit(); + for (&self.finishes) |*x| x.deinit(); + self.mod.unload(); + } + + pub fn add(self: *Canons, a: std.mem.Allocator, list: *std.ArrayList(g.Kernel)) !void { + for (&self.trees) |*t| try list.append(a, t.kernel()); + for (&self.finishes) |*x| try list.append(a, x.kernel()); + } +}; 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..a3e63f11dd --- /dev/null +++ b/zig/tests/cuda/contract_flashnext_exl3_suite.zig @@ -0,0 +1,45 @@ +//! 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"); +const canon_ = @import("contract_flashnext_exl3_canon.zig"); +const writeback = @import("contract_flashnext_writeback.zig"); + +pub const Suite = struct { + a: std.mem.Allocator, + writebacks: [3]writeback.Writeback, + masks: tp_.Masks, + canons: canon_.Canons, + + pub fn create(a: std.mem.Allocator, gpu: check.Gpu, s: cuda.Stream, kernels: *std.ArrayList(g.Kernel), hc: cuda.Module, controls: cuda.Module) !*Suite { + const self = try a.create(Suite); + errdefer a.destroy(self); + self.a = a; + for (std.enums.values(writeback.Mode), &self.writebacks) |mode, *w| { + const symbol: [:0]const u8 = switch (mode) { + .sums => "fn_hc_writeback0", + .moe => "fn_hc_writeback_moe", + .slices => "fn_hc_writeback_slices", + }; + w.* = .{ .rig = .{ .gpu = gpu, .s = s }, .f = try hc.function(symbol), .mode = mode }; + } + errdefer for (&self.writebacks) |*w| w.deinit(); + for (&self.writebacks) |*w| try kernels.append(a, w.kernel()); + self.masks = try tp_.Masks.init(gpu, s); + errdefer self.masks.deinit(); + try self.masks.add(a, kernels); + self.canons = try canon_.Canons.init(gpu, s, controls); + errdefer self.canons.deinit(); + try self.canons.add(a, kernels); + return self; + } + + pub fn destroy(self: *Suite) void { + self.canons.deinit(); + self.masks.deinit(); + for (&self.writebacks) |*w| w.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 7b6c80bcb4..47b185fab3 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 &.{}; return error.UnknownContractImage; diff --git a/zig/tests/cuda/contract_tests.zig b/zig/tests/cuda/contract_tests.zig index 92c1d691a6..0ed2a0a835 100644 --- a/zig/tests/cuda/contract_tests.zig +++ b/zig/tests/cuda/contract_tests.zig @@ -17,7 +17,6 @@ const mx_k = @import("contract_mxfp8.zig"); const half_k = @import("contract_lane_half.zig"); const wide_k = @import("contract_projection_f32.zig"); const rank_k = @import("contract_rank.zig"); -const writeback_k = @import("contract_flashnext_writeback.zig"); const hcnorm_k = @import("contract_hc_normalize.zig"); const hcmix_k = @import("contract_hc_mix.zig"); const utility_k = @import("contract_utility.zig"); @@ -62,7 +61,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").*; @@ -242,17 +240,8 @@ pub fn run(gpu: check.Gpu) !void { try kernels.append(a, rank.kernel()); var hc = try cuda.Module.load(gpu.d, kk.flashnext_hc); defer hc.unload(); - var writebacks: [3]writeback_k.Writeback = undefined; - for (std.enums.values(writeback_k.Mode), &writebacks) |mode, *w| { - const symbol: [:0]const u8 = switch (mode) { - .sums => "fn_hc_writeback0", - .moe => "fn_hc_writeback_moe", - .slices => "fn_hc_writeback_slices", - }; - w.* = .{ .rig = .{ .gpu = gpu, .s = s }, .f = try hc.function(symbol), .mode = mode }; - } - 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, hc, mods[5]); + 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 6e7d343a9a..6764a8523b 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 @@ -53,7 +53,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; @@ -87,7 +87,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); } }; @@ -163,33 +163,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] @@ -245,7 +220,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 afb8f48546..5f79c031e7 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 6f40415f77..9d80056405 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..c1b5811d87 --- /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, within 5e-4 of the row's peak + 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..6347604ded --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/tp/flashnext_exl3_tp_tests.zig @@ -0,0 +1,347 @@ +//! flashnext-exl3-tp: EXL3 expert parallel (cuda_exl3_tp.zig) on two ranks, one process a GPU: layer, head, memory. +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 canonical tree adds in one order on one rank and on two) + 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 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, .{ .format = .exl3 }); + 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, "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..7c0b3019f6 --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/tp/flashnext_exl3_tpe_engine.zig @@ -0,0 +1,215 @@ +//! 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); + try flash.exl3.kv.setup(g.e, true); // the server's load hook: EXL3 scratch format, chunk and head rows + 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 80f80fbd51..cdfd5c7a13 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 48673199cf..4dc7ec9e75 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 ecb4a69e87..e5f0e39a2c 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -105,10 +105,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"); @@ -151,7 +154,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 @@ -247,6 +252,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 @@ -374,7 +380,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, "flashnext-attn-lanes")) return @import("contract_fn_attn_lanes.zig").run(gpu); if (std.mem.eql(u8, cmd, "flashnext-readout-cut")) return @import("flashnext_readout_cut.zig").run(gpu); @@ -434,8 +442,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-shapes")) return exl3_plan_bench.profileShapes(gpu, if (rest.len > 0) rest[0] else null); if (std.mem.eql(u8, cmd, "exl3-midm")) return exl3_plan_bench.midm(gpu); From 5cef4bb82ff3a444ef574ef9ab77d75916ba842d Mon Sep 17 00:00:00 2001 From: BlivionIaG <15323295+BlivionIaG@users.noreply.github.com> Date: Tue, 3 Nov 2026 08:36:26 +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 tf_exl3_promptd_0, _1 and _2 as replacements of linear_body, each bit-equal to what it replaces over the row counts and widths it serves and through gates 1 to 4 of the kernel contract, with planted controls that must go red. Gate 4 runs zero trellis words for the expert tiles (one dyadic weight makes every dot an exact fp32 integer) and all-ones states for the dense tiles (an fp16 sum overflows). tf-cuda-test exl3-prefill-contract and exl3-promptd-contract run those gates alone. The TP prefill's tf_exl3pf_rotate, tf_exl3pf_tp_combine, tf_exl3pf_tp_finish and tf_exl3pf_tp_place carry contract kernels with planted controls as well, so all 13 EXL3 prefill and tensor-parallel launches run gates 1 to 4; their host references run in test-flashnext-tppf-contract-host, and exl3-prefill mixed and resume device checks cover segments and states kept at any row. --- docs/kernel-contract.md | 26 +- docs/recipes/flash-next-cuda.md | 9 + zig/build/cuda.zig | 4 +- zig/build/cuda_kernels.zig | 2 + zig/build/exl3_host.zig | 5 +- zig/build/exl3_prefill_body.zig | 16 + zig/build/flashnext_tppf_host.zig | 9 + zig/kernels/cuda/contract_controls.cu | 3 + zig/kernels/cuda/contract_tppf_controls.cuh | 286 +++++++++ zig/kernels/cuda/exl3_candidate.cu | 114 ++++ .../flashnext_exl3/flashnext_exl3_prefill.cu | 478 +++++++++++++++ .../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 | 26 +- zig/src/families/flashnext/cuda_batch_moe.zig | 29 + .../families/flashnext/cuda_collective.zig | 20 +- zig/src/families/flashnext/cuda_lanes.zig | 10 +- .../families/flashnext/cuda_lanes_load.zig | 13 +- 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 | 15 +- .../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 | 338 +++++++++++ .../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 | 296 +++++++++ .../exl3/prefill/cuda_exl3_prefill_test.zig | 303 ++++++++++ .../exl3/prefill/cuda_exl3_prefill_tiles.zig | 49 ++ .../cuda_exl3_prefill_writeback_test.zig | 131 ++++ .../exl3/prefill/exl3_prefill_window.zig | 96 +++ zig/tests/cuda/contract_arch_test.zig | 1 + .../cuda/contract_flashnext_exl3_prefill.zig | 420 +++++++++++++ .../cuda/contract_flashnext_exl3_promptd.zig | 453 ++++++++++++++ .../cuda/contract_flashnext_exl3_suite.zig | 20 + .../cuda/contract_flashnext_exl3_tppf.zig | 201 +++++++ .../cuda/contract_flashnext_exl3_tppf_ref.zig | 556 +++++++++++++++++ .../contract_flashnext_exl3_tppf_rows_ref.zig | 370 ++++++++++++ .../contract_flashnext_exl3_tppf_test.zig | 254 ++++++++ zig/tests/cuda/contract_limits.zig | 2 +- .../prefill/exl3_prefill_dense.zig | 43 +- .../prefill/exl3_prefill_engine.zig | 361 +++++++++++ .../prefill/exl3_prefill_experts.zig | 131 +++- .../prefill/exl3_prefill_segments.zig | 274 +++++++++ .../prefill/exl3_prefill_tests.zig | 12 + .../tp/exl3_tpprefill_tests.zig | 340 +++++++++++ zig/tests/cuda/main.zig | 12 + zig/tests/exl3_prefill_body_compile.zig | 5 + 64 files changed, 8337 insertions(+), 66 deletions(-) create mode 100644 zig/build/exl3_prefill_body.zig create mode 100644 zig/build/flashnext_tppf_host.zig create mode 100644 zig/kernels/cuda/contract_tppf_controls.cuh 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/contract_flashnext_exl3_promptd.zig create mode 100644 zig/tests/cuda/contract_flashnext_exl3_tppf.zig create mode 100644 zig/tests/cuda/contract_flashnext_exl3_tppf_ref.zig create mode 100644 zig/tests/cuda/contract_flashnext_exl3_tppf_rows_ref.zig create mode 100644 zig/tests/cuda/contract_flashnext_exl3_tppf_test.zig create mode 100644 zig/tests/cuda/flashnext_exl3/prefill/exl3_prefill_engine.zig create mode 100644 zig/tests/cuda/flashnext_exl3/prefill/exl3_prefill_segments.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 ce6c2338e7..251a6b6b23 100644 --- a/docs/kernel-contract.md +++ b/docs/kernel-contract.md @@ -311,6 +311,28 @@ sum of |y w| with n = 11 and u = 2^-24, not the contract's depth 17: a term meet of c and one per tree level whose sibling holds terms, at most 11 - c, so n <= 11, as for the old one chain over the 11 slots (the bound is the chain's, not a lower one). On router-like and wide-range rows the tree's worst output is no worse than the old chain's or a rank-split sum's and its mean within a quarter of the best of them. +CUDA also checks Flash Next's EXL3 deep prefill: its rotate and gate-up pass (`tf_exl3pf_rotate`, `tf_exl3pf_gu`) and +its down pass (`tf_exl3pf_down`), on rows of ten routed picks and the shared expert at 1 to 1,408 rows, each bit-equal +to the grouped decode form of `exl3_experts` it replaces (the planted control runs the gate-up pass against a decode +form with the other SwiGLU roundings); gate 4 runs zero trellis words, whose one dyadic weight makes every dot an exact +FP32 integer. Its dense prompt tile (`exl3_candidate.cu`'s `tf_exl3_promptd_0`, `_1` and `_2`, one a codebook) is +bit-equal to `linear_body` at 1 to 2,048 rows (the 64-row block's and the 256-row chunk's edges included) at every +width of 1 to 8 bits, on each form `plan()` gives (a K split or none, 4 or 8 warps) and on a family's fixed plans (one +split of 4 warps, two of 8); gate 4 runs all-ones states, one weight at every width, where an fp16 sum overflows. Four +controls: the tile at another K split, one a codebook, and one that zeroes x's NaN and infinities (red at gate 1). +CUDA also checks the deep prefill's rotate and its expert-parallel reduction on their own (`flashnext_exl3_prefill.cu`, +`tf-cuda-test exl3-tppf-contract` runs them alone): `tf_exl3pf_rotate` (the tables of 16 experts, 1 to 1,408 rows of ten +picks and the shared expert at width 2,560) against a host fp32 Hadamard in the kernel's butterfly order, bit for bit +(NaN payloads aside), and `tf_exl3pf_tp_combine` (every rank's partial blocks into the all-gather layout at degrees 1 +to 8, one region or 256-row sub-blocks, to 2,048 rows, picks masked, foreign or past the table), `tf_exl3pf_tp_finish` +(the gathered partials at degrees 2, 3, 5 and 8, to 4,096 rows) and `tf_exl3pf_tp_place` (y[r, 0] and the weights, the +other slots untouched, to 1,024 rows), each byte-equal to the host `exl3_canon.zig` order, with float64's class and its +bound on every output. Gate 4 runs a transform whose two big outputs leave the rest cancelling (rotate), mirrored +integer products (combine) and shuffled pairs of 2^19 to 2^20 (finish); place has no sums (N/A). Fourteen controls: the +rotate with NaN zeroed, its input written, the up table's signs from the gate table, and bf16 butterfly adds; the +combine as one chain, with NaN products dropped, its input written, and bf16 tree adds; the finish as one chain, with +bf16 merges and its input written; the place with NaN zeroed, its input written, and every weight 1. The host test +`test-flashnext-tppf-contract-host` runs the same oracles through the gates and the gathered ranks equal to one rank. CUDA also checks Nemotron's pocket (`nemotron_ops.cu`, the learned neurons the living-weights learner adds beside the shared expert): relu^2 up into act rows and down added into the shared expert's slot, a block row a row and tiled, the down to fp32 and bf16, against float64, the neurons past the count NaN so a read of one shows (two controls: the up's @@ -355,7 +377,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. The first, `fn_exl3_split_f16`, also runs gates 1 to 4 on the CUDA runner (`contract_flashnext_exl3.zig`): an fp16 row splits into bf16 hi = bf16(x) and lo = bf16(x - hi), every output's bits must equal float64's (NaN stays NaN; an @@ -363,6 +385,8 @@ infinity gives hi = Inf and lo = NaN), and the device's hi + lo must equal x for hold all 65,536 fp16 patterns and partial last blocks (n = 1,000 and 65,537). Three controls go red: one flushing NaN and Inf to zero (gate 1), one launching whole blocks only (gate 3, at every tail) and one writing hi over its input (gate 2). Gate 4 is N/A: it has no sums. +The other 12 launches of the set also run gates 1 to 4 on the CUDA runner (see the coverage above): the canonical tree +and finish, the pick mask, the prefill's rotate, gate-up, down, combine, finish and place, and the prompt tiles. 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/recipes/flash-next-cuda.md b/docs/recipes/flash-next-cuda.md index b0351aae65..2fe3f96570 100644 --- a/docs/recipes/flash-next-cuda.md +++ b/docs/recipes/flash-next-cuda.md @@ -127,6 +127,15 @@ starts a rank. These switches tune it: | `TF_FLASHNEXT_EXL3_ROUTED` | host | `device` keeps every routed expert on the device. | | `TF_FLASHNEXT_EXL3_ARCH` | unset | `unqualified` serves on a device that is not sm_120, without a claim. | | `TF_FLASHNEXT_EXL3_MTP_RESIDENT` | `0` | `1` keeps the MTP layer's experts on the device instead of a pool (one rank). | +| `TF_FLASHNEXT_EXL3_PREFILL` | on | `0` turns the deep-chunk prefill path off. | +| `TF_FLASHNEXT_EXL3_PREFILL_ROWS` | 8,192 | Rows of a deep chunk; also what the server reserves at load. | +| `TF_FLASHNEXT_EXL3_PREFILL_RING` | on | `0` stops streaming the coming layers' experts into one ring of the pools' lines. | +| `TF_FLASHNEXT_EXL3_PREFILL_RING_AHEAD` | 1 | Layers ahead the ring prefetches. | +| `TF_FLASHNEXT_EXL3_PREFILL_OVERLAP` | on | `0` (two ranks) stops overlapping the all-reduce with a chunk's compute. | +| `TF_FLASHNEXT_EXL3_PREFILL_REDUCED_BRANCH` | `0` | Experimental: `1` (two ranks) reduces the weighted branch in fp32, written back with unit weights. | +| `TF_FLASHNEXT_EXL3_PREFILL_TWO_CHUNK` | `0` | Experimental: `1` (two ranks, ring) runs a second deep chunk beside the first when memory admits it. | +| `TF_FLASHNEXT_EXL3_PREFILL_SNAP_GRID` | on | Diagnostics only: `0` keeps prompt-cache states at any row instead of on the serving step's starts, to test that a resumed prompt equals a fresh one. | +| `TF_FLASHNEXT_PROFILE` | `0` | Diagnostics only: `1` times the ranks' combine, all-reduce and place of a deep chunk, with a wait each. | | `TF_FLASHNEXT_PLE_READ` | `map` | `pread` reads the n-gram embedding rows by positional read. | | `TF_FLASHNEXT_SNAP_HOST` | `1` above 32,768 tokens of context, else `0` | Whether kept prompt-cache states live in pinned host memory. | | `TF_FLASHNEXT_CACHE_TELEMETRY` | `0` | Diagnostics only: `1` logs `EXL3-CACHE` expert-cache counters per phase, with a sync at each boundary. | diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index 328eb6db11..935fac84a4 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -386,7 +386,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; @@ -473,7 +473,7 @@ 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, mods.core); } - @import("exl3_host.zig").steps(b, flashnextModule(b, host, .debug, cuda, mods.core, mods.lanes, mods.qwen_vision, mods.api)); + @import("exl3_host.zig").steps(b, mods.core, flashnextModule(b, host, .debug, cuda, mods.core, mods.lanes, mods.qwen_vision, mods.api), &.{ step, flash_contract }); // the tokenizer's own tests (memo, steps, regex, unicode) run here: no other test binary reaches them 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:"; diff --git a/zig/build/cuda_kernels.zig b/zig/build/cuda_kernels.zig index 0bd612df8b..62476b939c 100644 --- a/zig/build/cuda_kernels.zig +++ b/zig/build/cuda_kernels.zig @@ -152,6 +152,8 @@ pub const kernels = [_]Kernel{ .{ .name = "flashnext_exl3", .src = "flashnext_exl3/flashnext_exl3", .flags = glue }, // ours: Flash Next's EXL3 expert-parallel pick mask .{ .name = "flashnext_exl3_tp", .src = "flashnext_exl3/flashnext_exl3_tp", .flags = glue }, + // ours: Flash Next EXL3 prefill experts, 64-member tiles any width + .{ .name = "flashnext_exl3_prefill", .src = "flashnext_exl3/flashnext_exl3_prefill", .flags = &.{ "-O3", "--expt-relaxed-constexpr" } }, }; /// 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_host.zig b/zig/build/exl3_host.zig index 3e46cb84d5..78d530adb2 100644 --- a/zig/build/exl3_host.zig +++ b/zig/build/exl3_host.zig @@ -1,6 +1,7 @@ -//! The EXL3 host tests of the Flash Next module. +//! The EXL3 host tests of the Flash Next module, and its TP prefill reduction launches through the contract gates. const std = @import("std"); -pub fn steps(b: *std.Build, flashnext: *std.Build.Module) void { +pub fn steps(b: *std.Build, core: *std.Build.Module, flashnext: *std.Build.Module, all: []const *std.Build.Step) void { + @import("flashnext_tppf_host.zig").steps(b, core, flashnext, all); 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); } 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/build/flashnext_tppf_host.zig b/zig/build/flashnext_tppf_host.zig new file mode 100644 index 0000000000..286a118234 --- /dev/null +++ b/zig/build/flashnext_tppf_host.zig @@ -0,0 +1,9 @@ +//! The TP prefill reduction launches' host cases through the real gates: oracles, stress and planted faults, no device. +const std = @import("std"); + +pub fn steps(b: *std.Build, core: *std.Build.Module, family: *std.Build.Module, all: []const *std.Build.Step) void { + const root = b.createModule(.{ .root_source_file = b.path("zig/tests/cuda/contract_flashnext_exl3_tppf_test.zig"), .target = b.graph.host, .optimize = .debug, .link_libc = true, .imports = &.{ .{ .name = "core", .module = core }, .{ .name = "flashnext", .module = family } } }); + const run = b.addRunArtifact(b.addTest(.{ .root_module = root })); + for (all) |step| step.dependOn(&run.step); + b.step("test-flashnext-tppf-contract-host", "The TP prefill launches' oracles, stress and planted faults through the gates").dependOn(&run.step); +} diff --git a/zig/kernels/cuda/contract_controls.cu b/zig/kernels/cuda/contract_controls.cu index 8fe55c7f3d..5185f52e8e 100644 --- a/zig/kernels/cuda/contract_controls.cu +++ b/zig/kernels/cuda/contract_controls.cu @@ -536,3 +536,6 @@ TF_CTL_EXL3_GROUPED(tf_ctl_exl3_grouped_stale, 2, tf_exl3::STALE) // Gates 1-4: the canonical expert reduction's leaf tree and finish, each with one planted fault. #include "contract_leaf_controls.cuh" + +// Gates 1-4: the TP prefill's rotate, combine, finish and place launches, each with planted faults. +#include "contract_tppf_controls.cuh" diff --git a/zig/kernels/cuda/contract_tppf_controls.cuh b/zig/kernels/cuda/contract_tppf_controls.cuh new file mode 100644 index 0000000000..dfd014e2c3 --- /dev/null +++ b/zig/kernels/cuda/contract_tppf_controls.cuh @@ -0,0 +1,286 @@ +// Planted faults for tf_exl3pf_rotate and the TP reduction's tf_exl3pf_tp_combine, _tp_finish and _tp_place. +#pragma once +#include +#include + +#include "exl3_decode.cuh" + +namespace tf_tppf_ctl { + +enum Fault { FLAT = 0, NAN_DROPPED = 1, WRITES_INPUT = 2, BF16_ADDS = 3, WRONG = 4 }; + +__device__ __forceinline__ float rb(float x) { return __bfloat162float(__float2bfloat16_rn(x)); } + +__device__ __forceinline__ float add(float a, float b, bool bf16) { + const float s = __fadd_rn(a, b); + return bf16 ? rb(s) : s; +} + +// fwht128 of exl3_decode.cuh with every add and subtract rounded to bf16 (gate 4). +__device__ __forceinline__ void fwht128_bf16(float (&v)[4], int lane) { + const float a = rb(v[0] + v[1]), b = rb(v[0] - v[1]), c = rb(v[2] + v[3]), d = rb(v[2] - v[3]); + v[0] = rb(a + c); + v[1] = rb(b + d); + v[2] = rb(a - c); + v[3] = rb(b - d); +#pragma unroll + for (int m = 1; m < 32; m <<= 1) { +#pragma unroll + for (int j = 0; j < 4; ++j) { + const float o = __shfl_xor_sync(0xffffffffu, v[j], m); + v[j] = rb((lane & m) ? o - v[j] : v[j] + o); + } + } +} + +// tf_exl3pf_rotate, then the fault: NAN_DROPPED zeroes a non-finite x (gate 1), WRITES_INPUT (2), WRONG, BF16_ADDS (4). +template +__device__ __forceinline__ void rotate(__nv_bfloat16* x, int x_stride, const int* pick, const int* members, int slots, + const half* suh_g, const half* suh_u, half* xg, half* 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 && F != WRONG) ? suh_u : suh_g) + static_cast(e) * D + blk * 128 + 4 * lane; + __nv_bfloat16* xr = x + static_cast(row) * x_stride + blk * 128 + 4 * lane; + float v[4]; +#pragma unroll + for (int j = 0; j < 4; ++j) { + float xv = __bfloat162float(xr[j]); + if (F == NAN_DROPPED && !isfinite(xv)) xv = 0.0f; + v[j] = xv * __half2float(suh[j]); + } + if (F == BF16_ADDS) tf_tppf_ctl::fwht128_bf16(v, lane); + else tf_exl3::fwht128(v, lane); + half* o = (mat ? xu : xg) + static_cast(c) * D + blk * 128 + 4 * lane; +#pragma unroll + for (int j = 0; j < 4; ++j) o[j] = __float2half_rn(v[j] * tf_exl3::HAD_SCALE); + if (F == WRITES_INPUT && c == 0 && blk == 0 && mat == 0 && lane == 0) xr[0] = __float2bfloat16(__bfloat162float(xr[0]) + 1.0f); +} + +// The maximal complete subtrees of the leaf run [lo, hi), left to right. +__device__ __forceinline__ int 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; +} + +// tf_exl3pf_tp_combine, then the fault: FLAT one chain a block (3), NAN_DROPPED (1), WRITES_INPUT (2), BF16_ADDS (4). +template +__device__ __forceinline__ void combine(float* y, const float* wts, const int* pick, float* 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 (static_cast(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 = 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 + static_cast(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; + const float yv = yr[static_cast(k) * D + d]; + if (F == NAN_DROPPED && __fmul_rn(yv, w_sh[k]) != __fmul_rn(yv, w_sh[k])) continue; + acc = fmaf(yv, 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; + float chain = 0.0f; + 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; + } + if (F == FLAT) { + for (int k = 0; k < slots; ++k) { + const int p = pick[r * slots + k]; + const int leaf = p >= 512 ? 63 : p / 8; + if (p >= 0 && leaf >= bst[j] && leaf < bst[j] + bsz[j]) chain = fmaf(yr[static_cast(k) * D + d], w_sh[k], chain); + } + } + 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] = add(v[i], v[i + 1], F == BF16_ADDS); + i += 2; + } else { + v[m] = v[i]; + ++i; + } + s[m] = start; + ++m; + } + cnt = m; + } + const float result = F == FLAT ? chain : (cnt ? v[0] : 0.0f); + out[(static_cast(b) * sub * world * M + static_cast(rank * M + j) * mb + rr) * D + d] = result; + } + } + __syncthreads(); + if (F == WRITES_INPUT && r == 0 && threadIdx.x == 0) y[0] = __fadd_rn(y[0], 1.0f); +} + +// tf_exl3pf_tp_finish, then the fault: FLAT adds every piece in one chain (gate 3), WRITES_INPUT (2), BF16_ADDS (4). +template +__device__ __forceinline__ void finish(float* gath, float* 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 = 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]; + float chain = 0.0f; + for (int i = 0; i < nb; ++i) { + v[i] = gath[(static_cast(bq_sh[i]) * m + rr) * D + d]; + s[i] = bs_sh[i]; + z[i] = bz_sh[i]; + chain = __fadd_rn(chain, v[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] = add(v[i], v[i + 1], F == BF16_ADDS); + 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[static_cast(rr) * D + d] = F == FLAT ? chain : v[0]; + } + __syncthreads(); + if (F == WRITES_INPUT && rr == 0 && threadIdx.x == 0) gath[0] = __fadd_rn(gath[0], 1.0f); +} + +// tf_exl3pf_tp_place, then the fault: NAN_DROPPED copies a NaN as 0 (gate 1), WRITES_INPUT (2), WRONG: weights 1 (3). +template +__device__ __forceinline__ void place(float* sum, float* y, float* wts, int slots, int D) { + const int r = blockIdx.x; + float* out = y + static_cast(r) * slots * D; + float* in = sum + static_cast(r) * D; + for (int d = threadIdx.x; d < D; d += blockDim.x) { + const float v = in[d]; + out[d] = (F == NAN_DROPPED && v != v) ? 0.0f : v; + } + if (static_cast(threadIdx.x) < slots) wts[r * slots + threadIdx.x] = (threadIdx.x == 0 || F == WRONG) ? 1.0f : 0.0f; + __syncthreads(); + if (F == WRITES_INPUT && r == 0 && threadIdx.x == 0) sum[0] = __fadd_rn(sum[0], 1.0f); +} + +} // namespace tf_tppf_ctl + +#define TF_CTL_TPPF_ROTATE(NAME, F) \ + extern "C" __global__ void NAME(__nv_bfloat16* x, int x_stride, const int* pick, const int* members, int slots, \ + const half* suh_g, const half* suh_u, half* xg, half* xu, int D) { \ + tf_tppf_ctl::rotate(x, x_stride, pick, members, slots, suh_g, suh_u, xg, xu, D); \ + } +TF_CTL_TPPF_ROTATE(tf_ctl_tppf_rotate_nan_dropped, tf_tppf_ctl::NAN_DROPPED) +TF_CTL_TPPF_ROTATE(tf_ctl_tppf_rotate_writes_input, tf_tppf_ctl::WRITES_INPUT) +TF_CTL_TPPF_ROTATE(tf_ctl_tppf_rotate_wrong, tf_tppf_ctl::WRONG) +TF_CTL_TPPF_ROTATE(tf_ctl_tppf_rotate_bf16_adds, tf_tppf_ctl::BF16_ADDS) +#undef TF_CTL_TPPF_ROTATE + +#define TF_CTL_TPPF_COMBINE(NAME, F) \ + extern "C" __global__ void __launch_bounds__(256) NAME(float* y, const float* wts, const int* pick, float* out, \ + int slots, int D, int n, int sub, int world, int rank, \ + int M) { \ + tf_tppf_ctl::combine(y, wts, pick, out, slots, D, n, sub, world, rank, M); \ + } +TF_CTL_TPPF_COMBINE(tf_ctl_tppf_combine_flat, tf_tppf_ctl::FLAT) +TF_CTL_TPPF_COMBINE(tf_ctl_tppf_combine_nan_dropped, tf_tppf_ctl::NAN_DROPPED) +TF_CTL_TPPF_COMBINE(tf_ctl_tppf_combine_writes_input, tf_tppf_ctl::WRITES_INPUT) +TF_CTL_TPPF_COMBINE(tf_ctl_tppf_combine_bf16_adds, tf_tppf_ctl::BF16_ADDS) +#undef TF_CTL_TPPF_COMBINE + +#define TF_CTL_TPPF_FINISH(NAME, F) \ + extern "C" __global__ void __launch_bounds__(256) NAME(float* gath, float* sum, int D, int m, int world, int M) { \ + tf_tppf_ctl::finish(gath, sum, D, m, world, M); \ + } +TF_CTL_TPPF_FINISH(tf_ctl_tppf_finish_flat, tf_tppf_ctl::FLAT) +TF_CTL_TPPF_FINISH(tf_ctl_tppf_finish_bf16_adds, tf_tppf_ctl::BF16_ADDS) +TF_CTL_TPPF_FINISH(tf_ctl_tppf_finish_writes_input, tf_tppf_ctl::WRITES_INPUT) +#undef TF_CTL_TPPF_FINISH + +#define TF_CTL_TPPF_PLACE(NAME, F) \ + extern "C" __global__ void __launch_bounds__(256) NAME(float* sum, float* y, float* wts, int slots, int D) { \ + tf_tppf_ctl::place(sum, y, wts, slots, D); \ + } +TF_CTL_TPPF_PLACE(tf_ctl_tppf_place_nan_dropped, tf_tppf_ctl::NAN_DROPPED) +TF_CTL_TPPF_PLACE(tf_ctl_tppf_place_writes_input, tf_tppf_ctl::WRITES_INPUT) +TF_CTL_TPPF_PLACE(tf_ctl_tppf_place_wrong, tf_tppf_ctl::WRONG) +#undef TF_CTL_TPPF_PLACE 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..341f273654 --- /dev/null +++ b/zig/kernels/cuda/flashnext_exl3/flashnext_exl3_prefill.cu @@ -0,0 +1,478 @@ +// Flash Next EXL3 prefill experts: upstream's transposed 64-member prompt tiles (exl3_prompt.cuh). + +#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; + +// the widest tile a stage holds (5 bits): Flash Next's routed 3-bit and shared 5-bit experts +constexpr int KMAX2 = 10; +// 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); + } + } + } + } + // a chain ends: into its split's warp sum; a split's last into the total (0 + each, as the epilogue) + if (--left == 0) { + 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 e98b7dd800..da56e13504 100644 --- a/zig/src/cuda/exl3.zig +++ b/zig/src/cuda/exl3.zig @@ -421,19 +421,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 = planOf(w); + 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..d730d84345 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 the weight's launch plan. + 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 e5e3f79f36..4f57de5bec 100644 --- a/zig/src/cuda/kernels.zig +++ b/zig/src/cuda/kernels.zig @@ -22,6 +22,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 66f467a760..ab88811050 100644 --- a/zig/src/families/flashnext/cuda_batch.zig +++ b/zig/src/families/flashnext/cuda_batch.zig @@ -122,6 +122,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. @@ -169,29 +170,7 @@ pub fn Pass(comptime Tri: type) type { // ---- the blocks ---------------------------------------------------------------------------------------- - /// The pending MoE output (or none: mode 0, the streams' sums alone) written back into rows [r0, r0 + n). - pub fn writeback(self: *Self, r0: usize, n: usize) !void { - const s = self.s; - const dummy = s.dummy; - const h = s.h + r0 * wide * 2; - const pss = s.pss + r0 * pss_n * 4; - const inj = s.inj + r0 * streams * 2; - const rs = self.rows * dims; - if (self.pending == null) { - try self.tri.hcWriteback(h, h, pss, dummy, "*bf16", dummy, dummy, "*bf16", dummy, "*bf16", rs, n, dims, streams, 0, 1, 1, 1); - return; - } - self.pending = null; - if (self.pending_f32) { - // the fp32 branch rounded once, as MODE 2 rounds its sum (a slice of one) - self.pending_f32 = false; - 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; - } - const y = s.y + r0 * slots * dims * 4; - const wts = s.wts + r0 * slots * 4; - try self.tri.hcWriteback(h, h, pss, dummy, "*bf16", inj, y, "*fp32", wts, "*fp32", rs, n, dims, streams, 2, 10, 11, 1); - } + pub const writeback = @import("cuda_batch_moe.zig").writeback; /// Prompt rows' read-out on flashnext_hcf.cu's tiles, fn_hc_down's and fn_hc_upmix's bits (K slices folded). fn readoutTiles(self: *Self, hc_w: HcW, r0: usize, n: usize) !void { @@ -494,6 +473,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_batch_moe.zig b/zig/src/families/flashnext/cuda_batch_moe.zig index 505bf69bc6..2c1e3551d7 100644 --- a/zig/src/families/flashnext/cuda_batch_moe.zig +++ b/zig/src/families/flashnext/cuda_batch_moe.zig @@ -56,3 +56,32 @@ pub fn kept(self: anytype) void { self.canon_f32 = true; self.pending_f32 = true; } + +/// The pending MoE output (or none: mode 0, the streams' sums alone) written back into rows [r0, r0 + n). +pub fn writeback(self: anytype, r0: usize, n: usize) !void { + const s = self.s; + const dummy = s.dummy; + const h = s.h + r0 * wt.wide * 2; + const pss = s.pss + r0 * batch.pss_n * 4; + const inj = s.inj + r0 * wt.streams * 2; + const rs = self.rows * dims; + if (self.pending == null) { + try self.tri.hcWriteback(h, h, pss, dummy, "*bf16", dummy, dummy, "*bf16", dummy, "*bf16", rs, n, dims, wt.streams, 0, 1, 1, 1); + return; + } + self.pending = null; + 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, wt.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, wt.streams, 4, 1, 1, 1); + return; + } + const y = s.y + r0 * slots * dims * 4; + const wts = s.wts + r0 * slots * 4; + try self.tri.hcWriteback(h, h, pss, dummy, "*bf16", inj, y, "*fp32", wts, "*fp32", rs, n, dims, wt.streams, 2, 10, 11, 1); +} diff --git a/zig/src/families/flashnext/cuda_collective.zig b/zig/src/families/flashnext/cuda_collective.zig index 70b52c916a..11017d1f62 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"; @@ -144,6 +145,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 25c7778b9f..574f7a3482 100644 --- a/zig/src/families/flashnext/cuda_lanes.zig +++ b/zig/src/families/flashnext/cuda_lanes.zig @@ -318,9 +318,10 @@ pub fn Namespace(comptime Tri: type) type { self.lanes.removeByPtr(gop.key_ptr); } if (drafting and from != .pool and l.part_seq != null) { - l.part_draft = try Parts.take(&self.parts.drafts, self.engine.driver, mtp.Draft.partsBytes(self.head.?.ids.len, self.engine.chunk)); - l.draft = try mtp.Draft.initParts(self.gpa, self.engine.driver, self.stream, cap, self.head.?.ids.len, self.engine.chunk, l.part_draft.?); - } else 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); + const rows = @import("exl3/prefill/cuda_exl3_prefill.zig").draftRows(self.engine); + l.part_draft = try Parts.take(&self.parts.drafts, self.engine.driver, mtp.Draft.partsBytes(self.head.?.ids.len, rows)); + l.draft = try mtp.Draft.initParts(self.gpa, self.engine.driver, self.stream, cap, self.head.?.ids.len, rows, l.part_draft.?); + } else 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)); if (s.images.len != 0) { const v = self.vision orelse return error.VisionNotEnabled; l.seq.media = try v.state(ids, s.images); @@ -340,7 +341,8 @@ 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); + // EXL3 deep prefill + 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)); 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 939ba5cc72..664ad0498e 100644 --- a/zig/src/families/flashnext/cuda_lanes_load.zig +++ b/zig/src/families/flashnext/cuda_lanes_load.zig @@ -105,12 +105,16 @@ 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; + // EXL3 offload: deep prefill rows before the pools + @import("exl3/prefill/cuda_exl3_prefill.zig").reserveFor(own.engine, o.context, o.tp); + 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; // an EXL3 rank 1 keeps no head cache 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)); + 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))); 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 @@ -134,11 +138,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, @@ -146,7 +151,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) @import("exl3/prefill/cuda_exl3_prefill.zig").snapGrid(own.engine, 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, diff --git a/zig/src/families/flashnext/cuda_mtp_ops.zig b/zig/src/families/flashnext/cuda_mtp_ops.zig index 97512aecef..41ac2a3848 100644 --- a/zig/src/families/flashnext/cuda_mtp_ops.zig +++ b/zig/src/families/flashnext/cuda_mtp_ops.zig @@ -124,6 +124,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; 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 91120debf8..c1ec9abf19 100644 --- a/zig/src/families/flashnext/cuda_prompt.zig +++ b/zig/src/families/flashnext/cuda_prompt.zig @@ -249,6 +249,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); @@ -413,6 +414,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 7a6be0aa9f..2971d81305 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,8 @@ 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(); + // prefill rows on the prompt tiles + if (try @import("../prefill/cuda_exl3_prefill.zig").tilesPart(de, p.stream, w, x, m, y, f32_out)) return; try de.decode(p.stream, w, x, .bf16, m, y, if (f32_out) .f32 else .bf16); }, .native => |w| { @@ -193,7 +195,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; @@ -243,6 +245,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 47540a2436..ab7a4235eb 100644 --- a/zig/src/families/flashnext/exl3/offload/cuda_exl3_pool.zig +++ b/zig/src/families/flashnext/exl3/offload/cuda_exl3_pool.zig @@ -244,6 +244,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..c79fca4a16 --- /dev/null +++ b/zig/src/families/flashnext/exl3/prefill/cuda_exl3_prefill.zig @@ -0,0 +1,338 @@ +//! 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; +} + +/// The grid the prompt cache keeps states on: the serving step's rows, or 0 (any row) with the option off. +pub fn snapGrid(e: *const prompt.Engine, base: usize) u32 { + return if (opts().snap_grid) @intCast(stepRows(e, base, null)) else 0; +} + +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..0dd403efa7 --- /dev/null +++ b/zig/src/families/flashnext/exl3/prefill/cuda_exl3_prefill_state.zig @@ -0,0 +1,296 @@ +//! 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 the prompt tiles (Dense.prompt, .tiles). + tiles: bool = true, + /// After a ring prefill, refill every pool with its layer's hottest experts of the prefills. + reseed: bool = false, + reduced_branch: bool = false, + two_chunk: bool = false, + /// The prompt cache keeps states on the serving step's starts only; false: at any row (a test of resume = fresh). + snap_grid: bool = true, +}; + +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; + if (std.c.getenv("TF_FLASHNEXT_EXL3_PREFILL_SNAP_GRID")) |v| o.snap_grid = !std.mem.eql(u8, std.mem.span(v), "0"); + 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, timed 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..c0b73c47da --- /dev/null +++ b/zig/src/families/flashnext/exl3/prefill/cuda_exl3_prefill_test.zig @@ -0,0 +1,303 @@ +//! 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.expect((Options{}).snap_grid); + 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..c34bc0b8f6 --- /dev/null +++ b/zig/src/families/flashnext/exl3/prefill/cuda_exl3_prefill_tiles.zig @@ -0,0 +1,49 @@ +//! Dense EXL3 projections of a deep prefill on the prompt tiles (the hook cuda_exl3_pass.zig calls), in decode's bytes. + +const std = @import("std"); +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; + de.promptExact(s, w, x, m, y, if (f32_out) .f32 else .bf16) catch |err| return if (unserved(err)) false else err; + return true; +} + +/// A weight the tiles have no plan or variant for: promptExact refuses it before any launch, so decode takes it. +fn unserved(err: anyerror) bool { + return err == error.Exl3Shape or err == error.UnsupportedExl3Width; +} + +test "exl3: a weight the prompt tiles refuse before launching falls back to decode, any other error aborts" { + try std.testing.expect(unserved(error.Exl3Shape)); + try std.testing.expect(unserved(error.UnsupportedExl3Width)); + try std.testing.expect(!unserved(error.OutOfMemory)); + try std.testing.expect(!unserved(error.CandidateFormat)); + try std.testing.expect(!unserved(error.Exl3InputsOverScratch)); +} + +/// 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_arch_test.zig b/zig/tests/cuda/contract_arch_test.zig index 8ca636d677..a8b1f8ce1c 100644 --- a/zig/tests/cuda/contract_arch_test.zig +++ b/zig/tests/cuda/contract_arch_test.zig @@ -141,4 +141,5 @@ test { _ = @import("contract_exl3_grouped_ref.zig"); _ = @import("contract_nvfp4_waves_test.zig"); _ = @import("contract_exl3.zig"); // its coverage of the EXL3 dispatch's forms + _ = @import("contract_flashnext_exl3_promptd.zig"); // its shapes and the stress weights of gate 4 } 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..3b349fb6c7 --- /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. +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_promptd.zig b/zig/tests/cuda/contract_flashnext_exl3_promptd.zig new file mode 100644 index 0000000000..b878fe7700 --- /dev/null +++ b/zig/tests/cuda/contract_flashnext_exl3_promptd.zig @@ -0,0 +1,453 @@ +//! Flash Next EXL3's prompt tile in decode's bytes (tf_exl3_promptd_0, _1, _2) under the contract. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const check = @import("check.zig"); +const rig_ = @import("contract_rig.zig"); + +const contract = core.contract; +const c = contract.spec; +const g = contract.gates; +const exl3 = cuda.exl3; +const fmt = core.exl3_format; + +/// A weight's launch: its own `plan()`, or the (K splits, warps) a family fixes (`Weight.sk` and `.wk`). +pub const Plan = enum { own, fixed_1_4, fixed_2_8 }; +/// A planted fault: another K split and warp count than the weight's, or x with its non-finite values zeroed. +pub const Fault = enum { none, split, sanitized }; + +pub const Reg = struct { book: exl3.Codebook = .mul1, bits: u32 = 4, plan: Plan = .own, fault: Fault = .none }; + +const max_rows = 2048; // the prompt chunk (256 rows) eight times over +const max_k = 4096; + +const why = "the K sums of decoded fp16 weights: no input makes them exact in fp32, so gate 4 proves fp16's miss (an overflow) alone"; +const rule = "linear_body's bytes on the same inputs: a NaN and both infinities in x reach the outputs as linear_body's do"; + +/// A K sum in any order (gate 4 simulates a sequential one), then the chains', splits' and output rotation's adds. +fn depth(s: c.Shape) usize { + return s.k + 64; +} + +fn sh(rows: usize, n: usize, k: usize) c.Shape { + return .{ .rows = rows, .n = n, .k = k }; +} + +// Rows past the 64-row block and the 256-row chunk; plan()'s forms: no split (K 512), a split, 8 warps (n 8192) +const own_served = [_]c.Shape{ sh(65, 1024, 512), sh(257, 1024, 4096), sh(256, 8192, 1024), sh(130, 8192, 512), sh(2048, 1024, 4096) }; +// A family's fixed plan on shapes whose own plan splits K differently. +const fixed_served = [_]c.Shape{ sh(65, 1024, 4096), sh(257, 2048, 2048), sh(129, 1024, 1024) }; + +fn served(plan: Plan) []const c.Shape { + const own: []const c.Shape = &own_served; + const fixed: []const c.Shape = &fixed_served; + return if (plan == .own) own else fixed; +} + +/// The weight a registration launches on a shape (device pointers filled by the case). +fn weightOf(r: Reg, shape: c.Shape) exl3.Weight { + const halves: u32 = 2 * r.bits; + const st = exl3.strides(shape.k, shape.n, halves, false); + var w: exl3.Weight = .{ .trellis = 0, .suh = 0, .svh = 0, .k = @intCast(shape.k), .n = @intCast(shape.n), .halves = halves, .codebook = r.book, .stride_k = st.k, .stride_nb = st.nb }; + switch (r.plan) { + .own => {}, + .fixed_1_4 => { + w.sk = 1; + w.wk = 4; + }, + .fixed_2_8 => { + w.sk = 2; + w.wk = 8; + }, + } + return w; +} + +/// Another valid (K splits, warps) of the same shape: the planted fault. +fn other(w: exl3.Weight) exl3.Weight { + const p = exl3.planOf(w); + var x = w; + x.wk = if (p.wk == 4) 8 else 4; + x.sk = if (p.wk == 4) @max(1, p.sk / 2) else p.sk * 2; + return x; +} + +fn divides(w: exl3.Weight) bool { + const p = exl3.planOf(w); + return (w.k / 16) % (p.sk * p.wk) == 0; +} + +/// Whether the launch takes the shape: whole 128-blocks, rows and K within the scratch, a plan K's tiles divide by. +fn takes(r: Reg, shape: c.Shape) bool { + if (shape.rows == 0 or shape.rows > max_rows or shape.k == 0 or shape.n == 0) return false; + if (shape.k % 128 != 0 or shape.n % 128 != 0 or shape.k > max_k) return false; + const w = weightOf(r, shape); + return divides(w) and (r.fault != .split or divides(other(w))); +} + +fn kernelName(buf: []u8, r: Reg) []const u8 { + const plan = switch (r.plan) { + .own => "", + .fixed_1_4 => ", fixed plan 1 split of 4 warps", + .fixed_2_8 => ", fixed plan 2 splits of 8 warps", + }; + const lead = switch (r.fault) { + .none => "cuda Flash Next EXL3 prompt tile", + .split => "control: Flash Next EXL3 prompt tile at another K split posing as", + .sanitized => "control: Flash Next EXL3 prompt tile with x's non-finite values zeroed posing as", + }; + return std.fmt.bufPrint(buf, "{s} tf_exl3_promptd_{d} ({s}, {d}-bit{s}, linear_body's bytes)", .{ lead, @backingInt(r.book), @tagName(r.book), r.bits, plan }) catch unreachable; +} + +/// The served shapes: every row edge of a 64-row block, the 128-column and 128-input blocks, and the largest shapes. +pub fn spec(r: Reg, name: []const u8) c.Spec { + return .{ + .name = name, + .accum = .f32, + .depth = depth, + .ops = 4, + .out = .bf16, + .replaces = "tf_exl3_linear (linear_body)", + .nan = .{ .stated = rule }, + .tiles = .{ .{ .size = 64 }, .{ .size = 128, .step = 128, .min = 128 }, .{ .size = 128, .step = 128, .min = 128 } }, + .served = served(r.plan), + .sums = .{ .f16 = 4, .bf16 = 0, .exact = false, .why = why }, + }; +} + +/// core/exl3_format.zig's name for a codebook of the device's numbering. +fn formatBook(b: exl3.Codebook) fmt.Codebook { + return switch (b) { + .inst3 => .inst3, + .mcg => .mcg, + .mul1 => .mul1, + }; +} + +fn random16(rng: std.Random, out: []u16) void { + for (out) |*v| v.* = @bitCast(@as(f16, @floatCast(rng.float(f32) * 2 - 1))); +} + +pub const Case = struct { + rig: rig_.Rig, + de: *exl3.Dense, + reg: Reg, + name_buf: [200]u8 = undefined, + shape: c.Shape = undefined, + w: exl3.Weight = undefined, + trellis: []i16 = &.{}, + suh: []u16 = &.{}, + svh: []u16 = &.{}, + x: []u16 = &.{}, + xd: u64 = 0, + xz: u64 = 0, // x with its non-finite values zeroed: what the sanitized control launches on + out: u64 = 0, + base: u64 = 0, + xh: u64 = 0, + msg: [160]u8 = undefined, + + pub fn deinit(self: *Case) void { + self.free(); + self.rig.deinit(); + } + + fn free(self: *Case) void { + const a = self.rig.gpu.gpa; + a.free(self.trellis); + a.free(self.suh); + a.free(self.svh); + a.free(self.x); + self.trellis = &.{}; + self.suh = &.{}; + self.svh = &.{}; + self.x = &.{}; + } + + pub fn kernel(self: *Case) g.Kernel { + const s = spec(self.reg, kernelName(&self.name_buf, self.reg)); + return .{ .spec = s, .ctx = self, .vt = &vt, .control = switch (self.reg.fault) { + .none => null, + .split => .shapes, + .sanitized => .nonfinite, + } }; + } + + const vt = g.hooks(Case); + + /// Random trellis words, scales in [-1, 1) and bf16 x; non-finite: x[1][3] NaN, x[2][5] +Inf, x[2][9] -Inf. + pub fn prepare(self: *Case, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (!takes(self.reg, shape)) return false; + self.free(); + self.rig.reset(); + self.shape = shape; + const a = self.rig.gpu.gpa; + var prng = std.Random.DefaultPrng.init(seed ^ 0x70726d64); + const rng = prng.random(); + const width = fmt.Width{ .halves = @intCast(2 * self.reg.bits) }; + self.trellis = try a.alloc(i16, (shape.k / 16) * (shape.n / 16) * width.tileWords()); + self.suh = try a.alloc(u16, shape.k); + self.svh = try a.alloc(u16, shape.n); + if (fill == .stress) { // one weight at every width: all-ones states, scales 1 + @memset(self.trellis, -1); + @memset(self.suh, 0x3c00); + @memset(self.svh, 0x3c00); + } else { + for (self.trellis) |*v| v.* = rng.int(i16); + random16(rng, self.suh); + random16(rng, self.svh); + } + self.x = try a.alloc(u16, shape.rows * shape.k); + for (self.x) |*v| v.* = c.toBf16(rng.float(f32) * 2 - 1); + switch (fill) { + .random => {}, + .nonfinite => if (shape.rows > 2) { + self.x[shape.k + 3] = 0x7fc0; + self.x[2 * shape.k + 5] = 0x7f80; + self.x[2 * shape.k + 9] = 0xff80; + }, + .stress => { // one input a block, rotated to 4096: K equal terms, past fp16's range + @memset(self.x, 0); + for (0..shape.rows) |r| for (0..shape.k / 128) |b| { + self.x[r * shape.k + b * 128] = c.toBf16(46336); + }; + }, + } + self.w = weightOf(self.reg, shape); + self.w.trellis = try self.rig.input(std.mem.sliceAsBytes(self.trellis)); + self.w.suh = try self.rig.input(std.mem.sliceAsBytes(self.suh)); + self.w.svh = try self.rig.input(std.mem.sliceAsBytes(self.svh)); + self.xd = try self.rig.input(std.mem.sliceAsBytes(self.x)); + self.xz = self.xd; + if (self.reg.fault == .sanitized) { + const z = try a.dupe(u16, self.x); + defer a.free(z); + for (z) |*v| if (v.* & 0x7f80 == 0x7f80) { + v.* = 0; + }; + self.xz = try self.rig.input(std.mem.sliceAsBytes(z)); + } + self.out = try self.rig.output(shape.rows * shape.n * 2); + self.base = try self.rig.scratch(shape.rows * shape.n * 2); + self.xh = try self.rig.scratch(shape.rows * shape.k * 2); + return true; + } + + pub fn run(self: *Case) !void { + const w = if (self.reg.fault == .split) other(self.w) else self.w; + try self.de.promptExact(self.rig.s, w, self.xz, self.shape.rows, self.out, .bf16); + try self.rig.s.synchronize(); + } + + pub fn count(self: *Case) usize { + return self.shape.rows * self.shape.n; + } + + /// linear_body on the same inputs: the decode kernel's launch in 128-row blocks, never the mid-M kernels. + pub fn replaced(self: *Case) ![]const u8 { + self.de.lin.midm = false; + defer self.de.lin.midm = true; + try self.de.decode(self.rig.s, self.w, self.xd, .bf16, self.shape.rows, self.base, .bf16); + return self.rig.bytes(self.base, self.count() * 2); + } + + /// Gate 1: linear_body's bytes, with some output non-finite to prove the planted values reach it. + pub fn check(self: *Case, _: g.Gate) !?[]const u8 { + const a = self.rig.gpu.gpa; + const got = try a.dupe(u8, try self.rig.outputBytes()); + defer a.free(got); + 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 linear_body's", .{ at, got.len }); + for (0..self.count()) |i| if (c.classOf(c.read(.bf16, got, i)) != .finite) return null; + return try std.fmt.bufPrint(&self.msg, "no output is non-finite: the planted values prove nothing", .{}); + } + + /// x * suh rotated per 128 by linear_body's own rot_in: the fp16 rows both kernels take, as float64. + fn rotated(self: *Case, xs: []f64) !void { + var args: cuda.Args = .{}; + args.add(self.xd); + args.add(self.w.suh); + args.add(self.xh); + args.add(@as(i32, @backingInt(exl3.DType.bf16))); + args.add(@as(i32, @intCast(self.shape.rows))); + args.add(@as(i32, @intCast(self.shape.k))); + try cuda.launch.launch(self.de.lin.rot_in, .{ .grid = .{ .x = @intCast((self.shape.k / 128 + 3) / 4), .y = @intCast(self.shape.rows), .z = 1 }, .block = .{ .x = 128, .y = 1, .z = 1 } }, self.rig.s, &args); + const raw = try self.rig.bytes(self.xh, xs.len * 2); + for (xs, 0..) |*v, q| v.* = @floatCast(@as(f16, @bitCast(std.mem.readInt(u16, raw[2 * q ..][0..2], .little)))); + } + + /// Float64 outputs and their terms' magnitude; with `acc`, the K sums simulated in that accumulator and order. + fn outputs(self: *Case, want: []f64, mag: ?[]f64, acc: ?c.Accum, order: g.Order) !void { + const a = self.rig.gpu.gpa; + const k = self.shape.k; + const n = self.shape.n; + const rows = self.shape.rows; + const xs = try a.alloc(f64, rows * k); + defer a.free(xs); + try self.rotated(xs); + const wq = try a.alloc(f16, k * n); + defer a.free(wq); + fmt.unpack(.{ .halves = @intCast(self.w.halves) }, formatBook(self.reg.book), self.trellis, k / 16, n / 16, wq); + const s1 = try a.alloc(f64, rows * n); + defer a.free(s1); + const t = try a.alloc(f64, k); + defer a.free(t); + for (0..rows) |r| for (0..n) |j| { + var abs: f64 = 0; + for (0..k) |q| { + t[q] = xs[r * k + q] * @as(f64, @floatCast(wq[q * n + j])); + abs += @abs(t[q]); + } + var sum: f64 = 0; + if (acc) |ac| { + sum = contract.rows.sumIn(ac, order, t); + } else for (t) |v| { + sum += v; + } + want[r * n + j] = sum; + s1[r * n + j] = abs; + }; + fmt.rotate(want, rows, n, 1); + for (0..rows) |r| for (0..n) |j| { + const sc: f64 = @floatCast(@as(f16, @bitCast(self.svh[j]))); + want[r * n + j] *= sc; + if (mag) |mg| { + var s: f64 = 0; + const b0 = j / 128 * 128; + for (b0..b0 + 128) |q| s += s1[r * n + q]; + mg[r * n + j] = @abs(sc) * s / @sqrt(128.0); + } + }; + } + + pub fn reference(self: *Case, want: []f64, mag: []f64) !void { + return self.outputs(want, mag, null, .sequential); + } + + pub fn simulate(self: *Case, acc: c.Accum, order: g.Order, out: []f64) !void { + return self.outputs(out, null, acc, order); + } +}; + +/// Every registration and control: the three codebooks at 1-8 bits, the fixed plans, and the planted faults. +pub const registrations = blk: { + const books = [_]exl3.Codebook{ .inst3, .mcg, .mul1 }; + var list: [books.len * 8 + 2 + books.len + 1]Reg = undefined; + var i = 0; + for (books) |b| for (1..9) |bits| { + list[i] = .{ .book = b, .bits = bits }; + i += 1; + }; + list[i] = .{ .plan = .fixed_1_4 }; + list[i + 1] = .{ .plan = .fixed_2_8 }; + i += 2; + for (books) |b| { + list[i] = .{ .book = b, .fault = .split }; + i += 1; + } + list[i] = .{ .fault = .sanitized }; + break :blk list; +}; + +/// The registrations over one Dense (its decode kernel and its candidate module), for the runner's list. +pub const Set = struct { + de: ?exl3.Dense = null, + cases: [registrations.len]Case = undefined, + n: usize = 0, + + pub fn add(self: *Set, gpu: check.Gpu, s: cuda.Stream, list: *std.ArrayList(g.Kernel)) !void { + self.de = try exl3.Dense.load(gpu.d); + try self.de.?.candidates(max_k); + for (registrations) |r| { + self.cases[self.n] = .{ .rig = .{ .gpu = gpu, .s = s }, .de = &self.de.?, .reg = r }; + self.n += 1; + } + for (self.cases[0..self.n]) |*cs| try list.append(gpu.gpa, cs.kernel()); + } + + pub fn deinit(self: *Set) void { + for (self.cases[0..self.n]) |*cs| cs.deinit(); + if (self.de) |*d| d.unload(); + } +}; + +/// tf-cuda-test exl3-promptd-contract: these registrations' gates 1-4 and controls alone (kernel-contract has them). +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 set = try gpu.gpa.create(Set); + defer gpu.gpa.destroy(set); + set.* = .{}; + defer set.deinit(); + try set.add(gpu, s, &list); + var t: g.Tally = .{}; + try g.runAll(gpu.gpa, list.items, 0x7064, &t); + std.debug.print("SUMMARY exl3 promptd {d} kernels: {d} pass gates 1-4, {d} fail; {d} controls: {d} red as planted, {d} missed\n", .{ t.kernels, t.passed, t.failed, t.controls, t.caught, t.missed }); + try check.expect(t.problems() == 0, "exl3-promptd-contract: {d} gate failures, {d} controls missed", .{ t.failed, t.missed }); +} + +test "promptd contract: every served shape and edge is one the kernel takes, and a control's other plan too" { + const a = std.testing.allocator; + for (registrations) |r| { + const sp = spec(r, "t"); + const list = try c.shapes(a, sp); + defer a.free(list); + for (sp.served) |s| try std.testing.expect(takes(r, s)); + // an edge the plan cannot divide is refused by the kernel and by prepare alike, never run + for (list) |s| if (takes(r, s)) { + const w = weightOf(r, s); + try std.testing.expect(divides(w)); + try std.testing.expect(@as(?usize, exl3.variantIndex(w.halves, w.codebook, exl3.planOf(w).wk)) != null); + }; + } +} + +const Form = struct { split: bool, wk: u32 }; + +fn formOf(r: Reg, s: c.Shape) Form { + const p = exl3.planOf(weightOf(r, s)); + return .{ .split = p.sk > 1, .wk = p.wk }; +} + +test "promptd contract: every plan() form at K and N to 16384 is a served shape's of the own-plan registrations" { + const a = std.testing.allocator; + const sp = spec(.{}, "t"); + const list = try c.shapes(a, sp); + defer a.free(list); + var have: std.ArrayList(Form) = .empty; + defer have.deinit(a); + for (list) |s| if (takes(.{}, s)) { + const f = formOf(.{}, s); + for (have.items) |x| { + if (std.meta.eql(x, f)) break; + } else try have.append(a, f); + }; + var k: usize = 128; + while (k <= 16384) : (k += 128) { + var n: usize = 128; + while (n <= 16384) : (n += 128) { + const f = formOf(.{}, .{ .rows = 1, .n = n, .k = k }); + for (have.items) |x| { + if (std.meta.eql(x, f)) break; + } else { + std.debug.print("unregistered promptd form: split {} {d} warps (n {d}, k {d})\n", .{ f.split, f.wk, n, k }); + return error.TestUnexpectedResult; + } + } + } +} + +test "promptd contract: gate 4's all-ones states decode to one weight that overflows an fp16 sum at every width" { + for (registrations) |r| { + if (r.plan != .own or r.fault != .none) continue; + const width = fmt.Width{ .halves = @intCast(2 * r.bits) }; + var words: [128]i16 = @splat(-1); + var tile: [256]f16 = undefined; + fmt.unpack(width, formatBook(r.book), words[0..width.tileWords()], 1, 1, &tile); + const w: f64 = @floatCast(tile[0]); + for (tile) |v| try std.testing.expectEqual(tile[0], v); + // K equal terms of 4,096 w at the stress shape's K: past fp16's range by the margin the spec asks + try std.testing.expect(@abs(w) * 4096 * max_k >= 4 * 65504); + } +} diff --git a/zig/tests/cuda/contract_flashnext_exl3_suite.zig b/zig/tests/cuda/contract_flashnext_exl3_suite.zig index a3e63f11dd..27c4e7f24c 100644 --- a/zig/tests/cuda/contract_flashnext_exl3_suite.zig +++ b/zig/tests/cuda/contract_flashnext_exl3_suite.zig @@ -6,12 +6,18 @@ const g = @import("core").contract.gates; const tp_ = @import("contract_flashnext_exl3_tp.zig"); const canon_ = @import("contract_flashnext_exl3_canon.zig"); const writeback = @import("contract_flashnext_writeback.zig"); +const prefill_ = @import("contract_flashnext_exl3_prefill.zig"); +const promptd_ = @import("contract_flashnext_exl3_promptd.zig"); +const tppf_ = @import("contract_flashnext_exl3_tppf.zig"); pub const Suite = struct { a: std.mem.Allocator, writebacks: [3]writeback.Writeback, masks: tp_.Masks, canons: canon_.Canons, + prefill: *prefill_.Group, + promptd: *promptd_.Set, + tppf: *tppf_.Set, pub fn create(a: std.mem.Allocator, gpu: check.Gpu, s: cuda.Stream, kernels: *std.ArrayList(g.Kernel), hc: cuda.Module, controls: cuda.Module) !*Suite { const self = try a.create(Suite); @@ -33,11 +39,25 @@ pub const Suite = struct { self.canons = try canon_.Canons.init(gpu, s, controls); errdefer self.canons.deinit(); try self.canons.add(a, kernels); + self.prefill = try prefill_.Group.create(a, gpu, s, kernels); + errdefer self.prefill.destroy(); + self.promptd = try a.create(promptd_.Set); + errdefer a.destroy(self.promptd); + self.promptd.* = .{}; + errdefer self.promptd.deinit(); + try self.promptd.add(gpu, s, kernels); + self.tppf = try tppf_.Set.create(a, gpu, s, controls); + errdefer self.tppf.destroy(); + try self.tppf.add(a, kernels); return self; } pub fn destroy(self: *Suite) void { + self.tppf.destroy(); self.canons.deinit(); + self.promptd.deinit(); + self.a.destroy(self.promptd); + self.prefill.destroy(); self.masks.deinit(); for (&self.writebacks) |*w| w.deinit(); self.a.destroy(self); diff --git a/zig/tests/cuda/contract_flashnext_exl3_tppf.zig b/zig/tests/cuda/contract_flashnext_exl3_tppf.zig new file mode 100644 index 0000000000..882425177f --- /dev/null +++ b/zig/tests/cuda/contract_flashnext_exl3_tppf.zig @@ -0,0 +1,201 @@ +//! Flash Next EXL3's TP prefill launches under the contract: rotate and the reduction's combine, finish and place. +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 ref = @import("contract_flashnext_exl3_tppf_ref.zig"); +const rows_ref = @import("contract_flashnext_exl3_tppf_rows_ref.zig"); + +const g = core.contract.gates; +const pk_ = flash.exl3.prefill_k; +const Fault = ref.Fault; + +const images = struct { + pub const controls align(16) = @embedFile("fatbin_contract_controls").*; +}; + +/// A case over a rig: its inputs uploaded, its outputs poisoned, the launch of `Case`'s kernel or of a control. +fn Dev(comptime Case: type, comptime launch: anytype) type { + return struct { + const Self = @This(); + rig: rig_.Rig, + f: cuda.Function, + case: Case, + ptrs: [Case.n_in + Case.n_out]u64 = undefined, + + pub fn deinit(self: *Self) void { + self.case.deinit(); + self.rig.deinit(); + } + + const vt = g.hooks(Self); + + pub fn kernel(self: *Self) g.Kernel { + return .{ .spec = self.case.spec(), .ctx = self, .vt = &vt, .control = ref.controlGate(self.case.fault) }; + } + + pub fn prepare(self: *Self, shape: core.contract.spec.Shape, fill: g.Fill, seed: u64) !bool { + if (!try self.case.prepare(shape, fill, seed)) return false; + self.rig.reset(); + var n: usize = 0; + for (self.case.inputs()) |bytes| { + self.ptrs[n] = try self.rig.input(bytes); + n += 1; + } + for (self.case.outputLens()) |len| { + self.ptrs[n] = try self.rig.output(len); + n += 1; + } + return true; + } + + pub fn run(self: *Self) !void { + try launch(self.f, self.rig.s, &self.ptrs, &self.case); + try self.rig.s.synchronize(); + } + + pub fn count(self: *Self) usize { + return self.case.count(); + } + + pub fn reference(self: *Self, want: []f64, mag: []f64) !void { + return self.case.reference(want, mag); + } + + pub fn simulate(self: *Self, acc: core.contract.spec.Accum, order: g.Order, out: []f64) !void { + return self.case.simulate(acc, order, out); + } + + /// Guards and inputs as prepared, then the case's own check of the outputs. + pub fn check(self: *Self, gate: g.Gate) !?[]const u8 { + if (!try self.rig.intact()) return "a payload guard changed"; + const seen = try self.rig.inputBytes(); + var at: usize = 0; + for (self.case.inputs()) |bytes| { + if (!std.mem.eql(u8, seen[at..][0..bytes.len], bytes)) return "an immutable input changed"; + at += bytes.len; + } + return self.case.verify(try self.rig.outputBytes(), gate); + } + }; +} + +fn i32of(v: usize) i32 { + return @intCast(v); +} + +/// One launch a rank, all into one gather buffer (the engine's launches of its ranks, one device). +fn launchCombine(f: cuda.Function, s: cuda.Stream, p: []const u64, x: *const ref.Combine) !void { + for (0..x.shape.k) |rank| { + var a: cuda.Args = .{}; + for (p[0..4]) |ptr| a.add(ptr); + for ([_]usize{ ref.slots, x.shape.n, x.shape.rows, x.sub, x.shape.k, rank, x.m }) |v| a.add(i32of(v)); + try cuda.launch.launch(f, .{ .grid = .{ .x = @intCast(x.shape.rows) }, .block = .{ .x = 256 } }, s, &a); + } +} + +fn launchFinish(f: cuda.Function, s: cuda.Stream, p: []const u64, x: *const ref.Finish) !void { + var a: cuda.Args = .{}; + a.add(p[0]); + a.add(p[1]); + for ([_]usize{ x.shape.n, x.shape.rows, x.shape.k, x.m }) |v| a.add(i32of(v)); + try cuda.launch.launch(f, .{ .grid = .{ .x = @intCast(x.shape.rows) }, .block = .{ .x = 256 } }, s, &a); +} + +fn launchPlace(f: cuda.Function, s: cuda.Stream, p: []const u64, x: *const rows_ref.Place) !void { + var a: cuda.Args = .{}; + for (p[0..3]) |ptr| a.add(ptr); + a.add(i32of(ref.slots)); + a.add(i32of(x.shape.n)); + try cuda.launch.launch(f, .{ .grid = .{ .x = @intCast(x.shape.rows) }, .block = .{ .x = 256 } }, s, &a); +} + +fn launchRotate(f: cuda.Function, s: cuda.Stream, p: []const u64, x: *const rows_ref.Rotate) !void { + var a: cuda.Args = .{}; + a.add(p[0]); + a.add(i32of(x.shape.n)); + a.add(p[1]); + a.add(p[2]); + a.add(i32of(ref.slots)); + a.add(p[3]); + a.add(p[4]); + a.add(p[5]); + a.add(p[6]); + a.add(i32of(x.shape.n)); + try cuda.launch.launch(f, .{ .grid = .{ .x = @intCast(x.plan.n_members), .y = @intCast(x.shape.n / 128), .z = 2 }, .block = .{ .x = 32 } }, s, &a); +} + +const CombineDev = Dev(ref.Combine, launchCombine); +const FinishDev = Dev(ref.Finish, launchFinish); +const PlaceDev = Dev(rows_ref.Place, launchPlace); +const RotateDev = Dev(rows_ref.Rotate, launchRotate); + +const Plan = struct { fault: Fault, symbol: [:0]const u8 }; + +const combine_plan = [_]Plan{ .{ .fault = .none, .symbol = "" }, .{ .fault = .none, .symbol = "" }, .{ .fault = .flat, .symbol = "tf_ctl_tppf_combine_flat" }, .{ .fault = .nan_dropped, .symbol = "tf_ctl_tppf_combine_nan_dropped" }, .{ .fault = .writes_input, .symbol = "tf_ctl_tppf_combine_writes_input" }, .{ .fault = .bf16_adds, .symbol = "tf_ctl_tppf_combine_bf16_adds" } }; +const finish_plan = [_]Plan{ .{ .fault = .none, .symbol = "" }, .{ .fault = .flat, .symbol = "tf_ctl_tppf_finish_flat" }, .{ .fault = .bf16_adds, .symbol = "tf_ctl_tppf_finish_bf16_adds" }, .{ .fault = .writes_input, .symbol = "tf_ctl_tppf_finish_writes_input" } }; +const place_plan = [_]Plan{ .{ .fault = .none, .symbol = "" }, .{ .fault = .nan_dropped, .symbol = "tf_ctl_tppf_place_nan_dropped" }, .{ .fault = .writes_input, .symbol = "tf_ctl_tppf_place_writes_input" }, .{ .fault = .wrong, .symbol = "tf_ctl_tppf_place_wrong" } }; +const rotate_plan = [_]Plan{ .{ .fault = .none, .symbol = "" }, .{ .fault = .nan_dropped, .symbol = "tf_ctl_tppf_rotate_nan_dropped" }, .{ .fault = .writes_input, .symbol = "tf_ctl_tppf_rotate_writes_input" }, .{ .fault = .wrong, .symbol = "tf_ctl_tppf_rotate_wrong" }, .{ .fault = .bf16_adds, .symbol = "tf_ctl_tppf_rotate_bf16_adds" } }; + +/// The four launches (combine in its two region layouts) and each one's planted faults, over the prefill image. +pub const Set = struct { + a: std.mem.Allocator, + k: pk_.Kernels, + combines: [combine_plan.len]CombineDev, + finishes: [finish_plan.len]FinishDev, + places: [place_plan.len]PlaceDev, + rotates: [rotate_plan.len]RotateDev, + + fn function(self: *const Set, controls: cuda.Module, comptime shipped: []const u8, plan: Plan) !cuda.Function { + return if (plan.fault == .none) @field(self.k, shipped) else controls.function(plan.symbol); + } + + pub fn create(a: std.mem.Allocator, gpu: check.Gpu, s: cuda.Stream, controls: cuda.Module) !*Set { + const self = try a.create(Set); + errdefer a.destroy(self); + self.a = a; + self.k = try pk_.Kernels.load(gpu.d); + errdefer self.k.unload(); + const rig: rig_.Rig = .{ .gpu = gpu, .s = s }; + for (combine_plan, &self.combines, 0..) |plan, *x, i| x.* = .{ .rig = rig, .f = try self.function(controls, "tp_combine", plan), .case = .{ .gpa = gpu.gpa, .blocks = i == 1, .fault = plan.fault } }; + for (finish_plan, &self.finishes) |plan, *x| x.* = .{ .rig = rig, .f = try self.function(controls, "tp_finish", plan), .case = .{ .gpa = gpu.gpa, .fault = plan.fault } }; + for (place_plan, &self.places) |plan, *x| x.* = .{ .rig = rig, .f = try self.function(controls, "tp_place", plan), .case = .{ .gpa = gpu.gpa, .fault = plan.fault } }; + for (rotate_plan, &self.rotates) |plan, *x| x.* = .{ .rig = rig, .f = try self.function(controls, "rotate", plan), .case = .{ .gpa = gpu.gpa, .fault = plan.fault } }; + return self; + } + + pub fn destroy(self: *Set) void { + for (&self.combines) |*x| x.deinit(); + for (&self.finishes) |*x| x.deinit(); + for (&self.places) |*x| x.deinit(); + for (&self.rotates) |*x| x.deinit(); + self.k.unload(); + self.a.destroy(self); + } + + pub fn add(self: *Set, a: std.mem.Allocator, list: *std.ArrayList(g.Kernel)) !void { + for (&self.rotates) |*x| try list.append(a, x.kernel()); + for (&self.combines) |*x| try list.append(a, x.kernel()); + for (&self.finishes) |*x| try list.append(a, x.kernel()); + for (&self.places) |*x| try list.append(a, x.kernel()); + } +}; + +/// tf-cuda-test exl3-tppf-contract: just these launches through the gates and their controls. +pub fn run(gpu: check.Gpu) !void { + var s = try cuda.Stream.init(gpu.d, false); + defer s.deinit(); + var controls = try cuda.Module.load(gpu.d, &images.controls); + defer controls.unload(); + var list: std.ArrayList(g.Kernel) = .empty; + defer list.deinit(gpu.gpa); + const set = try Set.create(gpu.gpa, gpu, s, controls); + defer set.destroy(); + try set.add(gpu.gpa, &list); + var t: g.Tally = .{}; + try g.runAll(gpu.gpa, list.items, 0x74707066, &t); + std.debug.print("SUMMARY exl3 tppf {d} kernels: {d} pass gates 1-4, {d} fail; {d} controls: {d} red as planted, {d} missed\n", .{ t.kernels, t.passed, t.failed, t.controls, t.caught, t.missed }); + try check.expect(t.problems() == 0, "exl3-tppf-contract: {d} gate failures, {d} controls missed", .{ t.failed, t.missed }); +} diff --git a/zig/tests/cuda/contract_flashnext_exl3_tppf_ref.zig b/zig/tests/cuda/contract_flashnext_exl3_tppf_ref.zig new file mode 100644 index 0000000000..cdf55274b5 --- /dev/null +++ b/zig/tests/cuda/contract_flashnext_exl3_tppf_ref.zig @@ -0,0 +1,556 @@ +//! Host cases of Flash Next EXL3's TP prefill reduction launches: inputs, float64 and bit oracles, stress, faults. +const std = @import("std"); +const core = @import("core"); +const flash = @import("flashnext"); + +const c = core.contract.spec; +const g = core.contract.gates; +const canon = flash.exl3_canon; + +pub const slots = canon.slots; +pub const dims = 2560; + +/// A planted fault of a control kernel (contract_tppf_controls.cuh); none is the shipped kernel. +pub const Fault = enum { none, flat, nan_dropped, writes_input, bf16_adds, wrong }; + +pub fn controlGate(f: Fault) ?g.Gate { + return switch (f) { + .none => null, + .flat, .wrong => .shapes, + .nan_dropped => .nonfinite, + .writes_input => .repeat, + .bf16_adds => .sums, + }; +} + +/// An output no launch wrote: the rig poisons outputs with 0xff bytes. +const unwritten: u32 = 0xffffffff; + +/// Bit-equal to the host's result, or both NaN (a payload is not part of the contract). +pub fn sameBits(got: f32, want: f32) bool { + return @as(u32, @bitCast(got)) == @as(u32, @bitCast(want)) or (std.math.isNan(got) and std.math.isNan(want)); +} + +pub fn bf16(x: f32) f32 { + return c.fromBf16(c.toBf16(x)); +} + +fn bitsOf(out: []const u8, i: usize) f32 { + return @bitCast(std.mem.readInt(u32, out[i * 4 ..][0..4], .little)); +} + +pub fn put(out: []u8, i: usize, v: f32) void { + std.mem.writeInt(u32, out[i * 4 ..][0..4], @bitCast(v), .little); +} + +/// Fp32 values that make fp16 sums overflow and bf16 ones round away their low bits: 2^19 to 2^20, either sign. +fn big(rng: std.Random) f32 { + const v: f32 = @floatFromInt((@as(u32, 1) << 19) + rng.uintLessThan(u32, 1 << 19)); + return if (rng.boolean()) v else -v; +} + +const stated = "the bytes of the host canonical tree (NaN payloads aside) and float64's class on every output"; + +/// Each output against the host's bits, float64's class and bound (gate 1: some finite and some not). +fn verifyF32(msg: []u8, s: c.Spec, shape: c.Shape, gate: g.Gate, out: []const u8, want: []const f64, mag: []const f64, ctx: anytype, comptime expected: anytype) !?[]const u8 { + var finite: usize = 0; + for (want, mag, 0..) |w, m, i| { + const got = bitsOf(out, i); + const host = expected(ctx, i); + if (!sameBits(got, host)) return try std.fmt.bufPrint(msg, "output {d}: {x:0>8} where the host canonical reduction gives {x:0>8}", .{ i, @as(u32, @bitCast(got)), @as(u32, @bitCast(host)) }); + const have: f64 = got; + if (c.classOf(have) != c.classOf(w)) return try std.fmt.bufPrint(msg, "output {d} is {t} where float64 gives {t}", .{ i, c.classOf(have), c.classOf(w) }); + if (std.math.isFinite(w)) { + finite += 1; + if (@abs(have - w) > c.bound(s, shape, w, m)) return try std.fmt.bufPrint(msg, "output {d}: {d} is outside float64's bound around {d}", .{ i, have, w }); + } + } + if (gate == .nonfinite and (finite == 0 or finite == want.len)) return "the planted values leave all outputs finite or none"; + return null; +} + +/// The planted non-finite rows of a 3-D f32 layout: a row each of NaN, +Inf, -Inf, both infinities and Inf x 0. +fn plantRows(rows: usize, n: usize, y: []f32, w: []f32, pieces: usize) void { + for (0..rows) |r| for (0..n) |col| { + if (col % 7 != 0) continue; + const k = (col / 7) % pieces; + const at = (r * pieces + k) * n + col; + switch (r % 5) { + 0 => y[at] = std.math.nan(f32), + 1 => y[at] = std.math.inf(f32), + 2 => y[at] = -std.math.inf(f32), + 3 => { + y[at] = std.math.inf(f32); + y[(r * pieces + (k + 1) % pieces) * n + col] = -std.math.inf(f32); + }, + else => { + y[at] = std.math.inf(f32); + w[r * pieces + k] = 0; + }, + } + }; +} + +/// A row's value of one block as the kernel's leaf chains and tree give it (absent leaves are +0, never -0). +pub fn blockSum(y: *const [slots]f32, w: *const [slots]f32, p: *const [slots]i32, start: usize, log: u8, fault: Fault) f32 { + var leaf: [canon.leaves]f32 = @splat(0); + const size = @as(usize, 1) << @intCast(log); + var chain: f32 = 0; + for (0..slots) |k| { + if (p[k] < 0) continue; + const l = canon.leafOf(@intCast(p[k])); + if (l < start or l >= start + size) continue; + if (fault == .nan_dropped and std.math.isNan(y[k] * w[k])) continue; + chain = @mulAdd(f32, y[k], w[k], chain); + leaf[l - start] = @mulAdd(f32, y[k], w[k], leaf[l - start]); + } + if (fault == .flat) return chain; + var n = size; + while (n > 1) : (n /= 2) for (0..n / 2) |i| { + const s = leaf[2 * i] + leaf[2 * i + 1]; + leaf[i] = if (fault == .bf16_adds) bf16(s) else s; + }; + return leaf[0]; +} + +/// The blocks rank `rank` of `world` computes, in leaf order. +pub fn rankBlocks(buf: *[canon.max_blocks]canon.Block, world: usize, rank: usize) usize { + const r = canon.run(world, rank); + return canon.blocks(buf, r.lo, r.hi); +} + +/// A leaf's fma chain (at most 11 slots) and the 6 levels of the tree above it. +fn combineDepth(_: c.Shape) usize { + return slots + canon.leaf_log; +} + +/// Each piece meets at most 6 adds on its way up the tree. +fn finishDepth(_: c.Shape) usize { + return canon.leaf_log; +} + +pub const edges: [3]c.Tile = .{ .{}, .{ .size = 256, .step = 256, .min = 256 }, .{} }; +const combine_region = [_]c.Shape{ + .{ .rows = 5, .n = dims, .k = 2 }, + .{ .rows = 1, .n = dims, .k = 1 }, + .{ .rows = 130, .n = dims, .k = 2 }, + .{ .rows = 1024, .n = dims, .k = 1 }, + .{ .rows = 2048, .n = dims, .k = 2 }, + .{ .rows = 17, .n = dims, .k = 3 }, + .{ .rows = 40, .n = dims, .k = 5 }, + .{ .rows = 6, .n = dims, .k = 8 }, +}; +const combine_blocks = [_]c.Shape{ + .{ .rows = 300, .n = dims, .k = 2 }, + .{ .rows = 512, .n = dims, .k = 2 }, + .{ .rows = 2048, .n = dims, .k = 2 }, + .{ .rows = 700, .n = dims, .k = 5 }, + .{ .rows = 257, .n = dims, .k = 8 }, +}; +pub const max_rows = 2048; +/// The all-gather's sub-block rows (the engine's chunk is a multiple of 256, one region of `sub` rows a block). +pub const sub_rows = 256; + +/// Rank partials of one launch per rank: every rank of k launches into one gather buffer, as the all-gather lays it. +pub const Combine = struct { + gpa: std.mem.Allocator, + blocks: bool, + fault: Fault = .none, + arena: ?std.heap.ArenaAllocator = null, + shape: c.Shape = .{ .rows = 0, .n = 0, .k = 0 }, + m: usize = 1, + sub: usize = 0, + y: []f32 = &.{}, + w: []f32 = &.{}, + p: []i32 = &.{}, + msg: [200]u8 = undefined, + + pub const n_in = 3; + pub const n_out = 1; + + pub fn deinit(self: *Combine) void { + if (self.arena) |*a| a.deinit(); + } + + pub fn spec(self: *const Combine) c.Spec { + const name = switch (self.fault) { + .none => if (self.blocks) "cuda Flash Next EXL3 prefill tf_exl3pf_tp_combine (256-row sub-blocks)" else "cuda Flash Next EXL3 prefill tf_exl3pf_tp_combine (one region)", + .flat => "control: tf_exl3pf_tp_combine as one chain, no tree", + .nan_dropped => "control: tf_exl3pf_tp_combine with NaN products dropped", + .writes_input => "control: tf_exl3pf_tp_combine writing its input", + .bf16_adds => "control: tf_exl3pf_tp_combine with bf16 tree adds", + else => unreachable, + }; + return .{ .name = name, .accum = .f32, .depth = combineDepth, .out = .f32, .nan = .{ .stated = stated }, .tiles = edges, .served = if (self.blocks) &combine_blocks else &combine_region }; + } + + pub fn accepts(shape: c.Shape) bool { + return shape.rows >= 1 and shape.rows <= max_rows and shape.n >= 256 and shape.n % 256 == 0 and shape.n <= dims and shape.k >= 1 and shape.k <= 8; + } + + pub fn prepare(self: *Combine, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (!accepts(shape)) return false; + self.deinit(); + self.arena = std.heap.ArenaAllocator.init(self.gpa); + const a = self.arena.?.allocator(); + self.shape = shape; + self.m = canon.maxPartials(shape.k); + self.sub = if (self.blocks and shape.rows > sub_rows) sub_rows else shape.rows; + self.y = try a.alloc(f32, shape.rows * slots * shape.n); + self.w = try a.alloc(f32, shape.rows * slots); + self.p = try a.alloc(i32, shape.rows * slots); + var prng = std.Random.DefaultPrng.init(seed ^ 0x636f6d62); + const rng = prng.random(); + if (fill == .stress) self.stress(rng) else self.random(rng, fill); + return true; + } + + pub fn inputs(self: *const Combine) [n_in][]const u8 { + return .{ std.mem.sliceAsBytes(self.y), std.mem.sliceAsBytes(self.w), std.mem.sliceAsBytes(self.p) }; + } + + pub fn outputLens(self: *const Combine) [n_out]usize { + return .{self.count() * 4}; + } + + pub fn count(self: *const Combine) usize { + return self.shape.k * self.m * self.shape.rows * self.shape.n; + } + + /// Router-like picks (clustered leaves, the shared expert last), some masked (-1) or past the table, small weights. + fn random(self: *Combine, rng: std.Random, fill: g.Fill) void { + for (0..self.shape.rows) |r| { + const base = [2]u16{ rng.uintLessThan(u16, 64) * 8, rng.uintLessThan(u16, 64) * 8 }; + for (0..slots) |k| { + var id: i32 = if (k == slots - 1) 512 else if (rng.boolean()) rng.uintLessThan(u16, 512) else base[rng.uintLessThan(usize, 2)] + rng.uintLessThan(u16, 8); + if (rng.uintLessThan(u8, 5) == 0) id = -1 else if (k == slots - 1 and rng.uintLessThan(u8, 9) == 0) id = 513 + @as(i32, rng.uintLessThan(u8, 9)); + self.p[r * slots + k] = id; + } + } + for (self.y) |*v| v.* = if (rng.uintLessThan(u8, 23) == 0) -0.0 else rng.floatNorm(f32) * 0.05; + for (self.w) |*v| v.* = if (rng.uintLessThan(u8, 29) == 0) -0.0 else rng.floatNorm(f32) * 0.2; + if (fill == .nonfinite) plantRows(self.shape.rows, self.shape.n, self.y, self.w, slots); + } + + /// Integers: ten routed picks in rank 0's first block, mirrored (same y, negated weight), the shared slot 3 x 1. + fn stress(self: *Combine, rng: std.Random) void { + var bl: [canon.max_blocks]canon.Block = undefined; + _ = rankBlocks(&bl, self.shape.k, 0); + const first: usize = bl[0].start; + const lo = first * canon.leaf_ids; + const span = @min(@as(usize, 1) << @intCast(bl[0].log), canon.leaves - first) * canon.leaf_ids; + const n = self.shape.n; + for (0..self.shape.rows) |r| { + for (0..slots - 1) |k| self.p[r * slots + k] = @intCast(lo + rng.uintLessThan(usize, span)); + self.p[r * slots + slots - 1] = 512; + for (0..(slots - 1) / 2) |k| { + const wv: f32 = @floatFromInt(512 + rng.uintLessThan(u32, 512)); + self.w[r * slots + k] = wv; + self.w[r * slots + slots - 2 - k] = -wv; + for (0..n) |col| { + const yv: f32 = @floatFromInt(512 + rng.uintLessThan(u32, 512)); + self.y[(r * slots + k) * n + col] = yv; + self.y[(r * slots + slots - 2 - k) * n + col] = yv; + } + } + self.w[r * slots + slots - 1] = 1; + for (0..n) |col| self.y[(r * slots + slots - 1) * n + col] = 3; + } + } + + /// Where output `i` of the gather layout comes from: its row, column and block (null: no rank writes it). + const At = struct { row: usize, col: usize, block: ?canon.Block }; + + fn locate(self: *const Combine, i: usize) At { + const n = self.shape.n; + const region = self.sub * self.shape.k * self.m * n; + const b = i / region; + const mb = @min(self.sub, self.shape.rows - b * self.sub); + const rem = i - b * region; + const slot = rem / (mb * n); + const row = b * self.sub + (rem / n) % mb; + var bl: [canon.max_blocks]canon.Block = undefined; + const nb = rankBlocks(&bl, self.shape.k, slot / self.m); + return .{ .row = row, .col = rem % n, .block = if (slot % self.m < nb) bl[slot % self.m] else null }; + } + + fn slotValues(self: *const Combine, at: At, y: *[slots]f32, w: *[slots]f32, p: *[slots]i32) void { + for (0..slots) |k| { + y[k] = self.y[(at.row * slots + k) * self.shape.n + at.col]; + w[k] = self.w[at.row * slots + k]; + p[k] = self.p[at.row * slots + k]; + } + } + + fn expectedAs(self: *const Combine, i: usize, fault: Fault) f32 { + const at = self.locate(i); + const blk = at.block orelse return @bitCast(unwritten); + var y: [slots]f32 = undefined; + var w: [slots]f32 = undefined; + var p: [slots]i32 = undefined; + self.slotValues(at, &y, &w, &p); + return blockSum(&y, &w, &p, blk.start, blk.log, fault); + } + + fn expected(self: *const Combine, i: usize) f32 { + return self.expectedAs(i, .none); + } + + /// Float64's sum of the products of the block's slots in slot order, and the sum of their magnitudes. + fn terms(self: *const Combine, i: usize, buf: *[slots]f64) ?usize { + const at = self.locate(i); + const blk = at.block orelse return null; + var y: [slots]f32 = undefined; + var w: [slots]f32 = undefined; + var p: [slots]i32 = undefined; + self.slotValues(at, &y, &w, &p); + var m: usize = 0; + for (0..slots) |k| { + if (p[k] < 0) continue; + const l = canon.leafOf(@intCast(p[k])); + if (l < blk.start or l >= blk.start + (@as(usize, 1) << @intCast(blk.log))) continue; + buf[m] = @as(f64, y[k]) * @as(f64, w[k]); + m += 1; + } + return m; + } + + pub fn reference(self: *const Combine, want: []f64, mag: []f64) !void { + var t: [slots]f64 = undefined; + for (want, mag, 0..) |*w, *m, i| { + const cnt = self.terms(i, &t) orelse { + w.* = std.math.nan(f64); + m.* = 0; + continue; + }; + w.* = 0; + m.* = 0; + for (t[0..cnt]) |x| { + w.* += x; + m.* += @abs(x); + } + } + } + + pub fn simulate(self: *const Combine, acc: c.Accum, order: g.Order, out: []f64) !void { + var t: [slots]f64 = undefined; + for (out, 0..) |*o, i| o.* = if (self.terms(i, &t)) |cnt| core.contract.rows.sumIn(acc, order, t[0..cnt]) else std.math.nan(f64); + } + + pub fn verify(self: *Combine, out: []const u8, gate: g.Gate) !?[]const u8 { + const a = self.arena.?.allocator(); + const want = try a.alloc(f64, self.count()); + const mag = try a.alloc(f64, self.count()); + try self.reference(want, mag); + return verifyF32(&self.msg, self.spec(), self.shape, gate, out, want, mag, self, Combine.expected); + } + + /// The host's emulation of a launch (a planted fault's too), for the gates' own host tests. + pub fn emulate(self: *Combine, fault: Fault, out: []u8) void { + for (0..self.count()) |i| put(out, i, self.expectedAs(i, fault)); + if (fault == .writes_input) self.y[0] += 1; + } +}; + +const finish_served = [_]c.Shape{ + .{ .rows = 5, .n = dims, .k = 2 }, + .{ .rows = 1, .n = dims, .k = 2 }, + .{ .rows = 44, .n = dims, .k = 2 }, + .{ .rows = 256, .n = dims, .k = 2 }, + .{ .rows = 4096, .n = dims, .k = 2 }, + .{ .rows = 7, .n = dims, .k = 3 }, + .{ .rows = 9, .n = dims, .k = 5 }, + .{ .rows = 6, .n = dims, .k = 8 }, +}; +pub const max_finish_rows = 4096; + +/// The tree finished from the gathered partials of one region: world x M partial blocks a row, merged in leaf order. +pub const Finish = struct { + gpa: std.mem.Allocator, + fault: Fault = .none, + arena: ?std.heap.ArenaAllocator = null, + shape: c.Shape = .{ .rows = 0, .n = 0, .k = 0 }, + m: usize = 1, + ps: [canon.max_pieces]canon.Piece = undefined, + np: usize = 0, + vals: []f32 = &.{}, + msg: [200]u8 = undefined, + + pub const n_in = 1; + pub const n_out = 1; + + pub fn deinit(self: *Finish) void { + if (self.arena) |*a| a.deinit(); + } + + pub fn spec(self: *const Finish) c.Spec { + const name = switch (self.fault) { + .none => "cuda Flash Next EXL3 prefill tf_exl3pf_tp_finish", + .flat => "control: tf_exl3pf_tp_finish as one chain, no tree", + .bf16_adds => "control: tf_exl3pf_tp_finish with bf16 merges", + .writes_input => "control: tf_exl3pf_tp_finish writing its input", + else => unreachable, + }; + return .{ .name = name, .accum = .f32, .depth = finishDepth, .out = .f32, .nan = .{ .stated = stated }, .tiles = edges, .served = &finish_served }; + } + + /// The gathered slot of a piece: rank * M + its index on the rank. + fn slotOf(self: *const Finish, p: canon.Piece) usize { + return @as(usize, p.rank) * self.m + p.idx; + } + + pub fn accepts(shape: c.Shape) bool { + return shape.rows >= 1 and shape.rows <= max_finish_rows and shape.n >= 256 and shape.n % 256 == 0 and shape.n <= dims and shape.k >= 1 and shape.k <= 8; + } + + pub fn prepare(self: *Finish, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (!accepts(shape)) return false; + self.deinit(); + self.arena = std.heap.ArenaAllocator.init(self.gpa); + const a = self.arena.?.allocator(); + self.shape = shape; + self.m = canon.maxPartials(shape.k); + self.np = canon.pieces(&self.ps, shape.k); + self.vals = try a.alloc(f32, shape.k * self.m * shape.rows * shape.n); + var prng = std.Random.DefaultPrng.init(seed ^ 0x66696e73); + const rng = prng.random(); + if (fill == .stress) self.stress(rng) else self.random(rng, fill); + self.poison(); + return true; + } + + pub fn inputs(self: *const Finish) [n_in][]const u8 { + return .{std.mem.sliceAsBytes(self.vals)}; + } + + pub fn outputLens(self: *const Finish) [n_out]usize { + return .{self.count() * 4}; + } + + pub fn count(self: *const Finish) usize { + return self.shape.rows * self.shape.n; + } + + fn at(self: *const Finish, piece: usize, i: usize) usize { + return self.slotOf(self.ps[piece]) * self.count() + i; + } + + /// Partials of different magnitudes (an fp32 sum of them depends on the order of the merges), a few negative zeros. + fn random(self: *Finish, rng: std.Random, fill: g.Fill) void { + for (self.vals) |*v| { + const scale: f32 = @floatFromInt(@as(u32, 1) << rng.uintLessThan(u5, 10)); + v.* = if (rng.uintLessThan(u8, 23) == 0) -0.0 else rng.floatNorm(f32) * scale; + } + if (fill != .nonfinite) return; + const n = self.shape.n; + for (0..self.shape.rows) |r| for (0..n) |col| { + if (col % 7 != 0) continue; + const piece = (col / 7) % self.np; + const i = r * n + col; + switch (r % 5) { + 0 => self.vals[self.at(piece, i)] = std.math.nan(f32), + 1 => self.vals[self.at(piece, i)] = std.math.inf(f32), + 2 => self.vals[self.at(piece, i)] = -std.math.inf(f32), + 3 => { + self.vals[self.at(piece, i)] = std.math.inf(f32); + self.vals[self.at((piece + 1) % self.np, i)] = -std.math.inf(f32); + }, + else => self.vals[self.at(self.np - 1, i)] = std.math.nan(f32), + } + }; + } + + /// Integer partials of 2^19 to 2^20 in shuffled negated pairs (odd one out 3): exact in fp32, ruinous in fp16/bf16. + fn stress(self: *Finish, rng: std.Random) void { + for (0..self.count()) |i| { + var v: [canon.max_pieces]f32 = undefined; + for (0..self.np / 2) |j| { + v[2 * j] = big(rng); + v[2 * j + 1] = -v[2 * j]; + } + if (self.np % 2 == 1) v[self.np - 1] = 3; + rng.shuffle(f32, v[0..self.np]); + for (v[0..self.np], 0..) |x, piece| self.vals[self.at(piece, i)] = x; + } + } + + /// Gather slots no rank fills (a rank with fewer blocks than M) hold a value no merge may read. + fn poison(self: *Finish) void { + var bl: [canon.max_blocks]canon.Block = undefined; + for (0..self.shape.k) |q| { + const nb = rankBlocks(&bl, self.shape.k, q); + for (nb..self.m) |j| @memset(self.vals[(q * self.m + j) * self.count() ..][0..self.count()], 1.0e30); + } + } + + fn pieceValues(self: *const Finish, i: usize, buf: *[canon.max_pieces]f32) []const f32 { + for (0..self.np) |piece| buf[piece] = self.vals[self.at(piece, i)]; + return buf[0..self.np]; + } + + fn expectedAs(self: *const Finish, i: usize, fault: Fault) f32 { + var buf: [canon.max_pieces]f32 = undefined; + const v = self.pieceValues(i, &buf); + if (fault == .flat) { + var chain: f32 = 0; + for (v) |x| chain += x; + return chain; + } + if (fault != .bf16_adds) return canon.finish(self.ps[0..self.np], v); + var top: usize = 0; + var s: [canon.max_pieces]f32 = undefined; + var lg: [canon.max_pieces]u8 = undefined; + var st: [canon.max_pieces]u8 = undefined; + for (self.ps[0..self.np], v) |p, x| { + s[top] = x; + lg[top] = p.log; + st[top] = p.start; + top += 1; + while (top >= 2 and lg[top - 1] == lg[top - 2] and st[top - 2] % (@as(usize, 2) << @intCast(lg[top - 1])) == 0) { + s[top - 2] = bf16(s[top - 2] + s[top - 1]); + lg[top - 2] += 1; + top -= 1; + } + } + return s[0]; + } + + fn expected(self: *const Finish, i: usize) f32 { + return self.expectedAs(i, .none); + } + + pub fn reference(self: *const Finish, want: []f64, mag: []f64) !void { + var buf: [canon.max_pieces]f32 = undefined; + for (want, mag, 0..) |*w, *m, i| { + w.* = 0; + m.* = 0; + for (self.pieceValues(i, &buf)) |v| { + w.* += v; + m.* += @abs(@as(f64, v)); + } + } + } + + pub fn simulate(self: *const Finish, acc: c.Accum, order: g.Order, out: []f64) !void { + var buf: [canon.max_pieces]f32 = undefined; + var t: [canon.max_pieces]f64 = undefined; + for (out, 0..) |*o, i| { + const v = self.pieceValues(i, &buf); + for (v, t[0..v.len]) |x, *y| y.* = x; + o.* = core.contract.rows.sumIn(acc, order, t[0..v.len]); + } + } + + pub fn verify(self: *Finish, out: []const u8, gate: g.Gate) !?[]const u8 { + const a = self.arena.?.allocator(); + const want = try a.alloc(f64, self.count()); + const mag = try a.alloc(f64, self.count()); + try self.reference(want, mag); + return verifyF32(&self.msg, self.spec(), self.shape, gate, out, want, mag, self, Finish.expected); + } + + pub fn emulate(self: *Finish, fault: Fault, out: []u8) void { + for (0..self.count()) |i| put(out, i, self.expectedAs(i, fault)); + if (fault == .writes_input) self.vals[0] += 1; + } +}; diff --git a/zig/tests/cuda/contract_flashnext_exl3_tppf_rows_ref.zig b/zig/tests/cuda/contract_flashnext_exl3_tppf_rows_ref.zig new file mode 100644 index 0000000000..8d19cb732b --- /dev/null +++ b/zig/tests/cuda/contract_flashnext_exl3_tppf_rows_ref.zig @@ -0,0 +1,370 @@ +//! Host cases of Flash Next EXL3's TP prefill rotate (tf_exl3pf_rotate) and place (tf_exl3pf_tp_place) launches. +const std = @import("std"); +const core = @import("core"); +const flash = @import("flashnext"); +const ref = @import("contract_flashnext_exl3_tppf_ref.zig"); + +const c = core.contract.spec; +const g = core.contract.gates; +const pk = flash.exl3.prefill_k; + +const Fault = ref.Fault; +const slots = ref.slots; +const dims = ref.dims; + +/// Experts of one layer's table: 15 routed and the shared one as id 15. +pub const table_n = 16; +const had32: f32 = 0.08838834764831845; +const had64: f64 = 1.0 / @sqrt(128.0); +const block = 128; + +/// The Hadamard transform of a block in the kernel's order: level h = 1, 2, ..., 64, each pair (a + b, a - b). +fn fwht(comptime T: type, v: *[block]T, fault: Fault) void { + var h: usize = 1; + while (h < block) : (h *= 2) for (0..block) |i| { + if (i & h != 0) continue; + const a = v[i]; + const b = v[i | h]; + v[i] = if (T == f32 and fault == .bf16_adds) ref.bf16(a + b) else a + b; + v[i | h] = if (T == f32 and fault == .bf16_adds) ref.bf16(a - b) else a - b; + }; +} + +/// Gate 4's inputs: Hadamard rows 3 and 5 summed (4,096 a sign) plus 64 m, m below 16: exact fp32 sums, most cancel. +fn stressX(k: usize) f32 { + const s3: f32 = if (@popCount(3 & k) % 2 == 1) -1 else 1; + const s5: f32 = if (@popCount(5 & k) % 2 == 1) -1 else 1; + return 4096 * (s3 + s5) + 64 * @as(f32, @floatFromInt((k * 5 + 3) % 16)); +} + +/// 7 adds, the product and one scale: the roundings a value passes (the 7 levels and the multiply). +fn rotateDepth(_: c.Shape) usize { + return 8; +} + +const rotate_served = [_]c.Shape{ + .{ .rows = 5, .n = dims, .k = block }, + .{ .rows = 1, .n = dims, .k = block }, + .{ .rows = 130, .n = dims, .k = block }, + .{ .rows = 1408, .n = dims, .k = block }, +}; +const max_rows = 1408; + +/// Member c of the pass: xg[c], xu[c] = fp16 of its row's bf16 x times suh[expert], rotated by the 128 Hadamard. +pub const Rotate = struct { + gpa: std.mem.Allocator, + fault: Fault = .none, + arena: ?std.heap.ArenaAllocator = null, + shape: c.Shape = .{ .rows = 0, .n = 0, .k = 0 }, + x: []u16 = &.{}, + picks: []i32 = &.{}, + plan: pk.Plan = undefined, + suh: [2][]u16 = .{ &.{}, &.{} }, + msg: [200]u8 = undefined, + + pub const n_in = 5; + pub const n_out = 2; + + pub fn deinit(self: *Rotate) void { + if (self.arena) |*a| a.deinit(); + } + + pub fn spec(self: *const Rotate) c.Spec { + const name = switch (self.fault) { + .none => "cuda Flash Next EXL3 prefill tf_exl3pf_rotate", + .nan_dropped => "control: tf_exl3pf_rotate with NaN and infinite inputs zeroed", + .writes_input => "control: tf_exl3pf_rotate writing its input", + .wrong => "control: tf_exl3pf_rotate taking the up table's signs from the gate table", + .bf16_adds => "control: tf_exl3pf_rotate with bf16 butterfly adds", + .flat => unreachable, + }; + return .{ + .name = name, + .accum = .f32, + .depth = rotateDepth, + // the normalization constant is rounded once and the product once + .ops = 2, + .out = .f16, + .nan = .{ .stated = "a NaN or infinite input reaches the 128 outputs of its block, as the host transform gives them" }, + .tiles = ref.edges, + .served = &rotate_served, + }; + } + + pub fn accepts(shape: c.Shape) bool { + return shape.rows >= 1 and shape.rows <= max_rows and shape.n >= 256 and shape.n % 256 == 0 and shape.n <= dims and shape.k == block; + } + + pub fn prepare(self: *Rotate, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (!accepts(shape)) return false; + self.deinit(); + self.arena = std.heap.ArenaAllocator.init(self.gpa); + const a = self.arena.?.allocator(); + self.shape = shape; + const rows = shape.rows; + const n = shape.n; + var prng = std.Random.DefaultPrng.init(seed ^ 0x726f7461); + const rng = prng.random(); + const stress = fill == .stress; + self.x = try a.alloc(u16, rows * n); + for (self.x, 0..) |*v, i| v.* = c.toBf16(if (stress) stressX(i % block) else (rng.float(f32) - 0.5) * 2); + if (fill == .nonfinite) { + self.x[(rows - 1) * n] = c.toBf16(std.math.nan(f32)); + if (rows > 1) self.x[(rows - 2) * n + 1] = c.toBf16(std.math.inf(f32)); + } + for (&self.suh) |*t| { + t.* = try a.alloc(u16, table_n * n); + for (t.*) |*h| { + const m: f16 = @floatCast(0.5 + rng.float(f32)); + h.* = @bitCast(if (stress) @as(f16, 1) else if (rng.boolean()) m else -m); + } + } + self.picks = try a.alloc(i32, rows * slots); + 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, rows * slots, table_n); + var counts: [table_n]u32 = undefined; + pk.build(self.picks, null, &counts, &self.plan); + return true; + } + + pub fn inputs(self: *const Rotate) [n_in][]const u8 { + return .{ std.mem.sliceAsBytes(self.x), std.mem.sliceAsBytes(self.picks), self.plan.memberBytes(), std.mem.sliceAsBytes(self.suh[0]), std.mem.sliceAsBytes(self.suh[1]) }; + } + + pub fn outputLens(self: *const Rotate) [n_out]usize { + return .{ self.plan.n_members * self.shape.n * 2, self.plan.n_members * self.shape.n * 2 }; + } + + pub fn count(self: *const Rotate) usize { + return 2 * self.plan.n_members * self.shape.n; + } + + /// The block's products x * suh (fp32 for the kernel, float64 for the reference) of member `m` in matrix `mat`. + fn products(self: *const Rotate, comptime T: type, fault: Fault, mat: usize, m: usize, blk: usize, v: *[block]T) void { + const n = self.shape.n; + const p: usize = @intCast(self.plan.members[m]); + const e: usize = @intCast(self.picks[p]); + const table = self.suh[if (fault == .wrong and mat == 1) 0 else mat]; + for (v, 0..) |*t, j| { + const col = blk * block + j; + var xv = c.fromBf16(self.x[p / slots * n + col]); + if (fault == .nan_dropped and !std.math.isFinite(xv)) xv = 0; + const sv: f32 = @as(f16, @bitCast(table[e * n + col])); + t.* = @as(T, xv) * @as(T, sv); + } + } + + /// The kernel's 128 fp16 outputs of one block, bit for bit. + fn outBlock(self: *const Rotate, fault: Fault, mat: usize, m: usize, blk: usize, bits: *[block]u16) void { + var v: [block]f32 = undefined; + self.products(f32, fault, mat, m, blk, &v); + fwht(f32, &v, fault); + for (bits, v) |*b, x| b.* = @bitCast(@as(f16, @floatCast(x * had32))); + } + + /// Float64's transform of the block's products, and the sum of their magnitudes (the bound's mag). + fn refBlock(self: *const Rotate, mat: usize, m: usize, blk: usize, want: *[block]f64, mag: *f64) void { + var t: [block]f64 = undefined; + self.products(f64, .none, mat, m, blk, &t); + mag.* = 0; + for (t) |x| mag.* += @abs(x) * had64; + fwht(f64, &t, .none); + for (want, t) |*w, x| w.* = x * had64; + } + + fn at(self: *const Rotate, mat: usize, m: usize, blk: usize) usize { + return (mat * self.plan.n_members + m) * self.shape.n + blk * block; + } + + pub fn reference(self: *const Rotate, want: []f64, mag: []f64) !void { + var w: [block]f64 = undefined; + var s: f64 = 0; + for (0..2) |mat| for (0..self.plan.n_members) |m| for (0..self.shape.n / block) |blk| { + self.refBlock(mat, m, blk, &w, &s); + @memcpy(want[self.at(mat, m, blk)..][0..block], &w); + @memset(mag[self.at(mat, m, blk)..][0..block], s); + }; + } + + /// Output `j` of a block summed in `acc` in `order`: the products with the transform's signs. + pub fn simulate(self: *const Rotate, acc: c.Accum, order: g.Order, out: []f64) !void { + var t: [block]f64 = undefined; + var s: [block]f64 = undefined; + for (0..2) |mat| for (0..self.plan.n_members) |m| for (0..self.shape.n / block) |blk| { + self.products(f64, .none, mat, m, blk, &t); + for (0..block) |j| { + for (&s, t, 0..) |*y, x, k| y.* = if (@popCount(j & k) % 2 == 1) -x else x; + out[self.at(mat, m, blk) + j] = core.contract.rows.sumIn(acc, order, &s) * had64; + } + }; + } + + /// Each output against the host's fp16 bits, float64's class and the declared bound (gate 1: finite and not). + pub fn verify(self: *Rotate, out: []const u8, gate: g.Gate) !?[]const u8 { + const s = self.spec(); + var bits: [block]u16 = undefined; + var want: [block]f64 = undefined; + var mag: f64 = 0; + var finite: usize = 0; + for (0..2) |mat| for (0..self.plan.n_members) |m| for (0..self.shape.n / block) |blk| { + self.outBlock(.none, mat, m, blk, &bits); + self.refBlock(mat, m, blk, &want, &mag); + for (bits, want, 0..) |host, w, j| { + const i = self.at(mat, m, blk) + j; + const got = std.mem.readInt(u16, out[2 * i ..][0..2], .little); + const same = got == host or (std.math.isNan(@as(f16, @bitCast(got))) and std.math.isNan(@as(f16, @bitCast(host)))); + if (!same) return try std.fmt.bufPrint(&self.msg, "output {d}: {x:0>4} where the host transform gives {x:0>4}", .{ i, got, host }); + const have: f64 = @as(f16, @bitCast(got)); + if (c.classOf(have) != c.classOf(w)) return try std.fmt.bufPrint(&self.msg, "output {d} is {t} where float64 gives {t}", .{ i, c.classOf(have), c.classOf(w) }); + if (!std.math.isFinite(w)) continue; + finite += 1; + if (@abs(have - w) > c.bound(s, self.shape, w, mag)) return try std.fmt.bufPrint(&self.msg, "output {d}: {d} is outside float64's bound around {d}", .{ i, have, w }); + } + }; + if (gate == .nonfinite and (finite == 0 or finite == self.count())) return "the planted values leave all outputs finite or none"; + return null; + } + + pub fn emulate(self: *Rotate, fault: Fault, out: []u8) void { + var bits: [block]u16 = undefined; + for (0..2) |mat| for (0..self.plan.n_members) |m| for (0..self.shape.n / block) |blk| { + self.outBlock(fault, mat, m, blk, &bits); + for (bits, 0..) |b, j| std.mem.writeInt(u16, out[2 * (self.at(mat, m, blk) + j) ..][0..2], b, .little); + }; + if (fault == .writes_input) self.x[0] +%= 1; + } +}; + +fn noSums(_: c.Shape) usize { + return 0; +} + +const place_served = [_]c.Shape{ + .{ .rows = 5, .n = dims, .k = 0 }, + .{ .rows = 1, .n = dims, .k = 0 }, + .{ .rows = 130, .n = dims, .k = 0 }, + .{ .rows = 1024, .n = dims, .k = 0 }, +}; +const max_place_rows = 1024; + +/// y[r, 0, :] = sum[r] and the weights {1, 0, ..., 0}: the other slots of y are left as they were (poison). +pub const Place = struct { + gpa: std.mem.Allocator, + fault: Fault = .none, + arena: ?std.heap.ArenaAllocator = null, + shape: c.Shape = .{ .rows = 0, .n = 0, .k = 0 }, + sum: []f32 = &.{}, + msg: [200]u8 = undefined, + + pub const n_in = 1; + pub const n_out = 2; + + pub fn deinit(self: *Place) void { + if (self.arena) |*a| a.deinit(); + } + + pub fn spec(self: *const Place) c.Spec { + const name = switch (self.fault) { + .none => "cuda Flash Next EXL3 prefill tf_exl3pf_tp_place", + .nan_dropped => "control: tf_exl3pf_tp_place with NaN zeroed", + .writes_input => "control: tf_exl3pf_tp_place writing its input", + .wrong => "control: tf_exl3pf_tp_place leaving every weight 1", + else => unreachable, + }; + return .{ + .name = name, + .accum = .none, + .depth = noSums, + .out = .raw, + .nan = .{ .stated = "y[r, 0] holds the branch's bits (NaN payloads and infinities included), weights are 1, 0, ..., 0" }, + .tiles = ref.edges, + .served = &place_served, + }; + } + + pub fn accepts(shape: c.Shape) bool { + return shape.rows >= 1 and shape.rows <= max_place_rows and shape.n >= 256 and shape.n % 256 == 0 and shape.n <= dims; + } + + pub fn prepare(self: *Place, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (!accepts(shape)) return false; + self.deinit(); + self.arena = std.heap.ArenaAllocator.init(self.gpa); + self.shape = shape; + self.sum = try self.arena.?.allocator().alloc(f32, shape.rows * shape.n); + var prng = std.Random.DefaultPrng.init(seed ^ 0x706c6163); + const rng = prng.random(); + for (self.sum) |*v| v.* = rng.floatNorm(f32) * 3; + // a row each of NaN (with a payload), +Inf and -Inf every 7th column + if (fill == .nonfinite) for (self.sum, 0..) |*v, i| if (i % 7 == 0) { + v.* = switch ((i / shape.n) % 3) { + 0 => @bitCast(@as(u32, 0x7fc01234)), + 1 => std.math.inf(f32), + else => -std.math.inf(f32), + }; + }; + return true; + } + + pub fn inputs(self: *const Place) [n_in][]const u8 { + return .{std.mem.sliceAsBytes(self.sum)}; + } + + fn yLen(self: *const Place) usize { + return self.shape.rows * slots * self.shape.n * 4; + } + + pub fn outputLens(self: *const Place) [n_out]usize { + return .{ self.yLen(), self.shape.rows * slots * 4 }; + } + + pub fn count(self: *const Place) usize { + return self.yLen() + self.shape.rows * slots * 4; + } + + /// The bytes y and the weights hold after a launch (a planted fault's too); everything else keeps the poison. + fn build(self: *const Place, fault: Fault, out: []u8) void { + @memset(out, 0xff); + const n = self.shape.n; + for (0..self.shape.rows) |r| for (0..n) |col| { + const v = self.sum[r * n + col]; + ref.put(out, r * slots * n + col, if (fault == .nan_dropped and std.math.isNan(v)) 0 else v); + }; + for (0..self.shape.rows * slots) |i| ref.put(out, self.yLen() / 4 + i, if (i % slots == 0 or fault == .wrong) 1 else 0); + } + + pub fn reference(self: *Place, want: []f64, mag: []f64) !void { + const bytes = try self.arena.?.allocator().alloc(u8, self.count()); + self.build(.none, bytes); + for (want, mag, bytes) |*w, *m, b| { + w.* = @floatFromInt(b); + m.* = w.*; + } + } + + pub fn simulate(_: *const Place, _: c.Accum, _: g.Order, _: []f64) !void { + return error.NoSums; + } + + /// The branch's bytes in y[r, 0, :], the weights, and poison elsewhere (gate 1: a non-finite value shows). + pub fn verify(self: *Place, out: []const u8, gate: g.Gate) !?[]const u8 { + const want = try self.arena.?.allocator().alloc(u8, self.count()); + self.build(.none, want); + if (std.mem.indexOfDiff(u8, out, want)) |at| return try std.fmt.bufPrint(&self.msg, "byte {d} of {d} is {x:0>2} where the branch's copy has {x:0>2}", .{ at, out.len, out[at], want[at] }); + if (gate != .nonfinite) return null; + var bad: usize = 0; + for (self.sum) |v| bad += @intFromBool(!std.math.isFinite(v)); + return if (bad == 0 or bad == self.sum.len) "the planted values leave all of the branch finite or none" else null; + } + + pub fn emulate(self: *Place, fault: Fault, out: []u8) void { + self.build(fault, out); + if (fault == .writes_input) self.sum[0] += 1; + } +}; diff --git a/zig/tests/cuda/contract_flashnext_exl3_tppf_test.zig b/zig/tests/cuda/contract_flashnext_exl3_tppf_test.zig new file mode 100644 index 0000000000..a7d31d88e7 --- /dev/null +++ b/zig/tests/cuda/contract_flashnext_exl3_tppf_test.zig @@ -0,0 +1,254 @@ +//! The TP prefill launches' host cases through the real gates: the shipped emulation passes, faults go red. +const std = @import("std"); +const core = @import("core"); +const flash = @import("flashnext"); +const ref = @import("contract_flashnext_exl3_tppf_ref.zig"); +const rows_ref = @import("contract_flashnext_exl3_tppf_rows_ref.zig"); + +const c = core.contract.spec; +const g = core.contract.gates; +const canon = flash.exl3_canon; +const t = std.testing; + +/// A device-less rig: the buffer the emulated launch writes and the inputs' bytes as the last launch left them. +const FakeRig = struct { + out: []u8 = &.{}, + seen: std.ArrayList(u8) = .empty, + + pub fn outputBytes(r: *FakeRig) ![]const u8 { + return r.out; + } + + pub fn inputBytes(r: *FakeRig) ![]const u8 { + return r.seen.items; + } + + pub fn intact(_: *FakeRig) !bool { + return true; + } + + pub fn overlap(_: *FakeRig) bool { + return false; + } +}; + +fn Fake(comptime Case: type) type { + return struct { + const Self = @This(); + rig: FakeRig = .{}, + case: Case, + small: []const c.Shape, + + fn deinit(self: *Self) void { + t.allocator.free(self.rig.out); + self.rig.seen.deinit(t.allocator); + self.case.deinit(); + } + + fn snapshot(self: *Self) !void { + self.rig.seen.clearRetainingCapacity(); + for (self.case.inputs()) |bytes| try self.rig.seen.appendSlice(t.allocator, bytes); + } + + const vt = g.hooks(Self); + + fn kernel(self: *Self) g.Kernel { + var k: g.Kernel = .{ .spec = self.case.spec(), .ctx = self, .vt = &vt, .control = ref.controlGate(self.case.fault) }; + k.spec.served = self.small; + return k; + } + + pub fn prepare(self: *Self, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (!try self.case.prepare(shape, fill, seed)) return false; + var total: usize = 0; + for (self.case.outputLens()) |len| total += len; + t.allocator.free(self.rig.out); + self.rig.out = try t.allocator.alloc(u8, total); + @memset(self.rig.out, 0xff); + try self.snapshot(); + return true; + } + + pub fn run(self: *Self) !void { + self.case.emulate(self.case.fault, self.rig.out); + try self.snapshot(); + } + + pub fn count(self: *Self) usize { + return self.case.count(); + } + + pub fn reference(self: *Self, want: []f64, mag: []f64) !void { + return self.case.reference(want, mag); + } + + pub fn simulate(self: *Self, acc: c.Accum, order: g.Order, out: []f64) !void { + return self.case.simulate(acc, order, out); + } + + pub fn check(self: *Self, gate: g.Gate) !?[]const u8 { + return self.case.verify(self.rig.out, gate); + } + }; +} + +/// The shipped case passes gates 1 to 4 (or N/A); a control goes red at its own gate. +fn gates(comptime Case: type, case: Case, small: []const c.Shape) !void { + var f: Fake(Case) = .{ .case = case, .small = small }; + defer f.deinit(); + const k = f.kernel(); + g.quiet = true; + defer g.quiet = false; + if (ref.controlGate(case.fault)) |gate| { + const v = try g.one(t.allocator, gate, k, 0x51); + if (v != .fail) std.debug.print("control {s} stayed {t} at {t}\n", .{ k.spec.name, v, gate }); + return t.expectEqual(g.Verdict.fail, v); + } + for ([_]g.Gate{ .nonfinite, .repeat, .shapes, .sums }) |gate| { + const v = try g.one(t.allocator, gate, k, 0x51); + if (v == .fail) { + g.quiet = false; + _ = try g.one(t.allocator, gate, k, 0x51); + return error.TestUnexpectedResult; + } + } +} + +const w = 256; +const combine_small = [_]c.Shape{ .{ .rows = 5, .n = w, .k = 2 }, .{ .rows = 1, .n = w, .k = 1 }, .{ .rows = 17, .n = w, .k = 3 }, .{ .rows = 40, .n = w, .k = 5 }, .{ .rows = 6, .n = w, .k = 8 } }; +const blocks_small = [_]c.Shape{ .{ .rows = 300, .n = w, .k = 2 }, .{ .rows = 700, .n = w, .k = 5 }, .{ .rows = 257, .n = w, .k = 8 } }; +const finish_small = [_]c.Shape{ .{ .rows = 5, .n = w, .k = 2 }, .{ .rows = 44, .n = w, .k = 2 }, .{ .rows = 7, .n = w, .k = 3 }, .{ .rows = 9, .n = w, .k = 5 }, .{ .rows = 6, .n = w, .k = 8 } }; +const rotate_small = [_]c.Shape{ .{ .rows = 5, .n = w, .k = 128 }, .{ .rows = 1, .n = w, .k = 128 }, .{ .rows = 30, .n = w, .k = 128 } }; +const place_small = [_]c.Shape{ .{ .rows = 5, .n = w, .k = 0 }, .{ .rows = 1, .n = w, .k = 0 }, .{ .rows = 130, .n = w, .k = 0 } }; + +test "tppf contract: combine passes gates 1 to 4 in both region layouts and its four controls go red" { + for ([_]bool{ false, true }) |blocks| { + const small: []const c.Shape = if (blocks) &blocks_small else &combine_small; + try gates(ref.Combine, .{ .gpa = t.allocator, .blocks = blocks }, small); + } + for ([_]ref.Fault{ .flat, .nan_dropped, .writes_input, .bf16_adds }) |f| try gates(ref.Combine, .{ .gpa = t.allocator, .blocks = false, .fault = f }, &combine_small); +} + +test "tppf contract: finish passes gates 1 to 4 and its three controls go red" { + try gates(ref.Finish, .{ .gpa = t.allocator }, &finish_small); + for ([_]ref.Fault{ .flat, .bf16_adds, .writes_input }) |f| try gates(ref.Finish, .{ .gpa = t.allocator, .fault = f }, &finish_small); +} + +test "tppf contract: place passes gates 1 to 3 (4 is N/A) and its three controls go red" { + try gates(rows_ref.Place, .{ .gpa = t.allocator }, &place_small); + for ([_]ref.Fault{ .nan_dropped, .writes_input, .wrong }) |f| try gates(rows_ref.Place, .{ .gpa = t.allocator, .fault = f }, &place_small); +} + +test "tppf contract: rotate passes gates 1 to 4 and its four controls go red" { + try gates(rows_ref.Rotate, .{ .gpa = t.allocator }, &rotate_small); + for ([_]ref.Fault{ .nan_dropped, .writes_input, .wrong, .bf16_adds }) |f| try gates(rows_ref.Rotate, .{ .gpa = t.allocator, .fault = f }, &rotate_small); +} + +/// Every shape the registration lists and its edges are ones the case takes; gates 1, 2 and 4 find their shapes in it. +fn takes(comptime Case: type, case: Case, ranked: bool) !void { + const s = case.spec(); + const list = try c.shapes(t.allocator, s); + defer t.allocator.free(list); + for (list) |x| try t.expect(Case.accepts(x)); + var widest: usize = 0; + for (s.served) |x| widest = @max(widest, x.k); + if (!ranked) return; + try t.expectEqual(@as(usize, 2), s.served[0].k); + try t.expectEqual(@as(usize, 8), widest); +} + +test "tppf contract: every served and edge shape of every registration is one its case takes" { + try takes(ref.Combine, .{ .gpa = t.allocator, .blocks = false }, true); + try takes(ref.Combine, .{ .gpa = t.allocator, .blocks = true }, true); + try takes(ref.Finish, .{ .gpa = t.allocator }, true); + try takes(rows_ref.Rotate, .{ .gpa = t.allocator }, false); + try takes(rows_ref.Place, .{ .gpa = t.allocator }, false); +} + +fn randomRow(rng: std.Random, y: *[ref.slots]f32, wt: *[ref.slots]f32, p: *[ref.slots]u16) void { + for (0..ref.slots) |k| { + y[k] = rng.floatNorm(f32) * 0.05; + wt[k] = rng.floatNorm(f32) * 0.2; + p[k] = if (k == ref.slots - 1) 512 else if (rng.boolean()) rng.uintLessThan(u16, 512) else 256 + rng.uintLessThan(u16, 24); + } +} + +test "tppf contract: the block oracle is canon.subtree, and a masked slot is the slot with weight 0" { + var prng = std.Random.DefaultPrng.init(0x7b10c); + const rng = prng.random(); + for (0..400) |_| { + var y: [ref.slots]f32 = undefined; + var wt: [ref.slots]f32 = undefined; + var pu: [ref.slots]u16 = undefined; + randomRow(rng, &y, &wt, &pu); + var p: [ref.slots]i32 = undefined; + for (&p, pu) |*d, v| d.* = v; + for (1..9) |world| for (0..world) |rank| { + var bl: [canon.max_blocks]canon.Block = undefined; + const nb = ref.rankBlocks(&bl, world, rank); + for (bl[0..nb]) |b| try t.expectEqual(@as(u32, @bitCast(canon.subtree(&y, &wt, &pu, b.start, b.log))), @as(u32, @bitCast(ref.blockSum(&y, &wt, &p, b.start, b.log, .none)))); + }; + const k = rng.uintLessThan(usize, ref.slots); + var masked = p; + masked[k] = -1; + var zeroed = wt; + zeroed[k] = 0; + try t.expectEqual(@as(u32, @bitCast(canon.subtree(&y, &zeroed, &pu, 0, canon.leaf_log))), @as(u32, @bitCast(ref.blockSum(&y, &wt, &masked, 0, canon.leaf_log, .none)))); + } +} + +/// Combine's gather buffer of `world` ranks, each region finished: the rows the finish gives. +fn gathered(world: usize, blocks: bool, shape: c.Shape, seed: u64, out: []f32) !void { + var cb: ref.Combine = .{ .gpa = t.allocator, .blocks = blocks }; + defer cb.deinit(); + var sh = shape; + sh.k = world; + try t.expect(try cb.prepare(sh, .random, seed)); + const gath = try t.allocator.alloc(f32, cb.count()); + defer t.allocator.free(gath); + cb.emulate(.none, std.mem.sliceAsBytes(gath)); + var at: usize = 0; + var r0: usize = 0; + while (r0 < sh.rows) : (r0 += cb.sub) { + const rows = @min(cb.sub, sh.rows - r0); + const region = rows * world * cb.m * sh.n; + var fin: ref.Finish = .{ .gpa = t.allocator }; + defer fin.deinit(); + fin.shape = .{ .rows = rows, .n = sh.n, .k = world }; + fin.m = cb.m; + fin.np = canon.pieces(&fin.ps, world); + fin.vals = gath[at..][0..region]; + at += region; + fin.emulate(.none, std.mem.sliceAsBytes(out[r0 * sh.n ..][0 .. rows * sh.n])); + } +} + +test "tppf contract: the ranks' partials gathered and finished equal the single rank's tree at every degree" { + const shape: c.Shape = .{ .rows = 300, .n = w, .k = 1 }; + const want = try t.allocator.alloc(f32, shape.rows * w); + defer t.allocator.free(want); + const got = try t.allocator.alloc(f32, shape.rows * w); + defer t.allocator.free(got); + for ([_]bool{ false, true }) |blocks| { + try gathered(1, blocks, shape, 0x9a7, want); + for (2..9) |world| { + try gathered(world, blocks, shape, 0x9a7, got); + for (want, got, 0..) |a, b, i| if (@as(u32, @bitCast(a)) != @as(u32, @bitCast(b))) { + std.debug.print("world {d} blocks {}: output {d}: {x:0>8} for {x:0>8}\n", .{ world, blocks, i, @as(u32, @bitCast(b)), @as(u32, @bitCast(a)) }); + return error.TestUnexpectedResult; + }; + } + } +} + +test "tppf contract: rotate's float64 transform is the signed sum of the products" { + var r: rows_ref.Rotate = .{ .gpa = t.allocator }; + defer r.deinit(); + try t.expect(try r.prepare(.{ .rows = 3, .n = w, .k = 128 }, .random, 0x5eed)); + const n = r.count(); + const buf = try t.allocator.alloc(f64, 3 * n); + defer t.allocator.free(buf); + try r.reference(buf[0..n], buf[n .. 2 * n]); + try r.simulate(.none, .sequential, buf[2 * n ..]); + for (buf[0..n], buf[2 * n ..]) |a, b| try t.expect(@abs(a - b) <= 1e-12 * @max(1, @abs(a))); +} diff --git a/zig/tests/cuda/contract_limits.zig b/zig/tests/cuda/contract_limits.zig index 47b185fab3..04116b8751 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 &.{}; return error.UnknownContractImage; 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..b2aa1b9ff6 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,42 @@ 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 }); + } + if (tiles_ok and exl3.planOf(w).sk != 1 and (k / 16) % 4 == 0) { + var fixed = w; + fixed.sk = 1; + fixed.wk = 4; + const rows: usize = 17; + de.implementation = .reference; + try de.decode(s, fixed, x.ptr, .bf16, rows, yr.ptr, .bf16); + try de.promptExact(s, fixed, x.ptr, rows, yt.ptr, .bf16); + try s.synchronize(); + const reference = try a.alloc(u16, rows * n); + defer a.free(reference); + const actual = try a.alloc(u16, rows * n); + defer a.free(actual); + try yr.download(0, std.mem.sliceAsBytes(reference)); + try yt.download(0, std.mem.sliceAsBytes(actual)); + try check.expect(std.mem.eql(u16, reference, actual), "{s}: fixed-plan prompt tiles differ from decode", .{name}); } } 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..0177901a1c --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/prefill/exl3_prefill_engine.zig @@ -0,0 +1,361 @@ +//! 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); +} + +pub 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); + try flash.exl3.kv.setup(g.e, true); // the server's load hook: EXL3 scratch format, chunk and head rows + 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); + } + + 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 }; + } + + 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", .{}); +} + +pub 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. +pub 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_segments.zig b/zig/tests/cuda/flashnext_exl3/prefill/exl3_prefill_segments.zig new file mode 100644 index 0000000000..56e577968d --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/prefill/exl3_prefill_segments.zig @@ -0,0 +1,274 @@ +//! exl3-prefill mixed|resume: deep and shallow segments side by side, and a state kept at any row then resumed. + +const std = @import("std"); +const cuda = @import("cuda"); +const check = @import("../../check.zig"); +const flash = @import("flashnext"); +const engine = @import("exl3_prefill_engine.zig"); +const Gpu = check.Gpu; +const Eng = engine.Eng; + +const xpre = flash.exl3.prefill; +const prompt = flash.prompt; +const vocab = prompt.vocabulary; +const max_segments = 256; +const n_tokens = 24; + +/// Shallow segments (under the deep path's 512 rows): a history mark's 7 rows, the longest that stays shallow, another. +const shorts = [_]usize{ 7, 511, 100 }; + +/// One prompt's cuts: the rows of each forward, and whether it took the deep path. +const Plan = struct { + len: [max_segments]usize = undefined, + deep: [max_segments]bool = undefined, + n: usize = 0, + + fn add(p: *Plan, rows: usize) !void { + if (p.n == max_segments) return error.TooManySegments; + p.len[p.n] = rows; + p.n += 1; + } + + /// Segments of `step` rows over [from, to): the server's steps (`stop = at + step`), the last one the remainder. + fn steps(from: usize, to: usize, step: usize) !Plan { + var p: Plan = .{}; + var at = from; + while (at < to) : (at += p.len[p.n - 1]) try p.add(@min(step, to - at)); + return p; + } + + /// Deep segments of `step` rows, a shallow one between, from either; `mark`: the last 7 rows cut. + fn mixed(n: usize, step: usize, shallow_first: bool, mark: bool) !Plan { + var p: Plan = .{}; + const end = if (mark) n - shorts[0] else n; + var at: usize = 0; + var is_deep = !shallow_first; + var k: usize = 0; + while (at < end) : (is_deep = !is_deep) { + const want = if (is_deep) step else shorts[k % shorts.len]; + k += @intFromBool(!is_deep); + try p.add(@min(want, end - at)); + at += p.len[p.n - 1]; + } + if (mark) try p.add(shorts[0]); + return p; + } + + fn count(p: *const Plan, deep: bool) usize { + var c: usize = 0; + for (p.deep[0..p.n]) |d| c += @intFromBool(d == deep); + return c; + } + + fn show(p: *const Plan) void { + for (p.len[0..p.n], p.deep[0..p.n]) |rows, d| std.debug.print(" {d}{s}", .{ rows, if (d) "d" else "s" }); + } +}; + +/// What one prompt gave: the last row's logits and the greedy tokens after it. +const Run = struct { + logits: []u16, + tokens: [n_tokens]u32 = undefined, + + fn init(a: std.mem.Allocator) !Run { + return .{ .logits = try a.alloc(u16, vocab) }; + } + + fn take(r: *Run, g: *Eng) !void { + @memcpy(std.mem.sliceAsBytes(r.logits), g.seq.logits); + try engine.greedy(g, &r.tokens); + } + + fn verdict(r: Run, want: Run) struct { logits: bool, tokens: bool } { + return .{ + .logits = std.mem.eql(u8, std.mem.sliceAsBytes(r.logits), std.mem.sliceAsBytes(want.logits)), + .tokens = std.mem.eql(u32, &r.tokens, &want.tokens), + }; + } + + /// Print how this run compares with `want`; true when both the logits bytes and the tokens are equal. + fn against(r: Run, want: Run, name: []const u8) bool { + const v = r.verdict(want); + std.debug.print("; vs {s}: logits {s}, {d} tokens {s}", .{ name, if (v.logits) "bit-equal" else "DIFFERENT", n_tokens, if (v.tokens) "identical" else "DIFFERENT" }); + return v.logits and v.tokens; + } +}; + +/// The options the next forwards read: the deep path on or off, its chunk, the cache's grid. +fn setOptions(deep_on: bool, step: usize, grid: bool) void { + var o = xpre.optionsFromEnv(); + o.enabled = deep_on; + o.rows = step; + o.snap_grid = grid; + xpre.state.options = o; +} + +/// `ids` from `at.*` through the engine, one forward a segment as the server steps a prompt; the last keeps its logits. +fn feed(g: *Eng, ids: []const u32, at: *usize, plan: *Plan) !void { + for (plan.len[0..plan.n], plan.deep[0..plan.n]) |rows, *deep| { + const stop = at.* + rows; + if (stop > ids.len) return error.UnexpectedTensor; + const before = xpre.state.stats.chunks; + try g.e.forward(flash.Tri, g.tri(), &g.seq, ids[at.*..stop], .{ .logits = if (stop == ids.len) .last else .none, .prompt = true }); + deep.* = xpre.state.stats.chunks != before; + at.* = stop; + } +} + +/// A fresh sequence through `plan`'s forwards, then the greedy tokens. +fn fresh(g: *Eng, ids: []const u32, plan: *Plan, out: *Run) !void { + try g.reset(); + var at: usize = 0; + try feed(g, ids, &at, plan); + try out.take(g); +} + +/// Every prompt row alone through a decode round, forced to the prompt's tokens (what the server's rounds run). +fn serial(g: *Eng, ids: []const u32, out: *Run) !void { + try g.reset(); + var buf = try cuda.DeviceBuffer.alloc(g.gpu.d, vocab * 2); + defer buf.free(); + for (ids) |id| { + var tok = id; + const q = [_]*prompt.Seq{&g.seq}; + try g.e.round(flash.Tri, g.tri(), &q, (&tok)[0..1], buf.ptr); + } + try out.take(g); +} + +fn parseList(text: []const u8, out: []isize) !usize { + var n: usize = 0; + var it = std.mem.splitScalar(u8, text, ','); + while (it.next()) |tok| : (n += 1) { + if (n == out.len) return error.TooManyValues; + out[n] = try std.fmt.parseInt(isize, tok, 10); + } + return n; +} + +fn parseLines(rest: []const [:0]const u8, at: usize) !?usize { + if (rest.len <= at or std.mem.eql(u8, rest[at], "-")) return null; + return try std.fmt.parseInt(usize, rest[at], 10); +} + +/// Shallow segments (under 512 rows) beside deep ones at each step size equal the batched pass and serial decode. +pub fn mixedRun(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 engine.readIds(gpu, rest[0]); + defer a.free(all_ids); + const n = try std.fmt.parseInt(usize, rest[1], 10); + if (n <= shorts[0] or n > all_ids.len) return error.NotEnoughIds; + var sizes: [8]isize = undefined; + const n_sizes = try parseList(rest[2], &sizes); + var max_step: usize = 0; + for (sizes[0..n_sizes]) |s| { + if (s < 512) return error.StepBelowDeepRows; + max_step = @max(max_step, @as(usize, @intCast(s))); + } + const g = try Eng.init(gpu, dir, try parseLines(rest, 3), n + 64, @min(max_step, (n + 255) / 256 * 256)); + defer g.deinit(); + const ids = all_ids[0..n]; + var batched = try Run.init(a); + defer a.free(batched.logits); + var forced = try Run.init(a); + defer a.free(forced.logits); + var got = try Run.init(a); + defer a.free(got.logits); + setOptions(false, 0, true); + var whole = try Plan.steps(0, n, n); + try fresh(g, ids, &whole, &batched); + try serial(g, ids, &forced); + std.debug.print("MIXED reference, {d} tokens: forced serial decode of every row", .{n}); + var all_equal = forced.against(batched, "the batched pass"); + std.debug.print("\n", .{}); + var cases: usize = 0; + const orders = [_]struct { shallow_first: bool, mark: bool }{ .{ .shallow_first = false, .mark = false }, .{ .shallow_first = true, .mark = false }, .{ .shallow_first = false, .mark = true } }; + for (sizes[0..n_sizes]) |size| for (orders) |order| { + const step: usize = @intCast(size); + setOptions(true, step, true); + var plan = try Plan.mixed(n, step, order.shallow_first, order.mark); + try fresh(g, ids, &plan, &got); + std.debug.print("MIXED {d} tokens, steps of {d} rows, cuts", .{ n, step }); + plan.show(); + const eq_batched = got.against(batched, "the batched pass"); + const eq_serial = got.against(forced, "forced serial decode"); + std.debug.print("; {d} deep, {d} shallow segments\n", .{ plan.count(true), plan.count(false) }); + try check.expect(plan.count(true) > 0 and plan.count(false) > 0, "the cut at steps of {d} rows mixes deep and shallow segments", .{step}); + all_equal = all_equal and eq_batched and eq_serial; + cases += 1; + }; + xpre.state.options = null; + try check.expect(all_equal, "mixed deep and shallow segments equal the batched pass and forced serial decode", .{}); + check.pass("{d} mixed cuts at {d} step sizes equal the batched pass and forced serial decode in logits bytes and tokens", .{ cases, n_sizes }); +} + +/// A state kept at any row (snap grid off), cold-continued or resumed, equals a fresh prompt; a cut < 0 is rows left. +pub fn resumeRun(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 engine.readIds(gpu, rest[0]); + defer a.free(all_ids); + const n = try std.fmt.parseInt(usize, rest[1], 10); + if (n < 2 or n > all_ids.len) return error.NotEnoughIds; + var sizes: [8]isize = undefined; + const n_sizes = try parseList(rest[2], &sizes); + var cuts: [16]isize = undefined; + const n_cuts = try parseList(rest[3], &cuts); + var max_step: usize = 0; + for (sizes[0..n_sizes]) |s| { + if (s <= 0) return error.BadStep; + max_step = @max(max_step, @as(usize, @intCast(s))); + } + const g = try Eng.init(gpu, dir, try parseLines(rest, 4), n + 64, @min(max_step, (n + 255) / 256 * 256)); + defer g.deinit(); + const ids = all_ids[0..n]; + var want = try Run.init(a); + defer a.free(want.logits); + var cold = try Run.init(a); + defer a.free(cold.logits); + var resumed = try Run.init(a); + defer a.free(resumed.logits); + var all_equal = true; + var cases: usize = 0; + for (sizes[0..n_sizes]) |size| { + const step: usize = @intCast(size); + setOptions(true, step, true); + const grid_on = xpre.snapGrid(g.e, g.e.chunk); + setOptions(true, step, false); + try check.expect(xpre.snapGrid(g.e, g.e.chunk) == 0, "the engine's snap grid is off with the option off", .{}); + std.debug.print("RESUME steps of {d} rows: snap grid {d} by default, 0 (any row) now\n", .{ step, grid_on }); + var whole = try Plan.steps(0, n, step); + try fresh(g, ids, &whole, &want); + for (cuts[0..n_cuts]) |c| { + const mark: usize = if (c < 0) n - @as(usize, @intCast(-c)) else @intCast(c); + if (mark == 0 or mark >= n) return error.CutOutsidePrompt; + var head = try Plan.steps(0, mark, step); + var tail = try Plan.steps(mark, n, step); + // the cold pass cuts at the mark (the history mark's last chunk), keeps the state, and goes on + try g.reset(); + var at: usize = 0; + try feed(g, ids, &at, &head); + const saved = try flash.snapshot.save(a, g.gpu.d, g.stream, &g.seq, null); + defer flash.snapshot.drop(a, g.stream, saved); + try feed(g, ids, &at, &tail); + try cold.take(g); + // the resumed pass: nothing of the prompt left on the device but the kept state + try g.reset(); + try flash.snapshot.restore(g.stream, &g.seq, null, saved); + try check.expect(g.seq.fed == mark, "the kept state resumes at row {d}", .{mark}); + at = mark; + try feed(g, ids, &at, &tail); + try resumed.take(g); + std.debug.print("RESUME {d} tokens, steps of {d} rows, state at row {d} ({d} past a chunk start, {d} rows after it, last chunk {s})", .{ n, step, mark, mark % step, n - mark, if (tail.deep[tail.n - 1]) "deep" else "shallow" }); + const eq_cold = cold.against(want, "fresh (cold pass with the mark)"); + const eq_resumed = resumed.against(want, "fresh (resumed pass)"); + std.debug.print("\n", .{}); + all_equal = all_equal and eq_cold and eq_resumed; + cases += 1; + } + } + xpre.state.options = null; + try check.expect(all_equal, "a state kept at any row, resumed or continued, equals a fresh prompt", .{}); + check.pass("{d} states kept at any row (snap grid off) resume to a fresh prompt's logits bytes and tokens", .{cases}); +} 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..d8f1fd5c0a 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,28 @@ 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 segments = @import("exl3_prefill_segments.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..]); + if (std.mem.eql(u8, mode, "mixed")) return segments.mixedRun(gpu, dir, rest[2..]); + if (std.mem.eql(u8, mode, "resume")) return segments.resumeRun(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 +44,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..5b7f85a443 --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/tp/exl3_tpprefill_tests.zig @@ -0,0 +1,340 @@ +//! 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); + try flash.exl3.kv.setup(g.e, true); // the server's load hook: EXL3 scratch format, chunk and head rows + 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 e5f0e39a2c..6e3cc52b4a 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -108,6 +108,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"); @@ -254,9 +255,16 @@ 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-prefill DIR mixed IDS N STEP[,STEP] [LINES] prompt cut into deep and shallow segments vs the batched pass and serial decode: logits bytes, tokens + \\ exl3-prefill DIR resume IDS N STEP[,STEP] CUT[,CUT] [LINES] state kept at any row (snap grid off), resumed vs a fresh prompt; CUT<0: rows left + \\ 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) + \\ exl3-promptd-contract the prompt tiles in decode's bytes (tf_exl3_promptd_*) alone through the contract + \\ exl3-tppf-contract the TP prefill's rotate, combine, finish and place launches alone through the contract \\ 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 @@ -379,6 +387,9 @@ 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, "exl3-promptd-contract")) return @import("contract_flashnext_exl3_promptd.zig").run(gpu); + if (std.mem.eql(u8, cmd, "exl3-tppf-contract")) return @import("contract_flashnext_exl3_tppf.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); @@ -445,6 +456,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-shapes")) return exl3_plan_bench.profileShapes(gpu, if (rest.len > 0) rest[0] else null); 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 d130763ece6bf3e6dae63792b858428a495c2de0 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 | 8 +-- 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(+), 24 deletions(-) diff --git a/zig/src/families/flashnext/cuda_kv_lanes.zig b/zig/src/families/flashnext/cuda_kv_lanes.zig index 31da875785..d2131c7a11 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 574f7a3482..87f35fea95 100644 --- a/zig/src/families/flashnext/cuda_lanes.zig +++ b/zig/src/families/flashnext/cuda_lanes.zig @@ -48,9 +48,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); @@ -305,7 +303,7 @@ pub fn Namespace(comptime Tri: type) type { // a contiguous lane's fixed-size buffers from the parts kept from an earlier stream var part_seq: ?cuda.DeviceBuffer = null; if (self.parts.on and self.ranks == null and from != .pool) part_seq = Parts.take(&self.parts.seqs, self.engine.driver, seq_parts.seqPartsBytes()) catch null; - gop.value_ptr.* = .{ .seq = (if (self.ranks) |r| r.create(cap) else if (from == .pool) self.kv.pool.?.open(cap) else if (part_seq) |b| seq_parts.newSeqIn(self.engine, cap, b) 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 if (part_seq) |b| seq_parts.newSeqIn(self.engine, cap, b) else self.engine.newSeq(cap)) catch |err| { if (part_seq) |b| Parts.give(&self.parts.seqs, self.gpa, b); self.lanes.removeByPtr(gop.key_ptr); return err; diff --git a/zig/src/families/flashnext/cuda_lanes_explain.zig b/zig/src/families/flashnext/cuda_lanes_explain.zig index 386bd56044..0ea9b4a3da 100644 --- a/zig/src/families/flashnext/cuda_lanes_explain.zig +++ b/zig/src/families/flashnext/cuda_lanes_explain.zig @@ -18,7 +18,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 664ad0498e..0aad0755db 100644 --- a/zig/src/families/flashnext/cuda_lanes_load.zig +++ b/zig/src/families/flashnext/cuda_lanes_load.zig @@ -30,8 +30,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(); - // the pool's pages are not wired to EXL3's K/V - if (page_tokens > 0 and (o.tp != 1 or o.slide or captured != null or mode == .exl3)) return error.FlashNextPagesUnsupported; + if (page_tokens > 0 and (o.slide or captured != null)) return error.FlashNextPagesUnsupported; // pending DeltaNet rows run on one GPU only: TF_FLASHNEXT_KEEP_LATE=1 with tensor parallel is refused const Draft = @import("cuda_mtp_draft.zig").Draft; const late_fits = o.tp == 1; @@ -118,8 +117,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); // 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 ba8b914f29..9d269a63f9 100644 --- a/zig/src/families/flashnext/cuda_mtp_draft.zig +++ b/zig/src/families/flashnext/cuda_mtp_draft.zig @@ -268,6 +268,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 41ac2a3848..1adac83ad3 100644 --- a/zig/src/families/flashnext/cuda_mtp_ops.zig +++ b/zig/src/families/flashnext/cuda_mtp_ops.zig @@ -138,7 +138,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; @@ -281,7 +281,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 c1ec9abf19..f2bcb4e330 100644 --- a/zig/src/families/flashnext/cuda_prompt.zig +++ b/zig/src/families/flashnext/cuda_prompt.zig @@ -486,7 +486,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 1adaad17c0..120172f228 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 { @@ -45,7 +49,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()) }; } /// How long a rank waits for its peer: a minute with a given digest, 30 when the peer may be hashing every byte. @@ -65,6 +69,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), @@ -123,7 +129,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); @@ -172,12 +178,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; @@ -188,12 +213,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; @@ -320,8 +345,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); }, @@ -376,7 +401,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 => { @@ -433,7 +458,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); @@ -472,10 +497,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 fd848a9b959491a88daefb8e3f2517b500b1c279 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 d8b47bc7c4..967551f1be 100644 --- a/zig/src/core/kv_pool.zig +++ b/zig/src/core/kv_pool.zig @@ -152,6 +152,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 }; @@ -170,6 +176,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); @@ -230,6 +237,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; } @@ -266,6 +274,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; } @@ -441,6 +450,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; } @@ -466,8 +476,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 81baafa111790617f3a946589eee1828998dc717 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 | 17 +- docs/recipes/flash-next-cuda.md | 4 + zig/build/cuda_kernels.zig | 2 + 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 | 315 ++++++++++ .../flashnext/exl3/model/cuda_exl3.zig | 3 + zig/tests/cuda/contract_arch_test.zig | 2 + .../cuda/contract_flashnext_exl3_kvt.zig | 480 +++++++++++++++ .../cuda/contract_flashnext_exl3_kvt_plan.zig | 553 ++++++++++++++++++ .../contract_flashnext_exl3_kvt_plan_test.zig | 50 ++ .../cuda/contract_flashnext_exl3_suite.zig | 16 + zig/tests/cuda/contract_fn_pages.zig | 22 +- zig/tests/cuda/contract_fn_pages_media.zig | 54 +- 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 | 4 + 42 files changed, 2539 insertions(+), 172 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/contract_flashnext_exl3_kvt.zig create mode 100644 zig/tests/cuda/contract_flashnext_exl3_kvt_plan.zig create mode 100644 zig/tests/cuda/contract_flashnext_exl3_kvt_plan_test.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 251a6b6b23..54185031c5 100644 --- a/docs/kernel-contract.md +++ b/docs/kernel-contract.md @@ -337,6 +337,21 @@ CUDA also checks Nemotron's pocket (`nemotron_ops.cu`, the learned neurons the l shared expert): relu^2 up into act rows and down added into the shared expert's slot, a block row a row and tiled, the down to fp32 and bf16, against float64, the neurons past the count NaN so a read of one shows (two controls: the up's relu by `fmaxf`, and the down launched one neuron short). +CUDA also checks Flash Next's EXL3 host tier (`flashnext_exl3_kvtier.cu`: a layer's K/V in lent frames, a block cache +over the host store). `kvt_fetch` is compared with a host model byte for byte over every lent frame (the NaN and +infinities planted in the host rows arrive as stored, nothing else is written, entries of the list past the count are +never read) and `kvt_remap` with its model over the key lists, counts and overflow flag, which it rewrites in place +(a missing block, a flag already set and rows without a selection included). `kvt_plan` places each miss in whichever +victim an atomic reaches first, so the harness reads an image of what does not depend on that: the counters, the blocks +needed and missed, the marks and the multiset of stamps, which the host's LRU model gives exactly, then checks the +placement is one an LRU may make (victims are the least recently used slots the call does not use, equally old ones in +any order; evicted blocks are unmapped, hit slots are stamped, nothing else moves). Its calls start from a steady +cache, every block a hit, a fresh one, heavy stamp ties, a wrapped clock, no selecting rows, and a flag left set, and at +the smallest slot count the lists overflow. Gate 4 does not apply (integers only), and gate 1 plants garbage and the bit +patterns of NaN and the infinities where the kernels never read. Three controls, red at gate 3: the table read one +entry late, slot rows one block late, the window one block early. What it does not check: which miss takes which of +several equally old slots, as that is the atomics' order, so a run of many calls is not compared byte for byte (the +device test, `exl3-kvtier`, holds its misses within 5% of the model over 60 drifting calls). Metal also checks DeepSeek V4.1's CSA2 attention (`csa2_attn.metal`): decode rows (ring or linear windows) and prompt rows (16-key tiles at absolute positions), fp32 and bf16 out, against float64's softmax over the window, the selected rows and @@ -377,7 +392,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 16 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. The first, `fn_exl3_split_f16`, also runs gates 1 to 4 on the CUDA runner (`contract_flashnext_exl3.zig`): an fp16 row splits into bf16 hi = bf16(x) and lo = bf16(x - hi), every output's bits must equal float64's (NaN stays NaN; an diff --git a/docs/recipes/flash-next-cuda.md b/docs/recipes/flash-next-cuda.md index 2fe3f96570..f55a9e6916 100644 --- a/docs/recipes/flash-next-cuda.md +++ b/docs/recipes/flash-next-cuda.md @@ -136,6 +136,10 @@ starts a rank. These switches tune it: | `TF_FLASHNEXT_EXL3_PREFILL_TWO_CHUNK` | `0` | Experimental: `1` (two ranks, ring) runs a second deep chunk beside the first when memory admits it. | | `TF_FLASHNEXT_EXL3_PREFILL_SNAP_GRID` | on | Diagnostics only: `0` keeps prompt-cache states at any row instead of on the serving step's starts, to test that a resumed prompt equals a fresh one. | | `TF_FLASHNEXT_PROFILE` | `0` | Diagnostics only: `1` times the ranks' combine, all-reduce and place of a deep chunk, with a wait each. | +| `TF_FLASHNEXT_KV_PAGES` | unset | `n`: streams beside the first take their attention caches from one pool of `n`-position pages. | +| `TF_FLASHNEXT_KV_HOST` | unset | With `TF_FLASHNEXT_KV_PAGES`: MiB of pinned host RAM holding the K/V plane of every page; `0` is no host tier. | +| `TF_FLASHNEXT_KV_FRAMES` | the pages the plan gave | With the host tier: at most this many device frames of the K/V plane. | +| `TF_FLASHNEXT_KV_SLOTS` | 128 | With the host tier: frames lent to a stream, 64 four-row slots a frame for each attention layer. | | `TF_FLASHNEXT_PLE_READ` | `map` | `pread` reads the n-gram embedding rows by positional read. | | `TF_FLASHNEXT_SNAP_HOST` | `1` above 32,768 tokens of context, else `0` | Whether kept prompt-cache states live in pinned host memory. | | `TF_FLASHNEXT_CACHE_TELEMETRY` | `0` | Diagnostics only: `1` logs `EXL3-CACHE` expert-cache counters per phase, with a sync at each boundary. | diff --git a/zig/build/cuda_kernels.zig b/zig/build/cuda_kernels.zig index 62476b939c..c54dad2802 100644 --- a/zig/build/cuda_kernels.zig +++ b/zig/build/cuda_kernels.zig @@ -154,6 +154,8 @@ pub const kernels = [_]Kernel{ .{ .name = "flashnext_exl3_tp", .src = "flashnext_exl3/flashnext_exl3_tp", .flags = glue }, // ours: Flash Next EXL3 prefill experts, 64-member tiles any width .{ .name = "flashnext_exl3_prefill", .src = "flashnext_exl3/flashnext_exl3_prefill", .flags = &.{ "-O3", "--expt-relaxed-constexpr" } }, + // ours: Flash Next K/V on the host tier, a device block cache in lent frames + .{ .name = "flashnext_exl3_kvtier", .src = "flashnext_exl3/flashnext_exl3_kvtier", .flags = glue }, }; /// 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_attn.cu b/zig/kernels/cuda/flashnext_attn.cu index 3d9f7bb7f7..7e8f5f4e6a 100644 --- a/zig/kernels/cuda/flashnext_attn.cu +++ b/zig/kernels/cuda/flashnext_attn.cu @@ -50,7 +50,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, @@ -58,13 +58,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; } @@ -149,19 +149,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])); @@ -171,7 +171,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)); } @@ -582,3 +582,6 @@ extern "C" __global__ void __launch_bounds__(256) } #include "flashnext_attn_lanes.cuh" + +// 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 4f57de5bec..4baa364aff 100644 --- a/zig/src/cuda/kernels.zig +++ b/zig/src/cuda/kernels.zig @@ -22,6 +22,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 ab88811050..9a2808e0d5 100644 --- a/zig/src/families/flashnext/cuda_batch.zig +++ b/zig/src/families/flashnext/cuda_batch.zig @@ -352,12 +352,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) { @@ -383,7 +385,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 d525b7ed1f..eb451a63e5 100644 --- a/zig/src/families/flashnext/cuda_batch_state.zig +++ b/zig/src/families/flashnext/cuda_batch_state.zig @@ -514,6 +514,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 445b9e4e88..f8afcbf594 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 { @@ -67,6 +85,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, @@ -162,6 +185,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"), @@ -245,8 +273,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); } @@ -259,7 +287,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; @@ -275,7 +303,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); } @@ -291,8 +319,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); } @@ -303,8 +331,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 @@ -550,9 +578,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 d2131c7a11..9ec9d215c7 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 0aad0755db..5991008156 100644 --- a/zig/src/families/flashnext/cuda_lanes_load.zig +++ b/zig/src/families/flashnext/cuda_lanes_load.zig @@ -77,6 +77,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 9d269a63f9..400f278974 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, @@ -270,7 +271,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 f2bcb4e330..94b771e867 100644 --- a/zig/src/families/flashnext/cuda_prompt.zig +++ b/zig/src/families/flashnext/cuda_prompt.zig @@ -517,6 +517,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, @@ -527,7 +529,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..1a3a33728a --- /dev/null +++ b/zig/src/families/flashnext/exl3/kv/cuda_exl3_kvtier.zig @@ -0,0 +1,315 @@ +//! 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); + // from its first call the MTP head holds pages back until it has written them + if (ai == attn_n) s.head_on = true; + 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_arch_test.zig b/zig/tests/cuda/contract_arch_test.zig index a8b1f8ce1c..e10e989668 100644 --- a/zig/tests/cuda/contract_arch_test.zig +++ b/zig/tests/cuda/contract_arch_test.zig @@ -142,4 +142,6 @@ test { _ = @import("contract_nvfp4_waves_test.zig"); _ = @import("contract_exl3.zig"); // its coverage of the EXL3 dispatch's forms _ = @import("contract_flashnext_exl3_promptd.zig"); // its shapes and the stress weights of gate 4 + _ = @import("contract_flashnext_exl3_kvt.zig"); // the host tier's fixtures and their models + _ = @import("contract_flashnext_exl3_kvt_plan_test.zig"); } diff --git a/zig/tests/cuda/contract_flashnext_exl3_kvt.zig b/zig/tests/cuda/contract_flashnext_exl3_kvt.zig new file mode 100644 index 0000000000..389e7e9b87 --- /dev/null +++ b/zig/tests/cuda/contract_flashnext_exl3_kvt.zig @@ -0,0 +1,480 @@ +//! Flash Next EXL3's host tier: kvt_fetch and kvt_remap against host models byte for byte; kvt_plan is in its own file. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const check = @import("check.zig"); +const rig_ = @import("contract_rig.zig"); +const plan_ = @import("contract_flashnext_exl3_kvt_plan.zig"); + +const c = core.contract.spec; +const g = core.contract.gates; +const page_rows = plan_.page_rows; +const ratio: usize = plan_.ratio; +const idw: usize = plan_.idw; +// a K or V row (2 heads of 256 dims, bf16) and a page of one layer's K then V (flash.exl3.pages' layout) +pub const row_bytes = 1024; +const row16 = row_bytes / 16; +pub const kv_page = 2 * page_rows * row_bytes; +pub const value_at = page_rows * row_bytes; +const hpages = 16; // logical pages in the host store +const vpage = 12; // the table entries before the lent frames +const win0 = 40_000; +const nblocks = 200 * page_rows / ratio; +const none = std.math.maxInt(u32); +const garbage = [_]u32{ 0x80000000, 0x7fffffff, 0xffffffff, 0x7fc00000, 0x7f800000, 0xff800000, 0 }; + +fn noSums(_: c.Shape) usize { + return 0; +} + +fn sh(rows: usize, n: usize, k: usize) c.Shape { + return .{ .rows = rows, .n = n, .k = k }; +} + +const fetch_rule = "bytes moved, never summed: a word's bits (the NaN and both infinities planted in the rows fetched) arrive as stored, " ++ + "and entries of the list past the count are never read"; +const remap_rule = "integer positions only: a row without a selection, and keys past a row's count, are never read"; + +// fetches (the 128 blocks a grid strides: the first 129 gives a block two), lent frames of 64 slots, pages in the store +const fetch_served = [_]c.Shape{ sh(128, 4, hpages), sh(1, 4, hpages), sh(200, 8, hpages), sh(512, 8, hpages) }; +// rows, keys a row (a row lists up to idw - 1, the key count the engine passes) +const remap_served = [_]c.Shape{ sh(16, idw - 1, 0), sh(1, idw - 1, 0), sh(512, idw - 1, 0), sh(4, 300, 0) }; + +fn fetchSpec(control: bool) c.Spec { + return .{ + .name = if (control) "control: Flash Next EXL3 kvt_fetch reading the table one entry late" else "cuda Flash Next EXL3 host-tier fetch kvt_fetch", + .accum = .none, + .depth = noSums, + .out = .raw, + .nan = .{ .stated = fetch_rule }, + .tiles = .{ .{ .size = 128 }, .{}, .{} }, + .served = &fetch_served, + }; +} + +fn remapSpec(control: bool) c.Spec { + return .{ + .name = if (control) "control: Flash Next EXL3 kvt_remap with the slot rows one block late" else "cuda Flash Next EXL3 host-tier remap kvt_remap", + .accum = .none, + .depth = noSums, + .out = .raw, + .nan = .{ .stated = remap_rule }, + .tiles = .{ .{}, .{ .size = 256 }, .{} }, + .served = &remap_served, + }; +} + +fn less(_: void, x: u32, y: u32) bool { + return x < y; +} + +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); +} + +/// Fills `out` with distinct values below `limit`, in a random order. +fn distinct(rng: std.Random, a: std.mem.Allocator, limit: usize, out: []u32) !void { + const all = try a.alloc(u32, limit); + defer a.free(all); + for (all, 0..) |*v, i| v.* = @intCast(i); + rng.shuffle(u32, all); + @memcpy(out, all[0..out.len]); +} + +pub const Fetch = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + f: cuda.Function, + control: bool = false, + shape: c.Shape = undefined, + fill: g.Fill = .random, + frames: usize = 0, + store: []u8 = &.{}, + lids: []u32 = &.{}, + pt: []u32 = &.{}, + list: []u32 = &.{}, + ctl: [8]u32 = @splat(0), + want: []u8 = &.{}, + d_store: u64 = 0, + d_lids: u64 = 0, + d_pt: u64 = 0, + d_list: u64 = 0, + d_ctl: u64 = 0, + d_arena: u64 = 0, + msg: [160]u8 = undefined, + + pub fn deinit(self: *Fetch) void { + self.free(); + self.rig.deinit(); + } + + fn free(self: *Fetch) void { + const a = self.a; + a.free(self.store); + a.free(self.want); + for ([_]*[]u32{ &self.lids, &self.pt, &self.list }) |p| { + a.free(p.*); + p.* = &.{}; + } + self.store = &.{}; + self.want = &.{}; + } + + const vt = g.hooks(Fetch); + + pub fn kernel(self: *Fetch) g.Kernel { + return .{ .spec = fetchSpec(self.control), .ctx = self, .vt = &vt, .control = if (self.control) .shapes else null }; + } + + /// Random store words; `n` distinct blocks into `n` distinct slots of the lent frames, shuffled frames behind them. + pub fn prepare(self: *Fetch, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (!try self.build(shape, fill, seed)) return false; + self.rig.reset(); + self.d_store = try self.rig.input(self.store); + self.d_lids = try self.rig.input(std.mem.sliceAsBytes(self.lids)); + self.d_pt = try self.rig.input(std.mem.sliceAsBytes(self.pt)); + self.d_list = try self.rig.input(std.mem.sliceAsBytes(self.list)); + self.d_ctl = try self.rig.input(std.mem.sliceAsBytes(&self.ctl)); + self.d_arena = try self.rig.output(self.want.len); + return true; + } + + fn build(self: *Fetch, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + const n = shape.rows; + self.frames = shape.n; + if (n == 0 or self.frames == 0 or n > self.frames * 64 or shape.k != hpages or n > hpages * 64) return false; + self.free(); + self.shape = shape; + self.fill = fill; + const a = self.a; + var prng = std.Random.DefaultPrng.init(seed ^ 0x6b767466); + const rng = prng.random(); + self.store = try a.alloc(u8, hpages * kv_page); + rng.bytes(self.store); + self.lids = try a.alloc(u32, hpages); + try distinct(rng, a, hpages, self.lids); + self.pt = try a.alloc(u32, vpage + self.frames + 1); + @memset(self.pt[0..vpage], 0); + try distinct(rng, a, 2 * self.frames, self.pt[vpage..]); + const slots = self.frames * 64; + self.list = try a.alloc(u32, 2 * (slots + 8)); + for (self.list, 0..) |*v, i| v.* = garbage[i % garbage.len]; + const blocks = try a.alloc(u32, n); + defer a.free(blocks); + const at = try a.alloc(u32, n); + defer a.free(at); + try distinct(rng, a, hpages * 64, blocks); + try distinct(rng, a, slots, at); + for (blocks, at, 0..) |b, s, i| { + self.list[2 * i] = b; + self.list[2 * i + 1] = s; + } + self.ctl = .{ 0x7fc00000, @intCast(n), 0xffffffff, 5, 6, 7, 8, 9 }; + self.want = try a.alloc(u8, 2 * self.frames * kv_page); + @memset(self.want, 0xff); + const per = ratio * row_bytes; + const planes = [_]usize{ 0, value_at }; + for (blocks, at) |b, s| { + const page = b * ratio / page_rows; + const within = b * ratio % page_rows; + const src = @as(usize, self.lids[page]) * kv_page + within * row_bytes; + const dst = @as(usize, self.pt[vpage + s / 64]) * kv_page + (s % 64) * per; + for (planes) |p| { + if (fill == .nonfinite) for ([_]u32{ 0x7fc00000, 0x7f800000, 0xff800000 }, 0..) |w, i| + @memcpy(self.store[p + src + 4 * i ..][0..4], std.mem.asBytes(&w)); + @memcpy(self.want[p + dst ..][0..per], self.store[p + src ..][0..per]); + } + } + return true; + } + + pub fn run(self: *Fetch) !void { + var f: cuda.Args = .{}; + const k: u64 = 0; + const v: u64 = value_at; + for ([_]u64{ self.d_store + k, self.d_store + v, self.d_arena + k, self.d_arena + v, self.d_ctl, self.d_list, self.d_lids, self.d_pt }) |p| f.add(p); + const table_at: usize = if (self.control) vpage + 1 else vpage; + f.add(@as(c_int, @intCast(table_at))); + f.add(@as(i64, @intCast(kv_page / 16))); + f.add(@as(c_int, @intCast(row16))); + f.add(@as(c_int, @intCast(ratio))); + try cuda.launch.launch(self.f, .{ .grid = .{ .x = 128 }, .block = .{ .x = 256 } }, self.rig.s, &f); + try self.rig.s.synchronize(); + } + + pub fn count(self: *Fetch) usize { + return self.want.len; + } + + pub fn reference(self: *Fetch, want: []f64, mag: []f64) !void { + for (want, mag, self.want) |*w, *m, b| { + w.* = @floatFromInt(b); + m.* = w.*; + } + } + + /// Every byte of the lent frames: the rows fetched where the table puts them, the rest as poisoned. + pub fn check(self: *Fetch, _: g.Gate) !?[]const u8 { + const got = try self.rig.outputBytes(); + if (std.mem.indexOfDiff(u8, got, self.want)) |at| return try std.fmt.bufPrint(&self.msg, "byte {d} of {d} of the lent frames differs from the host store's row", .{ at, got.len }); + if (self.fill != .nonfinite) return null; + for (std.mem.bytesAsSlice(u32, got)) |w| if (w == 0x7f800000) return null; + return "no fetched word is an infinity: the planted values prove nothing"; + } +}; + +/// The slot rows start after the stream's pages. +const vrow: u32 = 200 * page_rows; + +pub const Remap = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + f: cuda.Function, + control: bool = false, + shape: c.Shape = undefined, + ids: []u32 = &.{}, + nkr: []u32 = &.{}, + spr: []u32 = &.{}, + b2s: []u32 = &.{}, + ctl: [8]u32 = @splat(0), + want: []u8 = &.{}, + d_ids: u64 = 0, + d_nk: u64 = 0, + d_spr: u64 = 0, + d_b2s: u64 = 0, + d_ctl: u64 = 0, + msg: [160]u8 = undefined, + + pub fn deinit(self: *Remap) void { + self.free(); + self.rig.deinit(); + } + + fn free(self: *Remap) void { + const a = self.a; + for ([_]*[]u32{ &self.ids, &self.nkr, &self.spr, &self.b2s }) |p| { + a.free(p.*); + p.* = &.{}; + } + a.free(self.want); + self.want = &.{}; + } + + const vt = g.hooks(Remap); + + pub fn kernel(self: *Remap) g.Kernel { + return .{ .spec = remapSpec(self.control), .ctx = self, .vt = &vt, .control = if (self.control) .shapes else null }; + } + + /// Rows of 512 sorted blocks' keys; where a block is missing every selecting row has one, so the rows agree on it. + pub fn prepare(self: *Remap, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (!try self.build(shape, fill, seed)) return false; + self.rig.reset(); + self.d_spr = try self.rig.input(std.mem.sliceAsBytes(self.spr)); + self.d_b2s = try self.rig.input(std.mem.sliceAsBytes(self.b2s)); + self.d_ids = try self.rig.output(self.ids.len * 4); + self.d_nk = try self.rig.output(shape.rows * 4); + self.d_ctl = try self.rig.output(32); + return true; + } + + fn build(self: *Remap, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + const rows = shape.rows; + const nk = shape.n; + if (rows == 0 or rows > 512 or nk == 0 or nk >= idw) return false; + self.free(); + self.shape = shape; + const a = self.a; + var prng = std.Random.DefaultPrng.init(seed ^ 0x6b767472); + const rng = prng.random(); + const before = win0 / ratio; + // 0: rows select, every block in a slot; 1: the flag set; 2: each selecting row misses one; 3: one row does + const variant = seed % 4; + self.b2s = try a.alloc(u32, nblocks); + for (self.b2s) |*s| s.* = rng.uintLessThan(u32, 8192); + self.ids = try a.alloc(u32, rows * idw); + self.nkr = try a.alloc(u32, rows); + self.spr = try a.alloc(u32, rows); + const only = rng.uintLessThan(usize, rows); + for (0..rows) |r| { + self.spr[r] = switch (variant) { + 0 => rng.uintLessThan(u32, 2), + 3 => @intFromBool(r == only), + else => 1, + }; + if (variant == 0 and r == 0) self.spr[r] = 1; + self.nkr[r] = if (self.spr[r] != 0) @intCast(nk) else 1 + rng.uintLessThan(u32, idw - 1); + const out = self.ids[r * idw ..][0..idw]; + const blocks = out[0 .. idw - 4]; + try distinct(rng, a, before, blocks[0 .. blocks.len / ratio]); + std.mem.sort(u32, blocks[0 .. blocks.len / ratio], {}, less); + var j = blocks.len / ratio; + while (j > 0) { + j -= 1; + const b = blocks[j]; + for (0..ratio) |t| blocks[j * ratio + t] = b * @as(u32, @intCast(ratio)) + @as(u32, @intCast(t)); + } + for (0..4) |t| out[idw - 4 + t] = win0 + 100 + @as(u32, @intCast(t)); + if (self.spr[r] != 0 and (variant == 2 or variant == 3)) self.b2s[out[rng.uintLessThan(usize, @min(nk, idw - 4))] / ratio] = none; + } + self.ctl = .{ 7, 8, @intFromBool(variant == 1), 9, 10, 11, 12, 13 }; + if (fill == .nonfinite) for (0..rows) |r| { + const out = self.ids[r * idw ..][0..idw]; + for (out, 0..) |*e, i| if (self.spr[r] == 0 or i >= self.nkr[r]) { + e.* = garbage[(r + i) % garbage.len]; + }; + }; + try self.model(); + return true; + } + + /// The remapped lists, counts and flag, as the three outputs lie side by side. + fn model(self: *Remap) !void { + const a = self.a; + const rows = self.shape.rows; + var ctl2 = self.ctl[2]; + const ids = try a.dupe(u32, self.ids); + defer a.free(ids); + const nkr = try a.dupe(u32, self.nkr); + defer a.free(nkr); + for (0..rows) |r| { + if (self.spr[r] == 0) continue; + const row = ids[r * idw ..][0..idw]; + var missing = ctl2 != 0; + for (row[0..nkr[r]]) |p| missing = missing or (p < win0 and self.b2s[p / ratio] == none); + if (missing) { + ctl2 = 1; + nkr[r] = 0; + continue; + } + for (row[0..nkr[r]]) |*p| if (p.* < win0) { + p.* = vrow + self.b2s[p.* / ratio] * @as(u32, @intCast(ratio)) + p.* % @as(u32, @intCast(ratio)); + }; + } + var ctl = self.ctl; + ctl[2] = ctl2; + a.free(self.want); + self.want = try a.alloc(u8, ids.len * 4 + nkr.len * 4 + 32); + @memcpy(self.want[0 .. ids.len * 4], std.mem.sliceAsBytes(ids)); + @memcpy(self.want[ids.len * 4 ..][0 .. nkr.len * 4], std.mem.sliceAsBytes(nkr)); + @memcpy(self.want[ids.len * 4 + nkr.len * 4 ..], std.mem.sliceAsBytes(&ctl)); + } + + pub fn run(self: *Remap) !void { + try put(&self.rig, self.d_ids, std.mem.sliceAsBytes(self.ids)); + try put(&self.rig, self.d_nk, std.mem.sliceAsBytes(self.nkr)); + try put(&self.rig, self.d_ctl, std.mem.sliceAsBytes(&self.ctl)); + var r: cuda.Args = .{}; + for ([_]u64{ self.d_ids, self.d_nk, self.d_spr }) |v| r.add(v); + for ([_]usize{ idw, win0, if (self.control) vrow + ratio else vrow }) |v| r.add(@as(c_int, @intCast(v))); + r.add(self.d_b2s); + r.add(@as(c_int, @intCast(ratio))); + r.add(self.d_ctl); + try cuda.launch.launch(self.f, .{ .grid = .{ .x = @intCast(self.shape.rows) }, .block = .{ .x = 256 } }, self.rig.s, &r); + try self.rig.s.synchronize(); + } + + pub fn count(self: *Remap) usize { + return self.want.len; + } + + pub fn reference(self: *Remap, want: []f64, mag: []f64) !void { + for (want, mag, self.want) |*w, *m, b| { + w.* = @floatFromInt(b); + m.* = w.*; + } + } + + /// The lists, counts and flag the model gives, byte for byte. + pub fn check(self: *Remap, _: g.Gate) !?[]const u8 { + const got = try self.rig.outputBytes(); + if (std.mem.indexOfDiff(u8, got, self.want)) |at| { + const ids = self.ids.len * 4; + const part = if (at < ids) "a key list" else if (at < ids + self.nkr.len * 4) "a key count" else "the overflow flag"; + return try std.fmt.bufPrint(&self.msg, "byte {d} of {d} ({s}) differs from the model's", .{ at, got.len, part }); + } + return null; + } +}; + +/// Fetch and remap with their controls, and the plan's, over one loaded image. +pub const Group = struct { + a: std.mem.Allocator, + mod: cuda.Module, + plan: plan_.Set, + fetch: [2]Fetch, + remap: [2]Remap, + + pub fn create(a: std.mem.Allocator, gpu: check.Gpu, s: cuda.Stream, list: *std.ArrayList(g.Kernel)) !*Group { + const self = try a.create(Group); + errdefer a.destroy(self); + self.a = a; + self.mod = try cuda.Module.load(gpu.d, cuda.kernels.flashnext_exl3_kvtier); + errdefer self.mod.unload(); + self.plan = try plan_.Set.init(gpu, s); + errdefer self.plan.deinit(); + const ff = try self.mod.function("kvt_fetch"); + const fr = try self.mod.function("kvt_remap"); + for (&self.fetch, [2]bool{ false, true }) |*x, control| x.* = .{ .rig = .{ .gpu = gpu, .s = s }, .a = a, .f = ff, .control = control }; + errdefer for (&self.fetch) |*x| x.deinit(); + for (&self.remap, [2]bool{ false, true }) |*x, control| x.* = .{ .rig = .{ .gpu = gpu, .s = s }, .a = a, .f = fr, .control = control }; + errdefer for (&self.remap) |*x| x.deinit(); + try self.plan.add(a, list); + for (&self.fetch) |*x| try list.append(a, x.kernel()); + for (&self.remap) |*x| try list.append(a, x.kernel()); + return self; + } + + pub fn destroy(self: *Group) void { + for (&self.remap) |*x| x.deinit(); + for (&self.fetch) |*x| x.deinit(); + self.plan.deinit(); + self.mod.unload(); + self.a.destroy(self); + } +}; + +/// tf-cuda-test exl3-kvt-contract: these kernels through the gates and their controls alone. +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, 0x6b7674, &t); + std.debug.print("SUMMARY exl3 kvt {d} kernels: {d} pass, {d} fail; {d} controls: {d} red as planted, {d} missed\n", .{ t.kernels, t.passed, t.failed, t.controls, t.caught, t.missed }); + try check.expect(t.problems() == 0, "exl3-kvt-contract: {d} gate failures, {d} controls missed", .{ t.failed, t.missed }); +} + +test "kvt fetch and remap contract: every served call moves rows, and plants reach the model" { + const a = std.testing.allocator; + for (fetch_served) |shape| for ([_]g.Fill{ .random, .nonfinite }) |fill| { + var cs: Fetch = .{ .a = a, .rig = undefined, .f = undefined }; + defer cs.free(); + try std.testing.expect(try cs.build(shape, fill, 7)); + var written: usize = 0; + for (cs.want) |b| written += @intFromBool(b != 0xff); + try std.testing.expect(written >= shape.rows * 2 * ratio * row_bytes / 2); // random words are mostly not 0xff + var inf = false; + for (std.mem.bytesAsSlice(u32, cs.want)) |w| inf = inf or w == 0x7f800000; + try std.testing.expectEqual(fill == .nonfinite, inf); + }; + var flagged = false; + var zeroed = false; + var moved = false; + for (remap_served) |shape| for (0..8) |seed| { + var cs: Remap = .{ .a = a, .rig = undefined, .f = undefined }; + defer cs.free(); + try std.testing.expect(try cs.build(shape, .nonfinite, seed)); + const ids = cs.ids.len * 4; + const out = std.mem.bytesAsSlice(u32, cs.want[0..ids]); + const nk = std.mem.bytesAsSlice(u32, cs.want[ids..][0 .. cs.nkr.len * 4]); + flagged = flagged or std.mem.readInt(u32, cs.want[ids + cs.nkr.len * 4 + 8 ..][0..4], .little) != 0; + for (nk, cs.nkr, cs.spr) |got, was, sparse| zeroed = zeroed or (sparse != 0 and got == 0 and was != 0); + for (out, cs.ids) |got, was| moved = moved or (got != was and got >= vrow); + }; + try std.testing.expect(flagged and zeroed and moved); +} diff --git a/zig/tests/cuda/contract_flashnext_exl3_kvt_plan.zig b/zig/tests/cuda/contract_flashnext_exl3_kvt_plan.zig new file mode 100644 index 0000000000..137d682e7a --- /dev/null +++ b/zig/tests/cuda/contract_flashnext_exl3_kvt_plan.zig @@ -0,0 +1,553 @@ +//! Flash Next EXL3's host-tier slot plan (kvt_plan) under the contract: its LRU choice against a host model, per call. +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; + +// the tier's geometry (the suite checks it against the engine's constants at compile time) +pub const ratio = 4; +pub const top = 512; +pub const idw = top * ratio + ratio; +pub const page_rows = 256; +const npages = 200; +const nblocks = npages * page_rows / ratio; +const win0 = 40_000; // the window starts here: blocks from win0 / ratio are never cached +pub const before = win0 / ratio; +pub const none = std.math.maxInt(u32); +const rule = "integers only: a row without a selection and the three keys after a block's first are never read, " ++ + "so extreme values and the bit patterns of NaN and the infinities planted there change nothing"; + +fn noSums(_: c.Shape) usize { + return 0; +} + +fn sh(rows: usize, slots: usize) c.Shape { + return .{ .rows = rows, .n = slots, .k = top }; +} + +// rows a call, slots a layer (128, 16 and 9 lent frames of 64): a call plans `slots / top` sparse rows at most +pub const served = [_]c.Shape{ sh(16, 8192), sh(1, 8192), sh(64, 8192), sh(2, 1024), sh(1, 1024), sh(8, 1024), sh(1, 576), sh(3, 576) }; + +/// The slot plan: counts, epoch, who is needed, who misses, the stamps' multiset; raw: an image the harness compares. +pub fn spec(control: bool) c.Spec { + return .{ + .name = if (control) "control: Flash Next EXL3 kvt_plan with the window one block early" else "cuda Flash Next EXL3 host-tier slot plan kvt_plan", + .accum = .none, + .depth = noSums, + .out = .raw, + .nan = .{ .stated = rule }, + .served = &served, + }; +} + +/// Where a call starts: a steady cache, all hits, a fresh one, heavy ties, a wrapped clock, no rows, a set flag. +pub const Variant = enum { steady, hits, fresh, ties, wrap, idle, sticky }; + +const garbage = [_]u32{ 0x80000000, 0x7fffffff, 0xffffffff, 0x7fc00000, 0x7f800000, 0xff800000, 0 }; + +/// What one call leaves, as far as the order the atomics take does not matter, and the victims' old stamps. +pub const Model = struct { + ctl: [8]u32 = @splat(0), + need: []u32 = &.{}, + miss: []u32 = &.{}, + stamp: []u32 = &.{}, + mark: []u32 = &.{}, + hit: []bool = &.{}, // by slot: a slot the call uses + victims: []u32 = &.{}, // the old stamps of the slots the misses take, ascending + assigned: usize = 0, + epoch: u32 = 0, +}; + +fn less(_: void, x: u32, y: u32) bool { + return x < y; +} + +pub const Case = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + control: bool = false, + shape: c.Shape = undefined, + ids: []u32 = &.{}, + spr: []u32 = &.{}, + b2s: []u32 = &.{}, + s2b: []u32 = &.{}, + stamp: []u32 = &.{}, + mark: []u32 = &.{}, + ctl: [8]u32 = @splat(0), + model: Model = .{}, + f: cuda.Function, + d_ids: u64 = 0, + d_nk: u64 = 0, + d_spr: u64 = 0, + d_b2s: u64 = 0, + d_s2b: u64 = 0, + d_stamp: u64 = 0, + d_mark: u64 = 0, + d_ctl: u64 = 0, + d_need: u64 = 0, + d_miss: u64 = 0, + d_fetch: u64 = 0, + cap: usize = 0, + msg: [200]u8 = undefined, + image_buf: std.ArrayList(u8) = .empty, + + pub fn deinit(self: *Case) void { + self.free(); + self.image_buf.deinit(self.a); + self.rig.deinit(); + } + + pub fn free(self: *Case) void { + const a = self.a; + for ([_]*[]u32{ &self.ids, &self.spr, &self.b2s, &self.s2b, &self.stamp, &self.mark, &self.model.need, &self.model.miss, &self.model.stamp, &self.model.mark, &self.model.victims }) |p| { + a.free(p.*); + p.* = &.{}; + } + a.free(self.model.hit); + self.model.hit = &.{}; + } + + /// The harness reads the image, not the raw lists: a miss lands in whichever victim an atomic reaches first. + const vt = blk: { + var v = g.hooks(Case); + v.output = image; + break :blk v; + }; + + pub fn kernel(self: *Case) g.Kernel { + return .{ .spec = spec(self.control), .ctx = self, .vt = &vt, .control = if (self.control) .shapes else null }; + } + + /// A call's rows over a state from `seed`; non-finite: garbage where the kernel never reads. + pub fn prepare(self: *Case, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (!try self.build(shape, fill, seed)) return false; + self.rig.reset(); + try self.stageDevice(shape.n); + return true; + } + + /// The call's host side: the state, the rows and the model of what the call leaves. + pub fn build(self: *Case, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + const slots = shape.n; + if (shape.rows == 0 or shape.rows > 512 or shape.k != top or slots % 64 != 0 or slots < top or slots > before) return false; + self.free(); + self.shape = shape; + const a = self.a; + var prng = std.Random.DefaultPrng.init(seed ^ 0x6b767470); + const rng = prng.random(); + const variant: Variant = @fromBackingInt(@intCast(seed % std.enums.values(Variant).len)); + const rows = shape.rows; + const per_plan = @max(1, slots / top); + // overflow: one sparse row more than the slots hold, where the lists still fit (capacity 2,048 or the slots) + const over = (per_plan + 1) * top <= @max(slots, 2048) and variant == .steady and slots < 8192; + const sparse = if (variant == .idle) 0 else @min(rows, if (over) per_plan + 1 else per_plan); + self.b2s = try a.alloc(u32, nblocks); + self.s2b = try a.alloc(u32, slots); + self.stamp = try a.alloc(u32, slots); + self.mark = try a.alloc(u32, nblocks); + @memset(self.b2s, none); + @memset(self.s2b, none); + @memset(self.stamp, 0); + const clock: u32 = if (variant == .wrap) none else 1000; + @memset(self.mark, 0); + for (self.mark) |*m| if (rng.uintLessThan(u32, 8) == 0) { + m.* = 1 + rng.uintLessThan(u32, 1000); + }; + // the blocks before the window in a random order: the first `have` of them sit in slots + const perm = try a.alloc(u32, before); + defer a.free(perm); + for (perm, 0..) |*p, i| p.* = @intCast(i); + rng.shuffle(u32, perm); + const have: usize = switch (variant) { + .fresh => 0, + .hits, .ties => slots, + else => slots * 3 / 4, + }; + const at = try a.alloc(u32, slots); + defer a.free(at); + for (at, 0..) |*p, i| p.* = @intCast(i); + rng.shuffle(u32, at); + for (perm[0..have], at[0..have]) |b, s| { + self.b2s[b] = s; + self.s2b[s] = b; + self.stamp[s] = if (variant == .ties) 1 + rng.uintLessThan(u32, 4) else 1 + rng.uintLessThan(u32, 1000); + } + self.ctl = @splat(0); + self.ctl[0] = clock; + self.ctl[2] = @intFromBool(variant == .sticky); + for (3..6) |i| self.ctl[i] = rng.uintLessThan(u32, 100); // counters add to what they held + self.ids = try a.alloc(u32, rows * idw); + self.spr = try a.alloc(u32, rows); + const pre = top - 6; // blocks before the window a row lists, then six inside it (never cached) + const row = try a.alloc(u32, top); + defer a.free(row); + for (0..rows) |r| { + self.spr[r] = @intFromBool(r < sparse); + const out = self.ids[r * idw ..][0..idw]; + const want_hits: usize = switch (variant) { + .hits => pre, + .fresh => 0, + else => pre * 6 / 10, + }; + const hits = @min(want_hits, have); + try sample(rng, a, perm[0..have], hits, row[0..hits]); + try sample(rng, a, perm[have..], pre - hits, row[hits..pre]); + // the last block before the window is in every row, so a window that starts early misses it + if (std.mem.indexOfScalar(u32, row[0..pre], before - 1) == null) row[0] = before - 1; + for (row[pre..], 0..) |*e, i| e.* = @intCast(before + 1 + i); + std.mem.sort(u32, row, {}, less); + for (row, 0..) |b, j| for (0..ratio) |t| { + out[j * ratio + t] = b * ratio + @as(u32, @intCast(t)); + }; + for (0..ratio) |t| out[top * ratio + t] = win0 + 100 + @as(u32, @intCast(t)); + } + if (fill == .nonfinite) for (0..rows) |r| { + const out = self.ids[r * idw ..][0..idw]; + for (out, 0..) |*e, i| if (self.spr[r] == 0 or (i < top * ratio and i % ratio != 0)) { + e.* = garbage[(r + i) % garbage.len]; + }; + }; + try self.plan(slots, clock); + return true; + } + + fn stageDevice(self: *Case, slots: usize) !void { + const rows = self.shape.rows; + self.cap = @max(slots, 2048); + self.d_ids = try self.rig.input(std.mem.sliceAsBytes(self.ids)); + self.d_spr = try self.rig.input(std.mem.sliceAsBytes(self.spr)); + const nk = try self.a.alloc(u32, rows); + defer self.a.free(nk); + @memset(nk, idw - 1); + self.d_nk = try self.rig.input(std.mem.sliceAsBytes(nk)); + self.d_b2s = try self.rig.output(nblocks * 4); + self.d_mark = try self.rig.output(nblocks * 4); + self.d_ctl = try self.rig.output(8 * 4); + self.d_s2b = try self.rig.output(slots * 4); + self.d_stamp = try self.rig.output(slots * 4); + self.d_need = try self.rig.output(self.cap * 4); + self.d_miss = try self.rig.output(self.cap * 4); + self.d_fetch = try self.rig.output(self.cap * 8); + } + + fn put(self: *Case, ptr: u64, bytes: []const u8) !void { + const b: cuda.DeviceBuffer = .{ .d = self.rig.gpu.d, .ptr = ptr, .len = bytes.len }; + try b.upload(0, bytes); + } + + pub fn run(self: *Case) !void { + try self.put(self.d_b2s, std.mem.sliceAsBytes(self.b2s)); + try self.put(self.d_mark, std.mem.sliceAsBytes(self.mark)); + try self.put(self.d_ctl, std.mem.sliceAsBytes(&self.ctl)); + try self.put(self.d_s2b, std.mem.sliceAsBytes(self.s2b)); + try self.put(self.d_stamp, std.mem.sliceAsBytes(self.stamp)); + for ([_]struct { u64, usize }{ .{ self.d_need, self.cap * 4 }, .{ self.d_miss, self.cap * 4 }, .{ self.d_fetch, self.cap * 8 } }) |l| { + const b: cuda.DeviceBuffer = .{ .d = self.rig.gpu.d, .ptr = l[0], .len = l[1] }; + try b.fill8(0xff, null); + } + var a: cuda.Args = .{}; + for ([_]u64{ self.d_ids, self.d_nk, self.d_spr }) |v| a.add(v); + const origin: usize = if (self.control) win0 - ratio else win0; // the fault: a window one block early + for ([_]usize{ self.shape.rows, idw, top, ratio, origin }) |v| a.add(@as(c_int, @intCast(v))); + a.add(self.d_ctl); + a.add(self.d_b2s); + a.add(self.d_b2s); + a.add(@as(c_int, 0)); + for ([_]u64{ self.d_s2b, self.d_stamp, self.d_mark }) |v| a.add(v); + a.add(@as(c_int, @intCast(self.shape.n))); + for ([_]u64{ self.d_ctl, self.d_need, self.d_miss, self.d_fetch }) |v| a.add(v); + try cuda.launch.launch(self.f, .{ .grid = .{ .x = 1 }, .block = .{ .x = 1024 } }, self.rig.s, &a); + try self.rig.s.synchronize(); + } + + pub fn count(self: *Case) usize { + return self.imageWords() * 4; + } + + /// The host model of one call on the prepared state (kvt_plan's body, one thread, any order of its atomics). + fn plan(self: *Case, slots: usize, clock: u32) !void { + const a = self.a; + var m: Model = .{ .ctl = self.ctl }; + m.stamp = try a.dupe(u32, self.stamp); + m.mark = try a.dupe(u32, self.mark); + m.hit = try a.alloc(bool, slots); + @memset(m.hit, false); + self.model = m; + const mo = &self.model; + if (clock == none) { + mo.ctl[1] = 0; + mo.ctl[2] = 1; + return; + } + mo.epoch = clock + 1; + mo.ctl[0] = mo.epoch; + mo.ctl[6] = mo.epoch; + mo.ctl[1] = 0; + if (self.ctl[2] != 0) return; + var any = false; + for (self.spr) |s| any = any or s != 0; + if (!any) return; + var need: std.ArrayList(u32) = .empty; + defer need.deinit(a); + for (0..self.shape.rows) |r| { + if (self.spr[r] == 0) continue; + for (0..top) |j| { + const b = self.ids[r * idw + j * ratio] / ratio; + if (b * ratio >= win0) continue; + if (mo.mark[b] != mo.epoch) { + mo.mark[b] = mo.epoch; + try need.append(a, b); + } + } + } + var miss: std.ArrayList(u32) = .empty; + defer miss.deinit(a); + for (need.items) |b| { + const s = self.b2s[b]; + if (s != none and s < slots) { + mo.stamp[s] = mo.epoch; + mo.hit[s] = true; + } else try miss.append(a, b); + } + var avail: usize = 0; + for (mo.stamp) |t| avail += @intFromBool(t != mo.epoch); + mo.ctl[3] += if (miss.items.len > avail) 0 else @as(u32, @intCast(miss.items.len)); + mo.ctl[4] += @intCast(need.items.len); + mo.ctl[5] += 1; + mo.ctl[1] = if (miss.items.len > avail) 0 else @intCast(miss.items.len); + if (miss.items.len > avail) mo.ctl[2] = 1; + mo.need = try a.dupe(u32, need.items); + mo.miss = try a.dupe(u32, miss.items); + std.mem.sort(u32, mo.need, {}, less); + std.mem.sort(u32, mo.miss, {}, less); + if (miss.items.len > avail or miss.items.len == 0) return; + // the n smallest stamps among the slots the call does not use: ties are equally old, which is the multiset + var olds: std.ArrayList(u32) = .empty; + defer olds.deinit(a); + for (mo.stamp) |t| if (t != mo.epoch) try olds.append(a, t); + std.mem.sort(u32, olds.items, {}, less); + mo.assigned = miss.items.len; + mo.victims = try a.dupe(u32, olds.items[0..mo.assigned]); + // the multiset after: each victim's stamp becomes the epoch + for (mo.victims) |t| mo.stamp[std.mem.indexOfScalar(u32, mo.stamp, t).?] = mo.epoch; + } + + fn imageWords(self: *const Case) usize { + const mo = &self.model; + return 8 + 1 + mo.need.len + 1 + mo.miss.len + mo.stamp.len + mo.mark.len + 1 + mo.assigned; + } + + fn word(buf: *std.ArrayList(u8), a: std.mem.Allocator, v: u32) !void { + try buf.appendSlice(a, std.mem.asBytes(&std.mem.nativeToLittle(u32, v))); + } + + /// The call's result by words: ctl, the needed blocks, the missing ones, the stamps and marks, the fetched blocks. + fn words(self: *Case, buf: *std.ArrayList(u8), ctl: [8]u32, need: []const u32, miss: []const u32, stamp: []const u32, mark: []const u32, fetched: []const u32) !void { + const a = self.a; + buf.clearRetainingCapacity(); + for (ctl) |v| try word(buf, a, v); + try word(buf, a, @intCast(need.len)); + for (need) |v| try word(buf, a, v); + try word(buf, a, @intCast(miss.len)); + for (miss) |v| try word(buf, a, v); + for (stamp) |v| try word(buf, a, v); + for (mark) |v| try word(buf, a, v); + try word(buf, a, @intCast(fetched.len)); + for (fetched) |v| try word(buf, a, v); + } + + /// The model's image: its stamps' multiset sorted, the lists as the call's sets. + fn modelImage(self: *Case, buf: *std.ArrayList(u8)) !void { + const mo = &self.model; + const sorted = try self.a.dupe(u32, mo.stamp); + defer self.a.free(sorted); + std.mem.sort(u32, sorted, {}, less); + try self.words(buf, mo.ctl, mo.need, mo.miss, sorted, mo.mark, mo.miss[0..mo.assigned]); + } + + /// The model's image, as float64 bytes for the harness's reference. + pub fn reference(self: *Case, want: []f64, mag: []f64) !void { + var buf: std.ArrayList(u8) = .empty; + defer buf.deinit(self.a); + try self.modelImage(&buf); + for (want, mag, buf.items) |*w, *m, b| { + w.* = @floatFromInt(b); + m.* = w.*; + } + } + + const Read = struct { + ctl: [8]u32 = undefined, + b2s: []u32, + s2b: []u32, + stamp: []u32, + mark: []u32, + need: []u32, + miss: []u32, + fetch: []u32, + a: std.mem.Allocator, + + fn deinit(r: Read) void { + for ([_][]u32{ r.b2s, r.s2b, r.stamp, r.mark, r.need, r.miss, r.fetch }) |s| r.a.free(s); + } + }; + + fn pull(self: *Case, ptr: u64, n: usize) ![]u32 { + const out = try self.a.alloc(u32, n); + errdefer self.a.free(out); + const b: cuda.DeviceBuffer = .{ .d = self.rig.gpu.d, .ptr = ptr, .len = n * 4 }; + try self.rig.s.synchronize(); + try b.download(0, std.mem.sliceAsBytes(out)); + return out; + } + + fn read(self: *Case) !Read { + const a = self.a; + var r: Read = .{ .b2s = &.{}, .s2b = &.{}, .stamp = &.{}, .mark = &.{}, .need = &.{}, .miss = &.{}, .fetch = &.{}, .a = a }; + errdefer r.deinit(); + const ctl = try self.pull(self.d_ctl, 8); + defer a.free(ctl); + @memcpy(&r.ctl, ctl); + r.b2s = try self.pull(self.d_b2s, nblocks); + r.s2b = try self.pull(self.d_s2b, self.shape.n); + r.stamp = try self.pull(self.d_stamp, self.shape.n); + r.mark = try self.pull(self.d_mark, nblocks); + r.need = try self.pull(self.d_need, self.cap); + r.miss = try self.pull(self.d_miss, self.cap); + r.fetch = try self.pull(self.d_fetch, 2 * self.cap); + return r; + } + + /// The device's image: the lists as long as the model says, each sorted (their order is the atomics'). + fn deviceImage(self: *Case, r: Read, buf: *std.ArrayList(u8)) !void { + const a = self.a; + const mo = &self.model; + const stamp = try a.dupe(u32, r.stamp); + defer a.free(stamp); + std.mem.sort(u32, stamp, {}, less); + const need = try a.dupe(u32, r.need[0..mo.need.len]); + defer a.free(need); + std.mem.sort(u32, need, {}, less); + const miss = try a.dupe(u32, r.miss[0..mo.miss.len]); + defer a.free(miss); + std.mem.sort(u32, miss, {}, less); + const fetched = try a.alloc(u32, mo.assigned); + defer a.free(fetched); + for (fetched, 0..) |*f, i| f.* = r.fetch[2 * i]; + std.mem.sort(u32, fetched, {}, less); + try self.words(buf, r.ctl, need, miss, stamp, r.mark, fetched); + } + + fn image(ctx: *anyopaque) anyerror![]const u8 { + const self: *Case = @ptrCast(@alignCast(ctx)); + const r = try self.read(); + defer r.deinit(); + try self.deviceImage(r, &self.image_buf); + return self.image_buf.items; + } + + /// Every gate: the image equals the model's, then the placement is one a correct LRU may make. + pub fn check(self: *Case, _: g.Gate) !?[]const u8 { + const a = self.a; + const mo = &self.model; + const r = try self.read(); + defer r.deinit(); + var got: std.ArrayList(u8) = .empty; + defer got.deinit(a); + try self.deviceImage(r, &got); + var want: std.ArrayList(u8) = .empty; + defer want.deinit(a); + try self.modelImage(&want); + if (std.mem.indexOfDiff(u8, got.items, want.items)) |at| return try std.fmt.bufPrint(&self.msg, "the call's counters, lists or stamps differ from the model's at byte {d} (word {d} of {d})", .{ at, at / 4, want.items.len / 4 }); + // entries past the lists stay as poisoned + for (r.need[mo.need.len..]) |v| if (v != none) return "kvt_plan wrote past its need list"; + for (r.miss[mo.miss.len..]) |v| if (v != none) return "kvt_plan wrote past its miss list"; + for (r.fetch[2 * mo.assigned ..]) |v| if (v != none) return "kvt_plan wrote past its fetch list"; + return self.placement(r); + } + + /// The slots hold what the plan fetched, the evicted blocks are unmapped, and nothing else moved. + fn placement(self: *Case, r: Read) !?[]const u8 { + const a = self.a; + const mo = &self.model; + const slots = self.shape.n; + const taken = try a.alloc(bool, slots); + defer a.free(taken); + @memset(taken, false); + const gone = try a.alloc(bool, nblocks); + defer a.free(gone); + @memset(gone, false); + const new = try a.alloc(bool, nblocks); + defer a.free(new); + @memset(new, false); + const olds = try a.alloc(u32, mo.assigned); + defer a.free(olds); + for (0..mo.assigned) |i| { + const b = r.fetch[2 * i]; + const s = r.fetch[2 * i + 1]; + if (s >= slots or b >= nblocks) return try std.fmt.bufPrint(&self.msg, "fetch entry {d} names block {d} slot {d}", .{ i, b, s }); + if (taken[s]) return try std.fmt.bufPrint(&self.msg, "slot {d} is given to two blocks", .{s}); + if (mo.hit[s]) return try std.fmt.bufPrint(&self.msg, "slot {d} is a slot this call uses and was taken", .{s}); + if (r.s2b[s] != b or r.b2s[b] != s or r.stamp[s] != mo.epoch) return try std.fmt.bufPrint(&self.msg, "block {d} and slot {d} disagree: block->slot {d}, slot->block {d}, stamp {d}", .{ b, s, r.b2s[b], r.s2b[s], r.stamp[s] }); + taken[s] = true; + new[b] = true; + olds[i] = self.stamp[s]; + if (self.s2b[s] != none) gone[self.s2b[s]] = true; + } + std.mem.sort(u32, olds, {}, less); + if (!std.mem.eql(u32, olds, mo.victims)) return "the slots it took are not the least recently used"; + const missed = try a.alloc(bool, nblocks); + defer a.free(missed); + @memset(missed, false); + for (mo.miss) |b| missed[b] = true; + for (0..nblocks) |b| { + const want: u32 = if (gone[b] or (missed[b] and !new[b])) none else self.b2s[b]; + if (new[b]) continue; + if (r.b2s[b] != want) return try std.fmt.bufPrint(&self.msg, "block {d} maps to slot {d}, not {d}", .{ b, r.b2s[b], want }); + } + for (0..slots) |s| { + if (taken[s]) continue; + if (r.s2b[s] != self.s2b[s]) return try std.fmt.bufPrint(&self.msg, "slot {d} holds block {d}, not {d}", .{ s, r.s2b[s], self.s2b[s] }); + const want = if (mo.hit[s]) mo.epoch else self.stamp[s]; + if (r.stamp[s] != want) return try std.fmt.bufPrint(&self.msg, "slot {d} has stamp {d}, not {d}", .{ s, r.stamp[s], want }); + } + return null; + } +}; + +/// `n` distinct picks from `pool` into `out`. +fn sample(rng: std.Random, a: std.mem.Allocator, pool: []const u32, n: usize, out: []u32) !void { + const work = try a.dupe(u32, pool); + defer a.free(work); + for (0..n) |i| { + const j = i + rng.uintLessThan(usize, work.len - i); + std.mem.swap(u32, &work[i], &work[j]); + out[i] = work[i]; + } +} + +/// The plan and its control over one loaded image, for the runner's list. +pub const Set = struct { + mod: cuda.Module, + each: [2]Case, + + pub fn init(gpu: check.Gpu, s: cuda.Stream) !Set { + var mod = try cuda.Module.load(gpu.d, cuda.kernels.flashnext_exl3_kvtier); + errdefer mod.unload(); + const f = try mod.function("kvt_plan"); + return .{ .mod = mod, .each = .{ .{ .a = gpu.gpa, .rig = .{ .gpu = gpu, .s = s }, .f = f }, .{ .a = gpu.gpa, .rig = .{ .gpu = gpu, .s = s }, .f = f, .control = true } } }; + } + + pub fn deinit(self: *Set) void { + for (&self.each) |*x| x.deinit(); + self.mod.unload(); + } + + pub fn add(self: *Set, a: std.mem.Allocator, list: *std.ArrayList(g.Kernel)) !void { + for (&self.each) |*x| try list.append(a, x.kernel()); + } +}; diff --git a/zig/tests/cuda/contract_flashnext_exl3_kvt_plan_test.zig b/zig/tests/cuda/contract_flashnext_exl3_kvt_plan_test.zig new file mode 100644 index 0000000000..1d79e76014 --- /dev/null +++ b/zig/tests/cuda/contract_flashnext_exl3_kvt_plan_test.zig @@ -0,0 +1,50 @@ +//! kvt_plan's fixtures on the host: the served calls reach each branch of the kernel, and every model is whole. +const std = @import("std"); +const core = @import("core"); +const plan = @import("contract_flashnext_exl3_kvt_plan.zig"); + +const g = core.contract.gates; + +test "kvt plan contract: the served calls reach each branch of the kernel, and every model is whole" { + const a = std.testing.allocator; + var assigned = false; + var overflow = false; + var all_hit = false; + var ties = false; + var wrapped = false; + var idle = false; + var sticky = false; + for (plan.served) |shape| { + var pinned = false; // some call needs the last block before the window, so the control's early window shows + for (0..14) |seed| { + var cs: plan.Case = .{ .a = a, .rig = undefined, .f = undefined }; + defer cs.free(); + for ([_]g.Fill{ .random, .nonfinite }) |fill| { + try std.testing.expect(try cs.build(shape, fill, seed)); + const mo = &cs.model; + const cap = @max(shape.n, 2048); + try std.testing.expect(mo.need.len <= cap and mo.miss.len <= cap); + var used: usize = 0; + for (mo.stamp) |t| used += @intFromBool(t == mo.epoch and mo.epoch != 0); + var hits: usize = 0; + for (mo.hit) |h| hits += @intFromBool(h); + if (mo.assigned > 0) { + assigned = true; + try std.testing.expectEqual(hits + mo.assigned, used); + try std.testing.expectEqual(mo.assigned, mo.miss.len); + try std.testing.expectEqual(@as(u32, @intCast(mo.assigned)), mo.ctl[1]); + for (mo.victims[1..], mo.victims[0 .. mo.victims.len - 1]) |hi, lo| ties = ties or hi == lo; + } + overflow = overflow or (mo.ctl[2] == 1 and mo.epoch != 0 and mo.miss.len > 0 and mo.assigned == 0 and cs.ctl[2] == 0); + all_hit = all_hit or (mo.need.len > 0 and mo.miss.len == 0); + wrapped = wrapped or mo.epoch == 0; + sticky = sticky or (cs.ctl[2] == 1 and mo.epoch != 0 and mo.need.len == 0); + idle = idle or (mo.epoch != 0 and cs.ctl[2] == 0 and mo.need.len == 0); + pinned = pinned or std.mem.indexOfScalar(u32, mo.need, plan.before - 1) != null; + try std.testing.expect(mo.ctl[0] == mo.epoch or mo.epoch == 0); + } + } + try std.testing.expect(pinned); + } + try std.testing.expect(assigned and overflow and all_hit and ties and wrapped and idle and sticky); +} diff --git a/zig/tests/cuda/contract_flashnext_exl3_suite.zig b/zig/tests/cuda/contract_flashnext_exl3_suite.zig index 27c4e7f24c..2875c7cd1e 100644 --- a/zig/tests/cuda/contract_flashnext_exl3_suite.zig +++ b/zig/tests/cuda/contract_flashnext_exl3_suite.zig @@ -9,6 +9,19 @@ const writeback = @import("contract_flashnext_writeback.zig"); const prefill_ = @import("contract_flashnext_exl3_prefill.zig"); const promptd_ = @import("contract_flashnext_exl3_promptd.zig"); const tppf_ = @import("contract_flashnext_exl3_tppf.zig"); +const kvt_ = @import("contract_flashnext_exl3_kvt.zig"); +const kvt_plan_ = @import("contract_flashnext_exl3_kvt_plan.zig"); +const flash = @import("flashnext"); + +comptime { + const wt = flash.exl3.prompt_wt; + const pages = flash.exl3.pages; + const rows = flash.exl3.kvtier.page_rows; + std.debug.assert(kvt_plan_.ratio == wt.index_ratio and kvt_plan_.top == wt.index_top and kvt_plan_.idw == wt.index_idw); + std.debug.assert(kvt_plan_.page_rows == rows and kvt_.row_bytes == pages.rowBytes(.key)); + std.debug.assert(kvt_.kv_page == pages.kvPageBytes(rows, 1) and kvt_.value_at == pages.planeOffset(rows, 0, .value)); + std.debug.assert(pages.planeOffset(rows, 0, .key) == 0); +} pub const Suite = struct { a: std.mem.Allocator, @@ -18,6 +31,7 @@ pub const Suite = struct { prefill: *prefill_.Group, promptd: *promptd_.Set, tppf: *tppf_.Set, + kvt: *kvt_.Group, pub fn create(a: std.mem.Allocator, gpu: check.Gpu, s: cuda.Stream, kernels: *std.ArrayList(g.Kernel), hc: cuda.Module, controls: cuda.Module) !*Suite { const self = try a.create(Suite); @@ -49,6 +63,7 @@ pub const Suite = struct { self.tppf = try tppf_.Set.create(a, gpu, s, controls); errdefer self.tppf.destroy(); try self.tppf.add(a, kernels); + self.kvt = try kvt_.Group.create(a, gpu, s, kernels); return self; } @@ -57,6 +72,7 @@ pub const Suite = struct { self.canons.deinit(); self.promptd.deinit(); self.a.destroy(self.promptd); + self.kvt.destroy(); self.prefill.destroy(); self.masks.deinit(); for (&self.writebacks) |*w| w.deinit(); diff --git a/zig/tests/cuda/contract_fn_pages.zig b/zig/tests/cuda/contract_fn_pages.zig index 5e7869d8c2..79d96835dc 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..af6cb4a523 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 } }, }; } @@ -59,56 +60,3 @@ test "paged media contract: real plans retain image axes, text suffixes and a ne } } } - -const Capture = struct { - var calls: usize = 0; - var symbols: [2]usize = undefined; - var 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 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].?))).*; - 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" { - var d: cuda.Driver = undefined; - d.api.cuLaunchKernel = Capture.launch; - const s: cuda.Stream = .{ .d = &d, .handle = null }; - const f = .{ - .prep_media = cuda.Function{ .d = &d, .handle = @ptrFromInt(1) }, - .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) }, - }; - 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| { - 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); - for (0..2) |i| { - try std.testing.expectEqual(pg, 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 04116b8751..544cfb8279 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 &.{}; return error.UnknownContractImage; 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 6e3cc52b4a..b5e68d0343 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -109,6 +109,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"); @@ -265,6 +266,7 @@ const usage = \\ exl3-prefill-contract the prefill expert tiles alone through the contract (gates 1-4 and the control) \\ exl3-promptd-contract the prompt tiles in decode's bytes (tf_exl3_promptd_*) alone through the contract \\ exl3-tppf-contract the TP prefill's rotate, combine, finish and place launches alone through the contract + \\ exl3-kvt-contract the host tier's kvt_plan, kvt_fetch and kvt_remap alone through the contract \\ 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 @@ -390,7 +392,9 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "exl3-prefill-contract")) return @import("contract_flashnext_exl3_prefill.zig").run(gpu); if (std.mem.eql(u8, cmd, "exl3-promptd-contract")) return @import("contract_flashnext_exl3_promptd.zig").run(gpu); if (std.mem.eql(u8, cmd, "exl3-tppf-contract")) return @import("contract_flashnext_exl3_tppf.zig").run(gpu); + if (std.mem.eql(u8, cmd, "exl3-kvt-contract")) return @import("contract_flashnext_exl3_kvt.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 36f2b75c3bf2c3150f52aa51f707f89bf210d8dd Mon Sep 17 00:00:00 2001 From: BlivionIaG <15323295+BlivionIaG@users.noreply.github.com> Date: Wed, 4 Nov 2026 00:32:14 +0100 Subject: [PATCH 7/7] EXL3 host tier: admit the bounded frame working set --- docs/recipes/flash-next-cuda.md | 10 +++- zig/src/families/flashnext/cuda_kv_lanes.zig | 49 +++++++++++++++++++- 2 files changed, 56 insertions(+), 3 deletions(-) diff --git a/docs/recipes/flash-next-cuda.md b/docs/recipes/flash-next-cuda.md index f55a9e6916..cf57c8b47e 100644 --- a/docs/recipes/flash-next-cuda.md +++ b/docs/recipes/flash-next-cuda.md @@ -128,7 +128,7 @@ starts a rank. These switches tune it: | `TF_FLASHNEXT_EXL3_ARCH` | unset | `unqualified` serves on a device that is not sm_120, without a claim. | | `TF_FLASHNEXT_EXL3_MTP_RESIDENT` | `0` | `1` keeps the MTP layer's experts on the device instead of a pool (one rank). | | `TF_FLASHNEXT_EXL3_PREFILL` | on | `0` turns the deep-chunk prefill path off. | -| `TF_FLASHNEXT_EXL3_PREFILL_ROWS` | 8,192 | Rows of a deep chunk; also what the server reserves at load. | +| `TF_FLASHNEXT_EXL3_PREFILL_ROWS` | 4,096 | Rows of a deep chunk; also what the server reserves at load. | | `TF_FLASHNEXT_EXL3_PREFILL_RING` | on | `0` stops streaming the coming layers' experts into one ring of the pools' lines. | | `TF_FLASHNEXT_EXL3_PREFILL_RING_AHEAD` | 1 | Layers ahead the ring prefetches. | | `TF_FLASHNEXT_EXL3_PREFILL_OVERLAP` | on | `0` (two ranks) stops overlapping the all-reduce with a chunk's compute. | @@ -137,8 +137,9 @@ starts a rank. These switches tune it: | `TF_FLASHNEXT_EXL3_PREFILL_SNAP_GRID` | on | Diagnostics only: `0` keeps prompt-cache states at any row instead of on the serving step's starts, to test that a resumed prompt equals a fresh one. | | `TF_FLASHNEXT_PROFILE` | `0` | Diagnostics only: `1` times the ranks' combine, all-reduce and place of a deep chunk, with a wait each. | | `TF_FLASHNEXT_KV_PAGES` | unset | `n`: streams beside the first take their attention caches from one pool of `n`-position pages. | +| `TF_FLASHNEXT_KV_PAGES_ALL` | unset | Diagnostics only: `1` puts every request on pool pages, including single-stream checks; without it the first stream keeps contiguous caches. | | `TF_FLASHNEXT_KV_HOST` | unset | With `TF_FLASHNEXT_KV_PAGES`: MiB of pinned host RAM holding the K/V plane of every page; `0` is no host tier. | -| `TF_FLASHNEXT_KV_FRAMES` | the pages the plan gave | With the host tier: at most this many device frames of the K/V plane. | +| `TF_FLASHNEXT_KV_FRAMES` | paged streams times their frame working set | With the host tier: cap the device frames of the K/V plane; fewer than one stream's required working set is refused. | | `TF_FLASHNEXT_KV_SLOTS` | 128 | With the host tier: frames lent to a stream, 64 four-row slots a frame for each attention layer. | | `TF_FLASHNEXT_PLE_READ` | `map` | `pread` reads the n-gram embedding rows by positional read. | | `TF_FLASHNEXT_SNAP_HOST` | `1` above 32,768 tokens of context, else `0` | Whether kept prompt-cache states live in pinned host memory. | @@ -146,6 +147,11 @@ starts a rank. These switches tune it: | `TF_FLASHNEXT_EXL3_LAYER_TELEMETRY` | `0` | Diagnostics only: `1`, with the switch above, adds per-layer counters. | | `TENSORFOLD_BIN`, `TF_SRC` | built binary, the checkout | Diagnostics only: the binary `tools/exl3_*.sh` run and the checkout `tools/exl3_*.py` read their corpus from. | +Host-tier admission promises logical pages for the request separately from its bounded device-frame working set: the first +pages, the current call and drafted window, the state run and the slot frames. It does not require the entire logical context +to be device-resident. A request waits when its logical pages, frame promise or pool lane cannot fit; released leases restore +admission. Leave `TF_FLASHNEXT_KV_FRAMES` unset for the backend's computed budget rather than guessing below that minimum. + ## What the core runs and what the family runs The family plugs into the core's pieces and adds no scheduler: diff --git a/zig/src/families/flashnext/cuda_kv_lanes.zig b/zig/src/families/flashnext/cuda_kv_lanes.zig index 9ec9d215c7..00d0ef1f97 100644 --- a/zig/src/families/flashnext/cuda_kv_lanes.zig +++ b/zig/src/families/flashnext/cuda_kv_lanes.zig @@ -255,7 +255,7 @@ pub const Pool = struct { const k = p.layout.runPages(); const n = p.layout.pages(cap); if (p.host) |h| { - for (p.lanes) |x| if (!x.used) return h.tier.fits(n, n + k); + for (p.lanes) |x| if (!x.used) return h.tier.fits(n, kvt.framesFor(k, p.slot_frames)); return false; } for (p.lanes) |x| if (!x.used) return p.pool.fits(n + k) and p.pool.findRun(k) != null; @@ -435,6 +435,53 @@ test "a stream's pages: its rows and a drafted window, no more than a full windo try std.testing.expect(l.pages(40960 + window_rows) <= mostPages(40960, 256)); } +test "exl3: host-tier admission follows frame reservations rather than logical context size" { + const t = std.testing; + const run_pages = 2; + const slots = 9; + const frames = comptime kvt.framesFor(run_pages, slots); + const logical = frames + run_pages + 8; + const page_bytes = 256; + var backing: [2 * logical * page_bytes]u8 = undefined; + var host: host_tier.Host = .{ .tier = undefined, .store = undefined, .kv_page = page_bytes }; + try host.tier.init(t.allocator, &backing, frames, page_bytes); + defer host.tier.deinit(); + var lanes = [_]Lane{.{ .seq = undefined }}; + const pool: Pool = .{ + .gpa = t.allocator, + .e = undefined, + .layout = .{ .fixed_bytes = 0, .token_bytes = 1, .page_tokens = page_bytes, .snap_bytes = 0, .run_bytes = run_pages * page_bytes }, + .capacity = logical * page_bytes, + .head = null, + .plain = undefined, + .pool = &host.tier.pool, + .host = &host, + .slot_frames = slots, + .arena = undefined, + .ibase = 0, + .lanes = &lanes, + }; + try t.expect(pool.fits(pool.capacity)); + { + var lease: kv.tier.Lease = .{}; + try host.tier.reserve(&lease, 1, frames); + defer host.tier.finish(&lease); + try t.expect(!pool.fits(pool.capacity)); + } + try t.expect(pool.fits(pool.capacity)); + { + var lease: kv.tier.Lease = .{}; + try host.tier.reserve(&lease, 2 * logical, 0); + defer host.tier.finish(&lease); + try t.expect(!pool.fits(pool.capacity)); + } + try t.expect(pool.fits(pool.capacity)); + lanes[0].used = true; + try t.expect(!pool.fits(pool.capacity)); + lanes[0].used = false; + try t.expect(!pool.fits(pool.capacity + window_rows + 1)); +} + test { _ = @import("cuda_kv_init_test.zig"); _ = @import("cuda_kv_lifetime_test.zig");