remove hardcoding

Author: brb-nv

examples/llm-api/quickstart_advanced.py

@@ -70,7 +70,6 @@ def add_llm_args(parser):
                         choices=["auto", "TorchSampler", "TRTLLMSampler"])
     parser.add_argument('--tp_size', type=int, default=1)
     parser.add_argument('--pp_size', type=int, default=1)
-    parser.add_argument('--cp_size', type=int, default=1)
     parser.add_argument('--moe_ep_size', type=int, default=-1)
     parser.add_argument('--moe_tp_size', type=int, default=-1)
     parser.add_argument('--moe_cluster_size', type=int, default=-1)
@@ -260,7 +259,6 @@ def setup_llm(args, **kwargs):
         attention_dp_config=attention_dp_config,
         tensor_parallel_size=args.tp_size,
         pipeline_parallel_size=args.pp_size,
-        context_parallel_size=args.cp_size,
         moe_expert_parallel_size=args.moe_ep_size,
         moe_tensor_parallel_size=args.moe_tp_size,
         moe_cluster_parallel_size=args.moe_cluster_size,

tensorrt_llm/commands/serve.py

@@ -130,8 +130,6 @@ def get_llm_args(
         except KeyError:
             raise ValueError(f"Invalid cp_type: {cp_config['cp_type']}. " \
                              f"Must be one of: {', '.join([t.name for t in CpType])}")
-        if cp_config["cp_type"] == CpType.HELIX:
-            cp_config['tokens_per_block'] = kv_cache_config.tokens_per_block
 
     llm_args = {
         "model": model,

tensorrt_llm/mapping.py

@@ -469,11 +469,6 @@ def __init__(
             attn_cp_size=-1,
             enable_attention_dp=False,
             enable_lm_head_tp_in_adp=False):
-        #################################################################
-        # TODO: Remove this hardcoding and obtain cp_config from llm_args.
-        if cp_size > 1:
-            cp_config = {"cp_type": CpType.HELIX}
-        #################################################################
         super().__init__(world_size=world_size,
                          rank=rank,
                          gpus_per_node=gpus_per_node,

tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2_gentp1cp2_deepseek_v3_lite_bf16_tllm_gen.yaml

@@ -11,6 +11,7 @@ context_servers:
   kv_cache_config:
     enable_block_reuse: False
     enable_partial_reuse: False
+    tokens_per_block: 32
   tensor_parallel_size: 2
   pipeline_parallel_size: 1
   cache_transceiver_config:
@@ -23,9 +24,13 @@ generation_servers:
   pipeline_parallel_size: 1
   context_parallel_size: 2
   enable_chunked_prefill: False
+  cp_config:
+    cp_type: HELIX
+    tokens_per_block: 32
   kv_cache_config:
     enable_block_reuse: False
     enable_partial_reuse: False
+    tokens_per_block: 32
   cache_transceiver_config:
     backend: "UCX"
   urls: