[None][chore] Integration tests for MoE lora & bugfixes

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<int>(out_hidden_size) : 0;
     problem_sizes_in[i] = cutlass::gemm::GemmCoord(1, rank, static_cast<int>(in_hidden_size));
-    problem_sizes_out[i] = cutlass::gemm::GemmCoord(1, static_cast<int>(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.

cpp/tests/unit_tests/kernels/moeLoraProblemBuilderTest.cu

@@ -70,8 +70,9 @@ RefOutputs cpuReference(std::vector<int32_t> const& ranks, std::vector<int64_t>
     for (int64_t i = 0; i < P; ++i)
     {
         int32_t const rank = ranks[i];
+        int const out_n = (rank > 0) ? static_cast<int>(out_hidden_size) : 0;
         r.problem_sizes_in[i] = cutlass::gemm::GemmCoord(1, rank, static_cast<int>(in_hidden_size));
-        r.problem_sizes_out[i] = cutlass::gemm::GemmCoord(1, static_cast<int>(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;

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()))

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

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

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:

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