Add sm_scale in ragged attention kernel (#8771)

Author: yaochengji

test/test_pallas.py

@@ -746,6 +746,7 @@ def _verify_ragged_paged_attention_with_dynamo(
       num_kv_pages_per_block,
       num_queries_per_block,
       pad_num_q_tokens=False,
+      sm_scale=1.0,
   ):
     num_seqs = len(seq_lens)
     q, k_pages, v_pages, page_indices, cu_q_lens, kv_lens = self._ragged_pagedattention_generate_qkv(
@@ -768,7 +769,8 @@ def _verify_ragged_paged_attention_with_dynamo(
     def ragged_paged_attention_wrapper(q, k_pages, v_pages, kv_lens,
                                        page_indices, cu_q_lens, num_seqs,
                                        num_kv_pages_per_block,
-                                       num_queries_per_block, use_kernel):
+                                       num_queries_per_block, use_kernel,
+                                       sm_scale):
       return torch.ops.xla.ragged_paged_attention(
           q,
           k_pages,
@@ -780,6 +782,7 @@ def ragged_paged_attention_wrapper(q, k_pages, v_pages, kv_lens,
           num_kv_pages_per_block,
           num_queries_per_block,
           use_kernel=use_kernel,
+          sm_scale=sm_scale,
       )
 
     compiled_paged_attention = torch.compile(
@@ -796,6 +799,7 @@ def ragged_paged_attention_wrapper(q, k_pages, v_pages, kv_lens,
         num_kv_pages_per_block=num_kv_pages_per_block,
         num_queries_per_block=num_queries_per_block,
         use_kernel=True,
+        sm_scale=sm_scale,
     )
 
     nonkernel_output = compiled_paged_attention(
@@ -809,6 +813,7 @@ def ragged_paged_attention_wrapper(q, k_pages, v_pages, kv_lens,
         num_kv_pages_per_block=num_kv_pages_per_block,
         num_queries_per_block=num_queries_per_block,
         use_kernel=False,
+        sm_scale=sm_scale,
     )
 
     kernel_output_cpu = kernel_output.cpu()
@@ -836,6 +841,7 @@ def ragged_paged_attention_wrapper(q, k_pages, v_pages, kv_lens,
                 num_seqs=num_seqs,
                 num_kv_pages_per_block=num_kv_pages_per_block,
                 num_queries_per_block=num_queries_per_block,
+                sm_scale=sm_scale,
             )[1]))
     jax_kernel_output_cpu = jax_kernel_output.cpu()
 
@@ -845,21 +851,21 @@ def ragged_paged_attention_wrapper(q, k_pages, v_pages, kv_lens,
           torch.allclose(
               kernel_output_cpu[:actual_num_q_tokens],
               nonkernel_output_cpu[:actual_num_q_tokens],
-              atol=2e-1,
+              atol=2e-2,
               rtol=1e-2))
       self.assertTrue(
           torch.allclose(
               kernel_output_cpu[:actual_num_q_tokens],
               jax_kernel_output_cpu[:actual_num_q_tokens],
-              atol=2e-1,
+              atol=2e-2,
               rtol=1e-2))
     else:
       self.assertTrue(
           torch.allclose(
-              kernel_output_cpu, nonkernel_output_cpu, atol=2e-1, rtol=1e-2))
+              kernel_output_cpu, nonkernel_output_cpu, atol=2e-2, rtol=1e-2))
       self.assertTrue(
           torch.allclose(
-              kernel_output_cpu, jax_kernel_output_cpu, atol=2e-1, rtol=1e-2))
+              kernel_output_cpu, jax_kernel_output_cpu, atol=2e-2, rtol=1e-2))
 
   @unittest.skipIf(xr.device_type() != 'TPU' or tpu.version() < 4,
                    "This test only works on TPUv4+.")
@@ -882,6 +888,7 @@ def test_ragged_paged_attention_wrapper_no_query_padding_with_dynamo(self):
     dtype = torch.float32
     page_size = 16
     num_pages = 32768
