vllm.models.kimi_k3.nvidia.ops.attn_res

Functions:

• get_attn_res_triton_warmup_profiles –

  Return the small-batch profiles that bypass the native kernel.

get_attn_res_triton_warmup_profiles(max_blocks)

Return the small-batch profiles that bypass the native kernel.

Source code in vllm/models/kimi_k3/nvidia/ops/attn_res.py

def get_attn_res_triton_warmup_profiles(
    max_blocks: int,
) -> tuple[tuple[int, bool, int, bool], ...]:
    """Return the small-batch profiles that bypass the native kernel."""
    profiles = [
        (num_blocks, False, -1, True) for num_blocks in range(2, max_blocks + 1)
    ]
    profiles.extend(
        (block_write_idx, True, block_write_idx, True)
        for block_write_idx in range(2, max_blocks)
    )
    profiles.append((max_blocks, True, -1, False))
    return tuple(profiles)