Skip to content

fix(driver): keep Qwen3.5/3.6 linear-attention decay params in f32 so exp() runs on Metal #465

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ingim merged 1 commit into from
Jul 10, 2026
Merged

fix(driver): keep Qwen3.5/3.6 linear-attention decay params in f32 so exp() runs on Metal #465

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 6 additions & 1 deletion driver/portable/src/graph_qwen3_5.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);

Expand Down
22 changes: 20 additions & 2 deletions driver/portable/src/model.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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
Expand All @@ -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) {
Expand Down
Loading