refactor: inherit TRTLLMHAAttnBackend directly from AttentionBackend (#28808)

python/sglang/srt/layers/attention/trtllm_mha_backend.py

@@ -28,10 +28,9 @@
     build_trtllm_mha_page_table,
 )
 from sglang.srt.environ import envs
-from sglang.srt.layers.attention.base_attn_backend import SharedReadEnds
-from sglang.srt.layers.attention.flashinfer_backend import (
-    FlashInferAttnBackend,
-    FlashInferMultiStepDraftBackend,
+from sglang.srt.layers.attention.base_attn_backend import (
+    AttentionBackend,
+    SharedReadEnds,
 )
 from sglang.srt.layers.attention.trtllm_mla_backend import (
     make_persistent_multi_ctas_kv_counter_buffer,
@@ -163,7 +162,7 @@ class TRTLLMMHAMetadata:
     mixed_prefill_tokens: int = 0
 
 
-class TRTLLMHAAttnBackend(FlashInferAttnBackend):
+class TRTLLMHAAttnBackend(AttentionBackend):
     """TRTLLM MHA attention kernel from flashinfer."""
 
     # Build the page table on-device from seq_lens (incl. the SWA-translated table
@@ -187,18 +186,26 @@ def __init__(
         kv_last_page_len_buf: Optional[torch.Tensor] = None,
         speculative_step_id: int = 0,
     ):
-        # Capture workspace size before super().__init__() to preserve user's
-        # SGLANG_FLASHINFER_WORKSPACE_SIZE setting (may be overridden by parent)
+        super().__init__()
+
+        # Capture workspace size before any other init to preserve user's
+        # SGLANG_FLASHINFER_WORKSPACE_SIZE setting.
         env_var = envs.SGLANG_FLASHINFER_WORKSPACE_SIZE
         workspace_size_bytes = (
             env_var.get()
             if env_var.is_set()
             else DEFAULT_WORKSPACE_SIZE_MB * 1024 * 1024
         )
 
-        super().__init__(
-            model_runner, skip_prefill, kv_indptr_buf, kv_last_page_len_buf
-        )
+        # Pool / translator / quant setup — previously inherited from
+        # FlashInferAttnBackend.__init__ but TRTLLM never needed the FlashInfer
+        # wrappers, updaters, or the dozens of tensors that parent allocates.
+        self.req_to_token_pool = model_runner.req_to_token_pool
+        self.token_to_kv_pool = model_runner.token_to_kv_pool
+        self.kv_index_translator = model_runner.kv_index_translator
+        self.kv_cache_quant_method = self.token_to_kv_pool.get_kv_cache_quant_method()
+        self.skip_prefill = skip_prefill
+
         self.prefill_kv_access = self.kv_cache_quant_method.resolve_attention_access(
             "prefill", "trtllm_mha"
         )
@@ -318,6 +325,7 @@ def __init__(
         # SWA hybrid models split the KV cache into full and SWA pools with
         # separate index spaces; SWA layers need a translated page_table.
         self._swa_kv_pool: Optional[SWAKVPool] = self._resolve_swa_kv_pool(model_runner)
+        self.use_sliding_window_kv_pool = self._swa_kv_pool is not None
         # Raw full->swa index mapping tensor for the fused cuda-graph
         # metadata kernel (gather + // page_size happen on device). The unified
         # pool has no token-level mapping, so this is a static-pool mechanism.
@@ -450,6 +458,11 @@ def _check_prefill_kv_access(self) -> None:
             f"trtllm_mha. Available prefill accesses: {available}."
         )
 
+    def _kv_write_scales(self, layer: RadixAttention):
+        if self.kv_cache_quant_method.needs_global_scale():
+            return None, None
+        return layer.k_scale, layer.v_scale
+
     def _nvfp4_output_view(self, q: torch.Tensor) -> torch.Tensor:
         if self._nvfp4_fp8_output is None:
             raise RuntimeError("Native NVFP4 output buffer was not initialized.")
@@ -1897,7 +1910,7 @@ def _trtllm_context_attn(
         return o.view(-1, layer.tp_q_head_num * layer.head_dim)
 
 
-class TRTLLMHAAttnMultiStepDraftBackend(FlashInferMultiStepDraftBackend):
+class TRTLLMHAAttnMultiStepDraftBackend:
     """Multi-step TRTLLM MHA attention kernel used by EAGLE."""
 
     # Per-step backends build the page table on-device (sync-free); mirror that so
@@ -1907,14 +1920,35 @@ class TRTLLMHAAttnMultiStepDraftBackend(FlashInferMultiStepDraftBackend):
     def __init__(
         self, model_runner: ModelRunner, topk: int, speculative_num_steps: int
     ):
-        super().__init__(model_runner, topk, speculative_num_steps)
+        from sglang.srt.utils import get_cuda_graph_max_batch_size
+
+        self.topk = topk
+        self.speculative_num_steps = speculative_num_steps
+
+        # Allocate shared kv_indptr/kv_last_page_len for TRTLLMHAAttnBackend
+        # child instances (FlashInferMultiStepDraftBackend 原本负责这里).
+        max_bs = get_cuda_graph_max_batch_size(
+            model_runner.req_to_token_pool.size * self.topk
+        )
+        self.kv_indptr = torch.zeros(
+            (self.speculative_num_steps, max_bs + 1),
+            dtype=torch.int32,
+            device=model_runner.device,
+        )
+        self.kv_last_page_len = torch.ones(
+            (max_bs,), dtype=torch.int32, device=model_runner.device
+        )
+
+        self.attn_backends = []
         for i in range(self.speculative_num_steps - 1):
-            self.attn_backends[i] = TRTLLMHAAttnBackend(
-                model_runner,
-                skip_prefill=True,
-                kv_indptr_buf=self.kv_indptr[i],
-                kv_last_page_len_buf=self.kv_last_page_len,
-                speculative_step_id=i,
+            self.attn_backends.append(
+                TRTLLMHAAttnBackend(
+                    model_runner,
+                    skip_prefill=True,
+                    kv_indptr_buf=self.kv_indptr[i],
+                    kv_last_page_len_buf=self.kv_last_page_len,
+                    speculative_step_id=i,
+                )
             )
 
     def init_forward_metadata(self, forward_batch: ForwardBatch):