update forward

Author: Linboyan-trc

fastdeploy/model_executor/models/deepseek_v3.py

@@ -362,18 +362,6 @@ def forward(
             fused_read_cache_and_interleave,
         )
 
-        need_do_prefill = forward_meta.max_len_tensor_cpu[1] > 0
-        need_do_decode = forward_meta.max_len_tensor_cpu[2] > 0
-
-        # Idle pass (e.g. CUDAGraph padding): skip all attention computation
-        if not need_do_prefill and not need_do_decode:
-            return self.o_proj(
-                paddle.zeros(
-                    [hidden_states.shape[0], self.num_attention_heads_tp * self.v_head_dim],
-                    dtype=hidden_states.dtype,
-                )
-            )
-
         attn_out = None
         if self.use_gated_attn:
             gate_out = self.gate(hidden_states)
@@ -1070,6 +1058,12 @@ def forward(
         residual: paddle.Tensor,
     ):
         """ """
+        need_do_prefill = forward_meta.max_len_tensor_cpu[1] > 0
+        need_do_decode = forward_meta.max_len_tensor_cpu[2] > 0
+
+        if not need_do_prefill and not need_do_decode:
+            return hidden_states
+
         if hidden_states.shape[0] > 0:
             hidden_states, residual = self.input_layernorm(
                 hidden_states, residual_input=residual, forward_meta=forward_meta

tests/deterministic/test_triton_decode_attention.py

@@ -63,8 +63,7 @@ def cosine_similarity(a, b):
 # ---------------------------------------------------------------------------
 # Reference implementation: naive decode attention (no paging)
 # ---------------------------------------------------------------------------
-def naive_decode_attention_ref(q, k_pages, v_pages, kv_indptr, kv_indices,
-                               sm_scale, kv_block_size):
+def naive_decode_attention_ref(q, k_pages, v_pages, kv_indptr, kv_indices, sm_scale, kv_block_size):
     """
     Naive Python reference for decode attention with paged KV cache.
 
@@ -154,7 +153,6 @@ def build_decode_test_data(
     np.random.seed(seed)
     paddle.seed(seed)
 
-    total_kv_len = sum(seq_lens)
     num_blocks_needed = sum((s + block_size - 1) // block_size for s in seq_lens)
     num_blocks = max(num_blocks_needed + 4, 8)
 
@@ -194,7 +192,7 @@ def build_decode_test_data(
     attn_lse = paddle.empty([batch_size, num_heads, max_kv_splits], dtype="float32")
     o = paddle.empty([batch_size, num_heads, Lv], dtype=dtype)
 
-    sm_scale = head_dim_k ** -0.5
+    sm_scale = head_dim_k**-0.5
 
     return {
         "q": q,
@@ -290,8 +288,7 @@ def test_empty(self):
     ids=[c[0] for c in _DECODE_CASES],
 )
 @pytest.mark.parametrize("dtype", ["float16", "bfloat16"])
-def test_decode_attention_correctness(name, batch, num_heads, kv_heads, Lk, Lv,
-                                      seq_lens, block_size, dtype):
+def test_decode_attention_correctness(name, batch, num_heads, kv_heads, Lk, Lv, seq_lens, block_size, dtype):
     """Triton decode attention output should match naive reference."""
     data = build_decode_test_data(
         batch_size=batch,
@@ -336,12 +333,10 @@ def test_decode_attention_correctness(name, batch, num_heads, kv_heads, Lk, Lv,
     cos_sim = cosine_similarity(triton_out, ref_out)
 
     atol = BF16_ATOL if dtype == "bfloat16" else FP16_ATOL
-    assert max_diff < atol, (
-        f"[{name}/{dtype}] max_diff={max_diff:.6f} exceeds atol={atol}"
-    )
-    assert cos_sim > COSINE_SIM_THRESHOLD, (
-        f"[{name}/{dtype}] cos_sim={cos_sim:.6f} below threshold={COSINE_SIM_THRESHOLD}"
-    )
+    assert max_diff < atol, f"[{name}/{dtype}] max_diff={max_diff:.6f} exceeds atol={atol}"
+    assert (
+        cos_sim > COSINE_SIM_THRESHOLD
+    ), f"[{name}/{dtype}] cos_sim={cos_sim:.6f} below threshold={COSINE_SIM_THRESHOLD}"
 
 
 # ===========================================================================
@@ -380,10 +375,7 @@ def test_decode_attention_determinism():
         results.append(o.astype("float32").numpy())
 
     for i in range(1, len(results)):
-        np.testing.assert_array_equal(
-            results[0], results[i],
-            err_msg=f"Run 0 vs run {i} differ — non-deterministic!"
-        )
+        np.testing.assert_array_equal(results[0], results[i], err_msg=f"Run 0 vs run {i} differ — non-deterministic!")
 
 
 # ===========================================================================

tests/deterministic/test_triton_mla_cache_kernel.py

@@ -411,10 +411,7 @@ def test_write_cache_determinism():
         results.append(cache.astype("float32").numpy())
 
     for i in range(1, len(results)):
-        np.testing.assert_array_equal(
-            results[0], results[i],
-            err_msg=f"Run 0 vs run {i} differ — non-deterministic!"
-        )
+        np.testing.assert_array_equal(results[0], results[i], err_msg=f"Run 0 vs run {i} differ — non-deterministic!")
 
 
 # ===========================================================================
@@ -429,17 +426,23 @@ def test_manual_baseline():
     latent_dim = kv_lora_rank + qk_rope_head_dim  # 6
 
     # 3 tokens, deterministic values
-    compressed_kv = paddle.to_tensor([
-        [1.0, 2.0, 3.0, 4.0],
-        [5.0, 6.0, 7.0, 8.0],
-        [9.0, 10.0, 11.0, 12.0],
-    ], dtype="float32")
-
-    k_pe = paddle.to_tensor([
-        [0.1, 0.2],
-        [0.3, 0.4],
-        [0.5, 0.6],
-    ], dtype="float32")
+    compressed_kv = paddle.to_tensor(
+        [
+            [1.0, 2.0, 3.0, 4.0],
+            [5.0, 6.0, 7.0, 8.0],
+            [9.0, 10.0, 11.0, 12.0],
+        ],
+        dtype="float32",
+    )
+
+    k_pe = paddle.to_tensor(
+        [
+            [0.1, 0.2],
+            [0.3, 0.4],
+            [0.5, 0.6],
+        ],
+        dtype="float32",
+    )
 
     latent_cache = paddle.zeros([num_blocks, 1, block_size, latent_dim], dtype="float32")