Conceptio › Archive › arXiv CS

arXiv CSopen access

HieraSparse: Hierarchical Semi-Structured Sparse KV Attention

Unknown · 2026 · arxiv_cs

arXiv CS · Papers · License: Open Access · 2026

Open Source ↗Direct PDF ↓

cloud distributedcomputing parallelcomputing

distributed computing, parallel computing, cloud

arXiv:2604.16864v1 [cs.DC] 18 Apr 2026

HieraSparse: Hierarchical Semi-Structured Sparse KV Attention Haoxuan Wang

Chen Wang

College of Computing and Data Science Nanyang Technological University Singapore, Singapore [email protected]

Abstract—The deployment of long-context Large Language Models (LLMs) poses significant challenges due to the intense computational cost of self-attention and the substantial memory overhead of the Key-Value Cache (KV Cache). In this paper, we introduce HieraSparse, a hierarchical KV Cache compression framework with acceleration kernels that leverage GPU sparse tensor cores to speed up semi-structured KV Cache attention for both the prefill and decode phases. With the hierarchical design, our method allows for a flexible quality-sparsity tradeoff and successfully converts sparsity into efficiency. Compared to the state-of-the-art decode method that utilizes unstructured sparsity, HieraSparse achieves 1.2× KV compression ratio and 4.57× attention speedup at the same sparsity level. Furthermore, we extended the semi-structured KV Cache pruning to the prefill stage, which demonstrated up to 1.85× attention speedup at the highest sparsity. Lastly, we evaluate the generation quality of HieraSparse with a simple magnitude-based pruning method, and the results show that 1.37× prefill speedup and 1.77× decode speedup can be achieved without significant quality drop. The codebase can be found at https://github.com/psl-ntu/HieraSparse. Index Terms—Sparse GEMM, Sparse Attention, KV Cache, Kernel Optimization, Semi-structured Sparsity

I. I NTRODUCTION Large language models are constantly evolving to support longer context windows, showing dominating capabilities in long-context natural language processing [1]–[4] and downstream tasks like code generation [5], text summarization [6], and logical reasoning [7]. Concurrently, application-level demand for long-context processing continues to grow. Emerging paradigms such as Retrieval-Augmented Generation (RAG) [8], In-Context Learning [9], and memory-augmented agentic systems [10] require models to process vast amounts of retrieved documents, extensive demonstrations, or long-term interaction histories. These demands are pushing context window requirements to hundreds of thousands or even millions of tokens [11]–[13], imposing significant pressure on both computation and memory. Computation. This massive increase in sequence length exposes the quadratic complexity O(n2 ) of the self-attention mechanism [14] as a critical bottleneck. As illustrated in Figure 1, attention computation can take 50% of prefill latency at the context length of 64K and continue growing to 80% at 192K, accounting for even minutes of time-to-first-token

Fig. 1: The latency breakdown of the prefill and decode phases under different context lengths. The attention mechanism gradually dominates the computation as the context length increases.

(TTFT) latency. The decode phase is similarly affected: as the model must attend over all preceding tokens and their associated KV Cache entries, attention can consume more than 60% of time-per-output-token (TPOT) latency under longcontext settings. These observations motivate the development of efficient attention mechanisms capable of scaling to long contexts. Memory. Beyond computation, the KV Cache introduces a memory footprint that scales linearly with sequence length. As a concrete example, serving Llama-3.1-8B-Instruct [2] at its maximum supported context length requires 16 GiB of KV Cache memory per request, already on par with the model weights themselves. With context extension techniques [11]– [13], context lengths can readily scale to over one million tokens (1048K), demanding upwards of 125 GiB of KV Cache memory per request, far exceeding the capacity of most modern GPUs. This underscores the pressing need for methods that reduce the memory footprint of KV Cache during LLM inference. Prior work [15]–[22], [24]–[27] has shown that the KV Cache includes many low-importance entries that can be pruned without significantly degrading generation quality. As summarized in Table I, such pruning techniques operate at varying levels of granularity. Coarse-grained methods prune at the level of tokens, channels, heads, or layers, achieving

TABLE I: KV Cache pruning scheme comparison. Speedup Mechanism

Target

Example Method

Coarse-grained

Token Channel Head Layer Head+Token

H2O [15], SnapKV [16] ThinK [17], LeanK [18] HeadKV [19], AdaKV [20] PyramidKV [21], DynamicKV [22] DuoAttention [23]

Low

Computation Skip

Fine-grained unstructured

Element

MUSTAFAR [24]

High

Yes

Hierarchical semi-structured

Block+Element

HieraSparse (Ours)

High

Yes

