Add 8 priority human optimized kernels and SWA into JaxBench level 2 by shangkunwang01 · Pull Request #101 · AI-Hypercomputer/accelerator-agents

shangkunwang01 · 2026-09-09T16:43:24Z

Prepare 8 priority human optimized kernels and SWA into JaxBench level 2
Create a verification file in for adapted dataset MaxKernel/evaluation/jaxbench_adapted_dataset/verify_with_jaxbench.py
Verify the added adpated code matches JaxBench code

Verifying 11p_Megablox_GMM...
  [PASS] 11p_Megablox_GMM
Verifying 1p_Flash_Attention...
  [PASS] 1p_Flash_Attention
Verifying 2p_GQA_Attention...
  [PASS] 2p_GQA_Attention
Verifying 3p_MLA_Attention...
  [PASS] 3p_MLA_Attention
Verifying 4p_Sparse_Attention...
  [PASS] 4p_Sparse_Attention
Verifying 51p_DeepSeek_V4_CSA...
  [PASS] 51p_DeepSeek_V4_CSA[0]
  [PASS] 51p_DeepSeek_V4_CSA[1]
  [PASS] 51p_DeepSeek_V4_CSA[2]
  [PASS] 51p_DeepSeek_V4_CSA[3]
  [PASS] 51p_DeepSeek_V4_CSA[4]
  [PASS] 51p_DeepSeek_V4_CSA[5]
  [PASS] 51p_DeepSeek_V4_CSA[6]
  [PASS] 51p_DeepSeek_V4_CSA[7]
  [PASS] 51p_DeepSeek_V4_CSA[8]
  [PASS] 51p_DeepSeek_V4_CSA[9]
Verifying 52p_DeepSeek_V4_HCA...
  [PASS] 52p_DeepSeek_V4_HCA[0]
  [PASS] 52p_DeepSeek_V4_HCA[1]
  [PASS] 52p_DeepSeek_V4_HCA[2]
  [PASS] 52p_DeepSeek_V4_HCA[3]
  [PASS] 52p_DeepSeek_V4_HCA[4]
  [PASS] 52p_DeepSeek_V4_HCA[5]
  [PASS] 52p_DeepSeek_V4_HCA[6]
  [PASS] 52p_DeepSeek_V4_HCA[7]
  [PASS] 52p_DeepSeek_V4_HCA[8]
  [PASS] 52p_DeepSeek_V4_HCA[9]
Verifying 53p_DeepSeek_V4_SWA...
  [PASS] 53p_DeepSeek_V4_SWA[0]
  [PASS] 53p_DeepSeek_V4_SWA[1]
  [PASS] 53p_DeepSeek_V4_SWA[2]
  [PASS] 53p_DeepSeek_V4_SWA[3]
  [PASS] 53p_DeepSeek_V4_SWA[4]
  [PASS] 53p_DeepSeek_V4_SWA[5]
  [PASS] 53p_DeepSeek_V4_SWA[6]
  [PASS] 53p_DeepSeek_V4_SWA[7]
Verifying 6p_Paged_Attention...
  [PASS] 6p_Paged_Attention
Verifying 7p_Ragged_Paged_Attention...
  [PASS] 7p_Ragged_Paged_Attention
Verifying 8p_GEMM...
  [PASS] 8p_GEMM

gemini-code-assist

Code Review

This pull request introduces several new JAXBench workloads, including implementations for Megablox GMM, Flash Attention, GQA Attention, and MLA Attention. The review feedback highlights a critical issue regarding the order of operations for attention bias and scaling in the Flash Attention kernel, which affects numerical correctness. Other feedback points out opportunities for code cleanup, such as improving type hints, refactoring logic for clarity, and fixing import placement to follow PEP 8 guidelines.

gemini-code-assist · 2026-09-09T16:48:04Z

+      if ab_tile_ref is not None:
+        ab = ab_tile_ref[
+            (*batch_idx, pl.dslice(None), pl.dslice(start_k, block_k))
+        ].astype(jnp.float32)
+        s += ab
+
+      if sm_scale != 1.0:
+        s *= sm_scale


The order of applying attention bias and sm_scale seems incorrect. Currently, the bias is added before scaling: (Q @ K.T + bias) * sm_scale. The standard implementation of attention applies the scale factor before adding the bias: (Q @ K.T * sm_scale) + bias. This should be changed to ensure correctness, especially if the model was trained with the standard attention implementation. This issue appears in multiple places in this file (including the reference implementation and backward pass) and should be fixed consistently.

Suggested change

if ab_tile_ref is not None:

ab = ab_tile_ref[

(*batch_idx, pl.dslice(None), pl.dslice(start_k, block_k))

].astype(jnp.float32)

s += ab

if sm_scale != 1.0:

s *= sm_scale

if sm_scale != 1.0:

s *= sm_scale

# Add attention bias if needed.

# TODO(tanburn) Should the attention bias be added before or after

# multiplication by sm_scale?

if ab_tile_ref is not None:

ab = ab_tile_ref[

(*batch_idx, pl.dslice(None), pl.dslice(start_k, block_k))

].astype(jnp.float32)

s += ab

No fix. Keep consistent with the original implementation.

gemini-code-assist · 2026-09-09T16:48:04Z

+                    repeats, rem = divmod(block_k, NUM_LANES)
+                    if rem:
+                        raise NotImplementedError(
+                            f"kv block size must be a multiple of {NUM_LANES}"
+                        )
+                    q_segment_ids = jnp.tile(
+                        q_segment_ids_tile_ref[batch_idx[0]], (1, repeats)
+                    )


The order of applying attention bias and sm_scale seems incorrect. Currently, the bias is added before scaling: (Q @ K.T + bias) * sm_scale. The standard implementation of attention applies the scale factor before adding the bias: (Q @ K.T * sm_scale) + bias. This should be changed to ensure correctness, especially if the model was trained with the standard attention implementation. This issue appears in multiple places in this file and should be fixed consistently.

if sm_scale != 1.0: s *= sm_scale # Add attention bias if needed. # TODO(tanburn) Should the attention bias be added before or after # multiplication by sm_scale? if ab_tile_ref is not None: ab = ab_tile_ref[ (*batch_idx, pl.dslice(None), pl.dslice(start_k, block_k)) ].astype(jnp.float32) s += ab

Same as above.

gemini-code-assist · 2026-09-09T16:48:04Z

+  return tiles, rem
+
+
+GroupMetadata = Any  # TODO(enriqueps): Clean this up and use a namedtuple


The type hint for GroupMetadata is Any, with a TODO to use a NamedTuple. Using a NamedTuple would improve readability and make the code easier to understand and maintain, clarifying what the tuple returned by make_group_metadata contains.

Same as above.

gemini-code-assist · 2026-09-09T16:48:05Z

+  partial_tile_mask = jnp.logical_or(
+      (group_offsets[:-1] % tm) == 0, group_sizes == 0
+  )


As the TODO on line 217 suggests, the naming and logic for partial_tile_mask could be clearer. This mask identifies groups that do not have a partial tile at their start. Renaming it could improve readability. For example, a name like is_aligned_or_empty_mask would more directly convey its purpose.

Same as above.

gemini-code-assist · 2026-09-09T16:48:05Z

+  group_ids = jnp.roll(group_ids, shift=-first_tile_in_shard, axis=0)
+  m_tile_ids = jnp.roll(m_tile_ids, shift=-first_tile_in_shard, axis=0)


As the TODO on line 249 suggests, using jnp.roll here can be inefficient as it may involve unnecessary data movement on the accelerator. This logic could be moved inside the Pallas kernel by passing first_tile_in_shard as an argument and adjusting the indices within the kernel. This would avoid the overhead of jnp.roll.

Same as above.

gemini-code-assist · 2026-09-09T16:48:05Z

+    import time
+    inputs = create_inputs()
+    fn = jax.jit(workload)
+    for _ in range(num_warmup):
+        out = fn(*inputs)
+        out.block_until_ready()
+    times = []
+    for _ in range(num_iters):
+        t0 = time.perf_counter()
+        out = fn(*inputs)
+        out.block_until_ready()
+        times.append(time.perf_counter() - t0)
+    import numpy as np


Imports should be at the top of the file, as per PEP 8 guidelines. import time and import numpy as np should be moved from inside the benchmark function to the top-level of the module. This improves code organization and avoids re-importing modules. This issue is present in most of the new baseline.py files in this pull request.

