From 8b79971bb9f97ca0dd59203c628f52de4664de10 Mon Sep 17 00:00:00 2001 From: Artur Fierka Date: Mon, 6 Jul 2026 15:01:21 +0200 Subject: [PATCH] attention: pass None for unused args in unified attention TD path (#43597) Signed-off-by: Artur Fierka Co-authored-by: quinnlp --- .../attention/ops/triton_unified_attention.py | 49 +++++++++++-------- 1 file changed, 29 insertions(+), 20 deletions(-) diff --git a/vllm/v1/attention/ops/triton_unified_attention.py b/vllm/v1/attention/ops/triton_unified_attention.py index a13ae96a7a9..93622957b55 100644 --- a/vllm/v1/attention/ops/triton_unified_attention.py +++ b/vllm/v1/attention/ops/triton_unified_attention.py @@ -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, ...]