[Fix] fix num_kv_split becomes 0

[xjmxyt]

Description

Previously, we use NUM_KV_SPLITS = NUM_SMS // (batch_size * num_kv_heads), it has possibility to become 0.
Now change to max(1, NUM_SMS // (batch_size * num_kv_heads), it is at least 1.

CI Configuration

config:
  build: true
  # valid options are "ops" and "benchmark"
  test: ["ops", "benchmark"]

Checklist

• Code formatted and imports sorted via repo specifications (./format.sh)
• Documentation updated (if needed)
• CI configuration reviewed

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[xjmxyt]

/ok to test 008393e

[xjmxyt]

/ok to test 8245eda