fix: prevent division by zero in soft_max vulkan shader

[gn64]

Fix division by zero error in soft_max vulkan shader

This PR fixes #2596 by adding a check for p.KY being zero in the soft_max compute shader.

Changes:

• Modified the calculation of rowy in soft_max.comp to handle the case where p.KY is 0
• Changed const uint rowy = rowx % p.KY; to const uint rowy = (p.KY > 0) ? (rowx % p.KY) : 0;

The original code would cause a division by zero error when p.KY is 0. This fix ensures that rowy is set to 0 in such cases, preventing the crash while maintaining the expected behavior in normal scenarios.

Testing:

• Confirmed the fix resolves the crash in my environment
• Existing functionality remains unchanged when p.KY > 0

[slaren]

The Vulkan fix looks correct, this can happen when using ggml_soft_max. @0cc4m please take a look.

[0cc4m]

Looks good, but what are the other two changes?

[ggerganov]

@gn64 Please update the PR to include just the Vulkan fix.

@0cc4m In which cases p.KY == 0? We should add a test to test-backend-ops to cover these.

[gn64]

Closing this PR in favor of a more focused fix in #2633.
The new PR contains only the Vulkan shader division by zero fix, making the changes easier to review and maintain.

[ggerganov]

Thank you. I'm still interested in which cases p.KY == 0 - I'm not able to deduce by looking at the code.

[slaren]

I believe it happens when using using a NULL mask, which happens when using the plain ggml_soft_max.

    ggml_vk_op_f32<vk_op_soft_max_push_constants>(ctx, subctx, src0, src1, nullptr, dst, GGML_OP_SOFT_MAX, {
        ncols,
        src1 != nullptr ? nrows_y : (uint32_t)0, // this is KY
        scale, max_bias,
        m0, m1,
        n_head_log2,
        nrows_x,
    }, dryrun);

[ggerganov]

That's what I thought as well, but we do have tests in test-backend-ops without a mask, so I am not sure why they are passing.

[slaren]

Division by zero is likely just undefined behavior, and in some GPUs it may just return zero. Still, can't rely on that behavior.