alignment enforced for q,k,v

Signed-off-by: Jinseok Lee <jindol21@rebellions.ai>

Author: rebel-jindol21

vllm_rbln/triton_kernels/attention.py

@@ -377,6 +377,10 @@ def _(
 
     output = torch.empty_like(query)
 
+    query = rblib.align_tensor_last_dim_to_64(query)
+    key = rblib.align_tensor_last_dim_to_64(key)
+    value = rblib.align_tensor_last_dim_to_64(value)
+
     NUM_HEAD = query.shape[1]
     NUM_GROUP = query.shape[2]
     HEAD_DIM = query.shape[-1]
@@ -435,6 +439,10 @@ def _(
 
     output = torch.empty_like(query)
 
+    query = rblib.align_tensor_last_dim_to_64(query)
+    key = rblib.align_tensor_last_dim_to_64(key)
+    value = rblib.align_tensor_last_dim_to_64(value)
+
     NUM_HEAD = query.shape[1]
     NUM_GROUP = query.shape[2]
     HEAD_DIM = query.shape[-1]

vllm_rbln/triton_kernels/causal_attention.py

@@ -519,6 +519,10 @@ def _(
 
     output = torch.empty_like(query)
 
+    query = rblib.align_tensor_last_dim_to_64(query)
+    key = rblib.align_tensor_last_dim_to_64(key)
+    value = rblib.align_tensor_last_dim_to_64(value)
+
     NUM_HEAD = query.shape[1]
     NUM_GROUP = query.shape[2]
     HEAD_DIM = query.shape[-1]
@@ -591,6 +595,10 @@ def _(
 
     output = torch.empty_like(query)
 
+    query = rblib.align_tensor_last_dim_to_64(query)
+    key = rblib.align_tensor_last_dim_to_64(key)
+    value = rblib.align_tensor_last_dim_to_64(value)
+
     NUM_HEAD = query.shape[1]
     NUM_GROUP = query.shape[2]
     HEAD_DIM = query.shape[-1]

vllm_rbln/triton_kernels/sliding_window_attention.py

@@ -453,6 +453,10 @@ def _(
     qk_scale = qk_scale.to(torch.float32)
     output = torch.empty_like(query)
 
+    query = rblib.align_tensor_last_dim_to_64(query)
+    key = rblib.align_tensor_last_dim_to_64(key)
+    value = rblib.align_tensor_last_dim_to_64(value)
+
     NUM_HEAD = query.shape[1]
     NUM_GROUP = query.shape[2]
     HEAD_DIM = query.shape[-1]
@@ -525,6 +529,10 @@ def _(
     qk_scale = qk_scale.to(torch.float32)
     output = torch.empty_like(query)
 
+    query = rblib.align_tensor_last_dim_to_64(query)
+    key = rblib.align_tensor_last_dim_to_64(key)
+    value = rblib.align_tensor_last_dim_to_64(value)
+
     NUM_HEAD = query.shape[1]
     NUM_GROUP = query.shape[2]
     HEAD_DIM = query.shape[-1]