fix kernel impl

Author: wangzaijun

lightllm/models/llama/triton_kernel/gqa_flash_decoding_stage1.py

@@ -79,7 +79,7 @@ def _fwd_kernel_flash_decode_stage1(
         ).to(tl.int64)
         off_k = k_loc[None, :] * stride_kbs + cur_kv_head * stride_kh + offs_d[:, None]
         k = tl.load(K + off_k, mask=offs_n_new[None, :] < cur_batch_end_index, other=0.0)
-        att_value = tl.dot(q, k)
+        att_value = tl.dot(q, k.to(q.dtype))
         att_value *= sm_scale
         att_value = tl.where(offs_n_new[None, :] < cur_batch_end_index, att_value, float("-inf"))
         v = tl.load(

lightllm/models/llama/triton_kernel/ppl_int8kv_flash_decoding_diverse_stage1.py

@@ -77,7 +77,7 @@ def _fwd_kernel_flash_decode_stage1(
 
     offs_n = cur_batch_start_index + tl.arange(0, BLOCK_N)
     Q_BATCH_HEAD_NUM: tl.constexpr = BLOCK_BATCH * BLOCK_HEAD
-    q = tl.load(Q + off_q, other=0.0).view(Q_BATCH_HEAD_NUM, BLOCK_HEADDIM)
+    q = tl.load(Q + off_q).reshape(Q_BATCH_HEAD_NUM, BLOCK_HEADDIM)
 
     sum_exp = tl.zeros([Q_BATCH_HEAD_NUM], dtype=tl.float32)
     max_logic = tl.zeros([Q_BATCH_HEAD_NUM], dtype=tl.float32) - float("inf")
@@ -88,17 +88,17 @@ def _fwd_kernel_flash_decode_stage1(
         n_mask = offs_n_new < cur_batch_end_index
         k_loc = tl.load(
             Req_to_tokens + stride_req_to_tokens_b * cur_batch_req_idx + offs_n_new,
-            mask=offs_n_new < cur_batch_end_index,
+            mask=n_mask,
             other=0,
         ).to(tl.int64)
         off_k = k_loc[None, :] * stride_kbs + cur_kv_head * stride_kh + offs_d[:, None]
         off_k_scale = off_k // KV_QUANT_GROUP_SIZE
         k = tl.load(K + off_k, mask=n_mask[None, :], other=0)
         k_scale = tl.load(K_scale + off_k_scale, mask=n_mask[None, :], other=0.0)
         k = k * k_scale
-        att_value = tl.dot(q, k)
+        att_value = tl.dot(q, k.to(q.dtype))
         att_value *= sm_scale
-        att_value = tl.where(offs_n_new[None, :] < cur_batch_end_index, att_value, -1000000000.0)
+        att_value = tl.where(n_mask[None, :], att_value, float("-inf"))
         v = tl.load(
             V + off_k.T,
             mask=n_mask[:, None],
@@ -117,7 +117,7 @@ def _fwd_kernel_flash_decode_stage1(
         exp_logic = tl.exp(att_value - new_max_logic[:, None])
         logic_scale = tl.exp(max_logic - new_max_logic)
         acc *= logic_scale[:, None]
-        acc += tl.dot(exp_logic.to(v.dtype), v)
+        acc += tl.dot(exp_logic.to(q.dtype), v.to(q.dtype))
 
         sum_exp = sum_exp * logic_scale + tl.sum(exp_logic, axis=1)
         max_logic = new_max_logic
@@ -135,11 +135,11 @@ def _fwd_kernel_flash_decode_stage1(
         )
         tl.store(
             Mid_O + off_mid_o,
-            (acc / sum_exp[:, None]).view(BLOCK_BATCH, BLOCK_HEAD, BLOCK_HEADDIM),
+            (acc / sum_exp[:, None]).reshape(BLOCK_BATCH, BLOCK_HEAD, BLOCK_HEADDIM),
         )
         tl.store(
             Mid_O_LogExpSum + off_mid_o_logexpsum,
-            (max_logic + tl.log(sum_exp)).view(BLOCK_BATCH, BLOCK_HEAD),
+            (max_logic + tl.log(sum_exp)).reshape(BLOCK_BATCH, BLOCK_HEAD),
         )
     return
 
@@ -169,6 +169,7 @@ def flash_decode_stage1(
     b_mark_shared_group 中每一个不为0的位置都代表其与前面多少个请求形成一个共享前缀组。属于
     同一个共享前缀组的请求, 其在对应的 b_shared_seq_len 中的内容必然相同。
     """
+    assert q.dim() == 3 and k.dim() == 3 and v.dim() == 3
     BLOCK_SEQ = block_seq
     BLOCK_N = 16
     assert BLOCK_SEQ % BLOCK_N == 0
@@ -182,6 +183,7 @@ def flash_decode_stage1(
     gqa_group_size = q.shape[1] // k.shape[1]
     assert triton.next_power_of_2(Lk) == Lk
     KV_QUANT_GROUP_SIZE = v.shape[-1] // v_scale.shape[-1]
+    assert KV_QUANT_GROUP_SIZE == 8
     BLOCK_HEAD = triton.next_power_of_2(gqa_group_size)
     BLOCK_BATCH = triton.next_power_of_2(max_batch_group_size)
     if BLOCK_HEAD * BLOCK_BATCH < 16:
@@ -198,7 +200,7 @@ def flash_decode_stage1(
         stride_kh=k.stride(1),
         stride_kd=k.stride(2),
         V=v,
-        V_scale=v,
+        V_scale=v_scale,
         stride_vbs=v.stride(0),
         stride_vh=v.stride(1),
         stride_vd=v.stride(2),

unit_tests/models/llama/test_ppl_int8kv_flash_decoding_diverse_stage1.py

@@ -0,0 +1,106 @@
+import pytest
+import torch
+from lightllm.models.llama.triton_kernel.ppl_int8kv_flash_decoding_diverse_stage1 import flash_decode_stage1
+
+
+@pytest.fixture
+def setup_tensors():
+    batch_size = 4
+    num_heads = 4
+    kv_head_num = 1
+    seq_len = 256
+    head_dim = 128
+    max_len_in_batch = seq_len
+    block_seq = 256
+    max_batch_group_size = 4
+    quant_group_size = 8
+
+    test_dtype = torch.float32
+
+    kv_shape = (batch_size * seq_len, kv_head_num, head_dim)
+    kv_scale_shape = (batch_size * seq_len, kv_head_num, head_dim // quant_group_size)
+
+    q = torch.randn(size=(batch_size, num_heads, head_dim), dtype=test_dtype, device="cuda")
+    k = torch.randint(low=-100, high=100, size=kv_shape, dtype=torch.int8, device="cuda")
+    k_scale = torch.ones(size=kv_scale_shape, dtype=test_dtype, device="cuda")
+    v = torch.randint(low=-100, high=100, size=kv_shape, dtype=torch.int8, device="cuda")
+    v_scale = torch.ones(size=kv_scale_shape, dtype=test_dtype, device="cuda")
+    Req_to_tokens = torch.arange(0, seq_len * batch_size, dtype=torch.int32, device="cuda").view(batch_size, seq_len)
+    B_req_idx = torch.arange(batch_size, dtype=torch.int32, device="cuda")
+    b_shared_seq_len = torch.full((batch_size,), seq_len, dtype=torch.int32, device="cuda")
+    b_mark_shared_group = torch.ones(batch_size, dtype=torch.int32, device="cuda")
+    mid_out = torch.zeros(
+        size=(batch_size, num_heads, (seq_len // block_seq) + 2, head_dim), dtype=q.dtype, device="cuda"
+    )
+    mid_out_logsumexp = torch.zeros(
+        size=(batch_size, num_heads, (seq_len // block_seq) + 2), dtype=q.dtype, device="cuda"
+    )
+
+    return {
+        "q": q,
+        "k": k,
+        "k_scale": k_scale,
+        "v": v,
+        "v_scale": v_scale,
+        "Req_to_tokens": Req_to_tokens,
+        "B_req_idx": B_req_idx,
+        "b_shared_seq_len": b_shared_seq_len,
+        "b_mark_shared_group": b_mark_shared_group,
+        "max_len_in_batch": max_len_in_batch,
+        "mid_out": mid_out,
+        "mid_out_logsumexp": mid_out_logsumexp,
+        "block_seq": block_seq,
+        "max_batch_group_size": max_batch_group_size,
+    }
+
+
+def test_flash_decode_stage1_execution(setup_tensors):
+    flash_decode_stage1(
+        q=setup_tensors["q"],
+        k=setup_tensors["k"],
+        k_scale=setup_tensors["k_scale"],
+        v=setup_tensors["v"],
+        v_scale=setup_tensors["v_scale"],
+        Req_to_tokens=setup_tensors["Req_to_tokens"],
+        B_req_idx=setup_tensors["B_req_idx"],
+        b_shared_seq_len=setup_tensors["b_shared_seq_len"],
+        b_mark_shared_group=setup_tensors["b_mark_shared_group"],
+        max_len_in_batch=setup_tensors["max_len_in_batch"],
+        mid_out=setup_tensors["mid_out"],
+        mid_out_logsumexp=setup_tensors["mid_out_logsumexp"],
+        block_seq=setup_tensors["block_seq"],
+        max_batch_group_size=setup_tensors["max_batch_group_size"],
+    )
+
+    q = setup_tensors["q"]
+    k = setup_tensors["k"]
+    v = setup_tensors["v"]
+    true_mid_out = torch.zeros_like(setup_tensors["mid_out"])
+    true_mid_out_logsumexp = torch.zeros_like(setup_tensors["mid_out_logsumexp"])
+    new_q = q
+    new_k = k.to(q.dtype)
+    new_v = v.to(q.dtype)
+
+    from lightllm.models.llama.triton_kernel.gqa_flash_decoding_stage1 import (
+        flash_decode_stage1 as gqa_flash_decode_stage1,
+    )
+
+    gqa_flash_decode_stage1(
+        q=new_q,
+        k=new_k,
+        v=new_v,
+        Req_to_tokens=setup_tensors["Req_to_tokens"],
+        B_req_idx=setup_tensors["B_req_idx"],
+        B_Seqlen=setup_tensors["b_shared_seq_len"],
+        max_len_in_batch=setup_tensors["max_len_in_batch"],
+        mid_out=true_mid_out,
+        mid_out_logsumexp=true_mid_out_logsumexp,
+        block_seq=setup_tensors["block_seq"],
+    )
+    print(setup_tensors["mid_out"][0:4, 0, 0, 0], true_mid_out[0:4, 0, 0, 0])
+    assert torch.allclose(
+        setup_tensors["mid_out"][0:4, 0, 0, 0], true_mid_out[0:4, 0, 0, 0], atol=1e-2
+    ), "Mid output does not match expected values"
+    assert torch.allclose(
+        setup_tensors["mid_out_logsumexp"], true_mid_out_logsumexp, atol=1e-2
+    ), "LogSumExp output does not match expected values"