A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 3, 19:00 UTC
Triton attention: add USE_TD constexpr for tensor descriptor Q/K/V load/store (#40327)
Signed-off-by: Artur Fierka <artur.fierka@intel.com> Co-authored-by: quinnlp <quinnlp@users.noreply.github.com>
Artur Fierka committed
May 13, 2026 at 13:57 UTC
3c413a548177c6b1a3d684ece1bd24725fac89d8
4 files changed
+575
-70
tests/kernels/attention/test_triton_unified_attention.py
+179
@@ -345,3 +345,182 @@ def test_triton_unified_attn_fp16_input_fp8_output(
345
torch.testing.assert_close(output_fp16, ref_output, atol=atol, rtol=rtol),
346
f"{torch.max(torch.abs(output_fp16 - ref_output))}",
347
)
348
+
349
+
350
+# USE_TD path covers two head-size regimes:
351
+# - pow2 (HEAD_SIZE == HEAD_SIZE_PADDED): full TD path including Q/O.
352
+# - non-pow2 (96, HEAD_SIZE_PADDED=128): gates USE_TD_QO off — Q load
353
+# and output store fall back to pointer path, KV tile TD load remains.
354
+# The non-pow2 case mirrors real models like Phi-3-mini (head_size=96).
355
+HEAD_SIZES_USE_TD = [128, 256, 96]
356
+
357
+
358
+def _run_use_td_case(
359
+ seq_lens: list[tuple[int, int]],
360
+ num_heads: tuple[int, int],
361
+ head_size: int,
362
+ block_size: int,
363
+ sliding_window: int | None,
364
+ soft_cap: float | None,
365
+ seq_threshold_3D: int,
366
+ dtype: torch.dtype = torch.bfloat16,
367
+ num_blocks: int = 2048,
368
+) -> None:
369
+ """Shared driver for the USE_TD test cases.
370
+
371
+ Runs ``unified_attention(..., use_td=True)`` and compares against the
372
+ reference paged-attention implementation that the sibling non-TD
373
+ tests use.
374
+ """
375
+ torch.set_default_device(DEVICE_TYPE)
376
+ set_random_seed(0)
377
+
378
+ num_seqs = len(seq_lens)
379
+ query_lens = [x[0] for x in seq_lens]
380
+ kv_lens = [x[1] for x in seq_lens]
381
+ num_query_heads, num_kv_heads = num_heads
382
+ assert num_query_heads % num_kv_heads == 0
383
+ max_query_len = max(query_lens)
384
+ max_kv_len = max(kv_lens)
385
+ window_size = (sliding_window - 1, 0) if sliding_window is not None else (-1, -1)
386
+ scale = head_size**-0.5
387
+
388
+ query = torch.randn(sum(query_lens), num_query_heads, head_size, dtype=dtype)
389
+ key_cache = torch.randn(
390
+ num_blocks, block_size, num_kv_heads, head_size, dtype=dtype
391
+ )
392
+ value_cache = torch.randn_like(key_cache)
393
+ cu_query_lens = torch.tensor([0] + query_lens, dtype=torch.int32).cumsum(
394
+ dim=0, dtype=torch.int32
395
+ )
396
+ kv_lens_tensor = torch.tensor(kv_lens, dtype=torch.int32)
397
+
398
+ max_num_blocks_per_seq = (max_kv_len + block_size - 1) // block_size
399
+ block_tables = torch.randint(
400
+ 0, num_blocks, (num_seqs, max_num_blocks_per_seq), dtype=torch.int32
401
+ )
402
+
403
+ output = torch.empty_like(query)
404
+
405
+ num_par_softmax_segments = 16
406
+ head_size_padded = next_power_of_2(head_size)
407
+ softmax_segm_output = torch.empty(
408
+ (seq_threshold_3D, num_query_heads, num_par_softmax_segments, head_size_padded),
409
+ dtype=torch.float32,
410
+ )
411
+ softmax_segm_max = torch.empty(
412
+ (seq_threshold_3D, num_query_heads, num_par_softmax_segments),
413
+ dtype=torch.float32,
414
+ )
415
+ softmax_segm_expsum = torch.empty(
416
+ (seq_threshold_3D, num_query_heads, num_par_softmax_segments),
417
+ dtype=torch.float32,
418
+ )
419
+
420
+ unified_attention(
421
+ q=query,
422
+ k=key_cache,
423
+ v=value_cache,
424
+ out=output,
425
+ cu_seqlens_q=cu_query_lens,
426
+ seqused_k=kv_lens_tensor,
427
+ max_seqlen_q=max_query_len,
428
+ max_seqlen_k=max_kv_len,
429
+ softmax_scale=scale,
430
+ causal=True,
431
+ window_size=window_size,
432
+ block_table=block_tables,
433
+ softcap=soft_cap if soft_cap is not None else 0,
434
+ q_descale=None,
435
+ k_descale=None,
436
+ v_descale=None,
437
+ seq_threshold_3D=seq_threshold_3D,
438
+ num_par_softmax_segments=num_par_softmax_segments,
439
+ softmax_segm_output=softmax_segm_output,
440
+ softmax_segm_max=softmax_segm_max,
441
+ softmax_segm_expsum=softmax_segm_expsum,
442
+ use_td=True,
443
+ )
444
+
445
+ ref_output = ref_paged_attn(
446
+ query=query,
447
+ key_cache=key_cache,
448
+ value_cache=value_cache,
449
+ query_lens=query_lens,
450
+ kv_lens=kv_lens,
451
+ block_tables=block_tables,
452
+ scale=scale,
453
+ sliding_window=sliding_window,
454
+ soft_cap=soft_cap,
455
+ )
456
+ torch.testing.assert_close(output, ref_output, atol=1.5e-2, rtol=1e-2)
457
+
458
+
459
+@pytest.mark.parametrize(
460
+ "seq_lens", [[(1, 1328), (5, 18), (129, 463)], [(1, 523), (1, 37), (1, 2011)]]
461
+)
462
+@pytest.mark.parametrize("num_heads", NUM_HEADS)
463
+@pytest.mark.parametrize("head_size", HEAD_SIZES_USE_TD)
464
+@pytest.mark.parametrize("block_size", BLOCK_SIZES)
465
+@pytest.mark.parametrize("sliding_window", [None, 128])
466
+@pytest.mark.parametrize("soft_cap", [None, 50.0])
467
+@pytest.mark.parametrize("num_blocks", NUM_BLOCKS)
468
+@pytest.mark.parametrize("seq_threshold_3D", SEQ_THRESHOLD_3D_VALUES)
469
+@torch.inference_mode()
470
+def test_triton_unified_attn_use_td(
471
+ seq_lens: list[tuple[int, int]],
472
+ num_heads: tuple[int, int],
473
+ head_size: int,
474
+ sliding_window: int | None,
475
+ block_size: int,
476
+ soft_cap: float | None,
477
+ num_blocks: int,
478
+ seq_threshold_3D: int,
479
+) -> None:
480
+ """Exercise the USE_TD (tensor-descriptor) Q/K/V load/store path.
481
+
482
+ Covers both 2D and 3D kernels via ``seq_threshold_3D``. Two routes
483
+ to the USE_TD_QO=False fallback (pointer path for Q/O with TD still
484
+ active for KV tile loads):
485
+
486
+ - non-pow2 ``num_queries_per_kv`` via ``NUM_HEADS`` entry ``(5, 1)``,
487
+ - non-pow2 ``head_size`` via ``HEAD_SIZES_USE_TD`` entry ``96``.
488
+ """
489
+ _run_use_td_case(
490
+ seq_lens=seq_lens,
491
+ num_heads=num_heads,
492
+ head_size=head_size,
493
+ block_size=block_size,
494
+ sliding_window=sliding_window,
495
+ soft_cap=soft_cap,
496
+ seq_threshold_3D=seq_threshold_3D,
497
+ num_blocks=num_blocks,
498
+ )
499
+
500
+
501
+# Prefill-heavy shape: long query drives the prefill kernel path where
502
+# ``_get_tile_size`` returns 32, which exceeds block_size=16 and must be
503
+# clamped by the fix in 'clamp TILE_SIZE to block_size when USE_TD'.
504
+# Only the prefill launch exercises the clamp, so parameterize only over
505
+# the (num_heads, seq_threshold_3D=0) combinations needed to cover it.
506
+@pytest.mark.parametrize("num_heads", [(4, 4), (5, 1)])
507
+@torch.inference_mode()
508
+def test_triton_unified_attn_use_td_tile_clamp(
509
+ num_heads: tuple[int, int],
510
+) -> None:
511
+ """Regression guard: ``USE_TD`` needs ``BLOCK_SIZE % TILE_SIZE == 0``.
512
+
513
+ With ``block_size=16`` and ``head_size=128`` (non-Gemma3),
514
+ ``_get_tile_size`` returns 32 for prefill, which violates the
515
+ ``USE_TD`` constraint unless clamped to ``block_size``. Without
516
+ the clamp the triton kernel ``static_assert`` fires at compile time.
517
+ """
518
+ _run_use_td_case(
519
+ seq_lens=[(256, 256), (128, 128)],
520
+ num_heads=num_heads,
521
+ head_size=128,
522
+ block_size=16,
523
+ sliding_window=None,
524
+ soft_cap=None,
525
+ seq_threshold_3D=0,
526
+ )
vllm/envs.py
+11
@@ -82,6 +82,7 @@ if TYPE_CHECKING:
82
VLLM_MAIN_CUDA_VERSION: str = "13.0"
83
VLLM_FLOAT32_MATMUL_PRECISION: Literal["highest", "high", "medium"] = "highest"
84
VLLM_BATCH_INVARIANT: bool = False
85
+ VLLM_TRITON_ATTN_USE_TD: bool | None = None
86
MAX_JOBS: str | None = None
87
NVCC_THREADS: str | None = None
88
VLLM_USE_PRECOMPILED: bool = False
@@ -525,6 +526,16 @@ environment_variables: dict[str, Callable[[], Any]] = {
526
# Enable batch-invariant mode: deterministic results regardless of
527
# batch composition. Requires NVIDIA GPU with compute capability >= 9.0.
528
"VLLM_BATCH_INVARIANT": lambda: bool(int(os.getenv("VLLM_BATCH_INVARIANT", "0"))),
529
+ # Use tensor descriptors for Q/K/V loads and output stores in the
530
+ # Triton unified-attention kernel. Enables HW 2D block reads on
531
+ # Intel Xe2/Xe3; the non-TD branch is dead-code-eliminated at Triton
532
+ # compile time so other platforms see no overhead. Tri-state override:
533
+ # unset (default) lets the `triton_attn` backend auto-select per
534
+ # platform (currently auto-enabled on XPU only); ``1`` forces TD on;
535
+ # ``0`` forces TD off. Useful for A/B benchmarking the TD path.
536
+ "VLLM_TRITON_ATTN_USE_TD": lambda: {"1": True, "0": False}.get(
537
+ os.getenv("VLLM_TRITON_ATTN_USE_TD", "").strip()
538
+ ),
539
# Maximum number of compilation jobs to run in parallel.
540
# By default this is the number of CPUs
541
"MAX_JOBS": lambda: os.getenv("MAX_JOBS", None),
vllm/v1/attention/backends/triton_attn.py
+17
@@ -7,6 +7,7 @@ from typing import ClassVar
7
8
import torch
9
10
+import vllm.envs as envs
11
from vllm._aiter_ops import rocm_aiter_ops
12
from vllm.config import CUDAGraphMode, VllmConfig
13
from vllm.config.cache import CacheDType
@@ -502,6 +503,21 @@ class TritonAttentionImpl(AttentionImpl):
503
self._kv_quant_mode = get_kv_quant_mode(kv_cache_dtype)
504
self._is_per_token_head_quant = self._kv_quant_mode.is_per_token_head
505
506
+ # Enable tensor descriptors for Q/K/V load/store on platforms that
507
+ # benefit from HW 2D block reads (Intel Xe2/Xe3). The dead branch
508
+ # is eliminated at Triton compile time, so other platforms see
509
+ # zero cost when TD is off.
510
+ #
511
+ # ``VLLM_TRITON_ATTN_USE_TD`` is tri-state:
512
+ # - unset (None): auto-select (TD on for XPU, off elsewhere),
513
+ # - ``1``: force TD on regardless of platform,
514
+ # - ``0``: force TD off regardless of platform (useful for A/B).
515
+ td_override = envs.VLLM_TRITON_ATTN_USE_TD
516
+ if td_override is None:
517
+ self.use_td = current_platform.is_xpu()
518
+ else:
519
+ self.use_td = td_override
520
+
521
def forward(
522
self,
523
layer: torch.nn.Module,
@@ -637,6 +653,7 @@ class TritonAttentionImpl(AttentionImpl):
653
k_scale_cache=k_scale_cache,
654
v_scale_cache=v_scale_cache,
655
chunk_lookback=self.chunk_lookback,
656
+ use_td=self.use_td,
657
)
658
659
return output
vllm/v1/attention/ops/triton_unified_attention.py
+368
-70
@@ -54,16 +54,132 @@ def _cast_kv_tile(data, Q, tensor_scale, KV_QUANT_MODE: tl.constexpr):
54
return data.to(Q.dtype)
55
56
57
+# ---------------------------------------------------------------------------
58
+# Tensor-descriptor (TD) helpers
59
+#
60
+# Used when the caller enables ``USE_TD`` (Intel Xe2/Xe3 HW 2D block reads).
61
+# When ``USE_TD`` is False the helpers are dead-code-eliminated at Triton
62
+# compile time, leaving the pointer-arithmetic load/store path untouched.
63
+# ---------------------------------------------------------------------------
64
+
65
+
66
+@triton.jit
67
+def _load_q_td(
68
+ query_ptr,
69
+ q_block_local_len,
70
+ query_stride_0: tl.int64,
71
+ query_stride_1: tl.int64,
72
+ cur_batch_in_all_start_index,
73
+ q_block_local_idx,
74
+ kv_head_idx,
75
+ num_queries_per_kv: tl.constexpr,
76
+ BLOCK_Q: tl.constexpr,
77
+ BLOCK_M: tl.constexpr,
78
+ HEAD_SIZE: tl.constexpr,
79
+ HEAD_SIZE_PADDED: tl.constexpr,
80
+):
81
+ """Load Q via a 2D tensor descriptor.
82
+
83
+ Caller guarantees (via the wrapper's ``use_td_qo`` gate):
84
+ * ``HEAD_SIZE == HEAD_SIZE_PADDED`` (head_size is a power of 2),
85
+ * ``num_queries_per_kv`` is a power of 2,
86
+ * the ``num_queries_per_kv`` heads of the current KV group are
87
+ contiguous in memory (``query_stride_1 == HEAD_SIZE``, which is
88
+ the default vLLM query layout).
89
+
90
+ Under those preconditions the inner two axes are flattened into one
91
+ row of size ``num_queries_per_kv * HEAD_SIZE`` with stride 1, which
92
+ avoids the non-power-of-2 ``block_shape`` error from the Triton
93
+ tensor-descriptor validator. Returns (BLOCK_M, HEAD_SIZE_PADDED).
94
+ """
95
+ q_base = (
96
+ query_ptr
97
+ + (cur_batch_in_all_start_index + q_block_local_idx * BLOCK_Q) * query_stride_0
98
+ + (kv_head_idx * num_queries_per_kv) * query_stride_1
99
+ )
100
+ q_desc = tl.make_tensor_descriptor(
101
+ base=q_base,
102
+ shape=(q_block_local_len, num_queries_per_kv * HEAD_SIZE),
103
+ strides=(query_stride_0, 1),
104
+ block_shape=(BLOCK_Q, num_queries_per_kv * HEAD_SIZE_PADDED),
105
+ )
106
+ return q_desc.load([0, 0]).reshape(BLOCK_M, HEAD_SIZE_PADDED)
107
+
108
+
109
+@triton.jit
110
+def _load_kv_tile_td(
111
+ cache_ptr,
112
+ physical_block_idx_scalar,
113
+ kv_head_idx,
114
+ offset_in_block,
115
+ stride_cache_0: tl.int64,
116
+ stride_cache_1: tl.int64,
117
+ stride_cache_2: tl.int64,
118
+ stride_cache_3: tl.constexpr,
119
+ BLOCK_SIZE: tl.constexpr,
120
+ TILE_SIZE: tl.constexpr,
121
+ HEAD_SIZE: tl.constexpr,
122
+ HEAD_SIZE_PADDED: tl.constexpr,
123
+):
124
+ """Load a KV cache tile via tensor descriptor.
125
+
126
+ Returns shape (TILE_SIZE, HEAD_SIZE_PADDED). Caller transposes for K.
127
+ Tensor descriptors zero-pad reads beyond the shape boundary, so
128
+ ``HEAD_SIZE_PADDED > HEAD_SIZE`` is handled correctly.
129
+ """
130
+ base = (
131
+ cache_ptr
132
+ + physical_block_idx_scalar * stride_cache_0
133
+ + kv_head_idx * stride_cache_2
134
+ )
135
+ desc = tl.make_tensor_descriptor(
136
+ base=base,
137
+ shape=(BLOCK_SIZE, HEAD_SIZE),
138
+ strides=(stride_cache_1, stride_cache_3),
139
+ block_shape=(TILE_SIZE, HEAD_SIZE_PADDED),
140
+ )
141
+ return desc.load([offset_in_block, 0])
142
+
143
+
144
+@triton.jit
145
+def _store_output_td(
146
+ base_ptr,
147
+ acc,
148
+ q_block_local_len,
149
+ stride_token: tl.int64,
150
+ stride_head: tl.int64,
151
+ num_queries_per_kv: tl.constexpr,
152
+ BLOCK_Q: tl.constexpr,
153
+ HEAD_SIZE: tl.constexpr,
154
+ HEAD_SIZE_PADDED: tl.constexpr,
155
+):
156
+ """Store an output tile via a tensor descriptor.
157
+
158
+ The 2D and 3D epilogues differ only in ``base_ptr`` and the
159
+ ``(stride_token, stride_head)`` pair: 2D writes directly to the
160
+ flat output buffer, 3D writes to a single per-segment slice of
161
+ ``segm_output_ptr``. Descriptor shape / block_shape / reshape
162
+ are the same in both modes, so share one helper.
163
+ """
164
+ acc = acc.to(base_ptr.dtype.element_ty)
165
+ output_desc = tl.make_tensor_descriptor(
166
+ base=base_ptr,
167
+ shape=(q_block_local_len, num_queries_per_kv, HEAD_SIZE),
168
+ strides=(stride_token, stride_head, 1),
169
+ block_shape=(BLOCK_Q, num_queries_per_kv, HEAD_SIZE_PADDED),
170
+ )
171
+ output_desc.store(
172
+ [0, 0, 0],
173
+ acc.reshape(BLOCK_Q, num_queries_per_kv, HEAD_SIZE_PADDED),
174
+ )
175
+
176
+
177
@triton.jit
178
def kernel_unified_attention(
59
- # Output destinations. In 2D mode we write the final result into
60
- # ``output_ptr``; in 3D mode we write per-segment partials into the
61
- # three ``segm_*`` tensors and ``output_ptr`` is unused (callers may
62
- # pass any non-null pointer).
179
+ # Output destination for the 2D path. In 3D mode per-segment partials
180
+ # go to the ``segm_*`` tensors (see bottom of signature) and
181
+ # ``output_ptr`` is unused (callers may pass any non-null pointer).
182
output_ptr,
64
- segm_output_ptr,
65
- segm_max_ptr,
66
- segm_expsum_ptr,
183
# Inputs
184
query_ptr,
185
key_cache_ptr,
@@ -73,10 +189,6 @@ def kernel_unified_attention(
189
seq_lens_ptr,
190
alibi_slopes_ptr,
191
qq_bias_ptr,
76
- # Per-(token, head) scale caches (used iff KV_QUANT_MODE in {2, 3}).
77
- # For other modes callers may pass any non-null pointer.
78
- k_scale_cache_ptr,
79
- v_scale_cache_ptr,
192
# Scalars
193
scale,
194
k_scale,
@@ -112,12 +224,6 @@ def kernel_unified_attention(
224
stride_v_cache_1: tl.int64, # int
225
stride_v_cache_2: tl.int64, # int
226
stride_v_cache_3: tl.constexpr, # int
115
- stride_ks_blk: tl.int64,
116
- stride_ks_slot: tl.int64,
117
- stride_ks_head: tl.int64,
118
- stride_vs_blk: tl.int64,
119
- stride_vs_slot: tl.int64,
120
- stride_vs_head: tl.int64,
227
query_start_len_ptr,
228
BLOCK_Q: tl.constexpr,
229
num_seqs: tl.int32,
@@ -129,6 +235,23 @@ def kernel_unified_attention(
235
# to ``[segm_idx, segm_idx+1) × tiles_per_segment`` and writes
236
# per-segment partials, finalized by ``reduce_segments``.
237
IS_3D: tl.constexpr,
238
+ # Parameters below default to None so Triton can skip materialising them
239
+ # on call sites where the corresponding constexpr branch is dead.
240
+ # Credit: @quinnlp identified this as a perf regression source in
241
+ # intel/intel-xpu-backend-for-triton#6758 (review comment r3204641104).
242
+ # Per-segment outputs: used in 3D mode; unused in 2D (IS_3D=False).
243
+ segm_output_ptr=None,
244
+ segm_max_ptr=None,
245
+ segm_expsum_ptr=None,
246
+ # Per-(token, head) scale caches: used iff KV_QUANT_MODE in {2, 3}.
247
+ k_scale_cache_ptr=None,
248
+ v_scale_cache_ptr=None,
249
+ stride_ks_blk: tl.int64 = None,
250
+ stride_ks_slot: tl.int64 = None,
251
+ stride_ks_head: tl.int64 = None,
252
+ stride_vs_blk: tl.int64 = None,
253
+ stride_vs_slot: tl.int64 = None,
254
+ stride_vs_head: tl.int64 = None,
255
# KV cache quantization mode handled inside this kernel via constexpr
256
# branches: NONE (0), FP8_PER_TENSOR (1), INT8_PER_TOKEN_HEAD (2),
257
# FP8_PER_TOKEN_HEAD (3).
@@ -140,9 +263,20 @@ def kernel_unified_attention(
263
# over ``SLIDING_WINDOW`` inside the helpers. ``-1`` disables.
264
CHUNK_LOOKBACK: tl.constexpr = -1,
265
CHUNK_SIZE: tl.constexpr = -1,
266
+ # Tensor-descriptor load/store for HW 2D block reads on Intel Xe2/Xe3.
267
+ # ``USE_TD`` gates KV tile loads; ``USE_TD_QO`` separately gates Q/output
268
+ # (see ``unified_attention`` wrapper for the gating rules).
269
+ USE_TD: tl.constexpr = False,
270
+ USE_TD_QO: tl.constexpr = False,
271
):
272
USE_PER_TOKEN_HEAD_SCALES: tl.constexpr = KV_QUANT_MODE >= 2
273
274
+ if USE_TD:
275
+ tl.static_assert(
276
+ BLOCK_SIZE % TILE_SIZE == 0,
277
+ "USE_TD requires BLOCK_SIZE to be a multiple of TILE_SIZE",
278
+ )
279
+
280
q_block_global_idx = tl.program_id(0)
281
kv_head_idx = tl.program_id(1)
282
segm_idx = tl.program_id(2) if IS_3D else 0
@@ -167,6 +301,12 @@ def kernel_unified_attention(
301
else:
302
tiles_per_segment = 0
303
304
+ # Number of valid query rows in this block (used by TD descriptor
305
+ # shapes, but always computed so the variable stays in scope).
306
+ q_block_local_len = tl.minimum(
307
+ BLOCK_Q, cur_batch_query_len - q_block_local_idx * BLOCK_Q
308
+ )
309
+
310
offs_m = tl.arange(0, BLOCK_M)
311
offs_d = tl.arange(0, HEAD_SIZE_PADDED)
312
offs_t = tl.arange(0, TILE_SIZE)
@@ -185,11 +325,27 @@ def kernel_unified_attention(
325
query_mask_1 = tl.where(query_offset_1 < num_query_heads, 1, 0).to(tl.int1)
326
327
# Q : (BLOCK_M, HEAD_SIZE_PADDED)
188
- Q = tl.load(
189
- query_ptr + query_offset,
190
- mask=dim_mask[None, :] & query_mask_0[:, None] & query_mask_1[:, None],
191
- other=0.0,
192
- )
328
+ if USE_TD_QO:
329
+ Q = _load_q_td(
330
+ query_ptr,
331
+ q_block_local_len,
332
+ query_stride_0,
333
+ query_stride_1,
334
+ cur_batch_in_all_start_index,
335
+ q_block_local_idx,
336
+ kv_head_idx,
337
+ num_queries_per_kv,
338
+ BLOCK_Q,
339
+ BLOCK_M,
340
+ HEAD_SIZE,
341
+ HEAD_SIZE_PADDED,
342
+ )
343
+ else:
344
+ Q = tl.load(
345
+ query_ptr + query_offset,
346
+ mask=dim_mask[None, :] & query_mask_0[:, None] & query_mask_1[:, None],
347
+ other=0.0,
348
+ )
349
350
block_table_offset = seq_idx * block_table_stride
351
@@ -237,31 +393,71 @@ def kernel_unified_attention(
393
block_tables_ptr + block_table_offset + seq_offset // BLOCK_SIZE
394
).to(tl.int64)
395
240
- v_offset = (
241
- physical_block_idx[:, None] * stride_v_cache_0
242
- + kv_head_idx * stride_v_cache_2
243
- + offs_d[None, :] * stride_v_cache_3
244
- + (seq_offset % BLOCK_SIZE)[:, None] * stride_v_cache_1
245
- )
246
- k_offset = (
247
- physical_block_idx[None, :] * stride_k_cache_0
248
- + kv_head_idx * stride_k_cache_2
249
- + offs_d[:, None] * stride_k_cache_3
250
- + (seq_offset % BLOCK_SIZE)[None, :] * stride_k_cache_1
251
- )
252
- # K : (HEAD_SIZE, TILE_SIZE)
253
- K_load = tl.load(
254
- key_cache_ptr + k_offset,
255
- mask=dim_mask[:, None] & tile_mask[None, :],
256
- other=0.0,
257
- )
396
+ if USE_TD:
397
+ # All TILE_SIZE slots within a single KV tile map to one
398
+ # physical block (guaranteed by ``BLOCK_SIZE % TILE_SIZE == 0``
399
+ # from the static_assert above), so load the block index as
400
+ # a scalar instead of a broadcast reduction.
401
+ offset_in_block = (j * TILE_SIZE) % BLOCK_SIZE
402
+ physical_block_scalar = tl.load(
403
+ block_tables_ptr + block_table_offset + (j * TILE_SIZE) // BLOCK_SIZE
404
+ ).to(tl.int64)
405
+ # K : (HEAD_SIZE, TILE_SIZE)
406
+ K_load = _load_kv_tile_td(
407
+ key_cache_ptr,
408
+ physical_block_scalar,
409
+ kv_head_idx,
410
+ offset_in_block,
411
+ stride_k_cache_0,
412
+ stride_k_cache_1,
413
+ stride_k_cache_2,
414
+ stride_k_cache_3,
415
+ BLOCK_SIZE,
416
+ TILE_SIZE,
417
+ HEAD_SIZE,
418
+ HEAD_SIZE_PADDED,
419
+ ).T
420
+ # V : (TILE_SIZE, HEAD_SIZE)
421
+ V_load = _load_kv_tile_td(
422
+ value_cache_ptr,
423
+ physical_block_scalar,
424
+ kv_head_idx,
425
+ offset_in_block,
426
+ stride_v_cache_0,
427
+ stride_v_cache_1,
428
+ stride_v_cache_2,
429
+ stride_v_cache_3,
430
+ BLOCK_SIZE,
431
+ TILE_SIZE,
432
+ HEAD_SIZE,
433
+ HEAD_SIZE_PADDED,
434
+ )
435
+ else:
436
+ v_offset = (
437
+ physical_block_idx[:, None] * stride_v_cache_0
438
+ + kv_head_idx * stride_v_cache_2
439
+ + offs_d[None, :] * stride_v_cache_3
440
+ + (seq_offset % BLOCK_SIZE)[:, None] * stride_v_cache_1
441
+ )
442
+ k_offset = (
443
+ physical_block_idx[None, :] * stride_k_cache_0
444
+ + kv_head_idx * stride_k_cache_2
445
+ + offs_d[:, None] * stride_k_cache_3
446
+ + (seq_offset % BLOCK_SIZE)[None, :] * stride_k_cache_1
447
+ )
448
+ # K : (HEAD_SIZE, TILE_SIZE)
449
+ K_load = tl.load(
450
+ key_cache_ptr + k_offset,
451
+ mask=dim_mask[:, None] & tile_mask[None, :],
452
+ other=0.0,
453
+ )
454
+ # V : (TILE_SIZE, HEAD_SIZE)
455
+ V_load = tl.load(
456
+ value_cache_ptr + v_offset,
457
+ mask=dim_mask[None, :] & tile_mask[:, None],
458
+ other=0.0,
459
+ )
460
K = _cast_kv_tile(K_load, Q, k_scale, KV_QUANT_MODE)
259
- # V : (TILE_SIZE, HEAD_SIZE)
260
- V_load = tl.load(
261
- value_cache_ptr + v_offset,
262
- mask=dim_mask[None, :] & tile_mask[:, None],
263
- other=0.0,
264
- )
461
V = _cast_kv_tile(V_load, Q, v_scale, KV_QUANT_MODE)
462
463
# Per-(token, head) scales for INT8 / FP8 per-token-head modes.
@@ -342,18 +538,45 @@ def kernel_unified_attention(
538
# ---- Epilogue ---------------------------------------------------------
539
if IS_3D:
540
# Store per-segment partials; finalized by ``reduce_segments``.
345
- segm_output_offset = (
346
- query_offset_0[:, None].to(tl.int64)
347
- * (num_query_heads * NUM_SEGMENTS_PER_SEQ * HEAD_SIZE_PADDED)
348
- + query_offset_1[:, None] * (NUM_SEGMENTS_PER_SEQ * HEAD_SIZE_PADDED)
349
- + segm_idx * HEAD_SIZE_PADDED
350
- + tl.arange(0, HEAD_SIZE_PADDED)[None, :]
351
- )
352
- tl.store(
353
- segm_output_ptr + segm_output_offset,
354
- acc,
355
- mask=dim_mask[None, :] & query_mask_0[:, None] & query_mask_1[:, None],
356
- )
541
+ if USE_TD_QO:
542
+ # 3D target: segm_output[token, head, segm_idx, :]. Advance
543
+ # the base to the correct (token-start, head-start, segm)
544
+ # slice; strides step between tokens / heads of the flattened
545
+ # (T, H, SEGS, PAD) layout.
546
+ segm_base = (
547
+ segm_output_ptr
548
+ + (cur_batch_in_all_start_index + q_block_local_idx * BLOCK_Q).to(
549
+ tl.int64
550
+ )
551
+ * (num_query_heads * NUM_SEGMENTS_PER_SEQ * HEAD_SIZE_PADDED)
552
+ + (kv_head_idx * num_queries_per_kv)
553
+ * (NUM_SEGMENTS_PER_SEQ * HEAD_SIZE_PADDED)
554
+ + segm_idx * HEAD_SIZE_PADDED
555
+ )
556
+ _store_output_td(
557
+ segm_base,
558
+ acc,
559
+ q_block_local_len,
560
+ num_query_heads * NUM_SEGMENTS_PER_SEQ * HEAD_SIZE_PADDED,
561
+ NUM_SEGMENTS_PER_SEQ * HEAD_SIZE_PADDED,
562
+ num_queries_per_kv,
563
+ BLOCK_Q,
564
+ HEAD_SIZE,
565
+ HEAD_SIZE_PADDED,
566
+ )
567
+ else:
568
+ segm_output_offset = (
569
+ query_offset_0[:, None].to(tl.int64)
570
+ * (num_query_heads * NUM_SEGMENTS_PER_SEQ * HEAD_SIZE_PADDED)
571
+ + query_offset_1[:, None] * (NUM_SEGMENTS_PER_SEQ * HEAD_SIZE_PADDED)
572
+ + segm_idx * HEAD_SIZE_PADDED
573
+ + tl.arange(0, HEAD_SIZE_PADDED)[None, :]
574
+ )
575
+ tl.store(
576
+ segm_output_ptr + segm_output_offset,
577
+ acc,
578
+ mask=dim_mask[None, :] & query_mask_0[:, None] & query_mask_1[:, None],
579
+ )
580
store_segm_reduce_scalars(
581
segm_max_ptr,
582
segm_expsum_ptr,
@@ -372,16 +595,38 @@ def kernel_unified_attention(
595
if USE_FP8:
596
acc = acc * tl.load(out_scale)
597
acc = tl.clamp(acc, FP8_MIN, FP8_MAX)
375
- output_offset = (
376
- query_offset_0[:, None] * output_stride_0
377
- + query_offset_1[:, None] * output_stride_1
378
- + offs_d[None, :]
379
- )
380
- tl.store(
381
- output_ptr + output_offset,
382
- acc,
383
- mask=dim_mask[None, :] & query_mask_0[:, None] & query_mask_1[:, None],
384
- )
598
+ if USE_TD_QO:
599
+ # 2D target: flat output[token, head, :]. Strides come
600
+ # straight from the caller (``output_stride_0`` per token,
601
+ # ``output_stride_1`` per head).
602
+ output_base = (
603
+ output_ptr
604
+ + (cur_batch_in_all_start_index + q_block_local_idx * BLOCK_Q)
605
+ * output_stride_0
606
+ + (kv_head_idx * num_queries_per_kv) * output_stride_1
607
+ )
608
+ _store_output_td(
609
+ output_base,
610
+ acc,
611
+ q_block_local_len,
612
+ output_stride_0,
613
+ output_stride_1,
614
+ num_queries_per_kv,
615
+ BLOCK_Q,
616
+ HEAD_SIZE,
617
+ HEAD_SIZE_PADDED,
618
+ )
619
+ else:
620
+ output_offset = (
621
+ query_offset_0[:, None] * output_stride_0
622
+ + query_offset_1[:, None] * output_stride_1
623
+ + offs_d[None, :]
624
+ )
625
+ tl.store(
626
+ output_ptr + output_offset,
627
+ acc,
628
+ mask=dim_mask[None, :] & query_mask_0[:, None] & query_mask_1[:, None],
629
+ )
630
631
632
@triton.jit
@@ -538,6 +783,11 @@ def unified_attention(
783
v_scale_cache=None, # [num_blocks, block_size, num_kv_heads] float32
784
# Chunked attention: restrict attention to aligned blocks with lookback.
785
chunk_lookback=-1,
786
+ # Tensor-descriptor mode: use ``tl.make_tensor_descriptor`` for Q/K/V
787
+ # loads and output stores. Enables HW 2D block reads on Intel Xe2/Xe3.
788
+ # The non-TD branch is dead-code-eliminated at Triton compile time so
789
+ # disabling this flag costs nothing.
790
+ use_td: bool = False,
791
):
792
assert causal, "Only causal attention is supported"
793
assert q_descale is None, "Q scales not supported"
@@ -608,6 +858,52 @@ def unified_attention(
858
head_size, sliding_window_val, q.element_size(), is_prefill=False
859
)
860
861
+ # USE_TD requires BLOCK_SIZE % TILE_SIZE == 0 (enforced by a
862
+ # ``tl.static_assert`` in the kernel). The default prefill tile
863
+ # size (32) is larger than a common ``block_size=16``, so clamp it
864
+ # down when TD is enabled. Zero overhead when disabled.
865
+ if use_td:
866
+ TILE_SIZE_PREFILL = min(TILE_SIZE_PREFILL, block_size)
867
+ TILE_SIZE_DECODE = min(TILE_SIZE_DECODE, block_size)
868
+
869
+ # Tensor descriptors for Q load / output store require every element
870
+ # of ``block_shape`` to be a power of 2. ``num_queries_per_kv`` is
871
+ # not always pow2 (e.g. Qwen2-7B: 28 / 4 = 7), so gate the Q/O paths
872
+ # separately from the KV tile loads (whose ``block_shape`` does not
873
+ # include ``num_queries_per_kv``).
874
+ #
875
+ # The Q/O descriptors also encode ``HEAD_SIZE_PADDED`` on the inner
876
+ # axis while the backing buffers (both flat output and per-segment
877
+ # output) are laid out with ``HEAD_SIZE``. When they differ (e.g.
878
+ # Phi-3's head_size=96 → HEAD_SIZE_PADDED=128) the store would spill
879
+ # padded lanes into neighbouring heads because tensor-descriptor
880
+ # stores don't mask the padded tail. Fall back to the pointer path
881
+ # for Q/O in that case — KV tile loads are unaffected because their
882
+ # ``shape`` already matches ``block_shape`` on the inner axis.
883
+ head_size_padded = triton.next_power_of_2(head_size)
884
+ _is_pow2_nq = (num_queries_per_kv & (num_queries_per_kv - 1)) == 0
885
+ _is_pow2_hs = head_size == head_size_padded
886
+ use_td_qo = use_td and _is_pow2_nq and _is_pow2_hs
887
+
888
+ # ``_load_q_td`` / ``_store_output_td`` flatten ``(num_queries_per_kv,
889
+ # HEAD_SIZE)`` into a single contiguous inner axis. That's only
890
+ # equivalent to the pointer path when the ``num_queries_per_kv`` heads
891
+ # for this KV group start at ``kv_head_idx * num_queries_per_kv`` and
892
+ # lie exactly HEAD_SIZE apart — i.e. ``query_stride_1 == HEAD_SIZE``
893
+ # and ``output_stride_1 == head_size``. This is the default vLLM
894
+ # query/output layout; assert it explicitly so we fail fast if a
895
+ # future caller passes a non-contiguous query tensor.
896
+ if use_td_qo:
897
+ assert q.stride(1) == head_size, (
898
+ f"USE_TD_QO requires contiguous query heads "
899
+ f"(q.stride(1) = {q.stride(1)} != head_size = {head_size}); "
900
+ f"set VLLM_TRITON_ATTN_USE_TD=0 or pad the query layout."
901
+ )
902
+ assert out.stride(1) == head_size, (
903
+ f"USE_TD_QO requires contiguous output heads "
904
+ f"(out.stride(1) = {out.stride(1)} != head_size = {head_size})."
905
+ )
906
+
907
# Launch the 2D kernel if
908
# 1. No intermediate tiled softmax buffers for the 3D kernel have been allocated, or
909
# 2. The batch includes at least one prefill request, or
@@ -689,7 +985,7 @@ def unified_attention(
985
BLOCK_SIZE=block_size,
986
TILE_SIZE=tile_size,
987
HEAD_SIZE=head_size,
692
- HEAD_SIZE_PADDED=triton.next_power_of_2(head_size),
988
+ HEAD_SIZE_PADDED=head_size_padded,
989
USE_ALIBI_SLOPES=use_alibi_slopes,
990
USE_ALIBI_SQRT=use_alibi_sqrt,
991
USE_QQ_BIAS=use_qq_bias,
@@ -723,6 +1019,8 @@ def unified_attention(
1019
KV_QUANT_MODE=kv_quant_mode,
1020
CHUNK_LOOKBACK=chunk_lookback,
1021
CHUNK_SIZE=chunk_size,
1022
+ USE_TD=use_td,
1023
+ USE_TD_QO=use_td_qo,
1024
)
1025
1026
if use_3d:
@@ -740,7 +1038,7 @@ def unified_attention(
1038
block_table_stride=block_table.stride(0),
1039
TILE_SIZE=TILE_SIZE_DECODE,
1040
HEAD_SIZE=head_size,
743
- HEAD_SIZE_PADDED=triton.next_power_of_2(head_size),
1041
+ HEAD_SIZE_PADDED=head_size_padded,
1042
query_start_len_ptr=cu_seqlens_q,
1043
BLOCK_Q=BLOCK_Q,
1044
NUM_SEGMENTS_PER_SEQ=num_par_softmax_segments,