From 45c52010decf7543aa5410237e7b7444abd0db9d Mon Sep 17 00:00:00 2001 From: Dudu Moshe Date: Thu, 21 Nov 2024 22:20:52 +0200 Subject: [PATCH] wip --- torch_xla/experimental/custom_kernel.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/torch_xla/experimental/custom_kernel.py b/torch_xla/experimental/custom_kernel.py index d1b47485f7e..9d8f27183d5 100644 --- a/torch_xla/experimental/custom_kernel.py +++ b/torch_xla/experimental/custom_kernel.py @@ -261,10 +261,10 @@ def forward(ctx, q, k, v, causal, q_segment_ids, kv_segment_ids, sm_scale, ab, q_segment_ids = xs.enable_manual_sharding( q_segment_ids, partition_spec[:q_segment_ids.ndim], mesh=mesh).global_tensor - if kv_segment_ids is not None: - kv_segment_ids = xs.enable_manual_sharding( - kv_segment_ids, partition_spec[:kv_segment_ids.ndim], - mesh=mesh).global_tensor + # if kv_segment_ids is not None: + # kv_segment_ids = xs.enable_manual_sharding( + # kv_segment_ids, partition_spec[:kv_segment_ids.ndim], + # mesh=mesh).global_tensor # It computes the shape and type of o, l, m. shapes = [q.shape] @@ -378,10 +378,10 @@ def backward(ctx, grad_output): q_segment_ids = xs.enable_manual_sharding( q_segment_ids, partition_spec[:q_segment_ids.ndim], mesh=mesh).global_tensor - if kv_segment_ids is not None: - kv_segment_ids = xs.enable_manual_sharding( - kv_segment_ids, partition_spec[:kv_segment_ids.ndim], - mesh=mesh).global_tensor + # if kv_segment_ids is not None: + # kv_segment_ids = xs.enable_manual_sharding( + # kv_segment_ids, partition_spec[:kv_segment_ids.ndim], + # mesh=mesh).global_tensor if ctx.needs_input_grad[0]: payload, _ = trace_pallas(