fix None workspace buffer issue

Signed-off-by: elvischenv <[email protected]>

Author: elvischenv

vllm/attention/backends/flashinfer.py

@@ -1163,7 +1163,7 @@ def forward(
                 )
             else:
                 workspace_buffer = (
-                    decode_meta.decode_wrapper._int_workspace_buffer)
+                    decode_meta.decode_wrapper._float_workspace_buffer)
                 assert FlashInferState.get_kv_cache_layout() == "HND"
                 decode_output = trtllm_batch_decode_with_kv_cache(
                     query=decode_query,

vllm/v1/attention/backends/flashinfer.py

@@ -194,7 +194,6 @@ class FlashInferMetadata:
     max_seq_len: int
     seq_lens: torch.Tensor
     block_table_tensor: torch.Tensor
-    workspace_buffer: torch.Tensor
 
     # For handling prefill decode split
     num_decodes: int
@@ -473,7 +472,6 @@ def build(self,
             max_seq_len=max_seq_len,
             seq_lens=seq_lens,
             block_table_tensor=block_table_tensor,
-            workspace_buffer=self._get_workspace_buffer(),
         )
 
         self._plan(num_prefills, num_decodes, attn_metadata)
@@ -641,11 +639,11 @@ def forward(
         if decode_wrapper := attn_metadata.decode_wrapper:
             decode_query = query[:num_decode_tokens]
             assert decode_query.shape[0] == num_decode_tokens
+            assert decode_wrapper is not None
             if not FlashInferBackend.use_trtllm_decode_attention(
                     attn_metadata.num_decodes, attn_metadata.max_seq_len,
                     self.kv_cache_dtype, attn_metadata.num_qo_heads,
                     attn_metadata.num_kv_heads, attn_metadata.head_dim):
-                assert decode_wrapper is not None
                 assert decode_wrapper._window_left == window_left
                 assert decode_wrapper._logits_soft_cap == (self.logits_soft_cap
                                                            or 0.0)
@@ -666,18 +664,20 @@ def forward(
                                                                            num_decode_tokens]
                     seq_lens_decode = attn_metadata.seq_lens[:
                                                              num_decode_tokens]
+                    workspace_buffer = decode_wrapper._float_workspace_buffer
 
                     assert get_kv_cache_layout() == "HND"
                     assert decode_query.is_contiguous()
                     assert kv_cache_permute.is_contiguous()
                     assert block_tables_decode.is_contiguous()
                     assert seq_lens_decode.is_contiguous()
+                    assert workspace_buffer.is_contiguous()
 
                     output[:num_decode_tokens] = (
                         trtllm_batch_decode_with_kv_cache(
                             query=decode_query,
                             kv_cache=kv_cache_permute,
-                            workspace_buffer=attn_metadata.workspace_buffer,
+                            workspace_buffer=workspace_buffer,
                             num_heads=self.num_heads,
                             num_kv_heads=self.num_kv_heads,
                             scale=self.scale,