wip

wip

Author: yaoyaoding

benchmarks/routines/attention.py

@@ -526,7 +526,7 @@ def run_backend_wrapper(backend):
                 workspace_buffer=workspace_buffer,
                 block_tables=block_tables,
                 seq_lens=actual_seq_lens_kv,
-                max_seq_len=s_kv,
+                max_kv_len=s_kv,
                 bmm1_scale=scale if k_scale is None else k_scale * scale,
                 bmm2_scale=1.0 if v_scale is None else v_scale,
             )

csrc/trtllm_fmha_kernel_launcher.cu

@@ -157,6 +157,9 @@ void trtllm_paged_attention_launcher(
         use_multi_block ? TileScheduler::Static : TileScheduler::Persistent;
     runner_params.mMultiCtasKvMode = use_multi_block;
 
+    runner_params.cumSeqLensQPtr = cum_seq_lens_q;
+    runner_params.cumSeqLensKvPtr = cum_seq_lens_kv;
+
     size_t max_batch_size = 8192;   // todo(Yingyi): get from dlfw
     size_t max_num_qo_heads = 256;  // todo(Yingyi): get from dlfw, in total 8MB
     size_t num_semaphores =
@@ -204,26 +207,26 @@ inline bool is_4bit(Data_type data_type) { return data_type == Data_type::DATA_T
 void trtllm_paged_attention_decode(TensorView out, Optional<TensorView> out_scale_factor,
                                    TensorView query, TensorView key_cache, TensorView value_cache,
                                    TensorView workspace_buffer, TensorView block_tables,
-                                   TensorView seq_lens, int64_t max_kv_len,
-                                   Variant<double, ffi::Tensor> bmm1_scale,
-                                   Variant<double, ffi::Tensor> bmm2_scale, double o_sf_scale,
-                                   int64_t o_sf_vec_size, int64_t o_sf_start_index,
-                                   int64_t window_left, int64_t sm_count, bool enable_pdl,
-                                   int64_t workspace_size, Optional<TensorView> attention_sinks) {
+                                   TensorView seq_lens, int64_t max_q_len, int64_t max_kv_len,
+                                   Variant<double, ffi::Tensor> bmm1_scale, Variant<double, ffi::Tensor> bmm2_scale, 
+                                   double o_sf_scale, int64_t o_sf_vec_size, int64_t o_sf_start_index, int64_t batch_size,
+                                   int64_t window_left, TensorView cum_seq_lens_q, TensorView cum_seq_lens_kv, int64_t sm_count, 
+                                   bool enable_pdl, int64_t workspace_size, Optional<TensorView> attention_sinks) {
   auto q_data_type = dl_dtype_to_tllm_data_type(query.dtype());
   auto kv_data_type = dl_dtype_to_tllm_data_type(key_cache.dtype());
   TVM_FFI_ICHECK_EQ(key_cache.ndim(), value_cache.ndim());
   for (int i = 0; i < key_cache.ndim(); i++) {
     TVM_FFI_ICHECK_EQ(key_cache.size(i), value_cache.size(i));
   }
   auto o_data_type = dl_dtype_to_tllm_data_type(out.dtype());
-  // NOTE(Zihao): query is [B, Q, H, D]
-  // where Q is the number of query tokens per request, used in MTP
+  // NOTE(Zihao): query is [S, H, D]
+  // where S is the sum of query tokens for all requests, used in MTP
   // based on profiled results, always use decode mode for MTP (q_len is small)
   // example: when kv_len = 10000, q < 200, decode mode is faster
-  int batch_size = query.size(0);
-  int q_len_per_request = query.size(1);
-  int sum_seq_q = batch_size * q_len_per_request;
+  // int batch_size = query.size(0);
+  // int q_len_per_request = query.size(1);
+  // int sum_seq_q = batch_size * q_len_per_request;
+  int sum_seq_q = query.size(0);
   int num_qo_heads = query.size(2);
   // Multiply by two for FP4 tensor as it is stored as UINT8 dtype. Assume the dim is even.
   int head_dim_k = is_4bit(kv_data_type) ? key_cache.size(-1) * 2 : key_cache.size(-1);
@@ -281,14 +284,14 @@ void trtllm_paged_attention_decode(TensorView out, Optional<TensorView> out_scal
       out.data_ptr(), output_sf_ptr, query.data_ptr(), key_cache.data_ptr(), value_cache.data_ptr(),
       workspace_buffer.data_ptr(), static_cast<int*>(block_tables.data_ptr()),
       static_cast<int*>(seq_lens.data_ptr()),
-      /*cum_seq_lens_q=*/nullptr,
-      /*cum_seq_lens_kv=*/nullptr, attention_sinks_ptr, q_data_type, kv_data_type, o_data_type,
-      TllmPagedAttentionMode::ForGen, batch_size, /*max_q_len=*/q_len_per_request, max_kv_len,
-      num_pages_in_mem_pool, num_qo_heads, num_kv_heads, head_dim_q, head_dim_o, page_size,
-      kv_stride_keys_values, kv_stride_heads, kv_stride_batch, max_num_blocks_per_seq,
-      bmm1_scale_value, bmm2_scale_value, bmm1_scale_log2_ptr, bmm2_scale_ptr, o_sf_scale,
-      o_sf_vec_size, o_sf_start_index, window_left, sum_seq_q, sm_count, enable_pdl, workspace_size,
-      stream);
+      /*cum_seq_lens_q=*/static_cast<int*>(cum_seq_lens_q.data_ptr()),
+      /*cum_seq_lens_kv=*/static_cast<int*>(cum_seq_lens_kv.data_ptr()), attention_sinks_ptr, 
+      q_data_type, kv_data_type, o_data_type, TllmPagedAttentionMode::ForGen, batch_size, 
+      max_q_len, max_kv_len, num_pages_in_mem_pool, num_qo_heads, num_kv_heads, head_dim_q, 
+      head_dim_o, page_size, kv_stride_keys_values, kv_stride_heads, kv_stride_batch, 
+      max_num_blocks_per_seq, bmm1_scale_value, bmm2_scale_value, bmm1_scale_log2_ptr, 
+      bmm2_scale_ptr, o_sf_scale, o_sf_vec_size, o_sf_start_index, window_left, sum_seq_q, sm_count, 
+      enable_pdl, workspace_size, stream);
 }
 
 void trtllm_paged_attention_context(

flashinfer/decode.py

@@ -2065,26 +2065,30 @@ def trtllm_batch_decode_with_kv_cache(
     workspace_buffer: torch.Tensor,
     block_tables: torch.Tensor,
     seq_lens: torch.Tensor,
-    max_seq_len: int,
-    bmm1_scale: Union[float, torch.Tensor] = 1.0,
-    bmm2_scale: Union[float, torch.Tensor] = 1.0,
+    max_q_len: int,
+    max_kv_len: int,
+    bmm1_scale: Union[float, torch.Tensor],
+    bmm2_scale: Union[float, torch.Tensor],
+    batch_size: int,
+    cum_seq_lens_q: torch.Tensor,
+    cum_seq_lens_kv: torch.Tensor,
     window_left: int = -1,
     out: Optional[Union[torch.Tensor, FP4Tensor]] = None,
     out_dtype: Optional[Union[torch.dtype, str]] = None,
     o_sf_scale: Optional[float] = None,
     o_sf_vec_size: Optional[int] = None,
-    sinks: Optional[List[torch.Tensor]] = None,
     kv_layout: str = "HND",
     enable_pdl: Optional[bool] = None,
+    sinks: Optional[List[torch.Tensor]] = None,
     backend: str = "auto",
-    q_len_per_req: Optional[int] = 1,
+    # the following args are xqa-specific
     o_scale: Optional[float] = 1.0,
 ) -> Union[torch.Tensor, FP4Tensor]:
     """
     Parameters
     ----------
     query : torch.Tensor
-        query tensor with shape [num_tokens, num_heads, head_dim], num_tokens = batch_size * q_len_per_request
+        query tensor with shape [num_tokens, num_heads, head_dim], num_tokens = total query tokens in the batch.
 
     kv_cache : Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]
         If kv_cache is a single tensor, it should be a tensor with shape [num_pages, 1 or 2, num_kv_heads, page_size, head_dim] if :attr:`kv_layout` is ``HND``,
@@ -2185,6 +2189,8 @@ def trtllm_batch_decode_with_kv_cache(
             raise ValueError("xqa backend does not support nvfp4 output")
         if o_sf_scale is not None or o_sf_vec_size is not None:
             raise ValueError("xqa backend does not support o_sf_scale or o_sf_vec_size")
+        if max_q_len != 1:
+            raise ValueError("xqa backend only supports max_q_len == 1")
 
         # Handle out and out_dtype
         if out_dtype is None:
@@ -2199,15 +2205,15 @@ def trtllm_batch_decode_with_kv_cache(
             workspace_buffer=workspace_buffer,
             block_tables=block_tables,
             seq_lens=seq_lens,
-            max_seq_len=max_seq_len,
+            max_seq_len=max_kv_len,
             bmm1_scale=bmm1_scale,
             bmm2_scale=bmm2_scale,
             window_left=window_left,
             out=out,
             sinks=sinks,
             kv_layout=kv_layout,
             enable_pdl=enable_pdl,
-            q_len_per_req=q_len_per_req,
+            q_len_per_req=1,
             o_scale=o_scale,
         )
     elif backend == "trtllm-gen":
@@ -2299,31 +2305,30 @@ def trtllm_batch_decode_with_kv_cache(
             bmm1_scale = bmm1_scale * log2e
         if isinstance(bmm2_scale, torch.Tensor):
             assert bmm2_scale.dtype == torch.float32
-
+        workspace_size = workspace_buffer.numel() * workspace_buffer.element_size()
         run_func(
             out,
             out_scale_factor,
-            query.view(
-                query.size(0) // q_len_per_req,
-                q_len_per_req,
-                query.size(1),
-                query.size(2),
-            ),
+            query,
             k_cache,
             v_cache,
             workspace_buffer,
             block_tables,
             seq_lens,
-            max_seq_len,
+            max_q_len,
+            max_kv_len,
             bmm1_scale,
             bmm2_scale,
             o_sf_scale or -1.0,
             o_sf_vec_size or -1,
             o_sf_start_index,
+            batch_size,
             window_left,
+            cum_seq_lens_q,
+            cum_seq_lens_kv,
             sm_count,
             enable_pdl,
-            workspace_buffer.numel() * workspace_buffer.element_size(),
+            workspace_size,
             sinks,
         )
 

tests/attention/test_trtllm_gen_attention.py

@@ -54,12 +54,12 @@ def generate_seq_lens_prefill(batch_size, max_q_len, max_in_kv_len):
     return q_lens, in_kv_lens, seq_lens
 
 
-def generate_seq_lens_decode(batch_size, q_len_per_req, max_in_kv_len):
-    q_lens = torch.full((batch_size,), q_len_per_req, dtype=torch.int32)
-    in_kv_lens = torch.randint(0, max_in_kv_len + 1, (batch_size,), dtype=torch.int)
-    in_kv_lens[-1] = max_in_kv_len
-    seq_lens = q_lens + in_kv_lens
-    return q_lens, in_kv_lens, seq_lens
+def generate_seq_lens_decode(batch_size, max_q_len, max_kv_len):
+    q_lens = torch.full((batch_size,), max_q_len, dtype=torch.int32)
+    kv_lens = torch.randint(0, max_kv_len + 1, (batch_size,), dtype=torch.int)
+    kv_lens[-1] = max_kv_len
+    seq_lens = q_lens + kv_lens
+    return q_lens, kv_lens, seq_lens
 
 
 def generate_cumsum_lens(lens):
@@ -667,7 +667,6 @@ def _test_trtllm_batch_decode(
     backend,
     kv_layout,
     batch_size,
-    q_len_per_req,
     page_size,
     num_kv_heads,
     head_grp_size,
@@ -677,9 +676,10 @@ def _test_trtllm_batch_decode(
     kv_dtype,
     enable_pdl,
     enable_sink,
-    max_in_kv_len,
     head_dim,
-    device_scale=False,
+    max_q_len,
+    max_kv_len,
+    device_scale,
 ):
     """
     Common function for testing trtllm-gen decode.
@@ -702,12 +702,12 @@ def _test_trtllm_batch_decode(
         pytest.skip("xqa backend only supports fp16 and bf16 query")
 
     # xqa backend doesn't support speculative decoding yet
-    if backend == "xqa" and q_len_per_req > 1:
+    if backend == "xqa" and max_q_len > 1:
         pytest.skip(
             "xqa backend does not support speculative decoding (q_len_per_req > 1) yet"
         )
 
-    if o_dtype == "nvfp4" and q_len_per_req > 1:
+    if o_dtype == "nvfp4" and max_q_len > 1:
         # todo(Yingyi): add support for nvfp4 with speculative decoding
         pytest.skip("nvfp4 is not supported for q_len_per_req > 1")
 
@@ -719,8 +719,8 @@ def _test_trtllm_batch_decode(
 
     # Generate random sequence lengths
     num_qo_heads = num_kv_heads * head_grp_size
-    q_lens, in_kv_lens, seq_lens = generate_seq_lens_decode(
-        batch_size, q_len_per_req, max_in_kv_len
+    q_lens, kv_lens, seq_lens = generate_seq_lens_decode(
+        batch_size, max_q_len, max_kv_len
     )
 
     # Create query tensor and related data
@@ -775,7 +775,7 @@ def _test_trtllm_batch_decode(
         "window_left": window_left,
     }
     if not enable_sink:
-        if q_len_per_req == 1:
+        if max_q_len == 1:
             wrapper_ref = flashinfer.decode.BatchDecodeWithPagedKVCacheWrapper(
                 workspace_buffer_ref, kv_layout, use_tensor_cores=True
             )
@@ -843,6 +843,8 @@ def _test_trtllm_batch_decode(
         workspace_buffer,
         page_table,
         seq_lens.to(GPU_DEVICE),
+        max_q_len,
+        max_kv_len,
         torch.max(seq_lens).item(),
         bmm1_scale,
         bmm2_scale,
@@ -855,7 +857,6 @@ def _test_trtllm_batch_decode(
         kv_layout=kv_layout,
         enable_pdl=enable_pdl,
         backend=backend,
-        q_len_per_req=q_len_per_req,
         o_scale=o_scale,
     )
     if backend == "trtllm-gen":
@@ -882,7 +883,7 @@ def _test_trtllm_batch_decode(
 
     # convert to float32 for fp8 is not supported by assert_close
     # relax rtol and atol for speculative decoding test
-    if q_len_per_req > 1:
+    if max_q_len > 1:
         rtol, atol = rtol * 2, atol * 2
 
     # Arbitary small mismatch rate
@@ -1224,6 +1225,65 @@ def test_trtllm_batch_decode_long_sequence_length(
     )
 
 
+@pytest.mark.parametrize("kv_layout", ["HND"])  # trtllm-gen only support HND
+@pytest.mark.parametrize(
+    "batch_size,q_len_per_req,page_size,num_kv_heads,head_grp_size",
+    [
+        (1, 1, 16, 8, 8),
+        (1, 1, 32, 8, 8),
+    ],
+)
+@pytest.mark.parametrize("window_left", [-1])
+@pytest.mark.parametrize(
+    "q_dtype,kv_dtype,o_dtype",
+    [
+        ("fp8", "fp8", "fp8"),
+    ],
+)
+@pytest.mark.parametrize("enable_pdl", [None])
+@pytest.mark.parametrize("enable_sink", [False])
+@pytest.mark.parametrize("max_in_kv_len", [4096, 8192])
+@pytest.mark.parametrize("head_dim", [128])
+@pytest.mark.parametrize("device_scale", [True, False])
+def test_trtllm_batch_decode_spec(
+    kv_layout,
+    batch_size,
+    q_len_per_req,
+    page_size,
+    num_kv_heads,
+    head_grp_size,
+    window_left,
+    q_dtype,
+    o_dtype,
+    kv_dtype,
+    enable_pdl,
+    enable_sink,
+    max_in_kv_len,
+    head_dim,
+    device_scale,
+):
+    # Small number of test cases for batch size 1
+    _test_trtllm_batch_decode(
+        "trtllm-gen",
+        kv_layout,
+        batch_size,
+        q_len_per_req,
+        page_size,
+        num_kv_heads,
+        head_grp_size,
+        window_left,
+        q_dtype,
+        o_dtype,
+        kv_dtype,
+        enable_pdl,
+        enable_sink,
+        max_in_kv_len,
+        head_dim,
+        device_scale,
+    )
+
+
+
 @pytest.mark.parametrize("batch_size", [4, 128, 256])
 @pytest.mark.parametrize("s_qo", [32, 64, 87])
 @pytest.mark.parametrize("s_kv", [32, 64, 87])