Skip to content

feat(qwen35): add paged-attention monolithic decode-only AR foundation - #555

Open
Graffioh wants to merge 4 commits into
Luce-Org:mainfrom
Graffioh:codex/qwen36-paged-attention
Open

feat(qwen35): add paged-attention monolithic decode-only AR foundation#555
Graffioh wants to merge 4 commits into
Luce-Org:mainfrom
Graffioh:codex/qwen36-paged-attention

Conversation

@Graffioh

@Graffioh Graffioh commented Jul 23, 2026

Copy link
Copy Markdown
Contributor

Summary

  • Add an exact fixed-block paged K/V allocator (transactional allocation, stale-handle detection, immediate release).
  • Add a CUDA/HIP GGML_OP_PAGED_ATTN decode op for D=256 GQA with F16, Q4_0, and Q8_0 K/V.
  • Optimize decode: one warp computes all GQA heads of a K/V group per K/V row loaded (batched q8_1 dots); 32-token score tiles with a skip-if-unchanged accumulator rescale and exp2f softmax; Q quantized once per step by a pre-pass; partition partials stored normalized in f16 (half the scratch traffic); occupancy-aware context partitioning with a stable split-softmax merge.
  • Keep paging metadata (block table + sequence length) resident in the persistent target cache next to the K/V pool: decode steps upload 4 bytes per fresh 16-token block instead of the whole live table every token.
  • Integrate 16-token paged autoregressive decode for dense Qwen3.5/Qwen3.6 behind --paged-attention (one CUDA or HIP device).
  • Fix first-token AR bookkeeping so the first sampled token is forwarded and written at its committed K/V and RoPE position.
  • For now: keep exact chunked prefill; reject incompatible snapshots, speculative decoding, DDTree, KVFlash, layer split, and finite FA windows.
  • Add allocator, nine-type numerical, block-boundary, uniform-kernel, and ragged-capacity tests.

Scope

Paged K/V is the memory-management foundation for heterogeneous 4–8-user serving. This PR establishes exact single-request Qwen integration and a multi-sequence GPU primitive first, following the block-table design in ggml-org/llama.cpp#21961.

HTTP integration stays one active sequence: Qwen3.6's Gated DeltaNet conv/SSM state is backend-global and the worker is run-to-completion. Sequence-indexed recurrent state, batched decode, and a continuous-batching scheduler are follow-up work. The primary benefit is memory capacity and request-lifecycle flexibility; the benchmarks below measure attention operations, not whole-model tokens/s.

Opt in with --paged-attention. Unsupported combinations fail explicitly. Requests reaching --max-ctx finish cleanly.

Benchmark results

Qwen attention dims (D=256, Hq=24, Hkv=4). Single metric: paged ÷ contiguous throughput — above 1.0x = paged faster.

Uniform batches (identical context per sequence, Q4_0 K/V, batches 1/2/4/8)

GPU / backend 4K 32K 128K
RTX 3090 / CUDA 1.11–1.19x 1.13–1.27x 1.27–2.02x
Radeon 8060S / HIP 4.28–5.60x 4.72–5.16x 5.00–6.42x

Both rows reflect the reworked kernel. Measured on the same GPUs, the pre-rework kernel gave 0.46–0.77x on CUDA and 2.06–2.92x on HIP, so the rework is worth roughly 1.5–2.6x on CUDA and 1.3–2.3x on HIP in absolute paged step time. The HIP ratios are inflated by a weak native HIP flash_attn_ext_vec baseline — don't extrapolate them to whole-model throughput.

Ragged 8-request workload (128K + 7 × 8K, Q4_0 K/V)

Standalone synthetic multi-sequence kernel benchmark (constructs the interleaved block table directly; does not use PagedKvPool or demonstrate concurrent requests — the server integration remains single-sequence). Paging stores only live, block-rounded tokens (188,416) instead of padding every sequence to 128K (1,048,576 slots).

GPU / backend contiguous paged ratio
RTX 3090 / CUDA 3.139 ms 2.234 ms 1.40x
Radeon 8060S / HIP 11.665 ms 2.389 ms 4.88x

Small-context ragged batches (e.g. 4K + 7×256) remain behind contiguous: short sequences carry dead partitions sized by the longest sequence. The stream-K/persistent-kernel rework staged with continuous batching addresses this.

K/V capacity (projection over Qwen3.6's 16 full-attention layers from the one measured layer; excludes weights, activations, recurrent state):

layout one-layer K/V 16-layer K/V
padded per-sequence contiguous 1152 MiB 18.000 GiB
paged 207 MiB 3.234 GiB

5.57x less K/V (82.03% saving), now alongside faster attention on CUDA rather than at its expense.

Roadmap

  1. This PR: decode-only ggml_paged_attn, block manager, CUDA/HIP integration, benchmarks.
  2. Scheduler groundwork: sequence-indexed DeltaNet/conv state, iteration-level scheduling, dynamic decode batches.
  3. Paged-aware chunk prefill: arbitrary block-table writes, multi-token queries over preceding paged K/V, causal attention within the chunk.
  4. Continuous admission and interleaving: token budgets, decode prioritization, admission control, cancellation, prefill/decode interleaving; plus the stream-K decode kernel for ragged batches.
  5. Variable-length batched prefill.
  6. Prefix caching/CoW: shared prefix blocks with reference counting and copy-on-write.

Validation

  • All nine F16/Q4_0/Q8_0 K/V combinations pass the numerical oracle on CUDA in partitioned and direct modes, including the narrower 3/1-heads-per-warp fallbacks and forced single-partition (max abs ≤ 1.2e-4 vs the 5e-4 bound; the increase over the pre-rework ~5e-6 is f16-partial rounding, below KV quantization error). Bench float64-oracle error ~5–8e-6, same magnitude as the contiguous kernel.
  • Generation-level A/B on Qwen3.6-27B (RTX 3090, greedy, 8 prompts × 4 configs): paged vs dense diverges on 3/8 long open-ended prompts — the same rate and prompts as the two controls (dense vs DFLASH_QWEN35_NO_KVPAD=1 dense, and paged vs GGML_CUDA_PAGED_ATTN_FORCE_PARTITIONS=1 paged), so divergence is greedy argmax FP chaos, not a kernel defect. All checkable answers correct; paged output run-to-run deterministic; decode tok/s at parity.
  • Persistent-metadata path exercised end-to-end: prefill → many block boundaries → pool reset across requests, including a 10K-token prompt.
  • Correctness covers empty sequences, page boundaries, shuffled blocks, long-context partitioning, and corrupt block-table bounds guards.
  • HIP gfx1151 (Radeon 8060S): all nine K/V combinations pass in partitioned and direct modes with identical error to CUDA, paged_kv_pool passes, and the full 4K/32K/128K plus ragged benchmark matrix completed. Note the native HIP contiguous path's own oracle error reaches 3.3e-3–4.5e-3 at long context (near the benchmark's 5e-3 gate) while paged stays at ~5e-6, so paged is also the numerically better path on this backend.
  • git diff --check passes.

🤖 Generated with Claude Code

https://claude.ai/code/session_01Vkkox5Yo66K2Rn1oPuVoX7

Review in cubic

@Graffioh
Graffioh force-pushed the codex/qwen36-paged-attention branch from 5f7b111 to e5027ab Compare July 23, 2026 09:27
@Graffioh Graffioh changed the title Add Qwen3.6 paged attention foundation feat(qwen35): add paged-attention foundation Jul 23, 2026
@Graffioh
Graffioh force-pushed the codex/qwen36-paged-attention branch from e5027ab to 5ff3bf7 Compare July 23, 2026 10:16
@Graffioh Graffioh changed the title feat(qwen35): add paged-attention foundation feat(qwen35): add paged-attention decode-only foundation Jul 23, 2026
@Graffioh Graffioh changed the title feat(qwen35): add paged-attention decode-only foundation feat(qwen35): add paged-attention monolithic decode-only foundation Jul 23, 2026
@Graffioh
Graffioh force-pushed the codex/qwen36-paged-attention branch 5 times, most recently from 51dda8f to 1a2e479 Compare July 24, 2026 10:12
@Graffioh
Graffioh marked this pull request as ready for review July 24, 2026 18:41

@cubic-dev-ai cubic-dev-ai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

All reported issues were addressed across 28 files

Reply with feedback, questions, or to request a fix.

Re-trigger cubic

Comment thread server/test/bench_paged_attention.cpp
Comment thread server/src/qwen35/qwen35_backend.cpp Outdated
Comment thread server/src/common/paged_kv_pool.cpp Outdated
Comment thread optimizations/paged_attention/README.md Outdated
Comment thread server/src/common/backend_factory.cpp Outdated
@Graffioh
Graffioh force-pushed the codex/qwen36-paged-attention branch 3 times, most recently from 2a36448 to f1851f9 Compare July 24, 2026 19:54
@Graffioh Graffioh changed the title feat(qwen35): add paged-attention monolithic decode-only foundation feat(qwen35): add paged-attention monolithic decode-only AR foundation Jul 25, 2026
Comment on lines -1858 to -1859
committed++;
cache_.cur_pos = committed;

