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 4, 08:45 UTC
[Bugfix] Fix degenerate KV cache stride causing TMA cudaErrorIllegalInstruction (#40737)
Signed-off-by: David Oy <david@baseten.co> Signed-off-by: David Oy <58150256+the-david-oy@users.noreply.github.com> Signed-off-by: David Oy <david.oy@baseten.co> Co-authored-by: David Oy <david@baseten.co> Co-authored-by: Claude <claude@anthropic.com> Co-authored-by: Vadim Gimpelson <156319763+vadiklyutiy@users.noreply.github.com>
David Oy committed
May 3, 2026 at 16:52 UTC
66dfee7121dfcbdfcce04ce92c117e5e10b25b14
6 files changed
+267
-22
tests/v1/attention/test_kv_head_stride_canonicalization.py
new
+162
@@ -0,0 +1,162 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Unit tests for canonicalize_singleton_dim_strides.
4
+
5
+Background
6
+----------
7
+When num_kv_heads_per_rank == 1 (e.g. Qwen3.5-397B with TP=8 → 1 KV head
8
+per rank), PyTorch's is_contiguous() returns True for *any* stride on the
9
+size-1 dimension. The KV cache allocator can therefore produce a tensor
10
+where that singleton dim has stride = 1 element (2 bytes for bf16) instead
11
+of the canonical product-of-remaining-dims value.
12
+
13
+CUDA TMA (used by FlashInfer XQA SM90 and Flash-Attention 3/4 on H100+)
14
+requires all non-outermost strides to be multiples of 16 bytes. A 2-byte
15
+stride triggers cudaErrorIllegalInstruction.
16
+
17
+canonicalize_singleton_dim_strides() patches degenerate strides on all
18
+size-1 dimensions via torch.as_strided — zero-copy.
19
+
20
+The degenerate stride manifests at different positions in different backends:
21
+- FlashInfer: stride(-3) after kv_cache.permute() → shape [..., 1, B, D]
22
+- FlashAttention: stride(-2) after kv_cache.unbind(0) → shape [N, B, 1, D]
23
+"""
24
+
25
+import torch
26
+
27
+from vllm.utils.torch_utils import canonicalize_singleton_dim_strides
28
+
29
+# ---------------------------------------------------------------------------
30
+# Helpers
31
+# ---------------------------------------------------------------------------
32
+
33
+
34
+def _inject_degenerate_stride(t: torch.Tensor, dim: int) -> torch.Tensor:
35
+ """Return a view of t with a degenerate (stride=1) on a size-1 dim."""
36
+ assert t.shape[dim] == 1, f"dim {dim} must have size 1"
37
+ strides = list(t.stride())
38
+ strides[dim] = 1 # inject the bug
39
+ return t.as_strided(t.shape, strides)
40
+
41
+
42
+# ---------------------------------------------------------------------------
43
+# Tests: canonicalize_singleton_dim_strides
44
+# ---------------------------------------------------------------------------
45
+
46
+
47
+class TestCanonicalizeSingletonDimStrides:
48
+ def test_flashinfer_layout_dim_neg3(self):
49
+ """FlashInfer path: degenerate stride at dim -3 (num_kv_heads)."""
50
+ # Shape after permute: [num_blocks, 2, num_kv_heads, block_size, head_size]
51
+ num_blocks, block_size, head_size = 64, 16, 128
52
+ t = torch.zeros(num_blocks, 2, 1, block_size, head_size, dtype=torch.bfloat16)
53
+ t_deg = _inject_degenerate_stride(t, dim=-3)
54
+
55
+ assert t_deg.stride(-3) == 1 # confirm degenerate
56
+ assert t_deg.is_contiguous() # PyTorch doesn't notice
57
+
58
+ fixed = canonicalize_singleton_dim_strides(t_deg)
59
+
60
+ assert fixed.stride(-3) == block_size * head_size # canonical = 2048
61
+ assert fixed.stride(-2) == head_size # inner dims unchanged
62
+ assert fixed.stride(-1) == 1
63
+
64
+ def test_flash_attn_layout_dim_neg2(self):
65
+ """FlashAttention path: degenerate stride at dim -2 (num_kv_heads)."""
66
+ # Shape after unbind(0): [num_blocks, block_size, num_kv_heads, head_size]
67
+ num_blocks, block_size, head_size = 64, 16, 128
68
+ t = torch.zeros(num_blocks, block_size, 1, head_size, dtype=torch.bfloat16)
69
+ t_deg = _inject_degenerate_stride(t, dim=-2)
70
+
71
+ assert t_deg.stride(-2) == 1
72
+ assert t_deg.is_contiguous()
73
+
74
+ fixed = canonicalize_singleton_dim_strides(t_deg)
75
+
76
+ assert fixed.stride(-2) == head_size # canonical = 128
77
+ assert fixed.stride(-1) == 1
78
+
79
+ def test_canonical_strides_returned_as_is(self):
80
+ """No degenerate strides → same object returned (no copy, no new view)."""
81
+ t = torch.zeros(64, 2, 1, 16, 128, dtype=torch.bfloat16)
82
+ result = canonicalize_singleton_dim_strides(t)
83
+ assert result is t
84
+
85
+ def test_multi_kv_heads_unchanged(self):
86
+ """num_kv_heads > 1 → strides are already canonical → unchanged."""
87
+ t = torch.zeros(16, 2, 4, 16, 128, dtype=torch.bfloat16)
88
+ original_strides = t.stride()
89
+ result = canonicalize_singleton_dim_strides(t)
90
+ assert result.stride() == original_strides
91
+
92
+ def test_data_pointer_preserved(self):
93
+ """Fix is zero-copy: same underlying storage."""
94
+ t = torch.zeros(8, 2, 1, 16, 128, dtype=torch.bfloat16)
95
+ t_deg = _inject_degenerate_stride(t, dim=-3)
96
+ fixed = canonicalize_singleton_dim_strides(t_deg)
97
+ assert fixed.data_ptr() == t_deg.data_ptr()
98
+ assert fixed.storage_offset() == t_deg.storage_offset()
99
+
100
+ def test_multiple_singleton_dims(self):
101
+ """All size-1 dims with degenerate strides are fixed."""
102
+ # Shape: [1, 1, 8, 32] — two size-1 dims
103
+ t = torch.zeros(1, 1, 8, 32, dtype=torch.float16)
104
+ # Both size-1 dims get degenerate strides
105
+ t_deg = t.as_strided(t.shape, (1, 1, 32, 1)) # both leading dims = 1
106
+
107
+ fixed = canonicalize_singleton_dim_strides(t_deg)
108
+
109
+ assert fixed.stride(0) == 1 * 8 * 32 # canonical: 256
110
+ assert fixed.stride(1) == 1 * 8 * 32 # canonical: 256 (same since size-1)
111
+ assert fixed.stride(2) == 32
112
+ assert fixed.stride(3) == 1
113
+
114
+ def test_various_shapes_flashinfer(self):
115
+ """Correctness across different block_size / head_size for FlashInfer layout."""
116
+ for block_size, head_size in [(16, 64), (16, 128), (32, 128), (16, 256)]:
117
+ t = torch.zeros(8, 2, 1, block_size, head_size, dtype=torch.bfloat16)
118
+ t_deg = _inject_degenerate_stride(t, dim=-3)
119
+ fixed = canonicalize_singleton_dim_strides(t_deg)
120
+ assert fixed.stride(-3) == block_size * head_size, (
121
+ f"Failed for block_size={block_size}, head_size={head_size}: "
122
+ f"got stride(-3)={fixed.stride(-3)}"
123
+ )
124
+
125
+ def test_various_shapes_flash_attn(self):
126
+ """Correctness across different shapes for FlashAttention layout."""
127
+ for block_size, head_size in [(16, 64), (16, 128), (32, 128)]:
128
+ t = torch.zeros(8, block_size, 1, head_size, dtype=torch.bfloat16)
129
+ t_deg = _inject_degenerate_stride(t, dim=-2)
130
+ fixed = canonicalize_singleton_dim_strides(t_deg)
131
+ assert fixed.stride(-2) == head_size, (
132
+ f"Failed for block_size={block_size}, head_size={head_size}: "
133
+ f"got stride(-2)={fixed.stride(-2)}"
134
+ )
135
+
136
+ def test_tma_alignment_satisfied_after_fix_bf16(self):
137
+ """After fix, all strides meet 16-byte TMA alignment for bf16."""
138
+ t = torch.zeros(64, 2, 1, 16, 128, dtype=torch.bfloat16)
139
+ t_deg = _inject_degenerate_stride(t, dim=-3)
140
+ fixed = canonicalize_singleton_dim_strides(t_deg)
141
+
142
+ element_size = fixed.element_size() # 2 bytes for bf16
143
+ for i, s in enumerate(fixed.stride()):
144
+ assert (s * element_size) % 16 == 0 or i == len(fixed.stride()) - 1, (
145
+ f"dim {i} stride {s} * {element_size} bytes not 16-byte aligned"
146
+ )
147
+
148
+ def test_non_contiguous_outer_dims_preserved(self):
149
+ """Outer (non-size-1) non-contiguous strides are left unchanged."""
150
+ # Simulate cross-layer unified allocation: num_blocks stride is non-canonical
151
+ # but the inner dims should be fixed.
152
+ base = torch.zeros(200, 2, 1, 16, 128, dtype=torch.bfloat16)
153
+ # Slice every 2nd block → non-canonical outer stride
154
+ t_sliced = base[::2] # shape [100, 2, 1, 16, 128], stride[0] = 2*canonical
155
+ t_deg = _inject_degenerate_stride(t_sliced, dim=-3)
156
+
157
+ fixed = canonicalize_singleton_dim_strides(t_deg)
158
+
159
+ # Outer stride should be unchanged (not a size-1 dim)
160
+ assert fixed.stride(0) == t_sliced.stride(0)
161
+ # Inner degenerate stride should be fixed
162
+ assert fixed.stride(-3) == 16 * 128
vllm/utils/cpu_resource_utils.py
+1
-1
@@ -125,7 +125,7 @@ def get_allowed_cpu_list() -> list[LogicalCPUInfo]:
125
if platform.system() == "Darwin":
126
return cpu_list
127
128
- global_allowed_cpu_id_list = os.sched_getaffinity(0)
128
+ global_allowed_cpu_id_list = os.sched_getaffinity(0) # type: ignore[attr-defined]
129
logical_cpu_list = [x for x in cpu_list if x.id in global_allowed_cpu_id_list]
130
131
return logical_cpu_list
vllm/utils/torch_utils.py
+26
@@ -110,6 +110,32 @@ def is_strictly_contiguous(t: torch.Tensor) -> bool:
110
return True
111
112
113
+def canonicalize_singleton_dim_strides(t: torch.Tensor) -> torch.Tensor:
114
+ """Fix degenerate strides on size=1 dimensions for CUDA TMA compatibility.
115
+
116
+ PyTorch allows any stride on a size=1 dim (is_contiguous() is always True
117
+ there), so a size=1 dim may have stride=1 (2 bytes for bf16) instead of
118
+ the canonical product(shape[i+1:]). CUDA TMA on H100+ requires all
119
+ non-outermost strides to be ≥16-byte aligned; stride=1 triggers
120
+ cudaErrorIllegalInstruction. Zero-copy: patches stride metadata only via
121
+ as_strided; returns t unchanged if all size=1 strides are already canonical.
122
+ """
123
+ if 1 not in t.shape:
124
+ return t
125
+ strides = list(t.stride())
126
+ shape = t.shape
127
+ prev_stride = 1
128
+ changed = False
129
+ for i in range(len(shape) - 1, -1, -1):
130
+ if shape[i] == 1 and strides[i] != prev_stride:
131
+ strides[i] = prev_stride
132
+ changed = True
133
+ prev_stride = strides[i] * shape[i]
134
+ if not changed:
135
+ return t
136
+ return t.as_strided(t.shape, strides)
137
+
138
+
139
@contextlib.contextmanager
140
def set_default_torch_dtype(dtype: torch.dtype):
141
"""Sets the default torch dtype to the given dtype."""
vllm/v1/attention/backends/flash_attn.py
+23
-1
@@ -11,7 +11,10 @@ import torch
11
12
from vllm.model_executor.layers.attention import Attention
13
from vllm.platforms import current_platform
14
-from vllm.utils.torch_utils import is_quantized_kv_cache
14
+from vllm.utils.torch_utils import (
15
+ canonicalize_singleton_dim_strides,
16
+ is_quantized_kv_cache,
17
+)
18
from vllm.v1.attention.backend import (
19
AttentionBackend,
20
AttentionImpl,
@@ -747,6 +750,23 @@ class FlashAttentionImpl(AttentionImpl):
750
751
# For decoder and cross-attention, use KV cache as before
752
key_cache, value_cache = kv_cache.unbind(0)
753
+ # Fix degenerate strides on size-1 dims (e.g. num_kv_heads=1 with TP).
754
+ # FA3/4 on H100+ uses TMA, which requires ≥16-byte stride alignment.
755
+ # See vllm.utils.torch_utils.canonicalize_singleton_dim_strides.
756
+ fixed_k = canonicalize_singleton_dim_strides(key_cache)
757
+ fixed_v = canonicalize_singleton_dim_strides(value_cache)
758
+ if fixed_k is not key_cache or fixed_v is not value_cache:
759
+ logger.debug(
760
+ "Canonicalized degenerate KV cache strides (FlashAttention): "
761
+ "shape=%s, key strides before=%s after=%s, "
762
+ "value strides before=%s after=%s",
763
+ key_cache.shape,
764
+ key_cache.stride(),
765
+ fixed_k.stride(),
766
+ value_cache.stride(),
767
+ fixed_v.stride(),
768
+ )
769
+ key_cache, value_cache = fixed_k, fixed_v
770
771
if is_quantized_kv_cache(self.kv_cache_dtype):
772
# queries are quantized in the attention layer
@@ -861,6 +881,8 @@ class FlashAttentionImpl(AttentionImpl):
881
# we use direct Q, K, V tensors without caching
882
return
883
884
+ # Scatter write into the KV cache using slot_mapping indices.
885
+ # No TMA kernel is invoked here, so stride canonicalization is not needed.
886
key_cache, value_cache = kv_cache.unbind(0)
887
888
# Reshape the input keys and values and store them in the cache.
vllm/v1/attention/backends/flash_attn_diffkv.py
+24
-1
@@ -4,7 +4,11 @@
4
5
import torch
6
7
-from vllm.utils.torch_utils import is_quantized_kv_cache
7
+from vllm.logger import init_logger
8
+from vllm.utils.torch_utils import (
9
+ canonicalize_singleton_dim_strides,
10
+ is_quantized_kv_cache,
11
+)
12
from vllm.v1.attention.backend import AttentionType
13
from vllm.v1.attention.backends.fa_utils import (
14
get_flash_attn_version,
@@ -25,6 +29,8 @@ from .flash_attn import (
29
cascade_attention,
30
)
31
32
+logger = init_logger(__name__)
33
+
34
35
class FlashAttentionDiffKVBackend(FlashAttentionBackend):
36
# Default to 128 for this backend
@@ -204,6 +210,23 @@ class FlashAttentionDiffKVImpl(FlashAttentionImpl):
210
# Different head_size for K and V
211
key_cache = kv_cache[..., : self.head_size]
212
value_cache = kv_cache[..., self.head_size :]
213
+ # Fix degenerate strides on size-1 dims (e.g. num_kv_heads=1 with TP).
214
+ # FA3/4 on H100+ uses TMA, which requires ≥16-byte stride alignment.
215
+ # See vllm.utils.torch_utils.canonicalize_singleton_dim_strides.
216
+ fixed_k = canonicalize_singleton_dim_strides(key_cache)
217
+ fixed_v = canonicalize_singleton_dim_strides(value_cache)
218
+ if fixed_k is not key_cache or fixed_v is not value_cache:
219
+ logger.debug(
220
+ "Canonicalized degenerate KV cache strides (FlashAttentionDiffKV): "
221
+ "shape=%s, key strides before=%s after=%s, "
222
+ "value strides before=%s after=%s",
223
+ key_cache.shape,
224
+ key_cache.stride(),
225
+ fixed_k.stride(),
226
+ value_cache.stride(),
227
+ fixed_v.stride(),
228
+ )
229
+ key_cache, value_cache = fixed_k, fixed_v
230
231
if is_quantized_kv_cache(self.kv_cache_dtype):
232
# queries are quantized in the attention layer
vllm/v1/attention/backends/flashinfer.py
+31
-19
@@ -43,6 +43,7 @@ from vllm.utils.flashinfer import (
43
from vllm.utils.math_utils import cdiv
44
from vllm.utils.platform_utils import is_pin_memory_available
45
from vllm.utils.torch_utils import (
46
+ canonicalize_singleton_dim_strides,
47
is_quantized_kv_cache,
48
is_strictly_contiguous,
49
nvfp4_kv_cache_full_dim,
@@ -1479,6 +1480,21 @@ class FlashInferImpl(AttentionImpl):
1480
1481
stride_order = FlashInferBackend.get_kv_cache_stride_order()
1482
kv_cache_permute = kv_cache.permute(*stride_order) # HND and contiguous
1483
+ # Fix degenerate strides on any size-1 dimension (e.g. num_kv_heads=1
1484
+ # with TP=8). PyTorch permits non-canonical strides on size-1 dims;
1485
+ # CUDA TMA requires ≥16-byte alignment on all non-outermost strides.
1486
+ # canonicalize_singleton_dim_strides patches metadata via as_strided —
1487
+ # zero-copy. See vllm.utils.torch_utils.
1488
+ fixed = canonicalize_singleton_dim_strides(kv_cache_permute)
1489
+ if fixed is not kv_cache_permute:
1490
+ logger.debug(
1491
+ "Canonicalized degenerate KV cache strides (FlashInfer): "
1492
+ "shape=%s, strides before=%s, strides after=%s",
1493
+ kv_cache_permute.shape,
1494
+ kv_cache_permute.stride(),
1495
+ fixed.stride(),
1496
+ )
1497
+ kv_cache_permute = fixed
1498
1499
# For NVFP4, the kv_cache last dim is full_dim (data + scale packed).
1500
# Split into correctly-strided data and scale views.
@@ -1568,10 +1584,11 @@ class FlashInferImpl(AttentionImpl):
1584
else:
1585
assert isinstance(attn_metadata.prefill, TRTLLMPrefill)
1586
# prefill_query may be non-contiguous or have degenerate strides
1571
- # First ensure memory contiguity, then fix degenerate strides
1572
- # with reshape. contiguous() alone doesn't fix degenerate
1573
- # strides when a dimension has size 1.
1574
- prefill_query = prefill_query.contiguous().reshape(prefill_query.shape)
1587
+ # on size=1 dims. contiguous() ensures memory layout; then
1588
+ # canonicalize_singleton_dim_strides fixes any remaining
1589
+ # degenerate strides on size=1 dims for TMA alignment.
1590
+ prefill_query = prefill_query.contiguous()
1591
+ prefill_query = canonicalize_singleton_dim_strides(prefill_query)
1592
workspace_buffer = _get_trtllm_gen_workspace_buffer()
1593
block_tables_prefill = attn_metadata.prefill.block_tables
1594
seq_lens_prefill = attn_metadata.prefill.seq_lens
@@ -1621,11 +1638,9 @@ class FlashInferImpl(AttentionImpl):
1638
# with fp8 kv cache, we can construct a mock block
1639
# and mock kv cache with BF16 KV involved in the prefill
1640
#
1624
- # The inner (block_size, head_size) dims must be
1625
- # contiguous; outer dims may have non-canonical strides
1626
- # (e.g. cross-layer unified allocation).
1627
- # Degenerate strides on outer dims break TMA descriptors
1628
- # (see flashinfer-ai/flashinfer#2232).
1641
+ kv_cache_permute = canonicalize_singleton_dim_strides(
1642
+ kv_cache_permute
1643
+ )
1644
kv_strides = kv_cache_permute.stride()
1645
assert (
1646
kv_strides[-1] == 1
@@ -1732,12 +1747,13 @@ class FlashInferImpl(AttentionImpl):
1747
if needs_fp8_out:
1748
output[:num_decode_tokens].copy_(out_decode.to(output.dtype))
1749
else:
1735
- # decode_query may be non-contiguous or have degenerate strides
1750
assert isinstance(attn_metadata.decode, TRTLLMDecode)
1737
- # First ensure memory contiguity, then fix degenerate strides
1738
- # with reshape. contiguous() alone doesn't fix degenerate
1739
- # strides when a dimension has size 1.
1740
- decode_query = decode_query.contiguous().reshape(decode_query.shape)
1751
+ # decode_query may be non-contiguous or have degenerate strides
1752
+ # on size=1 dims. contiguous() ensures memory layout; then
1753
+ # canonicalize_singleton_dim_strides fixes any remaining
1754
+ # degenerate strides on size=1 dims for TMA alignment.
1755
+ decode_query = decode_query.contiguous()
1756
+ decode_query = canonicalize_singleton_dim_strides(decode_query)
1757
workspace_buffer = _get_trtllm_gen_workspace_buffer()
1758
block_tables_decode = attn_metadata.decode.block_tables
1759
seq_lens_decode = attn_metadata.decode.seq_lens
@@ -1748,11 +1764,7 @@ class FlashInferImpl(AttentionImpl):
1764
assert is_strictly_contiguous(workspace_buffer)
1765
assert is_strictly_contiguous(block_tables_decode)
1766
assert is_strictly_contiguous(seq_lens_decode)
1751
- # kv_cache outer dims may be non-contiguous (e.g.
1752
- # cross-layer unified allocation), but inner dims
1753
- # (block_size, head_size) must be contiguous and
1754
- # strides must be canonical to avoid TMA descriptor
1755
- # failures (see flashinfer-ai/flashinfer#2232).
1767
+ kv_cache_permute = canonicalize_singleton_dim_strides(kv_cache_permute)
1768
kv_strides = kv_cache_permute.stride()
1769
assert (
1770
kv_strides[-1] == 1 and kv_strides[-2] == kv_cache_permute.shape[-1]