attention: pass None for unused args in unified attention TD path (#43597)
Signed-off-by: Artur Fierka <artur.fierka@intel.com> Co-authored-by: quinnlp <quinnlp@users.noreply.github.com>
Artur Fierka committed
Jul 6, 2026 at 15:01 UTC
8b79971bb9f97ca0dd59203c628f52de4664de10
1 file changed
+29
-20
vllm/v1/attention/ops/triton_unified_attention.py
+29
-20
@@ -254,12 +254,19 @@ def kernel_unified_attention(
254
# Per-(token, head) scale caches: used iff KV_QUANT_MODE in {2, 3}.
255
k_scale_cache_ptr=None,
256
v_scale_cache_ptr=None,
257
- stride_ks_blk: tl.int64 = None,
258
- stride_ks_slot: tl.int64 = None,
259
- stride_ks_head: tl.int64 = None,
260
- stride_vs_blk: tl.int64 = None,
261
- stride_vs_slot: tl.int64 = None,
262
- stride_vs_head: tl.int64 = None,
257
+ # ``tl.int64`` cannot be combined with a ``None`` default — Triton's JIT
258
+ # rejects ``Optional[tl.int64]`` / ``tl.int64 | None`` at trace time, and
259
+ # plain ``tl.int64 = None`` raises ``TypeError: 'NoneType' object cannot
260
+ # be interpreted as an integer`` when callers omit these arguments.
261
+ # ``int | None`` is the only annotation that lets the wrapper pass
262
+ # ``None`` here so Triton can skip materialising the strides when the
263
+ # ``USE_PER_TOKEN_HEAD_SCALES`` branch is dead.
264
+ stride_ks_blk: int | None = None,
265
+ stride_ks_slot: int | None = None,
266
+ stride_ks_head: int | None = None,
267
+ stride_vs_blk: int | None = None,
268
+ stride_vs_slot: int | None = None,
269
+ stride_vs_head: int | None = None,
270
# KV cache quantization mode handled inside this kernel via constexpr
271
# branches: NONE (0), FP8_PER_TENSOR (1), INT8_PER_TOKEN_HEAD (2),
272
# FP8_PER_TOKEN_HEAD (3). Sub-byte INT4 (4) uses its own
@@ -283,7 +290,10 @@ def kernel_unified_attention(
290
# original (causal AND SW) OR mm_prefix behavior for all other models.
291
MM_PREFIX_CLAMP_SW: tl.constexpr = False,
292
):
286
- USE_PER_TOKEN_HEAD_SCALES: tl.constexpr = KV_QUANT_MODE >= 2
293
+ # Per-(token, head) scale caches: used iff KV_QUANT_MODE in {2, 3}.
294
+ USE_PER_TOKEN_HEAD_SCALES: tl.constexpr = (KV_QUANT_MODE >= 2) and (
295
+ KV_QUANT_MODE <= 3
296
+ )
297
USE_FP8_Q_DESCALE: tl.constexpr = KV_QUANT_MODE == 1 and Q_IS_FP8
298
299
if USE_TD:
@@ -1041,9 +1051,9 @@ def unified_attention(
1051
1052
# The kernel signature is the same for 2D and 3D — only the launch
1053
# grid + a handful of constexpr toggles differ. Per-token-head scale
1044
- # caches and their strides are required arguments; non-per-token-head
1045
- # modes pass dummy zeros (the code path is dead-code eliminated by
1046
- # the ``USE_PER_TOKEN_HEAD_SCALES`` constexpr branch in the kernel).
1054
+ # caches and their strides are passed as ``None`` when the
1055
+ # ``USE_PER_TOKEN_HEAD_SCALES`` branch is dead so Triton can skip
1056
+ # materialising those arguments and the associated registers.
1057
if use_per_token_head_scales:
1058
ks_strides = k_scale_cache.stride()
1059
vs_strides = v_scale_cache.stride()
@@ -1052,16 +1062,15 @@ def unified_attention(
1062
k_scale_ptr = k_scale_cache
1063
v_scale_ptr = v_scale_cache
1064
else:
1055
- ks_blk = ks_slot = ks_head = 0
1056
- vs_blk = vs_slot = vs_head = 0
1057
- # Pass the K cache as a stand-in pointer; never dereferenced.
1058
- k_scale_ptr = k
1059
- v_scale_ptr = v
1060
- # 3D needs real segm tensors; 2D never touches them but Triton wants
1061
- # a non-null pointer. Reuse ``out`` as the placeholder.
1062
- segm_output_ptr = softmax_segm_output if use_3d else out
1063
- segm_max_ptr = softmax_segm_max if use_3d else out
1064
- segm_expsum_ptr = softmax_segm_expsum if use_3d else out
1065
+ ks_blk = ks_slot = ks_head = None
1066
+ vs_blk = vs_slot = vs_head = None
1067
+ k_scale_ptr = None
1068
+ v_scale_ptr = None
1069
+ # 3D needs real segm tensors; 2D never touches them. Pass ``None`` in
1070
+ # 2D mode so Triton can skip materialising these pointer arguments.
1071
+ segm_output_ptr = softmax_segm_output if use_3d else None
1072
+ segm_max_ptr = softmax_segm_max if use_3d else None
1073
+ segm_expsum_ptr = softmax_segm_expsum if use_3d else None
1074
num_segments = num_par_softmax_segments if use_3d else 1
1075
1076
grid: tuple[Any, ...]