[KDA] Fused intra-chunk prefill path (`chunk_kda_fwd_intra(fuse_diagonal=True)`) collapses for strong per-channel decays because of the ±126 clamp in the exp2 factorization

Summary

chunk_kda_fwd selects chunk_kda_fwd_intra(fuse_diagonal=True, fuse_recompute=True) whenever B * NT * H <= 256 (i.e. for every short prefill). Inside chunk_kda_fwd_kernel_inter_solve_fused the diagonal sub-blocks factorize exp2(g_i - g_j) = exp2(g_i - g_n) * exp2(g_n - g_j) and clamp both factors:

b_gm0 = tl.clamp(b_g0 - b_gn0[None, :], -126.0, 126.0)
b_gq0 = tl.where(m_tc0[:, None], exp2(b_gm0), 0.0)
b_gk0 = tl.where(m_tc0[:, None], exp2(-b_gm0), 0.0)

When a channel decays by more than 2^126 inside a 16-token sub-block, the clamped factors no longer cancel and the block's exp2(g_i - g_j) is effectively replaced by exp2(0): keys that should have been forgotten are read at full weight. Kimi Linear's gates are mild enough not to trigger this; models with sharper per-channel decays do (our model: cumulative log2 decays inside a 64-token chunk down to -485, exp(A_log) up to 72).

Measurements (H200, Triton kernels from main @ 2026-08-03, one real layer's q/k/v/gate/beta, T = 298)

rel-L2 of the chunk output vs an fp32 torch recurrence on identical inputs:

End to end (40-layer hybrid model, per-layer comparison against the HF implementation): the first KDA layer is 15 % off with the fused path and 0.23 % with the unfused one; prompt top-1 agreement 95.3 % vs 97.0 % (the HF fla-vs-torch reference pair itself agrees 97.3 %).

The effect is gate-magnitude driven: clamping the cumulative gate at -8 (natural log) before the kernel makes fused == unfused (2.2e-2 on the pre-#31904 snapshot), i.e. the fused path is exact whenever the ±126 clamp is inactive.

Proposed mitigation

Per-layer opt-out of the fused path, threaded from the attention layer to the kernel:

• chunk_kda(..., fused_intra: Optional[bool] = None) / chunk_kda_fwd(...): _fused_intra = _small_grid if fused_intra is None else bool(fused_intra); pass it as fuse_diagonal / fuse_recompute.
• TritonKDAKernel.extend(...) forwards fused_intra=kwargs.get("fused_intra").
• KDAAttnBackend.forward_extend passes fused_intra=getattr(layer, "kda_fused_intra", None).
• A model whose gates exceed the clamp sets self.attn.kda_fused_intra = False on its RadixLinearAttention.

Default behaviour (heuristic) is unchanged for all existing models. A kernel-side alternative would be to detect |g_i - g_n| > 126 for the block and fall back to the direct difference form, but that costs every model.

Related

• #31904 fixed the exponent-base mix (natural-log cumsum vs exp2 intra kernels) in the same file family; snapshots before it also show a 2 % floor on the unfused path for such gates (units mismatch, not a precision limit).

Second (small) issue: LinearAttnModelSpec.support_mamba_cache_extra_buffer is ignored

arg_groups/overrides.py::supports_mamba_cache_extra_buffer decides from the hardcoded _MAMBA_EXTRA_BUFFER_ARCHS only, so a model registered through linear_attn_model_registry with support_mamba_cache_extra_buffer=True still resolves --mamba-radix-cache-strategy auto to no_buffer (which also forces disable_overlap_schedule), and an explicit extra_buffer is rejected by _validate_mamba_extra_buffer ("extra_buffer is not supported for ; use no_buffer"). Proposed fix (what the BerryLM PR carries): consult the registry first —

spec = get_linear_attn_spec_by_arch(model_arch)
if spec is not None and spec.support_mamba_cache_extra_buffer:
    return view.linear_attn_backend == "triton"
if model_arch in _MAMBA_EXTRA_BUFFER_ARCHS:
    return view.linear_attn_backend == "triton"
return False

Verified on BerryLM (KDA backend, TP1): with extra_buffer the radix cache hits a repeated 6k prompt (6016 cached), a sibling prefix, an 8k prompt (8000 cached) and a multi-turn continuation (6080 cached) with generations identical to the uncached runs and unchanged HF teacher-forced log-probs.

3. attn_backend_wrapper discards the registry's matched config (main, 2026-09)

python/sglang/srt/layers/attention/attention_registry.py, registry branch of attn_backend_wrapper (commit 6b333382):

spec_result = get_linear_attn_config(runner.model_config.hf_config)
if spec_result is not None:
    spec, _ = spec_result
    cfg = runner.model_config
    ...
full_attn_layers = cfg.full_attention_layer_ids

get_linear_attn_config returns (spec, matched_hf_config); the matched config is thrown away and cfg is set to the ModelConfig, which only gets full_attention_layer_ids in _derive_hybrid_model for hybrid-SWA architectures. Every model that reaches the hybrid backend through register_linear_attn_model() therefore dies at startup with AttributeError: 'ModelConfig' object has no attribute 'full_attention_layer_ids' (our BerryLM server smoke on the 2026-09-16 main tree; the 2026-08-03 tree still took the ids from the matched config). Fix: spec, cfg = spec_result.