Seeking technical feedback on SHBF: sparse long-context LLM inference with based candidate selection

Accelerated Computing CUDA CUDA Programming and Performance
1

Hi NVIDIA Developer Community,

This is not a bug report. I am looking for technical feedback on an experimental long-context LLM inference optimization I am working on, called SHBF.

At a high level, SHBF is a sparse attention path that uses a Hamming-style coarse retrieval stage over compact Q/K binary signatures, followed by exact reranking and sparse attention over a bounded candidate set.

The goal is to reduce quadratic attention/memory pressure in long-context inference while preserving quality.

Some current prototype observations:

  • In one Qwen-based long-context test, native attention at 8192 tokens used about 10.9 GB, while the SHBF path used about 5.6 GB.
  • At 16384 tokens, the native path reached OOM, while SHBF still ran at about 8.1 GB.
  • In the best validated path, quality degradation was small, around ~1% PPL delta in some long-context tests.
  • A recent SHBF v10 candidate showed about ~86–88% estimated attention-side memory saving in a limited validation harness, not total model VRAM reduction.
  • Runtime latency is still the main open problem. The current prototype is slower than optimized dense attention because candidate selection, reranking, sparse attention, and Python/PyTorch overhead are not yet implemented as optimized GPU kernels.

The approach is not designed to be specific to Qwen. In principle, it targets the attention mechanism itself and should be applicable to decoder-only Transformer LLMs that expose standard Q/K/V attention tensors. I have been using Qwen models mainly as convenient validation targets, not because the method depends on Qwen-specific weights.

I am trying to understand the best NVIDIA implementation path for this kind of candidate-selection + sparse-attention workload.

My main questions are:

  1. Would this kind of popcount/Hamming + top-k candidate selector be better implemented as a custom CUDA kernel, TensorRT plugin, Triton kernel, CUTLASS-style component, or another NVIDIA stack component?

  2. Are there known pitfalls with sparse candidate gathering, reranking, and sparse attention on NVIDIA GPUs?

  3. What would be the fairest benchmark methodology against dense attention, sliding-window attention, FlashAttention-style dense baselines, or existing sparse-attention approaches?

Environment

Current experiments are prototype-level.

GPU Type: Tesla T4 for Colab validation
CUDA Version: 12.8 in recent Colab run
PyTorch Version: 2.11.0+cu128 in recent Colab run
Models tested: Qwen2.5-0.5B and reduced tests with Qwen2.5-1.5B-Instruct
TensorRT Version: N/A
TensorFlow Version: N/A

Any guidance on the right NVIDIA stack, benchmark design, or relevant examples would be appreciated.

Thanks,
João Vitor

3

Hi @joaofags, thanks for laying out the prototype and the current measurements clearly.

This is not a TensorRT support issue yet, since the current path is PyTorch/prototype-level and TensorRT Version is N/A. The right first discussion is kernel design: popcount/Hamming candidate selection, top-k, sparse gather, reranking, and sparse attention scheduling on NVIDIA GPUs.

I would split the implementation path like this:

  1. Prove the candidate-selection and sparse-attention primitives outside TensorRT first, either as custom CUDA kernels or another kernel-level implementation.
  2. Benchmark against dense attention, sliding-window attention, and FlashAttention-style baselines with the same model, sequence length, batch size, precision, KV-cache policy, and warmup.
  3. Once the op contract is stable, use a TensorRT plugin only if you need this custom op to live inside a TensorRT engine.

For that first stage, CUDA Programming and Performance is the better category:

I’m moving the thread there so the CUDA folks can comment on the kernel and benchmark design. You do not need to repost, but a minimal kernel sketch or profiling breakdown would help them give more concrete advice.

Thanks,
Atharva

4

Related to a custom CUDA implementation, you could reuse cub::DeviceTopk. There are other initiatives for Block-level TopK, SegmentedTop-K, and multi-key variants. While for popcount/Hamming, you could use a custom operator with cuda::std::popcount + DeviceTransform