karylab_agents/vllm
1
0
Fork 0
forked from Karylab-cklius/vllm
Code Pull Requests 2 Actions 1 Packages Activity

attention: pass None for unused args in unified attention TD path (#43597)

Browse Source 
Signed-off-by: Artur Fierka <artur.fierka@intel.com>
Co-authored-by: quinnlp <quinnlp@users.noreply.github.com>
This commit is contained in:
Artur Fierka
2026-07-06 21:01:21 +08:00
committed by GitHub
co-authored by quinnlp
parent f676808ba0
commit 8b79971bb9
1 changed files with 29 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files
vllm/v1/attention/ops/triton_unified_attention.py
+29 -20
View File
@@ -254,12 +254,19 @@ def kernel_unified_attention(
# Per-(token, head) scale caches: used iff KV_QUANT_MODE in {2, 3}.
k_scale_cache_ptr=None,
v_scale_cache_ptr=None,
stride_ks_blk: tl.int64 = None,
stride_ks_slot: tl.int64 = None,
stride_ks_head: tl.int64 = None,
stride_vs_blk: tl.int64 = None,
stride_vs_slot: tl.int64 = None,
stride_vs_head: tl.int64 = None,
# ``tl.int64`` cannot be combined with a ``None`` default — Triton's JIT
# rejects ``Optional[tl.int64]`` / ``tl.int64 | None`` at trace time, and
# plain ``tl.int64 = None`` raises ``TypeError: 'NoneType' object cannot
# be interpreted as an integer`` when callers omit these arguments.
# ``int | None`` is the only annotation that lets the wrapper pass
# ``None`` here so Triton can skip materialising the strides when the
# ``USE_PER_TOKEN_HEAD_SCALES`` branch is dead.
stride_ks_blk: int | None = None,
stride_ks_slot: int | None = None,
stride_ks_head: int | None = None,
stride_vs_blk: int | None = None,
stride_vs_slot: int | None = None,
stride_vs_head: int | None = None,
# KV cache quantization mode handled inside this kernel via constexpr
# branches: NONE (0), FP8_PER_TENSOR (1), INT8_PER_TOKEN_HEAD (2),
# FP8_PER_TOKEN_HEAD (3). Sub-byte INT4 (4) uses its own
@@ -283,7 +290,10 @@ def kernel_unified_attention(
# original (causal AND SW) OR mm_prefix behavior for all other models.
MM_PREFIX_CLAMP_SW: tl.constexpr = False,
):
USE_PER_TOKEN_HEAD_SCALES: tl.constexpr = KV_QUANT_MODE >= 2
# Per-(token, head) scale caches: used iff KV_QUANT_MODE in {2, 3}.
USE_PER_TOKEN_HEAD_SCALES: tl.constexpr = (KV_QUANT_MODE >= 2) and (
KV_QUANT_MODE <= 3
)
USE_FP8_Q_DESCALE: tl.constexpr = KV_QUANT_MODE == 1 and Q_IS_FP8
if USE_TD:
@@ -1041,9 +1051,9 @@ def unified_attention(
# The kernel signature is the same for 2D and 3D — only the launch
# grid + a handful of constexpr toggles differ. Per-token-head scale
# caches and their strides are required arguments; non-per-token-head
# modes pass dummy zeros (the code path is dead-code eliminated by
# the ``USE_PER_TOKEN_HEAD_SCALES`` constexpr branch in the kernel).
# caches and their strides are passed as ``None`` when the
# ``USE_PER_TOKEN_HEAD_SCALES`` branch is dead so Triton can skip
# materialising those arguments and the associated registers.
if use_per_token_head_scales:
ks_strides = k_scale_cache.stride()
vs_strides = v_scale_cache.stride()
@@ -1052,16 +1062,15 @@ def unified_attention(
k_scale_ptr = k_scale_cache
v_scale_ptr = v_scale_cache
else:
ks_blk = ks_slot = ks_head = 0
vs_blk = vs_slot = vs_head = 0
# Pass the K cache as a stand-in pointer; never dereferenced.
k_scale_ptr = k
v_scale_ptr = v
# 3D needs real segm tensors; 2D never touches them but Triton wants
# a non-null pointer. Reuse ``out`` as the placeholder.
segm_output_ptr = softmax_segm_output if use_3d else out
segm_max_ptr = softmax_segm_max if use_3d else out
segm_expsum_ptr = softmax_segm_expsum if use_3d else out
ks_blk = ks_slot = ks_head = None
vs_blk = vs_slot = vs_head = None
k_scale_ptr = None
v_scale_ptr = None
# 3D needs real segm tensors; 2D never touches them. Pass ``None`` in
# 2D mode so Triton can skip materialising these pointer arguments.
segm_output_ptr = softmax_segm_output if use_3d else None
segm_max_ptr = softmax_segm_max if use_3d else None
segm_expsum_ptr = softmax_segm_expsum if use_3d else None
num_segments = num_par_softmax_segments if use_3d else 1
grid: tuple[Any, ...]