diff --git a/atom/plugin/sglang/models/deepseek_mla.py b/atom/plugin/sglang/models/deepseek_mla.py
index a8d4deb1a1..a7165b7112 100644
--- a/atom/plugin/sglang/models/deepseek_mla.py
+++ b/atom/plugin/sglang/models/deepseek_mla.py
@@ -3,19 +3,19 @@
 
 """Model-level DeepSeek MLA patching for SGLang plugin mode.
 
-This module owns the monkey-patch entrypoints that adapt DeepSeek MLA models to
-SGLang plugin mode. The heavy DeepSeek-specific forward and weight helpers live
-in `atom.plugin.sglang.models.deepseek_mla_forward`.
+This module owns the install-time hooks that adapt DeepSeek MLA models to
+SGLang plugin mode. The heavy DeepSeek-specific runtime helpers live in
+`atom.plugin.sglang.models.deepseek_mla_forward`.
 """
 
 from __future__ import annotations
 
 from typing import TYPE_CHECKING, Any
 
-import torch
-
+from atom.plugin.sglang.models.deepseek_mla_attention import (
+    SGLangDeepseekMLAAttention,
+)
 from atom.plugin.sglang.models.deepseek_mla_forward import (
-    forward_sgl_plugin_mode,
     init_sgl_attrs,
     process_mla_kv_b_proj_after_loading,
 )
@@ -56,20 +56,8 @@ def _patch_mla_attention_for_sglang(
 ) -> None:
     """Patch one DeepSeek MLA layer for SGLang plugin mode."""
     init_sgl_attrs(attn, config, kv_cache_dtype)
-
-    def patched_forward(
-        positions: torch.Tensor,
-        hidden_states: torch.Tensor,
-        **kwargs: Any,
-    ) -> torch.Tensor:
-        from atom.plugin.sglang.models.base_model_wrapper import (
-            get_current_forward_batch,
-        )
-
-        kwargs["forward_batch"] = get_current_forward_batch()
-        return forward_sgl_plugin_mode(attn, positions, hidden_states, **kwargs)
-
-    attn.forward = patched_forward
+    if not isinstance(attn.mla_attn, SGLangDeepseekMLAAttention):
+        attn.mla_attn = SGLangDeepseekMLAAttention(attn, attn.mla_attn)
     attn.process_weights_after_loading = lambda: process_mla_kv_b_proj_after_loading(
         attn
     )
