[CuTe] Route forward through typed configs

[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

Agent Note

Summary

This wires the typed config model into the real forward path without changing the public wrappers. _flash_attn_fwd becomes the one place that extracts host metadata, selects or validates one config, and launches it.

config=None preserves normal policy selection. Passing an explicit FwdConfig means use exactly that config or raise. Generic SM100 and dedicated HD256 consume the resolved named register allocation directly, so kernel code cannot silently replace an explicit choice. SplitKV callers may also provide reusable partial O/LSE workspaces for allocation-free replay.

Host-path cost

The SM-count query is now lazy. Only automatic SplitKV needs the target SM count; fixed and explicit split requests store num_sms=0 and skip the environment/device query. The fixed path measured 1.341→0.030 us, while automatic SplitKV still resolves the matching local GPU or explicit cross-compilation target exactly as before.

Together with the tuple-backed keys in #2732, constructing selector inputs plus a warm cache hit measured 5.734→1.378 us. This matters to eager and generated Flex runtime wrappers but is intentionally absent from the CUDA-graph kernel measurements in the policy PRs.

A complete warm host-dispatch check used real preallocated CUDA tensors, a hit-only compile cache, and a no-op compiled callable:

Shape	Before cleanup	After cleanup	Saved
Decode	36.856 us	21.482 us	15.374 us (41.7%)
Short attention	36.267 us	21.149 us	15.118 us (41.7%)
Prefill	36.271 us	21.289 us	14.982 us (41.3%)

“Before” is exact stack head 69285d1c; “after” is 60fe2c4c. The benchmark excludes GPU execution and compilation, so it isolates the wrapper work changed by #2732/#2733.

Implementation and validation

• Replaces the inline forward scheduling policy with FwdHeuristicInputs construction followed by selection or exact validation.
• Main and SplitKV-combine kernels use the independent typed identities defined in the preceding config PR.
• Generic SM100/SM110 and dedicated HD256 constructors consume FwdRegisterAllocation; the hidden kernel-side register tables are removed.
• Exact SplitKV counts can share one main specialization while compiling only the combine capacity buckets they require.
• Optional out_partial and lse_partial buffers are shape/device/dtype validated and accepted only for a resolved SplitKV config.
• has_lse is derived from preallocated LSE, return_lse, or autograd requirements, keeping output-only and LSE-producing policy separate.
• The standalone exp2 sweep remains exp2-only; register experiments use explicit configs.
• Public functions and autograd wrappers are unchanged.

The fuzz-found alignment, static-layout, aux-ABI, fake-SM, and combine-operand corrections live in the base PR rather than obscuring this routing change. At this boundary, interface.py is 19 lines smaller than its parent. The standalone boundary passes 47 host config/cache tests, and focused GB300 runtime tests pass for register forcing, SplitKV specialization reuse, CLC, MLA cache identity, and selector metadata. Emitted PTX confirms the resolved register immediates are consumed.