Add tests

Author: tdene

tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py

@@ -64,6 +64,7 @@ def setup_model(
         sequence_parallel: bool = False,
         expert_model_parallel_size: int = 1,
         num_moe_experts: int = None,
+        sampling_backend: str = 'torch',
     ):
         Utils.initialize_model_parallel(
             tensor_model_parallel_size=tensor_model_parallel_size,
@@ -142,6 +143,7 @@ def setup_model(
                     block_size_tokens=block_size_tokens,
                     enable_prefix_caching=enable_prefix_caching,
                     max_requests=max_requests,
+                    sampling_backend=sampling_backend,
                 ),
             )
 
@@ -249,47 +251,62 @@ def detokenize(self, inp, skip_special_tokens=False):
             sampled_logits >= expected_min_value
         ), f"The sampled logits should all be greater than {expected_min_value} but its {sampled_logits}"
 
-    @pytest.mark.parametrize("backend", ["torch"])
+    @pytest.mark.parametrize("use_cuda_graph", [False, True])
+    @pytest.mark.parametrize("backend", ["torch", "flashinfer"])
     @pytest.mark.parametrize("materialize_only_last_token_logits", [True, False])
     def test_sample_from_dynamic_logits(
-        self, backend: str, materialize_only_last_token_logits: bool
+        self, backend: str, materialize_only_last_token_logits: bool, use_cuda_graph: bool
     ):
-        batch_size = 12
+        if backend == "flashinfer":
+            pytest.importorskip("flashinfer")
+        if use_cuda_graph and backend != "flashinfer":
+            pytest.skip("CUDA graph sampling only applies to the flashinfer backend")
+        batch_size = 15
         self.setup_model(
             torch.float32,
             batch_size=batch_size,
             static=False,
             materialize_only_last_token_logits=materialize_only_last_token_logits,
+            sampling_backend=backend,
         )
         self.mock_tokenizer.eod = self.vocab_size
+        device = torch.cuda.current_device()
 
         context = self.text_generation_controller.inference_wrapped_model.inference_context
+        controller = self.text_generation_controller
 
-        # Prepare sampling params in human-readable format, to aid with test maintenance.
+        # Prepare sampling params.  The torch backend doesn't support simultaneous
+        # top_k + top_p, so the 4th case varies by backend.
         sampling_test_cases: List[Tuple[SamplingParams, List[int]]] = [
             (SamplingParams(temperature=0.1, top_p=0.01), [9, 6, 10]),
             (SamplingParams(temperature=5.0, top_k=15), [0, 3, 2]),
             (SamplingParams(top_p=0.8), [4, 1, 7]),
-            (SamplingParams(temperature=10.0, top_k=5), [11, 5, 8]),
+            (
+                SamplingParams(temperature=10.0, top_k=5, top_p=0.8 if backend != "torch" else 0.0),
+                [11, 5, 8],
+            ),
+            (SamplingParams(), [12, 13, 14]),  # plain (no filtering)
         ]
-        # For non-torch backends, test simultaneous top_k and top_p sampling.
-        if backend != "torch":
-            sampling_test_cases[3][0].top_p = 0.8
 
-        # Convert sampling params to non-readable format.
         rev_sampling_dict: List[SamplingParams] = [None] * batch_size
         for sampling_params, indices in sampling_test_cases:
             for idx in indices:
                 rev_sampling_dict[idx] = sampling_params
 
-        # Prepare metadata for sample bookkeeping.
-        temp_values = torch.Tensor([s.temperature for s in rev_sampling_dict])
-        top_k_values = torch.Tensor([s.top_k for s in rev_sampling_dict]).to(torch.int32)
-        top_p_values = torch.Tensor([s.top_p for s in rev_sampling_dict])
+        # Prepare metadata for sample bookkeeping. The original values are
+        # kept around for the assertion math below; we write the
+        # backend-encoded form into the context (simulating what add_request
+        # would do if we were going through the real request lifecycle).
+        temp_values = torch.tensor([s.temperature for s in rev_sampling_dict])
+        top_k_values = torch.tensor([s.top_k for s in rev_sampling_dict], dtype=torch.int32)
+        top_p_values = torch.tensor([s.top_p for s in rev_sampling_dict])
         context.active_request_metadata["temperature"][:batch_size].copy_(temp_values)
-        context.active_request_metadata["top_k"][:batch_size].copy_(top_k_values)
-        context.active_request_metadata["top_p"][:batch_size].copy_(top_p_values)
-        self.text_generation_controller._sampling_backend = backend
+        context.active_request_metadata["top_k"][:batch_size].copy_(
+            context._encode_sampling_metadata("top_k", top_k_values)
+        )
+        context.active_request_metadata["top_p"][:batch_size].copy_(
+            context._encode_sampling_metadata("top_p", top_p_values)
+        )
 
         context.padded_active_token_count = batch_size
         context.request_query_lengths = torch.ones(batch_size, dtype=torch.int32)
@@ -303,45 +320,152 @@ def test_sample_from_dynamic_logits(
         )
         context.paused_request_count = 0
         context.total_request_count = batch_size
+        context.num_prefill_requests = 0
+        context.padded_active_request_count = batch_size
+        context._using_cuda_graph_this_step = use_cuda_graph
 
-        # Bookkeeping.
-        self.text_generation_controller._dynamic_step_sample_bookkeeping()
-
-        # Sampling.
+        # Set up logits: ascending [0, 1, ..., vocab_size-1] per request.
         logits = torch.arange(0, self.vocab_size).repeat(batch_size, 1).unsqueeze(0).float().cuda()
-        self.text_generation_controller._all_logits_cuda = logits
-        self.text_generation_controller._dynamic_step_sample_logits()
-        sampled_logits = self.text_generation_controller._sampled_tokens_cuda[:batch_size]
-        vocab_indices = torch.arange(self.vocab_size).cuda()
+        if controller._sampling_backend == "flashinfer":
+            # Static buffer; copy in place so the address stays stable for
+            # any graph capture that follows.
+            controller._all_logits_cuda[:, :batch_size, :].copy_(logits)
+        else:
+            controller._all_logits_cuda = logits
+        # _enable_cuda_graph now only gates sampling-graph capture/replay (the
+        # static buffer is allocated regardless); flip it post-init to exercise
+        # graph mode without reconstructing the model with cuda_graph_impl.
+        controller._enable_cuda_graph = use_cuda_graph
+
+        # Two passes so graph mode exercises both capture and replay. Eager mode
+        # runs twice harmlessly.
+        for _ in range(2):
+            controller._dynamic_step_sample_bookkeeping()
+            controller._dynamic_step_sample_logits()
+            sampled = controller._sampled_tokens_cuda[:batch_size]
+
+            # All sampled tokens must be valid vocab indices.
+            assert torch.all(sampled >= 0) and torch.all(
+                sampled < self.vocab_size
+            ), f"Sampled tokens out of range: {sampled}"
+
+            # top_k constraint.
+            top_k_gpu = top_k_values.cuda()
+            top_k_gpu[top_k_gpu == 0] = self.vocab_size
+            assert torch.all(
+                sampled >= self.vocab_size - top_k_gpu
+            ), f"top_k violated: sampled {sampled}, min allowed {self.vocab_size - top_k_gpu}"
+
+            # Combined top_k + top_p constraint: compute the minimum token value
+            # that passes both filters given temperature-scaled probabilities.
+            l = logits.squeeze(0)
+            scaled_probs = l.div(temp_values.cuda().unsqueeze(1)).softmax(dim=-1)
+            vocab_indices = torch.arange(self.vocab_size, device=device)
+            top_k_mask = vocab_indices.unsqueeze(0) < (self.vocab_size - top_k_gpu.unsqueeze(1))
+            scaled_probs.masked_fill_(top_k_mask, 0.0)
+
+            top_p_gpu = top_p_values.cuda()
+            top_p_mask = scaled_probs.cumsum(dim=-1) > top_p_gpu.unsqueeze(1)
+            first_excluded = torch.where(
+                top_p_mask.any(dim=-1),
+                top_p_mask.float().argmax(dim=-1),
+                torch.full((batch_size,), self.vocab_size, device=device),
+            )
+            last_included = torch.clamp(first_excluded - 1, min=0)
+            start_idx = torch.clamp(self.vocab_size - top_k_gpu, min=0).long()
+            last_included = torch.max(last_included, start_idx)
+            expected_min_values = l.gather(1, last_included.unsqueeze(1)).squeeze(1)
+            assert torch.all(
+                sampled >= expected_min_values
+            ), f"Sampled below expected min: sampled {sampled}, expected_min {expected_min_values}"
+
+        if use_cuda_graph:
+            # First pass should have populated the graph cache; second should have replayed.
+            assert len(controller._sampling_cuda_graphs._graphs) >= 1, (
+                "Expected at least one captured sampling graph, got "
+                f"{len(controller._sampling_cuda_graphs._graphs)}"
+            )
 
-        # Move tensors to GPU for assertion checks.
-        temp_values = temp_values.cuda()
-        top_k_values = top_k_values.cuda()
-        top_p_values = top_p_values.cuda()
+    @pytest.mark.parametrize("use_cuda_graph", [False, True])
+    def test_flashinfer_plain_bucket(self, use_cuda_graph: bool):
+        """Exercise the plain (no top_k/top_p) FlashInfer sampling path.
 
-        # Assert correct sampled values.
-        top_k_values[top_k_values == 0] = self.vocab_size
-        assert torch.all(
-            sampled_logits >= self.vocab_size - top_k_values
-        ), f"The sampled logits should all be greater than {self.vocab_size - top_k_values} but its {sampled_logits}"
-        l = logits.squeeze(0)
-        sampled_l = l.div(temp_values.unsqueeze(1)).softmax(dim=-1)
-        top_k_mask = vocab_indices.unsqueeze(0) < (self.vocab_size - top_k_values.unsqueeze(1))
-        sampled_l.masked_fill_(top_k_mask, 0.0)
-        top_p_mask = sampled_l.cumsum(dim=-1) > top_p_values.unsqueeze(1)
-
-        first_excluded = torch.where(
-            top_p_mask.any(dim=-1),
-            top_p_mask.float().argmax(dim=-1),
-            torch.full((batch_size,), self.vocab_size, device=top_p_mask.device),
-        )
-        last_included = torch.clamp(first_excluded - 1, min=0)
-        start_idx = torch.clamp(self.vocab_size - top_k_values, min=0).long()
-        last_included = torch.max(last_included, start_idx)
-        expected_min_values = l.gather(1, last_included.unsqueeze(1)).squeeze(1)
-        assert torch.all(
-            sampled_logits >= expected_min_values
-        ), f"The sampled logits should all be greater than {expected_min_values} but its {sampled_logits}"
+        The mixed-params test above always has at least one filtered request,
+        so _fi_any_filtered_gpu is always True and the plain kernel never
+        runs. This test uses all-default SamplingParams across the batch so
+        the plain kernel graph is exercised.
+        """
+        pytest.importorskip("flashinfer")
+        batch_size = 8
+        self.setup_model(
+            torch.float32,
+            batch_size=batch_size,
+            static=False,
+            materialize_only_last_token_logits=True,
+            sampling_backend="flashinfer",
+        )
+        self.mock_tokenizer.eod = self.vocab_size
+        device = torch.cuda.current_device()
+
+        context = self.text_generation_controller.inference_wrapped_model.inference_context
+        controller = self.text_generation_controller
+
+        # All requests use default SamplingParams (top_k=0, top_p=0.0 -> both
+        # disabled). After add_request-style encoding this becomes the
+        # flashinfer passthrough state where _fi_any_filtered_gpu is False.
+        temp_values = torch.ones(batch_size, dtype=torch.float32)
+        top_k_values = torch.zeros(batch_size, dtype=torch.int32)
+        top_p_values = torch.zeros(batch_size, dtype=torch.float32)
+        context.active_request_metadata["temperature"][:batch_size].copy_(temp_values)
+        context.active_request_metadata["top_k"][:batch_size].copy_(
+            context._encode_sampling_metadata("top_k", top_k_values)
+        )
+        context.active_request_metadata["top_p"][:batch_size].copy_(
+            context._encode_sampling_metadata("top_p", top_p_values)
+        )
+
+        context.padded_active_token_count = batch_size
+        context.request_query_lengths = torch.ones(batch_size, dtype=torch.int32)
+        context.active_request_query_lengths[:batch_size].fill_(1)
+        context.active_request_last_token_idxs[:batch_size].copy_(
+            torch.arange(
+                batch_size,
+                dtype=context.active_request_last_token_idxs.dtype,
+                device=context.active_request_last_token_idxs.device,
+            )
+        )
+        context.paused_request_count = 0
+        context.total_request_count = batch_size
+        context.num_prefill_requests = 0
+        context.padded_active_request_count = batch_size
+        context._using_cuda_graph_this_step = use_cuda_graph
+
+        logits = torch.arange(0, self.vocab_size).repeat(batch_size, 1).unsqueeze(0).float().cuda()
+        controller._all_logits_cuda[:, :batch_size, :].copy_(logits)
+        controller._enable_cuda_graph = use_cuda_graph
+
+        for _ in range(2):
+            controller._dynamic_step_sample_bookkeeping()
+            # With all-default params the filtered check must be False so we
+            # actually hit the plain kernel path.
+            assert not controller._fi_any_filtered_gpu.item(), (
+                "Expected _fi_any_filtered_gpu to be False for all-default params"
+            )
+            controller._dynamic_step_sample_logits()
+            sampled = controller._sampled_tokens_cuda[:batch_size]
+            assert torch.all(sampled >= 0) and torch.all(
+                sampled < self.vocab_size
+            ), f"Sampled tokens out of range: {sampled}"
+
+        if use_cuda_graph:
+            # Exactly the plain graph (filtered=False) should be captured.
+            graphs = controller._sampling_cuda_graphs._graphs
+            assert len(graphs) == 1, f"Expected 1 captured graph (plain), got {len(graphs)}"
+            # Cache key is (n, filtered=False).
+            assert (batch_size, False) in graphs, (
+                f"Expected (batch_size={batch_size}, filtered=False) in graph cache, "
+                f"got keys {list(graphs.keys())}"
+            )
 
     @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
     @pytest.mark.parametrize(