[None][chore] cap RoPE cos/sin table to runtime max_seq_len by lancelly · Pull Request #14241 · NVIDIA/TensorRT-LLM

lancelly · 2026-05-18T07:10:36Z

Summary

Split out only the RoPE cos/sin table cap from #14053 / #14241 as an independent, always-on optimization.

Change

RoPE cos/sin table cap (`attention_backend/interface.py`, `models/modeling_deepseekv4.py`)

RopeParams.from_config auto-caps max_positions to max_seq_len when the latter is set, positive, and smaller than max_positions. YaRN cos/sin at position p depends only on p and the YaRN scaling constants (which use original_max_position_embeddings, NOT the table size), so truncating to max_seq_len yields bitwise-identical values for indices < cap. The cap is applied via field mutation so it participates in the cache key (preserving dedup semantics) and downstream kernel-facing fields stay consistent.

modeling_deepseekv4.py plumbs model_config.max_seq_len onto the HF PretrainedConfig so from_config can see it (DeepSeek HF configs don't have a max_seq_len attribute by default). The plumb mirrors the existing pattern in modeling_phi3.py.

Runtime safety net: the existing TrtllmAttentionWrapper.ensure_rope_table_size expands the table on demand if a kernel asks for a larger max_positions (e.g., MTP draft tokens push position index above the cap). So the cap is a memory-only optimization, not a correctness contract. At worst the extended-MTP edge case triggers one resize per attention wrapper (≈4 MiB × num_layers, paid once); the cap still saves vs the un-capped HF default (~20× smaller table for DSv4-Pro).

For DSv4-Pro on GB300: HF default max_position_embeddings = 163840, runtime max_seq_len = 8192 → table shrinks ≈20×.

Auto-cap RopeParams.max_positions to runtime max_seq_len whenever the latter is smaller than the model's default max_position_embeddings. YaRN cos/sin at position p depends only on p and the YaRN scaling constants (which use original_max_position_embeddings, not the table size), so truncating to max_seq_len yields bitwise-identical values for indices < cap. The cap is applied via field mutation on RopeParams so it participates in the cache key (preserving dedup semantics) and downstream kernel-facing fields stay consistent. modeling_deepseekv4.py plumbs ModelConfig.max_seq_len onto the HF PretrainedConfig so RopeParams.from_config sees it; HF DeepSeek configs lack a max_seq_len attribute by default, otherwise the auto-cap path would be a no-op for DSv4. Runtime safety net: TrtllmAttentionWrapper.ensure_rope_table_size already expands the table on demand if a kernel asks for a larger max_positions (e.g., MTP draft tokens push position index above the cap). So the cap is a memory-only optimization, not a correctness contract. For DSv4-Pro on GB300: HF default max_position_embeddings = 163840, runtime max_seq_len = 8192 -> table shrinks ~20x. Split out from PR NVIDIA#14241 (RoPE cap + MHC pool routing): this PR contains only the RoPE cap, applied unconditionally without the TRTLLM_DSV4_MEM_OPTS env gate. Signed-off-by: Lanyu Liao <lancelly@users.noreply.github.com>

Auto-cap RopeParams.max_positions to runtime max_seq_len whenever the latter is smaller than the model's default max_position_embeddings. YaRN cos/sin at position p depends only on p and the YaRN scaling constants (which use original_max_position_embeddings, not the table size), so truncating to max_seq_len yields bitwise-identical values for indices < cap. The cap is applied via field mutation on RopeParams so it participates in the cache key (preserving dedup semantics) and downstream kernel-facing fields stay consistent. modeling_deepseekv4.py plumbs ModelConfig.max_seq_len (plus MTP draft headroom of predicted_tokens_per_seq - 1) onto the HF PretrainedConfig so RopeParams.from_config sees it; HF DeepSeek configs lack a max_seq_len attribute by default, otherwise the auto-cap path would be a no-op for DSv4. The draft headroom ensures speculative draft token positions, which can overshoot the accepted boundary by up to max_draft_len, still fit within the capped table. Runtime safety net: TrtllmAttentionWrapper.ensure_rope_table_size already expands the table on demand if a kernel still asks for a larger max_positions, so the cap is a memory-only optimization, not a correctness contract. For DSv4-Pro on GB300: HF default max_position_embeddings = 163840, runtime max_seq_len = 8192 -> table shrinks ~20x. Split out from PR NVIDIA#14241 (RoPE cap + MHC pool routing): this PR contains only the RoPE cap, applied unconditionally without the TRTLLM_DSV4_MEM_OPTS env gate. Signed-off-by: Lanyu Liao <lancelly@users.noreply.github.com>

lancelly · 2026-05-19T07:52:25Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-05-19T07:59:39Z

PR_Github #49143 [ run ] triggered by Bot. Commit: 27e7eb8 Link to invocation

Add an opt-in runtime_max_seq_len kwarg to RopeParams.from_config and have DSv4 pass an explicit cap derived from model_config (without mutating the shared HF PretrainedConfig). When the cap is supplied and smaller than the model's default max_position_embeddings, from_config truncates max_positions so the precomputed cos/sin table matches the actual position range used at inference time. Cap is opt-in because RotaryScalingType.longrope folds num_pos into its scaling_factor (scale = num_pos / original_max_pos), where truncation would change cos/sin values. For YaRN / plain / linear / llama3 / dynamic, cos/sin at position p depends only on p and the scaling constants (not on the table size), so capping yields bitwise-identical values for indices < cap. The cap participates in the dataclass hash so per-layer RopeParams still dedup into a single cos/sin tensor in extra_attrs["rope_const_params"]. DSv4 derives the cap by mirroring py_executor_creator.py:599-606: - Spec decoding adds (tokens_per_gen_step - 1) for draft positions plus get_num_extra_kv_tokens(spec_config) extra KV slots. - The overlap scheduler (when enabled with spec) adds another tokens_per_gen_step - 1 positions. disable_overlap_scheduler isn't reachable at attention construction time, so the overlap headroom is included unconditionally when spec is on; the overshoot is at most max_draft_len positions, negligible relative to the cos/sin table savings. Runtime safety net: TrtllmAttentionWrapper.ensure_rope_table_size already expands the table on demand if a kernel asks for a larger max_positions, so the cap is a memory-only optimization, not a correctness contract. For DSv4-Pro on GB300: HF default max_position_embeddings = 163840, runtime max_seq_len = 8192 -> table shrinks ~20x. Split out from PR NVIDIA#14241 (RoPE cap + MHC pool routing): this PR contains only the RoPE cap, applied via an explicit kwarg on RopeParams.from_config without an env switch. Signed-off-by: Lanyu Liao <lancelly@users.noreply.github.com>

Add an opt-in runtime_max_seq_len kwarg to RopeParams.from_config and have DSv4 pass an explicit cap derived from model_config (without mutating the shared HF PretrainedConfig). When the cap is supplied and smaller than the model's default max_position_embeddings, from_config truncates max_positions so the precomputed cos/sin table matches the actual position range used at inference time. The kwarg only triggers the cap; rope_params.max_seq_len keeps its original semantics (read from config) so non-DSv4 callers are unaffected. Cap is opt-in because RotaryScalingType.longrope folds num_pos into its scaling_factor (scale = num_pos / original_max_pos), where truncation would change cos/sin values. For YaRN / plain / linear / llama3 / dynamic, cos/sin at position p depends only on p and the scaling constants (not on the table size), so capping yields bitwise-identical values for indices < cap. The cap participates in the dataclass hash so per-layer RopeParams still dedup into a single cos/sin tensor in extra_attrs["rope_const_params"]. DSv4 derives the cap by mirroring the spec/overlap headroom that py_executor_creator applies to model_engine_max_seq_len before sizing the KV cache: (tokens_per_gen_step - 1) for draft positions, plus get_num_extra_kv_tokens(spec_config), plus another (tokens_per_gen_step - 1) for the overlap scheduler. The overlap branch is included unconditionally because disable_overlap_scheduler isn't reachable at attention construction time; the overshoot when overlap is actually off is at most predicted_tokens_per_seq positions (= 1 + max_total_draft_tokens, which can exceed max_draft_len for tree-shaped drafts), negligible relative to the cos/sin table savings. Runtime safety net: TrtllmAttentionWrapper.ensure_rope_table_size already expands the table on demand if a kernel asks for a larger max_positions, so the cap is a memory-only optimization, not a correctness contract. For DSv4-Pro on GB300: HF default max_position_embeddings = 163840, runtime max_seq_len = 8192 -> table shrinks ~20x. Split out from PR NVIDIA#14241 (RoPE cap + MHC pool routing): this PR contains only the RoPE cap, applied via an explicit kwarg on RopeParams.from_config without an env switch. Signed-off-by: Lanyu Liao <lancelly@users.noreply.github.com>

lancelly · 2026-05-19T09:50:21Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-05-19T09:56:33Z

PR_Github #49171 [ run ] triggered by Bot. Commit: c316cc3 Link to invocation

tensorrt-cicd · 2026-05-19T09:56:36Z

PR_Github #49143 [ run ] completed with state ABORTED. Commit: 27e7eb8

Link to invocation

…s pool Add an opt-in path in _FusedHcWorkspaceCache.get that routes the four output buffers (residual_cur, post_mix_cur, comb_mix_cur, layer_input_cur) and the three internal scratch buffers (y_acc_ws, r_acc_ws, done_counter_ws) through get_memory_buffers() so MHC scratch can share backing memory with other layer-transient buffers (e.g., the attention workspace). The env switch TRTLLM_DSV4_MEM_OPTS is read once at class definition (cached as _USE_POOL) so the lookup doesn't run on every workspace request. Pool routing is gated to decode-sized B (B <= _CACHE_MAX_B). For B > _CACHE_MAX_B (prefill), the code falls through to fresh torch.empty per call, matching the upstream uncached path; this preserves prefill's no-aliasing semantics where layer N's *_prev inputs and layer N+1's *_cur outputs live in distinct tensors. At decode, upstream's existing per-runner LRU already returns the same tensor across layers, so the kernel is already in-place safe; pool routing inherits that aliasing without widening it. With TRTLLM_DSV4_MEM_OPTS off (default), behavior is byte-identical to current source. Split out from PR NVIDIA#14053 (DSv4 mem-opts master switch) and from PR NVIDIA#14241 (RoPE cap + MHC pool): this PR contains only the MHC pool routing piece, still gated on TRTLLM_DSV4_MEM_OPTS so it can be validated in isolation. Signed-off-by: Lanyu Liao <lancelly@users.noreply.github.com>

tensorrt-cicd · 2026-05-19T14:28:30Z

PR_Github #49171 [ run ] completed with state SUCCESS. Commit: c316cc3
/LLM/main/L0_MergeRequest_PR pipeline #38850 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

Link to invocation

qiaoxj07 · 2026-05-20T00:44:53Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-05-20T00:51:42Z

PR_Github #49289 [ run ] triggered by Bot. Commit: c316cc3 Link to invocation

tensorrt-cicd · 2026-05-20T07:55:39Z

PR_Github #49289 [ run ] completed with state SUCCESS. Commit: c316cc3
/LLM/main/L0_MergeRequest_PR pipeline #38951 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

Add an opt-in runtime_max_seq_len kwarg to RopeParams.from_config and have DSv4 pass an explicit cap derived from model_config (without mutating the shared HF PretrainedConfig). When the cap is supplied and smaller than the model's default max_position_embeddings, from_config truncates max_positions so the precomputed cos/sin table matches the actual position range used at inference time. The kwarg only triggers the cap; rope_params.max_seq_len keeps its original semantics (read from config) so non-DSv4 callers are unaffected. Cap is opt-in because RotaryScalingType.longrope folds num_pos into its scaling_factor (scale = num_pos / original_max_pos), where truncation would change cos/sin values. For YaRN / plain / linear / llama3 / dynamic, cos/sin at position p depends only on p and the scaling constants (not on the table size), so capping yields bitwise-identical values for indices < cap. The cap participates in the dataclass hash so per-layer RopeParams still dedup into a single cos/sin tensor in extra_attrs["rope_const_params"]. DSv4 derives the cap by mirroring the spec/overlap headroom that py_executor_creator applies to model_engine_max_seq_len before sizing the KV cache: (tokens_per_gen_step - 1) for draft positions, plus get_num_extra_kv_tokens(spec_config), plus another (tokens_per_gen_step - 1) for the overlap scheduler. The overlap branch is included unconditionally because disable_overlap_scheduler isn't reachable at attention construction time; the overshoot when overlap is actually off is at most predicted_tokens_per_seq positions (= 1 + max_total_draft_tokens, which can exceed max_draft_len for tree-shaped drafts), negligible relative to the cos/sin table savings. Runtime safety net: TrtllmAttentionWrapper.ensure_rope_table_size already expands the table on demand if a kernel asks for a larger max_positions, so the cap is a memory-only optimization, not a correctness contract. For DSv4-Pro on GB300: HF default max_position_embeddings = 163840, runtime max_seq_len = 8192 -> table shrinks ~20x. Split out from PR NVIDIA#14241 (RoPE cap + MHC pool routing): this PR contains only the RoPE cap, applied via an explicit kwarg on RopeParams.from_config without an env switch. Signed-off-by: Lanyu Liao <lancelly@users.noreply.github.com>

_deepseek_v4_pos_embd_params previously required predicted_tokens_per_seq as a positional arg before layer_idx, which broke the existing 3-arg test calls (config, model_config, layer_idx) introduced in the rope params test. Move predicted_tokens_per_seq after layer_idx and default it to 1 (= no speculative headroom). Update the production call site to match. The default is safe: predicted_tokens_per_seq is only consumed when model_config.spec_config is not None, and the runtime safety net in TrtllmAttentionWrapper.ensure_rope_table_size expands the table on demand if a kernel asks for a larger max_positions, so an undersized cap is at worst a missed memory optimization, not a correctness issue. Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com>

lancelly · 2026-05-20T15:38:19Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-05-20T15:45:22Z

PR_Github #49445 [ run ] triggered by Bot. Commit: 83ac162 Link to invocation

tensorrt-cicd · 2026-05-20T20:02:28Z

PR_Github #49445 [ run ] completed with state SUCCESS. Commit: 83ac162
/LLM/main/L0_MergeRequest_PR pipeline #39090 completed with status: 'SUCCESS'

CI Report

Link to invocation

…4241) Signed-off-by: Lanyu Liao <lancelly@users.noreply.github.com> Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> Co-authored-by: Lanyu Liao <lancelly@users.noreply.github.com> Signed-off-by: Fanrong Li <lfr-0531@users.noreply.github.com>

…4241) Signed-off-by: Lanyu Liao <lancelly@users.noreply.github.com> Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> Co-authored-by: Lanyu Liao <lancelly@users.noreply.github.com>

…4241) Signed-off-by: Lanyu Liao <lancelly@users.noreply.github.com> Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> Co-authored-by: Lanyu Liao <lancelly@users.noreply.github.com> Signed-off-by: Fanrong Li <lfr-0531@users.noreply.github.com>

github-actions Bot assigned lancelly May 18, 2026

lfr-0531 added the deepseek-v4 label May 19, 2026

lancelly closed this May 19, 2026

lancelly reopened this May 19, 2026

lancelly force-pushed the dsv4-mem-opts-rope-mhc branch from 54181f0 to 207225b Compare May 19, 2026 07:33

lancelly changed the title ~~[None][perf] DSv4 mem-opts: RoPE cap + MHC pool routing (split from #14053)~~ [None][perf] DSv4: cap RoPE cos/sin table to runtime max_seq_len May 19, 2026

lancelly changed the title ~~[None][perf] DSv4: cap RoPE cos/sin table to runtime max_seq_len~~ [None][perf] cap RoPE cos/sin table to runtime max_seq_len May 19, 2026

lancelly changed the title ~~[None][perf] cap RoPE cos/sin table to runtime max_seq_len~~ [None][chore] cap RoPE cos/sin table to runtime max_seq_len May 19, 2026

lancelly force-pushed the dsv4-mem-opts-rope-mhc branch from 207225b to ef2a4eb Compare May 19, 2026 07:37

lancelly marked this pull request as ready for review May 19, 2026 07:38

lancelly requested review from a team as code owners May 19, 2026 07:38

lancelly requested review from brb-nv and pengbowang-nv and removed request for a team May 19, 2026 07:38

lancelly force-pushed the dsv4-mem-opts-rope-mhc branch from ef2a4eb to 01d391a Compare May 19, 2026 07:43

lancelly force-pushed the dsv4-mem-opts-rope-mhc branch from 01d391a to 27e7eb8 Compare May 19, 2026 07:47

lancelly requested a review from yuxianq May 19, 2026 07:53

pengbowang-nv requested changes May 19, 2026

View reviewed changes

Comment thread tensorrt_llm/_torch/models/modeling_deepseekv4.py Outdated

lancelly force-pushed the dsv4-mem-opts-rope-mhc branch from 27e7eb8 to c3d0ecc Compare May 19, 2026 08:50

pengbowang-nv reviewed May 19, 2026

View reviewed changes

Comment thread tensorrt_llm/_torch/models/modeling_deepseekv4.py Outdated

Comment thread tensorrt_llm/_torch/models/modeling_deepseekv4.py

lancelly force-pushed the dsv4-mem-opts-rope-mhc branch from a8b1352 to c316cc3 Compare May 19, 2026 09:44

lancelly mentioned this pull request May 19, 2026

[None][chore] Route MHC fused-HC workspace through global Buffers pool #14305

Closed

lancelly mentioned this pull request May 19, 2026

[None][chore] memory observability probes for DSv4 mem-opts validation #14308

Merged

lancelly and others added 2 commits May 20, 2026 08:28

lancelly force-pushed the dsv4-mem-opts-rope-mhc branch from c316cc3 to 83ac162 Compare May 20, 2026 15:35

lancelly merged commit 94668b8 into NVIDIA:feat/deepseek_v4 May 21, 2026
6 checks passed

Uh oh!

Conversation

lancelly commented May 18, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Change

RoPE cos/sin table cap (attention_backend/interface.py, models/modeling_deepseekv4.py)

Uh oh!

lancelly commented May 19, 2026

Uh oh!

tensorrt-cicd commented May 19, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

lancelly commented May 19, 2026

Uh oh!

tensorrt-cicd commented May 19, 2026

Uh oh!

tensorrt-cicd commented May 19, 2026

Uh oh!

tensorrt-cicd commented May 19, 2026

Uh oh!

qiaoxj07 commented May 20, 2026

Uh oh!

tensorrt-cicd commented May 20, 2026

Uh oh!

tensorrt-cicd commented May 20, 2026

Uh oh!

lancelly commented May 20, 2026

Uh oh!

tensorrt-cicd commented May 20, 2026

Uh oh!

tensorrt-cicd commented May 20, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

lancelly commented May 18, 2026 •

edited

Loading

RoPE cos/sin table cap (`attention_backend/interface.py`, `models/modeling_deepseekv4.py`)