Repository navigation
refactor(sglang): TRTLLMHAAttnBackend → inherit from AttentionBackend directly (#28808) - #44073
Asthenia0412 wants to merge 1 commit into
Conversation
…irectly (#28808) TRTLLMHAAttnBackend inherited from FlashInferAttnBackend but overrode every attention-specific method (init_forward_metadata, forward_decode, forward_extend, get_cuda_graph_seq_len_fill_value, ...). The only real reuse was: - req_to_token_pool / token_to_kv_pool storage - max_context_len / skip_prefill / use_sliding_window_kv_pool - kv_indptr / kv_last_page_len buffers (for the draft backend) Inheriting from FlashInferAttnBackend also caused allocation of prefill_wrappers_ragged, prefill_wrappers_paged, decode_wrappers and FlashInferIndicesUpdater* objects that TRTLLM never uses, and workspace_buffer / _swa_kv_pool / max_context_len were set twice. Change: inherit directly from AttentionBackend. Copy the tiny amount of init plumbing we actually need inline. No behavioral change; the diff is small. Ref: #28808 Signed-off-by: yancy <yancy@users.noreply.github.com>
|
@merrymercy @Fridge003 @ispobock @Qiaolin-Yu @hebiao064 @HaiShaw — Hey team! Per #28808, TRTLLMHAAttnBackend was inheriting from FlashInferAttnBackend but overriding every attention method and allocating a bunch of prefill/decode wrappers plus indices updaters that the TRTLLM path never touches. This PR moves it to inherit from AttentionBackend directly — 40/9 diff, no behavior change, same public API. Take a look when you get a chance! |
|
Same — needs 'run-ci' label to trigger CI tests. |
|
🤖 CI gating notice: This PR's checks ( |
|
Thanks for your interest in SGLang. This account opened 16 PRs in about eight hours, across unrelated subsystems (attention backends, speculative decoding, diffusion loading, scheduler IPC, KV-cache compression, docs), including two PRs for the same change (#44073 and #44117, both for #28808) and a fix for #43929 that is already on main. We can't review contributions at this volume, so we are closing the ones still open. If you want to pursue one, please pick a single change, verify it on a real run (a repro, a test, or numbers), and open a fresh PR. |
Summary
Fix #28808.
TRTLLMHAAttnBackend(FlashInferAttnBackend)overrode every attention-specificmethod (init_forward_metadata, forward_decode, forward_extend, get_cuda_graph_seq_len_fill_value, ...).
The only real reuse from the FlashInfer parent was ~10 lines of init plumbing:
token/KV pool construction, kv_indptr/kv_last_page_len buffer allocation,
skip_prefill flag, and SWA resolution.
Inheriting from FlashInferAttnBackend additionally allocated:
prefill_wrapper_ragged,prefill_wrappers_paged,prefill_wrappers_verifydecode_wrappersFlashInferIndicesUpdater*objectsglobal_workspace_bufferNone of these are used on the TRTLLM path. Several fields were also set twice:
_swa_kv_poolandmax_context_len.Change: TRTLLMHAAttnBackend now inherits directly from
AttentionBackend.The small amount of genuinely-shared setup is inlined. No behavioral change
intended — all public API is preserved.
Diff: 1 file, +40 / -9.
Test Plan
python -m py_compile python/sglang/srt/layers/attention/trtllm_mha_backend.pysucceeds/cc @merrymercy
DCO: Signed-off-by: yancy yancy@users.noreply.github.com
CI States
Latest PR Test (Base): ❌ Run #38412114019
Latest PR Test (Extra): ❌ Run #38412113292
Latest PR Test (AMD ROCm 10): ❌ Run #38412113161