From 00fb0d9fdb53f089ce91b1d2a7dea32fbd4a1412 Mon Sep 17 00:00:00 2001 From: Shicheng Li Date: Wed, 29 Apr 2026 18:07:33 -0700 Subject: [PATCH 1/2] [None][feat] Add MLA dependency-aware overlap on DSv4 (compressor || q_b_proj+norm) In forward_dsa_proj, the MLA compressor only depends on hidden_states and writes its own KV-cache slot, while q_b_proj + q_b_layernorm only read q_a_layernorm(...). They are mutually independent on the data graph. Today they execute in series on the main stream. Wrap the compressor + (q_b_proj + q_b_layernorm) pair in maybe_execute_in_parallel with the existing aux_stream, gated behind TRTLLM_MLA_EXTRA_OVERLAP=1 (default off). Mathematically identical; only the runtime stream schedule changes. Bench (DSv4-Flash, gb200, 1280 prompts, side-by-side same-shape A/B): TEP4 c128 (kv=0.85, bs=256, conc=128, mtp=0): - (Track 1 ON) overlap=0: 4011.5 sys tok/s, 1002.87 per-GPU - (Track 1 ON) overlap=1: 4099.9 sys tok/s, 1024.98 per-GPU (+2.20%) DEP4 c128 (kv=0.7, bs=64, conc=128, mtp=0), standalone Track C without Track 1, on the prior mingyang-fix sqsh: - overlap=0: 4592.1 sys tok/s, 1148.0 per-GPU - overlap=1: 4690.2 sys tok/s, 1172.6 per-GPU (+2.14%) Two independent shapes, two different sqshs, near-identical +2.2% delta - the structural overlap is shape-agnostic across the MLA path. Accuracy gate (gsm8k 5-shot via lm-eval, custom_tokenizer deepseek_v4, chat-template + system_prompt; combined with Track 1 in same A/B): - baseline: flexible 95.98 +- 0.54 / strict 96.06 +- 0.54 - Track1+TrackC: flexible 96.44 +- 0.51 / strict 96.51 +- 0.51 Both deltas inside the +-0.5pp stderr band - no math change, only stream scheduling. Signed-off-by: Shicheng Li --- tensorrt_llm/_torch/modules/attention.py | 75 +++++++++++++++++++----- 1 file changed, 61 insertions(+), 14 deletions(-) diff --git a/tensorrt_llm/_torch/modules/attention.py b/tensorrt_llm/_torch/modules/attention.py index 5df9859f8fde..abc3b2915c4b 100644 --- a/tensorrt_llm/_torch/modules/attention.py +++ b/tensorrt_llm/_torch/modules/attention.py @@ -2023,22 +2023,69 @@ def forward_impl_with_deepseek_v4(self, qr = q latent_cache = torch.concat([compressed_kv, k_pe], dim=-1) - q = self.q_b_proj(q) - # Per-head RMS: view as [N*n_heads, head_dim] so RMSNorm reduces per-head. - q = self.q_b_layernorm(q.view(-1, self.qk_head_dim)).view_as(q) - # Indexer - topk_indices = None - if self.indexer is not None: - topk_indices = self.indexer( - qr, - hidden_states, - attn_metadata, - position_ids, + # TRTLLM_MLA_EXTRA_OVERLAP=1 reorders the V4 attention prologue so the + # outer compressor (writes its own KV-cache slot, only reads + # hidden_states) executes on the auxiliary CUDA stream concurrently + # with q_b_proj + q_b_layernorm on the default stream. Both branches + # are entirely independent (no shared inputs and disjoint writes), + # so this is a pure dependency-aware reorder. Falls back to the + # serial schedule when the env-var is unset, when the aux stream is + # unavailable, or when multi-stream mode is off. + _v4_extra_overlap = (os.environ.get("TRTLLM_MLA_EXTRA_OVERLAP", "0") + == "1" and self.compressor is not None + and self.aux_stream is not None) + + if _v4_extra_overlap: + + def _q_branch(): + q_proj = self.q_b_proj(q) + # Per-head RMS: view as [N*n_heads, head_dim] so RMSNorm + # reduces per-head. + return self.q_b_layernorm(q_proj.view( + -1, self.qk_head_dim)).view_as(q_proj) + + def _compressor_branch(): + self.compressor(hidden_states, attn_metadata) + return None + + q, _ = maybe_execute_in_parallel( + _q_branch, + _compressor_branch, + self.ln_events[0], + self.ln_events[1], + self.aux_stream, ) - # Compressor - if self.compressor is not None: - self.compressor(hidden_states, attn_metadata) + # Indexer (depends on qr / hidden_states only). Kept serial + # because it internally reuses self.aux_stream for its own + # multi-stream q-proj || weights-proj split; running it + # concurrently with q_b_proj would create a stream-aliasing + # hazard. + topk_indices = None + if self.indexer is not None: + topk_indices = self.indexer( + qr, + hidden_states, + attn_metadata, + position_ids, + ) + else: + q = self.q_b_proj(q) + # Per-head RMS: view as [N*n_heads, head_dim] so RMSNorm reduces per-head. + q = self.q_b_layernorm(q.view(-1, self.qk_head_dim)).view_as(q) + # Indexer + topk_indices = None + if self.indexer is not None: + topk_indices = self.indexer( + qr, + hidden_states, + attn_metadata, + position_ids, + ) + + # Compressor + if self.compressor is not None: + self.compressor(hidden_states, attn_metadata) assert q.shape[ 0] == num_tokens, f"Expect q.shape[0] to be {num_tokens}, but got {q.shape[0]}" From bcd2d42fc420eda33e8063227f3bdd4311f9506d Mon Sep 17 00:00:00 2001 From: Shicheng Li Date: Wed, 29 Apr 2026 19:12:28 -0700 Subject: [PATCH 2/2] [None][refactor] Hoist indexer call out of overlap/fallback if-else Per @yuxianq's review on #13629: the self.indexer(...) call was duplicated in both arms of the TRTLLM_MLA_EXTRA_OVERLAP if/else. Indexer is independent of both q_b_proj and the compressor's KV slot writes, so it can run after either schedule. Hoisting it removes 9 lines of duplicated code. Net behaviour: - Overlap path: q_b_proj || compressor (parallel) -> indexer (unchanged). - Fallback path was: q_b_proj+norm -> indexer -> compressor. Now: q_b_proj+norm -> compressor -> indexer. The compressor/indexer order swap is safe; the two ops have disjoint read/write sets (compressor reads hidden_states, writes its own KV slot; indexer reads qr/hidden_states/position_ids and writes only topk_indices). The overlap path already runs them in this relative order. Signed-off-by: Shicheng Li --- tensorrt_llm/_torch/modules/attention.py | 39 +++++++++--------------- 1 file changed, 14 insertions(+), 25 deletions(-) diff --git a/tensorrt_llm/_torch/modules/attention.py b/tensorrt_llm/_torch/modules/attention.py index abc3b2915c4b..3eba48b895a2 100644 --- a/tensorrt_llm/_torch/modules/attention.py +++ b/tensorrt_llm/_torch/modules/attention.py @@ -2055,38 +2055,27 @@ def _compressor_branch(): self.ln_events[1], self.aux_stream, ) - - # Indexer (depends on qr / hidden_states only). Kept serial - # because it internally reuses self.aux_stream for its own - # multi-stream q-proj || weights-proj split; running it - # concurrently with q_b_proj would create a stream-aliasing - # hazard. - topk_indices = None - if self.indexer is not None: - topk_indices = self.indexer( - qr, - hidden_states, - attn_metadata, - position_ids, - ) else: q = self.q_b_proj(q) # Per-head RMS: view as [N*n_heads, head_dim] so RMSNorm reduces per-head. q = self.q_b_layernorm(q.view(-1, self.qk_head_dim)).view_as(q) - # Indexer - topk_indices = None - if self.indexer is not None: - topk_indices = self.indexer( - qr, - hidden_states, - attn_metadata, - position_ids, - ) - - # Compressor if self.compressor is not None: self.compressor(hidden_states, attn_metadata) + # Indexer is independent of both q_b_proj and the compressor's KV-cache + # write, so it runs after either schedule. Kept serial because it + # internally reuses self.aux_stream for its own multi-stream q-proj || + # weights-proj split; running it concurrently with q_b_proj would + # create a stream-aliasing hazard. + topk_indices = None + if self.indexer is not None: + topk_indices = self.indexer( + qr, + hidden_states, + attn_metadata, + position_ids, + ) + assert q.shape[ 0] == num_tokens, f"Expect q.shape[0] to be {num_tokens}, but got {q.shape[0]}"