diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py index ca8019802db7..4e2878fbe435 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py +++ b/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( diff --git a/tests/integration/defs/disaggregated/test_disaggregated.py b/tests/integration/defs/disaggregated/test_disaggregated.py index f35b49961440..b76be7ed246e 100644 --- a/tests/integration/defs/disaggregated/test_disaggregated.py +++ b/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,8 +1095,8 @@ 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): @@ -1103,11 +1104,16 @@ def run_disaggregated_test(example_dir, 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 diff --git a/tests/unittest/_torch/attention/test_attention_mla.py b/tests/unittest/_torch/attention/test_attention_mla.py index 3ab1cfae388d..64ef045d3046 100644 --- a/tests/unittest/_torch/attention/test_attention_mla.py +++ b/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,