substantial speedups by skipping computation entirely, though at the cost of reduced flexibility and coarse control over the sparsity-quality tradeoff. Fine-grained, unstructured elementwise pruning, by contrast, offers precise control over which matrix entries are removed, but it fails to translate sparsity into practical efficiency gains. This is because existing sparse execution frameworks [28]–[30] adopt a load-as-sparse, computeas-dense scheme: while memory bandwidth is reduced, the compute cost remains unchanged, rendering such approaches ineffective for compute-bound operations such as large-scale GEMM and prefill attention [31], [32]. Furthermore, the irregular structure inherent in unstructured pruning introduces non-trivial compression overhead: approximately 12% additional latency during prefill in our experiments. Overall, these constraints create a fundamental gap between achieved sparsity and realized efficiency in existing systems. N:M semi-structured sparsity has emerged as a promising middle ground, where exactly N out of every M consecutive elements are non-zero. Modern GPU architectures (e.g., NVIDIA Ampere, Hopper, and AMD MI300X) provide dedicated sparse tensor cores that can accelerate N:M sparse computations, effectively doubling throughput [33]– [35]. Given the limitations of unstructured-sparsity systems outlined above, applying semi-structured sparsity to the KV Cache is a natural and compelling direction. However, doing so in practice introduces several non-trivial system challenges. First, to preserve generation quality, sparsity must be applied selectively: critical attention regions—such as attention sinks, heavy hitters, and local windows [15], [25], [36]—must remain dense to retain complete information. This necessitates a hierarchical sparsity design capable of efficiently mixing dense and sparse regions. Second, since sparse tensor cores only compress the first matrix operand, realizing hardware acceleration across both the prefill and decode phases requires a redesign of the attention computation. Third, performing online KV Cache compression risks a compression latency tax, where format conversion overhead negates the gains of sparse kernels. Avoiding this requires a highly efficient compressor that can integrate seamlessly into the prefill phase without inflating TTFT. To address these challenges, we present HieraSparse, a framework that enables hardware-accelerated semi-structured sparsity for KV Cache compression. To the best of our knowledge, HieraSparse is the first system to leverage GPU

Flexibility

Algorithm Agnostic

Granularity

Speedup

Supported Phase Prefill

Decode

High

× × × × ✓

✓ ✓ ✓ ✓ ✓

Load-As-Sparse Compute-As-Dense

Low

✓

Sparse Tensor Core

High

✓

sparse tensor cores to accelerate attention computation, and the first to extend semi-structured sparse KV cache compression to the prefill phase. Our key contributions are: • We propose a hierarchical block-based memory management scheme that supports mixing of dense and sparse KV Cache blocks, which allows a flexible qualityefficiency trade-off. • We design highly optimized sparse attention kernels that integrate sparse tensor cores, accelerating both prefill and decode phases with N:M structured sparsity. We also performed a detailed analysis of theoretical and actual speedup under different sparsity settings. • We implement highly efficient compression kernels that enable near-zero-overhead online KV Cache sparsification. II. BACKGROUND AND R ELATED W ORKS Attention The core component of the transformer is the scaled dot-product attention mechanism [14]. Given three matrices: Query Q ∈ Rn×d , Key K ∈ Rn×d , and Value V ∈ Rn×d , where n is the number of tokens and d is the hidden dimension, the output O ∈ Rn×d is computed as: P

}| {   S z }| {  T  Q × K  O = softmax √  ×V  d  z

(1)

where × denotes a general matrix multiplication (GEMM) operation, P and S denote the intermediate results. Naive attention implementations require O(n2 ) space complexity to materialize the intermediate attention score matrix P . To reduce this memory footprint, FlashAttention [37], [38] fuses two GEMMs and softmax operations into a single kernel by tiling the computation. This approach maintains row-wise statistics Li (log-sum-exp) and Mi (max) for the online safesoftmax algorithm, avoiding the need to store the full attention matrix. Additionally, the underlying GEMM operations leverage dedicated hardware acceleration, such as tensor core or matrix core, for maximum throughput. PagedAttention [39] as a memory-efficient attention variant designed for model inference, solved memory fragmentation by partitioning the key and value tensors into smaller blocks that can be loaded and processed individually.

the nonzero elements of matrix A, together with additional metadata E that encodes their positions. Importantly, both dense and sparse instructions ideally require the same number of hardware cycles. As a result, sparse tensor cores can achieve up to 2× computational throughput while reducing memory traffic for matrix A to approximately 75% of the dense case. III. M ETHOD Figure 3 demonstrates the overall workflow of HieraSparse. Given the KV Cache that is divided into sparse and dense regions, the caches are further split into blocks. For dense blocks, they are directly stored in the dense cache memory pool; for sparse blocks, they are further pruned and compressed into non-zero data and metadata, then stored in the respective memory pools. A block index mapping is created accordingly, representing the block-level sparsity pattern of KV Cache and tracking the offset of respective blocks in the memory pool. During attention computation, the block index mapping and the memory pools are fed to attention kernels, which utilize the sparse tensor core for acceleration. After the prefill phase, the dense cache can be further pruned and compressed, which allows different sparsity settings for the prefill and decode phases.

Map

Sequence Length (Split by Block Size) +

Key Value

Offset

d ea

Block Indices

Dim H

KV Cache Optimization There have been various works that optimize the KV Cache from different perspectives. For example, KV Cache is often offloaded to CPU memory or NVMe device for future reuse when multiple requests share the same prefix (i.e., multi-turn conversation, common document/knowledge question-answering (QA), in-context learning), which has been widely adopted in popular inference engines like vLLM and SGLang [39], [40]. Frontier works like CacheBlend [41], PromptCache [42], and CacheLink [43] also try to reuse KV Cache at a finer-grained scope, especially when the sharing-prefix condition can not be met. Another batch of works, including different methods of KV Cache quantization and pruning [44], [45], tried to optimize the KV Cache by utilizing its internal numeric characteristics. As shown in Figure 2, some channels in keys consistently exhibit large magnitude across all tokens, while value cache tends to have more uniform small magnitude without a distinguishable pattern. ThinK [17] leveraged this observation to remove trivial channels in both query and value matrices, reducing the overall computation and memory bandwidth. MUSTAFAR [24] further explored unstructured KV Cache pruning with a magnitudebased algorithm and performed various experiments to analyze the effect of different pruning dimensions, showing that key cache is suitable for per-token pruning due to its outlier channels, while per-token and per-channel pruning show minimal difference for value cache. ZipKV [46] proposed a new area where compressed KV Cache can be shared among different prompts that have the same prefixes, which can be done offline and introduces zero runtime overhead.

Prune and Compress

N:M Group

Store Memory Pool

Load

Dense K/V

Prefill Attention Kernel

(a) 19th layer key cache.

(b) 19th layer value cache.

Fig. 2: Visualizations of the key and value cache absolute value from Llama-3.1-8B-Instruct using text sliced from LongBench. The key cache shows outlier ridges in certain channels and is consistent across all tokens, while the value cache appears more random without distinguishable patterns. Sparse Tensor Core Since the introduction of the NVIDIA Ampere architecture (and AMD MI300X matrix cores), sparse tensor cores have been designed to accelerate semi-structured GEMM operations of the form D = A × B + C, where matrix A exhibits N:M structured sparsity along the reduction dimension. For example, alongside the dense tensor core instruction mma.m16n8k16, there exists a corresponding sparse instruction, mma.sp.m16n8k32, which operates on 2:4 sparse matrices. The sparse instruction processes only

(SK,SV)

Sparse Key (S'K,S'V) Optionally Prune

Sparse Value Decode Attention Kernel

Fig. 3: The overall workflow of HieraSparse, which supports flexibility at different levels: i) Different sparsity patterns for the prefill and decode phases. ii) Different sparsity patterns for the key and value caches. iii) Different block-level sparsity pattern. iv) Different element-level sparsity pattern. There are three main function components in HieraSparse: i) Hierarchical cache pruner. ii) Bundled compression kernels that take pruned tensors and output respective non-zero values and metadata. iii) GPU acceleration kernels for both prefill and decode phases. Each of them will be introduced in the following subsections. A. Hierarchical Cache Pruner The cache pruner takes a pruning algorithm that produces multi-level masks: a block-level pruning mask M and an element-level pruning mask m to generate the compressed caches. The block-level mask M determines which blocks are kept dense or pruned to sparse, while the element-level mask

m determines which elements within the sparse blocks are pruned to zero. HieraSparse is designed to be agnostic to the pruning algorithms. For example, the system can be configured with masks produced by ThinK or LeanK for key cache channel pruning when the mask satisfies or can be permuted to a channelwise N:M pattern. For simplicity and comparability with unstructured methods, we adopt a straightforward magnitudebased pruning method to generate the block and element-level masks by default. We first divide the KV Cache into blocks along the sequence dimension by a size of B, each block is denoted by Ki and Vi . Within each block, we select every N out of M elements with the lowest magnitude, forming the element-level mask m. For each block, we compute the magnitude loss sum Li incurred by applying the element-level mask and sort the blocks based on Li . Finally, we select a portion S of blocks with the lowest Li to prune to sparse, while the rest are kept dense. Formally, the pruning process can be described as Equations 2: Ti = topN :M (|Xi |) mXi = I(|Xi | ≥ Ti ) LXi = ∥Xi ⊙ (1 − mXi )∥1 MX = I(LXi ≥ TSX ), TS = topSX ({LX })

(2a) (2b) (2c) (2d)

Where X ∈ {K, V }, I(·) is the indicator function, topN :M (·) returns the threshold value to keep the top N elements out of every M elements. Finally, topS (·) determines the loss threshold to maintain the target block sparsity SX . B. Cache Compressor Given the block-level mask M and element-level mask m produced by the hierarchical cache pruner, the cache compressor transforms the original dense KV cache into mixed KV Cache blocks suitable for the acceleration kernels. Following memory management strategies from PageAttention, we organize dense data, nonzero data, and metadata into separate memory pools, and maintain a block-level index map to track them: we first allocate memory pools for dense blocks, nonzero data blocks, and metadata blocks, along with a block index map. Two counters are initialized to track the number of dense and sparse blocks, respectively. For each block in the original KV Cache, we consult the block-level mask M . If a block is marked as dense, it is directly copied into the dense memory pool, and the dense counter is incremented. Otherwise, the compressor applies the element-level mask m to extract nonzero elements and their corresponding metadata, which are stored in the sparse data and metadata pools, respectively, and the sparse counter is incremented. The block index map encodes both the type and location of each block using a sign convention: a positive index denotes a dense block and points to its offset in the dense pool, while a negative index denotes a sparse block and points to its offset in the sparse pool. During attention computation, kernels can efficiently dispatch to the appropriate representation by consulting this index map. It is worth noting that compression itself incurs

a latency cost that may partially offset the throughput gains from the acceleration kernels. We address this in Section IV-B, where we describe an efficient implementation that reduces compression overhead to near zero. C. GPU Accleration Kernels 1) Fully Semi-Structured Sparse Computation Workflow: We first describe the design that computes on sparse cache blocks, and later extend it to support mixed dense and sparse blocks. Since the N:M sparse tensor core requires the first input matrix to be semi-structured sparse along the reduction dimension, careful kernel design is required to determine which matrices should be compressed. Table II compares potential strategies based on transposing the two consequent GEMM operations to switch the sparse operands, highlighting their algorithmic support and ideal speedup. TABLE II: Design space exploration to apply sparse tensor core in attention. (·) is to denote the · has to be same as the shape of GEMM1 output. Config Naive Trans-K Trans-V Trans-Both

GEMM1 S = Q × KT T S = K × QT S = Q × KT ST = K × QT

GEMM2 O =P ×V T O = (P T ) × V OT = V T × (P )T OT = VT × PT

Sparse Ops Q, P K, P Q, V K, V

Algo. Support Online-Only Mixed Mixed Online/Offline

Prefill 2× 2× 2× 2×

Decode 1.0× 1.5× 1.5× 2×

We analyze the four combinations of the GEMM orientation: • Naive: Allows Q and P to be sparse. Since both Q and attention scores P are dynamic activations generated at runtime, they cannot be pre-compressed. While the sparse GEMM itself is 2× faster, the end-to-end speedup is limited, and there is no decode benefit. • Trans-K: Transposing GEMM1 allows key compression, but GEMM2 still requires online compression for P . It achieves theoretical 2× prefill speedup but limited decode speedup as value cache is uncompressed. • Trans-V: Similar to Trans-K, it allows offline value compression but requires online processing for key. Prefill is accelerated (2×), but decode benefits are constrained by the dense key cache. • Trans-Both: By transposing both operations (calculating S T then OT ), we make K and V T the sparse operands. This allows us to use standard KV Cache compression algorithms that operate directly on K and V , ensuring full compatibility with existing workflows. This design delivers robust 2× speedup for both prefill and decode. From the analysis above, we choose the Trans-Both configuration. The attention computation can be reformulated as Algorithm 1, where Xnnz and Xe represent non-zero and metadata counterparts of a dense cache X. Although the Trans-Both scheme appears mathematically straightforward, implementing it effectively on GPUs is nontrivial because the operands of tensor cores are often stored in private registers per thread, which requires a re-layout between two GEMMs. In Section IV-C2, we will describe the detailed implementation of RelayoutFragment without using expensive communication like shared memory or warp shuffling.

Algorithm 1 The algorithm for fully sparse attention Require: Q, Knnz , Ke , Vnnz , Ve , block sizes Br , Bc Let Tr = ⌈n/Br ⌉ and Tc = ⌈n/Bc ⌉, divide Q, Knnz , Ke , Vnnz , Ve into corresponding blocks. for i = 1 to Tr in parallel do Initialize OiT , Li , Mi QTi ← LoadDense(Q, i) for j = 1 to Tc do Knnzj , Kej ← LoadSparse(Knnz , Ke , j) T ← SparseGEMM((Knnzj , Kej ), QTi ) Sji T T , OiT , Li , Mi ) , OiT , Li , Mi ← OnlineSoftmax(Sji Pji T T Pji ← RelayoutFragment(Pji ) T Vnnz , VeTj ← LoadSparse(Vnnz , Ve , j) j T T T Oi ← OiT + SparseGEMM((Vsp , VeTj ), Pji ) j end for OiT ← Scale(OiT , Li ) Store OiT to HBM as Oi end for return O

2) Mixed Semi-Structured Sparse Computation Workflow: With the support of block index map and memory pool design, in each iteration of key and value blocks, we first check the block index map to determine whether the current block is dense or sparse. If the block is dense, we load it from the dense memory pool and perform a dense GEMM operation. Otherwise, we load the nonzero elements and metadata from their respective memory pools and perform a sparse GEMM operation. The rest of the attention computation remains unchanged.

L B Sizeden = L · D · (1 − SK ) + L · D · (1 − SV ) 1 1 Sizennz = · L · D · SK + · L · D · SV 2 2 1 1 Sizee = · L · D · SK + · L · D · SV 16 16 Sizeidx = 2 ·

Speeduppref ill =

Sizebaseline Size idx + Size den + Size nnz + Size e block indices

(3)

(5c) (5d)

Tbaseline T dense + T sparse

(7)

dense blocks sparse blocks

During the prefill phase, attention is typically computationbound because of its quadratic complexity. The two GEMM operations each cost 2L2 D FLOPs, totaling 4L2 D FLOPs (softmax and scaling are comparatively small). Let C denote the dense tensor-core throughput in FLOPs/s. The baseline time is:

D. Efficiency Analysis

rcomp =

(5b)

Substituting Equations. (4) and (5) into Equation (3) yields: 1 rcomp = 1 1 − 0.21875 · (SK + SV ) + B·D (6) 1 ≈ 1 − 0.21875 · (SK + SV ) Given that block size B and hidden dimension D are usually large enough (e.g., B = 64, D = 128), the term 1 B·D is negligible in practice. Using the approximation, SK = 0.5, SV = 1.0 yields rcomp ≈ 1.49×, while SK = SV = 1.0 gives rcomp ≈ 1.78×. 2) Prefill Speedup: We define the ideal speedup for prefill phase as:

Tbaseline =

We provide a theoretical analysis of the memory compression ratio and speedup for prefill and decode phases when using HieraSparse with float16 2:4 sparse tensor core. Under this hardware setting, the metadata is 1/16 of the original dense tensor size. We ignore the batch and head dimension for simplicity, and denote the sequence length as L, hidden dimension as D, block size B, key block sparsity as SK ∈ [0, 1], and value block sparsity as SV ∈ [0, 1]. 1) Compression Ratio: The overall compression rate of HieraSparse is defined as the ratio between the original dense KV Cache size and the compressed representation:

(5a)

4L2 D C

(8)

With dense tensor-core throughput C FLOPs/s and sparse throughput 2C FLOPs/s, the time decomposes into dense and sparse parts: 2L2 D · (1 − SK ) 2L2 D · (1 − SV ) + C C 2L2 D · SK 2L2 D · SV Tsparse = + 2C 2C Tdense =

(9a) (9b)

Substituting Equations (8) and (9) into Equation (7) yields: Speeduppref ill =

4 4 − (SK + SV )

(10)

3) Decode Speedup: During the Decode phase, the attention computation is often memory-bound due to frequent KV Cache accesses. Given the compression ratio rcomp derived above, the ideal speedup can be calculated as:

dense blocks non-zero blocks metadata

Speedupdecode = rcomp =

The total size of the original dense KV Cache is: Sizebaseline = 2 · L · D

(4)

When applying HieraSparse together with block sparsity SK , SV , the size of each component is:

1 1 − 0.21875 · (SK + SV )

(11)

Together, these formulas quantify the ideal gains from compressing key and value cache. For example, SK = 0.5 and SV = 1.0 give Speeduppref ill = 1.6× and Speedupdecode ≈ 1.49×, while SK = SV = 1.0 yields 2.0× prefill speedup and 1.78× decode speedup.

IV. I MPLEMENTATION We implement HieraSparse kernels with TileLang [47], a domain-specific language for efficient acceleration kernel development and prototyping. We integrate our kernels and memory management scheme with the popular deep learning frameworks PyTorch [48]. A. Sparse Tensor Core Integration We modify the CUDA backend of TileLang to support sparse tensor core. We first support float16 sparse tensor core computation with mma.sp.m16n8k32, and merge two dense mma.m16n8k16 as a logical equivalent atom. This makes sure that: i) Operands can be loaded from shared memory to register using ldmatrix.x4. ii) Both atoms share the same fragment layout for P T matrix as we need to dynamically switch between mma and mma.sp. Every 8 groups of 2-bit metadata data are padded to int16 data type, which can be further vectorized into int16x8 to fully utilize the 128-bit vectorized memory instruction. B. Compression Kernel Implementation We provide a suite of CUDA kernels for hierarchical compression. We use an int16 index map, which can support up to 4M tokens with a block size of 64 tokens, satisfying most practical use cases. The data type can be lifted to int32 to support 256B tokens, far exceeding the context length of any current LLM. Firstly, the compression kernel iterates over M sequentially along the sequence dimension and maintains block counters: for the dense part, the kernel simply applies a vectorized memory copy from the original KV Cache to the dense pool; for the sparse part, we assume that element-level mask m has been applied to the original KV Cache. Each thread first loads the original elements and selects the nonzero elements into registers, then generates the 2-bit metadata by checking the position of non-zero elements in each group. Finally, the non-zero elements and metadata are stored in the sparse pool. Specifically for magnitude-based compression, we provide a compression kernel that fuses the mask generation and compression process. Instead of non-zero elements when compressing the sparse part, we directly select the top-2 magnitude elements in each 2:4 group, which further reduces the pruning and compression overhead. C. Computation Kernel Implementation Besides the sparse tensor core integration, we also implement several optimization techniques to further boost the performance of the attention kernel. For the memory-bound decode kernel, the KV Cache compression is sufficient to achieve the expected throughput. Additionally, we have a splitKV design where each thread block only processes a subset of the key and value blocks to increase parallelism. Each block outputs a partial output together with its own log-sum-exp, which are later combined in a lightweight post-processing kernel. The kernel was also optimized for Grouped-Query Attention (GQA) [49], [50], where multiple queries attending

Fig. 4: The performance gain of different optimizations for prefill kernel, measured with 32K context and batch size of 8 with Llama-3.1-8B-Instruct attention setting. to the same KV Cache head are viewed as a short duplicated sequence and reduce the padding overhead. For the prefill phase, three main additional optimizations are implemented to fully exploit the sparse tensor core. The ablated performance gain of these techniques is shown in Figure 4, and details are as follows: 1) Asynchronous Pipelining: We manually control the pipelining with a double buffering between key/value loading and two GEMMs. Specifically, we utilize the cp.async instruction to asynchronously load the key and value blocks from HBM to shared memory. While the tensor core is performing GEMM operations on the current tile resident in one shared memory buffer, the memory unit simultaneously fetches the next tile data into the other buffer. This buffering strategy effectively hides the memory access latency, ensuring that the compute units remain highly utilized throughout the attention computation. 2) In-fragment Re-layout: Unlike dense attention, which takes P as the first matrix operand of P × V , HieraSparse takes P T as the second matrix operand in V T × P T and T , VeT )×P T . This leads to a mandatory re-layout since the (Vnnz GEMM1 output operand layout is not the same as the GEMM2 input operand, which usually requires extra shared memory or multiple warp shuffling operations to transpose the data among threads. We utilize the movmatrix instruction to avoid this overhead. movmatrix is a warp-level communication instruction that operates on an 8 × 8 matrix atom stored in registers of all threads in a warp, allowing a row-major to column-major transpose (or vice versa) without accessing shared memory. To apply this to a larger size matrix as shown in Figure 5, we first partition the fragment into multiple 8 × 8 matrix atoms, then apply the instruction multiple times among the atoms at the same location in the fragment. With this method, we can achieve the re-layout with less overhead compared to shared memory or shuffle-based methods. 3) On-chip Memory Allocation and Specialized Kernel: Similar to FlashAttention, we store the tiled query matrix in shared memory to maximize data reuse and mitigate register pressure. For key and value tensors, we first load the required blocks from global memory to shared memory, then transfer them to registers before computation. For sparse blocks, we also load the corresponding metadata to shared memory and

