diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_lora_problem_builder.cu b/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_lora_problem_builder.cu index f1b6da4ca189..e6b3f9b98fc0 100644 --- a/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_lora_problem_builder.cu +++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_lora_problem_builder.cu @@ -68,8 +68,16 @@ __global__ void moeLoraProblemBuilderKernel(int32_t const* __restrict__ ranks, i // Problem sizes: each permuted token gets its own (M=1) GEMM. This matches // worst-case scheduling with no run-length aggregation; a future // optimization can aggregate consecutive identical-adapter tokens. + // + // Rank-0 rows carry no active adapter (base/no-LoRA request, padding, or + // warmup) and have null A/B pointers, so their delta is zero and the caller + // pre-zeroes the output. The in-GEMM already collapses to N=0 (rank is its + // N) and is skipped, but the out-GEMM's N is out_hidden_size; forcing it to + // zero here lets the grouped GEMM skip these rows too instead of launching + // tiles that dereference the null B pointer. + int const out_n = (rank > 0) ? static_cast(out_hidden_size) : 0; problem_sizes_in[i] = cutlass::gemm::GemmCoord(1, rank, static_cast(in_hidden_size)); - problem_sizes_out[i] = cutlass::gemm::GemmCoord(1, static_cast(out_hidden_size), rank); + problem_sizes_out[i] = cutlass::gemm::GemmCoord(1, out_n, rank); // Pointer rows. dtype_bytes scales the per-row stride so the same // builder serves bf16/fp16/fp32 adapters without templating. diff --git a/cpp/tests/unit_tests/kernels/moeLoraProblemBuilderTest.cu b/cpp/tests/unit_tests/kernels/moeLoraProblemBuilderTest.cu index e9a414e837ca..26891eb532f3 100644 --- a/cpp/tests/unit_tests/kernels/moeLoraProblemBuilderTest.cu +++ b/cpp/tests/unit_tests/kernels/moeLoraProblemBuilderTest.cu @@ -70,8 +70,9 @@ RefOutputs cpuReference(std::vector const& ranks, std::vector for (int64_t i = 0; i < P; ++i) { int32_t const rank = ranks[i]; + int const out_n = (rank > 0) ? static_cast(out_hidden_size) : 0; r.problem_sizes_in[i] = cutlass::gemm::GemmCoord(1, rank, static_cast(in_hidden_size)); - r.problem_sizes_out[i] = cutlass::gemm::GemmCoord(1, static_cast(out_hidden_size), rank); + r.problem_sizes_out[i] = cutlass::gemm::GemmCoord(1, out_n, rank); int64_t const in_row_stride = in_hidden_size * dtype_bytes; int64_t const work_row_stride = max_lora_rank * dtype_bytes; diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py index de5a2bef27b5..85fdac23361d 100755 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py +++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py @@ -548,7 +548,10 @@ def _extract_moe_lora_tensors( if entry is None: continue kernel_ranks[slot] = entry["adapter_size"] - kernel_ptrs[slot] = entry["weight_pointers"] + # weight_pointers is built flat ([num_seqs * 3], row-major (A, B, + # DoRA) per seq) in PyTorchModelEngine._build_lora_params; the MoE op + # expects a [num_seqs, 3] table, so restore that shape. + kernel_ptrs[slot] = entry["weight_pointers"].reshape(-1, 3) try: active_max_rank = max(active_max_rank, int(entry["adapter_size"].max().item())) diff --git a/tensorrt_llm/_torch/peft/lora/cuda_graph_lora_params.py b/tensorrt_llm/_torch/peft/lora/cuda_graph_lora_params.py index 7c43c05bdd5e..5a7b3bd9ba9d 100644 --- a/tensorrt_llm/_torch/peft/lora/cuda_graph_lora_params.py +++ b/tensorrt_llm/_torch/peft/lora/cuda_graph_lora_params.py @@ -308,6 +308,37 @@ def zero_out_weight_pointers(slot_id: int): layer_param.d_b_ptrs.copy_(layer_param.h_b_ptrs, non_blocking=True) layer_param.d_b_prime_ptrs.copy_(layer_param.h_b_prime_ptrs, non_blocking=True) + # The routed-expert MoE LoRA path reads its slot weight-pointer table + # from pinned buffers that get_moe_slot_inputs caches for stable + # addresses but only refreshes during graph capture. The captured H2D + # copy reads them by address at replay, so refresh them in place here; + # otherwise the ranks update but the pointers stay stale and an active + # (rank>0) slot dereferences a stale/null pointer at replay. + self._refresh_moe_slot_ptr_cache() + + def _refresh_moe_slot_ptr_cache(self) -> None: + """Re-pack cached MoE slot weight-pointer tables from the current + per-layer host pointers so CUDA-graph replay reads up-to-date pointers. + + No-op until get_moe_slot_inputs has created cache entries. The cached + pinned buffers are updated in place to keep their addresses stable (the + captured H2D copy reads them by address at replay). + """ + cache = getattr(self, "_moe_slot_ptrs_cache", None) + if not cache: + return + for (layer_idx, module_id), packed in cache.items(): + key = self.layer_module2key.get((layer_idx, module_id)) + if key is None: + continue + layer_param = self.layer_params.get(key) + if layer_param is None: + continue + local_module_id = key.module_ids.index(module_id) + packed[:, 0].copy_(layer_param.h_b_ptrs[local_module_id].to(torch.int64)) + packed[:, 1].copy_(layer_param.h_b_prime_ptrs[local_module_id].to(torch.int64)) + # Column 2 (DoRA magnitude) stays zero. + @staticmethod def get_offset_from_counts( counts: torch.Tensor, full: bool = False, out: torch.Tensor = None diff --git a/tests/integration/test_lists/test-db/l0_a100.yml b/tests/integration/test_lists/test-db/l0_a100.yml index 4884bc1c6633..77176b68d7e5 100644 --- a/tests/integration/test_lists/test-db/l0_a100.yml +++ b/tests/integration/test_lists/test-db/l0_a100.yml @@ -17,6 +17,7 @@ l0_a100: - unittest/llmapi/test_llm_pytorch.py -m "part1" - unittest/llmapi/test_llm_pytorch.py -m "part2" - unittest/llmapi/test_llm_pytorch.py -m "part3" + - unittest/llmapi/test_llm_pytorch.py::test_qwen_moe_routed_expert_multi_lora_varying_ranks TIMEOUT (90) - unittest/llmapi/test_llm_encode.py - unittest/llmapi/test_mpi_session.py ISOLATION - unittest/llmapi/test_memory_profiling.py::test_profile_kvcache # profile kvcache for vision encoder diff --git a/tests/integration/test_lists/test-db/l0_h100.yml b/tests/integration/test_lists/test-db/l0_h100.yml index 6fd8c15bf5bf..c12e322ed194 100644 --- a/tests/integration/test_lists/test-db/l0_h100.yml +++ b/tests/integration/test_lists/test-db/l0_h100.yml @@ -219,6 +219,7 @@ l0_h100: - unittest/llmapi/test_llm_pytorch.py -m "part1" - unittest/llmapi/test_llm_pytorch.py -m "part2" - unittest/llmapi/test_llm_pytorch.py -m "part3" + - unittest/llmapi/test_llm_pytorch.py::test_qwen_moe_routed_expert_multi_lora_varying_ranks TIMEOUT (90) - unittest/llmapi/test_async_llm.py -m "not (gpu2 or gpu4)" - examples/test_ray.py::test_llm_inference_async_ray - condition: diff --git a/tests/unittest/llmapi/test_llm_pytorch.py b/tests/unittest/llmapi/test_llm_pytorch.py index ece8990f955c..b09a011d17d9 100644 --- a/tests/unittest/llmapi/test_llm_pytorch.py +++ b/tests/unittest/llmapi/test_llm_pytorch.py @@ -1,5 +1,6 @@ import asyncio import json +import os import pathlib import random import time @@ -12,7 +13,8 @@ from tensorrt_llm.disaggregated_params import DisaggregatedParams from tensorrt_llm.executor import GenerationExecutorWorker, RequestError from tensorrt_llm.executor.rpc_proxy import GenerationExecutorRpcProxy -from tensorrt_llm.llmapi import CacheTransceiverConfig, KvCacheConfig +from tensorrt_llm.llmapi import (CacheTransceiverConfig, CudaGraphConfig, + KvCacheConfig) from tensorrt_llm.llmapi.llm_args import (NGramDecodingConfig, PeftCacheConfig, SchedulerConfig, WaitingQueuePolicy) from tensorrt_llm.llmapi.tokenizer import TransformersTokenizer @@ -40,6 +42,7 @@ skip_gpu_memory_less_than_138gb, skip_ray) from utils.llm_data import llm_models_root from tensorrt_llm.lora_helper import LoraConfig +from tensorrt_llm.llmapi.llm_args import MoeConfig from tensorrt_llm.executor.request import LoRARequest import tempfile @@ -67,7 +70,7 @@ @force_ampere -@pytest.mark.parametrize("enable_chunked_prefill,", [False, True]) +@pytest.mark.parametrize("enable_chunked_prefill", [False, True]) @pytest.mark.part2 def test_tinyllama_logits_processor(enable_chunked_prefill): tinyllama_logits_processor_test_harness( @@ -1094,6 +1097,194 @@ def test_qwen_moe_shared_expert_lora(): llm.shutdown() +def _write_routed_expert_lora_adapter(save_dir: str, *, moe_layers: list[int], + num_experts: int, hidden_size: int, + moe_intermediate_size: int, rank: int, + lora_alpha: float, seed: int) -> None: + """Fabricate a per-expert routed-expert HF LoRA adapter on disk. + + Current transformers stores Qwen2-MoE routed experts as fused 3D parameters + (experts.gate_up_proj, experts.down_proj) rather than per-expert Linears, so + PEFT cannot emit the per-expert adapter keys the TRT-LLM loader expects. + This writes those keys directly: for every MoE layer and expert it creates + random lora_A/lora_B for gate_proj (moe_h_to_4h), up_proj (moe_gate) and + down_proj (moe_4h_to_h), keyed as .../mlp.experts.{e}.{proj}.lora_{A,B}.weight. + lora_B is non-zero so each adapter perturbs the routed-expert output. + """ + generator = torch.Generator().manual_seed(seed) + + def randn(rows, cols, std=0.02): + weight = torch.randn(rows, + cols, + generator=generator, + dtype=torch.float32) + return (weight * std).to(torch.bfloat16) + + # (projection name, in_features, out_features) for a single expert. + projections = ( + ("gate_proj", hidden_size, moe_intermediate_size), + ("up_proj", hidden_size, moe_intermediate_size), + ("down_proj", moe_intermediate_size, hidden_size), + ) + + state_dict = {} + for layer_idx in moe_layers: + prefix = f"base_model.model.model.layers.{layer_idx}.mlp.experts" + for expert_idx in range(num_experts): + for proj, in_features, out_features in projections: + key = f"{prefix}.{expert_idx}.{proj}" + state_dict[f"{key}.lora_A.weight"] = randn(rank, in_features) + state_dict[f"{key}.lora_B.weight"] = randn(out_features, rank) + + os.makedirs(save_dir, exist_ok=True) + torch.save(state_dict, os.path.join(save_dir, "adapter_model.bin")) + adapter_config = { + "peft_type": "LORA", + "r": int(rank), + "lora_alpha": float(lora_alpha), + "target_modules": ["gate_proj", "up_proj", "down_proj"], + "bias": "none", + "task_type": "CAUSAL_LM", + "use_rslora": False, + } + with open(os.path.join(save_dir, "adapter_config.json"), "w") as f: + json.dump(adapter_config, f) + + +@skip_gpu_memory_less_than_80gb +@pytest.mark.parametrize("moe_lora_mode", [ + "host_path", + "device_path_eager", + "device_path_cudagraph", +]) +def test_qwen_moe_routed_expert_multi_lora_varying_ranks( + moe_lora_mode: str, monkeypatch) -> None: + """Routed-expert MoE LoRA on Qwen1.5-MoE with the PyTorch CUTLASS backend. + + Five dummy adapters of varying rank target the routed experts (moe_h_to_4h, + moe_gate, moe_4h_to_h). The same workload runs through each of the three + routed-expert LoRA execution paths, selected by moe_lora_mode: + + - host_path: eager, legacy host path (per-request D2H pointer expand). + - device_path_eager: eager, capture-safe device path forced on via + TLLM_MOE_LORA_USE_DEVICE_PATH (per-request schema, no CUDA graph). + - device_path_cudagraph: CUDA graph decode, which always takes the + slot-indexed device path; all adapters share one captured graph, + exercising the slot-indexed device path and per-slot rank handling. + + Current transformers stores the routed experts as fused 3D parameters, so + PEFT cannot produce per-expert adapter weights; the adapters are fabricated + directly in the per-expert key layout the TRT-LLM loader expects. An + explicit module mapping is supplied because the default map only knows + w1/w2/w3 for routed experts. lora_B is non-zero so each adapter perturbs the + routed-expert output, letting the test assert the LoRA is actually applied. + """ + # Select the execution path. The eager device path is forced via env var + # (read once at FusedMoeRunner construction); the CUDA-graph path always + # takes the slot-indexed device path, so it needs no env opt-in. + cuda_graph_config = None + if moe_lora_mode == "device_path_eager": + monkeypatch.setenv("TLLM_MOE_LORA_USE_DEVICE_PATH", "1") + elif moe_lora_mode == "device_path_cudagraph": + cuda_graph_config = CudaGraphConfig(max_batch_size=10) + + model_dir = f"{llm_models_root()}/Qwen1.5-MoE-A2.7B-Chat" + + # Five adapters with varying ranks; max_lora_rank must cover the largest. + ranks = [8, 16, 32, 16, 64] + max_rank = max(ranks) + + # HF expert-projection names -> routed-expert TRT-LLM module ids. The + # default map only carries w1/w2/w3, so the gate/up/down names need an + # explicit entry. (gate_proj->w1->moe_h_to_4h, up_proj->w3->moe_gate, + # down_proj->w2->moe_4h_to_h.) + target_modules = ["moe_h_to_4h", "moe_gate", "moe_4h_to_h"] + trtllm_modules_to_hf_modules = { + "moe_h_to_4h": "gate_proj", + "moe_gate": "up_proj", + "moe_4h_to_h": "down_proj", + } + + # Derive expert dims and the set of MoE layers from the model config so the + # fabricated adapter matches the served model. + with open(f"{model_dir}/config.json") as f: + cfg = json.load(f) + num_experts = cfg["num_experts"] + hidden_size = cfg["hidden_size"] + moe_intermediate_size = cfg["moe_intermediate_size"] + num_hidden_layers = cfg["num_hidden_layers"] + decoder_sparse_step = cfg.get("decoder_sparse_step", 1) + mlp_only_layers = cfg.get("mlp_only_layers") or [] + moe_layers = [ + layer_idx for layer_idx in range(num_hidden_layers) + if layer_idx not in mlp_only_layers and num_experts > 0 and + (layer_idx + 1) % decoder_sparse_step == 0 + ] + + with tempfile.TemporaryDirectory() as lora_dir: + lora_paths = [] + for i, r in enumerate(ranks): + lora_path = f"{lora_dir}/lora_{i}" + _write_routed_expert_lora_adapter( + lora_path, + moe_layers=moe_layers, + num_experts=num_experts, + hidden_size=hidden_size, + moe_intermediate_size=moe_intermediate_size, + rank=r, + lora_alpha=2 * r, + seed=1000 + i, + ) + lora_paths.append(lora_path) + + lora_config = LoraConfig( + lora_dir=lora_paths, + lora_target_modules=target_modules, + trtllm_modules_to_hf_modules=trtllm_modules_to_hf_modules, + max_lora_rank=max_rank, + max_loras=len(ranks), + max_cpu_loras=len(ranks), + ) + llm = LLM(model=model_dir, + lora_config=lora_config, + moe_config=MoeConfig(backend="CUTLASS"), + kv_cache_config=global_kvcache_config, + cuda_graph_config=cuda_graph_config) + try: + sampling_params = SamplingParams(max_tokens=20, temperature=0.0) + prompt = "What is your name?" + + base_tokens = list( + llm.generate([prompt], sampling_params, + lora_request=None)[0].outputs[0].token_ids) + + lora_requests = [ + LoRARequest(f"moe-lora-{i}", i, path) + for i, path in enumerate(lora_paths) + ] + + # One batch mixes a no-LoRA (rank-0) request with every adapter so + # the rank-0 skip path and all adapters run through a single + # (captured, when enabled) decode graph. + requests = [None] + lora_requests + outputs = llm.generate([prompt] * len(requests), + sampling_params, + lora_request=requests) + out_tokens = [list(o.outputs[0].token_ids) for o in outputs] + + # The no-LoRA row (index 0) must run (rank-0 skip path) and produce + # output. + assert out_tokens[0], ( + "No-LoRA row in the mixed batch produced no tokens.") + # Every adapter -- not just one -- must change the output vs base. + for i in range(len(lora_requests)): + assert out_tokens[i + 1] != base_tokens, ( + f"Routed-expert MoE LoRA adapter {i} produced output " + "identical to the base model; it was not applied.") + finally: + llm.shutdown() + + class TestLlmError: @pytest.mark.part3