Skip to content

[None][feat] Attach feature and sparse metadata state as extensions #21194

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
yihwang-nv wants to merge 3 commits into
base: main
Choose a base branch
from
+7,845 −6,101
Draft

[None][feat] Attach feature and sparse metadata state as extensions #21194

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
f069dd1
[None][feat] Attach feature and sparse metadata state as extensions
yihwang-nv Oct 26, 2026
840cac8
[None][feat] Resolve the metadata-extension port onto newer main
yihwang-nv Oct 28, 2026
976948a
[None][test] Use extension-aware metadata stubs in ported tests
yihwang-nv Oct 28, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion docs/source/developer-guide/sparse-attention-development-guide.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -146,7 +146,7 @@ Dynamic generation-phase KV eviction is tracked as future work.
| Public config / lowering | `CSA2SparseAttentionConfig(algorithm="csa2")` → `CSA2Params`, `CSA2MetadataParams`; checkpoint-owned `CSA2Layout` |
| Registry | TRTLLM sparse backend → `CSA2TrtllmAttention`; cache manager → `CSA2CacheManager` (V2) |
| Module | Typed MLA adapter in `sparse/csa2/module.py`; compression, shared KV/index source routing and per-phase prediction |
| Metadata | `CSA2TrtllmMetadata`; SWA slots/visibility, owner page tables, compressor inputs, graph-stable index descriptors and query tiles |
| Metadata | `CSA2Metadata` extension; SWA slots/visibility, owner page tables, compressor inputs, graph-stable index descriptors and query tiles |
| Compute | SM100-family trtllm-gen; SM90 FlashMLA BF16 and SM120/121 FlashInfer BF16 helpers |

CSA2 GLOBAL pages couple 288-byte main and 68-byte index records under one
Expand Down
3 changes: 2 additions & 1 deletion ...h/_experimental/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_10x__default/modeling.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -83,6 +83,7 @@
from tensorrt_llm._torch._experimental.modeling_v2.catalog.quantization.fp4_quantize import (
fp4_quantize,
)
from tensorrt_llm._torch.attention.backends.features.spec_decoding import SpecDecodingMetadata
from tensorrt_llm._torch.attention.backends.interface import AttentionMetadata
from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata
from tensorrt_llm._torch.model_config import ModelConfig
Expand Down Expand Up @@ -131,7 +132,7 @@ def _build_step_args(md: TrtllmAttentionMetadata) -> dict:
num_ctx_tokens=md.num_ctx_tokens,
trtllm_gen_jit_warmup=md.trtllm_gen_jit_warmup,
max_num_sequences=md.max_num_sequences,
force_prepare_spec_dec_tree_mask=md.force_prepare_spec_dec_tree_mask,
force_prepare_spec_dec_tree_mask=md.ext(SpecDecodingMetadata).force_prepare_tree_mask,
)


Expand Down
54 changes: 30 additions & 24 deletions ...orch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/modeling.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,14 +22,16 @@
import torch
from transformers import PretrainedConfig

