refactor(qwen): rename int8-tq3 → tq3-* family + add no-MTP + Genesis variants
Compose layout for Qwen 3.6 27B + TQ3 KV now reflects the three real operating points after today's TQ3+MTP investigation: - dual/tq3-mtp.yml — TQ3+MTP attempt without Genesis, TOMBSTONED. Needs 4 of 5 missing upstream PRs (Genesis P64/P65/P66/P68/P69 equivalents); only PR #40914 has a community analog. Re-test when upstream catches up. - dual/tq3-nomtp.yml — TQ3 without MTP. Validated working on pure upstream nightly: 1.73M KV pool, 6.59× concurrency at 262K, verify-stress 7/7 pass. The deployable Genesis-free TQ3 path today. - dual/tq3-mtp-genesis.yml — TQ3+MTP via Genesis, matched-config 2-stream sibling of dual/turbo.yml's 4-stream production compose. Plus retained research artifacts in patches/: - vllm-pr40798-rebased/ — partial upstream workspace fix - vllm-pr40914-k1-only/ — manually-rebased K+1 dispatch (still has open question per vllm-issue#40880 closure history) Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>
This commit is contained in:
noonghunna
co-authored by
Claude Opus 4.7
Claude Opus 4.7
parent
a3b66c489c
commit
6182922225
@@ -0,0 +1,267 @@
|
||||
# ===========================================================================
|
||||
# Profile (at-a-glance):
|
||||
# Model: Qwen3.6-27B (Lorbus AutoRound INT4 + BF16 mtp.fc preserved)
|
||||
# Topology: Dual 3090 PCIe (TP=2, no NVLink)
|
||||
# Drafter: MTP n=3 (built-in)
|
||||
# KV: turboquant_3bit_nc (TQ3, 0.375 bytes/token)
|
||||
# Vision: yes
|
||||
# Max ctx: 262K
|
||||
# Streams: max-num-seqs=2 (matched-config sibling of dual/int8.yml)
|
||||
# Genesis: v7.72.2 (full P64/P65/P66/P68/P69 fix for TQ3+MTP — see vllm-issue#40880)
|
||||
# Status: ✅ Working (with Genesis)
|
||||
# Best for: Apples-to-apples 2-stream comparison against dual/int8.yml at the
|
||||
# same vLLM pin + matched ctx/seqs settings. For 4-stream serving
|
||||
# throughput, use dual/turbo.yml (which this file is derived from).
|
||||
# ---------------------------------------------------------------------------
|
||||
# Why this exists (separate from dual/turbo.yml):
|
||||
#
|
||||
# dual/turbo.yml = 4 streams, TQ3, MTP n=3, Genesis. Production multi-tenant.
|
||||
# This file = 2 streams, TQ3, MTP n=3, Genesis. Matched-config for the
|
||||
# head-to-head matrix vs dual/int8.yml (Qwen INT8 PTH 2 seqs).
|
||||
#
|
||||
# Why NOT dual/tq3-mtp.yml (Genesis-free attempt):
|
||||
# TQ3 + MTP on Qwen3-Next needs 5 upstream fixes. Only 1 of 5 has a
|
||||
# community PR analog today (vllm#40914). The other 4 ship as Genesis
|
||||
# patches P64/P65/P66/P68/P69. The Genesis-free path is tombstoned in
|
||||
# dual/tq3-mtp.yml — re-test when the upstream PR landscape catches up.
|
||||
#
|
||||
# ---------------------------------------------------------------------------
|
||||
# Dual-card Turbo — TP=2 + TurboQuant KV (turboquant_3bit_nc) + MTP n=3 + Genesis v7.69 dev.
|
||||
#
|
||||
# 2026-05-02 PM note: setup.sh now pins Genesis to v7.69 (commit 2db18df) which
|
||||
# closes the Cliff 1 mech B issue this header used to warn about (PN25 worker
|
||||
# self-install lands default-on in v7.69). The patch_workspace_lock_disable
|
||||
# sidecar mounted below is now redundant on v7.69 — Genesis PN34 covers the
|
||||
# same surface via env-gate. Drop is queued for the next dual-card validation
|
||||
# pass (Cliff 2 doesn't apply on TP=2 so the mount is harmless meanwhile).
|
||||
#
|
||||
# The 4-stream concurrent-serving variant. Trades ~22% per-stream TPS vs the
|
||||
# fp8 default for ~9× the KV pool capacity, unlocking 4 concurrent streams at
|
||||
# full 262K context. Aggregate throughput exceeds fp8 above ~3 concurrent users.
|
||||
#
|
||||
# Measured (re-bench 2026-04-28 on club-3090, dev205 + Genesis v7.51-stable, TQ3 KV):
|
||||
# 53.65 narr (CV 2.7%) / 72.93 code (CV 2.7%) TPS, AL 3.41, VRAM 24.1 GB/card.
|
||||
#
|
||||
# Genesis v7.14 P65 (cudagraph downgrade for spec-decode) makes MTP × TurboQuant
|
||||
# work on dual-card TP=2 with vision + tools + 262K. Other relevant patches:
|
||||
# - P4 hybrid turboquant support (replaces standalone PR #39931 patches)
|
||||
# - P5 KV page-size unification for hybrid models
|
||||
# - P64 streaming MTP tool-call edge case
|
||||
# - P66 cudagraph_capture_sizes divisibility filter
|
||||
# Plus patch_tolist_cudagraph.py for the .tolist() crash class.
|
||||
#
|
||||
# KV cache: turboquant_3bit_nc — 3-bit symmetric K and V (~3 bits average per
|
||||
# token). Aligned with the single-card v714 default config we test extensively.
|
||||
# The predecessor dual-3090 file shipped `turboquant_k8v4` (~6 bits avg) due to
|
||||
# config drift from the documented + tested target; corrected here to match
|
||||
# what's actually validated. Cross-rig users may see slightly higher concurrency
|
||||
# at 262K than the predecessor's published 4.59× (TQ3 is more compact than k8v4).
|
||||
#
|
||||
# To run:
|
||||
# cd <repo>/models/qwen3.6-27b/vllm/compose
|
||||
# docker compose -f dual/tq3-mtp-genesis.yml up -d
|
||||
# ===========================================================================
|
||||
services:
|
||||
vllm-qwen36-27b-dual-tq3-mtp-genesis:
|
||||
image: vllm/vllm-openai:nightly-1acd67a795ebccdf9b9db7697ae9082058301657
|
||||
container_name: vllm-qwen36-27b-dual-tq3-mtp-genesis
|
||||
restart: "no"
|
||||
ports:
|
||||
- "${BIND_HOST:-0.0.0.0}:${PORT:-8015}:8000"
|
||||
volumes:
|
||||
- ${MODEL_DIR:-../../../../../models-cache}:/root/.cache/huggingface
|
||||
# torch.compile + Triton kernel caches — first boot warms (~60-90 sec);
|
||||
# subsequent boots reuse cached graphs. Pattern from Sander's PROD launch.
|
||||
# Closes club-3090 #22.
|
||||
- ../../cache/torch_compile:/root/.cache/vllm/torch_compile_cache
|
||||
- ../../cache/triton:/root/.triton/cache
|
||||
# Marlin pad-sub-tile-n (vLLM PR #40361) — vendored in this repo at
|
||||
# ../../patches/vllm-marlin-pad/. Drops out when vllm#40361 lands upstream.
|
||||
- ../../patches/vllm-marlin-pad/marlin.py:/usr/local/lib/python3.12/dist-packages/vllm/model_executor/kernels/linear/mixed_precision/marlin.py:ro
|
||||
- ../../patches/vllm-marlin-pad/MPLinearKernel.py:/usr/local/lib/python3.12/dist-packages/vllm/model_executor/kernels/linear/mixed_precision/MPLinearKernel.py:ro
|
||||
# Genesis modular package (auto-checked-out at pin in scripts/setup.sh).
|
||||
# As of v7.72.2: PN35 supersedes patch_inputs_embeds_optional.py,
|
||||
# P78 supersedes patch_tolist_cudagraph.py, PN34 supersedes
|
||||
# patch_workspace_lock_disable.py — those local sidecars dropped
|
||||
# from this compose 2026-05-05.
|
||||
- ../../patches/genesis/vllm/_genesis:/usr/local/lib/python3.12/dist-packages/vllm/_genesis:ro
|
||||
- ../../patches/local/qwen3coder_tool_parser_deferred_commit.py:/patches/qwen3coder_tool_parser_deferred_commit.py:ro
|
||||
environment:
|
||||
- HUGGING_FACE_HUB_TOKEN=${HF_TOKEN:-}
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- NCCL_CUMEM_ENABLE=0
|
||||
- NCCL_P2P_DISABLE=1
|
||||
- VLLM_NO_USAGE_STATS=1
|
||||
- VLLM_USE_FLASHINFER_SAMPLER=1
|
||||
- OMP_NUM_THREADS=1
|
||||
# expandable_segments:True crashes boot on some setups (likely cuMemMap path).
|
||||
# Known: JusefPol on NVLink (PR #31), WSL2 single-card 3090 Ti.
|
||||
# Override via .env: PYTORCH_CUDA_ALLOC_CONF=expandable_segments:False
|
||||
- PYTORCH_CUDA_ALLOC_CONF=${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True,max_split_size_mb:512}
|
||||
- VLLM_ALLOW_LONG_MAX_MODEL_LEN=1
|
||||
- VLLM_MARLIN_USE_ATOMIC_ADD=1
|
||||
- TRITON_CACHE_DIR=/root/.triton/cache
|
||||
- VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS=0
|
||||
- VLLM_FLOAT32_MATMUL_PRECISION=high
|
||||
# VLLM_SSM_CONV_STATE_LAYOUT=DS — RE-ENABLED with our local PN30
|
||||
# dst-shaped temp fix (patch_pn30_dst_shaped_temp_fix.py, applied at
|
||||
# setup time). Without our fix, Sander's PN30 a9977d8 corrupts DS row
|
||||
# strides on spec-decode AL>1. With our fix, PN30 builds destination-
|
||||
# shaped temp via collect_mamba_copy_meta. +6% TPS retained.
|
||||
- VLLM_SSM_CONV_STATE_LAYOUT=DS
|
||||
- VLLM_USE_FUSED_MOE_GROUPED_TOPK=1
|
||||
- CUDA_DEVICE_MAX_CONNECTIONS=8
|
||||
# FULL Genesis v7.69 PROD env-var set per Sandermage's
|
||||
# bare_metal_27b_int4_TQ_k8v4.sh. Validated 2026-05-01 PM dual-3090:
|
||||
# 116.59 code wall_TPS / 92.12 narrative wall_TPS (vs Sander's A5000
|
||||
# 89.23 reference — +30.7% over).
|
||||
- GENESIS_ENABLE_P4=1
|
||||
- GENESIS_ENABLE_P58_ASYNC_PLACEHOLDER_FIX=1
|
||||
- GENESIS_ENABLE_P60_GDN_NGRAM_FIX=1
|
||||
- GENESIS_ENABLE_P60B_TRITON_KERNEL=1
|
||||
- GENESIS_ENABLE_P61_QWEN3_MULTI_TOOL=1
|
||||
- GENESIS_ENABLE_P61B_STREAMING_OVERLAP=1
|
||||
- GENESIS_ENABLE_P62_STRUCT_OUT_SPEC_TIMING=1
|
||||
- GENESIS_ENABLE_P64_QWEN3CODER_MTP_STREAMING=1
|
||||
# P65 dropped 2026-05-03 (mutually exclusive with P67/P67b in v7.69)
|
||||
- GENESIS_ENABLE_P66_CUDAGRAPH_SIZE_FILTER=1
|
||||
- GENESIS_ENABLE_P67_TQ_MULTI_QUERY_KERNEL=1
|
||||
- GENESIS_ENABLE_P68_AUTO_FORCE_TOOL=1
|
||||
- GENESIS_ENABLE_P69_LONG_CTX_TOOL_REMINDER=1
|
||||
- GENESIS_P68_P69_LONG_CTX_THRESHOLD_CHARS=50000
|
||||
- GENESIS_ENABLE_P72_PROFILE_RUN_CAP=1
|
||||
- GENESIS_PROFILE_RUN_CAP_M=4128
|
||||
- GENESIS_ENABLE_P74_CHUNK_CLAMP=1
|
||||
- GENESIS_ENABLE_P83=1
|
||||
# P85 dropped 2026-05-03 (requires P84 which we don't enable in v7.69)
|
||||
# P87 (marlin pad-sub-tile-n text-patch) disabled because the same fix is
|
||||
# already vendored at ../../patches/vllm-marlin-pad/marlin.py and RO-mounted
|
||||
# over the target file (lines 53-54). Letting Genesis re-do the patch fails
|
||||
# with [Errno 30] read-only filesystem and `set -e` propagates exit-1 from
|
||||
# `apply_all` before `vllm serve` runs (club-3090 #49).
|
||||
- GENESIS_ENABLE_P87=0
|
||||
- GENESIS_ENABLE_P91=1
|
||||
- GENESIS_ENABLE_P94=1
|
||||
- GENESIS_ENABLE_P98=1
|
||||
# PN34: active env-opt-in workspace-lock relaxation. P98 above auto-skips
|
||||
# on v0.20 (UNIFORM_SINGLE_TOKEN_DECODE drift-marker false-positive — see
|
||||
# docs/UPSTREAM.md), so PN34 is what's actually firing today. Belt+suspenders
|
||||
# pattern matches default + long-text composes. Propagated from #82 audit.
|
||||
- GENESIS_ENABLE_PN34_WORKSPACE_LOCK_RELAX=1
|
||||
- GENESIS_ENABLE_P99=1
|
||||
- GENESIS_ENABLE_P100=1
|
||||
- GENESIS_ENABLE_P101=1
|
||||
- GENESIS_ENABLE_P103=1
|
||||
- GENESIS_ENABLE_PN8_MTP_DRAFT_ONLINE_QUANT=1
|
||||
- GENESIS_ENABLE_PN9_INDEPENDENT_DRAFTER_ATTN=1
|
||||
- GENESIS_ENABLE_PN11_GDN_AB_CONTIGUOUS=1
|
||||
- GENESIS_ENABLE_PN12_FFN_INTERMEDIATE_POOL=1
|
||||
- GENESIS_ENABLE_PN13_CUDA_GRAPH_LAMBDA_ARITY=1
|
||||
- GENESIS_ENABLE_PN14_TQ_DECODE_OOB_CLAMP=1
|
||||
- GENESIS_ENABLE_PN17_FA2_LSE_CLAMP=1
|
||||
- GENESIS_ENABLE_PN19_SCOPED_MAX_SPLIT=1
|
||||
- GENESIS_ENABLE_PN22_LOCAL_ARGMAX_TP=1
|
||||
- GENESIS_ENABLE_PN26_SPARSE_V=1
|
||||
- GENESIS_ENABLE_PN59_STREAMING_GDN=1
|
||||
- GENESIS_PN26_SPARSE_V_BLOCK_KV=8
|
||||
- GENESIS_PN26_SPARSE_V_NUM_WARPS=4
|
||||
- GENESIS_PN26_SPARSE_V_THRESHOLD=0.01
|
||||
- GENESIS_ENABLE_P38B_COMPILE_SAFE=1
|
||||
- GENESIS_ENABLE_P15B_FA_VARLEN_CLAMP=1
|
||||
- GENESIS_ENABLE_PN25_SILU_INDUCTOR_SAFE=1
|
||||
# PN30 — RE-ENABLED with our local dst-shaped temp fix
|
||||
# (patch_pn30_dst_shaped_temp_fix.py, applied during setup.sh).
|
||||
- GENESIS_ENABLE_PN30_DS_LAYOUT_SPEC_DECODE=1
|
||||
- GENESIS_PREALLOC_TOKEN_BUDGET=4128
|
||||
- GENESIS_BUFFER_MODE=shared
|
||||
# P40 — TQ k8v4 GQA grouping kernel (+15-30% on compute-regime GPUs, L2>=24MB).
|
||||
# Default off (RTX 3090: 6MB L2 = no gain). Enable on RTX 5090/A100/H100.
|
||||
- GENESIS_ENABLE_P40=${GENESIS_ENABLE_P40:-0}
|
||||
# PN54 — GDN contiguous-call deduplication (Cliff 2b OOM mitigation).
|
||||
- GENESIS_ENABLE_PN54=${GENESIS_ENABLE_PN54:-0}
|
||||
# Explicit OFFs to match Sandermage's PROD env-var set verbatim:
|
||||
# P78 (P78_TOLIST_CAPTURE_GUARD) — superseded by our patch_tolist_cudagraph.py
|
||||
# P81 (FP8 block-scaled M<=8) — FP8-specific, no-op on our TQ3 path
|
||||
# P82 — biased on small-batch single-stream Lorbus INT4 + MTP K=3 (Sander PROD)
|
||||
- GENESIS_ENABLE_P78_TOLIST_CAPTURE_GUARD=0
|
||||
- GENESIS_ENABLE_P81_FP8_BLOCK_SCALED_M_LE_8=0
|
||||
- GENESIS_ENABLE_P82=${GENESIS_ENABLE_P82:-0}
|
||||
- GENESIS_P82_THRESHOLD_SINGLE=0.3
|
||||
shm_size: "16gb"
|
||||
ipc: host
|
||||
deploy:
|
||||
resources:
|
||||
reservations:
|
||||
devices:
|
||||
- driver: nvidia
|
||||
count: all
|
||||
capabilities: [gpu]
|
||||
entrypoint:
|
||||
- /bin/bash
|
||||
- -c
|
||||
- |
|
||||
set -e
|
||||
pip install xxhash pandas scipy -q
|
||||
# Pre-flight: Genesis patches must be populated. Empty volume mount
|
||||
# = silent no-op apply_all = boot fails later with cryptic upstream
|
||||
# error (e.g. "TurboQuant KV not supported for hybrid models", #13).
|
||||
if [ ! -f /usr/local/lib/python3.12/dist-packages/vllm/_genesis/patches/apply_all.py ]; then
|
||||
echo "ERROR: Genesis patches missing — host volume models/qwen3.6-27b/vllm/patches/genesis/ is empty." >&2
|
||||
echo " Run from repo root before 'docker compose up':" >&2
|
||||
echo " bash scripts/setup.sh qwen3.6-27b" >&2
|
||||
exit 1
|
||||
fi
|
||||
python3 -m vllm._genesis.patches.apply_all
|
||||
# Tool-parser deferred-commit fix for qwen3coder SSE-silence bug (issue #72).
|
||||
# Drops out when vllm-project/vllm lands the upstream fix.
|
||||
python3 /patches/qwen3coder_tool_parser_deferred_commit.py
|
||||
# were previously invoked here; superseded by Genesis natives in
|
||||
# v7.72.2 (P78 + PN34). Mounts and invocations dropped 2026-05-05.
|
||||
# VLLM_ENFORCE_EAGER=1 in compose/.env disables CUDA graphs — use on
|
||||
# hardware where Cliff 2 GDN activation spikes occur at runtime.
|
||||
exec vllm serve ${VLLM_ENFORCE_EAGER:+--enforce-eager} "$@"
|
||||
- --
|
||||
command:
|
||||
- --model
|
||||
- /root/.cache/huggingface/qwen3.6-27b-autoround-int4
|
||||
- --served-model-name
|
||||
- qwen3.6-27b-autoround
|
||||
- --quantization
|
||||
- auto_round
|
||||
- --dtype
|
||||
- float16
|
||||
- --tensor-parallel-size
|
||||
- "2"
|
||||
- --disable-custom-all-reduce
|
||||
- --max-model-len
|
||||
- "${MAX_MODEL_LEN:-262144}"
|
||||
- --gpu-memory-utilization
|
||||
- "${GPU_MEMORY_UTILIZATION:-0.85}"
|
||||
- --max-num-seqs
|
||||
- "2"
|
||||
- --max-num-batched-tokens
|
||||
- "4128"
|
||||
# TQ3 is the right pick on 24 GB / 3090 (smaller KV pool → more concurrency).
|
||||
# On 20 GB Ampere (modded 3080 / cap'd 3090) override to fp8_e5m2 — TQ3's
|
||||
# activation peak during DeltaNet GDN forward exceeds the per-card budget
|
||||
# after TP=2 split and Cliff 2 fires at 90K. fp8_e5m2 trades KV-pool
|
||||
# capacity for activation headroom on the smaller-VRAM sub-class. See
|
||||
# docs/HARDWARE.md "Note for sub-24 GB cards" + #47 for cross-rig data.
|
||||
- --kv-cache-dtype
|
||||
- "${KV_CACHE_DTYPE:-turboquant_3bit_nc}"
|
||||
- --trust-remote-code
|
||||
- --reasoning-parser
|
||||
- qwen3
|
||||
- --default-chat-template-kwargs
|
||||
- '{"enable_thinking": false}'
|
||||
- --enable-auto-tool-choice
|
||||
- --tool-call-parser
|
||||
- qwen3_coder
|
||||
- --enable-prefix-caching
|
||||
- --enable-chunked-prefill
|
||||
- --speculative-config
|
||||
- '{"method":"mtp","num_speculative_tokens":3}'
|
||||
- --host
|
||||
- 0.0.0.0
|
||||
- --port
|
||||
- "8000"
|
||||
@@ -0,0 +1,164 @@
|
||||
# ===========================================================================
|
||||
# Profile (at-a-glance):
|
||||
# Model: Qwen3.6-27B (Lorbus AutoRound INT4 + BF16 mtp.fc preserved)
|
||||
# Topology: Dual 3090 PCIe (TP=2, no NVLink)
|
||||
# Drafter: MTP n=3 (built-in)
|
||||
# KV: turboquant_3bit_nc (TQ3) — 3-bit per element
|
||||
# Max ctx: 262144 (262K)
|
||||
# Status: ⛔ TOMBSTONED 2026-05-11 — NOT shipping. See gap below.
|
||||
# Best for: Reference / re-test surface once the missing 4 of 5 Genesis
|
||||
# equivalent upstream PRs land.
|
||||
# ---------------------------------------------------------------------------
|
||||
# 2026-05-11 UPSTREAM-FIX GAP (no Genesis):
|
||||
#
|
||||
# TQ3 + MTP on Qwen3-Next requires 5 distinct fixes to work correctly.
|
||||
# Sandermage root-caused all 5 on 2026-04-25 and shipped them as
|
||||
# Genesis v7.14 patches P64 / P65 / P66 / P68 / P69, closing the issue
|
||||
# I filed at vllm-issue#40880 (now closed).
|
||||
#
|
||||
# Today's upstream-PR landscape covers only 1 of those 5:
|
||||
# - PR #40914 (OPEN) — K+1 spec-verify routing (analog of one piece)
|
||||
# - PR #40798 (OPEN) — TQ decode workspace pre-allocation (separate
|
||||
# workspace-lock bug; partially required for boot)
|
||||
# - PR #42215 (OPEN) — TQ decode kernel warmup (overlapping with #40798)
|
||||
#
|
||||
# The other 4 Genesis patches do not have upstream PR analogs we can
|
||||
# find:
|
||||
# - chunked-prefill `cu_seqlens_k = cu_seqlens_q` first-chunk assumption
|
||||
# - decode-kernel warmup coverage (broader than #42215)
|
||||
# - multi-tool / streaming overlap
|
||||
# - long-context tool reminders / force-tool
|
||||
#
|
||||
# Tracking issue vllm#40069 lists "Speculative decoding / Eagle" as
|
||||
# "currently disabled or unverified" for the TurboQuant backend —
|
||||
# upstream acknowledges this is incomplete work.
|
||||
#
|
||||
# We attempted a partial stack on 2026-05-11:
|
||||
# - marlin-pad (vllm#40361)
|
||||
# - PR #40798 `gpu_model_runner.py` ONLY (workspace pre-allocation)
|
||||
# - PR #40914 K+1 dispatch block manually rebased onto post-#41434 main
|
||||
#
|
||||
# Result: MTP acceptance fully stabilized (100% sustained, AL=4.0,
|
||||
# per-position [1.000, 1.000, 1.000]). But the other 4 missing fixes
|
||||
# surface as: `'crimson axolotl 36'` → `'cr!!!!!!!!!!!!!!!!!!!!!! '`
|
||||
# on long-context needle recall + tool-prefill / multi-turn timeouts +
|
||||
# bench prompts hanging on non-terminating corrupted output. Combined
|
||||
# with 100% drafter-target agreement, this means the K+1 dispatch
|
||||
# path is computing the same wrong attention output on both — i.e.
|
||||
# one fix doesn't compensate for the missing four.
|
||||
#
|
||||
# Deployable alternatives without Genesis:
|
||||
# - dual/tq3-nomtp.yml ← validated: 1.73M KV pool, 6.59× concurrency,
|
||||
# verify-stress all pass, no MTP acceleration
|
||||
# - dual/int8.yml ← shipping baseline: 605K KV pool, 2.31×, MTP working
|
||||
#
|
||||
# Deployable alternative WITH Genesis:
|
||||
# - dual/tq3-mtp-genesis.yml ← uses the full P64/P65/P66/P68/P69 stack
|
||||
# via Genesis modular package
|
||||
#
|
||||
# Drop the tombstone when the 4 missing upstream PRs land OR PR #40914
|
||||
# is rebased + fully extended to cover the gap.
|
||||
# ===========================================================================
|
||||
# Run (only when all 5 upstream fixes are available):
|
||||
# docker compose -f dual/tq3-mtp.yml up -d
|
||||
# ===========================================================================
|
||||
services:
|
||||
vllm-qwen36-27b-dual-int8-tq3:
|
||||
# Same vLLM nightly as gemma-4-31b/vllm/compose/dual/int8-tq3.yml — 2026-05-08
|
||||
# cut, post Gemma 4 MTP merge + INT8 PTH KV path (vllm#40391). Pin matches
|
||||
# so head-to-head TPS deltas reflect model intrinsics, not vLLM-version
|
||||
# drift between Qwen and Gemma legs.
|
||||
image: vllm/vllm-openai:nightly-1acd67a795ebccdf9b9db7697ae9082058301657
|
||||
container_name: vllm-qwen36-27b-dual-int8-tq3
|
||||
restart: "no"
|
||||
ports:
|
||||
- "${BIND_HOST:-0.0.0.0}:${PORT:-8013}:8000"
|
||||
volumes:
|
||||
- ${MODEL_DIR:-../../../../../models-cache}:/root/.cache/huggingface
|
||||
# Reuse the same torch.compile + Triton kernel caches as dual.yml so
|
||||
# second-boot warm time stays short. The kernel cache is keyed by
|
||||
# (model, kv-dtype, vllm-version) so we share with dual.yml only when
|
||||
# all three match — for int8 vs fp8 the cache won't collide.
|
||||
- ../../cache/torch_compile:/root/.cache/vllm/torch_compile_cache
|
||||
- ../../cache/triton:/root/.triton/cache
|
||||
# Marlin pad-sub-tile-n patch (vLLM PR #40361) — still required for
|
||||
# TP=2 on AutoRound W4A16 (drops out when upstream merges).
|
||||
- ../../patches/vllm-marlin-pad/marlin.py:/usr/local/lib/python3.12/dist-packages/vllm/model_executor/kernels/linear/mixed_precision/marlin.py:ro
|
||||
- ../../patches/vllm-marlin-pad/MPLinearKernel.py:/usr/local/lib/python3.12/dist-packages/vllm/model_executor/kernels/linear/mixed_precision/MPLinearKernel.py:ro
|
||||
# TOMBSTONED: the partial PR #40798 (gpu_model_runner.py) +
|
||||
# manually-rebased PR #40914 (K+1 dispatch) overlays were validated
|
||||
# 2026-05-11 but produce corrupted attention output on long-ctx
|
||||
# (only 1 of 5 needed Genesis fixes is upstream). Patches are
|
||||
# retained in models/qwen3.6-27b/vllm/patches/{vllm-pr40798-rebased,
|
||||
# vllm-pr40914-k1-only} for re-test once upstream PR landscape
|
||||
# extends to cover the gap.
|
||||
environment:
|
||||
- HUGGING_FACE_HUB_TOKEN=${HF_TOKEN:-}
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- NCCL_CUMEM_ENABLE=0
|
||||
- NCCL_P2P_DISABLE=1
|
||||
- VLLM_NO_USAGE_STATS=1
|
||||
- VLLM_USE_FLASHINFER_SAMPLER=1
|
||||
- OMP_NUM_THREADS=1
|
||||
- PYTORCH_CUDA_ALLOC_CONF=${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True,max_split_size_mb:512}
|
||||
shm_size: "16gb"
|
||||
ipc: host
|
||||
deploy:
|
||||
resources:
|
||||
reservations:
|
||||
devices:
|
||||
- driver: nvidia
|
||||
count: all
|
||||
capabilities: [gpu]
|
||||
entrypoint:
|
||||
- bash
|
||||
- -c
|
||||
- |
|
||||
exec vllm serve ${VLLM_ENFORCE_EAGER:+--enforce-eager} "$@"
|
||||
- --
|
||||
command:
|
||||
- --model
|
||||
- /root/.cache/huggingface/qwen3.6-27b-autoround-int4
|
||||
- --served-model-name
|
||||
- qwen3.6-27b-autoround
|
||||
- --quantization
|
||||
- auto_round
|
||||
- --dtype
|
||||
- float16
|
||||
- --tensor-parallel-size
|
||||
- "2"
|
||||
- --disable-custom-all-reduce
|
||||
- --max-model-len
|
||||
- "${MAX_MODEL_LEN:-262144}"
|
||||
- --gpu-memory-utilization
|
||||
- "${GPU_MEMORY_UTILIZATION:-0.92}"
|
||||
- --max-num-seqs
|
||||
- "${MAX_NUM_SEQS:-2}"
|
||||
- --max-num-batched-tokens
|
||||
- "8192"
|
||||
# INT8 per-token-head KV via vllm#40391 (same as gemma-int8-tq3.yml).
|
||||
# If this errors at boot on Qwen3-Next DeltaNet hybrid attention,
|
||||
# the model may not yet be compatible — fall back to fp8_e5m2 dual.yml.
|
||||
- --kv-cache-dtype
|
||||
- "${KV_DTYPE:-turboquant_3bit_nc}"
|
||||
- --trust-remote-code
|
||||
- --reasoning-parser
|
||||
- qwen3
|
||||
- --default-chat-template-kwargs
|
||||
- '{"enable_thinking": false}'
|
||||
- --enable-auto-tool-choice
|
||||
- --tool-call-parser
|
||||
- qwen3_coder
|
||||
- --enable-prefix-caching
|
||||
- --enable-chunked-prefill
|
||||
# MTP n — architectural target for Qwen3-Next built-in heads is the
|
||||
# number of physical MTP layers in the model (3 for Qwen3.6 AutoRound
|
||||
# at time of writing). Higher n re-uses the same MTP layer for the
|
||||
# extra draft steps and emits a vLLM warning about reduced acceptance
|
||||
# rate. Use SPEC_N_MAX=4 to A/B against Gemma's n=4.
|
||||
- --speculative-config
|
||||
- '{"method":"mtp","num_speculative_tokens":${SPEC_N_MAX:-3}}'
|
||||
- --host
|
||||
- 0.0.0.0
|
||||
- --port
|
||||
- "8000"
|
||||
+51
-22
@@ -1,36 +1,51 @@
|
||||
# ===========================================================================
|
||||
# Profile (at-a-glance):
|
||||
# Model: Qwen3.6-27B (Lorbus AutoRound INT4 + BF16 mtp.fc preserved)
|
||||
# Model: Qwen3.6-27B (Lorbus AutoRound INT4)
|
||||
# Topology: Dual 3090 PCIe (TP=2, no NVLink)
|
||||
# Drafter: MTP n=3 (built-in, architectural max for Qwen3-Next)
|
||||
# KV: int8_per_token_head (PTH, via vllm#40391) — 1 byte/token like fp8
|
||||
# but uses PR #40391 quantized-pool path that lets Gemma run at
|
||||
# the same KV class. Created for **matched-config head-to-head
|
||||
# against gemma-int8-tq3.yml** — same vLLM pin, same KV class, same
|
||||
# 262K ctx.
|
||||
# Vision: yes
|
||||
# Max ctx: 262144 (262K — same as fp8 dual.yml)
|
||||
# Genesis: none (matches dual.yml — intentionally Genesis-free)
|
||||
# Status: 🧪 Experimental — INT8 PTH KV compatibility on Qwen3-Next DeltaNet
|
||||
# hybrid not previously validated on this stack. Boot test + verify-
|
||||
# full + bench before treating as a shipping path.
|
||||
# Best for: Apples-to-apples comparison with Gemma 4 31B int8-tq3.yml on the
|
||||
# same vLLM nightly + KV class.
|
||||
# Drafter: none — MTP intentionally disabled
|
||||
# KV: turboquant_3bit_nc (TQ3, 0.375 bytes/token)
|
||||
# Max ctx: 262144 (262K)
|
||||
# Streams: max-num-seqs=2 (matched-config sibling of dual/int8.yml)
|
||||
# Genesis: none — pure upstream nightly + only the unavoidable marlin-pad overlay
|
||||
# Status: ✅ Working — validated 2026-05-11
|
||||
# Best for: Long-context multi-tenant agentic workloads where TQ3's KV-pool
|
||||
# expansion matters more than MTP-driven TPS acceleration.
|
||||
# ---------------------------------------------------------------------------
|
||||
# Run:
|
||||
# docker compose -f dual/int8-tq3.yml up -d
|
||||
# Why this compose exists:
|
||||
#
|
||||
# TQ3 + MTP on Qwen3-Next requires 5 upstream fixes (Sander's Genesis
|
||||
# P64/P65/P66/P68/P69; see vllm-issue#40880). Only 1 of 5 has a
|
||||
# community-PR analog today (vllm#40914). The Genesis-free TQ3+MTP path is
|
||||
# therefore tombstoned in dual/tq3-mtp.yml — re-test when upstream catches up.
|
||||
#
|
||||
# In the meantime, TQ3 KV WITHOUT MTP works perfectly on vanilla upstream
|
||||
# vLLM nightly + only marlin-pad (which we already vendor for AutoRound TP=2
|
||||
# booting). Verify-stress 7/7 pass on the 7-step boundary ladder including
|
||||
# 60K and 91K needle recall. The trade is: no MTP TPS acceleration (~67 TPS
|
||||
# raw decode vs ~85 TPS with MTP working at AL=2.31× on INT8 PTH baseline)
|
||||
# but in exchange:
|
||||
# KV pool = 1,727,935 tokens at 262K (6.59× concurrency)
|
||||
# vs INT8 PTH: 605K (2.31×) — 2.86× more KV budget per VRAM byte
|
||||
#
|
||||
# For workloads where the agent fleet's total in-flight context dominates
|
||||
# per-token latency (multi-tenant code review, long-doc analysis, RAG with
|
||||
# many concurrent sessions), this is the deployable TQ3 path until upstream
|
||||
# bundles the 5 Genesis fixes.
|
||||
#
|
||||
# To run:
|
||||
# docker compose -f dual/tq3-nomtp.yml up -d
|
||||
# ===========================================================================
|
||||
services:
|
||||
vllm-qwen36-27b-dual-int8-tq3:
|
||||
vllm-qwen36-27b-dual-int8-tq3-nomtp:
|
||||
# Same vLLM nightly as gemma-4-31b/vllm/compose/dual/int8-tq3.yml — 2026-05-08
|
||||
# cut, post Gemma 4 MTP merge + INT8 PTH KV path (vllm#40391). Pin matches
|
||||
# so head-to-head TPS deltas reflect model intrinsics, not vLLM-version
|
||||
# drift between Qwen and Gemma legs.
|
||||
image: vllm/vllm-openai:nightly-1acd67a795ebccdf9b9db7697ae9082058301657
|
||||
container_name: vllm-qwen36-27b-dual-int8-tq3
|
||||
container_name: vllm-qwen36-27b-dual-int8-tq3-nomtp
|
||||
restart: "no"
|
||||
ports:
|
||||
- "${BIND_HOST:-0.0.0.0}:${PORT:-8013}:8000"
|
||||
- "${BIND_HOST:-0.0.0.0}:${PORT:-8014}:8000"
|
||||
volumes:
|
||||
- ${MODEL_DIR:-../../../../../models-cache}:/root/.cache/huggingface
|
||||
# Reuse the same torch.compile + Triton kernel caches as dual.yml so
|
||||
@@ -43,6 +58,22 @@ services:
|
||||
# TP=2 on AutoRound W4A16 (drops out when upstream merges).
|
||||
- ../../patches/vllm-marlin-pad/marlin.py:/usr/local/lib/python3.12/dist-packages/vllm/model_executor/kernels/linear/mixed_precision/marlin.py:ro
|
||||
- ../../patches/vllm-marlin-pad/MPLinearKernel.py:/usr/local/lib/python3.12/dist-packages/vllm/model_executor/kernels/linear/mixed_precision/MPLinearKernel.py:ro
|
||||
# vLLM PR #40798 overlay — TurboQuant share decode scratch workspace
|
||||
# across layers + reserve max workspace before cudagraph lock.
|
||||
# Bug: vllm-issue#41565 / #41726 / #40420.
|
||||
# WE INTENTIONALLY DO NOT MOUNT THE PR's turboquant_attn.py — that
|
||||
# file in PR #40798's branch is pre-#41434 and would revert the
|
||||
# merged `.tolist()` → CPU-resident-metadata fixes, reintroducing
|
||||
# "Cannot copy CPU/CUDA during cudagraph capture" crashes. We keep
|
||||
# nightly's turboquant_attn.py (which has #41434 fixes) and only
|
||||
# mount the workspace-pre-allocation deltas from PR #40798's other
|
||||
# 2 files. See ../../patches/vllm-pr40798-rebased/README.md.
|
||||
# NOT mounting PR #40798's triton_turboquant_decode.py either — it's
|
||||
# also pre-#41434 and has the old function signature without the
|
||||
# `buf_holder` kwarg that nightly's turboquant_attn.py now passes.
|
||||
# Only the workspace-pre-allocation logic in gpu_model_runner.py is
|
||||
# used; nightly already has the kernel + caller for the rest.
|
||||
- ../../patches/vllm-pr40798-rebased/v1/worker/gpu_model_runner.py:/usr/local/lib/python3.12/dist-packages/vllm/v1/worker/gpu_model_runner.py:ro
|
||||
environment:
|
||||
- HUGGING_FACE_HUB_TOKEN=${HF_TOKEN:-}
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
@@ -107,8 +138,6 @@ services:
|
||||
# at time of writing). Higher n re-uses the same MTP layer for the
|
||||
# extra draft steps and emits a vLLM warning about reduced acceptance
|
||||
# rate. Use SPEC_N_MAX=4 to A/B against Gemma's n=4.
|
||||
- --speculative-config
|
||||
- '{"method":"mtp","num_speculative_tokens":${SPEC_N_MAX:-3}}'
|
||||
- --host
|
||||
- 0.0.0.0
|
||||
- --port
|
||||
@@ -0,0 +1,63 @@
|
||||
# vLLM PR #40798 overlay — TurboQuant max-workspace pre-allocation
|
||||
|
||||
Vendored 2026-05-11 to unblock `turboquant_3bit_nc` KV + MTP on Qwen 3.6 27B
|
||||
(`dual/int8-tq3.yml`).
|
||||
|
||||
## Source
|
||||
|
||||
- Upstream PR: <https://github.com/vllm-project/vllm/pull/40798>
|
||||
- Title: `[TurboQuant] Share decode scratch workspace across layers`
|
||||
- Head SHA: `cbe823e9b5e32b3590705f9ee86ab57db887754d`
|
||||
- State at vendor time: OPEN, MERGEABLE.
|
||||
- 3 source files: `turboquant_attn.py` (−1), `triton_turboquant_decode.py`
|
||||
(+15/−8), `gpu_model_runner.py` (+57/−1).
|
||||
|
||||
## What it fixes
|
||||
|
||||
Without this PR, the MTP drafter's `_decode_attention` and the long-context
|
||||
`_continuation_prefill` paths request workspace allocations AFTER the
|
||||
cudagraph capture has locked it (per vllm#39226). On vLLM nightly `1acd67a7`
|
||||
that surfaces as:
|
||||
|
||||
```
|
||||
AssertionError: Workspace is locked but allocation from
|
||||
'turboquant_attn.py:747:_continuation_prefill' requires 8.06 MB,
|
||||
current size is 0.76 MB. Workspace growth is not allowed after locking.
|
||||
```
|
||||
|
||||
Bug tracked at vllm-issue#41565 / #41726 / #40420.
|
||||
|
||||
This PR moves TurboQuant decode scratch allocation into the v1 workspace
|
||||
manager so scratch tensors are shared across layers, AND reserves the
|
||||
maximum TurboQuant decode workspace before CUDA graph capture locks the
|
||||
workspace. Both the decode and continuation-prefill paths get the right
|
||||
workspace size up front, so the strict lock invariant from vllm#39226 holds.
|
||||
|
||||
## Why not Genesis
|
||||
|
||||
The Sandermage Genesis package's PN34 solves the same problem (and many
|
||||
others) via a separate draft-model workspace pool. We've intentionally
|
||||
kept this compose Genesis-free to validate the upstream fix in isolation.
|
||||
|
||||
## Why not PR #42215
|
||||
|
||||
PR #42215 (Warm up decode kernels) attacks the same bug from a different
|
||||
angle: it warms up the decode kernel during `kernel_warmup()` so the
|
||||
workspace high-water mark is captured before lock. We tried that overlay
|
||||
first on 2026-05-11 and confirmed it fixes the decode path (line 879)
|
||||
but NOT the continuation-prefill path (line 747). PR #40798 reserves the
|
||||
max workspace across all TQ paths, which covers both.
|
||||
|
||||
## Drop trigger
|
||||
|
||||
```
|
||||
gh api repos/vllm-project/vllm/pulls/40798 --jq '.state, .merged_at'
|
||||
```
|
||||
|
||||
reports `MERGED`. Then bump the nightly pin past the merge commit.
|
||||
|
||||
## Verified on
|
||||
|
||||
- vLLM nightly: `1acd67a7`
|
||||
- Compose: `dual/int8-tq3.yml` (Qwen 3.6 27B AutoRound INT4, TQ3 KV, MTP n=3, 262K, 2 streams)
|
||||
- 2026-05-11
|
||||
+637
@@ -0,0 +1,637 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Triton fused TurboQuant decode attention.
|
||||
|
||||
Decode path: Triton stage1 (split-KV tiled attention scoring + value
|
||||
accumulation) + stage2 (log-sum-exp reduction across splits).
|
||||
|
||||
Supports FP8 (E4M3) keys, 3-bit and 4-bit uniform quantized values.
|
||||
"""
|
||||
|
||||
import math
|
||||
|
||||
import torch
|
||||
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.triton_utils import tl, triton
|
||||
from vllm.v1.attention.ops.triton_decode_attention import (
|
||||
_fwd_kernel_stage2,
|
||||
)
|
||||
|
||||
_FP8_E4B15: dict[int, int] = {}
|
||||
|
||||
|
||||
def _use_fp8_e4b15(device: int = 0) -> int:
|
||||
"""Return 1 if device needs fp8e4b15 (Ampere/Ada, SM < 8.9), else 0.
|
||||
On non-CUDA platforms (e.g. XPU), always returns 0 (use e4nv format).
|
||||
"""
|
||||
if device not in _FP8_E4B15:
|
||||
if current_platform.is_cuda_alike():
|
||||
cap = torch.cuda.get_device_capability(device)
|
||||
_FP8_E4B15[device] = 1 if cap < (8, 9) else 0
|
||||
else:
|
||||
_FP8_E4B15[device] = 0
|
||||
return _FP8_E4B15[device]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Stage 1: Fused TQ score + value accumulation (BLOCK_KV tiled)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _tq_decode_stage1(
|
||||
# Precomputed query projection
|
||||
Q_rot_ptr, # [B, Hq, D] float32
|
||||
# Compressed KV cache (combined K+V)
|
||||
KV_cache_ptr, # [num_blocks, block_size, Hk, padded_slot] uint8
|
||||
# Block table and sequence info
|
||||
Block_table_ptr, # [B, max_num_blocks] int32
|
||||
Seq_lens_ptr, # [B] int32
|
||||
# TQ parameters
|
||||
Centroids_ptr, # [n_centroids] float32
|
||||
# Output (intermediate for stage2)
|
||||
Mid_o_ptr, # [B, Hq, NUM_KV_SPLITS, D+1] float32
|
||||
# Strides
|
||||
stride_qb,
|
||||
stride_qh, # Q strides: [B, Hq, D]
|
||||
stride_cache_block,
|
||||
stride_cache_pos,
|
||||
stride_cache_head, # KV cache
|
||||
stride_bt_b, # block_table stride per batch
|
||||
stride_mid_b,
|
||||
stride_mid_h,
|
||||
stride_mid_s, # mid_o strides
|
||||
# Constexpr dims
|
||||
NUM_KV_HEADS: tl.constexpr,
|
||||
HEAD_DIM: tl.constexpr,
|
||||
BLOCK_SIZE: tl.constexpr, # KV cache block_size (pages)
|
||||
NUM_KV_SPLITS: tl.constexpr,
|
||||
KV_GROUP_SIZE: tl.constexpr, # Hq // Hk
|
||||
# TQ layout constants
|
||||
MSE_BITS: tl.constexpr, # 3 or 4
|
||||
MSE_BYTES: tl.constexpr, # ceil(D * mse_bits / 8)
|
||||
KPS: tl.constexpr, # key_packed_size
|
||||
VQB: tl.constexpr, # value_quant_bits (4 or 8=FP8)
|
||||
VAL_DATA_BYTES: tl.constexpr, # ceil(D * vqb / 8) or D for FP8
|
||||
# Score constants
|
||||
ATTN_SCALE: tl.constexpr, # 1/sqrt(D)
|
||||
# Block tile sizes
|
||||
BLOCK_D: tl.constexpr, # next_power_of_2(HEAD_DIM)
|
||||
BLOCK_KV: tl.constexpr, # tokens per tile (16)
|
||||
KEY_FP8: tl.constexpr, # 1 if K is stored as FP8
|
||||
NORM_CORRECTION: tl.constexpr = 0, # 1 = re-normalize centroids
|
||||
FP8_E4B15: tl.constexpr = 0, # 1 = use e4b15 (Ampere/Ada), 0 = e4nv (Hopper+)
|
||||
):
|
||||
bid = tl.program_id(0) # batch index
|
||||
hid = tl.program_id(1) # q_head index
|
||||
sid = tl.program_id(2) # kv_split index
|
||||
|
||||
kv_head = hid // KV_GROUP_SIZE
|
||||
|
||||
# Sequence length for this batch
|
||||
seq_len = tl.load(Seq_lens_ptr + bid)
|
||||
|
||||
# KV split range
|
||||
split_len = tl.cdiv(seq_len, NUM_KV_SPLITS)
|
||||
split_start = split_len * sid
|
||||
split_end = tl.minimum(split_start + split_len, seq_len)
|
||||
|
||||
if split_start >= split_end:
|
||||
return
|
||||
|
||||
# Dimension offsets
|
||||
d_offs = tl.arange(0, BLOCK_D)
|
||||
d_mask = d_offs < HEAD_DIM
|
||||
kv_range = tl.arange(0, BLOCK_KV)
|
||||
|
||||
# Load query vector: q_rot — [BLOCK_D] float32
|
||||
q_base = bid * stride_qb + hid * stride_qh
|
||||
q_rot = tl.load(Q_rot_ptr + q_base + d_offs, mask=d_mask, other=0.0).to(tl.float32)
|
||||
|
||||
# Precompute byte/bit index vectors for MSE gather loads
|
||||
if not KEY_FP8:
|
||||
mse_bit_off = d_offs * MSE_BITS
|
||||
mse_byte_idx = mse_bit_off // 8
|
||||
mse_bit_shift = mse_bit_off % 8
|
||||
mse_mask = (1 << MSE_BITS) - 1
|
||||
|
||||
# Precompute value bit/byte index vectors (loop-invariant)
|
||||
if VQB == 3:
|
||||
val_bit_off = d_offs * 3
|
||||
val_byte_idx = val_bit_off // 8
|
||||
val_bit_shift = val_bit_off % 8
|
||||
|
||||
# Online softmax accumulators
|
||||
m_prev = -float("inf")
|
||||
l_prev = 0.0
|
||||
acc = tl.zeros([BLOCK_D], dtype=tl.float32)
|
||||
|
||||
bt_base = bid * stride_bt_b
|
||||
|
||||
# ================================================================
|
||||
# TILED LOOP: process BLOCK_KV tokens per iteration
|
||||
# ================================================================
|
||||
for start_n in range(split_start, split_end, BLOCK_KV):
|
||||
kv_offs = start_n + kv_range
|
||||
kv_mask = kv_offs < split_end
|
||||
|
||||
page_idx = kv_offs // BLOCK_SIZE
|
||||
page_off = kv_offs % BLOCK_SIZE
|
||||
block_nums = tl.load(
|
||||
Block_table_ptr + bt_base + page_idx,
|
||||
mask=kv_mask,
|
||||
other=0,
|
||||
).to(tl.int64)
|
||||
|
||||
slot_bases = (
|
||||
block_nums * stride_cache_block
|
||||
+ page_off.to(tl.int64) * stride_cache_pos
|
||||
+ tl.cast(kv_head, tl.int64) * stride_cache_head
|
||||
)
|
||||
|
||||
# ============================================================
|
||||
# COMPUTE ATTENTION SCORES: [BLOCK_KV]
|
||||
# ============================================================
|
||||
if KEY_FP8:
|
||||
k_addrs = slot_bases[:, None] + d_offs[None, :]
|
||||
k_raw = tl.load(
|
||||
KV_cache_ptr + k_addrs,
|
||||
mask=kv_mask[:, None] & d_mask[None, :],
|
||||
other=0,
|
||||
)
|
||||
if FP8_E4B15:
|
||||
k_float = k_raw.to(tl.float8e4b15, bitcast=True).to(tl.float32)
|
||||
else:
|
||||
k_float = k_raw.to(tl.float8e4nv, bitcast=True).to(tl.float32)
|
||||
scores = (
|
||||
tl.sum(
|
||||
tl.where(d_mask[None, :], q_rot[None, :] * k_float, 0.0),
|
||||
axis=1,
|
||||
)
|
||||
* ATTN_SCALE
|
||||
)
|
||||
scores = tl.where(kv_mask, scores, -float("inf"))
|
||||
else:
|
||||
# MSE unpack + norms
|
||||
mse_addrs0 = slot_bases[:, None] + mse_byte_idx[None, :]
|
||||
mse_raw0 = tl.load(
|
||||
KV_cache_ptr + mse_addrs0,
|
||||
mask=kv_mask[:, None] & d_mask[None, :],
|
||||
other=0,
|
||||
).to(tl.int32)
|
||||
mse_raw1 = tl.load(
|
||||
KV_cache_ptr + mse_addrs0 + 1,
|
||||
mask=kv_mask[:, None] & d_mask[None, :],
|
||||
other=0,
|
||||
).to(tl.int32)
|
||||
raw16 = mse_raw0 | (mse_raw1 << 8)
|
||||
mse_idx = (raw16 >> mse_bit_shift[None, :]) & mse_mask
|
||||
|
||||
# Centroid gather + dot product
|
||||
c_vals = tl.load(
|
||||
Centroids_ptr + mse_idx,
|
||||
mask=kv_mask[:, None] & d_mask[None, :],
|
||||
other=0.0,
|
||||
)
|
||||
|
||||
# Norm correction: re-normalize centroid vector to unit norm
|
||||
if NORM_CORRECTION:
|
||||
c_norm_sq = tl.sum(
|
||||
tl.where(d_mask[None, :], c_vals * c_vals, 0.0),
|
||||
axis=1,
|
||||
)
|
||||
c_inv_norm = 1.0 / tl.sqrt(c_norm_sq + 1e-16)
|
||||
c_vals = c_vals * c_inv_norm[:, None]
|
||||
|
||||
term1 = tl.sum(
|
||||
tl.where(d_mask[None, :], q_rot[None, :] * c_vals, 0.0),
|
||||
axis=1,
|
||||
)
|
||||
|
||||
# Load norms (fp16 -> fp32): norms are at MSE_BYTES offset
|
||||
norm_bases = slot_bases + MSE_BYTES
|
||||
n_lo = tl.load(KV_cache_ptr + norm_bases, mask=kv_mask, other=0).to(
|
||||
tl.uint16
|
||||
)
|
||||
n_hi = tl.load(KV_cache_ptr + norm_bases + 1, mask=kv_mask, other=0).to(
|
||||
tl.uint16
|
||||
)
|
||||
vec_norms = (n_lo | (n_hi << 8)).to(tl.float16, bitcast=True).to(tl.float32)
|
||||
|
||||
scores = vec_norms * term1 * ATTN_SCALE
|
||||
scores = tl.where(kv_mask, scores, -float("inf"))
|
||||
|
||||
# ============================================================
|
||||
# ONLINE SOFTMAX UPDATE (block-level)
|
||||
# ============================================================
|
||||
n_e_max = tl.maximum(tl.max(scores, 0), m_prev)
|
||||
re_scale = tl.exp(m_prev - n_e_max)
|
||||
p = tl.exp(scores - n_e_max)
|
||||
|
||||
# ============================================================
|
||||
# VALUE LOAD + DEQUANTIZE: [BLOCK_KV, BLOCK_D]
|
||||
# ============================================================
|
||||
val_bases = slot_bases + KPS
|
||||
|
||||
if VQB == 3:
|
||||
val_addrs0 = val_bases[:, None] + val_byte_idx[None, :]
|
||||
val_raw0 = tl.load(
|
||||
KV_cache_ptr + val_addrs0,
|
||||
mask=kv_mask[:, None] & d_mask[None, :],
|
||||
other=0,
|
||||
).to(tl.int32)
|
||||
val_raw1 = tl.load(
|
||||
KV_cache_ptr + val_addrs0 + 1,
|
||||
mask=kv_mask[:, None] & d_mask[None, :],
|
||||
other=0,
|
||||
).to(tl.int32)
|
||||
raw16 = val_raw0 | (val_raw1 << 8)
|
||||
v_idx = ((raw16 >> val_bit_shift[None, :]) & 0x7).to(tl.float32)
|
||||
|
||||
sc_bases = val_bases + VAL_DATA_BYTES
|
||||
sc_lo = tl.load(KV_cache_ptr + sc_bases, mask=kv_mask, other=0).to(
|
||||
tl.uint16
|
||||
)
|
||||
sc_hi = tl.load(KV_cache_ptr + sc_bases + 1, mask=kv_mask, other=0).to(
|
||||
tl.uint16
|
||||
)
|
||||
v_scales = (
|
||||
(sc_lo | (sc_hi << 8)).to(tl.float16, bitcast=True).to(tl.float32)
|
||||
)
|
||||
zr_lo = tl.load(KV_cache_ptr + sc_bases + 2, mask=kv_mask, other=0).to(
|
||||
tl.uint16
|
||||
)
|
||||
zr_hi = tl.load(KV_cache_ptr + sc_bases + 3, mask=kv_mask, other=0).to(
|
||||
tl.uint16
|
||||
)
|
||||
v_zeros = (zr_lo | (zr_hi << 8)).to(tl.float16, bitcast=True).to(tl.float32)
|
||||
values = v_idx * v_scales[:, None] + v_zeros[:, None]
|
||||
else: # VQB == 4
|
||||
vb_idx = d_offs // 2
|
||||
vb_shift = (d_offs % 2) * 4
|
||||
val_addrs = val_bases[:, None] + vb_idx[None, :]
|
||||
val_raw = tl.load(
|
||||
KV_cache_ptr + val_addrs,
|
||||
mask=kv_mask[:, None] & d_mask[None, :],
|
||||
other=0,
|
||||
).to(tl.int32)
|
||||
v_idx = ((val_raw >> vb_shift[None, :]) & 0xF).to(tl.float32)
|
||||
|
||||
sc_bases = val_bases + VAL_DATA_BYTES
|
||||
sc_lo = tl.load(KV_cache_ptr + sc_bases, mask=kv_mask, other=0).to(
|
||||
tl.uint16
|
||||
)
|
||||
sc_hi = tl.load(KV_cache_ptr + sc_bases + 1, mask=kv_mask, other=0).to(
|
||||
tl.uint16
|
||||
)
|
||||
v_scales = (
|
||||
(sc_lo | (sc_hi << 8)).to(tl.float16, bitcast=True).to(tl.float32)
|
||||
)
|
||||
zr_lo = tl.load(KV_cache_ptr + sc_bases + 2, mask=kv_mask, other=0).to(
|
||||
tl.uint16
|
||||
)
|
||||
zr_hi = tl.load(KV_cache_ptr + sc_bases + 3, mask=kv_mask, other=0).to(
|
||||
tl.uint16
|
||||
)
|
||||
v_zeros = (zr_lo | (zr_hi << 8)).to(tl.float16, bitcast=True).to(tl.float32)
|
||||
values = v_idx * v_scales[:, None] + v_zeros[:, None]
|
||||
|
||||
# ============================================================
|
||||
# WEIGHTED VALUE ACCUMULATION
|
||||
# ============================================================
|
||||
acc = acc * re_scale + tl.sum(p[:, None] * values, 0)
|
||||
l_prev = l_prev * re_scale + tl.sum(p, 0)
|
||||
m_prev = n_e_max
|
||||
|
||||
# Store partial result
|
||||
out_base = bid * stride_mid_b + hid * stride_mid_h + sid * stride_mid_s
|
||||
safe_l = tl.where(l_prev > 0.0, l_prev, 1.0)
|
||||
tl.store(Mid_o_ptr + out_base + d_offs, acc / safe_l, mask=d_mask)
|
||||
lse = m_prev + tl.log(safe_l)
|
||||
tl.store(Mid_o_ptr + out_base + HEAD_DIM, lse)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Pre-dequant kernel: Bulk dequant K (MSE+norms) and V to fp16
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _tq_full_dequant_kv(
|
||||
KV_cache_ptr,
|
||||
Block_table_ptr,
|
||||
Centroids_ptr,
|
||||
K_out_ptr, # [B, Hk, max_seq, D] float16
|
||||
V_out_ptr, # [B, Hk, max_seq, D] float16
|
||||
stride_ko_b,
|
||||
stride_ko_h,
|
||||
stride_ko_s,
|
||||
stride_vo_b,
|
||||
stride_vo_h,
|
||||
stride_vo_s,
|
||||
stride_cache_block,
|
||||
stride_cache_pos,
|
||||
stride_cache_head,
|
||||
stride_bt_b,
|
||||
HEAD_DIM: tl.constexpr,
|
||||
BLOCK_SIZE: tl.constexpr,
|
||||
NUM_KV_HEADS: tl.constexpr,
|
||||
MSE_BYTES: tl.constexpr,
|
||||
KPS: tl.constexpr,
|
||||
VQB: tl.constexpr,
|
||||
VAL_DATA_BYTES: tl.constexpr,
|
||||
MSE_BITS: tl.constexpr,
|
||||
KEY_FP8: tl.constexpr,
|
||||
BLOCK_D: tl.constexpr,
|
||||
NORM_CORRECTION: tl.constexpr = 0,
|
||||
FP8_E4B15: tl.constexpr = 0, # 1 = use e4b15 (Ampere/Ada), 0 = e4nv (Hopper+)
|
||||
):
|
||||
"""Full dequant: reconstruct K (MSE centroids * norm or FP8) and V to fp16."""
|
||||
pos = tl.program_id(0)
|
||||
bh = tl.program_id(1)
|
||||
bid = bh // NUM_KV_HEADS
|
||||
hid = bh % NUM_KV_HEADS
|
||||
|
||||
page_idx = pos // BLOCK_SIZE
|
||||
page_off = pos % BLOCK_SIZE
|
||||
block_num = tl.load(Block_table_ptr + bid * stride_bt_b + page_idx).to(tl.int64)
|
||||
slot_base = (
|
||||
block_num * stride_cache_block
|
||||
+ tl.cast(page_off, tl.int64) * stride_cache_pos
|
||||
+ tl.cast(hid, tl.int64) * stride_cache_head
|
||||
)
|
||||
|
||||
d_offs = tl.arange(0, BLOCK_D)
|
||||
d_mask = d_offs < HEAD_DIM
|
||||
|
||||
# === K dequant ===
|
||||
ko_base = bid * stride_ko_b + hid * stride_ko_h + pos * stride_ko_s
|
||||
if KEY_FP8:
|
||||
k_raw = tl.load(KV_cache_ptr + slot_base + d_offs, mask=d_mask, other=0)
|
||||
if FP8_E4B15:
|
||||
k_recon = k_raw.to(tl.float8e4b15, bitcast=True).to(tl.float32)
|
||||
else:
|
||||
k_recon = k_raw.to(tl.float8e4nv, bitcast=True).to(tl.float32)
|
||||
tl.store(K_out_ptr + ko_base + d_offs, k_recon.to(tl.float16), mask=d_mask)
|
||||
else:
|
||||
# MSE unpack (3-bit or 4-bit) + norms
|
||||
mse_bit_off = d_offs * MSE_BITS
|
||||
mse_byte_idx = mse_bit_off // 8
|
||||
mse_bit_shift = mse_bit_off % 8
|
||||
mse_umask = (1 << MSE_BITS) - 1
|
||||
|
||||
mse_raw0 = tl.load(
|
||||
KV_cache_ptr + slot_base + mse_byte_idx, mask=d_mask, other=0
|
||||
).to(tl.int32)
|
||||
mse_raw1 = tl.load(
|
||||
KV_cache_ptr + slot_base + mse_byte_idx + 1, mask=d_mask, other=0
|
||||
).to(tl.int32)
|
||||
raw16_key = mse_raw0 | (mse_raw1 << 8)
|
||||
mse_idx = (raw16_key >> mse_bit_shift) & mse_umask
|
||||
|
||||
k_mse = tl.load(Centroids_ptr + mse_idx, mask=d_mask, other=0.0)
|
||||
|
||||
# Norm correction: re-normalize centroid vector to unit norm
|
||||
if NORM_CORRECTION:
|
||||
c_norm_sq = tl.sum(tl.where(d_mask, k_mse * k_mse, 0.0), axis=0)
|
||||
c_inv_norm = 1.0 / tl.sqrt(c_norm_sq + 1e-16)
|
||||
k_mse = k_mse * c_inv_norm
|
||||
|
||||
# Norms at MSE_BYTES offset (no QJL bytes)
|
||||
norm_base = slot_base + MSE_BYTES
|
||||
n_lo = tl.load(KV_cache_ptr + norm_base).to(tl.uint16)
|
||||
n_hi = tl.load(KV_cache_ptr + norm_base + 1).to(tl.uint16)
|
||||
vec_norm = (n_lo | (n_hi << 8)).to(tl.float16, bitcast=True).to(tl.float32)
|
||||
|
||||
k_recon = vec_norm * k_mse
|
||||
tl.store(K_out_ptr + ko_base + d_offs, k_recon.to(tl.float16), mask=d_mask)
|
||||
|
||||
# === V dequant ===
|
||||
val_base = slot_base + KPS
|
||||
if VQB == 4:
|
||||
vb_idx = d_offs // 2
|
||||
vb_shift = (d_offs % 2) * 4
|
||||
val_raw = tl.load(KV_cache_ptr + val_base + vb_idx, mask=d_mask, other=0).to(
|
||||
tl.int32
|
||||
)
|
||||
v_idx = ((val_raw >> vb_shift) & 0xF).to(tl.float32)
|
||||
|
||||
sc_base = val_base + VAL_DATA_BYTES
|
||||
sc_lo = tl.load(KV_cache_ptr + sc_base).to(tl.uint16)
|
||||
sc_hi = tl.load(KV_cache_ptr + sc_base + 1).to(tl.uint16)
|
||||
v_scale = (sc_lo | (sc_hi << 8)).to(tl.float16, bitcast=True).to(tl.float32)
|
||||
zr_lo = tl.load(KV_cache_ptr + sc_base + 2).to(tl.uint16)
|
||||
zr_hi = tl.load(KV_cache_ptr + sc_base + 3).to(tl.uint16)
|
||||
v_zero = (zr_lo | (zr_hi << 8)).to(tl.float16, bitcast=True).to(tl.float32)
|
||||
v_vals = v_idx * v_scale + v_zero
|
||||
elif VQB == 3:
|
||||
# 3-bit value unpack: 8 values per 3 bytes
|
||||
val_bit_off = d_offs * 3
|
||||
val_byte_idx = val_bit_off // 8
|
||||
val_bit_shift = val_bit_off % 8
|
||||
val_raw0 = tl.load(
|
||||
KV_cache_ptr + val_base + val_byte_idx, mask=d_mask, other=0
|
||||
).to(tl.int32)
|
||||
val_raw1 = tl.load(
|
||||
KV_cache_ptr + val_base + val_byte_idx + 1, mask=d_mask, other=0
|
||||
).to(tl.int32)
|
||||
raw16_val = val_raw0 | (val_raw1 << 8)
|
||||
v_idx = ((raw16_val >> val_bit_shift) & 0x7).to(tl.float32)
|
||||
|
||||
sc_base = val_base + VAL_DATA_BYTES
|
||||
sc_lo = tl.load(KV_cache_ptr + sc_base).to(tl.uint16)
|
||||
sc_hi = tl.load(KV_cache_ptr + sc_base + 1).to(tl.uint16)
|
||||
v_scale = (sc_lo | (sc_hi << 8)).to(tl.float16, bitcast=True).to(tl.float32)
|
||||
zr_lo = tl.load(KV_cache_ptr + sc_base + 2).to(tl.uint16)
|
||||
zr_hi = tl.load(KV_cache_ptr + sc_base + 3).to(tl.uint16)
|
||||
v_zero = (zr_lo | (zr_hi << 8)).to(tl.float16, bitcast=True).to(tl.float32)
|
||||
v_vals = v_idx * v_scale + v_zero
|
||||
else:
|
||||
v_vals = tl.zeros([BLOCK_D], dtype=tl.float32)
|
||||
|
||||
vo_base = bid * stride_vo_b + hid * stride_vo_h + pos * stride_vo_s
|
||||
tl.store(V_out_ptr + vo_base + d_offs, v_vals.to(tl.float16), mask=d_mask)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Stage 2: Reuse from triton_decode_attention.py
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Launcher — cached constants + fused GEMM
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_layout_cache: dict = {}
|
||||
|
||||
|
||||
def _get_layout(D, mse_bits, value_quant_bits, key_packed_size):
|
||||
"""Get cached layout constants."""
|
||||
key = (D, mse_bits, value_quant_bits, key_packed_size)
|
||||
cfg = _layout_cache.get(key)
|
||||
if cfg is None:
|
||||
val_data_bytes = math.ceil(D * value_quant_bits / 8)
|
||||
cfg = {
|
||||
"mse_bytes": math.ceil(D * mse_bits / 8),
|
||||
"val_data_bytes": val_data_bytes,
|
||||
"mse_bits": mse_bits,
|
||||
"n_centroids": 2**mse_bits,
|
||||
"BLOCK_D": triton.next_power_of_2(D),
|
||||
}
|
||||
_layout_cache[key] = cfg
|
||||
return cfg
|
||||
|
||||
|
||||
def triton_turboquant_decode_attention(
|
||||
query: torch.Tensor, # [B, Hq, D] — original query
|
||||
kv_cache: torch.Tensor, # [num_blocks, block_size, Hk, padded_slot] uint8
|
||||
block_table: torch.Tensor, # [B, max_num_blocks] int32
|
||||
seq_lens: torch.Tensor, # [B] int32
|
||||
Pi: torch.Tensor, # [D, D] float32
|
||||
centroids: torch.Tensor, # [n_centroids] float32
|
||||
scale: float,
|
||||
mse_bits: int,
|
||||
key_packed_size: int,
|
||||
value_quant_bits: int,
|
||||
key_fp8: bool = False,
|
||||
norm_correction: bool = False,
|
||||
PiT: torch.Tensor | None = None, # [D, D] pre-computed Pi.T contiguous
|
||||
# Pre-allocated buffers (optional, avoids per-call allocation)
|
||||
mid_o_buf: torch.Tensor | None = None,
|
||||
output_buf: torch.Tensor | None = None,
|
||||
lse_buf: torch.Tensor | None = None,
|
||||
max_num_kv_splits: int = 32, # fixed split count (must be constant for cudagraph)
|
||||
) -> torch.Tensor:
|
||||
"""Launch fused TQ decode attention (Triton stage1 + stage2).
|
||||
|
||||
Returns: output tensor [B, Hq, D] in query's dtype.
|
||||
"""
|
||||
B, Hq, D = query.shape
|
||||
Hk = kv_cache.shape[2]
|
||||
block_size = kv_cache.shape[1]
|
||||
kv_group_size = Hq // Hk
|
||||
device = query.device
|
||||
|
||||
cfg = _get_layout(D, mse_bits, value_quant_bits, key_packed_size)
|
||||
|
||||
# Compute q_rot = q @ Pi.T (rotated query for MSE key scoring)
|
||||
# FP8 path: pass query directly (float16); kernel casts inline.
|
||||
# MSE path: still needs external GEMM (cuBLAS), so q_rot is float32.
|
||||
if key_fp8:
|
||||
q_rot = query.contiguous()
|
||||
else:
|
||||
q_float = query.float()
|
||||
if PiT is None:
|
||||
PiT = Pi.T.contiguous()
|
||||
q_rot = (q_float @ PiT).contiguous()
|
||||
|
||||
NUM_KV_SPLITS = max_num_kv_splits
|
||||
|
||||
if mid_o_buf is None or output_buf is None or lse_buf is None:
|
||||
from vllm.v1.worker.workspace import (
|
||||
current_workspace_manager,
|
||||
is_workspace_manager_initialized,
|
||||
)
|
||||
|
||||
if is_workspace_manager_initialized():
|
||||
mid_o_buf, output_buf, lse_buf = (
|
||||
current_workspace_manager().get_simultaneous(
|
||||
((B, Hq, NUM_KV_SPLITS, D + 1), torch.float32),
|
||||
((B, Hq, D), query.dtype),
|
||||
((B, Hq), torch.float32),
|
||||
)
|
||||
)
|
||||
|
||||
if (
|
||||
mid_o_buf is not None
|
||||
and mid_o_buf.shape[0] >= B
|
||||
and mid_o_buf.shape[2] >= NUM_KV_SPLITS
|
||||
):
|
||||
mid_o = mid_o_buf[:B, :Hq, :NUM_KV_SPLITS, :]
|
||||
else:
|
||||
mid_o = torch.empty(
|
||||
B,
|
||||
Hq,
|
||||
NUM_KV_SPLITS,
|
||||
D + 1,
|
||||
dtype=torch.float32,
|
||||
device=device,
|
||||
)
|
||||
|
||||
# Stage 1: split-KV tiled attention scoring + value accumulation
|
||||
fp8_e4b15 = _use_fp8_e4b15(device.index or 0)
|
||||
BLOCK_KV = 4
|
||||
grid = (B, Hq, NUM_KV_SPLITS)
|
||||
_tq_decode_stage1[grid](
|
||||
q_rot,
|
||||
kv_cache,
|
||||
block_table,
|
||||
seq_lens,
|
||||
centroids,
|
||||
mid_o,
|
||||
q_rot.stride(0),
|
||||
q_rot.stride(1),
|
||||
kv_cache.stride(0),
|
||||
kv_cache.stride(1),
|
||||
kv_cache.stride(2),
|
||||
block_table.stride(0),
|
||||
mid_o.stride(0),
|
||||
mid_o.stride(1),
|
||||
mid_o.stride(2),
|
||||
NUM_KV_HEADS=Hk,
|
||||
HEAD_DIM=D,
|
||||
BLOCK_SIZE=block_size,
|
||||
NUM_KV_SPLITS=NUM_KV_SPLITS,
|
||||
KV_GROUP_SIZE=kv_group_size,
|
||||
MSE_BITS=mse_bits,
|
||||
MSE_BYTES=cfg["mse_bytes"],
|
||||
KPS=key_packed_size,
|
||||
VQB=value_quant_bits,
|
||||
VAL_DATA_BYTES=cfg["val_data_bytes"],
|
||||
ATTN_SCALE=scale,
|
||||
BLOCK_D=cfg["BLOCK_D"],
|
||||
BLOCK_KV=BLOCK_KV,
|
||||
KEY_FP8=1 if key_fp8 else 0,
|
||||
NORM_CORRECTION=1 if norm_correction else 0,
|
||||
FP8_E4B15=fp8_e4b15,
|
||||
num_warps=1,
|
||||
num_stages=1,
|
||||
)
|
||||
|
||||
# Stage 2: Reduce across KV splits
|
||||
# Output in query dtype — eliminates float16_copy kernel after stage2
|
||||
out_dtype = query.dtype
|
||||
if (
|
||||
output_buf is not None
|
||||
and output_buf.shape[0] >= B
|
||||
and output_buf.dtype == out_dtype
|
||||
):
|
||||
output = output_buf[:B, :Hq, :D]
|
||||
else:
|
||||
output = torch.empty(B, Hq, D, dtype=out_dtype, device=device)
|
||||
if lse_buf is not None and lse_buf.shape[0] >= B:
|
||||
lse = lse_buf[:B, :Hq]
|
||||
else:
|
||||
lse = torch.empty(B, Hq, dtype=torch.float32, device=device)
|
||||
|
||||
grid2 = (B, Hq)
|
||||
_fwd_kernel_stage2[grid2](
|
||||
mid_o,
|
||||
output,
|
||||
lse,
|
||||
seq_lens,
|
||||
mid_o.stride(0),
|
||||
mid_o.stride(1),
|
||||
mid_o.stride(2),
|
||||
output.stride(0),
|
||||
output.stride(1),
|
||||
lse.stride(0),
|
||||
NUM_KV_SPLITS=NUM_KV_SPLITS,
|
||||
BLOCK_DV=cfg["BLOCK_D"],
|
||||
Lv=D,
|
||||
OUTPUT_FP16=1 if out_dtype == torch.float16 else 0,
|
||||
num_warps=4,
|
||||
num_stages=2,
|
||||
)
|
||||
|
||||
return output # already in query dtype
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,71 @@
|
||||
# vLLM PR #40914 K+1-only overlay — manually rebased onto post-#41434 main
|
||||
|
||||
Vendored 2026-05-11 to unblock TQ3 + MTP on Qwen 3.6 27B (`dual/int8-tq3.yml`)
|
||||
with proper spec-verify routing.
|
||||
|
||||
## Source
|
||||
|
||||
- Upstream PR: <https://github.com/vllm-project/vllm/pull/40914>
|
||||
("[Bugfix][Spec-Decode] TurboQuant K+1 spec-verify routing (fixes #40880)")
|
||||
- Upstream head SHA: `0ee9b859bbb2bbb6e33a461e7fd1fee1fa4792cc` (2026-04-29)
|
||||
- State at vendor time: OPEN, MERGEABLE.
|
||||
|
||||
## Why this is a manual rebase
|
||||
|
||||
PR #40914 was forked from main BEFORE merged PR #41434 ("Eliminate
|
||||
GPU↔CPU syncs in attention impls", merged 2026-05-08). Its full
|
||||
`turboquant_attn.py` overlay would REVERT #41434's `query_start_loc_cpu`
|
||||
/ `seq_lens_cpu` CPU-resident-metadata fixes and reintroduce the
|
||||
`.tolist()` cudagraph crash family.
|
||||
|
||||
We need ONLY the additive K+1 spec-verify dispatch block (~86 lines)
|
||||
inserted between `num_decode_tokens = attn_metadata.num_decode_tokens`
|
||||
and `if not attn_metadata.is_prefill:`. All other diff hunks from PR
|
||||
#40914 are reverts of #41434 and must be skipped.
|
||||
|
||||
## What this overlay contains
|
||||
|
||||
`turboquant_attn.py` = current `origin/main` (post-#41434) + ONLY the
|
||||
additive K+1 dispatch block from PR #40914 inserted at the right location.
|
||||
|
||||
Result:
|
||||
- `query_start_loc_cpu` / `seq_lens_cpu` fields and accessors: PRESERVED
|
||||
- `.tolist()` → CPU-resident metadata fix: PRESERVED
|
||||
- K+1 spec-verify dispatch from #40914: ADDED
|
||||
- `buf_holder=layer` call site: matches nightly's `triton_turboquant_decode.py`
|
||||
signature (which post-#41434 accepts that kwarg)
|
||||
|
||||
## What this fixes
|
||||
|
||||
Bimodal MTP acceptance + "first word right then breaks" needle failure
|
||||
on TQ3 + MTP. Symptom on the previous stack (PR #40798 partial only):
|
||||
|
||||
```
|
||||
12:22:02 AL=3.73 accept 91.1% per-position [0.962, 0.930, 0.841] ← GOOD
|
||||
12:22:32 AL=2.30 accept 43.4% per-position [1.000, 0.303, 0.000] ← BAD
|
||||
12:23:42 AL=2.00 accept 33.2% per-position [0.397, 0.306, 0.294] ← VERY BAD
|
||||
```
|
||||
|
||||
Verify-stress needle 10K/30K/60K/90K all failed with `expected 'golden
|
||||
chinchilla 38' got 'golden '` pattern — first token recalled, then
|
||||
attention diverged because the verify pass was attending only to
|
||||
current-chunk K/V instead of prior cached compressed K/V.
|
||||
|
||||
This dispatch routes uniform K+1 spec-verify batches through the decode
|
||||
kernel (which natively reads prior cached K/V), restoring correctness.
|
||||
|
||||
## Drop trigger
|
||||
|
||||
```
|
||||
gh api repos/vllm-project/vllm/pulls/40914 --jq '.state, .merged_at'
|
||||
```
|
||||
|
||||
reports `MERGED` AND the merge commit is post-#41434 (i.e., the PR was
|
||||
rebased before merge so it doesn't revert #41434).
|
||||
|
||||
## Verified on
|
||||
|
||||
- vLLM nightly: `1acd67a7` (includes merged #41434)
|
||||
- Compose: `dual/int8-tq3.yml`
|
||||
- Qwen 3.6 27B AutoRound INT4, TQ3 KV, MTP n=3, 262K × 2 streams
|
||||
- 2026-05-11
|
||||
+988
@@ -0,0 +1,988 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""TurboQuant attention backend for vLLM.
|
||||
|
||||
Prefill: Standard scaled dot-product attention on uncompressed K/V,
|
||||
then quantize K and store K+V into combined cache slot.
|
||||
Decode: Compute TQ attention scores from compressed cache,
|
||||
unpack FP16 values, softmax + weighted sum.
|
||||
|
||||
Cache layout (no leading 2 dimension):
|
||||
(num_blocks, block_size, num_kv_heads, slot_size)
|
||||
where slot_size = key_packed_size + value_fp16_size
|
||||
|
||||
Per-head per-position slot layout:
|
||||
[key_packed (kps bytes) | value_fp16 (D*2 bytes)]
|
||||
For turboquant_k3v4_nc head_dim=256: [100 bytes key | 512 bytes value] = 612
|
||||
"""
|
||||
|
||||
import functools
|
||||
import math
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, ClassVar
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
from vllm.config import get_current_vllm_config
|
||||
from vllm.config.cache import CacheDType
|
||||
from vllm.model_executor.layers.quantization.turboquant.centroids import (
|
||||
get_centroids,
|
||||
)
|
||||
from vllm.triton_utils import triton
|
||||
from vllm.v1.attention.backend import (
|
||||
AttentionBackend,
|
||||
AttentionCGSupport,
|
||||
AttentionImpl,
|
||||
AttentionLayer,
|
||||
AttentionMetadata,
|
||||
AttentionMetadataBuilder,
|
||||
AttentionType,
|
||||
CommonAttentionMetadata,
|
||||
MultipleOf,
|
||||
)
|
||||
from vllm.v1.attention.backends.fa_utils import (
|
||||
get_flash_attn_version,
|
||||
is_flash_attn_varlen_func_available,
|
||||
)
|
||||
from vllm.v1.attention.backends.utils import split_decodes_and_prefills
|
||||
from vllm.v1.attention.ops.triton_turboquant_decode import (
|
||||
_tq_full_dequant_kv,
|
||||
_use_fp8_e4b15,
|
||||
triton_turboquant_decode_attention,
|
||||
)
|
||||
from vllm.v1.attention.ops.triton_turboquant_store import triton_turboquant_store
|
||||
from vllm.v1.worker.workspace import (
|
||||
current_workspace_manager,
|
||||
is_workspace_manager_initialized,
|
||||
)
|
||||
|
||||
_HAS_FLASH_ATTN = is_flash_attn_varlen_func_available()
|
||||
if _HAS_FLASH_ATTN:
|
||||
from vllm.v1.attention.backends.fa_utils import flash_attn_varlen_func
|
||||
|
||||
# Continuation prefill: for small continuation chunks (q_len ≤ threshold),
|
||||
# use the TQ decode kernel directly instead of full-dequant + flash_attn.
|
||||
# do_kv_cache_update already stored all tokens to TQ cache, so the decode
|
||||
# kernel can read them efficiently. This avoids O(cached_len) dequant work
|
||||
# per continuation, eliminating the O(N²/chunk_size) collapse at long context.
|
||||
_CONTINUATION_DECODE_THRESHOLD = 128
|
||||
|
||||
|
||||
def _build_hadamard(d: int, device_str: str) -> torch.Tensor:
|
||||
"""Orthonormal Hadamard matrix (Sylvester construction), cached per (d, device).
|
||||
|
||||
Precomputed D×D matrix enables matmul-based WHT — single cuBLAS GEMM
|
||||
instead of log2(D) butterfly kernel launches. 64KB for D=128.
|
||||
"""
|
||||
# Normalize device string so "cuda" and "cuda:0" hit the same cache entry.
|
||||
return _build_hadamard_cached(d, str(torch.device(device_str)))
|
||||
|
||||
|
||||
@functools.cache
|
||||
def _build_hadamard_cached(d: int, device_str: str) -> torch.Tensor:
|
||||
H = torch.tensor([[1.0]])
|
||||
while H.shape[0] < d:
|
||||
H = torch.cat([torch.cat([H, H], 1), torch.cat([H, -H], 1)], 0)
|
||||
return (H / math.sqrt(d)).to(torch.device(device_str))
|
||||
|
||||
|
||||
class TurboQuantAttentionBackend(AttentionBackend):
|
||||
"""Attention backend using TurboQuant KV-cache compression."""
|
||||
|
||||
accept_output_buffer: bool = True
|
||||
forward_includes_kv_cache_update: bool = False
|
||||
|
||||
supported_dtypes: ClassVar[list[torch.dtype]] = [
|
||||
torch.float16,
|
||||
torch.bfloat16,
|
||||
]
|
||||
supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
|
||||
"turboquant_k8v4",
|
||||
"turboquant_4bit_nc",
|
||||
"turboquant_k3v4_nc",
|
||||
"turboquant_3bit_nc",
|
||||
]
|
||||
|
||||
@staticmethod
|
||||
def get_name() -> str:
|
||||
return "TURBOQUANT"
|
||||
|
||||
@staticmethod
|
||||
def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
|
||||
return [16, 32, 64, 128]
|
||||
|
||||
@classmethod
|
||||
def supports_attn_type(cls, attn_type: str) -> bool:
|
||||
return attn_type == AttentionType.DECODER
|
||||
|
||||
@classmethod
|
||||
def supports_per_head_quant_scales(cls) -> bool:
|
||||
return False
|
||||
|
||||
@staticmethod
|
||||
def get_impl_cls() -> type["TurboQuantAttentionImpl"]:
|
||||
return TurboQuantAttentionImpl
|
||||
|
||||
@staticmethod
|
||||
def get_builder_cls() -> type["TurboQuantMetadataBuilder"]:
|
||||
return TurboQuantMetadataBuilder
|
||||
|
||||
@staticmethod
|
||||
def get_kv_cache_shape(
|
||||
num_blocks: int,
|
||||
block_size: int,
|
||||
num_kv_heads: int,
|
||||
head_size: int,
|
||||
cache_dtype_str: str = "turboquant_4bit_nc",
|
||||
) -> tuple[int, ...]:
|
||||
"""Combined K+V cache shape — no leading 2 dimension.
|
||||
|
||||
Standard attention backends use (2, num_blocks, block_size, num_kv_heads,
|
||||
head_dim) with a leading 2 to separate K and V. TurboQuant packs K+V
|
||||
into a single interleaved slot per head per position, so the cache is:
|
||||
|
||||
(num_blocks, block_size, num_kv_heads, slot_size_aligned)
|
||||
|
||||
Each slot = [key_packed | value_packed | padding].
|
||||
This is safe because TQ has its own get_kv_cache_shape override and
|
||||
never shares cache tensors with other backends. Layers that fall back
|
||||
to native dtype via kv_cache_dtype_skip_layers get their own
|
||||
standard-shaped cache allocation.
|
||||
|
||||
head_size is the model's real head_dim. slot_size_aligned is computed
|
||||
from the TQ config to ensure correct cache allocation for all head dims.
|
||||
"""
|
||||
from vllm.model_executor.layers.quantization.turboquant.config import (
|
||||
TurboQuantConfig,
|
||||
)
|
||||
|
||||
tq_config = TurboQuantConfig.from_cache_dtype(cache_dtype_str, head_size)
|
||||
return (num_blocks, block_size, num_kv_heads, tq_config.slot_size_aligned)
|
||||
|
||||
@classmethod
|
||||
def supports_kv_cache_dtype(cls, kv_cache_dtype: CacheDType | None) -> bool:
|
||||
if kv_cache_dtype is None:
|
||||
return False
|
||||
return kv_cache_dtype.startswith("turboquant_")
|
||||
|
||||
@classmethod
|
||||
def supports_head_size(cls, head_size: int) -> bool:
|
||||
# head_size from spec is effective_head_size (padded_slot//2),
|
||||
# not the model's actual head_dim. Accept any positive value.
|
||||
return head_size > 0
|
||||
|
||||
|
||||
@dataclass
|
||||
class TurboQuantMetadata(AttentionMetadata):
|
||||
"""Metadata for TurboQuant attention."""
|
||||
|
||||
seq_lens: torch.Tensor # (num_reqs,) — total context length per request
|
||||
slot_mapping: torch.Tensor # (num_tokens,) — cache slot for each token
|
||||
block_table: torch.Tensor # (num_reqs, max_num_blocks)
|
||||
query_start_loc: torch.Tensor # (num_reqs + 1,) — cu_seqlens for queries
|
||||
num_actual_tokens: int = 0 # actual tokens (excluding padding)
|
||||
max_query_len: int = 0 # longest query in batch
|
||||
max_seq_len: int = 0 # longest context in batch
|
||||
is_prefill: bool = False
|
||||
num_decodes: int = 0 # number of decode requests (first in batch)
|
||||
num_decode_tokens: int = 0 # tokens from decode requests
|
||||
# CPU-resident copies used by the prefill path for per-request iteration
|
||||
# without per-step D2H syncs.
|
||||
query_start_loc_cpu: torch.Tensor | None = None
|
||||
seq_lens_cpu: torch.Tensor | None = None
|
||||
|
||||
|
||||
class TurboQuantMetadataBuilder(AttentionMetadataBuilder[TurboQuantMetadata]):
|
||||
"""Builds TurboQuantMetadata from scheduler output."""
|
||||
|
||||
_cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
|
||||
|
||||
def __init__(self, kv_cache_spec, layer_names, vllm_config, device):
|
||||
super().__init__(kv_cache_spec, layer_names, vllm_config, device)
|
||||
self._init_reorder_batch_threshold(1, supports_spec_as_decode=False)
|
||||
|
||||
def build_for_cudagraph_capture(
|
||||
self, common_attn_metadata: CommonAttentionMetadata
|
||||
) -> TurboQuantMetadata:
|
||||
attn_metadata = self.build(0, common_attn_metadata)
|
||||
# Set seq_lens to 1 so CUDA graph capture is fast
|
||||
# (real seq_lens are filled at replay time).
|
||||
attn_metadata.seq_lens.fill_(1)
|
||||
return attn_metadata
|
||||
|
||||
def build(self, common_prefix_len, common_attn_metadata, fast_build=False):
|
||||
"""Build TurboQuantMetadata from common attention metadata."""
|
||||
cam = common_attn_metadata
|
||||
|
||||
# With reorder_batch_threshold=1, the model runner guarantees
|
||||
# decodes come first in the batch. split_decodes_and_prefills
|
||||
# finds the boundary (operates on CPU tensors — no GPU sync).
|
||||
assert self.reorder_batch_threshold is not None
|
||||
num_decodes, num_prefills, num_decode_tokens, _ = split_decodes_and_prefills(
|
||||
cam, decode_threshold=self.reorder_batch_threshold
|
||||
)
|
||||
|
||||
return TurboQuantMetadata(
|
||||
seq_lens=cam.seq_lens,
|
||||
slot_mapping=cam.slot_mapping,
|
||||
block_table=cam.block_table_tensor,
|
||||
query_start_loc=cam.query_start_loc,
|
||||
num_actual_tokens=cam.num_actual_tokens,
|
||||
max_query_len=cam.max_query_len,
|
||||
max_seq_len=cam.max_seq_len,
|
||||
is_prefill=(cam.max_query_len > 1),
|
||||
num_decodes=num_decodes,
|
||||
num_decode_tokens=num_decode_tokens,
|
||||
query_start_loc_cpu=cam.query_start_loc_cpu,
|
||||
seq_lens_cpu=cam.seq_lens_cpu_upper_bound,
|
||||
)
|
||||
|
||||
|
||||
class TurboQuantAttentionImpl(AttentionImpl["TurboQuantMetadata"]):
|
||||
"""TurboQuant attention implementation.
|
||||
|
||||
Vectorized PyTorch: batch quantize/store, vectorized bit-unpack
|
||||
decode with einsum scores and value gather.
|
||||
"""
|
||||
|
||||
supports_quant_query_input: bool = False
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
num_heads: int,
|
||||
head_size: int,
|
||||
scale: float,
|
||||
num_kv_heads: int | None = None,
|
||||
alibi_slopes: list[float] | None = None,
|
||||
sliding_window: int | None = None,
|
||||
kv_cache_dtype: str = "auto",
|
||||
logits_soft_cap: float | None = None,
|
||||
attn_type: str = AttentionType.DECODER,
|
||||
kv_sharing_target_layer_name: str | None = None,
|
||||
**kwargs,
|
||||
):
|
||||
self.num_heads = num_heads
|
||||
self.head_size = head_size
|
||||
self.scale = scale
|
||||
self.num_kv_heads = num_kv_heads if num_kv_heads is not None else num_heads
|
||||
self.num_kv_groups = num_heads // self.num_kv_heads
|
||||
self.kv_cache_dtype = kv_cache_dtype
|
||||
|
||||
from vllm.model_executor.layers.quantization.turboquant.config import (
|
||||
TurboQuantConfig,
|
||||
)
|
||||
|
||||
self.tq_config = TurboQuantConfig.from_cache_dtype(kv_cache_dtype, head_size)
|
||||
|
||||
# Pre-compute kernel constants from config (avoid repeated arithmetic)
|
||||
cfg = self.tq_config
|
||||
self._mse_bytes = (
|
||||
math.ceil(head_size * cfg.key_mse_bits / 8)
|
||||
if not cfg.key_fp8
|
||||
else head_size
|
||||
)
|
||||
self._val_data_bytes = math.ceil(head_size * cfg.effective_value_quant_bits / 8)
|
||||
self._n_centroids = cfg.n_centroids if not cfg.key_fp8 else 1
|
||||
|
||||
# Detect flash-attn version (FA2/3/4) for prefill paths.
|
||||
self.fa_version = get_flash_attn_version(head_size=head_size)
|
||||
|
||||
# Fixed NUM_KV_SPLITS (grid dims must be constant for cudagraph,
|
||||
# and benchmarks show no regression vs dynamic in eager mode).
|
||||
vllm_config = get_current_vllm_config()
|
||||
self.max_num_kv_splits = (
|
||||
vllm_config.attention_config.tq_max_kv_splits_for_cuda_graph
|
||||
)
|
||||
|
||||
def _flash_attn_varlen(
|
||||
self,
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
cu_seqlens_q: torch.Tensor,
|
||||
cu_seqlens_k: torch.Tensor,
|
||||
max_seqlen_q: int,
|
||||
max_seqlen_k: int,
|
||||
) -> torch.Tensor:
|
||||
# fa_utils.get_flash_attn_version() returns None on backends that
|
||||
# should not pass an explicit fa_version kwarg.
|
||||
if self.fa_version is None:
|
||||
return flash_attn_varlen_func(
|
||||
q=q,
|
||||
k=k,
|
||||
v=v,
|
||||
cu_seqlens_q=cu_seqlens_q,
|
||||
cu_seqlens_k=cu_seqlens_k,
|
||||
max_seqlen_q=max_seqlen_q,
|
||||
max_seqlen_k=max_seqlen_k,
|
||||
softmax_scale=self.scale,
|
||||
causal=True,
|
||||
)
|
||||
return flash_attn_varlen_func(
|
||||
q=q,
|
||||
k=k,
|
||||
v=v,
|
||||
cu_seqlens_q=cu_seqlens_q,
|
||||
cu_seqlens_k=cu_seqlens_k,
|
||||
max_seqlen_q=max_seqlen_q,
|
||||
max_seqlen_k=max_seqlen_k,
|
||||
softmax_scale=self.scale,
|
||||
causal=True,
|
||||
fa_version=self.fa_version,
|
||||
)
|
||||
|
||||
def _ensure_on_device(self, layer, device):
|
||||
"""One-time derivation of TQ buffers (rotation matrix, midpoints).
|
||||
|
||||
The Hadamard rotation is shared across all layers: random sign
|
||||
flips do not improve Lloyd-Max quantization quality because the
|
||||
quantizer is symmetric around zero (sign-flipping a coordinate
|
||||
maps it to the mirror centroid with identical distortion).
|
||||
"""
|
||||
if not hasattr(layer, "_tq_cached"):
|
||||
D = self.head_size
|
||||
|
||||
# Pure Hadamard: orthonormal + symmetric (H = H^T), enabling
|
||||
# in-kernel butterfly fusion and trivial inverse for continuation.
|
||||
H = _build_hadamard(D, str(device))
|
||||
layer._tq_PiT = H
|
||||
layer._tq_Pi = H
|
||||
# fp16 copy for rotation in continuation prefill path
|
||||
layer._tq_Pi_half = H.to(torch.float16)
|
||||
|
||||
# Centroids for Lloyd-Max quantization.
|
||||
layer._tq_centroids = get_centroids(D, self.tq_config.centroid_bits).to(
|
||||
device=device, dtype=torch.float32
|
||||
)
|
||||
|
||||
c_sorted, _ = layer._tq_centroids.sort()
|
||||
layer._tq_midpoints = (c_sorted[:-1] + c_sorted[1:]) / 2
|
||||
layer._tq_cached = True
|
||||
|
||||
def do_kv_cache_update(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
kv_cache: torch.Tensor,
|
||||
slot_mapping: torch.Tensor,
|
||||
) -> None:
|
||||
"""Store compressed K/V into the combined TQ cache.
|
||||
|
||||
Called as a separate custom op (unified_kv_cache_update) BEFORE
|
||||
the attention forward, matching FlashAttention's split pattern.
|
||||
slot_mapping is already sliced to num_actual_tokens by the caller.
|
||||
"""
|
||||
N = slot_mapping.shape[0]
|
||||
if N <= 0:
|
||||
return
|
||||
|
||||
device = key.device
|
||||
self._ensure_on_device(layer, device)
|
||||
|
||||
k = key[:N].view(N, self.num_kv_heads, self.head_size)
|
||||
v = value[:N].view(N, self.num_kv_heads, self.head_size)
|
||||
self._store_kv(k, v, kv_cache, slot_mapping, layer)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
layer: AttentionLayer,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
kv_cache: torch.Tensor,
|
||||
attn_metadata: "TurboQuantMetadata",
|
||||
output: torch.Tensor | None = None,
|
||||
output_scale: torch.Tensor | None = None,
|
||||
output_block_scale: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
num_tokens = query.shape[0]
|
||||
|
||||
if output is None:
|
||||
output = torch.zeros(
|
||||
num_tokens,
|
||||
self.num_heads * self.head_size,
|
||||
dtype=query.dtype,
|
||||
device=query.device,
|
||||
)
|
||||
|
||||
if attn_metadata is None:
|
||||
return output.fill_(0)
|
||||
|
||||
# Slice to actual tokens
|
||||
N = attn_metadata.num_actual_tokens
|
||||
if N <= 0:
|
||||
return output.fill_(0)
|
||||
|
||||
q = query[:N].view(N, self.num_heads, self.head_size)
|
||||
|
||||
# Get TQ buffers, ensure on device (one-time migration).
|
||||
# Use Any-typed alias for dynamic _tq_* attrs set by _ensure_on_device.
|
||||
tq_layer: Any = layer
|
||||
device = q.device
|
||||
self._ensure_on_device(tq_layer, device)
|
||||
Pi = tq_layer._tq_Pi
|
||||
PiT = tq_layer._tq_PiT
|
||||
centroids = tq_layer._tq_centroids
|
||||
|
||||
# Compute attention (KV cache was already updated by do_kv_cache_update)
|
||||
# With reorder_batch_threshold=1, decodes come first in the batch.
|
||||
# num_decodes/num_decode_tokens from metadata give the split point.
|
||||
num_decodes = attn_metadata.num_decodes
|
||||
num_decode_tokens = attn_metadata.num_decode_tokens
|
||||
|
||||
# ════════════════════════════════════════════════════════════════════
|
||||
# club3090 PR #40914 K+1 spec-verify routing — manually rebased onto
|
||||
# post-#41434 main 2026-05-11. The upstream PR's full file is pre-
|
||||
# #41434 and would revert the merged CPU-resident-metadata fixes;
|
||||
# this overlay ports ONLY the additive K+1 dispatch block.
|
||||
#
|
||||
# Fixes vllm-issue#40880 (degenerate token cascade on TQ + MTP).
|
||||
#
|
||||
# When MTP is active (num_speculative_tokens=K>0), the verify pass
|
||||
# produces uniform-query batches with max_query_len=K+1 (e.g. K=3 →
|
||||
# q_len=4) where max_seq_len > max_query_len (prior cached KV).
|
||||
# The default _prefill_attention continuation branch would read
|
||||
# query_start_loc.tolist() (fixed in #41434 elsewhere) and attend
|
||||
# only to current-chunk K/V instead of prior cached KV, causing
|
||||
# the bimodal MTP acceptance + first-word-right-then-breaks
|
||||
# needle failure we observed.
|
||||
#
|
||||
# Fix: detect uniform K+1 batches and route them through the
|
||||
# decode kernel (which natively handles compressed K+V cache
|
||||
# lookup and is cudagraph-safe). Synthesizes seq_lens/block_table
|
||||
# to mirror _continuation_prefill's pattern, all on GPU.
|
||||
#
|
||||
# Drop this overlay when PR #40914 lands upstream rebased onto
|
||||
# post-#41434 main.
|
||||
# ════════════════════════════════════════════════════════════════════
|
||||
_spec_verify_eligible = (
|
||||
attn_metadata.is_prefill
|
||||
and num_decodes == 0
|
||||
and 1 < attn_metadata.max_query_len <= 16
|
||||
and attn_metadata.max_seq_len > attn_metadata.max_query_len
|
||||
and N > 0
|
||||
and (N % attn_metadata.max_query_len) == 0
|
||||
and attn_metadata.query_start_loc is not None
|
||||
)
|
||||
if _spec_verify_eligible:
|
||||
K_PLUS_1 = attn_metadata.max_query_len
|
||||
B = N // K_PLUS_1
|
||||
if attn_metadata.query_start_loc.shape[0] == B + 1:
|
||||
# Build synth args mirroring _continuation_prefill's pattern:
|
||||
# synth_seq_lens[req*K1+i] = base_seq_lens[req] - K1 + 1 + i
|
||||
# synth_block_table[req*K1+i] = block_table[req]
|
||||
# All GPU ops — cudagraph-safe.
|
||||
_q_flat = q[:N].view(N, self.num_heads, self.head_size)
|
||||
_offs = torch.arange(
|
||||
K_PLUS_1, device=q.device,
|
||||
dtype=attn_metadata.seq_lens.dtype,
|
||||
)
|
||||
_synth_seq_lens = (
|
||||
attn_metadata.seq_lens[:B, None] - K_PLUS_1 + 1 + _offs[None, :]
|
||||
).reshape(-1)
|
||||
_synth_block_table = attn_metadata.block_table[:B].repeat_interleave(
|
||||
K_PLUS_1, dim=0,
|
||||
)
|
||||
# Reuse cached decode buffers from the layer to avoid
|
||||
# per-call torch.empty allocations — these would break
|
||||
# CUDA graph replay (the very thing this PR restores).
|
||||
# Per gemini-code-assist review on this PR.
|
||||
_mid_o_buf = getattr(layer, "_tq_mid_o_buf", None)
|
||||
_output_buf = getattr(layer, "_tq_output_buf", None)
|
||||
_lse_buf = getattr(layer, "_tq_lse_buf", None)
|
||||
attn_out = triton_turboquant_decode_attention(
|
||||
query=_q_flat,
|
||||
kv_cache=kv_cache,
|
||||
block_table=_synth_block_table,
|
||||
seq_lens=_synth_seq_lens,
|
||||
Pi=Pi,
|
||||
centroids=centroids,
|
||||
scale=self.scale,
|
||||
mse_bits=self.tq_config.key_mse_bits,
|
||||
key_packed_size=self.tq_config.key_packed_size,
|
||||
value_quant_bits=self.tq_config.effective_value_quant_bits,
|
||||
key_fp8=self.tq_config.key_fp8,
|
||||
norm_correction=self.tq_config.norm_correction,
|
||||
PiT=PiT,
|
||||
mid_o_buf=_mid_o_buf,
|
||||
output_buf=_output_buf,
|
||||
lse_buf=_lse_buf,
|
||||
buf_holder=layer,
|
||||
max_num_kv_splits=self.max_num_kv_splits,
|
||||
)
|
||||
return attn_out
|
||||
|
||||
if not attn_metadata.is_prefill:
|
||||
# Pure decode batch — fast path
|
||||
attn_out = self._decode_attention(
|
||||
q, kv_cache, attn_metadata, Pi, centroids, PiT, layer
|
||||
)
|
||||
elif num_decodes == 0:
|
||||
# Pure prefill batch
|
||||
k = key[:N].view(N, self.num_kv_heads, self.head_size)
|
||||
v = value[:N].view(N, self.num_kv_heads, self.head_size)
|
||||
attn_out = self._prefill_attention(
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
kv_cache,
|
||||
attn_metadata,
|
||||
Pi,
|
||||
centroids,
|
||||
PiT,
|
||||
layer=layer,
|
||||
)
|
||||
else:
|
||||
# Mixed batch: decodes first (guaranteed by reorder_batch).
|
||||
attn_out = torch.zeros(
|
||||
N, self.num_heads, self.head_size, device=device, dtype=q.dtype
|
||||
)
|
||||
|
||||
# --- Decode portion (first num_decodes requests) ---
|
||||
# Use full-batch max_seq_len as safe upper bound (no GPU sync).
|
||||
decode_meta = TurboQuantMetadata(
|
||||
seq_lens=attn_metadata.seq_lens[:num_decodes],
|
||||
slot_mapping=attn_metadata.slot_mapping[:num_decode_tokens],
|
||||
block_table=attn_metadata.block_table[:num_decodes],
|
||||
query_start_loc=attn_metadata.query_start_loc[: num_decodes + 1],
|
||||
num_actual_tokens=num_decode_tokens,
|
||||
max_query_len=1,
|
||||
max_seq_len=attn_metadata.max_seq_len,
|
||||
is_prefill=False,
|
||||
)
|
||||
attn_out[:num_decode_tokens] = self._decode_attention(
|
||||
q[:num_decode_tokens], kv_cache, decode_meta, Pi, centroids, PiT, layer
|
||||
)
|
||||
|
||||
# --- Prefill portion (remaining requests) ---
|
||||
# CRITICAL: use prefill-specific max_seq_len so flash_attn's
|
||||
# fast path (max_query_len == max_seq_len) triggers for
|
||||
# first-chunk prefills. Using full-batch max_seq_len breaks
|
||||
# this because decode requests inflate max_seq_len.
|
||||
prefill_seq_lens = attn_metadata.seq_lens[num_decodes:]
|
||||
# Use the CPU-resident `seq_lens` upper-bound from the metadata
|
||||
# (populated in the builder) to compute the prefill sub-batch
|
||||
# max without a GPU→CPU sync.
|
||||
if attn_metadata.seq_lens_cpu is not None:
|
||||
prefill_max_seq = int(attn_metadata.seq_lens_cpu[num_decodes:].max())
|
||||
else:
|
||||
prefill_max_seq = attn_metadata.max_seq_len
|
||||
prefill_qsl = (
|
||||
attn_metadata.query_start_loc[num_decodes:] - num_decode_tokens
|
||||
)
|
||||
prefill_qsl_cpu = None
|
||||
if attn_metadata.query_start_loc_cpu is not None:
|
||||
prefill_qsl_cpu = (
|
||||
attn_metadata.query_start_loc_cpu[num_decodes:] - num_decode_tokens
|
||||
)
|
||||
prefill_meta = TurboQuantMetadata(
|
||||
seq_lens=prefill_seq_lens,
|
||||
slot_mapping=attn_metadata.slot_mapping[num_decode_tokens:N],
|
||||
block_table=attn_metadata.block_table[num_decodes:],
|
||||
query_start_loc=prefill_qsl,
|
||||
num_actual_tokens=N - num_decode_tokens,
|
||||
max_query_len=attn_metadata.max_query_len,
|
||||
max_seq_len=prefill_max_seq,
|
||||
is_prefill=True,
|
||||
query_start_loc_cpu=prefill_qsl_cpu,
|
||||
seq_lens_cpu=attn_metadata.seq_lens_cpu[num_decodes:]
|
||||
if attn_metadata.seq_lens_cpu is not None
|
||||
else None,
|
||||
)
|
||||
k = key[:N].view(N, self.num_kv_heads, self.head_size)
|
||||
v = value[:N].view(N, self.num_kv_heads, self.head_size)
|
||||
attn_out[num_decode_tokens:] = self._prefill_attention(
|
||||
q[num_decode_tokens:],
|
||||
k[num_decode_tokens:],
|
||||
v[num_decode_tokens:],
|
||||
kv_cache,
|
||||
prefill_meta,
|
||||
Pi,
|
||||
centroids,
|
||||
PiT,
|
||||
layer=layer,
|
||||
)
|
||||
|
||||
# Write into output buffer: attn_out is (N, Hq, D)
|
||||
# output may be 2D (N, Hq*D) or 3D (N, Hq, D)
|
||||
if output.ndim == 3:
|
||||
output[:N] = attn_out.to(output.dtype)
|
||||
else:
|
||||
output[:N] = attn_out.reshape(N, -1).to(output.dtype)
|
||||
return output
|
||||
|
||||
# ------------------------------------------------------------------ #
|
||||
# Store K/V into combined cache (vectorized) #
|
||||
# ------------------------------------------------------------------ #
|
||||
def _store_kv(
|
||||
self,
|
||||
key: torch.Tensor, # (N, Hk, D)
|
||||
value: torch.Tensor, # (N, Hk, D)
|
||||
kv_cache: torch.Tensor, # (num_blocks, block_size, Hk, slot_size)
|
||||
slot_mapping: torch.Tensor,
|
||||
layer: Any,
|
||||
):
|
||||
"""Quantize + store via fused Triton kernel."""
|
||||
triton_turboquant_store(
|
||||
key,
|
||||
value,
|
||||
kv_cache,
|
||||
slot_mapping,
|
||||
layer._tq_PiT,
|
||||
layer._tq_midpoints,
|
||||
mse_bits=self.tq_config.key_mse_bits,
|
||||
key_packed_size=self.tq_config.key_packed_size,
|
||||
value_quant_bits=self.tq_config.effective_value_quant_bits,
|
||||
key_fp8=self.tq_config.key_fp8,
|
||||
)
|
||||
|
||||
# ------------------------------------------------------------------ #
|
||||
# Prefill: SDPA on raw Q/K/V with causal mask #
|
||||
# ------------------------------------------------------------------ #
|
||||
def _prefill_attention(
|
||||
self,
|
||||
query: torch.Tensor, # (N, Hq, D)
|
||||
key: torch.Tensor, # (N, Hk, D)
|
||||
value: torch.Tensor, # (N, Hk, D)
|
||||
kv_cache: torch.Tensor, # (num_blocks, block_size, Hk, slot_size)
|
||||
attn_metadata: TurboQuantMetadata,
|
||||
Pi: torch.Tensor,
|
||||
centroids: torch.Tensor,
|
||||
PiT: torch.Tensor | None = None,
|
||||
layer: Any = None,
|
||||
) -> torch.Tensor:
|
||||
N, Hq, D = query.shape
|
||||
|
||||
# Fast path: use flash_attn for first-chunk prefills (all K/V in batch).
|
||||
# max_query_len == max_seq_len means no request has prior cached KV.
|
||||
# Both are Python ints — no GPU sync.
|
||||
if _HAS_FLASH_ATTN and attn_metadata.max_query_len == attn_metadata.max_seq_len:
|
||||
return self._flash_attn_varlen(
|
||||
q=query,
|
||||
k=key,
|
||||
v=value,
|
||||
cu_seqlens_q=attn_metadata.query_start_loc,
|
||||
cu_seqlens_k=attn_metadata.query_start_loc,
|
||||
max_seqlen_q=attn_metadata.max_query_len,
|
||||
max_seqlen_k=attn_metadata.max_query_len,
|
||||
)
|
||||
|
||||
# Continuation or no flash_attn: per-request attention.
|
||||
# For continuation chunks (seq_len > q_len), we must attend to
|
||||
# previously cached K/V from the TQ cache, not just the current
|
||||
# chunk's raw K/V.
|
||||
Hk = key.shape[1]
|
||||
use_gqa = Hk < Hq
|
||||
query_start_loc = attn_metadata.query_start_loc
|
||||
num_reqs = query_start_loc.shape[0] - 1
|
||||
|
||||
output = torch.zeros(N, Hq, D, device=query.device, dtype=query.dtype)
|
||||
|
||||
# Prefer the CPU-resident copies from the metadata if populated —
|
||||
# otherwise `.tolist()` on GPU tensors forces a synchronizing copy.
|
||||
if attn_metadata.query_start_loc_cpu is not None:
|
||||
qsl = attn_metadata.query_start_loc_cpu.tolist()
|
||||
else:
|
||||
qsl = query_start_loc.tolist()
|
||||
if attn_metadata.seq_lens_cpu is not None:
|
||||
seq_lens_list = attn_metadata.seq_lens_cpu.tolist()
|
||||
else:
|
||||
seq_lens_list = attn_metadata.seq_lens.tolist()
|
||||
|
||||
# Pre-allocate cu_seqlens for single-request flash_attn calls
|
||||
# to avoid per-request host→device tensor creation.
|
||||
if not hasattr(self, "_cu_2"):
|
||||
self._cu_2 = torch.zeros(2, device=query.device, dtype=torch.int32)
|
||||
# Cache arange on self (avoid per-call kernel launch).
|
||||
_max_seq = attn_metadata.max_seq_len
|
||||
_ac: torch.Tensor | None = getattr(self, "_arange_cache", None)
|
||||
if _ac is None or _ac.shape[0] <= _max_seq:
|
||||
_ac = torch.arange(
|
||||
0, _max_seq + 1, device=query.device, dtype=attn_metadata.seq_lens.dtype
|
||||
)
|
||||
self._arange_cache = _ac
|
||||
_arange_cache: torch.Tensor = _ac
|
||||
|
||||
for i in range(num_reqs):
|
||||
q_start = qsl[i]
|
||||
q_end = qsl[i + 1]
|
||||
q_len = q_end - q_start
|
||||
if q_len <= 0:
|
||||
continue
|
||||
|
||||
seq_len = seq_lens_list[i]
|
||||
q_seq = query[q_start:q_end] # (q_len, Hq, D)
|
||||
k_seq = key[q_start:q_end] # (q_len, Hk, D)
|
||||
v_seq = value[q_start:q_end] # (q_len, Hk, D)
|
||||
|
||||
if q_len == seq_len:
|
||||
# First-chunk prefill: all K/V are in the current batch.
|
||||
if _HAS_FLASH_ATTN:
|
||||
# Assign to slice to avoid gpu/cpu sync.
|
||||
self._cu_2[1:2] = q_len
|
||||
cu = self._cu_2
|
||||
out = self._flash_attn_varlen(
|
||||
q=q_seq,
|
||||
k=k_seq,
|
||||
v=v_seq,
|
||||
cu_seqlens_q=cu,
|
||||
cu_seqlens_k=cu,
|
||||
max_seqlen_q=q_len,
|
||||
max_seqlen_k=q_len,
|
||||
)
|
||||
else:
|
||||
q_t = q_seq.transpose(0, 1).contiguous()
|
||||
k_t = k_seq.transpose(0, 1).contiguous()
|
||||
v_t = v_seq.transpose(0, 1).contiguous()
|
||||
out = F.scaled_dot_product_attention(
|
||||
q_t,
|
||||
k_t,
|
||||
v_t,
|
||||
is_causal=True,
|
||||
scale=self.scale,
|
||||
enable_gqa=use_gqa,
|
||||
).transpose(0, 1)
|
||||
output[q_start:q_end] = out.to(query.dtype)
|
||||
else:
|
||||
# Continuation chunk: tokens already stored to TQ cache
|
||||
# by do_kv_cache_update. Use decode kernel directly to
|
||||
# avoid O(cached_len) full-dequant per continuation.
|
||||
# For large continuations, fall back to _continuation_prefill.
|
||||
cached_len = seq_len - q_len
|
||||
if q_len <= _CONTINUATION_DECODE_THRESHOLD:
|
||||
# Fast path: treat each query as a decode request
|
||||
# with incremental seq_lens for causal masking.
|
||||
# Slice from pre-built arange (no kernel launch)
|
||||
synth_seq_lens = _arange_cache[cached_len + 1 : seq_len + 1]
|
||||
synth_bt = attn_metadata.block_table[i : i + 1].expand(q_len, -1)
|
||||
out = triton_turboquant_decode_attention(
|
||||
query=q_seq,
|
||||
kv_cache=kv_cache,
|
||||
block_table=synth_bt,
|
||||
seq_lens=synth_seq_lens,
|
||||
Pi=Pi,
|
||||
centroids=centroids,
|
||||
scale=self.scale,
|
||||
mse_bits=self.tq_config.key_mse_bits,
|
||||
key_packed_size=self.tq_config.key_packed_size,
|
||||
value_quant_bits=(self.tq_config.effective_value_quant_bits),
|
||||
key_fp8=self.tq_config.key_fp8,
|
||||
norm_correction=self.tq_config.norm_correction,
|
||||
PiT=PiT,
|
||||
)
|
||||
else:
|
||||
# Large continuation: dequant cached K/V and use
|
||||
# flash_attn for better throughput.
|
||||
out = self._continuation_prefill(
|
||||
layer,
|
||||
q_seq,
|
||||
k_seq,
|
||||
v_seq,
|
||||
kv_cache,
|
||||
attn_metadata.block_table[i : i + 1],
|
||||
cached_len,
|
||||
seq_len,
|
||||
Pi,
|
||||
centroids,
|
||||
)
|
||||
output[q_start:q_end] = out.to(query.dtype)
|
||||
|
||||
return output
|
||||
|
||||
def _continuation_prefill(
|
||||
self,
|
||||
layer: Any,
|
||||
query: torch.Tensor, # (q_len, Hq, D)
|
||||
key_chunk: torch.Tensor, # (q_len, Hk, D)
|
||||
val_chunk: torch.Tensor, # (q_len, Hk, D)
|
||||
kv_cache: torch.Tensor, # (num_blocks, block_size, Hk, slot_size)
|
||||
block_table: torch.Tensor, # (1, max_num_blocks)
|
||||
cached_len: int,
|
||||
seq_len: int,
|
||||
Pi: torch.Tensor,
|
||||
centroids: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
"""Handle continuation chunk by dequanting cached K/V from TQ cache.
|
||||
|
||||
Dequants previously cached K/V, concatenates with the current
|
||||
chunk's raw K/V, then runs flash_attn with causal masking.
|
||||
"""
|
||||
q_len, Hq, D = query.shape
|
||||
Hk = key_chunk.shape[1]
|
||||
device = query.device
|
||||
block_size = kv_cache.shape[1]
|
||||
BLOCK_D = triton.next_power_of_2(D)
|
||||
|
||||
mse_bytes = self._mse_bytes
|
||||
val_data_bytes = self._val_data_bytes
|
||||
|
||||
# Dequant cached K/V from TQ cache
|
||||
# Allocate slightly over to align to block_size for the grid.
|
||||
# Reuse cached buffers to avoid per-call allocation (~16MB at 8K).
|
||||
alloc_len = math.ceil(cached_len / block_size) * block_size
|
||||
buf_shape = (1, Hk, alloc_len, D)
|
||||
# Use WorkspaceManager for dequant buffers.
|
||||
# Shared across all layers — saves 60× memory at long context.
|
||||
# Required for CUDA Graph capture (per-layer growth incompatible with CG).
|
||||
k_buf, v_buf = current_workspace_manager().get_simultaneous(
|
||||
(buf_shape, torch.float16),
|
||||
(buf_shape, torch.float16),
|
||||
)
|
||||
# Skip .zero_() — kernel writes all positions up to cached_len,
|
||||
# and we only read [:cached_len] afterwards.
|
||||
k_cached = k_buf[:, :, :alloc_len, :]
|
||||
v_cached = v_buf[:, :, :alloc_len, :]
|
||||
|
||||
grid = (alloc_len, 1 * Hk)
|
||||
_tq_full_dequant_kv[grid](
|
||||
kv_cache,
|
||||
block_table,
|
||||
centroids,
|
||||
k_cached,
|
||||
v_cached,
|
||||
k_cached.stride(0),
|
||||
k_cached.stride(1),
|
||||
k_cached.stride(2),
|
||||
v_cached.stride(0),
|
||||
v_cached.stride(1),
|
||||
v_cached.stride(2),
|
||||
kv_cache.stride(0),
|
||||
kv_cache.stride(1),
|
||||
kv_cache.stride(2),
|
||||
block_table.stride(0),
|
||||
HEAD_DIM=D,
|
||||
BLOCK_SIZE=block_size,
|
||||
NUM_KV_HEADS=Hk,
|
||||
MSE_BYTES=mse_bytes,
|
||||
KPS=self.tq_config.key_packed_size,
|
||||
VQB=self.tq_config.effective_value_quant_bits,
|
||||
VAL_DATA_BYTES=val_data_bytes,
|
||||
MSE_BITS=self.tq_config.key_mse_bits,
|
||||
KEY_FP8=1 if self.tq_config.key_fp8 else 0,
|
||||
BLOCK_D=BLOCK_D,
|
||||
NORM_CORRECTION=1 if self.tq_config.norm_correction else 0,
|
||||
FP8_E4B15=_use_fp8_e4b15(device.index or 0),
|
||||
num_warps=4,
|
||||
)
|
||||
|
||||
# Inverse-rotate MSE keys back to original space
|
||||
if not self.tq_config.key_fp8:
|
||||
# fp16 matmul for rotation (2× less bandwidth, uses fp16 tensor cores)
|
||||
Pi_half = layer._tq_Pi_half
|
||||
k_flat = k_cached[0, :, :cached_len, :].reshape(-1, D)
|
||||
k_flat = k_flat @ Pi_half
|
||||
k_cached_trim = k_flat.reshape(Hk, cached_len, D).transpose(
|
||||
0, 1
|
||||
) # (cached_len, Hk, D) — already fp16
|
||||
else:
|
||||
k_cached_trim = k_cached[0, :, :cached_len, :].transpose(
|
||||
0, 1
|
||||
) # (cached_len, Hk, D)
|
||||
|
||||
# Skip .contiguous() — the copy into k_full/v_full handles layout
|
||||
v_cached_trim = v_cached[0, :, :cached_len, :].transpose(0, 1)
|
||||
|
||||
# Concatenate cached + current chunk K/V (match query dtype)
|
||||
# Pre-allocate full K/V buffer, copy into slices (no cat alloc)
|
||||
qdtype = query.dtype
|
||||
k_full = torch.empty(seq_len, Hk, D, dtype=qdtype, device=device)
|
||||
v_full = torch.empty(seq_len, Hk, D, dtype=qdtype, device=device)
|
||||
k_full[:cached_len] = k_cached_trim.to(qdtype)
|
||||
k_full[cached_len:] = key_chunk
|
||||
v_full[:cached_len] = v_cached_trim.to(qdtype)
|
||||
v_full[cached_len:] = val_chunk
|
||||
|
||||
# Attention: q_len queries attending to seq_len K/V with causal mask
|
||||
if _HAS_FLASH_ATTN:
|
||||
# Reuse pre-allocated cu_seqlens (avoid host→device transfer)
|
||||
if not hasattr(self, "_cu_2_q"):
|
||||
self._cu_2_q = torch.zeros(2, device=device, dtype=torch.int32)
|
||||
self._cu_2_k = torch.zeros(2, device=device, dtype=torch.int32)
|
||||
# Assigning to slice uses fill_ which avoids cpu/gpu sync.
|
||||
self._cu_2_q[1:2] = q_len
|
||||
self._cu_2_k[1:2] = seq_len
|
||||
cu_seqlens_q = self._cu_2_q
|
||||
cu_seqlens_k = self._cu_2_k
|
||||
return self._flash_attn_varlen(
|
||||
q=query,
|
||||
k=k_full,
|
||||
v=v_full,
|
||||
cu_seqlens_q=cu_seqlens_q,
|
||||
cu_seqlens_k=cu_seqlens_k,
|
||||
max_seqlen_q=q_len,
|
||||
max_seqlen_k=seq_len,
|
||||
)
|
||||
else:
|
||||
# SDPA fallback: expand KV for GQA, build causal mask
|
||||
q_t = query.transpose(0, 1).unsqueeze(0) # (1, Hq, q_len, D)
|
||||
k_t = k_full.transpose(0, 1).unsqueeze(0) # (1, Hk, seq_len, D)
|
||||
v_t = v_full.transpose(0, 1).unsqueeze(0) # (1, Hk, seq_len, D)
|
||||
# Build causal mask: query position p can attend to K position j
|
||||
# where j <= cached_len + p (p is 0-indexed within chunk)
|
||||
q_pos = torch.arange(q_len, device=device).unsqueeze(1) + cached_len
|
||||
k_pos = torch.arange(seq_len, device=device).unsqueeze(0)
|
||||
mask = k_pos <= q_pos # (q_len, seq_len)
|
||||
out = F.scaled_dot_product_attention(
|
||||
q_t,
|
||||
k_t,
|
||||
v_t,
|
||||
attn_mask=mask,
|
||||
scale=self.scale,
|
||||
enable_gqa=(Hk < Hq),
|
||||
) # (1, Hq, q_len, D)
|
||||
return out[0].transpose(0, 1) # (q_len, Hq, D)
|
||||
|
||||
# ------------------------------------------------------------------ #
|
||||
# Decode: Triton TQ decode attention #
|
||||
# ------------------------------------------------------------------ #
|
||||
def _decode_attention(
|
||||
self,
|
||||
query: torch.Tensor, # (B, Hq, D)
|
||||
kv_cache: torch.Tensor, # (num_blocks, block_size, Hk, slot_size)
|
||||
attn_metadata: TurboQuantMetadata,
|
||||
Pi: torch.Tensor,
|
||||
centroids: torch.Tensor,
|
||||
PiT: torch.Tensor | None = None,
|
||||
layer: torch.nn.Module | None = None,
|
||||
) -> torch.Tensor:
|
||||
# Acquire shared decode scratch buffers from WorkspaceManager.
|
||||
# Layers execute sequentially so one set of buffers is sufficient.
|
||||
# Falls back to kernel-internal allocation if workspace unavailable.
|
||||
B = query.shape[0]
|
||||
D = self.head_size
|
||||
S = self.max_num_kv_splits
|
||||
Hq = self.num_heads
|
||||
mid_o_buf = output_buf = lse_buf = None
|
||||
if is_workspace_manager_initialized():
|
||||
# output_buf in query dtype — matches the in-kernel fp16 cast in stage2.
|
||||
mid_o_buf, output_buf, lse_buf = (
|
||||
current_workspace_manager().get_simultaneous(
|
||||
((B, Hq, S, D + 1), torch.float32),
|
||||
((B, Hq, D), query.dtype),
|
||||
((B, Hq), torch.float32),
|
||||
)
|
||||
)
|
||||
|
||||
result = triton_turboquant_decode_attention(
|
||||
query=query,
|
||||
kv_cache=kv_cache,
|
||||
block_table=attn_metadata.block_table,
|
||||
seq_lens=attn_metadata.seq_lens,
|
||||
Pi=Pi,
|
||||
centroids=centroids,
|
||||
scale=self.scale,
|
||||
mse_bits=self.tq_config.key_mse_bits,
|
||||
key_packed_size=self.tq_config.key_packed_size,
|
||||
value_quant_bits=self.tq_config.effective_value_quant_bits,
|
||||
key_fp8=self.tq_config.key_fp8,
|
||||
norm_correction=self.tq_config.norm_correction,
|
||||
PiT=PiT,
|
||||
mid_o_buf=mid_o_buf,
|
||||
output_buf=output_buf,
|
||||
lse_buf=lse_buf,
|
||||
buf_holder=layer,
|
||||
max_num_kv_splits=self.max_num_kv_splits,
|
||||
)
|
||||
return result
|
||||
Reference in New Issue
Block a user