Fix token attention kernel

[flyinglandlord]

问题描述

token attention的stage1和stage2中间的att_m_tensor的类型应当是torch.float32而不是torch.float16

测试用例

data = {
        "inputs": """<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n<|im_start|>user\nWhat is the correct answer to this question: trans-cinnamaldehyde was treated with methylmagnesium bromide, forming product 1.\n\n1 was treated with pyridinium chlorochromate, forming product 2.\n\n3 was treated with (dimethyl(oxo)-l6-sulfaneylidene)methane in DMSO at elevated temperature, forming product 3.\n\nhow many carbon atoms are there in product 3?\nChoices:\n(A)12\n(B)14\n(C)11\n(D)10\nFormat your response as follows: "The correct answer is (insert answer here)"<|im_end|>\n<|im_start|>assistant\n""",
        "parameters": {
            "temperature": 1,
            "max_new_tokens": 300,
            "stop_sequences": [
            "<|endofblock|>",
            "<|endofblock|><|im_end|>",
            "<|endoftext|>",
            "<|im_start|>"
            ],
            # "repetition_penalty": 1.05,
            "top_k": 1,
            "best_of": 1,
            "do_sample": True
        }

输出为

{'generated_text': ['The correct answer is (A)12\n'], 'count_output_tokens': 12, 'finish_reason': 'stop', 'prompt_tokens': 149}