[3/5] refactor(attention): add provider backends and explicit selection

lightllm/common/attention/backends/_flashinfer.py

@@ -0,0 +1,10 @@
+import os
+
+import torch
+
+
+def set_flashinfer_envs():
+    capability = torch.cuda.get_device_capability()
+    arch = f"{capability[0]}.{capability[1]}"
+    os.environ["TORCH_CUDA_ARCH_LIST"] = f"{arch}{'+PTX' if arch == '9.0' else ''}"
+    return

lightllm/common/attention/backends/_shared/linear/gdn.py

@@ -0,0 +1,56 @@
+"""GDN tensor-parallel dimensions and input views shared by native and replay kernels."""
+
+import torch
+
+from lightllm.utils.envs_utils import get_env_start_args
+
+
+class GDNMetadata:
+    def _init_linear_layer_metadata(self, network_config, tp_world_size):
+        self.mtp_step = get_env_start_args().mtp_step
+        self.num_v_heads = network_config["linear_num_value_heads"]
+        self.num_k_heads = network_config["linear_num_key_heads"]
+        self.head_k_dim = network_config["linear_key_head_dim"]
+        self.head_v_dim = network_config["linear_value_head_dim"]
+        self.key_dim = self.head_k_dim * self.num_k_heads
+        self.value_dim = self.head_v_dim * self.num_v_heads
+        self.conv_kernel_dim = network_config["linear_conv_kernel_dim"]
+        self.activation = network_config["hidden_act"]
+        self.tp_qkvz_dim = (self.key_dim * 2 + self.value_dim * 2) // tp_world_size
+        self.tp_ba_dim = self.num_v_heads * 2 // tp_world_size
+        self.tp_num_k_heads = self.num_k_heads // tp_world_size
+        self.tp_num_v_heads = self.num_v_heads // tp_world_size
+        self.tp_key_dim = self.key_dim // tp_world_size
+        self.tp_value_dim = self.value_dim // tp_world_size
+        if self.num_v_heads % self.num_k_heads:
+            raise ValueError("num_v_heads must be divisible by num_k_heads")
+        self.num_v_heads_per_k_head = self.num_v_heads // self.num_k_heads
+        ssm_dtype_dict = {"bfloat16": torch.bfloat16, "float32": torch.float32}
+        self.ssm_state_dtype = ssm_dtype_dict.get(get_env_start_args().linear_att_ssm_data_type, torch.bfloat16)
+        return
+
+    def _split_qkvzba(self, mixed_qkvzba):
+        qkv_dim = self.tp_key_dim * 2 + self.tp_value_dim
+        z_end = qkv_dim + self.tp_value_dim
+        b_end = z_end + self.tp_num_v_heads
+        mixed_qkv = mixed_qkvzba[:, :qkv_dim]
+        z = mixed_qkvzba[:, qkv_dim:z_end].view(-1, self.tp_num_v_heads, self.head_v_dim)
+        b = mixed_qkvzba[:, z_end:b_end]
+        a = mixed_qkvzba[:, b_end:]
+        return (mixed_qkv, z, b, a)
+
+    def _rearrange_mixed_qkv(self, mixed_qkv, decode=False):
+        if decode:
+            query, key, value = torch.split(mixed_qkv, [self.tp_key_dim, self.tp_key_dim, self.tp_value_dim], dim=-1)
+            batch_size = mixed_qkv.shape[0]
+            query = query.view(batch_size, 1, self.tp_num_k_heads, self.head_k_dim)
+            key = key.view(batch_size, 1, self.tp_num_k_heads, self.head_k_dim)
+            value = value.view(batch_size, 1, self.tp_num_v_heads, self.head_v_dim)
+            return (query, key, value)
+        else:
+            query, key, value = torch.split(mixed_qkv, [self.tp_key_dim, self.tp_key_dim, self.tp_value_dim], dim=-1)
+            seq_len = query.shape[0]
+            query = query.view(1, seq_len, self.tp_num_k_heads, self.head_k_dim)
+            key = key.view(1, seq_len, self.tp_num_k_heads, self.head_k_dim)
+            value = value.view(1, seq_len, self.tp_num_v_heads, self.head_v_dim)
+            return (query, key, value)

lightllm/common/attention/backends/_shared/linear/layout.py

@@ -0,0 +1,67 @@
+from dataclasses import dataclass, field
+
+import torch
+
+from lightllm.common.attention.backends.base import AttentionExecutionContext
+
+from .mtp_state_params import build_dynamic_mtp_layout
+
+
+@dataclass(slots=True)
+class LinearDecodeContext(AttentionExecutionContext):
+    b_conv_buffer_idx: torch.Tensor | None = None
+    b_ssm_buffer_idx: torch.Tensor | None = None
+    b1_mtp_cu_q_seq_len: torch.Tensor | None = None
+    b_cu_seqlens: torch.Tensor | None = None
+    accepted_by_state: dict[torch.Tensor, torch.Tensor] = field(default_factory=dict)
+
+    def accepted_tokens(self, state_index: torch.Tensor) -> torch.Tensor:
+        """Gather once per State per update; Graph captures the gather for every replay.
+
+        The cache retains the source tensors until the next update, keeping identity
+        keys valid for this context. HOLD sequences must select the initial convolution state.
+        """
+        key = state_index
+        if key not in self.accepted_by_state:
+            self.accepted_by_state[key] = torch.where(
+                self.b_conv_buffer_idx == self.batch.hold_req_id,
+                1,
+                state_index[self.b_conv_buffer_idx] + 1,
+            )
+        return self.accepted_by_state[key]
+
+
+def update_linear_decode(context: LinearDecodeContext, draft_step: int, dynamic: bool, *, replay: bool) -> None:
+    batch = context.batch
+    context.accepted_by_state.clear()
+    if draft_step == 0:
+        context.b_conv_buffer_idx = batch.b_req_idx
+        context.b_ssm_buffer_idx = batch.b_req_idx
+        context.b1_mtp_cu_q_seq_len = None
+        if replay:
+            context.b_cu_seqlens = torch.arange(batch.batch_size + 1, dtype=torch.int32, device=batch.device)
+        return
+    width = draft_step + 1
+    if dynamic:
+        context.b1_mtp_cu_q_seq_len, context.b_conv_buffer_idx = build_dynamic_mtp_layout(
+            b_req_idx=batch.b_req_idx,
+            b_mtp_index=batch.b_mtp_index,
+            hold_req_id=batch.hold_req_id,
+        )
+    else:
+        if batch.batch_size % width:
+            raise ValueError("linear decode batch size must be divisible by draft_step + 1")
+        context.b1_mtp_cu_q_seq_len = torch.arange(
+            0,
+            batch.batch_size + 1,
+            width,
+            dtype=torch.int32,
+            device=batch.device,
+        )
+        context.b_conv_buffer_idx = batch.b_req_idx.view(-1, width)[:, 0].contiguous()
+    if replay:
+        context.b_ssm_buffer_idx = context.b_conv_buffer_idx
+        context.b_cu_seqlens = context.b1_mtp_cu_q_seq_len
+    else:
+        offsets = torch.arange(width, dtype=batch.b_req_idx.dtype, device=batch.device)
+        context.b_ssm_buffer_idx = context.b_conv_buffer_idx[:, None] * width + offsets

lightllm/common/attention/backends/_shared/linear/mtp_state_params.py

@@ -14,6 +14,7 @@ def _build_dynamic_mtp_linear_att_state_params_kernel(
     batch_size,
     hold_req_id,
     BLOCK_SIZE: tl.constexpr,
+    SELECT_ACCEPTED: tl.constexpr = True,
 ):
     offsets = tl.arange(0, BLOCK_SIZE)
     token_mask = offsets < batch_size
