Optimize K2 recurrent-state dataflow for SM103/B300

[ICtuvk123]

Summary

This PR adds an opt-in SM103/B300-specific optimization for FlashKDA K2 recurrent execution.

It introduces two optimized K2 implementations:

• V1a keeps the recurrent state persistently in per-thread registers instead of the canonical shared-memory state buffer.
• V1aE preserves the V1a recurrence and stages the final state through shared memory before using the existing TMA store path.

A workload-aware dispatcher can select between the baseline, V1a, and V1aE implementations on supported SM103 workloads.

The optimization targets recurrent-state dataflow rather than replacing the existing MMA instruction path.

Motivation

The recurrent K2 kernel repeatedly consumes and updates a 128x128 state matrix.

Keeping this state in shared memory introduces persistent shared-memory pressure. V1a instead keeps the state in registers across recurrence chunks, while V1aE improves the one-time final-state egress path.

The goal is to reduce state-management overhead on B300/SM103 without changing numerical behavior.

Implementation

Runtime selection is controlled through:

FLASH_KDA_K2_IMPL=baseline
FLASH_KDA_K2_IMPL=v1a
FLASH_KDA_K2_IMPL=v1ae
FLASH_KDA_K2_IMPL=auto