diff --git a/csrc/libtorch_stable/quantization/vectorization_utils.cuh b/csrc/libtorch_stable/quantization/vectorization_utils.cuh index 98b491b7e23f..0cc89bf289d4 100644 --- a/csrc/libtorch_stable/quantization/vectorization_utils.cuh +++ b/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(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(out); + + // fast path when input and output are both fully aligned. The vector + // load/store below go through vec_n_t, 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]); diff --git a/tests/kernels/attention/test_cache.py b/tests/kernels/attention/test_cache.py index 9b022a042c81..4cbeb7a0b97c 100644 --- a/tests/kernels/attention/test_cache.py +++ b/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)