diff --git a/atom/plugin/sglang/models/deepseek_mla_attention.py b/atom/plugin/sglang/models/deepseek_mla_attention.py
new file mode 100644
index 0000000000..bda5ef87d0
--- /dev/null
+++ b/atom/plugin/sglang/models/deepseek_mla_attention.py
@@ -0,0 +1,332 @@
+# SPDX-License-Identifier: MIT
+# Copyright (C) 2024-2025, Advanced Micro Devices, Inc. All rights reserved.
+
+"""DeepSeek MLA wrapper for SGLang plugin mode.
+
+This adapter keeps the model-side entry at ``self.mla_attn(...)`` and owns the
+SGLang-specific runtime dispatch for DeepSeek MLA. It is intentionally shaped
+closer to the vLLM plugin path than the older model-side monkey-patched entry.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+import torch
+from torch import nn
+
+if TYPE_CHECKING:
+    from atom.models.deepseek_v2 import DeepseekV2MLAAttention
+
+
+class SGLangDeepseekMLAAttention(nn.Module):
+    """Enter SGLang DeepSeek MLA runtime through ``self.mla_attn(...)``."""
+
+    def __init__(
+        self,
+        owner_attn: "DeepseekV2MLAAttention",
+        base_attn: nn.Module,
+    ) -> None:
+        super().__init__()
+        self.owner_attn = owner_attn
+        self.base_attn = base_attn
+
+    @property
+    def attn(self):
+        return getattr(self.base_attn, "attn", self.base_attn)
+
+    def _get_forward_batch(self, kwargs: dict[str, Any]):
+        forward_batch = kwargs.get("forward_batch", None)
+        if forward_batch is None:
+            from atom.plugin.sglang.models.base_model_wrapper import (
+                get_current_forward_batch,
+            )
+
+            forward_batch = get_current_forward_batch()
+            kwargs["forward_batch"] = forward_batch
+        if forward_batch is None:
+            raise RuntimeError(
+                "forward_batch is required for SGLang DeepSeek MLA wrapper"
+            )
+        return forward_batch
+
+    def _infer_total_tokens(self, forward_batch, tensor: torch.Tensor) -> int:
+        if hasattr(forward_batch, "input_ids") and forward_batch.input_ids is not None:
+            return int(forward_batch.input_ids.shape[0])
+        if hasattr(forward_batch, "positions") and forward_batch.positions is not None:
+            return int(forward_batch.positions.shape[0])
+        if hasattr(forward_batch, "seq_lens_sum"):
+            return int(forward_batch.seq_lens_sum)
+        return int(tensor.shape[0])
+
+    def _maybe_all_gather(
+        self,
+        tensor: torch.Tensor | None,
+        *,
+        total_tokens: int,
+        input_scattered: bool,
+    ):
+        if tensor is None or not input_scattered:
+            return tensor
+        from sglang.srt.distributed import get_tp_group
+
+        output = tensor.new_empty((total_tokens, *tensor.shape[1:]))
+        get_tp_group().all_gather_into_tensor(output, tensor)
+        return output
+
+    def _gather_runtime_inputs(
+        self,
+        q_input: torch.Tensor,
+        kv_c_normed: torch.Tensor,
+        k_pe: torch.Tensor,
+        positions: torch.Tensor,
+        q_scale: torch.Tensor | None,
+        *,
+        forward_batch,
+        input_scattered: bool,
+    ):
+        total_tokens = self._infer_total_tokens(forward_batch, q_input)
+        q_input = self._maybe_all_gather(
+            q_input,
+            total_tokens=total_tokens,
+            input_scattered=input_scattered,
+        )
+        kv_c_normed = self._maybe_all_gather(
+            kv_c_normed,
+            total_tokens=total_tokens,
+            input_scattered=input_scattered,
+        )
+        k_pe = self._maybe_all_gather(
+            k_pe,
+            total_tokens=total_tokens,
+            input_scattered=input_scattered,
+        )
+        positions = self._maybe_all_gather(
+            positions,
+            total_tokens=total_tokens,
+            input_scattered=input_scattered,
+        )
+        q_scale = self._maybe_all_gather(
+            q_scale,
+            total_tokens=total_tokens,
+            input_scattered=input_scattered,
+        )
+        return q_input, kv_c_normed, k_pe, positions, q_scale
+
+    def _project_q(
+        self,
+        q_input: torch.Tensor,
+        q_scale: torch.Tensor | None,
+    ) -> torch.Tensor:
+        attn = self.owner_attn
+        from atom.plugin.sglang.models.deepseek_mla_forward import _unwrap_linear_output
+
+        if attn.q_lora_rank is not None:
+            q = (
+                attn.q_b_proj(q_input, q_scale)
+                if q_scale is not None
+                else attn.q_b_proj(q_input)
+            )
+        else:
+            q = (
+                attn.q_proj(q_input, q_scale)
+                if q_scale is not None
+                else attn.q_proj(q_input)
+            )
+        return _unwrap_linear_output(q).view(
+            -1, attn.num_local_heads, attn.qk_head_dim
+        )
+
+    def _forward_absorbed(
+        self,
+        q_input: torch.Tensor,
+        kv_c_normed: torch.Tensor,
+        k_pe: torch.Tensor,
+        positions: torch.Tensor,
+        q_scale: torch.Tensor | None,
+        *,
+        forward_batch,
+    ) -> torch.Tensor:
+        attn = self.owner_attn
+        from aiter import dtypes
+        from atom.model_ops.attention_mla import fused_qk_rope_concat_and_cache_mla
+        from atom.plugin.sglang.models.deepseek_mla_forward import (
+            _get_sglang_radix_attn,
+            mla_absorbed_bmm,
+            mla_v_up_proj,
+        )
+        from sglang.srt.layers.attention.nsa.utils import nsa_use_prefill_cp
+
+        q = self._project_q(q_input, q_scale)
+        k_nope = kv_c_normed.unsqueeze(1)
+        k_pe = k_pe.unsqueeze(1)
+        q_nope, q_pe = q.split(
+            [attn.qk_nope_head_dim, attn.qk_rope_head_dim], dim=-1
+        )
+        q_nope_out = mla_absorbed_bmm(
+            attn, q_nope, attn.w_kc, attn.w_scale, attn.w_scale_k, attn.kv_lora_rank
+        )
+
+        if attn.rotary_emb is not None and not attn.use_fused_qk_rope_concat_and_cache_mla:
+            q_pe, k_pe = attn.rotary_emb(positions, q_pe, k_pe)
+
+        if nsa_use_prefill_cp(forward_batch):
+            latent_cache = torch.cat([k_nope.squeeze(1), k_pe.squeeze(1)], dim=-1)
+            k_nope, k_pe = attn.rebuild_cp_kv_cache(
+                latent_cache, forward_batch, k_nope, k_pe
+            )
+
+        save_kv_cache = True
+        if attn.use_fused_qk_rope_concat_and_cache_mla:
+            mla_attn = _get_sglang_radix_attn(self.base_attn)
+            kv_cache = forward_batch.token_to_kv_pool.get_key_buffer(mla_attn.layer_id)
+            q_out_dtype = (
+                dtypes.fp8
+                if attn.kv_cache_dtype == "fp8_e4m3"
+                else q_nope_out.dtype
+            )
+            q = torch.empty(
+                (
+                    q_nope_out.shape[0],
+                    attn.num_local_heads,
+                    attn.kv_lora_rank + attn.qk_rope_head_dim,
+                ),
+                dtype=q_out_dtype,
+                device=q_nope_out.device,
+            )
+            fused_qk_rope_concat_and_cache_mla(
+                q_nope_out,
+                q_pe,
+                k_nope,
+                k_pe,
+                kv_cache,
+                q,
+                forward_batch.out_cache_loc,
+                mla_attn.k_scale,
+                mla_attn.k_scale,
+                positions,
+                attn.rotary_emb.cos_cache,
+                attn.rotary_emb.sin_cache,
+                is_neox=attn.rotary_emb.is_neox_style,
+                is_nope_first=True,
+            )
+            k = None
+            v = None
+            save_kv_cache = False
+        else:
+            q = torch.cat([q_nope_out, q_pe], dim=-1)
+            k = torch.cat([k_nope, k_pe], dim=-1)
+            v = k_nope
+
+        attn_output = self.base_attn(
+            q,
+            k,
+            v,
+            forward_batch=forward_batch,
+            save_kv_cache=save_kv_cache,
+        )
+        attn_output = attn_output.view(-1, attn.num_local_heads, attn.kv_lora_rank)
+        attn_bmm_output = mla_v_up_proj(
+            attn, attn_output, attn.w_vc, attn.w_scale, attn.w_scale_v, attn.v_head_dim
+        )
+        return attn.o_proj(attn_bmm_output)
+
+    def _forward_non_absorbed(
+        self,
+        q_input: torch.Tensor,
+        kv_c_normed: torch.Tensor,
+        k_pe: torch.Tensor,
+        positions: torch.Tensor,
+        q_scale: torch.Tensor | None,
+        *,
+        forward_batch,
+    ) -> torch.Tensor:
+        attn = self.owner_attn
+        from atom.plugin.sglang.models.deepseek_mla_forward import (
+            _set_mla_kv_buffer_for_non_absorbed,
+            _unwrap_linear_output,
+        )
+
+        q = self._project_q(q_input, q_scale)
+        _, q_pe = q.split([attn.qk_nope_head_dim, attn.qk_rope_head_dim], dim=-1)
+
+        kv_a = kv_c_normed
+        k_pe = k_pe.unsqueeze(1)
+        if attn.rotary_emb is not None:
+            q_pe, k_pe = attn.rotary_emb(positions, q_pe, k_pe)
+        q[..., attn.qk_nope_head_dim :] = q_pe
+
+        _set_mla_kv_buffer_for_non_absorbed(attn, kv_a, k_pe, forward_batch)
+
+        kv = _unwrap_linear_output(attn.kv_b_proj(kv_a)).view(
+            -1, attn.num_local_heads, attn.qk_nope_head_dim + attn.v_head_dim
+        )
+        k_nope = kv[..., : attn.qk_nope_head_dim]
+        v = kv[..., attn.qk_nope_head_dim :]
+        k = torch.cat(
+            [k_nope, k_pe.expand(-1, attn.num_local_heads, -1)],
+            dim=-1,
+        )
+
+        attn_output = attn.attn_non_absorbed(
+            q,
+            k,
+            v,
+            forward_batch=forward_batch,
+            save_kv_cache=False,
+        )
+        attn_output = attn_output.reshape(-1, attn.num_local_heads * attn.v_head_dim)
+        return attn.o_proj(attn_output)
+
+    def forward(
+        self,
+        q_input: torch.Tensor,
+        kv_c_normed: torch.Tensor,
+        k_pe: torch.Tensor,
+        positions: torch.Tensor,
+        q_scale: torch.Tensor | None = None,
+        **kwargs: Any,
+    ) -> torch.Tensor:
+        attn = self.owner_attn
+        forward_batch = self._get_forward_batch(kwargs)
+
+        from atom.plugin.sglang.models.deepseek_mla_forward import (
+            _can_run_non_absorbed_mla_now,
+        )
+        from sglang.srt.layers.communicator import get_attn_tp_context
+
+        attn_tp_context = get_attn_tp_context()
+        with attn_tp_context.maybe_input_scattered(forward_batch):
+            q_input, kv_c_normed, k_pe, positions, q_scale = self._gather_runtime_inputs(
+                q_input,
+                kv_c_normed,
+                k_pe,
+                positions,
+                q_scale,
+                forward_batch=forward_batch,
+                input_scattered=attn_tp_context.input_scattered,
+            )
+
+            if forward_batch.forward_mode.is_extend_without_speculative():
+                if _can_run_non_absorbed_mla_now(attn, forward_batch):
+                    attn.current_sgl_plugin_attn_path = "non_absorbed"
+                    return self._forward_non_absorbed(
+                        q_input,
+                        kv_c_normed,
+                        k_pe,
+                        positions,
+                        q_scale,
+                        forward_batch=forward_batch,
+                    )
+                attn.current_sgl_plugin_attn_path = "absorbed_fallback"
+            else:
+                attn.current_sgl_plugin_attn_path = "absorbed"
+
+            return self._forward_absorbed(
+                q_input,
+                kv_c_normed,
+                k_pe,
+                positions,
+                q_scale,
+                forward_batch=forward_batch,
+            )
diff --git a/atom/plugin/sglang/models/deepseek_mla_forward.py b/atom/plugin/sglang/models/deepseek_mla_forward.py
index 25f1ef79aa..f88aa762f8 100644
--- a/atom/plugin/sglang/models/deepseek_mla_forward.py
+++ b/atom/plugin/sglang/models/deepseek_mla_forward.py
@@ -1,38 +1,26 @@
 # SPDX-License-Identifier: MIT
 # Copyright (C) 2024-2025, Advanced Micro Devices, Inc. All rights reserved.
 
-"""Model-specific DeepSeek MLA helpers for SGLang plugin mode.
+"""Helper functions for DeepSeek MLA in SGLang plugin mode.
 
-DeepSeek MLA (Multi-Latent Attention) forward logic for SGLang plugin mode:
-absorbed BMM computation, MHA/MLA path dispatch (prefill -> MHA, decode -> MLA),
-and kv_b_proj weight splitting (w_kc/w_vc).
-
-This module lives under ``atom.plugin.sglang.models`` because the logic is
-DeepSeek-model-specific rather than a generic SGLang attention backend.
-
-TODO: rewrite this file once sglang's attention flow is unified into ATOM's
-attention layer - the MLA absorbed path and MHA dispatch will then be handled
-natively by ATOM's attention ops, making this sglang-specific module
-unnecessary.
+This module now contains only the low-level helpers that are still shared by
+the SGLang DeepSeek MLA wrapper and the install-time weight hooks:
+absorbed BMM math, small utility helpers, non-absorbed cache staging, and
+kv_b_proj post-load processing.
 """
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, NamedTuple, Optional
+from typing import TYPE_CHECKING, Any, Optional
 
 import torch
 from aiter import dtypes
