[Fix] Non-16-aligned D_qk scale + Autograd non-contiguous gradient (#1669)

[wch0810]

Summary

Fixes #1669.

Bug 1: Non-16-aligned D_qk returns wrong gradients

run_bwd used padded dim_qk_padded (e.g. 144) for softmax scale instead of logical D_qk (e.g. 129), causing sm_scale = 1/sqrt(144) instead of 1/sqrt(129).

Fix: run_bwd accepts explicit dim_qk parameter. Phase 2 (softmax) uses logical dim_qk for sm_scale; Phase 1/3 use dim_qk_padded for buffer shapes. The Autograd wrapper passes ctx.D_qk through.

Bug 2: PyTorch Autograd non-contiguous gradient

output.sum().backward() produces a zero-stride gradient view. The wrapper did not materialize it, causing ValueError: Input tensor must be contiguous.

Fix: backward() entry adds if not do.is_contiguous(): do = do.contiguous().

Regression Tests Added

1. D_qk=129: gradient comparison via public attention wrapper (padded to 144, scale uses 129)
2. output.sum().backward(): Autograd compatibility with zero-stride gradient

Verification

Test	Result
L0 (6 cases)	6/6 PASS
L1 (10 cases)	10/10 PASS
L2 (3 negative)	3 correctly rejected
Boundary (4)	3 PASS + 1 WARN
Regression #1669	D_qk=129 PASS + Autograd PASS
do_bench perf	18423 us (no regression)
ruff format + check	All checks passed

Based on #1623 (Developer-mode 3-sub-kernel rewrite).

[github-actions]

👋 Hi! Thank you for contributing to the TileLang project.

Please remember to run bash format.sh in the root directory of the project to ensure your changes are properly linted and formatted. This will help ensure your contribution passes the format check.

We appreciate you taking this step! Our team will review your contribution, and we look forward to your awesome work!

🚀