[TRTLLM-15030][fix] CuteDSL MLA decode follow-ups: bucket AutoTuner fallback, autotune + disagg tests

tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py

@@ -9417,13 +9417,27 @@ def default_tactic(
         ) -> Tuple[Tuple[int, int], Tuple[int, int], int, bool]:
             """Fallback 4-tuple tactic ``(mma_qk, mma_pv, split_kv,
             is_persistent)`` for when the AutoTuner cache is not warmed and
-            ``choose_one`` returns its ``-1`` sentinel."""
+            ``choose_one`` returns its ``-1`` sentinel.
+
+            ``batch_size`` is rounded down to its tuning bucket
+            (``last_positive_power_of_2`` -- the same mapping the tuning
+            config uses) before deriving ``split_kv``: tuning profiles (and
+            therefore ``cute.compile``s) exactly the bucket-derived
+            ``split_kv`` variants, so a bucket-aligned fallback reuses an
+            already-compiled kernel where one exists instead of JIT-compiling
+            a fresh raw-batch ``split_kv`` variant in the serving loop. The
+            ``is_persistent`` choice is unaffected by the rounding (its
+            threshold is a power of two, so rounding down to a power of two
+            never crosses it), and both candidates are compiled during tuning
+            anyway."""
             mma_qk_tiler_mn = (128, 128)
             mma_pv_tiler_mn = (128, 256)
             max_active_blocks = self._get_max_active_blocks()
-            split_kv = self.get_default_split_kv(batch_size, self.seq_len_q,
+            bucketed_batch_size = last_positive_power_of_2(batch_size)
+            split_kv = self.get_default_split_kv(bucketed_batch_size,
+                                                 self.seq_len_q,
                                                  max_active_blocks)
-            is_persistent = self.get_default_is_persistent(batch_size)
+            is_persistent = self.get_default_is_persistent(bucketed_batch_size)
             return (mma_qk_tiler_mn, mma_pv_tiler_mn, split_kv, is_persistent)
 
         def forward(

tests/integration/defs/disaggregated/test_disaggregated.py

@@ -1031,8 +1031,8 @@ def run_disaggregated_test(example_dir,
     """Run disaggregated test using service discovery instead of MPI.
 
     If assert_gen_log_contains is set, the generation-worker logs are captured and, after the
-    client tests, at least one of them must contain that substring (used to prove the KV-cache
-    bounce path actually engaged instead of silently falling back to the per-fragment path).
+    client tests, at least one of them must contain that substring (used to prove an intended
+    code path actually engaged instead of silently falling back to another one).
     """
     if mpi_disabled():
         pytest.skip(
@@ -1061,6 +1061,7 @@ def run_disaggregated_test(example_dir,
 
     server_host = config.get("hostname", "localhost")
 
+    success = False
     try:
         server_url = f"http://{server_host}:{server_port}"
 
@@ -1094,20 +1095,25 @@ def run_disaggregated_test(example_dir,
         if post_client_test is not None:
             post_client_test(server_url)
         if assert_gen_log_contains is not None:
-            # Fail loudly if the marker is absent: the transfer silently fell back to the
-            # per-fragment path, so the bounce path we meant to exercise never ran.
+            # Fail loudly if the marker is absent: the code path the test means to
+            # exercise never ran and something else silently took its place.
             logs = []
             for w in gen_workers:
                 if w.log_path and os.path.exists(w.log_path):
                     with open(w.log_path, 'r', errors='replace') as f:
                         logs.append(f.read())
             assert any(assert_gen_log_contains in log for log in logs), (
                 f"expected marker {assert_gen_log_contains!r} in a generation-worker log, "
-                f"but none of {len(logs)} log(s) contained it (bounce did not engage)"
-            )
+                f"but none of {len(logs)} log(s) contained it "
+                f"(the intended code path did not engage)")
+        success = True
     finally:
         terminate(*ctx_workers, *gen_workers, disagg_server)
-        shutil.rmtree(work_dir, ignore_errors=True)
+        # When the marker assertion is active the worker logs are file-based
+        # (save_log=True). Preserve work_dir on the failure path so the first
+        # failures on the newly enabled stages arrive with logs to read.
+        if success or assert_gen_log_contains is None:
+            shutil.rmtree(work_dir, ignore_errors=True)
 
 
 @pytest.mark.parametrize("llama_model_root", ['TinyLlama-1.1B-Chat-v1.0'],
@@ -2064,10 +2070,9 @@ def test_disaggregated_deepseek_v3_lite_fp8_ctxtp2ep2pp2_gentp4_one_mtp_block_re
         cwd=llm_venv.get_working_directory())
 
 
-@skip_no_hopper
-@skip_arm
-@skip_no_hopper
 @skip_arm
+@skip_pre_hopper
+@pytest.mark.skip_less_device(4)
 @pytest.mark.parametrize("deepseek_v3_model_root", ['DeepSeek-V3-Lite-fp8'],
                          indirect=True)
 def test_disaggregated_deepseek_v3_lite_fp8_nixl(disaggregated_test_root,
@@ -2081,11 +2086,35 @@ def test_disaggregated_deepseek_v3_lite_fp8_nixl(disaggregated_test_root,
     env["TRTLLM_USE_NIXL_KVCACHE"] = "1"
     env["UCX_TLS"] = get_ucx_tls()
     env["UCX_MM_ERROR_HANDLING"] = "y"
+
+    # @skip_pre_hopper (SM >= 90), not @skip_no_hopper (SM == 90): placement is
+    # controlled by the test lists (l0_dgx_h100, l0_dgx_b200 pre_merge,
+    # l0_dgx_b300), which cover Hopper and Blackwell. The old Hopper-only
+    # @skip_no_hopper silently skipped this test on its B200/B300 registrations;
+    # skip_pre_hopper keeps those live while still gating out pre-Hopper.
+    #
+    # On SM100/103 this test doubles as the decode-only smoke for the CuTe DSL
+    # MLA decode FMHA lib: a disagg generation server runs decode-only batches,
+    # and gen TP2 yields the 16 heads/rank the lib's bf16-KV path admits at any
+    # batch size (the fp8 checkpoint keeps a bf16 KV cache), so the lib takes
+    # essentially every gen forward. Require its kernel-compile marker (logged
+    # at INFO) in a generation-worker log: correct client output alone would
+    # not distinguish the CuTe DSL path from a silent fallback to another FMHA
+    # library. TLLM_FMHA_LIBS=-cute_dsl_mla on the generation server is the
+    # documented off switch.
+    gen_env = None
+    assert_gen_log_contains = None
+    if get_sm_version() in (100, 103):
+        gen_env = {"TLLM_LOG_LEVEL": "INFO"}
+        assert_gen_log_contains = "CuteDSL MLA decode: compiling kernel variant"
+
     run_disaggregated_test(disaggregated_example_root,
                            "deepseek_v3_lite_fp8_nixl",
                            env=env,
+                           gen_env=gen_env,
                            model_path=deepseek_v3_model_root,
-                           cwd=llm_venv.get_working_directory())
+                           cwd=llm_venv.get_working_directory(),
+                           assert_gen_log_contains=assert_gen_log_contains)
 
 
 @skip_no_hopper

tests/unittest/_torch/attention/test_attention_mla.py

@@ -591,6 +591,161 @@ def test_attention_mla_flashinfer(scenario: Scenario,
                           v2_kv_cache)
 
 
+@pytest.mark.parametrize("v2_kv_cache", [True, False],
+                         ids=["v2_kv_cache", "v1_kv_cache"])
+def test_attention_mla_cute_dsl_autotune(v2_kv_cache: bool) -> None:
+    """Cover the CuTe DSL MLA decode AutoTuner path.
+
+    The plain test_attention_mla runs with the autotuner off, so the op
+    always takes the ``default_tactic`` (-1 sentinel) branch. This test
+    drives the tuning path instead: a tuning-mode pass must profile the
+    tactic space (split_kv and is_persistent tactic elements), and a
+    subsequent serving-mode pass must reuse the tuned kernels without
+    triggering any runtime ``cute.compile`` (a compile outside the tuning
+    window stalls the serving loop). A final pass forces the ``choose_one``
+    -1 sentinel at a non-power-of-2 batch and checks that the
+    ``default_tactic`` fallback reuses a bucket-aligned compiled kernel.
+    """
+    from unittest import mock
+
+    from tensorrt_llm._torch.autotuner import AutoTuner, autotune
+    from tensorrt_llm._torch.cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE
+    from tensorrt_llm._utils import get_sm_version
+
+    if get_sm_version() not in (100, 103):
+        pytest.skip("CuTe DSL MLA decode requires SM100 or SM103")
+    if not IS_CUTLASS_DSL_AVAILABLE:
+        pytest.skip("nvidia-cutlass-dsl is not installed")
+
+    from tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops import \
+        CuteDSLNVMlaDecodeBlackwellRunner
+
+    # fp8 KV with (num_heads=128, seq_len_q=1) is admitted by the CuTe DSL
+    # perf gate from batch_size >= 64, so a 64-request decode batch routes
+    # the generation phase through cute_dsl_mla in the default lib order.
+    scenario = Scenario(kv_cache_dtype=torch.float8_e4m3fn,
+                        num_layers=1,
+                        kv_cache_tokens_per_block=tokens_per_block)
+    rope_config = RopeConfig(
+        hidden_size=scenario.hidden_size,
+        num_attention_heads=scenario.num_heads,
+        rope_scaling={
+            "beta_fast": scenario.rope_beta_fast,
+            "beta_slow": scenario.rope_beta_slow,
+            "factor": scenario.rope_factor,
+            "mscale": scenario.rope_mscale,
+            "mscale_all_dim": scenario.rope_mscale_all_dim,
+            "original_max_position_embeddings":
+            scenario.rope_original_max_position_embeddings,
+            "type": scenario.rope_type,
+        },
+        max_position_embeddings=scenario.max_position_embeddings,
+        rope_theta=scenario.rope_theta,
+        qk_rope_head_dim=scenario.qk_rope_head_dim,
+        model_type=scenario.model_type,
+    )
+
+    def run_once(batch_size: int = 64) -> None:
+        # Numerics vs the reference implementation are asserted inside.
+        _run_test_for_backend("TRTLLM", scenario.num_heads,
+                              scenario.num_kv_heads, scenario.num_layers,
+                              scenario.q_lora_rank, scenario.kv_lora_rank,
+                              scenario.qk_nope_head_dim,
+                              scenario.qk_rope_head_dim, scenario.v_head_dim,
+                              rope_config, scenario.kv_cache_tokens_per_block,
+                              torch.device('cuda'), scenario.dtype,
+                              scenario.kv_cache_dtype, [10] * batch_size, 1, 2,
+                              v2_kv_cache)
+
+    AutoTuner.get().clear_cache()
+    CuteDSLNVMlaDecodeBlackwellRunner.kernel_cache.clear()
+
+    with autotune():
+        run_once()
+
+    tuned_ops = {key[0] for key in AutoTuner.get().profiling_cache.cache}
+    assert any("cute_dsl_mla_decode" in str(op) for op in tuned_ops), (
+        f"tuning-mode pass did not tune any cute_dsl_mla_decode op; "
+        f"tuned ops: {tuned_ops}")
+
+    kernel_keys = list(CuteDSLNVMlaDecodeBlackwellRunner.kernel_cache)
+    assert kernel_keys, "tuning-mode pass compiled no CuTe DSL MLA kernels"
+    # Tactic layout: unique_id + (out_dtype, mma_qk, mma_pv, split_kv,
+    # is_persistent); both tactic elements chosen by the tuner must have
+    # been exercised during profiling.
+    persistent_variants = {key[-1] for key in kernel_keys}
+    assert persistent_variants == {
+        True, False
+    }, (f"expected both is_persistent tactic candidates to be profiled, "
+        f"got {persistent_variants}")
+    split_kv_variants = {key[-2] for key in kernel_keys}
+    assert split_kv_variants, "no split_kv tactic variant was profiled"
+
+    # Serving-mode pass: tuned tactics must be reused as-is -- any new
+    # kernel_cache entry means a runtime cute.compile happened post-tuning.
+    num_compiled = len(CuteDSLNVMlaDecodeBlackwellRunner.kernel_cache)
+    run_once()
+    assert len(CuteDSLNVMlaDecodeBlackwellRunner.kernel_cache) == \
+        num_compiled, (
+        "serving-mode pass cute.compiled new kernel variants after tuning: "
+        f"{set(CuteDSLNVMlaDecodeBlackwellRunner.kernel_cache) - set(kernel_keys)}"
+    )
+
+    # Fallback pass: serve a non-power-of-2 batch (65) with the AutoTuner
+    # cache cleared, so ``choose_one`` misses and returns its -1 sentinel and
+    # the op must take the ``default_tactic`` branch. The fallback rounds the
+    # batch down to its tuning bucket (64), so it must land on a kernel
+    # variant the tuning pass already compiled; deriving split_kv from the
+    # raw batch could cute.compile a fresh variant in the serving loop.
+    AutoTuner.get().clear_cache()
+    fallback_calls = []
+    orig_default_tactic = CuteDSLNVMlaDecodeBlackwellRunner.default_tactic
+
+    def spying_default_tactic(self, batch_size: int):
+        tactic = orig_default_tactic(self, batch_size)
+        fallback_calls.append((self, batch_size, tactic))
+        return tactic
+
+    with mock.patch.object(CuteDSLNVMlaDecodeBlackwellRunner, "default_tactic",
+                           spying_default_tactic):
+        run_once(batch_size=65)
+
+    assert fallback_calls, (
+        "batch-65 serving pass never reached default_tactic: with the "
+        "AutoTuner cache cleared, choose_one must miss and return its -1 "
+        "sentinel")
+    runner = fallback_calls[0][0]
+    for _, batch_size, tactic in fallback_calls:
+        assert batch_size == 65, (
+            f"default_tactic saw batch {batch_size}, expected 65")
+        assert tactic == runner.default_tactic(64), (
+            f"batch-65 fallback tactic {tactic} does not match batch-64's "
+            f"{runner.default_tactic(64)}")
+    assert len(CuteDSLNVMlaDecodeBlackwellRunner.kernel_cache) == \
+        num_compiled, (
+        "batch-65 default_tactic fallback cute.compiled new kernel variants "
+        "instead of reusing a bucket-64 one from tuning: "
+        f"{set(CuteDSLNVMlaDecodeBlackwellRunner.kernel_cache) - set(kernel_keys)}"
+    )
+
+    # The assertions above can hold even for an unbucketed fallback when this
+    # GPU's occupancy makes split_kv(65) == split_kv(64), so also check the
+    # bucketing itself with the occupancy ceiling pinned to a value where the
+    # raw batch and its bucket disagree: 256 // 65 // 2 == 1, while bucket 64
+    # gives 256 // 64 // 2 == 2.
+    with mock.patch.object(CuteDSLNVMlaDecodeBlackwellRunner,
+                           "_cute_dsl_max_active_blocks",
+                           256,
+                           create=True):
+        pinned_tactic = runner.default_tactic(65)
+        assert pinned_tactic == runner.default_tactic(64), (
+            f"default_tactic no longer buckets the batch: got {pinned_tactic} "
+            f"for batch 65 vs {runner.default_tactic(64)} for batch 64")
+        assert pinned_tactic[2] == 2, (
+            f"expected bucket-64 split_kv 2 with max_active_blocks pinned to "
+            f"256, got {pinned_tactic[2]}")
+
+
 def _run_test_for_backend(backend_name, num_heads, num_kv_heads, num_layers,
                           q_lora_rank, kv_lora_rank, qk_nope_head_dim,
                           qk_rope_head_dim, v_head_dim, rope_config,