Add asynchronous prefill scheduling

.github/oncall_schedule.json

@@ -1,6 +1,6 @@
 [
     {
-        "user": "dimapihtar",
+        "user": "Phlip79",
         "date": "2026-07-15"
     },
     {

.gitlab-ci.yml

@@ -154,6 +154,7 @@ stages:
   - integration_tests
   - functional_tests
   - publish
+  - triage
 
 default:
   interruptible: true
@@ -268,7 +269,21 @@ variables:
       - "upgrade-dependencies"
     description: Type of publish (freeze or final release)
 
+  RUN_LINEAR_STATUS:
+    value: "True"
+    options:
+      - "True"
+      - "False"
+    description: Reconcile functional-test failures against Linear
+  RUN_LINEAR_WRITE:
+    value: "True"
+    options:
+      - "True"
+      - "False"
+    description: Apply proposed Linear issue opens, updates, and closes
+
   # CI wide variables
+  NEMO_CI_TRIAGE_CONFIG: .gitlab/nemo-ci-triage.yml
   CI_MCORE_LTS_IMAGE: ${GITLAB_ENDPOINT}:5005/adlr/megatron-lm/mcore_ci_lts
   CI_MCORE_DEV_IMAGE: ${GITLAB_ENDPOINT}:5005/adlr/megatron-lm/mcore_ci_dev
   CI_NEMO_IMAGE: ${GITLAB_ENDPOINT}:5005/adlr/megatron-lm/nemo_ci
@@ -282,3 +297,4 @@ include:
   - .gitlab/stages/03.integration-tests.yml
   - .gitlab/stages/04.functional-tests.yml
   - .gitlab/stages/05.publish.yml
+  - .gitlab/stages/06.triage.yml

.gitlab/nemo-ci-triage.yml

@@ -0,0 +1,14 @@
+# Megatron-LM configuration for nemo-ci-triage.
+
+gitlab:
+  project_id: 19378
+  repo_name: ADLR/megatron-lm
+
+modules:
+  megatron_lm:
+    build_module: megatron-lm
+    team_key: MCORE
+    project_template: "MCore CI Testing"
+    enable_linear_open: true
+    enable_linear_modify: true
+    enable_linear_close: true

.gitlab/stages/02.test.yml

@@ -195,8 +195,9 @@ test:unit_tests_notify:
       fi
     - export RO_API_TOKEN=${PROJECT_ACCESS_TOKEN_MCORE}
     - export GITLAB_ENDPOINT
-    - export TAG_TEAM=$([[ "$CI_COMMIT_BRANCH" == "main" ]] && echo "1" || "0")
+    - export TAG_TEAM=$([[ "$CI_COMMIT_BRANCH" == "main" ]] && echo "1" || echo "0")
     - export TEAM_SLUG=$SLACK_ADMIN
+    - export PYTHONPATH=$(pwd)
     - |
       python tests/test_utils/python_scripts/notify.py \
         --pipeline-id "${CI_PIPELINE_ID}" \

.gitlab/stages/04.functional-tests.yml

@@ -450,6 +450,7 @@ functional:smoke_notify:
       fi
     - export RO_API_TOKEN=${PROJECT_ACCESS_TOKEN_MCORE}
     - export GITLAB_ENDPOINT
+    - export PYTHONPATH=$(pwd)
     - |
       python tests/test_utils/python_scripts/notify.py \
         --pipeline-id "${CI_PIPELINE_ID}" \
@@ -494,19 +495,28 @@ functional:x_notify:
     - export RO_API_TOKEN=${PROJECT_ACCESS_TOKEN_MCORE}
     - export GITLAB_ENDPOINT
     - export CONTEXT=$FUNCTIONAL_TEST_SCOPE
-    - export TAG_TEAM=$([[ "$CI_COMMIT_BRANCH" == "main" ]] && echo "1" || "0")
+    - export TAG_TEAM=$([[ "$CI_COMMIT_BRANCH" == "main" ]] && echo "1" || echo "0")
     - export TEAM_SLUG=$SLACK_ADMIN
+    - export PYTHONPATH=$(pwd)
     - |
       python tests/test_utils/python_scripts/notify.py \
         --pipeline-id "${CI_PIPELINE_ID}" \
         --check-for functional-tests \
         --pipeline-context $CONTEXT \
-        --pipeline-created-at "${CI_PIPELINE_CREATED_AT}"
+        --pipeline-created-at "${CI_PIPELINE_CREATED_AT}" \
+        --summary-output pipeline_summaries.json \
+        --failure-buckets-output failure_buckets.json \
+        --slack-output slack_notification.json
 
   artifacts:
     when: always
     paths:
       - scripts
+      - pipeline_summaries.json
+      - failure_buckets.json
+      - slack_notification.json
+      - inference_metrics.json
+      - agent_formatter_debug.txt
   rules:
     - if: ($CI_PIPELINE_SOURCE == "schedule" || $CI_COMMIT_BRANCH == "main") && $FUNCTIONAL_TEST == "yes"
       when: always

.gitlab/stages/06.triage.yml

@@ -0,0 +1,95 @@
+.linear_reconcile_rules:
+  rules:
+    - if: >-
+        ($CI_PIPELINE_SOURCE == "schedule" || $CI_COMMIT_BRANCH == "main") &&
+        $FUNCTIONAL_TEST == "yes" &&
+        $RUN_LINEAR_STATUS == "True"
+      when: always
+    - when: never
+
+.linear_triage_job:
+  stage: triage
+  image: ${UTILITY_IMAGE}:${CI_PIPELINE_ID}
+  tags:
+    - arch/amd64
+    - env/prod
+    - origin/jet-fleet
+    - owner/jet-core
+    - purpose/utility
+    - team/megatron
+
+triage:linear_reconcile:
+  extends: [.linear_triage_job, .linear_reconcile_rules]
+  needs:
+    - job: functional:x_notify
+      artifacts: true
+  script:
+    - >-
+      nemo-ci-linear status
+      --config "${NEMO_CI_TRIAGE_CONFIG}"
+      --build-module-regex '^megatron-lm$'
+      --output linear_status_report.json
+    - >-
+      nemo-ci-linear reconcile
+      --failure-buckets failure_buckets.json
+      --linear-report linear_status_report.json
+      --pipeline-summaries pipeline_summaries.json
+      --output linear_action_plan.json
+  artifacts:
+    when: always
+    paths:
+      - linear_status_report.json
+      - linear_action_plan.json
+      - inference_metrics.json
+
+triage:linear_write:
+  extends: [.linear_triage_job]
+  needs:
+    - job: triage:linear_reconcile
+      artifacts: true
+  allow_failure: true
+  script:
+    - >-
+      nemo-ci-linear write
+      --config "${NEMO_CI_TRIAGE_CONFIG}"
+      --plan linear_action_plan.json
+      --output linear_action_plan_post.json
+  artifacts:
+    when: always
+    paths:
+      - linear_action_plan.json
+      - linear_action_plan_post.json
+      - linear_status_report.json
+  rules:
+    - if: >-
+        ($CI_PIPELINE_SOURCE == "schedule" || $CI_COMMIT_BRANCH == "main") &&
+        $FUNCTIONAL_TEST == "yes" &&
+        $RUN_LINEAR_STATUS == "True" &&
+        $RUN_LINEAR_WRITE == "True"
+      when: always
+    - when: never
+
+triage:slack_linear_followup:
+  extends: [.linear_triage_job]
+  needs:
+    - job: functional:x_notify
+      artifacts: true
+    - job: triage:linear_write
+      artifacts: true
+  allow_failure: true
+  script:
+    - >-
+      nemo-ci-notify
+      --pipeline-summary slack_notification.json
+      --linear-plan linear_action_plan_post.json
+      --slack-bot-token "${MCORE_SLACK_BOT_TOKEN:-${ALERTMANAGER_TOKEN}}"
+      --slack-channel-id "${MCORE_SLACK_CHANNEL_ID}"
+  rules:
+    # Post the applied Linear actions under the functional-test notification.
+    - if: >-
+        ($CI_PIPELINE_SOURCE == "schedule" || $CI_COMMIT_BRANCH == "main") &&
+        $FUNCTIONAL_TEST == "yes" &&
+        $RUN_LINEAR_STATUS == "True" &&
+        $RUN_LINEAR_WRITE == "True"
+      when: always
+    - when: never

docker/Dockerfile.linting

@@ -24,7 +24,7 @@ RUN --mount=type=secret,id=JET_INDEX_URLS \
 
 # Keep this in the internal-only stage so public CI has no internal service dependency.
 ARG CI_SERVER_URL
-ARG NEMO_CI_TRIAGE_COMMIT=8e65fa4ae20b58578d0e0f20ebea37ee7d92c8ea
+ARG NEMO_CI_TRIAGE_COMMIT=5474f95417758c76c75523ae5319727a1e703437
 RUN --mount=type=secret,id=NEMO_CI_TRIAGE_TOKEN \
       GIT_CONFIG_COUNT=1 \
       GIT_CONFIG_KEY_0=http.extraHeader \

examples/inference/advanced/gpt_dynamic_inference.py

@@ -119,7 +119,7 @@ def _add_request():
         nonlocal num_requests_added
         _request = requests[num_requests_added]
         engine.add_request(num_requests_added, _request.prompt_text, _request.sampling_params)
-        _request.time_start = get_curr_time()
+        _request.time_start = get_curr_time(do_broadcast=False)
         _request.state = "started"
         num_requests_added += 1
         tbar.update(1)
@@ -148,14 +148,14 @@ def _process_step_result(result):
                 step_times["prefill"].append(step_time)
 
             # Append output tokens.
-            output_start = get_curr_time()
+            output_start = get_curr_time(do_broadcast=False)
             for finished_request_record in finished_request_records:
 
                 finished_request = finished_request_record.merge()
 
                 # Update local request object.
                 request = requests[finished_request.request_id]
-                request.time_end = get_curr_time()
+                request.time_end = get_curr_time(do_broadcast=False)
                 request.state = "finished"
                 request.request_id = finished_request.request_id
                 request.events = finished_request.events
@@ -185,16 +185,16 @@ def _process_step_result(result):
                 if not finished_request.sampling_params.skip_prompt_log_probs:
                     request.prompt_top_n_logprobs = finished_request.prompt_top_n_logprobs
                 num_requests_finished += 1
-            output_times.append(get_curr_time() - output_start)
+            output_times.append(get_curr_time(do_broadcast=False) - output_start)
 
     if batch_ranges is not None:
         # Batch-drain mode: add all requests in a batch, drain, then next batch.
         for batch_idx, (batch_start, batch_end) in enumerate(batch_ranges):
             # Add all requests in current batch.
-            add_start = get_curr_time()
+            add_start = get_curr_time(do_broadcast=False)
             while num_requests_added < batch_end:
                 _add_request()
-            add_times.append(get_curr_time() - add_start)
+            add_times.append(get_curr_time(do_broadcast=False) - add_start)
 
             # Step until all active requests finish (drain).
             while engine.has_unfinished_requests():
@@ -212,7 +212,7 @@ def _process_step_result(result):
         # Original mode: add requests per step based on arrival time or count.
         while True:
             # Add requests.
-            add_start = get_curr_time()
+            add_start = get_curr_time(do_broadcast=False)
             if args.incoming_requests_per_step is None:
                 # Add requests with 'earlier' arrival time.
                 while num_requests_added < num_requests_total:
@@ -225,7 +225,7 @@ def _process_step_result(result):
                     min(args.incoming_requests_per_step, num_requests_total - num_requests_added)
                 ):
                     _add_request()
-            add_times.append(get_curr_time() - add_start)
+            add_times.append(get_curr_time(do_broadcast=False) - add_start)
 
             # Step inference engine (i.e., generate a token for each active request).
             # Before step, we haven't done the scheduling, so we cannot know the is_decode_only
@@ -241,7 +241,10 @@ def _process_step_result(result):
 
                 # Suspend.
                 if attempted_step_count % args.suspend_resume_interval == 0:
-                    print("**** step %d/%d ... suspend." % (engine.context.step_count, attempted_step_count))
+                    print(
+                        "**** step %d/%d ... suspend."
+                        % (engine.context.step_count, attempted_step_count)
+                    )
                     engine.suspend()
 
                 # Resume, 0+ attempted steps later.
@@ -251,7 +254,10 @@ def _process_step_result(result):
                     % args.suspend_resume_interval
                     == 0
                 ):
