Skip to content

viable/strict/1787086487

viable/strict/1787086487: Add missing XPU: _safe_softmax_xpu dispatch for _safe_softmax (#193786)
@msnliu msnliu tagged this 18 Aug 15:37 
### Descriptions:

**Root cause**: Commit [`08c29108`](https://github.com/intel/torch-xpu-ops/commit/08c29108) ("Remove the obsolete folder yaml/") in torch-xpu-ops migrated XPU dispatch registrations from `torch-xpu-ops/yaml/native/native_functions.yaml` to PyTorch's `aten/src/ATen/native/native_functions.yaml`, but missed the `XPU: _safe_softmax_xpu` line.

**Impact**: `aten::_safe_softmax` lost its XPU dispatch in 2.14, falling back to `CompositeExplicitAutograd` decomposition. Each call produces 5 separate XPU kernels (`isneginf` + `all` + `where` + `fill` + `softmax`) instead of 1 fused kernel. This causes 10-27% eager-path performance regression on 11 model configurations (issue https://github.com/intel/torch-xpu-ops/issues/4917).

torch-xpu-ops has an optimized XPU implementation of _safe_softmax (_safe_softmax_xpu in SoftMax.cpp and _safe_softmax_kernel in SoftMaxKernels.cpp) that fuses isneginf + all + where + softmax into a single kernel. This is used by the SDPA math backend. Without the dispatch registration, _safe_softmax falls back to CompositeExplicitAutograd, producing 5 separate kernels per call instead of 1.

YAML diff (what [`08c29108`](https://github.com/intel/torch-xpu-ops/commit/08c29108) removed from torch-xpu-ops, not added to PyTorch):
```diff
- func: _safe_softmax(Tensor self, int dim, ScalarType? dtype=None) -> Tensor
-  dispatch:
-    XPU: _safe_softmax_xpu
-    NestedTensorXPU: _safe_softmax
```

E2E regression reproduced (2.13 release vs 2.14 nightly, 11 configs, all within ±0.03)

### Fix

In `aten/src/ATen/native/native_functions.yaml`, add `XPU: _safe_softmax_xpu`:

```yaml
- func: _safe_softmax(Tensor self, int dim, ScalarType? dtype=None) -> Tensor
  dispatch:
    CompositeExplicitAutograd: _safe_softmax
+   XPU: _safe_softmax_xpu
    NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA, NestedTensorXPU: _safe_softmax
```

The XPU kernel implementation (`_safe_softmax_xpu` in `torch-xpu-ops/src/ATen/native/xpu/SoftMax.cpp` and `_safe_softmax_kernel` in `SoftMaxKernels.cpp`) already exists in the 2.14 wheel — only the dispatch registration was lost.

### Related

- Issue: https://github.com/intel/torch-xpu-ops/issues/4917
- torch-xpu-ops commit that removed the yaml: https://github.com/intel/torch-xpu-ops/commit/08c29108
- torch-xpu-ops commit that later cleaned up dead code: https://github.com/intel/torch-xpu-ops/commit/0276701d19c625221d2527a7374140ca162a1a28 (not in main)
Pull Request resolved: https://github.com/pytorch/pytorch/pull/193786
Approved by: https://github.com/guangyey, https://github.com/Skylion007
Assets 2
Loading