Skip to content

[TRTLLM-14855][chore] Use the canonical modules instead of the layout compatibility shims #22086

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

Open
YihuiLu512 wants to merge 1 commit into
base: main
Choose a base branch
from
+8 −9
Open

[TRTLLM-14855][chore] Use the canonical modules instead of the layout compatibility shims #22086

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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 cpp/tensorrt_llm/thop/moe/communication/moeAlltoAllOp.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -74,7 +74,7 @@ static constexpr size_t CACHELINE_ALIGNMENT = 128;
//
// The MoE all-to-all workspace setup used tensorrt_llm::mpi::MpiComm directly.
// mpiUtils.h throws "MPI is disabled, DON'T USE MPI" whenever TLLM_DISABLE_MPI=1,
// which the Ray executor sets for every worker (executor/ray_executor.py), so the
// which the Ray executor sets for every worker (tensorrt_llm/executor/ray/executor.py), so the
// NVLinkOneSided strategy always failed there and silently degraded to
// NVLinkTwoSided. Route through the Torch process group in that case; the same
// pattern is already used by cacheTransceiver / allreduceOp via pg_utils.
Expand Down
4 changes: 1 addition & 3 deletions tensorrt_llm/_torch/attention/backends/fp4_mla/cache_update.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -787,9 +787,7 @@ def scatter_fp4_mla_kv_cache(
and ``positions[token_offset : token_offset + N]``. Callers must pass a
latent_cache pre-sliced to the current phase (context or generation) so
that ``shape[0]`` matches the number of index entries they intend to
consume. ``MLA.forward_impl`` (tensorrt_llm/_torch/modules/attention.py)
slices ``latent_cache[:num_ctx_tokens]`` for context and
``latent_cache[num_ctx_tokens:]`` for generation before dispatching.
consume.

@sunnyqgg sunnyqgg Nov 11, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[Nice to have] Only the path in the removed sentence was stale. MLA now lives in tensorrt_llm/_torch/attention/mla.py, and MLA.forward_impl still slices latent_cache[:num_ctx_tokens, ...] for context (mla.py:1117) and latent_cache[num_ctx_tokens:, ...] for generation (mla.py:1148) before dispatching. That pre-slicing is what the contract above requires of callers, so I'd keep the pointer and fix only the path:

Suggested change
consume.
consume. ``MLA.forward_impl`` (tensorrt_llm/_torch/attention/mla.py)
slices ``latent_cache[:num_ctx_tokens]`` for context and
``latent_cache[num_ctx_tokens:]`` for generation before dispatching.

@YihuiLu512 YihuiLu512 Nov 11, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks — the path was indeed the only stale part. I dropped the sentence rather than re-pointing it, for two reasons:

  1. The slicing lives in the private MLA._forward_impl (mla.py:1117/1148), not forward_impl, which skips it on the sparse-hooks path; with external RoPE the context latent is rebuilt via torch.cat, and latent_cache_gen can be supplied by the caller.
  2. The direct callers of this helper are the backends (e.g. fmha/fp4_mla.py slices latent_cache[:num_tokens] itself), so the caller contract above already states the requirement without a cross-file pointer that can drift again.


Callers must pass ``phase``, ``local_layer``, and ``v_head_dim``. Context
scatter writes the final FP4 tile representation directly. Dimensions
Expand Down
2 changes: 1 addition & 1 deletion tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/mamba_cache_manager_v2.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -93,7 +93,7 @@
)

if TYPE_CHECKING:
from tensorrt_llm._torch.attention_backend.interface import AttentionMetadata
from tensorrt_llm._torch.attention.backends.interface import AttentionMetadata
from tensorrt_llm.llmapi.llm_args import DecodingBaseConfig
from tensorrt_llm.sampling_params import SamplingParams

Expand Down
3 changes: 2 additions & 1 deletion tensorrt_llm/_torch/speculative/dflash.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1277,7 +1277,8 @@ def _init_ctx_block_tables(
``kv_cache_block_offsets`` entries are ``pool_block_index *
num_pool_layers * kv_factor`` plus the K/V field index. What that has to
be divided by to index a ``get_buffers`` view depends on the manager --
see the stride check below; attention_backend/utils.py hardcodes the V1
see the stride check below;
``append_mla_latent_cache_generation_cuda_graph_safe`` hardcodes the V1
divisor and is not a general contract.
"""
mgr = draft_kv_cache_manager
Expand Down
4 changes: 2 additions & 2 deletions tests/unittest/_torch/ray_orchestrator/multi_gpu/test_mnnvl_allreduce.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,7 @@
try:
import ray
except ModuleNotFoundError:
from tensorrt_llm import ray_stub as ray
from tensorrt_llm.executor.ray import stub as ray

from tensorrt_llm._torch.distributed.communicator import TorchDist
from tensorrt_llm.functional import AllReduceFusionOp, AllReduceStrategy
Expand Down Expand Up @@ -60,7 +60,7 @@ def __init__(self, rank: int, world_size: int):

assert len(ray.get_gpu_ids()) == 1
self.gpu = int(ray.get_gpu_ids()[0])
from tensorrt_llm.executor.ray_gpu_worker import RayWorkerWrapper
from tensorrt_llm.executor.ray.gpu_worker import RayWorkerWrapper

torch.cuda.set_device(RayWorkerWrapper.physical_to_local_id(self.gpu))

Expand Down
2 changes: 1 addition & 1 deletion tests/unittest/_torch/ray_orchestrator/multi_gpu/test_refit_compile_isolation.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -92,8 +92,8 @@
from utils.util import skip_pre_rubin

from tensorrt_llm import LLM
from tensorrt_llm._ray_utils import control_action_decorator
from tensorrt_llm._torch.utils import get_device_uuid
from tensorrt_llm.executor.ray.utils import control_action_decorator
from tensorrt_llm.llmapi import KvCacheConfig, MoeConfig, SamplingParams, TorchCompileConfig
from tensorrt_llm.llmapi.rlhf_utils import WorkerExtension

Expand Down
Loading