@Graffioh Graffioh Jul 25, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

the first emitted output token from prefill writes KV row only during the subsequent forward (that's right)

BUT here on the first decode step, we increment committed by 1, even though for cache we are still considering the first emitted token from prefill, leading to a gap of 1

was there a reason for this?

(in case this needs to be fixed also in other model's backend)

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01QwBCUqGbQV4LSm6qUkZS3H
@Graffioh
Graffioh force-pushed the codex/qwen36-paged-attention branch 2 times, most recently from d91a8fe to ab11bb8 Compare July 25, 2026 11:10
Graffioh and others added 2 commits July 27, 2026 06:37
Rework the decode kernel: one warp now computes all GQA heads of a K/V
group per K/V row loaded (batched q8_1 dots), scores are tiled per 32
tokens with one skip-if-unchanged accumulator rescale, and softmax uses
exp2f with log2(e) folded into the Q prescale. Q is quantized once per
step by a small pre-pass instead of per (warp x partition), and
partition partials are stored normalized in f16, halving scratch
traffic. The block table and sequence length move into the persistent
target cache next to the K/V pool, so decode steps upload four bytes
per fresh block instead of the whole live table every token.

On an RTX 3090 (q4_0 K/V) vs the padded-contiguous fattn decode:
uniform 4K-32K batches go from +30..104% overhead to 10-22% faster,
128K to 21-51% faster, ragged 128K+7x8K to 29% faster. Short ragged
batches stay behind until the stream-K rework staged with continuous
batching.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01QwBCUqGbQV4LSm6qUkZS3H
Claude-Session: https://claude.ai/code/session_01Vkkox5Yo66K2Rn1oPuVoX7
@Graffioh
Graffioh force-pushed the codex/qwen36-paged-attention branch from ab11bb8 to e28f8d3 Compare July 27, 2026 06:37
Adopt the centralized feature/architecture gate from Luce-Org#558 and express
--paged-attention through it, deleting PagedAttentionOptions and
validate_paged_attention_options.

- model_capabilities.h gains a paged_attn column: Monolithic for qwen35,
  Never everywhere else. The arch and placement halves of the old rule now
  both read that row.
- The remaining rules (no draft/DDTree, --fa-window 0, no PFlash, no
  KVFlash, max_ctx bounds) move into check_feature_compatibility(), so the
  factory is the single admission owner for paged attention too.
- BackendFeatureConfig gains kvflash_enabled: KVFlash is requested through
  the environment and the gate is linked without the pool sizing code, so
  server_main resolves it the same way it resolves pflash_enabled.
- The "requires a CUDA or HIP build" rule is dropped rather than ported. It
  was already dead on both sides: backend_factory.cpp compiled it inside
  dflash_common, which CMake always gives DFLASH27B_BACKEND_CUDA or
  DFLASH27B_BACKEND_HIP, and server_main tested compiled_placement_backend()
  against Cuda/Hip, which are the only values it returns.
- paged_attn is cross-checked against Qwen35Config outside the
  DFLASH_CHECK_ARCH bundle, because qwen35moe shares that struct and the
  field-presence trait cannot tell the two apart. Its row is pinned by its
  own static_assert instead: Qwen35MoeBackend inherits the init() that
  builds the pool but overrides run_ar_decode_path() with a pipelined path
  that never reads a block table.
- Qwen35Backend::init() drops its copy of the rules but keeps a KVFlash
  guard, because create_backend(args) defaults kvflash_enabled to false and
  so skips that rule for every caller that is not the server.
- server_main keeps the prefix/prefill/disk-cache shutoff, now after
  prepare_backend(): it rewrites ServerConfig rather than rejecting a
  launch, which is not the gate's job.

paged_attention_config.h keeps only the block geometry.

Verified on an RTX 3090 with Qwen3.6-27B-Q4_K_M: all seven gate rules
reject as expected against real GGUFs; --paged-attention serves single and
multi-request chat completions and SSE streams with a correct 512x16 pool,
matching dense output on short prompts; --prefix-cache-slots is warned about
and zeroed. check: 22/22 tests pass.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01RKW2iJ94fznSMfkuvKQEaX
@Graffioh
Graffioh force-pushed the codex/qwen36-paged-attention branch from 235ef11 to f5a1078 Compare July 27, 2026 12:36

namespace dflash::common {

constexpr int PAGED_BLOCK_SIZE = 16;

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

we want this as user options or we fix as 16? (vllm default)

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.

1 participant