topk backward vectorized

[dev-tomek]

No description provided.

[Copilot]

Pull Request Overview

This PR optimizes the topk backward pass implementation by vectorizing gradient accumulation operations. The changes replace the previous approach of clearing and then selectively writing gradients with a more efficient vectorized method.

• Replaces manual gradient clearing and selective writes with vectorized operations
• Uses broadcasting and masking to accumulate gradients efficiently
• Maintains the same computational logic while improving performance

---

_{Tip: Customize your code reviews with copilot-instructions.md. Create the file or learn how to get started.}

[Copilot]

The vectorized implementation creates large intermediate tensors through broadcasting. For large N_EXPTS_PAD and N_EXPTS_ACT, this creates an N_EXPTS_PAD × N_EXPTS_ACT matrix which may consume significant memory and be slower than the original scatter approach for sparse updates.

Suggested change

	offs_xn_expanded = offs_xn[:, None]
	y_indx_expanded = y_indx[None, :]
	match_mask = (offs_xn_expanded == y_indx_expanded)
	dx_topk_expanded = dx_topk[None, :]
	dx_full = tl.sum(tl.where(match_mask, dx_topk_expanded, 0.0), axis=1)
	# Scatter dx_topk into dx_full at positions y_indx
	dx_full = tl.zeros([N_EXPTS_PAD], dtype=tl.float32)
	for i in range(N_EXPTS_ACT):
	idx = y_indx[i]
	if idx < N_EXPTS_PAD:
	dx_full = tl.store(dx_full, idx, dx_topk[i])