Skip to content

glm: keep rope models on the indexed causal attention kernel (fixes #932) #936

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

Open
garnetlyx wants to merge 1 commit into
base: main
Choose a base branch
from
Open

glm: keep rope models on the indexed causal attention kernel (fixes #932) #936

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
6 changes: 6 additions & 0 deletions ds4.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -45275,6 +45275,12 @@ static bool glm_graph_use_flash_attention_prefill(uint32_t n_tokens) {

static DS4_MAYBE_UNUSED bool glm_graph_use_dense_compact_attention_prefill(
uint32_t n_tokens) {
/* The dense compact path scores the shared latent cache directly and
* carries no separate RoPE contribution, so it is only exact for models
* whose DSA cache row has no RoPE tail (GLM-5.3 Flash, n_rot == 0).
* Models with n_rot == 64 (GLM-5.2, full GLM-5.3) must keep the indexed
* kernel, which adds the rotated RoPE term from k_rope_cache. */
if (DS4_N_ROT != 0u) return false;
if (!glm_graph_use_flash_attention_prefill(n_tokens)) return false;
#if !defined(__APPLE__) && !defined(DS4_ROCM_BUILD) && !defined(DS4_NO_GPU)
/* The CUDA GEMM setup crosses over the scalar online kernel near 256
Expand Down