From cf8afdded3d000e701a30f2b9a745313c42cd2ec Mon Sep 17 00:00:00 2001 From: Brian Nguyen Date: Tue, 4 Aug 2026 16:06:10 -0700 Subject: [PATCH 1/9] [None][fix] CuteDSL MLA decode: bucket the fallback tactic's batch size When the AutoTuner returns its -1 sentinel (cache miss at serving time), the op falls back to default_tactic, which derived split_kv from the raw runtime batch size. Tuning only ever profiles (and cute.compiles) the split_kv derived from each power-of-2 tuning bucket, so a raw-batch fallback almost always names a never-compiled kernel variant and JIT-compiles it inside the serving loop. Round the batch down to its tuning bucket (the same last_positive_power_of_2 mapping the tuning config uses) before deriving split_kv: a fallback on a tuned runner now reuses an already-compiled kernel, and on an untuned runner the number of distinct fallback variants is bounded by the bucket count instead of one per distinct batch size. The is_persistent choice is unchanged: its threshold (64) is a power of two, so rounding down to a power of two never crosses it. Signed-off-by: Brian Nguyen --- .../_torch/custom_ops/cute_dsl_custom_ops.py | 20 ++++++++++++++++--- 1 file changed, 17 insertions(+), 3 deletions(-) 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( From 8f4bf644b01bc3c77d2774dd23d248331a2b5dd6 Mon Sep 17 00:00:00 2001 From: Brian Nguyen Date: Tue, 4 Aug 2026 16:06:39 -0700 Subject: [PATCH 2/9] [None][test] CuteDSL MLA decode: cover the AutoTuner tuning path test_attention_mla runs with the autotuner off, so the CuTe DSL MLA decode op only ever exercises its default_tactic (-1) branch. Add a tuning-mode test on the fp8-KV DeepSeek decode geometry that asserts: - a tuning-mode pass profiles the op and both tactic elements the tuner owns (split_kv and both is_persistent candidates are compiled), and - a subsequent serving-mode pass reuses the tuned kernels with no new runtime cute.compile (which would stall the serving loop), while matching the reference output. The l0_b200 list already collects unittest/_torch/attention as a directory, so the new test runs in pre-merge B200 CI without a test-list change. Signed-off-by: Brian Nguyen --- .../_torch/attention/test_attention_mla.py | 98 +++++++++++++++++++ 1 file changed, 98 insertions(+) diff --git a/tests/unittest/_torch/attention/test_attention_mla.py b/tests/unittest/_torch/attention/test_attention_mla.py index 3ab1cfae388d..3c35d1931401 100644 --- a/tests/unittest/_torch/attention/test_attention_mla.py +++ b/tests/unittest/_torch/attention/test_attention_mla.py @@ -591,6 +591,104 @@ 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): + """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). + """ + 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) + ctx_lens = [10] * 64 + 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(): + # 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, ctx_lens, 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)}" + ) + + 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, From b2ede2836613d94addb7a391a897f0a8f1d1eb38 Mon Sep 17 00:00:00 2001 From: Brian Nguyen Date: Tue, 4 Aug 2026 16:13:29 -0700 Subject: [PATCH 3/9] [None][test] Add disagg decode-only smoke for the CuTe DSL MLA FMHA lib A disaggregated generation server runs decode-only batches, so the decode-only CuTe DSL MLA lib takes essentially every forward there, yet no disagg test covered it and the only off switch (TLLM_FMHA_LIBS) is unset in every checked-in disagg config. Add one smoke: DeepSeek-V3-Lite bf16 on a ctxTP1+genTP2 cluster (gen TP2 yields the 16 heads/rank the bf16 path admits at any batch size), asserting client output and the lib's kernel-compile marker in a generation-worker log so a silent fallback to the next FMHA lib fails the test. Signed-off-by: Brian Nguyen --- ...gentp2_deepseek_v3_lite_bf16_cute_dsl.yaml | 21 ++++++++++ .../defs/disaggregated/test_disaggregated.py | 40 +++++++++++++++++++ 2 files changed, 61 insertions(+) create mode 100644 tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp1_gentp2_deepseek_v3_lite_bf16_cute_dsl.yaml diff --git a/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp1_gentp2_deepseek_v3_lite_bf16_cute_dsl.yaml b/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp1_gentp2_deepseek_v3_lite_bf16_cute_dsl.yaml new file mode 100644 index 000000000000..1bc8e327ffb6 --- /dev/null +++ b/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp1_gentp2_deepseek_v3_lite_bf16_cute_dsl.yaml @@ -0,0 +1,21 @@ +hostname: localhost +model: DeepSeek-V3-Lite/bf16 +free_gpu_memory_fraction: 0.1 +backend: pytorch +cuda_graph_config: null +disable_overlap_scheduler: true +context_servers: + num_instances: 1 + tensor_parallel_size: 1 + pipeline_parallel_size: 1 + cache_transceiver_config: + backend: DEFAULT +# gen TP2 gives 16 attention heads per rank: the CuTe DSL MLA decode FMHA +# lib admits the bf16 path only for exactly 16 heads, so this is the +# smallest disagg layout on which the lib takes the decode forwards. +generation_servers: + num_instances: 1 + tensor_parallel_size: 2 + pipeline_parallel_size: 1 + cache_transceiver_config: + backend: DEFAULT diff --git a/tests/integration/defs/disaggregated/test_disaggregated.py b/tests/integration/defs/disaggregated/test_disaggregated.py index f35b49961440..109fd1fbf196 100644 --- a/tests/integration/defs/disaggregated/test_disaggregated.py +++ b/tests/integration/defs/disaggregated/test_disaggregated.py @@ -403,6 +403,8 @@ def get_test_config(test_desc, example_dir, test_root): f"{test_configs_root}/disagg_config_ctxtp2ep2pp2_gentp4_deepseek_v3_lite_one_mtp_block_reuse_chunked.yaml", "deepseek_v3_lite_bf16_empty_batch": f"{test_configs_root}/disagg_config_deepseek_v3_lite_empty_batch.yaml", + "deepseek_v3_lite_bf16_gentp2_cute_dsl": + f"{test_configs_root}/disagg_config_ctxtp1_gentp2_deepseek_v3_lite_bf16_cute_dsl.yaml", "llama4_kv_cache_overflow": f"{test_configs_root}/disagg_config_llama4_kv_cache_overflow.yaml", "deepseek_v3_lite_bf16_tllm_gen_helix": @@ -2006,6 +2008,44 @@ def test_disaggregated_deepseek_v3_lite_fp8_tp1_single_gpu_mtp( cwd=llm_venv.get_working_directory()) +@pytest.mark.skip_less_device(3) +@pytest.mark.skipif( + get_sm_version() not in (100, 103), + reason="CuTe DSL MLA decode FMHA lib requires SM100 or SM103") +@pytest.mark.parametrize("deepseek_v3_model_root", ['DeepSeek-V3-Lite-bf16'], + indirect=True) +def test_disaggregated_deepseek_v3_lite_bf16_gentp2_cute_dsl_mla_smoke( + disaggregated_test_root, disaggregated_example_root, llm_venv, + deepseek_v3_model_root): + """Decode-only smoke for the CuTe DSL MLA decode FMHA lib in disagg. + + A disaggregated generation server runs decode-only batches, so this + decode-only lib takes essentially every forward there (vs a fraction in + aggregated serving) and has no coverage from the aggregated tests. Run a + minimal ctxTP1+genTP2 disagg cluster on DeepSeek MLA geometry (gen TP2 + yields the 16 heads/rank the bf16 path admits at any batch size) and + require the lib's kernel-compile marker in a generation-worker log: + correct client output alone would not distinguish the CuTe DSL path from + a silent fallback to flashinfer_trtllm_gen. The lib stays enabled by + default; TLLM_FMHA_LIBS=-cute_dsl_mla on the generation server is the + documented off switch. + """ + setup_model_symlink(llm_venv, deepseek_v3_model_root, + "DeepSeek-V3-Lite/bf16") + + env = llm_venv._new_env.copy() + # The kernel-compile marker is logged at INFO level. + env["TLLM_LOG_LEVEL"] = "INFO" + + run_disaggregated_test( + disaggregated_example_root, + "deepseek_v3_lite_bf16_gentp2_cute_dsl", + env=env, + model_path=deepseek_v3_model_root, + cwd=llm_venv.get_working_directory(), + assert_gen_log_contains="CuteDSL MLA decode: compiling kernel variant") + + @pytest.mark.skip_less_device(4) @skip_no_hopper @pytest.mark.parametrize("deepseek_v3_model_root", ['DeepSeek-V3-Lite-fp8'], From ee23b97bde5604f54d6879109b5fa5b800f32b64 Mon Sep 17 00:00:00 2001 From: Brian Nguyen Date: Sun, 16 Aug 2026 22:47:14 -0500 Subject: [PATCH 4/9] [TRTLLM-15030][test] register disagg CuteDSL MLA smoke in l0_dgx_b200 The disagg decode-only smoke for the CuTe DSL MLA FMHA lib was added without a test-list entry, so it never ran in pre-merge/post-merge CI. Register it in the 8-GPU B200 post_merge block next to the peer bf16 DeepSeek disagg test; the smoke needs 3 GPUs (ctxTP1+genTP2) and gates on SM100/SM103, both satisfied there. Signed-off-by: Brian Nguyen --- tests/integration/test_lists/test-db/l0_dgx_b200.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/integration/test_lists/test-db/l0_dgx_b200.yml b/tests/integration/test_lists/test-db/l0_dgx_b200.yml index 884eb3e26bbf..780aeea20a16 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_b200.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_b200.yml @@ -244,6 +244,7 @@ l0_dgx_b200: - disaggregated/test_disaggregated.py::test_disaggregated_deepseek_v3_lite_bf16_tllm_gen_helix[DeepSeek-V3-Lite-bf16-short_prompt] TIMEOUT (60) - accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype_with_helix[fifo_v2-cudagraph:with_padding-pp1dp2cp2-overlap_off] TIMEOUT (60) - accuracy/test_disaggregated_serving.py::TestQwen3_8B::test_auto_dtype_with_helix[fifo_v2-cudagraph:with_padding-pp1tp2cp2-overlap_off] TIMEOUT (60) + - disaggregated/test_disaggregated.py::test_disaggregated_deepseek_v3_lite_bf16_gentp2_cute_dsl_mla_smoke[DeepSeek-V3-Lite-bf16] TIMEOUT (60) - accuracy/test_llm_api_pytorch.py::TestDeepSeekR1::test_nvfp4_multi_gpus_corner_case TIMEOUT (60) - accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_fp8_blockscale[baseline_fp8kv] TIMEOUT (60) - accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_fp8_blockscale[latency] TIMEOUT (60) From 6d4d6cc0b9718451ea5206233423d6e0a8db3633 Mon Sep 17 00:00:00 2001 From: Brian Nguyen Date: Sun, 16 Aug 2026 21:04:46 -0700 Subject: [PATCH 5/9] Address trivial review comments Signed-off-by: Brian Nguyen --- tests/integration/defs/disaggregated/test_disaggregated.py | 2 +- tests/unittest/_torch/attention/test_attention_mla.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/integration/defs/disaggregated/test_disaggregated.py b/tests/integration/defs/disaggregated/test_disaggregated.py index 109fd1fbf196..d38579473451 100644 --- a/tests/integration/defs/disaggregated/test_disaggregated.py +++ b/tests/integration/defs/disaggregated/test_disaggregated.py @@ -2016,7 +2016,7 @@ def test_disaggregated_deepseek_v3_lite_fp8_tp1_single_gpu_mtp( indirect=True) def test_disaggregated_deepseek_v3_lite_bf16_gentp2_cute_dsl_mla_smoke( disaggregated_test_root, disaggregated_example_root, llm_venv, - deepseek_v3_model_root): + deepseek_v3_model_root) -> None: """Decode-only smoke for the CuTe DSL MLA decode FMHA lib in disagg. A disaggregated generation server runs decode-only batches, so this diff --git a/tests/unittest/_torch/attention/test_attention_mla.py b/tests/unittest/_torch/attention/test_attention_mla.py index 3c35d1931401..dbab4d2a37d5 100644 --- a/tests/unittest/_torch/attention/test_attention_mla.py +++ b/tests/unittest/_torch/attention/test_attention_mla.py @@ -593,7 +593,7 @@ def test_attention_mla_flashinfer(scenario: Scenario, @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): +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 @@ -642,7 +642,7 @@ def test_attention_mla_cute_dsl_autotune(v2_kv_cache: bool): model_type=scenario.model_type, ) - def run_once(): + def run_once() -> None: # Numerics vs the reference implementation are asserted inside. _run_test_for_backend("TRTLLM", scenario.num_heads, scenario.num_kv_heads, scenario.num_layers, From b011d90677e6c12f5c92d7af68d43c07f721bbcc Mon Sep 17 00:00:00 2001 From: Brian Nguyen Date: Tue, 18 Aug 2026 09:16:19 -0500 Subject: [PATCH 6/9] Exercise the batch-65 default_tactic fallback in the CuTe DSL MLA autotune test The autotune test tuned at batch 64 and re-ran at batch 64, which never executes the bucketed default_tactic fallback this PR changes (a bucket-aligned batch always hits the tuned cache). Add a serving pass at batch 65 with the AutoTuner cache cleared so choose_one returns its -1 sentinel: assert (via a default_tactic spy) that the fallback ran at batch 65, returned batch-64's tactic, and reused a tuning-compiled kernel variant (no new cute.compile), with numerics still checked against the reference inside the run. Since split_kv(65) can coincidentally equal split_kv(64) on some GPUs, a pinned-occupancy check (max_active_blocks patched to 256, where raw batch 65 and bucket 64 disagree) makes the bucketing regression detection hardware-independent. Addresses the review comment at test_attention_mla.py:624. Signed-off-by: Brian Nguyen --- .../_torch/attention/test_attention_mla.py | 65 +++++++++++++++++-- 1 file changed, 61 insertions(+), 4 deletions(-) diff --git a/tests/unittest/_torch/attention/test_attention_mla.py b/tests/unittest/_torch/attention/test_attention_mla.py index dbab4d2a37d5..64ef045d3046 100644 --- a/tests/unittest/_torch/attention/test_attention_mla.py +++ b/tests/unittest/_torch/attention/test_attention_mla.py @@ -602,8 +602,12 @@ def test_attention_mla_cute_dsl_autotune(v2_kv_cache: bool) -> None: 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). + 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 @@ -622,7 +626,6 @@ def test_attention_mla_cute_dsl_autotune(v2_kv_cache: bool) -> None: scenario = Scenario(kv_cache_dtype=torch.float8_e4m3fn, num_layers=1, kv_cache_tokens_per_block=tokens_per_block) - ctx_lens = [10] * 64 rope_config = RopeConfig( hidden_size=scenario.hidden_size, num_attention_heads=scenario.num_heads, @@ -642,7 +645,7 @@ def test_attention_mla_cute_dsl_autotune(v2_kv_cache: bool) -> None: model_type=scenario.model_type, ) - def run_once() -> None: + 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, @@ -651,7 +654,7 @@ def run_once() -> None: 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, ctx_lens, 1, 2, + scenario.kv_cache_dtype, [10] * batch_size, 1, 2, v2_kv_cache) AutoTuner.get().clear_cache() @@ -688,6 +691,60 @@ def run_once() -> None: 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, From a29a1b64764f08a6d5f475405bdaf5a91c59e1ed Mon Sep 17 00:00:00 2001 From: Brian Nguyen Date: Tue, 18 Aug 2026 09:23:00 -0500 Subject: [PATCH 7/9] [TRTLLM-15030][test] Fold disagg CuteDSL MLA smoke into the existing fp8_nixl test Reviewers asked to reuse an existing disagg test instead of adding a new cluster config (addresses @Shixiaowei02's and @reasonsolo's review). test_disaggregated_deepseek_v3_lite_fp8_nixl already runs ctxTP2+genTP2 DeepSeek-V3-Lite and is registered on l0_dgx_b200 (pre_merge), l0_dgx_b300 and l0_dgx_h100: gen TP2 yields the 16 heads/rank the CuTe DSL MLA decode lib's bf16-KV path admits at any batch size (the fp8 block-scale checkpoint keeps a bf16 KV cache), so on SM100/103 the lib takes essentially every generation forward there. Add the SM100/103-gated kernel-compile-log assertion (with INFO logging on the gen workers only) to that test, and drop the dedicated smoke test, its cluster config, and its l0_dgx_b200 registration. The test carried a stale Hopper-only @skip_no_hopper gate (duplicated, along with @skip_arm, by the consolidation in #16614) that silently skipped it on its pre-existing B200/B300 registrations. Drop the Hopper gate and dedupe @skip_arm: placement is controlled by the test lists, and NIXL disagg runs on Blackwell in CI today (e.g. TestGLM52NVFP4's test_nvfp4_nixl on l0_dgx_b200). This makes the B200/B300 registrations live; this PR's B200 pre_merge CI validates the test passes there. Net CI delta: one fewer disagg cluster spin-up in the 8-GPU B200 post_merge stage; no new configs anywhere. Signed-off-by: Brian Nguyen --- ...gentp2_deepseek_v3_lite_bf16_cute_dsl.yaml | 21 ----- .../defs/disaggregated/test_disaggregated.py | 80 +++++++------------ .../test_lists/test-db/l0_dgx_b200.yml | 1 - 3 files changed, 30 insertions(+), 72 deletions(-) delete mode 100644 tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp1_gentp2_deepseek_v3_lite_bf16_cute_dsl.yaml diff --git a/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp1_gentp2_deepseek_v3_lite_bf16_cute_dsl.yaml b/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp1_gentp2_deepseek_v3_lite_bf16_cute_dsl.yaml deleted file mode 100644 index 1bc8e327ffb6..000000000000 --- a/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp1_gentp2_deepseek_v3_lite_bf16_cute_dsl.yaml +++ /dev/null @@ -1,21 +0,0 @@ -hostname: localhost -model: DeepSeek-V3-Lite/bf16 -free_gpu_memory_fraction: 0.1 -backend: pytorch -cuda_graph_config: null -disable_overlap_scheduler: true -context_servers: - num_instances: 1 - tensor_parallel_size: 1 - pipeline_parallel_size: 1 - cache_transceiver_config: - backend: DEFAULT -# gen TP2 gives 16 attention heads per rank: the CuTe DSL MLA decode FMHA -# lib admits the bf16 path only for exactly 16 heads, so this is the -# smallest disagg layout on which the lib takes the decode forwards. -generation_servers: - num_instances: 1 - tensor_parallel_size: 2 - pipeline_parallel_size: 1 - cache_transceiver_config: - backend: DEFAULT diff --git a/tests/integration/defs/disaggregated/test_disaggregated.py b/tests/integration/defs/disaggregated/test_disaggregated.py index d38579473451..a5f5428a8841 100644 --- a/tests/integration/defs/disaggregated/test_disaggregated.py +++ b/tests/integration/defs/disaggregated/test_disaggregated.py @@ -403,8 +403,6 @@ def get_test_config(test_desc, example_dir, test_root): f"{test_configs_root}/disagg_config_ctxtp2ep2pp2_gentp4_deepseek_v3_lite_one_mtp_block_reuse_chunked.yaml", "deepseek_v3_lite_bf16_empty_batch": f"{test_configs_root}/disagg_config_deepseek_v3_lite_empty_batch.yaml", - "deepseek_v3_lite_bf16_gentp2_cute_dsl": - f"{test_configs_root}/disagg_config_ctxtp1_gentp2_deepseek_v3_lite_bf16_cute_dsl.yaml", "llama4_kv_cache_overflow": f"{test_configs_root}/disagg_config_llama4_kv_cache_overflow.yaml", "deepseek_v3_lite_bf16_tllm_gen_helix": @@ -1033,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( @@ -1096,8 +1094,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): @@ -1105,8 +1103,8 @@ 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)") finally: terminate(*ctx_workers, *gen_workers, disagg_server) shutil.rmtree(work_dir, ignore_errors=True) @@ -2008,44 +2006,6 @@ def test_disaggregated_deepseek_v3_lite_fp8_tp1_single_gpu_mtp( cwd=llm_venv.get_working_directory()) -@pytest.mark.skip_less_device(3) -@pytest.mark.skipif( - get_sm_version() not in (100, 103), - reason="CuTe DSL MLA decode FMHA lib requires SM100 or SM103") -@pytest.mark.parametrize("deepseek_v3_model_root", ['DeepSeek-V3-Lite-bf16'], - indirect=True) -def test_disaggregated_deepseek_v3_lite_bf16_gentp2_cute_dsl_mla_smoke( - disaggregated_test_root, disaggregated_example_root, llm_venv, - deepseek_v3_model_root) -> None: - """Decode-only smoke for the CuTe DSL MLA decode FMHA lib in disagg. - - A disaggregated generation server runs decode-only batches, so this - decode-only lib takes essentially every forward there (vs a fraction in - aggregated serving) and has no coverage from the aggregated tests. Run a - minimal ctxTP1+genTP2 disagg cluster on DeepSeek MLA geometry (gen TP2 - yields the 16 heads/rank the bf16 path admits at any batch size) and - require the lib's kernel-compile marker in a generation-worker log: - correct client output alone would not distinguish the CuTe DSL path from - a silent fallback to flashinfer_trtllm_gen. The lib stays enabled by - default; TLLM_FMHA_LIBS=-cute_dsl_mla on the generation server is the - documented off switch. - """ - setup_model_symlink(llm_venv, deepseek_v3_model_root, - "DeepSeek-V3-Lite/bf16") - - env = llm_venv._new_env.copy() - # The kernel-compile marker is logged at INFO level. - env["TLLM_LOG_LEVEL"] = "INFO" - - run_disaggregated_test( - disaggregated_example_root, - "deepseek_v3_lite_bf16_gentp2_cute_dsl", - env=env, - model_path=deepseek_v3_model_root, - cwd=llm_venv.get_working_directory(), - assert_gen_log_contains="CuteDSL MLA decode: compiling kernel variant") - - @pytest.mark.skip_less_device(4) @skip_no_hopper @pytest.mark.parametrize("deepseek_v3_model_root", ['DeepSeek-V3-Lite-fp8'], @@ -2104,9 +2064,6 @@ 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 @pytest.mark.parametrize("deepseek_v3_model_root", ['DeepSeek-V3-Lite-fp8'], indirect=True) @@ -2121,11 +2078,34 @@ 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" + + # No arch gate: placement is controlled by the test lists (l0_dgx_h100, + # l0_dgx_b200 pre_merge, l0_dgx_b300). A stale Hopper-only @skip_no_hopper + # used to silently skip this test on its B200/B300 registrations; dropping + # it makes them live. + # + # 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/integration/test_lists/test-db/l0_dgx_b200.yml b/tests/integration/test_lists/test-db/l0_dgx_b200.yml index 780aeea20a16..884eb3e26bbf 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_b200.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_b200.yml @@ -244,7 +244,6 @@ l0_dgx_b200: - disaggregated/test_disaggregated.py::test_disaggregated_deepseek_v3_lite_bf16_tllm_gen_helix[DeepSeek-V3-Lite-bf16-short_prompt] TIMEOUT (60) - accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype_with_helix[fifo_v2-cudagraph:with_padding-pp1dp2cp2-overlap_off] TIMEOUT (60) - accuracy/test_disaggregated_serving.py::TestQwen3_8B::test_auto_dtype_with_helix[fifo_v2-cudagraph:with_padding-pp1tp2cp2-overlap_off] TIMEOUT (60) - - disaggregated/test_disaggregated.py::test_disaggregated_deepseek_v3_lite_bf16_gentp2_cute_dsl_mla_smoke[DeepSeek-V3-Lite-bf16] TIMEOUT (60) - accuracy/test_llm_api_pytorch.py::TestDeepSeekR1::test_nvfp4_multi_gpus_corner_case TIMEOUT (60) - accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_fp8_blockscale[baseline_fp8kv] TIMEOUT (60) - accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_fp8_blockscale[latency] TIMEOUT (60) From ce8de03373d78ff3c302df25c9ae012e5ec04ea8 Mon Sep 17 00:00:00 2001 From: Brian Nguyen Date: Wed, 19 Aug 2026 07:27:28 +0000 Subject: [PATCH 8/9] Address trivial review comments Signed-off-by: Brian Nguyen --- .../integration/defs/disaggregated/test_disaggregated.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/tests/integration/defs/disaggregated/test_disaggregated.py b/tests/integration/defs/disaggregated/test_disaggregated.py index a5f5428a8841..86d062b3764d 100644 --- a/tests/integration/defs/disaggregated/test_disaggregated.py +++ b/tests/integration/defs/disaggregated/test_disaggregated.py @@ -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}" @@ -1105,9 +1106,14 @@ def run_disaggregated_test(example_dir, f"expected marker {assert_gen_log_contains!r} in a generation-worker log, " 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'], @@ -2065,6 +2071,7 @@ def test_disaggregated_deepseek_v3_lite_fp8_ctxtp2ep2pp2_gentp4_one_mtp_block_re @skip_arm +@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, From 0fb5545953211058dc7948142501200bd1d0080e Mon Sep 17 00:00:00 2001 From: Brian Nguyen Date: Thu, 20 Aug 2026 09:57:37 -0500 Subject: [PATCH 9/9] Gate fp8_nixl disagg test with skip_pre_hopper The prior change dropped the stale Hopper-only @skip_no_hopper from test_disaggregated_deepseek_v3_lite_fp8_nixl to unblock its B200/B300 registrations, but that left the test runnable on pre-Hopper hardware. Add @skip_pre_hopper (SM >= 90) so Hopper and Blackwell stay live while pre-Hopper is gated out. Signed-off-by: Brian Nguyen --- .../defs/disaggregated/test_disaggregated.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/tests/integration/defs/disaggregated/test_disaggregated.py b/tests/integration/defs/disaggregated/test_disaggregated.py index 86d062b3764d..b76be7ed246e 100644 --- a/tests/integration/defs/disaggregated/test_disaggregated.py +++ b/tests/integration/defs/disaggregated/test_disaggregated.py @@ -2071,6 +2071,7 @@ def test_disaggregated_deepseek_v3_lite_fp8_ctxtp2ep2pp2_gentp4_one_mtp_block_re @skip_arm +@skip_pre_hopper @pytest.mark.skip_less_device(4) @pytest.mark.parametrize("deepseek_v3_model_root", ['DeepSeek-V3-Lite-fp8'], indirect=True) @@ -2086,10 +2087,11 @@ def test_disaggregated_deepseek_v3_lite_fp8_nixl(disaggregated_test_root, env["UCX_TLS"] = get_ucx_tls() env["UCX_MM_ERROR_HANDLING"] = "y" - # No arch gate: placement is controlled by the test lists (l0_dgx_h100, - # l0_dgx_b200 pre_merge, l0_dgx_b300). A stale Hopper-only @skip_no_hopper - # used to silently skip this test on its B200/B300 registrations; dropping - # it makes them live. + # @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,