fix(qwen3.5): cast prefill SSM state to cache dtype

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
 
 

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))