@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / f81daf8880
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

[Attention] add triton diff-kv backend for mimo (#41797)

Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>

Jiangyun Zhu committed Jun 11, 2026 at 23:36 UTC f81daf8880632eea46590a8222c082a1e27fd11f
8 files changed +1041 -9
.buildkite/test_areas/kernels.yaml
+13
@@ -75,6 +75,19 @@ steps:
75 - pytest -v -s kernels/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
76 parallelism: 2
77
78 +- label: Kernels Attention DiffKV Test (H100)
79 + key: kernels-attention-diffkv-test-h100
80 + timeout_in_minutes: 20
81 + device: h100
82 + num_devices: 1
83 + source_file_dependencies:
84 + - vllm/v1/attention/ops/triton_unified_attention_diffkv.py
85 + - vllm/v1/attention/backends/triton_attn_diffkv.py
86 + - vllm/v1/attention/backends/flash_attn_diffkv.py
87 + - tests/kernels/attention/test_triton_unified_attention_diffkv.py
88 + commands:
89 + - pytest -v -s kernels/attention/test_triton_unified_attention_diffkv.py
90 +
91 - label: Kernels Quantization Test %N
92 key: kernels-quantization-test
93 timeout_in_minutes: 90
docs/design/attention_backends.md
+1
@@ -181,6 +181,7 @@ Priority is **1 = highest** (tried first).
181 | `ROCM_AITER_UNIFIED_ATTN` | | fp16, bf16 | `auto` | %16 | Any | ✅ | ❌ | ✅ | ❌ | All | N/A |
182 | `ROCM_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | 32, 64, 80, 96, 128, 160, 192, 224, 256 | ❌ | ✅ | ✅ | ❌ | Decoder, Encoder, Encoder Only | N/A |
183 | `TRITON_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `int8_per_token_head`, `fp8_per_token_head` | %16 | Any | ✅ | ❌ | ✅ | ❌ | All | Any |
184 +| `TRITON_ATTN_DIFFKV` | | fp16, bf16 | `auto`, `bfloat16` | Any | Any | ❌ | ❌ | ❌ | ❌ | Decoder | Any |
185 | `TURBOQUANT` | | fp16, bf16 | `turboquant_k8v4`, `turboquant_4bit_nc`, `turboquant_k3v4_nc`, `turboquant_3bit_nc` | 16, 32, 64, 128 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | Any |
186
187 > **†** FlashInfer uses TRTLLM attention on Blackwell (SM100), which supports sinks. Disable via `--attention-config.use_trtllm_attention=0`.
tests/kernels/attention/test_triton_unified_attention_diffkv.py new
+189
@@ -0,0 +1,189 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +Unit tests for the Triton DiffKV unified-attention kernel.
5 +"""
6 +
7 +import pytest
8 +import torch
9 +
10 +from vllm.platforms import current_platform
11 +from vllm.utils.math_utils import next_power_of_2
12 +from vllm.utils.torch_utils import (
13 + canonicalize_singleton_dim_strides,
14 + set_random_seed,
15 +)
16 +from vllm.v1.attention.backends.fa_utils import (
17 + get_flash_attn_version,
18 + is_flash_attn_varlen_func_available,
19 +)
20 +from vllm.v1.attention.ops.triton_unified_attention_diffkv import (
21 + unified_attention_diffkv,
22 +)
23 +
24 +DEVICE_TYPE = current_platform.device_type
25 +
26 +# (num_query_heads, num_kv_heads): MHA, GQA, and the num_kv_heads==1
27 +# (degenerate-stride) case.
28 +NUM_HEADS = [(4, 4), (8, 2), (5, 1)]
29 +# (head_size_qk, head_size_v). (192, 128) is the canonical asymmetric
30 +# DiffKV shape; FA4 on Blackwell only supports head_size>128 when it is
31 +# 192, and FA3 on Hopper supports it too -- so this pair is runnable on
32 +# both. (128, 128) keeps the equal-dim path covered through the DiffKV
33 +# kernel.
34 +HEAD_SIZES = [(128, 128), (192, 128)]
35 +BLOCK_SIZES = [16]
36 +DTYPES = [torch.bfloat16]
37 +
38 +NUM_BLOCKS = 2048
39 +
40 +# 0: 2D decode kernel; 8: 3D (split-KV) decode kernel.
41 +SEQ_THRESHOLD_3D_VALUES = [0, 8]
42 +
43 +NUM_PAR_SOFTMAX_SEGMENTS = 16
44 +
45 +
46 +def _alloc_segm_buffers(seq_threshold_3D: int, num_query_heads: int, head_size_v: int):
47 + """Allocate the split-KV softmax scratch (last dim == head_size_v)."""
48 + head_size_v_padded = next_power_of_2(head_size_v)
49 + segm_output = torch.empty(
50 + (
51 + seq_threshold_3D,
52 + num_query_heads,
53 + NUM_PAR_SOFTMAX_SEGMENTS,
54 + head_size_v_padded,
55 + ),
56 + dtype=torch.float32,
57 + )
58 + segm_max = torch.empty(
59 + (seq_threshold_3D, num_query_heads, NUM_PAR_SOFTMAX_SEGMENTS),
60 + dtype=torch.float32,
61 + )
62 + segm_expsum = torch.empty(
63 + (seq_threshold_3D, num_query_heads, NUM_PAR_SOFTMAX_SEGMENTS),
64 + dtype=torch.float32,
65 + )
66 + return segm_output, segm_max, segm_expsum
67 +
68 +
69 +@pytest.mark.parametrize(
70 + "seq_lens",
71 + [
72 + [(1, 1328), (5, 18), (129, 463)], # mixed prefill + decode
73 + [(1, 523), (1, 37), (1, 2011)], # decode-only (exercises 3D path)
74 + ],
75 +)
76 +@pytest.mark.parametrize("num_heads", NUM_HEADS)
77 +@pytest.mark.parametrize("head_sizes", HEAD_SIZES)
78 +@pytest.mark.parametrize("block_size", BLOCK_SIZES)
79 +@pytest.mark.parametrize("sliding_window", [None, 128])
80 +@pytest.mark.parametrize("soft_cap", [None, 50.0])
81 +@pytest.mark.parametrize("dtype", DTYPES)
82 +@pytest.mark.parametrize("seq_threshold_3D", SEQ_THRESHOLD_3D_VALUES)
83 +@torch.inference_mode()
84 +def test_triton_unified_attn_diffkv_vs_fa(
85 + seq_lens: list[tuple[int, int]],
86 + num_heads: tuple[int, int],
87 + head_sizes: tuple[int, int],
88 + sliding_window: int | None,
89 + soft_cap: float | None,
90 + dtype: torch.dtype,
91 + block_size: int,
92 + seq_threshold_3D: int,
93 +) -> None:
94 + head_size_qk, head_size_v = head_sizes
95 +
96 + # DiffKV requires FA3 (Hopper) / FA4 (Blackwell) as the reference.
97 + fa_version = get_flash_attn_version(head_size=head_size_qk, head_size_v=head_size_v)
98 + if not is_flash_attn_varlen_func_available() or fa_version not in (3, 4):
99 + pytest.skip(f"FA DiffKV needs FA3/FA4 (got version {fa_version}).")
100 +
101 + from vllm.v1.attention.backends.fa_utils import flash_attn_varlen_func
102 +
103 + torch.set_default_device(DEVICE_TYPE)
104 + set_random_seed(0)
105 +
106 + num_seqs = len(seq_lens)
107 + query_lens = [x[0] for x in seq_lens]
108 + kv_lens = [x[1] for x in seq_lens]
109 + num_query_heads, num_kv_heads = num_heads
110 + assert num_query_heads % num_kv_heads == 0
111 + max_query_len = max(query_lens)
112 + max_kv_len = max(kv_lens)
113 + window_size = (sliding_window - 1, 0) if sliding_window is not None else (-1, -1)
114 + scale = head_size_qk**-0.5
115 +
116 + query = torch.randn(sum(query_lens), num_query_heads, head_size_qk, dtype=dtype)
117 + # Packed KV cache: [num_blocks, block_size, num_kv_heads, hqk + hv].
118 + kv_cache = torch.randn(
119 + NUM_BLOCKS,
120 + block_size,
121 + num_kv_heads,
122 + head_size_qk + head_size_v,
123 + dtype=dtype,
124 + )
125 + key_cache = kv_cache[..., :head_size_qk]
126 + value_cache = kv_cache[..., head_size_qk:]
127 +
128 + cu_query_lens = torch.tensor([0] + query_lens, dtype=torch.int32).cumsum(
129 + dim=0, dtype=torch.int32
130 + )
131 + kv_lens_t = torch.tensor(kv_lens, dtype=torch.int32)
132 +
133 + max_num_blocks_per_seq = (max_kv_len + block_size - 1) // block_size
134 + block_tables = torch.randint(
135 + 0, NUM_BLOCKS, (num_seqs, max_num_blocks_per_seq), dtype=torch.int32
136 + )
137 +
138 + # ---- FlashAttention DiffKV (ground truth) ---------------------------
139 + # Mirror the backend: fix degenerate strides on size-1 dims so FA's
140 + # TMA path sees ≥16-byte-aligned strides (matters for num_kv_heads==1).
141 + fa_k = canonicalize_singleton_dim_strides(key_cache)
142 + fa_v = canonicalize_singleton_dim_strides(value_cache)
143 + fa_out = torch.empty(sum(query_lens), num_query_heads, head_size_v, dtype=dtype)
144 + flash_attn_varlen_func(
145 + q=query,
146 + k=fa_k,
147 + v=fa_v,
148 + out=fa_out,
149 + cu_seqlens_q=cu_query_lens,
150 + max_seqlen_q=max_query_len,
151 + seqused_k=kv_lens_t,
152 + max_seqlen_k=max_kv_len,
153 + softmax_scale=scale,
154 + causal=True,
155 + window_size=list(window_size),
156 + block_table=block_tables,
157 + softcap=soft_cap if soft_cap is not None else 0,
158 + fa_version=fa_version,
159 + )
160 +
161 + # ---- Triton DiffKV --------------------------------------------------
162 + segm_output, segm_max, segm_expsum = _alloc_segm_buffers(
163 + seq_threshold_3D, num_query_heads, head_size_v
164 + )
165 + triton_out = torch.empty(sum(query_lens), num_query_heads, head_size_v, dtype=dtype)
166 + unified_attention_diffkv(
167 + q=query,
168 + k=key_cache,
169 + v=value_cache,
170 + out=triton_out,
171 + cu_seqlens_q=cu_query_lens,
172 + seqused_k=kv_lens_t,
173 + softmax_scale=scale,
174 + causal=True,
175 + window_size=window_size,
176 + block_table=block_tables,
177 + softcap=soft_cap if soft_cap is not None else 0,
178 + max_seqlen_q=max_query_len,
179 + seq_threshold_3D=seq_threshold_3D,
180 + num_par_softmax_segments=NUM_PAR_SOFTMAX_SEGMENTS,
181 + softmax_segm_output=segm_output,
182 + softmax_segm_max=segm_max,
183 + softmax_segm_expsum=segm_expsum,
184 + )
185 +
186 + (
187 + torch.testing.assert_close(triton_out, fa_out, atol=2e-2, rtol=2e-2),
188 + f"triton vs FA max abs diff: {torch.max(torch.abs(triton_out - fa_out))}",
189 + )
vllm/model_executor/models/mimo_v2.py
+21 -7
@@ -47,9 +47,7 @@ from vllm.model_executor.model_loader.weight_utils import (
47 from vllm.model_executor.models.utils import sequence_parallel_chunk
48 from vllm.sequence import IntermediateTensors
49 from vllm.v1.attention.backend import AttentionType
50 -from vllm.v1.attention.backends.flash_attn_diffkv import (
51 - FlashAttentionDiffKVBackend,
52 -)
50 +from vllm.v1.attention.backends.registry import AttentionBackendEnum
51
52 from .interfaces import MixtureOfExperts, SupportsPP
53 from .utils import (
@@ -292,11 +290,27 @@ class MiMoV2Attention(nn.Module):
290
291 sliding_window = sliding_window_size if sliding_window_size > -1 else None
292
295 - # Use DiffKV backend when V has a different head dim than K
293 + # Use DiffKV backend when V has a different head dim than K.
294 + # Auto-pick FA-DiffKV when FA3/4 is usable on this device, else fall
295 + # back to TRITON_ATTN_DIFFKV. Users can force a choice via
296 + # `--attention-backend <FLASH_ATTN_DIFFKV|TRITON_ATTN_DIFFKV>`.
297 if self.v_head_dim != self.head_dim:
297 - FlashAttentionDiffKVBackend.set_head_size_v(self.v_head_dim)
298 - attn_backend = FlashAttentionDiffKVBackend
299 - logger.info_once("Using FlashAttentionDiffKVBackend for attention.")
298 + requested = get_current_vllm_config().attention_config.backend
299 + if requested is not None and requested.name.endswith("_DIFFKV"):
300 + backend_enum = requested
301 + else:
302 + fa_backend = AttentionBackendEnum.FLASH_ATTN_DIFFKV.get_class()
303 + if fa_backend.is_supported_on_current_device(
304 + head_size=self.head_dim,
305 + head_size_v=self.v_head_dim,
306 + has_sinks=self.attention_sink_bias is not None,
307 + ):
308 + backend_enum = AttentionBackendEnum.FLASH_ATTN_DIFFKV
309 + else:
310 + backend_enum = AttentionBackendEnum.TRITON_ATTN_DIFFKV
311 + attn_backend = backend_enum.get_class()
312 + attn_backend.set_head_size_v(self.v_head_dim)
313 + logger.info_once("Using %s for attention.", attn_backend.get_name())
314 else:
315 attn_backend = None
316
vllm/v1/attention/backends/flash_attn_diffkv.py
+24 -2
@@ -41,6 +41,30 @@ class FlashAttentionDiffKVBackend(FlashAttentionBackend):
41 def set_head_size_v(cls, head_size_v: int) -> None:
42 cls.head_size_v = head_size_v
43
44 + @classmethod
45 + def is_supported_on_current_device(
46 + cls,
47 + head_size: int,
48 + head_size_v: int,
49 + has_sinks: bool,
50 + ) -> bool:
51 + """Check whether FA3/4 with this DiffKV config is usable here.
52 +
53 + DiffKV (hdim_qk != hdim_v) requires FA3 or FA4
54 + """
55 + if not is_flash_attn_varlen_func_available():
56 + return False
57 + try:
58 + version = get_flash_attn_version(
59 + requires_alibi=False,
60 + head_size=head_size,
61 + head_size_v=head_size_v,
62 + has_sinks=has_sinks,
63 + )
64 + except Exception:
65 + return False
66 + return version in (3, 4)
67 +
68 @staticmethod
69 def get_name() -> str:
70 return "FLASH_ATTN_DIFFKV"
@@ -49,8 +73,6 @@ class FlashAttentionDiffKVBackend(FlashAttentionBackend):
73 def get_impl_cls() -> type["FlashAttentionImpl"]:
74 return FlashAttentionDiffKVImpl
75
52 - # Do not modify the interface of get_kv_cache_shape,
53 - # but consider head_size_v when returning result.
76 @staticmethod
77 def get_kv_cache_shape(
78 num_blocks: int,
vllm/v1/attention/backends/registry.py
+3
@@ -46,6 +46,9 @@ class AttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta):
46 "vllm.v1.attention.backends.flash_attn_diffkv.FlashAttentionDiffKVBackend"
47 )
48 TRITON_ATTN = "vllm.v1.attention.backends.triton_attn.TritonAttentionBackend"
49 + TRITON_ATTN_DIFFKV = (
50 + "vllm.v1.attention.backends.triton_attn_diffkv.TritonAttentionDiffKVBackend"
51 + )
52 ROCM_ATTN = "vllm.v1.attention.backends.rocm_attn.RocmAttentionBackend"
53 ROCM_AITER_MLA = "vllm.v1.attention.backends.mla.rocm_aiter_mla.AiterMLABackend"
54 ROCM_AITER_TRITON_MLA = (
vllm/v1/attention/backends/triton_attn_diffkv.py new
+261
@@ -0,0 +1,261 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Triton attention backend with different K/V head dimensions (DiffKV).
4 +
5 +The KV cache layout is identical to ``FlashAttentionDiffKVBackend`` — K
6 +and V are packed along the last dim:
7 +
8 + [num_blocks, block_size, num_kv_heads, head_size_qk + head_size_v]
9 +
10 +so existing helpers (``triton_reshape_and_cache_flash_diffkv``) are reused.
11 +"""
12 +
13 +from typing import ClassVar
14 +
15 +import torch
16 +
17 +from vllm.config import VllmConfig
18 +from vllm.config.cache import CacheDType
19 +from vllm.logger import init_logger
20 +from vllm.utils.math_utils import next_power_of_2
21 +from vllm.utils.torch_utils import is_quantized_kv_cache
22 +from vllm.v1.attention.backend import AttentionLayer, AttentionType
23 +from vllm.v1.attention.backends.triton_attn import (
24 + TritonAttentionBackend,
25 + TritonAttentionImpl,
26 + TritonAttentionMetadata,
27 + TritonAttentionMetadataBuilder,
28 +)
29 +from vllm.v1.attention.backends.utils import get_kv_cache_layout
30 +from vllm.v1.attention.ops.triton_reshape_and_cache_flash import (
31 + triton_reshape_and_cache_flash_diffkv,
32 +)
33 +from vllm.v1.attention.ops.triton_unified_attention_diffkv import (
34 + unified_attention_diffkv,
35 +)
36 +from vllm.v1.kv_cache_interface import AttentionSpec
37 +
38 +logger = init_logger(__name__)
39 +
40 +
41 +class TritonAttentionDiffKVMetadataBuilder(TritonAttentionMetadataBuilder):
42 + """Override the parent's softmax buffer last-dim to head_size_v.
43 +
44 + The parent allocates ``softmax_segm_output`` with last-dim sized to
45 + ``next_power_of_2(head_size)`` (== Q/K head size). For DiffKV the
46 + accumulator and per-segment partial outputs are V-shaped, so we
47 + re-allocate with ``next_power_of_2(head_size_v)`` instead.
48 + """
49 +
50 + def __init__(
51 + self,
52 + kv_cache_spec: AttentionSpec,
53 + layer_names: list[str],
54 + vllm_config: VllmConfig,
55 + device: torch.device,
56 + ):
57 + super().__init__(kv_cache_spec, layer_names, vllm_config, device)
58 +
59 + head_size_v = TritonAttentionDiffKVBackend.head_size_v
60 + head_size_v_padded = next_power_of_2(head_size_v)
61 + self.softmax_segm_output = torch.empty(
62 + (
63 + self.seq_threshold_3D,
64 + self.num_heads_q,
65 + self.num_par_softmax_segments,
66 + head_size_v_padded,
67 + ),
68 + dtype=torch.float32,
69 + device=device,
70 + )
71 +
72 +
73 +class TritonAttentionDiffKVBackend(TritonAttentionBackend):
74 + # V head dim — set per layer via ``set_head_size_v`` before instantiation.
75 + head_size_v: int = 128
76 +
77 + # No FP8 / int8 KV cache for the DiffKV path yet; require fp16/bf16/fp32.
78 + supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
79 + "auto",
80 + "bfloat16",
81 + ]
82 +
83 + @classmethod
84 + def set_head_size_v(cls, head_size_v: int) -> None:
85 + cls.head_size_v = head_size_v
86 +
87 + @staticmethod
88 + def get_name() -> str:
89 + return "TRITON_ATTN_DIFFKV"
90 +
91 + @staticmethod
92 + def get_impl_cls() -> type["TritonAttentionDiffKVImpl"]:
93 + return TritonAttentionDiffKVImpl
94 +
95 + @staticmethod
96 + def get_builder_cls() -> type["TritonAttentionDiffKVMetadataBuilder"]:
97 + return TritonAttentionDiffKVMetadataBuilder
98 +
99 + @staticmethod
100 + def get_kv_cache_shape(
101 + num_blocks: int,
102 + block_size: int,
103 + num_kv_heads: int,
104 + head_size: int,
105 + cache_dtype_str: str = "auto",
106 + ) -> tuple[int, ...]:
107 + if block_size % 16 != 0:
108 + raise ValueError("Block size must be a multiple of 16.")
109 + return (
110 + num_blocks,
111 + block_size,
112 + num_kv_heads,
113 + head_size + TritonAttentionDiffKVBackend.head_size_v,
114 + )
115 +
116 + @staticmethod
117 + def get_kv_cache_stride_order(
118 + include_num_layers_dimension: bool = False,
119 + ) -> tuple[int, ...]:
120 + cache_layout = get_kv_cache_layout()
121 + if cache_layout == "NHD" and include_num_layers_dimension:
122 + # (num_blocks, num_layers, block_size,
123 + # num_kv_heads, head_size + head_size_v)
124 + return (1, 0, 2, 3, 4)
125 + elif cache_layout == "NHD":
126 + return (0, 1, 2, 3)
127 + elif cache_layout == "HND" and include_num_layers_dimension:
128 + # (num_blocks, num_kv_heads, num_layers,
129 + # block_size, head_size + head_size_v)
130 + return (1, 3, 0, 2, 4)
131 + elif cache_layout == "HND":
132 + return (0, 2, 1, 3)
133 + else:
134 + raise ValueError(f"Unknown cache layout format {cache_layout}.")
135 +
136 + @classmethod
137 + def supports_head_size(cls, head_size: int) -> bool:
138 + # DiffKV K head sizes (e.g. 192 for MiMo-V2.5) need to be allowed.
139 + return head_size >= 32
140 +
141 + @classmethod
142 + def supports_attn_type(cls, attn_type: str) -> bool:
143 + # DiffKV only implements decoder self-attention. Unlike the parent
144 + # TritonAttentionBackend (which advertises all types), encoder
145 + # attention is not supported, so gate it here at backend selection.
146 + return attn_type == AttentionType.DECODER
147 +
148 +
149 +class TritonAttentionDiffKVImpl(TritonAttentionImpl):
150 + """Triton attention impl for the DiffKV packed KV cache layout."""
151 +
152 + def __init__(self, *args, **kwargs) -> None:
153 + super().__init__(*args, **kwargs)
154 + if is_quantized_kv_cache(self.kv_cache_dtype):
155 + raise NotImplementedError(
156 + "TritonAttentionDiffKVBackend does not yet support quantized "
157 + f"KV cache (got kv_cache_dtype={self.kv_cache_dtype!r})."
158 + )
159 + if self._is_per_token_head_quant:
160 + raise NotImplementedError(
161 + "TritonAttentionDiffKVBackend does not support per-token-head "
162 + "quantization."
163 + )
164 + if self.chunk_lookback > -1:
165 + raise NotImplementedError(
166 + "TritonAttentionDiffKVBackend does not support chunked "
167 + "attention with lookback."
168 + )
169 +
170 + def do_kv_cache_update(
171 + self,
172 + layer: AttentionLayer,
173 + key: torch.Tensor,
174 + value: torch.Tensor,
175 + kv_cache: torch.Tensor,
176 + slot_mapping: torch.Tensor,
177 + ) -> None:
178 + # Cache is packed [..., head_size_qk + head_size_v]; the diffkv
179 + # reshape kernel writes K to [..., :head_size_qk] and V to
180 + # [..., head_size_qk:hqk+hv].
181 + triton_reshape_and_cache_flash_diffkv(
182 + key,
183 + value,
184 + kv_cache,
185 + slot_mapping,
186 + self.kv_cache_dtype,
187 + layer._k_scale,
188 + layer._v_scale,
189 + )
190 +
191 + def fused_rope_kvcache_supported(self):
192 + # The fused rope+cache path assumes the standard 2-tensor layout.
193 + return False
194 +
195 + def forward(
196 + self,
197 + layer: torch.nn.Module,
198 + query: torch.Tensor,
199 + key: torch.Tensor,
200 + value: torch.Tensor,
201 + kv_cache: torch.Tensor,
202 + attn_metadata: TritonAttentionMetadata,
203 + output: torch.Tensor,
204 + output_scale: torch.Tensor | None = None,
205 + output_block_scale: torch.Tensor | None = None,
206 + ) -> torch.Tensor:
207 + """Forward pass.
208 +
209 + Shapes:
210 + query: [num_tokens, num_heads, head_size_qk]
211 + key: [num_tokens, num_kv_heads, head_size_qk]
212 + value: [num_tokens, num_kv_heads, head_size_v]
213 + kv_cache: [num_blocks, block_size, num_kv_heads,
214 + head_size_qk + head_size_v]
215 + output: [num_tokens, num_heads, head_size_v]
216 + """
217 + if output_scale is not None or output_block_scale is not None:
218 + raise NotImplementedError(
219 + "fused output quantization is not supported for "
220 + "TritonAttentionDiffKVImpl"
221 + )
222 +
223 + if attn_metadata is None:
224 + return output.fill_(0)
225 +
226 + assert attn_metadata.use_cascade is False, (
227 + "Cascade attention not supported for TritonAttentionDiffKVImpl"
228 + )
229 +
230 + num_actual_tokens = attn_metadata.num_actual_tokens
231 + head_size_qk = self.head_size
232 + head_size_v = TritonAttentionDiffKVBackend.head_size_v
233 +
234 + # Slice the packed cache into K / V views. Strides on dims 0/1/2
235 + # match the original cache; dim 3 stays contiguous (stride 1).
236 + key_cache = kv_cache[..., :head_size_qk]
237 + value_cache = kv_cache[..., head_size_qk : head_size_qk + head_size_v]
238 +
239 + unified_attention_diffkv(
240 + q=query[:num_actual_tokens],
241 + k=key_cache,
242 + v=value_cache,
243 + out=output[:num_actual_tokens],
244 + cu_seqlens_q=attn_metadata.query_start_loc,
245 + seqused_k=attn_metadata.seq_lens,
246 + softmax_scale=self.scale,
247 + causal=True,
248 + alibi_slopes=self.alibi_slopes,
249 + use_alibi_sqrt=self.use_alibi_sqrt,
250 + window_size=self.sliding_window,
251 + block_table=attn_metadata.block_table,
252 + softcap=self.logits_soft_cap,
253 + sinks=self.sinks,
254 + max_seqlen_q=attn_metadata.max_query_len,
255 + seq_threshold_3D=attn_metadata.seq_threshold_3D,
256 + num_par_softmax_segments=attn_metadata.num_par_softmax_segments,
257 + softmax_segm_output=attn_metadata.softmax_segm_output,
258 + softmax_segm_max=attn_metadata.softmax_segm_max,
259 + softmax_segm_expsum=attn_metadata.softmax_segm_expsum,
260 + )
261 + return output
vllm/v1/attention/ops/triton_unified_attention_diffkv.py new
+529
@@ -0,0 +1,529 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Triton unified attention with different K/V head dimensions (DiffKV).
4 +
5 +This is a slimmed fork of ``triton_unified_attention.py`` for models like
6 +MiMo-V2.5 where the V tensor's head dimension differs from K's. The KV cache
7 +is the same packed layout used by ``FlashAttentionDiffKVBackend``:
8 +
9 + kv_cache: [num_blocks, block_size, num_kv_heads, head_size_qk + head_size_v]
10 +
11 +We slice ``key_cache = kv_cache[..., :head_size_qk]`` and
12 +``value_cache = kv_cache[..., head_size_qk:]`` on the host, so the kernel
13 +takes two cache pointers but with two distinct head sizes.
14 +
15 +Both 2D and 3D launches are supported:
16 + - 2D: one program per (q-block, kv-head); tile-loop walks the full KV
17 + sequence; final output written directly. Used for prefill and large
18 + decode batches.
19 + - 3D: one program per (q-block, kv-head, segm); each program covers a
20 + KV slice and writes per-segment partials (max/expsum/output). A
21 + follow-up ``kernel_reduce_segments_diffkv`` combines them. Selected
22 + for decode-only batches whose 2D grid would under-fill the GPU.
23 +"""
24 +
25 +from typing import Any
26 +
27 +import torch
28 +
29 +import vllm.envs as envs
30 +from vllm.logger import init_logger
31 +from vllm.triton_utils import tl, triton
32 +from vllm.v1.attention.ops.triton_attention_helpers import (
33 + apply_alibi_to_score,
34 + apply_softcap,
35 + cdiv_fn,
36 + compute_kv_seq_mask,
37 + compute_tile_loop_bounds,
38 + find_seq_idx,
39 + init_softmax_M,
40 + resolve_seq_and_query_len,
41 + softmax_step,
42 + store_segm_reduce_scalars,
43 +)
44 +
45 +logger = init_logger(__name__)
46 +
47 +is_batch_invariant = envs.VLLM_BATCH_INVARIANT
48 +
49 +
50 +@triton.jit
51 +def kernel_unified_attention_diffkv(
52 + # Output destinations. In 2D mode we write the final result into
53 + # ``output_ptr``; in 3D mode we write per-segment partials into
54 + # ``segm_*`` and ``output_ptr`` is unused (callers may pass any
55 + # non-null pointer).
56 + output_ptr,
57 + segm_output_ptr,
58 + segm_max_ptr,
59 + segm_expsum_ptr,
60 + query_ptr,
61 + key_cache_ptr, # view of packed cache: [..., :head_size_qk]
62 + value_cache_ptr, # view of packed cache: [..., head_size_qk:hqk+hv]
63 + sink_ptr,
64 + block_tables_ptr,
65 + seq_lens_ptr,
66 + alibi_slopes_ptr,
67 + scale,
68 + softcap,
69 + num_query_heads: tl.constexpr,
70 + num_queries_per_kv: tl.constexpr,
71 + block_table_stride: tl.int64,
72 + query_stride_0: tl.int64,
73 + query_stride_1: tl.int64, # == HEAD_SIZE_QK
74 + output_stride_0: tl.int64,
75 + output_stride_1: tl.int64, # == HEAD_SIZE_V
76 + BLOCK_SIZE: tl.constexpr,
77 + TILE_SIZE: tl.constexpr,
78 + HEAD_SIZE_QK: tl.constexpr,
79 + HEAD_SIZE_QK_PADDED: tl.constexpr,
80 + HEAD_SIZE_V: tl.constexpr,
81 + HEAD_SIZE_V_PADDED: tl.constexpr,
82 + USE_ALIBI_SLOPES: tl.constexpr,
83 + USE_ALIBI_SQRT: tl.constexpr,
84 + USE_SOFTCAP: tl.constexpr,
85 + USE_SINKS: tl.constexpr,
86 + SLIDING_WINDOW: tl.constexpr,
87 + # Strides for both cache views (they share the same packed buffer, so
88 + # dims 0/1/2 strides match; only the per-head extent differs).
89 + stride_k_cache_0: tl.int64,
90 + stride_k_cache_1: tl.int64,
91 + stride_k_cache_2: tl.int64,
92 + stride_k_cache_3: tl.constexpr,
93 + stride_v_cache_0: tl.int64,
94 + stride_v_cache_1: tl.int64,
95 + stride_v_cache_2: tl.int64,
96 + stride_v_cache_3: tl.constexpr,
97 + query_start_len_ptr,
98 + BLOCK_Q: tl.constexpr,
99 + num_seqs: tl.int32,
100 + BLOCK_M: tl.constexpr,
101 + NUM_SEGMENTS_PER_SEQ: tl.constexpr,
102 + # ``IS_3D`` toggles between 2D layout (one program walks the full KV
103 + # sequence) and 3D layout (split-KV / FlashDecoding-style: per-segm
104 + # programs write partials, finalized by ``kernel_reduce_segments_diffkv``).
105 + IS_3D: tl.constexpr,
106 +):
107 + q_block_global_idx = tl.program_id(0)
108 + kv_head_idx = tl.program_id(1)
109 + segm_idx = tl.program_id(2) if IS_3D else 0
110 +
111 + (
112 + seq_idx,
113 + q_block_local_idx,
114 + cur_batch_in_all_start_index,
115 + cur_batch_query_len,
116 + seq_len,
117 + ) = resolve_seq_and_query_len(
118 + query_start_len_ptr, seq_lens_ptr, q_block_global_idx, num_seqs, BLOCK_Q
119 + )
120 +
121 + if q_block_local_idx * BLOCK_Q >= cur_batch_query_len:
122 + return
123 +
124 + if IS_3D:
125 + tiles_per_segment = cdiv_fn(seq_len, NUM_SEGMENTS_PER_SEQ * TILE_SIZE)
126 + if segm_idx * tiles_per_segment * TILE_SIZE >= seq_len:
127 + return
128 + else:
129 + tiles_per_segment = 0
130 +
131 + offs_m = tl.arange(0, BLOCK_M)
132 + offs_d_qk = tl.arange(0, HEAD_SIZE_QK_PADDED)
133 + offs_d_v = tl.arange(0, HEAD_SIZE_V_PADDED)
134 + offs_t = tl.arange(0, TILE_SIZE)
135 + query_pos = q_block_local_idx * BLOCK_Q + offs_m // num_queries_per_kv
136 +
137 + query_offset_0 = cur_batch_in_all_start_index + query_pos
138 + query_offset_1 = kv_head_idx * num_queries_per_kv + offs_m % num_queries_per_kv
139 + query_offset = (
140 + query_offset_0[:, None] * query_stride_0
141 + + query_offset_1[:, None] * query_stride_1
142 + + offs_d_qk[None, :]
143 + )
144 +
145 + dim_mask_qk = tl.where(offs_d_qk < HEAD_SIZE_QK, 1, 0).to(tl.int1)
146 + dim_mask_v = tl.where(offs_d_v < HEAD_SIZE_V, 1, 0).to(tl.int1)
147 + query_mask_0 = tl.where(query_pos < cur_batch_query_len, 1, 0).to(tl.int1)
148 + query_mask_1 = tl.where(query_offset_1 < num_query_heads, 1, 0).to(tl.int1)
149 +
150 + # Q : (BLOCK_M, HEAD_SIZE_QK_PADDED)
151 + Q = tl.load(
152 + query_ptr + query_offset,
153 + mask=dim_mask_qk[None, :] & query_mask_0[:, None] & query_mask_1[:, None],
154 + other=0.0,
155 + )
156 +
157 + block_table_offset = seq_idx * block_table_stride
158 +
159 + M = init_softmax_M(
160 + sink_ptr, query_offset_1, query_mask_1, segm_idx, BLOCK_M, USE_SINKS, IS_3D
161 + )
162 + L = tl.full([BLOCK_M], 1.0, dtype=tl.float32)
163 + # acc : (BLOCK_M, HEAD_SIZE_V_PADDED)
164 + acc = tl.zeros([BLOCK_M, HEAD_SIZE_V_PADDED], dtype=tl.float32)
165 +
166 + context_len = seq_len - cur_batch_query_len
167 +
168 + if USE_ALIBI_SLOPES:
169 + alibi_slope = tl.load(
170 + alibi_slopes_ptr + query_offset_1, mask=query_mask_1, other=0.0
171 + )
172 +
173 + loop_lo, loop_hi, max_seq_prefix_len = compute_tile_loop_bounds(
174 + context_len,
175 + seq_len,
176 + cur_batch_query_len,
177 + q_block_local_idx,
178 + segm_idx,
179 + tiles_per_segment,
180 + TILE_SIZE,
181 + BLOCK_M,
182 + BLOCK_Q,
183 + num_queries_per_kv,
184 + SLIDING_WINDOW,
185 + False, # USE_MM_PREFIX
186 + IS_3D,
187 + )
188 +
189 + for j in range(loop_lo, loop_hi):
190 + seq_offset = j * TILE_SIZE + offs_t
191 + tile_mask = seq_offset < max_seq_prefix_len
192 +
193 + physical_block_idx = tl.load(
194 + block_tables_ptr + block_table_offset + seq_offset // BLOCK_SIZE
195 + ).to(tl.int64)
196 +
197 + v_offset = (
198 + physical_block_idx[:, None] * stride_v_cache_0
199 + + kv_head_idx * stride_v_cache_2
200 + + offs_d_v[None, :] * stride_v_cache_3
201 + + (seq_offset % BLOCK_SIZE)[:, None] * stride_v_cache_1
202 + )
203 + k_offset = (
204 + physical_block_idx[None, :] * stride_k_cache_0
205 + + kv_head_idx * stride_k_cache_2
206 + + offs_d_qk[:, None] * stride_k_cache_3
207 + + (seq_offset % BLOCK_SIZE)[None, :] * stride_k_cache_1
208 + )
209 + # K : (HEAD_SIZE_QK_PADDED, TILE_SIZE)
210 + K_load = tl.load(
211 + key_cache_ptr + k_offset,
212 + mask=dim_mask_qk[:, None] & tile_mask[None, :],
213 + other=0.0,
214 + )
215 + K = K_load.to(Q.dtype)
216 + # V : (TILE_SIZE, HEAD_SIZE_V_PADDED)
217 + V_load = tl.load(
218 + value_cache_ptr + v_offset,
219 + mask=dim_mask_v[None, :] & tile_mask[:, None],
220 + other=0.0,
221 + )
222 + V = V_load.to(Q.dtype)
223 +
224 + query_abs_pos = context_len + query_pos[:, None]
225 + seq_mask = compute_kv_seq_mask(
226 + query_abs_pos,
227 + seq_offset,
228 + seq_idx,
229 + None, # mm_prefix_range_ptr
230 + SLIDING_WINDOW,
231 + False, # USE_MM_PREFIX
232 + 0, # MAX_MM_RANGES
233 + )
234 +
235 + # S : (BLOCK_M, TILE_SIZE)
236 + S = tl.zeros(shape=(BLOCK_M, TILE_SIZE), dtype=tl.float32)
237 + S += scale * tl.dot(Q, K)
238 +
239 + if USE_SOFTCAP:
240 + S = apply_softcap(S, softcap)
241 +
242 + S = tl.where(
243 + query_mask_1[:, None] & query_mask_0[:, None] & seq_mask, S, float("-inf")
244 + )
245 +
246 + if USE_ALIBI_SLOPES:
247 + S = apply_alibi_to_score(
248 + S, alibi_slope, seq_offset, context_len, query_pos, USE_ALIBI_SQRT
249 + )
250 +
251 + M, L, P, alpha = softmax_step(S, M, L)
252 + acc = acc * alpha[:, None]
253 +
254 + if SLIDING_WINDOW:
255 + qpos_lo = q_block_local_idx * BLOCK_Q
256 + V = tl.where(
257 + (context_len + qpos_lo - seq_offset[:, None]) < SLIDING_WINDOW,
258 + V,
259 + 0.0,
260 + )
261 + acc += tl.dot(P.to(V.dtype), V)
262 +
263 + # ---- Epilogue --------------------------------------------------------
264 + if IS_3D:
265 + # Store per-segment partials; finalized by reduce_segments_diffkv.
266 + segm_output_offset = (
267 + query_offset_0[:, None].to(tl.int64)
268 + * (num_query_heads * NUM_SEGMENTS_PER_SEQ * HEAD_SIZE_V_PADDED)
269 + + query_offset_1[:, None] * (NUM_SEGMENTS_PER_SEQ * HEAD_SIZE_V_PADDED)
270 + + segm_idx * HEAD_SIZE_V_PADDED
271 + + tl.arange(0, HEAD_SIZE_V_PADDED)[None, :]
272 + )
273 + tl.store(
274 + segm_output_ptr + segm_output_offset,
275 + acc,
276 + mask=dim_mask_v[None, :] & query_mask_0[:, None] & query_mask_1[:, None],
277 + )
278 + store_segm_reduce_scalars(
279 + segm_max_ptr,
280 + segm_expsum_ptr,
281 + query_offset_0,
282 + query_offset_1,
283 + segm_idx,
284 + M,
285 + L,
286 + query_mask_0,
287 + query_mask_1,
288 + num_query_heads,
289 + NUM_SEGMENTS_PER_SEQ,
290 + )
291 + else:
292 + acc = acc / L[:, None]
293 + output_offset = (
294 + query_offset_0[:, None] * output_stride_0
295 + + query_offset_1[:, None] * output_stride_1
296 + + offs_d_v[None, :]
297 + )
298 + tl.store(
299 + output_ptr + output_offset,
300 + acc,
301 + mask=dim_mask_v[None, :] & query_mask_0[:, None] & query_mask_1[:, None],
302 + )
303 +
304 +
305 +@triton.jit
306 +def kernel_reduce_segments_diffkv(
307 + output_ptr, # [num_tokens, num_query_heads, head_size_v]
308 + segm_output_ptr,
309 + # [num_tokens, num_query_heads, max_num_segments, head_size_v]
310 + segm_max_ptr, # [num_tokens, num_query_heads, max_num_segments]
311 + segm_expsum_ptr, # [num_tokens, num_query_heads, max_num_segments]
312 + seq_lens_ptr, # [num_seqs]
313 + num_seqs,
314 + num_query_heads: tl.constexpr,
315 + output_stride_0: tl.int64,
316 + output_stride_1: tl.int64, # == HEAD_SIZE_V
317 + TILE_SIZE: tl.constexpr,
318 + HEAD_SIZE_V: tl.constexpr,
319 + HEAD_SIZE_V_PADDED: tl.constexpr,
320 + query_start_len_ptr, # [num_seqs+1]
321 + BLOCK_Q: tl.constexpr,
322 + NUM_SEGMENTS_PER_SEQ: tl.constexpr,
323 +):
324 + """Combine per-segment partials into the final softmax output.
325 +
326 + Mirrors ``reduce_segments`` from triton_unified_attention.py but
327 + indexes V's head size (``HEAD_SIZE_V``) instead of the shared one.
328 + """
329 + query_token_idx = tl.program_id(0)
330 + query_head_idx = tl.program_id(1)
331 +
332 + seq_idx = find_seq_idx(
333 + query_start_len_ptr, query_token_idx, num_seqs, BLOCK_Q, False
334 + )
335 + seq_len = tl.load(seq_lens_ptr + seq_idx)
336 +
337 + tiles_per_segment = cdiv_fn(seq_len, NUM_SEGMENTS_PER_SEQ * TILE_SIZE)
338 + act_num_segments = cdiv_fn(seq_len, tiles_per_segment * TILE_SIZE)
339 + segm_mask = tl.arange(0, NUM_SEGMENTS_PER_SEQ) < tl.full(
340 + [NUM_SEGMENTS_PER_SEQ], act_num_segments, dtype=tl.int32
341 + )
342 + dim_mask = tl.where(tl.arange(0, HEAD_SIZE_V_PADDED) < HEAD_SIZE_V, 1, 0).to(
343 + tl.int1
344 + )
345 +
346 + segm_offset = (
347 + query_token_idx.to(tl.int64) * (num_query_heads * NUM_SEGMENTS_PER_SEQ)
348 + + query_head_idx * NUM_SEGMENTS_PER_SEQ
349 + + tl.arange(0, NUM_SEGMENTS_PER_SEQ)
350 + )
351 + segm_max = tl.load(segm_max_ptr + segm_offset, mask=segm_mask, other=float("-inf"))
352 + overall_max = tl.max(segm_max)
353 +
354 + segm_expsum = tl.load(segm_expsum_ptr + segm_offset, mask=segm_mask, other=0.0)
355 + segm_expsum = segm_expsum * tl.exp(segm_max - overall_max)
356 + overall_expsum = tl.sum(segm_expsum)
357 +
358 + segm_output_offset = (
359 + query_token_idx.to(tl.int64)
360 + * (num_query_heads * NUM_SEGMENTS_PER_SEQ * HEAD_SIZE_V_PADDED)
361 + + query_head_idx * (NUM_SEGMENTS_PER_SEQ * HEAD_SIZE_V_PADDED)
362 + + tl.arange(0, NUM_SEGMENTS_PER_SEQ)[:, None] * HEAD_SIZE_V_PADDED
363 + + tl.arange(0, HEAD_SIZE_V_PADDED)[None, :]
364 + )
365 + segm_output = tl.load(
366 + segm_output_ptr + segm_output_offset,
367 + mask=segm_mask[:, None] & dim_mask[None, :],
368 + other=0.0,
369 + )
370 + segm_output *= tl.exp(segm_max - overall_max)[:, None]
371 + acc_sum = tl.sum(segm_output, axis=0)
372 + acc = tl.where(overall_expsum == 0.0, 0.0, acc_sum / overall_expsum)
373 +
374 + output_offset = (
375 + query_token_idx * output_stride_0
376 + + query_head_idx * output_stride_1
377 + + tl.arange(0, HEAD_SIZE_V_PADDED)
378 + )
379 + tl.store(output_ptr + output_offset, acc, mask=dim_mask)
380 +
381 +
382 +def unified_attention_diffkv(
383 + q, # [num_tokens, num_query_heads, head_size_qk]
384 + k, # view: [num_blocks, block_size, num_kv_heads, head_size_qk]
385 + v, # view: [num_blocks, block_size, num_kv_heads, head_size_v]
386 + out, # [num_tokens, num_query_heads, head_size_v]
387 + cu_seqlens_q,
388 + seqused_k,
389 + softmax_scale,
390 + causal,
391 + window_size,
392 + block_table,
393 + softcap,
394 + max_seqlen_q: int = 1,
395 + alibi_slopes=None,
396 + sinks=None,
397 + use_alibi_sqrt=False,
398 + # 3D / split-KV softmax buffers. When all four are provided and the
399 + # batch is decode-only with few sequences, the 3D path is taken.
400 + seq_threshold_3D: int | None = None,
401 + num_par_softmax_segments: int | None = None,
402 + softmax_segm_output: torch.Tensor | None = None,
403 + softmax_segm_max: torch.Tensor | None = None,
404 + softmax_segm_expsum: torch.Tensor | None = None,
405 +):
406 + assert causal, "Only causal attention is supported"
407 +
408 + if sinks is not None:
409 + assert sinks.shape[0] == q.shape[1], "Sinks must be num_query_heads size"
410 +
411 + use_alibi_slopes = alibi_slopes is not None
412 +
413 + block_size = v.shape[1]
414 + num_seqs = len(seqused_k)
415 + num_query_heads = q.shape[1]
416 + num_kv_heads = k.shape[2]
417 + num_queries_per_kv = num_query_heads // num_kv_heads
418 + head_size_qk = q.shape[2]
419 + head_size_v = v.shape[3]
420 +
421 + BLOCK_M = (
422 + 16 if num_queries_per_kv <= 16 else triton.next_power_of_2(num_queries_per_kv)
423 + )
424 + BLOCK_Q = BLOCK_M // num_queries_per_kv
425 +
426 + total_num_q_blocks = q.shape[0] // BLOCK_Q + num_seqs
427 +
428 + sliding_window_val = 1 + window_size[0] if window_size[0] >= 0 else 0
429 +
430 + # Decide between 2D and 3D launch. Mirrors the standard launcher:
431 + # 3D requires preallocated softmax buffers, decode-only batches, and
432 + # a small number of sequences (otherwise 2D already saturates the SM).
433 + use_3d = not (
434 + seq_threshold_3D is None
435 + or num_par_softmax_segments is None
436 + or softmax_segm_output is None
437 + or softmax_segm_max is None
438 + or softmax_segm_expsum is None
439 + or max_seqlen_q > 1
440 + or num_seqs > seq_threshold_3D
441 + or is_batch_invariant
442 + )
443 +
444 + # Tile size: 32 for prefill-class kernels. Decode (small Q) prefers
445 + # smaller tiles to expose more parallelism along the KV dim.
446 + tile_size = 32 if not use_3d else (16 if q.element_size() >= 2 else 32)
447 +
448 + grid: tuple[Any, ...]
449 + if use_3d:
450 + grid = (total_num_q_blocks, num_kv_heads, num_par_softmax_segments)
451 + segm_output_ptr = softmax_segm_output
452 + segm_max_ptr = softmax_segm_max
453 + segm_expsum_ptr = softmax_segm_expsum
454 + num_segments = num_par_softmax_segments
455 + else:
456 + grid = (total_num_q_blocks, num_kv_heads)
457 + # 2D never touches the segm tensors but Triton wants a non-null
458 + # pointer; reuse ``out``.
459 + segm_output_ptr = out
460 + segm_max_ptr = out
461 + segm_expsum_ptr = out
462 + num_segments = 1
463 +
464 + kernel_unified_attention_diffkv[grid](
465 + output_ptr=out,
466 + segm_output_ptr=segm_output_ptr,
467 + segm_max_ptr=segm_max_ptr,
468 + segm_expsum_ptr=segm_expsum_ptr,
469 + query_ptr=q,
470 + key_cache_ptr=k,
471 + value_cache_ptr=v,
472 + sink_ptr=sinks,
473 + block_tables_ptr=block_table,
474 + seq_lens_ptr=seqused_k,
475 + alibi_slopes_ptr=alibi_slopes,
476 + scale=softmax_scale,
477 + softcap=softcap,
478 + num_query_heads=num_query_heads,
479 + num_queries_per_kv=num_queries_per_kv,
480 + block_table_stride=block_table.stride(0),
481 + query_stride_0=q.stride(0),
482 + query_stride_1=q.stride(1),
483 + output_stride_0=out.stride(0),
484 + output_stride_1=out.stride(1),
485 + BLOCK_SIZE=block_size,
486 + TILE_SIZE=tile_size,
487 + HEAD_SIZE_QK=head_size_qk,
488 + HEAD_SIZE_QK_PADDED=triton.next_power_of_2(head_size_qk),
489 + HEAD_SIZE_V=head_size_v,
490 + HEAD_SIZE_V_PADDED=triton.next_power_of_2(head_size_v),
491 + USE_ALIBI_SLOPES=use_alibi_slopes,
492 + USE_ALIBI_SQRT=use_alibi_sqrt,
493 + USE_SOFTCAP=(softcap > 0),
494 + USE_SINKS=(sinks is not None),
495 + SLIDING_WINDOW=sliding_window_val,
496 + stride_k_cache_0=k.stride(0),
497 + stride_k_cache_1=k.stride(1),
498 + stride_k_cache_2=k.stride(2),
499 + stride_k_cache_3=k.stride(3),
500 + stride_v_cache_0=v.stride(0),
501 + stride_v_cache_1=v.stride(1),
502 + stride_v_cache_2=v.stride(2),
503 + stride_v_cache_3=v.stride(3),
504 + query_start_len_ptr=cu_seqlens_q,
505 + BLOCK_Q=BLOCK_Q,
506 + num_seqs=num_seqs,
507 + BLOCK_M=BLOCK_M,
508 + NUM_SEGMENTS_PER_SEQ=num_segments,
509 + IS_3D=use_3d,
510 + )
511 +
512 + if use_3d:
513 + kernel_reduce_segments_diffkv[(q.shape[0], num_query_heads)](
514 + output_ptr=out,
515 + segm_output_ptr=softmax_segm_output,
516 + segm_max_ptr=softmax_segm_max,
517 + segm_expsum_ptr=softmax_segm_expsum,
518 + seq_lens_ptr=seqused_k,
519 + num_seqs=num_seqs,
520 + num_query_heads=num_query_heads,
521 + output_stride_0=out.stride(0),
522 + output_stride_1=out.stride(1),
523 + TILE_SIZE=tile_size,
524 + HEAD_SIZE_V=head_size_v,
525 + HEAD_SIZE_V_PADDED=triton.next_power_of_2(head_size_v),
526 + query_start_len_ptr=cu_seqlens_q,
527 + BLOCK_Q=BLOCK_Q,
528 + NUM_SEGMENTS_PER_SEQ=num_par_softmax_segments,
529 + )