Fix zero division risk in embedding _compute_grid when M=0 and update div

Author: ftgreat

src/flag_gems/runtime/backend/_ascend/ops/div.py

@@ -141,17 +141,6 @@ def true_div_func_scalar_tensor(x, y):
     return x / y
 
 
-def _can_use_manual_kernel(A, B):
-    """Check if both tensors have the same shape (no broadcasting needed) and are float types."""
-    if not isinstance(A, torch.Tensor) or not isinstance(B, torch.Tensor):
-        return False
-    if A.shape != B.shape:
-        return False
-    if not A.dtype.is_floating_point:
-        return False
-    return True
-
-
 def true_divide(A, B):
     logger.debug("GEMS_ASCEND TRUE_DIVIDE")
     if isinstance(A, torch.Tensor) and isinstance(B, torch.Tensor):

src/flag_gems/runtime/backend/_ascend/ops/embedding.py

@@ -142,6 +142,8 @@ def _compute_block_sizes(N):
 
 
 def _compute_grid(M):
+    if M == 0:
+        return 0, 0
     ncore = min(M, NUM_VECTOR_CORES)
     rows_per_core = triton.cdiv(M, ncore)
     return ncore, rows_per_core
@@ -158,6 +160,18 @@ def forward(
         M = math.prod(indices.shape)
         N = weight.shape[-1]
 
+        if M == 0:
+            ctx.M = 0
+            ctx.N = N
+            ctx.num_weights = weight.shape[0]
+            ctx.padding_idx = padding_idx
+            ctx.scale_grad_by_freq = scale_grad_by_freq
+            ctx.sparse = sparse
+            ctx.indices = indices
+            return torch.empty(
+                (*indices.shape, N), device=indices.device, dtype=weight.dtype
+            )
+
         BLOCK_SIZE, NUM_ITERS = _compute_block_sizes(N)
         ncore, rows_per_core = _compute_grid(M)
 
@@ -188,6 +202,14 @@ def backward(ctx, grad_outputs):
         logger.debug("GEMS_ASCEND EMBEDDING BACKWARD")
         assert not ctx.sparse, "Currently do not support sparse format"
 
+        if ctx.M == 0:
+            grad_inputs = torch.zeros(
+                (ctx.num_weights, grad_outputs.shape[-1]),
+                device=grad_outputs.device,
+                dtype=grad_outputs.dtype,
+            )
+            return grad_inputs, None, None, None, None
+
         grad_inputs = torch.zeros(
             (ctx.num_weights, grad_outputs.shape[-1]),
             device=grad_outputs.device,