From 81c61d909b66f9f70c8e70596078fd87781a8b72 Mon Sep 17 00:00:00 2001 From: junq <22017000+QiJune@users.noreply.github.com> Date: Thu, 2 Apr 2026 20:31:21 +0800 Subject: [PATCH 1/7] add attention developer guide Signed-off-by: junq <22017000+QiJune@users.noreply.github.com> --- AGENTS.md | 1 + .../modules/ATTENTION_DEVELOPER_GUIDE.md | 665 ++++++++++++++++++ 2 files changed, 666 insertions(+) create mode 100644 tensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.md diff --git a/AGENTS.md b/AGENTS.md index 39ea7e4fb6f7..64b83ce5972a 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -83,6 +83,7 @@ HuggingFace Model → LLM API → Executor (PyTorch/AutoDeploy/TensorRT) | `tensorrt_llm/executor/executor.py` | Execution abstraction (`GenerationExecutor`) | | `tensorrt_llm/models/automodel.py` | Auto-discovery and model registry | | `tensorrt_llm/_torch/models/` | PyTorch backend model implementations (distinct from `models/` used by TensorRT backend) | +| `tensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.md` | Attention, MLA, backend families, sparse backends, metadata contracts, and KV-cache behavior - **read before modifying `tensorrt_llm/_torch/modules/attention.py` or `_torch/attention_backend/`** | | `tensorrt_llm/_torch/modules/fused_moe/MOE_DEVELOPER_GUIDE.md` | MoE architecture, backends, communication, development patterns — **read before modifying MoE code** | | `CODING_GUIDELINES.md` | C++ and Python coding standards (referenced throughout, must read before contributing) | diff --git a/tensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.md b/tensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.md new file mode 100644 index 000000000000..5ea3c112c400 --- /dev/null +++ b/tensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.md @@ -0,0 +1,665 @@ +# Attention Developer Guide + +## Scope + +This guide covers the TRT-LLM PyTorch attention stack: + +- `tensorrt_llm/_torch/modules/attention.py` +- `tensorrt_llm/_torch/attention_backend/` +- `tensorrt_llm/_torch/attention_backend/sparse/` + +Use it when modifying the current implementation or adding a new model's +attention behavior. It covers standard `Attention`, MLA, dense backends, and +sparse backends. It does not cover +`tensorrt_llm/_torch/attention_backend/star_flashinfer.py`, which is planned +for deprecation. + +## How to Read the Stack + +Attention in TRT-LLM is split across four layers: + +1. module wrapper (`Attention` or `MLA`) +2. backend class selected by `config.attn_backend` +3. metadata subtype and runtime buffers +4. KV-cache manager and decode-time cache semantics + +Keep these four questions separate: + +1. What math happens at the module layer around the backend call? +2. Which backend family can execute the core attention path? +3. Which metadata subtype and runtime contract does that backend require? +4. What KV-cache ownership and decode-time semantics does the path assume? + +The same module math can still require a different backend, metadata subtype, +KV-cache manager, or a fallback from `TRTLLM` to `VANILLA`. Attention work is +not only score computation. It also includes how the backend reads, writes, +appends, and reuses KV cache, especially during decode. + +## 1. Module Layer Reference + +### 1.1 `Attention`: the module wrapper around the backend + +`Attention` is not just the backend call. It owns the logic around the backend: + +- QKV projection and output projection +- TP/CP reshaping and mapping setup +- fused or split QKV handling +- optional unfused RoPE +- optional output gating +- optional LoRA injection +- passing masks, sinks, and metadata into the backend + +At a high level: + +```text +hidden_states + -> qkv_proj + -> optional LoRA + -> optional gate split + -> optional unfused RoPE + -> fused/split QKV conversion + -> backend.forward(...) + -> optional output gate + -> o_proj +``` + +Important extension points in `Attention`: + +- `apply_rope()` +- `apply_qk_norm()` +- `convert_qkv()` + +If a source model needs extra Q/K processing, gating, scaling, or projection +layout logic, the first question is whether it can stay at this module layer +without changing the outer runtime contract. + +### 1.2 `MLA`: a separate module on top of the same backend system + +`MLA` is a separate module in `attention.py`. It keeps module-level +projections and tensor transforms in `MLA`, delegates core execution to backend +objects, and depends on metadata and KV-cache contract. It owns: + +- low-rank Q decomposition +- low-rank KV decomposition +- absorbed MLA path +- DSA-specific dispatch +- short-seq MHA routing +- MLA-specific RoPE and latent-cache flow + +`MLA` has two projection layouts. + +Non-lite MLA (`is_lite == False`): + +```text +MLA(nn.Module) +├── kv_a_proj_with_mqa +├── q_a_layernorm +├── q_b_proj +├── kv_a_layernorm +├── kv_b_proj +├── mha # dense backend, used by short-seq MHA path +├── mqa # sparse / DSA backend +└── short_seq_mha_threshold +``` + +Lite MLA (`is_lite == True`): + +```text +MLA(nn.Module) +├── kv_a_proj_with_mqa +├── q_proj # also assigned to q_b_proj +├── kv_a_layernorm +├── kv_b_proj +├── mha +├── mqa +└── short_seq_mha_threshold +``` + +In lite mode there is no separate `q_a_proj`, `q_a_layernorm`, or `kv_a_proj`. +`q_proj` is used as `q_b_proj`. + +### 1.3 MLA dispatch reference + +For DSA-style MLA models, the dispatch is: + +```text +forward() + -> forward_impl_with_dsa() + -> forward_dsa_proj() + -> forward_dsa_attn() + -> context tokens? + -> yes: forward_context_dsa() + -> short-seq gate? + -> yes: forward_context() + -> forward_context_default() + -> forward_context_with_cached_kv() + -> forward_context_with_chunked_prefill() + -> no: absorption / sparse MLA path + -> generation tokens? + -> yes: forward_generation_dsa() +``` + +For MLA maintenance: + +- the short-seq MHA path should stay inside `forward_context()` +- do not bypass the dispatcher and call `forward_context_default()` directly + +`forward_context_default()` only handles fresh context. It does not cover +cached-KV or chunked-context cases. + +### 1.4 Practical MLA notes + +Structure and dispatch: + +- `is_lite` changes the projection structure, not just a small code path. +- `self.is_dsa == True` means the DSA path is active. +- `self.mqa` is the sparse DSA backend. +- `self.mha` is a dense backend used only for the short-seq dense context path. +- `self.mha` being present does not mean DSA is disabled. +- `_should_use_short_mha()` uses `max_ctx_kv_len` when available, not just the + new-token count. + +`torch.compile` path: + +- the compiled path may use custom-op based execution paths +- under `torch.compile`, `_should_use_short_mha()` returns `False`, so the + split DSA path is always used + +Helix CP path: + +- `_helix_cp_allgather_input()` runs before the attention body on layers after + the first +- `_helix_cp_output_projection()` runs after the attention body + +## 2. Backend Layer Reference + +### 2.1 Backend selection + +Backends are created by: + +- `get_attention_backend()` +- `create_attention()` + +The backend is chosen from: + +- `config.attn_backend` +- optional `sparse_attention_config` +- optional MLA parameters + +Base backend families: + +| Backend name | Class | Metadata subtype | Notes | +|---|---|---|---| +| `TRTLLM` | `TrtllmAttention` | `TrtllmAttentionMetadata` | Standard backend path | +| `VANILLA` | `VanillaAttention` | `VanillaAttentionMetadata` | Torch fallback path | +| `FLASHINFER` | `FlashInferAttention` | `FlashInferAttentionMetadata` | FlashInfer planning/runtime path | + +### 2.2 Sparse backend families + +Sparse attention is not selected by a separate top-level module. It is resolved +through `sparse_attention_config` on top of a base backend family. + +Current sparse registrations: + +| Base backend | Sparse algorithm | Resulting backend class | Metadata subtype | KV-cache manager | +|---|---|---|---|---| +| `TRTLLM` | `rocket` | `RocketTrtllmAttention` | `RocketTrtllmAttentionMetadata` | `RocketKVCacheManager` | +| `TRTLLM` | `dsa` | `DSATrtllmAttention` | `DSAtrtllmAttentionMetadata` | `DSACacheManager` | +| `TRTLLM` | `skip_softmax` | `TrtllmAttention` | `TrtllmAttentionMetadata` | standard `KVCacheManager` | +| `VANILLA` | `rocket` | `RocketVanillaAttention` | `RocketVanillaAttentionMetadata` | `RocketKVCacheManager` | +| `VANILLA` | `dsa` | unsupported | none | none | +| `VANILLA` | `skip_softmax` | unsupported | none | none | +| `FLASHINFER` | sparse variants | unsupported | none | none | + +### 2.3 Backend contract + +All backends implement the `AttentionBackend` interface. + +The core contract is: + +- `forward(q, k, v, metadata, attention_mask=..., **kwargs)` +- `Metadata` subtype +- coarse capability hooks: + - `support_fused_rope()` + - `support_fused_qkv()` + - `support_mla()` + +Those capability hooks are coarse checks. They do not prove that every +required operator or sparse path already exists. + +### 2.4 Capability reference + +The current coarse capability picture is: + +| Backend family | Fused RoPE | Fused QKV input | MLA | +|---|---|---|---| +| `TrtllmAttention` | yes | yes | yes | +| `VanillaAttention` | no | no | no | +| `FlashInferAttention` | no | no | no | + +Sparse subclasses inherit the base backend family and then add sparse-specific +metadata and cache behavior. + +### 2.5 `TRTLLM` internal kernel paths + +`TrtllmAttention` can dispatch to `trtllm_gen.py` for supported dense cases. +That is an internal fast path, not a separate top-level backend selection. + +It only applies to a subset of dense cases. If it does not apply, +`TrtllmAttention` stays on its regular runtime path. + +## 3. Runtime Contract Reference + +### 3.1 Metadata families + +All backend metadata types inherit from `AttentionMetadata`. + +#### 3.1.1 Base `AttentionMetadata` + +The common contract includes: + +- sequence-length and request-level runtime state +- KV-cache manager and KV-cache parameters +- runtime feature flags +- optional sparse-attention state +- optional CUDA-graph buffer management + +The base metadata also supports cross attention through `seq_lens_kv` and +cross sub-metadata. Backend support is not uniform: + +- `FlashInferAttention` has explicit cross-attention handling +- `TrtllmAttention` currently asserts that cross attention is not supported + +#### 3.1.2 `TrtllmAttentionMetadata` + +`TrtllmAttentionMetadata` is the main metadata family for the standard TRTLLM +path. + +It extends the base metadata with: + +- paged-KV block information +- request and sequence state for TRTLLM runtime execution +- runtime state for chunked prefill, speculative decode, and Helix +- MLA-specific runtime state when MLA is active + +If a source attention implementation needs paged KV, chunked prefill, +FlashMLA, speculative decoding, or Helix-aware execution, the fit question is +mostly a `TrtllmAttentionMetadata` fit question. + +#### 3.1.3 `VanillaAttentionMetadata` + +`VanillaAttentionMetadata` mainly prepares: + +- base attention metadata +- simple cache-index information for torch-side cache access + +The actual attention computation is mostly done in torch. Use it when the +current `Attention` module boundary still fits but the fused TRTLLM path is +too restrictive. + +#### 3.1.4 `FlashInferAttentionMetadata` + +`FlashInferAttentionMetadata` adds a planning-oriented runtime contract: + +- workspace and planning state +- page-table style KV metadata +- prefill and decode wrapper state + +This backend requires a planning step and a paged-table runtime shape. + +#### 3.1.5 Sparse metadata families + +Sparse backends extend the base metadata family rather than inventing an +unrelated attention interface. + +`DSAtrtllmAttentionMetadata` extends `TrtllmAttentionMetadata` with: + +- indexer-side sparse runtime state +- top-k and token-to-request routing state +- extra state for sparse context and generation flows + +`RocketTrtllmAttentionMetadata` and `RocketVanillaAttentionMetadata` add: + +- sparse-window and routing state +- sparse offsets and sequence state +- KT-cache related state + +### 3.2 KV-cache and decode-time semantics + +The main question is not just "does the backend read K and V?" but: + +- who owns the cache +- what cache layout the backend assumes +- how new tokens are appended +- whether decode updates happen in place +- how pages or blocks are indexed +- whether cached KV can be revisited during context +- whether sparse state must be maintained alongside KV + +A backend may support the score computation you want, but still be the wrong +fit because it assumes a different KV-cache layout or a different decode-time +update pattern. + +#### 3.2.1 Common paged-KV model + +All current `_torch` backends are built on top of paged KV cache. + +At the cache-manager level, `KVCacheManager.get_buffers()` exposes a per-layer +view of the primary pool in two common layouts: + +- `NHD`: `[num_pages, kv_factor, tokens_per_block, num_kv_heads, head_dim]` +- `HND`: `[num_pages, kv_factor, num_kv_heads, tokens_per_block, head_dim]` + +For standard dense attention, `kv_factor = 2`, which means separate K and V +planes. + +For MLA-style cache, `kv_factor = 1`. The cache stores one latent-cache tensor +per token rather than separate K and V planes. + +The main differences across backends are: + +- which tensor view or block table the backend expects +- whether the backend writes cache internally or Python writes it explicitly +- whether extra side cache is required in addition to the main KV cache + +#### 3.2.2 `TRTLLM` backend + +At the C++ boundary, the main contract is a block-offset table plus pool +pointers, not only a reshaped cache tensor view. + +In the Python metadata path, this appears as fields such as: + +- `kv_cache_block_offsets` +- `host_kv_cache_pool_pointers` +- `host_kv_cache_pool_mapping` +- `cache_indirection` + +Read and write behavior: + +- decode and prefill cache updates are backend-managed +- dense cache writes go through backend ops such as `qkv_preprocessing` +- some sparse updates go through backend postprocessing +- the regular path reads cached KV through the block-table and pool-pointer + contract + +Python does not update the main KV cache with `index_copy_` in the regular +`TRTLLM` path. + +#### 3.2.3 `TRTLLM` internal `trtllm_gen` path + +`trtllm_gen.py` is not a separate top-level backend, but it has its own KV +view assumptions. + +It still starts from the same paged cache state, but it bridges the TRTLLM +block-offset format into the page-table shape expected by the FlashInfer +`trtllm_gen` kernels. + +In practice: + +- cache writes still begin from TRTLLM preprocessing +- the fast path converts K-side block offsets into shared page indices +- it reads cache again through `KVCacheManager.get_buffers(...)` + +This path is narrower than the main `TRTLLM` backend: + +- dense only +- no MLA +- no sparse attention +- no cross attention +- fused QKV only + +If `trtllm_gen` does not fit, that does not rule out the main `TRTLLM` +backend. + +#### 3.2.4 `FlashInfer` backend + +`FlashInfer` also uses paged KV cache, but its runtime contract is more +directly page-table oriented. It reads the cache through +`kv_cache_manager.get_buffers(...)` using the layout requested by +`metadata.kv_layout`. + +Read and write behavior: + +- Python explicitly appends current K and V into paged cache +- FlashInfer wrappers then read from that paged cache using page-table metadata + +#### 3.2.5 `VANILLA` backend + +`VANILLA` gets a paged cache tensor from `kv_cache_manager.get_buffers()` and +request block ids from `block_ids_per_seq`. + +Read and write behavior: + +- Python writes K and V directly into cache +- Python slices the same cache tensor to rebuild key and value states +- sparse token filtering, if any, also happens around this Python-side path + +#### 3.2.6 MLA cached-context semantics + +MLA adds a different cache shape and different decode-time assumptions. + +The main difference is that MLA cached state is not regular dense K and V. +The paged cache stores latent-cache state, and backend ops handle: + +- appending latent cache into paged storage +- applying RoPE as part of that flow +- loading paged cached state back for attention use + +MLA fit cannot be judged from attention math alone. The module and backend also +have to agree on: + +- latent-cache layout +- paged-KV read path +- paged-KV write path +- cached-context and chunked-context behavior + +The short-seq MHA path is only correct if cached-KV behavior stays inside the +top-level `forward_context()` dispatcher. + +#### 3.2.7 Sparse side-cache semantics + +Sparse backends may change more than the attention score path. They may also +add side caches. + +`skip_softmax` keeps the standard `KVCacheManager`. + +`DSA` uses `DSACacheManager`: + +- the main cache is still the paged MLA-style cache +- DSA also adds `indexer_k_cache` +- this side cache is used for sparse indexing and top-k related work + +`Rocket` uses `RocketKVCacheManager`: + +- the main cache still follows the base backend family +- Rocket also adds `KT` cache +- this side cache is used for sparse routing and sparse block selection + +When evaluating a new sparse attention implementation, check: + +- the main KV-cache contract +- the side-cache contract + +## 4. Evaluating New Attention + +### 4.1 First-pass fit against the current `TRTLLM` backend + +For a new model, first compare it against the current `TRTLLM` backend surface +in four parts. + +| What to compare | Current `TRTLLM` backend surface | +|---|---| +| Module-layer math around the backend call | `Attention` and `MLA` already handle QKV projection, fused or split QKV conversion, optional unfused RoPE, Q/K normalization hooks, output gating, LoRA, TP/CP handling, and Helix wrappers. | +| Backend-side execution | `TrtllmAttention` handles the standard dense path. It supports fused RoPE, fused QKV input, MLA mode, sliding-window style `attention_window_size`, mRoPE config, attention sinks, and sparse variants through registered sparse backends. | +| Metadata and runtime contract | The direct path expects `TrtllmAttentionMetadata` or one of its sparse subclasses. That means the new attention must still fit request-based metadata, paged-KV metadata, runtime feature flags, and any extra sparse or MLA buffers required by the path. | +| KV-cache semantics | The direct path assumes paged-KV ownership. Dense TRTLLM uses the standard `KVCacheManager`. Sparse paths can swap in `DSACacheManager` or `RocketKVCacheManager`. Backend choice also implies assumptions about KV-cache layout, decode-time reads, and inplace append or update behavior. If the new attention needs a different cache ownership model or different decode-time cache semantics, it is not a direct fit. | + +A practical check: + +1. Write down the new attention in four buckets: + - module math before or after the backend call + - backend features it needs + - metadata and runtime state it needs + - KV-cache ownership, layout assumptions, and update rules it needs +2. Compare each bucket against the current `TRTLLM` backend surface above. +3. Treat the first mismatch as the current blocker. + +For direct fit, ask these questions in order: + +1. Can `Attention` or `MLA` express the new math with module-side code only? +2. Can `TrtllmAttention.forward(...)` express the backend call shape that is + needed? + - fused or split QKV + - dense or sparse path + - RoPE or mRoPE handling + - windowed masking + - MLA or non-MLA +3. Can the runtime state be stored in `TrtllmAttentionMetadata` or a known + sparse metadata subclass? +4. Can the KV-cache behavior stay inside the current paged-KV and cache-manager + model? + +If the answer stays "yes" through all four questions, start with the current +`TRTLLM` backend and then check the runtime contract in Section 3. + +### 4.2 Detailed checklist + +#### 4.2.1 Module-level math + +Check all math around the backend call: + +- Q/K/V layout +- fused or split QKV input +- MQA or GQA structure +- Q/K normalization +- extra QK scaling terms +- output gating +- pre-backend and post-backend transforms + +The first question is whether this math can stay at the module layer by +overriding or extending `Attention` / `MLA`, without changing the outer runtime +contract. + +#### 4.2.2 Backend capability mapping + +Check which backend family can actually run the source behavior: + +- `TRTLLM` +- `VANILLA` +- `FLASHINFER` +- sparse variants on top of those families + +Then check capability assumptions explicitly: + +- fused RoPE +- fused QKV input +- MLA support +- sparse operator support +- chunked-context support + +Do not use backend name alone as proof of support. For a first-pass fit check, +use Section 4.1 before looking at other backends. + +#### 4.2.3 Positional embedding and mask contract + +Check: + +- whether RoPE is applied outside the backend or fused into it +- standard RoPE vs mRoPE +- causal, full, sliding-window, or custom masks +- any sink-token or mask-side special logic + +#### 4.2.4 Metadata and runtime contract + +Check which metadata subtype is required: + +- `TrtllmAttentionMetadata` +- `VanillaAttentionMetadata` +- `FlashInferAttentionMetadata` +- sparse metadata subclasses + +Then check whether the source behavior needs runtime state such as: + +- request-level state +- KV-cache manager handles and KV-cache parameters +- runtime feature flags +- planning or sparse buffers +- CUDA-graph assumptions + +#### 4.2.5 KV-cache ownership and decode-time semantics + +Check: + +- how K/V are appended +- what KV-cache layout the backend expects +- how tokens are indexed in cache +- whether cached KV is revisited during context +- whether decode updates require inplace cache writes +- whether cache reuse or chunked prefill is required +- whether speculative decoding assumptions are involved +- whether sparse state is attached to cache ownership + +### 4.3 Bring-up order + +Start with the `TRTLLM` backend when the new attention fits the existing +runtime contract or only needs limited changes. If that path is too costly for +initial bring-up or quick experiments, use `VANILLA` to validate the math and +outer module behavior first, then re-evaluate whether the implementation +should move back to `TRTLLM`. + +Working rules: + +- Stay on `Attention` or `MLA` plus an existing backend family when possible. +- Extend the `TRTLLM` backend path before adding a new backend. +- Extend module-level hooks before adding a new backend. +- Follow an existing sparse family pattern before adding a new sparse + abstraction. +- Treat cache-manager mismatch as a real blocker. + +## 5. Key File Map + +| File | Role | +|------|------| +| `tensorrt_llm/_torch/modules/attention.py` | Standard attention and MLA module logic | +| `tensorrt_llm/_torch/attention_backend/interface.py` | Backend contract, base metadata, capability hooks | +| `tensorrt_llm/_torch/attention_backend/utils.py` | Backend and sparse-backend selection | +| `tensorrt_llm/_torch/attention_backend/trtllm.py` | TRTLLM backend and metadata | +| `tensorrt_llm/_torch/attention_backend/trtllm_gen.py` | Internal dense fast path used by `TrtllmAttention` on supported Blackwell configurations | +| `tensorrt_llm/_torch/attention_backend/vanilla.py` | Torch fallback backend and metadata | +| `tensorrt_llm/_torch/attention_backend/flashinfer.py` | FlashInfer backend and metadata | +| `tensorrt_llm/_torch/attention_backend/sparse/dsa.py` | DSA sparse backend, metadata, indexer, cache manager | +| `tensorrt_llm/_torch/attention_backend/sparse/kernel.py` | Triton helper kernels used by sparse attention implementations | +| `tensorrt_llm/_torch/attention_backend/sparse/rocket.py` | Rocket sparse backends, metadata, cache manager | +| `tensorrt_llm/_torch/attention_backend/sparse/utils.py` | Sparse backend and sparse cache-manager registration | +| `tensorrt_llm/_torch/models/modeling_deepseekv3.py` | DeepSeek weight loading, including `kv_b_proj` layout transforms | + +## 6. Testing Notes + +- `mla.to(device)` does not move `mla.mqa.indexer` weights automatically. +- Copy indexer weights explicitly in A/B tests. +- Initialize `kv_b_proj` weights in loaded TRT-LLM layout, not HuggingFace + layout. +- Test lite and non-lite MLA separately when changing projection logic. +- Test eager and compiled paths separately when changing DSA MLA dispatch or + custom-op behavior. +- Test fresh context, cached context, chunked context, and generation + separately. +- Any dispatch change touching `forward_context()` needs chunked-context tests. + +Useful tests: + +- `tests/unittest/_torch/attention/test_attention.py` +- `tests/unittest/_torch/attention/test_attention_mla.py` +- `tests/unittest/_torch/attention/test_vanilla_attention.py` +- `tests/unittest/_torch/attention/test_flashinfer_attention.py` +- `tests/unittest/_torch/attention/sparse/test_short_seq_mha.py` +- `tests/unittest/_torch/attention/sparse/test_sparse_mla_forward.py` + +## 7. Anti-Patterns + +- Do not treat attention work as "math only". +- Do not treat backend choice as independent from metadata choice. +- Do not treat KV-cache semantics as a small implementation detail. +- Do not call `forward_context_default()` directly for chunked MLA context. +- Do not duplicate RoPE handling before checking the fused path. +- Do not assume `self.mha is not None` means DSA is disabled. From d318cfb2c7cee1f84b82dd38e21cf88ecb406dfe Mon Sep 17 00:00:00 2001 From: junq <22017000+QiJune@users.noreply.github.com> Date: Fri, 3 Apr 2026 14:41:24 +0800 Subject: [PATCH 2/7] simplify doc Signed-off-by: junq <22017000+QiJune@users.noreply.github.com> --- .../modules/ATTENTION_DEVELOPER_GUIDE.md | 505 ++++-------------- 1 file changed, 107 insertions(+), 398 deletions(-) diff --git a/tensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.md b/tensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.md index 5ea3c112c400..a3b380d9fdf2 100644 --- a/tensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.md +++ b/tensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.md @@ -86,90 +86,25 @@ objects, and depends on metadata and KV-cache contract. It owns: - short-seq MHA routing - MLA-specific RoPE and latent-cache flow -`MLA` has two projection layouts. +`MLA` has two projection layouts: non-lite (`is_lite == False`) and lite +(`is_lite == True`). In lite mode there is no separate Q low-rank compression +stage. `is_lite` changes the projection structure, not just a small code path. -Non-lite MLA (`is_lite == False`): +**DSA dispatch.** DSA-style MLA uses a multi-stage dispatch: projection and +attention are separated, context and generation paths are separated, and a +short-seq gate inside the context path can route to a dense MHA fallback. The +short-seq MHA path should stay inside `forward_context()`. Do not bypass the +dispatcher and call `forward_context_default()` directly — it only handles +fresh context, not cached-KV or chunked-context cases. -```text -MLA(nn.Module) -├── kv_a_proj_with_mqa -├── q_a_layernorm -├── q_b_proj -├── kv_a_layernorm -├── kv_b_proj -├── mha # dense backend, used by short-seq MHA path -├── mqa # sparse / DSA backend -└── short_seq_mha_threshold -``` - -Lite MLA (`is_lite == True`): - -```text -MLA(nn.Module) -├── kv_a_proj_with_mqa -├── q_proj # also assigned to q_b_proj -├── kv_a_layernorm -├── kv_b_proj -├── mha -├── mqa -└── short_seq_mha_threshold -``` - -In lite mode there is no separate `q_a_proj`, `q_a_layernorm`, or `kv_a_proj`. -`q_proj` is used as `q_b_proj`. - -### 1.3 MLA dispatch reference - -For DSA-style MLA models, the dispatch is: - -```text -forward() - -> forward_impl_with_dsa() - -> forward_dsa_proj() - -> forward_dsa_attn() - -> context tokens? - -> yes: forward_context_dsa() - -> short-seq gate? - -> yes: forward_context() - -> forward_context_default() - -> forward_context_with_cached_kv() - -> forward_context_with_chunked_prefill() - -> no: absorption / sparse MLA path - -> generation tokens? - -> yes: forward_generation_dsa() -``` +**Practical notes:** -For MLA maintenance: - -- the short-seq MHA path should stay inside `forward_context()` -- do not bypass the dispatcher and call `forward_context_default()` directly - -`forward_context_default()` only handles fresh context. It does not cover -cached-KV or chunked-context cases. - -### 1.4 Practical MLA notes - -Structure and dispatch: - -- `is_lite` changes the projection structure, not just a small code path. -- `self.is_dsa == True` means the DSA path is active. -- `self.mqa` is the sparse DSA backend. -- `self.mha` is a dense backend used only for the short-seq dense context path. - `self.mha` being present does not mean DSA is disabled. - `_should_use_short_mha()` uses `max_ctx_kv_len` when available, not just the new-token count. - -`torch.compile` path: - -- the compiled path may use custom-op based execution paths -- under `torch.compile`, `_should_use_short_mha()` returns `False`, so the - split DSA path is always used - -Helix CP path: - -- `_helix_cp_allgather_input()` runs before the attention body on layers after - the first -- `_helix_cp_output_projection()` runs after the attention body +- Under `torch.compile`, `_should_use_short_mha()` returns `False`, so the + split DSA path is always used. +- Helix CP wraps the attention body with allgather/output-projection helpers. ## 2. Backend Layer Reference @@ -197,19 +132,19 @@ Base backend families: ### 2.2 Sparse backend families Sparse attention is not selected by a separate top-level module. It is resolved -through `sparse_attention_config` on top of a base backend family. +through `sparse_attention_config` on top of a base backend family. Sparse +selection can change the backend class, metadata subtype, and KV-cache manager. Current sparse registrations: -| Base backend | Sparse algorithm | Resulting backend class | Metadata subtype | KV-cache manager | -|---|---|---|---|---| -| `TRTLLM` | `rocket` | `RocketTrtllmAttention` | `RocketTrtllmAttentionMetadata` | `RocketKVCacheManager` | -| `TRTLLM` | `dsa` | `DSATrtllmAttention` | `DSAtrtllmAttentionMetadata` | `DSACacheManager` | -| `TRTLLM` | `skip_softmax` | `TrtllmAttention` | `TrtllmAttentionMetadata` | standard `KVCacheManager` | -| `VANILLA` | `rocket` | `RocketVanillaAttention` | `RocketVanillaAttentionMetadata` | `RocketKVCacheManager` | -| `VANILLA` | `dsa` | unsupported | none | none | -| `VANILLA` | `skip_softmax` | unsupported | none | none | -| `FLASHINFER` | sparse variants | unsupported | none | none | +| Base backend | Sparse algorithm | Resulting backend class | KV-cache manager | +|---|---|---|---| +| `TRTLLM` | `rocket` | `RocketTrtllmAttention` | `RocketKVCacheManager` | +| `TRTLLM` | `dsa` | `DSATrtllmAttention` | `DSACacheManager` | +| `TRTLLM` | `skip_softmax` | `TrtllmAttention` | standard `KVCacheManager` | +| `VANILLA` | `rocket` | `RocketVanillaAttention` | `RocketKVCacheManager` | +| `VANILLA` | `dsa` | unsupported | — | +| `FLASHINFER` | sparse variants | unsupported | — | ### 2.3 Backend contract @@ -229,8 +164,6 @@ required operator or sparse path already exists. ### 2.4 Capability reference -The current coarse capability picture is: - | Backend family | Fused RoPE | Fused QKV input | MLA | |---|---|---|---| | `TrtllmAttention` | yes | yes | yes | @@ -252,77 +185,27 @@ It only applies to a subset of dense cases. If it does not apply, ### 3.1 Metadata families -All backend metadata types inherit from `AttentionMetadata`. +All backend metadata types inherit from `AttentionMetadata`. The base contract +includes sequence-length and request-level state, KV-cache manager and +parameters, runtime feature flags, optional sparse state, and optional +CUDA-graph buffer management. -#### 3.1.1 Base `AttentionMetadata` +**`TrtllmAttentionMetadata`** is the main metadata family. It adds paged-KV +block information, TRTLLM runtime state, chunked-prefill/speculative-decode/Helix +state, and MLA-specific state. If a source attention needs paged KV, chunked +prefill, FlashMLA, speculative decoding, or Helix-aware execution, the fit +question is mostly a `TrtllmAttentionMetadata` fit question. -The common contract includes: +**`VanillaAttentionMetadata`** is lighter — base metadata plus simple +cache-index information. Use it when the `Attention` module boundary fits but +the fused TRTLLM path is too restrictive. -- sequence-length and request-level runtime state -- KV-cache manager and KV-cache parameters -- runtime feature flags -- optional sparse-attention state -- optional CUDA-graph buffer management +**`FlashInferAttentionMetadata`** adds a planning-oriented contract with +workspace, page-table KV metadata, and prefill/decode wrapper state. -The base metadata also supports cross attention through `seq_lens_kv` and -cross sub-metadata. Backend support is not uniform: - -- `FlashInferAttention` has explicit cross-attention handling -- `TrtllmAttention` currently asserts that cross attention is not supported - -#### 3.1.2 `TrtllmAttentionMetadata` - -`TrtllmAttentionMetadata` is the main metadata family for the standard TRTLLM -path. - -It extends the base metadata with: - -- paged-KV block information -- request and sequence state for TRTLLM runtime execution -- runtime state for chunked prefill, speculative decode, and Helix -- MLA-specific runtime state when MLA is active - -If a source attention implementation needs paged KV, chunked prefill, -FlashMLA, speculative decoding, or Helix-aware execution, the fit question is -mostly a `TrtllmAttentionMetadata` fit question. - -#### 3.1.3 `VanillaAttentionMetadata` - -`VanillaAttentionMetadata` mainly prepares: - -- base attention metadata -- simple cache-index information for torch-side cache access - -The actual attention computation is mostly done in torch. Use it when the -current `Attention` module boundary still fits but the fused TRTLLM path is -too restrictive. - -#### 3.1.4 `FlashInferAttentionMetadata` - -`FlashInferAttentionMetadata` adds a planning-oriented runtime contract: - -- workspace and planning state -- page-table style KV metadata -- prefill and decode wrapper state - -This backend requires a planning step and a paged-table runtime shape. - -#### 3.1.5 Sparse metadata families - -Sparse backends extend the base metadata family rather than inventing an -unrelated attention interface. - -`DSAtrtllmAttentionMetadata` extends `TrtllmAttentionMetadata` with: - -- indexer-side sparse runtime state -- top-k and token-to-request routing state -- extra state for sparse context and generation flows - -`RocketTrtllmAttentionMetadata` and `RocketVanillaAttentionMetadata` add: - -- sparse-window and routing state -- sparse offsets and sequence state -- KT-cache related state +**Sparse metadata** families extend the base backend metadata with +sparse-specific runtime state (indexer buffers, routing state, side-cache +state). ### 3.2 KV-cache and decode-time semantics @@ -343,269 +226,105 @@ update pattern. #### 3.2.1 Common paged-KV model All current `_torch` backends are built on top of paged KV cache. +`KVCacheManager.get_buffers()` exposes a per-layer view of the primary pool: -At the cache-manager level, `KVCacheManager.get_buffers()` exposes a per-layer -view of the primary pool in two common layouts: - -- `NHD`: `[num_pages, kv_factor, tokens_per_block, num_kv_heads, head_dim]` -- `HND`: `[num_pages, kv_factor, num_kv_heads, tokens_per_block, head_dim]` - -For standard dense attention, `kv_factor = 2`, which means separate K and V -planes. - -For MLA-style cache, `kv_factor = 1`. The cache stores one latent-cache tensor -per token rather than separate K and V planes. +- For standard dense attention, `kv_factor = 2` (separate K and V planes). +- For MLA-style cache, `kv_factor = 1` (one latent-cache tensor per token). -The main differences across backends are: +The main differences across backends: -- which tensor view or block table the backend expects -- whether the backend writes cache internally or Python writes it explicitly -- whether extra side cache is required in addition to the main KV cache - -#### 3.2.2 `TRTLLM` backend - -At the C++ boundary, the main contract is a block-offset table plus pool -pointers, not only a reshaped cache tensor view. - -In the Python metadata path, this appears as fields such as: - -- `kv_cache_block_offsets` -- `host_kv_cache_pool_pointers` -- `host_kv_cache_pool_mapping` -- `cache_indirection` - -Read and write behavior: - -- decode and prefill cache updates are backend-managed -- dense cache writes go through backend ops such as `qkv_preprocessing` -- some sparse updates go through backend postprocessing -- the regular path reads cached KV through the block-table and pool-pointer - contract - -Python does not update the main KV cache with `index_copy_` in the regular -`TRTLLM` path. - -#### 3.2.3 `TRTLLM` internal `trtllm_gen` path - -`trtllm_gen.py` is not a separate top-level backend, but it has its own KV -view assumptions. - -It still starts from the same paged cache state, but it bridges the TRTLLM -block-offset format into the page-table shape expected by the FlashInfer -`trtllm_gen` kernels. - -In practice: - -- cache writes still begin from TRTLLM preprocessing -- the fast path converts K-side block offsets into shared page indices -- it reads cache again through `KVCacheManager.get_buffers(...)` - -This path is narrower than the main `TRTLLM` backend: - -- dense only -- no MLA -- no sparse attention -- no cross attention -- fused QKV only - -If `trtllm_gen` does not fit, that does not rule out the main `TRTLLM` -backend. - -#### 3.2.4 `FlashInfer` backend - -`FlashInfer` also uses paged KV cache, but its runtime contract is more -directly page-table oriented. It reads the cache through -`kv_cache_manager.get_buffers(...)` using the layout requested by -`metadata.kv_layout`. - -Read and write behavior: - -- Python explicitly appends current K and V into paged cache -- FlashInfer wrappers then read from that paged cache using page-table metadata - -#### 3.2.5 `VANILLA` backend - -`VANILLA` gets a paged cache tensor from `kv_cache_manager.get_buffers()` and -request block ids from `block_ids_per_seq`. - -Read and write behavior: +| Backend | Cache write | Cache read | Notes | +|---|---|---|---| +| `TRTLLM` | Backend-managed (C++ ops like `qkv_preprocessing`) | Block-table + pool pointers | Python does not call `index_copy_` in the regular path | +| `VANILLA` | Python-side (`index_copy_`) | Python-side slicing from same cache tensor | Most direct path for inspecting what is written to cache | +| `FlashInfer` | Python-side (explicit append) | Page-table metadata | Requires a planning step and its own page-table shape | -- Python writes K and V directly into cache -- Python slices the same cache tensor to rebuild key and value states -- sparse token filtering, if any, also happens around this Python-side path +#### 3.2.2 `TRTLLM` internal `trtllm_gen` path -#### 3.2.6 MLA cached-context semantics +`trtllm_gen.py` is not a separate backend, but it has its own KV view +assumptions. It bridges the TRTLLM block-offset format into the page-table +shape expected by its kernels. This path is narrower than the main `TRTLLM` +backend (dense only, no MLA, no sparse, fused QKV only). If `trtllm_gen` does +not fit, that does not rule out the main `TRTLLM` backend. -MLA adds a different cache shape and different decode-time assumptions. +#### 3.2.3 MLA cached-context semantics -The main difference is that MLA cached state is not regular dense K and V. -The paged cache stores latent-cache state, and backend ops handle: +MLA cached state is not regular dense K and V. The paged cache stores +latent-cache state, and backend ops handle: - appending latent cache into paged storage - applying RoPE as part of that flow - loading paged cached state back for attention use -MLA fit cannot be judged from attention math alone. The module and backend also -have to agree on: +MLA fit cannot be judged from attention math alone. The module and backend must +agree on latent-cache layout, paged-KV read/write paths, and cached/chunked +context behavior. The short-seq MHA path is only correct if cached-KV behavior +stays inside the top-level `forward_context()` dispatcher. -- latent-cache layout -- paged-KV read path -- paged-KV write path -- cached-context and chunked-context behavior +#### 3.2.4 Sparse side-cache semantics -The short-seq MHA path is only correct if cached-KV behavior stays inside the -top-level `forward_context()` dispatcher. +Sparse backends may add side caches beyond the main KV cache: -#### 3.2.7 Sparse side-cache semantics +- `skip_softmax` keeps the standard `KVCacheManager`, no side cache. +- `DSA` uses `DSACacheManager` and adds `indexer_k_cache` for sparse indexing. +- `Rocket` uses `RocketKVCacheManager` and adds `KT` cache for sparse routing. -Sparse backends may change more than the attention score path. They may also -add side caches. - -`skip_softmax` keeps the standard `KVCacheManager`. - -`DSA` uses `DSACacheManager`: - -- the main cache is still the paged MLA-style cache -- DSA also adds `indexer_k_cache` -- this side cache is used for sparse indexing and top-k related work - -`Rocket` uses `RocketKVCacheManager`: - -- the main cache still follows the base backend family -- Rocket also adds `KT` cache -- this side cache is used for sparse routing and sparse block selection - -When evaluating a new sparse attention implementation, check: - -- the main KV-cache contract -- the side-cache contract +When evaluating new sparse attention, check both the main KV-cache contract +and the side-cache contract. ## 4. Evaluating New Attention -### 4.1 First-pass fit against the current `TRTLLM` backend - -For a new model, first compare it against the current `TRTLLM` backend surface -in four parts. - -| What to compare | Current `TRTLLM` backend surface | -|---|---| -| Module-layer math around the backend call | `Attention` and `MLA` already handle QKV projection, fused or split QKV conversion, optional unfused RoPE, Q/K normalization hooks, output gating, LoRA, TP/CP handling, and Helix wrappers. | -| Backend-side execution | `TrtllmAttention` handles the standard dense path. It supports fused RoPE, fused QKV input, MLA mode, sliding-window style `attention_window_size`, mRoPE config, attention sinks, and sparse variants through registered sparse backends. | -| Metadata and runtime contract | The direct path expects `TrtllmAttentionMetadata` or one of its sparse subclasses. That means the new attention must still fit request-based metadata, paged-KV metadata, runtime feature flags, and any extra sparse or MLA buffers required by the path. | -| KV-cache semantics | The direct path assumes paged-KV ownership. Dense TRTLLM uses the standard `KVCacheManager`. Sparse paths can swap in `DSACacheManager` or `RocketKVCacheManager`. Backend choice also implies assumptions about KV-cache layout, decode-time reads, and inplace append or update behavior. If the new attention needs a different cache ownership model or different decode-time cache semantics, it is not a direct fit. | - -A practical check: - -1. Write down the new attention in four buckets: - - module math before or after the backend call - - backend features it needs - - metadata and runtime state it needs - - KV-cache ownership, layout assumptions, and update rules it needs -2. Compare each bucket against the current `TRTLLM` backend surface above. -3. Treat the first mismatch as the current blocker. +### 4.1 First-pass fit -For direct fit, ask these questions in order: +For a new model, compare it against the current stack in four parts: -1. Can `Attention` or `MLA` express the new math with module-side code only? -2. Can `TrtllmAttention.forward(...)` express the backend call shape that is - needed? - - fused or split QKV - - dense or sparse path - - RoPE or mRoPE handling - - windowed masking - - MLA or non-MLA -3. Can the runtime state be stored in `TrtllmAttentionMetadata` or a known - sparse metadata subclass? -4. Can the KV-cache behavior stay inside the current paged-KV and cache-manager - model? +1. **Module math**: can `Attention` or `MLA` express the new math with + module-side code only? +2. **Backend execution**: can `TrtllmAttention.forward(...)` handle the needed + call shape (fused/split QKV, dense/sparse, RoPE/mRoPE, MLA/non-MLA)? +3. **Metadata**: can the runtime state fit in `TrtllmAttentionMetadata` or a + known sparse subclass? +4. **KV-cache**: can the cache behavior stay inside the current paged-KV and + cache-manager model? -If the answer stays "yes" through all four questions, start with the current -`TRTLLM` backend and then check the runtime contract in Section 3. +If yes to all four, start with the `TRTLLM` backend. Treat the first mismatch +as the current blocker. -### 4.2 Detailed checklist +### 4.2 Checklist -#### 4.2.1 Module-level math +#### Module-level math -Check all math around the backend call: +- Q/K/V layout, fused or split QKV, MQA/GQA structure +- Q/K normalization, extra scaling, output gating +- Pre-backend and post-backend transforms -- Q/K/V layout -- fused or split QKV input -- MQA or GQA structure -- Q/K normalization -- extra QK scaling terms -- output gating -- pre-backend and post-backend transforms +#### Backend capability -The first question is whether this math can stay at the module layer by -overriding or extending `Attention` / `MLA`, without changing the outer runtime -contract. +- Which backend family can run the source behavior +- Fused RoPE, fused QKV, MLA, sparse, chunked-context support +- Do not use backend name alone as proof of support -#### 4.2.2 Backend capability mapping +#### Positional embedding and masking -Check which backend family can actually run the source behavior: +- RoPE applied outside vs fused, standard RoPE vs mRoPE +- Causal, full, sliding-window, or custom masks -- `TRTLLM` -- `VANILLA` -- `FLASHINFER` -- sparse variants on top of those families +#### Metadata and runtime contract -Then check capability assumptions explicitly: - -- fused RoPE -- fused QKV input -- MLA support -- sparse operator support -- chunked-context support - -Do not use backend name alone as proof of support. For a first-pass fit check, -use Section 4.1 before looking at other backends. - -#### 4.2.3 Positional embedding and mask contract - -Check: - -- whether RoPE is applied outside the backend or fused into it -- standard RoPE vs mRoPE -- causal, full, sliding-window, or custom masks -- any sink-token or mask-side special logic - -#### 4.2.4 Metadata and runtime contract - -Check which metadata subtype is required: - -- `TrtllmAttentionMetadata` -- `VanillaAttentionMetadata` -- `FlashInferAttentionMetadata` -- sparse metadata subclasses - -Then check whether the source behavior needs runtime state such as: - -- request-level state -- KV-cache manager handles and KV-cache parameters -- runtime feature flags -- planning or sparse buffers +- Which metadata subtype, what runtime state is needed - CUDA-graph assumptions -#### 4.2.5 KV-cache ownership and decode-time semantics - -Check: +#### KV-cache ownership -- how K/V are appended -- what KV-cache layout the backend expects -- how tokens are indexed in cache -- whether cached KV is revisited during context -- whether decode updates require inplace cache writes -- whether cache reuse or chunked prefill is required -- whether speculative decoding assumptions are involved -- whether sparse state is attached to cache ownership +- How K/V are appended, what layout, how indexed +- Cache reuse, chunked prefill, speculative decoding assumptions +- Sparse side-cache requirements ### 4.3 Bring-up order -Start with the `TRTLLM` backend when the new attention fits the existing -runtime contract or only needs limited changes. If that path is too costly for -initial bring-up or quick experiments, use `VANILLA` to validate the math and -outer module behavior first, then re-evaluate whether the implementation -should move back to `TRTLLM`. +Start with `TRTLLM` when the new attention fits or only needs limited changes. +Use `VANILLA` for quick bring-up or experiments when the module boundary fits +but the fused path is too costly to change initially. Working rules: @@ -624,36 +343,26 @@ Working rules: | `tensorrt_llm/_torch/attention_backend/interface.py` | Backend contract, base metadata, capability hooks | | `tensorrt_llm/_torch/attention_backend/utils.py` | Backend and sparse-backend selection | | `tensorrt_llm/_torch/attention_backend/trtllm.py` | TRTLLM backend and metadata | -| `tensorrt_llm/_torch/attention_backend/trtllm_gen.py` | Internal dense fast path used by `TrtllmAttention` on supported Blackwell configurations | +| `tensorrt_llm/_torch/attention_backend/trtllm_gen.py` | Internal dense fast path | | `tensorrt_llm/_torch/attention_backend/vanilla.py` | Torch fallback backend and metadata | | `tensorrt_llm/_torch/attention_backend/flashinfer.py` | FlashInfer backend and metadata | -| `tensorrt_llm/_torch/attention_backend/sparse/dsa.py` | DSA sparse backend, metadata, indexer, cache manager | -| `tensorrt_llm/_torch/attention_backend/sparse/kernel.py` | Triton helper kernels used by sparse attention implementations | -| `tensorrt_llm/_torch/attention_backend/sparse/rocket.py` | Rocket sparse backends, metadata, cache manager | -| `tensorrt_llm/_torch/attention_backend/sparse/utils.py` | Sparse backend and sparse cache-manager registration | -| `tensorrt_llm/_torch/models/modeling_deepseekv3.py` | DeepSeek weight loading, including `kv_b_proj` layout transforms | +| `tensorrt_llm/_torch/attention_backend/sparse/` | DSA, Rocket sparse backends, metadata, cache managers | ## 6. Testing Notes -- `mla.to(device)` does not move `mla.mqa.indexer` weights automatically. -- Copy indexer weights explicitly in A/B tests. -- Initialize `kv_b_proj` weights in loaded TRT-LLM layout, not HuggingFace - layout. - Test lite and non-lite MLA separately when changing projection logic. -- Test eager and compiled paths separately when changing DSA MLA dispatch or - custom-op behavior. +- Test eager and compiled paths separately when changing DSA MLA dispatch. - Test fresh context, cached context, chunked context, and generation separately. - Any dispatch change touching `forward_context()` needs chunked-context tests. -Useful tests: +Key test files: - `tests/unittest/_torch/attention/test_attention.py` - `tests/unittest/_torch/attention/test_attention_mla.py` - `tests/unittest/_torch/attention/test_vanilla_attention.py` - `tests/unittest/_torch/attention/test_flashinfer_attention.py` -- `tests/unittest/_torch/attention/sparse/test_short_seq_mha.py` -- `tests/unittest/_torch/attention/sparse/test_sparse_mla_forward.py` +- `tests/unittest/_torch/attention/sparse/` ## 7. Anti-Patterns From d6fb559bd988efa814ce2f17627a32781f6cbadd Mon Sep 17 00:00:00 2001 From: junq <22017000+QiJune@users.noreply.github.com> Date: Fri, 3 Apr 2026 17:44:48 +0800 Subject: [PATCH 3/7] fix comments Signed-off-by: junq <22017000+QiJune@users.noreply.github.com> --- AGENTS.md | 2 +- .../modules/ATTENTION_DEVELOPER_GUIDE.md | 32 +++++++++++-------- 2 files changed, 20 insertions(+), 14 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 64b83ce5972a..20019a4d8cc9 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -83,7 +83,7 @@ HuggingFace Model → LLM API → Executor (PyTorch/AutoDeploy/TensorRT) | `tensorrt_llm/executor/executor.py` | Execution abstraction (`GenerationExecutor`) | | `tensorrt_llm/models/automodel.py` | Auto-discovery and model registry | | `tensorrt_llm/_torch/models/` | PyTorch backend model implementations (distinct from `models/` used by TensorRT backend) | -| `tensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.md` | Attention, MLA, backend families, sparse backends, metadata contracts, and KV-cache behavior - **read before modifying `tensorrt_llm/_torch/modules/attention.py` or `_torch/attention_backend/`** | +| `tensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.md` | Attention, MLA, backend families, sparse backends, metadata contracts, and KV-cache behavior - **read before modifying `tensorrt_llm/_torch/modules/attention.py` or `tensorrt_llm/_torch/attention_backend/`** | | `tensorrt_llm/_torch/modules/fused_moe/MOE_DEVELOPER_GUIDE.md` | MoE architecture, backends, communication, development patterns — **read before modifying MoE code** | | `CODING_GUIDELINES.md` | C++ and Python coding standards (referenced throughout, must read before contributing) | diff --git a/tensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.md b/tensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.md index a3b380d9fdf2..b2e3ba9e9461 100644 --- a/tensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.md +++ b/tensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.md @@ -90,20 +90,25 @@ objects, and depends on metadata and KV-cache contract. It owns: (`is_lite == True`). In lite mode there is no separate Q low-rank compression stage. `is_lite` changes the projection structure, not just a small code path. -**DSA dispatch.** DSA-style MLA uses a multi-stage dispatch: projection and -attention are separated, context and generation paths are separated, and a -short-seq gate inside the context path can route to a dense MHA fallback. The -short-seq MHA path should stay inside `forward_context()`. Do not bypass the -dispatcher and call `forward_context_default()` directly — it only handles -fresh context, not cached-KV or chunked-context cases. +**Non-DSA MLA dispatch.** Non-DSA MLA uses absorption: Q, K, V are projected +through low-rank decomposition and then absorbed into the backend call. The +context path still dispatches through `forward_context()`, which handles fresh +context, cached-KV context, and chunked prefill as separate cases. Chunked +prefill routing depends on architecture (SM90 vs SM100+). Do not assume all +context goes through the same handler. + +**DSA dispatch.** DSA-style MLA adds a further split: projection and attention +are separated, context and generation paths are separated, and a short-seq +gate inside the context path can route to a dense MHA fallback. + +**For both DSA and non-DSA MLA:** the short-seq MHA path should stay inside +`forward_context()`. Do not bypass the dispatcher and call +`forward_context_default()` directly — it only handles fresh context, not +cached-KV or chunked-context cases. **Practical notes:** - `self.mha` being present does not mean DSA is disabled. -- `_should_use_short_mha()` uses `max_ctx_kv_len` when available, not just the - new-token count. -- Under `torch.compile`, `_should_use_short_mha()` returns `False`, so the - split DSA path is always used. - Helix CP wraps the attention body with allgather/output-projection helpers. ## 2. Backend Layer Reference @@ -176,10 +181,11 @@ metadata and cache behavior. ### 2.5 `TRTLLM` internal kernel paths `TrtllmAttention` can dispatch to `trtllm_gen.py` for supported dense cases. -That is an internal fast path, not a separate top-level backend selection. +That is an internal fast path, not a separate top-level backend selection. It +is disabled by default and gated by `TRTLLM_ENABLE_TRTLLM_GEN_ATTENTION`. -It only applies to a subset of dense cases. If it does not apply, -`TrtllmAttention` stays on its regular runtime path. +It only applies to a narrow subset of dense cases. If it does not apply or is +not enabled, `TrtllmAttention` stays on its regular runtime path. ## 3. Runtime Contract Reference From d19c1213848d90f2e67a4cccc5f60bb80ee04b8b Mon Sep 17 00:00:00 2001 From: junq <22017000+QiJune@users.noreply.github.com> Date: Fri, 3 Apr 2026 18:17:50 +0800 Subject: [PATCH 4/7] simplify Signed-off-by: junq <22017000+QiJune@users.noreply.github.com> --- .../modules/ATTENTION_DEVELOPER_GUIDE.md | 60 +++++++------------ 1 file changed, 23 insertions(+), 37 deletions(-) diff --git a/tensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.md b/tensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.md index b2e3ba9e9461..d3403acddf40 100644 --- a/tensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.md +++ b/tensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.md @@ -101,10 +101,10 @@ context goes through the same handler. are separated, context and generation paths are separated, and a short-seq gate inside the context path can route to a dense MHA fallback. -**For both DSA and non-DSA MLA:** the short-seq MHA path should stay inside -`forward_context()`. Do not bypass the dispatcher and call -`forward_context_default()` directly — it only handles fresh context, not -cached-KV or chunked-context cases. +For both DSA and non-DSA MLA, keep context routing inside `forward_context()`. +For DSA, the short-seq MHA fallback should also stay inside that dispatcher. Do +not bypass it and call `forward_context_default()` directly; that handler only +covers fresh context, not cached-KV or chunked-context cases. **Practical notes:** @@ -115,16 +115,8 @@ cached-KV or chunked-context cases. ### 2.1 Backend selection -Backends are created by: - -- `get_attention_backend()` -- `create_attention()` - -The backend is chosen from: - -- `config.attn_backend` -- optional `sparse_attention_config` -- optional MLA parameters +The backend is chosen from `config.attn_backend`, optional +`sparse_attention_config`, and optional MLA parameters. Base backend families: @@ -178,15 +170,6 @@ required operator or sparse path already exists. Sparse subclasses inherit the base backend family and then add sparse-specific metadata and cache behavior. -### 2.5 `TRTLLM` internal kernel paths - -`TrtllmAttention` can dispatch to `trtllm_gen.py` for supported dense cases. -That is an internal fast path, not a separate top-level backend selection. It -is disabled by default and gated by `TRTLLM_ENABLE_TRTLLM_GEN_ATTENTION`. - -It only applies to a narrow subset of dense cases. If it does not apply or is -not enabled, `TrtllmAttention` stays on its regular runtime path. - ## 3. Runtime Contract Reference ### 3.1 Metadata families @@ -231,27 +214,30 @@ update pattern. #### 3.2.1 Common paged-KV model -All current `_torch` backends are built on top of paged KV cache. -`KVCacheManager.get_buffers()` exposes a per-layer view of the primary pool: +When KV cache is enabled, all current `_torch` backends use paged KV cache. +`VanillaAttention` also has a separate no-KV-cache path for models that do not +use cache. `KVCacheManager.get_buffers()` exposes a per-layer view of the +primary pool: - For standard dense attention, `kv_factor = 2` (separate K and V planes). - For MLA-style cache, `kv_factor = 1` (one latent-cache tensor per token). The main differences across backends: -| Backend | Cache write | Cache read | Notes | -|---|---|---|---| -| `TRTLLM` | Backend-managed (C++ ops like `qkv_preprocessing`) | Block-table + pool pointers | Python does not call `index_copy_` in the regular path | -| `VANILLA` | Python-side (`index_copy_`) | Python-side slicing from same cache tensor | Most direct path for inspecting what is written to cache | -| `FlashInfer` | Python-side (explicit append) | Page-table metadata | Requires a planning step and its own page-table shape | +| Backend | Cache write | Cache read | +|---|---|---| +| `TRTLLM` | Backend-managed (C++ ops) | Block-table + pool pointers | +| `VANILLA` | Python-side | Python-side slicing | +| `FlashInfer` | Python-side (explicit append) | Page-table metadata | #### 3.2.2 `TRTLLM` internal `trtllm_gen` path -`trtllm_gen.py` is not a separate backend, but it has its own KV view -assumptions. It bridges the TRTLLM block-offset format into the page-table -shape expected by its kernels. This path is narrower than the main `TRTLLM` -backend (dense only, no MLA, no sparse, fused QKV only). If `trtllm_gen` does -not fit, that does not rule out the main `TRTLLM` backend. +`trtllm_gen.py` integrates trtllm-gen kernels from FlashInfer into the +`TRTLLM` backend. It is not a separate backend — it is an internal fast path +disabled by default (`TRTLLM_ENABLE_TRTLLM_GEN_ATTENTION`). It bridges the +TRTLLM block-offset format into the page-table shape expected by those kernels. +This path is under active development and its supported scope is expanding. If +it does not apply, `TrtllmAttention` stays on its regular runtime path. #### 3.2.3 MLA cached-context semantics @@ -286,8 +272,8 @@ For a new model, compare it against the current stack in four parts: 1. **Module math**: can `Attention` or `MLA` express the new math with module-side code only? -2. **Backend execution**: can `TrtllmAttention.forward(...)` handle the needed - call shape (fused/split QKV, dense/sparse, RoPE/mRoPE, MLA/non-MLA)? +2. **Backend execution**: can the current `TRTLLM` backend family handle the + needed call shape (fused/split QKV, dense/sparse, RoPE/mRoPE, MLA/non-MLA)? 3. **Metadata**: can the runtime state fit in `TrtllmAttentionMetadata` or a known sparse subclass? 4. **KV-cache**: can the cache behavior stay inside the current paged-KV and From cf251e96fbb8606001365377cfa6660f4c143b9b Mon Sep 17 00:00:00 2001 From: junq <22017000+QiJune@users.noreply.github.com> Date: Tue, 7 Apr 2026 10:05:01 +0800 Subject: [PATCH 5/7] update Signed-off-by: junq <22017000+QiJune@users.noreply.github.com> --- .../modules/ATTENTION_DEVELOPER_GUIDE.md | 143 ++++++++++-------- 1 file changed, 79 insertions(+), 64 deletions(-) diff --git a/tensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.md b/tensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.md index d3403acddf40..2f78cc9c5214 100644 --- a/tensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.md +++ b/tensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.md @@ -9,11 +9,26 @@ This guide covers the TRT-LLM PyTorch attention stack: - `tensorrt_llm/_torch/attention_backend/sparse/` Use it when modifying the current implementation or adding a new model's -attention behavior. It covers standard `Attention`, MLA, dense backends, and -sparse backends. It does not cover +attention behavior. It covers standard `Attention`, Multi-head Latent +Attention (MLA), dense backends, and sparse backends. It does not cover `tensorrt_llm/_torch/attention_backend/star_flashinfer.py`, which is planned for deprecation. +## Glossary + +| Acronym | Meaning | +|---|---| +| MLA | Multi-head Latent Attention | +| DSA | DeepSeek Sparse Attention | +| MHA | Multi-Head Attention | +| MQA | Multi-Query Attention | +| GQA | Grouped-Query Attention | +| RoPE | Rotary Position Embedding | +| mRoPE | Multimodal Rotary Position Embedding | +| TP | Tensor Parallelism | +| CP | Context Parallelism | +| KV | Key/Value | + ## How to Read the Stack Attention in TRT-LLM is split across four layers: @@ -21,7 +36,7 @@ Attention in TRT-LLM is split across four layers: 1. module wrapper (`Attention` or `MLA`) 2. backend class selected by `config.attn_backend` 3. metadata subtype and runtime buffers -4. KV-cache manager and decode-time cache semantics +4. key/value (KV) cache manager and decode-time cache semantics Keep these four questions separate: @@ -42,9 +57,9 @@ appends, and reuses KV cache, especially during decode. `Attention` is not just the backend call. It owns the logic around the backend: - QKV projection and output projection -- TP/CP reshaping and mapping setup +- tensor parallelism (TP) / context parallelism (CP) reshaping and mapping setup - fused or split QKV handling -- optional unfused RoPE +- optional unfused Rotary Position Embedding (RoPE) - optional output gating - optional LoRA injection - passing masks, sinks, and metadata into the backend @@ -75,41 +90,44 @@ without changing the outer runtime contract. ### 1.2 `MLA`: a separate module on top of the same backend system -`MLA` is a separate module in `attention.py`. It keeps module-level -projections and tensor transforms in `MLA`, delegates core execution to backend -objects, and depends on metadata and KV-cache contract. It owns: +`MLA` (Multi-head Latent Attention) is a separate module in `attention.py`. +Like `Attention`, it keeps module-level projection logic in the module, +delegates core execution to a backend object, and depends on metadata and +KV-cache contract. At a high level, it owns: - low-rank Q decomposition - low-rank KV decomposition - absorbed MLA path -- DSA-specific dispatch -- short-seq MHA routing - MLA-specific RoPE and latent-cache flow +- integration points for sparse attention paths `MLA` has two projection layouts: non-lite (`is_lite == False`) and lite (`is_lite == True`). In lite mode there is no separate Q low-rank compression stage. `is_lite` changes the projection structure, not just a small code path. -**Non-DSA MLA dispatch.** Non-DSA MLA uses absorption: Q, K, V are projected -through low-rank decomposition and then absorbed into the backend call. The -context path still dispatches through `forward_context()`, which handles fresh -context, cached-KV context, and chunked prefill as separate cases. Chunked -prefill routing depends on architecture (SM90 vs SM100+). Do not assume all -context goes through the same handler. +Dense MLA and current sparse MLA variants still follow the same four layers +described above: module-layer logic, backend execution, metadata/runtime +contract, and KV-cache semantics. Sparse-specific routing may currently pass +through `MLA`, but that should be treated as an implementation detail rather +than a stable design boundary. -**DSA dispatch.** DSA-style MLA adds a further split: projection and attention -are separated, context and generation paths are separated, and a short-seq -gate inside the context path can route to a dense MHA fallback. +When reviewing an MLA-related task, the main questions are: -For both DSA and non-DSA MLA, keep context routing inside `forward_context()`. -For DSA, the short-seq MHA fallback should also stay inside that dispatcher. Do -not bypass it and call `forward_context_default()` directly; that handler only -covers fresh context, not cached-KV or chunked-context cases. +1. Does the new behavior fit the current MLA projection structure, including + the lite vs non-lite split? +2. Can the core execution stay on an existing backend family? +3. Can the runtime state stay in an existing metadata family? +4. Can the cache behavior stay within the current latent-cache and paged-KV + contract? -**Practical notes:** +If the answer to those questions is yes, the task usually stays within the +existing MLA stack. This applies whether the task is adding behavior, +improving performance, or checking whether a new model's attention can be +supported by the current MLA path. -- `self.mha` being present does not mean DSA is disabled. -- Helix CP wraps the attention body with allgather/output-projection helpers. +If the task depends on current sparse helper-level control flow, read +`attention.py` and the relevant sparse backend code directly. This guide is a +high-level reference, not the source of truth for dispatch details. ## 2. Backend Layer Reference @@ -141,6 +159,7 @@ Current sparse registrations: | `TRTLLM` | `skip_softmax` | `TrtllmAttention` | standard `KVCacheManager` | | `VANILLA` | `rocket` | `RocketVanillaAttention` | `RocketKVCacheManager` | | `VANILLA` | `dsa` | unsupported | — | +| `VANILLA` | `skip_softmax` | unsupported | — | | `FLASHINFER` | sparse variants | unsupported | — | ### 2.3 Backend contract @@ -228,16 +247,15 @@ The main differences across backends: |---|---|---| | `TRTLLM` | Backend-managed (C++ ops) | Block-table + pool pointers | | `VANILLA` | Python-side | Python-side slicing | -| `FlashInfer` | Python-side (explicit append) | Page-table metadata | +| `FLASHINFER` | Python-side (explicit append) | Page-table metadata | #### 3.2.2 `TRTLLM` internal `trtllm_gen` path `trtllm_gen.py` integrates trtllm-gen kernels from FlashInfer into the -`TRTLLM` backend. It is not a separate backend — it is an internal fast path +`TRTLLM` backend. It is not a separate backend. It is an internal fast path disabled by default (`TRTLLM_ENABLE_TRTLLM_GEN_ATTENTION`). It bridges the TRTLLM block-offset format into the page-table shape expected by those kernels. -This path is under active development and its supported scope is expanding. If -it does not apply, `TrtllmAttention` stays on its regular runtime path. +If it does not apply, `TrtllmAttention` stays on its regular runtime path. #### 3.2.3 MLA cached-context semantics @@ -251,7 +269,7 @@ latent-cache state, and backend ops handle: MLA fit cannot be judged from attention math alone. The module and backend must agree on latent-cache layout, paged-KV read/write paths, and cached/chunked context behavior. The short-seq MHA path is only correct if cached-KV behavior -stays inside the top-level `forward_context()` dispatcher. +stays inside the module's top-level context dispatch rather than bypassing it. #### 3.2.4 Sparse side-cache semantics @@ -268,51 +286,48 @@ and the side-cache contract. ### 4.1 First-pass fit -For a new model, compare it against the current stack in four parts: +When evaluating a new attention path, compare it against the current stack in +the same four layers used throughout this guide: -1. **Module math**: can `Attention` or `MLA` express the new math with - module-side code only? -2. **Backend execution**: can the current `TRTLLM` backend family handle the - needed call shape (fused/split QKV, dense/sparse, RoPE/mRoPE, MLA/non-MLA)? -3. **Metadata**: can the runtime state fit in `TrtllmAttentionMetadata` or a - known sparse subclass? -4. **KV-cache**: can the cache behavior stay inside the current paged-KV and - cache-manager model? +1. **Module layer**: can `Attention` or `MLA` express the required math with + module-side changes only? +2. **Backend layer**: can the current `TRTLLM` backend family handle the + required execution shape? +3. **Runtime contract**: can the state fit in an existing metadata family? +4. **KV-cache semantics**: can the cache behavior stay within the current + paged-KV and cache-manager model? If yes to all four, start with the `TRTLLM` backend. Treat the first mismatch as the current blocker. -### 4.2 Checklist - -#### Module-level math - -- Q/K/V layout, fused or split QKV, MQA/GQA structure -- Q/K normalization, extra scaling, output gating -- Pre-backend and post-backend transforms - -#### Backend capability - -- Which backend family can run the source behavior -- Fused RoPE, fused QKV, MLA, sparse, chunked-context support -- Do not use backend name alone as proof of support +### 4.2 What to check -#### Positional embedding and masking +Check the following at a high level: -- RoPE applied outside vs fused, standard RoPE vs mRoPE -- Causal, full, sliding-window, or custom masks +- **Module layer** + Q/K/V layout, fused or split QKV, MQA/GQA structure, Q/K normalization, + extra scaling, output gating, and pre-backend or post-backend transforms. -#### Metadata and runtime contract +- **Backend layer** + Which backend family can run the source behavior, and whether it needs fused + RoPE, fused QKV, MLA, sparse, or chunked-context support. Do not use backend + name alone as proof of support. -- Which metadata subtype, what runtime state is needed -- CUDA-graph assumptions +- **Positional embedding and masking** + Whether RoPE is applied outside or fused, whether the path needs mRoPE, and + whether masking fits the current causal, full, sliding-window, or custom + paths. -#### KV-cache ownership +- **Runtime contract** + Which metadata subtype is needed, what runtime state must be carried, and + whether the path depends on CUDA-graph assumptions. -- How K/V are appended, what layout, how indexed -- Cache reuse, chunked prefill, speculative decoding assumptions -- Sparse side-cache requirements +- **KV-cache semantics** + How K/V are appended, what layout is assumed, how cached state is indexed and + reused, whether chunked prefill or speculative decoding matters, and whether + sparse side caches are required. -### 4.3 Bring-up order +### 4.3 Default bring-up order Start with `TRTLLM` when the new attention fits or only needs limited changes. Use `VANILLA` for quick bring-up or experiments when the module boundary fits From 62a4802c76f0f5936f325f1296c8e5b8460fb93a Mon Sep 17 00:00:00 2001 From: junq <22017000+QiJune@users.noreply.github.com> Date: Tue, 7 Apr 2026 10:34:55 +0800 Subject: [PATCH 6/7] polish Signed-off-by: junq <22017000+QiJune@users.noreply.github.com> --- .../modules/ATTENTION_DEVELOPER_GUIDE.md | 43 ++++++------------- 1 file changed, 14 insertions(+), 29 deletions(-) diff --git a/tensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.md b/tensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.md index 2f78cc9c5214..faecfe6a6b8f 100644 --- a/tensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.md +++ b/tensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.md @@ -105,29 +105,17 @@ KV-cache contract. At a high level, it owns: (`is_lite == True`). In lite mode there is no separate Q low-rank compression stage. `is_lite` changes the projection structure, not just a small code path. -Dense MLA and current sparse MLA variants still follow the same four layers -described above: module-layer logic, backend execution, metadata/runtime -contract, and KV-cache semantics. Sparse-specific routing may currently pass -through `MLA`, but that should be treated as an implementation detail rather -than a stable design boundary. - -When reviewing an MLA-related task, the main questions are: - -1. Does the new behavior fit the current MLA projection structure, including - the lite vs non-lite split? -2. Can the core execution stay on an existing backend family? -3. Can the runtime state stay in an existing metadata family? -4. Can the cache behavior stay within the current latent-cache and paged-KV - contract? - -If the answer to those questions is yes, the task usually stays within the -existing MLA stack. This applies whether the task is adding behavior, -improving performance, or checking whether a new model's attention can be -supported by the current MLA path. - -If the task depends on current sparse helper-level control flow, read -`attention.py` and the relevant sparse backend code directly. This guide is a -high-level reference, not the source of truth for dispatch details. +Dense MLA and current sparse MLA variants still use the same module/backend/ +metadata/KV-cache split described above. Sparse-specific routing may currently +pass through `MLA`, but that should be treated as an implementation detail +rather than a stable design boundary. + +For MLA-related tasks, first check whether the work fits the current +projection structure, can stay on an existing backend and metadata family, and +can preserve the current latent-cache / paged-KV contract. If it can, the +task usually stays within the existing MLA stack. If it depends on sparse +helper-level control flow, read `attention.py` and the relevant sparse +backend code directly. ## 2. Backend Layer Reference @@ -268,8 +256,7 @@ latent-cache state, and backend ops handle: MLA fit cannot be judged from attention math alone. The module and backend must agree on latent-cache layout, paged-KV read/write paths, and cached/chunked -context behavior. The short-seq MHA path is only correct if cached-KV behavior -stays inside the module's top-level context dispatch rather than bypassing it. +context behavior. #### 3.2.4 Sparse side-cache semantics @@ -286,8 +273,8 @@ and the side-cache contract. ### 4.1 First-pass fit -When evaluating a new attention path, compare it against the current stack in -the same four layers used throughout this guide: +When evaluating a new attention path, compare it against the same four layers +used throughout this guide: 1. **Module layer**: can `Attention` or `MLA` express the required math with module-side changes only? @@ -302,8 +289,6 @@ as the current blocker. ### 4.2 What to check -Check the following at a high level: - - **Module layer** Q/K/V layout, fused or split QKV, MQA/GQA structure, Q/K normalization, extra scaling, output gating, and pre-backend or post-backend transforms. From 11d72f6e8c08a914f1c4603af03de6e62baebac8 Mon Sep 17 00:00:00 2001 From: junq <22017000+QiJune@users.noreply.github.com> Date: Tue, 7 Apr 2026 10:47:20 +0800 Subject: [PATCH 7/7] remove more implementation details Signed-off-by: junq <22017000+QiJune@users.noreply.github.com> --- .../modules/ATTENTION_DEVELOPER_GUIDE.md | 46 +++++++------------ 1 file changed, 16 insertions(+), 30 deletions(-) diff --git a/tensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.md b/tensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.md index faecfe6a6b8f..8dc2ee193ab5 100644 --- a/tensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.md +++ b/tensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.md @@ -138,17 +138,8 @@ Sparse attention is not selected by a separate top-level module. It is resolved through `sparse_attention_config` on top of a base backend family. Sparse selection can change the backend class, metadata subtype, and KV-cache manager. -Current sparse registrations: - -| Base backend | Sparse algorithm | Resulting backend class | KV-cache manager | -|---|---|---|---| -| `TRTLLM` | `rocket` | `RocketTrtllmAttention` | `RocketKVCacheManager` | -| `TRTLLM` | `dsa` | `DSATrtllmAttention` | `DSACacheManager` | -| `TRTLLM` | `skip_softmax` | `TrtllmAttention` | standard `KVCacheManager` | -| `VANILLA` | `rocket` | `RocketVanillaAttention` | `RocketKVCacheManager` | -| `VANILLA` | `dsa` | unsupported | — | -| `VANILLA` | `skip_softmax` | unsupported | — | -| `FLASHINFER` | sparse variants | unsupported | — | +Sparse registrations are defined in `attention_backend/sparse/utils.py`. Check +that file for the current supported combinations, as they may change over time. ### 2.3 Backend contract @@ -168,11 +159,10 @@ required operator or sparse path already exists. ### 2.4 Capability reference -| Backend family | Fused RoPE | Fused QKV input | MLA | -|---|---|---|---| -| `TrtllmAttention` | yes | yes | yes | -| `VanillaAttention` | no | no | no | -| `FlashInferAttention` | no | no | no | +Check each backend's capability hooks (`support_fused_rope()`, +`support_fused_qkv()`, `support_mla()`) directly in the code. `TrtllmAttention` +currently supports all three; other backends may not. These capabilities can +change over time. Sparse subclasses inherit the base backend family and then add sparse-specific metadata and cache behavior. @@ -248,26 +238,23 @@ If it does not apply, `TrtllmAttention` stays on its regular runtime path. #### 3.2.3 MLA cached-context semantics MLA cached state is not regular dense K and V. The paged cache stores -latent-cache state, and backend ops handle: - -- appending latent cache into paged storage -- applying RoPE as part of that flow -- loading paged cached state back for attention use +latent-cache state rather than separate K and V planes. Backend ops handle +appending, RoPE application, and loading cached state for attention use. MLA fit cannot be judged from attention math alone. The module and backend must agree on latent-cache layout, paged-KV read/write paths, and cached/chunked -context behavior. +context behavior. Read the MLA section of `attention.py` and the relevant +backend code for the current implementation details. #### 3.2.4 Sparse side-cache semantics -Sparse backends may add side caches beyond the main KV cache: - -- `skip_softmax` keeps the standard `KVCacheManager`, no side cache. -- `DSA` uses `DSACacheManager` and adds `indexer_k_cache` for sparse indexing. -- `Rocket` uses `RocketKVCacheManager` and adds `KT` cache for sparse routing. +Sparse backends may add side caches beyond the main KV cache. Some sparse +algorithms keep the standard cache manager; others replace it with a +sparse-aware cache manager that adds side caches for indexing or routing. When evaluating new sparse attention, check both the main KV-cache contract -and the side-cache contract. +and the side-cache contract. See `attention_backend/sparse/` for the current +sparse cache managers and their side-cache structures. ## 4. Evaluating New Attention @@ -361,6 +348,5 @@ Key test files: - Do not treat attention work as "math only". - Do not treat backend choice as independent from metadata choice. - Do not treat KV-cache semantics as a small implementation detail. -- Do not call `forward_context_default()` directly for chunked MLA context. +- Do not bypass MLA's context dispatcher for chunked or cached-KV cases. - Do not duplicate RoPE handling before checking the fused path. -- Do not assume `self.mha is not None` means DSA is disabled.