From 3314cd426009b8e62a12953929cbc88799233047 Mon Sep 17 00:00:00 2001 From: sufubao Date: Thu, 30 Jul 2026 18:50:34 +0800 Subject: [PATCH] feat: add FlashQLA GDN prefill backend --- docker/Dockerfile | 2 + .../common/basemodel/attention/linear/gdn.py | 5 +- .../linear_att/gdn_prefill_backend.py | 70 +++++++++++++++++++ lightllm/server/api_cli.py | 9 +++ lightllm/server/core/objs/start_args_type.py | 1 + 5 files changed, 85 insertions(+), 2 deletions(-) create mode 100644 lightllm/common/basemodel/triton_kernel/linear_att/gdn_prefill_backend.py diff --git a/docker/Dockerfile b/docker/Dockerfile index 115b0c48e..3688f2023 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -10,6 +10,7 @@ ARG DEEPGEMM_REF=891d57b4db1071624b5c8fa0d1e51cb317fa709f ARG DEEPEP_REF=099d5f2bad488b9c534ea785062b12f2e91d1d41 ARG DEEPEP_NCCL_VERSION=2.30.4 ARG DEEPEP_NVSHMEM_VERSION=3.3.24 +ARG FLASHQLA_VERSION=0.1.1 ARG TARGETPLATFORM ARG ENABLE_DEEPEP=1 ARG ENABLE_NIXL=1 @@ -58,6 +59,7 @@ RUN pip install --no-cache-dir \ RUN pip install -r /lightllm/requirements.txt --no-cache-dir \ -i https://pypi.org/simple \ --extra-index-url https://download.pytorch.org/whl/cu130 +RUN PIP_NO_INDEX=0 pip install --no-cache-dir "flash-qla==${FLASHQLA_VERSION}" -i https://pypi.org/simple RUN export CPATH=/usr/local/cuda/targets/x86_64-linux/include/cccl:/usr/local/cuda/targets/x86_64-linux/include${CPATH:+:${CPATH}} && \ git clone https://github.com/deepseek-ai/FlashMLA.git /root/FlashMLA && \ cd /root/FlashMLA && \ diff --git a/lightllm/common/basemodel/attention/linear/gdn.py b/lightllm/common/basemodel/attention/linear/gdn.py index c4c3e15cb..f313d2e9b 100644 --- a/lightllm/common/basemodel/attention/linear/gdn.py +++ b/lightllm/common/basemodel/attention/linear/gdn.py @@ -5,8 +5,8 @@ from lightllm.utils.envs_utils import get_env_start_args, get_llm_data_type from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d import causal_conv1d_fn from lightllm.common.basemodel.triton_kernel.linear_att.fused_gdn_gating import fused_gdn_gating -from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops import chunk_gated_delta_rule from lightllm.common.basemodel.triton_kernel.linear_att.gdn_decode_pack import conv_pack_gdn_decode_inputs +from lightllm.common.basemodel.triton_kernel.linear_att.gdn_prefill_backend import get_gdn_prefill_chunk_fn from lightllm.common.basemodel.triton_kernel.linear_att.mtp_fused_recurrent import ( mtp_fused_recurrent_gated_delta_rule, ) @@ -58,6 +58,7 @@ def _init_linear_layer_metadata(self, network_config, tp_world_size): # GDN kernel output dtype is self.data_type # Conversion needed only if SSM state uses different dtype self.needs_ssm_dtype_conversion = get_llm_data_type() != self.ssm_state_dtype + self._gdn_prefill_chunk = get_gdn_prefill_chunk_fn() return def _split_qkvzba(self, mixed_qkvzba): @@ -175,7 +176,7 @@ def _gdn_prefill_kernel( query, key, value = backend._rearrange_mixed_qkv(mixed_qkv) initial_state = ssm_states[self.b_ssm_buffer_idx] # g and beta have shape (total_tokens, num_heads), need to unsqueeze to get (1, total_tokens, num_heads) - core_attn_out, last_recurrent_state = chunk_gated_delta_rule( + core_attn_out, last_recurrent_state = backend._gdn_prefill_chunk( q=query, k=key, v=value, diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/gdn_prefill_backend.py b/lightllm/common/basemodel/triton_kernel/linear_att/gdn_prefill_backend.py new file mode 100644 index 000000000..764228671 --- /dev/null +++ b/lightllm/common/basemodel/triton_kernel/linear_att/gdn_prefill_backend.py @@ -0,0 +1,70 @@ +import functools + +import torch + +from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops import ( + chunk_gated_delta_rule as _fla_chunk_gated_delta_rule, +) +from lightllm.utils.envs_utils import get_env_start_args +from lightllm.utils.log_utils import init_logger + +logger = init_logger(__name__) + + +@functools.lru_cache(maxsize=1) +def get_gdn_prefill_chunk_fn(): + backend = getattr(get_env_start_args(), "gdn_prefill_backend", "fla") + if backend != "flashqla": + return _fla_chunk_gated_delta_rule + + capability = torch.cuda.get_device_capability() + if capability[0] < 9: + logger.warning( + f"gdn_prefill_backend=flashqla requires Hopper (SM90+), got SM{capability[0]}{capability[1]}; " + "falling back to the FLA triton kernel." + ) + return _fla_chunk_gated_delta_rule + + try: + from flash_qla import chunk_gated_delta_rule as flashqla_chunk_gated_delta_rule + except Exception as exc: + logger.warning( + f"gdn_prefill_backend=flashqla but importing flash_qla failed ({exc!r}); " + "falling back to the FLA triton kernel. " + "Install FlashQLA (https://github.com/QwenLM/FlashQLA)." + ) + return _fla_chunk_gated_delta_rule + + flashqla_initialized = False + use_flashqla = True + + def flashqla_chunk(q, k, v, **kwargs): + nonlocal flashqla_initialized, use_flashqla + + if not use_flashqla: + return _fla_chunk_gated_delta_rule(q=q, k=k, v=v, **kwargs) + + flashqla_q = q.contiguous() + flashqla_k = k.contiguous() + flashqla_v = v.contiguous() + try: + result = flashqla_chunk_gated_delta_rule( + q=flashqla_q, + k=flashqla_k, + v=flashqla_v, + **kwargs, + ) + except Exception as exc: + if flashqla_initialized: + raise + use_flashqla = False + logger.warning( + f"FlashQLA failed during its first invocation ({exc!r}); " "falling back to the FLA triton kernel." + ) + return _fla_chunk_gated_delta_rule(q=q, k=k, v=v, **kwargs) + + flashqla_initialized = True + return result + + logger.info("GDN chunked-prefill backend: FlashQLA (TileLang, Hopper).") + return flashqla_chunk diff --git a/lightllm/server/api_cli.py b/lightllm/server/api_cli.py index 37d7a0d56..3f203aaf1 100644 --- a/lightllm/server/api_cli.py +++ b/lightllm/server/api_cli.py @@ -857,6 +857,15 @@ def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: default="float32", help="the data type of linear att smm data type", ) + parser.add_argument( + "--gdn_prefill_backend", + type=str, + choices=["fla", "flashqla"], + default="fla", + help="""GDN chunked-prefill kernel backend for hybrid linear-attention models. + 'fla' uses the vendored flash-linear-attention Triton kernel. 'flashqla' uses the + TileLang FlashQLA kernel on SM90+ and falls back to 'fla' if FlashQLA is unavailable.""", + ) parser.add_argument( "--disable_linear_att_small_page_cpu_cache", action="store_true", diff --git a/lightllm/server/core/objs/start_args_type.py b/lightllm/server/core/objs/start_args_type.py index 516672258..1f0781814 100644 --- a/lightllm/server/core/objs/start_args_type.py +++ b/lightllm/server/core/objs/start_args_type.py @@ -228,3 +228,4 @@ class StartArgs: disable_linear_att_small_page_cpu_cache: bool = field(default=False) linear_att_cache_size: Optional[int] = field(default=None) linear_att_ssm_data_type: Optional[str] = field(default="float32", metadata={"choices": ["bfloat16", "float32"]}) + gdn_prefill_backend: str = field(default="fla", metadata={"choices": ["fla", "flashqla"]})