Feat/strict embedding backward

[leonardHONG]

Depends on #339. Please merge #339 before this PR.

Summary

• Replace the vocab × hidden backward grid with a stable sorted-segment reduction over active token rows.
• Keep duplicate-token accumulation deterministic, atomic-free, and bitwise aligned with the ascending FP32 left-fold reference.
• Canonicalize strict Qwen embedding gradients by logical row order.

The launch grid is reduced from vocab × hidden to tokens × ceil(hidden / 128).

Performance

The following measurements time exactly one Triton backward kernel launch. Input sorting, allocation, and output zeroing are excluded.

H100 80GB, BF16, PyTorch 2.9.1+cu128, Triton 3.5.1:

Tokens	Vocab	Hidden	Previous kernel	New kernel	Speedup
8	32,000	1,024	19.725 ms	0.0256 ms	771.5×
32	151,936	1,024	93.581 ms	0.0240 ms	3,904.4×
32	151,936	2,048	187.142 ms	0.0237 ms	7,892.3×
32	32,000	4,096	78.846 ms	0.0244 ms	3,231.4×
32	128,256	4,096	315.886 ms	0.0234 ms	13,504.0×
32	151,936	4,096	374.177 ms	0.0232 ms	16,161.8×
128	151,936	4,096	944.681 ms	0.0236 ms	40,083.2×
512	151,936	4,096	3,698.262 ms	0.0270 ms	136,851.0×

Validation

• Bitwise equal to the per-token FP32 left-fold reference in original row order across all benchmark configurations.
• Unused embedding rows remain zero after output initialization in the full backward path.
• Covers duplicate IDs, empty inputs, invalid IDs, layout changes, repeated execution, and canonical chunk/permutation invariance.

[coderabbitai]

Important

Review skipped

Auto reviews are disabled on base/target branches other than the default branch.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro Plus

Run ID: b1fe84f3-0fc2-4085-9fb7-b6f0cf52b3a6

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

---

_{Comment @coderabbitai help to get the list of available commands.}