transfer it to registers together with the non-zero elements. In a single iteration, either the shared memory of dense blocks or sparse blocks is active, so we manually overlap the allocation of them by setting the shared memory buffer pointer to the same address, which saves shared memory budget for sparse blocks. Furthermore, we provide specialized kernels for cases where the key/value cache is fully dense or fully sparse (e.g., initial phase of generation or specific layers). In these specialized kernels, we completely remove the shared memory allocation and loading logic for the unused format (e.g., no sparse buffer allocation for fully dense kernels), allowing for larger tile sizes and better occupancy.

movmatrix

Source Layout

𝑃𝑇 atom as D matrix of the first GEMM

Destination Layout

movmatrix(Src[0], Dst[0]); movmatrix(Src[2], Dst[4]); movmatrix(Src[1], Dst[1]); // ... movmatrix(Src[5], Dst[3]); movmatrix(Src[7], Dst[7]);

𝑃𝑇 atom as B matrix of the second GEMM

Fig. 5: The illustration of P T fragment re-layout. The source layout consists of multiple 16 × 8 D-matrix atoms, and the destination layout consists of multiple 32 × 8 B-matrix atoms, both in row-major. They are both partitioned into 8 × 8 atoms, and multiple movmatrix are issued to perform the re-layout without shared memory access.

V. E VALUATION We evaluate HieraSparse to demonstrate its effectiveness in balancing generation quality and system efficiency. Our evaluation is divided into two main parts: Quality-Sparsity Evaluation, which explores how different pruning strategies and sparsity settings affect the quality across various LLMs and generation stages; and Sparsity-Efficiency Evaluation, which assesses (i) the computational speedup in attention kernels, (ii) the memory compression efficiency, and (iii) the end-to-end performance gains. We compare our method with MUSTAFAR, the state-of-the-art fine-grained KV Cache pruning approach. MUSTAFAR is most closely related to our work in that it performs fine-grained element-wise KV pruning. However, it is limited to the decode phase and employs unstructured sparsity with a load-as-sparse, compute-as-dense scheme, which prevents it from fully translating sparsity into efficiency. All experiments are conducted on NVIDIA L40S with 48GiB DRAM. Python 3.10.19, PyTorch 2.10.0, and CUDA 12.8 are used as the software environment. The reproduction

scripts for both evaluations can be found in our open-source repository. A. Quality-Sparsity Evaluation We first evaluated the generation quality of HieraSparse under different sparsity settings to find the best trade-off. To understand the impact of pruning on different inference phases, we organized our evaluation into three experimental setups: i) Applying pruning exclusively to the decode stage under various sparsity settings. ii) Applying pruning to both the prefill and decode stages using uniform sparsity settings. iii) Applying pruning to both stages with different sparsity settings, where sparsity configurations are independently chosen for prefill and decode based on the insights from the previous two setups. The generation quality was measured using LongBench [51], a popular comprehensive benchmark for long-context LLM evaluation, which consists of 16 tasks across 6 categories, including single-document QA, multidocument QA, summarization, few-shot learning, synthetic tasks, and code completion. We tested our method on three popular models to make sure the conclusion is general: Llama3.1-8B-Instruct, Mistral-7B-Instruct-v0.2, and Qwen3-8B. We didn’t include MUSTAFAR results for Qwen3-8B since the model was not supported by the official implementation. For space reasons, we report the results of 6 categories and an averaged score; complete results of 16 tasks can be found in our repository. Following the common settings in prior work [17], [18], [23], [24], we kept the first 64 “sink” tokens and the last 256 “local window” tokens dense, and pruned the remaining tokens with different sparsity. 1) Pruning on Decode Stage Only: MUSTAFAR has performed a comprehensive evaluation on KV Cache during the decode phase, showing that the generation quality can be well preserved even pruned to 50% sparsity. Hence, we compared the LongBench scores between dense, MUSTAFAR, HieraSparse under different sparsity settings. To ensure a fair comparison across methods, we standardize hyperparameters (SK and SV for HieraSparse, and K and V for MUSTAFAR) by selecting configurations that achieve the same sparsity level, defined as the proportion of zero entries in the key and value caches. As shown in Table III, HieraSparse is able to maintain generation quality compared to the dense baselines across different models. For instance, under 50% Value sparsity, the average score drops by only 0.06 on Llama-3.1-8B-Instruct and 0.32 on Qwen3-8B. Even when configuring both key and value caches to 50% sparsity, the models retain strong performance with manageable degradation (e.g., a 2.08 drop on Llama-3.1-8B-Instruct). As expected, the semi-structured nature of our method introduces a slight quality degradation compared to the unstructured MUSTAFAR (e.g., trailing by 0.04 to 1.70 on Llama-3.1-8B-Instruct depending on the sparsity configuration). However, these quality trade-offs are well justified by a significant efficiency gain of up to 4.57× against MUSTAFAR. We also surprisingly observed that MUSTAFAR exhibits lower performance than the dense baseline, and the detailed efficiency evaluation will be presented in Section V-B.

TABLE III: LongBench score under different methods and sparsity settings when applying to decode stage.

Additionally, we observed that Mistral-7B-Instruct-v0.2 is more sensitive to our structural pruning than the other two models, showing a larger drop of up to 4.31 against the dense baseline when directly applying the same setting, which could be mitigated by tuning the sparsity settings.

(a) The average LongBench scores with different sparsity settings when extending pruned cache to prefill stage.

(b) The distribution of the key and value magnitude and respective average loss when pruning to semi-structured.

Fig. 6: The quality evaluation of HieraSparse when extended to prefill stage. 2) Uniformed Sparsity Pruning on Both Prefill and Decode: As shown in Figure 6a, we measured the overall LongBench scores in two settings: i) Keep all value cache as dense,

. .S pd At tn

1.0× 1.0× 1.8× 1.8×

Av g.

0% 50% 0% 50%

Co de

0% 0% 50% 50%

— SK 0.0, SV 1.0 SK 1.0, SV 0.0 SK 1.0, SV 1.0

Sy nt h

Dense HieraSparse HieraSparse HieraSparse

1.0× 1.0× 1.0× 1.5× 1.8× 1.5× 1.8×

Fe wSh

0% 50% 50% 0% 0% 50% 50%

mm .

0% 0% 0% 50% 50% 50% 50%

Llama-3.1-8B-Instruct 1.0× 42.87 44.65 1.5× 42.89 44.62 1.8× 43.60 44.69 1.0× 43.06 44.26 1.0× 40.90 43.23 1.5× 42.95 44.02 1.8× 40.96 43.21 Mistral-7B-Instruct-v0.2 1.0× 32.27 25.78 1.5× 35.88 30.18 1.8× 32.22 25.64 1.0× 36.32 30.23 1.0× 29.78 23.61 1.5× 36.28 30.40 1.8× 29.41 21.96 Qwen3-8B 1.0× 42.49 46.30 1.8× 42.04 46.77 1.0× 39.09 43.68 1.8× 37.26 42.92

— K0.0, V 0.5 SK 0.0, SV 1.0 K0.5, V 0.0 SK 1.0, SV 0.0 K0.5, V 0.5 SK 1.0, SV 1.0

Do c

Dense MUSTAFAR HieraSparse MUSTAFAR HieraSparse MUSTAFAR HieraSparse

lti-

1.0× 1.0× 1.0× 1.5× 1.8× 1.5× 1.8×

0% 50% 50% 0% 0% 50% 50%

Do c

Ke y

0% 0% 0% 50% 50% 50% 50%

gle -

Va lue

— K0.0, V 0.5 SK 0.0, SV 1.0 K0.5, V 0.0 SK 1.0, SV 0.0 K0.5, V 0.5 SK 1.0, SV 1.0

LongBench Score

S in

Ke y

Dense MUSTAFAR HieraSparse MUSTAFAR HieraSparse MUSTAFAR HieraSparse

Va lu e

Hy per pa ram

Comp. Rate

Me tho d

Sparsity

29.21 28.95 28.47 28.83 26.13 28.35 26.03

69.31 69.51 69.23 69.19 68.16 69.17 67.69

53.71 53.73 53.91 53.94 52.19 53.33 52.83

60.03 59.95 59.67 59.97 57.07 59.69 56.57

49.96 49.94 49.90 49.88 47.95 49.58 47.88

1.00× 0.32× 1.28× 0.32× 1.27× 0.37× 1.71×

27.91 27.58 27.31 27.96 25.06 27.84 24.37

66.72 66.73 66.87 66.70 66.15 66.65 65.36

46.45 42.31 43.99 43.52 37.66 41.92 33.92

54.96 54.80 54.71 54.91 53.24 54.79 53.19

42.34 42.91 41.79 43.27 39.25 42.98 38.03

1.00× 0.32× 1.28× 0.32× 1.27× 0.37× 1.71×

27.54 26.29 23.04 22.59

68.78 68.08 65.37 64.54

50.75 50.50 50.75 50.75

65.83 66.10 64.93 64.11

50.28 49.96 47.81 47.03

1.00× 1.28× 1.27× 1.71×

