Refactors dynamic mask function to improve clarity

[LoserCheems]

Renames dt_states parameter to zoh_states for better semantic meaning.

Updates variable names from attn_mask to attn_bias to more accurately
reflect the additive bias nature of the operation.

Improves code organization by moving attn_mask creation logic and
adding explicit handling for sequences shorter than keep_window_size.

Enhances documentation with clearer parameter descriptions and
adds return type specification.

[Copilot]

Pull Request Overview

This PR refactors the dynamic mask function in the benchmark_grad.py file to improve code clarity and semantics. The changes focus on renaming variables for better meaning and improving code organization.

• Renames dt_states parameter to zoh_states for clearer semantic meaning
• Updates variable names from attn_mask to attn_bias to better reflect additive bias operations
• Improves code organization by moving mask creation logic and adding explicit handling for short sequences

[Copilot]

[nitpick] Moving the q_vec assignment before k_vecs and v_vecs creates an inconsistent ordering. The original placement after k_vecs and v_vecs was more logical as it groups the vector extractions together before their usage.

Suggested change

	q_vec = query_states[b_idx, h_idx, q_idx, :] # [head_dim]
	k_vecs = key_states[b_idx, h_idx, non_mask_indices, :] # [keep_window_size, head_dim]
	v_vecs = value_states[b_idx, h_idx, non_mask_indices, :] # [keep_window_size, head_dim]
	k_vecs = key_states[b_idx, h_idx, non_mask_indices, :] # [keep_window_size, head_dim]
	v_vecs = value_states[b_idx, h_idx, non_mask_indices, :] # [keep_window_size, head_dim]
	q_vec = query_states[b_idx, h_idx, q_idx, :] # [head_dim]

[Copilot]

The parameter description lacks explanation of what 'zoh_states' represents functionally. Consider adding a brief description of its purpose in the dynamic masking process.

Suggested change

	zoh_states: zoh_states of shape (batch_size, num_kv_heads, key_sequence_length)
	zoh_states: Binary tensor indicating zones of high attention, of shape
	(batch_size, num_kv_heads, key_sequence_length). Used to guide the
	dynamic masking process by identifying key positions to retain.