The new parametrization of test_mla_backends.py over MLAPrefillBackendEnum
exercised FLASHINFER MLA prefill for the first time, exposing two test-only
issues. Both are fixed entirely in the test:
1. MockMLAAttentionLayer now subclasses MLAAttention (skipping
MLAAttention.__init__ to avoid its self-registration side effects).
FlashInferPrefillBackend._resolve_global_hyperparameters filters
static_forward_context by isinstance(layer, MLAAttention); the mock
previously only inherited from AttentionLayerBase, so the filter found
zero layers and infer_global_hyperparameters asserted.
2. run_attention_backend now passes 1/sqrt(qk_head_dim) (the prefill scale)
to impl_cls, matching production where impl.scale and prefill backend
scale are the same. The decode SDPA reference uses prefill_scale too.
Without this, FlashInfer would plan with the wrong sm_scale (read back
from impl.scale via global_hyperparameters) and prefill numerics would
drift past the rtol=1e-2, atol=5e-1 tolerance.
Verified locally on B300: all 144 previously-failing FLASHINFER cases pass
(144 passed in 387s).
Co-authored-by: Claude
Signed-off-by: Roger Wang <hey@rogerw.io>
