Function worker

test/kernel/deepseekv2_gqa_decode_tuning.py:130–164

(
    q_nope_shape: List[int],
    q_rope_shape: List[int],
    kv_nope_shape: List[int],
    kv_rope_shape: List[int],
    test_seq_len: int,
    dtype: torch.dtype,
    test_count: int,
    test_configs,
    queue,
)

def worker(
    q_nope_shape: List[int],
    q_rope_shape: List[int],
    kv_nope_shape: List[int],
    kv_rope_shape: List[int],
    test_seq_len: int,
    dtype: torch.dtype,
    test_count: int,
    test_configs,
    queue,
):
    dog = Watchdog(timeout=10)
    dog.start()

    try:
        for index in range(len(test_configs)):
            tuning_config = test_configs[index]
            cost_time = test_decode_attentions(
                q_nope_shape=q_nope_shape,
                q_rope_shape=q_rope_shape,
                kv_nope_shape=kv_nope_shape,
                kv_rope_shape=kv_rope_shape,
                test_seq_len=test_seq_len,
                dtype=dtype,
                test_count=test_count,
                **tuning_config,
            )
            dog.heartbeat()
            queue.put(cost_time)  # Put result in queue
    except Exception as ex:
        logger.error(str(ex) + f"config {tuning_config}")
        import sys

        sys.exit(-1)
        pass


def get_test_configs(split_id, split_count):