sgl-project/sglang

Refactor: TRTLLMHAAttnBackend should not inherit from FlashInferAttnBackend when reuse is small

Open

#28,808 opened on Jun 20, 2026

View on GitHub
 (1 comment) (0 reactions) (0 assignees)Python (6,216 forks)auto 404
good first issue

Repository metrics

Stars
 (28,442 stars)
PR merge metrics
 (Avg merge 2d 1h) (1,000 merged PRs in 30d)

Description

Summary

TRTLLMHAAttnBackend(FlashInferAttnBackend) (in python/sglang/srt/layers/attention/trtllm_mha_backend.py) inherits from FlashInferAttnBackend, but it reuses very little from the parent. Almost all attention-specific behavior is overridden, and some inherited setup is allocated-but-never-used. We should reconsider this inheritance: if the reuse is small, we should not inherit from FlashInferAttnBackend — inheriting directly from the base AttentionBackend (and pulling in only the small bits of plumbing that are actually shared) would be cleaner.

What is actually reused

The only real reuse comes from super().__init__(...):

  • req_to_token_pool, token_to_kv_pool
  • max_context_len, skip_prefill, use_sliding_window_kv_pool
  • the kv_indptr / kv_last_page_len buffers (used by the companion TRTLLMHAAttnMultiStepDraftBackend draft steps)
  • forward() dispatch — which actually comes from the base AttentionBackend, not from FlashInfer.

What is overridden (not reused)

Essentially all the attention logic:

  • init_forward_metadata, init_cuda_graph_state, init_forward_metadata_out_graph
  • forward_decode, forward_extend (call flashinfer.decode.trtllm_batch_decode_with_kv_cache / flashinfer.prefill.trtllm_batch_context_with_kv_cache directly)
  • get_cuda_graph_seq_len_fill_value, _resolve_swa_kv_pool
  • it uses its own TRTLLMMHAMetadata and its own page-table builder.

Inherited-but-unused / redundant work

Because the TRTLLM kernels use a completely different API, part of the parent __init__ is wasted:

  • Parent builds prefill_wrapper_ragged, prefill_wrappers_paged, decode_wrappers, and the FlashInferIndicesUpdater* objects — TRTLLM-MHA never uses them.
  • Parent sets self.workspace_buffer to the shared global_workspace_buffer, but the child replaces it with its own zero-init global_zero_init_workspace_buffer.
  • max_context_len and _swa_kv_pool are set by the parent and then re-set by the child.

Proposal

  • Evaluate inheriting TRTLLMHAAttnBackend directly from AttentionBackend instead of FlashInferAttnBackend.
  • Move the small amount of genuinely shared init plumbing into a thin shared helper/mixin (e.g. pool/buffer setup) rather than dragging in the full FlashInfer wrapper/indices-updater allocation.
  • Avoid allocating FlashInfer wrappers and updaters that the TRTLLM path never uses.
  • Apply the same review to TRTLLMHAAttnMultiStepDraftBackend(FlashInferMultiStepDraftBackend).

As a general guideline: do not inherit from FlashInferAttnBackend when the reuse is little — prefer inheriting from the base AttentionBackend and sharing only what is actually common.

Good first issue

This is a good first issue: it is well-scoped, localized to trtllm_mha_backend.py (plus the base classes), and mostly a refactor with clear before/after behavior to validate against existing tests.

Contributor guide