diff --git a/server/src/qwen35/qwen35_backend.cpp b/server/src/qwen35/qwen35_backend.cpp index 4c9a727da..71a1cb6c6 100644 --- a/server/src/qwen35/qwen35_backend.cpp +++ b/server/src/qwen35/qwen35_backend.cpp @@ -1322,7 +1322,7 @@ void Qwen35Backend::kvflash_upload_mask() { const size_t need = (size_t)sg_.attn_mask->ne[0] * sg_.attn_mask->ne[1]; if (kvflash_mask_buf_.size() != need || kvflash_pager_.epoch() != kvflash_mask_epoch_) { kvflash_mask_buf_.assign(need, F16_NEG_INF); - kvflash_pager_.fill_slot_mask(kvflash_mask_buf_.data()); // q row 0 + kvflash_pager_.fill_slot_mask(kvflash_mask_buf_.data()); // q row 0; pager caches kvflash_mask_epoch_ = kvflash_pager_.epoch(); } // Upload before EVERY compute: the input tensor's buffer region is @@ -1532,6 +1532,17 @@ bool Qwen35Backend::do_ar_decode(int committed, int n_gen, } // AR decode loop for remaining tokens + // + // Graph reuse: build_target_step rebuilds the ggml cgraph every step, + // which resets the ggml-cuda CUDA-graph warmup counter and prevents + // replay. The only shape that varies across steps is the FA view window + // (win_len_padded = round_up(committed+1, 256)), which only advances + // every 256 decode steps. So we rebuild only when committed/256 crosses + // a boundary; within a bucket we reuse the cached graph and just update + // the mutable input tensors. + const bool ar_graph_reuse = std::getenv("DFLASH_AR_NO_REUSE") == nullptr; + ar_decode_fa_bucket_ = -1; // force first-step build + for (int i = initial_emitted; i < n_gen; i++) { int32_t tok = out_tokens.back(); @@ -1543,17 +1554,23 @@ bool Qwen35Backend::do_ar_decode(int committed, int n_gen, // kvflash: graph carries a slot-validity mask alongside the // step-invariant set_rows write; the FA span clamps to the pool. const bool pool = kvflash_active(); - if (!build_target_step(sg_, w_, cache_, target_backend_, - /*kv_start=*/committed, /*n_tokens=*/1, - /*with_mask=*/pool, /*capture=*/false, - /*capture_delta_intermediate=*/false, - /*fa_window=*/0, - /*last_token_logits_only=*/false, - cfg_.kq_stride_pad, - should_capture_moe_router(), - /*kvflash_mask=*/pool, - /*capture_qk=*/pool && kvflash_qk_policy_)) { - return false; + + const int fa_bucket = committed >> 8; // committed / 256 + const bool need_rebuild = !ar_graph_reuse || (fa_bucket != ar_decode_fa_bucket_); + if (need_rebuild) { + if (!build_target_step(sg_, w_, cache_, target_backend_, + /*kv_start=*/committed, /*n_tokens=*/1, + /*with_mask=*/pool, /*capture=*/false, + /*capture_delta_intermediate=*/false, + /*fa_window=*/0, + /*last_token_logits_only=*/false, + cfg_.kq_stride_pad, + should_capture_moe_router(), + /*kvflash_mask=*/pool, + /*capture_qk=*/pool && kvflash_qk_policy_)) { + return false; + } + ar_decode_fa_bucket_ = fa_bucket; } // Fill kv_write_rows with this step's cache slot for set_rows: diff --git a/server/src/qwen35/qwen35_backend.h b/server/src/qwen35/qwen35_backend.h index bc4fcdc0b..4d664a7bc 100644 --- a/server/src/qwen35/qwen35_backend.h +++ b/server/src/qwen35/qwen35_backend.h @@ -231,6 +231,8 @@ class Qwen35Backend : public ModelBackend { StepGraph draft_sg_; // draft forward StepGraph proj_sg_; // lm-head projection (remote-lm-head mode) + int ar_decode_fa_bucket_ = -1; + // ── Draft feature mirror (cross-GPU feature transfer) ──────────── DraftFeatureMirror feature_mirror_; DFlashDraftIpcClient remote_draft_; diff --git a/server/src/qwen35/qwen35_graph_options.h b/server/src/qwen35/qwen35_graph_options.h new file mode 100644 index 000000000..d98af6fa8 --- /dev/null +++ b/server/src/qwen35/qwen35_graph_options.h @@ -0,0 +1,30 @@ +#pragma once + +#include "ggml.h" + +#include +#include + +namespace dflash::common { + +inline bool qwen35_env_flag_is_one(const char * name) { + const char * value = std::getenv(name); + return value != nullptr && std::strcmp(value, "1") == 0; +} + +inline bool qwen35_should_use_graph_wht_k_rotation(ggml_type kv_k_type) { + if (std::getenv("DFLASH_NO_WHT") != nullptr) { + return false; + } + if (kv_k_type == GGML_TYPE_TQ3_0) { + return false; + } + if (qwen35_env_flag_is_one("DFLASH_FORCE_WHT")) { + return true; + } + // q4_0 on Ampere is already fast enough without graph-level FWHT; avoiding + // it removes extra prefill/decode kernels in the parity path. + return kv_k_type != GGML_TYPE_Q4_0; +} + +} // namespace dflash::common diff --git a/server/src/qwen35/qwen35_target_graph.cpp b/server/src/qwen35/qwen35_target_graph.cpp index 22a8e3a61..2803dd183 100644 --- a/server/src/qwen35/qwen35_target_graph.cpp +++ b/server/src/qwen35/qwen35_target_graph.cpp @@ -33,6 +33,7 @@ #include "internal.h" #include "delta_net_chunked.h" #include "kv_quant.h" +#include "qwen35/qwen35_graph_options.h" #include "qwen35_ops.h" #include "qwen35moe_ffn.h" @@ -128,10 +129,10 @@ bool create_target_cache_partial(const TargetWeights & w, out.kv_k_type = kv_k_type; out.kv_v_type = kv_v_type; - // Graph-level FWHT K-rotation (TurboQuant-style outlier spreading with - // standard quant types that keep fast FA kernel paths on all arches). - // Skip for TQ3_0 K cache — that type already applies WHT during quantization. - out.kv_k_rotated = (kv_k_type != GGML_TYPE_TQ3_0); + // Graph-level FWHT K-rotation (TurboQuant-style outlier spreading). + // q4_0 on Ampere skips it by default for parity with llama.cpp's leaner + // q4_0 path; DFLASH_FORCE_WHT=1 restores the old graph-level rotation. + out.kv_k_rotated = qwen35_should_use_graph_wht_k_rotation(kv_k_type); const bool needs_256_stride = kv_k_type == GGML_TYPE_TQ3_0 || kv_v_type == GGML_TYPE_TQ3_0; @@ -789,7 +790,10 @@ static ggml_tensor * build_delta_net_block( // qkv_mixed currently is [conv_channels, n_tokens, n_seqs]; we need // [n_tokens, conv_channels, n_seqs] to concat on dim 0. - ggml_tensor * qkv_T = ggml_transpose(ctx, qkv_mixed); + // Materialize the transpose once so both concat inputs are contiguous + // → concat_cont fast path fires (saves 1 ggml_cont per layer vs wrapping + // both inputs individually). + ggml_tensor * qkv_T = ggml_cont(ctx, ggml_transpose(ctx, qkv_mixed)); ggml_tensor * conv_input = ggml_concat(ctx, conv_states_r, qkv_T, 0); // conv_input: [kernel-1 + n_tokens, conv_channels, n_seqs] @@ -864,9 +868,14 @@ static ggml_tensor * build_delta_net_block( q_c = ggml_l2_norm(ctx, q_c, w.rms_eps); k_c = ggml_l2_norm(ctx, k_c, w.rms_eps); - // Repeat Q and K from num_k_heads to num_v_heads so they match V's layout - // (only needed if not using the fused op's broadcast support). - if (num_k_heads != num_v_heads) { + // Repeat Q and K from num_k_heads to num_v_heads so they match V's layout. + // In pure AR (cap==nullptr, parent_ids==nullptr) the kernel's fastmodulo + // broadcast handles the head mismatch natively — skip the repeat to save + // one tensor allocation and a memory-bandwidth copy per SSM layer per step. + // In tree/capture mode keep the repeat: the DFS parent-reload in the kernel + // indexes intermediate states by h_idx (v-head), so Q/K must already be + // in the expanded [num_v_heads] layout before the kernel runs. + if (num_k_heads != num_v_heads && (cap != nullptr || parent_ids != nullptr)) { q_c = ggml_repeat_4d(ctx, q_c, head_k_dim, num_v_heads, n_seq_tokens, n_seqs); k_c = ggml_repeat_4d(ctx, k_c, head_k_dim, num_v_heads, n_seq_tokens, n_seqs); } @@ -1006,12 +1015,12 @@ static ggml_tensor * build_delta_net_block( ggml_build_forward_expand(gf, ggml_cpy(ctx, new_state, s)); } - // ── Gated output norm: rms_norm(output) * silu(z_4d) + // ── Gated output norm: rms_norm(output) * weight * silu(z_4d) + // fused: swiglu_split(z_4d, normed) = silu(z_4d) * normed (1 GLU kernel vs 2) ggml_tensor * z_4d = ggml_reshape_4d(ctx, z, head_v_dim, num_v_heads, n_seq_tokens, n_seqs); ggml_tensor * output_n = ggml_rms_norm(ctx, rms_norm_input_f32(ctx, output), w.rms_eps); output_n = ggml_mul(ctx, output_n, L.ssm_norm); - ggml_tensor * z_silu = ggml_silu(ctx, z_4d); - output_n = ggml_mul(ctx, output_n, z_silu); + output_n = ggml_swiglu_split(ctx, z_4d, output_n); // Reshape to [d_inner, n_tokens] ggml_tensor * flat = ggml_reshape_3d(ctx, output_n,