feat: Triattention support#356
Conversation
|
/gemini review |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request integrates the external TriAttention package with SGLang and vLLM serving engines in kvcached. It introduces patch files, documentation, auto-patching logic, and benchmark comparison scripts (compare_sglang_triattention.py and compare_vllm_triattention.py) to evaluate serving performance. The review feedback highlights critical issues in both benchmark scripts where the use of preexec_fn in subprocess.Popen within a multi-threaded context could lead to deadlocks, and potential exceptions during process startup could cause log file resource leaks. It is recommended to replace preexec_fn with start_new_session=True and wrap the process creation in a try...except block to ensure proper resource cleanup.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
…ero-output bench runs as failures
TriAttention's vLLM integration replaces GPUModelRunner._prepare_inputs with a
wrapper expecting `num_scheduled_tokens`, a calling convention that only exists
in vLLM >= 0.13.0. On vLLM <= 0.11.x the engine crashes on the first decode
step with a cryptic "missing 1 required positional argument: 'num_scheduled_tokens'"
and produces zero output tokens.
- triattention_external: check vllm.__version__ before installing the vLLM
hooks; raise a clear error below 0.13.0 (instead of crashing mid-decode) and
warn above the last validated version (0.21.0).
- compare_{vllm,sglang}_triattention: a stream that completes with 0 content
tokens is now a failure, not success=N/N — so a crashed engine no longer
reports a (misleadingly low) GPU peak as a win.
Verified: vLLM 0.11.0 now fails fast at startup with a clear message; vLLM
0.13.0 (SGLang 0.5.10) still activate and compress normally.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Summary
This PR adds optional external TriAttention integration for kvcached without vendoring the TriAttention source tree.
TriAttention is loaded from a separately installed checkout and is enabled only when
ENABLE_TRIATTENTION=1is set. Existing kvcached users are unaffected by default.What Changed
kvcached.integration.triattention_external.ENABLE_TRIATTENTION=1is requested but TriAttention activation fails.tools/compare_vllm_triattention.pytools/compare_sglang_triattention.pydecode-stressbenchmark profile, since SGLang TriAttention compression is decode-time driven.Differences between vLLM and SGLang Workloads
The vLLM benchmark uses a long-prompt workload because the vLLM integration can reclaim KV blocks through
block_pool.free_blocks(...), which is routed through kvcached'sElasticBlockPool. As a result, reclaim can directly reduce physical KV memory pressure and show up clearly in peak GPU memory.SGLang behaves differently. Its TriAttention compression path is decode-time driven, so a very long prompt alone mostly creates a prefill-time memory peak before compression has a chance to help. For SGLang, this PR uses a
decode-stressworkload: short prompt, long forced decode,min_tokens=max_tokens, andignore_eos=true. This makes compression and reclaim happen while decode memory is growing, which better exposes the effect in GPU peak memory.Benchmark Results
Model: Qwen3-8B
GPU memory sampled with
nvidia-smi.vLLM Long-Context Workload
Workload:
The vLLM path shows a large peak GPU memory reduction, especially at higher concurrency, because reclaimed KV blocks are returned through kvcached's physical allocator path.
SGLang Decode-Stress Workload
Workload:
decode-stressSGLang peak GPU memory was reduced by 3660 MiB, or 14.3%.
These happen when a request has already been released by the time a late compression attempt is observed. The successful compression and reclaim path is confirmed by compression_events=16 and freed_blocks=34316.
Validation
pre-commit run --all-files
Passed.
Additional checks:
vLLM benchmark CSV parsed successfully.
SGLang decode-stress benchmark CSV parsed successfully.
SGLang benchmark script syntax check passed.