perf(gemm):Optimize deterministic GEMM tensor layouts

[frank-2077]

Summary

This PR optimizes deterministic GEMM layouts used by the Qwen FFN:

• Added layout-aware GEMM entry points that consume canonical [out, in] weights directly.
• Removed Python-side weight transposes and redundant CUDA round-trip copies.
• Made weight-gradient GEMMs write contiguous [out, in] gradients directly, removing three large output transposes.
• Reused the existing physical transpose in the input-gradient path for TMA.
• Added handling for 16-byte TMA alignment, M-padding/cropping, and scalar fallback.
• Preserved the original MMA/K-tree order, legacy GEMM APIs, and TP/CP/SP communication and reduction order.

Performance

Measured on H100 with BF16 and T=128, H=4096, I=12288:

• Forward: 4.0343 ms -> 1.7180 ms (2.35x)
• Forward + backward: 9.4847 ms -> 5.7756 ms (1.64x)
• direct_copy_kernel launches: 18 -> 9
• Forward copy launches: 6 -> 0

Correctness

• FFN output, input gradient, and all three weight gradients are bitwise identical to the baseline.
• Weight gradients remain contiguous with parameter-matching shapes and strides.
• Single-GPU result: 56 passed, 14 skipped
• 2-GPU TP/CP/SP result: 14 passed

Tests

Build the CUDA extension:

KERNEL_ALIGN_DET_GEMM_SM90=1 MAX_JOBS=2 \
python setup.py build_ext --inplace

Run the single-GPU tests:

CUDA_VISIBLE_DEVICES=0 python -m pytest -q \
  tests/test_det_gemm.py \
  tests/distributed/test_det_gemm_simulated_tp.py \
  tests/test_qwen_ffn.py

Run the 2-GPU TP/CP/SP tests:

CUDA_VISIBLE_DEVICES=0,1 python -m pytest -q \
  tests/test_qwen_ffn.py::test_qwen_ffn_tp_correctness_and_batch_invariance \
  tests/test_qwen_ffn.py::test_qwen_ffn_tp_sp_correctness_and_batch_invariance \
  tests/test_qwen_ffn.py::test_qwen_ffn_world2_tp_sp_and_cp_match_tp1_cp1_bitwise

[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: 7473df17-79ea-4dae-9863-fe400123d41a

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

[Flink-ddd]

LGTM, Thank you for update.