vllm.model_executor.layers.mamba.ops.ssd_chunk_scan

_drop_racy_wgmma_configs(configs)

Drop the configs that Triton miscompiles on sm90.

With whole warpgroups (num_warps % 4 == 0) and BLOCK_SIZE_M % 64 == 0, tl.dot(cb, x) lowers to wgmma with cb in registers. With num_stages > 1, Triton leaves that wgmma running across the loop back-edge while the next iteration overwrites the registers of cb, which writes NaN/Inf into whole output rows.

Source code in vllm/model_executor/layers/mamba/ops/ssd_chunk_scan.py

def _drop_racy_wgmma_configs(configs: list[triton.Config]) -> list[triton.Config]:
    """Drop the configs that Triton miscompiles on sm90.

    With whole warpgroups (num_warps % 4 == 0) and BLOCK_SIZE_M % 64 == 0,
    `tl.dot(cb, x)` lowers to wgmma with `cb` in registers. With
    num_stages > 1, Triton leaves that wgmma running across the loop
    back-edge while the next iteration overwrites the registers of `cb`,
    which writes NaN/Inf into whole output rows.
    """
    if not (current_platform.is_cuda() and current_platform.is_device_capability(90)):
        return configs
    return [
        c
        for c in configs
        if c.num_stages == 1
        or c.num_warps % 4 != 0
        or c.kwargs["BLOCK_SIZE_M"] % 64 != 0
    ]