fix(dsa): skip invalid slots in sparse indexer backward

[QiZhangNV]

Summary

This PR fixes severe atomic contention caused by padded slots in the cuDNN sparse indexer backward path.

It:

• Marks invalid/padded top-k indices with -1, allowing the kernel's bounds guard to skip their K/dK work.
• Keeps scores for invalid slots at zero, preserving the existing mathematical behavior.
• Replaces torch.where(..., torch.zeros_like(predict)) with an out-of-place clamp_min() followed by an in-place masked_fill_(), avoiding unnecessary full-size temporary tensors.

Problem

The caller previously assigned index 0 to invalid slots. Although their scores were zero, index 0 passed the kernel's bounds check, so the backward kernel still issued atomic dK updates to dK[0, :].

When many padded slots shared index 0, these atomics became heavily contended.

Fix

Invalid positions now use index -1. The sparse backward kernel rejects these positions through its existing bounds guard and therefore skips their K/dK updates.

Valid indices, scores, and gradients are unchanged. The new masked_fill_() operates on the output of clamp_min(), so the original predict tensor is not modified.

Validation

Validated with the GLM-5.2 proxy model on 16 GB200 GPUs:

• Sequence length: 4096
• TP1 / PP1 / EP16 / CP1
• MXFP8
• 28 training iterations
• First 5 iterations excluded from performance statistics
• 0 skipped iterations
• 0 NaN iterations
• Loss and gradient norm remained aligned with the dev baseline

Rank-0 NSys results for the affected backward kernel:

• Before: approximately 12.824 ms per call
• After: approximately 1.273 ms per call
• Improvement: approximately 10x

End-to-end comparison using the same 4K recipe without NSys:

Metric	Before	After	Change
Median iteration latency	3109.9 ms	2802.3 ms	-9.9%
Mean iteration latency	3291.8 ms	2994.3 ms	-9.0%
Median TFLOP/s/GPU	453.1	502.8	+11.0%
Mean TFLOP/s/GPU	431.7	476.1	+10.3%
Peak allocated memory	129708.63 MB	129708.63 MB	unchanged
Peak reserved memory	132596 MB	132636 MB	+40 MB

The small reserved-memory difference is allocator-level run-to-run variation; peak allocated memory is unchanged.

[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.