[#6507][fix] Fix precision issue due to KV layout mismatch for split/concat kernels (#6917)

Signed-off-by: ZhangGe6 <sjtu.zg123@gmail.com>
Co-authored-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>

Author: ZhangGe6

tensorrt_llm/_torch/attention_backend/flashinfer.py

@@ -56,7 +56,10 @@ class FlashInferWrappers:
 class FlashInferAttentionMetadata(AttentionMetadata):
     workspace_buffer: Optional[torch.Tensor] = None
 
-    kv_layout: Literal["NHD", "HND"] = "NHD"
+    # cache concat/split kernels when using PD disaggregation
+    # expects KV cache in [max_num_pages, 2, num_kv_heads, page_size, head_dim] layout,
+    # so set kv_layout as "HND" here
+    kv_layout: Literal["NHD", "HND"] = "HND"
 
     paged_kv_indptr_decode: torch.Tensor = field(init=False)
     paged_kv_indptr_prefill: torch.Tensor = field(init=False)
@@ -506,7 +509,8 @@ def forward_impl(
         q = q.view(-1, self.num_heads, self.head_dim)
 
         # Key and Value
-        kv_cache = metadata.kv_cache_manager.get_buffers(self.layer_idx)
+        kv_cache = metadata.kv_cache_manager.get_buffers(
+            self.layer_idx, kv_layout=metadata.kv_layout)
 
         if k is not None and v is not None:
             k = k.view(-1, self.num_kv_heads, self.head_dim)

tensorrt_llm/_torch/attention_backend/star_flashinfer.py

@@ -331,7 +331,8 @@ def forward(self,
         num_ctx_tokens = metadata.num_ctx_tokens
         num_qry_tokens = metadata.num_qry_tokens
 
-        kv_cache = metadata.kv_cache_manager.get_buffers(self.layer_idx)
+        kv_cache = metadata.kv_cache_manager.get_buffers(
+            self.layer_idx, kv_layout=metadata.kv_layout)
         if self.quant_config and self.quant_config.layer_quant_mode.has_any_quant(
         ):
             qc = self.quant_config

tensorrt_llm/_torch/pyexecutor/resource_manager.py

@@ -813,16 +813,43 @@ def get_num_available_tokens(self, max_num_draft_tokens: int = 0) -> int:
         return (self.get_num_free_blocks() * self.tokens_per_block -
                 self.num_extra_kv_tokens - max_num_draft_tokens)
 
-    def get_buffers(self, layer_idx: int) -> Optional[torch.Tensor]:
+    def get_buffers(self,
+                    layer_idx: int,
+                    kv_layout: str = "NHD") -> Optional[torch.Tensor]:
+        ''' Slice KV tensor for a specified layer and reshape it.
+
+        1. Slice:
+            [max_num_pages, num_layers, kv_factor, page_size * num_kv_heads * head_dim] ->
+            [max_num_pages, kv_factor, page_size * num_kv_heads * head_dim]
+
+        2. Reshape:
+            kv_layout = "NHD" -> [max_num_pages, kv_factor, page_size, num_kv_heads, head_dim]
+            kv_layout = "HND" -> [max_num_pages, kv_factor, num_kv_heads, page_size, head_dim]
+
+        Note that different attention backend/implementation can have different KV layouts,
+        "kv_layout" should be set accordingly to avoid surprises.
+        '''
         layer_offset = self.layer_offsets[layer_idx]
         result = self.impl.get_primary_pool_data(layer_offset)
-        return result.reshape(
-            result.shape[0],
-            self.kv_factor,
-            self.tokens_per_block,
-            self.num_kv_heads_per_layer[layer_offset],
-            self.head_dim,
-        )
+
+        assert kv_layout in ["NHD",
+                             "HND"], f"Unsupported kv_layout: {kv_layout}"
+        if kv_layout == "NHD":
+            return result.reshape(
+                result.shape[0],
+                self.kv_factor,
+                self.tokens_per_block,
+                self.num_kv_heads_per_layer[layer_offset],
+                self.head_dim,
+            )
+        else:
+            return result.reshape(
+                result.shape[0],
+                self.kv_factor,
+                self.num_kv_heads_per_layer[layer_offset],
+                self.tokens_per_block,
+                self.head_dim,
+            )
 
     def get_indexer_k_cache_pool_data(self, layer_idx: int) -> torch.Tensor:
         result = self.impl.get_indexer_k_cache_pool_data(layer_idx)

tests/unittest/_torch/attention/test_attention.py

@@ -438,12 +438,13 @@ def test_attention_backend(s: Scenario):
     flashinfer_kv_cache = torch.randn(num_layers,
                                       s.max_num_pages,
                                       2,
-                                      page_size,
                                       num_kv_heads,
+                                      page_size,
                                       head_dim,
                                       device="cuda").to(s.kvcache_dtype)
-    ref_kv_cache = flashinfer_kv_cache.transpose(1, 2).contiguous().view(
-        num_layers, 2, batch_size, kv_cache_len, num_kv_heads, head_dim)
+    ref_kv_cache = flashinfer_kv_cache.transpose(1, 2).transpose(
+        3, 4).contiguous().view(num_layers, 2, batch_size, kv_cache_len,
+                                num_kv_heads, head_dim)
     kv = torch.randn(num_layers,
                      2,
                      nnz_kv,
@@ -588,12 +589,13 @@ def test_attention_backend_ifb(s: PagedScenario):
     flashinfer_kv_cache = torch.randn(num_layers,
                                       s.max_num_pages,
                                       2,
-                                      page_size,
                                       num_kv_heads,
+                                      page_size,
                                       head_dim,
                                       device="cuda").to(s.kvcache_dtype)
-    ref_kv_cache = flashinfer_kv_cache.transpose(1, 2).contiguous().view(
-        num_layers, 2, batch_size, kv_cache_len, num_kv_heads, head_dim)
+    ref_kv_cache = flashinfer_kv_cache.transpose(1, 2).transpose(
+        3, 4).contiguous().view(num_layers, 2, batch_size, kv_cache_len,
+                                num_kv_heads, head_dim)
     vanilla_kv_cache = ref_kv_cache.transpose(1, 2).contiguous()
     kv = torch.randn(num_layers,
                      2,

tests/unittest/_torch/attention/test_flashinfer_attention.py

@@ -227,7 +227,10 @@ def test_flashinfer_attention(self, scenario: Scenario):
                              sum(context_sequence_lengths) + num_gens)
 
             # validate kv cache was updated expectedly
-            cache_buf = kv_cache_manager.get_buffers(flashinfer_attn.layer_idx)
+            cache_buf = kv_cache_manager.get_buffers(
+                flashinfer_attn.layer_idx, kv_layout=attn_metadata.kv_layout)
+            if attn_metadata.kv_layout == "HND":
+                cache_buf = cache_buf.transpose(2, 3).contiguous()
             assert cache_buf is not None
             num_kv_heads = cache_buf.size(-2)
 

tests/unittest/_torch/attention/test_flashinfer_star_attn.py

@@ -312,7 +312,10 @@ def test_flashinfer_star_attention(self, scenario: Scenario):
                 num_gens)
 
             # validate kv cache was updated expectedly
-            cache_buf = kv_cache_manager.get_buffers(star_attn.layer_idx)
+            cache_buf = kv_cache_manager.get_buffers(
+                star_attn.layer_idx, kv_layout=attn_metadata.kv_layout)
+            if attn_metadata.kv_layout == "HND":
+                cache_buf = cache_buf.transpose(2, 3).contiguous()
             assert cache_buf is not None
             num_kv_heads = cache_buf.size(-2)