feat(deepseek-v4): implement HiSA hierarchical sparse attention

[inureyes]

Summary

DeepSeek-V4 attention runs over two KV streams per layer: a sliding-window local stream at full token resolution, and (on layers with a nonzero compress ratio) a pooled stream in which a compressor summarizes every r consecutive tokens into one pooled token (r is the layer's entry in compress_ratios, either 4 or 128). On the r == 4 layers the pooled prefix grows without bound, so the model does not attend to all of it: a per-layer indexer selects the index_topk most relevant pooled positions per query, and attention runs over the local window plus only the selected pooled tokens. HiSA (hierarchical indexing for sparse attention) is the selection algorithm: a coarse block-level filter followed by a fine within-block top-k, replacing a flat O(Np) scan over all pooled positions.

This issue covers the indexer end to end: the indexer module and its weights, the flat top-k baseline, both HiSA hierarchical paths (single-token decode and batched prefill), and the sparse attention combine that consumes the selected indices.

Dependency

Blocked by #523 (base DeepSeek-V4 backbone). The plan below targets the module layout #523 will introduce: once #523 lands, the model lives in src/models/deepseek_v4.rs with the following pieces this issue builds on:

• The Compressor (gated pooling + RMSNorm + RoPE) and its pooled-stream cache. The indexer owns a private compressor instance, so the module must be reusable with a caller-chosen output head dim.
• The per-layer attention dispatch on compress_ratios (0: local only, 128: local + dense pooled, 4: local + sparse pooled).
• The dense local+pooled combine with attention sinks, which the sparse combine below generalizes.

Until this issue is done, r == 4 layers should fall back to dense attention over all pooled tokens. That fallback is numerically exact while the pooled prefix has at most index_topk entries and an approximation beyond that (the model is specified to attend sparsely there).

Current state in mlxcel

No DeepSeek-V4 module exists yet (#523 is open). The in-tree prior art is the DeepSeek-V3.2 sparse-attention indexer (added for #509):

• src/models/deepseek_v32_indexer.rs: a per-layer indexer that scores every cached indexer key against the current query with relu(q @ k^T), combines heads with a learned per-head weight weights_proj(x) * n_heads^-0.5 * head_dim^-0.5, applies the additive causal mask, and takes a flat argpartition top-index_topk. Indexer::load (line 66) returns Ok(None) when the checkpoint has no indexer weights, and DENSE_FALLBACK_ENV (MLXCEL_DSA_DENSE, line 41) force-disables it for A/B runs. The scoring core is the pure function indexer_top_indices (line 199), factored out for unit testing.
• src/models/deepseek_v32.rs: consumption patterns. Decode (L == 1) gathers the selected KV rows directly with take_along_axis (gather_topk_kv, line 497) instead of building a mask; prefill restricts scores with a sparse additive mask (apply_sparse_prefill_mask). Short-context parity: selection returns None when kv_len <= index_topk and the dense path runs unchanged.
• src/models/deepseek_v32_tests.rs: unit-test structure (synthetic weights, Indexer::load round trip, selector assertions, short-context None behavior).
• src/models/glm_moe_dsa.rs: a second consumer of the same indexer, showing how the module is shared across architectures.

How the V4 indexer differs from that one:

1. It scores pooled compressed keys, not per-token keys. The keys come from the indexer's own compressor (its own pooled cache), so there is no wk/k_norm per-token key projection; there are compressor.{wkv,wgate,ape,norm} weights instead.
2. Selection is hierarchical (block means first, then within-block token scores), not one flat top-k, and it has distinct decode and batched formulations.
3. The selected indices address pooled positions, and the consuming attention combines a sliding-window local stream with the gathered pooled rows under a single split softmax with per-head sinks, instead of feeding one homogeneous KV stream.

What carries over: the relu-scored per-head weighted combine, argpartition top-k, take_along_axis gather, the Ok(None)-when-absent load contract, the env kill-switch pattern, and the "pure selector function + unit tests" factoring. All required ops already exist in mlxcel-core (argpartition, take_along_axis, relu, logsumexp_axis, logaddexp, minimum, maximum, clip), so no FFI additions are expected.

Algorithm

Conventions

• All tensors are written with explicit axis order; axis 0 is batch. Activations are (B, L, hidden), attention tensors (B, H, L, D) with heads on axis 1.
• Linear weights are stored (out_features, in_features); linear(x) = x @ W^T.
• t0 is the cache offset: the number of tokens already consumed by this layer's local cache before the current step. Query row j in [0, L) has absolute token position t = t0 + j (0-based).
• All selection scoring is computed in f32 regardless of model dtype.
• argpartition top-k returns an unordered index set. Attention is order-invariant, so this is correct, but tests comparing selections must sort both sides.

Configuration

From the checkpoint config.json (values for DeepSeek-V4-Flash):

key	value	meaning
compress_ratios	list, entries in {0, 4, 128}	per-layer ratio; the checkpoint carries 44 entries for 43 layers, use the first num_hidden_layers entries
index_topk	512	pooled positions kept per query (k)
index_n_heads	64	indexer heads (H_i)
index_head_dim	128	indexer head dim (D_i)
sliding_window	128	local stream window
q_lora_rank	1024	rank of the shared query residual
compress_rope_theta	160000.0	RoPE base for pooled-stream layers
qk_rope_head_dim	64	rotated dims per head

Two keys are architectural defaults not present in the checkpoint config and must get serde defaults:

key	default	meaning
index_block	64	HiSA block size b, in pooled rows; 0 disables the hierarchy (flat scan only)
index_keep	16	HiSA blocks kept by the coarse stage (Kb upper bound)

Only r == 4 layers have an indexer. Checkpoint weight keys and logical shapes (the 4-bit checkpoint at /home/inureyes/models/deepseek-v4-flash-4bit stores linear weights quantized, U32-packed with scales; shapes below are dequantized logical shapes):

model.layers.{i}.attn.indexer.wq_b.weight            (8192, 1024)  = (H_i * D_i, q_lora_rank)
model.layers.{i}.attn.indexer.weights_proj.weight    (64, 4096)    = (H_i, hidden_size)
model.layers.{i}.attn.indexer.compressor.wkv.weight  (256, 4096)   = (2 * D_i, hidden_size)
model.layers.{i}.attn.indexer.compressor.wgate.weight(256, 4096)   = (2 * D_i, hidden_size)
model.layers.{i}.attn.indexer.compressor.ape         (4, 256)      = (r, 2 * D_i)
model.layers.{i}.attn.indexer.compressor.norm.weight (128,)        = (D_i,)

The 2 * D_i output width is the overlap compressor variant used at r == 4 (each pooled row also draws on the previous window; r == 128 layers use the plain variant with output width equal to the head dim, see #523).

The pooled stream and positions

• The indexer's compressor turns the layer input x (B, L, hidden) into pooled indexer keys pooled (B, Np, D_i), appended to its own pooled cache. Pooled row p (axis 1, 0-based) is anchored at the token window [p*r, p*r + r - 1]; the overlap variant additionally draws on tokens [(p-1)*r, p*r - 1].
• Tokens that do not yet fill a complete window sit in the compressor's remainder buffer and produce no pooled row (feat(models): port DeepSeek-V4 (MLA + group-limited MoE backbone) #523 owns this cache mechanics). The pooled cache advances only in whole windows.
• RoPE: pooled row p is rotated as absolute token position p * r; indexer queries are rotated at their absolute token positions t. Both use base compress_rope_theta = 160000 and rotate only a 64-dim subspace of the 128-dim head: with interleaved pairing (pair m = dims (2m, 2m+1), m in [0, 64)), pairs 0..31 (dims 0..63) are position-independent and pairs 32..63 (dims 64..127) rotate.
• Causal visibility: the query at absolute position t may see pooled row p iff p < floor((t + 1) / r). Because visibility is a contiguous prefix of the pooled axis, it is fully described by a count: V(t) = min(Np, floor((t + 1) / r)). During single-token decode every cached pooled row satisfies this (the newest window closes at the current token), so decode needs no pooled mask.

The indexer query and per-head weights

Inputs shared with the main attention path: q_residual (B, L, 1024) = rms_norm(wq_a(x)) (the main attention's query LoRA bottleneck; the indexer must receive this tensor, not recompute it).

q  = wq_b(q_residual)                       # (B, L, H_i * D_i)
q  = reshape(q, (B, L, H_i, D_i))
q  = transpose(q, (0, 2, 1, 3))             # (B, H_i, L, D_i)
q  = rope(q, offset = t0)                   # per row j, position t0 + j
w  = weights_proj(x) * H_i^(-1/2)           # (B, L, H_i), f32

Let scale = D_i^(-1/2). Every path below combines heads the same way: per-head relu score, times scale, times w[.., j, h], summed over the head axis.

Flat selection (baseline path)

Used when the hierarchy is disabled (index_block == 0) or the prefix is too short for it (thresholds below). With k = min(index_topk, Np):

s[b,h,j,p]   = relu( sum_d q[b,h,j,d] * pooled[b,p,d] )        # (B, H_i, L, Np)
score[b,j,p] = sum_h s[b,h,j,p] * scale * w[b,j,h]             # (B, L, Np)
score[b,j,p] = NEG if p >= V(t0 + j)                           # causal, skip when L == 1
I            = argpartition_desc(score, k, axis=-1)[..., :k]   # (B, L, k), unordered

NEG is a large negative finite constant (dtype minimum or -1e30), never -inf: scores are relu-nonnegative, so any finite very-negative value sorts strictly last and argpartition never sees NaN.

HiSA decode path (L == 1)

Gate conditions, all required: L == 1, the pooled visibility mask is absent (always true for L == 1), index_block > 0, Np >= index_block * index_keep, and index_block * index_keep >= k (enough fine candidates to fill the top-k).

Let b = index_block, nb = floor(Np / b), Kb = min(index_keep, nb), C = Kb * b. Block beta in [0, nb) covers pooled rows [beta*b, (beta+1)*b - 1] on the pooled axis. The tail rows [nb*b, Np) are not candidates in this path; this is intended (they summarize the newest tokens, which the local window already covers at full resolution).

# coarse: score block-mean representatives
rep[b_,beta,d]  = mean over rows p in block beta of pooled[b_,p,d]        # (B, nb, D_i)
cs[b_,h,0,beta] = relu( sum_d q[b_,h,0,d] * rep[b_,beta,d] )              # (B, H_i, 1, nb)
cscore[b_,0,beta] = sum_h cs[b_,h,0,beta] * scale * w[b_,0,h]             # (B, 1, nb)
top_blk = argpartition_desc(cscore, Kb, axis=-1)[..., :Kb]                # (B, 1, Kb)

# fine: expand retained blocks to a candidate table and rescore exactly
pos[b_,0,c] = top_blk[b_,0, c // b] * b + (c % b)      for c in [0, C)    # (B, 1, C) absolute pooled row ids
cand[b_,c,d] = pooled[b_, pos[b_,0,c], d]                                 # (B, C, D_i) via take_along_axis on axis 1
fs[b_,h,0,c] = relu( sum_d q[b_,h,0,d] * cand[b_,c,d] )                   # (B, H_i, 1, C)
fscore[b_,0,c] = sum_h fs[b_,h,0,c] * scale * w[b_,0,h]                   # (B, 1, C)
sel = argpartition_desc(fscore, k, axis=-1)[..., :k]                      # (B, 1, k) indices into the candidate axis
I[b_,0,kk] = pos[b_,0, sel[b_,0,kk]]                                      # (B, 1, k) absolute pooled row ids

No causal masking anywhere in this path: during decode all Np pooled rows are visible. Cost is O(nb + Kb*b) scores per step instead of O(Np).

HiSA batched path (L > 1, prefill and multi-token verify)

Gate conditions: L > 1, index_block > 0, Np >= index_block * index_keep, index_block * index_keep >= k. This path handles causality explicitly, so it runs with or without a pending visibility mask.

valid_len[b_,j] = V(t0 + j) = min(Np, floor((t0 + j + 1) / r))            # (B, L) i32, visible-prefix length per query row

# coarse (same rep/cscore math as decode, but per query row)
cscore[b_,j,beta]                                                          # (B, L, nb)
cscore[b_,j,beta] = -1e30  if beta * b >= valid_len[b_,j]                  # block eligible iff its FIRST row is visible
top_blk = argpartition_desc(cscore, Kb, axis=-1)[..., :Kb]                 # (B, L, Kb), i32

# fine, tiled over the query axis in chunks of fine_chunk = 512 rows
# (bounds the (B, chunk, C, D_i) candidate gather; memory stays flat in L)
for each chunk [s, e):
    pos_c[b_,j,c] = top_blk[b_,s+j, c // b] * b + (c % b)                  # (B, e-s, C)
    cand[b_,j,c,d] = pooled[b_, pos_c[b_,j,c], d]                          # (B, e-s, C, D_i), gather on the pooled axis
    q_bl = transpose(q[:, :, s:e, :], (0, 2, 1, 3))                        # (B, e-s, H_i, D_i): batch the matmul over (B, chunk)
    fs[b_,j,h,c] = relu( sum_d q_bl[b_,j,h,d] * cand[b_,j,c,d] )           # (B, e-s, H_i, C)
    oc[b_,j,c] = sum_h fs[b_,j,h,c] * scale * w[b_,s+j,h]                  # (B, e-s, C)
    oc[b_,j,c] = -1e30  if pos_c[b_,j,c] >= valid_len[b_,s+j]              # per-row causal mask (handles partially visible blocks)

fscore = concat(chunks, axis=1)                                            # (B, L, C)
sel    = argpartition_desc(fscore, k, axis=-1)[..., :k]                    # (B, L, k)
I[b_,j,kk] = pos[b_,j, sel[b_,j,kk]]                                       # (B, L, k) absolute pooled row ids

Coarse eligibility uses only the block's first row; rows of a partially visible block are re-masked in the fine stage. Early query rows can have fewer than k visible rows; the top-k then contains masked ids, which is safe because the consuming attention re-masks them (below) so they get zero weight.

Consuming the selection: sparse combine

The attention branch on an r == 4 layer is three-way on the attention-side pooled length Np (the attention compressor and the indexer compressor pool with the same r, so their pooled axes are position-aligned):

1. Np == 0: local-only attention (as feat(models): port DeepSeek-V4 (MLA + group-limited MoE backbone) #523).
2. 0 < Np <= index_topk: dense over local + all pooled rows (as feat(models): port DeepSeek-V4 (MLA + group-limited MoE backbone) #523); byte-identical to what sparse selection would produce, since every row would be selected.
3. Np > index_topk: sparse combine with I (B, L, k) from the indexer, k = index_topk:

# streams (Dh = 512 = head_dim, W = current local KV length, K = V in both streams)
q         (B, 64, L, Dh), already roped; q_s = q * Dh^(-1/2)
local_kv  (B, 1, W, Dh)                 # single shared KV head, broadcast over the 64 query heads
pooled_kv (B, Np, Dh)                   # the ATTENTION compressor's pooled stream, not the indexer's

# gather per-query pooled rows
P[b_,j,kk,d] = pooled_kv[b_, I[b_,j,kk], d]                                # (B, L, k, Dh)

local_scores  = q_s @ local_kv^T                                           # (B, 64, L, W), + sliding-window causal mask
pooled_scores = transpose(transpose(q_s,(0,2,1,3)) @ P^T, (0,2,1,3))       # (B, 64, L, k)
pooled_scores[b_,h,j,kk] = -inf  if I[b_,j,kk] >= V(t0 + j)                # re-mask gathered invisible rows (prefill only)

n = logaddexp(logsumexp(local_scores, -1), logsumexp(pooled_scores, -1))   # (B, 64, L, 1)
n = logaddexp(n, sink[h])                                                  # per-head attention sink joins the normalizer
w_local  = exp(local_scores  - n)                                          # sink absorbs mass, contributes no value
w_pooled = exp(pooled_scores - n)
out[b_,h,j,:] = w_local @ local_kv + w_pooled[b_,h,j,:] @ P[b_,j,:,:]      # (B, 64, L, Dh)

This is one softmax over the union of both streams plus the sink, computed via logaddexp so the two score tensors never have to be concatenated per query.

Engagement thresholds (defaults, r = 4)

pooled rows Np	tokens (approx.)	path
Np <= 512 (index_topk)	<= ~2048	dense over pooled; selection output unused
512 < Np < 1024 (index_block * index_keep)	~2048 to ~4096	flat scan top-512
Np >= 1024	>= ~4096	HiSA hierarchy (decode or batched)

Operational requirement: the indexer's compressor must run and advance its pooled cache on every forward of an r == 4 layer, even while the dense branch is taken, otherwise the pooled keys do not exist when Np crosses index_topk. Only the scoring/selection may be skipped below the threshold.

Implementation plan

1. src/models/deepseek_v4_indexer.rs (new)
   • Indexer struct: wq_b, weights_proj (UnifiedLinear), a private Compressor instance (from feat(models): port DeepSeek-V4 (MLA + group-limited MoE backbone) #523, output head dim index_head_dim), and n_heads, head_dim, index_topk, index_block, index_keep, scale.
   • Indexer::load(weights, args, attn_prefix) reading {attn_prefix}.indexer.*; return Ok(None) when the keys are absent, mirroring src/models/deepseek_v32_indexer.rs:66.
   • Pure selector functions with explicit shapes, factored for unit tests (the indexer_top_indices pattern at src/models/deepseek_v32_indexer.rs:199): flat_top_indices, hisa_select_decode, hisa_select_batched.
   • Env kill-switches following the MLXCEL_DSA_DENSE precedent: one that disables the indexer entirely (dense fallback) and one that forces the flat scan (A/B against the hierarchy).
2. src/models/deepseek_v4.rs (once feat(models): port DeepSeek-V4 (MLA + group-limited MoE backbone) #523 lands)
   • Give r == 4 layers a third cache slot for the indexer's pooled stream, next to the local cache and the attention pooled cache.
   • Wire the three-way branch (local-only / dense pooled / sparse) and the sparse combine; keep branch 2 sharing feat(models): port DeepSeek-V4 (MLA + group-limited MoE backbone) #523's dense code so parity is structural.
   • Decode gather via take_along_axis (prior art gather_topk_kv, src/models/deepseek_v32.rs:497).
3. Config (ModelArgs for deepseek_v4): add index_block (serde default 64) and index_keep (serde default 16); confirm index_topk, index_n_heads, index_head_dim are parsed (they are present in the checkpoint config).
4. Tests: src/models/deepseek_v4_indexer_tests.rs (or #[cfg(test)] in the module), patterned on src/models/deepseek_v32_tests.rs: synthetic f32 weights, load round trip, selector assertions.

Validation and acceptance criteria

Selector unit tests (pure functions, synthetic tensors):

• Flat top-k matches a brute-force sorted argsort selection (as sets) on random scores, with and without a visibility cutoff.
• Hierarchy equivalence: plant score mass so the true top-k rows all lie inside the top-Kb blocks; assert hisa_select_decode and hisa_select_batched return the same index set as the flat scan (sort both sides).
• Causality: for the batched path with a synthetic valid_len, every selected index is < valid_len[j] whenever at least k rows are visible for row j.
• Tail exclusion: rows in [nb*b, Np) never appear in hierarchical selections.
• Gate behavior: Np < index_block * index_keep or index_block == 0 falls back to the flat scan; Np <= index_topk on the attention side keeps the dense branch.

Parity vs a full-attention oracle at short context:

• With Np <= index_topk, assert the sparse-capable layer produces bit-identical output to feat(models): port DeepSeek-V4 (MLA + group-limited MoE backbone) #523's dense pooled path (branch selection, structural).
• With a small synthetic config where k == Np (every row selected), assert the sparse combine (gather + split softmax + sink) reproduces dense local+pooled attention within fp tolerance. This is the end-to-end check of the logaddexp algebra.

Real-checkpoint smoke (/home/inureyes/models/deepseek-v4-flash-4bit):

• Short prompt: greedy output identical to the feat(models): port DeepSeek-V4 (MLA + group-limited MoE backbone) #523 baseline (sparse code present but not engaged).
• Long prompt (>= ~4300 tokens so Np >= 1024 at r == 4): prefill (batched path) and decode (fast path) both engage, generation is coherent, no NaN/inf in logits.

Long-context sanity:

• Retrieval probe: a fact planted early in a filler prompt long enough to cross the sparse threshold is still answered correctly.
• Decode step time stays roughly flat as context grows (selection cost is O(nb + Kb*b), not O(Np)).

Out of scope

• Everything feat(models): port DeepSeek-V4 (MLA + group-limited MoE backbone) #523 owns: HyperConnections, the compressor and pooled caches, hash-routed MoE, the grouped output projection, sliding-window local attention, weight loading/sanitization.
• The shared-expert SwiGLU clamp (feat(deepseek-v4): clamp shared-expert SwiGLU to swiglu_limit #550).
• Multi-token-prediction / speculative-decode integration and pooled-cache snapshot/rollback.
• Fused GPU kernels for selection or gather; the first implementation composes existing mlxcel-core ops.
• Distributed sharding of the indexer.