from tensorrt_llm._torch.attention.backends.features.spec_decoding import SpecDecodingMetadata
from tensorrt_llm._torch.attention.backends.interface import (
AttentionForwardArgs,
AttentionInputType,
AttentionMetadata,
)
from tensorrt_llm._torch.attention.backends.sparse.dsa import DSAMetadata
from tensorrt_llm._torch.attention.backends.sparse.dsa.module import forward_generation_sparse_attn
from tensorrt_llm._torch.attention.backends.sparse.dsa.params import DSABackendForwardArgs
from tensorrt_llm._torch.attention.backends.sparse.params import MTPIndexShareMetadata
from tensorrt_llm._torch.attention.backends.sparse.params import MTPIndexShareExtension
from tensorrt_llm._torch.distributed import AllReduceFusionOp, AllReduceParams
from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl
from tensorrt_llm._torch.model_config import ModelConfig
Expand All @@ -43,7 +45,6 @@
to_trtllm_gen_act_type,
)
from tensorrt_llm._torch.speculative import MTPEagleWorker, SpecMetadata
from tensorrt_llm._torch.speculative.eagle3 import _reset_mtp_index_share
from tensorrt_llm._torch.speculative.interface import (
_FORCE_ACCEPT_RNG_COUNTER_STRIDE,
_FORCE_ACCEPT_RNG_POOL_SIZE,
Expand Down Expand Up @@ -627,6 +628,7 @@ def draft_first_step(
among the verify rows; and the selection the later steps attend to, as KV pool rows,
or None where they attend to every position (a short context) or no step follows.
"""
dsa = attn_metadata.ext(DSAMetadata)
draft = self._draft
layer = draft.layer
rows = hidden_states.shape[0]
Expand All @@ -651,7 +653,7 @@ def draft_first_step(
)
attn_out = self._attention(layer, position_ids, hidden_states, attn_metadata, True)
pool_rows = None
if draft_len > 1 and not attn_metadata.skip_indexer_for_gen_reqs:
if draft_len > 1 and not dsa.skip_indexer_for_gen_reqs:
pool_rows = self._stash_selection(attn_metadata, num_accepted, draft_len + 1)
hidden_states, residual = self._post_attention_hop(layer, attn_out, residual)
gather_ids, hidden_states, residual = glm_gather_last_accepted_rows(
Expand All @@ -675,11 +677,12 @@ def draft_next_step(
indexer, to every position up to their own, the dense selection that metadata lists
(refreshed for this step by the loop). Returns the MTP layer's output rows.
"""
dsa = attn_metadata.ext(DSAMetadata)
draft = self._draft
layer = draft.layer
hidden_states, residual = eh_proj_output
if attn_metadata.skip_indexer_for_gen_reqs:
positions = attn_metadata.topk_indices_buffer[: hidden_states.shape[0]]
if dsa.skip_indexer_for_gen_reqs:
positions = dsa.topk_indices_buffer[: hidden_states.shape[0]]
attn_out = self._attention(
layer, None, hidden_states, attn_metadata, True, positions=positions
)
Expand Down Expand Up @@ -854,6 +857,7 @@ def _stash_selection(self, attn_metadata, num_accepted: torch.Tensor, next_n: in
stream after draft step 0's attention: each request's top-k of its last accepted row
(the MTP index share), as KV pool rows of the MTP layer. The first later step joins
it (``selection_ready``)."""
dsa = attn_metadata.ext(DSAMetadata)
draft = self._draft
manager = attn_metadata.kv_cache_manager
requests = num_accepted.shape[0]
Expand All @@ -872,7 +876,7 @@ def _stash_selection(self, attn_metadata, num_accepted: torch.Tensor, next_n: in
)
pool_rows = convert_req_index_to_global(
draft.request_rows[:requests],
attn_metadata.block_table[:requests],
dsa.block_table[:requests],
stashed,
tokens_per_block,
self.index_topk,
Expand Down Expand Up @@ -938,6 +942,7 @@ def _shard_indexer_keys(self, attn_metadata) -> bool:
Applies to one generation request whose shard fits the top-k kernel; returns whether it
did, and otherwise every indexer layer scores the whole context with the built-in
indexer."""
dsa = attn_metadata.ext(DSAMetadata)
manager = attn_metadata.kv_cache_manager
if self._hist_workspace is None or attn_metadata.num_seqs != 1:
return False
Expand Down Expand Up @@ -975,8 +980,8 @@ def _shard_indexer_keys(self, attn_metadata) -> bool:
)
buffers = self._hist_buffers
hist_shard_meta(
attn_metadata.gen_indexer_kv_lens_cuda_runtime,
attn_metadata.indexer_k_cache_block_offsets[0],
dsa.gen_indexer_kv_lens_cuda_runtime,
dsa.indexer_k_cache_block_offsets[0],
buffers["shard_ctx"],
buffers["shard_block_table"],
self.mapping.tp_size,
Expand Down Expand Up @@ -1004,6 +1009,7 @@ def _attention(
On an MTP draft step after the first (``sharded``) the rows attend to a given
selection and the indexer does not run: ``pool_rows``, KV pool rows, or
``positions``, sequence positions the flash decode maps to pool rows itself."""
dsa = attn_metadata.ext(DSAMetadata)
mla = layer.self_attn
mqa = mla.mqa
indexer = mqa.indexer
Expand Down Expand Up @@ -1138,7 +1144,7 @@ def kv_half_then_q_b():
output,
seq_len_q,
self.index_topk,
attn_metadata.block_table[: attn_metadata.num_seqs],
dsa.block_table[: attn_metadata.num_seqs],
manager.tokens_per_block,
page_index_scale,
layer_offset,
Expand Down Expand Up @@ -1190,6 +1196,7 @@ def kv_half_then_q_b():
def _select_rows(self, indexer, kv_a, latent_end, q_a, position_ids, attn_metadata, rows, mla):
"""The indexer of a layer that runs one, on this rank's key shard: writes the exact
top-2048 cache positions of every row into the shared selection buffer."""
dsa = attn_metadata.ext(DSAMetadata)
buffers = self._hist_buffers
manager = attn_metadata.kv_cache_manager
rope = indexer.rotary_emb
Expand Down Expand Up @@ -1217,8 +1224,8 @@ def _select_rows(self, indexer, kv_a, latent_end, q_a, position_ids, attn_metada
buffers["k_fp8"][:rows],
buffers["k_scale"][:rows],
k_cache,
attn_metadata.slot_mapping_fp8[:rows],
attn_metadata.slot_mapping_scale[:rows],
dsa.slot_mapping_fp8[:rows],
dsa.slot_mapping_scale[:rows],
rows,
True,
True,
Expand Down Expand Up @@ -1260,7 +1267,7 @@ def _select_rows(self, indexer, kv_a, latent_end, q_a, position_ids, attn_metada
)
hist_topk(
logits,
attn_metadata.gen_indexer_kv_lens_cuda_runtime,
dsa.gen_indexer_kv_lens_cuda_runtime,
buffers["topk"][:rows],
self._hist_workspace.comm_buffer,
self._hist_workspace.state,
Expand Down Expand Up @@ -1773,6 +1780,7 @@ def _draft_on_target(
dense selection is refreshed from the new lengths. The host side of draft step 0's
switch to one row per request comes before that.
"""
dsa = attn_metadata.ext(DSAMetadata)
draft_len = spec_metadata.runtime_draft_len
# The accepted tokens' rows: draft step 0's input ids (prepare_1st_drafter_inputs).
accepted_tokens = inputs["input_ids"].view(batch_size, draft_len + 1)
Expand All @@ -1781,12 +1789,12 @@ def _draft_on_target(
draft_tokens = accepted_tokens.new_empty((batch_size, draft_len))
positions = inputs["position_ids"]
eh_proj_output = gather_ids = pool_rows = next_new_tokens = None
uses_mtp_index_share = isinstance(attn_metadata, MTPIndexShareMetadata)
index_share = attn_metadata.find_ext(MTPIndexShareExtension)
with contextlib.ExitStack() as draft_scope:
if uses_mtp_index_share:
attn_metadata.set_in_mtp_draft_loop(True)
attn_metadata.set_mtp_num_accepted(num_accepted_tokens)
draft_scope.callback(_reset_mtp_index_share, attn_metadata)
if index_share is not None:
index_share.set_in_mtp_draft_loop(True)
index_share.set_mtp_num_accepted(num_accepted_tokens)
draft_scope.callback(index_share.reset_mtp_index_share)
attn_metadata = draft_scope.enter_context(
self.draft_kv_cache_context(attn_metadata, draft_kv_cache_manager)
)
Expand All @@ -1795,8 +1803,8 @@ def _draft_on_target(
seq_lens = attn_metadata.seq_lens_cuda[:batch_size]
kv_lens = attn_metadata.kv_lens_cuda[:batch_size]
for i in range(draft_len):
if uses_mtp_index_share:
attn_metadata.set_skip_topk(i > 0)
if index_share is not None:
index_share.set_skip_topk(i > 0)
if i == 0:
hidden_states, gather_ids, pool_rows = decode.draft_first_step(
attn_metadata,
Expand Down Expand Up @@ -1843,12 +1851,10 @@ def _draft_on_target(
kv_lens,
rewind_accepted,
)
if not last and attn_metadata.skip_indexer_for_gen_reqs:
attn_metadata.prepare_dense_topk_indices(
attn_metadata.kv_lens_cuda, device=True
)
if not last and dsa.skip_indexer_for_gen_reqs:
dsa.prepare_dense_topk_indices(attn_metadata.kv_lens_cuda, device=True)
if i == 0:
attn_metadata.use_spec_decoding = False
attn_metadata.ext(SpecDecodingMetadata).use_spec_decoding = False
self._drafted_rows = (draft_tokens, next_new_tokens)
return draft_tokens

Expand Down
62 changes: 36 additions & 26 deletions ..._torch/_experimental/modeling_v2/models/gpt_oss/gpt_oss_120b__sm_10x__default/modeling.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -75,6 +75,8 @@
from tensorrt_llm._torch._experimental.modeling_v2.catalog.quantization.mxfp8_quantize import (
mxfp8_quantize,
)
from tensorrt_llm._torch.attention.backends.features import FlashMlaMetadata
from tensorrt_llm._torch.attention.backends.features.spec_decoding import SpecDecodingMetadata
from tensorrt_llm._torch.attention.backends.interface import AttentionMetadata
from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata
from tensorrt_llm._torch.model_config import ModelConfig
Expand Down Expand Up @@ -153,6 +155,12 @@
_MOE_MAX_T = 8192


def _flash_mla_field(md: TrtllmAttentionMetadata, name: str):
"""A FlashMLA metadata field, or None when FlashMLA is not attached."""
flash_mla = md.find_ext(FlashMlaMetadata)
return getattr(flash_mla, name) if flash_mla is not None else None


def _in_chunks(experts: Callable, x: torch.Tensor, L: "_LayerWeights") -> torch.Tensor:
"""`experts(x, L)` over at most `_MOE_MAX_T` rows at a time, re-joined in order."""
parts = [experts(c, L) for c in torch.split(x, _MOE_MAX_T, 0)]
Expand Down Expand Up @@ -529,6 +537,7 @@ def _check_step_contract(self, md) -> None:

def step_args(self, md: TrtllmAttentionMetadata) -> dict:
"""Project the prepared metadata onto thop_attention's explicit batch state, once per forward."""
spec = md.find_ext(SpecDecodingMetadata)
return dict(
sequence_length=md.kv_lens_cuda_runtime,
host_past_key_value_lengths=md.kv_lens_runtime,
Expand All @@ -553,21 +562,21 @@ def step_args(self, md: TrtllmAttentionMetadata) -> dict:
block_ids_per_seq=md.block_ids_per_seq,
max_context_q_len_override=md.max_context_q_len_override,
is_cross=md.is_cross,
is_spec_decoding_enabled=md.is_spec_decoding_enabled,
use_spec_decoding=md.use_spec_decoding,
is_spec_dec_tree=md.is_spec_dec_tree,
spec_decoding_generation_lengths=md.spec_decoding_generation_lengths,
spec_decoding_position_offsets_for_cpp=md.spec_decoding_position_offsets_for_cpp,
spec_decoding_packed_mask=md.spec_decoding_packed_mask,
spec_decoding_bl_tree_mask_offset=md.spec_decoding_bl_tree_mask_offset,
spec_decoding_bl_tree_mask=md.spec_decoding_bl_tree_mask,
spec_decoding_target_max_draft_tokens=md.max_total_draft_tokens,
spec_bl_tree_first_sparse_mask_offset_kv=md.spec_bl_tree_first_sparse_mask_offset_kv,
is_spec_decoding_enabled=spec.is_spec_decoding_enabled,
use_spec_decoding=spec.use_spec_decoding,
is_spec_dec_tree=spec.is_spec_dec_tree,
spec_decoding_generation_lengths=spec.generation_lengths,
spec_decoding_position_offsets_for_cpp=spec.position_offsets_for_cpp,
spec_decoding_packed_mask=spec.packed_mask,
spec_decoding_bl_tree_mask_offset=spec.bl_tree_mask_offset,
spec_decoding_bl_tree_mask=spec.bl_tree_mask,
spec_decoding_target_max_draft_tokens=spec.max_total_draft_tokens,
spec_bl_tree_first_sparse_mask_offset_kv=spec.bl_tree_first_sparse_mask_offset_kv,
num_sparse_topk=md.num_sparse_topk,
flash_mla_tile_scheduler_metadata=md.flash_mla_tile_scheduler_metadata,
flash_mla_num_splits=md.flash_mla_num_splits,
flash_mla_tile_scheduler_metadata=_flash_mla_field(md, "tile_scheduler_metadata"),
flash_mla_num_splits=_flash_mla_field(md, "num_splits"),
max_num_sequences=md.max_num_sequences,
force_prepare_spec_dec_tree_mask=md.force_prepare_spec_dec_tree_mask,
force_prepare_spec_dec_tree_mask=spec.force_prepare_tree_mask,
)

def forward(
Expand Down Expand Up @@ -766,6 +775,7 @@ def _check_step_contract(self, md) -> None:

def step_args(self, md: TrtllmAttentionMetadata) -> dict:
"""Project the prepared metadata onto thop_attention's explicit batch state, once per forward."""
spec = md.find_ext(SpecDecodingMetadata)
return dict(
sequence_length=md.kv_lens_cuda_runtime,
host_past_key_value_lengths=md.kv_lens_runtime,
Expand All @@ -790,21 +800,21 @@ def step_args(self, md: TrtllmAttentionMetadata) -> dict:
block_ids_per_seq=md.block_ids_per_seq,
max_context_q_len_override=md.max_context_q_len_override,
is_cross=md.is_cross,
is_spec_decoding_enabled=md.is_spec_decoding_enabled,
use_spec_decoding=md.use_spec_decoding,
is_spec_dec_tree=md.is_spec_dec_tree,
spec_decoding_generation_lengths=md.spec_decoding_generation_lengths,
spec_decoding_position_offsets_for_cpp=md.spec_decoding_position_offsets_for_cpp,
spec_decoding_packed_mask=md.spec_decoding_packed_mask,
spec_decoding_bl_tree_mask_offset=md.spec_decoding_bl_tree_mask_offset,
spec_decoding_bl_tree_mask=md.spec_decoding_bl_tree_mask,
spec_decoding_target_max_draft_tokens=md.max_total_draft_tokens,
spec_bl_tree_first_sparse_mask_offset_kv=md.spec_bl_tree_first_sparse_mask_offset_kv,
is_spec_decoding_enabled=spec.is_spec_decoding_enabled,
use_spec_decoding=spec.use_spec_decoding,
is_spec_dec_tree=spec.is_spec_dec_tree,
spec_decoding_generation_lengths=spec.generation_lengths,
spec_decoding_position_offsets_for_cpp=spec.position_offsets_for_cpp,
spec_decoding_packed_mask=spec.packed_mask,
spec_decoding_bl_tree_mask_offset=spec.bl_tree_mask_offset,
spec_decoding_bl_tree_mask=spec.bl_tree_mask,
spec_decoding_target_max_draft_tokens=spec.max_total_draft_tokens,
spec_bl_tree_first_sparse_mask_offset_kv=spec.bl_tree_first_sparse_mask_offset_kv,
num_sparse_topk=md.num_sparse_topk,
flash_mla_tile_scheduler_metadata=md.flash_mla_tile_scheduler_metadata,
flash_mla_num_splits=md.flash_mla_num_splits,
flash_mla_tile_scheduler_metadata=_flash_mla_field(md, "tile_scheduler_metadata"),
flash_mla_num_splits=_flash_mla_field(md, "num_splits"),
max_num_sequences=md.max_num_sequences,
force_prepare_spec_dec_tree_mask=md.force_prepare_spec_dec_tree_mask,
force_prepare_spec_dec_tree_mask=spec.force_prepare_tree_mask,
)

def forward(
Expand Down
Loading
Loading