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 10, 09:00 UTC
[MISC] Add strict contiguity check for FlashInfer attention tensors (#32008)
Signed-off-by: Vadim Gimpelson <vadim.gimpelson@gmail.com> Signed-off-by: Vadim Gimpelson <156319763+vadiklyutiy@users.noreply.github.com>
Vadim Gimpelson committed
Jan 11, 2026 at 00:40 UTC
e15a5ff07b89646090458ce6ae7908bcbae82b90
2 files changed
+41
-10
vllm/utils/torch_utils.py
+30
@@ -64,6 +64,36 @@ MODELOPT_TO_VLLM_KV_CACHE_DTYPE_MAP = {
64
T = TypeVar("T")
65
66
67
+def is_strictly_contiguous(t: torch.Tensor) -> bool:
68
+ """
69
+ Check if tensor is contiguous AND has no degenerate strides.
70
+
71
+ A degenerate stride occurs when a dimension has size 1 but the stride
72
+ doesn't match the canonical contiguous layout. This can cause issues
73
+ in some CUDA kernels that rely on stride values for memory access.
74
+
75
+ For a C-contiguous tensor of shape (d0, d1, ..., dn), the expected
76
+ strides are: stride[i] = product(shape[i+1:]) for all i, with stride[-1]=1.
77
+
78
+ Example with torch.Size([16, 1, 8, 32]):
79
+ - Canonical strides: (256, 256, 32, 1)
80
+ - Degenerate strides: (256, 1, 32, 1) # dim=1 has size=1, allowing
81
+ # non-canonical stride in dim=0
82
+ """
83
+ if not t.is_contiguous():
84
+ return False
85
+
86
+ # Check that strides match canonical contiguous layout
87
+ shape = t.shape
88
+ strides = t.stride()
89
+ expected_stride = 1
90
+ for i in range(len(shape) - 1, -1, -1):
91
+ if strides[i] != expected_stride:
92
+ return False
93
+ expected_stride *= shape[i]
94
+ return True
95
+
96
+
97
@contextlib.contextmanager
98
def set_default_torch_dtype(dtype: torch.dtype):
99
"""Sets the default torch dtype to the given dtype."""
vllm/v1/attention/backends/flashinfer.py
+11
-10
@@ -40,6 +40,7 @@ from vllm.utils.flashinfer import (
40
)
41
from vllm.utils.math_utils import cdiv
42
from vllm.utils.platform_utils import is_pin_memory_available
43
+from vllm.utils.torch_utils import is_strictly_contiguous
44
from vllm.v1.attention.backend import (
45
AttentionBackend,
46
AttentionImpl,
@@ -1392,11 +1393,11 @@ class FlashInferImpl(AttentionImpl):
1393
1394
# This path needs to be enabled with VLLM_KV_CACHE_LAYOUT = HND
1395
assert get_kv_cache_layout() == "HND"
1395
- assert prefill_query.is_contiguous()
1396
- assert kv_cache_permute.is_contiguous()
1397
- assert workspace_buffer.is_contiguous()
1398
- assert block_tables_prefill.is_contiguous()
1399
- assert seq_lens_prefill.is_contiguous()
1396
+ assert is_strictly_contiguous(prefill_query)
1397
+ assert is_strictly_contiguous(kv_cache_permute)
1398
+ assert is_strictly_contiguous(workspace_buffer)
1399
+ assert is_strictly_contiguous(block_tables_prefill)
1400
+ assert is_strictly_contiguous(seq_lens_prefill)
1401
1402
if output.dtype == FP4_DTYPE:
1403
assert self.o_sf_scale is not None
@@ -1503,11 +1504,11 @@ class FlashInferImpl(AttentionImpl):
1504
1505
# This path needs to be enabled with VLLM_KV_CACHE_LAYOUT = HND
1506
assert get_kv_cache_layout() == "HND"
1506
- assert decode_query.is_contiguous()
1507
- assert kv_cache_permute.is_contiguous()
1508
- assert workspace_buffer.is_contiguous()
1509
- assert block_tables_decode.is_contiguous()
1510
- assert seq_lens_decode.is_contiguous()
1507
+ assert is_strictly_contiguous(decode_query)
1508
+ assert is_strictly_contiguous(kv_cache_permute)
1509
+ assert is_strictly_contiguous(workspace_buffer)
1510
+ assert is_strictly_contiguous(block_tables_decode)
1511
+ assert is_strictly_contiguous(seq_lens_decode)
1512
1513
if output.dtype == FP4_DTYPE:
1514
assert self.o_sf_scale is not None