[sanmai]

Overview

This is just a prototype. No review is needed.

Additional information

	Query chunking (2048 vs off)	without kq_mask
CUDA0 compute buffer, ub 8192	about -2.4 GB	10957 -> 6885 MiB (-4.0 GiB)
CUDA0 compute buffer, ub 4096	not recorded	6112 -> 4076 MiB (-2.0 GiB)
Pinned host buffer, ub 8192	-4.6 GB on CPU	5526 -> 1430 MiB (-4.0 GiB)
Largest MoE cache, ub 8192	9000 -> 11500 (+2500)	11500 -> 15750 (+4250)
Largest MoE cache, ub 4096	14700 -> 16700 (+2000)	16700 -> 18700 (+2000)

#   ub    cache MiB      prefill t/s  decode t/s  VRAM peak MiB
#   4096  18700 (max)    1130-1197    48.3-54.7   32062
#   8192  15750 (max)    1740-1760    45.4-48.8   32050

Before 72e03e1:

  ub    chunk  cache MiB      prefill t/s  decode t/s  VRAM peak MiB
  2048  -      17500          786          50.6        30176
  2048  -      18500          796          53.8        31178
  2048  -      19000 (max)    804          53.5        31680
  2048  -      16200          767          48.5        28876
  2048  -      11500          711          38.4        24182
  4096  2048   15700          1167         47.3        31100
  4096  2048   16200          1186         47.5        31598
  4096  2048   16700 (max)    1180         48.5        32102
  4096  1024   15700 (max)    1167         45.4        31392
  4096  0      14700          1156         45.1        31066
  8192  2048   11500 (max)    1662         38.5        31892
  8192  1024   10500          1640         37.3        31968
  8192  0      9000           1593         32.5        31732
  8192  2048   18000 -c 131072  1753       52.1        31976

Requirements

• I have read and agree with the contributing guidelines
• AI usage disclosure: YES

[sanmai]

@dagbdagb Replying here. Asked one of the agents to benchmark and find the best setting.

[am17an]

Is the main problem here at long context? i.e. uploading the kq mask?

[dagbdagb]

Not entirely sure if that question was directed at me or not.

If it is: The help-text for --moe-cache-mib N has 0 guidance for sizing N. Unless one is very deep into the matter or developed the feature, one can only resort to guess-work. Or automated testing. Certainly not the only such parameter in the README, but still.

[am17an]

@dagbdagb it was not for you. Re the guidance, yes I will make some changes there soon

[sanmai]

These are ideas pulled straight from Strata, with the goal to reduce memory footprint of work buffers:

• QSA query chunking (c391f0f) goes like this:

  • Take the first 2048 tokens of the batch.
  • Build the mask, only a quarter of full 8192 batch
  • Run the attention
  • Free that mask and reuse
  • Join the outputs, and since chunks attend to the same kv cache the result is same

• Removal the kq_mask from QSA (2020bb6, 72e03e1) that frees about 2 GiB at ub 4096 and about 4 GiB at ub 8192, so we all can use a larger expert cache and get better overall perf at ubatch 8192

Shrinking pool_mask is also an option.

And there's a bug with CUDA top_k tie ordering - haven't looked into it.

ub	Before #30640	Now │
ub 4096	1180 / 48.5 t/s at cache 16700	1160 / 50 t/s at cache 18700 │
ub 8192	1660 / 38.5 t/s at cache 11500	1750 / 47 t/s at cache 15500 │

1750 t/s pp is the major win here

[sanmai]

@am17an Feel free to slice and dice this thing into pieces as much as you want or need.

[sanmai]

2048 was a sweet spot on 5090

at 1024 the launch tax starts to show up

[am17an]

Okay so this is because the kq_mask being the largest user of the compute buffer at these ub sizes. But do we need to allocate the kq_mask for this model since we're using sparse attention? cc @ServeurpersoCom if you have a good idea.

EDIT: maybe I'm not correct and need to measure again.

[sanmai]

all 12 attention layers are QSA so we should not need kq_mask

[am17an]

@dagbdagb maybe try the script in #30966

[ServeurpersoCom]

Okay so this is because the kq_mask being the largest user of the compute buffer at these ub sizes. But do we need to allocate the kq_mask for this model since we're using sparse attention? cc @ServeurpersoCom if you have a good idea.

You are right, it is not needed for qwen4exp. The QSA selection is already causal and sequence local by construction: pool_mask only keeps pools of the token's own sequence that end at or before its rank, the tail holds its own cells up to itself, and dead picks or n_kv sentinels are skipped by top_k_mask. So kq_mask just multiplies the kept cells by zero, and in the real GGUF all 12 attention layers are QSA (the other 36 are GDN), so nothing else reads it.

sanmai's numbers are correct, I get exactly the same on the real UD-Q4_K_XL: feeding top_k_mask a broadcast row of zeros instead of kq_mask takes the pinned host buffer from 5526 to 1430 MiB and the compute buffer down by 4 GiB at -c 262144 -ub 8192, which is the mask itself (262144 x 8192 x 2 bytes in F16). The logits at -c 16384 are byte identical to master.

If we want to go all the way, what remains is the dense [n_kv, n_tokens] selection mask FA consumes, which the CUDA sparse FA path already turns back into indices for the shapes it covers. The query chunking in this PR bounds it, and a possible follow-up would be letting flash_attn_ext take the selected cell indices directly, which would remove it for qwen4exp, glm5-next and the DSA models.

[sanmai]

Just QSA without kq_mask for a pair of backends:

• glm5-next : QSA without kq_mask (+ DSA) #31009
• qwen4exp : build the QSA selection without the kq_mask #30965
  (I'll merge the last one into this PR in a moment so they follow the same approach)

I extracted them because they pack the most punch (4 GB saved on GPU at ub 8192), but if that's a preference to keep this PR as a whole, I have nothing against it.

[am17an]

I'm wondering how m-rope interacts with this whole thing. I don't have a good understanding of it at the moment in relation on how it interacts with the kq mask

[sanmai]

Looking into the build failure.

Tests are passing on my end.

[ServeurpersoCom]

I'm wondering how m-rope interacts with this whole thing. I don't have a good understanding of it at the moment in relation on how it interacts with the kq mask

M-RoPE only matters for image tokens, which share pos[0]. kq_mask orders them by (y, x) through is_2d_gt, while the k-pool layout orders a sequence by (pos, cell) in kpool_rank, so ties go by cell index. An image is normally written in token order into increasing cells, so both orders agree and the selection is the same with or without kq_mask; a single image test on the real model gives identical output.

They can only disagree when find_slot wraps around the ring in the middle of an image, so a later patch lands in a lower cell. That mismatch is already in the indexer on master: kq_mask only hides the half where the selection takes a later patch, it cannot bring back an earlier one the indexer dropped. The clean fix is to sort the k-pool layout by (pos, y, x), the ext positions are already read there for new_pool_pos, and then kq_mask is not needed in that case either.

[am17an]

Okay then this should we merge this?

[ServeurpersoCom]

Okay then this should we merge this?

I would start with #30965, which isolates the kq_mask part. I tested that exact diff on the real UD-Q4_K_XL: logits byte identical to master at -c 16384, -4 GiB on both the compute and pinned host buffers at -c 262144 -ub 8192, and the new test-backend-ops cases pass. The chunking here bounds the selection mask but keeps it dense, so I would hold it until we know whether the index based FA mentioned above is the way to go.

[am17an]

it looks it may help with large ubatch on dsv4.1/glm5.3 as well