Optimize Qwen3.8 inference on Apple silicon
This commit is contained in:
1
Cargo.lock
generated
1
Cargo.lock
generated
@@ -825,6 +825,7 @@ dependencies = [
|
|||||||
"pulldown-cmark 0.12.2",
|
"pulldown-cmark 0.12.2",
|
||||||
"regex",
|
"regex",
|
||||||
"rfd",
|
"rfd",
|
||||||
|
"rustc-hash 2.1.3",
|
||||||
"serde",
|
"serde",
|
||||||
"serde_json",
|
"serde_json",
|
||||||
"serde_norway",
|
"serde_norway",
|
||||||
|
|||||||
@@ -26,6 +26,7 @@ png = "0.17.16"
|
|||||||
pulldown-cmark = "0.12.2"
|
pulldown-cmark = "0.12.2"
|
||||||
regex = "1.13.1"
|
regex = "1.13.1"
|
||||||
rfd = { version = "0.15.4", default-features = false }
|
rfd = { version = "0.15.4", default-features = false }
|
||||||
|
rustc-hash = "2.1.3"
|
||||||
serde = { version = "1.0.228", features = ["derive"] }
|
serde = { version = "1.0.228", features = ["derive"] }
|
||||||
serde_json = { version = "1.0.149", features = ["preserve_order", "raw_value"] }
|
serde_json = { version = "1.0.149", features = ["preserve_order", "raw_value"] }
|
||||||
serde_norway = "0.9.42"
|
serde_norway = "0.9.42"
|
||||||
|
|||||||
146
docs/QWEN38_PERFORMANCE.md
Normal file
146
docs/QWEN38_PERFORMANCE.md
Normal file
@@ -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.
|
||||||
18
docs/qwen38-issue-99.csv
Normal file
18
docs/qwen38-issue-99.csv
Normal file
@@ -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
|
||||||
|
3334
metal/qwen38.metal
3334
metal/qwen38.metal
File diff suppressed because it is too large
Load Diff
@@ -43,8 +43,10 @@ typedef struct {
|
|||||||
|
|
||||||
int ds4_gpu_init(void);
|
int ds4_gpu_init(void);
|
||||||
void ds4_gpu_cleanup(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(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_alloc_managed(uint64_t bytes);
|
||||||
ds4_gpu_tensor *ds4_gpu_tensor_view(const ds4_gpu_tensor *base, uint64_t offset, 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);
|
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);
|
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 {
|
typedef struct {
|
||||||
|
const ds4_gpu_tensor *tensor;
|
||||||
const void *map;
|
const void *map;
|
||||||
uint64_t size;
|
uint64_t size;
|
||||||
uint64_t offset;
|
uint64_t offset;
|
||||||
@@ -147,6 +150,7 @@ int ds4_gpu_qwen_dispatch(
|
|||||||
const ds4_gpu_tensor *a,
|
const ds4_gpu_tensor *a,
|
||||||
const ds4_gpu_tensor *b,
|
const ds4_gpu_tensor *b,
|
||||||
const ds4_gpu_tensor *c,
|
const ds4_gpu_tensor *c,
|
||||||
|
const ds4_gpu_tensor *d,
|
||||||
const ds4_gpu_qwen_weight_view *weights,
|
const ds4_gpu_qwen_weight_view *weights,
|
||||||
uint32_t weight_count,
|
uint32_t weight_count,
|
||||||
const ds4_gpu_qwen_kernel_args *args,
|
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);
|
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_quality(bool quality);
|
||||||
void ds4_gpu_set_glm_model(bool enabled);
|
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_ssd_streaming(bool enabled);
|
||||||
void ds4_gpu_set_glm_streaming_prefill_full_layer(bool enabled);
|
void ds4_gpu_set_glm_streaming_prefill_full_layer(bool enabled);
|
||||||
#ifdef __APPLE__
|
#ifdef __APPLE__
|
||||||
|
|||||||
@@ -50,6 +50,7 @@ enum {
|
|||||||
static id<MTLDevice> g_device;
|
static id<MTLDevice> g_device;
|
||||||
static id<MTLCommandQueue> g_queue;
|
static id<MTLCommandQueue> g_queue;
|
||||||
static id<MTLLibrary> g_library;
|
static id<MTLLibrary> g_library;
|
||||||
|
static id<MTLLibrary> g_qwen_library;
|
||||||
static id<MTLCommandBuffer> g_batch_cb;
|
static id<MTLCommandBuffer> g_batch_cb;
|
||||||
static id<MTLComputeCommandEncoder> g_batch_enc;
|
static id<MTLComputeCommandEncoder> g_batch_enc;
|
||||||
static BOOL g_batch_encoder_concurrent;
|
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 uint64_t g_model_residency_count;
|
||||||
static int g_model_residency_added_to_queue;
|
static int g_model_residency_added_to_queue;
|
||||||
static int g_glm_model_mode;
|
static int g_glm_model_mode;
|
||||||
|
static int g_qwen_model_mode;
|
||||||
static int g_ssd_streaming_mode;
|
static int g_ssd_streaming_mode;
|
||||||
static int g_glm_streaming_prefill_full_layer_runtime;
|
static int g_glm_streaming_prefill_full_layer_runtime;
|
||||||
static int g_metal4_runtime_available;
|
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) {
|
static MTLResourceOptions ds4_gpu_model_resource_options(void) {
|
||||||
MTLResourceOptions options = MTLResourceStorageModeShared;
|
MTLResourceOptions options = MTLResourceStorageModeShared;
|
||||||
if (getenv("DS4_METAL_MODEL_UNTRACKED") != NULL) {
|
if (g_qwen_model_mode || getenv("DS4_METAL_MODEL_UNTRACKED") != NULL) {
|
||||||
options |= MTLResourceHazardTrackingModeUntracked;
|
options |= MTLResourceHazardTrackingModeUntracked;
|
||||||
}
|
}
|
||||||
return options;
|
return options;
|
||||||
@@ -2238,7 +2240,10 @@ static id<MTLComputePipelineState> ds4_gpu_get_pipeline(
|
|||||||
|
|
||||||
NSError *error = nil;
|
NSError *error = nil;
|
||||||
NSString *name = [NSString stringWithUTF8String:function_name];
|
NSString *name = [NSString stringWithUTF8String:function_name];
|
||||||
id<MTLFunction> fn = [g_library newFunctionWithName:name];
|
id<MTLLibrary> library = strncmp(function_name, "kernel_qwen_", 12) == 0
|
||||||
|
? g_qwen_library
|
||||||
|
: g_library;
|
||||||
|
id<MTLFunction> fn = [library newFunctionWithName:name];
|
||||||
if (!fn) {
|
if (!fn) {
|
||||||
fprintf(stderr, "ds4: Metal %s function not found\n", function_name);
|
fprintf(stderr, "ds4: Metal %s function not found\n", function_name);
|
||||||
return nil;
|
return nil;
|
||||||
@@ -2369,6 +2374,10 @@ static int ds4_gpu_mpp_available(void) {
|
|||||||
return g_metal4_tensor_api_enabled && !g_quality_mode;
|
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.
|
* Retained Metal4 defaults live here instead of behind user-visible options.
|
||||||
* The public runtime has one automatic accelerated path plus the global
|
* 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;
|
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) {
|
void ds4_gpu_set_ssd_streaming(bool enabled) {
|
||||||
g_ssd_streaming_mode = enabled ? 1 : 0;
|
g_ssd_streaming_mode = enabled ? 1 : 0;
|
||||||
ds4_gpu_stream_expert_cache_clear_all(1);
|
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_BF16_SOURCE", @"metal/glm53_bf16.metal"],
|
||||||
@[@"DS4_METAL_GLM53_VISION_SOURCE", @"metal/glm53_vision.metal"],
|
@[@"DS4_METAL_GLM53_VISION_SOURCE", @"metal/glm53_vision.metal"],
|
||||||
@[@"DS4_METAL_GLM53_KDA_SOURCE", @"metal/glm53_kda.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_MOE_SOURCE", @"metal/moe.metal"],
|
||||||
@[@"DS4_METAL_DSV4_HC_SOURCE", @"metal/dsv4_hc.metal"],
|
@[@"DS4_METAL_DSV4_HC_SOURCE", @"metal/dsv4_hc.metal"],
|
||||||
@[@"DS4_METAL_UNARY_SOURCE", @"metal/unary.metal"],
|
@[@"DS4_METAL_UNARY_SOURCE", @"metal/unary.metal"],
|
||||||
@@ -4418,6 +4430,31 @@ static NSString *ds4_gpu_full_source(void) {
|
|||||||
return source;
|
return source;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
static NSString *ds4_gpu_qwen_source(void) {
|
||||||
|
const char *override_path = getenv("DS4_METAL_QWEN38_SOURCE");
|
||||||
|
NSMutableArray<NSString *> *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 {
|
typedef struct {
|
||||||
int32_t ne00t;
|
int32_t ne00t;
|
||||||
int32_t ne00;
|
int32_t ne00;
|
||||||
@@ -6507,6 +6544,37 @@ int ds4_gpu_init(void) {
|
|||||||
}
|
}
|
||||||
g_library = library;
|
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<MTLLibrary> 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<MTLFunction> fn = [library newFunctionWithName:@"kernel_get_rows_f32"];
|
id<MTLFunction> fn = [library newFunctionWithName:@"kernel_get_rows_f32"];
|
||||||
if (!fn) {
|
if (!fn) {
|
||||||
fprintf(stderr, "ds4: Metal kernel_get_rows_f32 function not found\n");
|
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;
|
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 (!g_initialized && !ds4_gpu_init()) return NULL;
|
||||||
if (bytes == 0 || bytes > (uint64_t)NSUIntegerMax) return NULL;
|
if (bytes == 0 || bytes > (uint64_t)NSUIntegerMax) return NULL;
|
||||||
|
|
||||||
@autoreleasepool {
|
@autoreleasepool {
|
||||||
DS4MetalTensor *tensor = [DS4MetalTensor new];
|
DS4MetalTensor *tensor = [DS4MetalTensor new];
|
||||||
tensor.buffer = [g_device newBufferWithLength:(NSUInteger)bytes
|
tensor.buffer = [g_device newBufferWithLength:(NSUInteger)bytes
|
||||||
options:MTLResourceStorageModeShared];
|
options:options];
|
||||||
if (!tensor.buffer) {
|
if (!tensor.buffer) {
|
||||||
return NULL;
|
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) {
|
ds4_gpu_tensor *ds4_gpu_tensor_alloc_managed(uint64_t bytes) {
|
||||||
return ds4_gpu_tensor_alloc(bytes);
|
return ds4_gpu_tensor_alloc(bytes);
|
||||||
}
|
}
|
||||||
@@ -10528,6 +10608,7 @@ void ds4_gpu_cleanup(void) {
|
|||||||
g_model_buffer_cache = nil;
|
g_model_buffer_cache = nil;
|
||||||
g_transient_buffers = nil;
|
g_transient_buffers = nil;
|
||||||
g_pending_cbs = nil;
|
g_pending_cbs = nil;
|
||||||
|
g_qwen_library = nil;
|
||||||
g_library = nil;
|
g_library = nil;
|
||||||
g_queue = nil;
|
g_queue = nil;
|
||||||
g_device = nil;
|
g_device = nil;
|
||||||
@@ -11744,13 +11825,14 @@ int ds4_gpu_qwen_dispatch(
|
|||||||
const ds4_gpu_tensor *a,
|
const ds4_gpu_tensor *a,
|
||||||
const ds4_gpu_tensor *b,
|
const ds4_gpu_tensor *b,
|
||||||
const ds4_gpu_tensor *c,
|
const ds4_gpu_tensor *c,
|
||||||
|
const ds4_gpu_tensor *d,
|
||||||
const ds4_gpu_qwen_weight_view *weights,
|
const ds4_gpu_qwen_weight_view *weights,
|
||||||
uint32_t weight_count,
|
uint32_t weight_count,
|
||||||
const ds4_gpu_qwen_kernel_args *args,
|
const ds4_gpu_qwen_kernel_args *args,
|
||||||
uint32_t grid_x,
|
uint32_t grid_x,
|
||||||
uint32_t grid_y) {
|
uint32_t grid_y) {
|
||||||
if (!kernel || !out || !args || grid_x == 0 || grid_y == 0 ||
|
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;
|
return 0;
|
||||||
}
|
}
|
||||||
id<MTLComputePipelineState> pipeline = ds4_gpu_get_pipeline(kernel);
|
id<MTLComputePipelineState> 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];
|
[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++) {
|
for (uint32_t i = 0; i < weight_count; i++) {
|
||||||
uint64_t inner = 0;
|
uint64_t inner = 0;
|
||||||
id<MTLBuffer> weight = ds4_gpu_wrap_model_exact_range(
|
id<MTLBuffer> weight = nil;
|
||||||
weights[i].map,
|
if (weights[i].tensor) {
|
||||||
weights[i].size,
|
const DS4MetalTensor *tensor = ds4_gpu_tensor_const_obj(weights[i].tensor);
|
||||||
weights[i].offset,
|
weight = tensor.buffer;
|
||||||
weights[i].bytes,
|
inner = tensor.offset;
|
||||||
&inner);
|
} 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) {
|
if (!weight) {
|
||||||
ds4_gpu_end_compute_encoder(cb, enc);
|
ds4_gpu_end_compute_encoder(cb, enc);
|
||||||
if (owned) [cb commit];
|
if (owned) [cb commit];
|
||||||
@@ -11791,16 +11893,34 @@ int ds4_gpu_qwen_dispatch(
|
|||||||
/* Metal resource offsets are four-byte aligned. Preserve an arbitrary
|
/* Metal resource offsets are four-byte aligned. Preserve an arbitrary
|
||||||
* safetensors data offset for the Qwen kernels to decode explicitly. */
|
* safetensors data offset for the Qwen kernels to decode explicitly. */
|
||||||
const uint64_t aligned_inner = inner & ~3ull;
|
const uint64_t aligned_inner = inner & ~3ull;
|
||||||
bound_args.u[13 + i] = (uint32_t)(inner - aligned_inner);
|
static const uint32_t offset_slots[6] = {13, 14, 15, 9, 6, 7};
|
||||||
[enc setBuffer:weight offset:(NSUInteger)aligned_inner atIndex:5 + i];
|
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];
|
[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 width = pipeline.threadExecutionWidth;
|
||||||
const NSUInteger max_threads = pipeline.maxTotalThreadsPerThreadgroup;
|
const NSUInteger max_threads = pipeline.maxTotalThreadsPerThreadgroup;
|
||||||
const NSUInteger threads = MIN(MAX(width, 1u), max_threads);
|
const NSUInteger requested_threads = bound_args.u[12];
|
||||||
[enc dispatchThreads:MTLSizeMake(grid_x, grid_y, 1)
|
if (requested_threads != 0) {
|
||||||
threadsPerThreadgroup:MTLSizeMake(threads, 1, 1)];
|
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);
|
ds4_gpu_end_compute_encoder(cb, enc);
|
||||||
return owned ? ds4_gpu_finish_command_buffer(cb, 1, kernel) : 1;
|
return owned ? ds4_gpu_finish_command_buffer(cb, 1, kernel) : 1;
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -24,6 +24,7 @@ impl ModelPreferences {
|
|||||||
Self {
|
Self {
|
||||||
reasoning_mode: model.default_reasoning_mode(),
|
reasoning_mode: model.default_reasoning_mode(),
|
||||||
speculative: SpeculativePreferences {
|
speculative: SpeculativePreferences {
|
||||||
|
glm_mtp: model == ModelChoice::Qwen38FlashNext,
|
||||||
dspark_enabled: model == ModelChoice::DeepSeekV4Flash0731,
|
dspark_enabled: model == ModelChoice::DeepSeekV4Flash0731,
|
||||||
..SpeculativePreferences::default()
|
..SpeculativePreferences::default()
|
||||||
},
|
},
|
||||||
@@ -308,7 +309,13 @@ impl Config {
|
|||||||
.map_err(|error| format!("Could not read {}: {error}", path.display()))?;
|
.map_err(|error| format!("Could not read {}: {error}", path.display()))?;
|
||||||
drop_legacy_model_settings(&mut value);
|
drop_legacy_model_settings(&mut value);
|
||||||
let migrated = migrate_deprecated_flash(&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 =
|
let explicit_reasoning =
|
||||||
MODEL_CHOICES.map(|model| (model, model_reasoning_mode_is_explicit(&value, model)));
|
MODEL_CHOICES.map(|model| (model, model_reasoning_mode_is_explicit(&value, model)));
|
||||||
let mut config: Self = serde_norway::from_value(value)
|
let mut config: Self = serde_norway::from_value(value)
|
||||||
@@ -331,6 +338,14 @@ impl Config {
|
|||||||
.speculative
|
.speculative
|
||||||
.dspark_enabled = true;
|
.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()?;
|
config.validate()?;
|
||||||
if migrated {
|
if migrated {
|
||||||
config.save(path)?;
|
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())))
|
.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 {
|
let Value::Mapping(root) = value else {
|
||||||
return false;
|
return false;
|
||||||
};
|
};
|
||||||
root.get(Value::String("model_profiles".into()))
|
root.get(Value::String("model_profiles".into()))
|
||||||
.and_then(Value::as_mapping)
|
.and_then(Value::as_mapping)
|
||||||
.and_then(|profiles| {
|
.and_then(|profiles| profiles.get(Value::String(model.id().into())))
|
||||||
profiles.get(Value::String(ModelChoice::DeepSeekV4Flash0731.id().into()))
|
|
||||||
})
|
|
||||||
.and_then(Value::as_mapping)
|
.and_then(Value::as_mapping)
|
||||||
.and_then(|profile| profile.get(Value::String("speculative".into())))
|
.and_then(|profile| profile.get(Value::String("speculative".into())))
|
||||||
.and_then(Value::as_mapping)
|
.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) {
|
fn drop_legacy_model_settings(value: &mut Value) {
|
||||||
@@ -603,6 +616,28 @@ mod tests {
|
|||||||
ReasoningMode::Medium
|
ReasoningMode::Medium
|
||||||
);
|
);
|
||||||
assert_eq!(loaded.active_generation().context_tokens, 131_072);
|
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();
|
fs::remove_dir_all(directory).unwrap();
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
378
src/engine.rs
378
src/engine.rs
@@ -792,6 +792,9 @@ impl Generator {
|
|||||||
let result = self.generate_inner(messages, settings, cancelled, &mut emit, &mut progress);
|
let result = self.generate_inner(messages, settings, cancelled, &mut emit, &mut progress);
|
||||||
self.publish_execution_stats();
|
self.publish_execution_stats();
|
||||||
let (mut output, prompt_complete) = result?;
|
let (mut output, prompt_complete) = result?;
|
||||||
|
if prompt_complete {
|
||||||
|
self.executor.finalize_pending_sampled_token()?;
|
||||||
|
}
|
||||||
let mut completed = messages.to_vec();
|
let mut completed = messages.to_vec();
|
||||||
completed.push(output.message.clone());
|
completed.push(output.message.clone());
|
||||||
output.previous_checkpoint_bytes =
|
output.previous_checkpoint_bytes =
|
||||||
@@ -870,6 +873,9 @@ impl Generator {
|
|||||||
let result = self.generate_inner(messages, settings, cancelled, &mut emit, &mut progress);
|
let result = self.generate_inner(messages, settings, cancelled, &mut emit, &mut progress);
|
||||||
self.publish_execution_stats();
|
self.publish_execution_stats();
|
||||||
let (mut output, prompt_complete) = result?;
|
let (mut output, prompt_complete) = result?;
|
||||||
|
if prompt_complete {
|
||||||
|
self.executor.finalize_pending_sampled_token()?;
|
||||||
|
}
|
||||||
let mut completed = messages.to_vec();
|
let mut completed = messages.to_vec();
|
||||||
completed.push(output.message.clone());
|
completed.push(output.message.clone());
|
||||||
output.previous_checkpoint_bytes = previous_checkpoint
|
output.previous_checkpoint_bytes = previous_checkpoint
|
||||||
@@ -1462,7 +1468,12 @@ impl Generator {
|
|||||||
self.executor.set_vision_overlays(overlays)?;
|
self.executor.set_vision_overlays(overlays)?;
|
||||||
self.metrics.kv_prefix_reused(reused);
|
self.metrics.kv_prefix_reused(reused);
|
||||||
progress(self.executor.position(), self.executor.context(), None);
|
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 reasoning = settings.reasoning_mode != ReasoningMode::Direct;
|
||||||
let mut generated = ChatTurn {
|
let mut generated = ChatTurn {
|
||||||
user: false,
|
user: false,
|
||||||
@@ -1525,14 +1536,32 @@ impl Generator {
|
|||||||
true,
|
true,
|
||||||
));
|
));
|
||||||
}
|
}
|
||||||
let token = sample(
|
let token = self
|
||||||
self.executor.logits(),
|
.executor
|
||||||
settings.temperature,
|
.take_pending_sampled_token()
|
||||||
settings.top_p,
|
.unwrap_or_else(|| {
|
||||||
settings.min_p,
|
if matches!(self.executor.model(), ModelRef::Qwen(_)) {
|
||||||
settings.top_k,
|
sample_probabilities_f64(
|
||||||
&mut rng,
|
&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
|
if self
|
||||||
.executor
|
.executor
|
||||||
.model()
|
.model()
|
||||||
@@ -1903,6 +1932,123 @@ fn sample(
|
|||||||
sample_probabilities(&probabilities, rng, None)
|
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::<f64>();
|
||||||
|
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))]
|
#[cfg(any(target_os = "macos", test))]
|
||||||
fn sampling_probabilities(
|
fn sampling_probabilities(
|
||||||
logits: &[f32],
|
logits: &[f32],
|
||||||
@@ -1938,20 +2084,29 @@ fn sampling_probabilities(
|
|||||||
top_p
|
top_p
|
||||||
};
|
};
|
||||||
let min_p = min_p.max(0.0);
|
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()
|
.iter()
|
||||||
|
.copied()
|
||||||
.enumerate()
|
.enumerate()
|
||||||
.filter(|(_, logit)| logit.is_finite())
|
.filter(|(_, logit)| logit.is_finite())
|
||||||
.map(|(index, logit)| (index, ((*logit - maximum) / temperature).exp()))
|
.collect::<Vec<_>>();
|
||||||
.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::<Vec<_>>();
|
||||||
if probabilities.is_empty() {
|
if probabilities.is_empty() {
|
||||||
return greedy();
|
return greedy();
|
||||||
}
|
}
|
||||||
if top_p < 1.0 || top_k > 0 || min_p > 0.0 {
|
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)));
|
probabilities.sort_unstable_by(rank);
|
||||||
if top_k > 0 {
|
|
||||||
probabilities.truncate(probabilities.len().min((top_k as usize).min(1024)));
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
let total: f32 = probabilities
|
let total: f32 = probabilities
|
||||||
.iter()
|
.iter()
|
||||||
@@ -2037,7 +2192,7 @@ fn exact_delta_sample(
|
|||||||
)
|
)
|
||||||
}
|
}
|
||||||
|
|
||||||
#[cfg(any(target_os = "macos", test))]
|
#[cfg(test)]
|
||||||
fn sample_from_logits(
|
fn sample_from_logits(
|
||||||
logits: &[f32],
|
logits: &[f32],
|
||||||
temperature: f32,
|
temperature: f32,
|
||||||
@@ -2046,14 +2201,13 @@ fn sample_from_logits(
|
|||||||
top_k: i32,
|
top_k: i32,
|
||||||
rng: &mut Rng,
|
rng: &mut Rng,
|
||||||
) -> i32 {
|
) -> i32 {
|
||||||
sample_probabilities(
|
sample_probabilities_f64(
|
||||||
&sampling_probabilities(logits, temperature, top_p, min_p, top_k),
|
&qwen_sampling_probabilities(logits, temperature, top_p, min_p, top_k),
|
||||||
rng,
|
rng,
|
||||||
None,
|
|
||||||
)
|
)
|
||||||
}
|
}
|
||||||
|
|
||||||
#[cfg(any(target_os = "macos", test))]
|
#[cfg(test)]
|
||||||
#[allow(clippy::too_many_arguments)]
|
#[allow(clippy::too_many_arguments)]
|
||||||
fn exact_speculative_sample(
|
fn exact_speculative_sample(
|
||||||
target_logits: &[f32],
|
target_logits: &[f32],
|
||||||
@@ -2065,19 +2219,53 @@ fn exact_speculative_sample(
|
|||||||
top_k: i32,
|
top_k: i32,
|
||||||
rng: &mut Rng,
|
rng: &mut Rng,
|
||||||
) -> (i32, bool) {
|
) -> (i32, bool) {
|
||||||
let mut target = sampling_probabilities(target_logits, temperature, top_p, min_p, top_k);
|
let proposal = qwen_sampling_probabilities(draft_logits, temperature, top_p, min_p, top_k);
|
||||||
let mut proposal = sampling_probabilities(draft_logits, temperature, top_p, min_p, top_k);
|
exact_speculative_sample_with_probabilities(
|
||||||
target.sort_unstable_by_key(|(token, _)| *token);
|
target_logits,
|
||||||
proposal.sort_unstable_by_key(|(token, _)| *token);
|
&proposal,
|
||||||
let probability = |values: &[(usize, f32)], token: usize| {
|
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
|
values
|
||||||
.binary_search_by_key(&token, |(candidate, _)| *candidate)
|
.binary_search_by_key(&token, |(candidate, _)| *candidate)
|
||||||
.ok()
|
.ok()
|
||||||
.map_or(0.0, |index| values[index].1)
|
.map_or(0.0, |index| values[index].1)
|
||||||
};
|
};
|
||||||
let target_probability = probability(&target, draft as usize);
|
let target_probability = probability(target, draft as usize);
|
||||||
let draft_probability = probability(&proposal, 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) {
|
if draft_probability > 0.0
|
||||||
|
&& rng.unit_f64() <= (target_probability / draft_probability).min(1.0)
|
||||||
|
{
|
||||||
return (draft, true);
|
return (draft, true);
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -2086,8 +2274,8 @@ fn exact_speculative_sample(
|
|||||||
.max_by(|left, right| left.1.total_cmp(&right.1))
|
.max_by(|left, right| left.1.total_cmp(&right.1))
|
||||||
.map_or(draft, |(token, _)| *token as i32);
|
.map_or(draft, |(token, _)| *token as i32);
|
||||||
let mut residual = Vec::with_capacity(target.len());
|
let mut residual = Vec::with_capacity(target.len());
|
||||||
for (token, target_probability) in target {
|
for &(token, target_probability) in target {
|
||||||
let remaining = target_probability - probability(&proposal, token);
|
let remaining = target_probability - probability(proposal, token);
|
||||||
if remaining > 0.0 {
|
if remaining > 0.0 {
|
||||||
residual.push((token, remaining));
|
residual.push((token, remaining));
|
||||||
}
|
}
|
||||||
@@ -2095,27 +2283,113 @@ fn exact_speculative_sample(
|
|||||||
if residual.is_empty() {
|
if residual.is_empty() {
|
||||||
return (target_best, false);
|
return (target_best, false);
|
||||||
}
|
}
|
||||||
(sample_probabilities(&residual, rng, None), false)
|
(sample_probabilities_f64(&residual, rng), false)
|
||||||
}
|
}
|
||||||
|
|
||||||
#[cfg(any(target_os = "macos", test))]
|
#[cfg(any(target_os = "macos", test))]
|
||||||
struct Rng(u64);
|
enum Rng {
|
||||||
|
Ds4(u64),
|
||||||
|
Qwen { state: u128, increment: u128 },
|
||||||
|
}
|
||||||
|
|
||||||
#[cfg(any(target_os = "macos", test))]
|
#[cfg(any(target_os = "macos", test))]
|
||||||
impl Rng {
|
impl Rng {
|
||||||
fn new(seed: u64) -> Self {
|
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 {
|
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 {
|
if value == 0 {
|
||||||
value = 0x9e37_79b9_7f4a_7c15;
|
value = 0x9e37_79b9_7f4a_7c15;
|
||||||
}
|
}
|
||||||
value ^= value >> 12;
|
value ^= value >> 12;
|
||||||
value ^= value << 25;
|
value ^= value << 25;
|
||||||
value ^= value >> 27;
|
value ^= value >> 27;
|
||||||
self.0 = value;
|
*state = value;
|
||||||
let value = value.wrapping_mul(0x2545_f491_4f6c_dd1d);
|
let value = value.wrapping_mul(0x2545_f491_4f6c_dd1d);
|
||||||
((value >> 40) & 0xff_ffff) as f32 / 16_777_216.0
|
((value >> 40) & 0xff_ffff) as f32 / 16_777_216.0
|
||||||
}
|
}
|
||||||
@@ -2125,6 +2399,21 @@ impl Rng {
|
|||||||
mod sampling_tests {
|
mod sampling_tests {
|
||||||
use super::*;
|
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]
|
#[test]
|
||||||
fn checkpoint_headers_identify_the_exact_model() {
|
fn checkpoint_headers_identify_the_exact_model() {
|
||||||
let path =
|
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::<Vec<_>>(),
|
||||||
|
[0, 1]
|
||||||
|
);
|
||||||
|
assert!(
|
||||||
|
probabilities
|
||||||
|
.iter()
|
||||||
|
.all(|(_, probability)| (*probability - 0.5).abs() < f64::EPSILON)
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn split_utf8_token_bytes_are_joined_before_decoding() {
|
fn split_utf8_token_bytes_are_joined_before_decoding() {
|
||||||
let mut generated = ChatTurn {
|
let mut generated = ChatTurn {
|
||||||
|
|||||||
@@ -13,10 +13,13 @@ use profile::ExpertProfile;
|
|||||||
use qwen::QwenExecutor;
|
use qwen::QwenExecutor;
|
||||||
|
|
||||||
use super::gguf::{BF16, F16, F32, Gguf, IQ2_XXS, MXFP4, Q2_K, Q4_K, Q8_0, Tensor as GgufTensor};
|
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::validation::{DsparkConfig, SupportKind, dspark_config};
|
||||||
use super::{
|
use super::{
|
||||||
LoadedModel, Model, ModelFamily, ModelRef, Rng, exact_delta_sample, exact_speculative_sample,
|
LoadedModel, Model, ModelFamily, ModelRef, Rng, exact_delta_sample,
|
||||||
sample_from_logits,
|
exact_speculative_sample_from_probabilities, qwen_sampling_probabilities,
|
||||||
|
sample_probabilities_f64,
|
||||||
};
|
};
|
||||||
use crate::model::ModelChoice;
|
use crate::model::ModelChoice;
|
||||||
use crate::settings::{
|
use crate::settings::{
|
||||||
@@ -4517,6 +4520,20 @@ impl Executor {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
pub(super) fn take_pending_sampled_token(&mut self) -> Option<i32> {
|
||||||
|
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> {
|
pub(super) fn eval(&mut self, token: i32) -> Result<(), String> {
|
||||||
match self {
|
match self {
|
||||||
Self::DeepSeek(executor) => executor.eval(token),
|
Self::DeepSeek(executor) => executor.eval(token),
|
||||||
|
|||||||
@@ -8,6 +8,7 @@ pub(super) struct GpuTensor {
|
|||||||
#[derive(Clone, Copy)]
|
#[derive(Clone, Copy)]
|
||||||
#[repr(C)]
|
#[repr(C)]
|
||||||
pub(super) struct QwenWeightView {
|
pub(super) struct QwenWeightView {
|
||||||
|
pub(super) tensor: *const GpuTensor,
|
||||||
pub(super) map: *const c_void,
|
pub(super) map: *const c_void,
|
||||||
pub(super) size: u64,
|
pub(super) size: u64,
|
||||||
pub(super) offset: u64,
|
pub(super) offset: u64,
|
||||||
@@ -90,6 +91,7 @@ pub(super) struct StreamExpertCacheStats {
|
|||||||
unsafe extern "C" {
|
unsafe extern "C" {
|
||||||
pub(super) fn ds4_gpu_init() -> i32;
|
pub(super) fn ds4_gpu_init() -> i32;
|
||||||
pub(super) fn ds4_gpu_cleanup();
|
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(
|
pub(super) fn ds4_gpu_set_model_map_range(
|
||||||
model_map: *const c_void,
|
model_map: *const c_void,
|
||||||
model_size: u64,
|
model_size: u64,
|
||||||
@@ -103,6 +105,7 @@ unsafe extern "C" {
|
|||||||
a: *const GpuTensor,
|
a: *const GpuTensor,
|
||||||
b: *const GpuTensor,
|
b: *const GpuTensor,
|
||||||
c: *const GpuTensor,
|
c: *const GpuTensor,
|
||||||
|
d: *const GpuTensor,
|
||||||
weights: *const QwenWeightView,
|
weights: *const QwenWeightView,
|
||||||
weight_count: u32,
|
weight_count: u32,
|
||||||
args: *const QwenKernelArgs,
|
args: *const QwenKernelArgs,
|
||||||
@@ -132,6 +135,7 @@ unsafe extern "C" {
|
|||||||
) -> i32;
|
) -> i32;
|
||||||
pub(super) fn ds4_gpu_set_quality(quality: bool);
|
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_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_ssd_streaming(enabled: bool);
|
||||||
pub(super) fn ds4_gpu_set_model_fd(fd: i32) -> i32;
|
pub(super) fn ds4_gpu_set_model_fd(fd: i32) -> i32;
|
||||||
pub(super) fn ds4_gpu_set_streaming_expert_cache_budget(experts: u32);
|
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_finish() -> i32;
|
||||||
pub(super) fn ds4_gpu_parallel_ffn_abort();
|
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(bytes: u64) -> *mut GpuTensor;
|
||||||
|
pub(super) fn ds4_gpu_tensor_alloc_untracked(bytes: u64) -> *mut GpuTensor;
|
||||||
pub(super) fn ds4_gpu_tensor_view(
|
pub(super) fn ds4_gpu_tensor_view(
|
||||||
base: *const GpuTensor,
|
base: *const GpuTensor,
|
||||||
offset: u64,
|
offset: u64,
|
||||||
@@ -1847,6 +1852,7 @@ impl Context {
|
|||||||
check(unsafe { ds4_gpu_init() }, "Metal initialization")?;
|
check(unsafe { ds4_gpu_init() }, "Metal initialization")?;
|
||||||
unsafe {
|
unsafe {
|
||||||
ds4_gpu_set_glm_model(model.shape.family == ModelFamily::Glm);
|
ds4_gpu_set_glm_model(model.shape.family == ModelFamily::Glm);
|
||||||
|
ds4_gpu_set_qwen_model(false);
|
||||||
ds4_gpu_set_ssd_streaming(ssd_streaming);
|
ds4_gpu_set_ssd_streaming(ssd_streaming);
|
||||||
// Decode enables this only around DS4's eligible resident pre-M5
|
// Decode enables this only around DS4's eligible resident pre-M5
|
||||||
// MXFP4 token path; all other work starts from the portable path.
|
// MXFP4 token path; all other work starts from the portable path.
|
||||||
@@ -1944,6 +1950,7 @@ impl Context {
|
|||||||
check(unsafe { ds4_gpu_init() }, "Metal initialization")?;
|
check(unsafe { ds4_gpu_init() }, "Metal initialization")?;
|
||||||
unsafe {
|
unsafe {
|
||||||
ds4_gpu_set_glm_model(false);
|
ds4_gpu_set_glm_model(false);
|
||||||
|
ds4_gpu_set_qwen_model(true);
|
||||||
ds4_gpu_set_ssd_streaming(false);
|
ds4_gpu_set_ssd_streaming(false);
|
||||||
ds4_gpu_set_decode_pipeline_fast_lookup(0);
|
ds4_gpu_set_decode_pipeline_fast_lookup(0);
|
||||||
}
|
}
|
||||||
@@ -2071,6 +2078,12 @@ impl Buffer {
|
|||||||
.ok_or_else(|| format!("Metal could not allocate {bytes} bytes"))
|
.ok_or_else(|| format!("Metal could not allocate {bytes} bytes"))
|
||||||
}
|
}
|
||||||
|
|
||||||
|
pub(super) fn untracked(bytes: u64) -> Result<Self, String> {
|
||||||
|
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<Self, String> {
|
pub(super) fn view(&self, offset: u64, bytes: u64) -> Result<Self, String> {
|
||||||
// SAFETY: `ds4_gpu_tensor_view` in native/metal/ds4_metal.m bounds-checks
|
// 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
|
// the view, ARC-retains base_obj.buffer, and marks the view non-owning, so
|
||||||
|
|||||||
File diff suppressed because it is too large
Load Diff
@@ -3,15 +3,19 @@ use super::{ChatTurn, ModelSummary};
|
|||||||
use crate::model::ModelChoice;
|
use crate::model::ModelChoice;
|
||||||
use crate::settings::ReasoningMode;
|
use crate::settings::ReasoningMode;
|
||||||
use memmap2::{Mmap, MmapOptions};
|
use memmap2::{Mmap, MmapOptions};
|
||||||
|
use rustc_hash::FxHashMap as HashMap;
|
||||||
use serde::de::{MapAccess, Visitor};
|
use serde::de::{MapAccess, Visitor};
|
||||||
use serde::{Deserialize, Deserializer};
|
use serde::{Deserialize, Deserializer};
|
||||||
use serde_json::Value;
|
use serde_json::Value;
|
||||||
use sha2::{Digest, Sha256};
|
use sha2::{Digest, Sha256};
|
||||||
use std::collections::{BTreeMap, BTreeSet, HashMap};
|
use std::collections::{BTreeMap, BTreeSet, HashSet, VecDeque};
|
||||||
use std::fmt;
|
use std::fmt;
|
||||||
use std::fs::{self, File};
|
use std::fs::{self, File};
|
||||||
use std::io::Read;
|
use std::io::Read;
|
||||||
|
use std::os::unix::fs::FileExt;
|
||||||
use std::path::{Path, PathBuf};
|
use std::path::{Path, PathBuf};
|
||||||
|
use std::sync::mpsc;
|
||||||
|
use std::thread;
|
||||||
|
|
||||||
const MAX_SAFETENSORS_HEADER: u64 = 16 * 1024 * 1024;
|
const MAX_SAFETENSORS_HEADER: u64 = 16 * 1024 * 1024;
|
||||||
const MANIFEST: &[u8] = include_bytes!("../../assets/models/qwen38-flash-next-bare-speed.json");
|
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 PLE_CONV_BYTES: u64 = 184_320;
|
||||||
const MTP_CAPTURE_HIDDEN_BYTES: u64 = 10_240 * 4;
|
const MTP_CAPTURE_HIDDEN_BYTES: u64 = 10_240 * 4;
|
||||||
const MTP_CAPTURE_LOGITS_BYTES: u64 = 248_320 * 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)]
|
#[derive(Deserialize)]
|
||||||
struct Manifest {
|
struct Manifest {
|
||||||
@@ -104,7 +111,9 @@ pub(super) struct ArtifactBindings {
|
|||||||
|
|
||||||
pub(super) struct QwenMap {
|
pub(super) struct QwenMap {
|
||||||
path: PathBuf,
|
path: PathBuf,
|
||||||
map: Mmap,
|
file: File,
|
||||||
|
bytes: u64,
|
||||||
|
map: Option<Mmap>,
|
||||||
}
|
}
|
||||||
|
|
||||||
#[derive(Clone)]
|
#[derive(Clone)]
|
||||||
@@ -123,9 +132,227 @@ pub(super) struct QwenModel {
|
|||||||
memory: MemoryPlan,
|
memory: MemoryPlan,
|
||||||
maps: Vec<QwenMap>,
|
maps: Vec<QwenMap>,
|
||||||
tensors: HashMap<String, QwenTensor>,
|
tensors: HashMap<String, QwenTensor>,
|
||||||
|
ple_reader: PleReader,
|
||||||
identity: [u8; 32],
|
identity: [u8; 32],
|
||||||
}
|
}
|
||||||
|
|
||||||
|
struct PleTask {
|
||||||
|
row: u64,
|
||||||
|
response: mpsc::Sender<(u64, Result<[u8; PLE_ROW_BYTES], String>)>,
|
||||||
|
}
|
||||||
|
|
||||||
|
struct PleReader {
|
||||||
|
workers: Vec<mpsc::Sender<Option<PleTask>>>,
|
||||||
|
threads: Vec<thread::JoinHandle<()>>,
|
||||||
|
hot: HashMap<u64, ([u8; PLE_ROW_BYTES], u64)>,
|
||||||
|
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<String, QwenTensor>) -> Result<Self, String> {
|
||||||
|
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::<Option<PleTask>>();
|
||||||
|
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<Vec<[u8; PLE_ROW_BYTES]>, 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::<Vec<_>>();
|
||||||
|
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>, 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 {
|
impl QwenModel {
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
pub(super) fn open(root: &Path, context: u32) -> Result<Self, String> {
|
pub(super) fn open(root: &Path, context: u32) -> Result<Self, String> {
|
||||||
@@ -150,14 +377,28 @@ impl QwenModel {
|
|||||||
paths.sort();
|
paths.sort();
|
||||||
paths.dedup();
|
paths.dedup();
|
||||||
let mut maps = Vec::with_capacity(paths.len());
|
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 {
|
for path in paths {
|
||||||
let file = File::open(&path).map_err(|error| format!("{}: {error}", path.display()))?;
|
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.
|
// SAFETY: verified managed artifacts remain read-only while the model owns each mapping.
|
||||||
let map = unsafe { MmapOptions::new().map(&file) }
|
let map = unsafe { MmapOptions::new().map(&file) }
|
||||||
.map_err(|error| format!("cannot map {}: {error}", path.display()))?;
|
.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());
|
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
|
let tensors = bindings
|
||||||
.into_iter()
|
.into_iter()
|
||||||
@@ -174,6 +415,7 @@ impl QwenModel {
|
|||||||
(binding.name, tensor)
|
(binding.name, tensor)
|
||||||
})
|
})
|
||||||
.collect::<HashMap<_, _>>();
|
.collect::<HashMap<_, _>>();
|
||||||
|
let ple_reader = PleReader::new(&maps, &tensors)?;
|
||||||
let mut hash = Sha256::new();
|
let mut hash = Sha256::new();
|
||||||
hash.update(b"DS4Server Qwen3.8 checkpoint identity v1");
|
hash.update(b"DS4Server Qwen3.8 checkpoint identity v1");
|
||||||
hash.update(MANIFEST);
|
hash.update(MANIFEST);
|
||||||
@@ -183,6 +425,7 @@ impl QwenModel {
|
|||||||
memory: loaded.memory,
|
memory: loaded.memory,
|
||||||
maps,
|
maps,
|
||||||
tensors,
|
tensors,
|
||||||
|
ple_reader,
|
||||||
identity,
|
identity,
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
@@ -193,12 +436,25 @@ impl QwenModel {
|
|||||||
.ok_or_else(|| format!("Qwen tensor is missing: {name}"))
|
.ok_or_else(|| format!("Qwen tensor is missing: {name}"))
|
||||||
}
|
}
|
||||||
|
|
||||||
|
pub(super) fn tensors(&self) -> impl Iterator<Item = &QwenTensor> {
|
||||||
|
self.tensors.values()
|
||||||
|
}
|
||||||
|
|
||||||
pub(super) fn map(&self, index: usize) -> (&[u8], &Path) {
|
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> {
|
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)
|
let start = usize::try_from(tensor.range.start)
|
||||||
.map_err(|_| format!("{} starts beyond this platform", tensor.name))?;
|
.map_err(|_| format!("{} starts beyond this platform", tensor.name))?;
|
||||||
let end = usize::try_from(tensor.range.end)
|
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))
|
.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<Vec<[u8; PLE_ROW_BYTES]>, 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] {
|
pub(super) fn checkpoint_identity(&self) -> [u8; 32] {
|
||||||
self.identity
|
self.identity
|
||||||
}
|
}
|
||||||
@@ -214,7 +512,7 @@ impl QwenModel {
|
|||||||
pub(super) fn summary(&self) -> ModelSummary {
|
pub(super) fn summary(&self) -> ModelSummary {
|
||||||
ModelSummary {
|
ModelSummary {
|
||||||
model: ModelChoice::Qwen38FlashNext,
|
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(),
|
tensor_count: self.tensors.len(),
|
||||||
vocabulary_size: self.tokenizer.vocab_size(),
|
vocabulary_size: self.tokenizer.vocab_size(),
|
||||||
support_loaded: false,
|
support_loaded: false,
|
||||||
@@ -284,12 +582,15 @@ impl QwenModel {
|
|||||||
let mut core = 0_u64;
|
let mut core = 0_u64;
|
||||||
let mut ple = 0_u64;
|
let mut ple = 0_u64;
|
||||||
for item in &self.maps {
|
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.
|
// SAFETY: each read-only mmap and residency vector remain valid for this call.
|
||||||
if unsafe {
|
if unsafe {
|
||||||
libc::mincore(
|
libc::mincore(
|
||||||
item.map.as_ptr().cast_mut().cast(),
|
map.as_ptr().cast_mut().cast(),
|
||||||
item.map.len(),
|
map.len(),
|
||||||
pages.as_mut_ptr(),
|
pages.as_mut_ptr(),
|
||||||
)
|
)
|
||||||
} != 0
|
} != 0
|
||||||
@@ -301,7 +602,7 @@ impl QwenModel {
|
|||||||
));
|
));
|
||||||
}
|
}
|
||||||
let bytes = (pages.iter().filter(|value| **value & 1 != 0).count() * page)
|
let bytes = (pages.iter().filter(|value| **value & 1 != 0).count() * page)
|
||||||
.min(item.map.len()) as u64;
|
.min(map.len()) as u64;
|
||||||
if item
|
if item
|
||||||
.path
|
.path
|
||||||
.file_name()
|
.file_name()
|
||||||
@@ -846,8 +1147,8 @@ mod tests {
|
|||||||
assert_eq!(
|
assert_eq!(
|
||||||
direct,
|
direct,
|
||||||
[
|
[
|
||||||
248_045, 846, 198, 12_675, 248_046, 198, 248_045, 74_455, 198, 13_314, 741, 29,
|
248_045, 846, 198, 12_675, 248_046, 198, 248_045, 74_455, 198, 248_068, 271,
|
||||||
271, 510, 26_003, 29, 271,
|
248_069, 271,
|
||||||
]
|
]
|
||||||
);
|
);
|
||||||
assert_eq!(
|
assert_eq!(
|
||||||
|
|||||||
@@ -62,7 +62,6 @@ struct QwenTokenizerFile {
|
|||||||
struct QwenAddedToken {
|
struct QwenAddedToken {
|
||||||
id: usize,
|
id: usize,
|
||||||
content: String,
|
content: String,
|
||||||
special: bool,
|
|
||||||
}
|
}
|
||||||
|
|
||||||
#[derive(Deserialize)]
|
#[derive(Deserialize)]
|
||||||
@@ -225,9 +224,7 @@ impl Tokenizer {
|
|||||||
));
|
));
|
||||||
}
|
}
|
||||||
tokens[token.id] = token.content.as_bytes().to_vec();
|
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) {
|
if tokens[..=maximum].iter().any(Vec::is_empty) {
|
||||||
return Err("Qwen tokenizer token ids are not contiguous".into());
|
return Err("Qwen tokenizer token ids are not contiguous".into());
|
||||||
|
|||||||
Reference in New Issue
Block a user