Sparton: Fast Kernel for Sparse Retrieval

• Sparton is a fast, memory-efficient Triton kernel designed to fuse tiled GEMM and online max-reduction, eliminating the need to materialize a large dense logit tensor in learned sparse retrieval.
• It leverages the monotonicity of the post-logit activation to reorder operations, cutting data movement from 16 GB to around 31 MiB and enabling larger batch sizes and longer sequences.
• Sparton uses a hybrid design that pairs vendor GEMM libraries with custom Triton kernels for both forward and backward passes, preserving retrieval effectiveness while boosting efficiency.

Searching arXiv for the specified paper and closely related learned sparse retrieval work. Sparton is a fast, memory-efficient Triton kernel for the language-model (LM) head used in learned sparse retrieval (LSR) systems such as Splade. It addresses a specific systems bottleneck in sparse lexical representation learning: the LM head projects transformer hidden states into a vocabulary-sized lexical space, producing an intermediate tensor of shape batch $\times$ sequence $\times$ vocabulary, and that tensor becomes prohibitively expensive to materialize for large batches, long sequences, or large vocabularies. The method introduced in "Sparton: Fast and Memory-Efficient Triton Kernel for Learned Sparse Retrieval" reorders and fuses the LM-head computation so that the dense logit tensor is never materialized in high-bandwidth memory (HBM), thereby reducing memory footprint and HBM traffic while preserving retrieval effectiveness (Nguyen et al., 26 Mar 2026). In the LSR context defined by models such as Splade, Sparton is therefore best understood as a kernel-level optimization of the lexical projection stage rather than as a new retrieval objective or ranking model; the similarly named SPARTan for PARAFAC2 is a distinct method in tensor mining and not part of the LSR literature (Perros et al., 2017).

1. Problem setting in learned sparse retrieval

State-of-the-art LSR models, including Splade, employ an LM head to map hidden states into a lexically anchored vocabulary space (Nguyen et al., 26 Mar 2026). Let $H \in \mathbb{R}^{B \times S \times D}$ denote transformer hidden states, $E \in \mathbb{R}^{|\mathcal{V}| \times D}$ the vocabulary embedding matrix, $b \in \mathbb{R}^{|\mathcal{V}|}$ the bias, and $M \in \{0,1\}^{B \times S}$ the attention mask. The sparse lexical representation is defined as

$$\mathbf{Y} = \max_{s} \left[ \log\left(1 + \text{ReLU}(H E^\top + b) \right) \odot M' \right]$$

where $M'$ is $M$ broadcast over the vocabulary dimension and $\max_s$ pools over the sequence dimension $\times$0 (Nguyen et al., 26 Mar 2026).

The standard PyTorch-style execution order is explicit: compute logits $\times$1, add bias, apply mask, apply ReLU, apply $\times$2 or Log1P, then max-pool across sequence to obtain $\times$3 (Nguyen et al., 26 Mar 2026). This decomposition is algorithmically straightforward, but it requires materializing the dense logit tensor $\times$4, which is the dominant source of memory blow-up and throughput degradation in LSR training (Nguyen et al., 26 Mar 2026).

The practical severity of this bottleneck is quantified by a concrete example. For $\times$5, $\times$6, and $\times$7, the intermediate logit tensor alone requires about 16 GB in half precision (Nguyen et al., 26 Mar 2026). This constrains batch size, sequence length, vocabulary size, and overall training speed. The paper’s characterization is that the bottleneck is not merely floating-point arithmetic; it is primarily data movement between HBM and on-chip SRAM, because standard frameworks execute GEMM and subsequent operators as separate stages and repeatedly read and write the same massive tensor (Nguyen et al., 26 Mar 2026).

This problem is structurally amplified in multilingual sparse retrievers. The vocabulary can range from 30,000 to over 250,000 tokens in recent models, so the cost of the LM head scales directly with lexical coverage (Nguyen et al., 26 Mar 2026). A plausible implication is that systems-level LM-head optimization becomes increasingly central as LSR moves from monolingual to multilingual backbones.

2. Computational bottleneck and memory-I/O pathology

The key systems diagnosis in Sparton is that standard execution materializes the full logit tensor in HBM and then subjects it to repeated operator-wise passes (Nguyen et al., 26 Mar 2026). GEMM writes the full output to HBM; ReLU reads and writes it again; Log1P reads and writes it again; max-pooling reads it again. Consequently, the LM head is dominated by repeated HBM traffic rather than by compute alone (Nguyen et al., 26 Mar 2026).

The paper emphasizes the architectural mismatch underlying this behavior. On-chip SRAM is tiny compared to HBM, and repeatedly moving the $\times$8 tensor between HBM and SRAM is inefficient (Nguyen et al., 26 Mar 2026). PyTorch compilation can fuse some elementwise operators, but it cannot fuse the matrix multiplication itself because GEMM is typically delegated to black-box vendor libraries such as cuBLAS or rocBLAS (Nguyen et al., 26 Mar 2026). As a result, the dense intermediate still has to be written out.

The resulting pathology has two coupled dimensions. The first is peak-memory inflation: storing the full dense logits consumes large amounts of memory even though the final representation requires only a sequence-wise maximum per $\times$9 pair. The second is bandwidth overhead: the same tensor is repeatedly streamed through HBM for masking, activation, and pooling (Nguyen et al., 26 Mar 2026). Sparton’s significance lies in reframing the LM head as an I/O-limited kernel whose optimization requires changing the order of operations, not merely accelerating existing operator calls.

This interpretation is consistent with broader trends in sparse retrieval engineering. Splade established the utility of lexicalized sparse expansions in neural retrieval, but its practical training pipeline inherits large-vocabulary costs from the LM projection stage. Sparton therefore targets a systems bottleneck inside an otherwise established LSR formulation rather than altering Splade’s sparse retrieval semantics (Nguyen et al., 26 Mar 2026).

3. Algebraic reformulation and fused-kernel design

Sparton’s central observation is that the post-logit activation

$H \in \mathbb{R}^{B \times S \times D}$0

is monotonically non-decreasing, so

$H \in \mathbb{R}^{B \times S \times D}$1

This permits max-pooling to be moved before ReLU and Log1P (Nguyen et al., 26 Mar 2026). The consequence is substantial: instead of applying ReLU and Log1P to a full $H \in \mathbb{R}^{B \times S \times D}$2 tensor, Sparton first reduces over the sequence dimension and then applies the activation only to a $H \in \mathbb{R}^{B \times S \times D}$3 tensor (Nguyen et al., 26 Mar 2026).

The paper gives a concrete numerical illustration. With $H \in \mathbb{R}^{B \times S \times D}$4, $H \in \mathbb{R}^{B \times S \times D}$5, and half precision, activation I/O drops from 16 GB per pass to about 31 MiB (Nguyen et al., 26 Mar 2026). This is the algebraic basis for the claimed memory and throughput gains.

Sparton implements this idea as a fused Triton kernel that combines tiled matrix multiplication output handling, masking, online max-reduction, ReLU, Log1P, and output writing (Nguyen et al., 26 Mar 2026). The defining mechanism is early online reduction on raw logit tiles: the kernel computes a tile of raw logits, immediately compares it with current running maxima, keeps only the best value per $H \in \mathbb{R}^{B \times S \times D}$6, optionally stores the argmax index, and discards the raw tile (Nguyen et al., 26 Mar 2026). The full logit tensor is therefore never materialized in HBM; only reduced maxima and, optionally, the indices of maximal sequence positions are stored (Nguyen et al., 26 Mar 2026).

This design is more precise than a generic statement about operator fusion. In Sparton, the critical fusion target is not simply “several elementwise operators” but the semantic collapse of the sequence dimension before expensive post-processing. A plausible implication is that Sparton’s effectiveness depends specifically on the monotonicity property of the post-logit transformation and the pooling structure of LSR, rather than on a universally applicable fusion heuristic.

4. Hybrid Triton–vendor GEMM implementation

The paper distinguishes between an ideal fully fused implementation and a practical hybrid realization (Nguyen et al., 26 Mar 2026). In principle, a fully custom Triton kernel could fuse GEMM, bias addition, masking, max reduction, and activations in one implementation. In practice, vendor GEMM libraries are faster than a custom Triton GEMM, so Sparton adopts a hybrid design (Nguyen et al., 26 Mar 2026).

The implemented workflow is: use cuBLAS or rocBLAS to compute tiled logits for a vocabulary tile, immediately pass the tile to a Triton kernel, and have Triton perform masked max-reduction over sequence, ReLU, Log1P, and writing only the reduced output (Nguyen et al., 26 Mar 2026). The fusion is therefore partial at the GEMM level but complete for the post-GEMM pipeline.

The forward algorithm is expressed conceptually per vocabulary tile. For each $H \in \mathbb{R}^{B \times S \times D}$7 and $H \in \mathbb{R}^{B \times S \times D}$8, the method loads $H \in \mathbb{R}^{B \times S \times D}$9, $E \in \mathbb{R}^{|\mathcal{V}| \times D}$0, and $E \in \mathbb{R}^{|\mathcal{V}| \times D}$1, computes

$E \in \mathbb{R}^{|\mathcal{V}| \times D}$2

applies max reduction over the sequence dimension,

$E \in \mathbb{R}^{|\mathcal{V}| \times D}$3

then applies $E \in \mathbb{R}^{|\mathcal{V}| \times D}$4, and writes the reduced outputs and max indices to HBM (Nguyen et al., 26 Mar 2026). The tiling strategy is along the batch and vocabulary dimensions, enabling block-wise GPU parallelism (Nguyen et al., 26 Mar 2026).

The paper’s justification for Triton is implementation-specific rather than theoretical: Triton enables a custom GPU kernel specialized for the LM head while retaining high-level expressiveness and sufficient control over fusion and tiling (Nguyen et al., 26 Mar 2026). This suggests that Sparton’s contribution lies at the intersection of retrieval systems and GPU kernel design. It is not a new sparse encoder architecture; it is a specialized execution strategy for a particular stage of sparse encoder training.

5. Backward pass and reduced activation state

Sparton also fuses the backward pass (Nguyen et al., 26 Mar 2026). From the forward pass it stores only the max value $E \in \mathbb{R}^{|\mathcal{V}| \times D}$5 and the argmax index $E \in \mathbb{R}^{|\mathcal{V}| \times D}$6 for each $E \in \mathbb{R}^{|\mathcal{V}| \times D}$7 (Nguyen et al., 26 Mar 2026). This compact saved state replaces the need to preserve the dense logit tensor for autograd.

Given upstream gradient $E \in \mathbb{R}^{|\mathcal{V}| \times D}$8, the derivative through $E \in \mathbb{R}^{|\mathcal{V}| \times D}$9 is handled as

$b \in \mathbb{R}^{|\mathcal{V}|}$0

Gradients are then routed only to the selected sequence position:

$b \in \mathbb{R}^{|\mathcal{V}|}$1

and

$b \in \mathbb{R}^{|\mathcal{V}|}$2

with atomic accumulation where needed (Nguyen et al., 26 Mar 2026).

The stated advantage is that backward no longer requires storing and revisiting the dense logit tensor. Instead of $b \in \mathbb{R}^{|\mathcal{V}|}$3 activation storage, the saved state is reduced to about $b \in \mathbb{R}^{|\mathcal{V}|}$4 (Nguyen et al., 26 Mar 2026). This is important because in large-batch training, backward activation state can dominate practical memory limits even when forward kernels are accelerated.

The backward design also clarifies the scope of Sparton’s exactness. The kernel does not approximate the sequence-wise maximum or sparsify gradients heuristically; it stores the exact max value and argmax index and routes gradient accordingly (Nguyen et al., 26 Mar 2026). This is relevant to the paper’s claim of “no effectiveness loss” in end-to-end retrieval experiments, because the training signal remains faithful to the original max-pooled LSR formulation (Nguyen et al., 26 Mar 2026).

6. Empirical behavior in kernel benchmarks and end-to-end training

The paper reports kernel-only and end-to-end results that are explicitly tied to the LM-head bottleneck (Nguyen et al., 26 Mar 2026). In isolation, Sparton achieves up to a 4.8× speedup and an order-of-magnitude reduction in peak memory compared to PyTorch baselines; relative to a compiled baseline, the best point reaches 4.8× speedup and 12× reduction in peak memory (Nguyen et al., 26 Mar 2026). When batch size, sequence length, and vocabulary size are varied, PyTorch baselines scale steeply, often linearly or worse, whereas Sparton remains much flatter, and the performance gap widens as inputs grow (Nguyen et al., 26 Mar 2026).

A backward-pass sequence-length comparison at $b \in \mathbb{R}^{|\mathcal{V}|}$5 and $b \in \mathbb{R}^{|\mathcal{V}|}$6 illustrates this scaling:

Sequence length	Tiled LM eager	Tiled LM compiled	Sparton
1024	279.5 ms, 8.51 GB	150.6 ms, 14.75 GB	71.2 ms, 0.99 GB
2048	531.1 ms, 16.82 GB	294.2 ms, 29.37 GB	118.6 ms, 1.55 GB
4096	1031.0 ms, 33.44 GB	OOM	212.7 ms, 2.68 GB
8192	OOM	OOM	399.6 ms, 5.13 GB

In this setting, Sparton is the only method that scales to sequence length 8192 (Nguyen et al., 26 Mar 2026). The table underscores the paper’s central thesis that memory traffic and activation storage, not only arithmetic throughput, determine practical scalability.

In end-to-end Splade training on an H200 GPU, the paper compares Sparton with a compiled PyTorch LM head at batch size 384. The compiled LM head requires 14.24 hours, 125.78 GB peak memory, and achieves NDCG@10 = 0.421; Sparton requires 12.38 hours, 96.83 GB peak memory, and achieves NDCG@10 = 0.416 (Nguyen et al., 26 Mar 2026). The paper interprets this as 14% faster training, nearly 30 GB less memory, and no meaningful effectiveness loss (Nguyen et al., 26 Mar 2026).

The reduced memory footprint also permits larger batches. With Sparton, batch size can be increased to 512, yielding 12.24 hours, 128.63 GB peak memory, and NDCG@10 = 0.427 (Nguyen et al., 26 Mar 2026). The reported interpretation is that larger batches become feasible and can slightly improve retrieval effectiveness (Nguyen et al., 26 Mar 2026). This suggests that kernel-level optimization can indirectly change the optimization regime of LSR by expanding the feasible training configuration space.

7. Relation to multilingual scaling, significance, and name ambiguity

The paper reports especially strong gains on a multilingual backbone, xlm-roberta-base, with vocabulary around 250k (Nguyen et al., 26 Mar 2026). In that setting, Sparton enables a 26× larger batch size, 420 versus 16, and 2.5× faster training, 67 versus 25 hours (Nguyen et al., 26 Mar 2026). Because the LM-head cost scales with vocabulary size, the advantage grows with larger multilingual vocabularies (Nguyen et al., 26 Mar 2026). This makes Sparton particularly relevant for multilingual sparse retrievers, where lexical coverage expands dramatically relative to monolingual encoders.

The method’s main contribution is therefore not only a faster kernel but a reformulation of the LM head that makes the LM head “no longer the dominant bottleneck in LSR training,” especially for Splade-like systems and multilingual sparse retrievers (Nguyen et al., 26 Mar 2026). The practical benefits listed by the paper are larger batch sizes, longer sequences, larger vocabularies, faster training, lower memory consumption, and preserved retrieval effectiveness (Nguyen et al., 26 Mar 2026).

A recurrent source of confusion is nomenclature. “Sparton” in (Nguyen et al., 26 Mar 2026) is a Triton kernel for LSR LM heads; “SPARTan” in "SPARTan: Scalable PARAFAC2 for Large & Sparse Data" is a scalable implementation of PARAFAC2 for large, sparse, irregular multi-subject data (Perros et al., 2017). The latter exploits the special structure of PARAFAC2 and MTTKRP to accelerate tensor decomposition and is unrelated to learned sparse retrieval. The shared name reflects a surface similarity in emphasis on scalability and memory efficiency, but the two methods address different models, different data structures, and different computational kernels.

Within the LSR literature, Sparton occupies a systems-oriented position. It does not replace Splade’s lexical expansion principle, nor does it introduce a new ranking loss or sparse regularizer. Instead, it specializes GPU execution for the exact LM-head computation already used in LSR and exploits monotonicity, tiling, and early online reduction to avoid dense intermediate materialization (Nguyen et al., 26 Mar 2026). This suggests a broader methodological lesson: in sparse retrieval pipelines with very large vocabularies, algebraic reorderings that reduce I/O and activation state can be as consequential as architectural changes at the model level.