fa3 cudagraph

Author: RunningLeon

lmdeploy/pytorch/backends/cuda/attention.py

@@ -42,6 +42,9 @@ class TritonAttentionMetadata(AttentionMetadata):
     num_splits: torch.Tensor = None
     cu_seqlens_q: torch.Tensor = None
     cu_seqlens_k: torch.Tensor = None
+    # flash attn
+    scheduler_metadata: torch.Tensor = None
+    max_kv_seqlen: int = None
 
 
 def _cdiv(a, b):
@@ -477,42 +480,41 @@ def forward(
                 v_scales_zeros=v_scales_zeros,
                 quant_policy=quant_policy,
             )
-        # sliding_window = (-1, -1) if self.sliding_window is None else self.sliding_window
-        # if isinstance(sliding_window, int):
-        #     sliding_window = (sliding_window, sliding_window)
-        # attn_output = self.flash_attn_with_kvcache_v3(
-        #     query,
-        #     k_cache,
-        #     v_cache,
-        #     cache_seqlens=attn_metadata.kv_seqlens.to(torch.int32),
-        #     cu_seqlens_q=attn_metadata.cu_seqlens_q,
-        #     cu_seqlens_k_new=attn_metadata.cu_seqlens_k,
-        #     max_seqlen_q=max_q_seqlen,
-        #     page_table=block_offsets,
-        #     softmax_scale=self.scale,
-        #     causal=self.causal,
-        #     window_size=sliding_window,
-        #     softcap=-1.0 if self.logit_softcapping is None else self.logit_softcapping,
-        # )
-        # return attn_output
         if is_decoding:
-            q_shape = query.shape
-            o_shape = q_shape[:-1] + (self.v_head_size, )
-            attn_output = query.new_empty(o_shape)
-            self.paged_attention_fwd(
+            sliding_window = (-1, -1) if self.sliding_window is None else self.sliding_window
+            if isinstance(sliding_window, int):
+                sliding_window = (sliding_window, sliding_window)
+            query = query.unflatten(0, (-1, max_q_seqlen))
+            attn_output = self.flash_attn_with_kvcache_v3(
                 query,
                 k_cache,
                 v_cache,
-                attn_output,
-                block_offsets,
-                kv_seqlens=kv_seqlens,
-                k_scales_zeros=k_scales_zeros,
-                v_scales_zeros=v_scales_zeros,
-                quant_policy=quant_policy,
-                window_size=self.sliding_window,
-                sm_scale=self.scale,
-                logit_softcapping=self.logit_softcapping,
+                cache_seqlens=attn_metadata.kv_seqlens.to(torch.int32),
+                max_seqlen_q=max_q_seqlen,
+                scheduler_metadata=attn_metadata.scheduler_metadata,
+                page_table=block_offsets,
+                softmax_scale=self.scale,
+                causal=self.causal,
+                window_size=sliding_window,
+                softcap=-1.0 if self.logit_softcapping is None else self.logit_softcapping,
             )
+            # q_shape = query.shape
+            # o_shape = q_shape[:-1] + (self.v_head_size, )
+            # attn_output = query.new_empty(o_shape)
+            # self.paged_attention_fwd(
+            #     query,
+            #     k_cache,
+            #     v_cache,
+            #     attn_output,
+            #     block_offsets,
+            #     kv_seqlens=kv_seqlens,
+            #     k_scales_zeros=k_scales_zeros,
+            #     v_scales_zeros=v_scales_zeros,
+            #     quant_policy=quant_policy,
+            #     window_size=self.sliding_window,
+            #     sm_scale=self.scale,
+            #     logit_softcapping=self.logit_softcapping,
+            # )
         else:
             # sliding_window = (-1, -1) if self.sliding_window is None else self.sliding_window
             # if isinstance(sliding_window, int):

lmdeploy/pytorch/backends/cuda/op_backend.py

@@ -1,5 +1,5 @@
 # Copyright (c) OpenMMLab. All rights reserved.
-from typing import Tuple
+from typing import Optional, Tuple
 
 import torch
 
@@ -19,6 +19,46 @@ def _get_meta_flashmla(kv_seqlens, num_attention_heads):
     return tile_scheduler_metadata, num_splits
 
 
+def _get_meta_flashattn(
+        batch_size: int,
+        max_seqlen_q: int,
+        max_seqlen_k: int,
+        num_heads_q: int,
+        num_heads_kv: int,
+        headdim: int,
+        cache_seqlens: torch.Tensor,
+        qkv_dtype=torch.bfloat16,
+        headdim_v=None,
+        cu_seqlens_q: Optional[torch.Tensor] = None,
+        cu_seqlens_k_new: Optional[torch.Tensor] = None,
+        page_size: Optional[int] = None,
+        causal=True,
+        window_size=(-1, -1),  # -1 means infinite context window
+        num_splits=0,
+):
+    """Get scheduler metadata for flash attn."""
+    from flash_attn_interface import get_scheduler_metadata
+
+    metadata = get_scheduler_metadata(
+        batch_size,
+        max_seqlen_q,
+        max_seqlen_k,
+        num_heads_q,
+        num_heads_kv,
+        headdim,
+        cache_seqlens,
+        qkv_dtype=qkv_dtype,
+        headdim_v=headdim_v,
+        cu_seqlens_q=cu_seqlens_q,
+        cu_seqlens_k_new=cu_seqlens_k_new,
+        page_size=page_size,
+        causal=causal,
+        window_size=window_size,
+        num_splits=num_splits,
+    )
+    return metadata
+
+
 class CudaOpsBackend(DefaultOpsBackend):
     """Cuda layer backend."""
 
@@ -121,6 +161,28 @@ def update_meta_flashmla(cls, attn_metadata, num_attention_heads):
         if attn_metadata.block_offsets.dtype != torch.int32:
             attn_metadata.block_offsets = attn_metadata.block_offsets.to(torch.int32)
 
+    @classmethod
+    def update_meta_flashattn(cls, attn_metadata, step_context):
+        batch_size = attn_metadata.q_seqlens.size(0)
+        max_seqlen_q = step_context.input_ids.size(1) // batch_size
+        block_size = step_context.kv_caches[0][0].size(1)
+        window_size = (step_context.model_config.sliding_window, ) * 2
+        scheduler_metadata = _get_meta_flashattn(
+            batch_size=batch_size,
+            max_seqlen_q=max_seqlen_q,
+            max_seqlen_k=step_context.max_kv_seqlen,
+            num_heads_q=step_context.model_config.num_attention_heads,
+            num_heads_kv=step_context.model_config.num_key_value_heads,
+            headdim=step_context.model_config.head_dim,
+            cache_seqlens=attn_metadata.kv_seqlens.to(torch.int32),
+            qkv_dtype=step_context.model_config.dtype,
+            page_size=block_size,
+            window_size=window_size,
+        )
+        attn_metadata.scheduler_metadata = scheduler_metadata
+        attn_metadata.max_kv_seqlen = step_context.max_kv_seqlen
+        return attn_metadata
+
     @classmethod
     def update_step_context(cls, step_context):
         """Update step context."""
@@ -135,9 +197,10 @@ def update_step_context(cls, step_context):
         cu_seqlens_q = None
         cu_seqlens_k = None
         if use_flash_mla or use_flash_attn3:
-            cu_seqlens_q = torch.nn.functional.pad(torch.cumsum(q_seqlens, dim=0, dtype=torch.int32), (1, 0))
-            cu_seqlens_k = torch.nn.functional.pad(torch.cumsum(kv_seqlens, dim=0, dtype=torch.int32), (1, 0))
             step_context.block_offsets = step_context.block_offsets.to(torch.int32)
+            if not step_context.is_decoding:
+                cu_seqlens_q = torch.nn.functional.pad(torch.cumsum(q_seqlens, dim=0, dtype=torch.int32), (1, 0))
+                cu_seqlens_k = torch.nn.functional.pad(torch.cumsum(kv_seqlens, dim=0, dtype=torch.int32), (1, 0))
 
         if not step_context.is_decoding:
             kv_start_loc = kv_seqlens.cumsum(0) - kv_seqlens
@@ -160,6 +223,10 @@ def update_step_context(cls, step_context):
                 cls.update_meta_flashmla(attn_metadata,
                                          step_context.model_config.num_attention_heads * decode_query_len)
 
+        if use_flash_attn3:
+            if step_context.is_decoding is True:
+                attn_metadata = cls.update_meta_flashattn(attn_metadata, step_context)
+
         cross_seqlens = step_context.cross_seqlens
         cross_kv_seqlens = step_context.cross_kv_seqlens
         cross_attn_metadata = None

lmdeploy/pytorch/engine/engine.py

@@ -854,7 +854,6 @@ def update_running_migration(self, running: SeqList, next_token_ids: np.ndarray,
 
     def _debug_spec_stats(self, batched_outputs: BatchedOutputs, is_decoding: bool = False):
         """Make spec stats."""
-        # if self.speculative_config is not None and (debug or self.engine_config.enable_metrics):
         if self.speculative_config is not None:
             if not hasattr(self, 'spec_stats'):
                 from lmdeploy.metrics.stats import SpeculativeDecodingStats
@@ -880,7 +879,7 @@ def _make_infer_outputs(
         logprobs = batched_outputs.logprobs
 
         # for debug
-        debug = True
+        debug = False
         if debug:
             self._debug_spec_stats(batched_outputs, is_decoding=is_decoding)
 
@@ -912,8 +911,13 @@ def _make_infer_outputs(
             cur_logprobs = None
             if num_logprobs >= 0:
                 cur_logprobs = (logprobs.vals[idx, :num_logprobs + 1], logprobs.indices[idx, :num_logprobs + 1])
-
-            req_metrics = RequestMetrics(new_token_timestamp, msg.engine_events, spec_info=None)
+            # get spec stats info
+            spec_info = None
+            if self.speculative_config is not None and is_decoding and self.engine_config.enable_metrics:
+                num_draft_tokens = self.speculative_config.num_speculative_tokens
+                num_accepted_tokens = (batched_outputs.next_token_ids[idx] > -1).sum() - 1
+                spec_info = dict(num_draft_tokens=num_draft_tokens, num_accepted_tokens=num_accepted_tokens)
+            req_metrics = RequestMetrics(new_token_timestamp, msg.engine_events, spec_info=spec_info)
             out = InferOutput(session_id=session_id,
                               resp=msg.resp,
                               finish=finish,

lmdeploy/pytorch/engine/model_agent.py

@@ -514,7 +514,7 @@ def get_output(self):
                                         self._output.numel() // self._output.size(-1),
                                         device=self._output.device,
                                         dtype=self._output.dtype)
-                    return strategy.slice_outputs(self._output, seqlen)
+                    return strategy.slice_outputs(self._output, seqlen), self._aux_output
                 torch.cuda.synchronize()
                 if self._aux_output is not None:
                     self._aux_output = self._aux_output.to(self._device)
@@ -796,13 +796,14 @@ async def __prepare_dp():
                     logger.debug(f'<ForwardTask> rank[{rank}]: Output [{idx}]')
                     extra_outputs = self.agent_strategy.make_extra_outputs(extra_inputs)
                     self._push_output(
-                        BatchedOutputs(next_token_ids=next_token_ids if self.spec_agent is None else extra_inputs.output_token_ids,
-                                       logits=logits if return_logits else None,
-                                       stopped=stopped,
-                                       stop_pos=stop_pos,
-                                       model_metas=model_metas,
-                                       logprobs=logprobs,
-                                       extra_outputs=extra_outputs))
+                        BatchedOutputs(
+                            next_token_ids=next_token_ids if self.spec_agent is None else extra_inputs.output_token_ids,
+                            logits=logits if return_logits else None,
+                            stopped=stopped,
+                            stop_pos=stop_pos,
+                            model_metas=model_metas,
+                            logprobs=logprobs,
+                            extra_outputs=extra_outputs))
             else:
                 # Avoid adding the ADInplaceOrView dispatch key to `next_token_ids`,
                 # as it can trigger recompilation on different ranks when using torch.compile.

lmdeploy/pytorch/model_inputs.py

@@ -314,6 +314,7 @@ class StepContext:
     kv_caches: List
     is_decoding: bool
     sum_kv_seqlen: int
+    max_kv_seqlen: int = None
     local_adapter_ids: torch.LongTensor = None
     input_embeddings: torch.Tensor = None
     input_embedding_indexing: torch.Tensor = None
@@ -388,6 +389,7 @@ def new(
             kv_caches=kv_caches,
             is_decoding=inputs.is_decoding,
             sum_kv_seqlen=inputs.sum_kv_seqlen,
+            max_kv_seqlen=inputs.max_kv_seqlen,
             local_adapter_ids=inputs.local_adapter_ids,
             vision_inputs=inputs.vision_inputs,
             kv_quant_policy=kv_quant_policy,

lmdeploy/pytorch/models/utils/cudagraph.py

@@ -69,8 +69,7 @@ def make_buffers_cudagraph(self, graph_meta: CudaGraphMeta, *args, **kwargs) ->
         seqlens_dtype = torch.int64
         use_flash_mla = getattr(self.config, 'use_flash_mla', False)
         use_flash_attn3 = getattr(self.config, 'use_flash_attn3', False)
-        if use_flash_attn3 and not graph_meta.is_decoding:
-            seqlens_dtype = torch.int32
+
         if use_flash_mla is True:
             import flash_mla
             if graph_meta.is_decoding:
@@ -79,6 +78,9 @@ def make_buffers_cudagraph(self, graph_meta: CudaGraphMeta, *args, **kwargs) ->
             input_buffers['tile_scheduler_metadata'], input_buffers['num_splits'] = flash_mla.get_mla_metadata(
                 torch.ones(max_batches, dtype=torch.int32, device=device),
                 self.config.num_attention_heads * decode_query_len, 1)
+        elif use_flash_attn3 is True:
+            seqlens_dtype = torch.int32
+            input_buffers['scheduler_metadata'] = torch.zeros(max_batches + 1, dtype=torch.int32, device=device)
 
         # flash_mla requires block_offsets and kv_lens int32
         input_buffers['block_offsets'] = torch.zeros((max_batches, num_blocks), dtype=seqlens_dtype, device=device)
@@ -129,7 +131,11 @@ def fill_buffers_cudagraph(self, graph_meta: CudaGraphMeta, input_ids: Tensor, p
         attn_metadata.q_start_loc = input_buffers['q_start_loc']
         attn_metadata.q_seqlens = input_buffers['q_seqlens']
         attn_metadata.kv_seqlens = input_buffers['kv_seqlens']
-        if getattr(self.config, 'use_flash_mla', False) is True:
+
+        use_flash_mla = getattr(self.config, 'use_flash_mla', False)
+        use_flash_attn3 = getattr(self.config, 'use_flash_attn3', False)
+
+        if use_flash_mla is True:
             import flash_mla
             tile_scheduler_metadata, num_splits = flash_mla.get_mla_metadata(
                 attn_metadata.kv_seqlens.to(torch.int32), self.config.num_attention_heads * decode_query_len, 1)
@@ -139,6 +145,25 @@ def fill_buffers_cudagraph(self, graph_meta: CudaGraphMeta, input_ids: Tensor, p
             attn_metadata.tile_scheduler_metadata = input_buffers['tile_scheduler_metadata']
             attn_metadata.num_splits = input_buffers['num_splits']
 
+        if use_flash_attn3:
+            from flash_attn_interface import get_scheduler_metadata
+            block_size = past_key_values[0][0].size(1)
+            # TODO may check tp>1?
+            scheduler_metadata = get_scheduler_metadata(
+                batch_size=batch_size,
+                max_seqlen_q=decode_query_len,
+                max_seqlen_k=attn_metadata.max_kv_seqlen,
+                num_heads_q=self.config.num_attention_heads,
+                num_heads_kv=self.config.num_key_value_heads,
+                headdim=self.config.head_dim,
+                cache_seqlens=attn_metadata.kv_seqlens.to(torch.int32),
+                qkv_dtype=self.config.torch_dtype,
+                page_size=block_size,
+            )
+            input_buffers['scheduler_metadata'].zero_()
+            input_buffers['scheduler_metadata'][:batch_size + 1].copy_(scheduler_metadata[:batch_size + 1])
+            attn_metadata.scheduler_metadata = input_buffers['scheduler_metadata']
+
         new_inputs = dict(
             past_key_values=past_key_values,
             attn_metadata=attn_metadata,