@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 162 lines 7.41 KB
Raw
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(1) → 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