def check_sgmv_expand(
    batches: int,
    num_loras: int,
    rank: int,
    hidden_size: int,
    nslices: int,
    dtype: torch.dtype,
    device: str,
    seq_length: int,
    add_inputs: bool,
    inductor: bool,
):
    """
    Compare outputs of vllm.sgmv_expand kernel against a reference
    implementation.
    """
    ipex_sgmv_expand_op = ipex.llm.functional.fusions.sgmv_expand
    ipex_sgmv_expand_slice_op = ipex.llm.functional.fusions.sgmv_expand_slice
    if inductor:
        torch._dynamo.reset()
        ipex_sgmv_expand_op = torch.compile(ipex_sgmv_expand_op, backend="inductor")
        ipex_sgmv_expand_slice_op = torch.compile(
            ipex_sgmv_expand_slice_op, backend="inductor"
        )

    def ipex_sgmv_expand(
        inputs: torch.Tensor,
        lora_b_weights: torch.Tensor,
        output_tensor: torch.Tensor,
        b_seq_start_loc: torch.Tensor,
        seq_len_tensor: torch.Tensor,
        lora_indices_tensor: torch.Tensor,
        batches: int,
        max_seq_length: int,
        token_nums: int,
        add_inputs: bool = False,
    ):
        ipex_sgmv_expand_op(
            inputs,
            lora_b_weights,
            output_tensor,
            b_seq_start_loc,
            seq_len_tensor,
            lora_indices_tensor,
            batches,
            max_seq_length,
            add_inputs,
        )

    def ipex_sgmv_expand_slice(
        inputs: torch.Tensor,
        lora_b_weights: torch.Tensor,
        output_tensor: torch.Tensor,
        b_seq_start_loc: torch.Tensor,
        seq_len_tensor: torch.Tensor,
        lora_indices_tensor: torch.Tensor,
        batches: int,
        max_seq_length: int,
        token_nums: int,
        slice_offset: int,
        slice_size: int,
        add_inputs: bool = False,
    ):
        ipex.llm.functional.fusions.sgmv_expand_slice(
            inputs,
            lora_b_weights,
            output_tensor,
            b_seq_start_loc,
            seq_len_tensor,
            lora_indices_tensor,
            batches,
            max_seq_length,
            slice_offset,
            slice_size,
            add_inputs,
        )

    def ipex_sgmv_expand_for_nslices(
        nslices: int,
        hidden_size: int,
        inputs_tensor: torch.Tensor,
        lora_weights_lst: list[torch.Tensor],
        out_tensor: torch.Tensor,
        b_seq_start_loc: torch.Tensor,
        seq_len_tensor: torch.Tensor,
        prompt_lora_mapping: torch.Tensor,
        batches: int,
        max_seq_length: int,
        num_tokens: int,
        add_inputs: bool,
    ) -> None:
        """
        Wrapper around sgmv_expand that handles any nslices.
        """
        if nslices == 1:
            # Verify the torch's sgmv_expand op
            ipex_sgmv_expand(
                inputs_tensor[0],
                lora_weights_lst[0],
                out_tensor,
                b_seq_start_loc,
                seq_len_tensor,
                prompt_lora_mapping,
                batches,
                max_seq_length,
                num_tokens,
                add_inputs=add_inputs,
            )
        else:
            slice_offset = 0
            for index in range(nslices):
                lora_weights = lora_weights_lst[index]
                ipex_sgmv_expand_slice(
                    inputs_tensor[index],
                    lora_weights,
                    out_tensor,
                    b_seq_start_loc,
                    seq_len_tensor,
                    prompt_lora_mapping,
                    batches,
                    max_seq_length,
                    num_tokens,
                    slice_offset,
                    hidden_size,
                    add_inputs=add_inputs,
                )
                slice_offset += hidden_size

    data: PunicaTensors = generate_data_for_nslices(
        batches,
        hidden_size,
        num_loras,
        rank,
        seq_length,
        nslices,
        dtype,
        "expand",
        device,
    )

    max_seq_length, token_nums = data.meta()

    data.our_out_tensor = data.our_out_tensor.to(dtype)
    data.ref_out_tensor = data.ref_out_tensor.to(dtype)
    ipex_sgmv_expand_for_nslices(
        nslices,
        hidden_size,
        data.inputs_tensor,
        data.lora_weights,
        data.our_out_tensor,
        data.b_seq_start_loc,
        data.seq_len_tensor,
        data.prompt_lora_mapping,
        batches,
        max_seq_length,
        token_nums,
        add_inputs=add_inputs,
    )

    sgmv_expand_for_nslices(
        nslices,
        hidden_size,
        data.inputs_tensor,
        data.lora_weights,
        data.ref_out_tensor,
        data.b_seq_start_loc,
        data.seq_len_tensor,
        data.prompt_lora_mapping,
        batches,
        max_seq_length,
        token_nums,
        add_inputs=add_inputs,
    )

    assert torch.allclose(
        data.ref_out_tensor, data.our_out_tensor, atol=ATOL, rtol=RTOL
    )

    Compare outputs of vllm.sgmv_expand kernel against a reference
    implementation.
    
tests/cpu/test_punica.py