+    sm_scale = head_dim**-0.5
 
     self._verify_ragged_paged_attention_with_dynamo(
         seq_lens,
@@ -892,6 +899,7 @@ def test_ragged_paged_attention_wrapper_no_query_padding_with_dynamo(self):
         dtype,
         num_kv_pages_per_block=128,
         num_queries_per_block=8,
+        sm_scale=sm_scale,
     )
 
   @parameterized.product(
@@ -910,6 +918,7 @@ def test_ragged_paged_attention_wrapper_with_query_padding_with_dynamo(
     dtype = torch.float32
     page_size = 16
     num_pages = 32768
+    sm_scale = head_dim**-0.5
 
     self._verify_ragged_paged_attention_with_dynamo(
         seq_lens,
@@ -921,6 +930,7 @@ def test_ragged_paged_attention_wrapper_with_query_padding_with_dynamo(
         num_kv_pages_per_block=128,
         num_queries_per_block=num_queries_per_block,
         pad_num_q_tokens=True,
+        sm_scale=sm_scale,
     )
 
   @unittest.skipIf(xr.device_type() != 'TPU' or tpu.version() < 4,

torch_xla/experimental/custom_kernel.py

@@ -708,6 +708,7 @@ def _ragged_paged_attention_nonkernel(
     page_indices,  # i32[num_tokens, pages_per_sequence]
     cu_q_lens,  # i32[num_tokens + 1]
     num_seqs,  # int
+    sm_scale,  # float
 ):
   _, num_q_heads, head_dim = queries.shape
   num_kv_heads, total_num_pages, page_size, _ = k_pages.shape
@@ -751,6 +752,7 @@ def _ragged_paged_attention_nonkernel(
     attn = torch.einsum("qhd,khd->hqk", q,
                         k)  # [num_query_heads, cur_q_len, kv_len]
     attn = attn.float()
+    attn = attn * sm_scale
     empty_mask = torch.ones(cur_q_len, cur_kv_len, device=attn.device)
     mask = torch.triu(empty_mask, diagonal=cur_kv_len - cur_q_len + 1).bool()
     attn.masked_fill_(mask, float("-inf"))
@@ -784,6 +786,7 @@ def ragged_paged_attention(
     num_kv_pages_per_block,
     num_queries_per_block,
     use_kernel=True,
+    sm_scale=1.0,
     # TODO(jevinjiang, xiowei): add attn_logits_soft_cap.
     # attn_logits_soft_cap: float | None = None,
 ):  # [batch_size, query_len, num_heads, head_dim]:
@@ -797,6 +800,7 @@ def ragged_paged_attention(
         page_indices,
         cu_q_lens,
         num_seqs,
+        sm_scale,
     )
 
   # Import JAX within the function such that we don't need to call the jax_import_guard()
@@ -813,11 +817,13 @@ def ragged_paged_attention(
       num_seqs=num_seqs,
       num_kv_pages_per_block=num_kv_pages_per_block,
       num_queries_per_block=num_queries_per_block,
+      sm_scale=sm_scale,
       static_argnames=[
           "num_kv_pages_per_block",
           "num_queries_per_block",
           "mask_value",
           "num_seqs",
+          "sm_scale",
       ],
   )
 
@@ -1541,28 +1547,27 @@ def multi_queries_paged_attention_non_xla(q: torch.Tensor,
 
 
 XLA_LIB.define(
-    "ragged_paged_attention(Tensor q, Tensor k_pages, Tensor v_pages, Tensor kv_lens, Tensor page_indices, Tensor cu_q_lens, int num_seqs, int num_kv_pages_per_block, int num_queries_per_block, bool use_kernel) -> Tensor",
+    "ragged_paged_attention(Tensor q, Tensor k_pages, Tensor v_pages, Tensor kv_lens, Tensor page_indices, Tensor cu_q_lens, int num_seqs, int num_kv_pages_per_block, int num_queries_per_block, bool use_kernel, float sm_scale) -> Tensor",
 )
 
 
 @impl(XLA_LIB, "ragged_paged_attention", "XLA")
-def ragged_paged_attention_xla(q: torch.Tensor, k_pages: torch.Tensor,
-                               v_pages: torch.Tensor, kv_lens: torch.Tensor,
-                               page_indices: torch.Tensor,
-                               cu_q_lens: torch.Tensor, num_seqs: int,
-                               num_kv_pages_per_block: int,
-                               num_queries_per_block: int, use_kernel: bool):
+def ragged_paged_attention_xla(
+    q: torch.Tensor, k_pages: torch.Tensor, v_pages: torch.Tensor,
+    kv_lens: torch.Tensor, page_indices: torch.Tensor, cu_q_lens: torch.Tensor,
+    num_seqs: int, num_kv_pages_per_block: int, num_queries_per_block: int,
+    use_kernel: bool, sm_scale: float):
   return ragged_paged_attention(q, k_pages, v_pages, kv_lens, page_indices,
                                 cu_q_lens, num_seqs, num_kv_pages_per_block,
-                                num_queries_per_block, use_kernel)
+                                num_queries_per_block, use_kernel, sm_scale)
 
 
 @impl(XLA_LIB, "ragged_paged_attention", "CompositeExplicitAutograd")
 def ragged_paged_attention_non_xla(
     q: torch.Tensor, k_pages: torch.Tensor, v_pages: torch.Tensor,
     kv_lens: torch.Tensor, page_indices: torch.Tensor, cu_q_lens: torch.Tensor,
     num_seqs: int, num_kv_pages_per_block: int, num_queries_per_block: int,
-    use_kernel: bool):
+    use_kernel: bool, sm_scale: float):
   return non_xla_attetion(q, k_pages, v_pages, "paged")
 
 

torch_xla/experimental/pallas_kernels/ragged_paged_attention_kernel.py

@@ -355,6 +355,7 @@ def _flash_attention(
     page_size: int,
     head_dim: int,
     num_q_heads_per_kv_head: int,
+    sm_scale: float,
 ):
   assert q_ref.shape == (num_q_heads_per_kv_head, num_queries_per_block,
                          head_dim)
@@ -405,6 +406,7 @@ def init_scratch_ref():  # pylint: disable=unused-variable
       'qd,td->qt', q, k,
       preferred_element_type=jnp.float32)  # [block_q, block_k]
   assert s.shape == (num_queries_per_block, kv_blk_size)
+  s = s * sm_scale
 
   # Modify the mask accordingly: first form the mask. Then move the mask up/down to the right place.
   cur_seq_idx = seq_ids[logical_q_blk_idx]
@@ -597,6 +599,7 @@ def paged_flash_attention_kernel(
     num_seqs: int,
     num_kv_pages_per_block: int,
     mask_value: float,
+    sm_scale: float,
 ):
   kv_head_idx, logical_q_blk_idx, kv_blk_idx = (
       pl.program_id(0),
@@ -704,6 +707,7 @@ def prefetch_next_block():  # pylint: disable=unused-variable
           page_size=page_size,
           head_dim=head_dim,
           num_q_heads_per_kv_head=num_q_heads_per_kv_head,
+          sm_scale=sm_scale,
       )
     step_ref[0] = step + 1
     # end of get_kv_and_run_flash_attention
@@ -724,6 +728,7 @@ def _round_up_to_multiple_of_tm(x, tm):
         "num_queries_per_block",
         "mask_value",
         "num_seqs",
+        "sm_scale",
     ],
 )
 def ragged_paged_attention(
@@ -738,6 +743,7 @@ def ragged_paged_attention(
     mask_value: float = DEFAULT_MASK_VALUE,
     num_kv_pages_per_block: int = 128,
     num_queries_per_block: int = 128,
+    sm_scale: float = 1.0,
 ) -> jax.Array:
   """Paged attention kernel with ragged input.
 
@@ -940,6 +946,7 @@ def next_kv_blk_page_indices_index_map(kv_head_idx, logical_q_blk_idx,
           num_seqs=num_seqs,
           num_kv_pages_per_block=num_kv_pages_per_block,
           mask_value=mask_value,
+          sm_scale=sm_scale,
       ),
       grid_spec=pltpu.PrefetchScalarGridSpec(
           num_scalar_prefetch=6,