[Bugfix][Kernel] Check output alignment in vectorize_with_alignment (fixes misaligned-address crash for non-multiple-of-8 head sizes) (#45466)
Signed-off-by: HumphreySun98 <humphreysun98@gmail.com>
Humphrey committed
Jun 18, 2026 at 11:58 UTC
4583630b562124c033551e5630a7ab3d607a6f03
2 files changed
+61
-6
csrc/libtorch_stable/quantization/vectorization_utils.cuh
+24
-6
@@ -24,13 +24,21 @@ __device__ inline void vectorize_with_alignment(
24
ScaOp&& scalar_op) { // InT -> OutT
25
static_assert(VEC_SIZE > 0 && (VEC_SIZE & (VEC_SIZE - 1)) == 0,
26
"VEC_SIZE must be a positive power-of-two");
27
- constexpr int WIDTH = VEC_SIZE * sizeof(InT); // eg: 64 B
27
+ constexpr int WIDTH = VEC_SIZE * sizeof(InT); // eg: 16 B
28
+ constexpr int OUT_WIDTH = VEC_SIZE * sizeof(OutT); // eg: 16 B
29
uintptr_t addr = reinterpret_cast<uintptr_t>(in);
29
-
30
- // fast path when the whole region is already aligned
31
- // Note: currently the output is guaranteed to be same as the input, so we
32
- // don't check it here, comments here just for future reference.
33
- bool can_vec = ((addr & (WIDTH - 1)) == 0) && ((len & (VEC_SIZE - 1)) == 0);
30
+ uintptr_t out_addr = reinterpret_cast<uintptr_t>(out);
31
+
32
+ // fast path when input and output are both fully aligned. The vector
33
+ // load/store below go through vec_n_t<T, VEC_SIZE>, declared
34
+ // __align__(VEC_SIZE * sizeof(T)), so each side must be aligned to its
35
+ // own vector width. out is NOT generally co-aligned with in: e.g.
36
+ // reshape_and_cache_flash writes KV-cache rows whose byte offset is a
37
+ // multiple of head_size, which for head sizes that are not a multiple
38
+ // of VEC_SIZE puts some rows off the vector-width boundary.
39
+ bool can_vec = ((addr & (WIDTH - 1)) == 0) &&
40
+ ((out_addr & (OUT_WIDTH - 1)) == 0) &&
41
+ ((len & (VEC_SIZE - 1)) == 0);
42
if (can_vec) {
43
int num_vec = len / VEC_SIZE;
44
@@ -55,6 +63,16 @@ __device__ inline void vectorize_with_alignment(
63
prefix_elems /= sizeof(InT);
64
prefix_elems = min(prefix_elems, len); // 0 ≤ prefix < 16
65
66
+ // the prefix below aligns in; if that does not also align out (their
67
+ // addresses differ modulo the vector width), vectorizing is impossible
68
+ // and the whole copy must stay scalar.
69
+ if (((out_addr + prefix_elems * sizeof(OutT)) & (OUT_WIDTH - 1)) != 0) {
70
+ for (int i = tid; i < len; i += stride) {
71
+ scalar_op(out[i], in[i]);
72
+ }
73
+ return;
74
+ }
75
+
76
// 1. prefill the when it is unsafe to vectorize
77
for (int i = tid; i < prefix_elems; i += stride) {
78
scalar_op(out[i], in[i]);
tests/kernels/attention/test_cache.py
+37
@@ -428,6 +428,43 @@ def test_reshape_and_cache_flash(
428
torch.testing.assert_close(value_cache_compact, cloned_value_cache)
429
430
431
+@pytest.mark.parametrize("dtype", DTYPES)
432
+@pytest.mark.parametrize("kv_cache_dtype", KV_CACHE_DTYPE)
433
+@pytest.mark.parametrize("kv_cache_layout", CACHE_LAYOUTS)
434
+@pytest.mark.parametrize("implementation", RESHAPE_FLASH_IMPLEMENTATIONS)
435
+@torch.inference_mode()
436
+def test_reshape_and_cache_flash_unaligned_rows(
437
+ kv_cache_factory_flashinfer,
438
+ dtype: torch.dtype,
439
+ kv_cache_dtype: str,
440
+ kv_cache_layout: str,
441
+ implementation: str,
442
+) -> None:
443
+ """Regression test for https://github.com/vllm-project/vllm/issues/41257.
444
+
445
+ head_size=46 with num_heads=13 places KV-cache rows at byte offsets
446
+ that are not a multiple of the vector width (NHD row pitch
447
+ 13*46*itemsize, HND head pitch 46*itemsize), unlike HEAD_SIZES above
448
+ which are all 16-byte multiples. The CUDA kernel used to issue
449
+ vectorized stores to those rows -> CUDA misaligned address.
450
+ """
451
+ test_reshape_and_cache_flash(
452
+ kv_cache_factory_flashinfer,
453
+ num_tokens=42,
454
+ num_heads=13,
455
+ head_size=46,
456
+ block_size=16,
457
+ num_blocks=128,
458
+ dtype=dtype,
459
+ seed=0,
460
+ device=CUDA_DEVICES[0],
461
+ kv_cache_dtype=kv_cache_dtype,
462
+ kv_cache_layout=kv_cache_layout,
463
+ kv_scale_type="tensor",
464
+ implementation=implementation,
465
+ )
466
+
467
+
468
@pytest.mark.parametrize("direction", COPYING_DIRECTION)
469
@pytest.mark.parametrize("num_mappings", NUM_MAPPINGS)
470
@pytest.mark.parametrize("num_heads", NUM_HEADS)