vllm.models.beam.nvidia.dflash2_ops

Functions:

• candidate_edges –

  Gather codebooks and score candidate edges without intermediate tensors.

candidate_edges(predecessor_table, successor_table, candidate_ids, unary_logits, hidden, anchor_token_ids, top_k)

Gather codebooks and score candidate edges without intermediate tensors.

Source code in vllm/models/beam/nvidia/dflash2_ops.py

def candidate_edges(
    predecessor_table: torch.Tensor,
    successor_table: torch.Tensor,
    candidate_ids: torch.Tensor,
    unary_logits: torch.Tensor,
    hidden: torch.Tensor,
    anchor_token_ids: torch.Tensor,
    top_k: int,
) -> torch.Tensor:
    """Gather codebooks and score candidate edges without intermediate tensors."""
    batch, steps, _ = candidate_ids.shape
    rank = hidden.shape[-1]
    output = torch.empty(
        (batch, steps, top_k, top_k), dtype=unary_logits.dtype, device=hidden.device
    )
    if batch * steps:
        _candidate_edges_kernel[(batch * steps, top_k)](
            predecessor_table,
            successor_table,
            candidate_ids,
            unary_logits,
            hidden,
            anchor_token_ids,
            output,
            anchor_token_ids.stride(0),
            steps,
            top_k,
            rank,
            triton.next_power_of_2(top_k),
            triton.next_power_of_2(rank),
            num_warps=4,
            enable_fp_fusion=False,
        )
    return output