From bd6353804bb988d1a735494cd830d17f5b5a84d2 Mon Sep 17 00:00:00 2001 From: Georg Bauer Date: Fri, 4 Sep 2026 23:27:09 +0200 Subject: [PATCH] Optimize Qwen3.8 inference on Apple silicon --- Cargo.lock | 1 + Cargo.toml | 1 + docs/QWEN38_PERFORMANCE.md | 146 + docs/qwen38-issue-99.csv | 18 + metal/qwen38.metal | 3334 ++++++++++++++++-- native/metal/ds4_gpu.h | 5 + native/metal/ds4_metal.m | 154 +- src/config.rs | 47 +- src/engine.rs | 378 +- src/engine/metal.rs | 21 +- src/engine/metal/gpu.rs | 13 + src/engine/metal/qwen.rs | 6844 +++++++++++++++++++++++++++++++----- src/engine/qwen.rs | 327 +- src/engine/tokenizer.rs | 5 +- 14 files changed, 10149 insertions(+), 1145 deletions(-) create mode 100644 docs/QWEN38_PERFORMANCE.md create mode 100644 docs/qwen38-issue-99.csv diff --git a/Cargo.lock b/Cargo.lock index 2ef2ac6..ef845ab 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -825,6 +825,7 @@ dependencies = [ "pulldown-cmark 0.12.2", "regex", "rfd", + "rustc-hash 2.1.3", "serde", "serde_json", "serde_norway", diff --git a/Cargo.toml b/Cargo.toml index 66c665e..2c9d2d4 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -26,6 +26,7 @@ png = "0.17.16" pulldown-cmark = "0.12.2" regex = "1.13.1" rfd = { version = "0.15.4", default-features = false } +rustc-hash = "2.1.3" serde = { version = "1.0.228", features = ["derive"] } serde_json = { version = "1.0.149", features = ["preserve_order", "raw_value"] } serde_norway = "0.9.42" diff --git a/docs/QWEN38_PERFORMANCE.md b/docs/QWEN38_PERFORMANCE.md new file mode 100644 index 0000000..5516dfd --- /dev/null +++ b/docs/QWEN38_PERFORMANCE.md @@ -0,0 +1,146 @@ +# Qwen3.8 Flash Next performance and parity receipt + +Issue #99 optimizes the pinned `Youssofal/Qwen3.8-Flash-Next-MTPLX-Bare-Speed` +text model without changing its artifact precision, sampler, or token behavior. +All application and inference work remains Rust; native additions are confined to +the existing Objective-C Metal bridge and Qwen `.metal` kernels. Integrated MTP +is the Qwen production default; an explicitly stored user opt-out remains +authoritative. + +## Result + +The acceptance B1 median is **60.825779 tok/s MTP**. A later B2 confirmation, +after adding complete repair telemetry and under the then-current machine load, +measured **58.019436 tok/s**. Both sequences produced the same 64-token SHA-256 +`46ce63601c44d80b5ef8297eaa49a636f6680040b916fe30cf7a154c05daef72`, with 44 +accepted drafts from 60 proposals in 20 verifier passes. The fresh matched MTPLX +median was 59.751170 tok/s, putting B1 1.8% above and B2 2.9% below the oracle; +both are inside the issue's 5% parity band. Raw per-run data is in +`docs/qwen38-issue-99.csv`. + +Plain AR reached a 50.671077 tok/s median with stable token hash +`15373689bcb1c667123accfcfbe49736466f28ce7beb89d2cbe5e4120416056d`. +The published 47 tok/s MTPLX figure is contextual rather than a fresh matched +AR baseline: after the valid MTP oracle run, the installed MTPLX package and +checked-out source drifted apart. The current CLI rejects this Flash Next pack +as unsupported and the source path fails at the missing +`DecoderLayer.input_layernorm`, so no invalid number is presented as a fresh AR +comparison. + +The published 75.9 tok/s MTP result is from a different MTPLX runtime/workload. +With identical pinned artifacts, 4,096 context, prompt, sampling, seed, and +64-token measurement window, the current usable oracle result is 59.751170 +tok/s. The remaining gap is therefore not explained by quantization drift or a +different output distribution. + +Maximum fan control could not be established because neither ThermalForge nor +TG Pro is installed; MTPLX reports its fan boost as unavailable. Both compared +engines ran under the same no-boost constraint. No system utility was installed +or substituted for the missing control. + +## Fixed workload + +- Hardware: MacBook Pro `Mac17,7`, Apple M5 Max, 18 CPU cores, 128 GiB unified + memory. +- OS: macOS 26.6.2, build 25G83, arm64. +- Model revision: `74559cdf34fbfc0b593de72d17e93f37fd4f9ea7`. +- Model manifest contract SHA-256: + `6f1172de47fa30b9602e13fc7ad14e578a813bb5b1fab8f0ef320041aac6c19a`. +- Tensor inventory SHA-256: + `b5731e6febcf865d276a0e7b144da02375f2d7e3129594f271ee9c6c351f4c8f`. +- `config.json` SHA-256: + `826302af9ff3293cfc9975c9a5a96546fa0aa9f0808505adfc30a62fef806965`. +- `mtplx_runtime.json` SHA-256: + `381c7666872f878542494a48a874bb0784c72a955d18cc077324d3c9e74b8cdb`. +- Prompt: `Write a concise Rust implementation of binary search and explain its invariants.` +- Prompt tokens: 26; context: 4,096; generated tokens: 64. +- Sampling: temperature 1.0, top-p 0.95, top-k 20, min-p 0, seed 12,345. +- One process, three runs: first run cold, following runs warm. + +The exact DS4Server commands were: + +```sh +DS4SERVER_QWEN38_SOURCE=/private/tmp/qwen38-bare-source \ +DS4SERVER_QWEN38_MEASURED=64 DS4SERVER_QWEN38_RUNS=3 \ +cargo test --release qwen_m5_performance_gate -- --ignored --nocapture + +DS4SERVER_QWEN38_SOURCE=/private/tmp/qwen38-bare-source \ +DS4SERVER_QWEN38_MTP=1 DS4SERVER_QWEN38_MEASURED=64 \ +DS4SERVER_QWEN38_RUNS=3 \ +cargo test --release qwen_m5_performance_gate -- --ignored --nocapture + +DS4SERVER_QWEN38_SOURCE=/private/tmp/qwen38-bare-source \ +cargo test --release qwen_mtp_matches_target_and_restores_its_cache \ + -- --ignored --nocapture +``` + +The MTPLX comparison used its normal `serve` path with the same local Bare Speed +directory, generation mode, prompt, context, sampler, seed, and single request +stream. A later retry is not a valid benchmark because the local MTPLX +source/package compatibility gate fails as described above. + +## MTPLX logic review + +The production path now follows the MTPLX/MLX graph rather than retaining the +temporary Qwen experiments used while earlier issues established correctness: + +- PLE uses the artifact's 4-bit, group-32 rows, streams them with `pread`, keeps + a bounded hot cache, and publishes hit, miss, eviction, byte, and read-time + counters. The warm MTP runs had 2,080 hits, no misses, and no reads. +- GDN uses fused input projection, MTPLX-equivalent conv+SiLU+L2 normalization, + BF16 rounding points, decay/delta recurrence, gated norm, and direct verifier + capture buffers. Rejected rows restore or advance authoritative recurrence + rather than inventing state. +- QSA uses the MTPLX `qsa_flash_skip` row ownership, online-softmax merge, and + native 256-thread/eight-SIMD layout. Long-context state remains linear and + admitted through 262,144 tokens without a gather-copy fallback. +- MoE uses the exact router ownership and merged gate/up projection. The custom + verifier MoE experiment and NAX path were removed because they are not the + oracle's promoted path. +- MTP uses the same checkpoint, history, proposal, verifier, rejection-repair, + full-accept bonus, and NumPy-compatible PCG64 sampling rules. The exact cache + restoration test passes after the final telemetry change. +- Sparse top-k sampling stays on device for the supported official sampler and + uses the existing exact CPU path only when the requested sampler is outside + that kernel's contract. Full production logits are not copied to the CPU. +- Capability fallbacks remain exact implementations, not alternate Qwen + behavior. Experiments that changed tokens (normal Metal math, wider Tensor + QMM, and recurrence-only repair) were removed rather than left dormant. + +The 4-bit/group-64 core, declared 8-bit router/QSA matrices, BF16 recurrent and +normalization tensors, 4-bit/group-32 PLE, and published 2/4/8-bit MTP tensors +are consumed as stored. No blanket up-conversion or replacement quantization is +used. Peak RSS in the final MTP sequence was 52.61 GB, down from the original +102.65 GB path; the 105.44 GB artifact size includes a 32.00 GB file-backed PLE +table that is not resident wholesale. + +## A/B evidence + +The optimization work used exact hashes or tensor comparisons before promotion: + +| Area | Before | After | Decision | +| --- | ---: | ---: | --- | +| Resident memory | 102.65 GB RSS | 52.61 GB RSS | Reused staging and streamed PLE promoted | +| Router | 49.7 ms GPU | 18.7 ms GPU, 23.4 ms wall | Row-owned reduced router promoted | +| 16K QSA AR | 2.6 tok/s | 45.7 tok/s | MTPLX flash-skip rewrite promoted | +| Shared verifier gate/up | 54.7 tok/s | 58.3 tok/s | Exact fused projection promoted | +| Direct GDN verifier capture | 58.5/58.7 tok/s | 60.8 tok/s | Direct capture buffers promoted | + +The direct-capture B1 acceptance was followed by the B2 confirmation in the CSV; +both retained the same token hash. Normal Metal math, Tensor QMM at verifier +width, and shortened recurrence repair failed exact-oracle checks and were +removed. Custom MoE verify, fused GDN output projection, and async AR were not +promoted because the MTPLX evidence did not show a win and DS4Server profiles did +not justify overriding that result. + +## Context and verification + +The long-context gate exercised 2,048, 16,384, 65,536, 131,072, and the native +262,144-token maximum. AR remained 43.488 tok/s at the maximum, and live state +grew from 58,195,968 bytes at 2K to 7,449,083,904 bytes at 262K without collapse +or over-admission. The CSV records every context result and state size. + +Final issue verification comprises the exact MTP cache/token oracle, affine +QMV/QMM MLX reference comparisons, focused Qwen tests, the long-context state +gate, the release performance gate, formatting, Clippy with warnings denied, the +macOS application bundle, and the full all-features Rust test suite. diff --git a/docs/qwen38-issue-99.csv b/docs/qwen38-issue-99.csv new file mode 100644 index 0000000..cb6f730 --- /dev/null +++ b/docs/qwen38-issue-99.csv @@ -0,0 +1,18 @@ +date,engine,mode,sequence,run,state,context,prompt_tokens,generated_tokens,temperature,top_p,top_k,seed,prefill_tps,first_token_ms,steady_tps,decode_tps,p50_cycle_ms,p95_cycle_ms,peak_rss_bytes,token_sha256,cycles,drafted,accepted,verifier_passes,verifier_ms,repair_ms,ple_entries,ple_hits,ple_misses,ple_evictions,ple_pread_bytes,ple_pread_ms,notes +2026-09-04,DS4Server,MTP,B1,1,cold,4096,26,64,1.0,0.95,20,12345,6.360708,59.484667,57.518099,58.042801,54.731042,57.838709,52720156672,46ce63601c44d80b5ef8297eaa49a636f6680040b916fe30cf7a154c05daef72,20,60,44,20,1000,0,1648,432,1648,0,164800,,repair timing counter not yet wired +2026-09-04,DS4Server,MTP,B1,2,warm,4096,26,64,1.0,0.95,20,12345,51.095361,52.229584,60.307983,61.119864,52.427709,54.273459,52720156672,46ce63601c44d80b5ef8297eaa49a636f6680040b916fe30cf7a154c05daef72,20,60,44,20,946,0,1648,2080,0,0,0,0,repair timing counter not yet wired +2026-09-04,DS4Server,MTP,B1,3,warm,4096,26,64,1.0,0.95,20,12345,51.152493,52.422375,60.014218,60.825779,52.679334,54.147250,52720156672,46ce63601c44d80b5ef8297eaa49a636f6680040b916fe30cf7a154c05daef72,20,60,44,20,951,0,1648,2080,0,0,0,0,acceptance median 60.825779 +2026-09-04,DS4Server,MTP,B2,1,cold,4096,26,64,1.0,0.95,20,12345,6.527240,62.605750,54.320764,54.834156,57.909625,60.828916,52605190144,46ce63601c44d80b5ef8297eaa49a636f6680040b916fe30cf7a154c05daef72,20,60,44,20,1056,30,1648,432,1648,0,164800,63,confirmation after repair telemetry +2026-09-04,DS4Server,MTP,B2,2,warm,4096,26,64,1.0,0.95,20,12345,50.178762,54.061500,57.196397,58.019436,55.002750,57.242875,52605190144,46ce63601c44d80b5ef8297eaa49a636f6680040b916fe30cf7a154c05daef72,20,60,44,20,993,29,1648,2080,0,0,0,0,confirmation median 58.019436 +2026-09-04,DS4Server,MTP,B2,3,warm,4096,26,64,1.0,0.95,20,12345,49.976289,55.289000,57.556598,58.301539,54.853000,56.619333,52605190144,46ce63601c44d80b5ef8297eaa49a636f6680040b916fe30cf7a154c05daef72,20,60,44,20,988,30,1648,2080,0,0,0,0,confirmation after repair telemetry +2026-09-04,DS4Server,AR,A1,1,cold,4096,26,64,1.0,0.95,20,12345,7.796453,21.087125,49.845055,49.805297,18.894042,19.426917,53962735616,15373689bcb1c667123accfcfbe49736466f28ce7beb89d2cbe5e4120416056d,0,0,0,0,0,0,1648,16,1424,0,,,plain autoregressive +2026-09-04,DS4Server,AR,A1,2,warm,4096,26,64,1.0,0.95,20,12345,51.386057,18.606375,51.034286,51.074537,18.392375,18.953542,53962735616,15373689bcb1c667123accfcfbe49736466f28ce7beb89d2cbe5e4120416056d,0,0,0,0,0,0,1648,,,,,,plain autoregressive +2026-09-04,DS4Server,AR,A1,3,warm,4096,26,64,1.0,0.95,20,12345,51.424682,18.256166,50.610874,50.671077,18.575958,19.129667,53962735616,15373689bcb1c667123accfcfbe49736466f28ce7beb89d2cbe5e4120416056d,0,0,0,0,0,0,1648,,,,,,plain autoregressive median 50.671077 +2026-09-04,DS4Server,AR,LONG,1,warm,2048,,1,,,,,0.778000,,,29.259000,,,52158349312,b6d28407,,,,,,,,,,,,,synthetic long-context state gate token=5 live_state_bytes=58195968 +2026-09-04,DS4Server,AR,LONG,1,warm,16384,,1,,,,,12.101000,,,46.486000,,,,79d71772,,,,,,,,,,,,,synthetic long-context state gate token=5 live_state_bytes=465567744 +2026-09-04,DS4Server,AR,LONG,1,warm,65536,,1,,,,,10.037000,,,45.739000,,,,0a1131fe,,,,,,,,,,,,,synthetic long-context state gate token=5 live_state_bytes=1862270976 +2026-09-04,DS4Server,AR,LONG,1,warm,131072,,1,,,,,3.277000,,,45.145000,,,,fb678d10,,,,,,,,,,,,,synthetic long-context state gate token=5 live_state_bytes=3724541952 +2026-09-04,DS4Server,AR,LONG,1,warm,262144,,1,,,,,0.777000,,,43.488000,,,,01ae755c,,,,,,,,,,,,,synthetic native-maximum state gate token=5 live_state_bytes=7449083904 +2026-09-04,MTPLX,MTP,ORACLE,1,warm,4096,26,64,1.0,0.95,20,12345,,,,59.170000,,,,,20,58,44,20,950,,,,,,,,fresh matched run before local source-package drift +2026-09-04,MTPLX,MTP,ORACLE,2,warm,4096,26,64,1.0,0.95,20,12345,,,,59.820000,,,,,20,58,44,20,950,,,,,,,,fresh matched run before local source-package drift +2026-09-04,MTPLX,MTP,ORACLE,3,warm,4096,26,64,1.0,0.95,20,12345,,,,59.751170,,,,,20,58,44,20,950,,,,,,,,matched median 59.751170 diff --git a/metal/qwen38.metal b/metal/qwen38.metal index 65972ef..2bab705 100644 --- a/metal/qwen38.metal +++ b/metal/qwen38.metal @@ -1,6 +1,13 @@ // Qwen3.8 Flash Next primitives. Rust owns the graph and all state lifetimes; // this file contains only the data-parallel kernels executed by Metal. +#include +#ifdef DS4_METAL_HAS_TENSOR +#include +#include +#endif +using namespace metal; + struct qwen_kernel_args { uint u[16]; float f[8]; @@ -16,6 +23,42 @@ static inline ushort qwen_to_bf16(float value) { return (ushort)(bits >> 16); } +static inline float qwen_round_bf16(float value) { + return qwen_bf16(qwen_to_bf16(value)); +} + +static inline bfloat qwen_simd_sum_bf16(bfloat value) { + float sum = (float)value; + for (ushort offset = 16u; offset >= 1u; offset >>= 1u) { + sum += simd_shuffle_down(sum, offset); + } + return (bfloat)sum; +} + +static inline float qwen_sigmoid_bf16(float value) { + const float input = qwen_round_bf16(value); + const float magnitude = qwen_round_bf16(input < 0.0f ? -input : input); + const float exponential = qwen_round_bf16(metal::precise::exp(magnitude)); + const float denominator = qwen_round_bf16(1.0f + exponential); + const float tail = qwen_round_bf16(1.0f / denominator); + return input < 0.0f ? tail : qwen_round_bf16(1.0f - tail); +} + +static inline float qwen_silu_sigmoid_bf16(float value) { + const bfloat input = (bfloat)value; + const bfloat magnitude = input < bfloat(0) ? -input : input; + const bfloat exponential = (bfloat)exp((float)magnitude); + const bfloat denominator = (bfloat)(1.0f + (float)exponential); + const bfloat tail = (bfloat)(1.0f / (float)denominator); + return (float)(input < bfloat(0) ? tail : (bfloat)(1.0f - (float)tail)); +} + +static inline float qwen_apply_rope(float value, float rotated, float theta) { + volatile float cosine_term = value * cos(theta); + volatile float sine_term = rotated * sin(theta); + return cosine_term + sine_term; +} + static inline ushort qwen_weight_u16( device const uchar *data, uint byte_offset, @@ -37,6 +80,24 @@ static inline uint qwen_weight_u32( ((uint)data[byte + 3u] << 24u); } +kernel void kernel_qwen_probe_affine_table( + constant qwen_kernel_args &args [[buffer(0)]], + device uint *out [[buffer(1)]], + device const uint *selected [[buffer(3)]], + device const uchar *packed [[buffer(5)]], + device const uchar *scales [[buffer(6)]], + device const uchar *biases [[buffer(7)]], + uint slot [[thread_position_in_grid]]) { + if (slot >= args.u[4]) return; + const uint expert = selected[slot]; + const ulong row = (ulong)expert * args.u[1]; + const ulong packed_index = row * (args.u[0] / (32u / args.u[2])); + const ulong parameter = row * (args.u[0] / args.u[3]); + out[slot * 3u] = qwen_weight_u32(packed, args.u[13], packed_index); + out[slot * 3u + 1u] = qwen_weight_u16(scales, args.u[14], parameter); + out[slot * 3u + 2u] = qwen_weight_u16(biases, args.u[15], parameter); +} + static inline float qwen_quant_weight( device const uchar *packed, device const uchar *scales, @@ -98,7 +159,1054 @@ kernel void kernel_qwen_affine_mv( args.u[13], args.u[14], args.u[15], row, column, in_dim, args.u[2], args.u[3]), x[column], sum); } - out[row] = sum; + out[row] = args.u[11] != 0u ? qwen_round_bf16(sum) : sum; +} + +// Decode uses one activation row. Each SIMD group owns four output rows so +// the packed-weight decode shares every activation load across those rows, +// matching the affine qmv shape used by the reference MLX backend. +kernel void kernel_qwen_affine_qmv( + constant qwen_kernel_args &args [[buffer(0)]], + device float *out [[buffer(1)]], + device const float *x [[buffer(2)]], + device const uchar *packed [[buffer(5)]], + device const uchar *scales [[buffer(6)]], + device const uchar *biases [[buffer(7)]], + uint2 group [[threadgroup_position_in_grid]], + uint simd_group [[simdgroup_index_in_threadgroup]], + uint lane [[thread_index_in_simdgroup]]) { + const uint in_dim = args.u[0]; + const uint out_dim = args.u[1]; + const uint bits = args.u[2]; + const uint group_size = args.u[3]; + const uint values_per_word = 32u / bits; + const uint packed_columns = in_dim / values_per_word; + const uint words_per_group = group_size / values_per_word; + const uint groups_per_row = in_dim / group_size; + const uint first_row = group.x * 8u + simd_group * 4u; + const uint mask = (1u << bits) - 1u; + float sums[4] = {0.0f, 0.0f, 0.0f, 0.0f}; + + for (uint word_index = lane; word_index < packed_columns; word_index += 32u) { + const uint first_column = word_index * values_per_word; + const uint quant_group = word_index / words_per_group; + uint words[4] = {0u, 0u, 0u, 0u}; + float scale[4] = {0.0f, 0.0f, 0.0f, 0.0f}; + float bias[4] = {0.0f, 0.0f, 0.0f, 0.0f}; + for (uint output = 0; output < 4u; output++) { + const uint row = first_row + output; + if (row < out_dim) { + words[output] = qwen_weight_u32( + packed, args.u[13], (ulong)row * packed_columns + word_index); + const ulong parameter = (ulong)row * groups_per_row + quant_group; + scale[output] = qwen_bf16(qwen_weight_u16(scales, args.u[14], parameter)); + bias[output] = qwen_bf16(qwen_weight_u16(biases, args.u[15], parameter)); + } + } + for (uint packed_value = 0; packed_value < values_per_word; packed_value++) { + const float activation = x[first_column + packed_value]; + const uint shift = packed_value * bits; + for (uint output = 0; output < 4u; output++) { + const float weight = fma((float)((words[output] >> shift) & mask), + scale[output], bias[output]); + sums[output] = fma(weight, activation, sums[output]); + } + } + } + for (uint output = 0; output < 4u; output++) { + const float sum = simd_sum(sums[output]); + if (lane == 0u && first_row + output < out_dim) { + out[first_row + output] = args.u[11] != 0u ? qwen_round_bf16(sum) : sum; + } + } +} + +template +static inline void qwen_affine_qmv_fast_impl( + constant qwen_kernel_args &args, + device float *out, + device const float *x, + device const uchar *packed, + device const uchar *scale_bytes, + device const uchar *bias_bytes, + uint packed_offset, + uint scale_offset, + uint bias_offset, + uint2 group, + uint simd_group, + uint lane, + uint expert, + uint input_slot, + uint output_slot) { + static_assert(bits == 4 || bits == 8, "Qwen affine QMV supports 4/8-bit packs"); + static_assert(group_size == 32 || group_size == 64, + "Qwen affine QMV supports group sizes 32/64"); + constexpr ushort packs_per_thread = bits == 2 ? 1 : 2; + constexpr ushort pack_factor = 32 / bits; + constexpr ushort values_per_thread = pack_factor * packs_per_thread; + constexpr ushort block_size = values_per_thread * 32; + constexpr ushort scale_step = group_size / values_per_thread; + const uint in_dim = args.u[0]; + const uint out_dim = args.u[1]; + const uint input_bytes_per_row = in_dim * bits / 8; + const uint groups_per_row = in_dim / group_size; + const uint first_row = group.x * 8u + simd_group * 4u; + const ulong table_row = (ulong)expert * out_dim + first_row; + device const uchar *weights = packed + packed_offset + + table_row * input_bytes_per_row + lane * packs_per_thread * 4u; + device const ushort *scales = + (device const ushort *)(scale_bytes + scale_offset); + device const ushort *biases = + (device const ushort *)(bias_bytes + bias_offset); + ulong parameter = table_row * groups_per_row + lane / scale_step; + device const float *input = x + + (ulong)input_slot * in_dim + lane * values_per_thread; + float result[4] = {0.0f, 0.0f, 0.0f, 0.0f}; + + for (uint column = 0; column < in_dim; column += block_size) { + float values[values_per_thread]; + float input_sum = 0.0f; + for (ushort i = 0; i < values_per_thread; i++) { + const float value = input[i]; + if constexpr (bits == 4) { + constexpr float divisors[4] = {1.0f, 16.0f, 256.0f, 4096.0f}; + values[i] = value / divisors[i & 3u]; + } else { + values[i] = value; + } + } + if constexpr (bits == 4) { + for (ushort i = 0; i < values_per_thread; i += 4u) { + const bfloat group_sum = + (bfloat)input[i] + (bfloat)input[i + 1u] + + (bfloat)input[i + 2u] + (bfloat)input[i + 3u]; + input_sum += (float)group_sum; + } + } else { + for (ushort i = 0; i < values_per_thread; i++) input_sum += input[i]; + } + for (ushort row = 0; row < 4; row++) { + device const uchar *row_weights = weights + row * input_bytes_per_row; + float dot = 0.0f; + if constexpr (bits == 4) { + for (ushort i = 0; i < values_per_thread / 4; i++) { + const ushort word = ((device const ushort *)row_weights)[i]; + dot += values[4 * i] * (word & 0x000fu) + + values[4 * i + 1] * (word & 0x00f0u) + + values[4 * i + 2] * (word & 0x0f00u) + + values[4 * i + 3] * (word & 0xf000u); + } + } else { + for (ushort i = 0; i < values_per_thread; i++) { + dot += values[i] * row_weights[packed_offset + i]; + } + } + const ulong row_parameter = parameter + (ulong)row * groups_per_row; + result[row] += + qwen_bf16(scales[row_parameter]) * dot + + input_sum * + qwen_bf16(biases[row_parameter]); + } + weights += block_size * bits / 8; + parameter += block_size / group_size; + input += block_size; + } + + for (ushort row = 0; row < 4; row++) { + const float sum = simd_sum(result[row]); + if (lane == 0u) { + out[(ulong)output_slot * out_dim + first_row + row] = + args.u[11] != 0u ? (float)(bfloat)sum : sum; + } + } +} + +template +kernel void kernel_qwen_affine_qmv_fast( + constant qwen_kernel_args &args [[buffer(0)]], + device float *out [[buffer(1)]], + device const float *x [[buffer(2)]], + device const uchar *packed [[buffer(5)]], + device const uchar *scales [[buffer(6)]], + device const uchar *biases [[buffer(7)]], + uint group [[threadgroup_position_in_grid]], + uint simd_group [[simdgroup_index_in_threadgroup]], + uint lane [[thread_index_in_simdgroup]]) { + qwen_affine_qmv_fast_impl( + args, out, x, packed, scales, biases, args.u[13], args.u[14], args.u[15], + group, simd_group, lane, + 0u, 0u, 0u); +} + +// MLX affine_qmv_wide for the verifier's S=2..4 matrices. M5-class GPUs +// route these shapes here instead of running one independent QMV per row. +template +static inline void qwen_affine_qmv_wide_impl( + constant qwen_kernel_args &args, + device float *out, + device const float *x, + device const uchar *packed, + device const uchar *scales, + device const uchar *biases, + uint group, + uint simd_group, + uint lane) { + constexpr uint k_lanes = 8u; + constexpr uint outputs_per_simdgroup = 4u; + constexpr uint outputs_per_group = 8u; + constexpr uint sub = 8u; + const uint k_lane = lane % k_lanes; + const uint simd_row = lane / k_lanes; + const uint output_row = group * outputs_per_group + + simd_group * outputs_per_simdgroup + simd_row; + const uint in_dim = args.u[0]; + const uint out_dim = args.u[1]; + const uint rows = args.u[4]; + const uint row = min(output_row, out_dim - 1u); + const uint groups_per_row = in_dim / group_size; + const ulong packed_row = (ulong)row * in_dim * bits / 8u; + const ulong parameter_row = (ulong)row * groups_per_row; + float result[4] = {0.0f, 0.0f, 0.0f, 0.0f}; + + for (uint quant_group = k_lane; quant_group < groups_per_row; + quant_group += k_lanes) { + const float scale = qwen_bf16(qwen_weight_u16( + scales, args.u[14], parameter_row + quant_group)); + const float bias = qwen_bf16(qwen_weight_u16( + biases, args.u[15], parameter_row + quant_group)); + for (uint chunk = 0u; chunk < group_size / sub; chunk++) { + const uint column = quant_group * group_size + chunk * sub; + const ulong weight_byte = (ulong)args.u[13] + packed_row + + (ulong)column * bits / 8u; + float weights[sub]; + for (uint index = 0u; index < sub; index++) { + uint quantized; + if constexpr (bits == 4) { + const uchar byte = packed[weight_byte + index / 2u]; + quantized = (byte >> ((index & 1u) * 4u)) & 15u; + } else { + quantized = packed[weight_byte + index]; + } + weights[index] = scale * (float)quantized + bias; + } + for (uint vector = 0u; vector < rows; vector++) { + float sum = 0.0f; + for (uint index = 0u; index < sub; index++) { + sum += x[(ulong)vector * in_dim + column + index] * weights[index]; + } + result[vector] += sum; + } + } + } + + for (uint vector = 0u; vector < rows; vector++) { + result[vector] += simd_shuffle_down(result[vector], 4u); + result[vector] += simd_shuffle_down(result[vector], 2u); + result[vector] += simd_shuffle_down(result[vector], 1u); + if (k_lane == 0u && output_row < out_dim) { + out[(ulong)vector * out_dim + output_row] = args.u[11] != 0u + ? qwen_round_bf16(result[vector]) + : result[vector]; + } + } +} + +template +kernel void kernel_qwen_affine_qmv_batch_fast( + constant qwen_kernel_args &args [[buffer(0)]], + device float *out [[buffer(1)]], + device const float *x [[buffer(2)]], + device const uchar *packed [[buffer(5)]], + device const uchar *scales [[buffer(6)]], + device const uchar *biases [[buffer(7)]], + uint2 group [[threadgroup_position_in_grid]], + uint simd_group [[simdgroup_index_in_threadgroup]], + uint lane [[thread_index_in_simdgroup]]) { + qwen_affine_qmv_wide_impl( + args, out, x, packed, scales, biases, group.x, simd_group, lane); +} + +template +kernel void kernel_qwen_affine_pair_qmv_wide( + constant qwen_kernel_args &args [[buffer(0)]], + device float *gate_out [[buffer(1)]], + device const float *x [[buffer(2)]], + device float *up_out [[buffer(4)]], + device const uchar *gate_packed [[buffer(5)]], + device const uchar *gate_scales [[buffer(6)]], + device const uchar *gate_biases [[buffer(7)]], + device const uchar *up_packed [[buffer(9)]], + device const uchar *up_scales [[buffer(10)]], + device const uchar *up_biases [[buffer(11)]], + uint2 group [[threadgroup_position_in_grid]], + uint simd_group [[simdgroup_index_in_threadgroup]], + uint lane [[thread_index_in_simdgroup]]) { + const uint projection_groups = (args.u[1] + 7u) / 8u; + const bool up = group.x >= projection_groups; + qwen_affine_qmv_wide_impl( + args, + up ? up_out : gate_out, + x, + up ? up_packed : gate_packed, + up ? up_scales : gate_scales, + up ? up_biases : gate_biases, + group.x - (up ? projection_groups : 0u), + simd_group, + lane); +} + +template +kernel void kernel_qwen_affine_gather_qmv_fast( + constant qwen_kernel_args &args [[buffer(0)]], + device float *out [[buffer(1)]], + device const float *x [[buffer(2)]], + device const uint *selected [[buffer(3)]], + device const uchar *packed [[buffer(5)]], + device const uchar *scales [[buffer(6)]], + device const uchar *biases [[buffer(7)]], + uint2 group [[threadgroup_position_in_grid]], + uint simd_group [[simdgroup_index_in_threadgroup]], + uint lane [[thread_index_in_simdgroup]]) { + const uint slot = group.y; + qwen_affine_qmv_fast_impl( + args, out, x, packed, scales, biases, args.u[13], args.u[14], args.u[15], + group, simd_group, lane, + selected[slot], args.u[5] != 0u ? slot : 0u, slot); +} + +template +kernel void kernel_qwen_affine_gather_qmv_batch_fast( + constant qwen_kernel_args &args [[buffer(0)]], + device float *out [[buffer(1)]], + device const float *x [[buffer(2)]], + device const uint *selected [[buffer(3)]], + device const uchar *packed [[buffer(5)]], + device const uchar *scales [[buffer(6)]], + device const uchar *biases [[buffer(7)]], + uint2 group [[threadgroup_position_in_grid]], + uint simd_group [[simdgroup_index_in_threadgroup]], + uint lane [[thread_index_in_simdgroup]]) { + const uint slot = group.y; + qwen_affine_qmv_fast_impl( + args, out, x, packed, scales, biases, args.u[13], args.u[14], args.u[15], + group, simd_group, lane, selected[slot], + args.u[5] != 0u ? slot : slot / args.u[8], slot); +} + +template +kernel void kernel_qwen_affine_pair_qmv_fast( + constant qwen_kernel_args &args [[buffer(0)]], + device float *gate_out [[buffer(1)]], + device const float *x [[buffer(2)]], + device const uint *selected [[buffer(3)]], + device float *up_out [[buffer(4)]], + device const uchar *gate_packed [[buffer(5)]], + device const uchar *gate_scales [[buffer(6)]], + device const uchar *gate_biases [[buffer(7)]], + device const uchar *up_packed [[buffer(9)]], + device const uchar *up_scales [[buffer(10)]], + device const uchar *up_biases [[buffer(11)]], + uint2 group [[threadgroup_position_in_grid]], + uint simd_group [[simdgroup_index_in_threadgroup]], + uint lane [[thread_index_in_simdgroup]]) { + const uint projection_groups = (args.u[1] + 7u) / 8u; + const bool up = group.x >= projection_groups; + const uint2 projection_group = uint2(group.x - (up ? projection_groups : 0u), group.y); + const uint slot = group.y; + qwen_affine_qmv_fast_impl( + args, + up ? up_out : gate_out, + x, + up ? up_packed : gate_packed, + up ? up_scales : gate_scales, + up ? up_biases : gate_biases, + up ? args.u[9] : args.u[13], + up ? args.u[6] : args.u[14], + up ? args.u[7] : args.u[15], + projection_group, + simd_group, + lane, + gather ? selected[slot] : 0u, + batch + ? (gather && args.u[5] == 0u ? slot / args.u[8] : slot) + : (gather && args.u[5] != 0u ? slot : 0u), + slot); +} + +typedef decltype(kernel_qwen_affine_qmv_fast<4, 32>) qwen_affine_qmv_fast_b4g32; +typedef decltype(kernel_qwen_affine_qmv_fast<4, 64>) qwen_affine_qmv_fast_b4g64; +typedef decltype(kernel_qwen_affine_qmv_fast<8, 64>) qwen_affine_qmv_fast_b8g64; +typedef decltype(kernel_qwen_affine_qmv_batch_fast<4, 32>) qwen_affine_qmv_batch_fast_b4g32; +typedef decltype(kernel_qwen_affine_qmv_batch_fast<4, 64>) qwen_affine_qmv_batch_fast_b4g64; +typedef decltype(kernel_qwen_affine_qmv_batch_fast<8, 64>) qwen_affine_qmv_batch_fast_b8g64; +typedef decltype(kernel_qwen_affine_gather_qmv_fast<4, 32>) qwen_affine_gather_qmv_fast_b4g32; +typedef decltype(kernel_qwen_affine_gather_qmv_fast<4, 64>) qwen_affine_gather_qmv_fast_b4g64; +typedef decltype(kernel_qwen_affine_gather_qmv_fast<8, 64>) qwen_affine_gather_qmv_fast_b8g64; +typedef decltype(kernel_qwen_affine_gather_qmv_batch_fast<4, 32>) qwen_affine_gather_qmv_batch_fast_b4g32; +typedef decltype(kernel_qwen_affine_gather_qmv_batch_fast<4, 64>) qwen_affine_gather_qmv_batch_fast_b4g64; +typedef decltype(kernel_qwen_affine_gather_qmv_batch_fast<8, 64>) qwen_affine_gather_qmv_batch_fast_b8g64; +typedef decltype(kernel_qwen_affine_pair_qmv_fast<4, 32, false>) qwen_affine_pair_qmv_fast_b4g32; +typedef decltype(kernel_qwen_affine_pair_qmv_fast<4, 64, false>) qwen_affine_pair_qmv_fast_b4g64; +typedef decltype(kernel_qwen_affine_pair_qmv_fast<8, 64, false>) qwen_affine_pair_qmv_fast_b8g64; +typedef decltype(kernel_qwen_affine_pair_qmv_fast<4, 32, true>) qwen_affine_gather_pair_qmv_fast_b4g32; +typedef decltype(kernel_qwen_affine_pair_qmv_fast<4, 64, true>) qwen_affine_gather_pair_qmv_fast_b4g64; +typedef decltype(kernel_qwen_affine_pair_qmv_fast<8, 64, true>) qwen_affine_gather_pair_qmv_fast_b8g64; +typedef decltype(kernel_qwen_affine_pair_qmv_fast<4, 32, false, true>) qwen_affine_pair_qmv_batch_fast_b4g32; +typedef decltype(kernel_qwen_affine_pair_qmv_fast<4, 64, false, true>) qwen_affine_pair_qmv_batch_fast_b4g64; +typedef decltype(kernel_qwen_affine_pair_qmv_fast<8, 64, false, true>) qwen_affine_pair_qmv_batch_fast_b8g64; +typedef decltype(kernel_qwen_affine_pair_qmv_fast<4, 32, true, true>) qwen_affine_gather_pair_qmv_batch_fast_b4g32; +typedef decltype(kernel_qwen_affine_pair_qmv_fast<4, 64, true, true>) qwen_affine_gather_pair_qmv_batch_fast_b4g64; +typedef decltype(kernel_qwen_affine_pair_qmv_fast<8, 64, true, true>) qwen_affine_gather_pair_qmv_batch_fast_b8g64; + +template [[host_name("kernel_qwen_affine_qmv_fast_b4g32")]] +kernel qwen_affine_qmv_fast_b4g32 kernel_qwen_affine_qmv_fast<4, 32>; +template [[host_name("kernel_qwen_affine_qmv_fast_b4g64")]] +kernel qwen_affine_qmv_fast_b4g64 kernel_qwen_affine_qmv_fast<4, 64>; +template [[host_name("kernel_qwen_affine_qmv_fast_b8g64")]] +kernel qwen_affine_qmv_fast_b8g64 kernel_qwen_affine_qmv_fast<8, 64>; +template [[host_name("kernel_qwen_affine_qmv_batch_fast_b4g32")]] +kernel qwen_affine_qmv_batch_fast_b4g32 kernel_qwen_affine_qmv_batch_fast<4, 32>; +template [[host_name("kernel_qwen_affine_qmv_batch_fast_b4g64")]] +kernel qwen_affine_qmv_batch_fast_b4g64 kernel_qwen_affine_qmv_batch_fast<4, 64>; +template [[host_name("kernel_qwen_affine_qmv_batch_fast_b8g64")]] +kernel qwen_affine_qmv_batch_fast_b8g64 kernel_qwen_affine_qmv_batch_fast<8, 64>; +typedef decltype(kernel_qwen_affine_pair_qmv_wide<4, 32>) qwen_affine_pair_qmv_wide_b4g32; +typedef decltype(kernel_qwen_affine_pair_qmv_wide<4, 64>) qwen_affine_pair_qmv_wide_b4g64; +typedef decltype(kernel_qwen_affine_pair_qmv_wide<8, 64>) qwen_affine_pair_qmv_wide_b8g64; +template [[host_name("kernel_qwen_affine_pair_qmv_wide_b4g32")]] +kernel qwen_affine_pair_qmv_wide_b4g32 kernel_qwen_affine_pair_qmv_wide<4, 32>; +template [[host_name("kernel_qwen_affine_pair_qmv_wide_b4g64")]] +kernel qwen_affine_pair_qmv_wide_b4g64 kernel_qwen_affine_pair_qmv_wide<4, 64>; +template [[host_name("kernel_qwen_affine_pair_qmv_wide_b8g64")]] +kernel qwen_affine_pair_qmv_wide_b8g64 kernel_qwen_affine_pair_qmv_wide<8, 64>; + +#ifdef DS4_METAL_HAS_TENSOR +using namespace mpp::tensor_ops; + +template +kernel void kernel_qwen_affine_qmm_mpp( + constant qwen_kernel_args &args [[buffer(0)]], + device float *out [[buffer(1)]], + device const float *x [[buffer(2)]], + device const uchar *packed [[buffer(5)]], + device const uchar *scales [[buffer(6)]], + device const uchar *biases [[buffer(7)]], + uint2 group [[threadgroup_position_in_grid]], + uint tid [[thread_index_in_threadgroup]]) { + constexpr uint tile_n = 64u; + constexpr uint tile_k = 64u; + constexpr uint threads = 128u; + threadgroup bfloat xs[tile_m * tile_k]; + threadgroup bfloat ws[tile_n * tile_k]; + const uint first_m = group.y * tile_m; + const uint first_n = group.x * tile_n; + constexpr auto descriptor = matmul2d_descriptor( + tile_m, tile_n, tile_k, false, true, false, + matmul2d_descriptor::mode::multiply_accumulate); + matmul2d> multiply; + auto accum = multiply.template get_destination_cooperative_tensor< + tensor, tensor_inline>, + tensor, tensor_inline>, + float>(); +#pragma clang loop unroll(full) + for (uint i = 0u; i < accum.get_capacity(); i++) { + accum[i] = 0.0f; + } + + for (uint first_k = 0u; first_k < args.u[0]; first_k += tile_k) { + for (uint index = tid; index < tile_m * tile_k; index += threads) { + const uint row = index / tile_k; + const uint column = index % tile_k; + const uint input_m = first_m + row; + xs[index] = input_m < args.u[4] + ? (bfloat)x[(ulong)input_m * args.u[0] + first_k + column] + : bfloat(0.0f); + } + for (uint index = tid; index < tile_n * tile_k; index += threads) { + const uint row = index / tile_k; + const uint column = index % tile_k; + const uint weight_n = first_n + row; + ws[index] = weight_n < args.u[1] + ? (bfloat)qwen_quant_weight( + packed, scales, biases, + args.u[13], args.u[14], args.u[15], + weight_n, first_k + column, + args.u[0], bits, group_size) + : bfloat(0.0f); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + auto a = tensor, tensor_inline>( + xs, dextents(tile_k, tile_m)); + auto b = tensor, tensor_inline>( + ws, dextents(tile_k, tile_n)); + multiply.run(a, b, accum); + threadgroup_barrier(mem_flags::mem_threadgroup); + } + +#pragma clang loop unroll(full) + for (uint i = 0u; i < accum.get_capacity(); i++) { + const auto index = accum.get_multidimensional_index(i); + const uint output_n = first_n + index[0]; + const uint output_m = first_m + index[1]; + if (output_n < args.u[1] && output_m < args.u[4]) { + const float value = accum[i]; + out[(ulong)output_m * args.u[1] + output_n] = args.u[11] != 0u + ? qwen_round_bf16(value) + : value; + } + } +} + +typedef decltype(kernel_qwen_affine_qmm_mpp<4, 32, 32>) qwen_affine_qmm_mpp_b4g32_bm32; +typedef decltype(kernel_qwen_affine_qmm_mpp<4, 64, 32>) qwen_affine_qmm_mpp_b4g64_bm32; +typedef decltype(kernel_qwen_affine_qmm_mpp<8, 64, 32>) qwen_affine_qmm_mpp_b8g64_bm32; +typedef decltype(kernel_qwen_affine_qmm_mpp<4, 32, 64>) qwen_affine_qmm_mpp_b4g32_bm64; +typedef decltype(kernel_qwen_affine_qmm_mpp<4, 64, 64>) qwen_affine_qmm_mpp_b4g64_bm64; +typedef decltype(kernel_qwen_affine_qmm_mpp<8, 64, 64>) qwen_affine_qmm_mpp_b8g64_bm64; + +template [[host_name("kernel_qwen_affine_qmm_mpp_b4g32_bm32")]] +kernel qwen_affine_qmm_mpp_b4g32_bm32 kernel_qwen_affine_qmm_mpp<4, 32, 32>; +template [[host_name("kernel_qwen_affine_qmm_mpp_b4g64_bm32")]] +kernel qwen_affine_qmm_mpp_b4g64_bm32 kernel_qwen_affine_qmm_mpp<4, 64, 32>; +template [[host_name("kernel_qwen_affine_qmm_mpp_b8g64_bm32")]] +kernel qwen_affine_qmm_mpp_b8g64_bm32 kernel_qwen_affine_qmm_mpp<8, 64, 32>; +template [[host_name("kernel_qwen_affine_qmm_mpp_b4g32_bm64")]] +kernel qwen_affine_qmm_mpp_b4g32_bm64 kernel_qwen_affine_qmm_mpp<4, 32, 64>; +template [[host_name("kernel_qwen_affine_qmm_mpp_b4g64_bm64")]] +kernel qwen_affine_qmm_mpp_b4g64_bm64 kernel_qwen_affine_qmm_mpp<4, 64, 64>; +template [[host_name("kernel_qwen_affine_qmm_mpp_b8g64_bm64")]] +kernel qwen_affine_qmm_mpp_b8g64_bm64 kernel_qwen_affine_qmm_mpp<8, 64, 64>; + +template +kernel void kernel_qwen_bf16_qmm_mpp( + constant qwen_kernel_args &args [[buffer(0)]], + device float *out [[buffer(1)]], + device const float *x [[buffer(2)]], + device const uchar *weights [[buffer(5)]], + uint2 group [[threadgroup_position_in_grid]], + uint tid [[thread_index_in_threadgroup]]) { + constexpr uint tile_n = 64u; + constexpr uint tile_k = 64u; + constexpr uint threads = 128u; + threadgroup bfloat xs[tile_m * tile_k]; + threadgroup bfloat ws[tile_n * tile_k]; + const uint first_m = group.y * tile_m; + const uint first_n = group.x * tile_n; + constexpr auto descriptor = matmul2d_descriptor( + tile_m, tile_n, tile_k, false, true, false, + matmul2d_descriptor::mode::multiply_accumulate); + matmul2d> multiply; + auto accum = multiply.template get_destination_cooperative_tensor< + tensor, tensor_inline>, + tensor, tensor_inline>, + float>(); +#pragma clang loop unroll(full) + for (uint i = 0u; i < accum.get_capacity(); i++) accum[i] = 0.0f; + + for (uint first_k = 0u; first_k < args.u[0]; first_k += tile_k) { + for (uint index = tid; index < tile_m * tile_k; index += threads) { + const uint row = index / tile_k; + const uint column = index % tile_k; + const uint input_m = first_m + row; + xs[index] = input_m < args.u[4] + ? (bfloat)x[(ulong)input_m * args.u[0] + first_k + column] + : bfloat(0.0f); + } + for (uint index = tid; index < tile_n * tile_k; index += threads) { + const uint row = index / tile_k; + const uint column = index % tile_k; + const uint weight_n = first_n + row; + ws[index] = weight_n < args.u[1] + ? (bfloat)qwen_bf16(qwen_weight_u16( + weights, args.u[13], + (ulong)weight_n * args.u[0] + first_k + column)) + : bfloat(0.0f); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + auto a = tensor, tensor_inline>( + xs, dextents(tile_k, tile_m)); + auto b = tensor, tensor_inline>( + ws, dextents(tile_k, tile_n)); + multiply.run(a, b, accum); + threadgroup_barrier(mem_flags::mem_threadgroup); + } + +#pragma clang loop unroll(full) + for (uint i = 0u; i < accum.get_capacity(); i++) { + const auto index = accum.get_multidimensional_index(i); + const uint output_n = first_n + index[0]; + const uint output_m = first_m + index[1]; + if (output_n < args.u[1] && output_m < args.u[4]) { + out[(ulong)output_m * args.u[1] + output_n] = + qwen_round_bf16(accum[i]); + } + } +} + +typedef decltype(kernel_qwen_bf16_qmm_mpp<32>) qwen_bf16_qmm_mpp_bm32; +typedef decltype(kernel_qwen_bf16_qmm_mpp<64>) qwen_bf16_qmm_mpp_bm64; + +template [[host_name("kernel_qwen_bf16_qmm_mpp_bm32")]] +kernel qwen_bf16_qmm_mpp_bm32 kernel_qwen_bf16_qmm_mpp<32>; +template [[host_name("kernel_qwen_bf16_qmm_mpp_bm64")]] +kernel qwen_bf16_qmm_mpp_bm64 kernel_qwen_bf16_qmm_mpp<64>; +#endif +template [[host_name("kernel_qwen_affine_gather_qmv_fast_b4g32")]] +kernel qwen_affine_gather_qmv_fast_b4g32 kernel_qwen_affine_gather_qmv_fast<4, 32>; +template [[host_name("kernel_qwen_affine_gather_qmv_fast_b4g64")]] +kernel qwen_affine_gather_qmv_fast_b4g64 kernel_qwen_affine_gather_qmv_fast<4, 64>; +template [[host_name("kernel_qwen_affine_gather_qmv_fast_b8g64")]] +kernel qwen_affine_gather_qmv_fast_b8g64 kernel_qwen_affine_gather_qmv_fast<8, 64>; +template [[host_name("kernel_qwen_affine_gather_qmv_batch_fast_b4g32")]] +kernel qwen_affine_gather_qmv_batch_fast_b4g32 kernel_qwen_affine_gather_qmv_batch_fast<4, 32>; +template [[host_name("kernel_qwen_affine_gather_qmv_batch_fast_b4g64")]] +kernel qwen_affine_gather_qmv_batch_fast_b4g64 kernel_qwen_affine_gather_qmv_batch_fast<4, 64>; +template [[host_name("kernel_qwen_affine_gather_qmv_batch_fast_b8g64")]] +kernel qwen_affine_gather_qmv_batch_fast_b8g64 kernel_qwen_affine_gather_qmv_batch_fast<8, 64>; +template [[host_name("kernel_qwen_affine_pair_qmv_fast_b4g32")]] +kernel qwen_affine_pair_qmv_fast_b4g32 kernel_qwen_affine_pair_qmv_fast<4, 32, false>; +template [[host_name("kernel_qwen_affine_pair_qmv_fast_b4g64")]] +kernel qwen_affine_pair_qmv_fast_b4g64 kernel_qwen_affine_pair_qmv_fast<4, 64, false>; +template [[host_name("kernel_qwen_affine_pair_qmv_fast_b8g64")]] +kernel qwen_affine_pair_qmv_fast_b8g64 kernel_qwen_affine_pair_qmv_fast<8, 64, false>; +template [[host_name("kernel_qwen_affine_gather_pair_qmv_fast_b4g32")]] +kernel qwen_affine_gather_pair_qmv_fast_b4g32 kernel_qwen_affine_pair_qmv_fast<4, 32, true>; +template [[host_name("kernel_qwen_affine_gather_pair_qmv_fast_b4g64")]] +kernel qwen_affine_gather_pair_qmv_fast_b4g64 kernel_qwen_affine_pair_qmv_fast<4, 64, true>; +template [[host_name("kernel_qwen_affine_gather_pair_qmv_fast_b8g64")]] +kernel qwen_affine_gather_pair_qmv_fast_b8g64 kernel_qwen_affine_pair_qmv_fast<8, 64, true>; +template [[host_name("kernel_qwen_affine_pair_qmv_batch_fast_b4g32")]] +kernel qwen_affine_pair_qmv_batch_fast_b4g32 kernel_qwen_affine_pair_qmv_fast<4, 32, false, true>; +template [[host_name("kernel_qwen_affine_pair_qmv_batch_fast_b4g64")]] +kernel qwen_affine_pair_qmv_batch_fast_b4g64 kernel_qwen_affine_pair_qmv_fast<4, 64, false, true>; +template [[host_name("kernel_qwen_affine_pair_qmv_batch_fast_b8g64")]] +kernel qwen_affine_pair_qmv_batch_fast_b8g64 kernel_qwen_affine_pair_qmv_fast<8, 64, false, true>; +template [[host_name("kernel_qwen_affine_gather_pair_qmv_batch_fast_b4g32")]] +kernel qwen_affine_gather_pair_qmv_batch_fast_b4g32 kernel_qwen_affine_pair_qmv_fast<4, 32, true, true>; +template [[host_name("kernel_qwen_affine_gather_pair_qmv_batch_fast_b4g64")]] +kernel qwen_affine_gather_pair_qmv_batch_fast_b4g64 kernel_qwen_affine_pair_qmv_fast<4, 64, true, true>; +template [[host_name("kernel_qwen_affine_gather_pair_qmv_batch_fast_b8g64")]] +kernel qwen_affine_gather_pair_qmv_batch_fast_b8g64 kernel_qwen_affine_pair_qmv_fast<8, 64, true, true>; + +template +static inline void qwen_affine_qmv_safe_impl( + constant qwen_kernel_args &args, + device float *out, + device const float *x, + device const uchar *packed, + device const uchar *scale_bytes, + device const uchar *bias_bytes, + uint packed_offset, + uint scale_offset, + uint bias_offset, + uint2 group, + uint simd_group, + uint lane, + uint expert, + uint input_slot, + uint output_slot) { + constexpr ushort pack_factor = 32 / bits; + constexpr ushort values_per_thread = pack_factor; + constexpr ushort block_size = values_per_thread * 32; + constexpr ushort scale_step = group_size / values_per_thread; + const uint in_dim = args.u[0]; + const uint out_dim = args.u[1]; + const uint input_bytes_per_row = in_dim * bits / 8; + const uint groups_per_row = in_dim / group_size; + const uint first_row = group.x * 8u + simd_group * 4u; + if (first_row >= out_dim) return; + const uint used_row = out_dim < 8u ? first_row : min(out_dim - 4u, first_row); + const ulong table_row = (ulong)expert * out_dim + used_row; + device const uchar *weights = packed + + table_row * input_bytes_per_row + lane * 4u; + ulong parameter = table_row * groups_per_row + lane / scale_step; + device const float *input = x + + (ulong)input_slot * in_dim + lane * values_per_thread; + float result[4] = {0.0f, 0.0f, 0.0f, 0.0f}; + + uint column = 0; + for (; column + block_size < in_dim; column += block_size) { + float values[values_per_thread]; + float input_sum = 0.0f; + for (ushort i = 0; i < values_per_thread; i++) { + if constexpr (bits == 4) { + constexpr float divisors[4] = {1.0f, 16.0f, 256.0f, 4096.0f}; + values[i] = input[i] / divisors[i & 3u]; + } else { + values[i] = input[i]; + } + } + if constexpr (bits == 4) { + for (ushort i = 0; i < values_per_thread; i += 4u) { + const bfloat input_group = + (bfloat)input[i] + (bfloat)input[i + 1u] + + (bfloat)input[i + 2u] + (bfloat)input[i + 3u]; + input_sum += (float)input_group; + } + } else { + for (ushort i = 0; i < values_per_thread; i++) input_sum += input[i]; + } + for (ushort row = 0; row < 4u && used_row + row < out_dim; row++) { + device const uchar *row_weights = weights + row * input_bytes_per_row; + float dot = 0.0f; + if constexpr (bits == 4) { + for (ushort i = 0; i < values_per_thread / 4u; i++) { + const ushort word = qwen_weight_u16(row_weights, packed_offset, i); + dot += values[4u * i] * (word & 0x000fu) + + values[4u * i + 1u] * (word & 0x00f0u) + + values[4u * i + 2u] * (word & 0x0f00u) + + values[4u * i + 3u] * (word & 0xf000u); + } + } else { + for (ushort i = 0; i < values_per_thread; i++) { + dot += values[i] * row_weights[packed_offset + i]; + } + } + const ulong row_parameter = parameter + (ulong)row * groups_per_row; + result[row] += + qwen_bf16(qwen_weight_u16(scale_bytes, scale_offset, row_parameter)) * dot + + input_sum * + qwen_bf16(qwen_weight_u16(bias_bytes, bias_offset, row_parameter)); + } + weights += block_size * bits / 8; + parameter += block_size / group_size; + input += block_size; + } + + const int remaining = clamp( + (int)in_dim - (int)column - (int)lane * (int)values_per_thread, + 0, + (int)values_per_thread); + if (remaining > 0) { + float values[values_per_thread]; + float input_sum = 0.0f; + for (ushort i = 0; i < values_per_thread; i++) { + if ((int)i < remaining) { + if constexpr (bits == 4) { + constexpr float divisors[4] = {1.0f, 16.0f, 256.0f, 4096.0f}; + values[i] = input[i] / divisors[i & 3u]; + } else { + values[i] = input[i]; + } + } else { + values[i] = 0.0f; + } + } + if constexpr (bits == 4) { + for (ushort i = 0; i < (ushort)remaining; i += 4u) { + const bfloat input_group = + (bfloat)input[i] + (bfloat)input[i + 1u] + + (bfloat)input[i + 2u] + (bfloat)input[i + 3u]; + input_sum += (float)input_group; + } + } else { + for (ushort i = 0; i < (ushort)remaining; i++) input_sum += input[i]; + } + for (ushort row = 0; row < 4u && used_row + row < out_dim; row++) { + device const uchar *row_weights = weights + row * input_bytes_per_row; + float dot = 0.0f; + if constexpr (bits == 4) { + for (ushort i = 0; i < (ushort)remaining / 4u; i++) { + const ushort word = qwen_weight_u16(row_weights, packed_offset, i); + dot += values[4u * i] * (word & 0x000fu) + + values[4u * i + 1u] * (word & 0x00f0u) + + values[4u * i + 2u] * (word & 0x0f00u) + + values[4u * i + 3u] * (word & 0xf000u); + } + } else { + for (ushort i = 0; i < (ushort)remaining; i++) { + dot += values[i] * row_weights[packed_offset + i]; + } + } + const ulong row_parameter = parameter + (ulong)row * groups_per_row; + result[row] += + qwen_bf16(qwen_weight_u16(scale_bytes, scale_offset, row_parameter)) * dot + + input_sum * + qwen_bf16(qwen_weight_u16(bias_bytes, bias_offset, row_parameter)); + } + } + + for (ushort row = 0; row < 4u && used_row + row < out_dim; row++) { + const float sum = simd_sum(result[row]); + if (lane == 0u) { + out[(ulong)output_slot * out_dim + used_row + row] = + args.u[11] != 0u ? (float)(bfloat)sum : sum; + } + } +} + +template +kernel void kernel_qwen_affine_qmv_safe( + constant qwen_kernel_args &args [[buffer(0)]], + device float *out [[buffer(1)]], + device const float *x [[buffer(2)]], + device const uchar *packed [[buffer(5)]], + device const uchar *scales [[buffer(6)]], + device const uchar *biases [[buffer(7)]], + uint2 group [[threadgroup_position_in_grid]], + uint simd_group [[simdgroup_index_in_threadgroup]], + uint lane [[thread_index_in_simdgroup]]) { + qwen_affine_qmv_safe_impl( + args, out, x, packed, scales, biases, args.u[13], args.u[14], args.u[15], + group, simd_group, lane, + 0u, 0u, 0u); +} + +template +kernel void kernel_qwen_affine_qmv_batch_safe( + constant qwen_kernel_args &args [[buffer(0)]], + device float *out [[buffer(1)]], + device const float *x [[buffer(2)]], + device const uchar *packed [[buffer(5)]], + device const uchar *scales [[buffer(6)]], + device const uchar *biases [[buffer(7)]], + uint2 group [[threadgroup_position_in_grid]], + uint simd_group [[simdgroup_index_in_threadgroup]], + uint lane [[thread_index_in_simdgroup]]) { + qwen_affine_qmv_wide_impl( + args, out, x, packed, scales, biases, group.x, simd_group, lane); +} + +template +kernel void kernel_qwen_affine_gather_qmv_safe( + constant qwen_kernel_args &args [[buffer(0)]], + device float *out [[buffer(1)]], + device const float *x [[buffer(2)]], + device const uint *selected [[buffer(3)]], + device const uchar *packed [[buffer(5)]], + device const uchar *scales [[buffer(6)]], + device const uchar *biases [[buffer(7)]], + uint2 group [[threadgroup_position_in_grid]], + uint simd_group [[simdgroup_index_in_threadgroup]], + uint lane [[thread_index_in_simdgroup]]) { + const uint slot = group.y; + qwen_affine_qmv_safe_impl( + args, out, x, packed, scales, biases, args.u[13], args.u[14], args.u[15], + group, simd_group, lane, + selected[slot], args.u[5] != 0u ? slot : 0u, slot); +} + +template +kernel void kernel_qwen_affine_gather_qmv_batch_safe( + constant qwen_kernel_args &args [[buffer(0)]], + device float *out [[buffer(1)]], + device const float *x [[buffer(2)]], + device const uint *selected [[buffer(3)]], + device const uchar *packed [[buffer(5)]], + device const uchar *scales [[buffer(6)]], + device const uchar *biases [[buffer(7)]], + uint2 group [[threadgroup_position_in_grid]], + uint simd_group [[simdgroup_index_in_threadgroup]], + uint lane [[thread_index_in_simdgroup]]) { + const uint slot = group.y; + qwen_affine_qmv_safe_impl( + args, out, x, packed, scales, biases, args.u[13], args.u[14], args.u[15], + group, simd_group, lane, selected[slot], + args.u[5] != 0u ? slot : slot / args.u[8], slot); +} + +template +kernel void kernel_qwen_affine_pair_qmv_safe( + constant qwen_kernel_args &args [[buffer(0)]], + device float *gate_out [[buffer(1)]], + device const float *x [[buffer(2)]], + device const uint *selected [[buffer(3)]], + device float *up_out [[buffer(4)]], + device const uchar *gate_packed [[buffer(5)]], + device const uchar *gate_scales [[buffer(6)]], + device const uchar *gate_biases [[buffer(7)]], + device const uchar *up_packed [[buffer(9)]], + device const uchar *up_scales [[buffer(10)]], + device const uchar *up_biases [[buffer(11)]], + uint2 group [[threadgroup_position_in_grid]], + uint simd_group [[simdgroup_index_in_threadgroup]], + uint lane [[thread_index_in_simdgroup]]) { + const uint projection_groups = (args.u[1] + 7u) / 8u; + const bool up = group.x >= projection_groups; + const uint2 projection_group = uint2(group.x - (up ? projection_groups : 0u), group.y); + const uint slot = group.y; + qwen_affine_qmv_safe_impl( + args, + up ? up_out : gate_out, + x, + up ? up_packed : gate_packed, + up ? up_scales : gate_scales, + up ? up_biases : gate_biases, + up ? args.u[9] : args.u[13], + up ? args.u[6] : args.u[14], + up ? args.u[7] : args.u[15], + projection_group, + simd_group, + lane, + gather ? selected[slot] : 0u, + batch + ? (gather && args.u[5] == 0u ? slot / args.u[8] : slot) + : (gather && args.u[5] != 0u ? slot : 0u), + slot); +} + +typedef decltype(kernel_qwen_affine_qmv_safe<4, 32>) qwen_affine_qmv_safe_b4g32; +typedef decltype(kernel_qwen_affine_qmv_safe<4, 64>) qwen_affine_qmv_safe_b4g64; +typedef decltype(kernel_qwen_affine_qmv_safe<8, 64>) qwen_affine_qmv_safe_b8g64; +typedef decltype(kernel_qwen_affine_qmv_batch_safe<4, 32>) qwen_affine_qmv_batch_safe_b4g32; +typedef decltype(kernel_qwen_affine_qmv_batch_safe<4, 64>) qwen_affine_qmv_batch_safe_b4g64; +typedef decltype(kernel_qwen_affine_qmv_batch_safe<8, 64>) qwen_affine_qmv_batch_safe_b8g64; +typedef decltype(kernel_qwen_affine_gather_qmv_safe<4, 32>) qwen_affine_gather_qmv_safe_b4g32; +typedef decltype(kernel_qwen_affine_gather_qmv_safe<4, 64>) qwen_affine_gather_qmv_safe_b4g64; +typedef decltype(kernel_qwen_affine_gather_qmv_safe<8, 64>) qwen_affine_gather_qmv_safe_b8g64; +typedef decltype(kernel_qwen_affine_gather_qmv_batch_safe<4, 32>) qwen_affine_gather_qmv_batch_safe_b4g32; +typedef decltype(kernel_qwen_affine_gather_qmv_batch_safe<4, 64>) qwen_affine_gather_qmv_batch_safe_b4g64; +typedef decltype(kernel_qwen_affine_gather_qmv_batch_safe<8, 64>) qwen_affine_gather_qmv_batch_safe_b8g64; +typedef decltype(kernel_qwen_affine_pair_qmv_safe<4, 32, false>) qwen_affine_pair_qmv_safe_b4g32; +typedef decltype(kernel_qwen_affine_pair_qmv_safe<4, 64, false>) qwen_affine_pair_qmv_safe_b4g64; +typedef decltype(kernel_qwen_affine_pair_qmv_safe<8, 64, false>) qwen_affine_pair_qmv_safe_b8g64; +typedef decltype(kernel_qwen_affine_pair_qmv_safe<4, 32, true>) qwen_affine_gather_pair_qmv_safe_b4g32; +typedef decltype(kernel_qwen_affine_pair_qmv_safe<4, 64, true>) qwen_affine_gather_pair_qmv_safe_b4g64; +typedef decltype(kernel_qwen_affine_pair_qmv_safe<8, 64, true>) qwen_affine_gather_pair_qmv_safe_b8g64; +typedef decltype(kernel_qwen_affine_pair_qmv_safe<4, 32, false, true>) qwen_affine_pair_qmv_batch_safe_b4g32; +typedef decltype(kernel_qwen_affine_pair_qmv_safe<4, 64, false, true>) qwen_affine_pair_qmv_batch_safe_b4g64; +typedef decltype(kernel_qwen_affine_pair_qmv_safe<8, 64, false, true>) qwen_affine_pair_qmv_batch_safe_b8g64; +typedef decltype(kernel_qwen_affine_pair_qmv_safe<4, 32, true, true>) qwen_affine_gather_pair_qmv_batch_safe_b4g32; +typedef decltype(kernel_qwen_affine_pair_qmv_safe<4, 64, true, true>) qwen_affine_gather_pair_qmv_batch_safe_b4g64; +typedef decltype(kernel_qwen_affine_pair_qmv_safe<8, 64, true, true>) qwen_affine_gather_pair_qmv_batch_safe_b8g64; + +template [[host_name("kernel_qwen_affine_qmv_safe_b4g32")]] +kernel qwen_affine_qmv_safe_b4g32 kernel_qwen_affine_qmv_safe<4, 32>; +template [[host_name("kernel_qwen_affine_qmv_safe_b4g64")]] +kernel qwen_affine_qmv_safe_b4g64 kernel_qwen_affine_qmv_safe<4, 64>; +template [[host_name("kernel_qwen_affine_qmv_safe_b8g64")]] +kernel qwen_affine_qmv_safe_b8g64 kernel_qwen_affine_qmv_safe<8, 64>; +template [[host_name("kernel_qwen_affine_qmv_batch_safe_b4g32")]] +kernel qwen_affine_qmv_batch_safe_b4g32 kernel_qwen_affine_qmv_batch_safe<4, 32>; +template [[host_name("kernel_qwen_affine_qmv_batch_safe_b4g64")]] +kernel qwen_affine_qmv_batch_safe_b4g64 kernel_qwen_affine_qmv_batch_safe<4, 64>; +template [[host_name("kernel_qwen_affine_qmv_batch_safe_b8g64")]] +kernel qwen_affine_qmv_batch_safe_b8g64 kernel_qwen_affine_qmv_batch_safe<8, 64>; +template [[host_name("kernel_qwen_affine_gather_qmv_safe_b4g32")]] +kernel qwen_affine_gather_qmv_safe_b4g32 kernel_qwen_affine_gather_qmv_safe<4, 32>; +template [[host_name("kernel_qwen_affine_gather_qmv_safe_b4g64")]] +kernel qwen_affine_gather_qmv_safe_b4g64 kernel_qwen_affine_gather_qmv_safe<4, 64>; +template [[host_name("kernel_qwen_affine_gather_qmv_safe_b8g64")]] +kernel qwen_affine_gather_qmv_safe_b8g64 kernel_qwen_affine_gather_qmv_safe<8, 64>; +template [[host_name("kernel_qwen_affine_gather_qmv_batch_safe_b4g32")]] +kernel qwen_affine_gather_qmv_batch_safe_b4g32 kernel_qwen_affine_gather_qmv_batch_safe<4, 32>; +template [[host_name("kernel_qwen_affine_gather_qmv_batch_safe_b4g64")]] +kernel qwen_affine_gather_qmv_batch_safe_b4g64 kernel_qwen_affine_gather_qmv_batch_safe<4, 64>; +template [[host_name("kernel_qwen_affine_gather_qmv_batch_safe_b8g64")]] +kernel qwen_affine_gather_qmv_batch_safe_b8g64 kernel_qwen_affine_gather_qmv_batch_safe<8, 64>; +template [[host_name("kernel_qwen_affine_pair_qmv_safe_b4g32")]] +kernel qwen_affine_pair_qmv_safe_b4g32 kernel_qwen_affine_pair_qmv_safe<4, 32, false>; +template [[host_name("kernel_qwen_affine_pair_qmv_safe_b4g64")]] +kernel qwen_affine_pair_qmv_safe_b4g64 kernel_qwen_affine_pair_qmv_safe<4, 64, false>; +template [[host_name("kernel_qwen_affine_pair_qmv_safe_b8g64")]] +kernel qwen_affine_pair_qmv_safe_b8g64 kernel_qwen_affine_pair_qmv_safe<8, 64, false>; +template [[host_name("kernel_qwen_affine_gather_pair_qmv_safe_b4g32")]] +kernel qwen_affine_gather_pair_qmv_safe_b4g32 kernel_qwen_affine_pair_qmv_safe<4, 32, true>; +template [[host_name("kernel_qwen_affine_gather_pair_qmv_safe_b4g64")]] +kernel qwen_affine_gather_pair_qmv_safe_b4g64 kernel_qwen_affine_pair_qmv_safe<4, 64, true>; +template [[host_name("kernel_qwen_affine_gather_pair_qmv_safe_b8g64")]] +kernel qwen_affine_gather_pair_qmv_safe_b8g64 kernel_qwen_affine_pair_qmv_safe<8, 64, true>; +template [[host_name("kernel_qwen_affine_pair_qmv_batch_safe_b4g32")]] +kernel qwen_affine_pair_qmv_batch_safe_b4g32 kernel_qwen_affine_pair_qmv_safe<4, 32, false, true>; +template [[host_name("kernel_qwen_affine_pair_qmv_batch_safe_b4g64")]] +kernel qwen_affine_pair_qmv_batch_safe_b4g64 kernel_qwen_affine_pair_qmv_safe<4, 64, false, true>; +template [[host_name("kernel_qwen_affine_pair_qmv_batch_safe_b8g64")]] +kernel qwen_affine_pair_qmv_batch_safe_b8g64 kernel_qwen_affine_pair_qmv_safe<8, 64, false, true>; +template [[host_name("kernel_qwen_affine_gather_pair_qmv_batch_safe_b4g32")]] +kernel qwen_affine_gather_pair_qmv_batch_safe_b4g32 kernel_qwen_affine_pair_qmv_safe<4, 32, true, true>; +template [[host_name("kernel_qwen_affine_gather_pair_qmv_batch_safe_b4g64")]] +kernel qwen_affine_gather_pair_qmv_batch_safe_b4g64 kernel_qwen_affine_pair_qmv_safe<4, 64, true, true>; +template [[host_name("kernel_qwen_affine_gather_pair_qmv_batch_safe_b8g64")]] +kernel qwen_affine_gather_pair_qmv_batch_safe_b8g64 kernel_qwen_affine_pair_qmv_safe<8, 64, true, true>; + +kernel void kernel_qwen_affine_gather_qmv( + constant qwen_kernel_args &args [[buffer(0)]], + device float *out [[buffer(1)]], + device const float *x [[buffer(2)]], + device const int *selected [[buffer(3)]], + device const uchar *packed [[buffer(5)]], + device const uchar *scales [[buffer(6)]], + device const uchar *biases [[buffer(7)]], + uint2 group [[threadgroup_position_in_grid]], + uint simd_group [[simdgroup_index_in_threadgroup]], + uint lane [[thread_index_in_simdgroup]]) { + const uint in_dim = args.u[0]; + const uint out_dim = args.u[1]; + const uint bits = args.u[2]; + const uint group_size = args.u[3]; + const uint slot = group.y; + const uint expert = (uint)selected[slot]; + const uint values_per_word = 32u / bits; + const uint packed_columns = in_dim / values_per_word; + const uint words_per_group = group_size / values_per_word; + const uint groups_per_row = in_dim / group_size; + const uint first_row = group.x * 8u + simd_group * 4u; + const uint mask = (1u << bits) - 1u; + float sums[4] = {0.0f, 0.0f, 0.0f, 0.0f}; + + for (uint word_index = lane; word_index < packed_columns; word_index += 32u) { + const uint first_column = word_index * values_per_word; + const uint quant_group = word_index / words_per_group; + uint words[4] = {0u, 0u, 0u, 0u}; + float scale[4] = {0.0f, 0.0f, 0.0f, 0.0f}; + float bias[4] = {0.0f, 0.0f, 0.0f, 0.0f}; + for (uint output = 0; output < 4u; output++) { + const uint row = first_row + output; + if (row < out_dim) { + const ulong table_row = (ulong)expert * out_dim + row; + words[output] = qwen_weight_u32( + packed, args.u[13], table_row * packed_columns + word_index); + const ulong parameter = table_row * groups_per_row + quant_group; + scale[output] = qwen_bf16(qwen_weight_u16(scales, args.u[14], parameter)); + bias[output] = qwen_bf16(qwen_weight_u16(biases, args.u[15], parameter)); + } + } + for (uint packed_value = 0; packed_value < values_per_word; packed_value++) { + const float activation = + x[(args.u[5] != 0u ? (ulong)slot * in_dim : 0u) + + first_column + packed_value]; + const uint shift = packed_value * bits; + for (uint output = 0; output < 4u; output++) { + const float weight = fma((float)((words[output] >> shift) & mask), + scale[output], bias[output]); + sums[output] = fma(weight, activation, sums[output]); + } + } + } + for (uint output = 0; output < 4u; output++) { + const float sum = simd_sum(sums[output]); + if (lane == 0u && first_row + output < out_dim) { + out[(ulong)slot * out_dim + first_row + output] = qwen_round_bf16(sum); + } + } +} + +kernel void kernel_qwen_weighted_sum10( + constant qwen_kernel_args &args [[buffer(0)]], + device float *out [[buffer(1)]], + device const float *experts [[buffer(2)]], + device const float *weights [[buffer(3)]], + uint index [[thread_position_in_grid]]) { + if (index >= args.u[0]) return; + const uint hidden = args.u[1] == 0u ? args.u[0] : args.u[1]; + const uint row = index / hidden; + const uint column = index % hidden; + const ulong experts_base = (ulong)row * 10u * hidden; + const ulong weights_base = (ulong)row * 10u; + float weighted[10]; + for (uint slot = 0; slot < 10u; slot++) { + weighted[slot] = qwen_round_bf16( + experts[experts_base + (ulong)slot * hidden + column] * + weights[weights_base + slot]); + } + // MLX's 10xhidden strided reduction uses eight rows per threadgroup: + // lanes 0 and 1 first fold rows 8 and 9, then lane 0 folds lanes 1..7. + float value = qwen_round_bf16(weighted[0] + weighted[8]); + value = qwen_round_bf16(value + qwen_round_bf16(weighted[1] + weighted[9])); + for (uint slot = 2; slot < 8u; slot++) { + value = qwen_round_bf16(value + weighted[slot]); + } + out[index] = value; } kernel void kernel_qwen_affine_embedding( @@ -109,9 +1217,25 @@ kernel void kernel_qwen_affine_embedding( device const uchar *biases [[buffer(7)]], uint column [[thread_position_in_grid]]) { if (column >= args.u[0]) return; - out[column] = qwen_quant_weight(packed, scales, biases, - args.u[13], args.u[14], args.u[15], - args.u[4], column, args.u[0], args.u[2], args.u[3]); + out[column] = qwen_round_bf16(qwen_quant_weight( + packed, scales, biases, args.u[13], args.u[14], args.u[15], + args.u[4], column, args.u[0], args.u[2], args.u[3])); +} + +kernel void kernel_qwen_affine_embedding_batch( + constant qwen_kernel_args &args [[buffer(0)]], + device float *out [[buffer(1)]], + device const int *tokens [[buffer(2)]], + device const uchar *packed [[buffer(5)]], + device const uchar *scales [[buffer(6)]], + device const uchar *biases [[buffer(7)]], + uint2 gid [[thread_position_in_grid]]) { + const uint column = gid.x; + const uint row = gid.y; + if (column >= args.u[0] || row >= args.u[4]) return; + out[(ulong)row * args.u[0] + column] = qwen_round_bf16(qwen_quant_weight( + packed, scales, biases, args.u[13], args.u[14], args.u[15], + (uint)tokens[row], column, args.u[0], args.u[2], args.u[3])); } kernel void kernel_qwen_ple_dequant( @@ -125,8 +1249,8 @@ kernel void kernel_qwen_ple_dequant( if (index >= dim * args.u[1]) return; const uint row = index / dim; const uint column = index % dim; - out[index] = qwen_quant_value(packed, scales, biases, row, column, - dim, args.u[2], args.u[3]); + out[index] = qwen_round_bf16(qwen_quant_value( + packed, scales, biases, row, column, dim, args.u[2], args.u[3])); } kernel void kernel_qwen_bf16_mv( @@ -142,7 +1266,389 @@ kernel void kernel_qwen_bf16_mv( (ulong)row * args.u[0] + column)), x[column], sum); } - out[row] = sum; + out[row] = qwen_round_bf16(sum); +} + +kernel void kernel_qwen_bf16_qmv( + constant qwen_kernel_args &args [[buffer(0)]], + device float *out [[buffer(1)]], + device const float *x [[buffer(2)]], + device const uchar *weights [[buffer(5)]], + uint group [[threadgroup_position_in_grid]], + uint simd_group [[simdgroup_index_in_threadgroup]], + uint lane [[thread_index_in_simdgroup]]) { + const uint in_dim = args.u[0]; + const uint out_dim = args.u[1]; + const uint first_row = group * 8u + simd_group * 4u; + float sums[4] = {0.0f, 0.0f, 0.0f, 0.0f}; + for (uint column = lane; column < in_dim; column += 32u) { + const float activation = x[column]; + for (uint output = 0; output < 4u; output++) { + const uint row = first_row + output; + if (row < out_dim) { + const float weight = qwen_bf16(qwen_weight_u16( + weights, args.u[13], (ulong)row * in_dim + column)); + sums[output] = fma(weight, activation, sums[output]); + } + } + } + for (uint output = 0; output < 4u; output++) { + const float sum = simd_sum(sums[output]); + if (lane == 0u && first_row + output < out_dim) { + out[first_row + output] = qwen_round_bf16(sum); + } + } +} + +kernel void kernel_qwen_bf16_gemv( + constant qwen_kernel_args &args [[buffer(0)]], + device float *out [[buffer(1)]], + device const float *x [[buffer(2)]], + device const uchar *weights [[buffer(5)]], + threadgroup float *partials [[threadgroup(0)]], + uint2 group [[threadgroup_position_in_grid]], + ushort lane [[thread_index_in_simdgroup]], + ushort simd_group [[simdgroup_index_in_threadgroup]]) { + const uint in_dim = args.u[0]; + const uint out_dim = args.u[1]; + const uint mode = args.u[4]; + device const float *input = x + (ulong)group.y * in_dim; + device float *output = out + (ulong)group.y * out_dim; + float sums[4] = {0.0f, 0.0f, 0.0f, 0.0f}; + + if (mode == 2u) { + const uint first_row = group.x * 4u; + uint column = (uint)simd_group * 128u + (uint)lane * 4u; + for (; column < in_dim; column += 1024u) { + for (uint row = 0; row < 4u; row++) { + if (first_row + row < out_dim) { + const ulong base = (ulong)(first_row + row) * in_dim + column; + sums[row] += qwen_bf16(qwen_weight_u16(weights, args.u[13], base)) * input[column]; + sums[row] += qwen_bf16(qwen_weight_u16(weights, args.u[13], base + 1u)) * input[column + 1u]; + sums[row] += qwen_bf16(qwen_weight_u16(weights, args.u[13], base + 2u)) * input[column + 2u]; + sums[row] += qwen_bf16(qwen_weight_u16(weights, args.u[13], base + 3u)) * input[column + 3u]; + } + } + } + for (uint row = 0; row < 4u; row++) { + for (ushort offset = 16u; offset >= 1u; offset >>= 1u) { + sums[row] += simd_shuffle_down(sums[row], offset); + } + if (lane == 0u) partials[(uint)simd_group * 8u + row] = sums[row]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (simd_group == 0u && lane == 0u) { + for (uint row = 0; row < 4u && first_row + row < out_dim; row++) { + float value = partials[row]; + for (uint part = 1u; part < 8u; part++) value += partials[part * 8u + row]; + output[first_row + row] = qwen_round_bf16(value); + } + } + return; + } + + if (mode == 3u) { + const uint row_group = lane / 4u; + const uint column_lane = lane % 4u; + const uint first_row = group.x * 32u + row_group * 4u; + for (uint column = column_lane * 4u; column < in_dim; column += 16u) { + const uint valid = min(4u, in_dim - column); + for (uint row = 0; row < 4u && first_row + row < out_dim; row++) { + const ulong base = (ulong)(first_row + row) * in_dim + column; + for (uint i = 0; i < valid; i++) { + sums[row] += qwen_bf16(qwen_weight_u16(weights, args.u[13], base + i)) * input[column + i]; + } + } + } + for (uint row = 0; row < 4u; row++) { + sums[row] += simd_shuffle_down(sums[row], 2u); + sums[row] += simd_shuffle_down(sums[row], 1u); + if (column_lane == 0u && first_row + row < out_dim) { + output[first_row + row] = qwen_round_bf16(sums[row]); + } + } + return; + } + + const uint rows_per_group = mode == 1u ? 32u : 16u; + const uint first_row = group.x * rows_per_group + (uint)simd_group * 4u; + for (uint column = (uint)lane * 4u; column < in_dim; column += 128u) { + const uint valid = min(4u, in_dim - column); + for (uint row = 0; row < 4u && first_row + row < out_dim; row++) { + const ulong base = (ulong)(first_row + row) * in_dim + column; + for (uint i = 0; i < valid; i++) { + sums[row] += qwen_bf16(qwen_weight_u16(weights, args.u[13], base + i)) * input[column + i]; + } + } + } + for (uint row = 0; row < 4u; row++) { + for (ushort offset = 16u; offset >= 1u; offset >>= 1u) { + sums[row] += simd_shuffle_down(sums[row], offset); + } + if (lane == 0u && first_row + row < out_dim) { + output[first_row + row] = qwen_round_bf16(sums[row]); + } + } +} + +// out[4, N] = x[4, K] @ weight[N, K]^T. This is the decode shape used by +// the MTP hidden projection: stream each BF16 weight row once for all four +// hyper streams and preserve MLX's vec4 accumulation and reduction order. +kernel void kernel_qwen_bf16_gemv4( + constant qwen_kernel_args &args [[buffer(0)]], + device float *out [[buffer(1)]], + device const float *x [[buffer(2)]], + device const uchar *weights [[buffer(5)]], + uint group [[threadgroup_position_in_grid]], + uint simd_group [[simdgroup_index_in_threadgroup]], + uint lane [[thread_index_in_simdgroup]]) { + const uint in_dim = args.u[0]; + const uint out_dim = args.u[1]; + const uint row = group * 4u + simd_group; + if (row >= out_dim) return; + + constexpr uint k_lanes = 32u; + constexpr uint unroll = 8u; + const uint vectors = in_dim / 4u; + const uint main = vectors - vectors % (k_lanes * unroll); + const ulong weight_row = (ulong)row * in_dim; + float result[4] = {0.0f, 0.0f, 0.0f, 0.0f}; + + for (uint base = 0u; base < main; base += k_lanes * unroll) { + float accumulated[4] = {0.0f, 0.0f, 0.0f, 0.0f}; + for (uint i = 0u; i < unroll; i++) { + const uint column = (base + i * k_lanes + lane) * 4u; + const float4 weight = float4( + qwen_bf16(qwen_weight_u16(weights, args.u[13], weight_row + column)), + qwen_bf16(qwen_weight_u16(weights, args.u[13], weight_row + column + 1u)), + qwen_bf16(qwen_weight_u16(weights, args.u[13], weight_row + column + 2u)), + qwen_bf16(qwen_weight_u16(weights, args.u[13], weight_row + column + 3u))); + for (uint stream = 0u; stream < 4u; stream++) { + accumulated[stream] += dot(weight, *((device const float4 *)( + x + (ulong)stream * in_dim + column))); + } + } + for (uint stream = 0u; stream < 4u; stream++) { + result[stream] += accumulated[stream]; + } + } + for (uint vector = main + lane; vector < vectors; vector += k_lanes) { + const uint column = vector * 4u; + const float4 weight = float4( + qwen_bf16(qwen_weight_u16(weights, args.u[13], weight_row + column)), + qwen_bf16(qwen_weight_u16(weights, args.u[13], weight_row + column + 1u)), + qwen_bf16(qwen_weight_u16(weights, args.u[13], weight_row + column + 2u)), + qwen_bf16(qwen_weight_u16(weights, args.u[13], weight_row + column + 3u))); + for (uint stream = 0u; stream < 4u; stream++) { + result[stream] += dot(weight, *((device const float4 *)( + x + (ulong)stream * in_dim + column))); + } + } + for (uint stream = 0u; stream < 4u; stream++) { + for (ushort offset = 16u; offset >= 1u; offset >>= 1u) { + result[stream] += simd_shuffle_down(result[stream], offset); + } + if (lane == 0u) { + out[(ulong)stream * out_dim + row] = qwen_round_bf16(result[stream]); + } + } +} + +// MLX affine quantization, specialized for the private 8-bit/group-64 hyper +// packs. Rust owns pack creation and lifetime; this kernel only transforms the +// source BF16 matrices already admitted by the model loader. +kernel void kernel_qwen_quantize_hyper_q8( + constant qwen_kernel_args &args [[buffer(0)]], + device uchar *pack [[buffer(1)]], + device const uchar *weight [[buffer(5)]], + uint2 group [[threadgroup_position_in_grid]], + uint lane [[thread_index_in_simdgroup]]) { + constexpr uint quant_group = 64u; + constexpr float bins = 255.0f; + const uint columns = args.u[0]; + const uint source_rows = args.u[1]; + const uint destination_row = args.u[2] + group.y; + const uint destination_rows = args.u[3]; + if (group.y >= source_rows || group.x >= columns / quant_group) return; + + float first = qwen_bf16(qwen_weight_u16( + weight, args.u[13], (ulong)group.y * columns + group.x * quant_group + lane)); + float second = qwen_bf16(qwen_weight_u16( + weight, args.u[13], (ulong)group.y * columns + group.x * quant_group + lane + 32u)); + const float minimum = simd_min(min(first, second)); + const float maximum = simd_max(max(first, second)); + float scale = max((maximum - minimum) / bins, 1e-7f); + const bool negative_side = abs(minimum) > abs(maximum); + scale = negative_side ? scale : -scale; + const float edge = negative_side ? minimum : maximum; + const float zero = round(edge / scale); + const bool at_zero = zero == 0.0f; + scale = at_zero ? scale : edge / zero; + const float bias = at_zero ? 0.0f : edge; + + const ulong row = destination_row; + const ulong q_index = row * columns + (ulong)group.x * quant_group + lane; + pack[q_index] = (uchar)clamp(round((first - bias) / scale), 0.0f, bins); + pack[q_index + 32u] = (uchar)clamp(round((second - bias) / scale), 0.0f, bins); + if (lane == 0u) { + const ulong values = (ulong)destination_rows * columns; + const ulong groups_per_row = columns / quant_group; + const ulong parameter = row * groups_per_row + group.x; + device ushort *scales = (device ushort *)(pack + values); + device ushort *biases = scales + destination_rows * groups_per_row; + scales[parameter] = qwen_to_bf16(scale); + biases[parameter] = qwen_to_bf16(bias); + } +} + +// Faithful port of MTPLX hyper_connection_v3 R1: four grouped RMS reductions, +// q8 down+inject projection, then the family's activation boundaries. +kernel void kernel_qwen_hyper_v3_r1( + constant qwen_kernel_args &args [[buffer(0)]], + device float *mix_out [[buffer(1)]], + device const float *x [[buffer(2)]], + device const uchar *pack [[buffer(3)]], + device float *inject_out [[buffer(4)]], + device const uchar *norm_weight [[buffer(5)]], + device float *rms_out [[buffer(8)]], + uint3 group [[threadgroup_position_in_grid]], + uint3 thread_position [[thread_position_in_threadgroup]], + uint lane [[thread_index_in_simdgroup]], + uint simd_group [[simdgroup_index_in_threadgroup]]) { + constexpr uint hidden = 2560u; + constexpr uint width = 10240u; + constexpr uint rank = 320u; + constexpr uint rows = 324u; + constexpr uint quant_group = 64u; + constexpr uint groups_per_row = width / quant_group; + const uint tid = thread_position.x; + const ulong batch_width = (ulong)group.y * width; + const ulong batch_rank = (ulong)group.y * rank; + const ulong batch_streams = (ulong)group.y * 4u; + threadgroup float group_sums[4]; + threadgroup float partial[32]; + + float sums[4] = {0.0f, 0.0f, 0.0f, 0.0f}; + for (uint index = tid; index < width; index += 1024u) { + const float value = x[batch_width + index]; + sums[index / hidden] += value * value; + } + for (uint stream = 0u; stream < 4u; stream++) { + const float value = simd_sum(sums[stream]); + if (lane == 0u) partial[simd_group] = value; + threadgroup_barrier(mem_flags::mem_threadgroup); + if (simd_group == 0u) { + float total = partial[lane]; + total = simd_sum(total); + if (lane == 0u) group_sums[stream] = total; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + float rms[4]; + for (uint stream = 0u; stream < 4u; stream++) { + rms[stream] = rsqrt(group_sums[stream] / (float)hidden + 1e-6f); + } + if (group.x == 0u && tid < 4u) rms_out[batch_streams + tid] = rms[tid]; + + const uint row = group.x * 32u + simd_group; + if (row >= rows) return; + constexpr ulong q_bytes = (ulong)rows * width; + constexpr ulong parameter_count = (ulong)rows * groups_per_row; + device const uint *quantized = (device const uint *)pack; + device const ushort *scales = (device const ushort *)(pack + q_bytes); + device const ushort *biases = scales + parameter_count; + device const ushort *norm = + (device const ushort *)(norm_weight + args.u[13]); + float accumulator = 0.0f; + for (uint quant = lane; quant < groups_per_row; quant += 32u) { + const float scale = qwen_bf16(scales[(ulong)row * groups_per_row + quant]); + const float bias = qwen_bf16(biases[(ulong)row * groups_per_row + quant]); + const ulong base = (ulong)row * width + (ulong)quant * quant_group; + const uint activation_base = quant * quant_group; + float quantized_sum = 0.0f; + float normalized_sum = 0.0f; + for (uint word_index = 0u; word_index < 16u; word_index++) { + const uint word = quantized[(base >> 2u) + word_index]; + for (uint byte = 0u; byte < 4u; byte++) { + const uint index = activation_base + word_index * 4u + byte; + const float normalized = x[batch_width + index] * + qwen_bf16(norm[index]) * + rms[index / hidden]; + quantized_sum += (float)((word >> (byte * 8u)) & 255u) * normalized; + normalized_sum += normalized; + } + } + accumulator += scale * quantized_sum + bias * normalized_sum; + } + accumulator = simd_sum(accumulator); + if (lane == 0u) { + const float value = accumulator * 0.25f; + if (row < rank) { + mix_out[batch_rank + row] = qwen_round_bf16(value / (1.0f + exp(-value))); + } else { + inject_out[batch_streams + row - rank] = + qwen_round_bf16(2.0f / (1.0f + exp(-value))); + } + } +} + +// Faithful port of MTPLX hyper_connection_v3 R2: q8 up projection followed by +// the four-stream gated mean. One SIMD group owns each output dimension. +kernel void kernel_qwen_hyper_v3_r2( + constant qwen_kernel_args &args [[buffer(0)]], + device float *out [[buffer(1)]], + device const float *x [[buffer(2)]], + device const uchar *pack [[buffer(3)]], + device const float *mix [[buffer(4)]], + device const uchar *norm_weight [[buffer(5)]], + device const float *rms [[buffer(8)]], + uint2 group [[threadgroup_position_in_grid]], + uint lane [[thread_index_in_simdgroup]], + uint simd_group [[simdgroup_index_in_threadgroup]]) { + constexpr uint rank = 320u; + constexpr uint hidden = 2560u; + constexpr uint rows = 10240u; + constexpr uint quant_group = 64u; + constexpr uint groups_per_row = rank / quant_group; + const uint dimension = group.x * 32u + simd_group; + if (dimension >= hidden) return; + constexpr ulong q_bytes = (ulong)rows * rank; + constexpr ulong parameter_count = (ulong)rows * groups_per_row; + device const uint *quantized = (device const uint *)pack; + device const ushort *scales = (device const ushort *)(pack + q_bytes); + device const ushort *biases = scales + parameter_count; + device const ushort *norm = + (device const ushort *)(norm_weight + args.u[13]); + float mixed = 0.0f; + for (uint stream = 0u; stream < 4u; stream++) { + const uint row = stream * hidden + dimension; + float dot = 0.0f; + for (uint quant = 0u; quant < groups_per_row; quant++) { + const float scale = qwen_bf16(scales[(ulong)row * groups_per_row + quant]); + const float bias = qwen_bf16(biases[(ulong)row * groups_per_row + quant]); + const ulong base = (ulong)row * rank + (ulong)quant * quant_group; + const uint activation_base = quant * quant_group; + float quantized_sum = 0.0f; + float activation_sum = 0.0f; + for (uint word_index = lane; word_index < 16u; word_index += 32u) { + const uint word = quantized[(base >> 2u) + word_index]; + for (uint byte = 0u; byte < 4u; byte++) { + const float value = mix[(ulong)group.y * rank + activation_base + word_index * 4u + byte]; + quantized_sum += (float)((word >> (byte * 8u)) & 255u) * value; + activation_sum += value; + } + } + dot += scale * quantized_sum + bias * activation_sum; + } + dot = simd_sum(dot); + const uint index = stream * hidden + dimension; + const float normalized = x[(ulong)group.y * rows + index] * + qwen_bf16(norm[index]) * rms[(ulong)group.y * 4u + stream]; + mixed += (1.0f / (1.0f + exp(-dot))) * normalized; + } + if (lane == 0u) { + out[(ulong)group.y * hidden + dimension] = qwen_round_bf16(mixed * 0.25f); + } } kernel void kernel_qwen_repeat4( @@ -150,7 +1656,20 @@ kernel void kernel_qwen_repeat4( device float *out [[buffer(1)]], device const float *x [[buffer(2)]], uint index [[thread_position_in_grid]]) { - if (index < args.u[0] * 4u) out[index] = x[index % args.u[0]]; + const uint width = args.u[0]; + const uint row_width = width * 4u; + const uint row = index / row_width; + if (index < max(args.u[4], 1u) * row_width) { + out[index] = x[(ulong)row * width + index % width]; + } +} + +kernel void kernel_qwen_copy_bf16( + constant qwen_kernel_args &args [[buffer(0)]], + device ushort *out [[buffer(1)]], + device const ushort *input [[buffer(2)]], + uint index [[thread_position_in_grid]]) { + if (index < args.u[0]) out[index] = input[index]; } kernel void kernel_qwen_zero_rms( @@ -158,18 +1677,43 @@ kernel void kernel_qwen_zero_rms( device float *out [[buffer(1)]], device const float *x [[buffer(2)]], device const uchar *weight [[buffer(5)]], - uint group [[thread_position_in_grid]]) { + threadgroup float *scratch [[threadgroup(0)]], + uint group [[threadgroup_position_in_grid]], + uint tid [[thread_index_in_threadgroup]], + uint threads [[threads_per_threadgroup]], + uint lane [[thread_index_in_simdgroup]], + uint simd_group [[simdgroup_index_in_threadgroup]]) { const uint width = args.u[0]; const uint group_size = args.u[1]; - if (group >= width / group_size) return; + if (group >= (width / group_size) * max(args.u[4], 1u)) return; const uint start = group * group_size; - float variance = 0.0f; - for (uint i = 0; i < group_size; i++) variance = fma(x[start + i], x[start + i], variance); - const float scale = rsqrt(variance / (float)group_size + args.f[0]); - for (uint i = 0; i < group_size; i++) { - const uint index = start + i; - out[index] = x[index] * scale * - (1.0f + qwen_bf16(qwen_weight_u16(weight, args.u[13], index))); + threadgroup float *scale = scratch; + threadgroup float *partial = scratch + 1u; + float sum = 0.0f; + for (uint base = tid * 4u; base < group_size; base += threads * 4u) { + for (uint i = 0; i < 4u && base + i < group_size; i++) { + const float value = x[start + base + i]; + sum += value * value; + } + } + sum = simd_sum(sum); + if (simd_group == 0u) partial[lane] = 0.0f; + threadgroup_barrier(mem_flags::mem_threadgroup); + if (lane == 0u) partial[simd_group] = sum; + threadgroup_barrier(mem_flags::mem_threadgroup); + if (simd_group == 0u) { + sum = simd_sum(partial[lane]); + if (lane == 0u) scale[0] = precise::rsqrt(sum / (float)group_size + args.f[0]); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint base = tid * 4u; base < group_size; base += threads * 4u) { + for (uint i = 0; i < 4u && base + i < group_size; i++) { + const uint index = start + base + i; + const float normalized = qwen_round_bf16(x[index] * scale[0]); + float multiplier = qwen_bf16(qwen_weight_u16(weight, args.u[13], index % width)); + if (args.u[11] != 0u) multiplier = qwen_round_bf16(multiplier + 1.0f); + out[index] = qwen_round_bf16(multiplier * normalized); + } } } @@ -179,8 +1723,8 @@ kernel void kernel_qwen_silu_div4( device const float *x [[buffer(2)]], uint index [[thread_position_in_grid]]) { if (index >= args.u[0]) return; - const float value = x[index] * 0.25f; - out[index] = value / (1.0f + exp(-value)); + const float value = qwen_round_bf16(x[index] * 0.25f); + out[index] = qwen_round_bf16(value * qwen_silu_sigmoid_bf16(value)); } kernel void kernel_qwen_sigmoid( @@ -188,7 +1732,7 @@ kernel void kernel_qwen_sigmoid( device float *out [[buffer(1)]], device const float *x [[buffer(2)]], uint index [[thread_position_in_grid]]) { - if (index < args.u[0]) out[index] = 1.0f / (1.0f + exp(-x[index])); + if (index < args.u[0]) out[index] = qwen_sigmoid_bf16(x[index]); } kernel void kernel_qwen_sigmoid2_div4( @@ -196,7 +1740,10 @@ kernel void kernel_qwen_sigmoid2_div4( device float *out [[buffer(1)]], device const float *x [[buffer(2)]], uint index [[thread_position_in_grid]]) { - if (index < args.u[0]) out[index] = 2.0f / (1.0f + exp(-x[index] * 0.25f)); + if (index < args.u[0]) { + const float value = qwen_round_bf16(x[index] * 0.25f); + out[index] = qwen_round_bf16(2.0f * qwen_sigmoid_bf16(value)); + } } kernel void kernel_qwen_hyper_mix( @@ -205,13 +1752,18 @@ kernel void kernel_qwen_hyper_mix( device const float *normalized [[buffer(2)]], device const float *mix [[buffer(3)]], uint index [[thread_position_in_grid]]) { - if (index >= args.u[0]) return; + const uint hidden = args.u[0]; + if (index >= hidden * max(args.u[4], 1u)) return; + const uint row = index / hidden; + const uint column = index % hidden; + const ulong row_base = (ulong)row * hidden * 4u; float value = 0.0f; for (uint stream = 0; stream < 4u; stream++) { - const uint offset = stream * args.u[0] + index; - value = fma(normalized[offset], mix[offset], value); + const ulong offset = row_base + (ulong)stream * hidden + column; + const float product = qwen_round_bf16(normalized[offset] * mix[offset]); + value = qwen_round_bf16(value + product); } - out[index] = value * 0.25f; + out[index] = qwen_round_bf16(value * 0.25f); } kernel void kernel_qwen_hyper_inject( @@ -222,8 +1774,14 @@ kernel void kernel_qwen_hyper_inject( device const float *gate [[buffer(4)]], uint index [[thread_position_in_grid]]) { const uint hidden = args.u[0]; - if (index >= hidden * 4u) return; - out[index] = residual[index] + block[index % hidden] * gate[index / hidden]; + const uint row_width = hidden * 4u; + if (index >= row_width * max(args.u[4], 1u)) return; + const uint row = index / row_width; + const uint local = index % row_width; + const float injected = qwen_round_bf16( + block[(ulong)row * hidden + local % hidden] * + gate[(ulong)row * 4u + local / hidden]); + out[index] = qwen_round_bf16(residual[index] + injected); } kernel void kernel_qwen_ple_gate( @@ -232,16 +1790,43 @@ kernel void kernel_qwen_ple_gate( device const float *key [[buffer(2)]], device const float *query [[buffer(3)]], device const float *value [[buffer(4)]], - uint stream [[thread_position_in_grid]]) { + threadgroup bfloat *partials [[threadgroup(0)]], + uint3 group [[threadgroup_position_in_grid]], + uint3 thread_position [[thread_position_in_threadgroup]], + uint lane [[thread_index_in_simdgroup]], + uint simd_group [[simdgroup_index_in_threadgroup]]) { const uint hidden = args.u[0]; + const uint tid = thread_position.x; + const uint stream = group.x; if (stream >= 4u) return; - const ulong base = (ulong)stream * hidden; - float score = 0.0f; - for (uint i = 0; i < hidden; i++) score = fma(key[base + i], query[base + i], score); - score *= rsqrt((float)hidden); - const float transformed = copysign(sqrt(max(abs(score), 1.0e-6f)), score); - const float gate = 1.0f / (1.0f + exp(-transformed)); - for (uint i = 0; i < hidden; i++) out[base + i] = value[i] * gate; + const ulong base = ((ulong)group.y * 4u + stream) * hidden; + const ulong value_base = (ulong)group.y * hidden; + bfloat score = 0.0bf; + for (uint i = 0; i < 4u; i++) { + const uint index = tid * 4u + i; + const bfloat product = (bfloat)((bfloat)key[base + index] * (bfloat)query[base + index]); + score = product + score; + } + score = qwen_simd_sum_bf16(score); + if (lane == 0u) partials[simd_group] = score; + threadgroup_barrier(mem_flags::mem_threadgroup); + if (simd_group == 0u) { + score = lane < 20u ? partials[lane] : bfloat(0); + score = qwen_simd_sum_bf16(score); + if (lane == 0u) partials[0] = score; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + const bfloat scale = (bfloat)sqrt((float)hidden); + const float scaled = (float)(bfloat)((float)partials[0] / (float)scale); + const float transformed = scaled == 0.0f + ? 0.0f + : qwen_round_bf16(copysign( + qwen_round_bf16(sqrt(max(abs(scaled), qwen_round_bf16(1.0e-6f)))), + scaled)); + const float gate = qwen_sigmoid_bf16(transformed); + for (uint i = tid; i < hidden; i += 640u) { + out[base + i] = qwen_round_bf16(value[value_base + i] * gate); + } } kernel void kernel_qwen_ple_conv( @@ -249,22 +1834,35 @@ kernel void kernel_qwen_ple_conv( device float *out [[buffer(1)]], device const float *gated [[buffer(2)]], device const float *normalized [[buffer(3)]], - device ushort *state [[buffer(4)]], + device const ushort *state [[buffer(4)]], device const uchar *weight [[buffer(5)]], + device ushort *state_out [[buffer(8)]], uint channel [[thread_position_in_grid]]) { if (channel >= args.u[0]) return; - device ushort *history = state + (ulong)channel * 9u; - float value = fma(qwen_bf16(history[0]), - qwen_bf16(qwen_weight_u16(weight, args.u[13], (ulong)channel * 4u)), - fma(qwen_bf16(history[3]), - qwen_bf16(qwen_weight_u16(weight, args.u[13], (ulong)channel * 4u + 1u)), - fma(qwen_bf16(history[6]), - qwen_bf16(qwen_weight_u16(weight, args.u[13], (ulong)channel * 4u + 2u)), - normalized[channel] * - qwen_bf16(qwen_weight_u16(weight, args.u[13], (ulong)channel * 4u + 3u))))); - for (uint i = 0; i < 8u; i++) history[i] = history[i + 1u]; - history[8] = qwen_to_bf16(normalized[channel]); - out[channel] = gated[channel] + value / (1.0f + exp(-value)); + const ulong width = args.u[0]; + for (uint row = 0; row < max(args.u[4], 1u); row++) { + const ulong index = (ulong)row * width + channel; + float value = 0.0f; + for (uint tap = 0; tap < 4u; tap++) { + const uint sequence = row + tap * 3u; + const float input = tap == 3u + ? normalized[index] + : sequence < 9u + ? qwen_bf16(state[(ulong)sequence * width + channel]) + : normalized[(ulong)(sequence - 9u) * width + channel]; + value += input * qwen_bf16(qwen_weight_u16( + weight, args.u[13], (ulong)channel * 4u + tap)); + } + value = qwen_round_bf16(value); + const float activated = qwen_round_bf16(value * qwen_silu_sigmoid_bf16(value)); + out[index] = qwen_round_bf16(gated[index] + activated); + } + for (uint i = 0; i < 9u; i++) { + const uint sequence = args.u[4] + i; + state_out[(ulong)i * width + channel] = sequence < 9u + ? state[(ulong)sequence * width + channel] + : qwen_to_bf16(normalized[(ulong)(sequence - 9u) * width + channel]); + } } kernel void kernel_qwen_add( @@ -273,7 +1871,14 @@ kernel void kernel_qwen_add( device const float *a [[buffer(2)]], device const float *b [[buffer(3)]], uint index [[thread_position_in_grid]]) { - if (index < args.u[0]) out[index] = a[index] + b[index]; + if (index < args.u[0]) out[index] = qwen_round_bf16(a[index] + b[index]); +} + +kernel void kernel_qwen_zero_words( + constant qwen_kernel_args &args [[buffer(0)]], + device uint *out [[buffer(1)]], + uint index [[thread_position_in_grid]]) { + if (index < args.u[0]) out[(ulong)args.u[1] + index] = 0u; } kernel void kernel_qwen_conv_silu( @@ -284,68 +1889,686 @@ kernel void kernel_qwen_conv_silu( device const uchar *weight [[buffer(5)]], uint channel [[thread_position_in_grid]]) { if (channel >= args.u[0]) return; - device ushort *history = state + (ulong)channel * 3u; - float value = fma(qwen_bf16(history[0]), - qwen_bf16(qwen_weight_u16(weight, args.u[13], (ulong)channel * 4u)), - fma(qwen_bf16(history[1]), - qwen_bf16(qwen_weight_u16(weight, args.u[13], (ulong)channel * 4u + 1u)), - fma(qwen_bf16(history[2]), - qwen_bf16(qwen_weight_u16(weight, args.u[13], (ulong)channel * 4u + 2u)), - x[channel] * - qwen_bf16(qwen_weight_u16(weight, args.u[13], (ulong)channel * 4u + 3u))))); - history[0] = history[1]; - history[1] = history[2]; - history[2] = qwen_to_bf16(x[channel]); - out[channel] = value / (1.0f + exp(-value)); + const ulong c = channel; + const ulong width = args.u[0]; + const float value = + qwen_bf16(state[c]) * + qwen_bf16(qwen_weight_u16(weight, args.u[13], c * 4u)) + + qwen_bf16(state[width + c]) * + qwen_bf16(qwen_weight_u16(weight, args.u[13], c * 4u + 1u)) + + qwen_bf16(state[2u * width + c]) * + qwen_bf16(qwen_weight_u16(weight, args.u[13], c * 4u + 2u)) + + x[c] * qwen_bf16(qwen_weight_u16(weight, args.u[13], c * 4u + 3u)); + state[c] = state[width + c]; + state[width + c] = state[2u * width + c]; + state[2u * width + c] = qwen_to_bf16(x[c]); + const float rounded = qwen_round_bf16(value); + out[c] = qwen_round_bf16(rounded * qwen_silu_sigmoid_bf16(rounded)); } -kernel void kernel_qwen_gdn_step( +// Exact MTPLX empty-cache decode path: fused depthwise conv, SiLU and +// per-head q/k normalization. Later rows use kernel_qwen_gdn_step_fused. +kernel void kernel_qwen_gdn_conv_norm( + constant qwen_kernel_args &args [[buffer(0)]], + device float *out [[buffer(1)]], + device const float *x [[buffer(2)]], + device ushort *state [[buffer(3)]], + device const uchar *weight [[buffer(5)]], + uint tid [[thread_index_in_threadgroup]], + uint lane [[thread_index_in_simdgroup]], + uint simd_group [[simdgroup_index_in_threadgroup]], + uint group [[threadgroup_position_in_grid]]) { + constexpr uint width = 10240u; + constexpr uint key_width = 2048u; + constexpr uint dim = 128u; + constexpr float inverse_scale = 0.08838834764831845f; + const uint channel = group * 1024u + tid; + if (channel >= width) return; + + threadgroup float values[1024]; + threadgroup float partial[32]; + const float convolved = + qwen_bf16(qwen_weight_u16(weight, args.u[13], (ulong)channel * 4u)) * + qwen_bf16(state[channel]) + + qwen_bf16(qwen_weight_u16(weight, args.u[13], (ulong)channel * 4u + 1u)) * + qwen_bf16(state[width + channel]) + + qwen_bf16(qwen_weight_u16(weight, args.u[13], (ulong)channel * 4u + 2u)) * + qwen_bf16(state[2u * width + channel]) + + qwen_bf16(qwen_weight_u16(weight, args.u[13], (ulong)channel * 4u + 3u)) * + x[channel]; + const float activated = convolved / (1.0f + exp(-convolved)); + state[channel] = state[width + channel]; + state[width + channel] = state[2u * width + channel]; + state[2u * width + channel] = qwen_to_bf16(x[channel]); + + if (channel >= 2u * key_width) { + out[channel] = qwen_round_bf16(activated); + return; + } + + values[tid] = activated; + float sum = simd_sum(activated * activated); + if (lane == 0u) partial[simd_group] = sum; + threadgroup_barrier(mem_flags::mem_threadgroup); + const uint first = (simd_group / 4u) * 4u; + sum = partial[first] + partial[first + 1u] + + partial[first + 2u] + partial[first + 3u]; + const float normalized = values[tid] * rsqrt(sum + 1.0e-6f); + out[channel] = channel < key_width + ? qwen_round_bf16(normalized * inverse_scale) + : qwen_round_bf16(normalized); +} + +// MTPLX fused_gdn_conv_norm_rows, including its in-window convolution tail. +kernel void kernel_qwen_gdn_conv_norm_rows( + constant qwen_kernel_args &args [[buffer(0)]], + device float *q_out [[buffer(1)]], + device const float *qkv [[buffer(2)]], + device float *k_out [[buffer(3)]], + device float *v_out [[buffer(4)]], + device const ushort *state [[buffer(5)]], + device const uchar *weight [[buffer(6)]], + device ushort *state_out [[buffer(8)]], + uint tid [[thread_index_in_threadgroup]], + uint lane [[thread_index_in_simdgroup]], + uint simd_group [[simdgroup_index_in_threadgroup]], + uint group [[threadgroup_position_in_grid]]) { + constexpr uint width = 10240u; + constexpr uint key_width = 2048u; + constexpr uint dim = 128u; + constexpr float inverse_scale = 0.08838834764831845f; + const uint channel = group * 1024u + tid; + if (channel >= width) return; + + threadgroup float values[1024]; + threadgroup float partial[32]; + const float w0 = qwen_bf16(qwen_weight_u16(weight, args.u[14], (ulong)channel * 4u)); + const float w1 = qwen_bf16(qwen_weight_u16(weight, args.u[14], (ulong)channel * 4u + 1u)); + const float w2 = qwen_bf16(qwen_weight_u16(weight, args.u[14], (ulong)channel * 4u + 2u)); + const float w3 = qwen_bf16(qwen_weight_u16(weight, args.u[14], (ulong)channel * 4u + 3u)); + const bool value_channel = channel >= 2u * key_width; + + for (uint row = 0u; row < args.u[4]; row++) { + const float x0 = row < 3u + ? qwen_bf16(state[(ulong)row * width + channel]) + : qkv[(ulong)(row - 3u) * width + channel]; + const float x1 = row + 1u < 3u + ? qwen_bf16(state[(ulong)(row + 1u) * width + channel]) + : qkv[(ulong)(row - 2u) * width + channel]; + const float x2 = row + 2u < 3u + ? qwen_bf16(state[(ulong)(row + 2u) * width + channel]) + : qkv[(ulong)(row - 1u) * width + channel]; + const float x3 = qkv[(ulong)row * width + channel]; + const float convolved = w0 * x0 + w1 * x1 + w2 * x2 + w3 * x3; + const float activated = convolved / (1.0f + exp(-convolved)); + if (value_channel) { + v_out[(ulong)row * (width - 2u * key_width) + channel - 2u * key_width] = + qwen_round_bf16(activated); + continue; + } + values[tid] = activated; + float sum = simd_sum(activated * activated); + if (lane == 0u) partial[simd_group] = sum; + threadgroup_barrier(mem_flags::mem_threadgroup); + const uint first = (simd_group / 4u) * 4u; + sum = partial[first] + partial[first + 1u] + + partial[first + 2u] + partial[first + 3u]; + const float normalized = values[tid] * rsqrt(sum + 1.0e-6f); + if (channel < key_width) { + q_out[(ulong)row * key_width + channel] = + qwen_round_bf16(normalized * inverse_scale); + } else { + k_out[(ulong)row * key_width + channel - key_width] = + qwen_round_bf16(normalized); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + + for (uint tail = 0u; tail < 3u; tail++) { + const uint sequence = args.u[4] + tail; + const float value = sequence < 3u + ? qwen_bf16(state[(ulong)sequence * width + channel]) + : qkv[(ulong)(sequence - 3u) * width + channel]; + state_out[(ulong)tail * width + channel] = qwen_to_bf16(value); + } +} + +kernel void kernel_qwen_gdn_l2norm( + constant qwen_kernel_args &args [[buffer(0)]], + device float *out [[buffer(1)]], + device const float *x [[buffer(2)]], + uint head [[threadgroup_position_in_grid]], + uint lane [[thread_index_in_simdgroup]]) { + constexpr uint dim = 128u; + constexpr uint key_width = 2048u; + if (head >= 16u) return; + const ulong qbase = (ulong)head * dim; + const ulong kbase = key_width + qbase; + float qv[4]; + float kv[4]; + float qsum = 0.0f; + float ksum = 0.0f; + for (uint i = 0; i < 4u; i++) { + const uint element = lane * 4u + i; + qv[i] = x[qbase + element]; + kv[i] = x[kbase + element]; + qsum += qv[i] * qv[i]; + ksum += kv[i] * kv[i]; + } + qsum = simd_sum(qsum); + ksum = simd_sum(ksum); + const float qinv = precise::rsqrt(qsum + 1.0e-6f); + const float kinv = precise::rsqrt(ksum + 1.0e-6f); + for (uint i = 0; i < 4u; i++) { + const uint element = lane * 4u + i; + const float normalized_q = qwen_round_bf16(qv[i] * qinv); + out[qbase + element] = qwen_round_bf16(normalized_q * 0.08837890625f); + out[kbase + element] = qwen_round_bf16(kv[i] * kinv); + } +} + +kernel void kernel_qwen_gdn_controls( + constant qwen_kernel_args &args [[buffer(0)]], + device float *controls [[buffer(1)]], + device const uchar *a_log [[buffer(5)]], + device const uchar *dt_bias [[buffer(6)]], + uint index [[thread_position_in_grid]]) { + constexpr uint heads = 48u; + constexpr uint controls_width = 6240u; + const uint rows = max(args.u[4], 1u); + if (index >= heads * rows) return; + const uint row = index / heads; + const uint head = index % heads; + const ulong base = (ulong)row * controls_width; + const float step = qwen_round_bf16( + controls[base + args.u[1] + head] + + qwen_bf16(qwen_weight_u16(dt_bias, args.u[14], head))); + const float maximum = max(step, 0.0f); + const float minimum = min(step, 0.0f); + const float exponential = qwen_round_bf16( + precise::exp(qwen_round_bf16(minimum - maximum))); + const float plus_one = 1.0f + exponential; + const float logarithm = plus_one == 1.0f + ? exponential + : exponential * (precise::log(plus_one) / (plus_one - 1.0f)); + const float softplus = qwen_round_bf16( + maximum + qwen_round_bf16(logarithm)); + const float rate = precise::exp( + qwen_bf16(qwen_weight_u16(a_log, args.u[13], head))); + controls[base + args.u[1] + head] = precise::exp(-rate * softplus); + controls[base + args.u[0] + head] = + qwen_sigmoid_bf16(controls[base + args.u[0] + head]); +} + +kernel void kernel_qwen_pack_gdn_controls( + constant qwen_kernel_args &args [[buffer(0)]], + device float *controls [[buffer(1)]], + device const float *z [[buffer(2)]], + device const float *b [[buffer(3)]], + device const float *a [[buffer(4)]], + uint index [[thread_position_in_grid]]) { + constexpr uint value_width = 6144u; + constexpr uint heads = 48u; + constexpr uint controls_width = value_width + 2u * heads; + if (index >= args.u[4] * controls_width) return; + const uint row = index / controls_width; + const uint column = index % controls_width; + if (column < value_width) { + controls[index] = z[(ulong)row * value_width + column]; + } else if (column < value_width + heads) { + controls[index] = b[(ulong)row * heads + column - value_width]; + } else { + controls[index] = a[(ulong)row * heads + column - value_width - heads]; + } +} + +kernel void kernel_qwen_gdn_delta_packed( constant qwen_kernel_args &args [[buffer(0)]], device float *out [[buffer(1)]], device const float *qkv [[buffer(2)]], device const float *controls [[buffer(3)]], device float *state [[buffer(4)]], - device const uchar *a_log [[buffer(5)]], - device const uchar *dt_bias [[buffer(6)]], - uint2 gid [[thread_position_in_grid]]) { - const uint value_index = gid.x; - const uint head = gid.y; - const uint dim = args.u[0]; - const uint key_heads = args.u[1]; - const uint value_heads = args.u[2]; - if (value_index >= dim || head >= value_heads) return; - const uint key_head = head / (value_heads / key_heads); - device const float *q_raw = qkv + (ulong)key_head * dim; - device const float *k_raw = qkv + (ulong)key_heads * dim + (ulong)key_head * dim; - device const float *value = qkv + (ulong)key_heads * dim * 2u + (ulong)head * dim; - float qsum = 0.0f; - float ksum = 0.0f; - for (uint i = 0; i < dim; i++) { - qsum = fma(q_raw[i], q_raw[i], qsum); - ksum = fma(k_raw[i], k_raw[i], ksum); + uint2 group [[threadgroup_position_in_grid]], + uint lane [[thread_index_in_simdgroup]], + uint simd_group [[simdgroup_index_in_threadgroup]]) { + constexpr uint dim = 128u; + constexpr uint key_width = 2048u; + constexpr uint rows_per_simdgroup = 8u; + constexpr uint values_per_lane = 32u; + constexpr uint partials_per_lane = 8u; + const uint head = group.y; + const uint row_group = group.x * 2u + simd_group; + const uint row_in_simdgroup = lane / 4u; + const uint lane_in_row = lane & 3u; + const uint value_index = row_group * rows_per_simdgroup + row_in_simdgroup; + const uint key_head = head / 3u; + const uint key_start = lane_in_row * values_per_lane; + device const float *q = qkv + (ulong)key_head * dim + key_start; + device const float *k = qkv + key_width + (ulong)key_head * dim + key_start; + device const float *v = qkv + 2u * key_width + (ulong)head * dim; + device float *row = state + ((ulong)head * dim + value_index) * dim + key_start; + float values[values_per_lane]; + float partial[partials_per_lane]; + const float decay = controls[args.u[1] + head]; + for (uint i = 0; i < values_per_lane; i++) values[i] = row[i]; + for (uint block = 0; block < partials_per_lane; block++) { + float sum = 0.0f; + for (uint i = 0; i < 4u; i++) { + const uint element = block * 4u + i; + values[element] *= decay; + sum += values[element] * k[element]; + } + partial[block] = sum; } - const float qscale = rsqrt(qsum + args.f[0]) * rsqrt((float)dim); - const float kscale = rsqrt(ksum + args.f[0]); - const float beta = 1.0f / (1.0f + exp(-controls[args.u[3] + head])); - const float step = controls[args.u[4] + head] + - qwen_bf16(qwen_weight_u16(dt_bias, args.u[14], head)); + float prediction = + ((partial[0] + partial[1]) + (partial[2] + partial[3])) + + ((partial[4] + partial[5]) + (partial[6] + partial[7])); + prediction += simd_shuffle_xor(prediction, 1u); + prediction += simd_shuffle_xor(prediction, 2u); + const float delta = + (v[value_index] - prediction) * controls[args.u[0] + head]; + for (uint block = 0; block < partials_per_lane; block++) { + float sum = 0.0f; + for (uint i = 0; i < 4u; i++) { + const uint element = block * 4u + i; + values[element] += k[element] * delta; + sum += values[element] * q[element]; + } + partial[block] = sum; + } + float result = + ((partial[0] + partial[1]) + (partial[2] + partial[3])) + + ((partial[4] + partial[5]) + (partial[6] + partial[7])); + result += simd_shuffle_xor(result, 1u); + result += simd_shuffle_xor(result, 2u); + for (uint i = 0; i < values_per_lane; i++) row[i] = values[i]; + if (lane_in_row == 0u) { + out[(ulong)head * dim + value_index] = qwen_round_bf16(result); + } +} + +// mlx-lm gated_delta_step_packed_btree for S=2..4 verifier rows. +kernel void kernel_qwen_gdn_delta_rows( + constant qwen_kernel_args &args [[buffer(0)]], + device float *values_out [[buffer(1)]], + device const float *q_rows [[buffer(2)]], + device const float *k_rows [[buffer(3)]], + device const float *controls [[buffer(4)]], + device const float *state [[buffer(5)]], + device float *state_out [[buffer(8)]], + uint2 group [[threadgroup_position_in_grid]], + uint lane [[thread_index_in_simdgroup]], + uint simd_group [[simdgroup_index_in_threadgroup]]) { + constexpr uint dim = 128u; + constexpr uint key_width = 2048u; + constexpr uint value_width = 6144u; + constexpr uint controls_width = 6240u; + constexpr uint rows_per_simdgroup = 8u; + constexpr uint values_per_lane = 32u; + constexpr uint partials_per_lane = 8u; + const uint head = group.y; + const uint row_group = group.x * 2u + simd_group; + const uint row_in_simdgroup = lane / 4u; + const uint lane_in_row = lane & 3u; + const uint value_index = row_group * rows_per_simdgroup + row_in_simdgroup; + const uint key_head = head / 3u; + const uint key_start = lane_in_row * values_per_lane; + device const float *q = q_rows + (ulong)key_head * dim + key_start; + device const float *k = k_rows + (ulong)key_head * dim + key_start; + device const float *state_row = + state + ((ulong)head * dim + value_index) * dim + key_start; + device float *state_out_row = + state_out + ((ulong)head * dim + value_index) * dim + key_start; + float values[values_per_lane]; + float partial[partials_per_lane]; + for (uint index = 0u; index < values_per_lane; index++) { + values[index] = state_row[index]; + } + + for (uint token = 0u; token < args.u[4]; token++) { + const ulong control_base = (ulong)token * controls_width; + const float decay = controls[control_base + args.u[1] + head]; + for (uint block = 0u; block < partials_per_lane; block++) { + float sum = 0.0f; + for (uint index = 0u; index < 4u; index++) { + const uint element = block * 4u + index; + values[element] *= decay; + sum += values[element] * k[element]; + } + partial[block] = sum; + } + float prediction = + ((partial[0] + partial[1]) + (partial[2] + partial[3])) + + ((partial[4] + partial[5]) + (partial[6] + partial[7])); + prediction += simd_shuffle_xor(prediction, 1u); + prediction += simd_shuffle_xor(prediction, 2u); + const ulong value_offset = + (ulong)token * value_width + (ulong)head * dim + value_index; + const float delta = (values_out[value_offset] - prediction) * + controls[control_base + args.u[0] + head]; + for (uint block = 0u; block < partials_per_lane; block++) { + float sum = 0.0f; + for (uint index = 0u; index < 4u; index++) { + const uint element = block * 4u + index; + values[element] += k[element] * delta; + sum += values[element] * q[element]; + } + partial[block] = sum; + } + float result = + ((partial[0] + partial[1]) + (partial[2] + partial[3])) + + ((partial[4] + partial[5]) + (partial[6] + partial[7])); + result += simd_shuffle_xor(result, 1u); + result += simd_shuffle_xor(result, 2u); + if (lane_in_row == 0u) { + values_out[value_offset] = qwen_round_bf16(result); + } + q += key_width; + k += key_width; + } + for (uint index = 0u; index < values_per_lane; index++) { + state_out_row[index] = values[index]; + } +} + +kernel void kernel_qwen_gdn_step_fused( + constant qwen_kernel_args &args [[buffer(0)]], + device float *out [[buffer(1)]], + device const float *qkv [[buffer(2)]], + device const float *controls [[buffer(3)]], + device float *state [[buffer(4)]], + device const uchar *conv_weight [[buffer(5)]], + device const uchar *a_log [[buffer(6)]], + device const uchar *dt_bias [[buffer(7)]], + device ushort *conv_state [[buffer(8)]], + device const uchar *norm_weight [[buffer(9)]], + threadgroup float *scratch [[threadgroup(0)]], + uint2 group [[threadgroup_position_in_grid]], + ushort tid [[thread_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]], + ushort simd_group [[simdgroup_index_in_threadgroup]]) { + constexpr uint dim = 128u; + constexpr uint key_width = 2048u; + constexpr uint conv_width = 10240u; + constexpr uint values_per_key = 3u; + constexpr float inverse_scale = 0.08838834764831845f; + const uint head = group.y; + const uint key_head = head / values_per_key; + threadgroup float *q = scratch; + threadgroup float *k = q + dim; + threadgroup float *v = k + dim; + threadgroup float *result = v + dim; + threadgroup float *partial = result + dim; + device ushort *conv_next = (device ushort *)(out + 6144u); + + for (uint local = tid; local < 3u * dim; local += 256u) { + uint channel; + if (local < dim) { + channel = key_head * dim + local; + } else if (local < 2u * dim) { + channel = key_width + key_head * dim + local - dim; + } else { + channel = 2u * key_width + head * dim + local - 2u * dim; + } + const float value = + qwen_bf16(qwen_weight_u16(conv_weight, args.u[13], (ulong)channel * 4u)) * + qwen_bf16(conv_state[channel]) + + qwen_bf16(qwen_weight_u16(conv_weight, args.u[13], (ulong)channel * 4u + 1u)) * + qwen_bf16(conv_state[conv_width + channel]) + + qwen_bf16(qwen_weight_u16(conv_weight, args.u[13], (ulong)channel * 4u + 2u)) * + qwen_bf16(conv_state[2u * conv_width + channel]) + + qwen_bf16(qwen_weight_u16(conv_weight, args.u[13], (ulong)channel * 4u + 3u)) * + qkv[channel]; + const float activated = value / (1.0f + exp(-value)); + if (local >= 2u * dim || head % values_per_key == 0u) { + conv_next[channel] = conv_state[conv_width + channel]; + conv_next[conv_width + channel] = conv_state[2u * conv_width + channel]; + conv_next[2u * conv_width + channel] = qwen_to_bf16(qkv[channel]); + } + if (local < dim) { + q[local] = activated; + } else if (local < 2u * dim) { + k[local - dim] = activated; + } else { + v[local - 2u * dim] = qwen_round_bf16(activated); + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + float squared = 0.0f; + if (tid < dim) { + squared = q[tid] * q[tid]; + } else if (tid < 2u * dim) { + squared = k[tid - dim] * k[tid - dim]; + } + squared = simd_sum(squared); + if (lane == 0u) partial[simd_group] = squared; + threadgroup_barrier(mem_flags::mem_threadgroup); + const float q_scale = + rsqrt(partial[0] + partial[1] + partial[2] + partial[3] + 1.0e-6f) * + inverse_scale; + const float k_scale = rsqrt(partial[4] + partial[5] + partial[6] + partial[7] + 1.0e-6f); + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid < dim) { + q[tid] = qwen_round_bf16(q[tid] * q_scale); + } else if (tid < 2u * dim) { + k[tid - dim] = qwen_round_bf16(k[tid - dim] * k_scale); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + const float step = controls[args.u[1] + head] + + qwen_bf16(qwen_weight_u16(dt_bias, args.u[15], head)); const float softplus = max(step, 0.0f) + log(1.0f + exp(-abs(step))); - const float decay = exp(-exp(qwen_bf16(qwen_weight_u16(a_log, args.u[13], head))) * softplus); - device float *column = state + ((ulong)head * dim * dim) + value_index; - float prediction = 0.0f; - for (uint i = 0; i < dim; i++) { - prediction = fma(column[(ulong)i * dim] * decay, k_raw[i] * kscale, prediction); + const float decay = exp( + -exp(qwen_bf16(qwen_weight_u16(a_log, args.u[14], head))) * softplus); + const float beta = 1.0f / (1.0f + exp(-controls[args.u[0] + head])); + + for (uint value_index = simd_group; value_index < dim; value_index += 8u) { + const ulong state_row = ((ulong)head * dim + value_index) * dim; + const uint column = 4u * lane; + float s0 = state[state_row + column] * decay; + float s1 = state[state_row + column + 1u] * decay; + float s2 = state[state_row + column + 2u] * decay; + float s3 = state[state_row + column + 3u] * decay; + const float k0 = k[column]; + const float k1 = k[column + 1u]; + const float k2 = k[column + 2u]; + const float k3 = k[column + 3u]; + float prediction = s0 * k0 + s1 * k1 + s2 * k2 + s3 * k3; + prediction = simd_sum(prediction); + const float delta = (v[value_index] - prediction) * beta; + s0 += k0 * delta; + s1 += k1 * delta; + s2 += k2 * delta; + s3 += k3 * delta; + float value = s0 * q[column] + s1 * q[column + 1u] + + s2 * q[column + 2u] + s3 * q[column + 3u]; + value = simd_sum(value); + state[state_row + column] = s0; + state[state_row + column + 1u] = s1; + state[state_row + column + 2u] = s2; + state[state_row + column + 3u] = s3; + if (lane == 0u) { + result[value_index] = qwen_round_bf16(value); + } } - const float delta = (value[value_index] - prediction) * beta; - float result = 0.0f; - for (uint i = 0; i < dim; i++) { - const ulong offset = (ulong)i * dim; - const float updated = column[offset] * decay + k_raw[i] * kscale * delta; - column[offset] = updated; - result = fma(updated, q_raw[i] * qscale, result); + threadgroup_barrier(mem_flags::mem_threadgroup); + + float output_squared = tid < dim ? result[tid] * result[tid] : 0.0f; + output_squared = simd_sum(output_squared); + if (lane == 0u && simd_group < 4u) partial[simd_group] = output_squared; + threadgroup_barrier(mem_flags::mem_threadgroup); + const float output_scale = rsqrt( + (partial[0] + partial[1] + partial[2] + partial[3]) / (float)dim + 1.0e-6f); + if (tid < dim) { + const float normalized = qwen_round_bf16( + result[tid] * output_scale * + qwen_bf16(qwen_weight_u16(norm_weight, args.u[9], tid))); + const float z = controls[(ulong)head * dim + tid]; + const float gate = 1.0f / (1.0f + exp(-z)); + out[(ulong)head * dim + tid] = qwen_round_bf16(gate * normalized); + } +} + +// MTPLX verify-width GDN: convolution is evaluated against the immutable +// pre-window tail plus all earlier rows, while gated-delta recurrence advances +// sequentially through S=2..4 in one head-owned threadgroup. +kernel void kernel_qwen_gdn_verify( + constant qwen_kernel_args &args [[buffer(0)]], + device float *out [[buffer(1)]], + device const float *qkv [[buffer(2)]], + device const float *controls [[buffer(3)]], + device float *state [[buffer(4)]], + device const uchar *conv_weight [[buffer(5)]], + device const uchar *a_log [[buffer(6)]], + device const uchar *dt_bias [[buffer(7)]], + device const ushort *conv_state [[buffer(8)]], + device const uchar *norm_weight [[buffer(9)]], + threadgroup float *scratch [[threadgroup(0)]], + uint head [[threadgroup_position_in_grid]], + ushort tid [[thread_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]], + ushort simd_group [[simdgroup_index_in_threadgroup]]) { + constexpr uint dim = 128u; + constexpr uint key_width = 2048u; + constexpr uint conv_width = 10240u; + constexpr uint values_per_key = 3u; + constexpr uint controls_width = 6240u; + constexpr uint value_width = 6144u; + constexpr float inverse_scale = 0.08838834764831845f; + const uint key_head = head / values_per_key; + threadgroup float *q = scratch; + threadgroup float *k = q + dim; + threadgroup float *v = k + dim; + threadgroup float *result = v + dim; + threadgroup float *partial = result + dim; + + for (uint token = 0u; token < args.u[4]; token++) { + for (uint local = tid; local < 3u * dim; local += 256u) { + uint channel; + if (local < dim) { + channel = key_head * dim + local; + } else if (local < 2u * dim) { + channel = key_width + key_head * dim + local - dim; + } else { + channel = 2u * key_width + head * dim + local - 2u * dim; + } + float convolved = 0.0f; + for (uint tap = 0u; tap < 4u; tap++) { + const uint sequence = token + tap; + const float input = sequence < 3u + ? qwen_bf16(conv_state[(ulong)sequence * conv_width + channel]) + : qkv[(ulong)(sequence - 3u) * conv_width + channel]; + convolved += input * qwen_bf16(qwen_weight_u16( + conv_weight, args.u[13], (ulong)channel * 4u + tap)); + } + const float activated = convolved / (1.0f + exp(-convolved)); + if (local < dim) { + q[local] = activated; + } else if (local < 2u * dim) { + k[local - dim] = activated; + } else { + v[local - 2u * dim] = qwen_round_bf16(activated); + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + float squared = 0.0f; + if (tid < dim) { + squared = q[tid] * q[tid]; + } else if (tid < 2u * dim) { + squared = k[tid - dim] * k[tid - dim]; + } + squared = simd_sum(squared); + if (lane == 0u) partial[simd_group] = squared; + threadgroup_barrier(mem_flags::mem_threadgroup); + const float q_scale = + rsqrt(partial[0] + partial[1] + partial[2] + partial[3] + 1.0e-6f) * + inverse_scale; + const float k_scale = + rsqrt(partial[4] + partial[5] + partial[6] + partial[7] + 1.0e-6f); + if (tid < dim) { + q[tid] = qwen_round_bf16(q[tid] * q_scale); + } else if (tid < 2u * dim) { + k[tid - dim] = qwen_round_bf16(k[tid - dim] * k_scale); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + const ulong control_base = (ulong)token * controls_width; + const float step = controls[control_base + args.u[1] + head] + + qwen_bf16(qwen_weight_u16(dt_bias, args.u[15], head)); + const float softplus = max(step, 0.0f) + log(1.0f + exp(-abs(step))); + const float decay = exp( + -exp(qwen_bf16(qwen_weight_u16(a_log, args.u[14], head))) * softplus); + const float beta = 1.0f / + (1.0f + exp(-controls[control_base + args.u[0] + head])); + + for (uint value_index = simd_group; value_index < dim; value_index += 8u) { + const ulong state_row = ((ulong)head * dim + value_index) * dim; + const uint column = 4u * lane; + float s0 = state[state_row + column] * decay; + float s1 = state[state_row + column + 1u] * decay; + float s2 = state[state_row + column + 2u] * decay; + float s3 = state[state_row + column + 3u] * decay; + const float k0 = k[column]; + const float k1 = k[column + 1u]; + const float k2 = k[column + 2u]; + const float k3 = k[column + 3u]; + float prediction = s0 * k0 + s1 * k1 + s2 * k2 + s3 * k3; + prediction = simd_sum(prediction); + const float delta = (v[value_index] - prediction) * beta; + s0 += k0 * delta; + s1 += k1 * delta; + s2 += k2 * delta; + s3 += k3 * delta; + float value = s0 * q[column] + s1 * q[column + 1u] + + s2 * q[column + 2u] + s3 * q[column + 3u]; + value = simd_sum(value); + state[state_row + column] = s0; + state[state_row + column + 1u] = s1; + state[state_row + column + 2u] = s2; + state[state_row + column + 3u] = s3; + if (lane == 0u) result[value_index] = qwen_round_bf16(value); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + float output_squared = tid < dim ? result[tid] * result[tid] : 0.0f; + output_squared = simd_sum(output_squared); + if (lane == 0u && simd_group < 4u) partial[simd_group] = output_squared; + threadgroup_barrier(mem_flags::mem_threadgroup); + const float output_scale = rsqrt( + (partial[0] + partial[1] + partial[2] + partial[3]) / + (float)dim + 1.0e-6f); + if (tid < dim) { + const float normalized = qwen_round_bf16( + result[tid] * output_scale * + qwen_bf16(qwen_weight_u16(norm_weight, args.u[9], tid))); + const float z = controls[control_base + (ulong)head * dim + tid]; + const float gate = 1.0f / (1.0f + exp(-z)); + out[(ulong)token * value_width + (ulong)head * dim + tid] = + qwen_round_bf16(gate * normalized); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } +} + +kernel void kernel_qwen_gdn_commit_conv_rows( + constant qwen_kernel_args &args [[buffer(0)]], + device ushort *state [[buffer(1)]], + device const float *qkv [[buffer(2)]], + uint channel [[thread_position_in_grid]]) { + const uint width = args.u[0]; + if (channel >= width) return; + float tail[3]; + for (uint tap = 0u; tap < 3u; tap++) { + const uint sequence = args.u[4] + tap; + tail[tap] = sequence < 3u + ? qwen_bf16(state[(ulong)sequence * width + channel]) + : qkv[(ulong)(sequence - 3u) * width + channel]; + } + for (uint tap = 0u; tap < 3u; tap++) { + state[(ulong)tap * width + channel] = qwen_to_bf16(tail[tap]); } - device float *head_out = out + (ulong)head * dim; - head_out[value_index] = result; } kernel void kernel_qwen_gdn_norm_gate( @@ -354,18 +2577,35 @@ kernel void kernel_qwen_gdn_norm_gate( device const float *x [[buffer(2)]], device const float *controls [[buffer(3)]], device const uchar *weight [[buffer(5)]], - uint head [[thread_position_in_grid]]) { - const uint dim = args.u[0]; - if (head >= args.u[1]) return; - device const float *row = x + (ulong)head * dim; - float variance = 0.0f; - for (uint i = 0; i < dim; i++) variance = fma(row[i], row[i], variance); - const float scale = rsqrt(variance / (float)dim + args.f[0]); - for (uint i = 0; i < dim; i++) { - const ulong index = (ulong)head * dim + i; - out[index] = row[i] * scale * - qwen_bf16(qwen_weight_u16(weight, args.u[13], i)) / - (1.0f + exp(-controls[index])); + uint2 group [[threadgroup_position_in_grid]], + uint lane [[thread_index_in_simdgroup]]) { + constexpr uint dim = 128u; + constexpr uint value_width = 6144u; + constexpr uint controls_width = 6240u; + const uint head = group.x; + const uint row = group.y; + if (head >= 48u) return; + const ulong base = (ulong)row * value_width + (ulong)head * dim; + const ulong control_base = (ulong)row * controls_width + (ulong)head * dim; + float values[4]; + float sum = 0.0f; + for (uint i = 0; i < 4u; i++) { + const uint element = lane * 4u + i; + values[i] = x[base + element]; + sum += values[i] * values[i]; + } + sum = simd_sum(sum); + const float scale = precise::rsqrt(sum / (float)dim + 1.0e-6f); + for (uint i = 0; i < 4u; i++) { + const uint element = lane * 4u + i; + const float normalized = qwen_round_bf16(values[i] * scale); + const float weighted = qwen_round_bf16( + qwen_bf16(qwen_weight_u16(weight, args.u[13], element)) * normalized); + const float z = controls[control_base + element]; + const float exponential = precise::exp(abs(z)); + const float tail = 1.0f / (1.0f + exponential); + const float gate = z < 0.0f ? tail : 1.0f - tail; + out[base + element] = qwen_round_bf16(gate * weighted); } } @@ -376,38 +2616,218 @@ kernel void kernel_qwen_swiglu( device const float *up [[buffer(3)]], uint index [[thread_position_in_grid]]) { if (index >= args.u[0]) return; - out[index] = gate[index] / (1.0f + exp(-gate[index])) * up[index]; + const float activated = qwen_round_bf16( + gate[index] * qwen_silu_sigmoid_bf16(gate[index])); + out[index] = qwen_round_bf16(activated * up[index]); +} + +kernel void kernel_qwen_unpack_gdn_inputs( + constant qwen_kernel_args &args [[buffer(0)]], + device float *qkv [[buffer(1)]], + device const float *fused [[buffer(2)]], + device float *controls [[buffer(3)]], + uint index [[thread_position_in_grid]]) { + constexpr uint qkv_width = 10240u; + constexpr uint controls_width = 6240u; + constexpr uint fused_width = qkv_width + controls_width; + if (index >= args.u[4] * fused_width) return; + const uint row = index / fused_width; + const uint column = index - row * fused_width; + if (column < qkv_width) { + qkv[(ulong)row * qkv_width + column] = fused[index]; + } else { + controls[(ulong)row * controls_width + column - qkv_width] = fused[index]; + } +} + +kernel void kernel_qwen_softmax_precise_512( + constant qwen_kernel_args &args [[buffer(0)]], + device float *out [[buffer(1)]], + device const float *input [[buffer(2)]], + uint row [[threadgroup_position_in_grid]], + uint lid [[thread_position_in_threadgroup]], + uint lane [[thread_index_in_simdgroup]], + uint simd_group [[simdgroup_index_in_threadgroup]]) { + constexpr uint reads = 4u; + input += (ulong)row * args.u[0]; + out += (ulong)row * args.u[0]; + threadgroup float local_max[32]; + threadgroup float local_normalizer[32]; + float values[reads]; + for (uint i = 0; i < reads; i++) values[i] = input[lid * reads + i]; + if (simd_group == 0u) { + local_max[lane] = -INFINITY; + local_normalizer[lane] = 0.0f; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + float maximum = -FLT_MAX; + for (uint i = 0; i < reads; i++) maximum = maximum < values[i] ? values[i] : maximum; + maximum = simd_max(maximum); + if (lane == 0u) local_max[simd_group] = maximum; + threadgroup_barrier(mem_flags::mem_threadgroup); + if (simd_group == 0u) { + maximum = simd_max(local_max[lane]); + if (lane == 0u) local_max[0] = maximum; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + maximum = local_max[0]; + float normalizer = 0.0f; + for (uint i = 0; i < reads; i++) { + values[i] = fast::exp(values[i] - maximum); + normalizer += values[i]; + } + normalizer = simd_sum(normalizer); + if (lane == 0u) local_normalizer[simd_group] = normalizer; + threadgroup_barrier(mem_flags::mem_threadgroup); + if (simd_group == 0u) { + normalizer = simd_sum(local_normalizer[lane]); + if (lane == 0u) local_normalizer[0] = normalizer; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + normalizer = 1.0f / local_normalizer[0]; + for (uint i = 0; i < reads; i++) { + out[lid * reads + i] = (float)(bfloat)(values[i] * normalizer); + } +} + +kernel void kernel_qwen_sparse_probabilities( + constant qwen_kernel_args &args [[buffer(0)]], + device float *out [[buffer(1)]], + device const float *logits [[buffer(2)]], + device const int *selected [[buffer(3)]], + uint row [[threadgroup_position_in_grid]], + uint lid [[thread_position_in_threadgroup]], + uint lane [[thread_index_in_simdgroup]], + uint simd_group [[simdgroup_index_in_threadgroup]]) { + threadgroup float partial_max[32]; + threadgroup float partial_sum[32]; + const uint width = args.u[0]; + const uint top_k = args.u[1]; + const uint threads = args.u[12]; + logits += (ulong)row * width; + selected += (ulong)row * top_k; + out += (ulong)row * top_k; + + if (simd_group == 0u) { + partial_max[lane] = -INFINITY; + partial_sum[lane] = 0.0f; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + float maximum = -INFINITY; + for (uint token = lid; token < width; token += threads) { + maximum = max(maximum, logits[token] * args.f[0]); + } + maximum = simd_max(maximum); + if (lane == 0u) partial_max[simd_group] = maximum; + threadgroup_barrier(mem_flags::mem_threadgroup); + if (simd_group == 0u) { + maximum = simd_max(partial_max[lane]); + if (lane == 0u) partial_max[0] = maximum; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + maximum = partial_max[0]; + + float sum = 0.0f; + for (uint token = lid; token < width; token += threads) { + sum += fast::exp(logits[token] * args.f[0] - maximum); + } + sum = simd_sum(sum); + if (lane == 0u) partial_sum[simd_group] = sum; + threadgroup_barrier(mem_flags::mem_threadgroup); + if (simd_group == 0u) { + sum = simd_sum(partial_sum[lane]); + if (lane == 0u) partial_sum[0] = sum; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (lid < top_k) { + const uint token = (uint)selected[lid]; + const float log_total = maximum + fast::log(partial_sum[0]); + out[lid] = fast::exp(logits[token] * args.f[0] - log_total); + } } kernel void kernel_qwen_route_top10( constant qwen_kernel_args &args [[buffer(0)]], device int *ids [[buffer(1)]], device float *weights [[buffer(2)]], - device const float *logits [[buffer(3)]], - uint gid [[thread_position_in_grid]]) { - if (gid != 0u) return; - float max_value = -INFINITY; - for (uint i = 0; i < args.u[0]; i++) max_value = max(max_value, logits[i]); - float sum = 0.0f; - for (uint i = 0; i < args.u[0]; i++) sum += exp(logits[i] - max_value); - float selected_sum = 0.0f; - for (uint slot = 0; slot < 10u; slot++) { - float best = -1.0f; - int best_id = -1; - for (uint i = 0; i < args.u[0]; i++) { - bool used = false; - for (uint j = 0; j < slot; j++) used = used || ids[j] == (int)i; - const float probability = exp(logits[i] - max_value) / sum; - if (!used && probability > best) { - best = probability; - best_id = (int)i; + device const float *probabilities [[buffer(3)]], + uint row [[threadgroup_position_in_grid]], + uint lane [[thread_index_in_simdgroup]], + uint simd_group [[simdgroup_index_in_threadgroup]]) { + constexpr uint top_k = 10u; + constexpr uint simd_groups = 16u; + constexpr uint local_candidates = simd_groups * top_k; + constexpr uint candidates_per_lane = local_candidates / 32u; + threadgroup float local_probabilities[local_candidates]; + threadgroup int local_indices[local_candidates]; + threadgroup float merged_probabilities[top_k]; + threadgroup int merged_indices[top_k]; + + ids += (ulong)row * 10u; + weights += (ulong)row * 10u; + probabilities += (ulong)row * args.u[0]; + + const int expert = (int)(simd_group * 32u + lane); + float candidate_probability = probabilities[expert]; + for (uint rank = 0; rank < top_k; rank++) { + const float winner_probability = simd_max(candidate_probability); + const int winner_index = (int)simd_max( + candidate_probability == winner_probability ? (float)expert : -1.0f); + if (lane == 0u) { + const uint destination = simd_group * top_k + rank; + local_probabilities[destination] = winner_probability; + local_indices[destination] = winner_index; + } + if (expert == winner_index) candidate_probability = -INFINITY; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + if (simd_group == 0u) { + float candidates[candidates_per_lane]; + int candidate_indices[candidates_per_lane]; + for (uint slot = 0; slot < candidates_per_lane; slot++) { + const uint source = lane + slot * 32u; + candidates[slot] = local_probabilities[source]; + candidate_indices[slot] = local_indices[source]; + } + + for (uint rank = 0; rank < top_k; rank++) { + float lane_probability = candidates[0]; + int lane_index = candidate_indices[0]; + uint lane_slot = 0u; + for (uint slot = 1; slot < candidates_per_lane; slot++) { + if (candidates[slot] > lane_probability || + (candidates[slot] == lane_probability && + candidate_indices[slot] > lane_index)) { + lane_probability = candidates[slot]; + lane_index = candidate_indices[slot]; + lane_slot = slot; + } + } + const float winner_probability = simd_max(lane_probability); + const int winner_index = (int)simd_max( + lane_probability == winner_probability ? (float)lane_index : -1.0f); + if (lane == 0u) { + merged_probabilities[rank] = winner_probability; + merged_indices[rank] = winner_index; + } + if (lane_index == winner_index) candidates[lane_slot] = -INFINITY; + } + + if (lane == 0u) { + float denominator = 0.0f; + for (uint slot = 0; slot < top_k; slot++) { + denominator = qwen_round_bf16( + denominator + merged_probabilities[top_k - 1u - slot]); + } + for (uint slot = 0; slot < top_k; slot++) { + const uint source = top_k - 1u - slot; + ids[slot] = merged_indices[source]; + weights[slot] = qwen_round_bf16( + merged_probabilities[source] / denominator); } } - ids[slot] = best_id; - weights[slot] = best; - selected_sum += best; } - for (uint slot = 0; slot < 10u; slot++) weights[slot] /= selected_sum; } kernel void kernel_qwen_accumulate( @@ -424,7 +2844,13 @@ kernel void kernel_qwen_accumulate_sigmoid_scalar( device const float *x [[buffer(2)]], device const float *gate [[buffer(3)]], uint index [[thread_position_in_grid]]) { - if (index < args.u[0]) out[index] += x[index] / (1.0f + exp(-gate[0])); + const uint hidden = args.u[0]; + if (index < hidden * max(args.u[4], 1u)) { + const uint row = index / hidden; + const float shared = qwen_round_bf16( + x[index] * qwen_sigmoid_bf16(gate[row])); + out[index] = qwen_round_bf16(out[index] + shared); + } } kernel void kernel_qwen_split_q_gate( @@ -435,11 +2861,15 @@ kernel void kernel_qwen_split_q_gate( uint index [[thread_position_in_grid]]) { const uint heads = args.u[0]; const uint dim = args.u[1]; - if (index >= heads * dim) return; - const uint head = index / dim; - const uint column = index % dim; - q[index] = packed[(ulong)head * dim * 2u + column]; - gate[index] = packed[(ulong)head * dim * 2u + dim + column]; + const uint row_width = heads * dim; + if (index >= row_width * max(args.u[4], 1u)) return; + const uint row = index / row_width; + const uint local = index % row_width; + const uint head = local / dim; + const uint column = local % dim; + const ulong packed_base = (ulong)row * row_width * 2u; + q[index] = packed[packed_base + (ulong)head * dim * 2u + column]; + gate[index] = packed[packed_base + (ulong)head * dim * 2u + dim + column]; } kernel void kernel_qwen_head_norm_rope( @@ -447,27 +2877,59 @@ kernel void kernel_qwen_head_norm_rope( device float *out [[buffer(1)]], device const float *x [[buffer(2)]], device const uchar *weight [[buffer(5)]], - uint2 gid [[thread_position_in_grid]]) { - const uint column = gid.x; - const uint head = gid.y; + uint3 group [[threadgroup_position_in_grid]], + uint3 thread_position [[thread_position_in_threadgroup]], + uint lane [[thread_index_in_simdgroup]], + uint simd_group [[simdgroup_index_in_threadgroup]]) { + const uint head = group.x; + const uint lid = thread_position.x; const uint dim = args.u[0]; const uint rotary = args.u[1]; - if (column >= dim || head >= args.u[2]) return; - device const float *row = x + (ulong)head * dim; + if (head >= args.u[2]) return; + const ulong batch_base = (ulong)group.y * args.u[2] * dim; + device const float *row = x + batch_base + (ulong)head * dim; + + threadgroup float scale[1]; + threadgroup float partial[32]; float variance = 0.0f; - for (uint i = 0; i < dim; i++) variance = fma(row[i], row[i], variance); - const float scale = rsqrt(variance / (float)dim + args.f[0]); - float value = row[column] * scale * - (1.0f + qwen_bf16(qwen_weight_u16(weight, args.u[13], column))); - if (column < rotary) { - const uint rotary_half = rotary / 2u; - const uint pair = column < rotary_half ? column + rotary_half : column - rotary_half; - const float paired = row[pair] * scale * - (1.0f + qwen_bf16(qwen_weight_u16(weight, args.u[13], pair))); - const float theta = (float)args.u[3] * pow(args.f[1], -2.0f * (float)(column % rotary_half) / (float)rotary); - value = value * cos(theta) + (column < rotary_half ? -paired : paired) * sin(theta); + float values[4]; + for (uint i = 0; i < 4u; i++) { + const uint column = lid * 4u + i; + values[i] = column < dim ? row[column] : 0.0f; + variance += values[i] * values[i]; + } + variance = simd_sum(variance); + if (simd_group == 0u) partial[lane] = 0.0f; + threadgroup_barrier(mem_flags::mem_threadgroup); + if (lane == 0u) partial[simd_group] = variance; + threadgroup_barrier(mem_flags::mem_threadgroup); + if (simd_group == 0u) { + variance = simd_sum(partial[lane]); + if (lane == 0u) { + scale[0] = metal::precise::rsqrt(variance / (float)dim + args.f[0]); + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint i = 0; i < 4u; i++) { + const uint column = lid * 4u + i; + if (column >= dim) continue; + const float normalized = qwen_round_bf16(values[i] * scale[0]); + float value = args.u[11] != 0u ? normalized : qwen_round_bf16( + normalized * qwen_bf16(qwen_weight_u16(weight, args.u[13], column))); + if (column < rotary) { + const uint rotary_half = rotary / 2u; + const uint pair = column < rotary_half ? column + rotary_half : column - rotary_half; + const float paired_normalized = qwen_round_bf16(row[pair] * scale[0]); + const float paired = args.u[11] != 0u ? paired_normalized : qwen_round_bf16( + paired_normalized * qwen_bf16(qwen_weight_u16(weight, args.u[13], pair))); + const float exponent = 2.0f * (float)(column % rotary_half) / (float)rotary; + const float inv_frequency = 1.0f / pow(args.f[1], exponent); + const float theta = (float)(args.u[3] + group.y) * inv_frequency; + value = qwen_apply_rope( + value, column < rotary_half ? -paired : paired, theta); + } + out[batch_base + (ulong)head * dim + column] = qwen_round_bf16(value); } - out[(ulong)head * dim + column] = value; } kernel void kernel_qwen_store_kv_bf16( @@ -475,12 +2937,15 @@ kernel void kernel_qwen_store_kv_bf16( device ushort *cache [[buffer(1)]], device const float *key [[buffer(2)]], device const float *value [[buffer(3)]], - uint index [[thread_position_in_grid]]) { + uint2 gid [[thread_position_in_grid]]) { const uint width = args.u[0]; + const uint index = gid.x; + const uint row = gid.y; if (index >= width) return; - const ulong base = (ulong)args.u[1] * width * 2u; - cache[base + index] = qwen_to_bf16(key[index]); - cache[base + width + index] = qwen_to_bf16(value[index]); + const ulong base = (ulong)(args.u[1] + row) * width * 2u; + const ulong input = (ulong)row * width + index; + cache[base + index] = qwen_to_bf16(key[input]); + cache[base + width + index] = qwen_to_bf16(value[input]); } kernel void kernel_qwen_qsa_store_raw( @@ -488,13 +2953,17 @@ kernel void kernel_qwen_qsa_store_raw( device ushort *raw [[buffer(1)]], device const float *projected [[buffer(2)]], device float *query [[buffer(3)]], - uint index [[thread_position_in_grid]]) { + uint2 gid [[thread_position_in_grid]]) { const uint dim = args.u[0]; const uint query_width = args.u[1] * dim; + const uint index = gid.x; + const uint row = gid.y; + const ulong projected_base = (ulong)row * (query_width + dim); if (index < query_width) { - query[index] = projected[index]; + query[(ulong)row * query_width + index] = projected[projected_base + index]; } else if (index < query_width + dim) { - raw[(ulong)args.u[2] * dim + index - query_width] = qwen_to_bf16(projected[index]); + raw[(ulong)(args.u[2] + row) * dim + index - query_width] = + qwen_to_bf16(projected[projected_base + index]); } } @@ -523,7 +2992,7 @@ kernel void kernel_qwen_qsa_pool_key( } const float scale = rsqrt(variance / (float)dim + args.f[0]); float value = qwen_bf16(mean_bf16) * scale * - (1.0f + qwen_bf16(qwen_weight_u16(weight, args.u[13], column))); + qwen_bf16(qwen_weight_u16(weight, args.u[13], column)); const uint rotary = 64u; if (column < rotary) { const uint rotary_half = rotary / 2u; @@ -534,10 +3003,11 @@ kernel void kernel_qwen_qsa_pool_key( } paired_mean = qwen_bf16(qwen_to_bf16(paired_mean / (float)args.u[2])); const float paired = paired_mean * scale * - (1.0f + qwen_bf16(qwen_weight_u16(weight, args.u[13], pair))); + qwen_bf16(qwen_weight_u16(weight, args.u[13], pair)); const float theta = (float)args.u[3] * pow(args.f[1], -2.0f * (float)(column % rotary_half) / (float)rotary); - value = value * cos(theta) + (column < rotary_half ? -paired : paired) * sin(theta); + value = qwen_apply_rope( + value, column < rotary_half ? -paired : paired, theta); } pooled[(ulong)args.u[1] * dim + column] = qwen_to_bf16(value); } @@ -579,67 +3049,184 @@ kernel void kernel_qwen_qsa_sort_blocks( } } -static inline uint qwen_qsa_token( - device const int *selected, - uint ordinal, - uint selected_count, - uint ratio, - uint tail_start) { - const uint selected_tokens = selected_count * ratio; - return ordinal < selected_tokens - ? (uint)selected[ordinal / ratio] * ratio + ordinal % ratio - : tail_start + ordinal - selected_tokens; -} - kernel void kernel_qwen_sparse_attention( constant qwen_kernel_args &args [[buffer(0)]], device float *out [[buffer(1)]], device const float *query [[buffer(2)]], device const ushort *cache [[buffer(3)]], device const int *selected [[buffer(4)]], - uint2 gid [[thread_position_in_grid]], - uint lane [[thread_index_in_threadgroup]]) { + uint head [[threadgroup_position_in_grid]], + uint tid [[thread_position_in_threadgroup]], + uint lane [[thread_index_in_simdgroup]]) { const uint heads = args.u[0]; const uint kv_heads = args.u[1]; const uint dim = args.u[2]; - const uint column = gid.x; - const uint head = gid.y; - if (column >= dim || head >= heads) return; + if (head >= heads) return; + const uint simd_group = tid / 32u; const uint kv_head = head / (heads / kv_heads); - const uint tokens = args.u[4] * args.u[5] + args.u[7]; - const float attention_scale = rsqrt((float)dim); - threadgroup float probabilities[2051]; - if (lane == 0u) { - float max_score = -INFINITY; - for (uint ordinal = 0; ordinal < tokens; ordinal++) { - const uint token = qwen_qsa_token(selected, ordinal, args.u[4], args.u[5], args.u[6]); - const ulong base = (ulong)token * kv_heads * dim * 2u + (ulong)kv_head * dim; + constexpr uint simdgroups = 8u; + constexpr uint values_per_lane = 8u; + threadgroup float q[256]; + threadgroup float max_scores[simdgroups]; + threadgroup float sum_exp_scores[simdgroups]; + threadgroup float outputs[simdgroups * 256u]; + const float attention_scale = precise::rsqrt((float)dim); + q[tid] = query[(ulong)head * dim + tid] * attention_scale; + threadgroup_barrier(mem_flags::mem_threadgroup); + + float result[values_per_lane] = {0.0f}; + float max_score = -FLT_MAX; + float sum_exp = 0.0f; + const uint column = lane * values_per_lane; + for (uint block = simd_group; block < args.u[4]; block += simdgroups) { + const uint token_start = (uint)selected[block] * args.u[5]; + for (uint inner = 0u; inner < args.u[5]; inner++) { + const uint token = token_start + inner; + const ulong base = (ulong)token * kv_heads * dim * 2u + + (ulong)kv_head * dim + column; float score = 0.0f; - for (uint i = 0; i < dim; i++) { - score = fma(query[(ulong)head * dim + i], qwen_bf16(cache[base + i]), score); + for (uint i = 0u; i < values_per_lane; i++) { + score += q[column + i] * qwen_bf16(cache[base + i]); } - probabilities[ordinal] = score * attention_scale; - max_score = max(max_score, probabilities[ordinal]); + score = simd_sum(score); + const float new_max = max(max_score, score); + const float factor = exp(max_score - new_max); + const float exp_score = exp(score - new_max); + sum_exp = fma(sum_exp, factor, exp_score); + for (uint i = 0u; i < values_per_lane; i++) { + result[i] = fma( + result[i], factor, + exp_score * qwen_bf16(cache[base + kv_heads * dim + i])); + } + max_score = new_max; } - float denominator = 0.0f; - for (uint ordinal = 0; ordinal < tokens; ordinal++) { - probabilities[ordinal] = exp(probabilities[ordinal] - max_score); - denominator += probabilities[ordinal]; + } + for (uint token = args.u[6] + simd_group; + token < args.u[3]; + token += simdgroups) { + const ulong base = (ulong)token * kv_heads * dim * 2u + + (ulong)kv_head * dim + column; + float score = 0.0f; + for (uint i = 0u; i < values_per_lane; i++) { + score += q[column + i] * qwen_bf16(cache[base + i]); } - for (uint ordinal = 0; ordinal < tokens; ordinal++) { - probabilities[ordinal] /= denominator; + score = simd_sum(score); + const float new_max = max(max_score, score); + const float factor = exp(max_score - new_max); + const float exp_score = exp(score - new_max); + sum_exp = fma(sum_exp, factor, exp_score); + for (uint i = 0u; i < values_per_lane; i++) { + result[i] = fma( + result[i], factor, + exp_score * qwen_bf16(cache[base + kv_heads * dim + i])); } + max_score = new_max; + } + + if (lane == 0u) { + max_scores[simd_group] = max_score; + sum_exp_scores[simd_group] = sum_exp; + } + for (uint i = 0u; i < values_per_lane; i++) { + outputs[simd_group * dim + column + i] = result[i]; } threadgroup_barrier(mem_flags::mem_threadgroup); - float value = 0.0f; - for (uint ordinal = 0; ordinal < tokens; ordinal++) { - const uint token = qwen_qsa_token(selected, ordinal, args.u[4], args.u[5], args.u[6]); - const ulong base = (ulong)token * kv_heads * dim * 2u + (ulong)kv_head * dim; - value = fma(probabilities[ordinal], - qwen_bf16(cache[base + kv_heads * dim + column]), - value); + float merged_max = max_scores[0]; + for (uint group = 1u; group < simdgroups; group++) { + merged_max = max(merged_max, max_scores[group]); + } + float merged_sum = 0.0f; + float merged_value = 0.0f; + for (uint group = 0u; group < simdgroups; group++) { + const float weight = exp(max_scores[group] - merged_max); + merged_sum += sum_exp_scores[group] * weight; + merged_value += outputs[group * dim + tid] * weight; + } + out[(ulong)head * dim + tid] = qwen_round_bf16(merged_value / merged_sum); +} + +kernel void kernel_qwen_dense_attention_masked( + constant qwen_kernel_args &args [[buffer(0)]], + device float *out [[buffer(1)]], + device const float *query [[buffer(2)]], + device const ushort *cache [[buffer(3)]], + device const uchar *mask [[buffer(4)]], + uint2 group [[threadgroup_position_in_grid]], + uint simd_group [[simdgroup_index_in_threadgroup]], + uint lane [[thread_index_in_simdgroup]]) { + const uint head = group.x; + const uint row = group.y; + const uint heads = args.u[0]; + const uint kv_heads = args.u[1]; + const uint dim = args.u[2]; + const uint tokens = args.u[3]; + if (head >= heads) return; + const uint kv_head = head / (heads / kv_heads); + constexpr uint simdgroups = 32u; + constexpr uint values_per_lane = 8u; + threadgroup float outputs[simdgroups * simdgroups]; + threadgroup float max_scores[simdgroups]; + threadgroup float sum_exp_scores[simdgroups]; + float q[values_per_lane]; + float k[values_per_lane]; + float result[values_per_lane] = {0.0f}; + const ulong query_base = + ((ulong)row * heads + head) * dim + lane * values_per_lane; + const float attention_scale = precise::rsqrt((float)dim); + for (uint i = 0; i < values_per_lane; i++) { + q[i] = attention_scale * query[query_base + i]; + } + + float max_score = -FLT_MAX; + float sum_exp = 0.0f; + for (uint token = simd_group; token < tokens; token += simdgroups) { + const bool use_key = mask[(ulong)row * tokens + token] != 0u; + if (use_key) { + const ulong base = (ulong)token * kv_heads * dim * 2u + + (ulong)kv_head * dim + lane * values_per_lane; + for (uint i = 0; i < values_per_lane; i++) { + k[i] = qwen_bf16(cache[base + i]); + } + float score = 0.0f; + for (uint i = 0; i < values_per_lane; i++) { + score = fma(q[i], k[i], score); + } + score = simd_sum(score); + const float new_max = max(max_score, score); + const float factor = fast::exp(max_score - new_max); + const float exp_score = fast::exp(score - new_max); + max_score = new_max; + sum_exp = fma(sum_exp, factor, exp_score); + for (uint i = 0; i < values_per_lane; i++) { + result[i] = fma( + result[i], factor, + exp_score * qwen_bf16(cache[base + kv_heads * dim + i])); + } + } + } + if (lane == 0u) { + max_scores[simd_group] = max_score; + sum_exp_scores[simd_group] = sum_exp; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + max_score = max_scores[lane]; + const float new_max = simd_max(max_score); + const float factor = fast::exp(max_score - new_max); + sum_exp = simd_sum(sum_exp_scores[lane] * factor); + for (uint i = 0; i < values_per_lane; i++) { + outputs[lane * simdgroups + simd_group] = result[i]; + threadgroup_barrier(mem_flags::mem_threadgroup); + result[i] = simd_sum(outputs[simd_group * simdgroups + lane] * factor); + if (sum_exp != 0.0f) result[i] /= sum_exp; + threadgroup_barrier(mem_flags::mem_threadgroup); + } + if (lane == 0u) { + const ulong output_base = + ((ulong)row * heads + head) * dim + simd_group * values_per_lane; + for (uint i = 0; i < values_per_lane; i++) { + out[output_base + i] = qwen_round_bf16(result[i]); + } } - out[(ulong)head * dim + column] = value; } kernel void kernel_qwen_dense_attention( @@ -647,41 +3234,309 @@ kernel void kernel_qwen_dense_attention( device float *out [[buffer(1)]], device const float *query [[buffer(2)]], device const ushort *cache [[buffer(3)]], - uint2 gid [[thread_position_in_grid]], - uint lane [[thread_index_in_threadgroup]]) { + uint head [[threadgroup_position_in_grid]], + uint simd_group [[simdgroup_index_in_threadgroup]], + uint lane [[thread_index_in_simdgroup]]) { const uint heads = args.u[0]; const uint kv_heads = args.u[1]; const uint dim = args.u[2]; const uint tokens = args.u[3]; - const uint column = gid.x; - const uint head = gid.y; - if (column >= dim || head >= heads) return; + if (head >= heads) return; const uint kv_head = head / (heads / kv_heads); - const float attention_scale = rsqrt((float)dim); - threadgroup float probabilities[2048]; + constexpr uint simdgroups = 32u; + constexpr uint values_per_lane = 8u; + threadgroup float outputs[simdgroups * simdgroups]; + threadgroup float max_scores[simdgroups]; + threadgroup float sum_exp_scores[simdgroups]; + float q[values_per_lane]; + float result[values_per_lane] = {0.0f}; + const ulong query_base = (ulong)head * dim + lane * values_per_lane; + const float attention_scale = precise::rsqrt((float)dim); + for (uint i = 0; i < values_per_lane; i++) { + q[i] = attention_scale * query[query_base + i]; + } + + float max_score = -FLT_MAX; + float sum_exp = 0.0f; + for (uint token = simd_group; token < tokens; token += simdgroups) { + const ulong base = (ulong)token * kv_heads * dim * 2u + + (ulong)kv_head * dim + lane * values_per_lane; + float score = 0.0f; + for (uint i = 0; i < values_per_lane; i++) { + score += q[i] * qwen_bf16(cache[base + i]); + } + score = simd_sum(score); + const float new_max = max(max_score, score); + const float factor = fast::exp(max_score - new_max); + const float exp_score = fast::exp(score - new_max); + max_score = new_max; + sum_exp = fma(sum_exp, factor, exp_score); + for (uint i = 0; i < values_per_lane; i++) { + result[i] = fma( + result[i], factor, + exp_score * qwen_bf16(cache[base + kv_heads * dim + i])); + } + } if (lane == 0u) { - float max_score = -INFINITY; - for (uint token = 0; token < tokens; token++) { - const ulong base = (ulong)token * kv_heads * dim * 2u + (ulong)kv_head * dim; - float score = 0.0f; - for (uint i = 0; i < dim; i++) score = fma(query[(ulong)head * dim + i], qwen_bf16(cache[base + i]), score); - probabilities[token] = score * attention_scale; - max_score = max(max_score, probabilities[token]); - } - float denominator = 0.0f; - for (uint token = 0; token < tokens; token++) { - probabilities[token] = exp(probabilities[token] - max_score); - denominator += probabilities[token]; - } - for (uint token = 0; token < tokens; token++) probabilities[token] /= denominator; + max_scores[simd_group] = max_score; + sum_exp_scores[simd_group] = sum_exp; } threadgroup_barrier(mem_flags::mem_threadgroup); - float value = 0.0f; - for (uint token = 0; token < tokens; token++) { - const ulong base = (ulong)token * kv_heads * dim * 2u + (ulong)kv_head * dim; - value = fma(probabilities[token], qwen_bf16(cache[base + kv_heads * dim + column]), value); + max_score = max_scores[lane]; + const float new_max = simd_max(max_score); + const float factor = fast::exp(max_score - new_max); + sum_exp = simd_sum(sum_exp_scores[lane] * factor); + for (uint i = 0; i < values_per_lane; i++) { + outputs[lane * simdgroups + simd_group] = result[i]; + threadgroup_barrier(mem_flags::mem_threadgroup); + result[i] = simd_sum(outputs[simd_group * simdgroups + lane] * factor); + if (sum_exp != 0.0f) result[i] /= sum_exp; + threadgroup_barrier(mem_flags::mem_threadgroup); + } + if (lane == 0u) { + const ulong output_base = (ulong)head * dim + simd_group * values_per_lane; + for (uint i = 0; i < values_per_lane; i++) { + out[output_base + i] = qwen_round_bf16(result[i]); + } + } +} + +// MLX deliberately leaves its fused SDPA path when query rows times the GQA +// factor exceeds 32. These three kernels preserve the fallback's BF16 +// boundaries: scaled Q @ K.t, precise softmax, then probabilities @ V. +kernel void kernel_qwen_attention_fallback_scores( + constant qwen_kernel_args &args [[buffer(0)]], + device ushort *scores [[buffer(1)]], + device const float *query [[buffer(2)]], + device const ushort *cache [[buffer(3)]], + uint2 group [[threadgroup_position_in_grid]], + uint lane [[thread_index_in_simdgroup]]) { + const uint head = group.x; + const uint row = group.y; + const uint heads = args.u[0]; + const uint kv_heads = args.u[1]; + const uint dim = args.u[2]; + const uint tokens = args.u[3]; + if (head >= heads || row >= args.u[4]) return; + const uint kv_head = head / (heads / kv_heads); + const bfloat scale = (bfloat)precise::rsqrt((float)dim); + const ulong query_base = ((ulong)row * heads + head) * dim; + const ulong score_base = ((ulong)head * args.u[4] + row) * tokens; + for (uint token = 0u; token < tokens; token++) { + if (token > args.u[5] + row) { + if (lane == 0u) scores[score_base + token] = 0xff7fu; + continue; + } + const ulong key_base = (ulong)token * kv_heads * dim * 2u + + (ulong)kv_head * dim; + float sum = 0.0f; + for (uint column = lane; column < dim; column += 32u) { + const float q = (float)(bfloat)((bfloat)query[query_base + column] * scale); + sum = fma(q, qwen_bf16(cache[key_base + column]), sum); + } + sum = simd_sum(sum); + if (lane == 0u) scores[score_base + token] = qwen_to_bf16(sum); + } +} + +kernel void kernel_qwen_attention_fallback_softmax( + constant qwen_kernel_args &args [[buffer(0)]], + device ushort *out [[buffer(1)]], + device const ushort *input [[buffer(2)]], + uint group [[threadgroup_position_in_grid]], + uint lid [[thread_position_in_threadgroup]], + uint lane [[thread_index_in_simdgroup]], + uint simd_group [[simdgroup_index_in_threadgroup]]) { + constexpr uint reads = 4u; + threadgroup float local_max[32]; + threadgroup float local_normalizer[32]; + const uint tokens = args.u[3]; + const ulong base = (ulong)group * tokens; + float values[reads]; + for (uint i = 0u; i < reads; i++) { + const uint token = lid * reads + i; + values[i] = token < tokens ? qwen_bf16(input[base + token]) : -INFINITY; + } + if (simd_group == 0u) { + local_max[lane] = -INFINITY; + local_normalizer[lane] = 0.0f; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + float maximum = -FLT_MAX; + for (uint i = 0u; i < reads; i++) maximum = max(maximum, values[i]); + maximum = simd_max(maximum); + if (lane == 0u) local_max[simd_group] = maximum; + threadgroup_barrier(mem_flags::mem_threadgroup); + if (simd_group == 0u) { + maximum = simd_max(local_max[lane]); + if (lane == 0u) local_max[0] = maximum; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + maximum = local_max[0]; + float normalizer = 0.0f; + for (uint i = 0u; i < reads; i++) { + values[i] = fast::exp(values[i] - maximum); + normalizer += values[i]; + } + normalizer = simd_sum(normalizer); + if (lane == 0u) local_normalizer[simd_group] = normalizer; + threadgroup_barrier(mem_flags::mem_threadgroup); + if (simd_group == 0u) { + normalizer = simd_sum(local_normalizer[lane]); + if (lane == 0u) local_normalizer[0] = normalizer; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + normalizer = 1.0f / local_normalizer[0]; + for (uint i = 0u; i < reads; i++) { + const uint token = lid * reads + i; + if (token < tokens) out[base + token] = qwen_to_bf16(values[i] * normalizer); + } +} + +kernel void kernel_qwen_attention_fallback_output( + constant qwen_kernel_args &args [[buffer(0)]], + device float *out [[buffer(1)]], + device const ushort *probabilities [[buffer(2)]], + device const ushort *cache [[buffer(3)]], + uint2 group [[threadgroup_position_in_grid]], + uint lane [[thread_index_in_simdgroup]]) { + const uint head = group.x; + const uint row = group.y; + const uint heads = args.u[0]; + const uint kv_heads = args.u[1]; + const uint dim = args.u[2]; + const uint tokens = args.u[3]; + if (head >= heads || row >= args.u[4]) return; + const uint kv_head = head / (heads / kv_heads); + const ulong probability_base = ((ulong)head * args.u[4] + row) * tokens; + const ulong output_base = ((ulong)row * heads + head) * dim; + for (uint column = lane; column < dim; column += 32u) { + float sum = 0.0f; + for (uint token = 0u; token < tokens; token++) { + const ulong value = (ulong)token * kv_heads * dim * 2u + + (ulong)kv_head * dim + kv_heads * dim + column; + sum = fma(qwen_bf16(probabilities[probability_base + token]), + qwen_bf16(cache[value]), sum); + } + out[output_base + column] = qwen_round_bf16(sum); + } +} + +kernel void kernel_qwen_causal_mask( + constant qwen_kernel_args &args [[buffer(0)]], + device uchar *mask [[buffer(1)]], + uint2 gid [[thread_position_in_grid]]) { + if (gid.x < args.u[0] && gid.y < args.u[4]) { + mask[(ulong)gid.y * args.u[0] + gid.x] = gid.x <= args.u[5] + gid.y; + } +} + +kernel void kernel_qwen_prepare_dense_kv( + constant qwen_kernel_args &args [[buffer(0)]], + device ushort *keys [[buffer(1)]], + device const ushort *cache [[buffer(2)]], + device ushort *values [[buffer(3)]], + uint2 gid [[thread_position_in_grid]]) { + const uint width = args.u[0]; + const uint dim = args.u[1]; + const uint tokens = args.u[3]; + if (gid.x >= width || gid.y >= tokens) return; + const uint head = gid.x / dim; + const uint column = gid.x % dim; + const ulong source = (ulong)gid.y * width * 2u + gid.x; + const ulong target = ((ulong)head * tokens + gid.y) * dim + column; + keys[target] = cache[source]; + values[target] = cache[source + width]; +} + +kernel void kernel_qwen_float_to_bf16( + constant qwen_kernel_args &args [[buffer(0)]], + device ushort *out [[buffer(1)]], + device const float *input [[buffer(2)]], + uint index [[thread_position_in_grid]]) { + if (index < args.u[0]) out[index] = qwen_to_bf16(input[index]); +} + +kernel void kernel_qwen_dense_attention_rows( + constant qwen_kernel_args &args [[buffer(0)]], + device float *out [[buffer(1)]], + device const ushort *queries [[buffer(2)]], + device const ushort *keys [[buffer(3)]], + device const ushort *values [[buffer(4)]], + device const uchar *mask [[buffer(8)]], + uint2 group [[threadgroup_position_in_grid]], + uint simd_group [[simdgroup_index_in_threadgroup]], + uint lane [[thread_index_in_simdgroup]]) { + const uint head = group.x; + const uint row = group.y; + const uint heads = args.u[0]; + const uint kv_heads = args.u[1]; + const uint dim = args.u[2]; + const uint tokens = args.u[3]; + if (head >= heads) return; + const uint kv_head = head / (heads / kv_heads); + constexpr uint simdgroups = 32u; + constexpr uint values_per_lane = 8u; + threadgroup float outputs[simdgroups * simdgroups]; + threadgroup float max_scores[simdgroups]; + threadgroup float sum_exp_scores[simdgroups]; + float q[values_per_lane]; + float k[values_per_lane]; + float result[values_per_lane] = {0.0f}; + queries += ((ulong)row * heads + head) * dim + lane * values_per_lane; + keys += ((ulong)kv_head * tokens + simd_group) * dim + lane * values_per_lane; + values += ((ulong)kv_head * tokens + simd_group) * dim + lane * values_per_lane; + mask += (ulong)row * tokens + simd_group; + const float attention_scale = precise::rsqrt((float)dim); + for (uint i = 0; i < values_per_lane; i++) { + q[i] = attention_scale * qwen_bf16(queries[i]); + } + + float max_score = -FLT_MAX; + float sum_exp = 0.0f; + for (uint token = simd_group; token < tokens; token += simdgroups) { + const bool use_key = mask[0] != 0u; + if (use_key) { + for (uint i = 0; i < values_per_lane; i++) k[i] = qwen_bf16(keys[i]); + float score = 0.0f; + for (uint i = 0; i < values_per_lane; i++) score += q[i] * k[i]; + score = simd_sum(score); + const float new_max = max(max_score, score); + const float factor = fast::exp(max_score - new_max); + const float exp_score = fast::exp(score - new_max); + max_score = new_max; + sum_exp = sum_exp * factor + exp_score; + for (uint i = 0; i < values_per_lane; i++) { + result[i] = result[i] * factor + exp_score * qwen_bf16(values[i]); + } + } + keys += simdgroups * dim; + values += simdgroups * dim; + mask += simdgroups; + } + if (lane == 0u) { + max_scores[simd_group] = max_score; + sum_exp_scores[simd_group] = sum_exp; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + max_score = max_scores[lane]; + const float new_max = simd_max(max_score); + const float factor = fast::exp(max_score - new_max); + sum_exp = simd_sum(sum_exp_scores[lane] * factor); + for (uint i = 0; i < values_per_lane; i++) { + outputs[lane * simdgroups + simd_group] = result[i]; + threadgroup_barrier(mem_flags::mem_threadgroup); + result[i] = simd_sum(outputs[simd_group * simdgroups + lane] * factor); + if (sum_exp != 0.0f) result[i] /= sum_exp; + threadgroup_barrier(mem_flags::mem_threadgroup); + } + if (lane == 0u) { + const ulong output_base = + ((ulong)row * heads + head) * dim + simd_group * values_per_lane; + for (uint i = 0; i < values_per_lane; i++) { + out[output_base + i] = qwen_round_bf16(result[i]); + } } - out[(ulong)head * dim + column] = value; } kernel void kernel_qwen_gate_attention( @@ -690,5 +3545,8 @@ kernel void kernel_qwen_gate_attention( device const float *attention [[buffer(2)]], device const float *gate [[buffer(3)]], uint index [[thread_position_in_grid]]) { - if (index < args.u[0]) out[index] = attention[index] / (1.0f + exp(-gate[index])); + if (index < args.u[0]) { + out[index] = qwen_round_bf16( + attention[index] * qwen_sigmoid_bf16(gate[index])); + } } diff --git a/native/metal/ds4_gpu.h b/native/metal/ds4_gpu.h index f848b9e..be309d1 100644 --- a/native/metal/ds4_gpu.h +++ b/native/metal/ds4_gpu.h @@ -43,8 +43,10 @@ typedef struct { int ds4_gpu_init(void); void ds4_gpu_cleanup(void); +int ds4_gpu_metal4_tensor_api_enabled(void); ds4_gpu_tensor *ds4_gpu_tensor_alloc(uint64_t bytes); +ds4_gpu_tensor *ds4_gpu_tensor_alloc_untracked(uint64_t bytes); ds4_gpu_tensor *ds4_gpu_tensor_alloc_managed(uint64_t bytes); ds4_gpu_tensor *ds4_gpu_tensor_view(const ds4_gpu_tensor *base, uint64_t offset, uint64_t bytes); void ds4_gpu_tensor_free(ds4_gpu_tensor *tensor); @@ -130,6 +132,7 @@ int ds4_gpu_set_aux_model_map_range(const void *model_map, int ds4_gpu_set_model_map_spans(const void *model_map, uint64_t model_size, const uint64_t *offsets, const uint64_t *sizes, uint32_t count, uint64_t max_tensor_bytes); typedef struct { + const ds4_gpu_tensor *tensor; const void *map; uint64_t size; uint64_t offset; @@ -147,6 +150,7 @@ int ds4_gpu_qwen_dispatch( const ds4_gpu_tensor *a, const ds4_gpu_tensor *b, const ds4_gpu_tensor *c, + const ds4_gpu_tensor *d, const ds4_gpu_qwen_weight_view *weights, uint32_t weight_count, const ds4_gpu_qwen_kernel_args *args, @@ -193,6 +197,7 @@ int ds4_gpu_preload_q4_expert_tables(const void *model_map, uint64_t model_size, int ds4_gpu_should_use_managed_kv_cache(uint64_t kv_cache_bytes, uint64_t context_bytes); void ds4_gpu_set_quality(bool quality); void ds4_gpu_set_glm_model(bool enabled); +void ds4_gpu_set_qwen_model(bool enabled); void ds4_gpu_set_ssd_streaming(bool enabled); void ds4_gpu_set_glm_streaming_prefill_full_layer(bool enabled); #ifdef __APPLE__ diff --git a/native/metal/ds4_metal.m b/native/metal/ds4_metal.m index 36ce2ad..40e325c 100644 --- a/native/metal/ds4_metal.m +++ b/native/metal/ds4_metal.m @@ -50,6 +50,7 @@ enum { static id g_device; static id g_queue; static id g_library; +static id g_qwen_library; static id g_batch_cb; static id g_batch_enc; static BOOL g_batch_encoder_concurrent; @@ -501,6 +502,7 @@ static int g_glm_stream_expert_addr_table_building; static uint64_t g_model_residency_count; static int g_model_residency_added_to_queue; static int g_glm_model_mode; +static int g_qwen_model_mode; static int g_ssd_streaming_mode; static int g_glm_streaming_prefill_full_layer_runtime; static int g_metal4_runtime_available; @@ -1255,7 +1257,7 @@ static int ds4_gpu_scratch_needs_cpu_access(const char *label) { static MTLResourceOptions ds4_gpu_model_resource_options(void) { MTLResourceOptions options = MTLResourceStorageModeShared; - if (getenv("DS4_METAL_MODEL_UNTRACKED") != NULL) { + if (g_qwen_model_mode || getenv("DS4_METAL_MODEL_UNTRACKED") != NULL) { options |= MTLResourceHazardTrackingModeUntracked; } return options; @@ -2238,7 +2240,10 @@ static id ds4_gpu_get_pipeline( NSError *error = nil; NSString *name = [NSString stringWithUTF8String:function_name]; - id fn = [g_library newFunctionWithName:name]; + id library = strncmp(function_name, "kernel_qwen_", 12) == 0 + ? g_qwen_library + : g_library; + id fn = [library newFunctionWithName:name]; if (!fn) { fprintf(stderr, "ds4: Metal %s function not found\n", function_name); return nil; @@ -2369,6 +2374,10 @@ static int ds4_gpu_mpp_available(void) { return g_metal4_tensor_api_enabled && !g_quality_mode; } +int ds4_gpu_metal4_tensor_api_enabled(void) { + return g_metal4_tensor_api_enabled; +} + /* * Retained Metal4 defaults live here instead of behind user-visible options. * The public runtime has one automatic accelerated path plus the global @@ -4228,6 +4237,10 @@ void ds4_gpu_set_glm_model(bool enabled) { g_glm_model_mode = enabled ? 1 : 0; } +void ds4_gpu_set_qwen_model(bool enabled) { + g_qwen_model_mode = enabled ? 1 : 0; +} + void ds4_gpu_set_ssd_streaming(bool enabled) { g_ssd_streaming_mode = enabled ? 1 : 0; ds4_gpu_stream_expert_cache_clear_all(1); @@ -4359,7 +4372,6 @@ static NSString *ds4_gpu_full_source(void) { @[@"DS4_METAL_GLM53_BF16_SOURCE", @"metal/glm53_bf16.metal"], @[@"DS4_METAL_GLM53_VISION_SOURCE", @"metal/glm53_vision.metal"], @[@"DS4_METAL_GLM53_KDA_SOURCE", @"metal/glm53_kda.metal"], - @[@"DS4_METAL_QWEN38_SOURCE", @"metal/qwen38.metal"], @[@"DS4_METAL_MOE_SOURCE", @"metal/moe.metal"], @[@"DS4_METAL_DSV4_HC_SOURCE", @"metal/dsv4_hc.metal"], @[@"DS4_METAL_UNARY_SOURCE", @"metal/unary.metal"], @@ -4418,6 +4430,31 @@ static NSString *ds4_gpu_full_source(void) { return source; } +static NSString *ds4_gpu_qwen_source(void) { + const char *override_path = getenv("DS4_METAL_QWEN38_SOURCE"); + NSMutableArray *paths = [NSMutableArray array]; + if (override_path && override_path[0]) { + [paths addObject:[NSString stringWithUTF8String:override_path]]; + } + [paths addObject:@"metal/qwen38.metal"]; + [paths addObject:@"./metal/qwen38.metal"]; + NSFileManager *fm = [NSFileManager defaultManager]; + for (NSString *path in paths) { + if (![fm fileExistsAtPath:path]) continue; + NSError *error = nil; + NSString *source = [NSString stringWithContentsOfFile:path + encoding:NSUTF8StringEncoding + error:&error]; + if (source) return source; + fprintf(stderr, "ds4: failed to read Metal source %s: %s\n", + [path UTF8String], [[error localizedDescription] UTF8String]); + return nil; + } + fprintf(stderr, + "ds4: Metal source metal/qwen38.metal not found (set DS4_METAL_QWEN38_SOURCE to override)\n"); + return nil; +} + typedef struct { int32_t ne00t; int32_t ne00; @@ -6507,6 +6544,37 @@ int ds4_gpu_init(void) { } g_library = library; + NSString *qwen_source = ds4_gpu_qwen_source(); + if (!qwen_source) { + g_library = nil; + g_queue = nil; + g_device = nil; + return 0; + } + MTLCompileOptions *qwen_options = [MTLCompileOptions new]; + if (@available(macOS 15.0, *)) { + qwen_options.mathMode = MTLMathModeSafe; + } else { +#pragma clang diagnostic push +#pragma clang diagnostic ignored "-Wdeprecated-declarations" + qwen_options.fastMathEnabled = NO; +#pragma clang diagnostic pop + } + qwen_options.preprocessorMacros = macros; + error = nil; + id qwen_library = [g_device newLibraryWithSource:qwen_source + options:qwen_options + error:&error]; + if (!qwen_library) { + fprintf(stderr, "ds4: Qwen Metal shader compilation failed: %s\n", + [[error localizedDescription] UTF8String]); + g_library = nil; + g_queue = nil; + g_device = nil; + return 0; + } + g_qwen_library = qwen_library; + id fn = [library newFunctionWithName:@"kernel_get_rows_f32"]; if (!fn) { fprintf(stderr, "ds4: Metal kernel_get_rows_f32 function not found\n"); @@ -8744,14 +8812,16 @@ void ds4_gpu_test_set_flags(uint32_t flags) { g_test_flags = flags; } -ds4_gpu_tensor *ds4_gpu_tensor_alloc(uint64_t bytes) { +static ds4_gpu_tensor *ds4_gpu_tensor_alloc_with_options( + uint64_t bytes, + MTLResourceOptions options) { if (!g_initialized && !ds4_gpu_init()) return NULL; if (bytes == 0 || bytes > (uint64_t)NSUIntegerMax) return NULL; @autoreleasepool { DS4MetalTensor *tensor = [DS4MetalTensor new]; tensor.buffer = [g_device newBufferWithLength:(NSUInteger)bytes - options:MTLResourceStorageModeShared]; + options:options]; if (!tensor.buffer) { return NULL; } @@ -8783,6 +8853,16 @@ ds4_gpu_tensor *ds4_gpu_tensor_alloc(uint64_t bytes) { } } +ds4_gpu_tensor *ds4_gpu_tensor_alloc(uint64_t bytes) { + return ds4_gpu_tensor_alloc_with_options(bytes, MTLResourceStorageModeShared); +} + +ds4_gpu_tensor *ds4_gpu_tensor_alloc_untracked(uint64_t bytes) { + return ds4_gpu_tensor_alloc_with_options( + bytes, + MTLResourceStorageModeShared | MTLResourceHazardTrackingModeUntracked); +} + ds4_gpu_tensor *ds4_gpu_tensor_alloc_managed(uint64_t bytes) { return ds4_gpu_tensor_alloc(bytes); } @@ -10528,6 +10608,7 @@ void ds4_gpu_cleanup(void) { g_model_buffer_cache = nil; g_transient_buffers = nil; g_pending_cbs = nil; + g_qwen_library = nil; g_library = nil; g_queue = nil; g_device = nil; @@ -11744,13 +11825,14 @@ int ds4_gpu_qwen_dispatch( const ds4_gpu_tensor *a, const ds4_gpu_tensor *b, const ds4_gpu_tensor *c, + const ds4_gpu_tensor *d, const ds4_gpu_qwen_weight_view *weights, uint32_t weight_count, const ds4_gpu_qwen_kernel_args *args, uint32_t grid_x, uint32_t grid_y) { if (!kernel || !out || !args || grid_x == 0 || grid_y == 0 || - weight_count > 3 || (weight_count != 0 && !weights)) { + weight_count > 6 || (weight_count != 0 && !weights)) { return 0; } id pipeline = ds4_gpu_get_pipeline(kernel); @@ -11775,14 +11857,34 @@ int ds4_gpu_qwen_dispatch( [enc setBuffer:tensors[i].buffer offset:(NSUInteger)tensors[i].offset atIndex:1 + i]; } } + if (d) { + const DS4MetalTensor *tensor = ds4_gpu_tensor_const_obj(d); + [enc setBuffer:tensor.buffer offset:(NSUInteger)tensor.offset atIndex:8]; + } for (uint32_t i = 0; i < weight_count; i++) { uint64_t inner = 0; - id weight = ds4_gpu_wrap_model_exact_range( - weights[i].map, - weights[i].size, - weights[i].offset, - weights[i].bytes, - &inner); + id weight = nil; + if (weights[i].tensor) { + const DS4MetalTensor *tensor = ds4_gpu_tensor_const_obj(weights[i].tensor); + weight = tensor.buffer; + inner = tensor.offset; + } else if (!g_qwen_model_mode && + ds4_gpu_model_views_cover_range(weights[i].map, + weights[i].size, + weights[i].offset, + weights[i].bytes)) { + weight = ds4_gpu_wrap_model_range(weights[i].map, + weights[i].size, + weights[i].offset, + weights[i].bytes, + &inner); + } else { + weight = ds4_gpu_wrap_model_exact_range(weights[i].map, + weights[i].size, + weights[i].offset, + weights[i].bytes, + &inner); + } if (!weight) { ds4_gpu_end_compute_encoder(cb, enc); if (owned) [cb commit]; @@ -11791,16 +11893,34 @@ int ds4_gpu_qwen_dispatch( /* Metal resource offsets are four-byte aligned. Preserve an arbitrary * safetensors data offset for the Qwen kernels to decode explicitly. */ const uint64_t aligned_inner = inner & ~3ull; - bound_args.u[13 + i] = (uint32_t)(inner - aligned_inner); - [enc setBuffer:weight offset:(NSUInteger)aligned_inner atIndex:5 + i]; + static const uint32_t offset_slots[6] = {13, 14, 15, 9, 6, 7}; + static const NSUInteger buffer_slots[6] = {5, 6, 7, 9, 10, 11}; + const uint32_t offset_slot = offset_slots[i]; + const NSUInteger buffer_slot = buffer_slots[i]; + bound_args.u[offset_slot] = (uint32_t)(inner - aligned_inner); + [enc setBuffer:weight offset:(NSUInteger)aligned_inner atIndex:buffer_slot]; } [enc setBytes:&bound_args length:sizeof(bound_args) atIndex:0]; + if (bound_args.u[10] != 0) { + [enc setThreadgroupMemoryLength:bound_args.u[10] atIndex:0]; + } const NSUInteger width = pipeline.threadExecutionWidth; const NSUInteger max_threads = pipeline.maxTotalThreadsPerThreadgroup; - const NSUInteger threads = MIN(MAX(width, 1u), max_threads); - [enc dispatchThreads:MTLSizeMake(grid_x, grid_y, 1) - threadsPerThreadgroup:MTLSizeMake(threads, 1, 1)]; + const NSUInteger requested_threads = bound_args.u[12]; + if (requested_threads != 0) { + if (requested_threads > max_threads || requested_threads % width != 0) { + ds4_gpu_end_compute_encoder(cb, enc); + if (owned) [cb commit]; + return 0; + } + [enc dispatchThreadgroups:MTLSizeMake(grid_x, grid_y, 1) + threadsPerThreadgroup:MTLSizeMake(requested_threads, 1, 1)]; + } else { + const NSUInteger threads = MIN(MAX(width, 1u), max_threads); + [enc dispatchThreads:MTLSizeMake(grid_x, grid_y, 1) + threadsPerThreadgroup:MTLSizeMake(threads, 1, 1)]; + } ds4_gpu_end_compute_encoder(cb, enc); return owned ? ds4_gpu_finish_command_buffer(cb, 1, kernel) : 1; } diff --git a/src/config.rs b/src/config.rs index 81d69aa..2f944dd 100644 --- a/src/config.rs +++ b/src/config.rs @@ -24,6 +24,7 @@ impl ModelPreferences { Self { reasoning_mode: model.default_reasoning_mode(), speculative: SpeculativePreferences { + glm_mtp: model == ModelChoice::Qwen38FlashNext, dspark_enabled: model == ModelChoice::DeepSeekV4Flash0731, ..SpeculativePreferences::default() }, @@ -308,7 +309,13 @@ impl Config { .map_err(|error| format!("Could not read {}: {error}", path.display()))?; drop_legacy_model_settings(&mut value); let migrated = migrate_deprecated_flash(&mut value); - let dspark_explicit = deepseek_0731_dspark_is_explicit(&value); + let dspark_explicit = model_speculative_field_is_explicit( + &value, + ModelChoice::DeepSeekV4Flash0731, + "dspark_enabled", + ); + let qwen_mtp_explicit = + model_speculative_field_is_explicit(&value, ModelChoice::Qwen38FlashNext, "glm_mtp"); let explicit_reasoning = MODEL_CHOICES.map(|model| (model, model_reasoning_mode_is_explicit(&value, model))); let mut config: Self = serde_norway::from_value(value) @@ -331,6 +338,14 @@ impl Config { .speculative .dspark_enabled = true; } + if !qwen_mtp_explicit { + config + .model_profiles + .get_mut(&ModelChoice::Qwen38FlashNext) + .expect("Qwen model profile is present") + .speculative + .glm_mtp = true; + } config.validate()?; if migrated { config.save(path)?; @@ -446,19 +461,17 @@ fn model_reasoning_mode_is_explicit(value: &Value, model: ModelChoice) -> bool { .is_some_and(|profile| profile.contains_key(Value::String("reasoning_mode".into()))) } -fn deepseek_0731_dspark_is_explicit(value: &Value) -> bool { +fn model_speculative_field_is_explicit(value: &Value, model: ModelChoice, field: &str) -> bool { let Value::Mapping(root) = value else { return false; }; root.get(Value::String("model_profiles".into())) .and_then(Value::as_mapping) - .and_then(|profiles| { - profiles.get(Value::String(ModelChoice::DeepSeekV4Flash0731.id().into())) - }) + .and_then(|profiles| profiles.get(Value::String(model.id().into()))) .and_then(Value::as_mapping) .and_then(|profile| profile.get(Value::String("speculative".into()))) .and_then(Value::as_mapping) - .is_some_and(|section| section.contains_key(Value::String("dspark_enabled".into()))) + .is_some_and(|section| section.contains_key(Value::String(field.into()))) } fn drop_legacy_model_settings(value: &mut Value) { @@ -603,6 +616,28 @@ mod tests { ReasoningMode::Medium ); assert_eq!(loaded.active_generation().context_tokens, 131_072); + assert!( + loaded + .runtime_for(ModelChoice::Qwen38FlashNext) + .speculative + .glm_mtp + ); + + let mut opted_out = loaded; + opted_out + .model_profiles + .get_mut(&ModelChoice::Qwen38FlashNext) + .unwrap() + .speculative + .glm_mtp = false; + opted_out.save(&path).unwrap(); + assert!( + !Config::load(&path) + .unwrap() + .runtime_for(ModelChoice::Qwen38FlashNext) + .speculative + .glm_mtp + ); fs::remove_dir_all(directory).unwrap(); } diff --git a/src/engine.rs b/src/engine.rs index 6927349..acc0d8d 100644 --- a/src/engine.rs +++ b/src/engine.rs @@ -792,6 +792,9 @@ impl Generator { let result = self.generate_inner(messages, settings, cancelled, &mut emit, &mut progress); self.publish_execution_stats(); let (mut output, prompt_complete) = result?; + if prompt_complete { + self.executor.finalize_pending_sampled_token()?; + } let mut completed = messages.to_vec(); completed.push(output.message.clone()); output.previous_checkpoint_bytes = @@ -870,6 +873,9 @@ impl Generator { let result = self.generate_inner(messages, settings, cancelled, &mut emit, &mut progress); self.publish_execution_stats(); let (mut output, prompt_complete) = result?; + if prompt_complete { + self.executor.finalize_pending_sampled_token()?; + } let mut completed = messages.to_vec(); completed.push(output.message.clone()); output.previous_checkpoint_bytes = previous_checkpoint @@ -1462,7 +1468,12 @@ impl Generator { self.executor.set_vision_overlays(overlays)?; self.metrics.kv_prefix_reused(reused); progress(self.executor.position(), self.executor.context(), None); - let mut rng = Rng::new(settings.seed.unwrap_or(0x4453_3453_4552_5645)); + let seed = settings.seed.unwrap_or(0x4453_3453_4552_5645); + let mut rng = if matches!(self.executor.model(), ModelRef::Qwen(_)) { + Rng::new_qwen(seed) + } else { + Rng::new(seed) + }; let mut reasoning = settings.reasoning_mode != ReasoningMode::Direct; let mut generated = ChatTurn { user: false, @@ -1525,14 +1536,32 @@ impl Generator { true, )); } - let token = sample( - self.executor.logits(), - settings.temperature, - settings.top_p, - settings.min_p, - settings.top_k, - &mut rng, - ); + let token = self + .executor + .take_pending_sampled_token() + .unwrap_or_else(|| { + if matches!(self.executor.model(), ModelRef::Qwen(_)) { + sample_probabilities_f64( + &qwen_sampling_probabilities( + self.executor.logits(), + settings.temperature, + settings.top_p, + settings.min_p, + settings.top_k, + ), + &mut rng, + ) + } else { + sample( + self.executor.logits(), + settings.temperature, + settings.top_p, + settings.min_p, + settings.top_k, + &mut rng, + ) + } + }); if self .executor .model() @@ -1903,6 +1932,123 @@ fn sample( sample_probabilities(&probabilities, rng, None) } +#[cfg(any(target_os = "macos", test))] +fn qwen_sampling_probabilities( + logits: &[f32], + temperature: f32, + top_p: f32, + min_p: f32, + top_k: i32, +) -> Vec<(usize, f64)> { + let greedy = || { + vec![( + logits + .iter() + .enumerate() + .max_by(|a, b| a.1.total_cmp(b.1)) + .map_or(0, |(index, _)| index), + 1.0, + )] + }; + if temperature <= 0.0 { + return greedy(); + } + let maximum = logits + .iter() + .copied() + .filter(|value| value.is_finite()) + .fold(f32::NEG_INFINITY, f32::max); + if !maximum.is_finite() { + return greedy(); + } + let top_p = if top_p <= 0.0 || top_p > 1.0 { + 1.0 + } else { + f64::from(top_p) + }; + let rank = |a: &(usize, f64), b: &(usize, f64)| b.1.total_cmp(&a.1).then_with(|| a.0.cmp(&b.0)); + let inverse_temperature = temperature.recip(); + let requested = usize::try_from(top_k.max(0)).unwrap_or(usize::MAX); + let mut probabilities = Vec::with_capacity(if requested > 0 { + requested.min(logits.len()) + } else { + logits.len() + }); + let mut total = 0.0_f32; + for (index, logit) in logits.iter().copied().enumerate() { + if !logit.is_finite() { + continue; + } + let probability = ((logit - maximum) * inverse_temperature).exp(); + total += probability; + let candidate = (index, f64::from(probability)); + if requested == 0 { + probabilities.push(candidate); + continue; + } + let slot = probabilities.partition_point(|current| { + current.1 > candidate.1 || (current.1 == candidate.1 && current.0 < candidate.0) + }); + if slot < requested { + probabilities.insert(slot, candidate); + if probabilities.len() > requested { + probabilities.pop(); + } + } + } + if probabilities.is_empty() || !total.is_finite() || total <= 0.0 { + return greedy(); + } + for (_, probability) in &mut probabilities { + *probability = f64::from((*probability as f32) / total); + } + let count = probabilities.len(); + if requested == 0 { + probabilities.sort_unstable_by(rank); + } + let minimum = f64::from(min_p.max(0.0)); + let mut kept = 0.0; + let mut end = 0; + for (_, probability) in probabilities.iter().take(count) { + if end > 0 && *probability < minimum * probabilities[0].1 { + break; + } + if kept >= top_p { + break; + } + kept += *probability; + end += 1; + } + probabilities.truncate(end); + if probabilities.is_empty() || !kept.is_finite() || kept <= 0.0 { + return greedy(); + } + for (_, probability) in &mut probabilities { + *probability /= kept; + } + probabilities.sort_unstable_by_key(|(token, _)| *token); + probabilities +} + +#[cfg(any(target_os = "macos", test))] +fn sample_probabilities_f64(probabilities: &[(usize, f64)], rng: &mut Rng) -> i32 { + let mut choice = rng.unit_f64() * probabilities.iter().map(|(_, p)| p).sum::(); + for (token, probability) in probabilities { + if !probability.is_finite() || *probability <= 0.0 { + continue; + } + choice -= probability; + if choice <= 0.0 { + return *token as i32; + } + } + probabilities + .iter() + .rev() + .find(|(_, probability)| probability.is_finite() && *probability > 0.0) + .map_or(0, |(token, _)| *token as i32) +} + #[cfg(any(target_os = "macos", test))] fn sampling_probabilities( logits: &[f32], @@ -1938,20 +2084,29 @@ fn sampling_probabilities( top_p }; let min_p = min_p.max(0.0); - let mut probabilities: Vec<(usize, f32)> = logits + let rank = |a: &(usize, f32), b: &(usize, f32)| b.1.total_cmp(&a.1).then_with(|| a.0.cmp(&b.0)); + let mut candidates = logits .iter() + .copied() .enumerate() .filter(|(_, logit)| logit.is_finite()) - .map(|(index, logit)| (index, ((*logit - maximum) / temperature).exp())) - .collect(); + .collect::>(); + if top_k > 0 { + let count = candidates.len().min((top_k as usize).min(1024)); + if count < candidates.len() { + candidates.select_nth_unstable_by(count, rank); + candidates.truncate(count); + } + } + let mut probabilities = candidates + .into_iter() + .map(|(index, logit)| (index, ((logit - maximum) / temperature).exp())) + .collect::>(); if probabilities.is_empty() { return greedy(); } if top_p < 1.0 || top_k > 0 || min_p > 0.0 { - probabilities.sort_unstable_by(|a, b| b.1.total_cmp(&a.1).then_with(|| a.0.cmp(&b.0))); - if top_k > 0 { - probabilities.truncate(probabilities.len().min((top_k as usize).min(1024))); - } + probabilities.sort_unstable_by(rank); } let total: f32 = probabilities .iter() @@ -2037,7 +2192,7 @@ fn exact_delta_sample( ) } -#[cfg(any(target_os = "macos", test))] +#[cfg(test)] fn sample_from_logits( logits: &[f32], temperature: f32, @@ -2046,14 +2201,13 @@ fn sample_from_logits( top_k: i32, rng: &mut Rng, ) -> i32 { - sample_probabilities( - &sampling_probabilities(logits, temperature, top_p, min_p, top_k), + sample_probabilities_f64( + &qwen_sampling_probabilities(logits, temperature, top_p, min_p, top_k), rng, - None, ) } -#[cfg(any(target_os = "macos", test))] +#[cfg(test)] #[allow(clippy::too_many_arguments)] fn exact_speculative_sample( target_logits: &[f32], @@ -2065,19 +2219,53 @@ fn exact_speculative_sample( top_k: i32, rng: &mut Rng, ) -> (i32, bool) { - let mut target = sampling_probabilities(target_logits, temperature, top_p, min_p, top_k); - let mut proposal = sampling_probabilities(draft_logits, temperature, top_p, min_p, top_k); - target.sort_unstable_by_key(|(token, _)| *token); - proposal.sort_unstable_by_key(|(token, _)| *token); - let probability = |values: &[(usize, f32)], token: usize| { + let proposal = qwen_sampling_probabilities(draft_logits, temperature, top_p, min_p, top_k); + exact_speculative_sample_with_probabilities( + target_logits, + &proposal, + draft, + temperature, + top_p, + min_p, + top_k, + rng, + ) +} + +#[cfg(test)] +#[allow(clippy::too_many_arguments)] +fn exact_speculative_sample_with_probabilities( + target_logits: &[f32], + proposal: &[(usize, f64)], + draft: i32, + temperature: f32, + top_p: f32, + min_p: f32, + top_k: i32, + rng: &mut Rng, +) -> (i32, bool) { + let target = qwen_sampling_probabilities(target_logits, temperature, top_p, min_p, top_k); + exact_speculative_sample_from_probabilities(&target, proposal, draft, rng) +} + +#[cfg(any(target_os = "macos", test))] +fn exact_speculative_sample_from_probabilities( + target: &[(usize, f64)], + proposal: &[(usize, f64)], + draft: i32, + rng: &mut Rng, +) -> (i32, bool) { + let probability = |values: &[(usize, f64)], token: usize| { values .binary_search_by_key(&token, |(candidate, _)| *candidate) .ok() .map_or(0.0, |index| values[index].1) }; - let target_probability = probability(&target, draft as usize); - let draft_probability = probability(&proposal, draft as usize); - if draft_probability > 0.0 && rng.unit() <= (target_probability / draft_probability).min(1.0) { + let target_probability = probability(target, draft as usize); + let draft_probability = probability(proposal, draft as usize); + if draft_probability > 0.0 + && rng.unit_f64() <= (target_probability / draft_probability).min(1.0) + { return (draft, true); } @@ -2086,8 +2274,8 @@ fn exact_speculative_sample( .max_by(|left, right| left.1.total_cmp(&right.1)) .map_or(draft, |(token, _)| *token as i32); let mut residual = Vec::with_capacity(target.len()); - for (token, target_probability) in target { - let remaining = target_probability - probability(&proposal, token); + for &(token, target_probability) in target { + let remaining = target_probability - probability(proposal, token); if remaining > 0.0 { residual.push((token, remaining)); } @@ -2095,27 +2283,113 @@ fn exact_speculative_sample( if residual.is_empty() { return (target_best, false); } - (sample_probabilities(&residual, rng, None), false) + (sample_probabilities_f64(&residual, rng), false) } #[cfg(any(target_os = "macos", test))] -struct Rng(u64); +enum Rng { + Ds4(u64), + Qwen { state: u128, increment: u128 }, +} #[cfg(any(target_os = "macos", test))] impl Rng { fn new(seed: u64) -> Self { - Self(seed.max(1)) + Self::Ds4(seed.max(1)) + } + + fn new_qwen(seed: u64) -> Self { + const INIT_A: u32 = 0x43b0_d7e5; + const MULT_A: u32 = 0x931e_8875; + const INIT_B: u32 = 0x8b51_f9dd; + const MULT_B: u32 = 0x58f3_8ded; + const MIX_MULT_L: u32 = 0xca01_f9dd; + const MIX_MULT_R: u32 = 0x4973_f715; + + fn hash_mix(value: u32, hash: &mut u32) -> u32 { + let mut value = value ^ *hash; + *hash = hash.wrapping_mul(MULT_A); + value = value.wrapping_mul(*hash); + value ^ (value >> 16) + } + + fn mix(left: u32, right: u32) -> u32 { + let value = MIX_MULT_L + .wrapping_mul(left) + .wrapping_sub(MIX_MULT_R.wrapping_mul(right)); + value ^ (value >> 16) + } + + let entropy = [seed as u32, (seed >> 32) as u32]; + let mut pool = [0_u32; 4]; + let mut hash = INIT_A; + for (index, slot) in pool.iter_mut().enumerate() { + *slot = hash_mix(entropy.get(index).copied().unwrap_or(0), &mut hash); + } + for source in 0..pool.len() { + for destination in 0..pool.len() { + if source != destination { + pool[destination] = mix(pool[destination], hash_mix(pool[source], &mut hash)); + } + } + } + + let mut words = [0_u32; 8]; + hash = INIT_B; + for (index, word) in words.iter_mut().enumerate() { + let mut value = pool[index % pool.len()] ^ hash; + hash = hash.wrapping_mul(MULT_B); + value = value.wrapping_mul(hash); + *word = value ^ (value >> 16); + } + let word64 = |index| u64::from(words[index]) | (u64::from(words[index + 1]) << 32); + let initial_state = u128::from(word64(2)) | (u128::from(word64(0)) << 64); + let stream = u128::from(word64(6)) | (u128::from(word64(4)) << 64); + let increment = stream.wrapping_mul(2) | 1; + let mut rng = Self::Qwen { + state: 0, + increment, + }; + rng.advance_qwen(); + if let Self::Qwen { state, .. } = &mut rng { + *state = state.wrapping_add(initial_state); + } + rng.advance_qwen(); + rng + } + + fn advance_qwen(&mut self) { + const MULTIPLIER: u128 = 47_026_247_687_942_121_848_144_207_491_837_523_525; + if let Self::Qwen { state, increment } = self { + *state = state.wrapping_mul(MULTIPLIER).wrapping_add(*increment); + } + } + + fn unit_f64(&mut self) -> f64 { + if matches!(self, Self::Qwen { .. }) { + self.advance_qwen(); + let Self::Qwen { state, .. } = self else { + unreachable!() + }; + let mixed = ((*state >> 64) as u64) ^ (*state as u64); + let raw = mixed.rotate_right((*state >> 122) as u32); + return (raw >> 11) as f64 / 9_007_199_254_740_992.0; + } + f64::from(self.unit()) } fn unit(&mut self) -> f32 { - let mut value = self.0; + let Self::Ds4(state) = self else { + return self.unit_f64() as f32; + }; + let mut value = *state; if value == 0 { value = 0x9e37_79b9_7f4a_7c15; } value ^= value >> 12; value ^= value << 25; value ^= value >> 27; - self.0 = value; + *state = value; let value = value.wrapping_mul(0x2545_f491_4f6c_dd1d); ((value >> 40) & 0xff_ffff) as f32 / 16_777_216.0 } @@ -2125,6 +2399,21 @@ impl Rng { mod sampling_tests { use super::*; + #[test] + fn qwen_rng_matches_numpy_pcg64() { + let mut rng = Rng::new_qwen(12_345); + let expected = [ + 0.227_336_022_467_169_66, + 0.316_758_339_709_752_87, + 0.797_365_457_332_734_1, + 0.676_254_670_750_974_6, + 0.391_109_550_601_909, + ]; + for expected in expected { + assert_eq!(rng.unit_f64(), expected); + } + } + #[test] fn checkpoint_headers_identify_the_exact_model() { let path = @@ -2229,6 +2518,23 @@ mod sampling_tests { ); } + #[test] + fn qwen_sampling_uses_mtplx_global_nucleus_mass_before_top_k() { + let probabilities = qwen_sampling_probabilities(&[0.0, 0.0, 0.0], 1.0, 0.4, 0.0, 2); + assert_eq!( + probabilities + .iter() + .map(|(token, _)| *token) + .collect::>(), + [0, 1] + ); + assert!( + probabilities + .iter() + .all(|(_, probability)| (*probability - 0.5).abs() < f64::EPSILON) + ); + } + #[test] fn split_utf8_token_bytes_are_joined_before_decoding() { let mut generated = ChatTurn { diff --git a/src/engine/metal.rs b/src/engine/metal.rs index 9e8649c..eec5623 100644 --- a/src/engine/metal.rs +++ b/src/engine/metal.rs @@ -13,10 +13,13 @@ use profile::ExpertProfile; use qwen::QwenExecutor; use super::gguf::{BF16, F16, F32, Gguf, IQ2_XXS, MXFP4, Q2_K, Q4_K, Q8_0, Tensor as GgufTensor}; +#[cfg(test)] +use super::sample_from_logits; use super::validation::{DsparkConfig, SupportKind, dspark_config}; use super::{ - LoadedModel, Model, ModelFamily, ModelRef, Rng, exact_delta_sample, exact_speculative_sample, - sample_from_logits, + LoadedModel, Model, ModelFamily, ModelRef, Rng, exact_delta_sample, + exact_speculative_sample_from_probabilities, qwen_sampling_probabilities, + sample_probabilities_f64, }; use crate::model::ModelChoice; use crate::settings::{ @@ -4517,6 +4520,20 @@ impl Executor { } } + pub(super) fn take_pending_sampled_token(&mut self) -> Option { + match self { + Self::Qwen(executor) => executor.take_pending_sampled_token(), + Self::DeepSeek(_) | Self::Glm(_) => None, + } + } + + pub(super) fn finalize_pending_sampled_token(&mut self) -> Result<(), String> { + match self { + Self::Qwen(executor) => executor.finalize_pending_sampled_token(), + Self::DeepSeek(_) | Self::Glm(_) => Ok(()), + } + } + pub(super) fn eval(&mut self, token: i32) -> Result<(), String> { match self { Self::DeepSeek(executor) => executor.eval(token), diff --git a/src/engine/metal/gpu.rs b/src/engine/metal/gpu.rs index 634c50a..79acd1d 100644 --- a/src/engine/metal/gpu.rs +++ b/src/engine/metal/gpu.rs @@ -8,6 +8,7 @@ pub(super) struct GpuTensor { #[derive(Clone, Copy)] #[repr(C)] pub(super) struct QwenWeightView { + pub(super) tensor: *const GpuTensor, pub(super) map: *const c_void, pub(super) size: u64, pub(super) offset: u64, @@ -90,6 +91,7 @@ pub(super) struct StreamExpertCacheStats { unsafe extern "C" { pub(super) fn ds4_gpu_init() -> i32; pub(super) fn ds4_gpu_cleanup(); + pub(super) fn ds4_gpu_metal4_tensor_api_enabled() -> i32; pub(super) fn ds4_gpu_set_model_map_range( model_map: *const c_void, model_size: u64, @@ -103,6 +105,7 @@ unsafe extern "C" { a: *const GpuTensor, b: *const GpuTensor, c: *const GpuTensor, + d: *const GpuTensor, weights: *const QwenWeightView, weight_count: u32, args: *const QwenKernelArgs, @@ -132,6 +135,7 @@ unsafe extern "C" { ) -> i32; pub(super) fn ds4_gpu_set_quality(quality: bool); pub(super) fn ds4_gpu_set_glm_model(enabled: bool); + pub(super) fn ds4_gpu_set_qwen_model(enabled: bool); pub(super) fn ds4_gpu_set_ssd_streaming(enabled: bool); pub(super) fn ds4_gpu_set_model_fd(fd: i32) -> i32; pub(super) fn ds4_gpu_set_streaming_expert_cache_budget(experts: u32); @@ -199,6 +203,7 @@ unsafe extern "C" { pub(super) fn ds4_gpu_parallel_ffn_finish() -> i32; pub(super) fn ds4_gpu_parallel_ffn_abort(); pub(super) fn ds4_gpu_tensor_alloc(bytes: u64) -> *mut GpuTensor; + pub(super) fn ds4_gpu_tensor_alloc_untracked(bytes: u64) -> *mut GpuTensor; pub(super) fn ds4_gpu_tensor_view( base: *const GpuTensor, offset: u64, @@ -1847,6 +1852,7 @@ impl Context { check(unsafe { ds4_gpu_init() }, "Metal initialization")?; unsafe { ds4_gpu_set_glm_model(model.shape.family == ModelFamily::Glm); + ds4_gpu_set_qwen_model(false); ds4_gpu_set_ssd_streaming(ssd_streaming); // Decode enables this only around DS4's eligible resident pre-M5 // MXFP4 token path; all other work starts from the portable path. @@ -1944,6 +1950,7 @@ impl Context { check(unsafe { ds4_gpu_init() }, "Metal initialization")?; unsafe { ds4_gpu_set_glm_model(false); + ds4_gpu_set_qwen_model(true); ds4_gpu_set_ssd_streaming(false); ds4_gpu_set_decode_pipeline_fast_lookup(0); } @@ -2071,6 +2078,12 @@ impl Buffer { .ok_or_else(|| format!("Metal could not allocate {bytes} bytes")) } + pub(super) fn untracked(bytes: u64) -> Result { + NonNull::new(unsafe { ds4_gpu_tensor_alloc_untracked(bytes) }) + .map(Self) + .ok_or_else(|| format!("Metal could not allocate {bytes} untracked bytes")) + } + pub(super) fn view(&self, offset: u64, bytes: u64) -> Result { // SAFETY: `ds4_gpu_tensor_view` in native/metal/ds4_metal.m bounds-checks // the view, ARC-retains base_obj.buffer, and marks the view non-owning, so diff --git a/src/engine/metal/qwen.rs b/src/engine/metal/qwen.rs index 0ee1570..c713aae 100644 --- a/src/engine/metal/qwen.rs +++ b/src/engine/metal/qwen.rs @@ -1,12 +1,15 @@ use super::checkpoint::{read_buffer, read_u32, write_buffer, write_u32}; use super::*; use crate::engine::qwen::{QwenModel, QwenTensor}; +use rustc_hash::FxHashMap as HashMap; use std::ffi::CStr; const HIDDEN: u32 = 2_560; const HC: u32 = 4; const HC_WIDTH: u32 = HIDDEN * HC; const HC_RANK: u32 = 320; +const HYPER_FOLDED_ROWS: u32 = HC_RANK + HC; +const HYPER_QUANT_GROUP: u32 = 64; const LAYERS: usize = 48; const GDN_HEADS_K: u32 = 16; const GDN_HEADS_V: u32 = 48; @@ -14,6 +17,7 @@ const HEAD_DIM: u32 = 128; const GDN_QKV: u32 = (GDN_HEADS_K * 2 + GDN_HEADS_V) * HEAD_DIM; const GDN_VALUE: u32 = GDN_HEADS_V * HEAD_DIM; const GDN_CONTROLS: u32 = GDN_VALUE + GDN_HEADS_V * 2; +const GDN_FUSED_INPUT: u32 = GDN_QKV + GDN_CONTROLS; const ATTN_HEADS: u32 = 24; const ATTN_KV_HEADS: u32 = 2; const ATTN_DIM: u32 = 256; @@ -23,14 +27,18 @@ const QSA_HEADS: u32 = 4; const QSA_KV_HEADS: u32 = 1; const QSA_DIM: u32 = 128; const QSA_WIDTH: u32 = (QSA_HEADS + QSA_KV_HEADS) * QSA_DIM; +const QSA_FUSED_QKV: u32 = ATTN_WIDTH * 2 + ATTN_KV_WIDTH * 2; +const QSA_FUSED_INPUT: u32 = QSA_FUSED_QKV + QSA_WIDTH; const QSA_RATIO: u32 = 4; const QSA_TOP_K: u32 = 512; const MTP_TOKEN_RESERVE: u32 = 3; const MTP_CAPTURE_ROWS: usize = MTP_TOKEN_RESERVE as usize + 1; +const VERIFY_ROWS: u64 = MTP_CAPTURE_ROWS as u64; const EXPERTS: u32 = 512; const EXPERTS_USED: usize = 10; const EXPERT_WIDTH: u32 = 640; const VOCAB: u32 = 248_320; +const SAMPLING_TOP_K: u32 = 32; const DENSE_BUDGET: u32 = 2_048; const PLE_HEADS: usize = 16; const PLE_HEAD_DIM: u32 = 160; @@ -39,8 +47,13 @@ const PLE_CONV_STATE: u32 = 9; const EOS_TOKEN: i32 = 248_044; const PLE_ROW_BYTES: usize = 100; const CHECKPOINT_MAGIC: &[u8; 8] = b"DS4QWN01"; -const CHECKPOINT_VERSION: u32 = 4; +const CHECKPOINT_VERSION: u32 = 5; const CHECKPOINT_CHUNK: usize = 8 * 1024 * 1024; +const PREFILL_CHUNK: usize = 2_048; + +type TokenDistribution = Vec<(usize, f64)>; +type TokenDistributions = Vec; +type DraftStep = (Vec, Option); #[derive(Clone, Copy)] struct Weight<'a> { @@ -69,16 +82,27 @@ enum LayerState { } struct Scratch { + rows: u32, + token_ids: Buffer, hidden: Buffer, hc: Buffer, + batch_final_hc: Buffer, hc_norm: Buffer, hc_mix: Buffer, rank: Buffer, block: Buffer, injection: Buffer, + hyper_rms: Buffer, qkv: Buffer, controls: Buffer, - gdn_raw: Buffer, + gdn_fused: Buffer, + gdn_fused_qkv: Buffer, + gdn_fused_controls: Buffer, + gdn_z: Buffer, + gdn_b: Buffer, + gdn_a: Buffer, + gdn_q: Buffer, + gdn_k: Buffer, gdn_out: Buffer, router: Buffer, route_ids: Buffer, @@ -90,6 +114,11 @@ struct Scratch { moe: Buffer, shared: Buffer, q_packed: Buffer, + qsa_fused: Buffer, + qsa_fused_q: Buffer, + qsa_fused_k: Buffer, + qsa_fused_v: Buffer, + qsa_fused_indexer: Buffer, q: Buffer, q_gate: Buffer, k: Buffer, @@ -100,6 +129,11 @@ struct Scratch { qsa_q: Buffer, qsa_scores: Buffer, qsa_selected: Buffer, + causal_mask: Buffer, + dense_query: Buffer, + dense_keys: Buffer, + dense_values: Buffer, + attention_scores: Buffer, ple_packed: Buffer, ple_scales: Buffer, ple_biases: Buffer, @@ -110,53 +144,99 @@ struct Scratch { ple_norm: Buffer, ple_output: Buffer, logits: Buffer, + sampling_selected: Buffer, + sampling_probabilities: Buffer, } impl Scratch { - fn new(context: u32) -> Result { + fn new(context: u32, capacity: u32) -> Result { let qsa_blocks = qsa_block_capacity(context)?; + let capacity = capacity.max(VERIFY_ROWS as u32); + let row_count = u64::from(capacity); + let rows = |width: u32| u64::from(width) * row_count; + let gdn_fused = Buffer::floats(rows(GDN_FUSED_INPUT))?; + let gdn_fused_qkv = gdn_fused.view(0, u64::from(GDN_QKV) * 4)?; + let gdn_fused_controls = + gdn_fused.view(u64::from(GDN_QKV) * 4, u64::from(GDN_CONTROLS) * 4)?; + let qsa_fused = Buffer::floats(QSA_FUSED_INPUT.into())?; + let qsa_fused_q = qsa_fused.view(0, u64::from(ATTN_WIDTH * 2) * 4)?; + let qsa_fused_k = + qsa_fused.view(u64::from(ATTN_WIDTH * 2) * 4, u64::from(ATTN_KV_WIDTH) * 4)?; + let qsa_fused_v = qsa_fused.view( + u64::from(ATTN_WIDTH * 2 + ATTN_KV_WIDTH) * 4, + u64::from(ATTN_KV_WIDTH) * 4, + )?; + let qsa_fused_indexer = qsa_fused.view( + u64::from(ATTN_WIDTH * 2 + ATTN_KV_WIDTH * 2) * 4, + u64::from(QSA_WIDTH) * 4, + )?; Ok(Self { - hidden: Buffer::floats(HIDDEN.into())?, - hc: Buffer::floats(HC_WIDTH.into())?, - hc_norm: Buffer::floats(HC_WIDTH.into())?, - hc_mix: Buffer::floats(HC_WIDTH.into())?, - rank: Buffer::floats(HC_RANK.into())?, - block: Buffer::floats(HIDDEN.into())?, - injection: Buffer::floats(HC.into())?, - qkv: Buffer::floats(GDN_QKV.into())?, - controls: Buffer::floats(GDN_CONTROLS.into())?, - gdn_raw: Buffer::floats(GDN_VALUE.into())?, - gdn_out: Buffer::floats(GDN_VALUE.into())?, - router: Buffer::floats(EXPERTS.into())?, - route_ids: Buffer::bytes((EXPERTS_USED * 4) as u64)?, - route_weights: Buffer::floats(EXPERTS_USED as u64)?, - gate: Buffer::floats(EXPERT_WIDTH.into())?, - up: Buffer::floats(EXPERT_WIDTH.into())?, - mid: Buffer::floats(EXPERT_WIDTH.into())?, - expert: Buffer::floats(HIDDEN.into())?, - moe: Buffer::floats(HIDDEN.into())?, - shared: Buffer::floats(HIDDEN.into())?, - q_packed: Buffer::floats((ATTN_WIDTH * 2).into())?, - q: Buffer::floats(ATTN_WIDTH.into())?, - q_gate: Buffer::floats(ATTN_WIDTH.into())?, - k: Buffer::floats(ATTN_KV_WIDTH.into())?, - k_rope: Buffer::floats(ATTN_KV_WIDTH.into())?, - v: Buffer::floats(ATTN_KV_WIDTH.into())?, - attention: Buffer::floats(ATTN_WIDTH.into())?, - qsa_qk: Buffer::floats(QSA_WIDTH.into())?, - qsa_q: Buffer::floats((QSA_HEADS * QSA_DIM).into())?, - qsa_scores: Buffer::floats(qsa_blocks.into())?, - qsa_selected: Buffer::bytes(u64::from(QSA_TOP_K) * 4)?, - ple_packed: Buffer::bytes((PLE_HEADS as u64) * 80)?, - ple_scales: Buffer::bytes((PLE_HEADS as u64) * 10)?, - ple_biases: Buffer::bytes((PLE_HEADS as u64) * 10)?, - ple_embedding: Buffer::floats(HIDDEN.into())?, - ple_key: Buffer::floats(HC_WIDTH.into())?, - ple_value: Buffer::floats(HIDDEN.into())?, - ple_gated: Buffer::floats(HC_WIDTH.into())?, - ple_norm: Buffer::floats(HC_WIDTH.into())?, - ple_output: Buffer::floats(HC_WIDTH.into())?, - logits: Buffer::floats(VOCAB.into())?, + rows: capacity, + token_ids: Buffer::bytes(row_count * 4)?, + hidden: Buffer::floats(rows(HIDDEN))?, + hc: Buffer::floats(rows(HC_WIDTH))?, + batch_final_hc: Buffer::floats(HC_WIDTH.into())?, + hc_norm: Buffer::floats(rows(HC_WIDTH))?, + hc_mix: Buffer::floats(rows(HC_WIDTH))?, + rank: Buffer::floats(rows(HC_RANK))?, + block: Buffer::floats(rows(HIDDEN))?, + injection: Buffer::floats(rows(HC))?, + hyper_rms: Buffer::floats(rows(HC))?, + qkv: Buffer::floats(rows(GDN_QKV))?, + controls: Buffer::floats(rows(GDN_CONTROLS))?, + gdn_fused, + gdn_fused_qkv, + gdn_fused_controls, + gdn_z: Buffer::floats(rows(GDN_VALUE))?, + gdn_b: Buffer::floats(rows(GDN_HEADS_V))?, + gdn_a: Buffer::floats(rows(GDN_HEADS_V))?, + gdn_q: Buffer::floats(rows(GDN_HEADS_K * HEAD_DIM))?, + gdn_k: Buffer::floats(rows(GDN_HEADS_K * HEAD_DIM))?, + gdn_out: Buffer::bytes(rows(GDN_VALUE) * 4 + u64::from(GDN_QKV) * 3 * 2)?, + router: Buffer::floats(rows(EXPERTS))?, + route_ids: Buffer::bytes(row_count * EXPERTS_USED as u64 * 4)?, + route_weights: Buffer::floats(row_count * EXPERTS_USED as u64)?, + gate: Buffer::floats(row_count * EXPERTS_USED as u64 * u64::from(EXPERT_WIDTH))?, + up: Buffer::floats(row_count * EXPERTS_USED as u64 * u64::from(EXPERT_WIDTH))?, + mid: Buffer::floats(row_count * EXPERTS_USED as u64 * u64::from(EXPERT_WIDTH))?, + expert: Buffer::floats(row_count * EXPERTS_USED as u64 * u64::from(HIDDEN))?, + moe: Buffer::floats(rows(HIDDEN))?, + shared: Buffer::floats(rows(HIDDEN))?, + q_packed: Buffer::floats(rows(ATTN_WIDTH * 2))?, + qsa_fused, + qsa_fused_q, + qsa_fused_k, + qsa_fused_v, + qsa_fused_indexer, + q: Buffer::floats(rows(ATTN_WIDTH))?, + q_gate: Buffer::floats(rows(ATTN_WIDTH))?, + k: Buffer::floats(rows(ATTN_KV_WIDTH))?, + k_rope: Buffer::floats(rows(ATTN_KV_WIDTH))?, + v: Buffer::floats(rows(ATTN_KV_WIDTH))?, + attention: Buffer::floats(rows(ATTN_WIDTH))?, + qsa_qk: Buffer::floats(rows(QSA_WIDTH))?, + qsa_q: Buffer::floats(rows(QSA_HEADS * QSA_DIM))?, + qsa_scores: Buffer::floats(u64::from(qsa_blocks) * row_count)?, + qsa_selected: Buffer::bytes(u64::from(QSA_TOP_K) * row_count * 4)?, + causal_mask: Buffer::bytes(u64::from(DENSE_BUDGET) * row_count)?, + dense_query: Buffer::bytes(rows(ATTN_WIDTH) * 2)?, + dense_keys: Buffer::bytes(u64::from(DENSE_BUDGET) * u64::from(ATTN_KV_WIDTH) * 2)?, + dense_values: Buffer::bytes(u64::from(DENSE_BUDGET) * u64::from(ATTN_KV_WIDTH) * 2)?, + attention_scores: Buffer::bytes( + u64::from(DENSE_BUDGET) * u64::from(ATTN_HEADS) * row_count * 2, + )?, + ple_packed: Buffer::bytes(row_count * PLE_HEADS as u64 * 80)?, + ple_scales: Buffer::bytes(row_count * PLE_HEADS as u64 * 10)?, + ple_biases: Buffer::bytes(row_count * PLE_HEADS as u64 * 10)?, + ple_embedding: Buffer::floats(rows(HIDDEN))?, + ple_key: Buffer::floats(rows(HC_WIDTH))?, + ple_value: Buffer::floats(rows(HIDDEN))?, + ple_gated: Buffer::floats(rows(HC_WIDTH))?, + ple_norm: Buffer::floats(rows(HC_WIDTH))?, + ple_output: Buffer::floats(rows(HC_WIDTH))?, + logits: Buffer::floats(u64::from(VOCAB) * row_count.min(VERIFY_ROWS))?, + sampling_selected: Buffer::bytes(row_count * u64::from(SAMPLING_TOP_K) * 4)?, + sampling_probabilities: Buffer::floats(row_count * u64::from(SAMPLING_TOP_K))?, }) } } @@ -172,6 +252,7 @@ struct PleState { conv: Buffer, } +#[cfg(test)] struct RecurrentSnapshot { gdn: Vec>, ple_conv: Buffer, @@ -182,9 +263,28 @@ struct RecurrentSnapshot { tokens: usize, } +struct VerifyCapture { + gdn: Vec>, + ple_conv: Buffer, + ple_input: Buffer, + hidden: Buffer, +} + +struct GdnVerifyCapture { + conv: Buffer, + recurrent: Buffer, + qkv: Buffer, + q: Buffer, + k: Buffer, + v: Buffer, + controls: Buffer, +} + struct QwenMtp { attention: LayerState, - captures: Vec, + #[cfg(test)] + snapshot: RecurrentSnapshot, + verify: VerifyCapture, timing: bool, cycles: u64, drafted: u64, @@ -198,11 +298,20 @@ struct QwenMtp { struct QwenDraft { token: i32, + #[cfg(test)] logits: Vec, + probabilities: Vec<(usize, f64)>, +} + +struct HyperPack { + r1: Buffer, + r2: Buffer, } pub(in crate::engine) struct QwenExecutor { model: QwenModel, + weights: HashMap, + weight_views: HashMap, states: Vec, ple_contract: PleContract, ple_state: PleState, @@ -213,56 +322,321 @@ pub(in crate::engine) struct QwenExecutor { context: u32, checkpoint_tag: [u8; 32], mtp: Option, + pending_sampled: Option, + folding_sampled: bool, + hyper_packs: HashMap, + packed_moe: bool, + fused_gdn: bool, + fused_gdn_input: bool, + fused_qsa_input: bool, _context: Context, } pub(in crate::engine) struct QwenResidentState { states: Vec, ple_state: PleState, + hidden: Buffer, logits: Vec, tokens: Vec, position: u32, checkpoint_tag: [u8; 32], mtp: Option, + pending_sampled: Option, + folding_sampled: bool, } impl QwenExecutor { #[cfg(test)] pub(super) fn open(model: QwenModel, context: u32) -> Result { - Self::open_configured(model, context, EngineSpeculativeSettings::default()) + Self::open_configured_with_hyper( + model, + context, + EngineSpeculativeSettings::default(), + true, + true, + ) } pub(super) fn open_configured( model: QwenModel, context: u32, speculative: EngineSpeculativeSettings, + ) -> Result { + Self::open_configured_with_hyper( + model, + context, + speculative, + qwen_feature_enabled("DS4SERVER_QWEN38_FUSED_HYPER"), + true, + ) + } + + fn open_configured_with_hyper( + mut model: QwenModel, + context: u32, + speculative: EngineSpeculativeSettings, + fused_hyper: bool, + materialize: bool, ) -> Result { let ple_contract = ple_contract(&model)?; let native = Context::open_qwen(model.memory().admission)?; + let weights = if materialize { + let weights = Self::materialize_weights(&model)?; + model.release_materialized_mappings(); + weights + } else { + HashMap::default() + }; + let weight_views = model + .tensors() + .filter_map(|tensor| { + weights.get(&tensor.name).map(|buffer| { + let bytes = tensor.range.end - tensor.range.start; + ( + tensor as *const QwenTensor as usize, + QwenWeightView { + tensor: buffer.raw().cast_const(), + map: std::ptr::null(), + size: bytes, + offset: 0, + bytes, + }, + ) + }) + }) + .collect(); let states = allocate_states(context)?; let mtp = speculative .glm_mtp .then(|| allocate_mtp(context, speculative.glm_mtp_timing)) .transpose()?; - Ok(Self { + let mut executor = Self { model, + weights, + weight_views, states, ple_contract, ple_state: allocate_ple_state()?, - scratch: Scratch::new(context)?, + scratch: Scratch::new(context, VERIFY_ROWS as u32)?, logits: vec![0.0; VOCAB as usize], tokens: Vec::new(), position: 0, context, checkpoint_tag: [0; 32], mtp, + pending_sampled: None, + folding_sampled: false, + hyper_packs: HashMap::default(), + packed_moe: qwen_feature_enabled("DS4SERVER_QWEN38_FUSED_GATE_UP"), + fused_gdn: qwen_feature_enabled("DS4SERVER_QWEN38_FUSED_GDN"), + fused_gdn_input: qwen_feature_enabled("DS4SERVER_QWEN38_FUSED_GDN_INPROJ"), + fused_qsa_input: qwen_feature_enabled("DS4SERVER_QWEN38_FUSED_QSA_QKV"), _context: native, - }) + }; + if fused_hyper { + executor.prepare_hyper_packs()?; + } + Ok(executor) + } + + fn materialize_weights(model: &QwenModel) -> Result, String> { + const COPY_CHUNK: usize = 8 * 1024 * 1024; + let mut tensors = model + .tensors() + .filter(|tensor| !tensor.name.starts_with("ngram.")) + .collect::>(); + tensors.sort_by(|left, right| left.name.cmp(&right.name)); + let mut weights = HashMap::with_capacity_and_hasher(tensors.len(), Default::default()); + let mut staging = vec![0; COPY_CHUNK]; + for tensor in tensors { + let bytes = tensor.range.end - tensor.range.start; + let target = Buffer::untracked(bytes)?; + for offset in (0..bytes).step_by(COPY_CHUNK) { + let count = usize::try_from((bytes - offset).min(COPY_CHUNK as u64)) + .map_err(|_| format!("{} chunk size exceeds this platform", tensor.name))?; + model.read_tensor_at(tensor, offset, &mut staging[..count])?; + target.write(offset, &staging[..count])?; + std::thread::yield_now(); + } + weights.insert(tensor.name.clone(), target); + } + Self::pack_gdn_inputs(model, &mut weights)?; + Self::pack_qsa_inputs(model, &mut weights)?; + Ok(weights) + } + + fn pack_gdn_inputs( + model: &QwenModel, + weights: &mut HashMap, + ) -> Result<(), String> { + for layer in (0..LAYERS).filter(|layer| layer % 4 != 3) { + let prefix = format!("language_model.model.layers.{layer}.linear_attn"); + for suffix in ["weight", "scales", "biases"] { + let names = ["in_proj_qkv", "in_proj_z", "in_proj_b", "in_proj_a"] + .map(|projection| format!("{prefix}.{projection}.{suffix}")); + let sizes = names + .each_ref() + .map(|name| { + model + .tensor(name) + .map(|tensor| tensor.range.end - tensor.range.start) + }) + .into_iter() + .collect::, _>>()?; + let total = sizes.iter().sum(); + let fused = Buffer::untracked(total)?; + let commands = Commands::begin()?; + let mut offset = 0; + for (name, bytes) in names.iter().zip(&sizes) { + fused.copy_from( + offset, + weights + .get(name) + .ok_or_else(|| format!("Qwen materialization missed {name}"))?, + 0, + *bytes, + "packing Qwen GDN input projections", + )?; + offset += bytes; + } + commands.finish()?; + offset = 0; + for (name, bytes) in names.into_iter().zip(sizes) { + weights.insert(name, fused.view(offset, bytes)?); + offset += bytes; + } + weights.insert(format!("{prefix}.in_proj_fused.{suffix}"), fused); + } + } + Ok(()) + } + + fn pack_qsa_inputs( + model: &QwenModel, + weights: &mut HashMap, + ) -> Result<(), String> { + for layer in (0..LAYERS).filter(|layer| layer % 4 == 3) { + let prefix = format!("language_model.model.layers.{layer}.self_attn"); + let q = model.tensor(&format!("{prefix}.q_proj.weight"))?; + let indexer = model.tensor(&format!("{prefix}.indexer.index_qk_proj.weight"))?; + let include_indexer = q.quant_bits == indexer.quant_bits + && q.group_size == indexer.group_size + && q.shape.last() == indexer.shape.last(); + for suffix in ["weight", "scales", "biases"] { + let mut names = ["q_proj", "k_proj", "v_proj"] + .map(|projection| format!("{prefix}.{projection}.{suffix}")) + .to_vec(); + if include_indexer { + names.push(format!("{prefix}.indexer.index_qk_proj.{suffix}")); + } + let sizes = names + .iter() + .map(|name| { + model + .tensor(name) + .map(|tensor| tensor.range.end - tensor.range.start) + }) + .collect::, _>>()?; + let total = sizes.iter().sum(); + let fused = Buffer::untracked(total)?; + let commands = Commands::begin()?; + let mut offset = 0; + for (name, bytes) in names.iter().zip(&sizes) { + fused.copy_from( + offset, + weights + .get(name) + .ok_or_else(|| format!("Qwen materialization missed {name}"))?, + 0, + *bytes, + "packing Qwen QSA input projections", + )?; + offset += bytes; + } + commands.finish()?; + offset = 0; + for (name, bytes) in names.into_iter().zip(sizes) { + weights.insert(name, fused.view(offset, bytes)?); + offset += bytes; + } + weights.insert(format!("{prefix}.qkv_fused.{suffix}"), fused); + } + } + Ok(()) + } + + fn prepare_hyper_packs(&mut self) -> Result<(), String> { + let mut prefixes = Vec::with_capacity(LAYERS * 2 + usize::from(self.mtp.is_some()) * 2); + for layer in 0..LAYERS { + for block in ["attn", "mlp"] { + prefixes.push(format!( + "language_model.model.layers.{layer}.{block}_hyper_connection" + )); + } + } + if self.mtp.is_some() { + prefixes.push("mtp.layers.0.attn_hyper_connection".to_owned()); + prefixes.push("mtp.layers.0.mlp_hyper_connection".to_owned()); + } + + let commands = Commands::begin()?; + for prefix in prefixes { + let pack = self.prepare_hyper_pack(&prefix)?; + self.hyper_packs.insert(prefix, pack); + } + commands.finish() + } + + fn prepare_hyper_pack(&self, prefix: &str) -> Result { + let down = self.weight(&format!("{prefix}.input_mix_weight_down.weight"))?; + let inject = self.weight(&format!("{prefix}.block_inject_weight.weight"))?; + let up = self.weight(&format!("{prefix}.input_mix_weight_up.weight"))?; + let r1 = Buffer::bytes(hyper_pack_bytes(HYPER_FOLDED_ROWS, HC_WIDTH))?; + let r2 = Buffer::bytes(hyper_pack_bytes(HC_WIDTH, HC_RANK))?; + self.quantize_hyper_rows(down, &r1, HC_WIDTH, HC_RANK, 0, HYPER_FOLDED_ROWS)?; + self.quantize_hyper_rows(inject, &r1, HC_WIDTH, HC, HC_RANK, HYPER_FOLDED_ROWS)?; + self.quantize_hyper_rows(up, &r2, HC_RANK, HC_WIDTH, 0, HC_WIDTH)?; + Ok(HyperPack { r1, r2 }) + } + + #[allow(clippy::too_many_arguments)] + fn quantize_hyper_rows( + &self, + weight: Weight<'_>, + pack: &Buffer, + columns: u32, + rows: u32, + destination_row: u32, + destination_rows: u32, + ) -> Result<(), String> { + if weight.tensor.dtype != "BF16" || weight.tensor.shape != [rows as u64, columns as u64] { + return Err(format!( + "{} has an incompatible BF16 matrix shape", + weight.tensor.name + )); + } + let mut values = args(); + values.u[0] = columns; + values.u[1] = rows; + values.u[2] = destination_row; + values.u[3] = destination_rows; + values.u[12] = 32; + self.dispatch( + c"kernel_qwen_quantize_hyper_q8", + pack, + None, + None, + None, + &[self.view(weight)], + &values, + columns / HYPER_QUANT_GROUP, + rows, + ) } pub(super) fn eval(&mut self, token: i32) -> Result<(), String> { if self.position > 0 && self.mtp.is_some() { - let _ = self.mtp_step(token, self.position - 1, false)?; + let _ = self.mtp_step(token, self.position - 1, false, None)?; } self.eval_target(token) } @@ -277,8 +651,8 @@ impl QwenExecutor { self.context )); } + let commands = Commands::begin()?; self.begin_token(token)?; - for layer in 0..LAYERS { if layer == 1 { self.ple(token)?; @@ -286,11 +660,194 @@ impl QwenExecutor { self.encode_layer(layer)?; } self.final_output()?; + commands.finish()?; + self.scratch.logits.read_f32(&mut self.logits)?; self.tokens.push(token); self.position += 1; Ok(()) } + fn eval_target_batch( + &mut self, + tokens: &[i32], + sampling: Option<(f32, f32, f32, i32)>, + ) -> Result<(Vec>, Option), String> { + if tokens.is_empty() || tokens.len() > MTP_CAPTURE_ROWS { + return Err("Qwen verifier width is outside the MTPLX S=2..4 contract".into()); + } + if self.position + tokens.len() as u32 > self.context { + return Err(format!( + "the Qwen executor supports {} tokens per session", + self.context + )); + } + let rows = tokens.len() as u32; + let profiled = std::env::var_os("DS4SERVER_QWEN38_MTP_PROFILE").is_some(); + let total_started = Instant::now(); + let commands = Commands::begin()?; + self.begin_tokens(tokens)?; + for layer in 0..LAYERS { + let prefix = format!("language_model.model.layers.{layer}"); + if layer == 1 { + if let Some(mtp) = &self.mtp { + mtp.verify.ple_input.copy_from( + 0, + &self.scratch.hc, + 0, + u64::from(rows * HC_WIDTH) * 4, + "capturing Qwen verifier PLE input rows", + )?; + } + let state_out = self + .mtp + .as_ref() + .ok_or("Qwen MTP is not configured")? + .verify + .ple_conv + .view(0, u64::from(HC_WIDTH) * PLE_CONV_STATE as u64 * 2)?; + self.ple_batch_with_state(tokens, Some(&state_out))?; + } + self.hyper_read_rows(&format!("{prefix}.attn_hyper_connection"), rows)?; + match &self.states[layer] { + LayerState::Gdn { .. } => self.gdn_rows(&prefix, layer, rows, true)?, + state @ LayerState::Attention { .. } => { + self.attention_rows(&prefix, state, self.position, rows)? + } + } + self.hyper_write_rows(rows)?; + self.encode_moe_rows(&prefix, rows)?; + } + self.final_output_rows(rows)?; + if let Some(mtp) = &self.mtp { + mtp.verify.hidden.copy_from( + 0, + &self.scratch.hc, + 0, + u64::from(rows) * u64::from(HC_WIDTH) * 4, + "capturing Qwen verifier hidden rows", + )?; + } + let sparse = if let Some((temperature, _, _, top_k)) = sampling { + self.encode_sparse_probabilities(&self.scratch.logits, rows, temperature, top_k)? + } else { + false + }; + let encoded = total_started.elapsed(); + commands.finish()?; + let executed = total_started.elapsed(); + let probabilities = if sparse { + let (_, top_p, min_p, top_k) = sampling.expect("sampling configuration disappeared"); + self.read_sparse_probabilities(rows, top_p, min_p, top_k)? + } else { + None + }; + let logits = if probabilities.is_some() { + Vec::new() + } else { + let mut flattened = vec![0.0; tokens.len() * VOCAB as usize]; + self.scratch.logits.read_f32(&mut flattened)?; + flattened + .chunks_exact(VOCAB as usize) + .map(<[f32]>::to_vec) + .collect::>() + }; + if profiled { + eprintln!( + "ds4: Qwen verifier rows={rows} encode={:.3} ms gpu={:.3} ms read={:.3} ms", + encoded.as_secs_f64() * 1_000.0, + (executed - encoded).as_secs_f64() * 1_000.0, + (total_started.elapsed() - executed).as_secs_f64() * 1_000.0, + ); + } + if let Some(last) = logits.last() { + self.logits.copy_from_slice(last); + } + self.tokens.extend_from_slice(tokens); + self.position += rows; + Ok((logits, probabilities)) + } + + fn read_logits_row(&mut self, row: usize) -> Result<(), String> { + self.logits = self.logits_row(row)?; + Ok(()) + } + + fn logits_row(&self, row: usize) -> Result, String> { + let mut logits = vec![0.0; VOCAB as usize]; + self.scratch + .logits + .view(row as u64 * u64::from(VOCAB) * 4, u64::from(VOCAB) * 4)? + .read_f32(&mut logits)?; + Ok(logits) + } + + fn eval_prefill_batch( + &mut self, + tokens: &[i32], + mtp_tokens: Option<&[i32]>, + ) -> Result<(), String> { + if tokens.is_empty() || tokens.len() > self.scratch.rows as usize { + return Err(format!( + "Qwen prefill batch width must be in 1..={}", + self.scratch.rows, + )); + } + if mtp_tokens.is_some_and(|next| next.len() != tokens.len()) { + return Err("Qwen MTP prompt-history rows do not match the target rows".into()); + } + if self.position + tokens.len() as u32 > self.context { + return Err(format!( + "the Qwen executor supports {} tokens per session", + self.context + )); + } + let rows = tokens.len() as u32; + let commands = Commands::begin()?; + self.begin_tokens(tokens)?; + for layer in 0..LAYERS { + let prefix = format!("language_model.model.layers.{layer}"); + if layer == 1 { + self.ple_batch(tokens)?; + } + self.hyper_read_rows(&format!("{prefix}.attn_hyper_connection"), rows)?; + match &self.states[layer] { + LayerState::Gdn { .. } => self.gdn_rows(&prefix, layer, rows, false)?, + state @ LayerState::Attention { .. } => { + self.attention_rows(&prefix, state, self.position, rows)? + } + } + self.hyper_write_rows(rows)?; + self.encode_moe_rows(&prefix, rows)?; + } + self.mix_output_rows("language_model.model.hyper_connection_mixer", rows)?; + let final_block = self.scratch.block.view( + u64::from(rows - 1) * u64::from(HIDDEN) * 4, + u64::from(HIDDEN) * 4, + )?; + self.affine_mv_into( + &self.affine("language_model.lm_head", HIDDEN, VOCAB, None)?, + &final_block, + &self.scratch.logits, + HIDDEN, + VOCAB, + )?; + self.scratch.batch_final_hc.copy_from( + 0, + &self.scratch.hc, + u64::from(rows - 1) * u64::from(HC_WIDTH) * 4, + u64::from(HC_WIDTH) * 4, + "preserving the final Qwen prefill hidden row", + )?; + commands.finish()?; + self.scratch.logits.read_f32(&mut self.logits)?; + if let Some(next) = mtp_tokens { + self.mtp_history_batch(next, self.position)?; + } + self.tokens.extend_from_slice(tokens); + self.position += rows; + Ok(()) + } + fn begin_token(&self, token: i32) -> Result<(), String> { let embedding = self.affine("language_model.model.embed_tokens", HIDDEN, VOCAB, None)?; let mut args = args(); @@ -298,7 +855,6 @@ impl QwenExecutor { args.u[2] = embedding.bits; args.u[3] = embedding.group; args.u[4] = token as u32; - let commands = Commands::begin()?; self.dispatch( c"kernel_qwen_affine_embedding", &self.scratch.hidden, @@ -326,34 +882,100 @@ impl QwenExecutor { HC_WIDTH, 1, )?; - commands.finish() + Ok(()) + } + + fn begin_tokens(&self, tokens: &[i32]) -> Result<(), String> { + if tokens.is_empty() || tokens.len() > self.scratch.rows as usize { + return Err(format!( + "Qwen batch width is outside the allocated {}-row capacity", + self.scratch.rows + )); + } + if let Some(&token) = tokens + .iter() + .find(|&&token| token < 0 || token as u32 >= VOCAB) + { + return Err(format!("token {token} is outside the Qwen vocabulary")); + } + self.scratch.token_ids.write_i32(tokens)?; + let embedding = self.affine("language_model.model.embed_tokens", HIDDEN, VOCAB, None)?; + let mut args = args(); + args.u[0] = HIDDEN; + args.u[2] = embedding.bits; + args.u[3] = embedding.group; + args.u[4] = tokens.len() as u32; + self.dispatch( + c"kernel_qwen_affine_embedding_batch", + &self.scratch.hidden, + Some(&self.scratch.token_ids), + None, + None, + &[ + self.view(embedding.packed), + self.view(embedding.scales), + self.view(embedding.biases), + ], + &args, + HIDDEN, + tokens.len() as u32, + )?; + args.u[0] = HIDDEN; + self.dispatch( + c"kernel_qwen_repeat4", + &self.scratch.hc, + Some(&self.scratch.hidden), + None, + None, + &[], + &args, + HIDDEN * HC * tokens.len() as u32, + 1, + ) } fn ple(&mut self, token: i32) -> Result<(), String> { - let rows = ple_rows(&self.ple_contract, self.ple_state.history, token)?; - let packed = self.model.tensor("ngram.weight")?; - let scales = self.model.tensor("ngram.scales")?; - let biases = self.model.tensor("ngram.biases")?; - let mut packed_stage = [0_u8; PLE_HEADS * 80]; - let mut scales_stage = [0_u8; PLE_HEADS * 10]; - let mut biases_stage = [0_u8; PLE_HEADS * 10]; - let packed_bytes = self.model.tensor_bytes(packed)?; - let scales_bytes = self.model.tensor_bytes(scales)?; - let biases_bytes = self.model.tensor_bytes(biases)?; - for (head, &row) in rows.iter().enumerate() { - let bytes = gather_ple_row(row, packed_bytes, scales_bytes, biases_bytes)?; - packed_stage[head * 80..][..80].copy_from_slice(&bytes[..80]); - scales_stage[head * 10..][..10].copy_from_slice(&bytes[80..90]); - biases_stage[head * 10..][..10].copy_from_slice(&bytes[90..]); + self.ple_batch(&[token]) + } + + fn ple_batch(&mut self, tokens: &[i32]) -> Result<(), String> { + self.ple_batch_with_state(tokens, None) + } + + fn ple_batch_with_state( + &mut self, + tokens: &[i32], + state_out: Option<&Buffer>, + ) -> Result<(), String> { + let read_started = Instant::now(); + let mut history = self.ple_state.history; + let mut requested = Vec::with_capacity(tokens.len() * PLE_HEADS); + for &token in tokens { + requested.extend(ple_rows(&self.ple_contract, history, token)?); + history = advance_ple_history(history, token); + } + let mut packed_stage = vec![0_u8; tokens.len() * PLE_HEADS * 80]; + let mut scales_stage = vec![0_u8; tokens.len() * PLE_HEADS * 10]; + let mut biases_stage = vec![0_u8; tokens.len() * PLE_HEADS * 10]; + for (index, bytes) in self.model.read_ple_rows(&requested)?.iter().enumerate() { + packed_stage[index * 80..][..80].copy_from_slice(&bytes[..80]); + scales_stage[index * 10..][..10].copy_from_slice(&bytes[80..90]); + biases_stage[index * 10..][..10].copy_from_slice(&bytes[90..]); + } + if std::env::var_os("DS4SERVER_QWEN38_PLE_PROFILE").is_some() { + eprintln!( + "DS4SERVER_QWEN38_PLE_READ rows={} us={:.3}", + tokens.len(), + read_started.elapsed().as_secs_f64() * 1_000_000.0 + ); } self.scratch.ple_packed.write(0, &packed_stage)?; self.scratch.ple_scales.write(0, &scales_stage)?; self.scratch.ple_biases.write(0, &biases_stage)?; - let commands = Commands::begin()?; let mut dequant = args(); dequant.u[0] = PLE_HEAD_DIM; - dequant.u[1] = PLE_HEADS as u32; + dequant.u[1] = PLE_HEADS as u32 * tokens.len() as u32; dequant.u[2] = 4; dequant.u[3] = 32; self.dispatch( @@ -364,32 +986,32 @@ impl QwenExecutor { Some(&self.scratch.ple_biases), &[], &dequant, - HIDDEN, + HIDDEN * tokens.len() as u32, 1, )?; let prefix = "language_model.model.layers.1.ple"; - self.bf16_mv( + self.bf16_batch( self.weight(&format!("{prefix}.key_proj.weight"))?, &self.scratch.ple_embedding, &self.scratch.ple_key, HIDDEN, HC_WIDTH, + tokens.len() as u32, )?; - self.bf16_mv( + self.bf16_batch( self.weight(&format!("{prefix}.value_proj.weight"))?, &self.scratch.ple_embedding, &self.scratch.ple_value, HIDDEN, HIDDEN, + tokens.len() as u32, )?; for (input, output, name) in [ (&self.scratch.ple_key, &self.scratch.ple_key, "norm_key"), (&self.scratch.hc, &self.scratch.hc_norm, "norm_query"), ] { - let mut norm = args(); - norm.u[0] = HC_WIDTH; - norm.u[1] = HIDDEN; - norm.f[0] = 1.0e-6; + let mut norm = rms_args(HC_WIDTH, HIDDEN); + norm.u[4] = tokens.len() as u32; self.dispatch( c"kernel_qwen_zero_rms", output, @@ -398,12 +1020,14 @@ impl QwenExecutor { None, &[self.view(self.weight(&format!("{prefix}.{name}.weight"))?)], &norm, - HC, + HC * tokens.len() as u32, 1, )?; } let mut gate = args(); gate.u[0] = HIDDEN; + gate.u[10] = 32 * 4; + gate.u[12] = 640; self.dispatch( c"kernel_qwen_ple_gate", &self.scratch.ple_gated, @@ -413,12 +1037,10 @@ impl QwenExecutor { &[], &gate, HC, - 1, + tokens.len() as u32, )?; - let mut norm = args(); - norm.u[0] = HC_WIDTH; - norm.u[1] = HIDDEN; - norm.f[0] = 1.0e-6; + let mut norm = rms_args(HC_WIDTH, HIDDEN); + norm.u[4] = tokens.len() as u32; self.dispatch( c"kernel_qwen_zero_rms", &self.scratch.ple_norm, @@ -427,24 +1049,26 @@ impl QwenExecutor { None, &[self.view(self.weight(&format!("{prefix}.norm_conv.weight"))?)], &norm, - HC, + HC * tokens.len() as u32, 1, )?; let mut conv = args(); conv.u[0] = HC_WIDTH; - self.dispatch( + conv.u[4] = tokens.len() as u32; + self.dispatch_with_d( c"kernel_qwen_ple_conv", &self.scratch.ple_output, Some(&self.scratch.ple_gated), Some(&self.scratch.ple_norm), Some(&self.ple_state.conv), + Some(state_out.unwrap_or(&self.ple_state.conv)), &[self.view(self.weight(&format!("{prefix}.conv_weight"))?)], &conv, HC_WIDTH, 1, )?; let mut add = args(); - add.u[0] = HC_WIDTH; + add.u[0] = HC_WIDTH * tokens.len() as u32; self.dispatch( c"kernel_qwen_add", &self.scratch.hc_norm, @@ -453,24 +1077,16 @@ impl QwenExecutor { None, &[], &add, - HC_WIDTH, + HC_WIDTH * tokens.len() as u32, 1, )?; - self.scratch.hc.copy_from( - 0, - &self.scratch.hc_norm, - 0, - u64::from(HC_WIDTH) * 4, - "committing Qwen PLE injection", - )?; - commands.finish()?; - self.ple_state.history = advance_ple_history(self.ple_state.history, token); + std::mem::swap(&mut self.scratch.hc, &mut self.scratch.hc_norm); + self.ple_state.history = history; Ok(()) } fn encode_layer(&mut self, layer: usize) -> Result<(), String> { let prefix = format!("language_model.model.layers.{layer}"); - let commands = Commands::begin()?; self.hyper_read(&format!("{prefix}.attn_hyper_connection"))?; match &self.states[layer] { LayerState::Gdn { .. } => self.gdn(&prefix, layer)?, @@ -479,22 +1095,41 @@ impl QwenExecutor { } } self.hyper_write()?; - commands.finish()?; self.encode_moe(&prefix, &layer.to_string()) } - fn encode_moe(&mut self, prefix: &str, layer_name: &str) -> Result<(), String> { - let commands = Commands::begin()?; - self.hyper_read(&format!("{prefix}.mlp_hyper_connection"))?; - self.affine_mv_into( + fn encode_moe(&mut self, prefix: &str, _layer_name: &str) -> Result<(), String> { + self.encode_moe_rows(prefix, 1) + } + + fn encode_moe_rows(&mut self, prefix: &str, rows: u32) -> Result<(), String> { + self.hyper_read_rows(&format!("{prefix}.mlp_hyper_connection"), rows)?; + self.affine_batch_into( &self.affine(&format!("{prefix}.mlp.gate"), HIDDEN, EXPERTS, None)?, &self.scratch.block, &self.scratch.router, HIDDEN, EXPERTS, + rows, + true, + )?; + let mut softmax_args = args(); + softmax_args.u[0] = EXPERTS; + softmax_args.u[12] = EXPERTS / 4; + self.dispatch( + c"kernel_qwen_softmax_precise_512", + &self.scratch.router, + Some(&self.scratch.router), + None, + None, + &[], + &softmax_args, + rows, + 1, )?; let mut route_args = args(); route_args.u[0] = EXPERTS; + route_args.u[12] = EXPERTS; self.dispatch( c"kernel_qwen_route_top10", &self.scratch.route_ids, @@ -503,34 +1138,14 @@ impl QwenExecutor { None, &[], &route_args, - 1, + rows, 1, )?; - commands.finish()?; - - let mut ids = [0_i32; EXPERTS_USED]; - let mut weights = [0.0_f32; EXPERTS_USED]; - self.scratch.route_ids.read_i32(&mut ids)?; - self.scratch.route_weights.read_f32(&mut weights)?; - self.scratch.moe.fill(0.0, HIDDEN.into())?; - let commands = Commands::begin()?; - for (&expert, &weight) in ids.iter().zip(&weights) { - if !(0..EXPERTS as i32).contains(&expert) || !weight.is_finite() || weight < 0.0 { - let mut router = vec![0.0; EXPERTS as usize]; - self.scratch.router.read_f32(&mut router)?; - let non_finite = router.iter().filter(|value| !value.is_finite()).count(); - let mut block = vec![0.0; HIDDEN as usize]; - self.scratch.block.read_f32(&mut block)?; - let block_non_finite = block.iter().filter(|value| !value.is_finite()).count(); - return Err(format!( - "Qwen layer {layer_name} router produced invalid expert {expert} with weight {weight} ({non_finite} non-finite logits, {block_non_finite} non-finite inputs)" - )); - } - self.expert(prefix, expert as u32, weight)?; - } - self.shared_expert(prefix)?; + self.routed_experts_rows(prefix, rows)?; + self.shared_expert_rows(prefix, rows)?; let mut inject_args = args(); inject_args.u[0] = HIDDEN; + inject_args.u[4] = rows; self.dispatch( c"kernel_qwen_hyper_inject", &self.scratch.hc_norm, @@ -539,28 +1154,68 @@ impl QwenExecutor { Some(&self.scratch.injection), &[], &inject_args, - HC_WIDTH, + HC_WIDTH * rows, 1, )?; - self.scratch.hc.copy_from( - 0, - &self.scratch.hc_norm, - 0, - u64::from(HC_WIDTH) * 4, - "committing Qwen MoE hyper streams", - )?; - commands.finish() + std::mem::swap(&mut self.scratch.hc, &mut self.scratch.hc_norm); + Ok(()) } fn hyper_read(&self, prefix: &str) -> Result<(), String> { + self.hyper_read_rows(prefix, 1) + } + + fn hyper_read_rows(&self, prefix: &str, rows: u32) -> Result<(), String> { let norm = self.weight(&format!("{prefix}.hc_norm.weight"))?; let down = self.weight(&format!("{prefix}.input_mix_weight_down.weight"))?; let up = self.weight(&format!("{prefix}.input_mix_weight_up.weight"))?; let inject = self.weight(&format!("{prefix}.block_inject_weight.weight"))?; - let mut rms = args(); - rms.u[0] = HC_WIDTH; - rms.u[1] = HIDDEN; - rms.f[0] = 1.0e-6; + for (weight, shape) in [ + (norm, &[HC_WIDTH as u64][..]), + (down, &[HC_RANK as u64, HC_WIDTH as u64][..]), + (up, &[HC_WIDTH as u64, HC_RANK as u64][..]), + (inject, &[HC as u64, HC_WIDTH as u64][..]), + ] { + if weight.tensor.dtype != "BF16" || weight.tensor.shape != shape { + return Err(format!( + "{} has an incompatible BF16 tensor shape", + weight.tensor.name + )); + } + } + if rows == 1 + && let Some(pack) = self.hyper_packs.get(prefix) + { + let mut values = args(); + values.u[12] = 1024; + self.dispatch_with_d( + c"kernel_qwen_hyper_v3_r1", + &self.scratch.rank, + Some(&self.scratch.hc), + Some(&pack.r1), + Some(&self.scratch.injection), + Some(&self.scratch.hyper_rms), + &[self.view(norm)], + &values, + HYPER_FOLDED_ROWS.div_ceil(32), + rows, + )?; + self.dispatch_with_d( + c"kernel_qwen_hyper_v3_r2", + &self.scratch.block, + Some(&self.scratch.hc), + Some(&pack.r2), + Some(&self.scratch.rank), + Some(&self.scratch.hyper_rms), + &[self.view(norm)], + &values, + HIDDEN.div_ceil(32), + rows, + )?; + return Ok(()); + } + let mut rms = rms_args(HC_WIDTH, HIDDEN); + rms.u[4] = rows; self.dispatch( c"kernel_qwen_zero_rms", &self.scratch.hc_norm, @@ -569,18 +1224,19 @@ impl QwenExecutor { None, &[self.view(norm)], &rms, - HC, + HC * rows, 1, )?; - self.bf16_mv( + self.bf16_batch( down, &self.scratch.hc_norm, &self.scratch.rank, HC_WIDTH, HC_RANK, + rows, )?; let mut unary = args(); - unary.u[0] = HC_RANK; + unary.u[0] = HC_RANK * rows; self.dispatch( c"kernel_qwen_silu_div4", &self.scratch.rank, @@ -589,17 +1245,18 @@ impl QwenExecutor { None, &[], &unary, - HC_RANK, + HC_RANK * rows, 1, )?; - self.bf16_mv( + self.bf16_batch( up, &self.scratch.rank, &self.scratch.hc_mix, HC_RANK, HC_WIDTH, + rows, )?; - unary.u[0] = HC_WIDTH; + unary.u[0] = HC_WIDTH * rows; self.dispatch( c"kernel_qwen_sigmoid", &self.scratch.hc_mix, @@ -608,11 +1265,12 @@ impl QwenExecutor { None, &[], &unary, - HC_WIDTH, + HC_WIDTH * rows, 1, )?; let mut mix = args(); mix.u[0] = HIDDEN; + mix.u[4] = rows; self.dispatch( c"kernel_qwen_hyper_mix", &self.scratch.block, @@ -621,17 +1279,18 @@ impl QwenExecutor { None, &[], &mix, - HIDDEN, + HIDDEN * rows, 1, )?; - self.bf16_mv( + self.bf16_batch( inject, &self.scratch.hc_norm, &self.scratch.injection, HC_WIDTH, HC, + rows, )?; - unary.u[0] = HC; + unary.u[0] = HC * rows; self.dispatch( c"kernel_qwen_sigmoid2_div4", &self.scratch.injection, @@ -640,14 +1299,19 @@ impl QwenExecutor { None, &[], &unary, - HC, + HC * rows, 1, ) } - fn hyper_write(&self) -> Result<(), String> { + fn hyper_write(&mut self) -> Result<(), String> { + self.hyper_write_rows(1) + } + + fn hyper_write_rows(&mut self, rows: u32) -> Result<(), String> { let mut values = args(); values.u[0] = HIDDEN; + values.u[4] = rows; self.dispatch( c"kernel_qwen_hyper_inject", &self.scratch.hc_norm, @@ -656,29 +1320,198 @@ impl QwenExecutor { Some(&self.scratch.injection), &[], &values, - HC_WIDTH, + HC_WIDTH * rows, 1, )?; - self.scratch.hc.copy_from( - 0, - &self.scratch.hc_norm, - 0, - u64::from(HC_WIDTH) * 4, - "committing Qwen attention hyper streams", - ) + std::mem::swap(&mut self.scratch.hc, &mut self.scratch.hc_norm); + Ok(()) } fn gdn(&self, prefix: &str, layer: usize) -> Result<(), String> { let LayerState::Gdn { conv, recurrent } = &self.states[layer] else { return Err("Qwen GDN graph received attention state".into()); }; + let (qkv, controls) = self.gdn_input(prefix)?; + if self.fused_gdn && self.position > 0 { + let mut fused = args(); + fused.u[0] = GDN_VALUE; + fused.u[1] = GDN_VALUE + GDN_HEADS_V; + fused.u[10] = (HEAD_DIM * 4 + 8) * 4; + fused.u[12] = 256; + self.dispatch_with_d( + c"kernel_qwen_gdn_step_fused", + &self.scratch.gdn_out, + Some(qkv), + Some(controls), + Some(recurrent), + Some(conv), + &[ + self.view(self.weight(&format!("{prefix}.linear_attn.conv1d.weight"))?), + self.view(self.weight(&format!("{prefix}.linear_attn.A_log"))?), + self.view(self.weight(&format!("{prefix}.linear_attn.dt_bias"))?), + self.view(self.weight(&format!("{prefix}.linear_attn.norm.weight"))?), + ], + &fused, + 1, + GDN_HEADS_V, + )?; + let next_conv = self + .scratch + .gdn_out + .view(u64::from(GDN_VALUE) * 4, u64::from(GDN_QKV) * 3 * 2)?; + let mut commit = args(); + commit.u[0] = GDN_QKV * 3; + self.dispatch( + c"kernel_qwen_copy_bf16", + conv, + Some(&next_conv), + None, + None, + &[], + &commit, + GDN_QKV * 3, + 1, + )?; + } else { + let mut conv_args = args(); + conv_args.u[0] = GDN_QKV; + conv_args.u[12] = 1024; + self.dispatch( + c"kernel_qwen_gdn_conv_norm", + qkv, + Some(qkv), + Some(conv), + None, + &[self.view(self.weight(&format!("{prefix}.linear_attn.conv1d.weight"))?)], + &conv_args, + GDN_QKV / 1024, + 1, + )?; + let mut control_args = args(); + control_args.u[0] = GDN_VALUE; + control_args.u[1] = GDN_VALUE + GDN_HEADS_V; + self.dispatch( + c"kernel_qwen_gdn_controls", + controls, + None, + None, + None, + &[ + self.view(self.weight(&format!("{prefix}.linear_attn.A_log"))?), + self.view(self.weight(&format!("{prefix}.linear_attn.dt_bias"))?), + ], + &control_args, + GDN_HEADS_V, + 1, + )?; + let mut delta_args = control_args; + delta_args.u[12] = 64; + self.dispatch( + c"kernel_qwen_gdn_delta_packed", + &self.scratch.gdn_out, + Some(qkv), + Some(controls), + Some(recurrent), + &[], + &delta_args, + HEAD_DIM / 16, + GDN_HEADS_V, + )?; + let mut gate_args = args(); + gate_args.u[12] = 32; + self.dispatch( + c"kernel_qwen_gdn_norm_gate", + &self.scratch.gdn_out, + Some(&self.scratch.gdn_out), + Some(controls), + None, + &[self.view(self.weight(&format!("{prefix}.linear_attn.norm.weight"))?)], + &gate_args, + GDN_HEADS_V, + 1, + )?; + } self.affine_mv_into( &self.affine( - &format!("{prefix}.linear_attn.in_proj_qkv"), + &format!("{prefix}.linear_attn.out_proj"), + GDN_VALUE, HIDDEN, - GDN_QKV, None, )?, + &self.scratch.gdn_out, + &self.scratch.hidden, + GDN_VALUE, + HIDDEN, + ) + } + + fn gdn_input<'a>(&'a self, prefix: &str) -> Result<(&'a Buffer, &'a Buffer), String> { + let qkv = self.affine( + &format!("{prefix}.linear_attn.in_proj_qkv"), + HIDDEN, + GDN_QKV, + None, + )?; + if self.fused_gdn_input { + for (name, width) in [ + ("in_proj_z", GDN_VALUE), + ("in_proj_b", GDN_HEADS_V), + ("in_proj_a", GDN_HEADS_V), + ] { + let projection = + self.affine(&format!("{prefix}.linear_attn.{name}"), HIDDEN, width, None)?; + if projection.bits != qkv.bits || projection.group != qkv.group { + return Err(format!( + "{prefix}.linear_attn has incompatible fused input quantization" + )); + } + } + let fused = |suffix: &str, bytes: u64| -> Result { + let name = format!("{prefix}.linear_attn.in_proj_fused.{suffix}"); + let buffer = self + .weights + .get(&name) + .ok_or_else(|| format!("Qwen materialization missed {name}"))?; + Ok(QwenWeightView { + tensor: buffer.raw().cast_const(), + map: std::ptr::null(), + size: bytes, + offset: 0, + bytes, + }) + }; + let packed = u64::from(GDN_FUSED_INPUT) * u64::from(HIDDEN) * u64::from(qkv.bits) / 8; + let parameters = u64::from(GDN_FUSED_INPUT) * u64::from(HIDDEN / qkv.group) * 2; + let mut values = args(); + values.u[0] = HIDDEN; + values.u[1] = GDN_FUSED_INPUT; + values.u[2] = qkv.bits; + values.u[3] = qkv.group; + values.u[11] = 1; + values.u[12] = 64; + self.dispatch( + self.affine_qmv_kernel(&qkv, false, HIDDEN, GDN_FUSED_INPUT)?, + &self.scratch.gdn_fused, + Some(&self.scratch.block), + None, + None, + &[ + fused("weight", packed)?, + fused("scales", parameters)?, + fused("biases", parameters)?, + ], + &values, + GDN_FUSED_INPUT.div_ceil(8), + 1, + )?; + return Ok(( + &self.scratch.gdn_fused_qkv, + &self.scratch.gdn_fused_controls, + )); + } + + self.affine_mv_into( + &qkv, &self.scratch.block, &self.scratch.qkv, HIDDEN, @@ -701,66 +1534,281 @@ impl QwenExecutor { width, )?; } - let mut conv_args = args(); - conv_args.u[0] = GDN_QKV; - self.dispatch( - c"kernel_qwen_conv_silu", - &self.scratch.qkv, - Some(&self.scratch.qkv), - Some(conv), - None, - &[self.view(self.weight(&format!("{prefix}.linear_attn.conv1d.weight"))?)], - &conv_args, - GDN_QKV, + Ok((&self.scratch.qkv, &self.scratch.controls)) + } + + fn gdn_rows( + &self, + prefix: &str, + layer: usize, + rows: u32, + capture_verify: bool, + ) -> Result<(), String> { + let LayerState::Gdn { conv, recurrent } = &self.states[layer] else { + return Err("Qwen GDN graph received attention state".into()); + }; + let verify_capture = if capture_verify { + Some( + self.mtp + .as_ref() + .ok_or("Qwen MTP is not configured")? + .verify + .gdn + .get(layer) + .and_then(Option::as_ref) + .ok_or("Qwen verifier GDN capture layout is invalid")?, + ) + } else { + None + }; + let qkv_out = verify_capture.map_or(&self.scratch.qkv, |capture| &capture.qkv); + let q_out = verify_capture.map_or(&self.scratch.gdn_q, |capture| &capture.q); + let k_out = verify_capture.map_or(&self.scratch.gdn_k, |capture| &capture.k); + let conv_out = verify_capture.map_or(conv, |capture| &capture.conv); + let recurrent_out = verify_capture.map_or(recurrent, |capture| &capture.recurrent); + if self.fused_gdn_input { + self.gdn_fused_input_rows(prefix, rows, qkv_out)?; + } else { + self.affine_batch_into( + &self.affine( + &format!("{prefix}.linear_attn.in_proj_qkv"), + HIDDEN, + GDN_QKV, + None, + )?, + &self.scratch.block, + qkv_out, + HIDDEN, + GDN_QKV, + rows, + true, + )?; + for (name, target, width) in [ + ("in_proj_z", &self.scratch.gdn_z, GDN_VALUE), + ("in_proj_b", &self.scratch.gdn_b, GDN_HEADS_V), + ("in_proj_a", &self.scratch.gdn_a, GDN_HEADS_V), + ] { + self.affine_batch_into( + &self.affine(&format!("{prefix}.linear_attn.{name}"), HIDDEN, width, None)?, + &self.scratch.block, + target, + HIDDEN, + width, + rows, + true, + )?; + } + let mut pack = args(); + pack.u[4] = rows; + self.dispatch_with_d( + c"kernel_qwen_pack_gdn_controls", + &self.scratch.controls, + Some(&self.scratch.gdn_z), + Some(&self.scratch.gdn_b), + Some(&self.scratch.gdn_a), + None, + &[], + &pack, + GDN_CONTROLS * rows, + 1, + )?; + } + let mut conv_norm = args(); + conv_norm.u[4] = rows; + conv_norm.u[12] = 1024; + self.dispatch_with_d( + c"kernel_qwen_gdn_conv_norm_rows", + q_out, + Some(qkv_out), + Some(k_out), + Some(&self.scratch.gdn_out), + Some(conv_out), + &[ + buffer_weight_view(conv, u64::from(GDN_QKV) * 3 * 2), + self.view(self.weight(&format!("{prefix}.linear_attn.conv1d.weight"))?), + ], + &conv_norm, + GDN_QKV / 1024, 1, )?; - let mut step = args(); - step.u[0] = HEAD_DIM; - step.u[1] = GDN_HEADS_K; - step.u[2] = GDN_HEADS_V; - step.u[3] = GDN_VALUE; - step.u[4] = GDN_VALUE + GDN_HEADS_V; - step.f[0] = 1.0e-6; + let mut controls = args(); + controls.u[0] = GDN_VALUE; + controls.u[1] = GDN_VALUE + GDN_HEADS_V; + controls.u[4] = rows; self.dispatch( - c"kernel_qwen_gdn_step", - &self.scratch.gdn_raw, - Some(&self.scratch.gdn_out), - Some(&self.scratch.controls), - Some(recurrent), + c"kernel_qwen_gdn_controls", + &self.scratch.controls, + None, + None, + None, &[ self.view(self.weight(&format!("{prefix}.linear_attn.A_log"))?), self.view(self.weight(&format!("{prefix}.linear_attn.dt_bias"))?), ], - &step, - HEAD_DIM, + &controls, + GDN_HEADS_V * rows, + 1, + )?; + if let Some(capture) = verify_capture { + for (target, source, width, message) in [ + ( + &capture.v, + &self.scratch.gdn_out, + GDN_VALUE, + "capturing Qwen verifier GDN value rows", + ), + ( + &capture.controls, + &self.scratch.controls, + GDN_CONTROLS, + "capturing Qwen verifier GDN control rows", + ), + ] { + target.copy_from(0, source, 0, u64::from(rows * width) * 4, message)?; + } + } + let mut delta = controls; + delta.u[12] = 64; + self.dispatch_with_d( + c"kernel_qwen_gdn_delta_rows", + &self.scratch.gdn_out, + Some(q_out), + Some(k_out), + Some(&self.scratch.controls), + Some(recurrent_out), + &[buffer_weight_view( + recurrent, + u64::from(GDN_HEADS_V) * HEAD_DIM as u64 * HEAD_DIM as u64 * 4, + )], + &delta, + HEAD_DIM / 16, GDN_HEADS_V, )?; - let mut gate = args(); - gate.u[0] = HEAD_DIM; - gate.u[1] = GDN_HEADS_V; - gate.f[0] = 1.0e-6; + let mut norm = controls; + norm.u[12] = 32; self.dispatch( c"kernel_qwen_gdn_norm_gate", - &self.scratch.gdn_out, - Some(&self.scratch.gdn_raw), + &self.scratch.gdn_z, + Some(&self.scratch.gdn_out), Some(&self.scratch.controls), None, &[self.view(self.weight(&format!("{prefix}.linear_attn.norm.weight"))?)], - &gate, + &norm, GDN_HEADS_V, - 1, + rows, )?; - self.affine_mv_into( + self.affine_batch_into( &self.affine( &format!("{prefix}.linear_attn.out_proj"), GDN_VALUE, HIDDEN, None, )?, - &self.scratch.gdn_out, + &self.scratch.gdn_z, &self.scratch.hidden, GDN_VALUE, HIDDEN, + rows, + true, + ) + } + + fn gdn_fused_input_rows( + &self, + prefix: &str, + rows: u32, + qkv_out: &Buffer, + ) -> Result<(), String> { + let qkv = self.affine( + &format!("{prefix}.linear_attn.in_proj_qkv"), + HIDDEN, + GDN_QKV, + None, + )?; + for (name, width) in [ + ("in_proj_z", GDN_VALUE), + ("in_proj_b", GDN_HEADS_V), + ("in_proj_a", GDN_HEADS_V), + ] { + let projection = + self.affine(&format!("{prefix}.linear_attn.{name}"), HIDDEN, width, None)?; + if projection.bits != qkv.bits || projection.group != qkv.group { + return Err(format!( + "{prefix}.linear_attn has incompatible fused input quantization" + )); + } + } + let fused = |suffix: &str, bytes: u64| -> Result { + let name = format!("{prefix}.linear_attn.in_proj_fused.{suffix}"); + let buffer = self + .weights + .get(&name) + .ok_or_else(|| format!("Qwen materialization missed {name}"))?; + Ok(QwenWeightView { + tensor: buffer.raw().cast_const(), + map: std::ptr::null(), + size: bytes, + offset: 0, + bytes, + }) + }; + let packed = u64::from(GDN_FUSED_INPUT) * u64::from(HIDDEN) * u64::from(qkv.bits) / 8; + let parameters = u64::from(GDN_FUSED_INPUT) * u64::from(HIDDEN / qkv.group) * 2; + let mut values = args(); + values.u[0] = HIDDEN; + values.u[1] = GDN_FUSED_INPUT; + values.u[2] = qkv.bits; + values.u[3] = qkv.group; + values.u[4] = rows; + values.u[11] = 1; + values.u[12] = 64; + let views = [ + fused("weight", packed)?, + fused("scales", parameters)?, + fused("biases", parameters)?, + ]; + if rows > VERIFY_ROWS as u32 { + if unsafe { ds4_gpu_metal4_tensor_api_enabled() } == 0 { + return Err("Qwen batched prefill requires the Metal 4 tensor API".to_string()); + } + values.u[10] = 0; + values.u[12] = 128; + self.dispatch( + self.affine_qmm_kernel(&qkv, rows)?, + &self.scratch.gdn_fused, + Some(&self.scratch.block), + None, + None, + &views, + &values, + GDN_FUSED_INPUT.div_ceil(64), + rows.div_ceil(if rows <= 32 { 32 } else { 64 }), + )?; + } else { + self.dispatch( + self.affine_qmv_batch_kernel(&qkv, HIDDEN, GDN_FUSED_INPUT)?, + &self.scratch.gdn_fused, + Some(&self.scratch.block), + None, + None, + &views, + &values, + GDN_FUSED_INPUT.div_ceil(8), + 1, + )?; + } + let mut unpack = args(); + unpack.u[4] = rows; + self.dispatch( + c"kernel_qwen_unpack_gdn_inputs", + qkv_out, + Some(&self.scratch.gdn_fused), + Some(&self.scratch.controls), + None, + &[], + &unpack, + GDN_FUSED_INPUT * rows, + 1, ) } @@ -773,18 +1821,7 @@ impl QwenExecutor { else { return Err("Qwen attention graph received GDN state".into()); }; - self.affine_mv_into( - &self.affine( - &format!("{prefix}.self_attn.indexer.index_qk_proj"), - HIDDEN, - QSA_WIDTH, - None, - )?, - &self.scratch.block, - &self.scratch.qsa_qk, - HIDDEN, - QSA_WIDTH, - )?; + let (q_packed, k, v, qsa_qk) = self.attention_input(prefix)?; let mut index = args(); index.u[0] = QSA_DIM; index.u[1] = QSA_HEADS; @@ -792,7 +1829,7 @@ impl QwenExecutor { self.dispatch( c"kernel_qwen_qsa_store_raw", qsa_raw, - Some(&self.scratch.qsa_qk), + Some(qsa_qk), Some(&self.scratch.qsa_q), None, &[], @@ -828,49 +1865,13 @@ impl QwenExecutor { 1, )?; } - self.affine_mv_into( - &self.affine( - &format!("{prefix}.self_attn.q_proj"), - HIDDEN, - ATTN_WIDTH * 2, - None, - )?, - &self.scratch.block, - &self.scratch.q_packed, - HIDDEN, - ATTN_WIDTH * 2, - )?; - self.affine_mv_into( - &self.affine( - &format!("{prefix}.self_attn.k_proj"), - HIDDEN, - ATTN_KV_WIDTH, - None, - )?, - &self.scratch.block, - &self.scratch.k, - HIDDEN, - ATTN_KV_WIDTH, - )?; - self.affine_mv_into( - &self.affine( - &format!("{prefix}.self_attn.v_proj"), - HIDDEN, - ATTN_KV_WIDTH, - None, - )?, - &self.scratch.block, - &self.scratch.v, - HIDDEN, - ATTN_KV_WIDTH, - )?; let mut split = args(); split.u[0] = ATTN_HEADS; split.u[1] = ATTN_DIM; self.dispatch( c"kernel_qwen_split_q_gate", &self.scratch.q, - Some(&self.scratch.q_packed), + Some(q_packed), Some(&self.scratch.q_gate), None, &[], @@ -886,7 +1887,7 @@ impl QwenExecutor { position, )?; self.head_norm_rope( - &self.scratch.k, + k, &self.scratch.k_rope, self.weight(&format!("{prefix}.self_attn.k_norm.weight"))?, ATTN_KV_HEADS, @@ -899,7 +1900,7 @@ impl QwenExecutor { c"kernel_qwen_store_kv_bf16", kv, Some(&self.scratch.k_rope), - Some(&self.scratch.v), + Some(v), None, &[], &store, @@ -934,6 +1935,424 @@ impl QwenExecutor { ) } + fn attention_input<'a>( + &'a self, + prefix: &str, + ) -> Result<(&'a Buffer, &'a Buffer, &'a Buffer, &'a Buffer), String> { + let q = self.affine( + &format!("{prefix}.self_attn.q_proj"), + HIDDEN, + ATTN_WIDTH * 2, + None, + )?; + if self.fused_qsa_input + && self + .weights + .contains_key(&format!("{prefix}.self_attn.qkv_fused.weight")) + { + for (name, width) in [("k_proj", ATTN_KV_WIDTH), ("v_proj", ATTN_KV_WIDTH)] { + let projection = + self.affine(&format!("{prefix}.self_attn.{name}"), HIDDEN, width, None)?; + if projection.bits != q.bits || projection.group != q.group { + return Err(format!( + "{prefix}.self_attn has incompatible fused input quantization" + )); + } + } + let indexer = self.affine( + &format!("{prefix}.self_attn.indexer.index_qk_proj"), + HIDDEN, + QSA_WIDTH, + None, + )?; + let include_indexer = indexer.bits == q.bits && indexer.group == q.group; + let out_dim = if include_indexer { + QSA_FUSED_INPUT + } else { + QSA_FUSED_QKV + }; + let fused = |suffix: &str, bytes: u64| -> Result { + let name = format!("{prefix}.self_attn.qkv_fused.{suffix}"); + let buffer = self + .weights + .get(&name) + .ok_or_else(|| format!("Qwen materialization missed {name}"))?; + Ok(QwenWeightView { + tensor: buffer.raw().cast_const(), + map: std::ptr::null(), + size: bytes, + offset: 0, + bytes, + }) + }; + let packed = u64::from(out_dim) * u64::from(HIDDEN) * u64::from(q.bits) / 8; + let parameters = u64::from(out_dim) * u64::from(HIDDEN / q.group) * 2; + let mut values = args(); + values.u[0] = HIDDEN; + values.u[1] = out_dim; + values.u[2] = q.bits; + values.u[3] = q.group; + values.u[11] = 1; + values.u[12] = 64; + self.dispatch( + self.affine_qmv_kernel(&q, false, HIDDEN, out_dim)?, + &self.scratch.qsa_fused, + Some(&self.scratch.block), + None, + None, + &[ + fused("weight", packed)?, + fused("scales", parameters)?, + fused("biases", parameters)?, + ], + &values, + out_dim.div_ceil(8), + 1, + )?; + if !include_indexer { + self.affine_mv_into( + &indexer, + &self.scratch.block, + &self.scratch.qsa_qk, + HIDDEN, + QSA_WIDTH, + )?; + } + return Ok(( + &self.scratch.qsa_fused_q, + &self.scratch.qsa_fused_k, + &self.scratch.qsa_fused_v, + if include_indexer { + &self.scratch.qsa_fused_indexer + } else { + &self.scratch.qsa_qk + }, + )); + } + + for (name, output, width) in [ + ("q_proj", &self.scratch.q_packed, ATTN_WIDTH * 2), + ("k_proj", &self.scratch.k, ATTN_KV_WIDTH), + ("v_proj", &self.scratch.v, ATTN_KV_WIDTH), + ("indexer.index_qk_proj", &self.scratch.qsa_qk, QSA_WIDTH), + ] { + self.affine_mv_into( + &self.affine(&format!("{prefix}.self_attn.{name}"), HIDDEN, width, None)?, + &self.scratch.block, + output, + HIDDEN, + width, + )?; + } + Ok(( + &self.scratch.q_packed, + &self.scratch.k, + &self.scratch.v, + &self.scratch.qsa_qk, + )) + } + + fn attention_rows( + &self, + prefix: &str, + state: &LayerState, + position: u32, + rows: u32, + ) -> Result<(), String> { + let LayerState::Attention { + kv, + qsa_raw, + qsa_pooled, + } = state + else { + return Err("Qwen attention graph received GDN state".into()); + }; + self.affine_batch_into( + &self.affine( + &format!("{prefix}.self_attn.indexer.index_qk_proj"), + HIDDEN, + QSA_WIDTH, + None, + )?, + &self.scratch.block, + &self.scratch.qsa_qk, + HIDDEN, + QSA_WIDTH, + rows, + true, + )?; + let mut index = args(); + index.u[0] = QSA_DIM; + index.u[1] = QSA_HEADS; + index.u[2] = position; + index.u[4] = rows; + self.dispatch( + c"kernel_qwen_qsa_store_raw", + qsa_raw, + Some(&self.scratch.qsa_qk), + Some(&self.scratch.qsa_q), + None, + &[], + &index, + QSA_WIDTH, + rows, + )?; + self.indexer_norm_rope_rows( + &self.scratch.qsa_q, + &self.scratch.qsa_qk, + self.weight(&format!("{prefix}.self_attn.indexer.q_layernorm.weight"))?, + QSA_HEADS, + position, + rows, + )?; + for row in 0..rows { + let row_position = position + row; + if (row_position + 1).is_multiple_of(QSA_RATIO) { + let mut pool = args(); + pool.u[0] = QSA_DIM; + pool.u[1] = row_position / QSA_RATIO; + pool.u[2] = QSA_RATIO; + pool.u[3] = row_position + 1 - QSA_RATIO; + pool.f[0] = 1.0e-6; + pool.f[1] = 10_000_000.0; + self.dispatch( + c"kernel_qwen_qsa_pool_key", + qsa_pooled, + Some(qsa_raw), + None, + None, + &[self.view( + self.weight(&format!("{prefix}.self_attn.indexer.k_layernorm.weight"))?, + )], + &pool, + QSA_DIM, + 1, + )?; + } + } + for (name, output, width) in [ + ("q_proj", &self.scratch.q_packed, ATTN_WIDTH * 2), + ("k_proj", &self.scratch.k, ATTN_KV_WIDTH), + ("v_proj", &self.scratch.v, ATTN_KV_WIDTH), + ] { + self.affine_batch_into( + &self.affine(&format!("{prefix}.self_attn.{name}"), HIDDEN, width, None)?, + &self.scratch.block, + output, + HIDDEN, + width, + rows, + true, + )?; + } + let mut split = args(); + split.u[0] = ATTN_HEADS; + split.u[1] = ATTN_DIM; + split.u[4] = rows; + self.dispatch( + c"kernel_qwen_split_q_gate", + &self.scratch.q, + Some(&self.scratch.q_packed), + Some(&self.scratch.q_gate), + None, + &[], + &split, + ATTN_WIDTH * rows, + 1, + )?; + self.head_norm_rope_rows( + &self.scratch.q, + &self.scratch.attention, + self.weight(&format!("{prefix}.self_attn.q_norm.weight"))?, + ATTN_HEADS, + position, + rows, + )?; + self.head_norm_rope_rows( + &self.scratch.k, + &self.scratch.k_rope, + self.weight(&format!("{prefix}.self_attn.k_norm.weight"))?, + ATTN_KV_HEADS, + position, + rows, + )?; + let mut store = args(); + store.u[0] = ATTN_KV_WIDTH; + store.u[1] = position; + store.u[4] = rows; + self.dispatch( + c"kernel_qwen_store_kv_bf16", + kv, + Some(&self.scratch.k_rope), + Some(&self.scratch.v), + None, + &[], + &store, + ATTN_KV_WIDTH, + rows, + )?; + if position + rows <= DENSE_BUDGET { + let mut dense = args(); + dense.u[0] = ATTN_HEADS; + dense.u[1] = ATTN_KV_HEADS; + dense.u[2] = ATTN_DIM; + dense.u[3] = position + rows; + dense.u[4] = rows; + dense.u[5] = position; + dense.u[12] = 1024; + if rows * (ATTN_HEADS / ATTN_KV_HEADS) > 32 { + dense.u[12] = 32; + self.dispatch( + c"kernel_qwen_attention_fallback_scores", + &self.scratch.attention_scores, + Some(&self.scratch.attention), + Some(kv), + None, + &[], + &dense, + ATTN_HEADS, + rows, + )?; + dense.u[12] = (position + rows).div_ceil(4).div_ceil(32) * 32; + self.dispatch( + c"kernel_qwen_attention_fallback_softmax", + &self.scratch.attention_scores, + Some(&self.scratch.attention_scores), + None, + None, + &[], + &dense, + ATTN_HEADS * rows, + 1, + )?; + dense.u[12] = 32; + self.dispatch( + c"kernel_qwen_attention_fallback_output", + &self.scratch.q, + Some(&self.scratch.attention_scores), + Some(kv), + None, + &[], + &dense, + ATTN_HEADS, + rows, + )?; + return self.finish_attention_rows(prefix, rows); + } + let mut mask = args(); + mask.u[0] = position + rows; + mask.u[4] = rows; + mask.u[5] = position; + self.dispatch( + c"kernel_qwen_causal_mask", + &self.scratch.causal_mask, + None, + None, + None, + &[], + &mask, + position + rows, + rows, + )?; + let mut prepare = args(); + prepare.u[0] = ATTN_KV_WIDTH; + prepare.u[1] = ATTN_DIM; + prepare.u[3] = position + rows; + self.dispatch( + c"kernel_qwen_prepare_dense_kv", + &self.scratch.dense_keys, + Some(kv), + Some(&self.scratch.dense_values), + None, + &[], + &prepare, + ATTN_KV_WIDTH, + position + rows, + )?; + let mut query = args(); + query.u[0] = ATTN_WIDTH * rows; + self.dispatch( + c"kernel_qwen_float_to_bf16", + &self.scratch.dense_query, + Some(&self.scratch.attention), + None, + None, + &[], + &query, + ATTN_WIDTH * rows, + 1, + )?; + self.dispatch_with_d( + c"kernel_qwen_dense_attention_rows", + &self.scratch.q, + Some(&self.scratch.dense_query), + Some(&self.scratch.dense_keys), + Some(&self.scratch.dense_values), + Some(&self.scratch.causal_mask), + &[], + &dense, + ATTN_HEADS, + rows, + )?; + } else { + let qsa_blocks = qsa_block_capacity(self.context)?; + for row in 0..rows { + let f32_view = |buffer: &Buffer, width: u32| { + buffer.view(u64::from(row) * u64::from(width) * 4, u64::from(width) * 4) + }; + let attention = f32_view(&self.scratch.attention, ATTN_WIDTH)?; + let q = f32_view(&self.scratch.q, ATTN_WIDTH)?; + let qsa_query = f32_view(&self.scratch.qsa_qk, QSA_HEADS * QSA_DIM)?; + let scores = f32_view(&self.scratch.qsa_scores, qsa_blocks)?; + let selected = self.scratch.qsa_selected.view( + u64::from(row) * u64::from(QSA_TOP_K) * 4, + u64::from(QSA_TOP_K) * 4, + )?; + self.attend_row( + kv, + qsa_pooled, + position + row, + &attention, + &q, + &qsa_query, + &scores, + &selected, + )?; + } + } + self.finish_attention_rows(prefix, rows) + } + + fn finish_attention_rows(&self, prefix: &str, rows: u32) -> Result<(), String> { + let mut gate = args(); + gate.u[0] = ATTN_WIDTH * rows; + self.dispatch( + c"kernel_qwen_gate_attention", + &self.scratch.attention, + Some(&self.scratch.q), + Some(&self.scratch.q_gate), + None, + &[], + &gate, + ATTN_WIDTH * rows, + 1, + )?; + self.affine_batch_into( + &self.affine( + &format!("{prefix}.self_attn.o_proj"), + ATTN_WIDTH, + HIDDEN, + None, + )?, + &self.scratch.attention, + &self.scratch.hidden, + ATTN_WIDTH, + HIDDEN, + rows, + true, + ) + } + fn indexer_norm_rope( &self, input: &Buffer, @@ -941,12 +2360,25 @@ impl QwenExecutor { weight: Weight<'_>, heads: u32, position: u32, + ) -> Result<(), String> { + self.indexer_norm_rope_rows(input, output, weight, heads, position, 1) + } + + fn indexer_norm_rope_rows( + &self, + input: &Buffer, + output: &Buffer, + weight: Weight<'_>, + heads: u32, + position: u32, + rows: u32, ) -> Result<(), String> { let mut values = args(); values.u[0] = QSA_DIM; values.u[1] = 64; values.u[2] = heads; values.u[3] = position; + values.u[12] = QSA_DIM.div_ceil(4).div_ceil(32) * 32; values.f[0] = 1.0e-6; values.f[1] = 10_000_000.0; self.dispatch( @@ -957,29 +2389,56 @@ impl QwenExecutor { None, &[self.view(weight)], &values, - QSA_DIM, heads, + rows, ) } fn attend(&self, kv: &Buffer, qsa_pooled: &Buffer, position: u32) -> Result<(), String> { + self.attend_row( + kv, + qsa_pooled, + position, + &self.scratch.attention, + &self.scratch.q, + &self.scratch.qsa_qk, + &self.scratch.qsa_scores, + &self.scratch.qsa_selected, + ) + } + + #[allow(clippy::too_many_arguments)] + fn attend_row( + &self, + kv: &Buffer, + qsa_pooled: &Buffer, + position: u32, + query: &Buffer, + output: &Buffer, + qsa_query: &Buffer, + qsa_scores: &Buffer, + qsa_selected: &Buffer, + ) -> Result<(), String> { let tokens = position + 1; let mut values = args(); values.u[0] = ATTN_HEADS; values.u[1] = ATTN_KV_HEADS; values.u[2] = ATTN_DIM; values.u[3] = tokens; + values.u[4] = 1; + values.u[5] = position; if tokens <= DENSE_BUDGET { + values.u[12] = 1024; return self.dispatch( c"kernel_qwen_dense_attention", - &self.scratch.q, - Some(&self.scratch.attention), + output, + Some(query), Some(kv), None, &[], &values, - ATTN_DIM, ATTN_HEADS, + 1, ); } @@ -991,8 +2450,8 @@ impl QwenExecutor { score.f[0] = (QSA_DIM as f32).sqrt().recip(); self.dispatch( c"kernel_qwen_qsa_scores", - &self.scratch.qsa_scores, - Some(&self.scratch.qsa_qk), + qsa_scores, + Some(qsa_query), Some(qsa_pooled), None, &[], @@ -1003,8 +2462,8 @@ impl QwenExecutor { call( unsafe { ds4_gpu_indexer_topk_tensor( - self.scratch.qsa_selected.raw(), - self.scratch.qsa_scores.raw(), + qsa_selected.raw(), + qsa_scores.raw(), complete_blocks, 1, QSA_TOP_K, @@ -1016,7 +2475,7 @@ impl QwenExecutor { order.u[0] = QSA_TOP_K; self.dispatch( c"kernel_qwen_qsa_sort_blocks", - &self.scratch.qsa_selected, + qsa_selected, None, None, None, @@ -1029,16 +2488,17 @@ impl QwenExecutor { values.u[5] = QSA_RATIO; values.u[6] = complete_blocks * QSA_RATIO; values.u[7] = tokens - complete_blocks * QSA_RATIO; + values.u[12] = 256; self.dispatch( c"kernel_qwen_sparse_attention", - &self.scratch.q, - Some(&self.scratch.attention), + output, + Some(query), Some(kv), - Some(&self.scratch.qsa_selected), + Some(qsa_selected), &[], &values, - ATTN_DIM, ATTN_HEADS, + 1, ) } @@ -1049,12 +2509,25 @@ impl QwenExecutor { weight: Weight<'_>, heads: u32, position: u32, + ) -> Result<(), String> { + self.head_norm_rope_rows(input, output, weight, heads, position, 1) + } + + fn head_norm_rope_rows( + &self, + input: &Buffer, + output: &Buffer, + weight: Weight<'_>, + heads: u32, + position: u32, + rows: u32, ) -> Result<(), String> { let mut values = args(); values.u[0] = ATTN_DIM; values.u[1] = 64; values.u[2] = heads; values.u[3] = position; + values.u[12] = ATTN_DIM.div_ceil(4).div_ceil(32) * 32; values.f[0] = 1.0e-6; values.f[1] = 10_000_000.0; self.dispatch( @@ -1065,31 +2538,55 @@ impl QwenExecutor { None, &[self.view(weight)], &values, - ATTN_DIM, heads, + rows, ) } - fn expert(&self, prefix: &str, expert: u32, weight: f32) -> Result<(), String> { - for (name, target) in [ - ("gate_proj", &self.scratch.gate), - ("up_proj", &self.scratch.up), - ] { - self.affine_mv_into( - &self.affine( - &format!("{prefix}.mlp.switch_mlp.{name}"), - HIDDEN, - EXPERT_WIDTH, - Some(expert), - )?, + #[cfg(test)] + fn routed_experts(&self, prefix: &str) -> Result<(), String> { + self.routed_experts_rows(prefix, 1) + } + + fn routed_experts_rows(&self, prefix: &str, rows: u32) -> Result<(), String> { + let gate = self.affine_table( + &format!("{prefix}.mlp.switch_mlp.gate_proj"), + HIDDEN, + EXPERT_WIDTH, + )?; + let up = self.affine_table( + &format!("{prefix}.mlp.switch_mlp.up_proj"), + HIDDEN, + EXPERT_WIDTH, + )?; + if self.packed_moe { + self.affine_pair_batch_into( + &gate, + &up, &self.scratch.block, - target, + &self.scratch.gate, + &self.scratch.up, HIDDEN, EXPERT_WIDTH, + true, + rows, + false, )?; + } else { + for (table, target) in [(&gate, &self.scratch.gate), (&up, &self.scratch.up)] { + self.affine_gather_batch_into( + table, + &self.scratch.block, + target, + HIDDEN, + EXPERT_WIDTH, + rows, + false, + )?; + } } let mut swiglu = args(); - swiglu.u[0] = EXPERT_WIDTH; + swiglu.u[0] = EXPERT_WIDTH * EXPERTS_USED as u32 * rows; self.dispatch( c"kernel_qwen_swiglu", &self.scratch.mid, @@ -1098,57 +2595,85 @@ impl QwenExecutor { None, &[], &swiglu, - EXPERT_WIDTH, + EXPERT_WIDTH * EXPERTS_USED as u32 * rows, 1, )?; - self.affine_mv_into( - &self.affine( - &format!("{prefix}.mlp.switch_mlp.down_proj"), - EXPERT_WIDTH, - HIDDEN, - Some(expert), - )?, + let table = self.affine_table( + &format!("{prefix}.mlp.switch_mlp.down_proj"), + EXPERT_WIDTH, + HIDDEN, + )?; + self.affine_gather_batch_into( + &table, &self.scratch.mid, &self.scratch.expert, EXPERT_WIDTH, HIDDEN, + rows, + true, )?; - let mut accumulate = args(); - accumulate.u[0] = HIDDEN; - accumulate.f[0] = weight; + let mut sum = args(); + sum.u[0] = HIDDEN * rows; + sum.u[1] = HIDDEN; self.dispatch( - c"kernel_qwen_accumulate", + c"kernel_qwen_weighted_sum10", &self.scratch.moe, Some(&self.scratch.expert), - None, + Some(&self.scratch.route_weights), None, &[], - &accumulate, - HIDDEN, + &sum, + HIDDEN * rows, 1, ) } + #[cfg(test)] fn shared_expert(&self, prefix: &str) -> Result<(), String> { - for (name, target) in [ - ("gate_proj", &self.scratch.gate), - ("up_proj", &self.scratch.up), - ] { - self.affine_mv_into( - &self.affine( - &format!("{prefix}.mlp.shared_expert.{name}"), - HIDDEN, - EXPERT_WIDTH, - None, - )?, + self.shared_expert_rows(prefix, 1) + } + + fn shared_expert_rows(&self, prefix: &str, rows: u32) -> Result<(), String> { + let gate = self.affine( + &format!("{prefix}.mlp.shared_expert.gate_proj"), + HIDDEN, + EXPERT_WIDTH, + None, + )?; + let up = self.affine( + &format!("{prefix}.mlp.shared_expert.up_proj"), + HIDDEN, + EXPERT_WIDTH, + None, + )?; + if self.packed_moe && rows <= VERIFY_ROWS as u32 { + self.affine_pair_batch_into( + &gate, + &up, &self.scratch.block, - target, + &self.scratch.gate, + &self.scratch.up, HIDDEN, EXPERT_WIDTH, + false, + rows, + false, )?; + } else { + for (weight, target) in [(&gate, &self.scratch.gate), (&up, &self.scratch.up)] { + self.affine_batch_into( + weight, + &self.scratch.block, + target, + HIDDEN, + EXPERT_WIDTH, + rows, + true, + )?; + } } let mut swiglu = args(); - swiglu.u[0] = EXPERT_WIDTH; + swiglu.u[0] = EXPERT_WIDTH * rows; self.dispatch( c"kernel_qwen_swiglu", &self.scratch.mid, @@ -1157,10 +2682,10 @@ impl QwenExecutor { None, &[], &swiglu, - EXPERT_WIDTH, + EXPERT_WIDTH * rows, 1, )?; - self.affine_mv_into( + self.affine_batch_into( &self.affine( &format!("{prefix}.mlp.shared_expert.down_proj"), EXPERT_WIDTH, @@ -1171,16 +2696,21 @@ impl QwenExecutor { &self.scratch.shared, EXPERT_WIDTH, HIDDEN, + rows, + true, )?; - self.affine_mv_into( + self.affine_batch_into( &self.affine(&format!("{prefix}.mlp.shared_expert_gate"), HIDDEN, 1, None)?, &self.scratch.block, &self.scratch.gate, HIDDEN, 1, + rows, + true, )?; let mut accumulate = args(); accumulate.u[0] = HIDDEN; + accumulate.u[4] = rows; self.dispatch( c"kernel_qwen_accumulate_sigmoid_scalar", &self.scratch.moe, @@ -1189,33 +2719,38 @@ impl QwenExecutor { None, &[], &accumulate, - HIDDEN, + HIDDEN * rows, 1, ) } fn final_output(&mut self) -> Result<(), String> { - let commands = Commands::begin()?; - self.mix_output("language_model.model.hyper_connection_mixer")?; - self.affine_mv_into( + self.final_output_rows(1) + } + + fn final_output_rows(&mut self, rows: u32) -> Result<(), String> { + self.mix_output_rows("language_model.model.hyper_connection_mixer", rows)?; + self.affine_batch_into( &self.affine("language_model.lm_head", HIDDEN, VOCAB, None)?, &self.scratch.block, &self.scratch.logits, HIDDEN, VOCAB, - )?; - commands.finish()?; - self.scratch.logits.read_f32(&mut self.logits) + rows, + true, + ) } fn mix_output(&self, prefix: &str) -> Result<(), String> { + self.mix_output_rows(prefix, 1) + } + + fn mix_output_rows(&self, prefix: &str, rows: u32) -> Result<(), String> { let norm = self.weight(&format!("{prefix}.hc_norm.weight"))?; let down = self.weight(&format!("{prefix}.input_mix_weight_down.weight"))?; let up = self.weight(&format!("{prefix}.input_mix_weight_up.weight"))?; - let mut rms = args(); - rms.u[0] = HC_WIDTH; - rms.u[1] = HIDDEN; - rms.f[0] = 1.0e-6; + let mut rms = rms_args(HC_WIDTH, HIDDEN); + rms.u[4] = rows; self.dispatch( c"kernel_qwen_zero_rms", &self.scratch.hc_norm, @@ -1224,18 +2759,19 @@ impl QwenExecutor { None, &[self.view(norm)], &rms, - HC, + HC * rows, 1, )?; - self.bf16_mv( + self.bf16_batch( down, &self.scratch.hc_norm, &self.scratch.rank, HC_WIDTH, HC_RANK, + rows, )?; let mut unary = args(); - unary.u[0] = HC_RANK; + unary.u[0] = HC_RANK * rows; self.dispatch( c"kernel_qwen_silu_div4", &self.scratch.rank, @@ -1244,17 +2780,18 @@ impl QwenExecutor { None, &[], &unary, - HC_RANK, + HC_RANK * rows, 1, )?; - self.bf16_mv( + self.bf16_batch( up, &self.scratch.rank, &self.scratch.hc_mix, HC_RANK, HC_WIDTH, + rows, )?; - unary.u[0] = HC_WIDTH; + unary.u[0] = HC_WIDTH * rows; self.dispatch( c"kernel_qwen_sigmoid", &self.scratch.hc_mix, @@ -1263,11 +2800,12 @@ impl QwenExecutor { None, &[], &unary, - HC_WIDTH, + HC_WIDTH * rows, 1, )?; let mut mix = args(); mix.u[0] = HIDDEN; + mix.u[4] = rows; self.dispatch( c"kernel_qwen_hyper_mix", &self.scratch.block, @@ -1276,7 +2814,7 @@ impl QwenExecutor { None, &[], &mix, - HIDDEN, + HIDDEN * rows, 1, ) } @@ -1346,6 +2884,174 @@ impl QwenExecutor { }) } + fn affine_table(&self, prefix: &str, in_dim: u32, out_dim: u32) -> Result, String> { + let mut table = self.affine(prefix, in_dim, out_dim, Some(0))?; + table.packed.expert = None; + table.scales.expert = None; + table.biases.expert = None; + Ok(table) + } + + fn affine_qmv_kernel( + &self, + weight: &Affine<'_>, + gather: bool, + in_dim: u32, + out_dim: u32, + ) -> Result<&'static CStr, String> { + let alignment = match weight.bits { + 4 => 512, + 8 => 256, + _ => 0, + }; + let fast = alignment != 0 + && out_dim.is_multiple_of(8) + && in_dim.is_multiple_of(alignment) + && self.affine_is_direct(weight); + match (gather, fast, weight.bits, weight.group) { + (false, true, 4, 32) => Ok(c"kernel_qwen_affine_qmv_fast_b4g32"), + (false, true, 4, 64) => Ok(c"kernel_qwen_affine_qmv_fast_b4g64"), + (false, true, 8, 64) => Ok(c"kernel_qwen_affine_qmv_fast_b8g64"), + (true, true, 4, 32) => Ok(c"kernel_qwen_affine_gather_qmv_fast_b4g32"), + (true, true, 4, 64) => Ok(c"kernel_qwen_affine_gather_qmv_fast_b4g64"), + (true, true, 8, 64) => Ok(c"kernel_qwen_affine_gather_qmv_fast_b8g64"), + (false, false, 4, 32) => Ok(c"kernel_qwen_affine_qmv_safe_b4g32"), + (false, false, 4, 64) => Ok(c"kernel_qwen_affine_qmv_safe_b4g64"), + (false, false, 8, 64) => Ok(c"kernel_qwen_affine_qmv_safe_b8g64"), + (true, false, 4, 32) => Ok(c"kernel_qwen_affine_gather_qmv_safe_b4g32"), + (true, false, 4, 64) => Ok(c"kernel_qwen_affine_gather_qmv_safe_b4g64"), + (true, false, 8, 64) => Ok(c"kernel_qwen_affine_gather_qmv_safe_b8g64"), + _ => Err(format!( + "Qwen affine QMV does not support {}-bit group-{} weights", + weight.bits, weight.group + )), + } + } + + fn affine_qmv_batch_kernel( + &self, + weight: &Affine<'_>, + in_dim: u32, + out_dim: u32, + ) -> Result<&'static CStr, String> { + let alignment = match weight.bits { + 4 => 512, + 8 => 256, + _ => 0, + }; + let fast = alignment != 0 + && out_dim.is_multiple_of(8) + && in_dim.is_multiple_of(alignment) + && self.affine_is_direct(weight); + match (fast, weight.bits, weight.group) { + (true, 4, 32) => Ok(c"kernel_qwen_affine_qmv_batch_fast_b4g32"), + (true, 4, 64) => Ok(c"kernel_qwen_affine_qmv_batch_fast_b4g64"), + (true, 8, 64) => Ok(c"kernel_qwen_affine_qmv_batch_fast_b8g64"), + (false, 4, 32) => Ok(c"kernel_qwen_affine_qmv_batch_safe_b4g32"), + (false, 4, 64) => Ok(c"kernel_qwen_affine_qmv_batch_safe_b4g64"), + (false, 8, 64) => Ok(c"kernel_qwen_affine_qmv_batch_safe_b8g64"), + _ => Err(format!( + "Qwen affine batch QMV does not support {}-bit group-{} weights", + weight.bits, weight.group + )), + } + } + + fn affine_qmm_kernel(&self, weight: &Affine<'_>, rows: u32) -> Result<&'static CStr, String> { + match (rows <= 32, weight.bits, weight.group) { + (true, 4, 32) => Ok(c"kernel_qwen_affine_qmm_mpp_b4g32_bm32"), + (true, 4, 64) => Ok(c"kernel_qwen_affine_qmm_mpp_b4g64_bm32"), + (true, 8, 64) => Ok(c"kernel_qwen_affine_qmm_mpp_b8g64_bm32"), + (false, 4, 32) => Ok(c"kernel_qwen_affine_qmm_mpp_b4g32_bm64"), + (false, 4, 64) => Ok(c"kernel_qwen_affine_qmm_mpp_b4g64_bm64"), + (false, 8, 64) => Ok(c"kernel_qwen_affine_qmm_mpp_b8g64_bm64"), + _ => Err(format!( + "Qwen affine QMM does not support {}-bit group-{} weights", + weight.bits, weight.group + )), + } + } + + fn affine_gather_batch_kernel( + &self, + weight: &Affine<'_>, + in_dim: u32, + out_dim: u32, + ) -> Result<&'static CStr, String> { + let alignment = match weight.bits { + 4 => 512, + 8 => 256, + _ => 0, + }; + let fast = alignment != 0 + && out_dim.is_multiple_of(8) + && in_dim.is_multiple_of(alignment) + && self.affine_is_direct(weight); + match (fast, weight.bits, weight.group) { + (true, 4, 32) => Ok(c"kernel_qwen_affine_gather_qmv_batch_fast_b4g32"), + (true, 4, 64) => Ok(c"kernel_qwen_affine_gather_qmv_batch_fast_b4g64"), + (true, 8, 64) => Ok(c"kernel_qwen_affine_gather_qmv_batch_fast_b8g64"), + (false, 4, 32) => Ok(c"kernel_qwen_affine_gather_qmv_batch_safe_b4g32"), + (false, 4, 64) => Ok(c"kernel_qwen_affine_gather_qmv_batch_safe_b4g64"), + (false, 8, 64) => Ok(c"kernel_qwen_affine_gather_qmv_batch_safe_b8g64"), + _ => Err(format!( + "Qwen affine gather batch does not support {}-bit group-{} weights", + weight.bits, weight.group + )), + } + } + + fn affine_pair_qmv_batch_kernel( + &self, + gate: &Affine<'_>, + up: &Affine<'_>, + gather: bool, + in_dim: u32, + out_dim: u32, + ) -> Result<&'static CStr, String> { + let alignment = match gate.bits { + 4 => 512, + 8 => 256, + _ => 0, + }; + let fast = alignment != 0 + && out_dim.is_multiple_of(8) + && in_dim.is_multiple_of(alignment) + && self.affine_is_direct(gate) + && self.affine_is_direct(up); + match (gather, fast, gate.bits, gate.group) { + (false, true, 4, 32) => Ok(c"kernel_qwen_affine_pair_qmv_batch_fast_b4g32"), + (false, true, 4, 64) => Ok(c"kernel_qwen_affine_pair_qmv_batch_fast_b4g64"), + (false, true, 8, 64) => Ok(c"kernel_qwen_affine_pair_qmv_batch_fast_b8g64"), + (true, true, 4, 32) => Ok(c"kernel_qwen_affine_gather_pair_qmv_batch_fast_b4g32"), + (true, true, 4, 64) => Ok(c"kernel_qwen_affine_gather_pair_qmv_batch_fast_b4g64"), + (true, true, 8, 64) => Ok(c"kernel_qwen_affine_gather_pair_qmv_batch_fast_b8g64"), + (false, false, 4, 32) => Ok(c"kernel_qwen_affine_pair_qmv_batch_safe_b4g32"), + (false, false, 4, 64) => Ok(c"kernel_qwen_affine_pair_qmv_batch_safe_b4g64"), + (false, false, 8, 64) => Ok(c"kernel_qwen_affine_pair_qmv_batch_safe_b8g64"), + (true, false, 4, 32) => Ok(c"kernel_qwen_affine_gather_pair_qmv_batch_safe_b4g32"), + (true, false, 4, 64) => Ok(c"kernel_qwen_affine_gather_pair_qmv_batch_safe_b4g64"), + (true, false, 8, 64) => Ok(c"kernel_qwen_affine_gather_pair_qmv_batch_safe_b8g64"), + _ => Err(format!( + "Qwen affine batch pair does not support {}-bit group-{} weights", + gate.bits, gate.group + )), + } + } + + fn affine_pair_qmv_wide_kernel( + &self, + gate: &Affine<'_>, + up: &Affine<'_>, + ) -> Result<&'static CStr, String> { + match (gate.bits, gate.group, up.bits, up.group) { + (4, 32, 4, 32) => Ok(c"kernel_qwen_affine_pair_qmv_wide_b4g32"), + (4, 64, 4, 64) => Ok(c"kernel_qwen_affine_pair_qmv_wide_b4g64"), + (8, 64, 8, 64) => Ok(c"kernel_qwen_affine_pair_qmv_wide_b8g64"), + _ => Err("Qwen gate/up projections use incompatible quantization layouts".into()), + } + } + fn weight(&self, name: &str) -> Result, String> { self.model.tensor(name).map(|tensor| Weight { tensor, @@ -1353,12 +3059,37 @@ impl QwenExecutor { }) } + fn affine_is_direct(&self, weight: &Affine<'_>) -> bool { + [&weight.packed, &weight.scales, &weight.biases] + .into_iter() + .all(|part| { + self.weight_views + .contains_key(&(part.tensor as *const QwenTensor as usize)) + || part.tensor.range.start.is_multiple_of(4) + }) + } + fn view(&self, weight: Weight<'_>) -> QwenWeightView { + if let Some(view) = self + .weight_views + .get(&(weight.tensor as *const QwenTensor as usize)) + { + let experts = weight.expert.map_or(1, |_| EXPERTS as u64); + let bytes = view.bytes / experts; + return QwenWeightView { + tensor: view.tensor, + map: view.map, + size: view.size, + offset: u64::from(weight.expert.unwrap_or(0)) * bytes, + bytes, + }; + } let (map, _) = self.model.map(weight.tensor.map); let experts = weight.expert.map_or(1, |_| EXPERTS as u64); let bytes = (weight.tensor.range.end - weight.tensor.range.start) / experts; let offset = weight.tensor.range.start + u64::from(weight.expert.unwrap_or(0)) * bytes; QwenWeightView { + tensor: std::ptr::null(), map: map.as_ptr().cast(), size: map.len() as u64, offset, @@ -1373,14 +3104,28 @@ impl QwenExecutor { output: &Buffer, in_dim: u32, out_dim: u32, + ) -> Result<(), String> { + self.affine_mv_with_precision(weight, input, output, in_dim, out_dim, true) + } + + fn affine_mv_with_precision( + &self, + weight: &Affine<'_>, + input: &Buffer, + output: &Buffer, + in_dim: u32, + out_dim: u32, + bf16_output: bool, ) -> Result<(), String> { let mut values = args(); values.u[0] = in_dim; values.u[1] = out_dim; values.u[2] = weight.bits; values.u[3] = weight.group; + values.u[11] = u32::from(bf16_output); + values.u[12] = 64; self.dispatch( - c"kernel_qwen_affine_mv", + self.affine_qmv_kernel(weight, false, in_dim, out_dim)?, output, Some(input), None, @@ -1391,11 +3136,182 @@ impl QwenExecutor { self.view(weight.biases), ], &values, - out_dim, + out_dim.div_ceil(8), 1, ) } + #[allow(clippy::too_many_arguments)] + fn affine_batch_into( + &self, + weight: &Affine<'_>, + input: &Buffer, + output: &Buffer, + in_dim: u32, + out_dim: u32, + rows: u32, + bf16_output: bool, + ) -> Result<(), String> { + if rows == 1 { + return self.affine_mv_with_precision( + weight, + input, + output, + in_dim, + out_dim, + bf16_output, + ); + } + if rows > VERIFY_ROWS as u32 { + if unsafe { ds4_gpu_metal4_tensor_api_enabled() } == 0 { + return Err("Qwen batched prefill requires the Metal 4 tensor API".to_string()); + } + let mut values = args(); + values.u[0] = in_dim; + values.u[1] = out_dim; + values.u[2] = weight.bits; + values.u[3] = weight.group; + values.u[4] = rows; + values.u[10] = 0; + values.u[11] = u32::from(bf16_output); + values.u[12] = 128; + return self.dispatch( + self.affine_qmm_kernel(weight, rows)?, + output, + Some(input), + None, + None, + &[ + self.view(weight.packed), + self.view(weight.scales), + self.view(weight.biases), + ], + &values, + out_dim.div_ceil(64), + rows.div_ceil(if rows <= 32 { 32 } else { 64 }), + ); + } + let mut values = args(); + values.u[0] = in_dim; + values.u[1] = out_dim; + values.u[2] = weight.bits; + values.u[3] = weight.group; + values.u[4] = rows; + values.u[11] = u32::from(bf16_output); + values.u[12] = 64; + self.dispatch( + self.affine_qmv_batch_kernel(weight, in_dim, out_dim)?, + output, + Some(input), + None, + None, + &[ + self.view(weight.packed), + self.view(weight.scales), + self.view(weight.biases), + ], + &values, + out_dim.div_ceil(8), + 1, + ) + } + + #[allow(clippy::too_many_arguments)] + fn affine_gather_batch_into( + &self, + weight: &Affine<'_>, + input: &Buffer, + output: &Buffer, + in_dim: u32, + out_dim: u32, + rows: u32, + per_expert_input: bool, + ) -> Result<(), String> { + let mut values = args(); + values.u[0] = in_dim; + values.u[1] = out_dim; + values.u[2] = weight.bits; + values.u[3] = weight.group; + values.u[4] = rows; + values.u[5] = u32::from(per_expert_input); + values.u[8] = EXPERTS_USED as u32; + values.u[11] = 1; + values.u[12] = 64; + self.dispatch( + self.affine_gather_batch_kernel(weight, in_dim, out_dim)?, + output, + Some(input), + Some(&self.scratch.route_ids), + None, + &[ + self.view(weight.packed), + self.view(weight.scales), + self.view(weight.biases), + ], + &values, + out_dim.div_ceil(8), + rows * EXPERTS_USED as u32, + ) + } + + #[allow(clippy::too_many_arguments)] + fn affine_pair_batch_into( + &self, + gate: &Affine<'_>, + up: &Affine<'_>, + input: &Buffer, + gate_output: &Buffer, + up_output: &Buffer, + in_dim: u32, + out_dim: u32, + gather: bool, + rows: u32, + per_expert_input: bool, + ) -> Result<(), String> { + if gate.bits != up.bits || gate.group != up.group { + return Err("Qwen gate/up projections use different quantization layouts".into()); + } + let mut values = args(); + values.u[0] = in_dim; + values.u[1] = out_dim; + values.u[2] = gate.bits; + values.u[3] = gate.group; + values.u[4] = rows; + values.u[5] = u32::from(per_expert_input); + values.u[8] = EXPERTS_USED as u32; + values.u[11] = 1; + values.u[12] = 64; + let wide = !gather && rows > 1; + self.dispatch( + if wide { + self.affine_pair_qmv_wide_kernel(gate, up)? + } else { + self.affine_pair_qmv_batch_kernel(gate, up, gather, in_dim, out_dim)? + }, + gate_output, + Some(input), + Some(&self.scratch.route_ids), + Some(up_output), + &[ + self.view(gate.packed), + self.view(gate.scales), + self.view(gate.biases), + self.view(up.packed), + self.view(up.scales), + self.view(up.biases), + ], + &values, + out_dim.div_ceil(8) * 2, + if wide { + 1 + } else if gather { + rows * EXPERTS_USED as u32 + } else { + rows + }, + ) + } + fn bf16_mv( &self, weight: Weight<'_>, @@ -1413,15 +3329,130 @@ impl QwenExecutor { let mut values = args(); values.u[0] = in_dim; values.u[1] = out_dim; + let (mode, rows_per_group, threads, threadgroup_bytes) = if in_dim <= 64 { + (3, 32, 32, 0) + } else if in_dim >= 16 * out_dim { + (2, 4, 256, 64 * 4) + } else if out_dim >= 4096 { + (1, 32, 256, 0) + } else { + (0, 16, 128, 0) + }; + values.u[4] = mode; + values.u[10] = threadgroup_bytes; + values.u[12] = threads; self.dispatch( - c"kernel_qwen_bf16_mv", + c"kernel_qwen_bf16_gemv", output, Some(input), None, None, &[self.view(weight)], &values, - out_dim, + out_dim.div_ceil(rows_per_group), + 1, + ) + } + + fn bf16_batch( + &self, + weight: Weight<'_>, + input: &Buffer, + output: &Buffer, + in_dim: u32, + out_dim: u32, + rows: u32, + ) -> Result<(), String> { + if rows == VERIFY_ROWS as u32 { + return self.bf16_mv4(weight, input, output, in_dim, out_dim); + } + if weight.tensor.dtype != "BF16" || weight.tensor.shape != [out_dim as u64, in_dim as u64] { + return Err(format!( + "{} has an incompatible BF16 matrix shape", + weight.tensor.name + )); + } + if rows > VERIFY_ROWS as u32 { + if unsafe { ds4_gpu_metal4_tensor_api_enabled() } == 0 { + return Err("Qwen batched prefill requires the Metal 4 tensor API".to_string()); + } + let mut values = args(); + values.u[0] = in_dim; + values.u[1] = out_dim; + values.u[4] = rows; + values.u[12] = 128; + let tile_m = if rows <= 32 { 32 } else { 64 }; + return self.dispatch( + if tile_m == 32 { + c"kernel_qwen_bf16_qmm_mpp_bm32" + } else { + c"kernel_qwen_bf16_qmm_mpp_bm64" + }, + output, + Some(input), + None, + None, + &[self.view(weight)], + &values, + out_dim.div_ceil(64), + rows.div_ceil(tile_m), + ); + } + let mut values = args(); + values.u[0] = in_dim; + values.u[1] = out_dim; + let (mode, rows_per_group, threads, threadgroup_bytes) = if in_dim <= 64 { + (3, 32, 32, 0) + } else if in_dim >= 16 * out_dim { + (2, 4, 256, 64 * 4) + } else if out_dim >= 4096 { + (1, 32, 256, 0) + } else { + (0, 16, 128, 0) + }; + values.u[4] = mode; + values.u[10] = threadgroup_bytes; + values.u[12] = threads; + self.dispatch( + c"kernel_qwen_bf16_gemv", + output, + Some(input), + None, + None, + &[self.view(weight)], + &values, + out_dim.div_ceil(rows_per_group), + rows, + ) + } + + fn bf16_mv4( + &self, + weight: Weight<'_>, + input: &Buffer, + output: &Buffer, + in_dim: u32, + out_dim: u32, + ) -> Result<(), String> { + if weight.tensor.dtype != "BF16" || weight.tensor.shape != [out_dim as u64, in_dim as u64] { + return Err(format!( + "{} has an incompatible BF16 matrix shape", + weight.tensor.name + )); + } + let mut values = args(); + values.u[0] = in_dim; + values.u[1] = out_dim; + values.u[12] = 128; + self.dispatch( + c"kernel_qwen_bf16_gemv4", + output, + Some(input), + None, + None, + &[self.view(weight)], + &values, + out_dim.div_ceil(4), 1, ) } @@ -1442,12 +3473,271 @@ impl QwenExecutor { dispatch_qwen(kernel, out, a, b, c, weights, values, grid_x, grid_y) } + #[allow(clippy::too_many_arguments)] + fn dispatch_with_d( + &self, + kernel: &CStr, + out: &Buffer, + a: Option<&Buffer>, + b: Option<&Buffer>, + c: Option<&Buffer>, + d: Option<&Buffer>, + weights: &[QwenWeightView], + values: &QwenKernelArgs, + grid_x: u32, + grid_y: u32, + ) -> Result<(), String> { + dispatch_qwen_with_d(kernel, out, a, b, c, d, weights, values, grid_x, grid_y) + } + + fn encode_sparse_probabilities( + &self, + logits: &Buffer, + rows: u32, + temperature: f32, + top_k: i32, + ) -> Result { + if temperature <= 0.0 || top_k <= 0 || top_k as u32 > SAMPLING_TOP_K { + return Ok(false); + } + let top_k = top_k as u32; + call( + unsafe { + ds4_gpu_indexer_topk_tensor( + self.scratch.sampling_selected.raw(), + logits.raw(), + VOCAB, + rows, + top_k, + ) + }, + "selecting Qwen sampling support", + )?; + let mut values = args(); + values.u[0] = VOCAB; + values.u[1] = top_k; + values.u[4] = rows; + values.u[12] = 256; + values.f[0] = temperature.recip(); + self.dispatch( + c"kernel_qwen_sparse_probabilities", + &self.scratch.sampling_probabilities, + Some(logits), + Some(&self.scratch.sampling_selected), + None, + &[], + &values, + rows, + 1, + )?; + Ok(true) + } + + fn read_sparse_probabilities( + &self, + rows: u32, + top_p: f32, + min_p: f32, + top_k: i32, + ) -> Result, String> { + if top_k <= 0 || top_k as u32 > SAMPLING_TOP_K { + return Ok(None); + } + let top_k = top_k as u32; + let count = rows as usize * top_k as usize; + let mut ids = vec![0; count]; + let mut probabilities = vec![0.0; count]; + self.scratch.sampling_selected.read_i32(&mut ids)?; + self.scratch + .sampling_probabilities + .read_f32(&mut probabilities)?; + let top_p = if top_p <= 0.0 || top_p > 1.0 { + 1.0 + } else { + f64::from(top_p) + }; + let minimum = f64::from(min_p.max(0.0)); + let mut distributions = Vec::with_capacity(rows as usize); + for row in 0..rows as usize { + let start = row * top_k as usize; + let mut support = ids[start..start + top_k as usize] + .iter() + .copied() + .zip(probabilities[start..start + top_k as usize].iter().copied()) + .filter(|(token, probability)| { + *token >= 0 && *token < VOCAB as i32 && probability.is_finite() + }) + .map(|(token, probability)| (token as usize, f64::from(probability))) + .collect::>(); + support.sort_unstable_by(|left, right| { + right + .1 + .total_cmp(&left.1) + .then_with(|| left.0.cmp(&right.0)) + }); + let mut kept = 0.0; + let mut end = 0; + for (_, probability) in &support { + if end > 0 && *probability < minimum * support[0].1 { + break; + } + if kept >= top_p { + break; + } + kept += *probability; + end += 1; + } + support.truncate(end); + if support.is_empty() || !kept.is_finite() || kept <= 0.0 { + return Ok(None); + } + for (_, probability) in &mut support { + *probability /= kept; + } + support.sort_unstable_by_key(|(token, _)| *token); + distributions.push(support); + } + Ok(Some(distributions)) + } + + fn mtp_history_batch(&mut self, tokens: &[i32], position: u32) -> Result<(), String> { + let commands = Commands::begin()?; + self.encode_mtp_history_batch(tokens, position)?; + commands.finish() + } + + fn encode_mtp_history_batch(&mut self, tokens: &[i32], position: u32) -> Result<(), String> { + if tokens.is_empty() || tokens.len() > self.scratch.rows as usize { + return Err("Qwen MTP prompt-history width exceeds the prefill batch".into()); + } + if position + tokens.len() as u32 > self.context + MTP_TOKEN_RESERVE { + return Err("Qwen MTP prompt history exceeds its bounded cache".into()); + } + if let Some(&token) = tokens + .iter() + .find(|&&token| token < 0 || token as u32 >= VOCAB) + { + return Err(format!("token {token} is outside the Qwen vocabulary")); + } + + let mtp = self.mtp.take().ok_or("Qwen MTP is not configured")?; + let result = (|| { + let rows = tokens.len() as u32; + + let mut hidden_norm = rms_args(HC_WIDTH, HC_WIDTH); + hidden_norm.u[4] = rows; + hidden_norm.u[11] = 1; + self.dispatch( + c"kernel_qwen_zero_rms", + &self.scratch.hc_norm, + Some(&self.scratch.hc), + None, + None, + &[self.view(self.weight("mtp.pre_fc_norm_hidden.weight")?)], + &hidden_norm, + HC * rows, + 1, + )?; + self.bf16_batch( + self.weight("mtp.fc_hidden.weight")?, + &self.scratch.hc_norm, + &self.scratch.hc_mix, + HIDDEN, + HIDDEN, + HC * rows, + )?; + + self.scratch.token_ids.write_i32(tokens)?; + let embedding = + self.affine("language_model.model.embed_tokens", HIDDEN, VOCAB, None)?; + let mut embed = args(); + embed.u[0] = HIDDEN; + embed.u[2] = embedding.bits; + embed.u[3] = embedding.group; + embed.u[4] = rows; + self.dispatch( + c"kernel_qwen_affine_embedding_batch", + &self.scratch.hidden, + Some(&self.scratch.token_ids), + None, + None, + &[ + self.view(embedding.packed), + self.view(embedding.scales), + self.view(embedding.biases), + ], + &embed, + HIDDEN, + rows, + )?; + let mut embedding_norm = rms_args(HIDDEN, HIDDEN); + embedding_norm.u[4] = rows; + embedding_norm.u[11] = 1; + self.dispatch( + c"kernel_qwen_zero_rms", + &self.scratch.block, + Some(&self.scratch.hidden), + None, + None, + &[self.view(self.weight("mtp.pre_fc_norm_embedding.weight")?)], + &embedding_norm, + rows, + 1, + )?; + self.bf16_batch( + self.weight("mtp.fc_embedding.weight")?, + &self.scratch.block, + &self.scratch.hidden, + HIDDEN, + HIDDEN, + rows, + )?; + let mut repeat = args(); + repeat.u[0] = HIDDEN; + repeat.u[4] = rows; + self.dispatch( + c"kernel_qwen_repeat4", + &self.scratch.hc_norm, + Some(&self.scratch.hidden), + None, + None, + &[], + &repeat, + HC_WIDTH * rows, + 1, + )?; + let mut add = args(); + add.u[0] = HC_WIDTH * rows; + self.dispatch( + c"kernel_qwen_add", + &self.scratch.hc, + Some(&self.scratch.hc_mix), + Some(&self.scratch.hc_norm), + None, + &[], + &add, + HC_WIDTH * rows, + 1, + )?; + + let prefix = "mtp.layers.0"; + self.hyper_read_rows(&format!("{prefix}.attn_hyper_connection"), rows)?; + self.attention_rows(prefix, &mtp.attention, position, rows)?; + self.hyper_write_rows(rows)?; + self.encode_moe_rows(prefix, rows)?; + Ok(()) + })(); + self.mtp = Some(mtp); + result + } + fn mtp_step( &mut self, token: i32, position: u32, draft_requested: bool, - ) -> Result>, String> { + sampling: Option<(f32, f32, f32, i32)>, + ) -> Result, String> { let mut mtp = self.mtp.take().ok_or("Qwen MTP is not configured")?; let started = Instant::now(); let result = (|| { @@ -1481,10 +3771,8 @@ impl QwenExecutor { HIDDEN, 1, )?; - let mut norm = args(); - norm.u[0] = HIDDEN; - norm.u[1] = HIDDEN; - norm.f[0] = 1.0e-6; + let mut norm = rms_args(HIDDEN, HIDDEN); + norm.u[11] = 1; self.dispatch( c"kernel_qwen_zero_rms", &self.scratch.block, @@ -1517,12 +3805,13 @@ impl QwenExecutor { 1, )?; let hidden_weight = self.weight("mtp.fc_hidden.weight")?; - for stream in 0..HC { - let offset = u64::from(stream * HIDDEN) * 4; - let input = self.scratch.hc_norm.view(offset, u64::from(HIDDEN) * 4)?; - let output = self.scratch.hc_mix.view(offset, u64::from(HIDDEN) * 4)?; - self.bf16_mv(hidden_weight, &input, &output, HIDDEN, HIDDEN)?; - } + self.bf16_mv4( + hidden_weight, + &self.scratch.hc_norm, + &self.scratch.hc_mix, + HIDDEN, + HIDDEN, + )?; let mut repeat = args(); repeat.u[0] = HIDDEN; self.dispatch( @@ -1549,20 +3838,16 @@ impl QwenExecutor { HC_WIDTH, 1, )?; - commands.finish()?; - let prefix = "mtp.layers.0"; - let commands = Commands::begin()?; self.hyper_read(&format!("{prefix}.attn_hyper_connection"))?; self.attention(prefix, &mtp.attention, position)?; self.hyper_write()?; - commands.finish()?; self.encode_moe(prefix, "MTP")?; if !draft_requested { + commands.finish()?; return Ok(None); } - let commands = Commands::begin()?; self.mix_output("mtp.hyper_connection_mixer")?; self.affine_mv_into( &self.affine("language_model.lm_head", HIDDEN, VOCAB, None)?, @@ -1571,8 +3856,12 @@ impl QwenExecutor { HIDDEN, VOCAB, )?; + let sparse = if let Some((temperature, _, _, top_k)) = sampling { + self.encode_sparse_probabilities(&self.scratch.logits, 1, temperature, top_k)? + } else { + false + }; commands.finish()?; - self.scratch.logits.read_f32(&mut self.logits)?; mtp.drafted += 1; if mtp.timing { eprintln!( @@ -1580,19 +3869,31 @@ impl QwenExecutor { started.elapsed().as_secs_f64() * 1000.0 ); } - Ok(Some(self.logits.clone())) + let probabilities = if sparse { + let (_, top_p, min_p, top_k) = + sampling.expect("sampling configuration disappeared"); + self.read_sparse_probabilities(1, top_p, min_p, top_k)? + .and_then(|mut rows| rows.pop()) + } else { + None + }; + let logits = if probabilities.is_some() { + Vec::new() + } else { + self.scratch.logits.read_f32(&mut self.logits)?; + self.logits.clone() + }; + Ok(Some((logits, probabilities))) })(); self.mtp = Some(mtp); result } - fn capture_recurrent(&mut self, slot: usize) -> Result<(), String> { + #[cfg(test)] + fn capture_recurrent(&mut self) -> Result<(), String> { let commands = Commands::begin()?; let mtp = self.mtp.as_mut().ok_or("Qwen MTP is not configured")?; - let snapshot = mtp - .captures - .get_mut(slot) - .ok_or("Qwen MTP capture slot is out of range")?; + let snapshot = &mut mtp.snapshot; for (state, saved) in self.states.iter().zip(&snapshot.gdn) { match (state, saved) { (LayerState::Gdn { conv, recurrent }, Some((saved_conv, saved_recurrent))) => { @@ -1636,13 +3937,11 @@ impl QwenExecutor { commands.finish() } - fn restore_recurrent(&mut self, slot: usize) -> Result<(), String> { + #[cfg(test)] + fn restore_recurrent(&mut self) -> Result<(), String> { let commands = Commands::begin()?; let mtp = self.mtp.as_ref().ok_or("Qwen MTP is not configured")?; - let snapshot = mtp - .captures - .get(slot) - .ok_or("Qwen MTP restore slot is out of range")?; + let snapshot = &mtp.snapshot; for (state, saved) in self.states.iter().zip(&snapshot.gdn) { match (state, saved) { (LayerState::Gdn { conv, recurrent }, Some((saved_conv, saved_recurrent))) => { @@ -1694,10 +3993,15 @@ impl QwenExecutor { { return 0; } - max_tokens + let limit = max_tokens .saturating_sub(1) .min(MTP_TOKEN_RESERVE) - .min(self.context.saturating_sub(self.position + 1)) + .min(self.context.saturating_sub(self.position + 1)); + #[cfg(test)] + if let Ok(depth) = std::env::var("DS4SERVER_QWEN38_DRAFT_DEPTH") { + return limit.min(depth.parse().expect("invalid Qwen test draft depth")); + } + limit } fn draft_greedy( @@ -1711,10 +4015,16 @@ impl QwenExecutor { let start = self.position - 1; for step in 0..count { let logits = self - .mtp_step(input, start + step, true)? - .ok_or("Qwen MTP did not return draft logits")?; + .mtp_step(input, start + step, true, None)? + .ok_or("Qwen MTP did not return draft logits")? + .0; let token = argmax(&logits); - drafts.push(QwenDraft { token, logits }); + drafts.push(QwenDraft { + token, + #[cfg(test)] + logits, + probabilities: Vec::new(), + }); input = token; if self.model.is_stop_token_for_reasoning(token, reasoning) { break; @@ -1739,11 +4049,24 @@ impl QwenExecutor { let mut input = token; let start = self.position - 1; for step in 0..count { - let logits = self - .mtp_step(input, start + step, true)? + let (logits, probabilities) = self + .mtp_step( + input, + start + step, + true, + Some((temperature, top_p, min_p, top_k)), + )? .ok_or("Qwen MTP did not return draft logits")?; - let token = sample_from_logits(&logits, temperature, top_p, min_p, top_k, rng); - drafts.push(QwenDraft { token, logits }); + let probabilities = probabilities.unwrap_or_else(|| { + qwen_sampling_probabilities(&logits, temperature, top_p, min_p, top_k) + }); + let token = sample_probabilities_f64(&probabilities, rng); + drafts.push(QwenDraft { + token, + #[cfg(test)] + logits, + probabilities, + }); input = token; if self.model.is_stop_token_for_reasoning(token, reasoning) { break; @@ -1752,48 +4075,249 @@ impl QwenExecutor { Ok(drafts) } - fn evaluate_drafts(&mut self, drafts: &[QwenDraft]) -> Result<(Vec>, bool), String> { - let mut captures_complete = self.capture_recurrent(0).is_ok(); - let mut target_logits = Vec::with_capacity(drafts.len()); - for (index, draft) in drafts.iter().enumerate() { - target_logits.push(self.logits.clone()); - self.eval_target(draft.token)?; - if captures_complete { - captures_complete = self.capture_recurrent(index + 1).is_ok(); - } - } - Ok((target_logits, captures_complete)) - } - - fn restore_or_replay( - &mut self, - slot: usize, - tokens: &[i32], - verified_end: u32, - mtp_end: u32, - captures_complete: bool, - ) -> Result<(), String> { - if captures_complete - && self.restore_recurrent(slot).is_ok() - && self.trim_rejected_attention(verified_end, mtp_end).is_ok() - { + fn canonicalize_batch_row(&self, row: usize) -> Result<(), String> { + if row == 0 { return Ok(()); } - self.replay_tokens(tokens) + let commands = Commands::begin()?; + self.scratch.hc.copy_from( + 0, + &self.scratch.hc, + row as u64 * u64::from(HC_WIDTH) * 4, + u64::from(HC_WIDTH) * 4, + "selecting the committed Qwen verify hidden row", + )?; + commands.finish() } - fn trim_rejected_attention(&self, verified_end: u32, mtp_end: u32) -> Result<(), String> { - for state in &self.states { - if matches!(state, LayerState::Attention { .. }) { - clear_attention_rows(state, self.position, verified_end)?; + fn commit_final_mtp_batch_row( + &mut self, + drafts: &[QwenDraft], + batch_start: u32, + ) -> Result<(), String> { + let rows = u32::try_from(drafts.len()).map_err(|_| "Qwen MTP draft width overflow")?; + if rows == 0 { + return Ok(()); + } + self.commit_verifier_recurrent_state()?; + let commands = Commands::begin()?; + self.scratch.batch_final_hc.copy_from( + 0, + &self.scratch.hc, + u64::from(rows) * u64::from(HC_WIDTH) * 4, + u64::from(HC_WIDTH) * 4, + "preserving the final Qwen verify hidden row", + )?; + self.scratch.hc.copy_from( + 0, + &self + .mtp + .as_ref() + .ok_or("Qwen MTP is not configured")? + .verify + .hidden, + 0, + u64::from(rows * HC_WIDTH) * 4, + "restoring authoritative Qwen MTP history rows", + )?; + let tokens = drafts.iter().map(|draft| draft.token).collect::>(); + self.encode_mtp_history_batch(&tokens, batch_start)?; + self.scratch.hc.copy_from( + 0, + &self.scratch.batch_final_hc, + 0, + u64::from(HC_WIDTH) * 4, + "restoring the final Qwen verify hidden row", + )?; + commands.finish() + } + + fn commit_verifier_recurrent_state(&mut self) -> Result<(), String> { + let mtp = self.mtp.as_mut().ok_or("Qwen MTP is not configured")?; + for (state, capture) in self.states.iter_mut().zip(&mut mtp.verify.gdn) { + match (state, capture) { + ( + LayerState::Gdn { conv, recurrent }, + Some(GdnVerifyCapture { + conv: verified_conv, + recurrent: verified_recurrent, + .. + }), + ) => { + std::mem::swap(conv, verified_conv); + std::mem::swap(recurrent, verified_recurrent); + } + (LayerState::Attention { .. }, None) => {} + _ => return Err("Qwen verifier GDN state layout is invalid".into()), } } - if let Some(mtp) = &self.mtp { - clear_attention_rows(&mtp.attention, self.position.saturating_sub(1), mtp_end)?; - } + std::mem::swap(&mut self.ple_state.conv, &mut mtp.verify.ple_conv); Ok(()) } + fn commit_verified_window_from_current( + &mut self, + committed: &[i32], + committed_logits: &[f32], + ) -> Result<(), String> { + let keep_rows = committed.len(); + if keep_rows == 0 || keep_rows > MTP_CAPTURE_ROWS { + return Err("Qwen verifier commit width is invalid".into()); + } + if committed_logits.len() != VOCAB as usize { + return Err("Qwen verifier commit logits are invalid".into()); + } + let batch_start = self.position; + + let rows = keep_rows as u32; + let commands = Commands::begin()?; + { + let ple = self.mtp.as_ref().ok_or("Qwen MTP is not configured")?; + self.scratch.hc.copy_from( + 0, + &ple.verify.ple_input, + 0, + u64::from(rows * HC_WIDTH) * 4, + "restoring captured Qwen verifier PLE input rows", + )?; + } + if rows == 1 { + self.ple(committed[0])?; + self.encode_layer(1)?; + } else { + self.ple_batch(committed)?; + let ple_layer = "language_model.model.layers.1"; + self.hyper_read_rows(&format!("{ple_layer}.attn_hyper_connection"), rows)?; + self.gdn_rows(ple_layer, 1, rows, false)?; + self.hyper_write_rows(rows)?; + self.encode_moe_rows(ple_layer, rows)?; + } + + for layer in 0..LAYERS { + if layer == 1 { + continue; + } + let LayerState::Gdn { conv, recurrent } = &self.states[layer] else { + continue; + }; + let capture = self + .mtp + .as_ref() + .and_then(|mtp| mtp.verify.gdn.get(layer)) + .and_then(Option::as_ref) + .ok_or("Qwen verifier GDN capture layout is invalid")?; + self.scratch.gdn_out.copy_from( + 0, + &capture.v, + 0, + u64::from(rows * GDN_VALUE) * 4, + "restoring captured Qwen verifier GDN value rows", + )?; + let mut verify = args(); + verify.u[0] = GDN_VALUE; + verify.u[1] = GDN_VALUE + GDN_HEADS_V; + verify.u[4] = rows; + verify.u[12] = 64; + self.dispatch_with_d( + c"kernel_qwen_gdn_delta_rows", + &self.scratch.gdn_out, + Some(&capture.q), + Some(&capture.k), + Some(&capture.controls), + Some(recurrent), + &[buffer_weight_view( + recurrent, + u64::from(GDN_HEADS_V) * HEAD_DIM as u64 * HEAD_DIM as u64 * 4, + )], + &verify, + HEAD_DIM / 16, + GDN_HEADS_V, + )?; + let mut commit = args(); + commit.u[0] = GDN_QKV; + commit.u[4] = rows; + self.dispatch( + c"kernel_qwen_gdn_commit_conv_rows", + conv, + Some(&capture.qkv), + None, + None, + &[], + &commit, + GDN_QKV, + 1, + )?; + } + + if rows > 1 { + let mtp = self.mtp.as_ref().ok_or("Qwen MTP is not configured")?; + self.scratch.hc.copy_from( + 0, + &mtp.verify.hidden, + 0, + u64::from((rows - 1) * HC_WIDTH) * 4, + "restoring authoritative Qwen MTP history rows", + )?; + self.encode_mtp_history_batch(&committed[1..], batch_start)?; + } + let mtp = self.mtp.as_ref().ok_or("Qwen MTP is not configured")?; + self.scratch.hc.copy_from( + 0, + &mtp.verify.hidden, + (u64::from(rows) - 1) * u64::from(HC_WIDTH) * 4, + u64::from(HC_WIDTH) * 4, + "selecting the committed Qwen verifier hidden row", + )?; + commands.finish()?; + + self.tokens.extend_from_slice(committed); + self.position = batch_start + rows; + self.logits.copy_from_slice(committed_logits); + Ok(()) + } + + #[cfg(test)] + fn commit_verified_window( + &mut self, + keep_rows: usize, + tokens: &[i32], + committed_logits: &[f32], + ) -> Result<(), String> { + self.restore_recurrent()?; + let baseline_tokens = self.tokens.len(); + if tokens.len() != baseline_tokens + keep_rows { + return Err("Qwen verifier commit token span is invalid".into()); + } + self.commit_verified_window_from_current(&tokens[baseline_tokens..], committed_logits) + } + + #[allow(clippy::too_many_arguments)] + fn restore_or_replay( + &mut self, + committed: &[i32], + committed_logits: &[f32], + batch_start: u32, + baseline_tokens: usize, + ple_history: [i32; PLE_HISTORY], + captures_complete: bool, + ) -> Result { + if captures_complete { + self.position = batch_start; + self.tokens.truncate(baseline_tokens); + self.ple_state.history = ple_history; + match self.commit_verified_window_from_current(committed, committed_logits) { + Ok(()) => return Ok(true), + Err(error) => { + eprintln!("ds4: Qwen capture-commit refused ({error}); falling back to replay") + } + } + } + let mut tokens = self.tokens[..baseline_tokens].to_vec(); + tokens.extend_from_slice(committed); + self.replay_tokens(&tokens)?; + Ok(false) + } + fn replay_tokens(&mut self, tokens: &[i32]) -> Result<(), String> { let tag = self.checkpoint_tag; let counters = self.mtp.as_ref().map(|mtp| { @@ -1828,25 +4352,6 @@ impl QwenExecutor { Ok(()) } - fn commit_final_mtp_row( - &mut self, - draft: &QwenDraft, - final_slot: usize, - captures_complete: bool, - ) -> Result<(), String> { - let tokens = self.tokens.clone(); - if captures_complete - && self.restore_recurrent(final_slot - 1).is_ok() - && self - .mtp_step(draft.token, self.position.saturating_sub(1), false) - .is_ok() - && self.restore_recurrent(final_slot).is_ok() - { - return Ok(()); - } - self.replay_tokens(&tokens) - } - fn note_mtp_verification(&mut self, accepted: usize, started: Instant) { if let Some(mtp) = &mut self.mtp { mtp.accepted = mtp.accepted.saturating_add(accepted as u64); @@ -1876,42 +4381,67 @@ impl QwenExecutor { let count = self.draft_limit(max_tokens, reasoning, token); if count == 0 || cancelled.load(std::sync::atomic::Ordering::Relaxed) { if self.position > 0 && !self.model.is_stop_token_for_reasoning(token, reasoning) { - let _ = self.mtp_step(token, self.position - 1, false)?; + let _ = self.mtp_step(token, self.position - 1, false, None)?; } self.eval_target(token)?; return Ok(vec![token]); } let drafts = self.draft_greedy(token, count, reasoning)?; - self.eval_target(token)?; if drafts.is_empty() || cancelled.load(std::sync::atomic::Ordering::Relaxed) { + self.eval_target(token)?; return Ok(vec![token]); } let started = Instant::now(); - let baseline_tokens = self.tokens.clone(); - let mtp_end = self.position.saturating_sub(2) + drafts.len() as u32; - let (target_logits, captures_complete) = self.evaluate_drafts(&drafts)?; - let verified_end = self.position; + let baseline_tokens = self.tokens.len(); + let batch_start = self.position; + let baseline_ple_history = self.ple_state.history; + let mut verify_tokens = Vec::with_capacity(drafts.len() + 1); + verify_tokens.push(token); + verify_tokens.extend(drafts.iter().map(|draft| draft.token)); + let (verify_logits, _) = self.eval_target_batch(&verify_tokens, None)?; + let target_logits = &verify_logits[..drafts.len()]; let accepted = drafts .iter() - .zip(&target_logits) + .zip(target_logits) .take_while(|(draft, target)| argmax(target) == draft.token) .count(); if accepted != drafts.len() { let repair = Instant::now(); - let mut tokens = baseline_tokens; - tokens.extend(drafts[..accepted].iter().map(|draft| draft.token)); - self.restore_or_replay(accepted, &tokens, verified_end, mtp_end, captures_complete)?; + let mut committed = Vec::with_capacity(accepted + 1); + committed.push(token); + committed.extend(drafts[..accepted].iter().map(|draft| draft.token)); + let _ = self.restore_or_replay( + &committed, + &verify_logits[accepted], + batch_start, + baseline_tokens, + baseline_ple_history, + true, + )?; if let Some(mtp) = &mut self.mtp { mtp.repair_ns = mtp .repair_ns .saturating_add(u64::try_from(repair.elapsed().as_nanos()).unwrap_or(u64::MAX)); } - } else if let Some(last) = drafts.last() - && !self + } else if drafts.last().is_some_and(|last| { + !self .model .is_stop_token_for_reasoning(last.token, reasoning) - { - self.commit_final_mtp_row(last, drafts.len(), captures_complete)?; + }) { + let commit = Instant::now(); + self.commit_final_mtp_batch_row(&drafts, batch_start)?; + self.read_logits_row(drafts.len())?; + if std::env::var_os("DS4SERVER_QWEN38_MTP_PROFILE").is_some() { + eprintln!( + "ds4: Qwen verifier full commit rows={} ms={:.3}", + drafts.len(), + commit.elapsed().as_secs_f64() * 1_000.0 + ); + } + } else { + self.commit_verifier_recurrent_state()?; + self.canonicalize_batch_row(drafts.len())?; + self.read_logits_row(drafts.len())?; } self.note_mtp_verification(accepted, started); let mut emitted = Vec::with_capacity(accepted + 1); @@ -1933,20 +4463,34 @@ impl QwenExecutor { rng: &mut Rng, cancelled: &std::sync::atomic::AtomicBool, ) -> Result, String> { + let folded_primary = self.folding_sampled; + self.folding_sampled = false; if self.mtp.is_none() { self.eval_target(token)?; - return Ok(vec![token]); + return Ok(if folded_primary { + Vec::new() + } else { + vec![token] + }); } if let Some(mtp) = &mut self.mtp { mtp.cycles = mtp.cycles.saturating_add(1); } - let count = self.draft_limit(max_tokens, reasoning, token); + let count = self.draft_limit( + max_tokens.saturating_add(u32::from(folded_primary)), + reasoning, + token, + ); if count == 0 || cancelled.load(std::sync::atomic::Ordering::Relaxed) { if self.position > 0 && !self.model.is_stop_token_for_reasoning(token, reasoning) { - let _ = self.mtp_step(token, self.position - 1, false)?; + let _ = self.mtp_step(token, self.position - 1, false, None)?; } self.eval_target(token)?; - return Ok(vec![token]); + return Ok(if folded_primary { + Vec::new() + } else { + vec![token] + }); } let drafts = self.draft_sampled( token, @@ -1958,26 +4502,36 @@ impl QwenExecutor { top_k, rng, )?; - self.eval_target(token)?; if drafts.is_empty() || cancelled.load(std::sync::atomic::Ordering::Relaxed) { - return Ok(vec![token]); + self.eval_target(token)?; + return Ok(if folded_primary { + Vec::new() + } else { + vec![token] + }); } let started = Instant::now(); - let baseline_tokens = self.tokens.clone(); - let mtp_end = self.position.saturating_sub(2) + drafts.len() as u32; - let (target_logits, captures_complete) = self.evaluate_drafts(&drafts)?; - let verified_end = self.position; + let baseline_tokens = self.tokens.len(); + let batch_start = self.position; + let baseline_ple_history = self.ple_state.history; + let mut verify_tokens = Vec::with_capacity(drafts.len() + 1); + verify_tokens.push(token); + verify_tokens.extend(drafts.iter().map(|draft| draft.token)); + let (verify_logits, verify_probabilities) = + self.eval_target_batch(&verify_tokens, Some((temperature, top_p, min_p, top_k)))?; + let verify_probabilities = verify_probabilities.unwrap_or_else(|| { + verify_logits + .iter() + .map(|logits| qwen_sampling_probabilities(logits, temperature, top_p, min_p, top_k)) + .collect() + }); let mut accepted = 0; let mut replacement = None; - for ((draft, target), index) in drafts.iter().zip(&target_logits).zip(0..) { - let (sampled, was_draft) = exact_speculative_sample( + for ((draft, target), index) in drafts.iter().zip(&verify_probabilities).zip(0..) { + let (sampled, was_draft) = exact_speculative_sample_from_probabilities( target, - &draft.logits, + &draft.probabilities, draft.token, - temperature, - top_p, - min_p, - top_k, rng, ); if !was_draft { @@ -1986,35 +4540,101 @@ impl QwenExecutor { } accepted += 1; } + if std::env::var_os("DS4SERVER_QWEN38_MTP_PROFILE").is_some() { + eprintln!( + "ds4: Qwen MTP cycle primary={token} folded={folded_primary} drafts={:?} accepted={accepted} replacement={replacement:?}", + drafts.iter().map(|draft| draft.token).collect::>() + ); + } let mut emitted = Vec::with_capacity(accepted + 2); emitted.push(token); emitted.extend(drafts[..accepted].iter().map(|draft| draft.token)); if let Some((slot, replacement)) = replacement { let repair = Instant::now(); - let mut tokens = baseline_tokens; - tokens.extend(drafts[..accepted].iter().map(|draft| draft.token)); - self.restore_or_replay(slot, &tokens, verified_end, mtp_end, captures_complete)?; - if !self - .model - .is_stop_token_for_reasoning(replacement, reasoning) - { - let _ = self.mtp_step(replacement, self.position - 1, false)?; + let mut committed = Vec::with_capacity(accepted + 1); + committed.push(token); + committed.extend(drafts[..accepted].iter().map(|draft| draft.token)); + let captured = if verify_logits.is_empty() { + let committed_logits = self.logits_row(slot)?; + self.restore_or_replay( + &committed, + &committed_logits, + batch_start, + baseline_tokens, + baseline_ple_history, + true, + )? + } else { + self.restore_or_replay( + &committed, + &verify_logits[slot], + batch_start, + baseline_tokens, + baseline_ple_history, + true, + )? + }; + if std::env::var_os("DS4SERVER_QWEN38_MTP_PROFILE").is_some() { + eprintln!( + "ds4: Qwen verifier commit rows={} ms={:.3}", + slot + 1, + repair.elapsed().as_secs_f64() * 1_000.0 + ); } - self.eval_target(replacement)?; emitted.push(replacement); + let newly_emitted = emitted.len() - usize::from(folded_primary); + if captured + && newly_emitted < max_tokens as usize + && !self + .model + .is_stop_token_for_reasoning(replacement, reasoning) + { + self.pending_sampled = Some(replacement); + } else { + if !self + .model + .is_stop_token_for_reasoning(replacement, reasoning) + { + let _ = self.mtp_step(replacement, self.position - 1, false, None)?; + } + self.eval_target(replacement)?; + } if let Some(mtp) = &mut self.mtp { mtp.repair_ns = mtp .repair_ns .saturating_add(u64::try_from(repair.elapsed().as_nanos()).unwrap_or(u64::MAX)); } - } else if let Some(last) = drafts.last() - && !self + } else if drafts.last().is_some_and(|last| { + !self .model .is_stop_token_for_reasoning(last.token, reasoning) - { - self.commit_final_mtp_row(last, drafts.len(), captures_complete)?; + }) { + let commit = Instant::now(); + self.commit_final_mtp_batch_row(&drafts, batch_start)?; + self.read_logits_row(drafts.len())?; + if std::env::var_os("DS4SERVER_QWEN38_MTP_PROFILE").is_some() { + eprintln!( + "ds4: Qwen verifier full commit rows={} ms={:.3}", + drafts.len(), + commit.elapsed().as_secs_f64() * 1_000.0 + ); + } + let newly_emitted = emitted.len() - usize::from(folded_primary); + if newly_emitted < max_tokens as usize { + let bonus = sample_probabilities_f64(&verify_probabilities[drafts.len()], rng); + emitted.push(bonus); + self.pending_sampled = Some(bonus); + } + } else { + self.commit_verifier_recurrent_state()?; + self.canonicalize_batch_row(drafts.len())?; + self.read_logits_row(drafts.len())?; } self.note_mtp_verification(accepted, started); + if folded_primary { + debug_assert_eq!(emitted.first(), Some(&token)); + emitted.remove(0); + } Ok(emitted) } @@ -2023,25 +4643,127 @@ impl QwenExecutor { tokens: &[i32], mut progress: impl FnMut(u32) -> bool, ) -> Result { - let mut completed = 0; - for &token in tokens { - if !progress(completed) { - break; - } - self.eval(token)?; - completed += 1; + let end = self + .position + .checked_add(u32::try_from(tokens.len()).map_err(|_| "Qwen prefill is too large")?) + .ok_or("Qwen prefill position overflow")?; + if end > self.context { + return Err(format!( + "the Qwen executor supports {} tokens per session", + self.context + )); } - progress(completed); - Ok(completed as usize) + if tokens.len() <= MTP_CAPTURE_ROWS { + let mut completed = 0; + for &token in tokens { + if !progress(completed) { + break; + } + self.eval(token)?; + completed += 1; + } + progress(completed); + return Ok(completed as usize); + } + + let has_mtp = self.mtp.is_some(); + if self.position > 0 && has_mtp { + let _ = self.mtp_step(tokens[0], self.position - 1, false, None)?; + } + let body = tokens.len() - 1; + let capacity = body.clamp(1, PREFILL_CHUNK) as u32; + let batch_scratch = Scratch::new(self.context, capacity)?; + let decode_scratch = std::mem::replace(&mut self.scratch, batch_scratch); + let mut completed = 0; + let result = (|| { + while completed < body { + if !progress(completed as u32) { + break; + } + let rows = (body - completed).min(PREFILL_CHUNK); + self.eval_prefill_batch( + &tokens[completed..completed + rows], + has_mtp.then_some(&tokens[completed + 1..completed + rows + 1]), + )?; + completed += rows; + } + if completed == body && progress(completed as u32) { + self.eval_target(tokens[body])?; + completed += 1; + let commands = Commands::begin()?; + self.scratch.batch_final_hc.copy_from( + 0, + &self.scratch.hc, + 0, + u64::from(HC_WIDTH) * 4, + "preserving the final scalar Qwen prefill hidden row", + )?; + commands.finish()?; + } + Ok::<(), String>(()) + })(); + let batch_scratch = std::mem::replace(&mut self.scratch, decode_scratch); + if completed != 0 { + let commands = Commands::begin()?; + self.scratch.hc.copy_from( + 0, + &batch_scratch.batch_final_hc, + 0, + u64::from(HC_WIDTH) * 4, + "restoring the final Qwen prefill hidden row", + )?; + commands.finish()?; + } + progress(completed as u32); + result.map(|()| completed) } pub(super) fn logits(&self) -> &[f32] { &self.logits } + pub(super) fn take_pending_sampled_token(&mut self) -> Option { + let pending = self.pending_sampled.take(); + self.folding_sampled = pending.is_some(); + pending + } + + pub(super) fn finalize_pending_sampled_token(&mut self) -> Result<(), String> { + self.folding_sampled = false; + let Some(token) = self.pending_sampled.take() else { + return Ok(()); + }; + let result = (|| { + if self.position > 0 { + let _ = self.mtp_step(token, self.position - 1, false, None)?; + } + self.eval_target(token) + })(); + if result.is_err() { + self.pending_sampled = Some(token); + } + result + } + pub(super) fn execution_stats(&self) -> ExecutionStats { - self.mtp - .as_ref() - .map_or_else(ExecutionStats::default, |mtp| ExecutionStats { + let ple = self.model.ple_stats(); + self.mtp.as_ref().map_or_else( + || ExecutionStats { + ssd_enabled: true, + ssd_cache_bytes: ple.cache_bytes, + ssd_cache_entries: ple.entries, + ssd_cache_hits: ple.hits, + ssd_cache_misses: ple.misses, + ssd_cache_evictions: ple.evictions, + ssd_pread_bytes: ple.read_bytes, + ssd_pread_ms: ple.read_ms, + ssd_selected_requests: ple.hits.saturating_add(ple.misses), + ssd_requested_bytes: ple + .hits + .saturating_add(ple.misses) + .saturating_mul(PLE_ROW_BYTES as u64), + ..ExecutionStats::default() + }, + |mtp| ExecutionStats { speculative_mode: 4, speculative_cycles: mtp.cycles, drafted_tokens: mtp.drafted, @@ -2051,8 +4773,22 @@ impl QwenExecutor { verifier_passes: mtp.verifier_passes, verifier_ms: mtp.verifier_ns / 1_000_000, repair_ms: mtp.repair_ns / 1_000_000, + ssd_enabled: true, + ssd_cache_bytes: ple.cache_bytes, + ssd_cache_entries: ple.entries, + ssd_cache_hits: ple.hits, + ssd_cache_misses: ple.misses, + ssd_cache_evictions: ple.evictions, + ssd_pread_bytes: ple.read_bytes, + ssd_pread_ms: ple.read_ms, + ssd_selected_requests: ple.hits.saturating_add(ple.misses), + ssd_requested_bytes: ple + .hits + .saturating_add(ple.misses) + .saturating_mul(PLE_ROW_BYTES as u64), ..ExecutionStats::default() - }) + }, + ) } pub(super) fn model(&self) -> &QwenModel { &self.model @@ -2082,6 +4818,8 @@ impl QwenExecutor { self.logits.fill(0.0); self.tokens.clear(); self.position = 0; + self.pending_sampled = None; + self.folding_sampled = false; self.checkpoint_tag = [0; 32]; Ok(()) } @@ -2097,6 +4835,7 @@ impl QwenExecutor { Ok(QwenResidentState { states: allocate_states(self.context)?, ple_state: allocate_ple_state()?, + hidden: Buffer::floats(HC_WIDTH.into())?, logits: vec![0.0; VOCAB as usize], tokens: Vec::new(), position: 0, @@ -2106,6 +4845,8 @@ impl QwenExecutor { .as_ref() .map(|mtp| allocate_mtp(self.context, mtp.timing)) .transpose()?, + pending_sampled: None, + folding_sampled: false, }) } @@ -2116,11 +4857,14 @@ impl QwenExecutor { let mut incoming = state.take().map_or_else(|| self.blank_resident(), Ok)?; std::mem::swap(&mut self.states, &mut incoming.states); std::mem::swap(&mut self.ple_state, &mut incoming.ple_state); + std::mem::swap(&mut self.scratch.hc, &mut incoming.hidden); std::mem::swap(&mut self.logits, &mut incoming.logits); std::mem::swap(&mut self.tokens, &mut incoming.tokens); std::mem::swap(&mut self.position, &mut incoming.position); std::mem::swap(&mut self.checkpoint_tag, &mut incoming.checkpoint_tag); std::mem::swap(&mut self.mtp, &mut incoming.mtp); + std::mem::swap(&mut self.pending_sampled, &mut incoming.pending_sampled); + std::mem::swap(&mut self.folding_sampled, &mut incoming.folding_sampled); *state = Some(incoming); Ok(()) } @@ -2161,6 +4905,14 @@ impl QwenExecutor { write_u32(&mut file, logit.to_bits())?; } let mut chunk = vec![0; CHECKPOINT_CHUNK]; + write_buffer( + &mut file, + &self.scratch.hc, + 0, + u64::from(HC_WIDTH) * 4, + &mut chunk, + progress, + )?; write_buffer( &mut file, &self.ple_state.conv, @@ -2211,7 +4963,8 @@ impl QwenExecutor { else { return Err("Qwen MTP checkpoint state is invalid".into()); }; - let [kv_bytes, raw_bytes, pooled_bytes] = attention_checkpoint_bytes(self.position); + let [kv_bytes, raw_bytes, pooled_bytes] = + attention_checkpoint_bytes(self.position.saturating_sub(1)); write_buffer(&mut file, kv, 0, kv_bytes, &mut chunk, progress)?; write_buffer(&mut file, qsa_raw, 0, raw_bytes, &mut chunk, progress)?; write_buffer(&mut file, qsa_pooled, 0, pooled_bytes, &mut chunk, progress)?; @@ -2285,6 +5038,14 @@ impl QwenExecutor { } self.reset()?; let mut chunk = vec![0; CHECKPOINT_CHUNK]; + read_buffer( + &mut file, + &self.scratch.hc, + 0, + u64::from(HC_WIDTH) * 4, + &mut chunk, + progress, + )?; read_buffer( &mut file, &self.ple_state.conv, @@ -2334,7 +5095,8 @@ impl QwenExecutor { else { return Err("Qwen MTP checkpoint state is invalid".into()); }; - let [kv_bytes, raw_bytes, pooled_bytes] = attention_checkpoint_bytes(position); + let [kv_bytes, raw_bytes, pooled_bytes] = + attention_checkpoint_bytes(position.saturating_sub(1)); read_buffer(&mut file, kv, 0, kv_bytes, &mut chunk, progress)?; read_buffer(&mut file, qsa_raw, 0, raw_bytes, &mut chunk, progress)?; read_buffer(&mut file, qsa_pooled, 0, pooled_bytes, &mut chunk, progress)?; @@ -2396,6 +5158,7 @@ fn allocate_attention_state(context: u32) -> Result { }) } +#[cfg(test)] fn allocate_recurrent_snapshot() -> Result { let gdn = (0..LAYERS) .map(|layer| { @@ -2420,12 +5183,40 @@ fn allocate_recurrent_snapshot() -> Result { }) } +fn allocate_verify_capture() -> Result { + let gdn = (0..LAYERS) + .map(|layer| { + if layer % 4 == 3 { + Ok(None) + } else { + Ok(Some(GdnVerifyCapture { + conv: Buffer::bytes(u64::from(GDN_QKV) * 3 * 2)?, + recurrent: Buffer::bytes( + u64::from(GDN_HEADS_V) * HEAD_DIM as u64 * HEAD_DIM as u64 * 4, + )?, + qkv: Buffer::floats(VERIFY_ROWS * u64::from(GDN_QKV))?, + q: Buffer::floats(VERIFY_ROWS * u64::from(GDN_HEADS_K * HEAD_DIM))?, + k: Buffer::floats(VERIFY_ROWS * u64::from(GDN_HEADS_K * HEAD_DIM))?, + v: Buffer::floats(VERIFY_ROWS * u64::from(GDN_VALUE))?, + controls: Buffer::floats(VERIFY_ROWS * u64::from(GDN_CONTROLS))?, + })) + } + }) + .collect::>()?; + Ok(VerifyCapture { + gdn, + ple_conv: Buffer::bytes(u64::from(HC_WIDTH) * PLE_CONV_STATE as u64 * 2)?, + ple_input: Buffer::floats(VERIFY_ROWS * u64::from(HC_WIDTH))?, + hidden: Buffer::floats(VERIFY_ROWS * u64::from(HC_WIDTH))?, + }) +} + fn allocate_mtp(context: u32, timing: bool) -> Result { Ok(QwenMtp { attention: allocate_attention_state(context)?, - captures: (0..MTP_CAPTURE_ROWS) - .map(|_| allocate_recurrent_snapshot()) - .collect::>()?, + #[cfg(test)] + snapshot: allocate_recurrent_snapshot()?, + verify: allocate_verify_capture()?, timing, cycles: 0, drafted: 0, @@ -2438,39 +5229,6 @@ fn allocate_mtp(context: u32, timing: bool) -> Result { }) } -fn clear_attention_rows(state: &LayerState, start: u32, end: u32) -> Result<(), String> { - if end <= start { - return Ok(()); - } - let LayerState::Attention { - kv, - qsa_raw, - qsa_pooled, - } = state - else { - return Err("Qwen attention trim received GDN state".into()); - }; - fn clear(buffer: &Buffer, row_bytes: u64, start: u32, end: u32) -> Result<(), String> { - let bytes = u64::from(end - start) - .checked_mul(row_bytes) - .ok_or_else(|| "Qwen attention trim size overflows".to_owned())?; - let zeros = vec![ - 0; - usize::try_from(bytes) - .map_err(|_| "Qwen attention trim exceeds addressable memory")? - ]; - buffer.write(u64::from(start) * row_bytes, &zeros) - } - clear(kv, u64::from(ATTN_KV_WIDTH) * 4, start, end)?; - clear(qsa_raw, u64::from(QSA_DIM) * 2, start, end)?; - clear( - qsa_pooled, - u64::from(QSA_DIM) * 2, - start / QSA_RATIO, - end / QSA_RATIO, - ) -} - fn qsa_block_capacity(context: u32) -> Result { context .checked_add(MTP_TOKEN_RESERVE) @@ -2633,6 +5391,7 @@ fn advance_ple_history(history: [i32; PLE_HISTORY], token: i32) -> [i32; PLE_HIS } } +#[cfg(test)] fn gather_ple_row( row: u64, packed: &[u8], @@ -2647,6 +5406,7 @@ fn gather_ple_row( Ok(value) } +#[cfg(test)] fn ple_row_slice(data: &[u8], row: usize, width: usize) -> Result<&[u8], String> { let start = row .checked_mul(width) @@ -2659,6 +5419,35 @@ fn args() -> QwenKernelArgs { QwenKernelArgs::default() } +fn qwen_feature_enabled(name: &str) -> bool { + std::env::var(name).map_or(true, |value| { + matches!( + value.trim().to_ascii_lowercase().as_str(), + "1" | "true" | "yes" | "on" + ) + }) +} + +fn hyper_pack_bytes(rows: u32, columns: u32) -> u64 { + let values = u64::from(rows) * u64::from(columns); + let groups = values / u64::from(HYPER_QUANT_GROUP); + values + groups * 4 +} + +fn rms_args(width: u32, group_size: u32) -> QwenKernelArgs { + let mut values = args(); + values.u[0] = width; + values.u[1] = group_size; + values.u[10] = 33 * 4; + values.u[12] = if group_size <= 4096 { + group_size.div_ceil(4).div_ceil(32) * 32 + } else { + 1024 + }; + values.f[0] = 1.0e-6; + values +} + #[allow(clippy::too_many_arguments)] fn dispatch_qwen( kernel: &CStr, @@ -2670,6 +5459,32 @@ fn dispatch_qwen( values: &QwenKernelArgs, grid_x: u32, grid_y: u32, +) -> Result<(), String> { + dispatch_qwen_with_d(kernel, out, a, b, c, None, weights, values, grid_x, grid_y) +} + +fn buffer_weight_view(buffer: &Buffer, bytes: u64) -> QwenWeightView { + QwenWeightView { + tensor: buffer.raw().cast_const(), + map: std::ptr::null(), + size: bytes, + offset: 0, + bytes, + } +} + +#[allow(clippy::too_many_arguments)] +fn dispatch_qwen_with_d( + kernel: &CStr, + out: &Buffer, + a: Option<&Buffer>, + b: Option<&Buffer>, + c: Option<&Buffer>, + d: Option<&Buffer>, + weights: &[QwenWeightView], + values: &QwenKernelArgs, + grid_x: u32, + grid_y: u32, ) -> Result<(), String> { call( unsafe { @@ -2679,6 +5494,7 @@ fn dispatch_qwen( a.map_or(std::ptr::null(), |buffer| buffer.raw().cast_const()), b.map_or(std::ptr::null(), |buffer| buffer.raw().cast_const()), c.map_or(std::ptr::null(), |buffer| buffer.raw().cast_const()), + d.map_or(std::ptr::null(), |buffer| buffer.raw().cast_const()), weights.as_ptr(), weights.len() as u32, values, @@ -2914,17 +5730,18 @@ mod tests { bytes.extend_from_slice(&0x3333_3333_u32.to_le_bytes()); } for value in [ - 0.5, -1.0, 0.0, 0.0, 1.0, 2.0, 3.0, 4.0, 1.0, 1.0, 0.0, 0.0, 0.0, 0.0, + 0.5, -1.0, 0.0, 0.0, 1.0, 2.0, 3.0, 4.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, ] { bytes.extend_from_slice(&bf16(value).to_le_bytes()); } let qsa_weight_offset = bytes.len() as u64; - bytes.extend_from_slice(&bf16(0.0).to_le_bytes().repeat(QSA_DIM as usize)); + bytes.extend_from_slice(&bf16(1.0).to_le_bytes().repeat(QSA_DIM as usize)); fs::write(&path, bytes).unwrap(); let file = File::open(&path).unwrap(); // SAFETY: this test owns the read-only file for the lifetime of the mapping. let map = unsafe { MmapOptions::new().map(&file).unwrap() }; let view = |offset, size| QwenWeightView { + tensor: std::ptr::null(), map: map.as_ptr().cast(), size: map.len() as u64, offset, @@ -2972,7 +5789,10 @@ mod tests { ) .unwrap(); conv_output.read_f32(&mut scalar).unwrap(); - close(scalar[0], 8.0 / (1.0 + (-8.0_f32).exp())); + close( + scalar[0], + f32::from_bits(u32::from(bf16(8.0 / (1.0 + (-8.0_f32).exp()))) << 16), + ); assert_eq!( { let mut state = [0; 6]; @@ -2982,69 +5802,32 @@ mod tests { [0, 0, 0, 0, 0, 64] ); - let qkv = Buffer::floats(6).unwrap(); - qkv.write_f32(&[3.0, 4.0, 0.0, 2.0, 5.0, 7.0]).unwrap(); - let controls = Buffer::floats(4).unwrap(); - controls.write_f32(&[0.0; 4]).unwrap(); - let recurrent = Buffer::floats(4).unwrap(); - recurrent.write_f32(&[1.0, 2.0, 3.0, 4.0]).unwrap(); - let raw = Buffer::floats(2).unwrap(); - let gated = Buffer::floats(2).unwrap(); - let mut step = args(); - step.u[0] = 2; - step.u[1] = 1; - step.u[2] = 1; - step.u[3] = 2; - step.u[4] = 3; - step.f[0] = 1.0e-6; - dispatch_qwen( - c"kernel_qwen_gdn_step", - &raw, - Some(&qkv), - Some(&controls), - Some(&recurrent), - &[view(37, 2), view(39, 2)], - &step, - 2, - 1, - ) - .unwrap(); - let mut raw_values = [0.0; 2]; - raw.read_f32(&mut raw_values).unwrap(); - close(raw_values[0], 2.0506096); - close(raw_values[1], 2.9698484); - let mut norm = args(); - norm.u[0] = 2; - norm.u[1] = 1; - norm.f[0] = 1.0e-6; - dispatch_qwen( - c"kernel_qwen_gdn_norm_gate", - &gated, - Some(&raw), - Some(&controls), - None, - &[view(49, 4)], - &norm, - 1, - 1, - ) - .unwrap(); - let mut gated_values = [0.0; 2]; - gated.read_f32(&mut gated_values).unwrap(); - let rms = ((raw_values[0].powi(2) + raw_values[1].powi(2)) * 0.5 + 1.0e-6).sqrt(); - close(gated_values[0], raw_values[0] / rms * 0.5); - close(gated_values[1], raw_values[1] / rms * 0.5); - let router = Buffer::floats(EXPERTS.into()).unwrap(); let mut logits = vec![-100.0; EXPERTS as usize]; for (index, value) in logits.iter_mut().take(EXPERTS_USED).enumerate() { *value = index as f32; } router.write_f32(&logits).unwrap(); + let mut softmax = args(); + softmax.u[0] = EXPERTS; + softmax.u[12] = EXPERTS / 4; + dispatch_qwen( + c"kernel_qwen_softmax_precise_512", + &router, + Some(&router), + None, + None, + &[], + &softmax, + 1, + 1, + ) + .unwrap(); let ids = Buffer::bytes((EXPERTS_USED * 4) as u64).unwrap(); let weights = Buffer::floats(EXPERTS_USED as u64).unwrap(); let mut route = args(); route.u[0] = EXPERTS; + route.u[12] = EXPERTS; dispatch_qwen( c"kernel_qwen_route_top10", &ids, @@ -3059,18 +5842,29 @@ mod tests { .unwrap(); let mut selected = [0; EXPERTS_USED]; ids.read_i32(&mut selected).unwrap(); - assert_eq!(selected, [9, 8, 7, 6, 5, 4, 3, 2, 1, 0]); + assert_eq!(selected, [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]); let mut probabilities = [0.0; EXPERTS_USED]; weights.read_f32(&mut probabilities).unwrap(); - close(probabilities.iter().sum(), 1.0); + assert_eq!( + probabilities, + [ + 7.8201294e-5, + 0.0002117157, + 0.0005760193, + 0.0015640259, + 0.004272461, + 0.01159668, + 0.03149414, + 0.08544922, + 0.23242188, + 0.6328125, + ] + ); let norm_input = Buffer::floats(4).unwrap(); norm_input.write_f32(&[1.0, 2.0, 3.0, 4.0]).unwrap(); let norm_output = Buffer::floats(4).unwrap(); - let mut zero_norm = args(); - zero_norm.u[0] = 4; - zero_norm.u[1] = 2; - zero_norm.f[0] = 1.0e-6; + let zero_norm = rms_args(4, 2); dispatch_qwen( c"kernel_qwen_zero_rms", &norm_output, @@ -3091,7 +5885,7 @@ mod tests { 3.0 / (12.5_f32 + 1.0e-6).sqrt(), 4.0 / (12.5_f32 + 1.0e-6).sqrt(), ]) { - close(actual, expected); + close(actual, f32::from_bits(u32::from(bf16(expected)) << 16)); } let hyper_input = Buffer::floats(8).unwrap(); @@ -3192,11 +5986,7 @@ mod tests { .unwrap(); let mut moe_values = [0.0; 2]; accumulated.read_f32(&mut moe_values).unwrap(); - close(moe_values[0], 1.0); - close( - moe_values[1], - 0.25 * (2.0 / (1.0 + (-2.0_f32).exp()) * 3.0) + 2.0, - ); + assert_eq!(moe_values, [1.0, 3.3125]); let rope_input = Buffer::floats(4).unwrap(); rope_input.write_f32(&[1.0, 2.0, 3.0, 4.0]).unwrap(); @@ -3220,78 +6010,9 @@ mod tests { 1, ) .unwrap(); - let rms = (7.5_f32 + 1.0e-6).sqrt(); - let theta = [1.0_f32, 0.01]; - let normalized = [1.0 / rms, 2.0 / rms, 3.0 / rms, 4.0 / rms]; - let expected_rope = [ - normalized[0] * theta[0].cos() - normalized[2] * theta[0].sin(), - normalized[1] * theta[1].cos() - normalized[3] * theta[1].sin(), - normalized[2] * theta[0].cos() + normalized[0] * theta[0].sin(), - normalized[3] * theta[1].cos() + normalized[1] * theta[1].sin(), - ]; let mut actual_rope = [0.0; 4]; rope_output.read_f32(&mut actual_rope).unwrap(); - for (actual, expected) in actual_rope.into_iter().zip(expected_rope) { - close(actual, expected); - } - - let cache = Buffer::bytes(16).unwrap(); - let key = Buffer::floats(2).unwrap(); - let value = Buffer::floats(2).unwrap(); - let mut store = args(); - store.u[0] = 2; - for (position, (keys, values)) in [([1.0, 0.0], [2.0, 3.0]), ([0.0, 1.0], [5.0, 7.0])] - .into_iter() - .enumerate() - { - key.write_f32(&keys).unwrap(); - value.write_f32(&values).unwrap(); - store.u[1] = position as u32; - dispatch_qwen( - c"kernel_qwen_store_kv_bf16", - &cache, - Some(&key), - Some(&value), - None, - &[], - &store, - 2, - 1, - ) - .unwrap(); - } - let query = Buffer::floats(2).unwrap(); - query.write_f32(&[1.0, 0.0]).unwrap(); - let attention = Buffer::floats(2).unwrap(); - let mut dense = args(); - dense.u[0] = 1; - dense.u[1] = 1; - dense.u[2] = 2; - dense.u[3] = 2; - dispatch_qwen( - c"kernel_qwen_dense_attention", - &attention, - Some(&query), - Some(&cache), - None, - &[], - &dense, - 2, - 1, - ) - .unwrap(); - let first = (1.0_f32 / 2.0_f32.sqrt()).exp(); - let probability = first / (first + 1.0); - let mut actual_attention = [0.0; 2]; - attention.read_f32(&mut actual_attention).unwrap(); - close( - actual_attention[0], - probability * 2.0 + (1.0 - probability) * 5.0, - ); - close( - actual_attention[1], - probability * 3.0 + (1.0 - probability) * 7.0, - ); + assert_eq!(actual_rope, [-0.72265625, 0.71484375, 0.8984375, 1.46875]); let qsa_projected = Buffer::floats(QSA_WIDTH.into()).unwrap(); let qsa_query = Buffer::floats((QSA_HEADS * QSA_DIM).into()).unwrap(); @@ -3393,87 +6114,6 @@ mod tests { selected_blocks.read_i32(&mut sorted).unwrap(); assert_eq!(sorted, [1, 2, 3]); - let sparse_cache = Buffer::bytes(9 * 2 * 2).unwrap(); - for token in 0..9 { - key.write_f32(&[0.0, 0.0]).unwrap(); - value - .write_f32(&[token as f32 + 1.0, token as f32 + 11.0]) - .unwrap(); - store.u[1] = token; - dispatch_qwen( - c"kernel_qwen_store_kv_bf16", - &sparse_cache, - Some(&key), - Some(&value), - None, - &[], - &store, - 2, - 1, - ) - .unwrap(); - } - query.write_f32(&[0.0, 0.0]).unwrap(); - selected_blocks.write_i32(&[1]).unwrap(); - let mut sparse = args(); - sparse.u[0] = 1; - sparse.u[1] = 1; - sparse.u[2] = 2; - sparse.u[3] = 9; - sparse.u[4] = 1; - sparse.u[5] = 4; - sparse.u[6] = 8; - sparse.u[7] = 1; - dispatch_qwen( - c"kernel_qwen_sparse_attention", - &attention, - Some(&query), - Some(&sparse_cache), - Some(&selected_blocks), - &[], - &sparse, - 2, - 1, - ) - .unwrap(); - attention.read_f32(&mut actual_attention).unwrap(); - close(actual_attention[0], 7.0); - close(actual_attention[1], 17.0); - - selected_blocks.write_i32(&[0, 1]).unwrap(); - sparse.u[4] = 2; - dispatch_qwen( - c"kernel_qwen_sparse_attention", - &attention, - Some(&query), - Some(&sparse_cache), - Some(&selected_blocks), - &[], - &sparse, - 2, - 1, - ) - .unwrap(); - let dense_all = Buffer::floats(2).unwrap(); - dense.u[3] = 9; - dispatch_qwen( - c"kernel_qwen_dense_attention", - &dense_all, - Some(&query), - Some(&sparse_cache), - None, - &[], - &dense, - 2, - 1, - ) - .unwrap(); - let mut sparse_all = [0.0; 2]; - let mut dense_values = [0.0; 2]; - attention.read_f32(&mut sparse_all).unwrap(); - dense_all.read_f32(&mut dense_values).unwrap(); - assert_eq!(sparse_all, dense_values); - let ple_packed = Buffer::bytes(80).unwrap(); ple_packed.write(0, &[0x33; 80]).unwrap(); let ple_scales = Buffer::bytes(10).unwrap(); @@ -3506,35 +6146,6 @@ mod tests { ple_embedding.read_f32(&mut embedding).unwrap(); assert!(embedding.into_iter().all(|value| value == 0.5)); - let ple_key = Buffer::floats(8).unwrap(); - ple_key.write_f32(&[1.0; 8]).unwrap(); - let ple_query = Buffer::floats(8).unwrap(); - ple_query.write_f32(&[1.0; 8]).unwrap(); - let ple_value = Buffer::floats(2).unwrap(); - ple_value.write_f32(&[2.0, 3.0]).unwrap(); - let ple_gated = Buffer::floats(8).unwrap(); - let mut gate = args(); - gate.u[0] = 2; - dispatch_qwen( - c"kernel_qwen_ple_gate", - &ple_gated, - Some(&ple_key), - Some(&ple_query), - Some(&ple_value), - &[], - &gate, - 4, - 1, - ) - .unwrap(); - let expected_gate = 1.0 / (1.0 + (-2.0_f32.sqrt().sqrt()).exp()); - let mut gated = [0.0; 8]; - ple_gated.read_f32(&mut gated).unwrap(); - for stream in 0..4 { - close(gated[stream * 2], 2.0 * expected_gate); - close(gated[stream * 2 + 1], 3.0 * expected_gate); - } - let ple_normalized = Buffer::floats(1).unwrap(); ple_normalized.write_f32(&[2.0]).unwrap(); let ple_gate_value = Buffer::floats(1).unwrap(); @@ -3557,7 +6168,7 @@ mod tests { ) .unwrap(); ple_output.read_f32(&mut scalar).unwrap(); - close(scalar[0], 0.5 + 8.0 / (1.0 + (-8.0_f32).exp())); + assert_eq!(scalar[0], 8.5); let mut ple_history = [0; 18]; ple_state.read(0, &mut ple_history).unwrap(); assert_eq!(&ple_history[16..], &bf16(2.0).to_le_bytes()); @@ -3668,6 +6279,7 @@ mod tests { sparse.u[5] = QSA_RATIO; sparse.u[6] = tail_start; sparse.u[7] = depth - tail_start; + sparse.u[12] = 256; let started = Instant::now(); dispatch_qwen( c"kernel_qwen_sparse_attention", @@ -3677,8 +6289,8 @@ mod tests { Some(&selected), &[], &sparse, - ATTN_DIM, ATTN_HEADS, + 1, ) .unwrap(); let mut values = vec![1.0; ATTN_WIDTH as usize]; @@ -3700,7 +6312,14 @@ mod tests { .expect("set DS4SERVER_QWEN38_SOURCE to the pinned artifact directory"); let model = QwenModel::open(&root, 262_144).unwrap(); assert_eq!(model.memory().kv_and_recurrent, 7_564_812_288); - let executor = QwenExecutor::open(model, 262_144).unwrap(); + let executor = QwenExecutor::open_configured_with_hyper( + model, + 262_144, + EngineSpeculativeSettings::default(), + false, + false, + ) + .unwrap(); assert_eq!(executor.context(), 262_144); assert_eq!(qsa_block_capacity(executor.context()).unwrap(), 65_537); } @@ -3757,28 +6376,28 @@ mod tests { let oracles = [ ( 2_048, - "4cd0daf267d046d71f23291e756f81b2ebcb549712b441b2d1e8074b41e1da9f", - 44_496, + "b6d28407fc87ebef94c56cf8ece0e0be472d45d8b7a8ba6f08d17420412d7411", + 5, ), ( 16_384, - "f2513cb75775f0704a52b66053e5e47060234db57b3339e9de5a7d1ba6898bae", - 44_496, + "79d71772c9cffd81e84ea21f9f5a0f3d7e63575ef3196fdea8512fec2a286e8d", + 5, ), ( 65_536, - "f104bc1aa6c0e44fbfff2746300374d7c8edf59a6ee5eb43a5f5469a0219e637", - 197_597, + "0a1131fe9164b01265781b54a6dcca7da625e9b384a00015a779337a9a3ed35d", + 5, ), ( 131_072, - "660c5f431a6271748c2e4cb326eed6af21a4e67f3937f71b78e635a25fce2669", - 180_094, + "fb678d10cca29f42322253cde1ed2d4bc9dcfa47b7b6383e5de0f509964677d7", + 5, ), ( 262_144, - "ead42c990931a88a4196bdc2593013fd471e4f1559e80ccb1779d662e2d08725", - 19_559, + "01ae755c3366805aa54ccaf11d226c1f90b216b8bd05e910807bfb552a810bc3", + 5, ), ]; let mut baseline = None; @@ -3828,9 +6447,18 @@ mod tests { attention_checkpoint_bytes(depth).into_iter().sum::() * 12; let prefill_tps = 3.0 / prefill_seconds; let decode_tps = 1.0 / decode_seconds; + eprintln!( + "Qwen observed depth {depth}: prefill {prefill_tps:.3} tok/s, AR {decode_tps:.3} tok/s" + ); if let Some((baseline_prefill, baseline_decode)) = baseline { - assert!(prefill_tps >= baseline_prefill * 0.9); - assert!(decode_tps >= baseline_decode * 0.9); + assert!( + prefill_tps >= baseline_prefill * 0.9, + "Qwen prefill collapsed at depth {depth}: {prefill_tps:.3} versus {baseline_prefill:.3} tok/s", + ); + assert!( + decode_tps >= baseline_decode * 0.9, + "Qwen AR collapsed at depth {depth}: {decode_tps:.3} versus {baseline_decode:.3} tok/s", + ); } else { baseline = Some((prefill_tps, decode_tps)); } @@ -3870,7 +6498,7 @@ mod tests { .unwrap(); assert_eq!( fs::metadata(&checkpoint).unwrap().len(), - 116_635_752 + u64::from(depth) * 4 + live_attention_bytes + 116_676_712 + u64::from(depth) * 4 + live_attention_bytes ); assert!(executor.load_checkpoint(&checkpoint, &mut |_| {}).unwrap()); assert_eq!(executor.position, depth); @@ -3928,6 +6556,23 @@ mod tests { residency_before, residency_after ); assert!(executor.logits.iter().all(|value| value.is_finite())); + let observed_next = executor + .logits + .iter() + .enumerate() + .max_by(|a, b| a.1.total_cmp(b.1)) + .unwrap() + .0; + eprintln!("Qwen token-1 argmax after projection dispatch: {observed_next}"); + eprintln!( + "Qwen token-1 logit samples: {:?}", + [ + executor.logits[0], + executor.logits[1], + executor.logits[1000], + executor.logits[VOCAB as usize - 1], + ] + ); let mut digest = Sha256::new(); for value in &executor.logits { digest.update(value.to_bits().to_le_bytes()); @@ -4092,6 +6737,2458 @@ mod tests { fs::remove_file(checkpoint).unwrap(); } + #[test] + #[ignore = "diagnostic trace against the pinned MLX-LM reference"] + fn qwen_reference_stateful_trace() { + configure_sources().unwrap(); + let root = std::env::var_os("DS4SERVER_QWEN38_SOURCE") + .map(PathBuf::from) + .expect("set DS4SERVER_QWEN38_SOURCE to the pinned artifact directory"); + let mut executor = QwenExecutor::open(QwenModel::open(&root, 4).unwrap(), 4).unwrap(); + let save = |name: &str, buffer: &Buffer, width: usize| { + let mut values = vec![0.0; width]; + buffer.read_f32(&mut values).unwrap(); + let bytes = values + .iter() + .flat_map(|value| value.to_le_bytes()) + .collect::>(); + std::fs::write(format!("/private/tmp/ds4_t2_{name}.f32"), bytes).unwrap(); + let mut hash = Sha256::new(); + for value in values { + hash.update(value.to_le_bytes()); + } + eprintln!("{name}: {:02x?}", hash.finalize()); + }; + + executor.eval_target(1).unwrap(); + executor.eval_target(2).unwrap(); + executor.begin_token(2972).unwrap(); + for layer in 0..LAYERS { + if layer == 1 { + let commands = Commands::begin().unwrap(); + executor + .dispatch( + c"kernel_qwen_zero_rms", + &executor.scratch.hc_mix, + Some(&executor.scratch.hc), + None, + None, + &[executor.view( + executor + .weight("language_model.model.layers.1.ple.norm_query.weight") + .unwrap(), + )], + &rms_args(HC_WIDTH, HIDDEN), + HC, + 1, + ) + .unwrap(); + commands.finish().unwrap(); + save( + "ple_query_before", + &executor.scratch.hc_mix, + HC_WIDTH as usize, + ); + executor.ple(2972).unwrap(); + save("after_ple", &executor.scratch.hc, HC_WIDTH as usize); + save( + "ple_output", + &executor.scratch.ple_output, + HC_WIDTH as usize, + ); + save( + "ple_embedding", + &executor.scratch.ple_embedding, + HIDDEN as usize, + ); + save("ple_key", &executor.scratch.ple_key, HC_WIDTH as usize); + save("ple_value", &executor.scratch.ple_value, HIDDEN as usize); + save("ple_query", &executor.scratch.hc_norm, HC_WIDTH as usize); + save("ple_gated", &executor.scratch.ple_gated, HC_WIDTH as usize); + save("ple_norm", &executor.scratch.ple_norm, HC_WIDTH as usize); + } + if layer == 3 { + let prefix = "language_model.model.layers.3"; + let commands = Commands::begin().unwrap(); + executor + .hyper_read(&format!("{prefix}.attn_hyper_connection")) + .unwrap(); + commands.finish().unwrap(); + save( + "layer3_hc_norm", + &executor.scratch.hc_norm, + HC_WIDTH as usize, + ); + save("layer3_hc_rank", &executor.scratch.rank, HC_RANK as usize); + save("layer3_hc_mix", &executor.scratch.hc_mix, HC_WIDTH as usize); + save("layer3_mixed", &executor.scratch.block, HIDDEN as usize); + save("layer3_inject", &executor.scratch.injection, HC as usize); + let commands = Commands::begin().unwrap(); + executor + .attention(prefix, &executor.states[3], executor.position) + .unwrap(); + commands.finish().unwrap(); + save( + "layer3_q_packed", + &executor.scratch.q_packed, + (ATTN_WIDTH * 2) as usize, + ); + save( + "layer3_q_gate", + &executor.scratch.q_gate, + ATTN_WIDTH as usize, + ); + save("layer3_k", &executor.scratch.k, ATTN_KV_WIDTH as usize); + save("layer3_v", &executor.scratch.v, ATTN_KV_WIDTH as usize); + save( + "layer3_k_rope", + &executor.scratch.k_rope, + ATTN_KV_WIDTH as usize, + ); + save( + "layer3_raw_attention", + &executor.scratch.q, + ATTN_WIDTH as usize, + ); + save( + "layer3_gated_attention", + &executor.scratch.attention, + ATTN_WIDTH as usize, + ); + save("layer3_output", &executor.scratch.hidden, HIDDEN as usize); + let commands = Commands::begin().unwrap(); + executor + .head_norm_rope( + &executor.scratch.k, + &executor.scratch.qsa_qk, + executor + .weight("language_model.model.layers.3.self_attn.k_norm.weight") + .unwrap(), + ATTN_KV_HEADS, + 0, + ) + .unwrap(); + commands.finish().unwrap(); + save( + "layer3_k_norm", + &executor.scratch.qsa_qk, + ATTN_KV_WIDTH as usize, + ); + let commands = Commands::begin().unwrap(); + let mut norm_only = args(); + norm_only.u[0] = ATTN_DIM; + norm_only.u[1] = 64; + norm_only.u[2] = ATTN_KV_HEADS; + norm_only.u[11] = 1; + norm_only.u[12] = 32; + norm_only.f[0] = 1.0e-6; + norm_only.f[1] = 10_000_000.0; + executor + .dispatch( + c"kernel_qwen_head_norm_rope", + &executor.scratch.qsa_qk, + Some(&executor.scratch.k), + None, + None, + &[executor.view( + executor + .weight("language_model.model.layers.3.self_attn.k_norm.weight") + .unwrap(), + )], + &norm_only, + 1, + ATTN_KV_HEADS, + ) + .unwrap(); + commands.finish().unwrap(); + save( + "layer3_k_unweighted", + &executor.scratch.qsa_qk, + ATTN_KV_WIDTH as usize, + ); + let commands = Commands::begin().unwrap(); + executor.hyper_write().unwrap(); + commands.finish().unwrap(); + save( + "layer3_after_attention", + &executor.scratch.hc, + HC_WIDTH as usize, + ); + executor.encode_moe(prefix, "3").unwrap(); + } else { + let commands = Commands::begin().unwrap(); + executor.encode_layer(layer).unwrap(); + commands.finish().unwrap(); + } + save( + &format!("layer{layer}"), + &executor.scratch.hc, + HC_WIDTH as usize, + ); + } + save("final_widened", &executor.scratch.hc, HC_WIDTH as usize); + executor.final_output().unwrap(); + save("final_mixed", &executor.scratch.block, HIDDEN as usize); + save("logits", &executor.scratch.logits, VOCAB as usize); + } + + #[test] + #[ignore = "diagnostic trace against the pinned MTPLX verifier"] + fn qwen_reference_verify_trace() { + configure_sources().unwrap(); + let root = std::env::var_os("DS4SERVER_QWEN38_SOURCE") + .map(PathBuf::from) + .expect("set DS4SERVER_QWEN38_SOURCE to the pinned artifact directory"); + let model = QwenModel::open_configured(&root, 8, true).unwrap(); + let settings = EngineSpeculativeSettings { + glm_mtp: true, + ..EngineSpeculativeSettings::default() + }; + let mut executor = QwenExecutor::open_configured(model, 8, settings).unwrap(); + let save = |name: &str, buffer: &Buffer, width: usize| { + let mut values = vec![0.0; width]; + buffer.read_f32(&mut values).unwrap(); + let bytes = values + .iter() + .flat_map(|value| value.to_le_bytes()) + .collect::>(); + std::fs::write(format!("/private/tmp/ds4_verify_{name}.f32"), bytes).unwrap(); + }; + let save_i32 = |name: &str, buffer: &Buffer, width: usize| { + let mut values = vec![0_i32; width]; + buffer.read_i32(&mut values).unwrap(); + let bytes = values + .iter() + .flat_map(|value| value.to_le_bytes()) + .collect::>(); + std::fs::write(format!("/private/tmp/ds4_verify_{name}.i32"), bytes).unwrap(); + }; + executor.eval_target(1).unwrap(); + executor.eval_target(2).unwrap(); + let tokens = [2972, 44934, 28986, 31524]; + executor.begin_tokens(&tokens).unwrap(); + save("embedding", &executor.scratch.hc, HC_WIDTH as usize * 4); + for layer in 0..LAYERS { + let prefix = format!("language_model.model.layers.{layer}"); + if layer == 1 { + executor.ple_batch(&tokens).unwrap(); + } + executor + .hyper_read_rows(&format!("{prefix}.attn_hyper_connection"), 4) + .unwrap(); + if layer == 0 { + save( + "layer0_attn_hc_norm", + &executor.scratch.hc_norm, + HC_WIDTH as usize * 4, + ); + save( + "layer0_attn_hc_rank", + &executor.scratch.rank, + HC_RANK as usize * 4, + ); + save( + "layer0_attn_hc_mix", + &executor.scratch.hc_mix, + HC_WIDTH as usize * 4, + ); + save( + "layer0_attn_mixed", + &executor.scratch.block, + HIDDEN as usize * 4, + ); + save( + "layer0_attn_inject", + &executor.scratch.injection, + HC as usize * 4, + ); + } + match &executor.states[layer] { + LayerState::Gdn { .. } => executor.gdn_rows(&prefix, layer, 4, true).unwrap(), + state @ LayerState::Attention { .. } => executor + .attention_rows(&prefix, state, executor.position, 4) + .unwrap(), + } + if layer == 11 { + save( + "layer11_attn_mixed", + &executor.scratch.block, + HIDDEN as usize * 4, + ); + save( + "layer11_q_packed", + &executor.scratch.q_packed, + (ATTN_WIDTH * 2) as usize * 4, + ); + save("layer11_k", &executor.scratch.k, ATTN_KV_WIDTH as usize * 4); + save("layer11_v", &executor.scratch.v, ATTN_KV_WIDTH as usize * 4); + save( + "layer11_index_qk", + &executor.scratch.qsa_qk, + QSA_WIDTH as usize * 4, + ); + save( + "layer11_attention", + &executor.scratch.q, + ATTN_WIDTH as usize * 4, + ); + save( + "layer11_gated", + &executor.scratch.attention, + ATTN_WIDTH as usize * 4, + ); + save( + "layer11_attn_output", + &executor.scratch.hidden, + HIDDEN as usize * 4, + ); + } + if layer == 3 { + save( + "layer3_attn_mixed", + &executor.scratch.block, + HIDDEN as usize * 4, + ); + save( + "layer3_q_packed", + &executor.scratch.q_packed, + (ATTN_WIDTH * 2) as usize * 4, + ); + save("layer3_k", &executor.scratch.k, ATTN_KV_WIDTH as usize * 4); + save("layer3_v", &executor.scratch.v, ATTN_KV_WIDTH as usize * 4); + save( + "layer3_index_qk", + &executor.scratch.qsa_qk, + QSA_WIDTH as usize * 4, + ); + save( + "layer3_attn_output", + &executor.scratch.hidden, + HIDDEN as usize * 4, + ); + save( + "layer3_attention_batch", + &executor.scratch.q, + ATTN_WIDTH as usize * 4, + ); + save( + "layer3_gated_batch", + &executor.scratch.attention, + ATTN_WIDTH as usize * 4, + ); + let LayerState::Attention { kv, .. } = &executor.states[layer] else { + unreachable!() + }; + let mut packed = vec![0_u8; 6 * ATTN_KV_WIDTH as usize * 4]; + kv.read(0, &mut packed).unwrap(); + for (name, value_offset) in [ + ("layer3_cache_keys", 0), + ("layer3_cache_values", ATTN_KV_WIDTH as usize), + ] { + let mut values = Vec::with_capacity(6 * ATTN_KV_WIDTH as usize); + for head in 0..ATTN_KV_HEADS as usize { + for token in 0..6_usize { + for column in 0..ATTN_DIM as usize { + let index = token * ATTN_KV_WIDTH as usize * 2 + + value_offset + + head * ATTN_DIM as usize + + column; + let byte = index * 2; + values.push(f32::from_bits( + u32::from(u16::from_le_bytes([packed[byte], packed[byte + 1]])) + << 16, + )); + } + } + } + let bytes = values + .iter() + .flat_map(|value| value.to_le_bytes()) + .collect::>(); + std::fs::write(format!("/private/tmp/ds4_verify_{name}.f32"), bytes).unwrap(); + } + let mut split = args(); + split.u[0] = ATTN_HEADS; + split.u[1] = ATTN_DIM; + split.u[4] = 4; + executor + .dispatch( + c"kernel_qwen_split_q_gate", + &executor.scratch.q, + Some(&executor.scratch.q_packed), + Some(&executor.scratch.q_gate), + None, + &[], + &split, + ATTN_WIDTH * 4, + 1, + ) + .unwrap(); + save("layer3_q_raw", &executor.scratch.q, ATTN_WIDTH as usize * 4); + save( + "layer3_q_gate", + &executor.scratch.q_gate, + ATTN_WIDTH as usize * 4, + ); + save( + "layer3_k_rope", + &executor.scratch.k_rope, + ATTN_KV_WIDTH as usize * 4, + ); + executor + .head_norm_rope_rows( + &executor.scratch.q, + &executor.scratch.attention, + executor + .weight(&format!("{prefix}.self_attn.q_norm.weight")) + .unwrap(), + ATTN_HEADS, + executor.position, + 4, + ) + .unwrap(); + save( + "layer3_q_rope", + &executor.scratch.attention, + ATTN_WIDTH as usize * 4, + ); + let LayerState::Attention { kv, qsa_pooled, .. } = &executor.states[layer] else { + unreachable!() + }; + for row in 0..4_u32 { + let view = |buffer: &Buffer, width: u32| { + buffer + .view(u64::from(row) * u64::from(width) * 4, u64::from(width) * 4) + .unwrap() + }; + executor + .attend_row( + kv, + qsa_pooled, + executor.position + row, + &view(&executor.scratch.attention, ATTN_WIDTH), + &view(&executor.scratch.q, ATTN_WIDTH), + &view(&executor.scratch.qsa_qk, QSA_HEADS * QSA_DIM), + &view( + &executor.scratch.qsa_scores, + qsa_block_capacity(executor.context).unwrap(), + ), + &executor + .scratch + .qsa_selected + .view( + u64::from(row) * u64::from(QSA_TOP_K) * 4, + u64::from(QSA_TOP_K) * 4, + ) + .unwrap(), + ) + .unwrap(); + } + save( + "layer3_attention", + &executor.scratch.q, + ATTN_WIDTH as usize * 4, + ); + let mut gate = args(); + gate.u[0] = ATTN_WIDTH * 4; + executor + .dispatch( + c"kernel_qwen_gate_attention", + &executor.scratch.attention, + Some(&executor.scratch.q), + Some(&executor.scratch.q_gate), + None, + &[], + &gate, + ATTN_WIDTH * 4, + 1, + ) + .unwrap(); + save( + "layer3_gated", + &executor.scratch.attention, + ATTN_WIDTH as usize * 4, + ); + } + if layer == 0 { + save( + "layer0_gdn_qkv", + &executor.scratch.qkv, + GDN_QKV as usize * 4, + ); + save( + "layer0_gdn_controls", + &executor.scratch.controls, + GDN_CONTROLS as usize * 4, + ); + save( + "layer0_gdn_q", + &executor.scratch.gdn_q, + (GDN_HEADS_K * HEAD_DIM) as usize * 4, + ); + save( + "layer0_gdn_k", + &executor.scratch.gdn_k, + (GDN_HEADS_K * HEAD_DIM) as usize * 4, + ); + save( + "layer0_gdn_delta", + &executor.scratch.gdn_out, + GDN_VALUE as usize * 4, + ); + save( + "layer0_gdn_gated", + &executor.scratch.gdn_z, + GDN_VALUE as usize * 4, + ); + save( + "layer0_attn_output", + &executor.scratch.hidden, + HIDDEN as usize * 4, + ); + } + executor.hyper_write_rows(4).unwrap(); + if layer == 11 { + save( + "layer11_after_attn", + &executor.scratch.hc, + HC_WIDTH as usize * 4, + ); + } + if layer == 3 { + save( + "layer3_after_attn", + &executor.scratch.hc, + HC_WIDTH as usize * 4, + ); + } + if layer == 0 { + save( + "layer0_after_attn", + &executor.scratch.hc, + HC_WIDTH as usize * 4, + ); + executor + .hyper_read_rows(&format!("{prefix}.mlp_hyper_connection"), 4) + .unwrap(); + save( + "layer0_mlp_mixed", + &executor.scratch.block, + HIDDEN as usize * 4, + ); + save( + "layer0_mlp_inject", + &executor.scratch.injection, + HC as usize * 4, + ); + } + if layer == 0 { + executor + .affine_batch_into( + &executor + .affine(&format!("{prefix}.mlp.gate"), HIDDEN, EXPERTS, None) + .unwrap(), + &executor.scratch.block, + &executor.scratch.router, + HIDDEN, + EXPERTS, + 4, + true, + ) + .unwrap(); + save( + "layer0_router_logits", + &executor.scratch.router, + EXPERTS as usize * 4, + ); + let mut softmax = args(); + softmax.u[0] = EXPERTS; + softmax.u[12] = EXPERTS / 4; + executor + .dispatch( + c"kernel_qwen_softmax_precise_512", + &executor.scratch.router, + Some(&executor.scratch.router), + None, + None, + &[], + &softmax, + 4, + 1, + ) + .unwrap(); + let mut route = args(); + route.u[0] = EXPERTS; + route.u[12] = EXPERTS; + executor + .dispatch( + c"kernel_qwen_route_top10", + &executor.scratch.route_ids, + Some(&executor.scratch.route_weights), + Some(&executor.scratch.router), + None, + &[], + &route, + 4, + 1, + ) + .unwrap(); + save( + "layer0_router_probabilities", + &executor.scratch.router, + EXPERTS as usize * 4, + ); + save_i32( + "layer0_route_ids", + &executor.scratch.route_ids, + EXPERTS_USED * 4, + ); + save( + "layer0_route_weights", + &executor.scratch.route_weights, + EXPERTS_USED * 4, + ); + executor.routed_experts_rows(&prefix, 4).unwrap(); + save( + "layer0_expert_gate", + &executor.scratch.gate, + EXPERT_WIDTH as usize * EXPERTS_USED * 4, + ); + save( + "layer0_expert_up", + &executor.scratch.up, + EXPERT_WIDTH as usize * EXPERTS_USED * 4, + ); + save( + "layer0_expert_mid", + &executor.scratch.mid, + EXPERT_WIDTH as usize * EXPERTS_USED * 4, + ); + save( + "layer0_expert_down", + &executor.scratch.expert, + HIDDEN as usize * EXPERTS_USED * 4, + ); + save("layer0_routed", &executor.scratch.moe, HIDDEN as usize * 4); + executor.shared_expert_rows(&prefix, 4).unwrap(); + save("layer0_shared_gate", &executor.scratch.gate, 4); + save( + "layer0_shared_up", + &executor.scratch.up, + EXPERT_WIDTH as usize * 4, + ); + save( + "layer0_shared_mid", + &executor.scratch.mid, + EXPERT_WIDTH as usize * 4, + ); + save( + "layer0_shared_down", + &executor.scratch.shared, + HIDDEN as usize * 4, + ); + save( + "layer0_mlp_output", + &executor.scratch.moe, + HIDDEN as usize * 4, + ); + let mut inject = args(); + inject.u[0] = HIDDEN; + inject.u[4] = 4; + executor + .dispatch( + c"kernel_qwen_hyper_inject", + &executor.scratch.hc_norm, + Some(&executor.scratch.hc), + Some(&executor.scratch.moe), + Some(&executor.scratch.injection), + &[], + &inject, + HC_WIDTH * 4, + 1, + ) + .unwrap(); + std::mem::swap(&mut executor.scratch.hc, &mut executor.scratch.hc_norm); + } else { + if layer == 3 { + executor + .hyper_read_rows(&format!("{prefix}.mlp_hyper_connection"), 4) + .unwrap(); + save( + "layer3_mlp_mixed", + &executor.scratch.block, + HIDDEN as usize * 4, + ); + } + executor.encode_moe_rows(&prefix, 4).unwrap(); + if layer == 3 { + save( + "layer3_mlp_output", + &executor.scratch.moe, + HIDDEN as usize * 4, + ); + } + if layer == 11 { + save( + "layer11_mlp_mixed", + &executor.scratch.block, + HIDDEN as usize * 4, + ); + save( + "layer11_router_probabilities", + &executor.scratch.router, + EXPERTS as usize * 4, + ); + save_i32( + "layer11_route_ids", + &executor.scratch.route_ids, + EXPERTS_USED * 4, + ); + save( + "layer11_route_weights", + &executor.scratch.route_weights, + EXPERTS_USED * 4, + ); + save( + "layer11_expert_down", + &executor.scratch.expert, + HIDDEN as usize * EXPERTS_USED * 4, + ); + save("layer11_routed", &executor.scratch.moe, HIDDEN as usize * 4); + save( + "layer11_shared_down", + &executor.scratch.shared, + HIDDEN as usize * 4, + ); + save( + "layer11_mlp_output", + &executor.scratch.moe, + HIDDEN as usize * 4, + ); + } + } + save( + &format!("layer{layer}"), + &executor.scratch.hc, + HC_WIDTH as usize * 4, + ); + } + executor.final_output_rows(4).unwrap(); + let mut logits = vec![0.0; VOCAB as usize * 4]; + executor.scratch.logits.read_f32(&mut logits).unwrap(); + let bytes = logits + .iter() + .flat_map(|value| value.to_le_bytes()) + .collect::>(); + std::fs::write("/private/tmp/ds4_verify_logits.f32", bytes).unwrap(); + } + + #[test] + #[ignore = "diagnostic comparison with captured MTPLX SDPA inputs"] + fn qwen_reference_sdpa_trace() { + configure_sources().unwrap(); + let read_f32 = |path: &str| { + fs::read(path) + .unwrap() + .chunks_exact(4) + .map(|bytes| f32::from_le_bytes(bytes.try_into().unwrap())) + .collect::>() + }; + let upload_bf16 = |values: &[f32]| { + let bytes = values + .iter() + .flat_map(|value| bf16(*value).to_le_bytes()) + .collect::>(); + let buffer = Buffer::bytes(bytes.len() as u64).unwrap(); + buffer.write(0, &bytes).unwrap(); + buffer + }; + let query = upload_bf16(&read_f32("/private/tmp/mtplx_verify_layer3_q_rope.f32")); + let keys = upload_bf16(&read_f32("/private/tmp/mtplx_verify_layer3_cache_keys.f32")); + let values = upload_bf16(&read_f32( + "/private/tmp/mtplx_verify_layer3_cache_values.f32", + )); + let mask_values = (0..4_u32) + .flat_map(|row| (0..6_u32).map(move |token| u8::from(token <= row + 2))) + .collect::>(); + let mask = Buffer::bytes(mask_values.len() as u64).unwrap(); + mask.write(0, &mask_values).unwrap(); + let out = Buffer::floats(u64::from(ATTN_WIDTH) * 4).unwrap(); + let mut values_args = args(); + values_args.u[0] = ATTN_HEADS; + values_args.u[1] = ATTN_KV_HEADS; + values_args.u[2] = ATTN_DIM; + values_args.u[3] = 6; + values_args.u[4] = 4; + values_args.u[5] = 2; + values_args.u[12] = 1024; + let commands = Commands::begin().unwrap(); + dispatch_qwen_with_d( + c"kernel_qwen_dense_attention_rows", + &out, + Some(&query), + Some(&keys), + Some(&values), + Some(&mask), + &[], + &values_args, + ATTN_HEADS, + 4, + ) + .unwrap(); + commands.finish().unwrap(); + let mut actual = vec![0.0; ATTN_WIDTH as usize * 4]; + out.read_f32(&mut actual).unwrap(); + let expected = read_f32("/private/tmp/mtplx_verify_layer3_attention.f32"); + let hash = |input: &[f32]| { + let mut digest = Sha256::new(); + for value in input { + digest.update(value.to_le_bytes()); + } + digest + .finalize() + .iter() + .map(|byte| format!("{byte:02x}")) + .collect::() + }; + eprintln!("expected={} actual={}", hash(&expected), hash(&actual)); + assert_eq!(actual, expected); + } + + #[test] + #[ignore = "diagnostic trace against the pinned MLX-LM reference"] + fn qwen_reference_layer_trace() { + configure_sources().unwrap(); + let root = std::env::var_os("DS4SERVER_QWEN38_SOURCE") + .map(PathBuf::from) + .expect("set DS4SERVER_QWEN38_SOURCE to the pinned artifact directory"); + let model = QwenModel::open(&root, 4).unwrap(); + let mut executor = QwenExecutor::open(model, 4).unwrap(); + let show = |label: &str, buffer: &Buffer, width: usize| { + let mut values = vec![0.0; width]; + buffer.read_f32(&mut values).unwrap(); + let mut digest = Sha256::new(); + for value in &values { + digest.update(value.to_bits().to_le_bytes()); + } + let digest = digest + .finalize() + .iter() + .map(|byte| format!("{byte:02x}")) + .collect::(); + let at = |index| values[index]; + if matches!( + label, + "layer-0-attn-hc-mix" + | "ple-embedding" + | "ple-key" + | "ple-value" + | "ple-query" + | "ple-gated" + | "ple-norm" + | "ple-output" + | "ple" + | "layer-0-attn-input" + | "layer-0-gdn-qkv" + | "layer-0-gdn-conv" + | "layer-0-gdn-controls" + | "layer-0-gdn-gated" + | "layer-0-attn-output" + | "layer-0-after-attn" + | "layer-0-mlp-hc-norm" + | "layer-0-mlp-hc-down" + | "layer-0-mlp-hc-rank" + | "layer-0-mlp-hc-up" + | "layer-0-mlp-hc-mix" + | "layer-0-mlp-hc-inject" + | "layer-0-mlp-mixed" + | "layer-0-router-logits" + | "layer-0-router-probabilities" + | "layer-0-expert-gate" + | "layer-0-expert-up" + | "layer-0-expert-mid" + | "layer-0-expert-down" + | "layer-0-routed" + | "layer-0-shared-gate-raw" + | "layer-0-shared-gate-proj" + | "layer-0-shared-up" + | "layer-0-shared-mid" + | "layer-0-shared-down" + | "layer-0-total" + | "layer-0-output" + | "layer-1" + | "layer-1-attn-mixed" + | "layer-1-attn-inject" + | "layer-1-gdn-qkv" + | "layer-1-gdn-z" + | "layer-1-gdn-b" + | "layer-1-gdn-a" + | "layer-1-attn-out" + | "layer-1-after-attn" + | "layer-1-mlp-hc-norm" + | "layer-1-mlp-hc-rank" + | "layer-1-mlp-hc-up" + | "layer-1-mlp-hc-mix" + | "layer-1-mlp-inject" + | "layer-1-mlp-mixed" + | "layer-1-mlp-out" + | "layer-2-attn-mixed" + | "layer-2-attn-inject" + | "layer-2-attn-out" + | "layer-2-after-attn" + | "layer-2-mlp-mixed" + | "layer-2-router-probabilities" + | "layer-2-route-weights" + | "layer-2-expert-gate" + | "layer-2-expert-up" + | "layer-2-expert-mid" + | "layer-2-expert-down" + | "layer-2-routed" + | "layer-2-shared-gate-raw" + | "layer-2-shared-up" + | "layer-2-shared-mid" + | "layer-2-shared-down" + | "layer-2-total" + | "layer-2-mlp-out" + | "layer-3-attn-mixed" + | "layer-3-attn-inject" + | "layer-3-attn-out" + | "layer-3-after-attn" + | "layer-3-mlp-mixed" + | "layer-3-mlp-out" + | "layer-3-q" + | "layer-3-gate" + | "layer-3-v" + | "layer-3-attention" + | "layer-3-hidden" + | "layer-2" + | "layer-3" + | "layer-47" + | "final-widened" + | "final-mixed" + | "logits" + ) { + let bytes = values + .iter() + .flat_map(|value| value.to_le_bytes()) + .collect::>(); + let path = match label { + "ple-embedding" => "/private/tmp/ds4_ple_embedding.f32", + "ple-key" => "/private/tmp/ds4_ple_key.f32", + "ple-value" => "/private/tmp/ds4_ple_value.f32", + "ple-query" => "/private/tmp/ds4_ple_query.f32", + "ple-gated" => "/private/tmp/ds4_ple_gated.f32", + "ple-norm" => "/private/tmp/ds4_ple_norm.f32", + "ple-output" => "/private/tmp/ds4_ple_output.f32", + "ple" => "/private/tmp/ds4_after_ple.f32", + "layer-0-attn-hc-mix" => "/private/tmp/ds4_hc_mix.f32", + "layer-0-attn-input" => "/private/tmp/ds4_attn_mixed.f32", + "layer-0-gdn-qkv" => "/private/tmp/ds4_gdn_qkv.f32", + "layer-0-gdn-conv" => "/private/tmp/ds4_gdn_conv.f32", + "layer-0-gdn-controls" => "/private/tmp/ds4_gdn_controls.f32", + "layer-0-gdn-gated" => "/private/tmp/ds4_gdn_gated.f32", + "layer-0-after-attn" => "/private/tmp/ds4_after_attn.f32", + "layer-0-mlp-hc-norm" => "/private/tmp/ds4_mlp_hc_norm.f32", + "layer-0-mlp-hc-down" => "/private/tmp/ds4_mlp_hc_down.f32", + "layer-0-mlp-hc-rank" => "/private/tmp/ds4_mlp_hc_rank.f32", + "layer-0-mlp-hc-up" => "/private/tmp/ds4_mlp_hc_up.f32", + "layer-0-mlp-hc-mix" => "/private/tmp/ds4_mlp_hc_mix.f32", + "layer-0-mlp-hc-inject" => "/private/tmp/ds4_mlp_hc_inject.f32", + "layer-0-mlp-mixed" => "/private/tmp/ds4_mlp_mixed.f32", + "layer-0-router-logits" => "/private/tmp/ds4_router_logits.f32", + "layer-0-router-probabilities" => "/private/tmp/ds4_router_probabilities.f32", + "layer-0-expert-gate" => "/private/tmp/ds4_moe_expert_gate.f32", + "layer-0-expert-up" => "/private/tmp/ds4_moe_expert_up.f32", + "layer-0-expert-mid" => "/private/tmp/ds4_moe_expert_mid.f32", + "layer-0-expert-down" => "/private/tmp/ds4_moe_expert_down.f32", + "layer-0-routed" => "/private/tmp/ds4_moe_routed.f32", + "layer-0-shared-gate-raw" => "/private/tmp/ds4_moe_shared_gate_raw.f32", + "layer-0-shared-gate-proj" => "/private/tmp/ds4_moe_shared_gate_proj.f32", + "layer-0-shared-up" => "/private/tmp/ds4_moe_shared_up.f32", + "layer-0-shared-mid" => "/private/tmp/ds4_moe_shared_mid.f32", + "layer-0-shared-down" => "/private/tmp/ds4_moe_shared_down.f32", + "layer-0-total" => "/private/tmp/ds4_moe_total.f32", + "layer-0-output" => "/private/tmp/ds4_layer0_output.f32", + "layer-1" => "/private/tmp/ds4_layer1_output.f32", + "layer-1-attn-mixed" => "/private/tmp/ds4_layer1_attn_mixed.f32", + "layer-1-attn-inject" => "/private/tmp/ds4_layer1_attn_inject.f32", + "layer-1-gdn-qkv" => "/private/tmp/ds4_layer1_gdn_qkv.f32", + "layer-1-gdn-z" => "/private/tmp/ds4_layer1_gdn_z.f32", + "layer-1-gdn-b" => "/private/tmp/ds4_layer1_gdn_b.f32", + "layer-1-gdn-a" => "/private/tmp/ds4_layer1_gdn_a.f32", + "layer-1-attn-out" => "/private/tmp/ds4_layer1_attn_out.f32", + "layer-1-after-attn" => "/private/tmp/ds4_layer1_after_attn.f32", + "layer-1-mlp-hc-norm" => "/private/tmp/ds4_layer1_mlp_hc_norm.f32", + "layer-1-mlp-hc-rank" => "/private/tmp/ds4_layer1_mlp_hc_rank.f32", + "layer-1-mlp-hc-up" => "/private/tmp/ds4_layer1_mlp_hc_up.f32", + "layer-1-mlp-hc-mix" => "/private/tmp/ds4_layer1_mlp_hc_mix.f32", + "layer-1-mlp-inject" => "/private/tmp/ds4_layer1_mlp_hc_inject.f32", + "layer-1-mlp-mixed" => "/private/tmp/ds4_layer1_mlp_mixed.f32", + "layer-1-mlp-out" => "/private/tmp/ds4_layer1_mlp_out.f32", + "layer-2-attn-mixed" => "/private/tmp/ds4_layer2_attn_mixed.f32", + "layer-2-attn-inject" => "/private/tmp/ds4_layer2_attn_inject.f32", + "layer-2-attn-out" => "/private/tmp/ds4_layer2_attn_out.f32", + "layer-2-after-attn" => "/private/tmp/ds4_layer2_after_attn.f32", + "layer-2-mlp-mixed" => "/private/tmp/ds4_layer2_mlp_mixed.f32", + "layer-2-router-probabilities" => { + "/private/tmp/ds4_layer2_router_probabilities.f32" + } + "layer-2-route-weights" => "/private/tmp/ds4_layer2_route_weights.f32", + "layer-2-expert-gate" => "/private/tmp/ds4_layer2_moe_expert_gate.f32", + "layer-2-expert-up" => "/private/tmp/ds4_layer2_moe_expert_up.f32", + "layer-2-expert-mid" => "/private/tmp/ds4_layer2_moe_expert_mid.f32", + "layer-2-expert-down" => "/private/tmp/ds4_layer2_moe_expert_down.f32", + "layer-2-routed" => "/private/tmp/ds4_layer2_moe_routed.f32", + "layer-2-shared-gate-raw" => "/private/tmp/ds4_layer2_moe_shared_gate_raw.f32", + "layer-2-shared-up" => "/private/tmp/ds4_layer2_moe_shared_up.f32", + "layer-2-shared-mid" => "/private/tmp/ds4_layer2_moe_shared_mid.f32", + "layer-2-shared-down" => "/private/tmp/ds4_layer2_moe_shared_down.f32", + "layer-2-total" => "/private/tmp/ds4_layer2_moe_total.f32", + "layer-2-mlp-out" => "/private/tmp/ds4_layer2_mlp_out.f32", + "layer-2" => "/private/tmp/ds4_layer2_output.f32", + "layer-3-attn-mixed" => "/private/tmp/ds4_layer3_attn_mixed.f32", + "layer-3-attn-inject" => "/private/tmp/ds4_layer3_attn_inject.f32", + "layer-3-attn-out" => "/private/tmp/ds4_layer3_attn_out.f32", + "layer-3-after-attn" => "/private/tmp/ds4_layer3_after_attn.f32", + "layer-3-mlp-mixed" => "/private/tmp/ds4_layer3_mlp_mixed.f32", + "layer-3-mlp-out" => "/private/tmp/ds4_layer3_mlp_out.f32", + "layer-3-q" => "/private/tmp/ds4_layer3_raw_attention.f32", + "layer-3-gate" => "/private/tmp/ds4_layer3_gate.f32", + "layer-3-v" => "/private/tmp/ds4_layer3_value.f32", + "layer-3-attention" => "/private/tmp/ds4_layer3_gated_attention.f32", + "layer-3-hidden" => "/private/tmp/ds4_layer3_manual_out.f32", + "layer-3" => "/private/tmp/ds4_layer3_output.f32", + "layer-47" => "/private/tmp/ds4_layer47_output.f32", + "final-widened" => "/private/tmp/ds4_final_widened.f32", + "final-mixed" => "/private/tmp/ds4_final_mixed.f32", + "logits" => "/private/tmp/ds4_logits.f32", + _ => "/private/tmp/ds4_attn_out.f32", + }; + std::fs::write(path, bytes).unwrap(); + } + eprintln!( + "{label}: [{:.9}, {:.9}, {:.9}, {:.9}] sha256={}", + at(0), + at(1.min(width - 1)), + at(1000.min(width - 1)), + at((HIDDEN as usize - 1).min(width - 1)), + digest, + ); + }; + let sigmoid_inputs = (0..=u16::MAX) + .filter_map(|bits| { + let value = f32::from_bits(u32::from(bits) << 16); + (value.is_finite() && value.abs() <= 32.0).then_some(value) + }) + .collect::>(); + let mut sigmoid_outputs = Vec::with_capacity(sigmoid_inputs.len()); + for input in sigmoid_inputs.chunks(HC_WIDTH as usize) { + executor.scratch.hc_mix.write_f32(input).unwrap(); + let commands = Commands::begin().unwrap(); + let mut unary = args(); + unary.u[0] = input.len() as u32; + executor + .dispatch( + c"kernel_qwen_sigmoid", + &executor.scratch.hc_norm, + Some(&executor.scratch.hc_mix), + None, + None, + &[], + &unary, + input.len() as u32, + 1, + ) + .unwrap(); + commands.finish().unwrap(); + let mut output = vec![0.0; input.len()]; + executor.scratch.hc_norm.read_f32(&mut output).unwrap(); + sigmoid_outputs.extend(output); + } + for (path, values) in [ + ("/private/tmp/ds4_sigmoid_inputs.f32", &sigmoid_inputs), + ("/private/tmp/ds4_sigmoid_outputs.f32", &sigmoid_outputs), + ] { + let bytes = values + .iter() + .flat_map(|value| value.to_bits().to_le_bytes()) + .collect::>(); + std::fs::write(path, bytes).unwrap(); + } + executor.begin_token(1).unwrap(); + eprintln!( + "ple-rows: {:?}", + ple_rows(&executor.ple_contract, executor.ple_state.history, 1).unwrap() + ); + show("embed", &executor.scratch.hc, HC_WIDTH as usize); + for layer in 0..LAYERS { + let layer_started = Instant::now(); + if layer == 1 { + executor.ple(1).unwrap(); + show( + "ple-embedding", + &executor.scratch.ple_embedding, + HIDDEN as usize, + ); + show("ple-key", &executor.scratch.ple_key, HC_WIDTH as usize); + show("ple-value", &executor.scratch.ple_value, HIDDEN as usize); + show("ple-query", &executor.scratch.hc_norm, HC_WIDTH as usize); + show("ple-gated", &executor.scratch.ple_gated, HC_WIDTH as usize); + show("ple-norm", &executor.scratch.ple_norm, HC_WIDTH as usize); + show( + "ple-output", + &executor.scratch.ple_output, + HC_WIDTH as usize, + ); + show("ple", &executor.scratch.hc, HC_WIDTH as usize); + } + if layer == 0 { + let prefix = "language_model.model.layers.0"; + let commands = Commands::begin().unwrap(); + let norm = executor + .weight(&format!("{prefix}.attn_hyper_connection.hc_norm.weight")) + .unwrap(); + let down = executor + .weight(&format!( + "{prefix}.attn_hyper_connection.input_mix_weight_down.weight" + )) + .unwrap(); + executor + .dispatch( + c"kernel_qwen_zero_rms", + &executor.scratch.hc_norm, + Some(&executor.scratch.hc), + None, + None, + &[executor.view(norm)], + &rms_args(HC_WIDTH, HIDDEN), + HC, + 1, + ) + .unwrap(); + executor + .bf16_mv( + down, + &executor.scratch.hc_norm, + &executor.scratch.rank, + HC_WIDTH, + HC_RANK, + ) + .unwrap(); + commands.finish().unwrap(); + show( + "layer-0-attn-hc-down-raw", + &executor.scratch.rank, + HC_RANK as usize, + ); + let commands = Commands::begin().unwrap(); + executor + .hyper_read(&format!("{prefix}.attn_hyper_connection")) + .unwrap(); + commands.finish().unwrap(); + show( + "layer-0-attn-hc-norm", + &executor.scratch.hc_norm, + HC_WIDTH as usize, + ); + show( + "layer-0-attn-hc-rank", + &executor.scratch.rank, + HC_RANK as usize, + ); + let commands = Commands::begin().unwrap(); + let up = executor + .weight(&format!( + "{prefix}.attn_hyper_connection.input_mix_weight_up.weight" + )) + .unwrap(); + executor + .bf16_mv( + up, + &executor.scratch.rank, + &executor.scratch.hc_mix, + HC_RANK, + HC_WIDTH, + ) + .unwrap(); + commands.finish().unwrap(); + show( + "layer-0-attn-hc-up-raw", + &executor.scratch.hc_mix, + HC_WIDTH as usize, + ); + let commands = Commands::begin().unwrap(); + let mut sigmoid = args(); + sigmoid.u[0] = HC_WIDTH; + executor + .dispatch( + c"kernel_qwen_sigmoid", + &executor.scratch.hc_mix, + Some(&executor.scratch.hc_mix), + None, + None, + &[], + &sigmoid, + HC_WIDTH, + 1, + ) + .unwrap(); + commands.finish().unwrap(); + show( + "layer-0-attn-hc-mix", + &executor.scratch.hc_mix, + HC_WIDTH as usize, + ); + show( + "layer-0-attn-hc-inject", + &executor.scratch.injection, + HC as usize, + ); + show( + "layer-0-attn-input", + &executor.scratch.block, + HIDDEN as usize, + ); + let commands = Commands::begin().unwrap(); + executor + .affine_mv_into( + &executor + .affine( + &format!("{prefix}.linear_attn.in_proj_qkv"), + HIDDEN, + GDN_QKV, + None, + ) + .unwrap(), + &executor.scratch.block, + &executor.scratch.qkv, + HIDDEN, + GDN_QKV, + ) + .unwrap(); + commands.finish().unwrap(); + show("layer-0-gdn-qkv", &executor.scratch.qkv, GDN_QKV as usize); + let commands = Commands::begin().unwrap(); + executor.gdn(prefix, layer).unwrap(); + commands.finish().unwrap(); + show( + "layer-0-gdn-gated", + &executor.scratch.gdn_out, + GDN_VALUE as usize, + ); + show("layer-0-gdn-conv", &executor.scratch.qkv, GDN_QKV as usize); + show( + "layer-0-gdn-controls", + &executor.scratch.controls, + GDN_CONTROLS as usize, + ); + show( + "layer-0-attn-output", + &executor.scratch.hidden, + HIDDEN as usize, + ); + let commands = Commands::begin().unwrap(); + executor.hyper_write().unwrap(); + commands.finish().unwrap(); + show( + "layer-0-after-attn", + &executor.scratch.hc, + HC_WIDTH as usize, + ); + let commands = Commands::begin().unwrap(); + executor + .hyper_read(&format!("{prefix}.mlp_hyper_connection")) + .unwrap(); + executor + .affine_mv_into( + &executor + .affine(&format!("{prefix}.mlp.gate"), HIDDEN, EXPERTS, None) + .unwrap(), + &executor.scratch.block, + &executor.scratch.router, + HIDDEN, + EXPERTS, + ) + .unwrap(); + commands.finish().unwrap(); + show( + "layer-0-mlp-hc-norm", + &executor.scratch.hc_norm, + HC_WIDTH as usize, + ); + show( + "layer-0-mlp-hc-rank", + &executor.scratch.rank, + HC_RANK as usize, + ); + show( + "layer-0-mlp-hc-mix", + &executor.scratch.hc_mix, + HC_WIDTH as usize, + ); + show( + "layer-0-mlp-hc-inject", + &executor.scratch.injection, + HC as usize, + ); + let down = executor + .weight(&format!( + "{prefix}.mlp_hyper_connection.input_mix_weight_down.weight" + )) + .unwrap(); + let up = executor + .weight(&format!( + "{prefix}.mlp_hyper_connection.input_mix_weight_up.weight" + )) + .unwrap(); + let commands = Commands::begin().unwrap(); + executor + .bf16_mv( + down, + &executor.scratch.hc_norm, + &executor.scratch.rank, + HC_WIDTH, + HC_RANK, + ) + .unwrap(); + commands.finish().unwrap(); + show( + "layer-0-mlp-hc-down", + &executor.scratch.rank, + HC_RANK as usize, + ); + let commands = Commands::begin().unwrap(); + let mut unary = args(); + unary.u[0] = HC_RANK; + executor + .dispatch( + c"kernel_qwen_silu_div4", + &executor.scratch.rank, + Some(&executor.scratch.rank), + None, + None, + &[], + &unary, + HC_RANK, + 1, + ) + .unwrap(); + commands.finish().unwrap(); + show( + "layer-0-mlp-hc-rank", + &executor.scratch.rank, + HC_RANK as usize, + ); + let commands = Commands::begin().unwrap(); + executor + .bf16_mv( + up, + &executor.scratch.rank, + &executor.scratch.hc_mix, + HC_RANK, + HC_WIDTH, + ) + .unwrap(); + commands.finish().unwrap(); + show( + "layer-0-mlp-hc-up", + &executor.scratch.hc_mix, + HC_WIDTH as usize, + ); + let commands = Commands::begin().unwrap(); + unary.u[0] = HC_WIDTH; + executor + .dispatch( + c"kernel_qwen_sigmoid", + &executor.scratch.hc_mix, + Some(&executor.scratch.hc_mix), + None, + None, + &[], + &unary, + HC_WIDTH, + 1, + ) + .unwrap(); + let mut mix = args(); + mix.u[0] = HIDDEN; + executor + .dispatch( + c"kernel_qwen_hyper_mix", + &executor.scratch.block, + Some(&executor.scratch.hc_norm), + Some(&executor.scratch.hc_mix), + None, + &[], + &mix, + HIDDEN, + 1, + ) + .unwrap(); + commands.finish().unwrap(); + show( + "layer-0-mlp-mixed", + &executor.scratch.block, + HIDDEN as usize, + ); + show( + "layer-0-router-logits", + &executor.scratch.router, + EXPERTS as usize, + ); + executor.encode_moe(prefix, "0").unwrap(); + show( + "layer-0-router-probabilities", + &executor.scratch.router, + EXPERTS as usize, + ); + let mut route_ids = [0_i32; EXPERTS_USED]; + let mut route_weights = [0.0_f32; EXPERTS_USED]; + executor.scratch.route_ids.read_i32(&mut route_ids).unwrap(); + executor + .scratch + .route_weights + .read_f32(&mut route_weights) + .unwrap(); + eprintln!("layer-0-routes: {route_ids:?} {route_weights:?}"); + let commands = Commands::begin().unwrap(); + executor.routed_experts(prefix).unwrap(); + commands.finish().unwrap(); + show( + "layer-0-expert-gate", + &executor.scratch.gate, + (EXPERT_WIDTH as usize) * EXPERTS_USED, + ); + show( + "layer-0-expert-up", + &executor.scratch.up, + (EXPERT_WIDTH as usize) * EXPERTS_USED, + ); + show( + "layer-0-expert-mid", + &executor.scratch.mid, + (EXPERT_WIDTH as usize) * EXPERTS_USED, + ); + show( + "layer-0-expert-down", + &executor.scratch.expert, + (HIDDEN as usize) * EXPERTS_USED, + ); + show("layer-0-routed", &executor.scratch.moe, HIDDEN as usize); + let commands = Commands::begin().unwrap(); + executor.shared_expert(prefix).unwrap(); + commands.finish().unwrap(); + show("layer-0-shared-gate-raw", &executor.scratch.gate, 1); + show( + "layer-0-shared-up", + &executor.scratch.up, + EXPERT_WIDTH as usize, + ); + show( + "layer-0-shared-mid", + &executor.scratch.mid, + EXPERT_WIDTH as usize, + ); + show( + "layer-0-shared-down", + &executor.scratch.shared, + HIDDEN as usize, + ); + show("layer-0-total", &executor.scratch.moe, HIDDEN as usize); + let commands = Commands::begin().unwrap(); + executor + .affine_mv_into( + &executor + .affine( + &format!("{prefix}.mlp.shared_expert.gate_proj"), + HIDDEN, + EXPERT_WIDTH, + None, + ) + .unwrap(), + &executor.scratch.block, + &executor.scratch.gate, + HIDDEN, + EXPERT_WIDTH, + ) + .unwrap(); + commands.finish().unwrap(); + show( + "layer-0-shared-gate-proj", + &executor.scratch.gate, + EXPERT_WIDTH as usize, + ); + show("layer-0-output", &executor.scratch.hc, HC_WIDTH as usize); + } else if layer == 1 { + let prefix = "language_model.model.layers.1"; + let commands = Commands::begin().unwrap(); + executor + .hyper_read(&format!("{prefix}.attn_hyper_connection")) + .unwrap(); + commands.finish().unwrap(); + show( + "layer-1-attn-mixed", + &executor.scratch.block, + HIDDEN as usize, + ); + show( + "layer-1-attn-inject", + &executor.scratch.injection, + HC as usize, + ); + let commands = Commands::begin().unwrap(); + executor.gdn(prefix, layer).unwrap(); + commands.finish().unwrap(); + show("layer-1-gdn-qkv", &executor.scratch.qkv, GDN_QKV as usize); + show( + "layer-1-gdn-z", + &executor.scratch.controls, + GDN_VALUE as usize, + ); + let b = executor + .scratch + .controls + .view(u64::from(GDN_VALUE) * 4, u64::from(GDN_HEADS_V) * 4) + .unwrap(); + let a = executor + .scratch + .controls + .view( + u64::from(GDN_VALUE + GDN_HEADS_V) * 4, + u64::from(GDN_HEADS_V) * 4, + ) + .unwrap(); + show("layer-1-gdn-b", &b, GDN_HEADS_V as usize); + show("layer-1-gdn-a", &a, GDN_HEADS_V as usize); + show( + "layer-1-attn-out", + &executor.scratch.hidden, + HIDDEN as usize, + ); + let commands = Commands::begin().unwrap(); + executor.hyper_write().unwrap(); + commands.finish().unwrap(); + show( + "layer-1-after-attn", + &executor.scratch.hc, + HC_WIDTH as usize, + ); + let commands = Commands::begin().unwrap(); + executor + .hyper_read(&format!("{prefix}.mlp_hyper_connection")) + .unwrap(); + commands.finish().unwrap(); + show( + "layer-1-mlp-hc-norm", + &executor.scratch.hc_norm, + HC_WIDTH as usize, + ); + show( + "layer-1-mlp-hc-rank", + &executor.scratch.rank, + HC_RANK as usize, + ); + let commands = Commands::begin().unwrap(); + let up = executor + .weight(&format!( + "{prefix}.mlp_hyper_connection.input_mix_weight_up.weight" + )) + .unwrap(); + executor + .bf16_mv( + up, + &executor.scratch.rank, + &executor.scratch.qkv, + HC_RANK, + HC_WIDTH, + ) + .unwrap(); + commands.finish().unwrap(); + show( + "layer-1-mlp-hc-up", + &executor.scratch.qkv, + HC_WIDTH as usize, + ); + show( + "layer-1-mlp-hc-mix", + &executor.scratch.hc_mix, + HC_WIDTH as usize, + ); + show( + "layer-1-mlp-inject", + &executor.scratch.injection, + HC as usize, + ); + executor.encode_moe(prefix, "1").unwrap(); + show( + "layer-1-mlp-mixed", + &executor.scratch.block, + HIDDEN as usize, + ); + show("layer-1-mlp-out", &executor.scratch.moe, HIDDEN as usize); + } else if layer == 2 { + let prefix = "language_model.model.layers.2"; + let commands = Commands::begin().unwrap(); + executor + .hyper_read(&format!("{prefix}.attn_hyper_connection")) + .unwrap(); + commands.finish().unwrap(); + show( + "layer-2-attn-mixed", + &executor.scratch.block, + HIDDEN as usize, + ); + show( + "layer-2-attn-inject", + &executor.scratch.injection, + HC as usize, + ); + let commands = Commands::begin().unwrap(); + executor.gdn(prefix, layer).unwrap(); + commands.finish().unwrap(); + show( + "layer-2-attn-out", + &executor.scratch.hidden, + HIDDEN as usize, + ); + let commands = Commands::begin().unwrap(); + executor.hyper_write().unwrap(); + commands.finish().unwrap(); + show( + "layer-2-after-attn", + &executor.scratch.hc, + HC_WIDTH as usize, + ); + let commands = Commands::begin().unwrap(); + executor + .hyper_read(&format!("{prefix}.mlp_hyper_connection")) + .unwrap(); + commands.finish().unwrap(); + show( + "layer-2-mlp-mixed", + &executor.scratch.block, + HIDDEN as usize, + ); + executor.encode_moe(prefix, "2").unwrap(); + show( + "layer-2-router-probabilities", + &executor.scratch.router, + EXPERTS as usize, + ); + show( + "layer-2-route-weights", + &executor.scratch.route_weights, + EXPERTS_USED, + ); + let mut route_ids = [0_i32; EXPERTS_USED]; + executor.scratch.route_ids.read_i32(&mut route_ids).unwrap(); + let route_bytes = route_ids + .iter() + .flat_map(|value| value.to_le_bytes()) + .collect::>(); + std::fs::write("/private/tmp/ds4_layer2_route_ids.i32", route_bytes).unwrap(); + eprintln!("layer-2-routes: {route_ids:?}"); + let up_table = executor + .affine_table( + "language_model.model.layers.2.mlp.switch_mlp.up_proj", + HIDDEN, + EXPERT_WIDTH, + ) + .unwrap(); + let mut probe_args = args(); + probe_args.u[0] = HIDDEN; + probe_args.u[1] = EXPERT_WIDTH; + probe_args.u[2] = up_table.bits; + probe_args.u[3] = up_table.group; + probe_args.u[4] = EXPERTS_USED as u32; + let commands = Commands::begin().unwrap(); + executor + .dispatch( + c"kernel_qwen_probe_affine_table", + &executor.scratch.hc_norm, + None, + Some(&executor.scratch.route_ids), + None, + &[ + executor.view(up_table.packed), + executor.view(up_table.scales), + executor.view(up_table.biases), + ], + &probe_args, + EXPERTS_USED as u32, + 1, + ) + .unwrap(); + commands.finish().unwrap(); + let mut gpu_probe = [0_i32; EXPERTS_USED * 3]; + executor.scratch.hc_norm.read_i32(&mut gpu_probe).unwrap(); + let tensor_word = |weight: Weight<'_>, index: usize, bytes: usize| { + let data = executor.model.tensor_bytes(weight.tensor).unwrap(); + let start = index * bytes; + match bytes { + 2 => u16::from_le_bytes(data[start..start + 2].try_into().unwrap()) as u32, + 4 => u32::from_le_bytes(data[start..start + 4].try_into().unwrap()), + _ => unreachable!(), + } + }; + let packed_columns = HIDDEN as usize / (32 / up_table.bits as usize); + let groups_per_row = HIDDEN as usize / up_table.group as usize; + let mut cpu_probe = [0_u32; EXPERTS_USED * 3]; + for (slot, expert) in route_ids.iter().copied().enumerate() { + let row = expert as usize * EXPERT_WIDTH as usize; + cpu_probe[slot * 3] = tensor_word(up_table.packed, row * packed_columns, 4); + cpu_probe[slot * 3 + 1] = tensor_word(up_table.scales, row * groups_per_row, 2); + cpu_probe[slot * 3 + 2] = tensor_word(up_table.biases, row * groups_per_row, 2); + } + eprintln!("layer-2-up-probe-cpu: {cpu_probe:?}"); + eprintln!("layer-2-up-probe-gpu: {gpu_probe:?}"); + let commands = Commands::begin().unwrap(); + executor.routed_experts(prefix).unwrap(); + commands.finish().unwrap(); + show( + "layer-2-expert-gate", + &executor.scratch.gate, + EXPERT_WIDTH as usize * EXPERTS_USED, + ); + show( + "layer-2-expert-up", + &executor.scratch.up, + EXPERT_WIDTH as usize * EXPERTS_USED, + ); + show( + "layer-2-expert-mid", + &executor.scratch.mid, + EXPERT_WIDTH as usize * EXPERTS_USED, + ); + show( + "layer-2-expert-down", + &executor.scratch.expert, + HIDDEN as usize * EXPERTS_USED, + ); + show("layer-2-routed", &executor.scratch.moe, HIDDEN as usize); + let commands = Commands::begin().unwrap(); + executor.shared_expert(prefix).unwrap(); + commands.finish().unwrap(); + show("layer-2-shared-gate-raw", &executor.scratch.gate, 1); + show( + "layer-2-shared-up", + &executor.scratch.up, + EXPERT_WIDTH as usize, + ); + show( + "layer-2-shared-mid", + &executor.scratch.mid, + EXPERT_WIDTH as usize, + ); + show( + "layer-2-shared-down", + &executor.scratch.shared, + HIDDEN as usize, + ); + show("layer-2-total", &executor.scratch.moe, HIDDEN as usize); + show("layer-2-mlp-out", &executor.scratch.moe, HIDDEN as usize); + } else if layer == 3 { + let prefix = "language_model.model.layers.3"; + let commands = Commands::begin().unwrap(); + executor + .hyper_read(&format!("{prefix}.attn_hyper_connection")) + .unwrap(); + executor + .attention(prefix, &executor.states[3], executor.position) + .unwrap(); + commands.finish().unwrap(); + show( + "layer-3-attn-mixed", + &executor.scratch.block, + HIDDEN as usize, + ); + show( + "layer-3-attn-inject", + &executor.scratch.injection, + HC as usize, + ); + show("layer-3-q", &executor.scratch.q, ATTN_WIDTH as usize); + show( + "layer-3-gate", + &executor.scratch.q_gate, + ATTN_WIDTH as usize, + ); + show( + "layer-3-k", + &executor.scratch.k_rope, + ATTN_KV_WIDTH as usize, + ); + show("layer-3-v", &executor.scratch.v, ATTN_KV_WIDTH as usize); + show( + "layer-3-attention", + &executor.scratch.attention, + ATTN_WIDTH as usize, + ); + show("layer-3-hidden", &executor.scratch.hidden, HIDDEN as usize); + show( + "layer-3-attn-out", + &executor.scratch.hidden, + HIDDEN as usize, + ); + let commands = Commands::begin().unwrap(); + executor.hyper_write().unwrap(); + commands.finish().unwrap(); + show( + "layer-3-after-attn", + &executor.scratch.hc, + HC_WIDTH as usize, + ); + let commands = Commands::begin().unwrap(); + executor + .hyper_read(&format!("{prefix}.mlp_hyper_connection")) + .unwrap(); + commands.finish().unwrap(); + show( + "layer-3-mlp-mixed", + &executor.scratch.block, + HIDDEN as usize, + ); + executor.encode_moe(prefix, "3").unwrap(); + show("layer-3-mlp-out", &executor.scratch.moe, HIDDEN as usize); + } else { + executor.encode_layer(layer).unwrap(); + } + if matches!(layer, 0 | 1 | 2 | 3 | 47) { + show( + &format!("layer-{layer}"), + &executor.scratch.hc, + HC_WIDTH as usize, + ); + } + eprintln!( + "layer-{layer}-ms: {:.3}", + layer_started.elapsed().as_secs_f64() * 1_000.0 + ); + } + show("final-widened", &executor.scratch.hc, HC_WIDTH as usize); + executor.final_output().unwrap(); + show("final-mixed", &executor.scratch.block, HIDDEN as usize); + show("logits", &executor.scratch.logits, VOCAB as usize); + } + + #[test] + #[ignore = "requires the pinned 105 GB Qwen artifact set and Apple Metal"] + fn qwen_affine_qmv_profile() { + configure_sources().unwrap(); + let root = std::env::var_os("DS4SERVER_QWEN38_SOURCE") + .map(PathBuf::from) + .expect("set DS4SERVER_QWEN38_SOURCE to the pinned artifact directory"); + let model = QwenModel::open(&root, 4).unwrap(); + let mut executor = QwenExecutor::open(model, 4).unwrap(); + executor.scratch.block.fill(0.0, HIDDEN.into()).unwrap(); + let projection = "language_model.model.layers.3.mlp.shared_expert.gate_proj"; + let weight = executor + .affine(projection, HIDDEN, EXPERT_WIDTH, None) + .unwrap(); + executor + .affine_mv_into( + &weight, + &executor.scratch.block, + &executor.scratch.gate, + HIDDEN, + EXPERT_WIDTH, + ) + .unwrap(); + let started = Instant::now(); + let commands = Commands::begin().unwrap(); + for _ in 0..1_000 { + executor + .affine_mv_into( + &weight, + &executor.scratch.block, + &executor.scratch.gate, + HIDDEN, + EXPERT_WIDTH, + ) + .unwrap(); + } + let encoded = started.elapsed(); + commands.finish().unwrap(); + eprintln!( + "DS4SERVER_QWEN38_QMV encode_us={:.3} total_us={:.3}", + encoded.as_secs_f64() * 1_000.0, + started.elapsed().as_secs_f64() * 1_000.0, + ); + let prefix = "language_model.model.layers.3.attn_hyper_connection"; + executor.scratch.hc.fill(0.25, HC_WIDTH.into()).unwrap(); + executor.hyper_read(prefix).unwrap(); + let started = Instant::now(); + let commands = Commands::begin().unwrap(); + for _ in 0..100 { + executor.hyper_read(prefix).unwrap(); + } + commands.finish().unwrap(); + eprintln!( + "DS4SERVER_QWEN38_HYPER_EXACT us={:.3}", + started.elapsed().as_secs_f64() * 10_000.0 + ); + let pack = executor.prepare_hyper_pack(prefix).unwrap(); + executor.hyper_packs.insert(prefix.to_owned(), pack); + executor.hyper_read(prefix).unwrap(); + let started = Instant::now(); + let commands = Commands::begin().unwrap(); + for _ in 0..100 { + executor.hyper_read(prefix).unwrap(); + } + commands.finish().unwrap(); + eprintln!( + "DS4SERVER_QWEN38_HYPER_V3 us={:.3}", + started.elapsed().as_secs_f64() * 10_000.0 + ); + } + + #[test] + #[ignore = "requires the pinned Qwen artifact set and Apple Metal"] + fn qwen_affine_qmv_and_qmm_match_mlx_reference() { + configure_sources().unwrap(); + let root = std::env::var_os("DS4SERVER_QWEN38_SOURCE") + .map(PathBuf::from) + .expect("set DS4SERVER_QWEN38_SOURCE to the pinned artifact directory"); + let model = QwenModel::open(&root, 32).unwrap(); + let executor = QwenExecutor::open_configured_with_hyper( + model, + 32, + EngineSpeculativeSettings::default(), + false, + false, + ) + .unwrap(); + let projection = "language_model.model.layers.0.linear_attn.in_proj_qkv"; + let input = Buffer::floats(u64::from(HIDDEN)).unwrap(); + let mut state = 0x1234_5678_u32; + let values = (0..HIDDEN) + .map(|_| { + state = state.wrapping_mul(1_664_525).wrapping_add(1_013_904_223); + ((state >> 24) as i32 - 128) as f32 / 32.0 + }) + .collect::>(); + input.write_f32(&values).unwrap(); + let batch = Buffer::floats(u64::from(GDN_QKV)).unwrap(); + { + let weight = executor.affine(projection, HIDDEN, GDN_QKV, None).unwrap(); + let commands = Commands::begin().unwrap(); + executor + .affine_batch_into(&weight, &input, &batch, HIDDEN, GDN_QKV, 1, true) + .unwrap(); + commands.finish().unwrap(); + } + let mut actual = vec![0.0; GDN_QKV as usize]; + batch.read_f32(&mut actual).unwrap(); + let digest = |values: &[f32]| { + let mut hash = Sha256::new(); + for value in values { + hash.update(value.to_le_bytes()); + } + hash.finalize() + .iter() + .map(|byte| format!("{byte:02x}")) + .collect::() + }; + assert_eq!( + digest(&actual), + "8952fadfdb1ef4450fd23b0060d62ea9b118367733014985a94c2d815a2be203" + ); + + for (rows, expected) in [ + ( + 32_u32, + "1d6076d65844f0c918d40085fea3af9873e86bf818f03ea4f705ac0211134fdf", + ), + ( + 64_u32, + "b595beefdc06cc3b8a8408770e0d2b9fad609a6f0701bb94a511afebe60963a8", + ), + ] { + let qmm_values = values + .iter() + .map(|&value| f32::from_bits(u32::from(bf16(value)) << 16)) + .cycle() + .take(rows as usize * HIDDEN as usize) + .collect::>(); + let input = Buffer::floats(u64::from(rows) * u64::from(HIDDEN)).unwrap(); + input.write_f32(&qmm_values).unwrap(); + let batch = Buffer::floats(u64::from(rows) * u64::from(GDN_QKV)).unwrap(); + { + let weight = executor.affine(projection, HIDDEN, GDN_QKV, None).unwrap(); + let commands = Commands::begin().unwrap(); + executor + .affine_batch_into(&weight, &input, &batch, HIDDEN, GDN_QKV, rows, true) + .unwrap(); + commands.finish().unwrap(); + } + let mut actual = vec![0.0; rows as usize * GDN_QKV as usize]; + batch.read_f32(&mut actual).unwrap(); + assert_eq!(digest(&actual), expected); + } + + let weight = executor + .weight("language_model.model.layers.1.ple.key_proj.weight") + .unwrap(); + for (rows, expected) in [ + ( + 32_u32, + "cd327b392235feb9837d0a18e3d4524f788d9a3b3cffaff03436e9280627a29e", + ), + ( + 64_u32, + "73f56c67ffb2edc6ee82c0491346223b6d8c0cd1b9fd59fa58447ba8bdcb03aa", + ), + ] { + let qmm_values = values + .iter() + .map(|&value| f32::from_bits(u32::from(bf16(value)) << 16)) + .cycle() + .take(rows as usize * HIDDEN as usize) + .collect::>(); + let input = Buffer::floats(u64::from(rows) * u64::from(HIDDEN)).unwrap(); + input.write_f32(&qmm_values).unwrap(); + let batch = Buffer::floats(u64::from(rows) * u64::from(GDN_QKV)).unwrap(); + { + let commands = Commands::begin().unwrap(); + executor + .bf16_batch(weight, &input, &batch, HIDDEN, GDN_QKV, rows) + .unwrap(); + commands.finish().unwrap(); + } + let mut actual = vec![0.0; rows as usize * GDN_QKV as usize]; + batch.read_f32(&mut actual).unwrap(); + assert_eq!(digest(&actual), expected); + } + } + + #[test] + #[ignore = "requires the pinned 105 GB Qwen artifact set and Apple Metal"] + fn qwen_layer_profile() { + configure_sources().unwrap(); + let root = std::env::var_os("DS4SERVER_QWEN38_SOURCE") + .map(PathBuf::from) + .expect("set DS4SERVER_QWEN38_SOURCE to the pinned artifact directory"); + let model = QwenModel::open(&root, 4).unwrap(); + let mut executor = QwenExecutor::open(model, 4).unwrap(); + executor.eval(1).unwrap(); + let commands = Commands::begin().unwrap(); + executor.begin_token(2).unwrap(); + commands.finish().unwrap(); + for layer in 0..LAYERS { + if layer == 1 { + let started = Instant::now(); + let commands = Commands::begin().unwrap(); + executor.ple(2).unwrap(); + commands.finish().unwrap(); + eprintln!( + "DS4SERVER_QWEN38_PLE ms={:.3}", + started.elapsed().as_secs_f64() * 1_000.0 + ); + } + let started = Instant::now(); + if layer == 20 { + let prefix = "language_model.model.layers.20"; + let stage_started = Instant::now(); + let commands = Commands::begin().unwrap(); + executor + .hyper_read(&format!("{prefix}.attn_hyper_connection")) + .unwrap(); + commands.finish().unwrap(); + eprintln!( + "DS4SERVER_QWEN38_STAGE layer={layer} stage=attn_hyper ms={:.3}", + stage_started.elapsed().as_secs_f64() * 1_000.0 + ); + let stage_started = Instant::now(); + let commands = Commands::begin().unwrap(); + executor.gdn(prefix, layer).unwrap(); + commands.finish().unwrap(); + eprintln!( + "DS4SERVER_QWEN38_STAGE layer={layer} stage=attention ms={:.3}", + stage_started.elapsed().as_secs_f64() * 1_000.0 + ); + let stage_started = Instant::now(); + let commands = Commands::begin().unwrap(); + executor.hyper_write().unwrap(); + commands.finish().unwrap(); + eprintln!( + "DS4SERVER_QWEN38_STAGE layer={layer} stage=attn_write ms={:.3}", + stage_started.elapsed().as_secs_f64() * 1_000.0 + ); + let stage_started = Instant::now(); + let commands = Commands::begin().unwrap(); + executor + .hyper_read(&format!("{prefix}.mlp_hyper_connection")) + .unwrap(); + commands.finish().unwrap(); + eprintln!( + "DS4SERVER_QWEN38_STAGE layer={layer} stage=mlp_hyper ms={:.3}", + stage_started.elapsed().as_secs_f64() * 1_000.0 + ); + let stage_started = Instant::now(); + let commands = Commands::begin().unwrap(); + executor + .affine_mv_into( + &executor + .affine(&format!("{prefix}.mlp.gate"), HIDDEN, EXPERTS, None) + .unwrap(), + &executor.scratch.block, + &executor.scratch.router, + HIDDEN, + EXPERTS, + ) + .unwrap(); + let mut softmax_args = args(); + softmax_args.u[0] = EXPERTS; + softmax_args.u[12] = EXPERTS / 4; + executor + .dispatch( + c"kernel_qwen_softmax_precise_512", + &executor.scratch.router, + Some(&executor.scratch.router), + None, + None, + &[], + &softmax_args, + 1, + 1, + ) + .unwrap(); + let mut route_args = args(); + route_args.u[0] = EXPERTS; + route_args.u[12] = EXPERTS; + executor + .dispatch( + c"kernel_qwen_route_top10", + &executor.scratch.route_ids, + Some(&executor.scratch.route_weights), + Some(&executor.scratch.router), + None, + &[], + &route_args, + 1, + 1, + ) + .unwrap(); + commands.finish().unwrap(); + eprintln!( + "DS4SERVER_QWEN38_STAGE layer={layer} stage=router ms={:.3}", + stage_started.elapsed().as_secs_f64() * 1_000.0 + ); + let stage_started = Instant::now(); + let commands = Commands::begin().unwrap(); + executor.routed_experts(prefix).unwrap(); + commands.finish().unwrap(); + eprintln!( + "DS4SERVER_QWEN38_STAGE layer={layer} stage=routed ms={:.3}", + stage_started.elapsed().as_secs_f64() * 1_000.0 + ); + let stage_started = Instant::now(); + let commands = Commands::begin().unwrap(); + executor.shared_expert(prefix).unwrap(); + commands.finish().unwrap(); + eprintln!( + "DS4SERVER_QWEN38_STAGE layer={layer} stage=shared ms={:.3}", + stage_started.elapsed().as_secs_f64() * 1_000.0 + ); + let stage_started = Instant::now(); + let commands = Commands::begin().unwrap(); + let mut inject_args = args(); + inject_args.u[0] = HIDDEN; + executor + .dispatch( + c"kernel_qwen_hyper_inject", + &executor.scratch.hc_norm, + Some(&executor.scratch.hc), + Some(&executor.scratch.moe), + Some(&executor.scratch.injection), + &[], + &inject_args, + HC_WIDTH, + 1, + ) + .unwrap(); + executor + .scratch + .hc + .copy_from( + 0, + &executor.scratch.hc_norm, + 0, + u64::from(HC_WIDTH) * 4, + "committing profiled Qwen MoE hyper streams", + ) + .unwrap(); + commands.finish().unwrap(); + eprintln!( + "DS4SERVER_QWEN38_STAGE layer={layer} stage=mlp_write ms={:.3}", + stage_started.elapsed().as_secs_f64() * 1_000.0 + ); + } else { + let commands = Commands::begin().unwrap(); + executor.encode_layer(layer).unwrap(); + commands.finish().unwrap(); + } + eprintln!( + "DS4SERVER_QWEN38_LAYER layer={layer} ms={:.3}", + started.elapsed().as_secs_f64() * 1_000.0 + ); + } + let started = Instant::now(); + let commands = Commands::begin().unwrap(); + executor.final_output().unwrap(); + commands.finish().unwrap(); + eprintln!( + "DS4SERVER_QWEN38_HEAD ms={:.3}", + started.elapsed().as_secs_f64() * 1_000.0 + ); + } + + #[test] + #[ignore = "requires the pinned 105 GB Qwen artifact set and Apple Metal"] + fn qwen_verify_layer_profile() { + configure_sources().unwrap(); + let root = std::env::var_os("DS4SERVER_QWEN38_SOURCE") + .map(PathBuf::from) + .expect("set DS4SERVER_QWEN38_SOURCE to the pinned artifact directory"); + let model = QwenModel::open_configured(&root, 8, true).unwrap(); + let settings = EngineSpeculativeSettings { + glm_mtp: true, + ..EngineSpeculativeSettings::default() + }; + let mut executor = QwenExecutor::open_configured(model, 8, settings).unwrap(); + executor.eval_target(1).unwrap(); + executor.eval_target(2).unwrap(); + let tokens = [2, 3, 4, 5]; + let commands = Commands::begin().unwrap(); + executor.begin_tokens(&tokens).unwrap(); + commands.finish().unwrap(); + for layer in 0..LAYERS { + let prefix = format!("language_model.model.layers.{layer}"); + let started = Instant::now(); + let commands = Commands::begin().unwrap(); + if layer == 1 { + executor.ple_batch(&tokens).unwrap(); + } + executor + .hyper_read_rows(&format!("{prefix}.attn_hyper_connection"), 4) + .unwrap(); + match &executor.states[layer] { + LayerState::Gdn { .. } => executor.gdn_rows(&prefix, layer, 4, true).unwrap(), + state @ LayerState::Attention { .. } => executor + .attention_rows(&prefix, state, executor.position, 4) + .unwrap(), + } + executor.hyper_write_rows(4).unwrap(); + commands.finish().unwrap(); + let attention_elapsed = started.elapsed(); + let commands = Commands::begin().unwrap(); + executor.encode_moe_rows(&prefix, 4).unwrap(); + commands.finish().unwrap(); + if layer == 20 { + eprintln!( + "DS4SERVER_QWEN38_VERIFY_STAGE layer={layer} attention_ms={:.3} moe_ms={:.3}", + attention_elapsed.as_secs_f64() * 1_000.0, + (started.elapsed() - attention_elapsed).as_secs_f64() * 1_000.0 + ); + } + eprintln!( + "DS4SERVER_QWEN38_VERIFY_LAYER layer={layer} ms={:.3}", + started.elapsed().as_secs_f64() * 1_000.0 + ); + } + let started = Instant::now(); + let commands = Commands::begin().unwrap(); + executor.final_output_rows(4).unwrap(); + commands.finish().unwrap(); + eprintln!( + "DS4SERVER_QWEN38_VERIFY_HEAD ms={:.3}", + started.elapsed().as_secs_f64() * 1_000.0 + ); + } + + #[test] + #[ignore = "requires the pinned 105 GB Qwen artifact set and Apple Metal"] + fn qwen_target_submission_profile() { + configure_sources().unwrap(); + let root = std::env::var_os("DS4SERVER_QWEN38_SOURCE") + .map(PathBuf::from) + .expect("set DS4SERVER_QWEN38_SOURCE to the pinned artifact directory"); + let model = QwenModel::open(&root, 4).unwrap(); + let mut executor = QwenExecutor::open(model, 4).unwrap(); + for token in [1, 2] { + let total = Instant::now(); + let commands = Commands::begin().unwrap(); + let encoding = Instant::now(); + executor.begin_token(token).unwrap(); + for layer in 0..LAYERS { + if layer == 1 { + executor.ple(token).unwrap(); + } + executor.encode_layer(layer).unwrap(); + } + executor.final_output().unwrap(); + let encoding_ms = encoding.elapsed().as_secs_f64() * 1_000.0; + let gpu = Instant::now(); + commands.finish().unwrap(); + let gpu_ms = gpu.elapsed().as_secs_f64() * 1_000.0; + let read = Instant::now(); + executor + .scratch + .logits + .read_f32(&mut executor.logits) + .unwrap(); + eprintln!( + "DS4SERVER_QWEN38_SUBMISSION token={token} encode_ms={encoding_ms:.3} gpu_ms={gpu_ms:.3} read_ms={:.3} total_ms={:.3}", + read.elapsed().as_secs_f64() * 1_000.0, + total.elapsed().as_secs_f64() * 1_000.0, + ); + } + } + + #[test] + #[ignore = "requires the pinned 105 GB Qwen artifact set and Apple Metal"] + fn qwen_m5_performance_gate() { + use crate::engine::ChatTurn; + use std::sync::atomic::AtomicBool; + + configure_sources().unwrap(); + let root = std::env::var_os("DS4SERVER_QWEN38_SOURCE") + .map(PathBuf::from) + .expect("set DS4SERVER_QWEN38_SOURCE to the pinned artifact directory"); + let context = std::env::var("DS4SERVER_QWEN38_CONTEXT") + .ok() + .map_or(4_096, |value| value.parse().unwrap()); + let measured = std::env::var("DS4SERVER_QWEN38_MEASURED") + .ok() + .map_or(128, |value| value.parse::().unwrap()); + let runs = std::env::var("DS4SERVER_QWEN38_RUNS") + .ok() + .map_or(3, |value| value.parse::().unwrap()); + assert!(runs > 0); + let mtp = std::env::var_os("DS4SERVER_QWEN38_MTP").is_some(); + let model = QwenModel::open_configured(&root, context, mtp).unwrap(); + let prompt = model.render_conversation( + "", + &[ChatTurn { + user: true, + tool: false, + system: false, + skip_previous_eos: false, + reasoning: None, + reasoning_complete: true, + content: "Write a concise Rust implementation of binary search and explain its invariants." + .into(), + }], + ReasoningMode::Direct, + ); + let settings = EngineSpeculativeSettings { + glm_mtp: mtp, + glm_mtp_timing: std::env::var_os("DS4SERVER_QWEN38_MTP_PROFILE").is_some(), + ..EngineSpeculativeSettings::default() + }; + let mut executor = QwenExecutor::open_configured(model, context, settings).unwrap(); + let cancelled = AtomicBool::new(false); + let mut prefill_rates = Vec::with_capacity(runs); + let mut decode_rates = Vec::with_capacity(runs); + let mut expected_hash = None; + for run in 1..=runs { + executor.reset().unwrap(); + let stats_before = executor.execution_stats(); + let prefill = Instant::now(); + assert_eq!(executor.prefill(&prompt, |_| true).unwrap(), prompt.len()); + let prefill_seconds = prefill.elapsed().as_secs_f64(); + let mut rng = Rng::new_qwen(12_345); + let decode = Instant::now(); + let mut generated = Vec::new(); + let mut cycle_ms = Vec::new(); + let mut cycle_tokens = Vec::new(); + while generated.len() < measured as usize { + let token = executor.take_pending_sampled_token().unwrap_or_else(|| { + sample_from_logits(executor.logits(), 1.0, 0.95, 0.0, 20, &mut rng) + }); + let cycle_started = Instant::now(); + let cycle = executor + .eval_speculative_sampled( + token, + measured - generated.len() as u32, + ReasoningMode::Direct, + 1.0, + 0.95, + 0.0, + 20, + &mut rng, + &cancelled, + ) + .unwrap(); + cycle_ms.push(cycle_started.elapsed().as_secs_f64() * 1_000.0); + cycle_tokens.push(cycle.len()); + generated.extend(cycle); + } + let decode_seconds = decode.elapsed().as_secs_f64(); + let hash = Sha256::digest( + generated + .iter() + .flat_map(|token| token.to_le_bytes()) + .collect::>(), + ) + .iter() + .map(|byte| format!("{byte:02x}")) + .collect::(); + if let Some(expected) = &expected_hash { + assert_eq!(&hash, expected, "same-process token hash drift"); + } else { + expected_hash = Some(hash.clone()); + } + let prefill_tps = prompt.len() as f64 / prefill_seconds; + let decode_tps = generated.len() as f64 / decode_seconds; + prefill_rates.push(prefill_tps); + decode_rates.push(decode_tps); + let stats = executor.execution_stats(); + let first_ms = cycle_ms[0]; + let first_tokens = cycle_tokens[0]; + let steady_seconds = (decode_seconds - first_ms / 1_000.0).max(f64::EPSILON); + let steady_tps = (generated.len() - first_tokens) as f64 / steady_seconds; + cycle_ms.sort_by(f64::total_cmp); + let p50_ms = cycle_ms[cycle_ms.len() / 2]; + let p95_ms = cycle_ms[cycle_ms.len().saturating_mul(95).div_ceil(100) - 1]; + eprintln!("DS4SERVER_QWEN38_TOKENS run={run} {generated:?}"); + eprintln!( + "DS4SERVER_QWEN38_PERF run={run} mtp={mtp} context={context} prompt={} prefill_tps={prefill_tps:.6} measured={} decode_tps={decode_tps:.6} first_ms={first_ms:.6} steady_tps={steady_tps:.6} p50_ms={p50_ms:.6} p95_ms={p95_ms:.6} peak_rss_bytes={} token_sha256={hash} cycles={} drafted={} accepted={} verifier_passes={} verifier_ms={} repair_ms={} ple_entries={} ple_hits={} ple_misses={} ple_evictions={} ple_pread_bytes={} ple_pread_ms={}", + prompt.len(), + generated.len(), + peak_rss_bytes(), + stats.speculative_cycles, + stats.drafted_tokens, + stats.accepted_draft_tokens, + stats.verifier_passes, + stats.verifier_ms, + stats.repair_ms, + stats.ssd_cache_entries, + stats + .ssd_cache_hits + .saturating_sub(stats_before.ssd_cache_hits), + stats + .ssd_cache_misses + .saturating_sub(stats_before.ssd_cache_misses), + stats + .ssd_cache_evictions + .saturating_sub(stats_before.ssd_cache_evictions), + stats + .ssd_pread_bytes + .saturating_sub(stats_before.ssd_pread_bytes), + stats.ssd_pread_ms.saturating_sub(stats_before.ssd_pread_ms), + ); + } + prefill_rates.sort_by(f64::total_cmp); + decode_rates.sort_by(f64::total_cmp); + eprintln!( + "DS4SERVER_QWEN38_MEDIAN mtp={mtp} runs={runs} context={context} prefill_tps={:.6} decode_tps={:.6} token_sha256={}", + prefill_rates[runs / 2], + decode_rates[runs / 2], + expected_hash.unwrap(), + ); + } + #[test] #[ignore = "requires the pinned 105 GB Qwen artifact set and Apple Metal"] fn qwen_mtp_matches_target_and_restores_its_cache() { @@ -4103,6 +9200,13 @@ mod tests { hash.finalize().into() } + fn digest_hex(values: &[f32]) -> String { + digest(values) + .iter() + .map(|byte| format!("{byte:02x}")) + .collect() + } + fn target_state_digest(executor: &QwenExecutor) -> [u8; 32] { fn buffer(hash: &mut Sha256, buffer: &Buffer, bytes: u64) { let mut values = vec![0; usize::try_from(bytes).unwrap()]; @@ -4178,27 +9282,34 @@ mod tests { hash.finalize().into() } + fn committed_mtp_attention_digests(executor: &QwenExecutor) -> [[u8; 32]; 3] { + let LayerState::Attention { + kv, + qsa_raw, + qsa_pooled, + } = &executor.mtp.as_ref().unwrap().attention + else { + unreachable!() + }; + let mut digests = [[0; 32]; 3]; + for ((buffer, bytes), digest) in [kv, qsa_raw, qsa_pooled] + .into_iter() + .zip(attention_checkpoint_bytes( + executor.position.saturating_sub(1), + )) + .zip(&mut digests) + { + let mut values = vec![0; usize::try_from(bytes).unwrap()]; + buffer.read(0, &mut values).unwrap(); + *digest = Sha256::digest(values).into(); + } + digests + } + configure_sources().unwrap(); let root = std::env::var_os("DS4SERVER_QWEN38_SOURCE") .map(PathBuf::from) .expect("set DS4SERVER_QWEN38_SOURCE to the pinned artifact directory"); - let mut baseline = QwenExecutor::open(QwenModel::open(&root, 8).unwrap(), 8).unwrap(); - baseline.eval(1).unwrap(); - baseline.eval(2).unwrap(); - let seeded = digest(baseline.logits()); - let mut token = argmax(baseline.logits()); - let mut target_tokens = Vec::new(); - let mut target_digests = Vec::new(); - let mut target_state_digests = Vec::new(); - for _ in 0..4 { - target_tokens.push(token); - baseline.eval(token).unwrap(); - target_digests.push(digest(baseline.logits())); - target_state_digests.push(target_state_digest(&baseline)); - token = argmax(baseline.logits()); - } - drop(baseline); - let settings = EngineSpeculativeSettings { glm_mtp: true, ..EngineSpeculativeSettings::default() @@ -4207,6 +9318,28 @@ mod tests { let mut executor = QwenExecutor::open_configured(model, 8, settings).unwrap(); executor.eval(1).unwrap(); executor.eval(2).unwrap(); + let seeded = digest(executor.logits()); + assert_eq!( + seeded, + [ + 0xc2, 0x00, 0xc4, 0x87, 0x20, 0x30, 0xe5, 0xbe, 0x98, 0xbf, 0x22, 0x1d, 0xa3, 0x52, + 0xf9, 0xc1, 0x9e, 0x9b, 0x15, 0xd5, 0x73, 0xb3, 0x12, 0x24, 0x00, 0xf1, 0xfe, 0x01, + 0xea, 0xba, 0xf8, 0x8e, + ], + "two-token target logits must match the direct MTPLX Family-Serve oracle", + ); + let mut token = argmax(executor.logits()); + let mut target_tokens = Vec::new(); + let mut target_state_digests = Vec::new(); + for _ in 0..4 { + target_tokens.push(token); + executor.eval_target(token).unwrap(); + target_state_digests.push(target_state_digest(&executor)); + token = argmax(executor.logits()); + } + executor.reset().unwrap(); + executor.eval(1).unwrap(); + executor.eval(2).unwrap(); assert_eq!(digest(executor.logits()), seeded); let LayerState::Attention { kv, .. } = &executor.mtp.as_ref().unwrap().attention else { unreachable!() @@ -4225,61 +9358,98 @@ mod tests { .collect::>(), vec![ ( - 1788, + 59, [ - 101, 126, 239, 116, 151, 56, 134, 52, 71, 6, 11, 42, 11, 2, 178, 196, 80, - 234, 100, 88, 71, 26, 168, 15, 170, 92, 98, 85, 194, 250, 179, 145, + 0x97, 0xc3, 0xcd, 0x28, 0x9b, 0x4f, 0xc6, 0xe4, 0x27, 0x3b, 0x4f, 0x9e, + 0xb7, 0x0c, 0xee, 0x84, 0x08, 0xb7, 0xd9, 0x25, 0xd3, 0x4f, 0x6f, 0xa6, + 0xf0, 0xcf, 0xe1, 0xfe, 0x5c, 0x90, 0x94, 0x25, ], ), ( - 1151, + 84, [ - 102, 84, 243, 59, 200, 12, 148, 234, 248, 175, 4, 153, 221, 246, 199, 149, - 56, 191, 105, 233, 72, 18, 86, 132, 62, 215, 116, 22, 151, 136, 56, 63, + 0xe1, 0x2f, 0x01, 0xed, 0x8e, 0x93, 0x97, 0x9e, 0xe9, 0x37, 0x2c, 0x29, + 0x86, 0x4f, 0x75, 0x62, 0x45, 0xe6, 0xa6, 0xc1, 0xd3, 0xdf, 0xf3, 0x94, + 0xd3, 0x39, 0x8b, 0xe7, 0x5b, 0xf7, 0x0f, 0x6d, ], ), ( - 8598, + 19, [ - 102, 99, 10, 59, 130, 251, 42, 204, 133, 69, 46, 151, 25, 67, 85, 230, 132, - 176, 220, 200, 168, 160, 46, 136, 204, 105, 122, 0, 58, 239, 198, 72, + 0x9a, 0x9c, 0xec, 0x9e, 0xdb, 0xad, 0xc6, 0xde, 0xf7, 0xdc, 0xfe, 0x46, + 0x97, 0x65, 0x84, 0xbd, 0xaa, 0xfa, 0xf2, 0x9c, 0x06, 0x79, 0x60, 0x15, + 0x72, 0xd0, 0x3b, 0x6a, 0x87, 0xca, 0xc7, 0x8c, ], ), ] ); - let mut mtp_hidden = vec![0.0; HIDDEN as usize]; - executor.scratch.block.read_f32(&mut mtp_hidden).unwrap(); + let mut mtp_mixed = vec![0.0; HIDDEN as usize]; + executor.scratch.block.read_f32(&mut mtp_mixed).unwrap(); assert_eq!( - digest(&mtp_hidden), + digest(&mtp_mixed), [ - 37, 149, 12, 95, 102, 7, 237, 89, 2, 170, 48, 69, 151, 191, 178, 241, 145, 208, - 147, 40, 132, 7, 25, 233, 103, 20, 249, 90, 86, 102, 255, 240, + 0x71, 0x7c, 0x9a, 0x14, 0xc5, 0x14, 0xc0, 0x17, 0x13, 0xb6, 0xe0, 0x5a, 0xc8, 0x4e, + 0x9b, 0x64, 0xf0, 0x39, 0xd4, 0xc8, 0xac, 0x07, 0xfc, 0xd9, 0x40, 0xcb, 0x87, 0x3b, + 0x73, 0x0d, 0x72, 0x7c, ] ); - executor.eval_target(target_tokens[0]).unwrap(); - let target_drafts = target_tokens[1..] - .iter() - .map(|&token| QwenDraft { - token, - logits: Vec::new(), - }) - .collect::>(); - let (_, complete) = executor.evaluate_drafts(&target_drafts).unwrap(); - assert!(complete); - let verified_end = executor.position; - for slot in (0..=target_drafts.len()).rev() { - executor.restore_recurrent(slot).unwrap(); - executor.trim_rejected_attention(verified_end, 4).unwrap(); + executor.capture_recurrent().unwrap(); + let (target_logits, _) = executor.eval_target_batch(&target_tokens, None).unwrap(); + assert_eq!( + target_logits + .iter() + .map(|logits| digest_hex(logits)) + .collect::>(), + [ + "16c3853061e2e7bec6a526d1b78b55f000c1c20ac41b46923ca35067fe96370f", + "4f6790bc259a5572e53098293e7f9d7d8fa3ed02945f0e708d12384026b1175d", + "b33d082c969dc86a4b1cf4fbf59f37d595e5ad1c19e41aad1ca6dfe868ce148d", + "f2f67462a96225d305594eb4d6158a0fe122daa41cdb450c59bd67a6459b30d5", + ], + "batched target logits must match the MTPLX verifier oracle", + ); + let expected_hidden = [ + "a7028c4ff1b3f198e9ffdca11fe2cd9d7818347de9eb6023cd74f54e628da625", + "e86befa3fccd4f09da7693f299d2435a102fcedfa7248375743cab50dff14f84", + "ae7e8b65a4586b28ebc651362513823adf91df3273c88cc8398f29197e9ebd16", + "f8a0138552c7160007f0f174a0be00739f98dca509556f405ee219cd9373e14c", + ]; + let expected_next = [ + "fed33afcd846e6150285b146f0a28bacad5b2dc2e6c588151398c24a08af5bc3", + "ad69cf9e54d0510b3d578f8fbceac555c69a4d96d8558181a0e272a48962fef0", + "bc6db53c482bd54bf0f2b12dee8dbb275371478d866c7dcec0694082d1186db0", + "52179bf6753eb79aa9e272d8c196aa64c7e9baa7b0d1aa5057e7d460d7b81ed6", + ]; + for slot in (0..target_tokens.len()).rev() { + let mut committed = vec![1, 2]; + committed.extend_from_slice(&target_tokens[..=slot]); + executor + .commit_verified_window(slot + 1, &committed, &target_logits[slot]) + .unwrap(); + assert_eq!(executor.tokens, committed); + assert_eq!(executor.position, 3 + slot as u32); + assert_eq!(digest(executor.logits()), digest(&target_logits[slot])); + let mut hidden = vec![0.0; HC_WIDTH as usize]; + executor.scratch.hc.read_f32(&mut hidden).unwrap(); + assert_eq!(digest_hex(&hidden), expected_hidden[slot]); + executor.eval_target(42).unwrap(); assert_eq!( - target_state_digest(&executor), - target_state_digests[slot], - "target state after accepting {slot} draft tokens", + digest_hex(executor.logits()), + expected_next[slot], + "committed verifier slot {slot}" ); } - let fallback_tokens = vec![1, 2, target_tokens[0], target_tokens[1]]; + let fallback_tokens = [target_tokens[0], target_tokens[1]]; executor - .restore_or_replay(1, &fallback_tokens, verified_end, 4, false) + .restore_or_replay( + &fallback_tokens, + &target_logits[1], + 2, + 2, + [EOS_TOKEN; PLE_HISTORY], + false, + ) .unwrap(); assert_eq!( target_state_digest(&executor), @@ -4306,6 +9476,19 @@ mod tests { assert_eq!(executor.checkpoint_tag(), [98; 32]); let first = target_tokens[0]; + let checkpoint_mtp = committed_mtp_attention_digests(&executor); + executor.eval(first).unwrap(); + let finalized_target = target_state_digest(&executor); + let finalized_mtp = committed_mtp_attention_digests(&executor); + assert!(executor.load_checkpoint(&checkpoint, &mut |_| {}).unwrap()); + assert_eq!(committed_mtp_attention_digests(&executor), checkpoint_mtp); + executor.pending_sampled = Some(first); + executor.finalize_pending_sampled_token().unwrap(); + assert!(executor.pending_sampled.is_none()); + assert_eq!(target_state_digest(&executor), finalized_target); + assert_eq!(committed_mtp_attention_digests(&executor), finalized_mtp); + assert!(executor.load_checkpoint(&checkpoint, &mut |_| {}).unwrap()); + let emitted = executor .eval_speculative_greedy( first, @@ -4315,7 +9498,10 @@ mod tests { ) .unwrap(); assert_eq!(emitted, target_tokens[..emitted.len()]); - assert_eq!(digest(executor.logits()), target_digests[emitted.len() - 1]); + assert_eq!( + digest(executor.logits()), + digest(&target_logits[emitted.len() - 1]) + ); let stats = executor.execution_stats(); assert_eq!(stats.speculative_cycles, 1); assert_eq!(stats.drafted_tokens, 3); diff --git a/src/engine/qwen.rs b/src/engine/qwen.rs index 0a9832e..7c95262 100644 --- a/src/engine/qwen.rs +++ b/src/engine/qwen.rs @@ -3,15 +3,19 @@ use super::{ChatTurn, ModelSummary}; use crate::model::ModelChoice; use crate::settings::ReasoningMode; use memmap2::{Mmap, MmapOptions}; +use rustc_hash::FxHashMap as HashMap; use serde::de::{MapAccess, Visitor}; use serde::{Deserialize, Deserializer}; use serde_json::Value; use sha2::{Digest, Sha256}; -use std::collections::{BTreeMap, BTreeSet, HashMap}; +use std::collections::{BTreeMap, BTreeSet, HashSet, VecDeque}; use std::fmt; use std::fs::{self, File}; use std::io::Read; +use std::os::unix::fs::FileExt; use std::path::{Path, PathBuf}; +use std::sync::mpsc; +use std::thread; const MAX_SAFETENSORS_HEADER: u64 = 16 * 1024 * 1024; const MANIFEST: &[u8] = include_bytes!("../../assets/models/qwen38-flash-next-bare-speed.json"); @@ -32,6 +36,9 @@ const GDN_CONV_BYTES: u64 = 2_211_840; const PLE_CONV_BYTES: u64 = 184_320; const MTP_CAPTURE_HIDDEN_BYTES: u64 = 10_240 * 4; const MTP_CAPTURE_LOGITS_BYTES: u64 = 248_320 * 4; +const PLE_ROW_BYTES: usize = 100; +const PLE_READERS: usize = 16; +const PLE_HOT_BYTES: usize = 1024 * 1024 * 1024; #[derive(Deserialize)] struct Manifest { @@ -104,7 +111,9 @@ pub(super) struct ArtifactBindings { pub(super) struct QwenMap { path: PathBuf, - map: Mmap, + file: File, + bytes: u64, + map: Option, } #[derive(Clone)] @@ -123,9 +132,227 @@ pub(super) struct QwenModel { memory: MemoryPlan, maps: Vec, tensors: HashMap, + ple_reader: PleReader, identity: [u8; 32], } +struct PleTask { + row: u64, + response: mpsc::Sender<(u64, Result<[u8; PLE_ROW_BYTES], String>)>, +} + +struct PleReader { + workers: Vec>>, + threads: Vec>, + hot: HashMap, + order: VecDeque<(u64, u64)>, + generation: u64, + capacity: usize, + hits: u64, + misses: u64, + evictions: u64, + read_bytes: u64, + read_ns: u64, +} + +#[derive(Clone, Copy, Default)] +pub(super) struct PleStats { + pub(super) entries: u64, + pub(super) cache_bytes: u64, + pub(super) hits: u64, + pub(super) misses: u64, + pub(super) evictions: u64, + pub(super) read_bytes: u64, + pub(super) read_ms: u64, +} + +impl PleReader { + fn new(maps: &[QwenMap], tensors: &HashMap) -> Result { + let packed = tensors + .get("ngram.weight") + .ok_or_else(|| "Qwen tensor is missing: ngram.weight".to_owned())?; + let scales = tensors + .get("ngram.scales") + .ok_or_else(|| "Qwen tensor is missing: ngram.scales".to_owned())?; + let biases = tensors + .get("ngram.biases") + .ok_or_else(|| "Qwen tensor is missing: ngram.biases".to_owned())?; + if scales.map != packed.map || biases.map != packed.map { + return Err("Qwen PLE tensors do not share their sidecar file".into()); + } + let ranges = [ + (packed.range.clone(), 80_u64), + (scales.range.clone(), 10_u64), + (biases.range.clone(), 10_u64), + ]; + let mut workers = Vec::with_capacity(PLE_READERS); + let mut threads = Vec::with_capacity(PLE_READERS); + for index in 0..PLE_READERS { + let file = maps[packed.map] + .file + .try_clone() + .map_err(|error| format!("cloning the Qwen PLE sidecar: {error}"))?; + let ranges = ranges.clone(); + let (sender, receiver) = mpsc::channel::>(); + let handle = thread::Builder::new() + .name(format!("qwen-ple-{index}")) + .spawn(move || { + while let Ok(Some(task)) = receiver.recv() { + let result = read_ple_row(&file, &ranges, task.row); + let _ = task.response.send((task.row, result)); + } + }) + .map_err(|error| format!("starting a Qwen PLE reader: {error}"))?; + workers.push(sender); + threads.push(handle); + } + Ok(Self { + workers, + threads, + hot: HashMap::default(), + order: VecDeque::new(), + generation: 0, + capacity: PLE_HOT_BYTES / PLE_ROW_BYTES, + hits: 0, + misses: 0, + evictions: 0, + read_bytes: 0, + read_ns: 0, + }) + } + + fn rows(&mut self, rows: &[u64]) -> Result, String> { + let mut unique = Vec::with_capacity(rows.len()); + let mut seen = HashSet::with_capacity(rows.len()); + for &row in rows { + if seen.insert(row) { + unique.push(row); + } + } + let missing = unique + .iter() + .copied() + .filter(|row| !self.hot.contains_key(row)) + .collect::>(); + self.hits = self + .hits + .saturating_add((unique.len() - missing.len()) as u64); + self.misses = self.misses.saturating_add(missing.len() as u64); + if !missing.is_empty() { + let started = std::time::Instant::now(); + let (response, received) = mpsc::channel(); + for (index, row) in missing.iter().copied().enumerate() { + self.workers[index % self.workers.len()] + .send(Some(PleTask { + row, + response: response.clone(), + })) + .map_err(|_| "Qwen PLE reader stopped unexpectedly".to_owned())?; + } + drop(response); + for _ in 0..missing.len() { + let (row, result) = received + .recv() + .map_err(|_| "Qwen PLE reader returned too few rows".to_owned())?; + self.insert(row, result?); + } + self.read_bytes = self + .read_bytes + .saturating_add((missing.len() * PLE_ROW_BYTES) as u64); + self.read_ns = self + .read_ns + .saturating_add(u64::try_from(started.elapsed().as_nanos()).unwrap_or(u64::MAX)); + } + rows.iter() + .map(|row| { + let value = self + .hot + .get(row) + .map(|(value, _)| *value) + .ok_or_else(|| format!("Qwen PLE row {row} was not cached"))?; + self.touch(*row); + Ok(value) + }) + .collect() + } + + fn insert(&mut self, row: u64, value: [u8; PLE_ROW_BYTES]) { + self.generation = self.generation.wrapping_add(1); + self.hot.insert(row, (value, self.generation)); + self.order.push_back((row, self.generation)); + while self.hot.len() > self.capacity { + let Some((old, generation)) = self.order.pop_front() else { + break; + }; + if self + .hot + .get(&old) + .is_some_and(|entry| entry.1 == generation) + { + self.hot.remove(&old); + self.evictions = self.evictions.saturating_add(1); + } + } + } + + fn touch(&mut self, row: u64) { + self.generation = self.generation.wrapping_add(1); + if let Some(entry) = self.hot.get_mut(&row) { + entry.1 = self.generation; + self.order.push_back((row, self.generation)); + } + } + + fn stats(&self) -> PleStats { + PleStats { + entries: self.hot.len() as u64, + cache_bytes: (self.hot.len() * PLE_ROW_BYTES) as u64, + hits: self.hits, + misses: self.misses, + evictions: self.evictions, + read_bytes: self.read_bytes, + read_ms: self.read_ns / 1_000_000, + } + } +} + +impl Drop for PleReader { + fn drop(&mut self) { + for worker in &self.workers { + let _ = worker.send(None); + } + for thread in self.threads.drain(..) { + let _ = thread.join(); + } + } +} + +fn read_ple_row( + file: &File, + ranges: &[(std::ops::Range, u64); 3], + row: u64, +) -> Result<[u8; PLE_ROW_BYTES], String> { + let mut output = [0_u8; PLE_ROW_BYTES]; + let mut cursor = 0; + for (range, width) in ranges { + let offset = range + .start + .checked_add( + row.checked_mul(*width) + .ok_or_else(|| "Qwen PLE row offset overflows".to_owned())?, + ) + .ok_or_else(|| "Qwen PLE row offset overflows".to_owned())?; + if offset + width > range.end { + return Err(format!("Qwen PLE row {row} is truncated")); + } + let end = cursor + *width as usize; + file.read_exact_at(&mut output[cursor..end], offset) + .map_err(|error| format!("reading Qwen PLE row {row}: {error}"))?; + cursor = end; + } + Ok(output) +} + impl QwenModel { #[cfg(test)] pub(super) fn open(root: &Path, context: u32) -> Result { @@ -150,14 +377,28 @@ impl QwenModel { paths.sort(); paths.dedup(); let mut maps = Vec::with_capacity(paths.len()); - let mut map_indices = HashMap::with_capacity(paths.len()); + let mut map_indices = HashMap::with_capacity_and_hasher(paths.len(), Default::default()); for path in paths { let file = File::open(&path).map_err(|error| format!("{}: {error}", path.display()))?; // SAFETY: verified managed artifacts remain read-only while the model owns each mapping. let map = unsafe { MmapOptions::new().map(&file) } .map_err(|error| format!("cannot map {}: {error}", path.display()))?; + if path + .file_name() + .is_some_and(|name| name == "ngram-table.safetensors") + { + // SAFETY: the read-only mapping remains valid for this advisory call. + unsafe { + libc::madvise(map.as_ptr().cast_mut().cast(), map.len(), libc::MADV_RANDOM); + } + } map_indices.insert(path.clone(), maps.len()); - maps.push(QwenMap { path, map }); + maps.push(QwenMap { + path, + file, + bytes: map.len() as u64, + map: Some(map), + }); } let tensors = bindings .into_iter() @@ -174,6 +415,7 @@ impl QwenModel { (binding.name, tensor) }) .collect::>(); + let ple_reader = PleReader::new(&maps, &tensors)?; let mut hash = Sha256::new(); hash.update(b"DS4Server Qwen3.8 checkpoint identity v1"); hash.update(MANIFEST); @@ -183,6 +425,7 @@ impl QwenModel { memory: loaded.memory, maps, tensors, + ple_reader, identity, }) } @@ -193,12 +436,25 @@ impl QwenModel { .ok_or_else(|| format!("Qwen tensor is missing: {name}")) } + pub(super) fn tensors(&self) -> impl Iterator { + self.tensors.values() + } + pub(super) fn map(&self, index: usize) -> (&[u8], &Path) { - (&self.maps[index].map, &self.maps[index].path) + ( + self.maps[index] + .map + .as_deref() + .expect("materialized Qwen weights must not use their released source mapping"), + &self.maps[index].path, + ) } pub(super) fn tensor_bytes<'a>(&'a self, tensor: &QwenTensor) -> Result<&'a [u8], String> { - let map = &self.maps[tensor.map].map; + let map = self.maps[tensor.map] + .map + .as_deref() + .ok_or_else(|| format!("{} source mapping was already released", tensor.name))?; let start = usize::try_from(tensor.range.start) .map_err(|_| format!("{} starts beyond this platform", tensor.name))?; let end = usize::try_from(tensor.range.end) @@ -207,6 +463,48 @@ impl QwenModel { .ok_or_else(|| format!("{} is outside its mapped artifact", tensor.name)) } + pub(super) fn read_tensor_at( + &self, + tensor: &QwenTensor, + offset: u64, + output: &mut [u8], + ) -> Result<(), String> { + let bytes = tensor.range.end - tensor.range.start; + let end = offset + .checked_add(output.len() as u64) + .ok_or_else(|| format!("{} read offset overflows", tensor.name))?; + if end > bytes { + return Err(format!("{} read exceeds its tensor range", tensor.name)); + } + self.maps[tensor.map] + .file + .read_exact_at(output, tensor.range.start + offset) + .map_err(|error| format!("reading {} at byte {offset}: {error}", tensor.name)) + } + + pub(super) fn release_materialized_mappings(&mut self) { + for (index, map) in self.maps.iter_mut().enumerate() { + if !self + .tensors + .values() + .any(|tensor| tensor.map == index && tensor.name.starts_with("ngram.")) + { + map.map = None; + } + } + } + + pub(super) fn read_ple_rows( + &mut self, + rows: &[u64], + ) -> Result, String> { + self.ple_reader.rows(rows) + } + + pub(super) fn ple_stats(&self) -> PleStats { + self.ple_reader.stats() + } + pub(super) fn checkpoint_identity(&self) -> [u8; 32] { self.identity } @@ -214,7 +512,7 @@ impl QwenModel { pub(super) fn summary(&self) -> ModelSummary { ModelSummary { model: ModelChoice::Qwen38FlashNext, - mapped_bytes: self.maps.iter().map(|item| item.map.len() as u64).sum(), + mapped_bytes: self.maps.iter().map(|item| item.bytes).sum(), tensor_count: self.tensors.len(), vocabulary_size: self.tokenizer.vocab_size(), support_loaded: false, @@ -284,12 +582,15 @@ impl QwenModel { let mut core = 0_u64; let mut ple = 0_u64; for item in &self.maps { - let mut pages = vec![0_i8; item.map.len().div_ceil(page)]; + let Some(map) = &item.map else { + continue; + }; + let mut pages = vec![0_i8; map.len().div_ceil(page)]; // SAFETY: each read-only mmap and residency vector remain valid for this call. if unsafe { libc::mincore( - item.map.as_ptr().cast_mut().cast(), - item.map.len(), + map.as_ptr().cast_mut().cast(), + map.len(), pages.as_mut_ptr(), ) } != 0 @@ -301,7 +602,7 @@ impl QwenModel { )); } let bytes = (pages.iter().filter(|value| **value & 1 != 0).count() * page) - .min(item.map.len()) as u64; + .min(map.len()) as u64; if item .path .file_name() @@ -846,8 +1147,8 @@ mod tests { assert_eq!( direct, [ - 248_045, 846, 198, 12_675, 248_046, 198, 248_045, 74_455, 198, 13_314, 741, 29, - 271, 510, 26_003, 29, 271, + 248_045, 846, 198, 12_675, 248_046, 198, 248_045, 74_455, 198, 248_068, 271, + 248_069, 271, ] ); assert_eq!( diff --git a/src/engine/tokenizer.rs b/src/engine/tokenizer.rs index f7f6bc6..aa78135 100644 --- a/src/engine/tokenizer.rs +++ b/src/engine/tokenizer.rs @@ -62,7 +62,6 @@ struct QwenTokenizerFile { struct QwenAddedToken { id: usize, content: String, - special: bool, } #[derive(Deserialize)] @@ -225,9 +224,7 @@ impl Tokenizer { )); } tokens[token.id] = token.content.as_bytes().to_vec(); - if token.special { - rendered_specials.push((token.content.into_bytes(), token.id as i32)); - } + rendered_specials.push((token.content.into_bytes(), token.id as i32)); } if tokens[..=maximum].iter().any(Vec::is_empty) { return Err("Qwen tokenizer token ids are not contiguous".into());