Resolve the rest of the rebase conflict

Author: tdene

megatron/core/inference/contexts/dynamic_context.py

@@ -928,16 +928,20 @@ def initialize_all_tensors(self) -> None:
         # transferred to GPU each step via transfer_bookkeeping_to_gpu().
         # Layout matches ContextGPUView._buf so a single cudaMemcpyAsync
         # suffices. Int64 token fields come first (8-byte aligned automatically),
-        # then int32 token fields, then int32 request-staging fields.
-        #   token_to_input_ids                         (int64, max_tokens)
-        #   token_to_pos_ids                           (int64, max_tokens)
-        #   token_to_block_idx                         (int32, max_tokens)
-        #   token_to_local_position_within_kv_block    (int32, max_tokens)
-        #   token_to_request_idx                       (int32, max_tokens)
-        #   token_to_position_in_request               (int32, max_tokens)
-        #   request_in_prefill_status  (staging)       (int32, max_requests)
-        #   request_query_lengths      (staging)       (int32, max_requests)
-        #   request_kv_length_offsets  (staging)       (int32, max_requests)
+        # then int32 token fields, then int32/float32 request-staging fields.
+        #   token_to_input_ids                         (int64,   max_tokens)
+        #   token_to_pos_ids                           (int64,   max_tokens)
+        #   token_to_block_idx                         (int32,   max_tokens)
+        #   token_to_local_position_within_kv_block    (int32,   max_tokens)
+        #   token_to_request_idx                       (int32,   max_tokens)
+        #   token_to_position_in_request               (int32,   max_tokens)
+        #   request_in_prefill_status        (staging) (int32,   max_requests)
+        #   request_query_lengths            (staging) (int32,   max_requests)
+        #   request_kv_length_offsets        (staging) (int32,   max_requests)
+        #   temperature                      (staging) (float32, max_requests)
+        #   top_k                            (staging) (int32,   max_requests)
+        #   top_p                            (staging) (float32, max_requests)
+        #   active_request_last_token_idxs   (alias)   (int32,   max_requests)
         #
         # Token fields are aliased with the source-of-truth attributes
         # (`self.token_to_input_ids`, etc.) because the forward pass reads
@@ -948,7 +952,8 @@ def initialize_all_tensors(self) -> None:
         # slice from the persistent `request_*` tensors above.
         _tok_int64_bytes = self.max_tokens * 8
         _tok_int32_bytes = self.max_tokens * 4
-        _req_int32_bytes = self.max_requests * 4
+        # Request-level fields are all 4 bytes wide (5 int32 + 2 float32 = 7 fields).
+        _req_4byte_bytes = self.max_requests * 4
         # MHA section: 5 fields (int32) shared between GraphedMHAMetadata and
         # NonGraphedMHAMetadata. max_bs == max_requests.
         _mha_query_lengths_bytes = self.max_requests * 4
@@ -983,7 +988,7 @@ def initialize_all_tensors(self) -> None:
         _total_bytes = (
             2 * _tok_int64_bytes
             + 4 * _tok_int32_bytes
-            + 3 * _req_int32_bytes
+            + 7 * _req_4byte_bytes
             + _mha_query_lengths_bytes
             + _mha_cu_query_seq_lengths_bytes
             + _mha_kv_seq_lengths_bytes
@@ -1043,19 +1048,40 @@ def initialize_all_tensors(self) -> None:
         # CPU (refreshed from persistent tensors in transfer_bookkeeping_to_gpu);
         # read-only on GPU via matching slots in ContextGPUView._buf.
         self._staging_request_in_prefill_status = self._cpu_bookkeeping_buf[
-            _off : _off + _req_int32_bytes
+            _off : _off + _req_4byte_bytes
         ].view(torch.int32)
-        _off += _req_int32_bytes
+        _off += _req_4byte_bytes
         self._staging_request_query_lengths = self._cpu_bookkeeping_buf[
-            _off : _off + _req_int32_bytes
+            _off : _off + _req_4byte_bytes
         ].view(torch.int32)
-        _off += _req_int32_bytes
+        _off += _req_4byte_bytes
         self._staging_request_kv_length_offsets = self._cpu_bookkeeping_buf[
-            _off : _off + _req_int32_bytes
+            _off : _off + _req_4byte_bytes
         ].view(torch.int32)
-        _off += _req_int32_bytes
+        _off += _req_4byte_bytes
 
-        self.active_request_last_token_idxs = torch.empty_like(self.request_query_lengths)
+        # Sampling-parameter staging slots, refreshed from `active_request_metadata`
+        # in transfer_bookkeeping_to_gpu(). FlashInfer reads these via
+        # `gpu_view.{temperature, top_k, top_p}`.
+        self._staging_temperature = self._cpu_bookkeeping_buf[_off : _off + _req_4byte_bytes].view(
+            torch.float32
+        )
+        _off += _req_4byte_bytes
+        self._staging_top_k = self._cpu_bookkeeping_buf[_off : _off + _req_4byte_bytes].view(
+            torch.int32
+        )
+        _off += _req_4byte_bytes
+        self._staging_top_p = self._cpu_bookkeeping_buf[_off : _off + _req_4byte_bytes].view(
+            torch.float32
+        )
+        _off += _req_4byte_bytes
+
+        # Per-request last-token row indices. Aliased with the matching gpu_view slot:
+        # build_active_slices/pad_active_slices populate this CPU view.
+        self.active_request_last_token_idxs = self._cpu_bookkeeping_buf[
+            _off : _off + _req_4byte_bytes
+        ].view(torch.int32)
+        _off += _req_4byte_bytes
 
         # Static tensor addresses to make `last_token_logits` graphable with speculative decoding.
         max_logit_idxs = self.max_requests * (self.num_speculative_tokens + 1)
@@ -2244,7 +2270,7 @@ def transfer_bookkeeping_to_gpu(self) -> None:
         All copies use non_blocking=True with pinned CPU memory. CUDA stream
         ordering guarantees the forward pass sees completed transfers.
 
-        The 9 bookkeeping fields are backed by one contiguous pinned CPU buffer
+        The bookkeeping fields are backed by one contiguous pinned CPU buffer
         and one contiguous GPU buffer; a single cudaMemcpyAsync suffices.
         Request-level staging slots are refreshed from the persistent CPU
         tensors immediately before the H2D (GPU reads them at `[:n_active]`
@@ -2255,16 +2281,25 @@ def transfer_bookkeeping_to_gpu(self) -> None:
         padded_active = max(n_active, self.padded_active_request_count)
 
         # Refresh request-level staging slots from the persistent CPU source.
-        # CPU-to-CPU slice assignment on pinned memory (~7.5 KB total for 3
-        # int32 fields at max_requests=624). Negligible vs. the launch overhead
-        # we save by merging 9 H2D memcpys into 1.
+        # CPU-to-CPU slice assignment on pinned memory (~15 KB total for 6
+        # 4-byte fields at max_requests=624). Negligible vs. the launch overhead
+        # we save by merging the H2D memcpys into 1.
         self._staging_request_in_prefill_status[:n_active] = self.request_in_prefill_status_tensor[
             active_slice
         ]
         self._staging_request_query_lengths[:n_active] = self.request_query_lengths[active_slice]
         self._staging_request_kv_length_offsets[:n_active] = self.request_kv_length_offsets[
             active_slice
         ]
+        # Sampling-parameter staging slots: read from `active_request_metadata`,
+        # which `build_active_slices` + `pad_active_slices` already populated for
+        # `[:padded_active]` (active values + neutral padding defaults).
+        self._staging_temperature[:padded_active] = self.active_request_metadata["temperature"][
+            :padded_active
+        ]
+        self._staging_top_k[:padded_active] = self.active_request_metadata["top_k"][:padded_active]
+        self._staging_top_p[:padded_active] = self.active_request_metadata["top_p"][:padded_active]
+
         # Full-iteration CUDA graphs may have captured GPU consumers with the
         # padded graph request count. Keep those padded staging rows bounded so
         # graph replay never builds indices from stale request lengths.

megatron/core/inference/contexts/gpu_view.py

@@ -17,11 +17,11 @@ class ContextGPUView:
         ``context.foo``      -> CPU (source of truth, used by bookkeeping)
         ``context.gpu_view.foo`` -> GPU (snapshot, used by forward pass)
 
-    Layout note: the 9 bookkeeping fields are backed by a single contiguous
+    Layout note: the bookkeeping fields are backed by a single contiguous
     ``uint8`` buffer (``self._buf``). Each field is a ``view(dtype)`` onto a
     slice of that buffer. This matches the pinned-CPU-buffer layout in
     :class:`DynamicInferenceContext` so that the per-step H2D transfer is a
-    single ``cudaMemcpyAsync`` instead of nine small ones.
+    single ``cudaMemcpyAsync`` instead of one per field.
     """
 
     def __init__(
@@ -39,7 +39,10 @@ def __init__(
         #   max_mamba_chunks == 0).
         tok_int64_bytes = max_tokens * 8  # 2 fields of int64 = 8 bytes/elem
         tok_int32_bytes = max_tokens * 4  # 4 fields of int32 = 4 bytes/elem
-        req_int32_bytes = max_requests * 4  # 3 fields of int32
+        # Request-level fields are all 4 bytes wide. 3 int32 (in_prefill_status,
+        # query_lengths, kv_length_offsets) + 1 int32 (top_k) + 2 float32
+        # (temperature, top_p) + 1 int32 (active_request_last_token_idxs) = 7 fields.
+        req_4byte_bytes = max_requests * 4
 
         # MHA section: 5 fields shared by both graphed and non-graphed MHAMetadata
         # (only one is active per step, so sharing storage is fine).
@@ -90,7 +93,7 @@ def __init__(
         total_bytes = (
             2 * tok_int64_bytes
             + 4 * tok_int32_bytes
-            + 3 * req_int32_bytes
+            + 7 * req_4byte_bytes
             + mha_query_lengths_bytes
             + mha_cu_query_seq_lengths_bytes
             + mha_kv_seq_lengths_bytes
@@ -128,12 +131,28 @@ def __init__(
         off += tok_int32_bytes
 
         # Request-level tensors (consumed by sampling, log-probs, speculative verification, MTP).
-        self.request_in_prefill_status = self._buf[off : off + req_int32_bytes].view(torch.int32)
-        off += req_int32_bytes
-        self.request_query_lengths = self._buf[off : off + req_int32_bytes].view(torch.int32)
-        off += req_int32_bytes
-        self.request_kv_length_offsets = self._buf[off : off + req_int32_bytes].view(torch.int32)
-        off += req_int32_bytes
+        self.request_in_prefill_status = self._buf[off : off + req_4byte_bytes].view(torch.int32)
+        off += req_4byte_bytes
+        self.request_query_lengths = self._buf[off : off + req_4byte_bytes].view(torch.int32)
+        off += req_4byte_bytes
+        self.request_kv_length_offsets = self._buf[off : off + req_4byte_bytes].view(torch.int32)
+        off += req_4byte_bytes
+        # Sampling parameters (consumed by FlashInfer sampling).
+        # Mirror the active slice of `active_request_metadata[{label}]`;
+        # padded slots get neutral defaults from `pad_active_slices` (T=1.0, top_k=0, top_p=0.0).
+        self.temperature = self._buf[off : off + req_4byte_bytes].view(torch.float32)
+        off += req_4byte_bytes
+        self.top_k = self._buf[off : off + req_4byte_bytes].view(torch.int32)
+        off += req_4byte_bytes
+        self.top_p = self._buf[off : off + req_4byte_bytes].view(torch.float32)
+        off += req_4byte_bytes
+        # Per-request last-token row indices (consumed by sampling kernels as `gather_indices`).
+        # The CPU side of this slot IS `context.active_request_last_token_idxs`,
+        # populated by `build_active_slices` and `pad_active_slices`.
+        self.active_request_last_token_idxs = self._buf[off : off + req_4byte_bytes].view(
+            torch.int32
+        )
+        off += req_4byte_bytes
 
         # MHA flash-attention metadata (shared between GraphedMHAMetadata and
         # NonGraphedMHAMetadata — only one is active per step).

megatron/core/inference/sampling/flashinfer_sampling.py

@@ -66,16 +66,18 @@ def sample_kernel(
         # CudaGraphManager consumes these args, if it exists.
         del eager, cache_key
 
-        # Read GPU sampling parameters directly from `context.active_request_metadata`.
-        md = context.active_request_metadata
+        # Read GPU sampling parameters from the per-step gpu_view mirror. The
+        # CPU source-of-truth (`active_request_metadata`) is pinned but resident
+        # on CPU, so reading it here would mix devices with `logits`.
+        gv = context.gpu_view
         if token_to_request_index is None:
-            temperature = md["temperature"][:n]
-            top_k = md["top_k"][:n]
-            top_p = md["top_p"][:n]
+            temperature = gv.temperature[:n]
+            top_k = gv.top_k[:n]
+            top_p = gv.top_p[:n]
         else:
-            temperature = md["temperature"][token_to_request_index]
-            top_k = md["top_k"][token_to_request_index]
-            top_p = md["top_p"][token_to_request_index]
+            temperature = gv.temperature[token_to_request_index]
+            top_k = gv.top_k[token_to_request_index]
+            top_p = gv.top_p[token_to_request_index]
 
         # Clamp temperature to avoid division by 0.
         temperature = temperature.clamp(min=1e-6)

megatron/core/inference/text_generation_controllers/text_generation_controller.py

@@ -1052,7 +1052,7 @@ def _dynamic_step_sample_logits(self):
         gather_indices = (
             None
             if context.config.materialize_only_last_token_logits
-            else context.active_request_last_token_idxs
+            else context.gpu_view.active_request_last_token_idxs
         )
         self._sampled_tokens_cuda = self._sampling.sample_kernel(
             self._all_logits_cuda.squeeze(0),