[None][feat] Run the Muse modeling_v2 composition with a TRTLLM DFlash drafter

docs/source/features/speculative-decoding.md

@@ -152,7 +152,7 @@ Reference: [DFlash: Distilled Flash Speculative Decoding](https://arxiv.org/pdf/
 * `speculative_model`: Path or HuggingFace model ID for the DFlash draft model.
 * `mask_token_id`: Token ID used as the mask token for parallel prediction. If not set, it is read from the draft model config.
 * `target_layer_ids`: List of target model layer indices whose hidden states are captured for cross-attention in the draft model. If not set, read from the draft model config.
-* `attention_backend`: Cross-attention backend for the draft model, independent of the drafter's standard attention modules. `"AUTO"` (the default) resolves to `"TRTLLM"` where it can run and every drafter layer is windowed (sliding-window drafters such as DFlash 2), and to `"VANILLA"` otherwise, including on a disaggregated worker with `TRTLLM_DISAGG_SHARED_DRAFT_KV=0` and under `modeling_v2`. `"VANILLA"` uses FlashAttention with a contiguous context K/V cache and runs anywhere. `"TRTLLM"` uses TRTLLM-Gen FMHA (via FlashInfer) over a paged context K/V cache held by the KV cache manager (see [Draft KV cache](#dflash-draft-kv-cache)) and supports SM100/SM103 only. `"FA4"` uses the flash-attn CuTe DSL kernels on the same paged cache and supports SM90 only.
+* `attention_backend`: Cross-attention backend for the draft model, independent of the drafter's standard attention modules. `"AUTO"` (the default) resolves to `"TRTLLM"` where it can run and every drafter layer is windowed (sliding-window drafters such as DFlash 2), and to `"VANILLA"` otherwise, including on a disaggregated worker with `TRTLLM_DISAGG_SHARED_DRAFT_KV=0`. `"VANILLA"` uses FlashAttention with a contiguous context K/V cache and runs anywhere. `"TRTLLM"` uses TRTLLM-Gen FMHA (via FlashInfer) over a paged context K/V cache held by the KV cache manager (see [Draft KV cache](#dflash-draft-kv-cache)) and supports SM100/SM103 only. `"FA4"` uses the flash-attn CuTe DSL kernels on the same paged cache and supports SM90 only.
 
 ```python
 from tensorrt_llm.llmapi import DFlashDecodingConfig

tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/dispatch.py

@@ -545,8 +545,8 @@ def dflash_linear(spec_config: Any) -> bool:
     """Whether ``spec_config`` is DFlash drafting a linear chain of ``DRAFT_LEN`` tokens.
 
     That is the drafter every composition's drafter, worker and head paths are
-    written for: one draft per position (no tree), a fixed draft length (no
-    schedule), and the drafter's VANILLA pooled-context attention.
+    written for: one draft per position (no tree) and a fixed draft length (no
+    schedule).
     """
     from tensorrt_llm.llmapi.llm_args import DFlashDecodingConfig
 
@@ -559,7 +559,6 @@ def dflash_linear(spec_config: Any) -> bool:
         and getattr(spec_config, "max_draft_len", None) == DRAFT_LEN
         and getattr(spec_config, "max_total_draft_tokens", None) == DRAFT_LEN
         and getattr(spec_config, "draft_len_schedule", None) is None
-        and getattr(spec_config, "attention_backend", None) == "VANILLA"
     )
 
 
@@ -572,8 +571,7 @@ def describe_speculation(spec_config: Any) -> str:
         f"{getattr(mode, 'name', type(spec_config).__name__)} "
         f"max_draft_len={getattr(spec_config, 'max_draft_len', None)} "
         f"max_total_draft_tokens={getattr(spec_config, 'max_total_draft_tokens', None)} "
-        f"draft_len_schedule={getattr(spec_config, 'draft_len_schedule', None)} "
-        f"attention_backend={getattr(spec_config, 'attention_backend', None)}"
+        f"draft_len_schedule={getattr(spec_config, 'draft_len_schedule', None)}"
     )
 
 
@@ -604,7 +602,7 @@ def served_compositions(compositions: Mapping[CompositionKey, TargetPaths]) -> s
     return (
         f"{keys}, each with DFlash drafting a linear chain of {DRAFT_LEN} tokens "
         "(max_draft_len = max_total_draft_tokens = "
-        f"{DRAFT_LEN}, no draft_len_schedule, VANILLA attention_backend) and an NVFP4 checkpoint"
+        f"{DRAFT_LEN}, no draft_len_schedule) and an NVFP4 checkpoint"
     )
 
 
@@ -1143,7 +1141,11 @@ def compose(
             f"{TAG} no fast path for this deployment ({reason}); running the built-in forward"
         )
         return NO_PATHS
-    windowed = windowed_context_arena(spec_worker)
+    from tensorrt_llm._torch.speculative.dflash import _PAGED_ATTENTION_BACKENDS
+
+    # The worker pages the drafter's context on these backends, and windows only a dense arena.
+    paged = getattr(spec_worker, "_dflash_attention_backend", None) in _PAGED_ATTENTION_BACKENDS
+    windowed = windowed_context_arena(spec_worker) and not paged
     paths, gated = capacity_gated(frozenset(served.served), max_batch_size)
     paths, yields = yielded(paths)
     if windowed:
@@ -1169,6 +1171,8 @@ def compose(
     plan = Plan(paths)
     if windowed:
         composed += f" with {_WINDOWED_ARENA}"
+    if paged:
+        composed += " with a paged drafter context"
     for gate in gated:
         composed += (
             f" without {gate.path} (max_batch_size {max_batch_size} above {gate.most_requests};"