From 4fe51a2210ae602ac1904e60f17a22bf1123f43c Mon Sep 17 00:00:00 2001 From: Xiao Wang <24860335+xwang233@users.noreply.github.com> Date: Mon, 26 Oct 2026 18:43:44 +0000 Subject: [PATCH 01/16] [None][feat] KVCacheManagerV2: opt-in deferred close of finished requests' caches With TRTLLM_KV_DEFER_CLOSE=1 the overlap executor arms deferred closing on a target-model KVCacheManagerV2. free_resources then releases a finished request's index slot at once and queues its cache; the close, which commits the last block and releases the pages, runs right after the next forward step is launched, before the executor waits on an empty queue, and at shutdown. Dummy requests always close at once. Off by default (TRTLLM_KV_DEFER_CLOSE unset or 0): the unset path is unchanged; the index-slot release is factored into _detach_index_slot. Tests: test_kv_cache_v2_deferred_close.py. Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com> --- .../kv_cache/kv_cache_manager_v2.py | 72 +++- tensorrt_llm/_torch/pyexecutor/py_executor.py | 23 ++ .../test_kv_cache_v2_deferred_close.py | 346 ++++++++++++++++++ 3 files changed, 435 insertions(+), 6 deletions(-) create mode 100644 tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_deferred_close.py diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py index 37fc9e0e3a6d..59667f435e08 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py @@ -1230,6 +1230,12 @@ class KVCacheManagerV2(BaseResourceManager): # Read by KvCacheCreator: a subclass that overrides the context # commit/history protocol opts out of generic reuse-match backoff. _supports_reuse_match_backoff = True + # Deferred close (opt-in, TRTLLM_KV_DEFER_CLOSE=1): see arm_deferred_close(). + # Declared on the class so free_resources and shutdown also work on + # instances built without running __init__. + _defer_close_enabled: bool = False + _defer_close_armed: bool = False + _deferred_closes: Tuple[Tuple[int, _KVCache], ...] = () def __init__( self, @@ -1928,6 +1934,9 @@ def create_cold_page_codec(cache_config: object) -> Optional[object]: index_mapper_capacity, max_beam_width, self.max_copy_beam_width ) self._early_freed_index_requests: set[int] = set() + self._defer_close_enabled = ( + os.environ.get("TRTLLM_KV_DEFER_CLOSE", "0") == "1" and not self.is_draft + ) # Requests whose pages a connector is still reading from, so the # release half of `preempt_request` has to wait. self._pending_preemption: Dict[int, LlmRequest] = {} @@ -5635,15 +5644,54 @@ def release_index_slot(self, request_id: int) -> None: # mirrored, and the target may release the same request twice. return if kv_cache is not None: - if self.sparse_metadata_batch is not None: - self.sparse_metadata_batch.record_read(self._stream.cuda_stream) - self.sparse_metadata_batch.remove(kv_cache) - for beam_idx in range(int(kv_cache.beam_width)): - for pool_idx in range(self.num_pools): - kv_cache.set_base_page_index_buf(BeamIndex(beam_idx), pool_idx, None) + self._detach_index_slot(kv_cache) self.index_mapper.remove_sequence(request_id) self._early_freed_index_requests.add(request_id) + def _detach_index_slot(self, kv_cache: _KVCache) -> None: + """Unbind a cache from its IndexMapper slot's page-index buffers (and + the sparse metadata batch) so the slot can be reused while the cache + stays open.""" + if self.sparse_metadata_batch is not None: + self.sparse_metadata_batch.record_read(self._stream.cuda_stream) + self.sparse_metadata_batch.remove(kv_cache) + for beam_idx in range(int(kv_cache.beam_width)): + for pool_idx in range(self.num_pools): + kv_cache.set_base_page_index_buf(BeamIndex(beam_idx), pool_idx, None) + + def arm_deferred_close(self, armed: bool) -> bool: + """Arm or disarm deferred closing of finished requests' caches. + + While armed, free_resources releases a finished (non-dummy) request's + IndexMapper slot at once (as release_index_slot does) and queues the + cache; its close(), which commits the last block and releases the + pages, runs in flush_deferred_closes(). The overlap executor arms this + only for its loop and flushes right after launching each forward step, + where the GPU still has queued work to cover the close, and before it + waits on an empty request queue with no request left. Takes effect + only with TRTLLM_KV_DEFER_CLOSE=1 on a target-model manager. Disarming + flushes. Returns whether deferred closing is armed. + """ + was_armed = self._defer_close_armed + self._defer_close_armed = bool(armed) and self._defer_close_enabled + if self._defer_close_armed and not was_armed: + logger.info( + "KVCacheManagerV2: deferred close armed (TRTLLM_KV_DEFER_CLOSE=1): finished " + "requests' caches close after the next forward step is launched" + ) + if not self._defer_close_armed: + self.flush_deferred_closes() + return self._defer_close_armed + + def flush_deferred_closes(self) -> int: + """Close the caches whose close() free_resources deferred, in the order + their requests finished. Returns how many were closed.""" + pending, self._deferred_closes = self._deferred_closes, () + for request_id, kv_cache in pending: + kv_cache.close() + self.impl.clear_stats_excluded(request_id) + return len(pending) + def free_resources(self, request: LlmRequest, pin_on_release: bool = False): # A request awaiting preemption can still be cancelled or fail while # its saves drain, and a dead entry here would block every later @@ -5680,6 +5728,16 @@ def free_resources(self, request: LlmRequest, pin_on_release: bool = False): self.impl.clear_stats_excluded(request.py_request_id) return kv_cache.discard_pending_stats() + # Dummy requests reuse their ids, which the stats-exclusion bookkeeping + # done after close() is keyed on, so they always close at once. + if self._defer_close_armed and not request.is_dummy: + if request.py_request_id in self._early_freed_index_requests: + self._early_freed_index_requests.discard(request.py_request_id) + else: + self._detach_index_slot(kv_cache) + self.index_mapper.remove_sequence(request.py_request_id) + self._deferred_closes = (*self._deferred_closes, (request.py_request_id, kv_cache)) + return kv_cache.close() self.impl.clear_stats_excluded(request.py_request_id) if request.py_request_id in self._early_freed_index_requests: @@ -5931,6 +5989,8 @@ def check_invalid_values_in_kv_cache(self, fill_with_zero: bool = False) -> bool return bool(has_invalid_values) def shutdown(self): + self._defer_close_armed = False + self.flush_deferred_closes() if self.sparse_metadata_batch is not None: self.sparse_metadata_batch.record_read(self._stream.cuda_stream) self.sparse_metadata_batch.close() diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 9711f72f6f11..fbf93ff7e01d 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -400,6 +400,10 @@ class PyExecutor: # If the number of micro batches is too large, the executor will spend too much host memory (No additional GPU memory is required). # 1024 in-flight micro batches can avoid synchronization in most cases and keep host memory usage low. MIN_ASYNC_MICRO_BATCH_NUM = 1024 + # True while the overlap loop has armed KVCacheManagerV2's opt-in deferred + # close (TRTLLM_KV_DEFER_CLOSE=1). Declared on the class so instances built + # without running __init__ read it too. + _kv_defer_close_armed = False def __init__( self, @@ -5222,6 +5226,8 @@ def _maybe_rebalance_kv_pools(self) -> None: torch.cuda.current_stream().synchronize() self._consume_previous_batch_for_rebalance() + if self._kv_defer_close_armed: + self.kv_cache_manager.flush_deferred_closes() self._rebalance_kv_pools_now() def _rebalance_kv_pools_now(self) -> None: @@ -5422,6 +5428,9 @@ def _executor_loop_overlap(self): torch.cuda.set_device(self.device_id) # ensure the context is created, otherwise, some MPI calls will fail. CUASSERT(cudart.cudaSetDevice(self.device_id)) + self._kv_defer_close_armed = ( + self._is_kv_manager_v2 + and self.kv_cache_manager.arm_deferred_close(True)) with self._profiler() as profile_step, self.hang_detector: iter_start_time = time.time() iter_stats = None @@ -5639,6 +5648,12 @@ def _executor_loop_overlap(self): self._maybe_prefetch_next_iter_mm_encoders( scheduled_batch) + if self._kv_defer_close_armed: + # Close the caches of requests that finished in the + # previous iteration here, after this iteration's forward + # step is queued and before waiting for the previous one. + self.kv_cache_manager.flush_deferred_closes() + if self.previous_batch is not None and should_process_previous_batch: self._update_requests(self.previous_batch.sample_state) # Turning off speculative decoding when Acceptance Rate is low. @@ -5773,6 +5788,9 @@ def _executor_loop_overlap(self): self.disagg.pace_idle() self.iter_counter += 1 + if self._kv_defer_close_armed: + self._kv_defer_close_armed = False + self.kv_cache_manager.arm_deferred_close(False) @nvtx_range("_accept_draft_tokens") def _accept_draft_tokens( @@ -6097,6 +6115,11 @@ def _fetch_and_enqueue_requests(self, waiting_queue: WaitingQueue, idle = (total_num_live_requests == 0 and len(waiting_queue) == 0 and not self.is_shutdown and not self._has_pending_connector_transfers()) + if idle and self._kv_defer_close_armed: + # No forward step is coming to flush the caches of the requests + # that just finished, so close them before waiting: the next + # request then sees their pages and last blocks when scheduled. + self.kv_cache_manager.flush_deferred_closes() if idle: # In Ray path (TLLM_DISABLE_MPI=1), use a periodic heartbeat timeout so rank 0 # reaches the broadcast path regularly to prevent trtllm-serve timeout when idle. diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_deferred_close.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_deferred_close.py new file mode 100644 index 000000000000..9ffa3a4bf991 --- /dev/null +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_deferred_close.py @@ -0,0 +1,346 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""KVCacheManagerV2's opt-in deferred close (TRTLLM_KV_DEFER_CLOSE=1). + +While armed, free_resources releases a finished request's IndexMapper slot at +once and leaves the cache's close() to flush_deferred_closes(); the overlap +executor arms it for its loop and flushes after launching each forward step +and before waiting for requests with none left. +""" + +import types +from unittest.mock import MagicMock, Mock + +from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 +from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor + + +def _manager(events, *, enabled=True, num_pools=2, sparse=False): + mgr = object.__new__(KVCacheManagerV2) + mgr._defer_close_enabled = enabled + mgr._pending_preemption = {} + mgr.kv_connector_manager = None + mgr.conversation_manager = None + mgr._allocated_draft_lens = {} + mgr._request_stats_enabled_ids = set() + mgr._fresh_pages_filled = {} + mgr._disagg_receive_ready = {} + mgr.kv_cache_map = {} + mgr._early_freed_index_requests = set() + mgr.num_pools = num_pools + mgr.is_draft = False + mgr._stream = types.SimpleNamespace(cuda_stream=0) + mgr.impl = Mock() + mgr.impl.clear_stats_excluded.side_effect = lambda rid: events.append(("clear_stats", rid)) + mgr.index_mapper = Mock() + mgr.index_mapper.remove_sequence.side_effect = lambda rid: events.append(("remove_seq", rid)) + if sparse: + mgr.sparse_metadata_batch = Mock() + mgr.sparse_metadata_batch.remove.side_effect = lambda kv: events.append( + ("sparse_remove", kv.rid) + ) + mgr.sparse_metadata_batch.close.side_effect = lambda: events.append(("sparse_close",)) + return mgr + + +def _cache(mgr, events, rid): + kv = Mock() + kv.rid = rid + kv.beam_width = 1 + kv.close.side_effect = lambda: events.append(("close", rid)) + kv.set_base_page_index_buf.side_effect = lambda beam, pool, buf: events.append( + ("unbind", rid, pool, buf) + ) + mgr.kv_cache_map[rid] = kv + return kv + + +def _request(rid, *, dummy=False): + return types.SimpleNamespace(py_request_id=rid, is_dummy=dummy) + + +def test_off_unless_enabled_closes_at_once(): + events = [] + mgr = _manager(events, enabled=False) + _cache(mgr, events, 7) + assert mgr.arm_deferred_close(True) is False + mgr.free_resources(_request(7)) + assert events == [("close", 7), ("clear_stats", 7), ("remove_seq", 7)] + assert mgr._deferred_closes == () + + +def test_class_defaults_keep_managers_built_without_init_unchanged(): + events = [] + mgr = _manager(events) + del mgr._defer_close_enabled + _cache(mgr, events, 3) + assert mgr.arm_deferred_close(True) is False + mgr.free_resources(_request(3)) + assert events == [("close", 3), ("clear_stats", 3), ("remove_seq", 3)] + + +def test_armed_releases_slot_now_and_closes_on_flush(): + events = [] + mgr = _manager(events, num_pools=2) + kv = _cache(mgr, events, 7) + assert mgr.arm_deferred_close(True) is True + mgr.free_resources(_request(7)) + # Slot released at once, after the cache is unbound from it; no close yet. + assert events == [("unbind", 7, 0, None), ("unbind", 7, 1, None), ("remove_seq", 7)] + kv.close.assert_not_called() + assert 7 not in mgr.kv_cache_map + events.clear() + assert mgr.flush_deferred_closes() == 1 + assert events == [("close", 7), ("clear_stats", 7)] + assert mgr.flush_deferred_closes() == 0 + + +def test_flush_closes_in_finish_order(): + events = [] + mgr = _manager(events, num_pools=1) + for rid in (5, 2, 9): + _cache(mgr, events, rid) + mgr.arm_deferred_close(True) + for rid in (5, 2, 9): + mgr.free_resources(_request(rid)) + events.clear() + assert mgr.flush_deferred_closes() == 3 + assert [e for e in events if e[0] == "close"] == [("close", 5), ("close", 2), ("close", 9)] + assert [e for e in events if e[0] == "clear_stats"] == [ + ("clear_stats", 5), + ("clear_stats", 2), + ("clear_stats", 9), + ] + + +def test_each_stats_exclusion_is_cleared_after_its_close(): + events = [] + mgr = _manager(events, num_pools=1) + for rid in (1, 4): + _cache(mgr, events, rid) + mgr.arm_deferred_close(True) + for rid in (1, 4): + mgr.free_resources(_request(rid)) + events.clear() + mgr.flush_deferred_closes() + assert events == [("close", 1), ("clear_stats", 1), ("close", 4), ("clear_stats", 4)] + + +def test_dummy_requests_close_at_once_while_armed(): + events = [] + mgr = _manager(events, num_pools=1) + _cache(mgr, events, 11) + mgr.arm_deferred_close(True) + mgr.free_resources(_request(11, dummy=True)) + assert events == [("close", 11), ("clear_stats", 11), ("remove_seq", 11)] + assert mgr._deferred_closes == () + + +def test_early_released_slot_is_not_released_twice(): + events = [] + mgr = _manager(events, num_pools=1) + _cache(mgr, events, 6) + mgr._early_freed_index_requests.add(6) + mgr.arm_deferred_close(True) + mgr.free_resources(_request(6)) + assert events == [] + assert 6 not in mgr._early_freed_index_requests + mgr.flush_deferred_closes() + assert events == [("close", 6), ("clear_stats", 6)] + + +def test_sparse_batch_is_left_before_the_slot_is_released(): + events = [] + mgr = _manager(events, num_pools=1, sparse=True) + _cache(mgr, events, 8) + mgr.arm_deferred_close(True) + mgr.free_resources(_request(8)) + assert events == [("sparse_remove", 8), ("unbind", 8, 0, None), ("remove_seq", 8)] + mgr.sparse_metadata_batch.record_read.assert_called_once_with(0) + + +def test_disarm_flushes_and_stops_deferring(): + events = [] + mgr = _manager(events, num_pools=1) + _cache(mgr, events, 1) + _cache(mgr, events, 2) + mgr.arm_deferred_close(True) + mgr.free_resources(_request(1)) + events.clear() + assert mgr.arm_deferred_close(False) is False + assert events == [("close", 1), ("clear_stats", 1)] + events.clear() + mgr.free_resources(_request(2)) + assert events == [("close", 2), ("clear_stats", 2), ("remove_seq", 2)] + + +def test_shutdown_closes_deferred_caches_before_the_sparse_batch(): + events = [] + mgr = _manager(events, num_pools=1, sparse=True) + _cache(mgr, events, 3) + mgr.arm_deferred_close(True) + mgr.free_resources(_request(3)) + events.clear() + mgr.event_manager = None + mgr.shutdown() + assert events.index(("close", 3)) < events.index(("sparse_close",)) + assert mgr._defer_close_armed is False + assert mgr._deferred_closes == () + + +def _overlap_loop_stub(kv_cache_manager): + executor = object.__new__(PyExecutor) + executor.device_id = 0 + profiler = MagicMock() + profiler.__enter__.return_value = Mock() + executor._profiler = Mock(return_value=profiler) + executor.hang_detector = MagicMock() + executor.enable_iter_perf_stats = False + executor._resource_governor_enabled = False + executor._is_kv_manager_v2 = True + executor.kv_cache_manager = kv_cache_manager + executor._can_pause_for_rebalance = Mock(return_value=False) + executor._wait_for_model_engine_input_copy = Mock() + executor._prepare_and_schedule_batch = Mock(return_value=(None, None)) + executor._mm_encoder_item_scheduling_enabled = False + executor.is_benchmark_disagg = False + executor._flush_pending_transfer_responses = Mock() + return executor + + +def _patch_loop_cuda(monkeypatch): + for name in ("torch.cuda.set_device", "cudart.cudaSetDevice", "CUASSERT"): + monkeypatch.setattr(f"tensorrt_llm._torch.pyexecutor.py_executor.{name}", Mock()) + + +def test_overlap_loop_arms_for_its_run_and_disarms_on_exit(monkeypatch): + mgr = Mock() + mgr.arm_deferred_close = Mock(side_effect=lambda armed: armed) + executor = _overlap_loop_stub(mgr) + _patch_loop_cuda(monkeypatch) + + PyExecutor._executor_loop_overlap(executor) + + assert [c.args for c in mgr.arm_deferred_close.call_args_list] == [(True,), (False,)] + assert executor._kv_defer_close_armed is False + + +def test_overlap_loop_leaves_the_manager_alone_when_not_enabled(monkeypatch): + mgr = Mock() + mgr.arm_deferred_close = Mock(return_value=False) + executor = _overlap_loop_stub(mgr) + _patch_loop_cuda(monkeypatch) + + PyExecutor._executor_loop_overlap(executor) + + mgr.arm_deferred_close.assert_called_once_with(True) + mgr.flush_deferred_closes.assert_not_called() + assert executor._kv_defer_close_armed is False + + +class _QueueWait(Exception): + pass + + +def _fetch_stub(kv_cache_manager, events, *, armed=True): + executor = object.__new__(PyExecutor) + executor.kv_cache_manager = kv_cache_manager + executor._kv_defer_close_armed = armed + executor.kv_connector_manager = None + executor.control_requests = [] + executor.is_shutdown = False + executor._has_pending_connector_transfers = Mock(return_value=False) + executor._disable_mpi = False + executor.dist = types.SimpleNamespace(rank=0) + executor.request_accumulated = [] + executor.hang_detector = MagicMock() + executor.executor_request_queue = Mock() + + def wait(timeout): + events.append(("queue_wait", timeout)) + raise _QueueWait + + executor.executor_request_queue.get_from_request_queue.side_effect = wait + return executor + + +def _fetch(executor, total_num_live_requests, waiting=()): + try: + PyExecutor._fetch_and_enqueue_requests(executor, list(waiting), total_num_live_requests) + except _QueueWait: + pass + + +def test_last_request_closes_before_the_idle_queue_wait(): + events = [] + mgr = _manager(events, num_pools=1) + _cache(mgr, events, 4) + mgr.arm_deferred_close(True) + mgr.free_resources(_request(4)) + events.clear() + executor = _fetch_stub(mgr, events) + + _fetch(executor, total_num_live_requests=0) + + # The blocking wait for the next request starts only after the close. + assert events == [("close", 4), ("clear_stats", 4), ("queue_wait", None)] + assert mgr._deferred_closes == () + + +def test_busy_fetch_leaves_deferred_closes_to_the_loop(): + events = [] + mgr = _manager(events, num_pools=1) + _cache(mgr, events, 4) + mgr.arm_deferred_close(True) + mgr.free_resources(_request(4)) + events.clear() + executor = _fetch_stub(mgr, events) + + _fetch(executor, total_num_live_requests=3) + _fetch(executor, total_num_live_requests=0, waiting=[object()]) + + assert [e[0] for e in events] == ["queue_wait", "queue_wait"] + assert len(mgr._deferred_closes) == 1 + + +def test_idle_fetch_does_not_flush_when_not_armed(): + events = [] + mgr = Mock() + executor = _fetch_stub(mgr, events, armed=False) + + _fetch(executor, total_num_live_requests=0) + + mgr.flush_deferred_closes.assert_not_called() + assert events == [("queue_wait", None)] + + +def test_rebalance_flushes_deferred_closes_before_adjusting(monkeypatch): + order = [] + executor = object.__new__(PyExecutor) + executor.kv_cache_manager = Mock() + executor.kv_cache_manager.flush_deferred_closes.side_effect = lambda: order.append("flush") + executor._agreed_need_adjustment = Mock(return_value=True) + executor._consume_previous_batch_for_rebalance = Mock( + side_effect=lambda: order.append("consume") + ) + executor._rebalance_kv_pools_now = Mock(side_effect=lambda: order.append("adjust")) + executor._kv_defer_close_armed = True + monkeypatch.setattr( + "tensorrt_llm._torch.pyexecutor.py_executor.torch.cuda.current_stream", Mock() + ) + + PyExecutor._maybe_rebalance_kv_pools(executor) + + assert order == ["consume", "flush", "adjust"] From 37b37ed6288c8e97d84d67947ad0598db4aaaaf7 Mon Sep 17 00:00:00 2001 From: Xiao Wang <24860335+xwang233@users.noreply.github.com> Date: Mon, 26 Oct 2026 18:43:54 +0000 Subject: [PATCH 02/16] [None][feat] KVCacheV2Scheduler: opt-in context page backoff and admission diagnostics TRTLLM_SCHED_CTX_PAGE_BACKOFF_PASSES=N: once a pass that schedules generation sees a first context chunk refused for KV pages, the rest of that pass and the next N generation passes try no more first chunks; one pass in N + 1 probes again. Chunk continuations and passes without generation are never held. The state depends only on the scheduling decisions, so ranks stay in lockstep. TRTLLM_SCHED_ADMISSION_DIAG=1 (window TRTLLM_SCHED_ADMISSION_DIAG_PASSES) logs the host time of context admission per phase. Both are off by default (backoff 0, diagnostics unset): no candidate is held back and no scheduler method is wrapped. Tests: test_kv_cache_v2_scheduler.py. Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com> --- .../pyexecutor/scheduler/scheduler_v2.py | 188 +++++++++++++++++ .../kv_cache/test_kv_cache_v2_scheduler.py | 196 ++++++++++++++++++ 2 files changed, 384 insertions(+) diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py index 01f338aa3725..95539c89adf4 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py @@ -15,6 +15,7 @@ import enum import os +import time from collections import Counter from typing import Callable, Optional @@ -249,6 +250,138 @@ def withdraw_if_not_run( return True +class ContextPageBackoff: + """Stop retrying context admission that the KV pool keeps refusing. + + Admitting a first context chunk is not a cheap check. ``prepare_context`` + creates the request's cache and claims its reusable prefix, the resume + brings that prefix onto the GPU (onboarding pages that live on a lower + tier and evicting others to make room), and when the chunk's own pages + then do not fit, the whole claim is rolled back and freed. The cost grows + with the prompt, not with the chunk. With the pool near full, every + pending request repeats this on every pass and is refused every time, so + host time per step grows with the number of pending requests while no + context work is admitted. + + With ``passes`` N, once a pass that schedules generation sees a context + admission refused for pages, the rest of that pass tries no more first + chunks, and the next N passes that schedule generation try none either. + One pass in N + 1 then probes whether pages have come back. Chunk + continuations of started prefills are always tried, and passes without + generation are never held: with nothing decoding, nothing else can use + the step. + + The state is a pure function of the sequence of scheduling decisions, so + ranks that schedule the same requests stay in lockstep. + """ + + def __init__(self, passes: int): + if passes <= 0: + raise ValueError(f"passes must be positive, got {passes}") + self.passes = passes + self._remaining = 0 + + @property + def remaining(self) -> int: + """Generation passes left before first chunks are tried again.""" + return self._remaining + + def holds_first_chunks(self, has_generation: bool, refused_this_pass: bool) -> bool: + """Whether this pass must skip first-chunk context admission.""" + return has_generation and (self._remaining > 0 or refused_this_pass) + + def record_pass(self, has_generation: bool, refused_this_pass: bool) -> None: + """Advance the backoff by one scheduling pass.""" + if not has_generation: + return + if refused_this_pass: + self._remaining = self.passes + elif self._remaining > 0: + self._remaining -= 1 + + +class _AdmissionDiagnostics: + """Opt-in host-time accounting of phase-2 context admission. + + Logs, every ``log_passes`` scheduling passes, how many context admissions + were attempted and refused for pages, and where the host time of the + admission phase went. A pass slower than ``slow_pass_ms`` is also logged + on its own, at most ``max_slow_logs`` times per window. + """ + + _SECTIONS = ("prepare", "allocate", "release") + + def __init__(self, log_passes: int = 1024, slow_pass_ms: float = 100.0, max_slow_logs: int = 4): + self.log_passes = log_passes + self.slow_pass_ms = slow_pass_ms + self.max_slow_logs = max_slow_logs + self._window: Counter = Counter() + self._pass: Counter = Counter() + self._pass_start = 0.0 + self._slow_logged = 0 + + def wrap(self, name: str, fn: Callable) -> Callable: + """Return *fn* timed into section *name* of the current pass.""" + + def timed(*args, **kwargs): + start = time.perf_counter() + try: + return fn(*args, **kwargs) + finally: + self._pass[f"{name}_ms"] += (time.perf_counter() - start) * 1e3 + self._pass[f"{name}_calls"] += 1 + + return timed + + def begin_pass(self) -> None: + self._pass.clear() + self._pass_start = time.perf_counter() + + def count(self, key: str, n: int = 1) -> None: + self._pass[key] += n + + def end_pass(self, num_pending: int, num_scheduled_ctx: int) -> None: + pass_ms = (time.perf_counter() - self._pass_start) * 1e3 + stats = self._pass + stats["phase2_ms"] = pass_ms + window = self._window + window["passes"] += 1 + window.update(stats) + window["max_phase2_ms"] = max(window["max_phase2_ms"], pass_ms) + if stats["page_refused"]: + window["refused_passes"] += 1 + window["refused_pass_ms"] += pass_ms + if pass_ms > self.slow_pass_ms: + window["slow_passes"] += 1 + window["slow_pass_ms"] += pass_ms + if self._slow_logged < self.max_slow_logs: + self._slow_logged += 1 + logger.info( + f"[V2Scheduler] slow admission pass {pass_ms:.1f} ms: pending {num_pending}, " + f"attempts {stats['attempts']}, page refusals {stats['page_refused']}, " + f"admitted {num_scheduled_ctx}, held by backoff {stats['held']}, " + + ", ".join( + f"{s} {stats[s + '_calls']} calls {stats[s + '_ms']:.1f} ms" + for s in self._SECTIONS + ) + ) + if window["passes"] >= self.log_passes: + logger.info( + f"[V2Scheduler] admission diagnostics over {window['passes']} passes: " + f"attempts {window['attempts']}, page refusals {window['page_refused']} " + f"in {window['refused_passes']} passes ({window['refused_pass_ms']:.0f} ms), " + f"held by backoff {window['held']}, phase-2 total {window['phase2_ms']:.0f} ms, " + f"max {window['max_phase2_ms']:.1f} ms, passes > {self.slow_pass_ms:.0f} ms " + f"{window['slow_passes']} ({window['slow_pass_ms']:.0f} ms), " + + ", ".join( + f"{s} {window[s + '_calls']} calls {window[s + '_ms']:.0f} ms" + for s in self._SECTIONS + ) + ) + window.clear() + self._slow_logged = 0 + + class ShortestPrefillFirstOrder: """Order pending context requests by the prefill work they still need. @@ -480,6 +613,31 @@ def __init__( _read_non_negative_int_env("TRTLLM_SCHED_SPF_MAX_WAIT_PASSES", 128) ) self._context_admission_stats: Counter = Counter() + + # Opt-in (default off): back off first-chunk context admission after + # the KV pool refuses one; see ContextPageBackoff. + backoff_passes = _read_non_negative_int_env("TRTLLM_SCHED_CTX_PAGE_BACKOFF_PASSES", 0) + self._ctx_page_backoff = ContextPageBackoff(backoff_passes) if backoff_passes > 0 else None + # Set by the context admission paths when the KV pool refused pages + # during the current pass. + self._ctx_page_refused = False + # Opt-in (default off) host-time accounting of context admission. + self._admission_diag = None + if os.environ.get("TRTLLM_SCHED_ADMISSION_DIAG", "0") == "1": + diag = _AdmissionDiagnostics( + log_passes=_read_non_negative_int_env("TRTLLM_SCHED_ADMISSION_DIAG_PASSES", 1024) + or 1024 + ) + self._prepare_context_pair = diag.wrap("prepare", self._prepare_context_pair) + self._try_allocate_context = diag.wrap("allocate", self._try_allocate_context) + self.release_context_admission = diag.wrap("release", self.release_context_admission) + self._admission_diag = diag + if self._ctx_page_backoff is not None or self._admission_diag is not None: + logger.info( + "KVCacheV2Scheduler context page backoff: " + f"passes={backoff_passes}, admission_diagnostics=" + f"{self._admission_diag is not None}" + ) if self._prefill_decode_interval is not None or self._shortest_prefill_first is not None: logger.info( "KVCacheV2Scheduler context admission: " @@ -722,6 +880,11 @@ def preempt_for_pages(req: LlmRequest) -> bool: ctx_candidates, deferred_by_interval = self._admissible_pending_ctx( pending_ctx, requests_list, has_generation ) + self._ctx_page_refused = False + backoff = self._ctx_page_backoff + diag = self._admission_diag + if diag is not None: + diag.begin_pass() # Prefix-aware skip: when several first-chunk context requests would # contribute the same not-yet-cached block, admitting them together @@ -754,6 +917,15 @@ def preempt_for_pages(req: LlmRequest) -> bool: continue if req.py_request_id in preempted_ids: continue + if ( + backoff is not None + and req.state_value == self._context_init_state_value + and req.is_first_context_chunk + and backoff.holds_first_chunks(has_generation, self._ctx_page_refused) + ): + if diag is not None: + diag.count("held") + continue # Fast path: a pending context request whose draft-mirror admission # is certainly refused this iteration cannot be scheduled, so the # prefix-reuse probe below (two O(prompt_len / tokens_per_block) @@ -796,6 +968,8 @@ def preempt_for_pages(req: LlmRequest) -> bool: scheduled_encoder.append(req) budget.commit(req, tokens, peft_pages) else: + if diag is not None: + diag.count("attempts") action, tokens, chunking_flag = self._try_schedule_context( req, budget, preempt_for_pages ) @@ -820,6 +994,10 @@ def preempt_for_pages(req: LlmRequest) -> bool: if first_new_block is not None: contributed_blocks.add(first_new_block) + if backoff is not None: + backoff.record_pass(has_generation, self._ctx_page_refused) + if diag is not None: + diag.end_pass(len(ctx_candidates), len(scheduled_ctx)) if self._prefill_decode_interval is not None: self._record_context_admission(scheduled_ctx, has_generation, deferred_by_interval) if self._shortest_prefill_first is not None: @@ -1137,6 +1315,7 @@ def _try_schedule_context_full( # before budget check. if not self._prepare_context_pair(req): logger.debug(f"prepare_context failed for context request {req.py_request_id}") + self._note_page_refusal() return ScheduleAction.STOP, 0, False context_tokens = req.context_remaining_length @@ -1159,6 +1338,7 @@ def _try_schedule_context_full( # prepare_context to resume it. if preempt_for_pages(req): return ScheduleAction.STOP, 0, False + self._note_page_refusal() return ScheduleAction.SKIP, 0, False cross_action = self._try_schedule_cross_context(req) @@ -1168,6 +1348,12 @@ def _try_schedule_context_full( return ScheduleAction.SCHEDULED, req_tokens, False + def _note_page_refusal(self) -> None: + """Record that the KV pool refused a context admission this pass.""" + self._ctx_page_refused = True + if self._admission_diag is not None: + self._admission_diag.count("page_refused") + def _has_context_chunk_budget(self, budget: BudgetTracker) -> bool: remaining = budget.remaining_tokens return remaining is None or ( @@ -1202,6 +1388,7 @@ def _try_schedule_context_chunked( # Prepare context (create _KVCache, block reuse, resume — no resize) if not self._prepare_context_pair(req): logger.debug(f"prepare_context failed for chunked context request {req.py_request_id}") + self._note_page_refusal() return ScheduleAction.SKIP, 0, False # Calculate chunk size from remaining budget @@ -1273,6 +1460,7 @@ def _try_schedule_context_chunked( # Out of pages, as in _try_schedule_context_full. if preempt_for_pages(req): return ScheduleAction.STOP, 0, False + self._note_page_refusal() return ScheduleAction.SKIP, 0, False cross_action = self._try_schedule_cross_context(req) diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_scheduler.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_scheduler.py index c3c45c0515ba..ef973b720d0d 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_scheduler.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_scheduler.py @@ -5366,3 +5366,199 @@ def test_started_prefill_may_still_preempt_under_spf(self, admission_env): # The protection applies to requests that have not started prefilling; # between started requests the usual newest-first preemption holds. assert ids(out.recompute_paused_requests) == [2] + + +# =========================================================================== +# Opt-in context page backoff and admission diagnostics +# =========================================================================== + + +_BACKOFF_ENV = "TRTLLM_SCHED_CTX_PAGE_BACKOFF_PASSES" +_DIAG_ENV = "TRTLLM_SCHED_ADMISSION_DIAG" +_DIAG_PASSES_ENV = "TRTLLM_SCHED_ADMISSION_DIAG_PASSES" + + +@pytest.fixture +def backoff_env(monkeypatch): + for name in ( + _INTERVAL_ENV, + _SPF_ENV, + _SPF_WAIT_ENV, + _BACKOFF_ENV, + _DIAG_ENV, + _DIAG_PASSES_ENV, + ): + monkeypatch.delenv(name, raising=False) + return monkeypatch + + +def _refusing_manager(refused_ids): + """Manager whose pool refuses the context pages of every id in *refused_ids*.""" + return make_kv_cache_manager( + resize_context_fn=lambda req, n: req.py_request_id not in refused_ids + ) + + +def _resize_attempts(mgr): + return [c.args[0].py_request_id for c in mgr.resize_context.call_args_list] + + +class TestContextPageBackoff: + def test_off_by_default_retries_every_pass(self, backoff_env): + mgr = _refusing_manager({1, 2}) + sched = make_scheduler(mgr) + assert sched._ctx_page_backoff is None + assert sched._admission_diag is None + gen = make_gen_request(0) + pending = [make_ctx_request(1, 100), make_ctx_request(2, 100)] + for _ in range(3): + out = sched.schedule_request([gen, *pending], set()) + assert ids(out.generation_requests) == [0] + assert ids(out.context_requests) == [] + assert _resize_attempts(mgr) == [1, 2] * 3 + + @pytest.mark.parametrize("value", ["-1", "x", "2.5"]) + def test_invalid_value_rejected(self, backoff_env, value): + backoff_env.setenv(_BACKOFF_ENV, value) + with pytest.raises(ValueError, match=_BACKOFF_ENV): + make_scheduler(make_kv_cache_manager()) + + def test_zero_disables(self, backoff_env): + backoff_env.setenv(_BACKOFF_ENV, "0") + assert make_scheduler(make_kv_cache_manager())._ctx_page_backoff is None + + @pytest.mark.parametrize("passes", [1, 3]) + def test_refusal_holds_first_chunks_for_n_generation_passes(self, backoff_env, passes): + backoff_env.setenv(_BACKOFF_ENV, str(passes)) + mgr = _refusing_manager({1, 2}) + sched = make_scheduler(mgr) + gen = make_gen_request(0) + pending = [make_ctx_request(1, 100), make_ctx_request(2, 100)] + attempted_passes = [] + for step in range(2 * (passes + 1)): + before = len(mgr.resize_context.call_args_list) + out = sched.schedule_request([gen, *pending], set()) + # Decoding never waits on the backoff. + assert ids(out.generation_requests) == [0] + assert ids(out.context_requests) == [] + attempts = _resize_attempts(mgr)[before:] + if attempts: + # The first refusal ends first-chunk admission for the pass. + assert attempts == [1] + attempted_passes.append(step) + assert attempted_passes == [0, passes + 1] + + def test_admission_resumes_once_pages_return(self, backoff_env): + backoff_env.setenv(_BACKOFF_ENV, "2") + refused = {1} + mgr = _refusing_manager(refused) + sched = make_scheduler(mgr) + gen = make_gen_request(0) + ctx = make_ctx_request(1, 100) + assert ids(sched.schedule_request([gen, ctx], set()).context_requests) == [] + refused.clear() + assert ids(sched.schedule_request([gen, ctx], set()).context_requests) == [] + assert ids(sched.schedule_request([gen, ctx], set()).context_requests) == [] + assert ids(sched.schedule_request([gen, ctx], set()).context_requests) == [1] + assert sched._ctx_page_backoff.remaining == 0 + + def test_successful_admission_does_not_arm(self, backoff_env): + backoff_env.setenv(_BACKOFF_ENV, "4") + sched = make_scheduler(_refusing_manager(set())) + gen = make_gen_request(0) + for request_id in (1, 2, 3): + out = sched.schedule_request([gen, make_ctx_request(request_id, 100)], set()) + assert ids(out.context_requests) == [request_id] + assert sched._ctx_page_backoff.remaining == 0 + + def test_chunk_continuations_are_not_held(self, backoff_env): + backoff_env.setenv(_BACKOFF_ENV, "4") + mgr = _refusing_manager({2}) + sched = make_scheduler(mgr, max_num_tokens=1000, ctx_chunk_config=(None, 64)) + gen = make_gen_request(0) + refused = make_ctx_request(2, 100) + sched.schedule_request([gen, refused], set()) + assert sched._ctx_page_backoff.remaining == 4 + started = make_ctx_request( + 1, context_remaining_length=500, prompt_len=5000, is_first_context_chunk=False + ) + fresh = make_ctx_request(3, 100) + out = sched.schedule_request([gen, started, fresh], set()) + assert ids(out.context_requests) == [1] + assert 3 not in _resize_attempts(mgr) + + def test_not_held_without_generation(self, backoff_env): + backoff_env.setenv(_BACKOFF_ENV, "4") + mgr = _refusing_manager({1}) + sched = make_scheduler(mgr) + gen = make_gen_request(0) + sched.schedule_request([gen, make_ctx_request(1, 100)], set()) + assert sched._ctx_page_backoff.remaining == 4 + # Nothing decodes: the step has no other use, so admission is tried. + out = sched.schedule_request([make_ctx_request(2, 100)], set()) + assert ids(out.context_requests) == [2] + assert sched._ctx_page_backoff.remaining == 4 + + def test_held_context_is_not_a_stall(self, backoff_env): + backoff_env.setenv(_BACKOFF_ENV, "64") + sched = make_scheduler(_refusing_manager({1})) + gen = make_gen_request(0) + for _ in range(64): + sched.schedule_request([gen, make_ctx_request(1, 100)], set()) + assert sched._stalled_schedules == 0 + + def test_composes_with_prefill_decode_interval(self, backoff_env): + backoff_env.setenv(_BACKOFF_ENV, "2") + backoff_env.setenv(_INTERVAL_ENV, "1") + mgr = _refusing_manager({1}) + sched = make_scheduler(mgr) + gen = make_gen_request(0) + ctx = make_ctx_request(1, 100) + attempted = [] + for step in range(6): + before = len(mgr.resize_context.call_args_list) + sched.schedule_request([gen, ctx], set()) + if _resize_attempts(mgr)[before:]: + attempted.append(step) + # The interval never arms (nothing is admitted); the backoff alone + # spaces the refused attempts. + assert attempted == [0, 3] + + +class TestAdmissionDiagnostics: + def test_window_summary_counts_attempts_and_refusals(self, backoff_env): + backoff_env.setenv(_DIAG_ENV, "1") + backoff_env.setenv(_DIAG_PASSES_ENV, "3") + mgr = _refusing_manager({1}) + sched = make_scheduler(mgr) + gen = make_gen_request(0) + with patch("tensorrt_llm._torch.pyexecutor.scheduler.scheduler_v2.logger") as log: + for request_id in (2, 3, 4): + sched.schedule_request( + [gen, make_ctx_request(1, 100), make_ctx_request(request_id, 10)], set() + ) + summaries = [ + c.args[0] for c in log.info.call_args_list if "admission diagnostics" in c.args[0] + ] + assert len(summaries) == 1 + assert "over 3 passes: attempts 6, page refusals 3 in 3 passes" in summaries[0] + assert "allocate 6 calls" in summaries[0] + assert "prepare 6 calls" in summaries[0] + + def test_diagnostics_do_not_change_decisions(self, backoff_env): + def run(): + mgr = _refusing_manager({1}) + sched = make_scheduler(mgr) + gen = make_gen_request(0) + return [ + ids( + sched.schedule_request( + [gen, make_ctx_request(1, 100), make_ctx_request(r, 10)], set() + ).context_requests + ) + for r in (2, 3, 4) + ] + + plain = run() + backoff_env.setenv(_DIAG_ENV, "1") + assert run() == plain From 53f7d9a2a763efdccd2a18f3228f94b91630672d Mon Sep 17 00:00:00 2001 From: Xiao Wang <24860335+xwang233@users.noreply.github.com> Date: Mon, 26 Oct 2026 18:43:55 +0000 Subject: [PATCH 03/16] [None][feat] musespark_open: prefix-tokenization cache in the omni input processor MuseSparkOpenOmniInputProcessor declares supports_tokenization_cache and accepts enable_tokenization_cache. A media-free prompt is encoded through the shared prefix-tokenization cache, whose spliced ids equal what the processor's own encode would produce. Prompts with media, and processors built without a tokenizer, keep the existing path. Off by default: enabled only by the enable_tokenization_cache LLM argument. Tests: test_musespark_open_tokenization_cache.py. Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com> --- .../models/modeling_musespark_open_omni.py | 16 +- .../test_musespark_open_tokenization_cache.py | 222 ++++++++++++++++++ 2 files changed, 237 insertions(+), 1 deletion(-) create mode 100644 tests/unittest/_torch/modeling/test_musespark_open_tokenization_cache.py diff --git a/tensorrt_llm/_torch/models/modeling_musespark_open_omni.py b/tensorrt_llm/_torch/models/modeling_musespark_open_omni.py index 378a0f33bf14..f4201ade494c 100644 --- a/tensorrt_llm/_torch/models/modeling_musespark_open_omni.py +++ b/tensorrt_llm/_torch/models/modeling_musespark_open_omni.py @@ -165,12 +165,15 @@ class MuseSparkOpenOmniInputProcessor( semantics rather than imported from a shim. """ + supports_tokenization_cache = True + def __init__( self, model_path: str, config: PretrainedConfig, tokenizer, trust_remote_code: bool = True, + enable_tokenization_cache: bool = False, **kwargs, ): super().__init__( @@ -183,6 +186,12 @@ def __init__( self._config = config self._tokenizer = tokenizer self._model_path = model_path + # A media-free prompt is ``tokenizer.encode(text, add_special_tokens=False)`` + # (``_encode``), which is exactly what the prefix-tokenization cache + # computes, so its spliced ids are the ids this processor would produce. + # A processor built without a tokenizer (the engine's geometry copy) + # never tokenizes a served prompt and gets no cache. + self._init_tokenization_cache(enable_tokenization_cache, tokenizer) # `config.json` is not the whole preprocessing config on newer exports: # they drop `vision_config.navit_config` and declare the same knobs in # `processor_config.json` under HuggingFace spellings. Without this the @@ -541,7 +550,12 @@ def call_with_text_prompt( mm_data = inputs.get("multi_modal_data") or {} if not mm_data: - return self._encode(text_prompt), None + ids = ( + self._encode_with_tokenization_cache(text_prompt, sampling_params) + if text_prompt + else None + ) + return (self._encode(text_prompt) if ids is None else ids), None # Every modality key must be consumed below. Anything else -- including # "audio" against a checkpoint whose config does not enable the diff --git a/tests/unittest/_torch/modeling/test_musespark_open_tokenization_cache.py b/tests/unittest/_torch/modeling/test_musespark_open_tokenization_cache.py new file mode 100644 index 000000000000..4bf20fe277b4 --- /dev/null +++ b/tests/unittest/_torch/modeling/test_musespark_open_tokenization_cache.py @@ -0,0 +1,222 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The MuseSparkOpen input processor with the prefix-tokenization cache. + +With ``enable_tokenization_cache`` the processor's text-only path splices the +token ids of a cached earlier prompt with a freshly tokenized tail. The ids +must equal what the uncached path produces, turn by turn, for the prompt +shapes an agentic conversation has under this checkpoint's chat template: +multi-turn growth, consecutive assistant tool-call messages (whose end token +changes once another message follows), reasoning blocks, tool results, and the +``Current date`` line at the top of the system message. + +Runs against the real checkpoint's tokenizer and template, like +``tests/unittest/llmapi/test_musespark_open_tokenizer.py``; it is skipped when +``MUSESPARK_REFERENCE_CKPT`` is unset. +""" + +import os + +import pytest + +from tensorrt_llm.sampling_params import SamplingParams + +pytestmark = pytest.mark.cpu_only + +CHECKPOINT = os.environ.get("MUSESPARK_REFERENCE_CKPT", "") + +TOOLS = [ + { + "type": "function", + "function": { + "name": "Bash", + "description": "Run a shell command and return its output. " * 20, + "parameters": { + "type": "object", + "properties": {"command": {"type": "string"}, "timeout": {"type": "number"}}, + "required": ["command"], + }, + }, + }, + { + "type": "function", + "function": { + "name": "Read", + "description": "Read a file from the local filesystem. " * 20, + "parameters": { + "type": "object", + "properties": {"file_path": {"type": "string"}}, + "required": ["file_path"], + }, + }, + }, +] + +FILLER = ( + "Line {i}: the build log reports 1234567 bytes copied in 0.25s; " + "naïve café ünïcode, tabs\tand double spaces, and a URL " + "https://example.com/a?b=c&d=e#f.\n" +) + + +@pytest.fixture(scope="module") +def tokenizer(): + if not CHECKPOINT or not os.path.isdir(CHECKPOINT): + pytest.skip("checkpoint not present; set MUSESPARK_REFERENCE_CKPT") + from tensorrt_llm.tokenizer import load_custom_tokenizer + + return load_custom_tokenizer("musespark_open", CHECKPOINT) + + +def _processor(tokenizer, enable): + from tensorrt_llm.inputs.registry import create_input_processor + + return create_input_processor( + CHECKPOINT, tokenizer, trust_remote_code=True, enable_tokenization_cache=enable + ) + + +def _render(tokenizer, messages, current_date="2026-10-20"): + return tokenizer.tokenizer.apply_chat_template( + messages, + tools=TOOLS, + add_generation_prompt=True, + tokenize=False, + current_date=current_date, + ) + + +def _filler(start, count): + return "".join(FILLER.format(i=i) for i in range(start, start + count)) + + +def _tool_call(call_id, name, **arguments): + return {"id": call_id, "type": "function", "function": {"name": name, "arguments": arguments}} + + +def _tool_result(call_id, step): + return {"role": "tool", "tool_call_id": call_id, "content": _filler(step * 7, 30)} + + +def _agentic_turns(): + """Message lists of one growing agentic conversation, one per request.""" + messages = [ + {"role": "system", "content": "You are a coding agent. " + "Be careful. " * 400}, + {"role": "user", "content": "Fix the failing test.\n" + _filler(0, 200)}, + ] + turns = [list(messages)] + for step in range(12): + if step % 4 == 1: + # Two tool calls emitted as consecutive assistant messages. The + # first one ends with <|eot|> while it is the last message and with + # <|eom|> once the second follows, so the next prompt does not + # extend this one character for character. + messages.append( + { + "role": "assistant", + "reasoning_content": f"I should list files (step {step}).", + "content": "", + "tool_calls": [_tool_call(f"a{step}", "Bash", command=f"ls -la dir{step}")], + } + ) + turns.append(list(messages)) + messages.append( + { + "role": "assistant", + "content": "", + "tool_calls": [_tool_call(f"b{step}", "Read", file_path=f"/src/f{step}.py")], + } + ) + messages.append(_tool_result(f"a{step}", step)) + messages.append(_tool_result(f"b{step}", step)) + elif step % 4 == 3: + # A reasoning block plus a final answer, then a new user turn. + messages.append( + { + "role": "assistant", + "reasoning_content": "Thinking " * 50, + "content": f"Done with step {step}.", + } + ) + messages.append({"role": "user", "content": f"Next: step{step + 1}" + "x" * step}) + else: + messages.append( + { + "role": "assistant", + "reasoning_content": f"Run the tests ({step}).", + "content": "Running the tests now.", + "tool_calls": [ + _tool_call(f"c{step}", "Bash", command="pytest -x", timeout=600) + ], + } + ) + messages.append(_tool_result(f"c{step}", step)) + turns.append(list(messages)) + return turns + + +def test_cache_is_wired_only_when_enabled(tokenizer): + from tensorrt_llm._torch.models.modeling_musespark_open_omni import ( + MuseSparkOpenOmniInputProcessor, + ) + + on = _processor(tokenizer, True) + off = _processor(tokenizer, False) + assert isinstance(on, MuseSparkOpenOmniInputProcessor) + assert on._prefix_token_cache is not None + assert off._prefix_token_cache is None + + +def test_agentic_turns_are_id_identical(tokenizer): + on = _processor(tokenizer, True) + off = _processor(tokenizer, False) + params = SamplingParams(add_special_tokens=False) + turns = _agentic_turns() + for turn, messages in enumerate(turns): + prompt = _render(tokenizer, messages) + expected = tokenizer.encode(prompt, add_special_tokens=False) + cached, extra = on({"prompt": prompt}, params) + uncached, _ = off({"prompt": prompt}, params) + assert extra is None + assert uncached == expected, f"turn {turn}: uncached path changed" + assert cached == expected, f"turn {turn}: cached ids differ" + cache = on._prefix_token_cache + assert not cache.disabled + # Every turn but the first and the three that follow a lone assistant + # tool call extends the previous prompt, so the cache must splice them. + assert cache.hits >= len(turns) - 4, cache.summary() + + +def test_date_change_misses_but_stays_identical(tokenizer): + on = _processor(tokenizer, True) + params = SamplingParams(add_special_tokens=False) + turns = _agentic_turns() + for date, messages in (("2026-10-20", turns[3]), ("2026-10-21", turns[4])): + prompt = _render(tokenizer, messages, current_date=date) + cached, _ = on({"prompt": prompt}, params) + assert cached == tokenizer.encode(prompt, add_special_tokens=False) + # The date sits at the top of the system message, so the second prompt + # does not extend the first one. + assert on._prefix_token_cache.hits == 0 + + +@pytest.mark.parametrize("stride_offset", [1, 7, 31, 97, 1009]) +def test_appends_starting_mid_word_are_identical(tokenizer, stride_offset): + on = _processor(tokenizer, True) + params = SamplingParams(add_special_tokens=False) + full = _render(tokenizer, _agentic_turns()[-1]) + for end in range(len(full) // 2, len(full), len(full) // 40 + stride_offset): + prompt = full[:end] + cached, _ = on({"prompt": prompt}, params) + assert cached == tokenizer.encode(prompt, add_special_tokens=False), end + assert on._prefix_token_cache.hits > 0 + + +def test_cache_is_bypassed_when_special_tokens_are_requested(tokenizer): + on = _processor(tokenizer, True) + prompt = _render(tokenizer, _agentic_turns()[2]) + ids, _ = on({"prompt": prompt}, SamplingParams(add_special_tokens=True)) + # The processor never adds special tokens itself; the template carries BOS. + assert ids == tokenizer.encode(prompt, add_special_tokens=False) + cache = on._prefix_token_cache + assert cache.hits + cache.misses == 0 From d6859f16dd8fd479f445e241946722351e725637 Mon Sep 17 00:00:00 2001 From: Xiao Wang <24860335+xwang233@users.noreply.github.com> Date: Mon, 26 Oct 2026 18:43:55 +0000 Subject: [PATCH 04/16] [None][feat] musespark_open: attention-input gather hook on sequence-parallel layers MuseSparkOpenDecoderLayer gains _sp_attn_input_gather. When a model-specific path sets it on a layer (with a QKV projection that consumes its result), the sequence-parallel context step gathers the attention input through it instead of the NCCL all-gather. Off by default: the class attribute is None, and the layer is unchanged unless a path sets it. Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com> --- tensorrt_llm/_torch/models/modeling_musespark_open.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/models/modeling_musespark_open.py b/tensorrt_llm/_torch/models/modeling_musespark_open.py index babc872064d3..0c2dd676ddec 100644 --- a/tensorrt_llm/_torch/models/modeling_musespark_open.py +++ b/tensorrt_llm/_torch/models/modeling_musespark_open.py @@ -2524,6 +2524,11 @@ def forward( class MuseSparkOpenDecoderLayer(DecoderLayer): + #: The sequence-parallel arm's attention-input all-gather, ``(attn_input, sp_num_tokens) -> + #: rows``; None (the default) runs the NCCL all-gather. A model-specific path may set it + #: on a layer, together with a QKV projection that consumes what it returns. + _sp_attn_input_gather = None + def __init__( self, model_config: ModelConfig[MuseSparkOpenConfig], @@ -2627,7 +2632,11 @@ def forward( and not torch.cuda.is_current_stream_capturing() ) if sp_num_tokens is not None: - attn_input = allgather(attn_input, self.mapping, dim=0)[:sp_num_tokens] + gather = self._sp_attn_input_gather + if gather is None: + attn_input = allgather(attn_input, self.mapping, dim=0)[:sp_num_tokens] + else: + attn_input = gather(attn_input, sp_num_tokens) kwargs = dict(kwargs, all_reduce_params=AllReduceParams(enable_allreduce=False)) attn_out = self.self_attn( position_ids=position_ids, From 92ce9363c007ce0d7a542fe9637cff6e1572d1a8 Mon Sep 17 00:00:00 2001 From: Xiao Wang <24860335+xwang233@users.noreply.github.com> Date: Mon, 26 Oct 2026 18:44:05 +0000 Subject: [PATCH 05/16] [None][feat] Rubin CuTe DSL MoE GEMMs: TRTLLM-gen weight layouts, permuted store and 64-deep MMA - The Rubin gather + grouped GEMM + SwiGLU kernel (FC1) and its custom op take trtllm_gen_gated_layout: the weights are TRTLLM-gen's gated FC1 tensors, alpha holds the gate scales followed by the output scales, and the epilogue follows TRTLLM-gen's SwiGLU and NVFP4 quantization arithmetic. - The Rubin grouped GEMM finalize kernel (FC2) and its custom op take store_permuted (one output row per permuted row, alpha * acc, no token scale and no scatter), trtllm_gen_n_shuffle (B and its block scales in TRTLLM-gen's shuffled row order; output columns in logical order) and mma_k64 (the 64-deep SM100-family MMA, whose fp32 accumulators equal those of SM100-family NVFP4 GEMMs). - Kernel and tuning-config caches are keyed on the new options; the tiled MMA construction is factored into _make_tiled_mma with the same arguments. All four options are constructor arguments, False by default; with every option False the kernels and ops build and run as before. Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com> --- .../_torch/custom_ops/cute_dsl_custom_ops.py | 51 +++- ...her_grouped_blockscaled_gemm_act_fusion.py | 156 ++++++++++++- ...rouped_blockscaled_gemm_finalize_fusion.py | 217 ++++++++++++------ 3 files changed, 337 insertions(+), 87 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 e641111adb4e..b0e365826d1f 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py @@ -14784,8 +14784,13 @@ def __init__( scaling_vector_size: int = 16, activation_type: ActivationType = ActivationType.Swiglu, situ_beta: Optional[float] = None, - situ_linear_beta: Optional[float] = None): + situ_linear_beta: Optional[float] = None, + trtllm_gen_gated_layout: bool = False): super().__init__() + # trtllm_gen_gated_layout: the weights are TRTLLM-gen's gated FC1 + # tensors, alpha is [gate scales | output scales] (2 * L values) + # and the epilogue follows TRTLLM-gen's arithmetic. + self.trtllm_gen_gated_layout = trtllm_gen_gated_layout self.num_experts = num_experts self.top_k = top_k self.num_local_experts = num_local_experts @@ -14819,6 +14824,7 @@ def unique_id(self): int(self.activation_type), self.situ_beta, self.situ_linear_beta, + self.trtllm_gen_gated_layout, ) def get_valid_tactics( @@ -14998,7 +15004,8 @@ def forward(self, inputs: List[torch.Tensor], assert b_sf.size(0) == l assert b_sf.size(1) == n assert b_sf.size(2) == scale_k - assert alpha.size(0) == l + assert alpha.size(0) == l * (2 if self.trtllm_gen_gated_layout + else 1) num_tiles = m // self.tile_size assert tile_idx_to_group_idx.dtype == torch.int32 @@ -15122,7 +15129,8 @@ def forward(self, inputs: List[torch.Tensor], cluster_shape_mn, raster_along_m, locality_domain_half_gemm, a_path, int(self.activation_type), self.situ_beta, - self.situ_linear_beta, max_active_clusters) + self.situ_linear_beta, max_active_clusters, + self.trtllm_gen_gated_layout) if cache_key not in self.__class__.kernel_cache: gemm = self.__class__.kernel_class( sf_vec_size=self.scaling_vector_size, @@ -15137,6 +15145,7 @@ def forward(self, inputs: List[torch.Tensor], activation_type=self.activation_type, situ_beta=self.situ_beta, situ_linear_beta=self.situ_linear_beta, + trtllm_gen_gated_layout=self.trtllm_gen_gated_layout, ) compiled_gemm = cute.compile( gemm.wrapper, @@ -16161,8 +16170,20 @@ def __init__(self, local_expert_offset: int, tile_size: int, output_dtype: torch.dtype, - scaling_vector_size: int = 16): + scaling_vector_size: int = 16, + store_permuted: bool = False, + trtllm_gen_n_shuffle: bool = False, + mma_k64: bool = False): super().__init__() + # store_permuted: C gets one row per permuted row (alpha * acc, + # no token scale, no scatter); see the kernel class. + # trtllm_gen_n_shuffle: the weights are TRTLLM-gen's FC2 tensors + # (shuffle_matrix_a row order); C columns come out in logical order. + # mma_k64: the 64-deep SM100-family MMA (fp32 accumulators + # equal to SM100-family NVFP4 GEMMs'); see the kernel class. + self.store_permuted = store_permuted + self.trtllm_gen_n_shuffle = trtllm_gen_n_shuffle + self.mma_k64 = mma_k64 self.num_experts = num_experts self.top_k = top_k self.num_local_experts = num_local_experts @@ -16192,6 +16213,9 @@ def unique_id(self): self.tile_size, self.output_dtype, self.scaling_vector_size, + self.store_permuted, + self.trtllm_gen_n_shuffle, + self.mma_k64, ) @staticmethod @@ -16291,13 +16315,17 @@ def get_tuning_config(self) -> TuningConfig: self.num_local_experts, self.local_expert_offset, self.tile_size) + # Stored permuted rows: C has as many rows as A. + c_rows = ((lambda shapes: shapes[0][0]) + if self.store_permuted else + helper.infer_shape_num_tokens) self.__class__.tuning_config_cache[key] = TuningConfig( dynamic_tensor_specs=(DynamicTensorSpec( 0, 0, helper.gen_tuning_buckets, helper.map_to_tuning_buckets), ), constraint_specs=( ConstraintSpec(2, 0, fp4_scale_infer_shape), - ConstraintSpec(5, 0, helper.infer_shape_num_tokens), + ConstraintSpec(5, 0, c_rows), ConstraintSpec(6, 0, helper.infer_shape_max_num_tiles), ConstraintSpec(7, 0, @@ -16342,6 +16370,9 @@ def forward(self, inputs: List[torch.Tensor], assert c.dim() == 2 num_tokens = c.size(0) assert c.size(1) == n or c.size(1) == n * 2 + if self.store_permuted: + assert num_tokens == m, ( + "store_permuted needs one C row per permuted row") num_tiles = m // self.tile_size assert tile_idx_to_group_idx.dtype == torch.int32 @@ -16354,7 +16385,8 @@ def forward(self, inputs: List[torch.Tensor], assert num_non_exiting_tiles.numel() == 1 assert token_final_scales.dtype == torch.float32 assert token_final_scales.dim() == 2 - assert token_final_scales.size() == (num_tokens, self.top_k) + if not self.store_permuted: + assert token_final_scales.size() == (num_tokens, self.top_k) locality_domain_id = get_current_locality_domain() locality_domain_half_gemm = locality_domain_id is not None @@ -16434,7 +16466,9 @@ def forward(self, inputs: List[torch.Tensor], cache_key = (self.scaling_vector_size, self.tile_size, self.top_k, mma_tiler, mma_inst_shape, cluster_shape_mn, raster_along_m, - locality_domain_half_gemm, max_active_clusters) + locality_domain_half_gemm, max_active_clusters, + self.store_permuted, self.trtllm_gen_n_shuffle, + self.mma_k64) if cache_key not in self.__class__.kernel_cache: gemm = self.__class__.kernel_class( sf_vec_size=self.scaling_vector_size, @@ -16443,6 +16477,9 @@ def forward(self, inputs: List[torch.Tensor], cluster_shape_mn=cluster_shape_mn, raster_along_m=raster_along_m, topK=self.top_k, + store_permuted=self.store_permuted, + trtllm_gen_n_shuffle=self.trtllm_gen_n_shuffle, + mma_k64=self.mma_k64, ) compiled_gemm = cute.compile( gemm.wrapper, diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/rubin/moe/rubin_contiguous_gather_grouped_blockscaled_gemm_act_fusion.py b/tensorrt_llm/_torch/cute_dsl_kernels/rubin/moe/rubin_contiguous_gather_grouped_blockscaled_gemm_act_fusion.py index 89351bf6c0d1..f38a1a4497b5 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/rubin/moe/rubin_contiguous_gather_grouped_blockscaled_gemm_act_fusion.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/rubin/moe/rubin_contiguous_gather_grouped_blockscaled_gemm_act_fusion.py @@ -203,8 +203,17 @@ def __init__( activation_type: ActivationType = ActivationType.Swiglu, situ_beta: Optional[float] = None, situ_linear_beta: Optional[float] = None, + trtllm_gen_gated_layout: bool = False, ): self.a_path = a_path + # trtllm_gen_gated_layout: B (and SFB) rows are in TRTLLM-gen's gated + # FC1 layout (see _trtllm_gen_gated_split) instead of the granularity-64 + # [up, gate] interleave, alpha holds two values per group (the gate + # scale at [g], the output scale at [l + g], see + # _apply_trtllm_gen_swiglu_epilogue), and the SwiGLU + NVFP4 + # quantization use TRTLLM-gen's arithmetic, so the output codes and + # block scales equal TRTLLM-gen's gated-activation GEMM bit for bit. + self.trtllm_gen_gated_layout = trtllm_gen_gated_layout # locality domain half-GEMM: two partitions write their N-half into a shared # full-width C/SFC buffer at a column offset (see wrapper/__call__). self.locality_domain_half_gemm = locality_domain_half_gemm @@ -231,6 +240,12 @@ def __init__( self.is_gated = is_gated_activation(self.activation_type) if locality_domain_half_gemm and not self.is_gated: raise ValueError("Rubin locality domain half-GEMM currently supports SwiGLU only") + if trtllm_gen_gated_layout and ( + self.activation_type != ActivationType.Swiglu or locality_domain_half_gemm + ): + raise ValueError( + "trtllm_gen_gated_layout supports SwiGLU without the locality domain half-GEMM only" + ) self.acc_dtype = cutlass.Float32 self.mma_inst_shape = mma_inst_shape self.mma_tiler = mma_tiler @@ -2934,7 +2949,12 @@ def kernel( # expert_idx = mma_tile_coord_mnl[2] - alpha_val = alpha[expert_idx] + if cutlass.const_expr(self.trtllm_gen_gated_layout): + # gate scale; output scale (applied by the quantization) + alpha_val = alpha[(expert_idx, 0)] + scale_c_val = alpha[(expert_idx, 1)] + else: + alpha_val = alpha[expert_idx] # # Slice to per mma tile index @@ -3046,7 +3066,12 @@ def kernel( acc_vec_up = tTR_rAcc_up.load() tCompute = cute.make_rmem_tensor(acc_vec_up.shape, self.acc_dtype) - if cutlass.const_expr(self.activation_type == ActivationType.Swiglu): + if cutlass.const_expr(self.trtllm_gen_gated_layout): + tUp, tGate = self._trtllm_gen_gated_split(tTR_rAcc_up, tTR_rAcc_gate) + self._apply_trtllm_gen_swiglu_epilogue( + tUp.load(), tGate.load(), alpha_val, tCompute + ) + elif cutlass.const_expr(self.activation_type == ActivationType.Swiglu): acc_vec_gate = tTR_rAcc_gate.load() self._apply_swiglu_epilogue( acc_vec_up, acc_vec_gate, alpha_val, tCompute @@ -3094,7 +3119,31 @@ def kernel( abs_acc_frg_ir, acc_frg.shape, acc_frg.dtype ) - if cutlass.const_expr(self.vectorized_f32): + if cutlass.const_expr(self.trtllm_gen_gated_layout): + # TRTLLM-gen: SF = e4m3((amax * scale_c) * (1/6)) + for vi in cutlass.range_constexpr(abs_acc_frg.shape[1]): + tCrSFC_pvscale[vi] = abs_acc_frg[None, vi].reduce( + cute.ReductionOp.MAX, + cutlass.Float32(0.0), + 0, # Use 0.0 as init for abs values + ) + for vi in cutlass.range_constexpr(0, abs_acc_frg.shape[1], 2): + tCrSFC_pvscale[vi], tCrSFC_pvscale[vi + 1] = ( + cute.arch.mul_packed_f32x2( + (tCrSFC_pvscale[vi], tCrSFC_pvscale[vi + 1]), + (scale_c_val, scale_c_val), + ) + ) + tCrSFC_pvscale[vi], tCrSFC_pvscale[vi + 1] = ( + cute.arch.mul_packed_f32x2( + (tCrSFC_pvscale[vi], tCrSFC_pvscale[vi + 1]), + ( + self.get_dtype_rcp_limits(self.c_dtype), + self.get_dtype_rcp_limits(self.c_dtype), + ), + ) + ) + elif cutlass.const_expr(self.vectorized_f32): for vi in cutlass.range_constexpr(abs_acc_frg.shape[1]): tCrSFC_pvscale[vi] = abs_acc_frg[None, vi].reduce( cute.ReductionOp.MAX, @@ -3146,7 +3195,31 @@ def kernel( # TODO: need to add f8x2 -> f32x2 conversion tCrSFC_qpvscale_up = tCrSFC.load().to(cutlass.Float32) fp32_max = cutlass.Float32(3.40282346638528859812e38) - if cutlass.const_expr(self.vectorized_f32): + if cutlass.const_expr(self.trtllm_gen_gated_layout): + # TRTLLM-gen: q = (x * rcp(SF)) * scale_c + for vi in cutlass.range_constexpr(0, cute.size(tCrSFC), 2): + acc_scale_min0 = fmin( + cute.arch.rcp_approx(tCrSFC_qpvscale_up[vi]), + fp32_max, + nan=True, + ) + acc_scale_min1 = fmin( + cute.arch.rcp_approx(tCrSFC_qpvscale_up[vi + 1]), + fp32_max, + nan=True, + ) + vec0 = tTR_rAcc_frg[None, vi] + vec1 = tTR_rAcc_frg[None, vi + 1] + for ei in cutlass.range_constexpr(self.sf_vec_size): + vec0[ei], vec1[ei] = cute.arch.mul_packed_f32x2( + (vec0[ei], vec1[ei]), + (acc_scale_min0, acc_scale_min1), + ) + vec0[ei], vec1[ei] = cute.arch.mul_packed_f32x2( + (vec0[ei], vec1[ei]), + (scale_c_val, scale_c_val), + ) + elif cutlass.const_expr(self.vectorized_f32): for vi in cutlass.range_constexpr(0, cute.size(tCrSFC), 2): acc_scale = cute.arch.mul_packed_f32x2( ( @@ -3253,6 +3326,76 @@ def kernel( cute.arch.griddepcontrol_launch_dependents() + @cute.jit + def _trtllm_gen_gated_split(self, tAccLo: cute.Tensor, tAccHi: cute.Tensor): + """Gather up / gate for 64 output columns from TRTLLM-gen gated FC1 accumulators. + + ``tAccLo`` / ``tAccHi`` are one thread's two consecutive 64-column + accumulator subtiles (128 B rows). In TRTLLM-gen's gated layout every + 32-row block of B covers 16 intermediate columns: rows 0-7 are up of the + even columns, 8-15 gate of the even columns, 16-23 up of the odd + columns and 24-31 gate of the odd columns. Output column ``j`` of the + 64-column output subtile therefore takes, in subtile ``j // 32`` with + ``r = j % 32 = 16 g + 2 kk + o``, up from column ``32 g + 16 o + kk`` + and gate from column ``32 g + 16 o + 8 + kk``. Only register indices + move (compile-time constants). + """ + tUp = cute.make_rmem_tensor(tAccLo.shape, self.acc_dtype) + tGate = cute.make_rmem_tensor(tAccLo.shape, self.acc_dtype) + for j in cutlass.range_constexpr(64): + r = j % 32 + g = r // 16 + o = r % 2 + kk = (r % 16) // 2 + src_up = 32 * g + 16 * o + kk + src = tAccLo if j < 32 else tAccHi + tUp[j] = src[src_up] + tGate[j] = src[src_up + 8] + return tUp, tGate + + @cute.jit + def _apply_trtllm_gen_swiglu_epilogue( + self, + acc_vec_up: cute.Tensor, + acc_vec_gate: cute.Tensor, + scale_gate, + tCompute: cute.Tensor, + ): + """SwiGLU in TRTLLM-gen's gated-activation GEMM arithmetic (fp32, in this order):: + + k = log2(e) * scale_gate + s = rcp(1 + exp2(-(gate[i] * k))) + tCompute[i] = (up[i] * scale_gate) * (gate[i] * s) + + i.e. ``scale_gate * up * silu(scale_gate * gate)``, with the gate scale + folded into the exponent constant and the up half. The output scale is + applied by the quantization (``SF = e4m3((amax * scale_c) / 6)``, + ``q = (x * rcp(SF)) * scale_c``). + """ + k = cutlass.Float32(1.4426950408889634) * cutlass.Float32(scale_gate) + for i in cutlass.range_constexpr(0, cute.size(acc_vec_up.shape), 2): + arg = cute.arch.mul_packed_f32x2( + (acc_vec_gate[i], acc_vec_gate[i + 1]), + (k, k), + ) + arg = cute.arch.mul_packed_f32x2(arg, (-1.0, -1.0)) + den = cute.arch.add_packed_f32x2( + ( + cute.math.exp2(arg[0], fastmath=True), + cute.math.exp2(arg[1], fastmath=True), + ), + (1.0, 1.0), + ) + act = cute.arch.mul_packed_f32x2( + (acc_vec_gate[i], acc_vec_gate[i + 1]), + (cute.arch.rcp_approx(den[0]), cute.arch.rcp_approx(den[1])), + ) + lin = cute.arch.mul_packed_f32x2( + (acc_vec_up[i], acc_vec_up[i + 1]), + (cutlass.Float32(scale_gate), cutlass.Float32(scale_gate)), + ) + tCompute[i], tCompute[i + 1] = cute.arch.mul_packed_f32x2(lin, act) + @cute.jit def _apply_swiglu_epilogue( self, @@ -4145,7 +4288,10 @@ def wrapper( order=(2, 1, 4, 0, 3, 5), ), ) - alpha = cute.make_tensor(alpha_ptr, layout=cute.make_layout((l,))) + if cutlass.const_expr(self.trtllm_gen_gated_layout): + alpha = cute.make_tensor(alpha_ptr, layout=cute.make_layout((l, 2), stride=(1, l))) + else: + alpha = cute.make_tensor(alpha_ptr, layout=cute.make_layout((l,))) tile_idx_to_group_idx = cute.make_tensor( tile_idx_to_group_idx_ptr, layout=cute.make_layout((num_tiles,)) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/rubin/moe/rubin_contiguous_grouped_blockscaled_gemm_finalize_fusion.py b/tensorrt_llm/_torch/cute_dsl_kernels/rubin/moe/rubin_contiguous_grouped_blockscaled_gemm_finalize_fusion.py index 7cf15e17f995..73412e40f528 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/rubin/moe/rubin_contiguous_grouped_blockscaled_gemm_finalize_fusion.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/rubin/moe/rubin_contiguous_grouped_blockscaled_gemm_finalize_fusion.py @@ -117,6 +117,24 @@ def blk_reduce_fp32(dst_gemm, src_smem, size, loc=None, ip=None): ) +@dsl_user_op +def blk_store(dst_gemm, src_smem, size, loc=None, ip=None): + """Bulk copy of ``size`` bytes from shared to global memory (cp.async.bulk).""" + llvm.inline_asm( + None, + [ + dst_gemm.iterator.llvm_ptr, + src_smem.iterator.llvm_ptr, + size.ir_value(), + ], + "cp.async.bulk.global.shared::cta.bulk_group [$0], [$1], $2;", + "l,l,r", + has_side_effects=True, + loc=loc, + ip=ip, + ) + + @dsl_user_op def blk_reduce_fp16(dst_gemm, src_smem, size, loc=None, ip=None): """Block reduce for FP16 using cp.reduce.async.bulk.""" @@ -171,9 +189,30 @@ def __init__( cluster_shape_mn: Tuple[int, int], raster_along_m: bool = False, topK: int = 1, + store_permuted: bool = False, + trtllm_gen_n_shuffle: bool = False, + mma_k64: bool = False, ): + # store_permuted: the epilogue stores each valid permuted row + # (alpha * acc, no token scale) to row ``permuted_row`` of C, which then + # has one row per permuted row, instead of scatter-adding the + # token-scaled row into its token's row. + # trtllm_gen_n_shuffle: B (and SFB) rows are in TRTLLM-gen's + # shuffle_matrix_a order (epilogue_tile_m 128: within every 32 rows, + # stored row p is logical row 4 * (p % 8) + p // 8); the epilogue writes + # C columns back in logical order. + # mma_k64: issue the 64-deep tcgen05 kind::mxf4nvf4 MMA (the SM100 family + # instruction) instead of the 128-deep SM107 one. The fp32 accumulators + # then equal those of SM100-family NVFP4 GEMMs (e.g. TRTLLM-gen's) bit + # for bit; the K=128 instruction rounds its partial sums differently. + # mma_tiler[2] stays 256 (four MMAs per k-tile). + self.store_permuted = store_permuted + self.trtllm_gen_n_shuffle = trtllm_gen_n_shuffle + self.mma_k64 = mma_k64 self.sf_vec_size = sf_vec_size self.acc_dtype = cutlass.Float32 + if mma_k64: + mma_inst_shape = (mma_inst_shape[0], mma_inst_shape[1], 64) self.mma_inst_shape = mma_inst_shape self.mma_tiler = mma_tiler self.cluster_shape_mn = cluster_shape_mn @@ -231,6 +270,49 @@ def __init__( # B-reuse pattern control self.enable_breuse = True if mma_tiler[0] // mma_inst_shape[0] == 2 else False + if mma_k64 and self.enable_breuse: + raise ValueError("mma_k64 does not support the B-reuse pattern") + + def _make_tiled_mma(self, cta_group, mma_inst_shape, permutation_mnk=None): + """The block-scaled tiled MMA (SM107 K=128, or the SM100-family K=64 one with mma_k64).""" + if cutlass.const_expr(self.mma_k64): + return sm100_utils.make_blockscaled_trivial_tiled_mma( + self.a_dtype, + self.b_dtype, + self.a_major_mode, + self.b_major_mode, + self.sf_dtype, + self.sf_vec_size, + cta_group, + mma_inst_shape[:2], + ) + if cutlass.const_expr(permutation_mnk is None): + return sm107_utils.make_blockscaled_trivial_tiled_mma( + self.a_dtype, + self.b_dtype, + self.a_major_mode, + self.b_major_mode, + self.sf_dtype, + self.sf_vec_size, + cta_group, + mma_inst_shape, + a_collector_op=CollectorOp.DISCARD, + b_collector_op=CollectorOp.DISCARD, + ) + return sm107_utils.make_blockscaled_trivial_tiled_mma( + self.a_dtype, + self.b_dtype, + self.a_major_mode, + self.b_major_mode, + self.sf_dtype, + self.sf_vec_size, + cta_group, + mma_inst_shape, + a_collector_op=CollectorOp.DISCARD, + b_collector_op=CollectorOp.DISCARD, + atom_layout_mnk=(1, 1, 1), + permutation_mnk=permutation_mnk, + ) def _get_mma_permutation_mnk(self): if cutlass.const_expr(self.use_2cta_instrs and self.enable_breuse): @@ -312,32 +394,12 @@ def _setup_attributes(self): self.mma_inst_shape[2], ) - tiled_mma = sm107_utils.make_blockscaled_trivial_tiled_mma( - self.a_dtype, - self.b_dtype, - self.a_major_mode, - self.b_major_mode, - self.sf_dtype, - self.sf_vec_size, - self.cta_group, - self.mma_inst_shape, - a_collector_op=CollectorOp.DISCARD, - b_collector_op=CollectorOp.DISCARD, - atom_layout_mnk=(1, 1, 1), - permutation_mnk=self._get_mma_permutation_mnk(), + tiled_mma = self._make_tiled_mma( + self.cta_group, self.mma_inst_shape, self._get_mma_permutation_mnk() ) - tiled_mma_sfb = sm107_utils.make_blockscaled_trivial_tiled_mma( - self.a_dtype, - self.b_dtype, - self.a_major_mode, - self.b_major_mode, - self.sf_dtype, - self.sf_vec_size, - cute.nvgpu.tcgen05.CtaGroup.ONE, - self.mma_inst_shape_sfb, - a_collector_op=CollectorOp.DISCARD, - b_collector_op=CollectorOp.DISCARD, + tiled_mma_sfb = self._make_tiled_mma( + cute.nvgpu.tcgen05.CtaGroup.ONE, self.mma_inst_shape_sfb ) # Compute mma/cluster/tile shapes @@ -633,35 +695,14 @@ def __call__( atom_layout_mnk = (1, 1, 1) permutation_mnk = self._get_mma_permutation_mnk() - tiled_mma = sm107_utils.make_blockscaled_trivial_tiled_mma( - self.a_dtype, - self.b_dtype, - self.a_major_mode, - self.b_major_mode, - self.sf_dtype, - self.sf_vec_size, - self.cta_group, - self.mma_inst_shape, - a_collector_op=CollectorOp.DISCARD, - b_collector_op=CollectorOp.DISCARD, - atom_layout_mnk=atom_layout_mnk, - permutation_mnk=permutation_mnk, - ) + assert atom_layout_mnk == (1, 1, 1) + tiled_mma = self._make_tiled_mma(self.cta_group, self.mma_inst_shape, permutation_mnk) tiled_mma.set(tcgen05.Field.NEGATE_A, False) tiled_mma.set(tcgen05.Field.NEGATE_B, False) - tiled_mma_sfb = sm107_utils.make_blockscaled_trivial_tiled_mma( - self.a_dtype, - self.b_dtype, - self.a_major_mode, - self.b_major_mode, - self.sf_dtype, - self.sf_vec_size, - cute.nvgpu.tcgen05.CtaGroup.ONE, - self.mma_inst_shape_sfb, - a_collector_op=CollectorOp.DISCARD, - b_collector_op=CollectorOp.DISCARD, + tiled_mma_sfb = self._make_tiled_mma( + cute.nvgpu.tcgen05.CtaGroup.ONE, self.mma_inst_shape_sfb ) tiled_mma_sfb.set(tcgen05.Field.NEGATE_A, False) tiled_mma_sfb.set(tcgen05.Field.NEGATE_B, False) @@ -1587,11 +1628,12 @@ def kernel( # Compute token info and scaled alpha for this row token_idx = cutlass.Int32(0) alpha_val_iter = alpha_val - if is_valid_row: - token_idx = expanded_idx // self.topK - topk_idx = expanded_idx % self.topK - token_scale = token_final_scales[(token_idx, topk_idx)] - alpha_val_iter = alpha_val * token_scale + if cutlass.const_expr(not self.store_permuted): + if is_valid_row: + token_idx = expanded_idx // self.topK + topk_idx = expanded_idx % self.topK + token_scale = token_final_scales[(token_idx, topk_idx)] + alpha_val_iter = alpha_val * token_scale for n_iter_idx in cutlass.range(n_subtile_cnt): # Load accumulator from tensor memory buffer to register @@ -1600,7 +1642,12 @@ def kernel( cute.copy(tiled_copy_t2r, tTR_tAcc_mn, tTR_rAcc) # For block reduce: retile then load, compute, store - acc_vec = tiled_copy_r2s.retile(tTR_rAcc).load() + if cutlass.const_expr(self.trtllm_gen_n_shuffle): + acc_vec = tiled_copy_r2s.retile( + self._trtllm_gen_n_unshuffle(tTR_rAcc) + ).load() + else: + acc_vec = tiled_copy_r2s.retile(tTR_rAcc).load() acc_vec_final = (alpha_val_iter * acc_vec).to(self.c_dtype) tRS_rC.store(acc_vec_final) if is_valid_row: @@ -1628,32 +1675,38 @@ def kernel( is_valid_row = permuted_row < tile_info[4] if is_valid_row: - expanded_idx = permuted_idx_to_expanded_idx[permuted_row] - token_idx = expanded_idx // self.topK coord_n = mma_tile_coord_mnl[1] * self.cta_tile_shape_mnk[1] - scatter_out_offset = cute.domain_offset((token_idx, coord_n, 0), c) - # sC row index: m_iter_idx * 128 + epi_tidx sC_row = m_iter_idx * 128 + epi_tidx - - if cutlass.const_expr(self.c_dtype == cutlass.BFloat16): - blk_reduce_bf16( - scatter_out_offset, - sC[sC_row, None, 0], - cutlass.Int32(self.copy_size), - ) - elif cutlass.const_expr(self.c_dtype == cutlass.Float32): - blk_reduce_fp32( - scatter_out_offset, - sC[sC_row, None, 0], - cutlass.Int32(self.copy_size), - ) - elif cutlass.const_expr(self.c_dtype == cutlass.Float16): - blk_reduce_fp16( - scatter_out_offset, + if cutlass.const_expr(self.store_permuted): + store_out_offset = cute.domain_offset((permuted_row, coord_n, 0), c) + blk_store( + store_out_offset, sC[sC_row, None, 0], cutlass.Int32(self.copy_size), ) + else: + expanded_idx = permuted_idx_to_expanded_idx[permuted_row] + token_idx = expanded_idx // self.topK + scatter_out_offset = cute.domain_offset((token_idx, coord_n, 0), c) + if cutlass.const_expr(self.c_dtype == cutlass.BFloat16): + blk_reduce_bf16( + scatter_out_offset, + sC[sC_row, None, 0], + cutlass.Int32(self.copy_size), + ) + elif cutlass.const_expr(self.c_dtype == cutlass.Float32): + blk_reduce_fp32( + scatter_out_offset, + sC[sC_row, None, 0], + cutlass.Int32(self.copy_size), + ) + elif cutlass.const_expr(self.c_dtype == cutlass.Float16): + blk_reduce_fp16( + scatter_out_offset, + sC[sC_row, None, 0], + cutlass.Int32(self.copy_size), + ) cute.arch.cp_async_bulk_commit_group() cute.arch.cp_async_bulk_wait_group(0, read=True) @@ -1676,6 +1729,20 @@ def kernel( self.epilog_sync_barrier.arrive_and_wait() tmem.free(acc_tmem_ptr) + @cute.jit + def _trtllm_gen_n_unshuffle(self, tAcc: cute.Tensor) -> cute.Tensor: + """One thread's accumulator columns back in logical order (trtllm_gen_n_shuffle). + + Column q of every 32-column group takes accumulator column + ``8 * (q % 4) + q // 4``. Only register indices move. + """ + tOut = cute.make_rmem_tensor(tAcc.shape, self.acc_dtype) + for q in cutlass.range_constexpr(cute.size(tAcc)): + base = q - q % 32 + r = q % 32 + tOut[q] = tAcc[base + 8 * (r % 4) + r // 4] + return tOut + @staticmethod def _compute_grid( gemm_shape: Tuple[int, int, int], From 314cd02c056e60d38be547ce6a32ae822d49f87f Mon Sep 17 00:00:00 2001 From: Xiao Wang <24860335+xwang233@users.noreply.github.com> Date: Mon, 26 Oct 2026 18:44:06 +0000 Subject: [PATCH 06/16] [None][feat] modeling_v2 musespark_open tep: mixed-step decode attention, per-SM prep and split-KV launches - mixed_step_gen_rows (Rubin, sm107 TEP4): on the full-attention layers, a mixed step with 9..64 generation requests runs its context rows on the thop op (context phase only) and its generation rows on multi_request_split_kv's decode stand-in, as decode steps do. Sliding-window layers, pure context steps and mixed steps outside the band keep the thop attention. - decode_kv_splits on mixed steps (GB300, sm103 TEP4): a mixed step with 9..64 generation requests arms the decode KV-split rule for the generation phase's trtllm-gen decode FMHA call, as decode steps of 72..512 rows already do. The split stand-in counts the calls it split and the calls it kept. - multi_request_prep_append: on sm103 a packed launch geometry per row count (items per program, warps, vector lane loads and stores) and an L2 prefetch of the next projection's weight slice; sm107 keeps the grid of one program per (row, head). Both write the same bytes. - multi_request_split_kv: the ring depth of KV pages a CTA holds in flight is chosen per SM version and number of Q terms (on sm103 a deeper ring for a one-term Q; sm107 keeps its depth). Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com> --- .../kernels/decode_kv_splits/mixed.py | 125 +++++ .../kernels/decode_kv_splits/splits.py | 6 + .../kernels/mixed_step_gen_rows/__init__.py | 3 + .../kernels/mixed_step_gen_rows/attention.py | 257 +++++++++++ .../multi_request_prep_append/attention.py | 11 +- .../attention/multi_request_prep_append.py | 428 +++++++++++++++++- .../ops/attention/multi_request_split_kv.py | 31 +- .../multi_request_split_kv_kernel.py | 2 +- 8 files changed, 845 insertions(+), 18 deletions(-) create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/decode_kv_splits/mixed.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/mixed_step_gen_rows/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/mixed_step_gen_rows/attention.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/decode_kv_splits/mixed.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/decode_kv_splits/mixed.py new file mode 100644 index 000000000000..d77863d03d36 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/decode_kv_splits/mixed.py @@ -0,0 +1,125 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""decode_kv_splits on the generation rows of mixed steps. + +A mixed step (context rows and generation rows in one batch) runs ``PrefillTarget``, +the built-in text-model forward. Each layer's built-in attention runs the context +phase (flashinfer's trtllm-gen context FMHA) and then the generation phase of the +batch's generation sequences (``run_generation`` with ``seq_offset = num_contexts``), +which launches flashinfer's trtllm-gen decode FMHA over those sequences. flashinfer +picks that call's KV split from the SM count exactly as on a decode step (``splits``'s +module docstring), so a mixed step of 20 or more generation requests on a TEP4 rank +runs the persistent one-CTA-per-request kernel on every full-attention layer, and its +time is set by the longest of those requests. + +This mixin, re-classed in front of the full-attention layers' attention, arms the +model's ``DecodeKvSplits`` for the duration of one attention call when the step has +context rows and ``takes_generation`` accepts its generation requests (9 to 64). The +split rule then acts on the generation phase's decode call exactly as on the decode +steps ``DecodeTarget`` arms: same SM-count rule, same own counter buffer, same +workspace check, same calls left alone. The context phase calls a different flashinfer +entry point (``flashinfer.prefill``), which the stand-in passes through unchanged. + +Why at the attention call and not around ``PrefillTarget``'s forward: under the +breakable prefill CUDA graph a replay never runs the model's forward; it replays the +captured segments and re-runs the eager attention op between them, which reaches this +module's ``_attn_impl`` with the step's metadata. Arming here applies the rule on every +mixed step, captured or eager. Every read is host-side metadata of the step. + +Decode steps (no context rows) are left to ``DecodeTarget``'s own arming; steps of one +to 8 generation requests and pure context steps run what they ran before. + +Numerics: those of ``decode_kv_splits`` on decode steps (the same trtllm-gen kernel +family, multi-CTA KV with the in-kernel fp32 merge; only the KV partition of each +request changes). +""" + +from __future__ import annotations + +import functools +from typing import Optional + +import torch + +from tensorrt_llm.logger import logger + +from .splits import MAX_ROWS, MIN_ROWS, ROWS_PER_REQUEST, DecodeKvSplits + +#: Generation requests of a mixed step the rule is armed for. +MIN_REQUESTS = MIN_ROWS // ROWS_PER_REQUEST +MAX_REQUESTS = MAX_ROWS // ROWS_PER_REQUEST + + +def takes_generation(attn_metadata) -> bool: + """Whether a step is a mixed step whose generation rows the rule is armed for. + + Context rows, and ``MIN_REQUESTS``..``MAX_REQUESTS`` generation sequences with beam + width 1. Host-side reads only. + """ + num_contexts = getattr(attn_metadata, "num_contexts", 0) + num_generations = getattr(attn_metadata, "num_generations", 0) + return ( + num_contexts > 0 + and MIN_REQUESTS <= num_generations <= MAX_REQUESTS + and getattr(attn_metadata, "beam_width", 1) == 1 + ) + + +class MixedStepKvSplitsAttention: + """Mixin re-classed in front of a full-attention layer's attention class. + + A call on a step ``takes_generation`` accepts runs the class behind it with the model's + ``DecodeKvSplits`` armed; any other call is exactly the class behind it. Adds one plain + attribute (``_mixed_kv_splits``). + """ + + _mixed_kv_splits: Optional[DecodeKvSplits] = None + + def _attn_impl(self, q, k, v, attn_metadata, *args, **kwargs): + rule = self._mixed_kv_splits + if rule is None or not takes_generation(attn_metadata): + return super()._attn_impl(q, k, v, attn_metadata, *args, **kwargs) + split, kept = rule.split_calls, rule.kept_calls + with rule.armed(): + result = super()._attn_impl(q, k, v, attn_metadata, *args, **kwargs) + if rule.split_calls > split: + outcome = "split" + elif rule.kept_calls > kept: + outcome = "kept flashinfer's split" + else: + outcome = "no trtllm-gen decode call reached (another FMHA library ran)" + logger.info_once( + f"[musespark_open] decode_kv_splits: a mixed step's generation rows " + f"({attn_metadata.num_generations} requests, layer {self.layer_idx}): {outcome}", + key=f"musespark_open.decode_kv_splits.mixed.{outcome}", + ) + return result + + +@functools.lru_cache(maxsize=None) +def _mixed_class(base: type) -> type: + """``base`` with ``MixedStepKvSplitsAttention`` in front, one class per base.""" + return type(f"{base.__name__}MixedKv", (MixedStepKvSplitsAttention, base), {}) + + +def install(model: torch.nn.Module, rule: DecodeKvSplits) -> int: + """Re-class the full-attention layers' attention in front of whatever class they carry, + so their mixed-step attention calls arm ``rule``. Returns the number of layers re-classed. + + Run after every other attention re-class: the mixin goes in front of all of them and + acts only on mixed steps, which none of the others intercept. + """ + chosen = [layer.self_attn for layer in model.layers if not layer.self_attn.is_sliding] + for attention in chosen: + if not isinstance(attention, MixedStepKvSplitsAttention): + attention.__class__ = _mixed_class(type(attention)) + attention._mixed_kv_splits = rule + logger.info( + f"[musespark_open] decode_kv_splits: the generation rows of mixed steps with " + f"{MIN_REQUESTS}..{MAX_REQUESTS} generation requests run the split rule on " + f"{len(chosen)} full-attention layers" + ) + return len(chosen) + + +__all__ = ["MixedStepKvSplitsAttention", "install", "takes_generation"] diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/decode_kv_splits/splits.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/decode_kv_splits/splits.py index c82710716294..306171d084b2 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/decode_kv_splits/splits.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/decode_kv_splits/splits.py @@ -164,6 +164,10 @@ def __init__( self.kv_heads = kv_heads self.q_heads = q_heads self.fallback: Optional[str] = None + #: Decode calls that reached the armed stand-in: how many were split, how many kept + #: flashinfer's own split. + self.split_calls = 0 + self.kept_calls = 0 #: The multi-CTA counters of every call the rule splits (zeroed once; the #: kernel resets what it used at the end of each launch). self.counter = torch.zeros( @@ -250,7 +254,9 @@ def __getattr__(self, name: str): def trtllm_batch_decode_with_kv_cache(self, *args, **kwargs): presented = None if args else self._splits.presented_sm_count(kwargs) if presented is None: + self._splits.kept_calls += 1 return self._real.trtllm_batch_decode_with_kv_cache(*args, **kwargs) + self._splits.split_calls += 1 kwargs["multi_ctas_kv_counter_buffer"] = self._splits.counter decode = self._real device_sm_count = decode.get_device_sm_count diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/mixed_step_gen_rows/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/mixed_step_gen_rows/__init__.py new file mode 100644 index 000000000000..e2541f51f4f7 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/mixed_step_gen_rows/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""mixed_step_gen_rows: the generation rows of mixed steps on multi_request_split_kv (sm_107).""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/mixed_step_gen_rows/attention.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/mixed_step_gen_rows/attention.py new file mode 100644 index 000000000000..baaf9cc37f21 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/mixed_step_gen_rows/attention.py @@ -0,0 +1,257 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The generation rows of mixed steps on ``multi_request_split_kv`` (sm_107 full-attention layers). + +A mixed step (context rows and generation rows in one batch) runs ``PrefillTarget``, the +built-in text-model forward. On sm_107 the public FMHA library list leaves flashinfer's +trtllm-gen FMHA out; ``attention.select_decode_fmha`` puts a ``DecodeTrtllmGenFmha`` in front +of every layer's libraries, which offers the generation phase only. With nothing in front of +the thop fallback offering the context phase, the FMHA manager hands a mixed step's whole +batch to the thop op, so its generation rows run thop's generation FMHA, not the decode-step +paths (``DecodeTrtllmGenFmha`` and, in front of it, ``multi_request_split_kv``). + +This path changes that on the full-attention layers in two parts: + +* ``ThopContextPhaseFmha``, inserted in each full-attention layer's FMHA manager just before + the thop fallback, offers the context phase of a mixed step that ``takes_generation`` + accepts (context rows, ``MIN_GENERATIONS``..``MAX_GENERATIONS`` generation sequences, beam + width 1). The manager then forms ``CombinedFmha(context=ThopContextPhaseFmha, + generation=DecodeTrtllmGenFmha)``: the context rows run the thop op with + ``attention_input_type=context_only`` (the same op, the same context launch as before) and the + generation rows run ``FlashInferTrtllmGenFmha.run_generation`` with ``seq_offset = + num_contexts``, as on sm_103. +* ``MixedStepGenRowsAttention``, re-classed in front of the full-attention layers' attention, + runs such a step's attention call with the decode FMHA call intercepted by a + ``MixedStepCall`` (``attention_forward.intercepting``), so the generation phase launches + ``multi_request_split_kv`` as a decode step's does. ``MixedStepCall`` is + ``multi_request_split_kv.MultiRequestCall`` with the sequence lengths re-aligned: the + generation phase's view of them starts at ``num_contexts``. A call the stand-in declines runs + flashinfer's trtllm-gen decode FMHA. + +The band is read on the host from the step's metadata. Every bound is a point of the FMHA +manager's batch-size cache grid, so a cached selection never crosses it. Sliding-window +layers, pure context steps, decode steps and mixed steps outside the band run what they ran +before. The arming sits at ``_attn_impl``, which the breakable prefill CUDA graph re-runs +eagerly on every replay (``eager_on_graph``), so captured and eager mixed steps both take it. + +Numerics: the context rows are those of the thop op. The generation rows are those of +``multi_request_split_kv`` on decode steps (not bit-identical to thop's generation FMHA; +deterministic). +""" + +from __future__ import annotations + +import dataclasses +import functools +from typing import Optional + +import torch + +from tensorrt_llm._torch.attention.backends.fmha.fallback import FallbackFmha +from tensorrt_llm._torch.attention.backends.fmha.interface import FmhaPhase +from tensorrt_llm._torch.attention.backends.fmha.phased import FmhaParams, PhasedFmha +from tensorrt_llm._torch.attention.backends.interface import AttentionInputType +from tensorrt_llm.logger import logger + +from ...attention import DecodeTrtllmGenFmha +from ...attention_forward import intercepting +from ..multi_request_split_kv import attention as _mr_attention + +#: Generation sequences of a mixed step this path takes. Both bounds and their predecessors +#: are points of the FMHA cache's batch-size grid (``fmha.manager``). +MIN_GENERATIONS = 9 +MAX_GENERATIONS = 64 + + +def takes_generation(attn_metadata) -> bool: + """Whether a step is a mixed step whose generation rows this path takes. + + Context rows, ``MIN_GENERATIONS``..``MAX_GENERATIONS`` generation sequences, beam width 1. + Host-side reads only. + """ + return ( + getattr(attn_metadata, "num_contexts", 0) > 0 + and MIN_GENERATIONS <= getattr(attn_metadata, "num_generations", 0) <= MAX_GENERATIONS + and getattr(attn_metadata, "beam_width", 1) == 1 + ) + + +class ThopContextPhaseFmha(PhasedFmha): + """The thop op's context phase, offered on the mixed steps ``takes_generation`` accepts. + + It never takes a whole batch or a generation phase, so the manager pairs it with the + generation library in front of it. ``run_context`` calls the thop fallback with the + context rows and ``attention_input_type=context_only``. + """ + + supports_workspace_reclamation = True + + def __init__(self, attn, fallback: FallbackFmha): + super().__init__(attn) + self._fallback = fallback + #: Context phases run, read by the served-path log. + self.context_calls = 0 + + def _is_supported(self, q, k, v, metadata, forward_args, *, phase=None) -> bool: + if phase != FmhaPhase.CONTEXT or not takes_generation(metadata): + return False + if forward_args.attention_input_type != AttentionInputType.mixed: + return False + if forward_args.output is None or forward_args.output_sf is not None: + return False + return self._fallback.is_supported(q, k, v, metadata, forward_args) + + def run_context(self, params: FmhaParams) -> None: + rows = params.num_tokens + attention_input = params.qkv_input if params.qkv_input is not None else params.query_input + forward_args = dataclasses.replace( + params.fwd, + output=params.fwd.output[:rows], + attention_input_type=AttentionInputType.context_only, + ) + self._fallback.forward( + attention_input, params.key_input, params.value_input, params.meta, forward_args + ) + self.context_calls += 1 + + def run_generation(self, params: FmhaParams) -> None: + raise RuntimeError("ThopContextPhaseFmha serves the context phase only") + + +def context_library(attn) -> Optional[ThopContextPhaseFmha]: + """The ``ThopContextPhaseFmha`` in ``attn``'s FMHA manager, or None.""" + for fmha in attn._fmha_manager.fmha_libs: + if isinstance(fmha, ThopContextPhaseFmha): + return fmha + return None + + +#: Tensors of the decode call ``multi_request_split_kv`` reads at a 16-byte aligned address. +_ALIGNED_KWARGS = ("query", "out", "block_tables") + + +class MixedStepCall(_mr_attention.MultiRequestCall): + """A mixed step's generation-phase decode call for ``multi_request_split_kv``. + + The generation phase reads the sequence lengths from ``seq_offset = num_contexts`` on, so + ``seq_lens`` is a view that starts 4 bytes per context sequence into the step's lengths. + The kernel reads it at a 16-byte aligned address; ``launch`` hands it an aligned copy of + those ``B`` lengths when the view is not aligned. The other tensors it reads are declined + when not aligned. + """ + + def decline(self, kwargs: dict) -> Optional[str]: + reason = super().decline(kwargs) + if reason is not None: + return reason + for name in _ALIGNED_KWARGS: + if kwargs[name].data_ptr() % 16: + return f"a {name} not 16-byte aligned" + return None + + def launch(self, kwargs: dict) -> None: + seq_lens = kwargs["seq_lens"] + if seq_lens.data_ptr() % 16: + kwargs = dict(kwargs, seq_lens=seq_lens.clone()) + super().launch(kwargs) + + +class MixedStepGenRowsAttention: + """Mixin re-classed in front of a full-attention layer's attention class. + + A call on a step ``takes_generation`` accepts runs the class behind it with the decode + FMHA call intercepted by a ``MixedStepCall`` of the layer's KV heads; any other call is + exactly the class behind it. Adds plain attributes only. + """ + + _mixed_kv_heads: int = 4 + + def _attn_impl(self, q, k, v, attn_metadata, *args, **kwargs): + if not takes_generation(attn_metadata): + return super()._attn_impl(q, k, v, attn_metadata, *args, **kwargs) + call = MixedStepCall(kv_heads=self._mixed_kv_heads, window=0) + with intercepting(call): + result = super()._attn_impl(q, k, v, attn_metadata, *args, **kwargs) + if call.launched: + outcome = "multi_request_split_kv" + elif call.fallback is not None: + outcome = f"flashinfer's trtllm-gen decode FMHA ({call.fallback})" + else: + outcome = "no trtllm-gen decode call reached (the thop op ran the whole batch)" + logger.info_once( + f"[musespark_open] mixed_step_gen_rows: a mixed step's generation rows " + f"({attn_metadata.num_generations} requests, layer {self.layer_idx}): {outcome}", + key=f"musespark_open.mixed_step_gen_rows.{outcome}", + ) + return result + + +@functools.lru_cache(maxsize=None) +def _mixed_class(base: type) -> type: + """``base`` with ``MixedStepGenRowsAttention`` in front, one class per base.""" + return type(f"{base.__name__}MixedGen", (MixedStepGenRowsAttention, base), {}) + + +def _refusal(attention) -> Optional[str]: + libs = attention.attn._fmha_manager.fmha_libs + if not libs or not isinstance(libs[0], DecodeTrtllmGenFmha): + return "no DecodeTrtllmGenFmha in front of its FMHA libraries" + if not any(isinstance(fmha, FallbackFmha) for fmha in libs): + return "no thop fallback among its FMHA libraries" + if not isinstance(attention, _mr_attention.MultiRequestSplitKvAttention): + return "its attention is not on multi_request_split_kv" + return None + + +def install(model: torch.nn.Module) -> int: + """Put a ``ThopContextPhaseFmha`` before the thop fallback of every full-attention layer's + FMHA manager and re-class those layers' attention with ``MixedStepGenRowsAttention`` in + front of whatever class they carry. Returns the number of layers changed. + + Run after ``attention.select_decode_fmha`` and the ``multi_request_split_kv`` install, and + after every other attention re-class: the mixin acts only on mixed steps, which none of the + others intercept. Every layer is checked before any changes; idempotent. + + Raises: + ValueError: if a full-attention layer's FMHA manager does not start with + ``DecodeTrtllmGenFmha``, holds no thop fallback, or its attention is not on + ``multi_request_split_kv``. + """ + chosen = [layer.self_attn for layer in model.layers if not layer.self_attn.is_sliding] + refused = [ + f"layer {attention.layer_idx}: {reason}" + for attention in chosen + if (reason := _refusal(attention)) is not None + ] + if refused: + raise ValueError( + "[musespark_open] mixed_step_gen_rows cannot route mixed steps' generation rows: " + + "; ".join(refused) + ) + for attention in chosen: + attn = attention.attn + if context_library(attn) is None: + libs = attn._fmha_manager.fmha_libs + index = next(i for i, fmha in enumerate(libs) if isinstance(fmha, FallbackFmha)) + libs.insert(index, ThopContextPhaseFmha(attn, libs[index])) + if not isinstance(attention, MixedStepGenRowsAttention): + attention.__class__ = _mixed_class(type(attention)) + attention._mixed_kv_heads = attention._mr_kv_heads + logger.info( + f"[musespark_open] mixed_step_gen_rows: the generation rows of mixed steps with " + f"{MIN_GENERATIONS}..{MAX_GENERATIONS} generation requests run multi_request_split_kv " + f"on {len(chosen)} full-attention layers (context rows on the thop op)" + ) + return len(chosen) + + +__all__ = [ + "MAX_GENERATIONS", + "MIN_GENERATIONS", + "MixedStepCall", + "MixedStepGenRowsAttention", + "ThopContextPhaseFmha", + "context_library", + "install", + "takes_generation", +] diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/multi_request_prep_append/attention.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/multi_request_prep_append/attention.py index a807e54158b0..9683a844e172 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/multi_request_prep_append/attention.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/multi_request_prep_append/attention.py @@ -83,6 +83,8 @@ from ...ops.attention.multi_request_prep_append import ( MAX_ROWS_PER_SEQ, MAX_SEQS, + PACKED_GEOMETRY, + launch_kind, multi_request_prep_append, ) @@ -690,10 +692,17 @@ class each layer carries by then, and acts only on decode steps of ``min_seqs`` attention._mr_append_folded = None attention._mr_append_census = census attention._mr_append_min_seqs = min_seqs + grid = ( + "packed grid (last rows, items per program): " + f"{[(last, items) for last, items, _, _ in PACKED_GEOMETRY]}, one item per program on " + "o_proj-prefetching steps" + if not torch.cuda.is_available() or launch_kind(torch.device("cuda")) == "packed" + else "grid of one program per (row, head)" + ) logger.info( f"[musespark_open] multi_request_prep_append: the K/V append, E4M3 query and FMHA scales " f"of {len(attentions)} layers run in the decode prep on steps of {min_seqs}..{MAX_SEQS} sequences " f"of up to {MAX_ROWS_PER_SEQ} rows each; their trtllm-gen generation preprocess launch " - "(decoder-info + append) is skipped" + f"(decoder-info + append) is skipped; {grid}" ) return len(attentions) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/attention/multi_request_prep_append.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/attention/multi_request_prep_append.py index eb503cac93b3..e9fd45808a1f 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/attention/multi_request_prep_append.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/attention/multi_request_prep_append.py @@ -51,13 +51,24 @@ ``prefetch`` (optional): a tensor whose bytes the kernel asks the L2 to fetch -- the layer's attention output projection weight, read after the FMHA. Every program issues one ``cp.async.bulk.prefetch.L2`` of an equal slice right after -its wait, so the requests leave every SM's TMA unit within the prep, before the -FMHA's own TMA loads of K/V start. It reads nothing back and writes nothing: the -values computed are the same with and without it. +its wait, and a prefetching call runs one program per (row, head) item -- the +built-in grid's program count, so the requests and their slice sizes are the +built-in grid's -- and they leave every SM's TMA unit within the prep, before the +FMHA's own TMA loads of K/V start. It reads nothing back and writes nothing: the values computed are the same +with and without it. + +Launch geometry, per SM version (``LAUNCH_KIND_BY_SM``). On sm103 one program +takes several consecutive (row, head) items of the row-major ``rows x (H + G)`` +grid (``PACKED_GEOMETRY``), each item running the per-head code of a +one-program-per-(row, head) grid element for element (the same loads, the same +additions in the same order, the same rounding points and stores), so the bytes +written do not depend on the geometry. On sm107 the grid is one program per +(row, head), one warp each (``_multi_request_prep_append_grid_kernel``). """ from __future__ import annotations +import functools from typing import Optional, Tuple import torch @@ -74,7 +85,9 @@ _load8, _normalize, _rotate, + _rsqrt_approx_ftz, is_rope_table, + libdevice, preattn_prep, ) @@ -83,16 +96,340 @@ #: Query rows per sequence the entry takes: one accepted token and up to 7 drafts. MAX_ROWS_PER_SEQ = 8 #: Sequences one call takes: 128 verify blocks of 8 rows (``preattn_prep``'s 1024-row cap), the -#: largest decode step of a generation worker of ``max_batch_size`` 128. The grid is one program -#: per (row, head), so a call's rows do not depend on how many sequences it carries. +#: largest decode step of a generation worker of ``max_batch_size`` 128. The grid covers the +#: (row, head) items, so a call's rows do not depend on how many sequences it carries. MAX_SEQS = 128 #: Rows one call takes. MAX_ROWS = MAX_SEQS * MAX_ROWS_PER_SEQ +#: The launch geometry by row count: ``(last row count, (row, head) items per program, warps per +#: program, vector lanes)``, the first entry whose row count reaches the call's. An item is one +#: head of one row; a program's items are consecutive in row-major (row, head) order. With vector +#: lanes a lane loads and stores its 8 consecutive values as one vector (the same bytes). The +#: entries above 32 rows are the fastest geometries of a GB300 sweep at 16..1024 rows; up to 32 +#: rows one item per program (the built-in grid) is faster in a served decode step on GB300. +PACKED_GEOMETRY = ( + (32, 1, 1, False), + (128, 2, 1, False), + (256, 4, 1, True), + (768, 8, 2, True), + (MAX_ROWS, 4, 1, True), +) + + +def packed_geometry(rows: int, heads: int, programs_at_least: int = 0) -> Tuple[int, int, bool]: + """``(items per program, warps, vector lanes)`` of a call of ``rows`` rows of ``heads`` heads. + + ``programs_at_least`` lowers the items per program (never below one) until the grid has at + least that many programs -- the prefetching call's one bulk request per program then leaves + from at least that many programs. + """ + items, warps, vec = next((i, w, v) for last, i, w, v in PACKED_GEOMETRY if rows <= last) + while items > 1 and -(-rows * heads // items) < programs_at_least: + items //= 2 + warps = max(1, min(warps, items * 8 // 32)) + return items, warps, vec + + +@triton.jit +def _fp8_store8(ptr, y0, y1, y2, y3, y4, y5, y6, y7, scale_bf16, mask): + """Eight BF16 values one norm thread owns, as E4M3 at ``ptr + 0..7`` where ``mask``.""" + tl.store(ptr + 0, _to_cache_fp8(y0, scale_bf16), mask=mask) + tl.store(ptr + 1, _to_cache_fp8(y1, scale_bf16), mask=mask) + tl.store(ptr + 2, _to_cache_fp8(y2, scale_bf16), mask=mask) + tl.store(ptr + 3, _to_cache_fp8(y3, scale_bf16), mask=mask) + tl.store(ptr + 4, _to_cache_fp8(y4, scale_bf16), mask=mask) + tl.store(ptr + 5, _to_cache_fp8(y5, scale_bf16), mask=mask) + tl.store(ptr + 6, _to_cache_fp8(y6, scale_bf16), mask=mask) + tl.store(ptr + 7, _to_cache_fp8(y7, scale_bf16), mask=mask) + + +@triton.jit +def _split8(v): + """The 8 elements of ``v``'s last axis (shape ``[ITEMS, 2, 2, 2, 8]``), in order, each + ``[ITEMS, 2, 2, 2]``.""" + v = tl.reshape(v, (v.shape[0], 2, 2, 2, 2, 2, 2)) + # Element k = 4 a + 2 b + c of the last three axes (a, b, c): split c, then b, then a. + c0, c1 = tl.split(v) + c0b0, c0b1 = tl.split(c0) + c1b0, c1b1 = tl.split(c1) + x0, x4 = tl.split(c0b0) + x2, x6 = tl.split(c0b1) + x1, x5 = tl.split(c1b0) + x3, x7 = tl.split(c1b1) + return x0, x1, x2, x3, x4, x5, x6, x7 + + +@triton.jit +def _join8(x0, x1, x2, x3, x4, x5, x6, x7): + """``_split8``'s inverse: ``[ITEMS, 2, 2, 2, 8]`` with element k = ``xk``.""" + c0 = tl.join(tl.join(x0, x4), tl.join(x2, x6)) + c1 = tl.join(tl.join(x1, x5), tl.join(x3, x7)) + v = tl.join(c0, c1) + return tl.reshape(v, (v.shape[0], 2, 2, 2, 8)) + + +@triton.jit +def _inverse_rms_per_item(x0, x1, x2, x3, x4, x5, x6, x7, eps): + """``_inverse_rms`` of each item of ``[ITEMS, 2, 2, 2]`` lane-shaped values, as ``[ITEMS]``. + + The same additions in the same order: in order within a lane (``add.rn``), then lane + bit 0, 1, 2, each a single commutative add of a size-2 axis. + """ + s = libdevice.add_rn(x0 * x0, x1 * x1) + s = libdevice.add_rn(s, x2 * x2) + s = libdevice.add_rn(s, x3 * x3) + s = libdevice.add_rn(s, x4 * x4) + s = libdevice.add_rn(s, x5 * x5) + s = libdevice.add_rn(s, x6 * x6) + s = libdevice.add_rn(s, x7 * x7) + total = tl.sum(tl.sum(tl.sum(s, axis=3), axis=2), axis=1) + return _rsqrt_approx_ftz(total / 64.0 + eps) + + +@triton.jit +def _multi_request_prep_append_kernel( + qkv_ptr, + qkv_out_ptr, + gate_out_ptr, + cos_sin_ptr, + positions_ptr, + q8_ptr, + pool_ptr, + bt_ptr, + bt_stride_seq, + bt_stride_kv, + seq_len_ptr, + orig_quant_ptr, + quant_orig_ptr, + bmm1_ptr, + bmm2_ptr, + cu_seqlens_ptr, + bmm1_host, + o_scale, + qkv_stride, + qkv_out_stride, + gate_out_stride, + eps, + q_scale, + num_seqs, + num_items, + pf_base, + pf_bytes, + pf_chunk, + NUM_Q_HEADS: tl.constexpr, + NUM_KV_HEADS: tl.constexpr, + QKV_PERMUTED: tl.constexpr, + APPLY_ROPE: tl.constexpr, + ROWS_PER_SEQ: tl.constexpr, + P: tl.constexpr, + LOG2_P: tl.constexpr, + SEQS_BLOCK: tl.constexpr, + LAUNCH_WITH_PDL: tl.constexpr, + ITEMS: tl.constexpr = 1, + LANES: tl.constexpr = 32, + VEC: tl.constexpr = False, + PREFETCH: tl.constexpr = False, +): + """Program ``p`` takes the ``ITEMS`` items ``p * ITEMS ..`` of the ``rows * (H + G)`` + (row, head) items in row-major order; every item runs the one-head code of the + one-program-per-(row, head) grid it replaces, element for element.""" + D: tl.constexpr = 64 + HALF: tl.constexpr = 32 + HEADS: tl.constexpr = NUM_Q_HEADS + NUM_KV_HEADS + PAGE: tl.constexpr = NUM_KV_HEADS * P * D + if LAUNCH_WITH_PDL: + # Before any global read: the projection is the previous kernel's output. + # No early trigger: a dependent (the FMHA) launches when this grid ends. + tl.extra.cuda.gdc_wait() + program = tl.program_id(0) + if PREFETCH: + # One bulk L2 prefetch of this program's slice of ``pf_bytes`` (element 0 issues it; + # ``LANES`` is the program's thread count, so no element is held twice). The asm + # yields 0, so the store never happens; it keeps the request live. + lane = tl.arange(0, LANES) + start = program.to(tl.int64) * pf_chunk + size = tl.minimum(pf_bytes - start, pf_chunk) + size = tl.where((lane == 0) & (size > 0), size, 0).to(tl.int32) + issued = _bulk_prefetch_l2(tl.full([LANES], 0, tl.int64) + pf_base + start, size) + tl.store(cu_seqlens_ptr + lane, issued, mask=issued != 0) + + # The program's items; a tail item past the last one computes item 0's addresses + # (every load stays in bounds) and stores nothing. + item = program * ITEMS + tl.arange(0, ITEMS) + valid = item < num_items + item = tl.where(valid, item, 0) + row = item // HEADS + head = item % HEADS + is_q = head < NUM_Q_HEADS + is_kv = head >= NUM_Q_HEADS + kv_head = head - NUM_Q_HEADS + row64 = row.to(tl.int64) + scale_bf16 = tl.load(orig_quant_ptr).to(tl.bfloat16) + + # The row's cache position: the sequence's last ROWS_PER_SEQ positions + # (seq_lens counts the step's rows), addressed as KVBlockArray does. + seq = row // ROWS_PER_SEQ + position = tl.load(seq_len_ptr + seq) - ROWS_PER_SEQ + row % ROWS_PER_SEQ + table = bt_ptr + seq.to(tl.int64) * bt_stride_seq + (position >> LOG2_P) + slot = (kv_head * P + (position & (P - 1))) * D + + # The pass-through half first, so its load is in flight during the norm: + # a query item copies its gate head, a key item its value head (and stores it + # into the cache's V page). + d = tl.arange(0, D)[None, :] + if QKV_PERMUTED: + gate_in = (NUM_Q_HEADS + 2 * NUM_KV_HEADS) * D + head * D + v_in = NUM_Q_HEADS * D + NUM_KV_HEADS * D + kv_head * D + else: + gate_in = head * 2 * D + D + v_in = 2 * NUM_Q_HEADS * D + NUM_KV_HEADS * D + kv_head * D + pass_in = tl.where(is_q, gate_in, v_in) + passed = tl.load(qkv_ptr + (row64 * qkv_stride + pass_in)[:, None] + d) + q_store = (valid & is_q)[:, None] + kv_store = (valid & is_kv)[:, None] + tl.store(gate_out_ptr + (row64 * gate_out_stride + head * D)[:, None] + d, passed, mask=q_store) + tl.store( + qkv_out_ptr + + (row64 * qkv_out_stride + (NUM_Q_HEADS + NUM_KV_HEADS) * D + kv_head * D)[:, None] + + d, + passed, + mask=kv_store, + ) + page_v = tl.load(table + bt_stride_kv, mask=valid & is_kv, other=0).to(tl.int64) + tl.store( + pool_ptr + (page_v * PAGE + slot)[:, None] + d, + _to_cache_fp8(passed, scale_bf16), + mask=kv_store, + ) + + if QKV_PERMUTED: + in_off = tl.where(is_q, head * D, NUM_Q_HEADS * D + kv_head * D) + else: + in_off = tl.where(is_q, head * 2 * D, 2 * NUM_Q_HEADS * D + kv_head * D) + out_off = tl.where(is_q, head * D, NUM_Q_HEADS * D + kv_head * D) + scale = tl.where(is_q, q_scale, 1.0)[:, None, None, None] + + lane8 = 8 * ( + 4 * tl.arange(0, 2)[None, :, None, None] + + 2 * tl.arange(0, 2)[None, None, :, None] + + tl.arange(0, 2)[None, None, None, :] + ) + in_lane = qkv_ptr + (row64 * qkv_stride + in_off)[:, None, None, None] + lane8 + if VEC: + # One 16-byte load per lane: the same 8 values. + x0, x1, x2, x3, x4, x5, x6, x7 = _split8( + tl.load(in_lane[:, :, :, :, None] + tl.arange(0, 8)[None, None, None, None, :]).to( + tl.float32 + ) + ) + else: + x0, x1, x2, x3, x4, x5, x6, x7 = _load8(in_lane) + inverse_rms = _inverse_rms_per_item(x0, x1, x2, x3, x4, x5, x6, x7, eps)[:, None, None, None] + y0 = _normalize(x0, inverse_rms, scale) + y1 = _normalize(x1, inverse_rms, scale) + y2 = _normalize(x2, inverse_rms, scale) + y3 = _normalize(x3, inverse_rms, scale) + y4 = _normalize(x4, inverse_rms, scale) + y5 = _normalize(x5, inverse_rms, scale) + y6 = _normalize(x6, inverse_rms, scale) + y7 = _normalize(x7, inverse_rms, scale) + + if APPLY_ROPE: + # Pair p = 4 * lane + i rotates elements 2p, 2p + 1 by table row + # [cos(32) | sin(32)] of this token's RoPE position. + rope_position = tl.load(positions_ptr + row).to(tl.int64) + rope = cos_sin_ptr + (rope_position * D)[:, None, None, None] + lane8 // 2 + y0, y1 = _rotate(y0, y1, tl.load(rope + 0), tl.load(rope + HALF + 0)) + y2, y3 = _rotate(y2, y3, tl.load(rope + 1), tl.load(rope + HALF + 1)) + y4, y5 = _rotate(y4, y5, tl.load(rope + 2), tl.load(rope + HALF + 2)) + y6, y7 = _rotate(y6, y7, tl.load(rope + 3), tl.load(rope + HALF + 3)) + + # The BF16 operand values: what the prep stores and what the append converts. + b0 = y0.to(tl.bfloat16) + b1 = y1.to(tl.bfloat16) + b2 = y2.to(tl.bfloat16) + b3 = y3.to(tl.bfloat16) + b4 = y4.to(tl.bfloat16) + b5 = y5.to(tl.bfloat16) + b6 = y6.to(tl.bfloat16) + b7 = y7.to(tl.bfloat16) + store = valid[:, None, None, None] + out = qkv_out_ptr + (row64 * qkv_out_stride + out_off)[:, None, None, None] + lane8 + q8 = q8_ptr + ((row64 * NUM_Q_HEADS + head) * D)[:, None, None, None] + lane8 + page_k = tl.load(table, mask=valid & is_kv, other=0).to(tl.int64) + k8 = pool_ptr + (page_k * PAGE + slot)[:, None, None, None] + lane8 + if VEC: + # One vector store per lane each: the same bytes. + e = tl.arange(0, 8)[None, None, None, None, :] + b = _join8(b0, b1, b2, b3, b4, b5, b6, b7) + tl.store(out[:, :, :, :, None] + e, b, mask=store[:, :, :, :, None]) + b8 = _to_cache_fp8(b, scale_bf16) + tl.store(q8[:, :, :, :, None] + e, b8, mask=(valid & is_q)[:, None, None, None, None]) + tl.store(k8[:, :, :, :, None] + e, b8, mask=(valid & is_kv)[:, None, None, None, None]) + else: + tl.store(out + 0, b0, mask=store) + tl.store(out + 1, b1, mask=store) + tl.store(out + 2, b2, mask=store) + tl.store(out + 3, b3, mask=store) + tl.store(out + 4, b4, mask=store) + tl.store(out + 5, b5, mask=store) + tl.store(out + 6, b6, mask=store) + tl.store(out + 7, b7, mask=store) + # Their E4M3 copies: Q into the FMHA's query buffer, K into the cache's K page. + _fp8_store8( + q8, b0, b1, b2, b3, b4, b5, b6, b7, scale_bf16, (valid & is_q)[:, None, None, None] + ) + _fp8_store8( + k8, b0, b1, b2, b3, b4, b5, b6, b7, scale_bf16, (valid & is_kv)[:, None, None, None] + ) + + if program == 0: + # The FMHA's scale words, in the built-in kernel's operation order: + # (s_q * s_k) * bmm1_host, then that * log2(e); o_scale * s_v. + s = tl.load(quant_orig_ptr) + bmm1 = s * s * bmm1_host + tl.store(bmm1_ptr, bmm1) + tl.store(bmm1_ptr + 1, bmm1 * 1.4426950408889634) + tl.store(bmm2_ptr, o_scale * s) + if ROWS_PER_SEQ > 1: + # The decoder-info kernel's cumulative query lengths for fixed query lengths + # (it runs only for more than one query row per sequence). + i = tl.arange(0, SEQS_BLOCK) + tl.store(cu_seqlens_ptr + i, i * ROWS_PER_SEQ, mask=i <= num_seqs) + + +#: The kernel ``multi_request_prep_append`` launches per SM version (``launch_kind``): on sm103 +#: the packed kernel above on ``PACKED_GEOMETRY``; on sm107 the grid of one program per +#: (row, head), one warp each (``_multi_request_prep_append_grid_kernel``). Both compute every +#: item with the same per-head code, so the bytes written are the same. +LAUNCH_KIND_BY_SM = {103: "packed", 107: "grid"} + + +@functools.lru_cache(maxsize=None) +def _sm_of(index: int) -> int: + major, minor = torch.cuda.get_device_capability(index) + return 10 * major + minor + + +def launch_kind(device: torch.device) -> str: + """``LAUNCH_KIND_BY_SM`` for ``device``. + + Raises: + ValueError: on a device ``LAUNCH_KIND_BY_SM`` does not list. + """ + sm = _sm_of(device.index if device.index is not None else torch.cuda.current_device()) + kind = LAUNCH_KIND_BY_SM.get(sm) + if kind is None: + raise ValueError( + f"multi_request_prep_append: no launch for sm{sm} " + f"(supported: {sorted(LAUNCH_KIND_BY_SM)})" + ) + return kind @triton.jit -def _fp8_store8(ptr, y0, y1, y2, y3, y4, y5, y6, y7, scale_bf16): - """Eight BF16 values one norm thread owns, as E4M3 at ``ptr + 0..7``.""" +def _fp8_store8_all(ptr, y0, y1, y2, y3, y4, y5, y6, y7, scale_bf16): + """Eight BF16 values one norm thread owns, as E4M3 at ``ptr + 0..7`` (no mask).""" tl.store(ptr + 0, _to_cache_fp8(y0, scale_bf16)) tl.store(ptr + 1, _to_cache_fp8(y1, scale_bf16)) tl.store(ptr + 2, _to_cache_fp8(y2, scale_bf16)) @@ -104,7 +441,7 @@ def _fp8_store8(ptr, y0, y1, y2, y3, y4, y5, y6, y7, scale_bf16): @triton.jit -def _multi_request_prep_append_kernel( +def _multi_request_prep_append_grid_kernel( qkv_ptr, qkv_out_ptr, gate_out_ptr, @@ -252,11 +589,11 @@ def _multi_request_prep_append_kernel( # Their E4M3 copies: Q into the FMHA's query buffer, K into the cache's K page. if is_q: q8 = q8_ptr + (row64 * NUM_Q_HEADS + head) * D + lane * 8 - _fp8_store8(q8, b0, b1, b2, b3, b4, b5, b6, b7, scale_bf16) + _fp8_store8_all(q8, b0, b1, b2, b3, b4, b5, b6, b7, scale_bf16) else: page_k = tl.load(table).to(tl.int64) k8 = pool_ptr + page_k * PAGE + slot + lane * 8 - _fp8_store8(k8, b0, b1, b2, b3, b4, b5, b6, b7, scale_bf16) + _fp8_store8_all(k8, b0, b1, b2, b3, b4, b5, b6, b7, scale_bf16) if (row == 0) & (head == 0): # The FMHA's scale words, in the built-in kernel's operation order: @@ -321,7 +658,10 @@ class MultiRequestPrepAppend(OpWrapper): every target layer (46 RoPE + 16 NoPE) on decode steps of 2..128 sequences of the same number of rows (fast path multi_request_prep_append), whose trtllm-gen generation preprocess launch (decoder-info + append kernels) - the target then skips. Grid (rows, H + G), one warp per program. + the target then skips. Grid on sm103: the rows * (H + G) (row, head) items in + row-major order, PACKED_GEOMETRY's items, warps and lanes per program by row count + (with a prefetch, one program per item); on sm107: (rows, H + G), one warp per + program. Bit-identical, gated by its own tests: qkv_out and gate_out against preattn_prep; the cache bytes, the E4M3 query bytes, both scale words and @@ -428,7 +768,11 @@ def __call__( ) elif position_ids is not None: raise ValueError("multi_request_prep_append: position_ids without a cos/sin table") + heads = num_q_heads + num_kv_heads + num_items = rows * heads + kind = launch_kind(qkv.device) pf_base, pf_bytes, pf_chunk = 0, 0, 16 + programs_at_least = 0 if prefetch is not None: pf_bytes = prefetch.numel() * prefetch.element_size() if not ( @@ -442,14 +786,66 @@ def __call__( f"of a multiple of 16 bytes, got {tuple(prefetch.shape)} {prefetch.dtype} at " f"{prefetch.data_ptr():#x}" ) - programs = rows * (num_q_heads + num_kv_heads) + # One bulk request per program, one program per (row, head) item: the o_proj + # weight is requested in the one-program-per-(row, head) grid's slices, from + # every SM, whatever the table's packing at these rows. + programs_at_least = num_items + if kind == "grid": + items, warps, vec = 1, 1, False + else: + items, warps, vec = packed_geometry(rows, heads, programs_at_least) + programs = -(-num_items // items) + if prefetch is not None: pf_base = prefetch.data_ptr() pf_chunk = ((pf_bytes + programs - 1) // programs + 15) // 16 * 16 qkv_out = torch.empty( (rows, (num_q_heads + 2 * num_kv_heads) * HEAD_DIM), dtype=qkv.dtype, device=qkv.device ) gate_out = torch.empty((rows, num_q_heads * HEAD_DIM), dtype=qkv.dtype, device=qkv.device) - _multi_request_prep_append_kernel[(rows, num_q_heads + num_kv_heads)]( + if kind == "grid": + _multi_request_prep_append_grid_kernel[(rows, heads)]( + qkv, + qkv_out, + gate_out, + cos_sin, + positions, + q_fp8.view(torch.uint8).view(torch.float8_e4m3fn), + kv_pool, + block_tables, + block_tables.stride(0), + block_tables.stride(1), + seq_lens, + kv_scale_orig_quant, + kv_scale_quant_orig, + bmm1_scale, + bmm2_scale, + cu_seqlens, + bmm1_host_scale, + o_scale, + qkv.stride(0), + qkv_out.stride(0), + gate_out.stride(0), + eps, + q_scale, + num_seqs, + pf_base, + pf_bytes, + pf_chunk, + NUM_Q_HEADS=num_q_heads, + NUM_KV_HEADS=num_kv_heads, + QKV_PERMUTED=qkv_permuted, + APPLY_ROPE=cos_sin is not None, + ROWS_PER_SEQ=rows_per_seq, + P=page, + LOG2_P=page.bit_length() - 1, + SEQS_BLOCK=triton.next_power_of_2(num_seqs + 1), + LAUNCH_WITH_PDL=enable_pdl, + PREFETCH=prefetch is not None, + num_warps=1, + launch_pdl=enable_pdl, + ) + return qkv_out, gate_out + _multi_request_prep_append_kernel[(programs,)]( qkv, qkv_out, gate_out, @@ -474,6 +870,7 @@ def __call__( eps, q_scale, num_seqs, + num_items, pf_base, pf_bytes, pf_chunk, @@ -486,8 +883,11 @@ def __call__( LOG2_P=page.bit_length() - 1, SEQS_BLOCK=triton.next_power_of_2(num_seqs + 1), LAUNCH_WITH_PDL=enable_pdl, + ITEMS=items, + LANES=32 * warps, + VEC=vec, PREFETCH=prefetch is not None, - num_warps=1, + num_warps=warps, launch_pdl=enable_pdl, ) return qkv_out, gate_out diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/attention/multi_request_split_kv.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/attention/multi_request_split_kv.py index a18b928bb409..1f514a082e42 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/attention/multi_request_split_kv.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/attention/multi_request_split_kv.py @@ -93,6 +93,30 @@ #: The build the target launches: O rescale in 16-column chunks (spill-free), two Q terms. RESCALE_COLS = 16 NQ = 2 +#: Pages (8 KB of K and V each) a CTA's ring holds in flight, by the number of E4M3 terms Q enters +#: as, per SM version (``ring_stages``). At deep contexts the split kernel streams the cache and +#: the ring depth sets its rate. sm103: each build takes the deepest ring one CTA's shared memory +#: holds next to its Q buffers, 20 pages with a one-term Q (the served E4M3 Q) and 16 with a +#: two-term Q (24 and 20 do not launch). sm107: 16 pages for either Q. +RING_STAGES_BY_SM = {103: {1: 20, 2: 16}, 107: {1: 16, 2: 16}} +#: sm103's ring depths. +RING_STAGES = RING_STAGES_BY_SM[103] + + +def ring_stages(device_index: int, q_terms: int) -> int: + """The ring depth of a build whose Q enters as ``q_terms`` E4M3 terms on the device. + + Raises: + ValueError: on a device ``RING_STAGES_BY_SM`` does not list. + """ + major, minor = torch.cuda.get_device_capability(device_index) + table = RING_STAGES_BY_SM.get(10 * major + minor) + if table is None: + raise ValueError( + f"multi_request_split_kv: no ring depth for sm{10 * major + minor} " + f"(supported: {sorted(RING_STAGES_BY_SM)})" + ) + return table[q_terms] def default_num_work(kv_heads: int, sm_count: int) -> int: @@ -163,10 +187,11 @@ def _compiled_kernel( rescale_cols: int = RESCALE_COLS, q_fp8: bool = False, device_scale: bool = False, - stages: int = 16, + stages: Optional[int] = None, max_requests: int = REQUEST_BUILDS[0], ): - """The compiled split kernel of one build.""" + """The compiled split kernel of one build (its ring depth from ``ring_stages`` unless + ``stages`` is given).""" import cuda.bindings.driver as cuda import cutlass import cutlass.cute as cute @@ -175,6 +200,8 @@ def _compiled_kernel( from . import multi_request_split_kv_kernel as mk assert (mk.MAX_ROWS, mk.REQUEST_BUILDS) == (MAX_ROWS, REQUEST_BUILDS) + if stages is None: + stages = ring_stages(device_index, 1 if q_fp8 else nq) kernel = mk.MultiRequestSplitKvKernel( enable_pdl=enable_pdl, kv_heads=kv_heads, diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/attention/multi_request_split_kv_kernel.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/attention/multi_request_split_kv_kernel.py index 7ebefe6e9481..a96e07990486 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/attention/multi_request_split_kv_kernel.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/attention/multi_request_split_kv_kernel.py @@ -9,7 +9,7 @@ The per-tile machinery is ``full_attention_split_kv_kernel``'s (the one-sequence kernel): a CTA serves one pair of KV heads (``blockIdx.y``) with 384 threads -- two softmax warpgroups that take alternating 128-key tiles, one MMA warp, one TMA loader warp, two other warps -- over a -16-stage ring of 32-key pages (K in two 2 KB boxes, one per head, V in one 4 KB box). Q' is +ring of stages 32-key pages (K in two 2 KB boxes, one per head, V in one 4 KB box). Q' is block-diagonal over the pair's two heads (row = head * 64 + token * 8 + q-head-in-group), so one M=128 MMA per Q term and K block scores a 128-key tile; P goes to TMEM in E4M3 in two halves; each warpgroup keeps its own running max / sum and TMEM O accumulator. From c079b478d325386d794066ed9506f433c1b53a3b Mon Sep 17 00:00:00 2001 From: Xiao Wang <24860335+xwang233@users.noreply.github.com> Date: Mon, 26 Oct 2026 18:44:06 +0000 Subject: [PATCH 07/16] [None][feat] modeling_v2 musespark_open tep: fixed decode expert tactics and Rubin CuTe DSL prefill experts - decode_expert_tactics: the stacked NVFP4 expert call of tep4 decode steps runs a fixed TRTLLM-gen [tile_n, config] tactic per (expert window, rows) from a per-GPU table instead of the autotuner's power-of-two bucket choice. GB300 (sm103): the 128-token tile for 256..504 rows; every pinned tactic computes the expert rows bitwise as the autotuned ones. Rubin (sm107): the 32-token tile for 256..384 rows, bitwise as the autotuned tactics with the same FC2 MMA K. Unlisted row counts and windows, tuning passes and configs not valid for the call's shape keep the autotuner; a GPU without a table is refused. - prefill_cutedsl_experts (Rubin, sm107 at tp 4 only, selected by serves(compute_capability, tp_size)): the unfinalized expert calls of 4096..8192 NVFP4 rows (prefill and mixed steps) of every NVFP4 MoE layer on the rank's contiguous expert block run moe_sort, the Rubin CuTe DSL FC1 (TRTLLM-gen gated layout) and FC2 (permuted store, TRTLLM-gen N shuffle, 64-deep MMA) on the served TRTLLM-gen tensors, and return the TRTLLM-gen call's triple to the model's per-expert post-norm finalize. Kernels are compiled at install with fixed tactics per row bucket; the install is agreed over the TP group. Rounding differs from TRTLLM-gen; expert selection is unchanged. Every other GPU, tp size and call keeps the TRTLLM-gen expert call. Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com> --- .../prefill_cutedsl_experts/__init__.py | 3 + .../prefill_cutedsl_experts/install.py | 142 +++++++++++ .../kernels/replicated_experts/ring_buddy.py | 45 ++++ .../ops/moe/decode_expert_tactics.py | 146 ++++++++++++ .../ops/moe/prefill_cutedsl_experts.py | 220 ++++++++++++++++++ 5 files changed, 556 insertions(+) create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/prefill_cutedsl_experts/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/prefill_cutedsl_experts/install.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/moe/decode_expert_tactics.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/moe/prefill_cutedsl_experts.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/prefill_cutedsl_experts/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/prefill_cutedsl_experts/__init__.py new file mode 100644 index 000000000000..b1ea75c49827 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/prefill_cutedsl_experts/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""prefill_cutedsl_experts: large prefill expert calls on the Rubin CuTe DSL grouped GEMMs.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/prefill_cutedsl_experts/install.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/prefill_cutedsl_experts/install.py new file mode 100644 index 000000000000..2e258d81509c --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/prefill_cutedsl_experts/install.py @@ -0,0 +1,142 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""``prefill_cutedsl_experts``: the NVFP4 layers' large expert calls on the Rubin CuTe DSL GEMMs (sm107). + +Installed on Rubin (sm107) at tp 4 only (``serves``); every other GPU and tp size keeps +the TRTLLM-gen expert calls. Every NVFP4 MoE layer whose leaf is +the native TRTLLM-gen one on this rank's contiguous expert block +(``ring_buddy.eligible``) gets a wrapper around its expert forward: an +unfinalized call (``do_finalize=False``) of ``MIN_ROWS..MAX_ROWS`` NVFP4 rows -- +the prefill and mixed steps of at least 4096 tokens -- runs +``ops/moe/prefill_cutedsl_experts`` on the leaf's own tensors and returns the +same ``(gemm2_output, weights, expanded_idx_to_permuted_idx)`` triple the +TRTLLM-gen call returns, which the model's per-expert post-norm finalize then +consumes. Every other call (decode steps, smaller prefill chunks, the BF16 layers +0 and 61) keeps the expert forward it had. + +The routing is the routing method's own ``apply`` (memoized on the logits, so the +finalize reads the same weights). Rounding-different from the TRTLLM-gen call; +expert selection is unchanged. + +Install after the weights are final and before any CUDA graph capture: the +kernels are compiled here (``PrefillCuteDslExperts.prepare``), so no capture +compiles them. Agreed over the TP group (``agreed_install``). +""" + +from __future__ import annotations + +import types +from typing import List, Optional, Tuple + +import torch + +from tensorrt_llm._torch.utils import Fp4QuantizedTensor +from tensorrt_llm.logger import logger + +from ... import dispatch as _tp_agreement +from ...ops.moe import prefill_cutedsl_experts as _op +from ...ops.moe.ring_buddy_assign import EXPERTS, TOP_K +from ..replicated_experts import ring_buddy as _ring_buddy + +#: The SM version and tp size this path is installed at: Rubin (sm107) tep4. +SERVED_SM = 107 +SERVED_TP_SIZE = 4 + +#: Rows of the expert calls this path serves. +MIN_ROWS = 4096 +MAX_ROWS = 8192 + + +def serves(compute_capability: Optional[Tuple[int, int]], tp_size: int) -> bool: + """Whether a model on ``compute_capability`` at ``tp_size`` installs this path.""" + return _tp_agreement.sm_of(compute_capability) == SERVED_SM and tp_size == SERVED_TP_SIZE + + +def _refusal(moe, mapping) -> Optional[str]: + reason = _ring_buddy.eligible(moe, mapping) + if reason is not None: + return reason + if _ring_buddy.expert_format(moe) != "nvfp4": + return "the experts are not the native NVFP4 TRTLLM-gen leaf" + backend = _ring_buddy._backend(moe) + if not _ring_buddy._is_plain_swiglu(backend): + return "the experts are not plain SwiGLU" + if getattr(backend, "fc2_input_scale", None) is None: + return "the leaf has no FC2 input scale" + return None + + +def _wrap(moe, op: "_op.PrefillCuteDslExperts") -> None: + experts = moe.experts + prior = experts.forward + routing = experts.routing_method + + def forward(self, x, router_logits, *args, do_finalize=True, **kwargs): + if ( + do_finalize + or args + or not isinstance(x, Fp4QuantizedTensor) + or x.is_sf_swizzled + or not MIN_ROWS <= x.fp4_tensor.shape[0] <= MAX_ROWS + or torch.compiler.is_compiling() + ): + return prior(x, router_logits, *args, do_finalize=do_finalize, **kwargs) + selected, weights = routing.apply(router_logits) + gemm2_output, expanded_idx_to_permuted_idx = op( + x.fp4_tensor, x.scaling_factor, selected, weights + ) + return gemm2_output, weights, expanded_idx_to_permuted_idx + + experts.forward = types.MethodType(forward, experts) + + +def install(text_model, mapping) -> List[int]: + """Wrap every eligible NVFP4 layer's expert forward; returns the layer indices. + + Collective over the TP group (the outcome is agreed). Raises ValueError on every + rank alike when some rank refuses (not sm107, no eligible layer) or cannot + compile the kernels. + """ + + def build() -> List[int]: + if torch.cuda.get_device_capability() != (10, 7): + raise ValueError( + "[musespark_open] prefill_cutedsl_experts is written for sm107 (Rubin) only" + ) + installed, skipped, ops = [], {}, [] + for index, layer in enumerate(text_model.layers): + moe = layer.feed_forward + reason = _refusal(moe, mapping) + if reason is not None: + skipped[index] = reason + continue + backend = _ring_buddy._backend(moe) + op = _op.PrefillCuteDslExperts(backend, EXPERTS, TOP_K, backend.slot_start) + ops.append((moe, op)) + installed.append(index) + if not installed: + raise ValueError( + "[musespark_open] prefill_cutedsl_experts found no eligible NVFP4 layer" + ) + try: + # One compile serves every layer (the kernels' shapes are runtime arguments). + ops[0][1].prepare(MAX_ROWS) + torch.cuda.synchronize() + # Hand the compile call's transient buffers back before the KV cache is sized. + torch.cuda.empty_cache() + except Exception as error: # noqa: BLE001 - the compile's failure must reach the agreement + raise RuntimeError( + f"[musespark_open] prefill_cutedsl_experts could not compile: {error!r}" + ) from error + for moe, op in ops: + _wrap(moe, op) + moe._prefill_cutedsl_experts = op + logger.info( + f"[musespark_open] prefill_cutedsl_experts: {len(installed)} NVFP4 MoE layers run " + f"their {MIN_ROWS}..{MAX_ROWS}-row expert calls on the Rubin CuTe DSL GEMMs " + f"(tile {_op.TILE}); built-in expert call kept for layers {sorted(skipped)}: " + f"{sorted(set(skipped.values()))}" + ) + return installed + + return _tp_agreement.agreed_install(mapping, build, (ValueError, RuntimeError)) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/replicated_experts/ring_buddy.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/replicated_experts/ring_buddy.py index bd29c483e9ac..0535bd2926ab 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/replicated_experts/ring_buddy.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/replicated_experts/ring_buddy.py @@ -458,6 +458,9 @@ def _nvfp4_call(backend, routing, persistent: bool = False, stacks=None, persist from tensorrt_llm._torch._experimental.modeling_v2.catalog.moe.fp4_block_scale_moe_runner import ( # noqa: E501 fp4_block_scale_moe_runner, ) + from tensorrt_llm._torch.autotuner import AutoTuner + + from ...ops.moe import decode_expert_tactics as _decode_tactics intermediate = backend.w3_w1_weight.shape[-2] // 2 tune_max_num_tokens = backend.max_num_tokens @@ -521,6 +524,48 @@ def call(stacks, x, local_ids, weights_bf16): stacks["fc31_alpha"], stacks["fc2_alpha"], ) + tactic = None + if not AutoTuner.get().is_tuning_mode: + tactic = _decode_tactics.fixed_tactic( + x.fp4_tensor.device, + act_type, + TOP_K, + 2 * x.fp4_tensor.shape[1], + intermediate, + window, + x.fp4_tensor.shape[0], + ) + if tactic is not None: + return _decode_tactics.run_fixed( + tactic, + [ + None, + None, + x.fp4_tensor, + x_sf, + stacks["w3_w1_weight"], + stacks["w3_w1_weight_scale"].view(torch.float8_e4m3fn), + None, + None, + None, + None, + stacks["w2_weight"], + stacks["w2_weight_scale"].view(torch.float8_e4m3fn), + None, + stacks["fc31_scale_c"], + stacks["fc31_alpha"], + stacks["fc2_alpha"], + weights_bf16, + local_ids, + ], + EXPERTS, + TOP_K, + intermediate, + window, + routing.routing_method_type, + act_type, + tune_max_num_tokens, + ) return fp4_block_scale_moe_runner( None, None, diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/moe/decode_expert_tactics.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/moe/decode_expert_tactics.py new file mode 100644 index 000000000000..b35712085d19 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/moe/decode_expert_tactics.py @@ -0,0 +1,146 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Fixed TRTLLM-gen tactics for the NVFP4 expert call of tep4 decode steps (sm103, sm107). + +``trtllm::fp4_block_scale_moe_runner`` takes its ``[tile_n, config]`` tactic from the +autotuner, which profiles one tactic per power-of-two row bucket at warmup, on a dummy +routing, separately on every rank: every row count of 256 to 511 runs the tactic +profiled at 256 rows, where it settles on 64-token tiles. + +This table pins the tactic of the stacked expert call (``ring_buddy._nvfp4_call``) per +``(expert window, rows)`` for listed decode row counts of tep4 (8 rows per request at the +DFlash K = 7 verify width), one table per GPU: + +* sm107: the 32-token tile the autotuner serves at 128 to 255 rows, for 256 to 384 rows. + It computes every expert row bitwise as the autotuned tactics whose FC2 has the same + MMA K; tactics whose FC2 MMA K differs can round a rare row differently in its last + place (the autotuner picks from both). +* sm103: the 128-token tile the autotuner serves at 512 rows, for 256 to 504 rows. Every + NVFP4 expert kernel sm103 loads has an MMA K of 64, and every tactic valid at these row + counts computes every expert row bitwise as the autotuned 64-token tactics. + +Row counts and windows not listed, and tuning passes, keep the autotuner. A GPU without a +table is not supported. + +The config index is an index into the TRTLLM-gen MoE runner's list of (FC1, FC2) kernel +pairs for one tile N, built from the kernels the running GPU loads, so an entry is only +meaningful for the GPU and kernel library it was measured on; ``fixed_tactic`` also +requires the entry to be a valid config of the call's shape (``get_valid_configs``) and +returns None otherwise. +""" + +from __future__ import annotations + +from functools import lru_cache +from typing import Dict, List, Optional, Tuple + +import torch + +#: The 32-token tile the autotuner serves at 128..255 rows (FC1 +#: ``t128x32x512u2_s5 … m256x32x128``, FC2 ``t128x32x512_s5 … m256x32x128``), pinned for the +#: 256..384-row decode steps (32 to 48 requests) of the 68-expert ``hot_replicas`` window. +#: Its FC2's MMA K is 128, as is that of the tile-64 tactics the autotuner picks there +#: most, so every expert row is bitwise theirs. +_TILE32_DECODE = (32, 236) + +#: ``(local_num_experts, rows) -> (tile_n, config)`` on sm107. Window 68 = the 64 own +#: experts plus 4 hot copies (``hot_replicas``, 33..512 rows). +SM107_TACTICS: Dict[Tuple[int, int], Tuple[int, int]] = { + (68, rows): _TILE32_DECODE for rows in range(256, 385, 8) +} + +#: The 128-token tile the autotuner serves at 512 rows on sm103 (FC1 +#: ``t128x128x512_s3x3x3x3x1x3 … m256x128x64``, FC2 ``t128x128x256_s6 … m256x128x64``), pinned +#: for the 256..504-row decode steps (32 to 63 requests) of the 68-expert window. +_SM103_TILE128_DECODE = (128, 4) + +#: ``(local_num_experts, rows) -> (tile_n, config)`` on sm103. +SM103_TACTICS: Dict[Tuple[int, int], Tuple[int, int]] = { + (68, rows): _SM103_TILE128_DECODE for rows in range(256, 505, 8) +} + +_TABLES = {(10, 3): SM103_TACTICS, (10, 7): SM107_TACTICS} + + +def table_for(capability: Tuple[int, int]) -> Dict[Tuple[int, int], Tuple[int, int]]: + """The table of a device capability: sm103 or sm107; any other GPU is not supported.""" + table = _TABLES.get(tuple(capability)) + if table is None: + raise NotImplementedError( + f"decode expert tactics: no table for compute capability {tuple(capability)} " + "(measured for sm103 and sm107 only)" + ) + return table + + +@lru_cache(maxsize=None) +def _valid(act_type: int, top_k: int, hidden: int, intermediate: int, window: int, rows: int): + from tensorrt_llm._torch.custom_ops.trtllm_gen_custom_ops import FP4BlockScaleMoERunner + + runner = FP4BlockScaleMoERunner.runner_dict.get((act_type,)) + if runner is None: + runner = torch.classes.trtllm.FP4BlockScaleMoERunner(act_type) + FP4BlockScaleMoERunner.runner_dict[(act_type,)] = runner + configs = runner.get_valid_configs(top_k, hidden, intermediate, window, rows) + return frozenset(tuple(int(v) for v in c) for c in configs) + + +def fixed_tactic( + device: torch.device, + act_type: int, + top_k: int, + hidden: int, + intermediate: int, + window: int, + rows: int, +) -> Optional[List[int]]: + """The pinned ``[tile_n, config]`` of an expert call, or None to leave it to the autotuner. + + None when the ``(window, rows)`` is not listed, the call is not the SwiGLU expert + (``act_type`` 0), or the listed tactic is not a valid config of the call's shape; a + device without a table raises (``table_for``). + """ + if act_type != 0: + return None + tactic = table_for(torch.cuda.get_device_capability(device)).get((window, rows)) + if tactic is None: + return None + if tactic not in _valid(act_type, top_k, hidden, intermediate, window, rows): + return None + return list(tactic) + + +def run_fixed( + tactic: List[int], + inputs: List[Optional[torch.Tensor]], + num_experts: int, + top_k: int, + intermediate: int, + window: int, + routing_method_type: int, + act_type: int, + tune_max_num_tokens: int, +) -> List[torch.Tensor]: + """The unfinalized pre-routed expert call on ``tactic``, without the autotuner. + + ``inputs`` are ``FP4BlockScaleMoEInputs``'s 18 tensors in order (``routing_logits`` + None, ``topk_weights`` / ``topk_ids`` the caller's). Returns what + ``trtllm::fp4_block_scale_moe_runner`` returns for ``do_finalize=False``. + """ + from tensorrt_llm._torch.custom_ops.trtllm_gen_custom_ops import FP4BlockScaleMoERunner + + runner = FP4BlockScaleMoERunner( + num_experts, + top_k, + None, + None, + intermediate, + 0, + window, + None, + routing_method_type, + False, + act_type, + tune_max_num_tokens=tune_max_num_tokens, + ) + return runner(inputs, tactic=tactic) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/moe/prefill_cutedsl_experts.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/moe/prefill_cutedsl_experts.py new file mode 100644 index 000000000000..f29b5cb9b895 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/moe/prefill_cutedsl_experts.py @@ -0,0 +1,220 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Prefill routed NVFP4 experts on the Rubin CuTe DSL grouped GEMMs, over TRTLLM-gen's tensors. + +The unfinalized expert call of a large (prefill or mixed) step: the same +``(gemm2_output, expanded_idx_to_permuted_idx)`` contract as TRTLLM-gen's +``fp4_block_scale_moe_runner`` with ``do_finalize=False``, computed by + +* ``torch.ops.trtllm.moe_sort`` (tile ``TILE`` rows), +* the Rubin gather + grouped GEMM + SwiGLU + NVFP4 quantization kernel (FC1), + in its ``trtllm_gen_gated_layout`` mode, +* the Rubin grouped GEMM (FC2) in its ``store_permuted`` + + ``trtllm_gen_n_shuffle`` + ``mma_k64`` mode: one BF16 row per permuted row, + no finalize, accumulated with the same 64-deep MMA as TRTLLM-gen's FC2. + +Both kernels read the served expert tensors as they are (TRTLLM-gen's shuffled +rows and 128x4 swizzled block scales), so no second weight copy exists. The +FC1 epilogue takes the served leaf's two FC1 scales as TRTLLM-gen's gated +GEMM does (gate scale ``fc31_alpha``, output scale ``fc31_scale_c``) and uses +TRTLLM-gen's SwiGLU and NVFP4 quantization arithmetic, so the intermediate +codes and block scales equal TRTLLM-gen's; the model's per-expert post-norm +then runs on the rows exactly as on TRTLLM-gen's. + +Graph capture: every kernel is compiled by ``prepare`` (outside any capture); +the tactics are fixed per row bucket (``TACTICS``), so no autotuning happens in +the served process. +""" + +from __future__ import annotations + +from typing import Dict, Optional, Tuple + +import torch + +#: Routing tile (``moe_sort`` tile_tokens_dim, the kernels' M tile). +TILE = 256 + +#: Static tactics: (FC1, FC2) per smallest row count they serve, largest first. +#: FC1 = (mma_tiler, mma_inst_shape, cluster_shape_mn, raster_along_m, a_path); +#: FC2 = (mma_tiler, mma_inst_shape, cluster_shape_mn, raster_along_m). +TACTICS: Tuple[Tuple[int, tuple, tuple], ...] = ( + ( + 0, + ((256, 256, 256), (256, 256, 128), (2, 1), False, "cpasync"), + ((256, 256, 256), (256, 256, 128), (2, 1), False), + ), +) + + +def tactics_for(rows: int, table=None) -> Tuple[tuple, tuple]: + """The (FC1, FC2) tactics of a call of ``rows`` routed tokens.""" + for min_rows, fc1, fc2 in table or TACTICS: + if rows >= min_rows: + return fc1, fc2 + raise ValueError(f"no tactic for {rows} rows") + + +class PrefillCuteDslExperts: + """One layer's prefill expert call over a TRTLLM-gen NVFP4 leaf's tensors. + + ``backend`` is the leaf (or any object with its attributes: ``w3_w1_weight``, + ``w3_w1_weight_scale``, ``w2_weight``, ``w2_weight_scale``, ``fc31_alpha``, + ``fc31_scale_c``, ``fc2_alpha``, ``fc2_input_scale``). The tensors are read + on every call, as views, so a later re-binding of the leaf's parameters + (e.g. as views of a larger expert stack) is followed and nothing is copied. + The FC1 alpha pair ``[fc31_alpha | fc31_scale_c]`` is built once per set + of scale tensors. + """ + + def __init__(self, backend, num_experts: int, top_k: int, local_expert_offset: int): + from tensorrt_llm._torch.custom_ops import cute_dsl_custom_ops as _ops + + self.backend = backend + self.local_num_experts = int(backend.w3_w1_weight.shape[0]) + self.num_experts = int(num_experts) + self.top_k = int(top_k) + self.local_expert_offset = int(local_expert_offset) + self.hidden = int(backend.w2_weight.shape[1]) + e = self.local_num_experts + self._fc1 = _ops.Sm107BlockScaledContiguousGatherGroupedGemmActFusionRunner( + self.num_experts, + self.top_k, + e, + self.local_expert_offset, + TILE, + trtllm_gen_gated_layout=True, + ) + self._fc2 = _ops.Sm107BlockScaledContiguousGroupedGemmFinalizeFusionRunner( + self.num_experts, + self.top_k, + e, + self.local_expert_offset, + TILE, + torch.bfloat16, + store_permuted=True, + trtllm_gen_n_shuffle=True, + mma_k64=True, + ) + self._scales_key = None + self._scales = None + self._dummy_scales: Dict[torch.device, torch.Tensor] = {} + + def operands(self): + """``(w1, s1, w2, s2, gate_up_alpha, fc2_alpha, global_sf)`` in the kernels' dtypes (views).""" + be = self.backend + e = self.local_num_experts + w31 = be.w3_w1_weight + n1 = int(w31.shape[1]) + w1 = w31.view(torch.uint8).view(torch.float4_e2m1fn_x2) + s1 = be.w3_w1_weight_scale.view(torch.uint8).reshape(e, n1, -1) + w2 = be.w2_weight.view(torch.uint8).view(torch.float4_e2m1fn_x2) + s2 = be.w2_weight_scale.view(torch.uint8).reshape(e, self.hidden, -1) + key = tuple( + (t.data_ptr(), t._version if not t.is_inference() else None) + for t in (be.fc31_alpha, be.fc31_scale_c, be.fc2_alpha, be.fc2_input_scale) + ) + if key != self._scales_key: + # The FC1 kernel takes the output scale from the alpha pair; global_sf + # is a required operand it does not read in this mode. + g2 = be.fc2_input_scale.reshape(1).float() + gate_up_alpha = torch.cat( + (be.fc31_alpha.reshape(e).float(), be.fc31_scale_c.reshape(e).float()) + ).contiguous() + self._scales = ( + gate_up_alpha, + be.fc2_alpha.reshape(e).float().contiguous(), + g2.contiguous(), + ) + self._scales_key = key + return (w1, s1, w2, s2) + self._scales + + def __call__( + self, + x_fp4: torch.Tensor, + x_sf: torch.Tensor, + topk_ids: torch.Tensor, + topk_weights: torch.Tensor, + tactics: Optional[Tuple[tuple, tuple]] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """``(gemm2_output [permuted rows, hidden] bf16, expanded_idx_to_permuted_idx)``. + + ``x_fp4`` is ``[T, hidden / 2]`` packed e2m1, ``x_sf`` the linear e4m3 + block scales (``T * hidden / 16`` bytes), ``topk_ids`` global int32 ids + ``[T, top_k]``, ``topk_weights`` their combine weights (only sorted with + the ids; the rows carry no weight). + """ + rows = int(x_fp4.shape[0]) + fc1_tactic, fc2_tactic = tactics if tactics is not None else tactics_for(rows) + ( + tile_idx_to_expert_idx, + tile_idx_to_mn_limit, + expanded_idx_to_permuted_idx, + permuted_idx_to_expanded_idx, + _total_num_padded_tokens, + num_non_exiting_tiles, + ) = torch.ops.trtllm.moe_sort( + token_selected_experts=topk_ids, + token_final_scales=topk_weights.float(), + num_experts=self.num_experts, + top_k=self.top_k, + local_expert_offset=self.local_expert_offset, + local_num_experts=self.local_num_experts, + tile_tokens_dim=TILE, + ) + w1, s1, w2, s2, gate_up_alpha, fc2_alpha, global_sf = self.operands() + x = x_fp4.view(torch.uint8).view(torch.float4_e2m1fn_x2) + xs = x_sf.view(torch.uint8).reshape(rows, -1) + h, h_sf = self._fc1( + [ + x, + w1, + xs, + s1, + gate_up_alpha, + tile_idx_to_expert_idx, + tile_idx_to_mn_limit, + permuted_idx_to_expanded_idx, + num_non_exiting_tiles, + global_sf, + None, + None, + ], + tactic=fc1_tactic, + ) + m = int(permuted_idx_to_expanded_idx.shape[0]) + out = torch.empty(m, self.hidden, dtype=torch.bfloat16, device=x_fp4.device) + scales = self._dummy_scales.get(x_fp4.device) + if scales is None: + # Unread in store_permuted mode; any valid pointer of the right dtype. + scales = torch.zeros(1, self.top_k, dtype=torch.float32, device=x_fp4.device) + self._dummy_scales[x_fp4.device] = scales + self._fc2( + [ + h, + w2, + h_sf, + s2, + fc2_alpha, + out, + tile_idx_to_expert_idx, + tile_idx_to_mn_limit, + permuted_idx_to_expanded_idx, + num_non_exiting_tiles, + scales, + ], + tactic=fc2_tactic, + ) + return out, expanded_idx_to_permuted_idx + + def prepare(self, rows: int, table=None) -> None: + """Compile both kernels for ``rows`` (outside any graph capture) with an all-local routing.""" + dev = self.backend.w3_w1_weight.device + ids = ( + torch.arange(rows * self.top_k, device=dev, dtype=torch.int32) % self.local_num_experts + + self.local_expert_offset + ).reshape(rows, self.top_k) + x = torch.zeros(rows, self.hidden // 2, dtype=torch.uint8, device=dev) + xs = torch.zeros(rows * self.hidden // 16, dtype=torch.uint8, device=dev) + w = torch.full((rows, self.top_k), 1.0 / self.top_k, device=dev) + self(x, xs, ids, w, tactics=tactics_for(rows, table)) From 19d50b1f6ba30d010e3cd2ee4c95ae808c3302b4 Mon Sep 17 00:00:00 2001 From: Xiao Wang <24860335+xwang233@users.noreply.github.com> Date: Mon, 26 Oct 2026 18:44:06 +0000 Subject: [PATCH 08/16] [None][feat] modeling_v2 musespark_open tep: shard routing and copy-engine gather on sequence-parallel context steps The sequence-parallel context path wires a per-SM set of parts (SHARD_PARTS_BY_SM) and refuses any other SM: - GB300 (sm103): the MoE routes each rank's own rows with the fused routing kernel and gathers the top-8 ids and FP32 weights instead of the router logits; the attention-side residual join also writes the router's FP32 input (router_input); and on steps of at least MIN_ROWS rows the copy engines gather the attention input into the peers' receive buffers while the QKV projection runs in row blocks behind per-peer flags (ce_gather). The gather is agreed on by every rank and falls back to the NCCL all-gather when its buffers cannot be built or a projection cannot run in row blocks; its host flag words use prefer_pinned(). - Rubin (sm107): the MoE re-class only; the router logits are gathered, every row of the step is routed, and the attention input is all-gathered. Routing ids and weights, the router input and the gathered rows are the built-in path's, bit for bit. Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com> --- .../context_sequence_parallel/ce_gather.py | 571 ++++++++++++++++++ .../context_sequence_parallel/forward.py | 104 +++- .../kernels/context_sequence_parallel/moe.py | 157 ++++- .../context_sequence_parallel/router_input.py | 226 +++++++ 4 files changed, 1028 insertions(+), 30 deletions(-) create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/context_sequence_parallel/ce_gather.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/context_sequence_parallel/router_input.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/context_sequence_parallel/ce_gather.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/context_sequence_parallel/ce_gather.py new file mode 100644 index 000000000000..75d2eeec934c --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/context_sequence_parallel/ce_gather.py @@ -0,0 +1,571 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""A sequence-parallel layer's attention input all-gathered by the copy engines, QKV in row blocks. + +What it replaces. On a sequence-parallel context step every layer all-gathers its +attention input (``[shard, 8192]`` BF16 per rank) with an NCCL ring kernel and then +runs the QKV projection over all rows. The ring kernel needs SMs, so it runs next to +nothing and the projection waits for its last byte. + +What runs instead, on steps of at least ``MIN_ROWS`` rows. Rank ``r`` gathers on a +second stream ("copy stream"): + +* after the attention input is produced, the copy engines copy it into rows + ``[r * shard, (r + 1) * shard)`` of each peer's receive buffer, peers ``r+1, r+2, r+3`` + in turn, and after each copy ``cuStreamWriteValue32`` sets the flag word that peer + keeps for source ``r`` (default flags: a fence before the write); +* the QKV projection runs in row blocks on the main stream: rank ``r``'s own rows from + its own attention input first, then each peer's rows from the receive buffer after a + ``cuStreamWaitValue32`` on that peer's flag, peers in the order their copies arrive + (``r-1, r-2, r-3``: every sender copies to its next rank first). After the last + wait the flags are reset (fenced writes) and the main stream joins the copy stream. + +The receive buffers and flag words are one ``McastGPUBuffer`` per rank (the allocation +of the projection-publish transport), used through its unicast peer mappings only. + +Ownership and progress (cuFlow): + +* One writer per region and layer: rank ``p``'s copy writes ``recv_r[p * shard...]``; + one reader: rank ``r``'s block GEMM for ``p``. +* Publication: the copy completes, then the fenced flag write; acquire: the stream-side + wait, then the GEMM. The reader resets its flags after its last wait. +* Visibility: the waits carry no ``CU_STREAM_WAIT_VALUE_FLUSH`` (the device reports + ``CAN_FLUSH_REMOTE_WRITES`` = 0). The copied rows are visible to the reader's GEMM + because the sender orders the copy before its fenced flag write on one stream and + the NVLink peer writes land in the reader's memory before the flag word does. This + rests on that ordering of the peer-write path, not on a documented acquire of the + wait; the unit and in-engine checks under graph replay are bit exact. +* Consumers: on a step this takes, the view ``gather`` returns never holds this + rank's own rows (``project`` reads them from the attention input itself), so + ``project`` must stay its only reader; ``settle`` completes the view in place for + any other reader. +* Reuse (one buffer for every layer): rank ``p`` writes ``recv_r`` again for layer + ``L + 1`` only after its own layer ``L`` passed the o_proj and experts + reduce-scatters, which need rank ``r``'s contributions, issued after ``r``'s layer + ``L`` block GEMMs and flag resets; between steps the forward's final all-gather does + the same. This needs one collective of all ranks between two gathers; every layer + has two. +* Progress: a wait targets a copy its peer enqueued right after its own producer of the + same layer, which depends only on collectives this rank already passed. Copy + engines take no SM and the waits hold none, so nothing assumes co-residency. A peer + that never copies leaves the stream waiting without an error: ``start_watch``'s + thread sees a forward that began and did not end for ``stall_s`` seconds + (``first_stall_s`` before this rank ended its first one), reads this rank's flag + words on another stream, logs the sources that never arrived and aborts. + +Numerics: bitwise. The gathered rows are the same bytes, and each row block's GEMM is +``torch.matmul`` with the projection's own weight, which computes every row as the +whole-step GEMM does on the cuBLAS this was certified with (the in-engine checksum +arm compares the hidden states). +""" + +from __future__ import annotations + +import os +import threading +import time +from typing import Callable, List, Optional, Tuple + +import torch +from cuda.bindings import driver as _cu +from cuda.bindings import runtime as _cudart + +from tensorrt_llm._torch.distributed import allgather +from tensorrt_llm._utils import prefer_pinned +from tensorrt_llm.logger import logger + +try: + import triton + import triton.language as tl +except ImportError: # pragma: no cover - the image ships Triton + triton = None + +#: Seconds a sequence-parallel forward may stay begun and not ended before the watch +#: aborts, once this rank has ended one (a served context step takes well under a second; +#: a later step may still load or compile a kernel). +STALL_S = 120.0 +#: The same bound while this rank has not yet ended any sequence-parallel forward: the +#: first one (the warmup) compiles and loads the step's kernels, which on a cold cache +#: takes longer than ``STALL_S``. +FIRST_STALL_S = 900.0 +#: The fewest rows of a step whose attention input is gathered here: below it the +#: flags' fixed per-flag cost eats the overlap (measured: next to no gain per layer at +#: 2048 rows, a clear and growing gain from 4096 rows up). +MIN_ROWS = 4096 +#: int32 words between two flag words (one 128-byte line each). +FLAG_STRIDE = 32 +#: Seconds between two reads of the forward counters by the watch thread. +WATCH_INTERVAL_S = 1.0 + +_ALIGN = 2 << 20 + + +def _check(result): + err = result[0] if isinstance(result, tuple) else result + if err != _cu.CUresult.CUDA_SUCCESS: + raise RuntimeError(f"[musespark_open] ce_gather: CUDA driver call failed: {err}") + + +def _write_value(address: int, value: int) -> None: + """``cuStreamWriteValue32`` on the current stream with the default (fenced) flags.""" + _check( + _cu.cuStreamWriteValue32( + _cu.CUstream(torch.cuda.current_stream().cuda_stream), + _cu.CUdeviceptr(address), + value, + _cu.CUstreamWriteValue_flags.CU_STREAM_WRITE_VALUE_DEFAULT, + ) + ) + + +def _wait_value(address: int, value: int) -> None: + """``cuStreamWaitValue32`` (equal) on the current stream.""" + _check( + _cu.cuStreamWaitValue32( + _cu.CUstream(torch.cuda.current_stream().cuda_stream), + _cu.CUdeviceptr(address), + value, + _cu.CUstreamWaitValue_flags.CU_STREAM_WAIT_VALUE_EQ, + ) + ) + + +class _HostWords: + """Pinned host int32 words a kernel updates at system scope and a host thread reads unsynchronized. + + Triton reads ``data_ptr()`` and ``dtype`` of a pointer argument; the words' device + address is their host address (unified addressing). + """ + + dtype = torch.int32 + + def __init__(self, count: int) -> None: + # Without pinned memory there is no mapped device pointer: the error below makes + # ``CopyEngineGather.build`` refuse, and the NCCL all-gather runs instead. + self.host = torch.zeros(count, dtype=torch.int32, pin_memory=prefer_pinned()) + status, pointer = _cudart.cudaHostGetDevicePointer(self.host.data_ptr(), 0) + if status != _cudart.cudaError_t.cudaSuccess: + raise RuntimeError(f"cudaHostGetDevicePointer of the forward counters: {status}") + self._pointer = int(pointer) + + def data_ptr(self) -> int: + return self._pointer + + +if triton is not None: + + @triton.jit + def _count_kernel(words_ptr, INDEX: tl.constexpr): + """Add one to ``words[INDEX]`` (pinned host memory) at system scope.""" + tl.atomic_add(words_ptr + INDEX, 1, sem="release", scope="sys") + + +#: The forward counters' slots. +_BEGUN, _ENDED = 0, 1 + + +def projection_refusal(linear) -> Optional[str]: + """Why ``linear`` (a layer's ``qkv_proj``) cannot run in row blocks here, or None.""" + from tensorrt_llm._torch.modules.linear import UnquantizedLinearMethod + + weight = getattr(linear, "weight", None) + if not isinstance(getattr(linear, "quant_method", None), UnquantizedLinearMethod): + return "the QKV projection is quantized" + if weight is None or weight.dtype != torch.bfloat16 or weight.dim() != 2: + return "the QKV projection's weight is not a 2-D BF16 matrix" + if getattr(linear, "bias", None) is not None: + return "the QKV projection has a bias" + if getattr(linear, "use_cute_dsl_bf16_gemm", False) or getattr( + linear, "use_custom_cublas_mm", False + ): + return "the QKV projection runs a custom GEMM" + if getattr(linear, "lora", None) is not None: + return "the QKV projection carries LoRA" + return None + + +class CopyEngineGather: + """The copy-engine gather of every sequence-parallel layer of one model, and its row-block QKV. + + Built (``build``, collective over the TP group) on the first eager prefill step, before + any CUDA graph is captured; until then, and on steps it does not take, ``gather`` + is the NCCL all-gather the built-in layer runs. + """ + + def __init__(self, mapping, max_rows: int, hidden: int, comm=None) -> None: + self.mapping = mapping + self.rank = mapping.tp_rank + self.world = mapping.tp_size + self.hidden = int(hidden) + # A step of ``rows`` rows gathers ``world * ceil(rows / world)`` rows. + self.max_rows = -(-int(max_rows) // self.world) * self.world + self._comm = comm + self.built = False + self.refused: Optional[str] = None + self.allocated = 0 + self._pending: Optional[Tuple[torch.Tensor, torch.Tensor, int, int]] = None + self._watch: Optional[threading.Thread] = None + self._watch_stop: Optional[threading.Event] = None + #: Set (instead of aborting) when the watch was started with ``abort=False`` (tests). + self.stalled: Optional[dict] = None + self._flag_snap: Optional[torch.Tensor] = None + self.gathers = 0 + + # ------------------------------------------------------------------ build + + def build(self) -> None: + """Allocate the receive buffers and flags on every rank (collective); refuse alike on failure.""" + from tensorrt_llm._torch.distributed.ops import ( + _get_mnnvl_workspace_comm, + _make_mnnvl_mcast_buffer, + _mnnvl_workspace_all_succeeded, + ) + + if self.built: + return + self.built = True + comm = self._comm if self._comm is not None else _get_mnnvl_workspace_comm(self.mapping) + refusal = None + if torch.cuda.is_current_stream_capturing(): + refusal = "the first prefill step is captured" + elif triton is None: + refusal = "Triton is not available" + if not _mnnvl_workspace_all_succeeded(comm, refusal is None): + self.refused = refusal or "another rank refused" + return + error: Optional[BaseException] = None + try: + recv_bytes = -(-self.max_rows * self.hidden * 2 // _ALIGN) * _ALIGN + flag_words = self.world * FLAG_STRIDE + total = recv_bytes + flag_words * 4 + torch.cuda.synchronize() + free_before = torch.cuda.mem_get_info()[0] + self._handle = _make_mnnvl_mcast_buffer( + comm, total, self.mapping, self.mapping.is_multi_node() + ) + self.recv: List[torch.Tensor] = [ + self._handle.get_uc_buffer(p, (self.max_rows, self.hidden), torch.bfloat16, 0) + for p in range(self.world) + ] + self.flags = self._handle.get_uc_buffer( + self.rank, (flag_words,), torch.int32, recv_bytes // 4 + ) + self._flag_base = [ + self._handle.get_uc_buffer( + p, (flag_words,), torch.int32, recv_bytes // 4 + ).data_ptr() + for p in range(self.world) + ] + self.flags.zero_() + self.copy_stream = torch.cuda.Stream() + self.watch_stream = torch.cuda.Stream() + self.counters = _HostWords(2) + # Launch both counter kernels and the flag read once now, while no stream waits: + # a kernel's first launch loads its module, and that load does not return while + # any stream of the device is blocked on a flag a peer has not yet written. + self._count(_BEGUN) + self._count(_ENDED) + self.flag_words() + torch.cuda.synchronize() + self.counters.host.zero_() + self.allocated = max(0, free_before - torch.cuda.mem_get_info()[0]) + except Exception as exc: # noqa: BLE001 - every failure must reach the agreement + error = exc + if not _mnnvl_workspace_all_succeeded(comm, error is None): + self.refused = f"the buffers could not be allocated: {error!r}" + return + logger.info( + f"[musespark_open] ce_gather: sequence-parallel steps of {MIN_ROWS}..{self.max_rows} rows " + f"gather the attention input on the copy engines ({total / 2**20:.1f} MiB requested, " + f"{self.allocated / 2**20:.1f} MiB of device memory on this rank)" + ) + + @property + def armed(self) -> bool: + return self.built and self.refused is None + + def takes(self, rows: int) -> bool: + """Whether a step of ``rows`` rows gathers here (host facts, alike on every rank).""" + return self.armed and MIN_ROWS <= rows <= self.max_rows + + # --------------------------------------------------------------- protocol + + def _flag_address(self, owner: int, source: int) -> int: + return self._flag_base[owner] + 4 * source * FLAG_STRIDE + + def arrival_order(self) -> List[int]: + """The sources in the order their copies reach this rank.""" + return [(self.rank - k) % self.world for k in range(1, self.world)] + + def gather(self, attn_input: torch.Tensor, rows: int) -> torch.Tensor: + """The step's attention input, ``[rows, hidden]``: the hook the built-in layer calls. + + On a step this takes, the returned view's peer rows are complete only after + the row-block projection's waits, so it is handed to that projection only + (``project``); every other step gets the NCCL all-gather. + """ + if not self.takes(rows): + return allgather(attn_input, self.mapping, dim=0)[:rows] + if self._pending is not None: + raise RuntimeError( + "[musespark_open] ce_gather: a gathered attention input was not projected" + ) + shard = int(attn_input.shape[0]) + if shard * self.world < rows or not attn_input.is_contiguous(): + raise ValueError( + f"[musespark_open] ce_gather: a [{shard}, ...] shard cannot hold {rows} rows" + ) + rank = self.rank + mine = slice(rank * shard, (rank + 1) * shard) + main = torch.cuda.current_stream() + self.copy_stream.wait_stream(main) + with torch.cuda.stream(self.copy_stream): + for k in range(1, self.world): + peer = (rank + k) % self.world + self.recv[peer][mine].copy_(attn_input) + _write_value(self._flag_address(peer, rank), 1) + view = self.recv[rank][:rows] + self._pending = (view, attn_input, shard, rows) + self.gathers += 1 + return view + + def is_pending(self, tensor: torch.Tensor) -> bool: + """Whether ``tensor`` is the view ``gather`` returned and nothing has projected yet.""" + pending = self._pending + return ( + pending is not None + and tensor.data_ptr() == pending[0].data_ptr() + and tensor.shape == pending[0].shape + ) + + def project(self, weight: torch.Tensor) -> torch.Tensor: + """``view @ weight.T`` for the pending view, in row blocks behind the flags.""" + view, own, shard, rows = self._pending + self._pending = None + out = torch.empty((rows, weight.shape[0]), dtype=view.dtype, device=view.device) + w_t = weight.t() + rank = self.rank + + def block(source: int) -> Optional[slice]: + lo, hi = source * shard, min((source + 1) * shard, rows) + return slice(lo, hi) if hi > lo else None + + own_rows = block(rank) + if own_rows is not None: + torch.matmul(own[: own_rows.stop - own_rows.start], w_t, out=out[own_rows]) + order = self.arrival_order() + for index, source in enumerate(order): + _wait_value(self._flag_address(rank, source), 1) + if index == len(order) - 1: + for done in order: + _write_value(self._flag_address(rank, done), 0) + rows_of = block(source) + if rows_of is not None: + torch.matmul(self.recv[rank][rows_of], w_t, out=out[rows_of]) + torch.cuda.current_stream().wait_stream(self.copy_stream) + return out + + def settle(self) -> None: + """Complete the pending view in place (every flag waited and reset, own rows copied) and join.""" + view, own, shard, rows = self._pending + self._pending = None + order = self.arrival_order() + for source in order: + _wait_value(self._flag_address(self.rank, source), 1) + for source in order: + _write_value(self._flag_address(self.rank, source), 0) + lo, hi = self.rank * shard, min((self.rank + 1) * shard, rows) + if hi > lo: + view[lo:hi].copy_(own[: hi - lo]) + torch.cuda.current_stream().wait_stream(self.copy_stream) + + # ------------------------------------------------------------- the watch + + def _count(self, index: int) -> None: + _count_kernel[(1,)](self.counters, INDEX=index, num_warps=1) + + def begin(self) -> None: + """Count a sequence-parallel forward as begun (a one-thread kernel; captured with the step).""" + if self.armed: + self._count(_BEGUN) + + def end(self) -> None: + """Count a sequence-parallel forward as ended.""" + if self.armed: + self._count(_ENDED) + + def flag_words(self, bound_s: float = 5.0) -> Optional[List[int]]: + """This rank's flag word per source, or None if not read within ``bound_s``. + + One contiguous device-to-host copy of the flag block on the watch stream: a + copy-engine transfer into a pinned buffer kept from the first call, with no kernel + and no allocation, so it completes while the compute streams wait on a flag. + """ + # Inference mode on every thread: the buffer is first made by ``build`` inside the + # engine's inference mode and later refilled from the watch's helper thread. + with torch.inference_mode(): + snap = self._flag_snap + if snap is None: + snap = torch.empty( + self.flags.numel(), dtype=torch.int32, pin_memory=prefer_pinned() + ) + self._flag_snap = snap + with torch.cuda.stream(self.watch_stream): + snap.copy_(self.flags, non_blocking=True) + done = torch.cuda.Event() + done.record() + start = time.monotonic() + while not done.query(): + if time.monotonic() - start > bound_s: + return None + time.sleep(0.001) + return [int(v) for v in snap[::FLAG_STRIDE]] + + def _flag_words_bounded(self, bound_s: float = 5.0) -> Optional[List[int]]: + """``flag_words`` from a helper thread, abandoned after ``bound_s`` even if a CUDA call blocks.""" + result: List[Optional[List[int]]] = [] + reader = threading.Thread( + target=lambda: result.append(self.flag_words(bound_s)), + name="musespark_open.ce_gather_flags", + daemon=True, + ) + reader.start() + reader.join(bound_s + 1.0) + return result[0] if result else None + + def start_watch( + self, + stall_s: Optional[float] = None, + interval_s: float = WATCH_INTERVAL_S, + first_stall_s: Optional[float] = None, + abort: bool = True, + on_stall: Optional[Callable[[dict], None]] = None, + ) -> bool: + """Watch the forward counters from a host thread; returns whether a thread runs. + + A forward that began and has not ended for ``stall_s`` seconds (default + ``STALL_S``) is waiting on a flag that + no peer set: the thread logs this rank's flag words (the sources still at 0 never + arrived) and aborts the process, or with ``abort=False`` records the report in + ``stalled`` and calls ``on_stall`` (tests). Until this rank has ended one forward + the bound is ``first_stall_s`` (default ``FIRST_STALL_S``), at least ``stall_s``. + """ + if not self.armed or self._watch is not None: + return self._watch is not None + if stall_s is None: + stall_s = STALL_S + if first_stall_s is None: + first_stall_s = FIRST_STALL_S + first_stall_s = max(first_stall_s, stall_s) + stop = threading.Event() + device = torch.cuda.current_device() + + def watch(): + torch.cuda.set_device(device) + seen_end, since = None, None + while not stop.wait(interval_s): + begun = int(self.counters.host[_BEGUN]) + ended = int(self.counters.host[_ENDED]) + now = time.monotonic() + if begun <= ended or ended != seen_end: + seen_end, since = ended, now + continue + if now - since < (stall_s if ended > 0 else first_stall_s): + continue + logger.error( + f"[musespark_open] ce_gather: rank {self.rank} has waited {now - since:.1f} s in a " + f"sequence-parallel forward (begun {begun}, ended {ended}); reading its flag words" + ) + words = self._flag_words_bounded() + missing = ( + None if words is None else [s for s in self.arrival_order() if words[s] == 0] + ) + report = dict( + rank=self.rank, + begun=begun, + ended=ended, + stalled_s=round(now - since, 3), + flags=words, + missing=missing, + ) + logger.error( + f"[musespark_open] ce_gather: rank {self.rank} has waited {now - since:.1f} s in a " + f"sequence-parallel forward (begun {begun}, ended {ended}); flag words per source " + f"{words}, sources that never arrived {missing}" + ) + if abort: + os.abort() + self.stalled = report + if on_stall is not None: + on_stall(report) + return + + self._watch_stop = stop + self._watch = threading.Thread( + target=watch, name="musespark_open.ce_gather_watch", daemon=True + ) + self._watch.start() + return True + + def stop_watch(self) -> None: + """Stop the watch thread, if one runs (tests).""" + if self._watch is not None: + self._watch_stop.set() + self._watch.join() + self._watch = None + + +class CopyEngineGatherQKV: + """Mixin for a layer's ``self_attn.qkv_proj``: the pending gathered view runs in row blocks.""" + + _sp_ce_gather: Optional[CopyEngineGather] = None + + def forward(self, input, *args, **kwargs): + gather = self._sp_ce_gather + if gather is not None and isinstance(input, torch.Tensor) and gather.is_pending(input): + if not args and not kwargs.get("lora_params"): + return gather.project(self.weight) + gather.settle() + return super().forward(input, *args, **kwargs) + + +_CLASSES = {} + + +def _with_mixin(base: type) -> type: + if issubclass(base, CopyEngineGatherQKV): + return base + if base not in _CLASSES: + _CLASSES[base] = type( + "CopyEngineGatherQKVOver" + base.__name__, + (CopyEngineGatherQKV, base), + {"__module__": __name__}, + ) + return _CLASSES[base] + + +def install(model) -> Optional[CopyEngineGather]: + """Hook every layer's attention-input gather and return the gather, or None. + + Changes nothing and returns None when a layer's projection cannot run in row blocks + (logged); the buffers are built later (``CopyEngineGather.build``). + """ + layers = list(model.layers) + reasons = {layer.layer_idx: projection_refusal(layer.self_attn.qkv_proj) for layer in layers} + refused = {idx: why for idx, why in reasons.items() if why is not None} + if refused: + logger.warning( + f"[musespark_open] ce_gather: not armed: {sorted(set(refused.values()))} " + f"(layers {sorted(refused)}); the attention input is all-gathered by NCCL" + ) + return None + config = model.model_config + gather = CopyEngineGather( + model._sp_mapping, config.max_num_tokens, config.pretrained_config.hidden_size + ) + for layer in layers: + qkv = layer.self_attn.qkv_proj + qkv.__class__ = _with_mixin(type(qkv)) + qkv._sp_ce_gather = gather + layer._sp_attn_input_gather = gather.gather + return gather diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/context_sequence_parallel/forward.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/context_sequence_parallel/forward.py index fbb1d93b1528..e3b1a3d083d0 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/context_sequence_parallel/forward.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/context_sequence_parallel/forward.py @@ -22,7 +22,14 @@ bit for bit) and reduce-scatters the experts' partial sums in place of their all-reduce. Attention and experts still see every row of the step, so the generation rows of a mixed step are served as the context rows are. After the last layer the -final norm runs on the shard and the step's rows are all-gathered. +final norm runs on the shard and the step's rows are all-gathered. On sm103 +(``SHARD_PARTS_BY_SM``) each rank also routes its own rows and gathers their top-8 ids +and weights instead of the router logits (``moe.shard_routing``), and the attention-side +residual join writes the router's FP32 input next to the pre-FFN norm output +(``router_input``); both are bitwise. There the attention input of steps of at least +``ce_gather.MIN_ROWS`` rows is gathered by the copy engines and the QKV projection runs +in row blocks behind per-peer flags (``ce_gather``); also bitwise. On sm107 the router +logits and the attention input are all-gathered. DFlash taps the residual stream after its capture layers, and that stream is sharded here. For the forward the step's capture buffer @@ -139,23 +146,52 @@ def refusal(model: MuseSparkOpenModel) -> Optional[str]: return None -def install(model: MuseSparkOpenModel, sm: Optional[int] = None) -> "ContextSequenceParallel": +#: The SM versions ``install`` wires, each with its own set of the path's parts +#: (``SHARD_PARTS_BY_SM``). +SHARD_PARTS_BY_SM = { + # sm103: the MoE re-class with routing on the shard, the router's FP32 input from the + # attention-side join, and the copy-engine gather of the attention input. + 103: ("moe", "shard_routing", "router_input", "ce_gather"), + # sm107: the MoE re-class only; the router logits are gathered and every row of the + # step is routed, as the built-in arm does, and the attention input is all-gathered. + 107: ("moe",), +} + + +def install(model: MuseSparkOpenModel, sm: Optional[int]) -> "ContextSequenceParallel": """The path's forward for ``model``, with every layer's MoE re-classed (``moe.install``). - ``sm`` is the SM version the model's composition was built for; it picks the eager - threshold (``EAGER_MIN_ROWS_BY_SM``). + ``sm`` is the SM version the model's composition was built for + (``dispatch.sm_of``), required: the parts differ per SM, so none is assumed. It + picks the eager threshold (``EAGER_MIN_ROWS_BY_SM``) and the parts installed + beside the MoE re-class (``SHARD_PARTS_BY_SM``). The MNNVL all-reduce workspaces are grown on the first prefill step (``ContextSequenceParallel.reserve``). Raises: ValueError: when ``refusal`` names a reason, a layer's MoE is not the built-in one, - or a layer's fused finalize all-reduce serves ``MIN_ROWS`` rows. + a layer's fused finalize all-reduce serves ``MIN_ROWS`` rows, or ``sm`` is + neither 103 nor 107 (None included). """ from . import moe as _moe + parts = SHARD_PARTS_BY_SM.get(sm) + if parts is None: + got = "no SM version" if sm is None else f"sm{sm}" + raise ValueError( + f"[musespark_open] context_sequence_parallel is written for sm103 and sm107, got {got}" + ) path = ContextSequenceParallel(model, sm) - _moe.install(model, min_rows=MIN_ROWS) + _moe.install(model, min_rows=MIN_ROWS, shard_routing="shard_routing" in parts) + if "router_input" in parts: + from . import router_input as _router_input + + _router_input.install(model, path.step) + if "ce_gather" in parts: + from . import ce_gather as _ce_gather + + path.ce_gather = _ce_gather.install(model) return path @@ -223,11 +259,19 @@ class ContextSequenceParallel: """ def __init__(self, model: MuseSparkOpenModel, sm: Optional[int] = None) -> None: + from .router_input import SequenceParallelStep + reason = refusal(model) if reason is not None: raise ValueError(f"[musespark_open] context_sequence_parallel: {reason}") self.mapping = model._sp_mapping self.reserved = False + # The copy-engine gather of the attention input (``ce_gather.install``); None when + # a layer's QKV projection cannot run in row blocks. + self.ce_gather = None + # Set while ``forward`` runs the layers: the attention-side joins then also write + # the router's FP32 input (``router_input``). + self.step = SequenceParallelStep() #: The fewest rows of an eager step that runs sequence-parallel on this SM. self.eager_min_rows = EAGER_MIN_ROWS_BY_SM.get(sm, EAGER_MIN_ROWS) @@ -247,6 +291,9 @@ def reserve(self, model: MuseSparkOpenModel) -> None: "on an eager prefill step before any capture" ) _reserve_all_reduce_workspaces(model) + if self.ce_gather is not None: + self.ce_gather.build() + self.ce_gather.start_watch() self.reserved = True def takes(self, attn_metadata: Any, rows: int, kwargs: Dict[str, Any]) -> bool: @@ -293,24 +340,39 @@ def forward( embeds = embeds * core.m_emb # The FP32 residual stream of this rank's rows, entered once. hidden_states = embeds.to(torch.float32) + ce_gather = ( + self.ce_gather if self.ce_gather is not None and self.ce_gather.takes(rows) else None + ) with _shard_capture(kwargs.get("spec_metadata"), shard) as (capture, step_capture): layers = core.layers precomputed_attn_input = None - for index, decoder_layer in enumerate(layers): - next_norm = layers[index + 1].pre_qkv_proj_norm if index + 1 < len(layers) else None - result = decoder_layer( - position_ids=position_ids, - hidden_states=hidden_states, - attn_metadata=attn_metadata, - precomputed_attn_input=precomputed_attn_input, - next_layer_norm=next_norm, - sp_num_tokens=rows, - **kwargs, - ) - if next_norm is None: - hidden_states = result - else: - hidden_states, precomputed_attn_input = result + self.step.active = True + if ce_gather is not None: + ce_gather.begin() + try: + for index, decoder_layer in enumerate(layers): + next_norm = ( + layers[index + 1].pre_qkv_proj_norm if index + 1 < len(layers) else None + ) + result = decoder_layer( + position_ids=position_ids, + hidden_states=hidden_states, + attn_metadata=attn_metadata, + precomputed_attn_input=precomputed_attn_input, + next_layer_norm=next_norm, + sp_num_tokens=rows, + **kwargs, + ) + if next_norm is None: + hidden_states = result + else: + hidden_states, precomputed_attn_input = result + finally: + self.step.active = False + # Also counted when a layer raised: the watch then sees the forward end once + # its queued work does, rather than a forward that never ends. + if ce_gather is not None: + ce_gather.end() output = core.norm(hidden_states.to(core.dtype)) if capture is None: return allgather(output, mapping, dim=0)[:rows] diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/context_sequence_parallel/moe.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/context_sequence_parallel/moe.py index 70df6ccbd953..e586044a8e0e 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/context_sequence_parallel/moe.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/context_sequence_parallel/moe.py @@ -30,6 +30,24 @@ store rounds each FP32 sum to BF16 to nearest even, as the cast does, so the partials are the built-in arm's bit for bit and the FP32 partials are never written or read back. The built-in finalize still decides which kernel runs. + +Routing on the shard (``shard_routing``). The built-in arm gathers every rank's FP32 +router logits (1 KiB per row) and then routes every row of the step on every rank. The +routing is per row, so where the built-in routing method would run its fused top-8 +kernel on the gathered logits, each rank runs that kernel on its own rows -- launched as +the call over all the step's rows would be (``TORCH_ORDER`` and warps chosen from the +step's row count) -- and the top-8 ids and FP32 weights (64 B per row) are gathered with +the NVFP4 payload instead of the logits. The routing method's one-slot memo then holds +the gathered ids and weights, keyed on a placeholder that stands for the logits (no +values); the experts and the finalize read the memo, so both see the ids and weights +the built-in arm computes, bit for bit. +Every consumer of a shard-routed call's ``router_logits`` must read through that memo +(``routing_method.apply``): the placeholder's values are uninitialized. + +The router's FP32 input. When the layer's attention-side residual join already wrote +the FP32 widening of this call's input (``router_input.SequenceParallelAttnJoin``, +``sp_router_input``), the router GEMM reads it instead of a fresh BF16 -> FP32 copy: +the same values, the same GEMM. """ from __future__ import annotations @@ -43,13 +61,99 @@ from tensorrt_llm._torch.models import modeling_musespark_open as _builtin from tensorrt_llm._torch.models.modeling_musespark_open import ( MuseSparkOpenMoE, + MuseSparkOpenRoutingMethod, _pad_sequence_parallel_tokens, + _version_or_none, ) from tensorrt_llm._torch.utils import Fp4QuantizedTensor from tensorrt_llm.logger import logger _CLASSES: Dict[Tuple[type, type], type] = {} +#: The routed experts and the top-k the fused routing kernel serves. +_EXPERTS, _TOP_K = 256, 8 + + +def shard_routing( + routing, logits: torch.Tensor, rows: int +) -> Optional[Tuple[torch.Tensor, torch.Tensor]]: + """The fused routing kernel's ``(selected, weights)`` for the shard's ``logits``, or None. + + Launched with the arguments ``MuseSparkOpenRoutingMethod.apply`` chooses for a call + over all ``rows`` rows of the step (``TORCH_ORDER`` above 128 rows when the method + takes it, one warp from 512 rows), one program per row, so every row's ids and + weights are that call's. None where ``apply`` on the step's logits would not run the + fused kernel, or where its memo is off (the experts then could not read the + gathered routing). + """ + bias = getattr(routing, "e_score_correction_bias", None) + if not ( + isinstance(routing, MuseSparkOpenRoutingMethod) + and routing._memo_enabled + and rows > 0 + and (rows <= 128 or routing._full_routing_all_batches) + and routing.top_k == _TOP_K + and routing.output_dtype == torch.float32 + and logits.is_cuda + and logits.dtype == torch.float32 + and logits.ndim == 2 + and logits.shape[0] > 0 + and logits.shape[1] == _EXPERTS + and logits.is_contiguous() + and bias is not None + and bias.is_cuda + and bias.dtype == torch.float32 + and bias.shape == (_EXPERTS,) + and bias.is_contiguous() + and not torch.compiler.is_compiling() + ): + return None + shard = logits.shape[0] + selected = torch.empty((shard, _TOP_K), device=logits.device, dtype=torch.int32) + weights = torch.empty_like(selected, dtype=torch.float32) + _builtin._full_routing_kernel[(shard,)]( + logits, + bias, + selected, + weights, + routing.renorm_eps, + EXPERTS=_EXPERTS, + TORCH_ORDER=routing._full_routing_torch_order and rows > 128, + num_warps=1 if rows >= 512 else 4, + ) + return selected, weights + + +def routed_placeholder( + routing, selected: torch.Tensor, weights: torch.Tensor, rows: int +) -> torch.Tensor: + """A ``[rows, 256]`` FP32 stand-in for the step's logits, memoized to ``(selected, weights)``. + + The stand-in is a zero-stride view of one element: it carries no values and is never + read, because every routing call on it returns the memo. + """ + placeholder = selected.new_empty((1, 1), dtype=torch.float32).expand(rows, _EXPERTS) + routing._routing_memo = (placeholder, _version_or_none(placeholder), selected, weights) + return placeholder + + +def _fp32_router_input_takes(moe, hidden_states: torch.Tensor) -> bool: + """Whether the built-in ``router_logits`` on ``hidden_states`` is ``router(hidden.float()) * m``. + + That is its last arm: more than 128 rows (the small-batch kernels end there), no + BF16x3 router and no armed ``router_gemm`` call. + """ + call = getattr(moe, "_router_gemm_call", None) + weight = getattr(moe.router, "weight", None) + return ( + hidden_states.ndim == 2 + and hidden_states.shape[0] > 128 + and not getattr(moe, "_bf16x3_router_enabled", False) + and (call is None or not call.is_armed) + and weight is not None + and weight.dtype == torch.float32 + ) + def _rows_of(tensor, rows: int, shape_tail: Tuple[int, ...]): """The first ``rows`` rows of a gathered ``[*, width]`` tensor, in the shard's own layout.""" @@ -144,6 +248,25 @@ class ContextSequenceParallelMoE: Acts on calls with ``sp_num_tokens`` only; every other call is the class behind it. """ + #: ``(norm output, its FP32 widening)`` left by the layer's attention-side join for + #: the next sequence-parallel call (``router_input``); consumed or dropped by it. + sp_router_input: Optional[Tuple[torch.Tensor, torch.Tensor]] = None + #: Whether a sequence-parallel call routes the shard's rows and gathers their ids and + #: weights (``shard_routing``); otherwise it gathers the router logits (``install``). + _sp_shard_routing: bool = True + + def _sp_router_logits(self, hidden_states: torch.Tensor) -> torch.Tensor: + """The router logits of the shard's rows, from the join's FP32 rows when it left them.""" + handed = self.sp_router_input + self.sp_router_input = None + if ( + handed is not None + and handed[0] is hidden_states + and _fp32_router_input_takes(self, hidden_states) + ): + return self.router(handed[1]) * self.metap_multiplier + return self.router_logits(hidden_states) + def forward(self, hidden_states, attn_metadata, *args, **kwargs): rows: Optional[int] = kwargs.get("sp_num_tokens") experts = self.experts @@ -153,11 +276,13 @@ def forward(self, hidden_states, attn_metadata, *args, **kwargs): or getattr(experts, "comm", None) is not None or rows > getattr(experts, "moe_max_num_tokens", rows) ): + self.sp_router_input = None return super().forward(hidden_states, attn_metadata, *args, **kwargs) mapping = self.mapping shard = hidden_states.shape[0] + routing_method = experts.routing_method with self.nvtx("router"): - router_logits = self.router_logits(hidden_states) + router_logits = self._sp_router_logits(hidden_states) with self.nvtx("pool"): prepared = self.prepare_expert_input(self.pooler(hidden_states)) if isinstance(prepared, Fp4QuantizedTensor): @@ -169,21 +294,32 @@ def forward(self, hidden_states, attn_metadata, *args, **kwargs): f"{tuple(scale.shape)} scales, swizzled {prepared.is_sf_swizzled}) " f"for {shard} rows" ) - logits, fp4_rows, scale_rows = allgather( - [router_logits, fp4.reshape(shard, -1), scale.reshape(shard, -1)], - mapping, - dim=0, + routed = ( + shard_routing(routing_method, router_logits, rows) + if self._sp_shard_routing + else None ) + head = [router_logits] if routed is None else list(routed) + gathered = allgather( + [*head, fp4.reshape(shard, -1), scale.reshape(shard, -1)], mapping, dim=0 + ) + fp4_rows, scale_rows = gathered[-2:] scale_tail = tuple(scale.shape[1:]) if scale.dim() > 1 else () pooled = Fp4QuantizedTensor( _rows_of(fp4_rows, rows, tuple(fp4.shape[1:])), scale_rows[:rows].reshape((rows, *scale_tail) if scale_tail else (-1,)), is_sf_swizzled=False, ) + if routed is None: + router_logits = gathered[0][:rows] + else: + router_logits = routed_placeholder( + routing_method, gathered[0][:rows], gathered[1][:rows], rows + ) else: logits, pooled = allgather([router_logits, prepared], mapping, dim=0) pooled = pooled[:rows] - router_logits = logits[:rows] + router_logits = logits[:rows] with self.nvtx("experts"): outputs = experts( pooled, @@ -201,7 +337,6 @@ def forward(self, hidden_states, attn_metadata, *args, **kwargs): f"expanded_idx_to_permuted_idx) triple, got {type(outputs).__name__}" ) gemm2_output, _, expanded_idx_to_permuted_idx = outputs - routing_method = experts.routing_method with self.nvtx("router"): _, token_final_scales = routing_method.apply(router_logits) routing_method.clear_routing_memo() @@ -235,10 +370,13 @@ def _with_mixin(base: type) -> type: return _CLASSES[key] -def install(model, min_rows: int) -> int: +def install(model, min_rows: int, shard_routing: bool = True) -> int: """Re-class every layer's MoE with the mixin in front of its class; returns how many changed. - ``min_rows`` is the fewest rows of a sequence-parallel step. Runs after every + ``min_rows`` is the fewest rows of a sequence-parallel step. ``shard_routing`` has a + sequence-parallel call route the shard's rows and gather their top-8 ids and weights + (``shard_routing``); without it the call gathers the router logits and routes every + row of the step, as the built-in arm does. Runs after every construct-time re-class of the MoE (``PrefillTarget`` is built last in ``MuseSparkOpenTepModelingV2.adopt``), so no later fixed re-class drops the mixin. @@ -265,4 +403,5 @@ def install(model, min_rows: int) -> int: if type(moe) is not mixed: moe.__class__ = mixed changed += 1 + moe._sp_shard_routing = bool(shard_routing) return changed diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/context_sequence_parallel/router_input.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/context_sequence_parallel/router_input.py new file mode 100644 index 000000000000..2e7639f2ce33 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/context_sequence_parallel/router_input.py @@ -0,0 +1,226 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The attention-side residual join of a sequence-parallel layer also writes the router's FP32 input. + +On a sequence-parallel context step each rank's MoE routes its own rows: the built-in +router reads the pre-FFN norm output (BF16, ``[shard, 8192]``) as FP32, so the step +runs a standalone BF16 -> FP32 copy of it in front of the router GEMM (``router_logits`` +-> ``self.router(hidden_states.to(torch.float32))``). That norm output is written by +the attention-side residual join (``MuseSparkOpenGatedResidual._add_after_rmsnorm_and_norm``, +one Triton kernel), which holds every value in registers when it rounds it to BF16. + +``SequenceParallelAttnJoin`` re-classes a layer's ``add_attn`` so that, while a +sequence-parallel forward runs (``SequenceParallelStep.active``), the join is this +module's copy of the built-in kernel with one more store: the BF16-rounded norm output +widened back to FP32, which is ``ffn_input.to(torch.float32)`` bit for bit. The pair is +handed to the layer's MoE (``moe.sp_router_input``), whose sequence-parallel arm +(``moe.ContextSequenceParallelMoE``) feeds the FP32 rows to the router GEMM when it is +called with that very norm output. Every other call is the built-in join. + +Numerics: bitwise. The kernel body is the built-in one (same block, warps, loads, +roundings and reduction), and the extra output is an exact widening of the BF16 value +the built-in stores. +""" + +from __future__ import annotations + +from typing import Optional, Tuple + +import torch +import triton +import triton.language as tl + +from tensorrt_llm._torch.models import modeling_musespark_open as _builtin +from tensorrt_llm._torch.models.modeling_musespark_open import MuseSparkOpenGatedResidual + +#: Rows at or below which the built-in join may take its cluster boundary kernel instead +#: of the Triton one; a sequence-parallel shard always has more. +_CLUSTER_BOUNDARY_MAX_ROWS = 32 + + +@triton.jit +def _join_with_router_input_kernel( + out_ptr, # [T, HIDDEN] fp32, the new residual stream + residual_ptr, # [T, HIDDEN] fp32, the old one + branch_ptr, # [T, HIDDEN] branch output, activation dtype + gain_ptr, # [HIDDEN] (1 + w) when HAS_GAIN, else unread + alpha_ptr, # [HIDDEN] fp32 + beta_ptr, # [HIDDEN] fp32 + eps, + next_out_ptr, # [T, HIDDEN] activation dtype, the next norm's output + next_weight_ptr, + next_eps, + router_ptr, # [T, HIDDEN] fp32, the next norm's output widened back from the activation dtype + HIDDEN: tl.constexpr, + BLOCK_H: tl.constexpr, + HAS_GAIN: tl.constexpr, + NORM_FP32_RESIDUAL: tl.constexpr, +): + """The built-in ``_gated_residual_after_rmsnorm_kernel`` with ``FUSE_NEXT_NORM`` set, plus one store. + + Statement for statement the built-in body up to the next norm's store; that store + rounds to the activation dtype explicitly (round to nearest even, the conversion the + built-in store applies) and the rounded value is also stored widened to FP32. + """ + token = tl.program_id(0) + offs = tl.arange(0, BLOCK_H) + mask = offs < HIDDEN + base = token * HIDDEN + offs + + branch = tl.load(branch_ptr + base, mask=mask, other=0.0) + scaled = branch.to(tl.float32) + if HAS_GAIN: + gain = tl.load(gain_ptr + offs, mask=mask, other=0.0).to(tl.float32) + scaled = (scaled * gain).to(branch.dtype).to(tl.float32) + variance = tl.sum(scaled * scaled, axis=0) / HIDDEN + normed = (scaled * tl.rsqrt(variance + eps)).to(branch.dtype).to(tl.float32) + + residual = tl.load(residual_ptr + base, mask=mask, other=0.0).to(tl.float32) + alpha = tl.load(alpha_ptr + offs, mask=mask, other=0.0) + beta = tl.load(beta_ptr + offs, mask=mask, other=0.0) + updated = alpha * residual + beta * normed + tl.store(out_ptr + base, updated, mask=mask) + if NORM_FP32_RESIDUAL: + next_input = updated + else: + next_input = updated.to(branch.dtype).to(tl.float32) + next_variance = tl.sum(next_input * next_input, axis=0) / HIDDEN + next_weight = tl.load(next_weight_ptr + offs, mask=mask, other=0.0).to(tl.float32) + next_output = next_input * tl.rsqrt(next_variance + next_eps) + next_output = next_output * (next_weight + 1.0) + rounded = next_output.to(next_out_ptr.dtype.element_ty) + tl.store(next_out_ptr + base, rounded, mask=mask) + tl.store(router_ptr + base, rounded.to(tl.float32), mask=mask) + + +def join_with_router_input( + join: MuseSparkOpenGatedResidual, + residual: torch.Tensor, + branch: torch.Tensor, + eps: float, + next_norm, + gain: Optional[torch.Tensor] = None, +) -> Optional[Tuple[torch.Tensor, torch.Tensor, torch.Tensor]]: + """``(updated, next_input, next_input.float())`` from one launch, or None. + + None wherever the built-in ``_add_after_rmsnorm_and_norm`` would not launch its + Triton kernel with this call's operands (its unfused or cluster-boundary arms); the + caller then runs the built-in join. + """ + if not ( + join._fused_enabled + and residual.is_cuda + and residual.dtype == torch.float32 + and branch.dtype == torch.bfloat16 + and residual.ndim == 2 + and residual.shape == branch.shape + and residual.shape[0] > _CLUSTER_BOUNDARY_MAX_ROWS + and residual.is_contiguous() + and branch.is_contiguous() + and next_norm.use_gemma + and next_norm.weight.dtype == branch.dtype + and next_norm.weight.is_contiguous() + and (gain is None or (gain.dtype == branch.dtype and gain.is_contiguous())) + and not torch.compiler.is_compiling() + ): + return None + alpha, beta = join.coefficients() + updated = torch.empty_like(residual) + next_input = torch.empty_like(branch) + router_input = torch.empty_like(residual) + hidden_size = residual.shape[-1] + _join_with_router_input_kernel[(residual.shape[0],)]( + updated, + residual, + branch, + gain if gain is not None else alpha, + alpha, + beta, + eps, + next_input, + next_norm.weight, + next_norm.variance_epsilon, + router_input, + HIDDEN=hidden_size, + BLOCK_H=triton.next_power_of_2(hidden_size), + HAS_GAIN=gain is not None, + NORM_FP32_RESIDUAL=join._preqkv_norm_fp32, + num_warps=_builtin._RESIDUAL_NUM_WARPS, + ) + return updated, next_input, router_input + + +class SequenceParallelStep: + """Whether a sequence-parallel forward is running (set by ``forward.ContextSequenceParallel``).""" + + def __init__(self) -> None: + self.active = False + + +class SequenceParallelAttnJoin: + """``MuseSparkOpenGatedResidual`` mixin for a layer's ``add_attn``. + + While ``_sp_step.active``, a join that the built-in would run as its Triton kernel + runs ``join_with_router_input`` and hands ``(next_input, its FP32 widening)`` to the + layer's MoE (``sp_router_input``). Any other call is the built-in join. + """ + + _sp_step: Optional[SequenceParallelStep] = None + _sp_moe = None + + def _add_after_rmsnorm_and_norm(self, residual, branch, eps, next_norm, gain=None): + step, moe = self._sp_step, self._sp_moe + if step is not None and step.active and moe is not None: + joined = join_with_router_input(self, residual, branch, eps, next_norm, gain) + if joined is not None: + updated, next_input, router_input = joined + moe.sp_router_input = (next_input, router_input) + return updated, next_input + return super()._add_after_rmsnorm_and_norm(residual, branch, eps, next_norm, gain) + + +_CLASSES = {} + + +def _with_mixin(base: type) -> type: + """``base`` with ``SequenceParallelAttnJoin`` in front, one class per base.""" + if issubclass(base, SequenceParallelAttnJoin): + return base + if base not in _CLASSES: + _CLASSES[base] = type( + "SequenceParallelAttnJoinOver" + base.__name__, + (SequenceParallelAttnJoin, base), + {"__module__": __name__}, + ) + return _CLASSES[base] + + +def install(model, step: SequenceParallelStep) -> int: + """Re-class every layer's ``add_attn`` and point it at ``step`` and the layer's MoE; returns how many. + + Raises: + ValueError: if a layer's ``add_attn`` is not a built-in ``MuseSparkOpenGatedResidual``. + """ + joins = [layer.add_attn for layer in model.layers] + wrong = [ + layer.layer_idx + for layer, join in zip(model.layers, joins) + if not isinstance(join, MuseSparkOpenGatedResidual) + ] + if wrong: + raise ValueError( + "[musespark_open] context_sequence_parallel re-classes the built-in gated residual " + f"of every layer; layers {wrong} carry another class" + ) + changed = 0 + for layer, join in zip(model.layers, joins): + mixed = _with_mixin(type(join)) + if type(join) is not mixed: + join.__class__ = mixed + changed += 1 + # Plain attributes, not registered submodules: ``nn.Module.__setattr__`` would file + # the MoE under the join's children, where attribute lookup finds the class's + # ``_sp_moe = None`` first, and the module tree would hold the MoE twice. + object.__setattr__(join, "_sp_step", step) + object.__setattr__(join, "_sp_moe", layer.feed_forward) + return changed From ec4655074925b65dbcc382d255be017c9482d605 Mon Sep 17 00:00:00 2001 From: Xiao Wang <24860335+xwang233@users.noreply.github.com> Date: Mon, 26 Oct 2026 18:44:15 +0000 Subject: [PATCH 09/16] [None][feat] modeling_v2 musespark_open tep: model_boundary fast path - Input: a decode step of at most replicated_embedding.MAX_ROWS ids runs the replicated embedding gather, embed_norm, the FP32 residual copy and layer 0's pre-QKV norm as one CuTe DSL kernel (one 128-thread CTA per row, flashinfer's RMSNorm body), and hands its two outputs to layer 0 as the residual stream and the precomputed attention input. PDL: wait before the first global read, trigger at exit. - Output: while armed for a decode forward, the last layer's post-FFN join stores BF16 directly, so the cast before the final norm is a no-op. - install checks the model it acts on (replicated table, m_emb == 1, the pre-QKV ordering, flashinfer's RMSNorm geometry, the last layer not a DFlash capture layer) and is agreed over the TP group. inputs_embeds calls, other row counts and prefill keep the built-in chain. Bitwise with the built-in chain. Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com> --- .../kernels/model_boundary/__init__.py | 3 + .../kernels/model_boundary/boundary.py | 224 +++++++++++++ .../kernels/model_boundary/input_boundary.py | 309 ++++++++++++++++++ 3 files changed, 536 insertions(+) create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/model_boundary/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/model_boundary/boundary.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/model_boundary/input_boundary.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/model_boundary/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/model_boundary/__init__.py new file mode 100644 index 000000000000..a0a97f10d1ff --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/model_boundary/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""model_boundary: the decode step's input boundary in one kernel, and a BF16 last-layer join.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/model_boundary/boundary.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/model_boundary/boundary.py new file mode 100644 index 000000000000..31e557cc3c3f --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/model_boundary/boundary.py @@ -0,0 +1,224 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Fast path ``model_boundary``: the decode step's model input and output boundaries, fewer launches. + +Input. ``DecodeTarget`` runs, before layer 0's QKV projection, the replicated embedding gather, +``embed_norm``, the FP32 residual copy, its BF16 copy and layer 0's ``pre_qkv_proj_norm``: five +launches. With this fast path a decode step of at most ``replicated_embedding.MAX_ROWS`` ids +runs ``input_boundary`` instead -- one launch, the same bits -- and hands its two outputs to +layer 0 as the residual stream and ``precomputed_attn_input``, as every later layer receives +them from the previous layer's fused boundary. + +Output. Layer 61 has no next norm, so its post-FFN join is the built-in +``add_after_rmsnorm`` (the Triton ``_gated_residual_after_rmsnorm_kernel``, FP32 out), then +``DecodeTarget`` casts the stream to BF16 for the final norm: a separate copy kernel. While +``ModelBoundary.armed`` holds, that join's kernel stores BF16 directly -- the same kernel, +its output pointer BF16, so Triton's store rounds the FP32 value to nearest even exactly as +``.to(torch.bfloat16)`` does -- and the cast is a no-op. The FP32 stream after layer 61 has +no reader: it is not a DFlash capture layer (``install`` refuses otherwise) and the final +norm reads its BF16 cast only. + +What this module does: + +* ``install`` (after ``replicated_embedding``, from ``load_weights``) checks the model it acts + on -- the replicated table, ``m_emb == 1``, the historical pre-QKV ordering, flashinfer's + RMSNorm geometry, layer 61 not captured -- builds the kernel (agreed over the TP group: a + refusal on one rank makes every rank raise) and re-classes layer 61's post-FFN join. +* ``ModelBoundary.takes`` / ``input`` serve the input boundary per call: shapes and dtypes only, + so every rank takes the same branch; ``inputs_embeds`` calls and other row counts keep the + built-in chain. +* ``ModelBoundary.armed`` arms layer 61's BF16 join for one decode forward and disarms it on + any exit; prefill never arms it. +""" + +from __future__ import annotations + +import contextlib +from typing import FrozenSet, Optional + +import torch + +from tensorrt_llm._torch.models import modeling_musespark_open as _builtin +from tensorrt_llm.logger import logger + +from ... import dispatch as _tp_agreement +from ..replicated_embedding import embedding as _replicated +from . import input_boundary as _input + + +class Bf16OutputJoin: + """Mixin for layer 61's post-FFN join: ``add_after_rmsnorm`` stores BF16 while armed. + + ``_bf16_out`` is set by ``ModelBoundary.armed`` for one decode forward. Unarmed, the + join is the class it was before ``install``. + """ + + _bf16_out: bool = False + + def add_after_rmsnorm( + self, + residual: torch.Tensor, + branch: torch.Tensor, + eps: float, + gain: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + if not self._bf16_out: + return super().add_after_rmsnorm(residual, branch, eps, gain) + num_tokens = residual.shape[0] + hidden_size = residual.shape[-1] + if not ( + self._fused_enabled + and residual.is_cuda + and num_tokens > 0 + and residual.is_contiguous() + and branch.is_contiguous() + and branch.dtype == torch.bfloat16 + and (gain is None or gain.dtype == branch.dtype) + ): + # The built-in composition, then the cast DecodeTarget would have made. + return super().add_after_rmsnorm(residual, branch, eps, gain).to(torch.bfloat16) + alpha, beta = self.coefficients() + out = torch.empty(residual.shape, dtype=torch.bfloat16, device=residual.device) + # The built-in launch (MuseSparkOpenGatedResidual.add_after_rmsnorm), with a BF16 output. + _builtin._gated_residual_after_rmsnorm_kernel[(num_tokens,)]( + out, + residual, + branch, + gain if gain is not None else alpha, + alpha, + beta, + eps, + out, + alpha, + eps, + HIDDEN=hidden_size, + BLOCK_H=_builtin.triton.next_power_of_2(hidden_size), + HAS_GAIN=gain is not None, + FUSE_NEXT_NORM=False, + NORM_FP32_RESIDUAL=False, + num_warps=_builtin._RESIDUAL_NUM_WARPS, + ) + return out + + +_JOIN_CLASSES: dict = {} + + +def _bf16_join_class(base: type) -> type: + cls = _JOIN_CLASSES.get(base) + if cls is None: + cls = type(f"ModelBoundaryBf16{base.__name__}", (Bf16OutputJoin, base), {}) + _JOIN_CLASSES[base] = cls + return cls + + +class ModelBoundary: + """What ``DecodeTarget`` holds under ``model_boundary`` (``core._model_boundary``).""" + + def __init__(self, model) -> None: + self.table = model.embed_tokens.__dict__[_replicated.TABLE_ATTR] + self.embed_norm = model.embed_norm + self.qkv_norm = model.layers[0].pre_qkv_proj_norm + self.last_join = model.layers[-1].add_ffn + + def takes( + self, input_ids: Optional[torch.Tensor], inputs_embeds: Optional[torch.Tensor] + ) -> bool: + """Whether this call's input boundary runs fused (shapes and dtypes only).""" + return ( + inputs_embeds is None + and input_ids is not None + and input_ids.dim() == 1 + and 1 <= input_ids.shape[0] <= _replicated.MAX_ROWS + and _input.takes(self.table, input_ids, self.embed_norm.weight, self.qkv_norm.weight) + ) + + def input(self, input_ids: torch.Tensor): + """``(FP32 residual, layer 0's pre-QKV norm output)`` of the step, one launch.""" + return _input.input_boundary( + self.table, + input_ids, + self.embed_norm.weight, + self.embed_norm.variance_epsilon, + self.qkv_norm.weight, + self.qkv_norm.variance_epsilon, + ) + + @contextlib.contextmanager + def armed(self): + """Layer 61's join stores BF16 for one decode forward; disarmed on any exit.""" + self.last_join._bf16_out = True + try: + yield + finally: + self.last_join._bf16_out = False + + +def _check(model, captured: FrozenSet[int]) -> None: + """``install``'s checks. Raises ValueError naming the first that fails.""" + reason = _input.available() + if reason is not None: + raise ValueError(f"[musespark_open] model_boundary: {reason}") + if _replicated.TABLE_ATTR not in model.embed_tokens.__dict__: + raise ValueError( + "[musespark_open] model_boundary gathers from the replicated table; replicated_embedding " + "is not installed" + ) + if model.m_emb != 1.0: + raise ValueError( + f"[musespark_open] model_boundary is built for m_emb == 1, got {model.m_emb}" + ) + layers = list(model.layers) + if layers[0].add_attn._preqkv_norm_fp32: + raise ValueError( + "[musespark_open] model_boundary computes layer 0's pre-QKV norm on the BF16 residual; " + "MUSESPARK_OPEN_PREQKV_NORM_FP32_RESIDUAL selects the FP32 ordering" + ) + for name, norm in ( + ("embed_norm", model.embed_norm), + ("pre_qkv_proj_norm", layers[0].pre_qkv_proj_norm), + ): + if norm.use_cuda_tile or norm.is_nvfp4 or norm.weight.dtype != torch.bfloat16: + raise ValueError( + f"[musespark_open] model_boundary: {name} is not flashinfer's BF16 RMSNorm here" + ) + if model.embed_norm.use_gemma or not layers[0].pre_qkv_proj_norm.use_gemma: + raise ValueError( + "[musespark_open] model_boundary expects a plain embed_norm and a gemma pre_qkv_proj_norm" + ) + last = len(layers) - 1 + if last in captured: + raise ValueError( + f"[musespark_open] model_boundary: layer {last} is a DFlash capture layer; its FP32 " + "residual stream has a reader" + ) + join = layers[-1].add_ffn + if not isinstance(join, _builtin.MuseSparkOpenGatedResidual): + raise ValueError( + f"[musespark_open] model_boundary: layer {last}'s post-FFN join is {type(join).__name__}" + ) + + +def install(model, captured: FrozenSet[int], mapping=None) -> ModelBoundary: + """Build the input-boundary kernel and re-class layer 61's join (agreed over ``mapping``'s TP group).""" + reason = None + try: + _check(model, captured) + # Compile now, with the step's id dtype (int32), so a compile failure is agreed here. + _input._compile(_input.HIDDEN, _input.EXIT, "int32", torch.cuda.current_device()) + except Exception as error: # noqa: BLE001 - every failure must reach the agreement + reason = str(error) + reason = _tp_agreement.agreed_refusal(mapping, reason) + if reason is not None: + raise ValueError(reason) + join = model.layers[-1].add_ffn + join.__class__ = _bf16_join_class(type(join)) + join._bf16_out = False + boundary = ModelBoundary(model) + model._model_boundary = boundary + logger.info_once( + "[musespark_open] model_boundary: decode input boundary in one launch, layer " + f"{len(model.layers) - 1} join stores BF16", + key="musespark_open_model_boundary", + ) + return boundary diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/model_boundary/input_boundary.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/model_boundary/input_boundary.py new file mode 100644 index 000000000000..4b60dea0e363 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/model_boundary/input_boundary.py @@ -0,0 +1,309 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The decode step's input boundary as one kernel: embedding gather, ``embed_norm``, residual, layer 0's pre-QKV norm. + +On a decode step ``DecodeTarget`` runs, before layer 0's QKV projection: + +1. ``replicated_embedding_gather`` (``ops/embedding/replicated_gather.py``): one row of the + replicated table per id, -0.0 written as +0.0, a +0.0 row for an id outside the table; +2. ``embed_norm``: flashinfer's CuTe DSL ``RMSNormKernel`` (``weight_bias = 0``) on the BF16 rows; +3. ``m_emb`` (1.0 for this checkpoint, so no kernel) and the FP32 copy, the residual stream; +4. layer 0's ``pre_qkv_proj_norm`` on ``residual.to(bf16)`` -- the BF16 rows of step 2 again, + since they are exactly representable -- flashinfer's ``RMSNormKernel`` with + ``weight_bias = 1`` (gemma). + +That is five launches (a gather, two norms, two copies) of 8..128 CTAs each. This kernel does all +of it in one launch of one 128-thread CTA per row: it gathers the row as step 1 does, then runs +flashinfer's ``RMSNormKernel`` body twice on the same register fragment -- the same thread-value +layout (``_make_tv_layout``), the same per-thread ``TensorSSA.reduce``, the same +``row_reduce_sum_multirow`` (imported from flashinfer, not copied), the same ``rsqrt`` and the same +product order ``x * rstd * (w + weight_bias)`` -- and stores the FP32 residual between the two. +So every output is the five-launch chain's, bit for bit, for every table value the +replicated-embedding install admits. + +PDL: every thread waits (``griddepcontrol.wait``) before its first global read (the ids are the +step's input, written by the kernels before this one). The served trigger is ``"exit"``, +flashinfer's placement after the last store, so every dependent the replaced chain's last norm +had sees the same launch order. ``"entry"`` releases the dependent at once; it is correct only +when every kernel that may launch behind this one waits before reading its outputs. +""" + +from __future__ import annotations + +import functools +from typing import Optional + +import torch + +try: + import cutlass + import cutlass.cute as cute + from cutlass import BFloat16, Float32, Int32, Int64 + from flashinfer.norm.kernels.rmsnorm import RMSNormKernel + from flashinfer.norm.utils import row_reduce_sum_multirow + + HAS_CUTE_DSL = True +except ImportError: # pragma: no cover - the served image has the CuTe DSL and flashinfer + HAS_CUTE_DSL = False + +#: Width of the residual stream this kernel is compiled for. +HIDDEN = 8192 +#: Where the grid lets its dependent launch. +EXIT, ENTRY = 0, 1 +TRIGGERS = {"exit": EXIT, "entry": ENTRY} +_ID_DTYPES = {torch.int32: "int32", torch.int64: "int64"} + + +if HAS_CUTE_DSL: + + class _InputBoundary: + """One CTA per row; flashinfer ``RMSNormKernel``'s geometry at ``HIDDEN`` columns.""" + + def __init__(self, hidden: int, trigger: int, id_dtype): + # The geometry flashinfer picks for a BF16 row of this width on this device. + ref = RMSNormKernel(BFloat16, hidden, 0.0) + if ref.cluster_n != 1 or ref.rows_per_block != 1 or ref.cols_per_tile != hidden: + raise ValueError( + f"[musespark_open] model_boundary: flashinfer's RMSNormKernel at H={hidden} uses " + f"cluster_n={ref.cluster_n}, rows_per_block={ref.rows_per_block}, " + f"cols_per_tile={ref.cols_per_tile}; the input boundary is built for one CTA " + "per row whose tile is the whole row" + ) + self.hidden = hidden + self.trigger = trigger + self.id_dtype = id_dtype + self.vec_size = ref.vec_size + self.copy_bits = ref.copy_bits + self.threads_per_row = ref.threads_per_row + self.num_threads = ref.num_threads + self.rows_per_block = ref.rows_per_block + self.warps_per_row = ref.warps_per_row + self.num_vec_blocks = ref.num_vec_blocks + self.cols_per_tile = ref.cols_per_tile + + @cute.jit + def __call__(self, mT, mIds, mWe, mWq, mR, mY, eps_e: Float32, eps_q: Float32, stream): + tv_shape, tv_stride = RMSNormKernel._make_tv_layout( + self.threads_per_row, self.rows_per_block, self.vec_size, self.num_vec_blocks + ) + tv_layout = cute.make_layout(tv_shape, stride=tv_stride) + tiler_mn = (self.rows_per_block, self.cols_per_tile) + self.kernel(mT, mIds, mWe, mWq, mR, mY, eps_e, eps_q, tv_layout, tiler_mn).launch( + grid=[mIds.shape[0], 1, 1], + block=[self.num_threads, 1, 1], + smem=self.rows_per_block * self.warps_per_row * 4, + stream=stream, + use_pdl=True, + ) + + @cute.kernel + def kernel( + self, + mT: cute.Tensor, + mIds: cute.Tensor, + mWe: cute.Tensor, + mWq: cute.Tensor, + mR: cute.Tensor, + mY: cute.Tensor, + eps_e: Float32, + eps_q: Float32, + tv_layout: cute.Layout, + tiler_mn: cute.Shape, + ): + tidx, _, _ = cute.arch.thread_idx() + row, _, _ = cute.arch.block_idx() + if cutlass.const_expr(self.trigger == ENTRY): + cute.arch.griddepcontrol_launch_dependents() + # Before any global read: the ids are written by the kernels before this one. + cute.arch.griddepcontrol_wait() + + H = self.hidden + threads_per_row = tv_layout.shape[0][0] + rows_per_block = tiler_mn[0] + warps_per_row = max(threads_per_row // 32, 1) + + smem = cutlass.utils.SmemAllocator() + reduction_buffer = smem.allocate_tensor( + Float32, cute.make_layout((rows_per_block, warps_per_row)), byte_alignment=4 + ) + + # The gather's row choice (ops/embedding/replicated_gather.py): an id outside the + # table reads nothing. + token = Int64(mIds[row]) + valid = (token >= Int64(0)) & (token < Int64(mT.shape[0])) + source = Int64(0) + if valid: + source = token + + gX = cute.local_tile(mT, tiler_mn, (source, 0)) + gR = cute.local_tile(mR, tiler_mn, (row, 0)) + gY = cute.local_tile(mY, tiler_mn, (row, 0)) + we_layout = cute.prepend(mWe.layout, cute.make_layout((tiler_mn[0],), stride=(0,))) + wq_layout = cute.prepend(mWq.layout, cute.make_layout((tiler_mn[0],), stride=(0,))) + gWe = cute.local_tile(cute.make_tensor(mWe.iterator, we_layout), tiler_mn, (0, 0)) + gWq = cute.local_tile(cute.make_tensor(mWq.iterator, wq_layout), tiler_mn, (0, 0)) + + copy_atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), BFloat16, num_bits_per_copy=self.copy_bits + ) + tiled_copy = cute.make_tiled_copy(copy_atom, tv_layout, tiler_mn) + thr_copy = tiled_copy.get_slice(tidx) + + tXgX = thr_copy.partition_S(gX) + tXrX = cute.make_fragment_like(tXgX) + tWgWe = thr_copy.partition_S(gWe) + tWrWe = cute.make_fragment_like(tWgWe) + tWgWq = thr_copy.partition_S(gWq) + tWrWq = cute.make_fragment_like(tWgWq) + tXgR = thr_copy.partition_D(gR) + tXgY = thr_copy.partition_D(gY) + tXrY = cute.make_fragment_like(tXgY) + tXrR = cute.make_fragment_like(tXgY, Float32) + tXrS = cute.make_fragment_like(tXgY, Float32) + + tXrX.store(cute.zeros_like(tXrX, dtype=BFloat16)) + if valid: + cute.copy(copy_atom, tXgX, tXrX) + cute.copy(copy_atom, tWgWe, tWrWe) + cute.copy(copy_atom, tWgWq, tWrWq) + + # The gather's sign rule: a zero of either sign is +0.0. + tXrS.store(tXrX.load().to(Float32)) + for i in cutlass.range_constexpr(cute.size(tXrS)): + if tXrS[i] == Float32(0.0): + tXrS[i] = Float32(0.0) + + # embed_norm: flashinfer RMSNormKernel, weight_bias 0. + x = tXrS.load() + x_sq = x * x + sum_sq = row_reduce_sum_multirow(x_sq, threads_per_row, reduction_buffer, None, 1) + mean_sq = sum_sq / Float32(H) + rstd = cute.math.rsqrt(mean_sq + eps_e, fastmath=True) + cute.arch.barrier() + w = tWrWe.load().to(Float32) + embeds = (x * rstd * (w + Float32(0.0))).to(BFloat16) + + # The FP32 residual stream (m_emb == 1). + tXrR.store(embeds.to(Float32)) + cute.autovec_copy(tXrR, tXgR) + + # Layer 0's pre_qkv_proj_norm on residual.to(bf16) == embeds: RMSNormKernel, weight_bias 1. + x2 = embeds.to(Float32) + x2_sq = x2 * x2 + sum_sq2 = row_reduce_sum_multirow(x2_sq, threads_per_row, reduction_buffer, None, 1) + mean_sq2 = sum_sq2 / Float32(H) + rstd2 = cute.math.rsqrt(mean_sq2 + eps_q, fastmath=True) + wq = tWrWq.load().to(Float32) + tXrY.store((x2 * rstd2 * (wq + Float32(1.0))).to(BFloat16)) + cute.copy(copy_atom, tXrY, tXgY) + + if cutlass.const_expr(self.trigger == EXIT): + cute.arch.griddepcontrol_launch_dependents() + + @functools.lru_cache(maxsize=None) + def _compile(hidden: int, trigger: int, id_dtype: str, device_index: int): + ids_type = Int32 if id_dtype == "int32" else Int64 + rows = cute.sym_int64() + vocab = cute.sym_int64() + table = cute.runtime.make_fake_compact_tensor( + BFloat16, (vocab, hidden), stride_order=(1, 0), assumed_align=16 + ) + ids = cute.runtime.make_fake_compact_tensor(ids_type, (rows,), assumed_align=4) + weight_e = cute.runtime.make_fake_compact_tensor(BFloat16, (hidden,), assumed_align=16) + weight_q = cute.runtime.make_fake_compact_tensor(BFloat16, (hidden,), assumed_align=16) + residual = cute.runtime.make_fake_compact_tensor( + Float32, (rows, hidden), stride_order=(1, 0), assumed_align=16 + ) + out = cute.runtime.make_fake_compact_tensor( + BFloat16, (rows, hidden), stride_order=(1, 0), assumed_align=16 + ) + stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) + with torch.cuda.device(device_index): + return cute.compile( + _InputBoundary(hidden, trigger, ids_type), + table, + ids, + weight_e, + weight_q, + residual, + out, + Float32(1e-6), + Float32(1e-6), + stream, + options="--enable-tvm-ffi", + ) + + +def available() -> Optional[str]: + """Why this kernel cannot run here, or ``None``.""" + return ( + None if HAS_CUTE_DSL else "the CuTe DSL or flashinfer's CuTe DSL RMSNorm is not importable" + ) + + +def takes( + table: torch.Tensor, ids: torch.Tensor, embed_weight: torch.Tensor, qkv_weight: torch.Tensor +) -> bool: + """Whether a call fits the kernel (shapes, dtypes, contiguity; no pointer values).""" + return ( + table.is_cuda + and table.dtype == torch.bfloat16 + and table.dim() == 2 + and table.shape[1] == HIDDEN + and table.is_contiguous() + and ids.device == table.device + and ids.dtype in _ID_DTYPES + and ids.dim() == 1 + and ids.shape[0] > 0 + and ids.is_contiguous() + and all( + w.device == table.device + and w.dtype == torch.bfloat16 + and tuple(w.shape) == (HIDDEN,) + and w.is_contiguous() + for w in (embed_weight, qkv_weight) + ) + ) + + +def input_boundary( + table: torch.Tensor, + ids: torch.Tensor, + embed_weight: torch.Tensor, + embed_eps: float, + qkv_weight: torch.Tensor, + qkv_eps: float, + trigger: str = "exit", +) -> tuple: + """``(residual FP32 [rows, H], layer 0's pre-QKV norm output BF16 [rows, H])`` in one launch. + + ``table`` is the whole replicated ``[vocab, H]`` BF16 embedding, ``ids`` the step's token + ids, ``embed_weight`` ``embed_norm``'s (unit) weight and ``qkv_weight`` layer 0's gemma + ``pre_qkv_proj_norm`` weight. The caller checks ``takes`` first. + """ + rows = ids.shape[0] + residual = torch.empty((rows, HIDDEN), dtype=torch.float32, device=table.device) + out = torch.empty((rows, HIDDEN), dtype=torch.bfloat16, device=table.device) + _compile(HIDDEN, TRIGGERS[trigger], _ID_DTYPES[ids.dtype], table.device.index)( + table, ids, embed_weight, qkv_weight, residual, out, float(embed_eps), float(qkv_eps) + ) + return residual, out + + +def reference( + table: torch.Tensor, + ids: torch.Tensor, + embed_weight: torch.Tensor, + embed_eps: float, + qkv_weight: torch.Tensor, + qkv_eps: float, +) -> tuple: + """The five-launch chain this kernel replaces, through the served ops.""" + from tensorrt_llm._torch.custom_ops import flashinfer_gemma_rmsnorm, flashinfer_rmsnorm + + from ...ops.embedding.replicated_gather import replicated_embedding_gather + + embeds = replicated_embedding_gather(table, ids) + embeds = flashinfer_rmsnorm(embeds, embed_weight, embed_eps) + residual = embeds.to(torch.float32) + return residual, flashinfer_gemma_rmsnorm(residual.to(torch.bfloat16), qkv_weight, qkv_eps) From 542490218093ea3777b1ab6dbf111252975aeaf7 Mon Sep 17 00:00:00 2001 From: Xiao Wang <24860335+xwang233@users.noreply.github.com> Date: Mon, 26 Oct 2026 18:44:16 +0000 Subject: [PATCH 10/16] [None][feat] modeling_v2 musespark_open tep: vocab-parallel verify head - vocab_parallel_head (GB300, sm103 TEP4): decode steps of VP_MIN_ROWS (two K=7 verify blocks) up to the head's widest rows sample and argmax on each rank's own vocabulary shard; the ranks exchange per-thread states, greedy words and candidates over the tray's multicast and agree on the token in the head_sampler JIT CUDA kernels, instead of all-gathering the logits. Greedy rows are the argmax of the soft-capped row exactly; sampled rows keep the shipped CDF order; filtered rows run the shipped body on the row gathered from the ranks. - target_head.install(llm, mapping, gather, vocab_parallel=False) builds the exchange buffers (rebuilt by widen for the wider steps); 8-row steps keep the lm_head_inbox gather. - Split-sampler routing for the vocab-parallel pending rows. Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com> --- .../kernels/split_sampler/target_head.py | 116 ++- .../ops/spec/_head_sampler/head_sampler.cu | 772 ++++++++++++++++++ .../ops/spec/head_sampler.py | 465 +++++++++++ 3 files changed, 1346 insertions(+), 7 deletions(-) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/split_sampler/target_head.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/split_sampler/target_head.py index 95be6fda9d1c..db66937ca604 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/split_sampler/target_head.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/split_sampler/target_head.py @@ -16,6 +16,14 @@ ``[8, shard]`` output straight into this rank's slot of the multicast inbox (``HeadGatherWorkspace``) and the sampler's first kernel gathers the row from the inbox (``gather_split_sample``), instead of the NCCL all-gather. +* ``vocab_parallel_head`` -- on decode steps of ``head_sampler.VP_MIN_ROWS`` (16) rows or more, + no gather at all: the processor runs only the local LM-head GEMM and leaves a + ``[rows, vocab]`` BF16 placeholder pending (allocated, not written); the worker samples + the step from the ranks' shards (``head_sampler.vocab_parallel_sample``: each rank folds + its own shard, the ranks exchange per-virtual-thread partial states and crossings over a + multicast buffer) and argmaxes greedy batches the same way + (``vocab_parallel_argmax``, ``torch.argmax``'s token exactly). Every other consumer + gathers the shards with the built-in all-gather first (``materialize``). The processor instance is re-classed as ``SplitSamplerLogitsProcessor`` (only the target's: the drafter has its own processor), and the worker as a class @@ -191,6 +199,11 @@ def context_kv_head_fork(lm_head) -> tuple: class TargetHead: """The state the re-classed processor and worker share: the step's pending BF16 row.""" + #: ``vocab_parallel_head``'s fabric buffers (None: the head gathers every row). + vp: Optional[_hs.VocabParallelWorkspace] = None + #: This rank's ``[rows, shard]`` LM-head output while a vocab-parallel row is pending. + pending_local: Optional[torch.Tensor] = None + def __init__( self, processor: MuseSparkOpenLogitsProcessor, @@ -198,6 +211,7 @@ def __init__( worker: DFlashWorker, mapping, gather: bool, + vocab_parallel: bool = False, ) -> None: self.multiplier = float(processor.logit_multiplier) self.soft_cap = float(processor.soft_cap) @@ -219,9 +233,19 @@ def __init__( self.max_rows = _hs.MAX_ROWS self.states = _hs.new_states(self.max_rows, self.device) self.wide_states: Optional[torch.Tensor] = None + # vocab_parallel_head: the fabric buffers for steps of VP_MIN_ROWS..max_rows rows + # (collective; rebuilt for the wider steps by ``widen``). + self.mapping = mapping + self.lm_head = lm_head + self.vp = None + if vocab_parallel: + self.vp = _hs.VocabParallelWorkspace( + mapping, int(lm_head.weight.shape[0]), self.max_rows + ) self.load_seconds = time.perf_counter() - start self.pending: Optional[torch.Tensor] = None self.pending_inbox = False + self.pending_local = None self._bypass_logged = False self._wide_capture_logged = False self.hook_handle = None @@ -298,6 +322,17 @@ def produce(self, hidden_states: torch.Tensor, lm_head: Linear) -> torch.Tensor: "drained within its step" ) rows = hidden_states.shape[0] + if self.vp is not None and _hs.VP_MIN_ROWS <= rows <= self.vp.max_rows: + # The built-in Linear.forward's local GEMM (column-parallel: apply_linear, then the + # all-gather this path leaves out), through the instance's apply_linear. + local = lm_head.apply_linear(hidden_states, lm_head.bias, None, None) + logits = torch.empty( + (rows, self.vp.vocab), dtype=torch.bfloat16, device=hidden_states.device + ) + self.pending_local = local + self.pending_inbox = False + self.pending = logits + return logits if self.workspace is not None and rows == GATHER_ROWS: admitted, fork = context_kv_head_fork(lm_head) if not admitted: @@ -347,9 +382,26 @@ def sample( offset: Optional[torch.Tensor], ) -> torch.Tensor: """The fused sampler's tokens for the pending row; consumes it.""" + local = self.pending_local inbox = self.pending_inbox self.pending = None self.pending_inbox = False + self.pending_local = None + if local is not None: + return _hs.vocab_parallel_sample( + self.vp.view, + logits, + local, + temperatures, + top_ks, + top_ps, + min_ps, + seed, + offset, + self.multiplier, + self.soft_cap, + self.pdl, + ) if inbox: ws = self.workspace return _hs.gather_split_sample( @@ -384,7 +436,29 @@ def sample( self.pdl, ) - def _gather_now(self, logits: torch.Tensor) -> None: + def vocab_parallel(self, logits) -> bool: + """Whether ``logits`` is the pending row of a vocab-parallel step.""" + return self.owns(logits) and self.pending_local is not None + + def argmax(self, logits: torch.Tensor) -> torch.Tensor: + """``torch.argmax`` of the built-in capped row of a vocab-parallel pending row (int64); consumes it.""" + local = self.pending_local + self.pending = None + self.pending_local = None + return _hs.vocab_parallel_argmax( + self.vp.view, local, self.multiplier, self.soft_cap, self.pdl + ) + + def _gather_now(self, logits: torch.Tensor) -> Optional[torch.Tensor]: + """Consume the pending row; returns the gathered BF16 row of a vocab-parallel one.""" + if self.pending_local is not None: + from tensorrt_llm._torch.distributed import allgather + + local = self.pending_local + self.pending = None + self.pending_local = None + # The built-in Linear.forward's gather of the column-parallel output. + return allgather(local, self.lm_head.mapping, sizes=self.lm_head.gather_output_sizes) if self.pending_inbox: ws = self.workspace empty = ws.status # unread by the gather-only kernel @@ -408,15 +482,22 @@ def _gather_now(self, logits: torch.Tensor) -> None: ) self.pending = None self.pending_inbox = False + return None def materialize(self, logits: torch.Tensor) -> torch.Tensor: """The built-in processor's FP32 capped logits for the pending row; consumes it.""" - self._gather_now(logits) - return _hs.reference_capped(logits, self.multiplier, self.soft_cap) + gathered = self._gather_now(logits) + return _hs.reference_capped( + logits if gathered is None else gathered, self.multiplier, self.soft_cap + ) def drain(self) -> None: - """Consume a pending row nothing sampled (resets its inbox words).""" - if self.pending is not None: + """Consume a pending row nothing sampled (resets its inbox words; a vocab-parallel row + has none, and is dropped without a gather).""" + if self.pending_local is not None: + self.pending = None + self.pending_local = None + elif self.pending is not None: self._gather_now(self.pending) @@ -475,6 +556,13 @@ def _sample_tokens_for_batch(self, logits, spec_metadata, num_contexts, batch_si head = self._split_sampler_head head.check_routed(logits, "_sample_tokens_for_batch") if head.owns(logits): + if ( + spec_metadata is not None + and spec_metadata.is_all_greedy_sample + and head.vocab_parallel(logits) + ): + # The built-in greedy branch: torch.argmax of the capped row. + return head.argmax(logits) if ( spec_metadata is not None and not spec_metadata.is_all_greedy_sample @@ -635,6 +723,13 @@ def widen(llm) -> bool: if head.pending is not None: raise RuntimeError("[musespark_open] wide_split_sampler: a head row is pending") head.max_rows = _hs.WIDE_MAX_ROWS + if head.vp is not None: + # Collective: every rank widens at the same point of the load. + head.vp = _hs.VocabParallelWorkspace(head.mapping, head.vp.layout.shard, head.max_rows) + logger.info( + f"[musespark_open] vocab_parallel_head: decode steps of {_hs.VP_MIN_ROWS}..{head.max_rows} " + f"rows; {head.vp.nbytes() / 2**20:.1f} MiB of fabric buffers on this rank" + ) logger.info( f"[musespark_open] wide_split_sampler: the split sampler takes decode steps of up to " f"{head.max_rows} rows (the shipped 512-thread block above {_hs.MAX_ROWS}); its state " @@ -643,7 +738,7 @@ def widen(llm) -> bool: return True -def install(llm, mapping, gather: bool) -> None: +def install(llm, mapping, gather: bool, vocab_parallel: bool = False) -> None: """Install the split sampler on the target's processor and DFlash worker. Idempotent. With ``gather`` (``lm_head_inbox``) the LM-head GEMM also writes into the multicast @@ -671,7 +766,7 @@ def install(llm, mapping, gather: bool) -> None: if reason is not None: raise ValueError(reason) processor, worker, lm_head = llm.logits_processor, llm.spec_worker, llm.lm_head - head = TargetHead(processor, lm_head, worker, mapping, gather) + head = TargetHead(processor, lm_head, worker, mapping, gather, vocab_parallel) processor.__class__ = SplitSamplerLogitsProcessor processor._split_head = head composed = attach_worker(worker, head) @@ -683,6 +778,13 @@ def install(llm, mapping, gather: bool) -> None: if gather else "" ) + + ( + f"; vocab_parallel_head: decode steps of {_hs.VP_MIN_ROWS}..{head.max_rows} rows sample " + f"each rank's own shard (no logits all-gather), {head.vp.nbytes() / 2**20:.1f} MiB of " + "fabric buffers on this rank" + if head.vp is not None + else "" + ) + f"; library {head.library} (load {head.load_seconds:.2f} s); " + f"worker class {composed.__name__}" ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/_head_sampler/head_sampler.cu b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/_head_sampler/head_sampler.cu index 3b251d689eab..b7aed7e11434 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/_head_sampler/head_sampler.cu +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/_head_sampler/head_sampler.cu @@ -60,6 +60,7 @@ #include #include +#include #include #include @@ -578,6 +579,569 @@ __global__ void splitWatchKernel(uint32_t const* buf, uint32_t* firstSeen, int64 } } +// ---------------------------------------------------------------------------------------- +// The vocabulary-parallel head (vp_*): each rank folds only its own LM-head shard. +// +// Virtual thread vt of the shipped VB-thread block owns the vectors v = vt (mod VB) of the +// whole vocabulary; rank r holds the vectors [r * shardVecs, (r + 1) * shardVecs). vp_fold +// folds, per (row, vt), the vectors of vt that lie in this rank's shard, in increasing +// order, with the shipped per-element statements, and publishes that partial state to +// every rank. vp_draw combines the R partials of its vt in rank order (the vocabulary +// order of vt's vectors), then runs the shipped first round (block reduce, scan, Philox +// draw, claimants); every rank walks its own portion of each claimant's elements, from +// the claimant's base plus the masses of the earlier ranks' portions, and publishes the +// crossing it finds (or none). vp_decide takes, per claimant, the first rank in rank order +// with a crossing, and kernel B's rule across claimants. Greedy batches: vp_fold reduces +// torch.argmax's comparator over each CTA and vp_argmax over the ranks' CTAs. +// +// Exchange buffer (one copy per rank, every 32-bit word holds kSplitSentinel between +// calls): states [ranks][stateStride] x 16 B, greedy words [ranks][greedyStride] x 8 B at +// greedyOffset, candidates [ranks][candStride] x 4 B at candOffset. Each value is stored +// through every publish target (the multicast address in service; every rank's copy in +// the single-GPU emulation); every rank reads and resets its own copy. + +//! Ranks the vocabulary-parallel head takes. +constexpr int kVpMaxRanks = 8; +//! Claimants of a row whose crossings the head exchanges (in thread order). +constexpr int kVpClaimants = 4; +//! A candidate word: no crossing in this rank's portion. +constexpr uint32_t kVpNone = 0x7FFFFFFEu; +//! Most vectors of one virtual thread in one shard (the walk's shared-memory stage). +constexpr int kVpWalkVectors = 64; +//! Threads per CTA of vp_fold. +constexpr int kVpFoldThreads = 128; +//! Greedy words per (rank, row): one per vp_fold CTA, at most kSplitBlock / kVpFoldThreads. +constexpr int kVpGreedySlots = kSplitBlock / kVpFoldThreads; +//! vp_draw's per-row record for vp_decide: [claimants or -1 (filtered), argmax, claimant tids]. +constexpr int kVpInfoWords = 2 + kVpClaimants; + +struct VpParams +{ + //! Row parameters; vocabSize is the whole vocabulary; logits is the pending + //! [numRows, vocabSize] BF16 row (written by vp_draw for filtered rows only). + FusedSamplingParams fs; + //! This rank's [numRows][shardVecs] local LM-head output. + uint2 const* local{nullptr}; + //! This rank's staging rows [rows][shardVecs] (fabric memory; filtered rows only). + uint2* staging{nullptr}; + //! Every rank's staging buffer, in rank order (peer addresses). + uint2 const* stagingOf[kVpMaxRanks]{}; + //! Exchange-buffer bases the values are stored through. + char* publish[kVpMaxRanks]{}; + int numPublish{0}; + //! This rank's copy of the exchange buffer. + char* copy{nullptr}; + int* info{nullptr}; + int* status{nullptr}; + int64_t* greedyTokens{nullptr}; + SplitCap cap{}; + int rank{0}; + int ranks{0}; + int shardVecs{0}; + long long stateStride{0}; + long long greedyOffset{0}; + long long greedyStride{0}; + long long candOffset{0}; + long long candStride{0}; + unsigned long long timeoutNs{0}; +}; + +__device__ __forceinline__ uint4* vpState(char* base, VpParams const& p, int rank, int row, int vt, int vb) +{ + return reinterpret_cast(base) + rank * p.stateStride + static_cast(row) * vb + vt; +} + +__device__ __forceinline__ uint2* vpGreedy(char* base, VpParams const& p, int rank, int row, int slot) +{ + return reinterpret_cast(base + p.greedyOffset) + rank * p.greedyStride + + static_cast(row) * kVpGreedySlots + slot; +} + +__device__ __forceinline__ uint32_t* vpCand(char* base, VpParams const& p, int rank, int row, int c) +{ + return reinterpret_cast(base + p.candOffset) + rank * p.candStride + + static_cast(row) * kVpClaimants + c; +} + +__device__ __forceinline__ uint4 loadRelaxedSys4(uint4 const* address) +{ + uint4 v; + asm volatile("ld.relaxed.sys.global.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w) + : "l"(address) + : "memory"); + return v; +} + +__device__ __forceinline__ void storeRelaxedSys4(uint4* address, uint4 v) +{ + asm volatile( + "st.relaxed.sys.global.v4.b32 [%0], {%1, %2, %3, %4};" ::"l"(address), "r"(v.x), "r"(v.y), "r"(v.z), "r"(v.w) + : "memory"); +} + +__device__ __forceinline__ uint32_t loadRelaxedSys1(uint32_t const* address) +{ + uint32_t v; + asm volatile("ld.relaxed.sys.global.b32 %0, [%1];" : "=r"(v) : "l"(address) : "memory"); + return v; +} + +__device__ __forceinline__ void storeRelaxedSys1(uint32_t* address, uint32_t v) +{ + asm volatile("st.relaxed.sys.global.b32 [%0], %1;" ::"l"(address), "r"(v) : "memory"); +} + +//! Records a poll that waited past the timeout and traps (a peer that never published). +__device__ __noinline__ void vpPollTimedOut(int* status) +{ + atomicOr(status, kSplitPollTimeout); + __threadfence_system(); + __trap(); +} + +//! Spin until none of the first `words` 32-bit words at `address` is the sentinel. +template +__device__ __forceinline__ uint4 vpAwait(void const* address, int* status, unsigned long long timeoutNs) +{ + unsigned long long start = 0; + for (int spins = 0;; ++spins) + { + uint4 v{kSplitSentinel, kSplitSentinel, kSplitSentinel, kSplitSentinel}; + if constexpr (WORDS == 1) + { + v.x = loadRelaxedSys1(static_cast(address)); + } + else if constexpr (WORDS == 2) + { + uint2 const w = loadRelaxedSys(static_cast(address)); + v.x = w.x; + v.y = w.y; + } + else + { + v = loadRelaxedSys4(static_cast(address)); + } + bool ready = v.x != kSplitSentinel; + if constexpr (WORDS >= 2) + { + ready = ready && v.y != kSplitSentinel; + } + if constexpr (WORDS >= 3) + { + ready = ready && v.z != kSplitSentinel; + } + if (ready) + { + return v; + } + if (spins == 64) + { + start = globalTimerNs(); + } + if (spins >= 64) + { + __nanosleep(64); + if (globalTimerNs() - start > timeoutNs) + { + vpPollTimedOut(status); + } + } + } +} + +//! torch.argmax's comparator (GreaterOrNan): whether (a, ia) is chosen over (b, ib). +__device__ __forceinline__ bool vpArgmaxBetter(float a, int ia, float b, int ib) +{ + if (isnan(a)) + { + return isnan(b) ? ia < ib : true; + } + return a == b ? ia < ib : a > b; +} + +struct VpMaxArg +{ + float value; + int index; +}; + +struct VpMaxArgOp +{ + __device__ __forceinline__ VpMaxArg operator()(VpMaxArg const& a, VpMaxArg const& b) const + { + return vpArgmaxBetter(a.value, a.index, b.value, b.index) ? a : b; + } +}; + +//! The first vector of virtual thread vt in the shard starting at global vector `first`, as +//! an offset into the shard. +__device__ __forceinline__ int vpLead(int vt, int first, int vb) +{ + return ((vt - first) % vb + vb) % vb; +} + +//! vp_fold: grid (numRows, VB / kVpFoldThreads), kVpFoldThreads threads. PDL: waits on +//! the LM-head GEMM before reading anything; releases its dependent after publishing. +template +__global__ void __launch_bounds__(kVpFoldThreads) vpFoldKernel(VpParams p) +{ + gridDependencyWait(); + int const row = blockIdx.x; + int const vt = blockIdx.y * kVpFoldThreads + threadIdx.x; + int const sv = p.shardVecs; + int const first = p.rank * sv; + int const lead = vpLead(vt, first, VB); + uint2 const* localRow = p.local + static_cast(row) * sv; + float tempInv = 1.0f; + bool stage = false; + if constexpr (!GREEDY) + { + RowParams const rp = loadRowParams(p.fs, row); + tempInv = rp.tempInv; + stage = rp.needTopK || rp.needTopP || rp.needMinP; + } + uint2* stagingRow = p.staging + static_cast(row) * sv; + OnlineSoftmax st{-FLT_MAX, 0.0f, 0}; + VpMaxArg best{-INFINITY, INT_MAX}; + for (int l0 = lead; l0 < sv; l0 += kSplitPrefetch * VB) + { + uint2 buf[kSplitPrefetch]; +#pragma unroll + for (int u = 0; u < kSplitPrefetch; ++u) + { + int const l = l0 + u * VB; + if (l < sv) + { + buf[u] = __ldg(localRow + l); + } + } +#pragma unroll + for (int u = 0; u < kSplitPrefetch; ++u) + { + int const l = l0 + u * VB; + if (l < sv) + { + if (stage) + { + stagingRow[l] = buf[u]; + } + __nv_bfloat16 const* elems = reinterpret_cast<__nv_bfloat16 const*>(&buf[u]); +#pragma unroll + for (int j = 0; j < 4; ++j) + { + float const logit = capOf(__bfloat162float(elems[j]), p.cap); + int const index = (first + l) * 4 + j; + if constexpr (GREEDY) + { + if (vpArgmaxBetter(logit, index, best.value, best.index)) + { + best = VpMaxArg{logit, index}; + } + } + else + { + st = combineOnlineSoftmax(st, OnlineSoftmax{__fmul_rn(logit, tempInv), 1.0f, index}); + } + } + } + } + } + if constexpr (GREEDY) + { + using BlockReduceMaxArg = cub::BlockReduce; + __shared__ typename BlockReduceMaxArg::TempStorage temp; + VpMaxArg const cta = BlockReduceMaxArg(temp).Reduce(best, VpMaxArgOp()); + if (threadIdx.x == 0) + { + uint2 const word = make_uint2(__float_as_uint(cta.value), static_cast(cta.index)); + for (int i = 0; i < p.numPublish; ++i) + { + *vpGreedy(p.publish[i], p, p.rank, row, blockIdx.y) = word; + } + } + } + else + { + if (stage) + { + // The staged row reaches the peers before this thread's state does. + __threadfence_system(); + } + uint4 const word + = make_uint4(__float_as_uint(st.max), __float_as_uint(st.sum), static_cast(st.argmax), 0u); + for (int i = 0; i < p.numPublish; ++i) + { + *vpState(p.publish[i], p, p.rank, row, vt, VB) = word; + } + } + gridDependencyTrigger(); +} + +//! vp_draw: one VB-thread CTA per row. PDL: waits on vp_fold before reading anything. +template +__global__ void __launch_bounds__(VB, VB == kSplitNarrowBlock ? 2 : 1) vpDrawKernel(VpParams p) +{ + __shared__ FusedSamplingShared shared; + __shared__ typename cub::BlockScan::TempStorage scanClaims; + __shared__ __align__(16) float wbuf[4 * kVpWalkVectors]; + __shared__ int sClaimTid[kVpClaimants]; + __shared__ float sClaimStart[kVpClaimants]; + __shared__ uint32_t sCand[kVpClaimants]; + int const row = blockIdx.x; + int const tid = threadIdx.x; + if (tid == 0) + { + sSplitCap = p.cap; + } + __syncthreads(); + gridDependencyWait(); + RowParams const rp = loadRowParams(p.fs, row); + bool const filtered = rp.needTopK || rp.needTopP || rp.needMinP; + // This thread's R partial states, combined in rank order. They stay in the copy until + // the end of the call (a claimant reads the earlier ranks' again for its walk offset). + OnlineSoftmax st{-FLT_MAX, 0.0f, 0}; + for (int r = 0; r < p.ranks; ++r) + { + uint4 const w = vpAwait<3>(vpState(p.copy, p, r, row, tid, VB), p.status, p.timeoutNs); + OnlineSoftmax const part{__uint_as_float(w.x), __uint_as_float(w.y), static_cast(w.z)}; + st = r == 0 ? part : combineOnlineSoftmax(st, part); + } + uint4 const sentinel4 = make_uint4(kSplitSentinel, kSplitSentinel, kSplitSentinel, kSplitSentinel); + int const sv = p.shardVecs; + int const vecCount = p.fs.vocabSize / 4; + int* const info = p.info + static_cast(row) * kVpInfoWords; + if (filtered) + { + // Every rank staged this row before publishing its states (vp_fold); gather it + // into the pending row and run the shipped body on it, as kernel B does. + __threadfence_system(); + __syncthreads(); + uint2* const pendingRow + = reinterpret_cast(const_cast(p.fs.logits)) + static_cast(row) * vecCount; + for (int v = tid; v < vecCount; v += VB) + { + int const r = v / sv; + pendingRow[v] = loadRelaxedSys(p.stagingOf[r] + static_cast(row) * sv + (v - r * sv)); + } + __syncthreads(); + fusedSamplingBody(p.fs, shared); + for (int r = 0; r < p.ranks; ++r) + { + storeRelaxedSys4(vpState(p.copy, p, r, row, tid, VB), sentinel4); + } + if (tid < kVpClaimants) + { + for (int i = 0; i < p.numPublish; ++i) + { + *vpCand(p.publish[i], p, p.rank, row, tid) = kVpNone; + } + } + if (tid == 0) + { + info[0] = -1; + } + gridDependencyTrigger(); + return; + } + using BlockReduceOnline = cub::BlockReduce; + using BlockScanF = cub::BlockScan; + OnlineSoftmax const rowStats = BlockReduceOnline(shared.temp.reduceOnline).Reduce(st, OnlineSoftmaxOp()); + if (tid == 0) + { + shared.maxScaled = rowStats.max; + shared.argmax = rowStats.argmax; + } + __syncthreads(); + float const maxScaled = shared.maxScaled; + float const localMass = st.sum * __expf(st.max - maxScaled); + if (tid == 0) + { + initRowRng(p.fs, row, &shared.rejectRng); + } + __syncthreads(); + float base = 0.0f; + float supportMass = 0.0f; + BlockScanF(shared.temp.scanF).ExclusiveSum(localMass, base, supportMass); + __syncthreads(); + if (tid == 0) + { + shared.target = curand_uniform(&shared.rejectRng) * supportMass; + } + __syncthreads(); + float const target = shared.target; + bool const claims = target >= base && target < base + localMass; + int claimIdx = 0; + int claimants = 0; + cub::BlockScan(scanClaims).ExclusiveSum(claims ? 1 : 0, claimIdx, claimants); + if (claims && claimIdx < kVpClaimants) + { + // The masses of this thread's portions on the ranks before this one. + float earlier = 0.0f; + for (int r = 0; r < p.rank; ++r) + { + uint4 const w = loadRelaxedSys4(vpState(p.copy, p, r, row, tid, VB)); + earlier += __uint_as_float(w.y) * __expf(__uint_as_float(w.x) - maxScaled); + } + sClaimTid[claimIdx] = tid; + sClaimStart[claimIdx] = base + earlier; + } + for (int r = 0; r < p.ranks; ++r) + { + storeRelaxedSys4(vpState(p.copy, p, r, row, tid, VB), sentinel4); + } + __syncthreads(); + int const first = p.rank * sv; + uint2 const* localRow = p.local + static_cast(row) * sv; + int const walkers = min(claimants, kVpClaimants); + for (int c = 0; c < walkers; ++c) + { + int const walker = sClaimTid[c]; + int const lead = vpLead(walker, first, VB); + int const walkVectors = lead < sv ? (sv - lead + VB - 1) / VB : 0; + if (tid < walkVectors) + { + uint2 const packed = __ldg(localRow + lead + tid * VB); + __nv_bfloat16 const* elems = reinterpret_cast<__nv_bfloat16 const*>(&packed); +#pragma unroll + for (int j = 0; j < 4; ++j) + { + wbuf[tid * 4 + j] = __expf(__fmul_rn(capOf(__bfloat162float(elems[j]), p.cap), rp.tempInv) - maxScaled); + } + } + __syncthreads(); + if (tid == 0) + { + float running = sClaimStart[c]; + uint32_t cand = kVpNone; + float4 const* w4 = reinterpret_cast(wbuf); + for (int q = 0; q < walkVectors && cand == kVpNone; ++q) + { + float4 const quad = w4[q]; + float const ws[4] = {quad.x, quad.y, quad.z, quad.w}; +#pragma unroll + for (int j = 0; j < 4; ++j) + { + float const w = ws[j]; + if (cand == kVpNone && w >= 0.0f && w > -1.0f) + { + running += w; + if (running > target) + { + cand = static_cast((first + lead + q * VB) * 4 + j); + } + } + } + } + sCand[c] = cand; + } + __syncthreads(); + } + if (tid < kVpClaimants) + { + uint32_t const cand = tid < walkers ? sCand[tid] : kVpNone; + for (int i = 0; i < p.numPublish; ++i) + { + *vpCand(p.publish[i], p, p.rank, row, tid) = cand; + } + } + if (tid == 0) + { + info[0] = claimants; + info[1] = shared.argmax; + for (int c = 0; c < kVpClaimants; ++c) + { + info[2 + c] = c < walkers ? sClaimTid[c] : -1; + } + } + gridDependencyTrigger(); +} + +//! vp_decide: one warp per row. PDL: waits on vp_draw before reading anything. +__global__ void __launch_bounds__(32) vpDecideKernel(VpParams p, int vb) +{ + gridDependencyWait(); + int const row = blockIdx.x; + int const lane = threadIdx.x; + uint32_t mine = kVpNone; + if (lane < p.ranks * kVpClaimants) + { + uint32_t* const slot = vpCand(p.copy, p, lane / kVpClaimants, row, lane % kVpClaimants); + mine = vpAwait<1>(slot, p.status, p.timeoutNs).x; + storeRelaxedSys1(slot, kSplitSentinel); + } + uint32_t words[kVpMaxRanks * kVpClaimants]; +#pragma unroll + for (int i = 0; i < kVpMaxRanks * kVpClaimants; ++i) + { + words[i] = __shfl_sync(0xffffffffu, mine, i); + } + if (lane == 0) + { + int const* info = p.info + static_cast(row) * kVpInfoWords; + int const claimants = info[0]; + if (claimants >= 0) + { + int token = info[1]; + unsigned long long bestKey = 0ull; + int const walkers = min(claimants, kVpClaimants); + for (int c = 0; c < walkers; ++c) + { + uint32_t cand = kVpNone; + for (int r = 0; r < p.ranks && cand == kVpNone; ++r) + { + cand = words[r * kVpClaimants + c]; + } + if (cand == kVpNone) + { + continue; + } + int const walker = info[2 + c]; + int const v = static_cast(cand) / 4; + int const position = ((v - walker) / vb) * 4 + static_cast(cand) % 4 + 1; + unsigned long long const key + = (static_cast(position) << 32) | static_cast(walker + 1); + if (walkers == 1 || key > bestKey) + { + bestKey = key; + token = static_cast(cand); + } + } + p.fs.outputTokens[row] = token; + } + } + gridDependencyTrigger(); +} + +//! vp_argmax: one warp per row; the greedy token as int64. PDL: waits on vp_fold first. +__global__ void __launch_bounds__(32) vpArgmaxKernel(VpParams p, int slots) +{ + gridDependencyWait(); + int const row = blockIdx.x; + int const lane = threadIdx.x; + VpMaxArg best{-INFINITY, INT_MAX}; + for (int i = lane; i < p.ranks * slots; i += 32) + { + uint2* const slot = vpGreedy(p.copy, p, i / slots, row, i % slots); + uint4 const w = vpAwait<2>(slot, p.status, p.timeoutNs); + storeRelaxedSys(slot, make_uint2(kSplitSentinel, kSplitSentinel)); + VpMaxArg const cand{__uint_as_float(w.x), static_cast(w.y)}; + best = VpMaxArgOp()(best, cand); + } +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + { + VpMaxArg const other{ + __shfl_xor_sync(0xffffffffu, best.value, offset), __shfl_xor_sync(0xffffffffu, best.index, offset)}; + best = VpMaxArgOp()(best, other); + } + if (lane == 0) + { + p.greedyTokens[row] = static_cast(best.index); + } + gridDependencyTrigger(); +} + } // namespace fusedSampling } // namespace kernels } // namespace tensorrt_llm @@ -837,6 +1401,193 @@ void watch(at::Tensor const& buf, at::Tensor const& firstSeen, int64_t blocks, i firstSeen.numel(), static_cast(timeoutMs) * 1000000ull); } +//! The vocabulary-parallel head's common parameters: rank geometry and the exchange buffer. +fs::VpParams vpLayout(int64_t rank, int64_t ranks, int64_t shardVecs, int64_t copy, at::IntArrayRef publish, + int64_t stateStride, int64_t greedyOffset, int64_t greedyStride, int64_t candOffset, int64_t candStride, + int64_t timeoutS) +{ + TORCH_CHECK(ranks >= 1 && ranks <= fs::kVpMaxRanks && rank >= 0 && rank < ranks, "rank / ranks out of range"); + TORCH_CHECK(publish.size() >= 1 && publish.size() <= static_cast(fs::kVpMaxRanks), "1..8 publish targets"); + TORCH_CHECK(shardVecs >= 1, "empty shard"); + fs::VpParams p; + p.rank = static_cast(rank); + p.ranks = static_cast(ranks); + p.shardVecs = static_cast(shardVecs); + p.copy = reinterpret_cast(copy); + for (size_t i = 0; i < publish.size(); ++i) + { + TORCH_CHECK(publish[i] % 16 == 0, "publish targets must be 16-byte aligned"); + p.publish[i] = reinterpret_cast(publish[i]); + } + p.numPublish = static_cast(publish.size()); + TORCH_CHECK(copy % 16 == 0 && greedyOffset % 16 == 0 && candOffset % 16 == 0, "exchange regions 16-byte aligned"); + p.stateStride = stateStride; + p.greedyOffset = greedyOffset; + p.greedyStride = greedyStride; + p.candOffset = candOffset; + p.candStride = candStride; + p.timeoutNs = static_cast(timeoutS) * 1000000000ull; + return p; +} + +void checkVpRows( + int64_t rows, int64_t shardVecs, int64_t ranks, int64_t stateStride, int64_t greedyStride, int64_t candStride) +{ + TORCH_CHECK(rows >= 1 && rows <= fs::kSplitMaxRows, "1..", fs::kSplitMaxRows, " rows"); + int const vb = fs::splitVirtualBlock(static_cast(rows)); + TORCH_CHECK((shardVecs + vb - 1) / vb <= fs::kVpWalkVectors, "shard too wide for the walk stage"); + TORCH_CHECK(rows * vb <= stateStride, "rows exceed the state slots"); + TORCH_CHECK(rows * fs::kVpGreedySlots <= greedyStride, "rows exceed the greedy slots"); + TORCH_CHECK(rows * fs::kVpClaimants <= candStride, "rows exceed the candidate slots"); + TORCH_CHECK(ranks * shardVecs * 4 <= INT32_MAX, "vocabulary too large"); +} + +//! vp_fold: this rank's shard folded per (row, virtual thread), published to every rank. +void vpFold(at::Tensor const& local, at::Tensor const& temperatures, at::Tensor const& topKs, at::Tensor const& topPs, + at::Tensor const& minPs, int64_t staging, at::IntArrayRef publish, int64_t rank, int64_t ranks, int64_t stateStride, + int64_t greedyOffset, int64_t greedyStride, int64_t candStride, bool greedy, double multiplier, double invCap, + double cap, bool pdl) +{ + TORCH_CHECK(local.scalar_type() == at::kBFloat16 && local.dim() == 2 && local.is_contiguous() && local.is_cuda(), + "local must be a contiguous CUDA BF16 [rows, shard] tensor"); + TORCH_CHECK(local.size(1) % 4 == 0, "the shard must be a whole number of 4-element vectors"); + TORCH_CHECK(reinterpret_cast(local.data_ptr()) % 8 == 0, "local must be 8-byte aligned"); + auto const rows = local.size(0); + auto const shardVecs = local.size(1) / 4; + checkVpRows(rows, shardVecs, ranks, stateStride, greedyStride, candStride); + at::cuda::CUDAGuard guard(local.device()); + auto p = vpLayout( + rank, ranks, shardVecs, publish[0], publish, stateStride, greedyOffset, greedyStride, 0, candStride, 0); + p.fs.numRows = static_cast(rows); + p.fs.vocabSize = static_cast(ranks * shardVecs * 4); + if (!greedy) + { + checkRowVector(temperatures, "temperatures", at::kFloat, rows); + checkRowVector(topKs, "top_ks", at::kInt, rows); + checkRowVector(topPs, "top_ps", at::kFloat, rows); + checkRowVector(minPs, "min_ps", at::kFloat, rows); + p.fs.temperatures = temperatures.data_ptr(); + p.fs.topKs = topKs.data_ptr(); + p.fs.topPs = topPs.data_ptr(); + p.fs.minPs = minPs.data_ptr(); + TORCH_CHECK(staging % 8 == 0, "staging must be 8-byte aligned"); + } + p.local = reinterpret_cast(local.data_ptr()); + p.staging = reinterpret_cast(staging); + p.cap = softCap(multiplier, invCap, cap); + int const vb = fs::splitVirtualBlock(static_cast(rows)); + dim3 const grid(static_cast(rows), vb / fs::kVpFoldThreads); + if (vb == fs::kSplitBlock) + { + if (greedy) + { + launch(fs::vpFoldKernel, grid, dim3(fs::kVpFoldThreads), pdl, p); + } + else + { + launch(fs::vpFoldKernel, grid, dim3(fs::kVpFoldThreads), pdl, p); + } + } + else if (greedy) + { + launch(fs::vpFoldKernel, grid, dim3(fs::kVpFoldThreads), pdl, p); + } + else + { + launch(fs::vpFoldKernel, grid, dim3(fs::kVpFoldThreads), pdl, p); + } +} + +//! vp_draw: combine, reduce, scan, draw, walk this rank's portions; publishes candidates. +//! Returns the token tensor vp_decide completes (filtered rows are written here). +at::Tensor vpDraw(at::Tensor const& pending, at::Tensor const& local, at::Tensor const& temperatures, + at::Tensor const& topKs, at::Tensor const& topPs, at::Tensor const& minPs, std::optional const& seed, + std::optional const& offset, at::Tensor const& info, at::Tensor const& status, int64_t copy, + at::IntArrayRef stagingOf, at::IntArrayRef publish, int64_t rank, int64_t ranks, int64_t stateStride, + int64_t greedyStride, int64_t candOffset, int64_t candStride, int64_t timeoutS, double multiplier, double invCap, + double cap, bool pdl) +{ + TORCH_CHECK(pending.scalar_type() == at::kBFloat16, "pending must be BF16"); + TORCH_CHECK( + local.scalar_type() == at::kBFloat16 && local.dim() == 2 && local.is_contiguous(), "local [rows, shard]"); + at::cuda::CUDAGuard guard(pending.device()); + auto fsParams = rowParams(pending, temperatures, topKs, topPs, minPs, seed, offset); + auto const rows = pending.size(0); + auto const shardVecs = local.size(1) / 4; + TORCH_CHECK(local.size(0) == rows && local.size(1) % 4 == 0 && pending.size(1) == ranks * local.size(1), + "pending is [rows, ranks * shard] of local [rows, shard]"); + checkVpRows(rows, shardVecs, ranks, stateStride, greedyStride, candStride); + TORCH_CHECK(static_cast(stagingOf.size()) == ranks, "one staging address per rank"); + TORCH_CHECK(info.is_cuda() && info.scalar_type() == at::kInt && info.is_contiguous() + && info.numel() >= rows * fs::kVpInfoWords, + "info holds kVpInfoWords int32 per row"); + TORCH_CHECK(status.is_cuda() && status.scalar_type() == at::kInt && status.numel() >= 1, "status int32[1]"); + auto p = vpLayout( + rank, ranks, shardVecs, copy, publish, stateStride, 0, greedyStride, candOffset, candStride, timeoutS); + p.fs = fsParams; + for (int64_t r = 0; r < ranks; ++r) + { + p.stagingOf[r] = reinterpret_cast(stagingOf[r]); + } + p.local = reinterpret_cast(local.data_ptr()); + p.info = info.data_ptr(); + p.status = status.data_ptr(); + p.cap = softCap(multiplier, invCap, cap); + auto tokens = at::empty({rows}, pending.options().dtype(at::kInt)); + p.fs.outputTokens = tokens.data_ptr(); + if (fs::splitVirtualBlock(static_cast(rows)) == fs::kSplitBlock) + { + launch(fs::vpDrawKernel, dim3(static_cast(rows)), dim3(fs::kSplitBlock), pdl, p); + } + else + { + launch(fs::vpDrawKernel, dim3(static_cast(rows)), dim3(fs::kSplitNarrowBlock), + pdl, p); + } + return tokens; +} + +//! vp_decide: every rank's candidates -> the token of each unfiltered row, into `tokens`. +void vpDecide(at::Tensor const& tokens, at::Tensor const& info, at::Tensor const& status, int64_t copy, int64_t rank, + int64_t ranks, int64_t shardVecs, int64_t stateStride, int64_t greedyStride, int64_t candOffset, int64_t candStride, + int64_t timeoutS, bool pdl) +{ + TORCH_CHECK(tokens.is_cuda() && tokens.scalar_type() == at::kInt && tokens.dim() == 1, "tokens int32 [rows]"); + auto const rows = tokens.size(0); + checkVpRows(rows, shardVecs, ranks, stateStride, greedyStride, candStride); + TORCH_CHECK(info.is_cuda() && info.scalar_type() == at::kInt && info.numel() >= rows * fs::kVpInfoWords, "info"); + TORCH_CHECK(status.is_cuda() && status.scalar_type() == at::kInt && status.numel() >= 1, "status int32[1]"); + at::cuda::CUDAGuard guard(tokens.device()); + int64_t const self[1] = {copy}; + auto p = vpLayout(rank, ranks, shardVecs, copy, at::IntArrayRef(self, 1), stateStride, 0, greedyStride, candOffset, + candStride, timeoutS); + p.fs.numRows = static_cast(rows); + p.fs.outputTokens = tokens.data_ptr(); + p.info = info.data_ptr(); + p.status = status.data_ptr(); + launch(fs::vpDecideKernel, dim3(static_cast(rows)), dim3(32), pdl, p, + fs::splitVirtualBlock(static_cast(rows))); +} + +//! vp_argmax: every rank's greedy words -> torch.argmax's int64 token per row. +at::Tensor vpArgmax(at::Tensor const& status, int64_t rows, int64_t copy, int64_t rank, int64_t ranks, + int64_t shardVecs, int64_t stateStride, int64_t greedyOffset, int64_t greedyStride, int64_t candStride, + int64_t timeoutS, bool pdl) +{ + checkVpRows(rows, shardVecs, ranks, stateStride, greedyStride, candStride); + TORCH_CHECK(status.is_cuda() && status.scalar_type() == at::kInt && status.numel() >= 1, "status int32[1]"); + at::cuda::CUDAGuard guard(status.device()); + int64_t const self[1] = {copy}; + auto p = vpLayout(rank, ranks, shardVecs, copy, at::IntArrayRef(self, 1), stateStride, greedyOffset, greedyStride, + 0, candStride, timeoutS); + p.status = status.data_ptr(); + auto tokens = at::empty({rows}, status.options().dtype(at::kLong)); + p.greedyTokens = tokens.data_ptr(); + launch(fs::vpArgmaxKernel, dim3(static_cast(rows)), dim3(32), pdl, p, + fs::splitVirtualBlock(static_cast(rows)) / fs::kVpFoldThreads); + return tokens; +} + } // namespace TORCH_LIBRARY(musespark_open_head_sampler, m) @@ -854,6 +1605,23 @@ TORCH_LIBRARY(musespark_open_head_sampler, m) m.def("ref_fold(Tensor capped, Tensor temperatures, Tensor(a!) states) -> ()"); m.def("cap(Tensor x, float multiplier, float inv_cap, float cap) -> Tensor"); m.def("watch(Tensor buf, Tensor(a!) first_seen, int blocks, int timeout_ms) -> ()"); + m.def( + "vp_fold(Tensor local, Tensor temperatures, Tensor top_ks, Tensor top_ps, Tensor min_ps, int staging, " + "int[] publish, int rank, int ranks, int state_stride, int greedy_offset, int greedy_stride, " + "int cand_stride, bool greedy, float multiplier, float inv_cap, float cap, bool pdl) -> ()"); + m.def( + "vp_draw(Tensor(a!) pending, Tensor local, Tensor temperatures, Tensor top_ks, Tensor top_ps, " + "Tensor min_ps, Tensor? seed, Tensor? offset, Tensor(b!) info, Tensor(c!) status, int copy, " + "int[] staging_of, int[] publish, int rank, int ranks, int state_stride, int greedy_stride, " + "int cand_offset, int cand_stride, int timeout_s, float multiplier, float inv_cap, float cap, " + "bool pdl) -> Tensor"); + m.def( + "vp_decide(Tensor(a!) tokens, Tensor info, Tensor(b!) status, int copy, int rank, int ranks, " + "int shard_vecs, int state_stride, int greedy_stride, int cand_offset, int cand_stride, int timeout_s, " + "bool pdl) -> ()"); + m.def( + "vp_argmax(Tensor(a!) status, int rows, int copy, int rank, int ranks, int shard_vecs, int state_stride, " + "int greedy_offset, int greedy_stride, int cand_stride, int timeout_s, bool pdl) -> Tensor"); } TORCH_LIBRARY_IMPL(musespark_open_head_sampler, CUDA, m) @@ -864,4 +1632,8 @@ TORCH_LIBRARY_IMPL(musespark_open_head_sampler, CUDA, m) m.impl("ref_fold", &refFold); m.impl("cap", &capRows); m.impl("watch", &watch); + m.impl("vp_fold", &vpFold); + m.impl("vp_draw", &vpDraw); + m.impl("vp_decide", &vpDecide); + m.impl("vp_argmax", &vpArgmax); } diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/head_sampler.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/head_sampler.py index f694783969ff..08a01afdab61 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/head_sampler.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/head_sampler.py @@ -598,3 +598,468 @@ def shared_gather_workspace(mapping: Mapping, shard: int) -> HeadGatherWorkspace workspace = HeadGatherWorkspace(mapping, shard) _SHARED[key] = workspace return workspace + + +# --------------------------------------------------------------------------------------- +# The vocabulary-parallel head (fast path ``vocab_parallel_head``): no logits all-gather. + +#: Fewest rows the vocabulary-parallel head takes (two K=7 verify blocks); 8-row steps keep +#: the ``lm_head_inbox`` gather. +VP_MIN_ROWS = 16 +#: Claimants of a row whose crossings the ranks exchange (``kVpClaimants``). +VP_CLAIMANTS = 4 +#: Greedy words per (rank, row): one per ``vp_fold`` CTA of the 1024-thread block. +VP_GREEDY_SLOTS = VIRTUAL_THREADS // FOLD_THREADS +#: ``vp_draw``'s per-row record for ``vp_decide`` (int32 words). +VP_INFO_WORDS = 2 + VP_CLAIMANTS +#: Most ranks (``kVpMaxRanks``). +VP_MAX_RANKS = 8 + + +class VocabParallelLayout: + """The exchange buffer of the vocabulary-parallel head for ``ranks`` shards of ``shard`` + elements and calls of at most ``max_rows`` rows. + + One copy per rank, every 32-bit word the sentinel ``0xFFFFFFFF`` between calls: + ``[ranks][state_stride]`` 16-byte states (``vp_fold``'s per-virtual-thread partials), then + ``[ranks][greedy_stride]`` 8-byte greedy words at ``greedy_offset``, then + ``[ranks][cand_stride]`` 4-byte candidates at ``cand_offset``; ``nbytes`` in all. + """ + + def __init__(self, ranks: int, shard: int, max_rows: int): + if not 1 <= ranks <= VP_MAX_RANKS: + raise ValueError(f"1..{VP_MAX_RANKS} ranks, got {ranks}") + if shard % 4: + raise ValueError(f"shard {shard} is not a whole number of 4-element vectors") + if not VP_MIN_ROWS <= max_rows <= WIDE_MAX_ROWS: + raise ValueError(f"max_rows must be in {VP_MIN_ROWS}..{WIDE_MAX_ROWS}, got {max_rows}") + self.ranks = int(ranks) + self.shard = int(shard) + self.shard_vecs = self.shard // 4 + self.vocab = self.ranks * self.shard + self.max_rows = int(max_rows) + widest = max( + -(-self.shard_vecs // virtual_threads(r)) for r in {min(max_rows, MAX_ROWS), max_rows} + ) + if widest > 64: + raise ValueError(f"shard {shard} is too wide for the walk stage") + self.state_stride = max(r * virtual_threads(r) for r in {min(max_rows, MAX_ROWS), max_rows}) + self.greedy_stride = self.max_rows * VP_GREEDY_SLOTS + self.cand_stride = self.max_rows * VP_CLAIMANTS + self.greedy_offset = self.ranks * self.state_stride * 16 + self.cand_offset = self.greedy_offset + self.ranks * self.greedy_stride * 8 + self.nbytes = -(-(self.cand_offset + self.ranks * self.cand_stride * 4) // 16) * 16 + self.staging_bytes = self.max_rows * self.shard * 2 + + def strides(self) -> dict: + return dict( + state_stride=self.state_stride, + greedy_stride=self.greedy_stride, + cand_stride=self.cand_stride, + ) + + +class _VocabParallelRank: + """One rank's view: its exchange copy, the publish targets, every rank's staging rows.""" + + def __init__( + self, + layout: VocabParallelLayout, + rank: int, + copy: int, + publish: list, + staging: int, + staging_of: list, + info: torch.Tensor, + status: torch.Tensor, + ): + self.layout = layout + self.rank = int(rank) + self.copy = int(copy) + self.publish = [int(x) for x in publish] + self.staging = int(staging) + self.staging_of = [int(x) for x in staging_of] + self.info = info + self.status = status + + def fold(self, local, temperatures, top_ks, top_ps, min_ps, multiplier, soft_cap, greedy, pdl): + m, inv, cap = cap_constants(multiplier, soft_cap) + lay = self.layout + torch.ops.musespark_open_head_sampler.vp_fold( + local, + temperatures, + top_ks, + top_ps, + min_ps, + self.staging, + self.publish, + self.rank, + lay.ranks, + lay.state_stride, + lay.greedy_offset, + lay.greedy_stride, + lay.cand_stride, + greedy, + m, + inv, + cap, + pdl, + ) + + def draw( + self, + pending, + local, + temperatures, + top_ks, + top_ps, + min_ps, + seed, + offset, + multiplier, + soft_cap, + pdl, + ): + m, inv, cap = cap_constants(multiplier, soft_cap) + lay = self.layout + return torch.ops.musespark_open_head_sampler.vp_draw( + pending, + local, + temperatures, + top_ks, + top_ps, + min_ps, + seed, + offset, + self.info, + self.status, + self.copy, + self.staging_of, + self.publish, + self.rank, + lay.ranks, + lay.state_stride, + lay.greedy_stride, + lay.cand_offset, + lay.cand_stride, + POLL_TIMEOUT_S, + m, + inv, + cap, + pdl, + ) + + def decide(self, tokens, pdl): + lay = self.layout + torch.ops.musespark_open_head_sampler.vp_decide( + tokens, + self.info, + self.status, + self.copy, + self.rank, + lay.ranks, + lay.shard_vecs, + lay.state_stride, + lay.greedy_stride, + lay.cand_offset, + lay.cand_stride, + POLL_TIMEOUT_S, + pdl, + ) + + def argmax(self, rows, pdl): + lay = self.layout + return torch.ops.musespark_open_head_sampler.vp_argmax( + self.status, + rows, + self.copy, + self.rank, + lay.ranks, + lay.shard_vecs, + lay.state_stride, + lay.greedy_offset, + lay.greedy_stride, + lay.cand_stride, + POLL_TIMEOUT_S, + pdl, + ) + + +def vocab_parallel_sample( + view: _VocabParallelRank, + pending: torch.Tensor, + local: torch.Tensor, + temperatures: torch.Tensor, + top_ks: torch.Tensor, + top_ps: torch.Tensor, + min_ps: torch.Tensor, + seed: Optional[torch.Tensor], + offset: Optional[torch.Tensor], + multiplier: float, + soft_cap: float, + pdl: bool = False, +) -> torch.Tensor: + """One rank's sampled tokens (int32 ``[rows]``) of the step whose shard ``local`` is. + + ``vp_fold``, ``vp_draw`` and ``vp_decide`` in order; every rank of the group makes the same + call with its own shard. ``pending`` is the ``[rows, vocab]`` BF16 row: written for rows with + a top-k / top-p / min-p filter only (gathered from the ranks' staging rows), left as it is + otherwise. + """ + rows = local.shape[0] + view.fold(local, temperatures, top_ks, top_ps, min_ps, multiplier, soft_cap, False, pdl) + tokens = view.draw( + pending, + local, + temperatures, + top_ks, + top_ps, + min_ps, + seed, + offset, + multiplier, + soft_cap, + pdl, + ) + view.decide(tokens[:rows], pdl) + return tokens + + +def vocab_parallel_argmax( + view: _VocabParallelRank, + local: torch.Tensor, + multiplier: float, + soft_cap: float, + pdl: bool = False, +) -> torch.Tensor: + """One rank's greedy tokens (int64 ``[rows]``): ``torch.argmax`` of the soft-capped FP32 row.""" + view.fold( + local, view.status, view.status, view.status, view.status, multiplier, soft_cap, True, pdl + ) + return view.argmax(local.shape[0], pdl) + + +def vocab_parallel_reference( + locals_: list, + temperatures: torch.Tensor, + top_ks: torch.Tensor, + top_ps: torch.Tensor, + min_ps: torch.Tensor, + seed: Optional[torch.Tensor], + offset: Optional[torch.Tensor], + multiplier: float, + soft_cap: float, + greedy: bool, +) -> torch.Tensor: + """The built-in head on the rank-order concatenation of the shards: ``torch.argmax`` of the + soft-capped FP32 row (greedy) or ``fused_sample_from_logits`` on it.""" + capped = reference_capped(torch.cat(list(locals_), dim=1), multiplier, soft_cap) + if greedy: + return torch.argmax(capped, dim=-1) + return torch.ops.trtllm.fused_sample_from_logits( + capped, temperatures, top_ks, top_ps, min_ps, seed, offset + ) + + +class VocabParallelEmulation: + """``ranks`` virtual ranks on one GPU: every value is stored into each rank's copy (instead + of one multicast store), every rank's staging rows are plain tensors. The kernels are the + served ones; a call runs every rank's ``vp_fold``, then every ``vp_draw``, then every + ``vp_decide`` (so no poll waits).""" + + def __init__(self, ranks: int, shard: int, max_rows: int, device): + load_library() + self.layout = VocabParallelLayout(ranks, shard, max_rows) + lay = self.layout + self.device = torch.device(device) + self.copies = [ + torch.full((lay.nbytes // 4,), -1, dtype=torch.int32, device=self.device) + for _ in range(ranks) + ] + self.staging = [ + torch.empty((max_rows, shard), dtype=torch.bfloat16, device=self.device) + for _ in range(ranks) + ] + publish = [c.data_ptr() for c in self.copies] + staging_of = [s.data_ptr() for s in self.staging] + self.status = torch.zeros(1, dtype=torch.int32, device=self.device) + self.views = [ + _VocabParallelRank( + lay, + r, + self.copies[r].data_ptr(), + publish, + staging_of[r], + staging_of, + torch.empty(max_rows * VP_INFO_WORDS, dtype=torch.int32, device=self.device), + self.status, + ) + for r in range(ranks) + ] + + def sample( + self, + locals_, + temperatures, + top_ks, + top_ps, + min_ps, + seed, + offset, + multiplier, + soft_cap, + pdl=False, + ): + """Every rank's tokens (list of int32 ``[rows]``) and pending rows.""" + rows = locals_[0].shape[0] + pendings = [ + torch.zeros((rows, self.layout.vocab), dtype=torch.bfloat16, device=self.device) + for _ in self.views + ] + for view, local in zip(self.views, locals_): + view.fold(local, temperatures, top_ks, top_ps, min_ps, multiplier, soft_cap, False, pdl) + tokens = [ + view.draw( + pending, + local, + temperatures, + top_ks, + top_ps, + min_ps, + seed, + offset, + multiplier, + soft_cap, + pdl, + ) + for view, pending, local in zip(self.views, pendings, locals_) + ] + for view, t in zip(self.views, tokens): + view.decide(t, pdl) + return tokens, pendings + + def argmax(self, locals_, multiplier, soft_cap, pdl=False): + for view, local in zip(self.views, locals_): + view.fold( + local, + self.status, + self.status, + self.status, + self.status, + multiplier, + soft_cap, + True, + pdl, + ) + return [view.argmax(locals_[0].shape[0], pdl) for view in self.views] + + def all_sentinel(self) -> bool: + return all(bool((c == -1).all()) for c in self.copies) + + +class VocabParallelWorkspace: + """The vocabulary-parallel head's fabric buffers for one tensor-parallel group. + + ``exchange``: the multicast exchange buffer (``VocabParallelLayout``), each rank's copy filled + with the sentinel; values are stored through the multicast address, read and reset in the + own copy. ``staging``: ``[max_rows, shard]`` BF16 per rank (unicast; peers read it), written + by ``vp_fold`` for rows with a filter only. + + Protocol: (1) every exchange slot has one writer per call (``vp_fold`` per (rank, row, + virtual thread) or (rank, row, CTA); ``vp_draw`` per (rank, row, claimant)); every rank + reads each slot of its own copy once per call and resets it. (2) At the start of every call + every word of every copy is the sentinel: after construction (fill, synchronize, agreement) + and restored by the call's readers. (3) Reuse across calls: a peer stores into this copy (and + into its staging rows) again only in a later call, which it reaches after a collective of + the group issued between the two calls (the decode step's all-reduces and hops, the + drafter's hops). This rank publishes into that collective only after this call's last + kernel has completed: every kernel between them is a plain launch or waits for its + predecessor (``griddepcontrol.wait``) before it completes, and the collective publishes + after its own wait -- so an all-PDL chain (the Lamport hops) orders it as an NCCL call + does. Two calls with no collective between them are not ordered: a peer's next fold can + overwrite a word this rank has not read. (4) Every rank makes the same calls: the path + depends only on rank-uniform facts. + + Construction is collective over the group and succeeds or raises on every rank. + """ + + def __init__(self, mapping: Mapping, shard: int, max_rows: int): + from tensorrt_llm._torch.distributed.ops import ( + _get_mnnvl_workspace_comm, + _make_mnnvl_mcast_buffer, + _mnnvl_device_index, + _mnnvl_workspace_all_succeeded, + ) + + self.ranks = mapping.tp_size + self.rank = mapping.tp_rank + self.device = torch.device("cuda", _mnnvl_device_index(mapping)) + self._comm = _get_mnnvl_workspace_comm(mapping) + self.library = "" + error: Optional[BaseException] = None + try: + self.layout = VocabParallelLayout(self.ranks, shard, max_rows) + if torch.cuda.get_device_capability(self.device) not in ((10, 3), (10, 7)): + raise RuntimeError("the vocabulary-parallel head runs on SM103 or SM107 only") + self.library = load_library() + except Exception as exc: # noqa: BLE001 - every failure must reach the agreement + error = exc + if not _mnnvl_workspace_all_succeeded(self._comm, error is None): + raise RuntimeError( + "[musespark_open] vocab_parallel_head: the head sampler library failed to load on at " + "least one rank" + f"{f' (this rank: {error!r})' if error is not None else ''}" + ) from error + lay = self.layout + try: + multi_node = mapping.is_multi_node() + self._exchange = _make_mnnvl_mcast_buffer(self._comm, lay.nbytes, mapping, multi_node) + self._staging = _make_mnnvl_mcast_buffer( + self._comm, lay.staging_bytes, mapping, multi_node + ) + self.copy = self._exchange.get_uc_buffer(self.rank, (lay.nbytes // 4,), torch.int32, 0) + self._mc = self._exchange.get_mc_buffer((lay.nbytes // 4,), torch.int32, 0) + self.staging_of = [ + self._staging.get_uc_buffer(peer, (lay.staging_bytes // 2,), torch.bfloat16, 0) + for peer in range(self.ranks) + ] + with torch.cuda.device(self.device): + self.copy.fill_(-1) + self.info = torch.empty( + lay.max_rows * VP_INFO_WORDS, dtype=torch.int32, device=self.device + ) + self.status = torch.zeros(1, dtype=torch.int32, device=self.device) + torch.cuda.synchronize(self.device) + except Exception as exc: # noqa: BLE001 - every failure must reach the agreement + error = exc + if not _mnnvl_workspace_all_succeeded(self._comm, error is None): + raise RuntimeError( + "[musespark_open] vocab_parallel_head: workspace failed on at least one rank" + f"{f' (this rank: {error!r})' if error is not None else ''}" + ) from error + self.view = _VocabParallelRank( + lay, + self.rank, + self.copy.data_ptr(), + [self._mc.data_ptr()], + self.staging_of[self.rank].data_ptr(), + [t.data_ptr() for t in self.staging_of], + self.info, + self.status, + ) + + @property + def max_rows(self) -> int: + return self.layout.max_rows + + @property + def vocab(self) -> int: + return self.layout.vocab + + def nbytes(self) -> int: + """Device memory this rank holds for the head (exchange copy + staging rows).""" + return self.layout.nbytes + self.layout.staging_bytes + + def errors(self) -> int: + """The poll status word (synchronizes).""" + return int(self.status.item()) From 95b8a9a5dec3b7bd7221d557a91079ef523063f5 Mon Sep 17 00:00:00 2001 From: Xiao Wang <24860335+xwang233@users.noreply.github.com> Date: Mon, 26 Oct 2026 18:44:16 +0000 Subject: [PATCH 11/16] [None][feat] modeling_v2 musespark_open tep: wide drafter hops, drafter GEMM tactics and per-GPU PDL chain widths - wide_drafter_hops: the fused drafter all-reduce hops (one-shot all-reduce, finish convolution and add-RMSNorm per site) also on block decodes above 8 rows, up to the device's served rows (GB300: 9..56, Rubin: 9..64; the blocks of several requests at K = 7), on a second, wide one-shot workspace. Forwards of at most 8 rows keep the first. - Drafter GEMMs: the fixed low-M split-K tables at TP4 on sm103 and sm107 list the blocks of 2..4 requests (16, 24 and 32 rows) next to the one-request block. - drafter_pdl_chain: the conv / merge / silu entry triggers act on every block width of 1..4 requests at which the drafter's tactic table covers every successor; the merge and silu widths are per GPU (sm103: 1..4 requests for both; sm107: merge on one request, silu on one and four). The block decode reads the widths the chain's install recorded. Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com> --- .../drafter_block_decode/block_decode.py | 22 +- .../kernels/drafter_hops/hops.py | 65 ++++- .../kernels/drafter_pdl_chain/chain.py | 99 +++++++- .../ops/allreduce/drafter_hop.py | 230 +++++++++++++++--- .../ops/gemm/fixed_low_m_gemm.py | 39 ++- 5 files changed, 403 insertions(+), 52 deletions(-) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/drafter_block_decode/block_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/drafter_block_decode/block_decode.py index 292e17633680..358c8e76a726 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/drafter_block_decode/block_decode.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/drafter_block_decode/block_decode.py @@ -68,9 +68,11 @@ #: Operand dtypes the fused convolution takes; anything else runs the built-in. SUPPORTED_DTYPES = (torch.bfloat16, torch.float16, torch.float32) -#: Block rows on which the PDL chain's conv / merge entry triggers apply (``chain.ROWS``): -#: one request of K + 1 = 8 rows, whose successor GEMMs run the ``drafter_fixed_gemms`` tactic. -_ENTRY_ROWS = 8 +#: Block rows on which the PDL chain's conv / merge entry triggers apply when the chain's +#: install did not record their widths (``_pdl_entry_rows`` / ``_pdl_merge_rows``, +#: ``chain.entry_rows`` / ``chain.MERGE_ENTRY_ROWS_BY_SM``): one request of K + 1 = 8 rows, whose +#: successor GEMMs run the ``drafter_fixed_gemms`` tactic. +_ENTRY_ROWS = (8,) class FusedDFlash2BlockConv(DFlash2BlockConv): @@ -95,9 +97,13 @@ def _convolve( or hidden_states.dtype not in SUPPORTED_DTYPES ): return super()._convolve(hidden_states, delta, side, block_size) - if side == 0 and hidden_states.shape[0] == _ENTRY_ROWS and self.__dict__.get("_pdl_conv"): - # PDL chain conv trigger (drafter_pdl_chain/chain.py): the prepare side of an 8-row block, whose - # successor GEMM loads only its weight before its own wait. + if ( + side == 0 + and self.__dict__.get("_pdl_conv") + and hidden_states.shape[0] in self.__dict__.get("_pdl_entry_rows", _ENTRY_ROWS) + ): + # PDL chain conv trigger (drafter_pdl_chain/chain.py): the prepare side of a block whose + # successor GEMM runs a fixed tactic, which loads only its weight before its own wait. return dflash2_block_conv( hidden_states, delta, base_kernel, block_size, self.group_size, True ) @@ -318,7 +324,9 @@ def windowed_block_decode( layers = self.model.layers prenormed = None # PDL chain merge trigger (drafter_pdl_chain/chain.py): the merge releases o_proj's GEMM at its entry. - merge_entry = bool(self.__dict__.get("_pdl_merge")) and B * block_size == _ENTRY_ROWS + merge_entry = bool(self.__dict__.get("_pdl_merge")) and B * block_size in self.__dict__.get( + "_pdl_merge_rows", _ENTRY_ROWS + ) # wide_drafter_window_attention: a block decode of several requests on the wide op. wide = self.__dict__.get("_wide_window_attention") if wide is not None and not wide.takes(B * block_size): diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/drafter_hops/hops.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/drafter_hops/hops.py index eabd009108a8..d6970d7f8eb1 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/drafter_hops/hops.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/drafter_hops/hops.py @@ -18,6 +18,11 @@ ``down_proj``) is the same on every rank and every forward; a forward takes the fused hops at all its sites or at none, decided on the block's row count (``DrafterHopCall.takes``), which every rank shares. + +With ``wide_drafter_hops`` (``install_wide``) a forward of ``MAX_ROWS + 1`` .. +``wide_served_rows()`` rows (the blocks of several requests at K = 7, per device in +``WIDE_SERVED_ROWS_BY_SM``) takes the same hops on a second, wide workspace; a forward of at +most ``MAX_ROWS`` rows keeps the first one. """ from __future__ import annotations @@ -46,18 +51,29 @@ class DrafterHopCall: def __init__(self, workspace: "_hop.DrafterHopWorkspace") -> None: self.workspace = workspace + #: The wide workspace (``wide_drafter_hops``), or None. + self.wide: Optional["_hop.DrafterHopWorkspace"] = None + #: The workspace of the forward ``takes`` admitted last. + self.active = workspace #: Hops issued by this rank (a test counts them to see a forward took the fused path). self.calls = 0 def takes(self, rows: int, block_size: int, dtype: torch.dtype) -> bool: - """Whether a forward of ``rows`` rows runs every site on the fused hop. Rank-uniform.""" - if dtype != torch.bfloat16 or not 1 <= rows <= _hop.MAX_ROWS or block_size < 1: + """Whether a forward of ``rows`` rows runs every site on the fused hop, and on which + workspace. Rank-uniform.""" + if dtype != torch.bfloat16 or block_size < 1 or rows < 1: return False - return True + if rows <= _hop.MAX_ROWS: + self.active = self.workspace + return True + if self.wide is not None and rows <= self.wide.max_rows: + self.active = self.wide + return True + return False def begin_forward(self) -> None: """A drafter forward starts; its hops name the sites from 0 (see the workspace's).""" - self.workspace.begin_forward() + self.active.begin_forward() def _site(self, site, partial, residual, coefficients, conv, norm, block_size): # ``coefficients`` is what ``prepare`` returned for ``finish``: [rows, taps, groups]. @@ -80,7 +96,7 @@ def _site(self, site, partial, residual, coefficients, conv, norm, block_size): norm.weight, norm.variance_epsilon, block_size, - self.workspace, + self.active, ) def attention( @@ -353,3 +369,42 @@ def install(draft_model: Optional[nn.Module], mapping=None) -> None: draft_model._drafter_hop = DrafterHopCall(workspace) _hop.log_once_enabled(2 * len(layers), False, early) logger.info("[musespark_open] drafter_hops: fused drafter hops installed") + + +def install_wide(draft_model: Optional[nn.Module], mapping=None) -> None: + """``wide_drafter_hops``: the fused hops also on drafter forwards of ``MAX_ROWS + 1`` .. + ``wide_served_rows()`` rows of this device, on a wide workspace of their own. Collective; + after ``install``. + + The calls the wide hops replace are the same three kernels per site as at 8 rows (the + MNNVL one-shot all-reduce ``request_mnnvl`` gave every row projection on one tray, the + finish convolution, flashinfer's fused add + RMSNorm); the hop reproduces them per row, + and the tap-1 mask per block (``row % block``). Idempotent. + + Raises: + ValueError: when ``drafter_hops`` is not installed on this drafter. + """ + if not installed(draft_model): + raise ValueError("[musespark_open] wide_drafter_hops needs drafter_hops installed first") + call: DrafterHopCall = draft_model._drafter_hop + if call.wide is not None: + return + narrow = call.workspace + layers = draft_model.model.layers + norms = [layer.post_attention_layernorm for layer in layers] + norms += [layer.input_layernorm for layer in layers[1:]] + [draft_model.model.norm] + eps = sorted({float(norm.variance_epsilon) for norm in norms}) + rows = _hop.wide_served_rows() + call.wide = _hop.DrafterHopWorkspace( + layers[0].self_attn.o_proj.mapping, + sites=narrow.sites, + eps=eps, + early_trigger=narrow.early_trigger, + pre_wait=[tensor for _, tensor in pre_wait_operands(draft_model)], + max_rows=rows, + ) + logger.info( + f"[musespark_open] wide_drafter_hops: fused drafter hops on forwards of {_hop.MAX_ROWS + 1}.." + f"{rows} rows ({call.wide.areas} areas, " + f"{call.wide.nbytes() / 2**20:.1f} MiB of fabric memory on this rank)" + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/drafter_pdl_chain/chain.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/drafter_pdl_chain/chain.py index 57c002131551..9a6e6aba94f4 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/drafter_pdl_chain/chain.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/drafter_pdl_chain/chain.py @@ -37,10 +37,19 @@ right before the draft LM-head GEMM and joined where the projection is called, after the key all-reduce, instead of on the main stream after it. -Scope: the one-request step, whose successors are known. ``pdl_conv`` / -``pdl_merge`` / ``pdl_silu`` act on the block decode's 8 rows (K + 1, +Scope: block decodes whose successors are known. ``pdl_conv`` / +``pdl_merge`` / ``pdl_silu`` act on the block decode's 8 rows (one request, K + 1, K = 7), where ``drafter_fixed_gemms`` has a fixed tactic for ``qkv_proj`` / ``gate_up_proj`` / -``o_proj`` / ``down_proj`` / the first convolution's ``kernel_projection``. +``o_proj`` / ``down_proj`` / the first convolution's ``kernel_projection``. ``pdl_conv`` +also acts on every other block width of ``ENTRY_ROW_CANDIDATES`` (2..4 requests: 16, 24 +and 32 rows) at which the drafter's table lists a fixed tactic for all four of those +successors (``entry_rows``); a width without one runs the base launches. ``pdl_silu`` +and ``pdl_merge`` act on the widths their device's table lists among those +(``SILU_ENTRY_ROWS_BY_SM`` / ``MERGE_ENTRY_ROWS_BY_SM``), each measured in the in-model +drafter chain on that GPU. On GB300 all three triggers together are the fastest chain at +1..4 requests. On Rubin an entry-released ``o_proj`` competes with the multi-request merge +it follows (the chain is slower with the merge trigger at 2..4 requests), and the silu +trigger on top of the conv trigger gains nothing at two requests and loses at three. ``pdl_head`` acts on a one-request step whose K slots are one contiguous row range (any K); its successor is the low-M dispatcher's runner for the LM head (cuBLAS through ``torch.mm``, which has no programmatic attribute, or @@ -93,12 +102,50 @@ from tensorrt_llm._torch.modules.linear import Linear, UnquantizedLinearMethod from tensorrt_llm.logger import logger +from ...ops.gemm.fixed_low_m_gemm import device_capability from ...ops.spec.drafter_silu_and_mul import drafter_silu_and_mul -#: The drafter block rows these launches act on: one request of K + 1 = 8 rows, -#: the only width ``drafter_fixed_gemms`` lists a tactic for. +#: The drafter block rows of one request, K + 1 = 8: the width ``pdl_conv`` / ``pdl_merge`` / +#: ``pdl_silu`` require a fixed tactic at, and the ``pdl_head`` / ``selector_fork`` step. ROWS = 8 +#: The block widths the conv / merge / silu entry triggers may act on: 1..4 requests at +#: K = 7, at most the fixed GEMM's ``MAX_M`` of 32 rows. A width is taken when the +#: drafter's tactic table covers every successor at it (``entry_rows``). +ENTRY_ROW_CANDIDATES = (8, 16, 24, 32) + +#: The block widths of the merge trigger per compute capability (``trigger_rows``), among +#: those of ``entry_rows``. sm103: one to four requests. sm107: the one-request block only. +MERGE_ENTRY_ROWS_BY_SM: dict[tuple[int, int], frozenset] = { + (10, 3): frozenset(ENTRY_ROW_CANDIDATES), + (10, 7): frozenset({ROWS}), +} + +#: The block widths of the silu trigger per compute capability (``trigger_rows``), among +#: those of ``entry_rows``. sm103: one to four requests. sm107: one and four requests. +SILU_ENTRY_ROWS_BY_SM: dict[tuple[int, int], frozenset] = { + (10, 3): frozenset(ENTRY_ROW_CANDIDATES), + (10, 7): frozenset({ROWS, 4 * ROWS}), +} + + +def trigger_rows(table: dict, capability: Optional[tuple] = None) -> frozenset: + """The block widths ``table`` (``MERGE_ENTRY_ROWS_BY_SM`` / ``SILU_ENTRY_ROWS_BY_SM``) + lists for ``capability`` (the drafter GEMMs' device, ``fixed_low_m_gemm.device_capability``, + when None). + + Raises: + ValueError: on a device the table does not list. + """ + key = tuple(device_capability() if capability is None else capability) + rows = table.get(key) + if rows is None: + raise ValueError( + f"drafter_pdl_chain: no trigger widths for SM {key} (supported: {sorted(table)})" + ) + return rows + + #: The draft head's rows for that request: its K = 7 draft slots. ``selector_fork`` #: forks the selector projection only for this many rows. HEAD_ROWS = ROWS - 1 @@ -160,8 +207,29 @@ def _merge_reason(draft_model: nn.Module) -> Optional[str]: return None +def entry_rows(draft_model: nn.Module) -> frozenset: + """The widths of ``ENTRY_ROW_CANDIDATES`` at which every successor of the conv / merge / + silu triggers (``qkv_proj``, ``gate_up_proj``, ``o_proj``, ``down_proj`` of every layer) + has a fixed tactic. A function of the drafter's shapes and TP table: rank-uniform.""" + rows = [] + for width in ENTRY_ROW_CANDIDATES: + if all( + _fixed_tactic_reason(linear, width, "") is None + for layer in draft_model.model.layers + for linear in ( + layer.self_attn.qkv_proj, + layer.mlp.gate_up_proj, + layer.self_attn.o_proj, + layer.mlp.down_proj, + ) + ): + rows.append(width) + return frozenset(rows) + + class EntryTriggerGatedMLP(GatedMLP): - """The built-in ``GatedMLP`` whose 8-row ``silu(gate) * up`` releases ``down_proj`` at entry. + """The built-in ``GatedMLP`` whose ``silu(gate) * up`` releases ``down_proj`` at entry on the + widths of ``_pdl_entry_rows`` (``entry_rows``; 8 rows when unset). Never constructed: ``install`` re-classes the drafter's ``GatedMLP`` modules in place. Every other call is the built-in method's. @@ -172,7 +240,7 @@ def _apply_activation(self, x, *, has_lora: bool = False): not has_lora and isinstance(x, torch.Tensor) and x.dim() == 2 - and x.shape[0] == ROWS + and x.shape[0] in self.__dict__.get("_pdl_entry_rows", (ROWS,)) and x.dtype == torch.bfloat16 and x.is_contiguous() ): @@ -385,6 +453,9 @@ def unsupported_reason(llm: nn.Module) -> Optional[str]: bound = draft_model.__dict__.get("dflash_forward") if getattr(bound, "__func__", None) is not _draft.windowed_block_decode: return "the drafter does not run the windowed block decode" + capability = device_capability() + if capability not in MERGE_ENTRY_ROWS_BY_SM or capability not in SILU_ENTRY_ROWS_BY_SM: + return f"no merge / silu trigger widths for SM {capability}" reason = _conv_reason(draft_model) if reason is not None: return f"pdl_conv: {reason}" @@ -435,6 +506,10 @@ def installed(llm: nn.Module) -> frozenset: return frozenset(found) +def _widths(widths) -> str: + return ",".join(str(w) for w in sorted(widths)) or "none" + + def install(llm: nn.Module, mapping=None) -> frozenset: """Install the five PDL chain launches; return the names this call installed. @@ -458,12 +533,18 @@ def install(llm: nn.Module, mapping=None) -> frozenset: raise ValueError(reason) before = installed(llm) draft_model = llm.draft_model + widths = entry_rows(draft_model) for conv in list(draft_model.attention_convs) + list(draft_model.mlp_convs): conv._pdl_conv = True + conv._pdl_entry_rows = widths + merge_rows = widths & trigger_rows(MERGE_ENTRY_ROWS_BY_SM) + silu_rows = widths & trigger_rows(SILU_ENTRY_ROWS_BY_SM) draft_model._pdl_merge = True + draft_model._pdl_merge_rows = merge_rows for layer in draft_model.model.layers: if type(layer.mlp) is GatedMLP: layer.mlp.__class__ = EntryTriggerGatedMLP + layer.mlp._pdl_entry_rows = silu_rows llm.spec_worker._pdl_head = True projection = draft_model.candidate_selector.hidden_projection if type(projection) is nn.Linear: @@ -476,6 +557,8 @@ def install(llm: nn.Module, mapping=None) -> frozenset: if changed: logger.info( f"[musespark_open] {','.join(sorted(changed))}: drafter PDL chain installed " - f"(in place: {','.join(sorted(now))})" + f"(in place: {','.join(sorted(now))}; block rows with the conv trigger " + f"{_widths(widths)}, the silu trigger {_widths(silu_rows)}, " + f"the merge trigger {_widths(merge_rows)})" ) return changed diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/allreduce/drafter_hop.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/allreduce/drafter_hop.py index 27494e01484c..21de86a7dee9 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/allreduce/drafter_hop.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/allreduce/drafter_hop.py @@ -72,6 +72,31 @@ calls interleave freely, and a GEMM captured into a graph can store into a fixed area. +Wide workspace (``max_rows`` above 8, fast path ``wide_drafter_hops``): the same kernel on +drafter blocks of up to ``WIDE_MAX_ROWS`` rows (several requests' blocks; the convolution's +tap-1 mask is per block, ``row % block``), with one slab per rank of the group and ``areas`` +receive areas used in site order modulo ``areas`` -- a divisor of ``sites`` of at least 3, so +every window of ``areas`` consecutive calls (across forwards too) names distinct areas and +the reset argument above holds with ``areas`` in place of ``sites``. A call resets the rows of +the previous area that the previous call can have written: all ``max_rows`` at site 0 (the +previous forward may have had any row count), at every other site the larger of its own row +count and the previous call's (a forward's calls share one row count when served). + +Progress: a wide call launches ``CLUSTER * rows`` CTAs of 128 threads. A cluster publishes its +row before it polls, and polls only the slots of rows ``t`` and ``t - 1`` of its column slice, +which only the clusters of rows ``t`` and ``t - 1`` of the peers write. No other compute +stream of the step is live during the drafter forward (the ``context_kv_fork`` side stream +joins the main stream before the drafter's first hop), so no kernel of this rank holds an SM +the hop's clusters wait for. Up to 32 rows (32 clusters = 256 CTAs) every cluster of the call +is resident together on a GB300 (at most two 128-thread CTAs per SM), so progress needs no +dispatch order. Above that (sm103's served 56 rows are 448 CTAs on 152 SMs, sm107's served 64 +rows are 512 CTAs on 212 SMs, and the qualification tests' 64..256 rows) not every cluster +need be resident at once, and progress +also relies on CTAs dispatched in grid order on every rank (the assumption the built-in MNNVL +all-reduce kernels and ``post_attn_hop_wide`` make, which CUDA does not guarantee): the +lowest-indexed unfinished cluster is then resident on every rank once every lower one has +finished, so it publishes and finishes. + Ordering contract: calls on one workspace must be issued in site order ``0, 1, ..., sites - 1, 0, ...`` on every rank, serialized on one stream (graph replay included), and every kernel launched with programmatic dependent launch @@ -109,6 +134,33 @@ HIDDEN = 8192 #: Rows of a drafter block the transport is sized for (K + 1 at K = 7). MAX_ROWS = 8 +#: Rows of a wide workspace: the drafter blocks of up to 32 requests at K = 7. +WIDE_MAX_ROWS = 256 +#: Rows of the served wide workspace (``wide_drafter_hops``) per compute capability +#: (``wide_served_rows``). sm103: the drafter blocks of 2..7 requests at K = 7, where the +#: one-shot hop is faster than the built-in all-reduce, convolution and add-RMSNorm on +#: GB300 (at 8 requests the two tie; above, the built-in two-shot all-reduce wins). sm107: +#: the blocks of 2..8 requests; on Rubin the one-shot hop is still the faster one at 64 rows +#: and the slower one from 128 rows on. +WIDE_SERVED_ROWS_BY_SM: dict[tuple[int, int], int] = {(10, 3): 56, (10, 7): 64} + + +def wide_served_rows(capability: Optional[Tuple[int, int]] = None) -> int: + """The served wide workspace's rows on ``capability`` (this process's device when None). + + Raises: + ValueError: on a device ``WIDE_SERVED_ROWS_BY_SM`` does not list. + """ + key = tuple(torch.cuda.get_device_capability() if capability is None else capability) + rows = WIDE_SERVED_ROWS_BY_SM.get(key) + if rows is None: + raise ValueError( + f"drafter_hop: no served wide rows for SM {key} " + f"(supported: {sorted(WIDE_SERVED_ROWS_BY_SM)})" + ) + return rows + + #: Largest tensor-parallel group; slots are laid out for it whatever the group size. MAX_RANKS = 8 #: BF16 values per thread. @@ -293,6 +345,24 @@ def _stale_bit(w0, w1, w2, w3, bit, *, loc=None, ip=None): ) return Uint32(result) + @dsl_user_op + def _globaltimer(*, loc=None, ip=None): + """The global nanosecond timer (bounded-poll variant only).""" + result = _asm(T.i64(), [], "mov.u64 $0, %globaltimer;", "=l", loc, ip) + return Int64(result) + + @dsl_user_op + def _trap_if(expired, *, loc=None, ip=None): + """Abort this rank's context when ``expired`` is non-zero (bounded-poll variant only).""" + _asm( + None, + [_u32(expired, loc, ip)], + "{\n.reg .pred p;\nsetp.ne.u32 p, $0, 0;\n@p trap;\n}", + "r", + loc, + ip, + ) + # ---- arithmetic: one PTX instruction each, as the replaced kernels run them ---- def _f32_asm(text, operands, loc, ip): @@ -424,12 +494,27 @@ def _pack_bf16x2(low, high, *, loc=None, ip=None): class _DrafterHopKernel: """All-reduce + DFlash 2 finish convolution + flashinfer add + RMSNorm, one row per cluster.""" - def __init__(self, ranks: int, eps: float, pdl: bool, early_trigger: bool, direct: bool): + def __init__( + self, + ranks: int, + eps: float, + pdl: bool, + early_trigger: bool, + direct: bool, + max_rows: int = MAX_ROWS, + timeout_ns: int = 0, + ): self.ranks = ranks self.eps = float(eps) self.pdl = pdl self.early_trigger = early_trigger self.direct = direct + # Rows of one slab (one source rank's rows of an area) and its bytes. + self.max_rows = int(max_rows) + self.slab_bytes = self.max_rows * ROW_BYTES + # 0: the polls spin without a bound (served); else a poll still stale after this + # many nanoseconds executes ``trap`` (qualification builds). + self.timeout_ns = int(timeout_ns) @cute.jit def __call__( @@ -449,6 +534,7 @@ def __call__( delta_row: Int32, delta_tap: Int32, dry: Int32, + reset_rows: Int32, stream: cuda.CUstream, ): self.kernel( @@ -467,6 +553,7 @@ def __call__( delta_row, delta_tap, dry, + reset_rows, ).launch( grid=[rows * CLUSTER, 1, 1], block=[THREADS, 1, 1], @@ -494,6 +581,7 @@ def kernel( delta_row: Int32, delta_tap: Int32, dry: Int32, + reset_rows: Int32, ): tid = cute.arch.thread_idx()[0] lane = tid % 32 @@ -548,7 +636,7 @@ def kernel( # Publish this rank's partial row into its slot on every rank # (one store to the multicast address reaches every member). _store_x4( - mcast + Int64(rank) * SLAB_BYTES + slot_offset, + mcast + Int64(rank) * self.slab_bytes + slot_offset, own[0], own[1], own[2], @@ -556,20 +644,36 @@ def kernel( ) # Retire the previous site's area: every row, every slot this # rank polls (rows t, t + rows, ... of this CTA's columns). - for i in cutlass.range_constexpr(MAX_ROWS): - reset_row = row + i * rows - in_area = Uint32(reset_row < MAX_ROWS) - for p in cutlass.range_constexpr(self.ranks): - reset_on = Uint32(1) - if cutlass.const_expr(not self.direct): - reset_on = Uint32(rank != p) - _reset_x4( - previous - + Int64(p) * SLAB_BYTES - + Int64(reset_row) * ROW_BYTES - + slot_bytes, - in_area & reset_on, - ) + if cutlass.const_expr(self.max_rows == MAX_ROWS): + for i in cutlass.range_constexpr(MAX_ROWS): + reset_row = row + i * rows + in_area = Uint32(reset_row < MAX_ROWS) + for p in cutlass.range_constexpr(self.ranks): + reset_on = Uint32(1) + if cutlass.const_expr(not self.direct): + reset_on = Uint32(rank != p) + _reset_x4( + previous + + Int64(p) * self.slab_bytes + + Int64(reset_row) * ROW_BYTES + + slot_bytes, + in_area & reset_on, + ) + else: + # A wide area: the rows the previous call can have written + # (``reset_rows``), rows t, t + rows, ... of this CTA's columns. + for reset_row in cutlass.range(row, reset_rows, rows, unroll=1): + for p in cutlass.range_constexpr(self.ranks): + reset_on = Uint32(1) + if cutlass.const_expr(not self.direct): + reset_on = Uint32(rank != p) + _reset_x4( + previous + + Int64(p) * self.slab_bytes + + Int64(reset_row) * ROW_BYTES + + slot_bytes, + reset_on, + ) r = _load_x4(residual + row_offset) group = Int64(column // GROUP_SIZE) delta_address = delta + (Int64(row) * Int64(delta_row) + group) * 2 @@ -592,7 +696,7 @@ def kernel( on = on & has_prev mine = own if half == 0 else own_prev v0, v1, v2, v3 = _poll_x4( - receive_row + Int64(p) * SLAB_BYTES - half * ROW_BYTES, + receive_row + Int64(p) * self.slab_bytes - half * ROW_BYTES, on, mine[0], mine[1], @@ -605,13 +709,16 @@ def kernel( values[slot, 3] = v3 enabled[slot] = on stale = stale | _stale_bit(v0, v1, v2, v3, on << Uint32(slot)) + start = Int64(0) + if cutlass.const_expr(self.timeout_ns > 0): + start = _globaltimer() while stale != Uint32(0): polled = stale for half in cutlass.range_constexpr(2): for p in cutlass.range_constexpr(self.ranks): slot = half * self.ranks + p v0, v1, v2, v3 = _poll_x4( - receive_row + Int64(p) * SLAB_BYTES - half * ROW_BYTES, + receive_row + Int64(p) * self.slab_bytes - half * ROW_BYTES, (polled >> Uint32(slot)) & Uint32(1), values[slot, 0], values[slot, 1], @@ -631,6 +738,9 @@ def kernel( values[slot, 3], enabled[slot] << Uint32(slot), ) + if cutlass.const_expr(self.timeout_ns > 0): + elapsed = _globaltimer() - start + _trap_if(Uint32((stale != Uint32(0)) & (elapsed > Int64(self.timeout_ns)))) # Early trigger: past the last transport access. if cutlass.const_expr(self.pdl and self.early_trigger): cute.arch.griddepcontrol_launch_dependents() @@ -718,9 +828,16 @@ def kernel( @functools.lru_cache(maxsize=64) def _compiled_kernel( - device_index: int, ranks: int, eps: float, pdl: bool, early_trigger: bool, direct: bool + device_index: int, + ranks: int, + eps: float, + pdl: bool, + early_trigger: bool, + direct: bool, + max_rows: int = MAX_ROWS, + timeout_ns: int = 0, ): - kernel = _DrafterHopKernel(ranks, eps, pdl, early_trigger, direct) + kernel = _DrafterHopKernel(ranks, eps, pdl, early_trigger, direct, max_rows, timeout_ns) with torch.cuda.device(device_index): stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) return cute.compile( @@ -732,6 +849,7 @@ def _compiled_kernel( Int32(0), Int32(0), Int32(1), + Int32(max_rows), stream, options="--enable-tvm-ffi", ) @@ -759,6 +877,13 @@ class DrafterHopWorkspace: use; construction checks that each is on this device, contiguous and 16-byte aligned. + ``max_rows`` above ``MAX_ROWS`` builds a wide workspace (module docstring): one slab per + rank of the group, ``areas`` receive areas (default: the smallest divisor of ``sites`` of + at least 3), multicast publication only. + + ``timeout_ns`` > 0 builds the bounded-poll variant for qualification: a poll still stale + after that long executes ``trap`` (aborting this rank's context) instead of spinning. + Construction is collective over the group (the ``pre_wait`` check, fabric allocation, sentinel fill, compilation and a dry launch per ``eps``) and either succeeds on every rank or raises on every rank. @@ -773,6 +898,9 @@ def __init__( early_trigger: bool = False, direct: bool = False, pre_wait: Sequence[torch.Tensor] = (), + max_rows: int = MAX_ROWS, + areas: Optional[int] = None, + timeout_ns: int = 0, ): from tensorrt_llm._torch.distributed.ops import ( _get_mnnvl_workspace_comm, @@ -791,10 +919,33 @@ def __init__( raise ValueError(f"the fused drafter hop serves 4 or 8 ranks, got {self.ranks}") if self.sites < MIN_SITES: raise ValueError(f"the fused drafter hop needs at least {MIN_SITES} sites") + self.max_rows = int(max_rows) + self.wide = self.max_rows > MAX_ROWS + self.timeout_ns = int(timeout_ns) + #: Rows of the last call issued (the reset bound of the next one inside a forward). + self._last_rows = self.max_rows + if self.wide: + if self.direct: + raise ValueError( + "a wide drafter hop workspace publishes through the multicast only" + ) + if not MAX_ROWS < self.max_rows <= WIDE_MAX_ROWS: + raise ValueError(f"a wide workspace takes {MAX_ROWS + 1}..{WIDE_MAX_ROWS} rows") + if areas is None: + areas = next(a for a in range(MIN_SITES, self.sites + 1) if self.sites % a == 0) + if areas < MIN_SITES or self.sites % areas: + raise ValueError( + f"areas must divide the {self.sites} sites and be at least {MIN_SITES}" + ) + self.areas = int(areas) + self.area_bytes = self.ranks * self.max_rows * ROW_BYTES + else: + self.areas = self.sites + self.area_bytes = AREA_BYTES self.device = torch.device("cuda", _mnnvl_device_index(mapping)) self._kernels = {} self._comm = _get_mnnvl_workspace_comm(mapping) - size = self.sites * AREA_BYTES + size = self.areas * self.area_bytes error: Optional[BaseException] = None try: if not HAS_CUTE_DSL: @@ -879,10 +1030,17 @@ def _kernel(self, eps: float): if torch.cuda.is_current_stream_capturing(): raise RuntimeError(f"drafter hop variant eps={eps} was not built before graph capture") compiled = _compiled_kernel( - self.device.index, self.ranks, eps, self.pdl, self.early_trigger, self.direct + self.device.index, + self.ranks, + eps, + self.pdl, + self.early_trigger, + self.direct, + self.max_rows, + self.timeout_ns, ) dummy = self._uc_base - compiled(*([dummy] * 9), 1, self.rank, 8, 0, 0, 1, self._stream()) + compiled(*([dummy] * 9), 1, self.rank, 8, 0, 0, 1, self.max_rows, self._stream()) self._kernels[eps] = compiled return compiled @@ -890,7 +1048,11 @@ def _stream(self): return cuda.CUstream(torch.cuda.current_stream(self.device).cuda_stream) def _area(self, base: int, site: int) -> int: - return base + (site % self.sites) * AREA_BYTES + return base + (site % self.areas) * self.area_bytes + + def nbytes(self) -> int: + """Fabric memory this rank holds for the transport.""" + return self.areas * self.area_bytes def direct_output(self, site: int, rows: int) -> torch.Tensor: """This rank's slot of ``site``'s area through the multicast address, ``[rows, HIDDEN]`` BF16. @@ -899,7 +1061,7 @@ def direct_output(self, site: int, rows: int) -> torch.Tensor: output here (every element once, its final value), and the hop then reads it from every rank's receive area. """ - if not self.direct: + if not self.direct or self.wide: raise RuntimeError("direct_output is only for a direct (poll-only) workspace") start = (site % self.sites) * AREA_BYTES + self.rank * SLAB_BYTES flat = self._mc.view(torch.bfloat16) @@ -928,6 +1090,11 @@ def launch( ) compiled = self._kernel(float(eps)) output = torch.empty_like(residual) + # The rows of the previous area its call can have written: any row count before the + # first site of a forward (the last forward that ran may not be the last one issued + # here: a capture that failed, a graph replayed after eager forwards), else the + # previous call's (wide areas only). + reset_rows = self.max_rows if site % self.sites == 0 else max(rows, self._last_rows) compiled( 0 if partial is None else partial.data_ptr(), residual.data_ptr(), @@ -944,9 +1111,11 @@ def launch( delta.stride(0), delta.stride(1), 0, + reset_rows, self._stream(), ) self.next_site = (site + 1) % self.sites + self._last_rows = rows return output @@ -957,6 +1126,7 @@ def hop_takes( base: torch.Tensor, weight: torch.Tensor, block_size: int, + max_rows: int = MAX_ROWS, ) -> bool: """Whether the hop computes this call: rows, dtypes, shapes and devices only. @@ -964,7 +1134,7 @@ def hop_takes( (``rows`` is the drafter block's row count), which is what lets a caller choose between this collective and the built-in all-reduce. """ - if not 1 <= rows <= MAX_ROWS or block_size < 1: + if not 1 <= rows <= max_rows or block_size < 1: return False if any(t.dtype != torch.bfloat16 for t in (residual, delta, base, weight)): return False @@ -1038,11 +1208,11 @@ def __call__( workspace: DrafterHopWorkspace, ) -> torch.Tensor: rows = residual.shape[0] if residual.dim() == 2 else 0 - if not hop_takes(rows, residual, delta, base, weight, block_size) or not hop_layout_ok( - partial, residual, delta, base, weight - ): + if not hop_takes( + rows, residual, delta, base, weight, block_size, workspace.max_rows + ) or not hop_layout_ok(partial, residual, delta, base, weight): raise ValueError( - "drafter_hop takes BF16 [1..8, 8192] partial / residual, [rows, 2, 512] " + f"drafter_hop takes BF16 [1..{workspace.max_rows}, 8192] partial / residual, [rows, 2, 512] " "coefficients with unit inner stride, a contiguous [2, 8192] base kernel and an " "8192 weight, 16-byte aligned" ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/gemm/fixed_low_m_gemm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/gemm/fixed_low_m_gemm.py index 577e8ffb338c..f5a9c02a4de5 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/gemm/fixed_low_m_gemm.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/gemm/fixed_low_m_gemm.py @@ -58,7 +58,10 @@ #: coefficient projection is the tep8 one. The fastest split-K tactic of each shape in #: chains of cold weights with PDL on; it beats cuBLAS on the qkv, gate_up and #: convolution shapes, ties it on down_proj and trails it slightly on o_proj in -#: isolation, where ``drafter_pdl_chain``'s early weight load is what it is for. +#: isolation, where ``drafter_pdl_chain``'s early weight load is what it is for. Besides the +#: one-request block (8 rows) the table lists the blocks of 2..4 requests (16, 24 and 32 +#: rows), each the fastest split-K tactic at that row count on GB300's 152 SMs over cold +#: weights with PDL on; every one is a single wave of at most 128 CTAs. SPLITK_TACTICS_TP4: dict[tuple[int, int, int], tuple[int, int, int, int]] = { # DFlash 2 convolution coefficient projection (replicated). (8, 8192, 2048): (64, 8, 4, 10), @@ -70,16 +73,48 @@ (8, 7168, 8192): (64, 8, 1, 11), # drafter gate_up_proj per rank. (8, 8192, 14336): (128, 8, 1, 6), + (16, 8192, 2048): (64, 16, 4, 8), + (16, 8192, 2560): (64, 16, 2, 10), + (16, 2048, 8192): (64, 16, 1, 10), + (16, 7168, 8192): (64, 16, 1, 10), + (16, 8192, 14336): (128, 16, 1, 6), + (24, 8192, 2048): (64, 16, 2, 11), + (24, 8192, 2560): (64, 32, 2, 9), + (24, 2048, 8192): (64, 32, 1, 9), + (24, 7168, 8192): (64, 32, 1, 9), + (24, 8192, 14336): (128, 32, 1, 5), + (32, 8192, 2048): (64, 16, 2, 11), + (32, 8192, 2560): (64, 32, 2, 9), + (32, 2048, 8192): (64, 32, 1, 9), + (32, 7168, 8192): (64, 32, 1, 9), + (32, 8192, 14336): (128, 32, 1, 5), } #: The tep4 drafter shapes on sm107 (212 SMs), each the fastest split-K tactic there in -#: the same chains; every one beats cuBLAS at its shape. +#: the same chains; every one beats cuBLAS at its shape. Besides the one-request block +#: (8 rows) the table lists the blocks of 2..4 requests (16, 24 and 32 rows), each the +#: fastest split-K tactic at that row count over cold weights with PDL on. SM107_SPLITK_TACTICS_TP4: dict[tuple[int, int, int], tuple[int, int, int, int]] = { (8, 8192, 2048): (64, 8, 4, 10), (8, 8192, 2560): (64, 8, 4, 12), (8, 2048, 8192): (128, 16, 2, 6), (8, 7168, 8192): (64, 8, 2, 10), (8, 8192, 14336): (64, 8, 2, 11), + (16, 8192, 2048): (64, 16, 4, 10), + (16, 8192, 2560): (64, 16, 4, 10), + (16, 2048, 8192): (64, 16, 1, 11), + (16, 7168, 8192): (64, 32, 1, 8), + (16, 8192, 14336): (128, 16, 1, 6), + (24, 8192, 2048): (64, 32, 4, 8), + (24, 8192, 2560): (64, 32, 4, 8), + (24, 2048, 8192): (64, 32, 1, 9), + (24, 7168, 8192): (64, 32, 1, 9), + (24, 8192, 14336): (128, 32, 1, 5), + (32, 8192, 2048): (64, 32, 4, 8), + (32, 8192, 2560): (64, 32, 4, 8), + (32, 2048, 8192): (64, 32, 1, 9), + (32, 7168, 8192): (64, 32, 2, 9), + (32, 8192, 14336): (64, 32, 2, 8), } #: The measured table of each ``(compute capability, drafter TP size)``. From a79699f8841fc064c880809701ade8b5386c1db7 Mon Sep 17 00:00:00 2001 From: Xiao Wang <24860335+xwang233@users.noreply.github.com> Date: Mon, 26 Oct 2026 18:44:17 +0000 Subject: [PATCH 12/16] [None][feat] modeling_v2 musespark_open tep: compose the new fast paths on GB300 and Rubin - Dispatch registers mixed_step_gen_rows, model_boundary, wide_drafter_hops and vocab_parallel_head with their dependencies. - GB300 (sm103) TEP4 composes vocab_parallel_head and wide_drafter_hops, and decode_kv_splits installs its mixed-step mixin. - Rubin (sm107) TEP4 composes mixed_step_gen_rows and wide_drafter_hops; load_weights passes the device's compute capability, so the CuTe DSL prefill experts install on sm107 at tp 4 only. - Both GPUs, TEP4 and TEP8, compose model_boundary (bitwise). - DecodeTarget runs the model_boundary input kernel and arms its BF16 last-layer join; adopt_model installs the mixed-step attention mixins after every other attention re-class; the drafter shell installs the wide hops and passes vocab_parallel to the target head. Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com> --- .../dispatch.py | 27 +++++++ .../drafter.py | 9 ++- .../layers.py | 21 +++++- .../modeling.py | 71 ++++++++++++++----- .../weights.py | 24 ++++++- 5 files changed, 131 insertions(+), 21 deletions(-) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/dispatch.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/dispatch.py index 3151233d545b..751757a6d5f4 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/dispatch.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/dispatch.py @@ -131,6 +131,13 @@ class FastPath: "multi-request decode FMHA over an over-subscribed KV split, 72..512 rows", "construct", ), + FastPath( + "mixed_step_gen_rows", + "full-attention generation rows of mixed steps with 9..64 generation requests on " + "multi_request_split_kv, context rows on the thop op", + "construct", + ("multi_request_split_kv",), + ), # Dense projections and prefetch. FastPath("dense_gemms", "dense projections on one named kernel per shape", "construct"), FastPath( @@ -305,6 +312,13 @@ class FastPath: ("persistent_experts", "moe_prologue"), ), FastPath("replicated_embedding", "token embedding gathered from a replicated table", "load"), + FastPath( + "model_boundary", + "decode input boundary (gather, embed norm, residual, layer 0 pre-QKV norm) in one launch; " + "last-layer join storing BF16", + "load", + ("replicated_embedding",), + ), # Drafter. FastPath("drafter_block_conv", "drafter block convolutions as one fused kernel", "draft"), FastPath( @@ -358,6 +372,13 @@ class FastPath: "draft", ("drafter_block_decode",), ), + FastPath( + "wide_drafter_hops", + "the drafter all-reduce hops of 9..56-row (sm107: 9..64-row) block decodes on a wide " + "one-shot transport", + "draft", + ("drafter_hops",), + ), FastPath( "drafter_pdl_chain", "drafter and draft head launches releasing their GEMMs early", @@ -383,6 +404,12 @@ class FastPath: "draft", ("split_sampler",), ), + FastPath( + "vocab_parallel_head", + "the verify head of 16..1024-row decode steps on each rank's own vocabulary shard (no logits all-gather)", + "draft", + ("split_sampler",), + ), # Executor. FastPath("deferred_input_prep", "a decode step's input writes as one launch", "executor"), FastPath("sampler_store", "the sampler's store update and snapshots as one launch", "executor"), diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/drafter.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/drafter.py index ac97cd478efa..f159f2459851 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/drafter.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/drafter.py @@ -71,6 +71,8 @@ def adopt_drafter(shell, plan: "_dispatch.Plan") -> None: from .kernels.drafter_hops import hops as _drafter_hops _drafter_hops.install(getattr(shell.llm, "draft_model", None), shell._mapping) + if plan.enabled("wide_drafter_hops"): + _drafter_hops.install_wide(getattr(shell.llm, "draft_model", None), shell._mapping) if plan.enabled("drafter_pdl_chain"): # After the drafter_fixed_gemms GEMM methods and the drafter_hops hops: # the entry triggers rely on both. @@ -82,7 +84,12 @@ def adopt_drafter(shell, plan: "_dispatch.Plan") -> None: if plan.enabled("split_sampler"): from .kernels.split_sampler import target_head as _target_head - _target_head.install(shell.llm, shell._mapping, plan.enabled("lm_head_inbox")) + _target_head.install( + shell.llm, + shell._mapping, + plan.enabled("lm_head_inbox"), + plan.enabled("vocab_parallel_head"), + ) if plan.enabled("wide_split_sampler"): _target_head.widen(shell.llm) # Last: the arena compaction derives its worker class from whatever class every diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/layers.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/layers.py index 41cf773307e6..bdedd1c71ca7 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/layers.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/layers.py @@ -1053,8 +1053,9 @@ def adopt_model( full_attention.install_multi(model) decode_kv_splits = None if plan.enabled("decode_kv_splits"): - # Re-classes nothing: DecodeTarget arms it around a forward of 72..512 rows. - # Its refusal is agreed over the TP group, as the other attention installs. + # DecodeTarget arms it around a decode forward of 72..512 rows; the mixed-step + # mixin is installed after every other attention re-class (below). Its refusal is + # agreed over the TP group, as the other attention installs. from .kernels.decode_kv_splits import splits as _decode_kv_splits decode_kv_splits = _tp_agreement.agreed_install( @@ -1094,6 +1095,13 @@ def adopt_model( _tp_agreement.agreed_install( model.layers[0].mapping, lambda: multi_request_append.install(model, min_seqs) ) + if decode_kv_splits is not None: + # After every other attention re-class: the mixed-step mixin goes in front of all of + # them on the full-attention layers and arms the split rule only on mixed steps, which + # none of the others intercept. + from .kernels.decode_kv_splits import mixed as _mixed_kv_splits + + _mixed_kv_splits.install(model, decode_kv_splits) if preattn_prep_on: logger.info( f"[musespark_open] preattn_prep: fused pre-attention prep on the decode path of all " @@ -1109,6 +1117,15 @@ def adopt_model( swa_attention_block.log_installed(model.layers, swa_block) if plan.enabled("wide_swa_block"): swa_attention_block.install_multi(model) + if plan.enabled("mixed_step_gen_rows"): + # After every other attention re-class and after select_decode_fmha: the mixin goes + # in front of all of them and acts only on mixed steps, which none of the others + # intercept. Agreed like multi_request_split_kv's. + from .kernels.mixed_step_gen_rows import attention as mixed_step_gen_rows + + _tp_agreement.agreed_install( + model.layers[0].mapping, lambda: mixed_step_gen_rows.install(model) + ) layer_kwargs = {} if plan.enabled("projection_gathers"): # Imported only when planned: with no fast path the target imports diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/modeling.py index 63bcc3295d14..28bd758f540e 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/modeling.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/modeling.py @@ -116,16 +116,30 @@ class acts only on sliding-window layers): their decode prep also writes the decode steps run it, while prefill and mixed steps keep the built-in selection (the ``thop`` attention). The drafter's attention is not a target layer and keeps the built-in selection. -* ``decode_kv_splits`` (tep4 on sm_103 only) re-classes nothing: ``DecodeTarget`` +* ``decode_kv_splits`` (tep4 on sm_103 only): ``DecodeTarget`` arms it for a forward of 72..512 rows (9 to 64 requests), and while armed the full-attention layers' trtllm-gen decode FMHA calls present flashinfer's KV-split selection with the SM count of a grid of about ``TARGET_CTAS_PER_SM`` CTAs per SM (``kernels/decode_kv_splits``), so those - steps keep multi-CTA KV at every batch size. Decode steps of up to 8 - requests and prefill steps run what they ran before. Where + steps keep multi-CTA KV at every batch size. It also re-classes the + full-attention layers' attention with ``MixedStepKvSplitsAttention`` in front + of all of the above (``kernels/decode_kv_splits/mixed``): an attention call of + a mixed step with 9 to 64 generation requests runs with the same rule armed, + so the generation phase's decode FMHA of those steps keeps multi-CTA KV too. + Decode steps of up to 8 requests, mixed steps of up to 8 generation requests + and pure prefill steps run what they ran before. Where ``multi_request_split_kv`` serves a full-attention layer, its stand-in sits inside this one and launches its own kernel, so this split only acts on the calls that path declines. +* ``mixed_step_gen_rows`` (tep4 on sm_107 only) takes the generation rows of + mixed steps with 9..64 generation requests off the thop attention on the + full-attention layers: a ``ThopContextPhaseFmha`` before each such layer's + thop fallback offers the context phase, so those steps run + ``CombinedFmha(context=thop, generation=DecodeTrtllmGenFmha)``, and + ``MixedStepGenRowsAttention`` in front of the layer's attention intercepts + the generation phase's decode call with ``multi_request_split_kv``'s + stand-in (``kernels/mixed_step_gen_rows``). Sliding-window layers, pure + context steps and mixed steps outside the band keep the thop attention. Dense projections and prefetch: @@ -380,7 +394,9 @@ class acts only on sliding-window layers): their decode prep also writes the write into a multicast inbox that the sampler gathers from (and fires the context-K/V fork's head fork point right after that GEMM); ``wide_split_sampler`` (tep4 on sm_103 only) widens the head to decode steps of 129..1024 rows, which the -sampler takes on the shipped sampler's 512-thread block (``target_head.widen``). +sampler takes on the shipped sampler's 512-thread block (``target_head.widen``); +``vocab_parallel_head`` (tep4 on sm_103 only) samples and argmaxes decode steps of +16 rows or more from each rank's own vocabulary shard, without the logits all-gather. """ from __future__ import annotations @@ -734,12 +750,18 @@ def forward( self.qkv_prefetch.graph_start( int((input_ids if input_ids is not None else inputs_embeds).shape[0]) ) - if inputs_embeds is None: - inputs_embeds = core.embed_tokens(input_ids) - embeds = core.embed_norm(inputs_embeds.to(core.dtype)) - if core.m_emb != 1.0: - embeds = embeds * core.m_emb - hidden_states = embeds.to(torch.float32) + # Installed at weight load, after this target was built. + model_boundary = getattr(core, "_model_boundary", None) + first_attn_input = None + if model_boundary is not None and model_boundary.takes(input_ids, inputs_embeds): + hidden_states, first_attn_input = model_boundary.input(input_ids) + else: + if inputs_embeds is None: + inputs_embeds = core.embed_tokens(input_ids) + embeds = core.embed_norm(inputs_embeds.to(core.dtype)) + if core.m_emb != 1.0: + embeds = embeds * core.m_emb + hidden_states = embeds.to(torch.float32) if self.preattn_prep is not None: kwargs = dict(kwargs, preattn_prep=self.preattn_prep) if self.decode_attention is not None: @@ -787,7 +809,7 @@ def forward( batch_armed = batch is not None and batch.takes(rows) layers = core.layers - precomputed_attn_input = None + precomputed_attn_input = first_attn_input if hop_armed: arming = hop.armed(hidden_states.shape[0], hidden_states) elif wide_hop_armed: @@ -806,6 +828,7 @@ def forward( two_shot.armed() if two_shot_armed else contextlib.nullcontext(), hot.armed() if hot_armed else contextlib.nullcontext(), batch.armed() if batch_armed else contextlib.nullcontext(), + model_boundary.armed() if model_boundary is not None else contextlib.nullcontext(), ): for index, decoder_layer in enumerate(layers): next_norm = layers[index + 1].pre_qkv_proj_norm if index + 1 < len(layers) else None @@ -858,7 +881,10 @@ def forward( #: prefetch and PDL norm-gate of 33..512-row decode steps, the drafter window attention of #: 33..1024-row block decodes on a multi-request split geometry, the draft head's rank top-k #: key exchange on draft steps of 33 or more rows, the split verify sampler on -#: 129..1024-row decode steps, and the context steps' pre-attention prep in one launch. +#: 129..1024-row decode steps, the context steps' pre-attention prep in one launch, the +#: verify head of 16..1024-row decode steps on each rank's own vocabulary shard over the +#: tray's multicast, and the drafter hops of 9..56-row block decodes on a wide one-shot +#: transport. _SM103_TEP4_ONLY_PATHS = frozenset( { "hot_replicas", @@ -878,10 +904,18 @@ def forward( "wide_draft_topk", "wide_split_sampler", "context_prep", + "vocab_parallel_head", + "wide_drafter_hops", } ) +#: The fast paths validated at a tep4 rank's shape on sm107 only: the generation rows of +#: mixed steps on the multi-request decode FMHA (on sm103 flashinfer's trtllm-gen FMHA +#: already serves both phases of a mixed step). +_SM107_TEP4_ONLY_PATHS = frozenset({"mixed_step_gen_rows"}) #: TEP8: every fast path but the TP4-only ones of either architecture. -_TEP8_PATHS = _dispatch.FAST_PATH_NAMES - _TEP4_ONLY_PATHS - _SM103_TEP4_ONLY_PATHS +_TEP8_PATHS = ( + _dispatch.FAST_PATH_NAMES - _TEP4_ONLY_PATHS - _SM103_TEP4_ONLY_PATHS - _SM107_TEP4_ONLY_PATHS +) #: The fast paths written for any rank count whose output is bitwise the built-in #: model's. _TEP4_BITWISE_PATHS = frozenset( @@ -892,6 +926,7 @@ def forward( "deferred_input_prep", "sampler_store", "replicated_embedding", + "model_boundary", "fused_accept", "verify_fold", "draft_head", @@ -1023,7 +1058,7 @@ def forward( #: decode FMHA of steps of five or more requests (``multi_request_split_kv``; steps of two to #: four requests stay on ``wide_full_attention``). _SM107_TEP4_ATTENTION_PATHS = frozenset({"multi_request_split_kv"}) -_SM107_TEP4_PATHS = _SM107_TEP4_PATHS | _SM107_TEP4_ATTENTION_PATHS +_SM107_TEP4_PATHS = _SM107_TEP4_PATHS | _SM107_TEP4_ATTENTION_PATHS | _SM107_TEP4_ONLY_PATHS _SM107_TEP4_PATHS |= _SM107_TEP4_COMM_PATHS #: The GB300 TEP4 GEMM and expert-balance paths also served at sm107 TEP4, tuned for #: R200: the router GEMM of decode steps of five or more requests (33..512 rows; 16..32 @@ -1040,13 +1075,15 @@ def forward( #: the two-to-four-request steps' append themselves), the drafter window attention of #: 33..1024-row block decodes on the multi-request split geometry (sized by the device's #: SM count), the draft head's rank top-k key exchange on draft steps of 33 or more rows, -#: and the split verify sampler on 129..1024-row decode steps. +#: the split verify sampler on 129..1024-row decode steps, and the drafter hops of +#: 9..64-row block decodes on a wide one-shot transport. _SM107_TEP4_DRAFTER_AT_BATCH_PATHS = frozenset( { "multi_request_prep_append", "wide_drafter_window_attention", "wide_draft_topk", "wide_split_sampler", + "wide_drafter_hops", } ) _SM107_TEP4_PATHS = _SM107_TEP4_PATHS | _SM107_TEP4_DRAFTER_AT_BATCH_PATHS @@ -1162,7 +1199,9 @@ def load_weights(self, weights, weight_mapper=None): plan = self._plan if not plan: return result - _weights.install_loaded_features(self.llm.model, self.model_config.mapping, plan) + _weights.install_loaded_features( + self.llm.model, self.model_config.mapping, plan, device_compute_capability() + ) # The token embedding after the experts, so the expert replication's # all-gather and this one's never interleave. _weights.install_embedding_features(self.llm.model, self.model_config.mapping, plan) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/weights.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/weights.py index d9bb985aaaf1..9181f8d25191 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/weights.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/weights.py @@ -11,7 +11,7 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Any, Dict, Optional +from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple from tensorrt_llm._torch.models.modeling_musespark_open import MuseSparkOpenModel from tensorrt_llm._torch.models.modeling_musespark_open_omni import ( @@ -38,7 +38,10 @@ def load( def install_loaded_features( - text_model: MuseSparkOpenModel, mapping, plan: "_dispatch.Plan" + text_model: MuseSparkOpenModel, + mapping, + plan: "_dispatch.Plan", + compute_capability: Optional[Tuple[int, int]] = None, ) -> None: """Install the plan's expert paths on a text model whose weights are final. @@ -60,6 +63,10 @@ def install_loaded_features( Then ``hot_replicas`` (``hot_replicas.install``: hooks only, at tp 4; the copies are gathered on the first forward of a deployment that serves 33..512-row decode steps), + then, on Rubin (sm107) at tp 4 (``prefill_cutedsl_experts.serves`` of + ``compute_capability``, the device's; None installs it nowhere), + ``prefill_cutedsl_experts`` (``prefill_cutedsl_experts.install``: the + 4096..8192-row expert calls of the NVFP4 layers on the Rubin CuTe DSL GEMMs), then ``early_expert_ids`` (``early_ids.install``, on the layers just installed), then ``pool_into_route`` (``DirectPool.install_sites``, whose routed kernel variants carry the placement and the early expert ids), each collective over the @@ -76,6 +83,13 @@ def install_loaded_features( # After the owner-window install, whose expert and routing hooks it wraps. _hot_replicas.install(text_model, mapping) + if compute_capability is not None: + from .kernels.prefill_cutedsl_experts import install as _prefill_cutedsl + + if _prefill_cutedsl.serves(compute_capability, mapping.tp_size): + # After the expert paths above, whose expert forwards it wraps for the large + # prefill calls only. + _prefill_cutedsl.install(text_model, mapping) if plan.enabled("early_expert_ids"): _install_early_expert_ids(text_model, mapping) if plan.enabled("pool_into_route"): @@ -180,3 +194,9 @@ def install_embedding_features( from .kernels.replicated_embedding import embedding as _replicated_embedding _replicated_embedding.install(text_model, mapping) + if plan.enabled("model_boundary"): + # After replicated_embedding, whose table it gathers from. Collective (agreed refusal). + from .kernels.boundary_early_release import triggers as _pdl_triggers + from .kernels.model_boundary import boundary as _model_boundary + + _model_boundary.install(text_model, _pdl_triggers.capture_layers(text_model), mapping) From 869e5e4192e09e78256eb1bd05011dfedc0cc4fb Mon Sep 17 00:00:00 2001 From: Xiao Wang <24860335+xwang233@users.noreply.github.com> Date: Mon, 26 Oct 2026 18:44:28 +0000 Subject: [PATCH 13/16] [None][feat] modeling_v2 musespark_open tep: half expert window at sm103 TEP4 and the wide projection table by device - half_window_experts joins the GB300 (sm103) TEP4 composition, as at Rubin (sm107) TEP4: the NVFP4 expert layers also hold their ring predecessor's upper expert half. Decode steps of at most 8 rows are placed by the prologue's routed kernel, those of 9..32 rows (two to four requests) by ring_buddy_assign's half flow, both balancing the selected experts over the four ranks. The window costs 45 GiB of KV cache per rank, so its capacity gate carries it only up to max_batch_size 8, and a plan carrying it leaves hot_replicas out (dispatch.YIELDS). Compositions above max_batch_size 8 are unchanged. - wide_projection_gathers takes its 9..32-row kernel table from wide_kernel_steps on the dense GEMM plan's compute capability: the sm103 table, and a refusal on sm107 (whose TEP4 composition runs those projections on wide_dense_gemms' own table) and on any other device. No served composition changes: only sm103 TEP4 serves the path. Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com> --- .../wide_projection_gathers/gathers.py | 38 +++++++++++++++++-- .../modeling.py | 15 +++++++- 2 files changed, 47 insertions(+), 6 deletions(-) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/wide_projection_gathers/gathers.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/wide_projection_gathers/gathers.py index 6644e4f977d0..8989f463b03c 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/wide_projection_gathers/gathers.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/wide_projection_gathers/gathers.py @@ -80,6 +80,34 @@ "pool": ((WIDE_ROWS, "splitk:64,16,4,10"),), "unpool": ((16, "splitk:64,16,4,10"), (WIDE_ROWS, "splitk:64,16,2,11")), } + + +def wide_kernel_steps( + capability: Optional[Tuple[int, int]], +) -> Mapping[str, Tuple[Tuple[int, str], ...]]: + """The 9..32-row table of the device of ``capability`` (None: the target's own sm103, + as ``dense_gemms.device_capability`` reads a host without CUDA). + + Raises: + ValueError: on sm107, whose TEP4 composition serves its 9..32-row projections on + ``wide_dense_gemms``' R200 table instead (``dense_gemms`` ``SM107_TP4_WIDE_KERNELS``), + and on any other device. + """ + capability = (10, 3) if capability is None else tuple(capability) + if capability == (10, 3): + return TP4_WIDE_KERNELS + elif capability == (10, 7): + raise ValueError( + "[musespark_open] wide_projection_gathers has no sm107 table: sm107 TEP4 runs its " + "9..32-row projections on wide_dense_gemms" + ) + else: + raise ValueError( + f"[musespark_open] wide_projection_gathers has no table for SM {capability} " + "(measured on sm103 only)" + ) + + #: The per-rank ``(K, N)`` of each site at TP4. TP4_SHAPES: Mapping[str, Tuple[int, int]] = { "qkv": (8192, 4608), @@ -166,7 +194,8 @@ def check(model: torch.nn.Module, plan: Optional[DenseGemmPlan]) -> None: """Whether ``install`` may run on ``model`` with ``dense_gemms``' ``plan``; changes nothing. Raises: - ValueError: if the model is not TP4, a layer's MoE or post-FFN join is not + ValueError: if the plan's device has no 9..32-row table (``wide_kernel_steps``: sm103 + only), the model is not TP4, a layer's MoE or post-FFN join is not ``projection_gathers``' class with both sites fused, ``dense_gemms`` is not installed, or its plan prefetches (``dense_gemm_prefetch``): that prefetch's row exemptions (``skip_prefetch_rows``) are derived from the 1..8-row table, so @@ -200,7 +229,7 @@ def check(model: torch.nn.Module, plan: Optional[DenseGemmPlan]) -> None: f"[musespark_open] wide_projection_gathers widens projection_gathers' unpool site, " f"which layer {layer.layer_idx} does not run ({type(join).__name__})" ) - for site, by_rows in wide_kernels().items(): + for site, by_rows in wide_kernels(wide_kernel_steps(plan.capability)).items(): k, n = TP4_SHAPES[site] for rows, kernel in enumerate(by_rows, start=WIDE_MIN_ROWS): error = kernel_error(kernel, rows, k, n) @@ -226,7 +255,8 @@ def install(model: torch.nn.Module, plan: DenseGemmPlan) -> WideProjections: pool_inbox = ProjectionInbox( mapping, row_words=_pool.local_columns(world) // 2, max_rows=WIDE_ROWS ) - plan.set_wide_kernels(wide_kernels()) + kernels = wide_kernels(wide_kernel_steps(plan.capability)) + plan.set_wide_kernels(kernels) wide = WideProjections(plan) for layer in layers: moe, join = layer.feed_forward, layer.add_ffn @@ -239,7 +269,7 @@ def install(model: torch.nn.Module, plan: DenseGemmPlan) -> WideProjections: f"in {len(layers)} layers: pool/unpool slices with their gathers fused into the " "pre-expert quantization and the post-FFN boundary (32-row inboxes); split-K for " + ", ".join( - f"{site} {'|'.join(dict.fromkeys(by_rows))}" for site, by_rows in wide_kernels().items() + f"{site} {'|'.join(dict.fromkeys(by_rows))}" for site, by_rows in kernels.items() ) ) return wide diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/modeling.py index 28bd758f540e..3544e9f35e61 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/modeling.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/modeling.py @@ -319,11 +319,13 @@ class acts only on sliding-window layers): their decode prep also writes the ``ring_buddy.install(windows=1)``), with nothing replicated. * ``three_window_placement`` stacks a third ring window on them. * ``half_window_experts`` (tp 4) stacks the ring predecessor's upper expert half on - them, which the prologue's routed kernel places onto; only up to the + them, which the prologue's routed kernel places onto (and, on calls of 9..32 rows + that reach the routing method, ``ring_buddy_assign``'s half flow); only up to the ``max_batch_size`` of its capacity gate (``dispatch.CAPACITY_GATES``). * ``bf16_routed_experts`` gives the two BF16 layers' stacked calls unpacked routing (on the owner-only window as well). -* ``hot_replicas`` (tep4 on sm_103 only) hooks the NVFP4 layers' expert and routing +* ``hot_replicas`` (tep4, in a plan without ``half_window_experts``: + ``dispatch.YIELDS``) hooks the NVFP4 layers' expert and routing calls after the owner-window install; the copies themselves are gathered on the first forward of a deployment whose largest decode step (``max_batch_size`` times the verify width) reaches 33 rows (``HotReplicas.decide``), and ``DecodeTarget`` @@ -952,11 +954,20 @@ def forward( #: decode steps on the tray's MNNVL two-shot; on decode steps of 9..1024 rows, the #: post-attention hop on its own one-shot transport; large context steps #: sequence-parallel over the tray; and the context steps' pre-attention prep in one launch. +#: Its NVFP4 expert layers also hold their ring predecessor's upper expert half +#: (``half_window_experts``): decode steps of at most 8 rows are placed by the +#: prologue's routed kernel, those of 9..32 rows (two to four requests, behind the +#: multi-request router GEMM and projection gathers) by ``ring_buddy_assign``'s half +#: flow, both balancing the selected experts over the four ranks. As on sm107 the +#: window costs 45 GiB of KV cache per rank and is carried only up to +#: ``max_batch_size`` 8 (``dispatch.CAPACITY_GATES``); a plan carrying it leaves +#: ``hot_replicas`` out (``dispatch.YIELDS``). _SM103_TEP4_PATHS = ( _TEP4_PATHS | _SM103_TEP4_ONLY_PATHS | frozenset( { + "half_window_experts", "swa_paged_decode", "swa_block", "full_attention_split_kv", From 25525dc68ca68719a9896f33333600727175777a Mon Sep 17 00:00:00 2001 From: Xiao Wang <24860335+xwang233@users.noreply.github.com> Date: Tue, 27 Oct 2026 06:44:10 +0000 Subject: [PATCH 14/16] [None][fix] modeling_v2 musespark_open tep: apply the drafter logit transform after the sharded draft top-k The Muse Spark Open DFlash 2 drafter declares a logit transform (output multiplier and softcap), which the DFlash worker applies to the draft logits right after the drafter's logits processor and which returns fp32. The vocab-sharded draft head takes only bf16 shard logits, so with the transform on every draft step fell back to the built-in selector: a full-vocab fp32 all-gather, torch.topk over the full vocabulary and the full-vocab fill, scatter and argmax, and the path-walk tail and the drafter PDL chain were lost. The transform never decreases, so each position's top-k is the same before and after it, except where it maps distinct bf16 logits to the same fp32 value (tanh saturates), which torch.topk then orders by vocab id. The draft head now keeps the bf16 shard and applies the transform after its top-k: - kernels/draft_head/logit_transform.py: per drafter, a canonical-key table (bf16 radix key -> key of its tie group's representative) and a values table (bf16 radix key -> the drafter's own transform of that value), built on the device by DFlashForCausalLM.apply_draft_logit_transform itself and checked exhaustively over the 65536 bf16 values (fp32 result, NaN and signed zeros kept, never decreasing, groups contiguous, canonical values comparing as the transformed ones, layout independence on a host-seeded probe of the runtime shape). - draft_shard_topk / draft_rank_topk take the canonical table and its identity band (only keys outside the band load it); draft_topk_merge takes the values table for the unary values. Without the tables every kernel is the previous one. Each op's reference and is_valid follow the new arguments, and the ops declare cells for the transform-on configurations. - draft_head.install puts a deferring transform on the drafter instance: it hands back the very shard tensor ShardDraftLogitsProcessor just returned and records it; ShardedDraftHeadDFlashWorker._refine_block_logits keeps those logits bf16 only when the step takes the sharded selector (which then owes the transform and must apply it before sampling), and applies the drafter's transform to everything else, so sampled drafts, d2t, other layouts and any other tensor keep the built-in computation bit for bit. Deferral also stops when a class composed above overrides the refinement or the selector, or when the drafter's attributes no longer match the tables. selector_draft_tokens takes the transform as a required keyword; the verify-fold tail walk passes it too. - The install builds and checks everything before changing anything, runs the tie-order probe also on logits whose transformed values tie (from the transform's own tie groups, and at many requests), and agrees its refusal and a digest of the tables over the TP group. wide_draft_topk's probe runs with the transform as well. No change outside the target: an identity transform installs nothing new and runs the same kernels as before. Tests (tests/unittest/_torch/modeling_v2/models/musespark_open/...): the tables against the drafter's method on every bf16 value and in the runtime layout; the ops on their cells; the merge and the end-to-end draft tokens against torch.topk and the built-in selector over the transformed logits (1 to 24 requests, TP4 and TP8, all three key layouts, tie-heavy and saturating logit families); the fast path engaged with the transform on (no vocabulary all-gather, no full-vocab torch.topk), under the fold tail and in CUDA-graph replay; every fallback; install atomicity, idempotence and the probe's teeth; and a source pin on DFlashWorker._forward_impl's statements between the logits processor and _refine_block_logits. Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com> --- .../kernels/draft_head/head.py | 509 ++++++++++-- .../kernels/draft_head/logit_transform.py | 303 ++++++++ .../kernels/verify_fold/fold.py | 4 + .../kernels/wide_draft_topk/exchange.py | 122 +-- .../ops/spec/draft_rank_topk.py | 43 +- .../ops/spec/draft_shard_topk.py | 129 ++- .../ops/spec/draft_topk_merge.py | 115 ++- .../kernels/draft_head/_draft_head_cases.py | 328 ++++++++ .../test_modeling_v2_draft_head_transform.py | 734 ++++++++++++++++++ .../test_modeling_v2_draft_logit_transform.py | 164 ++++ .../test_modeling_v2_draft_transform_pin.py | 149 ++++ .../spec/test_modeling_v2_draft_rank_topk.py | 82 ++ .../spec/test_modeling_v2_draft_shard_topk.py | 100 +++ .../spec/test_modeling_v2_draft_topk_merge.py | 128 +++ 14 files changed, 2771 insertions(+), 139 deletions(-) create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/draft_head/logit_transform.py create mode 100644 tests/unittest/_torch/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/draft_head/_draft_head_cases.py create mode 100644 tests/unittest/_torch/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/draft_head/test_modeling_v2_draft_head_transform.py create mode 100644 tests/unittest/_torch/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/draft_head/test_modeling_v2_draft_logit_transform.py create mode 100644 tests/unittest/_torch/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/draft_head/test_modeling_v2_draft_transform_pin.py create mode 100644 tests/unittest/_torch/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/test_modeling_v2_draft_rank_topk.py create mode 100644 tests/unittest/_torch/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/test_modeling_v2_draft_shard_topk.py create mode 100644 tests/unittest/_torch/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/test_modeling_v2_draft_topk_merge.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/draft_head/head.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/draft_head/head.py index 0d9cc4b80b97..9f9be5c470fc 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/draft_head/head.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/draft_head/head.py @@ -30,22 +30,44 @@ full-vocab logits; for those the worker all-gathers the shard logits itself and runs the built-in selector, which is the built-in computation. +A drafter that declares a logit transform (``DFlashForCausalLM. +apply_draft_logit_transform``: a multiplier and a tanh softcap, in fp32) is +transformed after the top-k instead of before it (``logit_transform.py``): +stages 2 to 4 run on the bf16 shard with canonical keys, which rank as the +transformed logits do, and the merge reads each candidate's transformed value +from a table the drafter's own method computed. The built-in worker calls the +transform right after the drafter's logits processor; the install puts a +``_DeferredTransform`` on the drafter instance in its place, which hands back +the very shard tensor ``ShardDraftLogitsProcessor`` just returned, untouched, +and records it. ``ShardedDraftHeadDFlashWorker._refine_block_logits``, the +next worker method the built-in forward calls on those logits, keeps them in +bf16 only when this step takes the sharded selector, and applies the +drafter's transform otherwise, so every other path (sampled drafts, d2t, +another layout, any other tensor) gets the built-in transform bit for bit. +With a transform installed, bf16 draft logits are untransformed ones: the +transform always returns fp32. + Installed from the shell's ``load_draft_weights`` (the drafter's selector and the borrowed head exist only after the drafter loads), or from the shell's post-load hooks on load paths that never call it (``LoadFormat.DUMMY``, which never loads the drafter, so there the install refuses: no DFlash 2 drafter): the built-in DFlash worker is re-classed in place to -``ShardedDraftHeadDFlashWorker`` and the drafter's logits processor is -wrapped. Nothing else changes. +``ShardedDraftHeadDFlashWorker``, the drafter's logits processor is wrapped +and, with a logit transform, the drafter instance gets the deferring +transform. Nothing else changes. The merge reproduces the tie order of the ``torch.topk`` build it was written against, which torch does not document. Before installing, ``tie_order_reason`` runs ``torch.topk`` and the sharded path on tie-heavy logits of the deployed -shape and refuses the install if they disagree. +shape (with a transform, also on logits whose transformed values tie) and +refuses the install if they disagree. """ from __future__ import annotations +import weakref +from typing import Optional + import torch from torch import nn @@ -56,8 +78,10 @@ from tensorrt_llm.logger import logger from ...ops.spec.draft_path_walk import DraftPathWalk -from ...ops.spec.draft_shard_topk import CHUNK, NUM_WARPS, DraftShardTopK +from ...ops.spec.draft_shard_topk import CHUNK, FINITE_KEYS, NUM_WARPS, DraftShardTopK from ...ops.spec.draft_topk_merge import SORT_SLOTS, DraftTopKMerge +from . import logit_transform as _logit_transform +from .logit_transform import DraftLogitTransform, ensure_transformed draft_shard_topk = DraftShardTopK() draft_topk_merge = DraftTopKMerge() @@ -66,6 +90,17 @@ #: Attribute of the zero-copy block-logits stand-in that carries the draft tokens. TOKENS_ATTR = "_draft_head_tokens" +#: Worker attribute holding the drafter's ``DraftLogitTransform`` (absent without a transform). +TRANSFORM_ATTR = "_draft_head_transform" +#: Worker attribute: this step's logits were kept untransformed for the sharded selector. +OWED_ATTR = "_draft_head_owed" +#: Processor attribute: a weak reference to the shard tensor its last call returned. +SHARD_ATTR = "_draft_head_shard" +#: Drafter attribute: ``(data_ptr, numel)`` of the logits the deferring transform handed back. +DEFERRED_ATTR = "_draft_head_deferred" +#: The drafter method the built-in worker calls on the draft logits. +TRANSFORM_METHOD = "apply_draft_logit_transform" + #: Largest shard the 16-bit column field of a shard key holds. MAX_SHARD = 65536 @@ -74,10 +109,18 @@ #: of this target does (one request, K + 1 = 8 block positions). PROBE_ROWS = 8 +#: Rows of the install-time probe of many requests (24 requests of K = 7 +#: positions), run with a drafter logit transform. +WIDE_PROBE_ROWS = 168 + #: Seed of the probe's logits, drawn on the host so no device RNG state moves. PROBE_SEED = 1234 +class DraftTransformStateError(RuntimeError): + """A step kept its draft logits untransformed and no selector applied the transform.""" + + def gathers_column_shards(lm_head: nn.Module) -> bool: """Whether ``lm_head`` is a column-parallel head whose forward is shard GEMM + all-gather. @@ -112,6 +155,18 @@ def key_allreduce(mapping) -> nn.Module | None: return allreduce if allreduce.mnnvl_allreduce is not None else None +def canon_kwargs(transform: Optional[DraftLogitTransform]) -> dict: + """``draft_shard_topk`` / ``draft_rank_topk`` keywords for ``transform`` (none without one).""" + if transform is None: + return {} + return {"canon": transform.canon, "canon_band": transform.band} + + +def values_kwargs(transform: Optional[DraftLogitTransform]) -> dict: + """``draft_topk_merge`` keywords for ``transform`` (none without one).""" + return {} if transform is None else {"values": transform.values} + + def selector_draft_tokens( local_logits: torch.Tensor, block_hidden_states: torch.Tensor, @@ -120,6 +175,8 @@ def selector_draft_tokens( mapping, allreduce: nn.Module | None = None, walk=None, + *, + transform: Optional[DraftLogitTransform], ): """Greedy DFlash 2 draft tokens ``[B, K]`` int32 from this rank's ``[B, K, shard]`` logits. @@ -130,23 +187,32 @@ def selector_draft_tokens( (``merge_launch``) on the steps it takes. ``walk`` stands in for ``draft_path_walk`` (same arguments; its result is returned as is): the ``path_walk_tail`` fast path passes a walk that also writes the step's tail. + + ``transform`` is the drafter's logit transform when ``local_logits`` are + the untransformed bf16 shard of a drafter that declares one, else None: + the keys are then canonical and the unary values the transformed ones, so + the tokens are the built-in ones on the transformed logits. Keyword-only + and required, so no call site drops it by accident. """ requests, positions, shard = local_logits.shape local = local_logits.reshape(requests * positions, shard) + canon = canon_kwargs(transform) merge_launch = {} if allreduce is not None: takes = getattr(allreduce, "takes", None) if takes is not None and takes(local.shape[0]): - packed = allreduce.keys(local, selector.top_k, mapping.tp_rank, mapping.tp_size) + packed = allreduce.keys( + local, selector.top_k, mapping.tp_rank, mapping.tp_size, **canon + ) merge_launch = allreduce.merge_launch(local.shape[0]) else: packed = draft_shard_topk( - local, selector.top_k, CHUNK, NUM_WARPS, mapping.tp_rank, mapping.tp_size + local, selector.top_k, CHUNK, NUM_WARPS, mapping.tp_rank, mapping.tp_size, **canon ) gathered = allreduce(packed) else: gathered = allgather( - draft_shard_topk(local, selector.top_k, CHUNK, NUM_WARPS), mapping, dim=0 + draft_shard_topk(local, selector.top_k, CHUNK, NUM_WARPS, **canon), mapping, dim=0 ) # The built-in selector's projection call, operand for operand. gate = selector.hidden_projection( @@ -164,6 +230,7 @@ def selector_draft_tokens( selector.top_k, mapping.tp_size, **merge_launch, + **values_kwargs(transform), ) # dflash2_score_edges' einsum on tensors of its shapes and strides: the same kernel. edges = torch.einsum("blpr,blcr->blpc", predecessors, successors) @@ -175,7 +242,9 @@ class ShardDraftLogitsProcessor(nn.Module): Only the call the DFlash worker makes (``return_context_logits=True`` on a head ``gathers_column_shards`` admits) is changed; anything else goes to - the wrapped processor. + the wrapped processor. Each call records a weak reference to the shard it + returned (``SHARD_ATTR``; cleared by every other call), which the + deferring drafter transform matches. """ def __init__(self, inner: nn.Module): @@ -183,19 +252,105 @@ def __init__(self, inner: nn.Module): self.inner = inner def forward(self, hidden_states, lm_head, attn_metadata, return_context_logits=False): + self.__dict__.pop(SHARD_ATTR, None) if return_context_logits and gathers_column_shards(lm_head): fork = self.__dict__.get("_selector_fork") if fork is not None: # The selector fork of the PDL chain (drafter_pdl_chain/chain.py): the # selector's projection of these rows on a side stream, under the LM-head GEMM. fork.launch(hidden_states) - return lm_head.apply_linear(hidden_states, lm_head.bias) + out = lm_head.apply_linear(hidden_states, lm_head.bias) + self.__dict__[SHARD_ATTR] = weakref.ref(out) + return out return self.inner(hidden_states, lm_head, attn_metadata, return_context_logits) +class _DeferredTransform: + """The drafter's ``apply_draft_logit_transform`` with the draft head's transform installed. + + Set on the drafter instance (``__dict__``), where it shadows the class + method. A call on the shard tensor the drafter's ``ShardDraftLogitsProcessor`` + has just returned, while the drafter still declares the transform the + tables hold and the worker class still routes those logits through + ``ShardedDraftHeadDFlashWorker._refine_block_logits`` (``_defers_transform``), + returns that tensor untouched and records its ``(data_ptr, numel)`` on the + drafter (``DEFERRED_ATTR``). Every other call returns the drafter's own + transform, the built-in computation. Holds weak references only. + """ + + def __init__(self, worker, draft_model): + self._worker = weakref.ref(worker) + self._model = weakref.ref(draft_model) + + def __call__(self, logits: torch.Tensor) -> torch.Tensor: + draft_model = self._model() + worker = self._worker() + if draft_model is None: + raise DraftTransformStateError("[musespark_open] draft_head: the drafter is gone") + transform = None if worker is None else worker.__dict__.get(TRANSFORM_ATTR) + draft_model.__dict__[DEFERRED_ATTR] = None + # The processor in place now, not the one of the install: a wrapper put + # around it later records nothing, and the transform then runs here. + processor = getattr(draft_model, "logits_processor", None) + shard = None + if isinstance(processor, ShardDraftLogitsProcessor): + shard = processor.__dict__.pop(SHARD_ATTR, None) + if transform is None: + return getattr(type(draft_model), TRANSFORM_METHOD)(draft_model, logits) + if ( + shard is not None + and shard() is logits + and logits.dtype == torch.bfloat16 + and transform.matches(draft_model) + and _defers_transform(type(worker)) + ): + draft_model.__dict__[DEFERRED_ATTR] = (logits.data_ptr(), logits.numel()) + return logits + return transform.real(logits) + + +def _owner(cls: type, name: str, after: Optional[type] = None) -> Optional[type]: + """The class of ``cls``'s MRO (past ``after``, when given) that defines ``name``.""" + mro = cls.__mro__ + if after is not None: + mro = mro[mro.index(after) + 1 :] + return next((klass for klass in mro if name in vars(klass)), None) + + +#: ``_defers_transform`` by worker class. +_DEFERS: dict[type, bool] = {} + + +def _defers_transform(cls: type) -> bool: + """Whether a worker of class ``cls`` takes deferred draft logits as this module expects. + + ``ShardedDraftHeadDFlashWorker._refine_block_logits`` must be the one the + built-in forward calls, followed in the MRO by the built-in identity, and the + ``_apply_dflash2_selector`` it reaches must be one that applies the + transform after the top-k (a class marked ``_draft_transform_aware``). + """ + known = _DEFERS.get(cls) + if known is None: + known = False + if issubclass(cls, ShardedDraftHeadDFlashWorker): + name = "_refine_block_logits" + selector = _owner(cls, "_apply_dflash2_selector") + known = ( + _owner(cls, name) is ShardedDraftHeadDFlashWorker + and _owner(cls, name, after=ShardedDraftHeadDFlashWorker) is DFlashWorker + and selector is not None + and bool(vars(selector).get("_draft_transform_aware", False)) + ) + _DEFERS[cls] = known + return known + + class ShardedDraftHeadDFlashWorker(DFlashWorker): """The built-in DFlash worker with the draft head on vocab shards (fast path ``draft_head``).""" + #: Its ``_apply_dflash2_selector`` applies a deferred drafter transform (``_defers_transform``). + _draft_transform_aware = True + def _draft_slot_ids(self, draft_model, num_gens, block_size, num_draft_tokens): """The built-in slot ids, built once per shape outside capture and reused. @@ -237,6 +392,42 @@ def _shard_selector_applies(self, gen_logits, spec_metadata, full_vocab) -> bool and self._d2t is None ) + def _refine_block_logits(self, draft_model, gen_logits, inputs, spec_metadata): + """The drafter's logit transform, unless this step's sharded selector applies it. + + The deferring transform handed back the processor's bf16 shard + untransformed (``DEFERRED_ATTR``). These logits stay bf16 only when they + are that tensor (the reshape keeps its address and size) and this + step takes the sharded selector, which then owes the transform + (``OWED_ATTR``); any other bf16 logits get the drafter's transform + here. Then the next class's refinement, the built-in identity. + """ + deferred = draft_model.__dict__.pop(DEFERRED_ATTR, None) + self.__dict__[OWED_ATTR] = False + transform = self.__dict__.get(TRANSFORM_ATTR) + if transform is not None and gen_logits.dtype == torch.bfloat16: + full_vocab = draft_model.candidate_selector.predecessor_codebook.shape[0] + if ( + deferred == (gen_logits.data_ptr(), gen_logits.numel()) + and gen_logits.shape[-1] != full_vocab + and self._shard_selector_applies(gen_logits, spec_metadata, full_vocab) + ): + self.__dict__[OWED_ATTR] = True + else: + gen_logits = ensure_transformed(transform, gen_logits) + return super()._refine_block_logits(draft_model, gen_logits, inputs, spec_metadata) + + def _draft_head_transform_for(self, gen_logits) -> Optional[DraftLogitTransform]: + """The transform the sharded selector applies to ``gen_logits`` (bf16 means untransformed).""" + self.__dict__[OWED_ATTR] = False + transform = self.__dict__.get(TRANSFORM_ATTR) + return transform if gen_logits.dtype == torch.bfloat16 else None + + def _draft_head_transformed(self, gen_logits) -> torch.Tensor: + """``gen_logits`` with the drafter's transform applied once, for a built-in selector.""" + self.__dict__[OWED_ATTR] = False + return ensure_transformed(self.__dict__.get(TRANSFORM_ATTR), gen_logits) + def _apply_dflash2_selector( self, draft_model, gen_logits, block_hidden_states, anchor_tokens, spec_metadata ): @@ -244,7 +435,11 @@ def _apply_dflash2_selector( full_vocab = selector.predecessor_codebook.shape[0] if gen_logits.shape[-1] == full_vocab: return super()._apply_dflash2_selector( - draft_model, gen_logits, block_hidden_states, anchor_tokens, spec_metadata + draft_model, + self._draft_head_transformed(gen_logits), + block_hidden_states, + anchor_tokens, + spec_metadata, ) if self._shard_selector_applies(gen_logits, spec_metadata, full_vocab): tokens = selector_draft_tokens( @@ -254,17 +449,27 @@ def _apply_dflash2_selector( selector, self.mapping, self.__dict__.get("_draft_head_allreduce"), + transform=self._draft_head_transform_for(gen_logits), ) stand_in = tokens.unsqueeze(-1).expand(*tokens.shape, full_vocab) setattr(stand_in, TOKENS_ATTR, tokens) return stand_in - # The built-in path: the head's own all-gather and .float(), then the built-in selector. - full_logits = allgather(gen_logits, self.mapping, dim=-1).float() + # The built-in path: the transform, the head's own all-gather and .float(), + # then the built-in selector. + full_logits = allgather( + self._draft_head_transformed(gen_logits), self.mapping, dim=-1 + ).float() return super()._apply_dflash2_selector( draft_model, full_logits, block_hidden_states, anchor_tokens, spec_metadata ) def sample_draft_tokens(self, logits, spec_metadata, batch_size, **kwargs): + if self.__dict__.get(OWED_ATTR): + self.__dict__[OWED_ATTR] = False + raise DraftTransformStateError( + "[musespark_open] draft_head: this step kept its draft logits untransformed for " + "the sharded selector, but no selector applied the drafter's logit transform" + ) tokens = getattr(logits, TOKENS_ATTR, None) if tokens is not None: return tokens @@ -300,6 +505,9 @@ def unsupported_reason(worker, draft_model) -> str | None: return f"selector {name} is not bf16" if selector.hidden_projection.weight.dtype != torch.bfloat16: return "selector hidden_projection is not bf16" + shadow = getattr(draft_model, "__dict__", {}).get(TRANSFORM_METHOD) + if shadow is not None and not isinstance(shadow, _DeferredTransform): + return f"the drafter's {TRANSFORM_METHOD} is replaced on the instance" return None @@ -340,73 +548,203 @@ def place(row, values): return torch.stack([probe[r % len(probe)] for r in range(rows)]).to(torch.bfloat16) +def transform_probe_logits( + vocab: int, shard: int, transform: DraftLogitTransform, rows: int = PROBE_ROWS +) -> torch.Tensor: + """Host bf16 ``[rows, vocab]`` logits whose top-k under ``transform`` is decided by its ties. + + Built from the transform's own tie groups (``logit_transform.tie_groups``), + so no threshold is assumed: the largest group above +0.0 (``A < B`` its + smallest and largest finite members) and the largest below -0.0 (``C < + D``). Rows, in turn: ten ``A`` at low ids and ten ``B`` at high ids above a + small grid; the same reversed; a row of the lower group's members with ten + ``C`` low and ten ``D`` high; five +inf among fifteen finite members of the + upper group; consecutive keys across the upper edge ``hi`` of the identity + band; members of the upper group on every shard's first and last columns; + a row of upper-group members only; consecutive keys across the lower edge + ``lo`` above lower-group members. A transform without ties gets the + largest finite values in their place, which tie nowhere. + """ + g = torch.Generator().manual_seed(PROBE_SEED + 1) + lowest, highest = FINITE_KEYS + groups = _logit_transform.tie_groups(transform) + + def largest(candidates, fallback): + return max(candidates, key=lambda grp: grp[1] - grp[0], default=fallback) + + upper = largest([grp for grp in groups if grp[0] >= 0x8000], (highest - 255, highest)) + lower = largest([grp for grp in groups if grp[1] <= 0x7FFF], (lowest, lowest + 255)) + + def finite(group): + return torch.arange(max(group[0], lowest + 1), min(group[1], highest - 1) + 1) + + def members(group, n, with_inf=False): + keys = torch.arange(group[0], group[1] + 1) if with_inf else finite(group) + return keys[torch.randint(0, keys.numel(), (n,), generator=g)] + + def grid(levels): + return (torch.randint(0, levels, (vocab,), generator=g).float() * 0.25 - 3.0).to( + torch.bfloat16 + ) + + def place(row, columns, keys): + row[columns] = _logit_transform.bf16_of_key16(keys) + return row + + def ends(group): + keys = finite(group) + return int(keys[0]), int(keys[-1]) + + low = torch.randperm(vocab // 4, generator=g)[:10] + high = vocab - 1 - torch.randperm(vocab // 4, generator=g)[:10] + a, b = ends(upper) + c, d = ends(lower) + ten = torch.ones(10, dtype=torch.int64) + top = min(transform.hi, highest) + bottom = max(transform.lo, lowest) + upper_ladder = torch.arange(max(top - 23, 0x8000), min(top + 8, highest) + 1) + lower_ladder = torch.arange(max(bottom - 8, lowest), min(bottom + 23, 0x7FFF) + 1) + edges = torch.tensor( + [c for r in range(vocab // shard) for c in (r * shard, (r + 1) * shard - 1)] + ) + mixed = torch.randperm(vocab, generator=g)[:20] + probe = [ + place(place(grid(8), low, a * ten), high, b * ten), + place(place(grid(8), low, b * ten), high, a * ten), + place( + place(place(grid(8), torch.arange(vocab), members(lower, vocab, True)), low, c * ten), + high, + d * ten, + ), + place(place(grid(8), mixed[:5], torch.full((5,), highest)), mixed[5:], members(upper, 15)), + place(grid(8), torch.randperm(vocab, generator=g)[: upper_ladder.numel()], upper_ladder), + place(grid(4), edges, members(upper, edges.numel())), + place(grid(8), torch.arange(vocab), members(upper, vocab, True)), + place( + place(grid(8), torch.arange(vocab), members(lower, vocab, True)), + torch.randperm(vocab, generator=g)[: lower_ladder.numel()], + lower_ladder, + ), + ] + return torch.stack([probe[r % len(probe)] for r in range(rows)]) + + +def tie_order_probes( + vocab: int, shard: int, transform: Optional[DraftLogitTransform], device +) -> list[torch.Tensor]: + """The install's probe logits on ``device``, bf16 ``[rows, vocab]`` each. + + ``tie_order_probe_logits``; with a transform also ``transform_probe_logits`` + and, at ``WIDE_PROBE_ROWS`` rows, both repeated. + """ + probes = [tie_order_probe_logits(vocab, shard).to(device)] + if transform is not None: + ties = transform_probe_logits(vocab, shard, transform).to(device) + both = torch.cat([probes[0], ties]) + repeats = -(-WIDE_PROBE_ROWS // both.shape[0]) + probes += [ties, both.repeat(repeats, 1)[:WIDE_PROBE_ROWS].contiguous()] + return probes + + +def probe_reference(full: torch.Tensor, top_k: int, transform, topk=torch.topk): + """``topk`` as the built-in selector runs it on ``full`` (bf16 ``[rows, vocab]``). + + On ``[1, rows, vocab]`` fp32: ``full.float()``, or the drafter's transform of it. + """ + rows, vocab = full.shape + logits = full.float() if transform is None else transform.real(full) + return topk(logits.view(1, rows, vocab), top_k, dim=-1) + + def _same_bits(a: torch.Tensor, b: torch.Tensor) -> bool: """Bitwise equality of fp32 tensors, NaN payloads aside.""" - nan = a.isnan() - if not torch.equal(nan, b.isnan()): - return False - return torch.equal( - torch.where(nan, 0.0, a).view(torch.int32), torch.where(nan, 0.0, b).view(torch.int32) - ) + return _logit_transform.same_bits(a, b) -def tie_order_reason(selector: nn.Module, num_ranks: int, topk=torch.topk) -> str | None: +def tie_order_reason( + selector: nn.Module, + num_ranks: int, + topk=torch.topk, + transform: Optional[DraftLogitTransform] = None, +) -> str | None: """Why the sharded candidates would differ from this torch's ``torch.topk``, or None. The merge copies the tie order of ``torch.topk`` (selection, arrangement and bitonic network), which torch does not document. This runs ``topk`` as the built-in selector does -- ``[1, rows, vocab]`` fp32 -- on - ``tie_order_probe_logits``, and the sharded path on the same logits split + ``tie_order_probes``, and the sharded path on the same logits split into ``num_ranks`` shards with the keys exchanged in both layouts (concatenated int32 keys and summed fp32 halves, which is what the all-gather and the all-reduce return), and compares ids and values bit for - bit. ``topk`` is a parameter so a test can stand in another order. + bit. With ``transform`` the reference transforms the logits first and the + sharded path runs on canonical keys with the merge's values table. ``topk`` + is a parameter so a test can stand in another order. """ predecessor_codebook = selector.predecessor_codebook vocab, rank_dim = predecessor_codebook.shape shard = vocab // num_ranks top_k = selector.top_k device = predecessor_codebook.device + canon = canon_kwargs(transform) with torch.cuda.device(device): - full = tie_order_probe_logits(vocab, shard).to(device) - rows = full.shape[0] - ref_values, ref_ids = topk(full.float().view(1, rows, vocab), top_k, dim=-1) - shards = [full[:, r * shard : (r + 1) * shard].contiguous() for r in range(num_ranks)] - layouts = { - "int32 all-gather": torch.cat( - [draft_shard_topk(s, top_k, CHUNK, NUM_WARPS) for s in shards], dim=0 - ), - "fp32 all-reduce": torch.stack( - [ - draft_shard_topk(s, top_k, CHUNK, NUM_WARPS, r, num_ranks) - for r, s in enumerate(shards) - ] - ).sum(dim=0), - } - gate = torch.zeros(1, rows, rank_dim, dtype=torch.bfloat16, device=device) - anchor = torch.zeros(1, dtype=torch.int64, device=device) - for layout, keys in layouts.items(): - ids, values, _, _ = draft_topk_merge( - keys, - shard, - 1, - rows, - anchor, - gate, - predecessor_codebook, - selector.successor_codebook, - top_k, - num_ranks, - ) - if not (torch.equal(ids, ref_ids) and _same_bits(values, ref_values)): - differ = (ids != ref_ids).any(dim=-1).view(-1).nonzero().view(-1).tolist() - return ( - f"torch {torch.__version__} orders top-k ties differently from the " - f"merge ({layout} keys; probe rows {differ} differ)" + for probe, full in enumerate(tie_order_probes(vocab, shard, transform, device)): + rows = full.shape[0] + ref_values, ref_ids = probe_reference(full, top_k, transform, topk) + shards = [full[:, r * shard : (r + 1) * shard].contiguous() for r in range(num_ranks)] + layouts = { + "int32 all-gather": torch.cat( + [draft_shard_topk(s, top_k, CHUNK, NUM_WARPS, **canon) for s in shards], dim=0 + ), + "fp32 all-reduce": torch.stack( + [ + draft_shard_topk(s, top_k, CHUNK, NUM_WARPS, r, num_ranks, **canon) + for r, s in enumerate(shards) + ] + ).sum(dim=0), + } + gate = torch.zeros(1, rows, rank_dim, dtype=torch.bfloat16, device=device) + anchor = torch.zeros(1, dtype=torch.int64, device=device) + for layout, keys in layouts.items(): + ids, values, _, _ = draft_topk_merge( + keys, + shard, + 1, + rows, + anchor, + gate, + predecessor_codebook, + selector.successor_codebook, + top_k, + num_ranks, + **values_kwargs(transform), ) + if not (torch.equal(ids, ref_ids) and _same_bits(values, ref_values)): + differ = (ids != ref_ids).any(dim=-1).view(-1).nonzero().view(-1).tolist() + return ( + f"torch {torch.__version__} orders top-k ties differently from the " + f"merge ({layout} keys; probe {probe} rows {differ} differ)" + ) return None +def agreed_reason(mapping, reason: str | None, transform: Optional[DraftLogitTransform]): + """``reason`` as every rank of ``mapping``'s TP group agrees, tables compared. + + Collective over the TP group (``dispatch.agreed_union`` and + ``dispatch.agreed_refusal``): every rank calls it at the same point of the + install, whatever its own reason. Ranks whose transform tables differ + (``logit_transform.digest``) refuse, since each rank's merge reads its own + tables and differing ones would draft different tokens per rank. + """ + from ... import dispatch as _tp_agreement + + mine = _logit_transform.digest(transform) + seen = _tp_agreement.agreed_union(mapping, [mine]) + if reason is None and seen != [mine]: + reason = f"the TP ranks differ in their drafter logit transform ({'; '.join(seen)})" + return _tp_agreement.agreed_refusal(mapping, reason) + + def installed(llm: nn.Module) -> bool: """Whether ``install`` has put the draft head on ``llm``'s DFlash worker.""" return isinstance(getattr(llm, "spec_worker", None), ShardedDraftHeadDFlashWorker) @@ -417,15 +755,19 @@ def install(llm: nn.Module, make_allreduce=key_allreduce) -> bool: ``make_allreduce(mapping)`` builds the key exchange (collective: every rank installs at the same point of the load). A second call on an installed - ``llm`` changes nothing. + ``llm`` changes nothing. Everything is built and checked before anything + changes, and the refusal is agreed over the TP group (``agreed_reason``), + so a refusing install leaves the worker and the drafter as they were on + every rank. Raises: ValueError: when there is no DFlash 2 drafter (no spec worker, or a drafter without a candidate selector), or its head layout is not one - the feature supports, or this torch's ``torch.topk`` orders ties - differently from the merge (``tie_order_reason``) -- running the - built-in path, or a path whose drafts differ, under the fast - path's name would misreport what runs. + the feature supports, or its logit transform is not one the draft + head can apply after the top-k (``logit_transform.unsupported_reason``), + or this torch's ``torch.topk`` orders ties differently from the merge + (``tie_order_reason``) -- running the built-in path, or a path whose + drafts differ, under the fast path's name would misreport what runs. """ if installed(llm): return True @@ -435,15 +777,23 @@ def install(llm: nn.Module, make_allreduce=key_allreduce) -> bool: raise ValueError( "[musespark_open] draft_head cannot run here: there is no DFlash 2 drafter" ) - reason = unsupported_reason(worker, draft_model) or tie_order_reason( - draft_model.candidate_selector, worker.mapping.tp_size - ) + selector = draft_model.candidate_selector + transform = None + reason = unsupported_reason(worker, draft_model) + if reason is None: + transform = _logit_transform.build(draft_model, selector.predecessor_codebook.device) + if transform is not None: + reason = _logit_transform.unsupported_reason( + transform, draft_model.lm_head.out_features + ) + if reason is None: + reason = tie_order_reason(selector, worker.mapping.tp_size, transform=transform) + reason = agreed_reason(worker.mapping, reason, transform) if reason is not None: raise ValueError(f"[musespark_open] draft_head cannot run here: {reason}") - worker.__class__ = ShardedDraftHeadDFlashWorker - if type(draft_model.logits_processor) is LogitsProcessor: - draft_model.logits_processor = ShardDraftLogitsProcessor(draft_model.logits_processor) - if "_draft_head_allreduce" not in worker.__dict__: + if "_draft_head_allreduce" in worker.__dict__: + allreduce = worker.__dict__["_draft_head_allreduce"] + else: allreduce = make_allreduce(worker.mapping) if allreduce is None: # The composition's precondition checked the group takes an MNNVL all-reduce. @@ -451,13 +801,28 @@ def install(llm: nn.Module, make_allreduce=key_allreduce) -> bool: "the vocab-sharded draft head: the MNNVL FP32 all-reduce of the key exchange " "could not be built on this TP group" ) - # Not registered as a submodule: it owns no parameters, only a fabric workspace. - worker.__dict__["_draft_head_allreduce"] = allreduce - exchange = "MNNVL all-reduce" if worker.__dict__["_draft_head_allreduce"] else "NCCL all-gather" + worker.__class__ = ShardedDraftHeadDFlashWorker + if type(draft_model.logits_processor) is LogitsProcessor: + draft_model.logits_processor = ShardDraftLogitsProcessor(draft_model.logits_processor) + # Not registered as a submodule: it owns no parameters, only a fabric workspace. + worker.__dict__["_draft_head_allreduce"] = allreduce + if transform is not None: + # Tables built outside capture, on the drafter's device; plain attributes, not buffers. + worker.__dict__[TRANSFORM_ATTR] = transform + draft_model.__dict__[TRANSFORM_METHOD] = _DeferredTransform(worker, draft_model) + exchange = "MNNVL all-reduce" if allreduce else "NCCL all-gather" + applied = ( + "" + if transform is None + else ( + f", drafter logit transform applied after the top-k (canonical keys outside " + f"[{transform.lo:#06x}, {transform.hi:#06x}])" + ) + ) logger.info( "[musespark_open] draft_head: DFlash 2 draft head on vocab shards " f"(shard {draft_model.lm_head.out_features}, tp {worker.mapping.tp_size}, " f"top_k {draft_model.candidate_selector.top_k}, keys by {exchange}, " - f"tie order checked against torch {torch.__version__})" + f"tie order checked against torch {torch.__version__}{applied})" ) return True diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/draft_head/logit_transform.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/draft_head/logit_transform.py new file mode 100644 index 000000000000..9c58afd615fe --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/draft_head/logit_transform.py @@ -0,0 +1,303 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The DFlash drafter's logit transform on the vocab-sharded draft head (``draft_head``). + +A DFlash drafter whose config declares ``output_multiplier`` or +``final_logit_softcapping`` is trained on transformed logits, and the built-in +worker applies ``DFlashForCausalLM.apply_draft_logit_transform`` to the draft +logits before the selector: in fp32 and in this order, ``.float()``, ``* +multiplier``, ``/ softcap``, ``tanh``, ``* softcap``. Call it ``t``. The +selector then takes ``torch.topk(t(logits), top_k)``: its ids, their order and +its values (the unary scores it adds to its pairwise ones). + +``t`` never decreases, so a top-k of the bf16 logits has the same candidates in +the same order -- except where ``t`` maps distinct bf16 values to the same fp32 +value. tanh saturates: with the Muse Spark Open drafter's multiplier and softcap +every logit of large enough magnitude transforms to exactly +-softcap. Those +compare equal for ``torch.topk``, which then keeps and orders them by vocab id, +while their bf16 values are distinct. So the draft head ranks canonical keys +and reads the unary values from a table, both built once per drafter here: + + canon [65536] int32: 16-bit radix key of a bf16 logit -> the key of its tie + group's representative. A tie group is the bf16 values whose + transformed values have the same bits; ``t`` never decreases, so a + group is a run of consecutive keys, and canonical keys compare as the + transformed values' radix keys do. The representative is the group's + smallest key, or its largest below +0.0 so that the -0.0 group keeps + -0.0 (``-0.0 == +0.0`` under the selector sort's comparison, as for + the transformed values). NaN (0xFFFF) maps to itself. + values [65536] fp32: key -> ``t`` of that bf16 value, computed by the + drafter's own method on the device, so it is the built-in transform + of every logit bit for bit (the transform is elementwise). + lo, hi the key interval around 0 on which ``canon`` is the identity: logits + there need no table load. + +``unsupported_reason`` checks all of it exhaustively over the 65536 keys before +the draft head installs it. +""" + +from __future__ import annotations + +import hashlib +import weakref +from typing import Callable, Optional + +import torch + +from tensorrt_llm._torch.models.modeling_dflash import DFlashForCausalLM + +from ...ops.spec.draft_shard_topk import CANON_ENTRIES, FINITE_KEYS, canon_reason, radix_key16 + +__all__ = [ + "DraftLogitTransform", + "bf16_of_key16", + "build", + "digest", + "ensure_transformed", + "same_bits", + "tie_groups", + "unsupported_reason", +] + +#: Rows of the install-time layout check: eight requests of K = 7 positions. +LAYOUT_ROWS = 56 + +#: Seed of the layout check's logits, drawn on the host so no device RNG state moves. +LAYOUT_SEED = 2718 + +#: Key16 of +0.0 and -0.0. +POSITIVE_ZERO, NEGATIVE_ZERO = 0x8000, 0x7FFF + + +def bf16_of_key16(keys: torch.Tensor) -> torch.Tensor: + """The bf16 values of 16-bit radix keys (any int dtype, values in [0, 0xFFFF]).""" + key = keys.to(torch.int32) + bits = torch.where(key >= 0x8000, key ^ 0x8000, key ^ 0xFFFF) + bits = torch.where(bits >= 0x8000, bits - 0x10000, bits).to(torch.int16) + return bits.view(torch.bfloat16) + + +def _radix32(values: torch.Tensor) -> torch.Tensor: + """torch's fp32 radix key of ``values`` as int64 in [0, 2**32); every NaN is the top key.""" + word = values.contiguous().view(torch.int32).to(torch.int64) & 0xFFFFFFFF + key = torch.where(word >= 0x80000000, word ^ 0xFFFFFFFF, word ^ 0x80000000) + return torch.where(values.isnan(), torch.full_like(key, 0xFFFFFFFF), key) + + +def same_bits(a: torch.Tensor, b: torch.Tensor) -> bool: + """Bitwise equality of fp32 tensors, NaN payloads aside.""" + nan = a.isnan() + if not torch.equal(nan, b.isnan()): + return False + return torch.equal( + torch.where(nan, 0.0, a).view(torch.int32), torch.where(nan, 0.0, b).view(torch.int32) + ) + + +class DraftLogitTransform: + """A drafter's logit transform as the draft head's tables (see the module docstring). + + ``canon`` [65536] int32 and ``values`` [65536] fp32 live on the drafter's + device; ``lo`` and ``hi`` are plain ints; ``multiplier`` and ``softcap`` are + the drafter attributes the tables were built from. ``real`` is the drafter's + own transform, the built-in computation. + """ + + def __init__( + self, + draft_model, + function: Callable, + canon: torch.Tensor, + values: torch.Tensor, + lo: int, + hi: int, + ): + self._model = weakref.ref(draft_model) + self.function = function + self.canon = canon + self.values = values + self.lo = lo + self.hi = hi + self.multiplier = getattr(draft_model, "logit_multiplier", None) + self.softcap = getattr(draft_model, "logit_softcap", None) + + @property + def band(self) -> tuple[int, int]: + """``(lo, hi)``: ``draft_shard_topk``'s ``canon_band``.""" + return self.lo, self.hi + + def matches(self, draft_model) -> bool: + """Whether ``draft_model`` still declares the transform these tables hold.""" + return ( + getattr(draft_model, "logit_multiplier", None), + getattr(draft_model, "logit_softcap", None), + ) == (self.multiplier, self.softcap) + + def real(self, logits: torch.Tensor) -> torch.Tensor: + """The drafter's transform of ``logits``: the built-in computation, fp32 out.""" + draft_model = self._model() + if draft_model is None: + raise RuntimeError("[musespark_open] draft_head: the drafter of this transform is gone") + return self.function(draft_model, logits) + + +def build(draft_model, device) -> Optional[DraftLogitTransform]: + """The tables of ``draft_model``'s logit transform on ``device``, or None for no transform. + + None when the drafter's class has no ``apply_draft_logit_transform`` or it + returns a bf16 probe unchanged (multiplier 1, no softcap). The probe is + bf16 because on an fp32 input the built-in transform runs in place and + returns the input itself even when it is on. Runs the drafter's method on + all 65536 bf16 values, so call it outside CUDA graph capture. + ``unsupported_reason`` says whether the result can be used. + """ + function = getattr(type(draft_model), "apply_draft_logit_transform", None) + if function is None: + return None + probe = torch.zeros(4, dtype=torch.bfloat16, device=device) + if function(draft_model, probe) is probe: + return None + keys = torch.arange(CANON_ENTRIES, dtype=torch.int32, device=device) + values = function(draft_model, bf16_of_key16(keys)) + canon = torch.arange(CANON_ENTRIES, dtype=torch.int32) + lo, hi = 0, CANON_ENTRIES - 1 + if values.dtype == torch.float32 and tuple(values.shape) == (CANON_ENTRIES,): + first_key, last_key = FINITE_KEYS + radix = _radix32(values.cpu())[first_key : last_key + 1] + finite = torch.arange(first_key, last_key + 1, dtype=torch.int64) + start = torch.ones_like(radix, dtype=torch.bool) + start[1:] = radix[1:] != radix[:-1] + end = torch.ones_like(start) + end[:-1] = start[1:] + group = torch.cumsum(start.to(torch.int64), dim=0) - 1 + first, last = finite[start], finite[end] + representative = torch.where(first >= POSITIVE_ZERO, first, last) + canon[first_key : last_key + 1] = representative[group].to(torch.int32) + identity = canon == torch.arange(CANON_ENTRIES, dtype=torch.int32) + above = (~identity[POSITIVE_ZERO:]).nonzero() + below = (~identity[:POSITIVE_ZERO]).nonzero() + hi = CANON_ENTRIES - 1 if above.numel() == 0 else POSITIVE_ZERO + int(above[0]) - 1 + lo = 0 if below.numel() == 0 else int(below[-1]) + 1 + return DraftLogitTransform(draft_model, function, canon.to(device), values.contiguous(), lo, hi) + + +def tie_groups(transform: DraftLogitTransform) -> list[tuple[int, int]]: + """``(first, last)`` key16 of every tie group of two or more values, in key order.""" + first_key, last_key = FINITE_KEYS + canon = transform.canon.cpu()[first_key : last_key + 1].to(torch.int64) + start = torch.ones_like(canon, dtype=torch.bool) + start[1:] = canon[1:] != canon[:-1] + begins = (start.nonzero().view(-1) + first_key).tolist() + ends = [b - 1 for b in begins[1:]] + [last_key] + return [(b, e) for b, e in zip(begins, ends) if e > b] + + +def _layout_probe(rows: int, shard: int) -> torch.Tensor: + """Host bf16 ``[rows, shard]`` logits: every bf16 value once, then values of many scales. + + The 65536 bf16 bit patterns (NaNs included) fill the first columns in a + random order; each remaining row is normal logits of one scale up to deep + saturation, with +-inf, +-0.0 and NaN scattered in. + """ + g = torch.Generator().manual_seed(LAYOUT_SEED) + scales = (1.0, 10.0, 100.0, 400.0, 1000.0, 3000.0, 1e5, 1e30) + x = torch.randn(rows, shard, generator=g) + x *= torch.tensor([scales[r % len(scales)] for r in range(rows)])[:, None] + special = torch.tensor([float("inf"), float("-inf"), 0.0, -0.0, float("nan")]) + picks = torch.randint(0, shard, (rows, 40), generator=g) + x.scatter_(1, picks, special[torch.randint(0, len(special), (rows, 40), generator=g)]) + x = x.to(torch.bfloat16) + every = torch.arange(CANON_ENTRIES, dtype=torch.int32)[ + torch.randperm(CANON_ENTRIES, generator=g) + ] + flat = x.view(-1) + count = min(CANON_ENTRIES, flat.numel()) + flat[:count] = bf16_of_key16(every[:count]) + return x + + +def unsupported_reason( + transform: DraftLogitTransform, shard: int, rows: int = LAYOUT_ROWS +) -> Optional[str]: + """Why the draft head cannot apply ``transform`` after its top-k, or None. + + Exhaustive over the 65536 bf16 values: the transform is the built-in + ``DFlashForCausalLM`` method, returns fp32, keeps NaN, +0.0 and -0.0, never + decreases in torch's radix order (which makes the tie groups runs of + consecutive keys), and the canonical values compare under ``>`` as the + transformed values do. ``canon`` is a table ``draft_shard_topk`` reads + correctly with band ``(lo, hi)``. Last, the table gives the drafter's own + transform of a ``[rows, shard]`` bf16 slice bit for bit (a host-seeded + probe holding every bf16 value). + """ + function = transform.function + if function is not DFlashForCausalLM.apply_draft_logit_transform: + return ( + f"the drafter's apply_draft_logit_transform is {function.__qualname__}, " + "not DFlashForCausalLM's" + ) + values = transform.values + if values.dtype != torch.float32 or tuple(values.shape) != (CANON_ENTRIES,): + return f"the drafter's logit transform returns {values.dtype}, not fp32" + table = values.cpu() + inputs = bf16_of_key16(torch.arange(CANON_ENTRIES, dtype=torch.int32)) + if not torch.equal(inputs.isnan(), table.isnan()): + return "the drafter's logit transform does not map exactly the NaN logits to NaN" + words = table.view(torch.int32) + if int(words[POSITIVE_ZERO]) != 0 or int(words[NEGATIVE_ZERO]) != -(2**31): + return "the drafter's logit transform does not keep +0.0 and -0.0" + first_key, last_key = FINITE_KEYS + radix = _radix32(table)[first_key : last_key + 1] + drops = (radix[1:] < radix[:-1]).nonzero().view(-1) + if drops.numel(): + at = float(inputs[first_key + int(drops[0]) + 1].float()) + return f"the drafter's logit transform decreases at the bf16 logit {at}" + canon = transform.canon.cpu() + reason = canon_reason(canon, transform.band, torch.device("cpu")) + if reason is not None: + return reason + finite = canon[first_key : last_key + 1] + start = torch.ones_like(finite, dtype=torch.bool) + start[1:] = finite[1:] != finite[:-1] + members = torch.arange(first_key, last_key + 1)[start] + if bool((radix[1:] == radix[:-1])[start[1:]].any()): + return "canon splits a tie group of the transform" + if not torch.equal(_radix32(table[finite.long()]), _radix32(table[first_key : last_key + 1])): + return "canon maps a key outside its tie group" + shown = bf16_of_key16(finite[start]).float() + shown_values = table[members] + if not torch.equal(shown[1:] > shown[:-1], shown_values[1:] > shown_values[:-1]): + return "canonical values compare differently from the transformed values" + probe = _layout_probe(rows, shard).to(values.device) + got = values[radix_key16(probe.view(torch.int16)).long()] + if not same_bits(got, transform.real(probe)): + return ( + f"the drafter's logit transform of a [{rows}, {shard}] slice differs from its " + "transform of the 65536 bf16 values" + ) + return None + + +def digest(transform: Optional[DraftLogitTransform]) -> str: + """A short digest of the tables, equal on ranks whose draft heads would agree.""" + if transform is None: + return "no drafter logit transform" + h = hashlib.sha256() + h.update(transform.canon.cpu().numpy().tobytes()) + h.update(transform.values.cpu().view(torch.int32).numpy().tobytes()) + h.update(repr((transform.lo, transform.hi, transform.multiplier, transform.softcap)).encode()) + return f"drafter logit transform tables {h.hexdigest()[:16]}" + + +def ensure_transformed( + transform: Optional[DraftLogitTransform], logits: torch.Tensor +) -> torch.Tensor: + """``logits`` with the drafter's transform applied once. + + With a transform installed, bf16 draft logits are untransformed (the + transform always returns fp32, which ``unsupported_reason`` checks), so a + bf16 ``logits`` gets ``transform.real`` and anything else is returned as is. + """ + if transform is None or logits.dtype != torch.bfloat16: + return logits + return transform.real(logits) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/verify_fold/fold.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/verify_fold/fold.py index cbb858ad7d0a..460cbb116506 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/verify_fold/fold.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/verify_fold/fold.py @@ -180,6 +180,9 @@ class SpecLoopFoldDFlashWorker(DFlashWorker): _fold_verify: bool = False _fold_path_walk_tail: bool = False + #: Its ``_apply_dflash2_selector`` applies the draft head's deferred drafter + #: logit transform (``draft_head.head._defers_transform``). + _draft_transform_aware = True # ------------------------------------------------------------------ # The per-step state @@ -759,6 +762,7 @@ def walk(edges, unary, ids): self.mapping, self.__dict__.get("_draft_head_allreduce"), walk=walk, + transform=self._draft_head_transform_for(gen_logits), ) step.tail = _Tail(tokens, next_tokens, kv_lens is not None, batch_indices) if kv_lens is not None: diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/wide_draft_topk/exchange.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/wide_draft_topk/exchange.py index 85b29dcce163..8025d17ddf5f 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/wide_draft_topk/exchange.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/wide_draft_topk/exchange.py @@ -27,8 +27,10 @@ the row count only, the same on every rank). Before wrapping it, ``tie_order_reason`` runs the rank top-k layout with every launch of this fast path against the running ``torch.topk`` on the draft head's tie-heavy -probe (``draft_head.tie_order_probe_logits``) and refuses the install if they -disagree. +probes (``draft_head.tie_order_probes``) and refuses the install if they +disagree. With the drafter's logit transform installed on the draft head +(``draft_head/logit_transform.py``) ``keys`` takes its canonical keys, and the +probes and their reference are the transformed ones. """ from __future__ import annotations @@ -42,8 +44,9 @@ from tensorrt_llm.logger import logger from ...ops.spec.draft_rank_topk import MAX_KEYS, draft_rank_topk -from ...ops.spec.draft_shard_topk import CHUNK +from ...ops.spec.draft_shard_topk import CHUNK, FULL_BAND from ..draft_head import head as _draft_head +from ..draft_head.logit_transform import DraftLogitTransform #: Fewest draft rows (requests x K) the exchange takes: five requests at K = 7. Fewer #: rows are the one-request path and the 9..32-row band. @@ -67,9 +70,23 @@ def takes(self, rows: int) -> bool: """Whether a draft step of ``rows`` rows exchanges the rank top-k (``keys``).""" return rows >= MIN_ROWS - def keys(self, local: torch.Tensor, top_k: int, rank: int, num_ranks: int) -> torch.Tensor: - """This rank's own top ``top_k`` keys of ``local`` in the all-reduce layout.""" - return draft_rank_topk(local, top_k, CHUNK, CHUNK_WARPS, rank, num_ranks) + def keys( + self, + local: torch.Tensor, + top_k: int, + rank: int, + num_ranks: int, + canon: Optional[torch.Tensor] = None, + canon_band: tuple[int, int] = FULL_BAND, + ) -> torch.Tensor: + """This rank's own top ``top_k`` keys of ``local`` in the all-reduce layout. + + ``canon`` / ``canon_band``: a drafter logit transform's canonical keys + (``draft_shard_topk``'s), passed by ``draft_head.selector_draft_tokens``. + """ + return draft_rank_topk( + local, top_k, CHUNK, CHUNK_WARPS, rank, num_ranks, canon=canon, canon_band=canon_band + ) def merge_launch(self, rows: int) -> dict: """``draft_topk_merge``'s launch geometry for a step of ``rows`` rows.""" @@ -82,14 +99,21 @@ def __call__(self, packed: torch.Tensor) -> torch.Tensor: return self.allreduce(packed) -def tie_order_reason(selector: nn.Module, num_ranks: int, topk=torch.topk) -> Optional[str]: +def tie_order_reason( + selector: nn.Module, + num_ranks: int, + topk=torch.topk, + transform: Optional[DraftLogitTransform] = None, +) -> Optional[str]: """Why the rank top-k exchange would differ from this torch's ``torch.topk``, or None. - ``draft_head.tie_order_reason``'s probe for the rank top-k layout: the + ``draft_head.tie_order_reason``'s probes for the rank top-k layout: the probe logits split into ``num_ranks`` shards, each shard's ``keys``, their sum (what the all-reduce returns) merged with every ``MERGE_LAUNCH`` geometry, ids and values compared bit for bit with ``topk`` as the built-in - selector runs it. ``topk`` is a parameter so a test can stand in another order. + selector runs it. With ``transform`` (the draft head's) the keys are + canonical, the merge reads its values table and the reference transforms + the logits first. ``topk`` is a parameter so a test can stand in another order. """ predecessor_codebook = selector.predecessor_codebook vocab, rank_dim = predecessor_codebook.shape @@ -97,42 +121,48 @@ def tie_order_reason(selector: nn.Module, num_ranks: int, topk=torch.topk) -> Op top_k = selector.top_k device = predecessor_codebook.device exchange = RankTopKExchange(None) + canon = _draft_head.canon_kwargs(transform) with torch.cuda.device(device): - full = _draft_head.tie_order_probe_logits(vocab, shard).to(device) - rows = full.shape[0] - ref_values, ref_ids = topk(full.float().view(1, rows, vocab), top_k, dim=-1) - keys = torch.stack( - [ - exchange.keys( - full[:, r * shard : (r + 1) * shard].contiguous(), top_k, r, num_ranks - ) - for r in range(num_ranks) - ] - ).sum(dim=0) - gate = torch.zeros(1, rows, rank_dim, dtype=torch.bfloat16, device=device) - anchor = torch.zeros(1, dtype=torch.int64, device=device) - for _, rank_block, num_warps in MERGE_LAUNCH: - ids, values, _, _ = _draft_head.draft_topk_merge( - keys, - shard, - 1, - rows, - anchor, - gate, - predecessor_codebook, - selector.successor_codebook, - top_k, - num_ranks, - rank_block=rank_block, - num_warps=num_warps, - ) - if not (torch.equal(ids, ref_ids) and _draft_head._same_bits(values, ref_values)): - differ = (ids != ref_ids).any(dim=-1).view(-1).nonzero().view(-1).tolist() - return ( - f"torch {torch.__version__} orders top-k ties differently from the merge over " - f"rank top-k keys (rank_block {rank_block}, {num_warps} warps; probe rows " - f"{differ} differ)" + for probe, full in enumerate(_draft_head.tie_order_probes(vocab, shard, transform, device)): + rows = full.shape[0] + ref_values, ref_ids = _draft_head.probe_reference(full, top_k, transform, topk) + keys = torch.stack( + [ + exchange.keys( + full[:, r * shard : (r + 1) * shard].contiguous(), + top_k, + r, + num_ranks, + **canon, + ) + for r in range(num_ranks) + ] + ).sum(dim=0) + gate = torch.zeros(1, rows, rank_dim, dtype=torch.bfloat16, device=device) + anchor = torch.zeros(1, dtype=torch.int64, device=device) + for _, rank_block, num_warps in MERGE_LAUNCH: + ids, values, _, _ = _draft_head.draft_topk_merge( + keys, + shard, + 1, + rows, + anchor, + gate, + predecessor_codebook, + selector.successor_codebook, + top_k, + num_ranks, + rank_block=rank_block, + num_warps=num_warps, + **_draft_head.values_kwargs(transform), ) + if not (torch.equal(ids, ref_ids) and _draft_head._same_bits(values, ref_values)): + differ = (ids != ref_ids).any(dim=-1).view(-1).nonzero().view(-1).tolist() + return ( + f"torch {torch.__version__} orders top-k ties differently from the merge " + f"over rank top-k keys (rank_block {rank_block}, {num_warps} warps; probe " + f"{probe} rows {differ} differ)" + ) return None @@ -168,7 +198,11 @@ def install(llm: nn.Module) -> bool: return False reason = unsupported_reason(llm) if reason is None: - reason = tie_order_reason(llm.draft_model.candidate_selector, worker.mapping.tp_size) + reason = tie_order_reason( + llm.draft_model.candidate_selector, + worker.mapping.tp_size, + transform=worker.__dict__.get(_draft_head.TRANSFORM_ATTR), + ) if reason is not None: raise ValueError(f"[musespark_open] wide_draft_topk cannot run here: {reason}") worker.__dict__[ATTRIBUTE] = RankTopKExchange(worker.__dict__[ATTRIBUTE]) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/draft_rank_topk.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/draft_rank_topk.py index 05a2c2ef7448..5a0db30bcdbe 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/draft_rank_topk.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/draft_rank_topk.py @@ -15,6 +15,11 @@ ``draft_topk_merge`` over the sum of all ranks' outputs gives the result it gives over the per-chunk keys, bit for bit (selection, tie order and the selector's operands), from 1/13 of the bytes at TP4. + +``canon`` and ``canon_band`` (a drafter logit transform's canonical keys) go to +the per-chunk stage unchanged; the rank's top ``top_k`` of canonical keys holds +every key of the global top ``top_k`` under the transformed order for the same +reason. """ from __future__ import annotations @@ -23,9 +28,9 @@ import triton import triton.language as tl -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper -from .draft_shard_topk import CHUNK, NUM_WARPS, TOP_K, DraftShardTopK +from .draft_shard_topk import CHUNK, FULL_BAND, NUM_WARPS, SHARD, TOP_K, DraftShardTopK #: Per-chunk keys one row program reduces. MAX_KEYS = 4096 @@ -78,6 +83,8 @@ class DraftRankTopK(OpWrapper): chunk columns per program of the per-chunk stage (a power of two >= top_k) num_warps warps per program of the per-chunk stage rank, num_ranks + canon, canon_band + ``draft_shard_topk``'s canonical-key table and its identity band Returns ``[M, num_ranks * top_k * 2]`` fp32: rank ``rank``'s keys at columns ``[rank * top_k * 2, (rank + 1) * top_k * 2)`` as (high, low) 16-bit halves, @@ -87,6 +94,20 @@ class DraftRankTopK(OpWrapper): """ ARCHS = frozenset({Arch.SM_103, Arch.SM_107}) + CELLS: tuple[Cell, ...] = ( + Cell( + "five requests at TP8 (35 rows, the fewest wide_draft_topk takes)", + {"rows": 35, "shard": SHARD, "transform": None}, + ), + Cell( + "five requests at TP8, the Muse Spark Open drafter transform's canon", + {"rows": 35, "shard": SHARD, "transform": "musespark"}, + ), + Cell( + "24 requests at TP4 (a 50512-column shard), the drafter transform", + {"rows": 168, "shard": 2 * SHARD, "transform": "musespark"}, + ), + ) note = ( "Used with draft_topk_merge after the draft head's MNNVL fp32 all-reduce " "(wide_draft_topk's RankTopKExchange, from 33 rows): the sum of " @@ -103,9 +124,13 @@ def __call__( num_warps: int = NUM_WARPS, rank: int = 0, num_ranks: int = 1, + canon: torch.Tensor | None = None, + canon_band: tuple[int, int] = FULL_BAND, ) -> torch.Tensor: rows = logits.shape[0] - keys = _draft_shard_topk(logits, top_k, chunk, num_warps) + keys = _draft_shard_topk( + logits, top_k, chunk, num_warps, canon=canon, canon_band=canon_band + ) out = torch.empty((rows, num_ranks * top_k * 2), dtype=torch.float32, device=logits.device) _draft_rank_topk_kernel[(rows,)]( keys, @@ -129,11 +154,15 @@ def reference( num_warps: int = NUM_WARPS, rank: int = 0, num_ranks: int = 1, + canon: torch.Tensor | None = None, + canon_band: tuple[int, int] = FULL_BAND, ) -> torch.Tensor: """The rank's top ``top_k`` keys over the whole shard, from torch ops; geometry ignored.""" rows = logits.shape[0] # One chunk spanning the shard: draft_shard_topk's encoding of the rank's top_k. - keys = _draft_shard_topk.reference(logits, top_k, 1 << 16).to(torch.int64) + keys = _draft_shard_topk.reference( + logits, top_k, 1 << 16, canon=canon, canon_band=canon_band + ).to(torch.int64) unsigned = torch.where(keys < 0, keys + 2**32, keys) halves = torch.stack(((unsigned >> 16) & 0xFFFF, unsigned & 0xFFFF), dim=-1).float() packed = torch.zeros((rows, num_ranks, top_k, 2), device=logits.device) @@ -148,9 +177,13 @@ def is_valid( num_warps: int = NUM_WARPS, rank: int = 0, num_ranks: int = 1, + canon: torch.Tensor | None = None, + canon_band: tuple[int, int] = FULL_BAND, **unused, ) -> None: - _draft_shard_topk.is_valid(logits, top_k, chunk, num_warps) + _draft_shard_topk.is_valid( + logits, top_k, chunk, num_warps, canon=canon, canon_band=canon_band + ) if not 0 <= rank < num_ranks: raise ValueError(f"rank {rank} outside [0, {num_ranks})") if triton.cdiv(logits.shape[1], chunk) * top_k > MAX_KEYS: diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/draft_shard_topk.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/draft_shard_topk.py index 1a7c336c10c6..65a206fc5c32 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/draft_shard_topk.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/draft_shard_topk.py @@ -20,6 +20,17 @@ whose signed order is (value desc, column asc) and which is unique per element. ``INT32_MIN`` marks an empty slot; no element encodes to it. + +With ``canon`` (the canonical-key table of a drafter logit transform ``t``, +``kernels/draft_head/logit_transform.py``) the keys order the slice as +``torch.topk`` orders the transformed fp32 logits ``t(x)``. Every key16 outside +the band ``canon_band = (lo, hi)`` is replaced by ``canon[key16]``, the key16 of +its tie group's representative (a tie group is the bf16 values whose +transformed values have the same bits). ``t`` never decreases, so a tie group +is a run of consecutive key16 values and ``canon`` is non-decreasing: two +canonical keys compare as the radix keys of the transformed values do, and +equal ones tie on the column as torch's do. Inside the band every group is a +single value and ``canon`` is the identity, so those elements load nothing. """ from __future__ import annotations @@ -28,7 +39,7 @@ import triton import triton.language as tl -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper #: The draft head this target runs: 202048 / 8 vocabulary rows per rank, K + 1 #: = 8 block positions with K <= 7 drafted, top 16 candidates. @@ -42,20 +53,33 @@ #: The empty-slot key. EMPTY_KEY = -(2**31) +#: Entries of a canonical-key table: one per 16-bit radix key. +CANON_ENTRIES = 1 << 16 + +#: The band of a call without a table: every key16 lies inside it. +FULL_BAND = (0, CANON_ENTRIES - 1) + +#: The non-NaN 16-bit radix keys, -inf (0x007F) to +inf (0xFF80). +FINITE_KEYS = (0x007F, 0xFF80) + @triton.jit def _draft_shard_topk_kernel( bits_ptr, keys_ptr, + canon_ptr, shard, bits_row_stride, keys_row_stride, rank, num_ranks, keys_per_rank, + canon_lo, + canon_hi, CHUNK: tl.constexpr, TOP_K: tl.constexpr, HALVES: tl.constexpr, + CANON: tl.constexpr, ): row = tl.program_id(0) chunk = tl.program_id(1) @@ -66,6 +90,10 @@ def _draft_shard_topk_kernel( key16 = tl.where((bits & 0x8000) != 0, bits ^ 0xFFFF, bits ^ 0x8000) is_nan = ((bits & 0x7F80) == 0x7F80) & ((bits & 0x7F) != 0) key16 = tl.where(is_nan, 0xFFFF, key16) + if CANON: + # Only the keys outside the identity band read the table. + far = ((key16 < canon_lo) | (key16 > canon_hi)) & in_shard + key16 = tl.where(far, tl.load(canon_ptr + key16, mask=far, other=0), key16) key = ((key16 ^ 0x8000) << 16) | (0xFFFF - local) key = tl.where(in_shard, key, -2147483648) @@ -91,34 +119,91 @@ def _draft_shard_topk_kernel( tl.store(keys_ptr + row.to(tl.int64) * keys_row_stride + chunk * TOP_K + slots, out) -def _key16(bits: torch.Tensor) -> torch.Tensor: - """torch's radix key of bf16 bit patterns, as int32 in [0, 0xFFFF].""" +def radix_key16(bits: torch.Tensor) -> torch.Tensor: + """torch's radix key of bf16 bit patterns, as int32 in [0, 0xFFFF]; every NaN is 0xFFFF.""" b = bits.to(torch.int32) & 0xFFFF key = torch.where((b & 0x8000) != 0, b ^ 0xFFFF, b ^ 0x8000) is_nan = ((b & 0x7F80) == 0x7F80) & ((b & 0x7F) != 0) return torch.where(is_nan, torch.full_like(key, 0xFFFF), key) +def canon_reason(canon: torch.Tensor, canon_band: tuple[int, int], device) -> str | None: + """Why ``canon`` is not a canonical-key table the kernels read correctly, or None. + + A table is ``[65536]`` int32, contiguous, on ``device``; holds 16-bit keys; + is non-decreasing over the non-NaN keys (``FINITE_KEYS``), so it keeps + their order; is the identity on ``canon_band`` (the kernels skip the load + there); and maps 0xFFFF (every NaN) to itself. + """ + if canon.dtype != torch.int32 or tuple(canon.shape) != (CANON_ENTRIES,): + return f"canon must be [{CANON_ENTRIES}] int32, got {canon.dtype} {tuple(canon.shape)}" + if not canon.is_contiguous() or canon.device != torch.device(device): + return "canon must be contiguous and on the logits' device" + lo, hi = canon_band + if not 0 <= lo <= hi < CANON_ENTRIES: + return f"canon_band {canon_band} is not an interval of 16-bit keys" + table = canon.cpu() + if int(table.min()) < 0 or int(table.max()) >= CANON_ENTRIES: + return "canon holds a value outside the 16-bit keys" + ordered = table[FINITE_KEYS[0] : FINITE_KEYS[1] + 1] + if bool((ordered[1:] < ordered[:-1]).any()): + return "canon is not non-decreasing over the non-NaN keys" + if not torch.equal(table[lo : hi + 1], torch.arange(lo, hi + 1, dtype=torch.int32)): + return f"canon is not the identity on its band {tuple(canon_band)}" + if int(table[CANON_ENTRIES - 1]) != CANON_ENTRIES - 1: + return "canon must map 0xFFFF (NaN) to itself" + return None + + class DraftShardTopK(OpWrapper): """Top ``top_k`` sort keys of every ``chunk`` columns of a bf16 logits slice. - logits [M, shard] bf16, stride(-1) == 1, row stride free; shard <= 65536 - top_k candidates kept per chunk (a power of two) - chunk columns per program (a power of two >= top_k) + logits [M, shard] bf16, stride(-1) == 1, row stride free; shard <= 65536 + top_k candidates kept per chunk (a power of two) + chunk columns per program (a power of two >= top_k) num_warps + rank, num_ranks + with ``rank`` set, the all-reduce layout (see below) + canon None, or a canonical-key table ``[65536]`` int32 (``canon_reason``) + canon_band ``(lo, hi)``: the key16 interval on which ``canon`` is the identity Returns ``[M, ceil(shard / chunk) * top_k]`` int32: chunk ``c``'s keys at columns ``[c * top_k, (c + 1) * top_k)``, largest first, ``EMPTY_KEY`` where - the chunk has fewer than ``top_k`` columns. Integer selection: the output - is exact, not approximate. + the chunk has fewer than ``top_k`` columns. With ``rank`` set, ``[M, + num_ranks * ceil(shard / chunk) * top_k * 2]`` fp32: this rank's keys as + (high, low) 16-bit halves at its slot, +0.0 in every other rank's. Integer + selection: the output is exact, not approximate. """ ARCHS = frozenset({Arch.SM_103}) + CELLS: tuple[Cell, ...] = ( + Cell( + "one request at TP8: 7 rows of a 25256-column shard, int32 keys for the all-gather", + {"rows": 7, "shard": SHARD, "halves": False, "transform": None}, + ), + Cell( + "one request at TP8, fp32 halves for the MNNVL all-reduce", + {"rows": 7, "shard": SHARD, "halves": True, "transform": None}, + ), + Cell( + "one request at TP8, fp32 halves, the Muse Spark Open drafter transform's canon", + {"rows": 7, "shard": SHARD, "halves": True, "transform": "musespark"}, + ), + Cell( + "eight requests at TP4: a 50512-column shard, int32 keys, the drafter transform", + {"rows": 56, "shard": 2 * SHARD, "halves": False, "transform": "musespark"}, + ), + Cell( + "24 requests at TP8, fp32 halves, the drafter transform", + {"rows": 168, "shard": SHARD, "halves": True, "transform": "musespark"}, + ), + ) note = ( "Used with draft_topk_merge, which decodes the keys; the pair is exact " - "against torch.topk over the gathered fp32 logits. NaN payloads are not " - "carried (a NaN logit ranks first, as in torch, and decodes to the " - "canonical NaN)." + "against torch.topk over the gathered fp32 logits, and with canon (and the " + "merge's values table) against torch.topk over the drafter's transformed " + "logits. NaN payloads are not carried (a NaN logit ranks first, as in " + "torch, and decodes to the canonical NaN)." ) def __call__( @@ -129,6 +214,8 @@ def __call__( num_warps: int = NUM_WARPS, rank: int | None = None, num_ranks: int = 1, + canon: torch.Tensor | None = None, + canon_band: tuple[int, int] = FULL_BAND, ) -> torch.Tensor: rows, shard = logits.shape chunks = triton.cdiv(shard, chunk) @@ -141,18 +228,23 @@ def __call__( else: keys = torch.empty((rows, keys_per_rank), dtype=torch.int32, device=logits.device) bits = logits.view(torch.int16) + lo, hi = canon_band _draft_shard_topk_kernel[(rows, chunks)]( bits, keys, + bits if canon is None else canon, shard, bits.stride(0), keys.stride(0), rank if halves else 0, num_ranks, keys_per_rank, + lo, + hi, CHUNK=chunk, TOP_K=top_k, HALVES=halves, + CANON=canon is not None, num_warps=num_warps, ) return keys @@ -165,11 +257,18 @@ def reference( num_warps: int = NUM_WARPS, rank: int | None = None, num_ranks: int = 1, + canon: torch.Tensor | None = None, + canon_band: tuple[int, int] = FULL_BAND, ) -> torch.Tensor: """The same keys from torch ops; ``num_warps`` is ignored.""" rows, shard = logits.shape local = torch.arange(shard, device=logits.device, dtype=torch.int64) - key16 = _key16(logits.contiguous().view(torch.int16)).to(torch.int64) + key16 = radix_key16(logits.contiguous().view(torch.int16)) + if canon is not None: + lo, hi = canon_band + far = (key16 < lo) | (key16 > hi) + key16 = torch.where(far, canon.to(key16.device)[key16.long()], key16) + key16 = key16.to(torch.int64) # The unsigned 32-bit key: (value desc, column asc) as plain integer order. key32 = (key16 << 16) | (0xFFFF - local) out = [] @@ -203,6 +302,8 @@ def is_valid( num_warps: int = NUM_WARPS, rank: int | None = None, num_ranks: int = 1, + canon: torch.Tensor | None = None, + canon_band: tuple[int, int] = FULL_BAND, **unused, ) -> None: if logits.dtype != torch.bfloat16 or logits.dim() != 2: @@ -217,3 +318,7 @@ def is_valid( ) if rank is not None and not 0 <= rank < num_ranks: raise ValueError(f"rank {rank} outside [0, {num_ranks})") + if canon is not None: + reason = canon_reason(canon, canon_band, logits.device) + if reason is not None: + raise ValueError(reason) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/draft_topk_merge.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/draft_topk_merge.py index 97117280ab01..9dbca5a955a7 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/draft_topk_merge.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/draft_topk_merge.py @@ -25,6 +25,16 @@ Both are exact: a bf16 widens to fp32 exactly and a bf16 * bf16 product fits fp32's significand, so these equal the built-in tensors bit for bit, and the built-in einsum over them is unchanged. + +With a drafter logit transform ``t`` the keys are canonical +(``draft_shard_topk``'s ``canon``) and ``values`` is the transform's table +(``kernels/draft_head/logit_transform.py``): ``values[key16]`` is ``t`` of the +bf16 value of ``key16``, computed by the drafter's own method. The selection, +the arrangement and the network above run on canonical keys and their decoded +bf16 values unchanged, which compare as the transformed values do, so they give +``torch.topk(t(logits))``'s ids and order; the unary values are then +``values[key16]``, bitwise ``t`` of every candidate, since all members of a tie +group transform to the same bits. """ from __future__ import annotations @@ -33,7 +43,9 @@ import triton import triton.language as tl -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from .draft_shard_topk import CANON_ENTRIES, SHARD, radix_key16 #: ``torch.topk`` sorts k <= 32 results with this many bitonic slots. SORT_SLOTS = 32 @@ -93,6 +105,7 @@ def _draft_topk_merge_kernel( unary_ptr, pred_out_ptr, succ_out_ptr, + values_ptr, rows_per_rank, num_ranks, num_keys, @@ -110,6 +123,7 @@ def _draft_topk_merge_kernel( SLOTS: tl.constexpr, RANK_BLOCK: tl.constexpr, HALVES: tl.constexpr, + TRANSFORM: tl.constexpr, ): # Every program of a row runs the (small) merge; each gathers one RANK_BLOCK # slice of the codebook rows, so the gathers spread over more SMs. @@ -194,7 +208,11 @@ def _draft_topk_merge_kernel( ids = result >> 16 if part == 0: tl.store(ids_ptr + row * TOP_K + kk, ids) - tl.store(unary_ptr + row * TOP_K + kk, _decode_value(result & 0xFFFF)) + if TRANSFORM: + # The drafter transform of each candidate: one table entry per tie group. + tl.store(unary_ptr + row * TOP_K + kk, tl.load(values_ptr + (result & 0xFFFF))) + else: + tl.store(unary_ptr + row * TOP_K + kk, _decode_value(result & 0xFFFF)) # Selector operands for this row (successors) and the next one (predecessors). out_rows = (row * TOP_K + kk[:, None]).to(tl.int64) * rank_dim @@ -252,15 +270,77 @@ class DraftTopKMerge(OpWrapper): successor_codebook [vocab, rank] bf16, stride(-1) == 1 top_k <= 32, a power of two num_ranks ranks in the gather (the tensor-parallel size) + rank_block, num_warps launch geometry + values None, or a drafter logit transform's table [65536] + fp32 (key16 -> transformed value); the keys are then + ``draft_shard_topk``'s canonical keys of that transform Returns ``(candidate_ids [B, K, top_k] int64, unary [B, K, top_k] fp32, predecessors [B, K, top_k, rank] fp32, successors [B, K, top_k, rank] fp32)``: ``candidate_ids`` / ``unary`` are ``torch.topk(full_logits.float(), top_k)``'s - indices / values, and the operands are ``dflash2_score_edges``'s + indices / values -- with ``values``, ``torch.topk(t(full_logits))``'s, ``t`` + the transform -- and the operands are ``dflash2_score_edges``'s ``predecessors * gate.float()[:, :, None]`` and ``successors``. Exact. """ ARCHS = frozenset({Arch.SM_103}) + CELLS: tuple[Cell, ...] = ( + Cell( + "one request at TP8, int32 all-gather keys", + { + "requests": 1, + "positions": 7, + "shard": SHARD, + "ranks": 8, + "layout": "int32", + "transform": None, + }, + ), + Cell( + "one request at TP8, fp32 all-reduce halves", + { + "requests": 1, + "positions": 7, + "shard": SHARD, + "ranks": 8, + "layout": "halves", + "transform": None, + }, + ), + Cell( + "one request at TP8, fp32 halves, the Muse Spark Open drafter transform", + { + "requests": 1, + "positions": 7, + "shard": SHARD, + "ranks": 8, + "layout": "halves", + "transform": "musespark", + }, + ), + Cell( + "two requests at TP4, int32 keys, the drafter transform", + { + "requests": 2, + "positions": 7, + "shard": 2 * SHARD, + "ranks": 4, + "layout": "int32", + "transform": "musespark", + }, + ), + Cell( + "24 requests at TP8, rank top-k halves, the drafter transform", + { + "requests": 24, + "positions": 7, + "shard": SHARD, + "ranks": 8, + "layout": "rank", + "transform": "musespark", + }, + ), + ) note = ( "Tie order is torch 2.14's (SortUtils.cuh bitonicSort, k <= 32 -> 32 " "slots), which torch does not document. draft_head.install checks it " @@ -282,6 +362,7 @@ def __call__( num_ranks: int, rank_block: int = RANK_BLOCK, num_warps: int = NUM_WARPS, + values: torch.Tensor | None = None, ): rows = num_requests * num_positions halves = gathered_keys.dtype == torch.float32 @@ -307,6 +388,7 @@ def __call__( unary, preds, succs, + unary if values is None else values, rows, num_ranks, num_keys, @@ -324,6 +406,7 @@ def __call__( SLOTS=SORT_SLOTS, RANK_BLOCK=rank_block, HALVES=halves, + TRANSFORM=values is not None, num_warps=num_warps, ) return ids, unary, preds, succs @@ -342,6 +425,7 @@ def reference( num_ranks: int, rank_block: int = RANK_BLOCK, num_warps: int = NUM_WARPS, + values: torch.Tensor | None = None, ): """``torch.topk`` over the candidates scattered into a -inf vocabulary row. @@ -349,21 +433,28 @@ def reference( Every element outside the candidate lists is below the k-th value or tied with it at a higher id, so replacing it with -inf changes neither - the selection nor the arrangement torch sorts. + the selection nor the arrangement torch sorts. With ``values`` the + candidates carry the decoded canonical values, which order and tie as + the transformed values do; each selected value is then replaced by + ``values`` at its key16. """ rows = num_requests * num_positions ranks = num_ranks if gathered_keys.dtype == torch.float32: gathered_keys = unpack_halves(gathered_keys, ranks) vocab_size = ranks * shard - vocab, values = _decode_keys(gathered_keys, ranks, shard) + vocab, decoded = _decode_keys(gathered_keys, ranks, shard) dense = torch.full((rows, vocab_size), float("-inf"), device=gathered_keys.device) row_index = torch.arange(rows, device=gathered_keys.device)[:, None].expand_as(vocab[:rows]) for r in range(ranks): - ids_r, vals_r = vocab[r * rows : (r + 1) * rows], values[r * rows : (r + 1) * rows] + ids_r, vals_r = vocab[r * rows : (r + 1) * rows], decoded[r * rows : (r + 1) * rows] keep = ids_r >= 0 dense[row_index[keep], ids_r[keep]] = vals_r[keep] unary, ids = torch.topk(dense, top_k, dim=-1) + if values is not None: + # Decoded canonical values are bf16-exact: their key16 indexes the table. + key16 = radix_key16(unary.to(torch.bfloat16).view(torch.int16)) + unary = values.to(unary.device)[key16.long()] ids = ids.view(num_requests, num_positions, top_k) unary = unary.view(num_requests, num_positions, top_k) successors = successor_codebook[ids].float() @@ -385,8 +476,20 @@ def is_valid( successor_codebook: torch.Tensor, top_k: int, num_ranks: int, + rank_block: int = RANK_BLOCK, + num_warps: int = NUM_WARPS, + values: torch.Tensor | None = None, **unused, ) -> None: + if values is not None and ( + values.dtype != torch.float32 + or tuple(values.shape) != (CANON_ENTRIES,) + or not values.is_contiguous() + or values.device != gathered_keys.device + ): + raise ValueError( + f"values must be a contiguous [{CANON_ENTRIES}] fp32 table on the keys' device" + ) rows = num_requests * num_positions ranks = num_ranks if gathered_keys.dtype == torch.int32: diff --git a/tests/unittest/_torch/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/draft_head/_draft_head_cases.py b/tests/unittest/_torch/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/draft_head/_draft_head_cases.py new file mode 100644 index 000000000000..1930de55d98b --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/draft_head/_draft_head_cases.py @@ -0,0 +1,328 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Inputs and built-in references shared by the draft-head tests (fast path ``draft_head``). + +Not a test module: the entry tests under ``ops/spec`` and the draft-head tests +here import it. Every reference transform is the drafter's own +``DFlashForCausalLM.apply_draft_logit_transform``, never a formula written +here: on the device that method divides by the softcap as a multiplication by +its fp32 reciprocal, which a hand-written ``x / softcap`` does not reproduce. +""" + +from __future__ import annotations + +import types +import zlib + +import torch +import torch.nn.functional as F + +from tensorrt_llm._torch._experimental.modeling_v2.models.musespark_open.muse_spark_open_816b_a42b_nvfp4__sm_103__tep.kernels.draft_head import ( # noqa: E501 + head as draft_head, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.musespark_open.muse_spark_open_816b_a42b_nvfp4__sm_103__tep.kernels.draft_head import ( # noqa: E501 + logit_transform, +) +from tensorrt_llm._torch.models.modeling_dflash import DFlash2CandidateSelector, DFlashForCausalLM + +draft_shard_topk = draft_head.draft_shard_topk +draft_topk_merge = draft_head.draft_topk_merge +draft_path_walk = draft_head.draft_path_walk + +VOCAB, HIDDEN, RANK, TOP_K, TP = 202048, 8192, 256, 16, 8 + +#: The Muse Spark Open DFlash 2 drafter's declared transform: (output_multiplier, softcap). +MUSESPARK = (0.1767766952966369, 20.0) + +#: Transform configurations the table tests cover besides MUSESPARK. +OTHER_TRANSFORMS = { + "multiplier only": (0.5, None), + "softcap only": (1.0, 20.0), + "softcap 30": (1.0, 30.0), +} + +#: Logit families of the end-to-end and merge matrices. +FAMILIES = ( + "gemm1", + "gemm10", + "gemm100", + "gemm400", + "grid", + "special_nan", + "straddle_up", + "straddle_down", + "ladder", + "saturated", + "edges", +) + + +def seed_of(text: str) -> int: + return zlib.crc32(text.encode()) + + +# --------------------------------------------------------------------------- drafters + + +_DRAFTERS: dict = {} +_TRANSFORMS: dict = {} + + +def drafter(multiplier: float = 1.0, softcap: float | None = None) -> DFlashForCausalLM: + """A ``DFlashForCausalLM`` holding only the logit transform attributes (kept alive here).""" + key = (multiplier, softcap) + if key not in _DRAFTERS: + model = DFlashForCausalLM.__new__(DFlashForCausalLM) + torch.nn.Module.__init__(model) + model.logit_multiplier, model.logit_softcap = multiplier, softcap + _DRAFTERS[key] = model + return _DRAFTERS[key] + + +def transform(config=MUSESPARK) -> logit_transform.DraftLogitTransform: + """The draft head's tables of ``config``'s transform on the current device (built once).""" + if config not in _TRANSFORMS: + built = logit_transform.build(drafter(*config), torch.device("cuda")) + assert built is not None + _TRANSFORMS[config] = built + return _TRANSFORMS[config] + + +def builtin_transform(config, logits: torch.Tensor) -> torch.Tensor: + """The built-in transform: the drafter's own method, the call the DFlash worker makes.""" + model = drafter(*config) + return DFlashForCausalLM.apply_draft_logit_transform(model, logits) + + +# --------------------------------------------------------------------------- logits + + +def _keys_values(keys) -> torch.Tensor: + return logit_transform.bf16_of_key16(torch.as_tensor(keys, device="cuda")) + + +def _groups(config): + """(upper, lower): the largest tie group above +0.0 and below -0.0 of ``config``'s transform.""" + groups = logit_transform.tie_groups(transform(config)) + upper = max((g for g in groups if g[0] >= 0x8000), key=lambda g: g[1] - g[0]) + lower = max((g for g in groups if g[1] <= 0x7FFF), key=lambda g: g[1] - g[0]) + return upper, lower + + +def logits(kind: str, rows: int, vocab: int, seed: int, shard: int | None = None, config=MUSESPARK): + """bf16 draft logits ``[rows, vocab]`` of one value family. + + The tie families take their values from ``config``'s tie groups: + ``straddle_up`` ten copies each of two distinct members of the upper + saturating group at low and at high ids (both orders across rows), + ``straddle_down`` the same for the lower group above a row of lower-group + members, ``ladder`` consecutive bf16 values across the upper edge of the + identity band, ``saturated`` whole rows of upper-group members with +inf, + ``edges`` upper-group members on every shard's first and last columns. + """ + g = torch.Generator(device="cuda").manual_seed(seed) + + def pick(n): + return torch.randint(0, vocab, (rows, n), generator=g, device="cuda") + + def grid(levels=8): + return ( + torch.randint(0, levels, (rows, vocab), generator=g, device="cuda").float() * 0.25 - 3.0 + ) + + def members(group, shape, with_inf=False): + first = group[0] if with_inf else max(group[0], 0x0080) + last = group[1] if with_inf else min(group[1], 0xFF7F) + keys = torch.randint(first, last + 1, shape, generator=g, device="cuda") + return _keys_values(keys) + + if kind.startswith("gemm"): + sigma = float(kind[4:]) + x = torch.randn(rows, vocab, generator=g, device="cuda") * sigma + elif kind == "logits": + x = torch.randn(rows, vocab, generator=g, device="cuda") * 3.0 + elif kind == "grid": + x = torch.randint(0, 40, (rows, vocab), generator=g, device="cuda").float() * 0.125 + elif kind == "zeros": + x = torch.zeros(rows, vocab, device="cuda") + elif kind in ("special", "special_nan"): + x = -(torch.rand(rows, vocab, generator=g, device="cuda") + 1.0) + for value, n in ((0.0, 12), (-0.0, 12), (float("inf"), 3), (float("-inf"), 2000)): + idx = pick(n) + x.scatter_(-1, idx, torch.full(idx.shape, value, device="cuda")) + if kind == "special_nan": + idx = pick(4) + x.scatter_(-1, idx, torch.full(idx.shape, float("nan"), device="cuda")) + elif kind in ("straddle_up", "straddle_down", "ladder", "saturated", "edges"): + upper, lower = _groups(config) + x = grid().to(torch.bfloat16) + quarter = vocab // 4 + low = torch.randperm(quarter, generator=g, device="cuda")[:10] + high = vocab - 1 - torch.randperm(quarter, generator=g, device="cuda")[:10] + if kind in ("straddle_up", "straddle_down"): + group = upper if kind == "straddle_up" else lower + if kind == "straddle_down": + x = members(lower, (rows, vocab), with_inf=True) + a = _keys_values([max(group[0], 0x0080)]) + b = _keys_values([min(group[1], 0xFF7F)]) + for r in range(rows): + first, second = (a, b) if r % 2 == 0 else (b, a) + x[r, low] = first + x[r, high] = second + elif kind == "ladder": + hi = transform(config).hi + ladder = _keys_values(torch.arange(hi - 23, min(hi + 9, 0xFF80))) + for r in range(rows): + x[r, torch.randperm(vocab, generator=g, device="cuda")[: ladder.numel()]] = ladder + elif kind == "saturated": + x = members(upper, (rows, vocab), with_inf=True) + else: + assert shard is not None, "the edges family places its ties on shard edges" + cols = torch.tensor( + [c for r in range(vocab // shard) for c in (r * shard, (r + 1) * shard - 1)], + device="cuda", + ) + x[:, cols] = members(upper, (rows, cols.numel())) + return x + else: + raise ValueError(kind) + return x.to(torch.bfloat16) + + +# --------------------------------------------------------------------------- key exchange + + +def split(full: torch.Tensor, ranks: int) -> list: + shard = full.shape[1] // ranks + return [full[:, r * shard : (r + 1) * shard].contiguous() for r in range(ranks)] + + +def shard_keys(full: torch.Tensor, ranks: int, t=None) -> torch.Tensor: + """Every rank's shard keys, concatenated rank-major: what the all-gather returns.""" + return torch.cat( + [draft_shard_topk(s, **draft_head.canon_kwargs(t)) for s in split(full, ranks)], dim=0 + ) + + +def packed_keys(full: torch.Tensor, ranks: int, t=None) -> torch.Tensor: + """Every rank's fp32 key halves summed: what the all-reduce returns.""" + packed = [ + draft_shard_topk(s, rank=r, num_ranks=ranks, **draft_head.canon_kwargs(t)) + for r, s in enumerate(split(full, ranks)) + ] + return torch.stack(packed).sum(dim=0) + + +def rank_keys(full: torch.Tensor, ranks: int, t=None) -> torch.Tensor: + """Every rank's own top-k halves (``wide_draft_topk``) summed: what its all-reduce returns.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.musespark_open.muse_spark_open_816b_a42b_nvfp4__sm_103__tep.kernels.wide_draft_topk import ( # noqa: E501 + exchange, + ) + + ex = exchange.RankTopKExchange(None) + return torch.stack( + [ + ex.keys(s, TOP_K, r, ranks, **draft_head.canon_kwargs(t)) + for r, s in enumerate(split(full, ranks)) + ] + ).sum(dim=0) + + +def layout_keys(layout: str, full: torch.Tensor, ranks: int, t=None) -> torch.Tensor: + return {"int32": shard_keys, "halves": packed_keys, "rank": rank_keys}[layout](full, ranks, t) + + +def merge_launch(layout: str, rows: int) -> dict: + """The merge geometry a step of ``rows`` rows runs with in ``layout``.""" + if layout != "rank": + return {} + from tensorrt_llm._torch._experimental.modeling_v2.models.musespark_open.muse_spark_open_816b_a42b_nvfp4__sm_103__tep.kernels.wide_draft_topk import ( # noqa: E501 + exchange, + ) + + return exchange.RankTopKExchange(None).merge_launch(rows) + + +def assert_same_floats(a: torch.Tensor, b: torch.Tensor) -> None: + """Bitwise equality, NaN payloads aside (both NaN counts as equal).""" + assert a.shape == b.shape and a.dtype == b.dtype, (a.shape, b.shape, a.dtype, b.dtype) + assert logit_transform.same_bits(a, b), "values differ in their bits" + + +# --------------------------------------------------------------------------- selector + + +_SELECTOR = None + + +def selector() -> DFlash2CandidateSelector: + """A random selector of the published drafter's shape.""" + global _SELECTOR + if _SELECTOR is None: + g = torch.Generator(device="cuda").manual_seed(7) + pred = (0.05 * torch.randn(VOCAB, RANK, generator=g, device="cuda")).bfloat16() + succ = (0.05 * torch.randn(VOCAB, RANK, generator=g, device="cuda")).bfloat16() + proj = (0.01 * torch.randn(RANK, HIDDEN, generator=g, device="cuda")).bfloat16() + _SELECTOR = DFlash2CandidateSelector(pred, succ, proj, TOP_K, VOCAB, RANK).cuda() + return _SELECTOR + + +def draft_step(requests, positions, kind, seed, ranks=TP): + """Per-rank bf16 shard logits of one draft step, plus its selector inputs.""" + g = torch.Generator(device="cuda").manual_seed(seed) + rows = requests * positions + hidden = torch.randn(rows, HIDDEN, generator=g, device="cuda").bfloat16() + shard = VOCAB // ranks + if kind.startswith("gemm") and requests * positions <= 56: + # The head's GEMM on a random weight of the given logit scale. + sigma = float(kind[4:]) + shards = [] + for r in range(ranks): + weight = ( + torch.randn(shard, HIDDEN, generator=g, device="cuda") * (sigma / HIDDEN**0.5) + ).bfloat16() + shards.append(F.linear(hidden, weight)) + del weight + else: + full = logits(kind, rows, VOCAB, seed, shard=shard) + shards = split(full, ranks) + anchor = torch.randint(0, VOCAB, (requests,), generator=g, device="cuda") + return shards, hidden.view(requests, positions, HIDDEN), anchor + + +class DraftStub: + """What the worker's selector path reads off the drafter.""" + + select_candidate_path = DFlashForCausalLM.select_candidate_path + + def __init__(self, candidate_selector): + self.candidate_selector = candidate_selector + + +def builtin_full_logits(shards, requests, positions, config=None) -> torch.Tensor: + """The built-in ``[B, K, vocab]`` fp32 logits: each shard transformed, then all-gathered.""" + if config is None: + parts = [s.float() for s in shards] + else: + parts = [builtin_transform(config, s) for s in shards] + return torch.cat(parts, dim=-1).float().view(requests, positions, VOCAB) + + +def builtin_candidates(shards, requests, positions, config=None): + """``(unary, ids)``: ``torch.topk`` as the built-in selector runs it.""" + return torch.topk(builtin_full_logits(shards, requests, positions, config), TOP_K, dim=-1) + + +def builtin_draft_tokens(shards, block_hidden, anchor, config=None): + """The built-in path on the gathered logits: topk, selector, walk, scatter, argmax.""" + requests, positions, _ = block_hidden.shape + full = builtin_full_logits(shards, requests, positions, config) + unary, ids = torch.topk(full, TOP_K, dim=-1) + block = DFlashForCausalLM.select_candidate_path( + DraftStub(selector()), full, ids, unary, block_hidden, anchor.long() + ) + return torch.argmax(block.reshape(-1, VOCAB), dim=-1).to(torch.int32).view(requests, positions) + + +def mapping(rank: int, ranks: int = TP): + return types.SimpleNamespace(tp_size=ranks, tp_rank=rank) diff --git a/tests/unittest/_torch/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/draft_head/test_modeling_v2_draft_head_transform.py b/tests/unittest/_torch/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/draft_head/test_modeling_v2_draft_head_transform.py new file mode 100644 index 000000000000..344042a7091b --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/draft_head/test_modeling_v2_draft_head_transform.py @@ -0,0 +1,734 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test: the vocab-sharded draft head with the DFlash drafter's logit transform on. + +The claim is that with a drafter that declares a logit transform (the Muse +Spark Open drafter's multiplier and softcap) the draft head's fast path still +runs -- the deferring transform hands the bf16 shard through, no full-vocab +all-gather and no full-vocab ``torch.topk`` -- and drafts the built-in tokens +bit for bit: the built-in path transforms each shard with the drafter's own +method, all-gathers, runs ``torch.topk`` and the selector. Then every path +that keeps the built-in computation (sampled drafts, d2t, another layout, +another tensor, a stale record) gets the built-in transform bit for bit, the +fold's tail walk takes the same inputs, CUDA-graph replay follows new logits, +and the install is atomic, idempotent and checked. + +TP ranks are emulated in one process: each rank's drafter, worker and head +shard are built here, and the key exchange returns what every rank's kernels +give for the same step. +""" + +import types + +import pytest +import torch + +__extra_import_path__ = ["."] +from _draft_head_cases import ( # noqa: E402 + FAMILIES, + HIDDEN, + MUSESPARK, + TOP_K, + TP, + VOCAB, + DraftStub, + assert_same_floats, + builtin_candidates, + builtin_draft_tokens, + builtin_full_logits, + builtin_transform, + draft_head, + draft_step, + logit_transform, + packed_keys, + rank_keys, + seed_of, + selector, + shard_keys, +) + +from tensorrt_llm._torch.models.modeling_dflash import DFlashForCausalLM # noqa: E402 +from tensorrt_llm._torch.modules.linear import TensorParallelMode # noqa: E402 +from tensorrt_llm._torch.modules.logits_processor import LogitsProcessor # noqa: E402 + +assert torch.cuda.is_available(), "the vocab-sharded draft head requires a CUDA device" + +K = 7 + + +# --------------------------------------------------------------------------- the emulated group + + +class StubHead: + """A gathered bf16 column-parallel head whose shard GEMM returns ``source`` (a copy).""" + + tp_mode = TensorParallelMode.COLUMN + gather_output = True + padding_size = 0 + gather_output_sizes = None + lora = None + enable_lm_head_tp_in_adp = False + bias = None + + def __init__(self, source: torch.Tensor, ranks: int): + self.source = source + self.tp_size = ranks + self.out_features = source.shape[-1] + self.weight = torch.zeros(1, dtype=torch.bfloat16, device="cuda") + + def apply_linear(self, hidden_states, bias): + return self.source.clone() + + +class Exchange: + """The key all-reduce, emulated: every rank's keys of ``full`` (read at call time) summed.""" + + def __init__(self, full, ranks, t): + self.full, self.ranks, self.t = full, ranks, t + self.calls = 0 + + def __call__(self, packed): + self.calls += 1 + chunk = packed_keys(self.full, self.ranks, self.t) + if packed.shape == chunk.shape: + return chunk + return rank_keys(self.full, self.ranks, self.t) + + +def make_drafter(config, lm_head): + model = DFlashForCausalLM.__new__(DFlashForCausalLM) + torch.nn.Module.__init__(model) + model.logit_multiplier, model.logit_softcap = config + model.candidate_selector = selector() + model.logits_processor = LogitsProcessor() + model.lm_head = lm_head + return model + + +def make_worker(rank, ranks=TP): + from tensorrt_llm._torch.speculative.dflash import DFlashWorker + from tensorrt_llm.mapping import Mapping + + config = types.SimpleNamespace(max_draft_len=K, attention_backend="VANILLA") + return DFlashWorker(config, Mapping(world_size=ranks, tp_size=ranks, rank=rank)) + + +def spec_metadata(requests=1, advanced=False): + return types.SimpleNamespace( + wants_advanced_draft_sampling=advanced, + is_all_greedy_sample=not advanced, + is_cuda_graph=False, + draft_vocab_size=VOCAB, + vocab_size=VOCAB, + batch_indices_cuda=torch.arange(requests, dtype=torch.int32, device="cuda"), + ) + + +@pytest.fixture +def single_process(monkeypatch): + """The install's TP agreement in one process: each emulated rank agrees with itself.""" + monkeypatch.setattr(draft_head, "agreed_reason", lambda mapping, reason, transform: reason) + + +def install_rank(rank, local, full, ranks=TP, config=MUSESPARK, layout="halves"): + """One emulated rank: its worker and drafter with the draft head installed.""" + worker = make_worker(rank, ranks) + drafter = make_drafter(config, StubHead(local, ranks)) + llm = types.SimpleNamespace(spec_worker=worker, draft_model=drafter) + exchange = Exchange(full, ranks, None) + assert draft_head.install(llm, make_allreduce=lambda mapping: exchange) + exchange.t = worker.__dict__.get(draft_head.TRANSFORM_ATTR) + if layout == "int32": + worker.__dict__["_draft_head_allreduce"] = None + elif layout == "rank": + from tensorrt_llm._torch._experimental.modeling_v2.models.musespark_open.muse_spark_open_816b_a42b_nvfp4__sm_103__tep.kernels.wide_draft_topk import ( # noqa: E501 + exchange as wide, + ) + + assert wide.install(llm) + return llm + + +def patch_allgather(monkeypatch, full, ranks, t, gathered=None): + """The head's NCCL all-gather: keys along dim 0; ``gathered`` (or a raise) along dim -1.""" + + def allgather(tensor, mapping, dim=-1, sizes=None): + if dim == 0: + return shard_keys(full, ranks, t) + if gathered is None: + raise AssertionError("the fast path must not all-gather the vocabulary") + return gathered.clone() + + monkeypatch.setattr(draft_head, "allgather", allgather) + + +def drafter_step(worker, drafter, hidden, anchor, md): + """``DFlashWorker._forward_impl`` from the drafter's logits processor to the draft tokens. + + The statements the pin test (``test_modeling_v2_draft_transform_pin.py``) + holds the built-in forward to. + """ + num_gens = hidden.shape[0] // K + gen_logits = drafter.logits_processor(hidden, drafter.lm_head, None, True) + apply_logit_transform = getattr(drafter, "apply_draft_logit_transform", None) + if apply_logit_transform is not None: + gen_logits = apply_logit_transform(gen_logits) + vocab_size = gen_logits.shape[-1] + gen_logits = gen_logits.reshape(num_gens, K, vocab_size) + gen_logits = worker._refine_block_logits(drafter, gen_logits, {}, md) + gen_logits = worker._apply_dflash2_selector( + drafter, gen_logits, hidden.reshape(num_gens, K, -1), anchor, md + ) + return worker.sample_draft_tokens(gen_logits, md, num_gens, num_contexts=0) + + +# --------------------------------------------------------------------------- end to end + + +@pytest.mark.parametrize("kind", FAMILIES) +@pytest.mark.parametrize("layout", ["int32", "halves", "rank"]) +@pytest.mark.parametrize("ranks", [4, 8]) +@pytest.mark.parametrize("requests", [1, 2, 8, 24]) +def test_selector_draft_tokens_equal_builtin(requests, ranks, layout, kind, monkeypatch): + """``selector_draft_tokens`` on every rank's bf16 shard with the transform = the built-in drafts.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.musespark_open.muse_spark_open_816b_a42b_nvfp4__sm_103__tep.kernels.wide_draft_topk import ( # noqa: E501 + exchange as wide, + ) + + keep = make_drafter(MUSESPARK, None) + t = logit_transform.build(keep, torch.device("cuda")) + shards, block_hidden, anchor = draft_step( + requests, K, kind, seed_of(f"e2e{requests}{ranks}{kind}"), ranks=ranks + ) + full = torch.cat(shards, dim=-1) + expected = builtin_draft_tokens(shards, block_hidden, anchor, MUSESPARK) + ref_unary, ref_ids = builtin_candidates(shards, requests, K, MUSESPARK) + patch_allgather(monkeypatch, full, ranks, t) + allreduce = { + "int32": None, + "halves": Exchange(full, ranks, t), + "rank": wide.RankTopKExchange(Exchange(full, ranks, t)), + }[layout] + for rank in range(ranks): + seen = {} + + def walk(edges, unary, ids, seen=seen): + seen["unary"], seen["ids"] = unary.clone(), ids.clone() + return draft_head.draft_path_walk(edges, unary, ids) + + got = draft_head.selector_draft_tokens( + shards[rank].view(requests, K, -1), + block_hidden, + anchor, + selector(), + types.SimpleNamespace(tp_size=ranks, tp_rank=rank), + allreduce, + walk=walk, + transform=t, + ) + assert torch.equal(got, expected), f"rank {rank}: {got.tolist()} != {expected.tolist()}" + assert torch.equal(seen["ids"], ref_ids) + assert_same_floats(seen["unary"], ref_unary) + + +def test_transform_is_a_required_keyword(): + with pytest.raises(TypeError): + draft_head.selector_draft_tokens(None, None, None, None, None) + + +@pytest.mark.parametrize("kind", ["gemm100", "straddle_up", "saturated"]) +@pytest.mark.parametrize("requests", [1, 8]) +def test_path_walk_tail_takes_the_builtin_inputs(requests, kind, monkeypatch): + """The fold's tail walk on the fast path's unary values = the tail on the built-in candidates.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.musespark_open.muse_spark_open_816b_a42b_nvfp4__sm_103__tep.ops.spec.draft_path_walk_tail import ( # noqa: E501 + draft_path_walk_tail, + ) + + keep = make_drafter(MUSESPARK, None) + t = logit_transform.build(keep, torch.device("cuda")) + shards, block_hidden, anchor = draft_step(requests, K, kind, seed_of(f"tail{requests}{kind}")) + full = torch.cat(shards, dim=-1) + g = torch.Generator(device="cuda").manual_seed(requests) + accepted = torch.randint(0, VOCAB, (requests, K + 1), generator=g, device="cuda").int() + num_accepted = torch.randint(1, K + 2, (requests,), generator=g, device="cuda").int() + batch_indices = torch.arange(requests, dtype=torch.int32, device="cuda") + kv_lens = torch.randint(10, 100, (requests,), generator=g, device="cuda").int() + rewind = (1 - num_accepted).int() + + # The built-in candidates and einsum operands, then the tail's reference. + sel = selector() + ref_unary, ref_ids = builtin_candidates(shards, requests, K, MUSESPARK) + gate = sel.hidden_projection(block_hidden.to(sel.hidden_projection.weight.dtype)) + successors = sel.successor_codebook[ref_ids].float() + predecessor_ids = torch.cat( + (anchor.long()[:, None, None].expand(-1, 1, TOP_K), ref_ids[:, :-1]), 1 + ) + predecessors = sel.predecessor_codebook[predecessor_ids].float() * gate.float()[:, :, None] + edges = torch.einsum("blpr,blcr->blpc", predecessors, successors) + ref_kv = kv_lens.clone() + ref_tokens, ref_next = draft_path_walk_tail.reference( + edges, ref_unary, ref_ids, accepted, num_accepted, batch_indices, ref_kv, rewind + ) + patch_allgather(monkeypatch, full, TP, t) + ours_kv = kv_lens.clone() + + def walk(edges, unary, ids): + return draft_path_walk_tail( + edges, unary, ids, accepted, num_accepted, batch_indices, ours_kv, rewind + ) + + tokens, next_tokens = draft_head.selector_draft_tokens( + shards[5].view(requests, K, -1), + block_hidden, + anchor, + sel, + types.SimpleNamespace(tp_size=TP, tp_rank=5), + Exchange(full, TP, t), + walk=walk, + transform=t, + ) + assert torch.equal(tokens, ref_tokens) and torch.equal(next_tokens, ref_next) + assert torch.equal(ours_kv, ref_kv) + assert torch.equal(tokens, builtin_draft_tokens(shards, block_hidden, anchor, MUSESPARK)) + + +# --------------------------------------------------------------------------- the gate + + +def _forbid_full_vocab_topk(monkeypatch): + real_topk = torch.topk + + def topk(x, *args, **kwargs): + if x.shape[-1] == VOCAB: + raise AssertionError("the fast path must not run torch.topk over the full vocabulary") + return real_topk(x, *args, **kwargs) + + monkeypatch.setattr(torch, "topk", topk) + + +@pytest.mark.parametrize("layout", ["int32", "halves", "rank"]) +@pytest.mark.parametrize("requests", [1, 5]) +def test_the_fast_path_runs_with_the_transform_on(requests, layout, monkeypatch, single_process): + shards, block_hidden, anchor = draft_step( + requests, K, "gemm100", seed_of(f"gate{layout}{requests}") + ) + full = torch.cat(shards, dim=-1) + expected = builtin_draft_tokens(shards, block_hidden, anchor, MUSESPARK) + hidden = block_hidden.view(requests * K, HIDDEN) + llms = [install_rank(rank, shards[rank], full, layout=layout) for rank in range(TP)] + t = llms[0].spec_worker.__dict__[draft_head.TRANSFORM_ATTR] + for llm in llms: + # The drafter instance carries the deferring transform; the class is unchanged. + assert isinstance( + vars(llm.draft_model)[draft_head.TRANSFORM_METHOD], draft_head._DeferredTransform + ) + assert DFlashForCausalLM.apply_draft_logit_transform is not None + patch_allgather(monkeypatch, full, TP, t) + _forbid_full_vocab_topk(monkeypatch) + calls = [] + real = draft_head.selector_draft_tokens + + def spy(*args, **kwargs): + calls.append(kwargs["transform"]) + return real(*args, **kwargs) + + monkeypatch.setattr(draft_head, "selector_draft_tokens", spy) + md = spec_metadata(requests) + for rank, llm in enumerate(llms): + worker, drafter = llm.spec_worker, llm.draft_model + out = drafter.logits_processor(hidden, drafter.lm_head, None, True) + assert out.dtype == torch.bfloat16 + handed = drafter.apply_draft_logit_transform(out) + assert handed is out, "the deferring transform must hand the shard through" + block = handed.reshape(requests, K, -1) + refined = worker._refine_block_logits(drafter, block, {}, md) + assert refined.dtype == torch.bfloat16 and refined.data_ptr() == out.data_ptr() + assert worker._shard_selector_applies(refined, md, VOCAB) + stand_in = worker._apply_dflash2_selector(drafter, refined, block_hidden, anchor, md) + tokens = worker.sample_draft_tokens(stand_in, md, requests, num_contexts=0) + assert torch.equal(tokens, expected), f"rank {rank}" + assert calls[-1] is worker.__dict__[draft_head.TRANSFORM_ATTR] + # The whole sequence once more, as the built-in forward runs it. + assert torch.equal(drafter_step(worker, drafter, hidden, anchor, md), expected) + + +def test_the_fast_path_runs_under_the_fold_tail(monkeypatch, single_process): + from tensorrt_llm._torch._experimental.modeling_v2.models.musespark_open.muse_spark_open_816b_a42b_nvfp4__sm_103__tep.kernels.verify_fold import ( # noqa: E501 + fold, + ) + + requests = 2 + shards, block_hidden, anchor = draft_step(requests, K, "straddle_up", seed_of("fold")) + full = torch.cat(shards, dim=-1) + expected = builtin_draft_tokens(shards, block_hidden, anchor, MUSESPARK) + hidden = block_hidden.view(requests * K, HIDDEN) + g = torch.Generator(device="cuda").manual_seed(3) + accepted = torch.randint(0, VOCAB, (requests, K + 1), generator=g, device="cuda").int() + num_accepted = torch.randint(1, K + 2, (requests,), generator=g, device="cuda").int() + llms = {rank: install_rank(rank, shards[rank], full) for rank in (0, 6)} + for llm in llms.values(): + assert fold.install(llm.spec_worker, verify_fold=False, path_walk_tail=True) + assert draft_head._defers_transform(type(llm.spec_worker)) + t = llms[0].spec_worker.__dict__[draft_head.TRANSFORM_ATTR] + patch_allgather(monkeypatch, full, TP, t) + _forbid_full_vocab_topk(monkeypatch) + for rank, llm in llms.items(): + worker, drafter = llm.spec_worker, llm.draft_model + md = spec_metadata(requests) + kv_lens = torch.full((requests,), 50, dtype=torch.int32, device="cuda") + attn = types.SimpleNamespace(num_contexts=0, num_seqs=requests, kv_lens_cuda=kv_lens) + worker._kv_rewind_amount = (1 - num_accepted).int() + worker._kv_rewind_pending, worker._kv_rewind_nc, worker._kv_rewind_bs = True, 0, requests + step = fold._Step(attn_metadata=attn, spec_metadata=md, draft_model=drafter, warmup=False) + step.accepted, step.num_accepted = accepted, num_accepted + worker.__dict__["_fold_step"] = step + try: + tokens = drafter_step(worker, drafter, hidden, anchor, md) + finally: + worker.__dict__["_fold_step"] = None + assert step.tail is not None and torch.equal(tokens, expected), f"rank {rank}" + assert torch.equal(kv_lens, (50 - (1 - num_accepted)).clamp(min=0).int()) + + +def test_the_real_lm_head_hands_its_shard_through(single_process): + """With the gathered column-parallel ``LMHead`` the processor's shard reaches the selector untouched.""" + from tensorrt_llm._torch.modules.embedding import LMHead + from tensorrt_llm.mapping import Mapping + + rank = 3 + head = LMHead( + VOCAB, + HIDDEN, + dtype=torch.bfloat16, + mapping=Mapping(world_size=TP, tp_size=TP, rank=rank), + tensor_parallel_mode=TensorParallelMode.COLUMN, + gather_output=True, + reduce_output=False, + ).cuda() + with torch.no_grad(): + head.weight.copy_((torch.randn_like(head.weight, dtype=torch.float32) * 0.05).bfloat16()) + worker = make_worker(rank) + drafter = make_drafter(MUSESPARK, head) + llm = types.SimpleNamespace(spec_worker=worker, draft_model=drafter) + assert draft_head.install(llm, make_allreduce=lambda mapping: object()) + hidden = torch.randn(K, HIDDEN, device="cuda").bfloat16() + out = drafter.logits_processor(hidden, head, None, True) + assert out.dtype == torch.bfloat16 and out.shape == (K, VOCAB // TP) + assert drafter.apply_draft_logit_transform(out) is out + refined = worker._refine_block_logits(drafter, out.reshape(1, K, -1), {}, spec_metadata()) + assert refined.dtype == torch.bfloat16 and worker.__dict__[draft_head.OWED_ATTR] + + +# --------------------------------------------------------------------------- built-in paths + + +def _one_rank(monkeypatch, kind="gemm400", requests=1, rank=2): + shards, block_hidden, anchor = draft_step(requests, K, kind, seed_of(f"fallback{kind}{rank}")) + full = torch.cat(shards, dim=-1) + llm = install_rank(rank, shards[rank], full) + worker, drafter = llm.spec_worker, llm.draft_model + hidden = block_hidden.view(requests * K, HIDDEN) + return shards, block_hidden, anchor, full, worker, drafter, hidden + + +def _builtin_shard(shard, requests): + return builtin_transform(MUSESPARK, shard).view(requests, K, -1) + + +def test_sampled_drafts_get_the_builtin_transform(monkeypatch, single_process): + shards, _, _, _, worker, drafter, hidden = _one_rank(monkeypatch) + out = drafter.logits_processor(hidden, drafter.lm_head, None, True) + handed = drafter.apply_draft_logit_transform(out) + refined = worker._refine_block_logits( + drafter, handed.reshape(1, K, -1), {}, spec_metadata(advanced=True) + ) + # What the inherited _sample_dflash2_block receives: the built-in fp32 shard, bit for bit. + assert refined.dtype == torch.float32 + assert_same_floats(refined, _builtin_shard(shards[2], 1)) + assert not worker.__dict__[draft_head.OWED_ATTR] + + +@pytest.mark.parametrize("why", ["d2t", "stride"]) +def test_layouts_the_selector_cannot_take_get_the_builtin_path(why, monkeypatch, single_process): + shards, block_hidden, anchor, full, worker, drafter, hidden = _one_rank(monkeypatch) + md = spec_metadata() + out = drafter.logits_processor(hidden, drafter.lm_head, None, True) + handed = drafter.apply_draft_logit_transform(out) + block = handed.reshape(1, K, -1) + if why == "d2t": + worker._d2t = torch.zeros(VOCAB, dtype=torch.int64, device="cuda") + else: + block = block.as_strided(block.shape, (block.shape[-1] * K, 1, K), 0) + assert block.stride(-1) != 1 and block.data_ptr() == out.data_ptr() + refined = worker._refine_block_logits(drafter, block, {}, md) + assert refined.dtype == torch.float32 + assert_same_floats(refined, builtin_transform(MUSESPARK, block)) + if why == "d2t": + gathered = builtin_full_logits(shards, 1, K, MUSESPARK) + patch_allgather(monkeypatch, full, TP, None, gathered=gathered) + got = worker._apply_dflash2_selector(drafter, refined, block_hidden, anchor, md) + ref = DFlashForCausalLM.select_candidate_path( + DraftStub(selector()), + gathered.clone(), + *reversed(torch.topk(gathered, TOP_K, dim=-1)), + block_hidden, + anchor.long(), + ) + assert torch.equal(got, ref) + + +def test_any_other_tensor_gets_the_builtin_transform(monkeypatch, single_process): + shards, _, _, _, worker, drafter, hidden = _one_rank(monkeypatch) + out = drafter.logits_processor(hidden, drafter.lm_head, None, True) + other = out.clone() + assert_same_floats( + drafter.apply_draft_logit_transform(other), builtin_transform(MUSESPARK, other) + ) + # The shard itself once the record is spent. + assert_same_floats(drafter.apply_draft_logit_transform(out), builtin_transform(MUSESPARK, out)) + # A full-vocab bf16 output of the wrapped processor is never deferred. + full_vocab = torch.randn(K, VOCAB, device="cuda").bfloat16() + + class Inner(torch.nn.Module): + def forward(self, *args): + return full_vocab + + drafter.logits_processor.inner = Inner() + assert drafter.logits_processor(hidden, None, None, True) is full_vocab + assert draft_head.SHARD_ATTR not in vars(drafter.logits_processor) + assert_same_floats( + drafter.apply_draft_logit_transform(full_vocab), builtin_transform(MUSESPARK, full_vocab) + ) + # A drafter whose attributes moved away from the tables' transform is not deferred. + out = drafter.logits_processor(hidden, drafter.lm_head, None, True) + drafter.logit_softcap = 30.0 + handed = drafter.apply_draft_logit_transform(out) + assert handed is not out and handed.dtype == torch.float32 + assert_same_floats(handed, builtin_transform((MUSESPARK[0], 30.0), out)) + + +def test_a_stale_record_changes_nothing(monkeypatch, single_process): + shards, _, _, _, worker, drafter, hidden = _one_rank(monkeypatch) + out = drafter.logits_processor(hidden, drafter.lm_head, None, True) + drafter.__dict__[draft_head.DEFERRED_ATTR] = (out.data_ptr() + 2, out.numel()) + refined = worker._refine_block_logits(drafter, out.reshape(1, K, -1), {}, spec_metadata()) + assert refined.dtype == torch.float32 + assert_same_floats(refined, _builtin_shard(shards[2], 1)) + + +def test_the_selector_applies_the_transform_without_refine(monkeypatch, single_process): + """Bypassing refine: bf16 logits at the selector are untransformed, so it still drafts exactly.""" + shards, block_hidden, anchor, full, worker, drafter, hidden = _one_rank( + monkeypatch, "saturated" + ) + t = worker.__dict__[draft_head.TRANSFORM_ATTR] + patch_allgather(monkeypatch, full, TP, t) + out = drafter.logits_processor(hidden, drafter.lm_head, None, True) + handed = drafter.apply_draft_logit_transform(out) + md = spec_metadata() + stand_in = worker._apply_dflash2_selector( + drafter, handed.reshape(1, K, -1), block_hidden, anchor, md + ) + tokens = worker.sample_draft_tokens(stand_in, md, 1, num_contexts=0) + assert torch.equal(tokens, builtin_draft_tokens(shards, block_hidden, anchor, MUSESPARK)) + + +def test_a_step_that_owes_the_transform_refuses_to_sample(monkeypatch, single_process): + shards, _, _, _, worker, drafter, hidden = _one_rank(monkeypatch) + out = drafter.logits_processor(hidden, drafter.lm_head, None, True) + handed = drafter.apply_draft_logit_transform(out) + refined = worker._refine_block_logits(drafter, handed.reshape(1, K, -1), {}, spec_metadata()) + assert refined.dtype == torch.bfloat16 + with pytest.raises(draft_head.DraftTransformStateError): + worker.sample_draft_tokens(refined, spec_metadata(), 1, num_contexts=0) + + +def test_a_class_composed_above_with_its_own_refine_is_not_deferred(monkeypatch, single_process): + shards, _, _, _, worker, drafter, hidden = _one_rank(monkeypatch) + + class Above(type(worker)): + def _refine_block_logits(self, draft_model, gen_logits, inputs, spec_metadata): + return super()._refine_block_logits(draft_model, gen_logits, inputs, spec_metadata) + + worker.__class__ = Above + assert not draft_head._defers_transform(Above) + out = drafter.logits_processor(hidden, drafter.lm_head, None, True) + handed = drafter.apply_draft_logit_transform(out) + assert handed.dtype == torch.float32 + assert_same_floats(handed, builtin_transform(MUSESPARK, out)) + + +# --------------------------------------------------------------------------- CUDA graph + + +@pytest.mark.parametrize("tail", [False, True]) +def test_cuda_graph_replay_follows_new_logits(tail, monkeypatch, single_process): + from tensorrt_llm._torch._experimental.modeling_v2.models.musespark_open.muse_spark_open_816b_a42b_nvfp4__sm_103__tep.kernels.verify_fold import ( # noqa: E501 + fold, + ) + + requests, rank = 1, 2 + shards, block_hidden, anchor = draft_step(requests, K, "gemm100", seed_of(f"graph{tail}")) + full = torch.cat(shards, dim=-1) + shard = VOCAB // TP + local = full[:, rank * shard : (rank + 1) * shard] + llm = install_rank(rank, local, full) + worker, drafter = llm.spec_worker, llm.draft_model + t = worker.__dict__[draft_head.TRANSFORM_ATTR] + tables = (t.canon.data_ptr(), t.values.data_ptr()) + patch_allgather(monkeypatch, full, TP, t) + md = spec_metadata(requests) + hidden = block_hidden.view(requests * K, HIDDEN) + kv_lens = torch.full((requests,), 50, dtype=torch.int32, device="cuda") + accepted = torch.zeros(requests, K + 1, dtype=torch.int32, device="cuda") + num_accepted = torch.ones(requests, dtype=torch.int32, device="cuda") + if tail: + assert fold.install(worker, verify_fold=False, path_walk_tail=True) + + def step(): + if tail: + attn = types.SimpleNamespace(num_contexts=0, num_seqs=requests, kv_lens_cuda=kv_lens) + worker._kv_rewind_amount = (1 - num_accepted).int() + worker._kv_rewind_pending, worker._kv_rewind_nc, worker._kv_rewind_bs = ( + True, + 0, + requests, + ) + s = fold._Step(attn_metadata=attn, spec_metadata=md, draft_model=drafter, warmup=False) + s.accepted, s.num_accepted = accepted, num_accepted + worker.__dict__["_fold_step"] = s + try: + return drafter_step(worker, drafter, hidden, anchor, md) + finally: + worker.__dict__["_fold_step"] = None + + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + step() + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = step() + assert worker.__dict__[draft_head.TRANSFORM_ATTR] is t + assert (t.canon.data_ptr(), t.values.data_ptr()) == tables + for replay, kind in enumerate(("gemm400", "straddle_up", "saturated", "ladder", "special_nan")): + new_shards, new_hidden, new_anchor = draft_step(requests, K, kind, 60 + replay) + full.copy_(torch.cat(new_shards, dim=-1)) + block_hidden.copy_(new_hidden) + anchor.copy_(new_anchor) + graph.replay() + expected = builtin_draft_tokens(new_shards, new_hidden, new_anchor, MUSESPARK) + assert torch.equal(out, expected), f"replay {replay} ({kind})" + + +# --------------------------------------------------------------------------- install + + +def test_install_is_idempotent(single_process): + shards, _, _ = draft_step(1, K, "gemm3", 5) + full = torch.cat(shards, dim=-1) + llm = install_rank(0, shards[0], full) + t = llm.spec_worker.__dict__[draft_head.TRANSFORM_ATTR] + shadow = vars(llm.draft_model)[draft_head.TRANSFORM_METHOD] + assert draft_head.install(llm, make_allreduce=lambda mapping: 1 / 0) + assert llm.spec_worker.__dict__[draft_head.TRANSFORM_ATTR] is t + assert vars(llm.draft_model)[draft_head.TRANSFORM_METHOD] is shadow + + +def test_an_identity_transform_installs_nothing_new(monkeypatch, single_process): + shards, block_hidden, anchor = draft_step(1, K, "gemm3", 6) + full = torch.cat(shards, dim=-1) + llm = install_rank(1, shards[1], full, config=(1.0, None)) + worker, drafter = llm.spec_worker, llm.draft_model + assert draft_head.TRANSFORM_ATTR not in worker.__dict__ + assert draft_head.TRANSFORM_METHOD not in vars(drafter) + seen = [] + real = draft_head.draft_shard_topk + + class Spy: + def __call__(self, *args, **kwargs): + seen.append(sorted(kwargs)) + return real(*args, **kwargs) + + monkeypatch.setattr(draft_head, "draft_shard_topk", Spy()) + patch_allgather(monkeypatch, full, TP, None) + md = spec_metadata() + tokens = drafter_step(worker, drafter, block_hidden.view(K, HIDDEN), anchor, md) + assert torch.equal(tokens, builtin_draft_tokens(shards, block_hidden, anchor)) + assert seen and all("canon" not in kw for kw in seen) + + +def test_tie_order_probe_agrees_with_this_torch_under_the_transform(): + keep = make_drafter(MUSESPARK, None) + t = logit_transform.build(keep, torch.device("cuda")) + assert draft_head.tie_order_reason(selector(), TP, transform=t) is None + assert draft_head.tie_order_reason(selector(), 4, transform=t) is None + + +def _stable_topk(x, k, dim=-1): + """A top-k that breaks ties by ascending index: torch's selection, another order.""" + values, ids = torch.sort(x, dim=dim, descending=True, stable=True) + return values.narrow(dim, 0, k), ids.narrow(dim, 0, k) + + +def test_tie_order_probe_has_teeth_under_the_transform(monkeypatch): + keep = make_drafter(MUSESPARK, None) + t = logit_transform.build(keep, torch.device("cuda")) + reason = draft_head.tie_order_reason(selector(), TP, topk=_stable_topk, transform=t) + assert reason is not None and "orders top-k ties differently" in reason + # Without canonical keys the transform probe's saturated ties come out wrong. + monkeypatch.setattr(draft_head, "canon_kwargs", lambda transform: {}) + reason = draft_head.tie_order_reason(selector(), TP, transform=t) + assert reason is not None and "probe 1" in reason + + +def test_transform_probe_logits_hold_the_tie_families(): + keep = make_drafter(MUSESPARK, None) + t = logit_transform.build(keep, torch.device("cuda")) + probe = draft_head.transform_probe_logits(VOCAB, VOCAB // TP, t) + assert probe.dtype == torch.bfloat16 and probe.shape == (draft_head.PROBE_ROWS, VOCAB) + again = draft_head.transform_probe_logits(VOCAB, VOCAB // TP, t) + assert torch.equal(probe.view(torch.int16), again.view(torch.int16)), "not deterministic" + transformed = t.real(probe.cuda()) + # Rows 0, 1, 2, 3, 6: more than top_k candidates share the row's largest transformed value. + for row in (0, 1, 2, 3, 6): + top = transformed[row].max() + assert (transformed[row] == top).sum() > TOP_K, row + tied = probe[row][transformed[row].cpu() == top.cpu()] + assert tied.view(torch.int16).unique().numel() > 1, row + + +def test_a_refused_install_changes_nothing(monkeypatch, single_process): + from tensorrt_llm._torch.speculative.dflash import DFlashWorker + + monkeypatch.setattr( + draft_head, "tie_order_reason", lambda *args, **kwargs: "probe rows [3] differ" + ) + worker = make_worker(0) + shards, _, _ = draft_step(1, K, "gemm3", 8) + drafter = make_drafter(MUSESPARK, StubHead(shards[0], TP)) + inner = drafter.logits_processor + llm = types.SimpleNamespace(spec_worker=worker, draft_model=drafter) + built = [] + with pytest.raises(ValueError, match="probe rows"): + draft_head.install(llm, make_allreduce=lambda mapping: built.append(mapping)) + assert type(worker) is DFlashWorker and drafter.logits_processor is inner and not built + assert draft_head.TRANSFORM_ATTR not in worker.__dict__ + assert draft_head.TRANSFORM_METHOD not in vars(drafter) + + +def test_ranks_whose_tables_differ_refuse(monkeypatch): + from tensorrt_llm._torch._experimental.modeling_v2.models.musespark_open.muse_spark_open_816b_a42b_nvfp4__sm_103__tep import ( # noqa: E501 + dispatch, + ) + + keep = make_drafter(MUSESPARK, None) + t = logit_transform.build(keep, torch.device("cuda")) + mine = logit_transform.digest(t) + monkeypatch.setattr(dispatch, "agreed_union", lambda mapping, local: [mine, "another"]) + monkeypatch.setattr(dispatch, "agreed_refusal", lambda mapping, local: local) + reason = draft_head.agreed_reason(object(), None, t) + assert reason is not None and "differ in their drafter logit transform" in reason + monkeypatch.setattr(dispatch, "agreed_union", lambda mapping, local: list(local)) + assert draft_head.agreed_reason(object(), None, t) is None diff --git a/tests/unittest/_torch/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/draft_head/test_modeling_v2_draft_logit_transform.py b/tests/unittest/_torch/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/draft_head/test_modeling_v2_draft_logit_transform.py new file mode 100644 index 000000000000..a070c92d7fae --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/draft_head/test_modeling_v2_draft_logit_transform.py @@ -0,0 +1,164 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the draft head's drafter logit transform tables (``kernels/draft_head/logit_transform.py``). + +The claim is that ``values`` is the drafter's own transform of every bf16 value +bit for bit, in any layout, and that ``canon`` groups exactly the bf16 values +whose transformed values have the same bits, in runs of consecutive keys, with +the identity on ``[lo, hi]``. Checked exhaustively over the 65536 bf16 values +on this device, for the Muse Spark Open drafter's transform and three others. +""" + +import pytest +import torch + +__extra_import_path__ = ["."] +from _draft_head_cases import ( # noqa: E402 + MUSESPARK, + OTHER_TRANSFORMS, + builtin_transform, + drafter, + logit_transform, + transform, +) + +assert torch.cuda.is_available(), "the vocab-sharded draft head requires a CUDA device" + +CONFIGS = {"musespark": MUSESPARK, **OTHER_TRANSFORMS} +ALL = torch.arange(1 << 16, dtype=torch.int32) + + +def _radix32(v: torch.Tensor) -> torch.Tensor: + word = v.view(torch.int32).to(torch.int64) & 0xFFFFFFFF + key = torch.where(word >= 0x80000000, word ^ 0xFFFFFFFF, word ^ 0x80000000) + return torch.where(v.isnan(), torch.full_like(key, 0xFFFFFFFF), key) + + +@pytest.mark.parametrize("name", list(CONFIGS)) +def test_values_are_the_builtin_transform_of_every_bf16_value(name): + config = CONFIGS[name] + t = transform(config) + x = logit_transform.bf16_of_key16(ALL.cuda()) + assert t.values.dtype == torch.float32 and t.values.shape == (1 << 16,) + assert logit_transform.same_bits(t.values, builtin_transform(config, x)) + + +@pytest.mark.parametrize("requests", [1, 2, 8, 24]) +@pytest.mark.parametrize("name", list(CONFIGS)) +def test_values_hold_in_the_runtime_layout(name, requests): + """``values[key16(x)]`` equals the transform of x laid out as ``[B * 7, shard]``, shuffled.""" + config = CONFIGS[name] + t = transform(config) + rows, shard = requests * 7, 25256 + g = torch.Generator(device="cuda").manual_seed(requests) + keys = torch.randint(0, 1 << 16, (rows, shard), generator=g, device="cuda", dtype=torch.int32) + keys.view(-1)[: 1 << 16] = ALL.cuda()[torch.randperm(1 << 16, generator=g, device="cuda")] + x = logit_transform.bf16_of_key16(keys) + from tensorrt_llm._torch._experimental.modeling_v2.models.musespark_open.muse_spark_open_816b_a42b_nvfp4__sm_103__tep.ops.spec.draft_shard_topk import ( # noqa: E501 + radix_key16, + ) + + got = t.values[radix_key16(x.view(torch.int16)).long()] + assert logit_transform.same_bits(got, builtin_transform(config, x)) + + +@pytest.mark.parametrize("name", list(CONFIGS)) +def test_canon_groups_exactly_the_ties(name): + t = transform(CONFIGS[name]) + canon, values = t.canon.cpu().long(), t.values.cpu() + radix = _radix32(values) + finite = torch.arange(0x007F, 0xFF81) + # Non-decreasing over the non-NaN keys, so the tie groups are runs of consecutive keys. + assert (radix[finite][1:] >= radix[finite][:-1]).all() + # Same canonical key exactly when the transformed values have the same bits. + assert torch.equal(radix[canon[finite]], radix[finite]) + split = canon[finite][1:] != canon[finite][:-1] + assert torch.equal(split, radix[finite][1:] != radix[finite][:-1]) + # The band: canon is the identity on [lo, hi] and on no larger interval around zero. + assert torch.equal(canon[t.lo : t.hi + 1], torch.arange(t.lo, t.hi + 1)) + assert t.lo == 0 or canon[t.lo - 1] != t.lo - 1 + assert t.hi == 0xFFFF or canon[t.hi + 1] != t.hi + 1 + # NaN maps to itself; +0.0 and -0.0 stay apart and keep their sign. + assert int(canon[0xFFFF]) == 0xFFFF + assert int(canon[0x8000]) == 0x8000 and int(canon[0x7FFF]) == 0x7FFF + assert values[0x8000].item() == 0.0 and not torch.signbit(values[0x8000]) + assert torch.signbit(values[0x7FFF]) and values[0x7FFF].item() == 0.0 + assert logit_transform.unsupported_reason(t, 25256) is None + + +def test_musespark_band_and_saturation_on_this_device(): + """The identity band holds every logit of magnitude up to 512; +-inf saturate to +-softcap.""" + t = transform(MUSESPARK) + lo_value = float(logit_transform.bf16_of_key16(torch.tensor([t.lo])).float()) + hi_value = float(logit_transform.bf16_of_key16(torch.tensor([t.hi])).float()) + groups = logit_transform.tie_groups(t) + first_up = min(g[0] for g in groups if g[0] >= 0x8000) + print( + f"identity band [{lo_value}, {hi_value}], first tie above zero at bf16 " + f"{float(logit_transform.bf16_of_key16(torch.tensor([first_up])).float())}, " + f"{len(groups)} tie groups" + ) + assert lo_value <= -512.0 and hi_value >= 512.0 + softcap = MUSESPARK[1] + assert t.values[0xFF80].item() == softcap and t.values[0x007F].item() == -softcap + + +def test_identity_transform_builds_nothing(): + assert logit_transform.build(drafter(1.0, None), torch.device("cuda")) is None + + +def test_an_fp32_probe_would_not_have_found_the_transform(): + """Why build probes in bf16: on fp32 the built-in transform runs in place and returns its input.""" + x = torch.ones(4, device="cuda") + assert builtin_transform(MUSESPARK, x) is x + assert logit_transform.build(drafter(*MUSESPARK), torch.device("cuda")) is not None + + +def test_ensure_transformed_applies_once(): + t = transform(MUSESPARK) + x = torch.randn(3, 9, device="cuda").bfloat16() + y = logit_transform.ensure_transformed(t, x) + assert y.dtype == torch.float32 and logit_transform.same_bits( + y, builtin_transform(MUSESPARK, x) + ) + assert logit_transform.ensure_transformed(t, y) is y + assert logit_transform.ensure_transformed(None, x) is x + + +def test_the_tables_follow_the_drafter_attributes(): + t = transform(MUSESPARK) + model = drafter(*MUSESPARK) + assert t.matches(model) + model.logit_softcap = 30.0 + try: + assert not t.matches(model) + finally: + model.logit_softcap = MUSESPARK[1] + + +def test_a_transform_that_decreases_is_refused(): + t = transform(MUSESPARK) + bad = logit_transform.DraftLogitTransform( + drafter(*MUSESPARK), t.function, t.canon, -t.values, t.lo, t.hi + ) + reason = logit_transform.unsupported_reason(bad, 25256) + assert reason is not None + + +def test_another_method_is_refused(): + t = transform(MUSESPARK) + + def other(model, logits): + return logits.float() + + bad = logit_transform.DraftLogitTransform( + drafter(*MUSESPARK), other, t.canon, t.values, t.lo, t.hi + ) + assert "not DFlashForCausalLM's" in logit_transform.unsupported_reason(bad, 25256) + + +def test_digest_tells_tables_apart(): + a = logit_transform.digest(transform(MUSESPARK)) + assert a == logit_transform.digest(transform(MUSESPARK)) + assert a != logit_transform.digest(transform(OTHER_TRANSFORMS["softcap 30"])) + assert logit_transform.digest(None) != a diff --git a/tests/unittest/_torch/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/draft_head/test_modeling_v2_draft_transform_pin.py b/tests/unittest/_torch/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/draft_head/test_modeling_v2_draft_transform_pin.py new file mode 100644 index 000000000000..77c819887bde --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/draft_head/test_modeling_v2_draft_transform_pin.py @@ -0,0 +1,149 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The built-in DFlash forward's statements the draft head's deferred logit transform relies on. + +The draft head (``kernels/draft_head/head.py``) hands the drafter's bf16 shard +logits through the drafter's ``apply_draft_logit_transform`` untransformed and +applies the transform after its top-k, in ``ShardedDraftHeadDFlashWorker``'s +``_refine_block_logits`` and selector. That is the built-in computation only +while nothing between the transform call and ``_refine_block_logits`` reads the +logits, and only while ``_refine_block_logits`` runs before the selector and +the sampler. This test reads ``DFlashWorker._forward_impl``'s source (no +import, no device) and fails when that sequence changes, so a change to it is +reviewed against the draft head instead of silently reaching a reader with +untransformed logits. +""" + +from __future__ import annotations + +import ast +import importlib.util +from pathlib import Path + +import pytest + +pytestmark = pytest.mark.cpu_only + +#: The statements from the drafter's logits processor to the selector, as ``ast.unparse`` prints them. +SEQUENCE = ( + "gen_logits = draft_model.logits_processor(gen_hidden_states, draft_model.lm_head, " + "attn_metadata, True)", + "apply_logit_transform = getattr(draft_model, 'apply_draft_logit_transform', None)", + "if apply_logit_transform is not None:\n gen_logits = apply_logit_transform(gen_logits)", + "vocab_size = gen_logits.shape[-1]", + "gen_logits = gen_logits.reshape(num_gens, K, vocab_size)", + "gen_logits = self._refine_block_logits(draft_model, gen_logits, inputs, spec_metadata)", + "has_selector = getattr(draft_model, 'has_candidate_selector', False)", +) + + +def _source(relative: str) -> str: + spec = importlib.util.find_spec("tensorrt_llm") + assert spec is not None and spec.submodule_search_locations, "tensorrt_llm is not importable" + return (Path(next(iter(spec.submodule_search_locations))) / relative).read_text() + + +def _forward_impl() -> ast.FunctionDef: + tree = ast.parse(_source("_torch/speculative/dflash.py")) + worker = next(n for n in tree.body if isinstance(n, ast.ClassDef) and n.name == "DFlashWorker") + return next( + n for n in worker.body if isinstance(n, ast.FunctionDef) and n.name == "_forward_impl" + ) + + +def _bodies(node: ast.AST): + for child in ast.walk(node): + for name in ("body", "orelse"): + body = getattr(child, name, None) + if isinstance(body, list) and body and isinstance(body[0], ast.stmt): + yield body + + +def test_nothing_reads_the_draft_logits_between_the_transform_and_refine(): + forward = _forward_impl() + found = [] + for body in _bodies(forward): + printed = [ast.unparse(s) for s in body] + if SEQUENCE[0] in printed: + start = printed.index(SEQUENCE[0]) + found.append(tuple(printed[start : start + len(SEQUENCE)])) + assert len(found) == 1, ( + "the drafter's logits processor call is not where the draft head expects it" + ) + assert found[0] == SEQUENCE, ( + "DFlashWorker._forward_impl changed between the drafter's logits processor and " + "_refine_block_logits; check the draft head's deferred transform against it:\n" + + "\n".join(found[0]) + ) + + +def test_refine_runs_before_the_selector_and_the_samplers(): + forward = _forward_impl() + calls = [ + (node.lineno, node.func.attr) + for node in ast.walk(forward) + if isinstance(node, ast.Call) + and isinstance(node.func, ast.Attribute) + and isinstance(node.func.value, ast.Name) + and node.func.value.id == "self" + and node.func.attr + in ( + "_refine_block_logits", + "_apply_dflash2_selector", + "_sample_dflash2_block", + "sample_draft_tokens", + ) + ] + names = [name for _, name in sorted(calls)] + assert names.count("_refine_block_logits") == 1, names + assert names[0] == "_refine_block_logits", names + + +def test_the_transform_has_one_caller(): + """The deferring transform is consulted only by that call.""" + root = Path(next(iter(importlib.util.find_spec("tensorrt_llm").submodule_search_locations))) + callers = sorted( + str(path.relative_to(root)) + for path in (root / "_torch").rglob("*.py") + if "apply_draft_logit_transform" in path.read_text(encoding="utf-8", errors="ignore") + and "_experimental" not in path.parts + ) + assert callers == ["_torch/models/modeling_dflash.py", "_torch/speculative/dflash.py"], callers + + +#: This target's own sources, relative to the ``tensorrt_llm`` package. +TARGET = ( + "_torch/_experimental/modeling_v2/models/musespark_open/" + "muse_spark_open_816b_a42b_nvfp4__sm_103__tep" +) + + +def test_the_target_composes_no_other_reader_of_deferred_logits(): + """Only the draft head and the fold take deferred logits; only the draft head wraps the processor. + + A worker class of another fast path composed above the draft head that + defined ``_refine_block_logits`` or ``_apply_dflash2_selector`` would see the + deferred bf16 logits first (the draft head then stops deferring for it: + ``head._defers_transform``), and a fast path that replaced the drafter's + logits processor would end the deferral without a word. Read off the + target's sources, so a new one of either is reviewed against the draft head. + """ + root = Path(next(iter(importlib.util.find_spec("tensorrt_llm").submodule_search_locations))) + owners = {"_refine_block_logits": [], "_apply_dflash2_selector": [], "processor": []} + for path in sorted((root / TARGET).rglob("*.py")): + source = path.read_text(encoding="utf-8", errors="ignore") + rel = str(path.relative_to(root / TARGET)) + for node in ast.walk(ast.parse(source)): + if isinstance(node, ast.FunctionDef) and node.name in owners: + owners[node.name].append(rel) + if isinstance(node, ast.Assign) and any( + isinstance(target, ast.Attribute) and target.attr == "logits_processor" + for target in node.targets + ): + owners["processor"].append(rel) + assert owners["_refine_block_logits"] == ["kernels/draft_head/head.py"], owners + assert sorted(owners["_apply_dflash2_selector"]) == [ + "kernels/draft_head/head.py", + "kernels/verify_fold/fold.py", + ], owners + assert owners["processor"] == ["kernels/draft_head/head.py"], owners diff --git a/tests/unittest/_torch/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/test_modeling_v2_draft_rank_topk.py b/tests/unittest/_torch/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/test_modeling_v2_draft_rank_topk.py new file mode 100644 index 000000000000..67920d43b88d --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/test_modeling_v2_draft_rank_topk.py @@ -0,0 +1,82 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for ``draft_rank_topk`` (fast path ``wide_draft_topk``). + +Gated on its cells with ``torch.equal`` against its torch reference, with and +without a drafter logit transform's canonical keys, and the merge over its +output against the merge over ``draft_shard_topk``'s per-chunk keys. +""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.musespark_open.muse_spark_open_816b_a42b_nvfp4__sm_103__tep.ops.spec import ( # noqa: E501 + draft_rank_topk as rank_mod, +) + +__extra_import_path__ = ["../../../../..", "../../kernels/draft_head"] +from _draft_head_cases import ( # noqa: E402 + MUSESPARK, + TOP_K, + VOCAB, + draft_head, + logits, + packed_keys, + rank_keys, + seed_of, + selector, + transform, +) +from _validating import certifying # noqa: E402,F401 + +assert torch.cuda.is_available(), "the vocab-sharded draft head requires a CUDA device" + +draft_rank_topk = rank_mod.draft_rank_topk +CELLS = rank_mod.DraftRankTopK.CELLS + + +@pytest.mark.parametrize("kind", ["gemm3", "special_nan", "straddle_up", "saturated"]) +@pytest.mark.parametrize("cell", CELLS, ids=[c.why[:48] for c in CELLS]) +def test_rank_topk_cells(cell, kind, certifying): # noqa: F811 + certifying(draft_rank_topk) + spec = cell.spec + x = logits(kind, spec["rows"], spec["shard"], seed_of(cell.why + kind)) + kwargs = {} if spec["transform"] is None else draft_head.canon_kwargs(transform(MUSESPARK)) + ranks = VOCAB // spec["shard"] + out = draft_rank_topk(x, TOP_K, 4096, 4, 1, ranks, **kwargs) + ref = draft_rank_topk.reference(x, TOP_K, 4096, 4, 1, ranks, **kwargs) + assert torch.equal(out, ref), f"{(out != ref).sum().item()} key halves differ" + + +@pytest.mark.parametrize("kind", ["gemm400", "straddle_down", "ladder"]) +def test_merge_over_rank_keys_is_the_merge_over_chunk_keys(kind): + t = transform(MUSESPARK) + rows, ranks = 35, 8 + full = logits(kind, rows, VOCAB, seed_of(kind)) + sel = selector() + gate = torch.zeros( + 1, rows, sel.predecessor_codebook.shape[1], dtype=torch.bfloat16, device="cuda" + ) + anchor = torch.zeros(1, dtype=torch.int64, device="cuda") + outs = [] + for keys in (packed_keys(full, ranks, t), rank_keys(full, ranks, t)): + outs.append( + draft_head.draft_topk_merge( + keys, + VOCAB // ranks, + 1, + rows, + anchor, + gate, + sel.predecessor_codebook, + sel.successor_codebook, + TOP_K, + ranks, + values=t.values, + ) + ) + for a, b in zip(*outs): + assert torch.equal( + a.view(torch.int32) if a.is_floating_point() else a, + b.view(torch.int32) if b.is_floating_point() else b, + ) diff --git a/tests/unittest/_torch/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/test_modeling_v2_draft_shard_topk.py b/tests/unittest/_torch/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/test_modeling_v2_draft_shard_topk.py new file mode 100644 index 000000000000..703fa48b68b8 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/test_modeling_v2_draft_shard_topk.py @@ -0,0 +1,100 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for ``draft_shard_topk`` (fast path ``draft_head``). + +Gated on its cells with ``torch.equal`` against its torch reference, in both +output layouts (the int32 keys the all-gather takes and the fp32 halves the +fabric all-reduce takes), with and without a drafter logit transform's +canonical keys. +""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.musespark_open.muse_spark_open_816b_a42b_nvfp4__sm_103__tep.ops.spec import ( # noqa: E501 + draft_shard_topk as shard_mod, +) + +__extra_import_path__ = ["../../../../..", "../../kernels/draft_head"] +from _draft_head_cases import ( # noqa: E402 + MUSESPARK, + draft_head, + draft_shard_topk, + logits, + seed_of, + transform, +) +from _validating import certifying, validating # noqa: E402,F401 + +assert torch.cuda.is_available(), "the vocab-sharded draft head requires a CUDA device" + +CELLS = shard_mod.DraftShardTopK.CELLS +KINDS = ("gemm3", "special_nan", "straddle_up", "straddle_down", "ladder", "saturated") + + +def _call_args(spec, x, rank=3, num_ranks=8): + kwargs = {} + if spec["halves"]: + kwargs.update(rank=rank, num_ranks=num_ranks) + if spec["transform"] is not None: + kwargs.update(draft_head.canon_kwargs(transform(MUSESPARK))) + return kwargs + + +@pytest.mark.parametrize("kind", KINDS) +@pytest.mark.parametrize("cell", CELLS, ids=[c.why[:48] for c in CELLS]) +def test_shard_topk_cells(cell, kind, certifying): # noqa: F811 + certifying(draft_shard_topk) + spec = cell.spec + x = logits( + kind, spec["rows"], spec["shard"], seed_of(cell.why + kind), shard=spec["shard"] // 2 + ) + kwargs = _call_args(spec, x) + out = draft_shard_topk(x, **kwargs) + ref = draft_shard_topk.reference(x, **kwargs) + assert out.dtype == ref.dtype and out.shape == ref.shape + assert torch.equal(out, ref), f"{(out != ref).sum().item()} keys differ" + + +def test_without_a_table_the_keys_are_the_plain_ones(): + """``canon=None`` is the table-free kernel; a table changes only keys outside its band.""" + t = transform(MUSESPARK) + x = logits("gemm100", 7, 25256, 11) + plain = draft_shard_topk(x) + assert torch.equal(plain, draft_shard_topk.reference(x)) + inside = x.clamp(-500, 500) + assert torch.equal( + draft_shard_topk(inside, **draft_head.canon_kwargs(t)), draft_shard_topk(inside) + ) + straddle = logits("straddle_up", 7, 25256, 12) + assert not torch.equal( + draft_shard_topk(straddle, **draft_head.canon_kwargs(t)), draft_shard_topk(straddle) + ) + + +def test_shard_topk_reads_a_strided_row_view(): + t = transform(MUSESPARK) + full = logits("gemm400", 7, 202048, 5) + view = full[:, 3 * 25256 : 4 * 25256] + assert view.stride(0) == 202048 + kwargs = draft_head.canon_kwargs(t) + assert torch.equal( + draft_shard_topk(view, **kwargs), draft_shard_topk(view.contiguous(), **kwargs) + ) + + +def test_shard_topk_guards_refuse_what_the_kernel_cannot_take(): + t = transform(MUSESPARK) + x = torch.zeros(2, 1024, device="cuda", dtype=torch.bfloat16) + with validating(draft_shard_topk): + with pytest.raises(ValueError): # wider than the 16-bit column field + draft_shard_topk(torch.zeros(2, 70000, device="cuda", dtype=torch.bfloat16)) + with pytest.raises(ValueError): # not bf16 + draft_shard_topk(torch.zeros(2, 1024, device="cuda", dtype=torch.float32)) + with pytest.raises(ValueError): # a table that is not the identity on its band + draft_shard_topk(x, canon=t.canon, canon_band=(0, 0xFFFF)) + with pytest.raises(ValueError): # a table of the wrong dtype + draft_shard_topk(x, canon=t.canon.long(), canon_band=t.band) + with pytest.raises(ValueError): # a table that does not keep the order + draft_shard_topk(x, canon=t.canon.flip(0).contiguous(), canon_band=(0x8000, 0x8000)) + draft_shard_topk(x, **draft_head.canon_kwargs(t)) diff --git a/tests/unittest/_torch/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/test_modeling_v2_draft_topk_merge.py b/tests/unittest/_torch/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/test_modeling_v2_draft_topk_merge.py new file mode 100644 index 000000000000..fb562330f2b0 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/test_modeling_v2_draft_topk_merge.py @@ -0,0 +1,128 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for ``draft_topk_merge`` (fast path ``draft_head``) with and without a drafter logit transform. + +Gated with ``torch.equal`` on the ids and bitwise on the values against what +the built-in path computes: ``torch.topk`` over the gathered fp32 logits, or, +with the transform, over the drafter's transform of them (each shard +transformed by the drafter's own method, then gathered). Then against the +entry's own reference, and the einsum operands against ``dflash2_score_edges``. +""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.musespark_open.muse_spark_open_816b_a42b_nvfp4__sm_103__tep.ops.spec import ( # noqa: E501 + draft_topk_merge as merge_mod, +) +from tensorrt_llm._torch.models.modeling_dflash import dflash2_score_edges + +__extra_import_path__ = ["../../../../..", "../../kernels/draft_head"] +from _draft_head_cases import ( # noqa: E402 + FAMILIES, + MUSESPARK, + RANK, + TOP_K, + VOCAB, + assert_same_floats, + builtin_candidates, + draft_topk_merge, + layout_keys, + logits, + merge_launch, + seed_of, + selector, + split, + transform, +) +from _validating import certifying # noqa: E402,F401 + +assert torch.cuda.is_available(), "the vocab-sharded draft head requires a CUDA device" + +CELLS = merge_mod.DraftTopKMerge.CELLS + + +def _run(requests, positions, ranks, layout, kind, config, seed, canon=True): + """(merge outputs, built-in (unary, ids), gate, anchor) for one draft step.""" + rows = requests * positions + t = None if config is None else transform(config) + full = logits(kind, rows, VOCAB, seed, shard=VOCAB // ranks) + g = torch.Generator(device="cuda").manual_seed(seed + 1) + gate = torch.randn(requests, positions, RANK, generator=g, device="cuda").bfloat16() + anchor = torch.randint(0, VOCAB, (requests,), generator=g, device="cuda") + keys = layout_keys(layout, full, ranks, t if canon else None) + sel = selector() + args = ( + keys, + VOCAB // ranks, + requests, + positions, + anchor, + gate, + sel.predecessor_codebook, + sel.successor_codebook, + TOP_K, + ranks, + ) + kwargs = dict(merge_launch(layout, rows)) + if t is not None: + kwargs["values"] = t.values + out = draft_topk_merge(*args, **kwargs) + ref = builtin_candidates(split(full, ranks), requests, positions, config) + return out, ref, args, kwargs, gate, anchor + + +@pytest.mark.parametrize("cell", CELLS, ids=[c.why[:48] for c in CELLS]) +def test_topk_merge_cells(cell, certifying): # noqa: F811 + certifying(draft_topk_merge) + spec = cell.spec + config = None if spec["transform"] is None else MUSESPARK + requests, positions, ranks = spec["requests"], spec["positions"], spec["ranks"] + for kind in ("gemm3", "special_nan", "straddle_up", "saturated"): + (ids, unary, preds, succs), (ref_unary, ref_ids), args, kwargs, gate, anchor = _run( + requests, positions, ranks, spec["layout"], kind, config, seed_of(cell.why + kind) + ) + assert torch.equal(ids, ref_ids), f"{kind}: candidate ids or their order differ" + assert_same_floats(unary, ref_unary) + r_ids, r_unary, r_preds, r_succs = draft_topk_merge.reference(*args, **kwargs) + assert torch.equal(ids, r_ids) + assert_same_floats(unary, r_unary) + assert torch.equal(preds, r_preds) and torch.equal(succs, r_succs) + sel = selector() + builtin = dflash2_score_edges( + sel.predecessor_codebook, sel.successor_codebook, ref_ids, ref_unary, gate, anchor + ) + ours = unary.float()[:, :, None] + torch.einsum("blpr,blcr->blpc", preds, succs) + assert_same_floats(ours, builtin) + + +@pytest.mark.parametrize("kind", FAMILIES) +@pytest.mark.parametrize("layout", ["int32", "halves", "rank"]) +@pytest.mark.parametrize("ranks", [4, 8]) +@pytest.mark.parametrize("requests", [1, 2, 8, 24]) +def test_transformed_candidates_are_the_builtin_ones(requests, ranks, layout, kind): + (ids, unary, _, _), (ref_unary, ref_ids), *_ = _run( + requests, 7, ranks, layout, kind, MUSESPARK, seed_of(f"{requests}{ranks}{layout}{kind}") + ) + differ = (ids != ref_ids).any(dim=-1).view(-1).nonzero().view(-1).tolist() + assert not differ, f"rows {differ} differ from torch.topk over the transformed logits" + assert_same_floats(unary, ref_unary) + + +@pytest.mark.parametrize("kind", ["gemm3", "grid", "special_nan", "edges"]) +@pytest.mark.parametrize("layout", ["int32", "halves", "rank"]) +@pytest.mark.parametrize("requests", [8, 24]) +def test_untransformed_control(requests, layout, kind): + """The same matrix without a transform: a failure here is the merge's, not the transform's.""" + (ids, unary, _, _), (ref_unary, ref_ids), *_ = _run( + requests, 7, 8, layout, kind, None, seed_of(f"control{requests}{layout}{kind}") + ) + assert torch.equal(ids, ref_ids) + assert_same_floats(unary, ref_unary) + + +@pytest.mark.parametrize("kind", ["straddle_up", "straddle_down", "saturated"]) +def test_without_canonical_keys_the_ties_go_wrong(kind): + """The negative control: bf16 keys order saturated logits by their bf16 values, not by id.""" + (ids, _, _, _), (_, ref_ids), *_ = _run(1, 7, 8, "halves", kind, MUSESPARK, 3, canon=False) + assert not torch.equal(ids, ref_ids) From fc6d7acc846ab73871ac22323899d5a20d7ce629 Mon Sep 17 00:00:00 2001 From: Xiao Wang <24860335+xwang233@users.noreply.github.com> Date: Tue, 27 Oct 2026 07:58:36 +0000 Subject: [PATCH 15/16] [None][fix] modeling_v2 musespark_open tep: draft head transform tables follow the drafter, warn on built-in steps The vocab-sharded draft head applies the drafter's logit transform after its top-k from tables built at install from the drafter's logit_multiplier and logit_softcap. Only the deferring transform checked that the drafter still declares those attributes; ShardedDraftHeadDFlashWorker._refine_block_logits and the selector's _draft_head_transform_for took bf16 logits as owing the installed tables. A drafter whose attributes changed to the identity (1.0 / None) after the install hands bf16 logits through, and the sharded selector then ranked and scored them with the stale tables. - _refine_block_logits keeps a step's logits untransformed only when the drafter still matches the tables, and otherwise applies the drafter's class method with its current attributes, the built-in computation. - A new _draft_head_takes decides the sharded route for both selectors (ShardedDraftHeadDFlashWorker and the verify-fold tail walk): the existing shard conditions plus, with a transform installed, tables that still match the drafter. A step it turns down takes the built-in path with the drafter's current transform. _draft_head_transform_for takes the drafter and raises DraftTransformStateError on tables that do not hold, and _draft_head_transformed applies the drafter's current transform. - Every way a step with a transform installed takes the built-in path warns once per reason through tensorrt_llm.logger.warning_once with the "[musespark_open] draft_head:" prefix: the deferring transform declining (logits processor replaced, another tensor, non-bf16 logits, changed attributes, a worker class that overrides the refinement or selector), and the refinement or selector turning a step down (full-vocab logits, layout, sampled drafts, d2t, changed attributes, logits the transform did not defer). The deferring transform records its reason on the drafter so a step warns once, not once per stage. - The install agrees the key exchange's refusal over the TP group too: a rank whose make_allreduce returns None and the ranks that built theirs all raise the same ValueError, before anything changes. Tests (test_modeling_v2_draft_head_transform.py): the drafter's multiplier and softcap changed after the install (to the identity, another softcap, a multiplier only, a softcap only), before the logits processor, after the deferring transform and after the refinement, at 1 and 3 requests and under the fold tail: the built-in transform of this rank's shard reaches the vocabulary all-gather bit for bit, the sharded selector never runs, the draft tokens equal the built-in ones of the new transform, and the tables serve again once the drafter declares their transform; the selector refuses tables that do not hold; each fallback reason warns exactly once over repeated steps, the fast path and a drafter without a transform never warn; the key-exchange refusal is agreed and leaves the worker and drafter as they were. Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com> --- .../kernels/draft_head/head.py | 216 +++++++--- .../kernels/verify_fold/fold.py | 4 +- .../test_modeling_v2_draft_head_transform.py | 369 +++++++++++++++++- 3 files changed, 535 insertions(+), 54 deletions(-) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/draft_head/head.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/draft_head/head.py index 9f9be5c470fc..41b0068da411 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/draft_head/head.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/draft_head/head.py @@ -45,7 +45,11 @@ drafter's transform otherwise, so every other path (sampled drafts, d2t, another layout, any other tensor) gets the built-in transform bit for bit. With a transform installed, bf16 draft logits are untransformed ones: the -transform always returns fp32. +transform always returns fp32. The tables hold only while the drafter's +``logit_multiplier`` and ``logit_softcap`` are the ones they were built from; +a step on which they differ takes the built-in path with the drafter's +current transform. Each reason a step with a transform installed takes the +built-in path is logged once, as a warning (``_warn_builtin``). Installed from the shell's ``load_draft_weights`` (the drafter's selector and the borrowed head exist only after the drafter loads), or from the shell's @@ -81,7 +85,7 @@ from ...ops.spec.draft_shard_topk import CHUNK, FINITE_KEYS, NUM_WARPS, DraftShardTopK from ...ops.spec.draft_topk_merge import SORT_SLOTS, DraftTopKMerge from . import logit_transform as _logit_transform -from .logit_transform import DraftLogitTransform, ensure_transformed +from .logit_transform import DraftLogitTransform draft_shard_topk = DraftShardTopK() draft_topk_merge = DraftTopKMerge() @@ -98,9 +102,16 @@ SHARD_ATTR = "_draft_head_shard" #: Drafter attribute: ``(data_ptr, numel)`` of the logits the deferring transform handed back. DEFERRED_ATTR = "_draft_head_deferred" +#: Drafter attribute: why the deferring transform's last call applied the transform in full. +DECLINED_ATTR = "_draft_head_declined" #: The drafter method the built-in worker calls on the draft logits. TRANSFORM_METHOD = "apply_draft_logit_transform" +#: Why a step takes the built-in path: the drafter's transform attributes moved off the tables. +STALE_TABLES = "the drafter's logit_multiplier or logit_softcap differs from the install's" +#: Why a step takes the built-in path: its bf16 logits are not the shard the transform deferred. +NOT_DEFERRED = "the draft logits are not the shard the deferring transform handed back" + #: Largest shard the 16-bit column field of a shard key holds. MAX_SHARD = 65536 @@ -121,6 +132,25 @@ class DraftTransformStateError(RuntimeError): """A step kept its draft logits untransformed and no selector applied the transform.""" +def _warn_builtin(reason: str) -> None: + """Warn, once per ``reason``, that a draft step with a transform installed takes the built-in path.""" + logger.warning_once( + "[musespark_open] draft_head: a draft step takes the built-in path, with the drafter's " + f"logit transform applied before the top-k ({reason})", + key=f"musespark_open.draft_head.builtin.{reason}", + ) + + +def _builtin_transform(draft_model, logits: torch.Tensor) -> torch.Tensor: + """The drafter's class ``apply_draft_logit_transform`` on ``logits``, with its current attributes. + + The call the built-in worker makes on a drafter without the deferring + transform: fp32 out when the drafter declares a transform, ``logits`` + itself when it declares none. + """ + return getattr(type(draft_model), TRANSFORM_METHOD)(draft_model, logits) + + def gathers_column_shards(lm_head: nn.Module) -> bool: """Whether ``lm_head`` is a column-parallel head whose forward is shard GEMM + all-gather. @@ -275,7 +305,9 @@ class _DeferredTransform: ``ShardedDraftHeadDFlashWorker._refine_block_logits`` (``_defers_transform``), returns that tensor untouched and records its ``(data_ptr, numel)`` on the drafter (``DEFERRED_ATTR``). Every other call returns the drafter's own - transform, the built-in computation. Holds weak references only. + transform with its current attributes, the built-in computation, records + why on the drafter (``DECLINED_ATTR``) and warns once per reason + (``_warn_builtin``). Holds weak references only. """ def __init__(self, worker, draft_model): @@ -289,6 +321,7 @@ def __call__(self, logits: torch.Tensor) -> torch.Tensor: raise DraftTransformStateError("[musespark_open] draft_head: the drafter is gone") transform = None if worker is None else worker.__dict__.get(TRANSFORM_ATTR) draft_model.__dict__[DEFERRED_ATTR] = None + draft_model.__dict__[DECLINED_ATTR] = None # The processor in place now, not the one of the install: a wrapper put # around it later records nothing, and the transform then runs here. processor = getattr(draft_model, "logits_processor", None) @@ -296,17 +329,26 @@ def __call__(self, logits: torch.Tensor) -> torch.Tensor: if isinstance(processor, ShardDraftLogitsProcessor): shard = processor.__dict__.pop(SHARD_ATTR, None) if transform is None: - return getattr(type(draft_model), TRANSFORM_METHOD)(draft_model, logits) - if ( - shard is not None - and shard() is logits - and logits.dtype == torch.bfloat16 - and transform.matches(draft_model) - and _defers_transform(type(worker)) - ): + return _builtin_transform(draft_model, logits) + if not isinstance(processor, ShardDraftLogitsProcessor): + reason = f"the drafter's logits processor is {type(processor).__name__}" + elif shard is None or shard() is not logits: + reason = "the transformed tensor is not the shard the logits processor returned" + elif logits.dtype != torch.bfloat16: + reason = f"the drafter's logits processor returned {logits.dtype} logits" + elif not transform.matches(draft_model): + reason = STALE_TABLES + elif not _defers_transform(type(worker)): + reason = ( + f"the worker class {type(worker).__name__} overrides the draft head's " + "refinement or selector" + ) + else: draft_model.__dict__[DEFERRED_ATTR] = (logits.data_ptr(), logits.numel()) return logits - return transform.real(logits) + draft_model.__dict__[DECLINED_ATTR] = reason + _warn_builtin(reason) + return _builtin_transform(draft_model, logits) def _owner(cls: type, name: str, after: Optional[type] = None) -> Optional[type]: @@ -381,67 +423,126 @@ def _pdl_head_rows(self, key, ids): rows = ranges[key] return ids if rows is None else _drafter_pdl.DraftRowRange(ids, *rows) + def _shard_reason(self, gen_logits, spec_metadata, full_vocab) -> Optional[str]: + """Why the sharded selector cannot take ``gen_logits`` on this step, or None.""" + if gen_logits.shape[-1] == full_vocab: + return "full-vocab draft logits" + if gen_logits.dtype != torch.bfloat16: + return f"{gen_logits.dtype} draft logits" + if ( + gen_logits.dim() != 3 + or gen_logits.stride(-1) != 1 + or gen_logits.shape[-1] * self.mapping.tp_size != full_vocab + ): + return "draft logits outside the [B, K, shard] layout" + if spec_metadata.wants_advanced_draft_sampling: + return "sampled drafts" + if self._d2t is not None: + return "a d2t draft vocab map" + return None + def _shard_selector_applies(self, gen_logits, spec_metadata, full_vocab) -> bool: - mapping = self.mapping - return ( - gen_logits.dim() == 3 - and gen_logits.dtype == torch.bfloat16 - and gen_logits.stride(-1) == 1 - and gen_logits.shape[-1] * mapping.tp_size == full_vocab - and not spec_metadata.wants_advanced_draft_sampling - and self._d2t is None - ) + return self._shard_reason(gen_logits, spec_metadata, full_vocab) is None + + def _draft_head_takes(self, draft_model, gen_logits, spec_metadata, full_vocab) -> bool: + """Whether the sharded selector drafts this step. + + ``_shard_selector_applies``, and, with a transform installed, tables + that still match ``draft_model``. A step it turns down with a transform + installed takes the built-in path, where bf16 logits get + ``_draft_head_transformed``, and warns once per reason. + """ + reason = self._shard_reason(gen_logits, spec_metadata, full_vocab) + transform = self.__dict__.get(TRANSFORM_ATTR) + if transform is None: + return reason is None + if reason is None and not transform.matches(draft_model): + reason = STALE_TABLES + if reason is None: + return True + if gen_logits.dtype == torch.bfloat16: + # Logits transformed upstream (fp32) were warned for where that happened. + _warn_builtin(reason) + return False def _refine_block_logits(self, draft_model, gen_logits, inputs, spec_metadata): """The drafter's logit transform, unless this step's sharded selector applies it. The deferring transform handed back the processor's bf16 shard untransformed (``DEFERRED_ATTR``). These logits stay bf16 only when they - are that tensor (the reshape keeps its address and size) and this - step takes the sharded selector, which then owes the transform - (``OWED_ATTR``); any other bf16 logits get the drafter's transform - here. Then the next class's refinement, the built-in identity. + are that tensor (the reshape keeps its address and size), this step + takes the sharded selector and the drafter still declares the + transform the tables hold; the selector then owes the transform + (``OWED_ATTR``). Any other bf16 logits get the drafter's current + transform here, with a one-time warning per reason. Then the next + class's refinement, the built-in identity. """ deferred = draft_model.__dict__.pop(DEFERRED_ATTR, None) + declined = draft_model.__dict__.pop(DECLINED_ATTR, None) self.__dict__[OWED_ATTR] = False transform = self.__dict__.get(TRANSFORM_ATTR) if transform is not None and gen_logits.dtype == torch.bfloat16: full_vocab = draft_model.candidate_selector.predecessor_codebook.shape[0] - if ( - deferred == (gen_logits.data_ptr(), gen_logits.numel()) - and gen_logits.shape[-1] != full_vocab - and self._shard_selector_applies(gen_logits, spec_metadata, full_vocab) - ): + # A declined call has applied the drafter's current transform already, + # which keeps bf16 logits only when that transform is the identity. + reason = declined + if reason is None and deferred != (gen_logits.data_ptr(), gen_logits.numel()): + reason = NOT_DEFERRED + if reason is None: + reason = self._shard_reason(gen_logits, spec_metadata, full_vocab) + if reason is None and not transform.matches(draft_model): + reason = STALE_TABLES + if reason is None: self.__dict__[OWED_ATTR] = True else: - gen_logits = ensure_transformed(transform, gen_logits) + _warn_builtin(reason) + gen_logits = self._draft_head_transformed(draft_model, gen_logits) return super()._refine_block_logits(draft_model, gen_logits, inputs, spec_metadata) - def _draft_head_transform_for(self, gen_logits) -> Optional[DraftLogitTransform]: - """The transform the sharded selector applies to ``gen_logits`` (bf16 means untransformed).""" + def _draft_head_transform_for(self, draft_model, gen_logits) -> Optional[DraftLogitTransform]: + """The transform the sharded selector applies to ``gen_logits`` (bf16 means untransformed). + + For a step ``_draft_head_takes`` admits, which checked that the tables + still match ``draft_model``. + """ self.__dict__[OWED_ATTR] = False transform = self.__dict__.get(TRANSFORM_ATTR) - return transform if gen_logits.dtype == torch.bfloat16 else None + if transform is None or gen_logits.dtype != torch.bfloat16: + return None + if not transform.matches(draft_model): + raise DraftTransformStateError( + "[musespark_open] draft_head: the sharded selector was asked for tables that no " + f"longer hold ({STALE_TABLES})" + ) + return transform + + def _draft_head_transformed(self, draft_model, gen_logits) -> torch.Tensor: + """``gen_logits`` with the drafter's current transform applied once, for a built-in path. - def _draft_head_transformed(self, gen_logits) -> torch.Tensor: - """``gen_logits`` with the drafter's transform applied once, for a built-in selector.""" + With a transform installed, bf16 logits are untransformed or were + transformed by an identity: both get the drafter's class method with + its current attributes, which is what the built-in worker computes. + """ self.__dict__[OWED_ATTR] = False - return ensure_transformed(self.__dict__.get(TRANSFORM_ATTR), gen_logits) + if self.__dict__.get(TRANSFORM_ATTR) is None or gen_logits.dtype != torch.bfloat16: + return gen_logits + return _builtin_transform(draft_model, gen_logits) def _apply_dflash2_selector( self, draft_model, gen_logits, block_hidden_states, anchor_tokens, spec_metadata ): selector = draft_model.candidate_selector full_vocab = selector.predecessor_codebook.shape[0] + takes = self._draft_head_takes(draft_model, gen_logits, spec_metadata, full_vocab) if gen_logits.shape[-1] == full_vocab: return super()._apply_dflash2_selector( draft_model, - self._draft_head_transformed(gen_logits), + self._draft_head_transformed(draft_model, gen_logits), block_hidden_states, anchor_tokens, spec_metadata, ) - if self._shard_selector_applies(gen_logits, spec_metadata, full_vocab): + if takes: tokens = selector_draft_tokens( gen_logits, block_hidden_states, @@ -449,7 +550,7 @@ def _apply_dflash2_selector( selector, self.mapping, self.__dict__.get("_draft_head_allreduce"), - transform=self._draft_head_transform_for(gen_logits), + transform=self._draft_head_transform_for(draft_model, gen_logits), ) stand_in = tokens.unsqueeze(-1).expand(*tokens.shape, full_vocab) setattr(stand_in, TOKENS_ATTR, tokens) @@ -457,7 +558,7 @@ def _apply_dflash2_selector( # The built-in path: the transform, the head's own all-gather and .float(), # then the built-in selector. full_logits = allgather( - self._draft_head_transformed(gen_logits), self.mapping, dim=-1 + self._draft_head_transformed(draft_model, gen_logits), self.mapping, dim=-1 ).float() return super()._apply_dflash2_selector( draft_model, full_logits, block_hidden_states, anchor_tokens, spec_metadata @@ -745,6 +846,17 @@ def agreed_reason(mapping, reason: str | None, transform: Optional[DraftLogitTra return _tp_agreement.agreed_refusal(mapping, reason) +def agreed_refusal(mapping, reason: str | None) -> str | None: + """``reason`` as every rank of ``mapping``'s TP group agrees (``dispatch.agreed_refusal``). + + Collective over the TP group: every rank calls it at the same point of the + install, whatever its own reason. + """ + from ... import dispatch as _tp_agreement + + return _tp_agreement.agreed_refusal(mapping, reason) + + def installed(llm: nn.Module) -> bool: """Whether ``install`` has put the draft head on ``llm``'s DFlash worker.""" return isinstance(getattr(llm, "spec_worker", None), ShardedDraftHeadDFlashWorker) @@ -756,9 +868,10 @@ def install(llm: nn.Module, make_allreduce=key_allreduce) -> bool: ``make_allreduce(mapping)`` builds the key exchange (collective: every rank installs at the same point of the load). A second call on an installed ``llm`` changes nothing. Everything is built and checked before anything - changes, and the refusal is agreed over the TP group (``agreed_reason``), - so a refusing install leaves the worker and the drafter as they were on - every rank. + changes, and each refusal is agreed over the TP group (``agreed_reason`` + for the checks, ``agreed_refusal`` for the key exchange ``make_allreduce`` + built), so a refusing install leaves the worker and the drafter as they + were on every rank. Raises: ValueError: when there is no DFlash 2 drafter (no spec worker, or a @@ -766,8 +879,10 @@ def install(llm: nn.Module, make_allreduce=key_allreduce) -> bool: the feature supports, or its logit transform is not one the draft head can apply after the top-k (``logit_transform.unsupported_reason``), or this torch's ``torch.topk`` orders ties differently from the merge - (``tie_order_reason``) -- running the built-in path, or a path whose - drafts differ, under the fast path's name would misreport what runs. + (``tie_order_reason``), or the key exchange's MNNVL all-reduce could + not be built on some rank -- running the built-in path, or a path + whose drafts differ, under the fast path's name would misreport what + runs. """ if installed(llm): return True @@ -797,10 +912,11 @@ def install(llm: nn.Module, make_allreduce=key_allreduce) -> bool: allreduce = make_allreduce(worker.mapping) if allreduce is None: # The composition's precondition checked the group takes an MNNVL all-reduce. - raise ValueError( - "the vocab-sharded draft head: the MNNVL FP32 all-reduce of the key exchange " - "could not be built on this TP group" - ) + reason = "the MNNVL FP32 all-reduce of the key exchange could not be built" + # Building the all-reduce can fail on some ranks only; every rank refuses alike. + reason = agreed_refusal(worker.mapping, reason) + if reason is not None: + raise ValueError(f"[musespark_open] draft_head cannot run here: {reason}") worker.__class__ = ShardedDraftHeadDFlashWorker if type(draft_model.logits_processor) is LogitsProcessor: draft_model.logits_processor = ShardDraftLogitsProcessor(draft_model.logits_processor) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/verify_fold/fold.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/verify_fold/fold.py index 460cbb116506..1634b8830a70 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/verify_fold/fold.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/verify_fold/fold.py @@ -739,7 +739,7 @@ def _apply_dflash2_selector( if ( self._fold_path_walk_tail and gen_logits.shape[-1] != full_vocab - and self._shard_selector_applies(gen_logits, spec_metadata, full_vocab) + and self._draft_head_takes(draft_model, gen_logits, spec_metadata, full_vocab) ): args = self._path_walk_tail_args(step, gen_logits, spec_metadata) if args is None: @@ -762,7 +762,7 @@ def walk(edges, unary, ids): self.mapping, self.__dict__.get("_draft_head_allreduce"), walk=walk, - transform=self._draft_head_transform_for(gen_logits), + transform=self._draft_head_transform_for(draft_model, gen_logits), ) step.tail = _Tail(tokens, next_tokens, kv_lens is not None, batch_indices) if kv_lens is not None: diff --git a/tests/unittest/_torch/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/draft_head/test_modeling_v2_draft_head_transform.py b/tests/unittest/_torch/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/draft_head/test_modeling_v2_draft_head_transform.py index 344042a7091b..86a2af50f802 100644 --- a/tests/unittest/_torch/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/draft_head/test_modeling_v2_draft_head_transform.py +++ b/tests/unittest/_torch/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/draft_head/test_modeling_v2_draft_head_transform.py @@ -11,7 +11,10 @@ that keeps the built-in computation (sampled drafts, d2t, another layout, another tensor, a stale record) gets the built-in transform bit for bit, the fold's tail walk takes the same inputs, CUDA-graph replay follows new logits, -and the install is atomic, idempotent and checked. +and the install is atomic, idempotent and checked. A drafter whose transform +attributes change after the install drafts the built-in tokens of its new +transform through the built-in path, at every point of the step, and each +reason a step takes the built-in path is warned once. TP ranks are emulated in one process: each rank's drafter, worker and head shard are built here, and the key exchange returns what every rank's kernels @@ -126,8 +129,9 @@ def spec_metadata(requests=1, advanced=False): @pytest.fixture def single_process(monkeypatch): - """The install's TP agreement in one process: each emulated rank agrees with itself.""" + """The install's TP agreements in one process: each emulated rank agrees with itself.""" monkeypatch.setattr(draft_head, "agreed_reason", lambda mapping, reason, transform: reason) + monkeypatch.setattr(draft_head, "agreed_refusal", lambda mapping, reason: reason) def install_rank(rank, local, full, ranks=TP, config=MUSESPARK, layout="halves"): @@ -732,3 +736,364 @@ def test_ranks_whose_tables_differ_refuse(monkeypatch): assert reason is not None and "differ in their drafter logit transform" in reason monkeypatch.setattr(dispatch, "agreed_union", lambda mapping, local: list(local)) assert draft_head.agreed_reason(object(), None, t) is None + + +def test_an_exchange_one_rank_cannot_build_is_refused_on_every_rank(monkeypatch, single_process): + """``make_allreduce``'s refusal is agreed over the TP group: a rank that built it refuses too.""" + from tensorrt_llm._torch.speculative.dflash import DFlashWorker + + agreed = [] + elsewhere = ( + "the MNNVL FP32 all-reduce of the key exchange could not be built (on TP ranks [3] of 8)" + ) + + def agreed_refusal(mapping, reason): + agreed.append(reason) + return elsewhere + + monkeypatch.setattr(draft_head, "agreed_refusal", agreed_refusal) + shards, _, _ = draft_step(1, K, "gemm3", 9) + for built in (object(), None): + worker = make_worker(0) + drafter = make_drafter(MUSESPARK, StubHead(shards[0], TP)) + inner = drafter.logits_processor + llm = types.SimpleNamespace(spec_worker=worker, draft_model=drafter) + with pytest.raises(ValueError, match="could not be built"): + draft_head.install(llm, make_allreduce=lambda mapping, built=built: built) + assert type(worker) is DFlashWorker and drafter.logits_processor is inner + assert "_draft_head_allreduce" not in worker.__dict__ + assert draft_head.TRANSFORM_ATTR not in worker.__dict__ + assert draft_head.TRANSFORM_METHOD not in vars(drafter) + # This rank's own reason went into the agreement: none, then its own refusal. + assert agreed[0] is None and "could not be built" in agreed[1] + + +def test_agreed_refusal_is_the_dispatch_agreement(monkeypatch): + from tensorrt_llm._torch._experimental.modeling_v2.models.musespark_open.muse_spark_open_816b_a42b_nvfp4__sm_103__tep import ( # noqa: E501 + dispatch, + ) + + seen = [] + + def agreed_refusal(mapping, local): + seen.append((mapping, local)) + return "agreed" + + monkeypatch.setattr(dispatch, "agreed_refusal", agreed_refusal) + mapping = object() + assert draft_head.agreed_refusal(mapping, "mine") == "agreed" + assert seen == [(mapping, "mine")] + + +# --------------------------------------------------------------------------- the tables' drafter + + +#: Transforms a drafter may declare after the install, the identity among them. +CHANGED = { + "identity": (1.0, None), + "softcap 30": (MUSESPARK[0], 30.0), + "multiplier only": (0.5, None), + "softcap only": (1.0, MUSESPARK[1]), +} + + +def staged_step(worker, drafter, hidden, anchor, md, change=None, when=None): + """``drafter_step`` up to the selector, with ``change()`` run at ``when``. + + ``when`` is ``before`` (the logits processor), ``after_defer`` (the + drafter's transform) or ``after_refine`` (``_refine_block_logits``). + Returns what the selector returns. + """ + num_gens = hidden.shape[0] // K + if when == "before": + change() + gen_logits = drafter.logits_processor(hidden, drafter.lm_head, None, True) + apply_logit_transform = getattr(drafter, "apply_draft_logit_transform", None) + if apply_logit_transform is not None: + gen_logits = apply_logit_transform(gen_logits) + if when == "after_defer": + change() + gen_logits = gen_logits.reshape(num_gens, K, gen_logits.shape[-1]) + gen_logits = worker._refine_block_logits(drafter, gen_logits, {}, md) + if when == "after_refine": + change() + return worker._apply_dflash2_selector( + drafter, gen_logits, hidden.reshape(num_gens, K, -1), anchor, md + ) + + +def patch_builtin_gather(monkeypatch, shards, rank, config, requests): + """The head's vocabulary all-gather, checking this rank hands it the built-in transform. + + ``config``'s built-in transform of this rank's shard must reach the + all-gather bit for bit (bf16 for the identity, fp32 otherwise); it returns + every rank's built-in transform, concatenated. A key exchange fails. + """ + local = builtin_transform(config, shards[rank]).view(requests, K, -1) + gathered = torch.cat([builtin_transform(config, s) for s in shards], dim=-1) + gathered = gathered.view(requests, K, VOCAB) + calls = [] + + def allgather(tensor, mapping, dim=-1, sizes=None): + assert dim == -1, "a step whose tables do not hold must not exchange shard keys" + assert tensor.dtype == local.dtype, (tensor.dtype, local.dtype) + assert_same_floats(tensor.float(), local.float()) + calls.append(dim) + return gathered.clone() + + monkeypatch.setattr(draft_head, "allgather", allgather) + return calls + + +def _spy_sharded_selector(monkeypatch): + calls = [] + real = draft_head.selector_draft_tokens + + def spy(*args, **kwargs): + calls.append(kwargs.get("transform")) + return real(*args, **kwargs) + + monkeypatch.setattr(draft_head, "selector_draft_tokens", spy) + return calls + + +def _tokens_of(worker, out): + """The step's draft tokens off the built-in selector's full-vocab block logits.""" + assert getattr(out, draft_head.TOKENS_ATTR, None) is None, "the sharded selector drafted" + assert not worker.__dict__[draft_head.OWED_ATTR] + requests, positions, vocab = out.shape + return torch.argmax(out.reshape(-1, vocab), dim=-1).to(torch.int32).view(requests, positions) + + +@pytest.mark.parametrize("when", ["before", "after_defer", "after_refine"]) +@pytest.mark.parametrize("name", list(CHANGED)) +@pytest.mark.parametrize("requests", [1, 3]) +def test_a_changed_transform_drafts_the_builtin_tokens( + requests, name, when, monkeypatch, single_process +): + """The drafter's multiplier / softcap changed after the install: no stale tables, built-in drafts.""" + config = CHANGED[name] + rank = 5 + shards, block_hidden, anchor = draft_step( + requests, K, "saturated", seed_of(f"changed{requests}{name}{when}") + ) + full = torch.cat(shards, dim=-1) + llm = install_rank(rank, shards[rank], full) + worker, drafter = llm.spec_worker, llm.draft_model + t = worker.__dict__[draft_head.TRANSFORM_ATTR] + expected = builtin_draft_tokens(shards, block_hidden, anchor, config) + gathers = patch_builtin_gather(monkeypatch, shards, rank, config, requests) + sharded = _spy_sharded_selector(monkeypatch) + md = spec_metadata(requests) + hidden = block_hidden.view(requests * K, HIDDEN) + + def change(): + drafter.logit_multiplier, drafter.logit_softcap = config + + out = staged_step(worker, drafter, hidden, anchor, md, change, when) + assert not sharded and gathers == [-1] + assert torch.equal(_tokens_of(worker, out), expected) + # The tables stay installed and serve again once the drafter declares their transform. + assert worker.__dict__[draft_head.TRANSFORM_ATTR] is t + drafter.logit_multiplier, drafter.logit_softcap = MUSESPARK + patch_allgather(monkeypatch, full, TP, t) + tokens = worker.sample_draft_tokens( + staged_step(worker, drafter, hidden, anchor, md), md, requests, num_contexts=0 + ) + assert sharded == [t] + assert torch.equal(tokens, builtin_draft_tokens(shards, block_hidden, anchor, MUSESPARK)) + + +@pytest.mark.parametrize("when", ["before", "after_refine"]) +@pytest.mark.parametrize("name", ["identity", "softcap 30"]) +def test_a_changed_transform_under_the_fold_tail(name, when, monkeypatch, single_process): + """The fold's tail walk turns the step down too: the built-in drafts of the new transform.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.musespark_open.muse_spark_open_816b_a42b_nvfp4__sm_103__tep.kernels.verify_fold import ( # noqa: E501 + fold, + ) + + config = CHANGED[name] + requests, rank = 2, 6 + shards, block_hidden, anchor = draft_step( + requests, K, "straddle_up", seed_of(f"fold{name}{when}") + ) + full = torch.cat(shards, dim=-1) + llm = install_rank(rank, shards[rank], full) + worker, drafter = llm.spec_worker, llm.draft_model + assert fold.install(worker, verify_fold=False, path_walk_tail=True) + expected = builtin_draft_tokens(shards, block_hidden, anchor, config) + gathers = patch_builtin_gather(monkeypatch, shards, rank, config, requests) + sharded = _spy_sharded_selector(monkeypatch) + md = spec_metadata(requests) + hidden = block_hidden.view(requests * K, HIDDEN) + g = torch.Generator(device="cuda").manual_seed(4) + accepted = torch.randint(0, VOCAB, (requests, K + 1), generator=g, device="cuda").int() + num_accepted = torch.randint(1, K + 2, (requests,), generator=g, device="cuda").int() + kv_lens = torch.full((requests,), 50, dtype=torch.int32, device="cuda") + attn = types.SimpleNamespace(num_contexts=0, num_seqs=requests, kv_lens_cuda=kv_lens) + worker._kv_rewind_amount = (1 - num_accepted).int() + worker._kv_rewind_pending, worker._kv_rewind_nc, worker._kv_rewind_bs = True, 0, requests + step = fold._Step(attn_metadata=attn, spec_metadata=md, draft_model=drafter, warmup=False) + step.accepted, step.num_accepted = accepted, num_accepted + worker.__dict__["_fold_step"] = step + + def change(): + drafter.logit_multiplier, drafter.logit_softcap = config + + try: + out = staged_step(worker, drafter, hidden, anchor, md, change, when) + finally: + worker.__dict__["_fold_step"] = None + assert not sharded and gathers == [-1] and step.tail is None + assert torch.equal(_tokens_of(worker, out), expected) + + +def test_the_selector_refuses_tables_that_do_not_hold(monkeypatch, single_process): + shards, _, _, _, worker, drafter, _ = _one_rank(monkeypatch) + logits = shards[2].view(1, K, -1) + assert worker._draft_head_transform_for(drafter, logits) is not None + drafter.logit_softcap = 30.0 + assert not worker._draft_head_takes(drafter, logits, spec_metadata(), VOCAB) + with pytest.raises(draft_head.DraftTransformStateError): + worker._draft_head_transform_for(drafter, logits) + + +# --------------------------------------------------------------------------- warnings + + +@pytest.fixture +def logged_warnings(monkeypatch): + """The draft head's warnings logged from here on, with a fresh once-only record.""" + logged = [] + + def log(severity, *msg): + message = " ".join(map(str, msg)) + if severity == draft_head.logger.WARNING and "] draft_head:" in message: + logged.append(message) + + monkeypatch.setattr(draft_head.logger, "_appeared_keys", set()) + monkeypatch.setattr(draft_head.logger, "log", log) + return logged + + +def _sampled(worker, drafter, md): + md.wants_advanced_draft_sampling = True + + +def _d2t(worker, drafter, md): + worker._d2t = torch.zeros(VOCAB, dtype=torch.int64, device="cuda") + + +def _softcap_30(worker, drafter, md): + drafter.logit_softcap = 30.0 + + +def _identity(worker, drafter, md): + drafter.logit_multiplier, drafter.logit_softcap = 1.0, None + + +def _wrapped_processor(worker, drafter, md): + class Wrapper(torch.nn.Module): + def __init__(self, inner): + super().__init__() + self.inner = inner + + def forward(self, *args): + return self.inner(*args) + + drafter.logits_processor = Wrapper(drafter.logits_processor) + + +def _fp32_head(worker, drafter, md): + drafter.lm_head.source = drafter.lm_head.source.float() + + +def _class_above(worker, drafter, md): + class Above(type(worker)): + def _refine_block_logits(self, draft_model, gen_logits, inputs, spec_metadata): + return super()._refine_block_logits(draft_model, gen_logits, inputs, spec_metadata) + + worker.__class__ = Above + + +#: Each way a step with the transform installed takes the built-in path: (set-up, message part). +FALLBACKS = { + "sampled": (_sampled, "sampled drafts"), + "d2t": (_d2t, "a d2t draft vocab map"), + "attributes": (_softcap_30, draft_head.STALE_TABLES), + "identity": (_identity, draft_head.STALE_TABLES), + "processor": (_wrapped_processor, "the drafter's logits processor is Wrapper"), + "fp32 logits": (_fp32_head, "returned torch.float32 logits"), + "class above": (_class_above, "the worker class Above overrides"), + "not deferred": (None, draft_head.NOT_DEFERRED), +} + + +@pytest.mark.parametrize("why", list(FALLBACKS)) +def test_a_fallback_warns_once(why, logged_warnings, monkeypatch, single_process): + setup, part = FALLBACKS[why] + shards, block_hidden, anchor, full, worker, drafter, hidden = _one_rank(monkeypatch) + md = spec_metadata() + if setup is not None: + setup(worker, drafter, md) + gathered = builtin_full_logits(shards, 1, K, MUSESPARK) + patch_allgather(monkeypatch, full, TP, None, gathered=gathered) + for _ in range(3): + if why == "sampled": + # Sampled drafts leave the step after the refinement, for the built-in sampler. + out = drafter.apply_draft_logit_transform( + drafter.logits_processor(hidden, drafter.lm_head, None, True) + ) + worker._refine_block_logits(drafter, out.reshape(1, K, -1), {}, md) + elif why == "not deferred": + # bf16 logits the deferring transform never saw. + out = drafter.logits_processor(hidden, drafter.lm_head, None, True).reshape(1, K, -1) + refined = worker._refine_block_logits(drafter, out, {}, md) + assert_same_floats(refined, builtin_transform(MUSESPARK, out)) + else: + out = staged_step(worker, drafter, hidden, anchor, md) + assert getattr(out, draft_head.TOKENS_ATTR, None) is None + assert len(logged_warnings) == 1, logged_warnings + assert logged_warnings[0].startswith( + "[musespark_open] draft_head: a draft step takes the built-in path" + ) + assert part in logged_warnings[0], logged_warnings[0] + + +def test_each_reason_warns_once_and_the_fast_path_never( + logged_warnings, monkeypatch, single_process +): + shards, block_hidden, anchor, full, worker, drafter, hidden = _one_rank(monkeypatch) + t = worker.__dict__[draft_head.TRANSFORM_ATTR] + patch_allgather(monkeypatch, full, TP, t) + md = spec_metadata() + for _ in range(2): + out = staged_step(worker, drafter, hidden, anchor, md) + assert getattr(out, draft_head.TOKENS_ATTR, None) is not None + assert logged_warnings == [] + softcap_30 = builtin_full_logits(shards, 1, K, (MUSESPARK[0], 30.0)) + patch_allgather(monkeypatch, full, TP, None, gathered=softcap_30) + drafter.logit_softcap = 30.0 + for _ in range(2): + staged_step(worker, drafter, hidden, anchor, md) + worker._d2t = torch.zeros(VOCAB, dtype=torch.int64, device="cuda") + drafter.logit_softcap = MUSESPARK[1] + musespark = builtin_full_logits(shards, 1, K, MUSESPARK) + patch_allgather(monkeypatch, full, TP, None, gathered=musespark) + for _ in range(2): + staged_step(worker, drafter, hidden, anchor, md) + assert len(logged_warnings) == 2, logged_warnings + assert draft_head.STALE_TABLES in logged_warnings[0] and "d2t" in logged_warnings[1] + + +def test_without_a_transform_nothing_warns(logged_warnings, monkeypatch, single_process): + shards, block_hidden, anchor = draft_step(1, K, "gemm3", 11) + full = torch.cat(shards, dim=-1) + llm = install_rank(1, shards[1], full, config=(1.0, None)) + worker, drafter = llm.spec_worker, llm.draft_model + worker._d2t = torch.zeros(VOCAB, dtype=torch.int64, device="cuda") + gathered = builtin_full_logits(shards, 1, K) + patch_allgather(monkeypatch, full, TP, None, gathered=gathered) + out = staged_step(worker, drafter, block_hidden.view(K, HIDDEN), anchor, spec_metadata()) + assert getattr(out, draft_head.TOKENS_ATTR, None) is None + assert logged_warnings == [] From 4525d88de6b97683528281c807b94afa59532591 Mon Sep 17 00:00:00 2001 From: Xiao Wang <24860335+xwang233@users.noreply.github.com> Date: Wed, 28 Oct 2026 16:39:35 +0000 Subject: [PATCH 16/16] [None][fix] KVCacheManagerV2: defer only terminated requests' closes; flush before control actions With TRTLLM_KV_DEFER_CLOSE=1, free_resources deferred the close of every non-dummy request while the overlap loop was armed, including preemption and admission rollbacks, whose callers count on the pages being free when the call returns. The executor also fired a drained control request before flushing, so the action could run while a finished request's cache was still open. - KVCacheManagerV2.deferring_closes(): free_resources defers a close only while armed and inside this scope; every other release closes at once. - PyExecutor._do_terminate_request releases inside the scope; preemption, recompute pauses and scheduler rollbacks keep closing at once. - PyExecutor._handle_control_request flushes the deferred closes before it fires the action. Off by default: with TRTLLM_KV_DEFER_CLOSE unset nothing changes. Tests: test_kv_cache_v2_deferred_close.py. Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com> --- .../kv_cache/kv_cache_manager_v2.py | 40 +++- tensorrt_llm/_torch/pyexecutor/py_executor.py | 10 +- .../test_kv_cache_v2_deferred_close.py | 177 ++++++++++++++++-- 3 files changed, 199 insertions(+), 28 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py index 59667f435e08..57560fe32204 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py @@ -12,6 +12,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +import contextlib import hashlib import math import os @@ -24,6 +25,7 @@ Callable, Dict, Iterable, + Iterator, List, NamedTuple, Optional, @@ -1235,6 +1237,7 @@ class KVCacheManagerV2(BaseResourceManager): # instances built without running __init__. _defer_close_enabled: bool = False _defer_close_armed: bool = False + _defer_close_scope: bool = False _deferred_closes: Tuple[Tuple[int, _KVCache], ...] = () def __init__( @@ -5662,15 +5665,20 @@ def _detach_index_slot(self, kv_cache: _KVCache) -> None: def arm_deferred_close(self, armed: bool) -> bool: """Arm or disarm deferred closing of finished requests' caches. - While armed, free_resources releases a finished (non-dummy) request's - IndexMapper slot at once (as release_index_slot does) and queues the - cache; its close(), which commits the last block and releases the - pages, runs in flush_deferred_closes(). The overlap executor arms this - only for its loop and flushes right after launching each forward step, - where the GPU still has queued work to cover the close, and before it - waits on an empty request queue with no request left. Takes effect - only with TRTLLM_KV_DEFER_CLOSE=1 on a target-model manager. Disarming - flushes. Returns whether deferred closing is armed. + While armed, free_resources called inside deferring_closes() releases + the (non-dummy) request's IndexMapper slot at once (as + release_index_slot does) and queues the cache; its close(), which + commits the last block and releases the pages, runs in + flush_deferred_closes(). Every other release (preemption, admission + rollbacks, reservation drops) closes at once, since its caller counts + on the pages being free when it returns. The overlap executor arms this + only for its loop, releases terminated requests inside + deferring_closes(), and flushes right after launching each forward + step, where the GPU still has queued work to cover the close, before it + waits on an empty request queue with no request left, and before it + fires a control request. Takes effect only with TRTLLM_KV_DEFER_CLOSE=1 + on a target-model manager. Disarming flushes. Returns whether deferred + closing is armed. """ was_armed = self._defer_close_armed self._defer_close_armed = bool(armed) and self._defer_close_enabled @@ -5683,6 +5691,18 @@ def arm_deferred_close(self, armed: bool) -> bool: self.flush_deferred_closes() return self._defer_close_armed + @contextlib.contextmanager + def deferring_closes(self) -> Iterator[None]: + """Let free_resources defer the close of the requests it releases for + the duration, while deferred closing is armed (arm_deferred_close). + For the release of requests that have terminated only.""" + outer = self._defer_close_scope + self._defer_close_scope = True + try: + yield + finally: + self._defer_close_scope = outer + def flush_deferred_closes(self) -> int: """Close the caches whose close() free_resources deferred, in the order their requests finished. Returns how many were closed.""" @@ -5730,7 +5750,7 @@ def free_resources(self, request: LlmRequest, pin_on_release: bool = False): kv_cache.discard_pending_stats() # Dummy requests reuse their ids, which the stats-exclusion bookkeeping # done after close() is keyed on, so they always close at once. - if self._defer_close_armed and not request.is_dummy: + if self._defer_close_armed and self._defer_close_scope and not request.is_dummy: if request.py_request_id in self._early_freed_index_requests: self._early_freed_index_requests.discard(request.py_request_id) else: diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index fbf93ff7e01d..3d18f15bbaec 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -4906,6 +4906,10 @@ def _handle_control_request(self): f"drain={pending.control_requires_drain} " f"active_requests={len(self.active_requests)} " f"waiting_queue={len(self.waiting_queue)}") + if self._kv_defer_close_armed: + # The action must not run while a finished request's cache is + # still open: its close() commits the last block and frees pages. + self.kv_cache_manager.flush_deferred_closes() # Quiesce the device before the action mutates GPU state. Under the # overlap scheduler a previous batch's forward/sample kernels may still # be in flight, so an in-place update_weights reload (or sleep/wakeup @@ -8459,7 +8463,11 @@ def _free_request_resources(self, request: LlmRequest) -> None: self.disagg.forget_request(request.py_request_id) def _do_terminate_request(self, request: LlmRequest) -> None: - self._free_request_resources(request) + # Only a terminated request's cache may close late; preemption and + # recompute pauses release theirs through other paths, at once. + with (self.kv_cache_manager.deferring_closes() + if self._kv_defer_close_armed else nullcontext()): + self._free_request_resources(request) # Cancellation and request-scoped failures can terminate before the # normal post-prefill release point, including with a partial buffer. _strip_py_multimodal_data_post_prefill(request) diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_deferred_close.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_deferred_close.py index 9ffa3a4bf991..5320e156e7fe 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_deferred_close.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_deferred_close.py @@ -14,10 +14,12 @@ # limitations under the License. """KVCacheManagerV2's opt-in deferred close (TRTLLM_KV_DEFER_CLOSE=1). -While armed, free_resources releases a finished request's IndexMapper slot at -once and leaves the cache's close() to flush_deferred_closes(); the overlap -executor arms it for its loop and flushes after launching each forward step -and before waiting for requests with none left. +While armed, free_resources inside deferring_closes() releases a terminated +request's IndexMapper slot at once and leaves the cache's close() to +flush_deferred_closes(); every other release closes at once. The overlap +executor arms it for its loop, releases terminated requests inside the scope, +and flushes after launching each forward step, before waiting for requests with +none left, and before firing a control request. """ import types @@ -71,12 +73,18 @@ def _request(rid, *, dummy=False): return types.SimpleNamespace(py_request_id=rid, is_dummy=dummy) +def _finish(mgr, request): + """Release *request* as the executor releases a terminated one.""" + with mgr.deferring_closes(): + mgr.free_resources(request) + + def test_off_unless_enabled_closes_at_once(): events = [] mgr = _manager(events, enabled=False) _cache(mgr, events, 7) assert mgr.arm_deferred_close(True) is False - mgr.free_resources(_request(7)) + _finish(mgr, _request(7)) assert events == [("close", 7), ("clear_stats", 7), ("remove_seq", 7)] assert mgr._deferred_closes == () @@ -87,7 +95,7 @@ def test_class_defaults_keep_managers_built_without_init_unchanged(): del mgr._defer_close_enabled _cache(mgr, events, 3) assert mgr.arm_deferred_close(True) is False - mgr.free_resources(_request(3)) + _finish(mgr, _request(3)) assert events == [("close", 3), ("clear_stats", 3), ("remove_seq", 3)] @@ -96,7 +104,7 @@ def test_armed_releases_slot_now_and_closes_on_flush(): mgr = _manager(events, num_pools=2) kv = _cache(mgr, events, 7) assert mgr.arm_deferred_close(True) is True - mgr.free_resources(_request(7)) + _finish(mgr, _request(7)) # Slot released at once, after the cache is unbound from it; no close yet. assert events == [("unbind", 7, 0, None), ("unbind", 7, 1, None), ("remove_seq", 7)] kv.close.assert_not_called() @@ -114,7 +122,7 @@ def test_flush_closes_in_finish_order(): _cache(mgr, events, rid) mgr.arm_deferred_close(True) for rid in (5, 2, 9): - mgr.free_resources(_request(rid)) + _finish(mgr, _request(rid)) events.clear() assert mgr.flush_deferred_closes() == 3 assert [e for e in events if e[0] == "close"] == [("close", 5), ("close", 2), ("close", 9)] @@ -132,7 +140,7 @@ def test_each_stats_exclusion_is_cleared_after_its_close(): _cache(mgr, events, rid) mgr.arm_deferred_close(True) for rid in (1, 4): - mgr.free_resources(_request(rid)) + _finish(mgr, _request(rid)) events.clear() mgr.flush_deferred_closes() assert events == [("close", 1), ("clear_stats", 1), ("close", 4), ("clear_stats", 4)] @@ -143,7 +151,7 @@ def test_dummy_requests_close_at_once_while_armed(): mgr = _manager(events, num_pools=1) _cache(mgr, events, 11) mgr.arm_deferred_close(True) - mgr.free_resources(_request(11, dummy=True)) + _finish(mgr, _request(11, dummy=True)) assert events == [("close", 11), ("clear_stats", 11), ("remove_seq", 11)] assert mgr._deferred_closes == () @@ -154,7 +162,7 @@ def test_early_released_slot_is_not_released_twice(): _cache(mgr, events, 6) mgr._early_freed_index_requests.add(6) mgr.arm_deferred_close(True) - mgr.free_resources(_request(6)) + _finish(mgr, _request(6)) assert events == [] assert 6 not in mgr._early_freed_index_requests mgr.flush_deferred_closes() @@ -166,7 +174,7 @@ def test_sparse_batch_is_left_before_the_slot_is_released(): mgr = _manager(events, num_pools=1, sparse=True) _cache(mgr, events, 8) mgr.arm_deferred_close(True) - mgr.free_resources(_request(8)) + _finish(mgr, _request(8)) assert events == [("sparse_remove", 8), ("unbind", 8, 0, None), ("remove_seq", 8)] mgr.sparse_metadata_batch.record_read.assert_called_once_with(0) @@ -177,12 +185,12 @@ def test_disarm_flushes_and_stops_deferring(): _cache(mgr, events, 1) _cache(mgr, events, 2) mgr.arm_deferred_close(True) - mgr.free_resources(_request(1)) + _finish(mgr, _request(1)) events.clear() assert mgr.arm_deferred_close(False) is False assert events == [("close", 1), ("clear_stats", 1)] events.clear() - mgr.free_resources(_request(2)) + _finish(mgr, _request(2)) assert events == [("close", 2), ("clear_stats", 2), ("remove_seq", 2)] @@ -191,7 +199,7 @@ def test_shutdown_closes_deferred_caches_before_the_sparse_batch(): mgr = _manager(events, num_pools=1, sparse=True) _cache(mgr, events, 3) mgr.arm_deferred_close(True) - mgr.free_resources(_request(3)) + _finish(mgr, _request(3)) events.clear() mgr.event_manager = None mgr.shutdown() @@ -288,7 +296,7 @@ def test_last_request_closes_before_the_idle_queue_wait(): mgr = _manager(events, num_pools=1) _cache(mgr, events, 4) mgr.arm_deferred_close(True) - mgr.free_resources(_request(4)) + _finish(mgr, _request(4)) events.clear() executor = _fetch_stub(mgr, events) @@ -304,7 +312,7 @@ def test_busy_fetch_leaves_deferred_closes_to_the_loop(): mgr = _manager(events, num_pools=1) _cache(mgr, events, 4) mgr.arm_deferred_close(True) - mgr.free_resources(_request(4)) + _finish(mgr, _request(4)) events.clear() executor = _fetch_stub(mgr, events) @@ -344,3 +352,138 @@ def test_rebalance_flushes_deferred_closes_before_adjusting(monkeypatch): PyExecutor._maybe_rebalance_kv_pools(executor) assert order == ["consume", "flush", "adjust"] + + +def test_armed_release_outside_the_terminate_scope_closes_at_once(): + # Admission rollbacks and reservation drops call free_resources directly + # and count on the pages being free when it returns. + events = [] + mgr = _manager(events, num_pools=1) + _cache(mgr, events, 7) + mgr.arm_deferred_close(True) + mgr.free_resources(_request(7)) + assert events == [("close", 7), ("clear_stats", 7), ("remove_seq", 7)] + assert mgr._deferred_closes == () + + +def test_preemption_releases_pages_at_once_while_armed(): + # preempt_request reports the pages released; the scheduler retries the + # allocation right after each victim, so the close cannot wait for a flush. + events = [] + mgr = _manager(events, num_pools=1) + _cache(mgr, events, 9) + mgr.arm_deferred_close(True) + assert mgr.preempt_request(_request(9)) is True + assert events == [("close", 9), ("clear_stats", 9), ("remove_seq", 9)] + assert mgr._deferred_closes == () + + +def test_preemption_inside_the_terminate_scope_still_closes_at_once(): + events = [] + mgr = _manager(events, num_pools=1) + _cache(mgr, events, 2) + _cache(mgr, events, 3) + mgr.arm_deferred_close(True) + with mgr.deferring_closes(): + mgr.free_resources(_request(2)) + with mgr.deferring_closes(): + pass + assert mgr._defer_close_scope is True + assert mgr._defer_close_scope is False + mgr.preempt_request(_request(3)) + assert ("close", 3) in events and ("close", 2) not in events + assert [rid for rid, _ in mgr._deferred_closes] == [2] + + +def _terminate_stub(kv_cache_manager, *, armed=True): + executor = object.__new__(PyExecutor) + executor.kv_cache_manager = kv_cache_manager + executor._kv_defer_close_armed = armed + executor.resource_manager = types.SimpleNamespace( + free_resources=kv_cache_manager.free_resources + ) + executor._prefetched_request_ids = set() + executor._disagg_coordinator = Mock() + executor.gather_all_responses = False + executor.dist = types.SimpleNamespace(rank=1) + executor.result_wait_queues = {} + return executor + + +def test_terminated_request_defers_and_recompute_pause_does_not(): + events = [] + mgr = _manager(events, num_pools=1) + _cache(mgr, events, 4) + _cache(mgr, events, 5) + mgr.arm_deferred_close(True) + executor = _terminate_stub(mgr) + + PyExecutor._do_terminate_request(executor, _request(4)) + assert ("close", 4) not in events + assert [rid for rid, _ in mgr._deferred_closes] == [4] + + # The recompute-pause teardown releases through _free_request_resources. + PyExecutor._free_request_resources(executor, _request(5)) + assert ("close", 5) in events + assert [rid for rid, _ in mgr._deferred_closes] == [4] + + +def test_terminate_closes_at_once_when_not_armed(): + events = [] + mgr = _manager(events, num_pools=1) + _cache(mgr, events, 6) + executor = _terminate_stub(mgr, armed=False) + + PyExecutor._do_terminate_request(executor, _request(6)) + + assert events == [("close", 6), ("clear_stats", 6), ("remove_seq", 6)] + + +def _control_stub(kv_cache_manager, events, *, active=()): + executor = object.__new__(PyExecutor) + executor.kv_cache_manager = kv_cache_manager + executor._kv_defer_close_armed = True + executor.control_requests = [ + types.SimpleNamespace(control_requires_drain=True, control_id=None) + ] + executor.active_requests = list(active) + executor.waiting_queue = [] + executor._has_pending_connector_transfers = Mock(return_value=False) + executor.control_request_barrier = Mock() + executor.control_request_barrier.set.side_effect = lambda: events.append(("fire",)) + executor.control_action_done = Mock() + return executor + + +def test_drained_control_request_fires_after_the_deferred_closes(monkeypatch): + events = [] + mgr = _manager(events, num_pools=1) + _cache(mgr, events, 8) + mgr.arm_deferred_close(True) + _finish(mgr, _request(8)) + events.clear() + monkeypatch.setattr( + "tensorrt_llm._torch.pyexecutor.py_executor.torch.cuda.synchronize", + Mock(side_effect=lambda: events.append(("sync",))), + ) + executor = _control_stub(mgr, events) + + PyExecutor._handle_control_request(executor) + + assert events == [("close", 8), ("clear_stats", 8), ("sync",), ("fire",)] + assert mgr._deferred_closes == () + + +def test_parked_control_request_leaves_deferred_closes_alone(): + events = [] + mgr = _manager(events, num_pools=1) + _cache(mgr, events, 8) + mgr.arm_deferred_close(True) + _finish(mgr, _request(8)) + events.clear() + executor = _control_stub(mgr, events, active=[object()]) + + PyExecutor._handle_control_request(executor) + + assert events == [] + assert len(mgr._deferred_closes) == 1