vllm.models.beam.nvidia.activation

Functions:

• beam_silu_mul –

  Compute SiLU and multiplication in FP32, rounding only the product.

• beam_softcap_logits –

  Softcap in FP32 and round back into the original logits buffer.

beam_silu_mul(x)

Compute SiLU and multiplication in FP32, rounding only the product.

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

def beam_silu_mul(x: torch.Tensor) -> torch.Tensor:
    """Compute SiLU and multiplication in FP32, rounding only the product."""
    width = x.shape[-1] // 2
    output = torch.empty((*x.shape[:-1], width), dtype=x.dtype, device=x.device)
    if output.numel():
        _silu_mul_kernel[(triton.cdiv(output.numel(), 1024),)](
            x, output, output.numel(), width, 1024, enable_fp_fusion=False
        )
    return output

beam_softcap_logits(logits, cap)

Softcap in FP32 and round back into the original logits buffer.

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

def beam_softcap_logits(logits: torch.Tensor, cap: float) -> torch.Tensor:
    """Softcap in FP32 and round back into the original logits buffer."""
    assert logits.ndim == 2 and logits.stride(1) == 1
    if logits.numel():
        _softcap_logits_kernel[(logits.shape[0], triton.cdiv(logits.shape[1], 1024))](
            logits,
            logits.shape[1],
            logits.stride(0),
            cap,
            1024,
            enable_fp_fusion=False,
            enable_reflect_ftz=False,
        )
    return logits