-from aiter.dist.parallel_state import get_tensor_model_parallel_world_size, get_tp_group
 from atom.model_ops.base_attention import Attention
 from atom.model_ops.attention_mla import (
     dynamic_per_batched_tensor_quant,
-    fused_qk_rope_concat_and_cache_mla,
 )
 from atom.models.utils import maybe_prefix
-from atom.models.deepseek_v2 import _fuse_rmsnorm_quant
 
-from sglang.srt.layers.communicator import AttentionInputs, get_attn_tp_context
-from sglang.srt.layers.attention.nsa.utils import nsa_use_prefill_cp
 from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode
 from sglang.srt.models.deepseek_common.utils import (
     _use_aiter_gfx95,
@@ -90,25 +78,6 @@ def bmm_fp8(A, B, A_scale, B_scale, dtype, out=None):
         raise RuntimeError("bmm_fp8 requires CUDA (sgl_kernel)")
 
 
-class SglPrepareResult(NamedTuple):
-    q_pe: torch.Tensor
-    k_pe: torch.Tensor
-    q_nope_out: torch.Tensor
-    k_nope: torch.Tensor
-    forward_batch: Any
-    zero_allocator: Any
-    positions: torch.Tensor
-    topk_indices: Optional[torch.Tensor]
-    llama_4_scaling: Optional[Any]
-
-
-class SglMhaPrepareResult(NamedTuple):
-    q: torch.Tensor
-    k: torch.Tensor
-    v: torch.Tensor
-    forward_batch: Any
-
-
 def _unwrap_linear_output(output: Any) -> torch.Tensor:
     """Normalize ATOM/public-SGLang linear outputs to a tensor."""
     if isinstance(output, tuple):
@@ -116,57 +85,6 @@ def _unwrap_linear_output(output: Any) -> torch.Tensor:
     return output
 
 
-def _linear_quant_type_value(linear: Any) -> Optional[int]:
-    quant_type = getattr(linear, "quant_type", None)
-    return None if quant_type is None else getattr(quant_type, "value", quant_type)
-
-
-def _fuse_qk_rmsnorm_and_q_quant(
-    attn: DeepseekV2MLAAttention,
-    q: torch.Tensor,
-    k_nope: torch.Tensor,
-    *,
-    output_unquantized_q: bool = False,
-) -> tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor], torch.Tensor]:
-    """Fuse q/k RMSNorm and q quant using ATOM's DeepSeek-V2 path."""
-
-    (q_quantized, q_scale), q_normed, k_nope_normed, _ = _fuse_rmsnorm_quant(
-        q,
-        attn.q_a_layernorm.weight,
-        attn.q_a_layernorm.eps,
-        k_nope,
-        attn.kv_a_layernorm.weight,
-        attn.kv_a_layernorm.eps,
-        None,
-        dtype_quant=attn.quant_dtype,
-        shuffle=False,
-        scale_shuffle_padding=False,
-        group_size=128,
-        quant_type=_linear_quant_type_value(attn.q_b_proj),
-        output_unquantized_inp1=output_unquantized_q,
-        transpose_scale=True,
-    )
-    return q_quantized, q_scale, q_normed, k_nope_normed
-
-
-def _fuse_qk_rmsnorm(
-    attn: DeepseekV2MLAAttention,
-    q: torch.Tensor,
-    k_nope: torch.Tensor,
-) -> tuple[torch.Tensor, torch.Tensor]:
-    """Fuse q/k RMSNorm without quantizing q."""
-    from atom.models.deepseek_v2 import _fused_qk_rmsnorm
-
-    return _fused_qk_rmsnorm(
-        q,
-        attn.q_a_layernorm.weight,
-        attn.q_a_layernorm.eps,
-        k_nope,
-        attn.kv_a_layernorm.weight,
-        attn.kv_a_layernorm.eps,
-    )
-
-
 def _prepare_weight_for_bmm(
     weight: torch.Tensor, in_dim: int, out_dim: int
 ) -> torch.Tensor:
@@ -199,7 +117,7 @@ def init_sgl_attrs(
     attn.w_scale = None
     attn.w_scale_k = None
     attn.w_scale_v = None
-    attn.attn_mha = Attention(
+    attn.attn_non_absorbed = Attention(
         num_heads=attn.num_local_heads,
         head_dim=attn.qk_head_dim,
         scale=attn.scaling,
@@ -208,10 +126,10 @@ def init_sgl_attrs(
         layer_num=attn.layer_num,
         use_mla=False,
         v_head_dim=attn.v_head_dim,
-        prefix=maybe_prefix(attn.prefix, "attn_mha"),
+        prefix=maybe_prefix(attn.prefix, "attn_non_absorbed"),
     )
-    if hasattr(attn.attn_mha, "attn"):
-        attn.attn_mha.attn.kv_b_proj = None
+    if hasattr(attn.attn_non_absorbed, "attn"):
+        attn.attn_non_absorbed.attn.kv_b_proj = None
 
 
 def mla_absorbed_bmm(
@@ -340,275 +258,31 @@ def mla_v_up_proj(
     return mla_absorbed_bmm(
         attn, inp, weight, weight_scale, weight_scale_k, out_dim
     ).flatten(1, 2)
-
-
-def forward_sgl_prepare(
-    attn: DeepseekV2MLAAttention,
-    positions: torch.Tensor,
-    hidden_states: torch.Tensor,
-    **model_kwargs,
-) -> SglPrepareResult:
-    """Prepare QKV for sglang MLA attention."""
-    hidden_states_scale = None
-    if isinstance(hidden_states, tuple):
-        hidden_states, hidden_states_scale = hidden_states
-
-    forward_batch = model_kwargs.get("forward_batch", None)
-    zero_allocator = model_kwargs.get("zero_allocator", None)
-    llama_4_scaling = model_kwargs.get("llama_4_scaling", None)
-    q_lora = None
-    topk_indices = None
-
-    if attn.q_lora_rank is not None:
-        q, latent_cache = (
-            get_attn_tp_context()
-            .fetch_qkv_latent()
-            .split(
-                [attn.q_lora_rank, attn.kv_lora_rank + attn.qk_rope_head_dim],
-                dim=-1,
-            )
-        )
-
-        if (
-            q.shape[0] != positions.shape[0]
-            and get_tensor_model_parallel_world_size() > 1
-        ):
-            qkv_lora = torch.cat([q, latent_cache], dim=-1)
-            qkv_lora = get_tp_group().all_gather(qkv_lora, dim=0)
-            if qkv_lora.shape[0] < positions.shape[0]:
-                raise RuntimeError(
-                    f"qkv_lora gather mismatch: got {qkv_lora.shape[0]}, "
-                    f"expected {positions.shape[0]}"
-                )
-            qkv_lora = qkv_lora[: positions.shape[0]]
-            q, latent_cache = torch.split(
-                qkv_lora,
-                [attn.q_lora_rank, attn.kv_lora_rank + attn.qk_rope_head_dim],
-                dim=-1,
-            )
-
-        k_nope = latent_cache[..., : attn.kv_lora_rank]
-        q_scale = None
-
-        if getattr(attn, "fuse_qknorm_quant", False):
-            q, q_scale, q_lora, k_nope = _fuse_qk_rmsnorm_and_q_quant(
-                attn,
-                q,
-                k_nope,
-                output_unquantized_q=attn.use_nsa,
-            )
-        elif getattr(attn, "fuse_qknorm", False):
-            q, k_nope = _fuse_qk_rmsnorm(attn, q, k_nope)
-        elif attn.alt_stream is not None and get_is_capture_mode():
-            current_stream = torch.cuda.current_stream()
-            attn.alt_stream.wait_stream(current_stream)
-            q = attn.q_a_layernorm(q)
-            with torch.cuda.stream(attn.alt_stream):
-                k_nope = attn.kv_a_layernorm(k_nope)
-            current_stream.wait_stream(attn.alt_stream)
-        else:
-            q = attn.q_a_layernorm(q)
-            k_nope = attn.kv_a_layernorm(k_nope)
-
-        if attn.use_nsa and q_lora is None:
-            q_lora = q
-
-        if (
-            attn.alt_stream is not None
-            and get_is_capture_mode()
-            and forward_batch.forward_mode.is_decode_or_idle()
-            and q_lora is not None
-        ):
-            current_stream = torch.cuda.current_stream()
-            attn.alt_stream.wait_stream(current_stream)
-            with torch.cuda.stream(attn.alt_stream):
-                k_nope = k_nope.unsqueeze(1)
-                q = _unwrap_linear_output(
-                    attn.q_b_proj(q, q_scale)
-                    if q_scale is not None
-                    else attn.q_b_proj(q)
-                ).view(-1, attn.num_local_heads, attn.qk_head_dim)
-            topk_indices = attn.indexer(
-                x=hidden_states,
-                q_lora=q_lora,
-                positions=positions,
-                forward_batch=forward_batch,
-                layer_id=attn.layer_num,
-            )
-            current_stream.wait_stream(attn.alt_stream)
-        else:
-            k_nope = k_nope.unsqueeze(1)
-            q = _unwrap_linear_output(
-                attn.q_b_proj(q, q_scale) if q_scale is not None else attn.q_b_proj(q)
-            ).view(-1, attn.num_local_heads, attn.qk_head_dim)
-            if q_lora is not None:
-                topk_indices = attn.indexer(
-                    x=hidden_states,
-                    q_lora=q_lora,
-                    positions=positions,
-                    forward_batch=forward_batch,
-                    layer_id=attn.layer_num,
-                )
-    else:
-        q = _unwrap_linear_output(attn.q_proj(hidden_states)).view(
-            -1, attn.num_local_heads, attn.qk_head_dim
-        )
-        latent_cache = _unwrap_linear_output(attn.kv_a_proj_with_mqa(hidden_states))
-        k_nope = latent_cache[..., : attn.kv_lora_rank]
-        k_nope = attn.kv_a_layernorm(k_nope).unsqueeze(1)
-
-    q_nope, q_pe = q.split([attn.qk_nope_head_dim, attn.qk_rope_head_dim], dim=-1)
-    k_pe = latent_cache[..., attn.kv_lora_rank :].unsqueeze(1)
-
-    q_nope_out = mla_absorbed_bmm(
-        attn, q_nope, attn.w_kc, attn.w_scale, attn.w_scale_k, attn.kv_lora_rank
-    )
-
-    if attn.rotary_emb is not None and not attn.use_fused_qk_rope_concat_and_cache_mla:
-        q_pe, k_pe = attn.rotary_emb(positions, q_pe, k_pe)
-
-    if nsa_use_prefill_cp(forward_batch):
-        k_nope, k_pe = attn.rebuild_cp_kv_cache(
-            latent_cache, forward_batch, k_nope, k_pe
-        )
-
-    return SglPrepareResult(
-        q_pe=q_pe,
-        k_pe=k_pe,
-        q_nope_out=q_nope_out,
-        k_nope=k_nope,
-        forward_batch=forward_batch,
-        zero_allocator=zero_allocator,
-        positions=positions,
-        topk_indices=topk_indices,
-        llama_4_scaling=llama_4_scaling,
-    )
-
-
-def forward_sgl_core(
-    attn: DeepseekV2MLAAttention,
-    prepared: SglPrepareResult,
-) -> torch.Tensor:
-    """Core MLA attention computation for sglang."""
-    save_kv_cache = True
-
-    if attn.use_fused_qk_rope_concat_and_cache_mla:
-        mla_attn = _get_sglang_radix_attn(attn.mla_attn)
-        kv_cache = prepared.forward_batch.token_to_kv_pool.get_key_buffer(
-            mla_attn.layer_id
-        )
-        q_out_dtype = (
-            dtypes.fp8
-            if attn.kv_cache_dtype == "fp8_e4m3"
-            else prepared.q_nope_out.dtype
-        )
-        q = torch.empty(
-            (
-                prepared.q_nope_out.shape[0],
-                attn.num_local_heads,
-                attn.kv_lora_rank + attn.qk_rope_head_dim,
-            ),
-            dtype=q_out_dtype,
-            device=prepared.q_nope_out.device,
-        )
-
-        fused_qk_rope_concat_and_cache_mla(
-            prepared.q_nope_out,
-            prepared.q_pe,
-            prepared.k_nope,
-            prepared.k_pe,
-            kv_cache,
-            q,
-            prepared.forward_batch.out_cache_loc,
-            mla_attn.k_scale,
-            mla_attn.k_scale,
-            prepared.positions,
-            attn.rotary_emb.cos_cache,
-            attn.rotary_emb.sin_cache,
-            is_neox=attn.rotary_emb.is_neox_style,
-            is_nope_first=True,
-        )
-        k = None
-        v = None
-        save_kv_cache = False
-    else:
-        q = torch.cat([prepared.q_nope_out, prepared.q_pe], dim=-1)
-        k = torch.cat([prepared.k_nope, prepared.k_pe], dim=-1)
-        v = prepared.k_nope
-
-    if prepared.llama_4_scaling is not None:
-        q = q * prepared.llama_4_scaling
-
-    extra_kwargs = {}
-    if prepared.topk_indices is not None:
-        extra_kwargs["topk_indices"] = prepared.topk_indices
-
-    attn_output = attn.mla_attn(
-        q,
-        k,
-        v,
-        forward_batch=prepared.forward_batch,
-        save_kv_cache=save_kv_cache,
-        **extra_kwargs,
-    )
-    attn_output = attn_output.view(-1, attn.num_local_heads, attn.kv_lora_rank)
-
-    attn_bmm_output = mla_v_up_proj(
-        attn, attn_output, attn.w_vc, attn.w_scale, attn.w_scale_v, attn.v_head_dim
-    )
-
-    return attn.o_proj(attn_bmm_output)
-
-
-def _dispatch_sgl_plugin_attn_path(forward_batch) -> str:
-    """Decide the attention algorithm for this batch based on forward_mode."""
-    if forward_batch.forward_mode.is_extend_without_speculative():
-        return "mha"
-    return "mla"
-
-
-def forward_sgl_plugin_mode_mla(
-    attn: DeepseekV2MLAAttention,
-    positions: torch.Tensor,
-    hidden_states: torch.Tensor,
-    **model_kwargs,
-) -> torch.Tensor:
-    prepared = forward_sgl_prepare(attn, positions, hidden_states, **model_kwargs)
-    from atom.utils.forward_context import get_forward_context
-
-    if get_forward_context().context.is_dummy_run:
-        base_hidden_states = (
-            hidden_states[0] if isinstance(hidden_states, tuple) else hidden_states
-        )
-        dummy_output = base_hidden_states.new_empty(
-            (base_hidden_states.shape[0], base_hidden_states.shape[-1])
-        )
-        return dummy_output
-    return forward_sgl_core(attn, prepared)
-
-
 def _get_sglang_radix_attn(attn_module):
     return attn_module.attn if hasattr(attn_module, "attn") else attn_module
 
 
-def _set_mla_kv_buffer_for_mha(
+def _set_mla_kv_buffer_for_non_absorbed(
     attn: DeepseekV2MLAAttention,
     kv_a: torch.Tensor,
     k_pe: torch.Tensor,
     forward_batch,
 ) -> None:
-    attn_mha = _get_sglang_radix_attn(attn.attn_mha)
+    attn_non_absorbed = _get_sglang_radix_attn(attn.attn_non_absorbed)
     cache_k = torch.cat([kv_a.unsqueeze(1), k_pe], dim=-1)
     forward_batch.token_to_kv_pool.set_kv_buffer(
-        attn_mha,
+        attn_non_absorbed,
         forward_batch.out_cache_loc,
         cache_k,
         cache_k,
     )
 
 
-def _can_run_sgl_mha_now(attn: DeepseekV2MLAAttention, forward_batch) -> bool:
-    """Check if the model configuration supports the MHA attention path."""
+def _can_run_non_absorbed_mla_now(
+    attn: DeepseekV2MLAAttention,
+    forward_batch,
+) -> bool:
+    """Check if the non-absorbed MLA path is allowed for this batch."""
     del forward_batch
     if attn.use_nsa:
         return False
@@ -617,238 +291,6 @@ def _can_run_sgl_mha_now(attn: DeepseekV2MLAAttention, forward_batch) -> bool:
     return True
 
 
-def forward_sgl_mha_prepare(
-    attn: DeepseekV2MLAAttention,
-    positions: torch.Tensor,
-    hidden_states: torch.Tensor,
-    **model_kwargs,
-) -> SglMhaPrepareResult:
-
-    forward_batch = model_kwargs.get("forward_batch", None)
-    if forward_batch is None:
-        raise RuntimeError("forward_batch is required in forward_sgl_mha_prepare")
-
-    hidden_states_scale = None
-    if isinstance(hidden_states, tuple):
-        hidden_states, hidden_states_scale = hidden_states
-
-    attn_mha = _get_sglang_radix_attn(attn.attn_mha)
-    if getattr(attn_mha, "kv_b_proj", None) is None:
-        attn_mha.kv_b_proj = attn.kv_b_proj
-
-    if attn.q_lora_rank is not None:
-        q, latent_cache = (
-            get_attn_tp_context()
-            .fetch_qkv_latent()
-            .split(
-                [attn.q_lora_rank, attn.kv_lora_rank + attn.qk_rope_head_dim],
-                dim=-1,
-            )
-        )
-
-        if (
-            q.shape[0] != positions.shape[0]
-            and get_tensor_model_parallel_world_size() > 1
-        ):
-            qkv_lora = torch.cat([q, latent_cache], dim=-1)
-            qkv_lora = get_tp_group().all_gather(qkv_lora, dim=0)
-            if qkv_lora.shape[0] < positions.shape[0]:
-                raise RuntimeError(
-                    f"qkv_lora gather mismatch: got {qkv_lora.shape[0]}, "
-                    f"expected {positions.shape[0]}"
-                )
-            qkv_lora = qkv_lora[: positions.shape[0]]
-            q, latent_cache = torch.split(
-                qkv_lora,
-                [attn.q_lora_rank, attn.kv_lora_rank + attn.qk_rope_head_dim],
-                dim=-1,
-            )
-
-        if _use_aiter_gfx95 and attn.q_b_proj.weight.dtype == torch.float8_e4m3fn:
-            (q, q_scale), _, _, _ = _fuse_rmsnorm_quant(
-                q,
-                attn.q_a_layernorm.weight,
-                attn.q_a_layernorm.eps,
-                None,
-                None,
-                None,
-                res1=None,
-                dtype_quant=torch.float8_e4m3fn,
-                group_size=128,
-                quant_type=_linear_quant_type_value(attn.q_b_proj),
-                output_unquantized_inp1=False,
-                transpose_scale=True,
-            )
-            q = _unwrap_linear_output(attn.q_b_proj(q, q_scale)).view(
-                -1, attn.num_local_heads, attn.qk_head_dim
-            )
-        else:
-            q = attn.q_a_layernorm(q)
-            q = _unwrap_linear_output(attn.q_b_proj(q)).view(
-                -1, attn.num_local_heads, attn.qk_head_dim
-            )
-    else:
-        q = _unwrap_linear_output(attn.q_proj(hidden_states, hidden_states_scale)).view(
-            -1, attn.num_local_heads, attn.qk_head_dim
-        )
-        latent_cache = _unwrap_linear_output(
-            attn.kv_a_proj_with_mqa(hidden_states, hidden_states_scale)
-        )
-
-    _, q_pe = q.split([attn.qk_nope_head_dim, attn.qk_rope_head_dim], dim=-1)
-    kv_a, _ = latent_cache.split([attn.kv_lora_rank, attn.qk_rope_head_dim], dim=-1)
-    latent_cache = latent_cache.unsqueeze(1)
-
-    if _use_aiter_gfx95 and attn.kv_b_proj.weight.dtype == torch.float8_e4m3fn:
-        (kv_a_quanted, kv_a_quanted_scale), kv_a, _, _ = _fuse_rmsnorm_quant(
-            kv_a,
-            attn.kv_a_layernorm.weight,
-            attn.kv_a_layernorm.eps,
-            None,
-            None,
-            None,
-            res1=None,
-            dtype_quant=torch.float8_e4m3fn,
-            group_size=128,
-            quant_type=_linear_quant_type_value(attn.kv_b_proj),
-            output_unquantized_inp1=True,
-            transpose_scale=True,
-        )
-    else:
-        kv_a_quanted = None
-        kv_a = attn.kv_a_layernorm(kv_a)
-
-    k_pe = latent_cache[:, :, attn.kv_lora_rank :]
-    if attn.rotary_emb is not None:
-        q_pe, k_pe = attn.rotary_emb(positions, q_pe, k_pe)
-    q[..., attn.qk_nope_head_dim :] = q_pe
-
-    _set_mla_kv_buffer_for_mha(attn, kv_a, k_pe, forward_batch)
-
-    if kv_a_quanted is not None:
-        kv = _unwrap_linear_output(attn.kv_b_proj(kv_a_quanted, kv_a_quanted_scale))
-    else:
-        kv = _unwrap_linear_output(attn.kv_b_proj(kv_a))
-    kv = kv.view(-1, attn.num_local_heads, attn.qk_nope_head_dim + attn.v_head_dim)
-    k_nope = kv[..., : attn.qk_nope_head_dim]
-    v = kv[..., attn.qk_nope_head_dim :]
-    k = torch.cat(
-        [k_nope, k_pe.expand(-1, attn.num_local_heads, -1)],
-        dim=-1,
-    )
-    return SglMhaPrepareResult(q=q, k=k, v=v, forward_batch=forward_batch)
-
-
-def forward_sgl_mha_core(
-    attn: DeepseekV2MLAAttention,
-    prepared: SglMhaPrepareResult,
-) -> torch.Tensor:
-    attn_output = attn.attn_mha(
-        prepared.q,
-        prepared.k,
-        prepared.v,
-        forward_batch=prepared.forward_batch,
-        save_kv_cache=False,
-    )
-    attn_output = attn_output.reshape(-1, attn.num_local_heads * attn.v_head_dim)
-    return attn.o_proj(attn_output)
-
-
-def forward_sgl_plugin_mode_mha(
-    attn: DeepseekV2MLAAttention,
-    positions: torch.Tensor,
-    hidden_states: torch.Tensor,
-    **model_kwargs,
-) -> torch.Tensor:
-    forward_batch = model_kwargs.get("forward_batch", None)
-    if forward_batch is None:
-        raise RuntimeError("forward_batch is required in forward_sgl_plugin_mode_mha")
-    if not _can_run_sgl_mha_now(attn, forward_batch):
-        attn.current_sgl_plugin_attn_path = "mla_fallback"
-        return forward_sgl_plugin_mode_mla(
-            attn,
-            positions,
-            hidden_states,
-            **model_kwargs,
-        )
-    prepared = forward_sgl_mha_prepare(attn, positions, hidden_states, **model_kwargs)
-    return forward_sgl_mha_core(attn, prepared)
-
-
-def prepare_qkv_latent(
-    attn: DeepseekV2MLAAttention,
-    hidden_states: torch.Tensor,
-    forward_batch,
-) -> torch.Tensor:
-    """Prepare QKV latent tensor for the sglang communicator."""
-    assert attn.q_lora_rank is not None
-    hidden_states_scale = None
-    if isinstance(hidden_states, tuple):
-        hidden_states, hidden_states_scale = hidden_states
-    qkv_lora = attn.fused_qkv_a_proj(hidden_states, hidden_states_scale)
-
-    expected_tokens = 0
-    if hasattr(forward_batch, "positions") and forward_batch.positions is not None:
-        expected_tokens = int(forward_batch.positions.shape[0])
-    if expected_tokens <= 0:
-        expected_tokens = int(getattr(forward_batch, "seq_lens_sum", 0) or 0)
-
-    if (
-        expected_tokens > 0
-        and qkv_lora.shape[0] != expected_tokens
-        and get_tensor_model_parallel_world_size() > 1
-    ):
-        qkv_lora = get_tp_group().all_gather(qkv_lora, dim=0)
-        if qkv_lora.shape[0] > expected_tokens:
-            qkv_lora = qkv_lora[:expected_tokens]
-        elif qkv_lora.shape[0] < expected_tokens:
-            raise RuntimeError(
-                f"prepare_qkv_latent gather mismatch: got {qkv_lora.shape[0]}, "
-                f"expected {expected_tokens}"
-            )
-    return qkv_lora
-
-
-def forward_sgl_plugin_mode(
-    attn: DeepseekV2MLAAttention,
-    positions: torch.Tensor,
-    hidden_states: torch.Tensor,
-    **model_kwargs,
-) -> torch.Tensor:
-    """Full MLA forward in sglang plugin mode."""
-    forward_batch = model_kwargs.get("forward_batch", None)
-    if forward_batch is None:
-        raise RuntimeError("forward_batch is required in forward_sgl_plugin_mode")
-
-    attn_tp_context = get_attn_tp_context()
-    with attn_tp_context.maybe_input_scattered(forward_batch):
-        if attn.q_lora_rank is not None:
-            attn_tp_context.set_attn_inputs(
-                AttentionInputs(
-                    hidden_states,
-                    forward_batch,
-                    lambda hs, fb: prepare_qkv_latent(attn, hs, fb),
-                )
-            )
-        attn_path = _dispatch_sgl_plugin_attn_path(forward_batch)
-        attn.current_sgl_plugin_attn_path = attn_path
-        if attn_path == "mha":
-            return forward_sgl_plugin_mode_mha(
-                attn,
-                positions,
-                hidden_states,
-                **model_kwargs,
-            )
-        if attn_path == "mla":
-            return forward_sgl_plugin_mode_mla(
-                attn,
-                positions,
-                hidden_states,
-                **model_kwargs,
-            )
-        raise ValueError(f"Unsupported plugin attention path: {attn_path}")
-
-
 def _read_kv_b_proj_weight(attn: DeepseekV2MLAAttention) -> torch.Tensor:
     """Read kv_b_proj weight, handling AWQ and fnuz dtypes."""
     if hasattr(attn.kv_b_proj, "qweight"):