Skip to content

DFlash2 leaves the paged graph fast path at every c>1 and re-attends the whole batch context every decode step, which is the c=4 -> c=8 scaling stall #2087

Description

@localai-bot

At every concurrency above one the DFlash2 draft leaves the paged CUDA-graph
fast path and runs a materialized re-attention over the whole batch's
context, every decode step
. This is a change of algorithmic complexity, not a
constant factor, and it is the shape of the measured c=4 -> c=8 stall.

The measured shape it explains

main after #1994/#1997/#2000/#2010, dgx:gpu0 (GB10), idle under rc,
Qwen3.8-27B NVFP4 + DFlash2 k=8, 1024 in / 512 out, --max-num-seqs 16,
vllm bench serve --dataset-name random --backend openai-chat:

c ours out tok/s vLLM SGLang
1 25.14 24.36 25.20
2 40.12 38.59 44.93
4 60.25 64.25 77.06
8 63.3 80.0 109.24

We LEAD at c=1 -- the one rung that takes the fast path -- and the deficit opens
from c=2 on, exactly where the route changes.

The route change

runner.cpp:3295 is the only production caller of
Qwen3DFlashModel::ForwardBlockLogitsWithDeviceKV, and it passes one store per
PROPOSING ROW (runner.cpp:3260-3273), so P == num proposing rows.

The fast path is gated on P == 1 (qwen3_dflash.cpp:1577). That branch runs
ForwardPagedBody: Tq = 1+k = 9 query rows against the request's persistent
paged store, through the shared paged seam, CUDA-graph captured and replayed
(qwen3_dflash.cpp:1577-1604, :1461-1508).

At P > 1 control reaches the fallback at qwen3_dflash.cpp:1888-1930, whose
own comment says it is "not capture-targeted". Per decode step it:

  1. allocates 2 x L context buffers of [C, kdim], C = sum of every proposing row's full context length (qwen3_dflash.cpp:1892-1895);
  2. gathers every context row out of every request's paged store with
    4 x P x L IndexSelect/IndexCopy launches (:1908-1922);
  3. calls ForwardWithCtxKVDev (qwen3_dflash.cpp:664), which per layer
    allocates qcb/kcb/vcb and acomb of [Ncomb, ...] where
    Ncomb = C + Tq (:792-794, :811), memsets qcb (:795), and calls
    vt::DFlashBlockAttention (:818).

The cost

vt::DFlashBlockAttention's CUDA grid is over t = query.shape[0], i.e.
over all Ncomb rows (cuda_ops.cu:2634, :2643, :2650). The draft
therefore computes an attention output for every context row of every request
in the batch
, and then throws all of them away -- qwen3_dflash.cpp:820-827
IndexSelects only the Tq block rows back out.

Every draft layer of the benchmarked checkpoint is non-causal
(z-lab/Qwen3.8-27B-DFlash2 config.json: is_causal false, 5 layers, hidden
5120; .agents/specs/dflash2-spec-decode.md), so jhi = qe - qs - 1 and the
window is not applied (cuda_ops.cu:1580-1582, :2402-2409, :2435; see the
separate window issue). Each of the Ncomb query rows therefore attends over
its own request's ENTIRE span.

Work per draft step, per layer:

route attention pairs
P == 1 paged (1+k) x C
P > 1 fallback sum_r (ctx_r + 1 + k)^2

At ctx_r ~ 1300 that is a ~150x blow-up per row, on top of ~O(C) bytes of
gather, memset and index traffic per layer. It enters at c=2 and grows with c
because C does.

The fallback also loses the CUDA-graph lane entirely: the capture and replay live
inside the P == 1 branch (qwen3_dflash.cpp:1602-1604, :1862-1886).

What this is NOT

#2010's author flagged qwen3_dflash.cpp:1888-1927 as a per-step device
allocation
burst. The allocations are not the cost: DBuf draws from the
shared size-class DevicePool (dense_device_glue.h:109-127,
device_pool.h), which is UNCAPPED on GB10
(platforms/interface.h:118), so these are pool hits after warm-up and not
cudaMalloc/cudaFree syncs. The cost is the WORK the fallback performs.

Scope

Batch the paged draft propose so P > 1 takes the same paged, block-only,
graph-captured route P == 1 takes: attention computed for P x (1+k) query
rows against per-request paged context, never for the context rows themselves.

Two shapes, both plausible:

  • Narrow. Keep the materialized combined K/V, but give
    vt::DFlashBlockAttention a separate query cu so Q is [Tq, ...] while
    K/V stay [Ncomb, ...]. This deletes the Ncomb-sized qcb/acomb, the
    memset, the two IndexCopy/IndexSelect round trips and ~99% of the
    attention work, and touches one op plus its CPU and CUDA kernels.
  • Full. One shared paged pool for every request's draft context (which is
    also what Attention and recurrent state are allocated from two separate pools, so c=32 at k=8 is unservable where vLLM and SGLang serve it #2007 needs) plus a batched block table and per-request seq_lens,
    so the batched propose is the P == 1 path with num_reqs > 1. This
    additionally deletes the whole O(C) gather.

Per-request looping over the P == 1 path is NOT a fix: it re-reads the draft's
~1.5 GB of weights and the ~0.72 GB packed head once per row.

Owed evidence

A device A/B on one binary at c=4 and c=8 with VT_SPEC_TRACE=2, whose
[spec-phase-dev] fwd= segment is the quantity this issue claims dominates, plus
the throughput ladder rerun.

Owner: SPEC-DFLASH2.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions