[TRTLLM-16005][feat] Store DFlash draft cross-attention context paged in the target KV cache manager - #16150
[TRTLLM-16005][feat] Store DFlash draft cross-attention context paged in the target KV cache manager#16150chungen04 wants to merge 10 commits into
Conversation
📝 WalkthroughWalkthroughThis PR adds an optional hybrid-context mode to DFlash speculative decoding that stores draft cross-attention K/V in the target model's paged KV-cache pools instead of dedicated dense buffers. It includes a new Triton paged cross-attention kernel, KV-cache sizing/head-mapping updates, DFlashWorker pool-write wiring, and a config flag. ChangesDFlash Hybrid Context Storage
Estimated code review effort: 5 (Critical) | ~120 minutes Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (1)
tensorrt_llm/_torch/speculative/dflash_hybrid_attn.py (1)
357-391: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd return annotations and a public entrypoint docstring.
The new non-JIT Python functions should be fully annotated, and
dflash_ctx_paged_attentionis imported outside this file, so prefer a Google-style docstring over inline argument comments. As per coding guidelines, “Always annotate functions” and “Use Google style docstrings for classes and functions in Python.” Based on learnings, Python 3.10+ type syntax is acceptable here.Also applies to: 450-458
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/speculative/dflash_hybrid_attn.py` around lines 357 - 391, The new Python helpers are missing explicit return annotations and the public entrypoint lacks a proper docstring. Add Python 3.10+ type hints, including return types, to the non-JIT helpers like _get_partial_bufs and _num_splits, and replace the inline argument comments on dflash_ctx_paged_attention with a Google-style docstring that describes its arguments and return value. Keep the function names and signatures easy to locate in the speculative attention module, and apply the same documentation/annotation treatment to the related block around the later context attention helper.Sources: Coding guidelines, Learnings
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tensorrt_llm/_torch/speculative/dflash_hybrid_attn.py`:
- Around line 357-366: The `_get_partial_bufs` helper in `dflash_hybrid_attn`
uses the ambiguous single-letter buffer name `l`, which triggers Ruff E741.
Rename that local buffer to a clearer name like `logsumexp` or `lse`, and update
the returned tuple and `_PARTIAL_BUFS` assignment so the new identifier is used
consistently wherever the partial buffers are created.
- Around line 72-78: The Triton kernel setup in dflash_hybrid_attn uses
tl.arange spans that must be powers of two, but the current validation only
checks for a minimum size, so non-power-of-two values for D or R will still
compile-fail. Update the shape validation and launch path around the tl.arange
calls and the speculative attention kernel entry points to either reject
non-power-of-two head_dim / Q*group sizes early or pad them to the next power of
two before invoking the Triton kernels.
---
Nitpick comments:
In `@tensorrt_llm/_torch/speculative/dflash_hybrid_attn.py`:
- Around line 357-391: The new Python helpers are missing explicit return
annotations and the public entrypoint lacks a proper docstring. Add Python 3.10+
type hints, including return types, to the non-JIT helpers like
_get_partial_bufs and _num_splits, and replace the inline argument comments on
dflash_ctx_paged_attention with a Google-style docstring that describes its
arguments and return value. Keep the function names and signatures easy to
locate in the speculative attention module, and apply the same
documentation/annotation treatment to the related block around the later context
attention helper.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: b8f5735c-2def-4f0e-8aa4-6f7737f54632
📒 Files selected for processing (6)
tensorrt_llm/_torch/models/modeling_speculative.pytensorrt_llm/_torch/pyexecutor/_util.pytensorrt_llm/_torch/speculative/dflash.pytensorrt_llm/_torch/speculative/dflash_hybrid_attn.pytensorrt_llm/_torch/speculative/utils.pytensorrt_llm/llmapi/llm_args.py
| d = tl.arange(0, D) | ||
| tok = tl.arange(0, BLOCK_N) | ||
| for n0 in range(start, end, BLOCK_N): | ||
| idx = n0 + tok | ||
| valid = idx < end | ||
| # Per-token page gather: page ids come from the block table row. | ||
| page = tl.load(blk_row_ptr + idx // TPB, mask=valid, other=0).to(tl.int64) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy lift
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Inspect whether tests/configs cover non-power-of-two DFlash tile shapes.
rg -n --type=py \
'dflash_ctx_paged_attention|dflash_ctx_paged_attention_ref|max_draft_len|head_dim|num_key_value_heads|num_attention_heads' \
tensorrt_llm testsRepository: NVIDIA/TensorRT-LLM
Length of output: 50376
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Map the relevant file and inspect the shape guards and tl.arange uses.
ast-grep outline tensorrt_llm/_torch/speculative/dflash_hybrid_attn.py --view expanded
echo '---'
nl -ba tensorrt_llm/_torch/speculative/dflash_hybrid_attn.py | sed -n '50,170p'
echo '---'
nl -ba tensorrt_llm/_torch/speculative/dflash_hybrid_attn.py | sed -n '330,430p'
echo '---'
# Find any explicit power-of-two padding or next_power_of_2 usage in this module.
rg -n 'next_power_of_2|power of two|power-of-two|tl\.arange|BLOCK_|R =|D =' tensorrt_llm/_torch/speculative/dflash_hybrid_attn.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 654
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo '== relevant lines =='
sed -n '340,430p' tensorrt_llm/_torch/speculative/dflash_hybrid_attn.py
echo '== tl.arange uses =='
grep -n 'tl.arange\|next_power_of_2\|power' tensorrt_llm/_torch/speculative/dflash_hybrid_attn.py
echo '== shape/guard references in module =='
grep -n 'Q \* group\|D >=\|R >=\|BLOCK_\|tokens_per_block\|num_splits\|partial' tensorrt_llm/_torch/speculative/dflash_hybrid_attn.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 5385
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo '== top-of-file context =='
sed -n '1,60p' tensorrt_llm/_torch/speculative/dflash_hybrid_attn.py
echo '== surrounding kernels with tl.arange =='
sed -n '60,180p' tensorrt_llm/_torch/speculative/dflash_hybrid_attn.py
echo '---'
sed -n '180,340p' tensorrt_llm/_torch/speculative/dflash_hybrid_attn.py
echo '== call sites / tests using dflash_ctx_paged_attention =='
rg -n 'dflash_ctx_paged_attention\(' tensorrt_llm testsRepository: NVIDIA/TensorRT-LLM
Length of output: 11827
🌐 Web query:
Triton tl.arange power of two endpoint requirements documentation
💡 Result:
In Triton, the tl.arange(start, end) function has specific requirements regarding its arguments and the resulting range [1]. Key requirements include: 1. Power of Two Requirement: The difference between the end and start arguments (the range) must be a power of two [1]. Specifically, if the value (end - start) is not a power of two, Triton will raise a ValueError [1]. Note that official documentation snippets may sometimes generalize or misstate this as a requirement for the start and end values themselves being powers of two [2], but the functional implementation specifically enforces that the range (end - start) must be a power of two [1]. 2. Constexpr Requirement: Both the start and end arguments must be tl.constexpr [1]. If you pass runtime values to tl.arange, the kernel compilation will fail with a ValueError [3][4][1]. Because tl.arange requires constexpr arguments, any block size or range parameter used within it must also be defined as tl.constexpr within the kernel signature [5]. 3. Integer Constraints: The arguments must fit within a 32-bit integer, and the end value must be strictly greater than the start value [1]. For best practices, ensure that any block size or range values used in indexing are marked as tl.constexpr in your Triton kernel signature to avoid compilation errors [5]. Top results: [2] https://triton-lang.org/main/python-api/generated/triton.language.arange.html [1] https://github.com/triton-lang/triton/blob/505bdb74/python/triton/language/semantic.py
Citations:
- 1: https://github.com/triton-lang/triton/blob/505bdb74/python/triton/language/semantic.py
- 2: https://triton-lang.org/main/python-api/generated/triton.language.arange.html
- 3: ValueError("arange's arguments must be of type tl.constexpr") triton-lang/triton#2735
- 4: ValueError("arange's arguments must be of type tl.constexpr") Error On Example Code triton-lang/triton#2604
- 5: [Question] What parameters should be
tl.constexprand are there guidelines for parallelizing? triton-lang/triton#1087
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo '== call site context =='
sed -n '1350,1405p' tensorrt_llm/_torch/models/modeling_speculative.py
echo '== any draft/block size constraints near the call site =='
grep -n 'max_draft_len\|draft\|block_size\|Q =' -n tensorrt_llm/_torch/models/modeling_speculative.py | head -n 50
echo '== any explicit power-of-two checks in dflash module =='
grep -n 'power\|pow2\|next_power_of_2\|assert .*16' tensorrt_llm/_torch/speculative/dflash_hybrid_attn.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 6433
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
from pathlib import Path
path = Path("tensorrt_llm/_torch/models/modeling_speculative.py")
text = path.read_text()
for needle in ["dflash_ctx_paged_attention(", "max_draft_len", "draft", "block_size", "speculative"]:
idx = text.find(needle)
print(f"NEEDLE {needle!r}: {idx}")
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 306
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo '== DFlash-related config classes / model usage =='
rg -n 'DFlash|dflash|head_dim|num_attention_heads|num_key_value_heads|qk_nope_head_dim|qk_rope_head_dim|v_head_dim' tensorrt_llm/_torch/models tensorrt_llm/_torch/speculative
echo '== specific DFlash draft model definitions =='
rg -n 'class .*DFlash|def .*DFlash|DFlashDraftModel|dflash_ctx_paged_attention' tensorrt_llm/_torch/models tensorrt_llm/_torch/speculative
echo '== head_dim constraints in tests for DFlash / hybrid speculative paths =='
rg -n 'DFlash|hybrid_block_idx|cache_seqlens|draft model|head_dim.*(64|72|128)|qk_nope_head_dim|v_head_dim' tests/unittest/_torch tests/unittest/othersRepository: NVIDIA/TensorRT-LLM
Length of output: 50376
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo '== DFlash-specific definitions and constraints =='
rg -n 'DFlashDraftModel|hybrid_k_bufs|dflash_ctx_paged_attention|head_dim|num_attention_heads|num_key_value_heads|qk_nope_head_dim|v_head_dim' \
tensorrt_llm/_torch/models tensorrt_llm/_torch/speculative | head -n 200
echo '== nearby DFlash model class / init context =='
sed -n '820,920p' tensorrt_llm/_torch/models/modeling_speculative.py
echo '== any explicit size constraints in DFlash/speculative docs =='
rg -n 'power of two|pow2|next_power_of_2|supported head_dim|supported.*head_dim|DFlash' \
tensorrt_llm/_torch/models tensorrt_llm/_torch/speculative | head -n 200Repository: NVIDIA/TensorRT-LLM
Length of output: 34992
Pad the tl.arange spans to powers of two.
tl.arange(0, D) and tl.arange(0, R) require a power-of-two span, but this path only checks >= 16. Non-power-of-two head_dim or Q * group will fail kernel compilation, so pad or reject those shapes before launching the Triton kernels.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tensorrt_llm/_torch/speculative/dflash_hybrid_attn.py` around lines 72 - 78,
The Triton kernel setup in dflash_hybrid_attn uses tl.arange spans that must be
powers of two, but the current validation only checks for a minimum size, so
non-power-of-two values for D or R will still compile-fail. Update the shape
validation and launch path around the tl.arange calls and the speculative
attention kernel entry points to either reject non-power-of-two head_dim /
Q*group sizes early or pad them to the next power of two before invoking the
Triton kernels.
| def _get_partial_bufs(B, NKV, S, R, D, device): | ||
| key = (B, NKV, S, R, D, device) | ||
| bufs = _PARTIAL_BUFS.get(key) | ||
| if bufs is None: | ||
| acc = torch.empty(B * NKV * S * R * D, dtype=torch.float32, device=device) | ||
| m = torch.empty(B * NKV * S * R, dtype=torch.float32, device=device) | ||
| l = torch.empty(B * NKV * S * R, dtype=torch.float32, device=device) | ||
| bufs = (acc, m, l) | ||
| _PARTIAL_BUFS[key] = bufs | ||
| return bufs |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Rename the ambiguous l buffer.
Ruff flags Line 363 with E741. Rename it to avoid blocking lint.
Proposed fix
- l = torch.empty(B * NKV * S * R, dtype=torch.float32, device=device)
- bufs = (acc, m, l)
+ l_buf = torch.empty(B * NKV * S * R, dtype=torch.float32, device=device)
+ bufs = (acc, m, l_buf)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| def _get_partial_bufs(B, NKV, S, R, D, device): | |
| key = (B, NKV, S, R, D, device) | |
| bufs = _PARTIAL_BUFS.get(key) | |
| if bufs is None: | |
| acc = torch.empty(B * NKV * S * R * D, dtype=torch.float32, device=device) | |
| m = torch.empty(B * NKV * S * R, dtype=torch.float32, device=device) | |
| l = torch.empty(B * NKV * S * R, dtype=torch.float32, device=device) | |
| bufs = (acc, m, l) | |
| _PARTIAL_BUFS[key] = bufs | |
| return bufs | |
| def _get_partial_bufs(B, NKV, S, R, D, device): | |
| key = (B, NKV, S, R, D, device) | |
| bufs = _PARTIAL_BUFS.get(key) | |
| if bufs is None: | |
| acc = torch.empty(B * NKV * S * R * D, dtype=torch.float32, device=device) | |
| m = torch.empty(B * NKV * S * R, dtype=torch.float32, device=device) | |
| l_buf = torch.empty(B * NKV * S * R, dtype=torch.float32, device=device) | |
| bufs = (acc, m, l_buf) | |
| _PARTIAL_BUFS[key] = bufs | |
| return bufs |
🧰 Tools
🪛 Ruff (0.15.20)
[error] 363-363: Ambiguous variable name: l
(E741)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tensorrt_llm/_torch/speculative/dflash_hybrid_attn.py` around lines 357 -
366, The `_get_partial_bufs` helper in `dflash_hybrid_attn` uses the ambiguous
single-letter buffer name `l`, which triggers Ruff E741. Rename that local
buffer to a clearer name like `logsumexp` or `lse`, and update the returned
tuple and `_PARTIAL_BUFS` assignment so the new identifier is used consistently
wherever the partial buffers are created.
Source: Linters/SAST tools
chungen04
changed the title
chungen04
changed the title
…ontext Add DFlashDecodingConfig.use_hybrid_context (prototype). When enabled, the draft cross-attention context layers are registered as spec layers of the target KV cache manager: get_num_spec_layers reports the draft layer count and KvCacheCreator budgets the extra per-token cost on the last PP rank, instead of allocating a dense per-request context buffer. Signed-off-by: chungen04 <cho322@gatech.edu>
Read the draft context K/V directly from the target KV cache manager's paged pool. Existing kernels do not fit: sequence lengths are device-resident (acceptance-dependent, updated inside CUDA graphs), pages are target-pool sized rather than 256-token aligned, and the dense noise suffix must be fused into the same softmax. Online-softmax with bf16 tensor-core dots and fp32 accumulation; supports fp8 context storage. Includes a torch reference implementation for testing. Signed-off-by: chungen04 <cho322@gatech.edu>
…ache manager When use_hybrid_context is enabled, the DFlash worker scatters the projected context K/V into the draft spec layers of the target manager's paged pool (position-aligned, so prefix-cache hits restore context), and the drafter reads it back with the paged cross-attention kernel instead of flash_attn_with_kvcache over a dense per-request buffer. Dense-path behavior is unchanged when the flag is off. Signed-off-by: chungen04 <cho322@gatech.edu>
The single-pass kernel launched (batch, num_kv_heads) CTAs and marched one 32-token page per tl.dot iteration; at small batch that is ~4 CTAs on a 148-SM GPU and cost ~+1 ms/step vs the dense-buffer flash_attn path (measured -24% output throughput at concurrency 1, Qwen3-8B TP2). Rework it flash-decoding style: - Small batch: split-KV. The context is partitioned across S CTAs that emit partial softmax states; a merge kernel folds the partials plus the dense noise suffix. S derives from the launch batch size (static under CUDA graph capture); per-split ranges derive from the device-resident ctx_len, so the fixed grid stays graph-safe. Partial buffers are cached per shape for stable capture addresses. - Large batch: keep the single-pass path (no partial round-trip). - Both paths tile BLOCK_N=128 context tokens (4 pages) per tl.dot with a per-token page-id gather instead of one page per iteration. Kernel time at B=1 (ctx 8192, 4 KV heads/rank): 0.284 ms -> 0.063 ms, now faster than dense flash_attn_with_kvcache (0.077 ms). End-to-end aiperf sweep (Qwen3-8B TP2, ISL 8000/OSL 512): hybrid ctx goes from -24%..-7% vs the dense baseline to +2% (c=1), +8% (c=8), +4..5% (c=16..64), with bit-identical greedy outputs. Signed-off-by: chungen04 <cho322@gatech.edu>
The split count S was a tl.constexpr derived from the launch batch size. Mixed context+generation iterations run eagerly at their actual (unpadded) batch size, so serving traffic kept discovering new S values and paying a blocking Triton JIT compile mid-request: p90 ITL jumped 3-4x and TTFT tails reached 13-17 s on the first sweep levels after a server start, disappearing once the in-process compile cache warmed (which made it look like measurement noise). S only enters index arithmetic and the merge loop trip count, so pass it as a runtime argument: each kernel compiles exactly once and no JIT ever runs on the serving path. Kernel parity and timings unchanged. Signed-off-by: chungen04 <cho322@gatech.edu>
…from the target The KV cache manager has a single head_dim per pool, so draft spec layers inherited the target's head_dim and the context store failed with a shape mismatch whenever the two differ (e.g. a 128-head_dim draft on a 256-head_dim target). Register draft spec layers in target-head_dim units (heads scaled so bytes stay exact) and view the pool buffers back to draft geometry in the worker; pure view, no copy. Raises a clear error when the draft KV row is not expressible in target head_dim units. Signed-off-by: chungen04 <cho322@gatech.edu>
…a targets The hybrid-ctx cost path passed the draft's plain-attention config into the target's Mamba-hybrid cache manager class, raising 'Qwen3Config is not a supported hybrid Mamba config' at startup on GDN/Mamba-hybrid targets (e.g. Qwen3.5). Resolve the manager class from the draft config instead, matching the external-drafter path. Signed-off-by: chungen04 <cho322@gatech.edu>
Signed-off-by: chungen04 <b09901027@ntu.edu.tw>
af69173 to
5df8b51
Compare
…e paged attention kernel tl.arange requires a power-of-two span, but the query-row count R = block_size x GQA group was used directly, so draft lengths not of the form 2^k - 1 (e.g. max_draft_len=5 -> R=24) or non-power-of-two GQA ratios failed Triton compilation at warmup. Every published DFlash draft happens to be power-of-two clean, which is why this never fired in benchmarks. Pad rows to R_PAD = next_power_of_2(R) with masked query loads and masked output stores; padded rows carry zeros through the softmax pipeline and are discarded at the store, so results are exact. The partial buffers of the split path are R_PAD-strided. head_dim and tokens_per_block keep explicit power-of-two checks with actionable errors (no known model needs padding there). Power-of-two shapes compile with the masks folded away, so the existing fast path is unchanged; kernel parity verified for R=24, R=20 (group 5), and R=6 against the torch reference, including CUDA-graph replay. Signed-off-by: chungen04 <cho322@gatech.edu>
|
Hi reviewers, this PR solves an issue on DFlash with higher serving concurrency. Please take a look, thank you!
|
JunyiXu-nv
left a comment
JunyiXu-nv
left a comment
There was a problem hiding this comment.
Please add unit tests for the new kernel and feature. Thanks!
There was a problem hiding this comment.
These newly added fields should also be added to the stability tests' yaml list.
There was a problem hiding this comment.
You've shipped a torch reference implementation specifically for parity testing, but nothing in this PR actually calls it. This kernel is exactly the kind of code that needs parity tests: 630 lines, two execution paths (single-pass vs split+merge), fp8, and several padding edge cases. Could you add a few cases to tests/unittest/_torch/speculative/hw_agnostic/test_dflash.py? At minimum: ctx_len=0, ctx_len not page-aligned, empty splits when S>1, and a shape like Q=6 where R_PAD padding kicks in. Otherwise this ref function is dead code.
There was a problem hiding this comment.
This is a new user-visible knob (prototype or not), but not any mention of it under docs/. Please add a short section to the DFlash part of the speculative-decoding docs covering the prerequisites
There was a problem hiding this comment.
This new hybrid branch is nearly line-for-line identical to the separate-draft branch at L431-437. Can we fold them into a single "draft layers need to be charged" condition?
There was a problem hiding this comment.
set_max_total_draft_tokens now also flips _allow_separate_draft_kv_cache off as a side gig, so the name no longer matches what it does. Please rename.
There was a problem hiding this comment.
check spec_dec_mode.is_dflash() first and then still do getattr(sc, 'use_hybrid_context', False) — but is_dflash() already guarantees this is a DFlashDecodingConfig, so the field always exists. Just use sc.use_hybrid_context directly.
|
closing as duplicate with #18343 |
3 participants
@coderabbitai summary
Description
Closes #16005
In #16005, it was described that current design for DFlash cross attention allocates large GPU buffer and can limit the admitting requests during serving. Reflecting vLLM's design, the key was that the hidden state from target injected as KV cache in the drafter should be allocated paged rather than a contiguous worst-case buffer.
This PR stores the draft context inside the target KV cache manager. The draft layers register as spec layers of the target manager, so the context lives in ordinary paged blocks and its memory scales with live requests. The projected context K/V is written with norm and RoPE already applied, and keeps the store path CUDA-graph safe.
A kernel (
tensorrt_llm/_torch/speculative/dflash_hybrid_attn.py) performs the paged cross-attention read. The per-token cost of the draft layers is added to the KV size estimation so the scheduler budgets for it. The feature is gated behind a new use_hybrid_context flag on DFlashDecodingConfig and defaults to off (can be set to default). The kernel uses a flash-decoding style split-KV path for small batches, where the context is partitioned across several CTAs that emit partial softmax states. Split ranges derive from the device-resident context lengths, so the fixed grid stays CUDA-graph safe. Both paths read 128 context tokens.Isolated kernel timings at the DFlash decode shape (per TP 2 rank, 16 query heads, 4 KV heads, head_dim 128, 16 queries, context 8192) compare as follows against reading the dense buffer with
flash_attn_with_kvcache.Two compatibility fixes ride along:
Results
aiperf. The workload is zai-org/LongBench , sized with the model tokenizer to about 16k and 32k prompt tokens, decoded greedily with natural stopping capped at 1024 new tokens.Testing results on serving models with different context and concurrency:
Observing the Prometheus counters, metrics related to speculative decoding shows the correctness of the implementation reflecting the original:
The paged path reads only live context pages while the dense path strides over the worst-case per-request buffer, so the advantage grows with context length and batch size on top of the memory savings.
Limitations
Test coverage
PR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
If PR introduces API changes, an appropriate PR label is added - either
api-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin the PR title.Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.