From 6270da0124945ea76de9595aa510335f14643d43 Mon Sep 17 00:00:00 2001 From: baishihao Date: Wed, 11 Nov 2026 22:27:19 +0800 Subject: [PATCH] perf(kda): fuse prefill convolution, Q/K norm and history update --- .../common/basemodel/attention/linear/kda.py | 31 ++- .../linear_att/causal_conv1d_qkv.py | 170 ++++++++++++++++ .../models/glm5_next/test_kda_conv_qkv.py | 186 ++++++++++++++++++ 3 files changed, 367 insertions(+), 20 deletions(-) create mode 100644 lightllm/common/basemodel/triton_kernel/linear_att/causal_conv1d_qkv.py create mode 100644 unit_tests/models/glm5_next/test_kda_conv_qkv.py diff --git a/lightllm/common/basemodel/attention/linear/kda.py b/lightllm/common/basemodel/attention/linear/kda.py index cd2f8afb39..11f801b62a 100644 --- a/lightllm/common/basemodel/attention/linear/kda.py +++ b/lightllm/common/basemodel/attention/linear/kda.py @@ -17,10 +17,8 @@ BaseDecodeAttState, BasePrefillAttState, ) -from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d import ( - causal_conv1d_fn, - causal_conv1d_update, -) +from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d import causal_conv1d_update +from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d_qkv import causal_conv1d_qkv from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.kda import chunk_kda_with_fused_gate from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.kda_decode import fused_recurrent_kda from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d_mtp import ( @@ -92,22 +90,15 @@ def prefill_att( ) conv_states, ssm_states = state.get_layer_view(logical_layer_id=layer_num) conv_states = conv_states[..., : backend.conv_kernel_size - 1] - mixed_qkv = causal_conv1d_fn( - mixed_qkv.transpose(0, 1), + q, k, v = causal_conv1d_qkv( + mixed_qkv, layer_weight.get_merged_kda_conv_weight(), - bias=None, - query_start_loc=self.infer_state.b1_cu_q_seq_len, - cache_indices=self.b_conv_buffer_idx, - has_initial_state=self.infer_state.b_ready_cache_len > 0, - conv_states=conv_states, - activation="silu", - ).transpose(0, 1) - - q, k, v = backend.split_qkv(mixed_qkv) - q = q.view(1, -1, backend.tp_num_heads, backend.head_dim) - k = k.view(1, -1, backend.tp_num_heads, backend.head_dim) - v = v.view(1, -1, backend.tp_num_heads, backend.head_dim) - raw_gate = raw_gate.view(1, -1, backend.tp_hidden_size) + conv_states, + self.infer_state.b1_cu_q_seq_len, + self.b_conv_buffer_idx, + self.infer_state.b_ready_cache_len, + self.infer_state.max_q_seq_len, + ) raw_beta = raw_beta.view(1, -1, backend.tp_num_heads) initial_state = ssm_states[self.b_ssm_buffer_idx].contiguous() @@ -121,7 +112,7 @@ def prefill_att( g_bias=layer_weight.linear_dt_bias.weight, initial_state=initial_state, output_final_state=True, - use_qk_l2norm_in_kernel=True, + use_qk_l2norm_in_kernel=False, cu_seqlens=self.infer_state.b1_cu_q_seq_len, safe_gate=True, lower_bound=backend.lower_bound, diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/causal_conv1d_qkv.py b/lightllm/common/basemodel/triton_kernel/linear_att/causal_conv1d_qkv.py new file mode 100644 index 0000000000..bb3cd783c0 --- /dev/null +++ b/lightllm/common/basemodel/triton_kernel/linear_att/causal_conv1d_qkv.py @@ -0,0 +1,170 @@ +# SPDX-License-Identifier: Apache-2.0 + +"""Width-four BF16 prefill convolution fused with Q/K norm and history update.""" + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _conv_qkv_kernel( + x, + weight, + conv_states, + query_start_loc, + cache_indices, + cache_seqlens, + q, + k, + v, + head_num: tl.constexpr, + seq_num: tl.constexpr, + PACKED: tl.constexpr, + head_dim: tl.constexpr, + stride_x_token: tl.constexpr, + stride_conv_state_seq: tl.constexpr, + stride_conv_state_dim: tl.constexpr, + stride_conv_state_tok: tl.constexpr, + BLOCK_T: tl.constexpr, +): + if PACKED: + tile_id = tl.program_id(1) + seq_id = tl.full((), 0, tl.int32) + # floor(token offset / BLOCK_T) + request id gives each request + # enough tiles, including unaligned starts and empty requests. + for shift in tl.static_range((seq_num - 1).bit_length() - 1, -1, -1): + candidate = seq_id + (1 << shift) + tile_start = tl.load(query_start_loc + candidate, candidate < seq_num, 0) // BLOCK_T + candidate + seq_id = tl.where((candidate < seq_num) & (tile_start <= tile_id), candidate, seq_id) + seq_start = tl.load(query_start_loc + seq_id) + tile_idx = tile_id - (seq_start // BLOCK_T + seq_id) + head_idx = tl.program_id(0) + else: + seq_id = tl.program_id(0) + seq_start = tl.load(query_start_loc + seq_id) + tile_idx = tl.program_id(1) + head_idx = tl.program_id(2) + token_idx = tile_idx * BLOCK_T + tl.arange(0, BLOCK_T) + qkv_id = head_idx // head_num + head_id = head_idx % head_num + dim_idx = tl.arange(0, head_dim) + channel_idx = head_idx * head_dim + dim_idx + seq_len = tl.load(query_start_loc + seq_id + 1) - seq_start + cache_idx = tl.load(cache_indices + seq_id) + has_history = tl.load(cache_seqlens + seq_id) > 0 + mask = token_idx[:, None] < seq_len + if tile_idx * BLOCK_T < seq_len: + if cache_idx == -1: + y = tl.load(x + (seq_start + token_idx[:, None]) * stride_x_token + channel_idx[None, :], mask, 0) + else: + history0 = tl.full((head_dim,), 0, tl.float32) + history1 = tl.full((head_dim,), 0, tl.float32) + history2 = tl.full((head_dim,), 0, tl.float32) + if has_history and tile_idx == 0: + state_ptr = conv_states + cache_idx * stride_conv_state_seq + channel_idx * stride_conv_state_dim + history0 = tl.load(state_ptr).to(tl.float32) + history1 = tl.load(state_ptr + stride_conv_state_tok).to(tl.float32) + history2 = tl.load(state_ptr + 2 * stride_conv_state_tok).to(tl.float32) + acc = tl.full((BLOCK_T, head_dim), 0, tl.float32) + for conv_idx in tl.static_range(4): + input_token = token_idx + conv_idx - 3 + x_val = tl.load( + x + (seq_start + input_token[:, None]) * stride_x_token + channel_idx[None, :], + mask & (input_token[:, None] >= 0), + 0, + ).to(tl.float32) + if conv_idx < 3: + if has_history and tile_idx == 0: + history = tl.where( + input_token[:, None] == -3, + history0[None, :], + tl.where(input_token[:, None] == -2, history1[None, :], history2[None, :]), + ) + x_val = tl.where(input_token[:, None] < 0, history, x_val) + w_val = tl.load(weight + channel_idx * 4 + conv_idx).to(tl.float32) + acc = tl.fma(x_val, w_val[None, :], acc) + # Preserve the BF16 rounding between convolution and Q/K norm. + y = (acc / (1 + tl.exp(-acc))).to(q.dtype.element_ty) + if qkv_id < 2: + qk = y.to(tl.float32) + norm = tl.rsqrt(tl.sum(qk * qk, 1)[:, None] + 1e-6) + y = (qk * norm).to(q.dtype.element_ty) + output = tl.where(qkv_id == 0, q, tl.where(qkv_id == 1, k, v)) + out_idx = ((seq_start + token_idx[:, None]) * head_num + head_id) * head_dim + dim_idx[None, :] + tl.store(output + out_idx, y, mask) + + if tile_idx == 0 and cache_idx >= 0: + # Only this program reads and updates these channels' initial history. + tl.debug_barrier() + state_pos = tl.arange(0, 4) + input_token = seq_len - 3 + state_pos + state_mask = state_pos[:, None] < 3 + tail = tl.load( + x + (seq_start + input_token[:, None]) * stride_x_token + channel_idx[None, :], + state_mask & (input_token[:, None] >= 0), + 0, + ) + history = tl.load( + conv_states + + cache_idx * stride_conv_state_seq + + channel_idx[None, :] * stride_conv_state_dim + + (input_token[:, None] + 3) * stride_conv_state_tok, + # Match SGL: empty requests clear history, including warm requests. + state_mask & (input_token[:, None] < 0) & has_history & (seq_len > 0), + 0, + ) + # Finish all history reads before any lane overwrites a cached value. + tl.debug_barrier() + state_ptr = ( + conv_states + + cache_idx * stride_conv_state_seq + + channel_idx[None, :] * stride_conv_state_dim + + state_pos[:, None] * stride_conv_state_tok + ) + tl.store(state_ptr, tl.where(input_token[:, None] >= 0, tail, history), state_mask) + + +def causal_conv1d_qkv(x, weight, conv_states, query_start_loc, cache_indices, cache_seqlens, max_q_seq_len): + """Convolve projected QKV and return contiguous Q/K/V with Q/K L2 norm. + + x is BF16 [tokens, 3 * head_num * 128] with unit channel stride; + weight is contiguous [3 * head_num * 128, 4]. conv_states and request + metadata follow the KDA backend contract. cache_seqlens gives each + request's cached token count; zero means no history. The first token tile + updates each head's raw-input history after consuming the previous history. + """ + block_t, num_warps = 4, 1 + tokens, qkv_dim = x.shape + head_num, head_dim = qkv_dim // (3 * 128), 128 + q, k, v = [torch.empty((1, tokens, head_num, head_dim), dtype=x.dtype, device=x.device) for _ in range(3)] + stride_conv_state_seq, stride_conv_state_dim, stride_conv_state_tok = conv_states.stride() + seq_num = cache_indices.numel() + # Avoid launching mostly empty tiles for highly uneven request lengths. + packed = max_q_seq_len * seq_num > 2 * tokens + if packed: + grid = (3 * head_num, tokens // block_t + seq_num) + else: + grid = (seq_num, max(1, triton.cdiv(max_q_seq_len, block_t)), 3 * head_num) + _conv_qkv_kernel[grid]( + x, + weight, + conv_states, + query_start_loc, + cache_indices, + cache_seqlens, + q, + k, + v, + head_num, + seq_num if packed else 1, # Only packed grids specialize on batch size. + packed, + head_dim, + x.stride(0), + stride_conv_state_seq, + stride_conv_state_dim, + stride_conv_state_tok, + BLOCK_T=block_t, + num_warps=num_warps, + ) + return q, k, v diff --git a/unit_tests/models/glm5_next/test_kda_conv_qkv.py b/unit_tests/models/glm5_next/test_kda_conv_qkv.py new file mode 100644 index 0000000000..b677477519 --- /dev/null +++ b/unit_tests/models/glm5_next/test_kda_conv_qkv.py @@ -0,0 +1,186 @@ +import weakref +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.common.basemodel.attention.base_att import AttControl +from lightllm.common.basemodel.attention.linear import kda as backend_module +from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d import causal_conv1d_fn +from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d_qkv import ( + causal_conv1d_qkv, +) +from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.kda import chunk_kda_with_fused_gate +from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.l2norm import l2norm_fwd, USE_DEFAULT_FLA_NORM + + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (9, 0) or USE_DEFAULT_FLA_NORM, + reason="requires Hopper and the LightLLM L2 norm", +) + + +def reference(x, weight, states, starts, slots, warm, heads): + mixed = causal_conv1d_fn( + x.clone().T, + weight, + query_start_loc=starts, + cache_indices=slots, + has_initial_state=warm, + conv_states=states, + ).T + q, k, v = [part.view(1, x.shape[0], heads, 128) for part in mixed.chunk(3, dim=-1)] + return l2norm_fwd(q.contiguous()), l2norm_fwd(k.contiguous()), v.contiguous() + + +@pytest.mark.parametrize("heads", [1, 2, 8, 64]) +@pytest.mark.parametrize("state_width", [3, 4, 7]) +@pytest.mark.parametrize("lengths", [(2049, 2, 1024), (257, 1, 1), (1, 2, 3), (0, 1, 0, 2), (1,)]) +def test_conv_norm_preserves_outputs_and_history(heads, state_width, lengths): + torch.manual_seed(42) + channels = heads * 128 * 3 + opts = dict(device="cuda", dtype=torch.bfloat16) + projection = torch.randn(sum(lengths), channels + heads + 256, **opts) + original_projection = projection.clone() + x = projection[:, :channels] + weight = torch.randn(channels, 4, **opts) * 0.125 + states = torch.randn(5, channels, state_width, **opts) + expected_states = states.clone() + starts = torch.tensor([0, *lengths], device="cuda", dtype=torch.int32).cumsum(0, dtype=torch.int32) + slots = torch.tensor([1, 2, -1, 3][: len(lengths)], device="cuda", dtype=torch.int32) + cache_seqlens = torch.tensor([0, 7, 0, 1][: len(lengths)], device="cuda", dtype=torch.int32) + expected = reference(x, weight, expected_states[..., :3], starts, slots, cache_seqlens > 0, heads) + actual = causal_conv1d_qkv(x, weight, states[..., :3], starts, slots, cache_seqlens, max(lengths)) + for a, b in zip(actual, expected): + assert a.is_contiguous() + assert torch.equal(a.view(torch.uint8), b.view(torch.uint8)) + assert len({a.untyped_storage().data_ptr() for a in actual}) == 3 + assert torch.equal(states, expected_states) + assert torch.equal(projection, original_projection) + + +@pytest.mark.parametrize("lengths", [(3, 2, 4), (7, 1, 1)]) +def test_conv_norm_cuda_graph_changing_boundaries(lengths): + torch.manual_seed(42) + heads, channels = 2, 3 * 2 * 128 + opts = dict(device="cuda", dtype=torch.bfloat16) + x = torch.randn(9, channels + 258, **opts)[:, :channels] + weight = torch.randn(channels, 4, **opts) * 0.125 + initial = torch.randn(5, channels, 4, **opts) + states = initial.clone() + initial_boundaries = [0, lengths[0], lengths[0] + lengths[1], sum(lengths)] + starts = torch.tensor(initial_boundaries, device="cuda", dtype=torch.int32) + slots = torch.tensor([1, 2, -1], device="cuda", dtype=torch.int32) + cache_seqlens = torch.tensor([4, 0, 0], device="cuda", dtype=torch.int32) + causal_conv1d_qkv(x, weight, states[..., :3], starts, slots, cache_seqlens, max(lengths)) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + actual = causal_conv1d_qkv(x, weight, states[..., :3], starts, slots, cache_seqlens, max(lengths)) + for boundaries, cached in ((initial_boundaries, [4, 0, 0]), ([0, 4, 8, 9], [0, 7, 0])): + starts.copy_(torch.tensor(boundaries, device="cuda", dtype=torch.int32)) + cache_seqlens.copy_(torch.tensor(cached, device="cuda", dtype=torch.int32)) + states.copy_(initial) + expected_states = initial.clone() + expected = reference(x, weight, expected_states[..., :3], starts, slots, cache_seqlens > 0, heads) + graph.replay() + for a, b in zip(actual, expected): + assert torch.equal(a.view(torch.uint8), b.view(torch.uint8)) + assert torch.equal(states, expected_states) + + +@pytest.mark.parametrize("heads", [2, 8]) +@pytest.mark.parametrize("state_width", [3, 7]) +def test_conv_norm_chained_prefills(heads, state_width): + torch.manual_seed(42) + channels = 3 * heads * 128 + opts = dict(device="cuda", dtype=torch.bfloat16) + weight = torch.randn(channels, 4, **opts) * 0.125 + states = torch.randn(5, channels, state_width, **opts) + expected_states = states.clone() + slots = torch.tensor([1, 2, -1, 3], device="cuda", dtype=torch.int32) + cached = torch.tensor([0, 7, 0, 1], device="cuda", dtype=torch.int32) + for lengths in [(1, 2, 0, 3), (0, 1, 2, 0), (7, 0, 1, 2), (2, 1, 0, 1)]: + x = torch.randn(sum(lengths), channels + heads + 256, **opts)[:, :channels] + starts = torch.tensor([0, *lengths], device="cuda", dtype=torch.int32).cumsum(0, dtype=torch.int32) + expected = reference(x, weight, expected_states[..., :3], starts, slots, cached > 0, heads) + actual = causal_conv1d_qkv(x, weight, states[..., :3], starts, slots, cached, max(lengths)) + for a, b in zip(actual, expected): + assert torch.equal(a.view(torch.uint8), b.view(torch.uint8)) + assert torch.equal(states, expected_states) + cached += torch.tensor(lengths, device="cuda", dtype=torch.int32) + + +@pytest.mark.parametrize("lengths", [(129, 2), (129, 1, 1), (1,)]) +def test_backend_fused_conv_output_and_state(lengths): + torch.manual_seed(42) + heads, channels, tokens = 2, 3 * 2 * 128, sum(lengths) + opts = dict(device="cuda", dtype=torch.bfloat16) + projection = torch.randn(tokens, channels + heads + 256, **opts) + weight = torch.randn(channels, 4, **opts) * 0.125 + conv_states = torch.randn(4, channels, 4, **opts) + ssm_states = torch.randn(4, heads, 128, 128, device="cuda") * 0.01 + linear_state = SimpleNamespace(get_layer_view=lambda **kwargs: (conv_states, ssm_states)) + infer_state = type("State", (), {})() + infer_state.att_state_manager = SimpleNamespace(get_state=lambda **kwargs: linear_state) + infer_state.att_state_role = "main" + infer_state.b1_cu_q_seq_len = torch.tensor([0, *lengths], device="cuda", dtype=torch.int32).cumsum( + 0, dtype=torch.int32 + ) + infer_state.b_ready_cache_len = torch.arange(len(lengths), device="cuda", dtype=torch.int32) % 2 * 7 + infer_state.max_q_seq_len = max(lengths) + backend = SimpleNamespace( + head_dim=128, + tp_num_heads=heads, + tp_hidden_size=heads * 128, + conv_kernel_size=4, + lower_bound=-5.0, + ) + att = backend_module.KDAPrefillAttState(backend=backend, _infer_state_weak_ref=weakref.ref(infer_state)) + att.b_conv_buffer_idx = torch.arange(1, len(lengths) + 1, device="cuda", dtype=torch.int32) + att.b_ssm_buffer_idx = att.b_conv_buffer_idx + layer_weight = SimpleNamespace( + get_merged_kda_conv_weight=lambda: weight, + linear_A_log=SimpleNamespace(weight=torch.zeros(heads, device="cuda")), + linear_dt_bias=SimpleNamespace(weight=torch.zeros(heads * 128, device="cuda")), + ) + control = AttControl( + linear_att_prefill=True, + linear_att_prefill_dict=dict( + layer_weight=layer_weight, + layer_num=0, + mixed_qkv=projection[:, :channels], + raw_gate=torch.randn(tokens, heads * 128, **opts) * 0.1, + raw_beta=torch.randn(tokens, heads, **opts), + ), + ) + expected_conv, expected_ssm = conv_states.clone(), ssm_states.clone() + q, k, v = reference( + projection[:, :channels], + weight, + expected_conv[..., :3], + infer_state.b1_cu_q_seq_len, + att.b_conv_buffer_idx, + infer_state.b_ready_cache_len > 0, + heads, + ) + params = control.linear_att_prefill_dict + expected, final_state = chunk_kda_with_fused_gate( + q=q, + k=k, + v=v, + raw_g=params["raw_gate"].view(1, tokens, heads, 128), + beta=params["raw_beta"].float().sigmoid().view(1, tokens, heads), + A_log=layer_weight.linear_A_log.weight, + g_bias=layer_weight.linear_dt_bias.weight, + initial_state=ssm_states[att.b_ssm_buffer_idx].contiguous(), + output_final_state=True, + use_qk_l2norm_in_kernel=False, + cu_seqlens=infer_state.b1_cu_q_seq_len, + safe_gate=True, + lower_bound=-5.0, + ) + expected_ssm[att.b_ssm_buffer_idx] = final_state.to(ssm_states.dtype) + actual = att.prefill_att(None, None, None, att_control=control) + assert torch.equal(actual.view(torch.uint8), expected.view(torch.uint8)) + assert torch.equal(conv_states, expected_conv) + assert torch.equal(ssm_states, expected_ssm)