gradually increase key block sparsity SK (green line). ii) Keep all key cache as dense, gradually increase value block sparsity SV (blue line). The results indicate that pruning the key cache is much more sensitive than pruning the value cache, given the same block sparsity. By keeping the key cache dense and pruning the value cache to sparse, the overall score of LongBench only drops around 1.8, while keeping value cache dense and pruning key cache results in a much larger drop of 11.1. By inspecting the numeric distribution of key and value in Figure 6b, we conclude that this is mainly due to: i) Key and value exhibit different magnitude distributions. Key cache consistently shows much greater magnitude than value cache, which can lead to up to 4.8× magnitude loss when pruning lower 50% elements. ii) The calculation of the attention score involves exponential operations in softmax, which amplifies the impact of errors in key states; in contrast, errors in value representations have a more direct and linear effect on the final output. Thus, for magnitude-based pruning during the prefill stage, we keep the key cache dense to maintain generation quality.

3) Differentiated Sparsity Pruning on Both Prefill and Decode: We fix the prefill sparsity to SK 0.0, SV 1.0 and test different decode sparsity by either keeping the sparsity same or further pruning all key cache to SK 1.0, SV 1.0. The results show that aggressively pruning the decode key cache further boosts decode speedup from 1.28× to 1.71×, at the cost of a moderate accuracy drop across all three models. This suggests that differentiated sparsity between prefill and decode stages offers a flexible speedup–quality tradeoff.

TABLE IV: Average LongBench score and attention speedup under different sparsity settings for prefill and decode. Hyperparam Prefill

Decode

Attn. Spd. Prefill

Avg. Score

Decode

SK 0.0, SV 1.0 SK 0.0, SV 1.0

Llama-3.1-8B-Instruct SK 0.0, SV 1.0 1.34× 1.28× SK 1.0, SV 1.0 1.34× 1.71×

47.32 45.59

SK 0.0, SV 1.0 SK 0.0, SV 1.0

Mistral-7B-Instruct-v0.2 SK 0.0, SV 1.0 1.34× 1.28× SK 1.0, SV 1.0 1.34× 1.71×

41.30 38.18

SK 0.0, SV 1.0 SK 0.0, SV 1.0

Qwen3-8B SK 0.0, SV 1.0 1.34× SK 1.0, SV 1.0 1.34×

48.01 45.10

1.28× 1.71×

speedup of 4.57×, successfully converting sparsity into expected efficiency. In our experiment, we found MUSTAFAR performed worse than the dense baseline. By inspecting their methodology, experiment results, and official implementation, we conclude the inefficiency is due to three major reasons: i) The load-as-sparse and compute-as-dense scheme requires a decompression procedure before each mma issuance. The procedure includes looping over a 64-bit bitmap and moving sparse data to respective registers, which cannot be vectorized due to unstructured sparsity. This creates an extra latency within the computation pipeline. ii) The bitmap-based unstructured sparsity has a lower compression rate (details about compression in Section V-B2), which limits its theoretical speedup during the memory-bounded decode phase. iii) The current implementation lacks optimizations like kernel fusion and vectorized memory operations.

Fig. 7: Comparison of attention kernel latency, including compression overhead. The overhead of our method (HS) is minimal and barely visible in the figure.

B. Sparsity-Efficiency Evaluation In this section, we perform a detailed efficiency analysis without considering the exact pruning algorithm to explore the potential performance gain from HieraSparse. We benchmark the performance from different perspectives: i) Kernel performance. ii) Memory compression efficiency. iii) Per-layer and end-to-end performance when considering other computations during inference. 1) Kernel Performance: We first benchmark the kernel under prefill and decode stages of Llama-3.1-8B-Instruct with a batch size of 8 and context length of 32K, and compare the execution latency, as shown in Figure 7 (the same kernel acceleration results are also reported in Table III and Table IV). We also included the compression overhead of both systems, which contributes to prefill latency. The results indicate that when applying to the prefill stage, HieraSparse achieved a maximum 1.85× speedup when pruning both key and value, and a 1.36× speedup when only pruning value; when applying to the decode stage, HieraSparse achieved 1.71× and 1.28× speedup respectively. In addition, the compression accounts for only 0.5% of the prefill attention latency, whereas MUSTAFAR incurs up to 11.7% under 50% sparsity for both key and value. Compared with the unstructured decode kernels under equivalent sparsity, our implementation achieved a significant

(a) Kernel speedup against dense baseline for different block sparsity.

(b) Comparison between HieraSparse and MUSTAFAR on compression rate.

Fig. 8: The efficiency evaluation of HieraSparse under different sparsity. We also benchmark kernel speedup across block sparsity levels, as shown in Figure 8a. The decode kernel speedup closely follows the theoretical curve, with a small gap because non-memory operation latencies are excluded from the model. In contrast, the prefill speedup is slightly offset: at low sparsity it exceeds the theoretical estimate, likely due to reduced memory traffic and higher L2 cache hit rate. The trend is also flatter than theory because, even after overlapping sparse and dense shared memory usage, the total shared memory is still at least the same as a dense counterpart kernel, which limits occupancy and thus speedup; this discrepancy disappears at 100% block sparsity, where a specialized kernel with much lower shared memory consumption can be used. 2) Memory Compression Efficiency: Besides computation speedup, another important benefit of pruning the KV Cache is the reduction of memory usage. As shown in Figure 8b, we measured the compression rate of HieraSparse under different

TABLE V: End-to-end performance comparison on Llama-3.1-8B-Instruct. Context

Prefill Method

TTFT (Speedup)

Decode Method

TPOT (Speedup)

Key Mem

Value Mem

Peak Mem

32k