feat(flydsl): support strided GDN decode state routing

aiter/ops/flydsl/kernels/gdr_decode.py

@@ -33,6 +33,9 @@ def create_vk_gdr_decode_kernel(
     num_v_heads: int,
     head_k_dim: int,
     head_v_dim: int,
+    q_strides: tuple,
+    k_strides: tuple,
+    v_strides: tuple,
     state_strides: tuple,
     a_strides: tuple,
     b_strides: tuple,
@@ -93,7 +96,8 @@ def gdr_decode_kernel(
         b: fx.Tensor,
         dt_bias: fx.Tensor,
         A_log: fx.Tensor,
-        indices: fx.Tensor,
+        read_indices: fx.Tensor,
+        write_indices: fx.Tensor,
         state: fx.Tensor,
         out: fx.Tensor,
         batch_size: fx.Int32,
@@ -128,17 +132,28 @@ def gdr_decode_kernel(
         warp_k_vec_start = w_tid % WARP_THREADS_K * VALUES_PER_THREAD_K
         global_v_start = tile_v_start + wid * WARP_TILE_V + w_tid // WARP_THREADS_K
 
-        indices_tensor = GTensor(indices, dtype=T.i32, shape=(-1,))
-        pool_idx = fx.Int32(indices_tensor[b_i])
+        read_indices_tensor = GTensor(read_indices, dtype=T.i32, shape=(-1,))
+        write_indices_tensor = GTensor(write_indices, dtype=T.i32, shape=(-1,))
+        read_pool_idx = fx.Int32(read_indices_tensor[b_i])
+        write_pool_idx = fx.Int32(write_indices_tensor[b_i])
 
         q_tensor = GTensor(
-            query, dtype=dtype_, shape=(-1, seq_length, num_k_heads, head_k_dim)
+            query,
+            dtype=dtype_,
+            shape=(-1, seq_length, num_k_heads, head_k_dim),
+            stride=q_strides,
         )
         k_tensor = GTensor(
-            key, dtype=dtype_, shape=(-1, seq_length, num_k_heads, head_k_dim)
+            key,
+            dtype=dtype_,
+            shape=(-1, seq_length, num_k_heads, head_k_dim),
+            stride=k_strides,
         )
         v_tensor = GTensor(
-            value, dtype=dtype_, shape=(-1, seq_length, num_v_heads, head_v_dim)
+            value,
+            dtype=dtype_,
+            shape=(-1, seq_length, num_v_heads, head_v_dim),
+            stride=v_strides,
         )
         a_tensor = GTensor(
             a,
@@ -157,12 +172,21 @@ def gdr_decode_kernel(
         out_tensor = GTensor(
             out, dtype=dtype_, shape=(-1, seq_length, num_v_heads, head_v_dim)
         )
-        state_tensor = GTensor(
+        read_state_tensor = GTensor(
+            state,
+            dtype=state_dtype_,
+            shape=(num_v_heads, head_v_dim, head_k_dim),
+            stride=(state_strides[1], state_strides[2], state_strides[3]),
+            static_bytes_offset_i64=fx.Int64(read_pool_idx)
+            * fx.Int64(state_strides[0])
+            * get_dtype_bytes(state_dtype),
+        )
+        write_state_tensor = GTensor(
             state,
             dtype=state_dtype_,
             shape=(num_v_heads, head_v_dim, head_k_dim),
             stride=(state_strides[1], state_strides[2], state_strides[3]),
-            static_bytes_offset_i64=fx.Int64(pool_idx)
+            static_bytes_offset_i64=fx.Int64(write_pool_idx)
             * fx.Int64(state_strides[0])
             * get_dtype_bytes(state_dtype),
         )
@@ -192,8 +216,10 @@ def _do_decode():
                 global_v_i = global_v_start + vi * WARP_GROUP_TILE_V
                 for ki in range_constexpr(WARP_TILE_K_ITERS):
                     warp_k_vec_i = warp_k_vec_start + ki * WARP_TILE_K
-                    state_vecs[vi * WARP_TILE_K_ITERS + ki] = state_tensor.vec_load(
-                        (hv_i, global_v_i, warp_k_vec_i), VALUES_PER_THREAD_K
+                    state_vecs[vi * WARP_TILE_K_ITERS + ki] = (
+                        read_state_tensor.vec_load(
+                            (hv_i, global_v_i, warp_k_vec_i), VALUES_PER_THREAD_K
+                        )
                     )
                     if const_expr("f32" in state_dtype):
                         pass
@@ -365,11 +391,11 @@ def _write_q(_sum_hq=sum_hq, _gv=global_v_i, _sq=sq_i):
                         out_vec = state_vecs[vi * WARP_TILE_K_ITERS + ki]
                     else:
                         out_vec = state_vecs[vi * WARP_TILE_K_ITERS + ki].truncf(vec_t)
-                    state_tensor.vec_store(
+                    write_state_tensor.vec_store(
                         (hv_i, global_v_i, warp_k_vec_i), out_vec, VALUES_PER_THREAD_K
                     )
 
-        if pool_idx >= 0:
+        if (read_pool_idx >= 0) & (write_pool_idx >= 0):
             _do_decode()
 
     @flyc.jit
@@ -381,7 +407,8 @@ def launch_gdr_decode_kernel(
         b: fx.Tensor,
         dt_bias: fx.Tensor,
         A_log: fx.Tensor,
-        indices: fx.Tensor,
+        read_indices: fx.Tensor,
+        write_indices: fx.Tensor,
         state: fx.Tensor,
         out: fx.Tensor,
         batch_size: fx.Int32,
@@ -397,7 +424,8 @@ def launch_gdr_decode_kernel(
             b,
             dt_bias,
             A_log,
-            indices,
+            read_indices,
+            write_indices,
             state,
             out,
             batch_size,

aiter/ops/flydsl/linear_attention_kernels.py

@@ -96,19 +96,45 @@ def flydsl_gdr_decode(
     use_qk_l2norm: bool,
     need_shuffle_state: bool,
     stream: torch.cuda.Stream = None,
+    read_indices: torch.Tensor | None = None,
+    write_indices: torch.Tensor | None = None,
 ):
     if stream is None:
         stream = torch.cuda.current_stream()
     device = query.device
     dtype = query.dtype
-    for input in [query, key, value, a, b, dt_bias, A_log, indices, out]:
+    read_indices = indices if read_indices is None else read_indices
+    write_indices = indices if write_indices is None else write_indices
+    for input in [
+        query,
+        key,
+        value,
+        a,
+        b,
+        dt_bias,
+        A_log,
+        read_indices,
+        write_indices,
+        out,
+    ]:
         assert input.device == device
     assert state.data_ptr() % 16 == 0
     for input in [key, value, a, b, dt_bias, out]:
         assert input.dtype == dtype
     assert state.dtype in [torch.float, torch.bfloat16]
     assert A_log.dtype in [torch.float, torch.bfloat16]
-    assert indices.dtype == torch.int32
+    assert read_indices.dtype == torch.int32
+    assert write_indices.dtype == torch.int32
+    if query.stride(-1) != 1:
+        raise ValueError(
+            "`query` must have a contiguous last dimension for vectorized loads; "
+            f"got stride {query.stride()}."
+        )
+    if key.stride(-1) != 1:
+        raise ValueError(
+            "`key` must have a contiguous last dimension for vectorized loads; "
+            f"got stride {key.stride()}."
+        )
 
     if need_shuffle_state:
         state_ = state.permute(0, 1, 3, 2).contiguous()
@@ -136,6 +162,9 @@ def flydsl_gdr_decode(
         num_v_heads,
         head_k_dim,
         head_v_dim,
+        query.stride(),
+        key.stride(),
+        value.stride(),
         state_.stride(),
         a.stride(),
         b.stride(),
@@ -145,14 +174,15 @@ def flydsl_gdr_decode(
     with torch.cuda.device(query.device.index):
         _run_compiled(
             exe,
-            query.contiguous(),
-            key.contiguous(),
-            value.contiguous(),
+            query,
+            key,
+            value,
             a,
             b,
             dt_bias.contiguous(),
             A_log.contiguous(),
-            indices.contiguous(),
+            read_indices.contiguous(),
+            write_indices.contiguous(),
             state_,
             out,
             batch_size,

aiter/ops/flydsl/test_flydsl_linear_attention.py

@@ -452,3 +452,105 @@ def test_flydsl_gdr_decode(args):
         maxdiff_state = (inouts[-2] - ref_inouts[-2]).abs().max()
         print(f"maxdiff_out:{maxdiff_out}\nmaxdiff_state:{maxdiff_state}")
         assert is_allclose
+
+
+@pytest.mark.parametrize("input_index,input_name", [(1, "query"), (2, "key")])
+def test_flydsl_gdr_decode_rejects_noncontiguous_vector_dimension(
+    input_index, input_name
+):
+    args = Args(
+        dtype=torch.bfloat16,
+        b=2,
+        sq=1,
+        num_k_heads=16,
+        num_v_heads=32,
+        head_k_dim=128,
+        head_v_dim=128,
+    )
+    inouts = list(create_inputs(args) + create_outputs(args))
+    tensor = inouts[input_index]
+    strided_storage = torch.randn(
+        *tensor.shape[:-1],
+        tensor.shape[-1] * 2,
+        dtype=tensor.dtype,
+        device=tensor.device,
+    )
+    inouts[input_index] = strided_storage[..., ::2]
+    assert inouts[input_index].shape == tensor.shape
+    assert inouts[input_index].stride(-1) == 2
+
+    with pytest.raises(
+        ValueError,
+        match=rf"`{input_name}` must have a contiguous last dimension",
+    ):
+        func(*inouts)
+
+
+@pytest.mark.parametrize(
+    "num_k_heads,num_v_heads",
+    [(2, 8), (4, 8), (4, 16), (8, 16), (8, 32), (16, 32), (16, 64)],
+)
+@pytest.mark.parametrize("state_dtype", [torch.float32, torch.bfloat16])
+def test_flydsl_gdr_decode_strided_inputs_and_split_state_indices(
+    num_k_heads, num_v_heads, state_dtype
+):
+    batch, seq_length, dim = 2, 1, 128
+    mixed_qkv = torch.randn(
+        batch,
+        seq_length,
+        num_k_heads * 2 + num_v_heads,
+        dim,
+        dtype=torch.bfloat16,
+        device="cuda",
+    )
+    query, key, value = torch.split(
+        mixed_qkv, [num_k_heads, num_k_heads, num_v_heads], dim=2
+    )
+    mixed_ba = torch.randn(batch, num_v_heads * 2, dtype=torch.bfloat16, device="cuda")
+    a, b = (
+        tensor.reshape(batch, seq_length, num_v_heads)
+        for tensor in torch.split(mixed_ba, num_v_heads, dim=-1)
+    )
+    dt_bias = torch.randn(num_v_heads, dtype=torch.bfloat16, device="cuda")
+    A_log = torch.randn(num_v_heads, dtype=torch.float32, device="cuda")
+    read_indices = torch.tensor([1, 3], dtype=torch.int32, device="cuda")
+    write_indices = torch.tensor([2, 4], dtype=torch.int32, device="cuda")
+    state = torch.randn(5, num_v_heads, dim, dim, dtype=state_dtype, device="cuda")
+    reference_state = state.clone()
+    reference_state[write_indices.long()] = reference_state[read_indices.long()]
+    output = torch.empty_like(value)
+    reference_output = torch.empty_like(value)
+
+    common_kwargs = {
+        "dt_bias": dt_bias,
+        "A_log": A_log,
+        "indices": write_indices,
+        "use_qk_l2norm": True,
+        "need_shuffle_state": False,
+    }
+    flydsl_gdr_decode(
+        query,
+        key,
+        value,
+        a,
+        b,
+        state=state,
+        out=output,
+        read_indices=read_indices,
+        write_indices=write_indices,
+        **common_kwargs,
+    )
+    flydsl_gdr_decode(
+        query.contiguous(),
+        key.contiguous(),
+        value.contiguous(),
+        a.contiguous(),
+        b.contiguous(),
+        state=reference_state,
+        out=reference_output,
+        **common_kwargs,
+    )
+    torch.cuda.synchronize()
+
+    torch.testing.assert_close(output, reference_output, rtol=0, atol=0)
+    torch.testing.assert_close(state, reference_state, rtol=0, atol=0)