fix

Author: yzh119

flashinfer/decode.py

@@ -984,7 +984,6 @@ def plan(
         else:
             kv_lens_arr_host = seq_lens.cpu()
         if self._backend == "trtllm-gen":
-            assert self._kv_layout == "HND"
             assert logits_soft_cap == 0.0
             self._max_kv_len = max(kv_lens_arr_host).item()
             self._kv_lens_buffer[: len(kv_lens_arr_host)].copy_(
@@ -1227,6 +1226,7 @@ def run(
         if enable_pdl is None:
             enable_pdl = device_support_pdl(q.device)
         k_cache, v_cache = _unpack_paged_kv_cache(paged_kv_cache, self._kv_layout)
+
         if self._kv_layout == "NHD":
             page_size = k_cache.shape[1]
         else:
@@ -1235,6 +1235,12 @@ 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]
+            k_cache = k_cache.transpose(-3, -2)
+            v_cache = v_cache.transpose(-3, -2)
+
         pos_encoding_mode = self._pos_encoding_mode
         window_left = self._window_left if window_left is None else window_left
         if self._backend != "trtllm-gen":
@@ -1997,7 +2003,6 @@ def paged_run(
             1.0,  # NOTE(Siyuan): update this to expose bmm2 scale
             workspace_size,
             window_left,
-            layout,
             enable_pdl,
             out=o,
             sinks=sinks,

flashinfer/prefill.py

@@ -582,7 +582,6 @@ def paged_run(
                 batch_size,
                 cum_seq_lens_q,
                 cum_seq_lens_kv,
-                layout,
                 enable_pdl,
                 workspace_size,
                 window_left,
@@ -2041,6 +2040,7 @@ def run(
         _check_cached_qkv_data_type(
             q, k_cache, self._cached_q_data_type, self._cached_kv_data_type
         )
+
         stride_block = k_cache.stride(0)
         if self._kv_layout == "NHD":
             page_size = k_cache.shape[1]
@@ -2088,6 +2088,12 @@ 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]
+            k_cache = k_cache.transpose(-3, -2)
+            v_cache = v_cache.transpose(-3, -2)
+
         if self._custom_mask_buf is not None:
             mask_mode = MaskMode.CUSTOM.value
         else:

tests/attention/test_trtllm_gen_attention.py

@@ -809,8 +809,9 @@ def test_trtllm_batch_decode(
         out_dtype=out_dtype,
         o_sf_scale=o_sf_scale,
         o_sf_vec_size=o_sf_vec_size,
-        enable_pdl=enable_pdl,
         sinks=(sink if enable_sink else None),
+        kv_layout=kv_layout,
+        enable_pdl=enable_pdl,
         q_len_per_req=q_len_per_req,
     )
     # check if the first 8192 * 256 * 4 bytes of workspace_buffer is zero