Repository navigation
refactor(attention): move linear-attention metadata into a LinearAttnBackend - #837
Merged
Merged
Conversation
This was referenced Oct 28, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Refs #828
The linear-attention layers (GatedDeltaNet and KDA) get a backend object of their own, next to the attention backend. No behavior change.
Summary
Today the per-forward
FLAMetadatais built by the scheduler (build_fla_metadata), the scheduler also stages a second copy of each decode request's state slot inBatch.linear_table_idx, and the buffers the CUDA graphs read (the constantcu_seqlensand the slot map) live in the engine'sGraphCaptureBuffer, filled with the padding slot byGraphRunnerat capture time. This PR moves all of it intoLinearAttnBackendinattention/linear.py, aBaseAttnBackendlaid out like the attention backends:FLAMetadataper forward fromprepare_metadata, and the graph buffers in anFLACaptureData(cu_seqlens,cache_indices) built byinit_capture_graph, pointed at byprepare_for_capture, filled from the forward's metadata byprepare_for_replayand dropped by the basereset_capture. ItsforwardraisesNotImplementedError: the layers run their own kernels.Engine.linear_backend,Noneotherwise) and passes it toGraphRunner. The scheduler callslinear_backend.prepare_metadata(batch)where it calledbuild_fla_metadata. A cache rebuild resets its capture buffers together with the attention backend's.GraphCaptureBufferkeeps only the model inputs: input ids, positions,out_locand mrope positions.FLAMetadata.cache_indicesis the only per-forward copy of the state slots.Batch.linear_table_idxand the scheduler code that staged it per cache mode are gone;build_fla_metadatatakes each request's slot (Req.linear_slot_idx, elseReq.table_idx) on decode as it already did on prefill, and qwen4_exp PLE readsfla_metadata.cache_indiceson decode.linear_slot_idx = padding_slotonly with the hybrid-radix cache. With--cache-type naivethe state is keyed bytable_idx, where slot 0 belongs to a live request, so the padding request now uses its own table row, the pool's spare slot, as the naive decode path on main already did.batch.fla_metadata. GDN and KDA share one backend because their metadata is the same; vLLM does the same, where Kimi's KDA layers use the GDN metadata builder (GDN_ATTN).ctx.attn_backend(a single backend, or the prefill/decodeHybridBackend) is untouched; the two backends sit side by side instead of one wrapping the other, since models call backend-specific methods such asqsa_forwardandmla_forwardonctx.attn_backend.Why now: the speculative verify step (#836) adds verify metadata and graph buffers for these layers. With this PR they go next to the rest of the linear-attention metadata, not into the engine's graph buffer.
Tests
No new tests.
tests/models/qwen4_exp/test_gdn.pyandtest_ple.pynow give their hand-built decode batches anFLAMetadatainstead oflinear_table_idx.Test environment
BaseAttnBackendwithFLACaptureDataand dropBatch.linear_table_idx.Unit tests
1673 passed, 202 skipped, 1 failed. The failure is
tests/models/test_quant_config.py::test_scheme_for_agrees_with_the_stored_tensors[apodex/Apodex-1.1-mini-NVFP4[index-only]](a local checkpoint withquant_method: modelopt_mixed); it fails the same way on main.KDA (
glm5_next) is covered by these unit tests only (tests/models/test_glm5_next_kda_op.py,test_glm5_next_kda_snapshot.py,test_glm5_next_model.py,tests/kernels/test_kda.py); no GLM-5.3 checkpoint was available for an end-to-end run.End-to-end A/B against main
AIME25 problem 1 through
freetoken.llm.LLM(model_path, dtype=torch.bfloat16, max_running_req=4)with the chat template, on main e328cf6 and on this branch, both with CUDA graphs.Greedy, one request at a time (
temperature=0,max_tokens=16384, natural end), compared token by token:Qwen/Qwen3.6-35B-A3B(GDN hybrid, MoE)Qwen/Qwen3.6-27B(GDN hybrid, dense)nvidia/Qwen3.8-Flash-Next-NVFP4(GDN hybrid with PLE and QSA)Greedy, problems 1-3 submitted together (
temperature=0,max_tokens=1024,ignore_eos), so decode runs at batch size 3 on the captured batch-size-4 graph with one padded row; each request compared token by token:Qwen/Qwen3.6-35B-A3BQwen/Qwen3.6-35B-A3B--cache-type naivenvidia/Qwen3.8-Flash-Next-NVFP4Decode speed,
Qwen/Qwen3.6-35B-A3B, recommended sampling fromgeneration_config.json(T=1.0, top_p=0.95, top_k=20), 1024 decode tokens withignore_eos, a 32-token warmup, median of 3 runs: main 159.0 tok/s, this branch 159.1 tok/s.Smoke test on this branch, passed: