[CuTe] Fix forward dynamic-shape correctness

[drisspg]

Stacked PRs:

• [CuTe] Add explicit forward config autotuning #2743
• [CuTe] Tune B200 forward scheduling #2741
• [CuTe] Tune GB300 forward scheduling #2738
• [CuTe] Add explicit forward config benchmark #2735
• [CuTe] Route forward through typed configs #2733
• [CuTe] Add typed forward configurations #2732
• ->[CuTe] Fix forward dynamic-shape correctness #2745

---

Human Note

This is baiscally all small tweaks to our dynamic shape infra; I did a lot of fuzzing of the stack to get a good sense for when recompiles happen and passing in a bunch of configs. These are all the result of that. While I was here I did smallish tweaks to those helpers to remove host overhead where applicable

Agent Note

Summary

This is the small correctness base for the rest of the stack. Dynamic-shape and layout fuzzing found seven pre-existing forward bugs; this PR fixes them without adding configs, tuning APIs, or selector policy.

• Canonicalize unsupported strides and effective pointer alignment before TVM-FFI launch; reject broadcast output/workspace layouts.
• Include static tensor broadcast patterns in the main compile key.
• Key aux tensors by dtype, assumed alignment, and static stride modes.
• Compile SplitKV dynamic-split and semaphore operands exactly as called while retaining the lightweight existing runtime checks.
• Use matching target-SKU SM metadata during fake selection, with an explicit cross-compilation override.
• Size aliased SM100 K/V shared storage for the larger staged layout.
• Ceil-divide non-TMA paged-loader entries so partial row waves receive page pointers.

The cache fixes still keep concrete batch and sequence lengths dynamic.

Host-path cost

The initial fuzz fix was correct but too defensive on the eager hot path. This revision keeps the same cache identities and kernel behavior while removing work unrelated to the seven root bugs:

• contiguous input/output layouts return before scanning every stride;
• no-broadcast metadata uses a C-level zero-stride fast path;
• the added same-device pass was removed, leaving the pre-existing CUDA-input contract untouched;
• internal combine no longer repeats rank/shape/dtype/device/layout checks after _flash_attn_fwd has allocated or validated those tensors;
• target-SM discovery runs only when automatic or rewritten diff-head SplitKV actually needs it;
• fake-mode state is resolved once per runtime call, and fake stream/aux CuTe conversion work runs only on compile misses;
• inference and training retain the original unconditional tensor detach behavior.

Nightly host microbenchmarks on the same GB300 machine:

Host path	Before	After	Saved
Q/K/V canonicalization	1.50 us	0.75 us	0.75 us
Explicit out + LSE layout checks	0.92 us	0.75 us	0.17 us
Common forward broadcast metadata	1.70 us	1.00 us	0.70 us
Internal combine validation	3.971 us	1.168 us	2.803 us
Added same-device pass	+0.94 us	removed	0.94 us
Fixed-split target-SM lookup	2.28 us	skipped	2.28 us

These are host-only measurements; no selector policy or generated kernel changed.

For a whole-wrapper check, I used real preallocated CUDA tensors, a warm hit-only compile cache, and a no-op compiled callable. This measures all Python dispatch through the point where the generated kernel would launch, without mixing in GPU execution:

Shape	main	This PR	Net cost
Decode: B1, Q1, K128, H32/Hkv8, D128	16.951 us	12.922 us	-4.029 us (23.8%)
Short: B2, Q64, K128, H16, D64	16.608 us	12.669 us	-3.939 us (23.7%)
Prefill: B1, Q257, K2048, H32, D64	16.993 us	12.873 us	-4.120 us (24.2%)

Despite adding the correctness checks, this PR is 3.9–4.1 us (23.7–24.2%) faster than main on the warm host path because it removes unconditional target-SM discovery, reuses fake-mode state, defers compile-only setup, and uses fast common-layout metadata paths. Real no-graph runs with the actual kernel also showed no latency regression.

Review follow-ups

• The architecture resolved by _flash_attn_fwd is forwarded to SplitKV combine, so CPU fake compilation with explicit _arch never probes unavailable hardware.
• Real tensors use their effective data_ptr() for alignment; FakeTensor retains metadata-only offset checking.
• Writable outputs and SplitKV workspaces reject every stride-zero mode, including singleton dimensions, so output ABI does not need another compile-key projection.
• The combine regression now exercises cached transitions across neither/dynamic/semaphore/both optional operands, and the unaligned-layout regression now covers public backward gradients.
• The two aux-layout review threads are resolved: CuTe owns conversion-time inference, while its __cache_key__ remains the regression-test oracle rather than a 9–10 us/tensor runtime conversion.

Validation

Focused host tests cover target-SM fallback, CuTe aux-key equivalence, external-base pointer alignment, singleton broadcast-output rejection, and explicit-architecture fake SplitKV. The combine regression now exercises one cache in the order neither optional operand → dynamic splits only → semaphore only → both, checking four specializations and numerical results. The existing unaligned-input test now checks public forward and dq/dk/dv for both padded-stride and unaligned-offset views.

The targeted GPU regressions pass on GB300/SM103 against independent FP32 references where applicable. Compute Sanitizer previously reported ERROR SUMMARY: 0 errors across the wider q-stage, paged-tail, SplitKV, and DV > D matrix. Ruff, formatting, py_compile, and git diff --check pass.