Implement all 5 Triton kernel labs

- vector_add: basic masked load/store with block indexing
- row_softmax: single-pass numerically stable softmax per row
- tiled_matmul: K-dimension tile loop with edge masking (IEEE precision)
- online_softmax: two-pass running max/sum recurrence across blocks
- flash_attention_fwd: blockwise Q/K/V with online softmax, causal support

All 26 tests pass on RTX 5090 (CUDA 12.8, Triton 3.6).
This commit is contained in:
2026-05-15 20:46:04 +08:00
parent 7fa69b1354
commit 165a1b0bd5
5 changed files with 170 additions and 31 deletions

View File

@@ -26,10 +26,10 @@ if TRITON_AVAILABLE:
pid = tl.program_id(axis=0)
offsets = pid * block_size + tl.arange(0, block_size)
mask = offsets < num_elements
# TODO(student): load x and y using masked tl.load calls.
# TODO(student): add the vectors.
# TODO(student): write the result with tl.store.
pass
x = tl.load(x_ptr + offsets, mask=mask)
y = tl.load(y_ptr + offsets, mask=mask)
out = x + y
tl.store(out_ptr + offsets, out, mask=mask)
def triton_vector_add(x: torch.Tensor, y: torch.Tensor, block_size: int = 1024) -> torch.Tensor:
@@ -40,5 +40,9 @@ def triton_vector_add(x: torch.Tensor, y: torch.Tensor, block_size: int = 1024)
raise ValueError(f"shape mismatch: {x.shape} vs {y.shape}")
if not x.is_cuda or not y.is_cuda:
raise ValueError("Triton kernels in this lab expect CUDA tensors.")
raise NotImplementedError("TODO(student): launch vector_add_kernel and return the output tensor.")
out = torch.empty_like(x)
num_elements = x.numel()
grid = ((num_elements + block_size - 1) // block_size,)
vector_add_kernel[grid](x, y, out, num_elements, block_size=block_size)
return out