Denton's code

[zzczzc20]

No description provided.

[mark14wu]

@codex review.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 9083918640

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[mark14wu]

Code review

Found 8 issues:

1. Wrong stride used for V block pointer in flash attention Triton kernel. V_block_ptr uses stride_k_seqlen (K's sequence stride) instead of stride_v_seqlen for V's row stride. Works by coincidence when K/V have identical layouts, but will silently produce wrong results if they differ.

Tilebench/benchmarks/operators/flash_attention/impl_triton.py

Lines 40 to 46 in 9083918

	V_block_ptr = tl.make_block_ptr(
	base=V + qkv_base_offset,
	shape=(SEQLEN, DIM),
	strides=(stride_k_seqlen, stride_v_dim),
	offsets=(0, 0),
	block_shape=(BLOCK_N, DIM),
	order=(1, 0),

2. Q's head stride used for K and V base offsets in flash attention Triton kernel. qkv_base_offset = off_bs_head * stride_q_head is applied to all three block pointers, but stride_k_head and stride_v_head are accepted as kernel parameters and never used. Will read K/V from wrong memory for GQA/MQA where head strides differ.

Tilebench/benchmarks/operators/flash_attention/impl_triton.py

Lines 22 to 42 in 9083918

	qkv_base_offset = off_bs_head * stride_q_head
	Q_block_ptr = tl.make_block_ptr(
	base=Q + qkv_base_offset,
	shape=(SEQLEN, DIM),
	strides=(stride_q_seqlen, stride_q_dim),
	offsets=(start_m * BLOCK_M, 0),
	block_shape=(BLOCK_M, DIM),
	order=(1, 0),
	)
	K_block_ptr = tl.make_block_ptr(
	base=K + qkv_base_offset,
	shape=(DIM, SEQLEN),
	strides=(stride_k_dim, stride_k_seqlen),
	offsets=(0, 0),
	block_shape=(DIM, BLOCK_N),
	order=(0, 1),
	)
	V_block_ptr = tl.make_block_ptr(
	base=V + qkv_base_offset,
	shape=(SEQLEN, DIM),

3. Global verification tolerance silently doubled. atol and rtol changed from 1e-3 to 2e-3 with no justification. The engine always calls verify() with defaults, so this makes all operator verification twice as permissive, potentially masking correctness regressions.

Tilebench/core/verifier.py

Lines 2 to 4 in 9083918

	def verify(output, reference, atol=2e-3, rtol=2e-3):
	try:

4. Stream-K matmul locks tensor can be size 0. locks = torch.zeros((total_tiles_streamk,), ...) creates a zero-element tensor when total_tiles_streamk == 0, but the kernel still attempts atomic operations on it. The cuTile version guards with max(1, total_tiles_streamk) but the Triton version does not.

Tilebench/benchmarks/operators/streamk_matmul/impl_triton.py

Lines 181 to 183 in 9083918

	c = torch.empty((M, N), device=device, dtype=a.dtype)
	locks = torch.zeros((total_tiles_streamk,), device=device, dtype=torch.int32)

5. grid_programs config parameter silently dropped. config.yaml specifies grid_programs: 108 but engine.py only forwards block_size to run(). The Triton impl defaults to 108 (correct by coincidence), but cuTile defaults to 16 -- producing wrong work partitioning on A100-class hardware.

Tilebench/benchmarks/operators/streamk_matmul/impl_cutile.py

Lines 150 to 153 in 9083918

	def run(a: torch.Tensor, b: torch.Tensor, block_size: int = None,
	grid_programs: int = 16, # Adjust based on target GPU SM count (e.g. 108 for A100)
	BLK_M: int = 128, BLK_N: int = 128, BLK_K: int = 32,
	two_tiles: bool = True):

Tilebench/core/engine.py

Lines 44 to 55 in 9083918

	# 3. Triton run
	triton_output = impl_triton.run(*inputs, block_size=block_size)
	torch.cuda.synchronize()
	triton_ok, triton_err = verify(triton_output, ref_output)
	if not triton_ok:
	print(f" Triton verification FAILED: {triton_err}")
	triton_ms = report_benchmark(impl_triton.run, inputs, kwargs={'block_size': block_size})['mean_time_ms'] if triton_ok else float('nan')
	# 4. cuTile run
	try:
	cutile_output = impl_cutile.run(*inputs, block_size=block_size)
	torch.cuda.synchronize()

6. Debug tl.static_print left in block sparse attention kernel. No other Triton kernel in the repo contains debug prints. This will print on every kernel compilation.

Tilebench/benchmarks/operators/block_sparse_attention/impl_triton.py

Lines 38 to 40 in 9083918

	):
	tl.static_print(f"{BLOCK_M=} {BLOCK_N=} {BLOCK_D=} {EVEN_M=} {EVEN_N=} {NUM_D_BLOCKS=}")

7. causal config parameter absorbed but never forwarded to run(). flash_attention/config.yaml specifies causal: True, but generate_flash_attn_inputs() accepts it in **kwargs and discards it. All three impls hardcode causal=True as default, so it works by coincidence. Setting causal: False in config would be silently ignored.

Tilebench/data/tensors.py

Lines 37 to 44 in 9083918

	return (x,)
	def generate_flash_attn_inputs(batch_size, n_heads, seq_len, head_dim, dtype=torch.float16, device='cuda', **kwargs):
	q = torch.randn(batch_size, n_heads, seq_len, head_dim, dtype=dtype, device=device)
	k = torch.randn(batch_size, n_heads, seq_len, head_dim, dtype=dtype, device=device)
	v = torch.randn(batch_size, n_heads, seq_len, head_dim, dtype=dtype, device=device)
	return (q.contiguous(), k.contiguous(), v.contiguous())

Tilebench/benchmarks/operators/flash_attention/config.yaml

Lines 7 to 9 in 9083918

	dtype: "float16"
	causal: True

8. Torch run() signatures missing block_size parameter. softmax/impl_torch.py, rope/impl_torch.py, matmul_int8/impl_torch.py, and matmul_fp16_fp8/impl_torch.py all omit block_size from their run() signatures. However, engine.py only passes block_size to Triton/cuTile (not torch), so this is a latent mismatch rather than a current crash -- but the inconsistency should be addressed.

Tilebench/benchmarks/operators/softmax/impl_torch.py

Lines 2 to 4 in 9083918

	def run(x):
	return torch.softmax(x, dim=-1)

Tilebench/core/engine.py

Lines 39 to 55 in 9083918

	# 2. Reference run (Torch)
	ref_output = impl_torch.run(*inputs)
	torch.cuda.synchronize()
	torch_ms = report_benchmark(impl_torch.run, inputs)['mean_time_ms']
	# 3. Triton run
	triton_output = impl_triton.run(*inputs, block_size=block_size)
	torch.cuda.synchronize()
	triton_ok, triton_err = verify(triton_output, ref_output)
	if not triton_ok:
	print(f" Triton verification FAILED: {triton_err}")
	triton_ms = report_benchmark(impl_triton.run, inputs, kwargs={'block_size': block_size})['mean_time_ms'] if triton_ok else float('nan')
	# 4. cuTile run
	try:
	cutile_output = impl_cutile.run(*inputs, block_size=block_size)
	torch.cuda.synchronize()

🤖 Generated with Claude Code

_{- If this code review was useful, please react with 👍. Otherwise, react with 👎.}