upd

Author: yzh119

csrc/trtllm_fmha_kernel_launcher.cu

@@ -232,11 +232,11 @@ void trtllm_paged_attention_decode(TensorView out, Optional<TensorView> out_scal
   bool is_shared_kv = key_cache.data_ptr() == value_cache.data_ptr();
   int num_pages_in_mem_pool = is_shared_kv ? key_cache.size(0) : key_cache.size(0) * 2;
 
-  // Assume NHD layout: [..., N, H, D]
-  int page_size = key_cache.size(-3);
-  int num_kv_heads = key_cache.size(-2);
-  int kv_stride_keys_values = key_cache.stride(-3);  // key/values
-  int kv_stride_heads = key_cache.stride(-2);        // head
+  // Assume NHD layout: [..., H, N, D]
+  int page_size = key_cache.size(-2);
+  int num_kv_heads = key_cache.size(-3);
+  int kv_stride_keys_values = key_cache.stride(-2);  // key/values
+  int kv_stride_heads = key_cache.stride(-3);        // head
 
   int kv_stride_batch = key_cache.stride(0);  // batch
 
@@ -294,11 +294,11 @@ void trtllm_paged_attention_context(TensorView out, Optional<TensorView> out_sca
   bool is_shared_kv = key_cache.data_ptr() == value_cache.data_ptr();
   int num_pages_in_mem_pool = is_shared_kv ? key_cache.size(0) : key_cache.size(0) * 2;
 
-  // Assume NHD layout: [..., N, H, D]
-  int page_size = key_cache.size(-3);
-  int num_kv_heads = key_cache.size(-2);
-  int kv_stride_keys_values = key_cache.stride(-3);  // key/values
-  int kv_stride_heads = key_cache.stride(-2);        // head
+  // Assume NHD layout: [..., H, N, D]
+  int page_size = key_cache.size(-2);
+  int num_kv_heads = key_cache.size(-3);
+  int kv_stride_keys_values = key_cache.stride(-2);  // key/values
+  int kv_stride_heads = key_cache.stride(-3);        // head
   int kv_stride_batch = key_cache.stride(0);         // batch
 
   const auto stream = get_stream(query.device());

flashinfer/decode.py

@@ -1235,9 +1235,9 @@ def run(
             q, k_cache, self._cached_q_data_type, self._cached_kv_data_type
         )
 
-        # Convert HND layout to NHD for trtllm-gen backend
-        if self._backend == "trtllm-gen" and self._kv_layout == "HND":
-            # For HND: [..., H, N, D] -> NHD: [..., N, H, D]
+        # Convert NHD layout to HND for trtllm-gen backend
+        if self._backend == "trtllm-gen" and self._kv_layout == "NHD":
+            # For NHD: [..., N, H, D] -> HND: [..., H, N, D]
             k_cache = k_cache.transpose(-3, -2)
             v_cache = v_cache.transpose(-3, -2)
 
@@ -2198,9 +2198,9 @@ def trtllm_batch_decode_with_kv_cache(
             q_len_per_req=q_len_per_req,
         )
     elif backend == "trtllm-gen":
-        # Convert HND layout to NHD if necessary (transpose only changes stride, not data)
-        if kv_layout == "HND":
-            # For HND: [..., H, N, D] -> NHD: [..., N, H, D]
+        # Convert NHD layout to HND if necessary (transpose only changes stride, not data)
+        if kv_layout == "NHD":
+            # For NHD: [..., N, H, D] -> HND: [..., H, N, D]
             k_cache = k_cache.transpose(-3, -2)
             v_cache = v_cache.transpose(-3, -2)
 
@@ -2431,7 +2431,9 @@ def xqa_batch_decode_with_kv_cache(
         page_size = k_cache.shape[2]
         head_dim = k_cache.shape[3]
 
-    workspace_0, workspace_1 = torch.chunk(workspace_buffer, 2, dim=0)
+    workspace_u8 = workspace_buffer.view(torch.uint8)
+    semaphore = workspace_u8[: round_up(4 * sm_count, 16)]
+    scratch = workspace_u8[round_up(4 * sm_count, 16) :]
     kv_scale_value = bmm2_scale
     q_scale_value = bmm1_scale / kv_scale_value * (head_dim**0.5)
 
@@ -2448,8 +2450,8 @@ def xqa_batch_decode_with_kv_cache(
         block_tables,
         seq_lens_new,
         out,
-        workspace_0,
-        workspace_1,
+        scratch,
+        semaphore,
         num_kv_heads,
         page_size,
         sinks=sinks_new,
@@ -2571,6 +2573,10 @@ def trtllm_batch_decode_with_kv_cache_mla(
     ):  # todo(Yingyi): add support for more block sizes?
         raise ValueError(f"Supported block_size are 32 and 64, got {block_size}")
 
+    print(
+        f"Running TRTLLM batch decode with KV cache: {query.shape}, {kv_cache.shape}, {workspace_buffer.shape}"
+    )
+
     _check_trtllm_gen_mla_shape(
         query,
         kv_cache,

flashinfer/prefill.py

@@ -2088,9 +2088,9 @@ def run(
                 out, q.shape[:-1] + v_cache.shape[-1:], q.dtype, q.device, "out"
             )
 
-        # Convert HND layout to NHD for trtllm-gen backend
-        if self._backend == "trtllm-gen" and self._kv_layout == "HND":
-            # For HND: [..., H, N, D] -> NHD: [..., N, H, D]
+        # Convert NHD layout to HND for trtllm-gen backend
+        if self._backend == "trtllm-gen" and self._kv_layout == "NHD":
+            # For NHD: [..., N, H, D] -> HND: [..., H, N, D]
             k_cache = k_cache.transpose(-3, -2)
             v_cache = v_cache.transpose(-3, -2)
 
@@ -3411,9 +3411,9 @@ def trtllm_batch_context_with_kv_cache(
             # it doesn't change underlying storage
             k_cache, v_cache = kv_cache.unbind(dim=1)
 
-    # Convert HND layout to NHD if necessary (transpose only changes stride, not data)
-    if kv_layout == "HND":
-        # For HND: [..., H, N, D] -> NHD: [..., N, H, D]
+    # Convert NHD layout to HND if necessary (transpose only changes stride, not data)
+    if kv_layout == "NHD":
+        # For NHD: [..., N, H, D] -> HND: [..., H, N, D]
         k_cache = k_cache.transpose(-3, -2)
         v_cache = v_cache.transpose(-3, -2)