Triton Seeded Fused Dropout
Problem Statement
Implement inverted dropout in a single kernel using Triton's counter-based RNG: drop each element with probability p, scale survivors by 1/(1-p). Use a runtime seed so the result is reproducible.
Background
tl.rand(seed, offsets) produces deterministic uniform values in [0, 1) per offset. Keep elements where rand > p. Because the RNG is seeded, two launches with the same seed must be identical, and survivors must be exactly x/(1-p).
Your Task
Implement dropout_kernel and run(n=4096, p=0.5, seed=123) that returns True when (a) two same-seed launches match and (b) surviving values equal x/(1-p).
How it is tested
Your solution must define a top-level function run(...) that allocates inputs on the GPU, launches your Triton kernel, and returns a boolean from torch.allclose(triton_out, torch_reference, ...). The grader prints run(...); the expected output is True.
Example:
n = 4096, p = 0.5, seed = 123
True
- The input values are
n = 4096,p = 0.5, andseed = 123, which are used to allocate inputs on the GPU and launch thedropout_kernel. - The
dropout_kerneluses Triton's counter-based RNG to generate uniform values in [0, 1) per offset, and drops each element with probabilityp = 0.5, scaling survivors by 1/(1−p)=1/(1−0.5)=2. - The kernel is launched twice with the same seed, ensuring that the results are identical due to the deterministic nature of the RNG.
- The final output is
Truebecause the two same-seed launches match and the surviving values equalx/(1-p), as verified bytorch.allclose(triton_out, torch_reference, ...).
Constraints:
- Use tl.rand(seed, offsets) for per-element randomness
- keep = rand > p; out = where(keep, x/(1-p), 0)
- Same seed must give identical output (verified in run)
Background Knowledge
Introduction to Dropout
Dropout is a regularization technique used in deep learning to prevent overfitting. It works by randomly dropping out (setting to zero) a fraction of the neurons during training, which helps the model learn more robust features. In this problem, we are implementing inverted dropout, where the surviving elements are scaled by a factor of 1/(1−p), where p is the dropout probability.
Understanding Triton's Counter-Based RNG
Triton's counter-based RNG is a random number generator that produces deterministic uniform values in the range [0, 1) per offset. This means that given a seed and an offset, the RNG will always produce the same random number. This property is useful for ensuring reproducibility in our implementation.
Importance of Reproducibility
Reproducibility is crucial in machine learning, as it allows us to verify the correctness of our implementation and ensure that the results are consistent across different runs. In this problem, we are required to implement a runtime seed that ensures the result is reproducible, meaning that two launches with the same seed must produce identical results.
Algorithm/Approach
The general approach to solving this problem involves implementing a kernel that performs the following steps:
- Generate random numbers using Triton's counter-based RNG
- Compare the generated random numbers with the dropout probability p to determine which elements to drop
- Scale the surviving elements by a factor of 1/(1−p)
- Ensure that the implementation is reproducible by using a runtime seed
Continue the full explanation
You're reading the free preview. Unlock the complete walkthrough, the code editor, test runner and reference solution with Premium.
Editor locked
The code editor is locked for Pro problems. It is only available for free problems. Please upgrade to gain access to the code editor for all problems.