Education

• 
  Nanyang Technological University

  Aug. 2020 - Jan. 2025

  College of Computing and Data Science

  B.Comp. (Highest Distinction), CGPA 4.71/5.00

  • Dean's List (2022–2023); fully funded by the Ministry of Education, Singapore (SM2 Scholarship and Tuition Grant)
  • Elective focuses in Artificial Intelligence and in Data Science & Analytics
• 
  National University of Singapore

  Jan. 2022 - May 2022

  School of Computing

  Student Exchange Program

Publications ^*first author   corresponding author

SC '26

HieraSparse: Hierarchical Semi-Structured Sparse KV Attention

Haoxuan Wang^*, Chen Wang

The International Conference for High Performance Computing, Networking, Storage, and Analysis (SC '26) 2026

[Paper] [Code]

SC '26

HieraSparse: Hierarchical Semi-Structured Sparse KV Attention

Haoxuan Wang^*, Chen Wang

The International Conference for High Performance Computing, Networking, Storage, and Analysis (SC '26) 2026

[Paper] [Code]

Research Experience

• 
  Parallel Systems Lab, Nanyang Technological University

  Nov. 2025 - present

  Research Assistant (Advisor - Prof. Chen Wang)

  • Proposed HieraSparse (SC '26), a 2:4 semi-structured sparse KV attention accelerated by Sparse Tensor Cores
  • Built the lab's GPU cluster (up to 5 nodes) from bare components and interconnected the nodes with InfiniBand through an NVIDIA Quantum-2 switch for NFS and MPI
  • Benchmarked the school's shared GPU infrastructure (~200 H200 GPUs) on typical ML workloads with Megatron-LM, TransformerEngine, and TensorRT-LLM
• 
  HPC-AI Lab, National University of Singapore

  Sep. 2023 - Jan. 2025

  Student Assistant (Advisor - Prof. Yang You)

  • Developed EasyDist with Ph.D. students and the Alibaba Cloud PAI team, designing and implementing PyTorch dispatcher extensions, graph capture and transformation, an ILP solver for optimal SPMD placement, and pipeline parallelism integration

Industry Experience

• 
  HPC-AI Technology (ColossalAI), Singapore

  Infrastructure Engineer (full-time)

  Mar. 2025 - Oct. 2025

  Infrastructure Intern (part-time)

  May 2024 - Nov. 2024

  • Contributed to Colossal-AI, including ZeRO optimizations, FP8 quantization-aware training, MoE training, and hybrid parallelism (PP/DP/EP/SP/TP)
  • Optimized diffusion transformer training and inference with techniques such as sparse attention, timestep embedding caching, and quantization, tuned for GPUs including RTX 4090, H200, and B200
• 
  TikTok Singapore

  Jan. 2024 - May 2024

  Applied Machine Learning Intern (part-time)

  • Built a health-monitoring and alerting system for GPU clusters using Grafana and Flask, and integrated it with a Lark chatbot to send real-time alerts to the on-call team
• 
  Huawei 2012 Laboratories (Shield Lab), Singapore

  Jul. 2023 - Dec. 2023

  Intern (full-time)

  • Built a pipeline to generate adversarial prompt datasets by augmenting public data with attack schemes such as role-play, code injection, and goal hijacking
  • Designed a multi-stage adversarial prompt detection pipeline combining a lexical sensitive-word filter with a lightweight semantic classifier (TextCNN over PCA-reduced word embeddings)

Selected Projects

• 
  HieraSparse Author

  Hierarchical KV-cache compression with semi-structured sparse attention for long-context LLMs. SC '26, with all three ACM reproducibility badges.

  • Designed a hierarchical KV-cache compression scheme with flexible quality–sparsity trade-offs, and wrote the prefill and decode attention kernels that run it on Sparse Tensor Cores
  • 4.57× attention speedup and 1.2× higher KV compression than state-of-the-art unstructured sparsity at the same sparsity level; up to 1.85× attention speedup in prefill
  • With simple magnitude-based pruning, 1.34× prefill and 1.71× decode attention speedup with minimal quality loss
• 
  EasyDist Core contributor PRs

  A graph compiler for automatically parallelizing PyTorch model training across multiple GPUs.

  • Hacked torch.dynamo and dispatcher to capture the computation graph of DL models and optimizers, inserting collective communication via graph transformation and generating distributed code
  • Designed and implemented a sharding solver that explores distributed tensor placements (sharded, replicated, partial) over an n-dimensional GPU mesh and finds the optimal assignment via ILP (OR-Tools)
  • Designed and implemented pipeline parallelism composable with SPMD parallelism, supporting GPipe, 1F1B, and asynchronous 1F1B schedules
• 
  Colossal-AI Core contributor PRs

  A PyTorch-based system for training and serving large models across many GPUs.

  • Optimized the ZeRO optimizer by prefetching parameter/gradient chunks and overlapping collective and PCIe communication (when offloading) with forward/backward computation and optimizer updates
  • Designed and implemented MoE expert parallelism with intra- and inter-node all-to-all, integrated with data parallelism, tensor parallelism, and the distributed ZeRO optimizer
  • Supported FP8 quantization-aware training for both collective communication (all-to-all and all-gather of quantized tensors and their scales) and computation (custom kernels and cuBLAS integration)
• 
  TensorNVMe Core contributor PRs

  A Python library that transfers PyTorch tensors between GPU, CPU, and NVMe.

  • Designed and implemented a custom thread-pool backend for asynchronous GPU/CPU/disk tensor transfer, outperforming the io_uring and aio backends through lower overhead and support for D2H copies
  • Added safetensors support with a custom C++ writer (ported from the original Rust implementation), avoiding arbitrary code execution and staying compatible with most training and inference frameworks
• 
  Open-Sora Contributor

  An open-source video-generation model with its training and inference stack. Contributions were made in a private mirror of the repository.

  • Integrated asynchronous checkpointing via TensorNVMe into training, fully overlapping checkpoint-saving latency (both D2H copy and disk write) with computation
  • Optimized model inference by adapting sparse attention (Sliding-Tile-Attention), timestep embedding caching (TeaCache), and quantization (SageAttention, SVDQuant) to Open-Sora
• 
  TileLang Community contributor PRs

  A DSL designed to streamline the development of high-performance GPU/CPU/accelerator kernels.

  • Implemented T.gemm_sp, a 2:4 semi-structured GEMM tile operator generated from CuTe/CUTLASS templates, achieving 1.5–2× speedup over dense GEMM on NVIDIA A100, A40, L40S, H20, and H200
  • Later extended this contribution into HieraSparse (SC '26), as the sole student author
• 
  LeetCUDA Community contributor PRs

  A collection of CUDA learning notes and optimized kernels.

  • Contributed CuTe implementation examples, including matrix transpose, GEMM, and FlashAttention