From dec555a13e800ba724ac0be4c7d9ec0204267f53 Mon Sep 17 00:00:00 2001 From: Zhiyao Ma Date: Fri, 10 Jul 2026 15:43:12 -0400 Subject: [PATCH] fix(driver): keep Qwen3.5/3.6 linear-attention decay params in f32 so exp() runs on Metal MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Running a bf16 Qwen3.6 checkpoint on Metal warned: 48 graph node(s) fell back to CPU (no Metal kernel): UNARYx48 One per linear-attention layer: the gated-delta-rule block computes `exp(A_log)` directly on the bf16 `linear_attn.A_log` weight, and ggml-metal has no bf16 unary kernel, so every `exp` node fell back to CPU and split the graph. `A_log`/`dt_bias` slipped past the load-time f32 promotion, whose suffix match only covers `*norm.weight`/`*.bias`. Fix, scoped to PieArch::Qwen3_5: - model.cpp: promote `linear_attn.{A_log,dt_bias}` to f32 at load time ([num_v_heads] floats per layer, ~9 KiB total on 27B), and exempt them from the f32→bf16 downcast. - graph_qwen3_5.cpp: defensive f32 cast before `exp(A_log)` for checkpoints that bypass the load-time upcast. Verified on Qwen3.6-27B (M4 Pro): the warning is gone and the whole graph schedules onto MTL0. Qwen3-0.6B smoke test unchanged. --- driver/portable/src/graph_qwen3_5.cpp | 7 ++++++- driver/portable/src/model.cpp | 22 ++++++++++++++++++++-- 2 files changed, 26 insertions(+), 3 deletions(-) diff --git a/driver/portable/src/graph_qwen3_5.cpp b/driver/portable/src/graph_qwen3_5.cpp index 1a949fd27..cdb0e0877 100644 --- a/driver/portable/src/graph_qwen3_5.cpp +++ b/driver/portable/src/graph_qwen3_5.cpp @@ -67,7 +67,12 @@ ggml_tensor* build_qwen3_5_linear_layer( : ggml_cast(ctx, L.lin_dt_bias, GGML_TYPE_F32); auto* alpha_dt = ggml_add(ctx, alpha, dt_bias); auto* sp_alpha = ggml_softplus(ctx, alpha_dt); - auto* A_pos = ggml_exp(ctx, L.lin_A_log); // [H] + // Metal has no bf16 unary kernel — cast to f32 before `exp` (the cast + // itself has a Metal CPY pipeline) so the node stays on the GPU. + auto* A_log = (L.lin_A_log->type == GGML_TYPE_F32) + ? L.lin_A_log + : ggml_cast(ctx, L.lin_A_log, GGML_TYPE_F32); + auto* A_pos = ggml_exp(ctx, A_log); // [H] auto* g = ggml_mul(ctx, sp_alpha, A_pos); // [H, n_tokens, n_seqs] g = ggml_scale(ctx, g, -1.0f); diff --git a/driver/portable/src/model.cpp b/driver/portable/src/model.cpp index aa5c9c5ad..bdc4bee6e 100644 --- a/driver/portable/src/model.cpp +++ b/driver/portable/src/model.cpp @@ -58,6 +58,20 @@ bool is_small_weight_for_upcast(const std::string& hf_name) { pie_driver_common::ends_with(hf_name, kBiasSuffix); } +// Qwen3.5/3.6-only companion to `is_small_weight_for_upcast`: the +// gated-delta-rule per-head decay params ([num_v_heads] floats per +// linear-attention layer). `A_log` feeds `ggml_exp` directly; Metal has +// no bf16 unary kernel, so leaving it bf16 forces one CPU-fallback node +// per linear-attention layer. Callers must gate on +// `hparams_.arch == PieArch::Qwen3_5` so other archs' load behavior is +// untouched. +bool is_qwen3_5_linear_decay_weight(const std::string& hf_name) { + static constexpr std::string_view kALogSuffix = "linear_attn.A_log"; + static constexpr std::string_view kDtBiasSuffix = "linear_attn.dt_bias"; + return pie_driver_common::ends_with(hf_name, kALogSuffix) || + pie_driver_common::ends_with(hf_name, kDtBiasSuffix); +} + ggml_type st_to_ggml_dtype(StDtype dt, const std::string& tensor_name) { switch (dt) { case StDtype::F32: return GGML_TYPE_F32; @@ -535,9 +549,13 @@ ggml_tensor* Model::declare_(const std::string& hf_name) { // implicit `ggml_cast` in `norm_scale` becomes a no-op. The actual // bf16 -> f32 byte conversion is a Rust storage-program TileMap. const bool gguf_typed = t.ggml_type_override >= 0; + const bool small_f32_weight = + is_small_weight_for_upcast(hf_name) || + (hparams_.arch == PieArch::Qwen3_5 && + is_qwen3_5_linear_decay_weight(hf_name)); const bool upcast = !gguf_typed && t.dtype == StDtype::BF16 - && is_small_weight_for_upcast(hf_name); + && small_f32_weight; // Some HF releases (notably Gemma-2-2b) store weights as F32. The F32 // mul_mat path is ~2x slower than BF16 on tensor-core hardware and // doubles VRAM. Downcast large matmul weights at load time, matching @@ -556,7 +574,7 @@ ggml_tensor* Model::declare_(const std::string& hf_name) { const bool keep_f32 = keep_f32_env || hparams_.arch == PieArch::Csm; const bool downcast = !gguf_typed && t.dtype == StDtype::F32 - && !is_small_weight_for_upcast(hf_name) + && !small_f32_weight && !keep_f32; const bool fp8_decode = !gguf_typed && t.dtype == StDtype::F8_E4M3; if (!gguf_typed && t.dtype == StDtype::F8_E5M2) {