[Bugfix][Kernel] Check output alignment in vectorize_with_alignment (fixes misaligned-address crash for non-multiple-of-8 head sizes)

csrc/libtorch_stable/quantization/vectorization_utils.cuh

@@ -24,13 +24,21 @@ __device__ inline void vectorize_with_alignment(
     ScaOp&& scalar_op) {  // InT -> OutT
   static_assert(VEC_SIZE > 0 && (VEC_SIZE & (VEC_SIZE - 1)) == 0,
                 "VEC_SIZE must be a positive power-of-two");
-  constexpr int WIDTH = VEC_SIZE * sizeof(InT);  // eg: 64 B
+  constexpr int WIDTH = VEC_SIZE * sizeof(InT);       // eg: 16 B
+  constexpr int OUT_WIDTH = VEC_SIZE * sizeof(OutT);  // eg: 16 B
   uintptr_t addr = reinterpret_cast<uintptr_t>(in);
-
-  // fast path when the whole region is already aligned
-  // Note: currently the output is guaranteed to be same as the input, so we
-  // don't check it here, comments here just for future reference.
-  bool can_vec = ((addr & (WIDTH - 1)) == 0) && ((len & (VEC_SIZE - 1)) == 0);
+  uintptr_t out_addr = reinterpret_cast<uintptr_t>(out);
+
+  // fast path when input and output are both fully aligned. The vector
+  // load/store below go through vec_n_t<T, VEC_SIZE>, declared
+  // __align__(VEC_SIZE * sizeof(T)), so each side must be aligned to its
+  // own vector width. out is NOT generally co-aligned with in: e.g.
+  // reshape_and_cache_flash writes KV-cache rows whose byte offset is a
+  // multiple of head_size, which for head sizes that are not a multiple
+  // of VEC_SIZE puts some rows off the vector-width boundary.
+  bool can_vec = ((addr & (WIDTH - 1)) == 0) &&
+                 ((out_addr & (OUT_WIDTH - 1)) == 0) &&
+                 ((len & (VEC_SIZE - 1)) == 0);
   if (can_vec) {
     int num_vec = len / VEC_SIZE;
 
@@ -55,6 +63,16 @@ __device__ inline void vectorize_with_alignment(
   prefix_elems /= sizeof(InT);
   prefix_elems = min(prefix_elems, len);  // 0 ≤ prefix < 16
 
+  // the prefix below aligns in; if that does not also align out (their
+  // addresses differ modulo the vector width), vectorizing is impossible
+  // and the whole copy must stay scalar.
+  if (((out_addr + prefix_elems * sizeof(OutT)) & (OUT_WIDTH - 1)) != 0) {
+    for (int i = tid; i < len; i += stride) {
+      scalar_op(out[i], in[i]);
+    }
+    return;
+  }
+
   // 1. prefill the when it is unsafe to vectorize
   for (int i = tid; i < prefix_elems; i += stride) {
     scalar_op(out[i], in[i]);

tests/kernels/attention/test_cache.py

@@ -428,6 +428,43 @@ def dequant_nvfp4_cache_nhd(data_cache, scale_cache, global_scale):
         torch.testing.assert_close(value_cache_compact, cloned_value_cache)
 
 
+@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize("kv_cache_dtype", KV_CACHE_DTYPE)
+@pytest.mark.parametrize("kv_cache_layout", CACHE_LAYOUTS)
+@pytest.mark.parametrize("implementation", RESHAPE_FLASH_IMPLEMENTATIONS)
+@torch.inference_mode()
+def test_reshape_and_cache_flash_unaligned_rows(
+    kv_cache_factory_flashinfer,
+    dtype: torch.dtype,
+    kv_cache_dtype: str,
+    kv_cache_layout: str,
+    implementation: str,
+) -> None:
+    """Regression test for https://github.com/vllm-project/vllm/issues/41257.
+
+    head_size=46 with num_heads=13 places KV-cache rows at byte offsets
+    that are not a multiple of the vector width (NHD row pitch
+    13*46*itemsize, HND head pitch 46*itemsize), unlike HEAD_SIZES above
+    which are all 16-byte multiples. The CUDA kernel used to issue
+    vectorized stores to those rows -> CUDA misaligned address.
+    """
+    test_reshape_and_cache_flash(
+        kv_cache_factory_flashinfer,
+        num_tokens=42,
+        num_heads=13,
+        head_size=46,
+        block_size=16,
+        num_blocks=128,
+        dtype=dtype,
+        seed=0,
+        device=CUDA_DEVICES[0],
+        kv_cache_dtype=kv_cache_dtype,
+        kv_cache_layout=kv_cache_layout,
+        kv_scale_type="tensor",
+        implementation=implementation,
+    )
+
+
 @pytest.mark.parametrize("direction", COPYING_DIRECTION)
 @pytest.mark.parametrize("num_mappings", NUM_MAPPINGS)
 @pytest.mark.parametrize("num_heads", NUM_HEADS)