[NPU] adaptation to support operator FA3 in deterministic inference on NPU. (#26724)

python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py

@@ -316,6 +316,7 @@ class AscendAttnBackend(AttentionBackend):
         self.token_to_kv_pool = model_runner.token_to_kv_pool
         self.req_to_token = model_runner.req_to_token_pool.req_to_token
         self.graph_mode = False
+        self.use_fa = get_bool_env_var("ASCEND_USE_FA", "False")
         self.use_fia = get_bool_env_var("ASCEND_USE_FIA", "False")
         self.enable_torch_compile = model_runner.server_args.enable_torch_compile
         self.speculative_num_draft_tokens = (
@@ -1401,6 +1402,44 @@ class AscendAttnBackend(AttentionBackend):
                     attn_output = attn_output.view(
                         -1, layer.tp_q_head_num * layer.v_head_dim
                     )
+            elif self.use_fa:
+                from flash_attn_npu_v3 import flash_attn_with_kvcache
+
+                q = q.reshape(-1, layer.tp_q_head_num, layer.qk_head_dim)
+                k = k_cache.view(
+                    -1, self.page_size, layer.tp_k_head_num, layer.qk_head_dim
+                )
+                v = v_cache.view(
+                    -1, self.page_size, layer.tp_v_head_num, layer.v_head_dim
+                )
+                extend_seq_lens = self.forward_metadata.extend_seq_lens_cpu_int.npu()
+                cu_seqlens_q = torch.cat(
+                    [
+                        torch.zeros(1, dtype=torch.int32).npu(),
+                        extend_seq_lens.cumsum(0).to(torch.int32),
+                    ]
+                )
+                max_seqlen_q = extend_seq_lens.max().item()
+                attn_output = flash_attn_with_kvcache(
+                    q,
+                    k,
+                    v,
+                    cache_seqlens=self.forward_metadata.seq_lens,
+                    page_table=self.forward_metadata.block_tables,
+                    cu_seqlens_q=cu_seqlens_q,
+                    max_seqlen_q=max_seqlen_q,
+                    softmax_scale=layer.scaling,
+                    causal=True,
+                    window_size=[-1, -1],
+                    softcap=0.0,
+                    rotary_interleaved=False,
+                    num_splits=0,
+                    sm_margin=0,
+                    return_softmax_lse=False,
+                )
+                attn_output = attn_output.view(
+                    -1, layer.tp_q_head_num * layer.v_head_dim
+                )
             else:
                 causal = True
                 if (
@@ -2543,6 +2582,26 @@ class AscendAttnBackend(AttentionBackend):
                         ],
                         dim=0,
                     )
+            elif self.use_fa:
+                from flash_attn_npu_v3 import flash_attn_with_kvcache
+
+                q = q.view(
+                    forward_batch.batch_size, -1, layer.tp_q_head_num, layer.qk_head_dim
+                )
+                k = k_cache.view(
+                    -1, self.page_size, layer.tp_k_head_num, layer.qk_head_dim
+                )
+                v = v_cache.view(
+                    -1, self.page_size, layer.tp_v_head_num, layer.v_head_dim
+                )
+                attn_output = flash_attn_with_kvcache(
+                    q,
+                    k,
+                    v,
+                    page_table=self.forward_metadata.block_tables,
+                    cache_seqlens=self.forward_metadata.seq_lens,
+                    softmax_scale=layer.scaling,
+                )
             # there are some accuracy issues in cross attention scene to use torch_npu._npu_flash_attention_qlens
             # forward_batch.encoder_lens is not None in cross attention scend, we add native attn to solve accuracy issues
             elif forward_batch.encoder_lens is None and layer.logit_cap == 0: