@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 147 lines 5.39 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Unit test for ``_memcpy_u64_tiled``.
4
5 Guards the head/body/tail alignment handling and the u64 body's tile
6 partitioning used by ``postprocess_mamba_fused_kernel`` and
7 ``precopy_mamba_align_fused_kernel``. The device function must be
8 byte-identical to a ``dst[dst_off:dst_off+copy_size] = src[src_off:src_off+
9 copy_size]`` slice copy for every combination of:
10
11 * ``copy_size``: the degenerate 0-byte case, sub-8B (head-only), 8B-aligned
12 bodies, and a multi-tile body that spans several ``COPY_BLOCK_SIZE``
13 iterations.
14 * ``(src_off, dst_off)`` pairs: aligned, shared sub-8B misalignment (fast
15 path exercising head/tail masking), and mismatched src/dst alignment
16 (byte-fallback path). The kernel branches on ``(src ^ dst) & 7``, so a
17 full 5x5 product would re-exercise the same two codepaths.
18 * ``NUM_TILES``: 1 (SD conv callsite, single-CTA memcpy) and
19 ``_TEMPORAL_TILES`` (temporal callsite, u64 range partitioned across
20 CTAs).
21 * ``COPY_BLOCK_SIZE``: the production value 1024 (single-tile at these
22 sizes) and a small value 8 so the existing ``copy_size`` cases actually
23 cross tile boundaries under ``NUM_TILES=_TEMPORAL_TILES``. The
24 partitioning math is COPY_BLOCK_SIZE-agnostic, so testing at a small
25 value covers boundary arithmetic for all values.
26 """
27
28 from __future__ import annotations
29
30 import torch
31
32 from vllm.platforms import current_platform
33 from vllm.triton_utils import tl, triton
34 from vllm.v1.worker.mamba_utils import _TEMPORAL_TILES, _memcpy_u64_tiled
35
36 try:
37 import pytest
38
39 pytestmark = pytest.mark.skipif(
40 not current_platform.is_cuda(),
41 reason="_memcpy_u64_tiled needs CUDA/Triton",
42 )
43 _parametrize = pytest.mark.parametrize
44 except ModuleNotFoundError: # allow running directly as ``python <thisfile>``
45 pytest = None
46
47 def _parametrize(_name, _values):
48 def _deco(fn):
49 return fn
50
51 return _deco
52
53
54 @triton.jit
55 def _memcpy_wrapper_kernel(
56 src_ptr,
57 dst_ptr,
58 copy_size,
59 src_off,
60 dst_off,
61 COPY_BLOCK_SIZE: tl.constexpr,
62 NUM_TILES: tl.constexpr,
63 ):
64 """Thin harness: one CTA per tile, exercising the device function."""
65 tile_idx = tl.program_id(0)
66 src_addr = src_ptr.to(tl.int64) + src_off
67 dst_addr = dst_ptr.to(tl.int64) + dst_off
68 _memcpy_u64_tiled(
69 src_addr,
70 dst_addr,
71 copy_size,
72 tile_idx,
73 COPY_BLOCK_SIZE=COPY_BLOCK_SIZE,
74 NUM_TILES=NUM_TILES,
75 )
76
77
78 # Copy sizes: 0/1/7 (all-head), 8 (all-body when dst-aligned), 15 (body+tail,
79 # tail-mask path), 16 (small tiled body), 1024 and 4 KiB (spans all 16 tiles
80 # at COPY_BLOCK_SIZE=8, with 4 KiB adding multi-iteration-per-tile coverage;
81 # both single-tile at COPY_BLOCK_SIZE=1024).
82 _COPY_SIZES = [0, 1, 7, 8, 15, 16, 1024, 4 * 1024]
83 # (src_off, dst_off) pairs. Torch tensors are 256B-aligned at data_ptr, so
84 # slicing by these bytes yields a controlled sub-8B alignment. The kernel
85 # branches on ``(src ^ dst) & 7``, so we cover both sides plus the aligned
86 # baseline; the full 5x5 product added no coverage.
87 _ALIGN_PAIRS = [
88 (0, 0), # fully aligned; head/tail masked out
89 (3, 3), # shared misalignment; fast path with head_bytes=5
90 (7, 7), # shared, extreme; head_bytes=1 with non-empty body
91 (1, 3), # mismatched; byte load/store fallback
92 (3, 1), # mismatched, opposite direction
93 ]
94 _MAX_ALIGN_OFF = max(o for pair in _ALIGN_PAIRS for o in pair)
95 # (NUM_TILES, COPY_BLOCK_SIZE) configs. NUM_TILES=1 is the SD conv callsite
96 # (single-CTA memcpy) where tile partitioning collapses, so COPY_BLOCK_SIZE
97 # is not observable — one value suffices. NUM_TILES=_TEMPORAL_TILES is the
98 # temporal callsite: 1024 matches the production launch; 8 shrinks per-tile
99 # rounding so ``_COPY_SIZES``' 1024/4096 cases cross tile boundaries
100 # end-to-end (partitioning math is COPY_BLOCK_SIZE-agnostic, so a small
101 # value proves it for all values).
102 _TILE_BLOCK_CONFIGS = [
103 (1, 1024),
104 (_TEMPORAL_TILES, 1024),
105 (_TEMPORAL_TILES, 8),
106 ]
107
108
109 @_parametrize("copy_size", _COPY_SIZES)
110 @_parametrize("src_off,dst_off", _ALIGN_PAIRS)
111 @_parametrize("num_tiles,copy_block_size", _TILE_BLOCK_CONFIGS)
112 def test_memcpy_u64_tiled_matches_slice_copy(
113 copy_size, src_off, dst_off, num_tiles, copy_block_size
114 ):
115 device = torch.device("cuda")
116 torch.manual_seed(0)
117
118 # Pad by the max possible offset so the reference slice is in-bounds.
119 slack = _MAX_ALIGN_OFF + 8
120 src = torch.randint(0, 256, (copy_size + slack,), dtype=torch.uint8, device=device)
121 # Start dst from a distinct random pattern so the "unchanged region"
122 # check catches any accidental out-of-range writes.
123 dst = torch.randint(0, 256, (copy_size + slack,), dtype=torch.uint8, device=device)
124 ref = dst.clone()
125 ref[dst_off : dst_off + copy_size] = src[src_off : src_off + copy_size]
126
127 _memcpy_wrapper_kernel[(num_tiles,)](
128 src,
129 dst,
130 copy_size,
131 src_off,
132 dst_off,
133 COPY_BLOCK_SIZE=copy_block_size,
134 NUM_TILES=num_tiles,
135 )
136 torch.accelerator.synchronize()
137
138 # Byte-identical: the memcpy is dtype-agnostic; any diff is a bug.
139 torch.testing.assert_close(dst, ref, rtol=0, atol=0)
140
141
142 if __name__ == "__main__":
143 for cs in _COPY_SIZES:
144 for so, do in _ALIGN_PAIRS:
145 for nt, cbs in _TILE_BLOCK_CONFIGS:
146 test_memcpy_u64_tiled_matches_slice_copy(cs, so, do, nt, cbs)
147 print("ok")