Skip to content

llama : build unified decode masks in one cache scan - #29769

Closed
joelagnel wants to merge 1 commit into
ggml-org:masterfrom
joelagnel:pr/kv-unified-decode-mask
Closed

joelagnel wants to merge 1 commit into
ggml-org:masterfrom
joelagnel:pr/kv-unified-decode-mask

Conversation

@joelagnel

Copy link
Copy Markdown
Contributor

llama : build unified decode masks in one cache scan

Use sequence membership bits for batches with one token per sequence.
Preserve causal masking and two-dimensional position ties, with the
existing path for other batches.

The fast path applies to causal masks without SWA or ALiBi, in a single
unified KV stream, for batches of at least four tokens with exactly one
token per sequence and sequence IDs below 64: parallel decode, and MTP
draft steps. It fills the mask once and then visits each cell once,
unmasking it for every active sequence that owns it. This is faster than
scanning every cell for each token when the tokens belong to different
sequences.

Measured on RTX 5090 (compute capability 12.0; Windows x64, driver
617.14, CUDA 12.9, static build) with Qwen3.8-27B Q4_K_M and an F16 KV
cache.

llama-batched-bench -c 131072 -b 2048 -ub 512 -npp 15872 -ntg 512
-npl 8 -ngl 99 -fa on -ctk f16 -ctv f16 -kvu

Medians without -> with; the change is the mean of the 10 paired
per-repetition changes:

274.54 -> 291.21 tok/s, +5.77%
(95% CI +5.11% to +6.43%)

For completeness, ran the following as well on 5090, these don't trigger
the fast path so were only sanity runs for merge requirements (on top of
master (4453b53), without -> with this patch):

  • llama-bench -ngl 99 -fa 1, two rounds of 5 repetitions in mirrored
    order: pp512 3731.18 -> 3718.85 tok/s (-0.33%), tg128 78.67 -> 78.49
    (-0.23%). llama-bench runs a single sequence.
  • llama-perplexity -c 512 --chunks 50 on wikitext-2 test: 6.1814 +/-
    0.13150 with and without the patch.

Assisted-by: Codex
Signed-off-by: Joel Fernandes joelagnelf@nvidia.com

Requirements

  • I have read and agree with the contributing guidelines
  • AI usage disclosure: YES - Used AI to code, refine and test.

cc @HanClinto @ORippler @gaugarg-nv

Use sequence membership bits for batches with one token per sequence.
Preserve causal masking and two-dimensional position ties, with the
existing path for other batches.

The fast path applies to causal masks without SWA or ALiBi, in a single
unified KV stream, for batches of at least four tokens with exactly one
token per sequence and sequence IDs below 64: parallel decode, and MTP
draft steps. It fills the mask once and then visits each cell once,
unmasking it for every active sequence that owns it. This is faster than
scanning every cell for each token when the tokens belong to different
sequences.

Measured on RTX 5090 (compute capability 12.0; Windows x64, driver
617.14, CUDA 12.9, static build) with Qwen3.8-27B Q4_K_M and an F16 KV
cache.

  llama-batched-bench -c 131072 -b 2048 -ub 512 -npp 15872 -ntg 512
      -npl 8 -ngl 99 -fa on -ctk f16 -ctv f16 -kvu

Medians without -> with; the change is the mean of the 10 paired
per-repetition changes:

  274.54 -> 291.21 tok/s, +5.77%
                   (95% CI +5.11% to +6.43%)

For completeness, ran the following as well on 5090, these don't trigger
the fast path so were only sanity runs for merge requirements (on top of
master (4453b53), without -> with this patch):

- llama-bench -ngl 99 -fa 1, two rounds of 5 repetitions in mirrored
  order: pp512 3731.18 -> 3718.85 tok/s (-0.33%), tg128 78.67 -> 78.49
  (-0.23%). llama-bench runs a single sequence.
- llama-perplexity -c 512 --chunks 50 on wikitext-2 test: 6.1814 +/-
  0.13150 with and without the patch.

Assisted-by: Codex
Signed-off-by: Joel Fernandes <joelagnelf@nvidia.com>
@ggml-gh-bot

ggml-gh-bot Bot commented Sep 30, 2026

Copy link
Copy Markdown

Hi @joelagnel, thanks for your contribution!

Per our contribution guidelines, the automated PR checker found the following issue(s) that need your attention:

  • Multiple open PRs from a new contributor: We limit new contributors (those without a previously merged PR) to 1 open PR at a time. You currently have 2 open PRs.

Please note that maintainers reserve the right to make final decisions on PRs. If you believe there is a mistake, please comment below.

@ggerganov ggerganov self-assigned this Sep 30, 2026
@joelagnel

Copy link
Copy Markdown
Contributor Author

With #29510, this patch no longer gives improvement as masks are built differently now, closing this PR in favor of #29510

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants