update

Author: Linboyan-trc

fastdeploy/model_executor/layers/attention/triton_mla_attention_backend.py

@@ -107,7 +107,7 @@ def __init__(
         self.causal: bool = getattr(fd_config.model_config, "causal", True)
 
         self.num_heads: int = num_heads
-        self.head_dim: int = fd_config.model_config.head_dim
+        self.head_dim: int = head_dim
         self.num_layers: int = fd_config.model_config.num_hidden_layers
 
         self.kv_lora_rank: int = fd_config.model_config.kv_lora_rank
@@ -124,7 +124,7 @@ def __init__(
         self.max_kv_splits: int = 32
 
         self.rank, self.device_id = init_rank_and_device_id(fd_config)
-        self.useless_tensor = paddle.randn([1]).cast("int32")
+        self.useless_tensor = paddle.zeros([1], dtype="int32")
 
         # Pre-allocate buffers for CUDAGraph compatibility (stable memory addresses)
         self.max_num_seqs = fd_config.scheduler_config.max_num_seqs
@@ -133,6 +133,12 @@ def __init__(
         self._kv_indices_buf = paddle.zeros([self.max_num_seqs * max_blocks_per_seq * self.block_size], dtype="int32")
         self._num_kv_splits_buf = paddle.ones([self.max_num_seqs], dtype="int32")
 
+        # Pre-allocate decode kernel intermediate buffers for CUDAGraph address stability
+        Lv = fd_config.model_config.kv_lora_rank
+        self._attn_logits_buf = paddle.empty([self.max_num_seqs, num_heads, self.max_kv_splits, Lv], dtype="float32")
+        self._attn_lse_buf = paddle.empty([self.max_num_seqs, num_heads, self.max_kv_splits], dtype="float32")
+        self._o_buf = paddle.empty([self.max_num_seqs, num_heads, Lv], dtype=paddle.get_default_dtype())
+
         if self.flash_attn_func is None:
             prop = paddle.device.cuda.get_device_properties()
             cc = prop.major * 10 + prop.minor
@@ -191,7 +197,10 @@ def init_attention_metadata(self, forward_meta: ForwardMeta):
             total_kv_len = int(paddle.sum(decode_seq_lens).item())
 
             build_kv_indices_from_block_tables(
-                decode_block_tables, decode_seq_lens, self.block_size, decode_bs,
+                decode_block_tables,
+                decode_seq_lens,
+                self.block_size,
+                decode_bs,
                 total_kv_len=total_kv_len,
                 kv_indptr_buf=self._kv_indptr_buf,
                 kv_indices_buf=self._kv_indices_buf,
@@ -200,11 +209,10 @@ def init_attention_metadata(self, forward_meta: ForwardMeta):
             # kv_indptr[decode_bs] = total_kv_len; positions beyond must equal the same
             # so that (kv_indptr[i+1] - kv_indptr[i]) = 0 for padded batches.
             if decode_bs < self.max_num_seqs:
-                self._kv_indptr_buf[decode_bs + 1:] = total_kv_len
+                self._kv_indptr_buf[decode_bs + 1 :] = total_kv_len
 
             # Compute num_kv_splits into the pre-allocated buffer
-            compute_num_kv_splits(decode_seq_lens, decode_bs, self.max_kv_splits,
-                                  out_buf=self._num_kv_splits_buf)
+            compute_num_kv_splits(decode_seq_lens, decode_bs, self.max_kv_splits, out_buf=self._num_kv_splits_buf)
             # Padded entries must be >= 1 to avoid division by zero in kernel
             if decode_bs < self.max_num_seqs:
                 self._num_kv_splits_buf[decode_bs:] = 1
@@ -346,14 +354,15 @@ def _run_decode_kernel(
         latent_dim = self.kv_lora_rank + self.qk_rope_head_dim
         q_reshaped = q.reshape([bs, self.num_heads, latent_dim])
 
-        attn_logits = paddle.empty([bs, self.num_heads, self.max_kv_splits, Lv], dtype="float32")
-        attn_lse = paddle.empty([bs, self.num_heads, self.max_kv_splits], dtype="float32")
-        o = paddle.empty([bs, self.num_heads, Lv], dtype=q.dtype)
+        # Use pre-allocated buffers sliced to current batch size for CUDAGraph address stability
+        attn_logits = self._attn_logits_buf[:bs]
+        attn_lse = self._attn_lse_buf[:bs]
+        o = self._o_buf[:bs]
 
         decode_attention_fwd(
             q_reshaped,
             latent_cache,
-            latent_cache[:, :, :, :self.kv_lora_rank],
+            latent_cache[:, :, :, : self.kv_lora_rank],
             o,
             metadata.kv_indptr,
             metadata.kv_indices,

fastdeploy/model_executor/layers/attention/triton_ops/decode_attention.py

@@ -28,16 +28,9 @@
     enable_compat_on_triton_kernel,
 )
 
-
 _MIN_BLOCK_KV = 32
 
 
-@enable_compat_on_triton_kernel
-@triton.jit
-def tanh(x):
-    return 2 * tl.sigmoid(2 * x) - 1
-
-
 @enable_compat_on_triton_kernel
 @triton.jit
 def _fwd_grouped_kernel_stage1(
@@ -104,13 +97,9 @@ def _fwd_grouped_kernel_stage1(
     if BLOCK_DPE > 0:
         offs_dpe = BLOCK_DMODEL + tl.arange(0, BLOCK_DPE)
         mask_dpe = offs_dpe < Lk
-        off_qpe = (
-            cur_batch * stride_qbs + cur_head[:, None] * stride_qh + offs_dpe[None, :]
-        )
+        off_qpe = cur_batch * stride_qbs + cur_head[:, None] * stride_qh + offs_dpe[None, :]
 
-    kv_len_per_split = (
-        tl.cdiv(tl.cdiv(cur_batch_seq_len, kv_splits), MIN_BLOCK_KV) * MIN_BLOCK_KV
-    )
+    kv_len_per_split = tl.cdiv(tl.cdiv(cur_batch_seq_len, kv_splits), MIN_BLOCK_KV) * MIN_BLOCK_KV
     split_kv_start = kv_len_per_split * split_kv_id
     split_kv_end = tl.minimum(split_kv_start + kv_len_per_split, cur_batch_seq_len)
 
@@ -121,9 +110,7 @@ def _fwd_grouped_kernel_stage1(
     if split_kv_end > split_kv_start:
         q = tl.load(Q + offs_q, mask=(mask_h[:, None]) & (mask_d[None, :]), other=0.0)
         if BLOCK_DPE > 0:
-            qpe = tl.load(
-                Q + off_qpe, mask=(mask_h[:, None]) & (mask_dpe[None, :]), other=0.0
-            )
+            qpe = tl.load(Q + off_qpe, mask=(mask_h[:, None]) & (mask_dpe[None, :]), other=0.0)
         for start_n in range(split_kv_start, split_kv_end, BLOCK_N):
             offs_n = start_n + tl.arange(0, BLOCK_N)
             kv_loc = tl.load(
@@ -163,9 +150,7 @@ def _fwd_grouped_kernel_stage1(
                 qk += tl.dot(qpe, kpe.to(qpe.dtype))
             qk *= sm_scale
 
-            qk = tl.where(
-                mask_h[:, None] & (offs_n[None, :] < split_kv_end), qk, float("-inf")
-            )
+            qk = tl.where(mask_h[:, None] & (offs_n[None, :] < split_kv_end), qk, float("-inf"))
 
             # Load V from paged cache
             offs_buf_v = (
@@ -202,11 +187,7 @@ def _fwd_grouped_kernel_stage1(
             mask=(mask_h[:, None]) & (mask_dv[None, :]),
         )
 
-        offs_mid_o_1 = (
-            cur_batch * stride_mid_ob
-            + cur_head * stride_mid_oh
-            + split_kv_id * stride_mid_os
-        ) // Lv
+        offs_mid_o_1 = (cur_batch * stride_mid_ob + cur_head * stride_mid_oh + split_kv_id * stride_mid_os) // Lv
 
         tl.store(
             Att_Lse + offs_mid_o_1,
@@ -239,9 +220,7 @@ def _fwd_kernel_stage2(
     cur_batch = tl.program_id(0)
     cur_head = tl.program_id(1)
 
-    cur_batch_seq_len = tl.load(kv_indptr + cur_batch + 1) - tl.load(
-        kv_indptr + cur_batch
-    )
+    cur_batch_seq_len = tl.load(kv_indptr + cur_batch + 1) - tl.load(kv_indptr + cur_batch)
     kv_splits = tl.load(num_kv_splits + cur_batch)
 
     offs_d = tl.arange(0, BLOCK_DV)
@@ -253,18 +232,14 @@ def _fwd_kernel_stage2(
 
     offs_v = cur_batch * stride_mid_ob + cur_head * stride_mid_oh + offs_d
     offs_logic = (cur_batch * stride_mid_ob + cur_head * stride_mid_oh) // Lv
-    kv_len_per_split = (
-        tl.cdiv(tl.cdiv(cur_batch_seq_len, kv_splits), MIN_BLOCK_KV) * MIN_BLOCK_KV
-    )
+    kv_len_per_split = tl.cdiv(tl.cdiv(cur_batch_seq_len, kv_splits), MIN_BLOCK_KV) * MIN_BLOCK_KV
 
     for split_kv_id in range(0, MAX_KV_SPLITS):
         split_kv_start = kv_len_per_split * split_kv_id
         split_kv_end = tl.minimum(split_kv_start + kv_len_per_split, cur_batch_seq_len)
 
         if split_kv_end > split_kv_start:
-            tv = tl.load(
-                Mid_O + offs_v + split_kv_id * stride_mid_os, mask=mask_d, other=0.0
-            )
+            tv = tl.load(Mid_O + offs_v + split_kv_id * stride_mid_os, mask=mask_d, other=0.0)
             tlogic = tl.load(Mid_O_1 + offs_logic + split_kv_id * stride_mid_os // Lv)
             n_e_max = tl.maximum(tlogic, e_max)
 
@@ -276,9 +251,11 @@ def _fwd_kernel_stage2(
             e_sum = e_sum * old_scale + exp_logic
             e_max = n_e_max
 
+    # Guard against e_sum==0 (empty sequences from CUDAGraph padding) to avoid NaN
+    safe_e_sum = tl.where(e_sum == 0.0, 1.0, e_sum)
     tl.store(
         O + cur_batch * stride_obs + cur_head * stride_oh + offs_d,
-        acc / e_sum,
+        tl.where(e_sum == 0.0, 0.0, acc / safe_e_sum),
         mask=mask_d,
     )
 

fastdeploy/model_executor/models/deepseek_v3.py

@@ -17,6 +17,7 @@
 from __future__ import annotations
 
 import math
+import os
 import re
 from typing import Dict
 
@@ -344,6 +345,9 @@ def __init__(self, fd_config: FDConfig, layer_id: int, prefix: str = "") -> None
 
         self.prefix = prefix
 
+        prop = paddle.device.cuda.get_device_properties()
+        self.prop = prop
+
     @staticmethod
     def yarn_get_mscale(scale=1, mscale=1):
         """ """
@@ -362,6 +366,8 @@ def forward(
             fused_read_cache_and_interleave,
         )
 
+        q_total_token_num = hidden_states.shape[0]
+
         attn_out = None
         if self.use_gated_attn:
             gate_out = self.gate(hidden_states)
@@ -439,6 +445,36 @@ def forward(
             attn_out = fmha_out
 
         if need_do_decode:  # max_dec_len_this_time
+
+            if int(os.getenv("USE_FLASH_MLA", "0")) == 0 and self.prop.major == 9:
+                pass
+            else:
+                from fastdeploy.model_executor.layers.attention.mla_attention_backend import (
+                    extract_decoder_token_from_q,
+                    insert_decoder_result_back,
+                )
+
+                decoder_query_nope, cache_seqlens = extract_decoder_token_from_q(
+                    query_nope.reshape([0, -1]),
+                    forward_meta.cu_seqlens_q,
+                    forward_meta.seq_lens_encoder,
+                    forward_meta.seq_lens_decoder,
+                )
+
+                decoder_query_pe, cache_seqlens = extract_decoder_token_from_q(
+                    query_pe.reshape([0, -1]),
+                    forward_meta.cu_seqlens_q,
+                    forward_meta.seq_lens_encoder,
+                    forward_meta.seq_lens_decoder,
+                )
+                assert decoder_query_nope.shape[0] == forward_meta.seq_lens_encoder.shape[0]
+                assert decoder_query_pe.shape[0] == forward_meta.seq_lens_encoder.shape[0]
+
+                forward_meta.cache_seqlens = cache_seqlens
+
+                query_nope = decoder_query_nope.reshape([0, -1, self.qk_nope_head_dim])
+                query_pe = decoder_query_pe.reshape([0, -1, self.qk_rope_head_dim])
+
             q_nope_out = self.kv_b_proj_bmm(query_nope.transpose([1, 0, 2]), proj_type="k").transpose([1, 0, 2])
 
             q_input = paddle.concat([q_nope_out, query_pe], axis=-1)
@@ -467,6 +503,17 @@ def forward(
                 .reshape_([-1, self.num_attention_heads_tp * self.v_head_dim])
             )
 
+            if int(os.getenv("USE_FLASH_MLA", "0")) == 0 and self.prop.major == 9:
+                pass
+            else:
+                fmqa_out = insert_decoder_result_back(
+                    fmqa_out.reshape([0, 1, self.num_attention_heads_tp, self.v_head_dim]),
+                    forward_meta.cu_seqlens_q,
+                    forward_meta.seq_lens_encoder,
+                    forward_meta.seq_lens_decoder,
+                    q_total_token_num,
+                )
+
             if need_do_prefill:
                 merge_prefill_decode_output(
                     attn_out,
@@ -1062,7 +1109,7 @@ def forward(
         need_do_decode = forward_meta.max_len_tensor_cpu[2] > 0
 
         if not need_do_prefill and not need_do_decode:
-            return hidden_states
+            return hidden_states, residual
 
         if hidden_states.shape[0] > 0:
             hidden_states, residual = self.input_layernorm(