-                    print("**** step %d/%d ... resume." % (engine.context.step_count, attempted_step_count))
+                    print(
+                        "**** step %d/%d ... resume."
+                        % (engine.context.step_count, attempted_step_count)
+                    )
                     engine.resume()
 
             # If engine suspended, continue to next iter.
@@ -469,7 +475,9 @@ def escape_str(s):
             # Attach peak memory metrics; the functional test only validates these
             # if the fields exist in the golden values.
             json_results.update(peak_mem_stats)
-            json_results["lifetime_prefill_token_count"] = engine.context.lifetime_prefill_token_count
+            json_results["lifetime_prefill_token_count"] = (
+                engine.context.lifetime_prefill_token_count
+            )
             json_results["async_sched_step_count"] = engine.context.async_sched_step_count
             json_results["async_sched_compaction_step_count"] = (
                 engine.context.async_sched_compaction_step_count

examples/inference/utils.py

@@ -34,10 +34,24 @@ def get_default_sampling_params(termination_id: int = None):
 
 
 def get_curr_time(do_broadcast: bool = True) -> float:
-    """Get synchronized time across ranks."""
+    """Get the current time, optionally synchronized across distributed ranks.
+
+    Args:
+        do_broadcast (bool): Whether multi-rank callers require a rank-zero
+            timestamp broadcast.
+
+    Returns:
+        float: Current time in seconds.
+    """
+    if (
+        not do_broadcast
+        or not torch.distributed.is_initialized()
+        or torch.distributed.get_world_size() == 1
+    ):
+        return time.time_ns() / 10**9
+
     curr_time = torch.cuda.LongTensor([time.time_ns()])
-    if torch.distributed.is_initialized() and do_broadcast:
-        torch.distributed.broadcast(curr_time, src=0)
+    torch.distributed.broadcast(curr_time, src=0)
     return curr_time.item() / 10**9
 
 
@@ -401,7 +415,9 @@ def dump_inference_results_to_json(
         lifetime_prefill_token_count (int): Total prefill tokens processed.
         async_sched_step_count (int): Number of async scheduling decode steps.
         async_sched_compaction_step_count (int): Number of async scheduling decode
-            steps where post-forward compaction discarded finished rows.
+            steps that discarded speculative rows for finished requests. This
+            includes identity-prefix and all-finished cases that require no GPU
+            gather.
     """
     if not args.output_path:
         return
@@ -442,9 +458,7 @@ def dump_inference_results_to_json(
     json_results.update(peak_mem_stats)
     json_results["lifetime_prefill_token_count"] = lifetime_prefill_token_count
     json_results["async_sched_step_count"] = async_sched_step_count
-    json_results["async_sched_compaction_step_count"] = (
-        async_sched_compaction_step_count
-    )
+    json_results["async_sched_compaction_step_count"] = async_sched_compaction_step_count
 
     print(f' Saving results to {args.output_path}')
     with open(args.output_path, "w") as fp: