Refactor: TRTLLMHAAttnBackend should not inherit from FlashInferAttnBackend when reuse is small
#28,808 opened on Jun 20, 2026
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_poolmax_context_len,skip_prefill,use_sliding_window_kv_pool- the
kv_indptr/kv_last_page_lenbuffers (used by the companionTRTLLMHAAttnMultiStepDraftBackenddraft steps) forward()dispatch — which actually comes from the baseAttentionBackend, not from FlashInfer.
What is overridden (not reused)
Essentially all the attention logic:
init_forward_metadata,init_cuda_graph_state,init_forward_metadata_out_graphforward_decode,forward_extend(callflashinfer.decode.trtllm_batch_decode_with_kv_cache/flashinfer.prefill.trtllm_batch_context_with_kv_cachedirectly)get_cuda_graph_seq_len_fill_value,_resolve_swa_kv_pool- it uses its own
TRTLLMMHAMetadataand 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 theFlashInferIndicesUpdater*objects — TRTLLM-MHA never uses them. - Parent sets
self.workspace_bufferto the sharedglobal_workspace_buffer, but the child replaces it with its own zero-initglobal_zero_init_workspace_buffer. max_context_lenand_swa_kv_poolare set by the parent and then re-set by the child.
Proposal
- Evaluate inheriting
TRTLLMHAAttnBackenddirectly fromAttentionBackendinstead ofFlashInferAttnBackend. - 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.