Preserve tiny sigmoid scores in biased fused routing

[philipcmonk]

Description

Gather selected sigmoid scores from the saved unbiased intermediate instead of subtracting expert bias after top-k. Adding the bias can round away tiny scores, which subtraction cannot recover.

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

When using sigmoid score function, use saved original score instead of (score + bias) - bias which can cause unnecessary rounding.

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

cc @deepakn94 @abhinav-khattar @RobMcH

[greptile-apps]

Confidence Score: 5/5

[Medium risk] Fixes score preservation in the fused routing kernel.

The PR appears safe to merge; no blocking issue was found.

What we checked:

• Saved scores are ready: Both kernels save every raw sigmoid score before selecting experts. Their __syncwarp() calls order the writes before the selected-score reads.

Summary

Biased sigmoid routing now gathers selected scores from the saved unbiased values instead of subtracting the bias. This preserves tiny scores that adding the bias can round away.

• Updates both forward kernels without changing expert selection.
• Adds checks for scores, routing layouts, histograms, and gradients.
• No actionable issues found. Tests were not run during this review.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  A[Raw sigmoid scores] --> B[Save unbiased scores]
  A --> C[Add expert bias]
  C --> D[Select experts]
  B --> E[Gather selected raw scores]
  D --> E
  E --> F[Normalize and scale]

Loading

_{Reviews (2) · Last reviewed commit: "Merge branch 'main' into fix-sigmoid-sco..."}

[ksivaman]

/te-ci pytorch L0