Skip to content

refactor(attention): move linear-attention metadata into a LinearAttnBackend - #837

Merged
jason-fxz merged 3 commits into
mainfrom
refactor/linear-attn-backend
Oct 28, 2026
Merged

jason-fxz merged 3 commits into
mainfrom
refactor/linear-attn-backend

Conversation

@jason-fxz

@jason-fxz jason-fxz commented Oct 28, 2026 •

Copy link
Copy Markdown
Collaborator

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 FLAMetadata is built by the scheduler (build_fla_metadata), the scheduler also stages a second copy of each decode request's state slot in Batch.linear_table_idx, and the buffers the CUDA graphs read (the constant cu_seqlens and the slot map) live in the engine's GraphCaptureBuffer, filled with the padding slot by GraphRunner at capture time. This PR moves all of it into LinearAttnBackend in attention/linear.py, a BaseAttnBackend laid out like the attention backends: FLAMetadata per forward from prepare_metadata, and the graph buffers in an FLACaptureData (cu_seqlens, cache_indices) built by init_capture_graph, pointed at by prepare_for_capture, filled from the forward's metadata by prepare_for_replay and dropped by the base reset_capture. Its forward raises NotImplementedError: the layers run their own kernels.

  • The engine creates it when the model has linear-attention layers (Engine.linear_backend, None otherwise) and passes it to GraphRunner. The scheduler calls linear_backend.prepare_metadata(batch) where it called build_fla_metadata. A cache rebuild resets its capture buffers together with the attention backend's.
  • GraphCaptureBuffer keeps only the model inputs: input ids, positions, out_loc and mrope positions.
  • FLAMetadata.cache_indices is the only per-forward copy of the state slots. Batch.linear_table_idx and the scheduler code that staged it per cache mode are gone; build_fla_metadata takes each request's slot (Req.linear_slot_idx, else Req.table_idx) on decode as it already did on prefill, and qwen4_exp PLE reads fla_metadata.cache_indices on decode.
  • The padding request carries linear_slot_idx = padding_slot only with the hybrid-radix cache. With --cache-type naive the state is keyed by table_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.
  • The GDN and KDA layers do not change: they still call their fla kernels and read 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/decode HybridBackend) is untouched; the two backends sit side by side instead of one wrapping the other, since models call backend-specific methods such as qsa_forward and mla_forward on ctx.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.py and test_ple.py now give their hand-built decode batches an FLAMetadata instead of linear_table_idx.

Test environment

  • GPU: 1x NVIDIA H100 80GB HBM3, driver 580.95.05
  • CPU: Intel Xeon Platinum 8480+ (26 vCPUs), 221 GB RAM, Ubuntu 22.04.5
  • CUDA 13.0, torch 2.14.1+cu130, Python 3.12.12
  • A/B against main e328cf6. The unit tests and every greedy run except the 27B one ran on the head, d6e1e11. The speed run, the 27B greedy run and the smoke test ran on the first commit (before a comment-only fix); the later commits make the backend a BaseAttnBackend with FLACaptureData and drop Batch.linear_table_idx.

Unit tests

FREETOKEN_QWEN36_MODEL=/path/to/Qwen3.6-27B PYTHONPATH=python pytest tests/scheduler tests/engine tests/kvcache tests/attention tests/moe/test_offload.py tests/server tests/tokenizer tests/models tests/kernels/test_kda.py

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 with quant_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:

Checkpoint main this branch Token by token
Qwen/Qwen3.6-35B-A3B (GDN hybrid, MoE) 16384 tokens (hit the limit) 16384 tokens (hit the limit) identical
Qwen/Qwen3.6-27B (GDN hybrid, dense) 15676 tokens, answer 70 15676 tokens, answer 70 identical
nvidia/Qwen3.8-Flash-Next-NVFP4 (GDN hybrid with PLE and QSA) 727 tokens, answer 70 727 tokens, answer 70 identical

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:

Checkpoint cache Token by token, all 3 requests
Qwen/Qwen3.6-35B-A3B hybrid-radix (default) identical
Qwen/Qwen3.6-35B-A3B --cache-type naive identical
nvidia/Qwen3.8-Flash-Next-NVFP4 hybrid-radix (default) identical

Decode speed, Qwen/Qwen3.6-35B-A3B, recommended sampling from generation_config.json (T=1.0, top_p=0.95, top_k=20), 1024 decode tokens with ignore_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:

FREETOKEN_TEST_MODEL=/path/to/Qwen3.6-35B-A3B FREETOKEN_AIME25_JSONL=/path/to/aime25.jsonl PYTHONPATH=python pytest tests/e2e/test_aime.py

@jason-fxz
jason-fxz merged commit a2ce636 into main Oct 28, 2026
1 check passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant