diff --git a/lightllm/common/basemodel/attention/linear/gdn.py b/lightllm/common/basemodel/attention/linear/gdn.py index c4c3e15cb..ae2816840 100644 --- a/lightllm/common/basemodel/attention/linear/gdn.py +++ b/lightllm/common/basemodel/attention/linear/gdn.py @@ -2,7 +2,7 @@ import torch from typing import TYPE_CHECKING from ..base_att import BaseAttBackend, BasePrefillAttState, BaseDecodeAttState, AttControl -from lightllm.utils.envs_utils import get_env_start_args, get_llm_data_type +from lightllm.utils.envs_utils import get_env_start_args 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 @@ -54,10 +54,6 @@ def _init_linear_layer_metadata(self, network_config, tp_world_size): start_args = get_env_start_args() self.ssm_state_dtype = ssm_dtype_dict.get(start_args.linear_att_ssm_data_type, torch.bfloat16) - # Pre-compute whether dtype conversion is needed - # 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 return def _split_qkvzba(self, mixed_qkvzba): @@ -187,10 +183,10 @@ def _gdn_prefill_kernel( head_first=False, use_qk_l2norm_in_kernel=True, ) - if backend.needs_ssm_dtype_conversion: - ssm_states[self.b_ssm_buffer_idx] = last_recurrent_state.to(backend.ssm_state_dtype, copy=False) - else: - ssm_states[self.b_ssm_buffer_idx] = last_recurrent_state + # The chunk kernel accumulates the recurrent state in float32 even when + # the state cache is configured as bfloat16. Advanced indexing does + # not perform an implicit dtype conversion for index_put. + ssm_states[self.b_ssm_buffer_idx] = last_recurrent_state.to(ssm_states.dtype, copy=False) return core_attn_out diff --git a/unit_tests/common/basemodel/attention/linear/test_gdn.py b/unit_tests/common/basemodel/attention/linear/test_gdn.py new file mode 100644 index 000000000..241e0a5d2 --- /dev/null +++ b/unit_tests/common/basemodel/attention/linear/test_gdn.py @@ -0,0 +1,56 @@ +from types import SimpleNamespace + +import pytest +import torch + +import lightllm.common.basemodel.attention.linear.gdn as gdn + + +@pytest.mark.parametrize("cache_dtype", [torch.bfloat16, torch.float32]) +def test_prefill_casts_final_state_to_cache_dtype(monkeypatch, cache_dtype): + ssm_states = torch.zeros((1, 1, 2, 2), dtype=cache_dtype) + final_state = torch.ones((1, 1, 2, 2), dtype=torch.float32) + + monkeypatch.setattr(gdn, "fused_gdn_gating", lambda _log, a, b, _bias: (a, b)) + monkeypatch.setattr(gdn, "causal_conv1d_fn", lambda mixed, *args, **kwargs: mixed) + monkeypatch.setattr( + gdn, + "chunk_gated_delta_rule", + lambda *args, **kwargs: (None, final_state), + ) + + qkv = torch.zeros((1, 3), dtype=cache_dtype) + q = torch.zeros((1, 1, 1, 1), dtype=cache_dtype) + backend = SimpleNamespace( + mtp_step=0, + activation="silu", + ssm_state_dtype=cache_dtype, + _rearrange_mixed_qkv=lambda mixed: (q, q, q), + ) + state = gdn.LinearAttPrefillAttState( + backend=backend, + infer_state=SimpleNamespace( + b1_cu_q_seq_len=torch.tensor([0, 1], dtype=torch.int32), + b_ready_cache_len=0, + ), + ) + state.b_conv_buffer_idx = torch.tensor([0], dtype=torch.int64) + state.b_ssm_buffer_idx = torch.tensor([0], dtype=torch.int64) + layer_weight = SimpleNamespace( + linear_A_log=SimpleNamespace(weight=None), + linear_dt_bias=SimpleNamespace(weight=None), + linear_conv1d=SimpleNamespace(mm_param=SimpleNamespace(weight=None), bias=None), + ) + + state._gdn_prefill_kernel( + qkv, + torch.zeros((1, 3), dtype=cache_dtype), + ssm_states, + torch.zeros((1, 1), dtype=cache_dtype), + torch.zeros((1, 1), dtype=cache_dtype), + state.infer_state, + layer_weight, + ) + + assert ssm_states.dtype == cache_dtype + assert torch.equal(ssm_states, final_state.to(cache_dtype))