sum operator crashes with ZeroDivisionError on empty tensors

[xufeng102403]

Describe the bug

Title: sum operator crashes with ZeroDivisionError on empty tensors

Summary

The sum() and sum_out() functions in flag_gems/ops/sum.py crash with a ZeroDivisionError when the input tensor is empty (i.e., numel() == 0). Native torch.sum() handles empty tensors gracefully — returning a scalar 0 tensor — but the flag_gems implementation divides by zero when computing block_size and mid_size.

Environment

• flag_gems version: 5.3.0
• Reproduction: torch.sum(torch.empty(1, 0, device='flagos')) or any tensor where numel() == 0

Minimal Reproduction

import torch
x = torch.empty(1, 0, dtype=torch.bool, device='flagos')
torch.sum(x)  # ZeroDivisionError

import torch
x = torch.empty(0, device='flagos')
torch.sum(x)  # ZeroDivisionError

Error Traceback

flag_gems/ops/sum.py:71: in sum
    block_size = triton.next_power_of_2(math.ceil(math.sqrt(M)))
triton/__init__.py:68: in cdiv
    return (x + y - 1) // y
E   ZeroDivisionError: integer division or modulo by zero

Root Cause

In sum() (line 62–81) and sum_out() (line 84–100), when the input tensor has numel() == 0, the following chain leads to a division by zero:

M = inp.numel()                                        # M = 0
block_size = triton.next_power_of_2(math.ceil(math.sqrt(M)))  # math.ceil(math.sqrt(0)) = 0
                                                               # triton.next_power_of_2(0) = 0
mid_size = triton.cdiv(M, block_size)                  # cdiv(0, 0) → (0 + 0 - 1) // 0 → ZeroDivisionError

The native PyTorch behavior is:

>>> torch.sum(torch.empty(1, 0, dtype=torch.bool))
tensor(0)   # dtype=int64

Where the Bug Is (code)

File: flag_gems/ops/sum.py

sum() function (lines 62–81):

def sum(inp, *, dtype=None):
    logger.debug("GEMS SUM")
    inp = inp.contiguous()
    M = inp.numel()                                    # ← M = 0 for empty tensors
    if dtype is None:
        dtype = inp.dtype
        if dtype is torch.bool:
            inp = inp.to(torch.int64)
            dtype = torch.int64
    block_size = triton.next_power_of_2(math.ceil(math.sqrt(M)))  # ← becomes 0
    mid_size = triton.cdiv(M, block_size)                          # ← cdiv(0, 0) crashes
    ...

sum_out() function (lines 84–100) — same issue:

def sum_out(inp, *, dtype=None, out):
    logger.debug("GEMS SUM_OUT")
    M = inp.numel()                                    # ← M = 0 for empty tensors
    if dtype is None:
        dtype = inp.dtype
        if dtype is torch.bool:
            inp = inp.to(torch.int64)
            dtype = torch.int64
    block_size = triton.next_power_of_2(math.ceil(math.sqrt(M)))  # ← becomes 0
    mid_size = triton.cdiv(M, block_size)                          # ← cdiv(0, 0) crashes
    ...

Note: sum_dim() (line 329) already has an empty-tensor guard at line 332 (if inp.numel() == 0: ...), so it is not affected. Only the scalar sum() and sum_out() are missing this guard.

Proposed Fix

Add an early-return guard for empty tensors in both sum() and sum_out(), matching the pattern already used in sum_dim().

sum() — insert after dtype resolution, before block_size computation:

def sum(inp, *, dtype=None):
    logger.debug("GEMS SUM")
    inp = inp.contiguous()
    M = inp.numel()
    if dtype is None:
        dtype = inp.dtype
        if dtype is torch.bool:
            inp = inp.to(torch.int64)
            dtype = torch.int64

    # NEW: empty tensor guard — torch.sum(empty) returns scalar 0
    if M == 0:
        return torch.zeros([], dtype=dtype, device=inp.device)

    block_size = triton.next_power_of_2(math.ceil(math.sqrt(M)))
    ...

sum_out() — insert after dtype resolution, before block_size computation:

def sum_out(inp, *, dtype=None, out):
    logger.debug("GEMS SUM_OUT")
    M = inp.numel()
    if dtype is None:
        dtype = inp.dtype
        if dtype is torch.bool:
            inp = inp.to(torch.int64)
            dtype = torch.int64

    # NEW: empty tensor guard — fill out with 0 and return
    if M == 0:
        out.fill_(0)
        return out

    block_size = triton.next_power_of_2(math.ceil(math.sqrt(M)))
    ...

Expected Behavior After Fix

>>> import torch
>>> torch.sum(torch.empty(1, 0, dtype=torch.bool, device='flagos'))
tensor(0, device='flagos:0')   # scalar 0, dtype=int64

>>> torch.sum(torch.empty(0, device='flagos'))
tensor(0, device='flagos:0')

Real-World Impact

This bug was discovered while running the HuggingFace Transformers test suite with TRANSFORMERS_TEST_BACKEND=torch_fl and TRANSFORMERS_TEST_DEVICE=flagos. In _assisted_decoding (transformers/generation/utils.py:3667), when the assistant model produces zero new tokens in a candidate round, the expression:

n_matches = ((~(candidate_new_tokens == selected_tokens[:, :-1])).cumsum(dim=-1) < 1).sum()

operates on an empty (1, 0) tensor, and the trailing .sum() call hits this bug. The same issue was also present in cumsum (already fixed separately).