[Common] Remove redundant grad_logits zero-initialization in fused router backward kernels

[roycho96]

Description

Both fused_topk_with_score_function_backward_kernel and fused_score_for_moe_aux_loss_backward_kernel zero-initialize grad_logits at the beginning of each loop iteration. However, all positions are unconditionally overwritten by the final write at the end of the same iteration, making the zero-initialization a dead store.
Removing it eliminates num_tokens * num_experts * sizeof(dtype) bytes of unnecessary HBM writes per backward kernel call.

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

• Remove redundant grad_logits zero-initialization loop in fused_topk_with_score_function_backward_kernel
• Remove redundant grad_logits zero-initialization loop in fused_score_for_moe_aux_loss_backward_kernel

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[greptile-apps]

Greptile Summary

This PR removes dead-store zero-initialization loops for grad_logits in both fused_topk_with_score_function_backward_kernel and fused_score_for_moe_aux_loss_backward_kernel. The optimization is correct: in both kernels, every element of grad_logits[pos_offset + i] (for all i in [0, num_experts)) is unconditionally overwritten by the final write loop at the end of each iteration, so the preceding zero-init was a wasted HBM write.

• The logic is verified correct: in fused_topk_with_score_function_backward_kernel, the "Backward of topk" section (lines 416–421) explicitly zeroes any non-routed positions in local_grad before the activation backward passes, and the final loop writes all positions. In fused_score_for_moe_aux_loss_backward_kernel the final write loop similarly covers the full [0, num_experts) range unconditionally.
• Both files have stale Init buffer section-header comments that still reference - Clear the global buffer which will accept the result of this round after the zero-init was removed — these should be cleaned up.

Confidence Score: 5/5

• This PR is safe to merge — the removed zero-initialization is verifiably a dead store in both kernels.
• The change is a pure dead-store elimination. Both backward kernels unconditionally write every element of grad_logits through the final loop, making the preceding zero-init redundant. No new logic is introduced and no existing paths are altered. The only minor issue is two stale comments in section headers that can be cleaned up independently.
• No files require special attention.

Important Files Changed

Filename	Overview
transformer_engine/common/fused_router/fused_topk_with_score_function.cu	Removes the dead-store zero-init loop for grad_logits in the backward kernel; the final unconditional write loop covers all num_experts positions, making the removal correct. A stale section-header comment still references the removed "Clear the global buffer" step.
transformer_engine/common/fused_router/fused_score_for_moe_aux_loss.cu	Removes the dead-store zero-init loop for grad_logits in the backward kernel; same pattern as the other file — the final unconditional write loop covers all num_experts positions. A stale section-header comment still references the removed "Clear the global buffer" step.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[Loop iteration begins\ntoken_offset_cur_warp] --> B[Load local_grad from grad_probs\nLoad local_act_from_fwd, routing_map\ninto shared memory]
    B --> C{score_function == 2?}
    C -- Yes --> D[Compute sqrtsoftplus output\ninto local_comp_buf]
    C -- No --> E
    D --> E{topk > 1 AND\nsigmoid or sqrtsoftplus?}
    E -- Yes --> F[Post-processing bwd\nrouted → update grad\nnon-routed → local_grad = 0.0]
    E -- No --> G
    F --> G{!use_pre_softmax AND\nscore_function == 1?}
    G -- Yes --> H[Softmax bwd on topk scores]
    G -- No --> I
    H --> I[Backward of topk:\nset non-routed positions to 0.0]
    I --> J{score_function branches}
    J -- Pre-softmax bwd --> K[apply_softmax_bwd_on_float]
    J -- Sigmoid bwd --> L[apply_sigmoid_bwd_on_float]
    J -- Sqrtsoftplus bwd --> M[apply_sqrtsoftplus_bwd_on_float]
    J -- None --> N
    K & L & M --> N[Final write: grad_logits for ALL num_experts positions\nREMOVED: prior zero-init of grad_logits was dead store here]
    style N fill:#d4edda,stroke:#28a745

Loading

Comments Outside Diff (2)

1. transformer_engine/common/fused_router/fused_topk_with_score_function.cu, line 324-329 (link)

   Stale section comment references removed operation

   The section header still mentions - Clear the global buffer which will accept the result of this round, but that zero-initialization loop was just removed by this PR. The comment should be updated to avoid misleading future readers.

       /***
            * Section: Init buffer
            * - Clear/Init the shmem buffer used by current warp this round
            * - Load the dgrad/output_from_fwd to shmem
            */
   
   
   

2. transformer_engine/common/fused_router/fused_score_for_moe_aux_loss.cu, line 207-213 (link)

   Stale section comment references removed operation

   The section header still mentions - Clear the global buffer which will accept the result of this round, but that zero-initialization loop was just removed by this PR. The comment should be updated to avoid misleading future readers.

       /***
            * Section: Init buffer
            * - Clear/Init the shmem buffer used by current warp this round
            * - Load the dgrad/output_from_fwd to shmem
            */
   
   

_{Last reviewed commit: 343c210}

[ksivaman]

Could you sign-off your commit (guide here)? Other than that, looks good, thanks! @roycho96

[roycho96]

Could you sign-off your commit (guide here)? Other than that, looks good, thanks! @roycho96

Thank you for the review!
I signed off my commit.

[yaox12]

LGTM

[yaox12]

@roycho96 You need to sign-off all your commits, probably via a force push.

[roycho96]

@roycho96 You need to sign-off all your commits, probably via a force push.

Squashed into a single commit with sign-off. Thanks for the review!

[yaox12]

/te-ci