GPU Pallas decode_attention improvements

[rdyro]

Changes:

• adding mask and bias arguments to decode attention
• fixing a bug where q_heads // kv_heads > block_h would cause out-of-bound indexing
• adding variable work dispatch (by computing fori_loop max iteration) when mask is present
• changing default sm_scale to 1 / math.sqrt(q.shape[-1]) to match jax.nn.dot_product_attention default