Skip to content

hrx: score each key once above 2048 keys, instead of on every value lane (engine #991) #214

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
bong-water-water-bong merged 1 commit into from
Nov 6, 2026
Merged

hrx: score each key once above 2048 keys, instead of on every value lane (engine #991) #214

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
18 changes: 15 additions & 3 deletions ggml/src/ggml-hrx/dispatch_registration/common/dispatch-small-rows.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,7 @@ static constexpr KernelCatalogRef kClampInplaceKernel = GGML_HRX_KERNEL_REF("lo
static constexpr KernelCatalogRef kCopyF32F16Kernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_copy_strided_f32_f16");
static constexpr KernelCatalogRef kAttentionStridedKernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_attention_strided_f32_f16");
static constexpr KernelCatalogRef kAttentionRowsKernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_attention_rows_f32_f16");
static constexpr KernelCatalogRef kAttentionRowsBlockedKernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_attention_rows_blocked_f32_f16");
static constexpr KernelCatalogRef kRopeRotateHalfKernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_rope_rotate_half_f32");
static constexpr KernelCatalogRef kGegluStridedKernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_geglu_strided_f32");
static constexpr KernelCatalogRef kMulMatSmallF16Kernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_small_f16_f32");
Expand Down Expand Up @@ -466,9 +467,20 @@ static bool match_attention_strided(const DispatchMatchContext & context, Dispat
return false;
}
Dispatch dispatch;
// scores once per (query, head) in workgroup memory when they fit (ONEBIT_HRX_ATTN_PER_LANE=1: the old way)
const bool rows = nkv <= 2048 && d <= 256 && dv <= 256 && std::getenv("ONEBIT_HRX_ATTN_PER_LANE") == nullptr;
dispatch.kernel = make_kernel_specialization(rows ? kAttentionRowsKernel : kAttentionStridedKernel);
// The rows kernel scores each key once with the whole score row in workgroup memory, which is what limits it to
// 2048 keys. "strided" has one lane per value dim and each lane recomputes the whole q.k dot product for every
// key, i.e. nkv*dv*d instead of nkv*(d + dv) - 80x the arithmetic at d = dv = 160. That was the only kernel left
// past 2048 keys, so a long prompt fell off a cliff exactly at the 2048-key switch (measured on Zamba2-2.7B:
// pp2048/pp512 = 0.68, pp2560 = 0.18, pp4096 = 0.065, and forcing this kernel at 512 tokens cost 78% of
// throughput). The blocked rows kernel walks the keys in blocks of 2048 with the online softmax, so it keeps the
// rows arithmetic at any length and "strided" stays for the head sizes it alone admits.
// (ONEBIT_HRX_ATTN_PER_LANE=1 forces the old per-lane kernel at every length: the A/B that established the above.)
const bool per_lane = std::getenv("ONEBIT_HRX_ATTN_PER_LANE") != nullptr;
if (per_lane || d > 256 || dv > 256) {
dispatch.kernel = make_kernel_specialization(kAttentionStridedKernel);
} else {
dispatch.kernel = make_kernel_specialization(nkv <= 2048 ? kAttentionRowsKernel : kAttentionRowsBlockedKernel);
}
auto & ip = dispatch.kernel.integer_parameters;
ip.emplace("qk_size", d);
ip.emplace("v_size", dv);
Expand Down
169 changes: 169 additions & 0 deletions ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/manifest.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7377,6 +7377,175 @@
},
"compile_dependencies": []
},
{
"name": "ggml_attention_rows_blocked_f32_f16",
"family": "loom_libs",
"symbol": "ggml_attention_rows_blocked_f32_f16",
"source": "ops/small_rows_f32.loom",
"workload_parameters": [
{
"name": "qk_size",
"type": "index"
},
{
"name": "v_size",
"type": "index"
},
{
"name": "q_count",
"type": "index"
},
{
"name": "kv_count",
"type": "index"
},
{
"name": "head_count",
"type": "index"
},
{
"name": "kv_head_count",
"type": "index"
},
{
"name": "q_s1",
"type": "index"
},
{
"name": "q_s2",
"type": "index"
},
{
"name": "k_s1",
"type": "index"
},
{
"name": "k_s2",
"type": "index"
},
{
"name": "v_s1",
"type": "index"
},
{
"name": "v_s2",
"type": "index"
},
{
"name": "m_s1",
"type": "index"
},
{
"name": "q_extent",
"type": "index"
},
{
"name": "k_extent",
"type": "index"
},
{
"name": "v_extent",
"type": "index"
},
{
"name": "m_extent",
"type": "index"
}
],
"launch_parameters": [
{
"name": "qk_size",
"type": "index"
},
{
"name": "v_size",
"type": "index"
},
{
"name": "q_count",
"type": "index"
},
{
"name": "kv_count",
"type": "index"
},
{
"name": "head_count",
"type": "index"
},
{
"name": "kv_head_count",
"type": "index"
},
{
"name": "q_s1",
"type": "index"
},
{
"name": "q_s2",
"type": "index"
},
{
"name": "k_s1",
"type": "index"
},
{
"name": "k_s2",
"type": "index"
},
{
"name": "v_s1",
"type": "index"
},
{
"name": "v_s2",
"type": "index"
},
{
"name": "m_s1",
"type": "index"
},
{
"name": "q_extent",
"type": "index"
},
{
"name": "k_extent",
"type": "index"
},
{
"name": "v_extent",
"type": "index"
},
{
"name": "m_extent",
"type": "index"
}
],
"bindings": [
"query",
"key",
"value",
"mask",
"output"
],
"binding_access": [
"read",
"read",
"read",
"read",
"read_write"
],
"target_selector": "",
"compile_recipe": {
"mode": "direct",
"primary_sources": [
"ops/small_rows_f32.loom"
],
"library_sources": []
},
"compile_dependencies": []
},
{
"name": "ggml_rope_rotate_half_f32",
"family": "loom_libs",
Expand Down
162 changes: 162 additions & 0 deletions ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/small_rows_f32.loom
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -971,6 +971,168 @@ kernel.def target(@ggml_small_rows_f32_gfx11_wave64) export("ggml_attention_rows
kernel.return
}

// The same one-score-per-key attention for contexts longer than one score row. The rows kernel above keeps the
// whole score row in workgroup memory, which is what caps it at 2048 keys; past that the registration used to fall
// back to the strided kernel, whose lanes each recompute the whole q.k dot product for every key (cost nkv*dv*d
// instead of nkv*(d + dv) - 80x the arithmetic at d = dv = 160, measured as a 15x prefill collapse on Zamba2 at
// -p 4096 versus -p 512). This kernel keeps the rows layout and walks the keys in blocks of 2048 with the online
// softmax of the decode kernels (running max and sum, output row rescaled per block), so the cost stays
// nkv*(d + dv) at any length. Workgroup memory: query row 1 KB + score row 8 KB + output row 1 KB.
kernel.def target(@ggml_small_rows_f32_gfx11_wave64) export("ggml_attention_rows_blocked_f32_f16") @ggml_attention_rows_blocked_f32_f16(%qk_size: index, %v_size: index, %q_count: index, %kv_count: index, %head_count: index, %kv_head_count: index, %q_s1: index, %q_s2: index, %k_s1: index, %k_s2: index, %v_s1: index, %v_s2: index, %m_s1: index, %q_extent: index, %k_extent: index, %v_extent: index, %m_extent: index) {
%one = index.constant 1 : index
%sixtyfour = index.constant 64 : index
%groups = index.mul %q_count, %head_count : index
kernel.launch.config workgroups(%groups, %one, %one) workgroup_size(%sixtyfour, %one, %one) : index
} launch(%qk_size: index, %v_size: index, %q_count: index, %kv_count: index, %head_count: index, %kv_head_count: index, %q_s1: index, %q_s2: index, %k_s1: index, %k_s2: index, %v_s1: index, %v_s2: index, %m_s1: index, %q_extent: index, %k_extent: index, %v_extent: index, %m_extent: index, %query: buffer, %key: buffer, %value: buffer, %mask: buffer, %output: buffer) {
%d = index.assume %qk_size [range(%qk_size, 1, 256)] : index
%dv = index.assume %v_size [range(%v_size, 1, 256)] : index
%nq = index.assume %q_count [range(%q_count, 1, 65536)] : index
%nkv = index.assume %kv_count [range(%kv_count, 1, 262144)] : index
%nh = index.assume %head_count [range(%head_count, 1, 1024)] : index
%nhkv = index.assume %kv_head_count [range(%kv_head_count, 1, 1024)] : index
%qe = index.assume %q_extent [range(%q_extent, 1, 268435456)] : index
%ke = index.assume %k_extent [range(%k_extent, 1, 268435456)] : index
%ve = index.assume %v_extent [range(%v_extent, 1, 268435456)] : index
%me = index.assume %m_extent [range(%m_extent, 1, 268435456)] : index
%scale = config.get @ggml.attention_strided.scale : f32
%c0 = index.constant 0 : index
%c1 = index.constant 1 : index
%c64 = index.constant 64 : index
%c2048 = index.constant 2048 : index
%zero = index.constant 0 : offset
%f0 = scalar.constant 0.0 : f32
%lowest = scalar.constant -3.40282347e+38 : f32
%q_bytes = index.constant 1024 : offset
%s_bytes = index.constant 8192 : offset
%o_bytes = index.constant 1024 : offset
%group = kernel.workgroup.id<x> : index
%lane = kernel.workitem.id<x> : index
%qi = index.div %group, %nh : index
%h = index.rem %group, %nh : index
%ratio = index.div %nh, %nhkv : index
%hkv = index.div %h, %ratio : index
%q_na, %k_na, %v_na, %m_na, %o_na = buffer.assume.noalias %query, %key, %value, %mask, %output : buffer, buffer, buffer, buffer, buffer
%q_view = buffer.view %q_na[%zero] : buffer -> view<[%qe]xf32>
%k_view = buffer.view %k_na[%zero] : buffer -> view<[%ke]xf16>
%v_view = buffer.view %v_na[%zero] : buffer -> view<[%ve]xf16>
%m_view = buffer.view %m_na[%zero] : buffer -> view<[%me]xf16>
%out_count0 = index.mul %dv, %nh : index
%out_count = index.mul %out_count0, %nq : index
%o_view = buffer.view %o_na[%zero] : buffer -> view<[%out_count]xf32>
%q_shared = buffer.alloca<workgroup> align(16) %q_bytes : buffer
%q_row = buffer.view %q_shared[%zero] : buffer -> view<256xf32>
%s_shared = buffer.alloca<workgroup> align(16) %s_bytes : buffer
%s_row = buffer.view %s_shared[%zero] : buffer -> view<2048xf32>
%o_shared = buffer.alloca<workgroup> align(16) %o_bytes : buffer
%o_row = buffer.view %o_shared[%zero] : buffer -> view<256xf32>
%qa = index.mul %qi, %q_s1 : index
%qb = index.mul %h, %q_s2 : index
%q_base = index.add %qa, %qb : index
%u0 = scf.for %t = [%lane to %d step %c64](%carry0 = %f0 : f32) -> (f32) {
%qt0 = index.add %q_base, %t : index
%qt = index.assume %qt0 [range(%qt0, 0, 268435455)] : index
%qv = view.load %q_view[%qt] : view<[%qe]xf32> -> f32
view.store %qv, %q_row[%t] : f32, view<256xf32>
scf.yield %carry0 : f32
}
// The running output row starts at zero so the first block's rescale (exp(-inf - max) = 0) multiplies a
// defined value, not whatever the workgroup allocation happened to hold.
%u1 = scf.for %dim0 = [%lane to %dv step %c64](%carry1 = %f0 : f32) -> (f32) {
view.store %f0, %o_row[%dim0] : f32, view<256xf32>
scf.yield %carry1 : f32
}
kernel.barrier<workgroup> scope(workgroup) ordering(acq_rel)
%kb = index.mul %hkv, %k_s2 : index
%m_base = index.mul %qi, %m_s1 : index
%vb = index.mul %hkv, %v_s2 : index
%oh = index.mul %h, %dv : index
%oq0 = index.mul %qi, %nh : index
%oq = index.mul %oq0, %dv : index
%o_base = index.add %oq, %oh : index
%last_key = index.sub %nkv, %c1 : index
%final_max, %final_sum = scf.for %block_start = [%c0 to %nkv step %c2048](%run_max = %lowest : f32, %run_sum = %f0 : f32) -> (f32, f32) {
// The inner loops always walk a whole 2048-slot block, so every workgroup-memory index has a constant
// bound; a key index past the end of the context is clamped to the last key and its score forced to
// %lowest, which makes its probability exactly zero in both the sum and the value pass. Clamping rather
// than masking the load is what lets the compiler prove the buffer footprint: deriving
// `block_start + slot < nkv` from `min(nkv - block_start, 2048)` is not something it can do (SUBRANGE/024
// on the key, mask and value loads), while `index.min %j0, %last_key` carries that fact directly.
%lane_max = scf.for %slot = [%lane to %c2048 step %c64](%mx = %lowest : f32) -> (f32) {
%j0 = index.add %block_start, %slot : index
%in_range = index.cmp ult, %j0, %nkv : index
%j = index.min %j0, %last_key : index
%kj = index.mul %j, %k_s1 : index
%k_base = index.add %kj, %kb : index
%dot = scf.for %t = [%c0 to %d step %c1](%sum = %f0 : f32) -> (f32) {
%kt0 = index.add %k_base, %t : index
%kt = index.assume %kt0 [range(%kt0, 0, 268435455)] : index
%qv = view.load %q_row[%t] : view<256xf32> -> f32
%kv16 = view.load %k_view[%kt] : view<[%ke]xf16> -> f16
%kv = scalar.extf %kv16 : f16 to f32
%next = scalar.fmaf %qv, %kv, %sum : f32
scf.yield %next : f32
}
%mi0 = index.add %m_base, %j : index
%mi = index.assume %mi0 [range(%mi0, 0, 268435455)] : index
%mask16 = view.load %m_view[%mi] : view<[%me]xf16> -> f16
%maskv = scalar.extf %mask16 : f16 to f32
%scaled = scalar.mulf %dot, %scale : f32
%score0 = scalar.addf %scaled, %maskv : f32
%score = scf.select %in_range, %score0, %lowest : f32
view.store %score, %s_row[%slot] : f32, view<2048xf32>
%next_max = scalar.maxnumf %mx, %score : f32
scf.yield %next_max : f32
}
%block_max = kernel.workgroup.reduce<maxnumf> %lane_max : f32
%new_max = scalar.maxnumf %run_max, %block_max : f32
%old_shift = scalar.subf %run_max, %new_max : f32
%old_rescale = scalar.expf<afn> %old_shift : f32
%lane_sum = scf.for %slot = [%lane to %c2048 step %c64](%acc = %f0 : f32) -> (f32) {
%score = view.load %s_row[%slot] : view<2048xf32> -> f32
%shifted = scalar.subf %score, %new_max : f32
%p = scalar.expf<afn> %shifted : f32
view.store %p, %s_row[%slot] : f32, view<2048xf32>
%next = scalar.addf %acc, %p : f32
scf.yield %next : f32
}
%block_sum = kernel.workgroup.reduce<addf> %lane_sum : f32
%kept_sum = scalar.mulf %run_sum, %old_rescale : f32
%new_sum = scalar.addf %kept_sum, %block_sum : f32
kernel.barrier<workgroup> scope(workgroup) ordering(acq_rel)
// rescale the running output row onto this block's maximum, then accumulate the block's values into it
%u2 = scf.for %dim = [%lane to %dv step %c64](%carry2 = %f0 : f32) -> (f32) {
%previous = view.load %o_row[%dim] : view<256xf32> -> f32
%rescaled = scalar.mulf %previous, %old_rescale : f32
%acc_v = scf.for %slot = [%c0 to %c2048 step %c1](%acc = %rescaled : f32) -> (f32) {
%j0 = index.add %block_start, %slot : index
%j = index.min %j0, %last_key : index
%p = view.load %s_row[%slot] : view<2048xf32> -> f32
%vj = index.mul %j, %v_s1 : index
%v_row = index.add %vj, %vb : index
%vi0 = index.add %v_row, %dim : index
%vi = index.assume %vi0 [range(%vi0, 0, 268435455)] : index
%v16 = view.load %v_view[%vi] : view<[%ve]xf16> -> f16
%vv = scalar.extf %v16 : f16 to f32
%next = scalar.fmaf %p, %vv, %acc : f32
scf.yield %next : f32
}
view.store %acc_v, %o_row[%dim] : f32, view<256xf32>
scf.yield %carry2 : f32
}
kernel.barrier<workgroup> scope(workgroup) ordering(acq_rel)
scf.yield %new_max, %new_sum : f32, f32
}
%u3 = scf.for %c = [%lane to %dv step %c64](%carry3 = %f0 : f32) -> (f32) {
%acc_o = view.load %o_row[%c] : view<256xf32> -> f32
%result = scalar.divf %acc_o, %final_sum : f32
%oi0 = index.add %o_base, %c : index
%oi = index.assume %oi0 [range(%oi0, 0, 268435455)] : index
view.store %result, %o_view[%oi] : f32, view<[%out_count]xf32>
scf.yield %carry3 : f32
}
kernel.return
}

// Rotate-half RoPE as graph compilers lower it (x * cos + concat(-x[half:], x[:half]) * sin, eight
// nodes), in one pass: x packed [d, T, H], cos and sin [d, T] (row strides given), broadcast over
// heads. One workitem per element.
Expand Down
8 changes: 8 additions & 0 deletions tests/test-backend-ops.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9754,6 +9754,14 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
}
}

// Cases past one workgroup score row (kv > 2048), the first time the suite reaches the blocked rows
// attention kernel; hsv a multiple of 32 is what lets common.attention_strided_f32_f16 take them at all.
for (int kv : { 2560, 4096, }) {
for (int hs : { 64, 160, 256, }) {
test_cases.emplace_back(new test_flash_attn_ext(hs, hs, 8, {4, 1}, kv, 32, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16));
}
}

for (int hsk : { 40, 64, 72, 80, 96, 128, 192, 256, 320, 512, 576 }) {
for (int hsv : { 40, 64, 72, 80, 96, 128, 192, 256, 512 }) {
if (hsk != 192 && hsk != 320 && hsk != 576 && hsk != hsv) continue;
Expand Down