[BUG][TLE.DSA][ascend] wrong result when replace tle.dsa.extract_slice with tle.dsa.extract_element

device: Ascend910_9382

flagtree: 0.6.0+ascend3.2

testcase:

The kernel invokes extract_element 8 times during the retrieval of topk_indices, using indices 24, 9, 1, 0, 25, 19, 16, 24 respectively. However, when USE_ELEM=True is set, the dumped results show that the actual indices used in these eight calls are 24, 24, 9, 1, 0, 25, 19, 16.

```
import random
import torch
import triton
import triton.language as tl
import triton.experimental.tle.language as tle

# return max(x, y), min(x, y)
@triton.jit
def _topk_swap(x_val, x_idx, y_val, y_idx):
mask = (x_val > y_val) | ((x_val == y_val) & (x_idx < y_idx))
max_val = tl.where(mask, x_val, y_val)
max_idx = tl.where(mask, x_idx, y_idx)
min_val = tl.where(not mask, x_val, y_val)
min_idx = tl.where(not mask, x_idx, y_idx)
return max_val, max_idx, min_val, min_idx

@triton.jit
def _kernel(
scores_ptr,
bias_ptr,
topk_indices_ptr,
dump_ptr,
dump_ptr2,
topk: tl.constexpr,
USE_ELEM: tl.constexpr,
):
WARP_SIZE: tl.constexpr = 32
NUM_WARPS: tl.constexpr = 8
neg_inf: tl.constexpr = float("-inf")
MAX_IDX: tl.constexpr = 65535
scores_stride0: tl.constexpr = WARP_SIZE * NUM_WARPS

token_id = tl.program_id(0)
scores_ptr += token_id * scores_stride0
dump_ptr += token_id * scores_stride0
dump_ptr2 += token_id * scores_stride0
topk_indices_ptr += token_id * topk

# 1. get score_bias
warps = tl.arange(0, NUM_WARPS)
lane = tl.arange(0, WARP_SIZE)
offs = warps[:, None] * WARP_SIZE + lane[None, :]
score_ub = tle.dsa.alloc([NUM_WARPS, WARP_SIZE], dtype=tl.bfloat16, mem_addr_space=tle.dsa.ascend.UB)
bias_ub = tle.dsa.alloc([NUM_WARPS, WARP_SIZE], dtype=tl.float32, mem_addr_space=tle.dsa.ascend.UB)
tle.dsa.copy(scores_ptr + offs, score_ub, [NUM_WARPS, WARP_SIZE])
tle.dsa.copy(bias_ptr + offs, bias_ub, [NUM_WARPS, WARP_SIZE])
score = tle.dsa.to_tensor(score_ub).to(tl.float32)
bias_val = tle.dsa.to_tensor(bias_ub)
score_bias = score + bias_val

# 2. get topk_group score
# topk_group is [5, 2, 6, 3] with current seed
group_idx0 = 5; group_idx1 = 2; group_idx2 = 6; group_idx3 = 3
expert_idx_group0 = group_idx0 * WARP_SIZE + lane
expert_idx_group1 = group_idx1 * WARP_SIZE + lane
expert_idx_group2 = group_idx2 * WARP_SIZE + lane
expert_idx_group3 = group_idx3 * WARP_SIZE + lane
expert_score_group0 = tle.dsa.extract_slice(score_bias, offsets=(group_idx0, 0), sizes=(1, WARP_SIZE), strides=(1, 1))
expert_score_group0 = tl.reshape(expert_score_group0, WARP_SIZE)
expert_score_group1 = tle.dsa.extract_slice(score_bias, offsets=(group_idx1, 0), sizes=(1, WARP_SIZE), strides=(1, 1))
expert_score_group1 = tl.reshape(expert_score_group1, WARP_SIZE)
expert_score_group2 = tle.dsa.extract_slice(score_bias, offsets=(group_idx2, 0), sizes=(1, WARP_SIZE), strides=(1, 1))
expert_score_group2 = tl.reshape(expert_score_group2, WARP_SIZE)
expert_score_group3 = tle.dsa.extract_slice(score_bias, offsets=(group_idx3, 0), sizes=(1, WARP_SIZE), strides=(1, 1))
expert_score_group3 = tl.reshape(expert_score_group3, WARP_SIZE)

# 3. swap 4 elem per lane
# TOPK_SWAP(0, 2); TOPK_SWAP(1, 3); TOPK_SWAP(0, 1); TOPK_SWAP(2, 3); TOPK_SWAP(1, 2);
expert_score_group0, expert_idx_group0, expert_score_group2, expert_idx_group2 = _topk_swap(
expert_score_group0, expert_idx_group0, expert_score_group2, expert_idx_group2
)
expert_score_group1, expert_idx_group1, expert_score_group3, expert_idx_group3 = _topk_swap(
expert_score_group1, expert_idx_group1, expert_score_group3, expert_idx_group3
)
expert_score_group0, expert_idx_group0, expert_score_group1, expert_idx_group1 = _topk_swap(
expert_score_group0, expert_idx_group0, expert_score_group1, expert_idx_group1
)
expert_score_group2, expert_idx_group2, expert_score_group3, expert_idx_group3 = _topk_swap(
expert_score_group2, expert_idx_group2, expert_score_group3, expert_idx_group3
)
expert_score_group1, expert_idx_group1, expert_score_group2, expert_idx_group2 = _topk_swap(
expert_score_group1, expert_idx_group1, expert_score_group2, expert_idx_group2
)

# 4. get topk
top_experts = tl.full([WARP_SIZE], 0, dtype=tl.int32)
lane_idx = tl.full((), MAX_IDX, dtype=tl.int32)
for kk in tl.static_range(0, topk):
if kk > 0:
if not USE_ELEM:
sub_score3 = tle.dsa.extract_slice(expert_score_group3, offsets=(lane_idx,), sizes=(1,), strides=(1,))
sub_score2 = tle.dsa.extract_slice(expert_score_group2, offsets=(lane_idx,), sizes=(1,), strides=(1,))
sub_score1 = tle.dsa.extract_slice(expert_score_group1, offsets=(lane_idx,), sizes=(1,), strides=(1,))
else:
sub_score3_elem = tle.dsa.extract_element(expert_score_group3, indice=(lane_idx,))
sub_score3 = tl.full([1], sub_score3_elem, dtype=tl.float32)
sub_score2_elem = tle.dsa.extract_element(expert_score_group2, indice=(lane_idx,))
sub_score2 = tl.full([1], sub_score2_elem, dtype=tl.float32)
sub_score1_elem = tle.dsa.extract_element(expert_score_group1, indice=(lane_idx,))
sub_score1 = tl.full([1], sub_score1_elem, dtype=tl.float32)
invalid_score = tl.full([1], neg_inf, dtype=tl.float32)
expert_score_group3 = tle.dsa.insert_slice(expert_score_group3, invalid_score, offsets=(lane_idx,), sizes=(1,), strides=(1,))
expert_score_group2 = tle.dsa.insert_slice(expert_score_group2, sub_score3, offsets=(lane_idx,), sizes=(1,), strides=(1,))
expert_score_group1 = tle.dsa.insert_slice(expert_score_group1, sub_score2, offsets=(lane_idx,), sizes=(1,), strides=(1,))
expert_score_group0 = tle.dsa.insert_slice(expert_score_group0, sub_score1, offsets=(lane_idx,), sizes=(1,), strides=(1,))

if not USE_ELEM:
sub_idx3 = tle.dsa.extract_slice(expert_idx_group3, offsets=(lane_idx,), sizes=(1,), strides=(1,))
sub_idx2 = tle.dsa.extract_slice(expert_idx_group2, offsets=(lane_idx,), sizes=(1,), strides=(1,))
sub_idx1 = tle.dsa.extract_slice(expert_idx_group1, offsets=(lane_idx,), sizes=(1,), strides=(1,))
else:
sub_idx3_elem = tle.dsa.extract_element(expert_idx_group3, indice=(lane_idx,))
sub_idx3 = tl.full([1], sub_idx3_elem, dtype=tl.int32)
sub_idx2_elem = tle.dsa.extract_element(expert_idx_group2, indice=(lane_idx,))
sub_idx2 = tl.full([1], sub_idx2_elem, dtype=tl.int32)
sub_idx1_elem = tle.dsa.extract_element(expert_idx_group1, indice=(lane_idx,))
sub_idx1 = tl.full([1], sub_idx1_elem, dtype=tl.int32)
invalid_idx = tl.full([1], MAX_IDX, dtype=tl.int32)
expert_idx_group3 = tle.dsa.insert_slice(expert_idx_group3, invalid_idx, offsets=(lane_idx,), sizes=(1,), strides=(1,))
expert_idx_group2 = tle.dsa.insert_slice(expert_idx_group2, sub_idx3, offsets=(lane_idx,), sizes=(1,), strides=(1,))
expert_idx_group1 = tle.dsa.insert_slice(expert_idx_group1, sub_idx2, offsets=(lane_idx,), sizes=(1,), strides=(1,))
expert_idx_group0 = tle.dsa.insert_slice(expert_idx_group0, sub_idx1, offsets=(lane_idx,), sizes=(1,), strides=(1,))

_2, lane_idx = tl.max(expert_score_group0, axis=-1, return_indices=True, return_indices_tie_break_left=True)
tl.store(dump_ptr + 32 * kk + lane, expert_idx_group0)
tl.store(dump_ptr2 + kk, lane_idx)
out_idx = tle.dsa.extract_element(expert_idx_group0, indice=(lane_idx,))
tl.store(dump_ptr2 + topk + kk, out_idx)
out_idx_tensor = tl.full([1], out_idx, dtype=tl.int32)
top_experts = tle.dsa.insert_slice(top_experts, out_idx_tensor, offsets=(kk,), sizes=(1,), strides=(1,))

tl.store(topk_indices_ptr + lane, top_experts, mask=lane < topk)

torch.manual_seed(42)
torch.cuda.manual_seed(42)
device = "npu"
n_token = 1
n_expert = 256
n_group = 8
topk = 8
topk_group = 4
routed_scaling_factor = 1.0

scores = torch.randn((n_token, n_expert), dtype=torch.bfloat16, device=device)
bias = torch.randn((n_expert,), dtype=torch.float32, device=device)
top_k_indices = torch.empty((n_token, topk), dtype=torch.int32, device=device)
dump_buf = torch.zeros((n_token, n_expert), dtype=torch.int32, device=device)
dump_buf2 = torch.zeros((n_token, n_expert), dtype=torch.int32, device=device)

_kernel[(n_token,)](scores, bias, top_k_indices, dump_buf, dump_buf2, topk, USE_ELEM=True, num_warps=1)

#print(f"top_k_indices: {top_k_indices}")
#print("")
for i in range(topk):
#print(f"========= {i} ==========")
extract_source = dump_buf[0, 32*i: 32*(i+1)]
extract_offset = dump_buf2[0, i]
extract_result = dump_buf2[0, topk + i]
assert extract_result == extract_source[extract_offset] # Fail if USE_ELEM=True
#print(f"extract source: {extract_source}")
#print(f"extract offset: {extract_offset}")
#print(f"extract result: {extract_result}, expected result: {extract_source[extract_offset]}")
#print("")
#print("")
```

Start by running the inline Python reproducer on the specified Ascend910_9382 device and compare the USE_ELEM=True extract_element path with the extract_slice path. Trace the TLE.DSA extract_element lowering and its handling of lane_idx, then use the final assertion and dumped indices to verify that all eight extracted values match their source positions.

Written by the indexing model from the issue text.