@@ -30,21 +31,18 @@ def _build_dynamic_mtp_linear_att_state_params_kernel(
     # replay arbitrary compact per-request widths at the same token batch size.
     tl.store(out_cu_q_seq_len + offsets, actual_token_num, mask=cu_mask)
     tl.store(out_conv_buffer_idx + offsets, hold_req_id, mask=token_mask)
-    tl.store(out_num_accepted_tokens + offsets, 1, mask=token_mask)
+    if SELECT_ACCEPTED:
+        tl.store(out_num_accepted_tokens + offsets, 1, mask=token_mask)
     tl.debug_barrier()
 
     # mtp_index restarts at zero on the first row of every request group.
     is_start = valid_row & (mtp_index == 0)
     sequence_index = tl.cumsum(tl.where(is_start, 1, 0), axis=0) - 1
-    accepted_state_index = tl.load(
-        req_to_mtp_state_index + req_idx,
-        mask=is_start,
-        other=0,
-    )
-
     tl.store(out_cu_q_seq_len + sequence_index, offsets, mask=is_start)
     tl.store(out_conv_buffer_idx + sequence_index, req_idx, mask=is_start)
-    tl.store(out_num_accepted_tokens + sequence_index, accepted_state_index + 1, mask=is_start)
+    if SELECT_ACCEPTED:
+        accepted_state_index = tl.load(req_to_mtp_state_index + req_idx, mask=is_start, other=0)
+        tl.store(out_num_accepted_tokens + sequence_index, accepted_state_index + 1, mask=is_start)
 
 
 def build_dynamic_mtp_linear_att_state_params(
@@ -85,11 +83,10 @@ def build_dynamic_mtp_linear_att_state_params(
     per-request query lengths.
     """
 
-    assert b_req_idx.is_cuda and b_mtp_index.is_cuda and req_to_mtp_state_index.is_cuda
-    assert b_req_idx.ndim == 1 and b_req_idx.shape == b_mtp_index.shape
-    assert b_req_idx.dtype == torch.int32 and b_mtp_index.dtype == torch.int32
+    _validate_layout_inputs(b_req_idx, b_mtp_index)
+    if req_to_mtp_state_index.device != b_req_idx.device or req_to_mtp_state_index.ndim != 1:
+        raise ValueError("MTP state indices must be a vector on the request device")
     batch_size = b_req_idx.shape[0]
-    assert batch_size > 0
 
     b1_cu_q_seq_len = torch.empty((batch_size + 1,), dtype=torch.int32, device=b_req_idx.device)
     b_conv_buffer_idx = torch.empty_like(b_req_idx)
@@ -109,3 +106,34 @@ def build_dynamic_mtp_linear_att_state_params(
         num_stages=1,
     )
     return b1_cu_q_seq_len, b_conv_buffer_idx, b_num_accepted_tokens
+
+
+def build_dynamic_mtp_layout(b_req_idx: torch.Tensor, b_mtp_index: torch.Tensor, hold_req_id: int):
+    _validate_layout_inputs(b_req_idx, b_mtp_index)
+    batch_size = b_req_idx.numel()
+    cu = torch.empty(batch_size + 1, dtype=torch.int32, device=b_req_idx.device)
+    requests = torch.empty_like(b_req_idx)
+    _build_dynamic_mtp_linear_att_state_params_kernel[(1,)](
+        b_req_idx=b_req_idx,
+        b_mtp_index=b_mtp_index,
+        req_to_mtp_state_index=None,
+        out_cu_q_seq_len=cu,
+        out_conv_buffer_idx=requests,
+        out_num_accepted_tokens=None,
+        batch_size=batch_size,
+        hold_req_id=hold_req_id,
+        BLOCK_SIZE=triton.next_power_of_2(batch_size + 1),
+        SELECT_ACCEPTED=False,
+        num_warps=8,
+        num_stages=1,
+    )
+    return cu, requests
+
+
+def _validate_layout_inputs(requests, positions):
+    if not requests.is_cuda or positions.device != requests.device:
+        raise ValueError("MTP layout requires inputs on the same CUDA device")
+    if requests.ndim != 1 or requests.shape != positions.shape or requests.numel() == 0:
+        raise ValueError("MTP layout requires matching nonempty vectors")
+    if requests.dtype != torch.int32 or positions.dtype != torch.int32:
+        raise ValueError("MTP layout requires int32 request and position indices")

lightllm/common/attention/backends/_shared/linear/replayssm_decode.py

@@ -10,7 +10,7 @@
 import triton
 import triton.language as tl
 
-from lightllm.common.attention.backends._shared.linear.replayssm_delta import _apply_accepted_prefix
+from .replayssm_delta import _apply_accepted_prefix
 
 
 def check_replay_sync_mark(

lightllm/common/attention/backends/_shared/linear/replayssm_fold.py

@@ -4,7 +4,7 @@
 import triton
 import triton.language as tl
 
-from lightllm.common.attention.backends._shared.linear.replayssm_delta import _apply_accepted_prefix
+from .replayssm_delta import _apply_accepted_prefix
 
 
 @triton.jit