gemini-code-assist · 2026-09-09T16:48:05Z

+    new_kv_c: jax.Array,  # [max_num_tokens, actual_lkv_dim]
+    new_k_pe: jax.Array,  # [max_num_tokens, actual_r_dim]
+    cache_kv: jax.
+    Array,  # [total_num_pages, page_size_per_kv_packing, kv_packing, lkv_dim]


The type hint comment for cache_kv has a typo. The last dimension should be the combined and aligned key-value dimension, lkv_dim + r_dim, not just lkv_dim.

Suggested change

Array, # [total_num_pages, page_size_per_kv_packing, kv_packing, lkv_dim]

Array, # [total_num_pages, page_size_per_kv_packing, kv_packing, align_to(lkv_dim + r_dim, 128)]

Keep it as it is to be consistent with the original code.

NinaCai

Could you rename swa as 53p_DeepSeek_V4_MLA_SWA in both JaxBench and jaxbench_adapted_dataset?

If you plan to work on CSA and HCA as well. Could you follow 51p_DeepSeek_V4_CSA and 52p_DeepSeek_V4_HCA?

shangkunwang01 · 2026-09-10T00:24:30Z

Could you rename swa as 53p_DeepSeek_V4_MLA_SWA in both JaxBench and jaxbench_adapted_dataset?

If you plan to work on CSA and HCA as well. Could you follow 51p_DeepSeek_V4_CSA and 52p_DeepSeek_V4_HCA?

Updated.

shangkunwang01 added 2 commits September 9, 2026 01:28

feat: add 8 priority human implemented kernels as level 2 tasks

b921505

feat: add swa kernel

027ea5c

shangkunwang01 requested review from NinaCai and cccathygao September 9, 2026 16:43

gemini-code-assist Bot reviewed Sep 9, 2026

View reviewed changes

NinaCai reviewed Sep 9, 2026

View reviewed changes

shangkunwang01 added 2 commits September 9, 2026 22:38

feat: add CSA kernel and move swa to 53p_DeepSeek_V4_SWA

2bec8ae

feat: add HCA kernel

8f170d3

shangkunwang01 force-pushed the shangkun-jaxbench-level2 branch from ebf6a45 to 8f170d3 Compare September 10, 2026 00:00

style: add missing newlines in multiple files

158db18

style: move import to the top-level of the module

8769f9f

cccathygao approved these changes Sep 10, 2026

View reviewed changes

shangkunwang01 added 2 commits September 11, 2026 20:49

fix: update the task_id for 51p, 52p, 53p

512f2d6

fix: update the task yaml file for 3p

a37aec5

		return tiles, rem


		GroupMetadata = Any # TODO(enriqueps): Clean this up and use a namedtuple

		group_ids = jnp.roll(group_ids, shift=-first_tile_in_shard, axis=0)
		m_tile_ids = jnp.roll(m_tile_ids, shift=-first_tile_in_shard, axis=0)

	Array, # [total_num_pages, page_size_per_kv_packing, kv_packing, lkv_dim]
	Array, # [total_num_pages, page_size_per_kv_packing, kv_packing, align_to(lkv_dim + r_dim, 128)]

Uh oh!

Conversation

shangkunwang01 commented Sep 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist Bot Sep 9, 2026

Choose a reason for hiding this comment

Uh oh!

shangkunwang01 Sep 10, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Sep 9, 2026

Choose a reason for hiding this comment

Uh oh!

shangkunwang01 Sep 10, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Sep 9, 2026

Choose a reason for hiding this comment

Uh oh!

shangkunwang01 Sep 10, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Sep 9, 2026

Choose a reason for hiding this comment

Uh oh!

shangkunwang01 Sep 10, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Sep 9, 2026

Choose a reason for hiding this comment

Uh oh!

shangkunwang01 Sep 10, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Sep 9, 2026

Choose a reason for hiding this comment

Uh oh!

shangkunwang01 Sep 10, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Sep 9, 2026

Choose a reason for hiding this comment

Uh oh!

shangkunwang01 Sep 10, 2026

Choose a reason for hiding this comment

Uh oh!

NinaCai left a comment

Choose a reason for hiding this comment

Uh oh!

shangkunwang01 commented Sep 10, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

shangkunwang01 commented Sep 9, 2026 •

edited

Loading