nvidia

/

Qwen3-8B-DMS-8x

like 37

Follow

NVIDIA 66.1k

Hi DMS team! Following up from a conversation on kvpress#184 with @alancucki and @SimJeg , who suggested continuing the discussion here.

I've been benchmarking DMS-8x inference on an RTX 4090 (24 GB) and wrote a set of Triton GPU kernels to optimize the prefill path. I also benchmarked against the flex_attention prefill from Model Optimizer. Sharing the results and a few questions.

Benchmark Results

Hardware: RTX 4090 24 GB, Qwen3-8B-DMS-8x, batch size 1, bf16, 128 output tokens.
Needle-in-a-haystack retrieval task, 3 iterations averaged per data point.

Prefill Speed (TTFT, no_think mode)

Peak VRAM

DMS-8x vs Vanilla vs kvpress (prefill + decode, 512 output tokens)

Decode speed is comparable across all variants (~20-42 tok/s), so prefill is the key differentiator.

Triton Kernels

Five kernels replacing the Python batch-loops in dms_attention.py and dms_cache.py:

1. left_pad_2d — replaces left_pad_one() serial loop
2. scatter_by_index — replaces restore_order() loop
3. bool_gather_left_pad — replaces convert_to_left_padding() inner loop
4. compact_by_bool — replaces get_contiguous_cache() compaction loop
5. _dms_flash_attn_fwd_kernel — single-pass flash attention with DMS causal + window + eviction masking, replacing the chunked SDPA loop

The key insight is that the eviction state can be pre-computed upfront — the sequential chunk dependency in dms_prefill_attention() is redundant since eviction is monotonic.

Full code: github.com/westers/kvcompress

Questions

1. FlexAttn integration path: The flex_attention prefill from Model Optimizer is clearly superior for speed. Are there plans to integrate it into the HuggingFace model files directly? Currently users need to clone Model Optimizer separately.

2. Decode-phase cache compression: DMS achieves 8x compression during prefill, but at shorter contexts (4-8K), the total time is still dominated by the prefill overhead vs vanilla. Is there work on reducing the prefill cost further, or is the focus on longer contexts where the VRAM savings compound?

3. Upstream interest: Would there be interest in upstreaming the Triton utility kernels (left_pad, scatter, compact) to the Hub model files? They're drop-in replacements with no API changes and improve prefill by 1.2-1.7x on their own.

Thanks for the great work on DMS — the 8x KV compression with minimal quality loss is impressive, and it's exciting to see it in a usable HuggingFace model.