hugo/DS4Server
1
0
Fork 0
Code Issues 2 Pull Requests Actions Packages Projects Releases Wiki Activity

Implement exact Qwen sparse attention

Browse Source
This commit is contained in:
Georg Bauer
2026-09-03 21:36:33 +02:00
parent 87ccf67d0c
commit dda70c3793
3 changed files with 909 additions and 84 deletions
Download Patch File Download Diff File Expand all files Collapse all files

35
src/engine/qwen.rs
View File

@@ -21,6 +21,12 @@ const CORE_BYTES: u64 = 71_742_682_599;
const PLE_BYTES: u64 = 32_000_154_008;
const MTP_BYTES: u64 = 1_672_575_532;
const KV_BYTES_PER_TOKEN: u64 = 24_576;
const QSA_RAW_BYTES_PER_TOKEN: u64 = 3_072;
const QSA_POOLED_BYTES_PER_BLOCK: u64 = 3_072;
const QSA_POOL_RATIO: u64 = 4;
const QSA_FIXED_SCRATCH_BYTES: u64 = 6_656;
const QSA_TOPK_SCRATCH_BYTES_PER_BLOCK: u64 = 8;
const MTP_TOKEN_RESERVE: u64 = 3;
const GDN_STATE_BYTES: u64 = 113_246_208;
const GDN_CONV_BYTES: u64 = 2_211_840;
const PLE_CONV_BYTES: u64 = 184_320;
@@ -634,14 +640,31 @@ pub(super) fn memory_plan(
if prefill_chunk == 0 {
return Err("Qwen prefill chunk must be positive".into());
}
let kv = KV_BYTES_PER_TOKEN
.checked_mul(u64::from(context))
let token_capacity = u64::from(context)
.checked_add(MTP_TOKEN_RESERVE)
.ok_or_else(|| "Qwen attention capacity overflows".to_owned())?;
let block_capacity = token_capacity.div_ceil(QSA_POOL_RATIO);
let topk_scratch = if context > 2_048 {
u64::from(context) / QSA_POOL_RATIO * QSA_TOPK_SCRATCH_BYTES_PER_BLOCK
} else {
0
};
let kv = (KV_BYTES_PER_TOKEN + QSA_RAW_BYTES_PER_TOKEN)
.checked_mul(token_capacity)
.and_then(|bytes| {
QSA_POOLED_BYTES_PER_BLOCK
.checked_mul(block_capacity)
.and_then(|pooled| bytes.checked_add(pooled))
})
.ok_or_else(|| "Qwen KV memory size overflows".to_owned())?;
let kv_and_recurrent = kv
.checked_add(GDN_STATE_BYTES + GDN_CONV_BYTES + PLE_CONV_BYTES)
.ok_or_else(|| "Qwen recurrent memory size overflows".to_owned())?;
let prefill_transient = u64::from(prefill_chunk)
.checked_mul((4 * 2_560 + 2_048 + 2_048 + 6_144 + 6_144) * 2)
.and_then(|bytes| bytes.checked_add(block_capacity * 4))
.and_then(|bytes| bytes.checked_add(QSA_FIXED_SCRATCH_BYTES))
.and_then(|bytes| bytes.checked_add(topk_scratch))
.ok_or_else(|| "Qwen prefill transient size overflows".to_owned())?;
let admitted_mtp = if enable_mtp { MTP_BYTES } else { 0 };
let admission = CORE_BYTES
@@ -681,12 +704,12 @@ mod tests {
assert_eq!(plan.resident_core, CORE_BYTES);
assert_eq!(plan.mapped_ple, PLE_BYTES);
assert_eq!(plan.optional_mtp, MTP_BYTES);
assert_eq!(plan.kv_and_recurrent, 6_558_093_312);
assert_eq!(plan.prefill_transient, 27_262_976);
assert_eq!(plan.admission, 80_000_614_419);
assert_eq!(plan.kv_and_recurrent, 7_564_812_288);
assert_eq!(plan.prefill_transient, 28_056_068);
assert_eq!(plan.admission, 81_008_126_487);
let without_mtp = memory_plan(262_144, false, 512).unwrap();
assert_eq!(without_mtp.optional_mtp, MTP_BYTES);
assert_eq!(without_mtp.admission, 78_328_038_887);
assert_eq!(without_mtp.admission, 79_335_550_955);
assert!(memory_plan(0, false, 512).is_err());
assert!(memory_plan(262_145, false, 512).is_err());
assert!(memory_plan(1, false, 0).is_err());